Repository navigation
perf(models): decide whether chunked GLA prefill becomes the bailing_moe_linear default #1040
Description
Summary
PR #1039 (bailing_moe_linear, the Ling / Ring linear-attention MoE family) ships two evaluations of the same GLA recurrence. The chunked one is 2.6x to 4x faster on prefill and, in isolation, the more accurate of the two, but it is opt-in. This issue is the data-backed decision on whether it becomes the default.
Background
Two paths, same recurrence, both in src/models/bailing_moe_linear.rs:
gla_sequentialis upstream mlx-lm's per-token loop transcribed step for step. It is the default.gla_chunkedis the closed form of that recurrence evaluated over 64-token chunks. It is opt-in behindMLXCEL_BAILING_LINEAR_CHUNKED_PREFILL=1.
Measured on mlx-community/Ring-mini-linear-2.0-4bit on an M1 Ultra, prefill only. Decode is identical on both paths, since both take the same single-step recurrence at L == 1.
| prompt tokens | sequential | chunked |
|---|---|---|
| 101 | 0.29s | 0.07s |
| 512 | 0.75s | 0.28s |
| 2048 | 2.81s | 1.07s |
Chunking is also the more accurate path in isolation, because the intra-chunk sum lands in a matmul accumulator instead of a bf16 running state.
Why it is not already the default. Reassociating the sum changes results by well under one bf16 ulp per layer, but this stack amplifies that. On a 101-token prompt, layer 0's mean absolute activation moves by 0.04%; layer 3 then multiplies its input magnitude by roughly 23x; and the 256-expert top-8 router converts a sub-ulp score difference into a different selected expert. The two paths therefore decode different (not worse, just different) continuations from the same checkpoint. Rather than assume the faster path is fine, PR #1039 shipped the reference's own arithmetic by default and left the promotion to a measurement.
What to measure
- Perplexity on a held-out corpus for both paths on
mlx-community/Ring-mini-linear-2.0-4bit, over enough tokens that the difference is outside run-to-run noise. If chunked is equal or better, the accuracy objection is answered. - Larger sibling. Repeat on
inclusionAI/Ring-flash-linear-2.0if it is reachable, since the amplification behaviour may scale with depth. - Prefill throughput at 512 / 2048 / 8192 tokens: median of repeated runs with dispersion reported. This box carries background load, so a single measurement is not trustworthy.
- Decode throughput unchanged, which it should be by construction. Confirm rather than assume.
Acceptance Criteria
- A decision is recorded in this issue, with the perplexity and throughput numbers behind it.
- If chunked is promoted: it becomes the default; the environment variable is inverted or removed; the
bailing_moe_linearentry indocs/supported-models.mdis updated (it currently documents the present default and the reasoning above); and the module docs insrc/models/bailing_moe_linear.rsare updated to match. - If chunked is not promoted: the reason is recorded and the environment variable stays as an opt-in escape hatch.
Technical Considerations
The default choice is a numerics decision, not a kernel-selection one: both paths are already implemented and tested, so the work here is measurement plus a documentation sweep. Any promotion changes decoded output for this family from the same checkpoint, so it should land as its own change rather than riding along with unrelated model work.
Activity
inureyes commentedon Aug 7, 2026
Decision: promote
gla_chunkedto the defaultThe accuracy objection that kept it opt-in is answered, and in the opposite direction from the one it anticipated. Chunked is not a different-but-equal answer traded for speed; it is measurably the better one.
All numbers on
mlx-community/Ring-mini-linear-2.0-4bit, M1 Ultra,--features metal,accelerate.Perplexity
Teacher-forced,
examples/perplexityovertests/fixtures/wikitext2_excerpt.txt. The harness is deterministic (no sampling), so a repeat of a configuration reproduces its number exactly; that was verified by accident when a shell quoting slip re-ran the default configuration three times and produced85.9565every time. Run-to-run noise is therefore not the thing to guard against here, and the robustness check that matters is varying the configuration.Chunked wins in every configuration. Within a fixed scoring shape the advantage grows with the window (128 to 1024 at one window; 256 to 512 across 32 windows), which is exactly what a bf16 running state compounding its error over a longer recurrence predicts. Absolute perplexity is not comparable across rows, since each scores a different corpus slice and window length; only the within-row comparison is.
That also corrects the framing in this issue's Background. "Well under one bf16 ulp per layer" describes the per-step difference, but the sequential path accumulates in bf16 across the whole prefill, so the end-of-sequence difference is not sub-ulp at all.
Prefill throughput
Wall time of
mlxcel generate -n 1 -t 0 --no-chat-template, median of 5 runs with the min-max spread, prompts taken from the same corpus. Wall includes a common ~0.15s load, so these slightly understate the pure-prefill ratio.Consistent with the 2.6x to 4x in the issue, at the lower end of it. The box carries heavy background load, which shows in the ~512 sequential row's spread; the two larger sizes are tight.
Decode throughput
Confirmed unchanged rather than assumed, since both paths take the same single-step recurrence at
L == 1: median 142.87 tok/s sequential against 143.58 chunked over 5 runs of 64 tokens.Larger sibling
inclusionAI/Ring-flash-linear-2.0was not measured. It is not on this box and pulling it was out of proportion to what it would settle: the mechanism is a per-layer accumulation error, the measured advantage already grows monotonically with sequence length within this checkpoint, and a deeper sibling would be expected to move in the same direction. Recorded as not done rather than assumed.What changed
chunked_prefill_enabled()now defaults to on. The variable inverted rather than being renamed, mirroringMLXCEL_FUSED_MOE:MLXCEL_BAILING_LINEAR_CHUNKED_PREFILL=0(alsofalse/off/no, case-insensitive) restores upstream's sequential recurrence, and a pre-existing=1still selects chunked instead of silently becoming an opt-out. The parsing is split into a purechunked_prefill_enabled_fromso the default is testable without mutating the process environment.chunked_prefill_is_on_unless_explicitly_switched_offpins the default and both spelling sets. That is the assertion that catches the default flipping back.gla_chunkeddoc comment,docs/supported-models.md, and a newdocs/environment-variables.mdrow all carry the measurement and the reason.The cost, stated plainly
A checkpoint decoded here no longer matches the same checkpoint decoded under mlx-lm token for token, because upstream implements the sequential path. That was the original reason for the default and it is a real loss, not a rounding detail. It is worth paying because the measurement says the divergence is toward the better answer, and because this family already could not be gated on token-exactness against mlx-lm: as
docs/supported-models.mdrecords, the reference disagrees with itself between its cached and full-recompute paths, flipping a one-ulp top-2 tie six tokens into a 12-token prompt. Set the variable to0when that comparability is what you need.docs(reports): add technical reports for the #1053/#1054/#1045/#1040 …
feat(models): qualify Qwen3.8-27B on the qwen3_5 path (#1174)
perf(models): make chunked GLA prefill the bailing_moe_linear default (…