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
python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py
+63 -1
View File
@@ -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
+31 -3
View File
@@ -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
+46 -6
View File
@@ -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