[Akshay Sonawane (apsonawane)]

Summary

• Optimize quantized MoE finalization and FP8 block-scaled projection decode.
• Preserve upstream packed-INT, native decode and FP32-output paths.
• Add split-output FP8 projections and FP8-to-BF16 dequantization with independent FP32 scales, without a tensor-PR dependency.

Clean main-based replacement

Supersedes #33463. Reconstructed directly on main (28225b8583474cf6a6e7f97222227ae668f80e8b), without the old divergent feature-branch history. This PR targets main directly and has no dependency on another operator PR. Family documentation is bundled here; the separate old documentation PR #33465 is retired.

Review and CI fixes

• Enable the FP4 QMoE internal regressions in both ordinary and plugin CUDA builds, guarded by the internal-test, FP4 and contrib options. Plugin-only implementation linkage is unchanged.
• Implement the missing Float8E4M3FN_BFloat16 CUDA factory, with FP32 scales kept separate from BF16 outputs. The kernel supports scalar and per-axis scales, negative axes, optional zero points and empty inputs, and rejects malformed quantization shapes.
• Add model-level CUDA regressions for FP32 scale precision, non-last negative axes, zero points, rank-one per-axis inputs, empty inputs and invalid scales. Successful new cases use zero absolute/relative output tolerance.
• Regenerate the serialized layout-transformation kernel-type resolver to separate DequantizeLinear's scale (TS) from its output (T2). The stale resolver caused KernelTypeStrResolverUtilsTest.VerifyLayoutTransformationRequiredOpsResolver failures across Linux, Windows/plugin, WebGPU, wasm, macOS, Android NNAPI, QNN and OpenVINO jobs.
• Regenerate the affected operator sections and CPU/CUDA/DML kernel rows using the official documentation renderers and source-derived schema/registration metadata.

The Windows CUDA log also confirms the missing FP8/BF16 factory as an unresolved external. The React Native Android failure is an npm ETIMEDOUT; the iOS log shows a Signal 11 crash in the existing JSI AsyncWorker::toPromise path. Neither mobile failure has been established as caused by this PR, and neither is claimed fixed by these operator changes. New-head CI must establish their current status.

Validation

• Used the qwen-flash Conda environment.
• Ran lintrunner -a before committing, followed by a passing check-only run on the full PR delta.
• Changed C++/CUDA files pass clang-format --dry-run --Werror; source/CMake whitespace checks pass.
• A compiler-free CMake fixture validates all 16 combinations of internal tests, FP4, contrib and plugin flags against the shared definition block.
• Decoded the regenerated FlatBuffer and verified all 51 operator maps: only com.microsoft:DequantizeLinear:1 changes, and its mapping matches the checked-out schema.
• Verified the new numerical reference data, including a case that distinguishes FP32 scale multiplication from premature BF16 scale rounding.
• Documentation provenance passes byte-for-byte against the official renderers; the metadata adapter also reproduces the affected origin/main documentation. Unrelated sections/rows are unchanged. Official generated whitespace is retained.
• Per request, no ONNX Runtime/GenAI rebuild or native/GPU test execution was performed for these review fixes. Full native --gen_doc validate, linking, and runtime regression execution are left to CI. The documentation review thread remains open pending native validation.
• Earlier targeted CUDA/host and test translation units were compiled before these review fixes; those historical checks do not validate this new head. Full main-based model/GPU requalification remains outstanding; keeping this PR in draft.
• No benchmark/profiling/evaluation scripts or local task metadata are included.

[Copilot]

🟡 Changes recommended

The FP8-to-BF16 kernel factory is undefined, regression coverage excludes non-plugin builds, and generated documentation is stale.

3 open issues

• Enable FP4 QMoE tests in non-plugin CUDA builds · New
• Implement and register FP8-to-BF16 CUDA kernel · New
• Regenerate operator and kernel documentation · New

What changed in this PR

Optimizes ONNX Runtime’s CUDA quantized MoE and block-scaled FP8 projections in a standalone, main-based replacement.

Changes:

• Adds small-decode routing and fused MoE finalization paths.
• Adds split-output FP8 projections and FP8-to-BF16 schema/registry support.
• Adds regression tests and CUDA test wiring; full GPU requalification remains outstanding.

