[Akshay Sonawane (apsonawane)]

This pull request primarily updates documentation and test coverage for the QMoE operator, focusing on clarifying and expanding support for block-scaled FP8 quantization. The documentation now clearly distinguishes between legacy per-expert global scaling and the new square-block scaling for FP8 weights, and provides detailed descriptions of tensor shapes and behaviors for all quantization modes. Additionally, a new test is added for FP8 block compaction.

Documentation updates for FP8 block scaling:

• Expanded the QMoE operator documentation (docs/ContribOperators.md, docs/contrib_ops/cuda/moe_qmoe.md) to describe FP8 block scaling: positive block_size now selects square blocks with floating-point scales, while nonpositive values use per-expert global scales. The documentation clarifies tensor shapes, required/optional inputs, and dequantization behavior for both modes. [1] [2] [3] [4] [5] [6] [7] [8] [9] [10]

• Updated tables, attribute descriptions, and input/output lists to include block-scaled FP8 as a supported mode for relevant inputs (e.g., fc*_scales). [1] [2]

• Clarified kernel fallback logic and device support for block-scaled FP8, noting that block scaling triggers a fallback even on SM90. [1] [2] [3]

Test coverage:

• Added a new unit test source file for FP8 compaction (qmoe_fp8_compaction_test.cc) to the CUDA plugin internal test suite.

Other clarifications:

• Updated documentation to clarify the behavior of SwiGLU fusion and FC3 handling in the context of block-scaled FP8.

These changes ensure that both developers and users have clear guidance on how to use and implement block-scaled FP8 quantization in QMoE, and that the new functionality is covered by internal tests.

[Copilot]

Copilot review overview

🟡 Changes recommended

The CUDA kernel registration rejects several newly documented and tested block-scale data types before execution.

Review effort: Balanced
Findings: 1

Open (1)

• CUDA QMoE rejects supported independent block scale types · New

What changed in this PR

Adds CUDA QMoE square-block FP8 support with routed-expert compaction to reduce dequantization scratch memory.

Changes:

• Implements FP8 block-scale validation, dequantization, expert compaction, and split/fused SwiGLU handling.
• Adds CUDA, row-tiling, and Python parity/validation tests.
• Documents tensor layouts, fallback behavior, and memory usage.

File	Description
onnxruntime/​test/​python/​transformers/​test_qmoe_fp8_cuda.py	Adds block-FP8 integration tests.
onnxruntime/​test/​contrib_ops/​qmoe_row_tiling_test.cc	Tests active-expert capacity bounds.
onnxruntime/​test/​contrib_ops/​cuda_kernels/​qmoe_fp8_compaction_test.cc	Tests compaction and selected dequantization.
onnxruntime/​core/​graph/​contrib_ops/​contrib_defs.cc	Extends the QMoE schema documentation.
onnxruntime/​contrib_ops/​cuda/​moe/​qmoe_row_tiling.h	Calculates compact expert capacity.
onnxruntime/​contrib_ops/​cuda/​moe/​qmoe_kernels.h	Declares compaction and extended FP8 kernels.
onnxruntime/​contrib_ops/​cuda/​moe/​qmoe_kernels.cu	Implements compaction and block dequantization.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_quantization.cc	Integrates block FP8 into CUDA QMoE execution.
onnxruntime/​contrib_ops/​cpu/​moe/​moe_helper.h	Adds block-scale validation.
docs/​ContribOperators.md	Updates generated operator documentation.
docs/​contrib_ops/​cuda/​moe_qmoe.md	Documents block FP8 behavior and constraints.
cmake/​onnxruntime_unittests.cmake	Registers the new CUDA kernel test.

---

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

[Copilot]

Copilot review overview

🟡 Changes recommended

A negative test follows the wrong validation path, generated kernel documentation is stale, and the schema documentation gives conflicting weight-layout guidance.

Review effort: Balanced
Findings: 1 · 2

Open (3)

• Transpose all weights to reach the row-major validation · New
• Regenerate kernel documentation for the widened T2 constraint · New
• Document block-scaled FP8 weights as row-major · New

Resolved since last review (1)

• CUDA QMoE rejects supported independent block scale types

[Copilot]

Copilot review overview

🔵 Needs a closer look

The new tests fail rather than skip when PyTorch lacks the FP8 dtype.

Review effort: Balanced
Findings: None

Resolved since last review (3)

• Transpose all weights to reach the row-major validation
• Document block-scaled FP8 weights as row-major
• Regenerate kernel documentation for the widened T2 constraint

Previously missed (1)

In code that hasn't changed since last review

Guard torch.float8_e4m3fn access for older PyTorch versions

onnxruntime/​test/​python/​transformers/​test_qmoe_fp8_cuda.py:335

This helper accesses torch.float8_e4m3fn without the availability guard used by the legacy FP8 helper above and by the NVFP4 tests. On CUDA test environments with an older PyTorch build, every new class test plus the standalone compaction test fails with AttributeError instead of being skipped. Raise SkipTest here so all callers are covered.

[Akshay Sonawane (apsonawane)]

Pushed bd24e9f to address review feedback and the PR-related CI failures:

• Skip block-FP8 tests cleanly when PyTorch lacks torch.float8_e4m3fn.
• Fix MoETest.QMoETest_MixedWidthCudaRejectsUnsafeFallbackInputs: raw Linux CUDA plugin logs confirm the input was correctly rejected by the earlier guard, but the test expected the older message (fc1_scales dtype to match the input dtype versus fc1_scales must match the activation type). Updated the expectation; the safety guard remains intact.
• Restore docs/ContribOperators.md exactly from the updated-docs artifact for run 36823911063. The mismatch was whitespace on one generated blank line. OperatorKernels.md already matches the artifact.

