# Copyright 2025 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """PyTorch Qwen3-Next model.""" from collections.abc import Callable import torch import torch.nn.functional as F from torch import nn from ... import initialization as init from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub_with_fallback, use_kernelized_func from ...integrations.accelerate import force_accelerate_hooks from ...masking_utils import create_causal_mask, create_recurrent_attention_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, is_torchdynamo_exporting, logging from ...utils.generic import merge_with_config_defaults, no_inherit_decorator from ...utils.output_capturing import OutputRecorder, capture_outputs from ..bamba.modeling_bamba import apply_mask_to_padding_states, apply_rotary_pos_emb from ..gemma2.modeling_gemma2 import Gemma2RotaryEmbedding from ..gemma3.modeling_gemma3 import Gemma3RMSNorm from ..llama.modeling_llama import ( LlamaForQuestionAnswering, LlamaForSequenceClassification, LlamaForTokenClassification, ) from ..mixtral.modeling_mixtral import MixtralForCausalLM from ..qwen2_moe.modeling_qwen2_moe import Qwen2MoeExperts, Qwen2MoeSparseMoeBlock, Qwen2MoeTopKRouter from ..qwen3_moe.modeling_qwen3_moe import ( Qwen3MoeAttention, Qwen3MoeDecoderLayer, Qwen3MoeMLP, eager_attention_forward, ) from .configuration_qwen3_next import Qwen3NextConfig logger = logging.get_logger(__name__) # NOTE: the FLA package does not re-cast to `input_dtype` in its implementation, maybe we should do the same @use_kernel_forward_from_hub("RMSNormGated") class Qwen3NextRMSNormGated(nn.Module): def __init__(self, hidden_size: int, eps: float = 1e-6, **kwargs) -> None: super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps self.activation = "silu" def forward(self, hidden_states: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: input_dtype = hidden_states.dtype hidden_states = hidden_states.to(torch.float32) variance = hidden_states.pow(2).mean(-1, keepdim=True) # Norm before gate hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) hidden_states = self.weight * hidden_states.to(input_dtype) hidden_states = hidden_states * ACT2FN[self.activation](gate.to(torch.float32)) return hidden_states.to(input_dtype) class Qwen3NextRotaryEmbedding(Gemma2RotaryEmbedding): def compute_default_rope_parameters(config: Qwen3NextConfig, device=None, **kwargs) -> tuple[torch.Tensor, float]: """ Computes the inverse frequencies according to the original RoPE implementation Args: config ([`~transformers.PreTrainedConfig`]): The model configuration. Returns: Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). """ base = config.rope_parameters["rope_theta"] partial_rotary_factor = config.rope_parameters.get("partial_rotary_factor", 1.0) head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads dim = int(head_dim * partial_rotary_factor) attention_factor = 1.0 # Unused in this type of RoPE # Compute the inverse frequencies inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float) / dim)) return inv_freq.to(device), attention_factor class Qwen3NextRMSNorm(Gemma3RMSNorm): pass @no_inherit_decorator class Qwen3NextAttention(Qwen3MoeAttention): def __init__(self, config: Qwen3NextConfig, layer_idx: int): super().__init__(config, layer_idx) self.q_proj = nn.Linear( config.hidden_size, config.num_attention_heads * self.head_dim * 2, bias=config.attention_bias ) del self.sliding_window def forward( self, hidden_states: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor], attention_mask: torch.Tensor | None, past_key_values: Cache | None = None, **kwargs: Unpack[FlashAttentionKwargs], ) -> tuple[torch.Tensor, torch.Tensor | None]: input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.head_dim) query_states, gate = torch.chunk( self.q_proj(hidden_states).view(*input_shape, -1, self.head_dim * 2), 2, dim=-1 ) gate = gate.reshape(*input_shape, -1) query_states = self.q_norm(query_states.view(hidden_shape)).transpose(1, 2) key_states = self.k_norm(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 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) if past_key_values is not None: key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( self.config._attn_implementation, eager_attention_forward ) attn_output, attn_weights = attention_interface( self, query_states, key_states, value_states, attention_mask, dropout=0.0 if not self.training else self.attention_dropout, scaling=self.scaling, **kwargs, ) attn_output = attn_output.reshape(*input_shape, -1).contiguous() attn_output = attn_output * torch.sigmoid(gate) attn_output = self.o_proj(attn_output) return attn_output, attn_weights @use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d") def causal_conv1d_update( hidden_states: torch.Tensor, conv_state: torch.Tensor, weight: nn.Parameter, bias: nn.Parameter | None = None, activation: str | None = None, ): _, hidden_size, seq_len = hidden_states.shape state_len = conv_state.shape[-1] hidden_states_new = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype) conv_state.copy_(hidden_states_new[:, :, -state_len:]) out = F.conv1d(hidden_states_new, weight.unsqueeze(1), bias, padding=0, groups=hidden_size) out = out[:, :, -seq_len:] if activation is not None: out = ACT2FN[activation](out) return out.to(hidden_states.dtype) @use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d") def causal_conv1d_fn( hidden_states: torch.Tensor, weight: nn.Parameter, bias: nn.Parameter | None = None, activation: str | None = None, **kwargs, ): _, hidden_size, seq_len = hidden_states.shape padding = weight.shape[-1] - 1 out = F.conv1d( hidden_states.to(weight.dtype), weight=weight.unsqueeze(1), bias=bias, padding=padding, groups=hidden_size, )[:, :, :seq_len] if activation is not None: out = ACT2FN[activation](out) return out.to(hidden_states.dtype) # NOTE: the FLA package computes `x / torch.sqrt((x * x).sum(dim=dim, keepdim=True) + eps)` instead, so if we align # with the GatedRMSNorm, maybe we can make that change as well. def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6): """This function is intended to align with the l2norm implementation in the FLA library.""" inv_norm = torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps) return x * inv_norm @use_kernel_func_from_hub_with_fallback("chunk_gated_delta_rule", "fla") def torch_chunk_gated_delta_rule( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, g: torch.Tensor, beta: torch.Tensor, chunk_size: int = 64, initial_state: torch.Tensor | None = None, output_final_state: bool = False, use_qk_l2norm_in_kernel: bool = False, **kwargs, ) -> tuple[torch.Tensor, torch.Tensor | None]: """Computes the gated delta rule, by chunking along the sequence dimension. Args: query: Query tensor of shape [batch_size, sequence_length, num_k_heads, k_head_dim] key: Key tensor of shape [batch_size, sequence_length, num_k_heads, k_head_dim] value: Value tensor of shape [batch_size, sequence_length, num_v_heads, v_head_dim]. num_v_heads can be equal to num_k_heads, same for v_head_dim and k_head_dim. g: Decay (in log space) tensor of shape [batch_size, sequence_length, num_v_heads]: the recurrent state is multiplied by exp(g) at each step, so entries must be <= 0. beta: Beta tensor of shape [batch_size, sequence_length, num_v_heads] chunk_size: Size of the chunks along the sequence dimension. initial_state: The recurrent state, an optional tensor of shape [batch_size, num_v_heads, k_head_dim, v_head_dim] output_final_state: Whether to output the new recurrent state along with the output. use_qk_l2norm_in_kernel: If this flag is set to True, query and key vectors are L2-normalized. Returns: - The output tensor of shape [batch_size, sequence_length, num_v_heads, v_head_dim] - Either None or the new recurrent state tensor of shape [batch_size, num_v_heads, k_head_dim, v_head_dim] """ initial_dtype = query.dtype batch_size, sequence_length, _, k_head_dim = key.shape num_v_heads, v_head_dim = value.shape[-2:] recurrent_state_shape = (batch_size, num_v_heads, k_head_dim, v_head_dim) padded_output_shape = (batch_size, num_v_heads, -1, v_head_dim) # -1 is the padded sequence length decay = g # rename for clarity: argument name must stay "g" to match flash_linear_attention's API # Make sure all tensors are fp32 and reshape them to [batch_size, num_*_heads, seqlen, ...] query, key, value, beta, decay = [ x.transpose(1, 2).to(torch.float32, memory_format=torch.contiguous_format) for x in (query, key, value, beta, decay) ] # If enabled, normalize query and key vectors (in fp32 to match the FLA library) if use_qk_l2norm_in_kernel: query = l2norm(query, dim=-1, eps=1e-6) key = l2norm(key, dim=-1, eps=1e-6) # And always normalize queries by the head dimension scaling = query.shape[-1] ** -0.5 query = query * scaling # Pad sequence length to be a multiple of chunk_size. Padding is described as (left_pad, right_pad) for each dim. pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size query, key, value = (F.pad(x, (0, 0, 0, pad_size)) for x in (query, key, value)) beta, decay = (F.pad(x, (0, pad_size)) for x in (beta, decay)) total_sequence_length = sequence_length + pad_size num_chunks = total_sequence_length // chunk_size # Apply beta to K and V, which is the "learning rate" of the recurrent state for a given token, ie. how much the new # state influence the old state. Beta is often normalized to (0, 1) where 0 = no update; 1 = overwrite old state. v_beta = value * beta.unsqueeze(-1) k_beta = key * beta.unsqueeze(-1) # Reshape all tensors to chunk the sequence dimension (adds a new dimension of size chunk_size) query, key, k_beta, v_beta = [ x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1]) for x in (query, key, k_beta, v_beta) ] decay = decay.reshape(decay.shape[0], decay.shape[1], -1, chunk_size) # Create a chunk-sized strictly upper triangular mask, ie. the mask of what a causal chunk may not attend to strictly_upper_mask = torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device).triu(1) # Cumulative decay within each chunk (dim 3 is the position inside the chunk). Since decay is in log space, # cum_decay[..., t] is the log of a product of decays between the start of the chunk and position t cum_decay = decay.cumsum(dim=3) # First phase: compute intra-chunk quantities. # The pairwise decays: pairwise_decay[..., i, j] = exp(cum_decay_i - cum_decay_j) is the decay accumulated between # positions j and i of a chunk. Positive values are masked to -inf before exp to avoid overflow pairwise_decay = cum_decay.unsqueeze(4) - cum_decay.unsqueeze(3) pairwise_decay = pairwise_decay.masked_fill(strictly_upper_mask, float("-inf")) pairwise_decay = pairwise_decay.exp() # with the exp, we exit log space, so we can apply this decay to the states # Compute auxiliary tensors: the Upper Triangular (ut) transform system and the intra-chunk attn (QK dot product) ut_system = (k_beta @ key.transpose(-1, -2)) * pairwise_decay intra_chunk_attn = (query @ key.transpose(-1, -2)) * pairwise_decay decayed_k_beta = k_beta * cum_decay.exp().unsqueeze(-1) # Gated delta attention uses a UT transform to condense several delta rule updates into a few matmuls. After the UT # system is solved, we can then compute the new_values (called "u" in the DeltaNet paper) and the decayed keys # reading the old state (k_cumdecay). In the update, the part of new_values that the old state already predicts is # subtracted out, so that only the correction is written to the recurrent state: this is the delta rule. # Not all export targets support the fast triangular solver, so we build the inverse by forward substitution then if not is_torchdynamo_exporting(): new_values = torch.linalg.solve_triangular(ut_system, v_beta, upper=False, unitriangular=True) k_cumdecay = torch.linalg.solve_triangular(ut_system, decayed_k_beta, upper=False, unitriangular=True) else: ut_system = -ut_system.tril(-1) # ut_system is masked to only keep the strictly lower triangle for i in range(1, chunk_size): row = ut_system[..., i, :i].clone() sub = ut_system[..., :i, :i].clone() ut_system[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2) ut_system = ut_system + torch.eye(chunk_size, dtype=ut_system.dtype, device=ut_system.device) new_values, k_cumdecay = ut_system @ v_beta, ut_system @ decayed_k_beta if initial_state is None: last_recurrent_state = torch.zeros(recurrent_state_shape, dtype=new_values.dtype, device=new_values.device) else: last_recurrent_state = initial_state.to(new_values) core_attn_out = torch.zeros_like(new_values) # Apply decay once rather than in each chunk query = query * cum_decay.exp().unsqueeze(-1) key = key * (cum_decay[..., -1:] - cum_decay).exp().unsqueeze(-1) chunk_decay = cum_decay[..., -1].exp()[..., None, None] # Second phase: the sequential scan over chunks for i in range(num_chunks): # Compute attention output for the current chunk: add the read of the previous recurrent state # (inter_chunk_attn) with the within-chunk attention (intra_chunk_attn) v_new = new_values[:, :, i] - k_cumdecay[:, :, i] @ last_recurrent_state inter_chunk_attn = query[:, :, i] @ last_recurrent_state core_attn_out[:, :, i] = inter_chunk_attn + intra_chunk_attn[:, :, i] @ v_new # Update the recurrent state: new recurrent state (S_t+1) = decayed old state (S_t * (I-βkk^T)) + update (βvk^T) last_recurrent_state = last_recurrent_state * chunk_decay[:, :, i] + key[:, :, i].transpose(-1, -2) @ v_new # Discard the final state if not requested last_recurrent_state = None if not output_final_state else last_recurrent_state # Reshape the output to the orignal shape: flatten the chunk dimension, then drop padding core_attn_out = core_attn_out.reshape(padded_output_shape) core_attn_out = core_attn_out[:, :, :sequence_length] # Convert back to the original shape [batch_size, sequence_length, num_v_heads, v_head_dim] and dtype core_attn_out = core_attn_out.transpose(1, 2).to(initial_dtype, memory_format=torch.contiguous_format) return core_attn_out, last_recurrent_state @use_kernel_func_from_hub_with_fallback("fused_recurrent_gated_delta_rule", "fla") def torch_recurrent_gated_delta_rule( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, g: torch.Tensor, beta: torch.Tensor, initial_state: torch.Tensor | None = None, output_final_state: bool = False, use_qk_l2norm_in_kernel: bool = False, **kwargs, ) -> tuple[torch.Tensor, torch.Tensor | None]: """Computes linear attention using the gated delta rule, by iterating over each token in the sequence dimension. Same args and return value as torch_chunk_gated_delta_rule, except for `chunk_size` because the sequence dim is not chunked.""" initial_dtype = query.dtype batch_size, sequence_length, _, k_head_dim = key.shape num_v_heads, v_head_dim = value.shape[-2:] decay = g # rename for clarity: argument name must stay "g" to match flash_linear_attention's API # Make sure all tensors are fp32 and reshape them to [batch_size, num_*_heads, seqlen, ...] query, key, value, beta, decay = [ x.transpose(1, 2).to(torch.float32, memory_format=torch.contiguous_format) for x in (query, key, value, beta, decay) ] # If enabled, normalize query and key vectors (done once in fp32 for better accuracy) if use_qk_l2norm_in_kernel: query = l2norm(query, dim=-1, eps=1e-6) key = l2norm(key, dim=-1, eps=1e-6) # And always normalize queries by the head dimension query = query / (query.shape[-1] ** 0.5) # Create the storage for the last recurrent state, which will be updated in place. If a previous state is provided, # it is the starting point, otherwise start with a zeroed buffer. if initial_state is None: recurrent_state_shape = (batch_size, num_v_heads, k_head_dim, v_head_dim) last_recurrent_state = torch.zeros(recurrent_state_shape, dtype=value.dtype, device=value.device) else: last_recurrent_state = initial_state.to(value) core_attn_out = torch.zeros_like(value) # Loop over each token and update the recurrent state for i in range(sequence_length): q_t, k_t, v_t = query[:, :, i], key[:, :, i], value[:, :, i] # Decay the recurrent state decay_t = decay[:, :, i].exp()[..., None, None] last_recurrent_state = last_recurrent_state * decay_t # Update the recurrent state with the current token beta_t = beta[:, :, i].unsqueeze(-1) kv_mem = (last_recurrent_state * k_t.unsqueeze(-1)).sum(dim=-2) delta = (v_t - kv_mem) * beta_t last_recurrent_state = last_recurrent_state + k_t.unsqueeze(-1) * delta.unsqueeze(-2) # And use it to compute the attention output for the current token core_attn_out[:, :, i] = (last_recurrent_state * q_t.unsqueeze(-1)).sum(dim=-2) # Discard the final state if not requested last_recurrent_state = None if not output_final_state else last_recurrent_state # Convert back to the original shape [batch_size, sequence_length, num_v_heads, v_head_dim] and dtype core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype) return core_attn_out, last_recurrent_state @use_kernelized_func( [torch_recurrent_gated_delta_rule, torch_chunk_gated_delta_rule, causal_conv1d_fn, causal_conv1d_update] ) class Qwen3NextGatedDeltaNet(nn.Module): def __init__(self, config: Qwen3NextConfig, layer_idx: int): super().__init__() self.hidden_size = config.hidden_size self.num_v_heads = config.linear_num_value_heads self.num_k_heads = config.linear_num_key_heads self.head_k_dim = config.linear_key_head_dim self.head_v_dim = config.linear_value_head_dim self.key_dim = self.head_k_dim * self.num_k_heads self.value_dim = self.head_v_dim * self.num_v_heads self.conv_kernel_size = config.linear_conv_kernel_dim self.layer_idx = layer_idx self.activation = config.hidden_act self.layer_norm_epsilon = config.rms_norm_eps # QKV self.conv_dim = self.key_dim * 2 + self.value_dim self.conv1d = nn.Conv1d( in_channels=self.conv_dim, out_channels=self.conv_dim, bias=False, kernel_size=self.conv_kernel_size, groups=self.conv_dim, padding=self.conv_kernel_size - 1, ) # projection of the input hidden states projection_size_qkvz = self.key_dim * 2 + self.value_dim * 2 projection_size_ba = self.num_v_heads * 2 self.in_proj_qkvz = nn.Linear(self.hidden_size, projection_size_qkvz, bias=False) self.in_proj_ba = nn.Linear(self.hidden_size, projection_size_ba, bias=False) # time step projection (discretization) # instantiate once and copy inv_dt in init_weights of PretrainedModel self.dt_bias = nn.Parameter(torch.ones(self.num_v_heads)) # Lower bound kept away from 0 so log(A) never becomes -inf A = torch.empty(self.num_v_heads).uniform_(0.01, 16) self.A_log = nn.Parameter(torch.log(A)) self.norm = Qwen3NextRMSNormGated(self.head_v_dim, eps=self.layer_norm_epsilon) self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False) self.layer_type = config.layer_types[layer_idx] def fix_query_key_value_ordering(self, mixed_qkvz, mixed_ba): """ Derives `query`, `key` and `value` tensors from `mixed_qkvz` and `mixed_ba`. """ new_tensor_shape_qkvz = mixed_qkvz.size()[:-1] + ( self.num_k_heads, 2 * self.head_k_dim + 2 * self.head_v_dim * self.num_v_heads // self.num_k_heads, ) new_tensor_shape_ba = mixed_ba.size()[:-1] + (self.num_k_heads, 2 * self.num_v_heads // self.num_k_heads) mixed_qkvz = mixed_qkvz.view(*new_tensor_shape_qkvz) mixed_ba = mixed_ba.view(*new_tensor_shape_ba) split_arg_list_qkvz = [ self.head_k_dim, self.head_k_dim, (self.num_v_heads // self.num_k_heads * self.head_v_dim), (self.num_v_heads // self.num_k_heads * self.head_v_dim), ] split_arg_list_ba = [self.num_v_heads // self.num_k_heads, self.num_v_heads // self.num_k_heads] query, key, value, z = torch.split(mixed_qkvz, split_arg_list_qkvz, dim=3) b, a = torch.split(mixed_ba, split_arg_list_ba, dim=3) # [b, sq, ng, np/ng * hn] -> [b, sq, np, hn] value = value.reshape(value.size(0), value.size(1), -1, self.head_v_dim) z = z.reshape(z.size(0), z.size(1), -1, self.head_v_dim) b = b.reshape(b.size(0), b.size(1), self.num_v_heads) a = a.reshape(a.size(0), a.size(1), self.num_v_heads) return query, key, value, z, b, a @force_accelerate_hooks("conv1d") def forward( self, hidden_states: torch.Tensor, cache_params: Cache | None = None, attention_mask: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ): hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) # Set up dimensions for reshapes later seq_len = hidden_states.shape[1] use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) projected_states_qkvz = self.in_proj_qkvz(hidden_states) projected_states_ba = self.in_proj_ba(hidden_states) query, key, value, z, b, a = self.fix_query_key_value_ordering(projected_states_qkvz, projected_states_ba) query, key, value = (x.reshape(x.shape[0], x.shape[1], -1) for x in (query, key, value)) mixed_qkv = torch.cat((query, key, value), dim=-1) mixed_qkv = mixed_qkv.transpose(1, 2) if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past: conv_state = cache_params.layers[self.layer_idx].conv_states[0] # Single-token cached decode: the fused per-step kernel updates the conv state in-place. mixed_qkv = causal_conv1d_update( mixed_qkv, conv_state, self.conv1d.weight.squeeze(1), self.conv1d.bias, self.activation, ) else: if cache_params is not None: mixed_qkv = cache_params.update_conv_state( mixed_qkv, self.layer_idx, conv_kernel_size=self.conv_kernel_size ) mixed_qkv = causal_conv1d_fn( mixed_qkv, self.conv1d.weight.squeeze(1), self.conv1d.bias, activation=self.activation, **kwargs, ) # Drop the additional previous states if cache_params is not None: mixed_qkv = mixed_qkv[:, :, -seq_len:] mixed_qkv = mixed_qkv.transpose(1, 2) query, key, value = torch.split( mixed_qkv, [ self.key_dim, self.key_dim, self.value_dim, ], dim=-1, ) query = query.reshape(query.shape[0], query.shape[1], -1, self.head_k_dim) key = key.reshape(key.shape[0], key.shape[1], -1, self.head_k_dim) value = value.reshape(value.shape[0], value.shape[1], -1, self.head_v_dim) beta = b.sigmoid() # If the model is loaded in fp16, without the .float() here, A might be -inf g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias) if self.num_v_heads // self.num_k_heads > 1: query = query.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2) key = key.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2) recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] if use_precomputed_states else None if use_precomputed_states and seq_len == 1: core_attn_out, last_recurrent_state = torch_recurrent_gated_delta_rule( query, key, value, g=g, beta=beta, initial_state=recurrent_state, output_final_state=cache_params is not None, use_qk_l2norm_in_kernel=True, cu_seqlens=kwargs.pop("cu_seq_lens_q", None), **kwargs, ) else: core_attn_out, last_recurrent_state = torch_chunk_gated_delta_rule( query, key, value, g=g, beta=beta, initial_state=recurrent_state, output_final_state=cache_params is not None, use_qk_l2norm_in_kernel=True, cu_seqlens=kwargs.pop("cu_seq_lens_q", None), **kwargs, ) # Update cache if cache_params is not None: cache_params.update_recurrent_state(last_recurrent_state, self.layer_idx) z_shape_og = z.shape # reshape input data into 2D tensor core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1]) z = z.reshape(-1, z.shape[-1]) core_attn_out = self.norm(core_attn_out, z) core_attn_out = core_attn_out.reshape(z_shape_og) core_attn_out = core_attn_out.reshape(core_attn_out.shape[0], core_attn_out.shape[1], -1) output = self.out_proj(core_attn_out) return output class Qwen3NextMLP(Qwen3MoeMLP): pass class Qwen3NextExperts(Qwen2MoeExperts): pass class Qwen3NextTopKRouter(Qwen2MoeTopKRouter): pass class Qwen3NextSparseMoeBlock(Qwen2MoeSparseMoeBlock): pass class Qwen3NextDecoderLayer(Qwen3MoeDecoderLayer): def __init__(self, config: Qwen3NextConfig, layer_idx: int): nn.Module.__init__(self) self.hidden_size = config.hidden_size # token mixer self.block_type = config.layer_types[layer_idx] if self.block_type == "linear_attention": self.linear_attn = Qwen3NextGatedDeltaNet(config, layer_idx) elif self.block_type == "full_attention": self.self_attn = Qwen3NextAttention(config, layer_idx) if (layer_idx not in config.mlp_only_layers) and ( config.num_experts > 0 and (layer_idx + 1) % config.decoder_sparse_step == 0 ): self.mlp = Qwen3NextSparseMoeBlock(config) else: self.mlp = Qwen3NextMLP(config, intermediate_size=config.intermediate_size) self.input_layernorm = Qwen3NextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = Qwen3NextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) def forward( self, hidden_states: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor], attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, **kwargs: Unpack[FlashAttentionKwargs], ) -> torch.FloatTensor: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) # Token Mixer if self.block_type == "linear_attention": hidden_states = self.linear_attn( hidden_states=hidden_states, cache_params=past_key_values, attention_mask=attention_mask, **kwargs, ) elif self.block_type == "full_attention": # Self Attention hidden_states, _ = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, position_embeddings=position_embeddings, **kwargs, ) hidden_states = residual + hidden_states # Fully Connected residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) # For the MoE layers, we need to unpack if isinstance(hidden_states, tuple): hidden_states, _ = hidden_states hidden_states = residual + hidden_states return hidden_states class Qwen3NextPreTrainedModel(PreTrainedModel): config: Qwen3NextConfig base_model_prefix = "model" supports_gradient_checkpointing = True _no_split_modules = ["Qwen3NextDecoderLayer"] _skip_keys_device_placement = ["past_key_values"] _supports_flash_attn = True _supports_sdpa = True _keys_to_ignore_on_load_unexpected = [r"^mtp.*"] _can_record_outputs = { "router_logits": OutputRecorder(Qwen3NextTopKRouter, index=0), "hidden_states": Qwen3NextDecoderLayer, "attentions": Qwen3NextAttention, } _is_stateful = True _can_compile_fullgraph = True @torch.no_grad() def _init_weights(self, module): super()._init_weights(module) if isinstance(module, Qwen3NextGatedDeltaNet): init.ones_(module.dt_bias) # Lower bound kept away from 0 so log(A) never becomes -inf init.copy_( module.A_log, torch.empty(module.num_v_heads, device=module.A_log.device).uniform_(0.01, 16).log_(), ) # We initialize with 0s to be 1 centered as the RMSNorm here does (1 + weight) elif isinstance(module, Qwen3NextRMSNorm): init.zeros_(module.weight) elif isinstance(module, Qwen3NextExperts): init.normal_(module.gate_up_proj, mean=0.0, std=self.config.initializer_range) init.normal_(module.down_proj, mean=0.0, std=self.config.initializer_range) elif isinstance(module, Qwen3NextSparseMoeBlock): init.normal_(module.gate.weight, mean=0.0, std=self.config.initializer_range) class Qwen3NextModel(Qwen3NextPreTrainedModel): def __init__(self, config: Qwen3NextConfig): super().__init__(config) self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, config.pad_token_id) self.layers = nn.ModuleList( [Qwen3NextDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] ) self.norm = Qwen3NextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = Qwen3NextRotaryEmbedding(config=config) self.gradient_checkpointing = False # Initialize weights and apply final processing self.post_init() @merge_with_config_defaults @capture_outputs @auto_docstring def forward( self, input_ids: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, inputs_embeds: torch.FloatTensor | None = None, use_cache: bool | None = None, **kwargs: Unpack[TransformersKwargs], ) -> MoeModelOutputWithPast: 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) if use_cache and past_key_values is None: past_key_values = DynamicCache(config=self.config) if position_ids is None: past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens position_ids = position_ids.unsqueeze(0) if not isinstance(causal_mask_mapping := attention_mask, dict): # Prepare mask arguments mask_kwargs = { "config": self.config, "inputs_embeds": inputs_embeds, "attention_mask": attention_mask, "past_key_values": past_key_values, "position_ids": position_ids, } # Create the masks causal_mask_mapping = { "full_attention": create_causal_mask(**mask_kwargs), "linear_attention": create_recurrent_attention_mask(**mask_kwargs), } hidden_states = inputs_embeds position_embeddings = self.rotary_emb(hidden_states, position_ids) for i, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]): hidden_states = decoder_layer( hidden_states, position_embeddings=position_embeddings, attention_mask=causal_mask_mapping[self.config.layer_types[i]], position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, **kwargs, ) hidden_states = self.norm(hidden_states) return MoeModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=past_key_values, ) class Qwen3NextForCausalLM(MixtralForCausalLM): def __init__(self, config): super().__init__(config) self.num_experts = config.num_experts def forward( self, input_ids: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, output_router_logits: bool | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], ) -> MoeCausalLMOutputWithPast: r""" Example: ```python >>> from transformers import AutoTokenizer, Qwen3NextForCausalLM >>> model = Qwen3NextForCausalLM.from_pretrained("Qwen/Qwen3-Next-80B-A3B-Instruct") >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-Next-80B-A3B-Instruct") >>> prompt = "Hey, are you conscious? Can you talk to me?" >>> inputs = tokenizer(prompt, return_tensors="pt") >>> # Generate >>> generate_ids = model.generate(inputs.input_ids, max_length=30) >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." ```""" return super().forward( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, labels=labels, use_cache=use_cache, output_router_logits=output_router_logits, logits_to_keep=logits_to_keep, **kwargs, ) class Qwen3NextForSequenceClassification(LlamaForSequenceClassification): pass class Qwen3NextForTokenClassification(LlamaForTokenClassification): pass class Qwen3NextForQuestionAnswering(LlamaForQuestionAnswering): pass __all__ = [ "Qwen3NextForCausalLM", "Qwen3NextForQuestionAnswering", "Qwen3NextModel", "Qwen3NextPreTrainedModel", "Qwen3NextForSequenceClassification", "Qwen3NextForTokenClassification", ]