Fix A2A3 repeated-run AICore stream regression - #1807
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (13)
📝 WalkthroughWalkthroughThe PR makes AICore stream reuse image-aware. Pipeline slots retain completed streams for the same image and recreate them after image transitions. Native preparation now provisions AICPU streams separately and passes callable image hashes through stream acquisition. Tests and documentation reflect the updated lifecycle. ChangesImage-aware AICore stream reuse
Estimated code review effort: 4 (Complex) | ~45 minutes 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 |
ChaoWao
left a comment
ChaoWao
left a comment
There was a problem hiding this comment.
Reviewed at a15632e3 against merge-base 50c06606. The state machine is clean and the UT coverage around it is genuinely good, so this is not a rewrite request — but I think the reuse key is the wrong one, and picking the right one makes the patch both smaller and safer.
The invalidation trigger has a single chokepoint
I audited every RT_MEMCPY_HOST_TO_DEVICE on the onboard host path:
| site | writes |
|---|---|
device_runner_base.cpp:692 (upload_chip_callable_buffer) |
AICore child code — the only one |
host_regs.cpp:210 |
MMIO register address table |
device_runner_helpers.cpp:76 / :105 |
runtime args / kernel args |
device_runner_base.cpp:146 (copy_to_device) |
tensor payload |
ensure_binaries_loaded |
AICPU SO dispatcher bootstrap, not AICore code |
So the instruction bytes in GM change in exactly one place, and only when the content-hash dedup misses (a dedup hit returns before the rtMemcpy, so no code is written).
That means the question a slot needs to answer is "has any code been published since this stream was created?" — not "which image did this slot run last time?" Those are different predicates, and keying on the second one is wrong in both directions.
Under-invalidates: A -> B -> A across two slots
bound_image is per-slot, but the instruction cache belongs to the cores, which both pipeline slots dispatch to. upload_chip_callable_buffer allocates via mem_alloc_.alloc → rtMalloc, and unregister_callable → release_chip_callable_buffer → rtFree releases the block, so an address is recyclable:
- register A → run on slot 0 →
S0retained,bound_image = H_A,complete - unregister A →
rtFree(X) - register B →
rtMallocreturnsX→ B's children occupy A's old code range - run B on slot 1 → slot 1 creates its own stream → cores fetch B at
X - unregister B; register A again →
rtMallocreturnsX; A's bytes are back atX - run A on slot 0 →
bound_image == H_A→ stream reused, no fresh-stream flush
Slot 0 never observed step 4 because it happened on the other slot. Under the exact hardware premise #1540 was built on, step 6 can execute B's cached lines. This is the case #1791 warned about in Additional Context:
A plain restoration of the old per-slot image-hash cache may reintroduce stale-instruction risk for an
A -> B -> Asequence across depth-two slots.
Over-invalidates: two resident images alternating
When both callables stay registered, they occupy two distinct rtMalloc blocks and no code byte changes between runs — yet every transition destroys and recreates. The PR's own test pins that cost:
# test_alternating_code_images_never_execute_stale_instructions
assert st_worker.run_stream_set_create_count == before + 1 # every runSo an alternating-resident-image workload pays the full #1540 penalty and recovers nothing. Worth checking against real serving traffic, where prefill and decode are separate callables on one worker — the dsv4 HCA/SWA benchmark in #1791 is a single callable in a loop, so it is the one shape this keying does fix.
Suggested shape
One bool per slot, using the vocabulary already in the tree (stale instructions is the wording in this file's own test name and in docs/task-flow.md:293):
// RunStreamSlots::Slot
bool stale{false};
void mark_all_stale() {
for (Slot &s : slots_) {
std::lock_guard<std::mutex> lock(s.mutex);
s.stale = true;
}
}
int acquire(unsigned slot) {
...
if (s.aicore != nullptr) {
if (!s.complete) return -1;
if (!s.stale) return 0; // no publication since creation: reuse
rc = destroy_(s.aicore); if (rc != 0) return rc;
s.aicore = nullptr;
}
rc = create_(&s.aicore); ...
s.stale = false;
}Hook it after the successful rtMemcpy in upload_chip_callable_buffer (not on the dedup-hit path). run_stream_slots_ lives on the a2a3 DeviceRunner while upload lives on the base, so a default-no-op virtual on the base that a2a3 overrides — the shape abandon_native_run_resources already uses at device_runner_base.h:524.
Two notes on why this needs no extra machinery:
- In-flight runs need no coordination. A resident callable holds a refcount and the allocator only hands out free blocks, so a publication can never overwrite the code of an executing image. A newly-set
staleflag only affects the nextacquire. - No counter or per-slot snapshot needed. Publication is rare and already does a full H2D plus stream sync, so taking the two slot mutexes is free, and those mutexes are the existing serialization point against
acquire/retire/poll. (I'd also avoid the word "generation" here — this repo already uses it for the pipeline lease{slot_id, generation}and theChipWorkerhigh-water mark.)
This also removes most of the current diff: acquire(slot) keeps its signature and stays in provision_native_run_resources, so the callable_aicore_image_hash() accessor, the ensure_aicpu / ensure_aicpu_locked split, the move of acquire into prepare_execution, and the aicore_image_hash == 0 failure path all become unnecessary.
Other findings
Must fix
-
docs/task-flow.md:293-301states the invariant this PR reverses and is not updated — "each run creates its own AICore stream ... and no record of which image a stream last ran is load-bearing." After this PR that record is load-bearing. Same-commit fix per.claude/rules/doc-consistency.md§1/§4. While rewriting: the paragraph's claim that "every slot publishes its image to the same GM code address" does not matchupload_chip_callable_buffer, whichrtMallocs percontent_hashwith refcounting — simultaneously-registered callables sit at different addresses, and an address is only reused after an unregister frees it. -
No test covers the hazard the policy weakens. Both
test_alternating_code_images_never_execute_stale_instructionsand the newtest_depth_two_slots_rebuild_on_independent_image_transitionsregister add and sub concurrently and only unregister infinally. Two live leases means two distinct blocks, so no code address is ever recycled and neither test can fail from stale instructions regardless of policy. The docstring correction from "at one reused GM code address" to "on one pipeline slot" is honest, but the PR description still lists A/B/A/B as coverage for this risk. The nearest existing coverage,dynamic_register/test_dynamic_register.py::test_unregister_last_handle_allows_reprepare_same_hashid, re-registers identical content and is safe by construction. A test that walks the six steps above would close it. -
All five self-hosted checks are red —
ut-a2a3,ut-a5,st-onboard-a2a3,st-onboard-a5,st-pod-onboard-a2a3, each annotated "The job was not acquired by Runner of type self-hosted even after multiple attempts." Infra starvation, not attributable to this diff — but it means no hardware validation ran in CI for a change that relaxes a hardware correctness invariant. Worth a re-run before merge. -
No performance evidence. #1791's acceptance list leads with "repeated same-image performance", and the description concedes "I did not report mixed-revision latency numbers as comparable performance evidence." The stream-count assertions prove the mechanism, not that the 3.81–4.36% is recovered. Happy to run the dsv4-flash HCA/SWA A/B if that helps.
Should fix
-
The reuse gate
s.completeis weaker than theaicore == nullptrgate it replaces.poll()setscomplete = truefrom a progress thread on a successful device query, before the executor drains. A secondacquireon a slot whose run has not been retired used to be hard-refused; it now succeeds and hands back the in-flight stream. Consider a flag set only byretire_aicore(Complete)and cleared byacquire/mark_submitted. -
retire_aicore(Complete)turns a benign no-op into an error. The old code fell throughif (s.aicore == nullptr) return 0;. It now returns −1 when the handle is gone, anddrain_executionpropagates that as the run's error after a successful device drain (reachable in principle viaabandon_all()on the device-reset path). Either keep it benign or document why absence must be an error. -
aicore_image_hash == 0is a hard prepare-time failure but is not validated at registration.record_device_orch_callable/record_host_orch_callablecheckchip_buffer_hashandchip_devbut not this field, so a registration path that omits it registers cleanly and then fails every run with "active callable has no AICore image identity", far from the cause. (Moot if thestaleshape is adopted.) -
test_depth_two_slots_own_separate_resourcesgained assertions that cannot fail —assert after_slot0 in (stream_sets, stream_sets + 1)andassert after_slot1 in (after_slot0, after_slot0 + 1)accept both branches. If slot warmth is order-dependent under the sharedst_worker, warm the slots explicitly and assert exact counts.
Consider
-
The description should surface the goal narrowing: #1791 recommended runner-wide, publication-aware reuse; this implements per-slot content-hash. Even with the above fixed, stating why is worth a line, and a
docs/investigations/entry next to2026-06-aicore-cold-start-warmup.mdwould help — the underlying mechanism is still unproven in both directions (#1791: "PMUicache_req/icache_missdata or a stream-only toggle is still needed"). -
_run_registeredand_run_registered_with_leaseeach inline(base+1)*(base+2), duplicatingcompute_golden. Parameterizingcompute_goldenonsubtractwould avoid the third copy.
Retain proven-complete per-slot AICore streams while no new child code has been uploaded. Mark all slots stale at the sole AICore H2D publication point, recreate stale streams before reuse, and preserve conservative cleanup for unproven runs and destroy failures.\n\nFixes hw-native-sys#1791
a15632e to
a9ec96b
Compare
|
@ChaoWao Addressed in The reuse policy is now runner-wide and publication-aware rather than keyed by a per-slot image hash:
Coverage now includes repeated no-publication reuse, deduplicated registration without invalidation, resident A/B alternation without recreation, the cross-slot unregister/register Validation on the new head: changed-file pre-commit passed all hooks and the I have not reported a mixed-revision latency number as performance evidence. Your offer to run the canonical dsv4-flash HCA/SWA A/B on this head would close that remaining evidence gap without conflating PyPTO/PTO-ISA revisions. |
2 participants
Summary
A -> B -> Aaddress-recycle publication, retirement gating, and exact stream-creation countsdocs/task-flow.mdwith the actual content-hash allocation and publication invariantFixes #1791
Validation
clang-format,clang-tidy,cpplint,ruff,pyright,markdownlint)RunStreamSlotsC++ unit tests: 19/19 passed with GCC 15 on LCWtask-submit(task_20260814_024716_37579914368)a9ec96bb2: 18/18 checks completed with no failures (17 success, expecteddeployskip), includingut-a2a3,ut-a5,st-onboard-a2a3,st-onboard-a5,st-pod-onboard-a2a3, both A2A3/A5 simulation matrices, packaging, profiling, and pre-commitAn additional LCW mixed-venv targeted run (
task_20260814_025329_67147427068) reached its 600-second pytest session timeout innative_run_lifecycle::test_runbefore any assertion result. It is not treated as product evidence; the clean GitHubst-onboard-a2a3run above passed the full suite on this exact head.The stream-creation assertions validate the mechanism. The issue-specific HCA/SWA latency benchmark pins older PyPTO/PTO-ISA revisions, so no mixed-revision latency number is reported as comparable performance evidence.