iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Feature] Triton INT4 per-token-head KV cache quantization (#40835)

Browse Source 
Signed-off-by: JartX <sagformas@epdcenter.es>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
JartX
2026-06-24 10:21:25 +00:00
committed by GitHub
co-authored by mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent d20dbf921b
commit 70749fdcca
10 changed files with 1462 additions and 111 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/design/attention_backends.md
+1 -1
View File
@@ -170,7 +170,7 @@ Priority is **1 = highest** (tried first).
| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ✅ | ✅ | ❌ | ❌ | Decoder | N/A |
| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |
| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A |
| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ✅ | ✅ | ❌ | All | Any |
| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int4_per_token_head`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ✅ | ✅ | ❌ | All | Any |
| `TRITON_ATTN_DIFFKV` | | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
| `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
tests/models/quantization/test_per_token_kv_cache.py
+2 -1
View File
@@ -31,7 +31,8 @@ from ..utils import check_logprobs_close
],
)
@pytest.mark.parametrize(
"kv_cache_dtype", ["int8_per_token_head", "fp8_per_token_head"]
"kv_cache_dtype",
["int4_per_token_head", "int8_per_token_head", "fp8_per_token_head"],
)
@pytest.mark.parametrize("max_tokens", [4])
@pytest.mark.parametrize("enforce_eager", [True])
tests/quantization/test_per_token_kv_cache.py
+184 -69
View File
@@ -1,6 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for per-token-head KV cache quantization (INT8 and FP8).
"""Tests for per-token-head KV cache quantization (INT4, INT8 and FP8).
Covers:
- Per-token-head Triton reshape-and-cache kernel
@@ -23,6 +23,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
)
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
@@ -81,11 +82,19 @@ FP8_CONFIG = QuantConfig(
kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD,
rounds_before_store=False,
)
QUANT_CONFIGS = [INT8_CONFIG, FP8_CONFIG]
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=["int8", "fp8"])
@pytest.fixture(params=QUANT_CONFIGS, ids=["int4", "int8", "fp8"])
def qcfg(request) -> QuantConfig:
return request.param
@@ -120,6 +129,9 @@ class TestIsQuantizedKvCache:
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")
@@ -132,6 +144,13 @@ class TestIsQuantizedKvCache:
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
@@ -171,15 +190,17 @@ def test_reshape_and_cache_per_token_head(
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, head_size, dtype=qcfg.cache_dtype
num_blocks, block_size, num_heads, cache_head_size, dtype=qcfg.cache_dtype
)
value_cache = torch.zeros(
num_blocks, block_size, num_heads, head_size, dtype=qcfg.cache_dtype
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)
@@ -197,45 +218,68 @@ def test_reshape_and_cache_per_token_head(
k_scale_cache,
v_scale_cache,
slot_mapping,
kv_quant_mode=qcfg.kv_quant_mode,
)
# Reference
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)
# 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)
# Compare dequantized values rather than raw quantized values.
# Triton and PyTorch reductions can differ at FP8 rounding boundaries
# (up to 32 in quantized domain for fp8_e4m3), but the dequantized
# error is bounded by the scale.
for i, slot in enumerate(slot_mapping.tolist()):
blk = slot // block_size
off = slot % block_size
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
)
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
)
# ===========================================================================
@@ -265,16 +309,18 @@ def test_per_token_head_round_trip_accuracy(
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, head_size, dtype=qcfg.cache_dtype
num_blocks, block_size, num_heads, cache_head_size, dtype=qcfg.cache_dtype
)
value_cache = torch.zeros(
num_blocks, block_size, num_heads, head_size, dtype=qcfg.cache_dtype
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)
@@ -289,8 +335,11 @@ def test_per_token_head_round_trip_accuracy(
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
@@ -300,18 +349,30 @@ def test_per_token_head_round_trip_accuracy(
("val", value, value_cache, v_scale_cache),
]:
for h in range(num_heads):
orig = data[i, h].float() # [head_size]
actual_q = cache[blk, off, h]
orig = data[i, h].float()
actual_sc = sc[blk, off, h]
actual_deq = actual_q.float() * actual_sc
# Round-trip: dequantized should be close to original
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=0.1,
rtol=0.1,
atol=rt_atol,
rtol=rt_atol,
)
@@ -347,6 +408,7 @@ def test_int8_per_token_head_raw_cache_matches_round_reference():
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)
@@ -377,15 +439,17 @@ def test_per_token_head_negative_slot_skipped(qcfg: QuantConfig):
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, head_size, dtype=qcfg.cache_dtype
num_blocks, block_size, num_heads, cache_head_size, dtype=qcfg.cache_dtype
)
value_cache = torch.zeros(
num_blocks, block_size, num_heads, head_size, dtype=qcfg.cache_dtype
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)
@@ -403,6 +467,7 @@ def test_per_token_head_negative_slot_skipped(qcfg: QuantConfig):
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)
@@ -420,7 +485,8 @@ def test_per_token_head_negative_slot_skipped(qcfg: QuantConfig):
# 5. process_weights_after_loading -- per-token-head early return
# ===========================================================================
@pytest.mark.parametrize(
"kv_cache_dtype", ["int8_per_token_head", "fp8_per_token_head"]
"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
@@ -454,7 +520,7 @@ def test_process_weights_sets_placeholder_scales(kv_cache_dtype: str):
# ===========================================================================
# 6. Triton unified_attention -- per-token-head scale cache (INT8 and FP8)
# 6. Triton unified_attention -- per-token-head scale cache (INT4/INT8/FP8)
# ===========================================================================
@pytest.mark.parametrize(
"seq_lens",
@@ -481,6 +547,8 @@ def test_triton_unified_attention_per_token_head_scale(
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]
@@ -499,22 +567,68 @@ def test_triton_unified_attention_per_token_head_scale(
)
value_cache_bf16 = torch.randn_like(key_cache_bf16)
# Per-token-head quantization: one scale per (block, slot, head)
k_absmax = key_cache_bf16.float().abs().amax(dim=-1) # [..., num_kv_heads]
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)
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)
scaled_k = key_cache_bf16.float() / k_scale_cache[:, :, :, None]
scaled_v = value_cache_bf16.float() / v_scale_cache[:, :, :, None]
if qcfg.rounds_before_store:
key_cache_q = (
scaled_k.round().clamp(qcfg.quant_min, qcfg.quant_max).to(qcfg.cache_dtype)
key_cache_q_full = (
(kf / k_scale_cache[..., None] + k_zp[..., None]).round().clamp(0, 15)
)
value_cache_q = (
scaled_v.round().clamp(qcfg.quant_min, qcfg.quant_max).to(qcfg.cache_dtype)
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
)
@@ -522,10 +636,6 @@ def test_triton_unified_attention_per_token_head_scale(
qcfg.cache_dtype
)
# Dequantized reference
key_cache_deq = key_cache_q.float() * k_scale_cache[:, :, :, None]
value_cache_deq = value_cache_q.float() * v_scale_cache[:, :, :, None]
cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
dim=0, dtype=torch.int32
)
@@ -605,4 +715,9 @@ def test_triton_unified_attention_per_token_head_scale(
softmax_segm_expsum=softmax_segm_expsum,
)
torch.testing.assert_close(output_q, output_ref, atol=5e-2, rtol=5e-2)
# 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)
vllm/config/cache.py
+1
View File
@@ -29,6 +29,7 @@ CacheDType = Literal[
"turboquant_4bit_nc",
"turboquant_k3v4_nc",
"turboquant_3bit_nc",
"int4_per_token_head",
"int8_per_token_head",
"fp8_per_token_head",
"nvfp4",
vllm/utils/torch_utils.py
+1
View File
@@ -39,6 +39,7 @@ STR_DTYPE_TO_TORCH_DTYPE = {
"fp8_e4m3": torch.uint8,
"fp8_e5m2": torch.uint8,
"int8": torch.int8,
"int4_per_token_head": torch.uint8,
"int8_per_token_head": torch.int8,
"fp8_per_token_head": torch.uint8,
"fp8_inc": torch.float8_e4m3fn,
vllm/v1/attention/backends/triton_attn.py
+33 -16
View File
@@ -258,6 +258,7 @@ class TritonAttentionBackend(AttentionBackend):
"fp8",
"fp8_e4m3",
"fp8_e5m2",
"int4_per_token_head",
"int8_per_token_head",
"fp8_per_token_head",
]
@@ -301,9 +302,11 @@ class TritonAttentionBackend(AttentionBackend):
if block_size % 16 != 0:
raise ValueError("Block size must be a multiple of 16.")
if kv_cache_uses_per_token_head_scales(cache_dtype_str):
# Pad head_size by sizeof(float32)/sizeof(cache_dtype) so
# the per-head scale fits inline. The backend extracts
# data[:head_size] and scale[head_size:] via typed views.
# Pad the head dim by sizeof(float32)/sizeof(cache_dtype) so the
# per-(token, head) scale fits inline after the quantized data;
# the backend extracts data[:head_size] and scale[head_size:] via
# typed views (see _ensure_scale_caches). INT4 packs two values
# per byte, so the data occupies only head_size // 2 bytes.
from vllm.utils.torch_utils import (
STR_DTYPE_TO_TORCH_DTYPE,
get_dtype_size,
@@ -311,7 +314,11 @@ class TritonAttentionBackend(AttentionBackend):
cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_dtype_str]
scale_pad = get_dtype_size(torch.float32) // get_dtype_size(cache_dtype)
return (num_blocks, 2, block_size, num_kv_heads, head_size + scale_pad)
if get_kv_quant_mode(cache_dtype_str) == KVQuantMode.INT4_PER_TOKEN_HEAD:
data_head_size = head_size // 2
else:
data_head_size = head_size
return (num_blocks, 2, block_size, num_kv_heads, data_head_size + scale_pad)
return (num_blocks, 2, block_size, num_kv_heads, head_size)
@staticmethod
@@ -587,18 +594,14 @@ class TritonAttentionImpl(AttentionImpl):
layer,
)
# Per-token-head quantized KV cache: use separate scale caches.
# Per-token-head quantized KV cache: handled by the core unified
# kernel, which dequantizes per-(token, head) inline via constexpr
# branches (INT8 / FP8) and dispatches to the packed INT4 kernel.
if self._is_per_token_head_quant:
self._ensure_scale_caches(kv_cache)
key_cache, value_cache = kv_cache.unbind(1)
if key_cache.dtype == torch.uint8:
key_cache = key_cache.view(self.fp8_dtype)
value_cache = value_cache.view(self.fp8_dtype)
q_descale = None
k_descale = None
v_descale = None
key_cache, value_cache = self._pth_key_value_caches(kv_cache)
k_scale_cache = self._k_scale_cache
v_scale_cache = self._v_scale_cache
q_descale = k_descale = v_descale = None
# FP8 per-tensor / auto path (original flow).
else:
key_cache, value_cache = kv_cache.unbind(1)
@@ -675,6 +678,17 @@ class TritonAttentionImpl(AttentionImpl):
return output
def _pth_key_value_caches(
self, kv_cache: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
"""Per-token-head K/V cache views (ensures scale caches; FP8 retyped)."""
self._ensure_scale_caches(kv_cache)
key_cache, value_cache = kv_cache.unbind(1)
if self._kv_quant_mode == KVQuantMode.FP8_PER_TOKEN_HEAD:
key_cache = key_cache.view(self.fp8_dtype)
value_cache = value_cache.view(self.fp8_dtype)
return key_cache, value_cache
def _forward_encoder_attention(
self,
query: torch.Tensor,
@@ -737,7 +751,9 @@ class TritonAttentionImpl(AttentionImpl):
if self._is_per_token_head_quant:
self._ensure_scale_caches(kv_cache)
key_cache, value_cache = kv_cache.unbind(1)
if key_cache.dtype == torch.uint8:
k_scale_cache = self._k_scale_cache
v_scale_cache = self._v_scale_cache
if self._kv_quant_mode == KVQuantMode.FP8_PER_TOKEN_HEAD:
key_cache = key_cache.view(self.fp8_dtype)
value_cache = value_cache.view(self.fp8_dtype)
triton_reshape_and_cache_flash_per_token_head_quant(
@@ -745,9 +761,10 @@ class TritonAttentionImpl(AttentionImpl):
value,
key_cache,
value_cache,
self._k_scale_cache,
self._v_scale_cache,
k_scale_cache,
v_scale_cache,
slot_mapping,
kv_quant_mode=self._kv_quant_mode,
)
return
# For decoder and cross-attention, use KV cache as before.
vllm/v1/attention/ops/int4_per_token_head.py
+1155
View File
File diff suppressed because it is too large Load Diff
vllm/v1/attention/ops/triton_reshape_and_cache_flash.py
+21 -2
View File
@@ -10,6 +10,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
from vllm.platforms import current_platform
from vllm.triton_utils import tl, triton
from vllm.utils.torch_utils import is_quantized_kv_cache
from vllm.v1.kv_cache_interface import KVQuantMode
FP8_MIN, FP8_MAX = get_fp8_min_max()
@@ -276,6 +277,7 @@ def triton_reshape_and_cache_flash_per_token_head_quant(
k_scale_cache: torch.Tensor, # [num_blocks, block_size, num_kv_heads] float32
v_scale_cache: torch.Tensor, # [num_blocks, block_size, num_kv_heads] float32
slot_mapping: torch.Tensor, # [num_tokens]
kv_quant_mode: KVQuantMode,
):
"""Quantize key/value per (token, head) and write to paged cache.
@@ -283,9 +285,26 @@ def triton_reshape_and_cache_flash_per_token_head_quant(
quantized data in key_cache/value_cache, and stores the float32
scale in k_scale_cache/v_scale_cache.
The quantization range (QUANT_MAX, QUANT_MIN) is derived from the
cache tensor dtype so the same code path works for int8 and fp8.
INT4 needs sub-byte packing + a Hadamard rotation, so it is handled by
its own kernel; INT8 / FP8 share this kernel, with the quantization
range (QUANT_MAX, QUANT_MIN) derived from the cache tensor dtype.
"""
if kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD:
from vllm.v1.attention.ops.int4_per_token_head import (
reshape_and_cache_int4,
)
reshape_and_cache_int4(
key,
value,
key_cache,
value_cache,
slot_mapping,
k_scale_cache=k_scale_cache,
v_scale_cache=v_scale_cache,
)
return
cache_dtype = key_cache.dtype
quant_params = _PER_TOKEN_HEAD_QUANT_PARAMS.get(cache_dtype)
if quant_params is None:
vllm/v1/attention/ops/triton_unified_attention.py
+43 -1
View File
@@ -259,7 +259,8 @@ def kernel_unified_attention(
stride_vs_head: tl.int64 = None,
# KV cache quantization mode handled inside this kernel via constexpr
# branches: NONE (0), FP8_PER_TENSOR (1), INT8_PER_TOKEN_HEAD (2),
# FP8_PER_TOKEN_HEAD (3).
# FP8_PER_TOKEN_HEAD (3). Sub-byte INT4 (4) uses its own
# int4_per_token_head kernel, not this one.
KV_QUANT_MODE: tl.constexpr = 0,
FP8_MIN: tl.constexpr = float8_info.min,
FP8_MAX: tl.constexpr = float8_info.max,
@@ -824,6 +825,47 @@ def unified_attention(
use_causal = bool(causal) if not use_per_seq_causal else True
per_seq_causal_ptr = causal if use_per_seq_causal else None
# Sub-byte packed mode (INT4) needs a bespoke kernel (split-dot +
# sub-byte unpack); everything else goes through the core kernel below.
if kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD:
assert use_causal and not use_per_seq_causal, (
"INT4_PER_TOKEN_HEAD only supports causal attention"
)
from vllm.v1.attention.ops.int4_per_token_head import (
unified_attention_int4,
)
if sinks is not None:
assert sinks.shape[0] == q.shape[1], "Sinks must be num_query_heads size"
unified_attention_int4(
q=q,
k_cache=k,
v_cache=v,
out=out,
cu_seqlens_q=cu_seqlens_q,
max_seqlen_q=max_seqlen_q,
seqused_k=seqused_k,
max_seqlen_k=max_seqlen_k,
softmax_scale=softmax_scale,
window_size=window_size,
block_table=block_table,
softcap=softcap,
sinks=sinks,
alibi_slopes=alibi_slopes,
use_alibi_sqrt=use_alibi_sqrt,
qq_bias=qq_bias,
output_scale=output_scale,
mm_prefix_range=mm_prefix_range,
k_scale_cache=k_scale_cache,
v_scale_cache=v_scale_cache,
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,
)
return
if sinks is not None:
assert sinks.shape[0] == q.shape[1], "Sinks must be num_query_heads size"
vllm/v1/kv_cache_interface.py
+21 -21
View File
@@ -41,7 +41,8 @@ class KVQuantMode(IntEnum):
FP8_PER_TENSOR = 1 # per-tensor scales (current fp8 path)
INT8_PER_TOKEN_HEAD = 2 # per-token-head dynamic scales for int8
FP8_PER_TOKEN_HEAD = 3 # per-token-head dynamic scales for fp8
NVFP4 = 4 # packed fp4 data + fp8 block scales
INT4_PER_TOKEN_HEAD = 4 # packed 2×int4/byte, RHT + asymmetric zp
NVFP4 = 5 # packed fp4 data + fp8 block scales
@property
def is_per_token_head(self) -> bool:
@@ -49,6 +50,7 @@ class KVQuantMode(IntEnum):
return self in (
KVQuantMode.INT8_PER_TOKEN_HEAD,
KVQuantMode.FP8_PER_TOKEN_HEAD,
KVQuantMode.INT4_PER_TOKEN_HEAD,
)
@property
@@ -59,6 +61,8 @@ class KVQuantMode(IntEnum):
def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode:
"""Map a ``kv_cache_dtype`` string to a :class:`KVQuantMode`."""
if kv_cache_dtype == "int4_per_token_head":
return KVQuantMode.INT4_PER_TOKEN_HEAD
if kv_cache_dtype == "int8_per_token_head":
return KVQuantMode.INT8_PER_TOKEN_HEAD
if kv_cache_dtype == "fp8_per_token_head":
@@ -184,19 +188,16 @@ class AttentionSpec(KVCacheSpec):
def real_page_size_bytes(self) -> int:
if self.kv_quant_mode.is_nvfp4:
# Packed layout: fp4 data + fp8 block scales per head.
full_dim = nvfp4_kv_cache_full_dim(self.head_size)
return (
2
* self.block_size
* self.num_kv_heads
* full_dim
* get_dtype_size(self.dtype)
)
head_dim = nvfp4_kv_cache_full_dim(self.head_size)
elif self.kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD:
head_dim = self.head_size // 2
else:
head_dim = self.head_size
return (
2
* self.block_size
* self.num_kv_heads
* self.head_size
* head_dim
* get_dtype_size(self.dtype)
)
@@ -314,17 +315,12 @@ class FullAttentionSpec(AttentionSpec):
last_dim = nvfp4_kv_cache_full_dim(
self.head_size
) + nvfp4_kv_cache_full_dim(self.head_size_v)
return (
self.block_size
* self.num_kv_heads
* last_dim
* get_dtype_size(self.dtype)
)
elif self.kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD:
last_dim = self.head_size // 2 + self.head_size_v // 2
else:
last_dim = self.head_size + self.head_size_v
return (
self.block_size
* self.num_kv_heads
* (self.head_size + self.head_size_v)
* get_dtype_size(self.dtype)
self.block_size * self.num_kv_heads * last_dim * get_dtype_size(self.dtype)
)
@@ -390,10 +386,14 @@ class MLAAttentionSpec(FullAttentionSpec):
# V3.2 main MLA: 656-byte custom layout (kv_lora_rank=512 +
# qk_rope_head_dim=64, head_size=576). See flashmla_sparse.py.
return self.block_size * 656
if self.kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD:
head_dim = self.head_size // 2
else:
head_dim = self.head_size
return (
self.storage_block_size
* self.num_kv_heads
* self.head_size
* head_dim
* get_dtype_size(self.dtype)
)