[ROCm][DSV4.1][Perf] Stride the DSA decode candidate mask over the live context (#56628)
Signed-off-by: Fangzhou Ai <fangzhou.ai@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Fangzhou Ai committed
Sep 14, 2026 at 16:47 UTC
dabc4362b47ad2665b0802b0b42b13f660efe634
2 files changed
+344
-2
tests/kernels/attention/test_rocm_aiter_candidate_mask.py
new
+226
@@ -0,0 +1,226 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""ROCm correctness tests for the strided DSA candidate mask.
4
+
5
+``_apply_candidate_mask_strided`` is a decode-path variant of the shared
6
+``apply_candidate_mask``. It differs in two ways that need pinning down:
7
+
8
+* it launches a fixed number of programs that stride the column axis, rather
9
+ than one program per 1024-column tile, so at ``max_model_len`` widths the
10
+ grid no longer scales with a workspace that is mostly padding;
11
+* it stops at each row's ``end`` instead of sanitizing the whole width.
12
+
13
+The second is only sound because every consumer bounds its scan by the same
14
+ends -- ``top_k_per_row_decode`` takes ``rowEnd`` from ``seq_lens``. So the
15
+contract is agreement with the shared kernel over ``[0, end)`` and nothing
16
+beyond it, which is what these tests assert rather than comparing whole rows.
17
+"""
18
+
19
+import pytest
20
+import torch
21
+
22
+from vllm.platforms import current_platform
23
+
24
+pytestmark = pytest.mark.skipif(
25
+ not current_platform.is_rocm(), reason="ROCm-specific tests"
26
+)
27
+
28
+BLOCK_SIZE = 8
29
+TOPK_BLOCKS = 64
30
+# Above 128 tiles the grid saturates and programs start striding; below it the
31
+# min() clamp applies instead. Both regimes need covering, so the widths here
32
+# deliberately straddle 128 * 1024.
33
+NARROW = 8192
34
+WIDE = 262144
35
+
36
+
37
+def _inputs(rows, width, ends, starts, *, block_size, col_stride, dtype):
38
+ """Logits, bounds and candidate blocks shared by both kernels."""
39
+ torch.manual_seed(0)
40
+ # A column stride > 1 keeps the kernels honest about stride_col; the decode
41
+ # caller slices a workspace, so contiguous rows are not guaranteed.
42
+ logits = torch.randn(rows, width * col_stride, device="cuda")[:, ::col_stride]
43
+ logits[0, : min(16, width)] = 0.0
44
+ if rows > 2:
45
+ logits[2, 5] = float("nan")
46
+ row_ke = torch.tensor(ends, device="cuda", dtype=dtype)
47
+ row_ks = torch.tensor(starts, device="cuda", dtype=dtype) if starts else None
48
+ nblocks = (width + block_size - 1) // block_size
49
+ candidates = torch.randint(
50
+ 0, nblocks, (rows, TOPK_BLOCKS), device="cuda", dtype=torch.int32
51
+ )
52
+ # -1 is the "no block" sentinel the selector emits for short rows.
53
+ candidates[:, ::7] = -1
54
+ return logits, row_ks, row_ke, candidates
55
+
56
+
57
+def _run_both(logits, row_ks, row_ke, candidates, block_size, row_repeat):
58
+ from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
59
+ apply_candidate_mask,
60
+ )
61
+ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
62
+ _apply_candidate_mask_strided,
63
+ )
64
+
65
+ ref = logits.clone()
66
+ got = logits.clone()
67
+ apply_candidate_mask(ref, row_ks, row_ke, candidates, block_size, row_repeat)
68
+ _apply_candidate_mask_strided(
69
+ got, row_ks, row_ke, candidates, block_size, row_repeat
70
+ )
71
+ return ref, got
72
+
73
+
74
+def _assert_agrees_within_ends(ref, got, row_ke, row_repeat):
75
+ """Compare only the region a consumer can read."""
76
+ for row in range(ref.shape[0]):
77
+ end = int(row_ke[row // row_repeat])
78
+ if end <= 0:
79
+ continue
80
+ torch.testing.assert_close(
81
+ got[row, :end],
82
+ ref[row, :end],
83
+ rtol=0,
84
+ atol=0,
85
+ equal_nan=True,
86
+ msg=lambda m, r=row, e=end: f"row {r} differs within [0, {e}): {m}",
87
+ )
88
+
89
+
90
+@pytest.mark.parametrize("width", [NARROW, WIDE])
91
+@pytest.mark.parametrize("col_stride", [1, 2])
92
+def test_matches_shared_kernel_ragged_ends(width, col_stride):
93
+ """Ends on and around tile boundaries, plus the empty and full rows."""
94
+ ends = [0, 1, 1023, 1024, 1025, width - 1, width, width // 2]
95
+ logits, row_ks, row_ke, candidates = _inputs(
96
+ len(ends),
97
+ width,
98
+ ends,
99
+ None,
100
+ block_size=BLOCK_SIZE,
101
+ col_stride=col_stride,
102
+ dtype=torch.int32,
103
+ )
104
+ ref, got = _run_both(logits, row_ks, row_ke, candidates, BLOCK_SIZE, 1)
105
+ _assert_agrees_within_ends(ref, got, row_ke, 1)
106
+
107
+
108
+@pytest.mark.parametrize("block_size", [8, 16])
109
+def test_matches_shared_kernel_with_starts(block_size):
110
+ """Non-zero starts: the decode caller passes zeros, prefill does not."""
111
+ ends = [0, 64, 4096, 8192]
112
+ starts = [0, 7, 1024, 3]
113
+ logits, row_ks, row_ke, candidates = _inputs(
114
+ len(ends),
115
+ NARROW,
116
+ ends,
117
+ starts,
118
+ block_size=block_size,
119
+ col_stride=1,
120
+ dtype=torch.int64,
121
+ )
122
+ ref, got = _run_both(logits, row_ks, row_ke, candidates, block_size, 1)
123
+ _assert_agrees_within_ends(ref, got, row_ke, 1)
124
+
125
+
126
+def test_matches_shared_kernel_row_repeat():
127
+ """Speculative decode: next_n logit rows share one bound."""
128
+ row_repeat = 3
129
+ ends = [0, 1025, WIDE]
130
+ logits, row_ks, row_ke, candidates = _inputs(
131
+ len(ends) * row_repeat,
132
+ WIDE,
133
+ ends,
134
+ None,
135
+ block_size=BLOCK_SIZE,
136
+ col_stride=1,
137
+ dtype=torch.int32,
138
+ )
139
+ ref, got = _run_both(logits, row_ks, row_ke, candidates, BLOCK_SIZE, row_repeat)
140
+ _assert_agrees_within_ends(ref, got, row_ke, row_repeat)
141
+
142
+
143
+def test_leaves_columns_past_end_untouched():
144
+ """The deliberate difference from the shared kernel, pinned explicitly.
145
+
146
+ Skipping the tail is where the speedup comes from, so a change that
147
+ quietly restores full-width sanitizing should fail here rather than just
148
+ get slower. Work stops at the tile boundary containing ``end``, not at
149
+ ``end`` itself, so that is the bound asserted.
150
+ """
151
+ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
152
+ _MASK_TILE,
153
+ _apply_candidate_mask_strided,
154
+ )
155
+
156
+ ends = [0, 1, 1025, WIDE // 2]
157
+ logits, row_ks, row_ke, candidates = _inputs(
158
+ len(ends),
159
+ WIDE,
160
+ ends,
161
+ None,
162
+ block_size=BLOCK_SIZE,
163
+ col_stride=1,
164
+ dtype=torch.int32,
165
+ )
166
+ sentinel = 1.5
167
+ logits.fill_(sentinel)
168
+ _apply_candidate_mask_strided(logits, row_ks, row_ke, candidates, BLOCK_SIZE, 1)
169
+ for row, end in enumerate(ends):
170
+ touched = min(-(-end // _MASK_TILE) * _MASK_TILE, WIDE)
171
+ tail = logits[row, touched:]
172
+ assert torch.equal(tail, torch.full_like(tail, sentinel)), (
173
+ f"row {row} (end {end}) was written past column {touched}"
174
+ )
175
+ # And the bound is tight: a row that ends early must not have cost a
176
+ # full-width pass, which is the whole point of the kernel.
177
+ assert touched <= end + _MASK_TILE
178
+
179
+
180
+def test_cudagraph_replay_tracks_changing_ends():
181
+ """The grid is static so a FULL capture stays valid when ends change.
182
+
183
+ Capturing bakes the grid, so a data-dependent program count would freeze
184
+ whatever lengths were resident at capture. Replaying against longer and
185
+ shorter contexts than the captured ones is the property that matters.
186
+ """
187
+ from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
188
+ apply_candidate_mask,
189
+ )
190
+ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
191
+ _apply_candidate_mask_strided,
192
+ )
193
+
194
+ rows = 4
195
+ captured_ends = [1024, 4096, 8192, 2048]
196
+ logits, row_ks, row_ke, candidates = _inputs(
197
+ rows,
198
+ WIDE,
199
+ captured_ends,
200
+ None,
201
+ block_size=BLOCK_SIZE,
202
+ col_stride=1,
203
+ dtype=torch.int32,
204
+ )
205
+ base = logits.clone()
206
+
207
+ # Warm up outside the capture: Triton compiles on first call, and autotune
208
+ # launches would otherwise happen inside the graph.
209
+ _apply_candidate_mask_strided(
210
+ logits.clone(), row_ks, row_ke, candidates, BLOCK_SIZE, 1
211
+ )
212
+ torch.accelerator.synchronize()
213
+
214
+ graph = torch.cuda.CUDAGraph()
215
+ with torch.cuda.graph(graph):
216
+ _apply_candidate_mask_strided(logits, row_ks, row_ke, candidates, BLOCK_SIZE, 1)
217
+
218
+ for replay_ends in ([1, 131072, WIDE, 0], [WIDE, 1025, 3, 65536]):
219
+ row_ke.copy_(torch.tensor(replay_ends, device="cuda", dtype=row_ke.dtype))
220
+ logits.copy_(base)
221
+ graph.replay()
222
+ torch.accelerator.synchronize()
223
+
224
+ ref = base.clone()
225
+ apply_candidate_mask(ref, row_ks, row_ke, candidates, BLOCK_SIZE, 1)
226
+ _assert_agrees_within_ends(ref, logits, row_ke, 1)
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+118
-2
@@ -819,6 +819,123 @@ def rocm_fp8_mqa_logits(
819
return fp8_mqa_logits_torch(q, kv, weights, cu_seqlen_ks, cu_seqlen_ke)
820
821
822
+# Programs along the column axis of the decode mask grid. The shared kernel in
823
+# model_executor/kernels/attention/dsa/candidate_blocks.py launches one program
824
+# per 1024-column tile, which is 1024 of them per row at max_model_len 1048576
825
+# with only the first few doing work. This is a fixed count that strides
826
+# instead, measured on gfx950; it is not a portable choice, which is why this
827
+# variant lives here rather than replacing the shared one.
828
+_MASK_GRID_COLS = 128
829
+# Columns each program handles per iteration. Work stops at the tile boundary
830
+# containing a row's end rather than at the end itself; the overshoot is
831
+# masked the same way the shared kernel masks it, so it costs a tile and
832
+# changes nothing.
833
+_MASK_TILE = 1024
834
+
835
+
836
+@triton.jit(do_not_specialize=["width", "nblocks"])
837
+def _mask_candidates_strided_kernel(
838
+ logits,
839
+ starts,
840
+ ends,
841
+ flags,
842
+ stride_row,
843
+ stride_col,
844
+ stride_start,
845
+ stride_end,
846
+ width,
847
+ nblocks,
848
+ BLOCK_SIZE: tl.constexpr,
849
+ HAS_STARTS: tl.constexpr,
850
+ ROW_REPEAT: tl.constexpr,
851
+ TILE: tl.constexpr,
852
+):
853
+ row = tl.program_id(0).to(tl.int64)
854
+ start = tl.load(starts + row // ROW_REPEAT * stride_start) if HAS_STARTS else 0
855
+ end = tl.load(ends + row // ROW_REPEAT * stride_end)
856
+ edge = tl.load(flags + row * (nblocks + 1) + nblocks)
857
+ step = tl.num_programs(1) * TILE
858
+ tile_start = tl.program_id(1) * TILE
859
+ # The shared kernel also sanitizes columns at or past `end`. Nothing reads
860
+ # them: top_k_per_row_decode bounds its scan by the same row ends passed
861
+ # here, and persistent_topk/cooperative_topk clamp by the same lengths.
862
+ # Stopping at `end` is what makes the cost track the live context rather
863
+ # than max_model_len.
864
+ while tile_start < end:
865
+ cols = tile_start + tl.arange(0, TILE)
866
+ valid = (cols >= start) & (cols < end) & (cols < width)
867
+ block = (cols - start) // BLOCK_SIZE
868
+ keep = tl.load(flags + row * (nblocks + 1) + block, valid, other=0)
869
+ keep = (keep != 0) | ((cols == width - 1) & (edge != 0))
870
+ tl.store(
871
+ logits + row * stride_row + cols * stride_col,
872
+ -float("inf"),
873
+ (cols < width) & ~(valid & keep),
874
+ )
875
+ tile_start += step
876
+
877
+
878
+def _apply_candidate_mask_strided(
879
+ logits: torch.Tensor,
880
+ row_ks: torch.Tensor | None,
881
+ row_ke: torch.Tensor,
882
+ candidate_blocks: torch.Tensor,
883
+ block_size: int,
884
+ row_repeat: int = 1,
885
+) -> None:
886
+ """ROCm decode variant of ``apply_candidate_mask``.
887
+
888
+ Same masking semantics over ``[0, end)``, but the grid is sized by a fixed
889
+ program count rather than by the logits width. Only worth using where the
890
+ width is the ``max_model_len`` workspace and the live context is far
891
+ shorter, i.e. the paged decode path below; the prefill chunks pass
892
+ chunk-sized logits and stay on the shared kernel.
893
+ """
894
+ from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
895
+ _candidate_flags_kernel,
896
+ )
897
+
898
+ rows, width = logits.shape
899
+ if not rows or not width:
900
+ return
901
+ nblocks = triton.cdiv(width, block_size)
902
+ flags = torch.empty((rows, nblocks + 1), device=logits.device, dtype=torch.uint8)
903
+ start_stride = row_ks.stride(0) if row_ks is not None else 0
904
+ _candidate_flags_kernel[(rows,)](
905
+ candidate_blocks,
906
+ row_ks,
907
+ flags,
908
+ *candidate_blocks.stride(),
909
+ start_stride,
910
+ width,
911
+ nblocks,
912
+ block_size,
913
+ candidate_blocks.shape[1],
914
+ row_ks is not None,
915
+ row_repeat,
916
+ )
917
+ # Derived from width, which is a tensor shape, so the grid stays static and
918
+ # a FULL cudagraph capture remains valid across replays; only the loop trip
919
+ # count inside the kernel is data-dependent. The min keeps narrow widths
920
+ # from launching programs that would only fall through.
921
+ grid_cols = min(_MASK_GRID_COLS, triton.cdiv(width, _MASK_TILE))
922
+ _mask_candidates_strided_kernel[(rows, grid_cols)](
923
+ logits,
924
+ row_ks,
925
+ row_ke,
926
+ flags,
927
+ *logits.stride(),
928
+ start_stride,
929
+ row_ke.stride(0),
930
+ width,
931
+ nblocks,
932
+ block_size,
933
+ row_ks is not None,
934
+ row_repeat,
935
+ _MASK_TILE,
936
+ )
937
+
938
+
939
def _max_decode_logits_rows(num_batched_tokens: int) -> int:
940
"""Upper bound on decode rows the paged-MQA logits buffer can ever hold.
941
@@ -1102,7 +1219,6 @@ def rocm_aiter_sparse_attn_indexer(
1219
1220
if candidate_blocks is not None:
1221
from vllm.model_executor.layers.sparse_attn_indexer import (
1105
- _apply_candidate_mask,
1222
_select_candidate_blocks,
1223
)
1224
@@ -1123,7 +1239,7 @@ def rocm_aiter_sparse_attn_indexer(
1239
decode_candidates,
1240
)
1241
else:
1126
- _apply_candidate_mask(
1242
+ _apply_candidate_mask_strided(
1243
logits,
1244
row_starts,
1245
visible,