From 9297d50da57ffec1dbb1dba0093e7c386e745423 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 22 Jul 2026 01:30:29 +0900 Subject: [PATCH] update(nemotron-h): add intra-MoE phase profiling Add opt-in forced-evaluation boundaries inside fused_moe_forward so router, routed-expert, combine, and shared-expert costs can be attributed independently. Keep normal decode unsynchronized, clarify why the routed squared-ReLU kernel was neutral, and document the M1 Ultra re-measurement and evidence gate for future fusion work. Validation: cargo fmt --all -- --check; cargo check -p mlxcel-core; cargo build -p mlxcel-core; cargo clippy -p mlxcel-core --lib -- -D warnings; cargo test -p mlxcel-core --lib ffi_tests::test_compiled_swiglu_activation -- --exact. The full library suite is not a valid backend-neutral gate on this host because existing paged-cache tests fail and a CUDA-only kernel aborts without CUDA. Refs #284 --- .../moe-decode-gap-investigation.md | 25 +++++- docs/environment-variables.md | 3 +- src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp | 89 +++++++++++++++++-- 3 files changed, 109 insertions(+), 8 deletions(-) diff --git a/docs/benchmark_results/moe-decode-gap-investigation.md b/docs/benchmark_results/moe-decode-gap-investigation.md index 0836da0b..ca48e944 100644 --- a/docs/benchmark_results/moe-decode-gap-investigation.md +++ b/docs/benchmark_results/moe-decode-gap-investigation.md @@ -85,7 +85,28 @@ fusion they largely collapse into adjacent kernels. This is the same class as the prior Gemma3n bf16 decode-gap investigation, where removing the AsType excess measured only ~+2.6%; the static node count overstates GPU kernels. -## 4. Dense gaps (#269): same overhead, milder +## 4. Nemotron-H intra-MoE attribution (#284) + +`fused_moe_forward` already owns the complete Nemotron-H MoE graph: router scoring and top-k selection, routed fc1/relu²/fc2, score-weighted combine, and the always-on shared up/relu²/down expert. The remaining useful profiling boundary is therefore inside that C++ graph, not another Rust-to-C++ fusion boundary. + +Set `MLXCEL_PROFILE_NEMOTRON_MOE=1` to force evaluation after each of those four phases. Each MoE call emits: + +```text +[MLXCEL_PROFILE_NEMOTRON_MOE] path=gather-qmm tokens=1 top_k=6 router_ms=... routed_ms=... combine_ms=... shared_ms=... total_ms=... +``` + +The values are attribution measurements, not production throughput: every phase boundary synchronizes the device, so the sum includes repeated scheduling and synchronization costs. Compare phase shares within the same run, then re-run the normal unsynchronized benchmark after changing any lever. The experimental `MLXCEL_FUSED_MOE_RELU2` path folds score multiplication into its down kernel, so its `routed_ms` includes that multiply and `combine_ms` covers only the final top-k reduction. + +Reproduce on the M1 Ultra with a warmed release binary: + +```bash +cargo build --release --features metal,accelerate --bin mlxcel-bench-decode +MLXCEL_PROFILE_NEMOTRON_MOE=1 ./target/release/mlxcel-bench-decode --model ./models/nemotron-h-30b-4bit --prompt "Hello, how are you today?" --max-tokens 16 --warmup-tokens 4 +``` + +This instrumentation sizes the remaining levers without claiming a speedup. The prior squared-ReLU kernel is neutral because it replaces only the routed fc1/fc2 GEMVs, while the larger always-on shared expert and router/combine stay unchanged. A dense shared-expert fusion or a router/gather/combine kernel should be built only if this phase split shows enough time to recover; each candidate must then pass greedy temp-0 parity and an unsynchronized `mlxcel-bench-decode` A/B against mlx-lm. + +## 5. Dense gaps (#269): same overhead, milder apertus (0.83×) and seed_oss (0.84×) have no model-specific bug. llama-8b is at parity, qwen3-8b (QK-norm) is 0.90×, apertus (QK-norm + xIELU) 0.83×. mlxcel's @@ -94,7 +115,7 @@ runtime `softplus` on GPU arrays). The residual tracks the per-layer extra small-op count, the same per-op execution overhead as the MoE case, just smaller. Subsumed by #268. -## 5. Conclusion and realistic paths +## 6. Conclusion and realistic paths The remaining gap is a single underlying phenomenon: mlxcel's per-op execution overhead, amplified by op density. Vanilla dense ≈ parity; QK-norm/xIELU add a diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 762b4890..4a62fb8b 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -248,7 +248,7 @@ recommended as normal deployment settings. | `MLXCEL_FUSED_MOE` | `0`/`false`/`off`/`no` disable; any other value or unset enables | on | Fused single-token decode-MoE kernel (#268), on by default since #282 (Metal) and #319 (CUDA, via `mx.fast.cuda_kernel`); validated on M1 Ultra, M5, and GB10. Set to `0` to force the proven `gather_qmm`/`SwitchGLU` path. Active for qwen3_moe, qwen3_next, dots.llm1, gemma4, qwen2_moe, mixtral, phimoe, lfm2, qwen3_vl_moe, and olmoe decode. | | `MLXCEL_FUSED_MOE_SGY` | `1`-`32` | `8` | Simdgroups (Metal) / warps-per-block (CUDA) per threadgroup for the fused decode-MoE kernel; tune per hardware. | | `MLXCEL_FUSED_MOE_MAX_DFF` | positive int | `4096` (Metal) / `8192` (CUDA) | Expert-intermediate (Dff) upper bound for the fused path; above it the caller falls back to `gather_qmm`. The fused path wins only while `gather_qmm` underutilizes the GPU (small experts), so the break-even is backend-dependent and the default is chosen from the live backend: `4096` on Metal (M1 Ultra tuning) and `8192` on CUDA (GB10 re-measured under MLX pin e9463bb, #626; fused wins through Dff 6400 and is break-even at 8192). An explicit value overrides the default on both backends: lower it to force `gather_qmm` sooner, raise it (e.g. `20000`) to force the fused kernel on larger experts such as mixtral (Dff 14336, where it is a slight net loss). | -| `MLXCEL_FUSED_MOE_RELU2` | presence enables | off | Enables the squared-ReLU fused MoE path for nemotron-class experts; performance-neutral on nemotron-h, kept for a future MoE-dominated squared-ReLU model. | +| `MLXCEL_FUSED_MOE_RELU2` | presence enables | off | Enables the squared-ReLU fused routed-expert path for nemotron-class experts. It is performance-neutral on nemotron-h even though MoE is its largest block because it replaces only the already-efficient routed fc1/fc2 GEMVs; the router, always-on shared expert, and combine remain outside this kernel. Kept for a future squared-ReLU model whose routed-expert slice dominates. | | `MLXCEL_GATHER_QMM_GROUPED` | `0` disables; any other value or unset enables | on | **CUDA only.** Gates the sorted MoE prefill fast path in `GatherQMM::eval_gpu` (issue #629). When the sorted-indices `M == 1` prefill contract holds (right-sorted, transpose, one activation row pre-gathered per (token, expert) pair, non-nvfp4, float activation dtype, `E <= 1024`) and the batch is large enough to amortize (see `MLXCEL_GATHER_QMM_GROUPED_MIN_ROWS`), the expert weight stack is dequantized once and routed through `cutlass_grouped_gemm_unaligned` instead of one 1-row `qmm_sm80` GEMM per (token, expert) pair. Fixes a 5-10x CUDA MoE prefill collapse relative to Metal M1 Ultra; measured 3.6-40x prefill speedup across mixtral-8x7b, phi-3.5-moe, llama-4-scout, minimax-m2, gpt-oss-20b, and solar-open-100b on GB10, with no decode-path regression. Set to `0` to force the legacy per-row dispatch (A/B, rollback). See `docs/benchmark_results/moe-prefill-grouped-gemm-gb10-2026-07-10.md`. | | `MLXCEL_GATHER_QMM_GROUPED_MIN_ROWS` | positive integer | `8` | **CUDA only.** Amortization threshold for `MLXCEL_GATHER_QMM_GROUPED`: the fast path activates only once the sorted batch size `B` reaches `min_rows * E` (E = expert count), the point past which the one-time expert dequant traffic is cheaper than the legacy per-row re-reads. Lower it to trigger the fast path on smaller sorted batches (e.g. high-concurrency batched decode crossing the same sort gate); raise it to keep more traffic on the legacy path for tuning or rollback. | | `MLXCEL_QMV_MULTIROW` | `0` disables; any other value or unset enables | on | **CUDA only.** Gates the weight-amortizing multirow qmv path (issue #725). With broadcast weights and `2 <= M*B <= 8` input rows (batched decode with B in [2,8), speculative-verify `[1, K]` forwards), one warp applies each dequantized weight tile to every input row instead of launching one weight-rereading block column per row, so weight DRAM traffic is O(1) in the row count. Per-row outputs are bit-identical to the stock per-row launches (pinned by `qmv_multirow_matches_per_row_qmv_bitwise`); classic `M*B == 1` decode is untouched. On GB10 this flips B=1 MTP speculative decoding from 0.52-0.77x to 1.31-1.46x and lifts 4-client aggregate serving decode from ~50 to ~74 tok/s on llama-3.1-8b-4bit. Set to `0` to force the stock per-row dispatch (A/B, rollback). See `docs/benchmark_results/qmv-multirow-gb10-2026-07-11.md`. | @@ -280,6 +280,7 @@ throughput measurements. Use them for diagnosis, not capacity planning. | `MLXCEL_PROFILE_PIPELINE_DETAIL` | presence enables | off | Adds per-step pipeline timing detail. | | `MLXCEL_PROFILE_BLOCKS` | presence enables | off | Emits per-block/model-family timing where implemented. | | `MLXCEL_PROFILE_FORWARD` | presence enables | off | Enables model-specific forward profiling where implemented. | +| `MLXCEL_PROFILE_NEMOTRON_MOE` | presence enables | off | Forces evaluation at the router, routed-expert, combine, and shared-expert boundaries inside `fused_moe_forward`, then emits one timing line per MoE call. Diagnostic only: the forced synchronization changes absolute throughput. On the experimental `MLXCEL_FUSED_MOE_RELU2` path, `routed_ms` includes the score multiply folded into the down kernel and `combine_ms` is only the final top-k reduction. | | `MLXCEL_PROFILE_QWEN3_MOE_DETAIL` | presence enables | off | Profiles Qwen3 MoE internals. | | `MLXCEL_PROFILE_MOE_INNER` | presence enables | off | Profiles Gemma 4 MoE sub-operations. | | `MLXCEL_PROFILE_PER_LAYER` | presence enables | off | Prints per-layer Gemma 4 timing. | diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp b/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp index a7bf9b04..a249b4cf 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp @@ -9,6 +9,8 @@ // header is backend-agnostic: builds without CUDA link the no_cuda stub. #include "mlx/backend/cuda/cuda.h" +#include + namespace mlx_cxx { // ── BitLinear ternary matmul (BitNet b1.58) ──────────────────────────────── @@ -1093,6 +1095,21 @@ std::unique_ptr fused_moe_forward( ) { using namespace mlx::core; + // Forced evaluation is intentionally opt-in: it serializes each phase and + // therefore measures attribution, not production throughput. Cache the + // switch so the normal decode path pays only this predictable branch. + static const bool profile_nemotron_moe = + std::getenv("MLXCEL_PROFILE_NEMOTRON_MOE") != nullptr; + using ProfileClock = std::chrono::steady_clock; + ProfileClock::time_point profile_checkpoint; + double router_ms = 0.0; + double routed_ms = 0.0; + double combine_ms = 0.0; + double shared_ms = 0.0; + if (profile_nemotron_moe) { + profile_checkpoint = ProfileClock::now(); + } + // 1. Gate: compiled sigmoid + topk + compiled normalize + scale auto gates = matmul(x.inner, transpose(gate_weight.inner)); @@ -1113,6 +1130,16 @@ std::unique_ptr fused_moe_forward( } else { topk_scores = multiply(topk_scores, array(scaling_factor)); } + if (profile_nemotron_moe) { + // Both outputs are consumed below. Evaluating them at the boundary + // charges gate matmul, scoring, top-k, and normalization to the router. + topk_indices.eval(); + topk_scores.eval(); + auto now = ProfileClock::now(); + router_ms = std::chrono::duration( + now - profile_checkpoint).count(); + profile_checkpoint = now; + } // 2. SwitchMLP: expand + gather_qmm(fc1) + relu² + gather_qmm(fc2) + squeeze. auto x_shape = x.inner.shape(); @@ -1121,11 +1148,12 @@ std::unique_ptr fused_moe_forward( // Experimental fused squared-ReLU decode path (#268), behind its own flag // MLXCEL_FUSED_MOE_RELU2 (NOT the default MLXCEL_FUSED_MOE): fc1 + relu² -> // act_g[K, Dff], then reuse moe_down for fc2 * score. Correct and - // byte-identical, but measured performance-NEUTRAL on nemotron-h-30b (its - // decode is dominated by Mamba2 + attention, so the MoE expert GEMV is a - // small, already-efficient slice). Kept wired behind the dedicated flag so - // the kernel stays referenceable for a future squared-ReLU model whose - // decode is MoE-dominated; the default path below stays on gather_qmm. + // byte-identical, but measured performance-NEUTRAL on nemotron-h-30b. MoE + // is its largest block, but this kernel replaces only the already-efficient + // routed fc1/fc2 GEMVs; the router, shared expert, and combine remain. Kept + // wired behind the dedicated flag so it stays referenceable for a model + // where that narrower routed-expert slice dominates; the default path below + // stays on gather_qmm. bool fused_relu2 = x_shape[0] == 1 && (bits == 4 || bits == 8) && std::getenv("MLXCEL_FUSED_MOE_RELU2"); array result = x.inner; // placeholder; overwritten in both branches @@ -1169,7 +1197,23 @@ std::unique_ptr fused_moe_forward( std::make_tuple(32, round_up(din, sgy), k), std::make_tuple(32, sgy, 1), ta, std::nullopt, false, {}); + if (profile_nemotron_moe) { + // moe_down folds score multiplication into fc2 on this experimental + // path. routed_ms includes it; combine_ms is the final K reduction. + rB[0].eval(); + auto now = ProfileClock::now(); + routed_ms = std::chrono::duration( + now - profile_checkpoint).count(); + profile_checkpoint = now; + } result = reshape(sum(rB[0], 0, false), {1, din}); + if (profile_nemotron_moe) { + result.eval(); + auto now = ProfileClock::now(); + combine_ms = std::chrono::duration( + now - profile_checkpoint).count(); + profile_checkpoint = now; + } } else { auto x_exp = reshape(x.inner, {x_shape[0], 1, 1, x_shape[1]}); @@ -1184,6 +1228,13 @@ std::unique_ptr fused_moe_forward( std::nullopt, topk_indices, true, group_size, bits, "affine", false); h = squeeze(h, -2); // [tokens, top_k, hidden] + if (profile_nemotron_moe) { + h.eval(); + auto now = ProfileClock::now(); + routed_ms = std::chrono::duration( + now - profile_checkpoint).count(); + profile_checkpoint = now; + } // 3. Score weighting: weighted sum over experts, cast back to input dtype // Cast scores to h's dtype to avoid mixed float32×float16 multiply @@ -1191,6 +1242,13 @@ std::unique_ptr fused_moe_forward( auto scores_exp = reshape(topk_scores, {topk_scores.shape()[0], top_k, 1}); auto scores_cast = astype(scores_exp, h.dtype()); result = astype(sum(multiply(h, scores_cast), -2, false), x.inner.dtype()); + if (profile_nemotron_moe) { + result.eval(); + auto now = ProfileClock::now(); + combine_ms = std::chrono::duration( + now - profile_checkpoint).count(); + profile_checkpoint = now; + } } // 4. Optional shared expert @@ -1205,6 +1263,27 @@ std::unique_ptr fused_moe_forward( shared_down_biases ? std::optional(shared_down_biases->inner) : std::nullopt, true, group_size, bits); result = add(result, shared_h); + if (profile_nemotron_moe) { + result.eval(); + auto now = ProfileClock::now(); + shared_ms = std::chrono::duration( + now - profile_checkpoint).count(); + profile_checkpoint = now; + } + } + + if (profile_nemotron_moe) { + std::cerr + << "[MLXCEL_PROFILE_NEMOTRON_MOE] path=" + << (fused_relu2 ? "fused-relu2" : "gather-qmm") + << " tokens=" << x_shape[0] + << " top_k=" << top_k + << " router_ms=" << router_ms + << " routed_ms=" << routed_ms + << " combine_ms=" << combine_ms + << " shared_ms=" << shared_ms + << " total_ms=" << (router_ms + routed_ms + combine_ms + shared_ms) + << std::endl; } return std::make_unique(std::move(result));