"""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()