perf(pd): overlap the decode-side KV inject copies across device pairs

tilert/models/deepseek_v3_2/generator.py

@@ -11,6 +11,7 @@
 from tilert.models.deepseek_v3_2.modules.end2end import ShowHandsDSALayer
 from tilert.models.deepseek_v3_2.temp_var_indices import Idx
 from tilert.tilert_init import tilert_init
+from tilert.utils import copy_by_device_pair
 
 __all__ = [
     "DSAv32Generator",
@@ -463,6 +464,7 @@ def inject_cache(
 
         num_devices = self.decode_layer.num_devices
 
+        copies = []
         for device_id in range(num_devices):
             _, caches, _, _ = self.decode_layer._get_device_result(device_id)
 
@@ -473,13 +475,11 @@ def inject_cache(
 
                 base_idx = layer_id * 3
 
-                ki_src = ki[:cache_len].to(f"cuda:{device_id}")
-                kv_src = kv[:cache_len].to(f"cuda:{device_id}")
-                pe_src = pe[:cache_len].to(f"cuda:{device_id}")
+                for _off, _src in ((0, ki), (1, kv), (2, pe)):
+                    _dst = caches[base_idx + _off][0, start_pos:end_pos, :]
+                    copies.append((_dst, _src[:cache_len]))
 
-                caches[base_idx + 0][0, start_pos:end_pos, :].copy_(ki_src)
-                caches[base_idx + 1][0, start_pos:end_pos, :].copy_(kv_src)
-                caches[base_idx + 2][0, start_pos:end_pos, :].copy_(pe_src)
+        copy_by_device_pair(copies, self.__dict__.setdefault("_inject_streams", {}))
 
         logger.info(f"Cache injection completed for {num_devices} devices")
 

tilert/models/glm_5/generator.py

@@ -11,6 +11,7 @@
 from tilert.models.glm_5.modules.end2end import ShowHandsDSALayer
 from tilert.models.glm_5.temp_var_indices import Idx
 from tilert.tilert_init import tilert_init
+from tilert.utils import copy_by_device_pair
 
 __all__ = [
     "GLM5Generator",
@@ -457,6 +458,7 @@ def inject_cache(
 
         num_devices = self.decode_layer.num_devices
 
+        copies = []
         for device_id in range(num_devices):
             _, caches, _, _ = self.decode_layer._get_device_result(device_id)
 
@@ -467,17 +469,13 @@ def inject_cache(
 
                 base_idx = layer_id * 3
 
-                ki_src = ki[:cache_len].to(f"cuda:{device_id}")
-                kv_src = kv[:cache_len].to(f"cuda:{device_id}")
-                pe_src = pe[:cache_len].to(f"cuda:{device_id}")
-
-                for _off, _src in ((0, ki_src), (1, kv_src), (2, pe_src)):
+                for _off, _src in ((0, ki), (1, kv), (2, pe)):
                     _dst = caches[base_idx + _off]
                     if _dst.size(1) < end_pos:
                         continue
-                    _dst[0, start_pos:end_pos, :].copy_(_src)
+                    copies.append((_dst[0, start_pos:end_pos, :], _src[:cache_len]))
 
-            torch.cuda.synchronize(device_id)
+        copy_by_device_pair(copies, self.__dict__.setdefault("_inject_streams", {}))
 
         logger.info(f"Cache injection completed for {num_devices} devices")
 

tilert/utils.py

@@ -88,3 +88,50 @@ def relative_l2_error(gt: torch.Tensor, out: torch.Tensor) -> Any:
         The relative L2 error.
     """
     return torch.norm(gt - out) / torch.norm(gt)
+
+
+def copy_by_device_pair(
+    copies: list[tuple[torch.Tensor, torch.Tensor]],
+    streams: dict[tuple[int, int], torch.cuda.Stream],
+) -> None:
+    """Run ``dst.copy_(src)`` for every pair, one stream pair per device pair.
+
+    torch fences a cross-device ``copy_`` against the current stream of both
+    devices, so issuing many cache copies on the default streams runs them one
+    at a time over a single link. Grouping them by (destination, source) device
+    and giving each group its own streams lets the pairs overlap. ``streams``
+    caches the streams between calls. Returns after every copy has finished.
+
+    Args:
+        copies: (destination, source) tensor pairs. Destinations must be CUDA
+            tensors; sources may be CUDA or CPU tensors.
+        streams: Cache of streams keyed by (device, peer device).
+    """
+
+    def stream(dev: int, peer: int) -> torch.cuda.Stream:
+        key = (dev, peer)
+        if key not in streams:
+            streams[key] = torch.cuda.Stream(device=dev)
+        return streams[key]
+
+    by_pair: dict[tuple[int, int], list[tuple[torch.Tensor, torch.Tensor]]] = {}
+    for dst, src in copies:
+        src_dev = src.device.index if src.is_cuda else -1
+        by_pair.setdefault((dst.device.index, src_dev), []).append((dst, src))
+    devices = set()
+    for (dst_dev, src_dev), group in by_pair.items():
+        devices.add(dst_dev)
+        dst_stream = stream(dst_dev, src_dev)
+        if src_dev < 0:
+            with torch.cuda.stream(dst_stream):
+                for dst, src in group:
+                    dst.copy_(src, non_blocking=True)
+            continue
+        devices.add(src_dev)
+        with torch.cuda.stream(dst_stream), torch.cuda.stream(stream(src_dev, dst_dev)):
+            for dst, src in group:
+                dst.copy_(src, non_blocking=True)
+    for st in streams.values():
+        st.synchronize()
+    for dev in devices:
+        torch.cuda.synchronize(dev)