generation_utils.py

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