From 693c41437559ca24dd895a4da09b70cc76b5900a Mon Sep 17 00:00:00 2001 From: nv-guomingz <137257613+nv-guomingz@users.noreply.github.com> Date: Tue, 7 Jul 2026 18:42:58 +0800 Subject: [PATCH] [https://nvbugs/6422332][fix] Keep SSM cache in weights dtype when mamba_ssm_cache_dtype is auto PR #14599 made the 'auto' resolution of mamba_ssm_cache_dtype also honor the checkpoint's mamba_ssm_dtype field. Qwen3.5 checkpoints declare mamba_ssm_dtype=float32 (SSM compute intent), so the GDN state cache silently flipped from bf16 to fp32. An fp32 state cache disables the FlashInfer bf16-state GDN decode kernel (falls back to the slower Triton recurrent kernel) and doubles state memory traffic and footprint, regressing qwen3_5_397b_fp4_dep4_1k1k perf-sanity output throughput by ~19.5% (25435.9 -> 20479.4 tok/s) on B200. Restore the pre-#14599 behavior: 'auto' resolves from the explicit mamba_ssm_cache_dtype field (top-level or text_config) and otherwise falls back to the model weights dtype. The checkpoint's mamba_ssm_dtype is no longer used for cache allocation; users can still opt in explicitly via kv_cache_config.mamba_ssm_cache_dtype. Signed-off-by: nv-guomingz <137257613+nv-guomingz@users.noreply.github.com> --- .../_torch/pyexecutor/config_utils.py | 33 +++++++++++-------- .../_torch/pyexecutor/model_loader.py | 4 +-- .../modeling/test_modeling_qwen3_5_vl_moe.py | 13 ++++++-- .../test_qwen_image_bench_modeling.py | 4 ++- 4 files changed, 36 insertions(+), 18 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py index 806c757ee40c..9da2f49e284c 100644 --- a/tensorrt_llm/_torch/pyexecutor/config_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py @@ -52,14 +52,21 @@ 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) @@ -67,10 +74,10 @@ def resolve_mamba_ssm_cache_dtype(config): 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) diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 93eb3436dc3f..27b66985cb2b 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/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: diff --git a/tests/unittest/_torch/modeling/test_modeling_qwen3_5_vl_moe.py b/tests/unittest/_torch/modeling/test_modeling_qwen3_5_vl_moe.py index dae1f1092290..e715c53a3ac2 100644 --- a/tests/unittest/_torch/modeling/test_modeling_qwen3_5_vl_moe.py +++ b/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: diff --git a/tests/unittest/_torch/modeling/test_qwen_image_bench_modeling.py b/tests/unittest/_torch/modeling/test_qwen_image_bench_modeling.py index 5ce22ec104fb..6a3962710cac 100644 --- a/tests/unittest/_torch/modeling/test_qwen_image_bench_modeling.py +++ b/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):