LiquidAI/masked-diffusion

1"""lfm2_diffusion.py — CPU (PyTorch) masked-diffusion inference for LFM2.5 Encoder.2 3Server-side replacement for the old in-browser WebGPU path. The model, decode4loop, and numerics are a 1:1 port of the reference JS (lfm2_webgpu_forward.js /5lfm2_bidir_mdlm.js): a *bidirectional* LFM2 forward, then the block-scheduled,6confidence-driven MDLM unmasking loop.7 8Speed comes from three things:9  * fast-dLLM "DualCache" (as in the WebGPU build): the first step of each block10    runs a full-canvas forward that fills a per-layer K/V cache (attention) and Bx11    cache (shortconv); every later step of that block only recomputes the active12    ~block-size slice, attending to / convolving against the frozen cache.13  * smoothed int8 dynamic quantization of every matmul (see `quantize_int8`).14  * one worker thread per available core (cgroup-aware, no oversubscription).15 16The checkpoint is pulled from the Hub model repo at boot (no local copy);17everything runs on CPU so nothing is downloaded to the client.18"""19 20import os21import math22import re23 24import torch25import torch.nn as nn26import torch.nn.functional as F27from huggingface_hub import hf_hub_download28from safetensors.torch import load_file29 30# The perfectblend diffusion-SFT checkpoint (f32 safetensors). It is the encoder31# turned into an MDLM chatbot — the *base* LFM2.5-Encoder-350M is a plain MLM and32# echoes instead of answering, so it is not interchangeable here.33MODEL_ID = os.environ.get("DIFFUSION_MODEL", "LiquidAI/LFM2.5-Encoder-350M-Diffusion")34WEIGHTS_FILE = os.environ.get("DIFFUSION_WEIGHTS", "model.safetensors")35 36# --- architecture constants (LFM2.5-Encoder-350M), mirrored from the JS ------37H = 1024            # hidden size38NL = 16             # layers39NH = 16             # attention heads40NKV = 8             # kv heads (GQA group size 2)41HD = 64             # head dim42FF = 4608           # ffn inner size43EPS = 1e-544ROPE_THETA = 1e645ATTN_LAYERS = frozenset({2, 5, 8, 10, 12, 14})   # the rest are shortconv layers46 47# --- MDLM / tokenizer constants (see the model card) -------------------------48MASK_ID = 16        # <|mask|>49REAL_VOCAB = 64402  # ignore logit ids >= REAL_VOCAB (padded vocab is 65536)50 51# --- int8 smoothing (see Lfm2Diffusion.quantize_int8) ------------------------52SMOOTH_ALPHA = 0.553# Fixed calibration set, only used to measure per-channel activation ranges (one54# forward each, ~1s at boot). app.py tokenizes these.55CALIB_PROMPTS = (56    "What is the capital of France?",57    "Write a haiku about the ocean.",58    "Give me two tips for better sleep.",59    "What is 12 times 8?",60)61 62 63def effective_cpus():64    """Cores actually available to this process (cgroup quota, not the host)."""65    try:  # cgroup v266        raw = open("/sys/fs/cgroup/cpu.max").read().split()67        if raw and raw[0] != "max":68            return max(1, round(int(raw[0]) / int(raw[1])))69    except (OSError, ValueError):70        pass71    try:  # cgroup v172        q = int(open("/sys/fs/cgroup/cpu/cpu.cfs_quota_us").read())73        p = int(open("/sys/fs/cgroup/cpu/cpu.cfs_period_us").read())74        if q > 0:75            return max(1, q // p)76    except (OSError, ValueError):77        pass78    try:79        return max(1, len(os.sched_getaffinity(0)))80    except AttributeError:81        return os.cpu_count() or 182 83 84def render_chat_prompt(user_text, system_text=None):85    """LFM2 chat template for a single user turn + add_generation_prompt."""86    s = ""87    if system_text:88        s += f"[SYS]\n{system_text.strip()}\n[/SYS]\n\n"89    s += f"[Question]\n{user_text.strip()}\n[/Question]\n\n[Answer]\n"90    return s91 92 93def trim(text):94    """Cut generated text at the answer terminator / role markers (app.py trim)."""95    for stop in ("[/Answer]", "[Question]", "[SYS]"):96        i = text.find(stop)97        if i != -1:98            text = text[:i]99    text = text.strip()100    text = re.sub(r"^(?:\[Answer\]|Answer\b)\s*:?\s*", "", text, flags=re.I)101    text = re.sub(r"\s*(?:\.{3}|…)\s*$", "", text)102    return text.strip()103 104 105def load_hf_tensors(path):106    """Load the safetensors checkpoint — weights are already f32 and [out, in], so107    there is nothing to dequantize or transpose. The dimension check makes a108    checkpoint that doesn't match the constants above fail loudly at boot rather109    than as silent garbage."""110    t = load_file(path)111    attn = frozenset(int(m.group(1)) for k in t112                     if (m := re.match(r"lfm2\.layers\.(\d+)\.self_attn\.", k)))113    a = min(attn)114    got = (t["lfm2.embed_tokens.weight"].shape[1],115           1 + max(int(m.group(1)) for k in t if (m := re.match(r"lfm2\.layers\.(\d+)\.", k))),116           attn,117           t["lfm2.layers.0.feed_forward.w1.weight"].shape[0],118           t[f"lfm2.layers.{a}.self_attn.q_proj.weight"].shape[0],119           t[f"lfm2.layers.{a}.self_attn.k_proj.weight"].shape[0])120    want = (H, NL, ATTN_LAYERS, FF, NH * HD, NKV * HD)121    if got != want:122        raise ValueError(f"{path}: (hidden, layers, attn layers, ffn, q, kv) is {got}, expected {want}")123    return t124 125 126class RowScaledLinear(nn.Module):127    """Per-token activations around a dynamically quantized Linear.128 129    The int8 kernel derives ONE activation scale per call, so a single loud token130    coarsens every other row in the batch. Normalising each row to unit max first131    (and putting the scale back on the output) makes that one scale exactly right132    for every row, which is per-token quantization at the cost of two elementwise133    passes."""134 135    def __init__(self, qlin):136        super().__init__()137        self.q = qlin138 139    def forward(self, x):140        r = x.abs().amax(-1, keepdim=True).clamp_min(1e-6)141        return self.q(x / r) * r142 143 144def rms_norm(x, w, eps=EPS):145    inv = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)146    return x * inv * w147 148 149def _linear(in_f, out_f, weight):150    m = nn.Linear(in_f, out_f, bias=False)151    m.weight.data = weight.contiguous()152    return m153 154 155class Lfm2Layer(nn.Module):156    def __init__(self, t, li):157        super().__init__()158        p = f"lfm2.layers.{li}."159        self.attn = li in ATTN_LAYERS160        self.register_buffer("op_norm", t[p + "operator_norm.weight"], persistent=False)161        self.register_buffer("ffn_norm", t[p + "ffn_norm.weight"], persistent=False)162        # fused gate+up: one GEMM instead of two163        self.w13 = _linear(H, 2 * FF, torch.cat(164            [t[p + "feed_forward.w1.weight"], t[p + "feed_forward.w3.weight"]], 0))165        self.w2 = _linear(FF, H, t[p + "feed_forward.w2.weight"])166        if self.attn:167            self.qkv = _linear(H, H + 2 * NKV * HD, torch.cat(168                [t[p + "self_attn.q_proj.weight"], t[p + "self_attn.k_proj.weight"],169                 t[p + "self_attn.v_proj.weight"]], 0))170            self.o = _linear(H, H, t[p + "self_attn.out_proj.weight"])171            self.register_buffer("q_norm", t[p + "self_attn.q_layernorm.weight"], persistent=False)172            self.register_buffer("k_norm", t[p + "self_attn.k_layernorm.weight"], persistent=False)173        else:174            self.register_buffer("conv", t[p + "conv.conv.weight"].reshape(H, 1, -1), persistent=False)175            self.in_proj = _linear(H, 3 * H, t[p + "conv.in_proj.weight"])176            self.out_proj = _linear(H, H, t[p + "conv.out_proj.weight"])177 178 179class Lfm2Diffusion(nn.Module):180    """Takes ownership of `tensors`: most weights are used without copying, and181    quantize_int8 rescales them in place."""182 183    def __init__(self, tensors):184        super().__init__()185        t = tensors186        self.register_buffer("emb", t["lfm2.embed_tokens.weight"].contiguous(), persistent=False)187        self.register_buffer("emb_norm", t["lfm2.embedding_norm.weight"], persistent=False)188        # tied logits, trimmed to the real vocab (ids >= REAL_VOCAB are suppressed189        # anyway, so there's no point computing the padded tail every step). Cloned,190        # not shared with `emb`: quantize_int8 rescales the head's columns in place.191        self.lm_head = _linear(H, REAL_VOCAB, t["lfm2.embed_tokens.weight"][:REAL_VOCAB].clone())192        self.layers = nn.ModuleList([Lfm2Layer(t, li) for li in range(NL)])193        self.eval()194 195        i = torch.arange(HD // 2, dtype=torch.float32)196        ang = torch.arange(1024, dtype=torch.float32)[:, None] * torch.pow(ROPE_THETA, -2.0 * i / HD)[None, :]197        self.rope_cos, self.rope_sin = torch.cos(ang), torch.sin(ang)198 199    @classmethod200    def from_hub(cls, token=None):201        path = hf_hub_download(MODEL_ID, WEIGHTS_FILE, token=token)202        return cls(load_hf_tensors(path))203 204    # --- int8 -----------------------------------------------------------------205    def _calibrate(self, calib_ids):206        """Per-channel |activation| maxima at every Linear input, from one forward207        per calibration prompt over the canvas a first decode step actually sees208        (prompt + an all-masked answer). Running whole generations instead is 3x209        slower at boot and no more accurate."""210        act = {}211 212        def hook(name):213            def f(_mod, inp, _out):214                a = inp[0].detach().abs().amax(0)215                act[name] = a if name not in act else torch.maximum(act[name], a)216            return f217 218        handles = [m.register_forward_hook(hook(n))219                   for n, m in self.named_modules() if isinstance(m, nn.Linear)]220        try:221            for ids in calib_ids:222                P, T = len(ids), len(ids) + 64223                x = torch.full((T,), MASK_ID, dtype=torch.long)224                x[:P] = torch.tensor(ids, dtype=torch.long)225                self.forward(x, list(range(P, T)), self.new_state(T))226        finally:227            for h in handles:228                h.remove()229        return act230 231    def quantize_int8(self, calib_ids):232        """Smooth per-channel activation outliers into the weights, then quantize.233 234        Dynamic int8 (onednn/fbgemm) is the only fast int8 path on CPU, but it235        scales activations per *tensor*, and this model's post-norm activations have236        per-channel outliers that hijack that single scale — plain int8 costs ~35%237        of the argmaxes, which iterative unmasking then compounds into visibly worse238        answers. So do SmoothQuant first: pull a per-channel factor s out of each239        Linear's input and into whatever produced it, which is exact algebra, and240        the activations quantize cleanly at no runtime cost."""241        act = self._calibrate(calib_ids)242 243        def smooth(lin, act_max, absorb, group=1):244            """s = max|x|^a / max|w|^(1-a), folded into lin's columns and, inverted,245            into the producer. `group` shares one factor across neighbouring input246            channels (GQA: several query heads read the same kv channel)."""247            w_max = lin.weight.abs().amax(0)248            if group > 1:249                act_max = act_max.view(-1, group, HD).amax(1)250                w_max = w_max.view(-1, group, HD).amax(1)251            s = (act_max.clamp_min(1e-5) ** SMOOTH_ALPHA /252                 w_max.clamp_min(1e-5) ** (1 - SMOOTH_ALPHA)).clamp(1e-2, 1e2)253            cols = s[:, None, :].expand(-1, group, HD).reshape(-1) if group > 1 else s254            lin.weight.data *= cols255            absorb(s)256 257        for li, layer in enumerate(self.layers):258            n = f"layers.{li}."259            smooth(layer.w13, act[n + "w13"], lambda s: layer.ffn_norm.div_(s))260            # w2 reads silu(gate) * up, so its factor folds into the up rows.261            smooth(layer.w2, act[n + "w2"], lambda s: layer.w13.weight.data[FF:].div_(s[:, None]))262            if layer.attn:263                smooth(layer.qkv, act[n + "qkv"], lambda s: layer.op_norm.div_(s))264                # o reads the attention output, i.e. the v channels repeated across265                # each GQA group, so one factor per kv channel folds into the v rows.266                smooth(layer.o, act[n + "o"], group=NH // NKV,267                       absorb=lambda s: layer.qkv.weight.data[H + NKV * HD:].div_(s.reshape(-1, 1)))268            else:269                smooth(layer.in_proj, act[n + "in_proj"], lambda s: layer.op_norm.div_(s))270                # out_proj reads C * conv(Bx); the conv is depthwise, so its271                # per-channel kernel absorbs the factor (the Bx cache is untouched).272                smooth(layer.out_proj, act[n + "out_proj"], lambda s: layer.conv.div_(s[:, None, None]))273        smooth(self.lm_head, act["lm_head"], lambda s: self.emb_norm.div_(s))274 275        for eng in ("fbgemm", "onednn", "qnnpack"):276            if eng in torch.backends.quantized.supported_engines:277                torch.backends.quantized.engine = eng278                break279        torch.ao.quantization.quantize_dynamic(280            self, {nn.Linear: torch.ao.quantization.per_channel_dynamic_qconfig},281            dtype=torch.qint8, inplace=True)282        # collect first: wrapping while walking would descend into the wrappers283        todo = [(mod, name, child) for mod in self.modules()284                for name, child in mod.named_children()285                if isinstance(child, torch.ao.nn.quantized.dynamic.Linear)]286        for mod, name, child in todo:287            setattr(mod, name, RowScaledLinear(child))288        return self289 290    def _rope(self, x, pos):291        cos = self.rope_cos[pos][:, None, :]292        sin = self.rope_sin[pos][:, None, :]293        x1, x2 = x[..., :HD // 2], x[..., HD // 2:]294        return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)295 296    def new_state(self, T):297        """Per-generation DualCache: frozen K/V (attn) and Bx (conv) per layer."""298        st = []299        for L in self.layers:300            if L.attn:301                st.append({"k": torch.zeros(T, NKV, HD), "v": torch.zeros(T, NKV, HD)})302            else:303                st.append({"bx": torch.zeros(T, H)})304        return st305 306    def _conv_slice(self, bx_cache, conv_w, s0, length):307        T = bx_cache.shape[0]308        left = bx_cache[s0 - 1] if s0 - 1 >= 0 else torch.zeros(H)309        right = bx_cache[s0 + length] if s0 + length < T else torch.zeros(H)310        padded = torch.cat([left[None], bx_cache[s0:s0 + length], right[None]], 0)311        out = F.conv1d(padded.t().unsqueeze(0), conv_w, padding=0, groups=H)312        return out.squeeze(0).t()313 314    @torch.inference_mode()315    def forward(self, ids, rows, st, s0=None, length=None):316        """Full refresh (s0 is None) fills the caches over all T; a slice step317        (s0, length) recomputes only [s0, s0+length) against the frozen cache.318        rows: absolute positions to return logits for -> (len(rows), REAL_VOCAB)."""319        T = ids.shape[0]320        if s0 is None:321            s0, length = 0, T322        pos = torch.arange(s0, s0 + length)323        h = self.emb[ids[s0:s0 + length]]324 325        for L, cache in zip(self.layers, st):326            xn = rms_norm(h, L.op_norm)327            if L.attn:328                qkv = L.qkv(xn)329                q = qkv[:, :H].view(length, NH, HD)330                k = qkv[:, H:H + NKV * HD].view(length, NKV, HD)331                v = qkv[:, H + NKV * HD:].view(length, NKV, HD)332                q = self._rope(rms_norm(q, L.q_norm), pos)333                k = self._rope(rms_norm(k, L.k_norm), pos)334                cache["k"][s0:s0 + length] = k335                cache["v"][s0:s0 + length] = v336                # GQA handled inside sdpa (query head h -> kv head h // GQA); no337                # need to materialize an expanded [T, NH, HD] copy of the cache.338                attn = F.scaled_dot_product_attention(339                    q.transpose(0, 1), cache["k"].transpose(0, 1), cache["v"].transpose(0, 1),340                    enable_gqa=True)341                h = h + L.o(attn.transpose(0, 1).reshape(length, H))342            else:343                bcx = L.in_proj(xn)344                B, C, x = bcx[:, :H], bcx[:, H:2 * H], bcx[:, 2 * H:]345                cache["bx"][s0:s0 + length] = B * x346                conv = self._conv_slice(cache["bx"], L.conv, s0, length)347                h = h + L.out_proj(C * conv)348            fn = rms_norm(h, L.ffn_norm)349            gu = L.w13(fn)350            h = h + L.w2(F.silu(gu[:, :FF]) * gu[:, FF:])351 352        h = rms_norm(h, self.emb_norm)353        if rows is not None:354            h = h[torch.as_tensor([r - s0 for r in rows], dtype=torch.long)]355        return self.lm_head(h)[:, :REAL_VOCAB]356 357    @staticmethod358    def _select(logits, temperature):359        if temperature and temperature > 0:360            u = torch.rand_like(logits).clamp_min_(1e-20)361            ids = (logits + temperature * (-torch.log(-torch.log(u)))).argmax(-1)362        else:363            ids = logits.argmax(-1)364        conf = torch.softmax(logits, dim=-1).gather(1, ids.unsqueeze(1)).squeeze(1)365        return ids, conf366 367 368# --- MDLM decode loop (port of diffuseGPU: block-scheduled adaptive reveal) ---369def diffuse(model, prompt_ids, max_new=64, steps=32, block_size=16,370            temperature=0.0, tau=0.9):371    """Generator yielding {step, gen_ids, revealed} frames. gen_ids is the372    generated region only (length max_new); revealed is the set of indices (into373    gen_ids) freshly committed this step."""374    P = len(prompt_ids)375    T = P + max_new376    x = torch.full((T,), MASK_ID, dtype=torch.long)377    x[:P] = torch.tensor(prompt_ids, dtype=torch.long)378 379    yield {"step": 0, "gen_ids": x[P:].tolist(), "revealed": set()}380 381    st = model.new_state(T)382    nblk = max(1, math.ceil(max_new / block_size))383    per_block = max(1, math.ceil(steps / nblk))384    step = 0385    for b in range(nblk):386        s0 = P + b * block_size387        s1 = min(P + (b + 1) * block_size, T)388        if not (x[s0:s1] == MASK_ID).any():389            continue390        block_steps_left = per_block391        fresh = True   # first step of each block = full refresh (rebuild caches)392        while True:393            rows = [i for i in range(s0, s1) if x[i].item() == MASK_ID]394            if not rows:395                break396            if fresh:397                logits = model.forward(x, rows, st)398                fresh = False399            else:400                logits = model.forward(x, rows, st, s0=s0, length=s1 - s0)401            ids, conf = model._select(logits, temperature)402            cand = sorted(403                ({"pos": rows[i], "id": int(ids[i]), "conf": float(conf[i])}404                 for i in range(len(rows))),405                key=lambda c: c["conf"], reverse=True)406            min_k = len(rows) if block_steps_left <= 1 else math.ceil(len(rows) / block_steps_left)407            revealed, taken, committed = set(), set(), 0408 409            def adj(pos):410                return (pos - 1) in taken or (pos + 1) in taken411 412            for idx, c in enumerate(cand):413                if idx == 0 or ((c["conf"] >= tau or committed < min_k) and not adj(c["pos"])):414                    x[c["pos"]] = c["id"]415                    revealed.add(c["pos"] - P)416                    taken.add(c["pos"])417                    committed += 1418            for c in cand:419                if committed >= min_k:420                    break421                if c["pos"] not in taken:422                    x[c["pos"]] = c["id"]423                    revealed.add(c["pos"] - P)424                    taken.add(c["pos"])425                    committed += 1426            block_steps_left = max(1, block_steps_left - 1)427            step += 1428            yield {"step": step, "gen_ids": x[P:].tolist(), "revealed": revealed}429