Extend CUDA QMoE packed INT2 prefill block sizes and weight combinations

[David Fan (jiafatom)]

Summary

Extend CUDA QMoE packed INT prefill on SM80+ from block64 with FC1/FC2 bits
(2,4) to block sizes 32, 64, 128 and bit-width combinations
(2,4), (2,2), (4,2).

• Pass the actual block size and per-projection bit widths through the prefill parameters.
• Dispatch each projection to the existing INT2 grouped GEMM or INT4 grouped GEMM runner.
• Keep (4,4) on its existing grouped MoE path.
• Extend INT2 grouped GEMM support checks and QMoE weight prepacking to the supported block sizes.
• Retain SM80+, symmetric quantization, interleaved fused SwiGLU, K/N alignment, expert-count and workspace restrictions. This does not enable arbitrary block sizes, zero points, INT8 combinations or new GEMV configurations.

Tests

• Parameterize direct FP16/BF16 INT2 grouped GEMM numerical tests across block32/64/128 and both row-tile shapes, including empty experts and partial tiles.
• Parameterize QMoE numerical integration coverage across all three bit-width combinations and all three block sizes, FP16/BF16, rows 1/33/257, initializer/runtime scales, and untiled/tiled execution.
• QMoE integration cases disable GEMV and set dense expert-dequantization scratch to 1 byte, so they must exercise packed prefill rather than decode or dense fallback.
• Preserve unsupported block16 and INT8-combination fallback coverage.

Validation Status

Before adding (2,2) / (4,2) dispatch, the (2,4) implementation with
block32/64/128 was incrementally built and passed the targeted block-size
integration test on A100 (SM80). The neighboring packed-prefill integration
suite passed 11 tests, including FP16/BF16, runtime scales, routing, row tiling,
workspace limits, disabled-path and expected-fallback tests.

Final head db8a7625b8baf3a5a9cee477fddc957933ada7fc was built and numerically
tested locally on October 5 using CUDA 12.8, Release, ENABLE_BF16 and
onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS=ON, on GPU1 (A100-SXM4-80GB, SM80).
The non-plugin CUDA provider, internal-test module and provider-test executable
completed compilation and linking using the existing CMake build directory.

• Discovery found 14 QMoE integration/regression tests and CUDA_EP_Unittest.All.
• Outer run: 15 passed, 0 failed, 0 skipped (14 QMoE tests plus the wrapper), 82.038 seconds.
• Internal module: 26 passed, 0 failed, 0 skipped (22 grouped GEMM tests plus 4 validation tests), 56.151 seconds. Three benchmark tests remained disabled.
• The packed-prefill matrix passed all 216 invocations: block32/64/128, (2,4)/(2,2)/(4,2), FP16/BF16, rows 1/33/257, initializer/runtime scales and untiled/tiled execution.
• Both direct FP16/BF16 block-size tests passed across block32/64/128 and tile rows 32/64, including empty experts and partial tiles.
• Block32 dense fallback with prefill disabled and the decode-fallback regression both passed, as did cached-scale decode-to-prefill and neighboring packed-prefill tests.

These are actual executions, not zero-match filters or discovery-only results.
This validates the non-plugin SM80 configuration, not plugin builds, other GPU
architectures, sanitizer coverage or performance improvement. The detached build
completed and linked the executable; a subsequent no-op check unexpectedly
restarted compilation after Ninja reported an incomplete build record. That
redundant rebuild was stopped, and the completed final-head binaries were tested.

Changes since the first revision, covered by the final-head run above:

• The existing block32 (2,2) decode-fallback test now disables packed prefill so it keeps testing decode fallback.
• Block size 32 is prepacked only for packed prefill; the GEMV-only prepack branch still accepts 64 and 128.
• Added QMoETest_MixedWidthCudaBlock32DenseFallbackWithoutPrefill (block32, prefill disabled, dense fallback must succeed).
• Prefill parameter defaults for block size and FC1/FC2 widths are now 0 (unset) instead of block64 / (2,4).

Suggested CUDA Validation

Build with onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS=ON, including
onnxruntime_providers_cuda_ut and onnxruntime_provider_test. Run the
provider-test executable with the CUDA internal module available:

