Repository navigation
Add INT3 QMoE CUDA execution and versioned private prepack - #33555
David Fan (jiafatom) wants to merge 9 commits into
Conversation
|
Pushed review/CI fixes in ead612c:
New hosted Windows/CodeQL/format checks must still finish; local Linux validation is not a Windows CI pass. No dependency pins or unrelated files changed. |
There was a problem hiding this comment.
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
7 open issues
Scratch buffers are allocated with GetComputeStream(context) but then initialized/used with stream… · New Scratch buffers are allocated with GetComputeStream(context) but then initialized/used with stream… · New Scratch buffers are allocated with GetComputeStream(context) but then initialized/used with stream… · New Weight() unconditionally dereferences projection.scales. If params.fc{1,2}.scales is ever null… · New PrePackInt3 calls cudaGetDeviceProperties for every prepacked initializer. This is relatively… · New PrePackInt3 launches packing on the legacy default stream (nullptr) and then synchronizes that… · New This call validates scales by passing rows=scales.shape.Size() and logical_k=block_size_. The… · New
What changed in this PR
Adds CUDA execution support for INT3 QMoE along with provider-private (versioned) packed-weight prepacking and a new CUDA-focused correctness test suite.
Changes:
- Introduces INT3 CUDA kernels (including versioned bit-plane packed weights) and validation helpers.
- Extends QMoE CUDA kernel to route INT3 paths via private prepack + runtime validation/caching.
- Adds integrated CUDA GTest coverage for INT3 QMoE correctness/regressions.
| File | Description |
|---|---|
| onnxruntime/test/contrib_ops/qmoe_int3_cuda_test.cc | Adds formal CUDA OpTester-based correctness and lifecycle regression tests for INT3 QMoE. |
| onnxruntime/contrib_ops/cuda/moe/qmoe_int3.h | Declares INT3 MoE kernel params and CUDA entrypoints (launch/pack/validate). |
| onnxruntime/contrib_ops/cuda/moe/qmoe_int3.cu | Implements INT3 fused CUDA kernels, bit-plane packing, and device-side content validation. |
| onnxruntime/contrib_ops/cuda/moe/qmoe_int3.cc | Implements QMoE INT3 compute path + provider-private prepack and cache identity checks. |
| onnxruntime/contrib_ops/cuda/moe/packed_code.cuh | Adds reusable helper for reading row-local packed codes on device. |
| onnxruntime/contrib_ops/cuda/moe/moe_quantization.h | Adds INT3-specific members (packed input cache) and helper method declarations to QMoE. |
| onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc | Enables INT3 bit-widths, routes Compute/PrePack to INT3-specific implementations, and adds INT3 constraints. |
Review effort: Lite
|
Pushed abcb535 addressing all seven latest review threads and Windows CI. Windows CUDA, plugin and TensorRT logs reported MSVC C4334/C2220 in correctness-test bit shifts; the operands now explicitly use int64_t. This is separate from the previously fixed launch-grid C4267. Prepack now uses a dedicated nonblocking stream and cached device properties; Compute explicitly shares its stream object/handle across scratch and kernels; scale validation uses its real shape; launcher rejects missing required pointers. Local CUDA13 provider/test rebuild and final rerun: 9 passed, zero failures/skips. Compiled missing-scale launch checks, pinned clang-format and whitespace checks also pass. Hosted Windows/MSVC and the new CI run remain to be confirmed. |
|
CI follow-up: pushed ffe0578 to fix the shared Windows CUDA / TensorRT / CUDA plugin build failure (MSVC C4244 in PackInt3TestCode). The test helper now keeps row/count/index/code and offset arithmetic in int64_t, with an explicit uint8_t conversion for the packed byte. Validation: rebuilt onnxruntime_provider_test; all 9 QMoEInt3CudaTest cases passed with zero skips; the actual helper also passed -Wall -Wextra -Wconversion -Werror, clang-format 20 dry-run, and git diff --check. Native MSVC validation is pending CI. The other three failures on the preceding head were bootstrap failures: Python format and Suggest fixes failed during Rust setup before lint ran; Java CodeQL failed while resolving Microsoft JDK 17 with connection timeouts/unreachable network. The new push triggers fresh checks; these failures do not yet report source findings. |
|
CI follow-up: fixed the remaining two MSVC C4334 diagnostics in the zero-point generation calls by shifting int64_t{1} rather than a 32-bit literal. These were introduced as warning sites when the packing helper code argument became int64_t. Audited all shift expressions in this test file. Local rebuild and all 9 formal QMoEInt3CudaTest cases passed (zero skips); clang-format and git diff --check passed. Native MSVC results are pending the new CI run. Separately, the wasm Debug failure was a 30-minute browser silence timeout after InferenceSessionTests.SessionLoggerOutlivesEPsOnDestruction started, not a compile error or a reported assertion; requested a targeted rerun of job 120689610465 to check reproducibility without changing that unrelated test. |
Signed-off-by: David Fan <jiafa@microsoft.com>
Signed-off-by: David Fan <jiafa@microsoft.com>
Signed-off-by: David Fan <jiafa@microsoft.com>
Signed-off-by: David Fan <jiafa@microsoft.com>
Signed-off-by: David Fan <jiafa@microsoft.com>
Signed-off-by: David Fan <jiafa@microsoft.com>
Signed-off-by: David Fan <jiafa@microsoft.com>
Signed-off-by: David Fan <jiafa@microsoft.com>
Signed-off-by: David Fan <jiafa@microsoft.com>
|
Rebased this branch onto freshly fetched origin/main (36f07f7). Rebase completed without conflicts, and git range-diff confirms all nine branch patches are unchanged. Updated head: d75ac16. Validation after rebase: rebuilt CUDA provider and provider tests; all 9 QMoEInt3CudaTest cases passed with zero skips; plugin-mode syntax compilation of qmoe_int3.cc passed; git diff --check passed. Pushed with an exact force-with-lease to avoid overwriting concurrent branch changes. |
870a8ef to
d75ac16
Compare
|
microsoft/olive-recipes#679 shows that #33569 has better performance, I just merged that PR, so close this one. |


