modeling_dhara_ar.py

#!/usr/bin/env python3
"""
Dhara-AR: LLaMA3-style autoregressive model with Canon Layer positions (ABCD).

Canon positions from "Physics of Language Models: Part 4.1" by Zeyuan Allen-Zhu:
- A: After input LayerNorm, before attention
- B: Inside attention, after Q/K/V projections
- C: After post-attention LayerNorm, before MLP
- D: Inside MLP, after gate/up projections
"""

import math
from typing import Optional, Tuple, List, Union

import torch
import torch.nn as nn
import torch.nn.functional as F

from transformers import PreTrainedModel, GenerationMixin, DynamicCache
from transformers.modeling_outputs import CausalLMOutputWithPast

from .configuration_dhara_ar import DharaARConfig

# Flash Attention for memory-efficient attention
try:
    from flash_attn import flash_attn_func
    FLASH_ATTN_AVAILABLE = True
except ImportError:
    FLASH_ATTN_AVAILABLE = False
    print("Warning: Flash Attention not available, falling back to standard attention")


def build_block_causal_mask(seq_len, block_len, device, dtype):
    """Additive (1,1,S,S) mask: causal across blocks, bidirectional within a
    block. block_len==1 -> standard causal; block_len>=S -> full bidirectional.
    Used by diffusion mode and the self-speculation draft."""
    idx = torch.arange(seq_len, device=device)
    blk = (idx // block_len).unsqueeze(0)
    blk_row = (idx // block_len).unsqueeze(1)
    allowed = blk <= blk_row
    bias = torch.zeros((seq_len, seq_len), device=device, dtype=dtype)
    bias = bias.masked_fill(~allowed, float("-inf"))
    return bias.unsqueeze(0).unsqueeze(0)


class CanonLayer(nn.Module):
    """
    Canon Layer: Causal 1D depthwise convolution for local context.
    """

    def __init__(
        self,
        hidden_size: int,
        kernel_size: int = 4,
        use_residual: bool = True,
        use_activation: bool = False,
        use_bias: bool = False,
    ):
        super().__init__()
        self.hidden_size = hidden_size
        self.kernel_size = kernel_size
        self.use_residual = use_residual
        self.use_activation = use_activation

        self.conv = nn.Conv1d(
            in_channels=hidden_size,
            out_channels=hidden_size,
            kernel_size=kernel_size,
            padding=kernel_size - 1,
            groups=hidden_size,
            bias=use_bias,
        )

        nn.init.normal_(self.conv.weight, mean=0.0, std=0.02)
        if use_bias:
            nn.init.zeros_(self.conv.bias)

        # Per-layer conv state for incremental decode: caches the last
        # (kernel_size-1) input timesteps so cached generation matches a full
        # forward (otherwise the causal conv would pad with zeros and corrupt
        # every decode step). Reset automatically on each prefill (past_len==0).
        self._conv_state = None

    def forward(self, hidden_states: torch.Tensor, use_cache: bool = False,
                past_len: int = 0) -> torch.Tensor:
        batch_size, seq_len, hidden_size = hidden_states.shape
        pad = self.kernel_size - 1

        x = hidden_states.transpose(1, 2)  # (B, H, S)

        if use_cache and past_len > 0 and self._conv_state is not None:
            # incremental decode: prepend cached real inputs (exact causal context)
            x_in = torch.cat([self._conv_state.to(dtype=x.dtype, device=x.device), x], dim=2)
        else:
            # prefill / full forward: standard causal left-pad with zeros
            x_in = F.pad(x, (pad, 0)) if pad > 0 else x

        # explicit padding=0 conv over the (history + current) window
        out = F.conv1d(x_in, self.conv.weight, self.conv.bias, padding=0, groups=self.hidden_size)
        out = out[:, :, :seq_len]

        if use_cache and pad > 0:
            self._conv_state = x_in[:, :, -pad:].detach()  # last K-1 real inputs
        elif not use_cache:
            self._conv_state = None

        if self.use_activation:
            out = F.silu(out)

        out = out.transpose(1, 2)

        if self.use_residual:
            out = hidden_states + out

        return out


class RMSNorm(nn.Module):
    """Root Mean Square Layer Normalization."""

    def __init__(self, hidden_size: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        variance = x.pow(2).mean(-1, keepdim=True)
        x = x * torch.rsqrt(variance + self.eps)
        return self.weight * x


class RotaryEmbedding(nn.Module):
    """Rotary Position Embedding (RoPE) with optional YaRN scaling.

    YaRN (Yet another RoPE extensioN) splits frequency dimensions into 3 groups:
    - High frequencies (local positions): no interpolation
    - Medium frequencies: partial interpolation via NTK-aware blend
    - Low frequencies (global positions): full interpolation
    This preserves short-range attention while extending long-range reach.

    Usage:
        # Standard RoPE (training)
        rope = RotaryEmbedding(dim=64, theta=8000000.0)

        # YaRN-extended RoPE (inference, 2x extension: 32K -> 64K)
        rope = RotaryEmbedding(dim=64, theta=8000000.0,
                               rope_scaling={"type": "yarn", "factor": 2.0})

        # YaRN 4x extension: 32K -> 128K
        rope = RotaryEmbedding(dim=64, theta=8000000.0,
                               rope_scaling={"type": "yarn", "factor": 4.0})
    """

    def __init__(
        self,
        dim: int,
        max_position_embeddings: int = 8192,
        theta: float = 10000.0,
        rope_scaling: dict = None,
    ):
        super().__init__()
        self.dim = dim
        self.max_position_embeddings = max_position_embeddings
        self.theta = theta
        self.rope_scaling = rope_scaling

        inv_freq = self._compute_inv_freq()
        self.register_buffer("inv_freq", inv_freq, persistent=True)
        self.cos_cached = None
        self.sin_cached = None
        self.max_seq_len_cached = 0

        # YaRN attention magnitude scaling factor
        self._yarn_attn_factor = 1.0
        if rope_scaling and rope_scaling.get("type") == "yarn":
            factor = rope_scaling.get("factor", 1.0)
            # Attention temperature correction: 0.1 * ln(factor) + 1.0
            self._yarn_attn_factor = 0.1 * math.log(factor) + 1.0

    def _compute_inv_freq(self) -> torch.Tensor:
        """Compute inverse frequencies, applying YaRN scaling if configured."""
        base_inv_freq = 1.0 / (self.theta ** (torch.arange(0, self.dim, 2).float() / self.dim))

        if not self.rope_scaling or self.rope_scaling.get("type") != "yarn":
            return base_inv_freq

        factor = self.rope_scaling.get("factor", 1.0)
        if factor <= 1.0:
            return base_inv_freq

        # YaRN parameters
        beta_fast = self.rope_scaling.get("beta_fast", 32)
        beta_slow = self.rope_scaling.get("beta_slow", 1)
        original_max_pos = self.rope_scaling.get("original_max_position_embeddings",
                                                  self.max_position_embeddings)

        # Compute wavelengths for each frequency dimension
        dim_indices = torch.arange(0, self.dim, 2).float()
        wavelengths = 2 * math.pi * self.theta ** (dim_indices / self.dim)

        # Boundaries for the 3 regions
        low_bound = original_max_pos / (beta_fast / (2 * math.pi))
        high_bound = original_max_pos / (beta_slow / (2 * math.pi))

        # Interpolation ramp: 0 = full interpolation, 1 = no interpolation
        ramp = (wavelengths - low_bound) / (high_bound - low_bound)
        ramp = ramp.clamp(0, 1)

        # NTK-aware interpolation: scale theta by factor for interpolated dims
        # High freq (ramp=1): use original inv_freq (no change)
        # Low freq (ramp=0): scale by 1/factor (full interpolation)
        # Medium freq: blend between the two
        inv_freq_interpolated = base_inv_freq / factor
        inv_freq_yarn = inv_freq_interpolated * (1 - ramp) + base_inv_freq * ramp

        return inv_freq_yarn

    def _build_cache(self, seq_len: int, device: torch.device, dtype: torch.dtype):
        self.max_seq_len_cached = seq_len
        t = torch.arange(seq_len, device=device, dtype=torch.float32)
        inv_freq = self.inv_freq.to(device=device, dtype=torch.float32)
        freqs = torch.outer(t, inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        cos = emb.cos().to(dtype)
        sin = emb.sin().to(dtype)
        # Apply YaRN attention scaling factor
        if self._yarn_attn_factor != 1.0:
            cos = cos * self._yarn_attn_factor
            sin = sin * self._yarn_attn_factor
        self.cos_cached = cos
        self.sin_cached = sin

    @torch.compiler.disable
    def _ensure_cache(self, seq_len: int, device: torch.device, dtype: torch.dtype):
        if self.cos_cached is None or seq_len > self.max_seq_len_cached or self.cos_cached.device != device:
            self._build_cache(max(seq_len, self.max_position_embeddings), device, dtype)

    def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        seq_len = position_ids.max().item() + 1
        self._ensure_cache(seq_len, x.device, x.dtype)

        cos = self.cos_cached[position_ids].unsqueeze(1)
        sin = self.sin_cached[position_ids].unsqueeze(1)
        return cos.to(x.dtype), sin.to(x.dtype)


def rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary_pos_emb(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed


class DharaARAttention(nn.Module):
    """Multi-head attention with GQA and Canon-B position."""

    def __init__(self, config: DharaARConfig, layer_idx: int):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        self.head_dim = config.hidden_size // config.num_attention_heads
        self.num_kv_groups = self.num_heads // self.num_kv_heads

        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)
        self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias)
        self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias)
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=config.attention_bias)

        # Canon-B: Inside attention
        self.canon_b_q = None
        self.canon_b_k = None
        self.canon_b_v = None
        if "B" in config.canon_set:
            self.canon_b_q = CanonLayer(
                hidden_size=self.num_heads * self.head_dim,
                kernel_size=config.canon_kernel,
                use_residual=config.canon_residual,
                use_activation=config.canon_activation,
                use_bias=config.canon_bias,
            )
            self.canon_b_k = CanonLayer(
                hidden_size=self.num_kv_heads * self.head_dim,
                kernel_size=config.canon_kernel,
                use_residual=config.canon_residual,
                use_activation=config.canon_activation,
                use_bias=config.canon_bias,
            )
            self.canon_b_v = CanonLayer(
                hidden_size=self.num_kv_heads * self.head_dim,
                kernel_size=config.canon_kernel,
                use_residual=config.canon_residual,
                use_activation=config.canon_activation,
                use_bias=config.canon_bias,
            )

        # QK Normalization
        self.q_norm = None
        self.k_norm = None
        if getattr(config, 'use_qk_norm', False):
            self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
            self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)

        self.rotary_emb = RotaryEmbedding(
            self.head_dim,
            max_position_embeddings=config.max_position_embeddings,
            theta=config.rope_theta,
            rope_scaling=getattr(config, 'rope_scaling', None),
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        past_key_value: Optional[DynamicCache] = None,
        use_cache: bool = False,
        layer_idx: int = 0,
        past_len: int = 0,
        trimode_bias: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, Optional[DynamicCache]]:
        batch_size, seq_len, _ = hidden_states.shape

        query = self.q_proj(hidden_states)
        key = self.k_proj(hidden_states)
        value = self.v_proj(hidden_states)

        # Canon-B: Apply after projections
        if self.canon_b_q is not None:
            query = self.canon_b_q(query, use_cache=use_cache, past_len=past_len)
            key = self.canon_b_k(key, use_cache=use_cache, past_len=past_len)
            value = self.canon_b_v(value, use_cache=use_cache, past_len=past_len)

        query = query.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        key = key.view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
        value = value.view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)

        cos, sin = self.rotary_emb(query, position_ids)
        query, key = apply_rotary_pos_emb(query, key, cos, sin)

        # QK Normalization (after RoPE)
        if self.q_norm is not None:
            query = self.q_norm(query)
            key = self.k_norm(key)

        if past_key_value is not None:
            key, value = past_key_value.update(key, value, layer_idx)

        # Tri-mode: block-diffusion / self-spec draft (additive block-causal bias).
        # Only active when a bias is passed; AR path below is byte-identical without it.
        if trimode_bias is not None:
            if self.num_kv_groups > 1:
                key = key.repeat_interleave(self.num_kv_groups, dim=1)
                value = value.repeat_interleave(self.num_kv_groups, dim=1)
            attn_output = F.scaled_dot_product_attention(
                query, key, value, attn_mask=trimode_bias, dropout_p=0.0, is_causal=False)
            attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, -1)
            return self.o_proj(attn_output), past_key_value

        # Use Flash Attention if available (much more memory efficient for long sequences)
        if FLASH_ATTN_AVAILABLE and not use_cache:
            # Flash attention expects (batch, seqlen, nheads, headdim)
            # Current shape is (batch, nheads, seqlen, headdim), so transpose
            query = query.transpose(1, 2).to(torch.bfloat16)
            key = key.transpose(1, 2).to(torch.bfloat16)
            value = value.transpose(1, 2).to(torch.bfloat16)

            # Flash attention with causal mask, handles GQA natively
            attn_output = flash_attn_func(
                query, key, value,
                causal=True,
                softmax_scale=1.0 / math.sqrt(self.head_dim),
            )
            # Output is (batch, seqlen, nheads, headdim)
            attn_output = attn_output.reshape(batch_size, seq_len, -1)
        else:
            # Fallback to standard attention (for inference with KV cache)
            # Repeat KV for GQA
            if self.num_kv_groups > 1:
                key = key.repeat_interleave(self.num_kv_groups, dim=1)
                value = value.repeat_interleave(self.num_kv_groups, dim=1)

            # Attention
            scale = 1.0 / math.sqrt(self.head_dim)
            attn_weights = torch.matmul(query, key.transpose(-2, -1)) * scale

            if attention_mask is not None:
                attn_weights = attn_weights + attention_mask

            attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
            attn_output = torch.matmul(attn_weights, value)

            attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, -1)
        attn_output = self.o_proj(attn_output)

        return attn_output, past_key_value


