[Performance] early_sync_start cohort always enters global drain even when a single scheduler thread owns enough cores (add single-owner Case A fast path)

[ChaoZheng109]

Platform

a2a3 (Ascend 910B/C hardware)

Runtime Variant

tensormap_and_ringbuffer

Summary

The sync_start SPMD dispatch has two lanes, and they are asymmetric in how they handle the "one scheduler thread already owns enough cores" case:

• Ready (producer done) sync_start — scheduler/scheduler_dispatch.cpp:417-434 — checks available >= logical_block_num on the popping thread's own cores. If it fits, it dispatches inline on its own cores (Case A); only if it does not fit does it enter_drain_mode (Case B).
• Early / gated sync_start (ED) — scheduler/scheduler_dispatch.cpp:822-833, the Tier-0 lane in try_early_dispatch — pops from the shape-agnostic early_sync_start_queue and unconditionally calls enter_drain_mode. There is no available >= block_num branch at all: every early sync cohort arms the global stop-the-world drain barrier, even when a single scheduler thread already owns enough idle+pending cores to host the entire cohort by itself.

enter_drain_mode stops all active scheduler threads (sync_start_pending != 0 -> every thread continues past dispatch, scheduler_dispatch.cpp:1138) so each thread can stage its own share of a cross-thread cohort. When the cohort fits on one owner, dragging the other threads through the ack-tree barrier + drain_stage_go handshake buys nothing: per-thread core ownership already guarantees no other thread can touch the owner's cores, so there is no race for the barrier to protect against, and there is no second stager to coordinate with for the rendezvous seed.

Proposal: add a single-owner Case A to the early_sync_start Tier-0 lane, mirroring the ready lane. On pop, if the popping thread's own idle + pending cores (gated availability; count_mix_split_clusters for MIX) already satisfy logical_block_num, stage the cohort locally, seed running_slot_count itself, and call retry_sync_start_rendezvous_after_drain without arming enter_drain_mode. Fall back to today's global drain only when the popping thread is insufficient. This is purely additive — the drain path is untouched for the genuinely cross-thread case.

Note the one ED-specific difference from the ready Case A: the producer is not done, so Case A here cannot ring immediately. It must gate-stage (prepare_block_for_dispatch(..., gated=true)), record staged_core_mask + the doorbell table, seed running_slot_count, then defer the ring to the rendezvous (producer-release half or pending->running promotion half). This is exactly what drain_stage_cores + the coordinator finalize already do per-thread — the fast path reuses those primitives for a single thread, minus the barrier.

Applies to both the a2a3 and a5 tensormap_and_ringbuffer runtimes (same scheduler structure).

Git Commit ID

f293be3

Host Platform

Linux (aarch64)

Reproduction

Observed on a pypto-lib l3_moe run (a2a3), captured in an L2 swimlane. The 8-block AIV SPMD task fused_pre_norm_cce_spmd (r1t4) is an early (gated) sync_start cohort whose two producers had not yet completed when the drain appeared.

# Any workload that submits an early-resolve-flagged producer feeding an
# 8-block (block_num <= per-thread AIV core count) sync_start AIV SPMD
# consumer reproduces it. l3_moe is one such workload.
# Capture with the L2 swimlane at SCHED_PHASES level and inspect the
# merged swimlane JSON for the drain(...) / drain_prepare / drain_publish bars
# on the Scheduler View (pid=2), plus the consumer's block launch timestamps.

Expected Performance

When a single scheduler thread owns enough idle+pending cores for the whole cohort:

• Only that one owner thread does drain work (or, ideally, a barrier-free local stage).
• The other scheduler threads keep dispatching / processing completions with zero stall.
• The cohort's atomic launch timestamp is no later than today (it should be no worse, and may be earlier by skipping the ack-tree handshake and any availability-retry spin).

Actual Performance

From the l3_moe swimlane (merged_swimlane_20260728_002506.json, Scheduler View, pid=2). The r1t4 cohort's 8 AIV blocks all launch atomically at ts=89.36 us, and all 8 blocks were staged by a single thread (Sched_2) — yet all three scheduler threads were stopped:

