onnxruntime
0c67e77b - Extend CUDA QMoE packed INT2 prefill block sizes and weight combinations (#33092)

Commit
34 days ago
Extend CUDA QMoE packed INT2 prefill block sizes and weight combinations (#33092) ## 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. The final changes adding `(2,2)` / `(4,2)` and the expanded test matrix have **not been CUDA-built or numerically executed**. The direct grouped GEMM tests also have not been executed; earlier filters matching zero tests are not counted as validation. A larger internal-test build was stopped at the requester's direction. Local diff/format checks passed. This PR is ready for review, but the final head has no local CUDA execution evidence for the newly added combinations: the earlier passing `(2,4)` results do not establish correctness of `(2,2)` / `(4,2)`, and no performance improvement is claimed. Local CUDA build and test results for the final head are still outstanding and are requested before merge; they will be posted on this PR. Changes since the first revision, also not executed locally: - 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: ```bash 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.
Author
Parents
Loading