[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