lablup / lablup/mlxcel

perf(core): host graph build at decode is serial with device execution and scales with op count

Open
#1,713 0 comments 0 reactions 0 assignees View on GitHub
arch:hybrid arch:moe area:benchmark area:core area:inference area:models platform:macos priority:high status:ready type:performance
Dominant language
Rust
Stars
467
Forks
54
Avg merge
4h 25m
Merged PRs (30d)
310

Description

## Problem / Background

Decode pays a host-side graph build that is serial with device execution, and its cost is proportional to the number of MLX ops a model's decode step constructs, not to parameter count. On small models and on SSM/MoE hybrids that build many ops per layer, that share is large enough to hold decode below mlx-lm.

Environment for every number below: mlxcel 0.7.0-beta.1 on branch `bench/0.7.0-refresh` at commit 74ffdf94, MacBook Pro M5 Max 128GB, macOS 26.6.2, MLX pin 9a795735, `cargo build --release --features metal,accelerate`, measured with `./target/release/mlxcel-bench-decode --prompt "Hello, how are you today?" --max-tokens 128 --prompt-tokens 512` and `MLXCEL_PROFILE_PIPELINE_DETAIL=1`, on an idle machine. Ratios are against `benchmarks/pylm_m5max_2026-09-06.csv`.

Commit 74ffdf94 wired the existing `[PIPELINE_DETAIL]` breakdown into `generate_with_stats` (`src/lib/mlxcel-core/src/generate.rs:2536-2649`), the path `mlxcel-bench-decode` takes, so the decode step splits into reshape, forward (host graph build), sample, async_eval (device work) and item_wait. Per token, in ms:

| Checkpoint | forward | async_eval | vs mlx-lm |
|---|--:|--:|--:|
| `qwen3-0.6b-4bit` | 0.110 | 1.622 | 91.0% |
| `qwen3.5-0.8b-4bit` | 0.203 | 1.700 | 93.5% |
| `stablelm-1.6b-4bit` | 0.143 | 3.311 | 93.9% |
| `hunyuan-1.8b-4bit` | 0.121 | 3.058 | 94.4% |
| `granite-4.0-h-350m-4bit` | 0.284 | 1.790 | 82.4% |
| `granite-4.0-h-tiny-4bit` | 0.517 | 4.247 | 88.5% |
| `qwen3-4b-4bit` | 0.165 | 5.273 | 100.3% |
| `qwen3-8b-4bit` | 0.148 | 8.832 | 99.4% |

Two facts. First, the components sum to the measured per-token time (qwen3-0.6b-4bit: 0.110 + 1.622 + 0.205 item_wait + 0.001 = 1.938 against 1.927 measured from 519.05 tok/s), so host build is serial with device execution rather than overlapped by the lookahead. Second, `forward` is roughly flat with model size while `async_eval` scales with it, so the same fixed host cost is 1.6% of a token on the 8B and 5.7% on the 0.6B. The per-token gap the low-ratio models show against mlx-lm, 0.125 to 0.216 ms, is the same magnitude as their `forward` term.

## Current Behavior

The cost is op count, not per-op cost. Exporting the first decode step's graph with `MLXCEL_EXPORT_DECODE_DOT` (`src/lib/mlxcel-core/src/generate.rs:1869`) gives 3207 edges for qwen3-8b-4bit against 9257 for granite-4.0-h-tiny-4bit, at 36 and 40 layers, so 89 against 231 nodes per layer. That 2.6x node ratio sits against the 3.5x ratio in `forward` time. The op-kind histogram puts granite's extra nodes in shape and constant work rather than arithmetic:

| Op kind | qwen3-8b-4bit | granite-4.0-h-tiny-4bit |
|---|--:|--:|
| Broadcast | 0 | 428 |
| AsType | 0 | 224 |
| ExpandDims | 0 | 120 |
| Arange | 0 | 120 |
| Full | 0 | 118 |
| Reshape | 147 | 471 |
| Slice | 180 | 416 |

## Ruled out (do not repeat this work)

- **The FFI bridge is not it.** `cargo run --release --example bridge_overhead_microbench` against `scripts/bridge_overhead_microbench_py.py` gives 0.11 to 0.34us per call in Rust against 0.13 to 0.40us in Python, with the eval paths agreeing within 1%.
- **The decode-specialized short conv is not it.** `src/models/lfm2.rs:386-389` deliberately leaves `decode_weight` unset on Metal because `conv1d` already dispatches a fast small-conv kernel there, so `src/models/conv_decode.rs` is a CUDA fix and was already evaluated for Metal.
- **Constant creation in the Mamba2 mixer is not it.** `src/models/granitemoehybrid.rs:510-512` builds the `time_step_limit` bounds with `full_f32` and `:505-526` casts to f32 repeatedly, but those sit in `ssm_step`, which serves prefill. Decode dispatches to `ssm_step_kernel` at `src/models/granitemoehybrid.rs:430` when `seq_len == 1` and the SSM Metal kernel is available, and that kernel takes the limits and `a_log` as scalars without building arrays.

## Proposed Solution

Cut the number of MLX ops the decode step constructs, starting with the shared MoE helper. Stated as a candidate, not a conclusion: `granitemoehybrid.rs` routes through `SwitchGLU` and `moe_weighted_sum` in `src/models/switch_layers.rs`, which stacks `expand_dims` several deep at `:768-769`, `:808-809` and `:814-824` and builds a `full_f32` at `:1261`. That would account for granite-4.0-h-tiny's extra ExpandDims and Full. Attribute the histogram rows to concrete call sites against the exported graph first, then hoist or remove the ops that turn out to be per-step constants or redundant rank adjustments. The fix belongs in the shared helper with a cross-family control, not in `granitemoehybrid.rs`, because any change there touches every MoE family.

