Add CudaQuantizer to onnxruntime.quantization - #29509
Merged
Merged
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR brings a CUDA-only QMoE weight quantizer into ONNX Runtime’s Python quantization tooling and exposes it as onnxruntime.quantization.MoeCudaQuantizer, replacing an out-of-repo helper previously loaded by absolute path in the QMoE CUDA tests.
Changes:
- Added
MoeCudaQuantizerimplementation with per-channel symmetric INT4/INT8 storage and optional CUDA mixed-GEMM prepacking. - Exported
MoeCudaQuantizerviaonnxruntime.quantizationpackage API. - Updated
test_qmoe_cuda.pyto useMoeCudaQuantizerfor per-channel quantization/prepack and added smoke coverage for storage contract and default-prepacked execution.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| onnxruntime/test/python/transformers/test_qmoe_cuda.py | Switches per-channel quantization/prepack to MoeCudaQuantizer and adds additional smoke tests (including default-prepacked layout). |
| onnxruntime/python/tools/quantization/qmoe_quantizer.py | Introduces the MoeCudaQuantizer implementation with lazy torch import and optional CUDA prepacking via _pybind_state. |
| onnxruntime/python/tools/quantization/init.py | Exports MoeCudaQuantizer from the public quantization package. |
tianleiwu
changed the title
tianleiwu
force-pushed
the
tlwu/qmoe_quantizer
branch
from
July 5, 2026 22:57
439fc78 to
2414fa4
Compare
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Restructure symmetric_per_channel_quantize so qmin/qmax/scale_divisor/ zero_point are unconditionally initialized (if bits==4 / else), matching the earlier bits validation. Silences CodeQL 'potentially uninitialized local variable' alerts on cuda_quantizer.py.
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
kunal-vaishnavi
approved these changes
Jul 6, 2026
kunal-vaishnavi
left a comment
kunal-vaishnavi
left a comment
Contributor
There was a problem hiding this comment.
Thanks for adding this support! Can we also update the model builder to use this quantizer?
tianleiwu
enabled auto-merge (squash)
July 6, 2026 21:29
Contributor
Author
I'm working on that. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Adds a CUDA-focused
CudaQuantizerhelper to ONNX Runtime's Python quantization tools and exports it fromonnxruntime.quantization. The helper centralizes the CUDA quantization paths needed by ONNX Runtime tests andonnxruntime-genaimodel building, covering both per-channel QMoE weight packing and MatMulNBits blockwise formats without duplicating that logic in test code.This also updates the CUDA MoE/QMoE transformer tests to consume the shared helper, including full-range symmetric offset storage and CUTLASS mixed-GEMM prepacked layouts.
Summary of Changes
Quantization API
onnxruntime/python/tools/quantization/cuda_quantizer.pyCudaQuantizerwith per-channel INT4/INT8 quantization, CUDA mixed-GEMM prepacking, MatMulNBits blockwise quantization, CUTLASS prepacked blockwise quantization, and a pure-PyTorch symmetric blockwise reference path.onnxruntime/python/tools/quantization/__init__.pyCudaQuantizerfromonnxruntime.quantization.onnxruntime/python/tools/quantization/qmoe_quantizer.pyTest Updates
onnxruntime/test/python/transformers/test_moe_cuda.pyCudaQuantizercalls.onnxruntime/test/python/transformers/test_qmoe_cuda.pyCudaQuantizerfor per-channel and blockwise QMoE test setup, including full-range unsigned offset storage coverage and default prepacked layout smoke tests.Testing
pytest transformers/test_qmoe_cuda.py::TestQMoEIntPrePackSmoke- 5 passed, 2 subtests passed.pytest transformers/test_qmoe_cuda.py- 97 passed, 12 skipped, 2 subtests passed.torch, and that torch-backed methods raise a clear error only when used.python -m py_compile onnxruntime/python/tools/quantization/cuda_quantizer.pyChecklist