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

Signed-off-by: storyicon <storyicon@foxmail.com>

python/sglang/multimodal_gen/runtime/layers/attention/layer.py

@@ -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
             )

python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py

@@ -1050,12 +1050,9 @@ class LTX2DenoisingStage(DenoisingStage):
         if valid is None:
             return None
         valid = int(valid)
-        if valid <= 0 or valid >= int(seq_len):
+        mask = torch.ones((batch_size, int(seq_len)), device=device, dtype=torch.bool)
-            return None
+        if valid < int(seq_len):
-        mask = torch.ones(
+            mask[:, max(0, valid) :] = False
-            (batch_size, int(seq_len)), device=device, dtype=torch.float32
-        )
-        mask[:, valid:] = 0.0
         return mask
 
     @staticmethod