perf(moe): record CUDA fused-MoE Dff cap provenance and decline test - #711
Merged
Conversation
The backend-aware fused-MoE Dff cap for #330 (Metal 4096 / CUDA 8192, the MLXCEL_FUSED_MOE_MAX_DFF override, the pure fused_moe_max_dff_from resolver, and the 2026-07-03 GB10 sweep) already landed in #643, measured under MLX pin e9463bb (#626), but that PR was never linked to #330. Since #643 the MLX pin advanced to 0.32.1 (#703/#704, commit 57c66cac), so the CUDA default is now a carried-forward documented-data value rather than one confirmed on current binaries. This completes #330 by recording that provenance honestly and pinning the decline boundary; the cap value 8192 is unchanged. Changes: - Expand the FUSED_MOE_MAX_DFF_CUDA doc comment with the pin-drift caveat: 8192 is the conservative crossover floor (break-even, not the 14336 regression edge) kept so MLX pin drift cannot flip the default into a regression, and a GB10 re-validation on 0.32.1 is pending. - Document on fused_moe_max_dff_from that the cap is family-agnostic (governs every SwitchGLU MoE model per the module Used by: list) and that the backend is resolved at runtime via metal_is_available() rather than a cfg!(feature = "cuda") switch, matching the fused kernel's own runtime dispatch where run_fused_moe_two_kernel selects the cuda_kernel port when metal::is_available() is false. - Add fused_moe_dff_above_cap_declines_and_at_cap_dispatches, a pure test pinning that forward_fused_kernel declines exactly when dff > max_dff at the cap boundary on both backends and under an explicit override. - Add the 2026-07-09 design-doc addendum (CUDA default provenance, why 8192 is safe under pin drift, and the exact mlxcel-bench-decode harness for the pending GB10 re-validation on MLX 0.32.1). Validated on Apple M1 Ultra (Metal): cargo fmt and clippy clean, the switch_layers unit tests pass (8/8 including the new decline test), and a qwen1.5-moe-a2.7b-4bit temp-0 decode smoke confirms the Metal MoE dispatch is unchanged. The diff is runtime-inert (doc comments, a doc file, and a test only), so Metal behavior is byte-identical. Refs #330
3 tasks
3 tasks
inureyes
added a commit
that referenced
this pull request
Jul 9, 2026
## Summary Re-validates the CUDA fused-MoE Dff cap sweep behind `FUSED_MOE_MAX_DFF_CUDA` on GB10 under the current MLX 0.32.1 pin, closing the "pending re-validation" note left by issue #330 / PR #711 in docs/benchmark_results/fused-moe-decode-kernel-design.md. ## Sweep table Hardware: GB10 (DGX Spark), sm_121, CUDA 13.0. MLX pin: `57c66cac7cb3e5b1eb350488a61f1506b40d39f8` (0.32.1, #703/#704). Harness: `mlxcel-bench-decode`, prompt "Hello, how are you today?", 100 decode tokens after a 20-token warmup, median of 3 runs per side, `MLXCEL_FUSED_MOE_MAX_DFF=1` (gather_qmm fallback) vs `=20000` (fused). | Model | Dff | fallback tok/s | fused tok/s | delta | |-------|----:|----------------:|------------:|------:| | lfm2-8b-a1b (control) | 1792 | 140.68 | 160.63 | +14.2% | | phi-3.5-moe | 6400 | 53.85 | 55.80 | +3.6% | | llama-4-scout-17b | 8192 | 21.64 | 21.32 | -1.5% | | mixtral-8x7b | 14336 | 28.26 | 27.78 | -1.7% | The ratio=1.0 crossover interpolates to ~7672 between the phi-3.5-moe and llama-4-scout points, close to the prior ~8000 measured on MLX pin e9463bb. The ~4% shift sits within run-to-run noise at this Dff point: the 2026-07-03 sweep already flagged Dff 8192 as noisy (one of its three fused runs matched the fallback), and this 2026-07-10 sweep's own llama-4-scout fallback spans 21.33-21.85 tok/s run to run, a 2.4% spread comparable to the shift. ## What changed - `FUSED_MOE_MAX_DFF_CUDA` in `src/models/switch_layers.rs` is confirmed at 8192, not re-tuned; this PR does not touch that file. - Adds a 2026-07-10 addendum to docs/benchmark_results/fused-moe-decode-kernel-design.md with the sweep table, per-model findings, and the confirm-not-retune conclusion. - Amends the 2026-07-09 addendum's "pending" note to point at the new addendum, keeping the original text as historical provenance. - Updates the `MLXCEL_FUSED_MOE_MAX_DFF` env-var table row to state the confirmed result instead of "pending". - Adds the raw per-run sweep data as `benchmarks/cuda_gb10_issue712_fused_moe_dff_2026-07-10.csv`. ## Test plan - [x] `cargo check --features cuda --lib --tests` (pinned 1.93.1 toolchain) - clean, `Finished dev profile` with no errors or warnings. - [x] `cargo test --features cuda --lib switch_layers` - 8 passed, 0 failed (unchanged since the constant was not modified). - [x] Real GB10 hardware sweep: 24 `mlxcel-bench-decode` runs (4 models x 2 arms x 3 repeats), one process per run, rc=0 on every run. Closes #712
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Complete the backend-aware fused-MoE
Dffcap for #330 by recording the honest provenance of the CUDA default and pinning the decline boundary with a test. The cap-resolution logic itself (Metal 4096 / CUDA 8192, theMLXCEL_FUSED_MOE_MAX_DFFoverride, the purefused_moe_max_dff_from, and the 2026-07-03 GB10 sweep) already landed in #643; that work was measured under MLX pine9463bb(#626) and was never linked to #330. Since #643 the MLX pin advanced to 0.32.1 (#703/#704, commit57c66cac), so the CUDA default is now a carried-forward documented-data value, not one confirmed on current binaries. This PR closes that honesty gap, documents the runtime-vs-cfg!backend choice, adds a decline-boundary test, and formally closes #330. The cap value 8192 is unchanged.What changed
src/models/switch_layers.rs: expanded theFUSED_MOE_MAX_DFF_CUDAdoc comment with a provenance caveat (the 8192 sweep ran on MLX pine9463bb/perf(cuda/moe): backend-specific default for the fused decode-MoE Dff threshold #626; the pin has since moved to 0.32.1/chore(mlx): bump upstream MLX pin to 0.32.1 #703/chore: bump MLX pin to 0.32.1 (#703) #704, commit57c66cac; 8192 is the conservative crossover floor kept so pin drift cannot flip it into a regression; a GB10 re-validation on 0.32.1 is pending).src/models/switch_layers.rs: documented onfused_moe_max_dff_fromthat the cap is family-agnostic (governs every SwitchGLU MoE model per the moduleUsed by:list) and that the backend is resolved at runtime viametal_is_available()rather than acfg!(feature = "cuda")switch, matching the fused kernel's own runtime dispatch (run_fused_moe_two_kernelpicks thecuda_kernelport whenmetal::is_available()is false).src/models/switch_layers.rs: addedfused_moe_dff_above_cap_declines_and_at_cap_dispatches, a pure test pinning thatforward_fused_kerneldeclines exactly whendff > max_dffat the cap boundary on both backends and under an explicit override.docs/benchmark_results/fused-moe-decode-kernel-design.md: added the 2026-07-09 addendum (CUDA default provenance, why keeping 8192 is safe under pin drift, and the pending GB10 re-validation on MLX 0.32.1 with the exactmlxcel-bench-decodeharness to run it), and cross-referenced it from theMLXCEL_FUSED_MOE_MAX_DFFtable row.Design notes
Test plan
Ran here (Apple M1 Ultra, Metal):
cargo fmt --all -- --checkcargo test --release --features metal,accelerate -p mlxcel --lib switch_layers(8/8 pass, including the newfused_moe_dff_above_cap_declines_and_at_cap_dispatches; the successful release build subsumescargo check --lib --tests)cargo clippy --features metal,accelerate -p mlxcel --lib --tests -- -D warnings(clean)mlxcel generate -m models/qwen1.5-moe-a2.7b-4bit -p "Hello" -n 32 --temp 0produced coherent output at 96.58 tok/s on Apple GPU (Metal), exercising the fused MoE decode path (Dff below the 4096 Metal cap, so it dispatches).Documented-data (not re-run in this PR): the CUDA default 8192 rests on the GB10 sweep in the 2026-07-03 addendum under MLX pin
e9463bb(#626), landed in #643.Pending (hardware-blocked): GB10 re-validation of the crossover on the current MLX 0.32.1 pin. No CUDA/GB10 hardware was reachable for this change; the how-to-run harness is recorded in the 2026-07-09 addendum.
Closes #330