[Tianlei Wu (tianleiwu)]

Summary

This PR fixes Windows WDDM and bounded-host-cache issues in QMoE CPU/GPU offload, adds tail-expert probability threshold pruning (session.moe.offload.min_expert_weight), and optimizes prefill and decode throughput for mixed-precision (INT4 / INT3 / INT2) QMoE models on consumer GPUs (e.g., RTX 4090 24 GB VRAM + 30 GiB RAM):

1. Fix Windows FileExpertStore::InitializeMapped & QMoEOffload::InitializeBoundedHostCache (qmoe_offload.h, qmoe_offload.cc, qmoe_shared_layout.cc):
   • Propagate external_file_path and external_file_offset from MoeExpertSourceDescriptor into QMoEOffload::InitializeBoundedHostCache().
   • On Windows (#elif defined(_WIN32) in FileExpertStore::InitializeMapped), accept committed MEM_MAPPED (or read_opt_enabled) mappings via VirtualQuery when source.file_path == nullptr.
2. Fix pinned memory budget accounting & persistent decode pinned starvation (qmoe_offload.cc):
   • Add SessionPromotionBudget::HeldBytes() and compute reserved = bounce_held + std::max(decode_held, reserve_target) in QMoEOffload::ComputeShared so promotion bounce allocations (16 MiB) are not conflated with persistent decode headroom (24 MiB).
   • In QMoEOffload::EnsurePersistentDecodeState, release idle promotion bounce and landing-ring pinned staging when try_reserve_decode_pinned() needs headroom.
3. Automatic machine-aware session planning (qmoe_session_controller.h):
   • On Windows hosts where ullTotalPhys < host_expert_bytes + 24 GiB, automatically prefer direct OS file-backed mapping (use_bounded_file_cache = false unless ORT_QMOE_FORCE_BOUNDED_HOST_CACHE=1), avoiding multi-GiB private heap commits and enabling 2-bank parallel pinned wave staging.
   • For full-model sessions (num_layers >= 24, total_vram >= 16 GiB, total_full_gpu_bytes >= 8 GiB), automatically raise implicit device_working_bytes to 560 MiB and implicit pinned_bytes to 160 MiB, and adapt admitted_gpu_expert_budget against platform_free_vram (DXGI local budget) and SessionBudgetAccount without spilling into WDDM non-local shared memory.
   • Enable EvictUploadedHostPages (VirtualUnlock) on Windows MEM_MAPPED views.
4. 8K single-tile streaming prefill, 2-bank double-buffered wave staging, and default Windows PrefetchVirtualMemory (qmoe_offload.cc):
   • Default stream_rows to 8192 across INT2, INT3, and full-budget INT4 layers.
   • Automatically halve wave_experts (32 -> 16) in ComputeShared when needed so INT3 layers fit stage_weight_banks = 2 (full double-buffered CPU staging + PCIe H2D DMA + GPU Tensor-Core GEMM overlap).
   • Extend ORT_QMOE_WIN_PREFETCH_MMAP (default 1 on Windows) to all file-mapped integer-bank layers.
5. Tail-Expert Probability Threshold Pruning (session.moe.offload.min_expert_weight, qmoe_kernels.cu, moe_kernels.cu, qmoe_offload.cc):
   • Add kOrtSessionOptionsMoeOffloadMinExpertWeight (session.moe.offload.min_expert_weight, env fallback ORT_QMOE_MIN_EXPERT_WEIGHT) to skip non-top-1 routed experts (rank >= 1) whose probability after SoftmaxTopK is strictly below min_expert_weight (topk_scales[route] = 0.0f, topk_indices[route] = -2), always keeping at least the top-1 expert (rank == 0).
   • Support ORT_QMOE_SKIP_TAIL_COLD_ONLY=1 to retain GPU-resident (slot >= 0) tail experts during decode at zero wall-time cost while pruning only CPU-cold (slot < 0) tail experts inside QMoEOffloadClassifyKernel.
   • Skip zero-weight routes in QMoEOffloadMergeFinalizeCombineKernel and finalizeMoeRoutingKernel without reading uninitialized contribution rows, and skip input_dev -> host_input_out PCIe copies in QMoEOffloadClassifyKernel when has_doorbell && cold_count == 0.
6. CPU AVX2 INT2 / INT3 cold-expert & SM89 dense INT8 Tensor-Core GEMV optimizations (qmoe_shared_layout_int2_avx2.cc, qmoe_shared_layout_int3_avx2.cc, dispatcher.h, paged_attention_impl.cu):
   • Add 8-column Int2Group8PrepackedAvx2 so all 8 columns of each 128-byte SM80 interleaved tile share one activation load while scaling weights before products.
   • Hoist activation vector loads outside step_col in Int3Group4Avx2 and implement a true 8-column Int3Group8Avx2.
   • Enable CanUseInt8TensorCore on params.device_sm >= 89 (SM89) and fix MSVC /WX C4459 variable shadowing in paged_attention_impl.cu.

Testing

• test_qmoe_offload_v4.py targeted unit tests: PASS.
• End-to-end Qwen3.8-Flash-Next RC5 (qwen_flash_next_int4_int8_rc5_contrib_dflash2_qmoe_2b375_x02_y14_z32_dp23_b64) benchmarks on RTX 4090 (24 GB VRAM) + Intel i9-13900 (30 GiB RAM):
  • 2K (2,048 tokens): TTFT 8.92 s (7.55 s at min_expert_weight=0.08), decode 12.15–23.43 tok/s (24.63–29.63 tok/s with min_expert_weight=0.065–0.08).
  • 8K (8,192 tokens): TTFT 14.42–15.32 s (10.84–12.30 s with min_expert_weight=0.065–0.08), decode 23.04 tok/s median / 29.49 tok/s peak.
  • 32K (32,768 tokens): TTFT 71.71–73.16 s, decode 8.22–9.48 tok/s.
  • 128K (131,072 tokens): TTFT 536.21–579.64 s, decode 7.42–15.27 tok/s, peak VRAM 21,345 MiB.
  • 100-Sample MMLU-Pro (shard 0 of 8, max-new 8192):
    • Default (min_expert_weight=0): 86/100 (86.0%), 0 truncated, 27.39 tok/s aggregate decode (3,196 s wall time).
    • Tail-skip (min_expert_weight=0.065, skip_cold_only=1): 84/100 (84.0%), 0 truncated, 34.70 tok/s aggregate decode (43.21 tok/s max, 2,099 s wall time; 43.10 tok/s avg / 53.42 tok/s max at 0.080 on 20-item subset).

[Copilot]

🟡 Changes recommended

Unresolved issues can falsely exhaust pinned memory, exceed explicit device-memory budgets, and override per-run warmup controls or terminate the process on malformed input, while the newly enabled SM89 path lacks parity and dispatch test coverage.

4 open issues

• Exercise the SM89 INT8 tensor-core path in parity tests · New
• Respect per-run warmup controls and handle malformed values safely · New
• Avoid retrying a decode that already succeeded · New
• Keep explicit expert-memory budgets as upper bounds · New

What changed in this PR

This PR improves ONNX Runtime’s mixed-precision QMoE CPU/GPU offload, targeting Windows WDDM memory pressure and prefill/decode performance on consumer GPUs.

Changes:

• Updates file-backed expert loading, session memory planning, and pinned-memory reclamation.
• Expands streaming tiles and double-buffered staging, with Windows prefetch enabled by default.
• Optimizes AVX2 INT2/INT3 expert computation, enables SM89 INT8 Tensor-Core dispatch, and resolves MSVC variable-shadowing warnings.

File	Description
onnxruntime/​core/​providers/​cuda/​qmoe_session_controller.h	Adds machine-aware budgets and Windows mapping policies.
onnxruntime/​contrib_ops/​cuda/​moe/​qmoe_offload.h	Stores expert-bank file paths and offsets.
onnxruntime/​contrib_ops/​cuda/​moe/​qmoe_offload.cc	Updates cache initialization, pinned budgets, streaming, prefetch, and warmup handling.
onnxruntime/​contrib_ops/​cuda/​llm/​fpA_intB_gemv/​dispatcher.h	Enables INT8 Tensor-Core dispatch on SM89 and newer devices.
onnxruntime/​contrib_ops/​cuda/​bert/​paged_attention_impl.cu	Renames local constants to avoid MSVC shadowing warnings.
onnxruntime/​contrib_ops/​cpu/​moe/​qmoe_shared_layout.cc	Accepts additional Windows mapped expert-bank sources.
onnxruntime/​contrib_ops/​cpu/​moe/​qmoe_shared_layout_int3_avx2.cc	Adds eight-column computation and reuses activation loads.
onnxruntime/​contrib_ops/​cpu/​moe/​qmoe_shared_layout_int2_avx2.cc	Adds an eight-column prepacked AVX2 path with weight scaling before products.

_{Review effort: Balanced}

[github-actions]

You can commit the suggested changes from lintrunner.

[github-actions]

You can commit the suggested changes from lintrunner.

[Baiju Meswani (baijumeswani)]

This PR improves QMoE CPU/GPU offload for Windows, updates memory planning and staging, widens the AVX2 INT2/INT3 CPU paths, enables the INT8 tensor-core path on SM89, and adds optional pruning of low-weight expert routes. These are useful directions for running larger models on limited GPU memory.

What looks good: the packed layouts and numerical behavior held up in extensive fresh CPU tests. The earlier warmup-control, duplicate pinned-reservation and SM89-test concerns are addressed. The ordinary explicit GPU-budget case is also fixed, although one boundary case remains below.

I do not think this is ready for approval yet. The six inline comments cover two memory-limit violations, two pruning configuration/path inconsistencies, and two repeatable CPU projection slowdowns. Please preserve explicit memory limits, honor an explicit zero pruning threshold, make pruning consistent across decode paths, and retain or improve the CPU paths that are faster in the measured configuration. Add regression tests for the limit and pruning cases, and repeat the CPU comparisons with production build settings.

Validation: two independent reviewers covered CPU math/performance and memory/planning, alongside CUDA/integration review. Fresh checks passed 16 CPU tests (one additional AVX512-only test skipped), 39 planning/store/retry tests, 56 production CUDA scenarios, and eight adapted INT8 tests. Targeted CUDA memory checks were clean, and native/plugin offload frontends compiled. The two Windows limit tests fail at this head and pass at the base and with small corrective changes in the test artifacts.

Limits: Windows planning was tested through exact-source simulations, not native WDDM execution. INT8 tests exercised pre-SM90 code on A100 with a review-only SM89 reporting adaptation, not actual SM89 hardware. CPU performance measurements are Linux/EPYC/GCC13 projection-helper timings, not full-model or Windows throughput. No full-model speed or quality result is claimed.

[Baiju Meswani (baijumeswani)]

Thanks for the updates. This PR improves QMoE CPU/GPU offload and memory planning, and adds optional pruning of low-weight expert routes. Five of the six earlier findings are now resolved: the explicit GPU limit, explicit zero pruning setting, callback/all-cold pruning behavior, and both CPU performance regressions.

The fixes look good in the focused checks. INT2 performance is back within measurement noise, and the tested INT3 projection is about 15% faster than the original base. Fresh validation passed 20 CPU tests (two additional AVX512-only tests skipped), 34 native policy/admission tests, 736 INT4 parity cases, 120 pruning-pipeline cases and 56 routing/finalization cases. Targeted CUDA memory checks were clean, and native/plugin offload frontends compiled.

I would still hold approval for the remaining host-memory issue in the inline comment. The planner now preserves the explicit host limit, but the backend constructor can disable the bounded cache afterward. Please preserve that admitted limit and add a regression covering the constructor and cache initialization, not only the planner.

The source-level Windows regression reproduces this issue and passes when the constructor override is removed. Windows behavior was simulated; native Windows, SM89 and AVX512 execution were not locally validated. No full-model speed or quality result is claimed.