Repository navigation
perf(moe): profile the batch-1 decode gap on M3 Ultra after #2310 #2479
Description
Context
mlxcel decodes the three large 256-expert top-6 MoE checkpoints at 16-24% of the mlx-vlm baseline on the M3 Ultra (512 GB), consistent across two architectures (deepseek_v4, inkling) and two modalities (text, VLM). Prefill reaches 55-82% of the baseline, so the gap is specific to batch-1 decode.
2026-10-20 rerun at 058d7f7 (MLX 243bdf09), 512-token synthetic prompt, 128 generated, decode tok/s (data merged in #2476):
| model | mlxcel | mlx-vlm | ratio |
|---|---|---|---|
| deepseek-v4-flash-4bit | 8.38 | 35.40 | 23.7% |
| inkling-small-mlx-4bit | 7.78 | 48.94 | 15.9% |
| inkling-small-nvfp4 | 6.69 | 36.12 | 18.5% |
The VLM protocol (60-token image prompt) shows 15.8% / 18.5% for the two Inklings. Every number is within ~3% of the 2026-09-08 pre-fix values on both sides.
Ruled out
- The unwired fused MoE kernel was the perf(moe): attribute the MoE decode gap at the M = 1 expert gather #1859 attribution, fixed by perf(moe): attribute M = 1 expert gather gap and wire Inkling/V4 MoE fast paths (#1859) #2310 (2026-10-14). The rerun after that fix shows the gap unchanged, and a same-build control with MLXCEL_FUSED_MOE=0 decodes at 7.80 tok/s vs 7.78 with the kernel enabled, so the kernel is not the lever on this host.
- Not background load: all runs are wrapped in scripts/with_indexers_paused.sh (SIGSTOP/SIGCONT quieting, nothing killed).
- Not protocol or EOS handling: both sides use the same synthetic corpus prompt, no chat template, EOS suppressed.
- Not quantization-specific: the mlx 4-bit and nvfp4 arms sit in the same 16-24% band, suggesting a shared mechanism.
Candidates to profile
- Per-layer graph overhead: node count, dispatch and allocation cost per token at M=1.
- Weight gather path: how routed expert weights are gathered and scaled before the GEMM, memory layout and copy volume per token.
- Scale application: where the top-6 weights multiply in, whether it forces extra passes over activations.
- Weight layout differences vs mlx-vlm for the same checkpoints (mlx-vlm reads the same MLX weight files).
Next steps
- Time per-layer forward on one architecture (inkling) with and without the MoE block at batch 1 on Metal, same build.
- Compare against mlx-vlm per-layer timing for the same checkpoint in the uv venv.
- Metal counters (gputrace or Instruments) on the decode loop to attribute bytes moved vs compute.
- Record findings in docs/benchmark_results/ following the m3ultra protocol; the fix follows the attribution.
Activity
inureyes commentedon Oct 20, 2026
Cause-profiling findings for the 4.2x-6.3x M3 Ultra decode gap (issue #2479). Method: paired controls, env-gated code skips in
InklingSparseMoe::forward(TEMP patch, inert without env vars), xctrace Metal System Trace on both runtimes, and weight-index bandwidth math. All numbers on the M3 Ultra 512GB host, cool, release build at 5d35b3b (mlx commit 243bdf09).What the controls exonerate
MLX_MAX_OPS_PER_BUFFERsweep 50 to 100000 is flat at ~7.9 tok/s: command-buffer batching is not the lever.MLXCEL_FUSED_MOE=0control: 7.80 vs 7.78 tok/s. The fused vs gather expert-kernel choice does not matter on this host.What the skip experiment isolated
With
MOE_PROFILE_SKIP_ROUTED=1the routed experts are replaced with zeros; the routing chain, shared expert, attention, and LM head still run. Decode moves 7.94 to 20.38-20.41 tok/s (126 to ~49 ms/token); the repeat and the interleaved control confirm it. The routed-expert block therefore costs ~77 ms/token, 61% of decode, and streams its 3.42GB/token at ~44 GB/s effective. WithMOE_PROFILE_SKIP_ROUTE_CHAIN=1added (constant stand-in routing), decode reaches ~46 ms/token: the ~700 tiny routing kernels per token cost only 3-5 ms and are not a lever.The remaining 46 ms non-MoE token (attention 3.7GB + LM head 1.65GB + shared + router) is itself 3x below this binary's dense-model speed (~123 GB/s vs ~428). Both blocks are slow at this scale.
The footprint cliff
Same binary, same MoE code path: inkling-0.6b-a0.6b (~2GB streamed/token) decodes at 229 tok/s = ~449 GB/s, and its skip-both arm at ~569 GB/s. The cliff is between 15GB and 143GB of model footprint and the baseline does not have it (48.94 tok/s = ~446 GB/s at 143GB).
What the residency work ruled out
ResidencySets(multi-set, #4211) and the refresh-interval var (#4633) exist in the baseline wheel (0.32.4) too; the refresh thread is opt-in and off by default.MLXCEL_WIRED_LIMIT=0gives 0.57 tok/s with and without the indexers paused (the baseline decodes unwired at 48.94). mlxcel's wire-at-startup default is load-bearing; the residency machinery is not what makes the wired run slow.Where the remaining time goes
The full-run sampler shows the main thread 80% blocked in
std::condition_variable::wait, ~15% of work samples in one mlx function that creates AGXBuffers (placementSparse/remoteStorageResource) with residency-set commits and mach_msg2 traps; the four async workers are 100% idle. The GPU execution-points table is sampled and its busy fraction is not trustworthy. The two suspects that survive the elimination are (1) per-op driver work that scales with the buffer population (AGXBuffer creation and residency-set registration per new temporary), and (2) macOS 26 sparse-placement behavior for a 143GB buffer population. Both explain why the baseline (older mlx, different allocation pattern) does not cliff.Next steps
MLX_RESIDENCY_SET_MAX_PCT(single-set vs 32) at 143GB.inureyes commentedon Oct 20, 2026
Refinement to the sampler picture above, from a second sample of the skip-both run (400 tokens, 20s window, 14132 main-thread samples): the main thread splits ~43% blocked in IOSurfaceSharedEvent waitUntilSignaledValue (waiting on GPU completions, via iokit_user_client_trap) and ~47% in the same stripped mlx evaluation path seen in the full run (the AGXBuffer-creation/residency branch), with the four async workers 100% idle throughout. So the non-MoE token is roughly half GPU-wait and half CPU-side evaluation work, and the CPU half runs through the code that registers and creates AGXBuffers per new temporary. The next diagnostic that follows from this is counting MTLBuffer creations and destructions per token in both runtimes: if mlxcel churns thousands of buffers per token at 143GB while the baseline reuses them from the allocator cache, the cliff is allocator-path, not kernel-path.
inureyes commentedon Oct 20, 2026
Root cause confirmed for the 4.2x-6.3x M3 Ultra decode gap (follows the two findings comments above). Method: env-gated counters patched into the pinned mlx allocator and device completion path, a one-graph micro-benchmark that pays the eval sync once, the v0.7.0 official release binary, and four control A/Bs. All numbers on the M3 Ultra 512GB host, cool, release build at 5d35b3b (mlx commit 243bdf09).
The verdict: GPU-bound small-kernel inefficiency
MTLCommandBuffer.GPUStartTime/GPUEndTimesums 15,482 ms of GPU busy time for a process whose measured decode wall is 12,564 ms (plus a ~3 s load). The 126 ms token is ~126 ms of GPU work; the main thread's 80 %condition_variable::waitshare is waiting for the GPU, and the ~18 ms/token of CPU dispatch overlaps it. Every host-side hypothesis (allocator, residency, pipeline, batching) is therefore out.gather_qmmops folded into one dependency, eval once, sync paid once) measures 29 us/op for the big model's shapes ([256, 2048, 4096] affine 4-bit/g64, effective 144 GB/s) and 22 us/op for a dense [4096, 4096] plane (376 GB/s). A 2-4 MB plane streams in ~7 us at wall bandwidth; the other ~20 us is per-kernel fixed cost.What the new controls exonerate
MLXCEL_FORCE_SYNC=1: 7.99-8.00 tok/s vs 7.97 async. The lookahead pipeline is not a lever.MLX_RESIDENCY_SET_MAX_PCT=100(single residency set) vs the 32-set default: 7.94-7.95 vs 7.96-7.97. The residency-set layout is not a lever.Why the baseline does not cliff
The python baseline's inkling decode uses fused custom kernels (seen in its Metal trace):
inkling_moe_route(one kernel for the whole routing chain),inkling_moe_down_combine(q4/g64 direct in-kernel dequant with the weighted sum folded), shapeless_swiglu_scaled, and a load-timeshared_experts_to_dense. That collapses the MoE block from ~756 gather/GEMV kernels plus the routing chain into roughly a hundred bigger kernels, so its ~500 ops/token at the same ~26 us fixed cost land at ~20 ms/token. mlxcel's graph pays the same fixed cost ~4,600 times.Next steps
inureyes commentedon Oct 20, 2026
Fix results: decode is now 44.0 tok/s on M3 Ultra, 5.4x faster, 90% of the mlx-lm baseline.
What the port alone showed. The fused decode kernels from the confirmed plan (single routing kernel, fused gate/up+swiglu activation, down projection with in-kernel q4/g64 dequant and folded weighted sum) rebuilt cleanly and passed their parity tests, but the benchmark barely moved: 8.10 tok/s vs 7.97 before the port, against the 48.94 tok/s mlx-lm baseline. The route kernel fired, but the down-combine kernel's gate silently refused at runtime, so the MoE stayed on the unfused path.
Root cause. The Inkling decoder layer was wired to a composed SDPA (matmul + softmax + matmul in
mlx_cxx_bridge.cpp, the only caller of that path in the tree, dating from #1540) instead of MLX'sfast::scaled_dot_product_attention.mlx.core.softmaxpromotes to float32, which poisoned the entire residual stream: attention output, both short convs, and the MoE inputs all ran in float32 while the expert scales are bfloat16. Two consequences: every bfloat16 projection in the layer promoted its weights to float32 per decode token, and the fused down-combine kernel's scales-dtype gate (which requiresscales.dtype == x.dtype, same condition as the reference implementation) refused every call.Fix. Switch
inkling_layerattention tofast_scaled_dot_product_attention(same signature, nullable mask, keeps the query dtype). The residual stream stays bfloat16 end to end, the fused MoE decode kernels now engage on both the routing and down-combine sides, and the shared experts run as plain dense SwiGLU matmuls.Numbers (inkling-small-mlx-4bit, 100 decode tokens, caffeinated, fastest of 3 runs):
The remaining ~10% is diffuse, consistent with the per-op execution overhead pattern documented for Gemma3n (#762); no single lever stands out in the op counts. TEMP instruments (MLX_ALLOC_STATS counters) stay in the working tree until the fix lands.
perf(moe): wire the switch gate/up kernel into SwitchGLU::forward
test(moe): validate the switch gate/up kernel on real OLMoE weights
fix(moe): guard switch gate-up dispatch and preserve fallback
perf(moe): fuse the routed gate/up pair into one Metal kernel (#2514)
perf(inkling): decode via fast SDPA and fused MoE kernels (lablup#2505)