[AlexCheema]

Builds on #2520 (Tencent Hy3), now merged into main. Its first commit is #2521 (pipeline ranks build their decode graph while the input is on its way), carried here so that the numbers include it; it drops out once #2521 lands. The second commit adds typing stubs for the nn.RoPE attributes the new code reads (the mlx.core.fast and nn.RMSNorm stubs are already in main).

Summary

A Hy3 decode step at batch 1 ran about 40 dependent GPU kernels per layer for ~140 MB of weights per layer. On a Mac Studio (M3 Ultra) that was 0.33 ms per layer, against a weight-streaming floor of ~0.17 ms. This PR runs the same arithmetic in fewer kernels. Under tensor parallelism over Ethernet, it also removes half of the all-reduces.

On one Mac Studio (M3 Ultra), with the 4-bit card, through exo's API:

• decode +15 % at 512 tokens of context (34.0 → 39.2 tok/s), +14 % at 4k (30.3 → 34.4) and +10 % at 16k (23.6 → 26.0);
• prefill +5 % at 4k (385 → 405 tok/s); time to first token at 512 tokens 3.42 → 2.99 s;
• aggregate decode with 4 and 8 concurrent requests +8 % and +4 %;
• memory unchanged.

On two Mac Studios joined by 1 Gb Ethernet:

• 2x pipeline: decode +22 %, +19 % and +15 % at 512, 4k and 16k tokens (31.0 → 37.9 tok/s at 512).
• 2x tensor: decode +59 % to +64 %, and prefill +60 % to +79 % (time to first token at 16k 286 → 160 s), from halving the all-reduces.
• Tensor parallelism over Ethernet stays slow for this model (8 tok/s); pipeline parallelism, or a single 256 GB machine, is the way to run it there.

On a single device and under pipeline parallelism, every change computes exactly what #2520 computes. The fused Metal kernels run on M1, M3 and M4 GPUs, the generations whose hardware ran them bit for bit. M2 and M5 GPUs use the separate ops until verified on that hardware. The stacked dense gate/up still runs there; on M5 that rests on MLX's dispatch rules, not on a run on M5 hardware. The stacked expert gathers run on M1 to M4. Other GPUs and MLX releases run #2520's code.

• Log-probs: logits and log-probs are the same bits at every prompt length and batch size checked. At full depth, all 4,960 teacher-forced positions of Add Tencent Hy3 #2520's correctness report are identical, so Add Tencent Hy3 #2520's paired gate against the float32 reference carries over unchanged.
• Sampling: with the same seed, temperature-1 sampling draws the same tokens.
• Unchanged: the weights, the quantization, the float32 router Add Tencent Hy3 #2520 added, and the model's math.

All the code is in the hy_v3 family folder; there are no MLX or mlx-lm changes. The kernels are mx.fast.metal_kernels that follow the approach of exo's fused kernels for Flint. Off Metal (CUDA, CPU), and wherever a kernel does not apply, #2520's code runs as written.

Decode: fewer, fused kernels (models/hy_v3/kernels.py, models/hy_v3/fused.py, models/hy_v3/model.py)

• The residual add with the next RMSNorm in one kernel.
  • After attention, it also writes the normed output widened to float32, which is the router's input (exact).
  • After the MoE block, the MoE combination and the residual add go into the same kernel as the next layer's input norm. Each layer hands the next one on its rank its normed input.
• The per-head query and key RMSNorms with RoPE, written straight into the layout attention reads, with the values copied alongside: 4 kernels to 1.
• The router after its logits, in one kernel:
  • the sigmoid and the expert bias;
  • the top 8 in exactly the order argpartition returns them on the GPU. MLX's merge sort is stable: highest first, ties to the lower index. The kernel ranks every expert against all the others;
  • the normalisation, summed left to right as MLX's small row reduction does, and the scaling factor.
  • Its float32 scales and biases are converted once instead of at every call (bf16 to float32 is exact).
• The MoE combination: 5 kernels to 1.
  • The weighted sum of the 8 experts in float32, summed in the order MLX's column reduction uses, with no fused multiply-add since MLX rounds the products first.
  • Then the rounding, plus the shared expert and the residual stream.
