From 35ba6fe19eca20e159d25762b296332d635396dd Mon Sep 17 00:00:00 2001 From: WenhaoZhang <42087078+niehen6174@users.noreply.github.com> Date: Fri, 5 Dec 2025 16:30:10 +0800 Subject: [PATCH] [diffusion] fix: fix CLIP text encoder attention mask not used (#14364) Co-authored-by: niehen6174 Co-authored-by: Mick --- .../configs/models/encoders/clip.py | 6 ++ .../runtime/models/encoders/clip.py | 61 ++++++++++++++++--- 2 files changed, 57 insertions(+), 10 deletions(-) diff --git a/python/sglang/multimodal_gen/configs/models/encoders/clip.py b/python/sglang/multimodal_gen/configs/models/encoders/clip.py index 6b36fc88b..ff9a90b32 100644 --- a/python/sglang/multimodal_gen/configs/models/encoders/clip.py +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/clip.py b/python/sglang/multimodal_gen/runtime/models/encoders/clip.py index ec80e387f..99db53a75 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/clip.py +++ b/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]