diff --git a/tests/models/glm5next/test_kda_state_dtype.py b/tests/models/glm5next/test_kda_state_dtype.py new file mode 100644 index 0000000000..b112e47cd5 --- /dev/null +++ b/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 diff --git a/vllm/models/glm5next/common/kda.py b/vllm/models/glm5next/common/kda.py index 141a8cc7cb..68d405b40c 100644 --- a/vllm/models/glm5next/common/kda.py +++ b/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( diff --git a/vllm/models/glm5next/common/model.py b/vllm/models/glm5next/common/model.py index e35887ca7c..1e58166054 100644 --- a/vllm/models/glm5next/common/model.py +++ b/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