[Akshay Sonawane (apsonawane)]

Description

Clarify and test packed token-major input support for the MoE and QMoE
operators on the CPU and CUDA execution providers.

Packed inputs use:

• Input: [total_tokens, hidden_size]
• Router probabilities: [total_tokens, num_experts]
• Output: [total_tokens, hidden_size]

Tokens from variable-length sequences can be concatenated without padding
because MoE routing is token-local and does not require sequence boundaries.

Changes

• Update the MoE and QMoE schema documentation.
• Document packed input behavior for CPU and CUDA.
• Add CPU parity tests for packed MoE and INT4 QMoE.
• Add CUDA parity tests for packed MoE and INT4 QMoE.
• Fix the CPU test harness to pass processed routing probabilities to MoE
  while preserving raw router logits for QMoE.

No runtime kernel changes are required because the existing CPU and CUDA
implementations already process rank-2 token-major inputs.

[Copilot]

🟡 Changes recommended

The new MoE cases cannot construct their graphs, routing differs from the reference, CPU tests may silently pass, and generated docs are stale.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Documents and tests packed token-major inputs for CPU/CUDA MoE and QMoE.

Changes:

• Clarifies packed-input schemas and provider documentation.
• Adds CPU and CUDA parity tests.
• Corrects CPU routing-input selection.

File summaries

File	Description
contrib_defs.cc	Updates MoE/QMoE schemas.
moe_qmoe.md	Documents CUDA packed inputs.
qmoe.md	Documents CPU packed inputs.
test_qmoe_cpu.py	Adds CPU tests and routing fix.
test_qmoe_cuda.py	Adds CUDA tests.

Review details

Suppressed comments (3)

onnxruntime/test/python/transformers/test_qmoe_cpu.py:1471

• The new quant_bits=0 case cannot construct this fixture. PhiMoESparseMoeBlock keeps float expert weights when quant_bits == 0, but create_cpu_moe_onnx_graph unconditionally asserts at lines 423-426 that both weights are uint8 and scales are present. Guard those QMoE-only assertions with use_quant so the intended MoE case can run.

    @parameterized.expand([(0,), (4,)])

onnxruntime/test/python/transformers/test_qmoe_cuda.py:1593

• After the graph-builder assertions are fixed, the new MoE (quant_bits=0) case still compares different routing algorithms: ort_forward uses masked_sampling_omp_inference (lines 765-782), whereas PhiMoESparseMoeBlock.forward uses a softmax over the top-k logits. These weights differ for ordinary non-tied logits, so this parity test will fail for routing rather than validating packed input. Generate the MoE router tensor with the same top-k softmax as the reference.

        packed_moe.parity_check()

onnxruntime/test/python/transformers/test_qmoe_cpu.py:1486

• This CPU test can pass without executing either operator: parity_check() silently returns when model recreation fails (lines 1072-1074) or ORT execution returns None (lines 1078-1081), while those helpers catch session/run exceptions. Make these paths raise/assert, as the CUDA helper already does, so this test actually guarantees packed-input coverage.

        packed_moe.parity_check()

• Files reviewed: 5/5 changed files
• Comments generated: 3
• Review effort level: Balanced

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Copilot review overview

🟡 Changes recommended

The new non-quantized CUDA parity case applies softmax twice and therefore tests incorrect routing semantics.

Get a fresh assessment by requesting another Copilot review.

Review effort: Balanced
Findings: 1 · 1

Open (2)

• Avoid double softmax in non-quantized CUDA MoE routing · New
• Update PR description to reflect CPU routing changes · New

Resolved since last review (3)

• The checked-in generated operator reference was not updated with these schema descriptions:…
• The new quant_bits=0 case cannot construct this fixture. PhiMoESparseMoeBlock keeps float…
• For the new MoE (quant_bits=0) case, router_input is produced by…

[Copilot]

Copilot review overview

🟢 Approval recommended

The implementation and focused parity coverage are consistent; only a non-blocking PR description correction remains.

Review effort: Balanced
Findings: 1

Open (1)

• Update PR description to reflect CPU routing changes

Resolved since last review (1)

• Avoid double softmax in non-quantized CUDA MoE routing

[Ti-Tai Wang (titaiwangms)]

The documented packed-token contract is consistent with the existing CPU/CUDA input validation, so I am approving.

Non-blocking test suggestion: the new test_packed_token_input_* cases still use the same flattened 2-D input path as the existing parity tests. A stronger regression test would compare padded 3-D input with concatenated packed input, ideally using sequences with different lengths, and assert token-for-token output parity. It would also be useful to retain coverage of normalize_routing_weights=0 in at least one case.