perf(a2a3/tmr): shrink dispatch cold-write cost on the sync_start drain#1328
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThe dispatch protocol replaces the ChangesDispatch protocol and execution flow
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant SchedulerContext
participant PTO2DispatchPayload
participant aicore_execute
SchedulerContext->>PTO2DispatchPayload: Populate args or set src_payload
SchedulerContext->>PTO2DispatchPayload: Publish dispatch
aicore_execute->>PTO2DispatchPayload: Check src_payload
aicore_execute->>PTO2DispatchPayload: Materialize gated args
aicore_execute->>PTO2DispatchPayload: Continue doorbell and acknowledgement flow
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 optimizes the task dispatch path by offloading kernel argument materialization to the idle AICore on the gated path, restructuring the PTO2DispatchPayload cache layout, and prefetching destination structures. Feedback focuses on adding bounds checks: first, validating tensor_count and scalar_count in aicore_executor.cpp to prevent out-of-bounds writes to args; second, validating that claim does not exceed the size of the claimed array in scheduler_completion.cpp.
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.
…eady path Records the experiment behind hw-native-sys#1328's asymmetric design (AICPU fills args for ready tasks, offloads only gated/early-dispatch ones): making the AICore fill its own args[] for every task adds ~1.0 us to each task's receive->start setup (paged_attention_unroll: 349 ns -> 1356 ns), because a ready task has no idle doorbell gate to hide the fill. Offload only pays off on the not_ready path, where the AICore is already spinning at the gate. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp (1)
1014-1019: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winStale init comment contradicts the new clearing model.
This block still states
build_payloadsetsnot_readyand thatdeferred_slab->count/error_code are reset inline on every dispatch. After this PR,not_readyis gone (folded intosrc_payload) and the slab is cleared once here plus re-cleared only in the completion path. The deinit comment (Lines 1078-1086) was updated; this one was missed.🤖 Prompt for 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. In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp` around lines 1014 - 1019, Update the stale initialization comment near the scheduler dispatch setup to reflect the current clearing model: remove references to build_payload setting not_ready and resetting deferred_slab count/error_code on every dispatch, and document that the slab is cleared once during initialization and re-cleared in the completion path while src_payload carries the readiness state.
🤖 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.
Outside diff comments:
In
`@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp`:
- Around line 1014-1019: Update the stale initialization comment near the
scheduler dispatch setup to reflect the current clearing model: remove
references to build_payload setting not_ready and resetting deferred_slab
count/error_code on every dispatch, and document that the slab is cleared once
during initialization and re-cleared in the completion path while src_payload
carries the readiness state.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 8464271f-de35-4a5d-8f15-e17e954ad8cc
📒 Files selected for processing (8)
docs/investigations/2026-07-aicore-fills-all-args-ready-path.mddocs/investigations/README.mdsrc/a2a3/runtime/tensormap_and_ringbuffer/aicore/aicore_executor.cppsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto2_dispatch_payload.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cppsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cppsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp
The sync_start drain launches every block of the barriered task serially on one elected scheduler thread, so its per-subtask cost is dominated by writes to cold, per-core structures. Cut that cost by removing writes from the dispatch path, folding the remaining ones onto one cache line, and offloading what can be offloaded: - PTO2DispatchPayload re-layout (still 512B): all per-dispatch control writes now land on cache line 0 (function_bin_addr + local_context + src_payload). The not_ready gate is folded into src_payload (0 = ready, non-zero = gated AND the source PTO2TaskPayload address — a payload pointer is never 0, so the flag is lossless), and the cold global_context moves to the tail. The explicit ABI pad is dropped (alignas(64) rounds to 512). - Init-time prefill: async_ctx's slab pointers + capacity and the two context-pointer args (args[48]/[49]) are per-(core,buf_idx) constants — set once in init() instead of every dispatch, so build_payload writes only CL0. - deferred_slab reset moved off the dispatch path: cleared once at init, then re-cleared on the completion path only after a task actually recorded a deferred completion (count > 0), which is rare. This removes a large, sparse (~222 KB, TLB-heavy) cold write from every dispatch — the biggest single win. - AICore-side arg materialization on the gated path: when a task is gated the AICPU writes only src_payload and the idle AICore fills args[0..num_args) from the source payload during its doorbell wait. Cross-toolchain offsets into PTO2TaskPayload are pinned with static_asserts. - Software-pipelined prefetch of the next block's cold per-core structures during the drain (read-prefetch CoreExecState first, then the dispatch buffer; write-prefetch measured to help on this shallow-store-buffer core). Measured (instrumented) on pypto-lib qwen3_14b 2-layer decode: sync_start drain burst ~38.9 us -> ~12 us/drain. The same build_payload path is shared with normal dispatch, so it benefits too. Correctness: tmr onboard ST suite passes (22 tests incl all spmd_sync_start*, SPMD multiblock, paged attention); decode_fwd --validate-fwd sample-match 16/16 with logits within tolerance. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…eady path Records the experiment behind hw-native-sys#1328's asymmetric design (AICPU fills args for ready tasks, offloads only gated/early-dispatch ones): making the AICore fill its own args[] for every task adds ~1.0 us to each task's receive->start setup (paged_attention_unroll: 349 ns -> 1356 ns), because a ready task has no idle doorbell gate to hide the fill. Offload only pays off on the not_ready path, where the AICore is already spinning at the gate. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
ChaoWao
force-pushed
the
perf/tmr-dispatch-drain-cold-writes
branch
from
bd908c8 to
d9f1d97
Compare
1 participant
What
The
sync_startdrain (drain_worker_dispatch) launches every block of the barriered task serially on one elected scheduler thread, so its per-subtask cost is dominated by writes to cold, per-core structures. This shrinks that cost by taking writes off the dispatch path, folding the rest onto one cache line, and offloading what can be offloaded.Changes
PTO2DispatchPayloadre-layout (still 512 B). All per-dispatch control writes now land on cache line 0 (function_bin_addr+local_context+src_payload). Thenot_readygate is folded intosrc_payload(0= ready, non-zero = gated and the sourcePTO2TaskPayloadaddress — a payload pointer is never 0, so the flag is lossless). Coldglobal_contextmoves to the tail; the explicit ABI pad is dropped (alignas(64)rounds to 512).async_ctx's slab pointers + capacity and the two context-pointer args (args[48]/[49]) are per-(core,buf_idx)constants — set once ininit()instead of every dispatch, sobuild_payloadwrites only CL0.deferred_slabreset moved off the dispatch path. Cleared once at init, then re-cleared on the completion path only after a task actually recorded a deferred completion (count > 0, rare). This removes a large, sparse (~222 KB, TLB-heavy) cold write from every dispatch — the biggest single win.src_payload; the idle AICore fillsargs[0..num_args)from the source payload during its doorbell wait. Cross-toolchain offsets intoPTO2TaskPayloadare pinned withstatic_asserts.CoreExecStatefirst, then the dispatch buffer).Measurement
Instrumented, on pypto-lib qwen3_14b 2-layer decode:
sync_startdrain burst ~38.9 µs → ~12 µs/drain. Thebuild_payloadpath is shared with normal dispatch, so it benefits too. (The AICore arg-fill only helps the gated/early-dispatch path, ~8% of dispatches here.)Correctness
spmd_sync_start*(drain), SPMD multiblock, paged attention, basics.decode_fwd --validate-fwdsample-match 16/16, logits within tolerance, across repeated runs (deferred-completion path exercised).🤖 Generated with Claude Code