fix(cpu): give the reserved dispatcher CPU a compute lane (#1746)

[justinchuby]

Closes #1746.

ONNX_GENAI_CPU_DECODE_THREADS=N builds N-1 compute lanes in production, for every N. At N=2 that is one lane, which trips the total_workers <= 1 serial short-circuit in dispatch_output_rows — so the whole persistent pool degenerates to serial dispatch on the engine thread and the knob is indistinguishable from =1.

Root cause: two correct pieces composing badly

Neither half is wrong on its own, which is why unit tests on either half could not see this.

1. provider.rs:340 — EpFactory::initialize(), the earliest per-session hook, calls bound_process_to_decode_budget(). With an explicit budget it confines the process to exactly N CPUs, so that "a user who caps cores disturbs at most N CPUs". Deliberate and documented.
2. Later, on first decode, node_shards(N) reads allowed_cpus() — now exactly N — and reserve_single_group_headroom(N, N) returns N-1, keeping one CPU free for the inline dispatcher. Also deliberate and measured: N spinning workers pinned across all N allowed CPUs starve the dispatcher and collapse throughput 20-60x (1.47 tok/s at 32 workers on taskset -c 0-31 vs ~29 tok/s once one CPU is spare).

reserve_single_group_headroom's docstring notes it only fires for "a user who sets ONNX_GENAI_CPU_DECODE_THREADS=N on an exactly-N-CPU cpuset". We build that cpuset ourselves in step 1, so it is not a corner case — it is the guaranteed outcome of every explicit budget.

Measured through the production sequence, under an outer taskset -c 0,2,...,30:

onnx-genai: CPU decode budget 2 confined the process to 2 CPUs [0, 2]
PROD budget=2 allowed_before=Some(16) allowed_after=Some(2) spmd_threads=1
PROD budget=4 allowed_before=Some(16) allowed_after=Some(4) spmd_threads=3
PROD budget=8 allowed_before=Some(16) allowed_after=Some(8) spmd_threads=7

allowed_before=16 -> allowed_after=N is the whole mechanism.

The dispatcher does not compute — it publishes and spins in SharedState::wait. So the reserved CPU is not merely unallocated, it is burned spinning.

Fix

Keep the reservation exactly as measured, and let the dispatcher compute the shard it was holding a CPU for.

At budget N: N-1 pinned worker threads (unchanged — the starvation cliff stays fixed) plus the dispatcher computing on the reserved CPU = N compute lanes on N CPUs, no thread oversubscribed. Strictly better than N-1 lanes plus a spinning CPU.

budget	pinned threads	compute lanes (before)	compute lanes (after)
2	1	1 (serial short-circuit)	2
4	3	3	4
8	7	7	8
16	15	15	16

Mechanically: publish counts down node_thread_counts (spawned threads only — the dispatcher's shard has no thread to wait for), the dispatcher runs the remaining shard inline, then wait()s. Partitioning already produces total_workers shards, so only the pending-count source and the shard-to-thread mapping change. dispatch_inline (the re-entrant fallback) already loops over all shards and stays correct.

Scope. Single-group layouts only, and only when the reservation actually fired. A group that already had headroom is untouched, or the pool would be one lane wider than the budget allows. On a NUMA split the dispatcher's node is not known at build time, so handing it a shard could pull that shard's weights across sockets; those layouts keep the previous behaviour.

catch_unwind is load-bearing, not defensive

The published Job holds a raw pointer to a closure borrowed off the dispatcher's own stack frame, and the workers read through it until the barrier drains. Now that the dispatcher also computes, a panic in its shard would unwind that frame while workers are still reading it — a use-after-free, not merely a hang. So: catch, complete the barrier, then resume_unwind.

Miri agrees. Removing the catch_unwind:

error: Undefined Behavior: Data race detected between (1) non-atomic read on
thread `onnx-genai-spmd` and (2) retag write of type `{closure@decode_spmd.rs}`
on thread `decode_spmd::te` at alloc1106927

decode_spmd's panic-safety test is added to the Miri lane, with the "verified both ways" note the workflow already uses.

Falsifiers

Every test was checked to fail without the fix.

#	Mutation	Result
F1	Dispatcher never takes a shard	3 unit tests fail (...restores_the_requested_width, ...fans_out_across_the_dispatcher_and_one_worker, ...covers_its_rows_exactly_once)
F2	Remove catch_unwind	panic test fails natively; Miri reports UB (above)
F3	publish the shard counts instead of thread counts	barrier never drains — hangs (exit 124)
F4	Dispatcher never takes a shard	end-to-end subprocess test fails: ONNX_GENAI_CPU_DECODE_THREADS=2 must buy 2 compute lanes, got 1 (1 pinned threads on 2 allowed CPUs)

The end-to-end test spawns one subprocess per budget — not stylistic. Both halves latch (PROCESS_BUDGET_BOUND and the pool are OnceLocks) and the child mutates process-wide CPU affinity, which would poison the test runner. It skips budgets above available_parallelism() and skips NUMA-split layouts, so it is meaningful on a 2-vCPU runner and inert where it cannot apply.

No performance claim

This PR claims a width restoration, proven categorically by thread and lane counts, not a speedup. The host is shared and has been above loadavg 60; per the standing protocol I am not quoting a throughput number I could not measure under control. The arithmetic ceiling is 2x at N=2 and 1.33x at N=4, but realised speedup depends on scaling efficiency and is not asserted here.

Corrects the record on #1740

The merged width-scaling benchmark lists this exact mechanism as vacuity case 2, but classifies it as "benchmarking inside a small container hits this".

#1740's measurements are correct and nothing in them is retracted. int4_decode_loop_ab never calls EpFactory::initialize() — it goes straight to with_decode_pool_scope — so the confinement never runs there, the process keeps its 16 taskset CPUs, reserve_single_group_headroom(2, 16) = 2, and the bench genuinely gets two busy workers. Its 1.96x and its per-thread attribution stand. That is also exactly why its non-vacuity check passed (w in gave w out) while production lost a worker.

The finding is the divergence: the decode bench does not reproduce the production thread topology, so a width sweep run through it is structurally unable to observe this class of defect.

One correction to that table: its t=16 row is a 15-worker measurement, since reserve_single_group_headroom(16, 16) = 15 fires even without the production confinement. The 1/2/4/8 checks could not catch it because the reservation only triggers at exactly full subscription. This understates the plateau rather than overstating it, so the conclusion drawn from it is unaffected.

Addendum written into docs/benchmarks/2026-08-22-decode-width-scaling.md rather than left in a PR comment.

I also withdrew my own first reading of this: I initially attributed Roy's =1/=2 bench timings to this defect. That attribution was wrong — he was on the bench binary — and is retracted in #1746 and on #1722.

Validation

• cargo test -p onnx-runtime-ep-cpu --lib — 1616 passed, 0 failed
• cargo test --workspace (excluding onnx-genai-bench) — 6303 passed, 0 failed
• cargo fmt --check; cargo clippy --all-targets -D warnings — clean
• Feature configs: --features mlas (incl. the Steal path, which keeps its previous width), --no-default-features — clean
• Cross: aarch64-unknown-linux-gnu clippy -D warnings — clean
• All seven Rust-quality lint scripts — pass
• Miri: decode_spmd::tests::a_panic_in_the_dispatcher — ok in ~23s; reports UB with F2 applied
• Latest origin/main merged in before this run

[codecov]

Codecov Report

❌ Patch coverage is 95.92760% with 9 lines in your changes missing coverage. Please review.
✅ Project coverage is 80.07%. Comparing base (7af550a) to head (5373ae2).
⚠️ Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
crates/onnx-runtime-ep-cpu/src/decode_spmd.rs	95.92%	8 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1748      +/-   ##
==========================================
+ Coverage   80.03%   80.07%   +0.03%     
==========================================
  Files         410      410              
  Lines      198269   198477     +208     
  Branches   198269   198477     +208     
==========================================
+ Hits       158685   158923     +238     
+ Misses      34195    34157      -38     
- Partials     5389     5397       +8     

Flag	Coverage Δ
cli-ort-linux	72.47% <ø> (ø)
cli-ort-windows	72.06% <ø> (+0.09%)	⬆️
mlas	85.20% <ø> (ø)
offline	80.19% <95.92%> (+0.03%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
...es/onnx-runtime-ep-cpu/src/kernels/matmul_nbits.rs	79.88% <ø> (+0.08%)	⬆️
crates/onnx-runtime-ep-cpu/src/decode_spmd.rs	90.17% <95.92%> (+0.87%)	⬆️

... and 2 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[justinchuby]

Adversarial review: claude-opus-4.8

Brief was written to falsify, not confirm — 8 numbered claims, each with an
explicit "what would prove this wrong" clause, plus the standing instruction to
report a finding even if it only might be real.

Verdict: no blocking findings, no should-fix findings. Two NITs.

Claims the reviewer independently checked and confirmed

• The scope guard cannot oversubscribe. dispatcher_owns_a_shard is
  shards.len() == 1 && shards[0].workers < requested. Checked both
  directions: NumaTopology::split_workers returns None unless
  len() >= 2 (decode_affinity.rs:364), so a multi-node pool can never trip
  it; allowed_count == 1 is unreachable because build_from_env bails first.
• catch_unwind is load-bearing, not defensive. The published Job holds
  data: *const () derived from std::ptr::from_ref(job), and job: &F lives
  on the dispatcher's stack frame. Unwinding out of the inline shard before the
  barrier drains is a use-after-free, and the ordering
  catch -> wait() -> drop(claim) -> resume_unwind is the only correct one.
  Falsifier F2 backs this: removing the catch makes Miri report
  Undefined Behavior: Data race between (1) non-atomic read on thread 'onnx-genai-spmd' and (2) retag write.
• publish must count threads, not shards. F3 (publishing shard counts)
  hangs the barrier — exit 124 — rather than failing loudly, which is exactly
  why it is a separate field instead of a derived value.

A risk the reviewer raised that I had not checked

run_mlas_shards (matmul_nbits.rs:2867) does shards.get(global_index) and
silently skips a missing shard, so a dispatcher index past the end of the
shard list would drop work with no error. Verified safe: build_mlas_shards
sizes from the same total_workers-derived segment list, so both sides grew
together and shards.get(total_threads) is always Some. The mlas
DecodeSchedule::Steal path additionally opts out of the dispatcher shard
entirely (dispatcher_shard: None). Worth recording because the failure mode
would have been silent wrong answers, not a panic.

NITs

1. Deferred worker poison. If the dispatcher's inline shard and a worker
   both panic on the same op, resume_unwind returns before
   panic_if_poisoned(), so the worker's poison surfaces on the next
   dispatch. Latched, not lost; no hang and no corruption. The reviewer called
   it "arguably correct-by-design" — the dispatcher's payload is the one
   carrying the caller's stack. Addressed in 5373ae2 by documenting the
   contract at the site so it reads as a decision rather than an accident.
2. Subprocess test degrades on small runners. Fully skipped on a 1-vCPU
   runner, and covers only budget=2 on 2-vCPU. Accepted: F1-F4 live in
   always-run unit tests, so the real coverage is not silently lost when the
   runner is small.