Skip to content

perf(moe): record CUDA fused-MoE Dff cap provenance and decline test - #711

Merged
inureyes merged 1 commit into
mainfrom
perf/330-backend-aware-fused-moe-dff-cap
Jul 9, 2026
Merged

perf(moe): record CUDA fused-MoE Dff cap provenance and decline test#711
inureyes merged 1 commit into
mainfrom
perf/330-backend-aware-fused-moe-dff-cap

Conversation

@inureyes

@inureyes inureyes commented Jul 9, 2026

Copy link
Copy Markdown
Member

Summary

Complete the backend-aware fused-MoE Dff cap 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, the MLXCEL_FUSED_MOE_MAX_DFF override, the pure fused_moe_max_dff_from, and the 2026-07-03 GB10 sweep) already landed in #643; that work was measured under MLX pin e9463bb (#626) and 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, 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 the FUSED_MOE_MAX_DFF_CUDA doc comment with a provenance caveat (the 8192 sweep ran on MLX pin e9463bb/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, commit 57c66cac; 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 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 (run_fused_moe_two_kernel picks the cuda_kernel port when metal::is_available() is false).
  • src/models/switch_layers.rs: added 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.
  • 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 exact mlxcel-bench-decode harness to run it), and cross-referenced it from the MLXCEL_FUSED_MOE_MAX_DFF table row.

Design notes

  • The value 8192 is unchanged. Under pin drift the most conservative floor of the measured crossover (break-even at 8192, regression only at 14336) is the correct choice, so no re-tuning is done without hardware re-validation.
  • Metal behavior is byte-identical by construction (still 4096). The diff is runtime-inert (doc comments, a doc file, and a test only); no compiled non-test code path changed.

Test plan

Ran here (Apple M1 Ultra, Metal):

  • cargo fmt --all -- --check
  • cargo test --release --features metal,accelerate -p mlxcel --lib switch_layers (8/8 pass, including the new fused_moe_dff_above_cap_declines_and_at_cap_dispatches; the successful release build subsumes cargo check --lib --tests)
  • cargo clippy --features metal,accelerate -p mlxcel --lib --tests -- -D warnings (clean)
  • Metal MoE decode smoke, dispatch unchanged: mlxcel generate -m models/qwen1.5-moe-a2.7b-4bit -p "Hello" -n 32 --temp 0 produced 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

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
@inureyes inureyes added type:performance Performance improvements priority:medium Medium priority area:models Model architectures, weights, loading, metadata area:core mlxcel-core: MLX FFI, primitives, KV cache, layers status:review Under review labels Jul 9, 2026
@inureyes inureyes added status:done Completed and removed status:review Under review labels Jul 9, 2026
@inureyes
inureyes merged commit b73be08 into main Jul 9, 2026
5 checks passed
@inureyes
inureyes deleted the perf/330-backend-aware-fused-moe-dff-cap branch July 9, 2026 08:49
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
@inureyes inureyes self-assigned this Aug 31, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:core mlxcel-core: MLX FFI, primitives, KV cache, layers area:models Model architectures, weights, loading, metadata priority:medium Medium priority status:done Completed type:performance Performance improvements

Projects

None yet

Development

Successfully merging this pull request may close these issues.

perf(moe): backend-aware fused-MoE Dff cap (CUDA crossover) and dispatch heuristic

1 participant