[ROCm][Perf][GLM-5.3-Flash] "Fit kpool top-k indices to AITER" with a single Triton kernel (#58008)

Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

tests/v1/attention/test_rocm_glm5next_sparse.py

@@ -8,6 +8,7 @@ import torch
 
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
+from vllm.utils.torch_utils import set_random_seed
 from vllm.v1.attention.backends.mla import rocm_aiter_mla_sparse as sparse_mod
 from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import (
     _use_rocm_sparse_triton,
@@ -26,6 +27,26 @@ def _store_sparse_kv_row_offset_kernel(slot_ptr, output_ptr, stride: tl.constexp
     tl.store(output_ptr, _sparse_kv_row_offset(slot, stride))
 
 
+def _fit_kpool_indices_reference(
+    token_indices: torch.Tensor, topk_tokens: int
+) -> torch.Tensor:
+    history = token_indices[:, :topk_tokens]
+    tail = token_indices[:, topk_tokens:]
+    valid_history = (history >= 0).sum(dim=1)
+    valid_tail = (tail >= 0).sum(dim=1)
+    keep_history = torch.minimum(valid_history, topk_tokens - valid_tail)
+
+    columns = torch.arange(topk_tokens, device=token_indices.device).unsqueeze(0)
+    tail_offsets = columns - keep_history.unsqueeze(1)
+    tail_values = torch.gather(
+        tail, 1, tail_offsets.clamp(min=0, max=tail.shape[1] - 1)
+    )
+    output = torch.where(columns < keep_history.unsqueeze(1), history, tail_values)
+    valid_output = columns < (keep_history + valid_tail).unsqueeze(1)
+    return torch.where(valid_output, output, -1)
+
+
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm required")
 def test_fit_kpool_indices_preserves_tail_and_best_history():
     token_indices = torch.tensor(
         [
@@ -34,6 +55,7 @@ def test_fit_kpool_indices_preserves_tail_and_best_history():
             [-1, -1, -1, -1, -1, -1, -1, -1],
         ],
         dtype=torch.int32,
+        device="cuda",
     )
 
     fitted = fit_kpool_indices_to_aiter(token_indices, topk_tokens=6)
@@ -45,6 +67,28 @@ def test_fit_kpool_indices_preserves_tail_and_best_history():
     ]
 
 
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm required")
+@pytest.mark.parametrize(
+    ("num_tokens", "topk_tokens", "tail_width"),
+    [(1, 4, 1), (7, 6, 2), (13, 128, 128), (3, 2048, 128)],
+)
+def test_fit_kpool_indices_matches_eager_reference(num_tokens, topk_tokens, tail_width):
+    """The Triton fit must reproduce the eager packing rule exactly."""
+    set_random_seed(topk_tokens + num_tokens)
+    width = topk_tokens + tail_width
+    token_indices = torch.randint(0, 4096, (num_tokens, width), dtype=torch.int32)
+    valid = torch.rand((num_tokens, width)) < 0.5
+    valid[0] = False
+    valid[-1] = True
+    token_indices = torch.where(valid, token_indices, -1).cuda()
+
+    fitted = fit_kpool_indices_to_aiter(token_indices, topk_tokens=topk_tokens)
+
+    torch.testing.assert_close(
+        fitted, _fit_kpool_indices_reference(token_indices, topk_tokens)
+    )
+
+
 def test_fit_kpool_indices_exact_width_is_noop():
     token_indices = torch.tensor([[3, 2, 1, -1]], dtype=torch.int32)
 

vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py

@@ -70,6 +70,50 @@ def _use_rocm_sparse_triton(
     )
 
 
+@triton.jit
+def _fit_kpool_indices_kernel(
+    token_indices_ptr,  # int32 [num_tokens, NUM_TOPK_TOKENS + TAIL_WIDTH]
+    out_ptr,  # int32 [num_tokens, NUM_TOPK_TOKENS]
+    ti_stride0,
+    ti_stride1,
+    out_stride0,
+    out_stride1,
+    NUM_TOPK_TOKENS: tl.constexpr,
+    TAIL_WIDTH: tl.constexpr,
+    BLOCK_T: tl.constexpr,
+    BLOCK_TAIL: tl.constexpr,
+):
+    row_ptr = token_indices_ptr + tl.program_id(0) * ti_stride0
+    columns = tl.arange(0, BLOCK_T)
+    column_mask = columns < NUM_TOPK_TOKENS
+    history = tl.load(row_ptr + columns * ti_stride1, mask=column_mask, other=-1)
+
+    tail_columns = tl.arange(0, BLOCK_TAIL)
+    tail = tl.load(
+        row_ptr + (NUM_TOPK_TOKENS + tail_columns) * ti_stride1,
+        mask=tail_columns < TAIL_WIDTH,
+        other=-1,
+    )
+
+    valid_history = tl.sum((history >= 0).to(tl.int32))
+    valid_tail = tl.sum((tail >= 0).to(tl.int32))
+    keep_history = tl.minimum(valid_history, NUM_TOPK_TOKENS - valid_tail)
+
+    tail_offsets = columns - keep_history
+    tail_values = tl.load(
+        row_ptr + (NUM_TOPK_TOKENS + tail_offsets) * ti_stride1,
+        mask=column_mask & (tail_offsets >= 0) & (tail_offsets < TAIL_WIDTH),
+        other=-1,
+    )
+    fitted = tl.where(columns < keep_history, history, tail_values)
+    fitted = tl.where(columns < keep_history + valid_tail, fitted, -1)
+    tl.store(
+        out_ptr + tl.program_id(0) * out_stride0 + columns * out_stride1,
+        fitted,
+        mask=column_mask,
+    )
+
+
 def fit_kpool_indices_to_aiter(
     token_indices: torch.Tensor, topk_tokens: int
 ) -> torch.Tensor:
@@ -79,20 +123,29 @@ def fit_kpool_indices_to_aiter(
     if token_indices.shape[1] == topk_tokens:
         return token_indices
 
-    history = token_indices[:, :topk_tokens]
-    tail = token_indices[:, topk_tokens:]
-    valid_history = (history >= 0).sum(dim=1)
-    valid_tail = (tail >= 0).sum(dim=1)
-    keep_history = torch.minimum(valid_history, topk_tokens - valid_tail)
-
-    columns = torch.arange(topk_tokens, device=token_indices.device).unsqueeze(0)
-    tail_offsets = columns - keep_history.unsqueeze(1)
-    tail_values = torch.gather(
-        tail, 1, tail_offsets.clamp(min=0, max=tail.shape[1] - 1)
+    num_tokens, width = token_indices.shape
+    tail_width = width - topk_tokens
+    output = torch.empty(
+        (num_tokens, topk_tokens),
+        dtype=token_indices.dtype,
+        device=token_indices.device,
     )
-    output = torch.where(columns < keep_history.unsqueeze(1), history, tail_values)
-    valid_output = columns < (keep_history + valid_tail).unsqueeze(1)
-    return torch.where(valid_output, output, -1)
+    if num_tokens == 0:
+        return output
+
+    _fit_kpool_indices_kernel[(num_tokens,)](
+        token_indices,
+        output,
+        token_indices.stride(0),
+        token_indices.stride(1),
+        output.stride(0),
+        output.stride(1),
+        NUM_TOPK_TOKENS=topk_tokens,
+        TAIL_WIDTH=tail_width,
+        BLOCK_T=triton.next_power_of_2(topk_tokens),
+        BLOCK_TAIL=triton.next_power_of_2(tail_width),
+    )
+    return output
 
 
 @triton.jit