[Feature] Let a host-side orchestrator read and write device-resident tensors

[ChaoWao]

Summary

A device-resident tensor cannot be read or written by a host-side orchestrator. Today host_build_graph's orchestrator can only reach a tensor's contents as a side effect of staging: runtime_maker.cpp:1232 registers the caller's host buffer as the read window while staging the tensor. A tensor that arrives already on the device takes the is_device_memory() pass-through at :1184 and skips that registration, so get_tensor_data / set_tensor_data hit report_fatal.

The existing error text already names the situation:

no host view for device address …: during host orchestration only tensors the runtime staged are readable, **not runtime-created or child-memory buffers**

Give a device-resident tensor an orchestrator-accessible view, established on demand and cached for the allocation's lifetime, so residency and data-dependent host orchestration stop being mutually exclusive.

Motivation / Use Case

Residency and data-dependent orchestration are mutually exclusive today, and both are wanted.

Keeping a tensor device-resident removes the entire per-run staging path — device_malloc (:1194), H2D (:1204), D2H copy-back, and device_free (:1449) — because the pass-through branch continues past all of it. #2041 already exploits this: the qwen3-14B decode driver keeps every parameter resident and runs under both runtimes.

But qwen's graph is static. 11 host_build_graph orchestrations are data-dependent — they read an input's contents on the host to decide the graph shape:

// tests/st/a2a3/host_build_graph/paged_attention/kernels/orchestration/paged_attention_orch.cpp
uint64_t cur_seq       = get_tensor_data<int32_t>(context_lens, 1, cl_idx);   // :155
uint64_t bn_this_batch = (cur_seq + block_size - 1) / block_size;             // ← how many tasks
uint64_t cur_block_idx = get_tensor_data<int32_t>(block_table, 2, bt_idx);    // :181
simpler::hbg::Tensor kj = key_cache.view(kv_shapes, kv_offsets);             // ← which KV slice

grep -rl get_tensor_data --include=*.cpp examples/ tests/st/ returns 30 sources, 11 under host_build_graph (the paged_attention / batch_paged_attention / paged_attention_unroll family, the two *_manual_scope cases, benchmark_bgemm, deepseek_v4_flash_decode, host_build_graph_validation). A further 6 hbg orchestrations call set_tensor_data, so writes are real too, not hypothetical.

This is not urgent, and the issue says so. Because the declaration is per-argument, these cases have a zero-cost workaround today: make the large tensors resident and leave the small control tensors host-staged.

	measured warm chip.run
all host-staged	1.006–1.013 ms
large tensors resident, control tensors staged	0.597–0.695 ms
control tensors also resident, with orchestrator views	0.566–0.597 ms

