[Rui Ren (rui-ren)]

Description

Extends com.microsoft::MoEPacked (CUDA) for ReLU-square experts and preselected
INT32 expert IDs with FP32 per-token mixture coefficients. New projected-sigmoid
routing is composable only: use the existing MoEGroupRouter extended in #33994,
then connect its expert IDs/mixture coefficients to MoEPacked inputs 5/6 and set
normalize_routing_weights=0.

The learned dense router matrix belongs to the router operator, not MoEPacked.
MoEPacked's expert_weights input remains necessary: it holds dynamic per-token
mixture coefficients, not learned router parameters or packed expert matrices.

Removed the proposed routing_function/router_scale attributes, inputs 7-9
(router_input, router_weight, router_bias), integrated sigmoid kernels and
their FP32 scores workspace before release. Preselected routing supports k <= 32;
legacy internal softmax/SwiGLU behavior retains k <= 16. ReLU-square expert
arithmetic, GGUF decoding, packed layouts and MatMulPacked precision are unchanged
by this review fix.

Dependency: #33962 in
native stack #33967.
The complete eight-member stack, including merged #33953, is preserved.

Summary of Changes

The MoEPacked schema, CUDA host/dispatch state and reference documentation now
expose only legacy softmax or preselected routing. Focused regressions preserve
existing expert/precision coverage and reject the removed integrated interface.
The dependent router composition test now validates the composable path directly
against its host reference.

Router separation verification

Owning-layer head: f3b1c3100df803646f9302175d924c1928219dbf. Combined-stack head (#34065): de9e2b2dc1766ee58e427850d34bf208729e5642.
The official production Release CUDA source build was run at both heads, with
Flash enabled, QUICK_BUILD disabled and no binary-object overlays.

• 19 focused owning-layer provider tests and 51 combined-stack
  provider/schema tests passed, with zero failures or skips.
• The compiled MoEPacked schema has 5..7 inputs and exactly six attributes.
  Removed router attributes and the former ten-input signature are rejected.
  Scoped schema/kernel documentation matches the repository generators.
• The chained router regression checks expert IDs, FP32 mixture coefficients
  and FP16/BF16 expert outputs against the independent host reference.
• Current frozen stack before/after controls matched all tokens, first-call
  emission counts and integer speculative counters in 18 complete request
  pairs across ordinary/k7 modes, generation, placement and short timing smoke checks.
  All three native capture/reset, candidate-prefix replay and opt-in R1/R8
  committed-state tests passed on both runtimes. Required operators ran on CUDA.
• Six admitted H100s were used in parallel; physical GPU1 and the foreign GPU0
  client were left untouched. Runtime library hashes and immutable payload hashes
  were checked. Existing fixture value-1 MatMulPacked rounding attributes were
  removed for both controls because the current parent already makes that the
  default; no production MatMulPacked precision code changed.
• Changed-source clang-format, whitespace and six-layer preservation audits passed.
  Expert CUDA kernels are byte-identical to their respective frozen layers.

Evidence: /home/mobius/bolt-moe-api-20261110-c8741c5a/.
These are focused source-bound checks, not independent HuggingFace parity,
statistically qualified performance results or full-CI approval.

Earlier review snapshots (not current-head acceptance)

Testing

Review-fix snapshot: 79067e34feddff33a1b49d1bb56cfd21d141004c.

Added independent nonuniform CPU-reference coverage, true BF16 NaN verification, discriminating FP32 route weights, 22 invalid-configuration cases and missing SwiGLU paths. Removed dead TC ReLU branches. relu2 is canonical; relu_squared remains a compatibility alias. Expanded the authoritative schema and generated reference.

Owning-layer validation: 10 actual focused CUDA tests (including 22 negative configurations); 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.

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

• Router-separation review addressed with focused regressions
• Documentation and supported contracts updated
• Existing expert arithmetic and precision behavior preserved
• Complete native stack and immediate-parent dependencies preserved
• Full CI and independent reference qualification (not claimed)

[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)]

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)]

Re-review of 97c77c1bbb.

Verdict: COMMENT — no blocking findings; earlier feedback is addressed. Two non-blocking items remain (one inline, one as a reply on the still-open bot thread).

Earlier threads, re-checked on this head

• Nonuniform IQ2 coverage — fixed (RandomIq2Relu2MatchesIndependentCpuReference against DequantizeGgufRowCpu, plus ReluIq2PreservesContiguousKCancellation).
• TensorCheck<BFloat16> NaN expectation that could not fail — fixed with a custom verifier that checks the actual NaNs / finite values.
• FP32 vs. BF16-rounded routing weights indistinguishable — fixed (.3116f/.5098f, plus an ASSERT_TRUE(distinguishes_narrowed_weights) guard on the fixture itself).
• Missing negative tests — fixed (RejectsInvalidRoutingContracts: 22 message-checked cases covering every group I listed).
• Dead relu_squared branches in the generic tensor-core SwiGLU kernels — removed; relu2 is canonical with relu_squared kept as an alias; runtime-only constraints are documented in the schema; default-route SwiGLU (rows 4/48/129) and SwiGLU + preselected routes (k = 2, 32) are covered.

Independent verification (stack tip 623a725f5c, which contains this layer; Linux, CUDA 13.0, sm_80;sm_90a, Release; A100-SXM4-80GB)

• onnxruntime_provider_test builds; all 14 MoEPackedRoutingTest cases at the tip (10 exist at this layer) and the rest of the stack's 60 provider tests pass, also 3x in shuffled order. The pre-existing *Packed* suites (655 tests) show 0 failures.
• compute-sanitizer synccheck, racecheck and initcheck on the MoEPacked tests: no findings.
• CI on this head has no failures so far.

Remaining

1. API surface of the integrated sigmoid router — see the inline comment on routing_function.
2. The FC2 reduction order is still not pinned by any fixture — see my reply on the open thread (no author reply yet).