diff --git a/python/sglang/srt/models/qwen3_next.py b/python/sglang/srt/models/qwen3_next.py index 28f36f661..ae17e4319 100644 --- a/python/sglang/srt/models/qwen3_next.py +++ b/python/sglang/srt/models/qwen3_next.py @@ -202,6 +202,7 @@ class Qwen3GatedDeltaNet(nn.Module): layer_id: int, quant_config: Optional[QuantizationConfig] = None, alt_stream: Optional[torch.cuda.Stream] = None, + prefix: str = "", ) -> None: super().__init__() self.config = config @@ -229,6 +230,7 @@ class Qwen3GatedDeltaNet(nn.Module): quant_config=None, tp_rank=self.attn_tp_rank, tp_size=self.attn_tp_size, + prefix=add_prefix("conv1d", prefix), ) self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1) projection_size_qkvz = self.key_dim * 2 + self.value_dim * 2 @@ -241,14 +243,16 @@ class Qwen3GatedDeltaNet(nn.Module): quant_config=quant_config, tp_rank=self.attn_tp_rank, tp_size=self.attn_tp_size, + prefix=add_prefix("in_proj_qkvz", prefix), ) self.in_proj_ba = ColumnParallelLinear( input_size=self.hidden_size, output_size=projection_size_ba, bias=False, - quant_config=None, + quant_config=quant_config, tp_rank=self.attn_tp_rank, tp_size=self.attn_tp_size, + prefix=add_prefix("in_proj_ba", prefix), ) query_key_settings = (self.key_dim, 0, False) @@ -297,6 +301,7 @@ class Qwen3GatedDeltaNet(nn.Module): reduce_results=False, tp_rank=self.attn_tp_rank, tp_size=self.attn_tp_size, + prefix=add_prefix("out_proj", prefix), ) def fix_query_key_value_ordering(self, mixed_qkvz, mixed_ba): @@ -452,7 +457,7 @@ class Qwen3GatedDeltaNet(nn.Module): z = z.reshape(-1, z.shape[-1]) # Add padding for DP-Attn - if is_dp_attention_enabled(): + if core_attn_out.shape != z.shape: core_attn_out_pad = torch.zeros_like(z) core_attn_out_pad[: core_attn_out.shape[0], :] = core_attn_out core_attn_out = core_attn_out_pad @@ -478,7 +483,7 @@ class Qwen3HybridLinearDecoderLayer(nn.Module): super().__init__() self.config = config self.linear_attn = Qwen3GatedDeltaNet( - config, layer_id, quant_config, alt_stream + config, layer_id, quant_config, alt_stream, prefix ) # Qwen3Next all layers are sparse and have no nextn now @@ -501,7 +506,7 @@ class Qwen3HybridLinearDecoderLayer(nn.Module): config=config, quant_config=quant_config, alt_stream=alt_stream, - prefix=add_prefix("mlp", prefix), + prefix=add_prefix("mlp", prefix.replace(".linear_attn", "")), ) else: self.mlp = Qwen2MoeMLP( @@ -509,6 +514,7 @@ class Qwen3HybridLinearDecoderLayer(nn.Module): intermediate_size=config.intermediate_size, hidden_act=config.hidden_act, quant_config=quant_config, + prefix=add_prefix("mlp", prefix.replace(".linear_attn", "")), ) self.input_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = GemmaRMSNorm( @@ -617,6 +623,7 @@ class Qwen3HybridAttentionDecoderLayer(nn.Module): quant_config=quant_config, tp_rank=self.attn_tp_rank, tp_size=self.attn_tp_size, + prefix=add_prefix("qkv_proj", prefix), ) self.o_proj = RowParallelLinear( @@ -627,6 +634,7 @@ class Qwen3HybridAttentionDecoderLayer(nn.Module): reduce_results=False, tp_rank=self.attn_tp_rank, tp_size=self.attn_tp_size, + prefix=add_prefix("o_proj", prefix), ) self.attn = RadixAttention( @@ -657,7 +665,7 @@ class Qwen3HybridAttentionDecoderLayer(nn.Module): config=config, quant_config=quant_config, alt_stream=alt_stream, - prefix=add_prefix("mlp", prefix), + prefix=add_prefix("mlp", prefix.replace(".self_attn", "")), ) else: self.mlp = Qwen2MoeMLP( @@ -665,6 +673,7 @@ class Qwen3HybridAttentionDecoderLayer(nn.Module): intermediate_size=config.intermediate_size, hidden_act=config.hidden_act, quant_config=quant_config, + prefix=add_prefix("mlp", prefix.replace(".self_attn", "")), ) self.input_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = GemmaRMSNorm( @@ -800,6 +809,10 @@ class Qwen3NextModel(nn.Module): def get_layer(idx: int, prefix: str): layer_class = ALL_DECODER_LAYER_TYPES[config.layers_block_type[idx]] + if config.layers_block_type[idx] == "attention": + prefix = add_prefix("self_attn", prefix) + else: + prefix = add_prefix("linear_attn", prefix) return layer_class( config, idx,