diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index de3cebabd..0ab8080ad 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -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, + ) diff --git a/python/sglang/srt/layers/attention/mamba/mamba_state_scatter_triton.py b/python/sglang/srt/layers/attention/mamba/mamba_state_scatter_triton.py new file mode 100644 index 000000000..0c1b7efac --- /dev/null +++ b/python/sglang/srt/layers/attention/mamba/mamba_state_scatter_triton.py @@ -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, + ) diff --git a/test/unit/test_mamba_state_scatter_triton.py b/test/unit/test_mamba_state_scatter_triton.py new file mode 100644 index 000000000..0778a091e --- /dev/null +++ b/test/unit/test_mamba_state_scatter_triton.py @@ -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()