diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index fb5fb0db1..1e90ca3a8 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py @@ -675,6 +675,10 @@ class PipelineConfig: def get_neg_prompt_embeds(self, batch): return batch.negative_prompt_embeds + def expand_conditioning_to_sample_batch(self, batch): + """Used for single-request multi-output generation case.""" + return batch + def post_denoising_loop(self, latents, batch): latents = maybe_unpad_latents(latents, batch) return latents diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py b/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py index 1bfb41e96..610267b7c 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py @@ -22,6 +22,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import ( from sglang.multimodal_gen.configs.post_training.pipeline_configs import ( QwenImageRolloutPipelineMixin, ) +from sglang.multimodal_gen.runtime.utils.condition_expansion import ( + PromptToSampleBatchExpander, +) from sglang.multimodal_gen.runtime.utils.vision import resize from sglang.multimodal_gen.utils import calculate_dimensions @@ -181,6 +184,24 @@ class QwenImagePipelineConfig(QwenImageRolloutPipelineMixin, ImagePipelineConfig ] ) + def expand_conditioning_to_sample_batch(self, batch): + expander = PromptToSampleBatchExpander.from_batch(batch) + if expander is None: + return batch + + for field_name in ( + "prompt_embeds", + "negative_prompt_embeds", + "prompt_attention_mask", + "negative_attention_mask", + "prompt_embeds_mask", + "negative_prompt_embeds_mask", + "prompt_seq_lens", + "negative_prompt_seq_lens", + ): + expander.expand_field(batch, field_name) + return batch + def tokenize_prompt(self, prompts: list[str], tokenizer, tok_kwargs) -> dict: tok_kwargs.setdefault("truncation", True) diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/wan.py b/python/sglang/multimodal_gen/configs/pipeline_configs/wan.py index 971a7b781..f8daf9bfc 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/wan.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/wan.py @@ -21,6 +21,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import ( from sglang.multimodal_gen.configs.pipeline_configs.model_deployment_config import ( ModelDeploymentConfig, ) +from sglang.multimodal_gen.runtime.utils.condition_expansion import ( + PromptToSampleBatchExpander, +) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) @@ -97,6 +100,20 @@ class WanT2V480PConfig(PipelineConfig): auto_dit_layerwise_offload=True, ) + def expand_conditioning_to_sample_batch(self, batch): + expander = PromptToSampleBatchExpander.from_batch(batch) + if expander is None: + return batch + + for field_name in ( + "prompt_embeds", + "negative_prompt_embeds", + "image_embeds", + "image_latent", + ): + expander.expand_field(batch, field_name) + return batch + def get_pos_prompt_embeds(self, batch): return batch.prompt_embeds[0] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 1df3b172e..c00cca589 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -754,6 +754,8 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin): Returns: A context object containing the invariant state for the denoising loop. """ + batch = server_args.pipeline_config.expand_conditioning_to_sample_batch(batch) + assert self.transformer is not None pipeline = self.pipeline() if self.pipeline else None scheduler = batch.scheduler diff --git a/python/sglang/multimodal_gen/runtime/utils/condition_expansion.py b/python/sglang/multimodal_gen/runtime/utils/condition_expansion.py new file mode 100644 index 000000000..2fcf185d1 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/utils/condition_expansion.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True) +class PromptToSampleBatchExpander: + """Expand selected conditioning from prompt order to sample order.""" + + prompt_batch_size: int + sample_batch_size: int + + @classmethod + def from_batch(cls, batch): + num_outputs = int(batch.num_outputs_per_prompt or 1) + if num_outputs <= 1: + return None + if isinstance(batch.prompt, list): + prompt_batch_size = len(batch.prompt) + elif batch.prompt is not None: + prompt_batch_size = 1 + else: + raise ValueError( + "Multi-output conditioning requires prompt text so the prompt " + "batch size is unambiguous." + ) + if prompt_batch_size <= 0: + raise ValueError("Multi-output conditioning requires at least one prompt.") + return cls(prompt_batch_size, prompt_batch_size * num_outputs) + + def _expand_tensor(self, value: torch.Tensor, name: str) -> torch.Tensor: + current_batch_size = value.shape[0] + if current_batch_size == self.sample_batch_size: + return value + if current_batch_size != self.prompt_batch_size: + raise ValueError( + f"{name} has batch dim {current_batch_size} (shape " + f"{tuple(value.shape)}); expected {self.prompt_batch_size} " + f"(per-prompt) or {self.sample_batch_size} (per-sample)." + ) + repeats = self.sample_batch_size // self.prompt_batch_size + return value.repeat_interleave(repeats, dim=0) + + def _expand_tensors(self, value, name: str): + """Expand a tensor or each tensor in a list, preserving its container.""" + if value is None: + return None + if isinstance(value, torch.Tensor): + return self._expand_tensor(value, name) + if not isinstance(value, list): + raise TypeError(f"{name} must be a tensor, list of tensors, or None.") + if any( + item is not None and not isinstance(item, torch.Tensor) for item in value + ): + raise TypeError(f"{name} entries must be tensors or None.") + return [ + self._expand_tensor(item, f"{name}[{index}]") if item is not None else None + for index, item in enumerate(value) + ] + + def _expand_sequence_lengths( + self, value: list[list[int] | None] | None, name: str + ) -> list[list[int] | None] | None: + if value is None: + return None + repeats = self.sample_batch_size // self.prompt_batch_size + expanded = [] + for index, sequence_lengths in enumerate(value): + if ( + sequence_lengths is None + or len(sequence_lengths) == self.sample_batch_size + ): + expanded.append(sequence_lengths) + elif len(sequence_lengths) == self.prompt_batch_size: + expanded.append( + [ + sequence_length + for sequence_length in sequence_lengths + for _ in range(repeats) + ] + ) + else: + raise ValueError( + f"{name}[{index}] has {len(sequence_lengths)} entries; expected " + f"{self.prompt_batch_size} (per-prompt) or " + f"{self.sample_batch_size} (per-sample)." + ) + return expanded + + def expand_field(self, batch, field_name: str) -> None: + """Expand one field in place, dispatching from its value type.""" + value = getattr(batch, field_name) + if value is None: + return + if isinstance(value, torch.Tensor) or ( + isinstance(value, list) + and all(item is None or isinstance(item, torch.Tensor) for item in value) + ): + expanded = self._expand_tensors(value, field_name) + elif isinstance(value, list) and all( + item is None or isinstance(item, list) for item in value + ): + expanded = self._expand_sequence_lengths(value, field_name) + else: + raise TypeError( + f"{field_name} must be a tensor, list of tensors, " + "list of sequence-length lists, or None." + ) + setattr(batch, field_name, expanded)