lfm2_diffusion.py

"""lfm2_diffusion.py — CPU (PyTorch) masked-diffusion inference for LFM2.5 Encoder.

Server-side replacement for the old in-browser WebGPU path. The model, decode
loop, and numerics are a 1:1 port of the reference JS (lfm2_webgpu_forward.js /
lfm2_bidir_mdlm.js): a *bidirectional* LFM2 forward, then the block-scheduled,
confidence-driven MDLM unmasking loop.

Speed comes from three things:
  * fast-dLLM "DualCache" (as in the WebGPU build): the first step of each block
    runs a full-canvas forward that fills a per-layer K/V cache (attention) and Bx
    cache (shortconv); every later step of that block only recomputes the active
    ~block-size slice, attending to / convolving against the frozen cache.
  * smoothed int8 dynamic quantization of every matmul (see `quantize_int8`).
  * one worker thread per available core (cgroup-aware, no oversubscription).

The checkpoint is pulled from the Hub model repo at boot (no local copy);
everything runs on CPU so nothing is downloaded to the client.
"""

import os
import math
import re

import torch
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

# The perfectblend diffusion-SFT checkpoint (f32 safetensors). It is the encoder
# turned into an MDLM chatbot — the *base* LFM2.5-Encoder-350M is a plain MLM and
# echoes instead of answering, so it is not interchangeable here.
MODEL_ID = os.environ.get("DIFFUSION_MODEL", "LiquidAI/LFM2.5-Encoder-350M-Diffusion")
WEIGHTS_FILE = os.environ.get("DIFFUSION_WEIGHTS", "model.safetensors")

# --- architecture constants (LFM2.5-Encoder-350M), mirrored from the JS ------
H = 1024            # hidden size
NL = 16             # layers
NH = 16             # attention heads
NKV = 8             # kv heads (GQA group size 2)
HD = 64             # head dim
FF = 4608           # ffn inner size
EPS = 1e-5
ROPE_THETA = 1e6
ATTN_LAYERS = frozenset({2, 5, 8, 10, 12, 14})   # the rest are shortconv layers

# --- MDLM / tokenizer constants (see the model card) -------------------------
MASK_ID = 16        # <|mask|>
REAL_VOCAB = 64402  # ignore logit ids >= REAL_VOCAB (padded vocab is 65536)

# --- int8 smoothing (see Lfm2Diffusion.quantize_int8) ------------------------
SMOOTH_ALPHA = 0.5
# Fixed calibration set, only used to measure per-channel activation ranges (one
# forward each, ~1s at boot). app.py tokenizes these.
CALIB_PROMPTS = (
    "What is the capital of France?",
    "Write a haiku about the ocean.",
    "Give me two tips for better sleep.",
    "What is 12 times 8?",
)


