[CUDA] QMoE: fused int2 mixed-input GEMM + GeGLU activation + fractional zero-point - #32198
Draft
Thiago Pereira Rocha (thpereir) wants to merge 2 commits into
Draft
[CUDA] QMoE: fused int2 mixed-input GEMM + GeGLU activation + fractional zero-point#32198Thiago Pereira Rocha (thpereir) wants to merge 2 commits into
Thiago Pereira Rocha (thpereir) wants to merge 2 commits into
Conversation
Add 2-bit (uint2b_t) weight support to the QMoE CUDA EP as a fused mixed-input GEMM: weights stay packed in HBM and are dequantized in-register inside the SM80 DqMma pipeline, mirroring the int4 path (no full-precision weight materialization). Includes the in-register interleaved converter, W2_A16 weight preprocessor (bias+interleave, LDSM permutation), LayoutDetailsB<uint2b_t>, DqMma/DefaultMma wiring, runner instantiations, decode GEMV, and the QMoE op wiring that replaces the phase-1 dequant fallback. The key correctness fix: the warp_frag_B double-buffer index was derived from the per-CTA-K-stage load offset, which stays coherent only when kWarpGemmIterationsForB is even (int8=4, int4=2). For int2 it is odd (=1), so every stage re-read stage-0's B fragment (tile0 counted twice, tile1 dropped), producing an ORT[m]=ref[m]+ref[m+64] K-fold for K>64. Fixed with a persistent cross-stage parity in all four DqMma variants (multistage/pipelined x finegrained/percol); the change is a no-op for even kWarpGemmIterationsForB (int4/int8). Also add the GeGLU gated activation (gelu-tanh gate), for Gemma4 MoE. Validation (H100, fp16, hidden=inter=2048, E=8, top2, block=128): - 2-bit parity max_diff 0.000488/0.003906 (tolerance 0.35). - Full test_qmoe_cuda.py: 122 passed, 0 failed (no int4/int8/fp4/fp8 regression). - int2 decode 1.019 -> 0.037 ms (tracks int4 0.036); prefill 1.047 -> 0.101 ms (beats int4 0.161).
Add an optional ``zero_point_offset`` float attribute to the QMoE op for
integer block-wise quantization when no uint8 fc*_zero_points tensor is
provided. It selects an asymmetric dequant (code - zero_point_offset) *
scale using a single fractional center, enabling balanced schemes whose
zero-point is not integer-representable by a uint8 tensor -- e.g. a 2-bit
checkpoint quantized around the midpoint 1.5, giving codes {0,1,2,3} ->
{-1.5,-0.5,0.5,1.5}*scale.
Implementation: the fused int2 weight converter already subtracts the
symmetric center 2^(bits-1); a constant per-element bias =
(2^(bits-1) - zero_point_offset) * scale corrects it to the requested
fractional center. The bias is built at PrePack time from the packed
scales (PrePackConstantBiasFromScales) and consumed via the same slot
int4-asymmetric uses; ComputeInternal has a fallback that builds it from
the live scales when prepacking is disabled.
- contrib_defs.cc: new OPTIONAL_VALUE ``zero_point_offset`` FLOAT attr.
- qmoe_kernels.cu/.h: LaunchQMoEConstantBias (float/half/bf16) computing
bias = delta * scale.
- moe_quantization.cc/.h: attr parse + ORT_ENFORCE(int && block_size>0),
PrePack wiring for fc1/fc2, ComputeInternal fallback branch.
- test_qmoe_cuda.py: TestQMoEFractionalZeroPoint 2-bit case (offset 1.5).
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
Contributor
There was a problem hiding this comment.
Pull request overview
Adds fused int2 QMoE CUDA execution, GeGLU activation, and fractional zero-point support.
Changes:
- Adds int2 preprocessing, GEMM/GEMV dispatch, conversion, and pipeline handling.
- Adds tanh-approximated GeGLU activation.
- Adds fractional zero-point handling, quantization support, and CUDA tests.
Reviewed changes
Copilot reviewed 32 out of 32 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
onnxruntime/test/python/transformers/test_qmoe_cuda.py |
Adds int2, GeGLU, and fractional-ZP tests. |
onnxruntime/python/tools/quantization/cuda_quantizer.py |
Adds 2-bit packing and quantization. |
onnxruntime/core/graph/contrib_ops/contrib_defs.cc |
Extends QMoE/MoE schemas. |
onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h |
Declares constant-bias launchers. |
onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu |
Implements fractional-ZP bias kernels. |
onnxruntime/contrib_ops/cuda/moe/moe_quantization.h |
Adds fractional-ZP state and helpers. |
onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc |
Wires int2 runners, preprocessing, and bias handling. |
onnxruntime/contrib_ops/cuda/moe/moe_base.h |
Parses GeGLU activation. |
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu |
Extends MoE dispatch to int2. |
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv.h |
Documents int2 GEMV support. |
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv.cu |
Implements int2 GEMV dispatch. |
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_template_dispatch.h |
Enables int2 GEMM templates. |
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_kernels_fp16_uint2.cu |
Instantiates FP16/int2 GEMM. |
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_kernels_bf16_uint2.cu |
Instantiates BF16/int2 GEMM. |
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_activation_kernels.cuh |
Uses tanh GELU for GeGLU. |
onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/details.h |
Adds int2 GEMV conversion traits. |
onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors.h |
Defines W2 quantization type. |
onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors_impl.h |
Adds W2 layout and permutation metadata. |
onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors_impl.cu |
Implements int2 transpose/interleave preprocessing. |
onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_adaptor.h |
Declares int2 transpose adaptor. |
onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_adaptor.cu |
Implements packed int2 transposition. |
onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/interleaved_numeric_conversion.h |
Adds int2-to-FP16/BF16 converters. |
onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/dq_mma_pipelined_percol.h |
Fixes B-fragment parity. |
onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/dq_mma_pipelined_finegrained.h |
Fixes fine-grained pipeline parity. |
onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/dq_mma_multistage_percol.h |
Fixes multistage B-fragment parity. |
onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/dq_mma_multistage_finegrained.h |
Fixes fine-grained multistage parity. |
onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_mma.h |
Adds FP16/int2 MMA specializations. |
onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_mma_bf16.h |
Adds BF16/int2 MMA specializations. |
onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_dq_mma_pipelined.h |
Allows int2 pipelined dequant MMA. |
onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_dq_mma_multistage.h |
Allows int2 multistage dequant MMA. |
onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/kernel/moe_cutlass_kernel.h |
Requires scales for int2 kernels. |
onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/kernel/mixed_gemm_B_layout.h |
Defines the int2 mixed-GEMM layout. |
Suppressed comments (1)
onnxruntime/python/tools/quantization/cuda_quantizer.py:573
- This comment incorrectly says int2 has no CUTLASS mixed-input GEMM and uses a dequant fallback. The raw tensor is needed because W2 is transformed by QMoE's internal PrePack rather than this offline helper; update the explanation and error text accordingly.
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+215
to
+216
| ORT_ENFORCE(expert_weight_bits_ == 8 || expert_weight_bits_ == 4 || expert_weight_bits_ == 2, | ||
| "expert_weight_bits must be 2, 4, or 8, but got ", expert_weight_bits_); |
Comment on lines
+1534
to
+1538
| .Attr("zero_point_offset", | ||
| "Only meaningful when quant_type='int' and block_size > 0 and no integer " | ||
| "fc*_zero_points are provided. A single fractional zero-point applied uniformly to " | ||
| "every weight code: dequant = (code - zero_point_offset) * scale. Enables balanced " | ||
| "asymmetric schemes whose zero-point is not integer-representable (e.g. the 1.5 " |
Comment on lines
+44
to
+47
| } else if (activation_type_str == "geglu") { | ||
| // GeGLU: gelu-gated GLU (down(gelu_tanh(gate) * up)). Gated like SwiGLU (interleaved fc1), | ||
| // differing only in the gate nonlinearity (tanh-approx GELU instead of SiLU). Used by Gemma4. | ||
| activation_type_ = ActivationType::Geglu; |
Comment on lines
+1540
to
+1542
| "When omitted, symmetric quantization centered on 2^(expert_weight_bits-1) is used.", | ||
| AttributeProto::FLOAT, | ||
| OPTIONAL_VALUE) |
Comment on lines
+1364
to
+1368
| } else if (!std::isnan(zero_point_offset_) && block_size_ > 0 && scales && eff_scale) { | ||
| // Fractional zero-point center, no uint8 zero-point tensor, and PrePack did NOT consume the | ||
| // scales (e.g. session.disable_prepacking) -- the packed_bias slot above was not populated, so | ||
| // build the constant bias here from the live scales. The normal path builds it at PrePack time | ||
| // (PrePackConstantBiasFromScales) and takes the ``packed_bias`` branch above. The weight |
Comment on lines
+357
to
+360
| if int(bits) == 2: | ||
| # No CUTLASS mixed-input GEMM exists for 2-bit; the CUDA QMoE kernel consumes the raw | ||
| # [N, K/4] storage via its dequant fallback, so prepacking is unsupported. | ||
| raise ValueError("QMoE 2-bit weights cannot be CUTLASS-prepacked; use prepack=False (raw storage).") |
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 three related capabilities to the QMoE CUDA execution provider:
cutlass::uint2b_t) mixed-input GEMM. 2-bit weights stay packed in HBM and are dequantized in-register inside the SM80 DqMma pipeline, mirroring the existing int4 (uint4b_t) path — no full-precision weight materialization. This replaces the phase-1 dequant-to-fp16 fallback that materialized[E, N, K]fp16 weights every call.zero_point_offsetattribute for integer block-wise quant, enabling balanced asymmetric schemes whose zero-point is not integer-representable by a uint8 tensor (e.g. a 2-bit checkpoint quantized around the 1.5 midpoint: codes{0,1,2,3}→{-1.5,-0.5,0.5,1.5}*scale).Motivation
The prior 2-bit path fell back to dequantizing the full weight tensor to fp16 in HBM, forfeiting the entire point of 2-bit (memory bandwidth). It ran ~24x slower than int4 at decode. The fused kernel keeps weights packed and dequantizes in-register.
Changes
In-register converter / layout / DqMma wiring (patterned on the
uint4b_tsites):cutlass_extensions/interleaved_numeric_conversion.h: from-scratchFastInterleavedAndBiasedNumericArrayConverter<half_t/bfloat16_t, uint2b_t, N>(16 codes/word, symmetric +2 bias, fp16/bf16 magic).cutlass_extensions/gemm/kernel/mixed_gemm_B_layout.h:LayoutDetailsB<uint2b_t>(ThreadblockK=64,ColumnsInterleaved=8).gemm/threadblock/default_mma*.h,default_dq_mma_*.h, the fourdq_mma_*variants:uint2b_tspecializations + static_assert extensions.Correctness fix (the crux): the
warp_frag_Bdouble-buffer index was derived from the per-CTA-K-stage load offset, which stays coherent only whenkWarpGemmIterationsForBis even (int8=4, int4=2). For int2 it is odd (=1), so every stage re-read stage-0's B fragment (tile0 counted twice, tile1 dropped) — anORT[m]=ref[m]+ref[m+64]K-fold for K>64. Fixed with a persistent cross-stage parity carried in all four DqMma variants; the change is a no-op for evenkWarpGemmIterationsForB(int4/int8 unaffected).Weight preprocessor / runner / decode GEMV:
fpA_intB_gemm_preprocessors*:W2_A16quant type, 2-bit subbyte transpose,add_bias_and_interleave_int2s, LDSM permutation map.moe_gemm/moe_gemm_kernels_{fp16,bf16}_uint2.cu,moe_kernels.cu,moe_gemm_template_dispatch.h:uint2b_trunner instantiations.fpA_intB_gemv/details.h,moe_gemv.cu: int2 decode GEMV path.QMoE op wiring:
moe/moe_quantization.cc/.h: construct the fuseduint2b_trunner for 2-bit; re-enable int2 PrePack (W2_A16); addzero_point_offsetparse +PrePackConstantBiasFromScales(bias =(2^(bits-1) - zero_point_offset) * scale) and a ComputeInternal fallback when prepacking is disabled.moe/qmoe_kernels.cu/.h:LaunchQMoEConstantBias(float/half/bf16).core/graph/contrib_ops/contrib_defs.cc:zero_point_offsetOPTIONAL_VALUE FLOAT attr on the QMoE schema.Quantizer / tests:
python/tools/quantization/cuda_quantizer.py: 2-bit packing support.test/python/transformers/test_qmoe_cuda.py: fused int2/int4/int8 GeGLU coverage +TestQMoEFractionalZeroPoint(2-bit, offset 1.5).Testing
Validated on H100 (SM90), fp16, hidden=inter=2048, E=8, top_k=2, block_size=128:
test_qmoe_cuda.py: 122 passed, 15 skipped, 0 failed (no int4/int8/fp4/fp8 regression).Opening as draft for review.