diff --git a/tests/unittest/_torch/modeling/test_modeling_qwen3_5_vl.py b/tests/unittest/_torch/modeling/test_modeling_qwen3_5_vl.py index fe097dabd79a..5c13c3a73853 100644 --- a/tests/unittest/_torch/modeling/test_modeling_qwen3_5_vl.py +++ b/tests/unittest/_torch/modeling/test_modeling_qwen3_5_vl.py @@ -143,15 +143,24 @@ def test_qwen35_dense_vl_resolves_mamba_ssm_cache_dtype( config = load_pretrained_config(str(_write_qwen35_dense_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_dense_vl_resolves_model_and_mapper(tmp_path: Path) -> None: