[David Fan (jiafatom)]

Description

Add a decode-only packed CUDA path for INT2 and mixed-width integer QMoE together with a bounded dense dequantization fallback for configurations that are not eligible for packed decode.

The packed path supports FC1/FC2 weight-width combinations (2,2), (2,4), and (4,2) for FP16 and BF16:

• Extend the existing symmetric MoE GEMV to cutlass::uint2b_t and use FP32 accumulation for INT2.
• Prepack FC1 and FC2 independently into the SM80 interleaved layout using each layer's actual weight width.
• Dispatch FC1 and FC2 independently to INT2 or INT4 GEMV.
• Keep raw weight initializers alive so the same session can use the dense fallback when a later invocation is not eligible for packed decode.

Packed-path eligibility

Packed decode is selected only when both FC GEMV shapes pass is_moe_gemv_supported and all of the following hold:

• SM80 or newer.
• Symmetric blockwise quantization with no zero points.
• Fused interleaved SwiGLU.
• Block size 64 or 128.
• Prepacked FC1 and FC2 buffers are available.
• Expanded routing rows satisfy num_rows * k <= 8.
• Routing is supported by the fused expert-map builder: at most 256 rows and experts, with k in {1,2,4,6,8}.

CUDA completion status

From the CUDA functionality perspective, the supported QMoE contract is complete: configurations accepted by the operator either use the packed decode path or the dense dequantization fallback. This PR does not complete the fully optimized CUDA kernel matrix. Prefill and packed-ineligible decode configurations still rely on the fallback, so they remain functionally correct but are not yet in their final optimized form.

Why packed decode is faster

The fallback performs the following work on every invocation:

1. Dequantize all experts for both FC1 and FC2 from INT2/INT4 into FP16/BF16 scratch.
2. Write those dense weights to global memory.
3. Read the dense weights again in the generic CUTLASS grouped MoE GEMM runner.
4. Execute routing finalization after FC2.

This fixed work is particularly expensive during decode: a single token selects only k experts, but fallback dequantization still expands both matrices for every expert. For W weight elements, it writes about 2W bytes of dense FP16/BF16 weights before GEMM begins, in addition to reading the quantized weights and scales.

The packed path instead prepacks each layer once at session initialization into the SM80 interleaved layout. At inference time it builds the routed-row/expert maps and runs FC1 and FC2 directly on packed weights. One CUDA block owns one routed row and an output-column tile; as it traverses K, it loads an activation tile, loads an INT2/INT4 weight tile and its block scale, dequantizes that tile in registers, and immediately accumulates the dot product. No full dense weight matrix is materialized.

This reduces weight traffic relative to FP16 by approximately 8x for INT2 or 4x for INT4, before the comparatively small scale overhead, and only accesses experts selected by routing. FC1 also fuses bias and interleaved SwiGLU into its epilogue and reads source token activations through the routing map, avoiding a full gate/up intermediate and explicit top-k activation replication. FC2 currently writes routed outputs and uses the existing finalize kernel for the routing-weighted top-k reduction.

The GEMV organization is intentional for decode and uses a conservative hard shape gate rather than a runtime GEMV/GEMM crossover: expanded_num_rows = num_rows * k must be at most 8. This permits at most 8, 4, 2, 1, or 1 input rows for k = 1, 2, 4, 6, or 8, respectively. The limit applies to the total top-k-expanded rows for the invocation, not an average per-expert M. Within this envelope, a general grouped-GEMM M tile is underfilled; the packed kernel uses CtaM=1 and exposes parallelism over output columns and K reduction instead. It also skips dense-runner tactic profiling and workspace allocation. For prefill, larger per-expert M allows grouped GEMM to reuse weight tiles across many rows and use Tensor Cores efficiently, so a future packed grouped-GEMM path is preferable to repeatedly reading weights through row-wise GEMV.

Why the dense fallback remains

