[CUDA] Skip FP4 QMoE fc1 activation expansion

[Tianlei Wu (tianleiwu)]

Description

Stacked on #31159; review only the top commit.

This removes the standalone FP4 QMoE fc1 activation expansion during GEMV decode. Instead, fc1 maps each permuted row back to its source token with permuted_row_to_source_row[row] % num_rows; fc2 remains unchanged because it consumes the expanded fc1 output.

Summary of Changes

• Pass the permuted-row-to-source-row mapping through all FP4 interleaved SwiGLU GEMV launch variants.
• Read fc1 activations directly from the original input while preserving PR 31159's SM80 pair-interleaved weight layout.
• Keep the legacy path available with ORT_DISABLE_FP4_GEMV_SKIP_EXPAND=1 for same-binary comparison.
• Add an exact-output parity test for the MTP shape (num_tokens=3, top_k=8) with skip-expand enabled and disabled.

Performance

H200, Qwen3.6 35B A3B NVFP4, MTP N=3, paired same-binary A/B:

• Removes 40 expandInputRows launches per decoding step.
• Removes 94.651 us/step of activation expansion.
• Adds 10.377 us/step to fc1 source-row lookup.
• Saves 84.274 us/step across named kernels.
• Saves 0.097695 ms/step median end-to-end (1.314%).

Testing

• lintrunner -a on the five changed files.
• Compiled moe_gemv_fp4.cu and moe_quantization.cc against the exact [CUDA] Speed up the NVFP4 QMoE decode GEMV and enable it for MTP verify #31159 head.
• New expanded-vs-skip-expand regression passes with exact tensor equality.
• 23 NVFP4 QMoE CUDA tests pass locally; the separately isolated large-input scaling test fails identically with skip-expand enabled and disabled on the pre-existing integration binary.

Checklist

• Tests added/updated
• No breaking changes
• No documentation changes required

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

[Ti-Tai Wang (titaiwangms)]

Multi-model review (5 reviewers: readability, correctness, adversarial, deep/spec, integration)

Verdict: the index math is correct — I tried hard to break it and could not. Three reviewers independently confirmed the % num_rows identity, and no reviewer found a Critical issue. Findings below are about a missing structural guard, docs/convention drift, and a test that doesn't actually pin down the behavior it claims to.

Confirmed correct (so nobody re-litigates it)

• The modulo identity holds. The producer writes unpermuted_row = i * num_tokens + token (moe_kernels.cu:424), the prose contract states the modulo recovery explicitly (moe_kernels.cu:1286-1289), and the deleted kernel's own gather used the identical unpermuted_row % num_tokens (moe_kernels.cu:1351-1353). The new inline gather is numerically identical to what expandInputRows did. / top_k would have been wrong; % num_rows is right.
• expandInputRows is a pure gather for this instantiation. With T = half|__nv_bfloat16, prequant_scales == nullptr, dispatch lands in BlockScalingType::NONE, PRE_QUANT_AWQ=false (moe_kernels.cu:1524-1526), so every fp4/mxfp8 quantize, SF-padding, AWQ-scaling and permuted_scales write is if constexpr-dead. Only the row copy at :1424 was live. Nothing is lost by skipping it.
• No OOB / sentinel hazard. The expert-validity early-return (moe_gemv_device.cuh:288-290) sits before the permuted_row_to_source_row[row] load at :314, grid X equals expanded_num_rows, and produced values are non-negative — so no negative-modulo and no read past the valid range.
• Padding invariant preserved. Both old and new paths leave p_fc1_buf rows [num_valid, expanded) untouched, and finalizeMoeRouting never addresses them.
• No contract drift in signatures. One external caller, two explicit instantiations, one declaration — all updated. The same-named templates in moe_gemv.cu are a separate TU with internal linkage, so no ODR/overload hazard.

---

Major

1. The static_assert(CtaM == 1) guard that makes this transformation safe was dropped.
moe_gemv_device.cuh:314

The already-merged INT-path equivalent guards the identical rewrite, and even documents the mapping:

// moe_gemv.cu:20-36
template <int CtaM>
__device__ __forceinline__ int act_source_row(const int* permuted_row_to_source_row, int num_rows,
                                              int row, int offset_m) {
  if (permuted_row_to_source_row == nullptr) return offset_m;
  static_assert(CtaM == 1, "source-row indirection assumes one expanded row per block");
  return permuted_row_to_source_row[row] % num_rows;
}

The FP4 version inlines the expression with neither the guard nor the comment. It is correct today only because CtaM is hardcoded to 1 47 lines earlier (moe_gemv_device.cuh:267) — the invariant "one expanded row per block" is now enforced by nothing. Since MoeGemvConfig already sweeps CtaN/Threads, promoting CtaM to a tiling knob is a natural next step, and it would silently produce wrong results for rows 1..CtaM-1 rather than a compile error.

