feat: TurboQuant KV cache (TBQ3_0/TBQ4_0) with CPU + CUDA support

[marksverdhei]

Summary

• Port upstream TurboQuant PR (ggml : add CPU TurboQuant KV cache types (TBQ3_0 / TBQ4_0) ggml-org/llama.cpp#21089) to ht branch with CPU + CUDA GPU support
• Add security hardening (malloc NULL checks, dead code removal)
• Add CUDA dequantize/quantize kernels with 128x128 Householder rotation matrix
• Fused on-device SET_ROWS kernel for CUDA graph compatibility

What is TurboQuant?

Google Research paper (ICLR 2026) for KV cache compression. Decomposes vectors into norm + unit direction, applies random orthogonal rotation (Householder QR), then scalar Lloyd-Max quantization. Achieves 3-5x KV cache compression with minimal quality loss.

Changes

Cherry-picked from upstream PR ggml-org#21089 (elusznik):

• GGML_TYPE_TBQ3_0 (3.06 bpw) and GGML_TYPE_TBQ4_0 (4.06 bpw) type definitions
• CPU quantize/dequantize with Householder rotation (256-element blocks, 2x128 sub-blocks)
• Block layouts: block_tbq3_0 (98 bytes/256 elem), block_tbq4_0 (130 bytes/256 elem)
• CLI args, llama-bench, quantize tool support
• CPU regression tests

Security hardening (ht-specific):

• Added GGML_ASSERT NULL checks after all malloc calls in turboq code
• Fixed partial-failure pattern in rotation/projection cache (allocate-then-swap)
• Removed unused QJL projection scaffolding (86 lines of dead code)

CUDA backend (ht-specific):

• Rotation matrix generated host-side via CPU code, uploaded to device global memory (lazy init)
• Dequantize kernels: 128 threads/CUDA block, cooperative inverse rotation via shared memory matvec
• Fused on-device SET_ROWS kernel: reads row indices on GPU, computes norm + rotation + quantization + packing in one kernel (no host-device sync, CUDA graph compatible)
• Quantize kernels: two-pass (norm reduction + rotate/quantize/pack) for CPY path
• Registered in type traits, CPY, convert, SET_ROWS dispatch

Benchmark Results (Qwen2.5-3B Q4_K_M, RTX 3090, FA=on, ngl=99)

Cache type	pp512 (t/s)	tg128 (t/s)	Compression vs f16
f16 (baseline)	4,486	73.5	1.0x
q4_0	4,456	72.0	3.6x
tbq4_0	2,109	7.7	3.9x
tbq3_0	1,666	7.8	5.2x

TBQ4_0 prompt processing is 2.1x slower than f16 baseline with 3.9x memory compression. Token generation is ~9.5x slower due to full KV cache dequant in the attention path.

Future optimization opportunities:

• Fused flash attention kernels with tile-by-tile TBQ dequant (avoids materializing full dequantized KV in global memory)
• Rotated-query vec_dot: dot(q, k) = dot(Q*q, codebook_values) — rotate query once, then cheap codebook dot products
• Native FA template instances for TBQ types

Security Audit

Foreign code (upstream PR ggml-org#21089) audited before integration:

• No critical/high severity issues found
• Codebook access always bounded by bitmasks (no OOB possible)
• Householder QR numerically standard
• All findings addressed (malloc checks, dead code removal)

References

• TurboQuant paper: https://arxiv.org/abs/2504.19874
• Upstream PR: ggml : add CPU TurboQuant KV cache types (TBQ3_0 / TBQ4_0) ggml-org/llama.cpp#21089
• Community discussion: TurboQuant - Extreme KV Cache Quantization ggml-org/llama.cpp#20969

Test plan

• CPU build compiles
• CUDA build compiles (RTX 3090, sm_86)
• llama-bench with tbq4_0/tbq3_0 KV cache (prompt + generation)
• CUDA graph capture works with TBQ SET_ROWS
• Perplexity comparison (tbq4_0 vs q4_0 vs f16)
• Server endpoint test with TBQ KV cache

🤖 Generated with Claude Code

[marksverdhei]

Performance & Quality Review — TurboQuant KV Cache

Speed Regression Root Cause (CRITICAL)

The 9.5x token generation slowdown (7.7 vs 73.5 t/s) is caused by full KV cache dequantization in global memory before attention (llama-graph.cpp:1793-1829):

k = ggml_cast(ctx0, k, tbq_attn_type);  // dequant ENTIRE K cache to f16
v = ggml_cast(ctx0, v, tbq_attn_type);  // dequant ENTIRE V cache to f16

Three compounding factors:

1. Every generated token dequants the entire KV cache — cost scales linearly with sequence length, paid per token
2. Dequant is O(n²) per block — 128×128 Householder inverse rotation = 16,384 FLOPs per 128 elements (vs ~zero for q4_0)
3. Rotation matrix reads are non-coalesced — column-strided access (stride 512 bytes) to Q_rot in global memory, near-zero coalescing
4. Memory bandwidth wasted — compressed KV advantage negated by writing full f16 copy, then reading again for attention

pp512 is "only" 2.1x slower because dequant cost amortizes across the batch. q4_0 loses only 2% because its dequant is one multiply per element with no rotation.

Task Quality: CORRECT

• Householder QR: Standard implementation with Haar sign correction (Mezzadri 2007) ✓
• Lloyd-Max codebooks: Symmetric, optimal for N(0,1) ✓
• CUDA/CPU consistency: CUDA generates rotation via CPU path → bit-identical ✓
• Seed handling: Fixed global rotation (consistent with paper) ✓
• Scale factor: 1/sqrt(256) consistent between CPU and CUDA ✓

Security: ADEQUATE

• malloc NULL checks with GGML_ASSERT ✓
• Codebook access bounded by bitmasks (no OOB) ✓
• Allocate-then-swap pattern prevents half-updated state ✓
• Dead QJL scaffolding removed (86 lines) ✓

Minor issues:

• d_turboq_norms buffer never freed (memory leak)
• Static d_turboq_Q pointer not multi-device safe (racy init)
• TLS scratch buffers grow but never shrink

Optimization Path Assessment

Optimization	Impact	Effort	Assessment
Rotated-query vec_dot (dot(Q*q, codebook_values))	Transformative — avoids 128×128 matvec in hot path entirely	1-2 weeks	HIGHEST PRIORITY — would make tg comparable to q4_0
Fused FA with tile-by-tile dequant	High — avoids materializing full dequantized cache	2-4 weeks	Realistic, rotation matrix fits in shared memory (64KB)
Native FA template instances	Clean integration	Depends on above	Low priority until fused FA exists

Unchecked Test Items

• Perplexity: needs llama-perplexity runs with --cache-type-k tbq4_0 vs q4_0 vs f16 on wikitext-2
• Server endpoint: needs llama-server with TBQ cache types, test multi-request memory growth and CUDA graph capture

Verdict

Solid correctness baseline — numerical implementation is right, security is adequate. Not production-ready due to the 9.5x tg regression. The rotated-query vec_dot optimization should be the next priority, as it would fundamentally change the performance profile. Recommend merging to ht as a foundation, with a follow-up PR for the optimized attention path.

---

🤖 Review by ht-fork-manager agent