[CUDA] Add fast INT3 QMoE kernels and offline SM80 INT2/INT4 prepack support (#33569)
## Description
Adds fast CUDA execution kernels for `INT3` and mixed-width `(FC1, FC2)
∈ {2, 3, 4} × {2, 3, 4}` `com.microsoft::QMoE` projections, and extends
the SM80 (`fpA_intB`) packed `INT2` and mixed `INT2`/`INT4` execution
paths (`moe_gemv` decode and `RunInt2MoePrefill` grouped GEMM prefill)
and single-CTA fused routing to support offline prepacked weights
(`weights_prepacked = 1`), `block_size = 32` 2-scale GEMV
dequantization, `num_experts <= 512`, `k = 10`, `expanded_num_rows <=
80`, and `kFusedRouteMaxRows = 8` ($S \le 8$, 7-draft DFlash2
speculative decoding).
## Summary of Changes
### Fast `INT3` & Mixed-Width `QMoE` CUDA Kernels
| File | Change |
|------|--------|
| `onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h` | Declares
`Int3Projection`, `Int3MoeParams`, `LaunchInt3Moe`, and
`ValidateInt3MoeContent`. |
| `onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu` | Implements
`Unpack8CenteredCodes` / `Unpack8CenteredCodes128` (aligned 32-bit
`__ldg` + hardware `__funnelshift_r` 24-bit 8-element `INT3` unpacking),
`FusedFc1FastKernel` (shared-memory `s_input` staging + 2-column `(gate,
up)` register reuse per warp), `Fc2RouteFastKernel` (parallel-route
`gridDim.y = rows * top_k` launch + shared-memory `s_act` staging +
16-lane half-warp 2-column tiling), and `FinalizeRoutesFastKernel`. |
| `onnxruntime/contrib_ops/cuda/moe/moe_quantization.h` | Adds
`ComputeInt3`, `has_int3_`, and per-instance `int3_content_validated_`
atomic flag. |
| `onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc` | Enables
`expert_weight_bits` and `fc{1,2,3}_expert_weight_bits` in `{2, 3, 4,
8}`, dispatches `INT3` projections to `ComputeInt3`, and caches one-time
content validation while skipping host synchronization during CUDA Graph
capture. |
### Offline SM80 `INT2` / Mixed `INT2`/`INT4` Prepacking
(`weights_prepacked = 1`), `block_size = 32` Decode `moe_gemv`, & $S \le
8$ Routing
| File | Change |
|------|--------|
| `onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv.h` | Raises
`kMaxProfiledExpandedRows` from `8` to `80` so $S \le 8, k = 10$
(`expanded_num_rows <= 80`) dispatches to `moe_gemv`. |
| `onnxruntime/contrib_ops/cuda/moe/qmoe_fused_route.cuh` | Widens
`kFusedRouteMaxRows` from `4` to `8` (`kFusedRouteThreads = 256`,
`s_sorted[256]`) so single-CTA `FusedSoftmaxTopKMapsKernel` handles up
to 8 rows ($S = 8$ DFlash2 verification). |
| `onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv.cu` | Adds
`dequantize_and_pack_to_vec2<Details, GroupSize, TypeA>` and
`scales_iterator_hi` across all 4 `moe_gemv` kernels so `Int2DetailsW`
(`StepK = 64`) supports 2-scale `GroupSize == 32` dequantization
alongside `GroupSize == 64` and `128`. |
| `onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc` | Accepts
`weights_prepacked = 1` for `INT2` (`(2, 2)`) and mixed `INT2`/`INT4`
(`(2, 4)`, `(4, 2)`), consuming offline SM80 `fpA_intB` weight tensors
directly in-place (`packed_int_fc1`, `packed_int_fc2`) for `block_size ∈
{32, 64, 128}`, `num_experts <= 512`, and `k ∈ {1, 2, 4, 6, 8, 10}`
while releasing untransposed `[E, N, K_blocks]` scales in `PrePack`. |
| `docs/contrib_ops/cuda/moe_qmoe.md` | Documents `INT3` QMoE CUDA
execution and offline SM80-prepacked `INT2` / mixed `INT2`/`INT4`
support. |
| `onnxruntime/test/contrib_ops/moe_test.cc` | Adds
`QMoETest_MixedWidthCudaPackedDecodeBlock32` and
`QMoETest_Int3AndMixedWidthCudaSwiGLU` unit tests. |
## Testing
- Verified `QMoETest_MixedWidthCudaPackedDecodeBlock32` and
`QMoETest_Int3AndMixedWidthCudaSwiGLU` across `(3, 3)`, `(3, 4)`, `(4,
3)`, `(2, 3)`, `(2, 4)`, `(4, 2)`, and `(2, 2)` with `block_size ∈ {32,
64}` and `num_rows ∈ {1, 4, 64}`.
- Evaluated $S = 8$ (7-draft DFlash2) and 9 lower-precision 48-layer
512-expert `QMoE` model packages on the full 800-sample MMLU-Pro
benchmark (`max_new = 8192`) on NVIDIA H200 GPUs (`361.14 tok/s` median,
`411.32 tok/s` P90, `453.30 tok/s` max at $S = 8$).
---------
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
onnxruntime
bd6010c0 - [CUDA] Add fast INT3 QMoE kernels and offline SM80 INT2/INT4 prepack support (#33569)
