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',