Add dequantize, 3D from_hp, transpose, and _grouped_mm support to Int4 tensors - #4731
Add dequantize, 3D from_hp, transpose, and _grouped_mm support to Int4 tensors#4731eryk-roch wants to merge 1 commit into
Conversation
…4 tensors Add support for MoE (Mixture of Experts) models with int4 weight-only quantization by enabling 3D expert weight stacks and grouped_mm dispatch. MoE models like Gemma-4 and Qwen3-Coder-Next store expert weights as 3D stacked tensors (e.g. [num_experts, N, K]) and use torch._grouped_mm for the forward pass. Previously, Int4PlainInt32Tensor only supported 2D tensors and neither Int4 tensor class implemented aten._grouped_mm.default, so these expert weights could not be quantized. Changes to Int4PlainInt32Tensor (XPU/NPU): - Add dequantize() using identity-matrix trick with existing int4pack kernel - Add 3D support to from_hp() via per-slice quantization - Add aten.transpose.int handler (logical-only, packed data stays in place) - Add aten._grouped_mm.default handler (dequant-then-native fallback) Changes to Int4Tensor (CUDA/mslk): - Add dequantize() with manual int4 unpack and affine dequantization - Add aten._grouped_mm.default handler (dequant-then-native fallback) Tests added for both tensor classes: dequantize roundtrip, 3D from_hp, and grouped_mm end-to-end with quantize_ API.
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/4731
Note: Links to docs will display an error until the docs builds have been completed. This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
Hi @eryk-roch! Thank you for your pull request and welcome to our community. Action RequiredIn order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you. ProcessIn order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA. Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks! |
1 participant
Add support for MoE (Mixture of Experts) models with int4 weight-only quantization by enabling 3D expert weight stacks and grouped_mm dispatch.
MoE models like Gemma-4 and Qwen3-Coder-Next store expert weights as 3D stacked tensors (e.g. [num_experts, N, K]) and use torch._grouped_mm for the forward pass. Previously, Int4PlainInt32Tensor only supported 2D tensors and neither Int4 tensor class implemented aten._grouped_mm.default, so these expert weights could not be quantized.
Changes to Int4PlainInt32Tensor (XPU/NPU):
Changes to Int4Tensor (CUDA/mslk):
Tests added for both tensor classes: dequantize roundtrip, 3D from_hp, and grouped_mm end-to-end with quantize_ API.