# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 import copy import torch import torch.nn as nn import torch.nn.functional as F from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec from megatron.core.transformer.spec_utils import ModuleSpec from megatron.core.transformer.transformer_block import get_num_layers_to_build from megatron.core.transformer.transformer_layer import get_transformer_layer_offset from transformers import AutoConfig from transformers.activations import ACT2FN try: from fla.modules import FusedRMSNormGated, ShortConvolution from fla.ops.gated_delta_rule import chunk_gated_delta_rule from transformers.models.qwen3_next.modeling_qwen3_next import Qwen3NextAttention, Qwen3NextRMSNorm except ImportError: pass from .hf_attention import HuggingfaceAttention # adapt from https://github.com/huggingface/transformers/blob/38a08b6e8ae35857109cedad75377997fecbf9d0/src/transformers/models/qwen3_next/modeling_qwen3_next.py#L564 class Qwen3NextGatedDeltaNet(nn.Module): """ Qwen3NextGatedDeltaNet with varlen support """ def __init__(self, config, 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.act = ACT2FN[config.hidden_act] self.layer_norm_epsilon = config.rms_norm_eps # QKV self.conv_dim = self.key_dim * 2 + self.value_dim self.conv1d = ShortConvolution( hidden_size=self.conv_dim, bias=False, kernel_size=self.conv_kernel_size, ) # 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)) A = torch.empty(self.num_v_heads).uniform_(0, 16) self.A_log = nn.Parameter(torch.log(A)) self.norm = FusedRMSNormGated( self.head_v_dim, eps=self.layer_norm_epsilon, activation=self.activation, device=torch.cuda.current_device(), dtype=config.dtype if config.dtype is not None else torch.get_current_dtype(), ) self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False) 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 def forward( self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor = None, ): 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, _ = self.conv1d( x=mixed_qkv, cu_seqlens=cu_seqlens, ) 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) core_attn_out, last_recurrent_state = chunk_gated_delta_rule( query, key, value, g=g, beta=beta, initial_state=None, output_final_state=False, use_qk_l2norm_in_kernel=True, ) 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 Attention(HuggingfaceAttention): def __init__( self, args, config, layer_number: int, cp_comm_type: str = "p2p", pg_collection=None, ): super().__init__( args, config, layer_number, cp_comm_type, pg_collection, ) if Qwen3NextAttention is None: raise ImportError("Please install transformers>=4.35.0 to use Qwen3NextAttention.") self.linear_attn = Qwen3NextGatedDeltaNet(self.hf_config, self.hf_layer_idx) self.input_layernorm = Qwen3NextRMSNorm(self.hf_config.hidden_size, eps=self.hf_config.rms_norm_eps) def hf_forward(self, hidden_states, packed_seq_params): hidden_states = self.input_layernorm(hidden_states) hidden_states = self.linear_attn( hidden_states=hidden_states, cu_seqlens=packed_seq_params.cu_seqlens_q, ) return hidden_states def get_qwen3_next_spec(args, config, vp_stage): # always use the moe path if not args.num_experts: config.moe_layer_freq = [0] * config.num_layers # Define the decoder block spec kwargs = { "use_transformer_engine": True, } if vp_stage is not None: kwargs["vp_stage"] = vp_stage transformer_layer_spec = get_gpt_decoder_block_spec(config, **kwargs) assert config.pipeline_model_parallel_layout is None, "not support this at the moment" # Slice the layer specs to only include the layers that are built in this pipeline stage. # Note: MCore layer_number starts at 1 num_layers_to_build = get_num_layers_to_build(config, vp_stage=vp_stage) offset = get_transformer_layer_offset(config, vp_stage=vp_stage) hf_config = AutoConfig.from_pretrained(args.hf_checkpoint, trust_remote_code=True) for layer_id in range(num_layers_to_build): if hf_config.layer_types[layer_id + offset] == "linear_attention": layer_specs = copy.deepcopy(transformer_layer_spec.layer_specs[layer_id]) layer_specs.submodules.self_attention = ModuleSpec( module=Attention, params={"args": args}, ) transformer_layer_spec.layer_specs[layer_id] = layer_specs return transformer_layer_spec