File	Description
onnxruntime/​test/​contrib_ops/​moe_test.cc	Tests column-tiled finalization.
onnxruntime/​test/​contrib_ops/​matmul_block_scaled_fp8_test.cc	Tests split outputs and invalid widths.
onnxruntime/​test/​contrib_ops/​dequantize_fp8_bf16_test.cc	Adds FP8-to-BF16 regression.
onnxruntime/​test/​contrib_ops/​cuda_kernels/​qmoe_fp8_compaction_test.cc	Adds NVFP4 routing and kernel regressions.
onnxruntime/​core/​graph/​contrib_ops/​quantization_defs.cc	Extends dequantization types and output inference.
onnxruntime/​core/​graph/​contrib_ops/​contrib_defs.cc	Defines split-output projection schema.
onnxruntime/​contrib_ops/​cuda/​quantization/​quantize_dequantize_linear.cc	Adds scale-type constraints.
onnxruntime/​contrib_ops/​cuda/​moe/​qmoe_kernels.h	Declares small-decode routing helper.
onnxruntime/​contrib_ops/​cuda/​moe/​qmoe_kernels.cu	Implements small-decode routing maps.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_quantization.cc	Integrates optimized routing and finalization.
onnxruntime/​contrib_ops/​cuda/​math/​matmul_block_scaled_fp8.h	Declares split-output interfaces.
onnxruntime/​contrib_ops/​cuda/​math/​matmul_block_scaled_fp8.cu	Routes tensor-core stores into split outputs.
onnxruntime/​contrib_ops/​cuda/​math/​matmul_block_scaled_fp8.cc	Allocates split outputs and handles fallback copies.
onnxruntime/​contrib_ops/​cuda/​llm/​moe_gemm/​moe_kernels.cu	Tiles decode finalization across columns.
onnxruntime/​contrib_ops/​cuda/​llm/​moe_gemm/​moe_gemv_fp4.h	Declares fused NVFP4 finalization.
onnxruntime/​contrib_ops/​cuda/​llm/​moe_gemm/​moe_gemv_fp4.cu	Implements fused finalization and adjusts tiling.
onnxruntime/​contrib_ops/​cuda/​cuda_contrib_kernels.cc	References FP8-to-BF16 kernel factory.
docs/​ContribOperators.md	Removes a trailing blank line only.
cmake/​onnxruntime_unittests.cmake	Enables NVFP4 regressions for plugin builds.

_{Review effort: Balanced}

[Copilot]

🟡 Changes recommended

Dequantization integration and edge-case handling remain incorrect, and the new NVFP4 tests need device-capability guards.

4 open issues

• Allow empty per-axis scales for empty input dimensions · New
• Exclude unsupported FP8 dequantization cases from QDQ propagation · New
• Constrain CPU scale types after separating TS from T2 · New
• Skip NVFP4 tests on devices below SM80 · New

3 resolved since last review

• Regenerate operator and kernel documentation
• Implement and register FP8-to-BF16 CUDA kernel
• Enable FP4 QMoE tests in non-plugin CUDA builds

_{Review effort: Balanced}

[Akshay Sonawane (apsonawane)]

Pushed 3a8c804 and addressed the four outstanding review findings:

• Permit an empty per-axis FP8-to-BF16 scale only for empty inputs, retaining axis-length validation.
• Prevent new MS-domain DQ cases from propagating into unsupported inverse QuantizeLinear nodes; add transformer and fully optimized CUDA-session regressions through Reshape/Transpose.
• Constrain scale types on all seven CPU MS-domain DQ registrations, regenerate the CPU kernel documentation row, and add CPU session success/rejection regressions.
• Skip the NVFP4 gate/up and fused-finalization tests on actual devices below SM80, reusing the existing capability query.

Full-PR lintrunner -a/check-only lint, git diff --check, source-derived official-renderer checks, baseline renderer selftest, and static contract/guard checks passed. No local builds, compilation, native tests, or model execution were performed; the new C++ regressions remain unexecuted locally. At inspection there were no failed head checks, with some CI still running. Fresh post-push CI is pending and is not claimed green.

[Copilot]

🟡 Changes recommended

