[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>
Simon Danielsson committed
Sep 30, 2026 at 16:45 UTC
2df122e65e87dc12870ce5f0180b49f012e7a53b
2 files changed
+109
-12
tests/v1/attention/test_rocm_glm5next_sparse.py
+44
@@ -8,6 +8,7 @@ import torch
8
9
from vllm.platforms import current_platform
10
from vllm.triton_utils import tl, triton
11
+from vllm.utils.torch_utils import set_random_seed
12
from vllm.v1.attention.backends.mla import rocm_aiter_mla_sparse as sparse_mod
13
from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import (
14
_use_rocm_sparse_triton,
@@ -26,6 +27,26 @@ def _store_sparse_kv_row_offset_kernel(slot_ptr, output_ptr, stride: tl.constexp
27
tl.store(output_ptr, _sparse_kv_row_offset(slot, stride))
28
29
30
+def _fit_kpool_indices_reference(
31
+ token_indices: torch.Tensor, topk_tokens: int
32
+) -> torch.Tensor:
33
+ history = token_indices[:, :topk_tokens]
34
+ tail = token_indices[:, topk_tokens:]
35
+ valid_history = (history >= 0).sum(dim=1)
36
+ valid_tail = (tail >= 0).sum(dim=1)
37
+ keep_history = torch.minimum(valid_history, topk_tokens - valid_tail)
38
+
39
+ columns = torch.arange(topk_tokens, device=token_indices.device).unsqueeze(0)
40
+ tail_offsets = columns - keep_history.unsqueeze(1)
41
+ tail_values = torch.gather(
42
+ tail, 1, tail_offsets.clamp(min=0, max=tail.shape[1] - 1)
43
+ )
44
+ output = torch.where(columns < keep_history.unsqueeze(1), history, tail_values)
45
+ valid_output = columns < (keep_history + valid_tail).unsqueeze(1)
46
+ return torch.where(valid_output, output, -1)
47
+
48
+
49
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm required")
50
def test_fit_kpool_indices_preserves_tail_and_best_history():
51
token_indices = torch.tensor(
52
[
@@ -34,6 +55,7 @@ def test_fit_kpool_indices_preserves_tail_and_best_history():
55
[-1, -1, -1, -1, -1, -1, -1, -1],
56
],
57
dtype=torch.int32,
58
+ device="cuda",
59
)
60
61
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():
67
]
68
69
70
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm required")
71
+@pytest.mark.parametrize(
72
+ ("num_tokens", "topk_tokens", "tail_width"),
73
+ [(1, 4, 1), (7, 6, 2), (13, 128, 128), (3, 2048, 128)],
74
+)
75
+def test_fit_kpool_indices_matches_eager_reference(num_tokens, topk_tokens, tail_width):
76
+ """The Triton fit must reproduce the eager packing rule exactly."""
77
+ set_random_seed(topk_tokens + num_tokens)
78
+ width = topk_tokens + tail_width
79
+ token_indices = torch.randint(0, 4096, (num_tokens, width), dtype=torch.int32)
80
+ valid = torch.rand((num_tokens, width)) < 0.5
81
+ valid[0] = False
82
+ valid[-1] = True
83
+ token_indices = torch.where(valid, token_indices, -1).cuda()
84
+
85
+ fitted = fit_kpool_indices_to_aiter(token_indices, topk_tokens=topk_tokens)
86
+
87
+ torch.testing.assert_close(
88
+ fitted, _fit_kpool_indices_reference(token_indices, topk_tokens)
89
+ )
90
+
91
+
92
def test_fit_kpool_indices_exact_width_is_noop():
93
token_indices = torch.tensor([[3, 2, 1, -1]], dtype=torch.int32)
94
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+65
-12
@@ -70,6 +70,50 @@ def _use_rocm_sparse_triton(
70
)
71
72
73
+@triton.jit
74
+def _fit_kpool_indices_kernel(
75
+ token_indices_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS + TAIL_WIDTH]
76
+ out_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS]
77
+ ti_stride0,
78
+ ti_stride1,
79
+ out_stride0,
80
+ out_stride1,
81
+ NUM_TOPK_TOKENS: tl.constexpr,
82
+ TAIL_WIDTH: tl.constexpr,
83
+ BLOCK_T: tl.constexpr,
84
+ BLOCK_TAIL: tl.constexpr,
85
+):
86
+ row_ptr = token_indices_ptr + tl.program_id(0) * ti_stride0
87
+ columns = tl.arange(0, BLOCK_T)
88
+ column_mask = columns < NUM_TOPK_TOKENS
89
+ history = tl.load(row_ptr + columns * ti_stride1, mask=column_mask, other=-1)
90
+
91
+ tail_columns = tl.arange(0, BLOCK_TAIL)
92
+ tail = tl.load(
93
+ row_ptr + (NUM_TOPK_TOKENS + tail_columns) * ti_stride1,
94
+ mask=tail_columns < TAIL_WIDTH,
95
+ other=-1,
96
+ )
97
+
98
+ valid_history = tl.sum((history >= 0).to(tl.int32))
99
+ valid_tail = tl.sum((tail >= 0).to(tl.int32))
100
+ keep_history = tl.minimum(valid_history, NUM_TOPK_TOKENS - valid_tail)
101
+
102
+ tail_offsets = columns - keep_history
103
+ tail_values = tl.load(
104
+ row_ptr + (NUM_TOPK_TOKENS + tail_offsets) * ti_stride1,
105
+ mask=column_mask & (tail_offsets >= 0) & (tail_offsets < TAIL_WIDTH),
106
+ other=-1,
107
+ )
108
+ fitted = tl.where(columns < keep_history, history, tail_values)
109
+ fitted = tl.where(columns < keep_history + valid_tail, fitted, -1)
110
+ tl.store(
111
+ out_ptr + tl.program_id(0) * out_stride0 + columns * out_stride1,
112
+ fitted,
113
+ mask=column_mask,
114
+ )
115
+
116
+
117
def fit_kpool_indices_to_aiter(
118
token_indices: torch.Tensor, topk_tokens: int
119
) -> torch.Tensor:
@@ -79,20 +123,29 @@ def fit_kpool_indices_to_aiter(
123
if token_indices.shape[1] == topk_tokens:
124
return token_indices
125
82
- history = token_indices[:, :topk_tokens]
83
- tail = token_indices[:, topk_tokens:]
84
- valid_history = (history >= 0).sum(dim=1)
85
- valid_tail = (tail >= 0).sum(dim=1)
86
- keep_history = torch.minimum(valid_history, topk_tokens - valid_tail)
126
+ num_tokens, width = token_indices.shape
127
+ tail_width = width - topk_tokens
128
+ output = torch.empty(
129
+ (num_tokens, topk_tokens),
130
+ dtype=token_indices.dtype,
131
+ device=token_indices.device,
132
+ )
133
+ if num_tokens == 0:
134
+ return output
135
88
- columns = torch.arange(topk_tokens, device=token_indices.device).unsqueeze(0)
89
- tail_offsets = columns - keep_history.unsqueeze(1)
90
- tail_values = torch.gather(
91
- tail, 1, tail_offsets.clamp(min=0, max=tail.shape[1] - 1)
136
+ _fit_kpool_indices_kernel[(num_tokens,)](
137
+ token_indices,
138
+ output,
139
+ token_indices.stride(0),
140
+ token_indices.stride(1),
141
+ output.stride(0),
142
+ output.stride(1),
143
+ NUM_TOPK_TOKENS=topk_tokens,
144
+ TAIL_WIDTH=tail_width,
145
+ BLOCK_T=triton.next_power_of_2(topk_tokens),
146
+ BLOCK_TAIL=triton.next_power_of_2(tail_width),
147
)
93
- output = torch.where(columns < keep_history.unsqueeze(1), history, tail_values)
94
- valid_output = columns < (keep_history + valid_tail).unsqueeze(1)
95
- return torch.where(valid_output, output, -1)
148
+ return output
149
150
151
@triton.jit