• The experts' gate and up projections in one gather.
  • They are stacked along their output rows once, at the first call. The modules then point into the stack, so it costs no extra memory.
  • Every gathered row is one token-expert pair. Which kernel MLX picks (its matrix-vector gather, or from 96 tokens the sorted gather) depends on the number of pairs and of experts, and the stack changes neither.
• SwiGLU inside the 4-bit down projection.
  • One kernel computes mlx-lm's swiglu for the 16 values each thread reads, and runs MLX's qmv_fast reduction on them unchanged.
  • It is used where MLX itself runs that kernel: the experts' down gathers whenever MLX takes the matrix-vector gather, and the shared expert's and dense layer's down projections for one row.
• Stacked gate and up in the dense layer and the shared expert: one matmul below MLX's matrix-vector batch limit; above it, each projection runs on its slice of the stack.

Tensor parallelism over the TCP ring (models/hy_v3/sharding.py)

• Over MLX's ring backend (Ethernet) an all-reduce costs about a millisecond, so the per-layer collectives set the speed.
• Attention and the dense first layer now stay whole on every rank there. Each layer then has one all-reduce (the MoE's) instead of two: 79 per token instead of 160.
• The cost: each rank reads all the attention weights (+1.7 GB per rank) and keeps the whole KV cache.
• Over RDMA (jaccl), attention stays split by heads, as in Add Tencent Hy3 #2520.

Pipeline parallelism (sharding/pipeline.py, #2521): a non-first rank builds its decode graph while the previous rank's activation is on its way.

Not changed here:

• Decode attention is MLX's. It already shares each KV load across the 8 query heads of a group from 8k keys.
• Prefill is mostly MLX's 4-bit gathered and dense matmuls.

Performance

Setup

• Hardware: two Mac Studios (M3 Ultra, 256 GB), macOS 27.0.1, joined only by 1 Gb Ethernet: tensor and pipeline parallelism run over MLX's ring backend (TCP). Single-device rows are on one of the two.
• Weights: XavierLocalAI/Hy3-4bit (affine 4-bit, group 64; 8-bit router).
• Software: before = Add Tencent Hy3 #2520's head; after = this branch. MLX 0.32.4 and mlx-lm 6b572ab as pinned.

Method: the exo correctness harness's perf runner (exo_harness.perf), as in #2520.

• Requests: one exo node per build. The runner places the model, then sends streaming /bench/chat/completions requests: prefix cache off, EOS banned, so every request generates 128 tokens, with the card's sampling defaults (temperature 0.9, top_p 1).
• Repetitions: the builds alternate on the same machines (before, after, before, after). Each session has a warm-up request and 2 timed runs per row, so the table pools 4 runs per build (8 requests per build for the concurrency rows).
• TTFT is client-side. Prefill and decode tok/s are exo's generation_stats.

placement	prompt tokens	concurrent requests	TTFT s (before → after)	prefill tok/s (before → after, ×)	decode tok/s per request (before → after, ×) [after min-max]	aggregate decode tok/s (before → after)
single device	512	1	3.42 → 2.99	157 → 182 (×1.16)	34.0 → 39.2 (×1.15) [39.1-39.2]	34.3 → 39.5
single device	4,096	1	10.88 → 10.32	385 → 405 (×1.05)	30.3 → 34.4 (×1.14) [34.2-34.5]	30.5 → 34.7
single device	16,384	1	48.14 → 46.81	343 → 352 (×1.03)	23.6 → 26.0 (×1.10) [25.9-26.1]	23.8 → 26.2
single device	512	4	7.78 → 7.13	401 → 416 (×1.04)	17.0 → 18.5 (×1.09)	56.9 → 61.3
single device	512	8	11.41 → 10.64	401 → 416 (×1.04)	9.6 → 10.0 (×1.04)	55.5 → 57.6
2x pipeline, ring (Ethernet)	512	1	2.16 → 2.07	266 → 273 (×1.03)	31.0 → 37.9 (×1.22) [37.8-37.9]	31.2 → 38.2
2x pipeline, ring (Ethernet)	4,096	1	11.27 → 10.97	372 → 381 (×1.02)	28.2 → 33.5 (×1.19) [33.5-33.6]	28.4 → 33.8
2x pipeline, ring (Ethernet)	16,384	1	30.60 → 30.12	543 → 553 (×1.02)	22.2 → 25.5 (×1.15) [25.5-25.5]	22.4 → 25.7
2x pipeline, ring (Ethernet)	512	4	6.58 → 6.31	380 → 393 (×1.04)	16.1 → 17.8 (×1.11)	53.8 → 59.2
2x pipeline, ring (Ethernet)	512	8	10.29 → 9.92	380 → 393 (×1.04)	9.2 → 9.6 (×1.05)	53.0 → 55.4
2x tensor, ring (Ethernet)	512	1	8.61 → 5.41	64 → 103 (×1.60)	5.1 → 8.0 (×1.59) [7.9-8.1]	5.1 → 8.1
2x tensor, ring (Ethernet)	4,096	1	68.28 → 38.30	61 → 108 (×1.79)	4.7 → 7.6 (×1.64) [7.5-7.7]	4.7 → 7.7
2x tensor, ring (Ethernet)	16,384	1	286.2 → 160.5	57 → 102 (×1.78)	4.1 → 6.6 (×1.59)	4.2 → 6.6
2x tensor, ring (Ethernet)	512	4	32.33 → 18.74	69 → 122 (×1.76)	3.7 → 5.6 (×1.51)	11.9 → 18.7
2x tensor, ring (Ethernet)	512	8	51.87 → 30.08	70 → 124 (×1.78)	2.2 → 3.4 (×1.52)	12.0 → 19.1

Runs per build: 4 for single-request rows (2 sessions x 2), 8 requests for the concurrency rows; for 2x tensor at 16,384 tokens, one run per build (each request takes minutes over 1 Gb Ethernet).

Memory. MLX's active memory after loading is unchanged on a single device and under pipeline parallelism (in-process runs: 166.0 GB single, 82.3 / 84.3 GB per pipeline rank, about half the model each, as in #2520). Under tensor parallelism over the ring, each rank holds the whole attention: 85.1 GB per rank instead of 83.3 GB, and 92.6 GB instead of 88.3 GB at the peak with a 14.7k-token prompt (the whole KV cache). exo's batch engine reports MLX's peak since the process started (loading and the first request's one-off weight stacking included): single 179.1 → 179.7 GB at 512 tokens and 186.1 → 185.9 GB at 16k, pipeline 89.2 → 89.8 GB per node, tensor 90.3 → 92.4 GB per node.

• Spread: for one request, decode varied by at most 0.65 tok/s (2 %) within a row, in both builds and every placement; the speedups are far outside it.
• Time to first token at 512 tokens includes re-wiring the model's memory after the runner's 3 s pause between requests: macOS drops the residency of idle wired memory, and the next request pays ~7 ms per GB. That is the same for both builds and is not addressed here.
• Concurrency rows: aggregate decode is all generated tokens over the time from the batch's first token to its last, as the harness reports it. With 8 requests most of a step reads the ~55 different experts the 8 tokens choose per layer, which this PR does not change.

Where the time goes

One Mac Studio (M3 Ultra, 96 GB), the first 8 layers of the 4-bit card: the dense layer and 7 MoE layers, real weights.

• Loop: a bare decode loop on mlx-lm's caches, with asynchronous evaluation of the next step and a temperature-1 sampler.
• Ablation: each part is the time it adds to a step, that is, the step with that part left out subtracted from the whole step (so rows overlap a little).

part (per MoE layer, 512 tokens of context)	before	after
attention (norm, q/k/v, q/k norms, RoPE, cache, SDPA, output projection, residual)	103 us	96 us
routed experts with the router (logits, selection, gathers, combination)	183 us	130 us
· the router alone (logits and selection)	34 us	23 us
shared expert (runs alongside the router)	29 us	30 us
whole step, 8 layers + LM head and sampler, 512 tokens of context	3.12 ms	2.73 ms
the same at 4,096 tokens of context	3.42 ms	2.99 ms
graph nodes per step that are not views	357	176

• Routed experts: they now take what their weights take to read. 8 experts x 3 projections x 3.5 MB = 85 MB per layer at ~800 GB/s is 106 us, and they measure 107 us without the router.
• The rest of a layer is mostly the attention projections (42 MB) and the latency of the kernels in between.
• Host side: building a step in Python takes longer (0.37 vs 0.21 ms for 8 layers: custom kernels cost more to issue) but stays well under the GPU's time, so it overlaps.
• Prefill of a 4,096-token chunk (8 layers) takes 1.353 s before and 1.331 s after. It is mostly MLX's gathered and dense 4-bit matmuls (the routed experts ~80 ms and attention ~52 ms per layer), which this PR leaves alone.

Correctness

Verdict: the same model, bit for bit, on a single device and under pipeline parallelism. Every comparison of log-probs, logits and sampled tokens below between the two builds is exact equality, not a tolerance. Tensor parallelism over the TCP ring changes its layout (attention is no longer all-reduced) and stays at bf16 noise from a single device, closer than #2520's layout.

1. Teacher-forced log-probs at full depth

#2520's correctness report scored 29 prompts (4,960 positions, up to 14.7k tokens) with exo in process on the 4-bit card, against Hugging Face transformers on the same weights. This branch was run the same way on the same kind of machine (a Mac Studio, M3 Ultra), teacher-forced on the same tokens:

• 4,960 of 4,960 positions are identical to Add Tencent Hy3 #2520's run: the top-64 log-probs and the tail mass at every position;
• greedy decoding gives the same tokens on all 29 prompts.

So every full-depth number in #2520's report holds for this branch unchanged:

• exo bf16 vs the float32 reference: mean KL 0.0119, against transformers' own bf16 at 0.0102;
• top-1 where the reference is at least 1 nat confident: 0.9964, including the narrow miss Add Tencent Hy3 #2520 analyses;
• per-layer drift at transformers' bf16 level;
• greedy decoding under the reference.

2. Chunk lengths, step sizes and batch sizes

MLX picks its kernels by shape, so a fixed prompt length could miss a deviation at one of its switch points:

• Its quantized matmul switches from the matrix-vector kernel (qmv, then qmv_wide from 2 rows on M3 and later) to matrix kernels at a batch limit that depends on the shapes and the GPU.
• Its expert gather switches to the sorted kernel at 16 index pairs and 4 per expert: 96 tokens for Hy3's 192 experts, top 8.
• mlx-lm sorts the pairs from 64 (8 tokens).

So, on the first 8 layers of the 4-bit card (real weights), #2520's code was run against this branch, comparing all logits at every position as bits:

• prefill of every length from 1 to 320 tokens, plus 383-385, 409-411, 511-513, 639-641, 819-821, 1023-1025, 2047-2049, 4096 and 4097 tokens, from an empty cache;
• lengths 1 to 64 and the same thresholds after 700 cached tokens;
• steps of 1 to 9 tokens after a 64-token prompt at batch sizes 1 to 9, and the step after each;
• through exo's batch engine: 1 to 9 concurrent greedy requests (tokens and log-probs) and 4 seeded temperature-1 requests.

592 of 592 logit cases and 49 of 49 engine cases are identical. Re-run after the GPU-generation change: 592 of 592 identical on M3 Ultra and M4 Max hardware, and with MLX dispatching as on M1, M2 and M4 Max.

3. Distributed placements

The first 8 layers, two ranks on one Mac Studio over the ring, with the harness's 29 prompts (5,612 teacher-forced positions) run through exo's loading, sharding and chunked prefill:

placement	this branch vs #2520	vs a single device
single device	5,612 / 5,612 positions identical, 29 / 29 greedy continuations identical
2x pipeline	5,612 / 5,612 identical, 29 / 29 greedy identical	bit-identical, as with #2520. The one exception is the 6k-token haystack, which the pipeline prefills in other chunks, the same way in both builds
2x tensor, ring	new layout (attention whole on each rank)	teacher-forced on the single device's tokens (28 prompts, 5,293 positions): mean KL 1.9e-4 (#2520's layout: 2.8e-4), top-1 0.955 (0.953), 1.000 where the single device is at least 1 nat confident (1.000)

