QMoE: prepack int4/int8 expert weights in PrePack hook (symmetric with MatMulNBits)

[justinchuby]

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 defaults to 1 (legacy — weights are already CUTLASS-prepacked), and
only weights_prepacked=0 triggers the new in-PrePack layout transform.

What changed

• New weights_prepacked attribute on the QMoE schema (default 1).
  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) preserves 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 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.

[justinchuby]

Built locally and ran the smoke test on H200 (sm_90):

$ pytest test_qmoe_cuda.py::TestQMoEIntPrePackParity -v
test_int4_swiglu_interleaved_small  PASSED
test_int4_swiglu_interleaved_medium PASSED

A few notes from local verification:

1. is_packed = false after PrePackIntExpertWeights — same trade-off the existing wfp4afp8 weight branch uses (line 970). moe_helper::CheckInputs still needs the original weight tensor's shape on every Compute call to infer moe_params, so the initializer has to stay alive. The prepacked bytes live in the new packed_fc{1,2}_weights_ buffers and the compute path prefers them over fc{1,2}_experts_weights->DataRaw().

2. Test is a smoke test, not bit-parity — first version I tried compared against the existing offline pack_weights_for_cuda_mixed_gemm path, but that comparison is invalid on SM>=90: the existing test harness in test_qmoe_cuda.py hardcodes force_arch=80, and on H100/H200 the existing test_swiglu_qmoe_parity_* cases all fail with max-diff > 1.0 on plain main (pre-dating this change). Filed as a separate observation if it's useful — looks like the harness needs to honour runtime SM.

3. Build verification — full library link succeeded on the dev box after working around a separate CUDA 13.2 / CCCL header bug in bias_softmax_impl.cu (unrelated to this PR). CI should hit no such issues.

[tianleiwu]

Review summary

Moving the CUTLASS fpA_intB layout transform for quant_type == "int" QMoE weights into the PrePack hook (mirroring MatMulNBits::PrePack_B) is a good direction, and the mechanics faithfully mirror the validated pack_weights_for_cuda_mixed_gemm pybind path. One blocking concern plus a couple of follow-ups.

High priority

Unconditional re-prepack is not backward compatible (no raw-vs-packed marker). The new dispatch runs PrePackIntExpertWeights for every int QMoE on slots 2/5, with no schema flag or version guard. Existing tooling — including this file's own quant_dequant_blockwise/preprocess_weights_for_mixed_gemm in test_qmoe_cuda.py — already emits weights in the CUTLASS layout and stores them under the logical [E, N, K/pack] shape that moe_helper::CheckInputs validates. Because raw and prepacked byte counts are identical (E*N*K/pack == E*K*N/pack), the declared shape, dtype, and size are indistinguishable. On the default path (prepacking enabled) those models get prepacked a second time and silently produce garbage — the same silent-failure class this PR set out to remove. The description's claim of "no behaviour change for callers that pre-prepacked their weights" only holds when session.disable_prepacking is set.

A safe fix needs an explicit signal that the weights are raw (e.g. a weights_prepacked attribute defaulting to legacy/prepacked, a com.microsoft opset bump, or a distinct marker). Until then this should be opt-in and the hard break documented. The existing prepacked-weight parity tests (test_qmoe_cuda.py ~L154-298) are not updated and would regress on an SM where they currently pass once the hook double-prepacks; the new test only covers the raw path so it won't catch that.

Suggestions

• Persistent weight memory ~2x for int QMoE. Keeping is_packed = false (so CheckInputs can still read the original shape) while also allocating persistent packed_fc{1,2}_weights_ means the original int weights and the prepacked int weights both stay resident ~= 2x the dominant weight memory for both FC layers. MatMulNBits::PrePack_B avoids this via is_packed = true. Consider caching just the (E, N, K) shape and releasing the source, or documenting the cost. Transient PrePack overhead per FC: E*N*K/pack host->device staging (CPU initializers only) + N*K/pack per-expert transpose scratch + 128 B perm map, freed after the sync.
• SM coverage. The offline packer restricts force_arch to {75,80,90} and warns arch>90 falls back to 80; the in-kernel path passes sm_ straight through (more correct, matches the device). Please confirm preprocess_weights_for_mixed_gemm_cuda is valid on SM100/120 if int QMoE is expected there, else add a guard/diagnostic.

Nitpick

• ORT_ENFORCE(bits != 4 || k % 2 == 0, ...) is always true for bits == 4 (k = k_packed*2).

Nice work on the motivation (decoupling QMoE export from a CUDA-built ORT) and the clear writeup.

[justinchuby]

Thanks for the thorough review! Addressed all 4 points in 2fcb940:

2. SM coverage — added the same SM75/80/90 fallback the offline pack_weights_for_cuda_mixed_gemm binding uses (clamp >90 → 80, also round SM86/89 → 80) inside PrePackIntExpertWeights before calling preprocess_weights_for_mixed_gemm_cuda. Defined layout on Blackwell+ via the Ampere tile tables, matching the offline packer's behaviour.

3. Nit — dropped the redundant ORT_ENFORCE(bits != 4 || k % 2 == 0) (tautology since k = k_packed * pack_factor).

4. 2x memory cost — documented inline. Kept is_packed = false so CheckInputs can still read the source shape per Compute. Folding the shape into member variables to release the source (matching MatMulNBits) is straightforward but touches all the other QMoE paths through CheckInputs; left as a clean-up follow-up to keep this PR focused.

H200 re-verification with the new attribute:

test_int4_swiglu_interleaved_small  PASSED
test_int4_swiglu_interleaved_medium PASSED

Test was also updated to set weights_prepacked=0 explicitly.

[justinchuby]

Followed up on the 2× weight memory concern in 5e1491c:

• Added TensorShape fc1_weights_shape_ / fc2_weights_shape_ members captured at PrePack time.
• Switched the CheckInputs call to the existing TensorShape* overload when the source is consumed, so we never need the live tensor for shape validation.
• PrePackIntExpertWeights now leaves is_packed = true (source initializer freed). Net persistent weight memory back to ~1× the int4/int8 footprint, matching the MatMulNBits / FP4-fallback profile.
• Compute path guards context->Input<Tensor>(2)/(5) and the check_weight_type dtype assertions for the consumed path.

H200 verification:

test_int4_swiglu_interleaved_small  PASSED   (opt-in raw path)
test_int4_swiglu_interleaved_medium PASSED   (opt-in raw path)
test_swiglu_qmoe_parity_0           PASSED   (legacy default path, max_diff ~0.001)

[tianleiwu]

Review summary

Thanks for reworking this into an opt-in path — the new weights_prepacked attribute (default 1 = legacy CUTLASS-prepacked) plus releasing the source initializer via cached fc1/fc2_weights_shape_ resolves both of my earlier blockers (silent double-prepack on existing models, and the ~2x weight-memory cost). Those two threads are resolved.

One new correctness issue remains in the compute path (inline), plus one design question:

Inline (correctness)

• ComputeInternal unconditionally overrides the weight pointers in the else if (is_int && !weights_prepacked_) branch, even when the prepack buffer is null. See inline comment.

packing_sm clamp vs. runner layout (question)

In PrePackIntExpertWeights, packing_sm is clamped (>90 → 80, and 86/89 → 80) before calling preprocess_weights_for_mixed_gemm_cuda. This diverges from MatMulNBits::PrePack_B, which passes sm_ unmodified. The risk: the CUTLASS MoE runner at compute time selects its expected weight layout from the actual sm_. If sm_ is e.g. SM100/120 (Blackwell) and PrePack lays the bytes out for SM80, the packer/runner layouts can disagree and silently produce garbage rather than erroring. Either (a) confirm the runner consumes the SM80 layout on those arches, or (b) make the clamp explicit/shared with the runner-side arch selection so the two cannot drift. Worth a short comment documenting that the compute side relies on the same clamped arch.

Nit

• The trailing cudaStreamSynchronize(stream) at the end of PrePackIntExpertWeights is redundant: preprocess_weights_for_mixed_gemm_cuda already synchronizes the stream internally on every per-expert call. Harmless, but can be dropped.

The schema doc, opt-in default, and the new parity smoke test all look good. Requesting changes only for the compute-path null-override.

[tianleiwu]

Correctness: this else if guard is weaker than the int_weights_consumed_by_prepack guard used above (is_int && !weights_prepacked_ && packed_fc1_weights_ != nullptr). When weights_prepacked_ == false but prepacking is disabled at the session level (session.disable_prepacking), PrePack never runs, so packed_fc{1,2}_weights_ stay null. In that case int_weights_consumed_by_prepack is false, the raw initializers are correctly read into fc1/fc2_weight_data above — but this branch still fires (is_int && !weights_prepacked_ is true) and overwrites them with packed_fc1_weights_.get() == nullptr, so the runner receives null weight pointers (crash / garbage). This also contradicts the PR description's stated "fall-through to the raw initializer ... for sessions that disable prepacking."

