core.py

from __future__ import annotations
from collections.abc import Callable
from dataclasses import dataclass
from typing import Optional, Union
import torch
from torch import nn
from transformers.utils.generic import check_model_inputs
from transformers.activations import ACT2FN
from transformers.cache_utils import Cache, DynamicCache
from transformers.integrations import use_kernel_forward_from_hub
from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
from transformers.modeling_layers import GradientCheckpointingLayer
from transformers.modeling_outputs import BaseModelOutputWithPast
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
from transformers.processing_utils import Unpack
from transformers.utils import TransformersKwargs, auto_docstring
from .configuration_nlduni import MinistralDLMConfig

def rotate_half(x):
    x1 = x[..., :x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2:]
    return torch.cat((-x2, x1), dim=-1)

def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
    cos = cos.unsqueeze(unsqueeze_dim)
    sin = sin.unsqueeze(unsqueeze_dim)
    q_embed = q * cos + rotate_half(q) * sin
    k_embed = k * cos + rotate_half(k) * sin
    return (q_embed, k_embed)

def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
    (batch, num_key_value_heads, slen, head_dim) = hidden_states.shape
    if n_rep == 1:
        return hidden_states
    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)

def eager_attention_forward(module: nn.Module, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attention_mask: Optional[torch.Tensor], scaling: float, dropout: float=0.0, **kwargs: Unpack[TransformersKwargs]):
    key_states = repeat_kv(key, module.num_key_value_groups)
    value_states = repeat_kv(value, module.num_key_value_groups)
    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
    if attention_mask is not None:
        causal_mask = attention_mask[:, :, :, :key_states.shape[-2]]
        attn_weights = attn_weights + causal_mask
    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=False)
    attn_output = torch.matmul(attn_weights, value_states)
    attn_output = attn_output.transpose(1, 2).contiguous()
    return (attn_output, attn_weights)

def _get_llama_4_attn_scale(positions_ids: torch.Tensor, beta: float, max_position_embeddings: int) -> torch.Tensor:
    scaling = 1 + beta * torch.log(1 + torch.floor(positions_ids / max_position_embeddings))
    return scaling.unsqueeze(-1)

@dataclass(frozen=True, slots=True)
class _MixedTokenRoutingPlan:
    mask: torch.Tensor
    understanding_indices: torch.Tensor
    generation_indices: torch.Tensor

    @classmethod
    def from_mask(cls, mask: torch.Tensor) -> '_MixedTokenRoutingPlan':
        flat_mask = mask.reshape(-1)
        return cls(mask=mask, understanding_indices=(~flat_mask).nonzero(as_tuple=False).flatten(), generation_indices=flat_mask.nonzero(as_tuple=False).flatten())
TokenRouting = Optional[Union[bool, _MixedTokenRoutingPlan]]

def _resolve_dual_tower_layers(config: MinistralDLMConfig) -> int:
    num_hidden_layers = config.num_hidden_layers
    configured = 18
    if configured is None:
        return num_hidden_layers
    if isinstance(configured, bool) or not isinstance(configured, int):
        raise TypeError('dual_tower_layers must be an integer')
    if not 0 <= configured <= num_hidden_layers:
        raise ValueError(f'dual_tower_layers must be between 0 and num_hidden_layers ({num_hidden_layers}), got {configured}')
    return configured

def _routing_mask(routing: TokenRouting, hidden_states: torch.Tensor) -> torch.Tensor:
    if routing is None:
        return torch.zeros(hidden_states.shape[:2], dtype=torch.bool, device=hidden_states.device)
    if routing is True:
        return torch.ones(hidden_states.shape[:2], dtype=torch.bool, device=hidden_states.device)
    return routing.mask