At full depth, on the two Mac Studios over Ethernet, with the harness's two ranks and #2520's 29 prompts (4,960 positions, the same tokens as in section 1):

placement	check	result
2x pipeline, this branch	vs this branch on a single device, both teacher-forced in 512-token chunks	4,960 / 4,960 positions identical
2x tensor, #2520's layout	vs a single device	mean KL 0.0125, top-1 0.967, 0.9957 where the single device is at least 1 nat confident
2x tensor, this branch	vs a single device	mean KL 0.0114, top-1 0.969, 0.9974
2x tensor, #2520's layout	paired gate vs the float32 reference (transformers bf16: 0.0102, 0.9978)	FAIL: mean KL 0.0132 (paired excess 3.0e-3 > 2e-3), 0.9969
2x tensor, this branch	paired gate vs the float32 reference	PASS: mean KL 0.0107 (paired excess 0.5e-3), 0.9978

• Why 512-token chunks for the pipeline: the harness's pipeline loop has every rank wait for each chunk's logits. With 4,096-token chunks over Ethernet, that wait hit macOS's GPU command-buffer timeout on the 6k-token prompt, the same limitation Add Tencent Hy3 #2520 reports. exo's own pipelined prefill does not wait that way.
• Memory per rank after loading the model: pipeline 82.3 and 84.3 GB, as in Add Tencent Hy3 #2520. Tensor 85.1 GB, against 83.3 GB with attention split, because attention is whole on each rank. Peaks with the 14.7k-token prompt: tensor 92.6 GB (88.3 GB with attention split), pipeline 85.4 / 87.4 GB.

