modeling_ministral_dlm.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.

import copy
from dataclasses import dataclass
from typing import Callable, Optional, Tuple, Union
import random
import os
import sys
import json
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 BlockMask, flex_attention, create_block_mask, or_masks

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 .chat_utils import generate_with_prefix_cache_block_diff
from .modeling_ministral import Ministral3Model, Ministral3PreTrainedModel, Ministral3Attention, apply_rotary_pos_emb, repeat_kv, _get_llama_4_attn_scale
from .configuration_ministral_dlm import MinistralDLMConfig

try:
    from flash_attn import flash_attn_func
except:
    print("flash attention not found, please install flash attention for better performance.")
__all__ = ["MinistralDiffEncoderModel", "MinistralFlexAttention"]

@dataclass
class MinistralDiffOutputWithPast(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(dynamic=True, mode="reduce-overhead")
# @torch.compile(mode="default")
# @torch.compile(fullgraph=True, mode="reduce-overhead", dynamic=False)
@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)


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 _extract_draft_kv_cache(past_key_values: DynamicCache, clean_len: int, block_length: int):
    """After quadratic decoding, extract only draft tokens (first of each block) from cache."""
    for layer_idx in range(len(past_key_values)):
        if hasattr(past_key_values, 'layers'):
            layer_cache = past_key_values.layers[layer_idx]
            k, v = layer_cache.keys, layer_cache.values
        else:
            k = past_key_values.key_cache[layer_idx]
            v = past_key_values.value_cache[layer_idx]

        clean_k, draft_k = k[:, :, :clean_len], k[:, :, clean_len::block_length + 1]
        clean_v, draft_v = v[:, :, :clean_len], v[:, :, clean_len::block_length + 1]
        new_k = torch.cat([clean_k, draft_k], dim=2)
        new_v = torch.cat([clean_v, draft_v], dim=2)

        if hasattr(past_key_values, 'layers'):
            layer_cache.keys = new_k
            layer_cache.values = new_v
        else:
            past_key_values.key_cache[layer_idx] = new_k
            past_key_values.value_cache[layer_idx] = new_v

    past_key_values._seen_tokens = clean_len + block_length