Summary
Combines PR2 (CUDA correctness) and PR3 (packed decode/provider-private prepack) from the INT3 delivery plan, tracked by #32657. Builds on merged #33394 and reuses the correctness baseline from #33519; this is the new combined implementation branch, not an assertion that #33519 itself supplies private prepack.
Validation
Fresh ORT provider/test objects built with consistent CUDA 13 in
jiafa-dev, A100:QMoEInt3CudaTest.*: 9 passed, 0 failed, 0 skipped, rerun after final runtime formatting.Build qualification boundary: dependency sources were reused but ORT/dependency objects were freshly compiled. Local ONNX source was
ee3ccbd2b2344299d3a4506c2954a47b2181a485, versus the repository pine1ac8e74b620bbdc7ba07f0440806bd93f651a28. No gitlink or dependency pin changes are included. Default-pin CI qualification remains open; this PR does not need a special ONNX commit. Concurrency, graph capture and every-GPU qualification are not claimed.Performance Analysis
A100, H2048/I768/E128/k8, rows1, block64; six width/type cases, 72 fresh Sessions and 2,160 timed calls, with exact finite nonzero output agreement. Constant/private-cache steady medians 0.469-0.486 ms versus all-dynamic canonical 0.660-0.736 ms, a 1.37-1.54x ratio. Validation/cache differences mean this is not an isolated universal layout speedup. Kernel-only raw results separately measure layout and conversion cost.
Prepack-inclusive initialization costs 649-892 ms. Sampled resident GPU increments were 774-834 MiB cached versus 264-312 MiB dynamic, including feeds and allocator retention, not peak memory. These are material setup/memory trade-offs. No routing speedup, full-model improvement or PR4 optimized-GEMM acceptance is claimed.
Experiment Archive
Reports and raw measurements are not part of this implementation diff. They remain accessible through the immutable snapshot
ead612cb60125b34bd61c4ed24c476fb9b8898cb(not the moving branch):Perf/benchmark and standalone smoke sources have also been removed from this diff. Reproduction harnesses remain accessible in the immutable experiment source snapshot, with standalone CUDA experiment sources. The archived report generator supports external data paths and optional
--kernel-results; generated results belong outside the source tree. Local compressed archive SHA256:6e49de22993370b8545d71225b3166b5089e4576f1f645bfe4d58a178ac60c1b.Scope And Follow-ups
Draft pending default-dependency CI and review of cache/prepack costs. Optimized GEMM prefill (PR4), CPU parity, producer merges (Olive microsoft/Olive#2749; Mobius onnxruntime/mobius#783), full-model quality/performance qualification, and padding-layout experiments (#33530) remain separate. This is not ONNX standard approval or production performance qualification.
Correctness Test Organization
The implementation diff adds only
onnxruntime/test/contrib_ops/qmoe_int3_cuda_test.ccunder test. It uses the existing provider-test GTest/OpTester infrastructure, analogous to INT2 QMoE cases inmoe_test.cc; existing INT2 grouped-kernel coverage is incuda_kernels/moe_gemm_int2_test.cc. The INT3 file is separate to keep its oracle and lifecycle helpers local. Latest local rerun after removing experimental files: 9 formal CUDA tests passed, zero skips; 223 existing reference tests passed. No new performance measurements are claimed by this cleanup.