4. Temperature 1

(a) Distribution check.

• Sampling: 12 harness prompts x 4 seeds x up to 96 tokens through exo's API, one request at a time, at temperature 1.0, top_p 1, no top_k or min_p, with each sampled token's log-prob and the top 20 at its position.
• Model: the full 4-bit model on one Mac Studio, before and after.
• Result: 48 of 48 sequences sampled the same tokens in both builds, so the per-token NLL distributions are identical (KS distance 0):

samples from	sequences	tokens	NLL per token, mean (95 % CI by sequence)
before	48	4,538	0.118 [0.087, 0.148]
after	48	4,538	0.118 [0.087, 0.148]

Calibration is identical in both builds, and calibrated: each position's top-20 probabilities, binned, against how often that token was the one sampled.

probability	mean p / how often sampled (n), both builds
0.01 - 0.05	0.025 / 0.021 (719)
0.05 - 0.10	0.072 / 0.089 (237)
0.10 - 0.20	0.146 / 0.159 (245)
0.20 - 0.30	0.243 / 0.168 (119)
0.30 - 0.50	0.414 / 0.443 (316)
0.50 - 0.70	0.595 / 0.567 (224)
0.70 - 0.90	0.816 / 0.814 (226)
0.90 - 1.00	0.996 / 0.994 (4,001)

