Port turbovec FastScan kernel into SharpInference.TurboQuant for KV-cache attention

[pekkah]

Background

turbovec (MIT-licensed Rust ANN index) and SharpInference.TurboQuant both implement the same underlying algorithm (Zandieh et al., arXiv:2504.19874): WHT rotation + Lloyd-Max scalar quantization. turbovec uses a FAISS-FastScan-derived SIMD inner loop (nibble-split byte-shuffle LUTs + u16 accumulators) that we can borrow to speed up the existing TurboQuant KV-cache hot paths.

This is option (2) from the integration analysis — borrow the kernel, not the library. Pure C# port, no FFI, NativeAOT-friendly.

Why this is worth doing

The current K-scoring loop at src/SharpInference.Engine/ForwardPass.cs:995-1002 calls KvCacheCompressor.DequantDot once per cached position. The AVX2 inner loop (TurboQuantOps.DequantDot4Avx2, src/SharpInference.TurboQuant/TurboQuantOps.cs:232) vectorizes along dim using Avx2.GatherVector256, which is the slow leg on every recent uarch (~2-3× slower than equivalent scalar throughput).

FastScan inverts the loop: vectorize along vectors (32 KV positions per tile) and replace gathers with vpshufb (1/clock). The V-aggregation at ForwardPass.cs:1016-1025 (currently full Decompress + scalar weighted sum) benefits from the same kernel transposed.

Kernel design

TurboQuant is scalar quantization (one Lloyd-Max codebook per layer, not PQ), so the FAISS layout doesn't port 1:1, but the byte-shuffle trick does. For 4-bit, dim=128:

Per-query setup (once per head):
  scale = max over (d,v) of |centroids[v] * rotatedQuery[d]| / 127
  For each d in 0..128:
    LUT[d][v] = quantize_i8( centroids[v] * rotatedQuery[d] / scale )

Per FastScan tile (32 KV positions):
  acc[0..32] : two YMM regs holding 32 × i16 = 0
  For d in 0..128:
    codes_lo = load 16 nibbles for dim d (positions 0..15)
    codes_hi = load 16 nibbles for dim d (positions 16..31)
    lut      = load LUT[d]
    acc.lo  += Ssse3.Shuffle(lut, codes_lo)   // saturating i16 add
    acc.hi  += Ssse3.Shuffle(lut, codes_hi)
  For t in 0..32:
    score[t] = (acc[t] * scale) * norm[t]

• Range: 128 dims × i8 ≈ ±16256, fits i16.
• AVX-512 path: VPDPBUSD fuses two-dim int8 multiply-add into i32 accumulators.
• V-path is symmetric — replace per-dim q[d] with per-position attention weight w[t], accumulate per-d instead of per-t.

Files & estimated LOC

File	Change	LOC
src/SharpInference.TurboQuant/FastScan.cs (new)	Block-of-32 packer, i8 LUT builder, AVX2 + AVX-512 K-score kernels, symmetric V-aggregation kernel	~500
src/SharpInference.TurboQuant/TurboQuantOps.cs	Add DequantDot32 / WeightedDecompress32 entry points; keep existing per-block fns for the partial-tile tail	+80
src/SharpInference.Engine/TurboQuantKvCache.cs:65-150	Storage: write FastScan-tiled tiles (32-position groups) directly; keep small per-position buffer only for in-flight partial tile	~150 modified
src/SharpInference.Engine/ForwardPass.cs:994-1031	Replace per-position for (t = 0; t < tqLen; t++) loops with tiled for (tile = 0; tile + 32 ≤ tqLen; tile += 32) + scalar tail	~50 modified
src/SharpInference.Engine/HybridForwardPass.cs:979	Mirror change	~50 modified
tests/SharpInference.Tests.TurboQuant/FastScanTests.cs (new)	Verify tile-scoring matches sum of per-block DequantDot within ~1e-3 (LUT i8 quant is the only drift source)	~150
benchmarks/SharpInference.Bench/TurboQuantBench.cs	New benches: q·KV for ctxLen ∈ {1K, 4K, 8K, 16K}	+60

Total: ~1000 LOC, one new public file, surgical engine edits.

Phased plan

Phase 1 — Kernel only, no engine changes

• Add FastScan.cs with K-score and V-aggregation kernels (AVX2 first; AVX-512 as a follow-up).
• Add tests verifying numerical equivalence with existing DequantDot within tolerance.
• Add standalone benchmark.
• Decision gate: if standalone speedup is <1.5× over existing AVX2 path, stop here.

Phase 2 — Engine integration (K-path)

• Switch TurboQuantKvCache to write tiled layout directly. Per-position layout only retained as staging for the in-flight partial tile (<32 positions).
• Rewrite ForwardPass.cs:995-1002 to call tiled DequantDot32 for full tiles + scalar tail.
• Mirror in HybridForwardPass.cs.
• Perplexity check on a held-out set to confirm no quality regression from i8 LUT quantization.

Phase 3 — V-aggregation path

• Replace the Decompress + weighted-sum loop at ForwardPass.cs:1016-1025 with the symmetric FastScan V-kernel.

