[diffusion] fix: scope the masked-path replicated guard to sp runs (#33953)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

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

@@ -850,17 +850,24 @@ class USPAttention(nn.Module):
 
         if attn_mask is not None or meta_only_pad:
             if (
-                num_replicated_prefix
+                (
-                or num_replicated_suffix
+                    num_replicated_prefix
-                or num_replicated_kv_prefix
+                    or num_replicated_suffix
+                    or num_replicated_kv_prefix
+                )
+                and not effective_skip_sp
+                and get_sequence_parallel_world_size() > 1
             ):
-                # This path shards every row through the all-to-all; a
+                # Under SP this path shards every row through the all-to-all;
-                # replicated prefix/suffix would be duplicated across ranks and
+                # a replicated prefix/suffix would be duplicated across ranks
-                # silently corrupt the output, so refuse loudly instead.
+                # and silently corrupt the output, so refuse loudly instead.
+                # On a single rank the mask already describes the full
+                # sequence and the replicated counts are meaningless, so the
+                # call is legal.
                 raise NotImplementedError(
                     "USPAttention's masked path does not support replicated "
-                    "prefix/suffix tokens; drop attn_mask/attn_mask_meta or "
+                    "prefix/suffix tokens under sequence parallelism; drop "
-                    "the replicated segment."
+                    "attn_mask/attn_mask_meta or the replicated segment."
                 )
 
             def _prepare_sdpa_mask(

python/sglang/multimodal_gen/test/unit/test_usp_attention_replicated_prefix.py

@@ -17,6 +17,7 @@ from unittest.mock import MagicMock, patch
 import torch
 
 from sglang.multimodal_gen.runtime.layers.attention.layer import USPAttention
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
 
 _LAYER = "sglang.multimodal_gen.runtime.layers.attention.layer"
 _SP = 2
@@ -127,3 +128,28 @@ class TestUSPAttentionMaskedReplicatedGuard(unittest.TestCase):
         ):
             with self.assertRaisesRegex(NotImplementedError, "replicated"):
                 obj.forward(q, q, q, attn_mask=mask, num_replicated_suffix=2)
+
+    def test_single_rank_masked_call_keeps_replicated_args(self):
+        # Without SP the mask describes the full sequence; the replicated
+        # counts are meaningless and must not be refused.
+        obj = USPAttention.__new__(USPAttention)
+        obj.attn_impl = _CaptureAttn()
+        obj.skip_sequence_parallel = False
+        obj.sp_attention_mode = "ulysses"
+        obj.sp_attention_mode_is_auto = False
+        obj.allow_cudnn_sdp = False
+        obj.softmax_scale = 0.5
+        obj.backend = AttentionBackendEnum.TORCH_SDPA
+        obj.causal = False
+        obj.dropout_p = 0.0
+        q = torch.randn(1, 6, 2, 4)
+        mask = torch.ones(1, 6, dtype=torch.bool)
+        with (
+            patch(
+                f"{_LAYER}.get_forward_context",
+                return_value=MagicMock(attn_metadata=None),
+            ),
+            patch(f"{_LAYER}.get_sequence_parallel_world_size", return_value=1),
+        ):
+            out = obj.forward(q, q, q, attn_mask=mask, num_replicated_prefix=2)
+        self.assertEqual(out.shape, q.shape)