def _create_split_tower_block_mask(block_mask, *, current_routing: torch.Tensor, cached_routing: torch.Tensor):
    from torch.nn.attention.flex_attention import BlockMask, create_block_mask
    if isinstance(block_mask, str):
        if block_mask != 'full':
            raise ValueError(f'Unsupported attention block-mask sentinel: {block_mask}')
        base_mask_mod = None
        query_length = current_routing.shape[1]
        key_value_length = cached_routing.shape[1] + query_length
        block_size = 128
    elif isinstance(block_mask, BlockMask):
        base_mask_mod = block_mask.mask_mod
        (query_length, key_value_length) = block_mask.seq_lengths
        block_size = block_mask.BLOCK_SIZE
    else:
        raise TypeError("partial dual-tower FlexAttention requires an explicit BlockMask or the 'full' sentinel")
    current_length = current_routing.shape[1]
    cached_length = cached_routing.shape[1]
    if query_length < current_length:
        raise ValueError(f'FlexAttention block-mask query length is shorter than routing: {query_length} < {current_length}')
    if key_value_length < cached_length + current_length:
        raise ValueError(f'FlexAttention block-mask key/value length is shorter than cached and current routing: {key_value_length} < {cached_length + current_length}')
    query_routing = torch.nn.functional.pad(current_routing, (0, query_length - current_length), value=False)
    key_value_routing = torch.cat((cached_routing, current_routing), dim=1)
    key_value_routing = torch.nn.functional.pad(key_value_routing, (0, key_value_length - key_value_routing.shape[1]), value=False)

    def split_tower_mask(b, h, q_idx, kv_idx):
        same_tower = query_routing[b, q_idx] == key_value_routing[b, kv_idx]
        same_tower = same_tower | (q_idx >= current_length)
        if base_mask_mod is None:
            return same_tower
        return base_mask_mod(b, h, q_idx, kv_idx) & same_tower
    return create_block_mask(split_tower_mask, B=current_routing.shape[0], H=None, Q_LEN=query_length, KV_LEN=key_value_length, device=current_routing.device, BLOCK_SIZE=block_size)

def _canonicalize_token_routing(routing: Optional[torch.Tensor], hidden_states: torch.Tensor, *, dual_tower: bool) -> TokenRouting:
    if routing is None:
        return None
    if not torch.is_tensor(routing) or routing.ndim != 2:
        raise TypeError('routing must be a rank-2 tensor with shape [batch, sequence]')
    if tuple(routing.shape) != tuple(hidden_states.shape[:2]):
        raise ValueError(f'routing must match the transformer token shape: expected {tuple(hidden_states.shape[:2])}, got {tuple(routing.shape)}')
    if routing.device != hidden_states.device:
        raise ValueError('routing and transformer hidden states must share a device')
    if routing.dtype == torch.bool:
        normalized = routing
    elif routing.dtype in {torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64}:
        if bool(((routing != 0) & (routing != 1)).any()):
            raise ValueError('routing values must be binary (0 for U, 1 for G)')
        normalized = routing.to(dtype=torch.bool)
    else:
        raise TypeError('routing must have boolean or integer dtype')
    has_generation_route = bool(normalized.any())
    if has_generation_route and (not dual_tower):
        raise ValueError('routing selects G but config.dual_tower is false')
    if not has_generation_route:
        return None
    if bool(normalized.all()):
        return True
    return _MixedTokenRoutingPlan.from_mask(normalized)

def _routed_module(hidden_states: torch.Tensor, routing: TokenRouting, understanding_module: Callable[[torch.Tensor], torch.Tensor], generation_module: Optional[Callable[[torch.Tensor], torch.Tensor]]) -> torch.Tensor:
    if routing is None:
        return understanding_module(hidden_states)
    if routing is True:
        if generation_module is None:
            raise RuntimeError('G routing requires generation-tower parameters')
        return generation_module(hidden_states)
    if generation_module is None:
        raise RuntimeError('mixed routing requires generation-tower parameters')
    if not isinstance(routing, _MixedTokenRoutingPlan):
        raise TypeError('mixed routing must use a canonical dispatch plan')
    flat_hidden = hidden_states.reshape(-1, hidden_states.shape[-1])
    understanding_output = understanding_module(flat_hidden.index_select(0, routing.understanding_indices))
    generation_output = generation_module(flat_hidden.index_select(0, routing.generation_indices))
    flat_output = understanding_output.new_empty((flat_hidden.shape[0], *understanding_output.shape[1:]))
    flat_output = flat_output.index_copy(0, routing.understanding_indices, understanding_output)
    flat_output = flat_output.index_copy(0, routing.generation_indices, generation_output)
    return flat_output.reshape(*hidden_states.shape[:-1], *flat_output.shape[1:])

