[Rui Ren (rui-ren)]

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.

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.

Historical PR description and earlier-head evidence (not current-head acceptance)

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 [CUDA] Extend MoEPacked for top-k, dispatch and fused MoE calculation #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.

[Copilot]

Note

Copilot was unable to run its full agentic suite in this review.

Copilot review overview

3 open issues

• The scratch buffer is allocated as uint8_t and then reinterpret_cast to float*. This assumes… · New
• This CU file uses expf, isfinite, and INFINITY but does not include a header that clearly… · New
• The 1e-20f additive constant is a magic number. Please replace it with a named constant (and… · New

What changed in this PR

Adds a new CUDA contrib operator MoERouterPacked to provide projected-sigmoid routing for MoEPacked, including schema/registration, shared CUDA implementation, tests, and documentation updates.

Changes:

• Introduce MoERouterPacked operator schema and CUDA kernel registration.
• Implement routing CUDA kernels (sigmoid scoring + deterministic top-k) and refactor MoEPacked sigmoid-routing path to reuse them.
• Add CUDA-focused unit tests and update operator documentation.

File	Description
onnxruntime/​test/​contrib_ops/​moe_router_packed_test.cc	Adds CUDA tests for routing behavior, tie-breaking, edge cases, and validation.
onnxruntime/​core/​graph/​contrib_ops/​ms_opset.h	Registers the new operator schema in the Microsoft opset.
onnxruntime/​core/​graph/​contrib_ops/​contrib_defs.cc	Defines the MoERouterPacked operator schema, constraints, and documentation.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_router_packed_impl.h	Declares shared CUDA launcher and packed-load helper.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_router_packed_impl.cu	Implements sigmoid projection and deterministic top-k selection CUDA kernels.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_router_packed.h	Declares the CUDA kernel wrapper class for the new op.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_router_packed.cc	Implements kernel validation, scratch allocation, and launch wiring.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_packed.cu	Refactors existing sigmoid-routing to call the shared router launcher; reuses packed-load helper.
onnxruntime/​contrib_ops/​cuda/​cuda_contrib_kernels.cc	Registers MoERouterPacked CUDA kernels for FP16/BF16.
docs/​OperatorKernels.md	Documents kernel availability/signature for MoERouterPacked.
docs/​ContribOperators.md	Adds generated operator documentation section for MoERouterPacked.

_{Review effort: Lite}

[Rui Ren (rui-ren)]

Stack #33967 review — parent 5809f2531dcb -> head ebdbb8548801.

Reviewed the standalone schema/kernel registration, FP16/BF16 projection, selection-only bias, finite/exceptional ranking and lower-ID ties, selected-rank FP32 normalization, and the integrated MoEPacked call into the same shared implementation. No new actionable correctness finding in this extraction.

The current byte-scratch alignment comment is not a demonstrated defect: GetScratchBuffer delegates to the same CUDA allocator, and changing T changes the allocation's count units, not its alignment. The normal CUDAAllocator uses cudaMalloc; the arena preserves allocator alignment. A typed float scratch buffer can improve clarity, but it does not itself repair an alignment failure. Likewise, naming the inherited 1e-20f stabilizer may clarify intent, but preserve its value and normalization order when doing so.

This is static review feedback, not a new clean build, standalone-op runtime test, or approval of outstanding CI.

[Tianlei Wu (tianleiwu)]

Reviewed this layer (stack #33967, layer 4) at ebdbb85488.

Extraction fidelity. I diffed the removed MoEPackedSigmoid* kernels in moe_packed.cu against moe_router_packed_impl.cu after renames: kernel bodies, launch geometry and error checks are identical, and MoEPacked now goes through the shared LaunchMoERouterPackedCuda. I also ran the unmodified kernels on an H200 against an independent host reference (1,200 randomized FP16/BF16 cases with ties, +/-inf and NaN bias and activations, 16-byte-misaligned pointers, widths 256-2048, experts 1-512, k 1-32) and found no selection or weight mismatches. Schema, kernel registration and op-set registration look correct.

No blocking issues. Main requests (all inline):

1. The new tests use all-zero (or single-element) data, so common kernel regressions still pass them.
2. The ContribOperators.md entry cannot be regenerated from the schema; doc-gen validation will report a diff once this PR targets main.
3. Smaller API-robustness items: a lumped validation message and missing shape inference.

Non-blocking design notes

• The standalone op rejects rows > 2048, while MoEPacked chunks internally at 2048 rows, so a graph that swaps in MoERouterPacked must pre-chunk long prefills itself. If that is intentional, say "callers must chunk" in the schema text and pin it with a rows = 2049 rejection test; otherwise the same chunk loop can be used here, since the kernels have no row limit other than grid.x (the cap only bounds the FP32 score scratch).
• LaunchMoERouterPackedCuda is now a shared entry point whose memory-safety preconditions (width % 256 == 0, experts <= 512, 1 <= k <= min(32, experts)) are enforced only by its callers (moe_router_packed.cc for the op, DispatchMoEPackedCuda and moe_packed.cc for the integrated path). A host-side guard in the launcher is cheap and keeps the two call sites from drifting.
• The BF16 SM80 check runs in ComputeInternal on every Run(), whereas MoEPacked enforces it in the constructor, so an unsupported device fails at the first inference rather than at session creation.

On the existing Copilot comments: I agree the scratch-alignment concern is not a defect (same allocator, and arena chunks are 256-byte aligned); the <cmath> and epsilon-naming items are optional.

[Rui Ren (rui-ren)]

Updated in c7f2ea8 to reuse the existing CPU/CUDA MoEGroupRouter from #33817, with optional ordinary dense projection and compact IDs/FP32 weights. Router weights are not quantized; MoERouterPacked is removed and the legacy grouped interface remains supported.

Fresh source-built validation passed 40 provider/schema tests, all six native GenAI capture/state tests (including R1/R8), 36 real component cases and 240 same-mode generation checks. The audited seven-H100 1k/2k/8k TTFT/TPS table is in the updated description; observed decode overhead versus integrated routing is 0.28–0.51% in this small sample. These rates are for this layer, not the downstream performance-stack tip.

Coordination: #34002/#34003 branches and the complete native stack membership are untouched. Please synchronize those active downstream layers with the updated parent before combined-stack qualification; I have not force-rebased or merged them.

[Tianlei Wu (tianleiwu)]

Re-review of 232a1dfa59.

Verdict: COMMENT — but two compile errors introduced in this layer must be fixed before merge: they break the Windows CUDA and wasm32 builds, and #34002 / #34003 inherit them.

The standalone MoERouterPacked op from the previous revision is gone, so my earlier threads on it are moot. Their equivalents on the new code are in place: observed-value diagnostics in ValidateMoEGroupRouterProjection, shape inference that validates every known operand, a randomized comparison against an independent host reference, and the schema text as the source of truth for docs/ContribOperators.md (I compared the generated MoEGroupRouter section with bert_defs.cc; it matches).

Blocking (inline)

1. moe_group_router_impl.cu:157 — int64_t grid size → MSVC C4244 / C2220 on every Windows CUDA job (example).
2. linear_attention_gates.cc:300,302 — int64_t index into an InlinedVector → -Werror,-Wshorten-64-to-32 on wasm_Debug, wasm_Release and wasm_Release_static_library (example).

In the #34003 Windows CUDA job (which contains this layer's code) the other new/changed CUDA translation units (moe_packed.cu, matmul_packed.cu, gather_packed.cu, fused_state_space_impl.cu) compiled without warnings, so these two lines look like the whole delta for those platforms.

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

• Build OK; the stack's 60 provider tests pass, including MoEGroupRouterTest (9), MoEGroupRouterProjectionShapeInferenceTest (2) and MoEGroupRouterProjectionTest (11: randomized reference, ties, non-finite inputs, chained-vs-integrated, runtime-geometry negatives); compute-sanitizer synccheck, racecheck and initcheck are clean on them.
• The CPU compact path runs in the same tests and also passed in the Windows CPU debug job of #34003 (onnxruntime_provider_test passed there; that job's only failure is onnxruntime_shared_lib_cpuinfo_dlopen_test, which this stack does not touch).

Non-blocking

• Test gap for the compact-path guards — inline on CompactOutputsMustBePaired.
• Public API surface: with compact outputs feeding MoEPacked inputs 5/6 (the "Graph wiring" in the description), MoEPacked's integrated sigmoid router (routing_function, router_scale, inputs 7–9, added in #33963) is a second public form of the same computation. See my inline comment on #33963; worth deciding before either ships.
• Nit: the CPU compact path allocates two InlinedVectors of num_experts elements per row inside the parallel-for (heap allocations once num_experts exceeds the inline capacity); per-batch scratch hoisted out of the row body would avoid them. Low priority for a reference path.