[inureyes]

Summary

Each gather_qmm / gather_mm call with no lhs_indices makes MLX add an Arange + Reshape + Broadcast to the graph, and a non-uint32 expert-id tensor adds an AsType. SwitchGLU paid for this once per projection on every step. It now builds those index arrays once per block and shares them across gate, up and down. Arithmetic is unchanged because the arrays it passes hold the same values MLX would build.

• prepare_gather_indices + SwitchLinear::forward_indexed: gate/up share one lhs_indices array (unsorted path), all three projections share one (sorted prefill path), and the uint32 ids are cast once. Down keeps MLX's default lhs in the unsorted path because its batch shape is [n, k].
• expand_dims pairs on the input and on the Inkling expert-scale sidecars become one expand_dims_multi node each.
• New arange_u32 bridge function builds the lhs exactly as MLX's indices_or_default does (arange(stop, uint32)), so no cast node is added.
• SwitchLinear::forward (used directly by 15 families) is unchanged and delegates with lhs = None.

Evidence (GB10, CUDA, MLXCEL_EXPORT_DECODE_DOT, first decode step)

Model	nodes before -> after	edges before -> after
granite-4.0-h-tiny-4bit (40 MoE layers)	3419 -> 3259	9057 -> 8777
qwen3-30b-a3b-4bit, MLXCEL_FUSED_MOE=0 (48 layers)	2895 -> 2703	7575 -> 7239
qwen3-30b-a3b-4bit, default (fused decode kernel)	2319 -> 2319	6327 -> 6327
granite-4.0-h-350m-4bit (dense MLP, no experts)	1667 -> 1667	4593 -> 4593

Per MoE layer: Arange, Broadcast, ExpandDims and Reshape each drop by 1, and no other op kind changes. Granite's remaining Full, Slice and AsType nodes come from the SSM and attention paths, which this PR does not change.

Greedy mlxcel generate --temp 0 -n 64 output is byte-identical between the base and new binaries on granite-4.0-h-350m, granite-4.0-h-tiny, qwen3-30b-a3b (fused on and off) and qwen3.5-35b-a3b (fused on and off). A second base run matched the first, so the comparison is deterministic. The prompt runs prefill through the sorted path (tokens x top_k >= 64), and decode runs through the unsorted path.

Throughput is informational only. With pp512/tg128, n=3 interleaved rounds per arm, and a null arm (a byte-identical copy of the base binary), the null arm moved up to -34% from base inside a single round. All base/new medians sit inside that spread (decode tok/s, base / new / null: granite-tiny 100.3 / 100.1 / 100.4, qwen3-30b-a3b default 80.1 / 75.1 / 75.4, qwen3.5-35b-a3b 59.2 / 62.3 / 60.1). On qwen3-30b-a3b the new binary matches the null copy, and decode there runs the fused kernel, which this PR does not touch. The forward ms/token medians are also flat (granite-tiny 1.44 vs 1.40, null 1.40). The 4.7% node cut predicts about 0.05 ms, which is below this host's noise. I make no throughput claim.

Test plan

• models::switch_layers on CUDA (36 tests, including mxfp_tests and 4 new gather_index_tests that pin bit-exactness against the per-call form in decode, small-prefill and sorted-prefill shapes with int32 and uint32 ids)
• Caller families on CUDA: granitemoehybrid, qwen3_moe, inkling, laguna, kimi_k3, qwen3_vl_moe, olmoe
• cargo clippy --release --lib --tests --features cuda -- -D warnings, cargo fmt --check

Not verified on this host

• M5 Max ratios against benchmarks/pylm_m5max_2026-09-06.csv and the Metal forward ms/token: these need Apple Silicon.
• The Metal nightly mxfp_tests ([nightly-verify] main is red #1599): this PR does not touch mxfp-specific dispatch, and those tests pass on CUDA.

Closes #1713