89 lines
3.1 KiB
Python
89 lines
3.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Tests for the UltraQuant full-KV dequant kernel (continuation prefill).
|
|
|
|
The kernel reconstructs cached K/V from FP4 codes plus UE8M0 group scales.
|
|
Because dequant is an exact table lookup times a power of two, the result
|
|
must match the reference to within the output dtype's rounding.
|
|
"""
|
|
|
|
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 (
|
|
ultraquant_dequant,
|
|
ultraquant_encode,
|
|
)
|
|
from vllm.v1.attention.ops.ultraquant.triton_dequant import (
|
|
ultraquant_full_dequant_kv,
|
|
)
|
|
from vllm.v1.attention.ops.ultraquant.triton_store import ultraquant_store
|
|
|
|
pytestmark = pytest.mark.skipif(
|
|
not torch.cuda.is_available(), reason="requires CUDA/HIP device"
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
|
@pytest.mark.parametrize("head_dim", [128, 256])
|
|
@pytest.mark.parametrize("num_kv_heads", [1, 4])
|
|
@pytest.mark.parametrize("cached_len", [32, 96])
|
|
def test_full_dequant_matches_reference(dtype, head_dim, num_kv_heads, cached_len):
|
|
device = torch.device("cuda")
|
|
set_random_seed(0xD4)
|
|
|
|
block_size = 32
|
|
alloc_len = cached_len
|
|
num_blocks = alloc_len // block_size + 1
|
|
|
|
key = torch.randn(alloc_len, num_kv_heads, head_dim, dtype=dtype, device=device)
|
|
value = torch.randn(alloc_len, num_kv_heads, head_dim, dtype=dtype, device=device)
|
|
|
|
kv_cache = torch.zeros(
|
|
num_blocks,
|
|
block_size,
|
|
num_kv_heads,
|
|
slot_size(head_dim),
|
|
dtype=torch.uint8,
|
|
device=device,
|
|
)
|
|
# Leave block 0 unused so the block table exercises a non-identity mapping.
|
|
slot_mapping = torch.arange(
|
|
block_size, block_size + alloc_len, dtype=torch.int64, device=device
|
|
)
|
|
ultraquant_store(key, value, kv_cache, slot_mapping)
|
|
|
|
block_table = torch.arange(
|
|
1, num_blocks, dtype=torch.int32, device=device
|
|
).unsqueeze(0)
|
|
|
|
k_out = torch.zeros(
|
|
1, num_kv_heads, alloc_len, head_dim, dtype=dtype, device=device
|
|
)
|
|
v_out = torch.zeros_like(k_out)
|
|
ultraquant_full_dequant_kv(
|
|
kv_cache=kv_cache,
|
|
block_table=block_table,
|
|
k_out=k_out,
|
|
v_out=v_out,
|
|
alloc_len=alloc_len,
|
|
)
|
|
torch.accelerator.synchronize()
|
|
|
|
# K is stored Hadamard-rotated, V unrotated; the kernel returns them as
|
|
# stored, so the reference must not inverse-rotate either.
|
|
ref_k = ultraquant_dequant(ultraquant_encode(key, rotate=True), head_dim)
|
|
ref_v = ultraquant_dequant(ultraquant_encode(value, rotate=False), head_dim)
|
|
|
|
got_k = k_out[0].transpose(0, 1).to(torch.float32)
|
|
got_v = v_out[0].transpose(0, 1).to(torch.float32)
|
|
|
|
# Reference is fp32; the kernel rounds to the output dtype once.
|
|
tol = 1e-2 if dtype is torch.bfloat16 else 1e-3
|
|
torch.testing.assert_close(got_k, ref_k.to(torch.float32), rtol=tol, atol=tol)
|
|
torch.testing.assert_close(got_v, ref_v.to(torch.float32), rtol=tol, atol=tol)
|