Suggest gating on the same condition, e.g.:

} else if (int_weights_consumed_by_prepack) {
  fc1_weight_data = packed_fc1_weights_.get();
  fc2_weight_data = packed_fc2_weights_.get();
}

so that when the prepack buffers are absent the code keeps the raw-initializer pointers.

[tianleiwu]

See https://github.com/microsoft/onnxruntime/blob/9f81e2dfc2d46e5b2a590a31129faa81d74011fa/docs/contrib_ops/cuda/moe_qmoe.md#7-cross-architecture-packing-compatibility

So packing_sm = (sm_ == 90) ? 90 : 80;

Skip packing when sm_ < 75 since those GPUs are not supported.

[Copilot]

Pull request overview

This PR extends the CUDA com.microsoft::QMoE integer-quantized path to optionally prepack raw int4/int8 expert weights inside ORT’s PrePack() hook (mirroring MatMulNBits), controlled via a new weights_prepacked attribute. It also adds a Python smoke test to validate that raw weights can execute through the CUDA QMoE kernel without producing NaN/Inf.

Changes:

• Add a new weights_prepacked attribute to the QMoE schema (default 1) to distinguish offline-prepacked vs raw quantized expert weights.
• Implement PrePackIntExpertWeights and wire it into QMoE’s CUDA PrePack() + ComputeInternal to use cached prepacked GPU buffers.
• Add a CUDA Python test that builds a minimal QMoE graph with raw int4 weights and checks output sanity.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 9 comments.

File	Description
onnxruntime/test/python/transformers/test_qmoe_cuda.py	Adds a CUDA test case for running QMoE with raw int4 weights and weights_prepacked=0.
onnxruntime/core/graph/contrib_ops/contrib_defs.cc	Adds the weights_prepacked schema attribute and documentation for QMoE.
onnxruntime/contrib_ops/cuda/moe/moe_quantization.h	Adds members/state and declares helper for int-weight prepacking.
onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc	Implements int expert-weight prepacking and uses prepacked buffers during compute.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[justinchuby]

Thanks for the automated pass. Summary of how I addressed the 9 comments (pushed in 3cbbf51):

Fixed

• PR description / schema note (Add doxygen generated website for the project #6): Rewrote the description; it now explicitly states weights_prepacked is a (backward-compatible) schema attribute addition and lists all changed files.
• Schema doc wording (Set up CI with Azure Pipelines #7): Changed "raw, row-major [E, N, K/pack]" → "raw, un-prepacked [E, N, K/pack]" so it no longer conflicts with the column-major schema docstring.
• Test name + comment (Fix build #8, Enable Mac pipeline #9): Renamed TestQMoEIntPrePackParity → TestQMoEIntPrePackSmoke and rewrote the module comment to describe the actual smoke test (single graph, finite + plausible-magnitude assertions, no bit-parity).

Declined, with rationale

• SM75 layout (Add doxygen web for the project #5): This is incorrect. getLayoutDetailsForTransform has a separate arch<80 branch, but the underlying cutlass::gemm::kernel::LayoutDetailsB for uint4b_t/uint8_t/half/bf16 is a single specialization gated on Arch::kMinComputeCapability >= 75 (mixed_gemm_B_layout.h:49–105), so SM75 and SM80 produce an identical packed layout. This matches docs/contrib_ops/cuda/moe_qmoe.md §7 (SM75 is in the universal Group A). packing_sm = (sm_ == 90) ? 90 : 80 is correct.
• Per-weight prepack tracking (Set up CI with Azure Pipelines #1–Incremental updates. #4): For a valid int QMoE, slots 2 (fc1) and 5 (fc2) are both constant initializers and are always prepacked together in the same PrePack pass; the "only one prepacked" case is unreachable for valid models (a non-initializer weight would already be rejected). int_weights_consumed_by_prepack requiring packed_fc1_weights_ != nullptr is sufficient, and the single-flag design was the approach agreed in the previous review round. Splitting into per-weight flags adds branching for a state that cannot occur.

Build + tests: onnxruntime_providers_cuda builds/links cleanly (nvcc 13.2, sm_90); TestQMoEIntPrePackSmoke and the existing int4 QMoE parity tests pass on H200, with the QMoE node confirmed on CUDAExecutionProvider via profiling.