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 3, 04:15 UTC
[Bugfix] Fix Qwen3.5 GatedDeltaNet in_proj_ba Marlin failure at TP>=2 (#36329)
Signed-off-by: Adi McM Sonus Flow <biuro@sonusflow.pl> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
sonusflow committed
May 21, 2026 at 06:22 UTC
e45df8c3f77572d03f638feded5b5efbccdbcc05
1 file changed
+34
-1
vllm/model_executor/layers/mamba/gdn_linear_attn.py
+34
-1
@@ -49,6 +49,9 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
49
causal_conv1d_update,
50
)
51
from vllm.model_executor.layers.quantization import QuantizationConfig
52
+from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig
53
+from vllm.model_executor.layers.quantization.awq_marlin import AWQMarlinConfig
54
+from vllm.model_executor.layers.quantization.inc import INCConfig
55
from vllm.model_executor.model_loader.weight_utils import (
56
sharded_weight_loader,
57
)
@@ -365,6 +368,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
368
quant_config=quant_config,
369
prefix=f"{prefix}.in_proj_ba",
370
)
371
+ self.disable_tp_for_ba_proj = self.maybe_disable_tp(quant_config)
372
373
query_key_settings = (self.key_dim, 0, False)
374
value_settings = (self.value_dim, 0, False)
@@ -479,8 +483,37 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
483
bias=False,
484
quant_config=quant_config,
485
prefix=prefix,
486
+ disable_tp=self.maybe_disable_tp(quant_config),
487
)
488
489
+ def maybe_disable_tp(self, quant_config: QuantizationConfig | None) -> bool:
490
+ """Whether to replicate ba_proj instead of TP-sharding it.
491
+
492
+ Marlin requires output_size_per_partition >= MIN_THREAD_N=64, which
493
+ the Qwen3.5 non-interleaved [num_v_heads]*2 layout violates at TP>=2
494
+ (e.g. num_v_heads=64, TP=4 -> 16). Replicating the projection keeps
495
+ each rank above the Marlin threshold; forward() then slices b/a to
496
+ the local TP partition. Qwen3-Next's interleaved [num_v_heads*2]
497
+ layout is unaffected and stays TP-sharded.
498
+
499
+ See https://github.com/vllm-project/vllm/issues/35924
500
+ """
501
+ return (
502
+ current_platform.is_cuda()
503
+ and not self.gqa_interleaved_layout
504
+ and isinstance(quant_config, (AWQMarlinConfig, AutoGPTQConfig, INCConfig))
505
+ )
506
+
507
+ def split_ba(self, ba: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
508
+ b, a = ba.chunk(2, dim=-1)
509
+ if self.disable_tp_for_ba_proj and self.tp_size > 1:
510
+ # ba_proj is replicated for Marlin; slice b/a to local TP rank.
511
+ ba_chunk = self.num_v_heads // self.tp_size
512
+ ba_start = self.tp_rank * ba_chunk
513
+ b = b[:, ba_start : ba_start + ba_chunk]
514
+ a = a[:, ba_start : ba_start + ba_chunk]
515
+ return b, a
516
+
517
def fix_query_key_value_ordering(
518
self,
519
mixed_qkvz: torch.Tensor,
@@ -778,7 +811,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
811
z_size = self.value_dim // self.tp_size
812
mixed_qkv, z = mixed_qkvz.split([qkv_size, z_size], dim=-1)
813
z = z.reshape(z.size(0), -1, self.head_v_dim)
781
- b, a = ba.chunk(2, dim=-1)
814
+ b, a = self.split_ba(ba)
815
b = b.contiguous()
816
a = a.contiguous()
817