diff --git a/docs_new/docs.json b/docs_new/docs.json index cd71ed913..7b150e01d 100644 --- a/docs_new/docs.json +++ b/docs_new/docs.json @@ -364,6 +364,10 @@ "source": "/diffusion/performance/attention_backends.html", "destination": "/docs/sglang-diffusion/attention_backends" }, + { + "source": "/diffusion/performance/dynamic_batching.html", + "destination": "/docs/sglang-diffusion/dynamic_batching" + }, { "source": "/diffusion/performance/cache/cache_dit.html", "destination": "/docs/sglang-diffusion/cache_dit" @@ -1170,6 +1174,12 @@ "docs/sglang-diffusion/performance-optimization", "docs/sglang-diffusion/ring_sp_performance", "docs/sglang-diffusion/attention_backends", + { + "group": "Inference Batching", + "pages": [ + "docs/sglang-diffusion/dynamic_batching" + ] + }, "docs/sglang-diffusion/profiling", "docs/sglang-diffusion/ci_perf" ] diff --git a/docs_new/docs/sglang-diffusion/dynamic_batching.mdx b/docs_new/docs/sglang-diffusion/dynamic_batching.mdx new file mode 100644 index 000000000..b05a6eb89 --- /dev/null +++ b/docs_new/docs/sglang-diffusion/dynamic_batching.mdx @@ -0,0 +1,143 @@ +--- +title: "Inference Batching" +description: "Batch compatible native SGLang-Diffusion requests during serving." +mode: wide +--- +Dynamic batching is an opt-in SGLang-Diffusion serving mode that merges compatible queued requests into one native pipeline batch. It is separate from LLM continuous batching and tokenizer batching. + +Use it for concurrent T2I or T2V traffic with the same model and sampling shape. Keep singleton serving for latency-sensitive or highly mixed traffic. + +## Enable + +Dynamic batching is disabled by default with `--batching-max-size 1`. + +```bash Command +sglang serve \ + --model-path black-forest-labs/FLUX.1-dev \ + --port 30010 \ + --batching-mode dynamic \ + --batching-max-size 8 \ + --batching-delay-ms 5 \ + --enable-batching-metrics +``` + +For request formats, see the [OpenAI-Compatible API](./api/openai_api). + +Use `--batching-config /path/to/batching_config.json` to load JSON rules when a model or resolution needs a lower cap than `--batching-max-size`: + +```json Config +{ + "schema_version": 1, + "rules": [ + { + "model_contains": "Qwen-Image", + "resolution": "1024x1024", + "max_batch_size": 1 + } + ] +} +``` + +## Compatibility + +An initial implementation of dynamic batching for T2I and T2V models can be found in [#18764](https://github.com/sgl-project/sglang/pull/18764). The current compatibility grid is below and will be updated as more coverage is added. See [Supported Models](./compatibility_matrix) for full model IDs. + +`✅` means supported, `❌` means not currently supported, `?` means untested, and `-` means not applicable. + +### Image + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
ModelT2II2I
FLUX.1-dev✅-
FLUX.2-dev✅❌
FLUX.2-dev-NVFP4??
FLUX.2-Klein-4B✅❌
FLUX.2-Klein-9B??
Z-Image?-
Z-Image-Turbo✅-
GLM-Image❌-
Qwen Image✅-
Qwen Image 2512✅-
Qwen Image Edit-❌
Qwen Image Edit 2509-?
Qwen Image Edit 2511-?
Qwen Image Layered??
SD3 Medium?-
SD3.5 Medium?-
SD3.5 Large?-
Hunyuan3D-2?-
SANA 1.5 1.6B✅-
SANA 1.5 4.8B✅-
SANA 1600M 1024px?-
SANA 600M 1024px?-
SANA 1600M 512px?-
SANA 600M 512px?-
FireRed-Image-Edit 1.0-?
FireRed-Image-Edit 1.1-?
ERNIE-Image?-
ERNIE-Image-Turbo?-
+ +### Video + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
ModelSupport
FastWan2.1 T2V 1.3B✅
FastWan2.2 TI2V 5B Full Attn❌
Wan2.2 TI2V 5B❌
Wan2.2 T2V A14B✅
Wan2.2 I2V A14B❌
HunyuanVideo❌
FastHunyuan❌
Wan2.1 T2V 1.3B✅
Wan2.1 T2V 14B✅
Wan2.1 I2V 480P?
Wan2.1 I2V 720P?
TurboWan2.1 T2V 1.3B✅
TurboWan2.1 T2V 14B✅
TurboWan2.1 T2V 14B 720P✅
TurboWan2.2 I2V A14B?
Wan2.1 Fun 1.3B InP?
Helios Base?
Helios Mid?
Helios Distilled?
LTX-2?
LTX-2.3?
+ +## Notes + +- Requests batch only when model inputs, sampling parameters, output handling, and any configured rules are compatible. +- There is no startup probing, runtime learning, OOM retry, or automatic fallback to singletons. If a merged batch fails or cannot be split, every request in that batch receives an error. +- Batch shape can change kernels, so singleton and dynamic outputs are not expected to be bit-exact. +- Use `--enable-batching-metrics` to inspect realized batches: + +```text +Dynamic batch dispatch: size=2/8, user_max=8, queue_wait=5.12ms, stop_reason=delay +Dynamic batch dispatch: size=1/8, user_max=8, queue_wait=0.04ms, stop_reason=config_cap:1 +Dynamic batch stats (last 5 dispatches): avg_size=2.80, merged_rate=60.0%, full_rate=20.0%, utilization=35.0%, wait_avg=3.21ms, wait_p95=5.12ms, top_rejects=none +``` diff --git a/docs_new/docs/sglang-diffusion/index.mdx b/docs_new/docs/sglang-diffusion/index.mdx index d473d1891..42a9ffbe1 100644 --- a/docs_new/docs/sglang-diffusion/index.mdx +++ b/docs_new/docs/sglang-diffusion/index.mdx @@ -34,6 +34,7 @@ sglang serve --model-path Qwen/Qwen-Image --port 30010 - [CLI](/docs/sglang-diffusion/api/cli): run one-off generation jobs or launch a persistent server - [OpenAI-Compatible API](/docs/sglang-diffusion/api/openai_api): send image and video requests to the HTTP server - [Attention Backends](/docs/sglang-diffusion/attention_backends): choose the best backend for your model and hardware +- [Inference Batching](/docs/sglang-diffusion/dynamic_batching): batch compatible native diffusion requests during serving - [Caching Acceleration](/docs/sglang-diffusion/caching-acceleration): use Cache-DiT or TeaCache to reduce denoising cost - [Quantization](/docs/sglang-diffusion/quantization): load quantized transformer checkpoints - [Contributing](/docs/sglang-diffusion/contributing): contribution workflow, adding new models, and CI perf baselines diff --git a/docs_new/docs/sglang-diffusion/performance-optimization.mdx b/docs_new/docs/sglang-diffusion/performance-optimization.mdx index bb5c15283..ef5a8950f 100644 --- a/docs_new/docs/sglang-diffusion/performance-optimization.mdx +++ b/docs_new/docs/sglang-diffusion/performance-optimization.mdx @@ -35,6 +35,11 @@ This section covers the main performance levers for SGLang Diffusion: attention Kernel Optimized attention implementations (FlashAttention, SageAttention, etc.) + + Inference Batching + Scheduler + Request batching for native diffusion serving + Profiling Diagnostics @@ -46,6 +51,7 @@ This section covers the main performance levers for SGLang Diffusion: attention ## Start Here - Use [Attention Backends](./attention_backends) to choose the best backend for your model and hardware. +- Use [Inference Batching](./dynamic_batching) to improve throughput for compatible concurrent requests. - Use [Caching Acceleration](./caching-acceleration) to reduce denoising cost with Cache-DiT or TeaCache. - Use [Profiling](./profiling) when you need to diagnose a bottleneck rather than guess. diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index da652a7ae..e3be31100 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/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": diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py b/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py index eed4f1166..b6dffd597 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py +++ b/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, ) } diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/hunyuan.py b/python/sglang/multimodal_gen/configs/pipeline_configs/hunyuan.py index abe421c89..7eba6e057 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/hunyuan.py +++ b/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 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 4fc40558d..d2b9c5a62 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py +++ b/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 diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/sana.py b/python/sglang/multimodal_gen/configs/pipeline_configs/sana.py index d008c0c62..234bc0066 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/sana.py +++ b/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,), ) diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py b/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py index 2bfc46d14..d26276863 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py +++ b/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"] diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index a5dad1fd7..5c8ea67ac 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py index 17028c5ff..19cd5ea09 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py index 3ba074181..bdc2d2fae 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py +++ b/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") diff --git a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/base.py b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/base.py index 6c42ca5d0..2e6088ee6 100644 --- a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/base.py +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/managers/dynamic_batch_admission.py b/python/sglang/multimodal_gen/runtime/managers/dynamic_batch_admission.py new file mode 100644 index 000000000..6d9153c75 --- /dev/null +++ b/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}") diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 765bfcbfc..5f667d892 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/managers/scheduler.py b/python/sglang/multimodal_gen/runtime/managers/scheduler.py index ce60cde55..07cf3dea5 100644 --- a/python/sglang/multimodal_gen/runtime/managers/scheduler.py +++ b/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() diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py index 5d78d0dd0..c5fe1bf12 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/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*): diff --git a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py index 5a95f5952..4a60cc516 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py index 4183fbcf3..e9a17c6e2 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py index 08957f1e0..f84f9fae3 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py +++ b/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]) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py index 7a7de8c4b..fdd9ad0d6 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py index f8ca29c44..bdd1cbc2c 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/qwen_image_layered.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/qwen_image_layered.py index 5bd25f31e..9f723d495 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/qwen_image_layered.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py index a8cb97ce6..be72efdae 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py +++ b/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": diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index cdc82c78f..8b83a02c6 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/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(): diff --git a/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py b/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py index dec8daee7..fbb6fce7b 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py +++ b/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)