perf(cpu): bound the decode dispatcher's yield loop and park it instead

[justinchuby]

The defect

SharedState::wait — the dispatcher side of the SPMD decode barrier — called thread::yield_now() on every iteration once past DISPATCHER_SPIN_BEFORE_YIELD (~10 µs), with no upper bound:

if spins >= DISPATCHER_SPIN_BEFORE_YIELD {
    if !yielded { yielded = true; /* counter */ }
    thread::yield_now();          // every iteration, forever
}

So the dispatcher hammered sched_yield for the entire remaining duration of every dispatch it did not win the race on. The cost is proportional to the shard length, not to the wake latency the yield backstop exists to hide. On llama the decode path issues 5 dispatches per token with millisecond-scale shards, so this is milliseconds of kernel time per token spent producing nothing.

Measurement

Interleaved within a single binary via the new env knob, so the two arms cannot be separated by host drift or by a rebuild. Zero inter-token gap, one worker per physical core, quiet host, PROBE_NULL=0, 300 tokens × 3 reps, 3 launches per arm plus an A/A partner.

width	arm	wall ms/token	total cpu ms/token	sys ms/token
2	before	35.72–35.87	37.36–37.63	15.74–16.10
2	after	35.81–35.93	35.93–36.85	0.59–0.66
3	before	23.92–24.06	37.22–38.67	11.99–12.33
3	after	23.93–24.07	36.50–36.99	1.88–2.21
4	before	17.96–18.10	37.87–38.21	9.85–10.33
4	after	18.05–18.09	37.26–37.63	2.39–2.51

• System CPU per token falls 24× / 6× / 4×, with no overlap between the arms at any width.
• Total CPU falls too, so this is not a sys→user relabel.
• Wall time is unchanged, inside the before-arm's own launch-to-launch spread (0.44% at width 2).

At width 2 the before-arm burns 3.2 ms of kernel time per dispatch against a ~7 ms dispatch — an entire core, in the scheduler, doing nothing.

Why this is not "spin harder"

Substituting spinning for yielding in worker_wait was previously measured to be a wash: kernel time down 5.5×, user time up by the same, total CPU flat 101.4 → 102.3. That result is real and it does not transfer here, for a reason worth stating explicitly: a worker that stops yielding still burns its core in user mode, so only the label changes. A dispatcher that parks burns nothing at all. That is why total CPU moves in this table and did not in that one.

The change

Escalation is now spin → a bounded number of yields → park:

• the yield window (default 128, ~100 µs) still covers the case the backstop exists for — a straggler that is merely descheduled and needs a core to finish a nearly-done shard — while bounding the cost when the shard is genuinely long;
• past it the dispatcher parks on a new completion_sense futex, bumped by the worker that retires the last outstanding shard.

The completing worker skips the wake syscall entirely unless the dispatcher is actually parked, so the fast path costs one load and one RMW per dispatch, not per worker. The dispatcher_parked / completion_sense handshake is the standard two-sided futex protocol and needs SeqCst on both store-then-load pairs; the reasoning is written out in park_until_complete.

ONNX_GENAI_CPU_DECODE_DISPATCHER_YIELDS overrides the budget. A very large value restores the exact previous behaviour — that is how the before-arm above was taken.

Tests

Two new tests, both verified by mutation (each fails under its own defect):

• a_parked_dispatcher_is_woken_by_the_last_worker_to_finish — asserted through recv_timeout on a channel rather than by calling wait on the test thread, because a lost wakeup is a hang, and a hang in CI reads as an infrastructure timeout rather than as this test failing. Deleting wake_all makes it fail in 10.01 s.
• repeated_parks_never_lose_a_wakeup — 16 consecutive ops, since completion_sense is monotonic and a stale observed value would strand exactly one iteration.

One honest note on the second: it was vacuous when first written. Without a sleep before signalling, the completion lands before the dispatcher exhausts its spin budget, park_until_complete returns on the pre-park re-check, and the loop asserts nothing about waking — confirmed by it passing unchanged with wake_all deleted. The sleep that forces a real park was added and the mutation then fails it. The comment records this so it is not re-broken.

