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)