Repository navigation
[CUDA] Improve QMoE offload decode throughput via AVX2 spill elimination and working-set retention - #33935
Conversation
… working-set retention - Eliminate AVX2 YMM register stack spills in Int2Group4Avx2, Int3Group8Avx2, Int4Group8Avx2, and Int4Group4Avx2 by using short-lived activation loads and sequential low/high half accumulation within 16 YMM registers (RT <= 2). - Add fast finite BF16 scale rounding paths to Int3LookupAvx2 and RoundInt4WeightsAvx2. - Lazily allocate CpuWork scratch buffers on demand in RunColdExperts, use direct 2-pass SpinPool execution when FileExpertStore demand loading is inactive, and keep both SessionPromotionBudget bounce epochs within remaining session pinned budget. - Trim dummy-token warmup pages from the Windows working set before the first live turn and gate PrefetchVirtualMemory on available physical RAM so subsequent prefills do not evict hot decode cold experts on <= 32 GB hosts. - Reduce 2K context tier WDDM extra reserve on < 28 GiB GPUs (5,742 -> 6,573 resident GPU slots, 58.71% decode hit rate). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Remove the Windows low-RAM override that zeroed host_cache_bytes_ after admission planned a finite bounded cache. This skipped InitializeBoundedHostCache and left cold experts on direct-source access with no working-set bound, while the saved plan still reported a bounded cache. Add a regression with 64 experts and a cache larger than 64 MiB that checks the bounded host cache serves cold experts. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
|
Addresses the unresolved review feedback from #33916 (#33916 (comment)), which merged before it could be fixed.
Validation: the new test and the neighbouring INT4 bounded-cache warmup test pass on an existing RTX 4090 Windows build. That build predates this change, and the override only fires on hosts with less than about 24 GiB of RAM, so the test did not exercise the fix itself on this 30 GiB machine. Native low-RAM Windows execution was not validated, and the benchmark numbers above were measured with the previous override active. |
There was a problem hiding this comment.
🟡 Changes recommended
The process-wide pending-trim flag can cause premature cross-session cache eviction, while asynchronous decode-sized warmups bypass flagging and miss the intended trim before live execution.
2 open issues
What changed in this PR
This follow-up to #33916 improves QMoE CPU/GPU offload decode throughput, targeting AVX2 register pressure and memory retention on resource-constrained Windows hosts.
Changes:
- Refactors INT2/INT3/INT4 AVX2 kernels to reduce register spills and accelerate BF16 rounding.
- Adjusts warmup trimming, prefetching, pinned-memory sizing, and cold-expert scratch allocation.
- Reduces the small-context WDDM reserve and adds bounded-host-cache regression coverage.
| File | Description |
|---|---|
onnxruntime/test/python/test_qmoe_offload_v4.py |
Adds regression coverage for retaining an admitted bounded host cache. |
onnxruntime/core/providers/cuda/qmoe_session_controller.h |
Reduces extra WDDM reserve for eligible smaller-GPU configurations. |
onnxruntime/contrib_ops/cuda/moe/qmoe_offload.cc |
Updates memory retention, staging budgets, and CPU cold-expert execution. |
onnxruntime/contrib_ops/cpu/moe/qmoe_shared_layout_int4_avx2.cc |
Reduces register pressure and adds a fast BF16 rounding path. |
onnxruntime/contrib_ops/cpu/moe/qmoe_shared_layout_int3_avx2.cc |
Shortens activation lifetimes and accelerates BF16 lookup construction. |
onnxruntime/contrib_ops/cpu/moe/qmoe_shared_layout_int2_avx2.cc |
Uses smaller row tiles and shorter-lived vectors to reduce spills. |
Review effort: Balanced
…async warmups The pending-trim flag was process-wide, so another session's live run could consume it and trim while a session was still warming up. Keep it on the session controller. Decode-sized warmups return through ComputeSharedAsync before the only place that set the flag, so mark them there as well. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Baiju Meswani (baijumeswani)
left a comment
There was a problem hiding this comment.
This PR speeds up the one- and two-row CPU expert paths, fixes the bounded host-cache initialization issue, and avoids allocating large temporary CPU buffers before they are needed. The memory-budget and warmup changes also look good in the focused checks.
I would hold approval until the two performance issues below are addressed. Removing the three- and four-row INT2/INT4 paths makes supported larger-row calls slower because they repeat work on the same weights. The inline comments include measurements at the default offload task width and a suggested fix. Please keep the faster one- and two-row paths while addressing these larger-row cases.
The measurements use fresh builds of this head and its base on Linux/GCC/EPYC. Three repeated, CPU-pinned comparisons gave matching outputs. These are CPU helper measurements, not Windows or full-model speed measurements.
Validation: 20 official CPU tests passed, with two AVX512-only skips; exhaustive rounding/range and sanitizer checks passed; 95 memory/policy tests passed, including eight real CUDA pinned-memory tests; 288 instrumented scheduling cases passed; native and plugin source compilation checks passed. I did not run a fully linked current CUDA provider, end-to-end model tests, or native Windows tests. The current CI snapshot has 105 successful checks and one neutral check.
…rnels Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Auto-merge was disabled for another pull request in the stack
…33935) (#33944) ### Co-authored-by: @BBaditi (original author of the tail substitution work in #33625; his three commits are kept as authored, and the follow-up commit carries a Co-authored-by trailer). Stacked on #33935 (base branch `tlwu/20261106/qmoe-offload-decode-avx2-opt`; retarget to `main` once #33935 merges). Cherry-picks the three commits of #33625 (opt-in tail expert substitution for hybrid QMoE decode) onto #33935, resolves conflicts with the merged `min_expert_weight` / `skip_tail_cold_only` options, and extends substitution to speculative-verify tiles. The settings below come from a sweep of `min_expert_weight`, `tail_substitution_threshold` and `tail_substitution_min_rank` on the real model. Changes beyond #33625: - Tail substitution now handles up to 16 decode rows (`kQMoETailSubstitutionMaxRows`) instead of 4, so S=8 speculative verify steps (80 routes per layer) are covered. Per-row telemetry arrays are sized accordingly. - The substitution kernel ignores routes that `SoftmaxTopK` already pruned (index `< 0`), so it is safe together with `skip_tail_cold_only=0`. - `QMoEOffloadClassifyKernel` keeps the `min_expert_weight` cold-prune arguments and takes a separate telemetry-only scales pointer. - Added `test_v4_tail_substitution_covers_speculative_verify_tiles` (8 rows) and tuning guidance in `docs/moe-cpu-offload/qmoe_offload.md`. ### Sweep results (Qwen3.8-Flash-Next RC5 2B375, k=10, 256 experts, 46.7% GPU residency, S=8 speculative decode, RTX 4090 + i9-13900 AVX2, 2K context for the sweep, 8K for MMLU-Pro) Token fidelity over 24 MMLU-Pro prompts x 128 greedy tokens (`div/100tok` = first-divergence events per 100 tokens of agreeing prefix, lower is closer to exact; baseline is deterministic): | Config (`min_expert_weight` / substitution threshold / min_rank) | Decode tok/s | Cold experts per step | `div/100tok` | | :--- | :---: | :---: | :---: | | exact | 24.5 | 785 | 0.00 | | skip 0.065 | 31.2 | 548 | 3.39 | | skip 0.080 | 44.1 | 322 | 7.44 | | sub 0.005 / r3 | 31.0 | 498 | 1.47 | | sub 0.010 / r3 | 44.1 | 343 | 3.14 | | sub 0.010 / r2 | 47.7 | 331 | 2.14 | | sub 0.020 / r2 | 62.6 | 196 | 3.18 | | sub 0.040 / r2 | 77.3 | 126 | 3.30 | | sub 0.020 / r3 + skip 0.050 | 60.3 | 217 | 2.33 | MMLU-Pro, 200 items (shards 0 and 1, `max-new 8192`): | Config | Accuracy | Decode tok/s | Wall time | | :--- | :---: | :---: | :---: | | exact | 162 / 200 (81.0%) | 29.6 / 30.9 | 6,937 s | | sub 0.020 / r2 | 162 / 200 (81.0%) | 96.9 / 103.1 | 2,631 s | | sub 0.040 / r2 | 165 / 200 (82.5%) | 108.5 / 117.1 | 2,376 s | | sub 0.020 / r3 + skip 0.050 | 168 / 200 (84.0%) | 86.0 / 86.4 | 2,630 s | Findings: - At equal speed, substitution diverges from the exact token stream less than tail skipping. - A low `min_rank` (2-3) matters: the `k`-relative default of 6 is close to a no-op on RC5. - `tail_substitution_threshold=0.04`, `tail_substitution_min_rank=2` gives about 3.7x decode throughput and 2.9x MMLU-Pro wall-time reduction with accuracy within noise of exact routing. `threshold=0.01`, `min_rank=2` is the conservative option (2.5x). - Adding `min_expert_weight` on top of substitution did not add speed and also prunes prefill routes. The defaults are unchanged (substitution disabled). Accuracy differences are within noise (one standard error is about 2.7% on 200 items); a teacher-forced logit KL on the real model was not available through the engine API. ### Testing - `pytest onnxruntime/test/python/test_qmoe_offload_v4.py`: 56 passed, 3 skipped (rebuilt CUDA plugin EP and Python bindings, RTX 4090). - Not run: the C++ unit-test targets and other platforms. ### Note on the host-cache change from #33935 With the Windows low-RAM override removed in #33935, the `consumer-gpu-24g-host-64g` profile (16 GiB `host_expert_bytes`) on a 29.66 GiB host creates a 16 GiB bounded host cache. Pinned, it fails model load with `CUDA error: out of memory`. Unpinned it decodes at about 11-12 tok/s instead of 17-19 tok/s with direct mapping. The sweeps above used a package profile with `host_expert_bytes` large enough to avoid the bounded cache. --------- Co-authored-by: BBaditi <bbaditi@microsoft.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
… offload optimizations (#33973) ## 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 - #33935 (`[CUDA] Improve QMoE offload decode throughput via AVX2 spill elimination and working-set retention`) - #33944 (`[CUDA] Tail expert substitution for QMoE offload decode (stacked on #33935)`) - #33877 (`Support direct external initializer buffers in AddExternalInitializers`) - #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. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Description
Follow-up to #33916. Profiling QMoE CPU/GPU offload decode on an RTX 4090 (24 GB VRAM, ~30 GiB host RAM, Windows) identified several decode bottlenecks. This PR addresses them:
Int2Group4Avx2,Int3Group8Avx2,Int4Group8Avx2,Int4Group4Avx2: captured activation arrays were spilled to the stack, andRT = 3, 4instantiations needed more than 16 YMM registers. Activations are now loaded with short-lived_mm256_loadu_psfolded into FMA, low/high halves are accumulated sequentially, andRT <= 2. Also adds fast finite BF16 scale rounding paths toInt3LookupAvx2andRoundInt4WeightsAvx2.PrefetchVirtualMemoryon later turns evicted hot decode experts on low-RAM hosts. The warmup working set is now trimmed before the first live turn (Windows, < 48 GiB RAM), andPrefetchVirtualMemoryis gated on available physical RAM.SessionPromotionBudgetbounce slots are now sized to fit the remaining session pinned budget, so both bounce epochs stay allocated during decode.CpuWorkallocation / unneeded spin-waits: scratch buffers are allocated lazily inRunColdExperts, and direct 2-passSpinPoolexecution is used whenFileExpertStoredemand loading is inactive.Results (RTX 4090, Qwen3.8-Flash-Next RC5, 3 reps, 64 new tokens)
8K outputs remain bit-exact (identical hashes) versus the baseline.
Motivation and Context
Improves decode throughput of QMoE CPU/GPU hybrid offload on consumer GPUs with limited VRAM and host RAM.