[Performance] host_build_graph: activate graph nodes incrementally as they materialize - #1804
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 graph scheduler now publishes materialized nodes incrementally. Atomic counters track publication and routing progress. Materialization registers or routes nodes as they become available, while graph activation routes published roots after the external dependency gate opens. ChangesIncremental graph activation
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant prepare_graph_task
participant graph_execution_materialize_slice
participant graph_incremental_publish
participant graph_route_ready_roots
participant GraphExecution
prepare_graph_task->>graph_execution_materialize_slice: materialize graph slice
graph_execution_materialize_slice->>GraphExecution: publish execution.nodes
prepare_graph_task->>graph_incremental_publish: publish new node range
graph_incremental_publish->>GraphExecution: register or route nodes
graph_route_ready_roots->>GraphExecution: claim published roots with route_cursor
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 |
![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/runtime/scheduler/graph_execution.cpp`:
- Around line 365-371: In both
src/a2a3/runtime/host_build_graph/runtime/scheduler/graph_execution.cpp:365-371
and
src/a5/runtime/host_build_graph/runtime/scheduler/graph_execution.cpp:365-371,
update the materialization flow so execution.nodes is assigned only before the
first range is published; remove subsequent slice assignments and the final
reassignment before PREPARED. Preserve byte-for-byte parity between both files.
🪄 Autofix
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: 264c2df8-6b8d-4833-ac36-a91a2c09e61b
📒 Files selected for processing (6)
src/a2a3/runtime/host_build_graph/runtime/graph_execution.hsrc/a2a3/runtime/host_build_graph/runtime/scheduler/graph_execution.cppsrc/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.hsrc/a5/runtime/host_build_graph/runtime/graph_execution.hsrc/a5/runtime/host_build_graph/runtime/scheduler/graph_execution.cppsrc/a5/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h
72f7d25 to
9643511
Compare
Review: incremental graph activationVerdict: Approve. Clean, well-scoped performance change; the two safety-critical assumptions both hold at the code level, verified below. One non-blocking test suggestion. What it doesInstead of fully materializing a The new hazard, and why it's safeDispatching a node before the whole graph is materialized opens a window the
The "producer completes before consumer registers" case is then handled by Should-fix (non-blocking): a deterministic test for the new interleavingThe four graph_execution scene tests + qwen3 golden mostly exercise the linear path; the new "producer completes mid-expansion" interleaving is only hit probabilistically (the 45× repeat raises the odds but doesn't force it). A deterministic end-to-end repro is genuinely hard (device timing isn't controllable without a runtime test hook), so I wouldn't ask for that. But a host-side cpput unit test is easy here and targets exactly the dangerous window — the fixtures in // producer already completed and its wake list already drained:
producerA.task_state.store(PTO2_TASK_COMPLETED);
producerA.wake_list_head.store(WAKE_LIST_SENTINEL);
// consumer registers only now — the "completes before register" order:
sched.register_graph_wake(exec, &producerA, &consumerC);
// assert C is routed to the ready queue, not lost.A companion case over Minor notes
|
9643511 to
062cd8b
Compare
Thanks for the thorough review — and especially for independently verifying both load-bearing assumptions (the bind_graph_topology topological-order check at graph_execution.cpp:222 and the single-owner publish via the prepare-queue). That's exactly the safety argument the change rests on. I've added the deterministic host-side tests you suggested (amended into the commit). Two TEST_Fs that force the dangerous "producer completes before its consumer registers" interleaving without any hardware or timing dependence:
One correction on placement: tests/ut/cpp/a2a3/test_scheduler_state.cpp compiles against the tensormap_and_ringbuffer PTO2SchedulerState, which has none of the graph-execution surface (GraphExecution / register_graph_wake / graph_incremental_publish). Those live only in host_build_graph, so the tests are in a new HBG test target, tests/ut/cpp/{a2a3,a5}/test_graph_activation.cpp (registered via add_a2a3_hbg_runtime_test / add_a5_hbg_runtime_test, so they run in ut-a2a3 / ut-a5). Mirrored byte-for-byte across a5, clang-format/clang-tidy clean, both pass on a2a3 and a5. |
|
This is blocking merge (the check won't go green), and it's worth pinning down before another re-run: the failure reproduces on the host_build_graph graph path this PR changes.
The failure (attempts 2 & 3, identical): The 1-2 tests failing right after it on the same xdist worker are just collateral ( Why I think it points at this change rather than a flake:
Could you check whether this is introduced by the incremental activation or already present on |
The on-device Graph scheduler expanded a GRAPH task into all its nodes before dispatching any of them: graph_execution_materialize_slice reached PREPARED only after the last node, and activate_prepared_graph routed the roots only at that point. For a large graph the AICores idle through the whole expansion — on the qwen3-14B 3-layer decode that front-loaded ramp is ~530 us per layer during which every core waits. Route roots (and any node whose producers have all completed) as soon as they materialize, once the outer Graph task's external-dependency gate has opened, so compute overlaps the remaining expansion. This is possible only because of the polling-completion scheduler: the per-slot completion_flags byte + task_state completion mirror introduced for tensormap_and_ringbuffer in hw-native-sys#1137 and adapted to host_build_graph in hw-native-sys#1435. That flag lets a consumer discover a producer that completed mid-expansion by reading its state instead of relying solely on the producer's wake list. Under the older pure wake-list model, dispatching before full materialization would have lost wakes if a producer finished before its consumer registered — which is exactly why the whole-graph PREPARED gate existed. - Materialization-time registration moves from register_initial_graph_waiter (a raw wake-list store that assumed no producer could complete during expansion) to register_graph_wake, which resolves a producer that completed mid-expansion through its task_state completion flag (the hw-native-sys#1137 / hw-native-sys#1435 mechanism) instead of losing the wake. The graph's topological node order guarantees every producer a node references is already constructed. - execution.nodes is published at MATERIALIZING (under materialize_busy) so the wake path can read producer slots while the graph is still expanding. - Roots reach the ready queue through route_cursor, a monotonic per-execution cursor that makes routing idempotent across the per-slice calls and the final call at the activation meet, so each root is pushed exactly once. Non-roots reach the queue only through their producers' wake list, unchanged. Because a node now dispatches before the graph reaches ACTIVE, it can complete while the graph is still MATERIALIZING or PREPARED. hw-native-sys#1767's complete_task rejects a node completion unless the graph is ACTIVE; that guard now accepts MATERIALIZING and PREPARED as well, rejecting only SUBMITTED (execution not yet localized) and COMPLETED (already retired). Without this, the largest graphs (longest materialization tail) fault with sched_error_code=5 INVALID_ARGS while the small graph_execution scene tests, which reach ACTIVE before any node completes, pass. Correctness: the three graph_execution scene tests (2D replay, AIC+AIV, MIX-SPMD) and the qwen3-14B 3-layer graph-execution golden pass, and the qwen3-14B 40-layer host-graph decode (the largest graph) passes 6/6 where it faulted 0/6 without the completion-state fix. Three host-side unit tests (tests/ut/cpp/{a2a3,a5}/test_graph_activation.cpp) cover the interleavings deterministically: a consumer registered after its producer completed and drained routes through the task_state re-scan; graph_incremental_publish routes all-producers-complete consumers and wake-chains the rest; and complete_task accepts a node completion in MATERIALIZING/PREPARED/ACTIVE and rejects it in SUBMITTED/COMPLETED. Applies the identical change to the a5 host_build_graph sibling (same graph scheduler; a5 is compile-verified, not perf-measured on this a2a3 box).
062cd8b to
2bd0c9b
Compare
3 participants
Human Summary
In HBG the task graph is sent to the device and fully expanded before computation starts. This PR pipelines the expansion and the computation so that AI cores can start doing useful work as soon as the first task is processed. This saves a bunch of initialization time.
AI Summary
The on-device Graph scheduler expanded a
GRAPHtask into all its nodes before dispatching any of them —graph_execution_materialize_slicereachedPREPAREDonly after the last node, andactivate_prepared_graphrouted the roots only at that point. For a large graph the AICores idle through the entire expansion. On the qwen3-14B 3-layer decode this front-loaded ramp is ~530 µs per layer with every core waiting.This routes roots (and any node whose producers have all completed) as soon as they materialize, once the outer
GRAPHtask's external-dependency gate has opened, so compute overlaps the remaining expansion.How it's made safe — and what makes it possible
Dispatching a node before the whole graph is materialized means a producer can complete while a later consumer is still being registered — the exact hazard the
PREPAREDgate existed to prevent. This change is only possible because of the polling-completion scheduler: the per-slotcompletion_flagsbyte +task_statecompletion mirror introduced fortensormap_and_ringbufferin #1137 and adapted tohost_build_graphin #1435. That model lets a consumer answer "has this producer finished?" by reading a flag, so a producer completing mid-expansion is discovered, not missed. Under the older pure wake-list model this would have been unsafe — which is precisely why the whole-graphPREPAREDgate existed. Three things close the hazard:register_initial_graph_waiter(a raw wake-list store that assumed no producer could complete during expansion) toregister_graph_wake, which resolves a producer that completed mid-expansion through itstask_statecompletion flag (the [Optimization] Replace wiring with polling-based task readiness test (~17% median device speedup) #1137/Add: polling completion scheduler for host_build_graph (a2a3) #1435 mechanism) instead of losing the wake. This machinery already existed on the live wake path; this change simply routes materialization through it.execution.nodesis published atMATERIALIZING(undermaterialize_busy) so the wake path can read producer slots during expansion.route_cursor, a monotonic per-execution cursor that makes routing idempotent across the per-slice calls and the final call at the activation meet — each root is pushed exactly once. Non-roots reach the queue only through their producers' wake list, unchanged.Testing
Correctness: the three
graph_executionscene tests (2D replay, AIC+AIV, MIX-SPMD) and the qwen3-14B 3-layer graph-execution golden pass, including a 3-round replay. A 45-run repeat of the three graph-exec tests passes 45/45 (initial rapid-repeat failures were host-sidehalMemCtlMMIO contention from process churn — a different binary, before the scheduler runs — and vanish once process starts are spaced).Performance (qwen3-14B 3-layer, a2a3 onboard, same-device A/B, 15 rounds):
The chip swimlane confirms compute begins overlapping graph expansion (cores busy at ~400 µs) instead of waiting for it (~600 µs baseline). The first-compute start is still bounded by the outer task's external-dependency latency, so the recovered time is the expansion tail, not an earlier start.
Scope
host_build_graph, same commit (identical graph scheduler). a5 is compile-verified only — the dev box is a2a3 silicon.