Repository navigation
[CUDA] Extend MoEPacked for top-k, dispatch and fused MoE calculation - #33963
Conversation
520f375 to
b9b776b
Compare
b9b776b to
c818493
Compare
736d86e to
a9bdbce
Compare
This comment was marked as resolved.
This comment was marked as resolved.
|
Standalone Bolt routing is now isolated in draft follow-up #33994 above #33969, extending native stack #33967 without rewriting this PR or the coworker's #32935. The existing integrated route remains compatible and uses the same shared CUDA implementation. The complete fresh production source build passed native tests, logits/cache/state and route comparisons, plus 240 exact same-mode generation/reset/counter checks. The repeated 12k/24k/32k/124k before/after timings and limitations are in #33994; those timings use the frozen extraction build, not a new benchmark of the production refactor. Preserved quantized control is not independent HF, and inherited ordinary-versus-k7 differences remain. |
|
The native kernel optimizations are now pushed in #34002 and registered as the Actual #33994-tip baseline versus the new native candidate, real 1028-1663-token The earlier ~92 tok/s private baseline already contained these optimizations: |
There was a problem hiding this comment.
🟡 Changes recommended
The BF16 IQ2 ReLU-square tests leave channel permutations and FC2 cancellation-sensitive accumulation order untested, requiring asymmetric reference cases and larger mixed-sign fixtures for both decode and prefill.
2 open issues
What changed in this PR
Extends the existing CUDA MoEPacked operator for Nemotron Bolt’s ReLU-square experts and routing requirements while retaining default SwiGLU behavior.
Changes:
- Adds preselected FP32 routes and projected sigmoid routing with selection-only correction bias.
- Introduces a BF16 IQ2 kernel preserving contiguous-K accumulation and rounding boundaries.
- Adds routing and precision regression tests and updates operator documentation.
| File | Description |
|---|---|
onnxruntime/test/contrib_ops/moe_packed_routing_test.cc |
Adds routing, malformed-route isolation, compatibility, and cancellation tests. |
onnxruntime/core/graph/contrib_ops/contrib_defs.cc |
Extends the schema with activation options and optional routing inputs. |
onnxruntime/contrib_ops/cuda/moe/moe_packed.h |
Adds activation and routing parameters. |
onnxruntime/contrib_ops/cuda/moe/moe_packed.cu |
Implements new routing paths, ReLU-square activation, and BF16 IQ2 arithmetic. |
onnxruntime/contrib_ops/cuda/moe/moe_packed.cc |
Validates new input combinations and configures CUDA dispatch. |
docs/OperatorKernels.md |
Updates the supported kernel inputs and types. |
docs/ContribOperators.md |
Documents the expanded operator schema. |
Review effort: Balanced
7b2b13a to
e50e349
Compare
Tianlei Wu (tianleiwu)
left a comment
There was a problem hiding this comment.
Reviewed this layer (stack #33967, layer 2) at e50e34972b.
Implementation. I traced the validation, workspace/chunk accounting and each new kernel, and ran the PR-head kernel code in a standalone H200 harness: random IQ2_XXS payloads against an independent dequantized reference (decode and prefill tiles), sigmoid top-k including ties and NaN/Inf up to 512 experts, and the preselected-route validator; compute-sanitizer (memcheck, initcheck, racecheck) was clean. I did not find an implementation defect. The remaining risk is in tests that cannot fail and in undocumented runtime-only constraints.
Should be fixed in this layer
MalformedRoutesPoisonOnlyTheirTokendoes not verify NaN poisoning in BF16:TensorCheck<BFloat16>checksisnan(expected)instead of the actual value (inline on the test).- The "FP32 weights" fixtures give identical BF16 expectations if the weights are rounded to BF16 (inline), so they cannot catch narrowing.
- The two coverage requests from the earlier review are still open (asymmetric/independent-reference fixtures, and an FC2-specific cancellation fixture with
inter_size >= 1024); the checklist item in the PR description is unchecked. I added details to the first thread. - Please clarify the
compute_75overlay result. The PR description and the author review say a compute_75-only PTX overlay passedReluIq2PreservesContiguousKCancellation. With compute_75 PTX the MMA body is compiled out (MmaBf16is guarded by__CUDA_ARCH__ >= 800),GetMoEPackedBf16KernelSupportreports false (ptxVersion < 80), and dispatch falls back to the SIMT kernels. Running the PR-head kernels on an H200, this fixture yields 1024 on the Tensor-Core path but 256 on the SIMT path for rows 1/8/17/129 (the SIMT FC1 lane tree adds the unit term at k=128 to2^25at k=0 before the cancellation, so FC1 is 1 instead of 2). So the expected value is specific to the Tensor-Core dispatch, the same situation asRoundedBf16ContiguousKCancellationin #33962, andHasCudaEnvironment(800)does not guarantee that path. Which dispatch did the overlay exercise?
Other requests (inline): negative tests for the ~10 new validation rules, naming (relu_squared vs the existing relu2), undocumented runtime-only constraints, the defaults-unchanged test not reaching prefill tiles, and unreachable relu_squared branches in the generic Tensor-Core SwiGLU kernels.
Minor / notes
- No test executes the >2048-row chunk loop (second-chunk offsets for preselected
indices/weightsand sigmoidrouter_input, and workspace reuse across chunks). I read the offset math as correct; a 2049-row case per routing mode (k=2 keeps it small) would pin it. - FP16
relu_squarednow rounds decoded weights to FP16 (p.relu_squared || is_bf16), whereas FP16 SwiGLU keeps FP32-decoded weights. That is a reasonable choice but is not mentioned in the PR text or the docs. - Error texts: with the default softmax routing and
router_input/router_weight/router_biassupplied, the check reports "sigmoid routing requires ...", which points at the wrong fix; and the constructor text "internally routed mode requires k <= 16" is only true for softmax (sigmoid routing is also internal but allows k <= 32). The// 1. Softmax + TopKcomment inmoe_packed.cuis stale now that it also covers sigmoid and preselected routing. docs/PackedWeightFormats.md(not part of this diff) still says MoEPacked always has three weight tensors with independently chosen formats. Withrelu_squaredthere is nouptensor and gate/down must both be IQ2_XXS, and the weight-rounding contract of the new BF16 IQ2 path (weights rounded to the activation dtype, FP32 accumulation with contiguous-K partitions, BF16 rounding of FC1/ReLU-square/FC2 outputs) is undocumented, unlike the MatMulPacked rounding table.
e50e349 to
79067e3
Compare
79067e3 to
97c77c1
Compare
Preserve default SwiGLU and internal routing. Add IQ2_XXS ReLU-square gate/down experts, unchanged preselected FP32 weights and expert IDs through k=32, and safe affected-token poisoning for malformed routes. Exercise analytic outputs, routing boundaries and existing defaults. Production-flag translation units compile. Runtime numerical, generation and throughput qualification remain pending. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: rui-ren <ruiren1225@gmail.com>
Add an explicit projected sigmoid routing mode with FP32 selection-only correction bias and unbiased normalized routing weights. Preserve the existing softmax/SwiGLU defaults and preselected routing. Retain the source reduction order, tie-breaking and route scaling, with FP16 and BF16 regression fixtures. Validation: all five focused MoE CUDA tests pass on H100; all eight combined projection/MoE tests pass compute-sanitizer memcheck with zero skips or errors. Source hashes match the tested v3 build. This is operator-level compatibility-overlay evidence; full-model strict bitwise parity currently fails and speculative/TPS/release acceptance remain outstanding. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: rui-ren <ruiren1225@gmail.com>
Use source-compatible FC1 and FC2 K partitions, MMA ordering, and BF16 store boundaries for the IQ2 ReLU-square path. Reuse existing expert routing and finalization, retain other dispatch paths, and exercise cancellation plus 32-route prefill subtiles. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: rui-ren <ruiren1225@gmail.com>
Keep unit contributions outside large-magnitude tensor-core partials so the exact expected result does not depend on Ampere versus Hopper internal MMA accumulation. Preserve the extreme cancellation input, kernel arithmetic and zero tolerance. The updated current packed test overlay executes all ten projection and MoE regressions on H100 without skips. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: rui-ren <ruiren1225@gmail.com>
Use the canonical CPU/CUDA/DirectML/WebGPU CI artifact for the MoE layer so schema and kernel documentation match the actual registrations. Preserve artifact provenance outside the repository. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: rui-ren <ruiren1225@gmail.com>
Keep upstream compiled-BF16 fallback, bounded launches, and checked workspace accounting while retaining contiguous ReLU-square reductions. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: f4d0e1a3-a389-4cc8-9c85-0fd0e3ffce5b Signed-off-by: rui-ren <ruiren1225@gmail.com>
Retain the upstream 16-expert softmax limit while allowing the existing preselected and projected-sigmoid paths to select up to 32 experts. Apply the same range to checked workspace sizing. Existing k=32 routing regressions expose the inherited guard mismatch. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: f4d0e1a3-a389-4cc8-9c85-0fd0e3ffce5b Signed-off-by: rui-ren <ruiren1225@gmail.com>
Address all eight PR33963 coworker review threads. Use independent CPU IQ2_XXS dequantization and double-precision reference matmuls with asymmetric FC1/FC2 dimensions and nonuniform packed payloads, activations and expert routes. Bound BF16 deviations to two ULPs and at most 0.1% differing elements (one for tiny outputs), without relaxing existing exact comparisons. Verify actual poisoned BF16 outputs and unaffected tokens, use routing fixtures that demonstrably distinguish FP32 from BF16 weights, add 22 negative contract cases, and cover default prefill plus SwiGLU preselected k=2/32. Canonicalize relu2 while retaining relu_squared, reject attribute-only incompatibilities at construction, remove unreachable generic Tensor-Core relu branches while retaining SIMT support, and document routing/format constraints in the schema and generated operator docs. Review discussions: 4517915471, 4517913786, 4517913804, 4517913813, 4517913820, 4517913826, 4517913833, 4517913839. Validation: isolated host compilation of moe_packed.cc, moe_packed_routing_test.cc and contrib_defs.cc using preserved focused overlay recipes; configured CLANGFORMAT lint and git diff --check pass. CUDA execution and sanitizer remain for the parent's coordinated combined rebuild; no GPU or shared build artifacts were used. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: f4d0e1a3-a389-4cc8-9c85-0fd0e3ffce5b Signed-off-by: rui-ren <ruiren1225@gmail.com>
The all-invalid routing case poisons every token, so a verifier that assumes token 1 is finite rejects correct NaN output. Follow the expected NaN pattern across the entire tensor, explicitly checking actual NaNs and retaining exact, finite-value comparisons for valid outputs. No tolerance changes. Validation: complete MoE host test compilation, eight extracted CPU verifier acceptance/rejection cases, configured CLANGFORMAT and diff checks pass. GPU rerun remains with the parent integrated overlay. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: f4d0e1a3-a389-4cc8-9c85-0fd0e3ffce5b Signed-off-by: rui-ren <ruiren1225@gmail.com>
Consume preselected expert IDs and FP32 mixture coefficients from a separate router instead of publishing a second integrated sigmoid API. Remove the new router inputs, attributes, kernels and score workspace while retaining legacy softmax routing and the existing expert arithmetic. Update the schema references and reject the removed interface explicitly. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: c8741c5a-3e9e-4860-b846-7d6fb0025040 Signed-off-by: rui-ren <ruiren1225@gmail.com>
The existing aligned-load helper also serves the IQ2 expert kernels. Retain its exact implementation while deleting only the sigmoid projection and selection kernels, so the owning MoEPacked layer remains independent of the later shared-router implementation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: c8741c5a-3e9e-4860-b846-7d6fb0025040 Signed-off-by: rui-ren <ruiren1225@gmail.com>
Port the qualified common routing delta onto the exact #33963 head. Preserve this layer’s native expert kernels rather than importing upper-layer optimization implementations. Keep CPU routing same-typed, expose independent FP32 CUDA routing, correct INT3 type inference, and update operator documentation and focused regressions. The actual CUDA translation unit compiled successfully. Packed/QMoE host, CPU registration, schema and test syntax checks passed with matching PR headers; configured C++ lint passed. Lower-layer runtime execution and full-model qualification remain pending. No claim of a clean complete runtime build. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: rui-ren <ruiren1225@gmail.com>
Restore the QMoE CUDA implementation, CPU registrations, schema and generated operator documentation to their exact pre-a42960a state. Independent FP32 routing and deterministic full-score semantics remain limited to MoEPacked, avoiding unintended changes to the CPU and WebGPU QMoE contracts. Retain the shared routing helper for its MoEPacked caller. Correct the stale MoEPacked kernel-table signature and remove the cross-operator normalization claim. The owning-layer probe passed 19 native tests with zero skips, including two same-typed QMoE routing rejection cases, and 70 profiled CUDA MoEPacked executions. The probe recompiles the actual owning sources against retained native support, not a clean full-source runtime build. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: rui-ren <ruiren1225@gmail.com>
Rename moe_routing.cu and moe_routing.h to moe_full_score_routing.cu and moe_full_score_routing.h, and update the MoEPacked and implementation includes. Keep routing arithmetic, operator schemas and public interfaces unchanged. CUDA contrib source discovery already uses CONFIGURE_DEPENDS globs, so no build-list change is needed. Both affected CUDA translation units compile with the renamed header. Exact source comparisons confirm the delta is limited to filenames and include paths; scoped clang-format 20.1.8 and whitespace checks pass. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: rui-ren <ruiren1225@gmail.com>
Move the full-score ranking kernels and typed dispatch into the anonymous namespace in moe_packed.cu. Remove the separate helper translation unit, header and externally visible launcher overloads because MoEPacked is their only consumer. Preserve all routing arithmetic, launch configurations and operator interfaces; QMoE is unchanged. Rebuild the owning MoEPacked CUDA/factory module and run the unchanged native routing tests: 19 passed with zero skips and 70 profiled CUDA executions. Retained test/schema and support artifacts are hash-verified; this is not a clean full-source runtime or full-model qualification. Pinned clang-format 20.1.8 and whitespace checks pass. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: rui-ren <ruiren1225@gmail.com>
Rebase MoEPacked on current main while retaining the upstream hybrid-offload implementation and its larger workspace bound. Reject CPU-resident expert banks for full-score, FP32-logit and ReLU-square modes instead of letting the legacy SwiGLU offload path silently reinterpret their routing inputs. GPU-resident routing and QMoE remain unchanged. Freshly compile the owning MoEPacked host/CUDA sources, upstream CPU/offload support, current schema and test harness against retained support archives. All 20 native tests pass without skips, including three CPU-bank incompatibility cases and 70 profiled CUDA executions. This is not a clean full-source runtime or full offload/full-model qualification. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: rui-ren <ruiren1225@gmail.com>
Drop the full-score QMoE helper and four tests left from the abandoned QMoE routing changes. The unchanged QMoE schema does not accept their mixed-dtype routing inputs. Preserve every MoEPacked test and keep this PR scoped to MoEPacked. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: rui-ren <ruiren1225@gmail.com>
Keep omitted normalize_routing_weights equivalent to 1 in the schema and CUDA defaults, preserving legacy softmax routing behavior. Full-score tests explicitly opt out of normalization when testing raw mixing values. Add an omitted-attribute k=1 regression and render the restored default in operator documentation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: rui-ren <ruiren1225@gmail.com>
Remove the hand-written default annotation from normalize_routing_weights so the checked-in document matches the current schema renderer and passes generated-document validation. Keep the schema and kernel defaults at 1; no operator behavior or generator changes are needed. Both complete generated documents were verified against updated-docs from current-head CI run 43111471918. OperatorKernels.md is already identical. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: f4d0e1a3-a389-4cc8-9c85-0fd0e3ffce5b Signed-off-by: rui-ren <ruiren1225@gmail.com>
51b5492 to
0375ed9
Compare
## Description Adds `com.microsoft::Mamba2Verify` version 1 (CUDA), a snapshot-aware Mamba2 operator for Nemotron Bolt. Fuses causal convolution, selective recurrence, and gated grouped RMS normalization. A device INT32 selector chooses the starting snapshot; each processed token writes its post-token convolution and SSM snapshots so the caller can select the accepted/final state on the next invocation. The retained contract uses BF16 activations/convolution, FP32 recurrent state, readout and normalization coefficients, head_dim=64, d_state=128, and 1-16 token rows. It preserves state aliasing, trailing snapshots, eight-row specialization, and invalid-selector non-finite output behavior. Longer prefills require chunking. This layer adds no projection, routing, or expert operator and requires no custom-op plugin. **Dependency:** #33963 in [native stack #33967](https://github.com/microsoft/onnxruntime/stack/33967); the eventual landing target is `main`. The slot-aware single-row replay readout contract is included in this introducing layer. ## Summary of Changes ### Schema and Kernel Registration | File | Change | |------|--------| | `onnxruntime/core/graph/contrib_ops/fused_state_space_defs.h` | Define Mamba2Verify inputs/outputs, type constraints, and shape inference. | | `onnxruntime/core/graph/contrib_ops/contrib_defs.cc, onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc` | Register the contrib schema and CUDA kernel. | ### Implementation, Tests, and Documentation | File | Change | |------|--------| | `onnxruntime/contrib_ops/cuda/sequence/fused_state_space.cc` | Validate geometry/snapshots and dispatch the CUDA implementation. | | `onnxruntime/contrib_ops/cuda/sequence/fused_state_space_impl.{cu,h}` | Implement convolution, recurrence, readout, normalization, and snapshot writes. | | `onnxruntime/test/contrib_ops/fused_state_space_op_test.cc, onnxruntime/test/contrib_ops/cuda_kernels/fused_state_space_test.cc` | Exercise provider-level state/selector behavior and internal CUDA kernels. | | `docs/contrib_ops/cuda/mamba2_verify.md, docs/ContribOperators.md, docs/OperatorKernels.md` | Describe the operator ABI, snapshot behavior, and CUDA support. | ## Testing ### Optional-output update - `3e40c0d4380572c4edd2aaf47b32883437bba418` - `xbc` (output 3) and `readout` (output 4) are independently optional; omitted intermediates use stream-aware ORT scratch workspace. Required outputs and all existing output positions are unchanged. - The scope remains one Bolt sequence per invocation. The public SSM layout remains `[slots, heads, d_state, head_dim]`. CUDA math, dtypes/cast boundaries, timestep clamping, state selection and the readout/replay guarantee are unchanged. - **PASS:** 19 repository gtests (5 schema, 10 CUDA provider, 4 internal CUDA) plus one fresh-kernel/schema guard: **20 checks, zero failures/skips**, on H100. Profiles record 66 `CUDAExecutionProvider/Mamba2Verify` node executions. - **PASS:** all four intermediate-output combinations, including an empty output-3 placeholder; state aliases, lower/upper invalid selectors and trailing/no-trailing snapshots. Eight capture cases (four combinations at one/eight rows) verify actual graph capture and bitwise replay. - **PASS:** 5 unwrapped CPU-only schema tests, compiled without `USE_CUDA` and run with `CUDA_VISIBLE_DEVICES` empty; zero failures/skips. - **PASS:** configured changed-range clang-format 20.1.8, whitespace checks and scoped official schema/kernel documentation rendering. The kernel-table row remains unchanged; generated schema documentation shows 3-5 outputs. Qualification freshly compiles the operator host/device code, schema and tests against ABI-matched retained ORT 1.32.0 infrastructure. It is **not a clean full-source build, full-CI approval, independent HF parity or full-model acceptance**. The initial mixed-infrastructure-header harness failure was diagnosed in `OpTester::BuildModel` and preserved; no production workaround or threshold relaxation was introduced. CUDA math/internal-test sources are byte-unchanged. ### Historical validation before the optional-output update **Review-fix snapshot:** `4f2bd30fe1f3017213ce9eb620b34ab507f5e392`. Fixed static rank/dimension validation with CPU-only schema tests, dt clamp/large-dt and NaN behavior, strict BF16 ULP/fraction checking and nonuniform 1024-wide normalization. The slot-aware replay correction and `MambaSingleRowReplayPreservesMultiRowReadout` regression are now owned by this introducing PR (moved from #34003), with both host references updated and the single-slot path preserved. The permanent contract is now `docs/contrib_ops/cuda/mamba2_verify.md`. **Owning-layer validation:** 8 provider + 3 CPU schema + 4 actual internal CUDA gtests; zero failures/skips. **Integrated tip regression:** 51 C++ provider/schema tests, 5 actual internal CUDA tests and 2 Python IOBinding tests passed, zero failures/skips. All owning-layer provider cases are retained. The final overlay recompiles the changed CUDA/host/schema/test sources over ABI-matched preserved ORT support; **not a clean full-source build, full-CI approval, independent HF parity or new end-to-end qualification**. Targeted configured clang-format/Ruff and whitespace checks passed. <details> <summary>Historical source provenance moved out of permanent documentation</summary> Origin: #33940, commit `491dc043952066267b2212ba979f0325da318573`, frozen external source `export-nemotron-pr-20261030/src/bolt_ops.cu`, SHA256 `1f587d9eacc356805c37d19e94e7aa4465dd50302576b2ca4b9579d7c5ae8ba5`. This identifies the original retained implementation, not a dependency in the merged tree or a promise that later dispatch fixes are absent. The slot-aware replay correction and current contract are documented in this PR and the permanent operator documentation. This is the Bolt precision contract, not independent HuggingFace model parity. </details> <details> <summary>Historical measurements and qualification evidence (not current-head acceptance)</summary> These previously reported results are retained for provenance. They have not been rerun on the rebased head above; earlier CI/build status and benchmark results do not establish current-head CI, accuracy, or merge readiness. ### Measured CUDA kernel latency - 2026-11-07 BF16 projection/convolution, FP32 SSM state; 128 heads, eight groups, head dimension 64, state size 128, channels 10240, projection stride 18560, four-tap convolution. One fused R-row production launch is compared with R serial single-row production calls with exact state handoff; every timed invocation starts from the same immutable initial state. This compares three fused kernels with 3*R serial kernels. Frozen **Bolt** precision contract: FP32 gate/normalization/gamma followed by BF16 output; this is not an acceptance test for the earlier Lightning cast-before-gamma contract. Conv snapshots, SSM snapshots and xBC are bitwise equal for all three cases in both runs. Y/readout pass normalized RMSE <=0.03 and cosine >=0.999; exact raw metrics are below. The one-row control is approximately unchanged; eight/sixteen-row fusion reduces latency. | Rows | Serial single-row calls us | Fused multi-row call us | Ratio | |---|---|---|---| | 1 | 8.91 | 8.91 | 1.00x | | 8 | 82.98 | 22.17 | 3.74x | | 16 | 167.01 | 41.41 | 4.03x | H100 80GB HBM3, CUDA 12.8.93, driver 580.159.04, GCC 11.4.0. CUDA-event median per dispatch; 64 dispatches per captured CUDA graph, 50 warmups, three graph warmups, and 15 paired timed batches alternating AB/BA. Four rotating allocation banks contain deterministic synthetic weights (seed 617); this is not a cold-cache guarantee or a checkpoint-weight benchmark. Raw p95 values describe batch-average latency, not individual-invocation tails. No profiler runs in the timing window. The assigned GPU had empty admission, one stable compute client during execution, and clean release afterward; clocks were not locked. Two independent 39-case runs passed, including the published replay runner. **Scope:** same-build production CUDA dispatch alternatives, not historical parent-commit binaries. The measured runtime is an isolated CUDA provider relink with preserved core/support, **not a clean full-source ORT build or end-to-end model qualification**. Each measured operator's CUDA source hash matches its owning PR head. A clean source build is a separate replication variant and must be measured again. Kernel-only timings do not establish model TTFT/tok/s, MMLU/LCB, speculative integration, or full-model readiness. [Complete primary/repeat results, numerical metrics, source hashes and benchmark code](#33969 (comment)). All 39 cases passed in both supervised runs; full-model work remains paused. No source push, stack change, merge, auto-merge or ready-for-review promotion is included. </details> ## Checklist - [x] Review feedback addressed with focused regressions - [x] Documentation and supported contracts updated - [x] Native stack order and immediate-parent dependencies preserved - [ ] Clean full-source/full-CI qualification (not claimed by focused overlays) --------- Signed-off-by: rui-ren <ruiren1225@gmail.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: f4d0e1a3-a389-4cc8-9c85-0fd0e3ffce5b
## Description Keep the existing CPU/CUDA `com.microsoft::MoEGroupRouter` and its legacy dense output 0 unchanged. Optional dense projection weights remain unquantized FP32/FP16/BF16 tensors. The routing interface now exposes only the full arrays consumed by native MoE: | Output | Name | Type / shape | |---|---|---| | 0 | `router_probs` | Legacy activation-dtype log-scores, optional | | 1 | `selection_scores` | FP32 `[tokens, experts]`: sigmoid plus correction bias | | 2 | `router_weights` | FP32 `[tokens, experts]`: unbiased sigmoid | Outputs 1/2 must be requested together, with output 0 omitted. The unused compact `expert_indices` / `expert_weights` outputs, `Tind`, `router_scale`, and compact CPU/CUDA routing implementation are removed. Full-score outputs move from 3/4 to 1/2. Full-score mode requires one group, at most 512 experts and `rows <= INT_MAX`; projection requires rank-2 activations and a positive width divisible by 256. **Graph wiring:** `MoEGroupRouter([router_input, correction_bias, dense_weight], outputs=["", selection_scores, router_weights], top_k=k)`. Connect `selection_scores` to MoEPacked's `router_probs` input and the unbiased array to its `router_weights` input. MoEPacked performs top-k selection and selected-weight normalization; apply model routing scale explicitly in the composed graph. The chained regression retains an external `Mul(2.5)`. Legacy dense routing and full-score sigmoid/projection arithmetic are unchanged. This layer does not change QMoE or MoEPacked's schema/packed payload. Dependency: #33969 in [native stack #33967](https://github.com/microsoft/onnxruntime/stack/33967). ## Current focused qualification Owning head: `c57ef4627f28ab8bf3dce260ffc382a25c26e299`. Synchronized descendants: #34002 `99d57c485bdac723f1f2c646bf239fc4e35e3f70`, #34003 `ea71ed9e9070ce15b630eaa24dae3277f0a06cb7`. Their individual net changes are preserved exactly; all eight original stack members and the lower PR heads are retained. - 25 production router/shape/projection tests plus one native metadata/factory guard passed, zero failures or skips, at both the owning and synchronized top layer. FP32, FP16 and BF16 execute on the intended CPU/CUDA paths. - Each run profiles 56 CPU router executions, 113 CUDA router executions and four CUDA MoEPacked executions. The chained FP16/BF16 tests retain the unchanged numerical thresholds and external model-scale operation. - Changed native router/consumer CUDA sources, CPU kernels, schemas and tests were compiled against retained ORT support libraries. This is **not a clean full-source build or full-CI acceptance**. - `ContribOperators.md` and CPU/CUDA `OperatorKernels.md` router sections are regenerated and checked with the official renderers using fresh native schema/kernel metadata. Stripped text is restored from the executing schema source. Unrelated sections are preserved. Full `--gen_doc validate` is NOT_RUN. - Changed-range clang-format and whitespace checks pass. No full-model, speculative-integration or performance acceptance is claimed by this update. Local immutable qualification receipts, source snapshots, binary hashes and provider profiles: `/root/.copilot/session-state/3b29a5de-77c7-4cf5-b77a-02adf4f73c89/files/pr33994-remove-compact-routing`. <details> <summary>Historical PR description and earlier-head evidence (not current-head acceptance)</summary> ### Description Extends the existing CPU and CUDA `com.microsoft::MoEGroupRouter` with optional dense router projection and two ways to expose routing results. The existing grouped dense-logit path and `router_probs` output remain available. No separate router operator is introduced, and this PR does not change QMoE. With `projection_weight`, the operator computes router logits from rank-2 activations and ordinary dense, unquantized weights. It can expose either: - **Compact routing:** paired `expert_indices` and `expert_weights` outputs. Selection uses sigmoid scores plus `expert_bias`; the returned weights use the unbiased sigmoid scores, with optional normalization and `router_scale`. This mode is currently restricted to `n_group=topk_group=1`, at most 512 experts, and `top_k <= 32`. - **Full-score routing:** paired FP32 `selection_scores` and `router_weights` outputs. `router_weights` contains unbiased sigmoid values; `selection_scores` contains those values plus `expert_bias`. This mode does not select top-k, normalize, scale, mask, or take logarithms. A downstream operator performs selection and weighting; the full-score outputs compose with `MoEPacked` from #33963. The compact outputs remain exposed by this PR, but currently have no in-tree consumer. The `MoEPacked` composition uses the full-score outputs. Tests cover CPU and CUDA routing, projection and shape inference, tie-breaking, selection bias versus mixing weights, boundary and invalid-input cases, and a chained full-score `MoEPacked` reference comparison. ### Motivation and Context Keep router computation reusable and separate from expert execution. The optional projection lets `MoEGroupRouter` own the router arithmetic, while its full-score interface allows `MoEPacked` to select experts and execute packed experts without adding a learned router projection to `MoEPacked`. This PR preserves the existing `MoEGroupRouter` schema and dense-logit behavior while adding optional outputs and projection support. It does not modify QMoE. </details> --------- Signed-off-by: rui-ren <ruiren1225@gmail.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: c8741c5a-3e9e-4860-b846-7d6fb0025040 Copilot-Session: f4d0e1a3-a389-4cc8-9c85-0fd0e3ffce5b
## Description Reduces CUDA packed decode tile/loader overhead above #34002. Narrows eligible eight-row IQ2 FC1 tiles, adds aligned record-window reconstruction for eligible IQ2 FC2 rows, and uses full-tile NVFP4 paths only when their shape/alignment guards hold. Unsupported geometries retain the existing fallback. **Dependency:** #34002 in [native stack #33967](https://github.com/microsoft/onnxruntime/stack/33967). No schema, packed layout, cache/mask precision workaround, or GenAI source change. The R8-prefill/R1-ordinary hybrid session/package work is in the companion microsoft/onnxruntime-genai#2916, **not implemented by this native-kernel PR**. ## Summary of Changes ### Packed Decode Specializations | File | Change | |------|--------| | `onnxruntime/contrib_ops/cuda/moe/moe_packed.cu` | Select narrower FC1 tiles and eligible aligned FC2 loader variants. | | `onnxruntime/contrib_ops/cuda/quantization/gguf_quant_cuda.cuh`, `onnxruntime/contrib_ops/cuda/moe/moe_packed_iq2.cuh` | Reuse the shared IQ2 decoder for aligned word windows, including two-byte-offset records. | | `onnxruntime/contrib_ops/cuda/quantization/matmul_packed.cu` | Specialize eligible full NVFP4 output tiles while preserving guarded fallbacks. | | `onnxruntime/test/contrib_ops/moe_packed_routing_test.cc` | Exercise eight-row width/routing boundaries and larger intermediate widths. | ## Testing **Review-fix snapshot:** `ed6e6e1991fd944f65f1d97472d3f29046d97361`. Bound aligned-window FC2 to four resident CTAs, kept the R8-specific narrow FC1 gate decode-only (`!use_prefill_tiles`), simplified the FC2 guard to its real row/base alignment invariant, and reused the lower shared decoder. Nonuniform regression protection is introduced in #33963/#34002. Mamba replay changes were moved into their owning PR #33969; this layer no longer changes Mamba numerics or tests. **Owning-layer validation:** Integrated stack regression run and exact-production-kernel H100 A/B below. **Integrated tip regression:** 51 C++ provider/schema tests, 5 actual internal CUDA tests and 2 Python IOBinding tests passed, zero failures/skips. All owning-layer provider cases are retained. The final overlay recompiles the changed CUDA/host/schema/test sources over ABI-matched preserved ORT support; **not a clean full-source build, full-CI approval, independent HF parity or new end-to-end qualification**. Targeted configured clang-format/Ruff and whitespace checks passed. ### Review-requested kernel A/B (H100, CUDA 12.8.93) Exact production kernel/helper text from the integrated head; random nonuniform IQ2, 512 experts, 16 rotating route sets, R=8. All variant outputs are bit-identical across all route sets. CUDA-event median: 20 warmups, 9 alternating batches x32 launches; no full-model/cold-cache guarantee. `k=9` is decode, `k=17` is four-subtile prefill; the latter remains excluded from the narrow dispatch. | Hidden | Inter | k | FC1 wide→narrow μs | FC2 plain / window / hinted μs | |---:|---:|---:|---:|---:| | 1024 | 1024 | 9 | 15.035 → 10.888 | 12.874 / 11.966 / 12.062 | | 1024 | 1024 | 17 | 16.264 → 14.610 | 18.149 / 15.968 / 16.188 | | 1024 | 2048 | 9 | 18.041 → 15.920 | 21.786 / 20.282 / 20.458 | | 1024 | 2048 | 17 | 26.279 → 23.879 | 32.182 / 28.613 / 28.801 | | 2048 | 1024 | 9 | 27.187 → 18.500 | 20.663 / 19.135 / 19.234 | | 2048 | 1024 | 17 | 28.614 → 25.400 | 29.040 / 25.962 / 26.176 | | 2048 | 2048 | 9 | 31.466 → 27.515 | 34.251 / 32.346 / 32.446 | | 2048 | 2048 | 17 | 48.462 → 42.877 | 51.977 / 45.898 / 46.270 | | 4096 | 1024 | 9 | 49.654 → 32.181 | 30.966 / 29.015 / 29.088 | | 4096 | 1024 | 17 | 53.387 → 46.004 | 50.201 / 45.097 / 45.371 | | 4096 | 2048 | 9 | 57.550 → 50.024 | 53.390 / 50.104 / 50.318 | | 4096 | 2048 | 17 | 91.847 → 79.396 | 90.444 / 80.277 / 80.838 | The local 12.8 compiler **did not reproduce** the reported 79-register/3-CTA window allocation: unhinted and hinted window kernels both use 64 registers/4 CTAs here. The selective hint is retained as a portability cap for the reviewer's newer-toolchain result, with its local timing tradeoff shown rather than hidden. These kernel-only measurements do not attribute the older hybrid end-to-end gains below to the later narrow-tile commit. <details> <summary>Historical measurements and qualification evidence (not current-head acceptance)</summary> These previously reported results are retained for provenance. They have not been rerun on the rebased head above; earlier CI/build status and benchmark results do not establish current-head CI, accuracy, or merge readiness. The following results include the companion GenAI runtime and hybrid package, not just this PR's native kernel changes. They use the original pinned binaries and package hashes below, not the current rebased ORT head. **Unprofiled frozen real30, six H100s, two reversed same-GPU passes, 240 complete measured generations; original warmups, mode order, EOS/512 cap.** | Actual input tokens | Before ordinary tok/s | Hybrid ordinary tok/s | Gain | Original R8/k7 tok/s | |---|---:|---:|---:|---:| | 178-254 | 94.40 | 119.95 | +27.1% | 361.46 → 361.57 | | 427-932 | 94.01 | 119.59 | +27.2% | 298.68 → 298.94 | | 1028-1663 | 93.89 | 118.90 | +26.6% | 324.25 → 324.31 | Ordinary TTFT medians: short 357.35→359.00ms, medium 826.97→831.29ms, long 2153.15→2161.92ms. The earlier static-R1 version's prefill regression is **not** the package being qualified here. Separate same-GPU synthetic-2048 CUDA profiles, 64 captured target passes: IQ2 expert work **3.259→1.496ms/pass**, Mamba **0.824→0.423ms/pass**. These are device-work attribution, not an additional throughput claim. **Correctness:** all240 exact same-mode sequences and complete counters match; 388 complete live committed-state boundaries (66.43GB compared), capture off/on, two resets, prefix through130; original R8 ordinary/k7 retains1782 complete finite tensor hashes. Three native kernel tests, seven Mamba operator tests, fifteen GenAI tests execute without skips. The new replay test also passes CUDA memcheck with zero errors (100 repetitions). **Scope:** hybrid is ordinary-only. Both speculative arms use the original R8 package. The benchmark explicitly keeps two per-mode models resident in each arm; it is not a unified hybrid-speculative package or a k7 speedup. No precision/cache/mask workaround, no weight repack, no independent HF rerun, no clean-build/full-CI claim. Reproduce the exact tested graph/config with: ```bash python examples/python/nemotron-bolt-hybrid.py --input <packed-r8> --output <fresh-same-filesystem-hybrid> ``` Requires onnx-ir/numpy/ml-dtypes and both updated runtimes. External-data validation stays enabled; payloads are hardlinked inside the output package, not copied. Pins: - ORT: `0357a7dfe47d3c383bf36be62015162b6a612d5e` - Compiled/common GenAI: `d73a6e627d6d20a5a1688b9fbd46237069a7be89` - Published GenAI: `db6ffd29fc96e319f645ff461efa05c643235f29` (converter/docs-only child) - CUDA provider: `fac33c75d0336236736aa96456268057b628d93d0b9901d49598e275bf335275` - GenAI shared library: `f977660a408bc0a01a31151c7280282ba443fa256785680ae3e841654eec9be7` - Tested hybrid graph: `aaf5d6ae06b6740ab5b3af523359581930d2b6c42c38a18529a22737c2cfe530` - Config: `77f2d920d2b04e17884c3d2ee691d9053225b957c047454ebd0a4ea3649dac44` - Unchanged packed payload: `06d3d508e8d11235d0b770484b6b59764c2f297c3383b6d0de75648d26c11867` Full commands, per-case exact tokens/counters, raw kernel traces, provider maps, UUID/admission receipts, and rejection evidence remain on the validation host under `/opt/bolt-runtime-fixes-20261108-v1/evidence/ordinary-r1-v1`. Sealed manifest SHA256: `412b3898f09565557b01e82906358f017f9df3227db2bf330850245a52d32bbc`. Root-disk binaries were reloaded and verified with full ordinary/k7 generation; the large model payload still resides on the existing tmpfs, not reboot-persistent. Original package/deployment untouched; no merge. </details> ## Checklist - [x] Review feedback addressed with focused regressions - [x] Documentation and supported contracts updated - [x] Native stack order and immediate-parent dependencies preserved - [ ] Clean full-source/full-CI qualification (not claimed by focused overlays) --------- Signed-off-by: rui-ren <ruiren1225@gmail.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: f4d0e1a3-a389-4cc8-9c85-0fd0e3ffce5b
This cherry-picks the following commits for the release: | Commit ID | PR Number | Commit Title | |-----------|-----------|-------------| | 85dcb54 | #32475 | Update Eigen to 5.0.1 + SSE2 fix (broken Clip & Sign) | | 1ea552d | #33976 | [CUDA/CPU] Add hybrid GPU/CPU expert offload for MoEPacked block-packed weights | | 3a8b8f2 | #34066 | [CUDA] Optimize multi-token MatMulPacked GGUF GEMV (M=2..8) and eliminate ReadU32Dev warp divergence | | 64506ea | #34098 | [CUDA] Guard tests requiring SM80 on older GPUs | | 74f861c | #34067 | [MoE] Optimize multi-token (M=4, 8) MoEPacked CPU hybrid cold routes and GPU FC2 decode | | 6aa9b27 | #33963 | [CUDA] Extend MoEPacked for top-k, dispatch and fused MoE calculation | | 8fb1b27 | #33969 | [CUDA] Add Mamba2Verify CUDA Kernel | | cedae42 | #33994 | [CUDA] Update `MoEGroupRouter` for routing to `MoEPacked` | | 158d0c0 | #34002 | [CUDA] Speed up native Bolt packed projections and IQ2 experts | | 38c6dce | #34003 | [CUDA] Reduce native packed decode tile and FC2 loader overhead | | d5704db | #34147 | [CUDA] Optimize MoEPacked, MoEGroupRouter, and MatMulPacked for Nemotron Bolt | | 7cca0e5 | #34126 | Fix Windows junction detection in model_package for non-MSVC standard libraries | | 2bf400a | #34099 | Fix external data validation for Hugging Face Hub shared blob store | | d2f77cd | #34164 | [CUDA] Accelerate Nemotron Bolt large-chunk prefill across contrib ops | | ca7ea48 | #28457 | Add set_default_logger_callback Python API for custom ORT logging | | fca9447 | #34117 | [CUDA] Optimize CUDA prefill and quantized paged attention | | 1804c6f | #34114| [CUDA] Optimize MoEPacked GGUF offload 3-stage prefill pipeline and F16C/Batch4/40-core cold decode | --------- Signed-off-by: Christian Bourjau <christian.bourjau@quantco.com> Signed-off-by: rui-ren <ruiren1225@gmail.com> Co-authored-by: Christian Bourjau <cbourjau@users.noreply.github.com> Co-authored-by: BBaditi <bbaditi@microsoft.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Rui Ren <ruiren1225@gmail.com> Co-authored-by: Mustapha Jaber <mustjab@users.noreply.github.com> Co-authored-by: Mustapha Jaber <mustjab+microsoft@microsoft.com> Co-authored-by: Copilot <198982749+Copilot@users.noreply.github.com> Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com> Copilot-Session: f4d0e1a3-a389-4cc8-9c85-0fd0e3ffce5b Copilot-Session: c8741c5a-3e9e-4860-b846-7d6fb0025040

Description
Extends CUDA
com.microsoft::MoEPackedwith ReLU-square experts and full-score routing. Routing scores and mixing weights can be FP32 independently of the FP16/BF16 expert activations. This PR does not change QMoE.When
router_weightsis provided,router_probscontains full selection scores androuter_weightscontains separate mixing values.MoEPackedselects the top-k experts from the scores, gathers their mixing values, optionally normalizes them, and combines the expert outputs. Withoutrouter_weights, the existing softmax-logit routing path is used.Routing tensors share a dtype: FP32 or the activation dtype. Full-score routing supports up to 32 experts per token; the legacy softmax path supports up to 16. Packed expert-weight layouts are unchanged.
The
relu2expert activation (relu_squaredis an alias) requires IQ2_XXS gate and down weights, with up weights omitted. Learned router projection remains outsideMoEPacked; this operator has no integrated router projection orrouter_scaleattribute.Motivation and Context
This allows a separate router to provide full-precision selection scores and mixing values to
MoEPackedwithout first narrowing them to the expert activation dtype. It also enables the supported IQ2_XXS ReLU-square expert path while keeping router computation and packed expert execution separate.Normalization-default compatibility correction
Commit
18da59fe12restoresnormalize_routing_weights=1in the schema and CUDA defaults, preserving legacy models that omit the attribute. Full-score callers wanting unnormalized mixing values must explicitly set0. The full-score tests now do so; a new omitted-attribute,k=1legacy regression protects the default. The generated operator documentation records default1.0.28125, actual0.140625. The explicit full-score opt-out still passed.CUDAExecutionProvider/MoEPackedexecutions. Changed-range clang-format 20.1.8, whitespace, and scoped official documentation-rendering checks passed.Both positive runs freshly rebuilt the operator, schema and tests against ABI-matched retained ORT support. This is not a clean full-source build, full CI approval, full hybrid-offload qualification, or full-model acceptance. Full
--gen_doc validatewas not run.Historical validation before the normalization-default correction
MoEPackedRoutingTestcases and three additional scope guards passed without a test filter: 20 tests, zero failures or skips.MoEPackedexecutions ran on CUDA.Head
87daac143fremoves the obsolete QMoE helper and four QMoE tests, preserving every MoEPacked test and leaving runtime code unchanged.The branch remains rebased onto main
0b22d27a85, preserving the hybrid-offload implementation from #33973. The new activation/routing modes remain GPU-resident; unsupported CPU-weight-bank combinations are rejected explicitly rather than entering the legacy hybrid-offload route.Qualification freshly rebuilt the current schema/test executable against the unchanged, previously qualified MoEPacked module and retained native support libraries. This is not a clean full-source runtime build, full hybrid-offload qualification, full CI approval, or independent full-model parity. Full
--gen_doc validatewas not run.