(b) GSM8K at temperature 1 with the card's sampling defaults (temperature 0.9, top_p 1; the template's default, answering without reasoning), through the API. 150 test items x 2 seeds, 8 requests in flight, so the batch generator decodes them together; full model, one Mac Studio:

build	items x seeds	accuracy (95 % CI)	per seed	mean completion tokens
before	150 x 2	0.973 ± 0.018	0.973, 0.973	153
after	150 x 2	0.980 ± 0.016	0.980, 0.980	159

• Paired difference (after − before), per item averaged over the seeds: +0.7 points, 95 % bootstrap CI [+0.0, +1.7]. No drop.
• Why the samples differ although the logits are the same bits: with 8 requests in flight the batches, and so the order of draws from the random stream, depend on timing.

(c) API end to end at temperature 1 (temperature 1.0, top_p 1), full model, one Mac Studio, each check 3 times:

check	before	after
a turn stops on its own, no role tokens	3/3	3/3
reasoning (enable_thinking) goes to reasoning_content	3/3	3/3
thinking off answers without reasoning	3/3	3/3
tool call returned as tool_calls	3/3	3/3
answer from a tool result	3/3	3/3
a JSON-valued string argument stays a string	3/3	3/3
streamed call with an integer argument	3/3	3/3
two parallel calls	3/3	3/3
developer message followed	3/3	3/3
prefix cache: same answer, cached prefix, partial hit	3/3	3/3
needle recall at ~4k, ~8k and ~16k tokens	9/9	9/9
8 concurrent requests through the batch generator, all correct	3/3	3/3

