QVAC-14555: TurboQuant (Vulkan): KV cache quantization (TBQ3_0 / TBQ4_0 / PQ3_0 / PQ4_0)

[jesusmb1995]

Summary

Implements TurboQuant KV cache quantization (Zandieh et al., ICLR 2026) for CPU and Vulkan backends with full Flash Attention support. Compresses KV cache to 3.25-4.25 bits per value, enabling ~4-5x larger context windows on the same hardware.

Paper: https://arxiv.org/pdf/2504.19874
Community discussion:

• TurboQuant - Extreme KV Cache Quantization ggml-org/llama.cpp#20969
• TurboQuant KV Cache Compression — Working Implementation Ready for Review ikawrakow/ik_llama.cpp#1509
  Related upstream PR: llama : rotate activations for better quantization ggml-org/llama.cpp#21038 (graph-level rotation for existing quant types)

Recommended configurations:

• High compression + speed: K=pq3_0 V=pq3_0 — codebook-only, no QJL overhead. Minimal PPL/speed loss at 3.25 bpw with a small retrieval quality trade-off on long contexts.
• High compression + and long-context quality: K=tbq3_0 V=pq3_0 — QJL-corrected keys with codebook-only values. Best retrieval accuracy at 3.75 avg bpw, with a moderate speed cost from QJL correction in the FA shader.

Features

• Full set of TurboQuant types: tbq3_0, tbq4_0, pq3_0, pq4_0 (and _64 variants)
• Automatic head_dim detection (64 vs 128) — user specifies pq3_0, internal type auto-selects
• Coopmat1 and Coopmat2 Flash Attention support (noticeable prefill speedup)
• Pre-compiled fused Flash Attention shaders for mixed K/V types (asymmetric compression)
• QJL Stage 2 correction in all FA paths (scalar, cm1, cm2)
• Comprehensive test/benchmark scripts (perplexity, throughput, RULER)
• Cooperative copy_to_quant Vulkan path for TBQ/PQ (faster KV writes)

How does TurboQuant work?

Random rotations spread values evenly across coordinates, preventing concentration on a few axes where zero-coordinates waste bits. In high dimensions, the marginal distribution of each coordinate of a unit-sphere vector follows a Beta distribution that converges to N(0, 1/d) as d grows. The algorithm exploits this by placing Lloyd-Max codebook centroids at optimal positions for this known distribution, minimizing MSE reconstruction error. Centroids are found by solving a continuous 1-dimensional k-means problem.

An additional QJL correction step (Stage 2) reduces bias in dot-product estimation. It quantizes the residual error from Stage 1 to 1-bit by storing only the signs of the residual vector after applying a random rotation (Hadamard × sign diagonal). Since only signs are stored (no centroid rounding), the paper proves this yields an unbiased dot-product estimator. This step is important for maintaining retrieval quality on long contexts.

Optimization details

• Hadamard instead of dense rotation: Rotations based on Hadamard use the butterfly pattern in O(d log d) instead of O(d²). Hadamard is deterministic, but applying a random sign diagonal preserves randomness while remaining orthogonal and invertible.

• Dense rotation for K/V/Q at graph level, FHT in shader for QJL: At block sizes d=64/128, O(d²) is negligible and utilizes better GPU parallelism for the graph-level rotation. The butterfly FHT is used inside the Flash Attention shader for the QJL projection, avoiding the need to copy a dense matrix into the shader (which would add memory pressure). Since there is no Q cache, the QJL projection of Q must be recomputed every step to apply corrections against the 1-bit signs stored in K blocks.

Type	Bits/val	Block size	Compression vs FP16	Description
q4_0	4.50	18 B	3.5x	Baseline: 16 linear values
pq3_0	3.25	52 B	4.9x	8 Lloyd-Max centroids
pq4_0	4.25	68 B	3.8x	16 Lloyd-Max centroids
tbq3_0	4.25	68 B	3.8x	8 centroids + QJL correction
tbq4_0	5.25	84 B	3.0x	16 centroids + QJL correction

Implementation overview

• vulkan-shaders-gen.cpp — orchestrates SPIR-V compilation of all variant combos
• ggml-vulkan.cpp — host-side: creates pipeline objects, dispatches compute

TurboQuant KV cache shader flow (TBQ/PQ is ONLY a KV cache type, never model weights):

STEP 1: Write to cache (same for all paths)

• copy_to_quant.comp: float K/V → TBQ/PQ quantized blocks
  • L2 norm, codebook binary search, 3/4-bit index packing
  • TBQ only: also computes QJL residual (qjl[], d_r)
  • PQ only: no QJL, smaller block, faster

