onnxruntime

1e1f574c - QMoE: prepack int4/int8 expert weights in PrePack hook (symmetric with MatMulNBits) (#28749)

QMoE: prepack int4/int8 expert weights in PrePack hook (symmetric with MatMulNBits) (#28749) ## Summary This PR lets the CUDA `com.microsoft::QMoE` operator prepack **raw** int4/int8 expert weights into the CUTLASS `fpA_intB` layout **inside ORT's `PrePack()` hook**, instead of requiring callers to run the layout transform offline via `pack_weights_for_cuda_mixed_gemm`. This makes integer QMoE symmetric with `MatMulNBits::PrePack_B`, and lets exporters ship the schema-conformant `[E, N, K/pack]` quantized weights produced by `quantize_matmul_{4,8}bits` directly, with no offline pre-pack step. The behaviour is **opt-in and backward compatible**: a new `weights_prepacked` attribute is a tri-state and defaults to `-1` (auto), which the CUDA EP treats as "already CUTLASS-prepacked" (today's behaviour). `1` forces the prepacked interpretation explicitly, and only `weights_prepacked=0` triggers the new in-`PrePack` layout transform. ## What changed - **New `weights_prepacked` attribute** on the QMoE schema (tri-state, default `-1`/auto). `-1` lets each execution provider pick its own backward-compatible default; the CUDA EP treats auto as prepacked. `1` = the int4/int8 `fc1`/`fc2` initializers are already in the CUTLASS `fpA_intB` layout (today's behaviour). `0` = the initializers are raw `[E, N, K/pack]` tensors and the kernel runs the layout transform itself in `PrePack()`. - **`PrePackIntExpertWeights`** — loops over the `E` experts and applies the per-expert transpose + CUTLASS `fpA_intB` row-permutation / column-interleave / bias / pair-interleave transform on the GPU, mirroring `pack_weights_for_cuda_mixed_gemm`. Architecture-aware packing per `docs/contrib_ops/cuda/moe_qmoe.md` §7 (SM90 is its own layout group; all other supported arches share the SM80 layout; SM75+ required). - **`PrePack()` dispatch** for the int weight slots (2 and 5) when `quant_type == "int"` and `weights_prepacked == 0`. The source initializers are released after their shapes are cached (`fc1/fc2_weights_shape_`), so peak weight memory stays ~1×. - **`ComputeInternal`** prefers the prepacked GPU buffers when the PrePack hook populated them (gated on `int_weights_consumed_by_prepack`), and otherwise falls through to the raw initializer pointers (e.g. for sessions that set `session.disable_prepacking`). ## Schema note This **does add a schema attribute** (`weights_prepacked`) to QMoE. It is backward compatible because the default (`-1`/auto) is interpreted by the CUDA EP as prepacked, preserving the existing offline-prepacked behaviour, but it is a schema surface-area change. ## Diff scope | File | Change | |---|---| | `onnxruntime/core/graph/contrib_ops/contrib_defs.cc` | New `weights_prepacked` schema attribute + docs | | `onnxruntime/contrib_ops/cuda/moe/moe_quantization.h` | New private method + prepack buffer / cached-shape members | | `onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc` | `PrePackIntExpertWeights` + PrePack dispatch + ComputeInternal hookup | | `onnxruntime/test/python/transformers/test_qmoe_cuda.py` | CUDA smoke test for the raw-weight `weights_prepacked=0` path | FP4 / FP8 / WFP4AFP8 paths are untouched, and there is no behaviour change for callers that pre-prepacked their weights. ## Testing - `onnxruntime_providers_cuda` builds and links cleanly (nvcc 13.2 / sm_90). - `TestQMoEIntPrePackSmoke` (`test_qmoe_cuda.py`) builds a QMoE graph with raw int4 weights and `weights_prepacked=0`, runs it through the CUDA kernel, and asserts the output is finite with a plausible magnitude. Verified on H200 (SM90); node placement confirmed on `CUDAExecutionProvider` via profiling. - Existing int4 QMoE parity tests (`phi3` / `swiglu`, fp16) pass on CUDA — no regression in the default `weights_prepacked=-1` (auto, prepacked) path. > Note: this is a smoke test, not a numerical parity check. The existing offline > pre-pack harness hardcodes `force_arch=80` and produces incorrect output on > SM≥90, so a bit-parity comparison against it is intentionally omitted until > that harness honours the runtime SM. --------- Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com> Co-authored-by: Justin Chu <11205048+justinchuby@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>