Add: early-dispatch for require_sync_start SPMD cohorts (a2a3 tmr)

[ChaoWao]

Summary

Let require_sync_start SPMD cohorts pre-stage onto their cores while the
producer still runs, then launch together via a unified doorbell, instead of
sitting idle until the producer fully retires. Cohort launch collapses from a
serial per-core spread to a single atomic rendezvous.

• Pre-stage gated — a flagged (allow_early_resolve) producer bumps
  consumers' dispatch_fanin; a sync_start block reaching fanin_actual_count
  stages onto its cores with not_ready=1 (STAGING), waiting on a
  DATA_MAIN_BASE doorbell rather than launching early.
• Rendezvous launch — running_slot_count counts staged gated cores
  (shape-agnostic); when it reaches popcount(staged_core_mask) and the
  producer has released, maybe_rendezvous_ring() rings every staged doorbell
  atomically so the cohort starts as one.
• MIX split placement — a MIX cluster (1 AIC + 2 AIV) stages per-core into
  whatever running/pending slot each core has free; gated, so partial placement
  is safe. count_mix_split_clusters gates availability.
• Parallel drain — handle_drain_mode runs the stop-the-world sync_start
  drain across the dispatch threads, replacing the serial stage that spread
  cohort launch across the round.
• pending_gated fix — a sync_start block stays gated after the producer
  releases (DISPATCHED), so pending_gated = STAGING || (DISPATCHED && requires_sync_start()), else a released-but-still-gated pending block is
  misclassified as a normal task and deadlocks on an ack that never comes.
• Deadlock avoidance (early-candidate gate) — propagate_dispatch_fanin
  no-ops until every producer block range has reserved a core slot
  (next_block_idx == logical_block_num). A flagged SPMD producer with more
  blocks than cores (for example, a 50-block AIC projection on 24 AICs) dispatches
  in waves; if its first wave triggered a downstream MIX cohort to gate every
  slot, the producer's unclaimed blocks would never dispatch and the cohort's
  rendezvous would never ring. next_block_idx is a reservation counter, not a
  publication counter; it covers normal dispatch, early staging, and drain claims.
• Deferred propagation capacity — the post-publish prop_list now matches
  the full per-thread batch capacity, so a final AIV claim after item 21 cannot
  lose the only full-reservation propagation retry.
• Swimlane — Drain / DrainPrepare / DrainPublish sched phases (gated
  on PTO2_PROFILING + SCHED_PHASES) make the drain visible.
• Docs — RUNTIME_LOGIC.md §8.6: two-source × sync-tier dispatch model,
  drain + rendezvous, early-candidate gate, MIX per-core placement.

Testing

• qwen3-14B batch-16 decode (pypto-serving) completes at 1GB and 2GB
  ring heap — previously deadlocked every run (Heap Exhausted at 1GB,
  running-stalled at 2GB).
• Targeted AIC50 -> MIX24 require_sync_start regression passes on
  a2a3sim and onboard a2a3; the pre-fix runtime stalls.
• C++ non-hardware unit suite: 54/54 passed, including partial-claim vs
  full-claim propagation coverage.
• Related sim regression: sync-start, MIX, starvation, and SPMD scenes pass;
  spmd_batch_dispatch_oob also passes.
• Onboard st (a2a3): spmd_basic, spmd_multiblock_aiv,
  spmd_multiblock_mix, spmd_sync_start, spmd_sync_start_aiv,
  spmd_sync_start_early_dispatch, spmd_sync_start_edge,
  spmd_sync_start_mix_spill, spmd_starvation,
  spmd_sync_start_stress, spmd_batch_dispatch_oob,
  spmd_paged_attention, spmd_paged_attention_highperf,
  mixed_example — 14 passed.
• Full CI (sim + onboard) on the latest PR head

Notes

• New st spmd_sync_start_mix_spill/ exercises the MIX per-core pending-spill
  rendezvous path.
• spmd_sync_start_early_dispatch/ is the full-reservation gate regression:
  a 50-block flagged AIC producer feeds a 24-block MIX sync-start consumer.
• Profiling additions do not gate behavior (diagnostic-only, level-gated).

---

Rebased onto #1308 / #1319 (per-core MIX placement + sync_start Tier-0 now
in main): reuses their per-core placement (prepare_block_for_dispatch with
to_pending=true) and shared run_staging_order skeleton.

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR adds a gated sync_start drain protocol with a doorbell rendezvous mechanism to the a2a3 scheduler: a new running_slot_count counter, ring_all_staged_doorbells/maybe_rendezvous_ring helpers, parallel core staging via drain_stage_cores, MIX split-placement tracking, updated drain state fields, profiling enum extensions for drain phases, and a new end-to-end SPMD test.

Changes

Sync-start drain rendezvous feature

