[Spec] Remove dead kernel params; fix stale comment in trtllm_mla (#25010)

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")