This PR optimizes a decode subset while providing broader correctness coverage through the bounded dense dequantization fallback. Accepted INT2 and mixed-width configurations that miss any packed-path condition continue through the bounded fallback, including prefill or unsupported GEMV shapes, asymmetric zero points, other supported block sizes, unsupported activation/fusion or routing configurations, older GPUs, and sessions where packed decode is disabled or unavailable.

The fallback independently dequantizes FC1 and FC2 to FP16/BF16 scratch and invokes the existing dense A16 CUTLASS MoE runner. Its scratch allocation remains protected by ep.cuda.qmoe_int_dequant_max_scratch_bytes. This is slower and uses more temporary memory than packed GEMV, but preserves correctness across the broader supported contract.

A future packed grouped-GEMM path can remove fallback usage for prefill and larger decode workloads. The fallback should remain until native packed kernels cover the complete supported configuration matrix or the remaining configurations are explicitly declared unsupported.

Testing

• Built onnxruntime_provider_test with CUDA 12.8.
• Passed all 18 related MoETest.QMoETest_MixedWidth* and MoETest.QMoETest_Int2* tests on NVIDIA A100 SM80.
• Packed-path tests cover FP16/BF16, (2,2), (2,4), (4,2), and block sizes 64/128.
• Packed tests use a one-byte dense scratch limit to prove dense dequantization is bypassed.
• clang-format --dry-run --Werror and git diff --check passed.

[Copilot]

Copilot review overview

🟡 Changes recommended

Unsupported routing configurations can enter the packed path with uninitialized maps, and block-size-128 coverage is missing.

Get a fresh assessment by requesting another Copilot review.

Review effort: Balanced
Findings: 1 · 1

Open (2)

• Packed path ignores unsupported routing-map build failures · New
• Add packed-decode coverage for block size 128 · New

What changed in this PR

Adds a CUDA packed INT2 QMoE decode path while retaining dense fallback behavior.

Changes:

• Adds INT2 weight prepacking and FP32-accumulating GEMV kernels.
• Routes eligible FP16/BF16 SwiGLU decode workloads through packed GEMV.
• Adds packed-path and fallback CUDA tests.

File	Description
onnxruntime/​test/​contrib_ops/​moe_test.cc	Adds CUDA packed-decode tests.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_quantization.h	Extends QMoE packed-weight state.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_quantization.cc	Implements prepacking, dispatch, and packed execution.
onnxruntime/​contrib_ops/​cuda/​llm/​moe_gemm/​moe_gemv.h	Documents INT2 GEMV support.
onnxruntime/​contrib_ops/​cuda/​llm/​moe_gemm/​moe_gemv.cu	Adds INT2 kernel traits and dispatch.

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Copilot review overview

🟡 Changes recommended

Packed decode retains unnecessary dense profiling and fallback-only prepacking overhead, while nontrivial routing remains untested.

Get a fresh assessment by requesting another Copilot review.

Review effort: Balanced
Findings: 2

Open (2)

• Skip dense tactic profiling and workspace allocation for packed decode · New
• Add multi-row, multi-expert packed decode coverage with top-k two · New

Resolved since last review (2)

• Packed path ignores unsupported routing-map build failures
• Add packed-decode coverage for block size 128

[Copilot]

Copilot review overview

🟡 Changes recommended

Prepacking can read beyond buffers for valid fallback-only INT2 initializer shapes.

Get a fresh assessment by requesting another Copilot review.

Review effort: Balanced
Findings: 1

Open (1)

• Prepacking fallback-ineligible INT2 weights can overread scratch buffers · New

Resolved since last review (2)

• Add multi-row, multi-expert packed decode coverage with top-k two
• Skip dense tactic profiling and workspace allocation for packed decode

[Ti-Tai Wang (titaiwangms)]

Reviewed the INT2/mixed-width packed QMoE path (five independent static reviews, then rechecked against head 72687ce). The K=32 prepack/fallback issue previously reported in this PR has been addressed by the new static eligibility guard; I am not re-reporting it as an open finding.

