[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>
This commit is contained in:
Andreas Karatzas
co-authored by
Micah Williamson
Micah Williamson
parent
7303c66f68
commit
31cd109f18
@@ -1,258 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
# Import AITER backend if on ROCm and aiter is available
|
||||
if current_platform.is_rocm():
|
||||
from vllm._aiter_ops import is_aiter_found_and_supported
|
||||
|
||||
if is_aiter_found_and_supported():
|
||||
import aiter
|
||||
|
||||
from vllm.v1.attention.backends.rocm_aiter_fa import cp_mha_gather_cache
|
||||
|
||||
NUM_HEADS = [(4, 4), (8, 2)]
|
||||
HEAD_SIZES = [128, 256]
|
||||
BLOCK_SIZES = [16]
|
||||
DTYPES = [torch.bfloat16]
|
||||
QDTYPES = [None]
|
||||
# one value large enough to test overflow in index calculation.
|
||||
# one value small enough to test the schema op check
|
||||
NUM_BLOCKS = [32768, 2048]
|
||||
|
||||
|
||||
def ref_paged_attn(
|
||||
query: torch.Tensor,
|
||||
key_cache: torch.Tensor,
|
||||
value_cache: torch.Tensor,
|
||||
query_lens: list[int],
|
||||
kv_lens: list[int],
|
||||
block_tables: torch.Tensor,
|
||||
scale: float,
|
||||
sliding_window: int | None = None,
|
||||
soft_cap: float | None = None,
|
||||
) -> torch.Tensor:
|
||||
num_seqs = len(query_lens)
|
||||
block_tables = block_tables.cpu().numpy()
|
||||
_, block_size, num_kv_heads, head_size = key_cache.shape
|
||||
|
||||
outputs: list[torch.Tensor] = []
|
||||
start_idx = 0
|
||||
for i in range(num_seqs):
|
||||
query_len = query_lens[i]
|
||||
kv_len = kv_lens[i]
|
||||
q = query[start_idx : start_idx + query_len]
|
||||
q *= scale
|
||||
|
||||
num_kv_blocks = (kv_len + block_size - 1) // block_size
|
||||
block_indices = block_tables[i, :num_kv_blocks]
|
||||
|
||||
k = key_cache[block_indices].view(-1, num_kv_heads, head_size)
|
||||
k = k[:kv_len]
|
||||
v = value_cache[block_indices].view(-1, num_kv_heads, head_size)
|
||||
v = v[:kv_len]
|
||||
|
||||
if q.shape[1] != k.shape[1]:
|
||||
k = torch.repeat_interleave(k, q.shape[1] // k.shape[1], dim=1)
|
||||
v = torch.repeat_interleave(v, q.shape[1] // v.shape[1], dim=1)
|
||||
attn = torch.einsum("qhd,khd->hqk", q, k).float()
|
||||
empty_mask = torch.ones(query_len, kv_len)
|
||||
mask = torch.triu(empty_mask, diagonal=kv_len - query_len + 1).bool()
|
||||
if sliding_window is not None:
|
||||
sliding_window_mask = (
|
||||
torch.triu(
|
||||
empty_mask, diagonal=kv_len - (query_len + sliding_window) + 1
|
||||
)
|
||||
.bool()
|
||||
.logical_not()
|
||||
)
|
||||
mask |= sliding_window_mask
|
||||
if soft_cap is not None:
|
||||
attn = soft_cap * torch.tanh(attn / soft_cap)
|
||||
attn.masked_fill_(mask, float("-inf"))
|
||||
attn = torch.softmax(attn, dim=-1).to(v.dtype)
|
||||
out = torch.einsum("hqk,khd->qhd", attn, v)
|
||||
|
||||
outputs.append(out)
|
||||
start_idx += query_len
|
||||
|
||||
return torch.cat(outputs, dim=0)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_platform.is_rocm(), reason="Only ROCm is supported")
|
||||
@pytest.mark.parametrize(
|
||||
"seq_lens", [[(10, 1328), (5, 18), (129, 463)], [(8, 523), (24, 37), (3, 2011)]]
|
||||
)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
|
||||
@pytest.mark.parametrize("sliding_window", [None, 256])
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("soft_cap", [None])
|
||||
@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
|
||||
@pytest.mark.parametrize("q_dtype", QDTYPES)
|
||||
@torch.inference_mode()
|
||||
def test_varlen_with_paged_kv(
|
||||
seq_lens: list[tuple[int, int]],
|
||||
num_heads: tuple[int, int],
|
||||
head_size: int,
|
||||
sliding_window: int | None,
|
||||
dtype: torch.dtype,
|
||||
block_size: int,
|
||||
soft_cap: float | None,
|
||||
num_blocks: int,
|
||||
q_dtype: torch.dtype | None,
|
||||
) -> None:
|
||||
from vllm._aiter_ops import is_aiter_found_and_supported
|
||||
|
||||
if not is_aiter_found_and_supported():
|
||||
pytest.skip("aiter package required for this test.")
|
||||
|
||||
torch.set_default_device("cuda")
|
||||
set_random_seed(0)
|
||||
num_seqs = len(seq_lens)
|
||||
query_lens = [x[0] for x in seq_lens]
|
||||
kv_lens = [x[1] for x in seq_lens]
|
||||
num_query_heads = num_heads[0]
|
||||
num_kv_heads = num_heads[1]
|
||||
assert num_query_heads % num_kv_heads == 0
|
||||
max_query_len = max(query_lens)
|
||||
max_kv_len = max(kv_lens)
|
||||
window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
|
||||
scale = head_size**-0.5
|
||||
|
||||
query = torch.randn(sum(query_lens), num_query_heads, head_size, dtype=dtype)
|
||||
key_cache = torch.randn(
|
||||
num_blocks, block_size, num_kv_heads, head_size, dtype=dtype
|
||||
)
|
||||
value_cache = torch.randn_like(key_cache)
|
||||
cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
|
||||
dim=0, dtype=torch.int32
|
||||
)
|
||||
|
||||
cu_seq_lens = torch.tensor([0] + kv_lens, dtype=torch.int32).cumsum(
|
||||
dim=0, dtype=torch.int32
|
||||
)
|
||||
# Save kv_lens as list before converting to tensor
|
||||
kv_lens_list = kv_lens
|
||||
kv_lens = torch.tensor(kv_lens, dtype=torch.int32)
|
||||
|
||||
max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
|
||||
block_tables = torch.randint(
|
||||
0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
|
||||
)
|
||||
|
||||
output = torch.empty_like(query)
|
||||
|
||||
maybe_quantized_query = query
|
||||
maybe_quantized_key_cache = key_cache
|
||||
maybe_quantized_value_cache = value_cache
|
||||
k_scale_tensor = None
|
||||
v_scale_tensor = None
|
||||
dequant = False
|
||||
|
||||
if q_dtype is not None:
|
||||
# QKV are drawn from N(0, 1): no need for a fp8 scaling factor
|
||||
maybe_quantized_query = query.to(q_dtype)
|
||||
maybe_quantized_key_cache = key_cache.to(q_dtype)
|
||||
maybe_quantized_value_cache = value_cache.to(q_dtype)
|
||||
dequant = True
|
||||
scale_shape = (num_seqs, num_kv_heads)
|
||||
|
||||
# For per-seq-per-head scales (matching AITER backend expectation)
|
||||
k_scale_tensor = torch.ones(scale_shape, dtype=torch.float32)
|
||||
v_scale_tensor = torch.ones(scale_shape, dtype=torch.float32)
|
||||
|
||||
# Prepare metadata for cp_mha_gather_cache
|
||||
# token_to_batch: maps each token to its batch index
|
||||
token_to_batch = torch.zeros(sum(kv_lens_list), dtype=torch.int32)
|
||||
seq_starts = torch.zeros(num_seqs, dtype=torch.int32)
|
||||
|
||||
token_idx = 0
|
||||
for batch_idx, kv_len in enumerate(kv_lens_list):
|
||||
token_to_batch[token_idx : token_idx + kv_len] = batch_idx
|
||||
seq_starts[batch_idx] = 0 # Assuming all sequences start at 0 in their blocks
|
||||
token_idx += kv_len
|
||||
|
||||
# Allocate buffers for gathered KV
|
||||
total_kv_tokens = sum(kv_lens_list)
|
||||
gathered_key = torch.empty(
|
||||
total_kv_tokens, num_kv_heads, head_size, dtype=maybe_quantized_key_cache.dtype
|
||||
)
|
||||
gathered_value = torch.empty(
|
||||
total_kv_tokens,
|
||||
num_kv_heads,
|
||||
head_size,
|
||||
dtype=maybe_quantized_value_cache.dtype,
|
||||
)
|
||||
|
||||
# Gather paged KV cache into contiguous tensors using triton kernel
|
||||
cp_mha_gather_cache(
|
||||
key_cache=maybe_quantized_key_cache,
|
||||
value_cache=maybe_quantized_value_cache,
|
||||
key=gathered_key,
|
||||
value=gathered_value,
|
||||
block_tables=block_tables,
|
||||
k_scales=k_scale_tensor
|
||||
if k_scale_tensor is not None
|
||||
else torch.ones(1, dtype=torch.float32),
|
||||
v_scales=v_scale_tensor
|
||||
if v_scale_tensor is not None
|
||||
else torch.ones(1, dtype=torch.float32),
|
||||
cu_seqlens_kv=cu_seq_lens,
|
||||
token_to_batch=token_to_batch,
|
||||
seq_starts=seq_starts,
|
||||
dequant=dequant,
|
||||
kv_cache_layout="NHD",
|
||||
total_tokens=total_kv_tokens,
|
||||
)
|
||||
|
||||
# Call aiter flash attention with gathered KV
|
||||
aiter.flash_attn_varlen_func(
|
||||
q=maybe_quantized_query,
|
||||
k=gathered_key,
|
||||
v=gathered_value,
|
||||
cu_seqlens_q=cu_query_lens,
|
||||
cu_seqlens_k=cu_seq_lens,
|
||||
max_seqlen_q=max_query_len,
|
||||
max_seqlen_k=max_kv_len,
|
||||
min_seqlen_q=1,
|
||||
dropout_p=0.0,
|
||||
softmax_scale=scale,
|
||||
causal=True,
|
||||
window_size=window_size,
|
||||
alibi_slopes=None,
|
||||
return_lse=False,
|
||||
out=output,
|
||||
)
|
||||
|
||||
ref_output = ref_paged_attn(
|
||||
query=query,
|
||||
key_cache=key_cache,
|
||||
value_cache=value_cache,
|
||||
query_lens=query_lens,
|
||||
kv_lens=kv_lens_list,
|
||||
block_tables=block_tables,
|
||||
scale=scale,
|
||||
sliding_window=sliding_window,
|
||||
soft_cap=soft_cap,
|
||||
)
|
||||
|
||||
atol, rtol = 2e-2, 2e-2
|
||||
if q_dtype is not None:
|
||||
atol, rtol = 1.5e-1, 1.5e-1
|
||||
(
|
||||
torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol),
|
||||
f"{torch.max(torch.abs(output - ref_output))}",
|
||||
)
|
||||
|
||||
# Log diff stats for tracking changes
|
||||
print(f"Max abs diff: {torch.max(torch.abs(output - ref_output))}")
|
||||
print(f"Mean diff: {torch.mean(torch.abs(output - ref_output))}")
|
||||
print(f"Min diff: {torch.std(torch.abs(output - ref_output))}")
|
||||
@@ -0,0 +1,906 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""ROCm-specific tests for the AITER Flash Attention backend.
|
||||
|
||||
This file owns the ROCm backend contract and the ROCm-specific execution paths
|
||||
wired through ``vllm.v1.attention.backends.rocm_aiter_fa``:
|
||||
- backend contract and env gating
|
||||
- representative prefill, multi-batch, decode, and FP8-KV execution
|
||||
- direct kernel stress for large block tables and sliding-window masking
|
||||
"""
|
||||
|
||||
import importlib
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.platforms.rocm import on_mi3xx
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not current_platform.is_rocm(), reason="ROCm-specific tests"
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _cuda_default_device():
|
||||
torch.set_default_device("cuda")
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
torch.set_default_device("cpu")
|
||||
|
||||
|
||||
DTYPES = [torch.bfloat16, torch.float16]
|
||||
HEAD_SIZES = [64, 128, 256]
|
||||
NUM_HEADS_PAIRS = [(8, 8), (16, 4)] # (num_q_heads, num_kv_heads) - tests GQA
|
||||
BLOCK_SIZE = 16
|
||||
NUM_BLOCKS = 2048
|
||||
DIRECT_NUM_BLOCKS = [2048, 32768]
|
||||
DIRECT_NUM_HEADS = (8, 2)
|
||||
DIRECT_HEAD_SIZE = 128
|
||||
DIRECT_DTYPE = torch.bfloat16
|
||||
HEAD_SIZE_TEST_NUM_HEADS_PAIRS = [(16, 16), (16, 4)]
|
||||
# Prefill seq lens: (query_len, kv_len). Exclude single-token decode (q=1)
|
||||
# because flash_attn_varlen_func is a prefill kernel; q_len=1 with short kv
|
||||
# triggers kernel limitations (MAE > 0.1 for head != 128 in BF16, all heads in FP16).
|
||||
# Single-token decode is covered separately in test_aiter_mha_decode_single_token.
|
||||
SEQ_LENS = [(8, 512), (32, 1024)]
|
||||
DIRECT_SEQ_LENS = [(10, 1328), (5, 18), (129, 463)]
|
||||
|
||||
|
||||
# Reference implementation -------------------------------------------------
|
||||
def ref_paged_attn(
|
||||
query: torch.Tensor,
|
||||
key_cache: torch.Tensor,
|
||||
value_cache: torch.Tensor,
|
||||
query_lens: list[int],
|
||||
kv_lens: list[int],
|
||||
block_tables: torch.Tensor,
|
||||
scale: float,
|
||||
sliding_window: int | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Naive reference paged attention using einsum."""
|
||||
num_seqs = len(query_lens)
|
||||
block_tables_np = block_tables.cpu().numpy()
|
||||
_, block_size, num_kv_heads, head_size = key_cache.shape
|
||||
|
||||
outputs = []
|
||||
start_idx = 0
|
||||
for i in range(num_seqs):
|
||||
query_len = query_lens[i]
|
||||
kv_len = kv_lens[i]
|
||||
q = query[start_idx : start_idx + query_len] * scale
|
||||
|
||||
num_kv_blocks = (kv_len + block_size - 1) // block_size
|
||||
block_indices = block_tables_np[i, :num_kv_blocks]
|
||||
|
||||
k = key_cache[block_indices].view(-1, num_kv_heads, head_size)[:kv_len]
|
||||
v = value_cache[block_indices].view(-1, num_kv_heads, head_size)[:kv_len]
|
||||
|
||||
if q.shape[1] != k.shape[1]:
|
||||
k = torch.repeat_interleave(k, q.shape[1] // k.shape[1], dim=1)
|
||||
v = torch.repeat_interleave(v, q.shape[1] // v.shape[1], dim=1)
|
||||
|
||||
attn = torch.einsum("qhd,khd->hqk", q, k).float()
|
||||
mask = torch.triu(
|
||||
torch.ones(query_len, kv_len), diagonal=kv_len - query_len + 1
|
||||
).bool()
|
||||
if sliding_window is not None:
|
||||
window_mask = (
|
||||
torch.triu(
|
||||
torch.ones(query_len, kv_len),
|
||||
diagonal=kv_len - (query_len + sliding_window) + 1,
|
||||
)
|
||||
.bool()
|
||||
.logical_not()
|
||||
)
|
||||
mask |= window_mask
|
||||
attn.masked_fill_(mask, float("-inf"))
|
||||
attn = torch.softmax(attn, dim=-1).to(v.dtype)
|
||||
out = torch.einsum("hqk,khd->qhd", attn, v)
|
||||
outputs.append(out)
|
||||
start_idx += query_len
|
||||
|
||||
return torch.cat(outputs, dim=0)
|
||||
|
||||
|
||||
# Small test helpers ------------------------------------------------------
|
||||
def _reload_envs():
|
||||
import vllm.envs as envs
|
||||
|
||||
return importlib.reload(envs)
|
||||
|
||||
|
||||
def _assert_aiter_supported() -> None:
|
||||
from vllm._aiter_ops import is_aiter_found_and_supported
|
||||
|
||||
assert is_aiter_found_and_supported(), (
|
||||
"aiter is required on supported ROCm hardware for this test"
|
||||
)
|
||||
|
||||
|
||||
def _format_observed_rate(count: int, total: int) -> str:
|
||||
return f"{count / total:.4%} ({count}/{total})"
|
||||
|
||||
|
||||
def _format_allowed_rate(rate: float, total: int) -> str:
|
||||
allowed_count = int(rate * total)
|
||||
return f"{rate:.4%} (<= {allowed_count}/{total})"
|
||||
|
||||
|
||||
def _print_close_stats(
|
||||
label: str,
|
||||
actual: torch.Tensor,
|
||||
expected: torch.Tensor,
|
||||
*,
|
||||
atol: float,
|
||||
rtol: float,
|
||||
) -> None:
|
||||
abs_diff = (actual - expected).abs().float().flatten()
|
||||
expected_abs = expected.abs().float().flatten()
|
||||
allowed = atol + rtol * expected_abs
|
||||
within = abs_diff <= allowed
|
||||
|
||||
total = abs_diff.numel()
|
||||
passed = int(within.sum().item())
|
||||
failed = total - passed
|
||||
|
||||
max_abs = abs_diff.max().item()
|
||||
mean_abs = abs_diff.mean().item()
|
||||
p99_abs = torch.quantile(abs_diff, 0.99).item()
|
||||
p999_abs = torch.quantile(abs_diff, 0.999).item()
|
||||
worst_ratio = (abs_diff / allowed.clamp_min(1e-12)).max().item()
|
||||
|
||||
print(
|
||||
"[rocm_aiter_fa] "
|
||||
f"{label}: "
|
||||
f"pass={passed / total:.4%} ({passed}/{total}) "
|
||||
f"fail={_format_observed_rate(failed, total)} "
|
||||
f"allowed_fail={_format_allowed_rate(0.0, total)} "
|
||||
f"atol={atol:g} "
|
||||
f"rtol={rtol:g} "
|
||||
f"max_abs={max_abs:.6g} "
|
||||
f"mean_abs={mean_abs:.6g} "
|
||||
f"p99_abs={p99_abs:.6g} "
|
||||
f"p999_abs={p999_abs:.6g} "
|
||||
f"worst_ratio={worst_ratio:.6g}"
|
||||
)
|
||||
|
||||
|
||||
# Direct kernel helper ----------------------------------------------------
|
||||
def _run_direct_flash_attn_case(
|
||||
*,
|
||||
seq_lens: list[tuple[int, int]],
|
||||
num_blocks: int,
|
||||
sliding_window: int | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
import aiter
|
||||
|
||||
from vllm.v1.attention.backends.rocm_aiter_fa import cp_mha_gather_cache
|
||||
|
||||
set_random_seed(0)
|
||||
|
||||
query_lens = [query_len for query_len, _ in seq_lens]
|
||||
kv_lens = [kv_len for _, kv_len in seq_lens]
|
||||
num_query_heads, num_kv_heads = DIRECT_NUM_HEADS
|
||||
total_query_tokens = sum(query_lens)
|
||||
total_kv_tokens = sum(kv_lens)
|
||||
max_query_len = max(query_lens)
|
||||
max_kv_len = max(kv_lens)
|
||||
scale = DIRECT_HEAD_SIZE**-0.5
|
||||
|
||||
query = torch.randn(
|
||||
total_query_tokens,
|
||||
num_query_heads,
|
||||
DIRECT_HEAD_SIZE,
|
||||
dtype=DIRECT_DTYPE,
|
||||
)
|
||||
key_cache = torch.randn(
|
||||
num_blocks,
|
||||
BLOCK_SIZE,
|
||||
num_kv_heads,
|
||||
DIRECT_HEAD_SIZE,
|
||||
dtype=DIRECT_DTYPE,
|
||||
)
|
||||
value_cache = torch.randn_like(key_cache)
|
||||
|
||||
cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
|
||||
dim=0, dtype=torch.int32
|
||||
)
|
||||
cu_seq_lens = torch.tensor([0] + kv_lens, dtype=torch.int32).cumsum(
|
||||
dim=0, dtype=torch.int32
|
||||
)
|
||||
max_num_blocks_per_seq = (max_kv_len + BLOCK_SIZE - 1) // BLOCK_SIZE
|
||||
block_tables = torch.randint(
|
||||
0, num_blocks, (len(seq_lens), max_num_blocks_per_seq), dtype=torch.int32
|
||||
)
|
||||
|
||||
token_to_batch = torch.empty(total_kv_tokens, dtype=torch.int32)
|
||||
seq_starts = torch.zeros(len(seq_lens), dtype=torch.int32)
|
||||
token_index = 0
|
||||
for batch_index, kv_len in enumerate(kv_lens):
|
||||
token_to_batch[token_index : token_index + kv_len] = batch_index
|
||||
token_index += kv_len
|
||||
|
||||
gathered_key = torch.empty(
|
||||
total_kv_tokens,
|
||||
num_kv_heads,
|
||||
DIRECT_HEAD_SIZE,
|
||||
dtype=DIRECT_DTYPE,
|
||||
)
|
||||
gathered_value = torch.empty_like(gathered_key)
|
||||
cp_mha_gather_cache(
|
||||
key_cache=key_cache,
|
||||
value_cache=value_cache,
|
||||
key=gathered_key,
|
||||
value=gathered_value,
|
||||
block_tables=block_tables,
|
||||
k_scales=torch.ones(1, dtype=torch.float32),
|
||||
v_scales=torch.ones(1, dtype=torch.float32),
|
||||
cu_seqlens_kv=cu_seq_lens,
|
||||
token_to_batch=token_to_batch,
|
||||
seq_starts=seq_starts,
|
||||
dequant=False,
|
||||
kv_cache_layout="NHD",
|
||||
total_tokens=total_kv_tokens,
|
||||
)
|
||||
|
||||
window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
|
||||
output = torch.empty_like(query)
|
||||
aiter.flash_attn_varlen_func(
|
||||
q=query,
|
||||
k=gathered_key,
|
||||
v=gathered_value,
|
||||
cu_seqlens_q=cu_query_lens,
|
||||
cu_seqlens_k=cu_seq_lens,
|
||||
max_seqlen_q=max_query_len,
|
||||
max_seqlen_k=max_kv_len,
|
||||
min_seqlen_q=1,
|
||||
dropout_p=0.0,
|
||||
softmax_scale=scale,
|
||||
causal=True,
|
||||
window_size=window_size,
|
||||
alibi_slopes=None,
|
||||
return_lse=False,
|
||||
out=output,
|
||||
)
|
||||
|
||||
ref = ref_paged_attn(
|
||||
query=query,
|
||||
key_cache=key_cache,
|
||||
value_cache=value_cache,
|
||||
query_lens=query_lens,
|
||||
kv_lens=kv_lens,
|
||||
block_tables=block_tables,
|
||||
scale=scale,
|
||||
sliding_window=sliding_window,
|
||||
)
|
||||
return output, ref
|
||||
|
||||
|
||||
# Backend contract tests --------------------------------------------------
|
||||
def test_aiter_mha_backend_contract():
|
||||
"""The ROCm backend advertises the dtypes, shapes, and attention types it
|
||||
is designed to handle."""
|
||||
from vllm.v1.attention.backend import AttentionType
|
||||
from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend
|
||||
|
||||
assert AiterFlashAttentionBackend.get_name() == "FLASH_ATTN"
|
||||
assert AiterFlashAttentionBackend.supported_dtypes == [
|
||||
torch.float16,
|
||||
torch.bfloat16,
|
||||
]
|
||||
assert AiterFlashAttentionBackend.supported_kv_cache_dtypes == [
|
||||
"auto",
|
||||
"float16",
|
||||
"bfloat16",
|
||||
"fp8",
|
||||
"fp8_e4m3",
|
||||
"fp8_e5m2",
|
||||
]
|
||||
assert AiterFlashAttentionBackend.get_supported_kernel_block_sizes() == [16, 32]
|
||||
assert AiterFlashAttentionBackend.get_supported_head_sizes() == [64, 128, 256]
|
||||
assert AiterFlashAttentionBackend.supports_attn_type(AttentionType.DECODER)
|
||||
assert not AiterFlashAttentionBackend.supports_attn_type(AttentionType.ENCODER)
|
||||
assert not AiterFlashAttentionBackend.supports_attn_type(
|
||||
AttentionType.ENCODER_DECODER
|
||||
)
|
||||
|
||||
|
||||
def test_aiter_mha_backend_validates_kv_cache_block_size():
|
||||
"""The backend should reject KV cache shapes that cannot be gathered
|
||||
correctly by the ROCm kernel."""
|
||||
from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend
|
||||
|
||||
assert AiterFlashAttentionBackend.get_kv_cache_shape(8, 16, 8, 128) == (
|
||||
8,
|
||||
8,
|
||||
16,
|
||||
256,
|
||||
)
|
||||
with pytest.raises(ValueError, match="Block size must be a multiple of 16"):
|
||||
AiterFlashAttentionBackend.get_kv_cache_shape(8, 15, 8, 128)
|
||||
|
||||
|
||||
def test_aiter_mha_backend_supports_compute_capability_matches_mi3xx_probe():
|
||||
"""The backend should trust the ROCm MI3xx probe instead of the raw torch
|
||||
capability tuple."""
|
||||
from vllm.platforms.interface import DeviceCapability
|
||||
from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend
|
||||
|
||||
assert (
|
||||
AiterFlashAttentionBackend.supports_compute_capability(DeviceCapability(0, 0))
|
||||
is on_mi3xx()
|
||||
)
|
||||
|
||||
|
||||
# Env and platform gate tests ---------------------------------------------
|
||||
@pytest.mark.skipif(not on_mi3xx(), reason="MI300/MI350 ROCm only")
|
||||
@pytest.mark.parametrize(
|
||||
("use_aiter", "use_mha", "expected"),
|
||||
[
|
||||
(True, True, True),
|
||||
(True, False, False),
|
||||
(False, True, False),
|
||||
(False, False, False),
|
||||
],
|
||||
)
|
||||
def test_aiter_mha_env_flags_control_enablement(
|
||||
use_aiter, use_mha, expected, monkeypatch
|
||||
):
|
||||
"""Both the global AITER flag and the MHA-specific flag must be enabled
|
||||
before the ROCm MHA path is considered active."""
|
||||
from vllm._aiter_ops import rocm_aiter_ops
|
||||
|
||||
_assert_aiter_supported()
|
||||
|
||||
with monkeypatch.context() as mp:
|
||||
mp.setenv("VLLM_ROCM_USE_AITER", "1" if use_aiter else "0")
|
||||
mp.setenv("VLLM_ROCM_USE_AITER_MHA", "1" if use_mha else "0")
|
||||
_reload_envs()
|
||||
rocm_aiter_ops.refresh_env_variables()
|
||||
assert rocm_aiter_ops.is_mha_enabled() is expected
|
||||
|
||||
_reload_envs()
|
||||
rocm_aiter_ops.refresh_env_variables()
|
||||
|
||||
|
||||
def test_aiter_mha_platform_gate_matches_install_and_arch():
|
||||
"""The global AITER availability check should only open on ROCm MI3xx when
|
||||
the aiter package is installed."""
|
||||
from vllm._aiter_ops import IS_AITER_FOUND, is_aiter_found_and_supported
|
||||
|
||||
assert is_aiter_found_and_supported() is (
|
||||
current_platform.is_rocm() and on_mi3xx() and IS_AITER_FOUND
|
||||
)
|
||||
|
||||
|
||||
# Kernel path tests -------------------------------------------------------
|
||||
def _run_single_seq_flash_attn_case(
|
||||
*,
|
||||
head_size: int,
|
||||
num_heads: tuple[int, int],
|
||||
seq_lens: tuple[int, int],
|
||||
dtype: torch.dtype,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
import aiter
|
||||
|
||||
from vllm.v1.attention.backends.rocm_aiter_fa import cp_mha_gather_cache
|
||||
|
||||
set_random_seed(0)
|
||||
|
||||
num_q_heads, num_kv_heads = num_heads
|
||||
query_len, kv_len = seq_lens
|
||||
scale = head_size**-0.5
|
||||
|
||||
query = torch.randn(query_len, num_q_heads, head_size, dtype=dtype)
|
||||
key_cache = torch.randn(
|
||||
NUM_BLOCKS, BLOCK_SIZE, num_kv_heads, head_size, dtype=dtype
|
||||
)
|
||||
value_cache = torch.randn_like(key_cache)
|
||||
|
||||
cu_query_lens = torch.tensor([0, query_len], dtype=torch.int32).cumsum(
|
||||
dim=0, dtype=torch.int32
|
||||
)
|
||||
cu_seq_lens = torch.tensor([0, kv_len], dtype=torch.int32).cumsum(
|
||||
dim=0, dtype=torch.int32
|
||||
)
|
||||
max_num_blocks = (kv_len + BLOCK_SIZE - 1) // BLOCK_SIZE
|
||||
block_tables = torch.randint(0, NUM_BLOCKS, (1, max_num_blocks), dtype=torch.int32)
|
||||
|
||||
token_to_batch = torch.zeros(kv_len, dtype=torch.int32)
|
||||
seq_starts = torch.zeros(1, dtype=torch.int32)
|
||||
gathered_key = torch.empty(kv_len, num_kv_heads, head_size, dtype=dtype)
|
||||
gathered_value = torch.empty_like(gathered_key)
|
||||
|
||||
cp_mha_gather_cache(
|
||||
key_cache=key_cache,
|
||||
value_cache=value_cache,
|
||||
key=gathered_key,
|
||||
value=gathered_value,
|
||||
block_tables=block_tables,
|
||||
k_scales=torch.ones(1, dtype=torch.float32),
|
||||
v_scales=torch.ones(1, dtype=torch.float32),
|
||||
cu_seqlens_kv=cu_seq_lens,
|
||||
token_to_batch=token_to_batch,
|
||||
seq_starts=seq_starts,
|
||||
dequant=False,
|
||||
kv_cache_layout="NHD",
|
||||
total_tokens=kv_len,
|
||||
)
|
||||
|
||||
output = torch.empty_like(query)
|
||||
aiter.flash_attn_varlen_func(
|
||||
q=query,
|
||||
k=gathered_key,
|
||||
v=gathered_value,
|
||||
cu_seqlens_q=cu_query_lens,
|
||||
cu_seqlens_k=cu_seq_lens,
|
||||
max_seqlen_q=query_len,
|
||||
max_seqlen_k=kv_len,
|
||||
min_seqlen_q=1,
|
||||
dropout_p=0.0,
|
||||
softmax_scale=scale,
|
||||
causal=True,
|
||||
window_size=(-1, -1),
|
||||
alibi_slopes=None,
|
||||
return_lse=False,
|
||||
out=output,
|
||||
)
|
||||
|
||||
ref = ref_paged_attn(
|
||||
query=query,
|
||||
key_cache=key_cache,
|
||||
value_cache=value_cache,
|
||||
query_lens=[query_len],
|
||||
kv_lens=[kv_len],
|
||||
block_tables=block_tables,
|
||||
scale=scale,
|
||||
)
|
||||
return output, ref
|
||||
|
||||
|
||||
@pytest.mark.skipif(not on_mi3xx(), reason="MI300/MI350 ROCm only")
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("num_heads", HEAD_SIZE_TEST_NUM_HEADS_PAIRS)
|
||||
@pytest.mark.parametrize("seq_lens", SEQ_LENS)
|
||||
def test_aiter_fa_head_sizes(head_size, dtype, num_heads, seq_lens):
|
||||
"""AITER flash attention should stay accurate across supported head sizes."""
|
||||
atol = 1.5e-2
|
||||
rtol = 1e-2
|
||||
_assert_aiter_supported()
|
||||
|
||||
output, ref = _run_single_seq_flash_attn_case(
|
||||
head_size=head_size,
|
||||
num_heads=num_heads,
|
||||
seq_lens=seq_lens,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
_print_close_stats(
|
||||
"head_sizes "
|
||||
f"dtype={dtype} head_size={head_size} "
|
||||
f"num_heads={num_heads} seq_lens={seq_lens}",
|
||||
output,
|
||||
ref,
|
||||
atol=atol,
|
||||
rtol=rtol,
|
||||
)
|
||||
torch.testing.assert_close(output, ref, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not on_mi3xx(), reason="MI300/MI350 ROCm only")
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS_PAIRS)
|
||||
@pytest.mark.parametrize("seq_lens", SEQ_LENS)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
def test_aiter_mha_varlen_paged_kv(head_size, num_heads, seq_lens, dtype):
|
||||
"""Test AITER flash attention varlen with paged KV cache.
|
||||
|
||||
Exercises: VLLM_ROCM_USE_AITER, VLLM_ROCM_USE_AITER_MHA
|
||||
"""
|
||||
atol = 1.5e-2
|
||||
rtol = 1e-2
|
||||
_assert_aiter_supported()
|
||||
|
||||
output, ref = _run_single_seq_flash_attn_case(
|
||||
head_size=head_size,
|
||||
num_heads=num_heads,
|
||||
seq_lens=seq_lens,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
_print_close_stats(
|
||||
"varlen_paged_kv "
|
||||
f"dtype={dtype} head_size={head_size} "
|
||||
f"num_heads={num_heads} seq_lens={seq_lens}",
|
||||
output,
|
||||
ref,
|
||||
atol=atol,
|
||||
rtol=rtol,
|
||||
)
|
||||
torch.testing.assert_close(output, ref, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not on_mi3xx(), reason="MI300/MI350 ROCm only")
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS_PAIRS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
def test_aiter_mha_multi_batch(num_heads, head_size, dtype):
|
||||
"""Test AITER flash attention with multiple sequences in a batch."""
|
||||
atol = 1.5e-2
|
||||
rtol = 1e-2
|
||||
_assert_aiter_supported()
|
||||
import aiter
|
||||
|
||||
from vllm.v1.attention.backends.rocm_aiter_fa import cp_mha_gather_cache
|
||||
|
||||
set_random_seed(42)
|
||||
|
||||
num_q_heads, num_kv_heads = num_heads
|
||||
seq_lens = [(4, 128), (2, 256), (8, 64)]
|
||||
query_lens = [q for q, _ in seq_lens]
|
||||
kv_lens = [k for _, k in seq_lens]
|
||||
num_seqs = len(seq_lens)
|
||||
scale = head_size**-0.5
|
||||
|
||||
total_q = sum(query_lens)
|
||||
total_kv = sum(kv_lens)
|
||||
|
||||
query = torch.randn(total_q, num_q_heads, head_size, dtype=dtype)
|
||||
key_cache = torch.randn(
|
||||
NUM_BLOCKS, BLOCK_SIZE, num_kv_heads, head_size, dtype=dtype
|
||||
)
|
||||
value_cache = torch.randn_like(key_cache)
|
||||
|
||||
cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
|
||||
dim=0, dtype=torch.int32
|
||||
)
|
||||
cu_seq_lens = torch.tensor([0] + kv_lens, dtype=torch.int32).cumsum(
|
||||
dim=0, dtype=torch.int32
|
||||
)
|
||||
|
||||
max_kv_len = max(kv_lens)
|
||||
max_num_blocks = (max_kv_len + BLOCK_SIZE - 1) // BLOCK_SIZE
|
||||
block_tables = torch.randint(
|
||||
0, NUM_BLOCKS, (num_seqs, max_num_blocks), dtype=torch.int32
|
||||
)
|
||||
|
||||
token_to_batch = torch.empty(total_kv, dtype=torch.int32)
|
||||
seq_starts = torch.zeros(num_seqs, dtype=torch.int32)
|
||||
tok_idx = 0
|
||||
for b, kl in enumerate(kv_lens):
|
||||
token_to_batch[tok_idx : tok_idx + kl] = b
|
||||
tok_idx += kl
|
||||
|
||||
gathered_key = torch.empty(total_kv, num_kv_heads, head_size, dtype=dtype)
|
||||
gathered_value = torch.empty_like(gathered_key)
|
||||
|
||||
cp_mha_gather_cache(
|
||||
key_cache=key_cache,
|
||||
value_cache=value_cache,
|
||||
key=gathered_key,
|
||||
value=gathered_value,
|
||||
block_tables=block_tables,
|
||||
k_scales=torch.ones(1, dtype=torch.float32),
|
||||
v_scales=torch.ones(1, dtype=torch.float32),
|
||||
cu_seqlens_kv=cu_seq_lens,
|
||||
token_to_batch=token_to_batch,
|
||||
seq_starts=seq_starts,
|
||||
dequant=False,
|
||||
kv_cache_layout="NHD",
|
||||
total_tokens=total_kv,
|
||||
)
|
||||
|
||||
output = torch.empty_like(query)
|
||||
aiter.flash_attn_varlen_func(
|
||||
q=query,
|
||||
k=gathered_key,
|
||||
v=gathered_value,
|
||||
cu_seqlens_q=cu_query_lens,
|
||||
cu_seqlens_k=cu_seq_lens,
|
||||
max_seqlen_q=max(query_lens),
|
||||
max_seqlen_k=max_kv_len,
|
||||
min_seqlen_q=1,
|
||||
dropout_p=0.0,
|
||||
softmax_scale=scale,
|
||||
causal=True,
|
||||
window_size=(-1, -1),
|
||||
alibi_slopes=None,
|
||||
return_lse=False,
|
||||
out=output,
|
||||
)
|
||||
|
||||
ref = ref_paged_attn(
|
||||
query=query,
|
||||
key_cache=key_cache,
|
||||
value_cache=value_cache,
|
||||
query_lens=query_lens,
|
||||
kv_lens=kv_lens,
|
||||
block_tables=block_tables,
|
||||
scale=scale,
|
||||
)
|
||||
|
||||
_print_close_stats(
|
||||
"multi_batch "
|
||||
f"dtype={dtype} head_size={head_size} "
|
||||
f"num_heads={num_heads} seq_lens={seq_lens}",
|
||||
output,
|
||||
ref,
|
||||
atol=atol,
|
||||
rtol=rtol,
|
||||
)
|
||||
torch.testing.assert_close(output, ref, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not on_mi3xx(), reason="MI300/MI350 ROCm only")
|
||||
@pytest.mark.parametrize("num_blocks", DIRECT_NUM_BLOCKS)
|
||||
def test_aiter_fa_large_block_table_matches_reference(num_blocks):
|
||||
"""The direct paged-KV path should stay stable for both normal and very
|
||||
large block tables."""
|
||||
atol = 2e-2
|
||||
rtol = 2e-2
|
||||
_assert_aiter_supported()
|
||||
|
||||
output, ref = _run_direct_flash_attn_case(
|
||||
seq_lens=DIRECT_SEQ_LENS,
|
||||
num_blocks=num_blocks,
|
||||
)
|
||||
_print_close_stats(
|
||||
f"direct_varlen_paged_kv num_blocks={num_blocks}",
|
||||
output,
|
||||
ref,
|
||||
atol=atol,
|
||||
rtol=rtol,
|
||||
)
|
||||
torch.testing.assert_close(output, ref, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not on_mi3xx(), reason="MI300/MI350 ROCm only")
|
||||
def test_aiter_fa_sliding_window_matches_reference():
|
||||
"""The direct kernel should respect the same sliding-window causal mask as
|
||||
the naive reference implementation."""
|
||||
atol = 2e-2
|
||||
rtol = 2e-2
|
||||
_assert_aiter_supported()
|
||||
|
||||
output, ref = _run_direct_flash_attn_case(
|
||||
seq_lens=[(8, 523), (24, 37), (3, 2011)],
|
||||
num_blocks=2048,
|
||||
sliding_window=256,
|
||||
)
|
||||
_print_close_stats(
|
||||
"sliding_window num_blocks=2048 window=256",
|
||||
output,
|
||||
ref,
|
||||
atol=atol,
|
||||
rtol=rtol,
|
||||
)
|
||||
torch.testing.assert_close(output, ref, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
# Decode path test --------------------------------------------------------
|
||||
@pytest.mark.skipif(not on_mi3xx(), reason="MI300/MI350 ROCm only")
|
||||
@pytest.mark.parametrize(
|
||||
"dtype",
|
||||
[
|
||||
torch.bfloat16,
|
||||
pytest.param(
|
||||
torch.float16,
|
||||
marks=pytest.mark.xfail(
|
||||
reason=(
|
||||
"aiter bug #2229: flash_attn_varlen_func currently "
|
||||
"miscomputes FP16 single-token decode on MI3xx "
|
||||
"(validated on gfx950; max abs diff 1.95 vs atol 1.5e-2). "
|
||||
"Remove xfail when the AITER decode kernel is fixed. "
|
||||
"https://github.com/ROCm/aiter/issues/2229"
|
||||
),
|
||||
),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_aiter_mha_decode_single_token(dtype):
|
||||
"""Test AITER MHA for decode (single query token per sequence).
|
||||
|
||||
BF16 is the working reference configuration here. FP16 remains xfail until
|
||||
the upstream AITER single-token decode bug is fixed.
|
||||
"""
|
||||
atol = 1.5e-2
|
||||
rtol = 1e-2
|
||||
_assert_aiter_supported()
|
||||
import aiter
|
||||
|
||||
from vllm.v1.attention.backends.rocm_aiter_fa import cp_mha_gather_cache
|
||||
|
||||
set_random_seed(0)
|
||||
|
||||
num_q_heads, num_kv_heads = 8, 8
|
||||
head_size = 128
|
||||
kv_len = 512
|
||||
scale = head_size**-0.5
|
||||
|
||||
query = torch.randn(1, num_q_heads, head_size, dtype=dtype)
|
||||
key_cache = torch.randn(
|
||||
NUM_BLOCKS, BLOCK_SIZE, num_kv_heads, head_size, dtype=dtype
|
||||
)
|
||||
value_cache = torch.randn_like(key_cache)
|
||||
|
||||
cu_query_lens = torch.tensor([0, 1], dtype=torch.int32)
|
||||
cu_seq_lens = torch.tensor([0, kv_len], dtype=torch.int32)
|
||||
|
||||
max_num_blocks = (kv_len + BLOCK_SIZE - 1) // BLOCK_SIZE
|
||||
block_tables = torch.randint(0, NUM_BLOCKS, (1, max_num_blocks), dtype=torch.int32)
|
||||
|
||||
token_to_batch = torch.zeros(kv_len, dtype=torch.int32)
|
||||
seq_starts = torch.zeros(1, dtype=torch.int32)
|
||||
gathered_key = torch.empty(kv_len, num_kv_heads, head_size, dtype=dtype)
|
||||
gathered_value = torch.empty_like(gathered_key)
|
||||
|
||||
cp_mha_gather_cache(
|
||||
key_cache=key_cache,
|
||||
value_cache=value_cache,
|
||||
key=gathered_key,
|
||||
value=gathered_value,
|
||||
block_tables=block_tables,
|
||||
k_scales=torch.ones(1, dtype=torch.float32),
|
||||
v_scales=torch.ones(1, dtype=torch.float32),
|
||||
cu_seqlens_kv=cu_seq_lens,
|
||||
token_to_batch=token_to_batch,
|
||||
seq_starts=seq_starts,
|
||||
dequant=False,
|
||||
kv_cache_layout="NHD",
|
||||
total_tokens=kv_len,
|
||||
)
|
||||
|
||||
output = torch.empty_like(query)
|
||||
aiter.flash_attn_varlen_func(
|
||||
q=query,
|
||||
k=gathered_key,
|
||||
v=gathered_value,
|
||||
cu_seqlens_q=cu_query_lens,
|
||||
cu_seqlens_k=cu_seq_lens,
|
||||
max_seqlen_q=1,
|
||||
max_seqlen_k=kv_len,
|
||||
min_seqlen_q=1,
|
||||
dropout_p=0.0,
|
||||
softmax_scale=scale,
|
||||
causal=True,
|
||||
window_size=(-1, -1),
|
||||
alibi_slopes=None,
|
||||
return_lse=False,
|
||||
out=output,
|
||||
)
|
||||
|
||||
ref = ref_paged_attn(
|
||||
query=query,
|
||||
key_cache=key_cache,
|
||||
value_cache=value_cache,
|
||||
query_lens=[1],
|
||||
kv_lens=[kv_len],
|
||||
block_tables=block_tables,
|
||||
scale=scale,
|
||||
)
|
||||
|
||||
_print_close_stats(
|
||||
f"decode_single_token dtype={dtype} kv_len={kv_len}",
|
||||
output,
|
||||
ref,
|
||||
atol=atol,
|
||||
rtol=rtol,
|
||||
)
|
||||
torch.testing.assert_close(output, ref, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
# FP8 KV cache test -------------------------------------------------------
|
||||
@pytest.mark.skipif(not on_mi3xx(), reason="MI300/MI350 ROCm only")
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
def test_aiter_mha_varlen_fp8_kv(dtype):
|
||||
"""AITER flash attention with FP8 KV cache matches reference on BF16-cast KV.
|
||||
|
||||
cp_mha_gather_cache is called with dequant=True to dequantize FP8 to dtype
|
||||
before passing to flash_attn_varlen_func. We compare to ref_paged_attn on
|
||||
the dtype-cast KV cache with FP8 quantization tolerance.
|
||||
|
||||
Exercises: VLLM_ROCM_USE_AITER, VLLM_ROCM_USE_AITER_MHA, FP8 KV path.
|
||||
"""
|
||||
atol = 0.15
|
||||
rtol = 0.05
|
||||
_assert_aiter_supported()
|
||||
if not current_platform.supports_fp8():
|
||||
pytest.skip("FP8 not supported on this hardware")
|
||||
|
||||
import aiter
|
||||
|
||||
from vllm.v1.attention.backends.rocm_aiter_fa import cp_mha_gather_cache
|
||||
|
||||
FP8_DTYPE = current_platform.fp8_dtype()
|
||||
|
||||
set_random_seed(10)
|
||||
|
||||
num_q_heads, num_kv_heads = 8, 8
|
||||
head_size = 128
|
||||
query_len, kv_len = 4, 128
|
||||
scale = head_size**-0.5
|
||||
|
||||
query = torch.randn(query_len, num_q_heads, head_size, dtype=dtype)
|
||||
# FP8 KV cache; clamp to stay in FP8 representable range
|
||||
key_cache_fp8 = torch.clamp(
|
||||
torch.randn(NUM_BLOCKS, BLOCK_SIZE, num_kv_heads, head_size), -1.0, 1.0
|
||||
).to(FP8_DTYPE)
|
||||
value_cache_fp8 = torch.clamp(
|
||||
torch.randn(NUM_BLOCKS, BLOCK_SIZE, num_kv_heads, head_size), -1.0, 1.0
|
||||
).to(FP8_DTYPE)
|
||||
|
||||
cu_query_lens = torch.tensor([0, query_len], dtype=torch.int32)
|
||||
cu_seq_lens = torch.tensor([0, kv_len], dtype=torch.int32)
|
||||
max_num_blocks = (kv_len + BLOCK_SIZE - 1) // BLOCK_SIZE
|
||||
block_tables = torch.randint(0, NUM_BLOCKS, (1, max_num_blocks), dtype=torch.int32)
|
||||
token_to_batch = torch.zeros(kv_len, dtype=torch.int32)
|
||||
seq_starts = torch.zeros(1, dtype=torch.int32)
|
||||
|
||||
# Gather and dequantize FP8 KV to dtype
|
||||
gathered_key = torch.empty(kv_len, num_kv_heads, head_size, dtype=dtype)
|
||||
gathered_value = torch.empty_like(gathered_key)
|
||||
k_scales = torch.ones(1, dtype=torch.float32)
|
||||
v_scales = torch.ones(1, dtype=torch.float32)
|
||||
|
||||
cp_mha_gather_cache(
|
||||
key_cache=key_cache_fp8,
|
||||
value_cache=value_cache_fp8,
|
||||
key=gathered_key,
|
||||
value=gathered_value,
|
||||
block_tables=block_tables,
|
||||
k_scales=k_scales,
|
||||
v_scales=v_scales,
|
||||
cu_seqlens_kv=cu_seq_lens,
|
||||
token_to_batch=token_to_batch,
|
||||
seq_starts=seq_starts,
|
||||
dequant=True,
|
||||
kv_cache_layout="NHD",
|
||||
total_tokens=kv_len,
|
||||
)
|
||||
|
||||
output = torch.empty_like(query)
|
||||
aiter.flash_attn_varlen_func(
|
||||
q=query,
|
||||
k=gathered_key,
|
||||
v=gathered_value,
|
||||
cu_seqlens_q=cu_query_lens,
|
||||
cu_seqlens_k=cu_seq_lens,
|
||||
max_seqlen_q=query_len,
|
||||
max_seqlen_k=kv_len,
|
||||
min_seqlen_q=1,
|
||||
dropout_p=0.0,
|
||||
softmax_scale=scale,
|
||||
causal=True,
|
||||
window_size=(-1, -1),
|
||||
alibi_slopes=None,
|
||||
return_lse=False,
|
||||
out=output,
|
||||
)
|
||||
|
||||
# Reference: ref_paged_attn on dtype-cast KV (simulates perfect dequant)
|
||||
key_cache_ref = key_cache_fp8.to(dtype)
|
||||
value_cache_ref = value_cache_fp8.to(dtype)
|
||||
ref = ref_paged_attn(
|
||||
query=query,
|
||||
key_cache=key_cache_ref,
|
||||
value_cache=value_cache_ref,
|
||||
query_lens=[query_len],
|
||||
kv_lens=[kv_len],
|
||||
block_tables=block_tables,
|
||||
scale=scale,
|
||||
)
|
||||
|
||||
# FP8 quantization + dequantization introduces noise
|
||||
_print_close_stats(
|
||||
f"varlen_fp8_kv dtype={dtype} query_len={query_len} kv_len={kv_len}",
|
||||
output,
|
||||
ref,
|
||||
atol=atol,
|
||||
rtol=rtol,
|
||||
)
|
||||
torch.testing.assert_close(output, ref, atol=atol, rtol=rtol)
|
||||
Reference in New Issue
Block a user