[Performance] HBG: up to 99% host-side overhead reduction

[SergioMartin86]

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.

• Up to 99% (depending on the case) of the host-side costs are reduced.
• This optimization is orthogonal to that of Add Graph Execution to host_build_graph #1444, both will help independently in ameliorating the host-side overhead.

AI Summary

PR: host_build_graph — init-on-write bounded SM upload + skip orch block

Branch hbg-sm-init-on-write, rebased onto upstream/main (0659745d, includes #1444
graph execution). One commit, +324/−38 across 12 files (core + a guardrail, a regression
test, and doc/comment updates from review).

TL;DR

HBG's per-dispatch wall is 96–99.6% host bind, and bind was dominated by
rebuilding and H2D-uploading full, ring-sized host structures every run — the
shared-memory mirror and the ~20 MB prebuilt runtime arena — even though a run touches
a tiny fraction of either. The device boots scheduler-only and reads no SM slot past
total_tasks, so the SM mirror is made init-on-write and shipped bounded to the task
count, and the arena's ~8.5 MB host-only orchestrator block is dropped from the
upload.

Result — host bind, A/B vs upstream (a2a3):

workload	baseline	this PR	Δ
bgemm	77.0 ms	~11.6 ms	−85%
paged_attention	402.7 ms	~7.8 ms	−98%

(bgemm/paged re-measured post-revert; matmul/vector were not re-measured but sit in the same
~8–12 ms small-kernel range — the SM-mirror bind floor once the payload/window uploads are
bounded.) Device time is unchanged (~0.05–0.1 ms) — the whole win drops into per-dispatch
latency. Note: an earlier revision also bounded the ready-queue upload, taking bgemm bind to
~4 ms; that was reverted for #1444 graph-execution compatibility (see Fix), so the queue
portion of the win is not in this PR.

Problem

Splitting bind (via [STRACE] markers) put nearly all of it in run_host_orchestration /
bind_callable_to_runtime_impl, in two structures that are both sized by ring capacity and
rebuilt+reuploaded every run:

• The shared-memory mirror (host_sm_buf(sm_size, 0)): 81 MB (default ring) to 651 MB
  (4 GB ring), of which ~97% is the payload segment.
• The prebuilt runtime arena (~20 MB, fixed): scheduler ready queues (65536×7) +
  tensormap (65536 entries).

A run uses a few dozen of the 16k–131k slots, so almost all of the alloc/zero/init and the
uploads is work on capacity that is never touched.

Fix

The device reads no slot past total_tasks. So:

Shared memory — descriptors and payloads are written per task at submit; slot_states and
completion_flags are reset per slot in orch::prepare_task as it is claimed (dropping the
window-wide reset loop in init_header_per_ring); only the header is zeroed on the host;
each segment is H2D-uploaded bounded to [0, total_tasks).

Runtime arena —

• Don't upload the orchestrator block (fanin_seen_epoch / scope / tensormap, ~8.5 MB): it is
  host-only dep-computation scratch the AICPU scheduler never reads. Everything from the
  scheduler block onward (ready queues, runtime header, mailbox) still ships whole.
• Drop the redundant per-entry stores in the tensormap reset (the preceding memset already
  zeroes the link pointers and producer_task_id).

The ready queues are not bounded to total_tasks. An earlier revision seeded/shipped them
bounded (with a +1 sentinel for the batched-dequeue boundary), but that is incompatible with
#1444 graph execution: a replayed GRAPH task is expanded by the device Scheduler into on-device
nodes that push into the ready queues past the host task count, so every queue slot must carry
a valid Vyukov sequence on the device. The queues therefore ship in full. (This cost the queue
portion of the arena win but leaves the dominant SM win intact — see below.)

total_tasks is range-checked before it sizes the SM copies. Single-ring (PTO2_MAX_RING_DEPTH == 1).

Why it is safe

Every scheduler-read SM field is explicitly initialized at submit, so nothing depends on the
removed blanket zeros; reads are bounded by current_task_index. The orchestrator block is
host-only (verified: zero AICPU-scheduler references); the tensormap-reset change is a pure
dedup of what memset already wrote. The ready queues ship in full, so graph execution's
on-device node expansion pushes into fully-seeded slots.

Validation

100-round golden × 4 workloads PASS (bgemm, matmul, vector, paged_attention). The full HBG
a2a3 scene-test suite is green — including the graph_execution cases (28 passed, 0 failed
per pass) and dep_gen — under pytest-xdist cross-run stress with 0 scheduler stalls / 0
op-execute timeouts, and ut-cpp is 83/83. (The graph_execution compatibility is load-bearing:
the queue-bounding revision failed those, which is why the ready queues now ship in full.)

Guardrails and tests (from review)

