[Performance] HBG: 80~95% host-side overhead reduction - #1659
[Performance] HBG: 80~95% host-side overhead reduction#1659SergioMartin86 wants to merge 1 commit into
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 Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThe host runtime now uses uniquely owned shared-memory storage, initializes only control regions, and uploads the live payload prefix plus complete control segments to the device. ChangesShared-memory staging
Estimated code review effort: 2 (Simple) | ~10 minutes 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: 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/host_build_graph/host/runtime_maker.cpp`:
- Around line 537-548: Validate total_tasks before computing payload_prefix_end
or performing relocation/copy operations: require it to be non-negative and no
greater than eff_task_window_sizes[0]. Reject invalid values early, preserving
the existing copy behavior only for valid task counts.
🪄 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: e9f8a9e0-e5d7-420b-a483-e1085fb0681e
📒 Files selected for processing (1)
src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp
run_host_orchestration allocated and zeroed a full ring-sized host mirror of shared memory (host_sm_buf(sm_size, 0)) and H2D-uploaded it whole on every run. sm_size scales with the ring task window (81 MB default, 651 MB at a 4 GB ring), not the workload, so nearly all of the alloc+zero and the upload was spent on slots a run never touches. This host bind path is 96-99.6% of the per-dispatch wall. The SM payload segment (~97% of sm_size) is written per task at submit and read only for [0, total_tasks); its tail is never read on device. Leave that segment allocated-but-untouched (no page faults), zero only the small control segments, and upload only the live prefix plus the whole (small) slot_states and completion_flags segments. Every scheduler-read field is already explicitly initialized per slot, so nothing depends on the removed blanket zero; reads are bounded by current_task_index. Single-ring (PTO2_MAX_RING_DEPTH == 1). bgemm 90.9 -> 16.0 ms (-82%), paged_attention 536 -> 32.5 ms (-94%) total run; device_wall unchanged; golden passes. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MqeALZTPEnDXTnbYfcnPfq
87874ee to
71923a6
Compare
1 participant
Human Summary
The host side of HBG was taking 20x more time than the device side. This was due to an unnecessary resetting (zeroing) and copying of the entire scheduling workspace. This PR reduces the H2D transfer to the minimal required information, and the work structures are initialized-on-use, rather than pre-zeroed.
AI Summary
PR: host_build_graph — init-on-write SM mirror (skip zero+upload of the unused tail)
Branch
hbg-sm-init-on-writeoffupstream/main(810fbcd1). One file, +25/−3.Problem
Profiling HBG on a2a3 (upstream/main, exclusive card, device time via
[STRACE])showed the per-dispatch wall is 96–99.6% host
bind; the on-device scheduler(
device_wall) is ~0.1%.bindruns in full on every launch (no amortization).Splitting
bindlocated the cost inrun_host_orchestration: it allocates+zeroes afull, ring-sized host mirror of the shared-memory region (
std::vector<uint8_t> host_sm_buf(sm_size, 0)), the orchestrator fills it, then it is H2D-uploaded whole.sm_sizescales with the ring task window, not the workload — 81 MB (default ring) to651 MB (4 GB ring). Measured floor:
The alloc+zero is page-fault-bound (fresh anonymous mmap, ~1.9 GB/s); the upload is the
full mirror. A run only uses a few dozen of the 16k–131k slots, so almost all of this is
work on capacity that is never touched.
Fix (init-on-write)
The SM is cleanly segmented (
ring_segment_offsets): header → descriptors → payloads→ slot_states → completion_flags. Payloads are ~97% of
sm_size. Each task fullyinitializes its own descriptor/payload/slot_state at submit, and the scheduler reads no
slot past
total_tasks, so:costs no page faults). Still zero the small control segments (header+descriptors,
slot_states+completion_flags), preserving every existing zero-assumption there.
[0, total_tasks)payload prefix, one for the whole (small) slot_states+completion segments. The payload
tail is never read on device, so it is not shipped.
Cost then tracks the run (tasks submitted), not the ring capacity.
Why it is safe
A completeness check enumerated every SM field the scheduler reads vs every per-slot init
site (
reset_for_reuse,PTO2TaskPayload::init,prepare_task, the control-segmentmemsets /
init_per_ring). Result: every scheduler-read byte is explicitly initializedat submit; nothing depends on the blanket zero. Reads are bounded by
current_task_index(=
total_tasks), so unsubmitted slots are never touched. The only bulk-zero-only bytesare unread (padding, array tails, slots past the high-water mark).
Single-ring (
PTO2_MAX_RING_DEPTH == 1), so no ring-loop; the offset math uses theexisting single-source-of-truth
ring_segment_offsets.Results (a2a3, exclusive card, medians)
The two SM-mirror ops collapse: bgemm 58.5 → 1.7 ms (−97%), paged 446 → 10.1 ms (−98%).
Device side is unchanged. Full HBG a2a3 scene-test suite golden: (filled after run).
Scope / follow-ups
host_build_graphonly. Orthogonal to Add streamed Graph Execution to host_build_graph #1444 (graph execution): that shrinks whatgoes into the SM; this fixes the SM buffer's alloc/zero/upload lifecycle. They compound.
slot_states) to
total_taskstoo — for the big ring these are still window-sized(~7 ms). The arena build/upload (~6 ms) is the next floor.
target/elem_sizein theop==NONEbranch (currently zero only via the old blanket zero,but gated unread by
has_predicate).