def effective_cpus():
    """Cores actually available to this process (cgroup quota, not the host)."""
    try:  # cgroup v2
        raw = open("/sys/fs/cgroup/cpu.max").read().split()
        if raw and raw[0] != "max":
            return max(1, round(int(raw[0]) / int(raw[1])))
    except (OSError, ValueError):
        pass
    try:  # cgroup v1
        q = int(open("/sys/fs/cgroup/cpu/cpu.cfs_quota_us").read())
        p = int(open("/sys/fs/cgroup/cpu/cpu.cfs_period_us").read())
        if q > 0:
            return max(1, q // p)
    except (OSError, ValueError):
        pass
    try:
        return max(1, len(os.sched_getaffinity(0)))
    except AttributeError:
        return os.cpu_count() or 1


def render_chat_prompt(user_text, system_text=None):
    """LFM2 chat template for a single user turn + add_generation_prompt."""
    s = ""
    if system_text:
        s += f"[SYS]\n{system_text.strip()}\n[/SYS]\n\n"
    s += f"[Question]\n{user_text.strip()}\n[/Question]\n\n[Answer]\n"
    return s


def trim(text):
    """Cut generated text at the answer terminator / role markers (app.py trim)."""
    for stop in ("[/Answer]", "[Question]", "[SYS]"):
        i = text.find(stop)
        if i != -1:
            text = text[:i]
    text = text.strip()
    text = re.sub(r"^(?:\[Answer\]|Answer\b)\s*:?\s*", "", text, flags=re.I)
    text = re.sub(r"\s*(?:\.{3}|…)\s*$", "", text)
    return text.strip()


def load_hf_tensors(path):
    """Load the safetensors checkpoint — weights are already f32 and [out, in], so
    there is nothing to dequantize or transpose. The dimension check makes a
    checkpoint that doesn't match the constants above fail loudly at boot rather
    than as silent garbage."""
    t = load_file(path)
    attn = frozenset(int(m.group(1)) for k in t
                     if (m := re.match(r"lfm2\.layers\.(\d+)\.self_attn\.", k)))
    a = min(attn)
    got = (t["lfm2.embed_tokens.weight"].shape[1],
           1 + max(int(m.group(1)) for k in t if (m := re.match(r"lfm2\.layers\.(\d+)\.", k))),
           attn,
           t["lfm2.layers.0.feed_forward.w1.weight"].shape[0],
           t[f"lfm2.layers.{a}.self_attn.q_proj.weight"].shape[0],
           t[f"lfm2.layers.{a}.self_attn.k_proj.weight"].shape[0])
    want = (H, NL, ATTN_LAYERS, FF, NH * HD, NKV * HD)
    if got != want:
        raise ValueError(f"{path}: (hidden, layers, attn layers, ffn, q, kv) is {got}, expected {want}")
    return t


class RowScaledLinear(nn.Module):
    """Per-token activations around a dynamically quantized Linear.

    The int8 kernel derives ONE activation scale per call, so a single loud token
    coarsens every other row in the batch. Normalising each row to unit max first
    (and putting the scale back on the output) makes that one scale exactly right
    for every row, which is per-token quantization at the cost of two elementwise
    passes."""

    def __init__(self, qlin):
        super().__init__()
        self.q = qlin

    def forward(self, x):
        r = x.abs().amax(-1, keepdim=True).clamp_min(1e-6)
        return self.q(x / r) * r


def rms_norm(x, w, eps=EPS):
    inv = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)
    return x * inv * w


def _linear(in_f, out_f, weight):
    m = nn.Linear(in_f, out_f, bias=False)
    m.weight.data = weight.contiguous()
    return m


class Lfm2Layer(nn.Module):
    def __init__(self, t, li):
        super().__init__()
        p = f"lfm2.layers.{li}."
        self.attn = li in ATTN_LAYERS
        self.register_buffer("op_norm", t[p + "operator_norm.weight"], persistent=False)
        self.register_buffer("ffn_norm", t[p + "ffn_norm.weight"], persistent=False)
        # fused gate+up: one GEMM instead of two
        self.w13 = _linear(H, 2 * FF, torch.cat(
            [t[p + "feed_forward.w1.weight"], t[p + "feed_forward.w3.weight"]], 0))
        self.w2 = _linear(FF, H, t[p + "feed_forward.w2.weight"])
        if self.attn:
            self.qkv = _linear(H, H + 2 * NKV * HD, torch.cat(
                [t[p + "self_attn.q_proj.weight"], t[p + "self_attn.k_proj.weight"],
                 t[p + "self_attn.v_proj.weight"]], 0))
            self.o = _linear(H, H, t[p + "self_attn.out_proj.weight"])
            self.register_buffer("q_norm", t[p + "self_attn.q_layernorm.weight"], persistent=False)
            self.register_buffer("k_norm", t[p + "self_attn.k_layernorm.weight"], persistent=False)
        else:
            self.register_buffer("conv", t[p + "conv.conv.weight"].reshape(H, 1, -1), persistent=False)
            self.in_proj = _linear(H, 3 * H, t[p + "conv.in_proj.weight"])
            self.out_proj = _linear(H, H, t[p + "conv.out_proj.weight"])