Unresolved optimizer compatibility, backend capability checks, and contrib-disabled test failures block approval.

1 open issue

• Guard propagation tests when contrib operators are disabled · New

4 previously missed issues

In code that hasn't changed since last review

Reject unsupported split outputs in WebGPU capability checks

onnxruntime/​core/​graph/​contrib_ops/​contrib_defs.cc:3211

[P2] Exclude unsupported split-output nodes from WebGPU capability selection, or implement splitting there too. The existing FP16 WebGPU kernel still allocates only output 0 with the full [..., N] shape (contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.cc:132–134), and WebGpuExecutionProvider::GetCapability has no split-output check. It therefore claims a node with split_sizes=[8,9] even though the kernel produces neither required output shape and never allocates output 1. This breaks execution and can prevent a lower-priority CUDA provider from handling the node. Add a capability or split-output regression.

Validate DirectML scale and output type compatibility

onnxruntime/​core/​graph/​contrib_ops/​quantization_defs.cc:206

[P2] Update DirectML's capability checks for the independent scale type. Its MS DequantizeLinear registration in Operators/OperatorRegistration.cpp:943 still constrains only T1/T2 and has no support query. An int8 input with FP16 scale and output_dtype=FLOAT therefore matches, but DmlOperatorElementwiseQLinear directly binds the FP16 scale and FP32 output to a DirectML operator that requires matching types. Initialization fails instead of falling back to the supported CPU kernel. Constrain TS to supported types and reject scale/output mismatches during capability selection, or implement the conversion, with regression coverage.

Preserve QDQ pairs with incompatible boundary types

onnxruntime/​core/​graph/​contrib_ops/​quantization_defs.cc:234

[P2] Check boundary types before QDQ final cleanup removes a pair. With session.enable_quant_qdq_cleanup=1, a FLOAT graph input → ONNX QuantizeLinear-19 producing FP8 → MS DequantizeLinear with output_dtype=BFLOAT16 → graph output passes IsQDQPairSupported when both nodes share constant scalar scale and zero point. That helper compares quantization parameters, not the input/output boundary types. qdq_final_cleanup.cc:127–130 then creates a FLOAT-to-BFLOAT16 Identity, which is invalid. Leave pairs with differing boundary types intact and add a cleanup regression.

Guard QDQ synthesis in transpose optimizer

onnxruntime/​core/​optimizer/​qdq_transformer/​qdq_propagation.cc:378

[P1] Apply the same eligibility restriction in the transpose optimizer's MakeQDQNodeUnit. A runtime FP8 input → MS DequantizeLinear with scalar FP32 scale, an explicit FP8 zero point and output_dtype=BFLOAT16 → Transpose → graph output still reaches that helper through TryFixTransposeMissingQ. The explicit zero point bypasses its zero-point-less guard, so it inserts an unsupported MS QuantizeLinear and a new DQ without output_dtype, making the optimized graph invalid. Guard QDQ synthesis there too and add a direct-transpose regression with an explicit zero point; the new propagation tests do not cover this path.

4 resolved since last review

• Skip NVFP4 tests on devices below SM80
• Constrain CPU scale types after separating TS from T2
• Exclude unsupported FP8 dequantization cases from QDQ propagation
• Allow empty per-axis scales for empty input dimensions

_{Review effort: Balanced}

[Akshay Sonawane (apsonawane)]

Pushed 63026c2: increase only the fused NVFP4 SwiGLU raw K-packed reduction loop unroll hint from 2 to 4. Non-fused paths remain at unroll1; arithmetic expressions and dispatch are unchanged. Rejected FP8 scheduling experiments and benchmark artifacts are not included.

Optimization evidence from the isolated deployed-v26 build on H200 (GPU3): 40 byte-exact FP16/BF16 operator cases, 4,608 benchmark token IDs matching the reference across 8K/16K/32K, three default-EOS prompts matching, and no MTP failures/fallbacks. Fresh deployed throughput was 198.63 / 213.22 / 211.76 tok/s, approximately +0.52% / +0.60% / +0.59% versus fresh controls.

Important scope: those GPU results qualify the deployed-v26 pre-integration source/binary, not this exact PR head, which has newer launcher/template integration changes. This push ports only the same scheduling hint. PR-head lintrunner and git diff --check pass; no native PR-head rebuild was performed. CI must validate the current integration.