# with reference to https://github.com/pytorch-labs/attention-gym/blob/main/examples/flex_attn.ipynb
class MinistralFlexAttention(Ministral3Attention):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.max_seq_length = getattr(self.config, 'max_seq_length', 4096)
        self.block_size_orig = self.config.block_size
        self.bidirectional_mask = None
        if self.config.dlm_paradigm == 'bidirectional':
            self.bidirectional_mask = self.compute_block_mask(mode='bidirectional')
        elif self.config.dlm_paradigm == 'autoregressive':
            self.autoregressive_mask = self.compute_block_mask(mode='autoregressive')
        elif self.config.dlm_paradigm == 'block_diff':
            self.block_diff_mask = None
        elif self.config.dlm_paradigm == 'sbd_block_diff':
            self.sbd_block_diff_mask = None
        else:
            raise ValueError(f"Unknown attention mode: {self.config.dlm_paradigm}")

        self.block_size = self.block_size_orig
        self.mode = self.config.dlm_paradigm
        self._quadratic_block_mask = {}

        import torch._dynamo.config as dcfg
        dcfg.cache_size_limit = 512


    def _get_sbd_inference_quadratic_decoding_block_mask(self, block_length: int):
        if block_length not in self._quadratic_block_mask:
            draft_len = block_length * (block_length + 1)

            def quadratic(b, h, q_idx, kv_idx):
                first_clean = torch.logical_and(
                    kv_idx % (block_length + 1) == 0,
                    kv_idx < draft_len,
                )
                first_clean = torch.logical_and(first_clean, q_idx >= kv_idx)
                block_q = q_idx // (block_length + 1)
                block_kv = kv_idx // (block_length + 1)
                same_block = torch.logical_and(block_q == block_kv, q_idx < draft_len)
                same_block_except_first = torch.logical_and(
                    same_block,
                    q_idx % (block_length + 1) != 0,
                )
                draft_part = torch.logical_or(first_clean, same_block_except_first)
                clean_part = kv_idx >= draft_len
                return torch.logical_or(draft_part, clean_part)

            block_mask = create_block_mask(
                quadratic,
                B=None,
                H=None,
                Q_LEN=draft_len,
                KV_LEN=draft_len + self.config.max_position_embeddings,
                device="cuda",
            )

            self._quadratic_block_mask[block_length] = block_mask

        return self._quadratic_block_mask[block_length]


    def set_attention_mode(self, mode, block_size=None):
        self.mode = mode
        self.block_size = block_size

    def compute_block_mask(self, mode, q_len=None, block_size=None):

        def bidirectional_mask(b, h, q, kv): 
            return (q >= kv) | (q < kv)
        
        def autoregressive_mask(b, h, q, kv):
            return (q >= kv)

        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_q == x0_flag_kv)

            # **2. Offset Block-Causal Mask (M_OBC) **
            offset_block_causal = (
                (block_q > block_kv)
                & (x0_flag_kv == 1)
                & (x0_flag_q == 0)
            )

            # **3. Block-Causal Mask (M_BC) **
            block_causal = (block_q >= block_kv) & (x0_flag_kv == 1) & (x0_flag_q == 1)

            # **4. Combine Masks **
            return block_diagonal | offset_block_causal | block_causal
        

        def sbd_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
        
        def modality_indices_based_mask(block_size, b, h, q_idx, kv_idx, image_doc_id):
            return  (image_doc_id[b, q_idx] > 0) & (image_doc_id[b, q_idx] == image_doc_id[b, kv_idx])

        if mode == 'bidirectional':
            attn_mask = bidirectional_mask
        elif mode == 'autoregressive':
            attn_mask = autoregressive_mask
        elif mode == 'block_diff':
            assert block_size is not None
            attn_mask = lambda b, h, q, kv: block_diff_mask(block_size, b, h, q, kv, self.max_seq_length)
        elif mode == 'sbd_block_diff':
            assert block_size is not None
            attn_mask = lambda b, h, q, kv: sbd_block_diff_mask(block_size, b, h, q, kv, self.max_seq_length)
        else:
            raise ValueError(f"Unknown attention mode: {mode}")

        if q_len is not None:
            Q_LEN = q_len
        else:
            if mode in ['block_diff', 'sbd_block_diff']:
                Q_LEN = self.max_seq_length * 2
            else:
                Q_LEN = self.max_seq_length

        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,
        overwrite_block_mask = None,
        overwrite_attn_impl = None,
        use_cache: Optional[bool] = False,
        **kwargs: Unpack[FlashAttentionKwargs],
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
        if overwrite_attn_impl == 'base':
            return super().forward(
                hidden_states=hidden_states,
                position_embeddings=position_embeddings,
                attention_mask=attention_mask,
                past_key_values=past_key_values,
                cache_position=cache_position,
                is_training=is_training,
                use_cache=use_cache,
                **kwargs,
            )
        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 self.mode in ['block_diff', 'sbd_block_diff'] and 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
            if use_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)
            else:  ## if use_cache == False, do not update cache
                old_k, old_v = past_key_values.layers[self.layer_idx].keys, past_key_values.layers[self.layer_idx].values
                key_states   = torch.cat([old_k, key_states], dim=-2)
                value_states = torch.cat([old_v, value_states], dim=-2)


        self_spec_inference_mode = getattr(self.config, "self_spec_inference_mode", None)
        if self_spec_inference_mode is not None:
            if self_spec_inference_mode == "quadratic":
                block_length = getattr(self.config, "block_length", None) or getattr(self.config, "block_size", None)
                if block_length is None:
                    raise ValueError("SBD quadratic decoding requires block_length in config.")
                if past_key_values is not None:
                    seq_len = key_states.shape[2]
                    draft_len = block_length * (block_length + 1)

                    clean_keys = key_states[:, :, :-draft_len]
                    draft_keys = key_states[:, :, -draft_len:]
                    clean_values = value_states[:, :, :-draft_len]
                    draft_values = value_states[:, :, -draft_len:]
                    key_states = torch.cat([draft_keys, clean_keys], dim=2)
                    value_states = torch.cat([draft_values, clean_values], dim=2)

                    block_mask: BlockMask = self._get_sbd_inference_quadratic_decoding_block_mask(
                        block_length=block_length
                    )
                    block_mask.seq_lengths = (draft_len, seq_len)
                else:
                    seq_len = query_states.shape[2]
                    draft_len = block_length * (block_length + 1)
                    clean_len = seq_len - draft_len

                    def _causal_mask(b, h, q_idx, kv_idx):
                        return torch.logical_and(q_idx >= kv_idx, q_idx < clean_len)

                    def _draft2clean_mask(b, h, q_idx, kv_idx):
                        full_clean = torch.logical_and(q_idx >= clean_len, kv_idx <= clean_len)
                        first_clean = torch.logical_and(
                            q_idx >= clean_len, (kv_idx - clean_len) % (block_length + 1) == 0
                        )
                        first_clean = torch.logical_and(first_clean, q_idx >= kv_idx)
                        return torch.logical_or(full_clean, first_clean)

                    def _draft_mask(b, h, q_idx, kv_idx):
                        block_q = (q_idx - clean_len) // (block_length + 1)
                        block_kv = (kv_idx - clean_len) // (block_length + 1)
                        quadrant = torch.logical_and(q_idx >= clean_len, kv_idx >= clean_len)
                        same_block = torch.logical_and(block_q == block_kv, quadrant)
                        same_block_except_first = torch.logical_and(
                            same_block,
                            (q_idx - clean_len) % (block_length + 1) != 0,
                        )
                        return torch.logical_and(block_q == block_kv, same_block_except_first)

                    mask = or_masks(_causal_mask, _draft2clean_mask)
                    mask = or_masks(mask, _draft_mask)

                    block_mask = create_block_mask(
                        mask, B=None, H=None, Q_LEN=seq_len, KV_LEN=seq_len,
                    )

                key_states = repeat_kv(key_states, self.num_key_value_groups)
                value_states = repeat_kv(value_states, self.num_key_value_groups)
                attn_output = 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

            elif self_spec_inference_mode == "default":
                block_length = getattr(self.config, "block_length", None) or getattr(self.config, "block_size", None)
                if block_length is None:
                    raise ValueError("SBD default decoding requires block_length in config.")
                seq_len = query_states.shape[2]
                prefix_len = seq_len - block_length

                def _clean_q_mask(b, h, q_idx, kv_idx):
                    return torch.logical_and(q_idx >= kv_idx, q_idx < prefix_len)

                def _noisy_q_mask(b, h, q_idx, kv_idx):
                    return q_idx >= prefix_len

                block_mask = create_block_mask(
                    or_masks(_clean_q_mask, _noisy_q_mask),
                    B=None,
                    H=None,
                    Q_LEN=seq_len,
                    KV_LEN=seq_len,
                )

                key_states = repeat_kv(key_states, self.num_key_value_groups)
                value_states = repeat_kv(value_states, self.num_key_value_groups)
                attn_output = 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
        
        else:
            key_states = repeat_kv(key_states, self.num_key_value_groups)
            value_states = repeat_kv(value_states, self.num_key_value_groups)

            if overwrite_block_mask is not None:
                block_mask = overwrite_block_mask
                if block_mask == 'full':
                    block_mask = None
            else:
                if self.mode == 'bidirectional':
                    block_mask = None
                    overwrite_attn_impl = 'flash_attn'
                    # if self.bidirectional_mask is None or q_len != self.bidirectional_mask.shape[-2]:
                    #     block_mask = self.compute_block_mask(mode='bidirectional', q_len=q_len)
                    # else:
                    #     block_mask = self.bidirectional_mask

                elif self.mode == 'autoregressive':
                    if self.autoregressive_mask is None or q_len != self.autoregressive_mask.shape[-2]:
                        block_mask = self.compute_block_mask(mode='autoregressive', q_len=q_len)
                    else:
                        block_mask = self.autoregressive_mask

                elif self.mode == 'block_diff':
                    if self.block_diff_mask is None or self.block_size != self.block_size_orig 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
                elif self.mode == 'sbd_block_diff':
                    if self.sbd_block_diff_mask is None or self.block_size != self.block_size_orig or q_len != self.sbd_block_diff_mask.shape[-2]:
                        block_mask = self.compute_block_mask(mode='sbd_block_diff', block_size=self.block_size, q_len=q_len)
                    else:
                        block_mask = self.sbd_block_diff_mask
                else:
                    raise ValueError(f"Unknown attention mode: {self.mode}")
            if overwrite_attn_impl == 'flash_attn':
                
    
                # FlashAttention expects (batch, seqlen, nheads, headdim)
                # Ensure your tensors are in this layout or permute them here
                #print(query_states.shape,key_states.shape,value_states.shape)
                if self.diffusion_lm:
                    causal = False
                else:
                    causal = True
                attn_output = flash_attn_func(
                    query_states.transpose(1,2), 
                    key_states.transpose(1,2), 
                    value_states.transpose(1,2), 
                    dropout_p=0.0,      # Set your dropout probability
                    softmax_scale=None, # Defaults to 1/sqrt(head_dim)
                    causal=causal         # Set to True if using a causal block_mask logic
                ).transpose(1,2)
                
            else:
                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


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


