Skip to content

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

Description

@inureyes

Summary

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).
  • The CUDA fused decode-MoE kernel itself now exists (PR perf(moe): implement the fused decode-MoE kernel on CUDA #319), so the cap policy is the remaining lever for CUDA mid-size experts.

Proposal

  • 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.

References

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:coremlxcel-core: MLX FFI, primitives, KV cache, layersarea:modelsModel architectures, weights, loading, metadatapriority:mediumMedium prioritytype:performancePerformance improvements

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions