[David Fan (jiafatom)]

Description

Implements Workstream 4, PR 4a of the delivery plan in #32657: an independently testable kernel foundation, not the QMoE prefill integration planned for PR 4b.

• Add packed symmetric W2A16 grouped GEMM for SM80, FP16/BF16 input/output, block size 64 and FP32 accumulation, without full global-memory A16 expert-weight materialization or caller-supplied GEMM workspace.
• Support explicit CTA 32x128x64 / warp 32x32x64 and CTA 64x128x64 / warp 32x64x64, both with three stages. Keep 32 rows as the default; no automatic tactic selection.
• Add scalar-reference parity tests, production W4/W8 regressions, GPT-OSS-20B-shaped balanced/skewed cases, and opt-in CUDA-event benchmarks against W4 and dense FP16 baselines.
• Register tests and implementation dependencies for legacy and plugin CUDA builds. Dense runner template instantiation also requires the existing five FP16 SM80 fused launchers.
• Fix stale routing_record API usage in packed INT decode, exposed by the clean legacy build, in a separate commit.

No QMoE prefill dispatch, bias/activation epilogue, SwiGLU/finalization integration, or additional architecture support is enabled. Mixed FC1/FC2 execution, lifetimes, runtime dispatch and model-level qualification remain PR 4b work. Validation notes are kept in this description rather than a repository Markdown artifact.

Motivation and Context

Provide the native packed INT2 grouped-GEMM building block needed to replace full expert-weight dequantization in mixed INT2/INT4 QMoE prefill. Related design: #32657; packed decode foundation: #32761.

Kernel Contract

For E experts, P routed rows, N output columns and K reduction elements:

Buffer	Layout
activations	FP16 or BF16 [P,K], contiguous rows grouped by expert
packed_weights	E consecutive SM80 fpA_intB prepacked INT2 matrices, K*N/4 bytes each
block_scales	Same dtype as activations [E,K/64,N]
expert_row_ends	E device int64 cumulative exclusive row ends
output	Same dtype as activations [P,N], same grouped row order

N and K must be positive multiples of 64; partial M/N tiles and empty experts are supported. P must be positive and fit in int32, and E must be positive. IsInt2GroupedGemmSupported checks metadata; RunInt2GroupedGemm also requires non-null buffers and checks 16-byte alignment of activations, packed_weights, block_scales, and output before invoking CUTLASS. The caller supplies the actual SM/SM count, keeps buffers alive until stream completion, ensures matching capacities and nonoverlapping output, and guarantees nonnegative, nondecreasing row ends terminating at P. Repeated ends represent empty experts; device offset contents are not validated on the host.

Raw unsigned codes represent (code - 2) * scale. Convert raw [E,N,K/4] weights using unpack_uint2_transposed_to_int8_direct_cuda and preprocess_weights_for_mixed_gemm_cuda with SM80/W2_A16. Preprocessing mutates its scratch input. GEMM converts tiles internally and performs no routing or finalization.

Local Validation

Latest alignment fix (d3d60ac): corrected the shared CUTLASS tensor_aligned pointer check to use alignment * sizeof_bits::value / 8 bytes, while retaining stride checks in elements and the entry-point 16-byte guard. Added positive packed-weight suballocation tests at +16/+32/+48 bytes for FP16/BF16 and both INT2 tactics, plus production INT4/INT8 offset regressions. The new INT2 tests reproduced Error Misaligned Operand before the fix and passed afterward on A100. Existing +8-byte rejection tests remain.

Rebuilt and validated in jiafa-dev on October 1, 2026. Both legacy and plugin passed all 23 enabled grouped tests and all 23 under Compute Sanitizer memcheck with zero errors. Forced recompilation of the MoE CUDA translation units because some cached nvcc objects recorded no header dependencies; the new INT4 offset test also detected the stale object before recompilation.

A100-SXM4-80GB, CUDA 12.8.61, GCC 13.3, Release, SM80, internal tests enabled, QUICK_BUILD=OFF; September 29-October 1, 2026. Separate build trees started empty without old objects or compiler caches, then fixes were rebuilt in those trees.