Fix: hoist act_source_row<CtaM> out of moe_gemv.cu into moe_gemv_device.cuh and call it from both paths — one definition, one comment, one guard. At minimum, add the static_assert inline.

2. The env knob is parsed on every ComputeInternal, breaking a contract the docs explicitly make.
moe_quantization.cc:66-68 and :1474

docs/contrib_ops/cuda/moe_qmoe.md:783 states: "The FP4 path is gated by several process-start environment variables (read once in the QMoE constructor)." Every sibling knob follows that — Fp4GemvAutotuneEnabled() / Fp4GemvAutotuneLogEnabled() are latched into enable_fp4_gemv_autotune_ / enable_fp4_gemv_autotune_log_ members. This one does a getenv + string construction + parse per forward pass, on the decode hot path this PR is optimizing for ~10 µs/step, and it lets a mid-session setenv change the numerics of an already-initialized session.

Fix: latch it in the constructor as disable_fp4_gemv_skip_expand_, matching the neighbors.

3. ORT_DISABLE_FP4_GEMV_SKIP_EXPAND is not in the env-var registry.
docs/contrib_ops/cuda/moe_qmoe.md §9.9

Every other ORT_FP4_GEMV_* / ORT_ENABLE_FP4_* knob is documented there. Without an entry (and without a CI leg that ever sets it), the legacy branch becomes undocumented, untested code with no stated removal plan.

4. The new parity test cannot distinguish the indirection it exists to protect.
test_qmoe_nvfp4_cuda.py:822-849

With top_k == num_experts == 8, every expert gets every token, so permuted_row == e * num_rows + t exactly. Enumerating what the test would catch:

candidate implementation	caught?
p_r2u[row] % num_rows (correct)	—
p_r2u[row] / top_k	✅ caught
identity (no reduction)	✅ caught
row % num_rows — ignoring p_r2u entirely	❌ not caught

The shape validates the modulus but not the indirection — which is the entire reason p_r2u is threaded through five files. The regime where p_r2u[row] genuinely diverges from row is sparse routing (top_k < num_experts), and it is untested. Adding top_k=2, num_experts=8, num_tokens=32 (the existing test_nvfp4_fp16_silu_basic shape at :561-568) plus a BFLOAT16 variant is cheap and makes the test load-bearing. BF16 matters independently: moe_gemv_fp4.cu:361 instantiates the changed template for __nv_bfloat16 and nothing covers it.

5. The test has no positive control — it passes green if the fused GEMV never dispatched.
test_qmoe_nvfp4_cuda.py:822

The dispatch gate at moe_quantization.cc:1385-1402 can fail silently (falling back to dequant) without raising. Both arms then take the same fallback and torch.equal trivially passes on a machine where the changed code never ran. Session/env caching is not the problem here — the test does create fresh sessions and reseed (:318-319, :373-407) — the gap is that nothing asserts the FP4 GEMV path was actually taken. Contrast the neighboring test_nvfp4_gemv_default_tiling_optout (:800-820), which shells out and inspects stdout. Suggest enabling ORT_FP4_GEMV_AUTOTUNE_LOG=1 and asserting the fc1 candidate line appears. Note this also gates finding #4: without a positive control, the added shapes aren't guaranteed to be exercised either.

6. The header doc now describes the wrong buffer.
moe_gemv_fp4.h:92-101

The doc block still inherits act: [expanded_num_rows, k] permuted activations from the sibling launch_moe_gemv_fp4_symmetric above it. When permuted_row_to_source_row != nullptr, act is instead the original [num_rows, k] unexpanded input — the opposite shape and semantics. The two new parameters are also undocumented: that nullptr selects the legacy path, and that num_rows here is the unexpanded token count while expanded_num_rows in the same signature is not. A reader working from the header will pass the wrong buffer.

---

Minor

