PR #7: perf(tq): fuse MSE store ops + inline decode Q cast + cleanup

• Repository: vibhavagarwal5/vllm
• Author: vibhavagarwal5
• State: closed | merged: True
• Link: https://github.com/vibhavagarwal5/vllm/pull/7

Description (problem / solution / changelog)

Summary

• WHT rotation: Replace QR-decomposed random orthogonal matrices with Walsh-Hadamard Transform + random sign flips for key/query rotation. Drop-in replacement (same D×D matmul), orthonormal + self-inverse, enables future in-kernel butterfly fusion
• Fused MSE store: Bucketize/centroid-gather/residual-norm fused into single Triton kernel (_tq_fused_store_mse), eliminating 4 PyTorch kernel launches per layer
• In-kernel FP8 cast: FP8 key cast moved from host-side torch.float8_e4m3fn to in-kernel tl.float8e4nv/tl.float8e4b15, removing a separate kernel launch
• Value quant dedup: Extracted shared _store_quantized_value Triton JIT helper, deduplicating ~60 lines between FP8 and MSE store kernels
• Prefill .tolist() optimization: Single CPU-GPU sync instead of per-request .item() calls in prefill loop
• CUDAGraph memory fix: Static NUM_KV_SPLITS=32 reduced estimated memory from 33 GiB → 8.7 GiB
• Dead code cleanup: Removed unused loggers, kernel constexprs (NUM_Q_HEADS, PADDED_SLOT, MAX_NUM_BLOCKS, N_CENTROIDS), value_packed_size params, stale QR matrix buffer

Benchmark results (Qwen3-4B, 4× RTX PRO 6000 Blackwell, cudagraphs+compile)

Quality remains constant

Throughput (output tok/s)

TPOT avg (ms) — lower is better

TTFT avg (ms) — lower is better

Improvement vs pre-optimization baseline (0408 run)

Decode overhead dropped from ~3x baseline to ~1.3-1.5x — a ~55% reduction in the TQ-vs-baseline gap:

Key takeaways

• k8v4 (FP8 keys + 4-bit values, ~2x compression): 79-100% of baseline throughput
• t4nc (4-bit MSE + 4-bit values, ~3.8x compression): 71-96% of baseline
• k3v4nc (3-bit MSE + 4-bit values, ~3.5x compression): 68-94% of baseline
• k8v4 long-prefill TPOT is faster than baseline (135.2ms vs 138.1ms) — compressed cache reduces memory bandwidth
• WHT rotation: No regression vs QR; consistent +0.5-2.5% improvement from structured Hadamard cache patterns

Test plan

• Full perf benchmark (6 scenarios × 5 configs) — no regressions on baseline
• All TQ configs produce correct output (k8v4, t4nc, k3v4nc, t3nc)
• CUDAGraph capture verified (51 FULL + 51 PIECEWISE graphs)
• WHT smoke test: coherent generation across all MSE configs
• Quality benchmark (PPL/GSM8K) sanity check

🤖 Generated with Claude Code

Changed files

• vllm/config/attention.py (modified, +5/-0)
• vllm/model_executor/layers/attention/attention.py (modified, +3/-3)
• vllm/model_executor/layers/quantization/turboquant/quantizer.py (modified, +18/-3)
• vllm/v1/attention/backends/turboquant_attn.py (modified, +65/-19)
• vllm/v1/attention/ops/triton_turboquant_decode.py (modified, +38/-62)
• vllm/v1/attention/ops/triton_turboquant_store.py (modified, +239/-154)

PR #38479: [Attention Backend] TurboQuant: 2-bit KV cache compression with 4x capacity

• Repository: vllm-project/vllm
• Author: vibhavagarwal5
• State: open | merged: False
• Link: https://github.com/vllm-project/vllm/pull/38479

Description (problem / solution / changelog)

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)

Baseline: 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

• k8v4 (FP8 keys + 4-bit values, ~2.6x compression): 79-100% of baseline throughput across all scenarios
• t4nc (4-bit MSE + NC, ~3.8x compression): 71-96% of baseline
• k8v4 TPOT is faster than baseline on long sequences (135.2ms vs 138.1ms) — compressed cache reduces memory bandwidth pressure
• Very-long-prefill at parity — 8K→64 shows 100% of baseline tok/s for k8v4

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 trivially H @ 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_e4m3fn to in-kernel tl.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). TQFullAttentionSpec properly overrides real_page_size_bytes with compact TQ slot bytes.

Shared value quant JIT helper — Extracted _store_quantized_value Triton 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_SPLITS grid 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

