[Tianlei Wu (tianleiwu)]

Summary

This PR consolidates and optimizes CPU/GPU hybrid expert offload for com.microsoft::MoEPacked (GGUF sub-byte block-packed experts) and com.microsoft::QMoE (RC5 mixed-precision INT4/INT3/INT2 experts) across 16 GB (RTX 5060 Ti), 24 GB (RTX 4090), 32 GB (RTX 5090), and H200 GPUs.

Stacked On

• [CUDA] Improve QMoE offload decode throughput via AVX2 spill elimination and working-set retention #33935 ([CUDA] Improve QMoE offload decode throughput via AVX2 spill elimination and working-set retention)
• [CUDA] Tail expert substitution for QMoE offload decode (stacked on #33935) #33944 ([CUDA] Tail expert substitution for QMoE offload decode (stacked on #33935))
• Support direct external initializer buffers in AddExternalInitializers #33877 (Support direct external initializer buffers in AddExternalInitializers)
• [CUDA] Add GatherPacked, MatMulPacked, and MoEPacked contrib ops for block-packed weights #33953 ([CUDA] Add GatherPacked, MatMulPacked, and MoEPacked contrib ops for block-packed weights)

Key Changes

1. End-to-End CPU/GPU Hybrid Offload for com.microsoft::MoEPacked (session_state.cc, moe_packed.h, moe_packed.cu, moe_packed.cc):
   • Places MoEPacked expert weight initializers (gate_weights, up_weights, down_weights) on CPU (zero-copy file mapping) and implements IMoeOffloadNode admission with QMoESessionController.
   • Implements CUDA-graph-capturable hybrid doorbell decode (MoEPackedOffloadClassifyAndBuildTilesKernel, MoEPackedDoorbellWaitKernel, MoEPackedOffloadFinalizeVec8Kernel, and LaunchQMoETailSubstitution integration) and 8,192-row double-buffered 32-expert wave-streaming prefill with D2D adaptive hot-expert promotion.
2. Lock-Free qmoe::SpinPool & Register-Accumulator AVX2+FMA+F16C GGUF Cold-Expert Kernels (moe_packed_cpu_avx2.cc, moe_packed.cc, cmake/onnxruntime_providers_cuda*.cmake):
   • Backs MoEPackedParallelPool with qmoe::SpinPoolRegistry (qmoe_spin_pool.h) and pins MoEPackedDoorbellService to CpuForWorker(0).
   • The AVX2+FMA+F16C kernels live in moe_packed_cpu_avx2.cc (ISA enabled per function; /arch:AVX2 for that file on MSVC x64) and are dispatched after qmoe_shared::CpuHasAvx2FmaF16c(), with a scalar fallback.
   • Implements pure register-accumulator 2-channel (GgufDualDot2ChAvx2Tpl<RT>) and 4-column (IQ4_NL/Q2_0) AVX2+FMA+F16C kernels covering IQ2_S, IQ2_XXS, IQ2_XS, IQ3_XXS, IQ4_XS, IQ4_NL, Q2_0, Q5_0, and mixed gate_type != up_type layers, cutting iq3_xxs CPU cold-expert latency 4.8x (245.5 us/exp -> 49.2 us/exp) and boosting iq3_xxs decode throughput 4.1x (16.2–20.8 tok/s -> 67.2–81.7 tok/s lossless, 179.96 tok/s with tail substitution) and iq1_s decode throughput up to 142.91 tok/s lossless (261.65 tok/s with tail substitution) on RTX 5090.
3. QMoE INT2 Prefill & AVX-512 Cold-Expert Optimizations + 16 GB / 24 GB / 32 GB WDDM Reserve Calibration (qmoe_offload.cc, moe_gemv_tc.cuh, qmoe_shared_layout_int2_avx512.cc, qmoe_session_controller.h):
   • Dispatches MoeGemvTcInt2Kernel on sorted expert_map prefill waves (+24% to +30% faster 8K/32K prefill on 2b375, reaching 2,560 tok/s).
   • Extends Int2Group8Avx512 to RT = 1..4 with factored scale FMA.
   • Calibrates WDDM post-session reserves across 16 GB, 24 GB, and 32 GB GPUs for both QMoE and MoEPacked and auto-detects 8-byte vs. 4-byte STRP v1 popularity profiles.
4. Numerical Parity and Run-Control Lifetime for MoEPacked Hybrid Decode (moe_packed.cc, moe_packed.cu, moe_packed.h, test_gguf_packed_ops.py):
   • Keeps the CPU cold-route partial sum in FP32 (MoEPackedDecodePinned::host_cold_output), so the hybrid combine rounds to BF16/FP16 once, as the GPU-only finalize does. Rounding the partial sum to the activation type first rounded the output twice.
   • Mirrors the CUDA FP16 rounding points on the CPU path (gate/up dot products, SwiGLU output, and per-route FC2 output), as the BF16 path already did.
   • Registers a per-node QMoEInvocationRecord with the kernel stream, and with the replay stream through the graph-replay hook, so the controller synchronizes it before EndRun() clears the warm-up flag. Previously, with device-resident outputs, 296 of 300 warm-up decode runs in an all-cold session computed the cold experts on the CPU; now none do.
   • Resolves the CPU pool size once per node instead of calling std::thread::hardware_concurrency() (a sysfs read, about 2.6 us on Linux) on every decode ParallelFor.

---

Accuracy & Performance Impact of Tail Expert Substitution (threshold = 0.02, min_rank = 2) on 2b375 (RTX 5090 32 GB)

Evaluated on qwen_flash_next_int4_int8_rc5_contrib_dflash2_qmoe_2b375_x02_y14_z32_dp23_b64_unified (xhigh reasoning, greedy decoding):

• MMLU-Pro (400 samples, max_new_tokens = 32,768):
  • Without TailSub (thr = 0.0): 82.75% (331/400, strict 81.75%), median 162.09 tok/s, token-weighted 230.99 tok/s
  • With TailSub (thr = 0.02, r = 2): 83.25% (333/400, strict 82.25%) (+0.50 pp, 12W / 10L, McNemar p = 0.8318), median 255.64 tok/s (1.58x faster), token-weighted 328.52 tok/s (1.42x faster)
• LiveCodeBench v6 (175 problems, 256K = 262,144 limit):
  • Paired 161 / 175 problems on RTX 5090: 83.85% (135/161) With TailSub vs. 83.23% (134/161) Without TailSub (+0.62 pp, 13W / 12L, McNemar p = 1.0000), median 217.27 tok/s (1.54x faster), token-weighted 248.82 tok/s (1.65x faster)
  • Full 175 / 175 problems (at 256K): 77.14% (135/175, 43/43 = 100% Easy, 42/52 = 80.8% Medium, 50/80 = 62.5% Hard, conditional 88.24% when code is emitted) vs. 82.86% (145/175) lossless H200 reference.

---

CPU vs. CUDA Numerical Parity of MoEPacked Hybrid Decode (H200)

Fraction of output elements that differ from the GPU-only MoEPacked output (session.moe.offload.mode=disabled), over eight GGUF format mixes that cover every CPU cold-expert kernel family and decode batches of 1, 3, and 16 rows:

Placement	Activation	Before item 4	After item 4
All experts on CPU	FP16	54–65%	0–6% (summation order only)
Half of the experts resident	FP16	50–63%	0–1.5%
All experts on CPU	BF16	38–77%	Unchanged (see below)

With BF16, the CPU path matches a NumPy emulation that keeps dequantized weights in FP32 bit-for-bit, and the CUDA path matches one that rounds them to BF16 bit-for-bit. The remaining difference comes from the CUDA BF16 MMA kernels rounding GGUF-dequantized weights to BF16. The CPU kernels keep the more precise FP32 weights: emulating the rounding would slow the AVX2 loops, and results would still depend on expert placement because the summation order differs.

Testing

• python -m pytest onnxruntime/test/python/transformers/test_gguf_packed_ops.py: 39 passed (857 subtests) on H200 (sm90). The offload tests also pass with ORT_QMOE_SHARED_CPU_THREADS=1 and 7.
• test_moe_packed_cuda_offload_cpu_kernels_match_reference runs every CPU cold-expert kernel family with both activation types against a reference with FP32 dequantized weights. test_moe_packed_cuda_offload_mixed_placement_matches_gpu compares FP16 hybrid decode with the GPU-only output, using aligned and offset buffers. Both fail before item 4 and pass now.

[Copilot]

🟡 Changes recommended

Unresolved critical defects can overflow scratch and stack buffers, trigger invalid CUDA accesses, and silently produce incorrect expert outputs.

16 open issues

• AVX2/FMA/F16C compilation can crash on unsupported x86 CPUs · New
• BF16 conversion mishandles NaN payloads as negative zero · New
• Staging singleton reuses CUDA resources across devices · New
• Registration can race with worker shutdown and resource teardown · New
• CUDA Graph replay uses a stale warmup latch for live-session filtering · New
• Backend ignores host-cache limits and reads complete initializer banks · New
• Unbounded inter_size can overflow the fallback dequantizer's stack buffer · New
• Prefill allocations exceed device and pinned-memory budgets · New
• Small-route layout exceeds the workspace sizing helper's allocation · New
• Finalizers read uninitialized contributions from pruned routes · New
• Cold-input copy uses uint4 without checking source alignment · New
• Offload decode finalizer misaligns element-aligned output stores · New
• Dispatchers ignore BF16 tensor-core support and launch empty MMA kernels · New
• FP16 offload always stages in shared memory, exceeding device limits · New
• Prefill finalizer needs an alignment check before uint4 stores · New
• Offload test lacks an SM80 capability guard · New

What changed in this PR

This PR extends ONNX Runtime’s existing QMoE offload infrastructure to MoEPacked, enabling CPU/GPU hybrid execution of GGUF-packed experts, and consolidates related offload optimizations.

Changes:

• Adds CPU placement, offload admission, hybrid decode, streamed prefill, and adaptive expert promotion for MoEPacked.
• Optimizes GGUF CPU kernels and QMoE INT2 execution, and adjusts memory reserves and popularity-profile handling.
• Adds offload parity coverage and expands INT2 overflow regression coverage.

File	Description
onnxruntime/​test/​python/​transformers/​test_gguf_packed_ops.py	Adds MoEPacked offload parity tests across row counts and activation types.
onnxruntime/​test/​contrib_ops/​qmoe_shared_layout_test.cc	Extends INT2 overflow regression coverage to multiple rows.
onnxruntime/​core/​providers/​cuda/​qmoe_session_controller.h	Adjusts memory reserves, profile parsing, admission ordering, and host-page eviction.
onnxruntime/​core/​framework/​session_state.cc	Validates MoEPacked offload eligibility and places expert initializers on CPU.
onnxruntime/​contrib_ops/​cuda/​moe/​qmoe_offload.cc	Optimizes compact CPU output writes, INT2 prefill dispatch, and Windows prefetching.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_packed.h	Declares offload state, interfaces, limits, and CUDA launch APIs.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_packed.cu	Adds hybrid decode classification, synchronization, grouped execution, and output combination.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_packed.cc	Implements CPU expert kernels, offload lifecycle, streamed prefill, and promotion.
onnxruntime/​contrib_ops/​cuda/​llm/​moe_gemm/​moe_gemv.cu	Handles the sorted-expert-map marker before fallback dispatch.
onnxruntime/​contrib_ops/​cuda/​llm/​moe_gemm/​moe_gemv_tc.cuh	Enables INT2 Tensor Core dispatch for larger sorted prefill waves.
onnxruntime/​contrib_ops/​cpu/​moe/​qmoe_shared_layout_int2_avx512.cc	Extends the eight-column INT2 kernel to four-row tiles.
cmake/​onnxruntime_providers_cuda.cmake	Adds AVX2/FMA/F16C compilation flags for MoEPacked host code.
cmake/​onnxruntime_providers_cuda_plugin.cmake	Applies corresponding compilation flags to the CUDA plugin build.

_{Review effort: Balanced}

[Baiju Meswani (baijumeswani)]

This PR lets block-packed GGUF MoE models keep cold experts on the CPU and run resident experts on the GPU, with streamed prefill and CUDA-graph decode. It also brings in QMoE INT2 and offload improvements.

The CPU fallback and ISA isolation, rounding fixes, FP32 cold-output accumulation, replay lifetime handling, and regression tests are good improvements. Independent CPU arithmetic and CUDA kernel reviews found no new actionable numerical or kernel-memory issues.

Local validation passed: native/plugin host compilation, CUDA compilation, 53 policy/admission/lifetime tests, 11 INT2 tests, 143,262 sanitized CPU comparisons, and 155 A100 reference cases. CUDA memcheck, racecheck, synccheck, and initcheck were clean. These were targeted fresh-source checks, not a complete CUDA wheel or end-to-end operator-suite rebuild; Windows/WDDM and AVX512 execution were not validated.

The four inline comments cover remaining integration issues: pinned-memory accounting, automatic GPU workspace limits, CPU worker-count settings, and the decode token threshold. These prevent the shared session policy from reliably controlling resource use and dispatch. Since the PR is already merged, please address them with focused regression tests in a follow-up.

[Baiju Meswani (baijumeswani)]

Please count this allocation against the session's pinned-memory limit. Admission accepts session.moe.offload.pinned_bytes=0, but this block still allocates 396,352 pinned bytes for every admitted node, including fully GPU-resident and CPU-disabled plans. The prefill routing buffers and shared staging banks also allocate outside the pinned-memory ledger, and larger banks remain allocated after a smaller request.

This means the configured limit does not bound actual pinned memory, and other allocations using the ledger cannot see this usage. Please account for all live decode and prefill pinned allocations through the controller's ledger, including retained shared capacity, and shrink, fall back, or return an error when the complete allocation cannot fit. Add tests for zero/small limits and for a smaller session following a larger one.

[Baiju Meswani (baijumeswani)]

Please enforce the resolved device working budget even when the caller did not set it explicitly. This condition skips the size check for automatic budgets, although admission still reserves a working budget when deciding how many experts fit on the GPU.

Calling the current workspace-size function with 8,192 rows, hidden size 4,096, intermediate size 1,536, k=8, 128 experts, and two-byte activations requires 704.9 MiB of scratch before either staging bank. That exceeds both the usual 256 MiB budget and the 560 MiB large-model allowance, but this branch accepts it. A model can therefore initialize successfully and run out of GPU memory during prefill.

Please fit the scratch and live staging capacity within every resolved budget, reducing the chunk/wave size or adjusting admission as needed. Add a regression using an automatic budget, not only an explicit one.

[Baiju Meswani (baijumeswani)]

Please use the admitted policy.cpu_threads value here when it is positive. ResolveThreads() only reads the hardware thread count and ORT_QMOE_SHARED_CPU_THREADS, so session.moe.offload.cpu_threads=1 is ignored. The resolver selects 24 workers on the test host, or 7 when the environment variable is 7, despite the session setting of 1.

This prevents callers from limiting CPU worker usage and gives the environment variable precedence over the explicit session setting. Use the positive policy value first, and only use the automatic/environment path when it is zero, as the QMoE backend does. Add a test with an explicit session value and a different environment value.

[Baiju Meswani (baijumeswani)]

Please include offload_policy_.async_max_tokens in this decode decision. The condition only checks the fixed 16-row limit, so session.moe.offload.async_max_tokens=0 still selects CPU doorbell decode for one row, and the default value of 8 still selects it for 9–16 rows.

Callers therefore cannot use the shared session setting to select GPU streaming for those batches, unlike the QMoE path. Require the row count to fit the configured threshold as well as the existing backend shape limits. Add tests for zero and for row counts just below, at, and above a configured threshold.