• Legacy (d3d60ac): all 23 grouped tests passed; all 23 passed under Compute Sanitizer memcheck with zero errors. Existing QMoE regression: 53 passed, 11 skipped.
• Plugin (d3d60ac): all 23 grouped tests passed; all 23 passed under Compute Sanitizer memcheck with zero errors. The combined Int2GroupedGemm*:QMoE run had 78 passed and 12 skipped (23 grouped tests plus 55 passing QMoE-related tests, including three directly registered FP4-to-FP8 internal tests).
• Plugin production compilation used local CMAKE_CUDA_FLAGS="--diag-suppress=970 --diag-suppress=2189" for Abseil/Protobuf diagnostics. This is not evidence that default-toolchain CI passes.
• Small cases compare every output against a scalar oracle, covering FP16 and BF16, minimum alignment, empty experts, partial tiles, dense parity, and the 32-row and 64-row candidates. Production W4/W8 block-64 paths are exercised directly.
• Model-shaped cases use E=32, top-k=4, K=2880, FC1 N=5760, FC2 N=2880, T=128/512/2048 and P=512/2048/8192. FP16 covers balanced and skewed routing at all three sizes; BF16 covers T=128. They check all outputs for finiteness and sampled first/middle/last expert rows and columns on both sides of 64-column boundaries against the oracle.
• Mixtral INT4 skips because its dimensions are too small; it is not counted as CUDA INT4 coverage. Plugin routing-statistics testing also skips because that EP does not support it. CPU-specific skips are not counted as validation.
• Full onnxruntime_test_all suite was not run. Three opt-in benchmarks are disabled by default. P is routed rows, not model input tokens T.

Run from the appropriate build directory with an available GPU:

# Legacy internal-test loader
GTEST_FILTER="Int2GroupedGemm*" ./onnxruntime_provider_test --gtest_filter=CUDA_EP_Unittest.All
./onnxruntime_provider_test --gtest_filter="*QMoE*"
GTEST_FILTER="Int2GroupedGemm*" compute-sanitizer --tool memcheck --error-exitcode 99 \
  ./onnxruntime_provider_test --gtest_filter=CUDA_EP_Unittest.All

# Plugin tests are directly registered
./onnxruntime_provider_test --gtest_filter="Int2GroupedGemm*:*QMoE*"
compute-sanitizer --tool memcheck --error-exitcode 99 \
  ./onnxruntime_provider_test --gtest_filter="Int2GroupedGemm*"

# Opt-in legacy benchmark; use DISABLED_GptOssKernelBenchmark for W4 comparison
GTEST_FILTER="Int2GroupedGemmTest.DISABLED_GptOssTacticBenchmark" \
GTEST_ALSO_RUN_DISABLED_TESTS=1 \
./onnxruntime_provider_test --gtest_filter=CUDA_EP_Unittest.All

Initial Performance Evidence

Five warmups after an initial launch; three trials of 30 launches on a nonblocking stream; median per-launch CUDA-event milliseconds. This measures stream elapsed time around host-submitted launches, not isolated profiler kernel duration. Allocation, packing, H2D/D2H transfers and the reference computation are outside GEMM timing.

First run, balanced routing, E=32, top-k=4, K=2880, FC1 N=5760, FC2 N=2880, tile32:

T	P	W2 FC1	W4 FC1	W4 FC2
128	512	0.403695	0.391612	0.212548
512	2048	0.760764	0.766020	0.323072
2048	8192	2.476240	2.570210	1.253720

W4 uses a different quantized value distribution, not an equal-quality model baseline. W2 and W4 are close; the T=2048 W2 trial range was 2.47624-2.98469 ms, so this does not establish a robust speedup.

Second run compares identical INT2 codes, activations, scales and offsets. The dense path uses production DequantizeNBits and MoeGemmRunner<half,half,half>. Cached dense dequantizes once before timing; Dequant+dense dequantizes all experts before every GEMM. Allocations are excluded. This is not an active-experts-only or tiled-dequantization baseline.

