[diffusion] fix: fix CLIP text encoder attention mask not used (#14364)

Co-authored-by: niehen6174 <niehen.6174@gmail.com>
Co-authored-by: Mick <mickjagger19@icloud.com>

python/sglang/multimodal_gen/configs/models/encoders/clip.py

@@ -9,6 +9,7 @@ from sglang.multimodal_gen.configs.models.encoders.base import (
     TextEncoderArchConfig,
     TextEncoderConfig,
 )
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
 
 
 def _is_transformer_layer(n: str, m) -> bool:
@@ -38,6 +39,11 @@ class CLIPTextArchConfig(TextEncoderArchConfig):
     bos_token_id: int = 49406
     eos_token_id: int = 49407
     text_len: int = 77
+    _supported_attention_backends: set[AttentionBackendEnum] = field(
+        default_factory=lambda: {
+            AttentionBackendEnum.TORCH_SDPA,  # Force TORCH_SDPA to support attention_mask
+        }
+    )
     stacked_params_mapping: list[tuple[str, str, str]] = field(
         default_factory=lambda: [
             # (param_name, shard_name, shard_id)

python/sglang/multimodal_gen/runtime/models/encoders/clip.py

@@ -33,6 +33,7 @@ from sglang.multimodal_gen.runtime.models.encoders.base import ImageEncoder, Tex
 from sglang.multimodal_gen.runtime.models.encoders.vision import (
     resolve_visual_encoder_outputs,
 )
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)
@@ -181,7 +182,7 @@ class CLIPAttention(nn.Module):
             self.head_dim,
             self.num_heads_per_partition,
             softmax_scale=self.scale,
-            causal=False,
+            causal=True,
             supported_attention_backends=config._supported_attention_backends,
         )
 
@@ -195,6 +196,7 @@ class CLIPAttention(nn.Module):
     def forward(
         self,
         hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
     ):
         """Input shape: Batch x Time x Channel"""
 
@@ -219,7 +221,36 @@ class CLIPAttention(nn.Module):
             self.num_heads_per_partition,
             self.head_dim,
         )
-        attn_output = self.attn(query_states, key_states, value_states)
+
+        if self.attn.backend == AttentionBackendEnum.TORCH_SDPA:
+            query_states = query_states.transpose(1, 2)  # [B, H, S, D]
+            key_states = key_states.transpose(1, 2)
+            value_states = value_states.transpose(1, 2)
+
+            if attention_mask is not None:
+                # SDPA requires [B, 1, 1, S] or [B, S, S] format mask
+                if attention_mask.dim() == 2:
+                    attn_mask = attention_mask[:, None, None, :].to(
+                        dtype=query_states.dtype
+                    )
+                    attn_mask = (1.0 - attn_mask) * torch.finfo(query_states.dtype).min
+                else:
+                    attn_mask = attention_mask
+            else:
+                attn_mask = None
+
+            attn_output = torch.nn.functional.scaled_dot_product_attention(
+                query_states,
+                key_states,
+                value_states,
+                attn_mask=attn_mask,
+                is_causal=True,
+                scale=self.scale,
+            )
+            attn_output = attn_output.transpose(1, 2)
+        else:
+            # Use LocalAttention (doesn't support attention_mask, but maintains compatibility)
+            attn_output = self.attn(query_states, key_states, value_states)
 
         attn_output = attn_output.reshape(
             attn_output.shape[0],
@@ -283,11 +314,18 @@ class CLIPEncoderLayer(nn.Module):
         self.mlp = CLIPMLP(config, quant_config=quant_config, prefix=f"{prefix}.mlp")
         self.layer_norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
 
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
         residual = hidden_states
 
         hidden_states = self.layer_norm1(hidden_states)
-        hidden_states, _ = self.self_attn(hidden_states=hidden_states)
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+        )
         hidden_states = residual + hidden_states
 
         residual = hidden_states
@@ -334,13 +372,19 @@ class CLIPEncoder(nn.Module):
         )
 
     def forward(
-        self, inputs_embeds: torch.Tensor, return_all_hidden_states: bool
+        self,
+        inputs_embeds: torch.Tensor,
+        return_all_hidden_states: bool,
+        attention_mask: torch.Tensor | None = None,
     ) -> torch.Tensor | list[torch.Tensor]:
         hidden_states_pool = [inputs_embeds]
         hidden_states = inputs_embeds
 
         for idx, encoder_layer in enumerate(self.layers):
-            hidden_states = encoder_layer(hidden_states)
+            hidden_states = encoder_layer(
+                hidden_states,
+                attention_mask=attention_mask,
+            )
             if return_all_hidden_states:
                 hidden_states_pool.append(hidden_states)
         # If we have multiple feature sample layers, we return all hidden
@@ -417,11 +461,8 @@ class CLIPTextTransformer(nn.Module):
 
         encoder_outputs = self.encoder(
             inputs_embeds=hidden_states,
-            # attention_mask=attention_mask,
-            # causal_attention_mask=causal_attention_mask,
-            # output_attentions=output_attentions,
             return_all_hidden_states=output_hidden_states,
-            # return_dict=return_dict,
+            attention_mask=attention_mask,
         )
 
         last_hidden_state = encoder_outputs[-1]