"""Block diffusion generation utilities for Dream models.""" from __future__ import annotations from dataclasses import dataclass from typing import List, Optional, Sequence, Union import torch from torch.nn import functional as F from transformers.cache_utils import DynamicCache from transformers.utils import ModelOutput def top_k_logits(logits: torch.Tensor, k: int) -> torch.Tensor: if k <= 0: return logits values, _ = torch.topk(logits, k) min_values = values[..., -1, None] return torch.where(logits < min_values, torch.full_like(logits, float('-inf')), logits) def top_p_logits(logits: torch.Tensor, p: float) -> torch.Tensor: sorted_logits, sorted_indices = torch.sort(logits, descending=True) cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) sorted_mask = cumulative_probs > p sorted_mask[..., 1:] = sorted_mask[..., :-1].clone() sorted_mask[..., 0] = False mask_indices = torch.scatter( torch.full_like(logits, False, dtype=torch.bool), -1, sorted_indices, sorted_mask, ) return logits.masked_fill(mask_indices, float('-inf')) def sample_with_temperature_topk_topp( logits: torch.Tensor, temperature: float = 1.0, top_k: int = 0, top_p: float = 1.0, ) -> tuple[torch.Tensor, torch.Tensor]: orig_shape = logits.shape[:-1] vocab_size = logits.shape[-1] logits = logits.reshape(-1, vocab_size) if temperature > 0: logits = logits / temperature if top_k > 0: logits = top_k_logits(logits, top_k) if top_p < 1.0: logits = top_p_logits(logits, top_p) probs = F.softmax(logits, dim=-1) if temperature > 0: token = torch.multinomial(probs, num_samples=1) else: token = probs.argmax(dim=-1, keepdim=True) token_prob = torch.gather(probs, -1, token) return token.view(*orig_shape), token_prob.view(*orig_shape) def get_num_transfer_tokens(block_length: int, steps: int) -> torch.Tensor: base = block_length // steps remainder = block_length % steps num_transfer_tokens = torch.zeros(steps, dtype=torch.int64) + base num_transfer_tokens[:remainder] += 1 return num_transfer_tokens def build_block_diffusion_attention_mask( num_blocks: int, block_length: int, device: torch.device, batch_size: int = 1, ) -> torch.Tensor: block_mask = torch.tril(torch.ones(num_blocks, num_blocks, device=device)) return block_mask.repeat_interleave(block_length, dim=0).repeat_interleave(block_length, dim=1).unsqueeze(0).expand( batch_size, -1, -1 ) def _resolve_stopping_ids(stopping_criteria_idx: Optional[Union[int, Sequence[int]]]) -> Optional[List[int]]: if stopping_criteria_idx is None: return None if isinstance(stopping_criteria_idx, int): return [stopping_criteria_idx] return list(stopping_criteria_idx) def _should_stop( generated_ids: torch.Tensor, prompt_length: int, stopping_criteria_idx: Optional[List[int]], ) -> bool: if not stopping_criteria_idx: return False gen_part = generated_ids[:, prompt_length:] return any((gen_part == stop_id).any().item() for stop_id in stopping_criteria_idx) def _default_use_kv_cache(model: torch.nn.Module) -> bool: """Dream / Dream1 models use prefix KV cache during block diffusion decode.""" model_type = getattr(model.config, 'model_type', None) if model_type is None: return False return model_type.lower() in ('dream', 'dream1') def _select_transfer_index( remasking_strategy: str, mask_index: torch.Tensor, x0: torch.Tensor, x0_p: torch.Tensor, num_transfer_tokens: torch.Tensor, step: int, confidence_threshold: float, eb_threshold: Optional[float], *, force_accept: bool = False, ) -> torch.Tensor: if force_accept: return mask_index.clone() if remasking_strategy == 'sequential': transfer_index = torch.zeros_like(x0, dtype=torch.bool) for j in range(x0.shape[0]): if not mask_index[j].any(): continue first_mask_index = mask_index[j].nonzero(as_tuple=True)[0].min().item() end = first_mask_index + int(num_transfer_tokens[step].item()) transfer_index[j, first_mask_index:end] = True return transfer_index if remasking_strategy == 'low_confidence_static': confidence = torch.where(mask_index, x0_p, -torch.inf) transfer_index = torch.zeros_like(x0, dtype=torch.bool) k = max(1, int(num_transfer_tokens[step].item())) for j in range(confidence.shape[0]): _, idx = torch.topk(confidence[j], k) transfer_index[j, idx] = True return transfer_index if remasking_strategy == 'low_confidence_dynamic': confidence = torch.where(mask_index, x0_p, -torch.inf) transfer_index = torch.zeros_like(x0, dtype=torch.bool) k = max(1, int(num_transfer_tokens[step].item())) for j in range(confidence.shape[0]): high_conf_mask = confidence[j] > confidence_threshold if int(high_conf_mask.sum().item()) >= k: transfer_index[j] = high_conf_mask else: _, idx = torch.topk(confidence[j], k) transfer_index[j, idx] = True return transfer_index if remasking_strategy == 'entropy_bounded': if eb_threshold is None: raise ValueError('eb_threshold is required for entropy_bounded remasking.') eps = 1e-12 entropies = -(x0_p.clamp_min(eps) * x0_p.clamp_min(eps).log()) entropies = torch.where(mask_index, entropies, torch.inf) ent_sorted, order = torch.sort(entropies, dim=1, descending=False) cumsum = torch.cumsum(ent_sorted, dim=1) transfer_index = torch.zeros_like(x0, dtype=torch.bool) for j in range(x0_p.shape[0]): k = torch.searchsorted( cumsum[j], torch.tensor(eb_threshold, device=x0_p.device), right=False ).item() k = max(1, min(k, int(mask_index[j].sum().item()))) transfer_index[j, order[j, :k]] = True return transfer_index raise ValueError(f'Unknown remasking strategy: {remasking_strategy}') def _denoise_current_block( model: torch.nn.Module, x: torch.Tensor, num_block: int, block_length: int, mask_id: int, block_diffusion_attention_mask: torch.Tensor, position_ids: torch.Tensor, denoising_steps: int, num_transfer_tokens: torch.Tensor, temperature: float, top_k: int, top_p: float, remasking_strategy: str, confidence_threshold: float, eb_threshold: Optional[float], *, use_kv_cache: bool, past_key_values: Optional[DynamicCache], ) -> tuple[torch.Tensor, Optional[DynamicCache], int]: block_start = num_block * block_length block_end = block_start + block_length cur_x = x[:, block_start:block_end].clone() nfe = 0 for step in range(denoising_steps + 1): mask_index = cur_x == mask_id if mask_index.sum() == 0: if use_kv_cache: cur_attn_mask = block_diffusion_attention_mask[:, block_start:block_end, :block_end] cur_position_ids = position_ids[:, block_start:block_end] model( cur_x, attention_mask=cur_attn_mask, position_ids=cur_position_ids, past_key_values=past_key_values, use_cache=True, store_kv=True, ) nfe += 1 break force_accept = step == denoising_steps - 1 if use_kv_cache: cur_attn_mask = block_diffusion_attention_mask[:, block_start:block_end, :block_end] cur_position_ids = position_ids[:, block_start:block_end] logits = model( cur_x, attention_mask=cur_attn_mask, position_ids=cur_position_ids, past_key_values=past_key_values, use_cache=True, store_kv=False, ).logits else: seq_end = block_end attn_mask = block_diffusion_attention_mask[:, :seq_end, :seq_end] pos_ids = position_ids[:, :seq_end] logits = model( x[:, :seq_end], attention_mask=attn_mask, position_ids=pos_ids, use_cache=False, ).logits[:, block_start:block_end] nfe += 1 x0, x0_p = sample_with_temperature_topk_topp( logits, temperature=temperature, top_k=top_k, top_p=top_p, ) x0 = torch.where(mask_index, x0, cur_x) transfer_index = _select_transfer_index( remasking_strategy, mask_index, x0, x0_p, num_transfer_tokens, step, confidence_threshold, eb_threshold, force_accept=force_accept, ) cur_x[transfer_index] = x0[transfer_index] if not use_kv_cache: x[:, block_start:block_end] = cur_x return cur_x, past_key_values, nfe @dataclass class BlockDiffusionOutput(ModelOutput): sequences: torch.LongTensor = None nfe: Optional[int] = None logits: Optional[tuple] = None @torch.no_grad() def block_diffusion_generate( model: torch.nn.Module, input_ids: torch.LongTensor, mask_id: int, gen_length: int = 128, block_length: Optional[int] = None, denoising_steps: Optional[int] = None, temperature: float = 0.0, top_k: int = 0, top_p: float = 1.0, remasking_strategy: str = 'low_confidence_dynamic', confidence_threshold: float = 0.9, eb_threshold: Optional[float] = 0.35, stopping_criteria_idx: Optional[Union[int, Sequence[int]]] = None, use_kv_cache: Optional[bool] = None, return_dict_in_generate: bool = False, ) -> Union[torch.LongTensor, BlockDiffusionOutput]: """Block-wise diffusion decoding with optional prefix KV cache.""" model.eval() if input_ids.dim() != 2: raise ValueError(f'input_ids must be 2D, got shape {tuple(input_ids.shape)}') device = input_ids.device batch_size, prompt_length = input_ids.shape block_length = block_length or getattr(model.config, 'block_size', 4) if denoising_steps is None: denoising_steps = 1 if remasking_strategy == 'low_confidence_static' else block_length stopping_criteria_idx = _resolve_stopping_ids(stopping_criteria_idx) if use_kv_cache is None: use_kv_cache = _default_use_kv_cache(model) num_blocks = (prompt_length + gen_length + block_length - 1) // block_length total_length = num_blocks * block_length block_diffusion_attention_mask = build_block_diffusion_attention_mask( num_blocks, block_length, device, batch_size=batch_size ) position_ids = torch.arange(total_length, device=device, dtype=torch.long).unsqueeze(0).expand(batch_size, -1) x = torch.full((batch_size, total_length), mask_id, dtype=input_ids.dtype, device=device) x[:, :prompt_length] = input_ids prefill_blocks = prompt_length // block_length prefill_length = prefill_blocks * block_length past_key_values = DynamicCache() if use_kv_cache else None nfe = 0 if use_kv_cache and prefill_length > 0: cur_x = x[:, :prefill_length] cur_attn_mask = block_diffusion_attention_mask[:, :prefill_length, :prefill_length] cur_position_ids = position_ids[:, :prefill_length] model( cur_x, attention_mask=cur_attn_mask, position_ids=cur_position_ids, past_key_values=past_key_values, use_cache=True, store_kv=True, ) nfe += 1 num_transfer_tokens = get_num_transfer_tokens(block_length, denoising_steps) for num_block in range(prefill_blocks, num_blocks): cur_x, past_key_values, block_nfe = _denoise_current_block( model, x, num_block, block_length, mask_id, block_diffusion_attention_mask, position_ids, denoising_steps, num_transfer_tokens, temperature, top_k, top_p, remasking_strategy, confidence_threshold, eb_threshold, use_kv_cache=use_kv_cache, past_key_values=past_key_values, ) nfe += block_nfe x[:, num_block * block_length:(num_block + 1) * block_length] = cur_x if _should_stop(x, prompt_length, stopping_criteria_idx): break output_length = min(total_length, prompt_length + gen_length) x = x[:, :output_length] if return_dict_in_generate: return BlockDiffusionOutput(sequences=x, nfe=nfe) return x # HF kwargs that ``generate()`` strips before calling ``block_diffusion_generate()``. _UNSUPPORTED_HF_KEYS = ( 'stopping_criteria', 'num_return_sequences', 'num_beams', 'num_beam_groups', 'penalty_alpha', 'use_cache', 'output_logits', 'output_scores', 'output_attentions', 'output_hidden_states', 'return_legacy_cache', 'synced_gpus', 'streamer', 'logits_processor', 'logits_warper', 'generation_config', 'tokenizer', 'min_length', 'min_new_tokens', 'pad_token_id', 'bos_token_id', 'eos_token_id', ) class BlockDiffusionGenerationMixin: def _resolve_generation_mode(self, generation_mode: Optional[str] = None) -> str: if generation_mode is not None: return generation_mode return getattr(self.config, 'generation_mode', 'block_diffusion') @torch.no_grad() def generate( self, input_ids: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None, generation_mode: Optional[str] = None, **kwargs, ) -> Union[torch.LongTensor, BlockDiffusionOutput, ModelOutput]: """HF-compatible ``generate`` entry point with pluggable decoding modes. Supported modes (``generation_mode`` kwarg or ``config.generation_mode``): - ``block_diffusion`` (default): calls :meth:`block_diffusion_generate`. - ``autoregressive``: delegates to ``GenerationMixin.generate``. """ mode = self._resolve_generation_mode(kwargs.pop('generation_mode', generation_mode)) return_dict_in_generate = kwargs.pop('return_dict_in_generate', False) if mode == 'autoregressive': return super().generate( input_ids=input_ids, attention_mask=attention_mask, return_dict_in_generate=return_dict_in_generate, **kwargs, ) if mode != 'block_diffusion': raise ValueError(f'Unknown generation_mode: {mode!r}. Supported: block_diffusion, autoregressive.') # HF-only normalizations; sampling params pass through to block_diffusion_generate. if kwargs.pop('do_sample', None) is False: kwargs['temperature'] = 0.0 if 'max_new_tokens' not in kwargs and (max_length := kwargs.pop('max_length', None)) is not None: kwargs['max_new_tokens'] = max(max_length - input_ids.shape[-1], 0) for key in _UNSUPPORTED_HF_KEYS: kwargs.pop(key, None) return self.block_diffusion_generate( input_ids=input_ids, return_dict_in_generate=return_dict_in_generate, **kwargs, ) @torch.no_grad() def block_diffusion_generate( self, input_ids: torch.LongTensor, max_new_tokens: int = 128, temperature: float = 0.0, top_k: int = 0, top_p: float = 1.0, return_dict_in_generate: bool = False, # block-diffusion specific block_length: Optional[int] = None, denoising_steps: Optional[int] = None, remasking_strategy: str = 'low_confidence_dynamic', confidence_threshold: float = 0.9, eb_threshold: Optional[float] = 0.35, use_kv_cache: Optional[bool] = None, mask_token_id: Optional[int] = None, ) -> Union[torch.LongTensor, BlockDiffusionOutput]: mask_token_id = mask_token_id if mask_token_id is not None else self.config.mask_token_id if mask_token_id is None: raise ValueError('mask_token_id must be provided or set in model.config.mask_token_id') stopping_criteria_idx = self.config.eos_token_id if getattr(self.config, 'eos_token_id', None) is not None else None return block_diffusion_generate( self, input_ids=input_ids, mask_id=mask_token_id, gen_length=max_new_tokens, block_length=block_length, denoising_steps=denoising_steps, temperature=temperature, top_k=top_k, top_p=top_p, remasking_strategy=remasking_strategy, confidence_threshold=confidence_threshold, eb_threshold=eb_threshold, stopping_criteria_idx=stopping_criteria_idx, use_kv_cache=use_kv_cache, return_dict_in_generate=return_dict_in_generate, )