minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Chi McIsaac <chixie.mcisaac@gmail.com>
Co-authored-by: Junhao Liu <junhaoliu2023@gmail.com>
This commit is contained in:
Chi McIsaac
2026-05-04 00:44:42 +08:00
committed by GitHub
co-authored by Junhao Liu
parent f2d1390909
commit 62265ca7fc
27 changed files with 2318 additions and 245 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs_new/docs.json
+10
View File
@@ -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"
]
docs_new/docs/sglang-diffusion/dynamic_batching.mdx
+143
View File
@@ -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
<table style={{display: "table", width: "100%", borderCollapse: "collapse", tableLayout: "fixed"}}>
<colgroup>
<col style={{width: "60%"}} />
<col style={{width: "20%"}} />
<col style={{width: "20%"}} />
</colgroup>
<thead>
<tr style={{borderBottom: "2px solid #d55816"}}>
<th style={{width: "60%", textAlign: "left", padding: "10px 12px", fontWeight: 700, whiteSpace: "nowrap", backgroundColor: "rgba(255,255,255,0.02)"}}>Model</th>
<th style={{width: "20%", textAlign: "center", padding: "10px 12px", fontWeight: 700, whiteSpace: "nowrap", backgroundColor: "rgba(255,255,255,0.05)"}}>T2I</th>
<th style={{width: "20%", textAlign: "center", padding: "10px 12px", fontWeight: 700, whiteSpace: "nowrap", backgroundColor: "rgba(255,255,255,0.02)"}}>I2I</th>
</tr>
</thead>
<tbody>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>FLUX.1-dev</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>✅</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>-</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>FLUX.2-dev</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>✅</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>❌</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>FLUX.2-dev-NVFP4</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>?</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>?</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>FLUX.2-Klein-4B</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>✅</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>❌</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>FLUX.2-Klein-9B</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>?</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>?</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Z-Image</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>?</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>-</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Z-Image-Turbo</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>✅</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>-</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>GLM-Image</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>-</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Qwen Image</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>✅</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>-</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Qwen Image 2512</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>✅</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>-</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Qwen Image Edit</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>-</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>❌</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Qwen Image Edit 2509</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>-</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>?</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Qwen Image Edit 2511</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>-</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>?</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Qwen Image Layered</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>?</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>?</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>SD3 Medium</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>?</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>-</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>SD3.5 Medium</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>?</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>-</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>SD3.5 Large</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>?</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>-</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Hunyuan3D-2</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>?</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>-</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>SANA 1.5 1.6B</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>✅</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>-</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>SANA 1.5 4.8B</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>✅</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>-</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>SANA 1600M 1024px</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>?</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>-</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>SANA 600M 1024px</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>?</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>-</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>SANA 1600M 512px</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>?</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>-</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>SANA 600M 512px</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>?</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>-</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>FireRed-Image-Edit 1.0</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>-</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>?</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>FireRed-Image-Edit 1.1</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>-</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>?</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>ERNIE-Image</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>?</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>-</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>ERNIE-Image-Turbo</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>?</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>-</td></tr>
</tbody>
</table>
### Video
<table style={{display: "table", width: "100%", borderCollapse: "collapse", tableLayout: "fixed"}}>
<colgroup>
<col style={{width: "74%"}} />
<col style={{width: "26%"}} />
</colgroup>
<thead>
<tr style={{borderBottom: "2px solid #d55816"}}>
<th style={{width: "65%", textAlign: "left", padding: "10px 12px", fontWeight: 700, whiteSpace: "nowrap", backgroundColor: "rgba(255,255,255,0.02)"}}>Model</th>
<th style={{width: "35%", textAlign: "center", padding: "10px 12px", fontWeight: 700, whiteSpace: "nowrap", backgroundColor: "rgba(255,255,255,0.05)"}}>Support</th>
</tr>
</thead>
<tbody>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>FastWan2.1 T2V 1.3B</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>✅</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>FastWan2.2 TI2V 5B Full Attn</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Wan2.2 TI2V 5B</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Wan2.2 T2V A14B</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>✅</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Wan2.2 I2V A14B</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>HunyuanVideo</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>FastHunyuan</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Wan2.1 T2V 1.3B</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>✅</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Wan2.1 T2V 14B</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>✅</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Wan2.1 I2V 480P</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>?</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Wan2.1 I2V 720P</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>?</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>TurboWan2.1 T2V 1.3B</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>✅</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>TurboWan2.1 T2V 14B</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>✅</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>TurboWan2.1 T2V 14B 720P</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>✅</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>TurboWan2.2 I2V A14B</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>?</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Wan2.1 Fun 1.3B InP</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>?</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Helios Base</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>?</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Helios Mid</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>?</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Helios Distilled</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>?</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>LTX-2</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>?</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>LTX-2.3</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>?</td></tr>
</tbody>
</table>
## 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
```
docs_new/docs/sglang-diffusion/index.mdx
+1
View File
@@ -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
docs_new/docs/sglang-diffusion/performance-optimization.mdx
+6
View File
@@ -35,6 +35,11 @@ This section covers the main performance levers for SGLang Diffusion: attention
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Kernel</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Optimized attention implementations (FlashAttention, SageAttention, etc.)</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Inference Batching</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Scheduler</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Request batching for native diffusion serving</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Profiling</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Diagnostics</td>
@@ -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.
python/sglang/multimodal_gen/configs/pipeline_configs/base.py
+160
View File
@@ -2,6 +2,7 @@
# SPDX-License-Identifier: Apache-2.0
import json
import math
import os
from collections.abc import Callable
from dataclasses import asdict, dataclass, field, fields
@@ -101,6 +102,44 @@ def postprocess_text(output: BaseEncoderOutput, _text_inputs) -> torch.tensor:
raise NotImplementedError
@dataclass(frozen=True)
class TextConditioningOutput:
"""Text embeddings and masks aligned to postprocessed sequence length.
`prompt_embeds_mask` and `prompt_seq_lens` describe real text tokens after
model-specific trimming or packing, not the raw tokenizer output.
"""
prompt_embeds: torch.Tensor
prompt_embeds_mask: torch.Tensor | None = None
prompt_seq_lens: list[int] | None = None
def pad_text_embeddings_with_mask(
text_embeds: list[torch.Tensor],
) -> TextConditioningOutput:
"""Pad variable-length text embeddings and return the valid-token mask."""
if not text_embeds:
raise ValueError("text_embeds must contain at least one tensor")
max_seq_len = max(e.size(0) for e in text_embeds)
prompt_embeds = torch.stack(
[
torch.cat([e, e.new_zeros(max_seq_len - e.size(0), e.size(1))])
for e in text_embeds
]
)
seq_lens = [int(e.size(0)) for e in text_embeds]
seq_lens_tensor = torch.tensor(
seq_lens,
device=prompt_embeds.device,
dtype=torch.long,
)
positions = torch.arange(max_seq_len, device=prompt_embeds.device).unsqueeze(0)
prompt_embeds_mask = positions < seq_lens_tensor.unsqueeze(1)
return TextConditioningOutput(prompt_embeds, prompt_embeds_mask, seq_lens)
def shard_rotary_emb_for_sp(emb):
"""
Shard rotary embeddings [S, D] along sequence for SP.
@@ -329,6 +368,41 @@ class PipelineConfig:
def allow_set_num_frames(self):
return False
def supports_dynamic_batching(self):
"""Return whether this pipeline can opt in to dynamic batching.
The scheduler still checks each request before merging it into a batch.
"""
return self.task_type in (ModelTaskType.T2I, ModelTaskType.T2V)
def estimate_request_cost(self, batch) -> float:
"""Return the relative cost used for batching admission caps.
This is compared with `max_cost` from the batching config; it is not a
memory estimate. The default cost is latent tokens times frames times
outputs; pipelines can override it for model-specific admission.
"""
latent_tokens = float(batch.n_tokens or 0)
if latent_tokens <= 0:
width = int(batch.width or 0)
height = int(batch.height or 0)
if width > 0 and height > 0:
vae_scale = getattr(
self.vae_config.arch_config, "vae_scale_factor", None
)
if vae_scale is None and hasattr(
self.vae_config, "get_vae_scale_factor"
):
vae_scale = self.vae_config.get_vae_scale_factor()
vae_scale = max(1, int(vae_scale or 1))
latent_tokens = math.ceil(width / vae_scale) * math.ceil(
height / vae_scale
)
num_frames = max(1, int(batch.num_frames or 1))
num_outputs = max(1, int(batch.num_outputs_per_prompt or 1))
return latent_tokens * num_frames * num_outputs
def get_decode_scale_and_shift(self, device, dtype, vae):
vae_arch_config = self.vae_config.arch_config
scaling_factor = getattr(vae_arch_config, "scaling_factor", None)
@@ -468,6 +542,92 @@ class PipelineConfig:
"""
return text_inputs.get("attention_mask")
def build_text_conditioning_mask(
self,
text_inputs: dict,
text_encoder_attention_mask: "torch.Tensor | None",
prompt_embeds: "torch.Tensor",
encoder_index: int,
) -> "torch.Tensor":
"""Return a mask aligned with post-processed prompt embeddings.
True values mark valid text tokens. Dynamic batching must carry
post-processed semantic text lengths explicitly; if a model-specific
postprocessor changes the sequence length, it must return
TextConditioningOutput with an embedding-aligned mask.
"""
if prompt_embeds.ndim < 2:
raise ValueError(
"prompt_embeds must have shape [batch, seq, ...] to build text conditioning mask"
)
if prompt_embeds.ndim == 2:
batch_size, embed_seq_len = 1, prompt_embeds.shape[0]
else:
batch_size, embed_seq_len = prompt_embeds.shape[:2]
device = prompt_embeds.device
if text_encoder_attention_mask is None:
return torch.ones(
(batch_size, embed_seq_len), dtype=torch.bool, device=device
)
raw_mask = text_encoder_attention_mask.to(device=device).bool()
if raw_mask.ndim != 2 or raw_mask.shape[0] != batch_size:
raise ValueError(
"text attention mask must have shape [batch, seq] matching prompt_embeds batch"
)
if raw_mask.shape[1] == embed_seq_len:
return raw_mask
if prompt_embeds.ndim == 2 and raw_mask.shape[0] == 1:
return torch.ones((1, embed_seq_len), dtype=torch.bool, device=device)
raise ValueError(
"text attention mask length does not match postprocessed prompt embeddings. "
"Postprocess functions that trim, pack, or otherwise change text sequence "
"length must return TextConditioningOutput with an embedding-aligned mask."
)
@staticmethod
def seq_lens_from_text_conditioning_mask(mask: "torch.Tensor") -> list[int]:
if mask.ndim != 2:
raise ValueError("text conditioning mask must have shape [batch, seq]")
return torch.count_nonzero(mask, dim=1).tolist()
def require_text_seq_lens(
self,
batch,
encoder_index: int,
*,
negative: bool = False,
expected_batch_size: int | None = None,
) -> list[int]:
"""Return postprocessed text lengths captured during text encoding.
Dynamic batches use these lengths for model masks, RoPE, and cache
sizing after text embeddings have been padded.
"""
seq_lens_by_encoder = (
batch.negative_prompt_seq_lens if negative else batch.prompt_seq_lens
)
kind = "negative" if negative else "positive"
if seq_lens_by_encoder is None or encoder_index >= len(seq_lens_by_encoder):
raise ValueError(
f"Missing {kind} prompt_seq_lens for text encoder {encoder_index}; "
"dynamic text conditioning requires explicit sequence lengths."
)
seq_lens = [int(x) for x in seq_lens_by_encoder[encoder_index]]
if expected_batch_size is not None and len(seq_lens) != int(
expected_batch_size
):
raise ValueError(
f"{kind} prompt_seq_lens for text encoder {encoder_index} has "
f"{len(seq_lens)} entries, expected {expected_batch_size}."
)
return seq_lens
def get_text_encoder_pooler_output(
self, outputs: "BaseEncoderOutput", encoder_index: int
) -> "torch.Tensor | None":
python/sglang/multimodal_gen/configs/pipeline_configs/flux.py
+95 -2
View File
@@ -91,6 +91,30 @@ class FluxPipelineConfig(ImagePipelineConfig):
# Flux v1 does not use attention masks for text encoders.
return None
def build_text_conditioning_mask(
self,
text_inputs: dict,
text_encoder_attention_mask: "torch.Tensor | None",
prompt_embeds: "torch.Tensor",
encoder_index: int,
) -> "torch.Tensor":
"""Use all-valid fixed-length masks for Flux v1 text embeddings."""
if prompt_embeds.ndim < 2:
raise ValueError(
"prompt_embeds must have shape [batch, seq, ...] or [seq, ...]"
)
if prompt_embeds.ndim == 2:
shape = (1, prompt_embeds.shape[0])
else:
shape = prompt_embeds.shape[:2]
return torch.ones(shape, dtype=torch.bool)
@staticmethod
def seq_lens_from_text_conditioning_mask(mask: "torch.Tensor") -> list[int]:
if mask.ndim != 2:
raise ValueError("text conditioning mask must have shape [batch, seq]")
return [int(mask.shape[1])] * int(mask.shape[0])
def get_text_encoder_pooler_output(self, outputs, encoder_index):
return outputs.pooler_output
@@ -143,7 +167,27 @@ class FluxPipelineConfig(ImagePipelineConfig):
return latent_image_ids
def get_freqs_cis(self, prompt_embeds, width, height, device, rotary_emb, batch):
@staticmethod
def _validate_fixed_text_seq_lens(prompt_embeds, txt_seq_lens):
if prompt_embeds.ndim < 3:
raise ValueError(
"Flux text conditioning expects prompt_embeds with shape [batch, seq, dim]"
)
batch_size, seq_len = prompt_embeds.shape[:2]
if len(txt_seq_lens) != batch_size:
raise ValueError(
f"Flux text sequence lengths have {len(txt_seq_lens)} entries, expected {batch_size}."
)
if any(int(seq_len_i) != seq_len for seq_len_i in txt_seq_lens):
raise ValueError(
"Flux currently requires fixed-length text conditioning; "
f"got seq_lens={txt_seq_lens}, expected all {seq_len}."
)
def get_freqs_cis(
self, prompt_embeds, width, height, device, rotary_emb, batch, txt_seq_lens
):
self._validate_fixed_text_seq_lens(prompt_embeds, txt_seq_lens)
txt_ids = torch.zeros(prompt_embeds.shape[1], 3, device=device)
img_ids = self._prepare_latent_image_ids(
original_height=height,
@@ -175,6 +219,21 @@ class FluxPipelineConfig(ImagePipelineConfig):
return latents
def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype):
"""Build Flux positive-conditioning kwargs from encoded text state.
Flux v1 uses encoder index 1 (the T5 encoder) as the token stream that
is concatenated with image tokens for rotary position embeddings. The
text encoding stage stores per-request sequence lengths in
batch.prompt_seq_lens; read them here instead of inferring from padded
embeddings so grouped multi-output requests preserve their explicit
text-conditioning contract.
"""
txt_seq_lens = self.require_text_seq_lens(
batch,
1,
negative=False,
expected_batch_size=batch.prompt_embeds[1].shape[0],
)
return {
"freqs_cis": self.get_freqs_cis(
batch.prompt_embeds[1],
@@ -183,6 +242,7 @@ class FluxPipelineConfig(ImagePipelineConfig):
device,
rotary_emb,
batch,
txt_seq_lens,
),
"pooled_projections": (
batch.pooled_embeds[0] if batch.pooled_embeds else None
@@ -190,6 +250,13 @@ class FluxPipelineConfig(ImagePipelineConfig):
}
def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype):
"""Build Flux negative-conditioning kwargs using T5 sequence lengths."""
txt_seq_lens = self.require_text_seq_lens(
batch,
1,
negative=True,
expected_batch_size=batch.negative_prompt_embeds[1].shape[0],
)
return {
"freqs_cis": self.get_freqs_cis(
batch.negative_prompt_embeds[1],
@@ -198,6 +265,7 @@ class FluxPipelineConfig(ImagePipelineConfig):
device,
rotary_emb,
batch,
txt_seq_lens,
),
"pooled_projections": (
batch.neg_pooled_embeds[0] if batch.neg_pooled_embeds else None
@@ -412,6 +480,14 @@ class Flux2PipelineConfig(FluxPipelineConfig):
# Flux2 does not use pooler output.
return None
def supports_dynamic_batching(self):
"""Allow batching for Flux2 text-only requests.
Flux2 is a TI2I pipeline, so image-input requests are rejected by the
scheduler's request-level batching checks.
"""
return True
def tokenize_prompt(self, prompts: list[str], tokenizer, tok_kwargs) -> dict:
messages = build_flux2_text_messages(prompts)
inputs = tokenizer.apply_chat_template(
@@ -513,7 +589,10 @@ class Flux2PipelineConfig(FluxPipelineConfig):
image_latent_ids = image_latent_ids.repeat(batch.batch_size, 1, 1)
batch.condition_image_latent_ids = image_latent_ids.to(get_local_torch_device())
def get_freqs_cis(self, prompt_embeds, width, height, device, rotary_emb, batch):
def get_freqs_cis(
self, prompt_embeds, width, height, device, rotary_emb, batch, txt_seq_lens
):
self._validate_fixed_text_seq_lens(prompt_embeds, txt_seq_lens)
txt_ids = _prepare_text_ids(prompt_embeds).to(device=device)
img_ids = batch.latent_ids
@@ -544,6 +623,19 @@ class Flux2PipelineConfig(FluxPipelineConfig):
return cos, sin
def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype):
"""Build Flux2 positive-conditioning kwargs from encoded text state.
Flux2 uses encoder index 0 for the Mistral text stream. The stored
sequence lengths are passed through to rotary-position preparation so
grouped requests use the same text-length metadata that was produced
during text encoding.
"""
txt_seq_lens = self.require_text_seq_lens(
batch,
0,
negative=False,
expected_batch_size=batch.prompt_embeds[0].shape[0],
)
return {
"freqs_cis": self.get_freqs_cis(
batch.prompt_embeds[0],
@@ -552,6 +644,7 @@ class Flux2PipelineConfig(FluxPipelineConfig):
device,
rotary_emb,
batch,
txt_seq_lens,
)
}
python/sglang/multimodal_gen/configs/pipeline_configs/hunyuan.py
+25 -6
View File
@@ -18,6 +18,7 @@ from sglang.multimodal_gen.configs.models.vaes import HunyuanVAEConfig
from sglang.multimodal_gen.configs.pipeline_configs.base import (
ModelTaskType,
PipelineConfig,
TextConditioningOutput,
)
PROMPT_TEMPLATE_ENCODE_VIDEO = (
@@ -46,19 +47,37 @@ def llama_preprocess_text(prompt: str) -> str:
return prompt_template_video["template"].format(prompt)
def llama_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.tensor:
def llama_postprocess_text(
outputs: BaseEncoderOutput, _text_inputs
) -> TextConditioningOutput:
hidden_state_skip_layer = 2
assert outputs.hidden_states is not None
hidden_states: tuple[torch.Tensor, ...] = outputs.hidden_states
last_hidden_state: torch.tensor = hidden_states[-(hidden_state_skip_layer + 1)]
last_hidden_state: torch.Tensor = hidden_states[-(hidden_state_skip_layer + 1)]
crop_start = prompt_template_video.get("crop_start", -1)
last_hidden_state = last_hidden_state[:, crop_start:]
return last_hidden_state
attention_mask = _text_inputs.attention_mask.to(
device=last_hidden_state.device, dtype=torch.bool
)
if crop_start < 0:
attention_mask = attention_mask[:, crop_start:]
else:
attention_mask = attention_mask[
:, crop_start : crop_start + last_hidden_state.shape[1]
]
seq_lens = [int(x) for x in attention_mask.to(torch.int64).sum(dim=1).tolist()]
return TextConditioningOutput(last_hidden_state, attention_mask, seq_lens)
def clip_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.tensor:
pooler_output: torch.tensor = outputs.pooler_output
return pooler_output
def clip_postprocess_text(
outputs: BaseEncoderOutput, _text_inputs
) -> TextConditioningOutput:
pooler_output: torch.Tensor = outputs.pooler_output
batch_size = int(pooler_output.shape[0])
prompt_embeds_mask = torch.ones(
(batch_size, 1), dtype=torch.bool, device=pooler_output.device
)
return TextConditioningOutput(pooler_output, prompt_embeds_mask, [1] * batch_size)
@dataclass
python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py
+157 -47
View File
@@ -16,6 +16,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import (
ImagePipelineConfig,
ModelTaskType,
maybe_unpad_latents,
pad_text_embeddings_with_mask,
shard_rotary_emb_for_sp,
)
from sglang.multimodal_gen.configs.post_training.pipeline_configs import (
@@ -45,34 +46,21 @@ def qwen_image_preprocess_text(prompt):
def qwen_image_postprocess_text(
outputs, _text_inputs, drop_idx=34, return_attention_mask=False
):
"""Postprocess Qwen text embeddings.
Returns padded embeddings by default, or TextConditioningOutput when
embedding-aligned masks are requested.
"""
# squeeze the batch dim
hidden_states = outputs.hidden_states[-1]
split_hidden_states = _extract_masked_hidden(
hidden_states, _text_inputs.attention_mask
)
split_hidden_states = [e[drop_idx:] for e in split_hidden_states]
attn_mask_list = [
torch.ones(e.size(0), dtype=torch.long, device=e.device)
for e in split_hidden_states
]
max_seq_len = max([e.size(0) for e in split_hidden_states])
prompt_embeds = torch.stack(
[
torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))])
for u in split_hidden_states
]
)
conditioning = pad_text_embeddings_with_mask(split_hidden_states)
if return_attention_mask:
encoder_attention_mask = torch.stack(
[
torch.cat([u, u.new_zeros(max_seq_len - u.size(0))])
for u in attn_mask_list
]
)
if encoder_attention_mask.all():
return prompt_embeds, None
return prompt_embeds, encoder_attention_mask
return prompt_embeds
return conditioning
return conditioning.prompt_embeds
def qwen_image_edit_postprocess_text(outputs, _text_inputs):
@@ -236,10 +224,14 @@ class QwenImagePipelineConfig(QwenImageRolloutPipelineMixin, ImagePipelineConfig
return {}
def get_vae_scale_factor(self):
return self.vae_config.arch_config.vae_scale_factor
return getattr(
self.vae_config.arch_config,
"vae_scale_factor",
self.vae_config.get_vae_scale_factor(),
)
def prepare_latent_shape(self, batch, batch_size, num_frames):
vae_scale_factor = self.vae_config.arch_config.vae_scale_factor
vae_scale_factor = self.get_vae_scale_factor()
height = 2 * (batch.height // (vae_scale_factor * 2))
width = 2 * (batch.width // (vae_scale_factor * 2))
num_channels_latents = self.dit_config.arch_config.in_channels // 4
@@ -247,10 +239,9 @@ class QwenImagePipelineConfig(QwenImageRolloutPipelineMixin, ImagePipelineConfig
return shape
def maybe_pack_latents(self, latents, batch_size, batch):
height = 2 * (
batch.height // (self.vae_config.arch_config.vae_scale_factor * 2)
)
width = 2 * (batch.width // (self.vae_config.arch_config.vae_scale_factor * 2))
vae_scale_factor = self.get_vae_scale_factor()
height = 2 * (batch.height // (vae_scale_factor * 2))
width = 2 * (batch.width // (vae_scale_factor * 2))
num_channels_latents = self.dit_config.arch_config.in_channels // 4
# pack latents
return _pack_latents(latents, batch_size, num_channels_latents, height, width)
@@ -294,11 +285,19 @@ class QwenImagePipelineConfig(QwenImageRolloutPipelineMixin, ImagePipelineConfig
txt_cos_sin_cache = torch.cat([txt_cos_half, txt_sin_half], dim=-1)
return img_cos_sin_cache, txt_cos_sin_cache
def _prepare_cond_kwargs(self, batch, prompt_embeds, rotary_emb, device, dtype):
def _prepare_cond_kwargs(
self, batch, prompt_embeds, rotary_emb, device, dtype, *, negative=False
):
"""Build Qwen DiT conditioning kwargs for positive or negative prompts.
The kwargs include text lengths for RoPE construction and optional
encoder masks for cross-attention.
"""
batch_size = prompt_embeds[0].shape[0]
text_seq_len = prompt_embeds[0].shape[1]
height = batch.height
width = batch.width
vae_scale_factor = self.vae_config.arch_config.vae_scale_factor
vae_scale_factor = self.get_vae_scale_factor()
img_shapes = [
[
@@ -309,14 +308,18 @@ class QwenImagePipelineConfig(QwenImageRolloutPipelineMixin, ImagePipelineConfig
)
]
] * batch_size
txt_seq_lens = [prompt_embeds[0].shape[1]]
txt_seq_lens, encoder_hidden_states_mask = self._prepare_text_conditioning(
batch, 0, text_seq_len, batch_size, negative=negative
)
if rotary_emb is None:
return {
cond_kwargs = {
"img_shapes": img_shapes,
"txt_seq_lens": txt_seq_lens,
"freqs_cis": None,
"encoder_hidden_states_mask": encoder_hidden_states_mask,
}
return cond_kwargs
freqs_cis = self.get_freqs_cis(
img_shapes, txt_seq_lens, rotary_emb, device, dtype
@@ -324,20 +327,105 @@ class QwenImagePipelineConfig(QwenImageRolloutPipelineMixin, ImagePipelineConfig
img_cache, txt_cache = freqs_cis
img_cache = shard_rotary_emb_for_sp(img_cache)
return {
cond_kwargs = {
"txt_seq_lens": txt_seq_lens,
"freqs_cis": (img_cache, txt_cache),
"img_shapes": img_shapes,
"encoder_hidden_states_mask": encoder_hidden_states_mask,
}
return cond_kwargs
def _prepare_text_conditioning(
self,
batch,
encoder_index: int,
text_seq_len: int,
batch_size: int,
*,
negative: bool = False,
):
"""Return Qwen text lengths and an optional DiT attention mask.
Single-request execution uses the full padded length. Batched execution
uses stored per-request lengths and masks from text encoding.
"""
if batch_size == 1:
return [text_seq_len], None
txt_seq_lens = self.require_text_seq_lens(
batch, encoder_index, negative=negative, expected_batch_size=batch_size
)
encoder_hidden_states_mask = self._prepare_encoder_hidden_states_mask(
batch,
encoder_index,
txt_seq_lens,
text_seq_len,
batch_size,
negative=negative,
)
return txt_seq_lens, encoder_hidden_states_mask
def _prepare_encoder_hidden_states_mask(
self,
batch,
encoder_index: int,
txt_seq_lens: list[int],
text_seq_len: int,
batch_size: int,
*,
negative: bool = False,
):
"""Return the text attention mask passed to the Qwen image DiT.
Qwen image batches can contain prompts with different semantic text
lengths after tokenization/postprocessing. The transformer still sees a
padded `encoder_hidden_states` tensor with shape [batch, text_seq_len,
dim], so we pass a [batch, text_seq_len] boolean mask to keep attention
on real text tokens and ignore padding.
If every request uses the full padded length, no mask is needed and this
returns None. Otherwise, prefer the embedding-aligned mask stored by the
text encoding stage. If that is unavailable, rebuild the same mask from
`txt_seq_lens`: position j is valid for row i when
`j < txt_seq_lens[i]`.
"""
if all(seq_len == text_seq_len for seq_len in txt_seq_lens):
return None
masks_by_encoder = (
batch.negative_prompt_embeds_mask if negative else batch.prompt_embeds_mask
)
if masks_by_encoder is not None and encoder_index < len(masks_by_encoder):
mask = masks_by_encoder[encoder_index]
if mask.shape != (batch_size, text_seq_len):
raise ValueError(
"QwenImage text conditioning mask has shape "
f"{tuple(mask.shape)}, expected {(batch_size, text_seq_len)}."
)
return mask
# TODO: cache positions by (device, text_seq_len) if this allocation shows up hot.
positions = torch.arange(text_seq_len, device=batch.prompt_embeds[0].device)
seq_lens = torch.tensor(
txt_seq_lens,
device=batch.prompt_embeds[0].device,
dtype=torch.long,
)
return positions.unsqueeze(0) < seq_lens.unsqueeze(1)
def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype):
return self._prepare_cond_kwargs(
batch, batch.prompt_embeds, rotary_emb, device, dtype
batch, batch.prompt_embeds, rotary_emb, device, dtype, negative=False
)
def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype):
return self._prepare_cond_kwargs(
batch, batch.negative_prompt_embeds, rotary_emb, device, dtype
batch,
batch.negative_prompt_embeds,
rotary_emb,
device,
dtype,
negative=True,
)
def post_denoising_loop(self, latents, batch):
@@ -363,10 +451,11 @@ class QwenImageEditPipelineConfig(QwenImagePipelineConfig):
)
def _prepare_edit_cond_kwargs(
self, batch, prompt_embeds, rotary_emb, device, dtype
self, batch, prompt_embeds, rotary_emb, device, dtype, *, negative=False
):
batch_size = batch.latents.shape[0]
assert batch_size == 1
text_seq_len = prompt_embeds[0].shape[1]
height = batch.height
width = batch.width
image_size = batch.original_condition_image_size
@@ -389,14 +478,18 @@ class QwenImageEditPipelineConfig(QwenImagePipelineConfig):
),
],
] * batch_size
txt_seq_lens = [prompt_embeds[0].shape[1]]
txt_seq_lens, encoder_hidden_states_mask = self._prepare_text_conditioning(
batch, 0, text_seq_len, batch_size, negative=negative
)
if rotary_emb is None:
return {
cond_kwargs = {
"img_shapes": img_shapes,
"txt_seq_lens": txt_seq_lens,
"freqs_cis": None,
"encoder_hidden_states_mask": encoder_hidden_states_mask,
}
return cond_kwargs
freqs_cis = QwenImagePipelineConfig.get_freqs_cis(
img_shapes, txt_seq_lens, rotary_emb, device, dtype
@@ -410,11 +503,13 @@ class QwenImageEditPipelineConfig(QwenImagePipelineConfig):
img_cache, txt_cache = _shard_qwen_edit_freqs_cis_for_sp(
freqs_cis, noisy_img_seq_len, device
)
return {
cond_kwargs = {
"txt_seq_lens": txt_seq_lens,
"freqs_cis": (img_cache, txt_cache),
"img_shapes": img_shapes,
"encoder_hidden_states_mask": encoder_hidden_states_mask,
}
return cond_kwargs
def preprocess_condition_image(
self, image, target_width, target_height, _vae_image_processor
@@ -453,12 +548,17 @@ class QwenImageEditPipelineConfig(QwenImagePipelineConfig):
def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype):
return self._prepare_edit_cond_kwargs(
batch, batch.prompt_embeds, rotary_emb, device, dtype
batch, batch.prompt_embeds, rotary_emb, device, dtype, negative=False
)
def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype):
return self._prepare_edit_cond_kwargs(
batch, batch.negative_prompt_embeds, rotary_emb, device, dtype
batch,
batch.negative_prompt_embeds,
rotary_emb,
device,
dtype,
negative=True,
)
def calculate_condition_image_size(self, image, width, height) -> tuple[int, int]:
@@ -562,10 +662,11 @@ class QwenImageEditPlusPipelineConfig(QwenImageEditPipelineConfig):
return batch
def _prepare_edit_cond_kwargs(
self, batch, prompt_embeds, rotary_emb, device, dtype
self, batch, prompt_embeds, rotary_emb, device, dtype, *, negative=False
):
batch_size = batch.latents.shape[0]
assert batch_size == 1
text_seq_len = prompt_embeds[0].shape[1]
height = batch.height
width = batch.width
@@ -584,7 +685,9 @@ class QwenImageEditPlusPipelineConfig(QwenImageEditPipelineConfig):
],
],
] * batch_size
txt_seq_lens = [prompt_embeds[0].shape[1]]
txt_seq_lens, encoder_hidden_states_mask = self._prepare_text_conditioning(
batch, 0, text_seq_len, batch_size, negative=negative
)
freqs_cis = QwenImageEditPlusPipelineConfig.get_freqs_cis(
img_shapes, txt_seq_lens, rotary_emb, device, dtype
@@ -595,13 +698,15 @@ class QwenImageEditPlusPipelineConfig(QwenImageEditPipelineConfig):
1 * (height // vae_scale_factor // 2) * (width // vae_scale_factor // 2)
)
return {
cond_kwargs = {
"txt_seq_lens": txt_seq_lens,
"freqs_cis": _shard_qwen_edit_freqs_cis_for_sp(
freqs_cis, noisy_img_seq_len, device
),
"img_shapes": img_shapes,
"encoder_hidden_states_mask": encoder_hidden_states_mask,
}
return cond_kwargs
@dataclass
@@ -625,17 +730,20 @@ class QwenImageLayeredPipelineConfig(QwenImageEditPipelineConfig):
return super().postprocess_cfg_noise(batch, noise_pred, noise_pred_cond)
def _prepare_edit_cond_kwargs(
self, batch, prompt_embeds, rotary_emb, device, dtype
self, batch, prompt_embeds, rotary_emb, device, dtype, *, negative=False
):
batch_size = batch.latents.shape[0]
assert batch_size == 1
text_seq_len = prompt_embeds[0].shape[1]
height = batch.height
width = batch.width
vae_scale_factor = self.get_vae_scale_factor()
img_shapes = batch.img_shapes
txt_seq_lens = [prompt_embeds[0].shape[1]]
txt_seq_lens, encoder_hidden_states_mask = self._prepare_text_conditioning(
batch, 0, text_seq_len, batch_size, negative=negative
)
freqs_cis = QwenImageEditPlusPipelineConfig.get_freqs_cis(
img_shapes, txt_seq_lens, rotary_emb, device, dtype
@@ -652,15 +760,17 @@ class QwenImageLayeredPipelineConfig(QwenImageEditPipelineConfig):
[noisy_img_cache, img_cache[noisy_img_seq_len:, :]], dim=0
).to(device=device)
return {
cond_kwargs = {
"txt_seq_lens": txt_seq_lens,
"img_shapes": img_shapes,
"freqs_cis": (img_cache, txt_cache),
"additional_t_cond": torch.tensor([0], device=device, dtype=torch.long),
"encoder_hidden_states_mask": encoder_hidden_states_mask,
}
return cond_kwargs
def _unpad_and_unpack_latents(self, latents, batch):
vae_scale_factor = self.vae_config.arch_config.vae_scale_factor
vae_scale_factor = self.get_vae_scale_factor()
channels = self.dit_config.arch_config.in_channels
batch_size = latents.shape[0]
layers = batch.num_frames
python/sglang/multimodal_gen/configs/pipeline_configs/sana.py
+9
View File
@@ -64,6 +64,15 @@ class SanaPipelineConfig(SpatialImagePipelineConfig):
text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16",))
text_encoder_extra_args: list[dict] = field(
default_factory=lambda: [
{
"padding": True,
"return_attention_mask": True,
}
]
)
preprocess_text_funcs: tuple[Callable[[str], str] | None, ...] = field(
default_factory=lambda: (None,),
)
python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py
+107 -7
View File
@@ -14,6 +14,8 @@ from sglang.multimodal_gen.configs.models.vaes.flux import FluxVAEConfig
from sglang.multimodal_gen.configs.pipeline_configs.base import (
ImagePipelineConfig,
ModelTaskType,
TextConditioningOutput,
pad_text_embeddings_with_mask,
)
from sglang.multimodal_gen.configs.post_training.pipeline_configs import (
ZImageRolloutPipelineMixin,
@@ -32,10 +34,24 @@ def zimage_preprocess_text(prompt: str):
return messages
def zimage_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tensor:
def zimage_postprocess_text(
outputs: BaseEncoderOutput, _text_inputs
) -> torch.Tensor | TextConditioningOutput:
"""Return unpadded Z-Image text embeddings.
Batched outputs return TextConditioningOutput to preserve per-prompt text
lengths.
"""
device = outputs.hidden_states[-2].device
prompt_mask = _text_inputs.attention_mask.to(device).bool()
return outputs.hidden_states[-2][0][prompt_mask[0]]
hidden_states = outputs.hidden_states[-2]
if hidden_states.shape[0] == 1:
return hidden_states[0][prompt_mask[0]]
split_hidden_states = [
hidden_states[idx][prompt_mask[idx]] for idx in range(hidden_states.shape[0])
]
return pad_text_embeddings_with_mask(split_hidden_states)
class TransformersModelConfig(EncoderConfig):
@@ -177,8 +193,72 @@ class ZImagePipelineConfig(ZImageRolloutPipelineMixin, ImagePipelineConfig):
plan = self._build_zimage_sp_plan(batch)
return plan
def _split_text_embeds_for_dit(self, batch, *, negative: bool = False):
"""Return per-request text tensors, trimming padded batched embeddings."""
embeds = batch.negative_prompt_embeds if negative else batch.prompt_embeds
if embeds is None:
return None
if isinstance(embeds, (list, tuple)):
if not embeds:
return []
embeds = embeds[0]
if not torch.is_tensor(embeds):
return embeds
if embeds.ndim == 2:
return [embeds]
if embeds.ndim != 3:
raise ValueError(
"Z-Image text embeddings must have shape [seq, dim] or [batch, seq, dim]"
)
seq_lens = self.require_text_seq_lens(
batch,
0,
negative=negative,
expected_batch_size=int(embeds.shape[0]),
)
return [
embeds[idx, :seq_len].contiguous() for idx, seq_len in enumerate(seq_lens)
]
def _caption_rope_length(self, prompt_embeds, batch, *, negative: bool = False):
"""Return the shared caption RoPE length for current text embeddings."""
if torch.is_tensor(prompt_embeds):
if prompt_embeds.ndim == 2:
return int(prompt_embeds.shape[0])
if prompt_embeds.ndim == 3:
seq_lens = self.require_text_seq_lens(
batch,
0,
negative=negative,
expected_batch_size=int(prompt_embeds.shape[0]),
)
return max(seq_lens) if seq_lens else int(prompt_embeds.shape[1])
if isinstance(prompt_embeds, (list, tuple)) and prompt_embeds:
first = prompt_embeds[0]
if torch.is_tensor(first):
if first.ndim == 3:
seq_lens = self.require_text_seq_lens(
batch,
0,
negative=negative,
expected_batch_size=int(first.shape[0]),
)
return max(seq_lens) if seq_lens else int(first.shape[1])
return max(int(item.shape[0]) for item in prompt_embeds)
raise ValueError("Unable to infer Z-Image caption length for rotary embeddings")
def get_pos_prompt_embeds(self, batch):
return batch.prompt_embeds
return self._split_text_embeds_for_dit(batch, negative=False)
def get_neg_prompt_embeds(self, batch):
return self._split_text_embeds_for_dit(batch, negative=True)
def get_latent_dtype(self, prompt_dtype: torch.dtype) -> torch.dtype:
# Match the official diffusers Z-Image pipeline, which samples latents in fp32
@@ -252,7 +332,23 @@ class ZImagePipelineConfig(ZImageRolloutPipelineMixin, ImagePipelineConfig):
return latents[:, :, 0, :, :]
return latents.view(bs, channels, height, width)
def get_freqs_cis(self, prompt_embeds, width, height, device, rotary_emb, batch):
def get_freqs_cis(
self,
prompt_embeds,
width,
height,
device,
rotary_emb,
batch,
*,
negative: bool = False,
):
"""Build caption and image RoPE caches for Z-Image conditioning.
Batched prompts use stored text lengths. SP mode builds image caches for
the local spatial shard.
"""
def create_coordinate_grid(size, start=None, device=None):
if start is None:
start = (0 for _ in size)
@@ -269,7 +365,9 @@ class ZImagePipelineConfig(ZImageRolloutPipelineMixin, ImagePipelineConfig):
# SP path: keep caption replicated on every rank and build local-only
# image freqs_cis matching the spatial shard.
plan = self._get_zimage_sp_plan(batch)
cap_ori_len = prompt_embeds.size(0)
cap_ori_len = self._caption_rope_length(
prompt_embeds, batch, negative=negative
)
cap_padding_len = (-cap_ori_len) % self.SEQ_LEN_MULTIPLE
# caption (replicated prefix)
@@ -304,7 +402,7 @@ class ZImagePipelineConfig(ZImageRolloutPipelineMixin, ImagePipelineConfig):
x_freqs_cis = rotary_emb(img_pos_ids)
return (cap_freqs_cis, x_freqs_cis)
cap_ori_len = prompt_embeds.size(0)
cap_ori_len = self._caption_rope_length(prompt_embeds, batch, negative=negative)
cap_padding_len = (-cap_ori_len) % self.SEQ_LEN_MULTIPLE
cap_padded_pos_ids = create_coordinate_grid(
size=(cap_ori_len + cap_padding_len, 1, 1),
@@ -361,9 +459,10 @@ class ZImagePipelineConfig(ZImageRolloutPipelineMixin, ImagePipelineConfig):
}
def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype):
use_negative_embeds = batch.negative_prompt_embeds is not None
prompt_embeds = (
batch.negative_prompt_embeds[0]
if batch.negative_prompt_embeds is not None
if use_negative_embeds
else batch.prompt_embeds[0]
)
return {
@@ -374,6 +473,7 @@ class ZImagePipelineConfig(ZImageRolloutPipelineMixin, ImagePipelineConfig):
device,
rotary_emb,
batch,
negative=use_negative_embeds,
),
"image_seq_len_target": (
self._get_zimage_sp_plan(batch)["img_seq_target"]
python/sglang/multimodal_gen/configs/sample/sampling_params.py
+31 -18
View File
@@ -12,7 +12,7 @@ import re
import time
import unicodedata
import uuid
from dataclasses import dataclass
from dataclasses import dataclass, field
from enum import Enum, auto
from typing import TYPE_CHECKING, Any, ClassVar
@@ -89,11 +89,14 @@ class DataType(Enum):
class SamplingParams:
"""
Sampling parameters for generation.
Dynamic batching compares these fields for compatibility, except fields
marked with `batch_sig_exclude`.
"""
data_type: DataType = DataType.VIDEO
request_id: str | None = None
request_id: str | None = field(default=None, metadata={"batch_sig_exclude": True})
# All fields below are copied from ForwardBatch
@@ -101,13 +104,17 @@ class SamplingParams:
image_path: str | list[str] | None = None
# Text inputs
prompt: str | list[str] | None = None
prompt: str | list[str] | None = field(
default=None, metadata={"batch_sig_exclude": True}
)
negative_prompt: str = (
"Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
)
prompt_path: str | None = None
output_path: str | None = None
output_file_name: str | None = None
prompt_path: str | None = field(default=None, metadata={"batch_sig_exclude": True})
output_path: str | None = field(default=None, metadata={"batch_sig_exclude": True})
output_file_name: str | None = field(
default=None, metadata={"batch_sig_exclude": True}
)
output_quality: str | None = "default"
output_compression: int | None = None
@@ -128,7 +135,7 @@ class SamplingParams:
# Batch info
num_outputs_per_prompt: int = 1
seed: int | list[int] = 42
seed: int | list[int] = field(default=42, metadata={"batch_sig_exclude": True})
generator_device: str | None = None # None means use the pipeline/model default
# Original dimensions (before VAE scaling)
@@ -147,9 +154,9 @@ class SamplingParams:
fps: int = 24
# Resolution validation
supported_resolutions: list[tuple[int, int]] | None = (
None # None means all resolutions allowed
)
supported_resolutions: list[tuple[int, int]] | None = field(
default=None, metadata={"batch_sig_exclude": True}
) # None means all resolutions allowed
# Denoising parameters
num_inference_steps: int = None
@@ -167,17 +174,21 @@ class SamplingParams:
)
# Profiling
profile: bool = False
num_profiled_timesteps: int = 5
profile_all_stages: bool = False
profile: bool = field(default=False, metadata={"batch_sig_exclude": True})
num_profiled_timesteps: int = field(default=5, metadata={"batch_sig_exclude": True})
profile_all_stages: bool = field(
default=False, metadata={"batch_sig_exclude": True}
)
# Debugging
debug: bool = False
perf_dump_path: str | None = None
debug: bool = field(default=False, metadata={"batch_sig_exclude": True})
perf_dump_path: str | None = field(
default=None, metadata={"batch_sig_exclude": True}
)
# Misc
save_output: bool = True
return_frames: bool = False
return_frames: bool = field(default=False, metadata={"batch_sig_exclude": True})
rollout: bool = False
rollout_sde_type: str = "sde"
rollout_noise_level: float = 0.7
@@ -197,11 +208,13 @@ class SamplingParams:
rollout_sde_step_indices: list[int] | None = None
rollout_return_step_indices: list[int] | None = None
# if True, disallow user params to override subclass-defined protected fields
no_override_protected_fields: bool = False
no_override_protected_fields: bool = field(
default=False, metadata={"batch_sig_exclude": True}
)
# whether to adjust num_frames for multi-GPU friendly splitting (default: True)
adjust_frames: bool = True
# if True, suppress verbose logging for this request
suppress_logs: bool = False
suppress_logs: bool = field(default=False, metadata={"batch_sig_exclude": True})
return_file_paths_only: bool = True
enable_sequence_shard: bool | None = None
python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py
+6 -1
View File
@@ -349,7 +349,12 @@ class DiffGenerator:
global_output_index += len(requests)
total_gen_time = time.perf_counter() - total_start_time
log_batch_completion(logger, len(results), total_gen_time)
if self.server_args.batching_max_size > 1:
log_batch_completion(
logger,
len(results),
total_gen_time,
)
self._log_summary(results)
if not results:
python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py
+9 -14
View File
@@ -22,7 +22,6 @@ from sglang.multimodal_gen.runtime.entrypoints.utils import (
SetLoraReq,
ShutdownReq,
UnmergeLoraWeightsReq,
expand_request_outputs,
format_lora_message,
save_outputs,
)
@@ -326,9 +325,8 @@ async def process_generation_batch(
batch,
) -> tuple[list[str], OutputBatch]:
total_start_time = time.perf_counter()
requests = expand_request_outputs(batch)
with trace_req(batch.trace_ctx), log_generation_timer(logger, batch.prompt):
result = await scheduler_client.forward(requests)
result = await scheduler_client.forward([batch])
if result.output is None and result.output_file_paths is None:
error_msg = result.error or "Unknown error"
@@ -338,22 +336,14 @@ async def process_generation_batch(
if result.output_file_paths:
save_file_path_list = result.output_file_paths
if len(save_file_path_list) < len(requests):
raise RuntimeError(
f"Expected at least {len(requests)} output paths, "
f"got {len(save_file_path_list)}"
)
else:
if len(result.output) != len(requests):
raise RuntimeError(
f"Expected {len(requests)} outputs, got {len(result.output)}"
)
num_outputs = len(result.output)
save_file_path_list = save_outputs(
result.output,
batch.data_type,
batch.fps,
batch.save_output,
lambda idx: str(requests[idx].output_file_path(1, 0)),
lambda idx: str(batch.output_file_path(num_outputs, idx)),
audio=result.audio,
audio_sample_rate=result.audio_sample_rate,
output_compression=batch.output_compression,
@@ -367,7 +357,12 @@ async def process_generation_batch(
)
total_time = time.perf_counter() - total_start_time
log_batch_completion(logger, len(save_file_path_list), total_time)
if get_global_server_args().batching_max_size > 1:
log_batch_completion(
logger,
len(save_file_path_list),
total_time,
)
if result.peak_memory_mb and result.peak_memory_mb > 0:
logger.info(f"Peak memory usage: {result.peak_memory_mb:.2f} MB")
python/sglang/multimodal_gen/runtime/layers/rotary_embedding/base.py
+2 -2
View File
@@ -77,14 +77,14 @@ class RotaryEmbedding(CustomOp):
cos, sin = cos_sin.chunk(2, dim=-1)
query_shape = query.shape
query = query.view(num_tokens, -1, self.head_size)
query = query.reshape(num_tokens, -1, self.head_size)
query_rot = query[..., : self.rotary_dim]
query_pass = query[..., self.rotary_dim :]
query_rot = _apply_rotary_emb(query_rot, cos, sin, self.is_neox_style)
query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
key_shape = key.shape
key = key.view(num_tokens, -1, self.head_size)
key = key.reshape(num_tokens, -1, self.head_size)
key_rot = key[..., : self.rotary_dim]
key_pass = key[..., self.rotary_dim :]
key_rot = _apply_rotary_emb(key_rot, cos, sin, self.is_neox_style)
python/sglang/multimodal_gen/runtime/managers/dynamic_batch_admission.py
+345
View File
@@ -0,0 +1,345 @@
# SPDX-License-Identifier: Apache-2.0
"""Admission control for native diffusion batching.
Native diffusion batching is model, resolution, device, and implementation
dependent. The scheduler treats `--batching-max-size` as the public ceiling;
`--batching-config` can apply stricter caps for specific model and shape
combinations.
"""
from __future__ import annotations
import json
import os
from dataclasses import dataclass
from difflib import get_close_matches
from typing import TYPE_CHECKING, Any
from sglang.multimodal_gen.runtime.loader.utils import BYTES_PER_GB
from sglang.multimodal_gen.runtime.pipelines_core import Req
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
if TYPE_CHECKING:
from sglang.multimodal_gen.runtime.server_args import ServerArgs
logger = init_logger(__name__)
_BATCHING_RULE_KEYS = frozenset(
{
"model",
"model_contains",
"resolution",
"device_memory_gb_min",
"device_memory_gb_max",
"offload",
"max_batch_size",
"max_cost",
# Free-form provenance/benchmark metadata. It is intentionally ignored
# by admission, but accepted so production configs can explain caps.
"calibration",
}
)
@dataclass(frozen=True)
class AdmissionLimit:
"""Effective batch size and cost caps after matching batching rules."""
max_batch_size: int
max_cost: float | None = None
cap_reason: str | None = None
def reject_reason(self, *, batch_size: int, batch_cost: float) -> str | None:
if batch_size > self.max_batch_size:
return self.cap_reason or f"config_cap:{self.max_batch_size}"
if self.max_cost is not None and batch_cost > self.max_cost:
return f"cost_budget:{batch_cost:.0f}>{self.max_cost:.0f}"
return None
def stop_reason_for_next_cost(self, next_batch_cost: float) -> str | None:
if self.max_cost is not None and next_batch_cost > self.max_cost:
return f"cost_budget_next:{next_batch_cost:.0f}>{self.max_cost:.0f}"
return None
@dataclass(frozen=True)
class BatchingRule:
"""One user-provided batching admission rule loaded from batching config."""
model: str | None = None
model_contains: str | None = None
resolution: str | None = None
device_memory_gb_min: float | None = None
device_memory_gb_max: float | None = None
offload: bool | None = None
max_batch_size: int = 1
max_cost: float | None = None
source: str = "user"
@classmethod
def from_dict(cls, data: dict[str, Any], *, source: str) -> "BatchingRule":
if not isinstance(data, dict):
raise ValueError(
f"batching config rule from {source} must be an object, "
f"got {type(data).__name__}"
)
_validate_rule_keys(data, source=source)
if "max_batch_size" not in data:
raise ValueError("batching config rule requires max_batch_size")
rule = cls(
model=_optional_str(data.get("model")),
model_contains=_optional_str(data.get("model_contains")),
resolution=_optional_str(data.get("resolution")),
device_memory_gb_min=_optional_float(data.get("device_memory_gb_min")),
device_memory_gb_max=_optional_float(data.get("device_memory_gb_max")),
offload=_optional_bool(data.get("offload")),
max_batch_size=int(data["max_batch_size"]),
max_cost=_optional_float(data.get("max_cost")),
source=source,
)
rule.validate()
return rule
def validate(self) -> None:
if self.model is not None and self.model_contains is not None:
raise ValueError(
"batching config rule cannot set both model and model_contains"
)
if self.model is None and self.model_contains is None:
raise ValueError("batching config rule requires model or model_contains")
if self.max_batch_size < 1:
raise ValueError("batching config rule max_batch_size must be >= 1")
if self.max_cost is not None and self.max_cost <= 0.0:
raise ValueError("batching config rule max_cost must be > 0")
if (
self.device_memory_gb_min is not None
and self.device_memory_gb_max is not None
and self.device_memory_gb_min > self.device_memory_gb_max
):
raise ValueError(
"batching config rule device_memory_gb_min must be <= device_memory_gb_max"
)
def matches(
self,
*,
model_path: str,
resolution: str | None,
device_memory_gb: float | None,
offload: bool,
) -> bool:
if self.model is not None and self.model != model_path:
return False
if self.model_contains is not None and self.model_contains not in model_path:
return False
if self.resolution not in (None, "*") and self.resolution != resolution:
return False
if self.offload is not None and self.offload != offload:
return False
if device_memory_gb is None:
return True
if (
self.device_memory_gb_min is not None
and device_memory_gb < self.device_memory_gb_min
):
return False
if (
self.device_memory_gb_max is not None
and device_memory_gb > self.device_memory_gb_max
):
return False
return True
class BatchAdmissionController:
"""Applies configured caps before adding requests to a batch."""
def __init__(self, server_args: "ServerArgs", gpu_id: int):
self._mode = getattr(server_args, "batching_mode", "dynamic")
self._user_max_batch_size = max(1, int(server_args.batching_max_size))
self._model_path = server_args.model_path
self._offload = bool(server_args.dit_layerwise_offload)
self._device_memory_gb = self._get_device_memory_gb(gpu_id)
self._rules = load_batching_config(server_args.batching_config)
self._pipeline_config = server_args.pipeline_config
if self.enabled:
logger.info(
"Batch admission enabled: user_max=%d, device_memory=%.1fGiB, rules=%d",
self._user_max_batch_size,
self._device_memory_gb or 0.0,
len(self._rules),
)
@property
def enabled(self) -> bool:
return self._mode == "dynamic" and self._user_max_batch_size > 1
def reject_reason_for_candidate(
self, current_reqs: list[Req], candidate_req: Req
) -> str | None:
if not self.enabled:
return None
proposed = current_reqs + [candidate_req]
limit = self.limit_for(proposed[0])
return limit.reject_reason(
batch_size=len(proposed),
batch_cost=self.estimate_batch_cost(proposed),
)
def batch_is_full(self, reqs: list[Req]) -> bool:
"""Return whether another roughly similar request would exceed the cap."""
if not self.enabled or not reqs:
return len(reqs) >= self._user_max_batch_size
limit = self.limit_for(reqs[0])
if len(reqs) >= limit.max_batch_size:
return True
next_cost = self.estimate_batch_cost(reqs + [reqs[0]])
return limit.max_cost is not None and next_cost > limit.max_cost
def limit_reason_for_batch(self, reqs: list[Req]) -> str | None:
if not self.enabled or not reqs:
return None
limit = self.limit_for(reqs[0])
if len(reqs) >= limit.max_batch_size:
return limit.cap_reason or f"config_cap:{limit.max_batch_size}"
next_cost = self.estimate_batch_cost(reqs + [reqs[0]])
return limit.stop_reason_for_next_cost(next_cost)
def max_admissible_batch_size(self, req: Req) -> int:
return self.limit_for(req).max_batch_size
def limit_for(self, req: Req) -> AdmissionLimit:
"""Return the effective admission limit for the request's model and shape."""
rules = self._matching_rules(req)
if not rules:
return AdmissionLimit(max_batch_size=self._user_max_batch_size)
config_cap = min(rule.max_batch_size for rule in rules)
max_batch_size = min(self._user_max_batch_size, config_cap)
cap_reason = (
f"config_cap:{max_batch_size}"
if max_batch_size < self._user_max_batch_size
else None
)
costs = [rule.max_cost for rule in rules if rule.max_cost is not None]
return AdmissionLimit(
max_batch_size=max(1, max_batch_size),
max_cost=min(costs) if costs else None,
cap_reason=cap_reason,
)
def estimate_batch_cost(self, reqs: list[Req]) -> float:
return sum(
float(self._pipeline_config.estimate_request_cost(req)) for req in reqs
)
def _matching_rules(self, req: Req) -> list[BatchingRule]:
return [
rule
for rule in self._rules
if rule.matches(
model_path=self._model_path,
resolution=req.resolution_key,
device_memory_gb=self._device_memory_gb,
offload=self._offload,
)
]
@staticmethod
def _get_device_memory_gb(gpu_id: int) -> float | None:
try:
return current_platform.get_device_total_memory(gpu_id) / BYTES_PER_GB
except Exception:
return None
def load_batching_config(path: str | None) -> list[BatchingRule]:
if path is None:
return []
with open(path, encoding="utf-8") as f:
payload = json.load(f)
source = os.path.abspath(path)
entries = _config_entries(payload)
rules = [BatchingRule.from_dict(entry, source=source) for entry in entries]
if not rules:
raise ValueError(f"batching config {source} does not contain any rules")
return rules
def _config_entries(payload: Any) -> list[dict[str, Any]]:
if isinstance(payload, dict) and payload.get("schema_version") not in (None, 1):
raise ValueError("batching config schema_version must be 1")
if isinstance(payload, dict) and isinstance(payload.get("rules"), list):
return payload["rules"]
if isinstance(payload, list):
return payload
if isinstance(payload, dict):
entries: list[dict[str, Any]] = []
for key, value in payload.items():
if key == "schema_version" or not isinstance(value, dict):
continue
model, _sep, resolution = key.partition("|")
entry = dict(value)
if model:
entry.setdefault("model", model)
if resolution:
entry.setdefault("resolution", resolution)
entries.append(entry)
return entries
raise ValueError(
"batching config must be a {'schema_version': 1, 'rules': [...]} object, "
"a list of rules, or a mapping keyed by model|resolution"
)
def _validate_rule_keys(data: dict[str, Any], *, source: str) -> None:
unknown = sorted(set(data) - _BATCHING_RULE_KEYS)
if not unknown:
return
hints = []
for key in unknown:
matches = get_close_matches(key, _BATCHING_RULE_KEYS, n=1)
if matches:
hints.append(f"{key!r} (did you mean {matches[0]!r}?)")
else:
hints.append(repr(key))
raise ValueError(
f"batching config rule from {source} contains unknown key(s): "
f"{', '.join(hints)}"
)
def _optional_str(value: Any) -> str | None:
if value is None:
return None
return str(value)
def _optional_float(value: Any) -> float | None:
if value is None:
return None
return float(value)
def _optional_bool(value: Any) -> bool | None:
if value is None:
return None
if isinstance(value, bool):
return value
if isinstance(value, str):
lowered = value.strip().lower()
if lowered in ("1", "true", "yes", "y", "on"):
return True
if lowered in ("0", "false", "no", "n", "off"):
return False
raise ValueError(f"cannot parse boolean batching config value: {value!r}")
python/sglang/multimodal_gen/runtime/managers/gpu_worker.py
+29 -1
View File
@@ -318,6 +318,7 @@ class GPUWorker:
return result
output_batch = self._to_output_batch(result)
self._record_output_peak_memory(output_batch)
output_metrics = self._iter_output_metrics(output_batch)
if self.rank == 0 and output_metrics and not current_platform.is_cpu():
@@ -377,8 +378,15 @@ class GPUWorker:
if output_batch is None:
output_batch = OutputBatch()
output_batch.error = f"Error executing {error_context}: {e}"
self._record_output_peak_memory(output_batch)
return output_batch
def _record_output_peak_memory(self, output_batch: OutputBatch) -> None:
if self.rank != 0 or current_platform.is_cpu():
return
peak_reserved_bytes = torch.get_device_module().max_memory_reserved()
output_batch.peak_memory_mb = peak_reserved_bytes / (1024**2)
def _forward_group(self, batch: list[Req]) -> OutputBatch:
assert self.pipeline is not None
results = self.pipeline.forward_batch(batch, self.server_args)
@@ -389,12 +397,32 @@ class GPUWorker:
if self.rank != 0 or output_batch.output is None:
return
dynamic_output_paths = None
if req.extra:
dynamic_output_paths = req.extra.get("dynamic_batch_output_paths")
if dynamic_output_paths is not None and (
len(dynamic_output_paths) != len(output_batch.output)
):
logger.warning(
"dynamic_batch_output_paths length mismatch (got=%d, expected=%d). "
"Falling back to merged request output file naming.",
len(dynamic_output_paths),
len(output_batch.output),
)
dynamic_output_paths = None
if dynamic_output_paths is not None:
build_output_path = lambda idx: dynamic_output_paths[idx]
else:
num_outputs = len(output_batch.output)
build_output_path = lambda idx: req.output_file_path(num_outputs, idx)
output_batch.output_file_paths = save_outputs(
output_batch.output,
req.data_type,
req.fps,
True,
lambda idx: req.output_file_path(len(output_batch.output), idx),
build_output_path,
audio=output_batch.audio,
audio_sample_rate=output_batch.audio_sample_rate,
output_compression=req.output_compression,
python/sglang/multimodal_gen/runtime/managers/scheduler.py
+673 -49
View File
@@ -2,10 +2,14 @@
# SPDX-License-Identifier: Apache-2.0
import asyncio
import dataclasses
import os
import pickle
import tempfile
import time
from collections import deque
from copy import deepcopy
from enum import Enum
from typing import Any, List
import zmq
@@ -32,9 +36,15 @@ from sglang.multimodal_gen.runtime.entrypoints.utils import (
UnmergeLoraWeightsReq,
)
from sglang.multimodal_gen.runtime.managers.cpu_worker import CPUWorker
from sglang.multimodal_gen.runtime.managers.dynamic_batch_admission import (
BatchAdmissionController,
)
from sglang.multimodal_gen.runtime.managers.gpu_worker import GPUWorker
from sglang.multimodal_gen.runtime.pipelines_core import Req
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import (
BatchMetricsWindow,
OutputBatch,
)
from sglang.multimodal_gen.runtime.server_args import (
PortArgs,
ServerArgs,
@@ -56,6 +66,9 @@ MINIMUM_PICTURE_BASE64_FOR_WARMUP = "data:image/jpg;base64,iVBORw0KGgoAAAANSUhEU
# _combine_cfg_parallel's all-reduce.
DEFAULT_PLACEHOLDER_PROMPT = "warmup"
_MAX_RECV_REQS_PER_POLL = 1024
_BATCH_METRICS_LOG_INTERVAL = 5
class Scheduler(SchedulerDisaggMixin):
"""
@@ -122,8 +135,16 @@ class Scheduler(SchedulerDisaggMixin):
GetWeightsChecksumReqInput: self._handle_get_weights_checksum,
}
# FIFO, new reqs are appended
self.waiting_queue: deque[tuple[bytes, Any]] = deque()
# FIFO queue entries: (identity, request, enqueue_ts_s)
self.waiting_queue: deque[tuple[bytes | None, Any, float]] = deque()
self._batching_max_size = server_args.batching_max_size
self._batching_delay_s = server_args.batching_delay_ms / 1000.0
self._batch_metrics_enabled = server_args.enable_batching_metrics
self._batch_metrics_window = BatchMetricsWindow()
self._batch_admission = BatchAdmissionController(server_args, gpu_id=local_rank)
self._poller = zmq.Poller()
if self.receiver is not None:
self._poller.register(self.receiver, zmq.POLLIN)
# whether we've send the necessary warmup reqs
self.warmed_up = False
@@ -139,6 +160,12 @@ class Scheduler(SchedulerDisaggMixin):
self._init_disagg_state(server_args, local_rank)
if self._batch_metrics_enabled:
logger.info(
"Dynamic batch metrics enabled; logging summary every %d dispatches.",
_BATCH_METRICS_LOG_INTERVAL,
)
def get_disagg_metrics(self) -> dict | None:
"""Return disagg role metrics snapshot, or None if not in disagg mode."""
if self._disagg_metrics is None:
@@ -216,31 +243,49 @@ class Scheduler(SchedulerDisaggMixin):
req = cls._first_generation_req(req_or_group)
return req.is_warmup if req is not None else False
def _dispatch_request(self, reqs: list[Any]) -> OutputBatch:
"""dispatch req to its registered handler"""
req_or_group = reqs[0]
def _dispatch_single_request(self, req_or_group: Any) -> OutputBatch:
if isinstance(req_or_group, list):
return self._handle_generation(reqs)
if not all(isinstance(req, Req) for req in req_or_group):
return OutputBatch(
error=f"Unknown request group type: {type(req_or_group)}"
)
return self._handle_generation(req_or_group, allow_dynamic_batching=False)
handler = self.request_handlers.get(type(req_or_group))
if handler is None:
return OutputBatch(error=f"Unknown request type: {type(req_or_group)}")
return handler(reqs)
return handler([req_or_group])
def _dispatch_items(
self, items: list[tuple[bytes | None, Any]]
) -> OutputBatch | list[OutputBatch]:
"""Dispatch ready queue items; several plain `Req`s form one dynamic batch."""
reqs = [item[1] for item in items]
if len(reqs) > 1 and all(isinstance(req, Req) for req in reqs):
return self._handle_generation(reqs, allow_dynamic_batching=True)
if len(reqs) > 1:
return [self._dispatch_single_request(req) for req in reqs]
return self._dispatch_single_request(reqs[0])
def _log_warmup_result(self, output_batch: OutputBatch, is_warmup: bool) -> None:
if not is_warmup:
return
if output_batch.error is None:
total_duration_s = (
output_batch.metrics.total_duration_s
if output_batch.metrics is not None
else 0.0
)
if self._warmup_total > 0:
logger.info(
f"Warmup req ({self._warmup_processed}/{self._warmup_total}) processed in {GREEN}%.2f{RESET} seconds",
output_batch.metrics.total_duration_s,
total_duration_s,
)
else:
logger.info(
f"Warmup req processed in {GREEN}%.2f{RESET} seconds",
output_batch.metrics.total_duration_s,
total_duration_s,
)
else:
if self._warmup_total > 0:
@@ -250,7 +295,10 @@ class Scheduler(SchedulerDisaggMixin):
else:
logger.info("Warmup req processing failed")
def _handle_generation(self, reqs: list[Any]):
def _handle_generation(
self, reqs: list[Any], *, allow_dynamic_batching: bool = True
):
"""Dispatch generation requests, merging compatible requests when allowed."""
reqs = self._normalize_generation_reqs(reqs)
warmup_reqs = [req for req in reqs if req.is_warmup]
if warmup_reqs:
@@ -262,8 +310,7 @@ class Scheduler(SchedulerDisaggMixin):
else:
logger.info("Processing warmup req...")
# Diffusion dispatches one generation request at a time, so reqs[0]
# always carries the trace context for the entire batch.
# Use the head request trace context for scheduler-side dispatch work.
req = reqs[0]
req.trace_ctx.rebuild_thread_context()
with trace_slice(
@@ -271,7 +318,285 @@ class Scheduler(SchedulerDisaggMixin):
DiffStage.SCHEDULER_DISPATCH,
thread_finish_flag=True,
):
return self.worker.execute_forward(reqs)
if len(reqs) == 1 or not allow_dynamic_batching:
return self.worker.execute_forward(reqs)
merged_req = self._try_merge_generation_reqs(reqs)
if merged_req is None:
return self._execute_generation_sequential(reqs)
batch_size = len(reqs)
try:
output_batch = self.worker.execute_forward([merged_req])
if output_batch.error:
logger.error(
"Dynamic batch execution returned error. Skipping sequential fallback and returning errors: %s",
output_batch.error,
)
return self._build_dynamic_batch_error_outputs(
reqs=reqs,
error_msg=output_batch.error,
)
split_outputs = self._split_batched_output(output_batch, reqs)
if split_outputs is None:
logger.error(
"Failed to split dynamic batched output cleanly. Skipping sequential fallback and returning errors."
)
return self._build_dynamic_batch_error_outputs(
reqs=reqs,
error_msg="Dynamic batching failed: could not split merged output.",
)
logger.info(
"Processed dynamic batch of %d/%d request(s) with max_delay=%.2fms",
batch_size,
self._batching_max_size,
self._batching_delay_s * 1000.0,
)
return split_outputs
except Exception as e:
logger.error(
"Dynamic batching failed (%s). Skipping sequential fallback and returning errors.",
e,
exc_info=True,
)
return self._build_dynamic_batch_error_outputs(
reqs=reqs,
error_msg=f"Dynamic batching failed: {e}",
)
def _execute_generation_sequential(self, reqs: List[Req]) -> List[OutputBatch]:
return [self.worker.execute_forward([req]) for req in reqs]
@staticmethod
def _percentile(values: list[float], percentile: float) -> float:
if not values:
return 0.0
ordered = sorted(values)
index = min(
len(ordered) - 1,
max(0, int(round((percentile / 100.0) * (len(ordered) - 1)))),
)
return ordered[index]
def _freeze_signature_value(self, value: Any):
"""Convert a value into a hashable, order-stable form for signature comparison."""
if isinstance(value, (str, int, float, bool, type(None))):
return value
if isinstance(value, Enum):
return value.value
if isinstance(value, dict):
return {
str(k): self._freeze_signature_value(v)
for k, v in sorted(value.items(), key=lambda kv: str(kv[0]))
}
if isinstance(value, (list, tuple)):
return tuple(self._freeze_signature_value(v) for v in value)
return repr(value)
def _sampling_param_signature_items(self, req: Req) -> list[tuple[str, Any]] | None:
"""Return per-field sampling-param signature items, skipping batch_sig_exclude fields."""
sp = req.sampling_params
if sp is None:
return None
try:
sp_fields = dataclasses.fields(sp)
except Exception:
return None
return [
(f.name, self._freeze_signature_value(getattr(sp, f.name, None)))
for f in sp_fields
if not f.metadata.get("batch_sig_exclude", False)
]
def _diffusers_kwargs_signature_value(self, req: Req) -> Any:
return self._freeze_signature_value((req.extra or {}).get("diffusers_kwargs"))
def _build_dynamic_batch_signature(self, req: Req) -> tuple[Any, ...] | None:
"""Build the request compatibility signature for dynamic batching.
The signature is built from `SamplingParams` fields, excluding fields
marked with `batch_sig_exclude`, plus generation-affecting
`extra.diffusers_kwargs`.
"""
signature_items = self._sampling_param_signature_items(req)
if signature_items is None:
return None
if req.extra:
diffusers_kwargs = req.extra.get("diffusers_kwargs")
if diffusers_kwargs:
signature_items.append(
(
"diffusers_kwargs",
self._freeze_signature_value(diffusers_kwargs),
)
)
return tuple(signature_items)
def _get_cached_signature(self, req: Req) -> tuple[Any, ...] | None:
cached = getattr(req, "_dynamic_batch_sig", None)
if cached is not None:
return cached
sig = self._build_dynamic_batch_signature(req)
req._dynamic_batch_sig = sig # type: ignore[attr-defined]
return sig
def _find_sampling_param_mismatch_field(
self, base_req: Req, candidate_req: Req
) -> str | None:
base_items = self._sampling_param_signature_items(base_req)
candidate_items = self._sampling_param_signature_items(candidate_req)
if base_items is None or candidate_items is None:
return None
if len(base_items) != len(candidate_items):
return "sampling_params"
for (name, base_value), (candidate_name, candidate_value) in zip(
base_items, candidate_items
):
if name != candidate_name:
return "sampling_params"
if base_value != candidate_value:
return f"sampling_params.{name}"
base_diffusers_kwargs = self._diffusers_kwargs_signature_value(base_req)
candidate_diffusers_kwargs = self._diffusers_kwargs_signature_value(
candidate_req
)
if base_diffusers_kwargs != candidate_diffusers_kwargs:
return "extra.diffusers_kwargs"
return None
def _get_dynamic_batch_reject_reason(
self, base_req: Req, candidate_req: Req
) -> str | None:
"""Return the first reason `candidate_req` cannot batch with `base_req`, or None."""
if self._can_dynamic_batch(base_req, candidate_req):
return None
if base_req.is_warmup or candidate_req.is_warmup:
return "warmup"
if not isinstance(base_req.prompt, str) or not isinstance(
candidate_req.prompt, str
):
return "prompt_type"
if base_req.image_path is not None or candidate_req.image_path is not None:
return "image_conditioning"
if base_req.return_file_paths_only != candidate_req.return_file_paths_only:
return "return_file_paths_only"
base_sig = self._get_cached_signature(base_req)
candidate_sig = self._get_cached_signature(candidate_req)
if base_sig is None or candidate_sig is None:
return "signature_unavailable"
return (
self._find_sampling_param_mismatch_field(base_req, candidate_req)
or "signature_mismatch"
)
def _can_dynamic_batch(self, base_req: Req, candidate_req: Req) -> bool:
"""Return whether `candidate_req` can be merged into a batch with `base_req`."""
if base_req.is_warmup or candidate_req.is_warmup:
return False
if not isinstance(base_req.prompt, str) or not isinstance(
candidate_req.prompt, str
):
return False
if base_req.image_path is not None or candidate_req.image_path is not None:
return False
if base_req.return_file_paths_only != candidate_req.return_file_paths_only:
return False
base_sig = self._get_cached_signature(base_req)
cand_sig = self._get_cached_signature(candidate_req)
return base_sig is not None and base_sig == cand_sig
def _record_batch_dispatch_metrics(
self,
batch_size: int,
queue_wait_ms: float,
effective_max_batch_size: int,
reject_reasons: list[str] | None = None,
stop_reason: str | None = None,
) -> None:
if not self._batch_metrics_enabled:
return
effective_max_batch_size = max(1, effective_max_batch_size)
logger.info(
"Dynamic batch dispatch: size=%d/%d, user_max=%d, queue_wait=%.2fms, stop_reason=%s",
batch_size,
effective_max_batch_size,
self._batching_max_size,
max(queue_wait_ms, 0.0),
stop_reason or "unspecified",
)
window = self._batch_metrics_window
window.dispatches += 1
window.total_requests += batch_size
window.total_capacity += effective_max_batch_size
if batch_size > 1:
window.merged_dispatches += 1
if self._dynamic_batching_enabled() and batch_size >= effective_max_batch_size:
window.full_dispatches += 1
window.wait_times_ms.append(max(queue_wait_ms, 0.0))
if reject_reasons:
window.reject_reasons.update(reject_reasons)
if window.dispatches >= _BATCH_METRICS_LOG_INTERVAL:
self._log_batch_metrics_summary()
def _log_batch_metrics_summary(self) -> None:
if not self._batch_metrics_enabled:
return
window = self._batch_metrics_window
if window.dispatches == 0:
return
avg_size = window.total_requests / window.dispatches
utilization = window.total_requests / max(1, window.total_capacity)
avg_wait_ms = sum(window.wait_times_ms) / len(window.wait_times_ms)
p95_wait_ms = self._percentile(window.wait_times_ms, 95.0)
merged_rate = window.merged_dispatches / window.dispatches
full_rate = window.full_dispatches / window.dispatches
top_rejects = ", ".join(
f"{reason}={count}"
for reason, count in window.reject_reasons.most_common(5)
)
if not top_rejects:
top_rejects = "none"
logger.info(
"Dynamic batch stats (last %d dispatches): avg_size=%.2f, merged_rate=%.1f%%, full_rate=%.1f%%, utilization=%.1f%%, wait_avg=%.2fms, wait_p95=%.2fms, top_rejects=%s",
window.dispatches,
avg_size,
merged_rate * 100.0,
full_rate * 100.0,
utilization * 100.0,
avg_wait_ms,
p95_wait_ms,
top_rejects,
)
self._batch_metrics_window = BatchMetricsWindow()
def _build_dynamic_batch_error_outputs(
self,
reqs: List[Req],
error_msg: str,
) -> List[OutputBatch]:
return [OutputBatch(error=error_msg) for _ in reqs]
def return_result(
self,
@@ -285,15 +610,265 @@ class Scheduler(SchedulerDisaggMixin):
if not is_warmup and self.receiver is not None and identity is not None:
self.receiver.send_multipart([identity, b"", pickle.dumps(output_batch)])
def get_next_batch_to_run(self) -> list[tuple[bytes, Any]] | None:
"""pull a req from waiting_queue"""
def _try_merge_generation_reqs(self, reqs: List[Req]) -> Req | None:
"""Create a batched generation request from compatible requests.
Per-request seeds and output paths are stored in `extra` so downstream
stages can preserve request ordering.
"""
if len(reqs) <= 1:
return reqs[0] if reqs else None
base_req = reqs[0]
for req in reqs[1:]:
if not self._can_dynamic_batch(base_req, req):
return None
merged_req = deepcopy(base_req)
merged_req.prompt = [req.prompt for req in reqs]
merged_req.extra = deepcopy(merged_req.extra)
merged_req.extra["dynamic_batch_seeds"] = [req.seed for req in reqs]
merged_req.return_file_paths_only = base_req.return_file_paths_only
if merged_req.return_file_paths_only:
dynamic_output_paths: list[str] = []
for req in reqs:
for output_idx in range(req.num_outputs_per_prompt):
dynamic_output_paths.append(
req.output_file_path(req.num_outputs_per_prompt, output_idx)
)
merged_req.extra["dynamic_batch_output_paths"] = dynamic_output_paths
merged_req.request_id = f"dynamic_batch::{merged_req.request_id}"
return merged_req
@staticmethod
def _count_first_dim(value: Any) -> int | None:
if value is None:
return None
if isinstance(value, (list, tuple)):
return len(value)
shape = getattr(value, "shape", None)
if shape is not None:
try:
if len(shape) > 0:
return int(shape[0])
except Exception:
return None
return None
def _slice_batched_value(
self, value: Any, start: int, end: int, total_items: int
) -> Any:
if value is None:
return None
if isinstance(value, (list, tuple)):
if len(value) == total_items:
sliced = value[start:end]
return list(sliced) if isinstance(value, list) else tuple(sliced)
return deepcopy(value)
value_items = self._count_first_dim(value)
if value_items == total_items:
try:
return value[start:end]
except Exception:
pass
# Scalar / non-batched metadata
return deepcopy(value)
def _split_batched_output(
self, output_batch: OutputBatch, reqs: List[Req]
) -> List[OutputBatch] | None:
"""Split a merged result only when outputs map one-to-one to requests."""
per_req_counts = [req.num_outputs_per_prompt for req in reqs]
total_items = sum(per_req_counts)
output_items = self._count_first_dim(output_batch.output)
output_path_items = self._count_first_dim(output_batch.output_file_paths)
if output_items is None and output_path_items is None:
logger.warning(
"Batched output has neither tensor outputs nor output_file_paths; cannot split safely."
)
return None
if output_items is not None and output_items != total_items:
logger.warning(
"Unexpected batched output size: got %s items, expected %s",
output_items,
total_items,
)
return None
if output_path_items is not None and output_path_items != total_items:
logger.warning(
"Unexpected batched output_file_paths size: got %s items, expected %s",
output_path_items,
total_items,
)
return None
outputs: list[OutputBatch] = []
start = 0
for req, req_count in zip(reqs, per_req_counts):
end = start + req_count
split = OutputBatch(
output=self._slice_batched_value(
output_batch.output, start, end, total_items
),
audio=self._slice_batched_value(
output_batch.audio, start, end, total_items
),
audio_sample_rate=output_batch.audio_sample_rate,
trajectory_timesteps=self._slice_batched_value(
output_batch.trajectory_timesteps, start, end, total_items
),
trajectory_latents=self._slice_batched_value(
output_batch.trajectory_latents, start, end, total_items
),
trajectory_decoded=self._slice_batched_value(
output_batch.trajectory_decoded, start, end, total_items
),
error=output_batch.error,
output_file_paths=self._slice_batched_value(
output_batch.output_file_paths, start, end, total_items
),
metrics=deepcopy(output_batch.metrics),
noise_pred=self._slice_batched_value(
output_batch.noise_pred, start, end, total_items
),
peak_memory_mb=output_batch.peak_memory_mb,
)
if split.metrics is not None:
split.metrics.request_id = req.request_id
outputs.append(split)
start = end
return outputs
def _dynamic_batching_enabled(self) -> bool:
"""Return whether this server and pipeline can use dynamic batching.
This is the coarse gate; request-level checks decide which requests can
actually be merged.
"""
pipeline_config = self.server_args.pipeline_config
supports_dynamic_batching = getattr(
pipeline_config, "supports_dynamic_batching", None
)
if callable(supports_dynamic_batching):
return self._batch_admission.enabled and supports_dynamic_batching()
return self._batch_admission.enabled
def get_next_batch_to_run(self) -> list[tuple[bytes | None, Any]] | None:
"""Return the next dispatchable queue item or dynamic batch.
Returns None when the head request is waiting for more compatible
requests within the configured batching delay.
"""
if not self.waiting_queue:
return None
# pop the first (earliest)
item = self.waiting_queue.popleft()
if not self._dynamic_batching_enabled():
identity, req, enqueue_time = self.waiting_queue.popleft()
if isinstance(req, Req):
self._record_batch_dispatch_metrics(
batch_size=1,
queue_wait_ms=(time.monotonic() - enqueue_time) * 1000.0,
effective_max_batch_size=1,
stop_reason="dynamic_disabled",
)
return [(identity, req)]
return [item]
identity, req, enqueue_time = self.waiting_queue[0]
if not isinstance(req, Req):
identity, req, _ = self.waiting_queue.popleft()
return [(identity, req)]
# If the head request itself is not eligible for dynamic batching
# (e.g., image-conditioned i2i request), dispatch it immediately.
if not self._can_dynamic_batch(req, req):
identity, req, head_enqueue_time = self.waiting_queue.popleft()
reject_reasons: list[str] = []
if self._batch_metrics_enabled:
reason = self._get_dynamic_batch_reject_reason(req, req)
if reason is not None:
reject_reasons.append(f"head:{reason}")
self._record_batch_dispatch_metrics(
batch_size=1,
queue_wait_ms=(time.monotonic() - head_enqueue_time) * 1000.0,
effective_max_batch_size=1,
reject_reasons=reject_reasons,
stop_reason=reject_reasons[0] if reject_reasons else "head_ineligible",
)
return [(identity, req)]
compatible_indices: list[int] = [0]
compatible_reqs: list[Req] = [req]
reject_reasons: list[str] = []
for idx in range(1, len(self.waiting_queue)):
if len(
compatible_indices
) >= self._batching_max_size or self._batch_admission.batch_is_full(
compatible_reqs
):
break
_identity, candidate_req, _enqueue_time = self.waiting_queue[idx]
if isinstance(candidate_req, Req) and self._can_dynamic_batch(
req, candidate_req
):
admission_reject = self._batch_admission.reject_reason_for_candidate(
compatible_reqs, candidate_req
)
if admission_reject is None:
compatible_indices.append(idx)
compatible_reqs.append(candidate_req)
elif self._batch_metrics_enabled:
reject_reasons.append(admission_reject)
elif self._batch_metrics_enabled and isinstance(candidate_req, Req):
reason = self._get_dynamic_batch_reject_reason(req, candidate_req)
if reason is not None:
reject_reasons.append(reason)
batch_len = len(compatible_indices)
oldest_wait_s = time.monotonic() - enqueue_time
should_wait_for_more = (
batch_len < self._batching_max_size
and not self._batch_admission.batch_is_full(compatible_reqs)
and oldest_wait_s < self._batching_delay_s
)
if should_wait_for_more:
return None
batch_items: list[tuple[bytes | None, Any]] = [None] * batch_len
for pos, idx in enumerate(reversed(compatible_indices)):
item_identity, item_req, _ = self.waiting_queue[idx]
batch_items[batch_len - 1 - pos] = (item_identity, item_req)
del self.waiting_queue[idx]
stop_reason = self._batch_admission.limit_reason_for_batch(compatible_reqs)
if stop_reason is None:
if batch_len >= self._batching_max_size:
stop_reason = "max_size"
elif reject_reasons:
stop_reason = reject_reasons[0]
elif oldest_wait_s >= self._batching_delay_s:
stop_reason = "delay"
else:
stop_reason = "ready"
self._record_batch_dispatch_metrics(
batch_size=batch_len,
queue_wait_ms=oldest_wait_s * 1000.0,
effective_max_batch_size=self._batch_admission.max_admissible_batch_size(
compatible_reqs[0]
),
reject_reasons=reject_reasons,
stop_reason=stop_reason,
)
return batch_items
def prepare_server_warmup_reqs(self):
if (
@@ -334,7 +909,7 @@ class Scheduler(SchedulerDisaggMixin):
req_kwargs["do_classifier_free_guidance"] = True
req = Req(**req_kwargs)
req.set_as_warmup(self.server_args.warmup_steps)
self.waiting_queue.append((None, req))
self.waiting_queue.append((None, req, time.monotonic()))
# if server is warmed-up, set this flag to avoid req-based warmup
self.warmed_up = True
@@ -408,34 +983,51 @@ class Scheduler(SchedulerDisaggMixin):
self.warmed_up = True
return recv_reqs
@staticmethod
def _normalize_received_payload(
identity: bytes, reqs: Any
) -> list[tuple[bytes, Any]]:
"""Normalize client payloads into queue entries.
A single-item `[Req]` is one request; a multi-item `list[Req]` remains
grouped as one logical request.
"""
if not isinstance(reqs, list):
return [(identity, reqs)]
if not reqs:
return []
if all(isinstance(req, Req) for req in reqs):
# AsyncSchedulerClient sends ordinary single requests as [Req].
# Only multi-item list[Req] payloads represent a grouped multi-output request.
if len(reqs) == 1:
return [(identity, reqs[0])]
return [(identity, reqs)]
return [(identity, req) for req in reqs]
def recv_reqs(self) -> List[tuple[bytes, Any]]:
"""
For non-main schedulers, reqs are broadcasted from main using broadcast_pyobj
"""
if self.receiver is not None:
try:
try:
# Accept valid REQ envelopes only, ignore malformed/probe frames.
parts = self.receiver.recv_multipart(zmq.NOBLOCK)
identity, payload = parts[0], parts[-1]
recv_reqs: list[tuple[bytes, Any]] = []
while len(recv_reqs) < _MAX_RECV_REQS_PER_POLL:
try:
# Accept valid REQ envelopes only, ignore malformed/probe frames.
parts = self.receiver.recv_multipart(zmq.NOBLOCK)
except zmq.Again:
break
# Ignore malformed probes or non-pickle data
recv_reqs = pickle.loads(payload) if len(parts) > 2 else []
except (zmq.Again, pickle.UnpicklingError, IndexError, EOFError):
recv_reqs = []
try:
identity, payload = parts[0], parts[-1]
reqs = pickle.loads(payload) if len(parts) > 2 else []
except (pickle.UnpicklingError, IndexError, EOFError):
continue
recv_reqs.extend(self._normalize_received_payload(identity, reqs))
except zmq.ZMQError:
# re-raise or handle appropriately to let the outer loop continue
raise
if recv_reqs:
if isinstance(recv_reqs, list) and all(
isinstance(req, Req) for req in recv_reqs
):
recv_reqs = [(identity, recv_reqs)]
else:
if not isinstance(recv_reqs, list):
recv_reqs = [recv_reqs]
recv_reqs = [(identity, req) for req in recv_reqs]
else:
recv_reqs = None
@@ -491,7 +1083,10 @@ class Scheduler(SchedulerDisaggMixin):
try:
new_reqs = self.recv_reqs()
new_reqs = self.process_received_reqs_with_req_based_warmup(new_reqs)
self.waiting_queue.extend(new_reqs)
now = time.monotonic()
self.waiting_queue.extend(
[(identity, req, now) for identity, req in new_reqs]
)
# Reset error count on success
self._consecutive_error_count = 0
except Exception as e:
@@ -515,33 +1110,62 @@ class Scheduler(SchedulerDisaggMixin):
# 2: execute, make sure a reply is always sent
items = self.get_next_batch_to_run()
if not items:
if self.waiting_queue and self._dynamic_batching_enabled():
oldest_ts = self.waiting_queue[0][2]
elapsed_ms = (time.monotonic() - oldest_ts) * 1000.0
remaining_ms = max(0, self._batching_delay_s * 1000.0 - elapsed_ms)
if remaining_ms > 0 and self.receiver is not None:
self._poller.poll(timeout=remaining_ms)
elif remaining_ms > 0:
time.sleep(remaining_ms / 1000.0)
continue
identities = [item[0] for item in items]
reqs = [item[1] for item in items]
try:
req_or_group = reqs[0]
is_warmup = self._is_warmup_item(req_or_group)
output_batch = self._dispatch_request(reqs)
handler_result = self._dispatch_items(items)
except Exception as e:
logger.error(
f"Error executing request in scheduler event loop: {e}",
exc_info=True,
)
output_batch = OutputBatch(error=str(e))
handler_result = OutputBatch(error=str(e))
if isinstance(handler_result, list):
output_batches = handler_result
else:
output_batches = [handler_result]
if len(output_batches) != len(items):
logger.error(
"Handler returned %d output(s) for %d request(s). Returning error for unmatched requests.",
len(output_batches),
len(items),
)
output_batches = [
OutputBatch(
error=(
f"Internal scheduler error: expected {len(items)} outputs, "
f"got {len(output_batches)}."
)
)
for _ in items
]
# 3. return results
try:
self._log_warmup_result(output_batch, is_warmup)
for (identity, processed_req), output_batch in zip(
items, output_batches, strict=True
):
is_warmup = self._is_warmup_item(processed_req)
self._log_warmup_result(output_batch, is_warmup)
# TODO: Support sending back to multiple identities if batched
self.return_result(output_batch, identities[0], is_warmup=is_warmup)
self.return_result(output_batch, identity, is_warmup=is_warmup)
except zmq.ZMQError as e:
# Reply failed; log and keep loop alive to accept future requests
logger.error(f"ZMQ error sending reply: {e}")
continue
self._log_batch_metrics_summary()
if self.receiver is not None:
self.receiver.close()
self._cleanup_disagg()
python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
+20 -1
View File
@@ -642,10 +642,19 @@ class QwenImageCrossAttention(nn.Module):
image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
**cross_attention_kwargs,
):
"""Run joint text-image attention.
`attn_mask` or `attention_mask` takes precedence. Otherwise,
`encoder_hidden_states_mask` keeps valid text tokens in the joint
text-image sequence.
"""
seq_len_txt = encoder_hidden_states.shape[1]
attn_mask = cross_attention_kwargs.get("attn_mask")
if attn_mask is None:
attn_mask = cross_attention_kwargs.get("attention_mask")
encoder_hidden_states_mask = cross_attention_kwargs.get(
"encoder_hidden_states_mask"
)
img_query, img_key, img_value, txt_query, txt_key, txt_value = (
_get_qkv_projections(self, hidden_states, encoder_hidden_states)
@@ -704,6 +713,16 @@ class QwenImageCrossAttention(nn.Module):
joint_query = torch.cat([txt_query, img_query], dim=1)
joint_key = torch.cat([txt_key, img_key], dim=1)
joint_value = torch.cat([txt_value, img_value], dim=1)
if attn_mask is None and encoder_hidden_states_mask is not None:
image_mask = torch.ones(
(hidden_states.shape[0], img_query.shape[1]),
device=encoder_hidden_states_mask.device,
dtype=torch.bool,
)
attn_mask = torch.cat(
[encoder_hidden_states_mask.to(dtype=torch.bool), image_mask],
dim=1,
)
# Compute joint attention
joint_hidden_states = self.attn(
@@ -1306,7 +1325,7 @@ class QwenImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_sequence_length, joint_attention_dim)`):
Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
encoder_hidden_states_mask (`torch.Tensor` of shape `(batch_size, text_sequence_length)`):
Mask of the input conditions.
Valid-token mask of the input conditions, where True keeps a text token.
timestep ( `torch.LongTensor`):
Used to indicate denoising step.
attention_kwargs (`dict`, *optional*):
python/sglang/multimodal_gen/runtime/models/dits/zimage.py
+104 -56
View File
@@ -780,65 +780,123 @@ class ZImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
f_patch_size: int,
image_seq_len_target: int | None = None,
):
assert len(all_image) == len(all_cap_feats) == 1
"""Patchify images and pad image/caption tokens to batch targets.
Each image is [C, F, H, W] and has one [L, D] caption. Returned tensors
are stacked as [B, S, D], while valid lengths keep track of real tokens
before learned pad tokens are restored. `image_seq_len_target`, when
set, is the SP-local padded image-token target.
"""
if len(all_image) != len(all_cap_feats):
raise ValueError(
f"Z-Image expects one caption embedding per image, got {len(all_image)} images and {len(all_cap_feats)} captions"
)
if not all_image:
raise ValueError("Z-Image batch must contain at least one image latent")
image = all_image[0] # C, F, H, W
cap_feat = all_cap_feats[0] # L, D
pH = pW = patch_size
pF = f_patch_size
device = image.device
all_image_out = []
all_image_size = []
all_cap_feats_out = []
all_image_valid_lens = []
all_cap_valid_lens = []
image_records = []
# ------------ Process Caption ------------
cap_ori_len = cap_feat.size(0)
cap_padding_len = (-cap_ori_len) % SEQ_MULTI_OF
# padded feature
cap_padded_feat = torch.cat(
[cap_feat, cap_feat[-1:].repeat(cap_padding_len, 1)],
dim=0,
cap_seq_len_target = max(
self._ceil_to_multiple(cap_feat.size(0), SEQ_MULTI_OF)
for cap_feat in all_cap_feats
)
all_cap_feats_out.append(cap_padded_feat)
all_cap_valid_lens.append(cap_ori_len)
# ------------ Process Image ------------
C, F, H, W = image.size()
all_image_size.append((F, H, W))
for cap_feat in all_cap_feats:
cap_ori_len = cap_feat.size(0)
cap_padding_len = cap_seq_len_target - cap_ori_len
cap_padded_feat = torch.cat(
[cap_feat, cap_feat[-1:].repeat(cap_padding_len, 1)],
dim=0,
)
all_cap_feats_out.append(cap_padded_feat)
all_cap_valid_lens.append(cap_ori_len)
F_tokens, H_tokens, W_tokens = F // pF, H // pH, W // pW
image = image.view(C, F_tokens, pF, H_tokens, pH, W_tokens, pW)
# "c f pf h ph w pw -> (f h w) (pf ph pw c)"
image = image.permute(1, 3, 5, 2, 4, 6, 0).reshape(
F_tokens * H_tokens * W_tokens, pF * pH * pW * C
)
image_ori_len = image.size(0)
min_image_seq_len = self._ceil_to_multiple(image_ori_len, SEQ_MULTI_OF)
if image_seq_len_target is None:
image_seq_len_target = min_image_seq_len
else:
image_seq_len_target = max(min_image_seq_len, image_seq_len_target)
image_padding_len = image_seq_len_target - image_ori_len
target_image_seq_len = image_seq_len_target or 0
for image in all_image:
# ------------ Process Image ------------
C, F, H, W = image.size()
image_size = (F, H, W)
# padded feature
image_padded_feat = torch.cat(
[image, image[-1:].repeat(image_padding_len, 1)],
dim=0,
)
all_image_out.append(image_padded_feat)
all_image_valid_lens.append(image_ori_len)
F_tokens, H_tokens, W_tokens = F // pF, H // pH, W // pW
image = image.view(C, F_tokens, pF, H_tokens, pH, W_tokens, pW)
# "c f pf h ph w pw -> (f h w) (pf ph pw c)"
image = image.permute(1, 3, 5, 2, 4, 6, 0).reshape(
F_tokens * H_tokens * W_tokens, pF * pH * pW * C
)
image_ori_len = image.size(0)
target_image_seq_len = max(
target_image_seq_len,
self._ceil_to_multiple(image_ori_len, SEQ_MULTI_OF),
)
image_records.append((image, image_size, image_ori_len))
for image, image_size, image_ori_len in image_records:
image_padding_len = target_image_seq_len - image_ori_len
image_padded_feat = torch.cat(
[image, image[-1:].repeat(image_padding_len, 1)],
dim=0,
)
all_image_out.append(image_padded_feat)
all_image_size.append(image_size)
all_image_valid_lens.append(image_ori_len)
return (
all_image_out,
all_cap_feats_out,
torch.stack(all_image_out, dim=0),
torch.stack(all_cap_feats_out, dim=0),
all_image_size,
all_image_valid_lens,
all_cap_valid_lens,
)
@staticmethod
def _as_image_list(hidden_states) -> list[torch.Tensor]:
"""Normalize 4D/5D image latents into per-sample tensors."""
if torch.is_tensor(hidden_states):
if hidden_states.dim() == 5:
return list(hidden_states.unbind(dim=0))
if hidden_states.dim() == 4:
return [hidden_states]
return list(hidden_states)
@staticmethod
def _as_caption_list(encoder_hidden_states) -> list[torch.Tensor]:
"""Normalize caption tensors into per-sample tensors."""
if torch.is_tensor(encoder_hidden_states):
if encoder_hidden_states.dim() == 3:
return list(encoder_hidden_states.unbind(dim=0))
if encoder_hidden_states.dim() == 2:
return [encoder_hidden_states]
cap_feats = list(encoder_hidden_states)
if len(cap_feats) == 1 and torch.is_tensor(cap_feats[0]):
if cap_feats[0].dim() == 3:
return list(cap_feats[0].unbind(dim=0))
if cap_feats[0].dim() == 2:
return cap_feats
return cap_feats
@staticmethod
def _replace_padding_with_token(
tensor: torch.Tensor,
valid_lens: list[int],
pad_token: torch.Tensor,
) -> torch.Tensor:
"""Replace padded token rows after each valid sequence length."""
positions = torch.arange(tensor.shape[1], device=tensor.device).unsqueeze(0)
lengths = torch.tensor(valid_lens, device=tensor.device).unsqueeze(1)
pad_mask = positions >= lengths
if pad_mask.any():
tensor = tensor.clone()
tensor[pad_mask] = pad_token.to(device=tensor.device, dtype=tensor.dtype)
return tensor
def forward(
self,
hidden_states: List[torch.Tensor],
@@ -854,14 +912,13 @@ class ZImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
assert patch_size in self.all_patch_size
assert f_patch_size in self.all_f_patch_size
x = hidden_states
cap_feats = encoder_hidden_states
x = self._as_image_list(hidden_states)
cap_feats = self._as_caption_list(encoder_hidden_states)
timestep = 1000.0 - timestep
t = timestep
bsz = 1
device = x[0].device
t = self.t_embedder(t)
adaln_input = t.type_as(x)
adaln_input = t.to(dtype=x[0].dtype)
(
x,
cap_feats,
@@ -876,29 +933,20 @@ class ZImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
image_seq_len_target=image_seq_len_target,
)
x = torch.cat(x, dim=0)
x, _ = self.all_x_embedder[f"{patch_size}-{f_patch_size}"](x)
if x_valid_lens[0] < x.shape[0]:
x[x_valid_lens[0] :] = self.x_pad_token.to(dtype=x.dtype)
x = self._replace_padding_with_token(x, x_valid_lens, self.x_pad_token)
x_freqs_cis = freqs_cis[1]
x = x.unsqueeze(0)
x_freqs_cis = x_freqs_cis
for layer_id, layer in enumerate(self.noise_refiner):
x = layer(x, x_freqs_cis, adaln_input)
cap_feats = torch.cat(cap_feats, dim=0)
cap_feats, _ = self.cap_embedder(cap_feats)
if cap_valid_lens[0] < cap_feats.shape[0]:
cap_feats[cap_valid_lens[0] :] = self.cap_pad_token.to(
dtype=cap_feats.dtype
)
cap_feats = self._replace_padding_with_token(
cap_feats, cap_valid_lens, self.cap_pad_token
)
cap_freqs_cis = freqs_cis[0]
cap_feats = cap_feats.unsqueeze(0)
cap_input_dtype = cap_feats.dtype
for layer_id, layer in enumerate(self.context_refiner):
cap_feats = layer(
cap_feats,
python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py
+30
View File
@@ -14,6 +14,7 @@ from __future__ import annotations
import logging
import os
import pprint
from collections import Counter
from copy import deepcopy
from dataclasses import MISSING, asdict, dataclass, field, fields
from typing import Any, Optional, Union
@@ -39,6 +40,23 @@ logger = init_logger(__name__)
SAMPLING_PARAMS_FIELDS = {f.name for f in fields(SamplingParams)}
@dataclass
class BatchMetricsWindow:
"""Counters accumulated between dynamic batching metric logs.
`total_capacity` uses each dispatch's effective admission cap, so
utilization reflects model/config limits instead of only the user max.
"""
dispatches: int = 0
total_requests: int = 0
total_capacity: int = 0
merged_dispatches: int = 0
full_dispatches: int = 0
wait_times_ms: list[float] = field(default_factory=list)
reject_reasons: Counter[str] = field(default_factory=Counter)
@dataclass(init=False)
class Req:
"""
@@ -71,6 +89,11 @@ class Req:
negative_prompt_embeds: list[torch.Tensor] | None = None
prompt_attention_mask: list[torch.Tensor | None] | None = None
negative_attention_mask: list[torch.Tensor | None] | None = None
# Masks and lengths aligned to postprocessed embeddings, one entry per text encoder.
prompt_embeds_mask: list[torch.Tensor | None] | None = None
negative_prompt_embeds_mask: list[torch.Tensor | None] | None = None
prompt_seq_lens: list[list[int]] | None = None
negative_prompt_seq_lens: list[list[int]] | None = None
clip_embedding_pos: list[torch.Tensor] | None = None
clip_embedding_neg: list[torch.Tensor] | None = None
@@ -269,6 +292,13 @@ class Req:
return None
return os.path.join(self.output_path, output_file_name)
@property
def resolution_key(self) -> str | None:
"""Return the batching config resolution key, e.g. "1024x1024"."""
if self.width is None or self.height is None:
return None
return f"{int(self.width)}x{int(self.height)}"
def set_as_warmup(self, warmup_steps: int = 1):
self.is_warmup = True
self.save_output = False
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py
+4 -1
View File
@@ -170,10 +170,13 @@ class DmdDenoisingStage(DenoisingStage):
**pos_cond_kwargs,
).permute(0, 2, 1, 3, 4)
video_timesteps = t_expand[:, None].expand(
-1, pred_noise.shape[1]
)
pred_video = pred_noise_to_pred_video(
pred_noise=pred_noise.flatten(0, 1),
noise_input_latent=noise_latents.flatten(0, 1),
timestep=t_expand,
timestep=video_timesteps,
scheduler=scheduler,
).unflatten(0, pred_noise.shape[:2])
python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py
+101 -6
View File
@@ -18,6 +18,7 @@ import torch
from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
from diffusers.models.modeling_outputs import AutoencoderKLOutput
from sglang.multimodal_gen.configs.pipeline_configs.base import TextConditioningOutput
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
from sglang.multimodal_gen.runtime.managers.component_manager import ComponentUse
from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
@@ -116,6 +117,10 @@ class ImageEncodingStage(PipelineStage):
"image_embeds",
"prompt_embeds",
"negative_prompt_embeds",
"prompt_embeds_mask",
"negative_prompt_embeds_mask",
"prompt_seq_lens",
"negative_prompt_seq_lens",
)
def __init__(
@@ -156,6 +161,30 @@ class ImageEncodingStage(PipelineStage):
return postprocess_funcs[0](outputs, image_inputs)
@staticmethod
def _split_text_conditioning_output(output):
if isinstance(output, TextConditioningOutput):
return (
output.prompt_embeds,
output.prompt_embeds_mask,
output.prompt_seq_lens,
)
return output, None, None
@staticmethod
def _full_text_seq_lens(prompt_embeds: torch.Tensor) -> list[int]:
if prompt_embeds.ndim == 2:
return [int(prompt_embeds.shape[0])]
return [int(prompt_embeds.shape[1])] * int(prompt_embeds.shape[0])
@staticmethod
def _default_text_mask(prompt_embeds: torch.Tensor) -> torch.Tensor:
if prompt_embeds.ndim == 2:
shape = (1, prompt_embeds.shape[0])
else:
shape = prompt_embeds.shape[:2]
return torch.ones(shape, dtype=torch.bool, device=prompt_embeds.device)
@torch.no_grad()
def forward(
self,
@@ -182,6 +211,10 @@ class ImageEncodingStage(PipelineStage):
all_prompt_embeds = []
all_neg_prompt_embeds = []
all_prompt_embeds_masks = []
all_neg_prompt_embeds_masks = []
all_prompt_seq_lens = []
all_neg_prompt_seq_lens = []
image_processor_call_params = inspect.signature(
self.image_processor.__call__
@@ -263,22 +296,84 @@ class ImageEncodingStage(PipelineStage):
output_hidden_states=True,
)
all_prompt_embeds.append(
self.encoding_image_edit(
outputs, image_inputs, server_args.pipeline_config
prompt_embeds, prompt_embeds_mask, prompt_seq_lens = (
self._split_text_conditioning_output(
self.encoding_image_edit(
outputs, image_inputs, server_args.pipeline_config
)
)
)
all_prompt_embeds.append(prompt_embeds)
all_prompt_embeds_masks.append(prompt_embeds_mask)
all_prompt_seq_lens.extend(
prompt_seq_lens
if prompt_seq_lens is not None
else self._full_text_seq_lens(prompt_embeds)
)
if batch.do_classifier_free_guidance:
all_neg_prompt_embeds.append(
self.encoding_image_edit(
neg_outputs, neg_image_inputs, server_args.pipeline_config
neg_prompt_embeds, neg_prompt_embeds_mask, neg_prompt_seq_lens = (
self._split_text_conditioning_output(
self.encoding_image_edit(
neg_outputs,
neg_image_inputs,
server_args.pipeline_config,
)
)
)
all_neg_prompt_embeds.append(neg_prompt_embeds)
all_neg_prompt_embeds_masks.append(neg_prompt_embeds_mask)
all_neg_prompt_seq_lens.extend(
neg_prompt_seq_lens
if neg_prompt_seq_lens is not None
else self._full_text_seq_lens(neg_prompt_embeds)
)
if all_prompt_embeds:
batch.prompt_embeds.append(torch.cat(all_prompt_embeds, dim=0))
if batch.prompt_embeds_mask is None:
batch.prompt_embeds_mask = []
batch.prompt_embeds_mask.append(
torch.cat(
[
(
mask
if mask is not None
else self._default_text_mask(prompt_embeds)
)
for prompt_embeds, mask in zip(
all_prompt_embeds, all_prompt_embeds_masks, strict=True
)
],
dim=0,
)
)
if batch.prompt_seq_lens is None:
batch.prompt_seq_lens = []
batch.prompt_seq_lens.append(all_prompt_seq_lens)
if all_neg_prompt_embeds:
batch.negative_prompt_embeds.append(torch.cat(all_neg_prompt_embeds, dim=0))
if batch.negative_prompt_embeds_mask is None:
batch.negative_prompt_embeds_mask = []
batch.negative_prompt_embeds_mask.append(
torch.cat(
[
(
mask
if mask is not None
else self._default_text_mask(neg_prompt_embeds)
)
for neg_prompt_embeds, mask in zip(
all_neg_prompt_embeds,
all_neg_prompt_embeds_masks,
strict=True,
)
],
dim=0,
)
)
if batch.negative_prompt_seq_lens is None:
batch.negative_prompt_seq_lens = []
batch.negative_prompt_seq_lens.append(all_neg_prompt_seq_lens)
return batch
python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py
+29 -3
View File
@@ -68,12 +68,32 @@ class InputValidationStage(PipelineStage):
return width, height
def _generate_seeds(self, batch: Req, server_args: ServerArgs):
"""Generate seeds for the inference"""
"""Generate deterministic per-output seeds.
Batched requests pass one base seed per prompt through `extra`; each
prompt expands to `num_outputs_per_prompt` consecutive seeds.
"""
seed = batch.seed
num_videos_per_prompt = batch.num_outputs_per_prompt
assert seed is not None
if isinstance(seed, list):
prompt_count = len(batch.prompt) if isinstance(batch.prompt, list) else 1
dynamic_batch_seeds = batch.extra.get("dynamic_batch_seeds")
if dynamic_batch_seeds is not None:
if (
not isinstance(dynamic_batch_seeds, list)
or len(dynamic_batch_seeds) != prompt_count
):
raise ValueError(
"dynamic_batch_seeds must be a list with one seed per prompt"
)
base_seeds = [int(item) for item in dynamic_batch_seeds]
seeds = []
for base_seed in base_seeds:
seeds.extend([base_seed + i for i in range(num_videos_per_prompt)])
elif isinstance(seed, list):
if len(seed) != num_videos_per_prompt:
raise ValueError(
f"seed list length must match num_outputs_per_prompt "
@@ -81,7 +101,13 @@ class InputValidationStage(PipelineStage):
)
seeds = [int(item) for item in seed]
else:
seeds = [int(seed) + i for i in range(num_videos_per_prompt)]
# Keep per-prompt seed streams deterministic and non-overlapping.
base_seeds = [
int(seed) + i * num_videos_per_prompt for i in range(prompt_count)
]
seeds = []
for base_seed in base_seeds:
seeds.extend([base_seed + i for i in range(num_videos_per_prompt)])
batch.seeds = seeds
# Create generators based on generator_device parameter
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/qwen_image_layered.py
+16 -11
View File
@@ -19,6 +19,15 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
def _seq_lens_from_optional_mask(
prompt_embeds: torch.Tensor, prompt_embeds_mask: torch.Tensor | None
) -> list[int]:
"""Return real text lengths, treating a missing mask as all tokens valid."""
if prompt_embeds_mask is None:
return [int(prompt_embeds.shape[1])] * int(prompt_embeds.shape[0])
return [int(x) for x in prompt_embeds_mask.sum(dim=1).tolist()]
# Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit_plus.calculate_dimensions
def calculate_dimensions(target_area, ratio):
width = math.sqrt(target_area * ratio)
@@ -521,22 +530,18 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
mu=mu,
)
txt_seq_lens = (
prompt_embeds_mask.sum(dim=1).tolist()
if prompt_embeds_mask is not None
else None
)
negative_txt_seq_lens = (
negative_prompt_embeds_mask.sum(dim=1).tolist()
if negative_prompt_embeds_mask is not None
else None
txt_seq_lens = _seq_lens_from_optional_mask(prompt_embeds, prompt_embeds_mask)
negative_txt_seq_lens = _seq_lens_from_optional_mask(
negative_prompt_embeds, negative_prompt_embeds_mask
)
is_rgb = torch.tensor([0]).to(device=device, dtype=torch.long)
batch.prompt_embeds = [prompt_embeds]
batch.prompt_attention_mask = [prompt_embeds_mask]
batch.prompt_embeds_mask = [prompt_embeds_mask]
batch.prompt_seq_lens = [txt_seq_lens]
batch.negative_prompt_embeds = [negative_prompt_embeds]
batch.negative_attention_mask = [negative_prompt_embeds_mask]
batch.negative_prompt_embeds_mask = [negative_prompt_embeds_mask]
batch.negative_prompt_seq_lens = [negative_txt_seq_lens]
batch.latents = latents
batch.image_latent = image_latents
batch.timesteps = timesteps
python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py
+155 -18
View File
@@ -14,6 +14,7 @@ from typing import Any
import torch
from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput
from sglang.multimodal_gen.configs.pipeline_configs.base import TextConditioningOutput
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
from sglang.multimodal_gen.runtime.managers.component_manager import ComponentUse
from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
@@ -53,6 +54,10 @@ class TextEncodingStage(PipelineStage):
"negative_prompt_embeds",
"prompt_attention_mask",
"negative_attention_mask",
"prompt_embeds_mask",
"negative_prompt_embeds_mask",
"prompt_seq_lens",
"negative_prompt_seq_lens",
"pooled_embeds",
"neg_pooled_embeds",
"clip_embedding_pos",
@@ -102,7 +107,13 @@ class TextEncodingStage(PipelineStage):
all_indices: list[int] = list(range(len(self.text_encoders)))
prompt_embeds_list, prompt_masks_list, pooler_embeds_list = self.encode_text(
(
prompt_embeds_list,
prompt_masks_list,
pooler_embeds_list,
prompt_embeds_masks_list,
prompt_seq_lens_list,
) = self.encode_text(
prompt_text,
server_args,
encoder_index=all_indices,
@@ -120,10 +131,23 @@ class TextEncodingStage(PipelineStage):
for am in prompt_masks_list:
batch.prompt_attention_mask.append(am)
batch.prompt_embeds_mask = []
batch.prompt_seq_lens = []
for mask in prompt_embeds_masks_list:
batch.prompt_embeds_mask.append(mask)
for seq_lens in prompt_seq_lens_list:
batch.prompt_seq_lens.append(seq_lens)
# Encode negative prompt if CFG is enabled
if batch.do_classifier_free_guidance:
assert isinstance(batch.negative_prompt, str)
neg_embeds_list, neg_masks_list, neg_pooler_embeds_list = self.encode_text(
(
neg_embeds_list,
neg_masks_list,
neg_pooler_embeds_list,
neg_embeds_masks_list,
neg_seq_lens_list,
) = self.encode_text(
batch.negative_prompt,
server_args,
encoder_index=all_indices,
@@ -132,16 +156,71 @@ class TextEncodingStage(PipelineStage):
assert batch.negative_prompt_embeds is not None
for ne in neg_embeds_list:
# A single negative prompt can be shared across positive prompts.
target_batch_sizes = [pe.shape[0] for pe in prompt_embeds_list]
def align_negative_batch_dim(
tensor: torch.Tensor, target_batch: int, name: str
) -> torch.Tensor:
if tensor.shape[0] == target_batch:
return tensor
if tensor.shape[0] == 1 and target_batch > 1:
return tensor.expand(target_batch, *tensor.shape[1:])
raise ValueError(
f"{name} batch dimension mismatch: got {tensor.shape[0]}, expected 1 or {target_batch}"
)
def align_negative_seq_lens(
seq_lens: list[int], target_batch: int, name: str
) -> list[int]:
if len(seq_lens) == target_batch:
return [int(x) for x in seq_lens]
if len(seq_lens) == 1 and target_batch > 1:
return [int(seq_lens[0])] * target_batch
raise ValueError(
f"{name} batch dimension mismatch: got {len(seq_lens)}, expected 1 or {target_batch}"
)
for idx, ne in enumerate(neg_embeds_list):
target_batch = target_batch_sizes[min(idx, len(target_batch_sizes) - 1)]
ne = align_negative_batch_dim(
ne, target_batch, "negative_prompt_embeds"
)
batch.negative_prompt_embeds.append(ne)
for pe in neg_pooler_embeds_list:
for idx, pe in enumerate(neg_pooler_embeds_list):
target_batch = target_batch_sizes[min(idx, len(target_batch_sizes) - 1)]
pe = align_negative_batch_dim(
pe, target_batch, "negative_pooled_embeds"
)
batch.neg_pooled_embeds.append(pe)
if batch.negative_attention_mask is None:
batch.negative_attention_mask = []
for nm in neg_masks_list:
for idx, nm in enumerate(neg_masks_list):
target_batch = target_batch_sizes[
min(idx, len(target_batch_sizes) - 1)
]
nm = align_negative_batch_dim(
nm, target_batch, "negative_attention_mask"
)
batch.negative_attention_mask.append(nm)
batch.negative_prompt_embeds_mask = []
batch.negative_prompt_seq_lens = []
for idx, nm in enumerate(neg_embeds_masks_list):
target_batch = target_batch_sizes[min(idx, len(target_batch_sizes) - 1)]
nm = align_negative_batch_dim(
nm, target_batch, "negative_prompt_embeds_mask"
)
batch.negative_prompt_embeds_mask.append(nm)
for idx, seq_lens in enumerate(neg_seq_lens_list):
target_batch = target_batch_sizes[min(idx, len(target_batch_sizes) - 1)]
batch.negative_prompt_seq_lens.append(
align_negative_seq_lens(
seq_lens, target_batch, "negative_prompt_seq_lens"
)
)
return batch
def build_dedup_fingerprint(
@@ -240,10 +319,14 @@ class TextEncodingStage(PipelineStage):
Returns:
Depending on return_type and return_attention_mask:
- list: List[Tensor] or (List[Tensor], List[Tensor])
- list: List[Tensor] or
(embeds, attention_masks, pooled_embeds, embeds_masks, seq_lens)
- dict: Dict[str, Tensor] or (Dict[str, Tensor], Dict[str, Tensor])
- stack: Tensor of shape [num_encoders, ...] or a tuple with stacked
attention masks
`embeds_masks` and `seq_lens` are aligned with postprocessed
embeddings for variable-length text conditioning.
"""
assert len(self.tokenizers) == len(self.text_encoders)
@@ -281,6 +364,8 @@ class TextEncodingStage(PipelineStage):
pooled_embeds_list: list[torch.Tensor] = []
attn_masks_list: list[torch.Tensor | None] = []
embeds_masks_list: list[torch.Tensor] = []
seq_lens_list: list[list[int]] = []
preprocess_funcs = server_args.pipeline_config.preprocess_text_funcs
postprocess_funcs = server_args.pipeline_config.postprocess_text_funcs
@@ -349,17 +434,34 @@ class TextEncodingStage(PipelineStage):
postprocess_kwargs["pipeline_config"] = server_args.pipeline_config
if "return_attention_mask" in postprocess_sig.parameters:
postprocess_kwargs["return_attention_mask"] = return_attention_mask
prompt_embeds = postprocess_func(outputs, text_inputs, **postprocess_kwargs)
has_postprocessed_attention_mask = False
postprocessed_attention_mask = None
if isinstance(prompt_embeds, tuple):
prompt_embeds, postprocessed_attention_mask = prompt_embeds
has_postprocessed_attention_mask = True
postprocess_result = postprocess_func(
outputs, text_inputs, **postprocess_kwargs
)
prompt_embeds_mask = None
prompt_seq_lens = None
if isinstance(postprocess_result, TextConditioningOutput):
prompt_embeds = postprocess_result.prompt_embeds
prompt_embeds_mask = postprocess_result.prompt_embeds_mask
prompt_seq_lens = postprocess_result.prompt_seq_lens
elif isinstance(postprocess_result, tuple):
if len(postprocess_result) != 2:
raise ValueError(
"Text postprocess tuple output must be (prompt_embeds, prompt_embeds_mask)"
)
prompt_embeds, prompt_embeds_mask = postprocess_result
else:
prompt_embeds = postprocess_result
if dtype is not None:
prompt_embeds = prompt_embeds.to(device=target_device, dtype=dtype)
else:
prompt_embeds = prompt_embeds.to(device=target_device)
if prompt_embeds_mask is not None:
prompt_embeds_mask = prompt_embeds_mask.to(
device=target_device, dtype=torch.bool
)
embeds_list.append(prompt_embeds)
pooled_output = server_args.pipeline_config.get_text_encoder_pooler_output(
@@ -369,11 +471,14 @@ class TextEncodingStage(PipelineStage):
pooled_embeds_list.append(pooled_output.to(device=target_device))
if return_attention_mask:
if has_postprocessed_attention_mask:
mask_to_store = (
postprocessed_attention_mask.to(device=target_device)
if postprocessed_attention_mask is not None
else None
if prompt_embeds_mask is not None:
mask_to_store = prompt_embeds_mask.to(
device=target_device,
dtype=(
attention_mask.dtype
if attention_mask is not None
else torch.long
),
)
elif attention_mask is not None and list(attention_mask.shape) == list(
prompt_embeds.shape[:2]
@@ -391,10 +496,42 @@ class TextEncodingStage(PipelineStage):
)
attn_masks_list.append(mask_to_store)
embeds_mask = prompt_embeds_mask
if embeds_mask is None:
embeds_mask = (
server_args.pipeline_config.build_text_conditioning_mask(
text_inputs,
attention_mask,
prompt_embeds,
i,
)
)
embeds_masks_list.append(embeds_mask)
if prompt_seq_lens is not None:
seq_lens_list.append([int(x) for x in prompt_seq_lens])
elif embeds_mask is not None:
seq_lens_list.append(
server_args.pipeline_config.seq_lens_from_text_conditioning_mask(
embeds_mask
)
)
elif prompt_embeds.ndim == 2:
seq_lens_list.append([int(prompt_embeds.shape[0])])
else:
seq_lens_list.append(
[int(prompt_embeds.shape[1])] * int(prompt_embeds.shape[0])
)
# Shape results according to return_type
if return_type == "list":
if return_attention_mask:
return embeds_list, attn_masks_list, pooled_embeds_list
return (
embeds_list,
attn_masks_list,
pooled_embeds_list,
embeds_masks_list,
seq_lens_list,
)
return embeds_list, pooled_embeds_list
if return_type == "dict":
python/sglang/multimodal_gen/runtime/server_args.py
+49
View File
@@ -225,6 +225,11 @@ class ServerArgs(DisaggArgsMixin):
webui_port: int | None = 12312
scheduler_port: int = 5555
batching_mode: str = "dynamic"
batching_max_size: int = 1
batching_delay_ms: float = 0.0
batching_config: str | None = None
enable_batching_metrics: bool = False
# Strict port mode: fail if requested port is unavailable instead of auto-selecting
strict_ports: bool = False
@@ -325,6 +330,7 @@ class ServerArgs(DisaggArgsMixin):
if not current_platform.is_cpu():
self._validate_parallelism()
self._validate_cfg_parallel()
self._validate_batching()
def _adjust_save_paths(self):
"""Normalize empty-string save paths to None (disabled)."""
@@ -1020,6 +1026,41 @@ class ServerArgs(DisaggArgsMixin):
default=ServerArgs.scheduler_port,
help="Port for the scheduler server.",
)
parser.add_argument(
"--batching-mode",
type=str,
default=ServerArgs.batching_mode,
choices=["dynamic"],
help="Request batching scheduler mode. Currently only 'dynamic' is implemented.",
)
parser.add_argument(
"--batching-max-size",
type=int,
default=ServerArgs.batching_max_size,
help="Maximum number of compatible generation requests to merge into one batch.",
)
parser.add_argument(
"--batching-delay-ms",
type=float,
default=ServerArgs.batching_delay_ms,
help="Maximum time (in ms) to wait for forming a larger batch before dispatch.",
)
parser.add_argument(
"--batching-config",
type=str,
default=ServerArgs.batching_config,
help=(
"Optional JSON file with {'schema_version': 1, 'rules': [...]} "
"batching admission rules that can cap model/resolution shapes "
"below --batching-max-size."
),
)
parser.add_argument(
"--enable-batching-metrics",
action="store_true",
default=ServerArgs.enable_batching_metrics,
help="Log periodic batch efficiency metrics such as realized batch size and queue wait time.",
)
parser.add_argument(
"--host",
type=str,
@@ -1461,6 +1502,14 @@ class ServerArgs(DisaggArgsMixin):
"CFG Parallelism is enabled via `--enable-cfg-parallel`, but num_gpus == 1"
)
def _validate_batching(self):
if self.batching_mode != "dynamic":
raise ValueError("batching_mode must be one of: dynamic")
if self.batching_max_size < 1:
raise ValueError("batching_max_size must be >= 1")
if self.batching_delay_ms < 0:
raise ValueError("batching_delay_ms must be >= 0")
def _set_default_attention_backend(self) -> None:
"""Configure ROCm defaults when users do not specify an attention backend."""
if current_platform.is_rocm():
python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py
+2 -2
View File
@@ -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)