[BBaditi]

Note: This PR is stacked on #33975 (qmoe-packed-gguf-cpu-ep) and #33976 (qmoe-packed-gguf-cpu-offload), and pairs with #34066 (qmoe-packed-gguf-cuda-matmul-m8-opt on main).

Summary

This PR optimizes com.microsoft::MoEPacked multi-token (M = 4, 8) speculative decoding verification across both the CPU hybrid cold-route path (gguf_vec_dot_cpu.h) and the GPU Tensor Core FC2 decode path (moe_packed.cu):

1. CPU Multi-Token Cold-Route Execution (ExecuteMoEPackedColdRoutesCpuWithScratch in gguf_vec_dot_cpu.h):

   • Zero-Heap Stack Buffers: Replaces per-layer std::vector<int> sorted_c and std::vector<ColdExpertGroup> groups allocations with stack buffers (sorted_c_stack[128], groups_stack[128]) for num_cold_routes <= 128, eliminating 96 heap allocations/frees per 48-layer verify step.
   • Parallel Multi-Token Input Q8 Quantization: Parallelizes BF16 -> FP32 + QuantizeRowQ8Cpu across num_unique_toks when num_unique_toks >= 2.
   • 2-Route Fused AVX2 Grouped Dot Products (VecDotGgufRow4xQ8Batch2Avx2 & VecDotGgufGateUpRow2xQ8Batch2Cpu): When $\ge 2$ tokens in an M = 4 or M = 8 verification batch route to the same cold expert (ColdExpertGroup with grp.count >= 2, where measured cold-route reuse $C/U$ is 1.39x-1.48x at M=4 and 1.86x-2.07x at M=8), each GGUF weight sub-block (IQ2_XXS, IQ2_XS, IQ2_S, IQ3_XXS, IQ4_NL, IQ4_XS, Q2_0) is unpacked once in AVX2 registers and multiplied against two routes' Q8 activations simultaneously with 0 ULP difference.

2. GPU Tensor Core FC2 Decode/Verify (MoEPackedFc2Bf16TcKernel<1> in moe_packed.cu):

   • Eliminates the 10.75 KB s_act_u4 shared-memory activation staging buffer, integer-division staging loop, and K-loop __syncthreads() barrier when kMaxSubTiles == 1 (total_routes <= 128), loading fc1_out directly via 128-bit __ldg.
   • Stages GGUF lookup tables (iq2xxs_grid, iq2xs_grid, iq2s_grid, iq3xxs_grid) in shared memory (s_lut, DecodeGgufSubBlock8SmemLutTpl<kDownType>) and switches decode CTA geometry from 8 warps (128 output channels/CTA) to 4 warps (64 output channels/CTA) in both DispatchMoEPackedHotExpertsCuda and LaunchMoEPackedFc2Bf16TcDispatch<1>, doubling CTA grid parallelism across SMs.

---

Measured Performance Impact (48-Layer Qwen3.8-Flash-Next + dflash2.onnx on NVIDIA A100-SXM4-80GB)

Table A: 48-Layer Verify Step Latency Progression (UD-IQ2_XS)

Metric (UD-IQ2_XS)	Stage 0 (Baseline)	Stage 1 (This PR: MoEPacked FC2 + CPU Batch-2)	Stage 2 (+ PR #34066 MatMulPacked & ReadU32Dev)	Total Gain
100% GPU MoEPacked (48 nodes, M=1 / M=4 / M=8)	10.74 / 15.82 / 20.29 ms	10.23 / 15.05 / 18.81 ms	10.28 / 15.32 / 19.24 ms	-0.5 to -1.5 ms
50/50 Hybrid MoEPacked (48 nodes, M=1 / M=4 / M=8)	17.40 / 27.02 / 37.81 ms	18.31 / 24.69 / 33.91 ms	16.98 / 24.83 / 34.33 ms	-2.3 to -3.9 ms
100% GPU T_verify(M=4)	34.48-35.10 ms	29.02-29.73 ms	26.46-27.32 ms	-7.9 ms/step (-22.8%)
100% GPU T_verify(M=8)	50.32-51.68 ms	35.76-38.05 ms	32.90-35.27 ms	-16.9 ms/step (-33.2%)
50/50 Hybrid T_verify(M=4)	38.68-52.97 ms	31.73-41.84 ms	29.64-39.80 ms	-11.1 ms/step (-24.2%)
50/50 Hybrid T_verify(M=8)	56.08-80.69 ms	41.46-56.03 ms	38.69-53.59 ms	-22.2 ms/step (-32.5%)

Table B: End-to-End dflash2.onnx Speculative Decoding (64 Generated Tokens / Prompt, Combined Stage 2)

Model	Placement	M=1 Decode (ms/tok / tok/s)	Fixed M=8 Peak (P1_structured)	Adaptive M={4,8} 3-Prompt Mean (tok/s / Speedup)	Greedy Token Parity
UD-IQ1_S	100% GPU	21.16 ms (47.3 tok/s)	185.3 tok/s (3.93x)	106.0 tok/s (86.2 harm., 2.24x)	64/64 (100.0%)
UD-IQ1_S	50/50 Hybrid (-5.21 GB)	24.27 ms (41.2 tok/s)	129.6 tok/s (3.43x)	90.2 tok/s (84.9 harm., 2.23x)	64/64 (100.0%)
UD-IQ2_XS	100% GPU	19.91 ms (50.2 tok/s)	197.4 tok/s (3.93x)	123.7 tok/s (111.3 harm., 2.46x)	64/64 (100.0%)
UD-IQ2_XS	50/50 Hybrid (-5.78 GB)	23.77 ms (42.1 tok/s)	126.8 tok/s (3.32x)	89.8 tok/s (82.7 harm., 2.17x)	64/64 (100.0%)
UD-IQ3_XXS	100% GPU	22.27 ms (44.9 tok/s)	177.5 tok/s (3.96x)	109.3 tok/s (98.1 harm., 2.44x)	64/64 (100.0%)
UD-IQ3_XXS	50/50 Hybrid (-7.04 GB)	27.77 ms (36.0 tok/s)	115.7 tok/s (3.30x)	80.1 tok/s (74.1 harm., 2.22x)	64/64 (100.0%)

---

Validation

• /usr/bin/clang-format --dry-run --Werror and git diff --check: 0 issues
• ctest --output-on-failure -V -R "GgufQuantPackedTest": 2/2 passed (100%)
• pytest onnxruntime/test/python/transformers/test_gguf_packed_ops.py -v: 33 passed, 811 subtests passed (100%)

[Copilot]

🟡 Changes recommended

Unresolved critical defects can cause out-of-bounds bias reads, incorrect CUDA results, execution hangs, and violations of configured offload memory limits.

20 open issues

• The optional biases are passed to a helper that indexes them by expert and channel without… · New
• The task count is packed into 20 bits without checking that it fits. A supported batch covering 512… · New
• This function-local static joins its worker threads from its destructor, so unloading the CUDA… · New
• The pinned slabs and retained scratch bypass pinned_bytes and host_working_bytes, so admission… · New
• The hybrid path caches these bias pointers without validating their dimensions; the offload… · New
• This flag lets admission accept a host-expert budget smaller than the cold-expert complement, but… · New
• Partial-residency plans execute every cold expert on CPU even when… · New
• With enable_cuda_graph=1, partially resident execution synchronizes an event recorded on the… · New
• The full-size cold-output buffer is allocated even for ordinary GPU-only execution, where it is… · New
• This changes an explicitly configured session.moe.offload.mode=disabled back to auto whenever… · New
• The MoEPacked-specific expert-count setting becomes an unrestricted wildcard, so it also changes… · New
• Malformed values such as abc or -1 are silently ignored for this expert-count option, unlike… · New
• The new biases are silently ignored by GPU-resident experts, so a biased model produces different… · New
• USE_CUDA only establishes that CUDA was built, not that the current GPU supports this test. All… · New
• The pool ignores session.moe.offload.cpu_threads: even an explicit value of 1 starts at least 15… · New
• Adaptive admission reserves spare slots and requests background promotion, but this backend… · New
• The new shared-LUT FC2 decode branches for IQ2_XS, IQ2_S and IQ3_XXS are not exercised by the C++… · New
• The new alias and the canonical device-expert-budget key can both be supplied with different… · New
• An explicit session.moe.offload.reserve_device_bytes=0 is treated as an automatic reserve, so… · New
• The per-format test checks only single-route helpers; batch-2 is covered indirectly at K=256/64… · New

What changed in this PR

This PR extends com.microsoft::MoEPacked CPU and hybrid GPU/CPU execution while optimizing CUDA multi-token decode. It builds on the stacked CPU-kernel and expert-offload work.

Changes:

• Adds CPU kernels and GGUF numerical/parity tests.
• Connects packed experts to offload admission, staging, and hot/cold routing.
• Reduces BF16 FC2 decode overhead and updates the schema and operator documentation.

File	Description
onnxruntime/​test/​contrib_ops/​gguf_quant_packed_test.cc	Adds GGUF dot-product and CPU/CUDA/hybrid parity tests.
onnxruntime/​core/​providers/​cuda/​qmoe_session_controller.h	Updates override matching and device-budget reserves.
onnxruntime/​core/​graph/​contrib_ops/​contrib_defs.cc	Adds optional biases and float activation support.
onnxruntime/​core/​framework/​session_state.cc	Adds MoEPacked eligibility and initializer placement for offload.
onnxruntime/​core/​framework/​moe_offload_plan.h	Adds budget alias and GPU-expert-count configuration handling.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_packed.h	Declares hybrid dispatch interfaces and persistent offload state.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_packed.cu	Optimizes FC2 decode and separates routing, hot execution, and finalization.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_packed.cc	Implements resident expert uploads, staging, and CPU-cold execution.
onnxruntime/​contrib_ops/​cpu/​moe/​moe_packed_cpu.h	Declares the CPU kernel.
onnxruntime/​contrib_ops/​cpu/​moe/​moe_packed_cpu.cc	Implements CPU routing, expert projections, and output accumulation.
onnxruntime/​contrib_ops/​cpu/​cpu_contrib_kernels.cc	Registers float, FP16, and BF16 CPU kernels.
docs/​OperatorKernels.md	Updates CPU availability and bias-input listings.
docs/​ContribOperators.md	Documents optional biases and expanded activation types.

_{Review effort: Max}

[Copilot]

🟡 Changes recommended

AVX2-enabled test builds fail because a void-returning helper is used as a condition and another helper is called despite being unavailable on native MSVC.

4 open issues

• AVX2 helper returns void, breaking the test build · New
• MSVC AVX2 build calls an undeclared quantization helper · New
• Optimized Q8 batching is not used by hybrid decode · New
• Decode tests miss trailing inactive FC2 warps · New

20 resolved since last review

• USE_CUDA only establishes that CUDA was built, not that the current GPU supports this test. All…
• The new biases are silently ignored by GPU-resident experts, so a biased model produces different…
• Malformed values such as abc or -1 are silently ignored for this expert-count option, unlike…
• The MoEPacked-specific expert-count setting becomes an unrestricted wildcard, so it also changes…
• This changes an explicitly configured session.moe.offload.mode=disabled back to auto whenever…
• The full-size cold-output buffer is allocated even for ordinary GPU-only execution, where it is…
• With enable_cuda_graph=1, partially resident execution synchronizes an event recorded on the…
• Partial-residency plans execute every cold expert on CPU even when…
• This flag lets admission accept a host-expert budget smaller than the cold-expert complement, but…
• The hybrid path caches these bias pointers without validating their dimensions; the offload…
• The pinned slabs and retained scratch bypass pinned_bytes and host_working_bytes, so admission…
• This function-local static joins its worker threads from its destructor, so unloading the CUDA…
• The task count is packed into 20 bits without checking that it fits. A supported batch covering 512…
• The optional biases are passed to a helper that indexes them by expert and channel without…
• The per-format test checks only single-route helpers; batch-2 is covered indirectly at K=256/64…
• An explicit session.moe.offload.reserve_device_bytes=0 is treated as an automatic reserve, so…
• The new alias and the canonical device-expert-budget key can both be supplied with different…
• The new shared-LUT FC2 decode branches for IQ2_XS, IQ2_S and IQ3_XXS are not exercised by the C++…
• Adaptive admission reserves spare slots and requests background promotion, but this backend…
• The pool ignores session.moe.offload.cpu_threads: even an explicit value of 1 starts at least 15…

_{Review effort: Max}