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 8, 08:45 UTC
[Bugfix] Fix mamba state dtype setting for Qwen3-Next and Qwen3.5 (#34489)
Signed-off-by: Roger Wang <hey@rogerw.io>
Roger Wang committed
Feb 12, 2026 at 22:48 UTC
eea3024f43e06ea4e037ec86464dcc249d0c0b44
4 files changed
+42
-6
vllm/model_executor/layers/mamba/mamba_utils.py
+4
-2
@@ -80,9 +80,11 @@ class MambaStateDtypeCalculator:
80
cls,
81
model_dtype: ModelDType | torch.dtype,
82
mamba_cache_dtype: MambaDType,
83
+ mamba_ssm_cache_dtype: MambaDType = "auto",
84
) -> tuple[torch.dtype, torch.dtype]:
84
- state_dtype = get_kv_cache_torch_dtype(mamba_cache_dtype, model_dtype)
85
- return (state_dtype, state_dtype)
85
+ return cls._mamba_state_dtype(
86
+ model_dtype, mamba_cache_dtype, mamba_ssm_cache_dtype
87
+ )
88
89
@classmethod
90
def kda_state_dtype(
vllm/model_executor/models/config.py
+29
@@ -582,6 +582,33 @@ class NemotronHForCausalLMConfig(VerifyAndUpdateConfig):
582
cache_config.mamba_ssm_cache_dtype = mamba_ssm_cache_dtype
583
584
585
+class Qwen3_5ForConditionalGenerationConfig(VerifyAndUpdateConfig):
586
+ @staticmethod
587
+ def verify_and_update_config(vllm_config: "VllmConfig") -> None:
588
+ """Update mamba_ssm_cache_dtype for Qwen3.5 models when set to 'auto'
589
+ (or not explicitly set), to the value specified in the HF config's
590
+ mamba_ssm_dtype field. Warn if the user explicitly overrides it to a
591
+ different value.
592
+ """
593
+ cache_config = vllm_config.cache_config
594
+ hf_text_config = vllm_config.model_config.hf_text_config
595
+ mamba_ssm_dtype = getattr(hf_text_config, "mamba_ssm_dtype", None)
596
+ if cache_config.mamba_ssm_cache_dtype == "auto":
597
+ if mamba_ssm_dtype is not None:
598
+ cache_config.mamba_ssm_cache_dtype = mamba_ssm_dtype
599
+ elif (
600
+ mamba_ssm_dtype is not None
601
+ and cache_config.mamba_ssm_cache_dtype != mamba_ssm_dtype
602
+ ):
603
+ logger.warning(
604
+ "Qwen3.5 model specifies mamba_ssm_dtype='%s' in its config, "
605
+ "but --mamba-ssm-cache-dtype='%s' was passed. "
606
+ "Using the user-specified value.",
607
+ mamba_ssm_dtype,
608
+ cache_config.mamba_ssm_cache_dtype,
609
+ )
610
+
611
+
612
class VoyageQwen3BidirectionalEmbedModelConfig(VerifyAndUpdateConfig):
613
@staticmethod
614
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
@@ -611,5 +638,7 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
638
"DeepseekV32ForCausalLM": DeepseekV32ForCausalLM,
639
"NemotronHForCausalLM": NemotronHForCausalLMConfig,
640
"NemotronHPuzzleForCausalLM": NemotronHForCausalLMConfig,
641
+ "Qwen3_5ForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig,
642
+ "Qwen3_5MoeForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig,
643
"VoyageQwen3BidirectionalEmbedModel": VoyageQwen3BidirectionalEmbedModelConfig,
644
}
vllm/model_executor/models/qwen3_5.py
+3
-2
@@ -870,9 +870,10 @@ class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration, IsHybrid)
870
cls,
871
vllm_config: "VllmConfig",
872
) -> tuple[torch.dtype, torch.dtype]:
873
- mamba_ssm_dtype = vllm_config.model_config.hf_text_config.mamba_ssm_dtype
873
return MambaStateDtypeCalculator.gated_delta_net_state_dtype(
875
- vllm_config.model_config.dtype, mamba_ssm_dtype
874
+ vllm_config.model_config.dtype,
875
+ vllm_config.cache_config.mamba_cache_dtype,
876
+ vllm_config.cache_config.mamba_ssm_cache_dtype,
877
)
878
879
@classmethod
vllm/model_executor/models/qwen3_next.py
+6
-2
@@ -341,7 +341,9 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
341
342
def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]:
343
return MambaStateDtypeCalculator.gated_delta_net_state_dtype(
344
- self.model_config.dtype, self.cache_config.mamba_cache_dtype
344
+ self.model_config.dtype,
345
+ self.cache_config.mamba_cache_dtype,
346
+ self.cache_config.mamba_ssm_cache_dtype,
347
)
348
349
def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]:
@@ -1372,7 +1374,9 @@ class Qwen3NextForCausalLM(
1374
vllm_config: "VllmConfig",
1375
) -> tuple[torch.dtype, torch.dtype]:
1376
return MambaStateDtypeCalculator.gated_delta_net_state_dtype(
1375
- vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype
1377
+ vllm_config.model_config.dtype,
1378
+ vllm_config.cache_config.mamba_cache_dtype,
1379
+ vllm_config.cache_config.mamba_ssm_cache_dtype,
1380
)
1381
1382
@classmethod