Quantized (fp16/q8_0) KV cache for the dense CUDA path → unlock long context (Gemma 4 12B first)

[pekkah]

Problem

The dense CUDA path stores the KV cache in fp32 (CudaForwardPass.cs:594-595 allocates _gpuKCache/_gpuVCache as fp32; GpuForwardPass.cs:199 — "K + V, fp32"). KvAppend, Attention, and AttentionSwa all read/write fp32 buffers with sizeof(float) strides. This applies to every dense model on the CUDA path, not just Gemma — fp32 KV is the single biggest barrier to long context on a 12 GB card.

SHARPI_KV_DTYPE already exists but only on the GDN hybrid path (CudaHybridGdnForwardPass.cs:672) and only supports fp32 | bf16 — it is not wired to the dense path and has no q8_0 mode. llama.cpp serves these models at 128K on the same class of card only because --cache-type-k/v q8_0 shrinks KV ~4×.

This is a general backend enhancement

Quantized KV storage + quant-aware append/attention kernels are a dense-CUDA-path capability, not a model feature. Every dense model benefits:

• Qwen3-8B Q4_K_M — KV-bound at long context just like Gemma; gets the same ~4× context headroom.
• Gemma 4 12B QAT — the driving use case (see VRAM math) and the validation target (its SWA ring was never exercised past -c 2048).
• Any future dense model on the CUDA path inherits the knob for free.

Two Gemma-specific behaviors must be preserved (constraints, not new work):

• SWA ring (SwaRingSize/SwaRingHeadroom) — quantized kernels must keep honoring the modulo wrap. Non-SWA models (Qwen3) use the full-context branch unchanged.
• attention_k_eq_v (12B global layers, K=V, store K only) — V's dequant-on-read reads the same quantized K blocks.

VRAM math — the driving case (Gemma 4 12B QAT, 12 GB card, 128K ctx)

48 layers, head_dim 256, ~40 SWA + 8 global; global layers MQA with attention_k_eq_v (K=V, store K only). SWA ring = window + SwaRingHeadroom (4096), capped at ctx.

Component	fp32 (today)	q8_0
Model (Q4_0 bulk + Q6_K tied embd)	~7.0 GB	~7.0 GB
SWA-layer KV (×40, ring)	~3.4 GB	~0.9 GB
Global-layer KV (full 128K, ×8, K=V)	~1.1 GB	~0.3 GB
Total	~11.5 GB	~8.2 GB

At fp32 the model + KV alone nearly fills the card with no room for activation scratch. q8_0 leaves ~3.5 GB of headroom and makes 128K viable.

Scope

