@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 226 lines 7.96 KB
Raw
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)