┌──────────────────────────────────────────────────────────────────┐
│  Store path (Triton)                                            │
│  K → WHT rotation → Lloyd-Max quantize → bit-pack ──┐          │
│  V → uniform quantize → bit-pack ────────────────────┤→ cache   │
│                                                      │          │
│  Decode path (Triton, split-KV)                      │          │
│  cache → unpack K → dequant → Q·K scores ──┐         │          │
│  cache → unpack V → dequant ──→ score·V ───┤→ output │          │
│                                            │         │          │
│  Prefill path (flash_attn_varlen_func)     │         │          │
│  Raw Q, K, V → flash attention → output    │         │          │
│  (continuation decode via TQ decode kernel)│         │          │
└──────────────────────────────────────────────────────────────────┘

Design Decisions

• Compact even-aligned slots — slots rounded to next even number (not pow2), eliminating up to 47% memory waste. Hybrid mamba+attention models are out of scope for this PR.
• Boundary layer protection — first/last N layers keep FP16 KV cache via kv_cache_dtype_skip_layers to protect embedding-adjacent representations. Also supports skipping "sliding_window" layers and arbitrary layer indices.
• TQFullAttentionSpec — proper spec subclass that overrides real_page_size_bytes with TQ slot bytes, with correct merge semantics for uniform-spec models. Passes UniformTypeKVCacheSpecs.is_uniform_type() check as a FullAttentionSpec subclass.
• No QJL — intentionally omitted per community consensus (5+ independent groups found it hurts attention quality by amplifying variance through softmax).
• Norm correction (NC) — re-normalizes centroid vectors to unit norm before inverse rotation during dequant, fixing quantization-induced norm distortion (~0.8% PPL improvement at 4-bit).
• Flash-attention prefill — uses flash_attn_varlen_func for memory-efficient O(N) prefill, with a continuation-decode threshold (128 tokens) routing small chunks directly through the TQ decode kernel.

Usage

# FP8 keys + 4-bit values (best quality/throughput trade-off)
vllm serve Qwen/Qwen3-4B --kv-cache-dtype turboquant_k8v4

# 4-bit MSE keys + 4-bit values + norm correction (3.8x compression)
vllm serve Qwen/Qwen3-4B --kv-cache-dtype turboquant_4bit_nc

# Maximum compression (4.9x)
vllm serve Qwen/Qwen3-4B --kv-cache-dtype turboquant_3bit_nc

# Skip specific layers (boundary protection)
vllm serve Qwen/Qwen3-4B --kv-cache-dtype turboquant_k8v4 \
  --kv-cache-dtype-skip-layers 0,1,34,35

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

• Full perf benchmark (6 scenarios × 5 configs) — no regressions on baseline
• All TQ configs produce correct output (k8v4, t4nc, k3v4nc, t3nc)
• CUDAGraph capture verified (51 FULL + 51 PIECEWISE graphs)
• WHT rotation: coherent generation across all MSE configs
• Quality benchmark: GSM8K + NIAH across all presets
• Mixed batch (decode+prefill) correct routing
• LM Eval harness integration test

Changed files

• .buildkite/test_areas/lm_eval.yaml (modified, +10/-0)
• tests/evals/gsm8k/configs/Qwen3-4B-TQ-k3v4nc.yaml (added, +5/-0)
• tests/evals/gsm8k/configs/Qwen3-4B-TQ-k8v4.yaml (added, +5/-0)
• tests/evals/gsm8k/configs/Qwen3-4B-TQ-t3nc.yaml (added, +5/-0)
• tests/evals/gsm8k/configs/Qwen3-4B-TQ-t4nc.yaml (added, +5/-0)
• tests/evals/gsm8k/configs/models-turboquant.txt (added, +4/-0)
• tests/quantization/test_turboquant.py (added, +328/-0)
• vllm/config/attention.py (modified, +5/-0)
• vllm/config/cache.py (modified, +4/-0)
• vllm/engine/arg_utils.py (modified, +18/-0)
• vllm/model_executor/layers/attention/attention.py (modified, +48/-0)
• vllm/model_executor/layers/quantization/turboquant/__init__.py (added, +13/-0)
• vllm/model_executor/layers/quantization/turboquant/centroids.py (added, +89/-0)
• vllm/model_executor/layers/quantization/turboquant/config.py (added, +185/-0)
• vllm/model_executor/layers/quantization/turboquant/quantizer.py (added, +39/-0)
• vllm/platforms/cuda.py (modified, +5/-0)
• vllm/utils/torch_utils.py (modified, +4/-0)
• vllm/v1/attention/backends/registry.py (modified, +4/-0)
• vllm/v1/attention/backends/turboquant_attn.py (added, +775/-0)
• vllm/v1/attention/ops/triton_turboquant_decode.py (added, +546/-0)
• vllm/v1/attention/ops/triton_turboquant_store.py (added, +381/-0)
• vllm/v1/core/single_type_kv_cache_manager.py (modified, +4/-2)
• vllm/v1/kv_cache_interface.py (modified, +27/-0)