class Ministral3Attention(nn.Module):

    def __init__(self, config: MinistralDLMConfig, layer_idx: int):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.split_tower = layer_idx >= _resolve_dual_tower_layers(config)
        self.head_dim = getattr(config, 'head_dim', None) or config.hidden_size // config.num_attention_heads
        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
        self.scaling = self.head_dim ** -0.5
        self.attention_dropout = config.attention_dropout
        self.is_causal = True
        self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
        self.q_proj_gen = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)
        self.k_proj_gen = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)
        self.v_proj_gen = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)
        self.o_proj_gen = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
        self.diffusion_lm = config.diffusion_lm

    def _project_qkv(self, hidden_states: torch.Tensor, routing: TokenRouting) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        return (_routed_module(hidden_states, routing, self.q_proj, getattr(self, 'q_proj_gen', None)), _routed_module(hidden_states, routing, self.k_proj, getattr(self, 'k_proj_gen', None)), _routed_module(hidden_states, routing, self.v_proj, getattr(self, 'v_proj_gen', None)))

    def _project_output(self, hidden_states: torch.Tensor, routing: TokenRouting) -> torch.Tensor:
        return _routed_module(hidden_states, routing, self.o_proj, getattr(self, 'o_proj_gen', None))

    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, use_cache: Optional[bool]=False, routing: TokenRouting=None, **kwargs: Unpack[FlashAttentionKwargs]) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
        input_shape = hidden_states.shape[:-1]
        hidden_shape = (*input_shape, -1, self.head_dim)
        (query_states, key_states, value_states) = self._project_qkv(hidden_states, routing)
        query_states = query_states.view(hidden_shape).transpose(1, 2)
        key_states = key_states.view(hidden_shape).transpose(1, 2)
        value_states = value_states.view(hidden_shape).transpose(1, 2)
        (cos, sin) = position_embeddings
        (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:
            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:
                (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)
        attention_interface: Callable = eager_attention_forward
        if self.config._attn_implementation != 'eager':
            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
        if self.diffusion_lm:
            (attn_output, attn_weights) = attention_interface(self, query_states, key_states, value_states, None, dropout=0.0, scaling=self.scaling, is_causal=False, **kwargs)
        else:
            (attn_output, attn_weights) = attention_interface(self, query_states, key_states, value_states, attention_mask, dropout=0.0, scaling=self.scaling, sliding_window=getattr(self.config, 'sliding_window', None), **kwargs)
        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
        attn_output = self._project_output(attn_output, routing)
        return (attn_output, attn_weights)

class Ministral3MLP(nn.Module):

    def __init__(self, config):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size
        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
        self.gate_proj_gen = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.up_proj_gen = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.down_proj_gen = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
        self.act_fn = ACT2FN[config.hidden_act]

    def _understanding_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.down_proj(self.act_fn(self.gate_proj(hidden_states)) * self.up_proj(hidden_states))

    def _generation_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.down_proj_gen(self.act_fn(self.gate_proj_gen(hidden_states)) * self.up_proj_gen(hidden_states))

    def forward(self, hidden_states: torch.Tensor, routing: TokenRouting=None) -> torch.Tensor:
        generation_module = self._generation_forward if hasattr(self, 'gate_proj_gen') else None
        return _routed_module(hidden_states, routing, self._understanding_forward, generation_module)

@use_kernel_forward_from_hub('RMSNorm')
class Ministral3RMSNorm(nn.Module):

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

    def forward(self, hidden_states):
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return self.weight * hidden_states.to(input_dtype)

    def extra_repr(self):
        return f'{tuple(self.weight.shape)}, eps={self.variance_epsilon}'

class Ministral3DecoderLayer(GradientCheckpointingLayer):

    def __init__(self, config: MinistralDLMConfig, layer_idx: int):
        super().__init__()
        self.hidden_size = config.hidden_size
        from .attention import MinistralFlexAttention
        attn_class = MinistralFlexAttention
        self.self_attn = attn_class(config=config, layer_idx=layer_idx)
        self.mlp = Ministral3MLP(config)
        self.input_layernorm = Ministral3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = Ministral3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.input_layernorm_gen = Ministral3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm_gen = Ministral3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor]=None, position_ids: Optional[torch.LongTensor]=None, past_key_values: Optional[Cache]=None, use_cache: Optional[bool]=False, cache_position: Optional[torch.LongTensor]=None, position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]]=None, routing: TokenRouting=None, **kwargs: Unpack[TransformersKwargs]) -> torch.Tensor:
        residual = hidden_states
        hidden_states = _routed_module(hidden_states, routing, self.input_layernorm, getattr(self, 'input_layernorm_gen', None))
        (hidden_states, _) = self.self_attn(hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, routing=routing, **kwargs)
        hidden_states = residual + hidden_states
        residual = hidden_states
        hidden_states = _routed_module(hidden_states, routing, self.post_attention_layernorm, getattr(self, 'post_attention_layernorm_gen', None))
        hidden_states = self.mlp(hidden_states, routing=routing)
        hidden_states = residual + hidden_states
        return hidden_states

