A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 1, 22:30 UTC
[ROCm][CI] Add test for Aiter unified attn kernel (#44436)
Signed-off-by: Andreas Karatzas <akaratza@amd.com> Signed-off-by: Divakar Verma <divakar.verma@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Divakar Verma committed
Jun 4, 2026 at 12:15 UTC
06f94633e786f8e871c927bef3d93aced8e26be4
1 file changed
+339
tests/kernels/attention/test_rocm_aiter_unified_attn.py
new
+339
@@ -0,0 +1,339 @@
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_MI3XX = True
19
+if current_platform.is_rocm():
20
+ from vllm.platforms.rocm import on_mi3xx
21
+
22
+ _SKIP_NON_MI3XX = not on_mi3xx()
23
+
24
+pytestmark = [
25
+ pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific tests"),
26
+ pytest.mark.skipif(_SKIP_NON_MI3XX, reason="MI300/MI350 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
+DEFAULT_ATOL, DEFAULT_RTOL = 1.5e-2, 1e-2
53
+FP8_ATOL, FP8_RTOL = 1.5e-1, 1.5e-1
54
+# Non-unity scale so q_descale handling is exercised explicitly.
55
+Q_SCALE = 0.75
56
+K_SCALE, V_SCALE = 0.5, 0.25
57
+
58
+Fp8Variant = Literal["fp8_kv", "fp8_query", "fp8_query_kv"]
59
+
60
+FP8_VARIANTS = [
61
+ pytest.param("fp8_kv", id="fp8_kv"),
62
+ pytest.param("fp8_query", id="fp8_query"),
63
+ pytest.param("fp8_query_kv", id="fp8_query_kv"),
64
+]
65
+
66
+FP8_SEQ_LENS = [
67
+ MIXED_SEQ_LENS[0],
68
+ DECODE_SEQ_LENS[0],
69
+ DECODE_SEQ_LENS[1],
70
+ PREFILL_SEQ_LENS[0],
71
+ PREFILL_SEQ_LENS[2],
72
+]
73
+
74
+
75
+def _require_aiter() -> None:
76
+ from vllm._aiter_ops import is_aiter_found_and_supported
77
+
78
+ if not is_aiter_found_and_supported():
79
+ pytest.skip("aiter is required on supported ROCm hardware for this test")
80
+
81
+
82
+def _make_case(
83
+ *,
84
+ seq_lens: list[tuple[int, int]],
85
+ head_size: int,
86
+ block_size: int,
87
+ dtype: torch.dtype,
88
+ num_blocks: int = 2048,
89
+ kv_cache_dtype: torch.dtype | None = None,
90
+ k_scale: float = 1.0,
91
+ v_scale: float = 1.0,
92
+ q_dtype: torch.dtype | None = None,
93
+ q_scale: float = Q_SCALE,
94
+) -> dict[str, Any]:
95
+ torch.set_default_device("cuda")
96
+
97
+ query_lens = [q for q, _ in seq_lens]
98
+ kv_lens = [k for _, k in seq_lens]
99
+ num_seqs = len(seq_lens)
100
+ max_query_len = max(query_lens)
101
+ max_kv_len = max(kv_lens)
102
+ scale = head_size**-0.5
103
+
104
+ query = torch.randn(sum(query_lens), NUM_Q_HEADS, head_size, dtype=dtype)
105
+ if kv_cache_dtype is None:
106
+ key_cache = torch.randn(
107
+ num_blocks, block_size, NUM_KV_HEADS, head_size, dtype=dtype
108
+ )
109
+ value_cache = torch.randn_like(key_cache)
110
+ else:
111
+ key_cache = torch.clamp(
112
+ torch.randn(num_blocks, block_size, NUM_KV_HEADS, head_size),
113
+ -1.0,
114
+ 1.0,
115
+ ).to(kv_cache_dtype)
116
+ value_cache = torch.clamp(
117
+ torch.randn(num_blocks, block_size, NUM_KV_HEADS, head_size),
118
+ -1.0,
119
+ 1.0,
120
+ ).to(kv_cache_dtype)
121
+
122
+ cu_seqlens_q = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
123
+ dim=0, dtype=torch.int32
124
+ )
125
+ seq_lens_tensor = torch.tensor(kv_lens, dtype=torch.int32)
126
+
127
+ max_num_blocks = (max_kv_len + block_size - 1) // block_size
128
+ block_tables = torch.randint(
129
+ 0, num_blocks, (num_seqs, max_num_blocks), dtype=torch.int32
130
+ )
131
+
132
+ descale_shape = (num_seqs, NUM_KV_HEADS)
133
+ k_descale = torch.full(descale_shape, k_scale, dtype=torch.float32, device="cuda")
134
+ v_descale = torch.full(descale_shape, v_scale, dtype=torch.float32, device="cuda")
135
+
136
+ kernel_query = query
137
+ q_descale = None
138
+ if q_dtype is not None:
139
+ q_descale = torch.tensor(q_scale, dtype=torch.float32, device="cuda")
140
+ kernel_query = (query / q_scale).to(q_dtype)
141
+
142
+ return {
143
+ "query": query,
144
+ "kernel_query": kernel_query,
145
+ "key_cache": key_cache,
146
+ "value_cache": value_cache,
147
+ "block_tables": block_tables,
148
+ "query_lens": query_lens,
149
+ "kv_lens": kv_lens,
150
+ "seq_lens_tensor": seq_lens_tensor,
151
+ "cu_seqlens_q": cu_seqlens_q,
152
+ "q_descale": q_descale,
153
+ "k_descale": k_descale,
154
+ "v_descale": v_descale,
155
+ "scale": scale,
156
+ "max_query_len": max_query_len,
157
+ "max_kv_len": max_kv_len,
158
+ "query_dtype": dtype,
159
+ "k_scale": k_scale,
160
+ "v_scale": v_scale,
161
+ }
162
+
163
+
164
+def _make_fp8_case(
165
+ *,
166
+ seq_lens: list[tuple[int, int]],
167
+ head_size: int,
168
+ block_size: int,
169
+ variant: Fp8Variant,
170
+) -> dict[str, Any]:
171
+ use_fp8_kv = variant in ("fp8_kv", "fp8_query_kv")
172
+ use_fp8_query = variant in ("fp8_query", "fp8_query_kv")
173
+ return _make_case(
174
+ seq_lens=seq_lens,
175
+ head_size=head_size,
176
+ block_size=block_size,
177
+ dtype=torch.bfloat16,
178
+ kv_cache_dtype=FP8_DTYPE if use_fp8_kv else None,
179
+ k_scale=K_SCALE if use_fp8_kv else 1.0,
180
+ v_scale=V_SCALE if use_fp8_kv else 1.0,
181
+ q_dtype=FP8_DTYPE if use_fp8_query else None,
182
+ )
183
+
184
+
185
+def _run_aiter_unified_attention(case: dict[str, Any]) -> torch.Tensor:
186
+ from aiter.ops.triton.unified_attention import unified_attention
187
+
188
+ kernel_query = case["kernel_query"]
189
+ # Kernel writes high-precision output even when Q is FP8 (matches vLLM usage).
190
+ output = torch.empty_like(case["query"])
191
+ unified_attention(
192
+ q=kernel_query,
193
+ k=case["key_cache"],
194
+ v=case["value_cache"],
195
+ out=output,
196
+ cu_seqlens_q=case["cu_seqlens_q"],
197
+ max_seqlen_q=case["max_query_len"],
198
+ seqused_k=case["seq_lens_tensor"],
199
+ max_seqlen_k=case["max_kv_len"],
200
+ softmax_scale=case["scale"],
201
+ causal=True,
202
+ alibi_slopes=None,
203
+ window_size=(-1, -1),
204
+ block_table=case["block_tables"],
205
+ softcap=0,
206
+ q_descale=case["q_descale"],
207
+ k_descale=case["k_descale"],
208
+ v_descale=case["v_descale"],
209
+ sinks=None,
210
+ output_scale=None,
211
+ )
212
+ return output
213
+
214
+
215
+def _ref_output(case: dict[str, Any]) -> torch.Tensor:
216
+ key_cache = case["key_cache"]
217
+ value_cache = case["value_cache"]
218
+ if key_cache.dtype != case["query_dtype"]:
219
+ key_cache = key_cache.to(case["query_dtype"]) * case["k_scale"]
220
+ value_cache = value_cache.to(case["query_dtype"]) * case["v_scale"]
221
+
222
+ return ref_paged_attn(
223
+ query=case["query"],
224
+ key_cache=key_cache,
225
+ value_cache=value_cache,
226
+ query_lens=case["query_lens"],
227
+ kv_lens=case["kv_lens"],
228
+ block_tables=case["block_tables"],
229
+ scale=case["scale"],
230
+ )
231
+
232
+
233
+def _assert_matches_reference(
234
+ case: dict[str, Any],
235
+ *,
236
+ atol: float = DEFAULT_ATOL,
237
+ rtol: float = DEFAULT_RTOL,
238
+) -> None:
239
+ output = _run_aiter_unified_attention(case)
240
+ output_ref = _ref_output(case)
241
+ torch.testing.assert_close(output, output_ref, atol=atol, rtol=rtol)
242
+
243
+
244
+@pytest.mark.parametrize("seq_lens", MIXED_SEQ_LENS)
245
+@pytest.mark.parametrize("head_size", HEAD_SIZES)
246
+@pytest.mark.parametrize("block_size", BLOCK_SIZES)
247
+@pytest.mark.parametrize("dtype", DTYPES)
248
+@torch.inference_mode()
249
+def test_aiter_unified_attn_mixed_batch(
250
+ seq_lens: list[tuple[int, int]],
251
+ head_size: int,
252
+ block_size: int,
253
+ dtype: torch.dtype,
254
+) -> None:
255
+ """Decode + prefill sequences in one batch (native dtypes)."""
256
+ _require_aiter()
257
+ set_random_seed(0)
258
+
259
+ case = _make_case(
260
+ seq_lens=seq_lens,
261
+ head_size=head_size,
262
+ block_size=block_size,
263
+ dtype=dtype,
264
+ )
265
+ _assert_matches_reference(case)
266
+
267
+
268
+@pytest.mark.parametrize("seq_lens", DECODE_SEQ_LENS)
269
+@pytest.mark.parametrize("head_size", HEAD_SIZES)
270
+@pytest.mark.parametrize("block_size", BLOCK_SIZES)
271
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
272
+@torch.inference_mode()
273
+def test_aiter_unified_attn_decode(
274
+ seq_lens: list[tuple[int, int]],
275
+ head_size: int,
276
+ block_size: int,
277
+ dtype: torch.dtype,
278
+) -> None:
279
+ """Single-token decode (native dtypes)."""
280
+ _require_aiter()
281
+ set_random_seed(0)
282
+
283
+ case = _make_case(
284
+ seq_lens=seq_lens,
285
+ head_size=head_size,
286
+ block_size=block_size,
287
+ dtype=dtype,
288
+ )
289
+ _assert_matches_reference(case)
290
+
291
+
292
+@pytest.mark.parametrize("seq_lens", PREFILL_SEQ_LENS)
293
+@pytest.mark.parametrize("head_size", [128])
294
+@pytest.mark.parametrize("block_size", [16])
295
+@torch.inference_mode()
296
+def test_aiter_unified_attn_prefill(
297
+ seq_lens: list[tuple[int, int]],
298
+ head_size: int,
299
+ block_size: int,
300
+) -> None:
301
+ """Prefill-only batches with query_len > 1 (native dtypes)."""
302
+ _require_aiter()
303
+ set_random_seed(0)
304
+
305
+ case = _make_case(
306
+ seq_lens=seq_lens,
307
+ head_size=head_size,
308
+ block_size=block_size,
309
+ dtype=torch.bfloat16,
310
+ )
311
+ _assert_matches_reference(case)
312
+
313
+
314
+@pytest.mark.skipif(
315
+ not current_platform.supports_fp8(),
316
+ reason="FP8 not supported on this hardware",
317
+)
318
+@pytest.mark.parametrize("variant", FP8_VARIANTS)
319
+@pytest.mark.parametrize("seq_lens", FP8_SEQ_LENS)
320
+@pytest.mark.parametrize("head_size", [128])
321
+@pytest.mark.parametrize("block_size", [16, 64])
322
+@torch.inference_mode()
323
+def test_aiter_unified_attn_fp8(
324
+ variant: Fp8Variant,
325
+ seq_lens: list[tuple[int, int]],
326
+ head_size: int,
327
+ block_size: int,
328
+) -> None:
329
+ """FP8 KV cache, FP8 query, or both; compared at bf16 reference precision."""
330
+ _require_aiter()
331
+ set_random_seed(0)
332
+
333
+ case = _make_fp8_case(
334
+ seq_lens=seq_lens,
335
+ head_size=head_size,
336
+ block_size=block_size,
337
+ variant=variant,
338
+ )
339
+ _assert_matches_reference(case, atol=FP8_ATOL, rtol=FP8_RTOL)