[Bugfix][GLM-5.3-Flash] Honor --mamba-ssm-cache-dtype for the KDA recurrent state (#59147)

Signed-off-by: Jared Wen <jaredwen@inferact.ai>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

tests/models/glm5next/test_kda_state_dtype.py

@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""``--mamba-ssm-cache-dtype`` must reach GLM-5.3-Flash's KDA recurrent state
+(it used to be dropped, leaving the state float32 unconditionally)."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.models.glm5next.common.kda import Glm5NextLinearAttention
+from vllm.models.glm5next.common.model import (
+    Glm5NextForCausalLM,
+    Glm5NextForConditionalGeneration,
+)
+
+
+@pytest.mark.parametrize(
+    ("ssm_dtype", "expected"),
+    [("auto", torch.float32), ("float32", torch.float32), ("bfloat16", torch.bfloat16)],
+)
+def test_kda_recurrent_state_dtype_follows_ssm_cache_dtype(ssm_dtype, expected):
+    model_config = SimpleNamespace(dtype=torch.bfloat16)
+    cache_config = SimpleNamespace(
+        mamba_cache_dtype="auto",
+        mamba_ssm_cache_dtype=ssm_dtype,
+        use_kda_recoverssm=False,
+    )
+    vllm_config = SimpleNamespace(model_config=model_config, cache_config=cache_config)
+    want = (torch.bfloat16, expected)
+
+    assert Glm5NextForCausalLM.get_mamba_state_dtype_from_config(vllm_config) == want
+    assert (
+        Glm5NextForConditionalGeneration.get_mamba_state_dtype_from_config(vllm_config)
+        == want
+    )
+    # The layer's own dtype must agree with the model-level one, since the
+    # KV-cache spec and the kernels read them separately.
+    layer = Glm5NextLinearAttention.__new__(Glm5NextLinearAttention)
+    layer.model_config, layer.cache_config = model_config, cache_config
+    assert layer.get_state_dtype() == want

vllm/models/glm5next/common/kda.py

@@ -181,7 +181,9 @@ class Glm5NextLinearAttention(GatedDeltaNetAttention):
         if self.model_config is None or self.cache_config is None:
             raise ValueError("model_config and cache_config must be set")
         return MambaStateDtypeCalculator.kda_state_dtype(
-            self.model_config.dtype, self.cache_config.mamba_cache_dtype
+            self.model_config.dtype,
+            self.cache_config.mamba_cache_dtype,
+            self.cache_config.mamba_ssm_cache_dtype,
         )
 
     def get_state_shape(

vllm/models/glm5next/common/model.py

@@ -1092,7 +1092,9 @@ class Glm5NextForCausalLM(
         vllm_config: "VllmConfig",
     ) -> tuple[torch.dtype, torch.dtype]:
         return MambaStateDtypeCalculator.kda_state_dtype(
-            vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype
+            vllm_config.model_config.dtype,
+            vllm_config.cache_config.mamba_cache_dtype,
+            vllm_config.cache_config.mamba_ssm_cache_dtype,
         )
 
     @classmethod