From d693ec0427bd70c8676316c634e00bd27514b7ec Mon Sep 17 00:00:00 2001 From: Ke Wen Date: Wed, 4 Dec 2024 17:26:00 -0800 Subject: [PATCH] Make torch TP composable with torch.compile (#2352) --- 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(