From ce1736fcc6cc9a599752199c441a99123f8a3bcf Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Mon, 11 May 2026 14:50:46 -0700 Subject: [PATCH] [Spec] Remove dead kernel params; fix stale comment in `trtllm_mla` (#25010) --- python/sglang/srt/layers/attention/fla/kda.py | 3 -- .../mamba/mamba_state_scatter_triton.py | 3 -- .../layers/attention/trtllm_mla_backend.py | 4 +-- .../layers/test_mamba_state_scatter_triton.py | 34 +++++++++---------- 4 files changed, 19 insertions(+), 25 deletions(-) diff --git a/python/sglang/srt/layers/attention/fla/kda.py b/python/sglang/srt/layers/attention/fla/kda.py index 2a6fd059b..06656ac39 100644 --- a/python/sglang/srt/layers/attention/fla/kda.py +++ b/python/sglang/srt/layers/attention/fla/kda.py @@ -52,7 +52,6 @@ def fused_recurrent_kda_fwd( inplace_final_state: bool = True, cu_seqlens: torch.LongTensor | None = None, # ssm_state_indices: torch.Tensor | None = None, - num_accepted_tokens: torch.Tensor | None = None, use_qk_l2norm_in_kernel: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, v.shape[-1] @@ -92,7 +91,6 @@ def fused_recurrent_kda_fwd( ht=final_state, cu_seqlens=cu_seqlens, # ssm_state_indices=ssm_state_indices, - # num_accepted_tokens=num_accepted_tokens, scale=scale, # N=N, T=T, @@ -155,7 +153,6 @@ def fused_recurrent_kda( inplace_final_state=inplace_final_state, cu_seqlens=cu_seqlens, # ssm_state_indices=ssm_state_indices, - num_accepted_tokens=None, use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, ) return o, final_state 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 index 0c1b7efac..a97320cac 100644 --- a/python/sglang/srt/layers/attention/mamba/mamba_state_scatter_triton.py +++ b/python/sglang/srt/layers/attention/mamba/mamba_state_scatter_triton.py @@ -18,8 +18,6 @@ def _fused_mamba_state_scatter_with_mask_kernel( # 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, @@ -176,7 +174,6 @@ def fused_mamba_state_scatter_with_mask( dst, dst_indices_raw, step_indices_raw, - total_requests, elem_per_entry, src_layer_stride, src_req_stride, diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index 58d7ab2f2..68a102d14 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mla_backend.py @@ -59,7 +59,7 @@ def pad_draft_extend_query_kernel( q_ptr, # Input query tensor [total_seq_len, num_heads, head_dim] padded_q_ptr, # Output padded query tensor [batch_size, max_seq_len, num_heads, head_dim] seq_lens_q_ptr, # Sequence lengths for each sequence [batch_size] - cumsum_ptr, # Cumulative sum of accept lengths [batch_size + 1] + cumsum_ptr, # Cumulative sum of sequence lengths [batch_size + 1] batch_size, max_seq_len, num_heads, @@ -78,7 +78,7 @@ def pad_draft_extend_query_kernel( if batch_id >= batch_size: return - # Load accept length for this batch + # Load sequence length for this batch seq_len = tl.load(seq_lens_q_ptr + batch_id) if seq_pos >= seq_len: diff --git a/test/registered/unit/layers/test_mamba_state_scatter_triton.py b/test/registered/unit/layers/test_mamba_state_scatter_triton.py index 1f20373b3..938251641 100644 --- a/test/registered/unit/layers/test_mamba_state_scatter_triton.py +++ b/test/registered/unit/layers/test_mamba_state_scatter_triton.py @@ -45,20 +45,20 @@ def _ref_update_like( intermediate_conv, *, state_indices_tensor, - accepted_steps, + step_indices_raw, mamba_track_indices=None, mamba_steps_to_track=None, ): """Reference implementation using PyTorch advanced indexing for correctness verification.""" - request_number = accepted_steps.shape[0] + total_requests = step_indices_raw.shape[0] intermediate_state_indices = torch.arange( - request_number, dtype=torch.int32, device=accepted_steps.device + total_requests, dtype=torch.int32, device=step_indices_raw.device ) - valid_mask = accepted_steps >= 0 + valid_mask = step_indices_raw >= 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) + last_steps = step_indices_raw[valid_mask].to(torch.int64) # Only scatter if there are valid indices (but don't early return - # mamba_track_indices processing is independent) @@ -110,7 +110,7 @@ def _fused_update_like( intermediate_conv, *, state_indices_tensor, - accepted_steps, + step_indices_raw, mamba_track_indices=None, mamba_steps_to_track=None, ): @@ -120,13 +120,13 @@ def _fused_update_like( ssm_states, intermediate_ssm, state_indices_tensor, - accepted_steps, + step_indices_raw, ) fused_mamba_state_scatter_with_mask( conv_states, intermediate_conv, state_indices_tensor, - accepted_steps, + step_indices_raw, ) if mamba_track_indices is not None: @@ -199,10 +199,10 @@ class TestMambaStateScatterCorrectness(unittest.TestCase): :B ].to(torch.int32) - accepted_steps = torch.randint(0, D, (B,), device=device, dtype=torch.int64) + step_indices_raw = 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 + step_indices_raw[invalid] = -1 # Optional track update mamba_track_indices = torch.randperm(C, device=device, dtype=torch.int64)[:B] @@ -223,7 +223,7 @@ class TestMambaStateScatterCorrectness(unittest.TestCase): conv_ref, intermediate_conv, state_indices_tensor=state_indices_tensor, - accepted_steps=accepted_steps, + step_indices_raw=step_indices_raw, mamba_track_indices=mamba_track_indices, mamba_steps_to_track=mamba_steps_to_track, ) @@ -233,7 +233,7 @@ class TestMambaStateScatterCorrectness(unittest.TestCase): conv_fused, intermediate_conv, state_indices_tensor=state_indices_tensor, - accepted_steps=accepted_steps, + step_indices_raw=step_indices_raw, mamba_track_indices=mamba_track_indices, mamba_steps_to_track=mamba_steps_to_track, ) @@ -290,10 +290,10 @@ class TestMambaStateScatterPerf(unittest.TestCase): 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) + step_indices_raw = 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 + step_indices_raw[invalid] = -1 mamba_track_indices = None mamba_steps_to_track = None @@ -314,7 +314,7 @@ class TestMambaStateScatterPerf(unittest.TestCase): conv_states, intermediate_conv, state_indices_tensor=state_indices_tensor, - accepted_steps=accepted_steps, + step_indices_raw=step_indices_raw, mamba_track_indices=mamba_track_indices, mamba_steps_to_track=mamba_steps_to_track, ) @@ -326,7 +326,7 @@ class TestMambaStateScatterPerf(unittest.TestCase): conv_states, intermediate_conv, state_indices_tensor=state_indices_tensor, - accepted_steps=accepted_steps, + step_indices_raw=step_indices_raw, mamba_track_indices=mamba_track_indices, mamba_steps_to_track=mamba_steps_to_track, ) @@ -339,7 +339,7 @@ class TestMambaStateScatterPerf(unittest.TestCase): ref_ms = _time_cuda_ms(ref_fn) fused_ms = _time_cuda_ms(fused_fn) - num_valid = int((accepted_steps >= 0).sum().item()) + num_valid = int((step_indices_raw >= 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")