Layer / File(s)	Summary
Drain phase profiling enum src/a2a3/platform/include/common/l2_swimlane_profiling.h, src/a5/platform/include/common/l2_swimlane_profiling.h, src/common/platform/shared/host/l2_swimlane_collector.cpp	Adds Drain, DrainPrepare, DrainPublish sched-phase discriminators and maps them to string labels in JSON export.
Rendezvous counter field src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h	Adds running_slot_count atomic to PTO2TaskPayload, reset to 0 during init.
CoreTracker MIX split helpers & drain state fields src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.h	Adds gated MIX split-placement helper methods and extends SyncStartDrainState with staging coordination atomics.
Doorbell rendezvous helpers src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h	Adds ring_all_staged_doorbells/maybe_rendezvous_ring, updates propagate_dispatch_fanin comments and try_speculative_release gating logic.
Scheduler context API contract src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h	Updates prepare_block_for_dispatch, count_global_available, handle_drain_mode, removes drain_worker_dispatch, adds drain_stage_cores.
Parallel staging drain protocol rewrite src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp, scheduler_cold_path.cpp	Rewrites decide_slot_transition, enter_drain_mode, count_global_available, and handle_drain_mode to coordinate election, parallel staging, and rendezvous; resets new drain fields at teardown.
Dispatch integration src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp	Passes split_placement through MIX dispatch, routes sync_start consumers through enter_drain_mode, adds profiling-instrumented drain timing.
SPMD sync-start early-dispatch test tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_early_dispatch/*	Adds AIV write kernel, orchestration entry, and Python test verifying the new drain rendezvous behavior.

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

Sequence Diagram(s)

sequenceDiagram
  participant Dispatcher as scheduler_dispatch
  participant Completion as SchedulerContext (drain)
  participant Tracker as CoreTracker
  participant Scheduler as PTO2SchedulerState
  participant Payload as PTO2TaskPayload

  Dispatcher->>Completion: enter_drain_mode(consumer, logical_block_num)
  Completion->>Completion: reset drain_stage_go, drain_stage_done_mask, drain_running_staged
  Completion->>Completion: elect drain worker, publish sync_start_pending
  Completion->>Completion: handle_drain_mode(thread_idx, out_ack_wait_cycles)
  Completion->>Tracker: get_mix_split_cluster_offset_states / core_offset_to_pending
  Completion->>Completion: drain_stage_cores(...) parallel staging
  Completion->>Payload: seed running_slot_count for gated drain
  Completion->>Completion: sync on drain_stage_done_mask, reopen gate
  Completion->>Scheduler: propagate_dispatch_fanin
  Scheduler->>Scheduler: maybe_rendezvous_ring(slot_state)
  Scheduler->>Scheduler: ring doorbells when running_slot_count matches staged core count

Loading

Possibly related PRs

• hw-native-sys/simpler#1074: Both PRs modify the sync_start drain protocol in SchedulerContext::handle_drain_mode, directly related to the rewritten drain-mode logic here.
• hw-native-sys/simpler#1079: Both PRs implement speculative early-dispatch gating via PTO2TaskPayload spec fields and pto_scheduler.h doorbell/rendezvous logic.
• hw-native-sys/simpler#1108: Both PRs extend L2SwimlaneSchedPhaseKind and the corresponding l2_swimlane_collector phase-name mappings for new scheduler phase kinds.

Poem

A rabbit taps its paw and waits,
for cores to stage through gated gates,
when all are counted, running true,
the doorbell rings — ding-ding, ka-choo! 🔔
Drain complete, the fan-in flows,
off to running-slot it goes! 🐇✨

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 47.37% 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	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly summarizes the main change: early-dispatch support for require_sync_start SPMD cohorts.
Description check	✅ Passed	The description is directly about the sync_start early-dispatch and drain changes in this PR.

---

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.}

[gemini-code-assist]

Code Review

This pull request parallelizes the sync_start drain staging and introduces a rendezvous mechanism to prevent deadlocks during early dispatch. It also adds split placement for MIX shape blocks, updates completion and dispatch logic, and introduces profiling support for the drain phases alongside a new system test. The review feedback highlights two important issues: a potential out-of-bounds memory write in drain_stage_cores due to a missing bounds check on the core ID, and undefined behavior from signed integer overflow in the test kernel's spin loop.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[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
`@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h`:
- Around line 351-353: The `handle_drain_mode` profiling output is mislabeled:
`out_ack_wait_cycles` in `scheduler_context.h` still reads like ack-barrier
time, but the implementation records stage/finalize wall time. Update the
parameter name and surrounding comments to match the actual value, or change the
implementation if ack-wait was intended, so the meaning is consistent across
`handle_drain_mode` and its caller-facing profiling data.

🪄 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

Run ID: 49a3f1c5-87c4-4fb6-9f7d-19acdf8ffe44

📥 Commits

Reviewing files that changed from the base of the PR and between 3a23d22 and d1aee9e.

📒 Files selected for processing (13)

• src/a2a3/platform/include/common/l2_swimlane_profiling.h
• src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h
• src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h
• src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp
• 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/platform/include/common/l2_swimlane_profiling.h
• src/common/platform/shared/host/l2_swimlane_collector.cpp
• tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_early_dispatch/kernels/aiv/kernel_spmd_write_slow.cpp
• tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_early_dispatch/kernels/orchestration/spmd_sync_start_early_dispatch_orch.cpp
• tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_early_dispatch/test_spmd_sync_start_early_dispatch.py