class Ministral3PreTrainedModel(PreTrainedModel):
    config: MinistralDLMConfig
    base_model_prefix = 'model'
    supports_gradient_checkpointing = True
    _no_split_modules = ['Ministral3DecoderLayer']
    _skip_keys_device_placement = ['past_key_values']
    _supports_flash_attn = True
    _supports_sdpa = True
    _supports_flex_attn = True
    _can_compile_fullgraph = True
    _supports_attention_backend = True
    _can_record_outputs = {'hidden_states': Ministral3DecoderLayer, 'attentions': Ministral3Attention}

class Ministral3RotaryEmbedding(nn.Module):
    inv_freq: torch.Tensor

    def __init__(self, config: MinistralDLMConfig, device=None):
        super().__init__()
        self.max_seq_len_cached = config.max_position_embeddings
        self.original_max_seq_len = config.max_position_embeddings
        self.config = config
        self.rope_type = self.config.rope_parameters['rope_type']
        rope_init_fn: Callable = self.compute_default_rope_parameters
        if self.rope_type != 'default':
            rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
        (inv_freq, self.attention_scaling) = rope_init_fn(self.config, device)
        self.register_buffer('inv_freq', inv_freq, persistent=False)
        self.original_inv_freq = inv_freq

    @staticmethod
    def compute_default_rope_parameters(config: Optional[MinistralDLMConfig]=None, device: Optional['torch.device']=None, seq_len: Optional[int]=None) -> tuple['torch.Tensor', float]:
        base = config.rope_parameters['rope_theta']
        dim = getattr(config, 'head_dim', None) or config.hidden_size // config.num_attention_heads
        attention_factor = 1.0
        inv_freq = 1.0 / base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
        return (inv_freq, attention_factor)

    @torch.no_grad()
    @dynamic_rope_update
    def forward(self, x, position_ids):
        with torch.autocast(device_type=x.device.type, enabled=False):
            inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
            position_ids_expanded = position_ids[:, None, :].float()
            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
            emb = torch.cat((freqs, freqs), dim=-1)
            cos = emb.cos() * self.attention_scaling
            sin = emb.sin() * self.attention_scaling
        return (cos.to(dtype=x.dtype), sin.to(dtype=x.dtype))

