Sparse decode attention for long contexts (on by default), Qwen3 14B split kernel, Gemma tokenizer fixes (1.4.0) by jkyberneees · Pull Request #10 · BackendStack21/mlx-runner

jkyberneees · 2026-11-06T14:20:25Z

Summary

Sparse decode attention for long contexts, on by default from 5,128 tokens. Also a faster dense decode for Qwen3 14B, and three Gemma fixes (tokenizer and BOS) found while measuring.

Sparse decode attention

What it does. Each full-attention layer keeps a mean key for every block of 4 keys. Each KV head scores those blocks against its query heads. It then attends to:

the first 4 keys;
the latest 512;
its 2,048 best-scoring earlier keys.

Prefill, speculative decoding and draft models stay dense. The design is training-free, based on Latent Sparse Attention (arXiv 2610.42147).

Coverage. Every architecture: Qwen3, Qwen2.5, Llama, Mistral, Phi-3 and Gemma 1–4. Sliding-window layers stay dense.

Speed on an M2 Pro (4-bit, 16 GB), against our own dense decode:

Context	Speed-up
8K	1.06–1.4×
16K	1.2–2.1×
30K	1.4–3.0×

For example, at 30K:

Model	Dense tok/s	Sparse tok/s
Qwen3 0.6B	43.8	103.5
Qwen3 8B	17.8	29.0
Mistral 7B	19.7	31.5

Against mlx-lm 0.32.2 (dense) it is level at 8K on 7–14B models and 1.4–2.4× faster at 30K.

It changes output past 5,128 tokens, which is an explicit exception to the token-identical gate (AGENTS.md is updated). Quality was measured over 512 teacher-forced positions of real text, with KL, text NLL and margin-bucketed agreement:

On most models the perplexity cost is at most about 3%, and sparse decode agrees with dense at over 99% of the positions where dense decode is decisive.
Gemma 4 E2B is clearly lossier: +8% perplexity at 30K, KL 0.07–0.19. Gemma 1.1 is lossier to a lesser degree. The doc says so; turning either off is one switch.

Settings:

App: the inspector's new Long context section (on/off, the long-range key budget, per-KV-head or shared selection), with the sparse layer count in the Model section.
CLI: mlxr run / mlxr serve --sparse <on|off|N>.
Environment: MLXR_SPARSE_ATTN (0/off disables it), MLXR_SPARSE_WINDOW, MLXR_SPARSE_MIN, MLXR_SPARSE_SHARED.

Fallbacks. If the index does not fit in memory, the model loads dense, and so it does if sparse setup fails on the GPU. The engine reports the reason.

Qwen3 14B dense decode, 20–27% faster

Its 5 query heads per KV head had no split-attention kernel. With one, 14B decodes at 17.5 tok/s instead of 14.6 at 8K, and 14.0 instead of 11.0 at 16K. Output is bit-identical. G = 8 (the 32B models) was left out because its split kernel was not bit-identical.

Gemma fixes

Whole-book tokenization now matches transformers on Gemma 1.1–4, Qwen3, Qwen2.5, Llama 3.2, Mistral and Phi-3.5.

Carriage returns. Foundation's JSONSerialization drops a U+FEFF that starts a string, so Gemma's U+FEFF + CR vocabulary entry replaced the one for CR. Every CR got the wrong token.
BPE merge order. BPE now merges in Hugging Face's order: one pair at a time, from a (rank, position) heap. Gemma's whitespace-run merges rank early, so runs of 32 or more spaces were split differently.
BOS on raw prompts. Gemma 4's raw completions (--raw, /v1/completions) now start with BOS. Its tokenizer files never ask for it.

Bench and tests