class MinistralDiffEncoderModel(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
    """

    # Shared/tied tensors that can appear dynamically based on config.
    # Registering these patterns lets save_pretrained() deduplicate safely.
    # _dynamic_tied_weights_keys = [
    #     r"encoder\.embed_tokens\.weight",
    #     r"diffusion_head\.weight",
    #     r"encoder\.vision_tower(?:\.vision_tower)?\.visual_bridge_model\.quantizer\.quantize\.codebooks\.\d+\.(?:embed|embed_ema|cluster_size_ema)",
    # ]

    def __init__(self, config: MinistralDLMConfig):
        super().__init__(config)

        self.mask_token_id = config.mask_token_id

        diffusion_config = copy.deepcopy(config)
        diffusion_config.diffusion_lm = True

        use_flex = getattr(config, 'enable_self_spec', False)

        if config.dlm_paradigm in ['block_diff', 'sbd_block_diff']:
            diffusion_config.attn_class = MinistralFlexAttention
        elif config.dlm_paradigm in ['bidirectional', 'autoregressive']:
            diffusion_config.attn_class = MinistralFlexAttention if use_flex else 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.current_iter_ratio = None

        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
            
        if self.config.adaptive_mask_rate:
            assert block_size is not None

            # --- simple linear window mapping ---
            bs_min = getattr(self.config, "t_bs_min", 16)
            bs_max = getattr(self.config, "t_bs_max", 128)
            w = getattr(self.config, "t_window_width", 0.6)  # fixed width

            # fraction in [0,1] (unclamped first)
            frac = (float(block_size) - float(bs_min)) / max(1.0, float(bs_max - bs_min))
            # upper bound decreases linearly from 1.0 -> 0.5
            u_max = 1.0 - w * frac
            # clamp to [0.6, 1.0] to handle bs outside [bs_min, bs_max]
            u_max = max(0.6, min(1.0, u_max))
            u_min = u_max - w  # ensures width = w

            # sample t ~ Uniform(u_min, u_max)
            t = u_min + (u_max - u_min) * torch.rand(b, device=device, generator=generator)
        else:
            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_process_exp(
        self,
        input_ids: torch.Tensor,
        eps: float = 1e-3,
        block_size: int | None = None,
        half_life_ratio: float = 0.25, # λ = ln 2 / (half_life_ratio·L)
        loss_mask: Optional[torch.Tensor] = None,
    ):
        """
        Two-stage corruption with optional per-block sampling.
        • Stage 1:  m ~ U(eps, 1)   →   k = round(m · len)  (exact budget).
        • Stage 2:  sample exactly k positions with weights
                    w_i(m) = exp[ λ · (1−m) · i ]   (late-heavy when m→0,
                                                     uniform when m→1).
          If `block_size` is given, the procedure is run *independently*
          inside each contiguous block of that length (last block may be shorter).
          When block_size is provided, m is sampled per-block and p_mask is per-block.
        Args
        ----
        input_ids : (B, L)  LongTensor
        eps       : minimum corruption ratio
        block_size: if not None, operate block-wise with per-block m sampling
        half_life_ratio : controls steepness when m→0
        """
        B, L = input_ids.shape
        device = input_ids.device
        dtype  = torch.float32

        masked_indices = torch.zeros((B, L), dtype=torch.bool, device=device)
        p_mask = torch.zeros((B, L), dtype=dtype, device=device)

        # ---------- Stage 1 & 2: whole-sentence or block-wise -------------------
        for b in range(B):
            if block_size is None:
                # ---------- Per-batch sampling (original behavior) ----------
                m = eps + (1.0 - eps) * torch.rand(1, device=device).item()   # scalar
                k_tot = int(round(m * L))
                k_tot = max(1, min(k_tot, L))  # clamp to [1, L]
                
                # Fill p_mask for this batch
                p_mask[b, :] = m
                
                slope = 1.0 - m          # ∈ [0,1]; 0 ⇒ uniform, 1 ⇒ late-heavy
                
                # ------- single pool over the whole sentence -------------
                lam_base = math.log(2.0) / (half_life_ratio * L) # base decay rate (λ when slope=1)

                pos   = torch.arange(L, device=device, dtype=dtype)
                log_w = (lam_base * slope * pos).clone()

                masked_indices[b] = gumbel_topk(log_w, k_tot)

            else:
                # ---------- Per-block sampling ----------
                num_blocks = math.ceil(L / block_size)
                lam_base = math.log(2.0) / (half_life_ratio * block_size) # base decay rate (λ when slope=1)

                for blk in range(num_blocks):
                    start = blk * block_size
                    end   = min((blk + 1) * block_size, L)
                    blk_len = end - start

                    # Sample m per block
                    m_blk = eps + (1.0 - eps) * torch.rand(1, device=device).item()
                    
                    # Fill p_mask for this block
                    p_mask[b, start:end] = m_blk
                    
                    # per-block budget
                    k_blk = int(round(m_blk * blk_len))
                    k_blk = max(0, min(k_blk, blk_len))
                    if k_blk == 0:
                        continue

                    slope = 1.0 - m_blk          # ∈ [0,1]; 0 ⇒ uniform, 1 ⇒ late-heavy

                    pos   = torch.arange(blk_len, device=device, dtype=dtype)
                    log_w = lam_base * slope * pos

                    blk_mask = gumbel_topk(log_w, k_blk)
                    masked_indices[b, start:end] = blk_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 = None,
        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,
        block_diff_ppl: bool                     = False,
        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,
        inputs_embeds: torch.Tensor = None,
        **kwargs,
    ) -> CausalLMOutputWithPast:

        if input_ids is None:
            if inputs_embeds is None:
                raise ValueError("Either `input_ids` or `inputs_embeds` must be provided.")
            batch_size, seq_len = inputs_embeds.shape[:2]
            if labels is not None:
                raise ValueError("`labels` training path requires `input_ids`.")
        else:
            batch_size, seq_len = input_ids.shape


        if self.config.dlm_paradigm == 'bidirectional' or self.config.dlm_paradigm == 'autoregressive':
            if labels is not None and torch.rand(1) < self.config.random_length_prob:
                raise NotImplementedError("Random length training not yet implemented for bidirectional/autoregressive paradigms.")
                random_length = torch.randint(2, input_ids.shape[1] + 1, (1,))
                input_ids = input_ids[:, :random_length]
                labels = labels[:, :random_length]
                
                if attention_mask is not None:
                    attention_mask = attention_mask[:, :random_length]
                if position_ids is not None:
                    position_ids = position_ids[:, :random_length]
                if loss_mask is not None:
                    loss_mask = loss_mask[:, :random_length]

        elif self.config.dlm_paradigm in ['block_diff', 'sbd_block_diff']:
            if labels is not None and block_size is None:
                if torch.rand(1) < self.config.random_length_prob:
                    block_size = torch.randint(1, 8, (1,)).item() * 4  ## [4, 32] divisible by 4
                else:
                    block_size = self.config.block_size

        else:
            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:
                if self.config.tok_mask_half_life_ratio is not None:
                    noisy_inputs, masked_indices, p_mask = self.forward_process_exp(input_ids, eps=eps, block_size=block_size, half_life_ratio=self.config.tok_mask_half_life_ratio, loss_mask=loss_mask)
                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

        if self.config.dlm_paradigm in ['block_diff', 'sbd_block_diff']:
            for layer in self.encoder.layers:
                if hasattr(layer.self_attn, 'set_attention_mode'):
                    layer.self_attn.set_attention_mode(self.config.dlm_paradigm, block_size=block_size)

        input_ids_len = noisy_inputs.shape[1] if noisy_inputs is not None else seq_len
        if labels is not None and self.config.dlm_paradigm in ['block_diff', 'sbd_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)

        if block_diff_ppl:
            if position_ids is None:
                position_ids = torch.arange(input_ids_len // 2, device=noisy_inputs.device).unsqueeze(0)

        enc_out  = self.encoder(
            past_key_values=past_key_values,
            input_ids=noisy_inputs,
            inputs_embeds=inputs_embeds if noisy_inputs is None else None,
            attention_mask=attention_mask,
            position_ids=position_ids,
            is_training=(labels is not None) or (block_diff_ppl),
            **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 in ['block_diff', 'sbd_block_diff']:
            if self.config.dlm_paradigm == 'sbd_block_diff':
                causal_logits = logits[:, input_ids_len:]
            else:
                causal_logits = None

            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:
                # Handle DREAM vs LLADA style losses
                if hasattr(self.config, 'dlm_type') and self.config.dlm_type == 'dream':
                    logits = logits[..., :-1, :].contiguous()
                    labels = labels[..., 1:].contiguous()
                    masked_indices = masked_indices[:, 1:]
                    p_mask = p_mask[:, 1:]

                if self.config.ada_perm_ratio_per_block is not None:
                    # Only compute loss for the top ada_perm_ratio_per_block tokens by confidence within each block
                    block_size = self.config.block_size
                    batch_size, seq_len = masked_indices.shape
                    num_blocks = seq_len // block_size
                    
                    # Get the max logit (confidence) for each position
                    confidence = logits.max(dim=-1).values.detach()  # (batch_size, seq_len)
                    
                    # Create a mask for tokens to include in loss
                    selected_mask = torch.zeros_like(masked_indices, dtype=torch.bool)
                    
                    for blk in range(num_blocks):
                        start = blk * block_size
                        end = min((blk + 1) * block_size, seq_len)
                        
                        # Get masked indices within this block
                        block_masked = masked_indices[:, start:end]  # (batch_size, block_len)
                        block_confidence = confidence[:, start:end]  # (batch_size, block_len)
                        
                        for b in range(batch_size):
                            # Get positions that are masked in this block for this batch
                            masked_positions = torch.where(block_masked[b])[0]
                            num_masked = len(masked_positions)
                            
                            if num_masked > 0:
                                # Number of tokens to keep (top by confidence)
                                k = min(max(1, int(block_size * self.config.ada_perm_ratio_per_block)), num_masked)
                                
                                # Get confidence values for masked positions
                                masked_confidence = block_confidence[b, masked_positions]
                                
                                # Get indices of top-k confident tokens
                                _, topk_indices = torch.topk(masked_confidence, k)
                                selected_positions = masked_positions[topk_indices]
                                
                                # Mark these positions in the selected mask
                                selected_mask[b, start + selected_positions] = True
                    
                    # Calculate loss only for selected positions
                    token_loss = torch.nn.functional.cross_entropy(
                        logits[selected_mask],
                        labels[selected_mask],
                        reduction='none'
                    ) / p_mask[selected_mask]

                    num_mask_tokens = selected_mask.sum()

                else:
                    # Calculate token-wise cross entropy loss for masked positions in B
                    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()

                if self.config.global_loss_avg:
                    loss = token_loss.sum()
                else:
                    loss = token_loss.sum() / num_mask_tokens
                
                if self.config.ada_dlm_loss_ratio is not None:
                    assert self.current_iter_ratio is not None
                    assert self.config.dlm_loss_weight is not None

                    dlm_loss_weight = min(self.config.dlm_loss_weight, self.current_iter_ratio / self.config.ada_dlm_loss_ratio * self.config.dlm_loss_weight)
                    loss = dlm_loss_weight * loss

                elif self.config.dlm_loss_weight is not None:
                    loss = self.config.dlm_loss_weight * loss

                if self.config.dlm_paradigm == 'sbd_block_diff':
                    causal_logits = causal_logits[..., :-1, :].contiguous()
                    causal_logits = causal_logits.view(-1, causal_logits.size(-1))

                    if hasattr(self.config, 'dlm_type') and self.config.dlm_type == 'dream':
                        causal_labels = labels.view(-1)
                    else: 
                        causal_labels = labels[..., 1:].contiguous().view(-1)
                    
                    if self.config.global_loss_avg:
                        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
                    else:
                        loss_fct = CrossEntropyLoss()
                        ar_loss = loss_fct(causal_logits, causal_labels)

                        self.loss_diffusion = loss.detach().item()
                        self.loss_ar = ar_loss.detach().item()

                        loss = loss + self.config.ar_loss_weight * ar_loss
                
                if self.config.global_loss_avg:
                    if self.config.dlm_paradigm == 'sbd_block_diff':
                        loss = (loss, num_mask_tokens + int(self.config.ar_loss_weight * seq_len))
                    else: 
                        loss = (loss, num_mask_tokens)

        return MinistralDiffOutputWithPast(
            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,
        )


    def generate_diffusion(self, prompt_ids, max_new_tokens=512, steps=512, block_length=32, shift_logits=False, threshold=0.9, causal_context=True, temperature=0, eos_token_id=None, max_thinking_tokens=None, end_think_token_id=None, step_ratio=None,prompt_embeds=None,**kwargs):
        if prompt_embeds is None and prompt_ids is not None and torch.is_floating_point(prompt_ids):
            prompt_embeds = prompt_ids
            prompt_ids = None

        if (prompt_ids is None) == (prompt_embeds is None):
            raise ValueError("Exactly one of `prompt_ids` or `prompt_embeds` must be provided.")

        if eos_token_id is None:
            eos_token_id = getattr(self.config, 'eos_token_id', None)
        if step_ratio is not None:
            steps_per_block = int(block_length * step_ratio)
            num_blocks = max_new_tokens // block_length
            steps = steps_per_block * num_blocks
        out_ids, nfe = generate_with_prefix_cache_block_diff(
                        model=self,
                        prompt=prompt_ids,
                        prompt_embeds=prompt_embeds,
                        gen_length=max_new_tokens,
                        steps=steps,
                        block_length=block_length,
                        remasking="low_confidence",
                        temperature=temperature,
                        mask_id=self.mask_token_id,
                        threshold=threshold,
                        shift_logits=shift_logits,
                        neg_entropy=False,
                        causal_context=causal_context,
                        eos_token_id=eos_token_id,
                        max_thinking_tokens=max_thinking_tokens,
                        end_think_token_id=end_think_token_id,
                    )

        return out_ids, nfe


    @torch.no_grad()
    def sbd_inference_diffusion_quadratic(
        self,
        clean_input_ids: Optional[torch.Tensor],
        draft_input_ids: torch.Tensor,
        block_length: int,
        draft_only: bool = False,
        past_key_values: Optional[Cache] = None,
        use_cache: bool = False,
    ):
        enc_config = self.encoder.config
        enc_config.use_sbd_objective = True
        enc_config.block_length = block_length

        if draft_only:
            assert clean_input_ids is not None

            if use_cache and past_key_values is None:
                past_key_values = DynamicCache()

            enc_config.self_spec_inference_mode = "default"
            input_ids = torch.cat([clean_input_ids, draft_input_ids], dim=-1)
            outputs = self.encoder(
                input_ids=input_ids,
                position_ids=None,
                past_key_values=past_key_values,
                use_cache=use_cache,
                is_training=False,
            )

            hidden_states = outputs.last_hidden_state
            logits = self.diffusion_head(hidden_states)

            past_key_values = getattr(outputs, "past_key_values", None)
            if use_cache and past_key_values is not None:
                _crop_dynamic_cache(past_key_values, clean_input_ids.shape[1])

            return logits, past_key_values
        else:
            enc_config.self_spec_inference_mode = "quadratic"

            draft_len = block_length * (block_length + 1)
            draft_input_ids = torch.cat(
                [
                    draft_input_ids.view(-1, block_length, 1),
                    torch.full(
                        (draft_input_ids.shape[0], block_length, block_length),
                        fill_value=self.config.mask_token_id,
                        device=draft_input_ids.device,
                    ),
                ],
                dim=-1,
            ).view(-1, draft_len)

            if use_cache:
                assert past_key_values is not None, (
                    "Past key values should be provided when using cache, e.g. run draft_only=True first."
                )
                assert clean_input_ids is None, (
                    "Clean input ids should already be in cache, thus none should be provided."
                )
                clean_len = past_key_values.get_seq_length()
                input_ids = draft_input_ids
            else:
                clean_len = clean_input_ids.shape[1]
                input_ids = torch.cat([clean_input_ids, draft_input_ids], dim=-1)

            per_block_position_ids = torch.arange(
                clean_len, clean_len + block_length + 1, device=draft_input_ids.device
            )[None,].repeat(block_length, 1)
            per_block_position_ids += torch.arange(block_length, device=draft_input_ids.device).view(-1, 1)

            if use_cache:
                position_ids = per_block_position_ids.view(-1)[None,]
            else:
                clean_position_ids = torch.arange(clean_len, device=draft_input_ids.device)
                position_ids = torch.cat([clean_position_ids, per_block_position_ids.view(-1)], dim=-1)[None,]

            outputs = self.encoder(
                input_ids=input_ids,
                position_ids=position_ids,
                past_key_values=past_key_values,
                use_cache=use_cache,
                is_training=False,
            )

            hidden_states = outputs.last_hidden_state
            logits = self.diffusion_head(hidden_states)
            past_key_values = getattr(outputs, "past_key_values", None)

            if use_cache and past_key_values is not None:
                _extract_draft_kv_cache(past_key_values, clean_len, block_length)

            return logits, past_key_values


    @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 MinistralDiffEncoderModel.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 self_spec_generate(
        self,
        prompt_ids: torch.Tensor,
        max_new_tokens: int = 128,
        steps: int = 128,
        block_length: int = 16,
        ar_mix_weight: Optional[float] = None,
        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,
    ):
        self.config.use_sbd_objective = True
        self.config.dlm_paradigm = "sbd"

        if prompt_ids.shape[0] != 1:
            raise ValueError("Self speculation quadratic decoding currently 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)

        x = torch.full(
            (1, prompt_ids.shape[1] + max_new_tokens + block_length * 2),
            token_mask_id,
            dtype=torch.long,
            device=prompt_ids.device,
        )
        x[:, : prompt_ids.shape[1]] = prompt_ids.clone()

        if max_new_tokens % block_length != 0:
            raise ValueError("max_new_tokens must be divisible by block_length")
        num_blocks = max_new_tokens // block_length
        if steps % num_blocks != 0:
            raise ValueError("steps must be divisible by (max_new_tokens // block_length)")

        prompt_len = prompt_ids.shape[1]
        nfe = 0
        nfe += 1
        logits, past_key_values = self.sbd_inference_diffusion_quadratic(
            clean_input_ids=x[:, :prompt_len],
            draft_input_ids=x[:, prompt_len : prompt_len + block_length],
            block_length=block_length,
            draft_only=True,
            use_cache=True,
        )

        logits_proposal = logits[:, prompt_len - 1 : prompt_len + block_length]
        logits_proposal[:, 1] = logits_proposal[:, 0]
        logits_proposal = logits_proposal[:, 1:]
        x0_proposal = torch.argmax(logits_proposal, dim=-1)
        x[:, prompt_len : prompt_len + block_length] = x0_proposal

        total_accept_token = 0
        while True:
            nfe += 1
            block_start = prompt_len + total_accept_token
            block_end = block_start + block_length
            draft_input_ids = x[:, block_start:block_end]

            logits, past_key_values = self.sbd_inference_diffusion_quadratic(
                clean_input_ids=None,
                draft_input_ids=draft_input_ids,
                block_length=block_length,
                draft_only=False,
                past_key_values=past_key_values,
                use_cache=True,
            )

            useful_token_logits = logits.view(1, block_length, block_length + 1, -1)
            if ar_mix_weight is None:
                useful_token_logits[:, :, 1] = useful_token_logits[:, :, 0]
            else:
                if not (0.0 <= ar_mix_weight <= 1.0):
                    raise ValueError("ar_mix_weight must be between 0 and 1")
                mix_logits = useful_token_logits[:, :, 0] * ar_mix_weight + useful_token_logits[:, :, 1] * (1 - ar_mix_weight)
                useful_token_logits[:, :, 0] = mix_logits
                useful_token_logits[:, :, 1] = mix_logits

            if temperature > 0:
                useful_token_logits = useful_token_logits / temperature

            useful_token_pred = torch.argmax(useful_token_logits, dim=-1)
            new_draft_input_ids = useful_token_pred[:, 0, 1:]
            accept_cnt = 1

            while accept_cnt < block_length:
                if useful_token_pred[:, accept_cnt - 1, 0].item() != draft_input_ids[:, accept_cnt].item():
                    break
                new_draft_input_ids = useful_token_pred[:, accept_cnt, 1:]
                accept_cnt += 1

            x[:, block_start : block_start + accept_cnt] = draft_input_ids[:, :accept_cnt]

            # EoS early stopping: all accepted tokens are finalized left-to-right,
            # so if any is EoS we can truncate and return immediately.
            if eos_token_id is not None:
                accepted = x[0, block_start : block_start + accept_cnt]
                eos_positions = (accepted == eos_token_id).nonzero(as_tuple=True)[0]
                if len(eos_positions) > 0:
                    first_eos_rel = eos_positions[0].item()
                    total_accept_token += first_eos_rel + 1
                    output_end = prompt_len + total_accept_token
                    return x[:, :output_end], nfe

            x[:, block_start + accept_cnt : block_start + accept_cnt + block_length] = new_draft_input_ids
            past_key_values.crop(block_start + accept_cnt)

            # ---- thinking budget enforcement ----
            # Insert end_think as the first token of the next draft block,
            # shifting all subsequent tokens right by 1 (discarding the last).
            # The first draft token is always accepted unconditionally, so
            # end_think is guaranteed to be finalized in the next iteration
            # without needing to re-encode or touch the KV cache.
            if end_think_token_id is not None and max_thinking_tokens is not None:
                tokens_so_far = total_accept_token + accept_cnt
                if tokens_so_far > max_thinking_tokens:
                    gen_so_far = x[0, prompt_len : prompt_len + tokens_so_far]
                    has_end_think = (gen_so_far == end_think_token_id).any()
                    if not has_end_think:
                        insert_pos = block_start + accept_cnt
                        x[0, insert_pos + 1:] = x[0, insert_pos:-1].clone()
                        x[0, insert_pos] = end_think_token_id

            total_accept_token += accept_cnt

            if total_accept_token >= max_new_tokens:
                break

        return x[:, : -(block_length * 2)], 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 verification.

        Each step:
          1. Draft: forward [last_accepted, mask, ...] with bidirectional attention
             (diffusion_lm=True, use_cache=False).  Shift AR logits to get
             per-position predictions; apply confidence filtering.
          2. Verify: forward the drafted block with causal attention
             (diffusion_lm=False, use_cache=True, use_causal_mask=True).
             Accept consecutive AR-matching tokens plus one bonus token.

        Args:
            prompt_ids: Input token IDs of shape (1, prompt_len).
            max_new_tokens: Maximum number of tokens to generate.
            block_length: Number of tokens per draft/verify block.
            temperature: Sampling temperature (0 = greedy).
            mask_token_id: Override for config.mask_token_id.
            eos_token_id: Override for config.eos_token_id.
            max_thinking_tokens: Budget for thinking tokens before forcing end_think.
            end_think_token_id: Token ID inserted when thinking budget is exceeded.
            threshold: Confidence threshold for accepting draft predictions.

        Returns:
            (output_ids, nfe): output_ids includes the prompt; nfe is the number
            of forward evaluations (matching self_spec_generate interface).
        """
        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
        prompt_len = prompt_ids.shape[1]
        dream_style = getattr(self.config, 'dlm_type', 'llada') == 'dream'

        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

        # ===== Prefill (causal) =====
        _set_diffusion_lm(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:
            probs = torch.softmax(last_logit / temperature, dim=-1)
            next_token = torch.multinomial(probs, 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:
            output_ids = torch.cat([prompt_ids, next_token], dim=1)
            return output_ids, nfe

        generated = [next_token]
        total_gen = 1

        # ===== Main loop =====
        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 (bidirectional, don't update cache) --------
            _set_diffusion_lm(True)
            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)
            if dream_style:
                # DREAM: logit[i] predicts position i+1 → shift to self-prediction
                draft_logits = torch.cat(
                    [draft_logits[:, :1, :], draft_logits[:, :-1, :]], dim=1
                )
            # LLaDA: logit[i] already 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)

            draft_conf = torch.gather(
                draft_probs, -1, draft_tokens.unsqueeze(-1)
            ).squeeze(-1)

            is_mask = block == token_mask_id
            draft_conf = torch.where(is_mask, draft_conf, -torch.inf)
            unmask = draft_conf > threshold

            if unmask.sum() > 0:
                block[unmask] = draft_tokens[unmask]
            else:
                raise AssertionError(
                    "No mask token above threshold for prediction"
                )

            # -------- Verify (causal, update cache) --------
            _set_diffusion_lm(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:
                verify_probs = torch.softmax(verify_logits / temperature, dim=-1)
                ar_tokens = torch.multinomial(
                    verify_probs.view(-1, verify_probs.shape[-1]), num_samples=1
                ).view(1, block_length)
            else:
                ar_tokens = verify_logits.argmax(dim=-1)

            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  # bonus token from AR verification

            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]

            # -------- EOS check --------
            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 --------
            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