[Rui Ren (rui-ren)]

Description

Extends com.microsoft::MoEPacked (CUDA) for Nemotron Bolt's ReLU-square experts and
routing contract without introducing a competing expert operator. Adds
activation_type="relu_squared", preselected expert IDs with FP32 mixture weights,
and integrated projected-sigmoid routing. Correction bias affects selection only;
mixture weights come from unbiased sigmoid affinities and post-normalization scaling.
The existing softmax/SwiGLU defaults remain supported.

Preselected and projected-sigmoid routing support up to 32 selected experts; the
existing internal softmax path retains its 16-expert top-k limit. Validation keeps
routing modes mutually exclusive and isolates malformed routes to their token.
BF16 IQ2 FC1/FC2 preserve complete FP32 scale reconstruction, contiguous-K partitions,
BF16 storage boundaries, and FP32 routing/final reduction.

Dependency: #33962 in
native stack #33967;
the eventual landing target is main. Packed weight layouts and GenAI source do not change.

Summary of Changes

Operator Contract and Documentation

File	Change
onnxruntime/core/graph/contrib_ops/contrib_defs.cc	Extend activation/routing attributes and optional preselected/projected-router inputs.
docs/ContribOperators.md, docs/OperatorKernels.md	Document routing modes, FP32 weights, and optional inputs.

CUDA Implementation and Regression Tests

File	Change
onnxruntime/contrib_ops/cuda/moe/moe_packed.{cc,h,cu}	Validate routes and workspace bounds; add ReLU-square, preselected and sigmoid routing; preserve decode and 32-route prefill expert arithmetic.
onnxruntime/test/contrib_ops/moe_packed_routing_test.cc	Cover FP16/BF16 preselected weights, malformed-token isolation, selection-only bias, unchanged SwiGLU defaults, and IQ2 cancellation.

Testing

Review snapshot (2026-11-09): e50e34972b245acf51d3f259efe992a37bd6b24f.

• Native filter: MoEPackedRoutingTest.* in a CUDA-enabled test build.
• A focused H100 ABI-matched overlay with this layer's MoEPacked TU compiled as
  compute_75-only PTX passed ReluIq2PreservesContiguousKCancellation.
• A separate stack-tip overlay at Reduce native packed decode tile and FC2 loader overhead #34003 (bbac6c023794) passed all 19
  projection/routing regressions with zero skips. That is integration evidence for
  the later stack, not a clean standalone build of this PR.
• Remaining coverage: add asymmetric channels/packed rows against an independent
  dequantized reference for decode and full 32-route prefill
  (thread),
  plus mixed-sign FC2 cancellation with inter_size >= 1024
  (thread).
  Existing fixtures do not establish those properties.
• No new clean full-source build, independent HF accuracy run, or current-head
  full-model/performance acceptance is claimed. Historical routing timings below
  compare routing alternatives in one build, not this entire PR against its parent.

Detailed review.

Historical measurements and qualification evidence (not current-head acceptance)

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, IQ2_XXS, ReLU-square experts, latent hidden dimension 1024, top-k 9. Internal softmax routing versus preselected FP32 indices/weights using identical routes and the same corrected expert arithmetic. Selected logits are +16 and others -16; selected weights are 1/9. Shared/distinct labels describe expert reuse across rows (the one-row variants are equivalent workload controls). Sixteen rows create 144 routes, crossing the 128-route prefill-tiles dispatch threshold.

The ratio measures routing-path overhead removed, not this entire PR versus its parent. All 24 complete BF16 outputs are bitwise equal in both runs. Benefits diminish for distinct-expert, bandwidth-heavy batches; all cases, including those smaller gains, are shown.