mlxr bench with MLXR_BENCH_CTX reports, besides dense and sparse speed:
- KL, text NLL and top-1 agreement bucketed by margin;
- several text offsets;
- passkey needles;
- the cost of a verification pass.
MLXR_KERNEL_TIMING=1 times the sparse kernels.
Tests:
- every sparse kernel against a CPU reference;
- full-budget sparse decode bit-identical to dense on every architecture;
- split vs single-kernel bit-identity, including G = 5;
- environment parsing, BOM shielding, Hugging Face merge order, and Gemma reference ids.
swift test: 170 tests in 29 suites pass. The Qwen3 14B and Gemma 4 12B cases need MLXR_TEST_LARGE=1.

Reviews

Two adversarial code reviews ran on the branch, and every validated finding is fixed. The fixes include:

prefill and speculative decoding are always dense;
MLXR_SPARSE_ATTN=off and unrecognised values turn it off;
a sparse setup failure falls back to dense;
the BOM shield is case-insensitive and string-aware;
one macro instantiates both the split and sparse kernels.

Not fixed: the memory estimate still counts index bytes for head shapes without a sparse kernel. No supported model has one.

Known limitations

These are in docs/sparse-decode-attention.md:

Prefill at long contexts is 1.4–2× behind mlx-lm. At 30K it is the bigger cost.
Gemma 4 E2B's sparse quality needs investigation.
Speculative decoding doesn't pay at long context until a multi-query verification kernel exists.

Reference: docs/sparse-decode-attention.md. Released as 1.4.0 (Version.swift and the CHANGELOG are bumped).

🤖 Generated with Claude Code

…piled) Training-free sparse decode attention after Latent Sparse Attention (arXiv 2610.42147), reduced to what pays at this engine's context sizes: every 4 keys of a full-attention layer get a pooled f16 index key; a decode query scores the blocks between the sink and the recent window (sum over heads, one selection for every head), keeps the exact top-p blocks and attends to sink + kept keys + window. No clustering, tree search or trained projections. Prefill and verification stay dense. Off unless MLXR_SPARSE_ATTN is set; the default code paths are unchanged. It changes greedy output, so it cannot ship by default under the project's correctness gates. - Kernels: sparse_index, sparse_score, sparse_select (radix select + stable compaction), attention_sparse (attention_split over a position list; bit-identical to it with the full list). - Transformer: init(directory:contextLength:sparse:), index maintenance after qkv_post, sparse branch of decode attention, `sparse` / `sparseAvailable`; the memory check counts the index. - mlxr bench: MLXR_BENCH_CTX=8192,30000 adds decode after long prompts, dense and sparse, with greedy and teacher-forced agreement (MLXR_BENCH_TEXT for a real-text prompt). - Tests: plan, index, scores, selection (ties, odd thread counts), attention over a list; full-budget sparse decode bit-identical to split decode on the installed models. Written without a Swift toolchain or Metal GPU: not compiled, run or benchmarked. See docs/sparse-decode-attention.md for the checklist. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011Px3ESajhJTwjR33meWGur

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011Px3ESajhJTwjR33meWGur

- KV-shared layers use their source's index only when its KV heads and head size match their own. - sparse_select assigns each thread's run of blocks in (simdgroup, lane) order, the order of its prefix sums. - A draft model always loads dense; the memory check counts the index for the target only. - Doc comment: a full budget is bit-identical to split decode. - docs: review and selection-simulation status. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011Px3ESajhJTwjR33meWGur

Compiled and ran on an M2 Pro (16 GB) with no source changes: the sparse unit tests, the full-budget bit-identical decode on 4 models and the full suite (161 tests, MLXR_SPARSE_ATTN unset) pass; start-up is unchanged. Dense vs sparse decode on Qwen3 0.6B and Mistral 7B at 8K/16K/30K with a real text prompt: 1.11-2.39x faster, teacher-forced agreement 89-100% at the default budget, 95% at 4096+1024 keys. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…panel plan Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Qwen3 14B (40 query / 8 KV heads, head size 128) had no attention_split instantiation, so long-context decode used the single-threadgroup kernel and MLXR_SPARSE_ATTN did nothing on it. Adds G = 5 and 8 for head size 128 (Qwen3 and Qwen2.5 14B / 32B). Split decode stays bit-identical to the single kernel (splitDecodeAttentionIsBitIdentical now covers Qwen3 14B), and so does sparse decode with a full budget. mlxr bench, long-context section (MLXR_BENCH_CTX): - teacher-forced on the text after the prompt (MLXR_BENCH_POSITIONS, default 512): KL(dense||sparse), NLL of the text under both, top-1 agreement by dense decode's top-1 margin; sparse runs first and dense rewrites each position, so the cache stays dense decode's - a dense control for the argmax agreement on dense decode's tokens - MLXR_BENCH_TEXT_OFFSETS: prompts from several places in the text - five passkey needles at 10-90% depth, recalled dense and sparse - the GPU cost of a dense verification pass of T = 1, 2, 5, 8 tokens Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

