[CUDA] Add fused Mamba verification and routing contrib operators by rui-ren · Pull Request #33898 · microsoft/onnxruntime

Rui Ren (rui-ren) · 2026-11-04T22:33:48Z

Summary

This PR adds fused CUDA contrib operators for Mamba2 verification, bias-corrected expert routing and squared ReLU. It migrates the existing Bolt fused computations into ONNX Runtime while preserving recurrent-state snapshots, routing semantics and source arithmetic:

Add fused Mamba verification (fused_state_space.cc, fused_state_space_impl.cu, fused_state_space_impl.h, fused_state_space_defs.h):
Register com.microsoft::Mamba2Verify with typed inputs/outputs, shape validation and CUDA dispatch. Preserve convolution, recurrence, readout and gated-normalization behavior for 1-16 rows, including the specialized eight-row verification path. Longer prompts continue to use caller-managed chunking; this is not a new arbitrary-length prefill implementation.
Preserve accepted-prefix state semantics (fused_state_space_impl.cu):
Read the caller-selected input snapshot and emit each row's post-token convolution/SSM state. Support corresponding input/output state aliases and preserve trailing snapshot slots. Keep accepted-prefix selection and commit/rollback policy with the caller rather than hiding them inside the operator.
Add deterministic bias-corrected routing and squared ReLU (fused_state_space.cc, fused_state_space_impl.cu):
Register BiasCorrectedMoERoute and ReluSquared. Use corrected scores for expert selection and unbiased sigmoid scores for routing weights. Retain stable lower-expert-ID tie breaking; sort NaN corrected scores after numeric scores, including negative infinity.
Harden device-boundary behavior without host round trips (fused_state_space_impl.cu, fused_state_space_defs.h):
Define non-finite output behavior for invalid Mamba selectors without dereferencing invalid state. Handle element-aligned, non-vector-aligned convolution/router bindings with scalar loads that preserve accumulation order. Cover large normalization group counts and make the ReluSquared launch-grid conversion explicit for Windows warning-as-error builds.
Integrate and document the operator contracts (cuda_contrib_kernels.cc, contrib_defs.cc, FusedVerification.md, generated operator docs):
Use the supplied ORT CUDA stream and capture-compatible device operations. Preserve the fused implementation rather than introducing host-side routing, selector inspection, per-call CUDA allocation or synchronization.
Optimize the router top-k hot loop (fused_state_space_impl.cu):
Make finite-score selection a warp-uniform fast path; precompute exact integer ranks for exceptional scores. Preserve NaN ordering, stable expert-ID ties, eligibility and FP32 routing-weight arithmetic instead of repeating generic comparisons for every selected expert.

Testing

Focused provider coverage added: FusedStateSpaceTest.* covers post-token recurrence/readout, nonzero snapshot selection, trailing slots, 1/3/8/16-row paths, grouped normalization, routing weights/ties, squared ReLU, invalid selectors and device-state aliasing.
H100 kernel regressions: PASS. Four focused state/routing tests passed on the reviewed combined stack: MambaNorm65536Groups, MambaOffsetConvolutionWeights, RouterOffsetBindingsAndIndependentReference, and RouterNaNsSortLastWithStableExpertTies.
Compilation checks: PASS. Changed host/CUDA units were compiled with cached dependencies; the review validation includes SM80/SM90a and minimal/contrib-disabled test configurations. The Windows grid-conversion issue was reproduced before its explicit-cast correction and compiled afterward without relaxing warning policy.
Combined-port evidence: the earlier pinned overlay passed 24 direct CUDA-provider cases and 24 capture/replay cases across all seven operators in the ORT stack. These are shared totals, not results from 24 tests of this PR alone.
Qualification: the targeted incremental production-provider rebuild and same-mode full-model replay below pass. They do not establish a clean from-scratch full ORT build, complete provider-suite acceptance, independent HF parity, lossless speculation, MMLU/LiveCodeBench accuracy, or release readiness. Other operator-family performance regressions remain.

Router top-k correction: production-provider qualification

Recompile the changed CUDA translation unit with the recorded production flags and relink the complete CUDA provider against unchanged cached objects. Run through normal ORT/GenAI dispatch; no custom Bolt libraries or kernel-launch replacement are loaded. The frozen runtime and model packages are unchanged.