STEP 2: Read cache at attention time (paths diverge here)

PATH A: Scalar Flash Attention (broad HW support, baseline)

• flash_attn.comp
• Includes: types.glsl, tq_utils.comp (via flash_attn_base.glsl), dequant_funcs.glsl
• Dequantizes K/V inline, element by element
• For TBQ/PQ K: uses centroid-gather optimization (reorders Q·K into per-centroid partial sums)
• For TBQ K only: applies QJL correction to attention scores
• Full fused kernel: QK^T → softmax → PV → output

PATH B: Cooperative matrix v1 Flash Attention (KHR, cross-vendor)

• flash_attn_cm1.comp
• K is fully dequantized into shared memory, then coopMatMulAdd for K·Q^T (subgroup-scope 16×16 tiles)
• P·V accumulation is still scalar with inline dequant
• Same QJL correction as scalar (applied to sfsh[] after coopmat store)

PATH C: Cooperative matrix v2 Flash Attention (NV only, most efficient)

• flash_attn_cm2.comp
• K and V loaded via coopMatLoadTensorNV with decode callback (dequant-on-load, no shared memory staging)
• Both K·Q^T and P·V use coopMatMulAdd (workgroup-scope matrices)
• QJL correction via raw byte reads from data_k[] with hardcoded byte offsets per type

PATH D: No-FA fallback, small N (MUL_MAT with N ≤ 8, e.g. decode)

• mul_mat_vec_tbq3_0.comp / mul_mat_vec_tbq4_0.comp
• Fused dequant + dot product, no centroid gather
• QJL correction applied in the same kernel

PATH E: No-FA fallback, large N (K·Q MUL_MAT with N > 8, e.g. prefill)

• This is the path exercised by -fa off with a TBQ/PQ K cache. Only the K·Q matmul is affected: V stays f16 under -fa off (upstream guard), so V·A stays on the existing f16 path.
• Stage 1: mul_mm.comp runs with TBQ/PQ load_a_to_shmem — centroid dequant × d into shared memory, then generic tiled matmul (scalar / cm1 pipelines; cm2 falls through to cm1/scalar since no _mat_f16 cm2 shader exists for TBQ/PQ).
• Stage 2 (TBQ only): mul_mm_tbq_qjl_correction.comp is dispatched after the main matmul as an additive pass — one workgroup per (row, col, batch), QUANT_K threads running the same Walsh–Hadamard + QJL dot product as the vec shader, accumulating d_r · √(π/2) / QUANT_K · sum_qjl(H(B)) into D.
• PQ has no Stage 2 (no qjl[] / d_r), so Stage 1 alone is exact.
• Requires B (src1) as f32; the scheduler is expected to feed f32 on this path. f16 src1 for standalone TBQ MUL_MAT reports not supported and falls back to CPU.
• Fixes external review Issue 3 on PR QVAC-14555: TurboQuant (Vulkan): KV cache quantization (TBQ3_0 / TBQ4_0 / PQ3_0 / PQ4_0) #115: before this patch supports_op claimed TBQ/PQ MUL_MAT on cm2 devices (RTX 5090) but had no pipeline behind it, so the correctness run segfaulted. tests/test-backend-ops.cpp now covers all 8 TBQ/PQ types × n ∈ {1,8,16,32} as a repro.
• Non-dim01-contiguous quantized src0 (permuted layouts) is now routed to the matrix path as well, so TBQ/PQ MUL_MAT works regardless of src0 stride pattern.

Example usage

llama-cli -m model.gguf --cache-type-k tbq3_0 --cache-type-v pq3_0
llama-cli -m model.gguf --cache-type-k pq3_0 --cache-type-v pq3_0

Works transparently with both head_dim=128 (Llama-3.1, Qwen, Mistral) and head_dim=64 (Llama-3.2-1B/3B) — the right block size is auto-selected.

Results / testing

Automatic CI/CTest should already cover the relevant backend regressions: test-backend-ops includes the TBQ/PQ backend-op cases, and test-copy-tbq-subgroups covers the Vulkan subgroup copy path. For a normal rebase/regression check, those automatic tests should be enough.

The shell scripts below are more useful for manual TurboQuant-focused testing and analysis, especially when you want to skip unrelated tests and compare report numbers more directly. The quickest manual TurboQuant sanity run is:

bash tests/test-turboquant.sh --full

That runs the TurboQuant correctness/sanity checks fairly quickly.

For a heavier but still manageable check against the simple report numbers, run the PPL and throughput scripts on the 5090 node from my checkout so the same in-place generated input files are reused:

