@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 31cd109f18

[ROCm][CI] Extend ROCm AITER MHA (FA) coverage (#40958)

Signed-off-by: Andreas Karatzas <akaratza@amd.com> Signed-off-by: Micah Williamson <micah.williamson@amd.com> Co-authored-by: Micah Williamson <micah.williamson@amd.com>

Andreas Karatzas committed Aug 10, 2026 at 02:35 UTC 31cd109f18f1ff49cef8dfd25329e7c49d0f2310
2 files changed +906 -258
tests/kernels/attention/test_aiter_flash_attn.py deleted
-258
@@ -1,258 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -
5 -import pytest
6 -import torch
7 -
8 -from vllm.platforms import current_platform
9 -from vllm.utils.torch_utils import set_random_seed
10 -
11 -# Import AITER backend if on ROCm and aiter is available
12 -if current_platform.is_rocm():
13 - from vllm._aiter_ops import is_aiter_found_and_supported
14 -
15 - if is_aiter_found_and_supported():
16 - import aiter
17 -
18 - from vllm.v1.attention.backends.rocm_aiter_fa import cp_mha_gather_cache
19 -
20 -NUM_HEADS = [(4, 4), (8, 2)]
21 -HEAD_SIZES = [128, 256]
22 -BLOCK_SIZES = [16]
23 -DTYPES = [torch.bfloat16]
24 -QDTYPES = [None]
25 -# one value large enough to test overflow in index calculation.
26 -# one value small enough to test the schema op check
27 -NUM_BLOCKS = [32768, 2048]
28 -
29 -
30 -def ref_paged_attn(
31 - query: torch.Tensor,
32 - key_cache: torch.Tensor,
33 - value_cache: torch.Tensor,
34 - query_lens: list[int],
35 - kv_lens: list[int],
36 - block_tables: torch.Tensor,
37 - scale: float,
38 - sliding_window: int | None = None,
39 - soft_cap: float | None = None,
40 -) -> torch.Tensor:
41 - num_seqs = len(query_lens)
42 - block_tables = block_tables.cpu().numpy()
43 - _, block_size, num_kv_heads, head_size = key_cache.shape
44 -
45 - outputs: list[torch.Tensor] = []
46 - start_idx = 0
47 - for i in range(num_seqs):
48 - query_len = query_lens[i]
49 - kv_len = kv_lens[i]
50 - q = query[start_idx : start_idx + query_len]
51 - q *= scale
52 -
53 - num_kv_blocks = (kv_len + block_size - 1) // block_size
54 - block_indices = block_tables[i, :num_kv_blocks]
55 -
56 - k = key_cache[block_indices].view(-1, num_kv_heads, head_size)
57 - k = k[:kv_len]
58 - v = value_cache[block_indices].view(-1, num_kv_heads, head_size)
59 - v = v[:kv_len]
60 -
61 - if q.shape[1] != k.shape[1]:
62 - k = torch.repeat_interleave(k, q.shape[1] // k.shape[1], dim=1)
63 - v = torch.repeat_interleave(v, q.shape[1] // v.shape[1], dim=1)
64 - attn = torch.einsum("qhd,khd->hqk", q, k).float()
65 - empty_mask = torch.ones(query_len, kv_len)
66 - mask = torch.triu(empty_mask, diagonal=kv_len - query_len + 1).bool()
67 - if sliding_window is not None:
68 - sliding_window_mask = (
69 - torch.triu(
70 - empty_mask, diagonal=kv_len - (query_len + sliding_window) + 1
71 - )
72 - .bool()
73 - .logical_not()
74 - )
75 - mask |= sliding_window_mask
76 - if soft_cap is not None:
77 - attn = soft_cap * torch.tanh(attn / soft_cap)
78 - attn.masked_fill_(mask, float("-inf"))
79 - attn = torch.softmax(attn, dim=-1).to(v.dtype)
80 - out = torch.einsum("hqk,khd->qhd", attn, v)
81 -
82 - outputs.append(out)
83 - start_idx += query_len
84 -
85 - return torch.cat(outputs, dim=0)
86 -
87 -
88 -@pytest.mark.skipif(not current_platform.is_rocm(), reason="Only ROCm is supported")
89 -@pytest.mark.parametrize(
90 - "seq_lens", [[(10, 1328), (5, 18), (129, 463)], [(8, 523), (24, 37), (3, 2011)]]
91 -)
92 -@pytest.mark.parametrize("num_heads", NUM_HEADS)
93 -@pytest.mark.parametrize("head_size", HEAD_SIZES)
94 -@pytest.mark.parametrize("block_size", BLOCK_SIZES)
95 -@pytest.mark.parametrize("sliding_window", [None, 256])
96 -@pytest.mark.parametrize("dtype", DTYPES)
97 -@pytest.mark.parametrize("soft_cap", [None])
98 -@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
99 -@pytest.mark.parametrize("q_dtype", QDTYPES)
100 -@torch.inference_mode()
101 -def test_varlen_with_paged_kv(
102 - seq_lens: list[tuple[int, int]],
103 - num_heads: tuple[int, int],
104 - head_size: int,
105 - sliding_window: int | None,
106 - dtype: torch.dtype,
107 - block_size: int,
108 - soft_cap: float | None,
109 - num_blocks: int,
110 - q_dtype: torch.dtype | None,
111 -) -> None:
112 - from vllm._aiter_ops import is_aiter_found_and_supported
113 -
114 - if not is_aiter_found_and_supported():
115 - pytest.skip("aiter package required for this test.")
116 -
117 - torch.set_default_device("cuda")
118 - set_random_seed(0)
119 - num_seqs = len(seq_lens)
120 - query_lens = [x[0] for x in seq_lens]
121 - kv_lens = [x[1] for x in seq_lens]
122 - num_query_heads = num_heads[0]
123 - num_kv_heads = num_heads[1]
124 - assert num_query_heads % num_kv_heads == 0
125 - max_query_len = max(query_lens)
126 - max_kv_len = max(kv_lens)
127 - window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
128 - scale = head_size**-0.5
129 -
130 - query = torch.randn(sum(query_lens), num_query_heads, head_size, dtype=dtype)
131 - key_cache = torch.randn(
132 - num_blocks, block_size, num_kv_heads, head_size, dtype=dtype
133 - )
134 - value_cache = torch.randn_like(key_cache)
135 - cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
136 - dim=0, dtype=torch.int32
137 - )
138 -
139 - cu_seq_lens = torch.tensor([0] + kv_lens, dtype=torch.int32).cumsum(
140 - dim=0, dtype=torch.int32
141 - )
142 - # Save kv_lens as list before converting to tensor
143 - kv_lens_list = kv_lens
144 - kv_lens = torch.tensor(kv_lens, dtype=torch.int32)
145 -
146 - max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
147 - block_tables = torch.randint(
148 - 0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
149 - )
150 -
151 - output = torch.empty_like(query)
152 -
153 - maybe_quantized_query = query
154 - maybe_quantized_key_cache = key_cache
155 - maybe_quantized_value_cache = value_cache
156 - k_scale_tensor = None
157 - v_scale_tensor = None
158 - dequant = False
159 -
160 - if q_dtype is not None:
161 - # QKV are drawn from N(0, 1): no need for a fp8 scaling factor
162 - maybe_quantized_query = query.to(q_dtype)
163 - maybe_quantized_key_cache = key_cache.to(q_dtype)
164 - maybe_quantized_value_cache = value_cache.to(q_dtype)
165 - dequant = True
166 - scale_shape = (num_seqs, num_kv_heads)
167 -
168 - # For per-seq-per-head scales (matching AITER backend expectation)
169 - k_scale_tensor = torch.ones(scale_shape, dtype=torch.float32)
170 - v_scale_tensor = torch.ones(scale_shape, dtype=torch.float32)
171 -
172 - # Prepare metadata for cp_mha_gather_cache
173 - # token_to_batch: maps each token to its batch index
174 - token_to_batch = torch.zeros(sum(kv_lens_list), dtype=torch.int32)
175 - seq_starts = torch.zeros(num_seqs, dtype=torch.int32)
176 -
177 - token_idx = 0
178 - for batch_idx, kv_len in enumerate(kv_lens_list):
179 - token_to_batch[token_idx : token_idx + kv_len] = batch_idx
180 - seq_starts[batch_idx] = 0 # Assuming all sequences start at 0 in their blocks
181 - token_idx += kv_len
182 -
183 - # Allocate buffers for gathered KV
184 - total_kv_tokens = sum(kv_lens_list)
185 - gathered_key = torch.empty(
186 - total_kv_tokens, num_kv_heads, head_size, dtype=maybe_quantized_key_cache.dtype
187 - )
188 - gathered_value = torch.empty(
189 - total_kv_tokens,
190 - num_kv_heads,
191 - head_size,
192 - dtype=maybe_quantized_value_cache.dtype,
193 - )
194 -
195 - # Gather paged KV cache into contiguous tensors using triton kernel
196 - cp_mha_gather_cache(
197 - key_cache=maybe_quantized_key_cache,
198 - value_cache=maybe_quantized_value_cache,
199 - key=gathered_key,
200 - value=gathered_value,
201 - block_tables=block_tables,
202 - k_scales=k_scale_tensor
203 - if k_scale_tensor is not None
204 - else torch.ones(1, dtype=torch.float32),
205 - v_scales=v_scale_tensor
206 - if v_scale_tensor is not None
207 - else torch.ones(1, dtype=torch.float32),
208 - cu_seqlens_kv=cu_seq_lens,
209 - token_to_batch=token_to_batch,
210 - seq_starts=seq_starts,
211 - dequant=dequant,
212 - kv_cache_layout="NHD",
213 - total_tokens=total_kv_tokens,
214 - )
215 -
216 - # Call aiter flash attention with gathered KV
217 - aiter.flash_attn_varlen_func(
218 - q=maybe_quantized_query,
219 - k=gathered_key,
220 - v=gathered_value,
221 - cu_seqlens_q=cu_query_lens,
222 - cu_seqlens_k=cu_seq_lens,
223 - max_seqlen_q=max_query_len,
224 - max_seqlen_k=max_kv_len,
225 - min_seqlen_q=1,
226 - dropout_p=0.0,
227 - softmax_scale=scale,
228 - causal=True,
229 - window_size=window_size,
230 - alibi_slopes=None,
231 - return_lse=False,
232 - out=output,
233 - )
234 -
235 - ref_output = ref_paged_attn(
236 - query=query,
237 - key_cache=key_cache,
238 - value_cache=value_cache,
239 - query_lens=query_lens,
240 - kv_lens=kv_lens_list,
241 - block_tables=block_tables,
242 - scale=scale,
243 - sliding_window=sliding_window,
244 - soft_cap=soft_cap,
245 - )
246 -
247 - atol, rtol = 2e-2, 2e-2
248 - if q_dtype is not None:
249 - atol, rtol = 1.5e-1, 1.5e-1
250 - (
251 - torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol),
252 - f"{torch.max(torch.abs(output - ref_output))}",
253 - )
254 -
255 - # Log diff stats for tracking changes
256 - print(f"Max abs diff: {torch.max(torch.abs(output - ref_output))}")
257 - print(f"Mean diff: {torch.mean(torch.abs(output - ref_output))}")
258 - print(f"Min diff: {torch.std(torch.abs(output - ref_output))}")
tests/kernels/attention/test_rocm_aiter_fa.py new
+906
@@ -0,0 +1,906 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""ROCm-specific tests for the AITER Flash Attention backend.
4 +
5 +This file owns the ROCm backend contract and the ROCm-specific execution paths
6 +wired through ``vllm.v1.attention.backends.rocm_aiter_fa``:
7 +- backend contract and env gating
8 +- representative prefill, multi-batch, decode, and FP8-KV execution
9 +- direct kernel stress for large block tables and sliding-window masking
10 +"""
11 +
12 +import importlib
13 +
14 +import pytest
15 +import torch
16 +
17 +from vllm.platforms import current_platform
18 +from vllm.platforms.rocm import on_mi3xx
19 +from vllm.utils.torch_utils import set_random_seed
20 +
21 +pytestmark = pytest.mark.skipif(
22 + not current_platform.is_rocm(), reason="ROCm-specific tests"
23 +)
24 +
25 +
26 +@pytest.fixture(autouse=True)
27 +def _cuda_default_device():
28 + torch.set_default_device("cuda")
29 + try:
30 + yield
31 + finally:
32 + torch.set_default_device("cpu")
33 +
34 +
35 +DTYPES = [torch.bfloat16, torch.float16]
36 +HEAD_SIZES = [64, 128, 256]
37 +NUM_HEADS_PAIRS = [(8, 8), (16, 4)] # (num_q_heads, num_kv_heads) - tests GQA
38 +BLOCK_SIZE = 16
39 +NUM_BLOCKS = 2048
40 +DIRECT_NUM_BLOCKS = [2048, 32768]
41 +DIRECT_NUM_HEADS = (8, 2)
42 +DIRECT_HEAD_SIZE = 128
43 +DIRECT_DTYPE = torch.bfloat16
44 +HEAD_SIZE_TEST_NUM_HEADS_PAIRS = [(16, 16), (16, 4)]
45 +# Prefill seq lens: (query_len, kv_len). Exclude single-token decode (q=1)
46 +# because flash_attn_varlen_func is a prefill kernel; q_len=1 with short kv
47 +# triggers kernel limitations (MAE > 0.1 for head != 128 in BF16, all heads in FP16).
48 +# Single-token decode is covered separately in test_aiter_mha_decode_single_token.
49 +SEQ_LENS = [(8, 512), (32, 1024)]
50 +DIRECT_SEQ_LENS = [(10, 1328), (5, 18), (129, 463)]
51 +
52 +
53 +# Reference implementation -------------------------------------------------
54 +def ref_paged_attn(
55 + query: torch.Tensor,
56 + key_cache: torch.Tensor,
57 + value_cache: torch.Tensor,
58 + query_lens: list[int],
59 + kv_lens: list[int],
60 + block_tables: torch.Tensor,
61 + scale: float,
62 + sliding_window: int | None = None,
63 +) -> torch.Tensor:
64 + """Naive reference paged attention using einsum."""
65 + num_seqs = len(query_lens)
66 + block_tables_np = block_tables.cpu().numpy()
67 + _, block_size, num_kv_heads, head_size = key_cache.shape
68 +
69 + outputs = []
70 + start_idx = 0
71 + for i in range(num_seqs):
72 + query_len = query_lens[i]
73 + kv_len = kv_lens[i]
74 + q = query[start_idx : start_idx + query_len] * scale
75 +
76 + num_kv_blocks = (kv_len + block_size - 1) // block_size
77 + block_indices = block_tables_np[i, :num_kv_blocks]
78 +
79 + k = key_cache[block_indices].view(-1, num_kv_heads, head_size)[:kv_len]
80 + v = value_cache[block_indices].view(-1, num_kv_heads, head_size)[:kv_len]
81 +
82 + if q.shape[1] != k.shape[1]:
83 + k = torch.repeat_interleave(k, q.shape[1] // k.shape[1], dim=1)
84 + v = torch.repeat_interleave(v, q.shape[1] // v.shape[1], dim=1)
85 +
86 + attn = torch.einsum("qhd,khd->hqk", q, k).float()
87 + mask = torch.triu(
88 + torch.ones(query_len, kv_len), diagonal=kv_len - query_len + 1
89 + ).bool()
90 + if sliding_window is not None:
91 + window_mask = (
92 + torch.triu(
93 + torch.ones(query_len, kv_len),
94 + diagonal=kv_len - (query_len + sliding_window) + 1,
95 + )
96 + .bool()
97 + .logical_not()
98 + )
99 + mask |= window_mask
100 + attn.masked_fill_(mask, float("-inf"))
101 + attn = torch.softmax(attn, dim=-1).to(v.dtype)
102 + out = torch.einsum("hqk,khd->qhd", attn, v)
103 + outputs.append(out)
104 + start_idx += query_len
105 +
106 + return torch.cat(outputs, dim=0)
107 +
108 +
109 +# Small test helpers ------------------------------------------------------
110 +def _reload_envs():
111 + import vllm.envs as envs
112 +
113 + return importlib.reload(envs)
114 +
115 +
116 +def _assert_aiter_supported() -> None:
117 + from vllm._aiter_ops import is_aiter_found_and_supported
118 +
119 + assert is_aiter_found_and_supported(), (
120 + "aiter is required on supported ROCm hardware for this test"
121 + )
122 +
123 +
124 +def _format_observed_rate(count: int, total: int) -> str:
125 + return f"{count / total:.4%} ({count}/{total})"
126 +
127 +
128 +def _format_allowed_rate(rate: float, total: int) -> str:
129 + allowed_count = int(rate * total)
130 + return f"{rate:.4%} (<= {allowed_count}/{total})"
131 +
132 +
133 +def _print_close_stats(
134 + label: str,
135 + actual: torch.Tensor,
136 + expected: torch.Tensor,
137 + *,
138 + atol: float,
139 + rtol: float,
140 +) -> None:
141 + abs_diff = (actual - expected).abs().float().flatten()
142 + expected_abs = expected.abs().float().flatten()
143 + allowed = atol + rtol * expected_abs
144 + within = abs_diff <= allowed
145 +
146 + total = abs_diff.numel()
147 + passed = int(within.sum().item())
148 + failed = total - passed
149 +
150 + max_abs = abs_diff.max().item()
151 + mean_abs = abs_diff.mean().item()
152 + p99_abs = torch.quantile(abs_diff, 0.99).item()
153 + p999_abs = torch.quantile(abs_diff, 0.999).item()
154 + worst_ratio = (abs_diff / allowed.clamp_min(1e-12)).max().item()
155 +
156 + print(
157 + "[rocm_aiter_fa] "
158 + f"{label}: "
159 + f"pass={passed / total:.4%} ({passed}/{total}) "
160 + f"fail={_format_observed_rate(failed, total)} "
161 + f"allowed_fail={_format_allowed_rate(0.0, total)} "
162 + f"atol={atol:g} "
163 + f"rtol={rtol:g} "
164 + f"max_abs={max_abs:.6g} "
165 + f"mean_abs={mean_abs:.6g} "
166 + f"p99_abs={p99_abs:.6g} "
167 + f"p999_abs={p999_abs:.6g} "
168 + f"worst_ratio={worst_ratio:.6g}"
169 + )
170 +
171 +
172 +# Direct kernel helper ----------------------------------------------------
173 +def _run_direct_flash_attn_case(
174 + *,
175 + seq_lens: list[tuple[int, int]],
176 + num_blocks: int,
177 + sliding_window: int | None = None,
178 +) -> tuple[torch.Tensor, torch.Tensor]:
179 + import aiter
180 +
181 + from vllm.v1.attention.backends.rocm_aiter_fa import cp_mha_gather_cache
182 +
183 + set_random_seed(0)
184 +
185 + query_lens = [query_len for query_len, _ in seq_lens]
186 + kv_lens = [kv_len for _, kv_len in seq_lens]
187 + num_query_heads, num_kv_heads = DIRECT_NUM_HEADS
188 + total_query_tokens = sum(query_lens)
189 + total_kv_tokens = sum(kv_lens)
190 + max_query_len = max(query_lens)
191 + max_kv_len = max(kv_lens)
192 + scale = DIRECT_HEAD_SIZE**-0.5
193 +
194 + query = torch.randn(
195 + total_query_tokens,
196 + num_query_heads,
197 + DIRECT_HEAD_SIZE,
198 + dtype=DIRECT_DTYPE,
199 + )
200 + key_cache = torch.randn(
201 + num_blocks,
202 + BLOCK_SIZE,
203 + num_kv_heads,
204 + DIRECT_HEAD_SIZE,
205 + dtype=DIRECT_DTYPE,
206 + )
207 + value_cache = torch.randn_like(key_cache)
208 +
209 + cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
210 + dim=0, dtype=torch.int32
211 + )
212 + cu_seq_lens = torch.tensor([0] + kv_lens, dtype=torch.int32).cumsum(
213 + dim=0, dtype=torch.int32
214 + )
215 + max_num_blocks_per_seq = (max_kv_len + BLOCK_SIZE - 1) // BLOCK_SIZE
216 + block_tables = torch.randint(
217 + 0, num_blocks, (len(seq_lens), max_num_blocks_per_seq), dtype=torch.int32
218 + )
219 +
220 + token_to_batch = torch.empty(total_kv_tokens, dtype=torch.int32)
221 + seq_starts = torch.zeros(len(seq_lens), dtype=torch.int32)
222 + token_index = 0
223 + for batch_index, kv_len in enumerate(kv_lens):
224 + token_to_batch[token_index : token_index + kv_len] = batch_index
225 + token_index += kv_len
226 +
227 + gathered_key = torch.empty(
228 + total_kv_tokens,
229 + num_kv_heads,
230 + DIRECT_HEAD_SIZE,
231 + dtype=DIRECT_DTYPE,
232 + )
233 + gathered_value = torch.empty_like(gathered_key)
234 + cp_mha_gather_cache(
235 + key_cache=key_cache,
236 + value_cache=value_cache,
237 + key=gathered_key,
238 + value=gathered_value,
239 + block_tables=block_tables,
240 + k_scales=torch.ones(1, dtype=torch.float32),
241 + v_scales=torch.ones(1, dtype=torch.float32),
242 + cu_seqlens_kv=cu_seq_lens,
243 + token_to_batch=token_to_batch,
244 + seq_starts=seq_starts,
245 + dequant=False,
246 + kv_cache_layout="NHD",
247 + total_tokens=total_kv_tokens,
248 + )
249 +
250 + window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
251 + output = torch.empty_like(query)
252 + aiter.flash_attn_varlen_func(
253 + q=query,
254 + k=gathered_key,
255 + v=gathered_value,
256 + cu_seqlens_q=cu_query_lens,
257 + cu_seqlens_k=cu_seq_lens,
258 + max_seqlen_q=max_query_len,
259 + max_seqlen_k=max_kv_len,
260 + min_seqlen_q=1,
261 + dropout_p=0.0,
262 + softmax_scale=scale,
263 + causal=True,
264 + window_size=window_size,
265 + alibi_slopes=None,
266 + return_lse=False,
267 + out=output,
268 + )
269 +
270 + ref = ref_paged_attn(
271 + query=query,
272 + key_cache=key_cache,
273 + value_cache=value_cache,
274 + query_lens=query_lens,
275 + kv_lens=kv_lens,
276 + block_tables=block_tables,
277 + scale=scale,
278 + sliding_window=sliding_window,
279 + )
280 + return output, ref
281 +
282 +
283 +# Backend contract tests --------------------------------------------------
284 +def test_aiter_mha_backend_contract():
285 + """The ROCm backend advertises the dtypes, shapes, and attention types it
286 + is designed to handle."""
287 + from vllm.v1.attention.backend import AttentionType
288 + from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend
289 +
290 + assert AiterFlashAttentionBackend.get_name() == "FLASH_ATTN"
291 + assert AiterFlashAttentionBackend.supported_dtypes == [
292 + torch.float16,
293 + torch.bfloat16,
294 + ]
295 + assert AiterFlashAttentionBackend.supported_kv_cache_dtypes == [
296 + "auto",
297 + "float16",
298 + "bfloat16",
299 + "fp8",
300 + "fp8_e4m3",
301 + "fp8_e5m2",
302 + ]
303 + assert AiterFlashAttentionBackend.get_supported_kernel_block_sizes() == [16, 32]
304 + assert AiterFlashAttentionBackend.get_supported_head_sizes() == [64, 128, 256]
305 + assert AiterFlashAttentionBackend.supports_attn_type(AttentionType.DECODER)
306 + assert not AiterFlashAttentionBackend.supports_attn_type(AttentionType.ENCODER)
307 + assert not AiterFlashAttentionBackend.supports_attn_type(
308 + AttentionType.ENCODER_DECODER
309 + )
310 +
311 +
312 +def test_aiter_mha_backend_validates_kv_cache_block_size():
313 + """The backend should reject KV cache shapes that cannot be gathered
314 + correctly by the ROCm kernel."""
315 + from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend
316 +
317 + assert AiterFlashAttentionBackend.get_kv_cache_shape(8, 16, 8, 128) == (
318 + 8,
319 + 8,
320 + 16,
321 + 256,
322 + )
323 + with pytest.raises(ValueError, match="Block size must be a multiple of 16"):
324 + AiterFlashAttentionBackend.get_kv_cache_shape(8, 15, 8, 128)
325 +
326 +
327 +def test_aiter_mha_backend_supports_compute_capability_matches_mi3xx_probe():
328 + """The backend should trust the ROCm MI3xx probe instead of the raw torch
329 + capability tuple."""
330 + from vllm.platforms.interface import DeviceCapability
331 + from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend
332 +
333 + assert (
334 + AiterFlashAttentionBackend.supports_compute_capability(DeviceCapability(0, 0))
335 + is on_mi3xx()
336 + )
337 +
338 +
339 +# Env and platform gate tests ---------------------------------------------
340 +@pytest.mark.skipif(not on_mi3xx(), reason="MI300/MI350 ROCm only")
341 +@pytest.mark.parametrize(
342 + ("use_aiter", "use_mha", "expected"),
343 + [
344 + (True, True, True),
345 + (True, False, False),
346 + (False, True, False),
347 + (False, False, False),
348 + ],
349 +)
350 +def test_aiter_mha_env_flags_control_enablement(
351 + use_aiter, use_mha, expected, monkeypatch
352 +):
353 + """Both the global AITER flag and the MHA-specific flag must be enabled
354 + before the ROCm MHA path is considered active."""
355 + from vllm._aiter_ops import rocm_aiter_ops
356 +
357 + _assert_aiter_supported()
358 +
359 + with monkeypatch.context() as mp:
360 + mp.setenv("VLLM_ROCM_USE_AITER", "1" if use_aiter else "0")
361 + mp.setenv("VLLM_ROCM_USE_AITER_MHA", "1" if use_mha else "0")
362 + _reload_envs()
363 + rocm_aiter_ops.refresh_env_variables()
364 + assert rocm_aiter_ops.is_mha_enabled() is expected
365 +
366 + _reload_envs()
367 + rocm_aiter_ops.refresh_env_variables()
368 +
369 +
370 +def test_aiter_mha_platform_gate_matches_install_and_arch():
371 + """The global AITER availability check should only open on ROCm MI3xx when
372 + the aiter package is installed."""
373 + from vllm._aiter_ops import IS_AITER_FOUND, is_aiter_found_and_supported
374 +
375 + assert is_aiter_found_and_supported() is (
376 + current_platform.is_rocm() and on_mi3xx() and IS_AITER_FOUND
377 + )
378 +
379 +
380 +# Kernel path tests -------------------------------------------------------
381 +def _run_single_seq_flash_attn_case(
382 + *,
383 + head_size: int,
384 + num_heads: tuple[int, int],
385 + seq_lens: tuple[int, int],
386 + dtype: torch.dtype,
387 +) -> tuple[torch.Tensor, torch.Tensor]:
388 + import aiter
389 +
390 + from vllm.v1.attention.backends.rocm_aiter_fa import cp_mha_gather_cache
391 +
392 + set_random_seed(0)
393 +
394 + num_q_heads, num_kv_heads = num_heads
395 + query_len, kv_len = seq_lens
396 + scale = head_size**-0.5
397 +
398 + query = torch.randn(query_len, num_q_heads, head_size, dtype=dtype)
399 + key_cache = torch.randn(
400 + NUM_BLOCKS, BLOCK_SIZE, num_kv_heads, head_size, dtype=dtype
401 + )
402 + value_cache = torch.randn_like(key_cache)
403 +
404 + cu_query_lens = torch.tensor([0, query_len], dtype=torch.int32).cumsum(
405 + dim=0, dtype=torch.int32
406 + )
407 + cu_seq_lens = torch.tensor([0, kv_len], dtype=torch.int32).cumsum(
408 + dim=0, dtype=torch.int32
409 + )
410 + max_num_blocks = (kv_len + BLOCK_SIZE - 1) // BLOCK_SIZE
411 + block_tables = torch.randint(0, NUM_BLOCKS, (1, max_num_blocks), dtype=torch.int32)
412 +
413 + token_to_batch = torch.zeros(kv_len, dtype=torch.int32)
414 + seq_starts = torch.zeros(1, dtype=torch.int32)
415 + gathered_key = torch.empty(kv_len, num_kv_heads, head_size, dtype=dtype)
416 + gathered_value = torch.empty_like(gathered_key)
417 +
418 + cp_mha_gather_cache(
419 + key_cache=key_cache,
420 + value_cache=value_cache,
421 + key=gathered_key,
422 + value=gathered_value,
423 + block_tables=block_tables,
424 + k_scales=torch.ones(1, dtype=torch.float32),
425 + v_scales=torch.ones(1, dtype=torch.float32),
426 + cu_seqlens_kv=cu_seq_lens,
427 + token_to_batch=token_to_batch,
428 + seq_starts=seq_starts,
429 + dequant=False,
430 + kv_cache_layout="NHD",
431 + total_tokens=kv_len,
432 + )
433 +
434 + output = torch.empty_like(query)
435 + aiter.flash_attn_varlen_func(
436 + q=query,
437 + k=gathered_key,
438 + v=gathered_value,
439 + cu_seqlens_q=cu_query_lens,
440 + cu_seqlens_k=cu_seq_lens,
441 + max_seqlen_q=query_len,
442 + max_seqlen_k=kv_len,
443 + min_seqlen_q=1,
444 + dropout_p=0.0,
445 + softmax_scale=scale,
446 + causal=True,
447 + window_size=(-1, -1),
448 + alibi_slopes=None,
449 + return_lse=False,
450 + out=output,
451 + )
452 +
453 + ref = ref_paged_attn(
454 + query=query,
455 + key_cache=key_cache,
456 + value_cache=value_cache,
457 + query_lens=[query_len],
458 + kv_lens=[kv_len],
459 + block_tables=block_tables,
460 + scale=scale,
461 + )
462 + return output, ref
463 +
464 +
465 +@pytest.mark.skipif(not on_mi3xx(), reason="MI300/MI350 ROCm only")
466 +@pytest.mark.parametrize("head_size", HEAD_SIZES)
467 +@pytest.mark.parametrize("dtype", DTYPES)
468 +@pytest.mark.parametrize("num_heads", HEAD_SIZE_TEST_NUM_HEADS_PAIRS)
469 +@pytest.mark.parametrize("seq_lens", SEQ_LENS)
470 +def test_aiter_fa_head_sizes(head_size, dtype, num_heads, seq_lens):
471 + """AITER flash attention should stay accurate across supported head sizes."""
472 + atol = 1.5e-2
473 + rtol = 1e-2
474 + _assert_aiter_supported()
475 +
476 + output, ref = _run_single_seq_flash_attn_case(
477 + head_size=head_size,
478 + num_heads=num_heads,
479 + seq_lens=seq_lens,
480 + dtype=dtype,
481 + )
482 +
483 + _print_close_stats(
484 + "head_sizes "
485 + f"dtype={dtype} head_size={head_size} "
486 + f"num_heads={num_heads} seq_lens={seq_lens}",
487 + output,
488 + ref,
489 + atol=atol,
490 + rtol=rtol,
491 + )
492 + torch.testing.assert_close(output, ref, atol=atol, rtol=rtol)
493 +
494 +
495 +@pytest.mark.skipif(not on_mi3xx(), reason="MI300/MI350 ROCm only")
496 +@pytest.mark.parametrize("head_size", HEAD_SIZES)
497 +@pytest.mark.parametrize("num_heads", NUM_HEADS_PAIRS)
498 +@pytest.mark.parametrize("seq_lens", SEQ_LENS)
499 +@pytest.mark.parametrize("dtype", DTYPES)
500 +def test_aiter_mha_varlen_paged_kv(head_size, num_heads, seq_lens, dtype):
501 + """Test AITER flash attention varlen with paged KV cache.
502 +
503 + Exercises: VLLM_ROCM_USE_AITER, VLLM_ROCM_USE_AITER_MHA
504 + """
505 + atol = 1.5e-2
506 + rtol = 1e-2
507 + _assert_aiter_supported()
508 +
509 + output, ref = _run_single_seq_flash_attn_case(
510 + head_size=head_size,
511 + num_heads=num_heads,
512 + seq_lens=seq_lens,
513 + dtype=dtype,
514 + )
515 +
516 + _print_close_stats(
517 + "varlen_paged_kv "
518 + f"dtype={dtype} head_size={head_size} "
519 + f"num_heads={num_heads} seq_lens={seq_lens}",
520 + output,
521 + ref,
522 + atol=atol,
523 + rtol=rtol,
524 + )
525 + torch.testing.assert_close(output, ref, atol=atol, rtol=rtol)
526 +
527 +
528 +@pytest.mark.skipif(not on_mi3xx(), reason="MI300/MI350 ROCm only")
529 +@pytest.mark.parametrize("num_heads", NUM_HEADS_PAIRS)
530 +@pytest.mark.parametrize("head_size", HEAD_SIZES)
531 +@pytest.mark.parametrize("dtype", DTYPES)
532 +def test_aiter_mha_multi_batch(num_heads, head_size, dtype):
533 + """Test AITER flash attention with multiple sequences in a batch."""
534 + atol = 1.5e-2
535 + rtol = 1e-2
536 + _assert_aiter_supported()
537 + import aiter
538 +
539 + from vllm.v1.attention.backends.rocm_aiter_fa import cp_mha_gather_cache
540 +
541 + set_random_seed(42)
542 +
543 + num_q_heads, num_kv_heads = num_heads
544 + seq_lens = [(4, 128), (2, 256), (8, 64)]
545 + query_lens = [q for q, _ in seq_lens]
546 + kv_lens = [k for _, k in seq_lens]
547 + num_seqs = len(seq_lens)
548 + scale = head_size**-0.5
549 +
550 + total_q = sum(query_lens)
551 + total_kv = sum(kv_lens)
552 +
553 + query = torch.randn(total_q, num_q_heads, head_size, dtype=dtype)
554 + key_cache = torch.randn(
555 + NUM_BLOCKS, BLOCK_SIZE, num_kv_heads, head_size, dtype=dtype
556 + )
557 + value_cache = torch.randn_like(key_cache)
558 +
559 + cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
560 + dim=0, dtype=torch.int32
561 + )
562 + cu_seq_lens = torch.tensor([0] + kv_lens, dtype=torch.int32).cumsum(
563 + dim=0, dtype=torch.int32
564 + )
565 +
566 + max_kv_len = max(kv_lens)
567 + max_num_blocks = (max_kv_len + BLOCK_SIZE - 1) // BLOCK_SIZE
568 + block_tables = torch.randint(
569 + 0, NUM_BLOCKS, (num_seqs, max_num_blocks), dtype=torch.int32
570 + )
571 +
572 + token_to_batch = torch.empty(total_kv, dtype=torch.int32)
573 + seq_starts = torch.zeros(num_seqs, dtype=torch.int32)
574 + tok_idx = 0
575 + for b, kl in enumerate(kv_lens):
576 + token_to_batch[tok_idx : tok_idx + kl] = b
577 + tok_idx += kl
578 +
579 + gathered_key = torch.empty(total_kv, num_kv_heads, head_size, dtype=dtype)
580 + gathered_value = torch.empty_like(gathered_key)
581 +
582 + cp_mha_gather_cache(
583 + key_cache=key_cache,
584 + value_cache=value_cache,
585 + key=gathered_key,
586 + value=gathered_value,
587 + block_tables=block_tables,
588 + k_scales=torch.ones(1, dtype=torch.float32),
589 + v_scales=torch.ones(1, dtype=torch.float32),
590 + cu_seqlens_kv=cu_seq_lens,
591 + token_to_batch=token_to_batch,
592 + seq_starts=seq_starts,
593 + dequant=False,
594 + kv_cache_layout="NHD",
595 + total_tokens=total_kv,
596 + )
597 +
598 + output = torch.empty_like(query)
599 + aiter.flash_attn_varlen_func(
600 + q=query,
601 + k=gathered_key,
602 + v=gathered_value,
603 + cu_seqlens_q=cu_query_lens,
604 + cu_seqlens_k=cu_seq_lens,
605 + max_seqlen_q=max(query_lens),
606 + max_seqlen_k=max_kv_len,
607 + min_seqlen_q=1,
608 + dropout_p=0.0,
609 + softmax_scale=scale,
610 + causal=True,
611 + window_size=(-1, -1),
612 + alibi_slopes=None,
613 + return_lse=False,
614 + out=output,
615 + )
616 +
617 + ref = ref_paged_attn(
618 + query=query,
619 + key_cache=key_cache,
620 + value_cache=value_cache,
621 + query_lens=query_lens,
622 + kv_lens=kv_lens,
623 + block_tables=block_tables,
624 + scale=scale,
625 + )
626 +
627 + _print_close_stats(
628 + "multi_batch "
629 + f"dtype={dtype} head_size={head_size} "
630 + f"num_heads={num_heads} seq_lens={seq_lens}",
631 + output,
632 + ref,
633 + atol=atol,
634 + rtol=rtol,
635 + )
636 + torch.testing.assert_close(output, ref, atol=atol, rtol=rtol)
637 +
638 +
639 +@pytest.mark.skipif(not on_mi3xx(), reason="MI300/MI350 ROCm only")
640 +@pytest.mark.parametrize("num_blocks", DIRECT_NUM_BLOCKS)
641 +def test_aiter_fa_large_block_table_matches_reference(num_blocks):
642 + """The direct paged-KV path should stay stable for both normal and very
643 + large block tables."""
644 + atol = 2e-2
645 + rtol = 2e-2
646 + _assert_aiter_supported()
647 +
648 + output, ref = _run_direct_flash_attn_case(
649 + seq_lens=DIRECT_SEQ_LENS,
650 + num_blocks=num_blocks,
651 + )
652 + _print_close_stats(
653 + f"direct_varlen_paged_kv num_blocks={num_blocks}",
654 + output,
655 + ref,
656 + atol=atol,
657 + rtol=rtol,
658 + )
659 + torch.testing.assert_close(output, ref, atol=atol, rtol=rtol)
660 +
661 +
662 +@pytest.mark.skipif(not on_mi3xx(), reason="MI300/MI350 ROCm only")
663 +def test_aiter_fa_sliding_window_matches_reference():
664 + """The direct kernel should respect the same sliding-window causal mask as
665 + the naive reference implementation."""
666 + atol = 2e-2
667 + rtol = 2e-2
668 + _assert_aiter_supported()
669 +
670 + output, ref = _run_direct_flash_attn_case(
671 + seq_lens=[(8, 523), (24, 37), (3, 2011)],
672 + num_blocks=2048,
673 + sliding_window=256,
674 + )
675 + _print_close_stats(
676 + "sliding_window num_blocks=2048 window=256",
677 + output,
678 + ref,
679 + atol=atol,
680 + rtol=rtol,
681 + )
682 + torch.testing.assert_close(output, ref, atol=atol, rtol=rtol)
683 +
684 +
685 +# Decode path test --------------------------------------------------------
686 +@pytest.mark.skipif(not on_mi3xx(), reason="MI300/MI350 ROCm only")
687 +@pytest.mark.parametrize(
688 + "dtype",
689 + [
690 + torch.bfloat16,
691 + pytest.param(
692 + torch.float16,
693 + marks=pytest.mark.xfail(
694 + reason=(
695 + "aiter bug #2229: flash_attn_varlen_func currently "
696 + "miscomputes FP16 single-token decode on MI3xx "
697 + "(validated on gfx950; max abs diff 1.95 vs atol 1.5e-2). "
698 + "Remove xfail when the AITER decode kernel is fixed. "
699 + "https://github.com/ROCm/aiter/issues/2229"
700 + ),
701 + ),
702 + ),
703 + ],
704 +)
705 +def test_aiter_mha_decode_single_token(dtype):
706 + """Test AITER MHA for decode (single query token per sequence).
707 +
708 + BF16 is the working reference configuration here. FP16 remains xfail until
709 + the upstream AITER single-token decode bug is fixed.
710 + """
711 + atol = 1.5e-2
712 + rtol = 1e-2
713 + _assert_aiter_supported()
714 + import aiter
715 +
716 + from vllm.v1.attention.backends.rocm_aiter_fa import cp_mha_gather_cache
717 +
718 + set_random_seed(0)
719 +
720 + num_q_heads, num_kv_heads = 8, 8
721 + head_size = 128
722 + kv_len = 512
723 + scale = head_size**-0.5
724 +
725 + query = torch.randn(1, num_q_heads, head_size, dtype=dtype)
726 + key_cache = torch.randn(
727 + NUM_BLOCKS, BLOCK_SIZE, num_kv_heads, head_size, dtype=dtype
728 + )
729 + value_cache = torch.randn_like(key_cache)
730 +
731 + cu_query_lens = torch.tensor([0, 1], dtype=torch.int32)
732 + cu_seq_lens = torch.tensor([0, kv_len], dtype=torch.int32)
733 +
734 + max_num_blocks = (kv_len + BLOCK_SIZE - 1) // BLOCK_SIZE
735 + block_tables = torch.randint(0, NUM_BLOCKS, (1, max_num_blocks), dtype=torch.int32)
736 +
737 + token_to_batch = torch.zeros(kv_len, dtype=torch.int32)
738 + seq_starts = torch.zeros(1, dtype=torch.int32)
739 + gathered_key = torch.empty(kv_len, num_kv_heads, head_size, dtype=dtype)
740 + gathered_value = torch.empty_like(gathered_key)
741 +
742 + cp_mha_gather_cache(
743 + key_cache=key_cache,
744 + value_cache=value_cache,
745 + key=gathered_key,
746 + value=gathered_value,
747 + block_tables=block_tables,
748 + k_scales=torch.ones(1, dtype=torch.float32),
749 + v_scales=torch.ones(1, dtype=torch.float32),
750 + cu_seqlens_kv=cu_seq_lens,
751 + token_to_batch=token_to_batch,
752 + seq_starts=seq_starts,
753 + dequant=False,
754 + kv_cache_layout="NHD",
755 + total_tokens=kv_len,
756 + )
757 +
758 + output = torch.empty_like(query)
759 + aiter.flash_attn_varlen_func(
760 + q=query,
761 + k=gathered_key,
762 + v=gathered_value,
763 + cu_seqlens_q=cu_query_lens,
764 + cu_seqlens_k=cu_seq_lens,
765 + max_seqlen_q=1,
766 + max_seqlen_k=kv_len,
767 + min_seqlen_q=1,
768 + dropout_p=0.0,
769 + softmax_scale=scale,
770 + causal=True,
771 + window_size=(-1, -1),
772 + alibi_slopes=None,
773 + return_lse=False,
774 + out=output,
775 + )
776 +
777 + ref = ref_paged_attn(
778 + query=query,
779 + key_cache=key_cache,
780 + value_cache=value_cache,
781 + query_lens=[1],
782 + kv_lens=[kv_len],
783 + block_tables=block_tables,
784 + scale=scale,
785 + )
786 +
787 + _print_close_stats(
788 + f"decode_single_token dtype={dtype} kv_len={kv_len}",
789 + output,
790 + ref,
791 + atol=atol,
792 + rtol=rtol,
793 + )
794 + torch.testing.assert_close(output, ref, atol=atol, rtol=rtol)
795 +
796 +
797 +# FP8 KV cache test -------------------------------------------------------
798 +@pytest.mark.skipif(not on_mi3xx(), reason="MI300/MI350 ROCm only")
799 +@pytest.mark.parametrize("dtype", DTYPES)
800 +def test_aiter_mha_varlen_fp8_kv(dtype):
801 + """AITER flash attention with FP8 KV cache matches reference on BF16-cast KV.
802 +
803 + cp_mha_gather_cache is called with dequant=True to dequantize FP8 to dtype
804 + before passing to flash_attn_varlen_func. We compare to ref_paged_attn on
805 + the dtype-cast KV cache with FP8 quantization tolerance.
806 +
807 + Exercises: VLLM_ROCM_USE_AITER, VLLM_ROCM_USE_AITER_MHA, FP8 KV path.
808 + """
809 + atol = 0.15
810 + rtol = 0.05
811 + _assert_aiter_supported()
812 + if not current_platform.supports_fp8():
813 + pytest.skip("FP8 not supported on this hardware")
814 +
815 + import aiter
816 +
817 + from vllm.v1.attention.backends.rocm_aiter_fa import cp_mha_gather_cache
818 +
819 + FP8_DTYPE = current_platform.fp8_dtype()
820 +
821 + set_random_seed(10)
822 +
823 + num_q_heads, num_kv_heads = 8, 8
824 + head_size = 128
825 + query_len, kv_len = 4, 128
826 + scale = head_size**-0.5
827 +
828 + query = torch.randn(query_len, num_q_heads, head_size, dtype=dtype)
829 + # FP8 KV cache; clamp to stay in FP8 representable range
830 + key_cache_fp8 = torch.clamp(
831 + torch.randn(NUM_BLOCKS, BLOCK_SIZE, num_kv_heads, head_size), -1.0, 1.0
832 + ).to(FP8_DTYPE)
833 + value_cache_fp8 = torch.clamp(
834 + torch.randn(NUM_BLOCKS, BLOCK_SIZE, num_kv_heads, head_size), -1.0, 1.0
835 + ).to(FP8_DTYPE)
836 +
837 + cu_query_lens = torch.tensor([0, query_len], dtype=torch.int32)
838 + cu_seq_lens = torch.tensor([0, kv_len], dtype=torch.int32)
839 + max_num_blocks = (kv_len + BLOCK_SIZE - 1) // BLOCK_SIZE
840 + block_tables = torch.randint(0, NUM_BLOCKS, (1, max_num_blocks), dtype=torch.int32)
841 + token_to_batch = torch.zeros(kv_len, dtype=torch.int32)
842 + seq_starts = torch.zeros(1, dtype=torch.int32)
843 +
844 + # Gather and dequantize FP8 KV to dtype
845 + gathered_key = torch.empty(kv_len, num_kv_heads, head_size, dtype=dtype)
846 + gathered_value = torch.empty_like(gathered_key)
847 + k_scales = torch.ones(1, dtype=torch.float32)
848 + v_scales = torch.ones(1, dtype=torch.float32)
849 +
850 + cp_mha_gather_cache(
851 + key_cache=key_cache_fp8,
852 + value_cache=value_cache_fp8,
853 + key=gathered_key,
854 + value=gathered_value,
855 + block_tables=block_tables,
856 + k_scales=k_scales,
857 + v_scales=v_scales,
858 + cu_seqlens_kv=cu_seq_lens,
859 + token_to_batch=token_to_batch,
860 + seq_starts=seq_starts,
861 + dequant=True,
862 + kv_cache_layout="NHD",
863 + total_tokens=kv_len,
864 + )
865 +
866 + output = torch.empty_like(query)
867 + aiter.flash_attn_varlen_func(
868 + q=query,
869 + k=gathered_key,
870 + v=gathered_value,
871 + cu_seqlens_q=cu_query_lens,
872 + cu_seqlens_k=cu_seq_lens,
873 + max_seqlen_q=query_len,
874 + max_seqlen_k=kv_len,
875 + min_seqlen_q=1,
876 + dropout_p=0.0,
877 + softmax_scale=scale,
878 + causal=True,
879 + window_size=(-1, -1),
880 + alibi_slopes=None,
881 + return_lse=False,
882 + out=output,
883 + )
884 +
885 + # Reference: ref_paged_attn on dtype-cast KV (simulates perfect dequant)
886 + key_cache_ref = key_cache_fp8.to(dtype)
887 + value_cache_ref = value_cache_fp8.to(dtype)
888 + ref = ref_paged_attn(
889 + query=query,
890 + key_cache=key_cache_ref,
891 + value_cache=value_cache_ref,
892 + query_lens=[query_len],
893 + kv_lens=[kv_len],
894 + block_tables=block_tables,
895 + scale=scale,
896 + )
897 +
898 + # FP8 quantization + dequantization introduces noise
899 + _print_close_stats(
900 + f"varlen_fp8_kv dtype={dtype} query_len={query_len} kv_len={kv_len}",
901 + output,
902 + ref,
903 + atol=atol,
904 + rtol=rtol,
905 + )
906 + torch.testing.assert_close(output, ref, atol=atol, rtol=rtol)