Native router regressions: PASS, 3/3, zero skips, linked against the production-compiled object. Include mixed numeric/exceptional-score ordering.
Exact historical 30-input replay: PASS. 120 fresh measured requests (30 questions x ordinary/k=7 x baseline/optimized), with two warmups per mode per worker. All 60 paired complete output sequences are identical; all 120 match the archived same-mode contrib reference. Natural EOS and the 512-token cap are retained. The fixtures are 19 LCB + 11 MMLU questions, not a new accuracy evaluation.
Controlled H100 measurements: the two variants use the same physical GPU for each question across two owned H100s; order is reversed across GPUs. This is a new paired assignment, not the archived eight-GPU assignment. Decode uses pooled (generated_tokens - 1) / decode_seconds.

Input tokens	Mode	Before tok/s	Optimized tok/s	Improvement
178-254	ordinary	101.75	103.65	1.87%
178-254	k=7	388.20	397.40	2.37%
427-932	ordinary	101.49	103.36	1.84%
427-932	k=7	318.90	326.68	2.44%
1,028-1,663	ordinary	101.35	103.20	1.83%
1,028-1,663	k=7	344.78	354.23	2.74%

Earlier controlled router profiling measured 12.24 -> 7.67 us (37.4% lower latency); isolated ablations and 4,204,800 live routing entries preserved ordering and weights with zero mismatches. These are router-only improvements, not recovery of the remaining projection/MoE regressions or a release-acceptance claim. Independent HF parity and ordinary-versus-speculative losslessness remain unestablished.

Stack: #33897 -> #33898 -> #33899. This PR targets the pinned 01-projections predecessor mirror. Its recorded mirror is older than #33897's latest follow-ups; integrate/synchronize predecessor updates before landing. Merge #33897 first, retarget this PR to main, and require main-branch CI acceptance. The pre-existing ordinary-versus-k7 discrepancy remains separate.

Port Mamba2Verify, BiasCorrectedMoERoute and ReluSquared with their existing shared implementation unit, schema contracts and tests. Preserve recurrent snapshots and arithmetic; commit/rollback policy remains with the caller. Signed-off-by: rui-ren <ruiren1225@gmail.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Windows CUDA CI promotes MSVC C4244 to an error when the int64 element count expression implicitly becomes dim3's unsigned grid size. Express the existing conversion explicitly without changing launch geometry, kernel arithmetic or source precision. Reproduced the narrowing before the fix and compiled the corrected unit for SM80 and SM90a with host conversion warnings treated as errors. Signed-off-by: rui-ren <ruiren1225@gmail.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot

🟡 Changes recommended

Unchecked state selectors and overflowing output offsets risk out-of-bounds GPU access, while missing generated operator documentation would fail CI validation.

7 open issues

What changed in this PR

Adds Mamba2Verify, BiasCorrectedMoERoute, and ReluSquared as built-in CUDA contrib operators, preserving fused arithmetic and state snapshots. Accepted-prefix selection and commit/rollback remain caller-controlled.

Changes:

Add operator schemas, runtime validation, and CUDA registrations.
Implement fused recurrence, grouped normalization, expert routing, and squared ReLU.
Add focused operator tests and contract documentation.

Packed-projection files below are inherited from part 1; this review focuses on the requested part-2 increment.

