From 40a14afbeb8c93d3047a8d8490b0a5ce0a70e813 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Wed, 23 Sep 2026 17:50:58 +0800 Subject: [PATCH] perf(pd): overlap the decode-side KV inject copies across device pairs inject_cache copied every layer's ki/kv/pe into every rank's cache one copy_ after another on the default streams, through a temporary .to() copy per tensor. torch fences a cross-device copy_ against the current stream of both devices, so all of those copies ran one at a time over a single link, and the injection sat directly on the PD request's TTFT. copy_by_device_pair (tilert/utils.py) groups the (destination, source) pairs by device pair and issues each group on its own stream pair, copying straight into the cache slice, then synchronises every stream and device it used. GLM5Generator and DSAv32Generator.inject_cache use it; the bytes copied and their destinations are unchanged. --- tilert/models/deepseek_v3_2/generator.py | 12 +++--- tilert/models/glm_5/generator.py | 12 +++--- tilert/utils.py | 47 ++++++++++++++++++++++++ 3 files changed, 58 insertions(+), 13 deletions(-) diff --git a/tilert/models/deepseek_v3_2/generator.py b/tilert/models/deepseek_v3_2/generator.py index fb7a467..64c7640 100644 --- a/tilert/models/deepseek_v3_2/generator.py +++ b/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") diff --git a/tilert/models/glm_5/generator.py b/tilert/models/glm_5/generator.py index 18c422a..9a7bb04 100644 --- a/tilert/models/glm_5/generator.py +++ b/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") diff --git a/tilert/utils.py b/tilert/utils.py index 4cc3b47..e2afce6 100644 --- a/tilert/utils.py +++ b/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)