feat(kv): fused quantized decode SDPA — make KV codecs live at decode (read quant store directly)

[Pushkinist]

Problem

18 of 25 KV-cache quant codecs are inert at decode. exit_prefill materialises a bf16 K/V decode seed (crates/rmlx-kv-quant/src/kvcache/update.rs:2245-2250); every per-codec update_* early-returns to update_decode_fp16 whenever the bf16 seed is present. Per the cache contract (update.rs:2920-2925), the quantized store is frozen at the prefill length and never read back at decode. Decode SDPA reads bf16.

Consequences, confirmed by the Gemma4 e2b/e4b dual-axis sweep:

• KV quant moves decode TPS within ±1% of none at every context (the quantized store is not on the decode read path).
• KV quant inflates resident KV (k8v4 1.20×, k8v8 1.25×, planar 1.47× vs none at 64k) because the full bf16 seed mirror coexists with the frozen quantized prefix + scales — a quantized global layer can be larger than bf16.

Only Mixed / RotK (fused quantized_matmul) genuinely read the quant store at decode; rot_k_tq4v reads it but the slow way (full-prefix dequant every step, O(seq)). The rest are decode-cosmetic.

Goal

A decode-time fused quantized SDPA path that reads the quantized KV store directly via MLX quantized_matmul (dequant inside the attention matmul), never materialising the full bf16 prefix. This is the only architecture in which a KV codec reduces decode-resident KV and decode KV-read bandwidth.

Scope

• Target the growing attention layers (full-attention / "global" layers; on Gemma4 these are the 7 global layers that grow KV O(ctx)). SWA windowed layers stay bf16 — their rotating ring is already window-bounded and tiny.
• Start with one codec (q8 or q4 global K/V) as the proof of mechanism, behind the existing codec plumbing — not all 25 at once.
• Per-step append = GPU-quantize the single new token's K/V into the quant ring; attention matmul runs over (quantized prefix + recent tokens).

Expected gain

• Memory: global KV stored q4/q8 instead of bf16 → ~2–4× smaller resident KV on the decode window (e2b 64k: 780 MB → ~200–390 MB).
• Bandwidth: decode reads q4/q8 KV instead of bf16 on the global layers → up to ~4× less KV-read → faster long-context decode. The only path to beat bf16-KV backends (mlx-lm) on long context rather than tie.
• Makes the "widest weight × KV quant matrix" a real decode feature instead of a prefill-only / cosmetic one.

References

• Working fused path to model after: Mixed / RotK routing at crates/rmlx-kv-quant/src/kvcache/update.rs:583, fused SDPA at crates/rmlx-kv-quant/src/kvcache/sdpa.rs:250.
• Anti-pattern (do not copy): rot_k_tq4v full-prefix dequant per step (sdpa.rs:269,295) — O(seq), materialises the bf16 prefix.
• Complementary, separate work: casting the none-path f32 global decode seed to bf16 (~74% of the current Gemma4 decode gap, codec-independent).

Acceptance criteria

• At least one codec stores global-layer KV quantized and runs fused quantized decode SDPA (no full-prefix bf16 materialisation).
• Coherence (temp=0 greedy) and a long-context retrieval check (NIAH-style) pass.
• Measured resident-KV reduction AND decode TPS ≥ bf16 at long context on a global-layer-heavy model (Gemma4 e2b/e4b).
• No regression on the SWA windowed path or on none.

Out of scope

• Per-step CPU dequant (the slow rot_k_tq4v pattern).
• SWA windowed layers (already bounded).
• The f32→bf16 global-seed change (tracked separately).

[Pushkinist]

Investigation lead — this may be narrower than "build fused decode from scratch"

A first read of the dispatcher suggests the fused mechanism already exists and the real work is removing a bf16 re-inflation on one specific path. Recorded as a lead to verify, not a confirmed diagnosis:

• The Mixed codec already does what this issue asks on the non-shared-KV path: a quantized K8/V4 decode ring with fused quantized_matmul, no bf16 prefix materialisation (routed at kvcache/update.rs:583, fused SDPA in mixed_quant/sdpa.rs).
• The gap is the Gemma4 shared-KV path: update_and_sdpa_returning_kv routes Mixed through update_and_sdpa_mixed_inner with want_kv=true, and that path keeps a bf16 accumulator (update_decode_fp16) to serve the cross-layer KV share — which re-inflates memory and puts bf16 back on the decode read path for the shared layers.

If this holds, the proof-of-mechanism is closer to: make the shared-KV global path return/serve quantized KV instead of maintaining the bf16 accumulator, reusing the existing Mixed fused decode — rather than implementing fused quantized SDPA anew.

To confirm: trace update_and_sdpa_returning_kv → update_and_sdpa_mixed_inner (want_kv=true) and the Gemma4 global-vs-SWA layer wiring in rmlx-models, and check whether the bf16 accumulator is required by the cross-layer share contract or is an avoidable copy.

[Pushkinist]

