Skip to content

perf(cache): route Turbo4Asym decode through dequant-SDPA - #369

Merged
inureyes merged 3 commits into
mainfrom
perf/issue-367-turbo4asym-dequant-sdpa
Jun 19, 2026
Merged

perf(cache): route Turbo4Asym decode through dequant-SDPA#369
inureyes merged 3 commits into
mainfrom
perf/issue-367-turbo4asym-dequant-sdpa

Conversation

@inureyes

@inureyes inureyes commented Jun 19, 2026 •
edited
Loading

Copy link
Copy Markdown
Member

Problem

Turbo4Asym (fp16-K + 4bit-V), the advisory --recommend-quant suggestion for medium/long context, had catastrophically slow decode (0.13-0.24x fp16) because it was hardwired to a lossy sparse-V weighted-sum path, while symmetric Turbo4 and Turbo4Delegated already had a fast dequant-SDPA path and Turbo4Asym had none. Found while running #354. Not a #353 regression (turbo4-asym was equally slow at MLX_MAX_OPS_PER_BUFFER=40 and 1000).

Change

Route Turbo4Asym decode through an exact dequant-then-SDPA path that uses the rotated-codec-basis identity rotate(SDPA(q,k,v)) == SDPA(q,k,rotate(v)) (attention scores depend only on Q/K). The new attention_turbo4_asym_dequant_sdpa (sparse_v.rs) dequantizes V into its WHT-rotated basis via the existing dequantize_v_turbo4_rotated_for_sdpa plus v_rescale, runs native MLX SDPA against the FP16 K (no Q/K rotation since K is already fp16, unlike symmetric Turbo4), and inverse-rotates only the small [B, Hq, Tq, D] output via turbo4_v_inverse_rotate. This skips the per-token inverse Walsh-Hadamard transform a full V dequant pays. The cache method update_and_turbo4_asym_dequant_sdpa_attention (cache.rs) uses v_rescale and mutates cache state identically to the sparse-V path (only self.update), so the packed buffers and offset are unchanged. A default-on flag MLXCEL_TURBO4_ASYM_DEQUANT_SDPA selects the new path; the sparse-V path stays reachable when the flag is off (A/B and fallback). Model dispatch updated in qwen3.rs and llama3.rs (qwen2/qwen2.5 re-export llama3's Attention; these are the only Turbo4Asym call sites). Other modes (Turbo4, Turbo4Delegated, int8, fp16) are untouched.

Result

Exact and ~3x faster. The parity test turbo4_asym_dequant_sdpa_matches_full_dequant_attention compares the new path against the full-dequant reference (update_and_fetch + native SDPA, GQA 4/2 heads) and passes at RMS 0.0 (bound < 5e-3): the rotated-basis result is bit-equal to the exact dequant attention, replacing the lossy sparse-V approximation.

M1 Ultra decode (qwen2.5-7b-4bit, greedy):

context sparse-V (old) this PR int8 fp16
4k 13 39 (0.40x) 75 98
16k ~10 30 (0.43x) 48 70

This raises Turbo4Asym from 0.14x to ~0.40x fp16 (3x), exact. It does not reach int8 / Turbo4Delegated throughput (0.7-0.8x): even in the rotated basis it materializes a full dequant V then runs attention as two passes, and the 4bit PolarQuant + WHT V dequant is ~2x costlier than int8's linear dequant. Reaching int8/delegated speed needs fusing the V dequant into the attention kernel like Turbo4Delegated (which never materializes the full V), tracked as a follow-up in #370.

Validated under the pinned toolchain (rust 1.93.1, matching CI): clippy -D warnings clean (root and mlxcel-core), the parity test passes, fmt clean.

Closes #367.

@inureyes
perf(cache): route Turbo4Asym decode through dequant-SDPA
ce39cbb 
Turbo4Asym (FP16 K + 4-bit V) decode was hardwired to the lossy sparse-V weighted-sum path in the qwen3 / llama3 single-token attention dispatch, the only Turbo mode without a dequant-then-SDPA route. At the default 1e-6 threshold sparse-V skips almost nothing, so it paid the full V dequant plus masking/gather overhead and ran at 0.13-0.24x fp16 (qwen2.5-7b at 4k on M1 Ultra), far below int8 (0.78x), symmetric Turbo4, and Turbo4Delegated (0.80x), which all dequantize transiently and run native SDPA.

Add a dequant-first native SDPA decode path for Turbo4Asym: KVCache::update_and_turbo4_asym_dequant_sdpa_attention slices the FP16 K, dequantizes the packed 4-bit V to FP16 with the same turbo::quant::dequantize_v_turbo4 the update_and_fetch path uses, and runs crate::layers::attention (native MLX SDPA, which handles the GQA head broadcast and the optional mask). This is the exact full-dequant attention replacing the lossy sparse-V approximation, so output is more accurate, not just faster.

Gate it behind MLXCEL_TURBO4_ASYM_DEQUANT_SDPA (turbo4_asym_dequant_sdpa_enabled, default-on via parse_env_default_on), mirroring the symmetric Turbo4 and Turbo4Delegated flags. The model dispatch in qwen3.rs and llama3.rs takes the dequant-SDPA branch when the flag is on and falls back to the sparse-V branch when it is off, keeping sparse-V reachable for A/B and as a fallback. qwen2 / qwen2.5 re-export llama3's Attention, so they are covered by the llama3 change. All other modes (Turbo4, Turbo4Delegated, int8, fp16) are untouched; only a higher-priority Turbo4Asym branch was added ahead of sparse-V.

Add a real-MLX parity test (turbo4_asym_dequant_sdpa_matches_full_dequant_attention) asserting the new path matches the exact update_and_fetch + native SDPA reference within RMS < 5e-3 over 16 decode steps with a grouped-query shape (4 query heads over 2 KV heads); observed max RMS 0.0. The test also drives the sparse-V route on an identical cache and asserts the visible offset stays in lockstep, confirming the V packing/trim is unchanged and only the attention computation differs.

Closes #367
@inureyes inureyes added type:performance Performance improvements priority:medium Medium priority area:core mlxcel-core: MLX FFI, primitives, KV cache, layers area:inference Generation, sampling, decoding (incl. speculative, DRY) status:review Under review labels Jun 19, 2026
inureyes added 2 commits June 19, 2026 14:45
@inureyes
perf(cache): use rotated-codec-basis SDPA for Turbo4Asym decode
b02369b 
The first cut dequantized the full visible V to FP16 before SDPA, which pays the per-token inverse Walsh-Hadamard transform and caps Turbo4Asym decode at ~0.33x fp16. Use the rotated-codec-basis identity instead (rotate(SDPA(q,k,v)) == SDPA(q,k,rotate(v)), since attention scores depend only on Q/K): dequantize V into its rotated basis via the existing dequantize_v_turbo4_rotated_for_sdpa plus v_rescale, run native SDPA against the FP16 K, and inverse-rotate only the small output. This raises Turbo4Asym decode to ~0.40x fp16 (about 3x over the old sparse-V path) and stays exact: the parity test against the full-dequant reference passes within the RMS bound. Reaching int8/delegated throughput (0.7-0.8x) would require fusing the V dequant into the attention kernel like Turbo4Delegated, which is tracked as a follow-up.
@inureyes
docs(cache): replace em dashes in Turbo4Asym doc comments
ec84d59
@inureyes inureyes added the status:done Completed label Jun 19, 2026
@inureyes
inureyes merged commit 7ef0a9d into main Jun 19, 2026
5 checks passed
@inureyes
inureyes deleted the perf/issue-367-turbo4asym-dequant-sdpa branch June 19, 2026 05:58
This was referenced Jun 20, 2026
Merged
Closed
inureyes added a commit that referenced this pull request Jun 20, 2026
@inureyes
fix(cache): align Turbo KV advisor, gates, and docs with measured swe…
cb80c39 
…ep (#376)

* fix(cache): align Turbo KV advisor, gates, and docs with measured sweep

The --recommend-quant KV advisor, the bench_kv_cache.sh acceptance gates, and
docs/turbo-kv-cache.md described an aspirational world where quantized KV is a
decode-speed win. The post-#369 M1 Ultra sweep (4 models x 6 modes x 4k/16k/32k
+ 8k prefill, recorded under benchmarks/turbo_kv/) shows the opposite: every
quantized KV mode decodes slower than fp16. Quantized KV trades cache footprint
for memory, not throughput.

Advisor (kv_cache_advisor.rs):
- Promote int8 and turbo4-delegated as the fast picks (medium leads with int8,
  also turbo4-delegated; long off-allowlist leads with turbo4-delegated, also
  fp16+turbo4). turbo4-asym becomes the exact-K memory pick, not the headline.
- Stop suggesting turbo3-asym anywhere: measured decode 0.02-0.07x of fp16
  (about 0.4 tok/s at 32K) makes it a documented memory-extremis last resort.
- Add Turbo4Delegated to the Walsh-Hadamard set so non-power-of-two head dims
  downgrade it to int8 like the other Turbo codecs.
- Safety invariants unchanged: symmetric turbo4 only for allowlisted families,
  MLA/pure-SSM never get a Turbo mode, KV modes never touch weight dtype.

bench_kv_cache.sh:
- Replace the unreachable "decode >=0.97x fp16" gate with the measured baselines
  as regression backstops, plus the turbo4-asym parity-RMS exactness gate.

docs/turbo-kv-cache.md:
- Add a "Measured speed and the memory trade-off" section with the sweep table.
- Update the recommended-starting-points and advisor tables to match.

The slowdown is inherent, consistent with upstream turboquant_plus: its +22.8%
decode-at-32K figure is vs turbo3's own full-dequant path, not fp16, and turbo3
still sits at 0.93x of int8. The fused kernel that would close the gap is #370.

* docs(cache): add ADR 0002 on Turbo KV split-dequant vs fused (sparse-V study)

Record the turboquant_plus cross-comparison (the #354 follow-up investigation):
why mlxcel's split Turbo KV decode does not reproduce the upstream sparse-V
+22.8% speedup, and why fused V dequant (#370) is the one lever.

Findings, backed by A/B on qwen3-30b-a3b @16k and the paper's exact model
qwen3.5-35b-a3b @32k:
- sparse-V is exact at the default tau=1e-6 (skipped weights are below FP16
  round-off); the earlier "lossy" note was a mischaracterization, not a bug.
- The sparse-V skip buys no decode speed: kernel vs full-dequant graph is +1.4%
  at 32K against the paper's +22.8%, and skipping the entire V weighted-sum
  (tau=1.0) equals full dequant. In a split design (graph scores+softmax + a
  separate V kernel) the V step is a negligible decode fraction.
- The +22.8% lives inside a fused flash-attention; mlxcel does not have one.
  #370 (fuse V dequant into native SDPA) is the only path that captures the
  long-context bandwidth win, a meaningful skip, and no materialize round-trip.

Links the ADR from docs/turbo-kv-cache.md and the ADR index.
inureyes added a commit that referenced this pull request Jun 20, 2026
@inureyes
release: v0.3.2
15b21de 
Audio I/O on the OpenAI /v1/audio/* surface (Whisper speech-to-text #371,
Kokoro-82M text-to-speech #374, endpoint plumbing #368), reasoning_content on
non-streaming chat (#359), pre-M5 decode +8-12% via a hardware-gated op cap
(#360), Turbo4Asym decode rerouted through dequant-SDPA (#369), batched
sampling fixes (#339, #340, #344, #356), and the Turbo KV advisor/docs aligned
to the measured sweep with ADR 0002 (#343, #376).

README: consolidate the v0.3.0 and v0.3.1 highlights under a single v0.3 entry
and add v0.3.2.
@inureyes inureyes mentioned this pull request Jun 21, 2026
Merged
@inureyes inureyes added this to the 0.3 milestone Jun 21, 2026
@inureyes inureyes self-assigned this Aug 31, 2026
@inureyes inureyes removed the status:review Under review label Sep 6, 2026
be-free-snowboy pushed a commit to freestylejs/snowbrain-mlxcel that referenced this pull request Oct 27, 2026
@inureyes
perf(cache): route Turbo4Asym decode through dequant-SDPA (lablup#369)
e05ed6a 
* perf(cache): route Turbo4Asym decode through dequant-SDPA

Turbo4Asym (FP16 K + 4-bit V) decode was hardwired to the lossy sparse-V weighted-sum path in the qwen3 / llama3 single-token attention dispatch, the only Turbo mode without a dequant-then-SDPA route. At the default 1e-6 threshold sparse-V skips almost nothing, so it paid the full V dequant plus masking/gather overhead and ran at 0.13-0.24x fp16 (qwen2.5-7b at 4k on M1 Ultra), far below int8 (0.78x), symmetric Turbo4, and Turbo4Delegated (0.80x), which all dequantize transiently and run native SDPA.

Add a dequant-first native SDPA decode path for Turbo4Asym: KVCache::update_and_turbo4_asym_dequant_sdpa_attention slices the FP16 K, dequantizes the packed 4-bit V to FP16 with the same turbo::quant::dequantize_v_turbo4 the update_and_fetch path uses, and runs crate::layers::attention (native MLX SDPA, which handles the GQA head broadcast and the optional mask). This is the exact full-dequant attention replacing the lossy sparse-V approximation, so output is more accurate, not just faster.

Gate it behind MLXCEL_TURBO4_ASYM_DEQUANT_SDPA (turbo4_asym_dequant_sdpa_enabled, default-on via parse_env_default_on), mirroring the symmetric Turbo4 and Turbo4Delegated flags. The model dispatch in qwen3.rs and llama3.rs takes the dequant-SDPA branch when the flag is on and falls back to the sparse-V branch when it is off, keeping sparse-V reachable for A/B and as a fallback. qwen2 / qwen2.5 re-export llama3's Attention, so they are covered by the llama3 change. All other modes (Turbo4, Turbo4Delegated, int8, fp16) are untouched; only a higher-priority Turbo4Asym branch was added ahead of sparse-V.

Add a real-MLX parity test (turbo4_asym_dequant_sdpa_matches_full_dequant_attention) asserting the new path matches the exact update_and_fetch + native SDPA reference within RMS < 5e-3 over 16 decode steps with a grouped-query shape (4 query heads over 2 KV heads); observed max RMS 0.0. The test also drives the sparse-V route on an identical cache and asserts the visible offset stays in lockstep, confirming the V packing/trim is unchanged and only the attention computation differs.

Closes lablup#367

* perf(cache): use rotated-codec-basis SDPA for Turbo4Asym decode

The first cut dequantized the full visible V to FP16 before SDPA, which pays the per-token inverse Walsh-Hadamard transform and caps Turbo4Asym decode at ~0.33x fp16. Use the rotated-codec-basis identity instead (rotate(SDPA(q,k,v)) == SDPA(q,k,rotate(v)), since attention scores depend only on Q/K): dequantize V into its rotated basis via the existing dequantize_v_turbo4_rotated_for_sdpa plus v_rescale, run native SDPA against the FP16 K, and inverse-rotate only the small output. This raises Turbo4Asym decode to ~0.40x fp16 (about 3x over the old sparse-V path) and stays exact: the parity test against the full-dequant reference passes within the RMS bound. Reaching int8/delegated throughput (0.7-0.8x) would require fusing the V dequant into the attention kernel like Turbo4Delegated, which is tracked as a follow-up.

* docs(cache): replace em dashes in Turbo4Asym doc comments
be-free-snowboy pushed a commit to freestylejs/snowbrain-mlxcel that referenced this pull request Oct 27, 2026
@inureyes
fix(cache): align Turbo KV advisor, gates, and docs with measured swe…
589ec38 
…ep (lablup#376)

* fix(cache): align Turbo KV advisor, gates, and docs with measured sweep

The --recommend-quant KV advisor, the bench_kv_cache.sh acceptance gates, and
docs/turbo-kv-cache.md described an aspirational world where quantized KV is a
decode-speed win. The post-lablup#369 M1 Ultra sweep (4 models x 6 modes x 4k/16k/32k
+ 8k prefill, recorded under benchmarks/turbo_kv/) shows the opposite: every
quantized KV mode decodes slower than fp16. Quantized KV trades cache footprint
for memory, not throughput.

Advisor (kv_cache_advisor.rs):
- Promote int8 and turbo4-delegated as the fast picks (medium leads with int8,
  also turbo4-delegated; long off-allowlist leads with turbo4-delegated, also
  fp16+turbo4). turbo4-asym becomes the exact-K memory pick, not the headline.
- Stop suggesting turbo3-asym anywhere: measured decode 0.02-0.07x of fp16
  (about 0.4 tok/s at 32K) makes it a documented memory-extremis last resort.
- Add Turbo4Delegated to the Walsh-Hadamard set so non-power-of-two head dims
  downgrade it to int8 like the other Turbo codecs.
- Safety invariants unchanged: symmetric turbo4 only for allowlisted families,
  MLA/pure-SSM never get a Turbo mode, KV modes never touch weight dtype.

bench_kv_cache.sh:
- Replace the unreachable "decode >=0.97x fp16" gate with the measured baselines
  as regression backstops, plus the turbo4-asym parity-RMS exactness gate.

docs/turbo-kv-cache.md:
- Add a "Measured speed and the memory trade-off" section with the sweep table.
- Update the recommended-starting-points and advisor tables to match.

The slowdown is inherent, consistent with upstream turboquant_plus: its +22.8%
decode-at-32K figure is vs turbo3's own full-dequant path, not fp16, and turbo3
still sits at 0.93x of int8. The fused kernel that would close the gap is lablup#370.

* docs(cache): add ADR 0002 on Turbo KV split-dequant vs fused (sparse-V study)

Record the turboquant_plus cross-comparison (the lablup#354 follow-up investigation):
why mlxcel's split Turbo KV decode does not reproduce the upstream sparse-V
+22.8% speedup, and why fused V dequant (lablup#370) is the one lever.

Findings, backed by A/B on qwen3-30b-a3b @16k and the paper's exact model
qwen3.5-35b-a3b @32k:
- sparse-V is exact at the default tau=1e-6 (skipped weights are below FP16
  round-off); the earlier "lossy" note was a mischaracterization, not a bug.
- The sparse-V skip buys no decode speed: kernel vs full-dequant graph is +1.4%
  at 32K against the paper's +22.8%, and skipping the entire V weighted-sum
  (tau=1.0) equals full dequant. In a split design (graph scores+softmax + a
  separate V kernel) the V step is a negligible decode fraction.
- The +22.8% lives inside a fused flash-attention; mlxcel does not have one.
  lablup#370 (fuse V dequant into native SDPA) is the only path that captures the
  long-context bandwidth win, a meaningful skip, and no materialize round-trip.

Links the ADR from docs/turbo-kv-cache.md and the ADR index.
be-free-snowboy pushed a commit to freestylejs/snowbrain-mlxcel that referenced this pull request Oct 27, 2026
@inureyes
release: v0.3.2
bcf6b87 
Audio I/O on the OpenAI /v1/audio/* surface (Whisper speech-to-text lablup#371,
Kokoro-82M text-to-speech lablup#374, endpoint plumbing lablup#368), reasoning_content on
non-streaming chat (lablup#359), pre-M5 decode +8-12% via a hardware-gated op cap
(lablup#360), Turbo4Asym decode rerouted through dequant-SDPA (lablup#369), batched
sampling fixes (lablup#339, lablup#340, lablup#344, lablup#356), and the Turbo KV advisor/docs aligned
to the measured sweep with ADR 0002 (lablup#343, lablup#376).

README: consolidate the v0.3.0 and v0.3.1 highlights under a single v0.3 entry
and add v0.3.2.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

@inureyes inureyes

Labels

area:core mlxcel-core: MLX FFI, primitives, KV cache, layers area:inference Generation, sampling, decoding (incl. speculative, DRY) priority:medium Medium priority status:done Completed type:performance Performance improvements

Projects

None yet

Milestone

0.3

Development

Successfully merging this pull request may close these issues.

perf(cache): Turbo4Asym decode uses slow sparse-V path (0.13-0.24x fp16); route through dequant-SDPA

1 participant

@inureyes