File	Description
`onnxruntime/test/contrib_ops/matmul_packed_projection_test.cc`	Inherited packed-projection regression tests.
`onnxruntime/test/contrib_ops/fused_state_space_op_test.cc`	Tests snapshots, routing, and squared ReLU.
`onnxruntime/core/graph/contrib_ops/ms_opset.h`	Inherited packed-projection schema entries.
`onnxruntime/core/graph/contrib_ops/fused_state_space_defs.h`	Defines the three operator schemas and shape inference.
`onnxruntime/core/graph/contrib_ops/contrib_defs.cc`	Registers fused schemas alongside inherited projection schemas.
`onnxruntime/contrib_ops/cuda/sequence/FusedVerification.md`	Documents contracts, provenance, and integration.
`onnxruntime/contrib_ops/cuda/sequence/fused_state_space.cc`	Adds CUDA adapters, validation, and snapshot-buffer handling.
`onnxruntime/contrib_ops/cuda/sequence/fused_state_space_impl.h`	Declares CUDA launch interfaces.
`onnxruntime/contrib_ops/cuda/sequence/fused_state_space_impl.cu`	Implements recurrence, normalization, routing, and activation kernels.
`onnxruntime/contrib_ops/cuda/quantization/packed_projection_impl.cuh`	Inherited packed decoding and matrix-multiply helpers.
`onnxruntime/contrib_ops/cuda/quantization/matmul_packed_projection.h`	Inherited projection kernel classes.
`onnxruntime/contrib_ops/cuda/quantization/matmul_packed_projection.cuh`	Inherited projection launch interface.
`onnxruntime/contrib_ops/cuda/quantization/matmul_packed_projection.cu`	Inherited projection kernels and dispatch.
`onnxruntime/contrib_ops/cuda/quantization/matmul_packed_projection.cc`	Inherited projection validation and registration.
`onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc`	Registers the three new CUDA kernels.
`docs/contrib_ops/cuda/packed_projection.md`	Inherited packed-layout and arithmetic documentation.
`cmake/onnxruntime_providers_cuda.cmake`	Inherited projection-specific CUDA arithmetic options.

_{Review effort: Max}

Address the review findings while retaining the source layouts and valid-input arithmetic. Add focused boundary and correctness regressions and preserve each stack installment's existing scope. Validation: changed-source host checks, CUDA SM80/SM90a compilation, and focused runtime regressions are recorded separately from clean-build and release qualification. Signed-off-by: rui-ren <ruiren1225@gmail.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

The all-files lintrunner job also checks C++ formatting. Adding the standalone preprocessor branch changes automatic language-standard detection for the existing RunCuda nested-template spelling. Apply its one-line formatter output; no test semantics or kernel arithmetic change. Signed-off-by: rui-ren <ruiren1225@gmail.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Rui Ren (rui-ren) · 2026-11-05T00:55:37Z

Review follow-up at 821706dcb9a0bfc1750b6e12c8184014a515ca5f: fresh GitHub readback shows all 17 current review threads resolved, with fixes published and individual evidence replies for the new findings.

Carried the projection fixes and added the contrib-disabled guard to the fused-state internal tests. Normal/minimal/contrib-disabled compilations pass; the excluded object has no production-launcher references. Five predecessor H100 kernel regressions pass without skips.

Whole-file clang-format-20 checks pass. These are exact changed-source builds and standalone production-kernel executions with cached unchanged dependencies, not a clean full ORT/internal-target build, linked provider-suite run or full-model qualification. Commits are signed off and additive; predecessor mirrors are synchronized. CI is not asserted green, and historical benchmark/ordinary-vs-k7 qualification remains separate.

Copilot

🟡 Changes recommended

Valid inputs can cause out-of-bounds or misaligned CUDA reads and incorrect router indices with NaNs, while regression coverage is insufficient to reliably catch these failures.

9 open issues

7 resolved since last review

_{Review effort: Max}

Promote convolution-weight offsets before multiplication and retain aligned vector loads with element-aligned scalar fallbacks. Use a single deterministic router comparator that places NaNs last. Preserve only trailing nonaliased state slots. Move the normalization-grid regression into the existing CUDA internal-test target. Add offset binding and independent K=768/E=511/512/topk=32 routing regressions, CUDA-minimal exclusions, and exact convolution-history tolerances. Validation: five host checks, exact production CUDA SM80/SM90a compilation, and four H100 internal-kernel regressions pass. Full linked ORT provider suite and full-model release qualification remain separate. Signed-off-by: rui-ren <ruiren1225@gmail.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Exclude fused-state CUDA internal regressions when DISABLE_CONTRIB_OPS is set, matching the absence of their production launchers. Compile the excluded object and verify that it has no launcher references. Retain CUDA-minimal exclusions. Carry the predecessor's SM52 guard and cancellation regressions forward without rewriting stack history. Five exact-source H100 predecessor-kernel regressions pass; full linked ORT provider tests and a clean internal-target build remain separate qualification. Signed-off-by: rui-ren <ruiren1225@gmail.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Use a warp-uniform finite-score fast path and exact precomputed ranks for exceptional scores. Preserve numeric-before-NaN ordering, stable expert-ID ties, removed-slot eligibility and the original FP32 weight arithmetic without repeating generic checks for every selection. Add independent-reference mixed-score coverage. H100 isolated tests and live routing comparisons pass; production-provider qualification is recorded separately before publication. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: rui-ren <ruiren1225@gmail.com>

