wehub-resource-sync
724 lines
26 KiB
Python
724 lines
26 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Tests for per-token-head KV cache quantization (INT4, INT8 and FP8).
|
|
|
|
Covers:
|
|
- Per-token-head Triton reshape-and-cache kernel
|
|
- Round-trip quantize/dequantize accuracy
|
|
- process_weights_after_loading early-return path
|
|
- End-to-end integration with Triton unified attention kernel
|
|
|
|
Run: pytest tests/quantization/test_per_token_kv_cache.py -v -s
|
|
"""
|
|
|
|
import random
|
|
from dataclasses import dataclass
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
|
get_fp8_min_max,
|
|
)
|
|
from vllm.platforms import current_platform
|
|
from vllm.utils.torch_utils import set_random_seed
|
|
from vllm.v1.attention.ops.int4_per_token_head import single_rht
|
|
from vllm.v1.kv_cache_interface import KVQuantMode, is_quantized_kv_cache
|
|
|
|
DEVICE_TYPE = current_platform.device_type
|
|
|
|
# Skip entire module if no CUDA/ROCm GPU available
|
|
pytestmark = [
|
|
pytest.mark.skipif(
|
|
current_platform.is_cpu(),
|
|
reason="Per-token-head KV cache tests require GPU.",
|
|
),
|
|
]
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Test parameters
|
|
# ---------------------------------------------------------------------------
|
|
NUM_TOKENS = [1, 7, 42]
|
|
NUM_KV_HEADS = [1, 4, 8]
|
|
HEAD_SIZES = [64, 128]
|
|
BLOCK_SIZES = [16]
|
|
SEEDS = [0]
|
|
|
|
# Platform-dependent FP8 dtype and range
|
|
FP8_DTYPE = current_platform.fp8_dtype()
|
|
FP8_MIN, FP8_MAX = get_fp8_min_max()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Per-dtype quantization config
|
|
# ---------------------------------------------------------------------------
|
|
@dataclass(frozen=True)
|
|
class QuantConfig:
|
|
"""Quantization parameters for a given cache dtype."""
|
|
|
|
cache_dtype: torch.dtype # torch.int8 or FP8_DTYPE
|
|
kv_cache_dtype_str: str # "int8_per_token_head" or "fp8_per_token_head"
|
|
quant_max: float
|
|
quant_min: float
|
|
kv_quant_mode: KVQuantMode
|
|
# INT8 rounds explicitly; FP8 relies on dtype cast rounding.
|
|
rounds_before_store: bool
|
|
|
|
|
|
INT8_CONFIG = QuantConfig(
|
|
cache_dtype=torch.int8,
|
|
kv_cache_dtype_str="int8_per_token_head",
|
|
quant_max=127.0,
|
|
quant_min=-128.0,
|
|
kv_quant_mode=KVQuantMode.INT8_PER_TOKEN_HEAD,
|
|
rounds_before_store=True,
|
|
)
|
|
FP8_CONFIG = QuantConfig(
|
|
cache_dtype=FP8_DTYPE,
|
|
kv_cache_dtype_str="fp8_per_token_head",
|
|
quant_max=FP8_MAX,
|
|
quant_min=FP8_MIN,
|
|
kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD,
|
|
rounds_before_store=False,
|
|
)
|
|
INT4_CONFIG = QuantConfig(
|
|
cache_dtype=torch.uint8,
|
|
kv_cache_dtype_str="int4_per_token_head",
|
|
quant_max=7.0,
|
|
quant_min=-8.0,
|
|
kv_quant_mode=KVQuantMode.INT4_PER_TOKEN_HEAD,
|
|
# Unused for int4 (handled by its own rint path); kept for the dataclass.
|
|
rounds_before_store=False,
|
|
)
|
|
QUANT_CONFIGS = [INT4_CONFIG, INT8_CONFIG, FP8_CONFIG]
|
|
|
|
|
|
@pytest.fixture(params=QUANT_CONFIGS, ids=["int4", "int8", "fp8"])
|
|
def qcfg(request) -> QuantConfig:
|
|
return request.param
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
def _quantize_per_token_head_ref(
|
|
data: torch.Tensor, # [num_tokens, num_heads, head_size]
|
|
cfg: QuantConfig,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Reference per-token-head quantization (one scale per token per head).
|
|
|
|
Returns (quantized, scales) where scales is [num_tokens, num_heads].
|
|
"""
|
|
absmax = data.float().abs().amax(dim=2) # [num_tokens, num_heads]
|
|
scales = (absmax / cfg.quant_max).clamp(min=1e-6)
|
|
scaled = data.float() * (1.0 / scales[:, :, None])
|
|
if cfg.rounds_before_store:
|
|
q = scaled.round().clamp(cfg.quant_min, cfg.quant_max).to(cfg.cache_dtype)
|
|
else:
|
|
q = scaled.clamp(cfg.quant_min, cfg.quant_max).to(cfg.cache_dtype)
|
|
return q, scales
|
|
|
|
|
|
# ===========================================================================
|
|
# 1. is_quantized_kv_cache / get_kv_quant_mode
|
|
# ===========================================================================
|
|
class TestIsQuantizedKvCache:
|
|
def test_fp8_variants(self):
|
|
assert is_quantized_kv_cache("fp8")
|
|
assert is_quantized_kv_cache("fp8_e4m3")
|
|
assert is_quantized_kv_cache("fp8_e5m2")
|
|
|
|
def test_int4_per_token_head(self):
|
|
assert is_quantized_kv_cache("int4_per_token_head")
|
|
|
|
def test_int8_per_token_head(self):
|
|
assert is_quantized_kv_cache("int8_per_token_head")
|
|
|
|
def test_fp8_per_token_head(self):
|
|
assert is_quantized_kv_cache("fp8_per_token_head")
|
|
|
|
def test_auto(self):
|
|
assert not is_quantized_kv_cache("auto")
|
|
|
|
def test_bfloat16(self):
|
|
assert not is_quantized_kv_cache("bfloat16")
|
|
|
|
def test_kv_quant_mode_int4(self):
|
|
from vllm.v1.kv_cache_interface import get_kv_quant_mode
|
|
|
|
assert (
|
|
get_kv_quant_mode("int4_per_token_head") == KVQuantMode.INT4_PER_TOKEN_HEAD
|
|
)
|
|
|
|
def test_kv_quant_mode_int8(self):
|
|
from vllm.v1.kv_cache_interface import get_kv_quant_mode
|
|
|
|
assert (
|
|
get_kv_quant_mode("int8_per_token_head") == KVQuantMode.INT8_PER_TOKEN_HEAD
|
|
)
|
|
|
|
def test_kv_quant_mode_fp8(self):
|
|
from vllm.v1.kv_cache_interface import get_kv_quant_mode
|
|
|
|
assert get_kv_quant_mode("fp8_per_token_head") == KVQuantMode.FP8_PER_TOKEN_HEAD
|
|
|
|
|
|
# ===========================================================================
|
|
# 2. Triton per-token-head kernel (reshape-and-cache)
|
|
# ===========================================================================
|
|
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
|
|
@pytest.mark.parametrize("num_heads", NUM_KV_HEADS)
|
|
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
|
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
|
|
@pytest.mark.parametrize("seed", SEEDS)
|
|
@torch.inference_mode()
|
|
def test_reshape_and_cache_per_token_head(
|
|
qcfg: QuantConfig,
|
|
num_tokens: int,
|
|
num_heads: int,
|
|
head_size: int,
|
|
block_size: int,
|
|
seed: int,
|
|
):
|
|
"""Test triton_reshape_and_cache_flash_per_token_head_quant kernel."""
|
|
from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
|
|
triton_reshape_and_cache_flash_per_token_head_quant,
|
|
)
|
|
|
|
set_random_seed(seed)
|
|
torch.set_default_device(DEVICE_TYPE)
|
|
|
|
num_blocks = (num_tokens + block_size - 1) // block_size + 4
|
|
is_int4 = qcfg.kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD
|
|
cache_head_size = head_size // 2 if is_int4 else head_size
|
|
|
|
key = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16)
|
|
value = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16)
|
|
|
|
key_cache = torch.zeros(
|
|
num_blocks, block_size, num_heads, cache_head_size, dtype=qcfg.cache_dtype
|
|
)
|
|
value_cache = torch.zeros(
|
|
num_blocks, block_size, num_heads, cache_head_size, dtype=qcfg.cache_dtype
|
|
)
|
|
k_scale_cache = torch.ones(num_blocks, block_size, num_heads, dtype=torch.float32)
|
|
v_scale_cache = torch.ones(num_blocks, block_size, num_heads, dtype=torch.float32)
|
|
|
|
num_slots = block_size * num_blocks
|
|
slot_mapping = torch.tensor(
|
|
random.sample(range(num_slots), num_tokens), dtype=torch.long
|
|
)
|
|
|
|
triton_reshape_and_cache_flash_per_token_head_quant(
|
|
key,
|
|
value,
|
|
key_cache,
|
|
value_cache,
|
|
k_scale_cache,
|
|
v_scale_cache,
|
|
slot_mapping,
|
|
kv_quant_mode=qcfg.kv_quant_mode,
|
|
)
|
|
|
|
# INT4 (RHT + asymmetric), INT8/FP8 have different dequant paths. Only
|
|
# INT8/FP8 can be compared to a PyTorch reference.
|
|
if not is_int4:
|
|
ref_k_quant, ref_k_scales = _quantize_per_token_head_ref(key, qcfg)
|
|
ref_v_quant, ref_v_scales = _quantize_per_token_head_ref(value, qcfg)
|
|
|
|
for i, slot in enumerate(slot_mapping.tolist()):
|
|
blk = slot // block_size
|
|
off = slot % block_size
|
|
|
|
if is_int4:
|
|
# Coarser quantization → wider tolerance.
|
|
deq_atol = deq_rtol = 0.5
|
|
for label, data, cache, sc in [
|
|
("key", key, key_cache, k_scale_cache),
|
|
("val", value, value_cache, v_scale_cache),
|
|
]:
|
|
packed_scale = sc[blk, off] # [num_heads] float32
|
|
scale_bits = packed_scale.view(torch.int32)
|
|
zp = (scale_bits & 0xF).to(torch.float32)
|
|
clean_scale = (scale_bits & -16).view(torch.float32)
|
|
|
|
packed = cache[blk, off]
|
|
lo = (packed & 0xF).to(torch.float32)
|
|
hi = ((packed >> 4) & 0xF).to(torch.float32)
|
|
full = torch.zeros(num_heads, head_size, dtype=torch.float32)
|
|
full[:, 0::2] = lo
|
|
full[:, 1::2] = hi
|
|
# Asymmetric dequant in RHT domain, then IRHT/d → original
|
|
deq_rht = (full - zp[:, None]) * clean_scale[:, None]
|
|
deq = single_rht(deq_rht, inverse=True) / head_size
|
|
ref_deq = data[i].float()
|
|
torch.testing.assert_close(deq, ref_deq, atol=deq_atol, rtol=deq_rtol)
|
|
else:
|
|
actual_k_scale = k_scale_cache[blk, off] # [num_heads]
|
|
k_deq = key_cache[blk, off].float() * actual_k_scale[:, None]
|
|
k_ref_deq = key[i].float()
|
|
torch.testing.assert_close(
|
|
k_deq,
|
|
k_ref_deq,
|
|
atol=0.1,
|
|
rtol=0.1,
|
|
)
|
|
actual_v_scale = v_scale_cache[blk, off] # [num_heads]
|
|
v_deq = value_cache[blk, off].float() * actual_v_scale[:, None]
|
|
v_ref_deq = value[i].float()
|
|
torch.testing.assert_close(
|
|
v_deq,
|
|
v_ref_deq,
|
|
atol=0.1,
|
|
rtol=0.1,
|
|
)
|
|
# Per-head scales: [num_heads]
|
|
torch.testing.assert_close(
|
|
k_scale_cache[blk, off], ref_k_scales[i], atol=1e-4, rtol=1e-3
|
|
)
|
|
torch.testing.assert_close(
|
|
v_scale_cache[blk, off], ref_v_scales[i], atol=1e-4, rtol=1e-3
|
|
)
|
|
|
|
|
|
# ===========================================================================
|
|
# 3. Per-token-head round-trip accuracy (quantize -> dequantize)
|
|
# ===========================================================================
|
|
@pytest.mark.parametrize("num_tokens", [1, 16])
|
|
@pytest.mark.parametrize("num_heads", [4])
|
|
@pytest.mark.parametrize("head_size", [128])
|
|
@pytest.mark.parametrize("block_size", [16])
|
|
@torch.inference_mode()
|
|
def test_per_token_head_round_trip_accuracy(
|
|
qcfg: QuantConfig,
|
|
num_tokens: int,
|
|
num_heads: int,
|
|
head_size: int,
|
|
block_size: int,
|
|
):
|
|
"""Verify per-token-head round-trip: kernel dequant matches reference.
|
|
|
|
INT8: round-to-nearest before int8 store.
|
|
FP8: hardware cast (clamp then cast).
|
|
"""
|
|
from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
|
|
triton_reshape_and_cache_flash_per_token_head_quant,
|
|
)
|
|
|
|
torch.set_default_device(DEVICE_TYPE)
|
|
set_random_seed(42)
|
|
|
|
is_int4 = qcfg.kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD
|
|
num_blocks = (num_tokens + block_size - 1) // block_size + 2
|
|
cache_head_size = head_size // 2 if is_int4 else head_size
|
|
|
|
key = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16) * 0.5
|
|
value = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16) * 0.5
|
|
|
|
key_cache = torch.zeros(
|
|
num_blocks, block_size, num_heads, cache_head_size, dtype=qcfg.cache_dtype
|
|
)
|
|
value_cache = torch.zeros(
|
|
num_blocks, block_size, num_heads, cache_head_size, dtype=qcfg.cache_dtype
|
|
)
|
|
k_scale_cache = torch.ones(num_blocks, block_size, num_heads, dtype=torch.float32)
|
|
v_scale_cache = torch.ones(num_blocks, block_size, num_heads, dtype=torch.float32)
|
|
|
|
slot_mapping = torch.arange(num_tokens, dtype=torch.long)
|
|
|
|
triton_reshape_and_cache_flash_per_token_head_quant(
|
|
key,
|
|
value,
|
|
key_cache,
|
|
value_cache,
|
|
k_scale_cache,
|
|
v_scale_cache,
|
|
slot_mapping,
|
|
kv_quant_mode=qcfg.kv_quant_mode,
|
|
)
|
|
|
|
rt_atol = 0.5 if is_int4 else 0.1
|
|
|
|
for i in range(num_tokens):
|
|
blk = i // block_size
|
|
off = i % block_size
|
|
|
|
for label, data, cache, sc in [
|
|
("key", key, key_cache, k_scale_cache),
|
|
("val", value, value_cache, v_scale_cache),
|
|
]:
|
|
for h in range(num_heads):
|
|
orig = data[i, h].float()
|
|
actual_sc = sc[blk, off, h]
|
|
if is_int4:
|
|
sc_bits = actual_sc.view(torch.int32)
|
|
zp = (sc_bits & 0xF).to(torch.float32)
|
|
clean_sc = (sc_bits & -16).view(torch.float32)
|
|
packed = cache[blk, off, h]
|
|
lo = (packed & 0xF).to(torch.float32)
|
|
hi = ((packed >> 4) & 0xF).to(torch.float32)
|
|
full = torch.zeros(head_size)
|
|
full[0::2] = lo
|
|
full[1::2] = hi
|
|
deq_rht = (full - zp) * clean_sc
|
|
actual_deq = (
|
|
single_rht(deq_rht.unsqueeze(0), inverse=True).squeeze(0)
|
|
/ head_size
|
|
)
|
|
else:
|
|
actual_deq = cache[blk, off, h].float() * actual_sc
|
|
torch.testing.assert_close(
|
|
actual_deq,
|
|
orig,
|
|
atol=rt_atol,
|
|
rtol=rt_atol,
|
|
)
|
|
|
|
|
|
@torch.inference_mode()
|
|
def test_int8_per_token_head_raw_cache_matches_round_reference():
|
|
"""INT8 cache writes should match round-to-nearest quantization exactly."""
|
|
from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
|
|
triton_reshape_and_cache_flash_per_token_head_quant,
|
|
)
|
|
|
|
torch.set_default_device(DEVICE_TYPE)
|
|
|
|
head_size = 8
|
|
block_size = 4
|
|
|
|
key = torch.tensor(
|
|
[[[-127.0, -2.6, -2.4, -1.6, -1.4, -0.6, -0.4, 127.0]]],
|
|
dtype=torch.bfloat16,
|
|
)
|
|
value = -key
|
|
|
|
key_cache = torch.zeros(1, block_size, 1, head_size, dtype=torch.int8)
|
|
value_cache = torch.zeros_like(key_cache)
|
|
k_scale_cache = torch.ones(1, block_size, 1, dtype=torch.float32)
|
|
v_scale_cache = torch.ones_like(k_scale_cache)
|
|
slot_mapping = torch.tensor([2], dtype=torch.long)
|
|
|
|
triton_reshape_and_cache_flash_per_token_head_quant(
|
|
key,
|
|
value,
|
|
key_cache,
|
|
value_cache,
|
|
k_scale_cache,
|
|
v_scale_cache,
|
|
slot_mapping,
|
|
kv_quant_mode=INT8_CONFIG.kv_quant_mode,
|
|
)
|
|
|
|
ref_k_quant, ref_k_scales = _quantize_per_token_head_ref(key, INT8_CONFIG)
|
|
ref_v_quant, ref_v_scales = _quantize_per_token_head_ref(value, INT8_CONFIG)
|
|
|
|
slot = slot_mapping.item()
|
|
blk = slot // block_size
|
|
off = slot % block_size
|
|
assert torch.equal(key_cache[blk, off], ref_k_quant[0])
|
|
assert torch.equal(value_cache[blk, off], ref_v_quant[0])
|
|
torch.testing.assert_close(k_scale_cache[blk, off], ref_k_scales[0])
|
|
torch.testing.assert_close(v_scale_cache[blk, off], ref_v_scales[0])
|
|
|
|
|
|
# ===========================================================================
|
|
# 4. Negative slot mapping (padding tokens should be skipped)
|
|
# ===========================================================================
|
|
@torch.inference_mode()
|
|
def test_per_token_head_negative_slot_skipped(qcfg: QuantConfig):
|
|
"""Tokens with slot_mapping=-1 should leave the cache unchanged."""
|
|
from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
|
|
triton_reshape_and_cache_flash_per_token_head_quant,
|
|
)
|
|
|
|
torch.set_default_device(DEVICE_TYPE)
|
|
num_tokens = 4
|
|
num_heads = 2
|
|
head_size = 64
|
|
block_size = 16
|
|
num_blocks = 2
|
|
is_int4 = qcfg.kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD
|
|
cache_head_size = head_size // 2 if is_int4 else head_size
|
|
|
|
key = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16)
|
|
value = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16)
|
|
|
|
key_cache = torch.zeros(
|
|
num_blocks, block_size, num_heads, cache_head_size, dtype=qcfg.cache_dtype
|
|
)
|
|
value_cache = torch.zeros(
|
|
num_blocks, block_size, num_heads, cache_head_size, dtype=qcfg.cache_dtype
|
|
)
|
|
k_scale_cache = torch.ones(num_blocks, block_size, num_heads, dtype=torch.float32)
|
|
v_scale_cache = torch.ones(num_blocks, block_size, num_heads, dtype=torch.float32)
|
|
|
|
slot_mapping = torch.tensor([0, -1, 1, -1], dtype=torch.long)
|
|
|
|
key_cache_before = key_cache.clone()
|
|
val_cache_before = value_cache.clone()
|
|
|
|
triton_reshape_and_cache_flash_per_token_head_quant(
|
|
key,
|
|
value,
|
|
key_cache,
|
|
value_cache,
|
|
k_scale_cache,
|
|
v_scale_cache,
|
|
slot_mapping,
|
|
kv_quant_mode=qcfg.kv_quant_mode,
|
|
)
|
|
|
|
# Slots 0 and 1 should have been written (tokens 0 and 2)
|
|
assert not torch.equal(key_cache[0, 0], key_cache_before[0, 0])
|
|
assert not torch.equal(key_cache[0, 1], key_cache_before[0, 1])
|
|
assert not torch.equal(value_cache[0, 0], val_cache_before[0, 0])
|
|
|
|
# All other slots should be unchanged
|
|
assert torch.equal(key_cache[0, 2:], key_cache_before[0, 2:])
|
|
assert torch.equal(key_cache[1], key_cache_before[1])
|
|
assert torch.equal(value_cache[0, 2:], val_cache_before[0, 2:])
|
|
|
|
|
|
# ===========================================================================
|
|
# 5. process_weights_after_loading -- per-token-head early return
|
|
# ===========================================================================
|
|
@pytest.mark.parametrize(
|
|
"kv_cache_dtype",
|
|
["int4_per_token_head", "int8_per_token_head", "fp8_per_token_head"],
|
|
)
|
|
def test_process_weights_sets_placeholder_scales(kv_cache_dtype: str):
|
|
"""Per-token-head should set _k_scale=1.0, _v_scale=1.0
|
|
and delete checkpoint attrs."""
|
|
from vllm.model_executor.layers.quantization.kv_cache import (
|
|
BaseKVCacheMethod,
|
|
)
|
|
|
|
layer = MagicMock()
|
|
layer.kv_cache_dtype = kv_cache_dtype
|
|
layer.calculate_kv_scales = False
|
|
layer.k_scale = torch.nn.Parameter(torch.tensor(-1.0), requires_grad=False)
|
|
layer.v_scale = torch.nn.Parameter(torch.tensor(-1.0), requires_grad=False)
|
|
layer.q_scale = torch.nn.Parameter(torch.tensor(-1.0), requires_grad=False)
|
|
layer.prob_scale = torch.nn.Parameter(torch.tensor(-1.0), requires_grad=False)
|
|
layer._k_scale = torch.tensor(0.0)
|
|
layer._v_scale = torch.tensor(0.0)
|
|
layer._k_scale_float = 0.0
|
|
layer._v_scale_float = 0.0
|
|
|
|
method = BaseKVCacheMethod.__new__(BaseKVCacheMethod)
|
|
method.quant_config = MagicMock()
|
|
method.process_weights_after_loading(layer)
|
|
|
|
assert layer._k_scale_float == 1.0
|
|
assert layer._v_scale_float == 1.0
|
|
assert not hasattr(layer, "k_scale")
|
|
assert not hasattr(layer, "v_scale")
|
|
assert not hasattr(layer, "q_scale")
|
|
assert not hasattr(layer, "prob_scale")
|
|
|
|
|
|
# ===========================================================================
|
|
# 6. Triton unified_attention -- per-token-head scale cache (INT4/INT8/FP8)
|
|
# ===========================================================================
|
|
@pytest.mark.parametrize(
|
|
"seq_lens",
|
|
[
|
|
[(1, 128)],
|
|
[(1, 64), (1, 32)],
|
|
],
|
|
)
|
|
@pytest.mark.parametrize("num_heads", [(4, 4)])
|
|
@pytest.mark.parametrize("head_size", [128])
|
|
@pytest.mark.parametrize("block_size", [16])
|
|
@torch.inference_mode()
|
|
def test_triton_unified_attention_per_token_head_scale(
|
|
qcfg: QuantConfig,
|
|
seq_lens: list[tuple[int, int]],
|
|
num_heads: tuple[int, int],
|
|
head_size: int,
|
|
block_size: int,
|
|
):
|
|
"""End-to-end: quantized KV with per-token-head scale caches."""
|
|
from vllm.utils.math_utils import next_power_of_2
|
|
from vllm.v1.attention.ops.triton_unified_attention import unified_attention
|
|
|
|
torch.set_default_device(DEVICE_TYPE)
|
|
set_random_seed(0)
|
|
|
|
is_int4 = qcfg.kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD
|
|
|
|
num_seqs = len(seq_lens)
|
|
query_lens = [s[0] for s in seq_lens]
|
|
kv_lens = [s[1] for s in seq_lens]
|
|
num_query_heads, num_kv_heads = num_heads
|
|
max_query_len = max(query_lens)
|
|
max_kv_len = max(kv_lens)
|
|
scale = head_size**-0.5
|
|
num_blocks = 2048
|
|
|
|
query = torch.randn(
|
|
sum(query_lens), num_query_heads, head_size, dtype=torch.bfloat16
|
|
)
|
|
|
|
key_cache_bf16 = torch.randn(
|
|
num_blocks, block_size, num_kv_heads, head_size, dtype=torch.bfloat16
|
|
)
|
|
value_cache_bf16 = torch.randn_like(key_cache_bf16)
|
|
|
|
if is_int4:
|
|
# Asymmetric quantization reference (matches the Triton kernel).
|
|
kf = key_cache_bf16.float()
|
|
vf = value_cache_bf16.float()
|
|
k_min = kf.amin(dim=-1)
|
|
k_max = kf.amax(dim=-1)
|
|
v_min = vf.amin(dim=-1)
|
|
v_max = vf.amax(dim=-1)
|
|
k_scale_cache = ((k_max - k_min) / 15.0).clamp(min=1e-6).to(torch.float32)
|
|
v_scale_cache = ((v_max - v_min) / 15.0).clamp(min=1e-6).to(torch.float32)
|
|
k_zp = (-k_min / k_scale_cache).round().clamp(0, 15)
|
|
v_zp = (-v_min / v_scale_cache).round().clamp(0, 15)
|
|
|
|
key_cache_q_full = (
|
|
(kf / k_scale_cache[..., None] + k_zp[..., None]).round().clamp(0, 15)
|
|
)
|
|
value_cache_q_full = (
|
|
(vf / v_scale_cache[..., None] + v_zp[..., None]).round().clamp(0, 15)
|
|
)
|
|
|
|
# Dequantized reference: x_hat = (q - zp) * scale
|
|
key_cache_deq = (key_cache_q_full - k_zp[..., None]) * k_scale_cache[..., None]
|
|
value_cache_deq = (value_cache_q_full - v_zp[..., None]) * v_scale_cache[
|
|
..., None
|
|
]
|
|
|
|
# Pack two uint4 values into one byte
|
|
def _pack_int4(data_float):
|
|
u = data_float.to(torch.uint8)
|
|
lo = u[..., 0::2]
|
|
hi = u[..., 1::2]
|
|
return (lo & 0xF) | ((hi & 0xF) << 4)
|
|
|
|
key_cache_q = _pack_int4(key_cache_q_full)
|
|
value_cache_q = _pack_int4(value_cache_q_full)
|
|
|
|
# Steganography: pack zp into low 4 bits of scale
|
|
k_zp_int = k_zp.to(torch.int32)
|
|
k_bits = k_scale_cache.view(torch.int32)
|
|
k_scale_cache = ((k_bits & -16) | (k_zp_int & 0xF)).view(torch.float32)
|
|
v_zp_int = v_zp.to(torch.int32)
|
|
v_bits = v_scale_cache.view(torch.int32)
|
|
v_scale_cache = ((v_bits & -16) | (v_zp_int & 0xF)).view(torch.float32)
|
|
else:
|
|
# Symmetric quantization for int8/fp8.
|
|
k_absmax = key_cache_bf16.float().abs().amax(dim=-1)
|
|
v_absmax = value_cache_bf16.float().abs().amax(dim=-1)
|
|
k_scale_cache = (k_absmax / qcfg.quant_max).clamp(min=1e-6).to(torch.float32)
|
|
v_scale_cache = (v_absmax / qcfg.quant_max).clamp(min=1e-6).to(torch.float32)
|
|
scaled_k = key_cache_bf16.float() / k_scale_cache[:, :, :, None]
|
|
scaled_v = value_cache_bf16.float() / v_scale_cache[:, :, :, None]
|
|
|
|
key_cache_q_full = scaled_k.round().clamp(qcfg.quant_min, qcfg.quant_max)
|
|
value_cache_q_full = scaled_v.round().clamp(qcfg.quant_min, qcfg.quant_max)
|
|
|
|
key_cache_deq = key_cache_q_full * k_scale_cache[:, :, :, None]
|
|
value_cache_deq = value_cache_q_full * v_scale_cache[:, :, :, None]
|
|
|
|
if not is_int4 and qcfg.rounds_before_store:
|
|
key_cache_q = key_cache_q_full.to(qcfg.cache_dtype)
|
|
value_cache_q = value_cache_q_full.to(qcfg.cache_dtype)
|
|
elif not is_int4:
|
|
key_cache_q = scaled_k.clamp(qcfg.quant_min, qcfg.quant_max).to(
|
|
qcfg.cache_dtype
|
|
)
|
|
value_cache_q = scaled_v.clamp(qcfg.quant_min, qcfg.quant_max).to(
|
|
qcfg.cache_dtype
|
|
)
|
|
|
|
cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
|
|
dim=0, dtype=torch.int32
|
|
)
|
|
kv_lens_t = 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
|
|
)
|
|
|
|
head_size_padded = next_power_of_2(head_size)
|
|
seq_threshold_3D = 0
|
|
num_par_softmax_segments = 16
|
|
softmax_segm_output = torch.empty(
|
|
(seq_threshold_3D, num_query_heads, num_par_softmax_segments, head_size_padded),
|
|
dtype=torch.float32,
|
|
)
|
|
softmax_segm_max = torch.empty(
|
|
(seq_threshold_3D, num_query_heads, num_par_softmax_segments),
|
|
dtype=torch.float32,
|
|
)
|
|
softmax_segm_expsum = torch.empty(
|
|
(seq_threshold_3D, num_query_heads, num_par_softmax_segments),
|
|
dtype=torch.float32,
|
|
)
|
|
|
|
output_q = torch.empty_like(query)
|
|
unified_attention(
|
|
q=query,
|
|
k=key_cache_q,
|
|
v=value_cache_q,
|
|
out=output_q,
|
|
cu_seqlens_q=cu_query_lens,
|
|
seqused_k=kv_lens_t,
|
|
max_seqlen_q=max_query_len,
|
|
max_seqlen_k=max_kv_len,
|
|
softmax_scale=scale,
|
|
causal=True,
|
|
window_size=(-1, -1),
|
|
block_table=block_tables,
|
|
softcap=0,
|
|
q_descale=None,
|
|
k_descale=None,
|
|
v_descale=None,
|
|
seq_threshold_3D=seq_threshold_3D,
|
|
num_par_softmax_segments=num_par_softmax_segments,
|
|
softmax_segm_output=softmax_segm_output,
|
|
softmax_segm_max=softmax_segm_max,
|
|
softmax_segm_expsum=softmax_segm_expsum,
|
|
kv_quant_mode=qcfg.kv_quant_mode,
|
|
k_scale_cache=k_scale_cache,
|
|
v_scale_cache=v_scale_cache,
|
|
)
|
|
|
|
output_ref = torch.empty_like(query)
|
|
unified_attention(
|
|
q=query,
|
|
k=key_cache_deq.to(torch.bfloat16),
|
|
v=value_cache_deq.to(torch.bfloat16),
|
|
out=output_ref,
|
|
cu_seqlens_q=cu_query_lens,
|
|
seqused_k=kv_lens_t,
|
|
max_seqlen_q=max_query_len,
|
|
max_seqlen_k=max_kv_len,
|
|
softmax_scale=scale,
|
|
causal=True,
|
|
window_size=(-1, -1),
|
|
block_table=block_tables,
|
|
softcap=0,
|
|
q_descale=None,
|
|
k_descale=None,
|
|
v_descale=None,
|
|
seq_threshold_3D=seq_threshold_3D,
|
|
num_par_softmax_segments=num_par_softmax_segments,
|
|
softmax_segm_output=softmax_segm_output,
|
|
softmax_segm_max=softmax_segm_max,
|
|
softmax_segm_expsum=softmax_segm_expsum,
|
|
)
|
|
|
|
# Coarser quantization → wider tolerance.
|
|
if is_int4:
|
|
atol, rtol = 0.5, 0.5
|
|
else:
|
|
atol, rtol = 5e-2, 5e-2
|
|
torch.testing.assert_close(output_q, output_ref, atol=atol, rtol=rtol)
|