• moe_quantization.cc:1417 — p_act_buf is still allocated (expanded * hidden * elt) but is dead on the default path. Its only two uses (:1477 expand, :1484 fc1 input) are both on the disabled branch. Compute skip_expand before the scratch allocations and size it to 0 when skipping — otherwise the PR keeps the peak allocator pressure it just stopped using.
• Naming: permuted_row_to_source_row is a third synonym for an existing concept. The buffer is p_r2u / permuted_row_to_unpermuted_row everywhere else (moe_util_kernels.h:58,66; moe_kernels.cu:426,641,933). Reusing the established name would make it obvious this is the same array already handed to expandInputRowsKernelLauncher.
• Fp4GemvSkipExpandDisabled() is a double negative (moe_quantization.cc:66), forcing const bool skip_expand = !Fp4GemvSkipExpandDisabled(); at :1474. Sibling knobs are all ...Enabled(). Prefer Fp4GemvSkipExpandEnabled() and use it directly.
• moe_gemv_device.cuh:314 has no comment explaining % num_rows. The modulo is doing real, non-obvious work (the value encodes k_slot * num_rows + token); as written a reader can't tell it from a bug masking an out-of-range index. The moe_gemv.cu:20-27 comment block already says it well — reuse it.
• Stale comment at moe_quantization.cc:1530: // fc1 reads p_act_buf (populated by the expand above). — false on the default path now.
• moe_gemv_fp4.cu:298 launcher takes int64_t num_rows, narrows to int, and the kernel uses it as a divisor. Safe today only because moe_quantization.cc:1402 enforces 0 < num_rows <= 256, an invariant the launcher neither documents nor enforces. A future caller could get modulo-by-zero or negative indexing.
• No rationale comment on the new test. It's the only parity test in the file using bit-exact torch.equal rather than a tolerance; worth stating why (skip-expand changes which pointer fc1 reads, not any arithmetic, so the paths must be identical, not merely close).

Questions

• const_cast<T*>(act) on a graph-input tensor (moe_gemv_fp4.cu:313,318,337): the kernel provably never writes through it (act_iterator is load-only, details.h:62), so no UB — but that pointer used to be ORT-owned scratch and is now a genuine upstream tensor, so the aliasing contract changed. Making the kernel's act parameter const TypeA* would encode the guarantee.
• num_rows is int64_t in moe_gemv.cu's wrappers but int in moe_gemv_device.cuh's (moe_gemv.cu:705,730,824 vs moe_gemv_device.cuh:401,452). Both narrow at the kernel boundary so behavior is identical — intentional, or drift that will bite whoever merges the duplicated machinery?
• fusedBuildExpertMapsSortFirstToken's bool return is discarded (moe_quantization.cc:1430). It returns false having launched nothing when num_experts > 511 (moe_kernels.cu:571-573), leaving p_r2u uninitialized. Pre-existing — but this PR promotes p_r2u from "input to a bounded gather" to "an index the GEMV dereferences directly", so it's worth confirming that case is unreachable here.

---

Nice result on the perf side, and the fact that the INT path already landed the same transformation makes this a low-risk port. The two things I'd genuinely want before merge are #1 (the guard — it's three lines and prevents a silent-wrong-answer regression) and #4/#5 (make the test actually able to fail).

_{🤖 Reviewed with a 5-model review team (Claude Opus 5 · GPT-5.3-Codex · GPT-5.6 · Gemini 3.1 Pro), findings verified against the source.}

[kunal-vaishnavi]

From Copilot:

Could we extend this parity test to cover at least one MXFP4 case and one BF16 case as well? The skip-expand implementation is shared across MXFP4/NVFP4 and FP16/BF16 activation paths, but this test currently exercises only FP16 NVFP4. A format- or dtype-specific scale/layout regression could therefore pass unnoticed.

[kunal-vaishnavi]

From Copilot:

Could we also validate both outputs against the existing dequantized reference, rather than relying only on torch.equal between the two runs? Since both executions share the same generated weights, routing inputs, and GEMV implementation, this assertion primarily verifies that the two paths address the activation rows identically. Comparing each result against the reference would independently confirm the numerical correctness of the skip-expand path as well. If exact equality is intentional here, please add a comment explaining why it is expected and what this assertion is meant to cover.

[kunal-vaishnavi]

From Copilot:

Could we either capture the skip_expand decision at session construction time or document that ORT_DISABLE_FP4_GEMV_SKIP_EXPAND is intentionally read at inference time? Most neighboring FP4 GEMV environment options are latched when the session is created, whereas this option can change between calls. Runtime mutability is useful for A/B testing, but the behavior should be explicit and consistent with the intended configuration model.

[kunal-vaishnavi]

From Copilot:

Could you update this comment to account for the skip-expand path? With skip-expand enabled, FC1 reads directly from the original input using p_r2u rather than from p_act_buf, so the current wording may imply that p_act_buf must be initialized during autotuning.

[kunal-vaishnavi]

From Copilot:

Could we make the routing inputs deterministic for this regression and include cases with repeated source rows and varied expert ordering? The current test relies on one random router result, so it may not exercise all % num_rows mappings. Explicitly covering repeated tokens routed to multiple experts and nontrivial expert orderings would better validate that the permutation map is populated and interpreted correctly.