[AlexCheema]

Builds on #2512 (GLM-5.3), now merged into main.

Summary

GLM-5.2/5.3 (glm_moe_dsa) is DeepSeek sparse attention (MLA with a DSA indexer shared across layers) and a 256-expert MoE. This PR speeds up its attention and MoE in exo without changing the weights, the quantization or what the model computes. Every change keeps the float32 paths the model PR needed, and two of them compute in float32 where mlx-lm rounded to bfloat16.

The work is in the GLM family folder (models/glm/): glm_moe_dsa.py, and two new modules, layers.py and kernels.py, with type stubs and tests. Where the Metal kernels can't run (no Metal GPU as the default device), every call goes to mlx-lm's own code.

Prefill

• Dense attention in one product. Up to 8192 keys of context, each head's 192 no-rope and 64 rope dimensions are scored in one 256-wide product with a boolean mask.
  • Before, the rope part was a separate product, materialised as an additive float mask over every [head, query, key] and added to the no-rope scores.
  • This saves a [64, L, N] score tensor and two passes over it.
• The first 2048 queries attend causally. In a prompt's first chunk, a query before position 2048 sees fewer keys than the indexer keeps, so its selection is every earlier key.
  • Those queries go through MLX's fused causal attention kernel, which skips the masked half.
  • Only the later queries take the masked path.
• Sparse attention from 8192 keys. From 8192 keys of context a chunk attends sparsely: each query gathers its 2048 selected keys.
  • mlx-lm starts at 16384 keys. On a Mac Studio (M3 Ultra) sparse attention already overtakes dense masked attention at about 10k keys.
  • The selected keys are gathered with 16-byte vector copies, a Metal kernel 3x faster than mx.take for the [1024, 2048, 576] gather.
  • The latent keys are concatenated once per call instead of once per query block, and the value up-projection runs once for all queries.
  • The score scaling, masking and float32 softmax are compiled into fewer passes (bit-identical).
• Indexer scores in float32, in one kernel. The DSA indexer's per-key score is relu(q·k) times each head's weight, summed over its 32 heads.
  • The new kernel sums them in float32 in one pass, with the mask applied.
  • mlx-lm adds the heads one at a time in bfloat16: 31 rounded additions of [L, N] tensors. That is 5x slower at 16k keys, and the rounding leaves many keys tied at the top-2048 boundary.
  • With float32 sums the selection matches a float32 indexer (see Correctness).

Decode

• Router selection in one kernel. For a router with a multiple of 32 experts (GLM: 256, top 8), one kernel selects them: sigmoid, bias, top-k, normalise, scale. mlx-lm runs a sort of all experts and five element-wise or gather kernels.
• Expert combine in one kernel. The experts' weighted outputs and the shared expert's output are added in one kernel, with mlx-lm's roundings (bit-identical in the unit test). mlx-lm runs a multiply, a sum, a cast and an add.
• Router weight converted once. The float32 router weight is converted once at load, instead of on every call.
• Decode attention as GEMMs. Single-token attention treats the 64 heads as the rows of one product with the shared latent keys (scores, then values), instead of one matrix-vector product per head. This helps batched decode the most.

Where the time goes

Measured on the first 8 layers of mlx-community/GLM-5.3-4bit (real weights), Mac Studio (M3 Ultra), by zeroing the attention or the MLP. Times are for the 8 layers.

	before	after
decode at 512 tokens of context, ms per token: attention / MLP / embedding, head and sampling	2.11 / 2.49 / 0.96	2.08 / 2.36 / 0.96
prefill of 16k tokens, s: attention / MLP	8.58 / 3.26	4.71 / 3.16
prefill of 4k tokens, s: attention / MLP	1.04 / 0.80	0.75 / 0.78

Before and after on the same 8 layers. This is mlx-lm's generation loop, with exo's chunking (4096-token chunks) and exo's batch caches for decode. Median of 2 interleaved rounds, after a warm-up:

