sgl-project/sglang · error · ValueError

Expected a 2D, 3D, or 4D attention mask, got {attention_mask

Error message

Expected a 2D, 3D, or 4D attention mask, got {attention_mask.ndim}D.

What it means

Raised by _prepare_mask when the incoming attention_mask tensor has fewer than 2 or more than 4 dimensions. The UNet only knows how to broadcast 2D [B, S], 3D [B, S, S], and 4D [B, H, S, S] masks into attention bias layout.

Source

Thrown at python/sglang/multimodal_gen/runtime/models/dits/stable_diffusion.py:169

        inner_dim = num_heads * head_dim
        context_dim = cross_attention_dim or query_dim
        self.heads = num_heads
        self.head_dim = head_dim
        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
        self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
        self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
        self.to_out = nn.ModuleList([nn.Linear(inner_dim, query_dim), nn.Dropout(0.0)])

    def _prepare_mask(self, attention_mask: torch.Tensor | None) -> torch.Tensor | None:
        if attention_mask is None:
            return None
        if attention_mask.ndim == 2:
            return attention_mask[:, None, None, :]
        if attention_mask.ndim == 3:
            return attention_mask[:, None, :, :]
        if attention_mask.ndim == 4:
            return attention_mask
        raise ValueError(
            f"Expected a 2D, 3D, or 4D attention mask, got {attention_mask.ndim}D."
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor | None = None,
        attention_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        context = (
            hidden_states if encoder_hidden_states is None else encoder_hidden_states
        )
        batch_size = hidden_states.shape[0]
        query = self.to_q(hidden_states).view(batch_size, -1, self.heads, self.head_dim)
        key = self.to_k(context).view(batch_size, -1, self.heads, self.head_dim)
        value = self.to_v(context).view(batch_size, -1, self.heads, self.head_dim)
        output = F.scaled_dot_product_attention(
            query.transpose(1, 2),

View on GitHub (pinned to 0132848349)

Solutions

1. Reshape the mask to [B, S] (2D), [B, S, S] (3D), or [B, H, S, S] (4D) before forward
2. If you have a 1D per-token mask, expand it: mask[None, None, None, :] to 4D

Example fix

# before
attn_mask = valid_token_bools  # shape [S]
# after
attn_mask = valid_token_bools[None, None, None, :]  # [1,1,1,S]

assert 2 <= attention_mask.ndim <= 4, f"bad mask ndim {attention_mask.ndim}"

def is_supported_mask(m: torch.Tensor) -> bool:
    return m.ndim in (2, 3, 4)

When it happens

Trigger: Calling the SD2 transformer block forward with an attention_mask of ndim 1 (a flat per-token mask) or ndim >= 5 (e.g. an already-expanded or wrongly stacked mask).

Common situations: Passing a boolean token-validity vector instead of an attention mask; passing a mask that was unsqueezed one time too many upstream.

Related errors

• The pointers must be multiple of 16 bytes.
• The last dimension ({input.shape[-1]}) x itemsize ({input.dt
• rope_pool_fused expects q/k/v to be 3-D
• rope_pool_fused expects positions/slots to be 1-D
• rope_pool_fused expects pool tensors to be 3-D

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