| 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 specs that cannot be gathered |
| 314 | correctly by the ROCm kernel.""" |
| 315 | from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend |
| 316 | from vllm.v1.kv_cache_interface import FullAttentionSpec |
| 317 | |
| 318 | with pytest.raises(ValueError, match="Block size must be a multiple of 16"): |
| 319 | AiterFlashAttentionBackend.customize_spec( |
| 320 | FullAttentionSpec( |
| 321 | block_size=15, num_kv_heads=8, head_size=128, dtype=torch.bfloat16 |
| 322 | ) |
| 323 | ) |
| 324 | # block_size == 1 is the per-token page-size probe and must pass through. |
| 325 | AiterFlashAttentionBackend.customize_spec( |
| 326 | FullAttentionSpec( |
| 327 | block_size=1, num_kv_heads=8, head_size=128, dtype=torch.bfloat16 |
| 328 | ) |
| 329 | ) |
| 330 | |
| 331 | |
| 332 | def test_aiter_mha_backend_supports_compute_capability_matches_mi3xx_probe(): |
| 333 | """The backend should trust the ROCm MI3xx probe instead of the raw torch |
| 334 | capability tuple.""" |
| 335 | from vllm.platforms.interface import DeviceCapability |
| 336 | from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend |
| 337 | |
| 338 | assert ( |
| 339 | AiterFlashAttentionBackend.supports_compute_capability(DeviceCapability(0, 0)) |
| 340 | is on_mi3xx() |
| 341 | ) |
| 342 | |
| 343 | |
| 344 | # Env and platform gate tests --------------------------------------------- |
| 345 | @pytest.mark.skipif(not on_mi3xx(), reason="MI300/MI350 ROCm only") |
| 346 | @pytest.mark.parametrize( |
| 347 | ("use_aiter", "use_mha", "expected"), |
| 348 | [ |
| 349 | (True, True, True), |
| 350 | (True, False, False), |
| 351 | (False, True, False), |
| 352 | (False, False, False), |
| 353 | ], |
| 354 | ) |
| 355 | def test_aiter_mha_env_flags_control_enablement( |
| 356 | use_aiter, use_mha, expected, monkeypatch |
| 357 | ): |
| 358 | """Both the global AITER flag and the MHA-specific flag must be enabled |
| 359 | before the ROCm MHA path is considered active.""" |
| 360 | from vllm._aiter_ops import rocm_aiter_ops |
| 361 | |
| 362 | _assert_aiter_supported() |
| 363 | |
| 364 | with monkeypatch.context() as mp: |
| 365 | mp.setenv("VLLM_ROCM_USE_AITER", "1" if use_aiter else "0") |
| 366 | mp.setenv("VLLM_ROCM_USE_AITER_MHA", "1" if use_mha else "0") |
| 367 | _reload_envs() |
| 368 | rocm_aiter_ops.refresh_env_variables() |
| 369 | assert rocm_aiter_ops.is_mha_enabled() is expected |
| 370 | |
| 371 | _reload_envs() |
| 372 | rocm_aiter_ops.refresh_env_variables() |
| 373 | |
| 374 | |
| 375 | def test_aiter_mha_platform_gate_matches_install_and_arch(): |
| 376 | """The global AITER availability check should only open on ROCm MI3xx when |
| 377 | the aiter package is installed.""" |
| 378 | from vllm._aiter_ops import IS_AITER_FOUND, is_aiter_found_and_supported |
| 379 | |
| 380 | assert is_aiter_found_and_supported() is ( |
| 381 | current_platform.is_rocm() and on_mi3xx() and IS_AITER_FOUND |
| 382 | ) |
| 383 | |
| 384 | |
| 385 | # Kernel path tests ------------------------------------------------------- |
| 386 | def _run_single_seq_flash_attn_case( |
| 387 | *, |
| 388 | head_size: int, |
| 389 | num_heads: tuple[int, int], |
| 390 | seq_lens: tuple[int, int], |
| 391 | dtype: torch.dtype, |
| 392 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 393 | import aiter |
| 394 | |
| 395 | from vllm.v1.attention.backends.rocm_aiter_fa import cp_mha_gather_cache |
| 396 | |
| 397 | set_random_seed(0) |
| 398 | |
| 399 | num_q_heads, num_kv_heads = num_heads |
| 400 | query_len, kv_len = seq_lens |
| 401 | scale = head_size**-0.5 |
| 402 | |
| 403 | query = torch.randn(query_len, num_q_heads, head_size, dtype=dtype) |
| 404 | key_cache = torch.randn( |
| 405 | NUM_BLOCKS, BLOCK_SIZE, num_kv_heads, head_size, dtype=dtype |
| 406 | ) |
| 407 | value_cache = torch.randn_like(key_cache) |
| 408 | |
| 409 | cu_query_lens = torch.tensor([0, query_len], dtype=torch.int32).cumsum( |
| 410 | dim=0, dtype=torch.int32 |
| 411 | ) |
| 412 | cu_seq_lens = torch.tensor([0, kv_len], dtype=torch.int32).cumsum( |
| 413 | dim=0, dtype=torch.int32 |
| 414 | ) |
| 415 | max_num_blocks = (kv_len + BLOCK_SIZE - 1) // BLOCK_SIZE |
| 416 | block_tables = torch.randint(0, NUM_BLOCKS, (1, max_num_blocks), dtype=torch.int32) |
| 417 | |
| 418 | token_to_batch = torch.zeros(kv_len, dtype=torch.int32) |
| 419 | seq_starts = torch.zeros(1, dtype=torch.int32) |
| 420 | gathered_key = torch.empty(kv_len, num_kv_heads, head_size, dtype=dtype) |
| 421 | gathered_value = torch.empty_like(gathered_key) |
| 422 | |
| 423 | cp_mha_gather_cache( |
| 424 | key_cache=key_cache, |
| 425 | value_cache=value_cache, |
| 426 | key=gathered_key, |
| 427 | value=gathered_value, |
| 428 | block_tables=block_tables, |
| 429 | k_scales=torch.ones(1, dtype=torch.float32), |
| 430 | v_scales=torch.ones(1, dtype=torch.float32), |
| 431 | cu_seqlens_kv=cu_seq_lens, |
| 432 | token_to_batch=token_to_batch, |
| 433 | seq_starts=seq_starts, |
| 434 | dequant=False, |
| 435 | kv_cache_layout="NHD", |
| 436 | total_tokens=kv_len, |
| 437 | ) |
| 438 | |
| 439 | output = torch.empty_like(query) |
| 440 | aiter.flash_attn_varlen_func( |
| 441 | q=query, |
| 442 | k=gathered_key, |
| 443 | v=gathered_value, |
| 444 | cu_seqlens_q=cu_query_lens, |
| 445 | cu_seqlens_k=cu_seq_lens, |
| 446 | max_seqlen_q=query_len, |
| 447 | max_seqlen_k=kv_len, |
| 448 | min_seqlen_q=1, |
| 449 | dropout_p=0.0, |
| 450 | softmax_scale=scale, |
| 451 | causal=True, |
| 452 | window_size=(-1, -1), |
| 453 | alibi_slopes=None, |
| 454 | return_lse=False, |
| 455 | out=output, |
| 456 | ) |
| 457 | |
| 458 | ref = ref_paged_attn( |
| 459 | query=query, |
| 460 | key_cache=key_cache, |
| 461 | value_cache=value_cache, |
| 462 | query_lens=[query_len], |
| 463 | kv_lens=[kv_len], |
| 464 | block_tables=block_tables, |
| 465 | scale=scale, |
| 466 | ) |
| 467 | return output, ref |
| 468 | |
| 469 | |
| 470 | @pytest.mark.skipif(not on_mi3xx(), reason="MI300/MI350 ROCm only") |
| 471 | @pytest.mark.parametrize("head_size", HEAD_SIZES) |
| 472 | @pytest.mark.parametrize("dtype", DTYPES) |
| 473 | @pytest.mark.parametrize("num_heads", HEAD_SIZE_TEST_NUM_HEADS_PAIRS) |
| 474 | @pytest.mark.parametrize("seq_lens", SEQ_LENS) |
| 475 | def test_aiter_fa_head_sizes(head_size, dtype, num_heads, seq_lens): |
| 476 | """AITER flash attention should stay accurate across supported head sizes.""" |
| 477 | atol = 1.5e-2 |
| 478 | rtol = 1e-2 |
| 479 | _assert_aiter_supported() |
| 480 | |
| 481 | output, ref = _run_single_seq_flash_attn_case( |
| 482 | head_size=head_size, |
| 483 | num_heads=num_heads, |
| 484 | seq_lens=seq_lens, |
| 485 | dtype=dtype, |
| 486 | ) |
| 487 | |
| 488 | _print_close_stats( |
| 489 | "head_sizes " |
| 490 | f"dtype={dtype} head_size={head_size} " |
| 491 | f"num_heads={num_heads} seq_lens={seq_lens}", |
| 492 | output, |
| 493 | ref, |
| 494 | atol=atol, |
| 495 | rtol=rtol, |
| 496 | ) |
| 497 | torch.testing.assert_close(output, ref, atol=atol, rtol=rtol) |
| 498 | |
| 499 | |
| 500 | @pytest.mark.skipif(not on_mi3xx(), reason="MI300/MI350 ROCm only") |
| 501 | @pytest.mark.parametrize("head_size", HEAD_SIZES) |
| 502 | @pytest.mark.parametrize("num_heads", NUM_HEADS_PAIRS) |
| 503 | @pytest.mark.parametrize("seq_lens", SEQ_LENS) |
| 504 | @pytest.mark.parametrize("dtype", DTYPES) |
| 505 | def test_aiter_mha_varlen_paged_kv(head_size, num_heads, seq_lens, dtype): |
| 506 | """Test AITER flash attention varlen with paged KV cache. |
| 507 | |
| 508 | Exercises: VLLM_ROCM_USE_AITER, VLLM_ROCM_USE_AITER_MHA |
| 509 | """ |
| 510 | atol = 1.5e-2 |
| 511 | rtol = 1e-2 |
| 512 | _assert_aiter_supported() |
| 513 | |
| 514 | output, ref = _run_single_seq_flash_attn_case( |
| 515 | head_size=head_size, |
| 516 | num_heads=num_heads, |
| 517 | seq_lens=seq_lens, |
| 518 | dtype=dtype, |
| 519 | ) |
| 520 | |
| 521 | _print_close_stats( |
| 522 | "varlen_paged_kv " |
| 523 | f"dtype={dtype} head_size={head_size} " |
| 524 | f"num_heads={num_heads} seq_lens={seq_lens}", |
| 525 | output, |
| 526 | ref, |
| 527 | atol=atol, |
| 528 | rtol=rtol, |
| 529 | ) |
| 530 | torch.testing.assert_close(output, ref, atol=atol, rtol=rtol) |
| 531 | |
| 532 | |
| 533 | @pytest.mark.skipif(not on_mi3xx(), reason="MI300/MI350 ROCm only") |
| 534 | @pytest.mark.parametrize("num_heads", NUM_HEADS_PAIRS) |
| 535 | @pytest.mark.parametrize("head_size", HEAD_SIZES) |
| 536 | @pytest.mark.parametrize("dtype", DTYPES) |
| 537 | def test_aiter_mha_multi_batch(num_heads, head_size, dtype): |
| 538 | """Test AITER flash attention with multiple sequences in a batch.""" |
| 539 | atol = 1.5e-2 |
| 540 | rtol = 1e-2 |
| 541 | _assert_aiter_supported() |
| 542 | import aiter |
| 543 | |
| 544 | from vllm.v1.attention.backends.rocm_aiter_fa import cp_mha_gather_cache |
| 545 | |
| 546 | set_random_seed(42) |
| 547 | |
| 548 | num_q_heads, num_kv_heads = num_heads |
| 549 | seq_lens = [(4, 128), (2, 256), (8, 64)] |
| 550 | query_lens = [q for q, _ in seq_lens] |
| 551 | kv_lens = [k for _, k in seq_lens] |
| 552 | num_seqs = len(seq_lens) |
| 553 | scale = head_size**-0.5 |
| 554 | |
| 555 | total_q = sum(query_lens) |
| 556 | total_kv = sum(kv_lens) |
| 557 | |
| 558 | query = torch.randn(total_q, num_q_heads, head_size, dtype=dtype) |
| 559 | key_cache = torch.randn( |
| 560 | NUM_BLOCKS, BLOCK_SIZE, num_kv_heads, head_size, dtype=dtype |
| 561 | ) |
| 562 | value_cache = torch.randn_like(key_cache) |
| 563 | |
| 564 | cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum( |
| 565 | dim=0, dtype=torch.int32 |
| 566 | ) |
| 567 | cu_seq_lens = torch.tensor([0] + kv_lens, dtype=torch.int32).cumsum( |
| 568 | dim=0, dtype=torch.int32 |
| 569 | ) |
| 570 | |
| 571 | max_kv_len = max(kv_lens) |
| 572 | max_num_blocks = (max_kv_len + BLOCK_SIZE - 1) // BLOCK_SIZE |
| 573 | block_tables = torch.randint( |
| 574 | 0, NUM_BLOCKS, (num_seqs, max_num_blocks), dtype=torch.int32 |
| 575 | ) |
| 576 | |
| 577 | token_to_batch = torch.empty(total_kv, dtype=torch.int32) |
| 578 | seq_starts = torch.zeros(num_seqs, dtype=torch.int32) |
| 579 | tok_idx = 0 |
| 580 | for b, kl in enumerate(kv_lens): |
| 581 | token_to_batch[tok_idx : tok_idx + kl] = b |
| 582 | tok_idx += kl |
| 583 | |
| 584 | gathered_key = torch.empty(total_kv, num_kv_heads, head_size, dtype=dtype) |
| 585 | gathered_value = torch.empty_like(gathered_key) |
| 586 | |
| 587 | cp_mha_gather_cache( |
| 588 | key_cache=key_cache, |
| 589 | value_cache=value_cache, |
| 590 | key=gathered_key, |
| 591 | value=gathered_value, |
| 592 | block_tables=block_tables, |
| 593 | k_scales=torch.ones(1, dtype=torch.float32), |
| 594 | v_scales=torch.ones(1, dtype=torch.float32), |
| 595 | cu_seqlens_kv=cu_seq_lens, |
| 596 | token_to_batch=token_to_batch, |
| 597 | seq_starts=seq_starts, |
| 598 | dequant=False, |
| 599 | kv_cache_layout="NHD", |
| 600 | total_tokens=total_kv, |
| 601 | ) |
| 602 | |
| 603 | output = torch.empty_like(query) |
| 604 | aiter.flash_attn_varlen_func( |
| 605 | q=query, |
| 606 | k=gathered_key, |
| 607 | v=gathered_value, |
| 608 | cu_seqlens_q=cu_query_lens, |
| 609 | cu_seqlens_k=cu_seq_lens, |
| 610 | max_seqlen_q=max(query_lens), |
| 611 | max_seqlen_k=max_kv_len, |
| 612 | min_seqlen_q=1, |
| 613 | dropout_p=0.0, |
| 614 | softmax_scale=scale, |
| 615 | causal=True, |
| 616 | window_size=(-1, -1), |
| 617 | alibi_slopes=None, |
| 618 | return_lse=False, |
| 619 | out=output, |
| 620 | ) |
| 621 | |
| 622 | ref = ref_paged_attn( |
| 623 | query=query, |
| 624 | key_cache=key_cache, |
| 625 | value_cache=value_cache, |
| 626 | query_lens=query_lens, |
| 627 | kv_lens=kv_lens, |
| 628 | block_tables=block_tables, |
| 629 | scale=scale, |
| 630 | ) |
| 631 | |
| 632 | _print_close_stats( |
| 633 | "multi_batch " |
| 634 | f"dtype={dtype} head_size={head_size} " |
| 635 | f"num_heads={num_heads} seq_lens={seq_lens}", |
| 636 | output, |
| 637 | ref, |
| 638 | atol=atol, |
| 639 | rtol=rtol, |
| 640 | ) |
| 641 | torch.testing.assert_close(output, ref, atol=atol, rtol=rtol) |
| 642 | |
| 643 | |
| 644 | @pytest.mark.skipif(not on_mi3xx(), reason="MI300/MI350 ROCm only") |
| 645 | @pytest.mark.parametrize("num_blocks", DIRECT_NUM_BLOCKS) |
| 646 | def test_aiter_fa_large_block_table_matches_reference(num_blocks): |
| 647 | """The direct paged-KV path should stay stable for both normal and very |
| 648 | large block tables.""" |
| 649 | atol = 2e-2 |
| 650 | rtol = 2e-2 |
| 651 | _assert_aiter_supported() |
| 652 | |
| 653 | output, ref = _run_direct_flash_attn_case( |
| 654 | seq_lens=DIRECT_SEQ_LENS, |
| 655 | num_blocks=num_blocks, |
| 656 | ) |
| 657 | _print_close_stats( |
| 658 | f"direct_varlen_paged_kv num_blocks={num_blocks}", |
| 659 | output, |
| 660 | ref, |
| 661 | atol=atol, |
| 662 | rtol=rtol, |
| 663 | ) |
| 664 | torch.testing.assert_close(output, ref, atol=atol, rtol=rtol) |
| 665 | |
| 666 | |
| 667 | @pytest.mark.skipif(not on_mi3xx(), reason="MI300/MI350 ROCm only") |
| 668 | def test_aiter_fa_sliding_window_matches_reference(): |
| 669 | """The direct kernel should respect the same sliding-window causal mask as |
| 670 | the naive reference implementation.""" |
| 671 | atol = 2e-2 |
| 672 | rtol = 2e-2 |
| 673 | _assert_aiter_supported() |
| 674 | |
| 675 | output, ref = _run_direct_flash_attn_case( |
| 676 | seq_lens=[(8, 523), (24, 37), (3, 2011)], |
| 677 | num_blocks=2048, |
| 678 | sliding_window=256, |
| 679 | ) |
| 680 | _print_close_stats( |
| 681 | "sliding_window num_blocks=2048 window=256", |
| 682 | output, |
| 683 | ref, |
| 684 | atol=atol, |
| 685 | rtol=rtol, |
| 686 | ) |
| 687 | torch.testing.assert_close(output, ref, atol=atol, rtol=rtol) |
| 688 | |
| 689 | |
| 690 | @pytest.mark.skipif(not on_mi3xx(), reason="MI300/MI350 ROCm only") |
| 691 | @pytest.mark.parametrize( |
| 692 | "shuffle,sliding_window", [(False, None), (False, 16), (True, None)] |
| 693 | ) |
| 694 | @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) |
| 695 | @pytest.mark.parametrize("batch", ["prefill", "extend", "mixed"]) |
| 696 | def test_aiter_fa_shared_kv_matches_reference( |
| 697 | monkeypatch, shuffle, sliding_window, kv_cache_dtype, batch |
| 698 | ): |
| 699 | """Shared layers read cached K/V for prefill and extend without modifying it.""" |
| 700 | from tests.v1.attention.utils import ( |
| 701 | BatchSpec, |
| 702 | create_common_attn_metadata, |
| 703 | create_standard_kv_cache_spec, |
| 704 | create_vllm_config, |
| 705 | ) |
| 706 | from vllm._aiter_ops import rocm_aiter_ops |
| 707 | from vllm.config import set_current_vllm_config |
| 708 | from vllm.model_executor.layers.attention import Attention |
| 709 | from vllm.v1.attention.backends.rocm_aiter_fa import ( |
| 710 | AiterFlashAttentionBackend, |
| 711 | AiterFlashAttentionMetadataBuilder, |
| 712 | ) |
| 713 | |
| 714 | _assert_aiter_supported() |
| 715 | monkeypatch.setattr(rocm_aiter_ops, "_SHUFFLE_KV_CACHE_ENABLED", shuffle) |
| 716 | set_random_seed(42) |
| 717 | batch_spec = { |
| 718 | "prefill": BatchSpec(seq_lens=[32, 24], query_lens=[32, 24]), |
| 719 | "extend": BatchSpec(seq_lens=[80, 56], query_lens=[32, 24]), |
| 720 | "mixed": BatchSpec(seq_lens=[48, 80, 24], query_lens=[1, 32, 24]), |
| 721 | }[batch] |
| 722 | config = create_vllm_config( |
| 723 | model_name="Qwen/Qwen3-0.6B", |
| 724 | dtype="bfloat16", |
| 725 | max_num_batched_tokens=1024, |
| 726 | ) |
| 727 | config.cache_config.cache_dtype = kv_cache_dtype |
| 728 | num_heads = config.model_config.get_num_attention_heads(config.parallel_config) |
| 729 | num_kv_heads = config.model_config.get_num_kv_heads(config.parallel_config) |
| 730 | head_size = config.model_config.get_head_size() |
| 731 | scale = head_size**-0.5 |
| 732 | with set_current_vllm_config(config): |
| 733 | target = Attention( |
| 734 | num_heads, |
| 735 | head_size, |
| 736 | scale, |
| 737 | num_kv_heads, |
| 738 | cache_config=config.cache_config, |
| 739 | per_layer_sliding_window=sliding_window, |
| 740 | prefix="target", |
| 741 | attn_backend=AiterFlashAttentionBackend, |
| 742 | ) |
| 743 | layer = Attention( |
| 744 | num_heads, |
| 745 | head_size, |
| 746 | scale, |
| 747 | num_kv_heads, |
| 748 | cache_config=config.cache_config, |
| 749 | per_layer_sliding_window=sliding_window, |
| 750 | kv_sharing_target_layer_name="target", |
| 751 | prefix="shared", |
| 752 | attn_backend=AiterFlashAttentionBackend, |
| 753 | ) |
| 754 | builder = AiterFlashAttentionMetadataBuilder( |
| 755 | create_standard_kv_cache_spec(config), |
| 756 | ["target", "shared"], |
| 757 | config, |
| 758 | torch.device("cuda"), |
| 759 | ) |
| 760 | common = create_common_attn_metadata( |
| 761 | batch_spec, BLOCK_SIZE, torch.device("cuda"), max_block_idx=32 |
| 762 | ) |
| 763 | cache_dtype = ( |
| 764 | current_platform.fp8_dtype() if kv_cache_dtype == "fp8" else torch.bfloat16 |
| 765 | ) |
| 766 | kv_cache = ( |
| 767 | torch.empty( |
| 768 | 2, 32, BLOCK_SIZE, num_kv_heads * head_size, dtype=cache_dtype |
| 769 | ).transpose(0, 1) |
| 770 | if shuffle |
| 771 | else torch.empty( |
| 772 | 32, num_kv_heads, BLOCK_SIZE, 2 * head_size, dtype=cache_dtype |
| 773 | ) |
| 774 | ) |
| 775 | target.kv_cache = kv_cache |
| 776 | layer.kv_cache = kv_cache |
| 777 | k_scale = 0.5 if kv_cache_dtype == "fp8" and not shuffle else 1.0 |
| 778 | v_scale = 0.25 if kv_cache_dtype == "fp8" and not shuffle else 1.0 |
| 779 | for attn_layer in (target, layer): |
| 780 | attn_layer._k_scale.fill_(k_scale) |
| 781 | attn_layer._v_scale.fill_(v_scale) |
| 782 | key = torch.randn( |
| 783 | 32 * BLOCK_SIZE, num_kv_heads, head_size, dtype=torch.bfloat16 |
| 784 | ) |
| 785 | value = torch.randn_like(key) |
| 786 | from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionImpl |
| 787 | |
| 788 | assert isinstance(target.impl, AiterFlashAttentionImpl) |
| 789 | target.impl.do_kv_cache_update( |
| 790 | target, |
| 791 | key, |
| 792 | value, |
| 793 | kv_cache, |
| 794 | torch.arange(32 * BLOCK_SIZE, dtype=torch.int64), |
| 795 | ) |
| 796 | key_cache = ((key / k_scale).to(cache_dtype).to(key.dtype) * k_scale).reshape( |
| 797 | 32, BLOCK_SIZE, num_kv_heads, head_size |
| 798 | ) |
| 799 | value_cache = ( |
| 800 | (value / v_scale).to(cache_dtype).to(value.dtype) * v_scale |
| 801 | ).reshape(32, BLOCK_SIZE, num_kv_heads, head_size) |
| 802 | with torch.device("cpu"): |
| 803 | metadata = builder.build(0, common) |
| 804 | cache_before = kv_cache.view(torch.uint8).clone() |
| 805 | query = torch.randn( |
| 806 | batch_spec.compute_num_tokens(), |
| 807 | num_heads, |
| 808 | head_size, |
| 809 | dtype=torch.bfloat16, |
| 810 | ) |
| 811 | output = torch.empty_like(query) |
| 812 | layer.impl.forward(layer, query, None, None, kv_cache, metadata, output) |
| 813 | expected = ref_paged_attn( |
| 814 | query, |
| 815 | key_cache.contiguous(), |
| 816 | value_cache.contiguous(), |
| 817 | batch_spec.query_lens, |
| 818 | batch_spec.seq_lens, |
| 819 | common.block_table_tensor, |
| 820 | scale, |
| 821 | sliding_window, |
| 822 | ) |
| 823 | num_decode_tokens = metadata.num_decode_tokens |
| 824 | if kv_cache_dtype == "fp8" and num_decode_tokens: |
| 825 | # Use the established FP8 decode tolerance; gathered prefill and |
| 826 | # extend tokens must still satisfy the tighter BF16 tolerance. |
| 827 | torch.testing.assert_close( |
| 828 | output[:num_decode_tokens], |
| 829 | expected[:num_decode_tokens], |
| 830 | atol=6e-2, |
| 831 | rtol=1e-1, |
| 832 | ) |
| 833 | output = output[num_decode_tokens:] |
| 834 | expected = expected[num_decode_tokens:] |
| 835 | torch.testing.assert_close(output, expected, atol=2e-2, rtol=2e-2) |
| 836 | torch.testing.assert_close( |
| 837 | kv_cache.view(torch.uint8), cache_before, atol=0, rtol=0 |
| 838 | ) |
| 839 | |
| 840 | |
| 841 | # Decode path test -------------------------------------------------------- |
| 842 | @pytest.mark.skipif(not on_mi3xx(), reason="MI300/MI350 ROCm only") |
| 843 | @pytest.mark.parametrize( |
| 844 | "dtype", |
| 845 | [ |
| 846 | torch.bfloat16, |
| 847 | pytest.param( |
| 848 | torch.float16, |
| 849 | marks=pytest.mark.xfail( |
| 850 | reason=( |
| 851 | "aiter bug #2229: flash_attn_varlen_func currently " |
| 852 | "miscomputes FP16 single-token decode on MI3xx " |
| 853 | "(validated on gfx950; max abs diff 1.95 vs atol 1.5e-2). " |
| 854 | "Remove xfail when the AITER decode kernel is fixed. " |
| 855 | "https://github.com/ROCm/aiter/issues/2229" |
| 856 | ), |
| 857 | ), |
| 858 | ), |
| 859 | ], |
| 860 | ) |
| 861 | def test_aiter_mha_decode_single_token(dtype): |
| 862 | """Test AITER MHA for decode (single query token per sequence). |
| 863 | |
| 864 | BF16 is the working reference configuration here. FP16 remains xfail until |
| 865 | the upstream AITER single-token decode bug is fixed. |
| 866 | """ |
| 867 | atol = 1.5e-2 |
| 868 | rtol = 1e-2 |
| 869 | _assert_aiter_supported() |
| 870 | import aiter |
| 871 | |
| 872 | from vllm.v1.attention.backends.rocm_aiter_fa import cp_mha_gather_cache |
| 873 | |
| 874 | set_random_seed(0) |
| 875 | |
| 876 | num_q_heads, num_kv_heads = 8, 8 |
| 877 | head_size = 128 |
| 878 | kv_len = 512 |
| 879 | scale = head_size**-0.5 |
| 880 | |
| 881 | query = torch.randn(1, num_q_heads, head_size, dtype=dtype) |
| 882 | key_cache = torch.randn( |
| 883 | NUM_BLOCKS, BLOCK_SIZE, num_kv_heads, head_size, dtype=dtype |
| 884 | ) |
| 885 | value_cache = torch.randn_like(key_cache) |
| 886 | |
| 887 | cu_query_lens = torch.tensor([0, 1], dtype=torch.int32) |
| 888 | cu_seq_lens = torch.tensor([0, kv_len], dtype=torch.int32) |
| 889 | |
| 890 | max_num_blocks = (kv_len + BLOCK_SIZE - 1) // BLOCK_SIZE |
| 891 | block_tables = torch.randint(0, NUM_BLOCKS, (1, max_num_blocks), dtype=torch.int32) |
| 892 | |
| 893 | token_to_batch = torch.zeros(kv_len, dtype=torch.int32) |
| 894 | seq_starts = torch.zeros(1, dtype=torch.int32) |
| 895 | gathered_key = torch.empty(kv_len, num_kv_heads, head_size, dtype=dtype) |
| 896 | gathered_value = torch.empty_like(gathered_key) |
| 897 | |
| 898 | cp_mha_gather_cache( |
| 899 | key_cache=key_cache, |
| 900 | value_cache=value_cache, |
| 901 | key=gathered_key, |
| 902 | value=gathered_value, |
| 903 | block_tables=block_tables, |
| 904 | k_scales=torch.ones(1, dtype=torch.float32), |
| 905 | v_scales=torch.ones(1, dtype=torch.float32), |
| 906 | cu_seqlens_kv=cu_seq_lens, |
| 907 | token_to_batch=token_to_batch, |
| 908 | seq_starts=seq_starts, |
| 909 | dequant=False, |
| 910 | kv_cache_layout="NHD", |
| 911 | total_tokens=kv_len, |
| 912 | ) |
| 913 | |
| 914 | output = torch.empty_like(query) |
| 915 | aiter.flash_attn_varlen_func( |
| 916 | q=query, |
| 917 | k=gathered_key, |
| 918 | v=gathered_value, |
| 919 | cu_seqlens_q=cu_query_lens, |
| 920 | cu_seqlens_k=cu_seq_lens, |
| 921 | max_seqlen_q=1, |
| 922 | max_seqlen_k=kv_len, |
| 923 | min_seqlen_q=1, |
| 924 | dropout_p=0.0, |
| 925 | softmax_scale=scale, |
| 926 | causal=True, |
| 927 | window_size=(-1, -1), |
| 928 | alibi_slopes=None, |
| 929 | return_lse=False, |
| 930 | out=output, |
| 931 | ) |
| 932 | |
| 933 | ref = ref_paged_attn( |
| 934 | query=query, |
| 935 | key_cache=key_cache, |
| 936 | value_cache=value_cache, |
| 937 | query_lens=[1], |
| 938 | kv_lens=[kv_len], |
| 939 | block_tables=block_tables, |
| 940 | scale=scale, |
| 941 | ) |
| 942 | |
| 943 | _print_close_stats( |
| 944 | f"decode_single_token dtype={dtype} kv_len={kv_len}", |
| 945 | output, |
| 946 | ref, |
| 947 | atol=atol, |
| 948 | rtol=rtol, |
| 949 | ) |
| 950 | torch.testing.assert_close(output, ref, atol=atol, rtol=rtol) |
| 951 | |
| 952 | |
| 953 | # FP8 KV cache test ------------------------------------------------------- |
| 954 | @pytest.mark.skipif(not on_mi3xx(), reason="MI300/MI350 ROCm only") |
| 955 | @pytest.mark.parametrize("dtype", DTYPES) |
| 956 | def test_aiter_mha_varlen_fp8_kv(dtype): |
| 957 | """AITER flash attention with FP8 KV cache matches reference on BF16-cast KV. |
| 958 | |
| 959 | The reference dequantizes the same FP8 KV tensors the kernel reads, so |
| 960 | e4m3 rounding cancels on both sides; the residual is bf16/fp16 kernel |
| 961 | rounding (~few millidecimals), not fp8-scale quantization error. |
| 962 | """ |
| 963 | if dtype == torch.bfloat16: |
| 964 | atol, rtol = 6e-3, 1e-2 |
| 965 | else: |
| 966 | atol, rtol = 7e-4, 1e-3 |
| 967 | _assert_aiter_supported() |
| 968 | if not current_platform.supports_fp8(): |
| 969 | pytest.skip("FP8 not supported on this hardware") |
| 970 | |
| 971 | import aiter |
| 972 | |
| 973 | from vllm.v1.attention.backends.rocm_aiter_fa import cp_mha_gather_cache |
| 974 | |
| 975 | FP8_DTYPE = current_platform.fp8_dtype() |
| 976 | |
| 977 | set_random_seed(10) |
| 978 | |
| 979 | num_q_heads, num_kv_heads = 8, 8 |
| 980 | head_size = 128 |
| 981 | query_len, kv_len = 4, 128 |
| 982 | scale = head_size**-0.5 |
| 983 | |
| 984 | query = torch.randn(query_len, num_q_heads, head_size, dtype=dtype) |
| 985 | # FP8 KV cache; clamp to stay in FP8 representable range |
| 986 | key_cache_fp8 = torch.clamp( |
| 987 | torch.randn(NUM_BLOCKS, BLOCK_SIZE, num_kv_heads, head_size), -1.0, 1.0 |
| 988 | ).to(FP8_DTYPE) |
| 989 | value_cache_fp8 = torch.clamp( |
| 990 | torch.randn(NUM_BLOCKS, BLOCK_SIZE, num_kv_heads, head_size), -1.0, 1.0 |
| 991 | ).to(FP8_DTYPE) |
| 992 | |
| 993 | cu_query_lens = torch.tensor([0, query_len], dtype=torch.int32) |
| 994 | cu_seq_lens = torch.tensor([0, kv_len], dtype=torch.int32) |
| 995 | max_num_blocks = (kv_len + BLOCK_SIZE - 1) // BLOCK_SIZE |
| 996 | block_tables = torch.randint(0, NUM_BLOCKS, (1, max_num_blocks), dtype=torch.int32) |
| 997 | token_to_batch = torch.zeros(kv_len, dtype=torch.int32) |
| 998 | seq_starts = torch.zeros(1, dtype=torch.int32) |
| 999 | |
| 1000 | # Gather and dequantize FP8 KV to dtype |
| 1001 | gathered_key = torch.empty(kv_len, num_kv_heads, head_size, dtype=dtype) |
| 1002 | gathered_value = torch.empty_like(gathered_key) |
| 1003 | k_scales = torch.ones(1, dtype=torch.float32) |
| 1004 | v_scales = torch.ones(1, dtype=torch.float32) |
| 1005 | |
| 1006 | cp_mha_gather_cache( |
| 1007 | key_cache=key_cache_fp8, |
| 1008 | value_cache=value_cache_fp8, |
| 1009 | key=gathered_key, |
| 1010 | value=gathered_value, |
| 1011 | block_tables=block_tables, |
| 1012 | k_scales=k_scales, |
| 1013 | v_scales=v_scales, |
| 1014 | cu_seqlens_kv=cu_seq_lens, |
| 1015 | token_to_batch=token_to_batch, |
| 1016 | seq_starts=seq_starts, |
| 1017 | dequant=True, |
| 1018 | kv_cache_layout="NHD", |
| 1019 | total_tokens=kv_len, |
| 1020 | ) |
| 1021 | |
| 1022 | output = torch.empty_like(query) |
| 1023 | aiter.flash_attn_varlen_func( |
| 1024 | q=query, |
| 1025 | k=gathered_key, |
| 1026 | v=gathered_value, |
| 1027 | cu_seqlens_q=cu_query_lens, |
| 1028 | cu_seqlens_k=cu_seq_lens, |
| 1029 | max_seqlen_q=query_len, |
| 1030 | max_seqlen_k=kv_len, |
| 1031 | min_seqlen_q=1, |
| 1032 | dropout_p=0.0, |
| 1033 | softmax_scale=scale, |
| 1034 | causal=True, |
| 1035 | window_size=(-1, -1), |
| 1036 | alibi_slopes=None, |
| 1037 | return_lse=False, |
| 1038 | out=output, |
| 1039 | ) |
| 1040 | |
| 1041 | # Reference: ref_paged_attn on dtype-cast KV (simulates perfect dequant) |
| 1042 | key_cache_ref = key_cache_fp8.to(dtype) |
| 1043 | value_cache_ref = value_cache_fp8.to(dtype) |
| 1044 | ref = ref_paged_attn( |
| 1045 | query=query, |
| 1046 | key_cache=key_cache_ref, |
| 1047 | value_cache=value_cache_ref, |
| 1048 | query_lens=[query_len], |
| 1049 | kv_lens=[kv_len], |
| 1050 | block_tables=block_tables, |
| 1051 | scale=scale, |
| 1052 | ) |
| 1053 | |
| 1054 | # FP8 quantization + dequantization cancels on the reference side; compare |
| 1055 | # kernel rounding at the working precision. |
| 1056 | _print_close_stats( |
| 1057 | f"varlen_fp8_kv dtype={dtype} query_len={query_len} kv_len={kv_len}", |
| 1058 | output, |
| 1059 | ref, |
| 1060 | atol=atol, |
| 1061 | rtol=rtol, |
| 1062 | ) |
| 1063 | torch.testing.assert_close(output, ref, atol=atol, rtol=rtol) |