feat: TurboQuant+ KV cache compression#16
Open
0xClandestine wants to merge 20 commits into
Open
Conversation
* feat: enable TurboQuant KV compression pipeline for alpha Core infrastructure for TurboQuant+ KV cache compression: - maybeQuantizeKVCache routes turbo schemes via parseTurboScheme() - isDonor flag on KVCache protocol — KV-sharing donors skip conversion - Boundary layer skip targets KV layers only (not GDN/Mamba) - Guard checks ANY cache offset (fixes hybrid architecture activation) - Full preset matrix: turbo2/3/4, turbo0v2/0v3/0v4/0v8, turbo4v2, etc. - Warmup bypassed for turbo configs - Gemma 4 newCache marks donor caches - kvScheme threaded through Evaluate + WiredMemoryUtils callers * feat: model-level TQ integration with compressed decode path Qwen3.5 newCache() creates TurboQuantKVCache directly when kvScheme set: - Phase 1 (prefill): raw fp16 via update(), zero overhead, standard SDPA - Phase 2 (decode): compressRawCache() frees raw buffers, compressedAttention() operates on packed codebook indices — no fp16 dequant buffer AttentionUtils dispatch: L>1 → raw SDPA (prefill), L=1 → compressedAttention. Real GPU memory savings — compressed data is the only storage during decode. * fix: Metal kernel integration + TurboFlash with conditional eval TurboFlash fused kernel re-enabled for L=1 decode. eval() barrier conditional: only for asymmetric (K≠V) or 2-bit K. Symmetric 3/4-bit works without eval — 12% decode regression vs 29% with eval. Profile: eval barrier was 28% of decode time. TQ kernels themselves are fast (0.86ms/layer, faster than SDPA). Bottleneck was sync. Decode (Qwen3.5 9B @ 4k): turbo3/4 77 tok/s (-12%), turbo2/4v2 62 (-29%). Also: 2-bit TurboFlash kernels, unit tests, llama.cpp centroids. * fix: tighten eval barrier for compatibility with fused gate changes After rebasing onto alpha (fused gate_up_proj #66, clipped-swiglu #71), turbo3 on 2B produced NaN. The fused gate changes alter the MLX compute graph enough to trigger the lazy eval fusion bug for 3-bit symmetric. Tighten the eval condition: only turbo4 (K=4,V=4) skips the eval barrier. All other configs (turbo3, turbo2, asymmetric) now force eval after TurboFlash output. Verified: 2B turbo3 PPL 3.0 (was NaN), turbo4 PPL 3.6 (unchanged). * feat: expose turboBoundarySkip as configurable parameter Boundary layer skipping (first N + last N KV layers stay fp16) was hardcoded at 2. Now exposed as GenerateParameters.turboBoundarySkip (default 2, matching llama.cpp mode 7). Set 0 to compress all layers. Threaded through TokenIterator, WiredMemoryUtils, and speculative decoding paths. * cleanup: remove dead N(0,1) codebook A/B test code The N(0,1) centroid path (env var TURBO_USE_N01_CENTROIDS, n01Centroids dict, n01Codebook function) was never called after switching to llama.cpp hardcoded reference centroids. Remove dead code and update docstring. * cleanup: migrate TQ profiling to MLX_BENCH_PROFILE=3 Replace SAM_TURBO_PROFILE/TQ_PROFILE env vars with the shared MLX_BENCH_PROFILE flag at level 3 (per-phase TQ decode profiling, forces eval per phase — invasive). * cleanup: reduce warmup from 2048 to 512 tokens Metal pipelines are warm after ~10 tokens. 2048 was overkill — 512 is enough to trigger all dispatch paths without wasting time. * cleanup: remove dead batch recompression code The batch recompression mechanism (pendingRawKeys/Values, flushPendingEncode, recompressInterval, TURBO_RECOMPRESS_INTERVAL env var) was from the old updateAndDequant path. The current compressedAttention path encodes per-token via encodeNewToken — nothing ever appended to pendingRawKeys.
…-explore#78) The inverse value rotation is already fused into the TurboFlash pass 2 kernel via valRotation parameter. The eval() barrier was protecting against MLX lazy eval fusion with a subsequent matmul(output, rotation) that no longer exists — the rotation happens inside the Metal kernel. Removing the barrier eliminates a CPU-GPU sync per decode step on all non-turbo4 configs. turbo3 decode goes from 63 to 78 tok/s on 9B (+25%), matching turbo4 speed. Verified coherent output and stable PPL across all models (2B, 9B, 35B MoE) up to 32K context. The separated-kernel path (asymmetric configs without TurboFlash) still has its own eval barrier since it does an explicit matmul for inverse rotation. Closes ml-explore#75
TurboFlash kernels only instantiated for specific (kb,vb) combos: (4,4), (4,2), (4,3), (3,2), (3,3), (8,4), (8,2), (8,8). turbo2 symmetric (2,2) and other combos now fall through to the separated score+softmax+value path instead of crashing with "Unable to load kernel". Also reverted Qwen3 newCache — Qwen3 head_dim=80 isn't validated for TurboQuant yet. Only Qwen3.5 uses TurboQuantKVCache. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-Authored-By: tturney@psyguard.ai
Two fixes: 1. Per-layer quantization path matching for Config-I models The sanitize chain was stripping 'language_model.' prefix from config keys, but quantize() uses Swift module paths that include the prefix. Now keeps both forms for reliable lookup. Enables thetom-ai/Qwen3.6-27B-ConfigI-MLX (mixed 4/8-bit). 2. TurboQuant compressed buffer overflow on step boundary Buffer was allocated exactly tokenCount slots during compress transition. The first decode token's encodeNewToken write overflowed when prefill length aligned with the step size. Fix: allocate tokenCount + step slots. Tested on Qwen3.6-27B-ConfigI-MLX: - All TQ schemes (turbo4, turbo4v2, turbo3, turbo3v2) working - turbo4 sym: PPL 1.29 (+11%), decode 20.8 tok/s (-9%), 3.2x compression Co-Authored-By: tturney@psyguard.ai
NemotronH.newCache() now reads kvScheme from GenerateParameters and creates TurboQuantKVCache for attention layers (Mamba layers still get MambaCache). Note: decode regression observed (-65%) — needs investigation. See ekryski#86. Co-Authored-By: tturney@psyguard.ai
…l-explore#87) Closes ml-explore#87. The "No eval barrier needed" comment in compressedAttention's TurboFlash path was wrong on Qwen3.5 nKVH=2 shapes. Without the barrier, MLX lazy-eval fusion of the TurboFlash output with downstream residual/MLP/next-layer operations produces garbage: Qwen3.5-0.8B / 2B / 35B-A3B with turbo4 or turbo4v2 emit `!!!!!` from decode token 1. Numerical A/B vs the separated mseScore + softmax + mseWeightedSum path proves TurboFlash itself is computing the right values within fp32 noise (relDiff ~1e-5, output values match to four decimals). The bug is in MLX's fusion of those values with subsequent ops, not in the Metal kernel. Adding `eval(output)` after TurboFlash fixes coherence. The barrier is gated on `nKVHeads < 4` because: - nKVH=4+ shapes (4B / 9B / 27B / Qwen3.6-27B) don't trip the fusion bug — empirically verified. - A sync eval costs ~40% decode tok/s on 4B at ctx=4096 (~57 → 33). Gating preserves perf where TurboFlash demonstrably works. - asyncEval doesn't break the fusion enough — it must be a sync. Verified post-fix on M1 Max: - 0.8B turbo4v2 ctx=1024: was garbage → coherent, 52 tok/s decode - 2B turbo4 ctx=1024: was garbage → coherent, 48 tok/s decode - 35B-A3B turbo4v2 ctx=1024: was garbage → coherent, 23.4 tok/s decode - 4B turbo4 ctx=4096 (control): coherent, 52 tok/s decode (no meaningful regression vs prior baseline of ~57) The proper long-term fix is in upstream MLX's lazy-eval fusion logic (should not silently corrupt outputs of fused custom kernels) or inside `turbo_flash_pass2` to make the kernel boundary unambiguous to the fusion pass. Out of scope for this repo. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…set tracks writes
The decode path in compressedAttention does:
offset += newKeys.dim(2)
let savedOffset = offset
offset = compressedWriteOffset
encodeNewToken(...)
compressedWriteOffset = offset // expects encodeNewToken to have bumped self.offset
offset = savedOffset
…but encodeNewToken never advanced self.offset, so compressedWriteOffset
was stuck at its post-prefill value forever. Every decode token
overwrote the same slot, leaving slots beyond it zero-initialized.
Reading those zero slots in subsequent attention passes can produce
division-by-zero in the dequant path (norm=0).
Adding `self.offset += numSteps` at the end of encodeNewToken makes
the caller's `compressedWriteOffset = offset` line capture the actual
new write position, matching the assumption already in the call site.
Also drops the unused `useCompressedAttention` field — declared but
never read anywhere; the decode path always goes through the
compressed-domain Metal kernel.
Verified on Qwen3.5-4B 4bit turbo4 ctx=1024 (was coherent before, still
coherent). Note: this fix alone does NOT resolve ml-explore#87 — the Qwen3.5-2B /
0.8B / 35B-A3B garbage-output bug has a separate root cause that
manifests as NaN at the second full-attention layer of decode token 1
even with offsets correct.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…-using models (ml-explore#85) GPT-OSS uses attention sinks, which TurboQuantKVCache.compressedAttention does not currently support. Pre-fix, --kv turbo* on GPT-OSS-20B silently produced garbage output and *increased* KV bytes vs fp16 (67MB vs 27MB at ctx=1024) because maybeQuantizeKVCache rewrapped RotatingKVCache entries with TurboQuantKVCache regardless of whether the model's attention path could route through compressedAttention. This change adds a clean opt-out: - LanguageModel: new var supportsTurboQuantization: Bool, default true via extension. Models with attention paths incompatible with TurboQuantKVCache (e.g. sinks) override to false. - GPTOSSModel: returns false. newCache also prints a one-shot warning when the caller asked for --kv turbo* so the fallback isn't silent. - TokenIterator (× 2 inits) + SpeculativeTokenIterator + the two WiredMemoryUtils.prefillOnly call sites + the budget estimator all gate parameters.kvScheme on model.supportsTurboQuantization, passing nil to maybeQuantizeKVCache when the model opts out. Speculative decoding requires both main and draft models to opt in. Verified post-fix: GPT-OSS-20B --kv turbo4 → KV 27MB (matches --kv none), coherent output, decode 71 tok/s (matches --kv none baseline). Closes ekryski#85. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…l-explore#84/ml-explore#88) Two unconditional eval(rotatedOutput) barriers were sitting after TurboQuantKernelOps.mseWeightedSum in the rawKeyMode path (line 1343) and the separated fallback path (line 1468). Both were defensive precautions added before ml-explore#87 was diagnosed; they ran on every decode step regardless of shape, contributing to the decode regression on Qwen3.5/3.6 (ml-explore#84) and the prefill regression on turbo8* (ml-explore#88). Drop-and-test confirmed: - Removing the rawKeyMode + separated-path barriers leaves all common configs coherent (those paths only fire on (kb,vb) combos without a TurboFlash kernel, which is rare). - The targeted gated barrier at line 1406 (gated on nKVHeads < 4) IS load-bearing — Qwen3.5-2B turbo4 produces `Here!!!!!!!!!!` without it. Restored with an updated comment that references ml-explore#92 (upstream MLX graph-compile fusion bug, the actual root cause). Effect: shapes with nKVH ≥ 4 (Qwen3.5-4B/9B/27B/3.6-27B) lose two unconditional sync evals per decode step. Small-nKVH shapes still pay the gated-barrier cost until ml-explore#92 lands an upstream fix that makes fast::CustomKernel outputs opaque to the graph fusion pass. Refs ekryski#84 ekryski#88 ekryski#92. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…009 α) Replaces TurboQuantKVCache's compressed-domain Metal attention with the classic dequant-to-FP16 + MLXFast.scaledDotProductAttention(... sinks:) path as the default. The compressed kernels remain reachable as opt-in β via TurboQuantKVCache.useCompressedAttention = true. Why: the compressed-domain decode path was shipping ~30% of `--kv none` on alpha (49 vs 166 tok/s on Qwen3.5-2B, 19 vs 97 on Gemma 4 E2B), and producing garbage on GPT-OSS-20B because compressedAttention had no sinks support. α restores the architecture from c977dae and matches llama.cpp's TurboQuant ports (TheTom Metal, spiritbuun CUDA), which both achieve 97-99% of baseline at zero quality cost. Verified on M1 Max, 4-bit weights, ctx=1024, 36/36 runs coherent: GPT-OSS-20B (sinks active): 93-96% of --kv none across all 8 turbo configs (turbo4/3/4v2/4v3/3v2/8/8v2/8v4) — was producing garbage before this change Gemma 4 E2B: 99-107% of --kv none across all 8 configs Qwen3.5-2B: 81-99% across all 8 configs (turbo8 hits 98.8%) Qwen3.5-0.8B: 85-105% across all 8 configs Closes ml-explore#85 (GPT-OSS sinks via native SDPA path), closes ml-explore#87 (Qwen3.5 small-nKVH garbage — graph-fuser bug only affected the β kernels), unblocks ml-explore#92 (β kept on the gated path with its existing eval barrier). Changes: Libraries/MLXLMCommon/AttentionUtils.swift - add `sinks: MLXArray? = nil` parameter - turbo decode (L=1) routes through updateAndDequant → SDPA(... sinks:) → inverseRotateOutput - β path gated behind useCompressedAttention=true; fatalErrors if combined with sinks - affine quantized path also fatalErrors on sinks Libraries/MLXLMCommon/TurboQuantKVCache.swift - add public useCompressedAttention flag (default false) + init arg - rewrite class-level docstring with α/β tradeoff and prefill-only state caveat - expand updateAndDequant docstring (rotate-Q / SDPA / inverse-rotate contract) - β compressedAttention left intact behind the gate (barriers stay until ml-explore#92 is isolated upstream) Libraries/MLXLLM/Models/GPTOSS.swift - collapse AttentionBlock.callAsFunction to a single attentionWithCacheUpdate(... sinks:) call - remove supportsTurboQuantization=false (sinks no longer require opt-out — reverts ee42b55) - remove kvScheme=turbo* warn-and-fallback in newCache Libraries/MLXLMCommon/LanguageModel.swift - update supportsTurboQuantization docstring to reflect that sinks no longer require opt-out Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
… Part 0) GPT-OSS-20B turbo decode produced garbage on summarization prompts ≥124 tokens — coherent for ~4 generated tokens, then collapsed into ellipses/dots. Spec 009 α verification matrix surfaced the regression across every (turbo*, ctx) cell that hits the sliding-attention layers. Four interlocking issues, all in TurboQuantKVCache: 1. `makeMask` fell through to `BaseKVCache.makeMask`, which ignores `windowSize` for n>1 and emits `.causal`. Sliding layers got a non-windowed causal mask after the prefill→turbo conversion. Override added that mirrors `RotatingKVCache.makeMask` semantics — windowed causal for prefill chunks, rolled mask for L=1 decode after the buffer wraps. 2. `update()` was `KVCacheSimple`-style: appended without trimming to `rotatingMaxSize`. For multi-chunk prefill on sliding layers (or prefill via TurboQuantKVCache after multi-turn conversion) the raw buffer would grow unbounded. Added a sliding-aware concat-and-trim branch that keeps the most-recent `maxSz` tokens. 3. `loadRawKV` clamped `offset = min(originalOffset, bufferTokens)`, losing absolute sequence position. RoPE on subsequent decode tokens then used the wrong index whenever the source `RotatingKVCache` had already wrapped (originalOffset > bufferTokens). Now `offset` preserves the original sequence position; internal buffer slicing remains gated by `min(offset, buffer.dim(2))`. 4. (additional, beyond spec) `updateAndDequant`'s rotating write logic reset `rotatingIdx = bufferTokens % maxSz` on every call, then used `(rotatingIdx - S + maxSz) % maxSz` as the write position. After the buffer first fills (offset == maxSz), this collapses to writing slot `maxSz-1` on every subsequent decode step — overwriting the most-recent token instead of evicting the oldest. This matches the observed pattern of "first ~4 decoded tokens coherent, then garbage" and is the proximate cause of GPT-OSS turbo failure at ctx=128. Fixed by tracking `rotatingIdx` as a true next-free-slot pointer: initialized once at the prefill→decode transition, advanced by S per write, wrapped to 0 on overflow. Mirrors `RotatingKVCache.updateInPlace` for keep=0. Verification: - Clean macOS build, no warnings on touched code. - Regression coverage deferred to the spec-009 verification matrix (4 models × 9 KV configs × 4 ctx points), which the user runs out of band — loading GPT-OSS-20B in unit tests is impractical. Refs: specs/010-turbo-kv-followups.md (Part 0) Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
8-bit turbo configs paid a ~1–4s codec-init cliff on first prefill window,
charged to the prefill metric. Root cause: `referenceCentroids128` only
covered bits ∈ {2, 3, 4}; (128, 8) fell through to the runtime
`generateCentroids` k-means (32768 grid × 256 levels × 100 iter ≈ 838M
inner-loop iterations of pure Swift).
Fix: pre-compute the (dim=128, bits=8) Lloyd-Max centroids and adjacent-pair
midpoints offline using the same algorithm the runtime path uses, paste
into the static reference tables. The existing `scaledCentroids` /
`scaledMidpoints` per-dim scale factor (sqrt(128/dim)) means a single
d=128 entry covers every standard model dim (64, 80, 96, 128).
Verification (M1 Max, summarization @ ctx=128):
| Model | KV | pre-fix | post-fix |
|---|---|---|---|
| Qwen3.5-0.8B | turbo8 | 27.2 tok/s | 729.3 tok/s (≈turbo4 735.5) |
| Qwen3.5-2B | turbo8 | 26.3 tok/s | 527.6 tok/s (≈turbo4 532.9) |
~27× / ~20× speedups; both within 1% of the matching bits=4 cell — passes
the "within 10% of turbo4" acceptance criterion. Output samples remain
coherent (Gatsby/Fitzgerald references in both pre- and post-fix runs).
Audit notes (A-2): `ensureCentroidsPopulated` and the `precomputed` cache
are dead code in the current codebase — no caller invokes them.
`codebook(dim:bits:)` and `boundaries(dim:bits:)` go directly through
`scaledCentroids` (table hit) or `generateCentroids` (cliff path), with no
lazy-populate intermediary. Pre-baking the table (A-1) therefore eliminates
the cliff for every standard dim without a separate warm-up step. Leaving
the dead code for a follow-up cleanup PR — out of scope for this fix.
Refs: specs/010-turbo-kv-followups.md (Part A)
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
`ensureCentroidsPopulated`, `precomputed`, `centroidLock`, `lazyDims`, and `lazyBits` had no callers — `codebook(dim:bits:)` and `boundaries(dim:bits:)` go directly through `scaledCentroids` (table hit) or `generateCentroids` (runtime k-means fallback), with no lazy-populate intermediary. Noted while implementing spec 010 Part A; deferred from that commit and cleaned up here. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
First turbo decode step paid ~80ms of MLX kernel JIT cost inside TTFT —
the codec's rotation matmul triggered shape-specialized kernel
compilation on its first dispatch (in `updateAndDequant`'s
prefill→decode transition). A second run in the same process paid no
JIT and matched no-quant prefill speed exactly, confirming the gap was
entirely first-call JIT.
Fix: hand the model's `headDim` to `TurboQuantKVCache.init` and
eagerly construct the codec(s) there (instead of waiting for first
decode). Codec construction now runs a small warmup matmul against the
rotation matrix and `eval`s it, forcing kernel compilation during model
load — outside the TTFT measurement window.
Two warmup shapes:
- `[1, 1, 1, dim]` — per-step decode rotation of one new token
- `[1, 8, 128, dim]` — one-shot rotation of prefill batch in transition
MLX dispatches different specialized matmul kernels by inner-batch dim;
warming both shapes covers the kernel families used in transition and
per-step decode.
Wired into the two models that construct `TurboQuantKVCache` directly
in `newCache(parameters:)`:
- Qwen3.5 (Qwen35.swift) — covers Qwen3.5-{0.8B,2B,4B,9B,27B,A3B/MoE}
- NemotronH
Verification (M1 Max, summarization @ ctx=128, median of 3+):
| Model | KV | pre-fix | post-fix |
|--------------|--------|-----------:|------------:|
| Qwen3.5-0.8B | turbo8 | 728 tok/s | ~935 tok/s |
| Qwen3.5-0.8B | TTFT | 156ms | ~118ms |
| Qwen3.5-2B | turbo8 | 528 tok/s | ~535 tok/s |
Best on small models (~30% prefill / ~24% TTFT); marginal on
Qwen3.5-2B where the per-shape JIT cost is amortized over more layers
and real compute dominates anyway.
Scope notes:
- The conversion-path models (Gemma 4 dense, GPT-OSS) construct
`TurboQuantKVCache` inside `maybeQuantizeKVCache` at the end of
step(0), not at model load — eager init there would just shift the
cost from step(1) to TTFT. Left untouched on that path with a
comment explaining why.
- The remaining gap to no-quant on a fully-warm second run goes to
zero, confirming MLX kernel JIT is the only meaningful overhead. A
more aggressive close requires either knowing the exact production
shape at codec construction or running a real warmup decode through
the cache — both larger API changes deferred.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…ption B) Today's α path always rotates K/V at the prefill→decode transition (and on every per-token append) so the workspace lives in the rotated space the compressed codebook expects. But α never actually populates that compressed buffer — only β (`useCompressedAttention=true`) does. With β off, the rotation is dead weight: SDPA on raw K/V is mathematically identical to SDPA on rotated→inverse-rotated K/V (the orthogonal rotation cancels in the SDPA product). Add a `rotationActive` flag (default `false`; `true` when β is on). When `false`: - updateAndDequant transition copies raw K/V into the workspace verbatim, no Π_K / Π_V matmul. - per-token append writes raw new K/V, no rotation. - prepareQueries / inverseRotateOutput pass through. - ensureCodecs is deferred — the rotation matrix and its prewarm matmul are never paid. Closes the prefill-side JIT gap that codec init opened on small models: the warmup matmul fires only when β is opted into. For ctx ≤ rotatingMaxSize (the common bench case) α now matches `--kv none` decode tok/s within JIT noise (smoke on qwen35-0.8b ctx=1024: none 170.7 / turbo4v2 172.7 / turbo8 172.5; prefill 2548 / 2258 / 2270 — the remaining prefill gap is update()'s concat-based growth, separate optimization). β path is untouched: `useCompressedAttention=true` flips rotationActive at init and the existing rotated/encode flow runs unchanged. State serialization in lazy α matches today's behaviour: `state` getter returned `[]` for any α post-transition cache (compressed buffers were never populated). No regression there; B-2 (mid-decode wrap-activation) is the follow-up that will plumb `state` through properly. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…d-domain decode (spec 010 B-5) Plumbs the spec 009 β path through the user-facing kvScheme parser and the benchmark harness so β can be exercised side-by-side with α from a single flag flip. Parser: - `parseTurboScheme` now returns `useCompressedAttention` alongside bits. Strip the optional `-compact` suffix, pass the flag through to `TurboQuantKVCache.init(useCompressedAttention:)`. - `turbo4` → α (default, dequant-FP16 + standard SDPA) - `turbo4-compact` → β (compressed-domain Metal kernels) - `turbo4v2-compact` → β with asymmetric kb=4, vb=2 Wiring: - Qwen35Model.newCache, NemotronHModel.newCache: thread the compact flag through to the cache constructor. - maybeQuantizeKVCache (the conversion path used by GPT-OSS / Gemma 4): same. - BenchEnv.kvConfig: replace the static switch with `parseTurboScheme` so any scheme — including `-compact` — works without a harness change. Sinks fallback: - α handles attention sinks natively (`MLXFast.scaledDotProductAttention(... sinks:)`); β kernels don't yet (spec 010 B-1 follow-up). Previously `attentionWithCacheUpdate` fatalError'd when β met sinks. Now silently falls through to α — α's rotation is already active in β-configured caches (init flips `rotationActive=true`), so the math is identical to pure α; we just skip the compressed-kernel call. GPT-OSS + `turbo*-compact` is no longer a crash. Smoke (qwen35-0.8b, ctx=1024, summarization): - turbo4v2-compact: prefill 2138 tok/s, decode 48 tok/s, KV cache 10MB (vs α's turbo4v2 KV cache 13MB and decode 172 tok/s) β decode is ~3.5× slower than α today — the WHT+turbo_flash path doesn't match SDPA throughput. β is shipped as a usable opt-in for memory-first long-context users; closing the speed gap is the spec 010 B-1/B-4 follow-up. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…e (β) Threads the `sinks:` parameter from `attentionWithCacheUpdate` down through `compressedAttention` → `turboFlashAttention[Causal]` → `MLXFast.turboFlashPass2[Fused]` so models with attention sinks (GPT-OSS) can use the β path natively instead of falling back to α. - AttentionUtils: drops the sinks-fallback gate. Sinks now flow through `compressedAttention(... sinks:)` directly. The 'fall back to α when sinks are non-nil' compromise is replaced by 'β kernel folds the sink into its pass2 softmax'. - TurboQuantKVCache.compressedAttention: gains `sinks: MLXArray? = nil` parameter, forwarded to the TurboFlashAttention call sites (decode L=1 and causal prefill paths). Separated path and rawKeyMode path do not yet plumb sinks (uncommon configs; α handles them). - TurboQuantKernels: `turboFlashAttention` / `turboFlashAttentionCausal` / `dispatchFlashPass2` gain `nQHeads` / `L` / `sinks` so the kernel can recover head index from the flat q_idx. Caveat: β compressedAttention does NOT yet honor the sliding-window mask required by GPT-OSS sliding layers and Gemma 4 sliding/global. Those configs run β kernels but attend to all tokens (not just the window), producing garbage on long prompts. Sliding-window support in β is the next follow-up; for now, sliding models should stay on α (the default). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…resolved
The pass2 sinks fold (spec 010 B-1) is wired end-to-end across mlx /
mlx-c / mlx-swift / mlx-swift-lm and produces correct output on the math
(sum exp(s-M)·V / (sum exp(s-M) + exp(sink-M)) under M = max(max_s, sink)),
but on GPT-OSS-20B turbo*-compact at any context size the visible output
is still incoherent ('<|channel|>analysis<|message|>We have a long,?…?').
α (turbo* without -compact) on the same model + same prompt produces
coherent text via MLXFast SDPA's native sinks support, so the gap is
between my pass2 fold and MLX's SDPA reference implementation — not a
sliding-window or graph-fuser issue (smoke at ctx=128 fails, where the
buffer doesn't wrap; eval barrier doesn't change behavior).
Until the bug is found, route sinks-using models through α as before.
The infrastructure stays in place so the path is callable end-to-end
(both for debugging via `MLX_TURBO_FORCE_BETA_SINKS=1` and for the
eventual fix). β with sinks plumbing covers:
- Pass2 metal kernel: fold sink into global softmax (turbo_quant.metal)
- C++ wrapper: optional sinks parameter (mlx fast.cpp/h)
- C ABI: sinks via mlx_array.ctx == nullptr sentinel (mlx-c fast.h/cpp)
- Swift wrapper: optional `sinks: MLXArray?` (MLXFast.swift)
- Compressed dispatch: forward sinks to pass2 (TurboQuantKernels.swift)
- Cache attention: forward sinks (TurboQuantKVCache.compressedAttention)
Also adds the unconditional eval barrier when sinks are folded (separate
from the nKVH<4 ml-explore#92 workaround) — left in place as it cost nothing on
the working non-sinks path and keeps the door open for the fix.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
GPT-OSS uses attention sinks. Turbo's compressed-domain decode path (path B) doesn't handle sinks coherently yet; tracked as the four-repo β-path stack in ml-explore#130. The α-path runs sinks-bearing attention without crashing (dequant V → standard SDPA fallback) but at valueBits=2 the output is silently incoherent — Harmony channel markers (<|channel|>, <|message|>) never appear, so downstream parsers see empty content. Two things conspire to make this *no longer reproduce* on alpha post-spec-006 (see PR ml-explore#166): 1. `maybeQuantizeKVCache` was deleted (PR 4); caches are now constructed up-front in each model's `newCache(parameters:)`. 2. GPT-OSS's `newCache` calls `makeAttentionCache(...)`, which handles `.affine` and `.none` but falls through to StandardKVCache for `.turbo` — so turbo silently doesn't reach GPT-OSS even when requested. Verified on the May 5 smoke log: `--kv turbo4v2` on GPT-OSS-20B emits coherent Harmony output (`<|channel|>analysis<|message|>...`). The bug described in ml-explore#171 was a v3 codebase artifact. This override makes the intent explicit at the model level so the contract survives future refactors of `makeAttentionCache` or GPT-OSS's `newCache`. Also lets `effectiveAlgorithm` (used for memory- budget estimation in WiredMemoryUtils) skip turbo when sizing wired memory tickets for GPT-OSS. Closes ml-explore#171. Path-B sinks β stack (ml-explore#130) tracked separately. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…xplore#185 for model-factory dispatch The `TurboQuantizedKVCache` rotating-window code path (`rotatingMaxSize` / `rotatingIdx` machinery, lines 750–1364 in `TurboQuantKVCache.swift`) has been implemented for a while but never exposed through the public factories — `makeKVCache(scheme:eviction:)` trapped on `.turbo + .window` with a precondition, and `makeAttentionCache(parameters:maxSize:keep:)` silently fell through to a raw `StandardKVCache` regardless of `compressionAlgorithm`. This commit unblocks the standalone factory and adds direct-construction test coverage. The model-factory dispatch is intentionally held back behind a one-model-specific bug (Gemma 4) tracked in ml-explore#185. Changes - `KVCacheTypes.swift` `makeKVCache(scheme:eviction:)`: drop the `.turbo + .window` precondition trap. The `.turbo` arm now reads `eviction` and constructs `TurboQuantizedKVCache(... maxSize:)` when windowed. Updated docstring. - `KVCacheTypes.swift` `makeAttentionCache(parameters:maxSize:keep:)`: unchanged in behaviour — still falls through to `StandardKVCache(maxSize:keep:)` for windowed turbo. Docstring updated to point at ml-explore#185 and explain why model dispatch stays on the safe path until the Gemma 4 interaction is fixed. - `KVCacheTests.swift` `testMakeKVCacheFactoryAllSchemes`: extended to cover `.turbo + .window(size:)` → `TurboQuantizedKVCache(maxSize:)`, including the symmetric-bits + asymmetric-bits variants. - `KVCacheTests.swift` `testMakeAttentionCacheTurboWindowedSafePath` (new): asserts the fall-through behaviour is preserved (so the safe path doesn't silently re-enable itself), with an inline note pointing at ml-explore#185. - `documentation/kv-cache.md` + `skills/mlx-swift-lm/references/kv-cache.md` algorithm matrix rows for `TurboQuantizedKVCache`: now state that windowed support is available via direct `makeKVCache(scheme:eviction:)` construction (working on Mistral 3 / Ministral 3 / Gemma 3) but model- factory dispatch through `makeAttentionCache` is held back pending ml-explore#185 (Gemma 4 specific). Constraints / footguns section updated to drop the "precondition trap" claim. Investigation summary (full repro + hypotheses in ml-explore#185) Direct-construction smoke runs via `--kv turbo4v2`: | Model | Result | |---|---| | Mistral 3 (Ministral 3 3B) | ✅ coherent: _"Hello! I'm an AI assistant—no name, just here to help…"_ | | Gemma 3 (Gemma 3 1B) | ✅ coherent: _"Hello there! I'm Gemma, and I'm here to help you with various tasks!"_ | | Qwen 3.5 (Qwen 3.5 0.8B, no sliding-window) | ✅ coherent: _"I am Qwen3.5, a highly capable language model…"_ | | **Gemma 4 (Gemma 4 E2B)** | ❌ incoherent: _"Except's. **SERVApa\\middleware narrowed\\immediately?**"_ | The general windowed-turbo code path is not broken — Gemma 4 has a specific cache-construction interaction (KV-shared layers + mixed sliding/full-attention layer types per `Gemma4.newCache(parameters:)`) that doesn't compose with the rotating compressed buffer. Hypotheses for the Gemma 4 root cause (donor / shared-cache `peek()` math; mid- prefill rotation under TurboQuant's two-phase prefill→decode transition; mixed compressed / uncompressed layer attention) all documented in ml-explore#185 with file-line references and suggested investigation steps. Verification - `swift build -c release` clean. - `swift test -c release --filter KVCache` — **58/58** pass, including the new dispatch-correctness tests. - Smoke (`./scripts/benchmark.sh --model gemma4-e2b --quant 4bit --kv turbo4v2 --method simple`): coherent output unchanged from alpha (StandardKVCache fallback path preserved). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Comment on lines
+362
to
+363
| @Test("StandardKVCache default init matches legacy StandardKVCache shape (typealias identity)") | ||
| func testStandardKVCacheUnboundedTypealiasIdentity() async throws { |
There was a problem hiding this comment.
can we get some CI actions setup that run these tests on every PR?
Member
Author
There was a problem hiding this comment.
unrelated to this pr, that's general housekeeping
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Cherry-picks all TurboQuant+ related commits from ekryski/mlx-swift-lm (
ek/turbo-kv-beta+alpha), preserving original authorship (Tom Turney / Eric Kryski).TurboQuantKVCache— 8-bit Lloyd-Max KV compression (Algorithm 1) with pre-baked centroids, WHT rotation, and MetalTurboFlashfused kernelcompressedAttentionpath inAttentionUtils(prefill → raw SDPA, decode → compressed-domain Metal kernels)supportsTurboQuantizationopt-out protocol for GPT-OSS (attention sinks)-compactCLI suffix, lazy-α rotation skip until buffer wraps, rotation JIT pre-warmKVCacheTypes.swift—makeKVCache(scheme:eviction:)now supports windowed-turbo constructionCommits (20 total, chronological)
20bb390bd94127131012f314cffdc4f732d82e1443daf230bb2c822d078de711be42f9914769c764669a91db58db8d14f7171afd86c50367-compactsuffix (spec 010 B-5)0c4f58d46fadc5dc1ded51b4b1b6Conflict resolution notes
previousKVsmechanism.Test plan
swift build -c release--kv turbo3/--kv turbo4on Qwen3.5 9B — verify PPL and decode tok/s--kv turbo3on NemotronH — verify no regression--kv turbo4on GPT-OSS — verify fallback warning + coherent output--kv turbo3 --kv-size 4096on Qwen3.5TurboQuantKernelTestssuite passes