Add: host-orchestrator access to child-memory tensors - #2212
Merged
ChaoWao merged 1 commit intoSep 13, 2026
Merged
Conversation
Fixes hw-native-sys#2205 A host_build_graph orchestration reads an input's bytes on the host to decide the graph shape: 11 orchestrations call get_tensor_data and 6 call set_tensor_data. That worked only as a side effect of staging, because the bind registered the caller's host buffer as the read window while staging the tensor. A tensor declared child_memory=True takes the is_device_memory() pass-through and skips that registration, so an access to it hit report_fatal and residency was mutually exclusive with data-dependent orchestration. The bind now claims each child-memory span with add_child_memory, which is a vector push and consults nothing. The first access landing inside a span resolves how to reach it: acquire_child_memory_host_view returns a mapping, or null and each access becomes a device copy. The mapped set is therefore exactly the set the orchestration touched, which is what hw-native-sys#1848 concluded after measuring unconditional registration at ~256 ms/run. The mapping covers the whole tracked allocation and is owned by the runner that owns device_malloc/device_free, so several tensors in one child buffer share it and free_tensor drops it before the pages go. Measured on a2a3 / CANN 9.0.0, a register/unregister pair costs ~5.2 us before any bytes are mapped and ~7.0 ms/GiB beyond that, so a two-tensor bind would pay ~24 us every run without that ownership -- see docs/investigations/2026-09-hbg-per-run-host-view-rebuild.md, which also corrects hw-native-sys#1848's reading that the cost is essentially all per-byte. The device-copy path holds no state, so a read cannot serve stale bytes and a write reaches the device immediately; it needs no staleness tracking. It is not an a5-only branch -- hw-native-sys#1531 refuses ordinary-page small allocations on 64 KiB-page hosts, which is this size class -- so its accesses are counted and reported as devcopy=N in the bind's host_view_close attributes. Runtime-created graph-heap tensors stay unreadable, and both report_fatal messages are updated to say so instead of naming child memory. Verified: the new paged_attention child-memory scene tests fail on main with the exact report_fatal and pass here; onboard a2a3 with 5 rounds establishes 2 mappings rather than 10, so the cost is amortized over the allocation. cpput 144/144, pyut 2285 passed / 18 skipped, a2a3sim and a5sim sweeps green.
|
Warning Review limit reachedNext included review available in 39 minutes. View limit detailsLimit details: You’ve used the included review currently available. You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository. Review configuration: ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Advanced Run ID: 📒 Files selected for processing (24)
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 |
5 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
A
host_build_graphorchestration runs on the host and reads an input's bytes to decide the graph shape — 11 orchestrations callget_tensor_data, 6 callset_tensor_data. That only ever worked as a side effect of staging: the bind registered the caller's host buffer as the read window while staging the tensor. A tensor declaredchild_memory=Truetakes theis_device_memory()pass-through and skips that registration, so an access to it hitreport_fatal— residency and data-dependent orchestration were mutually exclusive.The bind now claims each child-memory span, and how to reach it is decided on first access:
Nothing is registered for a tensor the orchestration never touches, so the mapped set is exactly the accessed set — which is what #1848 concluded after measuring unconditional registration at ~256 ms/run on qwen.
Who owns the mapping, and why it isn't the bind
The mapping covers the whole tracked allocation and is owned by the runner that owns
device_malloc/device_free. Two consequences: several tensors and views inside one child buffer share it, andfree_tensordrops it before the pages go — a cached host VA outliving its allocation would otherwise hand out a live mapping into a returned page.I initially wanted to skip the cache and rebuild per run, which is less code and has three arguments in its favour that hold regardless of cost. The measurement said no, on a rule fixed before the number was seen:
There is a ~5.2 µs per-call floor, flat from 64 B to 4 KiB. #1848 read the cost as essentially all per-byte, inferring it from the size asymmetry of one aggregate point — its own text says the split was never measured. Above ~1 MiB the slope is 7.0 ms/GiB, which does confirm it; both readings are right in their own regime. A realistic
paged_attentionbind reads two control tensors for 24 µs/run, ~4% of a 0.6 ms warmchip.run, against a 10 µs budget. Full write-up, including why register is per mapped page (so 256 KiB on ordinary pages costs more than 4 MiB on huge ones), indocs/investigations/2026-09-hbg-per-run-host-view-rebuild.md.The three rebuild arguments survive as constraints on the cache rather than reasons to skip it: invalidation is why it sits next to
mem_alloc_and is dropped infree_tensor/finalize, and the monotonic-union and pinned-page concerns are why held bytes are logged.No mirror on the fallback
Where no mapping is available the fallback re-copies per access rather than holding a host mirror. That keeps the path stateless — no allocation, no staleness flag, no refresh policy, no write-back ordering — so a read cannot serve stale bytes and a write reaches the device immediately.
It is not an a5-only branch: a5 onboard has no host-map path (
device_runner_base.h:177returnsnullptr), and #1531 refuses ordinary-page small allocations on 64 KiB-page hosts, which is exactly this size class. The cost is one PCIe round trip per access, so accesses are counted and surface asdevcopy=Nin the bind'shost_view_closeattributes;docs/testing.mdsays when to leave a heavily-read tensor host-staged instead.Runtime-created graph-heap tensors stay unreadable. Both
report_fatalmessages are updated to say that instead of naming child memory.Testing
mainwith the exactreport_fatalfrom the issue —get_tensor_data: FATAL(code=5): no host view for device address 0xaaaadb2ee3d0 (4 bytes)— and pass here. Built both trees to confirm.--rounds 5establishes 2 mappings, not 10.task-submit:small1_child_memoryPASSED, both control tensors mapped viahalHostRegister.close().a2a3simanda5simsweeps overexamples tests/st: 42 + 40 L3 cases each, 0 failures.check_retired_names.clang_tidy.pyonly indexes sim compile databases and silently skips files absent from them, so the onboard-only changes were run against the onboard database by hand — clean.Fixes #2205