(measurements from paged_attention_unroll_manual_scope on A2/A3, reported in #1854)

The mixed configuration captures ~90% of the win. The trigger for doing this work is a tensor that is both large and orchestrator-accessed — today none exists: every tensor actually read or written by an orchestration is ≤ 256 KiB, while the bulk tensors in the same cases are ~512 MiB.

tensor	shape	bytes	accesses
block_table (pa Case1)	(256, 256) i32	256 KiB	16,384
block_table (pa Case2)	(64, 512) i32	128 KiB	8,192
context_lens	(batch,) i32	256 B – 1 KiB	64 – 256
num_tokens_per_owner (dsv4)	(N_RANKS,) i32	~32–64 B	N_RANKS
ext_config (bgemm)	a few i64	~32 B	4
shape / layout / external	i32 slots	a few bytes	2–4 (writes)

Proposed API / Behavior

No user-facing API. Correctness should not depend on a declaration — a user who forgets one should not get a program that fails to run. Establish the view on demand.

bind:  for each device-memory tensor, push a lazy entry {base, span, means = none}
       — a plain vector push, no decision, no cost

read/write miss (the orchestration actually touched it):
       lazy entry with means == none → establish now, cache at allocation scope
       → serve the access

The set that gets a view is then exactly the set that is actually accessed — better than any declaration could be, and precisely what #1848 concluded: "registering only the tensors the orchestrator may actually read".

Two mechanisms, SVM preferred:

condition	mechanism	per-access cost
host map available	SVM mapping	0
host map unavailable (platform / size / host page config)	one bulk rtMemcpy mirror	0 (reads hit the mirror)

The fallback must be a bulk mirror, not per-access read-through: block_table in paged_attention Case1 is read 16,384 times at 4 bytes each, so read-through would be ~16k PCIe round trips (16–33 ms) against tens of µs for one 256 KiB copy.

The existing accessor already expresses both mechanisms — no new API is needed:

add(base, span, nullptr);   // SVM: calls register_device_memory_to_host itself;
                            //      needs_push_back = false
add(base, span, mirror);    // mirror: needs_push_back = true, so a write does
                            //      memcpy + copy_to_device automatically

The one missing step is filling the mirror with a copy_from_device before add — add() assumes the fallback view already holds the bytes, which is true on the staged path (it was just uploaded) but not for a resident tensor.

Cost is amortized over the allocation, not the bind. This is the whole lesson of #1848: hbg once registered every entry tensor per run, and on qwen that was 20 tensors / 38.05 GiB with ~256 ms per run spent in teardown alone — 5.7x the device execution — because the cost is O(mapped bytes). The means must therefore be held by whoever owns device_malloc / device_free (DeviceRunnerBase and its mem_alloc_, alongside the existing per-slot retained tables), established once and released before the allocation is freed. HostTensorAccessor::close() must stop tearing mappings down per bind.

Coherence:

mechanism	before a read	after a write
SVM mapping	nothing (coherent by construction)	nothing
mirror	refresh by D2H when stale	memcpy into the mirror + immediate copy_to_device

Staleness is decided by the existing CallableState::signature[i] (device_runner_base.cpp:932/:970 write it, :1258 reads it): the kernel's direction for this run says whether the device may have advanced. No new field on the callable.

Two axes, kept separate: the lazy entry decides whether a means is needed; signature[i] decides whether it can go stale. Direction must not be used as an access guard — a tensor that is IN to the kernel can legitimately be read and written by the orchestrator (the orchestration writes a config value the kernel only reads), which is what the 6 set_tensor_data orchestrations do.

Host addresses never leave the worker. The means is established inside the process that runs bind and orchestration, and only the device address identifies it. At L2 that process is the caller's; at L3+ it is the forked child, which establishes its own means. Nothing that crosses a fork carries a host VA, so this works at L3 without a separate mechanism.

tensormap_and_ringbuffer needs nothing. Its orchestrator runs on the AICPU and reads device memory directly, so the requirement is satisfied by construction. The difference is carried by which runtime_maker reads these fields — no runtime-conditional belongs in the framework or the Python layer.

Observability: record which mechanism was selected. On a host affected by #1531 every tensor in this size class falls back to the mirror, and "why is this machine slower" needs an answer. A one-line warning when establishing a view over ~1 MiB also keeps #1848's failure mode visible — today's maximum is 256 KiB, so that is 4x headroom and changes no behaviour.

Alternatives Considered

1. Leave control tensors host-staged (works today, costs nothing). Since residency is declared per argument, a case can keep key_cache resident and context_lens / block_table staged. This is the middle row of the measurement table and captures most of the win. It is the reason this issue is not urgent, and it is what #1848 already prescribed: "any change here must keep control tensors host-side."

2. Have the caller pass a host view address alongside the device tensor. Rejected: an address is only valid in the process that created it, so it forces an explicit L3 refusal, and it puts a host pointer into structures that otherwise cross a process boundary. It also pins the cost to the bind rather than to the allocation, which is exactly the shape #1848 measured at 256 ms/run.

3. Register host views unconditionally, as hbg once did. Measured and rejected — see #1848.

4. Declare orchestrator access explicitly and refuse when absent. Rejected as the default: it makes correctness depend on the user getting a declaration right, and users will over-declare, which drifts back toward alternative 3. An optional declaration remains useful as a tuning knob (move establishment from the orchestration window to allocation time), but not as a correctness requirement.

Additional Context

Related: #1848 (closed — the measurement that rules out unconditional registration, and the source of the "keep control tensors host-side" guidance), #1531 (open — halHostRegister returns DRV_ERROR_NOT_SUPPORT for ordinary-page small allocations on 64 KiB-page hosts, which is exactly this size class, so mechanism selection must branch on the per-allocation return value and never on the architecture), #2041 (merged — full parameter residency for qwen decode, the precedent this builds on), #1841 (adjacent: device-side allocation dominating the same bind phase), #1854 (open — implements residency plus a caller-supplied host view; this issue records the alternative shape for the orchestrator-access half).

Platform capability for reference:

	a2a3 onboard	a5 onboard	sim
register_device_memory_to_host	halHostRegister(DEV_SVM_MAP_HOST) (a2a3/platform/onboard/host/device_runner.cpp:889)	nullptr (base default, device_runner_base.h:164)	identity (sim/host/device_runner_base.h:224)