[diffusion] fix: fix ragged-caption dynamic-batching accuracy bug in ernie-Image (#30241)

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):

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)

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()