[diffusion] feat: initial support for dynamic batching (#18764)

Signed-off-by: Chi McIsaac <chixie.mcisaac@gmail.com>
Co-authored-by: Junhao Liu <junhaoliu2023@gmail.com>

python/sglang/multimodal_gen/configs/pipeline_configs/base.py

@@ -2,6 +2,7 @@
 
 # SPDX-License-Identifier: Apache-2.0
 import json
+import math
 import os
 from collections.abc import Callable
 from dataclasses import asdict, dataclass, field, fields
@@ -101,6 +102,44 @@ def postprocess_text(output: BaseEncoderOutput, _text_inputs) -> torch.tensor:
     raise NotImplementedError
 
 
+@dataclass(frozen=True)
+class TextConditioningOutput:
+    """Text embeddings and masks aligned to postprocessed sequence length.
+
+    `prompt_embeds_mask` and `prompt_seq_lens` describe real text tokens after
+    model-specific trimming or packing, not the raw tokenizer output.
+    """
+
+    prompt_embeds: torch.Tensor
+    prompt_embeds_mask: torch.Tensor | None = None
+    prompt_seq_lens: list[int] | None = None
+
+
+def pad_text_embeddings_with_mask(
+    text_embeds: list[torch.Tensor],
+) -> TextConditioningOutput:
+    """Pad variable-length text embeddings and return the valid-token mask."""
+    if not text_embeds:
+        raise ValueError("text_embeds must contain at least one tensor")
+
+    max_seq_len = max(e.size(0) for e in text_embeds)
+    prompt_embeds = torch.stack(
+        [
+            torch.cat([e, e.new_zeros(max_seq_len - e.size(0), e.size(1))])
+            for e in text_embeds
+        ]
+    )
+    seq_lens = [int(e.size(0)) for e in text_embeds]
+    seq_lens_tensor = torch.tensor(
+        seq_lens,
+        device=prompt_embeds.device,
+        dtype=torch.long,
+    )
+    positions = torch.arange(max_seq_len, device=prompt_embeds.device).unsqueeze(0)
+    prompt_embeds_mask = positions < seq_lens_tensor.unsqueeze(1)
+    return TextConditioningOutput(prompt_embeds, prompt_embeds_mask, seq_lens)
+
+
 def shard_rotary_emb_for_sp(emb):
     """
     Shard rotary embeddings [S, D] along sequence for SP.
@@ -329,6 +368,41 @@ class PipelineConfig:
     def allow_set_num_frames(self):
         return False
 
+    def supports_dynamic_batching(self):
+        """Return whether this pipeline can opt in to dynamic batching.
+
+        The scheduler still checks each request before merging it into a batch.
+        """
+        return self.task_type in (ModelTaskType.T2I, ModelTaskType.T2V)
+
+    def estimate_request_cost(self, batch) -> float:
+        """Return the relative cost used for batching admission caps.
+
+        This is compared with `max_cost` from the batching config; it is not a
+        memory estimate. The default cost is latent tokens times frames times
+        outputs; pipelines can override it for model-specific admission.
+        """
+        latent_tokens = float(batch.n_tokens or 0)
+        if latent_tokens <= 0:
+            width = int(batch.width or 0)
+            height = int(batch.height or 0)
+            if width > 0 and height > 0:
+                vae_scale = getattr(
+                    self.vae_config.arch_config, "vae_scale_factor", None
+                )
+                if vae_scale is None and hasattr(
+                    self.vae_config, "get_vae_scale_factor"
+                ):
+                    vae_scale = self.vae_config.get_vae_scale_factor()
+                vae_scale = max(1, int(vae_scale or 1))
+                latent_tokens = math.ceil(width / vae_scale) * math.ceil(
+                    height / vae_scale
+                )
+
+        num_frames = max(1, int(batch.num_frames or 1))
+        num_outputs = max(1, int(batch.num_outputs_per_prompt or 1))
+        return latent_tokens * num_frames * num_outputs
+
     def get_decode_scale_and_shift(self, device, dtype, vae):
         vae_arch_config = self.vae_config.arch_config
         scaling_factor = getattr(vae_arch_config, "scaling_factor", None)
@@ -468,6 +542,92 @@ class PipelineConfig:
         """
         return text_inputs.get("attention_mask")
 
+    def build_text_conditioning_mask(
+        self,
+        text_inputs: dict,
+        text_encoder_attention_mask: "torch.Tensor | None",
+        prompt_embeds: "torch.Tensor",
+        encoder_index: int,
+    ) -> "torch.Tensor":
+        """Return a mask aligned with post-processed prompt embeddings.
+
+        True values mark valid text tokens. Dynamic batching must carry
+        post-processed semantic text lengths explicitly; if a model-specific
+        postprocessor changes the sequence length, it must return
+        TextConditioningOutput with an embedding-aligned mask.
+        """
+        if prompt_embeds.ndim < 2:
+            raise ValueError(
+                "prompt_embeds must have shape [batch, seq, ...] to build text conditioning mask"
+            )
+
+        if prompt_embeds.ndim == 2:
+            batch_size, embed_seq_len = 1, prompt_embeds.shape[0]
+        else:
+            batch_size, embed_seq_len = prompt_embeds.shape[:2]
+        device = prompt_embeds.device
+        if text_encoder_attention_mask is None:
+            return torch.ones(
+                (batch_size, embed_seq_len), dtype=torch.bool, device=device
+            )
+
+        raw_mask = text_encoder_attention_mask.to(device=device).bool()
+        if raw_mask.ndim != 2 or raw_mask.shape[0] != batch_size:
+            raise ValueError(
+                "text attention mask must have shape [batch, seq] matching prompt_embeds batch"
+            )
+
+        if raw_mask.shape[1] == embed_seq_len:
+            return raw_mask
+
+        if prompt_embeds.ndim == 2 and raw_mask.shape[0] == 1:
+            return torch.ones((1, embed_seq_len), dtype=torch.bool, device=device)
+
+        raise ValueError(
+            "text attention mask length does not match postprocessed prompt embeddings. "
+            "Postprocess functions that trim, pack, or otherwise change text sequence "
+            "length must return TextConditioningOutput with an embedding-aligned mask."
+        )
+
+    @staticmethod
+    def seq_lens_from_text_conditioning_mask(mask: "torch.Tensor") -> list[int]:
+        if mask.ndim != 2:
+            raise ValueError("text conditioning mask must have shape [batch, seq]")
+        return torch.count_nonzero(mask, dim=1).tolist()
+
+    def require_text_seq_lens(
+        self,
+        batch,
+        encoder_index: int,
+        *,
+        negative: bool = False,
+        expected_batch_size: int | None = None,
+    ) -> list[int]:
+        """Return postprocessed text lengths captured during text encoding.
+
+        Dynamic batches use these lengths for model masks, RoPE, and cache
+        sizing after text embeddings have been padded.
+        """
+        seq_lens_by_encoder = (
+            batch.negative_prompt_seq_lens if negative else batch.prompt_seq_lens
+        )
+        kind = "negative" if negative else "positive"
+        if seq_lens_by_encoder is None or encoder_index >= len(seq_lens_by_encoder):
+            raise ValueError(
+                f"Missing {kind} prompt_seq_lens for text encoder {encoder_index}; "
+                "dynamic text conditioning requires explicit sequence lengths."
+            )
+
+        seq_lens = [int(x) for x in seq_lens_by_encoder[encoder_index]]
+        if expected_batch_size is not None and len(seq_lens) != int(
+            expected_batch_size
+        ):
+            raise ValueError(
+                f"{kind} prompt_seq_lens for text encoder {encoder_index} has "
+                f"{len(seq_lens)} entries, expected {expected_batch_size}."
+            )
+        return seq_lens
+
     def get_text_encoder_pooler_output(
         self, outputs: "BaseEncoderOutput", encoder_index: int
     ) -> "torch.Tensor | None":

python/sglang/multimodal_gen/configs/pipeline_configs/flux.py

@@ -91,6 +91,30 @@ class FluxPipelineConfig(ImagePipelineConfig):
         # Flux v1 does not use attention masks for text encoders.
         return None
 
+    def build_text_conditioning_mask(
+        self,
+        text_inputs: dict,
+        text_encoder_attention_mask: "torch.Tensor | None",
+        prompt_embeds: "torch.Tensor",
+        encoder_index: int,
+    ) -> "torch.Tensor":
+        """Use all-valid fixed-length masks for Flux v1 text embeddings."""
+        if prompt_embeds.ndim < 2:
+            raise ValueError(
+                "prompt_embeds must have shape [batch, seq, ...] or [seq, ...]"
+            )
+        if prompt_embeds.ndim == 2:
+            shape = (1, prompt_embeds.shape[0])
+        else:
+            shape = prompt_embeds.shape[:2]
+        return torch.ones(shape, dtype=torch.bool)
+
+    @staticmethod
+    def seq_lens_from_text_conditioning_mask(mask: "torch.Tensor") -> list[int]:
+        if mask.ndim != 2:
+            raise ValueError("text conditioning mask must have shape [batch, seq]")
+        return [int(mask.shape[1])] * int(mask.shape[0])
+
     def get_text_encoder_pooler_output(self, outputs, encoder_index):
         return outputs.pooler_output
 
@@ -143,7 +167,27 @@ class FluxPipelineConfig(ImagePipelineConfig):
 
         return latent_image_ids
 
-    def get_freqs_cis(self, prompt_embeds, width, height, device, rotary_emb, batch):
+    @staticmethod
+    def _validate_fixed_text_seq_lens(prompt_embeds, txt_seq_lens):
+        if prompt_embeds.ndim < 3:
+            raise ValueError(
+                "Flux text conditioning expects prompt_embeds with shape [batch, seq, dim]"
+            )
+        batch_size, seq_len = prompt_embeds.shape[:2]
+        if len(txt_seq_lens) != batch_size:
+            raise ValueError(
+                f"Flux text sequence lengths have {len(txt_seq_lens)} entries, expected {batch_size}."
+            )
+        if any(int(seq_len_i) != seq_len for seq_len_i in txt_seq_lens):
+            raise ValueError(
+                "Flux currently requires fixed-length text conditioning; "
+                f"got seq_lens={txt_seq_lens}, expected all {seq_len}."
+            )
+
+    def get_freqs_cis(
+        self, prompt_embeds, width, height, device, rotary_emb, batch, txt_seq_lens
+    ):
+        self._validate_fixed_text_seq_lens(prompt_embeds, txt_seq_lens)
         txt_ids = torch.zeros(prompt_embeds.shape[1], 3, device=device)
         img_ids = self._prepare_latent_image_ids(
             original_height=height,
@@ -175,6 +219,21 @@ class FluxPipelineConfig(ImagePipelineConfig):
         return latents
 
     def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype):
+        """Build Flux positive-conditioning kwargs from encoded text state.
+
+        Flux v1 uses encoder index 1 (the T5 encoder) as the token stream that
+        is concatenated with image tokens for rotary position embeddings. The
+        text encoding stage stores per-request sequence lengths in
+        batch.prompt_seq_lens; read them here instead of inferring from padded
+        embeddings so grouped multi-output requests preserve their explicit
+        text-conditioning contract.
+        """
+        txt_seq_lens = self.require_text_seq_lens(
+            batch,
+            1,
+            negative=False,
+            expected_batch_size=batch.prompt_embeds[1].shape[0],
+        )
         return {
             "freqs_cis": self.get_freqs_cis(
                 batch.prompt_embeds[1],
@@ -183,6 +242,7 @@ class FluxPipelineConfig(ImagePipelineConfig):
                 device,
                 rotary_emb,
                 batch,
+                txt_seq_lens,
             ),
             "pooled_projections": (
                 batch.pooled_embeds[0] if batch.pooled_embeds else None
@@ -190,6 +250,13 @@ class FluxPipelineConfig(ImagePipelineConfig):
         }
 
     def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype):
+        """Build Flux negative-conditioning kwargs using T5 sequence lengths."""
+        txt_seq_lens = self.require_text_seq_lens(
+            batch,
+            1,
+            negative=True,
+            expected_batch_size=batch.negative_prompt_embeds[1].shape[0],
+        )
         return {
             "freqs_cis": self.get_freqs_cis(
                 batch.negative_prompt_embeds[1],
@@ -198,6 +265,7 @@ class FluxPipelineConfig(ImagePipelineConfig):
                 device,
                 rotary_emb,
                 batch,
+                txt_seq_lens,
             ),
             "pooled_projections": (
                 batch.neg_pooled_embeds[0] if batch.neg_pooled_embeds else None
@@ -412,6 +480,14 @@ class Flux2PipelineConfig(FluxPipelineConfig):
         # Flux2 does not use pooler output.
         return None
 
+    def supports_dynamic_batching(self):
+        """Allow batching for Flux2 text-only requests.
+
+        Flux2 is a TI2I pipeline, so image-input requests are rejected by the
+        scheduler's request-level batching checks.
+        """
+        return True
+
     def tokenize_prompt(self, prompts: list[str], tokenizer, tok_kwargs) -> dict:
         messages = build_flux2_text_messages(prompts)
         inputs = tokenizer.apply_chat_template(
@@ -513,7 +589,10 @@ class Flux2PipelineConfig(FluxPipelineConfig):
         image_latent_ids = image_latent_ids.repeat(batch.batch_size, 1, 1)
         batch.condition_image_latent_ids = image_latent_ids.to(get_local_torch_device())
 
-    def get_freqs_cis(self, prompt_embeds, width, height, device, rotary_emb, batch):
+    def get_freqs_cis(
+        self, prompt_embeds, width, height, device, rotary_emb, batch, txt_seq_lens
+    ):
+        self._validate_fixed_text_seq_lens(prompt_embeds, txt_seq_lens)
         txt_ids = _prepare_text_ids(prompt_embeds).to(device=device)
 
         img_ids = batch.latent_ids
@@ -544,6 +623,19 @@ class Flux2PipelineConfig(FluxPipelineConfig):
         return cos, sin
 
     def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype):
+        """Build Flux2 positive-conditioning kwargs from encoded text state.
+
+        Flux2 uses encoder index 0 for the Mistral text stream. The stored
+        sequence lengths are passed through to rotary-position preparation so
+        grouped requests use the same text-length metadata that was produced
+        during text encoding.
+        """
+        txt_seq_lens = self.require_text_seq_lens(
+            batch,
+            0,
+            negative=False,
+            expected_batch_size=batch.prompt_embeds[0].shape[0],
+        )
         return {
             "freqs_cis": self.get_freqs_cis(
                 batch.prompt_embeds[0],
@@ -552,6 +644,7 @@ class Flux2PipelineConfig(FluxPipelineConfig):
                 device,
                 rotary_emb,
                 batch,
+                txt_seq_lens,
             )
         }
 

python/sglang/multimodal_gen/configs/pipeline_configs/hunyuan.py

@@ -18,6 +18,7 @@ from sglang.multimodal_gen.configs.models.vaes import HunyuanVAEConfig
 from sglang.multimodal_gen.configs.pipeline_configs.base import (
     ModelTaskType,
     PipelineConfig,
+    TextConditioningOutput,
 )
 
 PROMPT_TEMPLATE_ENCODE_VIDEO = (
@@ -46,19 +47,37 @@ def llama_preprocess_text(prompt: str) -> str:
     return prompt_template_video["template"].format(prompt)
 
 
-def llama_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.tensor:
+def llama_postprocess_text(
+    outputs: BaseEncoderOutput, _text_inputs
+) -> TextConditioningOutput:
     hidden_state_skip_layer = 2
     assert outputs.hidden_states is not None
     hidden_states: tuple[torch.Tensor, ...] = outputs.hidden_states
-    last_hidden_state: torch.tensor = hidden_states[-(hidden_state_skip_layer + 1)]
+    last_hidden_state: torch.Tensor = hidden_states[-(hidden_state_skip_layer + 1)]
     crop_start = prompt_template_video.get("crop_start", -1)
     last_hidden_state = last_hidden_state[:, crop_start:]
-    return last_hidden_state
+    attention_mask = _text_inputs.attention_mask.to(
+        device=last_hidden_state.device, dtype=torch.bool
+    )
+    if crop_start < 0:
+        attention_mask = attention_mask[:, crop_start:]
+    else:
+        attention_mask = attention_mask[
+            :, crop_start : crop_start + last_hidden_state.shape[1]
+        ]
+    seq_lens = [int(x) for x in attention_mask.to(torch.int64).sum(dim=1).tolist()]
+    return TextConditioningOutput(last_hidden_state, attention_mask, seq_lens)
 
 
-def clip_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.tensor:
-    pooler_output: torch.tensor = outputs.pooler_output
-    return pooler_output
+def clip_postprocess_text(
+    outputs: BaseEncoderOutput, _text_inputs
+) -> TextConditioningOutput:
+    pooler_output: torch.Tensor = outputs.pooler_output
+    batch_size = int(pooler_output.shape[0])
+    prompt_embeds_mask = torch.ones(
+        (batch_size, 1), dtype=torch.bool, device=pooler_output.device
+    )
+    return TextConditioningOutput(pooler_output, prompt_embeds_mask, [1] * batch_size)
 
 
 @dataclass

python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py

@@ -16,6 +16,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import (
     ImagePipelineConfig,
     ModelTaskType,
     maybe_unpad_latents,
+    pad_text_embeddings_with_mask,
     shard_rotary_emb_for_sp,
 )
 from sglang.multimodal_gen.configs.post_training.pipeline_configs import (
@@ -45,34 +46,21 @@ def qwen_image_preprocess_text(prompt):
 def qwen_image_postprocess_text(
     outputs, _text_inputs, drop_idx=34, return_attention_mask=False
 ):
+    """Postprocess Qwen text embeddings.
+
+    Returns padded embeddings by default, or TextConditioningOutput when
+    embedding-aligned masks are requested.
+    """
     # squeeze the batch dim
     hidden_states = outputs.hidden_states[-1]
     split_hidden_states = _extract_masked_hidden(
         hidden_states, _text_inputs.attention_mask
     )
     split_hidden_states = [e[drop_idx:] for e in split_hidden_states]
-    attn_mask_list = [
-        torch.ones(e.size(0), dtype=torch.long, device=e.device)
-        for e in split_hidden_states
-    ]
-    max_seq_len = max([e.size(0) for e in split_hidden_states])
-    prompt_embeds = torch.stack(
-        [
-            torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))])
-            for u in split_hidden_states
-        ]
-    )
+    conditioning = pad_text_embeddings_with_mask(split_hidden_states)
     if return_attention_mask:
-        encoder_attention_mask = torch.stack(
-            [
-                torch.cat([u, u.new_zeros(max_seq_len - u.size(0))])
-                for u in attn_mask_list
-            ]
-        )
-        if encoder_attention_mask.all():
-            return prompt_embeds, None
-        return prompt_embeds, encoder_attention_mask
-    return prompt_embeds
+        return conditioning
+    return conditioning.prompt_embeds
 
 
 def qwen_image_edit_postprocess_text(outputs, _text_inputs):
@@ -236,10 +224,14 @@ class QwenImagePipelineConfig(QwenImageRolloutPipelineMixin, ImagePipelineConfig
             return {}
 
     def get_vae_scale_factor(self):
-        return self.vae_config.arch_config.vae_scale_factor
+        return getattr(
+            self.vae_config.arch_config,
+            "vae_scale_factor",
+            self.vae_config.get_vae_scale_factor(),
+        )
 
     def prepare_latent_shape(self, batch, batch_size, num_frames):
-        vae_scale_factor = self.vae_config.arch_config.vae_scale_factor
+        vae_scale_factor = self.get_vae_scale_factor()
         height = 2 * (batch.height // (vae_scale_factor * 2))
         width = 2 * (batch.width // (vae_scale_factor * 2))
         num_channels_latents = self.dit_config.arch_config.in_channels // 4
@@ -247,10 +239,9 @@ class QwenImagePipelineConfig(QwenImageRolloutPipelineMixin, ImagePipelineConfig
         return shape
 
     def maybe_pack_latents(self, latents, batch_size, batch):
-        height = 2 * (
-            batch.height // (self.vae_config.arch_config.vae_scale_factor * 2)
-        )
-        width = 2 * (batch.width // (self.vae_config.arch_config.vae_scale_factor * 2))
+        vae_scale_factor = self.get_vae_scale_factor()
+        height = 2 * (batch.height // (vae_scale_factor * 2))
+        width = 2 * (batch.width // (vae_scale_factor * 2))
         num_channels_latents = self.dit_config.arch_config.in_channels // 4
         # pack latents
         return _pack_latents(latents, batch_size, num_channels_latents, height, width)
@@ -294,11 +285,19 @@ class QwenImagePipelineConfig(QwenImageRolloutPipelineMixin, ImagePipelineConfig
         txt_cos_sin_cache = torch.cat([txt_cos_half, txt_sin_half], dim=-1)
         return img_cos_sin_cache, txt_cos_sin_cache
 
-    def _prepare_cond_kwargs(self, batch, prompt_embeds, rotary_emb, device, dtype):
+    def _prepare_cond_kwargs(
+        self, batch, prompt_embeds, rotary_emb, device, dtype, *, negative=False
+    ):
+        """Build Qwen DiT conditioning kwargs for positive or negative prompts.
+
+        The kwargs include text lengths for RoPE construction and optional
+        encoder masks for cross-attention.
+        """
         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.vae_config.arch_config.vae_scale_factor
+        vae_scale_factor = self.get_vae_scale_factor()
 
         img_shapes = [
             [
@@ -309,14 +308,18 @@ class QwenImagePipelineConfig(QwenImageRolloutPipelineMixin, ImagePipelineConfig
                 )
             ]
         ] * batch_size
-        txt_seq_lens = [prompt_embeds[0].shape[1]]
+        txt_seq_lens, encoder_hidden_states_mask = self._prepare_text_conditioning(
+            batch, 0, text_seq_len, batch_size, negative=negative
+        )
 
         if rotary_emb is None:
-            return {
+            cond_kwargs = {
                 "img_shapes": img_shapes,
                 "txt_seq_lens": txt_seq_lens,
                 "freqs_cis": None,
+                "encoder_hidden_states_mask": encoder_hidden_states_mask,
             }
+            return cond_kwargs
 
         freqs_cis = self.get_freqs_cis(
             img_shapes, txt_seq_lens, rotary_emb, device, dtype
@@ -324,20 +327,105 @@ class QwenImagePipelineConfig(QwenImageRolloutPipelineMixin, ImagePipelineConfig
 
         img_cache, txt_cache = freqs_cis
         img_cache = shard_rotary_emb_for_sp(img_cache)
-        return {
+        cond_kwargs = {
             "txt_seq_lens": txt_seq_lens,
             "freqs_cis": (img_cache, txt_cache),
             "img_shapes": img_shapes,
+            "encoder_hidden_states_mask": encoder_hidden_states_mask,
         }
+        return cond_kwargs
+
+    def _prepare_text_conditioning(
+        self,
+        batch,
+        encoder_index: int,
+        text_seq_len: int,
+        batch_size: int,
+        *,
+        negative: bool = False,
+    ):
+        """Return Qwen text lengths and an optional DiT attention mask.
+
+        Single-request execution uses the full padded length. Batched execution
+        uses stored per-request lengths and masks from text encoding.
+        """
+        if batch_size == 1:
+            return [text_seq_len], None
+
+        txt_seq_lens = self.require_text_seq_lens(
+            batch, encoder_index, negative=negative, expected_batch_size=batch_size
+        )
+        encoder_hidden_states_mask = self._prepare_encoder_hidden_states_mask(
+            batch,
+            encoder_index,
+            txt_seq_lens,
+            text_seq_len,
+            batch_size,
+            negative=negative,
+        )
+        return txt_seq_lens, encoder_hidden_states_mask
+
+    def _prepare_encoder_hidden_states_mask(
+        self,
+        batch,
+        encoder_index: int,
+        txt_seq_lens: list[int],
+        text_seq_len: int,
+        batch_size: int,
+        *,
+        negative: bool = False,
+    ):
+        """Return the text attention mask passed to the Qwen image DiT.
+
+        Qwen image batches can contain prompts with different semantic text
+        lengths after tokenization/postprocessing. The transformer still sees a
+        padded `encoder_hidden_states` tensor with shape [batch, text_seq_len,
+        dim], so we pass a [batch, text_seq_len] boolean mask to keep attention
+        on real text tokens and ignore padding.
+
+        If every request uses the full padded length, no mask is needed and this
+        returns None. Otherwise, prefer the embedding-aligned mask stored by the
+        text encoding stage. If that is unavailable, rebuild the same mask from
+        `txt_seq_lens`: position j is valid for row i when
+        `j < txt_seq_lens[i]`.
+        """
+        if all(seq_len == text_seq_len for seq_len in txt_seq_lens):
+            return None
+
+        masks_by_encoder = (
+            batch.negative_prompt_embeds_mask if negative else batch.prompt_embeds_mask
+        )
+        if masks_by_encoder is not None and encoder_index < len(masks_by_encoder):
+            mask = masks_by_encoder[encoder_index]
+            if mask.shape != (batch_size, text_seq_len):
+                raise ValueError(
+                    "QwenImage text conditioning mask has shape "
+                    f"{tuple(mask.shape)}, expected {(batch_size, text_seq_len)}."
+                )
+            return mask
+
+        # TODO: cache positions by (device, text_seq_len) if this allocation shows up hot.
+        positions = torch.arange(text_seq_len, device=batch.prompt_embeds[0].device)
+        seq_lens = torch.tensor(
+            txt_seq_lens,
+            device=batch.prompt_embeds[0].device,
+            dtype=torch.long,
+        )
+        return positions.unsqueeze(0) < seq_lens.unsqueeze(1)
 
     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 post_denoising_loop(self, latents, batch):
@@ -363,10 +451,11 @@ class QwenImageEditPipelineConfig(QwenImagePipelineConfig):
     )
 
     def _prepare_edit_cond_kwargs(
-        self, batch, prompt_embeds, rotary_emb, device, dtype
+        self, batch, prompt_embeds, rotary_emb, device, dtype, *, negative=False
     ):
         batch_size = batch.latents.shape[0]
         assert batch_size == 1
+        text_seq_len = prompt_embeds[0].shape[1]
         height = batch.height
         width = batch.width
         image_size = batch.original_condition_image_size
@@ -389,14 +478,18 @@ class QwenImageEditPipelineConfig(QwenImagePipelineConfig):
                 ),
             ],
         ] * batch_size
-        txt_seq_lens = [prompt_embeds[0].shape[1]]
+        txt_seq_lens, encoder_hidden_states_mask = self._prepare_text_conditioning(
+            batch, 0, text_seq_len, batch_size, negative=negative
+        )
 
         if rotary_emb is None:
-            return {
+            cond_kwargs = {
                 "img_shapes": img_shapes,
                 "txt_seq_lens": txt_seq_lens,
                 "freqs_cis": None,
+                "encoder_hidden_states_mask": encoder_hidden_states_mask,
             }
+            return cond_kwargs
 
         freqs_cis = QwenImagePipelineConfig.get_freqs_cis(
             img_shapes, txt_seq_lens, rotary_emb, device, dtype
@@ -410,11 +503,13 @@ class QwenImageEditPipelineConfig(QwenImagePipelineConfig):
         img_cache, txt_cache = _shard_qwen_edit_freqs_cis_for_sp(
             freqs_cis, noisy_img_seq_len, device
         )
-        return {
+        cond_kwargs = {
             "txt_seq_lens": txt_seq_lens,
             "freqs_cis": (img_cache, txt_cache),
             "img_shapes": img_shapes,
+            "encoder_hidden_states_mask": encoder_hidden_states_mask,
         }
+        return cond_kwargs
 
     def preprocess_condition_image(
         self, image, target_width, target_height, _vae_image_processor
@@ -453,12 +548,17 @@ class QwenImageEditPipelineConfig(QwenImagePipelineConfig):
 
     def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype):
         return self._prepare_edit_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_edit_cond_kwargs(
-            batch, batch.negative_prompt_embeds, rotary_emb, device, dtype
+            batch,
+            batch.negative_prompt_embeds,
+            rotary_emb,
+            device,
+            dtype,
+            negative=True,
         )
 
     def calculate_condition_image_size(self, image, width, height) -> tuple[int, int]:
@@ -562,10 +662,11 @@ class QwenImageEditPlusPipelineConfig(QwenImageEditPipelineConfig):
         return batch
 
     def _prepare_edit_cond_kwargs(
-        self, batch, prompt_embeds, rotary_emb, device, dtype
+        self, batch, prompt_embeds, rotary_emb, device, dtype, *, negative=False
     ):
         batch_size = batch.latents.shape[0]
         assert batch_size == 1
+        text_seq_len = prompt_embeds[0].shape[1]
         height = batch.height
         width = batch.width
 
@@ -584,7 +685,9 @@ class QwenImageEditPlusPipelineConfig(QwenImageEditPipelineConfig):
                 ],
             ],
         ] * batch_size
-        txt_seq_lens = [prompt_embeds[0].shape[1]]
+        txt_seq_lens, encoder_hidden_states_mask = self._prepare_text_conditioning(
+            batch, 0, text_seq_len, batch_size, negative=negative
+        )
 
         freqs_cis = QwenImageEditPlusPipelineConfig.get_freqs_cis(
             img_shapes, txt_seq_lens, rotary_emb, device, dtype
@@ -595,13 +698,15 @@ class QwenImageEditPlusPipelineConfig(QwenImageEditPipelineConfig):
             1 * (height // vae_scale_factor // 2) * (width // vae_scale_factor // 2)
         )
 
-        return {
+        cond_kwargs = {
             "txt_seq_lens": txt_seq_lens,
             "freqs_cis": _shard_qwen_edit_freqs_cis_for_sp(
                 freqs_cis, noisy_img_seq_len, device
             ),
             "img_shapes": img_shapes,
+            "encoder_hidden_states_mask": encoder_hidden_states_mask,
         }
+        return cond_kwargs
 
 
 @dataclass
@@ -625,17 +730,20 @@ class QwenImageLayeredPipelineConfig(QwenImageEditPipelineConfig):
         return super().postprocess_cfg_noise(batch, noise_pred, noise_pred_cond)
 
     def _prepare_edit_cond_kwargs(
-        self, batch, prompt_embeds, rotary_emb, device, dtype
+        self, batch, prompt_embeds, rotary_emb, device, dtype, *, negative=False
     ):
         batch_size = batch.latents.shape[0]
         assert batch_size == 1
+        text_seq_len = prompt_embeds[0].shape[1]
         height = batch.height
         width = batch.width
 
         vae_scale_factor = self.get_vae_scale_factor()
 
         img_shapes = batch.img_shapes
-        txt_seq_lens = [prompt_embeds[0].shape[1]]
+        txt_seq_lens, encoder_hidden_states_mask = self._prepare_text_conditioning(
+            batch, 0, text_seq_len, batch_size, negative=negative
+        )
 
         freqs_cis = QwenImageEditPlusPipelineConfig.get_freqs_cis(
             img_shapes, txt_seq_lens, rotary_emb, device, dtype
@@ -652,15 +760,17 @@ class QwenImageLayeredPipelineConfig(QwenImageEditPipelineConfig):
             [noisy_img_cache, img_cache[noisy_img_seq_len:, :]], dim=0
         ).to(device=device)
 
-        return {
+        cond_kwargs = {
             "txt_seq_lens": txt_seq_lens,
             "img_shapes": img_shapes,
             "freqs_cis": (img_cache, txt_cache),
             "additional_t_cond": torch.tensor([0], device=device, dtype=torch.long),
+            "encoder_hidden_states_mask": encoder_hidden_states_mask,
         }
+        return cond_kwargs
 
     def _unpad_and_unpack_latents(self, latents, batch):
-        vae_scale_factor = self.vae_config.arch_config.vae_scale_factor
+        vae_scale_factor = self.get_vae_scale_factor()
         channels = self.dit_config.arch_config.in_channels
         batch_size = latents.shape[0]
         layers = batch.num_frames

python/sglang/multimodal_gen/configs/pipeline_configs/sana.py

@@ -64,6 +64,15 @@ class SanaPipelineConfig(SpatialImagePipelineConfig):
 
     text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16",))
 
+    text_encoder_extra_args: list[dict] = field(
+        default_factory=lambda: [
+            {
+                "padding": True,
+                "return_attention_mask": True,
+            }
+        ]
+    )
+
     preprocess_text_funcs: tuple[Callable[[str], str] | None, ...] = field(
         default_factory=lambda: (None,),
     )

python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py

@@ -14,6 +14,8 @@ from sglang.multimodal_gen.configs.models.vaes.flux import FluxVAEConfig
 from sglang.multimodal_gen.configs.pipeline_configs.base import (
     ImagePipelineConfig,
     ModelTaskType,
+    TextConditioningOutput,
+    pad_text_embeddings_with_mask,
 )
 from sglang.multimodal_gen.configs.post_training.pipeline_configs import (
     ZImageRolloutPipelineMixin,
@@ -32,10 +34,24 @@ def zimage_preprocess_text(prompt: str):
     return messages
 
 
-def zimage_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tensor:
+def zimage_postprocess_text(
+    outputs: BaseEncoderOutput, _text_inputs
+) -> torch.Tensor | TextConditioningOutput:
+    """Return unpadded Z-Image text embeddings.
+
+    Batched outputs return TextConditioningOutput to preserve per-prompt text
+    lengths.
+    """
     device = outputs.hidden_states[-2].device
     prompt_mask = _text_inputs.attention_mask.to(device).bool()
-    return outputs.hidden_states[-2][0][prompt_mask[0]]
+    hidden_states = outputs.hidden_states[-2]
+    if hidden_states.shape[0] == 1:
+        return hidden_states[0][prompt_mask[0]]
+
+    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)
 
 
 class TransformersModelConfig(EncoderConfig):
@@ -177,8 +193,72 @@ class ZImagePipelineConfig(ZImageRolloutPipelineMixin, ImagePipelineConfig):
             plan = self._build_zimage_sp_plan(batch)
         return plan
 
+    def _split_text_embeds_for_dit(self, batch, *, negative: bool = False):
+        """Return per-request text tensors, trimming padded batched embeddings."""
+        embeds = batch.negative_prompt_embeds if negative else batch.prompt_embeds
+        if embeds is None:
+            return None
+
+        if isinstance(embeds, (list, tuple)):
+            if not embeds:
+                return []
+            embeds = embeds[0]
+
+        if not torch.is_tensor(embeds):
+            return embeds
+
+        if embeds.ndim == 2:
+            return [embeds]
+
+        if embeds.ndim != 3:
+            raise ValueError(
+                "Z-Image text embeddings must have shape [seq, dim] or [batch, seq, dim]"
+            )
+
+        seq_lens = self.require_text_seq_lens(
+            batch,
+            0,
+            negative=negative,
+            expected_batch_size=int(embeds.shape[0]),
+        )
+        return [
+            embeds[idx, :seq_len].contiguous() for idx, seq_len in enumerate(seq_lens)
+        ]
+
+    def _caption_rope_length(self, prompt_embeds, batch, *, negative: bool = False):
+        """Return the shared caption RoPE length for current text embeddings."""
+        if torch.is_tensor(prompt_embeds):
+            if prompt_embeds.ndim == 2:
+                return int(prompt_embeds.shape[0])
+            if prompt_embeds.ndim == 3:
+                seq_lens = self.require_text_seq_lens(
+                    batch,
+                    0,
+                    negative=negative,
+                    expected_batch_size=int(prompt_embeds.shape[0]),
+                )
+                return max(seq_lens) if seq_lens else int(prompt_embeds.shape[1])
+
+        if isinstance(prompt_embeds, (list, tuple)) and prompt_embeds:
+            first = prompt_embeds[0]
+            if torch.is_tensor(first):
+                if first.ndim == 3:
+                    seq_lens = self.require_text_seq_lens(
+                        batch,
+                        0,
+                        negative=negative,
+                        expected_batch_size=int(first.shape[0]),
+                    )
+                    return max(seq_lens) if seq_lens else int(first.shape[1])
+                return max(int(item.shape[0]) for item in prompt_embeds)
+
+        raise ValueError("Unable to infer Z-Image caption length for rotary embeddings")
+
     def get_pos_prompt_embeds(self, batch):
-        return batch.prompt_embeds
+        return self._split_text_embeds_for_dit(batch, negative=False)
+
+    def get_neg_prompt_embeds(self, batch):
+        return self._split_text_embeds_for_dit(batch, negative=True)
 
     def get_latent_dtype(self, prompt_dtype: torch.dtype) -> torch.dtype:
         # Match the official diffusers Z-Image pipeline, which samples latents in fp32
@@ -252,7 +332,23 @@ class ZImagePipelineConfig(ZImageRolloutPipelineMixin, ImagePipelineConfig):
             return latents[:, :, 0, :, :]
         return latents.view(bs, channels, height, width)
 
-    def get_freqs_cis(self, prompt_embeds, width, height, device, rotary_emb, batch):
+    def get_freqs_cis(
+        self,
+        prompt_embeds,
+        width,
+        height,
+        device,
+        rotary_emb,
+        batch,
+        *,
+        negative: bool = False,
+    ):
+        """Build caption and image RoPE caches for Z-Image conditioning.
+
+        Batched prompts use stored text lengths. SP mode builds image caches for
+        the local spatial shard.
+        """
+
         def create_coordinate_grid(size, start=None, device=None):
             if start is None:
                 start = (0 for _ in size)
@@ -269,7 +365,9 @@ class ZImagePipelineConfig(ZImageRolloutPipelineMixin, ImagePipelineConfig):
             # SP path: keep caption replicated on every rank and build local-only
             # image freqs_cis matching the spatial shard.
             plan = self._get_zimage_sp_plan(batch)
-            cap_ori_len = prompt_embeds.size(0)
+            cap_ori_len = self._caption_rope_length(
+                prompt_embeds, batch, negative=negative
+            )
             cap_padding_len = (-cap_ori_len) % self.SEQ_LEN_MULTIPLE
 
             # caption (replicated prefix)
@@ -304,7 +402,7 @@ class ZImagePipelineConfig(ZImageRolloutPipelineMixin, ImagePipelineConfig):
             x_freqs_cis = rotary_emb(img_pos_ids)
             return (cap_freqs_cis, x_freqs_cis)
 
-        cap_ori_len = prompt_embeds.size(0)
+        cap_ori_len = self._caption_rope_length(prompt_embeds, batch, negative=negative)
         cap_padding_len = (-cap_ori_len) % self.SEQ_LEN_MULTIPLE
         cap_padded_pos_ids = create_coordinate_grid(
             size=(cap_ori_len + cap_padding_len, 1, 1),
@@ -361,9 +459,10 @@ class ZImagePipelineConfig(ZImageRolloutPipelineMixin, ImagePipelineConfig):
         }
 
     def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype):
+        use_negative_embeds = batch.negative_prompt_embeds is not None
         prompt_embeds = (
             batch.negative_prompt_embeds[0]
-            if batch.negative_prompt_embeds is not None
+            if use_negative_embeds
             else batch.prompt_embeds[0]
         )
         return {
@@ -374,6 +473,7 @@ class ZImagePipelineConfig(ZImageRolloutPipelineMixin, ImagePipelineConfig):
                 device,
                 rotary_emb,
                 batch,
+                negative=use_negative_embeds,
             ),
             "image_seq_len_target": (
                 self._get_zimage_sp_plan(batch)["img_seq_target"]

python/sglang/multimodal_gen/configs/sample/sampling_params.py

@@ -12,7 +12,7 @@ import re
 import time
 import unicodedata
 import uuid
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from enum import Enum, auto
 from typing import TYPE_CHECKING, Any, ClassVar
 
@@ -89,11 +89,14 @@ class DataType(Enum):
 class SamplingParams:
     """
     Sampling parameters for generation.
+
+    Dynamic batching compares these fields for compatibility, except fields
+    marked with `batch_sig_exclude`.
     """
 
     data_type: DataType = DataType.VIDEO
 
-    request_id: str | None = None
+    request_id: str | None = field(default=None, metadata={"batch_sig_exclude": True})
 
     # All fields below are copied from ForwardBatch
 
@@ -101,13 +104,17 @@ class SamplingParams:
     image_path: str | list[str] | None = None
 
     # Text inputs
-    prompt: str | list[str] | None = None
+    prompt: str | list[str] | None = field(
+        default=None, metadata={"batch_sig_exclude": True}
+    )
     negative_prompt: str = (
         "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
     )
-    prompt_path: str | None = None
-    output_path: str | None = None
-    output_file_name: str | None = None
+    prompt_path: str | None = field(default=None, metadata={"batch_sig_exclude": True})
+    output_path: str | None = field(default=None, metadata={"batch_sig_exclude": True})
+    output_file_name: str | None = field(
+        default=None, metadata={"batch_sig_exclude": True}
+    )
     output_quality: str | None = "default"
     output_compression: int | None = None
 
@@ -128,7 +135,7 @@ class SamplingParams:
 
     # Batch info
     num_outputs_per_prompt: int = 1
-    seed: int | list[int] = 42
+    seed: int | list[int] = field(default=42, metadata={"batch_sig_exclude": True})
     generator_device: str | None = None  # None means use the pipeline/model default
 
     # Original dimensions (before VAE scaling)
@@ -147,9 +154,9 @@ class SamplingParams:
     fps: int = 24
 
     # Resolution validation
-    supported_resolutions: list[tuple[int, int]] | None = (
-        None  # None means all resolutions allowed
-    )
+    supported_resolutions: list[tuple[int, int]] | None = field(
+        default=None, metadata={"batch_sig_exclude": True}
+    )  # None means all resolutions allowed
 
     # Denoising parameters
     num_inference_steps: int = None
@@ -167,17 +174,21 @@ class SamplingParams:
     )
 
     # Profiling
-    profile: bool = False
-    num_profiled_timesteps: int = 5
-    profile_all_stages: bool = False
+    profile: bool = field(default=False, metadata={"batch_sig_exclude": True})
+    num_profiled_timesteps: int = field(default=5, metadata={"batch_sig_exclude": True})
+    profile_all_stages: bool = field(
+        default=False, metadata={"batch_sig_exclude": True}
+    )
 
     # Debugging
-    debug: bool = False
-    perf_dump_path: str | None = None
+    debug: bool = field(default=False, metadata={"batch_sig_exclude": True})
+    perf_dump_path: str | None = field(
+        default=None, metadata={"batch_sig_exclude": True}
+    )
 
     # Misc
     save_output: bool = True
-    return_frames: bool = False
+    return_frames: bool = field(default=False, metadata={"batch_sig_exclude": True})
     rollout: bool = False
     rollout_sde_type: str = "sde"
     rollout_noise_level: float = 0.7
@@ -197,11 +208,13 @@ class SamplingParams:
     rollout_sde_step_indices: list[int] | None = None
     rollout_return_step_indices: list[int] | None = None
     # if True, disallow user params to override subclass-defined protected fields
-    no_override_protected_fields: bool = False
+    no_override_protected_fields: bool = field(
+        default=False, metadata={"batch_sig_exclude": True}
+    )
     # whether to adjust num_frames for multi-GPU friendly splitting (default: True)
     adjust_frames: bool = True
     # if True, suppress verbose logging for this request
-    suppress_logs: bool = False
+    suppress_logs: bool = field(default=False, metadata={"batch_sig_exclude": True})
 
     return_file_paths_only: bool = True
     enable_sequence_shard: bool | None = None

python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py

@@ -349,7 +349,12 @@ class DiffGenerator:
                 global_output_index += len(requests)
 
         total_gen_time = time.perf_counter() - total_start_time
-        log_batch_completion(logger, len(results), total_gen_time)
+        if self.server_args.batching_max_size > 1:
+            log_batch_completion(
+                logger,
+                len(results),
+                total_gen_time,
+            )
         self._log_summary(results)
 
         if not results:

python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py

@@ -22,7 +22,6 @@ from sglang.multimodal_gen.runtime.entrypoints.utils import (
     SetLoraReq,
     ShutdownReq,
     UnmergeLoraWeightsReq,
-    expand_request_outputs,
     format_lora_message,
     save_outputs,
 )
@@ -326,9 +325,8 @@ async def process_generation_batch(
     batch,
 ) -> tuple[list[str], OutputBatch]:
     total_start_time = time.perf_counter()
-    requests = expand_request_outputs(batch)
     with trace_req(batch.trace_ctx), log_generation_timer(logger, batch.prompt):
-        result = await scheduler_client.forward(requests)
+        result = await scheduler_client.forward([batch])
 
         if result.output is None and result.output_file_paths is None:
             error_msg = result.error or "Unknown error"
@@ -338,22 +336,14 @@ async def process_generation_batch(
 
         if result.output_file_paths:
             save_file_path_list = result.output_file_paths
-            if len(save_file_path_list) < len(requests):
-                raise RuntimeError(
-                    f"Expected at least {len(requests)} output paths, "
-                    f"got {len(save_file_path_list)}"
-                )
         else:
-            if len(result.output) != len(requests):
-                raise RuntimeError(
-                    f"Expected {len(requests)} outputs, got {len(result.output)}"
-                )
+            num_outputs = len(result.output)
             save_file_path_list = save_outputs(
                 result.output,
                 batch.data_type,
                 batch.fps,
                 batch.save_output,
-                lambda idx: str(requests[idx].output_file_path(1, 0)),
+                lambda idx: str(batch.output_file_path(num_outputs, idx)),
                 audio=result.audio,
                 audio_sample_rate=result.audio_sample_rate,
                 output_compression=batch.output_compression,
@@ -367,7 +357,12 @@ async def process_generation_batch(
             )
 
     total_time = time.perf_counter() - total_start_time
-    log_batch_completion(logger, len(save_file_path_list), total_time)
+    if get_global_server_args().batching_max_size > 1:
+        log_batch_completion(
+            logger,
+            len(save_file_path_list),
+            total_time,
+        )
 
     if result.peak_memory_mb and result.peak_memory_mb > 0:
         logger.info(f"Peak memory usage: {result.peak_memory_mb:.2f} MB")

python/sglang/multimodal_gen/runtime/layers/rotary_embedding/base.py

@@ -77,14 +77,14 @@ class RotaryEmbedding(CustomOp):
         cos, sin = cos_sin.chunk(2, dim=-1)
 
         query_shape = query.shape
-        query = query.view(num_tokens, -1, self.head_size)
+        query = query.reshape(num_tokens, -1, self.head_size)
         query_rot = query[..., : self.rotary_dim]
         query_pass = query[..., self.rotary_dim :]
         query_rot = _apply_rotary_emb(query_rot, cos, sin, self.is_neox_style)
         query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
 
         key_shape = key.shape
-        key = key.view(num_tokens, -1, self.head_size)
+        key = key.reshape(num_tokens, -1, self.head_size)
         key_rot = key[..., : self.rotary_dim]
         key_pass = key[..., self.rotary_dim :]
         key_rot = _apply_rotary_emb(key_rot, cos, sin, self.is_neox_style)

python/sglang/multimodal_gen/runtime/managers/dynamic_batch_admission.py

@@ -0,0 +1,345 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Admission control for native diffusion batching.
+
+Native diffusion batching is model, resolution, device, and implementation
+dependent. The scheduler treats `--batching-max-size` as the public ceiling;
+`--batching-config` can apply stricter caps for specific model and shape
+combinations.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass
+from difflib import get_close_matches
+from typing import TYPE_CHECKING, Any
+
+from sglang.multimodal_gen.runtime.loader.utils import BYTES_PER_GB
+from sglang.multimodal_gen.runtime.pipelines_core import Req
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+if TYPE_CHECKING:
+    from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
+logger = init_logger(__name__)
+
+_BATCHING_RULE_KEYS = frozenset(
+    {
+        "model",
+        "model_contains",
+        "resolution",
+        "device_memory_gb_min",
+        "device_memory_gb_max",
+        "offload",
+        "max_batch_size",
+        "max_cost",
+        # Free-form provenance/benchmark metadata. It is intentionally ignored
+        # by admission, but accepted so production configs can explain caps.
+        "calibration",
+    }
+)
+
+
+@dataclass(frozen=True)
+class AdmissionLimit:
+    """Effective batch size and cost caps after matching batching rules."""
+
+    max_batch_size: int
+    max_cost: float | None = None
+    cap_reason: str | None = None
+
+    def reject_reason(self, *, batch_size: int, batch_cost: float) -> str | None:
+        if batch_size > self.max_batch_size:
+            return self.cap_reason or f"config_cap:{self.max_batch_size}"
+        if self.max_cost is not None and batch_cost > self.max_cost:
+            return f"cost_budget:{batch_cost:.0f}>{self.max_cost:.0f}"
+        return None
+
+    def stop_reason_for_next_cost(self, next_batch_cost: float) -> str | None:
+        if self.max_cost is not None and next_batch_cost > self.max_cost:
+            return f"cost_budget_next:{next_batch_cost:.0f}>{self.max_cost:.0f}"
+        return None
+
+
+@dataclass(frozen=True)
+class BatchingRule:
+    """One user-provided batching admission rule loaded from batching config."""
+
+    model: str | None = None
+    model_contains: str | None = None
+    resolution: str | None = None
+    device_memory_gb_min: float | None = None
+    device_memory_gb_max: float | None = None
+    offload: bool | None = None
+    max_batch_size: int = 1
+    max_cost: float | None = None
+    source: str = "user"
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any], *, source: str) -> "BatchingRule":
+        if not isinstance(data, dict):
+            raise ValueError(
+                f"batching config rule from {source} must be an object, "
+                f"got {type(data).__name__}"
+            )
+        _validate_rule_keys(data, source=source)
+        if "max_batch_size" not in data:
+            raise ValueError("batching config rule requires max_batch_size")
+
+        rule = cls(
+            model=_optional_str(data.get("model")),
+            model_contains=_optional_str(data.get("model_contains")),
+            resolution=_optional_str(data.get("resolution")),
+            device_memory_gb_min=_optional_float(data.get("device_memory_gb_min")),
+            device_memory_gb_max=_optional_float(data.get("device_memory_gb_max")),
+            offload=_optional_bool(data.get("offload")),
+            max_batch_size=int(data["max_batch_size"]),
+            max_cost=_optional_float(data.get("max_cost")),
+            source=source,
+        )
+        rule.validate()
+        return rule
+
+    def validate(self) -> None:
+        if self.model is not None and self.model_contains is not None:
+            raise ValueError(
+                "batching config rule cannot set both model and model_contains"
+            )
+        if self.model is None and self.model_contains is None:
+            raise ValueError("batching config rule requires model or model_contains")
+        if self.max_batch_size < 1:
+            raise ValueError("batching config rule max_batch_size must be >= 1")
+        if self.max_cost is not None and self.max_cost <= 0.0:
+            raise ValueError("batching config rule max_cost must be > 0")
+        if (
+            self.device_memory_gb_min is not None
+            and self.device_memory_gb_max is not None
+            and self.device_memory_gb_min > self.device_memory_gb_max
+        ):
+            raise ValueError(
+                "batching config rule device_memory_gb_min must be <= device_memory_gb_max"
+            )
+
+    def matches(
+        self,
+        *,
+        model_path: str,
+        resolution: str | None,
+        device_memory_gb: float | None,
+        offload: bool,
+    ) -> bool:
+        if self.model is not None and self.model != model_path:
+            return False
+        if self.model_contains is not None and self.model_contains not in model_path:
+            return False
+        if self.resolution not in (None, "*") and self.resolution != resolution:
+            return False
+        if self.offload is not None and self.offload != offload:
+            return False
+        if device_memory_gb is None:
+            return True
+        if (
+            self.device_memory_gb_min is not None
+            and device_memory_gb < self.device_memory_gb_min
+        ):
+            return False
+        if (
+            self.device_memory_gb_max is not None
+            and device_memory_gb > self.device_memory_gb_max
+        ):
+            return False
+        return True
+
+
+class BatchAdmissionController:
+    """Applies configured caps before adding requests to a batch."""
+
+    def __init__(self, server_args: "ServerArgs", gpu_id: int):
+        self._mode = getattr(server_args, "batching_mode", "dynamic")
+        self._user_max_batch_size = max(1, int(server_args.batching_max_size))
+        self._model_path = server_args.model_path
+        self._offload = bool(server_args.dit_layerwise_offload)
+        self._device_memory_gb = self._get_device_memory_gb(gpu_id)
+        self._rules = load_batching_config(server_args.batching_config)
+        self._pipeline_config = server_args.pipeline_config
+
+        if self.enabled:
+            logger.info(
+                "Batch admission enabled: user_max=%d, device_memory=%.1fGiB, rules=%d",
+                self._user_max_batch_size,
+                self._device_memory_gb or 0.0,
+                len(self._rules),
+            )
+
+    @property
+    def enabled(self) -> bool:
+        return self._mode == "dynamic" and self._user_max_batch_size > 1
+
+    def reject_reason_for_candidate(
+        self, current_reqs: list[Req], candidate_req: Req
+    ) -> str | None:
+        if not self.enabled:
+            return None
+        proposed = current_reqs + [candidate_req]
+        limit = self.limit_for(proposed[0])
+        return limit.reject_reason(
+            batch_size=len(proposed),
+            batch_cost=self.estimate_batch_cost(proposed),
+        )
+
+    def batch_is_full(self, reqs: list[Req]) -> bool:
+        """Return whether another roughly similar request would exceed the cap."""
+        if not self.enabled or not reqs:
+            return len(reqs) >= self._user_max_batch_size
+
+        limit = self.limit_for(reqs[0])
+        if len(reqs) >= limit.max_batch_size:
+            return True
+
+        next_cost = self.estimate_batch_cost(reqs + [reqs[0]])
+        return limit.max_cost is not None and next_cost > limit.max_cost
+
+    def limit_reason_for_batch(self, reqs: list[Req]) -> str | None:
+        if not self.enabled or not reqs:
+            return None
+
+        limit = self.limit_for(reqs[0])
+        if len(reqs) >= limit.max_batch_size:
+            return limit.cap_reason or f"config_cap:{limit.max_batch_size}"
+
+        next_cost = self.estimate_batch_cost(reqs + [reqs[0]])
+        return limit.stop_reason_for_next_cost(next_cost)
+
+    def max_admissible_batch_size(self, req: Req) -> int:
+        return self.limit_for(req).max_batch_size
+
+    def limit_for(self, req: Req) -> AdmissionLimit:
+        """Return the effective admission limit for the request's model and shape."""
+        rules = self._matching_rules(req)
+        if not rules:
+            return AdmissionLimit(max_batch_size=self._user_max_batch_size)
+
+        config_cap = min(rule.max_batch_size for rule in rules)
+        max_batch_size = min(self._user_max_batch_size, config_cap)
+        cap_reason = (
+            f"config_cap:{max_batch_size}"
+            if max_batch_size < self._user_max_batch_size
+            else None
+        )
+        costs = [rule.max_cost for rule in rules if rule.max_cost is not None]
+        return AdmissionLimit(
+            max_batch_size=max(1, max_batch_size),
+            max_cost=min(costs) if costs else None,
+            cap_reason=cap_reason,
+        )
+
+    def estimate_batch_cost(self, reqs: list[Req]) -> float:
+        return sum(
+            float(self._pipeline_config.estimate_request_cost(req)) for req in reqs
+        )
+
+    def _matching_rules(self, req: Req) -> list[BatchingRule]:
+        return [
+            rule
+            for rule in self._rules
+            if rule.matches(
+                model_path=self._model_path,
+                resolution=req.resolution_key,
+                device_memory_gb=self._device_memory_gb,
+                offload=self._offload,
+            )
+        ]
+
+    @staticmethod
+    def _get_device_memory_gb(gpu_id: int) -> float | None:
+        try:
+            return current_platform.get_device_total_memory(gpu_id) / BYTES_PER_GB
+        except Exception:
+            return None
+
+
+def load_batching_config(path: str | None) -> list[BatchingRule]:
+    if path is None:
+        return []
+
+    with open(path, encoding="utf-8") as f:
+        payload = json.load(f)
+
+    source = os.path.abspath(path)
+    entries = _config_entries(payload)
+    rules = [BatchingRule.from_dict(entry, source=source) for entry in entries]
+    if not rules:
+        raise ValueError(f"batching config {source} does not contain any rules")
+    return rules
+
+
+def _config_entries(payload: Any) -> list[dict[str, Any]]:
+    if isinstance(payload, dict) and payload.get("schema_version") not in (None, 1):
+        raise ValueError("batching config schema_version must be 1")
+    if isinstance(payload, dict) and isinstance(payload.get("rules"), list):
+        return payload["rules"]
+    if isinstance(payload, list):
+        return payload
+    if isinstance(payload, dict):
+        entries: list[dict[str, Any]] = []
+        for key, value in payload.items():
+            if key == "schema_version" or not isinstance(value, dict):
+                continue
+            model, _sep, resolution = key.partition("|")
+            entry = dict(value)
+            if model:
+                entry.setdefault("model", model)
+            if resolution:
+                entry.setdefault("resolution", resolution)
+            entries.append(entry)
+        return entries
+    raise ValueError(
+        "batching config must be a {'schema_version': 1, 'rules': [...]} object, "
+        "a list of rules, or a mapping keyed by model|resolution"
+    )
+
+
+def _validate_rule_keys(data: dict[str, Any], *, source: str) -> None:
+    unknown = sorted(set(data) - _BATCHING_RULE_KEYS)
+    if not unknown:
+        return
+
+    hints = []
+    for key in unknown:
+        matches = get_close_matches(key, _BATCHING_RULE_KEYS, n=1)
+        if matches:
+            hints.append(f"{key!r} (did you mean {matches[0]!r}?)")
+        else:
+            hints.append(repr(key))
+    raise ValueError(
+        f"batching config rule from {source} contains unknown key(s): "
+        f"{', '.join(hints)}"
+    )
+
+
+def _optional_str(value: Any) -> str | None:
+    if value is None:
+        return None
+    return str(value)
+
+
+def _optional_float(value: Any) -> float | None:
+    if value is None:
+        return None
+    return float(value)
+
+
+def _optional_bool(value: Any) -> bool | None:
+    if value is None:
+        return None
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, str):
+        lowered = value.strip().lower()
+        if lowered in ("1", "true", "yes", "y", "on"):
+            return True
+        if lowered in ("0", "false", "no", "n", "off"):
+            return False
+    raise ValueError(f"cannot parse boolean batching config value: {value!r}")

python/sglang/multimodal_gen/runtime/managers/gpu_worker.py

@@ -318,6 +318,7 @@ class GPUWorker:
                 return result
 
             output_batch = self._to_output_batch(result)
+            self._record_output_peak_memory(output_batch)
 
             output_metrics = self._iter_output_metrics(output_batch)
             if self.rank == 0 and output_metrics and not current_platform.is_cpu():
@@ -377,8 +378,15 @@ class GPUWorker:
             if output_batch is None:
                 output_batch = OutputBatch()
             output_batch.error = f"Error executing {error_context}: {e}"
+            self._record_output_peak_memory(output_batch)
         return output_batch
 
+    def _record_output_peak_memory(self, output_batch: OutputBatch) -> None:
+        if self.rank != 0 or current_platform.is_cpu():
+            return
+        peak_reserved_bytes = torch.get_device_module().max_memory_reserved()
+        output_batch.peak_memory_mb = peak_reserved_bytes / (1024**2)
+
     def _forward_group(self, batch: list[Req]) -> OutputBatch:
         assert self.pipeline is not None
         results = self.pipeline.forward_batch(batch, self.server_args)
@@ -389,12 +397,32 @@ class GPUWorker:
         if self.rank != 0 or output_batch.output is None:
             return
 
+        dynamic_output_paths = None
+        if req.extra:
+            dynamic_output_paths = req.extra.get("dynamic_batch_output_paths")
+        if dynamic_output_paths is not None and (
+            len(dynamic_output_paths) != len(output_batch.output)
+        ):
+            logger.warning(
+                "dynamic_batch_output_paths length mismatch (got=%d, expected=%d). "
+                "Falling back to merged request output file naming.",
+                len(dynamic_output_paths),
+                len(output_batch.output),
+            )
+            dynamic_output_paths = None
+
+        if dynamic_output_paths is not None:
+            build_output_path = lambda idx: dynamic_output_paths[idx]
+        else:
+            num_outputs = len(output_batch.output)
+            build_output_path = lambda idx: req.output_file_path(num_outputs, idx)
+
         output_batch.output_file_paths = save_outputs(
             output_batch.output,
             req.data_type,
             req.fps,
             True,
-            lambda idx: req.output_file_path(len(output_batch.output), idx),
+            build_output_path,
             audio=output_batch.audio,
             audio_sample_rate=output_batch.audio_sample_rate,
             output_compression=req.output_compression,

python/sglang/multimodal_gen/runtime/managers/scheduler.py

@@ -2,10 +2,14 @@
 
 # SPDX-License-Identifier: Apache-2.0
 import asyncio
+import dataclasses
 import os
 import pickle
 import tempfile
+import time
 from collections import deque
+from copy import deepcopy
+from enum import Enum
 from typing import Any, List
 
 import zmq
@@ -32,9 +36,15 @@ from sglang.multimodal_gen.runtime.entrypoints.utils import (
     UnmergeLoraWeightsReq,
 )
 from sglang.multimodal_gen.runtime.managers.cpu_worker import CPUWorker
+from sglang.multimodal_gen.runtime.managers.dynamic_batch_admission import (
+    BatchAdmissionController,
+)
 from sglang.multimodal_gen.runtime.managers.gpu_worker import GPUWorker
 from sglang.multimodal_gen.runtime.pipelines_core import Req
-from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import (
+    BatchMetricsWindow,
+    OutputBatch,
+)
 from sglang.multimodal_gen.runtime.server_args import (
     PortArgs,
     ServerArgs,
@@ -56,6 +66,9 @@ MINIMUM_PICTURE_BASE64_FOR_WARMUP = "data:image/jpg;base64,iVBORw0KGgoAAAANSUhEU
 # _combine_cfg_parallel's all-reduce.
 DEFAULT_PLACEHOLDER_PROMPT = "warmup"
 
+_MAX_RECV_REQS_PER_POLL = 1024
+_BATCH_METRICS_LOG_INTERVAL = 5
+
 
 class Scheduler(SchedulerDisaggMixin):
     """
@@ -122,8 +135,16 @@ class Scheduler(SchedulerDisaggMixin):
             GetWeightsChecksumReqInput: self._handle_get_weights_checksum,
         }
 
-        # FIFO, new reqs are appended
-        self.waiting_queue: deque[tuple[bytes, Any]] = deque()
+        # FIFO queue entries: (identity, request, enqueue_ts_s)
+        self.waiting_queue: deque[tuple[bytes | None, Any, float]] = deque()
+        self._batching_max_size = server_args.batching_max_size
+        self._batching_delay_s = server_args.batching_delay_ms / 1000.0
+        self._batch_metrics_enabled = server_args.enable_batching_metrics
+        self._batch_metrics_window = BatchMetricsWindow()
+        self._batch_admission = BatchAdmissionController(server_args, gpu_id=local_rank)
+        self._poller = zmq.Poller()
+        if self.receiver is not None:
+            self._poller.register(self.receiver, zmq.POLLIN)
 
         # whether we've send the necessary warmup reqs
         self.warmed_up = False
@@ -139,6 +160,12 @@ class Scheduler(SchedulerDisaggMixin):
 
         self._init_disagg_state(server_args, local_rank)
 
+        if self._batch_metrics_enabled:
+            logger.info(
+                "Dynamic batch metrics enabled; logging summary every %d dispatches.",
+                _BATCH_METRICS_LOG_INTERVAL,
+            )
+
     def get_disagg_metrics(self) -> dict | None:
         """Return disagg role metrics snapshot, or None if not in disagg mode."""
         if self._disagg_metrics is None:
@@ -216,31 +243,49 @@ class Scheduler(SchedulerDisaggMixin):
         req = cls._first_generation_req(req_or_group)
         return req.is_warmup if req is not None else False
 
-    def _dispatch_request(self, reqs: list[Any]) -> OutputBatch:
-        """dispatch req to its registered handler"""
-        req_or_group = reqs[0]
+    def _dispatch_single_request(self, req_or_group: Any) -> OutputBatch:
         if isinstance(req_or_group, list):
-            return self._handle_generation(reqs)
+            if not all(isinstance(req, Req) for req in req_or_group):
+                return OutputBatch(
+                    error=f"Unknown request group type: {type(req_or_group)}"
+                )
+            return self._handle_generation(req_or_group, allow_dynamic_batching=False)
 
         handler = self.request_handlers.get(type(req_or_group))
         if handler is None:
             return OutputBatch(error=f"Unknown request type: {type(req_or_group)}")
-        return handler(reqs)
+        return handler([req_or_group])
+
+    def _dispatch_items(
+        self, items: list[tuple[bytes | None, Any]]
+    ) -> OutputBatch | list[OutputBatch]:
+        """Dispatch ready queue items; several plain `Req`s form one dynamic batch."""
+        reqs = [item[1] for item in items]
+        if len(reqs) > 1 and all(isinstance(req, Req) for req in reqs):
+            return self._handle_generation(reqs, allow_dynamic_batching=True)
+        if len(reqs) > 1:
+            return [self._dispatch_single_request(req) for req in reqs]
+        return self._dispatch_single_request(reqs[0])
 
     def _log_warmup_result(self, output_batch: OutputBatch, is_warmup: bool) -> None:
         if not is_warmup:
             return
 
         if output_batch.error is None:
+            total_duration_s = (
+                output_batch.metrics.total_duration_s
+                if output_batch.metrics is not None
+                else 0.0
+            )
             if self._warmup_total > 0:
                 logger.info(
                     f"Warmup req ({self._warmup_processed}/{self._warmup_total}) processed in {GREEN}%.2f{RESET} seconds",
-                    output_batch.metrics.total_duration_s,
+                    total_duration_s,
                 )
             else:
                 logger.info(
                     f"Warmup req processed in {GREEN}%.2f{RESET} seconds",
-                    output_batch.metrics.total_duration_s,
+                    total_duration_s,
                 )
         else:
             if self._warmup_total > 0:
@@ -250,7 +295,10 @@ class Scheduler(SchedulerDisaggMixin):
             else:
                 logger.info("Warmup req processing failed")
 
-    def _handle_generation(self, reqs: list[Any]):
+    def _handle_generation(
+        self, reqs: list[Any], *, allow_dynamic_batching: bool = True
+    ):
+        """Dispatch generation requests, merging compatible requests when allowed."""
         reqs = self._normalize_generation_reqs(reqs)
         warmup_reqs = [req for req in reqs if req.is_warmup]
         if warmup_reqs:
@@ -262,8 +310,7 @@ class Scheduler(SchedulerDisaggMixin):
             else:
                 logger.info("Processing warmup req...")
 
-        # Diffusion dispatches one generation request at a time, so reqs[0]
-        # always carries the trace context for the entire batch.
+        # Use the head request trace context for scheduler-side dispatch work.
         req = reqs[0]
         req.trace_ctx.rebuild_thread_context()
         with trace_slice(
@@ -271,7 +318,285 @@ class Scheduler(SchedulerDisaggMixin):
             DiffStage.SCHEDULER_DISPATCH,
             thread_finish_flag=True,
         ):
-            return self.worker.execute_forward(reqs)
+            if len(reqs) == 1 or not allow_dynamic_batching:
+                return self.worker.execute_forward(reqs)
+
+            merged_req = self._try_merge_generation_reqs(reqs)
+            if merged_req is None:
+                return self._execute_generation_sequential(reqs)
+
+            batch_size = len(reqs)
+            try:
+                output_batch = self.worker.execute_forward([merged_req])
+                if output_batch.error:
+                    logger.error(
+                        "Dynamic batch execution returned error. Skipping sequential fallback and returning errors: %s",
+                        output_batch.error,
+                    )
+                    return self._build_dynamic_batch_error_outputs(
+                        reqs=reqs,
+                        error_msg=output_batch.error,
+                    )
+
+                split_outputs = self._split_batched_output(output_batch, reqs)
+                if split_outputs is None:
+                    logger.error(
+                        "Failed to split dynamic batched output cleanly. Skipping sequential fallback and returning errors."
+                    )
+                    return self._build_dynamic_batch_error_outputs(
+                        reqs=reqs,
+                        error_msg="Dynamic batching failed: could not split merged output.",
+                    )
+
+                logger.info(
+                    "Processed dynamic batch of %d/%d request(s) with max_delay=%.2fms",
+                    batch_size,
+                    self._batching_max_size,
+                    self._batching_delay_s * 1000.0,
+                )
+                return split_outputs
+            except Exception as e:
+                logger.error(
+                    "Dynamic batching failed (%s). Skipping sequential fallback and returning errors.",
+                    e,
+                    exc_info=True,
+                )
+                return self._build_dynamic_batch_error_outputs(
+                    reqs=reqs,
+                    error_msg=f"Dynamic batching failed: {e}",
+                )
+
+    def _execute_generation_sequential(self, reqs: List[Req]) -> List[OutputBatch]:
+        return [self.worker.execute_forward([req]) for req in reqs]
+
+    @staticmethod
+    def _percentile(values: list[float], percentile: float) -> float:
+        if not values:
+            return 0.0
+        ordered = sorted(values)
+        index = min(
+            len(ordered) - 1,
+            max(0, int(round((percentile / 100.0) * (len(ordered) - 1)))),
+        )
+        return ordered[index]
+
+    def _freeze_signature_value(self, value: Any):
+        """Convert a value into a hashable, order-stable form for signature comparison."""
+        if isinstance(value, (str, int, float, bool, type(None))):
+            return value
+        if isinstance(value, Enum):
+            return value.value
+        if isinstance(value, dict):
+            return {
+                str(k): self._freeze_signature_value(v)
+                for k, v in sorted(value.items(), key=lambda kv: str(kv[0]))
+            }
+        if isinstance(value, (list, tuple)):
+            return tuple(self._freeze_signature_value(v) for v in value)
+        return repr(value)
+
+    def _sampling_param_signature_items(self, req: Req) -> list[tuple[str, Any]] | None:
+        """Return per-field sampling-param signature items, skipping batch_sig_exclude fields."""
+        sp = req.sampling_params
+        if sp is None:
+            return None
+
+        try:
+            sp_fields = dataclasses.fields(sp)
+        except Exception:
+            return None
+
+        return [
+            (f.name, self._freeze_signature_value(getattr(sp, f.name, None)))
+            for f in sp_fields
+            if not f.metadata.get("batch_sig_exclude", False)
+        ]
+
+    def _diffusers_kwargs_signature_value(self, req: Req) -> Any:
+        return self._freeze_signature_value((req.extra or {}).get("diffusers_kwargs"))
+
+    def _build_dynamic_batch_signature(self, req: Req) -> tuple[Any, ...] | None:
+        """Build the request compatibility signature for dynamic batching.
+
+        The signature is built from `SamplingParams` fields, excluding fields
+        marked with `batch_sig_exclude`, plus generation-affecting
+        `extra.diffusers_kwargs`.
+        """
+        signature_items = self._sampling_param_signature_items(req)
+        if signature_items is None:
+            return None
+
+        if req.extra:
+            diffusers_kwargs = req.extra.get("diffusers_kwargs")
+            if diffusers_kwargs:
+                signature_items.append(
+                    (
+                        "diffusers_kwargs",
+                        self._freeze_signature_value(diffusers_kwargs),
+                    )
+                )
+
+        return tuple(signature_items)
+
+    def _get_cached_signature(self, req: Req) -> tuple[Any, ...] | None:
+        cached = getattr(req, "_dynamic_batch_sig", None)
+        if cached is not None:
+            return cached
+        sig = self._build_dynamic_batch_signature(req)
+        req._dynamic_batch_sig = sig  # type: ignore[attr-defined]
+        return sig
+
+    def _find_sampling_param_mismatch_field(
+        self, base_req: Req, candidate_req: Req
+    ) -> str | None:
+        base_items = self._sampling_param_signature_items(base_req)
+        candidate_items = self._sampling_param_signature_items(candidate_req)
+        if base_items is None or candidate_items is None:
+            return None
+
+        if len(base_items) != len(candidate_items):
+            return "sampling_params"
+
+        for (name, base_value), (candidate_name, candidate_value) in zip(
+            base_items, candidate_items
+        ):
+            if name != candidate_name:
+                return "sampling_params"
+            if base_value != candidate_value:
+                return f"sampling_params.{name}"
+
+        base_diffusers_kwargs = self._diffusers_kwargs_signature_value(base_req)
+        candidate_diffusers_kwargs = self._diffusers_kwargs_signature_value(
+            candidate_req
+        )
+        if base_diffusers_kwargs != candidate_diffusers_kwargs:
+            return "extra.diffusers_kwargs"
+
+        return None
+
+    def _get_dynamic_batch_reject_reason(
+        self, base_req: Req, candidate_req: Req
+    ) -> str | None:
+        """Return the first reason `candidate_req` cannot batch with `base_req`, or None."""
+        if self._can_dynamic_batch(base_req, candidate_req):
+            return None
+
+        if base_req.is_warmup or candidate_req.is_warmup:
+            return "warmup"
+        if not isinstance(base_req.prompt, str) or not isinstance(
+            candidate_req.prompt, str
+        ):
+            return "prompt_type"
+        if base_req.image_path is not None or candidate_req.image_path is not None:
+            return "image_conditioning"
+        if base_req.return_file_paths_only != candidate_req.return_file_paths_only:
+            return "return_file_paths_only"
+
+        base_sig = self._get_cached_signature(base_req)
+        candidate_sig = self._get_cached_signature(candidate_req)
+        if base_sig is None or candidate_sig is None:
+            return "signature_unavailable"
+
+        return (
+            self._find_sampling_param_mismatch_field(base_req, candidate_req)
+            or "signature_mismatch"
+        )
+
+    def _can_dynamic_batch(self, base_req: Req, candidate_req: Req) -> bool:
+        """Return whether `candidate_req` can be merged into a batch with `base_req`."""
+        if base_req.is_warmup or candidate_req.is_warmup:
+            return False
+
+        if not isinstance(base_req.prompt, str) or not isinstance(
+            candidate_req.prompt, str
+        ):
+            return False
+
+        if base_req.image_path is not None or candidate_req.image_path is not None:
+            return False
+        if base_req.return_file_paths_only != candidate_req.return_file_paths_only:
+            return False
+
+        base_sig = self._get_cached_signature(base_req)
+        cand_sig = self._get_cached_signature(candidate_req)
+        return base_sig is not None and base_sig == cand_sig
+
+    def _record_batch_dispatch_metrics(
+        self,
+        batch_size: int,
+        queue_wait_ms: float,
+        effective_max_batch_size: int,
+        reject_reasons: list[str] | None = None,
+        stop_reason: str | None = None,
+    ) -> None:
+        if not self._batch_metrics_enabled:
+            return
+
+        effective_max_batch_size = max(1, effective_max_batch_size)
+        logger.info(
+            "Dynamic batch dispatch: size=%d/%d, user_max=%d, queue_wait=%.2fms, stop_reason=%s",
+            batch_size,
+            effective_max_batch_size,
+            self._batching_max_size,
+            max(queue_wait_ms, 0.0),
+            stop_reason or "unspecified",
+        )
+
+        window = self._batch_metrics_window
+        window.dispatches += 1
+        window.total_requests += batch_size
+        window.total_capacity += effective_max_batch_size
+        if batch_size > 1:
+            window.merged_dispatches += 1
+        if self._dynamic_batching_enabled() and batch_size >= effective_max_batch_size:
+            window.full_dispatches += 1
+        window.wait_times_ms.append(max(queue_wait_ms, 0.0))
+        if reject_reasons:
+            window.reject_reasons.update(reject_reasons)
+
+        if window.dispatches >= _BATCH_METRICS_LOG_INTERVAL:
+            self._log_batch_metrics_summary()
+
+    def _log_batch_metrics_summary(self) -> None:
+        if not self._batch_metrics_enabled:
+            return
+
+        window = self._batch_metrics_window
+        if window.dispatches == 0:
+            return
+
+        avg_size = window.total_requests / window.dispatches
+        utilization = window.total_requests / max(1, window.total_capacity)
+        avg_wait_ms = sum(window.wait_times_ms) / len(window.wait_times_ms)
+        p95_wait_ms = self._percentile(window.wait_times_ms, 95.0)
+        merged_rate = window.merged_dispatches / window.dispatches
+        full_rate = window.full_dispatches / window.dispatches
+        top_rejects = ", ".join(
+            f"{reason}={count}"
+            for reason, count in window.reject_reasons.most_common(5)
+        )
+        if not top_rejects:
+            top_rejects = "none"
+
+        logger.info(
+            "Dynamic batch stats (last %d dispatches): avg_size=%.2f, merged_rate=%.1f%%, full_rate=%.1f%%, utilization=%.1f%%, wait_avg=%.2fms, wait_p95=%.2fms, top_rejects=%s",
+            window.dispatches,
+            avg_size,
+            merged_rate * 100.0,
+            full_rate * 100.0,
+            utilization * 100.0,
+            avg_wait_ms,
+            p95_wait_ms,
+            top_rejects,
+        )
+        self._batch_metrics_window = BatchMetricsWindow()
+
+    def _build_dynamic_batch_error_outputs(
+        self,
+        reqs: List[Req],
+        error_msg: str,
+    ) -> List[OutputBatch]:
+        return [OutputBatch(error=error_msg) for _ in reqs]
 
     def return_result(
         self,
@@ -285,15 +610,265 @@ class Scheduler(SchedulerDisaggMixin):
         if not is_warmup and self.receiver is not None and identity is not None:
             self.receiver.send_multipart([identity, b"", pickle.dumps(output_batch)])
 
-    def get_next_batch_to_run(self) -> list[tuple[bytes, Any]] | None:
-        """pull a req from waiting_queue"""
+    def _try_merge_generation_reqs(self, reqs: List[Req]) -> Req | None:
+        """Create a batched generation request from compatible requests.
+
+        Per-request seeds and output paths are stored in `extra` so downstream
+        stages can preserve request ordering.
+        """
+        if len(reqs) <= 1:
+            return reqs[0] if reqs else None
+
+        base_req = reqs[0]
+        for req in reqs[1:]:
+            if not self._can_dynamic_batch(base_req, req):
+                return None
+
+        merged_req = deepcopy(base_req)
+        merged_req.prompt = [req.prompt for req in reqs]
+
+        merged_req.extra = deepcopy(merged_req.extra)
+        merged_req.extra["dynamic_batch_seeds"] = [req.seed for req in reqs]
+        merged_req.return_file_paths_only = base_req.return_file_paths_only
+        if merged_req.return_file_paths_only:
+            dynamic_output_paths: list[str] = []
+            for req in reqs:
+                for output_idx in range(req.num_outputs_per_prompt):
+                    dynamic_output_paths.append(
+                        req.output_file_path(req.num_outputs_per_prompt, output_idx)
+                    )
+            merged_req.extra["dynamic_batch_output_paths"] = dynamic_output_paths
+        merged_req.request_id = f"dynamic_batch::{merged_req.request_id}"
+
+        return merged_req
+
+    @staticmethod
+    def _count_first_dim(value: Any) -> int | None:
+        if value is None:
+            return None
+        if isinstance(value, (list, tuple)):
+            return len(value)
+
+        shape = getattr(value, "shape", None)
+        if shape is not None:
+            try:
+                if len(shape) > 0:
+                    return int(shape[0])
+            except Exception:
+                return None
+        return None
+
+    def _slice_batched_value(
+        self, value: Any, start: int, end: int, total_items: int
+    ) -> Any:
+        if value is None:
+            return None
+
+        if isinstance(value, (list, tuple)):
+            if len(value) == total_items:
+                sliced = value[start:end]
+                return list(sliced) if isinstance(value, list) else tuple(sliced)
+            return deepcopy(value)
+
+        value_items = self._count_first_dim(value)
+        if value_items == total_items:
+            try:
+                return value[start:end]
+            except Exception:
+                pass
+
+        # Scalar / non-batched metadata
+        return deepcopy(value)
+
+    def _split_batched_output(
+        self, output_batch: OutputBatch, reqs: List[Req]
+    ) -> List[OutputBatch] | None:
+        """Split a merged result only when outputs map one-to-one to requests."""
+        per_req_counts = [req.num_outputs_per_prompt for req in reqs]
+        total_items = sum(per_req_counts)
+        output_items = self._count_first_dim(output_batch.output)
+        output_path_items = self._count_first_dim(output_batch.output_file_paths)
+
+        if output_items is None and output_path_items is None:
+            logger.warning(
+                "Batched output has neither tensor outputs nor output_file_paths; cannot split safely."
+            )
+            return None
+
+        if output_items is not None and output_items != total_items:
+            logger.warning(
+                "Unexpected batched output size: got %s items, expected %s",
+                output_items,
+                total_items,
+            )
+            return None
+        if output_path_items is not None and output_path_items != total_items:
+            logger.warning(
+                "Unexpected batched output_file_paths size: got %s items, expected %s",
+                output_path_items,
+                total_items,
+            )
+            return None
+
+        outputs: list[OutputBatch] = []
+        start = 0
+        for req, req_count in zip(reqs, per_req_counts):
+            end = start + req_count
+            split = OutputBatch(
+                output=self._slice_batched_value(
+                    output_batch.output, start, end, total_items
+                ),
+                audio=self._slice_batched_value(
+                    output_batch.audio, start, end, total_items
+                ),
+                audio_sample_rate=output_batch.audio_sample_rate,
+                trajectory_timesteps=self._slice_batched_value(
+                    output_batch.trajectory_timesteps, start, end, total_items
+                ),
+                trajectory_latents=self._slice_batched_value(
+                    output_batch.trajectory_latents, start, end, total_items
+                ),
+                trajectory_decoded=self._slice_batched_value(
+                    output_batch.trajectory_decoded, start, end, total_items
+                ),
+                error=output_batch.error,
+                output_file_paths=self._slice_batched_value(
+                    output_batch.output_file_paths, start, end, total_items
+                ),
+                metrics=deepcopy(output_batch.metrics),
+                noise_pred=self._slice_batched_value(
+                    output_batch.noise_pred, start, end, total_items
+                ),
+                peak_memory_mb=output_batch.peak_memory_mb,
+            )
+            if split.metrics is not None:
+                split.metrics.request_id = req.request_id
+            outputs.append(split)
+            start = end
+
+        return outputs
+
+    def _dynamic_batching_enabled(self) -> bool:
+        """Return whether this server and pipeline can use dynamic batching.
+
+        This is the coarse gate; request-level checks decide which requests can
+        actually be merged.
+        """
+        pipeline_config = self.server_args.pipeline_config
+        supports_dynamic_batching = getattr(
+            pipeline_config, "supports_dynamic_batching", None
+        )
+        if callable(supports_dynamic_batching):
+            return self._batch_admission.enabled and supports_dynamic_batching()
+        return self._batch_admission.enabled
+
+    def get_next_batch_to_run(self) -> list[tuple[bytes | None, Any]] | None:
+        """Return the next dispatchable queue item or dynamic batch.
+
+        Returns None when the head request is waiting for more compatible
+        requests within the configured batching delay.
+        """
         if not self.waiting_queue:
             return None
 
-        # pop the first (earliest)
-        item = self.waiting_queue.popleft()
+        if not self._dynamic_batching_enabled():
+            identity, req, enqueue_time = self.waiting_queue.popleft()
+            if isinstance(req, Req):
+                self._record_batch_dispatch_metrics(
+                    batch_size=1,
+                    queue_wait_ms=(time.monotonic() - enqueue_time) * 1000.0,
+                    effective_max_batch_size=1,
+                    stop_reason="dynamic_disabled",
+                )
+            return [(identity, req)]
 
-        return [item]
+        identity, req, enqueue_time = self.waiting_queue[0]
+        if not isinstance(req, Req):
+            identity, req, _ = self.waiting_queue.popleft()
+            return [(identity, req)]
+
+        # If the head request itself is not eligible for dynamic batching
+        # (e.g., image-conditioned i2i request), dispatch it immediately.
+        if not self._can_dynamic_batch(req, req):
+            identity, req, head_enqueue_time = self.waiting_queue.popleft()
+            reject_reasons: list[str] = []
+            if self._batch_metrics_enabled:
+                reason = self._get_dynamic_batch_reject_reason(req, req)
+                if reason is not None:
+                    reject_reasons.append(f"head:{reason}")
+            self._record_batch_dispatch_metrics(
+                batch_size=1,
+                queue_wait_ms=(time.monotonic() - head_enqueue_time) * 1000.0,
+                effective_max_batch_size=1,
+                reject_reasons=reject_reasons,
+                stop_reason=reject_reasons[0] if reject_reasons else "head_ineligible",
+            )
+            return [(identity, req)]
+
+        compatible_indices: list[int] = [0]
+        compatible_reqs: list[Req] = [req]
+        reject_reasons: list[str] = []
+        for idx in range(1, len(self.waiting_queue)):
+            if len(
+                compatible_indices
+            ) >= self._batching_max_size or self._batch_admission.batch_is_full(
+                compatible_reqs
+            ):
+                break
+            _identity, candidate_req, _enqueue_time = self.waiting_queue[idx]
+            if isinstance(candidate_req, Req) and self._can_dynamic_batch(
+                req, candidate_req
+            ):
+                admission_reject = self._batch_admission.reject_reason_for_candidate(
+                    compatible_reqs, candidate_req
+                )
+                if admission_reject is None:
+                    compatible_indices.append(idx)
+                    compatible_reqs.append(candidate_req)
+                elif self._batch_metrics_enabled:
+                    reject_reasons.append(admission_reject)
+            elif self._batch_metrics_enabled and isinstance(candidate_req, Req):
+                reason = self._get_dynamic_batch_reject_reason(req, candidate_req)
+                if reason is not None:
+                    reject_reasons.append(reason)
+
+        batch_len = len(compatible_indices)
+
+        oldest_wait_s = time.monotonic() - enqueue_time
+
+        should_wait_for_more = (
+            batch_len < self._batching_max_size
+            and not self._batch_admission.batch_is_full(compatible_reqs)
+            and oldest_wait_s < self._batching_delay_s
+        )
+        if should_wait_for_more:
+            return None
+
+        batch_items: list[tuple[bytes | None, Any]] = [None] * batch_len
+        for pos, idx in enumerate(reversed(compatible_indices)):
+            item_identity, item_req, _ = self.waiting_queue[idx]
+            batch_items[batch_len - 1 - pos] = (item_identity, item_req)
+            del self.waiting_queue[idx]
+        stop_reason = self._batch_admission.limit_reason_for_batch(compatible_reqs)
+        if stop_reason is None:
+            if batch_len >= self._batching_max_size:
+                stop_reason = "max_size"
+            elif reject_reasons:
+                stop_reason = reject_reasons[0]
+            elif oldest_wait_s >= self._batching_delay_s:
+                stop_reason = "delay"
+            else:
+                stop_reason = "ready"
+        self._record_batch_dispatch_metrics(
+            batch_size=batch_len,
+            queue_wait_ms=oldest_wait_s * 1000.0,
+            effective_max_batch_size=self._batch_admission.max_admissible_batch_size(
+                compatible_reqs[0]
+            ),
+            reject_reasons=reject_reasons,
+            stop_reason=stop_reason,
+        )
+        return batch_items
 
     def prepare_server_warmup_reqs(self):
         if (
@@ -334,7 +909,7 @@ class Scheduler(SchedulerDisaggMixin):
                     req_kwargs["do_classifier_free_guidance"] = True
                 req = Req(**req_kwargs)
                 req.set_as_warmup(self.server_args.warmup_steps)
-                self.waiting_queue.append((None, req))
+                self.waiting_queue.append((None, req, time.monotonic()))
             # if server is warmed-up, set this flag to avoid req-based warmup
             self.warmed_up = True
 
@@ -408,34 +983,51 @@ class Scheduler(SchedulerDisaggMixin):
             self.warmed_up = True
         return recv_reqs
 
+    @staticmethod
+    def _normalize_received_payload(
+        identity: bytes, reqs: Any
+    ) -> list[tuple[bytes, Any]]:
+        """Normalize client payloads into queue entries.
+
+        A single-item `[Req]` is one request; a multi-item `list[Req]` remains
+        grouped as one logical request.
+        """
+        if not isinstance(reqs, list):
+            return [(identity, reqs)]
+        if not reqs:
+            return []
+        if all(isinstance(req, Req) for req in reqs):
+            # AsyncSchedulerClient sends ordinary single requests as [Req].
+            # Only multi-item list[Req] payloads represent a grouped multi-output request.
+            if len(reqs) == 1:
+                return [(identity, reqs[0])]
+            return [(identity, reqs)]
+        return [(identity, req) for req in reqs]
+
     def recv_reqs(self) -> List[tuple[bytes, Any]]:
         """
         For non-main schedulers, reqs are broadcasted from main using broadcast_pyobj
         """
         if self.receiver is not None:
             try:
-                try:
-                    # Accept valid REQ envelopes only, ignore malformed/probe frames.
-                    parts = self.receiver.recv_multipart(zmq.NOBLOCK)
-                    identity, payload = parts[0], parts[-1]
+                recv_reqs: list[tuple[bytes, Any]] = []
+                while len(recv_reqs) < _MAX_RECV_REQS_PER_POLL:
+                    try:
+                        # Accept valid REQ envelopes only, ignore malformed/probe frames.
+                        parts = self.receiver.recv_multipart(zmq.NOBLOCK)
+                    except zmq.Again:
+                        break
 
-                    # Ignore malformed probes or non-pickle data
-                    recv_reqs = pickle.loads(payload) if len(parts) > 2 else []
-                except (zmq.Again, pickle.UnpicklingError, IndexError, EOFError):
-                    recv_reqs = []
+                    try:
+                        identity, payload = parts[0], parts[-1]
+                        reqs = pickle.loads(payload) if len(parts) > 2 else []
+                    except (pickle.UnpicklingError, IndexError, EOFError):
+                        continue
+
+                    recv_reqs.extend(self._normalize_received_payload(identity, reqs))
             except zmq.ZMQError:
                 # re-raise or handle appropriately to let the outer loop continue
                 raise
-
-            if recv_reqs:
-                if isinstance(recv_reqs, list) and all(
-                    isinstance(req, Req) for req in recv_reqs
-                ):
-                    recv_reqs = [(identity, recv_reqs)]
-                else:
-                    if not isinstance(recv_reqs, list):
-                        recv_reqs = [recv_reqs]
-                    recv_reqs = [(identity, req) for req in recv_reqs]
         else:
             recv_reqs = None
 
@@ -491,7 +1083,10 @@ class Scheduler(SchedulerDisaggMixin):
             try:
                 new_reqs = self.recv_reqs()
                 new_reqs = self.process_received_reqs_with_req_based_warmup(new_reqs)
-                self.waiting_queue.extend(new_reqs)
+                now = time.monotonic()
+                self.waiting_queue.extend(
+                    [(identity, req, now) for identity, req in new_reqs]
+                )
                 # Reset error count on success
                 self._consecutive_error_count = 0
             except Exception as e:
@@ -515,33 +1110,62 @@ class Scheduler(SchedulerDisaggMixin):
             # 2: execute, make sure a reply is always sent
             items = self.get_next_batch_to_run()
             if not items:
+                if self.waiting_queue and self._dynamic_batching_enabled():
+                    oldest_ts = self.waiting_queue[0][2]
+                    elapsed_ms = (time.monotonic() - oldest_ts) * 1000.0
+                    remaining_ms = max(0, self._batching_delay_s * 1000.0 - elapsed_ms)
+                    if remaining_ms > 0 and self.receiver is not None:
+                        self._poller.poll(timeout=remaining_ms)
+                    elif remaining_ms > 0:
+                        time.sleep(remaining_ms / 1000.0)
                 continue
 
-            identities = [item[0] for item in items]
-            reqs = [item[1] for item in items]
-
             try:
-                req_or_group = reqs[0]
-                is_warmup = self._is_warmup_item(req_or_group)
-                output_batch = self._dispatch_request(reqs)
+                handler_result = self._dispatch_items(items)
             except Exception as e:
                 logger.error(
                     f"Error executing request in scheduler event loop: {e}",
                     exc_info=True,
                 )
-                output_batch = OutputBatch(error=str(e))
+                handler_result = OutputBatch(error=str(e))
+
+            if isinstance(handler_result, list):
+                output_batches = handler_result
+            else:
+                output_batches = [handler_result]
+
+            if len(output_batches) != len(items):
+                logger.error(
+                    "Handler returned %d output(s) for %d request(s). Returning error for unmatched requests.",
+                    len(output_batches),
+                    len(items),
+                )
+                output_batches = [
+                    OutputBatch(
+                        error=(
+                            f"Internal scheduler error: expected {len(items)} outputs, "
+                            f"got {len(output_batches)}."
+                        )
+                    )
+                    for _ in items
+                ]
 
             # 3. return results
             try:
-                self._log_warmup_result(output_batch, is_warmup)
+                for (identity, processed_req), output_batch in zip(
+                    items, output_batches, strict=True
+                ):
+                    is_warmup = self._is_warmup_item(processed_req)
+                    self._log_warmup_result(output_batch, is_warmup)
 
-                # TODO: Support sending back to multiple identities if batched
-                self.return_result(output_batch, identities[0], is_warmup=is_warmup)
+                    self.return_result(output_batch, identity, is_warmup=is_warmup)
             except zmq.ZMQError as e:
                 # Reply failed; log and keep loop alive to accept future requests
                 logger.error(f"ZMQ error sending reply: {e}")
                 continue
 
+        self._log_batch_metrics_summary()
+
         if self.receiver is not None:
             self.receiver.close()
         self._cleanup_disagg()

python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py

@@ -642,10 +642,19 @@ class QwenImageCrossAttention(nn.Module):
         image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
         **cross_attention_kwargs,
     ):
+        """Run joint text-image attention.
+
+        `attn_mask` or `attention_mask` takes precedence. Otherwise,
+        `encoder_hidden_states_mask` keeps valid text tokens in the joint
+        text-image sequence.
+        """
         seq_len_txt = encoder_hidden_states.shape[1]
         attn_mask = cross_attention_kwargs.get("attn_mask")
         if attn_mask is None:
             attn_mask = cross_attention_kwargs.get("attention_mask")
+        encoder_hidden_states_mask = cross_attention_kwargs.get(
+            "encoder_hidden_states_mask"
+        )
 
         img_query, img_key, img_value, txt_query, txt_key, txt_value = (
             _get_qkv_projections(self, hidden_states, encoder_hidden_states)
@@ -704,6 +713,16 @@ class QwenImageCrossAttention(nn.Module):
         joint_query = torch.cat([txt_query, img_query], dim=1)
         joint_key = torch.cat([txt_key, img_key], dim=1)
         joint_value = torch.cat([txt_value, img_value], dim=1)
+        if attn_mask is None and encoder_hidden_states_mask is not None:
+            image_mask = torch.ones(
+                (hidden_states.shape[0], img_query.shape[1]),
+                device=encoder_hidden_states_mask.device,
+                dtype=torch.bool,
+            )
+            attn_mask = torch.cat(
+                [encoder_hidden_states_mask.to(dtype=torch.bool), image_mask],
+                dim=1,
+            )
 
         # Compute joint attention
         joint_hidden_states = self.attn(
@@ -1306,7 +1325,7 @@ class QwenImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
             encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_sequence_length, joint_attention_dim)`):
                 Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
             encoder_hidden_states_mask (`torch.Tensor` of shape `(batch_size, text_sequence_length)`):
-                Mask of the input conditions.
+                Valid-token mask of the input conditions, where True keeps a text token.
             timestep ( `torch.LongTensor`):
                 Used to indicate denoising step.
             attention_kwargs (`dict`, *optional*):

python/sglang/multimodal_gen/runtime/models/dits/zimage.py

@@ -780,65 +780,123 @@ class ZImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
         f_patch_size: int,
         image_seq_len_target: int | None = None,
     ):
-        assert len(all_image) == len(all_cap_feats) == 1
+        """Patchify images and pad image/caption tokens to batch targets.
+
+        Each image is [C, F, H, W] and has one [L, D] caption. Returned tensors
+        are stacked as [B, S, D], while valid lengths keep track of real tokens
+        before learned pad tokens are restored. `image_seq_len_target`, when
+        set, is the SP-local padded image-token target.
+        """
+        if len(all_image) != len(all_cap_feats):
+            raise ValueError(
+                f"Z-Image expects one caption embedding per image, got {len(all_image)} images and {len(all_cap_feats)} captions"
+            )
+        if not all_image:
+            raise ValueError("Z-Image batch must contain at least one image latent")
 
-        image = all_image[0]  # C, F, H, W
-        cap_feat = all_cap_feats[0]  # L, D
         pH = pW = patch_size
         pF = f_patch_size
-        device = image.device
         all_image_out = []
         all_image_size = []
         all_cap_feats_out = []
         all_image_valid_lens = []
         all_cap_valid_lens = []
+        image_records = []
 
-        # ------------ Process Caption ------------
-        cap_ori_len = cap_feat.size(0)
-        cap_padding_len = (-cap_ori_len) % SEQ_MULTI_OF
-
-        # padded feature
-        cap_padded_feat = torch.cat(
-            [cap_feat, cap_feat[-1:].repeat(cap_padding_len, 1)],
-            dim=0,
+        cap_seq_len_target = max(
+            self._ceil_to_multiple(cap_feat.size(0), SEQ_MULTI_OF)
+            for cap_feat in all_cap_feats
         )
-        all_cap_feats_out.append(cap_padded_feat)
-        all_cap_valid_lens.append(cap_ori_len)
 
-        # ------------ Process Image ------------
-        C, F, H, W = image.size()
-        all_image_size.append((F, H, W))
+        for cap_feat in all_cap_feats:
+            cap_ori_len = cap_feat.size(0)
+            cap_padding_len = cap_seq_len_target - cap_ori_len
+            cap_padded_feat = torch.cat(
+                [cap_feat, cap_feat[-1:].repeat(cap_padding_len, 1)],
+                dim=0,
+            )
+            all_cap_feats_out.append(cap_padded_feat)
+            all_cap_valid_lens.append(cap_ori_len)
 
-        F_tokens, H_tokens, W_tokens = F // pF, H // pH, W // pW
-        image = image.view(C, F_tokens, pF, H_tokens, pH, W_tokens, pW)
-        # "c f pf h ph w pw -> (f h w) (pf ph pw c)"
-        image = image.permute(1, 3, 5, 2, 4, 6, 0).reshape(
-            F_tokens * H_tokens * W_tokens, pF * pH * pW * C
-        )
-        image_ori_len = image.size(0)
-        min_image_seq_len = self._ceil_to_multiple(image_ori_len, SEQ_MULTI_OF)
-        if image_seq_len_target is None:
-            image_seq_len_target = min_image_seq_len
-        else:
-            image_seq_len_target = max(min_image_seq_len, image_seq_len_target)
-        image_padding_len = image_seq_len_target - image_ori_len
+        target_image_seq_len = image_seq_len_target or 0
+        for image in all_image:
+            # ------------ Process Image ------------
+            C, F, H, W = image.size()
+            image_size = (F, H, W)
 
-        # padded feature
-        image_padded_feat = torch.cat(
-            [image, image[-1:].repeat(image_padding_len, 1)],
-            dim=0,
-        )
-        all_image_out.append(image_padded_feat)
-        all_image_valid_lens.append(image_ori_len)
+            F_tokens, H_tokens, W_tokens = F // pF, H // pH, W // pW
+            image = image.view(C, F_tokens, pF, H_tokens, pH, W_tokens, pW)
+            # "c f pf h ph w pw -> (f h w) (pf ph pw c)"
+            image = image.permute(1, 3, 5, 2, 4, 6, 0).reshape(
+                F_tokens * H_tokens * W_tokens, pF * pH * pW * C
+            )
+            image_ori_len = image.size(0)
+            target_image_seq_len = max(
+                target_image_seq_len,
+                self._ceil_to_multiple(image_ori_len, SEQ_MULTI_OF),
+            )
+            image_records.append((image, image_size, image_ori_len))
+
+        for image, image_size, image_ori_len in image_records:
+            image_padding_len = target_image_seq_len - image_ori_len
+            image_padded_feat = torch.cat(
+                [image, image[-1:].repeat(image_padding_len, 1)],
+                dim=0,
+            )
+            all_image_out.append(image_padded_feat)
+            all_image_size.append(image_size)
+            all_image_valid_lens.append(image_ori_len)
 
         return (
-            all_image_out,
-            all_cap_feats_out,
+            torch.stack(all_image_out, dim=0),
+            torch.stack(all_cap_feats_out, dim=0),
             all_image_size,
             all_image_valid_lens,
             all_cap_valid_lens,
         )
 
+    @staticmethod
+    def _as_image_list(hidden_states) -> list[torch.Tensor]:
+        """Normalize 4D/5D image latents into per-sample tensors."""
+        if torch.is_tensor(hidden_states):
+            if hidden_states.dim() == 5:
+                return list(hidden_states.unbind(dim=0))
+            if hidden_states.dim() == 4:
+                return [hidden_states]
+        return list(hidden_states)
+
+    @staticmethod
+    def _as_caption_list(encoder_hidden_states) -> list[torch.Tensor]:
+        """Normalize caption tensors into per-sample tensors."""
+        if torch.is_tensor(encoder_hidden_states):
+            if encoder_hidden_states.dim() == 3:
+                return list(encoder_hidden_states.unbind(dim=0))
+            if encoder_hidden_states.dim() == 2:
+                return [encoder_hidden_states]
+
+        cap_feats = list(encoder_hidden_states)
+        if len(cap_feats) == 1 and torch.is_tensor(cap_feats[0]):
+            if cap_feats[0].dim() == 3:
+                return list(cap_feats[0].unbind(dim=0))
+            if cap_feats[0].dim() == 2:
+                return cap_feats
+        return cap_feats
+
+    @staticmethod
+    def _replace_padding_with_token(
+        tensor: torch.Tensor,
+        valid_lens: list[int],
+        pad_token: torch.Tensor,
+    ) -> torch.Tensor:
+        """Replace padded token rows after each valid sequence length."""
+        positions = torch.arange(tensor.shape[1], device=tensor.device).unsqueeze(0)
+        lengths = torch.tensor(valid_lens, device=tensor.device).unsqueeze(1)
+        pad_mask = positions >= lengths
+        if pad_mask.any():
+            tensor = tensor.clone()
+            tensor[pad_mask] = pad_token.to(device=tensor.device, dtype=tensor.dtype)
+        return tensor
+
     def forward(
         self,
         hidden_states: List[torch.Tensor],
@@ -854,14 +912,13 @@ class ZImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
         assert patch_size in self.all_patch_size
         assert f_patch_size in self.all_f_patch_size
 
-        x = hidden_states
-        cap_feats = encoder_hidden_states
+        x = self._as_image_list(hidden_states)
+        cap_feats = self._as_caption_list(encoder_hidden_states)
         timestep = 1000.0 - timestep
         t = timestep
-        bsz = 1
         device = x[0].device
         t = self.t_embedder(t)
-        adaln_input = t.type_as(x)
+        adaln_input = t.to(dtype=x[0].dtype)
         (
             x,
             cap_feats,
@@ -876,29 +933,20 @@ class ZImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
             image_seq_len_target=image_seq_len_target,
         )
 
-        x = torch.cat(x, dim=0)
         x, _ = self.all_x_embedder[f"{patch_size}-{f_patch_size}"](x)
-        if x_valid_lens[0] < x.shape[0]:
-            x[x_valid_lens[0] :] = self.x_pad_token.to(dtype=x.dtype)
+        x = self._replace_padding_with_token(x, x_valid_lens, self.x_pad_token)
         x_freqs_cis = freqs_cis[1]
 
-        x = x.unsqueeze(0)
-        x_freqs_cis = x_freqs_cis
         for layer_id, layer in enumerate(self.noise_refiner):
             x = layer(x, x_freqs_cis, adaln_input)
 
-        cap_feats = torch.cat(cap_feats, dim=0)
-
         cap_feats, _ = self.cap_embedder(cap_feats)
-        if cap_valid_lens[0] < cap_feats.shape[0]:
-            cap_feats[cap_valid_lens[0] :] = self.cap_pad_token.to(
-                dtype=cap_feats.dtype
-            )
+        cap_feats = self._replace_padding_with_token(
+            cap_feats, cap_valid_lens, self.cap_pad_token
+        )
 
         cap_freqs_cis = freqs_cis[0]
 
-        cap_feats = cap_feats.unsqueeze(0)
-        cap_input_dtype = cap_feats.dtype
         for layer_id, layer in enumerate(self.context_refiner):
             cap_feats = layer(
                 cap_feats,

python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py

@@ -14,6 +14,7 @@ from __future__ import annotations
 import logging
 import os
 import pprint
+from collections import Counter
 from copy import deepcopy
 from dataclasses import MISSING, asdict, dataclass, field, fields
 from typing import Any, Optional, Union
@@ -39,6 +40,23 @@ logger = init_logger(__name__)
 SAMPLING_PARAMS_FIELDS = {f.name for f in fields(SamplingParams)}
 
 
+@dataclass
+class BatchMetricsWindow:
+    """Counters accumulated between dynamic batching metric logs.
+
+    `total_capacity` uses each dispatch's effective admission cap, so
+    utilization reflects model/config limits instead of only the user max.
+    """
+
+    dispatches: int = 0
+    total_requests: int = 0
+    total_capacity: int = 0
+    merged_dispatches: int = 0
+    full_dispatches: int = 0
+    wait_times_ms: list[float] = field(default_factory=list)
+    reject_reasons: Counter[str] = field(default_factory=Counter)
+
+
 @dataclass(init=False)
 class Req:
     """
@@ -71,6 +89,11 @@ class Req:
     negative_prompt_embeds: list[torch.Tensor] | None = None
     prompt_attention_mask: list[torch.Tensor | None] | None = None
     negative_attention_mask: list[torch.Tensor | None] | None = None
+    # Masks and lengths aligned to postprocessed embeddings, one entry per text encoder.
+    prompt_embeds_mask: list[torch.Tensor | None] | None = None
+    negative_prompt_embeds_mask: list[torch.Tensor | None] | None = None
+    prompt_seq_lens: list[list[int]] | None = None
+    negative_prompt_seq_lens: list[list[int]] | None = None
     clip_embedding_pos: list[torch.Tensor] | None = None
     clip_embedding_neg: list[torch.Tensor] | None = None
 
@@ -269,6 +292,13 @@ class Req:
             return None
         return os.path.join(self.output_path, output_file_name)
 
+    @property
+    def resolution_key(self) -> str | None:
+        """Return the batching config resolution key, e.g. "1024x1024"."""
+        if self.width is None or self.height is None:
+            return None
+        return f"{int(self.width)}x{int(self.height)}"
+
     def set_as_warmup(self, warmup_steps: int = 1):
         self.is_warmup = True
         self.save_output = False

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py

@@ -170,10 +170,13 @@ class DmdDenoisingStage(DenoisingStage):
                                 **pos_cond_kwargs,
                             ).permute(0, 2, 1, 3, 4)
 
+                        video_timesteps = t_expand[:, None].expand(
+                            -1, pred_noise.shape[1]
+                        )
                         pred_video = pred_noise_to_pred_video(
                             pred_noise=pred_noise.flatten(0, 1),
                             noise_input_latent=noise_latents.flatten(0, 1),
-                            timestep=t_expand,
+                            timestep=video_timesteps,
                             scheduler=scheduler,
                         ).unflatten(0, pred_noise.shape[:2])
 

python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py

@@ -18,6 +18,7 @@ import torch
 from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
 from diffusers.models.modeling_outputs import AutoencoderKLOutput
 
+from sglang.multimodal_gen.configs.pipeline_configs.base import TextConditioningOutput
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
 from sglang.multimodal_gen.runtime.managers.component_manager import ComponentUse
 from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
@@ -116,6 +117,10 @@ class ImageEncodingStage(PipelineStage):
         "image_embeds",
         "prompt_embeds",
         "negative_prompt_embeds",
+        "prompt_embeds_mask",
+        "negative_prompt_embeds_mask",
+        "prompt_seq_lens",
+        "negative_prompt_seq_lens",
     )
 
     def __init__(
@@ -156,6 +161,30 @@ class ImageEncodingStage(PipelineStage):
 
         return postprocess_funcs[0](outputs, image_inputs)
 
+    @staticmethod
+    def _split_text_conditioning_output(output):
+        if isinstance(output, TextConditioningOutput):
+            return (
+                output.prompt_embeds,
+                output.prompt_embeds_mask,
+                output.prompt_seq_lens,
+            )
+        return output, None, None
+
+    @staticmethod
+    def _full_text_seq_lens(prompt_embeds: torch.Tensor) -> list[int]:
+        if prompt_embeds.ndim == 2:
+            return [int(prompt_embeds.shape[0])]
+        return [int(prompt_embeds.shape[1])] * int(prompt_embeds.shape[0])
+
+    @staticmethod
+    def _default_text_mask(prompt_embeds: torch.Tensor) -> torch.Tensor:
+        if prompt_embeds.ndim == 2:
+            shape = (1, prompt_embeds.shape[0])
+        else:
+            shape = prompt_embeds.shape[:2]
+        return torch.ones(shape, dtype=torch.bool, device=prompt_embeds.device)
+
     @torch.no_grad()
     def forward(
         self,
@@ -182,6 +211,10 @@ class ImageEncodingStage(PipelineStage):
 
         all_prompt_embeds = []
         all_neg_prompt_embeds = []
+        all_prompt_embeds_masks = []
+        all_neg_prompt_embeds_masks = []
+        all_prompt_seq_lens = []
+        all_neg_prompt_seq_lens = []
 
         image_processor_call_params = inspect.signature(
             self.image_processor.__call__
@@ -263,22 +296,84 @@ class ImageEncodingStage(PipelineStage):
                                 output_hidden_states=True,
                             )
 
-                all_prompt_embeds.append(
-                    self.encoding_image_edit(
-                        outputs, image_inputs, server_args.pipeline_config
+                prompt_embeds, prompt_embeds_mask, prompt_seq_lens = (
+                    self._split_text_conditioning_output(
+                        self.encoding_image_edit(
+                            outputs, image_inputs, server_args.pipeline_config
+                        )
                     )
                 )
+                all_prompt_embeds.append(prompt_embeds)
+                all_prompt_embeds_masks.append(prompt_embeds_mask)
+                all_prompt_seq_lens.extend(
+                    prompt_seq_lens
+                    if prompt_seq_lens is not None
+                    else self._full_text_seq_lens(prompt_embeds)
+                )
                 if batch.do_classifier_free_guidance:
-                    all_neg_prompt_embeds.append(
-                        self.encoding_image_edit(
-                            neg_outputs, neg_image_inputs, server_args.pipeline_config
+                    neg_prompt_embeds, neg_prompt_embeds_mask, neg_prompt_seq_lens = (
+                        self._split_text_conditioning_output(
+                            self.encoding_image_edit(
+                                neg_outputs,
+                                neg_image_inputs,
+                                server_args.pipeline_config,
+                            )
                         )
                     )
+                    all_neg_prompt_embeds.append(neg_prompt_embeds)
+                    all_neg_prompt_embeds_masks.append(neg_prompt_embeds_mask)
+                    all_neg_prompt_seq_lens.extend(
+                        neg_prompt_seq_lens
+                        if neg_prompt_seq_lens is not None
+                        else self._full_text_seq_lens(neg_prompt_embeds)
+                    )
 
         if all_prompt_embeds:
             batch.prompt_embeds.append(torch.cat(all_prompt_embeds, dim=0))
+            if batch.prompt_embeds_mask is None:
+                batch.prompt_embeds_mask = []
+            batch.prompt_embeds_mask.append(
+                torch.cat(
+                    [
+                        (
+                            mask
+                            if mask is not None
+                            else self._default_text_mask(prompt_embeds)
+                        )
+                        for prompt_embeds, mask in zip(
+                            all_prompt_embeds, all_prompt_embeds_masks, strict=True
+                        )
+                    ],
+                    dim=0,
+                )
+            )
+            if batch.prompt_seq_lens is None:
+                batch.prompt_seq_lens = []
+            batch.prompt_seq_lens.append(all_prompt_seq_lens)
         if all_neg_prompt_embeds:
             batch.negative_prompt_embeds.append(torch.cat(all_neg_prompt_embeds, dim=0))
+            if batch.negative_prompt_embeds_mask is None:
+                batch.negative_prompt_embeds_mask = []
+            batch.negative_prompt_embeds_mask.append(
+                torch.cat(
+                    [
+                        (
+                            mask
+                            if mask is not None
+                            else self._default_text_mask(neg_prompt_embeds)
+                        )
+                        for neg_prompt_embeds, mask in zip(
+                            all_neg_prompt_embeds,
+                            all_neg_prompt_embeds_masks,
+                            strict=True,
+                        )
+                    ],
+                    dim=0,
+                )
+            )
+            if batch.negative_prompt_seq_lens is None:
+                batch.negative_prompt_seq_lens = []
+            batch.negative_prompt_seq_lens.append(all_neg_prompt_seq_lens)
 
         return batch
 

python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py

@@ -68,12 +68,32 @@ class InputValidationStage(PipelineStage):
         return width, height
 
     def _generate_seeds(self, batch: Req, server_args: ServerArgs):
-        """Generate seeds for the inference"""
+        """Generate deterministic per-output seeds.
+
+        Batched requests pass one base seed per prompt through `extra`; each
+        prompt expands to `num_outputs_per_prompt` consecutive seeds.
+        """
         seed = batch.seed
         num_videos_per_prompt = batch.num_outputs_per_prompt
 
         assert seed is not None
-        if isinstance(seed, list):
+
+        prompt_count = len(batch.prompt) if isinstance(batch.prompt, list) else 1
+        dynamic_batch_seeds = batch.extra.get("dynamic_batch_seeds")
+
+        if dynamic_batch_seeds is not None:
+            if (
+                not isinstance(dynamic_batch_seeds, list)
+                or len(dynamic_batch_seeds) != prompt_count
+            ):
+                raise ValueError(
+                    "dynamic_batch_seeds must be a list with one seed per prompt"
+                )
+            base_seeds = [int(item) for item in dynamic_batch_seeds]
+            seeds = []
+            for base_seed in base_seeds:
+                seeds.extend([base_seed + i for i in range(num_videos_per_prompt)])
+        elif isinstance(seed, list):
             if len(seed) != num_videos_per_prompt:
                 raise ValueError(
                     f"seed list length must match num_outputs_per_prompt "
@@ -81,7 +101,13 @@ class InputValidationStage(PipelineStage):
                 )
             seeds = [int(item) for item in seed]
         else:
-            seeds = [int(seed) + i for i in range(num_videos_per_prompt)]
+            # Keep per-prompt seed streams deterministic and non-overlapping.
+            base_seeds = [
+                int(seed) + i * num_videos_per_prompt for i in range(prompt_count)
+            ]
+            seeds = []
+            for base_seed in base_seeds:
+                seeds.extend([base_seed + i for i in range(num_videos_per_prompt)])
         batch.seeds = seeds
 
         # Create generators based on generator_device parameter

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/qwen_image_layered.py

@@ -19,6 +19,15 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 logger = init_logger(__name__)
 
 
+def _seq_lens_from_optional_mask(
+    prompt_embeds: torch.Tensor, prompt_embeds_mask: torch.Tensor | None
+) -> list[int]:
+    """Return real text lengths, treating a missing mask as all tokens valid."""
+    if prompt_embeds_mask is None:
+        return [int(prompt_embeds.shape[1])] * int(prompt_embeds.shape[0])
+    return [int(x) for x in prompt_embeds_mask.sum(dim=1).tolist()]
+
+
 # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit_plus.calculate_dimensions
 def calculate_dimensions(target_area, ratio):
     width = math.sqrt(target_area * ratio)
@@ -521,22 +530,18 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
             mu=mu,
         )
 
-        txt_seq_lens = (
-            prompt_embeds_mask.sum(dim=1).tolist()
-            if prompt_embeds_mask is not None
-            else None
-        )
-        negative_txt_seq_lens = (
-            negative_prompt_embeds_mask.sum(dim=1).tolist()
-            if negative_prompt_embeds_mask is not None
-            else None
+        txt_seq_lens = _seq_lens_from_optional_mask(prompt_embeds, prompt_embeds_mask)
+        negative_txt_seq_lens = _seq_lens_from_optional_mask(
+            negative_prompt_embeds, negative_prompt_embeds_mask
         )
         is_rgb = torch.tensor([0]).to(device=device, dtype=torch.long)
 
         batch.prompt_embeds = [prompt_embeds]
-        batch.prompt_attention_mask = [prompt_embeds_mask]
+        batch.prompt_embeds_mask = [prompt_embeds_mask]
+        batch.prompt_seq_lens = [txt_seq_lens]
         batch.negative_prompt_embeds = [negative_prompt_embeds]
-        batch.negative_attention_mask = [negative_prompt_embeds_mask]
+        batch.negative_prompt_embeds_mask = [negative_prompt_embeds_mask]
+        batch.negative_prompt_seq_lens = [negative_txt_seq_lens]
         batch.latents = latents
         batch.image_latent = image_latents
         batch.timesteps = timesteps

python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py

@@ -14,6 +14,7 @@ from typing import Any
 import torch
 
 from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput
+from sglang.multimodal_gen.configs.pipeline_configs.base import TextConditioningOutput
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
 from sglang.multimodal_gen.runtime.managers.component_manager import ComponentUse
 from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
@@ -53,6 +54,10 @@ class TextEncodingStage(PipelineStage):
         "negative_prompt_embeds",
         "prompt_attention_mask",
         "negative_attention_mask",
+        "prompt_embeds_mask",
+        "negative_prompt_embeds_mask",
+        "prompt_seq_lens",
+        "negative_prompt_seq_lens",
         "pooled_embeds",
         "neg_pooled_embeds",
         "clip_embedding_pos",
@@ -102,7 +107,13 @@ class TextEncodingStage(PipelineStage):
 
         all_indices: list[int] = list(range(len(self.text_encoders)))
 
-        prompt_embeds_list, prompt_masks_list, pooler_embeds_list = self.encode_text(
+        (
+            prompt_embeds_list,
+            prompt_masks_list,
+            pooler_embeds_list,
+            prompt_embeds_masks_list,
+            prompt_seq_lens_list,
+        ) = self.encode_text(
             prompt_text,
             server_args,
             encoder_index=all_indices,
@@ -120,10 +131,23 @@ class TextEncodingStage(PipelineStage):
             for am in prompt_masks_list:
                 batch.prompt_attention_mask.append(am)
 
+        batch.prompt_embeds_mask = []
+        batch.prompt_seq_lens = []
+        for mask in prompt_embeds_masks_list:
+            batch.prompt_embeds_mask.append(mask)
+        for seq_lens in prompt_seq_lens_list:
+            batch.prompt_seq_lens.append(seq_lens)
+
         # Encode negative prompt if CFG is enabled
         if batch.do_classifier_free_guidance:
             assert isinstance(batch.negative_prompt, str)
-            neg_embeds_list, neg_masks_list, neg_pooler_embeds_list = self.encode_text(
+            (
+                neg_embeds_list,
+                neg_masks_list,
+                neg_pooler_embeds_list,
+                neg_embeds_masks_list,
+                neg_seq_lens_list,
+            ) = self.encode_text(
                 batch.negative_prompt,
                 server_args,
                 encoder_index=all_indices,
@@ -132,16 +156,71 @@ class TextEncodingStage(PipelineStage):
 
             assert batch.negative_prompt_embeds is not None
 
-            for ne in neg_embeds_list:
+            # A single negative prompt can be shared across positive prompts.
+            target_batch_sizes = [pe.shape[0] for pe in prompt_embeds_list]
+
+            def align_negative_batch_dim(
+                tensor: torch.Tensor, target_batch: int, name: str
+            ) -> torch.Tensor:
+                if tensor.shape[0] == target_batch:
+                    return tensor
+                if tensor.shape[0] == 1 and target_batch > 1:
+                    return tensor.expand(target_batch, *tensor.shape[1:])
+                raise ValueError(
+                    f"{name} batch dimension mismatch: got {tensor.shape[0]}, expected 1 or {target_batch}"
+                )
+
+            def align_negative_seq_lens(
+                seq_lens: list[int], target_batch: int, name: str
+            ) -> list[int]:
+                if len(seq_lens) == target_batch:
+                    return [int(x) for x in seq_lens]
+                if len(seq_lens) == 1 and target_batch > 1:
+                    return [int(seq_lens[0])] * target_batch
+                raise ValueError(
+                    f"{name} batch dimension mismatch: got {len(seq_lens)}, expected 1 or {target_batch}"
+                )
+
+            for idx, ne in enumerate(neg_embeds_list):
+                target_batch = target_batch_sizes[min(idx, len(target_batch_sizes) - 1)]
+                ne = align_negative_batch_dim(
+                    ne, target_batch, "negative_prompt_embeds"
+                )
                 batch.negative_prompt_embeds.append(ne)
 
-            for pe in neg_pooler_embeds_list:
+            for idx, pe in enumerate(neg_pooler_embeds_list):
+                target_batch = target_batch_sizes[min(idx, len(target_batch_sizes) - 1)]
+                pe = align_negative_batch_dim(
+                    pe, target_batch, "negative_pooled_embeds"
+                )
                 batch.neg_pooled_embeds.append(pe)
             if batch.negative_attention_mask is None:
                 batch.negative_attention_mask = []
-                for nm in neg_masks_list:
+                for idx, nm in enumerate(neg_masks_list):
+                    target_batch = target_batch_sizes[
+                        min(idx, len(target_batch_sizes) - 1)
+                    ]
+                    nm = align_negative_batch_dim(
+                        nm, target_batch, "negative_attention_mask"
+                    )
                     batch.negative_attention_mask.append(nm)
 
+            batch.negative_prompt_embeds_mask = []
+            batch.negative_prompt_seq_lens = []
+            for idx, nm in enumerate(neg_embeds_masks_list):
+                target_batch = target_batch_sizes[min(idx, len(target_batch_sizes) - 1)]
+                nm = align_negative_batch_dim(
+                    nm, target_batch, "negative_prompt_embeds_mask"
+                )
+                batch.negative_prompt_embeds_mask.append(nm)
+            for idx, seq_lens in enumerate(neg_seq_lens_list):
+                target_batch = target_batch_sizes[min(idx, len(target_batch_sizes) - 1)]
+                batch.negative_prompt_seq_lens.append(
+                    align_negative_seq_lens(
+                        seq_lens, target_batch, "negative_prompt_seq_lens"
+                    )
+                )
+
         return batch
 
     def build_dedup_fingerprint(
@@ -240,10 +319,14 @@ class TextEncodingStage(PipelineStage):
 
         Returns:
             Depending on return_type and return_attention_mask:
-            - list: List[Tensor] or (List[Tensor], List[Tensor])
+            - list: List[Tensor] or
+              (embeds, attention_masks, pooled_embeds, embeds_masks, seq_lens)
             - dict: Dict[str, Tensor] or (Dict[str, Tensor], Dict[str, Tensor])
             - stack: Tensor of shape [num_encoders, ...] or a tuple with stacked
               attention masks
+
+            `embeds_masks` and `seq_lens` are aligned with postprocessed
+            embeddings for variable-length text conditioning.
         """
 
         assert len(self.tokenizers) == len(self.text_encoders)
@@ -281,6 +364,8 @@ class TextEncodingStage(PipelineStage):
         pooled_embeds_list: list[torch.Tensor] = []
 
         attn_masks_list: list[torch.Tensor | None] = []
+        embeds_masks_list: list[torch.Tensor] = []
+        seq_lens_list: list[list[int]] = []
 
         preprocess_funcs = server_args.pipeline_config.preprocess_text_funcs
         postprocess_funcs = server_args.pipeline_config.postprocess_text_funcs
@@ -349,17 +434,34 @@ class TextEncodingStage(PipelineStage):
                 postprocess_kwargs["pipeline_config"] = server_args.pipeline_config
             if "return_attention_mask" in postprocess_sig.parameters:
                 postprocess_kwargs["return_attention_mask"] = return_attention_mask
-            prompt_embeds = postprocess_func(outputs, text_inputs, **postprocess_kwargs)
-            has_postprocessed_attention_mask = False
-            postprocessed_attention_mask = None
-            if isinstance(prompt_embeds, tuple):
-                prompt_embeds, postprocessed_attention_mask = prompt_embeds
-                has_postprocessed_attention_mask = True
+            postprocess_result = postprocess_func(
+                outputs, text_inputs, **postprocess_kwargs
+            )
+            prompt_embeds_mask = None
+            prompt_seq_lens = None
+            if isinstance(postprocess_result, TextConditioningOutput):
+                prompt_embeds = postprocess_result.prompt_embeds
+                prompt_embeds_mask = postprocess_result.prompt_embeds_mask
+                prompt_seq_lens = postprocess_result.prompt_seq_lens
+            elif isinstance(postprocess_result, tuple):
+                if len(postprocess_result) != 2:
+                    raise ValueError(
+                        "Text postprocess tuple output must be (prompt_embeds, prompt_embeds_mask)"
+                    )
+                prompt_embeds, prompt_embeds_mask = postprocess_result
+            else:
+                prompt_embeds = postprocess_result
+
             if dtype is not None:
                 prompt_embeds = prompt_embeds.to(device=target_device, dtype=dtype)
             else:
                 prompt_embeds = prompt_embeds.to(device=target_device)
 
+            if prompt_embeds_mask is not None:
+                prompt_embeds_mask = prompt_embeds_mask.to(
+                    device=target_device, dtype=torch.bool
+                )
+
             embeds_list.append(prompt_embeds)
 
             pooled_output = server_args.pipeline_config.get_text_encoder_pooler_output(
@@ -369,11 +471,14 @@ class TextEncodingStage(PipelineStage):
                 pooled_embeds_list.append(pooled_output.to(device=target_device))
 
             if return_attention_mask:
-                if has_postprocessed_attention_mask:
-                    mask_to_store = (
-                        postprocessed_attention_mask.to(device=target_device)
-                        if postprocessed_attention_mask is not None
-                        else None
+                if prompt_embeds_mask is not None:
+                    mask_to_store = prompt_embeds_mask.to(
+                        device=target_device,
+                        dtype=(
+                            attention_mask.dtype
+                            if attention_mask is not None
+                            else torch.long
+                        ),
                     )
                 elif attention_mask is not None and list(attention_mask.shape) == list(
                     prompt_embeds.shape[:2]
@@ -391,10 +496,42 @@ class TextEncodingStage(PipelineStage):
                     )
                 attn_masks_list.append(mask_to_store)
 
+                embeds_mask = prompt_embeds_mask
+                if embeds_mask is None:
+                    embeds_mask = (
+                        server_args.pipeline_config.build_text_conditioning_mask(
+                            text_inputs,
+                            attention_mask,
+                            prompt_embeds,
+                            i,
+                        )
+                    )
+                embeds_masks_list.append(embeds_mask)
+                if prompt_seq_lens is not None:
+                    seq_lens_list.append([int(x) for x in prompt_seq_lens])
+                elif embeds_mask is not None:
+                    seq_lens_list.append(
+                        server_args.pipeline_config.seq_lens_from_text_conditioning_mask(
+                            embeds_mask
+                        )
+                    )
+                elif prompt_embeds.ndim == 2:
+                    seq_lens_list.append([int(prompt_embeds.shape[0])])
+                else:
+                    seq_lens_list.append(
+                        [int(prompt_embeds.shape[1])] * int(prompt_embeds.shape[0])
+                    )
+
         # Shape results according to return_type
         if return_type == "list":
             if return_attention_mask:
-                return embeds_list, attn_masks_list, pooled_embeds_list
+                return (
+                    embeds_list,
+                    attn_masks_list,
+                    pooled_embeds_list,
+                    embeds_masks_list,
+                    seq_lens_list,
+                )
             return embeds_list, pooled_embeds_list
 
         if return_type == "dict":

python/sglang/multimodal_gen/runtime/server_args.py

@@ -225,6 +225,11 @@ class ServerArgs(DisaggArgsMixin):
     webui_port: int | None = 12312
 
     scheduler_port: int = 5555
+    batching_mode: str = "dynamic"
+    batching_max_size: int = 1
+    batching_delay_ms: float = 0.0
+    batching_config: str | None = None
+    enable_batching_metrics: bool = False
 
     # Strict port mode: fail if requested port is unavailable instead of auto-selecting
     strict_ports: bool = False
@@ -325,6 +330,7 @@ class ServerArgs(DisaggArgsMixin):
         if not current_platform.is_cpu():
             self._validate_parallelism()
         self._validate_cfg_parallel()
+        self._validate_batching()
 
     def _adjust_save_paths(self):
         """Normalize empty-string save paths to None (disabled)."""
@@ -1020,6 +1026,41 @@ class ServerArgs(DisaggArgsMixin):
             default=ServerArgs.scheduler_port,
             help="Port for the scheduler server.",
         )
+        parser.add_argument(
+            "--batching-mode",
+            type=str,
+            default=ServerArgs.batching_mode,
+            choices=["dynamic"],
+            help="Request batching scheduler mode. Currently only 'dynamic' is implemented.",
+        )
+        parser.add_argument(
+            "--batching-max-size",
+            type=int,
+            default=ServerArgs.batching_max_size,
+            help="Maximum number of compatible generation requests to merge into one batch.",
+        )
+        parser.add_argument(
+            "--batching-delay-ms",
+            type=float,
+            default=ServerArgs.batching_delay_ms,
+            help="Maximum time (in ms) to wait for forming a larger batch before dispatch.",
+        )
+        parser.add_argument(
+            "--batching-config",
+            type=str,
+            default=ServerArgs.batching_config,
+            help=(
+                "Optional JSON file with {'schema_version': 1, 'rules': [...]} "
+                "batching admission rules that can cap model/resolution shapes "
+                "below --batching-max-size."
+            ),
+        )
+        parser.add_argument(
+            "--enable-batching-metrics",
+            action="store_true",
+            default=ServerArgs.enable_batching_metrics,
+            help="Log periodic batch efficiency metrics such as realized batch size and queue wait time.",
+        )
         parser.add_argument(
             "--host",
             type=str,
@@ -1461,6 +1502,14 @@ class ServerArgs(DisaggArgsMixin):
                 "CFG Parallelism is enabled via `--enable-cfg-parallel`, but num_gpus == 1"
             )
 
+    def _validate_batching(self):
+        if self.batching_mode != "dynamic":
+            raise ValueError("batching_mode must be one of: dynamic")
+        if self.batching_max_size < 1:
+            raise ValueError("batching_max_size must be >= 1")
+        if self.batching_delay_ms < 0:
+            raise ValueError("batching_delay_ms must be >= 0")
+
     def _set_default_attention_backend(self) -> None:
         """Configure ROCm defaults when users do not specify an attention backend."""
         if current_platform.is_rocm():

python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py

@@ -81,7 +81,7 @@ class TestWarmupReqCfgParallel(unittest.TestCase):
         scheduler.prepare_server_warmup_reqs()
 
         self.assertEqual(len(scheduler.waiting_queue), 1)
-        _, req = scheduler.waiting_queue[0]
+        _, req, _ = scheduler.waiting_queue[0]
         self.assertIs(req.do_classifier_free_guidance, True)
         self.assertEqual(req.negative_prompt, DEFAULT_PLACEHOLDER_PROMPT)
 
@@ -95,7 +95,7 @@ class TestWarmupReqCfgParallel(unittest.TestCase):
         scheduler.prepare_server_warmup_reqs()
 
         self.assertEqual(len(scheduler.waiting_queue), 1)
-        _, req = scheduler.waiting_queue[0]
+        _, req, _ = scheduler.waiting_queue[0]
         self.assertIs(req.do_classifier_free_guidance, False)
         self.assertNotEqual(req.negative_prompt, DEFAULT_PLACEHOLDER_PROMPT)