GTEST_FILTER='Int2GroupedGemmTest.*:Int2GroupedGemmValidationTest.*' \
./onnxruntime_provider_test --gtest_filter='CUDA_EP_Unittest.All:MoETest.QMoETest_MixedWidthCudaPackedPrefill*:MoETest.QMoETest_Int2CudaCachedScalesDecodeThenPrefill:MoETest.QMoETest_MixedWidthCudaBlock32DenseFallbackWithoutPrefill:MoETest.QMoETest_Int2CudaPackedDecodeFallback'

In the non-plugin build, the command-line filter selects the outer wrapper and
integration tests; the environment filter selects the direct tests registered
inside the CUDA internal-test module. Verify that the direct tests are discovered
and that nonzero, non-skipped direct-test counts are reported. Listing tests alone
does not establish numerical correctness.

No performance improvement is claimed by this PR.

[github-actions]

You can commit the suggested changes from lintrunner.

[Copilot]

Copilot review overview

🟡 Changes recommended

Enabling block32 prefill makes an existing SM80 fallback test unexpectedly succeed.

Review effort: Balanced
Findings: 1

Open (1)

• Update Int2 CUDA packed decode fallback test for block32 behavior · New

What changed in this PR

Extends CUDA QMoE packed prefill to additional block sizes and INT2/INT4 projection combinations.

Changes:

• Dispatches each projection according to its bit width.
• Supports block sizes 32, 64, and 128.
• Expands grouped-GEMM and QMoE test coverage.

File	Description
onnxruntime/​test/​contrib_ops/​moe_test.cc	Expands QMoE integration cases and fallbacks.
onnxruntime/​test/​contrib_ops/​cuda_kernels/​moe_gemm_int2_test.cc	Tests grouped GEMM across block sizes.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_quantization.cc	Broadens prefill gating and propagates parameters.
onnxruntime/​contrib_ops/​cuda/​llm/​moe_gemm/​moe_int2_prefill.h	Adds block-size and bit-width parameters.
onnxruntime/​contrib_ops/​cuda/​llm/​moe_gemm/​moe_int2_prefill.cu	Dispatches INT2 or INT4 grouped GEMM per projection.
onnxruntime/​contrib_ops/​cuda/​llm/​moe_gemm/​moe_gemm_kernels_fp16_uint2.cu	Extends INT2 kernel support validation.

---

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

[Copilot]

Copilot review overview

🔵 Needs a closer look

Block-128 prefill can route valid partial-block QMoE shapes into an INT2 kernel that rejects them.

Review effort: Balanced
Findings: None

Resolved since last review (1)

• Update Int2 CUDA packed decode fallback test for block32 behavior

Previously missed (1)

In code that hasn't changed since last review

Validate reduction divisibility before packed-prefill buffer creation

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

Block size 128 is now admitted here, but the packed-prefill branch below still accepts any reduction size divisible by 64. QMoE's input validation allows a partial final quantization block (ceiling division), so a valid (2,4) model with hidden_size=64 and block_size=128 gets prepacked and selects this path, then RunInt2GroupedGemm rejects it because reduction_size % block_size != 0 instead of using the dense fallback. Require each packed-prefill reduction dimension to be divisible by the selected block size before creating its packed buffer.

[Akshay Sonawane (apsonawane)]

Before merging, please compile the final head with onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS=ON and execute the direct INT2 grouped-GEMM tests and expanded packed-prefill integration matrix on SM80+. The previously passing (2,4) results do not validate the newly added (2,2) and (4,2) dispatch paths. Please include the test results and confirm that the filters discover and execute the intended tests, rather than matching zero tests.

[David Fan (jiafatom)]

Re the Copilot note on reduction divisibility at moe_quantization.cc:2559: no code change, because the described route is not reachable. For every INT2 or mixed-width QMoE (use_int_dequant_fallback is true for all configurations this PR routes to packed prefill), ComputeInternal rejects hidden_size % block_size != 0 and inter_size % block_size != 0 at moe_quantization.cc:839-845, before the packed-prefill decision at line 1147. hidden_size and inter_size are exactly the FC1 and FC2 reduction dimensions, so a partial final block (for example hidden_size=64 with block_size=128) returns that explicit error and never reaches RunInt2GroupedGemm. The dense fallback has the same divisibility requirement, so this shape failed the same way before this PR; it is not a regression. The kernel-level reduction_size % block_size check in IsInt2GroupedGemmSupportedImpl remains as defense in depth. Not tested locally (source inspection only).

