Add stochastic rounding for FP16 Mamba SSM cache (#26929)

Signed-off-by: Daniel Afrimi <dafrimi@login-lyris01.lyris.clusters.nvidia.com>
Co-authored-by: Daniel Afrimi <dafrimi@login-lyris01.lyris.clusters.nvidia.com>

docs_new/cookbook/autoregressive/NVIDIA/Nemotron3-Ultra.mdx

@@ -74,6 +74,10 @@ The generator only emits a runnable command for combinations that NVIDIA / SGLan
 
     The SSM state dtype defaults to the model config value. Set `--mamba-ssm-dtype float16` to store the Mamba states in FP16, which reduces mamba cache memory without significant accuracy loss.
 
+- **Mamba SSM stochastic rounding**:
+
+    When storing the Mamba states in FP16, add `--enable-mamba-cache-stochastic-rounding` to round SSM cache writes stochastically and reduce accumulation bias. It requires `--mamba-ssm-dtype float16` and CUDA; with the default `--mamba-backend triton` it additionally requires SM100. Use `--mamba-cache-philox-rounds` to control the number of Philox rounds (`0` uses the backend default).
+
 - **TP support**:
 
     To set tp size, use `--tp <4|8|16>`. Recommended pairings:

docs_new/docs/advanced_features/server_arguments.mdx

@@ -1911,6 +1911,18 @@ Please consult the documentation below and [server_args.py](https://github.com/s
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>The data type of the SSM states in mamba cache. If not set, read from the model config.</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Auto (from model config)</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>float32</code>, <code>bfloat16</code>, <code>float16</code></td>
+    </tr>
+        <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--enable-mamba-cache-stochastic-rounding`</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Enable stochastic rounding when writing FP16 Mamba SSM cache states. Requires <code>--mamba-ssm-dtype float16</code> and CUDA. With <code>--mamba-backend triton</code>, requires SM100.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`False`</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Type: bool</td>
+    </tr>
+        <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--mamba-cache-philox-rounds`</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Number of Philox rounds to use for stochastic rounding of FP16 Mamba SSM cache writes. Triton uses the Triton default when set to 0; FlashInfer uses 10 rounds when set to 0.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`0`</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Type: int</td>
     </tr>
         <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--mamba-full-memory-ratio`</td>

docs_new/src/snippets/autoregressive/nemotron3-ultra-deployment.jsx

@@ -172,6 +172,18 @@ export const Nemotron3UltraDeployment = () => {
       ],
       commandRule: (value) => value === 'float16' ? '--mamba-ssm-dtype float16' : null
     },
+    mambastochasticrounding: {
+      name: 'mambastochasticrounding',
+      title: 'Mamba Stochastic Rounding',
+      items: [
+        { id: 'disabled', label: 'Disabled', default: true  },
+        { id: 'enabled',  label: 'Enabled',  subtitle: 'FP16 SSM' }
+      ],
+      commandRule: (value, state) =>
+        value === 'enabled' && state.mambassmdtype === 'float16'
+          ? '--enable-mamba-cache-stochastic-rounding'
+          : null
+    },
     thinking: {
       name: 'thinking',
       title: 'Reasoning Parser',

python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py

@@ -32,6 +32,21 @@ else:
         return dt
 
 
+@triton.jit
+def convert_rs_fp16x2(x: tl.tensor, rand: tl.tensor) -> tl.tensor:
+    y = tl.inline_asm_elementwise(
+        asm="""{
+cvt.rs.f16x2.f32 $0, $2, $1, $3;
+}""",
+        constraints="=r,r,r,r,r",
+        args=(x, rand),
+        dtype=tl.float16,
+        is_pure=True,
+        pack=2,
+    )
+    return y
+
+
 @triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None})
 @triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None})
 @triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None})
@@ -85,6 +100,7 @@ def _selective_scan_update_kernel(
     cache_steps,
     retrieve_parent_token_ptr,
     intermediate_state_indices_ptr,
+    rand_seed_ptr,
     # Matrix dimensions
     batch,
     T,
@@ -143,6 +159,8 @@ def _selective_scan_update_kernel(
     HAS_EAGLE_TREE_CUSTOM_ATTN_MASK: tl.constexpr,
     HAS_INTERMEDIATE_STATE_INDICES: tl.constexpr,
     BLOCK_SIZE_DSTATE: tl.constexpr,
+    USE_RS_ROUNDING: tl.constexpr,
+    PHILOX_ROUNDS: tl.constexpr,
     USE_GDC: tl.constexpr = False,
 ):
     if USE_GDC:
@@ -300,7 +318,31 @@ def _selective_scan_update_kernel(
             z_ptr += stride_z_T
 
     if not DISABLE_STATE_UPDATE:
-        tl.store(state_ptrs, state.to(state_ptrs.dtype.element_ty), mask=mask)
+        if USE_RS_ROUNDING:
+            rand_seed = tl.load(rand_seed_ptr)
+            if HAS_STATE_BATCH_INDICES:
+                rand_offsets = (
+                    state_batch_idx * stride_state_batch + pid_h * stride_state_head
+                )
+            else:
+                rand_offsets = pid_b * stride_state_batch + pid_h * stride_state_head
+            rand_offsets += (
+                offs_m[:, None] * stride_state_dim
+                + offs_n[None, :] * stride_state_dstate
+            )
+            if PHILOX_ROUNDS > 0:
+                rand = tl.randint(rand_seed, rand_offsets, PHILOX_ROUNDS)
+            else:
+                rand = tl.randint(rand_seed, rand_offsets)
+            state_to_store = convert_rs_fp16x2(state, rand)
+            tl.static_assert(state_to_store.dtype == tl.float16, "state must be fp16")
+            tl.static_assert(
+                state_ptrs.dtype.element_ty == tl.float16,
+                "Stochastic rounding only supports fp16 state stores",
+            )
+        else:
+            state_to_store = state.to(state_ptrs.dtype.element_ty)
+        tl.store(state_ptrs, state_to_store, mask=mask)
 
     if USE_GDC:
         tl.extra.cuda.gdc_launch_dependents()
@@ -325,6 +367,8 @@ def selective_state_update(
     cache_steps=None,
     retrieve_parent_token=None,
     intermediate_state_indices=None,
+    enable_stochastic_rounding=False,
+    cache_philox_rounds=0,
 ):
     """
     Argument:
@@ -351,7 +395,17 @@ def selective_state_update(
         retrieve_parent_token: (batch, T) tensor of parent token indices for EAGLE tree attention
         intermediate_state_indices: (batch,) tensor of indices for intermediate_states_buffer operations.
             If provided, uses these indices instead of state_batch_indices for the buffer.
+        enable_stochastic_rounding: Whether to stochastically round final FP16 SSM cache writes.
+        cache_philox_rounds: Number of Philox rounds to use when stochastic rounding is enabled.
     """
+    if cache_philox_rounds < 0:
+        raise ValueError("cache_philox_rounds must be non-negative.")
+    if enable_stochastic_rounding and state.dtype != torch.float16:
+        raise ValueError(
+            "Stochastic rounding for the Mamba SSM cache requires state dtype "
+            f"torch.float16, got {state.dtype}."
+        )
+
     if state.dim() == 3:
         state = state.unsqueeze(1)
     if x.dim() == 2:
@@ -435,6 +489,11 @@ def selective_state_update(
         if retrieve_parent_token is not None
         else (0, 0)
     )
+    rand_seed = (
+        torch.randint(0, 2**32, (1,), device=state.device)
+        if enable_stochastic_rounding
+        else None
+    )
 
     pdl_kwargs = {"USE_GDC": True, "launch_pdl": True} if is_arch_support_pdl() else {}
 
@@ -456,6 +515,7 @@ def selective_state_update(
             cache_steps if cache_steps is not None else 0,
             retrieve_parent_token,
             intermediate_state_indices,
+            rand_seed,
             batch,
             T,
             nheads,
@@ -501,6 +561,8 @@ def selective_state_update(
             tie_hdim,
             BLOCK_SIZE_M,
             DISABLE_STATE_UPDATE=disable_state_update,
+            USE_RS_ROUNDING=enable_stochastic_rounding,
+            PHILOX_ROUNDS=cache_philox_rounds,
             num_warps=num_warps,
             **pdl_kwargs,
         )

python/sglang/srt/layers/attention/mamba/ops/ssu_dispatch.py

@@ -45,12 +45,19 @@ class MambaSSUBackend(ABC):
 class TritonSSUBackend(MambaSSUBackend):
     """Triton-based selective-state-update backend."""
 
-    def __init__(self) -> None:
+    def __init__(
+        self,
+        *,
+        enable_stochastic_rounding: bool = False,
+        cache_philox_rounds: int = 0,
+    ) -> None:
         from sglang.srt.layers.attention.mamba.ops.mamba_ssm import (
             selective_state_update,
         )
 
         self._kernel = selective_state_update
+        self._enable_stochastic_rounding = enable_stochastic_rounding
+        self._cache_philox_rounds = cache_philox_rounds
 
     @property
     def name(self) -> str:
@@ -96,16 +103,25 @@ class TritonSSUBackend(MambaSSUBackend):
             cache_steps=cache_steps,
             retrieve_parent_token=retrieve_parent_token,
             intermediate_state_indices=intermediate_state_indices,
+            enable_stochastic_rounding=self._enable_stochastic_rounding,
+            cache_philox_rounds=self._cache_philox_rounds,
         )
 
 
 class FlashInferSSUBackend(MambaSSUBackend):
     """FlashInfer-based selective-state-update backend."""
 
-    def __init__(self) -> None:
+    def __init__(
+        self,
+        *,
+        enable_stochastic_rounding: bool = False,
+        cache_philox_rounds: int = 0,
+    ) -> None:
         from flashinfer.mamba import selective_state_update
 
         self._kernel = selective_state_update
+        self._enable_stochastic_rounding = enable_stochastic_rounding
+        self._cache_philox_rounds = cache_philox_rounds
 
     @property
     def name(self) -> str:
@@ -137,6 +153,11 @@ class FlashInferSSUBackend(MambaSSUBackend):
                 "FlashInfer backend does not support retrieve_parent_token. "
                 "Use --mamba-backend triton for EAGLE tree attention."
             )
+        rand_seed = (
+            torch.randint(0, 2**32, (1,), device=state.device)
+            if self._enable_stochastic_rounding
+            else None
+        )
         # FlashInfer expects cache_steps as an int (0 when unused).
         self._kernel(
             state,
@@ -156,6 +177,8 @@ class FlashInferSSUBackend(MambaSSUBackend):
             intermediate_states_buffer=intermediate_states_buffer,
             cache_steps=0 if cache_steps is None else cache_steps,
             intermediate_state_indices=intermediate_state_indices,
+            rand_seed=rand_seed,
+            philox_rounds=self._cache_philox_rounds or 10,
         )
 
 
@@ -190,7 +213,12 @@ def initialize_mamba_selective_state_update_backend(server_args: ServerArgs) ->
         )
 
     try:
-        _mamba_ssu_backend = backend_cls()
+        _mamba_ssu_backend = backend_cls(
+            enable_stochastic_rounding=(
+                server_args.enable_mamba_cache_stochastic_rounding
+            ),
+            cache_philox_rounds=server_args.mamba_cache_philox_rounds,
+        )
     except ImportError:
         raise ValueError(
             f"Mamba backend '{requested}' requested but its dependencies are not "

python/sglang/srt/server_args.py

@@ -1793,6 +1793,14 @@ class ServerArgs:
             choices=["float32", "bfloat16", "float16"],
         ),
     ] = None
+    enable_mamba_cache_stochastic_rounding: A[
+        bool,
+        "Enable stochastic rounding when writing FP16 Mamba SSM cache states. Requires --mamba-ssm-dtype float16 and CUDA. With --mamba-backend triton, requires SM100.",
+    ] = False
+    mamba_cache_philox_rounds: A[
+        int,
+        "Number of Philox rounds to use for stochastic rounding of FP16 Mamba SSM cache writes. Triton uses the Triton default when set to 0; FlashInfer uses 10 rounds when set to 0.",
+    ] = 0
     mamba_full_memory_ratio: A[
         float,
         "The ratio of mamba state memory to full kv cache memory.",
@@ -4984,20 +4992,52 @@ class ServerArgs:
             self.grammar_backend = "xgrammar"
 
     def _handle_mamba_backend(self):
+        if self.mamba_cache_philox_rounds < 0:
+            raise ValueError("--mamba-cache-philox-rounds must be non-negative.")
+
+        if self.enable_mamba_cache_stochastic_rounding:
+            if self.mamba_ssm_dtype != "float16":
+                raise ValueError(
+                    "Stochastic rounding for the Mamba SSM cache requires "
+                    f"--mamba-ssm-dtype float16, got {self.mamba_ssm_dtype!r}. "
+                    "Run with --mamba-ssm-dtype float16 or disable "
+                    "--enable-mamba-cache-stochastic-rounding."
+                )
+            if not is_cuda():
+                raise ValueError(
+                    "Stochastic rounding for the Mamba SSM cache is only "
+                    "supported on NVIDIA CUDA platforms. Disable "
+                    "--enable-mamba-cache-stochastic-rounding on this platform."
+                )
+            if self.mamba_backend == "triton" and not is_sm100_supported():
+                raise ValueError(
+                    "Stochastic rounding for the Mamba SSM cache with "
+                    "--mamba-backend triton requires SM100 with CUDA >= 12.8 "
+                    "because it uses the cvt.rs.f16x2.f32 PTX instruction. On "
+                    "H100/SM90, run with --mamba-backend flashinfer "
+                    "--mamba-ssm-dtype float16, or disable "
+                    "--enable-mamba-cache-stochastic-rounding."
+                )
+
         if self.mamba_backend == "flashinfer":
+            flashinfer_error = (
+                "FlashInfer mamba module not available, please check the "
+                "FlashInfer installation."
+            )
+            if self.enable_mamba_cache_stochastic_rounding:
+                flashinfer_error += (
+                    " Stochastic rounding with --mamba-backend flashinfer "
+                    "requires FlashInfer Mamba and --mamba-ssm-dtype float16."
+                )
             if is_flashinfer_available():
                 try:
                     import flashinfer.mamba  # noqa: F401
 
                     logger.info("Successfully imported FlashInfer mamba module")
                 except (ImportError, AttributeError):
-                    raise ValueError(
+                    raise ValueError(flashinfer_error)
-                        "FlashInfer mamba module not available, please check flashinfer installation."
-                    )
             else:
-                raise ValueError(
+                raise ValueError(flashinfer_error)
-                    "FlashInfer mamba module not available, please check flashinfer installation."
-                )
 
     def _handle_int8_mamba_checkpoint(self):
         # The int8 mamba checkpoint pool is only wired into the built-in

test/registered/layers/mamba/test_mamba_ssm.py

@@ -14,8 +14,8 @@ import torch.nn.functional as F
 from einops import rearrange, repeat
 
 from sglang.srt.layers.attention.mamba.causal_conv1d_triton import PAD_SLOT_ID
-from sglang.srt.layers.attention.mamba.ops import selective_state_update
+from sglang.srt.layers.attention.mamba.ops.mamba_ssm import selective_state_update
-from sglang.srt.utils import get_device
+from sglang.srt.utils import get_device, is_sm100_supported
 
 
 def selective_state_update_ref(
@@ -297,6 +297,67 @@ def test_selective_state_update_with_heads_with_batch_indices(
     assert torch.allclose(out, out_ref, rtol=rtol, atol=atol)
 
 
+@pytest.mark.skipif(
+    not is_sm100_supported(),
+    reason=(
+        "Triton stochastic rounding uses cvt.rs.f16x2.f32 and requires "
+        "SM100-family Blackwell with CUDA >= 12.8"
+    ),
+)
+@pytest.mark.parametrize("philox_rounds", [0, 4])
+@pytest.mark.parametrize("has_z", [False, True])
+@pytest.mark.parametrize("dstate", [16, 64])
+@pytest.mark.parametrize("dim", [2048, 4096])
+def test_selective_state_update_stochastic_rounding(dim, dstate, has_z, philox_rounds):
+    device = "cuda"
+    torch.manual_seed(0)
+
+    batch_size = 2
+    state = torch.randn(batch_size, dim, dstate, dtype=torch.float16, device=device)
+    x = torch.randn(batch_size, dim, device=device, dtype=torch.bfloat16)
+    out = torch.empty_like(x)
+    dt = torch.randn(batch_size, dim, device=device, dtype=torch.bfloat16)
+    dt_bias = torch.rand(dim, device=device) - 4.0
+    A = -torch.rand(dim, dstate, device=device) - 1.0
+    B = torch.randn(batch_size, dstate, device=device)
+    C = torch.randn(batch_size, dstate, device=device)
+    D = torch.randn(dim, device=device)
+    z = torch.randn_like(x) if has_z else None
+    state_ref = state.float()
+
+    selective_state_update(
+        state,
+        x,
+        dt,
+        A,
+        B,
+        C,
+        D=D,
+        z=z,
+        dt_bias=dt_bias,
+        dt_softplus=True,
+        out=out,
+        enable_stochastic_rounding=True,
+        cache_philox_rounds=philox_rounds,
+    )
+    out_ref = selective_state_update_ref(
+        state_ref,
+        x,
+        dt,
+        A,
+        B,
+        C,
+        D=D,
+        z=z,
+        dt_bias=dt_bias,
+        dt_softplus=True,
+    )
+
+    assert state.dtype == torch.float16
+    assert torch.allclose(state, state_ref.to(torch.float16), rtol=5e-3, atol=1e-1)
+    assert torch.allclose(out, out_ref, rtol=5e-3, atol=1e-1)
+
+
 if __name__ == "__main__":
     import sys
 

test/registered/unit/server_args/test_server_args.py

@@ -71,6 +71,42 @@ class TestPrepareServerArgs(CustomTestCase):
             os.unlink(config_file)
 
 
+class TestMambaCacheStochasticRounding(unittest.TestCase):
+    def test_rejects_fp32_ssm_cache(self):
+        server_args = ServerArgs(
+            model_path="dummy",
+            mamba_ssm_dtype="float32",
+            enable_mamba_cache_stochastic_rounding=True,
+        )
+
+        with self.assertRaisesRegex(ValueError, "--mamba-ssm-dtype float16"):
+            server_args._handle_mamba_backend()
+
+    @patch("sglang.srt.server_args.is_cuda", return_value=False)
+    def test_rejects_non_cuda(self, _mock_is_cuda):
+        server_args = ServerArgs(
+            model_path="dummy",
+            mamba_ssm_dtype="float16",
+            enable_mamba_cache_stochastic_rounding=True,
+        )
+
+        with self.assertRaisesRegex(ValueError, "NVIDIA CUDA"):
+            server_args._handle_mamba_backend()
+
+    @patch("sglang.srt.server_args.is_cuda", return_value=True)
+    @patch("sglang.srt.server_args.is_sm100_supported", return_value=False)
+    def test_rejects_triton_without_sm100(self, _mock_sm100, _mock_is_cuda):
+        server_args = ServerArgs(
+            model_path="dummy",
+            mamba_ssm_dtype="float16",
+            mamba_backend="triton",
+            enable_mamba_cache_stochastic_rounding=True,
+        )
+
+        with self.assertRaisesRegex(ValueError, "requires SM100"):
+            server_args._handle_mamba_backend()
+
+
 class TestLoadBalanceMethod(unittest.TestCase):
     def test_non_pd_defaults_to_round_robin(self):
         server_args = ServerArgs(model_path="dummy", disaggregation_mode="null")