On the two Mac Studios (one trial each; the needles at 4k, 8k and 16k tokens included): 14 / 14 checks on 2x pipeline and 14 / 14 on 2x tensor, for both builds.

Test plan

• Checks: basedpyright (0 errors), ruff check (clean), nix fmt (no changes), and the CI test command python -m pytest src -m "not slow" --import-mode=importlib: 2,030 passed and 3 skipped, on a Mac Studio (M3 Ultra) and on a Mac with an M4 Max. With the slow tests, the Hy3 family's and Pipeline parallelism: build a rank's decode graph while its input is still on the way #2521's tests: 103 passed on both.
• GPU generations: The Hy3 family's tests (122) pass on a Mac Studio (M3 Ultra) and a Mac with an M4 Max, natively and with MLX dispatching as on M1, M2, M3 and M4 (MLX_METAL_GPU_ARCH). The full suite is green on the macOS CI runner, whose M1-class hardware runs every kernel bit for bit under M1 dispatch.
• New tests (hy_v3/tests/test_hy_v3_kernels.py, 81 cases), bit equality with the ops each kernel replaces:
  • the add + RMSNorm (with the float32 copy), for 1 to 513 rows;
  • the query/key norms with RoPE: int offsets, one-element and per-row offset arrays, 1 to 37 tokens and batches of 3;
  • the router: 192 experts (top 8) and 16 (top 4), 1 to 777 tokens, with and without exact ties;
  • the MoE combination, with and without the residual and the float32 combine;
  • SwiGLU + 4-bit down projection: for gathers, for a prefix of a stack of experts, and for one dense row;
  • stacked projections against their parts at every row count from 1 to 2 past MLX's batch limit, at Hy3's q/k/v and gate/up shapes;
  • the experts with gate and up stacked, against mlx-lm's SwitchGLU at 1, 2, 7, 8, 9, 63, 64, 95, 96, 97, 128 and 129 tokens (192 experts, top 8: mlx-lm sorts from 8 tokens, MLX's sorted gather starts at 96);
  • a whole small Hy3 (head size 128; 4-bit with an 8-bit router, and bf16) through the fused path against Add Tencent Hy3 #2520's path: prefill, then decode steps of 1 and 2 tokens, batch 1 and 3, logits bit for bit.
• Sharding (test_hy_v3_sharding.py): tensor parallelism in both layouts (attention whole over the ring, split by heads over RDMA) against a single device; pipeline parallelism bit-identical, as before.
• Pipeline parallelism: build a rank's decode graph while its input is still on the way #2521's tests (test_pipeline_decode_receive.py), carried with its commit.
• Models and hardware: Hy3 4-bit, full model on Mac Studios (M3 Ultra, 256 GB), and its first 8 layers on a Mac Studio (M3 Ultra, 96 GB).

Reproduce

# speed: one exo node per build (alternate the builds), the harness's perf runner against it
python -m exo_harness.perf run --base-url http://<node>:52415 --model XavierLocalAI/Hy3-4bit \
  --tokenizer <checkpoint> --configs single --pp 512,4096,16384 --tg 128 --repeat 2 --warmup 1 \
  --concurrency 1 --keep-instance --out perf.json
python -m exo_harness.perf run ... --pp 512 --concurrency 4,8 --reuse-instance ...
python -m exo_harness.perf run ... --configs pp2,tp2 --tp-meta ring --pp-meta ring ...

# teacher-forced log-probs on #2520's tokens, compared bit for bit with #2520's run
python -m exo_harness.exo_inprocess run --model-id XavierLocalAI/Hy3-4bit --model-path <checkpoint> \
  --prompts data/prompts.jsonl data/prompts_hy3.jsonl --skip chat-thinking-01 --tasks tf \
  --continuations-from runs/exo_final.json --out runs/after.json
# the same with exo_harness.launch and --parallel tensor / pipeline for two ranks

# kernels against the ops they replace, and the fused model against #2520's path
uv run pytest src/exo/worker/engines/mlx/models/hy_v3/tests

How exactness was established

