feat(runtime): add copy_stacked_from D2H read-back for StackedDeviceTensor#1954
Conversation
…ensor Symmetric to alloc_stacked_tensor: reads every resident shard of a StackedDeviceTensor back to a host [B, *tail] buffer, filling host[i] from shard i's owning worker. Like the upload path, the D2H copy runs in the forked chip worker, so the host buffer must be a CPU, contiguous, shared-memory tensor allocated before prepare(). Extract the shared forked-worker host-buffer validation into _require_forked_host_buffer, now used by both the upload (read) and read-back (write) paths. Add ut coverage (mocked worker) and an st read-back proof, and document the API in the en/zh-cn guides.
|
Caution Review failedAn error occurred during the review process. Please try again later. 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 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces the copy_stacked_from API to read back device contents of a StackedDeviceTensor to a host tensor (D2H copy). It includes documentation updates in both English and Chinese, a helper method _require_forked_host_buffer to validate host buffers, and comprehensive system and unit tests for the new functionality. There are no review comments to address, and the changes look solid.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
2 participants
Summary
Adds
DistributedWorker.copy_stacked_from(stacked, host)— the read-back symmetric toalloc_stacked_tensor. It reads every resident shard of aStackedDeviceTensorback to the host in one call, fillinghost[i]from shardi's owning worker (D2H). This lets callers recover the current device contents of a resident stack (e.g. a KV cache at the end of an L3 step), which the per-dispatch path otherwise skips.Like the upload path, the D2H copy runs inside the forked chip worker, so
hostmust be a CPU, contiguous, shared-memory[B, *tail]tensor allocated beforeprepare()(matching the stack's shape/dtype); this is validated up front.Changes
copy_stacked_from: per-shard D2H read-back overstacked.worker_ids, mirroring the existingalloc_stacked_tensor/free_stacked_tensorper-shard loops._require_forked_host_buffer, now used by both the upload (read) and read-back (write) paths — removes the duplicatedis_shared()/is_contiguous()/cpucheck.Testing
tests/ut/runtime/test_distributed_worker.py— newTestCopyStackedFrom(mocked worker): happy path, permuted worker_ids, and shape/dtype/type/non-shared/non-contiguous/after-close rejections. 80 passed.tests/st/distributed/test_l3_stacked_device_tensor.py— read-back proof: host source is zeroed after upload, so a correct10+iread-back confirms the D2H path reads device memory, not stale host state.