diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index d88df2eee..9641ee820 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/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 - or num_replicated_kv_prefix + ( + num_replicated_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 - # replicated prefix/suffix would be duplicated across ranks and - # silently corrupt the output, so refuse loudly instead. + # Under SP this path shards every row through the all-to-all; + # a replicated prefix/suffix would be duplicated across ranks + # 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 " - "the replicated segment." + "prefix/suffix tokens under sequence parallelism; drop " + "attn_mask/attn_mask_meta or the replicated segment." ) def _prepare_sdpa_mask( diff --git a/python/sglang/multimodal_gen/test/unit/test_usp_attention_replicated_prefix.py b/python/sglang/multimodal_gen/test/unit/test_usp_attention_replicated_prefix.py index 40ce07404..12cb5263d 100644 --- a/python/sglang/multimodal_gen/test/unit/test_usp_attention_replicated_prefix.py +++ b/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)