class DharaARMLP(nn.Module):
    """SwiGLU MLP with Canon-D position."""

    def __init__(self, config: DharaARConfig):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size

        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.mlp_bias)
        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.mlp_bias)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=config.mlp_bias)

        # Canon-D: Inside MLP
        self.canon_d = None
        if "D" in config.canon_set:
            self.canon_d = CanonLayer(
                hidden_size=config.intermediate_size,
                kernel_size=config.canon_kernel,
                use_residual=config.canon_residual,
                use_activation=config.canon_activation,
                use_bias=config.canon_bias,
            )

    def forward(self, x: torch.Tensor, use_cache: bool = False, past_len: int = 0) -> torch.Tensor:
        gate = F.silu(self.gate_proj(x))
        up = self.up_proj(x)

        if self.canon_d is not None:
            intermediate = gate * up
            intermediate = self.canon_d(intermediate, use_cache=use_cache, past_len=past_len)
            return self.down_proj(intermediate)
        else:
            return self.down_proj(gate * up)


class DharaARDecoderLayer(nn.Module):
    """Decoder layer with all 4 Canon positions."""

    def __init__(self, config: DharaARConfig, layer_idx: int):
        super().__init__()
        self.config = config

        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

        # Canon-A: Before attention
        self.canon_a = None
        if "A" in config.canon_set:
            self.canon_a = CanonLayer(
                hidden_size=config.hidden_size,
                kernel_size=config.canon_kernel,
                use_residual=config.canon_residual,
                use_activation=config.canon_activation,
                use_bias=config.canon_bias,
            )

        self.self_attn = DharaARAttention(config, layer_idx)
        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

        # Canon-C: Before MLP
        self.canon_c = None
        if "C" in config.canon_set:
            self.canon_c = CanonLayer(
                hidden_size=config.hidden_size,
                kernel_size=config.canon_kernel,
                use_residual=config.canon_residual,
                use_activation=config.canon_activation,
                use_bias=config.canon_bias,
            )

        self.mlp = DharaARMLP(config)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        past_key_value: Optional[DynamicCache] = None,
        use_cache: bool = False,
        layer_idx: int = 0,
        past_len: int = 0,
        trimode_bias: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, Optional[DynamicCache]]:
        residual = hidden_states

        hidden_states = self.input_layernorm(hidden_states)

        if self.canon_a is not None:
            hidden_states = self.canon_a(hidden_states, use_cache=use_cache, past_len=past_len)

        hidden_states, present_kv = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_value=past_key_value,
            use_cache=use_cache,
            layer_idx=layer_idx,
            past_len=past_len,
            trimode_bias=trimode_bias,
        )
        hidden_states = residual + hidden_states

        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)

        if self.canon_c is not None:
            hidden_states = self.canon_c(hidden_states, use_cache=use_cache, past_len=past_len)

        hidden_states = self.mlp(hidden_states, use_cache=use_cache, past_len=past_len)
        hidden_states = residual + hidden_states

        return hidden_states, present_kv


