bf16 + q8_0 KV cache for the dense CUDA path → 128K on Gemma 4 12B in 12 GB (#179)

[pekkah]

Adds opt-in narrowed KV-cache dtypes (--kv-type bf16|q8_0 / SHARPI_KV_DTYPE) to the dense CUDA forward pass, plus a host-allocation fix that was the real context ceiling. Takes the Gemma 4 12B QAT (-g -1) from a ~32K ceiling to -c 131072 (128K) within 12 GB, coherent. Default KV stays fp32; both narrowed dtypes are opt-in.

Result (Gemma 4 12B QAT Q4_0, RTX 4070 Ti 12 GB, -g -1)

KV dtype	Max ctx in 12 GB	Decode @128k
fp32 (default)	~32K (cudaMalloc-fails at 64K)	—
bf16	128K	18.8 t/s (embed-table cliff)
q8_0	128K	53.3 t/s — full speed (2.8× over bf16)

At 128K the tied ~818 MB Q6_K embed/LM-head table can't stay GPU-resident under bf16's KV (~350 MB free) and spills to host per token → decode cliffs. q8_0's block-quantized cache (~¼ fp32, ~½ bf16) frees ~1.1 GB more, keeps the table resident, and holds full speed. Same coherent output across all three dtypes.

What's in here (squashes #179 increments 1 → 2)

• Increment 1 — bf16 KV storage + SWA/global bf16 kernels, per-token decode/prefill dispatch, --kv-type CLI flag (rejects bf16+TQ / explicit-SnapKV).
• Increment 1.5a/b — bf16 scalar batched prefill (≤4096) + bf16 Tc2 tensor-core flash prefill incl. chunked >4096.
• Host-cache fix — KvCache.CreateBookkeepingOnly: the GPU passes kept the real KV in VRAM but still allocated a full host KV buffer used only for the position counter (~50 GB at 64K on the 12B → managed OutOfMemoryException). This was the hidden ceiling, masking the bf16 win.
• Increment 2 — q8_0 KV (this PR's headline). ggml block_q8_0 layout (32 int8 + fp16 scale = 34 B/block). One sharpi_kvload(const block_q8_0*) overload is the single variation point; the 6 bf16 read kernels were templated on <KV> with bf16/q8_0 thunks (bf16 thunk byte-identical, since sharpi_kvload(unsigned short) == sharpi_bf16_to_fp32). New store kernels quantize per 32-lane warp (amax → d=amax/127 → rintf clamp ±127, mirroring llm_quantize_q8_1). CudaBackend.Allocate now sizes via DTypeInfo.ByteSize (identical for scalar, adds block-quantized).

Validation

• Parity oracles (logit-level top-5 argmax-stable vs an independent fp32 trajectory, teacher-forced so KV dtype is the only variable, per feedback_cross_backend_parity_test):
  • bf16: per-token + batched (Qwen3-8B, E4B) + chunked (E4B) — all green, preserved through the templating.
  • q8_0: per-token + batched (Qwen3-8B, E4B, 12B) + chunked >4096 (E4B and 12B/k_eq_v) — all green.
  • DTypeInfo.ByteSize(.., Q8_0)==34 CPU unit assert.
• CLI A/B at -c 131072 on the 12B: q8_0 53.3 t/s vs bf16 18.8 t/s, both coherent.

Review (pr-review-toolkit)

code-reviewer: no critical/important issues. Addressed: a fail-loud guard for the q8_0 kvDim % 32 invariant (was a latent silent under-allocation; all current models satisfy it), the 12B chunked test above, the ByteSize unit test, and 6 stale kernel-header comments. Knowingly deferred: the _ => fp32 dtype-selector fall-through (future-proofing only — _kvDType is centrally validated to {fp32,bf16,q8_0}) and a pre-existing EnsureImageKernels logging nit.

Notes / follow-ups

• Non-SWA scalar batched q8 kernels + the wave (>4096 global) q8 kernel are wired but only reachable by a non-%64 head_dim model (none on disk) — same latent-but-unexercised status as their bf16 twins.
• Open follow-ups: auto-narrow default when ctx won't fit fp32; Tc/half2-flash q8 thunks; q8_0 draft-model headroom (GPU draft-MTP speculative decoding for Gemma 4 12B (decode 54 → ~70 t/s) #178).

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request introduces support for narrowed KV-cache element types (bf16 and q8_0) in the CUDA backend to significantly reduce VRAM footprint and enable longer context lengths. It adds command-line options, templates the CUDA attention and flash kernels to support the new dtypes, and implements a bookkeeping-only host KvCache to prevent host OOM errors. Additionally, comprehensive parity and long-context tests are added. Feedback on the changes points out a potential NaN generation issue in the q8_0 quantization kernel when dealing with extremely small subnormal floats, which could cause zero values to be incorrectly quantized to -127, and suggests using a small threshold instead of a strict zero check.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

When d is extremely small (e.g., a subnormal float in the range [1.4e-45, 2.94e-39]), 1.f / d will overflow to infinity. For any lane where val is 0.f, this results in 0.f * infinity = NaN. Casting NaN to int yields an indefinite integer value (typically -2147483648), which gets clamped to -127 instead of the expected 0.\n\nTo prevent NaN generation and ensure robust quantization for extremely small values, consider using a small threshold (e.g., 1e-30f) instead of a strict 0.f check.

    float d    = a / 127.f;
    float invd = (d < 1e-30f) ? 0.f : (1.f / d);
    int   q    = (int)rintf(val * invd);

[pekkah]

Added the narrowed-KV option to the API server too (4bdbd83), so it isn't CLI-only:

• SharpInferenceServerOptions.KvType (fp32/bf16/q8_0) — bindable from appsettings.json (SharpInference:KvType), inline Configure, or the SHARPI_KV_DTYPE env override (mirrors the SHARPI_MODEL/SHARPI_BACKEND pattern).
• InferenceEngineLoader forwards it to SHARPI_KV_DTYPE before building the forward pass (only when explicitly set, so an externally-set env var still works); validated at model load.
• Has effect only on the CUDA dense path; ignored by CPU/Vulkan and rejected with TurboQuant / explicit SnapKV — same semantics as the CLI --kv-type.
• Tests: KvType default + appsettings binding (q8_0).

[pekkah]

Good catch on the subnormal-scale edge — applied in 003bacf. A sub-block whose amax is subnormal makes d=amax/127 subnormal, 1/d overflows to +inf, and a zero lane's 0*inf = NaN → (int)NaN clamps to -127. Switched to the d < 1e-30f threshold as suggested. It's far below any real KV scale (d ≈ amax/127 ≈ 1e-3..1e-1), so real blocks are unaffected and the parity oracles are unchanged; an all-near-zero block now correctly quantizes to all-zeros.