fix(runtime): avoid global drain for local early sync-start

[vegetabledoww]

Summary

This PR fixes #1548 by adding a local-owner fast path for early require_sync_start cohorts. When one scheduler's tracker can place every block using its local idle running slots and free pending slots, that scheduler stages the complete cohort locally. Cohorts that do not fit on one owner continue to use the existing global stop-the-world drain.

The behavior is implemented consistently for:

• A2A3 host_build_graph
• A2A3 tensormap_and_ringbuffer
• A5 tensormap_and_ringbuffer

Problem

Early sync_start staging must be all-or-nothing. Partially staging a gated cohort can strand blocks before the rendezvous reaches block_num, so the previous implementation conservatively routed every early sync_start cohort through the global drain protocol.

That was unnecessarily expensive when the scheduler that claimed the cohort already had enough local capacity. In the representative 8-block AIV case from #1548, one scheduler could hold the entire cohort, but all scheduler threads still had to stop, acknowledge the drain generation, scan capacity, prepare blocks, and publish them globally.

The minimal regression shows the behavioral difference clearly:

• Before: drain_prepare=1, drain_publish=1, drain=1, with 8/8 blocks published by the drain.
• After: early_dispatch=8, drain=0, with all blocks staged by one local owner.

Root cause

The Tier-0 early sync_start path claimed a per-task owner and immediately attempted enter_drain_mode(...). It had no exact local-capacity check and no local all-or-nothing staging branch, even though the existing tracker state already described idle and pending placement capacity.

What changed

• Added shared count_available_blocks(...) accounting for AIC, AIV, and MIX shapes.
  • AIC/AIV count idle cores and, for gated early staging, running cores with free pending slots.
  • MIX counts active-mask-compatible logical clusters rather than individual cores.
  • Occupied pending slots are excluded.
• Generalized the staging primitive as stage_sync_start_cores(...) so local and global paths use the same placement semantics.
• Added an exact-fit local-owner path before global drain entry.
  • The owner must fit the entire logical_block_num locally.
  • Staging remains all-or-nothing; partial local staging is not allowed.
  • running_slot_count is seeded from directly staged running cores, then the producer/stager rendezvous is retried.
• Preserved the original global drain as the fallback whenever local capacity is insufficient or a global drain is already pending.
• Preserved one-shot launch and doorbell semantics across producer-first and stager-first races.
• Added a distinct early_dispatch scheduler phase on A5 and updated runtime documentation/comments to distinguish local staging from global drain.

Correctness and concurrency guarantees

The fast path keeps the same safety properties as the global path:

• ownership is claimed before publishing gated payloads;
• the local capacity check covers the complete cohort;
• the staged block count is asserted to equal logical_block_num;
• producer release and staging completion may arrive in either order;
• launch ownership remains one-shot, preventing duplicate doorbell ringing or fanout propagation;
• capacity-short cohorts never partially stage locally before falling back.

Test coverage

Unit tests cover both A2A3 and A5 implementations, including:

• AIC/AIV idle and pending-slot capacity;
• occupied pending-slot exclusion;
• MIX logical-cluster accounting and active-mask handling;
• launch ownership under concurrent claimants;
• producer-first and stager-first rendezvous ordering;
• idempotent retry and completion/release interleavings.

Scene regressions now cover three topologies on both A2A3 and A5:

1. one scheduler, 8-block idle-capacity cohort;
2. three schedulers, 4-block cohort that fits any local owner;
3. one scheduler, 8-block pending-only cohort with every AIV running slot occupied.

The pending-only scene uses blocker start markers to ensure the consumer is submitted after blocker ACKs are observable. Its trace assertion also proves local staging finishes before any blocker releases a running slot, so the test exercises real pending-slot placement rather than an idle fallback.

Validation results

• A2A3 simulator: all three regression cases passed.
• A5 simulator: all three regression cases passed.
• Real A2A3 hardware: all three cases passed on device 1.

Real A2A3 case	Scheduler owners observed	early_dispatch blocks	Drain records
Single scheduler, idle	1	8	0
Three schedulers, idle	3	4	0
Single scheduler, pending-only	1	8	0

Additional representative l3_moe validation passed golden checks on both ranks and showed no drain. A same-card 100-round ABBA comparison measured a median of 750.7 us for this PR versus 753.3 us for clean main (about -0.35%, within measurement noise), with no observed performance regression.

Relevant C++ scheduler-state and wiring tests pass for A2A3 and A5, and all repository pre-commit checks pass (clang-format, clang-tidy, cpplint, Ruff, Pyright, header/license checks, and whitespace checks).

Expected impact

Workloads whose early sync_start cohort fits one scheduler avoid unnecessary cross-thread coordination and stop-the-world drain overhead. Larger cohorts retain the existing global aggregation behavior. There is no public API change.