Each kernel was written against MLX 0.32.4's own sources for the ops it replaces, then checked bit for bit: by unit tests at Hy3's shapes, and on the real checkpoint. Bit equality also depends on the GPU's code generation and arithmetic, so the kernels run only on generations whose hardware ran them bit for bit: M1 (the CI runner), M3 and M4.

• rms_norm (nn.RMSNorm): rms_single_row reads 4 values per thread over ceil(D/4) threads, rounded up to a simdgroup. It sums the squares with simd_sum within and across simdgroups, normalises in float32, rounds to the model dtype and multiplies by the weight in that dtype. The kernels do the same; for a 128-value head the second simd_sum adds zeros.
• mx.fast.rope (traditional=False, scale 1): the same angle, exp2(-d * log2(base)) with log2 of the float32 base, and fast::cos/sin. The position arithmetic is the same too: integer for an int offset, unsigned for exo's per-row offset arrays.
• The router's selection:
  • On the GPU, argpartition runs MLX's merge sort: an odd-even sort within each thread, then merges that take the left run on ties. That is a stable sort, so the first k of argpartition(-scores) are the k highest scores, ties to the lower index, in that order. The kernel ranks each expert against all the others with the same rule.
  • The 8 chosen scores are summed as row_reduce_small does: one thread per row, left to right from 0.
  • They are then divided by the sum plus 1e-20 and multiplied by the scaling factor, as the compiled ops do. The sigmoid is MLX's Sigmoid.
• (experts * weights[..., None]).sum(-2) over 8 experts:
  • col_reduce_small gives each expert row its own thread and adds the rows in order onto 0, after the products are rounded in their own kernel. So the combination kernel disables fused multiply-add.
  • MLX's RMSNorm gives the same bits with and without contraction, so the combination and the next norm can share one kernel.
• Quantized matmuls:
  • Below the batch limit, MLX runs qmv for one row and qmv_wide for a few, and both give every output row its own reduction over K. The limit depends on the shapes and the GPU generation: 6 to 33 rows across M1-M5, 12 to 32 on an M3 Ultra. Stacked projections are used only below it, on the generations MLX 0.32.4's dispatch names, and a test checks the limits against MLX's dispatch for each.
  • The expert gather takes gather_qmm_rhs for sorted pairs from 16 pairs and 4 per expert, and otherwise gather_qmv. That choice depends only on the number of pairs and experts, and gather_qmv runs the same qmv_fast_impl as qmv. Stacking gate and up along the output rows changes neither.
• SwiGLU:
  • mx.compile's kernel for mlx-lm's swiglu computes a sigmoid and two multiplies in the model dtype. The fused down projection does the same for the 16 values each qmv_fast thread reads.
  • It is used only where MLX runs qmv_fast / gather_qmv_fast: one row per matmul, 4-bit, group size 64, an inner size that is a multiple of 512, on M1, M3 and M4 GPUs. M2 and M5 GPUs use the separate ops until verified on that hardware.

🤖 Generated with Claude Code

[AlexCheema]

The fused path on every Apple GPU generation

The macOS CI job failed four of this PR's tests: test_swiglu_down[1], test_swiglu_down[3], test_stacked_projection_matches_its_parts_at_every_row_count and test_fused_model_matches_the_reference_path[True].

What actually failed. No result was wrong. The fused model's logits matched #2520's bit for bit on CI. All four failures are a fast path that turned itself off while the test asserted it ran. CI's GPU is a virtual machine's: mx.device_info() reports "Apple Paravirtual device", architecture air64_v27. MLX itself treats that as an unknown GPU (generation 2) and uses its fallback dispatch tables.

• The stacked projections look up MLX's matrix-vector row limit for the GPU (qmv_batch_limit). That returned None for an unknown GPU, so nothing was stacked, which is correct. The tests asserted a stack.
• swiglu_down read the generation by concatenating the digits of the name: air64_v27 became generation 6427, which is above 16, so the kernel stayed off. That was right by accident. The tests asserted it was on.

What was actually exposed.

