| """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 |
| self.use_prefix_cache = None |
|
|
| @property |
| def causal_model(self): |
| return self.model |
|
|
| |
| @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) |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| @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: |
| 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) |
|
|
| |
| 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: |
| 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) |
|
|
| |
| 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)] |
| 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() |
| 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) |
|
|
| 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 |
| |
| 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 |
| |
| 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() |
|
|