Validation — CONFIRMED; true scope clarified (don't build fused SDPA from scratch)

Code validation confirms the body's inert-at-decode finding and the comment's narrowing, and sharpens the scope.

Confirmed (body): the bf16-seed early-return is real — 22 sites of if self.decode_fp16_k.is_some() { return self.update_decode_fp16(...) }, bf16 seed materialised at exit_prefill (update.rs:2245-2250), frozen-store contract verbatim (update.rs:2914-2930). (The exact "18/25" count was not enumerated; the qualitative "vast majority inert" is solid.)

Confirmed (comment lead): the fused mechanism already exists — uses_mixed_path() (Mixed + RotK) dispatches to update_and_sdpa_mixed (kvcache/sdpa.rs:569-572), and mixed_quantized_sdpa runs two quantized_matmul ops directly over the quant codes/scales (mixed_quant/sdpa.rs:123-211) — no bf16 prefix on the non-shared path. This is the fused decode the body asks to "build" — it does not need to be written from scratch.

Confirmed (the gap): on the Gemma4 shared-KV path, update_and_sdpa_returning_kv (sdpa.rs:738) routes Mixed through update_and_sdpa_mixed_inner(..., want_kv=true), which — at decode — additionally rebuilds the full-length bf16 prefix via update_decode_fp16 (sdpa.rs:181-195). That bf16 accumulator co-resident with the quant store is the source of the reported KV inflation.

True minimal scope (narrower than the body, one notch broader than the comment):

1. Route Gemma4 global / full-attention producer layers through Mixed (or one q8/q4 codec) on the shared-KV path.
2. Drop the redundant full-length bf16 accumulator on the want_kv=true decode branch (sdpa.rs:181-195). It is avoidable, not contract-required.
3. The piece the comment understated: the shared-KV consumer layer currently runs plain bf16 SDPA (crates/rmlx-models/src/gemma4/layers/mod.rs:422-460) — it must be taught to attend the quant store (via mixed_quantized_sdpa) or a per-step dequant. Without this, dropping the producer's bf16 only moves the dequant.

What is contract-required: the consumer must receive K/V it can attend (today bf16). Re-pointing it at the quant tuples removes the need for the persistent producer-side bf16 mirror — that double-storage is the avoidable part.

Reuse mixed_quantized_sdpa; do not reimplement. SWA layers stay bf16 (already forced, sdpa.rs:760-772). Open question: whether mixed_quantized_sdpa drops into the consumer without mask/GQA rework.

[Pushkinist]

Implementation findings — mechanism proven, long-context win blocked on a kernel (not merged)

A full implementation slice was built on branch feat/fused-quant-decode-global (commit 1686b83, foundation reference, not merged). It wires the Mixed (K8V4) codec to serve quantized KV on the Gemma4 shared-KV global layers at decode: the producer surfaces live quant 3-tuples (SharedKvOut::MixedQuant), drops the redundant full-length bf16 accumulator on the want_kv decode branch, and the consumer attends the quant store via mixed_quantized_sdpa. Scoped to decode-only (seq==1, empty global mask) to avoid the causal/GQA mask rework. Coherent at temp=0 on e2b/e4b. GPU unit test asserts MixedQuant at decode / Bf16 at prefill.

Result — goal met only in the decode-dominated regime, NOT at long context

• Decode-dominated (short prompt, long generation, ~1540 filled): Mixed 10.5 MB vs none 15.7 MB (−33%), decode 141 vs 138 TPS. The mechanism delivers here.
• Prompt-heavy long-ctx (17k prompt) — the actual target: Mixed 155 MB vs none 112 MB (BIGGER) and decode 30 vs 120 TPS (4× slower). Fails both axes.

Two blockers, neither fixable by the wiring

1. Dead exit_prefill bf16 seed. Mixed's feeds_bf16_k_at_decode() keeps a prompt-sized bf16 mirror resident even though the Mixed SDPA never reads it → prompt-heavy KV inflates instead of shrinking. Dropping it touches shared exit_prefill (Bonsai/non-shared Mixed blast radius) — needs careful gating + Bonsai regression.
2. Per-step mixed_quantized_sdpa < MLX bf16 flash-attention at scale. Two quantized_matmul + softmax + sparse-V per step is slower than MLX's bandwidth-optimal bf16 flash kernel over a large store. The KV-read bandwidth win does NOT convert to decode speed without a dedicated fused flash-decode-over-quant MSL kernel. This is the core remaining work and may be hard to beat MLX flash at all.

Also surfaced

• Spec-decode (forward_hidden_states_shared_kv) bf16-extract is incompatible with the quant share — currently loud-fails on MixedQuant. A graceful bf16 fallback for the drafter is needed before Mixed+spec is usable.

Status / next step to actually finish #45

The path is proven but not shippable as a long-context win. To close #45: (1) a fused flash-decode-over-quant kernel that beats bf16 flash at long ctx, (2) drop the dead exit_prefill bf16 seed for the Mixed path (gated, Bonsai-safe), (3) graceful spec-decode fallback. Foundation branch retained; not merged to avoid a 4× long-ctx decode regression + the spec break. Decode-dominated-only is too niche to justify the shared-code footprint as a default.