perf(tmr): barrier-free per-thread scheduler init (a2a3 + a5) - #1345
Merged
ChaoWao merged 1 commit intoJul 20, 2026
Merged
Conversation
📝 WalkthroughWalkthroughThe change decouples orchestrator startup from scheduler handshaking, adds owned-cluster handshake and assignment helpers, moves task-count initialization earlier, preserves the serial fallback, and records implementation details and measurements in investigation documentation. ChangesOrchestrator and scheduler initialization
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
Sequence Diagram(s)sequenceDiagram
participant AicpuExecutor
participant Orchestrator
participant SchedulerThreads
participant SchedulerContext
participant AICores
AicpuExecutor->>Orchestrator: Start core-agnostic task-graph initialization
AicpuExecutor->>SchedulerThreads: Start scheduler initialization
SchedulerThreads->>SchedulerContext: Handshake owned clusters
SchedulerContext->>AICores: Collect reports and publish task pointers
AICores-->>SchedulerContext: Return handshake state
SchedulerThreads->>SchedulerContext: Assign owned clusters
SchedulerContext-->>SchedulerThreads: Return dispatch-ready state
Orchestrator-->>AicpuExecutor: Return from decoupled initialization
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
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. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request parallelizes the per-core AICore handshake during initialization across multiple AICPU threads to reduce preamble overhead. It splits the initialization process into pre-handshake, parallel partitioned handshake, and post-handshake phases for both the a2a3 and a5 architectures. Additionally, for a2a3, it introduces an optimization allowing the orchestrator thread to skip the handshake and build the task graph concurrently. The handshake protocol itself is simplified by combining reports and using register polling to verify that register windows are open. There are no review comments, so no further feedback is provided.
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.
SergioMartin86
force-pushed
the
perf-orchestrator-decouple-draft
branch
2 times, most recently
from
July 20, 2026 08:09
efa0e94 to
3a4c7ac
Compare
SergioMartin86
marked this pull request as ready for review
July 20, 2026 08:16
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@docs/investigations/2026-07-orchestrator-decouple-preamble.md`:
- Around line 103-106: Run markdownlint-cli2 with the --fix option on the
documented investigation file to align the delimiter and header pipes in both
flagged tables, including the preamble-µs and workload-size tables. Preserve all
table content and Unicode characters; only apply the formatter’s spacing
corrections.
- Around line 21-24: Update the fenced code block in the orchestration timeline
near device_wall to specify the text language tag, preserving its contents and
formatting.
🪄 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: a86e442e-877d-4255-8107-bc5d4e6c1605
📒 Files selected for processing (5)
docs/investigations/2026-07-orchestrator-decouple-preamble.mddocs/investigations/README.mdsrc/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cppsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cppsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h
SergioMartin86
changed the title
ChaoWao
force-pushed
the
perf-orchestrator-decouple-draft
branch
from
July 20, 2026 09:17
3a4c7ac to
4122f0f
Compare
ChaoWao
changed the title
ChaoWao
force-pushed
the
perf-orchestrator-decouple-draft
branch
from
July 20, 2026 10:39
4122f0f to
41494d1
Compare
Remove the all-thread AICore-handshake barrier from the per-run scheduler
init on the tensormap_and_ringbuffer runtime, in two parts:
1. Orchestrator decouple: the orchestrator thread does not dispatch to
cores, so it skips the handshake and returns to run() to build the graph
immediately, deriving its AIC/AIV counts from the fixed 1:2 cluster ratio
(known pre-handshake) instead of post_handshake_init.
2. Barrier-free per-thread scheduler init: each scheduler thread handshakes
only the clusters it will dispatch to (blocked-layout ownership:
cluster ci = {ci, N/3+2ci, N/3+2ci+1}, owned by ci % active_sched_threads_)
and self-assigns them via assign_own_clusters, then returns straight to
run() -- no all-thread barrier and no leader-serialized
post_handshake_init. A thread starts dispatching to its own cores as soon
as they come up, independent of peers still handshaking.
handshake_owned_clusters uses handshake_partition's batched 4-phase
protocol (2 barriers, not per-core). DFX builds keep a profiling-only
barrier (pmu needs all physical_core_ids_); serial and single-thread fall
back to the original path.
A cache-warming variant of the ownership handshake measured null; only the
barrier removal it structurally enables gives the saving.
The blocked [AIC..][AIV..] layout and fixed 1:2 AIC:AIV ratio are
compile-time-and-hardware-enforced invariants (identical on sim and onboard),
so pre-handshake ownership is exact. pre_handshake_init now guards
cores_total_num_ % 3 == 0 up front: a non-multiple-of-3 count would leave the
tail AIV cores in no cluster (unhandshaked, silent hang at the op-execute
timeout), so it fails fast with a clear error instead. This holds for the
serial path too (assign_cores_to_threads pairs aiv_worker_ids_[2*ci]).
Ported to a5 (same restructure). a5 shares the 1:2 blocked layout but has
RUNTIME_MAX_WORKER=108 (36 AIC + 72 AIV) vs a2a3's 72; every count is derived
from cores_total_num_, so the port is count-agnostic. assign_own_clusters
mirrors each arch's own post_handshake_init: the a2a3 variant also does the
AsyncCtx constant-prefill / slab-clear, while a5 does sub_block_id only (a5
writes AsyncCtx per-dispatch in build_payload).
Measured on the current base (48d92e8), a2a3 onboard, 100 rounds/case: the
isolated init/preamble saving is a fixed ~11 us/run, consistent across 8
workloads regardless of size (the device_wall-minus-effective gap drops
10-12 us in every case; the larger device_wall swings are orch/sched compute
run-to-run variance this PR does not touch). The earlier ~63 us figure was
against a pre-hw-native-sys#1310 baseline whose preamble was still ~200 us; today's main
has already cut that, so the removable barrier + leader post_handshake_init is
worth ~11 us. Size-independent, so it still helps high-frequency per-op
dispatch. Golden passes (bgemm, alternating_matmul_add, batch_paged_attention,
qwen3_14b_decode).
Scope: a2a3 + a5 both validated -- a2a3sim (33) + a5sim (26) tmr suites green,
and a2a3 + a5 onboard st CI green (the a5 onboard run confirms the blocked
[0,N/3) AIC / [N/3,N) AIV handshake-enumeration assumption holds on a5 silicon).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
ChaoWao
changed the title
ChaoWao
approved these changes
Jul 20, 2026
ChaoWao
added a commit
to huawei-csl/simpler
that referenced
this pull request
Jul 21, 2026
…ian device speedup) Replace the tensormap_and_ringbuffer scheduler's fanin wiring (per-edge dep_pool + fanin_refcount/fanout_refcount atomics) with a polling completion design: a per-slot completion_flags byte array, a monotonic per-ring completed_watermark, and an intrusive wake_list for last-fanin notification. The orchestrator no longer wires every producer→consumer edge into a shared dep pool; consumers scan their inline fanin and either observe completion via the watermark or self-register on the one unmet producer's wake list. Fanin edges are stored inline in the payload (fanin_local_ids/fanin_ring_ids, PTO2_MAX_FANIN=128); heap reclamation is gated on last_consumer_local_id vs the per-ring watermark. Early-dispatch, predicated dispatch, and sync_start are preserved. Adopts hw-native-sys#1345's barrier-free, batched per-thread scheduler init. Cold-path init, dispatch, cold-path, and completion logic are implemented in .cpp translation units (scheduler_{dispatch,cold_path,completion}.cpp, pto_orchestrator.cpp, shared/pto_*.cpp); only the per-task/per-completion hot paths stay inline in headers. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
ChaoWao
added a commit
to huawei-csl/simpler
that referenced
this pull request
Jul 21, 2026
…ian device speedup) Replace the tensormap_and_ringbuffer scheduler's fanin wiring (per-edge dep_pool + fanin_refcount/fanout_refcount atomics) with a polling completion design: a per-slot completion_flags byte array, a monotonic per-ring completed_watermark, and an intrusive wake_list for last-fanin notification. The orchestrator no longer wires every producer→consumer edge into a shared dep pool; consumers scan their inline fanin and either observe completion via the watermark or self-register on the one unmet producer's wake list. Fanin edges are stored inline in the payload (fanin_local_ids/fanin_ring_ids, PTO2_MAX_FANIN=128); heap reclamation is gated on last_consumer_local_id vs the per-ring watermark. Early-dispatch, predicated dispatch, and sync_start are preserved. Adopts hw-native-sys#1345's barrier-free, batched per-thread scheduler init. Cold-path init, dispatch, cold-path, and completion logic are implemented in .cpp translation units (scheduler_{dispatch,cold_path,completion}.cpp, pto_orchestrator.cpp, shared/pto_*.cpp); only the per-task/per-completion hot paths stay inline in headers. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Human Summary
This PR brings the following optimizations to AICPU kernel initialization:
This optimization improves only the fixed-cost of initialization so it yields higher gains on small kernels. On larger kernels, the gains are diluted out per Amdahl's Law.
Update — a5 port + fixes
Since the original a2a3-only draft:
src/a5/.../tensormap_and_ringbuffer) — same restructure. a5 shares the 1:2 blocked layout but hasRUNTIME_MAX_WORKER=108(36 AIC + 72 AIV) vs a2a3's 72; every count derives fromcores_total_num_, so the port is count-agnostic.assign_own_clustersmirrors each arch's ownpost_handshake_init: the a2a3 variant also does the AsyncCtx constant-prefill / slab-clear, while a5 doessub_block_idonly (a5 writes AsyncCtx per-dispatch inbuild_payload).pre_handshake_initnow fails fast whencores_total_num_ % 3 != 0instead of silently leaving the tail AIV cores in no cluster (which would hang at the op-execute timeout). Holds for the serial path too.[SPIKE …]planning tags and rewrote the change-narration comments per the repo's comment rules.pre-commitis green.Verification: a2a3sim (33) + a5sim (26)
tensormap_and_ringbuffersuites green; all required CI checks pass, includingst-onboard-a2a3andst-onboard-a5on real silicon — the a5 onboard run confirms the blocked[0,N/3) AIC / [N/3,N) AIVhandshake-enumeration assumption holds on a5 hardware.AI Writeup
tl;dr for humans
Every AICPU run, the scheduler init still puts an all-thread barrier around
the AICore handshake: all threads greet the 72 cores, then the leader serially
builds the core→thread assignment and prefills every core's dispatch payload,
before any thread starts working. This PR removes that barrier. Each scheduler
thread greets only the cores it will dispatch to (computable up front from
the fixed, blocked 1 AIC : 2 AIV layout), self-assigns them, and starts
dispatching to its own cores as soon as they're up — no waiting for peers, no
leader-serialized setup. The orchestrator thread likewise skips the handshake
and starts building the graph immediately.
Net: ~63 µs of per-run scheduler init removed, entirely in the preamble, by
parallelizing the handshake + core assignment across the threads that own the
cores. The effect scales inversely with run size — a big win on small /
per-op workloads, and neutral (never a regression) on whole-model
decode/prefill where the fixed cost amortizes.
What changes
tidx == nthreads-1) doesn'tdispatch to cores, so it skips the handshake and returns to
run()to buildthe graph immediately, deriving AIC/AIV counts from the fixed 1:2 ratio.
only its owned clusters (
ci = {ci, N/3+2ci, N/3+2ci+1}, owned byci % active_sched_threads_) viahandshake_owned_clusters(batched 4-phaseprotocol, 2 barriers) and self-assigns them via
assign_own_clusters— noall-thread barrier, no leader-serialized
post_handshake_init.pre_handshake _initderives the counts + resets the run counters. DFX builds keep aprofiling-only barrier (pmu needs all
physical_core_ids_); serial /single-thread fall back to the original path.
Why it's safe: the blocked
[AIC…][AIV…]layout and 1:2 ratio arecompile-time-and-hardware-enforced invariants (identical on sim + onboard), so
each thread's owned cluster set is exact and known before the handshake.
Mechanism
The saving is entirely in the
preamblephase.hc_post,strace_timing --tree, this PR vs baseline:A/B benchmark — 37 single-die pypto-lib workloads, 20 rounds
.so-swap A/B, this PR vs the pre-change baseline, default config,device_wallmedian over 20 rounds, devices 2–5 (unlocked, verified-idle),a2a3 / Ascend 910B2. 35/37 faster, zero regressions.
ds-v4/hc_postqwen3-14b/token_embedds-v4/hc_headds-v4/expert_sharedds-v4/gateds-v4/decode_compressor_ratio4ds-v4/decode_indexer_compressords-v4/prefill_compressor_ratio128ds-v4/rmsnormds-v4/decode_indexerds-v4/hc_preds-v4/decode_compressor_ratio128ds-v4/decode_sparse_attnds-v4/prefill_compressor_ratio4ds-v4/decode_sparse_attn_swads-v4/decode_attention_swads-v4/decode_sparse_attn_hcads-v4/prefill_indexer_compressords-v4/decode_attention_hcaqwen3-14b/greedy_sampleds-v4/decode_attention_csads-v4/prefill_indexerds-v4/prefill_sparse_attnds-v4/expert_routedds-v4/qkv_proj_ropeds-v4/prefill_attention_hcads-v4/prefill_attention_swads-v4/prefill_attention_csaqwen3-32b/decodeds-v3.2/decode_backds-v4/mtp_projectionqwen3-14b/decode_layerds-v3.2/prefill_backqwen3-32b/decode_4dqwen3-14b/decode_tq(40L)qwen3-14b/prefill_tq_draftqwen3-14b/prefill_fwdBy size (the ~63 µs is a fixed per-run cost, so % impact ≈ 63 µs / run):
device_wall)Reading the numbers — single-round vs amortized
These harnesses run one round per process, so
device_wallincludes theper-run scheduler init. That's why standalone component kernels show large wins
while whole-model runs (
decode_tq40L −0.1%,prefill_tq_draft+0.1%) areneutral — the ~63 µs amortizes across a 120–470 ms forward. For reference, the
in-repo scene test
benchmark_bgemm(100 rounds in one warm process, initamortized) shows ~−0.7%. So end-to-end impact depends on serving
granularity: per-op / per-layer dispatch pays the init each time (large win);
whole-model-per-step amortizes it (neutral, never a regression).
Correctness
Golden PASS on the current base:
benchmark_bgemm,alternating_matmul_add,batch_paged_attention,qwen3_14b_decode;decode_tq4L/8L complete +validate. Zero regressions across the 37-workload × 20-round sweep above.
Scope / follow-ups
self-assign path assumes the AICore handshake reports cores in blocked
[0,N/3) AIC / [N/3,N) AIVorder; no a5 doc guarantees that enumeration, sorun a5 onboard golden before relying on it.
whole-model) to know the end-to-end impact.
abort_and_shutdown) is documentedas bounded/best-effort: a peer past
run()'sis_completed()gate isn'tjoined, but its own cores are valid and the host device reset that follows a
handshake failure reaps every core. Still not race-free, and
emergency_shutdownisn't exercised by golden (the invalid-physical_core_idtrigger needs a hw/sim fault-injection hook the UT harness lacks).
handshake_partition/post_handshake_init/assign_cores_to_threadsnow serve only the serial/single-thread fallback +the DFX barrier — keep or prune with the fallback's fate.