class Ministral3Model(Ministral3PreTrainedModel):

    def __init__(self, config: MinistralDLMConfig):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size
        self.dual_tower_layers = _resolve_dual_tower_layers(config)
        self._layer_loop_config = None
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
        self.layers = nn.ModuleList([Ministral3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)])
        self.norm = Ministral3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.norm_gen = Ministral3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.rotary_emb = Ministral3RotaryEmbedding(config=config)
        self.gradient_checkpointing = False
        self.post_init()

    @check_model_inputs
    def forward(self, input_ids: Optional[torch.LongTensor]=None, attention_mask: Optional[torch.Tensor]=None, position_ids: Optional[torch.LongTensor]=None, past_key_values: Optional[Cache]=None, inputs_embeds: Optional[torch.FloatTensor]=None, use_cache: Optional[bool]=None, cache_position: Optional[torch.LongTensor]=None, input_embeddings: Optional[torch.FloatTensor]=None, routing: Optional[torch.Tensor]=None, modality_indices: Optional[torch.Tensor]=None, loop_k: Optional[int]=None, loop_depths: Optional[torch.LongTensor]=None, **kwargs: Unpack[TransformersKwargs]) -> BaseModelOutputWithPast:
        if input_embeddings is not None:
            inputs_embeds = input_embeddings
            input_embeddings = None
        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError('You must specify exactly one of input_ids or inputs_embeds')
        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)
        single_stream_call = kwargs.get('single_stream_image_diffusion', False)
        if single_stream_call is None:
            single_stream_call = False
        if not isinstance(single_stream_call, bool):
            raise TypeError('single_stream_image_diffusion must be boolean')
        use_layer_loop = self._layer_loop_config is not None and single_stream_call
        if loop_k is not None or loop_depths is not None:
            raise ValueError('loop_k and loop_depths require an active single-stream looped_layer_execution call')
        resolved_loop_depths = None
        loop_visits = 1
        dual_tower = True
        partial_dual_tower = dual_tower and self.dual_tower_layers < self.config.num_hidden_layers
        if dual_tower and modality_indices is not None:
            modality_routing = _canonicalize_token_routing(modality_indices, inputs_embeds, dual_tower=dual_tower)
            if routing is None:
                canonical_routing = modality_routing
            else:
                canonical_routing = _canonicalize_token_routing(routing, inputs_embeds, dual_tower=dual_tower)
                if not torch.equal(routing.to(dtype=torch.bool), modality_indices.to(dtype=torch.bool)):
                    raise ValueError('routing and deprecated modality_indices alias must select identical U/G tokens when both are provided')
        else:
            canonical_routing = _canonicalize_token_routing(routing, inputs_embeds, dual_tower=dual_tower)
        if use_cache and past_key_values is None:
            past_key_values = DynamicCache()
        past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None and (partial_dual_tower or cache_position is None) else 0
        stored_cached_routing = getattr(past_key_values, '_nlduni_dual_tower_routing', None)
        track_dual_tower_routing = partial_dual_tower and (canonical_routing is not None or stored_cached_routing is not None)
        current_routing = _routing_mask(canonical_routing, inputs_embeds) if track_dual_tower_routing else None
        cached_routing = stored_cached_routing
        if track_dual_tower_routing and cached_routing is None:
            cached_routing = torch.zeros(current_routing.shape[0], past_seen_tokens, dtype=torch.bool, device=current_routing.device)
        elif track_dual_tower_routing and tuple(cached_routing.shape) != (current_routing.shape[0], past_seen_tokens):
            raise ValueError(f'cached dual-tower routing must match the KV cache shape: expected {(current_routing.shape[0], past_seen_tokens)}, got {tuple(cached_routing.shape)}')
        elif track_dual_tower_routing and cached_routing.device != current_routing.device:
            raise ValueError('cached dual-tower routing and inputs must share a device')
        if track_dual_tower_routing and past_key_values is not None:
            past_key_values._nlduni_dual_tower_routing = cached_routing
            past_key_values._nlduni_dual_tower_has_generation = bool(getattr(past_key_values, '_nlduni_dual_tower_has_generation', False))
        if cache_position is None:
            cache_position = torch.arange(past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device)
        if position_ids is None:
            position_ids = cache_position.unsqueeze(0)
        if kwargs.get('use_causal_mask', False):
            mask_function = create_causal_mask if self.config.sliding_window is None else create_sliding_window_causal_mask
            causal_mask = mask_function(config=self.config, input_embeds=inputs_embeds, attention_mask=attention_mask, cache_position=cache_position, past_key_values=past_key_values, position_ids=position_ids)
        else:
            causal_mask = None
        hidden_states = inputs_embeds
        position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
        split_layer_kwargs = None
        has_split_layers = track_dual_tower_routing
        if has_split_layers and kwargs.get('overwrite_attn_impl') != 'flash_attn' and (kwargs.get('overwrite_block_mask') is not None):
            split_layer_kwargs = dict(kwargs)
            split_layer_kwargs['overwrite_block_mask'] = _create_split_tower_block_mask(kwargs['overwrite_block_mask'], current_routing=current_routing, cached_routing=cached_routing)
        for (layer_idx, decoder_layer) in enumerate(self.layers[:self.config.num_hidden_layers]):
            layer_kwargs = split_layer_kwargs if split_layer_kwargs is not None and layer_idx >= self.dual_tower_layers else kwargs
            hidden_states = decoder_layer(hidden_states, attention_mask=causal_mask, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, routing=canonical_routing, **layer_kwargs)
        if track_dual_tower_routing and use_cache and (past_key_values is not None):
            past_key_values._nlduni_dual_tower_routing = torch.cat((cached_routing, current_routing), dim=1)
            past_key_values._nlduni_dual_tower_has_generation = past_key_values._nlduni_dual_tower_has_generation or canonical_routing is not None
        hidden_states = _routed_module(hidden_states, canonical_routing, self.norm, getattr(self, 'norm_gen', None))
        return BaseModelOutputWithPast(last_hidden_state=hidden_states, past_key_values=past_key_values if use_cache else None)