lablup / lablup/mlxcel

perf(moe): attribute the MoE decode gap at the M = 1 expert gather

Open
#1,859 1 comment 0 reactions 0 assignees View on GitHub
area:core area:inference priority:medium status:ready type:performance
Dominant language
Rust
Stars
467
Forks
54
Avg merge
4h 25m
Merged PRs (30d)
310

Description

## Problem / Background

Two MoE families reach `gather_qmm` with `M = 1` on every decode token and never reach the fused decode-MoE kernel. That much is a code fact, verified below against `60341873`. Whether it is the dominant term in a measured decode gap is **not** established, and this issue asks for a profiled attribution before any change is written.

The shape. `SwitchGLU::forward` (`src/models/switch_layers.rs:939-963`) expands `x` with two `expand_dims` at `switch_layers.rs:946-947`, producing `[n, 1, 1, K]`, so `gather_qmm` sees `M = 1` and `B = n * top_k`. The GB10 attribution record for #1799 states this and names its consequence on CUDA. `docs/benchmark_results/laguna-dflash-verify-cost-gb10-2026-09-11.md:96`: "`SwitchGLU::forward` expands `x` to `[n, 1, 1, K]`, so `gather_qmm` sees `M = 1, B = 8n` and the routed experts take `qmm_sm80` at every width including classic". `qmm_sm80` is the prefill GEMM; `docs/benchmark_results/qmm-sm121-tile-tuning-gb10-2026-07-10.md:24` records that single-sequence decode is supposed to take `qmv`. So on CUDA the routed experts are already known to run a tiled GEMM at one row per (token, expert) pair during classic decode, on a different backend and a different checkpoint from the observation below.

What is specific to these two families is that neither has the fused decode-MoE kernel as an escape hatch. Eighteen families wire `SwitchGLU::forward_fused_kernel` (`grep -rln forward_fused_kernel src/models/`, excluding the definition and a test file); `deepseek_v4` and `inkling` are not among them.

- `deepseek_v4` does not use `SwitchGLU` at all. It carries its own `LimitedSwitchGlu::forward` (`src/models/deepseek_v4_moe.rs:242-252`) with the same two `expand_dims` at `deepseek_v4_moe.rs:243-244`, and it passes `sorted: false` unconditionally. The module doc is explicit at `deepseek_v4_moe.rs:49-50`: "The fused MoE decode kernel (`MLXCEL_FUSED_MOE`) is deliberately not wired: it has no clamp stage."
- `inkling` calls `SwitchGLU::forward_with_expert_scales` (`src/models/inkling/mlp.rs:217-222`). That method has no fused-kernel branch at all: it delegates to `SwitchGLU::forward` only when both sidecars are `None` (`switch_layers.rs:982-984`) and otherwise runs the unsorted gather (`switch_layers.rs:986-989`), documented as deliberate at `switch_layers.rs:970-972`.

Sorting does not rescue decode either. `SwitchGLU::forward` sorts only when `n_tokens * top_k >= 64` (`switch_layers.rs:940-944`), which one decode token never reaches, and both `forward_with_expert_scales` and `LimitedSwitchGlu::forward` are unsorted by construction.

## Corroborating observation (Metal, one session, unprofiled)

An Apple Silicon session on a Mac Studio M3 Ultra benchmarked `deepseek-v4-flash-4bit`, `inkling-small-mlx-4bit` and `inkling-small-nvfp4` against mlx-vlm on 2026-09-08 and reported it today. Matched shape on both sides: 512-token prompt, 128 generated, warmup pass, EOS suppressed, no chat template. **mlxcel decode landed at 16 to 23 percent of the mlx-vlm baseline while prefill was 55 to 82 percent**, consistent across both architectures and across text and image modalities. Decode being the affected phase while prefill is much less so points away from the GEMMs and toward the per-token expert gather, which is why this is recorded rather than dropped. It is a number plus an inference, not a mechanism, so it supports the section above rather than standing on its own.

