112 lines
3.6 KiB
Python
112 lines
3.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Tests for the UltraQuant Triton unified-attention fallback."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm.utils.torch_utils import set_random_seed
|
|
from vllm.v1.attention.ops.ultraquant.format import slot_size
|
|
from vllm.v1.attention.ops.ultraquant.reference import reference_ultraquant_attention
|
|
from vllm.v1.attention.ops.ultraquant.triton_store import ultraquant_store
|
|
from vllm.v1.attention.ops.ultraquant.triton_unified_attention import (
|
|
ultraquant_unified_attention,
|
|
)
|
|
|
|
pytestmark = pytest.mark.skipif(
|
|
not torch.cuda.is_available(), reason="requires CUDA/HIP device"
|
|
)
|
|
|
|
|
|
def _build_paged_cache(key, value, block_size):
|
|
batch_size, seq_len, num_kv_heads, head_size = key.shape
|
|
blocks_per_seq = (seq_len + block_size - 1) // block_size
|
|
kv_cache = torch.zeros(
|
|
batch_size * blocks_per_seq,
|
|
block_size,
|
|
num_kv_heads,
|
|
slot_size(head_size),
|
|
dtype=torch.uint8,
|
|
device=key.device,
|
|
)
|
|
block_table = torch.arange(
|
|
batch_size * blocks_per_seq,
|
|
dtype=torch.int32,
|
|
device=key.device,
|
|
).view(batch_size, blocks_per_seq)
|
|
slot_mapping = torch.arange(
|
|
batch_size * blocks_per_seq * block_size,
|
|
dtype=torch.int64,
|
|
device=key.device,
|
|
).view(batch_size, -1)[:, :seq_len]
|
|
ultraquant_store(
|
|
key.reshape(-1, num_kv_heads, head_size),
|
|
value.reshape(-1, num_kv_heads, head_size),
|
|
kv_cache,
|
|
slot_mapping.reshape(-1),
|
|
)
|
|
return kv_cache, block_table
|
|
|
|
|
|
def _attention_quality(actual: torch.Tensor, expected: torch.Tensor) -> dict:
|
|
a = actual.float().reshape(-1)
|
|
b = expected.float().reshape(-1)
|
|
diff = (a - b).abs()
|
|
return {
|
|
"cos_sim": torch.nn.functional.cosine_similarity(a, b, dim=0).item(),
|
|
"mean_abs": diff.mean().item(),
|
|
"max_abs": diff.max().item(),
|
|
"norm_ratio": (a.norm() / b.norm()).item(),
|
|
}
|
|
|
|
|
|
def _unified_decode(query, kv_cache, block_table, seq_lens, scale, sinks=None):
|
|
batch_size = query.shape[0]
|
|
return ultraquant_unified_attention(
|
|
query=query,
|
|
kv_cache=kv_cache,
|
|
block_table=block_table,
|
|
seq_lens=seq_lens,
|
|
query_start_loc=torch.arange(
|
|
batch_size + 1, dtype=torch.int32, device=query.device
|
|
),
|
|
scale=scale,
|
|
max_query_len=1,
|
|
max_seq_len=int(seq_lens.max().item()),
|
|
sinks=sinks,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("query_group_size", [6, 8])
|
|
def test_unified_decode_d256_matches_reference(query_group_size):
|
|
"""The D=256 fallback covers both target Qwen GQA layouts."""
|
|
device = torch.device("cuda")
|
|
set_random_seed(0xD256 + query_group_size)
|
|
|
|
B, Hk, N, D = 2, 1, 256, 256
|
|
Hq = query_group_size * Hk
|
|
block_size = 32
|
|
dtype = torch.bfloat16
|
|
scale = D**-0.5
|
|
|
|
query = torch.randn(B, Hq, D, dtype=dtype, device=device)
|
|
key = torch.randn(B, N, Hk, D, dtype=dtype, device=device)
|
|
value = torch.randn_like(key)
|
|
kv_cache, block_table = _build_paged_cache(key, value, block_size)
|
|
seq_lens = torch.full((B,), N, dtype=torch.int32, device=device)
|
|
|
|
actual = _unified_decode(query, kv_cache, block_table, seq_lens, scale)
|
|
expected = reference_ultraquant_attention(
|
|
query=query,
|
|
key=key,
|
|
value=value,
|
|
scale=scale,
|
|
)
|
|
|
|
quality = _attention_quality(actual, expected)
|
|
assert quality["cos_sim"] > 0.98, quality
|
|
assert quality["mean_abs"] < 0.08, quality
|
|
assert quality["max_abs"] < 0.8, quality
|