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

Perf/fuse mamba state scatter mtp verify (#18088)

Browse Source
This commit is contained in:
xutizhou
2026-02-25 15:40:55 +08:00
committed by GitHub
parent ab7071b545
commit a1b39c1c26
3 changed files with 583 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py
+39 -37
View File
@@ -32,6 +32,9 @@ from sglang.srt.layers.attention.mamba.mamba2_metadata import (
ForwardMetadata,
Mamba2Metadata,
)
from sglang.srt.layers.attention.mamba.mamba_state_scatter_triton import (
fused_mamba_state_scatter_with_mask,
)
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.layers.radix_linear_attention import RadixLinearAttention
from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, MambaPool
@@ -1699,6 +1702,15 @@ class HybridLinearAttnBackend(AttentionBackend):
mamba_steps_to_track: Optional[torch.Tensor],
model,
):
"""
Update mamba states after MTP verify using fully fused Triton kernel.
This replaces the original advanced indexing operations with a single fused
gather-scatter kernel that also handles masking internally, avoiding:
- index_elementwise_kernel from tensor[bool_mask]
- index_select kernel launches
- nonzero kernel launches
"""
request_number = accepted_steps.shape[0]
state_indices_tensor = (
@@ -1706,9 +1718,6 @@ class HybridLinearAttnBackend(AttentionBackend):
:request_number
]
)
intermediate_state_indices = torch.arange(
request_number, dtype=torch.int32, device=state_indices_tensor.device
)
mamba_caches = (
self.linear_attn_backend.req_to_token_pool.get_speculative_mamba2_params_all_layers()
@@ -1719,41 +1728,34 @@ class HybridLinearAttnBackend(AttentionBackend):
intermediate_state_cache = mamba_caches.intermediate_ssm
intermediate_conv_window_cache = mamba_caches.intermediate_conv_window[0]
# Compute common indices once to avoid duplication
valid_mask = accepted_steps >= 0
dst_state_indices = state_indices_tensor[valid_mask].to(torch.int64) # [N]
src_state_indices = intermediate_state_indices[valid_mask].to(
torch.int64
) # [N]
last_steps = accepted_steps[valid_mask].to(torch.int64) # [N]
# scatter into ssm_states at the chosen cache lines
ssm_states[:, dst_state_indices, :] = intermediate_state_cache[
:, src_state_indices, last_steps
].to(ssm_states.dtype, copy=False)
# Scatter into conv_states at the chosen cache lines
conv_states[:, dst_state_indices, :] = intermediate_conv_window_cache[
:, src_state_indices, last_steps
].to(conv_states.dtype, copy=False)
# Use fully fused kernel that handles masking internally
# This avoids separate nonzero() and index_select() calls
fused_mamba_state_scatter_with_mask(
ssm_states,
intermediate_state_cache,
state_indices_tensor,
accepted_steps,
)
fused_mamba_state_scatter_with_mask(
conv_states,
intermediate_conv_window_cache,
state_indices_tensor,
accepted_steps,
)
# Track indices used for tracking mamba states for prefix cache
if mamba_track_indices is not None:
assert mamba_steps_to_track is not None
track_mask = mamba_steps_to_track >= 0
track_steps = mamba_steps_to_track[track_mask].to(torch.int64) # [N]
if track_steps.numel() == 0:
# No track indices to update
return
dst_track_indices = mamba_track_indices[track_mask].to(torch.int64)
src_track_indices = intermediate_state_indices[track_mask].to(torch.int64)
# scatter into ssm_states at the chosen track states
ssm_states[:, dst_track_indices, :] = intermediate_state_cache[
:, src_track_indices, track_steps
].to(ssm_states.dtype, copy=False)
# scatter into conv_states at the chosen track states
conv_states[:, dst_track_indices, :] = intermediate_conv_window_cache[
:, src_track_indices, track_steps
].to(conv_states.dtype, copy=False)
# Use fully fused kernel for track scatter operations
fused_mamba_state_scatter_with_mask(
ssm_states,
intermediate_state_cache,
mamba_track_indices,
mamba_steps_to_track,
)
fused_mamba_state_scatter_with_mask(
conv_states,
intermediate_conv_window_cache,
mamba_track_indices,
mamba_steps_to_track,
)
python/sglang/srt/layers/attention/mamba/mamba_state_scatter_triton.py
+190
View File
@@ -0,0 +1,190 @@
"""
Fused Triton kernel for Mamba state scatter operations.
This kernel replaces the expensive advanced indexing operations in
`update_mamba_state_after_mtp_verify` with a single fused gather-scatter kernel,
avoiding multiple `index_elementwise_kernel` launches.
"""
import torch
import triton
import triton.language as tl
@triton.jit
def _fused_mamba_state_scatter_with_mask_kernel(
src_ptr,
dst_ptr,
# Raw index arrays (before index_select)
dst_indices_raw_ptr, # [total_requests] - state_indices_tensor
step_indices_raw_ptr, # [total_requests] - accepted_steps or mamba_steps_to_track
# Total number of requests
total_requests,
elem_per_entry: tl.constexpr,
src_layer_stride,
src_req_stride,
src_step_stride,
dst_layer_stride,
dst_req_stride,
src_req_size,
src_step_size,
dst_req_size,
BLOCK_SIZE: tl.constexpr,
):
"""
Fused gather-scatter kernel with built-in masking.
This kernel fuses the index_select operations by:
1. Iterating over all requests (pid_req from 0 to total_requests-1)
2. Checking if step_indices_raw[pid_req] >= 0 (valid mask)
3. If valid, performing the scatter:
dst[l, dst_indices_raw[pid_req], :] = src[l, pid_req, step_indices_raw[pid_req], :]
Grid: (total_requests, num_layers, ceil(elem_per_entry / BLOCK_SIZE))
"""
pid_req = tl.program_id(0)
pid_layer = tl.program_id(1).to(tl.int64)
pid_block = tl.program_id(2).to(tl.int64)
# Load step index to check validity (step >= 0 means valid)
step_idx = tl.load(step_indices_raw_ptr + pid_req).to(tl.int64)
# Early exit if this request is not valid (step < 0)
if step_idx < 0:
return
# Load destination index
dst_idx = tl.load(dst_indices_raw_ptr + pid_req).to(tl.int64)
# Source index is just the request index itself
src_idx = pid_req
# Bounds check to avoid illegal memory access
if not (
(dst_idx >= 0)
& (dst_idx < dst_req_size)
& (src_idx >= 0)
& (src_idx < src_req_size)
& (step_idx < src_step_size)
):
return
# Compute base offsets
src_offset = (
pid_layer * src_layer_stride
+ src_idx * src_req_stride
+ step_idx * src_step_stride
)
dst_offset = pid_layer * dst_layer_stride + dst_idx * dst_req_stride
# Compute element range for this block
start = pid_block * BLOCK_SIZE
offsets = start + tl.arange(0, BLOCK_SIZE)
mask = offsets < elem_per_entry
# Load from source and store to destination
data = tl.load(src_ptr + src_offset + offsets, mask=mask)
tl.store(dst_ptr + dst_offset + offsets, data, mask=mask)
def fused_mamba_state_scatter_with_mask(
dst: torch.Tensor, # [num_layers, cache_size, *state_shape]
src: torch.Tensor, # [num_layers, spec_size, draft_tokens, *state_shape]
dst_indices_raw: torch.Tensor, # [total_requests] - raw indices (e.g., state_indices_tensor)
step_indices_raw: torch.Tensor, # [total_requests] - raw step indices (step >= 0 means valid)
):
"""
Fully fused gather-scatter with built-in masking for mamba state updates.
This function fuses the following operations into a single kernel:
1. valid_mask = step_indices_raw >= 0
2. valid_indices = valid_mask.nonzero()
3. dst_indices = dst_indices_raw[valid_indices] (index_select)
4. step_indices = step_indices_raw[valid_indices] (index_select)
5. for each valid i: dst[:, dst_indices[i], :] = src[:, i, step_indices[i], :]
Args:
dst: Destination tensor [num_layers, cache_size, *state_shape]
src: Source tensor [num_layers, spec_size, draft_tokens, *state_shape]
dst_indices_raw: Raw destination indices for all requests [total_requests]
step_indices_raw: Raw step indices; entry >= 0 means valid [total_requests]
"""
total_requests = step_indices_raw.shape[0]
if total_requests == 0:
return
if dst.device != src.device:
raise ValueError(
f"dst and src must be on the same device. {dst.device=} {src.device=}"
)
if not dst.is_cuda or not src.is_cuda:
raise ValueError(
"fused_mamba_state_scatter_with_mask only supports CUDA tensors."
)
if dst.ndim < 2 or src.ndim < 3:
raise ValueError(f"Unexpected tensor ranks: {dst.ndim=} {src.ndim=}")
if dst.shape[0] != src.shape[0]:
raise ValueError(
f"Layer dimension mismatch: {dst.shape[0]=} vs {src.shape[0]=}"
)
if dst.shape[2:] != src.shape[3:]:
raise ValueError(
f"Trailing dims mismatch: {dst.shape[2:]=} vs {src.shape[3:]=}"
)
if dst_indices_raw.ndim != 1 or step_indices_raw.ndim != 1:
raise ValueError(
f"indices must be 1D: {dst_indices_raw.shape=} {step_indices_raw.shape=}"
)
if dst_indices_raw.shape[0] != step_indices_raw.shape[0]:
raise ValueError(
f"indices length mismatch: {dst_indices_raw.shape[0]=} vs {step_indices_raw.shape[0]=}"
)
num_layers = dst.shape[0]
src_req_size = src.shape[1]
src_step_size = src.shape[2]
dst_req_size = dst.shape[1]
# Flatten trailing dimensions: number of elements per (layer, cache_line) entry.
elem_per_entry = dst.numel() // (dst.shape[0] * dst.shape[1])
# Get strides (in elements, not bytes)
src_layer_stride = src.stride(0)
src_req_stride = src.stride(1)
src_step_stride = src.stride(2)
dst_layer_stride = dst.stride(0)
dst_req_stride = dst.stride(1)
# Ensure indices are int32 and contiguous
dst_indices_raw = dst_indices_raw.to(torch.int32).contiguous()
step_indices_raw = step_indices_raw.to(torch.int32).contiguous()
# Ensure tensors are contiguous
if not dst.is_contiguous():
raise ValueError("dst tensor must be contiguous")
if not src.is_contiguous():
raise ValueError("src tensor must be contiguous")
# Block size for copying elements
BLOCK_SIZE = 1024
# Grid over all requests - invalid ones will early-exit in the kernel
grid = (total_requests, num_layers, triton.cdiv(elem_per_entry, BLOCK_SIZE))
_fused_mamba_state_scatter_with_mask_kernel[grid](
src,
dst,
dst_indices_raw,
step_indices_raw,
total_requests,
elem_per_entry,
src_layer_stride,
src_req_stride,
src_step_stride,
dst_layer_stride,
dst_req_stride,
src_req_size,
src_step_size,
dst_req_size,
BLOCK_SIZE=BLOCK_SIZE,
)
test/unit/test_mamba_state_scatter_triton.py
+354
View File
@@ -0,0 +1,354 @@
import os
import unittest
import torch
try:
from sglang.srt.layers.attention.mamba.mamba_state_scatter_triton import (
fused_mamba_state_scatter_with_mask,
)
_FUSED_IMPORT_ERROR = None
except Exception as e: # pragma: no cover
fused_mamba_state_scatter_with_mask = None
_FUSED_IMPORT_ERROR = e
def _dtype_from_str(name: str) -> torch.dtype:
mapping = {
"bfloat16": torch.bfloat16,
"float16": torch.float16,
"float32": torch.float32,
}
if name not in mapping:
raise ValueError(
f"Unsupported dtype string {name!r}. Supported: {sorted(mapping.keys())}"
)
return mapping[name]
def _ref_scatter(dst, src, dst_indices, src_indices, step_indices):
"""Reference implementation using PyTorch advanced indexing."""
# dst: [L, C, E]
# src: [L, S, D, E]
dst[:, dst_indices] = src[:, src_indices, step_indices].to(dst.dtype, copy=False)
def _ref_update_like(
ssm_states,
intermediate_ssm,
conv_states,
intermediate_conv,
*,
state_indices_tensor,
accepted_steps,
mamba_track_indices=None,
mamba_steps_to_track=None,
):
"""Reference implementation using PyTorch advanced indexing for correctness verification."""
request_number = accepted_steps.shape[0]
intermediate_state_indices = torch.arange(
request_number, dtype=torch.int32, device=accepted_steps.device
)
valid_mask = accepted_steps >= 0
dst_state_indices = state_indices_tensor[valid_mask].to(torch.int64)
src_state_indices = intermediate_state_indices[valid_mask].to(torch.int64)
last_steps = accepted_steps[valid_mask].to(torch.int64)
# Only scatter if there are valid indices (but don't early return -
# mamba_track_indices processing is independent)
if dst_state_indices.numel() > 0:
_ref_scatter(
ssm_states,
intermediate_ssm,
dst_state_indices,
src_state_indices,
last_steps,
)
_ref_scatter(
conv_states,
intermediate_conv,
dst_state_indices,
src_state_indices,
last_steps,
)
if mamba_track_indices is not None:
assert mamba_steps_to_track is not None
track_mask = mamba_steps_to_track >= 0
if not track_mask.any():
return
dst_track_indices = mamba_track_indices[track_mask].to(torch.int64)
src_track_indices = intermediate_state_indices[track_mask].to(torch.int64)
track_steps = mamba_steps_to_track[track_mask].to(torch.int64)
_ref_scatter(
ssm_states,
intermediate_ssm,
dst_track_indices,
src_track_indices,
track_steps,
)
_ref_scatter(
conv_states,
intermediate_conv,
dst_track_indices,
src_track_indices,
track_steps,
)
def _fused_update_like(
ssm_states,
intermediate_ssm,
conv_states,
intermediate_conv,
*,
state_indices_tensor,
accepted_steps,
mamba_track_indices=None,
mamba_steps_to_track=None,
):
"""Matches the fully fused logic that avoids index_select and nonzero calls."""
# Use fully fused kernel that handles masking internally
fused_mamba_state_scatter_with_mask(
ssm_states,
intermediate_ssm,
state_indices_tensor,
accepted_steps,
)
fused_mamba_state_scatter_with_mask(
conv_states,
intermediate_conv,
state_indices_tensor,
accepted_steps,
)
if mamba_track_indices is not None:
assert mamba_steps_to_track is not None
fused_mamba_state_scatter_with_mask(
ssm_states,
intermediate_ssm,
mamba_track_indices,
mamba_steps_to_track,
)
fused_mamba_state_scatter_with_mask(
conv_states,
intermediate_conv,
mamba_track_indices,
mamba_steps_to_track,
)
def _time_cuda_ms(fn, iters=50, warmup=10):
"""Measure average CUDA time (ms) using CUDA events."""
for _ in range(warmup):
fn()
torch.cuda.synchronize()
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
for _ in range(iters):
fn()
end.record()
torch.cuda.synchronize()
return start.elapsed_time(end) / iters
class TestMambaStateScatterCorrectness(unittest.TestCase):
@unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for this test.")
def test_fused_matches_reference(self):
"""Test that fused_mamba_state_scatter_with_mask matches the reference."""
if fused_mamba_state_scatter_with_mask is None:
self.skipTest(
f"fused_mamba_state_scatter_with_mask import failed: {_FUSED_IMPORT_ERROR}"
)
torch.manual_seed(42)
device = torch.device("cuda")
# Keep sizes moderate so this test is quick.
L = 8
B = 32
C = 49
D = 5
ssm_elems = 1024
conv_elems = 512
ssm_states0 = torch.empty(
(L, C, ssm_elems), device=device, dtype=torch.bfloat16
)
conv_states0 = torch.empty(
(L, C, conv_elems), device=device, dtype=torch.bfloat16
)
intermediate_ssm = torch.randn(
(L, B, D, ssm_elems), device=device, dtype=torch.bfloat16
)
intermediate_conv = torch.randn(
(L, B, D, conv_elems), device=device, dtype=torch.bfloat16
)
# unique cache lines (no duplicates) to avoid nondeterministic write order
state_indices_tensor = torch.randperm(C, device=device, dtype=torch.int64)[
:B
].to(torch.int32)
accepted_steps = torch.randint(0, D, (B,), device=device, dtype=torch.int64)
# set ~10% invalid
invalid = torch.rand((B,), device=device) < 0.1
accepted_steps[invalid] = -1
# Optional track update
mamba_track_indices = torch.randperm(C, device=device, dtype=torch.int64)[:B]
mamba_steps_to_track = torch.randint(
0, D, (B,), device=device, dtype=torch.int64
)
track_invalid = torch.rand((B,), device=device) < 0.7
mamba_steps_to_track[track_invalid] = -1
ssm_ref = ssm_states0.clone()
conv_ref = conv_states0.clone()
ssm_fused = ssm_states0.clone()
conv_fused = conv_states0.clone()
_ref_update_like(
ssm_ref,
intermediate_ssm,
conv_ref,
intermediate_conv,
state_indices_tensor=state_indices_tensor,
accepted_steps=accepted_steps,
mamba_track_indices=mamba_track_indices,
mamba_steps_to_track=mamba_steps_to_track,
)
_fused_update_like(
ssm_fused,
intermediate_ssm,
conv_fused,
intermediate_conv,
state_indices_tensor=state_indices_tensor,
accepted_steps=accepted_steps,
mamba_track_indices=mamba_track_indices,
mamba_steps_to_track=mamba_steps_to_track,
)
torch.testing.assert_close(ssm_fused, ssm_ref)
torch.testing.assert_close(conv_fused, conv_ref)
class TestMambaStateScatterPerf(unittest.TestCase):
@unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for this test.")
def test_perf_report_old_vs_fused(self):
"""Optional microbenchmark comparing baseline vs fused kernel.
Enable with: SGLANG_RUN_MAMBA_SCATTER_PERF_TEST=1
"""
if os.environ.get("SGLANG_RUN_MAMBA_SCATTER_PERF_TEST", "0") != "1":
self.skipTest("Set SGLANG_RUN_MAMBA_SCATTER_PERF_TEST=1 to run perf test.")
if fused_mamba_state_scatter_with_mask is None:
self.skipTest(
f"fused_mamba_state_scatter_with_mask import failed: {_FUSED_IMPORT_ERROR}"
)
torch.manual_seed(0)
device = torch.device("cuda")
# Parameterize sizes via env vars so we can match a real model more closely.
L = int(os.environ.get("SGLANG_MAMBA_SCATTER_LAYERS", "32"))
B = int(os.environ.get("SGLANG_MAMBA_SCATTER_BATCH", "48"))
C = int(os.environ.get("SGLANG_MAMBA_SCATTER_CACHE", "49"))
D = int(os.environ.get("SGLANG_MAMBA_SCATTER_DRAFT_TOKENS", "5"))
ssm_elems = int(os.environ.get("SGLANG_MAMBA_SCATTER_SSM_ELEMS", "4096"))
conv_elems = int(os.environ.get("SGLANG_MAMBA_SCATTER_CONV_ELEMS", "512"))
invalid_ratio = float(
os.environ.get("SGLANG_MAMBA_SCATTER_INVALID_RATIO", "0.0")
)
track_ratio = float(os.environ.get("SGLANG_MAMBA_SCATTER_TRACK_RATIO", "0.0"))
ssm_dtype = _dtype_from_str(
os.environ.get("SGLANG_MAMBA_SCATTER_SSM_DTYPE", "bfloat16")
)
conv_dtype = _dtype_from_str(
os.environ.get("SGLANG_MAMBA_SCATTER_CONV_DTYPE", "bfloat16")
)
# Use zeros for dst so each iteration overwrites the same memory.
ssm_states = torch.zeros((L, C, ssm_elems), device=device, dtype=ssm_dtype)
conv_states = torch.zeros((L, C, conv_elems), device=device, dtype=conv_dtype)
intermediate_ssm = torch.randn(
(L, B, D, ssm_elems), device=device, dtype=ssm_dtype
)
intermediate_conv = torch.randn(
(L, B, D, conv_elems), device=device, dtype=conv_dtype
)
state_indices_tensor = torch.randperm(C, device=device, dtype=torch.int64)[
:B
].to(torch.int32)
accepted_steps = torch.randint(0, D, (B,), device=device, dtype=torch.int64)
if invalid_ratio > 0:
invalid = torch.rand((B,), device=device) < invalid_ratio
accepted_steps[invalid] = -1
mamba_track_indices = None
mamba_steps_to_track = None
if track_ratio > 0:
mamba_track_indices = torch.randperm(C, device=device, dtype=torch.int64)[
:B
]
mamba_steps_to_track = torch.randint(
0, D, (B,), device=device, dtype=torch.int64
)
track_invalid = torch.rand((B,), device=device) >= track_ratio
mamba_steps_to_track[track_invalid] = -1
def ref_fn():
_ref_update_like(
ssm_states,
intermediate_ssm,
conv_states,
intermediate_conv,
state_indices_tensor=state_indices_tensor,
accepted_steps=accepted_steps,
mamba_track_indices=mamba_track_indices,
mamba_steps_to_track=mamba_steps_to_track,
)
def fused_fn():
_fused_update_like(
ssm_states,
intermediate_ssm,
conv_states,
intermediate_conv,
state_indices_tensor=state_indices_tensor,
accepted_steps=accepted_steps,
mamba_track_indices=mamba_track_indices,
mamba_steps_to_track=mamba_steps_to_track,
)
# Warm up JIT compilation for triton kernels (and caches for torch indexing)
ref_fn()
fused_fn()
torch.cuda.synchronize()
ref_ms = _time_cuda_ms(ref_fn)
fused_ms = _time_cuda_ms(fused_fn)
num_valid = int((accepted_steps >= 0).sum().item())
ratio = fused_ms / ref_ms if ref_ms > 0 else float("inf")
speedup = ref_ms / fused_ms if fused_ms > 0 else float("inf")
# Print a concise report
print(
"\n[MambaStateScatterPerf]\n"
f" shapes: L={L} B={B} C={C} D={D} ssm_elems={ssm_elems} conv_elems={conv_elems}\n"
f" dtypes: ssm={ssm_dtype} conv={conv_dtype}\n"
f" valid: {num_valid}/{B} invalid_ratio={invalid_ratio} track_ratio={track_ratio}\n"
f" ref_total_ms (baseline): {ref_ms:.4f}\n"
f" fused_total_ms: {fused_ms:.4f} (ratio={ratio:.3f}x, speedup={speedup:.2f}x)\n"
)
if __name__ == "__main__": # pragma: no cover
unittest.main()