[Akshay Sonawane (apsonawane)]

Pushed 25d286b with the recent qualified INT2 runtime changes:

• Default-compatible 512-expert/top-10 packed execution and up to 20 expanded INT2 GEMV rows.
• Lifetime-safe sharing of runtime-packed INT2 GPU weights across sessions, with explicit opt-out.
• Explicit uniform INT2 weights_prepacked=1 contract: consume pool-owned SM80 CUTLASS W2_A16 bytes directly, preserving logical dimensions and ordinary block scales.
• Reject unsupported prepacked configurations instead of interpreting packed bytes through raw dequantization fallback.
• Schema/docs updates and regression coverage for 512-expert/top-10 parity and unsupported prepacked fallback.

The integration preserves this PR's newer FP8 execution paths. No unpromoted paired-column NVFP4 kernel, local adapter workaround, model binaries, generated data, or workspace packaging/profiling artifacts were included. The automatic NVFP4 packed-prefill default is updated separately in #33393, where that implementation lives.

Prior isolated H200 qualification (not a build/test claim for this exact PR head):

• Efficient raw INT2: about 201.59 / 214.35 / 212.90 decode tokens/s at 8K / 16K / 32512 contexts, versus about 173.94 / 183.78 / 183.01 before eligibility/sharing improvements.
• Single-prepacked INT2: 201.59 / 214.12 / 212.63 tokens/s, peak VRAM 90.17 -> 62.04 GiB, removing 28.125 GiB of retained raw-plus-packed duplication.
• 20 actual-weight INT2 operator cases and 40 existing regression cases matched exactly. Concurrent shared-initializer session lifetime/recreation checks passed.
• 13,824 measured benchmark tokens and all 164 HumanEval completions matched the previously qualified INT2 configuration. HumanEval remains 156/164 versus NVFP4 157/164 (0.64% relative pass-rate drop on that public benchmark, not a general coding guarantee).
• Latest compact-package reprofile at 16K: 214.25 tokens/s, 6.05 s TTFT, 62.04 GiB peak, exact reference tokens and healthy MTP.

Per request, no PR-local builds, tests, linters or CI watchers were run. The added repository tests are unexecuted locally; exact integrated-head validation remains for CI. The experiment provider used an older source snapshot/ABI, so prior GPU qualification must not be interpreted as qualifying this exact newer PR head.

[Akshay Sonawane (apsonawane)]

Added scoped FP8 and small-decode routing optimizations on top of the current PR head, preserving the native-FP8, decode-image, and packed-weight review fixes.

Published commit: 661351202b50d00864f4d93ba0496d7737640121.

• Opt-in ORT_FP8_GEMV_PER_COLUMN_S2=1: specialize weight-only two-row MMA projections with one scale per output column. Preserve existing MMA/K-split reduction, bias, split outputs, and unsupported-shape/QDQ fallbacks. No unrelated skinny-GEMM or PDL experiment was included.
• Fuse small-decode stable softmax/top-K with expert-map construction, including the declaration and actual packed-INT/FP4 host dispatch wiring.
• Opt-in ORT_MOE_ROUTING_WARP_SELECT=1: normalized 512-expert/top-K-10 selection for one or two rows using one warp per row and a 128-thread stable map builder. Other shapes retain fused-sort/generic fallback.
• Added exact FP16/BF16 FP8 outputs and routing dtype/tie/infinity/map/graph-replay coverage; documented both options.
• No expert-cache/offload, multi-draft runtime, or unrelated accumulated experiments were included.

Validation of the publishing sources: changed CUDA code and the actual MoE host caller compile, along with matching PR-tip FP8 host/GEMV dependencies. Process-isolated native suites pass with new switches both off and on: 20 FP8 matmul tests and 12 routing tests per mode. Validation uses retained build dependencies, not a clean full upstream build; CI remains required.

Across all three scoped kernel updates, the focused run passed 175 tests per mode (350 total, no failures). The broader GDN empty-request regression is not clean in the incremental build and also fails with its unmodified PR-head kernel; it is excluded from the focused GDN subset, not silently counted as passing.