Routing	T	W2 tile32	W2 tile64	Cached dense tile32	Cached dense tile64	Dequant+dense tile32	Dequant+dense tile64
Balanced	128	0.404	0.602	0.651	0.663	1.538	1.548
Balanced	512	0.761	0.605	0.884	0.702	1.765	1.586
Balanced	2048	2.728	1.905	2.375	1.905	3.678	2.765
Skewed	128	0.233	0.228	0.217	0.189	1.097	1.071
Skewed	512	0.798	0.641	0.737	0.488	1.510	1.353
Skewed	2048	2.819	1.979	2.532	1.700	3.139	2.585

Balanced experts receive T/8 rows each. Skewed counts are 1,31,T-32,T,T,T, with the other experts empty; total 4T. Each benchmark case passed sampled parity and finiteness checks. GPU 1 had no competing compute clients at the first measurement start. Clocks were unlocked, execution order fixed, and CPU plugin compilation ran concurrently during the second run. Repeat controlled measurements before selecting dispatch thresholds or asserting stable speedups.

Tile64 helps larger per-expert M but is slower at balanced T=128; default remains tile32. Cached dense GEMM can be faster, especially with skew. A tile128 W2 candidate was rejected by CUTLASS thread-map compile assertions and is not exposed. These are not full prefill, routing, activation, finalization or TTFT measurements.

Prepack and Memory Accounting

One-pass prepack event times were 23.1-24.6 ms (W2 FC1), 24.1-24.3 ms (W4 FC1) and 10.2-12.5 ms (W4 FC2), including transpose/preprocessing launches and possible host submission gaps, not steady-state GEMM overhead.

Buffer	W2 FC1 bytes	W4 FC1 bytes	W4 FC2 bytes
Packed weights	132710400	265420800	132710400
A16 scales	16588800	16588800	8294400
Prepack scratch	4147456	8294656	4147456

Tests also retain a raw GPU copy equal to packed-weight size, activations, outputs and offsets. Dense FC1 adds 1,061,683,200 bytes (1012.5 MiB). These are buffer capacities, not measured peak process memory. W2 halves FC1 packed-weight bytes, not total model memory.

Remaining Qualification

• Supported-platform CI, including the default plugin toolchain configuration, remains separate from local validation.
• Controlled timing repeats are needed before automatic tactic selection or performance gates.
• Profiler evidence and peak allocation measurements are needed before end-to-end memory/prefill claims.
• Current commit 2edffb5 passed FP16/BF16 grouped regression, alignment rejection coverage, and grouped memcheck in both legacy and plugin builds. The full test suite and opt-in benchmarks were not rerun for this update.
• PR 4b will integrate routing, mixed FC1/FC2, SwiGLU/finalization, lifetimes and model-level evaluation.

[Copilot]

Copilot review overview

🔵 Needs a closer look

Low-level CUTLASS kernel changes and pending default-toolchain CI require human validation, and one diagnostic remains inaccurate.

Review effort: Balanced
Findings: 1

Open (1)

• Update missing-scale error to mention UINT2 · New

What changed in this PR

Adds the SM80 FP16/INT2 grouped-GEMM foundation for future QMoE prefill integration.

Changes:

• Adds 32-row and 64-row packed INT2 GEMM tactics.
• Adds correctness, regression, validation, and opt-in benchmark coverage.
• Registers legacy/plugin tests and updates routing snapshot usage.

File	Description
moe_gemm_int2_test.cc	Adds kernel tests and benchmarks.
moe_quantization.cc	Updates routing snapshot capture/consume flow.
moe_gemm_template_dispatch.h	Enables constrained UINT2 dispatch.
moe_gemm_kernels_fp16_uint2.cu	Implements validation and SM80 launchers.
moe_gemm_int2.h	Defines the kernel entry-point contract.
moe_cutlass_kernel.h	Enables UINT2 scale validation.
moe_int2_grouped_gemm.md	Documents usage, validation, and measurements.
onnxruntime_unittests.cmake	Registers plugin tests and dependencies.

---

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

[Copilot]

Copilot review overview

🟢 Approval recommended

The implementation and coverage match the scoped kernel contract; only a non-blocking diagnostic mismatch remains.

Review effort: Balanced
Findings: 1

