[Attention Backend] TurboQuant: 2-bit KV cache compression with 4x capacity - #38479
Conversation
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
There was a problem hiding this comment.
Code Review
This pull request implements TurboQuant, a near-optimal KV-cache quantization scheme for vLLM, including new cache types (tq3, tq4), a dedicated attention backend, and optimized Triton/CUDA kernels. The implementation covers centroids calculation, bit-packing quantizers, and fused store/decode operations. Feedback identifies critical bugs in the 3-bit unpacking logic for values crossing byte boundaries and the fallback FP8 conversion for older CUDA architectures, both of which require more precise bit manipulation to ensure numerical correctness.
|
This pr made extensive changes, which may not get concerned as review is very hard. |
|
This pull request has merge conflicts that must be resolved before it can be |
|
I see your approach now. Adding a new EDIT: After discussion during #sig-quantization meeting, the consensus was that it might be best to adopt this standalone |
|
Running a quick smoke test on H100 results in 0% gsm8k If I run |
| @staticmethod | ||
| def get_kv_cache_shape( | ||
| num_blocks: int, | ||
| block_size: int, | ||
| num_kv_heads: int, | ||
| head_size: int, | ||
| cache_dtype_str: str = "tq3", | ||
| ) -> tuple[int, ...]: | ||
| """Combined K+V cache shape — no leading 2 dimension. | ||
|
|
||
| Layout: (num_blocks, block_size, num_kv_heads, padded_slot_size) | ||
| Each slot = [key_packed | value_fp16 | padding]. | ||
|
|
||
| Note: head_size here is the *effective* head_size from the spec | ||
| (= padded_slot // 2), NOT the model's actual head_dim. | ||
| So padded_slot = head_size * 2. | ||
| """ | ||
| return (num_blocks, block_size, num_kv_heads, head_size * 2) |
There was a problem hiding this comment.
I'm not sure if removing the separate k and v will cause issues with other features - I know there are several places like kv transfer where if we don't find a leading 2 dim, then we assume it is MLA cache
There was a problem hiding this comment.
cc @NickLucche do you know if this will break anything re: the above?
|
Tested this on DGX Spark (GB10, SM121, 128 GB unified memory, aarch64) with Nemotron-3-Nano-30B-A3B-NVFP4. Applied the patch to vllm-node (eugr community build, vLLM 0.18.1rc1 with prebuilt SM121 FlashInfer). Used Triton fallback path — the CUDA kernels aren't compiled for SM121 yet. Results with
*FP8 short-context affected by warmup. Also tested 240K context (256K max_model_len) — needle-in-haystack recall working at 64 GB, zero memory creep. Patch details and full results: https://github.com/Sggin1/spark-ai-containers/tree/main/turboquant Six files needed minor patches to work with 0.18.1rc1 (dtype lookups, backend candidate list). Happy to share specifics if helpful. Disclaimer: Limited testing on a single model/hardware config. I'm a hobbyist, not an ML engineer — these results may not generalize. Sharing in case it's useful for SM121 validation. |
|
@vibhavagarwal5 Given this is a new feature and there are no unit tests, please disclose your local accuracy test e.g. NIAH accuracy scores. CC @mgoin |
Hi, @mgoin I noticed there are a few PRs implementing this feature. do we have a clear roadmap for moving it forward? |
Ampere (SM86) Compatibility + Quality FixTested this PR on 8x RTX A4000 (SM86) with Nemotron-Cascade-2-30B-A3B (hybrid Mamba+MoE+Attention, head_dim=128). A few findings: FP8 Ampere FixThe Triton kernels use Quality at Different Value Bit-widthsThis directly addresses @mgoin's 0% gsm8k finding. The default 2-bit value quantization destroys reasoning quality. FP8 values restore it completely:
Hybrid Mamba Model SupportNemotron-Cascade-2 is a 52-layer hybrid (Mamba2 + MoE + Attention). It loaded and served successfully — no page size errors. Only the 8 full-attention layers use TQ compression; Mamba/MoE layers pass through unchanged. Norm CorrectionAlso tested TheTom's norm correction (storing |
|
On the 0% gsm8k with tq3 The quality failure on The fix is asymmetric K/V bit allocation: K at 4-bit (preserves routing), V at 3-bit (weighted sum tolerates noise). @varjoranta's H100 benchmarks on the feature request thread confirm this — asymmetric K4/V3 scored highest (4.75) on Qwen3-235B AWQ, above both symmetric turbo4 and FP16 baseline. Quality validation tooling We ship a pip install turboquant-vllm
python -m turboquant_vllm.verify --model Qwen/Qwen3-4B --bits 4This catches quality regressions before full eval suites. We've validated 7 model families (Molmo2, Mistral, Llama, Qwen2.5, Phi-3/4, Gemma 2/3) across head_dim 96/128/256 — all pass >0.99 cosine at tq4. head_dim compatibility note Re @MidasMining's Ampere findings — we hit similar issues with non-pow2 head dimensions (Phi at 96, Gemma at 256). The fix is padding to next power of 2 in the Triton kernel with masked loads/stores. Happy to share the approach if useful for this PR. |
|
Re @albertocodesdev's asymmetric K/V suggestion — our data supports asymmetric allocation but points in a different direction on which axis needs precision. Alberto recommends K4/V3 (more bits on keys). Our testing on Nemotron-Cascade-2 (hybrid Mamba, head_dim=128) showed the opposite: values are the precision bottleneck for reasoning tasks, not keys.
Keys at 3-bit were fine across all configurations — the failures tracked exactly with value precision. The 4 checks that fail at 2-bit/4-bit values are the hardest multi-step reasoning tasks (race condition detection, memory leak identification, float precision bugs). This matters because cosine similarity can pass >0.99 while reasoning still breaks. Our 2-bit value config likely has high cosine similarity (the reconstruction is close in aggregate) but the small errors in value vectors compound through the residual stream across multiple reasoning steps. Cosine is a necessary but not sufficient quality metric for reasoning models. The reconciliation with @varjoranta's K4/V3 result may be model-dependent — Qwen3-235B AWQ has different attention patterns than a Mamba hybrid. Worth testing both K4/V3 and K3/V8 across model families to see if the optimal allocation varies. |
…antizer Add TurboQuant KV cache compression support: - Register tq3, tq4, tq_k4v3 in CacheDType - Add TURBOQUANT to AttentionBackendEnum - Route tq* dtypes to TURBOQUANT backend in CUDA platform - Add vllm/turboquant/ module: config, centroids (Lloyd-Max), quantizer (PolarQuant + WHT) Aligned with PR vllm-project#38479 structure. Adds asymmetric K/V support (tq_k4v3) which is not in the original PR. Attention backend implementation (turboquant_attn.py) in next commit.
|
On asymmetric K/V being model-architecture-dependent Your Nemotron-Cascade-2 data is compelling — and I think we're both right for different architectures. Our v1.4.0 shipped asymmetric K/V (
K4/V3 produced identical text output on the three models tested for generation quality (Qwen2.5-3B, Gemma 2 2B, Gemma 3 4B). But you're right that cosine is necessary-not-sufficient — we haven't validated K4/V3 on multi-step reasoning benchmarks like your 14-check suite. The reconciliation is likely what you suggested: optimal K/V allocation varies by architecture. Standard attention models tolerate V compression well; hybrid Mamba models with different information flow through the residual stream may need V preserved. We haven't tested hybrid Mamba models — our validation covers pure transformer attention only. For anyone wanting to find the right config for their model: pip install turboquant-vllm
python -m turboquant_vllm.verify --model <your-model> --k-bits 4 --v-bits 3 --threshold 0.97 |
|
Following this with interest for RTX 5090 (SM120) + Qwen3.5-9B hybrid workloads. The 4x KV cache on hybrid attention is exactly what we need for batch throughput scaling. A few observations from our testing setup (170 tok/s single-user, ~5K tok/s batch on NVFP4):
— Chris Klaus (AutoKernel) |
|
qwen3.5 dense models not supported? Error TurboQuant KV cache is not supported for hybrid (attention + Mamba) models. Boundary layer protection requires uniform attention layers. |
It does not, you have to use old models for this to even work. Defeats the purpose. |
|
@gaby I am following a forked repo that did implement it with the kernel for Qwen3.5/3.6 but only for blackwell so far. |
|
@xyehya Could you point me the repo? |
|
@vibhavagarwal5 Great to see this work land in vLLM! Starting to look at vLLM myself today. Looking forward to making improvements! |
|
@TheTom llama.cpp TQ finalized I guess? 😊 |
|
Can you support A100 GPU + |
|
same here : v0.20.0 , pro4500. |
|
@seasoncool @sweihub Report here #39931 |
|
Hi @vibhavagarwal5 thank you so much for this. I maintain turboquant-101 ([repo link]), a teaching and verification repo built on your merged code path: it walks beginners through running I have two questions for you as I update my repo....
Happy to fix anything you catch. My repo credits this PR as the foundation of its whole CUDA path. Thank you again! |
Summary
TurboQuant adds online KV cache compression to vLLM's v1 attention backend using PolarQuant (WHT rotation + Lloyd-Max scalar quantization) for keys and uniform quantization for values. All quantization happens at store time via fused Triton kernels — no offline calibration, model changes, or weight modifications required. Just set
--kv-cache-dtype turboquant_k8v4.Compression Presets (Qwen3-4B, head_dim=128)
turboquant_k8v4turboquant_4bit_ncturboquant_k3v4_ncturboquant_3bit_ncBaseline: GSM8K 0.900, NIAH 100%. Measured on Qwen/Qwen3-4B with 5-shot GSM8K (200q) and NIAH (512-32K, 77 probes).
Performance (Qwen3-4B, 4x RTX PRO 6000 Blackwell, cudagraphs+compile)
Throughput (output tok/s)
TPOT (ms) — lower is better
TTFT (ms) — lower is better
Key Takeaways
Technical Innovations
Walsh-Hadamard Transform (WHT) rotation — Replaced QR-decomposed random orthogonal matrices with WHT + random sign flips. Orthonormal, self-inverse (
H = H^T = H^{-1}), enabling future in-kernel butterfly fusion. Same D×D matmul API, zero quality regression, consistent +0.5-2.5% improvement from structured Hadamard cache patterns. Continuation-prefill inversion is triviallyH @ x(no transpose needed).Fused MSE store kernel — Bucketize, centroid gather, residual norm, index packing, and value quantization fused into a single Triton kernel (
_tq_fused_store_mse), eliminating 4 separate PyTorch kernel launches per layer. Result: +18-21% decode throughput, -10-12% prefill TTFT.In-kernel FP8 cast — FP8 key cast moved from host-side
torch.float8_e4m3fnto in-kerneltl.float8e4nv/tl.float8e4b15, removing a separate kernel launch. Auto-detects SM capability for Ampere vs Hopper FP8 formats.Compact slot sizes — Slots are rounded to next even number instead of power-of-2, eliminating up to 47% padding waste (t4nc: 136B vs 256B).
TQFullAttentionSpecproperly overridesreal_page_size_byteswith compact TQ slot bytes.Shared value quant JIT helper — Extracted
_store_quantized_valueTriton JIT function, deduplicating ~60 lines between FP8 and MSE store kernels for both 3-bit and 4-bit value paths.Prefill
.tolist()optimization — Single CPU-GPU sync via.tolist()instead of per-request.item()calls in the prefill loop.CUDAGraph memory fix — Static
NUM_KV_SPLITSgrid dimension (configurable, default 32) enables CUDAGraph capture. Estimated GPU memory reduced from 33 GiB → 8.7 GiB.Stream overlap — KV store runs on a secondary CUDA stream so it can overlap with the next layer's forward pass (disabled during CUDAGraph capture).
Architecture
Design Decisions
kv_cache_dtype_skip_layersto protect embedding-adjacent representations. Also supports skipping"sliding_window"layers and arbitrary layer indices.TQFullAttentionSpec— proper spec subclass that overridesreal_page_size_byteswith TQ slot bytes, with correct merge semantics for uniform-spec models. PassesUniformTypeKVCacheSpecs.is_uniform_type()check as aFullAttentionSpecsubclass.flash_attn_varlen_funcfor memory-efficient O(N) prefill, with a continuation-decode threshold (128 tokens) routing small chunks directly through the TQ decode kernel.Usage
Scope
Supports full-attention and uniform sliding-window transformer models. Hybrid architectures (mamba+attention, interleaved SWA) are planned for a follow-up PR.
Test Plan