Thread	Event	ts (us)	dur (us)	Staged blocks?
Sched_0 (30000)	drain(135)	76.60	15.24	none (pure barrier stall)
Sched_1 (30010)	drain(127)	82.80	7.92	none (pure barrier stall)
Sched_2 (30020)	drain(247)	83.36	7.34	(barrier)
Sched_2 (30020)	drain_prepare(8)	87.56	1.80	all 8
Sched_2 (30020)	drain_publish(8)	89.36	0.68	all 8

Sched_0 lost ~15 us and Sched_1 ~8 us of dispatch bandwidth to a barrier in which they staged nothing, while the only useful staging (Sched_2 prepare+publish) totaled ~2.5 us. The drain was armed at ts=76.6 but staging did not land until 87.56 — the ~11 us gap is availability-retry spin (all three threads idling in the barrier) while the producers were still freeing cores. In a step whose whole scheduler timeline is ~90 us, freezing two of three scheduler threads for this long is a material throughput loss, repeated for every early sync_start cohort that fits a single owner.

Profiling Data (Optional)

• r1t4 = fused_pre_norm_cce_spmd, AIV shape, block_num = 8, launched on AIV cores 28/29/34/35/40/41/46/47 (all Sched_2-owned), atomic launch at ts=89.36.
• Only Sched_2 emitted drain_prepare(8) / drain_publish(8); Sched_0 and Sched_1 emitted bare drain(...) bars = barrier participation with no staging.

Additional Context

Code anchors:

• Ready lane fast path (Case A/B): src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp:417-434
• Early lane, unconditional drain (no Case A): .../scheduler_dispatch.cpp:822-833
• Drain stop-the-world gate: .../scheduler_dispatch.cpp:1138
• Per-thread gated staging + rendezvous seed to reuse: .../scheduler/scheduler_completion.cpp:560 (drain_stage_cores), :789-818 (coordinator finalize / seed / retry_sync_start_rendezvous_after_drain)
• Rendezvous latch: .../scheduler/pto_scheduler.h:843 (maybe_rendezvous_ring), :761 (try_claim_early_dispatch_launch), :789 (try_claim_early_sync_drain)

Strict testing / validation required before merge (this touches a concurrency-sensitive path that already had one ABA regression — see docs/investigations/2026-07-sync-start-drain-retry-aba.md / #1455, so correctness proof must precede any perf claim):

1. Mutual exclusion: the local fast path must still claim ownership via try_claim_early_sync_drain and must not race a concurrent global drain or a second popper of the same cohort; early_sync_drain_state / sync_start_pending must be left consistent on both the fast-path and fallback branches.
2. Rendezvous race (highest risk): the fast path bypasses the barrier's natural serialization point, so the staged_core_mask.fetch_or (seq_cst) -> running_slot_count.store (seq_cst) -> retry_sync_start_rendezvous_after_drain ordering must exactly match the drain finalize, so a producer that releases concurrently and the local seed hand off through the NONE->RINGING->COMPLETE launch latch with exactly one ringer (no double-ring, no missed ring).
3. Coverage: 1/2/3/4 active scheduler threads; AIC, AIV, and MIX (split-placement) shapes; cohorts that fit a single owner (fast path) AND cohorts that must span threads (fallback to global drain, unchanged); the gated pending->running promotion half firing the rendezvous.
4. Regression barriers: existing spmd_sync_start normal + stress scene tests, the a2a3/a5 tmr and host_build_graph runtimes, plus the generation-token unit tests from [Bug] sync_start drain can deadlock when generation-less ack/election state is reused across retry attempts #1455 must all pass. Add a scene/unit test that asserts the single-owner early cohort launches without entering drain (e.g. no drain bar on the non-owner threads).
5. Perf validation: re-run the l3_moe swimlane and confirm the non-owner threads' drain(...) bars disappear and r1t4's atomic launch ts does not regress; ABBA 100-round comparison (per the [Bug] sync_start drain can deadlock when generation-less ack/election state is reused across retry attempts #1455 methodology) on a locked a2a3 device showing no regression on non-drain cases.

Question for maintainers: is this fast path acceptable to add, or is there a reason (beyond mechanism reuse) the early sync lane was deliberately funneled entirely through the drain barrier?

Related: #1455