feat(runtime): add copy_stacked_from D2H read-back for StackedDeviceTensor

docs/en/user/00-getting_started.md

@@ -433,6 +433,16 @@ Internally each shard `host_w[i]` becomes a worker-resident `DeviceTensor`, so t
 generated `x[r]` indexing skips the H2D upload (`child_memory`). Shards are
 auto-freed on `close()` if not released earlier via `free_stacked_tensor`.
 
+Like a single `DeviceTensor`, a `StackedDeviceTensor` is never copied back
+automatically. To read the current device contents of every shard back to the
+host in one call — e.g. a resident KV cache at the end of a step — use
+`rt.copy_stacked_from(w, host_out)`, the read-back symmetric of
+`alloc_stacked_tensor`. `host_out` is filled in place (`host_out[i]` receives
+shard `i`) and, like the upload source, must be a CPU, contiguous, **shared-memory**
+`[B, *tail]` tensor matching the stack's shape and dtype, allocated before
+`prepare()` (call `.share_memory_()`): the D2H copy runs in the forked chip worker,
+which can only write host memory it inherited at fork.
+
 The leading dimension is the shard dimension and `B` must equal the number of
 cards the program dispatches to. By default shard `i` lands on worker `i`
 (matching `device=r`). If the program uses a **non-identity** placement — a

docs/zh-cn/user/00-getting_started.md

@@ -411,6 +411,14 @@ with compiled.prepare() as rt:
 `x[r]` 取下标会跳过 H2D 上传(`child_memory`)。分片在 `close()` 时自动释放,也可提前用
 `free_stacked_tensor` 释放。
 
+和单个 `DeviceTensor` 一样,`StackedDeviceTensor` 也不会被自动拷回。若要一次把每个分片
+当前的设备内容读回主机——例如某一步结束时读回常驻的 KV cache——可用
+`rt.copy_stacked_from(w, host_out)`,即 `alloc_stacked_tensor` 的对称读回接口。`host_out`
+原地填充(`host_out[i]` 接收第 `i` 片);与上传源一样,它必须是形状和 dtype 与该 stack
+匹配、且在 `prepare()` **之前**分配的 CPU、连续、**共享内存** `[B, *tail]` 张量
+(调用 `.share_memory_()`):D2H 拷贝在 fork 出的 chip worker 中执行,只能写它在 fork
+时继承的主机内存。
+
 首维就是分片维,`B` 必须等于程序分发到的卡数。默认第 `i` 片落在第 `i` 个 worker 上
 (对应 `device=r`)。如果程序用的是**非恒等**放置——置换或子集卡(如 `device=2*r`,或字面量
 `device=1` / `device=0`)——就要传匹配的 `worker_ids`,其中 `worker_ids[i]` 是程序提交

python/pypto/runtime/distributed_runner.py

@@ -1091,20 +1091,41 @@ def _require_ready(self, op: str) -> None:
         # Worker ABC hook: device-memory ops are valid until close().
         self._require_open(op)
 
-    def _prepare_init(self, init: torch.Tensor) -> torch.Tensor:
-        # Worker ABC hook: the upload (``copy_to``) runs **inside the
-        # forked chip worker**, so ``init`` must be a CPU, contiguous,
-        # shared-memory tensor allocated **before**
-        # :meth:`DistributedCompiledProgram.prepare` (call ``.share_memory_()``).
-        # Unlike L2 we cannot make a defensive ``.cpu().contiguous()`` copy: that
-        # copy would live only in the parent and be invisible to the child.
-        if not (init.is_shared() and init.is_contiguous() and init.device.type == "cpu"):
+    @staticmethod
+    def _require_forked_host_buffer(tensor: torch.Tensor, api: str, access: str) -> None:
+        """Validate *tensor* is a host buffer the forked chip worker can ``access``.
+
+        Every H2D/D2H copy runs **inside the forked chip worker**, which can only
+        touch host memory it inherited at fork. So *tensor* must be a CPU,
+        contiguous, **shared-memory** tensor allocated **before**
+        :meth:`DistributedCompiledProgram.prepare` (call ``.share_memory_()``); a
+        buffer allocated after ``prepare()`` — or a non-shared one — is invisible
+        to the child.
+
+        Args:
+            tensor: The host buffer to validate.
+            api: The calling API signature, woven into the error message
+                (e.g. ``"copy_stacked_from(host=...)"``).
+            access: The child's access verb — ``"read"`` for uploads, ``"write"``
+                for read-backs.
+
+        Raises:
+            ValueError: If *tensor* is not CPU, contiguous, and shared-memory.
+        """
+        if not (tensor.is_shared() and tensor.is_contiguous() and tensor.device.type == "cpu"):
             raise ValueError(
-                "DistributedWorker.alloc_tensor(init=...) requires a CPU, contiguous, "
-                "shared-memory tensor allocated BEFORE prepare() (call .share_memory_()). "
-                "The upload runs in the forked chip worker, which can only read host "
-                "memory it inherited at fork."
+                f"{api} requires a CPU, contiguous, shared-memory tensor allocated "
+                f"BEFORE prepare() (call .share_memory_()). The copy runs in the forked "
+                f"chip worker, which can only {access} host memory it inherited at fork."
             )