Rui Ren (rui-ren) · 2026-11-09T05:55:18Z

Follow-up qualification is recorded in
#33940's common corrected-GenAI results.

With one identical corrected GenAI runtime in both arms, the current private
deployment versus the equivalent #33940 public route measured, for real
1028-1663-token inputs:

Ordinary: 92.02 -> 109.16 tok/s.
k=7: 318.91 -> 369.61 tok/s.

Both reversed-order passes completed 240 measured generations; full sequences
and within-mode counters matched, and packed-value/full-state checks passed.
This is not a fresh benchmark of the latest three predecessor heads together,
nor an isolated performance gain attributable to this individual PR.

#33940 remains an alternative/consolidation, not another PR to stack above
#33897 -> #33898 -> #33899. This comment does not change this PR's source or its
original qualification. The common GenAI correctness source is published
separately as microsoft/onnxruntime-genai#2916;
the measured ORT gain is not attributed to R16 or a change in accepted drafts.

…xes (#2921) ## Summary Complete ordinary target-only generation for native Nemotron Bolt while preserving the existing fixed-eight-row target graph and greedy DFlash2 speculative path. Also fix configuration handling, shared-initializer allocation metadata, and generation statistics. The supported runtime remains batch-one, single-CUDA-device greedy generation through the Model/Generator API, not paged Engine batching. ## Stack and dependencies | Layer | PR | Scope | Base | |---|---|---|---| | 1/4 | #2920 | DFlash2 padding safety and API compatibility | `main` | | 2/4 | **#2921** | Target-only generation, configuration, statistics, and shared weights | #2920 | | 3/4 | #2922 | Native drafter policy for partial-acceptance graph replay | #2921 | | 4/4 | #2923 | Verification arithmetic and opt-in one-token ordinary decode | #2922 | Merge in this order. After a parent lands, retarget its dependent PR to `main` and rerun required CI. No merge is requested here. **Current head:** `2e4886ef0046415dbbf583ee5822f087ff3c294c`. This is the in-repository replacement for #2868. The original PR retains earlier review/qualification history; the existing stack was not rebased or rewritten. ## Changes in this layer - **Target-only decoding:** accept `speculative.max_draft_tokens=0` or `7`. A model configured with `0` skips drafter loading and drafter-only state allocation. Selecting `0` on a model configured with a drafter keeps its weights resident; these are different memory-measurement configurations. - **Capture controls:** allow ordinary decoding without target graph capture. A capture request still requires a capture-enabled target session; speculative decoding requires the compatible drafter and capture-enabled sessions. - **Configuration and public API:** reject unsupported package draft counts before loading sessions, parse target-only configuration independently of JSON member order, and allow the public speculative setter's zero-draft mode for Bolt without changing unrelated model behavior. - **Optimizer settings:** add `GatherSliceToSplitFusion` to the drafter's disabled optimizers without discarding caller exclusions. - **Shared initializers:** use allocator memory metadata instead of the CPU device-interface method that throws. Keep allocation identity device-aware and test that shared buffers survive either session's destruction, then expire after both sessions are gone. - **State and statistics:** handle an absent drafter during prefill/reset, report effective draft count zero in ordinary mode, and make zero/partial/full acceptance buckets mutually exclusive. Cover EOS, output budgets, correction/bonus tokens, and public-token commit/reset behavior. - **Runtime documentation:** clarify target-only versus speculative execution, custom-op versus contrib-op packages, and the limits of migration comparisons. ## Review fix included `2e4886ef` declares the pinned `nlohmann_json` v3.12.0 dependency and links `nlohmann_json::nlohmann_json` privately to `engine_unit_tests`. The unconditional JSON include no longer relies on a system-installed header. The dependency fix is merged forward through #2922 and #2923 without rewriting history. The corresponding review thread has been answered and resolved. ## Validation and limits - Native test configuration was attempted during the review fix, but this Windows host has no `nmake` or C/C++ compiler. **No fresh local native/GPU build is claimed for the review fix.** - CI is **not yet green** at this description update; build failures and pending jobs are visible in the [live checks](https://github.com/microsoft/onnxruntime-genai/pull/2921/checks). - This layer does not establish ordinary/speculative token equality, independent HuggingFace parity, or a full accuracy/performance campaign. #2923 adds a scoped arithmetic-boundary correction and opt-in ordinary path; its qualification must not be attributed to this layer. - Sampling, adaptive draft counts, dynamic batching, arbitrary rewind, and multi-device execution remain outside scope. <details> <summary>Historical qualification and provenance</summary> The original reconciled #2867/#2868 stack recorded eight implementation/test units compiled, **16 CPU tests passed**, and **two H100 tests passed**, using cached dependencies. The CPU coverage included acceptance/budget logic, EOS, invalid search/configuration rejection, target-only parsing/setter behavior, optimizer exclusions, public-token reset, and shared-buffer lifetime. The GPU tests covered shared-initializer lifetime in both session-destruction orders and the predecessor's PagedAttention padding/cache policy; they were not a new generation-throughput campaign. Original source anchor: `237754099f35f093c41e5d0cf3ce252d2dade854`. These are historical combined-stack results, not a clean build or rerun of current head `2e4886ef`. The earlier custom-stack benchmark numbers in #2868 are historical, not an isolated speedup from this PR. Its ordinary/k7 mismatch was not resolved by this layer. Migration comparisons must distinguish `custom k0 == contrib k0` and `custom k7 == contrib k7` from the stronger ordinary/speculative equality requirement. Plugin-free converted packages also depend on microsoft/onnxruntime#33897, microsoft/onnxruntime#33898, and microsoft/onnxruntime#33899. See #2868 for the original evidence and runtime discussion. </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

Rui Ren (rui-ren) requested a max review from Copilot November 4, 2026 23:22

Copilot started reviewing on behalf of Rui Ren (rui-ren) November 4, 2026 23:24

View session

Rui Ren (rui-ren) changed the base branch from main to rui-ren/bolt-stack-20261104/01-projections November 4, 2026 23:32

Copilot AI reviewed Nov 4, 2026

View reviewed changes

github-advanced-security Bot found potential problems Nov 5, 2026

View reviewed changes

Comment thread onnxruntime/test/contrib_ops/fused_state_space_op_test.cc Fixed

Rui Ren (rui-ren) requested a max review from Copilot November 5, 2026 00:54

Rui Ren (rui-ren) added the model: nemotron-bolt label Nov 5, 2026

Copilot started reviewing on behalf of Rui Ren (rui-ren) November 5, 2026 00:55

View session

Copilot AI reviewed Nov 5, 2026

View reviewed changes

Rui Ren (rui-ren) requested a review from Tianlei Wu (tianleiwu) November 5, 2026 05:09

Rui Ren (rui-ren) marked this pull request as ready for review November 5, 2026 19:30

Rui Ren (rui-ren) mentioned this pull request Nov 6, 2026

[CUDA] Consolidate source-exact BF16 decoding into existing contrib operators #33940

Draft

Rui Ren (rui-ren) marked this pull request as draft November 6, 2026 19:29

Rui Ren (rui-ren) changed the title ~~[2/3] [CUDA] Add fused Mamba verification and routing contrib operators~~ [CUDA] Add fused Mamba verification and routing contrib operators Nov 6, 2026

Rui Ren (rui-ren) mentioned this pull request Nov 9, 2026

[CUDA] Complete Bolt target-only decoding and shared-weight review fixes microsoft/onnxruntime-genai#2921

Merged

Conversation

Rui Ren (rui-ren) commented Nov 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Testing

Router top-k correction: production-provider qualification

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Rui Ren (rui-ren) commented Nov 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Rui Ren (rui-ren) commented Nov 9, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Rui Ren (rui-ren) commented Nov 4, 2026 •

edited

Loading

Rui Ren (rui-ren) commented Nov 5, 2026 •

edited

Loading