Major: validate INT2/mixed-width scale dtypes before the packed path. QMoE T2 permits Float8 scales for other quantization modes, while the INT2/mixed-width common validation at moe_quantization.cc:825-856 checks scale presence and shape but not dtype. PrePack intentionally skips the scale inputs for INT2/mixed-width (:2379-2381). For an otherwise packed-eligible FP16 INT2 model (e.g. E=1, hidden=inter=512, block=64, fused SwiGLU, k=1), Float8E4M3FN [1,1024,8] FC1 scales can therefore reach prepare_scale_zp (:1584-1590): it allocates the one-byte-per-element SizeInBytes() buffer but selects the float32 transpose for every type other than FP16/BF16 (:1394-1409), accessing four bytes per logical element on both input and output. The packed GEMVs also treat the resulting scales as the activation dtype (:1814-1826). The FP16/BF16 checks in the dense dequantization branch occur too late to protect this path. Please reject non-matching FC1/FC2 scale dtypes in the common INT2/mixed-width validation before scale preparation; make the transpose default reject unsupported types rather than assume float32. This is a static memory-bounds finding, not a CUDA-sanitizer result.

Test coverage (non-blocking): the new multirow/multiexpert packed-vs-dense test exercises (2,2), but (2,4)/(4,2) packed tests are still single-row/single-expert, and the new scale tensors are all 1.0 (moe_test.cc:2090-2091,2317-2318). A mixed-width multirow/multiexpert case with distinct per-block scales would check transpose, expert strides, and routing together.

No CUDA build, sanitizer, performance measurement, or real Olive-produced checkpoint was part of this static review.

[David Fan (jiafatom)]

Ti-Tai Wang (@titaiwangms) Addressed the scale-dtype bounds finding in 8a80e2c. The common INT2/mixed-width validation now requires FC1/FC2 scales to match the FP16/BF16 activation dtype before scale preparation, and the generic transpose fallback explicitly accepts only float32 instead of treating every unknown dtype as float32. I also added a regression case with schema-valid Float8E4M3FN FC1/FC2 scales that expects the new validation error. The targeted onnxruntime_provider_test target builds and lintrunner passes in jiafa-dev. Runtime execution remains blocked because CUDA initialization in that container currently returns cudaErrorNoDevice. I kept the mixed-width multirow/distinct-scale coverage suggestion out of this safety fix since it was marked non-blocking.

[David Fan (jiafatom)]

Validation update: restarting jiafa-dev restored CUDA visibility (8x A100, SM80). The new QMoETest_MixedWidthCudaRejectsUnsafeFallbackInputs Float8 scale rejection case now runs and passes. The packed/fallback regression set also passes 4/4: QMoETest_Int2CudaPackedDecode, QMoETest_Int2CudaPackedDecodeBlock128, QMoETest_Int2CudaPackedDecodeMultiRowExpertParity, and QMoETest_Int2CudaPackedDecodeFallback.

[Tianlei Wu (tianleiwu)]

Reviewed all five changed files at 8a80e2c. The earlier routing, preprocessing, profiling, and scale-dtype fixes are present. One additional performance concern remains in packed decode scale preparation; see the inline comment. No additional correctness defect was established. Validation was static review plus byte-count checks and git diff --check; no local CUDA operator tests or benchmark were run.

[Tianlei Wu (tianleiwu)]

Reviewed all five changed files at the current head and rechecked the prior scale-cache thread, which is resolved by e0be11a. One blocking CUDA Plugin EP compatibility issue remains in the new cached-scale path; see the inline comment. The current-head Linux CUDA Plugin EP run reproduces it in all newly added packed INT2 tests. git diff --check origin/main...HEAD passes.

[Tianlei Wu (tianleiwu)]

Re-reviewed all five changed files at the latest head. The earlier routing, preprocessing eligibility, scale-cache, and CUDA Plugin EP dtype issues are resolved, and their threads are already closed. The packed path now has coherent shape/routing gates and preserves the dense fallback inputs; the latest standard CUDA and Linux/Windows CUDA Plugin EP checks pass. No remaining actionable findings.