PR #36893: [Feature] Kvcache Int8 per-token scale on TRITON_ATTN continue of #34327 thanks EricccYang

• Repository: vllm-project/vllm
• Author: JartX
• State: closed | merged: False
• Link: https://github.com/vllm-project/vllm/pull/36893

Description (problem / solution / changelog)

FIX https://github.com/vllm-project/vllm/issues/37319

This PR adds INT8 per-token KV cache quantization to vLLM for the Triton attention backend.

Unlike the existing FP8 path, this mode uses dynamic per-token scales computed at cache-write time. In practice, this reduces KV cache memory roughly by half compared to bf16/fp16, offering a "plug and play" zero-config option without requiring pre-calibrated checkpoint scales.

This change introduces a shared INT8 KV cache contract for attention backends: dispatch via the KVQuantMode enum, declare auxiliary buffers through AttentionSpec.auxiliary_buffer_specs, allocate them upfront in the model runner, bind them via AttentionImpl.bind_auxiliary_buffers(), and account for their memory accurately. Triton is the first backend implementing this contract.

What changed

• Added int8_per_token as a supported kv_cache_dtype.
• Introduced the KVQuantMode enum so backends dispatch on semantic quantization mode instead of string matching prefixes.
• Replaced lazy allocation with explicit upfront allocation of scale caches via _allocate_auxiliary_buffers() and bind_auxiliary_buffers().
• Added a dedicated Triton INT8 reshape-and-cache kernel for per-token dynamic quantization.
• Extended Triton unified attention to perform an explicit 3-way dispatch and read per-token INT8 scale caches during dequantization.
• Accounted for the explicit scale-cache memory through auxiliary_buffer_specs in the KV block-count calculation.

Design notes

The main idea is to offer a memory-efficient KV cache option that requires no additional configuration. The user only needs to set kv_cache_dtype="int8_per_token".

Instead of extending the existing FP8 kernel, INT8 uses a separate reshape kernel because it requires dynamic scale computation. The kernel computes an absmax across all heads for a given token, deriving a single scale. This shifts the auxiliary memory footprint to strictly per-token (rather than per-token per-head), optimizing the buffer size while maintaining bounded accuracy costs.

Backend Extensibility

While Triton is the first target backend in this PR, the redesign lays the groundwork for future quantization modes and other backends without requiring core architectural rewrites.

The new KVQuantMode enum already declares PER_TOKEN_GROUP and NVFP4 as future modes. Thanks to the explicit auxiliary_buffer_specs abstraction, supporting these will be straightforward:

• PER_TOKEN_GROUP can easily declare and allocate two float32 buffers with shape (block_size, num_kv_heads, ceil(head_size / group_size)).
• NVFP4 can request its required block-scales and global-scales through the same unified interface.

Triton implementation

When the user sets --kv-cache-dtype int8_per_token, process_weights_after_loading() forces checkpoint scales to 1.0 and deletes them, returning early because actual scales are dynamically generated.

For Triton, the write path uses _reshape_cache_per_token. It launches one instance per token, computes a global absmax over heads, derives scale = max(absmax / QUANT_MAX, 1e-6), and writes the quantized data alongside the float32 scales into the auxiliary scale buffers.

On the read path, forward() performs a dispatch based on self.kv_quant_mode. For PER_TOKEN, it passes k_scale_cache and v_scale_cache into unified_attention(), and the Triton attention kernels fetch and apply these per-token scales during the dot product accumulation.

Memory accounting

Until now, KV cache auxiliary memory was treated too flatly. Now, KVCacheSpec exposes auxiliary_buffer_specs.

For PER_TOKEN, it explicitly declares two float32 buffers with shape (block_size,). kv_cache_utils.py is updated to include this specific memory footprint when deciding how many pages fit into the available budget, ensuring the block count reflects the actual footprint of both the quantized tensor and the explicit float32 buffers.

Benchmark Results

Hardware: AMD RX 7900 XTX (gfx1100, no FP8 support)
Backend: vLLM 0.17.1rc1 + TRITON_ATTN + ROCm
Task: ChartQA — 2500 samples, 0-shot (lm-evaluation-harness)
Model weights: GPTQ W4A16-G32 for both models

Qwen 3.5 A3B — fp16 vs int8 kvcache vllm bench benchmark

35B model

27B model

Analysis:

On the dense model the accuracy delta is essentially noise — less than 0.005 across all metrics, well within the statistical margin (stderr ~0.007). In practice you wouldn't notice the difference.

The MoE model is a bit more sensitive to quantization (~0.010–0.016 delta), which makes sense given that expert routing produces more heterogeneous KV distributions. Still very much within acceptable bounds for a 50% memory saving, and notably better than what you'd expect from a per-tensor INT8 scheme.

Both models required no calibration — scales are computed dynamically at runtime with no offline prep.

Changed files