Phase 4 — AVX-512 VNNI variant

• VPDPBUSD-based kernel for Sapphire Rapids / Zen 4+.

Gotchas

1. i8 LUT quantization is the only quality-risk surface. TurboQuant centroids are bounded (~±3 after WHT) and rotated queries are unit-ish, so the per-query scale factor should be well-behaved. Fallback if perplexity regresses: i16 LUTs with vpmaddwd accumulation — still much faster than the gather path.
2. CPU/Vulkan-CPU-fallback only. The CUDA path (CudaForwardPass, CudaHybridForwardPass) already has GPU TQ kernels and is out of scope.
3. Partial-tile tail. For tqLen not a multiple of 32, the existing per-position DequantDot loop handles the last <32 positions. The per-position layout doesn't disappear entirely — it's the staging area for the in-flight tile.
4. Expected end-to-end gain. K-scoring + V-aggregation are ~30-50% of TQ-enabled decode at long context. A 2-3× phase-level speedup translates to ~15-30% end-to-end at ctxLen=8K. Bigger context → bigger win.
5. No public API breakage. KvCacheCompressor surface can stay identical; new tiled entry points are additive.

References

• turbovec: https://github.com/RyanCodrai/turbovec (MIT)
• TurboQuant paper: arXiv:2504.19874
• FAISS FastScan: André et al., "Accelerated Nearest Neighbor Search with Quick ADC" (ICMR 2017) and "Quicker ADC" (TPAMI 2021)
• Existing impl: src/SharpInference.TurboQuant/TurboQuantOps.cs, src/SharpInference.Engine/TurboQuantKvCache.cs

[pekkah]

Status update — 2026-05-28

Phases 1, 2, 3, and the TQ-aware batched prefill follow-up have all shipped:

Commit	Phase / scope
7fbdaa7	Phase 1 — FastScan K + V kernels (4-bit + 3-bit, scalar + AVX2, dim 128 & 256) + tile/V-tile packers + tests + standalone bench
6e88dcf	Phase 2 — Engine K-path: TurboQuantKvCache keys switch to FastScan tile + 32-slot staging; new ComputeKScores API; the three CPU forward-pass K-loops collapse to one call. Drive-by fix for a pre-existing CPU+TQ+dense bug (Prefill populated _kvCache but decode read _tqKvCache, causing OOB past the FP32 window).
4e0ba64	Phase 3 — Engine V-path: V-tile storage, ComputeVAggregation with deferred sign-flip + inverse Walsh-Hadamard transform (one IWHT per kv-head instead of tqLength).
c37dfdb	TQ-aware batched prefill: new PrefillCoreTq mirrors PrefillCore but uses ResetTotalLengthForBatchedPrefill between layers (per-layer TQ tile/staging/FP32 state is independent across layers — only the global position counter needs to snap back).
3df3ec3, feb3d2b, 554d596, 18992c1	README perf-table + narrative updates with measured numbers.

Measured impact

Standalone kernel (Ryzen 9 7900X, dim=128, K+V combined per layer/kv-head):

TQ positions	per-block K+V	FastScan K+V	speedup
1 024	479 µs	26 µs	18×
4 096	1 931 µs	98 µs	20×
8 192	3 936 µs	193 µs	20×
16 384	8 216 µs	390 µs	21×

End-to-end (Qwen3-8B Q4_K_M CPU --tq):

context	decode t/s
30	12.0
850	10.9
1 650	10.7
3 250	10.2
6 050	9.4

At 6 K positions the per-block AVX2 baseline would add ~120 ms/token of K+V work, collapsing decode to ~5 t/s — so FastScan delivers roughly 1.9× decode end-to-end at 6 K on this workload.

Phase 4 (AVX-512 VNNI) — blocked on .NET BCL

Hardware-wise Zen 4 (Ryzen 9 7900X here) ships AVX-512 including the VNNI extensions (VPDPBUSD, VPDPWSSD), so the kernel-level win from the fused i8 dot-product is sitting on the table. The blocker is the .NET runtime: as of .NET 10 there is still no System.Runtime.Intrinsics.X86.Avx512Vnni class exposing these instructions — tracked at dotnet/runtime#61817.

Workarounds either go native (P/Invoke into a small C/asm helper, breaks NativeAOT story) or wait for the runtime to land the intrinsics. Neither is "just code it up" work today, and the Phase 1–3 AVX2 path is already ~20× over per-block, so the marginal gain from VNNI on top is small relative to the K vs V transposition speedup already in the bag.

Marking this issue blocked and parking Phase 4 until Avx512Vnni lands in the BCL.

[pekkah]

Closing — shipped in commits c37dfdb (feat: TQ-aware batched prefill (#34 follow-up)) and 3df3ec3 (bench: FastScan K+V at 1K/4K/8K/16K + README perf table (#34)). FastScan kernels are live in the CPU --tq path with the README perf table updated; long-context decode at 6K ctx holds 9.4 t/s vs the ~5 t/s pre-FastScan baseline.