+
+    def _prepare_init(self, init: torch.Tensor) -> torch.Tensor:
+        # Worker ABC hook: the upload (``copy_to``) runs **inside the forked chip
+        # worker**, so ``init`` must be a CPU, contiguous, shared-memory tensor
+        # allocated **before** prepare() (see _require_forked_host_buffer). Unlike
+        # L2 we cannot make a defensive ``.cpu().contiguous()`` copy: that copy
+        # would live only in the parent and be invisible to the child.
+        self._require_forked_host_buffer(init, "DistributedWorker.alloc_tensor(init=...)", "read")
         return init
 
     def alloc_stacked_tensor(
@@ -1194,6 +1215,47 @@ def free_stacked_tensor(self, stacked: StackedDeviceTensor) -> None:
         for shard, w in zip(stacked.shards, stacked.worker_ids, strict=True):
             self.free_tensor(shard, worker_id=w)
 
+    def copy_stacked_from(self, stacked: StackedDeviceTensor, host: torch.Tensor) -> None:
+        """Read every shard of *stacked* back to *host* (D2H) — the read-back
+        symmetric to :meth:`alloc_stacked_tensor`.
+
+        Because a :class:`~pypto.runtime.StackedDeviceTensor` skips the
+        per-dispatch D2H copy, callers that want the shards' current device
+        contents (e.g. a resident KV cache at the end of an L3 step) must read
+        them back explicitly. Shard ``i`` is copied from its owning worker
+        ``stacked.worker_ids[i]`` into ``host[i]``.
+
+        Args:
+            stacked: The resident stacked tensor to read back.
+            host: A CPU, contiguous, **shared-memory** ``[B, *tail]`` tensor
+                allocated BEFORE :meth:`~DistributedCompiledProgram.prepare`
+                (call ``.share_memory_()``), whose shape and dtype match
+                ``stacked.full_shape`` / ``stacked.dtype``. Filled in place
+                (``host[i]`` receives shard ``i``). The D2H copy runs in the
+                forked chip worker, which can only write to host memory it
+                inherited at fork — a buffer allocated after ``prepare()`` (or a
+                non-shared one) would leave *host* untouched.
+        """
+        self._require_open("copy_stacked_from")
+        if not isinstance(stacked, StackedDeviceTensor):
+            raise TypeError(
+                f"copy_stacked_from(stacked=...) expects a StackedDeviceTensor, got {type(stacked).__name__}"
+            )
+        if not isinstance(host, torch.Tensor):
+            raise TypeError(f"copy_stacked_from(host=...) expects a torch.Tensor, got {type(host).__name__}")
+        if tuple(host.shape) != stacked.full_shape:
+            raise ValueError(
+                f"host shape {tuple(host.shape)} does not match stacked full_shape {stacked.full_shape}"
+            )
+        if host.dtype != stacked.dtype:
+            raise ValueError(f"host dtype {host.dtype} does not match stacked dtype {stacked.dtype}")
+        self._require_forked_host_buffer(host, "copy_stacked_from(host=...)", "write")
+        for i, (shard, w) in enumerate(zip(stacked.shards, stacked.worker_ids, strict=True)):
+            # host is contiguous + shared, so host[i] is a contiguous view at the
+            # right offset into the same shm segment the child inherited at fork;
+            # host[i].data_ptr() is therefore the correct cross-process D2H dst.
+            self.copy_from(host[i].data_ptr(), shard.data_ptr, shard.nbytes, worker_id=w)
+
     # ------------------------------------------------------------------
     # Dispatch
     # ------------------------------------------------------------------

tests/st/distributed/test_l3_stacked_device_tensor.py

@@ -145,6 +145,9 @@ def _run_stacked(test_config, device_ids, program, worker_ids):
     host_b = torch.empty((N_RANKS, DIM, DIM), dtype=torch.float32).share_memory_()
     for i in range(N_RANKS):
         host_b[i].fill_(10.0 + i)  # shard i holds 10 + i
+    # D2H read-back destination — like every host buffer the forked worker writes,
+    # it must be shared memory allocated BEFORE prepare() (see copy_stacked_from).
+    host_readback = torch.zeros((N_RANKS, DIM, DIM), dtype=torch.float32).share_memory_()
 
     with compiled.prepare() as rt:
         weight = rt.alloc_stacked_tensor(host_b, worker_ids=worker_ids)
@@ -162,6 +165,16 @@ def _run_stacked(test_config, device_ids, program, worker_ids):
                 for i in range(N_RANKS):
                     expected[i].fill_(host_a_val + 10.0 + i)  # a + (10 + i)
                 torch.testing.assert_close(host_f, expected, rtol=1e-5, atol=1e-5)
+
+            # Read the resident stack back to the host (D2H). The kernel only
+            # reads ``b``, so each shard still holds ``10 + i`` even though the
+            # host source was zeroed after upload — a clean copy_stacked_from proof.
+            host_readback.zero_()
+            rt.copy_stacked_from(weight, host_readback)
+            for i in range(N_RANKS):
+                torch.testing.assert_close(
+                    host_readback[i], torch.full((DIM, DIM), 10.0 + i), rtol=1e-5, atol=1e-5
+                )
         finally:
             rt.free_stacked_tensor(weight)
 

tests/ut/runtime/test_distributed_worker.py

@@ -290,18 +290,18 @@ def test_alloc_tensor_forwards_nonzero_worker_id(self, patched_setup):
         rt.close()
 
 
+def _compiled_2cards():
+    compiled = _fake_compiled([_param("b", [2, 4, 4])], [])
+    compiled._distributed_config = DistributedConfig(device_ids=[0, 1])
+    return compiled
+
+
 class TestAllocStackedTensor:
     """``alloc_stacked_tensor`` uploads each leading-dim shard to its worker once."""
 
-    @staticmethod
-    def _compiled_2cards():
-        compiled = _fake_compiled([_param("b", [2, 4, 4])], [])
-        compiled._distributed_config = DistributedConfig(device_ids=[0, 1])
-        return compiled
-
     def test_identity_uploads_shard_per_worker(self, patched_setup):
         patched_setup["worker"]._orch.malloc.side_effect = [0xA000, 0xB000]
-        rt = DistributedWorker(self._compiled_2cards())
+        rt = DistributedWorker(_compiled_2cards())
         host = torch.arange(2 * 4 * 4, dtype=torch.float32).view(2, 4, 4).share_memory_()
 
         stacked = rt.alloc_stacked_tensor(host)  # default worker_ids = range(2)
@@ -323,7 +323,7 @@ def test_identity_uploads_shard_per_worker(self, patched_setup):
 
     def test_permuted_worker_ids_place_shards(self, patched_setup):
         patched_setup["worker"]._orch.malloc.side_effect = [0xA000, 0xB000]
-        rt = DistributedWorker(self._compiled_2cards())
+        rt = DistributedWorker(_compiled_2cards())
         host = torch.zeros(2, 4, 4, dtype=torch.float32).share_memory_()
 
         stacked = rt.alloc_stacked_tensor(host, worker_ids=[1, 0])
@@ -340,7 +340,7 @@ def test_permuted_worker_ids_place_shards(self, patched_setup):
 
     def test_free_stacked_tensor_releases_each_shard(self, patched_setup):
         patched_setup["worker"]._orch.malloc.side_effect = [0xA000, 0xB000]
-        rt = DistributedWorker(self._compiled_2cards())
+        rt = DistributedWorker(_compiled_2cards())
         host = torch.zeros(2, 4, 4, dtype=torch.float32).share_memory_()
         stacked = rt.alloc_stacked_tensor(host, worker_ids=[1, 0])
 
@@ -356,7 +356,7 @@ def test_free_stacked_tensor_releases_each_shard(self, patched_setup):
 
     def test_close_auto_frees_stacked_shards(self, patched_setup):
         patched_setup["worker"]._orch.malloc.side_effect = [0xA000, 0xB000]
-        rt = DistributedWorker(self._compiled_2cards())
+        rt = DistributedWorker(_compiled_2cards())
         host = torch.zeros(2, 4, 4, dtype=torch.float32).share_memory_()
         rt.alloc_stacked_tensor(host)  # leak — close() must release both shards
 
@@ -367,7 +367,7 @@ def test_close_auto_frees_stacked_shards(self, patched_setup):
         orch.free.assert_any_call(1, 0xB000)
 
     def test_worker_ids_out_of_range_rejected(self, patched_setup):
-        rt = DistributedWorker(self._compiled_2cards())
+        rt = DistributedWorker(_compiled_2cards())
         host = torch.zeros(2, 4, 4, dtype=torch.float32).share_memory_()
         with pytest.raises(ValueError, match="out of range"):
             rt.alloc_stacked_tensor(host, worker_ids=[0, 5])
@@ -376,23 +376,23 @@ def test_worker_ids_out_of_range_rejected(self, patched_setup):
     def test_empty_leading_dim_rejected(self, patched_setup):
         # B == 0 must fail cleanly (before any malloc), not build an empty
         # StackedDeviceTensor that IndexErrors on .dtype / __repr__.
-        rt = DistributedWorker(self._compiled_2cards())
+        rt = DistributedWorker(_compiled_2cards())
         host = torch.zeros(0, 4, 4, dtype=torch.float32).share_memory_()
         with pytest.raises(ValueError, match="at least one shard"):
             rt.alloc_stacked_tensor(host)
         patched_setup["worker"]._orch.malloc.assert_not_called()
         rt.close()
 
     def test_worker_ids_length_mismatch_rejected(self, patched_setup):
-        rt = DistributedWorker(self._compiled_2cards())
+        rt = DistributedWorker(_compiled_2cards())
         host = torch.zeros(2, 4, 4, dtype=torch.float32).share_memory_()
         with pytest.raises(ValueError, match="entries"):
             rt.alloc_stacked_tensor(host, worker_ids=[0])
         rt.close()
 
     def test_non_shared_host_rejected_and_rolled_back(self, patched_setup):
         patched_setup["worker"]._orch.malloc.side_effect = [0xA000, 0xB000]
-        rt = DistributedWorker(self._compiled_2cards())
+        rt = DistributedWorker(_compiled_2cards())
         host = torch.zeros(2, 4, 4, dtype=torch.float32)  # NOT shared
 
         with pytest.raises(ValueError, match="shared-memory"):
@@ -402,6 +402,91 @@ def test_non_shared_host_rejected_and_rolled_back(self, patched_setup):
         rt.close()
 
 
+class TestCopyStackedFrom:
+    """``copy_stacked_from`` reads each resident shard back into host[i] (D2H)."""
+
+    def _make_stacked(self, patched_setup, worker_ids=None):
+        patched_setup["worker"]._orch.malloc.side_effect = [0xA000, 0xB000]
+        rt = DistributedWorker(_compiled_2cards())
+        host = torch.zeros(2, 4, 4, dtype=torch.float32).share_memory_()
+        stacked = rt.alloc_stacked_tensor(host, worker_ids=worker_ids)
+        patched_setup["worker"]._orch.copy_from.reset_mock()
+        return rt, stacked
+
+    def test_reads_each_shard_back(self, patched_setup):
+        rt, stacked = self._make_stacked(patched_setup)  # worker_ids == (0, 1)
+        out = torch.zeros(2, 4, 4, dtype=torch.float32).share_memory_()
+
+        rt.copy_stacked_from(stacked, out)
+
+        orch = patched_setup["worker"]._orch
+        nbytes = 4 * 4 * 4
+        # Facade arg order: copy_from(worker_id, dst_host_ptr, src_dev_ptr, nbytes).
+        orch.copy_from.assert_any_call(0, out[0].data_ptr(), 0xA000, nbytes)
+        orch.copy_from.assert_any_call(1, out[1].data_ptr(), 0xB000, nbytes)
+        assert orch.copy_from.call_count == 2
+        rt.close()
+
+    def test_permuted_worker_ids(self, patched_setup):
+        rt, stacked = self._make_stacked(patched_setup, worker_ids=[1, 0])
+        out = torch.zeros(2, 4, 4, dtype=torch.float32).share_memory_()
+
+        rt.copy_stacked_from(stacked, out)
+
+        orch = patched_setup["worker"]._orch
+        nbytes = 4 * 4 * 4
+        # shard 0 resides on worker 1, shard 1 on worker 0.
+        orch.copy_from.assert_any_call(1, out[0].data_ptr(), 0xA000, nbytes)
+        orch.copy_from.assert_any_call(0, out[1].data_ptr(), 0xB000, nbytes)
+        rt.close()
+
+    def test_shape_mismatch_rejected(self, patched_setup):
+        rt, stacked = self._make_stacked(patched_setup)
+        out = torch.zeros(3, 4, 4, dtype=torch.float32).share_memory_()
+        with pytest.raises(ValueError, match="does not match stacked full_shape"):
+            rt.copy_stacked_from(stacked, out)
+        rt.close()
+
+    def test_dtype_mismatch_rejected(self, patched_setup):
+        rt, stacked = self._make_stacked(patched_setup)
+        out = torch.zeros(2, 4, 4, dtype=torch.float16).share_memory_()
+        with pytest.raises(ValueError, match="does not match stacked dtype"):
+            rt.copy_stacked_from(stacked, out)
+        rt.close()
+
+    def test_non_shared_host_rejected(self, patched_setup):
+        # A plain (non-shared) host buffer is invisible to the forked worker's
+        # D2H write — reject it up front rather than silently returning zeros.
+        rt, stacked = self._make_stacked(patched_setup)
+        out = torch.zeros(2, 4, 4, dtype=torch.float32)  # NOT shared
+        with pytest.raises(ValueError, match="shared-memory"):
+            rt.copy_stacked_from(stacked, out)
+        rt.close()
+
+    def test_non_contiguous_host_rejected(self, patched_setup):
+        rt, stacked = self._make_stacked(patched_setup)
+        # Shared but transposed -> non-contiguous; still rejected.
+        out = torch.zeros(2, 4, 4, dtype=torch.float32).share_memory_().transpose(1, 2)
+        assert not out.is_contiguous()
+        with pytest.raises(ValueError, match="shared-memory"):
+            rt.copy_stacked_from(stacked, out)
+        rt.close()
+
+    def test_wrong_type_rejected(self, patched_setup):
+        rt, _stacked = self._make_stacked(patched_setup)
+        out = torch.zeros(2, 4, 4, dtype=torch.float32).share_memory_()
+        with pytest.raises(TypeError, match="expects a StackedDeviceTensor"):
+            rt.copy_stacked_from(object(), out)  # type: ignore[arg-type]  # runtime guard under test
+        rt.close()
+
+    def test_after_close_raises(self, patched_setup):
+        rt, stacked = self._make_stacked(patched_setup)
+        out = torch.zeros(2, 4, 4, dtype=torch.float32).share_memory_()
+        rt.close()
+        with pytest.raises(RuntimeError, match="called after close"):
+            rt.copy_stacked_from(stacked, out)
+
+
 class TestLifecycle:
     def test_close_idempotent_and_closes_worker(self, patched_setup):
         compiled = _fake_compiled([_param("a", [16, 16])], [])