diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/ernie_image.py b/python/sglang/multimodal_gen/configs/pipeline_configs/ernie_image.py index 1017f5dc7..4f4315f72 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/ernie_image.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/ernie_image.py @@ -12,6 +12,7 @@ from sglang.multimodal_gen.configs.models.vaes.ernie_image import ErnieImageVAEC from sglang.multimodal_gen.configs.pipeline_configs.base import ( ImagePipelineConfig, ModelTaskType, + pad_text_embeddings_with_mask, shard_rotary_emb_for_sp, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger @@ -20,8 +21,19 @@ logger = init_logger(__name__) def ernie_image_postprocess_text(outputs, _text_inputs, hidden_layer_index=-2): + """Return Ernie-Image text embeddings, re-padded from real token spans. + + Batched requests can have different real caption lengths after + tokenization; extract each request's real (unpadded) span via the + tokenizer's attention mask and re-pad, so TextConditioningOutput carries + the true per-request lengths instead of the tokenizer's padded length. + """ hidden_states = outputs.hidden_states[hidden_layer_index] - return hidden_states + prompt_mask = _text_inputs.attention_mask.to(hidden_states.device).bool() + split_hidden_states = [ + hidden_states[idx][prompt_mask[idx]] for idx in range(hidden_states.shape[0]) + ] + return pad_text_embeddings_with_mask(split_hidden_states) def _patchify_latents(latents: torch.Tensor) -> torch.Tensor: @@ -72,7 +84,12 @@ class ErnieImagePipelineConfig(ImagePipelineConfig): text_encoder_extra_args: list[dict] = field( default_factory=lambda: [ dict( - padding=False, + # "longest" (not False): dynamic batches can merge requests with + # different real caption lengths into one tokenize() call, and a + # ragged, un-padded batch can't be stacked into a tensor. The real + # per-request lengths are recovered from the attention mask in + # ernie_image_postprocess_text. + padding="longest", truncation=True, max_length=None, add_special_tokens=True, @@ -143,8 +160,33 @@ class ErnieImagePipelineConfig(ImagePipelineConfig): sin = freqs.imag.to(dtype=torch.float32).contiguous() return torch.cat([cos, sin], dim=-1) - def _prepare_cond_kwargs(self, batch, prompt_embeds, rotary_emb, device, dtype): + def _prepare_encoder_hidden_states_mask( + self, + batch, + txt_seq_lens: list[int], + text_seq_len: int, + device, + ): + """Return a `[batch, text_seq_len]` mask over real (non-padded) text tokens. + + Dynamic batches can merge requests whose captions have different real + lengths after tokenization; the DiT still sees one padded + `encoder_hidden_states` tensor of shape `[batch, text_seq_len, dim]`, so + we need a mask to keep attention off the padding. Returns None when every + request already fills the full padded length (no mask needed). + """ + if all(seq_len == text_seq_len for seq_len in txt_seq_lens): + return None + + positions = torch.arange(text_seq_len, device=device) + seq_lens = torch.tensor(txt_seq_lens, device=device, dtype=torch.long) + return positions.unsqueeze(0) < seq_lens.unsqueeze(1) + + def _prepare_cond_kwargs( + self, batch, prompt_embeds, rotary_emb, device, dtype, *, negative: bool = False + ): batch_size = prompt_embeds[0].shape[0] + text_seq_len = prompt_embeds[0].shape[1] height = batch.height width = batch.width vae_scale_factor = self.get_vae_scale_factor() @@ -158,13 +200,19 @@ class ErnieImagePipelineConfig(ImagePipelineConfig): ) ] ] * batch_size - txt_seq_lens = [prompt_embeds[0].shape[1]] + txt_seq_lens = self.require_text_seq_lens( + batch, 0, negative=negative, expected_batch_size=batch_size + ) + encoder_hidden_states_mask = self._prepare_encoder_hidden_states_mask( + batch, txt_seq_lens, text_seq_len, device + ) if rotary_emb is None: return { "img_shapes": img_shapes, "txt_seq_lens": txt_seq_lens, "freqs_cis": None, + "encoder_hidden_states_mask": encoder_hidden_states_mask, } freqs_cis = self.get_freqs_cis( @@ -180,16 +228,22 @@ class ErnieImagePipelineConfig(ImagePipelineConfig): "txt_seq_lens": txt_seq_lens, "freqs_cis": freqs_cis, "img_shapes": img_shapes, + "encoder_hidden_states_mask": encoder_hidden_states_mask, } def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype): return self._prepare_cond_kwargs( - batch, batch.prompt_embeds, rotary_emb, device, dtype + batch, batch.prompt_embeds, rotary_emb, device, dtype, negative=False ) def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype): return self._prepare_cond_kwargs( - batch, batch.negative_prompt_embeds, rotary_emb, device, dtype + batch, + batch.negative_prompt_embeds, + rotary_emb, + device, + dtype, + negative=True, ) def _check_vae_has_bn(self, vae): diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py b/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py index 5881220bf..e12bd874d 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py @@ -25,7 +25,10 @@ from sglang.multimodal_gen.configs.models.dits.ernie_image import ( from sglang.multimodal_gen.runtime.distributed import ( get_tp_world_size, ) -from sglang.multimodal_gen.runtime.layers.attention.layer import USPAttention +from sglang.multimodal_gen.runtime.layers.attention.layer import ( + USPAttention, + build_varlen_mask_meta, +) from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm, apply_qk_norm from sglang.multimodal_gen.runtime.layers.linear import ( ColumnParallelLinear, @@ -148,6 +151,8 @@ class ErnieImageSelfAttention(nn.Module): self, x: torch.Tensor, rotary_pos_emb: torch.Tensor, + attn_mask: torch.Tensor | None = None, + attn_mask_meta: dict | None = None, ) -> torch.Tensor: B, S, H = x.shape @@ -171,7 +176,9 @@ class ErnieImageSelfAttention(nn.Module): q = _apply_rotary_bshd(q, rotary_pos_emb) k = _apply_rotary_bshd(k, rotary_pos_emb) - attn_out = self.attn(q, k, v) + attn_out = self.attn( + q, k, v, attn_mask=attn_mask, attn_mask_meta=attn_mask_meta + ) attn_out = attn_out.reshape(B, S, self.num_local_heads * self.head_dim) out, _ = self.to_out[0](attn_out) return out @@ -244,10 +251,14 @@ class ErnieImageSharedAdaLNBlock(nn.Module): shift_mlp: torch.Tensor, scale_mlp: torch.Tensor, gate_mlp: torch.Tensor, + attn_mask: torch.Tensor | None = None, + attn_mask_meta: dict | None = None, ) -> torch.Tensor: residual = x x = self.adaLN_sa_ln(x) * (1 + scale_msa) + shift_msa - x = residual + gate_msa * self.self_attention(x, rotary_pos_emb) + x = residual + gate_msa * self.self_attention( + x, rotary_pos_emb, attn_mask=attn_mask, attn_mask_meta=attn_mask_meta + ) residual = x x = self.adaLN_mlp_ln(x) * (1 + scale_mlp) + shift_mlp @@ -382,6 +393,7 @@ class ErnieImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin) timestep: torch.LongTensor, encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, guidance=None, + encoder_hidden_states_mask: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: """ @@ -445,6 +457,18 @@ class ErnieImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin) all_ids = torch.cat([image_ids, text_ids], dim=1) rotary_pos_emb = self.pos_embed(all_ids) + attn_mask = attn_mask_meta = None + if encoder_hidden_states_mask is not None: + image_mask = torch.ones((B, N_img), dtype=torch.bool, device=device) + attn_mask = torch.cat( + [ + image_mask, + encoder_hidden_states_mask.to(device=device, dtype=torch.bool), + ], + dim=1, + ) + attn_mask_meta = build_varlen_mask_meta(attn_mask) + t_emb = self.time_proj(timestep.to(dtype)) c = self.time_embedding(t_emb.to(dtype=dtype)) @@ -463,6 +487,8 @@ class ErnieImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin) shift_mlp, scale_mlp, gate_mlp, + attn_mask=attn_mask, + attn_mask_meta=attn_mask_meta, ) scale, shift = self.final_norm["linear"](c).chunk(2, dim=-1) diff --git a/python/sglang/multimodal_gen/test/unit/test_ernie_image_pipeline_config.py b/python/sglang/multimodal_gen/test/unit/test_ernie_image_pipeline_config.py new file mode 100644 index 000000000..06c027c60 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_ernie_image_pipeline_config.py @@ -0,0 +1,85 @@ +import unittest +from types import SimpleNamespace + +import torch + +from sglang.multimodal_gen.configs.models.encoders.base import BaseEncoderOutput +from sglang.multimodal_gen.configs.pipeline_configs.base import TextConditioningOutput +from sglang.multimodal_gen.configs.pipeline_configs.ernie_image import ( + ErnieImagePipelineConfig, + ernie_image_postprocess_text, +) + + +class TestErnieImagePostprocessText(unittest.TestCase): + def test_single_request_returns_full_length_conditioning(self) -> None: + hidden_states = torch.randn(1, 5, 4) + outputs = BaseEncoderOutput(hidden_states=(hidden_states, hidden_states)) + text_inputs = SimpleNamespace(attention_mask=torch.ones(1, 5, dtype=torch.long)) + + result = ernie_image_postprocess_text(outputs, text_inputs) + + self.assertIsInstance(result, TextConditioningOutput) + self.assertEqual(result.prompt_seq_lens, [5]) + self.assertEqual(result.prompt_embeds.shape, (1, 5, 4)) + self.assertTrue(bool(result.prompt_embeds_mask.all())) + + def test_ragged_batch_preserves_real_lengths(self) -> None: + hidden_states = torch.randn(2, 6, 4) + outputs = BaseEncoderOutput(hidden_states=(hidden_states, hidden_states)) + # Row 0 has 4 real tokens + 2 pad, row 1 has 6 real tokens (no pad). + mask = torch.tensor([[1, 1, 1, 1, 0, 0], [1, 1, 1, 1, 1, 1]]) + text_inputs = SimpleNamespace(attention_mask=mask) + + result = ernie_image_postprocess_text(outputs, text_inputs) + + self.assertIsInstance(result, TextConditioningOutput) + self.assertEqual(result.prompt_seq_lens, [4, 6]) + self.assertEqual(result.prompt_embeds.shape, (2, 6, 4)) + self.assertEqual( + result.prompt_embeds_mask.tolist(), + [[True, True, True, True, False, False], [True] * 6], + ) + + +class TestErnieImagePrepareCondKwargs(unittest.TestCase): + def _make_batch(self, prompt_embeds, prompt_seq_lens): + return SimpleNamespace( + height=256, + width=256, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=prompt_embeds, + prompt_seq_lens=prompt_seq_lens, + negative_prompt_seq_lens=prompt_seq_lens, + ) + + def test_uniform_lengths_need_no_mask(self) -> None: + config = ErnieImagePipelineConfig() + prompt_embeds = [torch.randn(2, 30, 4)] + batch = self._make_batch(prompt_embeds, [[30, 30]]) + + cond_kwargs = config.prepare_pos_cond_kwargs( + batch, device=torch.device("cpu"), rotary_emb=None, dtype=torch.float32 + ) + + self.assertIsNone(cond_kwargs["encoder_hidden_states_mask"]) + self.assertEqual(cond_kwargs["txt_seq_lens"], [30, 30]) + + def test_ragged_lengths_build_boundary_mask(self) -> None: + config = ErnieImagePipelineConfig() + prompt_embeds = [torch.randn(2, 30, 4)] + batch = self._make_batch(prompt_embeds, [[18, 30]]) + + cond_kwargs = config.prepare_pos_cond_kwargs( + batch, device=torch.device("cpu"), rotary_emb=None, dtype=torch.float32 + ) + + mask = cond_kwargs["encoder_hidden_states_mask"] + self.assertEqual(mask.shape, (2, 30)) + self.assertTrue(bool(mask[0, 17])) + self.assertFalse(bool(mask[0, 18])) + self.assertTrue(bool(mask[1].all())) + + +if __name__ == "__main__": + unittest.main()