[Bugfix][Mamba] Reject stochastic rounding with Triton speculation

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"):

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