Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 11 additions & 2 deletions tests/unittest/_torch/modeling/test_modeling_qwen3_5_vl.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
Loading