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 active end-to-end fast path is implemented and validated for:

• A2A3 tensormap_and_ringbuffer
• A5 tensormap_and_ringbuffer

A2A3 host_build_graph carries the same scheduler-side capacity and staging machinery for implementation parity, but its propagate_dispatch_fanin path is currently a stub. It therefore does not populate the early queues, and the HBG fast path remains dormant rather than being an active end-to-end behavior.

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.
• Updated runtime documentation and code comments to distinguish local staging from the global drain fallback and to document HBG's dormant status.

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 scheduler 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 the A2A3 and A5 tensormap_and_ringbuffer runtimes:

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

In the active A2A3 and A5 tensormap_and_ringbuffer early-propagation paths, 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.

A2A3 host_build_graph contains the scheduler-side implementation but does not currently expose this optimization end to end because early producer propagation is disabled there; no active HBG performance benefit is claimed by this PR.

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]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: d72d9bb4-25b6-4b20-981f-7723fabdbefc

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 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 summarizes the main change: using a local early sync-start path instead of always entering global drain.
Description check	✅ Passed	The description is strongly related to the changeset and accurately describes the local fast path, fallback drain behavior, and added tests.

---

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

[ChaoZheng109]

文档一致性:重命名后有 2 处 drain_stage_cores 旧名残留

本 PR 已将运行时函数 drain_stage_cores 重命名为 stage_sync_start_cores(三棵树的实现与调用点均已同步,👍)。但 spmd_sync_start_mix_spill 测试目录下仍有两处散文注释引用旧函数名,建议在本 PR 一并更新,避免遗留失效引用(参见 .claude/rules/doc-consistency.md §1:重命名后需审计全仓引用):

• tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/kernels/orchestration/spmd_sync_start_mix_spill_orch.cpp:20

  * cores' gated pending slots (drain_stage_cores takes the to_pending=true split path,

• tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/test_spmd_sync_start_mix_spill.py:14

  (drain_stage_cores to_pending=true, mix_cluster_idle_core_count=1/cluster + Case 3.3 promote for

这两个文件不在本 PR 的 diff 内,故无法作为行内评论,统一在此说明。属于轻微(minor)问题,不阻塞合入。

其余部分 LGTM:Case A 逻辑在 a2a3 host_build_graph / a2a3 tmr / a5 tmr 三棵树逐字一致,发布顺序严格镜像 drain finalize,maybe_rendezvous_ring 内存序已加固,单元/scene 测试覆盖到位(scene 断言非 owner 线程无 drain 相位)。

[vegetabledoww]

@ChaoZheng109 感谢指出。两处散文注释中的旧函数名已在最新提交 6f7141fc 中更新为 stage_sync_start_cores：

• tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/kernels/orchestration/spmd_sync_start_mix_spill_orch.cpp
• tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/test_spmd_sync_start_mix_spill.py

同时已按文档一致性规则重新执行全仓搜索，drain_stage_cores 当前为零命中；相关定向 pre-commit 检查也已全部通过。