From 744b1c9e6ff862ed0d43ddcd0794e42fa666fc10 Mon Sep 17 00:00:00 2001 From: Kaixi Date: Tue, 17 Mar 2026 22:44:38 +0100 Subject: [PATCH] Added fallback to individual `copy_` (#20683) --- .../sglang/srt/model_executor/cuda_graph_runner.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/model_executor/cuda_graph_runner.py b/python/sglang/srt/model_executor/cuda_graph_runner.py index e3840daa7..ebe4c41d7 100644 --- a/python/sglang/srt/model_executor/cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/cuda_graph_runner.py @@ -93,9 +93,19 @@ logger = logging.getLogger(__name__) if TYPE_CHECKING: from sglang.srt.model_executor.model_runner import ModelRunner +_has_foreach_copy = hasattr(torch, "_foreach_copy_") + def _grouped_foreach_copy_(dsts: List[torch.Tensor], srcs: List[torch.Tensor]) -> None: """Call torch._foreach_copy_ grouped by (dst_dtype, src_dtype) pairs.""" + + def foreach_copy(dsts: List[torch.Tensor], srcs: List[torch.Tensor]) -> None: + if _has_foreach_copy: + torch._foreach_copy_(dsts, srcs) + else: + for dst, src in zip(dsts, srcs): + dst.copy_(src) + groups: Dict[Tuple[torch.dtype, torch.dtype], Tuple[List, List]] = {} for dst, src in zip(dsts, srcs): key = (dst.dtype, src.dtype) @@ -104,7 +114,7 @@ def _grouped_foreach_copy_(dsts: List[torch.Tensor], srcs: List[torch.Tensor]) - groups[key][0].append(dst) groups[key][1].append(src) for group_dsts, group_srcs in groups.values(): - torch._foreach_copy_(group_dsts, group_srcs) + foreach_copy(group_dsts, group_srcs) @dataclass