A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 00:15 UTC
[Bugfix] Fix Qwen3-Next in_proj_ba weight sharding with TP > 1 (#36242)
Signed-off-by: AjAnubolu <anuboluajay@gmail.com>
Ajay Anubolu committed
Mar 9, 2026 at 19:16 UTC
4e95ec111cd179f2ab0f6931bf57663f828a51ec
2 files changed
+45
-6
vllm/model_executor/models/qwen3_5.py
+18
@@ -145,6 +145,24 @@ class Qwen3_5GatedDeltaNet(Qwen3NextGatedDeltaNet):
145
prefix=prefix,
146
)
147
148
+ def create_ba_proj(
149
+ self,
150
+ hidden_size: int,
151
+ num_v_heads: int,
152
+ quant_config: QuantizationConfig | None,
153
+ prefix: str,
154
+ ) -> MergedColumnParallelLinear:
155
+ # Qwen3.5 has separate in_proj_b and in_proj_a weights in the
156
+ # checkpoint, which are loaded into the fused in_proj_ba parameter
157
+ # via stacked_params_mapping with shard_id 0 and 1 respectively.
158
+ return MergedColumnParallelLinear(
159
+ input_size=hidden_size,
160
+ output_sizes=[num_v_heads] * 2,
161
+ bias=False,
162
+ quant_config=quant_config,
163
+ prefix=prefix,
164
+ )
165
+
166
def forward(
167
self,
168
hidden_states: torch.Tensor,
vllm/model_executor/models/qwen3_next.py
+27
-6
@@ -412,12 +412,11 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
412
prefix=f"{prefix}.in_proj_qkvz",
413
)
414
# ba_proj doesn't support blockwise fp8 quantization.
415
- # # in_proj_ba is defined as MergedColumnParallelLinear for
416
- # compatibility with Qwen3_5.
417
- self.in_proj_ba = MergedColumnParallelLinear(
418
- input_size=self.hidden_size,
419
- output_sizes=[self.num_v_heads] * 2,
420
- bias=False,
415
+ # Qwen3-Next and Qwen3.5 have different in_proj_ba checkpoint
416
+ # layouts, so we use a factory method to create the projection.
417
+ self.in_proj_ba = self.create_ba_proj(
418
+ hidden_size=self.hidden_size,
419
+ num_v_heads=self.num_v_heads,
420
quant_config=quant_config,
421
prefix=f"{prefix}.in_proj_ba",
422
)
@@ -497,6 +496,28 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
496
prefix=prefix,
497
)
498
499
+ def create_ba_proj(
500
+ self,
501
+ hidden_size: int,
502
+ num_v_heads: int,
503
+ quant_config: QuantizationConfig | None,
504
+ prefix: str,
505
+ ) -> MergedColumnParallelLinear:
506
+ # Qwen3-Next stores in_proj_ba as a single fused weight with an
507
+ # interleaved GQA layout: [b_g0, a_g0, b_g1, a_g1, ...] where
508
+ # each group corresponds to a key-head group. We must use a single
509
+ # output shard so that ColumnParallel sharding preserves this
510
+ # interleaved structure across TP ranks.
511
+ # Qwen3.5 overrides this to use [num_v_heads, num_v_heads] since
512
+ # its checkpoint has separate in_proj_b and in_proj_a weights.
513
+ return MergedColumnParallelLinear(
514
+ input_size=hidden_size,
515
+ output_sizes=[num_v_heads * 2],
516
+ bias=False,
517
+ quant_config=quant_config,
518
+ prefix=prefix,
519
+ )
520
+
521
def fix_query_key_value_ordering(
522
self,
523
mixed_qkvz: torch.Tensor,