• The fused kernels claim bit equality with MLX's ops. That equality depends on the GPU's code generation and arithmetic (M1 has no bfloat16 units), not only on MLX's kernel choice.
• swiglu_down claimed "checked on M1 to M4", but it had never run on M1 hardware: on CI it was off.
• The other kernels ran on any Metal GPU, including M2 and M5, where nobody had run them.
• The architecture check accepted any name with digits, and nothing tied the dispatch rules to the MLX release they were read from.

Fix

• GPU and MLX detection: strict. Only MLX 0.32.4 (PyPI, or exo's pin of the same commit) and GPUs named applegpu_g13* to applegpu_g17* (size letter g/s/d: M1 to M5) are known. Anything else (another MLX release, an unknown or future GPU, CI's paravirtual GPU, CUDA/CPU) runs Add Tencent Hy3 #2520's ops.
• The Metal kernels (add + RMSNorm, q/k norm + RoPE, router selection, MoE combine, SwiGLU in the down projection) run on the generations whose hardware ran them bit for bit: M1, M3 and M4. M2 and M5 GPUs use Add Tencent Hy3 #2520's ops until verified on that hardware.
• The stacked dense gate/up is a choice of MLX kernels, not new arithmetic. It follows MLX 0.32.4's dispatch rules and runs on every generation those rules name (M1 to M5), with or without the kernels. On M5 this rests on MLX's dispatch rules alone (the matrix-vector kernels below the row limit, no matrix-unit kernel for those shapes): it was not run on M5 hardware, and a whole-model M5 emulation isn't possible on other GPUs.
• The stacked expert gathers run on M1 to M4. On M5, MLX runs large sorted gathers on the GPU's matrix units (gather_qmm_rhs_nax), which isn't checked here, so M5 keeps the separate gathers.

Verification

On CI's M1-class hardware, the fused-path tests now also run in subprocesses with MLX_METAL_GPU_ARCH set, so MLX dispatches as on M1, M2, M3 and M4 Max:

• With M1 dispatch, the runner's M1 silicon ran every kernel, swiglu_down included, against MLX's ops, bit for bit. That is the M1 verification.
• With M2 dispatch (stacks without the kernels), the fused model is Add Tencent Hy3 #2520's logits bit for bit.
• On CI's own GPU name the tests check the fallback.

The same tests pass natively on M3 Ultra and M4 Max, and on M3 Ultra under every emulated generation 13-16 (including CI's air64_v27).

More checks:

• qmv_batch_limit is checked against MLX's actual dispatch for 15 architectures (generations 13-17 × g/s/d), one subprocess each. One row below the limit MLX must run the matrix-vector kernel, at the limit it must not. An MLX upgrade that moves a limit fails this test, and another test fails until the new release is added to the checked list.
• The stacked projections are checked in float32, where any reordering shows, at Hy3's shapes: shared expert 1536+1536, dense 13312+13312, and their 2x tensor-parallel shards 768+768 and 6656+6656.
• The 8-layer cut of XavierLocalAI/Hy3-4bit (real weights) ran through the PR's 592-case sweep (prefill 1-320 and ±1 around MLX's thresholds, the same after 700 cached tokens, steps of 1-9 tokens at batch sizes 1-9). Each run compares this PR with Add Tencent Hy3 #2520 under the same GPU or emulated dispatch. 592 of 592 identical in every run: on M3 Ultra and M4 Max hardware, and on M4 Max with MLX dispatching as on M1 (kernels and stacks), M2 (stacks without the kernels) and M4 Max. The last change, keeping M5 off the expert stacks, touches only M5's paths.

Which GPUs run what

GPU	Metal kernels	stacked dense gate/up	stacked expert gathers
M1, M3, M4 (any size)	yes	yes	yes
M2	no (#2520's ops until verified on M2 hardware)	yes	yes
M5	no (until verified on M5 hardware)	yes (by MLX's dispatch rules)	no
other GPUs, other MLX releases, CUDA/CPU	no	no	no

Performance

• M3 Ultra and M4 Max run exactly the paths they ran before, so every number in the PR stands. M1 does too.
• M2 and M5 GPUs no longer run the Metal kernels and keep only the stacks (M5 only the dense one). Most of the decode speedup comes from the kernels, so on M2 and M5 this PR is close to Add Tencent Hy3 #2520's speed until the kernels are verified there.

🤖 Generated with Claude Code