The WebGPU job separately failed configuring Dawn because third_party/spirv-tools/src lacked CMakeLists.txt; the iOS job was cancelled. Neither was changed in this patch. New-head CI should re-evaluate these; no stale-head reruns were requested.

Validation: Ruff check/format, Python missing-dtype guard execution, changed-line C++ formatting, and exact generated-doc comparison passed. No local builds or GPU tests were run. CI results for the new commit remain pending.

[Copilot]

Copilot review overview

🔵 Needs a closer look

The CUDA routing, workspace, and dequantization changes require final human review and GPU execution validation.

Review effort: Balanced
Findings: None

[Ti-Tai Wang (titaiwangms)]

The block-FP8 layout/scale indexing and routed-expert compaction look consistent, but the added integer-scale validation regresses the existing prepacked INT4/INT8 path. All four failing CUDA/TensorRT test jobs report the same Missing Input: s1 failure in the existing QMoE expert-counting tests. Please address the inline finding.

Review used pinned source and existing CI logs; no local GPU execution. Performance, all-NaN routing behavior, and the named external checkpoint's layout were not validated.

[Copilot]

Warning

Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.

Copilot review overview

Review effort: Lite
Findings: 1 · 2 · 1

Open (4)

• helper.make_tensor(..., TensorProto.BFLOAT16, ..., vals, raw=False) is not reliably representable… · New
• This introduces a CUDA-only hard requirement that INT fc*_scales match the activation type… · New
• This assigns FLOAT8E4M3FN to *_zero_points tensors, but zero-points are UINT8 in the QMoE… · New
• cudaFree errors (including async failures from earlier work) are dropped here. For internal CUDA… · New

[github-actions]

You can commit the suggested changes from lintrunner.

[Ti-Tai Wang (titaiwangms)]

The previous INT4/INT8 prepacked-scale blocker is fixed in source: dtype validation reuses the guarded fc1_scales/fc2_scales pointers instead of retrieving consumed initializer inputs again. Reviewed the additional NVFP4 compaction, original-versus-compact expert indexing, bias gathering, deferred dequant lambda lifetimes, and expert-map fallback against their direct consumers. No blocking introduced defect found; one nonblocking test-cleanup concern is inline.

CUDA/TensorRT runtime checks have not yet completed at this head, so this approval does not claim that the previous Missing Input failure has been runtime-refuted or waive those checks. No local GPU execution/benchmarking. Tiled-prefill performance, NaN-routing behavior, and unrelated changes brought in by the main merge are excluded.

[Tianlei Wu (tianleiwu)]

The guarded scale-pointer fix addresses the consumed-initializer regression; I confirmed it in PrePack/Compute and resolved that thread. One new test failure remains because the runtime diagnostic changed without updating its two assertions, detailed inline. No additional introduced defect was substantiated in the authored compaction/expert-map follow-ups. The existing supported-GPU evidence request remains open; current CUDA CI is pending and reported static checks do not establish parity execution. No local build/test was run.

[kunal-vaishnavi]

Pushed 20aa6ff to fix the current CI failures and code-review findings:

• All six CUDA/CUDA-plugin/TensorRT test jobs first failed in QMoETest_MixedWidthCudaPackedPrefillWorkspaceLimit with an illegal memory access, then cascaded through the poisoned CUDA context. Scratch was sized using a reduced row tile, but the execution plan disabled tiling for all packed INT routes. Prefill now honors the bounded tiles; only packed INT GEMV remains untiled. Existing workspace-limit and row-tiling regression tests cover this path.
• Preserve the integer-scale diagnostic substring used by both C++ and Python negative tests, without changing dtype enforcement or consumed-initializer guards.
• Make CudaBuffer cleanup explicitly noexcept and log cudaFree errors instead of throwing during unwinding.
• Accept the precise shared-T1 schema rejection in the FP8 zero-point negative test.
• Match ContribOperators.md exactly to the generated-doc artifact from run 37059832750, including generator-owned whitespace. OperatorKernels.md already matched.

Validation: focused compiled checks of the actual cleanup helper, prefill-plan expression, and integer-scale guard; six row-tiling gtests from the prior exact-head CI binary; changed-line clang-format; Ruff check/format; and exact generated-doc comparison passed. The FP8 Python subset passed 18 tests plus 36 subtests on A100 using the existing dev wheel, with official 512-expert cases enabled. That wheel is not PR-head certified; the downloaded PR CI wheel has FP8 disabled and skips those cases, so the GPU evidence thread remains open with this limitation documented. No ORT source rebuild was performed, and patched runtime CUDA gtests await new-head CI. The push has started fresh checks; no stale-SHA reruns were requested.

[Ti-Tai Wang (titaiwangms)]

Reviewed the incremental changes. CudaBuffer cleanup is now noexcept and reports cudaFree errors without throwing during unwinding, addressing the previous nonblocking comment. The row-tiling gate now disables tiling for packed INT GEMV only, preserving the packed prefill plan used to size and execute its bounded scratch; the prefill loop advances input/output/routing offsets per tile. The dtype-error wording and negative-test acceptance change are consistent with validation occurring either at schema binding or in the kernel. No actionable introduced defect found. Current captured CI is complete with no failures, including CUDA/TensorRT tests and WASM checks. No local GPU execution or performance/NaN-routing measurement.

PR review comments have been addressed and resolved.