Open (1)

• Update can_implement diagnostic to list all supported scale-bearing types · New

Resolved since last review (1)

• Update missing-scale error to mention UINT2

[David Fan (jiafatom)]

CI triage: the only outstanding non-success/non-neutral check is React Native CI iOS E2E Tests (cancelled). Both the original attempt and the manual rerun show the same simulator initialization/cleanup failure signature:

• iOS app build succeeds.
• Detox starts the iPhone 15 / iOS 17.4 simulator and invokes xcrun simctl bootstatus.
• About four minutes later, cleanup reports CoreSimulator.SimError, code=405: Unable to shutdown device in current state: Shutdown.
• Jest reports that it did not exit because asynchronous operations remain, and the job stays alive until cancelled. The E2E steps lasted roughly 79 minutes and 53 minutes respectively.

Original job: https://github.com/microsoft/onnxruntime/actions/runs/36667939436/job/109741418017
Rerun job: https://github.com/microsoft/onnxruntime/actions/runs/36667939436/job/109906898589

The available logs do not show application tests or ORT inference reaching a failure. This points to simulator setup/teardown rather than the CUDA changes, but the underlying simulator boot failure is not established by these logs. Other recent workflow runs have succeeded, so this is not evidence of a platform-wide outage.

Could a React Native CI maintainer check the macOS runner/CoreSimulator environment and advise on a targeted rerun on a healthy runner? Leaving the CUDA code and check requirements unchanged; another identical rerun without environment diagnosis may reproduce the same hang. Separate CI hardening could add simulator boot diagnostics and bounded cleanup timeouts.

[Copilot]

Copilot review overview

🔵 Needs a closer look

The low-level architecture-specific CUDA kernel and toolchain-sensitive plugin build require final human and CI validation.

Review effort: Balanced
Findings: None

Resolved since last review (1)

• Update can_implement diagnostic to list all supported scale-bearing types

[Copilot]

Copilot review overview

🟡 Changes recommended

The new entry point can admit insufficiently aligned caller buffers into vectorized CUTLASS accesses.

Review effort: Balanced
Findings: 1

Open (1)

• Validate buffer alignment in bytes before vectorized kernel launch · New

[Copilot]

Copilot review overview

🔵 Needs a closer look

The low-level CUDA kernel and latest alignment changes still require completed legacy and plugin build/runtime validation.

Review effort: Balanced
Findings: None

Resolved since last review (1)

• Validate buffer alignment in bytes before vectorized kernel launch

[Tianlei Wu (tianleiwu)]

Reviewed commit 2edffb5ee626c5c4ca9e702cd040f7952360c3db. Requesting changes for one remaining P2 alignment-contract issue, detailed in the follow-up on the existing alignment thread.

Validation: a compiled host-side harness exercised the actual CUTLASS can_implement for FP16/BF16 and both tactics. Current GitHub CUDA/plugin checks pass. I did not rerun the SM80 numerical or sanitizer suite because the local GPUs are H200/SM90.

[Tianlei Wu (tianleiwu)]

Re-reviewed d3d60acf26. The earlier alignment P2 is resolved: tensor_aligned now checks pointer alignment in bytes (alignment * sizeof_bits<Element> / 8) and stride in elements. This loosens the over-strict UINT4 (32 B) and UINT2 (64 B) pointer checks to the actual 16 B access width and tightens FP16/BF16 A/C/D from 8 B to the 16 B that the vectorized loads/stores need, so no valid production configuration is newly rejected. Per-expert B offsets already use int64_t.

Remaining items are non-blocking (inline):

• Coverage: the numerical tests require device 8.0 exactly, but CUDA and plugin CI run on A10 (SM86), so they all skip in CI, including the W4/W8 regressions for the shared alignment change.
• Build: the plugin test object library recompiles about 15 CUTLASS TUs, mostly for test-only baselines, and adds diagnostic suppressions without the CUDA 13 gate.
• Rebase: the moe_quantization.cc change is already on main via #32932.
• Nit: the new files are missing the Microsoft copyright/MIT license header.

I did not rerun the SM80 numerical or sanitizer suites for this round.