class DharaARPreTrainedModel(PreTrainedModel):
    config_class = DharaARConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True

    def _init_weights(self, module):
        std = self.config.initializer_range if hasattr(self.config, 'initializer_range') else 0.02
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=std)


class DharaARModel(DharaARPreTrainedModel):
    """Dhara-AR transformer model."""

    def __init__(self, config: DharaARConfig):
        super().__init__(config)
        self.config = config

        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList([
            DharaARDecoderLayer(config, layer_idx)
            for layer_idx in range(config.num_hidden_layers)
        ])
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

        self.gradient_checkpointing = False
        self.post_init()

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        past_key_values: Optional[List[Tuple[torch.Tensor, ...]]] = None,
        use_cache: bool = False,
        trimode_bias: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, Optional[List[Tuple[torch.Tensor, ...]]]]:
        batch_size, seq_len = input_ids.shape

        hidden_states = self.embed_tokens(input_ids)

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

        past_len = past_key_values.get_seq_length() if past_key_values is not None else 0

        if position_ids is None:
            position_ids = torch.arange(past_len, past_len + seq_len, device=input_ids.device)
            position_ids = position_ids.unsqueeze(0).expand(batch_size, -1)

        # Causal mask
        total_len = past_len + seq_len

        causal_mask = torch.triu(
            torch.full((seq_len, total_len), float('-inf'), device=input_ids.device, dtype=hidden_states.dtype),
            diagonal=past_len + 1
        )
        causal_mask = causal_mask.unsqueeze(0).unsqueeze(0)

        for layer_idx, layer in enumerate(self.layers):
            if self.gradient_checkpointing and self.training:
                hidden_states, _ = torch.utils.checkpoint.checkpoint(
                    layer, hidden_states, causal_mask, position_ids, None, False,
                    use_reentrant=False,
                )
            else:
                hidden_states, _ = layer(
                    hidden_states=hidden_states,
                    attention_mask=causal_mask,
                    position_ids=position_ids,
                    past_key_value=past_key_values if use_cache else None,
                    use_cache=use_cache,
                    layer_idx=layer_idx,
                    past_len=past_len,
                    trimode_bias=trimode_bias,
                )

        hidden_states = self.norm(hidden_states)

        return hidden_states, past_key_values if use_cache else None


