diff --git a/tests/test_config.py b/tests/test_config.py index a2e75b10725d..c27b57e66a51 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -4,6 +4,7 @@ import json import logging import os +from contextlib import nullcontext from dataclasses import MISSING, Field, asdict, dataclass, field from datetime import timedelta from pathlib import Path @@ -42,7 +43,7 @@ from vllm.config.compilation import CompilationMode, CUDAGraphMode, PassConfig from vllm.config.kernel import IrOpPriorityConfig from vllm.config.load import LoadConfig -from vllm.config.mamba import MambaBackendEnum +from vllm.config.mamba import MambaBackendEnum, MambaConfig from vllm.config.speculative import _validate_qwen3_omni_dspark from vllm.config.utils import get_field from vllm.config.vllm import OPTIMIZATION_LEVEL_TO_CONFIG, OptimizationLevel @@ -283,6 +284,39 @@ def test_kda_recoverssm_derivation_is_revalidated(): VllmConfig.validate_mamba_cached_kernel(config) +@pytest.mark.parametrize( + ("backend", "enable_sr", "with_speculation", "reject"), + [ + ("triton", True, True, True), + ("flashinfer", True, True, False), + ("triton", False, True, False), + ("triton", True, False, False), + ], +) +def test_mamba_stochastic_rounding_with_speculation( + monkeypatch, backend, enable_sr, with_speculation, reject +): + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + monkeypatch.setattr(current_platform, "is_device_capability_family", lambda _: True) + with ( + pytest.raises(ValueError, match="Use `--mamba-backend flashinfer`") + if reject + else nullcontext() + ): + VllmConfig( + device_config=DeviceConfig(device="cpu"), + cache_config=CacheConfig(mamba_ssm_cache_dtype="float16"), + mamba_config=MambaConfig( + backend=backend, enable_stochastic_rounding=enable_sr + ), + speculative_config=( + SpeculativeConfig(method="ngram", num_speculative_tokens=1) + if with_speculation + else None + ), + ) + + def test_mamba_cache_mode_all_is_rejected(): """The removed 'all' mode must fail validation instead of being ignored.""" with pytest.raises(ValidationError, match="mamba_cache_mode"): diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 281666a63392..2e140a9c586f 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1490,6 +1490,17 @@ def __post_init__(self): "`--enable-mamba-cache-stochastic-rounding`." ) + if ( + self.mamba_config.enable_stochastic_rounding + and self.mamba_config.backend == MambaBackendEnum.TRITON + and self.speculative_config is not None + ): + raise ValueError( + "Triton Mamba stochastic rounding is not supported with " + "speculative decoding. Use `--mamba-backend flashinfer`, " + "or disable `--enable-mamba-cache-stochastic-rounding`." + ) + if self.quant_config is None and self.model_config is not None: self.quant_config = VllmConfig._get_quantization_config( self.model_config, self.load_config