perf(runtime): cut TMARB preamble ~72us -> ~9us via report-first + batched handshake#1323
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 A2A3 and A5 runtimes now use a report-first AICPU–AICore handshake. Worker completion flags are reset before launch, register-window setup is batched after AICore reports, compact core metadata supports post-handshake initialization, and dispatch slots no longer receive per-boot clearing. ChangesAICPU–AICore handshake protocol
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant DeviceRunner
participant AICoreExecutor
participant SchedulerContext
DeviceRunner->>AICoreExecutor: launch with cleared aicore_done
AICoreExecutor->>SchedulerContext: publish core identity and aicore_done
SchedulerContext->>SchedulerContext: collect and batch reported cores
SchedulerContext->>AICoreExecutor: publish task pointer and open register window
AICoreExecutor->>SchedulerContext: poll DATA_MAIN_BASE and fetch dispatch payload
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 implements a "report-first" handshake protocol between the AICore and AICPU across the a2a3 and a5 platforms. Instead of waiting for an AICPU initialization signal, the AICore now immediately reports its physical core ID and core type on launch. The host clears aicore_done before kernel launch to ensure freshness. Additionally, the scheduler's handshake partitioning and core discovery are optimized through batched phases, prefetching, and the use of a compact core-type array to reduce HBM write stalls and eliminate unnecessary zeroing. As there are no review comments provided, I have no additional feedback to offer.
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.
ChaoWao
force-pushed
the
perf/preamble-reportfirst-batched-handshake
branch
from
33ad605 to
b157389
Compare
Reworks the AICPU init / AICore handshake, applied symmetrically to both
arches. Measured on a2a3 onboard (spmd_paged_attention_highperf s8192,
per-step preamble from the device_wall STRACE markers); a5 rides CI.
1. Drop the per-round 256KB post_handshake_init memset (~55us -> 0).
It zeroed payload_per_core_ + deferred_slab_per_core_ every boot, but
prepare_subtask_to_core rewrites every field a slot exposes before the
AICore is told to read it, so a prior round's bytes are structurally
unreachable.
2. Report-first AICore handshake (pypto-style), removing the AICPU->AICore
round-trip (hs_ready2first ~5us -> ~0.3us). The AICore publishes
{physical_core_id, core_type, aicore_done} on launch without waiting for an
aicpu_ready gate; launched before the AICPU, its report is already in GM
when the AICPU sweeps. Window-open replaces aicpu_ready as the sync point.
- task is written per core AFTER aicore_done is observed (else the core's
CACHELINE_OUT report flush clobbers it); the core dcci-reads it after
window-open.
- Freshness: the host zeros each worker's aicore_done before
launch_aicore_kernel (workers persist across runs in the pooled arena),
so the AICPU never opens a window on a stale physical_core_id.
3. Batch the AICPU's per-core servicing (the handshake sweep, ~5us -> ~1.6us).
Report-first means all cores are reported when the AICPU sweeps
(n_passes==1), so instead of servicing each core inline (interleaved GM
store-through + posted MMIO STR + barrier, each stalling the next) the sweep
runs in phases: collect + prefetch CoreExecState lines, publish all task
pointers under ONE barrier, open all windows back-to-back (pipelined nGnRE
STRs), then publish each CoreExecState with a single struct store.
4. Batch the worker-id list build in post_handshake_init (~3.7us -> ~0.3us).
core_type is read from a compact contiguous array the sweep fills; the
aic/aiv lists are built in local buffers and published with two wide memcpys
(direct 4-byte stores into the GM members cost ~50ns each, store-through).
Residual ~9us: the front-matter (pre_handshake_init, barrier wait,
post_handshake_init tail) plus the ~2 MMIO opens/core that cannot move off the
AICPU (AICore cannot self-open its SPR window). Getting below this needs
persistent workers, out of scope here.
Validated: a2a3 onboard 20 rounds + 6 examples (paged_attention_highperf,
spmd_multiblock_mix, multi_round_paged_attention, spmd_basic, dummy_task,
alternating_matmul_add) green; a5 sim (spmd_multiblock_mix, dummy_task) green.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 participant
Summary
Reworks the AICPU init / AICore handshake to cut the TMARB preamble from ~72µs to ~9µs/step (a2a3 onboard,
spmd_paged_attention_highperfs8192, from thedevice_wallSTRACE markers). Applied symmetrically to a2a3 and a5.Drop the per-round 256KB
post_handshake_initmemset (~55µs → 0). It zeroedpayload_per_core_+deferred_slab_per_core_every boot, butprepare_subtask_to_corerewrites every field a slot exposes before the AICore is told to read it, so a prior round's bytes are structurally unreachable.Report-first AICore handshake (pypto-style), removing the AICPU→AICore round-trip (
hs_ready2first~5µs → ~0.3µs). The AICore publishes{physical_core_id, core_type, aicore_done}on launch without waiting for anaicpu_readygate; launched before the AICPU, its report is already in GM when the AICPU sweeps. Window-open replacesaicpu_readyas the sync point.taskis written per core afteraicore_doneis observed (else the core'sCACHELINE_OUTreport flush clobbers it); the core dcci-reads it after window-open.aicore_donebeforelaunch_aicore_kernel(workers persist across runs in the pooled arena), so the AICPU never opens a window on a stalephysical_core_id.Batch the AICPU's per-core servicing in the handshake sweep (~5µs → ~1.6µs). Report-first means all cores are reported when the AICPU sweeps (single pass), so instead of servicing each core inline (interleaved store-through GM write + posted MMIO STR + barrier, each stalling the next) the sweep runs in phases: collect + prefetch
CoreExecStatelines → publish alltaskpointers under one barrier → open all windows back-to-back (pipelined nGnRE STRs) → publish eachCoreExecStatewith a single struct store.Batch the worker-id list build in
post_handshake_init(~3.7µs → ~0.3µs).core_typeis read from a compact contiguous array the sweep fills; the aic/aiv lists are built in local buffers and published with two widememcpys (direct 4-byte stores into the GM members cost ~50ns each, store-through to HBM).Residual ~9µs is the front-matter (
pre_handshake_init, barrier wait,post_handshake_inittail) plus the ~2 MMIO opens/core that cannot move off the AICPU (the AICore cannot self-open its SPR window). Getting below this needs persistent workers, out of scope here.Testing
spmd_paged_attention_highperfs8192 + 6 examples (spmd_paged_attention_highperf,spmd_multiblock_mix,multi_round_paged_attention,spmd_basic,dummy_task,alternating_matmul_add) passspmd_multiblock_mix,dummy_taskpass