I also wrote and then deleted a third test asserting bounds on the default constant: clippy correctly flagged it as a constant-valued assertion, which is exactly the near-vacuous kind of coverage that should not be added.

The SeqCst pair is deliberately not claimed as covered — x86 is TSO, so a weakened ordering still passes here and would only fail on a weakly ordered target. It is argued for, not asserted, and the docstring says so.

Scope and limits

• Full crate suite green: 1689 passed, 0 failed. Clippy clean, cargo fmt applied.
• The win is largest exactly when the dispatcher waits longest. In runs where the machine sits in a faster regime the before-arm barely yields (dy/token 1.3 vs 4.2) and there is correspondingly little to reclaim — the change is neutral there, not negative. I am separately chasing what selects that regime; it is not caused by, and does not depend on, this change.
• Widths 8/16, non-zero gaps and 2/4 concurrent sessions were also run. sys is consistently lower in the after-arm across all of them, but launch-to-launch wall variance at those points is far larger than any arm difference, so I am not claiming a wall result there and the table above is confined to the widths where the arms separate cleanly.

---

Update: adversarial review found a blocking bug in the first commit

bb8b21f fixes a deadlock, not a slowdown, and it is worth reading before the rest.

Each node's last worker reaches signal_completion having just stored zero to its own pending counter, then loads every other node's counter to decide whether it owes the dispatcher a wakeup. With only the release store and acquire loads that is the store-buffering litmus — in a two-node pool both last workers may each read the other's pre-decrement value, so neither passes the gate, nobody bumps completion_sense, and a dispatcher that has already parked sleeps until the process exits. Parking before completion is the normal case for this feature, so the only rare ingredient is the two cross-node decrements racing.

x86_64 hides it entirely: the lock-prefixed decrement is already a full barrier, so at least one worker always observes all-zero. aarch64 does not. The full suite passing on this host was therefore not evidence.

The fix is a SeqCst fence between the decrement and the scan. That puts every arrival into one total order, so the worker whose fence is last in it is sequenced after all the other zero stores and cannot miss them — one guaranteed signaller for any node count, not just two. Pairwise reasoning is not enough here: with three nodes a cycle of pairwise misses is self-consistent, and only the total-order argument rules it out. It runs once per node per op, on the path that was already about to signal, never on the per-worker fast path.

I had also written the park_until_complete doc as if the dispatcher_parked/completion_sense pair were the whole proof. It is necessary but not sufficient — it orders the dispatcher against a worker that reaches the bump and says nothing about whether one does. Both halves are now written down and cross-referenced.

Also from the review

• New test a_parked_dispatcher_waits_for_every_node_not_just_the_first. No previous test reached the cross-node scan at all — both build a single node. It asserts that draining the first of two nodes does not release the dispatcher and that draining the second does. Mutation-verified: deleting wake_all makes it time out at 10.30 s. It explicitly does not claim to cover the ordering hazard — on x86 the fence can be deleted and it still passes, and the docstring says exactly that.
• repeated_parks_never_lose_a_wakeup carried a rationale I had not falsified — that a stale observed would strand an iteration. It would not: a stale observed fails the completion_sense == observed guard, so the wait is skipped and the caller's loop still exits on all_workers_done. Confirmed by freezing observed in a mutant — the test passes. That is the second unfalsified claim I have had to retract on this test, so I replaced it with coverage that is actually unique to it: completion_sense advances exactly once per op, and dispatcher_parked is cleared on the way out. The latter is mutation-verified (removing the reset fails it immediately). Its pre-signal sleep also went 5 ms → 50 ms to match the single-op test; a loaded runner could otherwise land the signal before the dispatcher parked and silently drop that iteration onto the pre-park re-check.
• barrier_only_shared_state now takes per-node pending counts rather than one count, so multi-node fixtures are expressible.
• Dropped "an entire core burned in the scheduler" for the measured split: the core is fully burned, but only ~45% of it is inside sched_yield and the rest is spinning between the calls.

