You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
Make the fused decode-MoE Dff cap backend-aware. The current single global cap (MLXCEL_FUSED_MOE_MAX_DFF, default 4096) is correct for Metal but conservative for CUDA, where measurements show a much higher crossover. Add a backend default, and optionally a per-model heuristic, so mid-size CUDA experts get the fused path.
Source: idea.md "3. Backend-Aware Fused-MoE Dispatch Policy" (P1, expected payoff medium on CUDA/MoE, risk medium).
Scope note (avoid overlap)
This issue is ONLY about the backend-aware Dff cap and dispatch policy. Wiring MiniMax into the fused decode dispatch is tracked separately by #304 under epic #307 and is OUT OF SCOPE here.
Current state
SwitchGLU::forward_fused_kernel() declines above MLXCEL_FUSED_MOE_MAX_DFF, which defaults to 4096 as a single global cap (src/models/switch_layers.rs, forward_fused_kernel near line 362, cap read near lines 402 to 406).
That cap is correct for Metal-era measurements: large experts such as phi-3.5-moe and Mixtral lost on M1 Ultra.
CUDA measurements show much larger fused-MoE gains and a higher crossover: qwen3-30b-a3b +55%, qwen3.5-35b-a3b +41%, with CUDA crossover around 13k to 14k Dff (docs/benchmark_results/fused-moe-decode-kernel-design.md).
Make the default Dff cap backend-aware: keep Metal at 4096; consider a CUDA default in the 8192 to 14336 range after validation on current binaries.
If the benchmark set shows model-specific cliffs, add a runtime/backend heuristic rather than a single global cap.
Keep MLXCEL_FUSED_MOE_MAX_DFF as a manual override.
Validation
Re-run fused-MoE cap sweeps on the current MLX pin BEFORE changing defaults (the policy must be based on current binaries, per idea.md suggested experiment 3).
Metal: confirm no regression on phi-3.5-moe, Mixtral, Qwen3-MoE, OLMoE.
CUDA/GB10: test qwen3-30b-a3b, qwen3.5-35b-a3b, lfm2-8b-a1b, phi-3.5-moe, Mixtral.
Track both decode tok/s and greedy parity or RMS jitter class.
Acceptance criteria
Default Dff cap is backend-aware (Metal unchanged at 4096; CUDA default raised to a validated value).
MLXCEL_FUSED_MOE_MAX_DFF still overrides the backend default.
Cap sweeps re-run on the current MLX pin; the chosen CUDA default is justified by current data, not the historical Metal number.
No Metal regression on the listed models; CUDA gains recorded for mid-size experts.
Design doc docs/benchmark_results/fused-moe-decode-kernel-design.md updated with the new policy and the sweep data.
Risk and priority
idea.md priority P1; risk medium. Smaller than the already-shipped small-expert path, but exposes free speed for mid-size CUDA experts and removes a conservative default.
Summary
Make the fused decode-MoE
Dffcap backend-aware. The current single global cap (MLXCEL_FUSED_MOE_MAX_DFF, default 4096) is correct for Metal but conservative for CUDA, where measurements show a much higher crossover. Add a backend default, and optionally a per-model heuristic, so mid-size CUDA experts get the fused path.Source: idea.md "3. Backend-Aware Fused-MoE Dispatch Policy" (P1, expected payoff medium on CUDA/MoE, risk medium).
Scope note (avoid overlap)
This issue is ONLY about the backend-aware Dff cap and dispatch policy. Wiring MiniMax into the fused decode dispatch is tracked separately by #304 under epic #307 and is OUT OF SCOPE here.
Current state
SwitchGLU::forward_fused_kernel()declines aboveMLXCEL_FUSED_MOE_MAX_DFF, which defaults to 4096 as a single global cap (src/models/switch_layers.rs,forward_fused_kernelnear line 362, cap read near lines 402 to 406).phi-3.5-moeand Mixtral lost on M1 Ultra.qwen3-30b-a3b+55%,qwen3.5-35b-a3b+41%, with CUDA crossover around 13k to 14k Dff (docs/benchmark_results/fused-moe-decode-kernel-design.md).Proposal
MLXCEL_FUSED_MOE_MAX_DFFas a manual override.Validation
phi-3.5-moe, Mixtral, Qwen3-MoE, OLMoE.qwen3-30b-a3b,qwen3.5-35b-a3b,lfm2-8b-a1b,phi-3.5-moe, Mixtral.Acceptance criteria
MLXCEL_FUSED_MOE_MAX_DFFstill overrides the backend default.docs/benchmark_results/fused-moe-decode-kernel-design.mdupdated with the new policy and the sweep data.Risk and priority
idea.md priority P1; risk medium. Smaller than the already-shipped small-expert path, but exposes free speed for mid-size CUDA experts and removes a conservative default.
References
src/models/switch_layers.rs(forward_fused_kernel,MLXCEL_FUSED_MOE_MAX_DFF).docs/benchmark_results/fused-moe-decode-kernel-design.md. Related: perf(moe): wire minimax to the fused decode-MoE kernel #304 (MiniMax wiring), epic(moe): wire the remaining unwired MoE families to the fused decode-MoE kernel #307 (epic), perf(moe): implement the fused decode-MoE kernel on CUDA #319 (CUDA fused-MoE kernel).