• Named error, not a silent hang. push_ready_routed checks the push return and
  latches PTO2_ERROR_READY_QUEUE_OVERFLOW on failure — a genuinely full queue (task
  window or graph-node expansion past the 65536 capacity) would otherwise drop a ready
  task and stall. Zero cost on the success path.
• Layout-order assert. bind_callable_to_runtime_impl always_asserts
  orch_start <= orch_end before slicing the host-only orchestrator block out of the
  upload, so a future runtime_reserve_layout reorder faults instead of shipping a
  misaligned image.
• Poison test. tests/ut/cpp/a2a3/test_hbg_submit_poison.cpp fills the SM window with
  0xAA, submits a representative mix (real mixed task with tensors + scalar, multi-fanin
  consumer, hidden-alloc, dummy), and asserts every device-read slot field carries a real
  value, not poison — so a future device-read field added without a submit-path write fails
  in-tree instead of non-deterministically on device.
• Docs/comments. The reset_for_reuse relocation (init → prepare_task) is
  reflected everywhere it was described, and the H2D contract is documented in
  RUNTIME_LOGIC.md §3.1.

Scope / not done

• a2a3 host_build_graph only.
• The ready-queue upload is not bounded (see Fix) — incompatible with Add Graph Execution to host_build_graph #1444 graph
  execution's on-device node expansion. Recovering that ~6.7 ms would need a device-side
  scheme that seeds the slots a graph run actually reaches; left as a follow-up.
• Fully bounding the tensormap reset (a further ~1–2 ms) is deliberately left out:
  print_stats scans the whole entry pool and dereferences entries, and the pool has a
  recycling free-list, so bounding its reset is high-risk for a sub-noise gain. The 8 MB
  memset stays.

[coderabbitai]

📝 Walkthrough

Walkthrough

The 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.

Changes

Shared-memory staging

Layer / File(s)	Summary
Selective initialization and device upload src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp	The runtime uses std::unique_ptr storage, selectively clears control regions, and splits device uploads between the live payload prefix and complete control segments.

Estimated code review effort: 2 (Simple) | ~10 minutes

Possibly related PRs

• hw-native-sys/simpler#1439: Changes shared-memory staging and metadata handling in the same runtime file.
• hw-native-sys/simpler#1618: Changes shared-memory layout and control-segment staging.
• hw-native-sys/simpler#1619: Changes completion-flag layout and shared-memory initialization semantics.

Poem

A rabbit hops through memory bright,
Leaves payload bytes untouched in flight.
Control flags clear, the slots align,
Two careful transfers cross the line.
“Efficient staging!” thumps my cheer.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly identifies the HBG performance optimization and its primary benefit, matching the changeset.
Description check	✅ Passed	The description directly explains the host-side overhead reduction, bounded uploads, initialization changes, validation, and scope.

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

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

📥 Commits

Reviewing files that changed from the base of the PR and between 71433ca and 87874ee.

📒 Files selected for processing (1)

• src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp

[ChaoZheng109]

The mechanism holds. The load-bearing invariant — device-side reads are bounded
by current_task_index everywhere — checks out at every site: classify_partition
iterates [0, submitted), update_completed_watermark walks bounded by
fc.current_task_index, fanin lookups go through fanin_local_ids[] (always
< total_tasks), and relocate_host_orch_image walks [0, count). The two claims
the PR body asserts without proof also hold:

• The skipped orch block is genuinely host-only. rt->orchestrator and
  rt->scheduler are embedded by value in PTO2Runtime at off_runtime, inside
  the "rest" region that still ships whole, so scheduler_cold_path.cpp:1082
  reading rt->orchestrator.inline_completed_tasks is safe. The scheduler/AICPU
  sources contain zero dereferences of tensor_map / scope_tasks /
  fanin_seen_epoch.
• Dropping active_mask = ActiveMask{} from init_header_per_ring is safe —
  prepare_task:442 rewrites it unconditionally. Between reset_for_reuse() and
  the explicit stores in prepare_task, every functional field of
  PTO2TaskSlotState is covered; only the _async_pad filler is left
  uninitialized, which is inert.

The +1 sentinel writeup in the PR description is excellent — "a lock-free queue's
read set extends one element past its write set" is exactly the kind of finding
that deserves to be recorded where the next person will hit it.

None of the items below is a correctness bug in the current code. They are
guardrails around the new invariants, plus the doc/comment drift the change leaves
behind.

---

1. The new "≤ total_tasks pushes per ready queue" invariant has no guardrail, and violating it fails silently

The whole bounding scheme rests on each task being pushed into a big ready queue at
most once. That holds across every push site today (push_ready_routed's three
branches, fed by the wake drain, register_wake re-classification, boot classify,
the cold path, and cancel_early_sync_drain). The problem is what happens when it
stops holding:

• push_tagged reads the target slot's Vyukov sequence. On an unseeded /
  un-uploaded slot that value is stale garbage, and diff < 0 makes it
  return false — which reads as "queue full" but here means "slot was never
  shipped". push_ready_routed ignores the return value, so a ready task is
  silently dropped and the run hangs.
• If the stale sequence lands on diff > 0 instead, the push spins forever.

Both surface as SCHEDULER_TIMEOUT / host-side 507018 — precisely the failure
this PR already had to debug once. This is not only about unseeded slots:
PTO2_READY_QUEUE_SIZE is 65536 while the task window can be larger
(paged_attention runs 131072), so a genuinely full queue is reachable too and is
discarded just as silently.

Minimum ask: check the push() return value in push_ready_routed and
report_fatal on false. That converts a silent hang into a named error at zero
layout cost. Optionally, record the seeded prefix length in PTO2ReadyQueue (there
is spare room in _pad0) and fault on enqueue_pos >= seeded_slots, which also
catches the spin case.

2. No regression test, though the harness is already in place

tests/ut/cpp/CMakeLists.txt already has add_a2a3_hbg_runtime_test() (currently
driving test_hbg_task_allocator and test_hbg_tensormap), and
tests/ut/cpp/a2a3/test_ready_queue.cpp already covers the Vyukov queue. A
test_hbg_ready_queue.cpp pinning this PR's core semantics is cheap:

• seed with slot_init_count = N+1, write a stale (too-large) sequence into slots
  [N+1, capacity), push N, then assert pop_batch returns N and
  terminates;
• the slot_init_count = N control case (the first cut described in the PR body)
  should be shown to hang, under a timeout assertion.

Per .claude/rules/discipline.md §3, the intermittent hang this PR fixed currently
has only a pytest-xdist stress run as its barrier, and nothing in-tree.

3. Doc / comment drift — reset_for_reuse() moved, four descriptions did not

.claude/rules/doc-consistency.md §4 wants these in the same commit:

Location	Current text	Status
docs/RUNTIME_LOGIC.md §8.4	"reset_for_reuse() runs once at init (pto_shared_memory.cpp)"	stale
runtime/pto_runtime2_types.h (reset_for_reuse doc block)	"Runs once per slot at init (pto_shared_memory.cpp)"	stale
runtime/shared/pto_runtime2_init.cpp:95-97	"Per-slot SM-side initialization … lives in init_header_per_ring so the AICPU performs it during SM reset"	stale
runtime/pto_shared_memory.h (completion_flags field)	"Zeroed host-side at init."	stale

Please also add the new H2D contract to RUNTIME_LOGIC.md: which segments are
shipped bounded to total_tasks, why the orch block is not shipped at all, and why
the ready-queue prefix is total_tasks + 1. That contract currently lives only in a
comment in runtime_maker.cpp, and it is the most easily broken thing in the
change.

4. Host-side style: prefer modern C++ over the C array idiom

.claude/rules/codestyle.md §8 — runtime_maker.cpp:933-955 uses
const size_t big_q_slots[] with sizeof(a) / sizeof(a[0]) in both loops.
std::array plus a range-for reads better and removes the two hand-written
divisions.