The exact historical fused_pre_norm_cce.py workload referenced by #1548 is not present on current public main, so the PR uses a minimal before/after regression to reproduce the scheduler condition and representative l3_moe runs for compatibility and performance validation.

Fixes #1548

[coderabbitai]

📝 Walkthrough

Walkthrough

The scheduler now supports capacity-aware local staging for early sync_start cohorts, returns structured staging results, tightens rendezvous publication ordering, updates runtime documentation, and adds unit, wiring, and single-owner swimlane coverage across a2a3, a5, and host-build-graph variants.

Changes

Sync-start staging flow

Layer / File(s)	Summary
Capacity and staging results src/a2a3/.../scheduler/*, src/a5/.../scheduler/*	CoreTracker exposes capacity counting, while stage_sync_start_cores returns staged-block and running-core counts used by drain handling.
Local dispatch and rendezvous gating src/a2a3/.../scheduler/*, src/a5/.../scheduler/*	Early sync-start dispatch checks full-cohort capacity before staging, and rendezvous logic loads the running-slot seed before staged-core publication.
Capacity and rendezvous validation tests/ut/cpp/a2a3/*, tests/ut/cpp/a5/*	Tests cover shape capacity, doorbell ownership, staging retries, producer-release ordering, and early-sync completion.
Runtime documentation and local-owner scene coverage src/**/docs/RUNTIME_LOGIC.md, src/common/platform/include/common/l2_swimlane_profiling.h, tests/st/**/sync_start_early_local_owner*	Documentation describes local/global staging and dormant producer propagation; scene tests verify local staging and swimlane phase output.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Scheduler
  participant CoreTracker
  participant Staging
  participant Rendezvous
  Scheduler->>CoreTracker: count_available_blocks
  CoreTracker-->>Scheduler: available logical blocks
  Scheduler->>Staging: stage_sync_start_cores
  Staging-->>Scheduler: staged_blocks and running_cores
  Scheduler->>Rendezvous: publish staged mask and running count
  Rendezvous->>Rendezvous: retry after staging

Loading

Possibly related PRs

• hw-native-sys/simpler#1423: Refactors the corresponding a2a3 sync-start staging and rendezvous semantics.
• hw-native-sys/simpler#1435: Stubs early producer propagation used by the updated scheduler documentation.
• hw-native-sys/simpler#1461: Updates related drain coordination and generation-safety machinery.

Poem

A rabbit watched the cores align,
Local blocks hopped into line.
The doorbell rang when counts were right,
No needless drain delayed the flight.
“Sync-start sprouts!” the rabbit sings.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 15.79% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	The PR appears to implement the #1548 fast path, local capacity check, rendezvous handling, and coverage requested by the issue.
Out of Scope Changes check	✅ Passed	The docs, runtime changes, and tests all support the same sync-start fast-path work and do not appear unrelated.
Title check	✅ Passed	The title clearly captures the main change: adding a local early sync-start path to avoid unnecessary global drain.
Description check	✅ Passed	The description is directly aligned with the changeset and accurately summarizes the new local-owner early sync-start fast path.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In
`@tests/st/a5/tensormap_and_ringbuffer/dfx/l2_swimlane/test_sync_start_early_local_owner.py`:
- Around line 94-106: Extend the assertions in the perf-phase validation near
phase_counts to verify that the local-owner path contains no records with phases
“drain”, “drain_prepare”, or “drain_publish”, mirroring the A2A3 test behavior.
Keep the existing early_dispatch count assertion and include the artifact path
in any new failure message.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 09dcf435-c281-4e1f-ad67-f88dabc62f8d

📥 Commits

Reviewing files that changed from the base of the PR and between 5bc11c2 and 972374b.

📒 Files selected for processing (27)

• src/a2a3/runtime/host_build_graph/docs/RUNTIME_LOGIC.md
• src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h
• src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp
• src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h
• src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp
• src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_types.h
• src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md
• src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h
• src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp
• src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h
• src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp
• src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.h
• src/a5/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md
• src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h
• src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp
• src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h
• src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp
• src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.h
• src/common/platform/include/common/l2_swimlane_profiling.h
• tests/st/a2a3/tensormap_and_ringbuffer/dfx/l2_swimlane/kernels/orchestration/sync_start_early_local_owner_orch.cpp
• tests/st/a2a3/tensormap_and_ringbuffer/dfx/l2_swimlane/test_sync_start_early_local_owner.py
• tests/st/a5/tensormap_and_ringbuffer/dfx/l2_swimlane/kernels/orchestration/sync_start_early_local_owner_orch.cpp
• tests/st/a5/tensormap_and_ringbuffer/dfx/l2_swimlane/test_sync_start_early_local_owner.py
• tests/ut/cpp/a2a3/test_scheduler_state.cpp
• tests/ut/cpp/a2a3/test_wiring.cpp
• tests/ut/cpp/a5/test_scheduler_state.cpp
• tests/ut/cpp/a5/test_wiring.cpp