[Tianlei Wu (tianleiwu)]

Reviewed db8a762. I did not identify a new kernel correctness defect; the earlier decode-fallback and block32 prepacking concerns are addressed in this head. The numerical matrix forces packed prefill rather than allowing GEMV/dense fallback to mask it.

One nonblocking correction to Suggested CUDA Validation in the PR description: in the non-plugin build, Int2GroupedGemmTest.* is registered in onnxruntime_providers_cuda_ut, not directly in onnxruntime_provider_test. The outer executable invokes that module only via CUDA_EP_Unittest.All (onnxruntime/test/providers/cuda/cuda_provider_test.cc:14). The module's TestAll() initializes its own Google Test instance without forwarding the outer command-line filter (onnxruntime/test/providers/cuda/test_cases/cuda_test_provider.cc:113). Consequently, the proposed combined filter runs the integration cases but does not invoke the direct INT2 block-size/tile-shape tests. An existing local non-plugin binary confirms the registration split: listing Int2GroupedGemmTest.*:CUDA_EP_Unittest.All exposes only the wrapper, and filtering to Int2GroupedGemmTest.* executes zero tests.

Please include the wrapper in the suggested command. To keep its internal run focused and include the new block32 dense-fallback regression:

GTEST_FILTER='Int2GroupedGemmTest.*:Int2GroupedGemmValidationTest.*' \
./onnxruntime_provider_test --gtest_filter='CUDA_EP_Unittest.All:MoETest.QMoETest_MixedWidthCudaPackedPrefill*:MoETest.QMoETest_Int2CudaCachedScalesDecodeThenPrefill:MoETest.QMoETest_MixedWidthCudaBlock32DenseFallbackWithoutPrefill:MoETest.QMoETest_Int2CudaPackedDecodeFallback'

The explicit command-line filter selects the outer wrapper/integration tests; the environment filter selects the module's direct tests. Verify nonzero, non-skipped direct-test counts. I have not CUDA-built or numerically executed this head; the local test-discovery check is not numerical validation. CUDA CI was still running when inspected, so current-head build and numerical results remain required before merge.

[David Fan (jiafatom)]

Thanks for identifying the registration split. Updated Suggested CUDA Validation to use your two-level filter: CUDA_EP_Unittest.All and the integration/regression cases in the outer command-line filter, and Int2GroupedGemmTest.* plus Int2GroupedGemmValidationTest.* via GTEST_FILTER for the internal module. The final-head local build is already underway with internal tests enabled; I will verify nonzero, non-skipped direct-test counts and report actual numerical results, not just discovery output.

[David Fan (jiafatom)]

Final-head local CUDA validation is complete for db8a762. CUDA 12.8, Release, ENABLE_BF16, onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS=ON, non-plugin build, GPU1 NVIDIA A100-SXM4-80GB (SM80). Built onnxruntime_providers_cuda, onnxruntime_providers_cuda_ut and onnxruntime_provider_test with cmake --build build/cuda-wheel-clean-20261002/Release --target onnxruntime_providers_cuda onnxruntime_providers_cuda_ut onnxruntime_provider_test --parallel 4. Discovery listed 14 QMoE tests plus CUDA_EP_Unittest.All. Executed the two-level filter now shown in Suggested CUDA Validation: outer run 15/15 passed (14 QMoE tests + wrapper), zero failures/skips, 82.038 s; internal module 26/26 passed (22 grouped GEMM + 4 validation), zero failures/skips, 56.151 s, 3 benchmarks disabled. PackedPrefillBlockSizes passed all 216 invocations covering (2,4)/(2,2)/(4,2), block32/64/128, FP16/BF16, rows 1/33/257, runtime/initializer scales and untiled/tiled execution. Direct FP16/BF16 block-size/tile-shape tests passed, along with the block32 dense-fallback and decode-fallback regressions and cached-scale decode-to-prefill. These are nonzero actual numerical executions, not listing-only checks. A disconnected build finished linking; a subsequent no-op check unexpectedly restarted compilation because Ninja reported an incomplete build record, so that redundant rebuild was stopped and the completed binaries were tested. PR description updated. No plugin/other-architecture/sanitizer/performance claim. Complete local output retained at /tmp/pr33092-final-head-cuda-tests.log.