From 70749fdcca7da31be7206101e6f6fc77de9e6839 Mon Sep 17 00:00:00 2001 From: JartX Date: Wed, 24 Jun 2026 12:21:25 +0200 Subject: [PATCH] [Feature] Triton INT4 per-token-head KV cache quantization (#40835) Signed-off-by: JartX Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- docs/design/attention_backends.md | 2 +- .../quantization/test_per_token_kv_cache.py | 3 +- tests/quantization/test_per_token_kv_cache.py | 253 +++- vllm/config/cache.py | 1 + vllm/utils/torch_utils.py | 1 + vllm/v1/attention/backends/triton_attn.py | 49 +- vllm/v1/attention/ops/int4_per_token_head.py | 1155 +++++++++++++++++ .../ops/triton_reshape_and_cache_flash.py | 23 +- .../attention/ops/triton_unified_attention.py | 44 +- vllm/v1/kv_cache_interface.py | 42 +- 10 files changed, 1462 insertions(+), 111 deletions(-) create mode 100644 vllm/v1/attention/ops/int4_per_token_head.py diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index a8d2439f82..f965127cbf 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -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 | diff --git a/tests/models/quantization/test_per_token_kv_cache.py b/tests/models/quantization/test_per_token_kv_cache.py index c581f01eb9..e18f38ffe8 100644 --- a/tests/models/quantization/test_per_token_kv_cache.py +++ b/tests/models/quantization/test_per_token_kv_cache.py @@ -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]) diff --git a/tests/quantization/test_per_token_kv_cache.py b/tests/quantization/test_per_token_kv_cache.py index b657c77a29..3715aaf3a5 100644 --- a/tests/quantization/test_per_token_kv_cache.py +++ b/tests/quantization/test_per_token_kv_cache.py @@ -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) diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 9b96c64513..2fb3358d55 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -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", diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py index 9269fbb44d..cf821a54ba 100644 --- a/vllm/utils/torch_utils.py +++ b/vllm/utils/torch_utils.py @@ -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, diff --git a/vllm/v1/attention/backends/triton_attn.py b/vllm/v1/attention/backends/triton_attn.py index 714c63ae3c..c88456aa1c 100644 --- a/vllm/v1/attention/backends/triton_attn.py +++ b/vllm/v1/attention/backends/triton_attn.py @@ -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. diff --git a/vllm/v1/attention/ops/int4_per_token_head.py b/vllm/v1/attention/ops/int4_per_token_head.py new file mode 100644 index 0000000000..322d32bb81 --- /dev/null +++ b/vllm/v1/attention/ops/int4_per_token_head.py @@ -0,0 +1,1155 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Sub-byte packed (INT4) per-token-head KV cache mode. + +INT4 packs two 4-bit values per cache byte, pre-rotates with a single RHT, +and hides a 4-bit zero-point in the scale's low mantissa bits — too +different from the core kernel to share it. Owns the whole mode: nibble +pack/unpack, the reshape (write) kernel, the split-dot attention (read) +kernel, the RHT transform, and the public ``reshape_and_cache_int4`` / +``unified_attention_int4`` entry points. +""" + +from __future__ import annotations + +from typing import Any + +import torch + +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton +from vllm.v1.attention.ops.triton_attention_helpers import ( + apply_alibi_to_score, + apply_softcap, + cdiv_fn, + compute_kv_seq_mask, + compute_tile_loop_bounds, + init_softmax_M, + load_qq_bias_tile, + resolve_seq_and_query_len, + softmax_step, + store_segm_reduce_scalars, +) +from vllm.v1.attention.ops.triton_unified_attention import reduce_segments + +float8_info = torch.finfo(current_platform.fp8_dtype()) + +# 2 x int4 packed per storage byte. +_INT4_PACKING_FACTOR = 2 + + +# ---------------------------------------------------------------------- +# Nibble pack / unpack (shared write+read format) +# ---------------------------------------------------------------------- + + +@triton.jit +def pack_int4_nibbles(lo, hi): + """Pack two uint8 values (each in [0, 15]) into one byte.""" + return (lo & 0xF) | ((hi & 0xF) << 4) + + +@triton.jit +def unpack_int4_nibbles(packed): + """Split one packed byte into the (low, high) nibble pair as uint8.""" + return packed & 0xF, (packed >> 4) & 0xF + + +# ---------------------------------------------------------------------- +# Write path: RHT + pack + per-(token, head) scale +# ---------------------------------------------------------------------- + + +@triton.jit +def _reshape_cache_int4_kernel( + key_ptr, + value_ptr, + key_cache_ptr, + value_cache_ptr, + k_scale_cache_ptr, + v_scale_cache_ptr, + slot_mapping_ptr, + stride_key_tok: tl.int64, + stride_key_head: tl.int64, + stride_val_tok: tl.int64, + stride_val_head: tl.int64, + stride_kc_blk: tl.int64, + stride_kc_slot: tl.int64, + stride_kc_head: tl.int64, + stride_vc_blk: tl.int64, + stride_vc_slot: tl.int64, + stride_vc_head: tl.int64, + stride_ks_blk: tl.int64, + stride_ks_slot: tl.int64, + stride_ks_head: tl.int64, + stride_vs_blk: tl.int64, + stride_vs_slot: tl.int64, + stride_vs_head: tl.int64, + block_size: tl.constexpr, + head_size: tl.constexpr, + head_size_v: tl.constexpr, + PACKED_HEAD_PADDED: tl.constexpr, +): + """INT4 asymmetric quantization with zero-point steganography.""" + tok = tl.program_id(0) + head = tl.program_id(1) + + slot = tl.load(slot_mapping_ptr + tok).to(tl.int64) + if slot < 0: + return + + blk = slot // block_size + slot_in_blk = slot % block_size + + half_offs = tl.arange(0, PACKED_HEAD_PADDED) + even_offs = half_offs * 2 + odd_offs = half_offs * 2 + 1 + + half_k = head_size // 2 + even_k_mask = even_offs < head_size + odd_k_mask = odd_offs < head_size + key_base = key_ptr + tok * stride_key_tok + head * stride_key_head + + k_even = tl.load(key_base + even_offs, mask=even_k_mask, other=0.0).to(tl.float32) + k_odd = tl.load(key_base + odd_offs, mask=odd_k_mask, other=0.0).to(tl.float32) + + k_min = tl.minimum( + tl.min(tl.where(even_k_mask, k_even, float("inf"))), + tl.min(tl.where(odd_k_mask, k_odd, float("inf"))), + ) + k_max = tl.maximum( + tl.max(tl.where(even_k_mask, k_even, float("-inf"))), + tl.max(tl.where(odd_k_mask, k_odd, float("-inf"))), + ) + k_scale = tl.maximum((k_max - k_min) / 15.0, 1e-6) + k_zp_f = tl.clamp( + tl.where( + -k_min / k_scale >= 0, + (-k_min / k_scale + 0.5).to(tl.int32), + (-k_min / k_scale - 0.5).to(tl.int32), + ).to(tl.float32), + 0.0, + 15.0, + ) + + inv_k = 1.0 / k_scale + k_even_s = k_even * inv_k + k_zp_f + k_odd_s = k_odd * inv_k + k_zp_f + k_even_q = tl.clamp( + tl.where( + k_even_s >= 0, + (k_even_s + 0.5).to(tl.int32), + (k_even_s - 0.5).to(tl.int32), + ).to(tl.float32), + 0.0, + 15.0, + ) + k_odd_q = tl.clamp( + tl.where( + k_odd_s >= 0, + (k_odd_s + 0.5).to(tl.int32), + (k_odd_s - 0.5).to(tl.int32), + ).to(tl.float32), + 0.0, + 15.0, + ) + + k_zp_int = k_zp_f.to(tl.int32) + k_scale_bits = k_scale.to(tl.int32, bitcast=True) + k_scale_packed = ((k_scale_bits & -16) | (k_zp_int & 0xF)).to( + tl.float32, bitcast=True + ) + + tl.store( + k_scale_cache_ptr + + blk * stride_ks_blk + + slot_in_blk * stride_ks_slot + + head * stride_ks_head, + k_scale_packed, + ) + + k_packed = pack_int4_nibbles(k_even_q.to(tl.uint8), k_odd_q.to(tl.uint8)) + tl.store( + key_cache_ptr + + blk * stride_kc_blk + + slot_in_blk * stride_kc_slot + + head * stride_kc_head + + half_offs, + k_packed, + mask=half_offs < half_k, + ) + + half_v = head_size_v // 2 + even_v_mask = even_offs < head_size_v + odd_v_mask = odd_offs < head_size_v + val_base = value_ptr + tok * stride_val_tok + head * stride_val_head + + v_even = tl.load(val_base + even_offs, mask=even_v_mask, other=0.0).to(tl.float32) + v_odd = tl.load(val_base + odd_offs, mask=odd_v_mask, other=0.0).to(tl.float32) + + v_min = tl.minimum( + tl.min(tl.where(even_v_mask, v_even, float("inf"))), + tl.min(tl.where(odd_v_mask, v_odd, float("inf"))), + ) + v_max = tl.maximum( + tl.max(tl.where(even_v_mask, v_even, float("-inf"))), + tl.max(tl.where(odd_v_mask, v_odd, float("-inf"))), + ) + v_scale = tl.maximum((v_max - v_min) / 15.0, 1e-6) + v_zp_f = tl.clamp( + tl.where( + -v_min / v_scale >= 0, + (-v_min / v_scale + 0.5).to(tl.int32), + (-v_min / v_scale - 0.5).to(tl.int32), + ).to(tl.float32), + 0.0, + 15.0, + ) + + inv_v = 1.0 / v_scale + v_even_s = v_even * inv_v + v_zp_f + v_odd_s = v_odd * inv_v + v_zp_f + v_even_q = tl.clamp( + tl.where( + v_even_s >= 0, + (v_even_s + 0.5).to(tl.int32), + (v_even_s - 0.5).to(tl.int32), + ).to(tl.float32), + 0.0, + 15.0, + ) + v_odd_q = tl.clamp( + tl.where( + v_odd_s >= 0, + (v_odd_s + 0.5).to(tl.int32), + (v_odd_s - 0.5).to(tl.int32), + ).to(tl.float32), + 0.0, + 15.0, + ) + + v_zp_int = v_zp_f.to(tl.int32) + v_scale_bits = v_scale.to(tl.int32, bitcast=True) + v_scale_packed = ((v_scale_bits & -16) | (v_zp_int & 0xF)).to( + tl.float32, bitcast=True + ) + + tl.store( + v_scale_cache_ptr + + blk * stride_vs_blk + + slot_in_blk * stride_vs_slot + + head * stride_vs_head, + v_scale_packed, + ) + + v_packed = pack_int4_nibbles(v_even_q.to(tl.uint8), v_odd_q.to(tl.uint8)) + tl.store( + value_cache_ptr + + blk * stride_vc_blk + + slot_in_blk * stride_vc_slot + + head * stride_vc_head + + half_offs, + v_packed, + mask=half_offs < half_v, + ) + + +def _run_reshape_kernel( + kernel, + *, + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + k_scale_cache: torch.Tensor, + v_scale_cache: torch.Tensor, + slot_mapping: torch.Tensor, + packing_factor: int, +) -> None: + """Launch the packed INT4 reshape kernel.""" + num_tokens, num_kv_heads, head_size = key.shape + head_size_v = value.shape[2] + assert head_size % packing_factor == 0 and head_size_v % packing_factor == 0 + packed_padded = triton.next_power_of_2( + max(head_size, head_size_v) // packing_factor + ) + if current_platform.is_rocm() or current_platform.is_xpu(): + num_warps = 4 + else: + num_warps = min(16, max(1, packed_padded // 32)) + + kernel[(num_tokens, num_kv_heads)]( + key_ptr=key, + value_ptr=value, + key_cache_ptr=key_cache, + value_cache_ptr=value_cache, + k_scale_cache_ptr=k_scale_cache, + v_scale_cache_ptr=v_scale_cache, + slot_mapping_ptr=slot_mapping, + stride_key_tok=key.stride(0), + stride_key_head=key.stride(1), + stride_val_tok=value.stride(0), + stride_val_head=value.stride(1), + stride_kc_blk=key_cache.stride(0), + stride_kc_slot=key_cache.stride(1), + stride_kc_head=key_cache.stride(2), + stride_vc_blk=value_cache.stride(0), + stride_vc_slot=value_cache.stride(1), + stride_vc_head=value_cache.stride(2), + stride_ks_blk=k_scale_cache.stride(0), + stride_ks_slot=k_scale_cache.stride(1), + stride_ks_head=k_scale_cache.stride(2), + stride_vs_blk=v_scale_cache.stride(0), + stride_vs_slot=v_scale_cache.stride(1), + stride_vs_head=v_scale_cache.stride(2), + block_size=key_cache.shape[1], + head_size=head_size, + head_size_v=head_size_v, + PACKED_HEAD_PADDED=packed_padded, + num_warps=num_warps, + ) + + +# ---------------------------------------------------------------------- +# Read path: split-dot attention over the packed cache +# ---------------------------------------------------------------------- + + +@triton.jit +def _attn_packed( + # Output destinations. In 2D mode the final result is written into + # ``output_ptr``; in 3D mode per-segment partials go into the three + # ``segm_*`` tensors and ``output_ptr`` is unused. + output_ptr, + segm_output_ptr, + segm_max_ptr, + segm_expsum_ptr, + query_ptr, + key_cache_ptr, + value_cache_ptr, + sink_ptr, + block_tables_ptr, + seq_lens_ptr, + alibi_slopes_ptr, + qq_bias_ptr, + scale, + out_scale, + softcap, + k_scale_cache_ptr, + v_scale_cache_ptr, + num_query_heads: tl.constexpr, + num_queries_per_kv: tl.constexpr, + block_table_stride: tl.int64, + query_stride_0: tl.int64, + query_stride_1: tl.int64, + output_stride_0: tl.int64, + output_stride_1: tl.int64, + qq_bias_stride_0: tl.int64, + BLOCK_SIZE: tl.constexpr, + TILE_SIZE: tl.constexpr, + HEAD_SIZE: tl.constexpr, + HEAD_SIZE_PADDED: tl.constexpr, + PACKED_HEAD_PADDED: tl.constexpr, # HEAD_SIZE / PACKING_FACTOR, rounded up + USE_ALIBI_SLOPES: tl.constexpr, + USE_ALIBI_SQRT: tl.constexpr, + USE_QQ_BIAS: tl.constexpr, + USE_SOFTCAP: tl.constexpr, + USE_SINKS: tl.constexpr, + SLIDING_WINDOW: tl.constexpr, + USE_MM_PREFIX: tl.constexpr, + MAX_MM_RANGES: tl.constexpr, + mm_prefix_range_ptr, + stride_k_cache_0: tl.int64, + stride_k_cache_1: tl.int64, + stride_k_cache_2: tl.int64, + stride_k_cache_3: tl.constexpr, + stride_v_cache_0: tl.int64, + stride_v_cache_1: tl.int64, + stride_v_cache_2: tl.int64, + stride_v_cache_3: tl.constexpr, + stride_ks_blk: tl.int64, + stride_ks_slot: tl.int64, + stride_ks_head: tl.int64, + stride_vs_blk: tl.int64, + stride_vs_slot: tl.int64, + stride_vs_head: tl.int64, + query_start_len_ptr, + BLOCK_Q: tl.constexpr, + num_seqs: tl.int32, + BLOCK_M: tl.constexpr, + NUM_SEGMENTS_PER_SEQ: tl.constexpr, + USE_FP8: tl.constexpr, + IS_3D: tl.constexpr, + # 2 → INT4 nibble pair (asymmetric + zp). The packed KV cache stores + # one byte per ``packed_offs``, holding PACKING_FACTOR values. + PACKING_FACTOR: tl.constexpr, + FP8_MIN: tl.constexpr = float8_info.min, + FP8_MAX: tl.constexpr = float8_info.max, +): + # Shared prologue: sequence lookup, q-block bounds, early returns. + q_block_global_idx = tl.program_id(0) + kv_head_idx = tl.program_id(1) + segm_idx = tl.program_id(2) if IS_3D else 0 + + ( + seq_idx, + q_block_local_idx, + cur_batch_in_all_start_index, + cur_batch_query_len, + seq_len, + ) = resolve_seq_and_query_len( + query_start_len_ptr, seq_lens_ptr, q_block_global_idx, num_seqs, BLOCK_Q + ) + + if q_block_local_idx * BLOCK_Q >= cur_batch_query_len: + return + + if IS_3D: + tiles_per_segment = cdiv_fn(seq_len, NUM_SEGMENTS_PER_SEQ * TILE_SIZE) + if segm_idx * tiles_per_segment * TILE_SIZE >= seq_len: + return + else: + tiles_per_segment = 0 + + offs_m = tl.arange(0, BLOCK_M) + offs_t = tl.arange(0, TILE_SIZE) + query_pos = q_block_local_idx * BLOCK_Q + offs_m // num_queries_per_kv + + query_offset_0 = cur_batch_in_all_start_index + query_pos + query_offset_1 = kv_head_idx * num_queries_per_kv + offs_m % num_queries_per_kv + + query_mask_0 = tl.where(query_pos < cur_batch_query_len, 1, 0).to(tl.int1) + query_mask_1 = tl.where(query_offset_1 < num_query_heads, 1, 0).to(tl.int1) + + # Split-Q prologue: PACKING_FACTOR interleaved streams of Q. + # INT4 uses 2 streams (even / odd). The packed KV cache stores one + # byte per ``packed_offs``, which holds PACKING_FACTOR values. + packed_offs = tl.arange(0, PACKED_HEAD_PADDED) + offs_s0 = packed_offs * PACKING_FACTOR + offs_s1 = packed_offs * PACKING_FACTOR + 1 + mask_s0 = tl.where(offs_s0 < HEAD_SIZE, 1, 0).to(tl.int1) + mask_s1 = tl.where(offs_s1 < HEAD_SIZE, 1, 0).to(tl.int1) + packed_dim_mask = tl.where(packed_offs < HEAD_SIZE // PACKING_FACTOR, 1, 0).to( + tl.int1 + ) + q_base = ( + query_offset_0[:, None] * query_stride_0 + + query_offset_1[:, None] * query_stride_1 + ) + q_mask = query_mask_0[:, None] & query_mask_1[:, None] + Q_s0 = tl.load( + query_ptr + q_base + offs_s0[None, :], + mask=mask_s0[None, :] & q_mask, + other=0.0, + ).to(tl.float32) + Q_s1 = tl.load( + query_ptr + q_base + offs_s1[None, :], + mask=mask_s1[None, :] & q_mask, + other=0.0, + ).to(tl.float32) + + # INT4 asymmetric correction needs sum(Q) per row. + Q_sum = tl.sum(Q_s0, axis=1) + tl.sum(Q_s1, axis=1) + + block_table_offset = seq_idx * block_table_stride + + # Online-softmax state + optional feature loads. + M = init_softmax_M( + sink_ptr, query_offset_1, query_mask_1, segm_idx, BLOCK_M, USE_SINKS, IS_3D + ) + L = tl.full([BLOCK_M], 1.0, dtype=tl.float32) + acc_s0 = tl.zeros([BLOCK_M, PACKED_HEAD_PADDED], dtype=tl.float32) + acc_s1 = tl.zeros([BLOCK_M, PACKED_HEAD_PADDED], dtype=tl.float32) + + context_len = seq_len - cur_batch_query_len + + if USE_ALIBI_SLOPES: + alibi_slope = tl.load( + alibi_slopes_ptr + query_offset_1, mask=query_mask_1, other=0.0 + ) + + if USE_QQ_BIAS: + qq_bias_row_ptrs = qq_bias_ptr + query_pos[:, None] * qq_bias_stride_0 + + loop_lo, loop_hi, max_seq_prefix_len = compute_tile_loop_bounds( + context_len, + seq_len, + cur_batch_query_len, + q_block_local_idx, + segm_idx, + tiles_per_segment, + TILE_SIZE, + BLOCK_M, + BLOCK_Q, + num_queries_per_kv, + SLIDING_WINDOW, + USE_MM_PREFIX, + IS_3D, + ) + + # Tile loop. Per-tile: load packed KV + scales, dequantize into + # PACKING_FACTOR streams, compute the split dot, run the shared + # softmax step, accumulate per stream. + for j in range(loop_lo, loop_hi): + seq_offset = j * TILE_SIZE + offs_t + tile_mask = seq_offset < max_seq_prefix_len + + physical_block_idx = tl.load( + block_tables_ptr + block_table_offset + seq_offset // BLOCK_SIZE + ).to(tl.int64) + + slot_in_blk = seq_offset % BLOCK_SIZE + k_off = ( + physical_block_idx[None, :] * stride_k_cache_0 + + kv_head_idx * stride_k_cache_2 + + packed_offs[:, None] * stride_k_cache_3 + + slot_in_blk[None, :] * stride_k_cache_1 + ) + K_packed = tl.load( + key_cache_ptr + k_off, + mask=packed_dim_mask[:, None] & tile_mask[None, :], + other=0, + ) + v_off = ( + physical_block_idx[:, None] * stride_v_cache_0 + + kv_head_idx * stride_v_cache_2 + + packed_offs[None, :] * stride_v_cache_3 + + slot_in_blk[:, None] * stride_v_cache_1 + ) + V_packed = tl.load( + value_cache_ptr + v_off, + mask=packed_dim_mask[None, :] & tile_mask[:, None], + other=0, + ) + # Dequantize KV. INT4 unpacks nibbles as plain uint [0..15]; + # the zero-point is applied on the score side. + K_s0_u, K_s1_u = unpack_int4_nibbles(K_packed) + K_s0 = K_s0_u.to(tl.float32) + K_s1 = K_s1_u.to(tl.float32) + V_s0_u, V_s1_u = unpack_int4_nibbles(V_packed) + V_s0 = V_s0_u.to(tl.float32) + V_s1 = V_s1_u.to(tl.float32) + + ks_idx = ( + physical_block_idx * stride_ks_blk + + slot_in_blk * stride_ks_slot + + kv_head_idx * stride_ks_head + ) + ks_raw = tl.load(k_scale_cache_ptr + ks_idx, mask=tile_mask, other=0) + vs_idx = ( + physical_block_idx * stride_vs_blk + + slot_in_blk * stride_vs_slot + + kv_head_idx * stride_vs_head + ) + vs_raw = tl.load(v_scale_cache_ptr + vs_idx, mask=tile_mask, other=0) + + # INT4 steganographs the 4-bit zero-point in the low 4 bits of + # the float32 scale's mantissa. + ks_bits = ks_raw.to(tl.int32, bitcast=True) + k_zp = (ks_bits & 0xF).to(tl.float32) + k_token_head_scales = (ks_bits & -16).to(tl.float32, bitcast=True) + vs_bits = vs_raw.to(tl.int32, bitcast=True) + v_zp = (vs_bits & 0xF).to(tl.float32) + v_token_head_scales = (vs_bits & -16).to(tl.float32, bitcast=True) + + query_abs_pos = context_len + query_pos[:, None] + seq_mask = compute_kv_seq_mask( + query_abs_pos, + seq_offset, + seq_idx, + seq_len, + mm_prefix_range_ptr, + SLIDING_WINDOW, + USE_MM_PREFIX, + MAX_MM_RANGES, + ) + + # Score: split-dot across the 2 INT4 streams; fused + # softmax_scale * per-(token, head) k_scale in one mul. INT4 + # subtracts the ``zp * sum(Q)`` correction term. + S = tl.zeros(shape=(BLOCK_M, TILE_SIZE), dtype=tl.float32) + raw_dot = tl.dot(Q_s0, K_s0) + tl.dot(Q_s1, K_s1) + S += (raw_dot - Q_sum[:, None] * k_zp[None, :]) * ( + scale * k_token_head_scales[None, :] + ) + + if USE_SOFTCAP: + S = apply_softcap(S, softcap) + + S = tl.where( + query_mask_1[:, None] & query_mask_0[:, None] & seq_mask, S, float("-inf") + ) + + if USE_ALIBI_SLOPES: + S = apply_alibi_to_score( + S, alibi_slope, seq_offset, context_len, query_pos, USE_ALIBI_SQRT + ) + + if USE_QQ_BIAS: + S += load_qq_bias_tile( + qq_bias_row_ptrs, seq_offset, context_len, qq_bias_stride_0 + ) + + M, L, P, alpha = softmax_step(S, M, L) + acc_s0 = acc_s0 * alpha[:, None] + acc_s1 = acc_s1 * alpha[:, None] + + if SLIDING_WINDOW: + qpos_lo = q_block_local_idx * BLOCK_Q + sw_mask = (context_len + qpos_lo - seq_offset) < SLIDING_WINDOW + V_s0 = tl.where(sw_mask[:, None], V_s0, 0.0) + V_s1 = tl.where(sw_mask[:, None], V_s1, 0.0) + + # Fuse v per-(token, head) scale into P. INT4 also subtracts + # the v-zero-point contribution from each stream once. + P_v = (P * v_token_head_scales[None, :]).to(tl.float32) + Pv_zp_sum = tl.sum(P_v * v_zp[None, :], axis=1) + acc_s0 += tl.dot(P_v, V_s0) - Pv_zp_sum[:, None] + acc_s1 += tl.dot(P_v, V_s1) - Pv_zp_sum[:, None] + + # Epilogue. 2D writes the final output with optional FP8 clamp; + # 3D writes the per-segment partials (output / max / expsum) for + # ``reduce_segments`` to finalize. Each stream writes its own + # stripe in the output layout. + out_mask = query_mask_0[:, None] & query_mask_1[:, None] + if IS_3D: + segm_base = ( + query_offset_0[:, None].to(tl.int64) + * (num_query_heads * NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED) + + query_offset_1[:, None] * (NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED) + + segm_idx * HEAD_SIZE_PADDED + ) + tl.store( + segm_output_ptr + segm_base + offs_s0[None, :], + acc_s0, + mask=mask_s0[None, :] & out_mask, + ) + tl.store( + segm_output_ptr + segm_base + offs_s1[None, :], + acc_s1, + mask=mask_s1[None, :] & out_mask, + ) + store_segm_reduce_scalars( + segm_max_ptr, + segm_expsum_ptr, + query_offset_0, + query_offset_1, + segm_idx, + M, + L, + query_mask_0, + query_mask_1, + num_query_heads, + NUM_SEGMENTS_PER_SEQ, + ) + else: + acc_s0 = acc_s0 / L[:, None] + acc_s1 = acc_s1 / L[:, None] + if USE_FP8: + out_s = tl.load(out_scale) + acc_s0 = tl.clamp(acc_s0 * out_s, FP8_MIN, FP8_MAX) + acc_s1 = tl.clamp(acc_s1 * out_s, FP8_MIN, FP8_MAX) + out_base = ( + query_offset_0[:, None] * output_stride_0 + + query_offset_1[:, None] * output_stride_1 + ) + tl.store( + output_ptr + out_base + offs_s0[None, :], + acc_s0, + mask=mask_s0[None, :] & out_mask, + ) + tl.store( + output_ptr + out_base + offs_s1[None, :], + acc_s1, + mask=mask_s1[None, :] & out_mask, + ) + + +def _launch_packed_attn( + *, + q, + k_cache, + v_cache, + out, + cu_seqlens_q, + max_seqlen_q, + seqused_k, + softmax_scale, + window_size, + block_table, + softcap, + sinks, + alibi_slopes, + use_alibi_sqrt, + qq_bias, + output_scale, + mm_prefix_range, + k_scale_cache, + v_scale_cache, + seq_threshold_3D, + num_par_softmax_segments, + softmax_segm_output, + softmax_segm_max, + softmax_segm_expsum, + packing_factor: int, +): + """Launch ``_attn_packed`` for one of the sub-byte modes. + + Handles 2D-vs-3D dispatch, placeholder pointers for the unused side + of that split, and the trailing ``reduce_segments`` pass. Writes + into ``out`` (directly for 2D; via the segm buffers for 3D). + """ + import vllm.envs as envs + from vllm.v1.attention.ops.triton_unified_attention import _get_tile_size + + is_batch_invariant = envs.VLLM_BATCH_INVARIANT + + use_mm_prefix = False + max_mm_ranges = 0 + if mm_prefix_range is not None: + assert mm_prefix_range.ndim == 3, ( + f"Unsupported mm_prefix_range shape: {mm_prefix_range.shape}" + ) + use_mm_prefix = True + max_mm_ranges = mm_prefix_range.shape[1] + + block_size = v_cache.shape[1] + num_seqs = len(seqused_k) + num_query_heads = q.shape[1] + num_kv_heads = k_cache.shape[2] + num_queries_per_kv = num_query_heads // num_kv_heads + head_size = q.shape[2] + + BLOCK_M = ( + 16 if num_queries_per_kv <= 16 else triton.next_power_of_2(num_queries_per_kv) + ) + BLOCK_Q = BLOCK_M // num_queries_per_kv + total_num_q_blocks = q.shape[0] // BLOCK_Q + num_seqs + sliding_window_val = 1 + window_size[0] if window_size[0] >= 0 else 0 + TILE_SIZE_PREFILL = _get_tile_size( + head_size, sliding_window_val, q.element_size(), is_prefill=True + ) + TILE_SIZE_DECODE = _get_tile_size( + head_size, sliding_window_val, q.element_size(), is_prefill=False + ) + + use_3d = not ( + seq_threshold_3D is None + or num_par_softmax_segments is None + or softmax_segm_output is None + or softmax_segm_max is None + or softmax_segm_expsum is None + or max_seqlen_q > 1 + or num_seqs > seq_threshold_3D + or is_batch_invariant + ) + + # 3D never reads ``output_ptr`` and 2D never reads the segm tensors, + # but Triton needs a non-null pointer everywhere; reuse ``out`` as + # the placeholder for the unused side. + segm_output_ptr = softmax_segm_output if use_3d else out + segm_max_ptr = softmax_segm_max if use_3d else out + segm_expsum_ptr = softmax_segm_expsum if use_3d else out + num_segments = num_par_softmax_segments if use_3d else 1 + + grid: tuple[Any, ...] + if use_3d: + grid = (total_num_q_blocks, num_kv_heads, num_par_softmax_segments) + tile_size = TILE_SIZE_DECODE + else: + grid = (total_num_q_blocks, num_kv_heads) + tile_size = TILE_SIZE_PREFILL + + _attn_packed[grid]( + output_ptr=out, + segm_output_ptr=segm_output_ptr, + segm_max_ptr=segm_max_ptr, + segm_expsum_ptr=segm_expsum_ptr, + query_ptr=q, + key_cache_ptr=k_cache, + value_cache_ptr=v_cache, + sink_ptr=sinks, + block_tables_ptr=block_table, + seq_lens_ptr=seqused_k, + alibi_slopes_ptr=alibi_slopes, + qq_bias_ptr=qq_bias, + scale=softmax_scale, + out_scale=1 / output_scale if output_scale is not None else 1.0, + softcap=softcap, + k_scale_cache_ptr=k_scale_cache, + v_scale_cache_ptr=v_scale_cache, + num_query_heads=num_query_heads, + num_queries_per_kv=num_queries_per_kv, + block_table_stride=block_table.stride(0), + query_stride_0=q.stride(0), + query_stride_1=q.stride(1), + output_stride_0=out.stride(0), + output_stride_1=out.stride(1), + qq_bias_stride_0=qq_bias.stride(0) if qq_bias is not None else 0, + BLOCK_SIZE=block_size, + TILE_SIZE=tile_size, + HEAD_SIZE=head_size, + HEAD_SIZE_PADDED=triton.next_power_of_2(head_size), + PACKED_HEAD_PADDED=triton.next_power_of_2(head_size) // packing_factor, + USE_ALIBI_SLOPES=alibi_slopes is not None, + USE_ALIBI_SQRT=use_alibi_sqrt, + USE_QQ_BIAS=qq_bias is not None, + USE_SOFTCAP=(softcap > 0), + USE_SINKS=(sinks is not None), + SLIDING_WINDOW=(1 + window_size[0]), + USE_MM_PREFIX=use_mm_prefix, + MAX_MM_RANGES=max_mm_ranges, + mm_prefix_range_ptr=mm_prefix_range, + stride_k_cache_0=k_cache.stride(0), + stride_k_cache_1=k_cache.stride(1), + stride_k_cache_2=k_cache.stride(2), + stride_k_cache_3=k_cache.stride(3), + stride_v_cache_0=v_cache.stride(0), + stride_v_cache_1=v_cache.stride(1), + stride_v_cache_2=v_cache.stride(2), + stride_v_cache_3=v_cache.stride(3), + stride_ks_blk=k_scale_cache.stride(0), + stride_ks_slot=k_scale_cache.stride(1), + stride_ks_head=k_scale_cache.stride(2), + stride_vs_blk=v_scale_cache.stride(0), + stride_vs_slot=v_scale_cache.stride(1), + stride_vs_head=v_scale_cache.stride(2), + query_start_len_ptr=cu_seqlens_q, + BLOCK_Q=BLOCK_Q, + num_seqs=num_seqs, + BLOCK_M=BLOCK_M, + NUM_SEGMENTS_PER_SEQ=num_segments, + USE_FP8=output_scale is not None, + IS_3D=use_3d, + PACKING_FACTOR=packing_factor, + ) + + if use_3d: + reduce_segments[(q.shape[0], num_query_heads)]( + output_ptr=out, + segm_output_ptr=softmax_segm_output, + segm_max_ptr=softmax_segm_max, + segm_expsum_ptr=softmax_segm_expsum, + seq_lens_ptr=seqused_k, + num_seqs=num_seqs, + num_query_heads=num_query_heads, + out_scale_inv=1 / output_scale if output_scale is not None else 1.0, + output_stride_0=out.stride(0), + output_stride_1=out.stride(1), + block_table_stride=block_table.stride(0), + TILE_SIZE=TILE_SIZE_DECODE, + HEAD_SIZE=head_size, + HEAD_SIZE_PADDED=triton.next_power_of_2(head_size), + query_start_len_ptr=cu_seqlens_q, + BLOCK_Q=BLOCK_Q, + NUM_SEGMENTS_PER_SEQ=num_par_softmax_segments, + USE_FP8=output_scale is not None, + ) + + +# ---------------------------------------------------------------------- +# Public entry points +# ---------------------------------------------------------------------- + + +def reshape_and_cache_int4( + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, + *, + k_scale_cache: torch.Tensor, + v_scale_cache: torch.Tensor, +) -> None: + """Pre-rotate (RHT), pack to INT4 and write into the paged cache.""" + key = single_rht(key.float()).to(key.dtype) + value = single_rht(value.float()).to(value.dtype) + _run_reshape_kernel( + _reshape_cache_int4_kernel, + key=key, + value=value, + key_cache=key_cache, + value_cache=value_cache, + k_scale_cache=k_scale_cache, + v_scale_cache=v_scale_cache, + slot_mapping=slot_mapping, + packing_factor=_INT4_PACKING_FACTOR, + ) + + +def unified_attention_int4( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + out: torch.Tensor, + *, + cu_seqlens_q: torch.Tensor, + max_seqlen_q: int, + seqused_k: torch.Tensor, + max_seqlen_k: int, + softmax_scale: float, + window_size: tuple[int, int], + block_table: torch.Tensor, + softcap: float, + sinks: torch.Tensor | None, + alibi_slopes: torch.Tensor | None, + use_alibi_sqrt: bool, + qq_bias: torch.Tensor | None, + output_scale: torch.Tensor | None, + mm_prefix_range: torch.Tensor | None, + k_scale_cache: torch.Tensor, + v_scale_cache: torch.Tensor, + seq_threshold_3D: int | None = None, + num_par_softmax_segments: int | None = None, + softmax_segm_output: torch.Tensor | None = None, + softmax_segm_max: torch.Tensor | None = None, + softmax_segm_expsum: torch.Tensor | None = None, +) -> None: + """Paged attention over the INT4 packed cache, writing into *out*. + + The forward RHT has norm ``sqrt(head_size)``, so ``softmax_scale`` is + divided by ``head_size`` and the inverse RHT divides the output by + ``head_size`` as well. + """ + q_orig_dtype = q.dtype + q = single_rht(q.float()).to(q_orig_dtype) + head_size = q.shape[2] + softmax_scale = softmax_scale / head_size + + _launch_packed_attn( + q=q, + k_cache=k_cache, + v_cache=v_cache, + out=out, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=max_seqlen_q, + seqused_k=seqused_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, + packing_factor=_INT4_PACKING_FACTOR, + ) + + out_f = single_rht(out.float(), inverse=True) / head_size + out.copy_(out_f.to(q_orig_dtype)) + + +# ---------------------------------------------------------------------- +# Randomized Hadamard Transform (RHT) — gaussianizes K/V before INT4 +# quantization; applied on write and to Q before the read kernels. +# ---------------------------------------------------------------------- + +# Hadacore (CUDA tensor core kernel) availability check +# Hadacore's CUDA impl is only registered when built for sm_80+, but the +# schema def is unconditional — on ROCm ``hasattr`` is True yet dispatch +# would crash, so we also gate on ``is_cuda()`` and the sm_80 capability. +_HADACORE_AVAILABLE: bool | None = None + + +def _hadacore_available() -> bool: + global _HADACORE_AVAILABLE + if _HADACORE_AVAILABLE is None: + _HADACORE_AVAILABLE = ( + current_platform.is_cuda() + and current_platform.has_device_capability(80) + and hasattr(torch.ops._C, "hadacore_transform") + ) + return _HADACORE_AVAILABLE + + +# Cached Hadamard matrices (one per (size, dtype, device) tuple) +_HADAMARD_MATRIX_CACHE: dict[tuple[int, torch.dtype, str], torch.Tensor] = {} + + +def _get_hadamard_matrix( + d: int, dtype: torch.dtype, device: torch.device +) -> torch.Tensor: + key = (d, dtype, str(device)) + cached = _HADAMARD_MATRIX_CACHE.get(key) + if cached is None: + H = torch.ones(1, 1, dtype=torch.float32, device=device) + while H.shape[0] < d: + H = torch.cat( + [ + torch.cat([H, H], dim=1), + torch.cat([H, -H], dim=1), + ], + dim=0, + ) + cached = H.to(dtype).contiguous() + _HADAMARD_MATRIX_CACHE[key] = cached + return cached + + +# Triton MMA Hadamard kernel (Tier 2) +@triton.jit +def _hadamard_mma_kernel( + x_ptr, + h_ptr, + out_ptr, + n_rows, + stride_x_row: tl.int64, + stride_x_col: tl.int64, + stride_o_row: tl.int64, + stride_o_col: tl.int64, + BLOCK_M: tl.constexpr, + D: tl.constexpr, +): + pid = tl.program_id(0) + rows = pid * BLOCK_M + tl.arange(0, BLOCK_M) + cols = tl.arange(0, D) + row_mask = rows < n_rows + + x = tl.load( + x_ptr + rows[:, None] * stride_x_row + cols[None, :] * stride_x_col, + mask=row_mask[:, None], + other=0.0, + ) + H = tl.load(h_ptr + cols[:, None] * D + cols[None, :]) + + out = tl.dot(x, H, out_dtype=tl.float32).to(x.dtype) + + tl.store( + out_ptr + rows[:, None] * stride_o_row + cols[None, :] * stride_o_col, + out, + mask=row_mask[:, None], + ) + + +# H is D×D bf16 = 2·D² bytes of LDS. AMD CDNA has 64 KiB LDS, so D ≤ 128 +# (32 KiB) leaves room for input + accumulator. Larger D falls back. +_TRITON_HADAMARD_MIN_D = 16 +_TRITON_HADAMARD_MAX_D = 128 + + +def _triton_hadamard_transform(x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] + orig_shape = x.shape + orig_dtype = x.dtype + + work_dtype = torch.bfloat16 if orig_dtype == torch.float32 else orig_dtype + x2d = x.contiguous().to(work_dtype).reshape(-1, d) + out2d = torch.empty_like(x2d) + n_rows = x2d.shape[0] + H_mat = _get_hadamard_matrix(d, work_dtype, x.device) + + BLOCK_M = 16 + grid = (triton.cdiv(n_rows, BLOCK_M),) + # num_stages=1: the kernel has no loop, so default 3-stage pipelining + # would triple-buffer H and blow the AMD LDS budget. + _hadamard_mma_kernel[grid]( + x2d, + H_mat, + out2d, + n_rows, + x2d.stride(0), + x2d.stride(1), + out2d.stride(0), + out2d.stride(1), + BLOCK_M=BLOCK_M, + D=d, + num_stages=1, + num_warps=4, + ) + return out2d.reshape(orig_shape).to(orig_dtype) + + +# Public API +def fast_hadamard_transform(x: torch.Tensor) -> torch.Tensor: + """Unnormalized Walsh-Hadamard Transform along the last dimension. + + H_d × x where H_d × H_d = d × I. Last dim must be a power of 2. + + Three-tier dispatch: + 1. Hadacore CUDA Tensor Core kernel (sm_80+). + 2. Triton MMA matmul kernel (CUDA fallback + ROCm MFMA/WMMA path). + 3. PyTorch butterfly (CPU and any GPU/dtype combo Triton can't take). + """ + d = x.shape[-1] + assert d & (d - 1) == 0, f"Requires power-of-2 dim, got {d}" + + # Tier 1 — hadacore on CUDA. + if _hadacore_available() and 0 < d <= (1 << 15): + from vllm import _custom_ops as ops + + # hadacore returns x @ (H/√d); rescale to the unnormalized H × x + # convention the INT4 scale math is calibrated to. + rescale = d**0.5 + if x.dtype in (torch.float16, torch.bfloat16): + y = ops.hadacore_transform(x.contiguous().clone(), inplace=True) + return y * rescale + # fp32 → bf16 round-trip; precision loss is irrelevant before + # INT4 quantization. + orig_dtype = x.dtype + x_bf16 = x.contiguous().to(torch.bfloat16) + y_bf16 = ops.hadacore_transform(x_bf16, inplace=True) + return y_bf16.to(orig_dtype) * rescale + + # Tier 2 — Triton MMA kernel (covers ROCm via MFMA/WMMA codegen, and + # also CUDA when hadacore is unavailable). + if ( + x.is_cuda + and _TRITON_HADAMARD_MIN_D <= d <= _TRITON_HADAMARD_MAX_D + and x.dtype in (torch.float16, torch.bfloat16, torch.float32) + ): + return _triton_hadamard_transform(x) + + # Tier 3 — PyTorch butterfly (CPU / unsupported dtype / D < 16). + h = 1 + while h < d: + xv = x.view(*x.shape[:-1], d // (2 * h), 2, h) + a = xv[..., 0, :] + b = xv[..., 1, :] + x = torch.stack([a + b, a - b], dim=-2).reshape(x.shape) + h <<= 1 + return x + + +# Randomized Hadamard Transform (used by INT4) +# Deterministic ±1 signs for Randomized Hadamard Transform. +# RHT = H × D × x (sign flip + Hadamard). Breaks residual structure +# in KV vectors, improving quantization quality. +_RHT_SIGNS_CACHE: dict[tuple[int, int, str], torch.Tensor] = {} + + +def _get_rht_signs(d: int, round_idx: int, device: torch.device) -> torch.Tensor: + """Return a cached deterministic ±1 sign vector of length *d*.""" + key = (d, round_idx, str(device)) + if key not in _RHT_SIGNS_CACHE: + gen = torch.Generator(device="cpu") + gen.manual_seed(0x9E3779B9 + round_idx * 0x517CC1B7) + signs = ( + 2.0 * torch.bernoulli(torch.full((d,), 0.5, device="cpu"), generator=gen) + - 1.0 + ) + _RHT_SIGNS_CACHE[key] = signs.to(device) + return _RHT_SIGNS_CACHE[key] + + +def single_rht(x: torch.Tensor, inverse: bool = False) -> torch.Tensor: + """Single Randomized Hadamard Transform: H × D₁ × x. + + Used by INT4 per-token-head quantization to gaussianize data + before asymmetric quantization. + """ + d = x.shape[-1] + d1 = _get_rht_signs(d, 0, x.device) + if inverse: + return fast_hadamard_transform(x) * d1 + else: + return fast_hadamard_transform(x * d1) diff --git a/vllm/v1/attention/ops/triton_reshape_and_cache_flash.py b/vllm/v1/attention/ops/triton_reshape_and_cache_flash.py index fb0c923055..0f0022c2cb 100644 --- a/vllm/v1/attention/ops/triton_reshape_and_cache_flash.py +++ b/vllm/v1/attention/ops/triton_reshape_and_cache_flash.py @@ -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: diff --git a/vllm/v1/attention/ops/triton_unified_attention.py b/vllm/v1/attention/ops/triton_unified_attention.py index f39e44286b..b6ba669a08 100644 --- a/vllm/v1/attention/ops/triton_unified_attention.py +++ b/vllm/v1/attention/ops/triton_unified_attention.py @@ -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" diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 5a2a5c5e29..b312a0fbee 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -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) )