Historical experimental-package routing profiling measured 7.93 -> 5.96 us, 256 -> 128 threads, and 8472 -> 80 shared bytes. Controlled 16K throughput improved by about 1.2% in both bracket orders; 8K did not show an order-independent gain. Earlier per-column FP8 brackets measured 218.12 -> 221.73 TPS at 8K and 236.86 -> 240.37 at 16K. These are not re-benchmarks of the transplanted PR tip. Full model accuracy remains deferred, with no sustained-250-TPS-at-8K claim.

Companion updates: GDN/convolution #33471, sparse attention/indexer #33472, and GenAI packaging microsoft/onnxruntime-genai#2759.

[Akshay Sonawane (apsonawane)]

Fresh-build integration validation update (no source changes pushed to this ORT PR in this update):

• Rebuilt ORT from scratch for H200 (sm90a), CUDA provider as a plugin, with FP4 QMoE, FP8/INT4 KV cache, GQA value layout, and FPA-INTB GEMM enabled. onnxruntime_QUICK_BUILD=OFF; no retained build objects were reused.
• Rebuilt GenAI against the fresh ORT headers/libraries, and rebuilt the fold custom library.
• Focused convolution/GDN/fused-tensor/hyper-connection/INT2 tests: 193 passed.
• Sparse/indexer/FP8 suite plus the isolated QSA capacity test: 140 additional passes.
• Known limitations: QSA and CSA MalformedStateLengthsAreClamped tests still fail, including with specialized top-k disabled. These are not being claimed as passing. This is selected-suite validation, not a clean full-suite qualification.
• Selected GenAI Qwen/executor/weight-pool tests: 69 passed, including the opt-in real-GPU mapped-weight test with proper CUDA plugin registration. One additional test is disabled.

The latest fused INT2 model completed one warmup plus three measured requests at each of 2K/8K/16K/32K, generating 512 tokens each. All 16 requests matched the previous-runtime token IDs exactly, with zero generation failures or standard fallbacks. Runtime-map inspection confirmed only the fresh ORT/GenAI libraries were used.

Prompt	Pooled decode TPS	Mean TTFT
2048	269.90	0.584 s
8192	250.17	2.448 s
16384	266.92	5.065 s
32768	264.12	10.578 s

Evidence scope: this is the full dirty experimental source stack containing this PR's optimizations plus other local work, not a standalone build/benchmark of this exact PR head. It is not an isolated speedup attribution, a fresh paired reference comparison, or proof of cross-process sustained 250 TPS. Full model accuracy qualification remains deferred.

The rebuild exposed local expert-cache function-pointer and INT2 direct-link test wiring issues, which were fixed locally. Those expert-cache extensions are not in the published PR and have deliberately not been mixed into this update.

[Akshay Sonawane (apsonawane)]

Pushed 8d2fc23: opt-in BF16 shared-expert FP8 projection/epilogue fusion, original-prefill numerical controls, and opt-in routing mode 2 using native warp reductions on SM80+. Existing routing modes/defaults and reviewer FP8 dispatch changes are preserved. Shared-expert helpers require companion #33470; tests also use its projected SigmoidGateMul. Scoped CUDA host/device, schema and test translation units compile against companion declarations. Combined-runtime qualification: 38 shared-expert native tests; 13 routing tests in each mode 0/1/2; 7 packed-decode integration tests in mode 2; 76 exact 512-token model-parity requests. These are not an independent branch build/link qualification. Routing mode 2 remains opt-in: combined bracketed gains +0.27% at 2K and +0.20% at 8K. Unrelated SkinnyGemm/prepacking experiments are excluded. Full LiveCodeBench accuracy evaluation is starting separately.

[Akshay Sonawane (apsonawane)]

Resolved conflicts with main (1f4b965) in merge commit fe74941.

• Preserved mainline FP4 PDL capability/routing/launch handling alongside this PR’s fused routing and NVFP4 fused finalizer.
• Corrected an automatic-merge interaction where adaptive block sizing overwrote the decode finalizer block size after its tiled grid was calculated. Decode grid and block are now chosen together after querying kernel caps.
• Compiled the merged native sources and focused tests, and linked an isolated CUDA provider/test executable against existing integration-build dependencies. The combined focused suite for these two merges ran on H200 GPU 1: 119 passed, 2 architecture/image-specific skips, 0 failures. Coverage includes fused routing, stable expert maps, fused NVFP4 finalization, sparse attention, and state-status handling.