• docs/design/attention_backends.md (modified, +1/-1)
• tests/models/quantization/test_int8_kv_cache.py (added, +91/-0)
• tests/quantization/test_int8_kv_cache.py (added, +732/-0)
• vllm/config/cache.py (modified, +11/-2)
• vllm/model_executor/layers/attention/attention.py (modified, +4/-0)
• vllm/model_executor/layers/attention/chunked_local_attention.py (modified, +2/-0)
• vllm/model_executor/layers/attention/cross_attention.py (modified, +6/-1)
• vllm/model_executor/layers/attention/static_sink_attention.py (modified, +2/-0)
• vllm/model_executor/layers/quantization/kv_cache.py (modified, +18/-1)
• vllm/utils/torch_utils.py (modified, +1/-0)
• vllm/v1/attention/backend.py (modified, +24/-2)
• vllm/v1/attention/backends/triton_attn.py (modified, +64/-17)
• vllm/v1/attention/ops/triton_reshape_and_cache_flash.py (modified, +208/-1)
• vllm/v1/attention/ops/triton_unified_attention.py (modified, +135/-34)
• vllm/v1/core/kv_cache_utils.py (modified, +16/-5)
• vllm/v1/kv_cache_interface.py (modified, +117/-1)
• vllm/v1/worker/gpu_model_runner.py (modified, +67/-0)

PR #5: [Feature] KV cache per-token-head INT4 quantization support

• Repository: lesj0610/vllm
• Author: lesj0610
• State: closed | merged: False
• Link: https://github.com/lesj0610/vllm/pull/5

Description (problem / solution / changelog)

Summary

This PR adds int4_per_token_head KV cache dtype, extending the per-token-head quantization design from #38378.

• packed int4 layout (2 channels per byte) with fp16 scales grouped by 32 channels, stored inline in the KV cache tensor
• Hadamard rotation on query and key before caching to reduce quantization error for Gaussian-like activations
• AttentionSpec.page_size_bytes / real_page_size_bytes updated to account for the int4 packed size and fp16 scale overhead
• fix profiling-time padded-page reshape to use as_strided instead of .view() when page sizes are padded

Related: #37319, #36893, #34327

Validation

pytest tests/quantization/test_per_token_kv_cache.py -v
TEST_PER_TOKEN_KV_CACHE_MODEL=meta-llama/Llama-3.2-1B-Instruct pytest tests/models/quantization/test_per_token_kv_cache.py -v -s
pytest tests/kernels/attention/test_attention_selector.py -k 'per_head_quant_scales_backend_selection or flash_attn_rejects_int4_kv_cache' -v
pytest tests/v1/worker/test_gpu_model_runner.py -k 'reshape_kv_cache_tensors_handles_padded_attention_pages or update_hybrid_attention_mamba_layout_with_num_block_2_rewrites_stride' -v
pytest tests/v1/core/test_kv_cache_utils.py -k 'page_size_padded or unify_kv_cache_spec_page_size' -v

Smoke-tested non-eager serving on Gemma4 26B and 31B (TP=2) with --kv-cache-dtype int4_per_token_head, VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=1, -O1, and --gpu-memory-utilization 0.90. Tested text, multilingual, long-context, MM, and multi-image inputs.

Qwen 3.5 Benchmark Summary

The first two columns below are the reference baseline numbers from #36893 on non-quantized Qwen weights.

The AutoRound columns are the local control benchmark added for this PR to show the existing int8_per_token_head path is still working and to compare it with the new int4_per_token_head path.

The FP16/INT8 baseline columns are quoted from #36893. Local reruns on 2x RTX 3090 focus on AutoRound INT8 vs AutoRound INT4.

Local AutoRound benchmark conditions:

• TP=2
• VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=1
• -O1
• prefix caching disabled
• random benchmark dataset, input_len=1024, output_len=128, num_prompts=10
• 35B MoE used --gpu-memory-utilization 0.85
• 27B dense used --gpu-memory-utilization 0.87 --max-cudagraph-capture-size 128

35B MoE (Qwen3.5-35B-A3B)

27B Dense (Qwen3.5-27B)

The slowdown is expected because AutoRound model quantization and int4 KV cache are both active at the same time.

Notes

• scale caches initialized with 1.0 so partially-written blocks dequantize to safe values instead of garbage
• profiling-time padded-page reshape now uses as_strided, same as the runtime path

Changed files

• tests/kernels/attention/test_attention_selector.py (modified, +37/-0)
• tests/models/quantization/test_per_token_kv_cache.py (modified, +19/-4)
• tests/quantization/test_per_token_kv_cache.py (modified, +523/-108)
• tests/v1/core/test_kv_cache_utils.py (modified, +58/-0)
• tests/v1/worker/test_gpu_model_runner.py (modified, +84/-0)
• vllm/config/cache.py (modified, +1/-0)
• vllm/utils/torch_utils.py (modified, +1/-0)
• vllm/v1/attention/backends/flash_attn.py (modified, +1/-1)
• vllm/v1/attention/backends/triton_attn.py (modified, +161/-37)
• vllm/v1/attention/ops/triton_reshape_and_cache_flash.py (modified, +288/-0)
• vllm/v1/attention/ops/triton_unified_attention.py (modified, +290/-14)
• vllm/v1/core/kv_cache_utils.py (modified, +9/-0)
• vllm/v1/kv_cache_interface.py (modified, +63/-11)
• vllm/v1/worker/gpu/attn_utils.py (modified, +15/-1)
• vllm/v1/worker/gpu_model_runner.py (modified, +20/-6)