minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Daniel Afrimi <dafrimi@login-lyris01.lyris.clusters.nvidia.com>
Co-authored-by: Daniel Afrimi <dafrimi@login-lyris01.lyris.clusters.nvidia.com>
This commit is contained in:
danielafrimi
2026-06-29 01:47:09 -07:00
committed by GitHub
co-authored by Daniel Afrimi
parent d5133e925b
commit a2b5ce2ed1
8 changed files with 267 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs_new/cookbook/autoregressive/NVIDIA/Nemotron3-Ultra.mdx
+4
View File
@@ -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: