[diffusion] fix: fix NCCL deadlock in ulysses sp when sequence length has remainder (#24694)

Signed-off-by: storyicon <storyicon@foxmail.com>
This commit is contained in:
storyicon
2026-05-09 11:05:37 +08:00
committed by GitHub
parent 50ed01674e
commit 590b13b513
2 changed files with 7 additions and 6 deletions
@@ -478,6 +478,10 @@ class USPAttention(nn.Module):
k = _usp_input_all_to_all(k, head_dim=2)
v = _usp_input_all_to_all(v, head_dim=2)
# If NCCL timeout/deadlock occurs here, check whether
# attn_mask is inconsistent across SP ranks (None on some, Tensor on
# others), which causes all_gather participant mismatch. Upstream
# mask builders must ensure all ranks produce the same mask type.
gathered_mask = sequence_model_parallel_all_gather(
attn_mask.contiguous(), dim=1
)
@@ -1050,12 +1050,9 @@ class LTX2DenoisingStage(DenoisingStage):
if valid is None:
return None
valid = int(valid)
if valid <= 0 or valid >= int(seq_len):
return None
mask = torch.ones(
(batch_size, int(seq_len)), device=device, dtype=torch.float32
)
mask[:, valid:] = 0.0
mask = torch.ones((batch_size, int(seq_len)), device=device, dtype=torch.bool)
if valid < int(seq_len):
mask[:, max(0, valid) :] = False
return mask
@staticmethod