[https://nvbugs/6422332][fix] Keep SSM cache in weights dtype when ma…

tensorrt_llm/_torch/pyexecutor/config_utils.py

@@ -52,25 +52,32 @@ def resolve_hf_torch_dtype(config):
     return None
 
 
-def resolve_mamba_ssm_cache_dtype(config):
-    """Return the dtype to use for hybrid Mamba/SSM cache allocations.
-
-    Qwen3.5-style configs may store this field on the top-level config or the
-    nested text_config, and may call it either mamba_ssm_cache_dtype or
-    mamba_ssm_dtype. This helper centralizes that lookup so cache creation does
-    not fail later with a missing dtype. An "auto" value in any field is
-    treated the same as missing.
+def resolve_ssm_cache_dtype(config):
+    """Return the dtype to use for hybrid SSM-style state cache allocations.
+
+    Covers all recurrent-state hybrids (Mamba SSM, Qwen3.5 GDN linear
+    attention); the "mamba" in the config field name is historical.
+    Qwen3.5-style configs may store the field on the top-level config or the
+    nested text_config. An "auto" value is treated the same as missing.
+
+    Only the explicit mamba_ssm_cache_dtype field is honored here. The
+    checkpoint's mamba_ssm_dtype field expresses the SSM *compute* dtype
+    intent, not the cache dtype: honoring it for cache allocation silently
+    switches Qwen3.5 GDN state caches to fp32, which disables the FlashInfer
+    bf16-state decode kernel and doubles state memory traffic (~20% serving
+    throughput loss). Users can opt in explicitly via
+    kv_cache_config.mamba_ssm_cache_dtype.
     """
     configs = [config]
     text_config = getattr(config, "text_config", None)
     if text_config is not None:
         configs.append(text_config)
 
     for candidate_config in configs:
-        for attr in ("mamba_ssm_cache_dtype", "mamba_ssm_dtype"):
-            coerced = _coerce_torch_dtype(getattr(candidate_config, attr, None))
-            if coerced is not None:
-                return coerced
+        coerced = _coerce_torch_dtype(
+            getattr(candidate_config, "mamba_ssm_cache_dtype", None))
+        if coerced is not None:
+            return coerced
     return None
 
 
@@ -307,7 +314,7 @@ def extract_mamba_kv_cache_params(
         mamba_ssm_cache_dtype = _coerce_torch_dtype(
             quant_config.mamba_ssm_cache_dtype)
     if mamba_ssm_cache_dtype is None:
-        mamba_ssm_cache_dtype = (resolve_mamba_ssm_cache_dtype(config)
+        mamba_ssm_cache_dtype = (resolve_ssm_cache_dtype(config)
                                  or resolve_hf_torch_dtype(config)
                                  or torch.bfloat16)
 

tensorrt_llm/_torch/pyexecutor/model_loader.py

@@ -35,7 +35,7 @@
     MoeLoadBalancer, maybe_create_moe_load_balancer)
 from ..virtual_memory import RestoreMode
 from ..virtual_memory import scope as virtual_memory_scope
-from .config_utils import resolve_hf_torch_dtype, resolve_mamba_ssm_cache_dtype
+from .config_utils import resolve_hf_torch_dtype, resolve_ssm_cache_dtype
 
 _KV_CACHE_MAP = {
     "fp8": QuantAlgo.FP8.value,
@@ -52,7 +52,7 @@ def validate_and_set_mamba_ssm_cache_dtype(
         mamba_ssm_philox_rounds: int = 10) -> None:
     if mamba_ssm_cache_dtype == "auto":
         mamba_ssm_cache_dtype = (
-            resolve_mamba_ssm_cache_dtype(config.pretrained_config)
+            resolve_ssm_cache_dtype(config.pretrained_config)
             or resolve_hf_torch_dtype(config.pretrained_config)
             or config.torch_dtype)
     else:

tests/unittest/_torch/modeling/test_modeling_qwen3_5_vl_moe.py

@@ -125,15 +125,24 @@ def test_qwen35_moe_vl_resolves_mamba_ssm_cache_dtype(
     config = load_pretrained_config(str(_write_qwen35_moe_vl_config(tmp_path)))
     model_config = ModelConfig(pretrained_config=config)
 
+    # "auto" keeps the SSM cache in the model weights dtype for performance:
+    # the checkpoint's mamba_ssm_dtype=float32 expresses SSM compute intent,
+    # and honoring it for cache allocation disables the FlashInfer bf16-state
+    # GDN decode kernel and doubles state memory traffic.
     validate_and_set_mamba_ssm_cache_dtype(model_config, "auto")
-    assert model_config.quant_config.mamba_ssm_cache_dtype is torch.float32
+    assert model_config.quant_config.mamba_ssm_cache_dtype is torch.bfloat16
 
     mamba_params = extract_mamba_kv_cache_params(
         config.text_config,
         quant_config=model_config.quant_config,
     )
     assert mamba_params.dtype is torch.bfloat16
-    assert mamba_params.mamba_ssm_cache_dtype is torch.float32
+    assert mamba_params.mamba_ssm_cache_dtype is torch.bfloat16
+
+    # Explicit opt-in honors the checkpoint's fp32 SSM state intent.
+    opt_in_config = ModelConfig(pretrained_config=config)
+    validate_and_set_mamba_ssm_cache_dtype(opt_in_config, "float32")
+    assert opt_in_config.quant_config.mamba_ssm_cache_dtype is torch.float32
 
 
 def test_qwen35_moe_vl_resolves_model_and_mapper(tmp_path: Path) -> None:

tests/unittest/_torch/modeling/test_qwen_image_bench_modeling.py

@@ -117,7 +117,9 @@ def test_qwen_image_bench_config_uses_internal_arch_and_normalizes_text(tmp_path
         "linear_attention",
         "full_attention",
     ]
-    assert extract_mamba_kv_cache_params(config.text_config).mamba_ssm_cache_dtype is torch.float32
+    # The declared mamba_ssm_dtype=float32 (SSM compute intent) is deliberately
+    # not used for cache allocation; the cache stays in the weights dtype.
+    assert extract_mamba_kv_cache_params(config.text_config).mamba_ssm_cache_dtype is torch.bfloat16
 
 
 def test_qwen3_5_conditional_text_config_does_not_use_image_bench_arch(tmp_path):