{"record":{"id":"b98825eea4290c5c","repo":"sgl-project/sglang","slug":"expected-a-2d-3d-or-4d-attention-mask-got-atte","errorCode":null,"errorMessage":"Expected a 2D, 3D, or 4D attention mask, got {attention_mask.ndim}D.","messagePattern":"Expected a 2D, 3D, or 4D attention mask, got (.+?)D\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/multimodal_gen/runtime/models/dits/stable_diffusion.py","lineNumber":169,"sourceCode":"        inner_dim = num_heads * head_dim\n        context_dim = cross_attention_dim or query_dim\n        self.heads = num_heads\n        self.head_dim = head_dim\n        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)\n        self.to_k = nn.Linear(context_dim, inner_dim, bias=False)\n        self.to_v = nn.Linear(context_dim, inner_dim, bias=False)\n        self.to_out = nn.ModuleList([nn.Linear(inner_dim, query_dim), nn.Dropout(0.0)])\n\n    def _prepare_mask(self, attention_mask: torch.Tensor | None) -> torch.Tensor | None:\n        if attention_mask is None:\n            return None\n        if attention_mask.ndim == 2:\n            return attention_mask[:, None, None, :]\n        if attention_mask.ndim == 3:\n            return attention_mask[:, None, :, :]\n        if attention_mask.ndim == 4:\n            return attention_mask\n        raise ValueError(\n            f\"Expected a 2D, 3D, or 4D attention mask, got {attention_mask.ndim}D.\"\n        )\n\n    def forward(\n        self,\n        hidden_states: torch.Tensor,\n        encoder_hidden_states: torch.Tensor | None = None,\n        attention_mask: torch.Tensor | None = None,\n    ) -> torch.Tensor:\n        context = (\n            hidden_states if encoder_hidden_states is None else encoder_hidden_states\n        )\n        batch_size = hidden_states.shape[0]\n        query = self.to_q(hidden_states).view(batch_size, -1, self.heads, self.head_dim)\n        key = self.to_k(context).view(batch_size, -1, self.heads, self.head_dim)\n        value = self.to_v(context).view(batch_size, -1, self.heads, self.head_dim)\n        output = F.scaled_dot_product_attention(\n            query.transpose(1, 2),","sourceCodeStart":151,"sourceCodeEnd":187,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/multimodal_gen/runtime/models/dits/stable_diffusion.py#L151-L187","documentation":"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.","triggerScenarios":"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).","commonSituations":"Passing a boolean token-validity vector instead of an attention mask; passing a mask that was unsqueezed one time too many upstream.","solutions":["Reshape the mask to [B, S] (2D), [B, S, S] (3D), or [B, H, S, S] (4D) before forward","If you have a 1D per-token mask, expand it: mask[None, None, None, :] to 4D"],"exampleFix":"# before\nattn_mask = valid_token_bools  # shape [S]\n# after\nattn_mask = valid_token_bools[None, None, None, :]  # [1,1,1,S]","handlingStrategy":"type-guard","validationCode":"assert 2 <= attention_mask.ndim <= 4, f\"bad mask ndim {attention_mask.ndim}\"","typeGuard":"def is_supported_mask(m: torch.Tensor) -> bool:\n    return m.ndim in (2, 3, 4)","tryCatchPattern":null,"preventionTips":["Normalize masks to 4D [B,1,1,S] right after creation","Add ndim asserts in pipeline glue code"],"tags":["attention-mask","shape-validation","stable-diffusion"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}