150 lines
5.1 KiB
Python
150 lines
5.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Tests for the Triton UltraQuant store kernel.
|
|
|
|
The Triton kernel must produce bit-identical codes and scale bytes vs
|
|
the PyTorch reference (``reference.ultraquant_encode``).
|
|
"""
|
|
|
|
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 (
|
|
k_codes_offset,
|
|
k_scales_offset,
|
|
n_groups,
|
|
slot_size,
|
|
v_codes_offset,
|
|
v_scales_offset,
|
|
)
|
|
from vllm.v1.attention.ops.ultraquant.reference import ultraquant_encode
|
|
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"
|
|
)
|
|
|
|
|
|
def _make_cache(
|
|
num_blocks: int,
|
|
block_size: int,
|
|
num_kv_heads: int,
|
|
head_dim: int,
|
|
device: torch.device,
|
|
) -> torch.Tensor:
|
|
slot_b = slot_size(head_dim)
|
|
return torch.zeros(
|
|
num_blocks, block_size, num_kv_heads, slot_b, dtype=torch.uint8, device=device
|
|
)
|
|
|
|
|
|
def _read_codes_and_scales(
|
|
kv_cache: torch.Tensor,
|
|
slot_mapping: torch.Tensor,
|
|
num_kv_heads: int,
|
|
head_dim: int,
|
|
):
|
|
"""Pull (k_codes, k_scale_bytes, v_codes, v_scale_bytes) for each
|
|
(token, head) from the AoS slot layout. Scales come back as uint8
|
|
(one E8M0 byte per group)."""
|
|
block_size = kv_cache.shape[1]
|
|
N = slot_mapping.shape[0]
|
|
Gk = n_groups(head_dim)
|
|
|
|
k_codes_off = k_codes_offset(head_dim)
|
|
k_scales_off = k_scales_offset(head_dim)
|
|
v_codes_off = v_codes_offset(head_dim)
|
|
v_scales_off = v_scales_offset(head_dim)
|
|
codes_bytes = head_dim // 2
|
|
|
|
k_codes = torch.zeros(
|
|
N, num_kv_heads, codes_bytes, dtype=torch.uint8, device=kv_cache.device
|
|
)
|
|
k_scale_bytes = torch.zeros(
|
|
N, num_kv_heads, Gk, dtype=torch.uint8, device=kv_cache.device
|
|
)
|
|
v_codes = torch.zeros_like(k_codes)
|
|
v_scale_bytes = torch.zeros_like(k_scale_bytes)
|
|
|
|
for ti in range(N):
|
|
slot = int(slot_mapping[ti].item())
|
|
if slot < 0:
|
|
continue
|
|
blk = slot // block_size
|
|
off = slot % block_size
|
|
for h in range(num_kv_heads):
|
|
slot_bytes = kv_cache[blk, off, h] # [slot_size_aligned] uint8
|
|
|
|
k_codes[ti, h] = slot_bytes[k_codes_off : k_codes_off + codes_bytes]
|
|
k_scale_bytes[ti, h] = slot_bytes[k_scales_off : k_scales_off + Gk]
|
|
v_codes[ti, h] = slot_bytes[v_codes_off : v_codes_off + codes_bytes]
|
|
v_scale_bytes[ti, h] = slot_bytes[v_scales_off : v_scales_off + Gk]
|
|
|
|
return k_codes, k_scale_bytes, v_codes, v_scale_bytes
|
|
|
|
|
|
@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("num_tokens", [1, 16, 17])
|
|
def test_store_matches_reference(dtype, head_dim, num_kv_heads, num_tokens):
|
|
"""Triton store output (codes + E8M0 scale bytes) matches the reference
|
|
bit-exactly."""
|
|
device = torch.device("cuda")
|
|
set_random_seed(0xF8)
|
|
|
|
block_size = 32
|
|
num_blocks = max(2, (num_tokens + block_size - 1) // block_size + 1)
|
|
|
|
key = torch.randn(num_tokens, num_kv_heads, head_dim, dtype=dtype, device=device)
|
|
value = torch.randn(num_tokens, num_kv_heads, head_dim, dtype=dtype, device=device)
|
|
|
|
slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
|
|
if num_tokens >= 4:
|
|
slot_mapping[1] = -1 # skip-path test
|
|
|
|
kv_cache = _make_cache(num_blocks, block_size, num_kv_heads, head_dim, device)
|
|
|
|
ultraquant_store(key, value, kv_cache, slot_mapping)
|
|
torch.accelerator.synchronize()
|
|
|
|
ref_k = ultraquant_encode(key, rotate=True)
|
|
ref_v = ultraquant_encode(value, rotate=False)
|
|
|
|
k_codes, k_scale_bytes, v_codes, v_scale_bytes = _read_codes_and_scales(
|
|
kv_cache, slot_mapping, num_kv_heads, head_dim
|
|
)
|
|
|
|
if num_tokens >= 4:
|
|
assert (kv_cache[0, 1, :, :] == 0).all(), (
|
|
"skipped slot's physical cache region was unexpectedly written"
|
|
)
|
|
|
|
for ti in range(num_tokens):
|
|
slot = int(slot_mapping[ti].item())
|
|
if slot < 0:
|
|
continue
|
|
|
|
assert torch.equal(k_codes[ti], ref_k.codes_packed[ti]), (
|
|
f"K codes mismatch at token {ti}\n"
|
|
f"got: {k_codes[ti].cpu().numpy().tolist()[:16]}\n"
|
|
f"expected: {ref_k.codes_packed[ti].cpu().numpy().tolist()[:16]}"
|
|
)
|
|
assert torch.equal(k_scale_bytes[ti], ref_k.scale_bytes[ti]), (
|
|
f"K scale bytes mismatch at token {ti}\n"
|
|
f"got: {k_scale_bytes[ti].cpu().numpy().tolist()}\n"
|
|
f"expected: {ref_k.scale_bytes[ti].cpu().numpy().tolist()}"
|
|
)
|
|
|
|
assert torch.equal(v_codes[ti], ref_v.codes_packed[ti]), (
|
|
f"V codes mismatch at token {ti}"
|
|
)
|
|
assert torch.equal(v_scale_bytes[ti], ref_v.scale_bytes[ti]), (
|
|
f"V scale bytes mismatch at token {ti}\n"
|
|
f"got: {v_scale_bytes[ti].cpu().numpy().tolist()}\n"
|
|
f"expected: {ref_v.scale_bytes[ti].cpu().numpy().tolist()}"
|
|
)
|