class DharaARForCausalLM(DharaARPreTrainedModel, GenerationMixin):
    """Dhara-AR for causal language modeling."""
    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}

    def __init__(self, config: DharaARConfig):
        super().__init__(config)
        self.model = DharaARModel(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        if config.tie_word_embeddings:
            self.lm_head.weight = self.model.embed_tokens.weight

        self.post_init()

    def get_input_embeddings(self):
        return self.model.embed_tokens

    def set_input_embeddings(self, value):
        self.model.embed_tokens = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        past_key_values: Optional[List[Tuple[torch.Tensor, ...]]] = None,
        labels: Optional[torch.Tensor] = None,
        use_cache: bool = False,
        return_dict: bool = True,
        trimode_bias: Optional[torch.Tensor] = None,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        hidden_states, present_key_values = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            use_cache=use_cache,
            trimode_bias=trimode_bias,
        )

        logits = self.lm_head(hidden_states)

        # Logit softcapping
        if getattr(self.config, 'use_logit_softcap', False) and self.config.logit_softcap > 0:
            cap = self.config.logit_softcap
            logits = cap * torch.tanh(logits / cap)

        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, self.config.vocab_size),
                shift_labels.view(-1),
                ignore_index=-100,
            )

        if not return_dict:
            output = (logits, present_key_values)
            return (loss,) + output if loss is not None else output

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=present_key_values,
        )

    def _eos_ids(self, eos_token_id):
        if eos_token_id is None:
            gc = getattr(self, "generation_config", None)
            eos_token_id = (gc.eos_token_id if gc is not None and gc.eos_token_id is not None
                            else self.config.eos_token_id)
        return list(eos_token_id) if isinstance(eos_token_id, (list, tuple)) else [eos_token_id]

    @torch.no_grad()
    def generate_diffusion(self, input_ids, block_len: int = 32, threshold: float = 0.5,
                           max_new_tokens: int = 128, eos_token_id=None):
        """Mode 2: block-diffusion decode. Appends a masked block and iteratively
        unmasks high-confidence (>=threshold) positions; repeats block by block.
        Returns the full sequence (prompt + generated)."""
        device = input_ids.device
        dtype = next(self.parameters()).dtype
        mask_id = int(self.config.mask_token_id)
        eos_ids = self._eos_ids(eos_token_id)
        cur = input_ids; gen = 0
        while gen < max_new_tokens:
            blk = torch.full((1, block_len), mask_id, device=device, dtype=cur.dtype)
            seq = torch.cat([cur, blk], 1); S = seq.shape[1]
            bias = build_block_causal_mask(S, block_len, device, dtype)
            for _ in range(block_len):
                mpos = (seq[0] == mask_id).nonzero(as_tuple=True)[0]
                if mpos.numel() == 0:
                    break
                logits = self(input_ids=seq, trimode_bias=bias).logits[0].float()
                conf, pred = F.softmax(logits[mpos], -1).max(-1)
                take = conf >= threshold
                if take.sum() == 0:
                    take[conf.argmax()] = True
                seq[0, mpos[take]] = pred[take]
            cur = seq; gen += block_len
            if any(t in eos_ids for t in cur[0, -block_len:].tolist()):
                break
        return cur

    @torch.no_grad()
    def generate_self_spec(self, input_ids, k: int = 8, block_len: int = 32,
                           max_new_tokens: int = 128, eos_token_id=None):
        """Mode 3: self-speculative decode. Diffusion drafts k tokens (1 forward);
        AR verifies (1 forward) and accepts the longest matching prefix + 1
        correction. Output is identical to AR greedy. Returns the full sequence."""
        device = input_ids.device
        dtype = next(self.parameters()).dtype
        mask_id = int(self.config.mask_token_id)
        eos_ids = self._eos_ids(eos_token_id)
        cur = input_ids; gen = 0
        while gen < max_new_tokens:
            n = cur.shape[1]
            blk = torch.full((1, k), mask_id, device=device, dtype=cur.dtype)
            seq = torch.cat([cur, blk], 1); S = seq.shape[1]
            bias = build_block_causal_mask(S, block_len, device, dtype)
            dl = self(input_ids=seq, trimode_bias=bias).logits[0].float()
            draft = dl[n:n + k].argmax(-1)
            cand = torch.cat([cur, draft.unsqueeze(0)], 1)
            al = self(input_ids=cand).logits[0].float()
            ar_pred = al[n - 1:n + k - 1].argmax(-1)
            match = (draft == ar_pred)
            m = int((~match).float().argmax().item()) if (~match).any() else k
            new = (torch.cat([draft[:m], ar_pred[m:m + 1]]) if m < k
                   else torch.cat([draft, al[n + k - 1:n + k].argmax(-1)]))
            cur = torch.cat([cur, new.unsqueeze(0)], 1); gen += new.numel()
            if any(t in eos_ids for t in new.tolist()):
                break
        return cur

    def prepare_inputs_for_generation(
        self,
        input_ids: torch.Tensor,
        past_key_values: Optional[List[Tuple[torch.Tensor, ...]]] = None,
        attention_mask: Optional[torch.Tensor] = None,
        **kwargs
    ):
        # transformers >=5 pre-creates an EMPTY cache and passes it at prefill,
        # so we must slice by the cache's actual length (not None-ness): keep the
        # full prompt at prefill (past_len==0), only the new tokens during decode.
        # position_ids is intentionally not propagated so the model recomputes it
        # from past_len (keeps RoPE + Canon conv-state aligned).
        past_len = past_key_values.get_seq_length() if past_key_values is not None else 0
        if past_len > 0:
            input_ids = input_ids[:, past_len:]

        return {
            "input_ids": input_ids,
            "past_key_values": past_key_values,
            "attention_mask": attention_mask,
            "use_cache": True,
        }