From 1be3163011dc14a4e688f1072aea2a7b96f5f5d8 Mon Sep 17 00:00:00 2001 From: Mick Date: Tue, 5 May 2026 00:08:48 +0800 Subject: [PATCH] [diffusion] fix: use direct all-to-all for USP collectives (#24366) --- python/sglang/multimodal_gen/runtime/layers/usp.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/layers/usp.py b/python/sglang/multimodal_gen/runtime/layers/usp.py index 3794605fd..7053a7de3 100644 --- a/python/sglang/multimodal_gen/runtime/layers/usp.py +++ b/python/sglang/multimodal_gen/runtime/layers/usp.py @@ -37,13 +37,12 @@ def _usp_all_to_all_single(x: torch.Tensor) -> torch.Tensor: ulysses_pg = get_sp_group().ulysses_group assert ulysses_pg is not None, "Ulysses process group is not initialized." x_shape = x.shape - x = x.flatten() - x = ft_c.all_to_all_single( - x, output_split_sizes=None, input_split_sizes=None, group=ulysses_pg - ) - x = _maybe_wait(x) - x = x.reshape(x_shape) - return x + x = x.flatten().contiguous() + output = torch.empty_like(x) + # USP calls this collective many times per denoising step and waits + # immediately, so avoid the extra wrapper overhead of functional collectives. + torch.distributed.all_to_all_single(output, x, group=ulysses_pg) + return output.reshape(x_shape) def _usp_input_all_to_all(x: torch.Tensor, head_dim: int = 1) -> torch.Tensor: