onnxruntime
a2c7c3bc - Fix QMoE CPU livelock by eliminating nested intra-op parallelism (#29081)

Commit
93 days ago
Fix QMoE CPU livelock by eliminating nested intra-op parallelism (#29081) ### Description `QMoECPU::ComputeCommon` runs an expert loop over the routed experts. Two issues hurt multi-threaded execution: 1. **Nested intra-op parallelism livelocked the pool.** When the expert loop fanned out over the shared session pool (`tp`), the per-expert body issued *further* work on the same pool (token copy, dequant blocks, activation, `MlasGemm` / `DirectQ4Gemm` / `TryRunLutGemm` / `DequantizeBlock`). This nested parallelism on ORT's Eigen pool intermittently livelocked (workers spinning at 100% CPU, never completing). 2. **Decode (batch=seq=1) was poorly parallelized.** The expert thread count was capped by `num_experts` (not *active* experts), and an unconditional work tier cap serialized the active experts. For decode each per-expert GEMM is effectively a GEMV (`M == 1`) that MLAS does not thread internally, so the cores were left idle. This change makes the expert loop parallelize across the **active** experts and keeps parallelism single-level so the livelock fix is preserved: ```cpp // Only experts that actually received tokens do work. int num_active_experts = 0; for (const auto& tokens : expert_token_map) { if (!tokens.empty()) ++num_active_experts; } // One expert (batch) per thread; cap by active experts so a single busy expert // does not look "parallel". int num_expert_threads = std::max(1, std::min(num_active_experts, max_expert_threads)); // Single-level parallelism preserves the no-livelock guarantee: // - expert loop multi-threaded -> inner ops run serially (inner_tp == nullptr) // - single active expert -> loop is serial, inner GEMM gets the full pool (inner_tp == tp) concurrency::ThreadPool* inner_tp = (num_expert_threads > 1) ? nullptr : tp; ``` Changes: - Count active experts and cap `num_expert_threads` by `num_active_experts`. - Derive `inner_tp` from the *active* fan-out, so a single-active-expert batch keeps full inner MLAS/GEMM parallelism (addresses review feedback). - Remove the work tier cap that serialized active experts. This keeps parallelism single-level (no nested pool use), so the original livelock fix is intact while decode and imbalanced routing run on all available cores. ### Motivation and Context A random `Run()` on a quantized MoE model (`com.microsoft.QMoE`) intermittently pinned one core at 100% inside `QMoECPU::Compute` and never returned under multi-threaded intra-op execution; `intra_op_num_threads = 1` avoided it at ~3× latency. The signature — non-deterministic, multi-threaded-only, identical stacks (main thread in the QMoE expert lambda, pool threads in `WorkerLoop`) — points to nested parallelism on the shared thread pool. Restricting QMoE to single-level parallelism removes the hang. Capping by *active* experts (rather than `num_experts`) and dropping the tier cap then restores — and substantially improves — throughput for decode and imbalanced routing. ### Experiment Results Latency in ms/inference (fp32, 4-bit, block size 32, default intra-op threads, AMD EPYC 7763 32 logical / 16 physical cores). **Decode (batch = seq = 1)** | Model (experts / top-k) | Before | After | Speedup | | --- | ---: | ---: | ---: | | gpt_oss_20b (e32 / top4) | 359 | 88 | ~4.1× | | qwen3 (e256 / top8) | 91 | 16 | ~5.6× | | gemma (e128 / top8) | 173 | 27 | ~6.4× | **Mid-range sequence lengths (gpt_oss)** | seq | Before | After | Speedup | | ---: | ---: | ---: | ---: | | 8 | 1776 | 146 | ~12× | | 16 | 1096 | 176 | ~6.2× | | 32 | 1165 | 188 | ~6.2× | | 64 | 705 | 190 | ~3.7× | **Prefill (neutral)** | seq | Before | After | | ---: | ---: | ---: | | 128 | 164 | 166 | | 512 | 192 | 194 | ### Validation - Parity tests: `test_qmoe_cpu.py` — 80 passed, 1 skipped. - Lint: clean (`lintrunner`). - Livelock stress: 8 threads × 300 iterations × 5 shapes = 2400 concurrent runs, no hang, no errors. --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: tlwu <tlwu@example.com>
Author
Parents
Loading