Fix QMoE CPU livelock by eliminating nested intra-op parallelism

[Copilot]

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:

// 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.

[tianleiwu]

I found one focused performance edge case in the threading change. The overall direction makes sense and the null-thread-pool fallbacks look safe, but sparse or highly imbalanced routing can still make the outer expert loop look parallel even when only one expert has real work. In that case this patch disables the inner MLAS/GEMM parallelism for the only substantial expert, so capping by active experts should preserve the livelock fix without serializing that workload.

[Copilot]

Pull request overview

This PR updates the QMoE CPU quantized implementation to avoid nested intra-op parallelism (which can livelock ORT’s Eigen-based thread pool) while improving CPU utilization for decode-style workloads by parallelizing across active experts.

Changes:

• Replace the previous “total-work” based expert thread selection with an active-expert count and cap num_expert_threads by num_active_experts.
• Introduce inner_tp so only one level uses the session threadpool: inner ops (GEMM/dequant/copies/activation) run with nullptr when the outer expert loop is multi-threaded, and use tp only when the outer loop is effectively serial.
• Remove the work-tier cap that previously forced additional serialization of active experts.

[hariharans29]

Review of PR #29081 — Fix QMoE CPU livelock by eliminating nested intra-op parallelism

Verdict: approve, with two design observations and one nit. The diagnosis (nested intra-op parallelism on the shared Eigen pool → livelock) is correct and well-known; the fix is exactly the right shape (single-level parallelism by construction); and the decode-throughput regression that the first round of feedback caught was addressed in commit 20285f3 the way @tianleiwu asked. 86/86 checks green, parity passes, 2400-run stress passes. The patch is also a net -1 LOC, which is a nice side effect.

Why the fix is right

The mechanical claim is: only one tier may use tp at a time. The PR enforces this statically with one local:

int num_expert_threads = std::max(1, std::min(num_active_experts, max_expert_threads));
concurrency::ThreadPool* inner_tp = (num_expert_threads > 1) ? nullptr : tp;

…then mechanically swaps every inner-loop tp over to inner_tp (8 sites: token copy, FC1 LUT-GEMM, FC1 direct-Q4 GEMM, FC1 dequant, FC1 MLAS-GEMM, SwiGLU activation, FC2 LUT/direct/MLAS, FC2 dequant). This is the only correct way to do this — anything dynamic (e.g. "use inner threads only if the outer is idle") would re-introduce the race that caused the hang. Sites I cross-checked against the diff at moe_quantization_cpu.cc:1346-1797 all swap correctly and the gating conditions stay coherent (num_expert_tokens >= 8 && num_blocks > 1 && inner_tp != nullptr for the token copy, the same for the activation, etc.). I don't see a stale tp left behind inside the per-expert lambda.

The single-active-expert corner that @tianleiwu raised is now also clean:

• 1 active expert → num_expert_threads = 1 → outer loop is effectively serial → inner_tp == tp → inner GEMMs get the full pool. Exactly what's needed for the "8 experts routed, only 1 got tokens" case.
• ≥2 active experts → outer pool fans out → inner ops serial. No nested use, no livelock.

The removal of the old tier cap (< 48 → 1, < 192 → 2, < 512 → 4) is also defensible: for decode the per-expert GEMV is M=1, MLAS does not internally thread it, so the only available parallelism is the outer expert loop. The benchmark numbers (4–6× decode wins on gpt_oss/qwen3/gemma; ≤1% prefill change) confirm the cap was hurting decode without helping prefill.

Two observations worth tracking (not blockers)

1. Highly imbalanced multi-active routing under-uses cores

The fix is binary at the outer level: either the outer loop is multi-threaded (and all inner ops are serial), or it's serial (and inner gets the pool). That's the right invariant for correctness, but it leaves an under-utilization case:

8 active experts with token counts [1000, 1, 1, 1, 1, 1, 1, 1].

The descending-size bin-packing at moe_quantization_cpu.cc:1314 puts the 1000-token expert alone on thread 0. Threads 1–7 finish in microseconds, then sit idle on the pool. Thread 0 runs its M=1000 GEMM with inner_tp == nullptr even though 7 cores are now free. Compared to the pre-PR (livelock-risky) version this is a regression for that distribution.

