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)