class Lfm2Diffusion(nn.Module):
    """Takes ownership of `tensors`: most weights are used without copying, and
    quantize_int8 rescales them in place."""

    def __init__(self, tensors):
        super().__init__()
        t = tensors
        self.register_buffer("emb", t["lfm2.embed_tokens.weight"].contiguous(), persistent=False)
        self.register_buffer("emb_norm", t["lfm2.embedding_norm.weight"], persistent=False)
        # tied logits, trimmed to the real vocab (ids >= REAL_VOCAB are suppressed
        # anyway, so there's no point computing the padded tail every step). Cloned,
        # not shared with `emb`: quantize_int8 rescales the head's columns in place.
        self.lm_head = _linear(H, REAL_VOCAB, t["lfm2.embed_tokens.weight"][:REAL_VOCAB].clone())
        self.layers = nn.ModuleList([Lfm2Layer(t, li) for li in range(NL)])
        self.eval()

        i = torch.arange(HD // 2, dtype=torch.float32)
        ang = torch.arange(1024, dtype=torch.float32)[:, None] * torch.pow(ROPE_THETA, -2.0 * i / HD)[None, :]
        self.rope_cos, self.rope_sin = torch.cos(ang), torch.sin(ang)

    @classmethod
    def from_hub(cls, token=None):
        path = hf_hub_download(MODEL_ID, WEIGHTS_FILE, token=token)
        return cls(load_hf_tensors(path))

    # --- int8 -----------------------------------------------------------------
    def _calibrate(self, calib_ids):
        """Per-channel |activation| maxima at every Linear input, from one forward
        per calibration prompt over the canvas a first decode step actually sees
        (prompt + an all-masked answer). Running whole generations instead is 3x
        slower at boot and no more accurate."""
        act = {}

        def hook(name):
            def f(_mod, inp, _out):
                a = inp[0].detach().abs().amax(0)
                act[name] = a if name not in act else torch.maximum(act[name], a)
            return f

        handles = [m.register_forward_hook(hook(n))
                   for n, m in self.named_modules() if isinstance(m, nn.Linear)]
        try:
            for ids in calib_ids:
                P, T = len(ids), len(ids) + 64
                x = torch.full((T,), MASK_ID, dtype=torch.long)
                x[:P] = torch.tensor(ids, dtype=torch.long)
                self.forward(x, list(range(P, T)), self.new_state(T))
        finally:
            for h in handles:
                h.remove()
        return act

    def quantize_int8(self, calib_ids):
        """Smooth per-channel activation outliers into the weights, then quantize.

        Dynamic int8 (onednn/fbgemm) is the only fast int8 path on CPU, but it
        scales activations per *tensor*, and this model's post-norm activations have
        per-channel outliers that hijack that single scale — plain int8 costs ~35%
        of the argmaxes, which iterative unmasking then compounds into visibly worse
        answers. So do SmoothQuant first: pull a per-channel factor s out of each
        Linear's input and into whatever produced it, which is exact algebra, and
        the activations quantize cleanly at no runtime cost."""
        act = self._calibrate(calib_ids)

        def smooth(lin, act_max, absorb, group=1):
            """s = max|x|^a / max|w|^(1-a), folded into lin's columns and, inverted,
            into the producer. `group` shares one factor across neighbouring input
            channels (GQA: several query heads read the same kv channel)."""
            w_max = lin.weight.abs().amax(0)
            if group > 1:
                act_max = act_max.view(-1, group, HD).amax(1)
                w_max = w_max.view(-1, group, HD).amax(1)
            s = (act_max.clamp_min(1e-5) ** SMOOTH_ALPHA /
                 w_max.clamp_min(1e-5) ** (1 - SMOOTH_ALPHA)).clamp(1e-2, 1e2)
            cols = s[:, None, :].expand(-1, group, HD).reshape(-1) if group > 1 else s
            lin.weight.data *= cols
            absorb(s)

        for li, layer in enumerate(self.layers):
            n = f"layers.{li}."
            smooth(layer.w13, act[n + "w13"], lambda s: layer.ffn_norm.div_(s))
            # w2 reads silu(gate) * up, so its factor folds into the up rows.
            smooth(layer.w2, act[n + "w2"], lambda s: layer.w13.weight.data[FF:].div_(s[:, None]))
            if layer.attn:
                smooth(layer.qkv, act[n + "qkv"], lambda s: layer.op_norm.div_(s))
                # o reads the attention output, i.e. the v channels repeated across
                # each GQA group, so one factor per kv channel folds into the v rows.
                smooth(layer.o, act[n + "o"], group=NH // NKV,
                       absorb=lambda s: layer.qkv.weight.data[H + NKV * HD:].div_(s.reshape(-1, 1)))
            else:
                smooth(layer.in_proj, act[n + "in_proj"], lambda s: layer.op_norm.div_(s))
                # out_proj reads C * conv(Bx); the conv is depthwise, so its
                # per-channel kernel absorbs the factor (the Bx cache is untouched).
                smooth(layer.out_proj, act[n + "out_proj"], lambda s: layer.conv.div_(s[:, None, None]))
        smooth(self.lm_head, act["lm_head"], lambda s: self.emb_norm.div_(s))

        for eng in ("fbgemm", "onednn", "qnnpack"):
            if eng in torch.backends.quantized.supported_engines:
                torch.backends.quantized.engine = eng
                break
        torch.ao.quantization.quantize_dynamic(
            self, {nn.Linear: torch.ao.quantization.per_channel_dynamic_qconfig},
            dtype=torch.qint8, inplace=True)
        # collect first: wrapping while walking would descend into the wrappers
        todo = [(mod, name, child) for mod in self.modules()
                for name, child in mod.named_children()
                if isinstance(child, torch.ao.nn.quantized.dynamic.Linear)]
        for mod, name, child in todo:
            setattr(mod, name, RowScaledLinear(child))
        return self

    def _rope(self, x, pos):
        cos = self.rope_cos[pos][:, None, :]
        sin = self.rope_sin[pos][:, None, :]
        x1, x2 = x[..., :HD // 2], x[..., HD // 2:]
        return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)

    def new_state(self, T):
        """Per-generation DualCache: frozen K/V (attn) and Bx (conv) per layer."""
        st = []
        for L in self.layers:
            if L.attn:
                st.append({"k": torch.zeros(T, NKV, HD), "v": torch.zeros(T, NKV, HD)})
            else:
                st.append({"bx": torch.zeros(T, H)})
        return st

    def _conv_slice(self, bx_cache, conv_w, s0, length):
        T = bx_cache.shape[0]
        left = bx_cache[s0 - 1] if s0 - 1 >= 0 else torch.zeros(H)
        right = bx_cache[s0 + length] if s0 + length < T else torch.zeros(H)
        padded = torch.cat([left[None], bx_cache[s0:s0 + length], right[None]], 0)
        out = F.conv1d(padded.t().unsqueeze(0), conv_w, padding=0, groups=H)
        return out.squeeze(0).t()

    @torch.inference_mode()
    def forward(self, ids, rows, st, s0=None, length=None):
        """Full refresh (s0 is None) fills the caches over all T; a slice step
        (s0, length) recomputes only [s0, s0+length) against the frozen cache.
        rows: absolute positions to return logits for -> (len(rows), REAL_VOCAB)."""
        T = ids.shape[0]
        if s0 is None:
            s0, length = 0, T
        pos = torch.arange(s0, s0 + length)
        h = self.emb[ids[s0:s0 + length]]

        for L, cache in zip(self.layers, st):
            xn = rms_norm(h, L.op_norm)
            if L.attn:
                qkv = L.qkv(xn)
                q = qkv[:, :H].view(length, NH, HD)
                k = qkv[:, H:H + NKV * HD].view(length, NKV, HD)
                v = qkv[:, H + NKV * HD:].view(length, NKV, HD)
                q = self._rope(rms_norm(q, L.q_norm), pos)
                k = self._rope(rms_norm(k, L.k_norm), pos)
                cache["k"][s0:s0 + length] = k
                cache["v"][s0:s0 + length] = v
                # GQA handled inside sdpa (query head h -> kv head h // GQA); no
                # need to materialize an expanded [T, NH, HD] copy of the cache.
                attn = F.scaled_dot_product_attention(
                    q.transpose(0, 1), cache["k"].transpose(0, 1), cache["v"].transpose(0, 1),
                    enable_gqa=True)
                h = h + L.o(attn.transpose(0, 1).reshape(length, H))
            else:
                bcx = L.in_proj(xn)
                B, C, x = bcx[:, :H], bcx[:, H:2 * H], bcx[:, 2 * H:]
                cache["bx"][s0:s0 + length] = B * x
                conv = self._conv_slice(cache["bx"], L.conv, s0, length)
                h = h + L.out_proj(C * conv)
            fn = rms_norm(h, L.ffn_norm)
            gu = L.w13(fn)
            h = h + L.w2(F.silu(gu[:, :FF]) * gu[:, FF:])

        h = rms_norm(h, self.emb_norm)
        if rows is not None:
            h = h[torch.as_tensor([r - s0 for r in rows], dtype=torch.long)]
        return self.lm_head(h)[:, :REAL_VOCAB]

    @staticmethod
    def _select(logits, temperature):
        if temperature and temperature > 0:
            u = torch.rand_like(logits).clamp_min_(1e-20)
            ids = (logits + temperature * (-torch.log(-torch.log(u)))).argmax(-1)
        else:
            ids = logits.argmax(-1)
        conf = torch.softmax(logits, dim=-1).gather(1, ids.unsqueeze(1)).squeeze(1)
        return ids, conf


# --- MDLM decode loop (port of diffuseGPU: block-scheduled adaptive reveal) ---
def diffuse(model, prompt_ids, max_new=64, steps=32, block_size=16,
            temperature=0.0, tau=0.9):
    """Generator yielding {step, gen_ids, revealed} frames. gen_ids is the
    generated region only (length max_new); revealed is the set of indices (into
    gen_ids) freshly committed this step."""
    P = len(prompt_ids)
    T = P + max_new
    x = torch.full((T,), MASK_ID, dtype=torch.long)
    x[:P] = torch.tensor(prompt_ids, dtype=torch.long)

    yield {"step": 0, "gen_ids": x[P:].tolist(), "revealed": set()}

    st = model.new_state(T)
    nblk = max(1, math.ceil(max_new / block_size))
    per_block = max(1, math.ceil(steps / nblk))
    step = 0
    for b in range(nblk):
        s0 = P + b * block_size
        s1 = min(P + (b + 1) * block_size, T)
        if not (x[s0:s1] == MASK_ID).any():
            continue
        block_steps_left = per_block
        fresh = True   # first step of each block = full refresh (rebuild caches)
        while True:
            rows = [i for i in range(s0, s1) if x[i].item() == MASK_ID]
            if not rows:
                break
            if fresh:
                logits = model.forward(x, rows, st)
                fresh = False
            else:
                logits = model.forward(x, rows, st, s0=s0, length=s1 - s0)
            ids, conf = model._select(logits, temperature)
            cand = sorted(
                ({"pos": rows[i], "id": int(ids[i]), "conf": float(conf[i])}
                 for i in range(len(rows))),
                key=lambda c: c["conf"], reverse=True)
            min_k = len(rows) if block_steps_left <= 1 else math.ceil(len(rows) / block_steps_left)
            revealed, taken, committed = set(), set(), 0

            def adj(pos):
                return (pos - 1) in taken or (pos + 1) in taken

            for idx, c in enumerate(cand):
                if idx == 0 or ((c["conf"] >= tau or committed < min_k) and not adj(c["pos"])):
                    x[c["pos"]] = c["id"]
                    revealed.add(c["pos"] - P)
                    taken.add(c["pos"])
                    committed += 1
            for c in cand:
                if committed >= min_k:
                    break
                if c["pos"] not in taken:
                    x[c["pos"]] = c["id"]
                    revealed.add(c["pos"] - P)
                    taken.add(c["pos"])
                    committed += 1
            block_steps_left = max(1, block_steps_left - 1)
            step += 1
            yield {"step": step, "gen_ids": x[P:].tolist(), "revealed": revealed}