[David Fan (jiafatom)]

Description

Implements the initial Workstream 4, PR 4b runtime integration from #32657. PR #32963 (4a) and PR #33045 (4a2) are merged. This PR targets main and is rebased onto f0843448e7, with only the PR4b commits replayed; no conflicts occurred.

• Add default-enabled packed INT2 FC1 + INT4 FC2 QMoE prefill for SM80 and later, FP16/BF16 activations/output, symmetric quantization, block size 64, and interleaved fused SwiGLU.
• Reuse existing routing and expert-map kernels, expand rows by expert, run the 4a INT2 grouped FC1, apply expert bias and configurable SwiGLU, run production INT4 grouped FC2, and finalize top-k weighted outputs with FC2 bias.
• Allocate aligned per-invocation workspace; reuse expanded-input storage for FC2 output after FC1 completes on the same stream. Avoid full A16 expert-weight materialization on the selected packed path.
• Reuse packed weights and transposed scales, including runtime scale inputs. Separate grouped-prefill prepack eligibility from decode-specific minimum dimensions.
• Keep packed decode priority when eligible. Prefill is enabled by default; no environment variable is required. Set ORT_ENABLE_QMOE_INT2_PREFILL=0 to opt out. Packed GEMV also defaults on. Unsupported configurations retain existing fallback behavior and scratch limits; debug logging and the one-byte scratch limit are validation controls, not normal runtime requirements.

Motivation and Context

Connect the independently tested grouped-GEMM foundation in #32963 to the mixed-width QMoE execution path. Delivery plan: #32657. Packed decode foundation: #32761.

Current Scope

• SM80 and later, including SM86, SM89, SM90, SM120 and SM121; INT2 FC1 and INT4 FC2; block size 64; no zero points; swiglu_fusion=1. Both production eligibility and grouped FC1 accept the actual device SM >= 80. FC1 retains the SM80 CUTLASS template and packed layout; non-TMA INT4 FC2 uses the existing Ampere fallback on newer architectures. This is not a new Hopper/Blackwell-native TMA implementation. The binary must contain suitable native code or compatible virtual PTX for the target GPU; an SM80 cubin alone is not cross-major compatible.
• Positive hidden/intermediate dimensions divisible by 64, bounded int-sized dimensions/routed-row count, and at most 256 experts.
• Current packed path requires cached prepacked weight buffers, obtained from weight initializers with weights_prepacked=0. Raw initializers are retained for fallback.
• Both stages use the fixed 32-row grouped tactic. FC2 uses four stages, supported by both QUICK_BUILD and full builds. No automatic threshold/tactic tuning.
• Eligible packed prefill is row-tiled to keep estimated temporary scratch within 256 MiB, including packed workspace, tile-local routing metadata and uncached scale transposes. The configured ep.cuda.qmoe_row_tile_size is an upper bound on tile rows. Full and final partial tiles are sized explicitly; the candidate tile size is halved until it fits. Tiles reuse packed weights and stream-ordered workspace without dense weight dequantization. Long prompts and explicit row tiling no longer select dense fallback solely because of their row count. If even one row plus fixed scale scratch exceeds the budget, execution reports a packed-scratch error instead of attempting larger dense allocations. This is not a whole-process GPU memory cap.

Latest Submission: SM80+ and Default Dispatch

• Corrected the earlier workspace-cap fallback after reviewer feedback: eligible long prompts remain on the packed path and run tile-by-tile, with input/output row offsets and tile-local routing maps. Dense fallback remains for ineligible configurations. The workspace-limit test now requires 8193-row FP16/BF16 inference to succeed with dense scratch capped at one byte; explicit-tiling tests likewise expect success. New multi-expert tiled tests cover top-k 1/2/4, channel-varying biases, cached/runtime scales and partial final tiles. These tests were added but not run locally, as requested; Docker formatting and diff checks passed. The historical 4096-token benchmark below predates this tiling implementation and is not validation of the fix.
• Revision 6c24651036 removes ORT_ENABLE_QMOE_INT2_GEMV. INT2 packed GEMV now follows the existing ORT_DISABLE_MOE_GEMV switch: unset/0 permits eligible GEMV dispatch, and 1 disables it for both INT2 and native INT4/INT8. Set it before starting the process. The INT2 prefill switch remains independent. Test overrides were migrated and the disabled-path test now includes a single-token case with a one-byte dense scratch limit to detect bypass of the shared switch. No local build or runtime tests were run; Docker clang-format and git diff --check passed. Old INT2 GEMV environment settings below describe historical benchmark revisions only and are no longer read by current code.
• Revision 78294ec898 fixes all eight clang-format suggestions and the BF16 mixed-width initialization failure reported by Linux/Windows plugin CI. QUICK_BUILD now constructs the already-built BF16 dense fallback runner for mixed-width/INT2 models; native BF16 INT4/INT8 runner restrictions remain unchanged. Existing packed-prefill tests cover the failed cases.
• No local build or runtime tests were run for this revision, as requested. Only clang-format 20.1.8 checks in Docker jiafa-dev and git diff --check were run and passed. New CI results are pending; no plugin runtime pass is claimed for this fix.
• Grouped numerical tests pass the actual device SM and no longer skip GPUs newer than SM86. Metadata coverage explicitly accepts SM80/86/89/90/100/103/110/120/121 and rejects pre-SM80 devices. Prefill operator tests now run on any SM80+ GPU, including BF16 in QUICK_BUILD; forced packed tests retain the one-byte dense scratch guard. Metadata checks are not hardware execution evidence.
• Added an unset-environment test covering FP16/BF16 at 1/33/257 rows with dense fallback blocked. Explicit opt-out and unsupported-configuration tests remain.
• No local build or runtime tests were run for the SM80+ extension, as requested. Docker clang-format and diff checks passed. SM86/89/90/120/121 execution of this extension remains unverified; CI results must be tied to the actual runner GPU before claiming architecture coverage. The following validation and benchmark records are historical SM80 results, not evidence for this new revision or its four-stage FC2/default-enable changes.
• Previous-head WebGPU CI (run 36932358189) remains unresolved: all compile jobs passed; the Windows novcpkg test lost its device during an Einsum test, causing subsequent failures and two result-summary jobs to fail. The vcpkg test passed. No WebGPU code fix is included here, and a transient-driver cause is not yet established. The new push will trigger fresh CI; old-head jobs are not manually duplicated.

BF16 QUICK_BUILD Guard: Rationale and Behavior Change

In moe_quantization.cc, changing the outer guard from #if !defined(ORT_QUICK_BUILD) && defined(ENABLE_BF16) to #if defined(ENABLE_BF16) does not enable every BF16 integer runner in quick builds. The #if !defined(ORT_QUICK_BUILD) restriction is retained inside that branch around the native BF16 INT4/INT8 runner construction. Only the dense BF16 runner used when use_int_dequant_fallback is true is made available in quick builds.

use_int_dequant_fallback applies to integer-quantized models with mixed FC weight widths or expert_weight_bits=2. This branch constructs CutlassMoeFCRunner<__nv_bfloat16, __nv_bfloat16, __nv_bfloat16>, whose dense BF16 implementation is already instantiated under ENABLE_BF16, including QUICK_BUILD. It does not require enabling an omitted native BF16 integer specialization.

Configuration	Before this guard change	After this guard change
BF16 + QUICK_BUILD + INT2/mixed-width	No runner was constructed; session initialization failed at m_moe_runner != nullptr, before packed-kernel dispatch.	Construct the existing dense BF16 fallback runner, removing that initialization failure. Runtime dispatch still checks packed-path eligibility.
BF16 + QUICK_BUILD + ordinary uniform INT4/INT8	Native BF16 integer runner construction excluded.	Unchanged: still excluded by the inner QUICK_BUILD guard.
BF16 + full build	Dense/native runners constructed according to the existing branch conditions.	Runner construction unchanged.
FP16, or builds without ENABLE_BF16	Existing FP16 behavior/BF16 exclusion.	Unchanged by this guard change.

Constructing a dense fallback runner does not mean dense execution is forced. The constructor requires a non-null runner even when inference later takes packed INT2 GEMV or packed INT2/INT4 prefill. Eligible calls use those packed paths; unsupported calls retain dense fallback and its scratch limits. This is an intentional behavior change for previously failing BF16 QUICK_BUILD INT2/mixed-width configurations, not a change to the existing pure INT4 default dispatch. The table describes the guard change specifically, separately from this PR's default-enabled packed prefill feature. Build success alone does not establish runtime correctness; no additional local build or runtime tests were run for this explanation.

Historical Validation

Local A100-SXM4-80GB / SM80, CUDA 12.8, GCC 13.3, Release, QUICK_BUILD=OFF, in Docker container jiafa-dev on September 30, 2026.

• Legacy build succeeded. All 6 focused packed-prefill tests passed; they also passed Compute Sanitizer memcheck with ERROR SUMMARY: 0 errors. After formatting the new implementation files, rebuilt and reran all 6 focused tests successfully.
• Legacy QMoE regression: 59 passed, 11 skipped. CPU-specific and unsupported small Mixtral cases are not counted as coverage.
• Forced packed tests disable INT2 GEMV and cap dense dequantization scratch at 1 byte, preventing silent decode/dense fallback from satisfying parity checks.
• FP16/BF16: 1, 33, and 257 rows. Independent scalar-reference multi-expert tests cover top-k 1/2/4, balanced and skewed routing, empty experts, unequal hidden/intermediate dimensions, nonuniform block/output scales, expert-specific biases, and nondefault SwiGLU alpha/beta/limit.
• Runtime-scales tests exercise transposition and 64/128-dimensional cases in both orientations. Negative cases verify fallback for zero points, block size 128, reversed weight widths, non-SwiGLU activation, runtime weights, and disabled prefill.
• git diff --check passed.
• Plugin rebuild and tests are still in progress; no plugin result is claimed for this commit. The existing plugin configuration uses local CMAKE_CUDA_FLAGS="--diag-suppress=970 --diag-suppress=2189"; it does not establish default-toolchain CI success.

Run from the build output directory on an available SM80 GPU:

./onnxruntime_provider_test --gtest_filter="MoETest.QMoETest_MixedWidthCudaPackedPrefill*"
compute-sanitizer --tool memcheck --error-exitcode 99 \
  ./onnxruntime_provider_test --gtest_filter="MoETest.QMoETest_MixedWidthCudaPackedPrefill*"
./onnxruntime_provider_test --gtest_filter="*QMoE*"

Full-Model Validation (October 1, 2026)

All model execution, dependency installation, benchmarking and NVML sampling below ran inside Docker jiafa-dev on an A100-SXM4-80GB. Used a clean non-plugin CUDA wheel, Release commit 463fa42cb5, CUDA 12.8, cuDNN 9.7, SM80, QUICK_BUILD=OFF. The installed wheel's build metadata was checked; all six focused packed-prefill tests passed in the clean build, and an installed-wheel CUDA MatMul smoke test passed.

• Full mixed-width GPT-OSS-20B export: 24 QMoE layers, INT2 FC1 / INT4 FC2, block size 64, top-k 4, fused SwiGLU, FP16 activations and no zero points.
• GEMM-only validation (ORT_ENABLE_QMOE_INT2_PREFILL=1, ORT_ENABLE_QMOE_INT2_GEMV=0): 19-token prefill followed by eight single-token decode calls; generated **Paris** is the capital of France. All logits/KV outputs were finite. ORT profiling recorded 216 QMoE CUDA events (24 layers x 9 runs).
• Combined-path validation (both flags set to 1): 309-token prefill followed by two single-token decode calls, ending naturally at EOS; generated Paris. Actual branch-log assertions confirmed 24 packed_int_prefill calls at 309 rows followed by 48 packed_int_gemv calls at one row. All logits/KV outputs were finite; profiling recorded 72 QMoE CUDA events. This confirms same-session long-prompt GEMM prefill -> GEMV decode, not repeated transitions or concurrency/capture safety.
• Both mixed validation runs kept ep.cuda.qmoe_int_dequant_max_scratch_bytes=1, preventing the full dense expert-weight dequantization fallback. This is not a one-byte cap on all workspace or GPU allocations.
• GPU I/O binding retains KV on device. Even-numbered sliding-window layers use 128-position cache capacity with shared past/present buffers; odd-numbered layers retain full-history KV. The 309-token test exceeds the sliding-window size.

Model Artifact Size

Measured logical file sizes of model.onnx plus every unique referenced external-weight file, excluding tokenizer/configuration files. Each export references one model.onnx.data file. GiB = bytes / 2^30.

Export	ONNX graph bytes	External weights bytes	Total bytes	Total GiB
Mixed INT2 FC1 / INT4 FC2	212,376	8,012,132,352	8,012,344,728	7.462
INT4 FC1 / INT4 FC2	209,568	11,794,870,272	11,795,079,840	10.985

The mixed export is 3,782,735,112 bytes (3.523 GiB, 32.07%) smaller on disk. This measures the whole existing export, not just QMoE weights. The two exports differ in QMoE block size as well as bit width, so this is not a bit-width-only ablation. Smaller artifacts did not translate into lower measured process GPU memory in these runs.

Initial Direct ORT Performance Baseline

Same GPU 0, installed wheel, harness and tokenizer for both exports; tokenizer files were byte-identical. Batch size 1, input lengths 128/512/2048/4096, 128 generated tokens, three warmups and ten measured requests per configuration. Each configuration ran in a fresh process. Each model completed 40 measured requests / 5,120 generated tokens; all ten token sequences within each configuration matched. Throughput was recomputed from the raw timing totals.

This is a synthetic direct-ORT benchmark, not tau2-bench. Prompts repeat Paris is the capital of France. Its museums and public parks attract visitors., followed by Summarize the context in detail. and GPT-OSS chat delimiters. Context tokens are repeated/truncated to obtain exact total input lengths. Greedy autoregressive generation ignores EOS to enforce the output length; this is not a quality evaluation.

Input tokens	Model	TTFT P50 ms	TTFT P95 ms	Prefill tokens/s	Decode tokens/s	Sampled process peak GiB
128	Mixed INT2/INT4	21.31	21.39	6,418.11	196.44	15.02
128	INT4	26.39	26.76	5,099.97	203.12	13.82
512	Mixed INT2/INT4	43.62	43.75	12,141.54	191.55	16.00
512	INT4	56.50	56.64	9,279.55	195.97	14.61
2048	Mixed INT2/INT4	137.01	137.11	15,118.97	180.10	16.04
2048	INT4	171.45	171.73	12,052.99	187.88	15.57
4096	Mixed INT2/INT4	268.66	269.06	15,336.09	169.55	18.01
4096	INT4	329.69	330.08	12,478.59	174.80	17.82

Measurement boundaries:

• TTFT: warm session, pretokenized input and initial GPU KV ready until the first greedy token is available on CPU. Includes binding, synchronized prefill, final-position logits D2H and CPU argmax. Excludes model loading, tokenization, initial KV reset, queuing and networking; not service-level TTFT.
• Prefill TPS: total input tokens / total synchronized ORT prefill run time. Excludes binding setup and token selection. The original graph still computes logits for every input position.
• Decode TPS: 1,270 tokens / total measured decode time across ten requests (127 decode steps per request). Excludes the first token produced by prefill; includes Python orchestration, binding, GPU execution, last-position logits transfer and CPU sampling.
• Memory: NVML process-resident bytes including weights, KV, workspace and retained ORT arena allocations. Requested sampling interval was 5 ms; observed maximum gaps were about 11-19 ms. These are sampled peaks and may miss brief allocations, not exact allocator high-water marks. Loaded-session residency was about 15.0 GiB for mixed versus 13.5-13.8 GiB for INT4. Device-wide peaks were 15.78/16.76/16.80/18.77 GiB for mixed and 14.57/15.37/16.33/18.57 GiB for INT4, respectively.
• Profiling and QMoE branch logging were disabled in measured runs. No full logits/KV CPU copies were performed; only the last-position logits were transferred for sampling. Both opt-in flags and the one-byte dense-dequantization limit were retained for both exports; enabled flags do not prove path selection.

Findings and limits:

• Mixed prefill throughput was 22.9%-30.8% higher, with 18.5%-22.8% lower TTFT P50, than this existing INT4 export. TTFT P50 reductions were 19.26%, 22.80%, 20.09%, and 18.51% for 128, 512, 2048, and 4096 input tokens, respectively, calculated from unrounded measurements as (INT4 TTFT P50 - mixed TTFT P50) / INT4 TTFT P50 * 100%. Mixed decode throughput was 2.3%-4.1% lower, not faster.
• Mixed sampled process peak memory was 0.20-1.39 GiB higher, despite the 32.07% smaller disk artifact. The cause has not been attributed; no runtime-memory reduction is claimed.
• Existing INT4 QMoE uses block size 32, while mixed uses 64. A separate INT4 diagnostic run confirmed grouped_moe for both prefill and decode (72 prefill calls and 216 decode calls across three four-token generations). Block size 32 fails the current packed INT GEMV eligibility gate. Mixed uses the packed prefill/GEMV routes verified above. This compares existing model configurations and execution paths, not isolated INT2-versus-INT4 kernels.
• Synthetic repeated text and differing greedy continuations can produce different MoE expert routing. Results do not establish real-workload throughput, tau2 task success or accuracy preservation.
• Measurements were sequential, not alternating A/B trials. GPU 0 had no other compute process, but other GPUs on the shared host were occupied and their utilization differed between model runs. Small decode differences require controlled repeated measurements before claiming a stable regression. P95 is based on only ten requests.

Local artifacts on the validation machine (not checked into this PR): gpt_oss_20b_pr4b_benchmark.py, pr4b-benchmark-20261001/input-*.json, pr4b-int4-benchmark-20261001/input-*.json, corresponding logs, and gpt-oss-20b-pr4b-combined.log, under /datadisks/disk1/jiafa/accuracy/. Export directories are int4_cuda_int2_int4_qmoe and int4_cuda_int4_qmoe under /sunghcho_data/jiafa/gpt-oss-20b-export/.

Example reproduction using the local harness; replace the model directory with the INT4 export for the comparison and repeat for each input length:

docker exec -w /tmp \
  -e CUDA_VISIBLE_DEVICES=0 \
  -e LD_LIBRARY_PATH=/datadisks/disk1/jiafa/cudnn9.7/lib:/datadisks/disk1/jiafa/cuda-12.8/lib64 \
  jiafa-dev /datadisks/disk1/jiafa/accuracy/onnxruntime/.venv/bin/python -u \
  /datadisks/disk1/jiafa/accuracy/gpt_oss_20b_pr4b_benchmark.py \
  --model-dir /sunghcho_data/jiafa/gpt-oss-20b-export/int4_cuda_int2_int4_qmoe \
  --input-length 512 --output-length 128 --warmups 3 --repeats 10 \
  --output /datadisks/disk1/jiafa/accuracy/rerun-mixed-512.json

Remaining Before Ready

• Complete plugin build, focused tests, regression and memcheck; run supported-platform CI. The October 1 full-model results above are non-plugin only.
• Extend the successful same-session prefill-to-decode smoke test to repeated transitions and concurrency/capture qualification.
• Run controlled alternating A/B measurements and matched-block-size comparisons, realistic workloads, accuracy/task-success evaluation, and exact peak-allocation attribution before broad performance claims or default dispatch changes.
• The full test suite and tau2-bench have not been run for this integration. The initial synthetic performance baseline above is complete but does not replace those qualifications.

[github-actions]

You can commit the suggested changes from lintrunner.

[Copilot]

Copilot review overview

🟡 Changes recommended

The two new packed-prefill files need the repository-standard copyright and MIT license headers.

Review effort: Balanced
Findings: 2

Open (2)

• Add copyright and MIT license header to source file · New
• Add copyright and MIT license header to header file · New

What changed in this PR

Enables default packed INT2/INT4 CUDA QMoE prefill on SM80/SM86.

Changes:

• Adds packed prefill routing, workspace, FC1/FC2, activation, and finalization.
• Extends INT2 grouped GEMM support to SM86.
• Adds default-dispatch, fallback, routing, runtime-scale, FP16, and BF16 tests.

File	Description
onnxruntime/​test/​contrib_ops/​moe_test.cc	Adds packed-prefill integration coverage.
onnxruntime/​test/​contrib_ops/​cuda_kernels/​moe_gemm_int2_test.cc	Tests SM80/SM86 grouped GEMM support.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_quantization.h	Stores prefill enablement state.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_quantization.cc	Integrates default packed-prefill dispatch and prepacking.
onnxruntime/​contrib_ops/​cuda/​llm/​moe_gemm/​moe_int2_prefill.h	Declares the packed-prefill interface.
onnxruntime/​contrib_ops/​cuda/​llm/​moe_gemm/​moe_int2_prefill.cu	Implements the packed-prefill pipeline.
onnxruntime/​contrib_ops/​cuda/​llm/​moe_gemm/​moe_gemm_kernels_fp16_uint2.cu	Enables grouped INT2 GEMM on SM86.
onnxruntime/​contrib_ops/​cuda/​llm/​moe_gemm/​moe_gemm_int2.h	Documents the supported architectures.

---

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

[Copilot]

Copilot review overview

🔵 Needs a closer look

The path is default-enabled despite explicitly outstanding platform, plugin, concurrency, and capture qualification.

Review effort: Balanced
Findings: None

Resolved since last review (2)

• Add copyright and MIT license header to header file
• Add copyright and MIT license header to source file

Previously missed (1)

In code that hasn't changed since last review

Enable new dispatch path only after qualification gates complete

onnxruntime/​contrib_ops/​cuda/​moe/​moe_quantization.cc:391

The PR's own “Remaining Before Ready” section says supported-platform/plugin qualification and repeated-transition/concurrency/capture validation are still required before default dispatch changes, but this default enables the new path for every eligible session now. Keep the path opt-in until those stated gates complete; otherwise unqualified SM86/plugin workloads can silently enter a new kernel path with no fallback after dispatch.

[Tianlei Wu (tianleiwu)]

Reviewed the current head, including packed-prefill eligibility, scratch budgeting, prepack/fallback ownership, scale preparation, routing maps, SwiGLU biases, fixed FC2 tactic availability, and the added tests. I did not establish a blocking correctness defect. The latest scratch-budget and channel-varying bias changes address the existing concerns, so I have not duplicated those threads. One nonblocking performance suggestion is inline. Validation: git diff --check passed; I did not build or run runtime tests for this head, and CUDA CI is still in progress.

[Tianlei Wu (tianleiwu)]

Re-reviewed head e0b833a. The FC2 runner now uses the supplied SM/SM-count metadata, which addresses my earlier comment. I found no correctness defect in the new prefill pipeline (routing maps, workspace layout/aliasing, SwiGLU bias handling, SM80+ gate, Ampere FC2 fallback on SM90+).

I left one inline comment on the 256 MiB workspace cap. When a prompt goes over the cap, the code falls back to the dense dequant path. With the default 1 GiB dense-scratch limit, that path cannot run GPT-OSS-20B-sized experts, so prompts longer than about 2.9k tokens fail. This is not a regression from main, but I suggest tiling the packed path by rows instead of falling back.

Validation: static review only; no local build or runtime tests.