That session reports both checkpoints as 256-expert top-6. Neither checkpoint is present on the GB10 host (`models/mlx` holds 190 checkpoints, none matching), so the expert counts and routing widths are unverified here and should be read from each `config.json` at re-measurement time.

### Limits, which are severe

- Metal only, one run per cell, not profiled, no kernel-level attribution.
- The two Inkling rows are **not** a quantization comparison. `inkling-small-mlx-4bit` carries `recipe: experts_only` while the NVFP4 build quantizes a different scope, so they are one checkpoint under two non-comparable recipes and must not be presented as 4-bit against NVFP4.
- Measured on mlxcel `b8fdddfa`, which is 124 commits behind `60341873`, against MLX pin `9a795735`. Main is now on `81ba1c6a`, and the Metal quantized kernels changed in between. `8ad62d4b` (#1772) rewrote 150 lines of `src/lib/mlx-cpp/patches/mlx/backend/metal/quantized.cpp`, including the tile choice, which went from `int bm = 64;` to `int bm = (transpose && M <= 32) ? 32 : 64;` under the upstream comment "Use smaller bm when one block covers all of M". Decode is `M = 1`, so that switch fires on exactly the shape this issue is about and did not exist when the measurement was taken. Eight diff lines in the same file touch `gather_qmm`: `gather_qmm_t_nax` and `gather_qmm_n_nax` gained a `global_scale` parameter and a `_hgs` kernel-name suffix. Do not write this off as "Metal was untouched". The mlxcel-side commits in the window are largely CUDA and ROCm, but the pin bump reaches every backend, and `src/models/switch_layers.rs` itself moved 231 lines (Laguna NVFP4 and Kimi K3 SiTU, from #1739 and #1741).
- The one thing that did **not** move: the `x_exp` expansion and the `do_sort` threshold are byte-identical across `b8fdddfa..60341873`, and the `deepseek_v4` and `inkling` call sites are unchanged. The hypothesis is still live on main; the kernels underneath it are not the ones that were measured. That is a reason to re-measure first, not a reason to discount the lead.

## Prior art, so it is not re-derived

- **#268** (closed, `docs/benchmark_results/moe-decode-gap-investigation.md`) is the general MoE decode gap against mlx-lm on Metal: worst case dots.llm1 at 0.51x, qwen3_moe 0.69x, nemotron-h 0.60x. It concluded the residual is diffuse per-op execution overhead with no single fusable lever, and it named a fused decode-MoE kernel as the highest-potential path. Every family in that table has that kernel today. 0.16 to 0.23x sits far outside its range, which is consistent with these two families lacking the escape hatch the others received, and is why this is not simply a duplicate of #268.
- **#1616** (closed) found MoE batched decode scaling to 1.55x against 3.25x on dense, traced it to the fused kernel being a B=1-only path, and already analysed the `do_sort` threshold and the one-row-per-(token, k) expert traffic. Its rejected alternative (forcing `supports_batching()` false for MoE) should not be relitigated.
- **#1799** is the source of the `M = 1, B = 8n` finding quoted above. **#1798** is the adjacent GB10 investigation from the same sweep, on the MLX graph ops budget; it is useful for the host, method and noise-floor conventions used there, but it is not where the gather-shape finding came from.

## Proposed Solution

Measure before changing anything, in this order.

1. **Re-measure on current main, on both backends.** Repeat the reporting session's matched shape on `60341873` or later on Metal, and add a CUDA arm on GB10, which requires fetching both checkpoints. Report n of at least 3 with per-arm min-to-max ranges and treat overlapping ranges as no effect, following the convention in `docs/benchmark_results/cuda-graph-budget-gb10-2026-09-12.md`. Include one family that does have the fused kernel as a control, so a general regression is distinguishable from a family-specific one.
2. **Attribute the gap by kernel, not by inference.** Confirm which kernel the routed experts actually take at decode on each backend (`MLXCEL_TRACE_ARCH` prints the quantized-matmul path on CUDA, per `docs/environment-variables.md:504`), and split the per-token cost into router, expert gather and combine with the `mlxcel-gpu-profiling` hooks already used in `docs/benchmark_results/moe-decode-gap-investigation.md`. The deliverable is a committed record under `docs/benchmark_results/` naming the kernels with their per-token launch counts and milliseconds.
3. **Only then decide what to build.** Candidates, in the order #268 and #1616 already ranked them: give these two families a fused path (which for `deepseek_v4` means a clamp stage the current kernel lacks, per `deepseek_v4_moe.rs:49-50`), deduplicate expert planes across the gathered rows, or re-tune the `do_sort` threshold. None should start before step 2 is committed.

**Named risk.** The obvious wrong turn is accepting the `M = 1` hypothesis because it is tidy and fits the numbers. The #1798 and #1799 investigations both opened with a plausible mechanism that the source then contradicted: #1799's own record corrects two of its issue's premises at `laguna-dflash-verify-cost-gb10-2026-09-11.md:96`, including a supposed `fp_qmv` to `qmm_sm80` family switch on the dense projections that turned out not to exist on that pairing because those projections are plain bf16. This issue requires a profiled attribution; a mechanism that merely fits is not an answer.

## Scope

**In scope:** a committed benchmark and profile record under `docs/benchmark_results/`; the kernel-level attribution on Metal and on CUDA; the decision about what, if anything, to build, written down with its evidence.

**Out of scope:** any edit to `src/models/switch_layers.rs`, `src/models/deepseek_v4_moe.rs` or `src/models/inkling/mlp.rs` before step 2 is committed; batched decode, which #1616 owns; the `bench_mlxlm.py` memory budget, filed separately.

## Acceptance Criteria

- [ ] A record under `docs/benchmark_results/` reports decode and prefill against the mlx-vlm baseline for both families on current main, on Metal and on CUDA, with host, MLX pin, tree commit, n and per-arm ranges.
- [ ] The record names the kernel the routed experts take at decode on each backend, with per-token launch counts and milliseconds, and gives the router / gather / combine split.
- [ ] The record includes a fused-kernel family as a control, so a family-specific gap is distinguishable from a general one.
- [ ] The record states explicitly whether the `M = 1` gather shape is or is not the dominant term, naming the measurement that settles it.
- [ ] Any resulting code change is reached by the real decode path used by `mlxcel-server` and `mlxcel generate`, not by a bench-only path, and either preserves greedy temp-0 token identity or documents exactly where it does not.
- [ ] A measured negative closes this issue, provided the measurement meets the criteria above.

## Verification

```bash
# Neither checkpoint is on the GB10 host; fetch before the CUDA arm.
ls models/mlx | grep -iE "deepseek-v4|inkling"

make release-cuda
# Metal: cargo build --release --features metal,accelerate --bin mlxcel-bench-decode

# Paired arms, same shape on both sides, n >= 3 after a discarded warm-up.
./target/release/mlxcel-bench-decode --model models/mlx/ --max-tokens 128 --warmup-tokens 20
MODELS_DIR=models/mlx ./scripts/bench_mlxlm.py models/mlx/

# Which quantized-matmul path the experts actually take (CUDA).
MLXCEL_TRACE_ARCH=1 ./target/release/mlxcel-bench-decode --model models/mlx/ --max-tokens 4

# Control: a family that does have the fused kernel, with it forced off.
MLXCEL_FUSED_MOE=0 ./target/release/mlxcel-bench-decode --model models/mlx/qwen3-30b-a3b-4bit --max-tokens 128
```

Contributor guide

Open the contributing guide

Research direction

Start with the verification commands in the issue, especially mlxcel-bench-decode and scripts/bench_mlxlm.py, after fetching the deepseek-v4 and inkling checkpoints. Read src/models/switch_layers.rs, src/models/deepseek_v4_moe.rs, src/models/inkling/mlp.rs, and the cited benchmark records before profiling both Metal and CUDA. Done means a committed docs/benchmark_results/ record with n≥3, kernel attribution, launch counts, router/gather/combine timings, and a measured conclusion about M = 1.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, rust, shell
Domain
backend, documentation, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.