[Performance] A5 TMR: batch progress publication every 16 advances - #1575
Conversation
📝 WalkthroughWalkthroughThe PR adds an OCCUPY-based AICPU topology fallback, throttles ring scheduler updates to shared memory, and introduces A5-specific benchmark mappings with a single-example benchmark wrapper. ChangesAICPU topology fallback
Ring scheduler synchronization
A5 benchmark selection
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant Probe as probe_aicpu_topology_uncached
participant CPU_TOPO
participant OCCUPY as OCCUPY bitmap
Probe->>CPU_TOPO: query_cpu_topo
CPU_TOPO-->>Probe: entries or failure
Probe->>OCCUPY: intersect reported CPU IDs
OCCUPY-->>Probe: usable IDs or empty result
Probe->>OCCUPY: synthesize pool when needed
Possibly related PRs
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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 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/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h`:
- Line 475: Update advance_ring_pointers() and its sync_to_sm() ordering so
ring->fc.last_task_alive is stored with last_task_alive once PUBLISH_INTERVAL_K
slots have been consumed, avoiding the extra-call throttle; when reuse resets
last_published_to_sm to 0, publish immediately after the local lifetime advance.
- Around line 467-478: Update the scheduler’s terminal, idle, and drained-work
paths around advance_ring_pointers() to force-publish the current
last_task_alive watermark, including sub-interval advances of 1–15 steps.
Preserve batched sync_to_sm() behavior for ongoing work, and add coverage for 1,
15, 16, and 17 advances to verify final publication and threshold batching.
In `@tools/benchmark_a5_case1.sh`:
- Around line 17-21: Update the option handling case for -d|--device and
-n|--rounds in the argument parser to validate that a following value exists
before reading $2 or executing shift 2. When either option is missing its value,
emit the script’s usage error and terminate consistently; preserve the existing
assignments and shifts for valid values.
🪄 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: f3f19174-72e0-4e19-bda6-6ddebea50fd0
📒 Files selected for processing (4)
src/a5/platform/onboard/host/aicpu_topology_probe.cppsrc/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.htools/benchmark_a5_case1.shtools/benchmark_rounds.sh
511a165 to
526eb8b
Compare
yanghaoran29
changed the title
A5 performance re-testRe-tested
Test details:
On this host, both CPU topology query paths return 65534, so unmodified main and the current PR cannot pass the AICPU affinity gate. For measurement only, the same OCCUPY-only topology fallback was applied to temporary baseline and PR worktrees. It was therefore a shared control variable; the only A/B code difference was the K=16 scheduler publication change. The fallback was not added back to this PR. |
d9e9fde to
144c30e
Compare
A2/A3 port benchmark follow-upI ported PR #1575's A5 scheduler progress-publication batching to A2/A3 locally: publish Benchmark setup:
Effective time improved in 5 workloads and regressed in 3. Every change stayed within the +/-2% noise margin; the unweighted mean Effective change was -0.24%. On A2/A3 this establishes no clear performance benefit and no significant regression, unlike the A5 result recorded in this PR. Validation of the port: targeted |
yanghaoran29
changed the title
Complete A5 benchmark evidenceMeasured on the same
Unweighted mean delta across the eight workloads:
For Qwen, two complete paired groups are included. A third group completed only the For comparison, the portable A2/A3 implementation produced an unweighted mean Effective change of |
144c30e to
ea16900
Compare
ReviewThe optimization is well-motivated and the A5-vs-A2/A3 asymmetry is properly measured — but I think the correctness argument has a hole that turns some previously-working runs into a fatal. Details below. Stated vs. real goalThey match. The diff does exactly what the body says: adds a Why the blast radius is bigger than 19 lines of core logic
Each has the same two detectors: a structural test (is the reclaim head the oldest task pinned by an open scope?) and a ~500 ms wall-clock backstop that latches a fatal. The load-bearing fact is in Pre-PR, unconditional publication gave the waiters an invariant they silently depend on: the orchestrator sees every slot the scheduler has actually reclaimed. This PR replaces that with two compensating force-publish paths. I believe one of them doesn't work. Must fix1. The "idle" force-publish never fires for a withheld batch
bool drain_pending_ring_advances() {
uint32_t pending = advance_pending_mask.load(std::memory_order_acquire);
if (pending == 0) return false; // <-- early-out
So this claim in the body and the doc is not accurate for the idle half:
The surviving backstop is 2. A blocked orchestrator can prevent the full-drain backstop it depends onCombining #1 with the scope-pinning fact above:
Pre-PR, step 3 published immediately, Note what does not bound the exposure: Green Suggested fix — make batching headroom-aware. The scheduler already reads // Publish immediately once the orchestrator is within one batch of the
// window limit: below that headroom the watermark is its allocation credit,
// not a progress statistic.
bool near_full = current_task_index - last_published_to_sm >=
ring->task_window_size - PUBLISH_INTERVAL_K;
sync_to_sm(force_publish || near_full || last_task_alive == current_task_index);This keeps the full win in the regime you measured (deep ring, plenty of headroom — where all 8 benchmarks live) and restores the pre-PR guarantee exactly where liveness depends on it. Alternatives: have the orchestrator set a relaxed 3. Both structural deadlock detectors lose their exactness
Whichever fix lands for #2, these detectors need a head they can trust, and the comments asserting the old reasoning ( Should fix4. Heap reclamation is throttled by the same withheld watermark. 5. The 6. The benchmark set omits the regime this PR changes. All 8 workloads run the default 16384 window, where the ring never fills and back-pressure never engages — which is why they all improved. 7. No test covers the interaction. The three new UTs are the right cases for the arithmetic — boundaries 1/15/16/17, drained-tail forcing, and 8. Doc and body wording. Besides the idle-path claim in #1: "The published watermark is a conservative lower bound" holds for validity (TensorMap) but not for credit — for a blocked allocator a lower bound isn't conservative, it's a false negative on available space. Also, the divergence doc's baseline blockquote was changed from a maintenance baseline to "compares Consider9. 10. K=16 is unmotivated — no sweep reported. If the win is mostly cache-line traffic, K=4 or 8 may capture most of it at a quarter of the lag, which also shrinks the hazard surface. 11. 12. Things I checked that are fineNoting these so they don't get re-litigated:
VerdictRequest changes — on #1/#2/#3, which are one fix. The performance work is sound and the measurements support it for the regime measured. The problem is that the stated safety argument names two backstops and one of them doesn't fire, leaving full-ring drain as the only one — which a blocked orchestrator can itself prevent, since scope-end only runs on the orchestrator. All eight benchmarks and green onboard CI are consistent with the bug being latent rather than absent: every measured workload uses the 16384 default window and never reaches back-pressure. Happy to be shown wrong on the reachability of #2 — if there's a path that force-publishes on a blocked allocator that I've missed, that resolves #1 and #2 together. |
ea16900 to
0b796df
Compare
|
@ChaoWao Thanks for the detailed review. I addressed the selected items as follows:
I did not add #10's K sweep in this correctness fix: K=16 is the already measured configuration, and changing it would require a new paired performance study. I also left #11's optional enum out; the existing Validation is clean: independent C++ suite 100/100, A5 simulator ringbuffer 258/258 checks, both locked A5 small-window runs above, and local pre-commit including clang-format, clang-tidy, cpplint, and markdownlint. |
Final implementation A5 performance re-testThis re-test supersedes the earlier performance summary for the current correctness implementation. Setup:
Effective improved in 7 of 8 workloads, with an unweighted mean change of -2.93%. Five improvements exceed the +/-2% noise band; manual Case2 and Qwen are effectively flat. The exception is |
0b796df to
9605f5e
Compare
|
Final split-mask Batch Paged Attention Case1 recheck on A5 device 0 (
Four-group means:
Per-group Effective changes: The earlier Correctness/quality evidence on the final code:
|
50f4e22 to
9847d12
Compare
Re-review (force-push
|
| # | Prior finding | Status |
|---|---|---|
| M1 | "Idle" force-publish never fires for a withheld batch | ✅ Fixed — dedicated publication_request_mask / publication_ack_mask, serviced by thread 0 from both productive and idle iterations. The old advance_pending_mask keeps its idle-only try-lock-retry semantics, correctly kept separate. |
| M2 | Blocked orchestrator can prevent the full-drain backstop | ✅ Fixed — 10 ms no-progress escape wired into all four waiters: alloc (slots and heap), PTO2DepListPool::ensure_space, PTO2FaninPool::ensure_space, ensure_tensormap_capacity. |
| M3 | Structural detectors lose exactness | ✅ Fixed, and better than I asked for. Gating on a fresh ack is the right move: an ack can only enable the check when the forced publish was a no-op (a productive ack raises the watermark → the last_alive > prev_last_alive branch fires → watermark_synchronized is cleared → check skipped). So the check runs precisely when published == scheduler-local head, which restores the pre-PR aliasing impossibility, since admission bounds live ids to [T, T+W) around an exact head. |
| S4 | Heap throttled by the same watermark | ✅ Fixed — same waiter, update_heap_tail refreshed on ack, dedicated HeapPressurePublishesWithheldProgress test. |
| S5 | Scope-size guard threshold now wrong | ✅ Resolved as a consequence — the escape restores the full effective window, so window_size() - 1 is the correct threshold again. |
| S6 | Benchmark set omits the back-pressure regime | 🟡 Partly — see should-fix 2. |
| S7 | No test covers the interaction | ✅ Fixed — four pressure tests spawn a real blocked allocator thread, plus independent-mask and arena-relocation tests. |
| S8 | Doc/body claims inaccurate | ✅ Fixed — the false idle-path claim is gone, the baseline blockquote is a date-verified maintenance baseline again, and the investigation doc records the first design's flaw. One new defect: see should-fix 1. |
| C9 | K constant / hardcoded "15" |
✅ Fixed — static constexpr PUBLISH_INTERVAL_K at struct scope; comment derives from it. |
| C12 | Cite the prior batched-publish investigation | ✅ Done, with an honest write-up of why the first design was incomplete. |
Verification I ran
- Configured and built
tests/ut/cpp;ctest -LE requires_hardware: 100/100 passed. - The four new pressure tests are coupled to two real wall-clock constants, so I stress-ran them:
--repeat until-fail:15 -j 4overtest_a5_{wiring,fanin_pool,orchestrator_fanin,task_allocator}→ 60/60 passed, no flakes. - I suspected a defect in the two pool waiters: their ack branch assigns the refreshed watermark straight into
prev_last_alive, which hides that advance from thecur > prevbranch that re-anchorsblock_cycle0— wherealloc()avoids this by using a separate local. I wrote a scratch A/B harness (ack-driven vs. unprompted publication, permanently starved pool, 2 s budget) to see whether a progressing waiter could reach the 500 ms backstop. It can't — any subsequent advance re-anchors, so only the single final advance before a genuine stall is affected, and declaring a deadlock there is correct. Not a defect; the asymmetry withalloc()is cosmetic.
Should fix
1. docs/MULTI_RING.md is now misclassified in the divergence doc — by this commit
The PR adds 2 lines to src/a5/.../docs/MULTI_RING.md, so the a2a3/a5 pair is no longer byte-identical, but docs/tensormap-and-ringbuffer-a2a3-vs-a5.md still lists it under "Byte-Identical Files". Measured, not derived:
$ diff -rqs src/a2a3/runtime/tensormap_and_ringbuffer \
src/a5/runtime/tensormap_and_ringbuffer | grep -c 'are identical'
24
$ cmp -s src/a2a3/runtime/tensormap_and_ringbuffer/docs/MULTI_RING.md \
src/a5/runtime/tensormap_and_ringbuffer/docs/MULTI_RING.md # → DIFFERSThe doc claims 25. The whole off-by-one is this one file: 24 identical + 29 differing = 53 matching pairs, which only reconciles with the doc's 25 + 11 + 17 once MULTI_RING.md moves. docs/RUNTIME_LOGIC.md was correctly moved into the functional list for the same reason, so this looks like a simple miss.
Fix: byte-identical 25 → 24, drop MULTI_RING.md from that list, add docs/MULTI_RING.md to the functional list (17 → 18 matching, 19 → 20 total). The doc's own definition of functional covers files that "document differences in runtime behavior", and its new maintenance note — "Recompute the counts and update the affected sections whenever the files or constants described here change" — is exactly the instruction that got missed.
2. Only one workload is measured on the final SHA, and the PR added a shared cache line to thread 0's productive dispatch path
The seven-row table is explicitly from the preceding SHA (thanks for labelling that honestly), which leaves Batch Case1 as the sole final-SHA measurement. That matters more than usual here, because drain_publication_requests() now runs on every thread-0 dispatch iteration:
if (thread_idx == 0 && made_progress && sched_->drain_publication_requests()) {and it opens with publication_request_mask.load(acquire) — a line the orchestrator writes. It should stay clean in thread 0's cache while no request is outstanding, and −3.83% on Case1 suggests it is. But a PR whose thesis is "remove a Scheduler↔Orchestrator cache line from the hot path" shouldn't leave a newly-added Scheduler↔Orchestrator line in the hot path validated on one workload. Please re-run the seven on the final SHA.
Relatedly: the 10 ms criterion is well justified for Case1 — 36 promotions → 0 is a genuinely good piece of profiling — but the configs where promotions should fire are the small-window ones, and those appear only as correctness passes. A promotion count plus an Effective number for paged_attention_ringbuffer at ring_task_window=64 would close this properly. Credit where due: the locked paged_attention_manual_scope at window 64 run is exactly the manual-scope × small-window conjunction I said CI never covers.
3. The wiring is fail-silent in the unsafe direction
If reclaim_request_mask_ / reclaim_ack_mask_ are ever left null, enabled() returns false, so watermark_synchronized is initialized true and never cleared. That simultaneously (a) disables the liveness escape and (b) lets the structural check run on a lagging watermark — i.e. it silently restores both bugs from the last round — while K-batching stays enabled. The fail-safe direction is inverted.
I checked the current wiring and it is correct: runtime_init_data_from_layout (which nulls via init) precedes runtime_wire_arena_pointers → set_scheduler on the host path; the AICPU path runs wire_arena_pointers then reset_for_reuse, and PTO2OrchestratorState::reset_for_reuse re-wires immediately after task_allocator.init, while both pools' reset_for_reuse preserve the pointers. So this is hardening, not a live bug. But the safe shape is for "escape not wired" to mean "don't batch" (or to latch a fatal at init), rather than "batch without an escape".
Consider
4. Compare task ids, not slot pointers, in the structural check. head_is_oldest_open_task compares oldest_open_task == &slot_states_[head_task_id & window_mask_]. The ack gate makes this exact at the instant of the ack, but the check runs up to 1024 spins later and the scheduler may have advanced locally by δ without publishing in that window; for δ ≥ 2, task P + W is live and aliases the published head's slot. It additionally needs the oldest open-scope task to sit exactly there, so it's a very narrow coincidence — but it's on a fatal path, and comparing the local task id closes it unconditionally.
5. ReclaimPublicationRequest computes 1u << ring_id directly, bypassing PTO2SchedulerState::ring_advance_pending_bit() and its static_assert(PTO2_MAX_RING_DEPTH <= 32). Same bit derivation in two places, only one guarded.
6. ensure_tensormap_capacity hand-rolls the handshake with a lambda over all 32 ring bits instead of reusing ReclaimPublicationRequest, and never consumes the ack. That is correct — it has no structural head check, so it only needs the watermark to move — but it reads as an oversight next to the other three call sites. A one-line comment ("no ack needed: no structural head check here") would settle it.
7. Document why a single servicer suffices. Thread 0 services from both loop branches, so the common case is covered. But thread 0 is not in the loop while inside a handle_drain_mode spin, which delays servicing; that resolves on its own (drain waits on cores freeing, independent of the orchestrator) and the 500 ms backstop leaves ~490 ms of margin. Worth stating at the call site, since "why thread 0 only" is the first question a reader will have.
8. Genuine open-scope deadlock classification is now ≥10 ms slower (was ~1024 spins). Fine for a fatal path, but docs/troubleshooting/device-error-codes.md and the running-onboard.md triage table describe the structural-vs-timeout distinction, so "A5 takes ≥10 ms to reach the structural verdict" is now part of that story.
9. The four pressure UTs depend on two real wall-clock constants. 15× clean at -j 4 here, and the 490 ms margin plus the latched-request design make them robust. Still, making PTO2_PUBLICATION_REQUEST_TIMEOUT_CYCLES injectable would decouple them from runner load and cut ~40 ms of hot spinning from the suite.
Verdict
Approve, conditional on st-onboard-a5 going green on 50f4e22c (CI was still mid-flight when I looked — pre-commit pending), with should-fix 1 requested before merge: it's a two-line doc correction and the count is verifiably wrong.
Should-fix 2 and 3 are worth doing but neither affects the correctness of what's here. The design is now the right shape: batching on the non-blocking path, an explicit receiver-visible escape when a consumer is actually starved, separate masks so scheduler lock contention cannot forge an acknowledgment, and structural classification gated on a watermark that is provably exact. The investigation-doc entry generalizing the lesson — a batched publication needs a receiver-visible escape whenever the receiver depends on it for forward progress — is the most valuable artifact in the diff.
9847d12 to
59a79f0
Compare
|
Implemented the requested follow-up and rebased the single PR commit onto current
Validation:
The final-SHA A5 hardware correctness/performance rerun is still blocked locally by the mandatory architecture gate: inside a single-device |
4863559 to
f66f976
Compare
Batch last_task_alive publication every 16 local advances on the non-blocking path. Keep scheduler try-lock retries on their original idle-only deferred mask, and isolate orchestrator publication requests and acknowledgments on a separate cache line. Promote task, heap, dependency, fanin, and TensorMap pressure to the productive-loop handshake only after 10 ms without reclaim progress. Service those requests from scheduler thread 0, require a fresh acknowledgment before structural deadlock classification, and preserve the 500 ms wall-clock backstop. Enable batching only after allocator, fanin, and dependency-pool request/ack wiring is validated against the current scheduler. Default initialization, arena relocation, and reuse to per-advance publication until validation succeeds. Compare exact ring-local task identities during structural classification, share the guarded ring-bit derivation, document the liveness and diagnostic contracts, and cover batching fallback, relocation, reuse, pressure, and slot-alias cases.
|
I pushed the two mechanical fixes from my review directly ( 1. Dropped three files unrelated to this change Restored from the merge-base CI's Correction to my review: I claimed dropping these would return the a2a3 jobs to $ git diff a59ffde7...HEAD --name-only \
| grep -vE '^(src/a5/|examples/a5/|tests/(st|ut/cpp)/a5/)' \
| grep -vE '<NON_CODE>'
tests/ut/cpp/common/test_scope_deadlock_detection.cppThat shared fixture had to change for the exact-identity head check, it belongs to no arch partition, and it is compiled against both runtimes — so it correctly flips 2. Rewrote the PR body It was two rebases stale (still citing
I did not invent numbers for the current head. Re-running the Case1 table on I deliberately did not rebase onto the newer Two notes on what I chose not to touch: I left the title alone (the commit subject didn't change), and I added no My approval from the previous comment stands, now with both should-fixes closed. |
2 participants
Summary
ring->fc.last_task_alivepublication every 16 Scheduler-local advances on the non-blocking path.advance_lockcontention on the originaladvance_pending_mask, drained only from no-progress iterations.publication_request_maskonly when an Orchestrator reclaim consumer has made no progress for 10 ms; Scheduler thread 0 then force-publishes and acknowledges from productive or no-progress iterations.Scope
This PR changes only:
tensormap_and_ringbufferScheduler progress publication and Orchestrator reclaim back-pressure;test_scope_deadlock_detectionfixture that the exact-identity head check requires; andIt does not change host-build-graph, A2/A3 runtime behavior, benchmark tooling, PTO-ISA checkout, or CI workflows.
Why A5 only
-0.24%; all results within approximately +/-2%-3.83%Any Scheduler may win
advance_lockand publish the shared watermark that the Orchestrator repeatedly reads. K=16 does not remove Scheduler-to-Scheduler contention; it reduces Scheduler-to-Orchestrator publication frequency by up to 16x on A5's distributed placement.The implementation is not unsafe or unsupported on A2/A3. A local A2/A3 port passed its tests, but batching added reclamation lag without a measurable payoff, so A2/A3 remains unchanged. See the A2/A3 benchmark follow-up.
Correctness
last_task_aliveis not a progress statistic: it is the allocation credit for task slots, heap bytes, dependency entries, fanin spill entries, and TensorMap entries. A blocked reclaim consumer can itself prevent an open scope from ending, so a withheld batch must never depend on full-ring drain to be published.last_task_aliveremains authoritative; the published watermark is a lower bound.PTO2TaskAllocator::alloc),PTO2DepListPool::ensure_space,PTO2FaninPool::ensure_space, andensure_tensormap_capacity.publication_batching_enableddefaults false at initialization, arena relocation, reuse, and teardown.runtime_wire_arena_pointersandruntime_reset_for_reuseenable it only after verifying every ring's allocator, fanin-pool, and dependency-pool request/ack pointers against the live Scheduler. Incomplete wiring degrades to pre-change per-advance publication rather than batching without a liveness escape.Validation
Rebased on
mainata59ffde75c2907819a45690eee065056857330dc.Measured on the current head:
Measured on the immediately preceding head, which differs from the current one only by removing three files unrelated to this change:
st-onboard-a5,st-sim-a5sim,ut-a5,ut-a2a3, andpackaging-matrix.--repeat until-fail:12 -j 4acrosswiring,fanin_pool,orchestrator_fanin,task_allocator,shared_memory,scheduler_state,scope_deadlock_detection): 96/96 passed, no flakes.Measured on earlier heads of this branch:
paged_attention_ringbuffer: 258/258 checks passed.paged_attention_ringbufferwithring_task_window=64,ring_heap=4 MiB,ring_dep_pool=256: passed.paged_attention_manual_scopewithPTO2_RING_TASK_WINDOW=64: 4/4 cases passed — the small-window plus manual-scope combination that exercises the pressure path.A5 performance
Batch Paged Attention Case1 on
Ascend950PRdevice 0. Four groups, each running baseline immediately followed by this branch, 100 measured rounds per side. Negative deltas mean this branch is faster.The four paired Effective changes were
-4.92%,-2.29%,-4.20%, and-3.91%. Profiling showed that an earlier implementation promoted 36 short-lived allocation-pressure episodes per Case1 run into exact publication handshakes; the 10 ms no-progress criterion promoted zero in this workload while preserving the liveness path for a genuinely blocked Orchestrator.Directional context from a seven-workload sweep:
Measurement provenance. Every number above was collected before the wiring-validation gate and the exact-identity head check landed; the seven-workload rows predate the Case1 table. They are not a single-SHA aggregate. Both later additions are off the steady-state path — a
boolread on a cache line the hot path already writes, and one extra indirection inside the 1024-spin cold block — so no steady-state change is expected, but that expectation is unmeasured. A re-run of the Case1 table on the current head would close it. See the full earlier evidence.