@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 346 lines 10.1 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """ROCm kernel correctness tests for AITER unified attention.
4
5 Compares ``aiter.ops.triton.unified_attention`` against ``ref_paged_attn`` under
6 decode, prefill, and mixed batches with varied shapes.
7 """
8
9 from typing import Any, Literal
10
11 import pytest
12 import torch
13
14 from tests.kernels.attention.test_triton_unified_attention import ref_paged_attn
15 from vllm.platforms import current_platform
16 from vllm.utils.torch_utils import set_random_seed
17
18 _SKIP_NON_CDNA_2_PLUS = True
19 if current_platform.is_rocm():
20 from vllm.platforms.rocm import get_cdna_version
21
22 _SKIP_NON_CDNA_2_PLUS = get_cdna_version() < 2
23
24 pytestmark = [
25 pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific tests"),
26 pytest.mark.skipif(_SKIP_NON_CDNA_2_PLUS, reason="CDNA 2+ ROCm only"),
27 ]
28
29 NUM_Q_HEADS = 8
30 NUM_KV_HEADS = 8
31 HEAD_SIZES = [128, 256]
32 BLOCK_SIZES = [16, 64]
33 DTYPES = [torch.bfloat16, torch.float16]
34 FP8_DTYPE = current_platform.fp8_dtype()
35
36 # (query_len, kv_len) per sequence
37 MIXED_SEQ_LENS = [
38 [(1, 128), (5, 18), (129, 463)],
39 [(10, 256), (5, 64), (32, 128)],
40 [(1, 1024), (5, 18), (129, 1328)],
41 ]
42 DECODE_SEQ_LENS = [
43 [(1, 128), (1, 256), (1, 384), (1, 512)],
44 [(1, 1024), (1, 1536), (1, 2048)],
45 ]
46 PREFILL_SEQ_LENS = [
47 [(256, 256), (128, 512)],
48 [(64, 128), (32, 256), (16, 512)],
49 [(256, 1024), (128, 2048)],
50 ]
51
52 Fp8Variant = Literal["fp8_kv", "fp8_query", "fp8_query_kv"]
53
54 DEFAULT_ATOL, DEFAULT_RTOL = 1.5e-2, 1e-2
55 # fp8_kv: reference dequantizes the same KV, so error is bf16-scale (~2e-3).
56 # fp8_query / fp8_query_kv: query and/or KV e4m3 rounding dominates (~0.1-0.4).
57 FP8_TOL_BY_VARIANT: dict[Fp8Variant, tuple[float, float]] = {
58 "fp8_kv": (2.5e-3, 1e-2),
59 "fp8_query": (1.5e-1, 1.5e-1),
60 "fp8_query_kv": (1.5e-1, 1.5e-1),
61 }
62 # Non-unity scale so q_descale handling is exercised explicitly.
63 Q_SCALE = 0.75
64 K_SCALE, V_SCALE = 0.5, 0.25
65
66 FP8_VARIANTS = [
67 pytest.param("fp8_kv", id="fp8_kv"),
68 pytest.param("fp8_query", id="fp8_query"),
69 pytest.param("fp8_query_kv", id="fp8_query_kv"),
70 ]
71
72 FP8_SEQ_LENS = [
73 MIXED_SEQ_LENS[0],
74 DECODE_SEQ_LENS[0],
75 DECODE_SEQ_LENS[1],
76 PREFILL_SEQ_LENS[0],
77 PREFILL_SEQ_LENS[2],
78 ]
79
80
81 def _require_aiter() -> None:
82 from vllm._aiter_ops import is_aiter_found_and_supported
83
84 if not is_aiter_found_and_supported():
85 pytest.skip("aiter is required on supported ROCm hardware for this test")
86
87
88 def _make_case(
89 *,
90 seq_lens: list[tuple[int, int]],
91 head_size: int,
92 block_size: int,
93 dtype: torch.dtype,
94 num_blocks: int = 2048,
95 kv_cache_dtype: torch.dtype | None = None,
96 k_scale: float = 1.0,
97 v_scale: float = 1.0,
98 q_dtype: torch.dtype | None = None,
99 q_scale: float = Q_SCALE,
100 ) -> dict[str, Any]:
101 torch.set_default_device("cuda")
102
103 query_lens = [q for q, _ in seq_lens]
104 kv_lens = [k for _, k in seq_lens]
105 num_seqs = len(seq_lens)
106 max_query_len = max(query_lens)
107 max_kv_len = max(kv_lens)
108 scale = head_size**-0.5
109
110 query = torch.randn(sum(query_lens), NUM_Q_HEADS, head_size, dtype=dtype)
111 if kv_cache_dtype is None:
112 key_cache = torch.randn(
113 num_blocks, block_size, NUM_KV_HEADS, head_size, dtype=dtype
114 )
115 value_cache = torch.randn_like(key_cache)
116 else:
117 key_cache = torch.clamp(
118 torch.randn(num_blocks, block_size, NUM_KV_HEADS, head_size),
119 -1.0,
120 1.0,
121 ).to(kv_cache_dtype)
122 value_cache = torch.clamp(
123 torch.randn(num_blocks, block_size, NUM_KV_HEADS, head_size),
124 -1.0,
125 1.0,
126 ).to(kv_cache_dtype)
127
128 cu_seqlens_q = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
129 dim=0, dtype=torch.int32
130 )
131 seq_lens_tensor = torch.tensor(kv_lens, dtype=torch.int32)
132
133 max_num_blocks = (max_kv_len + block_size - 1) // block_size
134 block_tables = torch.randint(
135 0, num_blocks, (num_seqs, max_num_blocks), dtype=torch.int32
136 )
137
138 descale_shape = (num_seqs, NUM_KV_HEADS)
139 k_descale = torch.full(descale_shape, k_scale, dtype=torch.float32, device="cuda")
140 v_descale = torch.full(descale_shape, v_scale, dtype=torch.float32, device="cuda")
141
142 kernel_query = query
143 q_descale = None
144 if q_dtype is not None:
145 q_descale = torch.tensor(q_scale, dtype=torch.float32, device="cuda")
146 kernel_query = (query / q_scale).to(q_dtype)
147
148 return {
149 "query": query,
150 "kernel_query": kernel_query,
151 "key_cache": key_cache,
152 "value_cache": value_cache,
153 "block_tables": block_tables,
154 "query_lens": query_lens,
155 "kv_lens": kv_lens,
156 "seq_lens_tensor": seq_lens_tensor,
157 "cu_seqlens_q": cu_seqlens_q,
158 "q_descale": q_descale,
159 "k_descale": k_descale,
160 "v_descale": v_descale,
161 "scale": scale,
162 "max_query_len": max_query_len,
163 "max_kv_len": max_kv_len,
164 "query_dtype": dtype,
165 "k_scale": k_scale,
166 "v_scale": v_scale,
167 }
168
169
170 def _make_fp8_case(
171 *,
172 seq_lens: list[tuple[int, int]],
173 head_size: int,
174 block_size: int,
175 variant: Fp8Variant,
176 ) -> dict[str, Any]:
177 use_fp8_kv = variant in ("fp8_kv", "fp8_query_kv")
178 use_fp8_query = variant in ("fp8_query", "fp8_query_kv")
179 return _make_case(
180 seq_lens=seq_lens,
181 head_size=head_size,
182 block_size=block_size,
183 dtype=torch.bfloat16,
184 kv_cache_dtype=FP8_DTYPE if use_fp8_kv else None,
185 k_scale=K_SCALE if use_fp8_kv else 1.0,
186 v_scale=V_SCALE if use_fp8_kv else 1.0,
187 q_dtype=FP8_DTYPE if use_fp8_query else None,
188 )
189
190
191 def _run_aiter_unified_attention(case: dict[str, Any]) -> torch.Tensor:
192 from aiter.ops.triton.unified_attention import unified_attention
193
194 kernel_query = case["kernel_query"]
195 # Kernel writes high-precision output even when Q is FP8 (matches vLLM usage).
196 output = torch.empty_like(case["query"])
197 unified_attention(
198 q=kernel_query,
199 k=case["key_cache"],
200 v=case["value_cache"],
201 out=output,
202 cu_seqlens_q=case["cu_seqlens_q"],
203 max_seqlen_q=case["max_query_len"],
204 seqused_k=case["seq_lens_tensor"],
205 max_seqlen_k=case["max_kv_len"],
206 softmax_scale=case["scale"],
207 causal=True,
208 alibi_slopes=None,
209 window_size=(-1, -1),
210 block_table=case["block_tables"],
211 softcap=0,
212 q_descale=case["q_descale"],
213 k_descale=case["k_descale"],
214 v_descale=case["v_descale"],
215 sinks=None,
216 output_scale=None,
217 )
218 return output
219
220
221 def _ref_output(case: dict[str, Any]) -> torch.Tensor:
222 key_cache = case["key_cache"]
223 value_cache = case["value_cache"]
224 if key_cache.dtype != case["query_dtype"]:
225 key_cache = key_cache.to(case["query_dtype"]) * case["k_scale"]
226 value_cache = value_cache.to(case["query_dtype"]) * case["v_scale"]
227
228 return ref_paged_attn(
229 query=case["query"],
230 key_cache=key_cache,
231 value_cache=value_cache,
232 query_lens=case["query_lens"],
233 kv_lens=case["kv_lens"],
234 block_tables=case["block_tables"],
235 scale=case["scale"],
236 )
237
238
239 def _assert_matches_reference(
240 case: dict[str, Any],
241 *,
242 atol: float = DEFAULT_ATOL,
243 rtol: float = DEFAULT_RTOL,
244 ) -> None:
245 output = _run_aiter_unified_attention(case)
246 output_ref = _ref_output(case)
247 torch.testing.assert_close(output, output_ref, atol=atol, rtol=rtol)
248
249
250 @pytest.mark.parametrize("seq_lens", MIXED_SEQ_LENS)
251 @pytest.mark.parametrize("head_size", HEAD_SIZES)
252 @pytest.mark.parametrize("block_size", BLOCK_SIZES)
253 @pytest.mark.parametrize("dtype", DTYPES)
254 @torch.inference_mode()
255 def test_aiter_unified_attn_mixed_batch(
256 seq_lens: list[tuple[int, int]],
257 head_size: int,
258 block_size: int,
259 dtype: torch.dtype,
260 ) -> None:
261 """Decode + prefill sequences in one batch (native dtypes)."""
262 _require_aiter()
263 set_random_seed(0)
264
265 case = _make_case(
266 seq_lens=seq_lens,
267 head_size=head_size,
268 block_size=block_size,
269 dtype=dtype,
270 )
271 _assert_matches_reference(case)
272
273
274 @pytest.mark.parametrize("seq_lens", DECODE_SEQ_LENS)
275 @pytest.mark.parametrize("head_size", HEAD_SIZES)
276 @pytest.mark.parametrize("block_size", BLOCK_SIZES)
277 @pytest.mark.parametrize("dtype", [torch.bfloat16])
278 @torch.inference_mode()
279 def test_aiter_unified_attn_decode(
280 seq_lens: list[tuple[int, int]],
281 head_size: int,
282 block_size: int,
283 dtype: torch.dtype,
284 ) -> None:
285 """Single-token decode (native dtypes)."""
286 _require_aiter()
287 set_random_seed(0)
288
289 case = _make_case(
290 seq_lens=seq_lens,
291 head_size=head_size,
292 block_size=block_size,
293 dtype=dtype,
294 )
295 _assert_matches_reference(case)
296
297
298 @pytest.mark.parametrize("seq_lens", PREFILL_SEQ_LENS)
299 @pytest.mark.parametrize("head_size", [128])
300 @pytest.mark.parametrize("block_size", [16])
301 @torch.inference_mode()
302 def test_aiter_unified_attn_prefill(
303 seq_lens: list[tuple[int, int]],
304 head_size: int,
305 block_size: int,
306 ) -> None:
307 """Prefill-only batches with query_len > 1 (native dtypes)."""
308 _require_aiter()
309 set_random_seed(0)
310
311 case = _make_case(
312 seq_lens=seq_lens,
313 head_size=head_size,
314 block_size=block_size,
315 dtype=torch.bfloat16,
316 )
317 _assert_matches_reference(case)
318
319
320 @pytest.mark.skipif(
321 not current_platform.supports_fp8(),
322 reason="FP8 not supported on this hardware",
323 )
324 @pytest.mark.parametrize("variant", FP8_VARIANTS)
325 @pytest.mark.parametrize("seq_lens", FP8_SEQ_LENS)
326 @pytest.mark.parametrize("head_size", [128])
327 @pytest.mark.parametrize("block_size", [16, 64])
328 @torch.inference_mode()
329 def test_aiter_unified_attn_fp8(
330 variant: Fp8Variant,
331 seq_lens: list[tuple[int, int]],
332 head_size: int,
333 block_size: int,
334 ) -> None:
335 """FP8 KV cache, FP8 query, or both; compared at bf16 reference precision."""
336 _require_aiter()
337 set_random_seed(0)
338
339 case = _make_fp8_case(
340 seq_lens=seq_lens,
341 head_size=head_size,
342 block_size=block_size,
343 variant=variant,
344 )
345 atol, rtol = FP8_TOL_BY_VARIANT[variant]
346 _assert_matches_reference(case, atol=atol, rtol=rtol)