Intermediate	Experts	Rows	Routes across rows	Internal routing us	Preselected us	Ratio
1536	240	1	distinct	54.71	41.35	1.32x
1536	240	1	shared	53.31	40.51	1.32x
1536	240	8	distinct	205.41	190.60	1.08x
1536	240	8	shared	56.53	43.36	1.30x
1536	240	16	distinct	396.04	381.20	1.04x
1536	240	16	shared	94.66	80.74	1.17x
2048	240	1	distinct	63.38	49.42	1.28x
2048	240	1	shared	62.54	48.42	1.29x
2048	240	8	distinct	252.80	238.01	1.06x
2048	240	8	shared	64.43	51.36	1.25x
2048	240	16	distinct	495.45	480.54	1.03x
2048	240	16	shared	112.66	99.46	1.13x
2304	432	1	distinct	79.56	57.62	1.38x
2304	432	1	shared	77.32	56.37	1.37x
2304	432	8	distinct	322.00	297.79	1.08x
2304	432	8	shared	82.72	61.52	1.34x
2304	432	16	distinct	670.85	653.92	1.03x
2304	432	16	shared	182.34	161.18	1.13x
2560	432	1	distinct	80.40	59.66	1.35x
2560	432	1	shared	78.78	58.72	1.34x
2560	432	8	distinct	325.76	308.10	1.06x
2560	432	8	shared	82.50	61.28	1.35x
2560	432	16	distinct	689.52	666.84	1.03x
2560	432	16	shared	196.47	175.37	1.12x

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

Checklist

• Tests added/updated
• No breaking schema changes (behavior changes documented above)
• Documentation updated (if applicable)
• Add the linked asymmetric-reference and FC2 cancellation coverage

[Rui Ren (rui-ren)]

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.

[Rui Ren (rui-ren)]

The native kernel optimizations are now pushed in #34002 and registered as the
new top of merge stack #33967.
This changes the intended merge stack, not alternative #33940. Existing layer
heads and their history are preserved.

Actual #33994-tip baseline versus the new native candidate, real 1028-1663-token
inputs: 54.25 -> 92.04 ordinary tok/s; 201.48 -> 318.91 k=7 tok/s.
Both arms used the same native package and corrected GenAI. All 240 complete
measured generations/counters, 1,782 full live tensor/state hashes per arm and
22 native CUDA regressions passed; no skips.

The earlier ~92 tok/s private baseline already contained these optimizations:
it was not the unmodified remote stack. The separate #33940 ~109/~370 result is
still not a result of this native candidate. That remaining gap is outstanding.
Full protocol, source pins and limits are in #34002; it remains draft.

[Copilot]

🟡 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

• Add asymmetric fixtures to catch channel permutations in BF16 IQ2 kernel · New
• Cover FC2 accumulation order with mixed-sign, cancellation-sensitive fixtures · New

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}

[Tianlei Wu (tianleiwu)]

Still open on e50e34972b, and a bit broader than channel permutations: the fixtures are also constant along K. Every block uses the same d, grid index 1 in all groups of all 8 sub-blocks, and each row's activation is a single value (factor / 32), so a mis-pairing of the A/B K fragments inside a thread's 16-element slice, or a wrong sub-block/half selection in the decode, would pass too.

For scoping: I ran the kernel text from this head outside the suite with random IQ2_XXS payloads (random bytes plus a random FP16 d per block, so all grid/sign/scale codes occur), random BF16 activations, and an independent reference (gguf_quant::DequantizeGgufRowCpu -> BF16 -> double accumulation -> the same BF16 rounding points), for hidden 256-2048, inter 256-1536 and k up to 32, in decode and prefill (4-subtile) tiles. FC1, FC2 and the finalize stage agree with the reference except for isolated BF16 rounding flips (about 0.02% of elements at most), and compute-sanitizer is clean, so I don't see a current permutation bug; this thread is about regression protection.

Suggested shape: any 66-byte pattern with a sane FP16 d is valid IQ2_XXS, so random packed bytes plus the host-side DequantizeGgufRowCpu (already in gguf_quant_tables.h) give an independent reference without extra infrastructure. Compare with a small ULP tolerance plus a cap on the fraction of mismatches: isolated elements differ by a few ULPs, mostly at BF16 rounding boundaries (and ReLU-square amplifies them in FC1).