The index query summed q . k over every head, so one list served all KV heads and a head's important keys could lose to the others'. Now sparse_score writes one score row per KV head (the sum over that head's query heads), sparse_select runs one threadgroup per row, and attention_sparse reads its KV head's list. Attention reads the same bytes; the per-head score maps each KV head to its own lanes, so it costs the same as the shared score (kernel timing at 30K: 216 vs 226 us per layer on Qwen3 0.6B shapes). MLXR_SPARSE_SHARED=1 keeps the shared selection. Qwen3 0.6B, Pride and Prejudice, 512 teacher-forced positions, two regions of the book (shared -> per-head): - 16K: KL 0.035 -> 0.004 and 0.012 -> 0.005 nats; top-1 91.4 -> 96.1%, 93.4 -> 96.5% - 30K: KL 0.033 -> 0.010 and 0.046 -> 0.013 nats; text NLL delta +0.049 -> +0.016 and +0.053 -> +0.012 nats/token; top-1 88.9 -> 93.9%, 88.3 -> 92.2% Decode speed unchanged (30K: 105.7 -> 104.2 tok/s, 16K 130.8 -> 128.2). Tests: per-head and shared scores against a CPU reference (including one lane per KV head and a KV head count that does not divide 32), independent selection rows, per-head lists in attention_sparse, and full-budget decode bit-identical to dense in both modes on every installed model. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Loaded next to the Gemma suite's models it pushes a 16 GB Mac into swap (the full suite stalled). Its split, sparse and prefill checks pass with MLXR_TEST_LARGE=1 swift test --filter <test>. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Records the stronger bench, the Qwen3 14B split kernel (+20% / +27% dense decode at 8K / 16K, bit-identical), per-head selection (better on Qwen3 0.6B and 8B, worse on Mistral 7B at 30K), the rejected kernel experiments with their numbers, the cost of a verification pass (T = 5 is 3.1-3.6x T = 1), and the re-bench of Qwen3 0.6B / 8B / 14B and Mistral 7B at 8K-30K. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…c caveats attentionOverTheList covers the new G = 5 and 8 instantiations (Qwen3 14B / 32B shapes) in the default suite, now that the 14B model tests need MLXR_TEST_LARGE=1. The doc notes the run-to-run exceptions, that the Mistral shared-vs-per-head comparison is a single run, that the larger models have probably memorised the bench text, and the verification cost as measured (3.1-3.6x) rather than as a full re-read per token. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

- On by default from 5,128 tokens of context (below that decode stays dense and outputs are unchanged). MLXR_SPARSE_ATTN=0 turns it off. - App: the chat inspector's Long context section (Sparse attention, Long-range keys, Key selection), applied at the next load like the context length; the Model section shows the layers that run sparse. - CLI: mlxr run / serve --sparse <on|off|N>, and a status line on load. - InferenceEngine.load(..., sparse:); when the index does not fit in memory the model loads with dense decode (LoadedModelInfo.sparseNote says why). - Every architecture: the full-budget bit-identity test now covers Qwen2.5, Gemma 1.1, 2 (soft-capped), 3 and 4 (KV-shared layers) besides Qwen3, Llama, Phi and Mistral (Gemma 4 12B and Qwen3 14B with MLXR_TEST_LARGE=1). Note: this changes greedy output past 5,128 tokens on the default path, an exception to AGENTS.md's token-identical gate made at the owner's request. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…B split kernel Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

mlx-lm 0.32.2 / mlx 0.32.4 on the same M2 Pro, model files and text: decode at 8K-30K (dense; its quantized KV cache at 30K) and prefill, beside our dense and sparse decode. Sparse decode is level at 8K on 7-14B models and 1.4-2.4x faster at 30K; our dense decode is 3-11% behind on 7-14B models; prefill is 1.4-2x slower at 30K. The README's limitations say so. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…s one Without it Gemma's agreement metrics were meaningless (Gemma 1.1 text NLL 15.8 nats, Gemma 4 E2B 5.8-6.2); Gemma 4's tokenizer does not ask for BOS, so the tokenizer's add_bos_token alone is not enough. Speed is unaffected. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Gemma 4's tokenizer files neither set add_bos_token nor add BOS in their post-processor, so raw completions (mlxr run --raw, /v1/completions) and the long-context bench fed it prompts without BOS, and it produced nonsense. A chat template that writes {{ bos_token }} now marks the model as expecting BOS on raw prompts too. Chat was unaffected (the template writes BOS itself). Tests: the template rule, and which installed models start raw prompts with BOS (Gemma 1.1-4, Llama, Mistral) or not (Qwen3, Phi-3.5). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

- Prefill is always dense, including a one-token tail (a regenerate's last prompt token ran sparse, so its first token could differ from a fresh run). forward() with several tokens is dense; a single token is a decode step. - Speculative decoding is dense throughout: the no-draft single step used to run sparse while verification was dense, so output depended on which drafts were accepted. Its output now equals dense decoding. - MLXR_SPARSE_ATTN: off / false / no (and anything not understood) turn it off; only an unset value, 1 / on / true / yes or a key budget turn it on. --sparse keeps MLXR_SPARSE_WINDOW / _MIN / _SHARED (SparseAttention.configured). - A sparse setup failure (pipeline, allocation, thread limit) leaves the model with dense decode instead of failing the load. No index is allocated when the context cannot reach minContext, and the memory check follows. - sparseLayerCount counts KV-shared layers that read their source's index. - The template BOS rule does not override an explicit add_bos_token: false. - G = 8 split / sparse kernels removed: attention_split_4_8 was not bit-identical to attention_4_32 (new kernel test). G = 5 (Qwen3 14B) is. - AGENTS.md states the sparse attention exception to the token-identical gate. - Tests: environment parsing; split vs single-threadgroup decode bitwise for head sizes up to 128, including G = 5. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…rder Two tokenizer bugs found by comparing whole-book tokenizations with the reference (transformers) tokenizer: - JSONSerialization drops a U+FEFF that starts a string, raw or escaped. Gemma lists "\u{FEFF}\r" next to "\r", so that entry took over "\r": every carriage return was encoded as id 141385 instead of 235316 on Gemma 1.1 (text NLL 17 nats on a CRLF book; the reference tokenization gives the same model far lower). U+FEFF in tokenizer.json is now swapped for a private-use stand-in the file does not contain and mapped back after parsing. Files without U+FEFF are untouched (one cheap scan). - BPE merged every occurrence of the best pair per pass (GPT-2's reference code). Hugging Face merges one pair at a time, lowest rank and leftmost first, so a pair the merge creates competes at once. They agree when the merge list is in training order; Gemma's is not (whitespace-run merges rank early), and 32 spaces came out as 16 + 16 instead of 31 + one joining the next word. mergePairs now follows Hugging Face with a (rank, position) heap, O(n log n). Whole-book tokenization now matches the reference on Gemma 1.1, 2, 3 and 4, Qwen3, Qwen2.5, Llama 3.2, Mistral and Phi-3.5. The inspector caption no longer says shared selection suits Mistral (that came from prompts without BOS). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

- Accuracy for Gemma 1.1-4, Llama 3.2, Phi-3.5 and Qwen2.5 beside their speeds (1.06-3.0x). Gemma 4 E2B is the outlier (KL 0.07-0.19, +8% perplexity at 30K) and Gemma 1.1 a lesser one; both are said so. - Mistral rows and the per-head/shared comparison redone with BOS: per-head is level or better on every model compared. The earlier Mistral reversal came from prompts without BOS. - Needle recall: one fewer than dense on Gemma 3 1B at 8K and 16K. - CHANGELOG: the Gemma tokenizer fixes and BOS on Gemma 4 raw prompts. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

- Speculative generations set the model's sparse mode to nil for their duration, so every step is dense: the prompt forward with a one-token suffix ran sparse, and prepared next-step buffers (encoded sparse) were thrown away on every no-draft step (lost CPU/GPU overlap). The per-call dense flag on that step is gone; prompt forwards are always dense. - Transformer.sparseOffReason says why a model asked for sparse attention decodes dense (context below minContext, no eligible layer, or a setup failure on this GPU); the engine's note and mlxr bench show it instead of "no layer of this model supports it". - The memory check takes the SparseAttention and counts the index only when the clamped context reaches minContext; a load checks once and retries without the index only when that fails. - tokenizer.json BOM shielding is one pass that tracks JSON strings: only a U+FEFF that starts a string (raw, or escaped in any case) is replaced, and the stand-in is a private-use scalar the file does not use. Pre-tokenizer patterns and the post-processor's BOS are mapped back too. - One macro instantiates each split shape's attention_split and attention_sparse, so the two lists cannot drift. Whole-book tokenization still matches the reference on all 9 tokenizers. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

cloudflare-workers-and-pages · 2026-11-06T14:20:44Z

🚀 Deploying Preview to Cloudflare 🚀

Preview Deployments by commit

Status	Deployment URL	Commit	Updated (UTC)	See this deployment's details
Build: Failed ❌ View logs ↗		`b65127f`	2026-11-06T15:54:56.902Z	View logs ↗
Build: Failed ❌ View logs ↗		`1b5893a`	2026-11-06T14:20:39.571Z	View logs ↗

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

claude and others added 24 commits November 5, 2026 18:18

Sparse attention doc: instructions for a local agent continuing on a Mac

8a7638c

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011Px3ESajhJTwjR33meWGur

Sparse attention doc: correct measured figures, budget sweep at 16K

f082a18

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Sparse attention doc: Qwen3 8B numbers, fixed-cost breakdown, review …

78f3376

…panel plan Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Opt-in GPU timing of the sparse decode kernels (MLXR_KERNEL_TIMING=1)

94a18d0

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Sparse attention doc: rewrite as reference documentation

65be2fe

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

README and CHANGELOG: sparse decode attention on by default, Qwen3 14…

871d48f

…B split kernel Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Inspector: link the Latent Sparse Attention paper under Long context

51ac7d8

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Landing page: sparse attention at long contexts

2450916

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Version 1.4.0

b65127f

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

jkyberneees changed the title ~~Sparse decode attention for long contexts (on by default), Qwen3 14B split kernel, Gemma tokenizer fixes~~ Sparse decode attention for long contexts (on by default), Qwen3 14B split kernel, Gemma tokenizer fixes (1.4.0) Nov 6, 2026

jkyberneees merged commit 97b86ec into main Nov 6, 2026
4 of 5 checks passed

Navigation Menu

Repository navigation

Sparse decode attention for long contexts (on by default), Qwen3 14B split kernel, Gemma tokenizer fixes (1.4.0) - #10

Sparse decode attention for long contexts (on by default), Qwen3 14B split kernel, Gemma tokenizer fixes (1.4.0)#10
jkyberneees merged 25 commits into
mainfrom
feat/sparse-decode-attention

jkyberneees commented Nov 6, 2026 •

edited

Loading

Uh oh!

cloudflare-workers-and-pages Bot commented Nov 6, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

jkyberneees commented Nov 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary