[CUDA] QMoE: fused int2 mixed-input GEMM + GeGLU activation + fractional zero-point

[Thiago Pereira Rocha (thpereir)]

Description

Adds three related capabilities to the QMoE CUDA execution provider:

1. Fused int2 (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.
2. GeGLU gated activation (gelu-tanh gate), for Gemma-style MoE blocks.
3. Optional fractional zero_point_offset attribute 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_t sites):

• cutlass_extensions/interleaved_numeric_conversion.h: from-scratch FastInterleavedAndBiasedNumericArrayConverter<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 four dq_mma_* variants: uint2b_t specializations + static_assert extensions.

Correctness fix (the crux): 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) — an ORT[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 even kWarpGemmIterationsForB (int4/int8 unaffected).

Weight preprocessor / runner / decode GEMV:

• fpA_intB_gemm_preprocessors*: W2_A16 quant 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_t runner instantiations.
• fpA_intB_gemv/details.h, moe_gemv.cu: int2 decode GEMV path.

QMoE op wiring:

• moe/moe_quantization.cc/.h: construct the fused uint2b_t runner for 2-bit; re-enable int2 PrePack (W2_A16); add zero_point_offset parse + 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_offset OPTIONAL_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:

• 2-bit parity max_diff 0.000488 / 0.003906 (tolerance 0.35) — effectively bit-exact vs the phase-1 fallback reference.
• Full test_qmoe_cuda.py: 122 passed, 15 skipped, 0 failed (no int4/int8/fp4/fp8 regression).
• Performance: int2 decode 1.019 → 0.037 ms/inference (tracks int4 0.036, ~27x speedup); prefill 1.047 → 0.101 ms (beats int4 0.161).

Opening as draft for review.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

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.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Thiago Pereira Rocha (thpereir)]

Heads-up on the latest force-push: I rebased this PR onto current main, which now includes the int2 fpA_intB primitives that landed on 2026-09-22:

• [CUDA] Add 2-bit fpA_intB GEMM and GEMV support #32699 — [CUDA] Add 2-bit fpA_intB GEMM and GEMV support (the low-level int2 GEMM/GEMV primitives + weight preprocessors)
• Add mixed-width weight contract for QMoE #32697 — Add mixed-width weight contract for QMoE (makes moe_quantization.cc accept expert_weight_bits == 2)

This PR originally introduced its own copies of those int2 primitives from scratch. Since they now exist upstream, I dropped the 15 duplicated primitive files in favor of main's versions to avoid duplicate class-template definitions, and reconciled the one call site to use main's unpack_uint2_transposed_to_int8_direct_cuda.

What remains here is the piece main still lacks: the int2 MoE runner (consumer) that wires those primitives into the QMoE op — the CutlassMoeFCRunner<*, cutlass::uint2b_t, ...> instantiations, GeGLU activation (GELU_taylor), int2 prepack path, and the fractional zero-point support. Final diff vs main is 19 files.

Validation on SM90: clean build (0 errors), full QMoE test suite 125 passed / 12 skipped / 0 failed.