[Tianlei Wu (tianleiwu)]

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

1. MalformedRoutesPoisonOnlyTheirToken does not verify NaN poisoning in BF16: TensorCheck<BFloat16> checks isnan(expected) instead of the actual value (inline on the test).
2. The "FP32 weights" fixtures give identical BF16 expectations if the weights are rounded to BF16 (inline), so they cannot catch narrowing.
3. 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.
4. Please clarify the compute_75 overlay result. The PR description and the author review say a compute_75-only PTX overlay passed ReluIq2PreservesContiguousKCancellation. With compute_75 PTX the MMA body is compiled out (MmaBf16 is guarded by __CUDA_ARCH__ >= 800), GetMoEPackedBf16KernelSupport reports 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 to 2^25 at 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 as RoundedBf16ContiguousKCancellation in #33962, and HasCudaEnvironment(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/weights and sigmoid router_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_squared now 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_bias supplied, 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 + TopK comment in moe_packed.cu is 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. With relu_squared there is no up tensor 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.

[Tianlei Wu (tianleiwu)]

The BF16 NaN expectation here cannot fail, so this test does not verify the poisoning.

TensorCheck<BFloat16> (checkers.cc#L545-L549) does if (std::isnan(f_expected[i])) { EXPECT_TRUE(std::isnan(f_expected[i])) ... }, i.e. it tests the expected value rather than f_actual[i] (the MLFloat16 checker a few lines above uses f_actual). MalformedRoutesPoisonOnlyTheirToken only instantiates the default T = BFloat16, so any output for token 0 (finite garbage, stale workspace, zeros) satisfies the NaN expectation, and the test effectively only asserts that token 1 is unaffected. The checker is unchanged up to #34003, so this is not fixed later in the stack.

Suggested fix: a custom verifier replaces the default comparison, so it has to check both tokens (the Mamba2Verify test in #33969 works around the same checker behavior this way):

test.SetCustomOutputVerifier([&](const std::vector<OrtValue>& fetches, const std::string&) {
  const auto out = fetches[0].Get<Tensor>().DataAsSpan<T>();
  for (int col = 0; col < hidden_size; ++col) {
    EXPECT_TRUE(std::isnan(out[col].ToFloat())) << "token 0 col " << col;
    EXPECT_EQ(out[hidden_size + col].ToFloat(), expected[hidden_size + col].ToFloat()) << "token 1 col " << col;
  }
});

(for the malformed != 0 cases), and/or also run the malformed cases for MLFloat16, whose checker is correct. Fixing TensorCheck<BFloat16> itself to test f_actual[i] would cover every BF16 test, but it may surface other latent expectations, so it is probably better as a separate PR.

[Tianlei Wu (tianleiwu)]

These fixtures cannot tell FP32 expert weights from BF16-rounded ones.

Replaying this helper's expected-value math with the real BFloat16/MLFloat16 conversions, the BF16 expectation is identical if weight is first rounded to T, for every case in ReluSquaredPreservesPreselectedFp32Routing (rows {1, 8, 9, 17} x k {2, 32}) and in the BF16 half of SigmoidCorrectionBiasUsesUnbiasedWeights: .333333f and .5001f are within a third of a BF16 ulp of their roundings, so the difference vanishes in the final BF16 rounding (the committed expectations are 8.1875 / 29.25 for k=2 and 119.5 / 474 for k=32 either way). A regression that narrows expert_weights (or the normalized sigmoid weights) to T would therefore pass. FP16 is only partly better: it differs for k=2 only.

Weights that do discriminate: {.3301f, .501f} give expected 8.1875 / 29.125 / 120 / 476 for (even row, k=2) / (odd row, k=2) / (even row, k=32) / (odd row, k=32) with FP32 weights, but 8.125 / 29 / 119.5 / 474 with BF16-rounded weights.

[Tianlei Wu (tianleiwu)]

The new validation matrix has no negative tests.

This block adds about ten distinct rejection rules (the constructor adds more), but moe_packed_routing_test.cc has no OpTester::ExpectResult::kExpectFailure case, so relaxing any of them would go unnoticed. Several guard memory safety or silent truncation (router_width % 256, where the scores kernel would otherwise drop or over-read the row tail; the shape checks; and k <= 16 for the internal softmax kernels). Suggested cases, each a few lines on top of the existing helpers: (1) router_probs together with expert_indices; (2) preselected routing with normalize_routing_weights=1; (3) expert_indices without expert_weights; (4) sigmoid without relu_squared, with router_probs, with router_width % 256 != 0, and with a wrong router_weight/router_bias shape; (5) relu_squared with up_experts_weights present, or with non-IQ2_XXS gate/down; (6) internal softmax with k > 16; (7) swiglu without up_type; (8) router_scale != 1 with softmax.

Minor: the attribute-only rules (relu_squared => IQ2_XXS, sigmoid => relu_squared + normalize=1, router_scale) can be checked in the constructor so they fail at session creation rather than at the first Run.

[Tianlei Wu (tianleiwu)]

Nit: this relu_squared branch (and the matching ones at L581, L720 and L753 in the generic Tensor-Core SwiGLU kernels) is unreachable. The constructor limits relu_squared to IQ2_XXS gate and down weights, and DispatchMoEPackedCuda sends exactly that combination to MoEPackedReluIq2ContiguousKKernel on the BF16 Tensor-Core path (or to the SIMT kernels), so the generic Tensor-Core kernels never see p.relu_squared == true. The dead code has no test and costs a runtime branch plus an 8-register zero-init (float u_vals[8] = {}) in the hot SwiGLU decode/prefill kernels. Either drop these edits (the SIMT ones are live and should stay) or add the dispatch path and the test that need them.

[Tianlei Wu (tianleiwu)]

Naming: MoE and QMoE already accept activation_type='relu2' for squared ReLU (contrib_defs.cc#L1711, parsed in moe_base.h#L52). Introducing relu_squared for MoEPacked spells the same activation two ways inside one operator family, and contrib schema strings are effectively frozen once released. Prefer relu2 (optionally also accepting relu_squared for readability), and update the kernel, tests and docs accordingly.

[Tianlei Wu (tianleiwu)]

Document the constraints that are only enforced at runtime. The schema and docs/ContribOperators.md do not say that: relu_squared currently requires IQ2_XXS gate and down weights (moe_packed.cc#L76-L78); router_width must be a positive multiple of 256 (moe_packed.cc#L138); the internal softmax path is limited to k <= 16 while sigmoid/preselected routing allow 32 (the k attribute text is unchanged); and router_scale/normalize_routing_weights are only meaningful with specific routing modes. Related wording: the Tweight constraint text ("Preselected routing weight type") is also used by router_bias, and the operator-level SetDoc still describes "gate, up, and down" blobs. Users otherwise only discover these through ORT_ENFORCE messages.

[Tianlei Wu (tianleiwu)]

This "defaults unchanged" test stays in the decode kernels, and SwiGLU + preselected routes is untested.

With k = 1, rows {4, 48} give 4 and 48 routes, both below the 128-route prefill threshold (use_prefill_tiles requires more than 128 routes), so only the decode Tensor-Core kernels run, while this PR also edits the prefill kernel MoEPackedFc1SwiGLUBf16TcKernel. Please add a case with >= 129 routes (e.g. rows = 129). Validation also admits SwiGLU with preselected routes and k up to 32 (nothing requires relu_squared for preselected mode), but no test exercises it; one small preselected SwiGLU case (k = 2 and k = 32) would pin that contract down.