Skip to content

Add INT3 QMoE CUDA execution and versioned private prepack - #33555

Closed
David Fan (jiafatom) wants to merge 9 commits into
mainfrom
jiafa/int3-qmoe-pr2-pr3
Closed

David Fan (jiafatom) wants to merge 9 commits into
mainfrom
jiafa/int3-qmoe-pr2-pr3

Conversation

@jiafatom

@jiafatom David Fan (jiafatom) commented Oct 23, 2026 •

Copy link
Copy Markdown
Contributor

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.

  • Supports FC1/FC2 widths (3,3), (3,4), (4,3), FP16/BF16, block32/64/128; fused interleaved SwiGLU with routing/bias support and bounded direct-packed prefill.
  • Adds version-1 provider-private K32 bit-plane weights: INT3 remains 12 bytes per 32 values, not INT4-expanded. Canonical model bytes, QMoE v1 schema and ONNX dependency are unchanged.
  • Kernel-owned caches record shape/type/width/block/layout/device identity. Initializer scales and INT3 zero points are validated once, consumed initializers are released, and dynamic inputs remain validated on every call. No global cross-session cache or exported private bytes.
  • Bounds routing tile offsets including expert count and avoids per-run device-property queries.
  • Includes integrated formal correctness regressions, following the existing INT2 QMoE OpTester approach. Standalone experimental tests and perf/benchmark harnesses are archived, not included in the implementation diff. Experiment reports and raw results are archived outside this PR diff via immutable commit links below; historical full-model evidence remains available in INT3 QMoE CUDA correctness baseline and Qwen experimental results #33519.

Validation

Fresh ORT provider/test objects built with consistent CUDA 13 in jiafa-dev, A100:

  • Formal QMoEInt3CudaTest.*: 9 passed, 0 failed, 0 skipped, rerun after final runtime formatting.
  • Covers constant/dynamic matrices and 108 valid partial-prepack combinations across blocks 32/64/128, all three width pairs, six independently dynamic inputs and both dtypes; malformed content; same-Session valid/invalid/valid mutations; initializer release and Session recreation; empty rows, malformed scale shapes and exact scratch limits.
  • Historical archived standalone checks: packed loader 80 cases passed; dual projection 216 cases passed, each repeated three times, with zero Compute Sanitizer errors. These experimental entry points are no longer in this diff; current CI coverage uses the formal GTest file.
  • Ruff, Python compilation, report link/JSON parsing and git whitespace checks passed.

Build qualification boundary: dependency sources were reused but ORT/dependency objects were freshly compiled. Local ONNX source was ee3ccbd2b2344299d3a4506c2954a47b2181a485, versus the repository pin e1ac8e74b620bbdc7ba07f0440806bd93f651a28. 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.cc under test. It uses the existing provider-test GTest/OpTester infrastructure, analogous to INT2 QMoE cases in moe_test.cc; existing INT2 grouped-kernel coverage is in cuda_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.

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You can commit the suggested changes from lintrunner.

@jiafatom

Copy link
Copy Markdown
Contributor Author

Pushed review/CI fixes in ead612c:

  • Windows CUDA/plugin/TensorRT/docs build failures were MSVC C4267 treated as errors: size_t launch-grid counts implicitly narrowed to unsigned int in qmoe_int3.cu. Explicit checked-by-existing-tile-bounds block-count types now avoid that conversion.
  • Applied pinned clang-format to session_perf.cc and fixed both mixed ONNX import CodeQL findings.
  • Local validation: current-source provider/test rebuild; 9 CUDA tests passed, 0 skips; 223 Python reference tests passed; 18 matched models generated/checked; performance harness C++ compilation and Ruff/format/whitespace checks passed.

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.

@jiafatom
David Fan (jiafatom) marked this pull request as ready for review October 23, 2026 17:17
Copilot AI balanced review requested due to automatic review settings October 23, 2026 17:17

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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
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

@jiafatom

Copy link
Copy Markdown
Contributor Author

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.

@jiafatom

Copy link
Copy Markdown
Contributor Author

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.

@jiafatom

Copy link
Copy Markdown
Contributor Author

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>
@jiafatom

Copy link
Copy Markdown
Contributor Author

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.

@jiafatom

Copy link
Copy Markdown
Contributor Author

microsoft/olive-recipes#679 shows that #33569 has better performance, I just merged that PR, so close this one.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants