SIGN IN SIGN UP

[CUDA] Enable SM80+ FP16/BF16 mixed INT2/INT4 QMoE prefill by default (#33005)

### 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:
```bash
./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:

```bash
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.
D
David Fan committed
2cdff3ca8e3b6d18dc8918e0e9701b03f1576fa7
Parent: e1b8996
Committed by GitHub <noreply@github.com> on 10/2/2026, 5:02:07 AM