[NPU] Adapt qwen3-next W8A8 on NPU (#16164)

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,