prompt / context tokens	512	4,095	8,191	16,383	32,767	65,535
prefill s, before → after	0.190 → 0.188	1.86 → 1.53	4.43 → 3.79	11.83 → 7.90	20.65 → 16.60	44.98 → 36.43
prefill speedup	1.01×	1.22×	1.17×	1.50×	1.24×	1.23×
decode ms/token, before → after	5.56 → 5.39	6.18 → 5.98	6.25 → 6.05	6.38 → 6.18	6.65 → 6.37	7.02 → 6.73
MLX peak GB, before → after	30.4 → 30.3	35.8 → 33.1	41.5 → 37.8	49.2 → 43.8	44.6 → 44.0	44.6 → 44.5

Batched decode at 4k tokens of context: 4 sequences 14.02 → 13.40 ms per step, 8 sequences 23.30 → 21.13 ms.

• Decode is bound by GPU latency, not by the host or by bandwidth. A layer runs about 25 small dependent kernels on top of ~285 MB of weight reads. The MoE kernels remove six of them per MoE layer.
• Prefill of long prompts is mostly attention. Dense masked attention materialises scores over every head, query and key, and the indexer's bfloat16 head loop alone took 44 ms per layer per 4096-token chunk at 16k keys (now 8 ms).

Performance

Full model, mlx-community/GLM-5.3-4bit (418 GB), through exo's API, on M3 Ultra Mac Studios (512 GB), the same machines before and after. The harness perf runner generates 128 tokens per request, with a 3 s cooldown between requests. Builds were interleaved before, after, before, after: 2 rounds, each with a warm-up and then 2 requests per row (2 batches for the concurrent rows). The rounds are pooled. TTFT is the median [min–max].

Single device

prompt tokens	concurrency	TTFT s, before → after	prefill tok/s, before → after	decode tok/s, before → after	speedup prefill / decode	MLX peak GB, before / after
512	1	7.84 [7.42–8.04] → 8.46 [7.57–8.52]	67.5 → 62.5	18.7 → 19.3	0.93× / 1.03×	450.3 / 450.4
4,096	1	26.50 [26.46–26.77] → 23.86 [23.37–24.16]	156.1 → 173.8	17.2 → 17.9	1.11× / 1.04×	456.3 / 453.6
16,384	1	129.67 [128.72–130.25] → 93.39 [93.30–93.84]	126.7 → 176.1	16.9 → 17.6	1.39× / 1.04×	471.8 / 466.2
512	4	16.91 → 16.98	193.8 → 197.1	28.0 → 28.3 (aggregate)	1.02× / 1.01×	471.8 / 466.2
512	8	24.52 → 22.62	193.8 → 197.1	26.3 → 27.6 (aggregate)	1.02× / 1.05×	471.8 / 466.2

• The 512-token row measures the GPU waking up, not the model. After a few seconds of GPU idle, macOS releases the wired residency of the 418 GB of weights. The next request spends most of its TTFT bringing it back, on both builds. The spread between requests is as wide as the gap between the builds.
  • Measured on this build with a 10 s cooldown: 8.45 s [8.36–9.28] with MLX's residency refresh off, as today, and 2.90 s [2.89–2.92] with it set to 500 ms (a separate PR).
  • On the 8-layer model, where this penalty is small, 512-token prefill is the same before and after (0.190 → 0.188 s).
• Concurrent decode is bound by expert weight reads. With 8 requests in a step, each MoE layer reads up to 64 distinct experts instead of 8, so aggregate decode rises only from ~19 to ~28 tok/s on both builds.
• Peak memory goes down at 16k (−5.6 GB): chunks with 8k to 16k keys of context now attend sparsely instead of building dense masked scores.