This is merged-source validation using existing integration-build dependencies, not a full clean standalone-branch build or WebGPU qualification. GitHub currently reports MERGEABLE (no merge conflicts); required checks/reviews remain separate. The active worktrees and GPU 3 LiveCodeBench runtime were not modified.

[github-actions]

You can commit the suggested changes from lintrunner.

[Akshay Sonawane (apsonawane)]

Follow-up validation: ran lintrunner -a on an isolated checkout of this PR against origin/main, reviewed and applied the formatter fixes, then reran both lintrunner -a and lintrunner successfully with no lint issues. Pushed formatting-only commit 2c6795d22892798acc55e673f91e2c4a654762df (5 files). Verified identical C++/CUDA token sequences before and after formatting, including preserved string/comment contents. No runtime behavior changes; active integration worktree and evaluation snapshot untouched.

[Copilot]

🟡 Changes recommended

Unresolved numerical and optimizer/provider compatibility issues, plus gaps in automated coverage of new execution paths, block approval.

9 open issues

• Guard transpose optimization for runtime FP8 dequantization · New
• Preserve INT2 GEMV eligibility through 80 expanded rows · New
• Restore NVFP4 fused finalization lane threshold · New
• Cover FP8 projection segment GEMMs and overrides · New
• Enable per-column S2 specialization in tests · New
• Test GPU initializer cache reuse and lifetime · New
• Handle full-row NaNs in MoE router normalization · New
• Enable warp-selection modes in fresh-process tests · New
• Reject unsupported DirectML scale and output type combinations · New

1 resolved since last review

• Guard propagation tests when contrib operators are disabled

_{Review effort: Balanced}

[Akshay Sonawane (apsonawane)]

[P2] Fall back when the fused finalizer exceeds CUDA’s grid-Y limit

This support check accepts widths for which the new launch geometry is invalid. For example, num_rows=1, experts_per_token=1, N=524288, and K=256 pass the guards, but this expression produces grid.y=65536, exceeding CUDA’s 65535 limit. The previous four-K-lane FC2 kernel needs only 16384 blocks in Y for this shape, so selecting the fused path introduces an invalid-configuration failure rather than preserving the existing fallback.

Please calculate the block count without integer overflow and return false before launching when the grid is unsupported, or flatten/tile the column dispatch. Add a boundary regression covering an accepted width immediately below the limit and N=524288.

The newly column-tiled generic finalizer needs the same protection: for two rows, top-K 10, and FP16/BF16 cols=16777216, its grid-Y calculation also produces 65536. That kernel already has a column grid-stride loop, so capping its grid-Y dimension can preserve coverage.

[Copilot]

🟡 Changes recommended

Scratch-budget enforcement and cache-lock scope need correction, and native linking and GPU validation remain outstanding.

1 open issue

• Enforce scratch budget for dense-prefill full-buffer allocation · New

1 previously missed issue

In code that hasn't changed since last review

Avoid global mutex during shared INT2 packing

onnxruntime/​contrib_ops/​cuda/​moe/​moe_quantization.cc:3760

[P2] With shared INT2 packing enabled, this process-wide mutex remains locked through allocation, all per-expert transforms, and cudaStreamSynchronize at line 3889. Concurrent session initialization therefore serializes packing even for unrelated initializers on different GPUs. Limit the global lock to map access and synchronize in-progress packing per key, preserving duplicate-allocation prevention without blocking independent model loads.

9 resolved since last review

• Guard transpose optimization for runtime FP8 dequantization
• Reject unsupported DirectML scale and output type combinations
• Enable warp-selection modes in fresh-process tests
• Handle full-row NaNs in MoE router normalization
• Test GPU initializer cache reuse and lifetime
• Enable per-column S2 specialization in tests
• Cover FP8 projection segment GEMMs and overrides
• Restore NVFP4 fused finalization lane threshold
• Preserve INT2 GEMV eligibility through 80 expanded rows

_{Review effort: Balanced}