Make torch TP composable with torch.compile (#2352)

sgl-project · Dec 5, 2024 · d693ec0 · d693ec0
1 parent 18ea841
commit d693ec0
Showing 1 changed file with 1 addition and 5 deletions.
diff --git 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(