diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index 800d04400..5ecdb0644 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py @@ -147,6 +147,26 @@ def _count_active_replicated_modes( ) +def _prepare_sdpa_mask( + mask: torch.Tensor, *, dtype: torch.dtype, device: torch.device +) -> torch.Tensor: + mask = mask.to(device=device) + if torch.is_floating_point(mask): + mask = mask.to(dtype=dtype) + if mask.dim() == 2: + mask = mask[:, None, None, :] + elif mask.dim() == 3: + mask = mask[:, None, :, :] + return mask + + mask = mask.to(dtype=dtype) + if mask.dim() == 2: + mask = mask[:, None, None, :] + elif mask.dim() == 3: + mask = mask[:, None, :, :] + return (mask - 1.0) * torch.finfo(dtype).max + + def build_varlen_mask_meta( key_mask: torch.Tensor, ) -> dict: @@ -905,6 +925,15 @@ class USPAttention(nn.Module): out = self.attn_impl.postprocess_output(out, ctx_attn_metadata) return _usp_output_all_to_all_varlen(out, seq_lens, head_dim=2) + # Under BCG each attention break point snapshots its own kwargs, so + # per-segment mask buffers differ while this meta object is shared by + # every segment; the replicated-prefix path memoizes its gathered + # joint mask on it (see _forward_with_replicated_prefix). + joint_mask_memo_host = ( + attn_mask_meta + if isinstance(attn_mask_meta, DynamicVarlenMaskMeta) + else attn_mask + ) if isinstance(attn_mask_meta, DynamicVarlenMaskMeta): attn_mask_meta = attn_mask_meta.resolve(attn_mask) @@ -975,36 +1004,41 @@ class USPAttention(nn.Module): and not effective_skip_sp and get_sequence_parallel_world_size() > 1 ): - # 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 under sequence parallelism; drop " - "attn_mask/attn_mask_meta or the replicated segment." + # A replicated prefix with a [B, S] key mask is served by the + # replicated-prefix flow: after its all-to-all the attention is + # rank-local over the full sequence, so the mask (identical on + # every rank for the prefix, gathered for the suffix) applies + # exactly as on a single rank. + masked_prefix_supported = ( + num_replicated_prefix > 0 + and not num_replicated_suffix + and not num_replicated_kv_prefix + and attn_mask is not None + and attn_mask.dim() == 2 + and not qkv_pre_all_to_all + and get_ring_parallel_world_size() == 1 + ) + if not masked_prefix_supported: + # Sharding every row through the all-to-all would duplicate + # a replicated segment across ranks and silently corrupt + # the output, so refuse loudly. 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 under sequence parallelism; drop " + "attn_mask/attn_mask_meta or the replicated segment." + ) + return self._forward_with_replicated_prefix( + q, + k, + v, + ctx_attn_metadata, + num_replicated_prefix, + attn_mask=attn_mask, + drop_masked_query_rows=attn_mask_meta is not None, + joint_mask_memo_host=joint_mask_memo_host, ) - - def _prepare_sdpa_mask( - mask: torch.Tensor, *, dtype: torch.dtype, device: torch.device - ) -> torch.Tensor: - mask = mask.to(device=device) - if torch.is_floating_point(mask): - mask = mask.to(dtype=dtype) - if mask.dim() == 2: - mask = mask[:, None, None, :] - elif mask.dim() == 3: - mask = mask[:, None, :, :] - return mask - - mask = mask.to(dtype=dtype) - if mask.dim() == 2: - mask = mask[:, None, None, :] - elif mask.dim() == 3: - mask = mask[:, None, :, :] - return (mask - 1.0) * torch.finfo(dtype).max sp_world_size = get_sequence_parallel_world_size() if effective_skip_sp or sp_world_size == 1: @@ -1596,6 +1630,9 @@ class USPAttention(nn.Module): v: torch.Tensor, ctx_attn_metadata, num_rep: int, + attn_mask: torch.Tensor | None = None, + drop_masked_query_rows: bool = False, + joint_mask_memo_host=None, ) -> torch.Tensor: """Ulysses attention where the first *num_rep* tokens are replicated across SP ranks (e.g. text tokens) and should NOT be duplicated by the @@ -1607,6 +1644,11 @@ class USPAttention(nn.Module): 3. Locally slice the replicated prefix to the same head shard. 4. Concatenate [prefix_h_local, gathered_suffix] and run attention. 5. Split output, all-to-all back the suffix, all-gather prefix heads. + + ``attn_mask`` is a [B, num_rep + S_local] key mask in the caller's + local layout: its prefix columns are replicated (identical on every + rank) and its suffix columns are gathered to match step 2's sequence + gather, so step 4 can apply it exactly as a single rank would. """ sp_size = get_ulysses_parallel_world_size() u_rank = get_ulysses_parallel_rank() @@ -1619,6 +1661,33 @@ class USPAttention(nn.Module): k_shard = _usp_input_all_to_all(k_shard, head_dim=2) v_shard = _usp_input_all_to_all(v_shard, head_dim=2) + joint_mask = joint_mask_meta = None + if attn_mask is not None: + cache_key = ( + num_rep, + sp_size, + tuple(attn_mask.shape), + get_current_replay_token(), + ) + host = ( + joint_mask_memo_host if joint_mask_memo_host is not None else attn_mask + ) + cached = getattr(host, "_sglang_usp_joint_mask_cache", None) + if cached is not None and cached[0] == cache_key: + _, joint_mask, joint_mask_meta = cached + else: + mask_suffix = sequence_model_parallel_all_gather( + attn_mask[:, num_rep:].contiguous(), dim=1 + ) + joint_mask = torch.cat([attn_mask[:, :num_rep], mask_suffix], dim=1) + if drop_masked_query_rows: + joint_mask_meta = build_varlen_mask_meta(joint_mask) + host._sglang_usp_joint_mask_cache = ( + cache_key, + joint_mask, + joint_mask_meta, + ) + # Q and KV can have different head counts (GQA), so slice each replicated # prefix by its own per-rank head shard to match the all-to-all'd suffix. # For MHA (kv heads == q heads) this is identical to the q shard. @@ -1632,7 +1701,14 @@ class USPAttention(nn.Module): q = torch.cat([q_rep, q_shard], dim=1) out = self._replicated_kv_attention( - q, k_shard, v_shard, k_rep, v_rep, ctx_attn_metadata + q, + k_shard, + v_shard, + k_rep, + v_rep, + ctx_attn_metadata, + attn_mask=joint_mask, + attn_mask_meta=joint_mask_meta, ) out_rep = out[:, :num_rep] @@ -1660,6 +1736,8 @@ class USPAttention(nn.Module): v_rep: torch.Tensor, ctx_attn_metadata, rep_first: bool = True, + attn_mask: torch.Tensor | None = None, + attn_mask_meta: dict | None = None, ) -> torch.Tensor: """Attention of q against replicated + ring-sharded KV. @@ -1669,8 +1747,12 @@ class USPAttention(nn.Module): parallelism the sharded KV rotates around the ring while the replicated KV contributes one extra local partial, LSE-merged with the ring result (exact up to float reordering). + + ``attn_mask`` is a [B, S] key mask over the concatenated KV order and + is only supported on the non-ring path (the caller gates ring out). """ if get_ring_parallel_world_size() > 1: + assert attn_mask is None, "masked replicated KV does not support ring" out_ring, lse_ring = ring_attn( q, k_shard, v_shard, self.attn_impl, return_softmax_lse=True ) @@ -1689,8 +1771,76 @@ class USPAttention(nn.Module): ) k = torch.cat(kv_parts[0], dim=1) v = torch.cat(kv_parts[1], dim=1) + if attn_mask is not None: + assert rep_first, "the joint key mask is built in [rep, shard] order" + return self._masked_local_attention( + q, k, v, attn_mask, attn_mask_meta=attn_mask_meta + ) return self.attn_impl.forward(q, k, v, ctx_attn_metadata) + def _masked_local_attention( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + attn_mask: torch.Tensor, + attn_mask_meta: dict | None = None, + ) -> torch.Tensor: + """Rank-local attention with a [B, S] key mask over full-sequence + q/k/v (heads may already be sharded). Mirrors the single-rank masked + path: varlen FA with precomputed meta when the caller opted into + dropped-query-row semantics (masked rows zero-filled on output), SDPA + with an additive mask otherwise. + """ + if ( + _VARLEN_FA_ENABLED + and attn_mask_meta is not None + and self.backend == AttentionBackendEnum.FA + and attn_mask.dtype in (torch.bool, torch.uint8, torch.int32, torch.int64) + and q.device.type == "cuda" + and attn_mask.device == q.device + and q.dtype in (torch.float16, torch.bfloat16) + and q.shape[:2] == attn_mask.shape == k.shape[:2] == v.shape[:2] + ): + bs, seq = q.shape[0], q.shape[1] + meta = attn_mask_meta + indices = meta["indices"] + if indices.shape[0] > 0: + q_unpad, k_unpad, v_unpad = fused_pack_qkv(q, k, v, indices) + out_unpad = flash_attn_varlen_func( + q=q_unpad, + k=k_unpad, + v=v_unpad, + cu_seqlens_q=meta["cu_seqlens"], + cu_seqlens_k=meta["cu_seqlens"], + max_seqlen_q=meta["max_seqlen"], + max_seqlen_k=meta["max_seqlen"], + softmax_scale=self.softmax_scale, + causal=False, + ver=_fa_backend.fa_ver, + ) + return fused_scatter_to_padded(out_unpad, meta["inv_indices"], bs, seq) + + q_ = q.transpose(1, 2) + k_ = k.transpose(1, 2) + v_ = v.transpose(1, 2) + mask = _prepare_sdpa_mask(attn_mask, dtype=q_.dtype, device=q_.device) + sdpa_context = ( + sdpa_kernel(_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS) + if self.allow_cudnn_sdp and q_.device.type == "cuda" + else nullcontext() + ) + with sdpa_context: + return torch.nn.functional.scaled_dot_product_attention( + q_, + k_, + v_, + attn_mask=mask, + dropout_p=0.0, + is_causal=False, + scale=self.softmax_scale, + ).transpose(1, 2) + def forward_with_replicated_kv_prefix( self, q: torch.Tensor, 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 12cb5263d..e2d19b697 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 @@ -67,7 +67,8 @@ class TestUSPAttentionReplicatedPrefix(unittest.TestCase): with ( patch(f"{_LAYER}.get_ulysses_parallel_world_size", return_value=_SP), - patch(f"{_LAYER}.get_sp_parallel_rank", return_value=sp_rank), + patch(f"{_LAYER}.get_ulysses_parallel_rank", return_value=sp_rank), + patch(f"{_LAYER}.get_ring_parallel_world_size", return_value=1), patch( f"{_LAYER}._usp_input_all_to_all", side_effect=_fake_input_all_to_all ), @@ -153,3 +154,155 @@ class TestUSPAttentionMaskedReplicatedGuard(unittest.TestCase): ): out = obj.forward(q, q, q, attn_mask=mask, num_replicated_prefix=2) self.assertEqual(out.shape, q.shape) + + +class TestUSPAttentionMaskedReplicatedPrefix(unittest.TestCase): + """A [B, S] key mask with a replicated prefix now rides the + replicated-prefix flow under ulysses SP instead of raising: after the + all-to-all the attention is rank-local over the full sequence, so the + output must match a plain full-sequence masked SDPA reference.""" + + _H, _D = 4, 8 + _NUM_REP, _S_LOCAL, _S_OTHER = 3, 4, 4 + + def _run_rank0(self, q_full, k_full, v_full, text_mask, other_suffix_mask): + + num_rep, s_local = self._NUM_REP, self._S_LOCAL + h_local = self._H // _SP + obj = USPAttention.__new__(USPAttention) + obj.attn_impl = _CaptureAttn() + obj.skip_sequence_parallel = False + obj.sp_attention_mode = "ulysses" + obj.sp_attention_mode_is_auto = True + obj.backend = None # not FA -> the masked branch takes SDPA + obj.softmax_scale = self._D**-0.5 + obj.allow_cudnn_sdp = False + obj.enable_packed_qkv_input_a2a = False + + local = slice(0, num_rep + s_local) + q, k, v = q_full[:, local], k_full[:, local], v_full[:, local] + other = { + "suffix": lambda full: full[:, num_rep + s_local :], + } + gathered = iter( + [ + torch.cat([q[:, num_rep:], other["suffix"](q_full)], dim=1), + torch.cat([k[:, num_rep:], other["suffix"](k_full)], dim=1), + torch.cat([v[:, num_rep:], other["suffix"](v_full)], dim=1), + ] + ) + + def fake_input_a2a(x, **_): + # rank0 view: gather sequence across ranks, keep the first head shard. + return next(gathered)[:, :, :h_local, :].contiguous() + + def fake_output_a2a(x, **_): + # inverse: keep rank0's sequence shard, duplicate the head shard. + return x[:, :s_local].repeat(1, 1, _SP, 1) + + def fake_mask_gather(m, **_): + return torch.cat([m, other_suffix_mask], dim=1) + + def fake_all_gather(out_list, tensor, **_): + for t in out_list: + t.copy_(tensor) + + sp_group = MagicMock() + sp_group.ulysses_group = None + local_mask = torch.cat( + [text_mask, torch.ones(1, s_local, dtype=torch.bool)], dim=1 + ) + with ( + patch( + f"{_LAYER}.get_forward_context", + return_value=MagicMock(attn_metadata=None), + ), + patch(f"{_LAYER}.get_sequence_parallel_world_size", return_value=_SP), + patch(f"{_LAYER}.get_ulysses_parallel_world_size", return_value=_SP), + patch(f"{_LAYER}.get_ring_parallel_world_size", return_value=1), + patch(f"{_LAYER}.get_ulysses_parallel_rank", return_value=0), + patch(f"{_LAYER}._usp_input_all_to_all", side_effect=fake_input_a2a), + patch(f"{_LAYER}._usp_output_all_to_all", side_effect=fake_output_a2a), + patch( + f"{_LAYER}.sequence_model_parallel_all_gather", + side_effect=fake_mask_gather, + ), + patch(f"{_LAYER}.get_sp_group", return_value=sp_group), + patch("torch.distributed.all_gather", side_effect=fake_all_gather), + ): + return obj.forward( + q, k, v, attn_mask=local_mask, num_replicated_prefix=num_rep + ) + + def test_matches_full_sequence_masked_sdpa(self): + torch.manual_seed(0) + seq = self._NUM_REP + self._S_LOCAL + self._S_OTHER + q = torch.randn(1, seq, self._H, self._D) + k = torch.randn(1, seq, self._H, self._D) + v = torch.randn(1, seq, self._H, self._D) + text_mask = torch.tensor([[True, True, False]]) + other_suffix_mask = torch.ones(1, self._S_OTHER, dtype=torch.bool) + + out = self._run_rank0(q, k, v, text_mask, other_suffix_mask) + + full_mask = torch.cat( + [text_mask, torch.ones(1, self._S_LOCAL + self._S_OTHER, dtype=torch.bool)], + dim=1, + ) + ref = torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2), + k.transpose(1, 2), + v.transpose(1, 2), + attn_mask=full_mask[:, None, None, :], + scale=self._D**-0.5, + ).transpose(1, 2) + + h_local = self._H // _SP + rows = self._NUM_REP + self._S_LOCAL + self.assertEqual(tuple(out.shape), (1, rows, self._H, self._D)) + torch.testing.assert_close( + out[:, :, :h_local, :], + ref[:, :rows, :h_local, :], + rtol=1e-5, + atol=1e-5, + ) + + def test_masked_prefix_still_rejects_ring_and_3d_masks(self): + obj = USPAttention.__new__(USPAttention) + obj.attn_impl = _CaptureAttn() + obj.skip_sequence_parallel = False + obj.sp_attention_mode = "ulysses" + obj.sp_attention_mode_is_auto = True + q = torch.randn(1, 6, 2, 4) + with ( + patch( + f"{_LAYER}.get_forward_context", + return_value=MagicMock(attn_metadata=None), + ), + patch(f"{_LAYER}.get_sequence_parallel_world_size", return_value=2), + patch(f"{_LAYER}.get_ring_parallel_world_size", return_value=2), + ): + with self.assertRaisesRegex(NotImplementedError, "replicated"): + obj.forward( + q, + q, + q, + attn_mask=torch.ones(1, 6, dtype=torch.bool), + num_replicated_prefix=2, + ) + with ( + patch( + f"{_LAYER}.get_forward_context", + return_value=MagicMock(attn_metadata=None), + ), + patch(f"{_LAYER}.get_sequence_parallel_world_size", return_value=2), + patch(f"{_LAYER}.get_ring_parallel_world_size", return_value=1), + ): + with self.assertRaisesRegex(NotImplementedError, "replicated"): + obj.forward( + q, + q, + q, + attn_mask=torch.ones(1, 6, 6, dtype=torch.bool), + num_replicated_prefix=2, + )