From 54fa11ce521095b04055f2ea3c5489ee8138a23c Mon Sep 17 00:00:00 2001 From: Ke Wen Date: Tue, 3 Dec 2024 14:22:21 -0800 Subject: [PATCH] Make TP composable with torch.compile --- python/sglang/srt/model_parallel.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/python/sglang/srt/model_parallel.py b/python/sglang/srt/model_parallel.py index 6817bce0235..778347b8ef3 100644 --- a/python/sglang/srt/model_parallel.py +++ b/python/sglang/srt/model_parallel.py @@ -54,11 +54,7 @@ def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_me )._prepare_output_fn( output_layouts, use_local_output, mod, outputs, device_mesh ) - # wait for the output to be ready - if isinstance(outputs, AsyncCollectiveTensor): - return outputs.wait() - else: - return outputs + return torch.distributed._functional_collectives.wait_tensor(outputs) def tensor_parallel(