Two comments are change-narration rather than present-tense facts
(.claude/rules/comments.md): pto_orchestrator.cpp:406 ("rather than pre-zeroing
the whole task window in init_header_per_ring") and pto_shared_memory.cpp:184
("not swept over the whole task window here"). Both would read better stated as the
invariant that holds now; the before/after belongs in the commit message.

5. The arena slicing silently depends on the reservation order

orch_start = layout.orch.off_fanin_seen_epoch and
rest_start = sq.off_early_dispatch_queue_slots[0] assume the order in
runtime_reserve_layout (sm_handle → orch → sched(big queues → small queues) →
runtime → mailbox). That order is correct today, but any future reordering ships a
wrong image with no diagnostic. A few always_asserts would pin it:

• orch_start < sq.off_ready_queue_slots[0]
• the seven big-queue offsets are strictly ascending
• rest_start >= last_big_queue_offset + capacity * sizeof(PTO2ReadyQueueSlot)

6. Nothing enforces the new "every device-read field is written at submit" contract

The SM buffer is no longer zero-filled, so the fields nobody writes —
payload.tensors[tensor_count..], fanin_local_ids[fanin_count..],
_fanin_reserved, slot_state._async_pad — read as zero by allocator accident,
not by construction. The exposure is a future one: add a field to
PTO2TaskDescriptor / PTO2TaskPayload / PTO2TaskSlotState, miss the
submit-path write, and it passes every test here and fails non-deterministically
on device.

Please pin the contract with a cpput test that owns its own SM buffer: fill it with
0xAA before init_per_ring, submit a representative mix (multi-fanin, zero-fanin
root, predicated, hidden-alloc via alloc_tensors, dummy), then assert that for
every slot in [0, total_tasks) no device-read field still holds poison —
descriptor task_id / kernel_id[] / packed_buffer_base / packed_buffer_end;
slot_state task_state / active_mask / task_attrs / total_required_subtasks /
logical_block_num / last_consumer_local_id / payload / task / wake-list;
payload tensor_count / scalar_count / fanin_count /
fanin_local_ids[0..fanin_count) / predicate.op / the early-dispatch atomics;
and completion_flags[slot].

test_orchestrator_fanin.cpp already has this shape against the tmr tree;
add_a2a3_hbg_runtime_test() needs pto_orchestrator.cpp and shared/*.cpp added
to link an HBG equivalent.

Nice work — the measurement, the A/B table and the root-cause writeup are all well
above the bar. The requests above are about making the new invariants enforceable by
the tree rather than by the reader.

[SergioMartin86]

Thanks @ChaoZheng109 for the thorough pass — especially for independently walking the current_task_index bound at every read site and confirming the orch-block and active_mask claims.

Heads-up on a significant change since the review. All six items were implemented as requested, but rebasing onto current main (which now includes #1444 graph execution) surfaced a deeper problem: bounding the ready-queue upload is fundamentally incompatible with graph execution. A replayed GRAPH task is expanded by the device Scheduler into on-device nodes that push into the ready queues past the host total_tasks, so any prefix-bounded queue seed (sentinel or not) leaves those slots un-seeded and the batched dequeue spins — exactly the SCHEDULER_TIMEOUT this whole thread was about, now triggered by graph nodes instead of cross-run reuse. It reproduced deterministically (clean main passes the graph_execution scene tests; the bounded-queue branch stalled them).

So I reverted the ready-queue bounding entirely. The queues now ship in full; the PR keeps the dominant, graph-safe part of the win — the init-on-write bounded SM upload plus dropping the host-only orchestrator block from the arena upload. That takes bgemm bind from 77 ms to ~11.6 ms and paged_attention from 403 ms to ~7.8 ms (−85% / −98%); the extra ~6.7 ms the queue bounding bought is gone. This reframes a few of your items — details inline.

1. Guardrail on queue pushes

Done, and now more relevant than ever. push_ready_routed checks the push() return and latches a named PTO2_ERROR_READY_QUEUE_OVERFLOW on failure. With the queues shipping full this can only fire on a genuinely full queue — a task window or a graph-node expansion exceeding the 65536 capacity — which is precisely the graph-execution edge worth failing loudly on rather than dropping silently. Zero cost on the success path.

2. Regression test for the sentinel — withdrawn

The sentinel (and its test_hbg_ready_queue) existed only to make the bounded queue safe. Since the queue bounding is reverted, both are gone — there's no longer a prefix boundary to pin. The graph-execution scene tests are now the regression barrier for the queue path, and they pass.

3. Doc / comment drift

Done. The three stale reset_for_reuse-location comments are re-pointed to orch::prepare_task, and the H2D contract is documented in RUNTIME_LOGIC.md §3.1 — updated to state that the SM ships bounded, the orch block is dropped, and the ready queues ship in full (with the graph-execution reason). Re: the §8.4 pointer you quoted — I still couldn't find that text; §8 here is "Scalar Access During Construction". Point me at it if it's in a different file and I'll fix it.

4. Host-side style

The std::array / range-for you flagged was the queue-offset array, which is gone with the bounding. The two change-narration comments on the SM path (pto_orchestrator.cpp / pto_shared_memory.cpp) are reworded to present-tense invariants and stay.

5. Arena slicing order asserts

Simplified along with the revert: the only slicing left is dropping the orchestrator block, so bind_callable_to_runtime_impl now always_asserts orch_start <= orch_end. The multi-queue ordering asserts went with the bounded upload.

6. SM-poison test

Done — tests/ut/cpp/a2a3/test_hbg_submit_poison.cpp, via a new orchestrator-linked target. It fills the SM window with 0xAA, submits a representative mix (real AIV task with an output tensor + scalar, multi-fanin consumer, hidden-alloc, dummy), and asserts every device-read field is written, not poison. It earned its keep: it flagged that the pre-completed hidden-alloc task legitimately leaves its dispatch-time predicate.op as poison (it's set_completion_flag'd at submit and never dispatched), so the test gates dispatch-field checks on PENDING — which doubles as documentation of that exception. Everything genuinely device-read is confirmed written, so removing the zero-fill is safe.

Validation: full HBG a2a3 scene-test suite green including graph_execution (28 passed / 0 failed per pass) and dep_gen, under xdist cross-run stress with 0 scheduler stalls; ut-cpp 83/83. Thanks again for the review — the queue-bounding line of work didn't survive the graph-execution rebase, but the review directly drove the guardrail and the poison test that make the remaining change trustworthy.