Re-validated after the fix: full lib suite 1658 passed, 0 failed, 22 ignored; clippy -D warnings clean on --all-targets; cargo fmt applied. The fence is one mfence per node per dispatch against a multi-millisecond dispatch, so it is not measurable in the table above.

---

Update 2: the "regime" caveat above is now explained, and it makes this change matter more

I described the win as regime-dependent without knowing what selected the regime. I do now, and it is not mysterious — it is CPU contention on the confined core set, which is the condition this change exists for.

ONNX_GENAI_CPU_DECODE_THREADS=N confines the whole process to N CPUs. At N=2 that is cpus [0, 2]. Because a dispatch is a barrier, a single foreign thread landing on one of those two CPUs halves the throughput of the whole dispatch — the healthy shard finishes and waits. The result is a clean 2× wall regression at unchanged CPU per token, which is exactly the "slow regime" signature I had been unable to attribute.

Controlled, reversible demonstration (width 2, one spin process pinned with taskset):

condition	ms/token	cpu ms/token
clean	17.64	35.37
spinner pinned to cpu 0 (inside the confined set)	35.86	36.45
spinner pinned to cpu 0 (repeat)	35.87	36.18
spinner removed	17.71	35.55
spinner pinned to cpu 5 (outside the confined set)	17.77	35.63

Turning it on and off flips the regime; placing the same load outside the confined set does nothing.

With the cause under control, the two arms separate cleanly and the picture inverts from "sometimes neutral" to "neutral when idle, large when contended":

condition	arm	ms/token	cpu ms/token	sys ms/token	dy/token
clean	before	17.70	35.78	0.22	0.53
clean	after	17.67	35.68	0.07	0.34
co-tenant on cpu 0	before	35.72	37.66	16.89	4.30
co-tenant on cpu 0	after	35.85	36.02	0.63	4.24
co-tenant on cpu 0	before	35.72	37.17	16.17	4.18
co-tenant on cpu 0	after	35.79	35.99	0.63	4.20

On an idle machine the dispatcher barely waits, so there is little to reclaim and the change is a small improvement. Under contention it reclaims 26× on sys and ~1.5 ms/token of total CPU — the yield loop was spending 45% of a core fighting the very co-tenant that was already the problem. That is the production case (concurrent sessions, co-tenanted hosts, oversubscribed containers), not the benchmark case, and it is where an unbounded sched_yield loop is worst: it burns most exactly when the machine has least to spare.

Note dy/token is ~4.2 in both arms under contention. The dispatcher still enters the yield path just as often; the change is that it now stops after a bounded number and parks instead of yielding for the rest of the shard.

This also retires the earlier caveat honestly rather than quietly: the wall-time numbers in the first table were taken under an uncontrolled co-tenant, so their absolute values reflect a contended host. The arm-to-arm sys and total-CPU deltas were always interleaved within one binary and are unaffected, and the table above re-establishes them with contention as a controlled variable rather than an unknown one.

[codecov]

Codecov Report

❌ Patch coverage is 97.41935% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 80.33%. Comparing base (add43fc) to head (4fadc2e).
⚠️ Report is 4 commits behind head on main.

Files with missing lines	Patch %	Lines
crates/onnx-runtime-ep-cpu/src/decode_spmd.rs	97.41%	2 Missing and 2 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1808      +/-   ##
==========================================
+ Coverage   80.30%   80.33%   +0.02%     
==========================================
  Files         411      411              
  Lines      200604   200895     +291     
  Branches   200604   200895     +291     
==========================================
+ Hits       161100   161387     +287     
+ Misses      34034    34028       -6     
- Partials     5470     5480      +10     

Flag	Coverage Δ
cli-ort-linux	72.47% <ø> (ø)
cli-ort-windows	72.06% <ø> (ø)
mlas	85.10% <ø> (-0.10%)	⬇️
offline	80.47% <97.41%> (+0.02%)	⬆️

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

Files with missing lines	Coverage Δ
crates/onnx-runtime-ep-cpu/src/decode_spmd.rs	90.63% <97.41%> (+0.72%)	⬆️

... and 5 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.