sgl-project/sglang · error · NotImplementedError

USPAttention's masked path does not support replicated prefi

Error message

USPAttention's masked path does not support replicated prefix/suffix tokens under sequence parallelism; drop attn_mask/attn_mask_meta or the replicated segment.

What it means

The masked path of USPAttention shards every row through the SP all-to-all. Replicated prefix/suffix tokens (e.g. shared text around an image) would be duplicated across ranks and silently corrupt the output, so under sequence parallelism the combination of an explicit attn_mask/attn_mask_meta and replicated tokens is refused. On a single rank it is legal.

Source

Thrown at python/sglang/multimodal_gen/runtime/layers/attention/layer.py:953

                raise NotImplementedError(unsupported)

        if attn_mask is not None or meta_only_pad:
            if (
                (
                    num_replicated_prefix
                    or num_replicated_suffix
                    or num_replicated_kv_prefix
                )
                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."
                )

            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)

View on GitHub (pinned to 0132848349)

Solutions

Drop the attn_mask/attn_mask_meta and rely on the layer's native prefix handling
Or remove the replicated prefix/suffix segment from the batch construction (let every rank own its tokens)
Or run with sequence parallel world size 1 if the masked+replicated combination is required

Example fix

# before
out = attn(q, k, v, attn_mask=mask, replicated_q=rq, replicated_k=rk, replicated_v=rv)
# after
out = attn(q, k, v, replicated_q=rq, replicated_k=rk, replicated_v=rv)  # no mask

Defensive patterns

Strategy: validation

Validate before calling

has_replicated = any(x is not None for x in (replicated_q, replicated_k, replicated_v))
has_mask = attn_mask is not None or attn_mask_meta is not None
if has_replicated and has_mask and get_sequence_parallel_world_size() > 1:
    raise ValueError("masked path + replicated tokens unsupported under SP; drop one")

Type guard

def masked_replicated_sp_ok(attn_mask, attn_mask_meta, rq, rk, rv, sp_ws: int) -> bool:
    if sp_ws == 1:
        return True
    has_rep = any(x is not None for x in (rq, rk, rv))
    return not (has_rep and (attn_mask is not None or attn_mask_meta is not None))

Prevention

Encode prefix masking via the layer's native prefix args, not explicit masks, under SP
Run the single-rank vs multi-rank output equality test when adding replicated segments

When it happens

Trigger: Calling USPAttention.forward with both an attn_mask or attn_mask_meta AND replicated prefix/suffix tokens, while sequence parallel world size > 1 (and not on the supported ring tail-pad path).

Common situations: Multimodal batches with a shared text prompt plus an explicit attention mask under SP; migrating a single-GPU masked attention path to multi-GPU SP without dropping the replicated segment; adding attn_mask_meta for prefix masking while keeping replicated_q/k/v arguments populated.

Related errors

AI-assisted analysis of sgl-project/sglang@0132848349 (2026-08-28). Data as JSON: /api/errors/4bb7de3f7e7cf936. Report an issue: GitHub.