Blame for modeling_diffusion_lfm.py

"""Optimized drop-in patch for summerAIO/diffusion-lfm2.5-24b-a2b (modeling_diffusion_lfm.py).

What changes vs the original sampler
-------------------------------------
1. Prefix reuse (exact): the prefix is causal and never attends to the canvas, so its
   per-layer hidden states are computed ONCE per block instead of once per denoising step.
   Each step then runs the MoE FFN / conv only on the canvas tokens. Attention layers still
   see the full prefix K/V, conv layers see the last few prefix tokens.
   -> a self-test compares against the original forward and falls back automatically if
      the internals don't match (e.g. a different transformers version).
2. Attention mask / rotary built once per block, not every step.
3. No GPU->CPU sync inside the step loop (commit count is computed arithmetically).
4. Bug fixes: eos_token_id may be a list; pad_token_id may be None; prompt is no longer
   silently cut to its last 256 tokens (prefix window is configurable, and the prompt
   itself is always kept); EOS inside a partially-committed canvas stops early.
5. Optional quality knobs (off by default = original behaviour):
   repetition_penalty, top_k / top_p, and remask_low_conf (re-masks the weakest committed
   tokens once before the final step).

Usage
-----
    from diffusion_lfm_fast import from_pretrained
    model = from_pretrained(repo, torch_dtype="bfloat16", device_map="auto")
    model.self_test(tokenizer)          # optional, also runs automatically on first generate
    tokens, stats = model.generate(prompt, max_new_tokens=1024, block_size=256, steps=8)
"""
import torch
from torch import nn


def _as_list(x):
    if x is None:
        return []
    return list(x) if isinstance(x, (list, tuple)) else [x]