## Scope

**In scope:** `src/models/switch_layers.rs` (op count in `SwitchGLU::forward`, `forward_with_expert_scales`, `moe_weighted_sum`, `group_mask_scores`), plus measurement through `MLXCEL_PROFILE_PIPELINE_DETAIL` and `MLXCEL_EXPORT_DECODE_DOT`.

**Out of scope:** the M5-gated per-mixer eval under investigation in #1685; CUDA-side conv decode in `src/models/conv_decode.rs`; overlapping host build with device execution, which is a pipelining change and needs its own issue if op-count reduction is not enough.

## Implementation Notes

- **Reuse**: measure with the existing `[PIPELINE_DETAIL]` instrumentation and `MLXCEL_EXPORT_DECODE_DOT`, documented at `docs/environment-variables.md:499` and `:509`. Do not add new counters.
- **Constraint**: `switch_layers.rs` is a shared hotspot per `docs/code-guidelines.md`. Update the `// Used by:` comment on any function changed and re-check its callers (`qwen3_moe.rs`, `deepseek_v32.rs`, `deepseek_v4_moe.rs`, `glm4_moe.rs`, `klear.rs`, `bailing_moe.rs`, `phixtral.rs`, `mellum.rs`, `solar_open.rs`, `qwen3_vl_moe.rs`).
- **A control is mandatory**: large MoE models are already above the reference (`qwen3-30b-a3b-4bit` at 119%, `qwen3.5-35b-a3b-4bit` at 107%), which says the helper is not slow in absolute terms and the problem is a fixed build cost against a short device step.
- **Structural floor**: `qwen3-0.6b-4bit` and `hunyuan-1.8b-4bit` already build at 3.9 and 3.8 us per layer, the same floor every dense transformer here reaches (`qwen3-8b` 4.1, `qwen3-4b` 4.6). Their ratio is low because the device step is short, not because their build is expensive, so they are not the target and may not move.
- **Arithmetic must not change**: removing shape and constant ops should be exactly output-preserving, so the bar is byte-identical greedy output rather than a disagreement budget.

## Acceptance Criteria

- [ ] Decode graph node count for at least one affected model is reduced, with before and after counts from `MLXCEL_EXPORT_DECODE_DOT` reported next to `forward` ms/token from `MLXCEL_PROFILE_PIPELINE_DETAIL`.
- [ ] Decode tok/s re-measured at pp512/tg128 for `granite-4.0-h-350m-4bit` and `granite-4.0-h-tiny-4bit`, reported as a ratio against `benchmarks/pylm_m5max_2026-09-06.csv`.
- [ ] Control: `qwen3-30b-a3b-4bit` and `qwen3.5-35b-a3b-4bit` re-measured at pp512/tg128 and shown not to regress.
- [ ] Greedy output unchanged at `--temp 0` on every model touched, including one non-Granite MoE family from the caller list.
- [ ] The change lands in the shared helper on the real decode path, not behind a benchmark-only or opt-in flag, and `cargo test --release`, `cargo clippy --all-targets -- -D warnings` and `cargo fmt --check` pass.

## Verification

```bash
cargo build --release --features metal,accelerate

MLXCEL_PROFILE_PIPELINE_DETAIL=1 ./target/release/mlxcel-bench-decode \
-m models/granite-4.0-h-tiny-4bit --prompt "Hello, how are you today?" \
--max-tokens 128 --prompt-tokens 512

MLXCEL_EXPORT_DECODE_DOT=/tmp/granite.dot ./target/release/mlxcel-bench-decode \
-m models/granite-4.0-h-tiny-4bit -p Hi -n 3 --warmup-tokens 3 \
--no-chat-template >/dev/null 2>&1
grep -c -- ' -> ' /tmp/granite.dot
grep -oE 'label ="[A-Za-z0-9_]+' /tmp/granite.dot | sed 's/label ="//' | sort | uniq -c | sort -rn | head -30

./target/release/mlxcel generate -m models/granite-4.0-h-tiny-4bit -p "Hello" -n 50 --temp 0

cargo test --release && cargo clippy --all-targets -- -D warnings && cargo fmt --check
```

A pass looks like: edge count down on the affected model, `forward` ms/token down by a matching share, decode tok/s up for both Granite checkpoints against the CSV, the two large MoE controls flat within run-to-run noise, and byte-identical `--temp 0` output before and after.

## Technical Considerations

Related to #1685 (FalconH1 and GraniteMoeHybrid decode on M5), which investigates the M5-gated per-mixer eval as a separate cause. The Granite ratios here (82.4% and 88.5%) are far above the 23% and 33% recorded there, so this covers the residual host-build share, not the same defect. Keep them apart: #1685 is a per-mixer eval question inside the SSM path, this is an op-count question about the graph the whole decode step builds.

Contributor guide

Open the contributing guide

Research direction

Start in src/models/switch_layers.rs, tracing SwitchGLU, forward_with_expert_scales, moe_weighted_sum, and group_mask_scores through the listed MoE callers. Use MLXCEL_EXPORT_DECODE_DOT and MLXCEL_PROFILE_PIPELINE_DETAIL to attribute graph nodes and forward time before changing the shared helper. Done means fewer decode nodes, improved Granite measurements without control regressions, byte-identical greedy output, and passing the listed cargo checks.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
machine-learning, performance
Issue type
Refactor
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.