Two Mac Studios: 2x tensor parallel over RDMA, 2x pipeline parallel over the ring (the same protocol, both configurations in every round, exo's own placement):

config	prompt tokens	TTFT s, before → after	prefill tok/s, before → after	decode tok/s, before → after	speedup prefill / decode	MLX peak GB per node, before / after
2x TP	512	3.36 [3.34–3.42] → 3.30 [3.26–3.37]	301.8 → 310.8	23.1 → 24.0	1.03× / 1.04×	227.2 / 227.2
2x TP	4,096	14.69 [14.61–14.83] → 12.88 [12.82–13.16]	316.6 → 366.3	21.2 → 22.5	1.16× / 1.06×	230.7 / 229.9
2x TP	16,384	73.64 [73.61–73.96] → 59.47 [59.29–59.65]	227.9 → 283.7	20.7 → 21.9	1.24× / 1.06×	243.5 / 240.8
2x PP	512	4.62 [4.48–8.93] → 5.16 [4.36–6.07]	117.7 → 108.0	17.7 → 18.4	0.92× / 1.04×	217.6 / 217.6
2x PP	4,096	26.93 [26.74–27.36] → 23.30 [23.28–23.33]	154.1 → 178.1	16.4 → 17.1	1.16× / 1.04×	223.6 / 220.7
2x PP	16,384	85.79 [84.66–93.68] → 58.00 [57.92–58.11]	199.5 → 285.6	16.0 → 16.7	1.43× / 1.04×	231.1 / 226.1

• 2x PP at 512 tokens has the same wake-up penalty as the single device: the slow requests (6–9 s) come after GPU idle on either build. In the first round, with no slow requests, it was 4.62 s before and 4.36 s after.
• 2x PP peak memory is from the second round of each build. In the first round of the before build, exo's placement split the layers differently (the split follows free memory at placement time), and one node held more layers.
• 2x TP gains less at 16k (1.24×, against 1.39× on one device and 1.43× over 2x PP): part of each tensor-parallel step is the all-reduces between the two nodes, which this PR doesn't change.

Long context, 2x TP, 32,768-token prompt (2 interleaved rounds per build, 2 requests each after a warm-up): TTFT 140.53 s [140.52–140.61] → 124.96 s [124.87–125.05], prefill 236.1 → 266.0 tok/s (1.13×), decode 20.0 → 21.3 tok/s (1.06×), MLX peak 242.2 / 242.6 GB per node.

Correctness

Verdict: PASS. This PR stays as close to the float32 reference as before, and past 2048 tokens it gets slightly closer, because the indexer now sums its scores in float32.

Teacher-forced against transformers in float32 (paired gate)

Setup, the same as #2512's:

• Weights: the first 8 layers of mlx-community/GLM-5.3-4bit (real weights; indexer layers 0, 1, 2 and 6, shared layers 3, 4, 5 and 7).
• Reference: transformers 5.22, float32, on the same weights dequantised.
• Scoring: exo's own loading and 4096-token prefill chunks, with the last 64 tokens of each prompt fed through single-token decode on batch caches.
• Positions: all 31,254, a float32 log-softmax over the reference's top 64 plus a tail bucket.

run	mean KL	p99 KL	top-1	top-1 at margin ≥ 1 nat	gate
transformers bf16 (noise floor, 3 prompts)	6.15e-5	6.8e-4	0.9687	1.0000	–
exo before (#2512)	6.09e-5	7.2e-4	0.9669	0.9998	PASS
exo after	5.66e-5	7.1e-4	0.9691	0.9998	PASS

By prompt, past 2048 tokens (top-1, mean KL):

prompt	before	after
haystack, 3,875 tokens	0.9639, 5.5e-5	0.9721, 5.6e-5
raw text, 5,999 tokens	0.9646, 7.4e-5	0.9663, 6.5e-5
raw text, 19,916 tokens (sparse prefill)	0.9680, 6.2e-5	0.9706, 5.8e-5

exo after against exo before (same positions), mean KL 4.5e-5 to 7.4e-5 and top-1 0.957 to 0.967. That is the size of bf16 rounding noise; the noise floor above is the same size against float32. Where it comes from:

• The float32 indexer scores select different keys at ties. Past 2048 tokens this is the change that moves the model, and it moves it toward the reference.
• The single 256-wide score product in dense prefill and the float32 score sum in decode round differently from the separate bf16 products they replace.
• The router logits keep the same float32 math in a different summation order.
• Bit-identical: prompts below 2048 tokens are bit-identical between the causal split (commit 3) and the commit before it.

Distributed (8 layers, two ranks on one Mac Studio over the TCP ring)

	before	after
2x pipeline, split at layer 4 (a shared layer) vs single device	bit-identical	bit-identical
2x tensor vs single device (mean KL / top-1, 2 prompts)	6.5e-5 to 8.2e-5 / 0.958	5.9e-5 to 6.7e-5 / 0.954 to 0.966
2x tensor vs transformers float32 (mean KL past 2048)	6.5e-5	6.9e-5

Full model, one Mac Studio

mlx-community/GLM-5.3-4bit on one M3 Ultra Mac Studio. Everything below except the teacher-forced check goes through exo's API.

Teacher-forced, before against after (in-process on one Mac Studio, exo's 4096-token prefill chunks and the last 64 tokens through single-token decode, the same four prompts as the 8-layer gate). There is no float32 reference for the full 418 GB model, so this compares the two builds directly:

prompt	below 2048 tokens: mean KL / top-1	past 2048 tokens: mean KL / top-1
raw text, 1,468 tokens	6.5e-3 / 0.965	–
haystack, 3,875 tokens	4.0e-3 / 0.916	3.1e-3 / 0.917
raw text, 5,999 tokens	2.7e-3 / 0.980	4.1e-4 / 0.990
raw text, 19,916 tokens (sparse prefill)	2.4e-3 / 0.966	7.9e-4 / 0.988

At 78 layers, any bfloat16-level change moves the full model this much: #2512 measured 2.8e-3 to 6.3e-3 mean KL below 2048 tokens between origin/main and #2512 on the full model, where only the float32 router, an epsilon and the MLX version differ. On the 8 layers, where a float32 reference exists, the same change moves the model toward the reference (above).

Temperature-1 samples: 17 chat prompts × 6 samples per build through the API, 192 tokens at most, temperature 1, top_p 1, no top-k or min-p. The prompts cover code, math, multilingual, multi-turn, JSON and thinking. Both sample sets are then scored by one fixed scorer, the before build in-process, on the exact prompts the API rendered:

	before	after
sampled tokens	15,732	15,780
mean NLL per token (95% bootstrap CI over sequences)	0.258 [0.225–0.294]	0.263 [0.230–0.299]
NLL quantiles, 50 / 90 / 99%	0.0005 / 0.70 / 3.91	0.0006 / 0.74 / 3.99
expected calibration error (top-1 probability against how often the sampled token was the top-1)	0.004	0.006

Two-sample KS test on the per-token NLL: D = 0.011, p = 0.31. No difference between the two builds' sampling distributions is detectable.

GSM8K (first 100 test items, 2 seeds, the card's sampling, thinking off):

	seed 0	seed 1	mean
before	98/100	99/100	0.985
after	99/100	98/100	0.985

The difference is 0.000, 95% CI [−0.015, +0.015].

End to end on the after build:

• Thinking: 3/3 correct (17 × 23 = 391), finished with stop.
• Tool calls: 3/3 calls of get_weather with valid JSON arguments, finished with tool_calls.
• Concurrency: 8 requests at once, 3 trials, 24/24 correct.
• Long-context needle, with the card's sampling: 4.5k and 16k tokens, 2 needles each, 2 trials. The 16k prompts take the sparse prefill path. 8/8 recalled both the code and the port.

Long-context needle on two Mac Studios (after build), with the card's sampling: 4.5k and 16k tokens, 2 needles each, 2 trials (the second trial reuses exo's prompt cache):

• 2x TP over RDMA: 8/8 recalled both the code and the port.
• 2x PP over the ring: 8/8.

Test plan

• uv run basedpyright (0 errors), uv run ruff check, nix fmt (no changes), and pytest src: 1929 passed, 3 skipped. The image-generation tests were left out because they need torch, and this PR doesn't touch them.
• New tests in models/glm/tests/test_glm_attention_paths.py:
  • Against mlx-lm's own GLM-5.x model with the same weights (tiny config, float32): dense prefill and decode past index_topk, sparse prefill, batched decode with left padding, and the fallback off Metal (bit-identical to mlx-lm).
  • The kernels against plain MLX: the indexer score reduction (with and without a mask), the row gather, the router selection (against group_expert_select) and the expert combine (bit-identical).
  • The router weight is stored in float32.
• The existing pipeline top-k forwarding tests (test_auto_parallel.py) pass with the new attention.
• Correctness and performance runs as above.

🤖 Generated with Claude Code