cd ~/jberlanga/llama.cpp
bash tests/test-kv-cache-quantization-perp.sh -c large -m "mistral-q4km"
bash tests/test-kv-cache-quantization-perf.sh -c huge -m "mistral-q4km"

Compare the resulting PPL vs F16 and TG% numbers against the simple report. These longer scripts are probably overkill for a simple backend regression, but they are useful for analysis/report validation.

The test scripts are in this PR, but the input text is downloaded or auto-generated the first time the scripts run. In theory, test-kv-cache-quantization-perp.sh should use the same text offsets because the slice seed is fixed, but existing cached wiki.test.offset_<n_ctx>.raw files are reused. The report also includes a zip of the same generated input files for anyone who wants to reproduce the exact input texts. RULER is more sensitive because its data is generated from the NVIDIA RULER repo and depends on the generated validation.jsonl, tokenizer, dependency versions, and source data. PPL can also vary across hardware/backends due to numerical differences. For strict reproducibility, use the same hardware and the same already-generated input folders.

Please see Asana for latest available data: https://app.asana.com/1/45238840754660/task/1214143691877486/comment/1214346089994897?focus=true

PR for testing integration on LLM Addon: tetherto/qvac#1564

Limitations

• head_dim must be 64 or 128. Codebooks and Hadamard transform are pre-computed for these dimensions.
• d=64 quality is poor on small models — expected, as KV cache quantization generally degrades more on small models.
• Metal shaders and vectorized CPU not yet implemented.
• Optimized Flash Attention shaders require K to be PQ or TBQ, and V to be PQ, TBQ, Q4, Q8, or F16.
• Quantized V with -fa off is not supported by this PR. Upstream llama_init_from_model rejects quantized V when flash attention is disabled ("V cache quantization requires flash_attn"), and that guard is intentionally left in place. The -fa off K·Q MUL_MAT fix in this PR would extend cleanly to A·V for a quantized V as well, but the v_trans V-cache layout used under -fa off is populated by ggml_set_rows with row_size=1, which corrupts any blck_size > 1 type at write time (reproducible on CPU as well, independent of backend). Fixing that is a KV-cache refactor out of scope here; the guard will be revisited once that lands.

TBQ / PQ Vulkan support matrix

What runs on GPU vs. is refused by the context, across FA on/off on dense and MoE models. The MoE-KV-cache rows behave the same as dense because attention itself is plain MUL_MAT / FLASH_ATTN_EXT, not MUL_MAT_ID; MoE routing (MUL_MAT_ID) only applies to the FFN weights, which are never stored as TBQ/PQ.

Scenario	FA	K-type	V-type	Path	Status
Dense / MoE — KV cache	on	tbq3/4_0 or pq3/4_0	pq/tbq/q4_0/q8_0/f16	Fused FA (scalar / cm1 / cm2), QJL in kernel	Full GPU
Dense / MoE — KV cache	on	tbq3/4_0 or pq3/4_0	other quantized (q5_0, q4_1, iq4_nl, k-quants, …)	No matching Vulkan FA pipeline → per-layer backend split	Runs, but attention falls back to CPU
Dense / MoE — KV cache	off	tbq3/4_0 or pq3/4_0	f16	K·Q via mul_mm.comp + QJL correction; V·A on the existing f16 path	Full GPU (Path E)
Dense / MoE — KV cache	off	tbq3/4_0 or pq3/4_0	any quantized type (incl. tbq/pq)	—	Context init refused (upstream FA-off rule: "V cache quantization requires flash_attn")

Notes:

• Head dimensions of both 128 and 64 are supported; the _64 block variants (tbq*_0_64, pq*_0_64) have their own pipelines, codebooks, and sign tables.
• MoE FFN weights are not in this table on purpose: TBQ/PQ are KV-cache quantizations only (llama-quantize has no TBQ/PQ target, and no GGUF stores FFN experts in those types), so MUL_MAT_ID never receives TBQ/PQ src0. Attention in MoE models is a plain MUL_MAT / FLASH_ATTN_EXT and therefore falls under the "KV cache" rows above.

Remaining work

• SIMD optimization — AVX2/NEON for CPU quantize/dequantize
• Metal shaders — Apple GPU backend support
• 2-bit variant — even higher compression
• Direct cosine similarity evaluation

[zoq]

Are you planning to merge this before the rebase to the latest version of llama.cpp?

[jesusmb1995]

Are you planning to merge this before the rebase to the latest version of llama.cpp?

Not particularly, if the rebase to latest version of llama.cpp will happen soon then I will change the target to the correct temp branch. I think its better if I target latest llama.cpp version.

Edit: @zoq Since it seems we want this merged in about 1-2 weeks, I would target this version for now. yes, planning to merge this before the rebase.