diff --git a/docs_new/cookbook/autoregressive/NVIDIA/Nemotron3-Ultra.mdx b/docs_new/cookbook/autoregressive/NVIDIA/Nemotron3-Ultra.mdx index 4cf1a513c..ab379faab 100644 --- a/docs_new/cookbook/autoregressive/NVIDIA/Nemotron3-Ultra.mdx +++ b/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: diff --git a/docs_new/docs/advanced_features/server_arguments.mdx b/docs_new/docs/advanced_features/server_arguments.mdx index 065d1d1f1..ee8af4482 100644 --- a/docs_new/docs/advanced_features/server_arguments.mdx +++ b/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 The data type of the SSM states in mamba cache. If not set, read from the model config. Auto (from model config) float32, bfloat16, float16 + + + `--enable-mamba-cache-stochastic-rounding` + Enable stochastic rounding when writing FP16 Mamba SSM cache states. Requires --mamba-ssm-dtype float16 and CUDA. With --mamba-backend triton, requires SM100. + `False` + Type: bool + + + `--mamba-cache-philox-rounds` + 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` + Type: int `--mamba-full-memory-ratio` diff --git a/docs_new/src/snippets/autoregressive/nemotron3-ultra-deployment.jsx b/docs_new/src/snippets/autoregressive/nemotron3-ultra-deployment.jsx index ee1e8ad43..0b4755eee 100644 --- a/docs_new/src/snippets/autoregressive/nemotron3-ultra-deployment.jsx +++ b/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', diff --git a/python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py b/python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py index 6aba2737e..09dbd73db 100644 --- a/python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py +++ b/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, ) diff --git a/python/sglang/srt/layers/attention/mamba/ops/ssu_dispatch.py b/python/sglang/srt/layers/attention/mamba/ops/ssu_dispatch.py index 2854e5fa4..df61896be 100644 --- a/python/sglang/srt/layers/attention/mamba/ops/ssu_dispatch.py +++ b/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 " diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index e5fbcaa83..29af32941 100644 --- a/python/sglang/srt/server_args.py +++ b/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( - "FlashInfer mamba module not available, please check flashinfer installation." - ) + raise ValueError(flashinfer_error) else: - raise ValueError( - "FlashInfer mamba module not available, please check flashinfer installation." - ) + raise ValueError(flashinfer_error) def _handle_int8_mamba_checkpoint(self): # The int8 mamba checkpoint pool is only wired into the built-in diff --git a/test/registered/layers/mamba/test_mamba_ssm.py b/test/registered/layers/mamba/test_mamba_ssm.py index 39c749eb4..165262178 100644 --- a/test/registered/layers/mamba/test_mamba_ssm.py +++ b/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.utils import get_device +from sglang.srt.layers.attention.mamba.ops.mamba_ssm import selective_state_update +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 diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index a06071e19..a6dd00a5f 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/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")