1. fp16/q8_0 KV storage on the dense CUDA path (CudaForwardPass.cs / GpuForwardPass.cs). The contained part is allocation; the real cost is the three kernels — KvAppend (quantize K/V on write), Attention / AttentionSwa (dequant-on-read, or dp4a-style q8 dot, inside the online-softmax loop). No q8_0 KV kernel precedent exists — the GDN bf16 path is a straight half-width store, not block-quantized, so it doesn't carry the scale handling. Must stay argmax-stable through the online softmax.
2. Extend SHARPI_KV_DTYPE to the dense path and honor --cache-type-k/v; add a q8_0 mode alongside the existing fp32 | bf16. Preserve the SWA ring and attention_k_eq_v aliasing.
3. Long-context validation (Gemma 4 12B QAT) — the SWA ring (SwaRingSize/SwaRingHeadroom) was validated for E4B in perf(cuda): close remaining Qwen3-8B Q4_K DECODE gap to llama.cpp — non-matvec cost (prefill handled by #167; kernel-efficiency in #149/#152) #162/perf(cuda): chunked batched prefill for Gemma 4 SWA via a real KV ring (#162) #164; the 12B was only run at -c 2048. Verify correctness past the SWA window at 8K / 32K / 128K (needle-in-haystack or argmax parity vs fp32 KV at short ctx).
4. Re-measure the README text-gen row at long context (-c 32768, -c 131072) and record decode degradation vs the 2K number.

Suggested increments

• fp16 KV first — half-width store, no block scales, trivially argmax-stable, already ~2× context. Lands as a checkpoint and de-risks the path.
• q8_0 second — block quantization + scale handling for the full ~4×.

Acceptance

• Dense CUDA path runs ≥32K (ideally 128K) context within 12 GB with q8_0 KV. Gemma 4 12B QAT -g -1 is the gating model; Qwen3-8B validated as the non-SWA case.
• Argmax-stable vs fp32 KV at short context; coherent long-context output. SHARPI_KV_DTYPE bisects.
• README text-gen row updated with a long-context decode figure.

[pekkah]

Progress — bf16 KV landed, 128K on the 12B in 12 GB ✅

Branch feat/bf16-kv-dense-cuda-179 (4 commits):

• Increment 1 — bf16 KV storage + SWA/global bf16 kernels + decode/per-token-prefill dispatch + --kv-type CLI flag. Argmax-stable parity vs fp32 (Qwen3-8B, Gemma E4B, Gemma 12B).
• Increment 1.5a — bf16 batched prefill (≤4096) via the scalar batched bf16 kernels.
• Increment 1.5b — bf16 Tc2 tensor-core flash prefill (templated kernel, fp32 byte-identical), incl. chunked prefill >4096. bf16 prefill is now full-speed at any length for head_dim%64 models (Gemma 256 / Qwen3 128).
• Host-cache fix — the GPU paths allocated a full host KvCache (layers × maxSeqLen × kvDim × 2 floats, ~50 GB on the 12B at 64K) used only for the length counter. Made it bookkeeping-only. This was the hidden ceiling — it OOM'd the host before VRAM did.

Acceptance check (Gemma 4 12B QAT, CUDA -g -1, 12 GB / RTX 4070 Ti)

Context	fp32 KV (default)	bf16 KV (--kv-type bf16)
32K	✅ fits	✅ fits
64K	❌ cudaMalloc failed	✅ fits, coherent
128K	❌ OOM	✅ fits, coherent

bf16 KV + the host-cache fix together take the 12B from a ~32K ceiling to 128K within 12 GB — closing the issue's headline ("≥32K, ideally 128K"). Output is coherent at 128K. SnapKV/TQ compose-guards in place; default stays fp32 (bf16 is opt-in via --kv-type bf16 / SHARPI_KV_DTYPE).

Remaining

• README text-gen row: add the long-context decode figure (steady-state decode deep in context still to be measured; a short-prompt 128K decode sample showed ~21 t/s vs 54 @ 2K — needs proper deep-context measurement; possible large-ctx decode-perf effect to confirm).
• q8_0 KV (Increment 2, ~4×) — now lower urgency since bf16 already reaches 128K; gives headroom for a draft model (GPU draft-MTP speculative decoding for Gemma 4 12B (decode 54 → ~70 t/s) #178) / >128K.
• Tc/half2-flash bf16 thunks — trivial follow-up, only a non-%64 head_dim model past 4096 needs them.
• Auto-bf16 default when the requested ctx won't fit at fp32 — now viable (prefill is full-speed).

[pekkah]

Done — shipped in #184 (merged to master, 8cd8e0c).

All acceptance criteria met:

• ✅ Dense CUDA path runs -c 131072 (128K) within 12 GB with q8_0 KV. Gemma 4 12B QAT -g -1 is coherent at 128K; Qwen3-8B validated as the non-SWA case.
• ✅ Argmax-stable vs fp32 at short context (logit-level top-5 oracle, teacher-forced); coherent long-context output. SHARPI_KV_DTYPE / --kv-type bisects.
• ✅ README 12B text-gen row updated with the long-context decode figure.

What shipped (vs the suggested increments):

• Half-width checkpoint — used bf16 rather than fp16. Strictly better for this: same footprint, fp32's exponent range (no overflow head-room loss), argmax-stable, and it reuses the proven GDN-path bf16 store. (Increments 1 / 1.5a / 1.5b: bf16 storage, scalar batched prefill, Tc2 tensor-core flash incl. chunked >4096.)
• q8_0 — ggml block_q8_0 (32 int8 + fp16 scale, 34 B/block ≈ ¼ fp32). One sharpi_kvload overload + the 6 read kernels templated on <KV>; warp-quantize store kernels. SWA ring + attention_k_eq_v aliasing preserved.
• Host-cache fix — the real ceiling: the GPU passes allocated a full host KV buffer used only for the position counter (~50 GB at 64K on the 12B → managed OOM before VRAM). KvCache.CreateBookkeepingOnly removes it.
• CLI + server — --kv-type fp32|bf16|q8_0 and SharpInferenceServerOptions.KvType (appsettings / SHARPI_KV_DTYPE).

Headline (12B, -g -1 -c 131072): at 128K, bf16 cliffs to ~19 t/s (the tied Q6_K embed table spills to host); q8_0 frees ~1.1 GB more, keeps it resident, holds ~53 t/s — full speed (2.8×), same coherent output.

Validation: q8_0 per-token + batched + chunked>4096 parity (Qwen3-8B, E4B, 12B); bf16 oracle preserved; DTypeInfo.ByteSize(.., Q8_0) unit; 128K CLI A/B.

Follow-ups (auto-narrow default, draft-model headroom #178, Tc/half2-flash q8 thunks) tracked separately.