import copy from dataclasses import dataclass from typing import Optional, Tuple import numpy as np import torch import torch.nn.functional as F from torch import nn from transformers.modeling_outputs import CausalLMOutputWithPast, BaseModelOutput from transformers.utils import ModelOutput from torch.nn.attention.flex_attention import flex_attention, create_block_mask from transformers.modeling_flash_attention_utils import FlashAttentionKwargs from transformers.processing_utils import Unpack from transformers.cache_utils import Cache, DynamicCache from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss from transformers.generation import GenerationMixin import math from .modeling_ministral import Ministral3Model, Ministral3PreTrainedModel, Ministral3Attention, apply_rotary_pos_emb, repeat_kv, _get_llama_4_attn_scale from .configuration_nemotron_labs_diffusion import NemotronLabsDiffusionConfig __all__ = ["NemotronLabsDiffusionModel", "NemotronLabsDiffusionFlexAttention"] @dataclass class NemotronLabsDiffusionOutputWithPast(ModelOutput): loss: torch.FloatTensor | None = None logits: torch.FloatTensor | None = None causal_logits: torch.FloatTensor | None = None past_key_values: Cache | None = None hidden_states: tuple[torch.FloatTensor, ...] | None = None attentions: tuple[torch.FloatTensor, ...] | None = None @torch.compile(fullgraph=True, mode="max-autotune-no-cudagraphs", dynamic=False) def fused_flex_attention(q, k, v, block_mask=None): return flex_attention(q, k, v, block_mask=block_mask) class NemotronLabsDiffusionFlexAttention(Ministral3Attention): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.block_size = self.config.block_size self.block_diff_mask = None import torch._dynamo.config as dcfg dcfg.cache_size_limit = 512 def compute_block_mask(self, mode, q_len, block_size=None): def block_diff_mask(block_size, b, h, q_idx, kv_idx, n): x0_flag_q = (q_idx >= n) x0_flag_kv = (kv_idx >= n) # Compute block indices block_q = torch.where(x0_flag_q == 1, (q_idx - n) // block_size, q_idx // block_size) block_kv = torch.where(x0_flag_kv == 1, (kv_idx - n) // block_size, kv_idx // block_size) # **1. Block Diagonal Mask (M_BD) ** block_diagonal = (block_q == block_kv) & (x0_flag_kv == 0) & (x0_flag_q == 0) # **2. Offset Block-Causal Mask (M_OBC) ** offset_block_causal = ( (block_q > block_kv) & (x0_flag_kv == 1) & (x0_flag_q == 0) ) # **3. Fully Causal Mask (M_BC) ** fully_causal = (q_idx >= kv_idx) & (x0_flag_kv == 1) & (x0_flag_q == 1) # **4. Combine Masks ** return block_diagonal | offset_block_causal | fully_causal attn_mask = lambda b, h, q, kv: block_diff_mask(block_size, b, h, q, kv, q_len//2) block_mask = create_block_mask( attn_mask, B=None, H=None, Q_LEN=q_len, KV_LEN=q_len ) return block_mask def forward( self, hidden_states: torch.Tensor, position_embeddings: Tuple[torch.Tensor, torch.Tensor], attention_mask: Optional[torch.Tensor], past_key_values: Optional[Cache] = None, cache_position: Optional[torch.LongTensor] = None, is_training: bool = True, **kwargs: Unpack[FlashAttentionKwargs], ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: bsz, q_len, _ = hidden_states.size() input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.head_dim) query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) cos, sin = position_embeddings if is_training: # Split query and key states in half along sequence length dimension q1, q2 = query_states.chunk(2, dim=2) k1, k2 = key_states.chunk(2, dim=2) # Apply RoPE independently to each half q1, k1 = apply_rotary_pos_emb(q1, k1, cos, sin) q2, k2 = apply_rotary_pos_emb(q2, k2, cos, sin) # Recombine the halves query_states = torch.cat([q1, q2], dim=2) key_states = torch.cat([k1, k2], dim=2) else: query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) query_states = query_states * _get_llama_4_attn_scale( cache_position, self.config.rope_parameters.get("llama_4_scaling_beta"), self.config.rope_parameters.get("original_max_position_embeddings"), ).to(query_states.dtype) if past_key_values is not None: # sin and cos are specific to RoPE models; cache_position needed for the static cache cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs) key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) if self.block_diff_mask is None or q_len != self.block_diff_mask.shape[-2]: block_mask = self.compute_block_mask(mode='block_diff', block_size=self.block_size, q_len=q_len) else: block_mask = self.block_diff_mask attn_output = fused_flex_attention(query_states, key_states, value_states, block_mask=block_mask) attn_output = attn_output.transpose(1, 2).reshape(*input_shape, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output, None class NemotronLabsDiffusionModel(Ministral3PreTrainedModel, GenerationMixin): """ A single model with: - a bidirectional encoder + diffusion‐LM head over A - a causal decoder + LM head over B, conditioned on F_A """ def __init__(self, config: NemotronLabsDiffusionConfig): super().__init__(config) self.mask_token_id = config.mask_token_id diffusion_config = copy.deepcopy(config) diffusion_config.diffusion_lm = True if config.dlm_paradigm == 'block_diff': diffusion_config.attn_class = NemotronLabsDiffusionFlexAttention elif config.dlm_paradigm in ['bidirectional', 'autoregressive']: diffusion_config.attn_class = Ministral3Attention if config.dlm_paradigm == 'autoregressive': diffusion_config.diffusion_lm = False else: raise ValueError(f"Unsupported DLM paradigm: {config.dlm_paradigm}") self.encoder = Ministral3Model(diffusion_config) self.diffusion_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) self.vocab_size = config.vocab_size self.post_init() def get_input_embeddings(self): return self.encoder.embed_tokens def set_input_embeddings(self, value): self.encoder.embed_tokens = value def get_output_embeddings(self): return self.diffusion_head def set_output_embeddings(self, new_embeddings): self.diffusion_head = new_embeddings def forward_process(self, input_ids, eps=1e-3, block_size=None, loss_mask=None): b, l = input_ids.shape device = input_ids.device if self.config.dp_varying_mask_ratio: # Enable different random seeds for each DP rank during sampling import torch.distributed as dist dp_rank = 0 if dist.is_initialized(): try: dp_rank = dist.get_rank() except Exception: dp_rank = 0 # Use a local generator to avoid affecting global RNG state generator = torch.Generator(device=device) generator.manual_seed(torch.seed() + dp_rank) else: generator = None t = torch.rand(b, device=device, generator=generator) p_mask = (1 - eps) * t + eps # shape: (b,) p_mask = p_mask[:, None].expand(-1, l) # shape: (b, l) masked_indices = torch.rand((b, l), device=device) < p_mask if loss_mask is not None: masked_indices[loss_mask == 0] = 0 noisy_batch = torch.where(masked_indices, self.mask_token_id, input_ids) return noisy_batch, masked_indices, p_mask def forward( self, input_ids: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, labels: Optional[torch.LongTensor] = None, split_len: Optional[int] = None, past_key_values: Optional[Cache] = None, block_size: Optional[int] = None, eps: float = 1e-3, is_teacher: bool = False, masked_indices: Optional[torch.Tensor] = None, p_mask: Optional[torch.Tensor] = None, teacher_logits: Optional[torch.Tensor] = None, masked_indices_teacher: Optional[torch.Tensor] = None, loss_mask: Optional[torch.Tensor] = None, ce_loss_weight: float = 1.0, output_last_hidden_states_only: bool = False, skip_loss: bool = False, **kwargs, ) -> CausalLMOutputWithPast: batch_size, seq_len = input_ids.shape if self.config.dlm_paradigm == 'block_diff': if labels is not None and block_size is None: block_size = self.config.block_size elif self.config.dlm_paradigm not in ('bidirectional', 'autoregressive'): raise ValueError(f"Unknown dLM paradigm: {self.config.dlm_paradigm}") if labels is not None and self.config.dlm_paradigm != 'autoregressive': if masked_indices is not None: # assert p_mask is not None if loss_mask is not None: masked_indices[loss_mask == 0] = 0 noisy_inputs = torch.where(masked_indices, self.mask_token_id, input_ids) else: noisy_inputs, masked_indices, p_mask = self.forward_process(input_ids, eps=eps, block_size=block_size, loss_mask=loss_mask) else: noisy_inputs = input_ids masked_indices = None p_mask = None input_ids_len = noisy_inputs.shape[1] if labels is not None and self.config.dlm_paradigm == 'block_diff': if position_ids is None: position_ids = torch.arange(input_ids_len, device=noisy_inputs.device).unsqueeze(0) noisy_inputs = torch.cat([noisy_inputs, input_ids], dim=1) enc_out = self.encoder( past_key_values=past_key_values, input_ids=noisy_inputs, attention_mask=attention_mask, position_ids=position_ids, is_training=(labels is not None), **kwargs, ) if output_last_hidden_states_only: return BaseModelOutput(last_hidden_state=enc_out.last_hidden_state) logits = self.diffusion_head(enc_out.last_hidden_state) # (batch, len_B, vocab) causal_logits = None if labels is not None and self.config.dlm_paradigm == 'block_diff': causal_logits = logits[:, input_ids_len:] logits = logits[:, :input_ids_len] loss = None if labels is not None and not skip_loss: if self.config.dlm_paradigm == 'autoregressive': shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() if loss_mask is None: loss_fct = CrossEntropyLoss() shift_logits = shift_logits.view(-1, shift_logits.size(-1)) shift_labels = shift_labels.view(-1) loss = loss_fct(shift_logits, shift_labels) else: loss_mask = loss_mask[..., 1:].contiguous() loss_fct = CrossEntropyLoss(reduction='none') shift_logits = shift_logits.view(-1, shift_logits.size(-1)) shift_labels = shift_labels.view(-1) shift_labels = shift_labels.to(shift_logits.device) token_losses = loss_fct(shift_logits, shift_labels) flat_loss_mask = loss_mask.reshape(-1) loss = token_losses[flat_loss_mask == 1].sum() / flat_loss_mask.sum() else: # LLaDA-style diffusion loss on masked positions. # Token-wise cross entropy loss on masked positions. token_loss = torch.nn.functional.cross_entropy( logits[masked_indices], labels[masked_indices], reduction='none' ) / p_mask[masked_indices] num_mask_tokens = masked_indices.sum() # global_loss_avg=True: loss is reduced externally by global token count. loss = token_loss.sum() if self.config.dlm_loss_weight is not None: loss = self.config.dlm_loss_weight * loss if self.config.dlm_paradigm == 'block_diff': # AR-side loss for block-diffusion paradigm. causal_logits = causal_logits[..., :-1, :].contiguous() causal_logits = causal_logits.view(-1, causal_logits.size(-1)) causal_labels = labels[..., 1:].contiguous().view(-1) loss_fct = CrossEntropyLoss(reduction='sum') ar_loss = loss_fct(causal_logits, causal_labels) self.loss_diffusion = loss.detach().item() / num_mask_tokens self.loss_ar = ar_loss.detach().item() / seq_len loss = loss + self.config.ar_loss_weight * ar_loss # global_loss_avg=True: return (sum_loss, token_count) for external mean. if self.config.dlm_paradigm == 'block_diff': loss = (loss, num_mask_tokens + int(self.config.ar_loss_weight * seq_len)) else: loss = (loss, num_mask_tokens) return NemotronLabsDiffusionOutputWithPast( loss=loss if not is_teacher else logits, logits=logits, causal_logits=causal_logits, past_key_values=enc_out.past_key_values, hidden_states=None, attentions=None, ) @torch.no_grad() def generate( self, prompt_ids: torch.Tensor, max_new_tokens: int, block_length: int, threshold: Optional[float] = None, causal_context: bool = True, temperature: float = 0.0, eos_token_id: Optional[int] = None, max_thinking_tokens: Optional[int] = None, end_think_token_id: Optional[int] = None, ): """Block-wise diffusion decoding with prefix-cached KV (LLaDA-style). Each block: append `block_length` mask tokens, then iteratively unmask by confidence top-k (with optional threshold). When `causal_context`, the KV cache and the next-block seed are produced via a causal forward between blocks (flipping `self_attn.diffusion_lm`), matching the AR objective at block boundaries. Returns (output_ids, nfe) — output_ids includes the prompt. """ if eos_token_id is None: eos_token_id = getattr(self.config, "eos_token_id", None) mask_id = self.mask_token_id x_accum = prompt_ids.clone() B = prompt_ids.shape[0] assert max_new_tokens % block_length == 0 num_blocks = max_new_tokens // block_length # one denoising step per generated token (matches legacy chat_utils call) steps_per_block = block_length nfe = 0 def _set_diffusion_lm(val: bool): for layer in self.encoder.layers: if hasattr(layer.self_attn, "diffusion_lm"): layer.self_attn.diffusion_lm = val # Initial causal prefill produces the KV cache and the next-block seed. if causal_context: _set_diffusion_lm(False) output = self(prompt_ids, use_cache=True, use_causal_mask=causal_context) past_key_values = output.past_key_values if causal_context: _set_diffusion_lm(True) next_token = None if causal_context: last_logit = output.logits[:, -1, :] if temperature > 0: next_token = torch.multinomial(torch.softmax(last_logit / temperature, dim=-1), num_samples=1) else: next_token = torch.argmax(last_logit, dim=-1, keepdim=True) for num_block in range(num_blocks): mask_block = torch.full( (B, block_length), mask_id, dtype=prompt_ids.dtype, device=prompt_ids.device, ) if causal_context: mask_block[:, 0] = next_token[:, 0] x_accum = torch.cat([x_accum, mask_block], dim=1) block_start = prompt_ids.size(1) + num_block * block_length block_slice = slice(block_start, block_start + block_length) # Thinking-budget enforcement: if we've passed max_thinking_tokens # without an end-think marker, inject one into this block. if end_think_token_id is not None and max_thinking_tokens is not None: tokens_before = num_block * block_length tokens_after = tokens_before + block_length if tokens_after > max_thinking_tokens: gen_so_far = x_accum[:, prompt_ids.size(1):block_start] has_end_think = ( (gen_so_far == end_think_token_id).any(dim=1) if gen_so_far.size(1) > 0 else torch.zeros(B, dtype=torch.bool, device=prompt_ids.device) ) if not has_end_think.all(): offset = max(0, max_thinking_tokens - tokens_before) inject_pos = block_start + offset for b in range(B): if not has_end_think[b]: x_accum[b, inject_pos] = end_think_token_id mask_block_idx0 = x_accum[:, block_slice] == mask_id num_transfer_tokens = _get_num_transfer_tokens(mask_block_idx0, steps_per_block) # Denoise the current block by repeated confidence-based unmasking. for i in range(steps_per_block): mask_block_idx = x_accum[:, block_slice] == mask_id if mask_block_idx.sum() == 0: break nfe += 1 logits_block = self( x_accum[:, block_slice], past_key_values=past_key_values, use_cache=False, ).logits x0, transfer_idx = _get_transfer_index( logits_block, temperature, mask_block_idx, x_accum[:, block_slice], num_transfer_tokens=num_transfer_tokens[:, i], threshold=threshold, ) cur = x_accum[:, block_slice].clone() cur[transfer_idx] = x0[transfer_idx] x_accum[:, block_slice] = cur if eos_token_id is not None: block_tokens = x_accum[:, block_slice] eos_mask = block_tokens == eos_token_id if eos_mask.any(dim=1).any(): after_eos = eos_mask.cumsum(dim=1).bool() mask_before = (block_tokens == mask_id) & ~after_eos if (eos_mask.any(dim=1) & ~mask_before.any(dim=1)).any(): break # Post-block: causal forward over the block to update the KV cache # and (when causal_context) sample the seed for the next block. if causal_context: _set_diffusion_lm(False) output = self( x_accum[:, block_slice], past_key_values=past_key_values, use_cache=True, use_causal_mask=causal_context, ) past_key_values = output.past_key_values nfe += 1 if causal_context: _set_diffusion_lm(True) last_logit = output.logits[:, -1, :] if temperature > 0: next_token = torch.multinomial(torch.softmax(last_logit / temperature, dim=-1), num_samples=1) else: next_token = torch.argmax(last_logit, dim=-1, keepdim=True) if eos_token_id is not None: gen_so_far = x_accum[:, prompt_ids.size(1):] is_eos = gen_so_far == eos_token_id if is_eos.any(dim=1).all(): first_eos = is_eos.to(torch.int64).argmax(dim=1) max_eos = first_eos.max().item() return x_accum[:, : prompt_ids.size(1) + max_eos + 1], nfe return x_accum, nfe @torch.no_grad() def ar_generate( self, prompt_ids: torch.Tensor, max_new_tokens: int = 128, temperature: float = 0.0, eos_token_id: Optional[int] = None, max_thinking_tokens: Optional[int] = None, end_think_token_id: Optional[int] = None, ) -> tuple: """Autoregressive generation calling the encoder directly (injected by build_hf_tidar_repo). Bypasses NemotronLabsDiffusionModel.forward() to avoid diffusion-specific code paths. Calls self.encoder (Ministral3Model) with explicit cache_position, position_ids, and use_cache so the KV cache and causal masking behave identically to MistralForCausalLM / vLLM. Returns: (output_ids, nfe) where output_ids includes the prompt. """ for layer in self.encoder.layers: if hasattr(layer.self_attn, 'diffusion_lm'): layer.self_attn.diffusion_lm = False if eos_token_id is None: eos_token_id = getattr(self.config, 'eos_token_id', None) device = prompt_ids.device batch_size, prompt_len = prompt_ids.shape past_key_values = DynamicCache() cache_position = torch.arange(prompt_len, device=device) position_ids = cache_position.unsqueeze(0).expand(batch_size, -1) enc_out = self.encoder( input_ids=prompt_ids, position_ids=position_ids, past_key_values=past_key_values, use_cache=True, cache_position=cache_position, ) past_key_values = enc_out.past_key_values next_logit = self.diffusion_head(enc_out.last_hidden_state[:, -1:, :]).squeeze(1) generated_tokens = [] nfe = 0 for step in range(max_new_tokens): nfe += 1 if temperature > 0: probs = torch.softmax(next_logit / temperature, dim=-1) next_token = torch.multinomial(probs, num_samples=1) else: next_token = torch.argmax(next_logit, dim=-1, keepdim=True) # ---- thinking budget enforcement ---- if end_think_token_id is not None and max_thinking_tokens is not None: if step >= max_thinking_tokens: if generated_tokens: gen_tensor = torch.cat(generated_tokens, dim=1) has_end_think = (gen_tensor == end_think_token_id).any(dim=1) else: has_end_think = torch.zeros(batch_size, dtype=torch.bool, device=device) for b in range(batch_size): if not has_end_think[b]: next_token[b] = end_think_token_id generated_tokens.append(next_token) if eos_token_id is not None and (next_token == eos_token_id).all(): break if step < max_new_tokens - 1: cur_pos = prompt_len + step step_cache_pos = torch.tensor([cur_pos], device=device) step_pos_ids = step_cache_pos.unsqueeze(0).expand(batch_size, -1) enc_out = self.encoder( input_ids=next_token, position_ids=step_pos_ids, past_key_values=past_key_values, use_cache=True, cache_position=step_cache_pos, ) past_key_values = enc_out.past_key_values next_logit = self.diffusion_head(enc_out.last_hidden_state[:, -1:, :]).squeeze(1) all_generated = torch.cat(generated_tokens, dim=1) output_ids = torch.cat([prompt_ids, all_generated], dim=1) return output_ids, nfe @torch.no_grad() def linear_spec_generate( self, prompt_ids: torch.Tensor, max_new_tokens: int = 128, block_length: int = 32, temperature: float = 0.0, mask_token_id: Optional[int] = None, eos_token_id: Optional[int] = None, max_thinking_tokens: Optional[int] = None, end_think_token_id: Optional[int] = None, threshold: float = 0.0, ): """Linear speculative decoding: diffusion draft + AR verify. Each iteration: (1) draft the next block under bidirectional attention, (2) verify the drafted block under causal attention, accept the longest prefix where draft matches AR + one bonus token, advance the KV cache. LoRA-aware: when a PEFT adapter is attached to the model (e.g. ``linear_spec_lora``), it is toggled ON for the bidirectional draft phase and OFF for the causal prefill / verify phases — so the adapter only specializes the diffusion-mode forward and AR semantics are preserved. With no adapter loaded the calls are no-ops. Returns ``(output_ids, nfe)`` — ``output_ids`` includes the prompt. """ if prompt_ids.shape[0] != 1: raise ValueError("Linear speculative decoding requires batch_size == 1") token_mask_id = mask_token_id if mask_token_id is not None else self.config.mask_token_id if eos_token_id is None: eos_token_id = getattr(self.config, "eos_token_id", None) device = prompt_ids.device def _set_diffusion_lm(val: bool): for layer in self.encoder.layers: if hasattr(layer.self_attn, "diffusion_lm"): layer.self_attn.diffusion_lm = val def _toggle_adapters(enable: bool): # No-op when no PEFT/LoRA modules are attached. for module in self.modules(): if hasattr(module, "_disable_adapters"): module._disable_adapters = not enable # Prefill (causal, LoRA OFF). _set_diffusion_lm(False) _toggle_adapters(False) enc_out = self.encoder( input_ids=prompt_ids, past_key_values=DynamicCache(), use_cache=True, use_causal_mask=True, ) past_key_values = enc_out.past_key_values last_logit = self.diffusion_head(enc_out.last_hidden_state[:, -1:, :]).squeeze(1) nfe = 1 if temperature > 0: next_token = torch.multinomial(torch.softmax(last_logit / temperature, dim=-1), num_samples=1) else: next_token = torch.argmax(last_logit, dim=-1, keepdim=True) if eos_token_id is not None and next_token.item() == eos_token_id: return torch.cat([prompt_ids, next_token], dim=1), nfe generated = [next_token] total_gen = 1 while total_gen < max_new_tokens: cache_len = past_key_values.get_seq_length() block = torch.full((1, block_length), token_mask_id, dtype=torch.long, device=device) block[0, 0] = next_token.item() # Draft phase (bidirectional, LoRA ON) — iterate at threshold>0 so # that even low-confidence blocks make progress. _set_diffusion_lm(True) _toggle_adapters(True) while True: is_mask = block == token_mask_id if not is_mask.any(): break enc_out = self.encoder(input_ids=block, past_key_values=past_key_values, use_cache=False) nfe += 1 draft_logits = self.diffusion_head(enc_out.last_hidden_state) # LLaDA: logit[i] directly predicts position i — no shift needed. if temperature > 0: draft_probs = torch.softmax(draft_logits / temperature, dim=-1) draft_tokens = torch.multinomial( draft_probs.view(-1, draft_probs.shape[-1]), num_samples=1 ).view(1, block_length) else: draft_tokens = draft_logits.argmax(dim=-1) draft_probs = torch.softmax(draft_logits, dim=-1) if threshold > 0: draft_conf = torch.gather(draft_probs, -1, draft_tokens.unsqueeze(-1)).squeeze(-1) draft_conf = torch.where(is_mask, draft_conf, -torch.inf) unmask = draft_conf >= threshold # Force progress even when every masked position is below threshold. if not unmask.any(): best_idx = draft_conf.view(-1).argmax() unmask = torch.zeros_like(is_mask, dtype=torch.bool) unmask.view(-1)[best_idx] = True block[unmask] = draft_tokens[unmask] else: block[is_mask] = draft_tokens[is_mask] break # Verify phase (causal, LoRA OFF). _set_diffusion_lm(False) _toggle_adapters(False) enc_out = self.encoder( input_ids=block, past_key_values=past_key_values, use_cache=True, use_causal_mask=True, ) past_key_values = enc_out.past_key_values nfe += 1 verify_logits = self.diffusion_head(enc_out.last_hidden_state) if temperature > 0: ar_tokens = torch.multinomial( torch.softmax(verify_logits / temperature, dim=-1).view(-1, verify_logits.shape[-1]), num_samples=1, ).view(1, block_length) else: ar_tokens = verify_logits.argmax(dim=-1) # Accept consecutive matches; AR also gives one bonus token at the tail. accepted = 0 for i in range(block_length - 1): if ar_tokens[0, i].item() == block[0, i + 1].item(): accepted += 1 else: break accepted += 1 accepted_toks = ar_tokens[:, :accepted] generated.append(accepted_toks) total_gen += accepted _crop_dynamic_cache(past_key_values, cache_len + accepted) next_token = ar_tokens[:, accepted - 1 : accepted] if eos_token_id is not None: eos_pos = (accepted_toks[0] == eos_token_id).nonzero(as_tuple=True)[0] if len(eos_pos) > 0: first_eos = eos_pos[0].item() generated[-1] = accepted_toks[:, : first_eos + 1] total_gen = total_gen - accepted + first_eos + 1 break # Thinking-budget enforcement: force end-think as next seed if budget exhausted. if end_think_token_id is not None and max_thinking_tokens is not None: if total_gen > max_thinking_tokens: all_gen = torch.cat(generated, dim=1) if not (all_gen == end_think_token_id).any(): next_token = torch.tensor([[end_think_token_id]], device=device) if total_gen >= max_new_tokens: break all_generated = torch.cat(generated, dim=1) output_ids = torch.cat([prompt_ids, all_generated], dim=1) return output_ids, nfe # ─── Module-level helpers used by `generate` and `linear_spec_generate` ── def _crop_dynamic_cache(past_key_values: DynamicCache, max_length: int): """Crop a DynamicCache to max_length, compatible with both old and new transformers.""" if hasattr(past_key_values, 'crop'): past_key_values.crop(max_length) else: for layer_idx in range(len(past_key_values)): past_key_values.key_cache[layer_idx] = past_key_values.key_cache[layer_idx][:, :, :max_length] past_key_values.value_cache[layer_idx] = past_key_values.value_cache[layer_idx][:, :, :max_length] past_key_values._seen_tokens = max_length def _add_gumbel_noise(logits, temperature): """Gumbel-max sampling in float64 (low-precision Gumbel hurts MDM quality).""" if temperature == 0: return logits logits = logits.to(torch.float64) noise = torch.rand_like(logits, dtype=torch.float64) gumbel_noise = (- torch.log(noise)) ** temperature return logits.exp() / gumbel_noise def _get_num_transfer_tokens(mask_index, steps: int): """Even split of masked positions across `steps`, with remainder front-loaded.""" mask_num = mask_index.sum(dim=1, keepdim=True) base = mask_num // steps remainder = mask_num % steps num_transfer_tokens = torch.zeros(mask_num.size(0), steps, device=mask_index.device, dtype=torch.int64) + base for i in range(mask_num.size(0)): num_transfer_tokens[i, : int(remainder[i])] += 1 return num_transfer_tokens def _get_transfer_index(logits, temperature, mask_index, x, num_transfer_tokens, threshold=None): """Pick which masked positions to commit this denoising step. Returns (x0, transfer_index): x0 is argmax tokens (clamped to original x at non-masked positions); transfer_index is a bool mask over positions to finalize, chosen by top-k confidence (and filtered by `threshold` if given). """ logits_with_noise = _add_gumbel_noise(logits, temperature=temperature) x0 = torch.argmax(logits_with_noise, dim=-1) p = F.softmax(logits, dim=-1) x0_p = torch.squeeze(torch.gather(p, dim=-1, index=torch.unsqueeze(x0, -1)), -1) x0 = torch.where(mask_index, x0, x) confidence = torch.where(mask_index, x0_p, -np.inf) transfer_index = torch.zeros_like(x0, dtype=torch.bool, device=x0.device) if threshold is not None: num_transfer_tokens = mask_index.sum(dim=1, keepdim=True) for j in range(confidence.shape[0]): _, select_index = torch.topk(confidence[j], k=num_transfer_tokens[j]) transfer_index[j, select_index] = True if threshold is not None: for k in range(1, num_transfer_tokens[j]): if confidence[j, select_index[k]] < threshold: transfer_index[j, select_index[k]] = False return x0, transfer_index def gumbel_topk(log_w: torch.Tensor, k: int) -> torch.Tensor: """Return a Bool mask of length len(log_w) with exactly k True.""" g = -torch.log(-torch.log(torch.rand_like(log_w) + 1e-9) + 1e-9) topk = torch.topk(log_w + g, k).indices mask = torch.zeros_like(log_w, dtype=torch.bool) mask[topk] = True return mask