update(nemotron-h): add intra-MoE phase profiling (#284)

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

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

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 <chrono>
+
 namespace mlx_cxx {
 
 // ── BitLinear ternary matmul (BitNet b1.58) ────────────────────────────────
@@ -1093,6 +1095,21 @@ std::unique_ptr<MlxArray> 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<MlxArray> 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<double, std::milli>(
+            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<MlxArray> 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<MlxArray> 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<double, std::milli>(
+                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<double, std::milli>(
+                now - profile_checkpoint).count();
+            profile_checkpoint = now;
+        }
     } else {
         auto x_exp = reshape(x.inner, {x_shape[0], 1, 1, x_shape[1]});
 
@@ -1184,13 +1228,27 @@ std::unique_ptr<MlxArray> 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<double, std::milli>(
+                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
         // which produces NaN on M5 Max (Metal GPU Family 4) NAx broadcast kernel.
         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<double, std::milli>(
+                now - profile_checkpoint).count();
+            profile_checkpoint = now;
+        }
     }
 
     // 4. Optional shared expert
@@ -1205,6 +1263,27 @@ std::unique_ptr<MlxArray> 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<double, std::milli>(
+                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<MlxArray>(std::move(result));