I don't think the PR should try to fix this — any "use inner threads if the outer is winding down" heuristic re-opens the door to the original race. But it's worth a follow-up TODO: classify "substantial" experts (e.g. token count above some fraction of the mean) and either cap num_expert_threads by substantial active experts, or pack such that substantial experts share threads with empty ones (so the substantial one is "alone" on a serial outer slot → gets inner_tp == tp). Out of scope for this PR; the livelock is the bigger problem and this is a knob for a later perf pass.

2. Per-thread workspace amplification

num_expert_threads now grows up to min(num_active_experts, max_expert_threads) rather than the (much smaller) tier-capped value. Five buffers scale with it:

• thread_local_outputs_ptr (num_tokens * hidden_size floats per thread)
• workspace_ptr (A1+C1+A2+C2+B1_dequant+B2_dequant per thread — the big one; B1_dequant_size = fc1_out_features * hidden_size)
• bias_conversion_buffers_ptr
• lut_packed_buffers_ptr
• lut_scale_conversion_buffers_ptr

For gpt_oss_20b sizes (4096 × 11008 fp32) B1_dequant_size alone is ~180 MB per thread; with num_expert_threads = 4 (top_k=4 decode) that's ~720 MB just for B1. Before the PR's removal of the tier cap, that decode path would have allocated 1–2 such workspaces. Now it's 4 (top_k=4) or 8 (top_k=8). Within budget for the workloads the PR targets (large server CPUs running MoEs) but not free. Not asking for a code change — flagging so the memory line item is on the record.

Nit

The new comment block at moe_quantization_cpu.cc:1137-1141 ends with (see PR #29081). Referencing the same PR's number from inside the merged code is awkward — once this lands, the reader is at HEAD and that pointer adds nothing. The why (nested parallelism on the Eigen pool livelocks) is the load-bearing part of the comment and that's already there; the PR-link tail can be dropped. The git blame will surface this PR anyway.

Things that look right

• All 8 tp → inner_tp sites are swapped consistently. I cross-checked each branch in the per-expert body and there is no stale tp capture inside the outer lambda.
• GetOptimalBlockSize(num_expert_tokens, inner_tp ? DegreeOfParallelism(inner_tp) : 1) correctly drops to a single-block sizing when inner is serial — the resulting num_blocks > 1 && inner_tp != nullptr gate then short-circuits to the unrolled-serial branch in the else. No dead pool calls.
• inner_tp is read-only inside the outer lambda and is set once before TrySimpleParallelFor — no closure-capture race.
• The tp == nullptr path collapses cleanly: max_expert_threads = 1, num_expert_threads = 1, inner_tp = nullptr. Same observable behavior as before for the single-threaded session config that the issue reporter used to work around the hang.
• The <numeric> include is correctly dropped now that std::accumulate is gone.

Bottom line

Approve. The fix kills a long-standing intermittent hang in a way that is provably single-level by construction, while also delivering the decode throughput numbers shown in the benchmark table. The two observations above are follow-up perf nits, not regressions on the workloads this PR targets, and the nit is a comment-style preference. Land it.

[yuslepukhin]

Findings (ordered by severity)

• Medium: missing deterministic regression test for the exact livelock scenario this PR fixes
  The PR changes only onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc, and adds no new test files or test cases. Existing parity coverage in onnxruntime/test/python/transformers/test_qmoe_cpu.py is broad for numerics, but it does not encode a targeted concurrent hang regression for nested intra-op pool use or the single-active-expert routing corner that motivated the follow-up commit.
  Why this matters: without a targeted stress/regression test, this can regress silently.

• Medium: increased per-thread workspace amplification can become a practical memory risk on large shapes
  After removing the old tier cap and setting threads by active experts at moe_quantization_cpu.cc#L1142, multiple large per-thread buffers are allocated using that multiplier at moe_quantization_cpu.cc#L1145, moe_quantization_cpu.cc#L1163, moe_quantization_cpu.cc#L1167.
  This is not a correctness bug by itself, but it is a meaningful operational risk for large models and high top-k/active-expert counts. I did not find a new overflow bug from this patch, but memory pressure can increase significantly relative to the prior capped behavior.