class DiffusionLFM(nn.Module):
    def __init__(self, causal_model, mask_embedding):
        super().__init__()
        self.model = causal_model
        self.mask_embedding = nn.Parameter(mask_embedding.detach().clone(), requires_grad=False)
        for layer in self.model.model.layers:
            if layer.is_attention_layer:
                layer.self_attn.is_causal = False
        cfg = self.model.config
        self.eos_ids = set(_as_list(cfg.eos_token_id))
        pad = getattr(cfg, "pad_token_id", None)
        self.pad_id = pad if pad is not None else (min(self.eos_ids) if self.eos_ids else 0)
        self.conv_tail = int(getattr(cfg, "conv_L_cache", 3)) + 1  # causal conv context
        self.use_prefix_cache = None  # decided by self_test()

    @property
    def causal_model(self):
        return self.model

    # ------------------------------------------------------------------ masks
    @staticmethod
    def _mask(length, p, device, dtype):
        q = torch.arange(length, device=device)[:, None]
        k = torch.arange(length, device=device)[None, :]
        allowed = torch.where(q < p, k <= q, torch.ones_like(k <= q))
        return torch.zeros(allowed.shape, device=device, dtype=dtype).masked_fill(
            ~allowed, torch.finfo(dtype).min)[None, None]

    def _embed(self, ids, masked):
        h = self.model.model.embed_tokens(ids)
        return torch.where(masked[..., None], self.mask_embedding.to(h.dtype), h)

    # ------------------------------------------------ original (reference) path
    def hidden(self, ids, masked, valid_or_p, prefix_lengths=None):
        """Backward compatible: hidden(ids, masked, valid, prefix_lengths) (original API)
        or hidden(ids, masked, p)."""
        m = self.model.model
        h = self._embed(ids, masked)
        pos = torch.arange(ids.shape[1], device=ids.device)[None]
        rot = m.pos_emb(h, pos)
        if prefix_lengths is None:
            p = int(valid_or_p)
            valid = torch.ones_like(ids, dtype=torch.bool)
        else:
            p = int(prefix_lengths.reshape(-1)[0])
            valid = valid_or_p.bool()
        attn = self._mask(ids.shape[1], p, ids.device, h.dtype)
        if not bool(valid.all()):
            attn = attn.masked_fill(~valid[:, None, None, :], torch.finfo(h.dtype).min)
        for layer in m.layers:
            h = layer(h, position_embeddings=rot, attention_mask=attn if layer.is_attention_layer else valid,
                      position_ids=pos, past_key_values=None)
        return m.embedding_norm(h)

    # ------------------------------------------------------ prefix-reuse path
    def _layer_parts(self, layer, x, rot, attn, ctx):
        """Run one decoder layer on canvas x, with ctx = this layer's prefix input."""
        n = x.shape[1]
        if layer.is_attention_layer:
            full = torch.cat((ctx, x), 1)
            out = layer.self_attn(hidden_states=layer.operator_norm(full),
                                  position_embeddings=rot, attention_mask=attn)
        else:
            full = torch.cat((ctx[:, -self.conv_tail:], x), 1)
            out = layer.conv(layer.operator_norm(full))
        out = out[0] if isinstance(out, tuple) else out
        h = out[:, -n:] + x
        ff = layer.feed_forward(layer.ffn_norm(h))
        ff = ff[0] if isinstance(ff, tuple) else ff
        return h + ff

    def _prefix_states(self, prefix):
        """Per-layer INPUT hidden states of the (causal) prefix, computed once."""
        m = self.model.model
        p = prefix.shape[1]
        h = m.embed_tokens(prefix)
        pos = torch.arange(p, device=prefix.device)[None]
        rot = m.pos_emb(h, pos)
        attn = self._mask(p, p, prefix.device, h.dtype)
        valid = torch.ones_like(prefix, dtype=torch.bool)
        states = []
        for layer in m.layers:
            states.append(h)
            h = layer(h, position_embeddings=rot, attention_mask=attn if layer.is_attention_layer else valid,
                      position_ids=pos, past_key_values=None)
        return states

    def _canvas_hidden(self, canvas_ids, masked, states, rot_canvas, rot_full, attn):
        m = self.model.model
        h = self._embed(canvas_ids, masked)
        for layer, ctx in zip(m.layers, states):
            h = self._layer_parts(layer, h, rot_full, attn, ctx)
        return m.embedding_norm(h)

    # ------------------------------------------------------------- self test
    @torch.inference_mode()
    def self_test(self, tokenizer=None, p=48, n=32, tol=5e-2):
        dev = self.model.device
        vocab = self.model.config.vocab_size
        g = torch.Generator(device="cpu").manual_seed(0)
        ids = torch.randint(100, min(vocab, 30000), (1, p + n), generator=g).to(dev)
        masked = torch.zeros_like(ids, dtype=torch.bool)
        masked[:, p + n // 2:] = True
        try:
            ref = self.hidden(ids, masked, p)[:, p:]
            got = self._fast_block_hidden(ids[:, :p], ids[:, p:], masked[:, p:])
            diff = (ref.float() - got.float()).abs().max().item()
            scale = ref.float().abs().max().item() + 1e-6
            self.use_prefix_cache = diff / scale < tol
            print(f"[diffusion_lfm_fast] prefix-reuse rel-diff={diff/scale:.2e} -> "
                  f"{'ENABLED' if self.use_prefix_cache else 'DISABLED (fallback)'}")
        except Exception as e:  # internals differ from what we expect
            self.use_prefix_cache = False
            print(f"[diffusion_lfm_fast] prefix-reuse unavailable ({type(e).__name__}: {e}); using reference path")
        return self.use_prefix_cache

    def _fast_block_hidden(self, prefix, canvas_ids, masked):
        """Helper used by self_test: one forward through the fast path."""
        m = self.model.model
        p, n = prefix.shape[1], canvas_ids.shape[1]
        states = self._prefix_states(prefix)
        dummy = m.embed_tokens(canvas_ids)
        rot_full = m.pos_emb(dummy, torch.arange(p + n, device=prefix.device)[None])
        attn = self._mask(p + n, p, prefix.device, dummy.dtype)[:, :, p:]
        return self._canvas_hidden(canvas_ids, masked, states, None, rot_full, attn)

    # ------------------------------------------------------------- sampling
    def _pick(self, logits, temperature, top_k, top_p):
        if not temperature:
            lp = logits.log_softmax(-1)
            conf, tok = lp.max(-1)
            return tok, conf
        logits = logits / temperature
        if top_k:
            kth = logits.topk(top_k, -1).values[..., -1:]
            logits = logits.masked_fill(logits < kth, -torch.inf)
        if top_p and top_p < 1.0:
            s, idx = logits.sort(-1, descending=True)
            cum = s.softmax(-1).cumsum(-1)
            s = s.masked_fill(cum - s.softmax(-1) > top_p, -torch.inf)
            logits = torch.full_like(logits, -torch.inf).scatter(-1, idx, s)
        probs = logits.softmax(-1)
        tok = torch.multinomial(probs, 1).squeeze(-1)
        return tok, probs.gather(-1, tok[:, None]).squeeze(-1).clamp_min(1e-12).log()

    @torch.inference_mode()
    def generate(self, prompt_ids, *, max_new_tokens=128, block_size=256, steps=3,
                 temperature=0.0, top_k=0, top_p=1.0, repetition_penalty=1.0,
                 prefix_window=256, remask_low_conf=0.0):
        if prompt_ids.ndim != 2 or prompt_ids.shape[0] != 1 or min(max_new_tokens, block_size, steps) < 1:
            raise ValueError("one prompt and positive generation settings required")
        if self.use_prefix_cache is None:
            self.self_test()
        m, dev = self.model.model, prompt_ids.device
        generated, calls = [], 0
        eos_t = torch.tensor(sorted(self.eos_ids) or [-1], device=dev)
        done = False

        while len(generated) < max_new_tokens and not done:
            count = min(block_size, max_new_tokens - len(generated))
            hist = prompt_ids.new_tensor([generated[-prefix_window:]]) if generated else prompt_ids[:, :0]
            prefix = torch.cat((prompt_ids, hist), 1)
            if prefix.shape[1] > prompt_ids.shape[1] + prefix_window:  # keep prompt, window history
                prefix = torch.cat((prompt_ids, hist[:, -prefix_window:]), 1)
            p = prefix.shape[1]

            canvas = prefix.new_full((1, count), self.pad_id)
            masked = torch.ones_like(canvas, dtype=torch.bool)

            # --- built once per block
            if self.use_prefix_cache:
                states = self._prefix_states(prefix)
                dummy = m.embed_tokens(canvas)
                rot_full = m.pos_emb(dummy, torch.arange(p + count, device=dev)[None])
                attn = self._mask(p + count, p, dev, dummy.dtype)[:, :, p:]
            penalty_ids = None
            if repetition_penalty != 1.0:
                penalty_ids = torch.unique(prefix[0, -prefix_window:])

            schedule = [count * (steps - s - 1) // steps for s in range(steps)]  # remaining masked
            for s in range(steps):
                if self.use_prefix_cache:
                    h = self._canvas_hidden(canvas, masked, states, None, rot_full, attn)
                else:
                    ids = torch.cat((prefix, canvas), 1)
                    mk = torch.cat((torch.zeros_like(prefix, dtype=torch.bool), masked), 1)
                    h = self.hidden(ids, mk, p)[:, p:]
                calls += 1
                logits = self.model.lm_head(h[0]).float()           # fixed shape [count, V]
                logits[:, self.pad_id] = -torch.inf
                if penalty_ids is not None:
                    pen = torch.cat((penalty_ids, canvas[0][~masked[0]]))
                    sel = logits[:, pen]
                    logits[:, pen] = torch.where(sel > 0, sel / repetition_penalty, sel * repetition_penalty)
                tok, conf = self._pick(logits, temperature, top_k, top_p)
                conf = conf.masked_fill(~masked[0], -torch.inf)       # only fill masked slots

                if s == steps - 1:
                    canvas[0] = torch.where(masked[0], tok, canvas[0])
                    masked.zero_()
                    break
                commit = (count if s == 0 else schedule[s - 1]) - schedule[s]
                if commit > 0:
                    chosen = conf.topk(commit).indices
                    canvas[0, chosen] = tok[chosen]
                    masked[0, chosen] = False
                # optional: re-mask weakest committed tokens once, right before the last step
                if remask_low_conf and s == steps - 2:
                    with_conf = conf.clone()
                    k = int(remask_low_conf * count)
                    if k > 0:
                        committed = (~masked[0]).nonzero().squeeze(1)
                        if committed.numel() > k:
                            weakest = committed[with_conf[committed].nan_to_num(-1e9).topk(k, largest=False).indices]
                            masked[0, weakest] = True
                # early exit: an EOS has been committed and everything before it is filled
                hit = torch.isin(canvas[0], eos_t) & ~masked[0]
                if hit.any():
                    first = int(hit.nonzero()[0])
                    if not masked[0, :first].any():
                        break

            out = canvas[0].tolist()
            for i, t in enumerate(out):
                if t in self.eos_ids:
                    out, done = out[:i], True
                    break
            generated.extend(out)
        return generated[:max_new_tokens], {"forward_calls": calls, "tokens": len(generated),
                                            "prefix_cache": bool(self.use_prefix_cache)}


def from_pretrained(model_id, **kwargs):
    from safetensors.torch import load_file
    from huggingface_hub import hf_hub_download
    from transformers import AutoModelForCausalLM
    model = AutoModelForCausalLM.from_pretrained(model_id, **kwargs)
    model.eval()
    weights = load_file(hf_hub_download(model_id, "diffusion.safetensors"), device=str(model.device))
    return DiffusionLFM(model, weights["mask_embedding"]).eval()