diff --git a/docs/cookbook/diffusion/Qwen-Image/Qwen-Image-2.1.mdx b/docs/cookbook/diffusion/Qwen-Image/Qwen-Image-2.1.mdx index 4aeb5774a..18a324199 100644 --- a/docs/cookbook/diffusion/Qwen-Image/Qwen-Image-2.1.mdx +++ b/docs/cookbook/diffusion/Qwen-Image/Qwen-Image-2.1.mdx @@ -22,14 +22,14 @@ the hardware picker selects a tested single-GPU recipe for the full checkpoint. Use **Setup** to select text-to-image, single-image editing, or multi-image editing. **Server** controls placement, attention, encoder scheduling, VAE -tiling, and graph execution. **Request** controls the background, resolution, +tiling, graph execution, and request batching. **Request** controls the background, resolution, steps, and output count. Set reference PNG paths under **Variables**; edits upload files from the machine running cURL, so they need not exist on the server. Hardware selection applies the recommended placement for that GPU. H200, B200, and RTX PRO 6000 96GB keep weights resident; RTX 5090 and RTX 4090 use offload to fit the full pipeline. -Custom two- and four-GPU topologies and unverified feature combinations remain selectable and are labeled +Untested topologies and feature combinations remain selectable and are labeled **Unverified**. Invalid topology combinations disable Copy. This integration currently uses the Python/source command; no published Docker image is verified. @@ -47,7 +47,117 @@ for i, item in enumerate(json.loads(Path("response.json").read_text())["data"]): PY ``` -### Platform measurements +### Current native recipes + +The picker defaults use native BF16/FP32 precision, exact attention, eager +execution, and full-image VAE decoding. The table below uses checkpoint +`840b4adb1e2c21c7d77967203188b55b678c535f` and source `1f89064b419`, +measured on 2026-09-20. These are the best configurations among the tested +candidates for this workload, not a claim of a global optimum. + +| GPU | Single-output recipe | Generation median | Edit median | Request-phase peak | +| --- | --- | --- | --- | --- | +| H200 141GB | Resident / FlashAttention | 4.48 s | 5.29 s | 38.4 GiB | +| B200 192GB | Resident / FlashAttention | 2.46 s | 3.02 s | 38.5 GiB | +| RTX PRO 6000 96GB | Resident / Torch SDPA | 8.03 s | 9.63 s | 38.4 GiB | +| RTX 4090 24GB | DiT offload, 8 resident layers, encoder CPU offload / FlashAttention | 23.25 s | 24.63 s | 23.1 GiB | + +Each shape uses two 1024px/40-step warmups, then five generation and five edit +requests; the RTX 4090 8-layer candidate uses three measured requests per mode. +All use seed 42, CFG 1, CPU noise, and one RGBA PNG per request. HTTP time +includes encoding and PNG serialization, excluding startup. Memory is the peak +sampled every 0.2 seconds during these requests, excluding startup. The H200, B200, and +RTX PRO 6000 batch-matrix servers use a ceiling of four images and a 20 ms batching +window; the single-output 4090 candidate has dynamic batching off. +PyTorch is 2.13.0+cu130, Transformers 5.12.1, and Diffusers 0.37.0; native +conditioning explicitly preserves the reference's Transformers 4.57.3 semantics. + +FlashAttention beats SDPA on H200 (4.48 vs 4.91 s generation; 5.29 vs 6.23 s +editing) and B200 (2.46 vs 2.91 s; 3.02 vs 4.03 s). SDPA uses three measured +requests after two warmups. On RTX 4090, retaining eight DiT layers reduces +generation from 24.50 to 23.25 s and editing from 25.19 to 24.63 s. Its generation, +editing, transparent generation, and transparent-input editing pixels match +the fully streamed recipe exactly in this comparison. The tradeoff is a larger +request-phase peak: 19.9 to 23.1 GiB. + +On RTX PRO 6000, resident weights are faster than DiT offload: 8.03 vs 10.24 s +for generation and 9.63 vs 10.59 s for editing. Offload lowers the request-phase +peak from 38.4 to 25.3 GiB; its timings use two warmups and three measurements. + +The picker applies the 8-layer flag only to a native, eager, single-output +RTX 4090 FlashAttention recipe with request batching off. Multiple outputs or +request batching select the fully streamed recipe instead; use the updated +**Server** command when switching. Multi-reference and other untested shapes +remain marked Unverified. + +The RTX 5090 recipe retains its earlier validation below; it has not been rerun +with this checkpoint. Both RTX 5090 and RTX PRO 6000 map FA selection to SDPA, so those +labels do not represent two different attention kernels. + +A 1024px H200 BCG server captured its warmup graph, but these prompts and image +prefixes missed that signature and ran eagerly. Text buckets do not pad this +model's condition KV layout. Keep eager execution as the default; declaring a +resolution alone does not establish graph replay or a speedup. + +### Batching + +Keep **Request batching → Off** for interactive use. For a concurrent +text-to-image workload on one RTX 4090, the tested two-image server command is: + +```bash Command +sglang serve \ + --model-path /models/qwen-image-2.1 --model-id Qwen-Image-2.1 \ + --num-gpus 1 --ulysses-degree 1 --encoder-parallel auto \ + --performance-mode manual \ + --dit-layerwise-offload true --text-encoder-cpu-offload true \ + --attention-backend fa \ + --batching-max-size 2 --batching-delay-ms 20 \ + --host 0.0.0.0 --port 30010 +``` + +Send requests concurrently to use dynamic batching. The merge limit counts +output images, including each request's `n`; a client that waits for one +response before sending the next does not supply concurrent work. Image-edit +requests are not merged across requests. Set **Outputs** (the HTTP `n` field) +to produce multiple images within one generation or edit request. + +| GPU | Generation, 1 / 2 / 4 outputs | Editing, 1 / 2 / 4 outputs | Recommendation | +| --- | --- | --- | --- | +| H200 | 4.48 / 9.13 / 18.32 s | 5.29 / 10.60 / 21.13 s | Off; no throughput gain | +| B200 | 2.46 / 5.17 / 9.91 s | 3.02 / 5.98 / 12.04 s | Off; no material throughput gain | +| RTX PRO 6000 | 8.03 / 16.19 / 32.55 s | 9.63 / 19.08 / 38.11 s | Off; no material throughput gain | +| RTX 4090 | 23.25 / 38.33 / untested | 24.63 / 44.01 / untested | Off for latency; consider 2 for throughput | + +The H200/B200/RTX PRO 6000 and two-output 4090 measurements use five requests after two +warmups per shape. The one-output 4090 numbers use the faster 8-layer recipe +above with three measurements; batches stream all DiT layers. Against that +single-output recipe, two outputs improve 4090 generation throughput by about +**21%** and editing throughput by about **12%**, while increasing request latency. +Two concurrent 4090 generation requests complete in **39.75 s** wall time +(median of three rounds with different prompts and seeds). Concurrent 2/4-request +batches did not improve resident H200/B200/RTX PRO 6000 throughput. All measurements use +the checkpoint, software, 1024px/40-step settings, and HTTP timing scope above. +Transparent generation and transparent-input editing also passed for these +batch sizes, with alpha values from 0 to 255. + +Each sample retains its own condition-prefix KV cache, prompt, seed, and output +position. DiT target projections and MLPs run as a batch, and layerwise offload +transfers each block once per batch. This amortizes transfers on consumer GPUs. +Resident weights do not incur those transfers, so batching is not automatically +faster there. Full-checkpoint TP2 and Ulysses2 generation, editing, alpha, and +dynamic two-request batching also passed on two B200s; those are functional +checks, not latency recommendations for every multi-GPU topology. + +Batching retains native BF16/FP32 precision but changes GEMM shapes and +floating-point rounding. Batched images are **not guaranteed to match singleton +pixels**, even with the same seeds. Use one output and leave request batching +off when reproducing a singleton image exactly. Quantization, SageAttention, +and approximate denoising caches remain separate options. Larger batches and +other prompts need their own memory and throughput measurements. See +[Inference batching](/docs/sglang-diffusion/dynamic_batching) for admission rules +and batch metrics. + +### Earlier platform measurements The following four-platform comparison and the fusion measurements below precede the training-template and VAE normalization corrections in `c2a31b2693c`; @@ -89,10 +199,11 @@ this complete recipe passed generation, editing, and transparent PNG output. These are measurements of this small workload, not universal latency or image quality guarantees. Different prompts, reference sizes, batching, and software -versions can change memory use and latency. Multi-reference and batched request -recipes retain their separate H200 verification scope in the picker. +versions can change memory use and latency. Multi-reference recipes retain their +separate H200 verification scope. See [Batching](#batching) for the current +multi-output and concurrent-request matrix. -### RTX PRO 6000 Blackwell 96GB +### Earlier RTX PRO 6000 Blackwell 96GB measurements The recommended single-GPU command keeps all weights resident and selects Torch SDPA. This is the 96GB Blackwell Server Edition (SM120). This runtime also maps @@ -265,9 +376,9 @@ The model predicts continuous alpha values, including partly transparent edges. No thresholding or background-removal postprocessing is applied. Transparent generation and transparent-input editing were compared against the reference at 1024 × 1024 and 40 steps; that check does not guarantee perfect cutouts for -every prompt. Transparent generation and single-image editing also passed on -the recommended one-H200 and one-RTX PRO 6000 servers at that resolution and -step count, with one output per request. +every prompt. The updated checkpoint also passed transparent generation and +transparent-input editing on H200, B200, RTX PRO 6000, and RTX 4090. See +[Batching](#batching) for the tested output counts and two-B200 topologies. ## 4. Offline requests @@ -531,7 +642,10 @@ passed dynamic loading, merging, and removal on one B200 and TP2 with encoder folding. Both removal paths restored the base image exactly. This verifies adapter application and lifecycle, not the quality of a trained LoRA. -Cache-DiT hooks operate on target-image transformer blocks. Breakable CUDA +Cache-DiT hooks operate on target-image transformer blocks; two-output generation, +editing, transparent generation, and transparent-input editing passed with it +enabled on B200 at 1024px/40 steps. This is a functional check of an approximate +cache, not a lossless recipe. Breakable CUDA Graph execution fills each request's prefix caches eagerly, then replays matching warmup graphs with those cache tensors as inputs. Warmup and request condition-prefix lengths must match, in addition to the output resolution; diff --git a/docs/docs/sglang-diffusion/compatibility_matrix.mdx b/docs/docs/sglang-diffusion/compatibility_matrix.mdx index fa28af250..1aa611353 100644 --- a/docs/docs/sglang-diffusion/compatibility_matrix.mdx +++ b/docs/docs/sglang-diffusion/compatibility_matrix.mdx @@ -37,7 +37,12 @@ bit-exact. In the tested image, spatial versus full decode gave RGBA PSNR 60.69 dB; this is one sample, not a general quality guarantee. Request-local prefix KV reuse and new-request graph replay have CUDA regression -tests. The deployment picker marks only its exact tested HTTP combinations as +tests, including batched targets and independent variable-length prefixes. +Compatible text-to-image requests support opt-in dynamic batching; image-edit +requests remain separate, while multiple outputs within one request are supported. +See the cookbook's [batching guidance](/cookbook/diffusion/Qwen-Image/Qwen-Image-2.1#batching) +for measured throughput and floating-point reproducibility limits. +The deployment picker marks only its exact tested HTTP combinations as verified, including H200, B200, RTX PRO 6000 96GB, RTX 5090, and RTX 4090. CLI-only combinations remain Unverified in the picker. diff --git a/docs/docs/sglang-diffusion/dynamic_batching.mdx b/docs/docs/sglang-diffusion/dynamic_batching.mdx index cd1ac61a7..1085b3282 100644 --- a/docs/docs/sglang-diffusion/dynamic_batching.mdx +++ b/docs/docs/sglang-diffusion/dynamic_batching.mdx @@ -77,6 +77,7 @@ An initial implementation of dynamic batching for T2I and T2V models can be foun GLM-Image (external AR)✅- Qwen Image✅- Qwen Image 2512✅- + Qwen Image 2.1✅❌ Qwen Image Edit-❌ Qwen Image Edit 2509-? Qwen Image Edit 2511-? @@ -98,6 +99,11 @@ An initial implementation of dynamic batching for T2I and T2V models can be foun +Qwen Image 2.1 supports merging compatible text-to-image requests. Image edits +are not merged across requests; `n > 1` still produces multiple outputs within +one edit request. See its [cookbook](/cookbook/diffusion/Qwen-Image/Qwen-Image-2.1#batching) +for platform measurements and when batching is useful. + ### Video diff --git a/docs/src/snippets/_deployment.jsx b/docs/src/snippets/_deployment.jsx index 1f061ac8f..42c127978 100644 --- a/docs/src/snippets/_deployment.jsx +++ b/docs/src/snippets/_deployment.jsx @@ -2014,7 +2014,7 @@ export const Deployment = ({ config, benchmarks }) => { const options = visibleOptions(dim, sel); const currentOption = selectedOption(dim); return ( -
+
{direct ? dim.title : `${dim.title} options`} diff --git a/docs/src/snippets/configs/Qwen/qwen-image-2.1.jsx b/docs/src/snippets/configs/Qwen/qwen-image-2.1.jsx index d98ca6e71..4c34f7ade 100644 --- a/docs/src/snippets/configs/Qwen/qwen-image-2.1.jsx +++ b/docs/src/snippets/configs/Qwen/qwen-image-2.1.jsx @@ -48,16 +48,19 @@ const config = { recommendedWhen: (s) => ["h200", "b200", "rtxpro6000"].includes(s.hw), disabled: (s) => ["rtx5090", "rtx4090"].includes(s.hw) && Number(s.gpus_per_node) === 1, disableReason: "The full resident pipeline exceeds one consumer GPU's memory. Select CPU offload.", - flags: (s) => [Number(s.gpus_per_node) === 1 ? "--performance-mode speed" : "--performance-mode manual"], + flags: ["--performance-mode speed"], description: "Keep all components on the GPU. Recommended for H200, B200, and RTX PRO 6000 96GB. RTX 5090 and RTX 4090 need offload.", }, { id: "offload", label: "CPU offload", - flags: (s) => ["--performance-mode manual", "--dit-layerwise-offload true", ...(s.hw === "rtx4090" ? ["--text-encoder-cpu-offload true"] : [])], + flags: (s) => ["--performance-mode manual", "--dit-layerwise-offload true", ...(s.hw === "rtx4090" ? ["--text-encoder-cpu-offload true"] : []), + ...(s.hw === "rtx4090" && Number(s.gpus_per_node) === 1 && effectiveAttention(s) === "fa" && s.precision === "native" && s.execution === "eager" + && ["text", "edit"].includes(s.mode) && Number(s.outputs) === 1 && (!s.batching || s.batching === "off") + ? ["--dit-layerwise-resident-layers 8"] : [])], recommendedWhen: (s) => ["rtx5090", "rtx4090"].includes(s.hw), soft: (s) => !["rtxpro6000", "rtx5090", "rtx4090"].includes(s.hw) || Number(s.gpus_per_node) !== 1, softReason: "This offload topology has not completed an HTTP verification run.", - description: "Streams DiT layers. RTX 4090 also offloads the encoder between requests to leave room for image editing. Requires sufficient host RAM.", + description: "Streams DiT layers. RTX 4090 also offloads the encoder. Its native single-output FlashAttention recipe keeps 8 DiT layers resident; batch recipes stream every layer for memory headroom. Requires sufficient host RAM.", }, { id: "all_offload", label: "All components layerwise", @@ -207,9 +210,32 @@ const config = { { id: "eager", label: "Eager", recommended: true }, { id: "bcg", label: "Breakable CUDA Graph", - flags: ["--enable-breakable-cuda-graph true", "--warmup-resolutions 512x512", "--bcg-text-buckets 64"], - soft: true, softReason: "Only a matching 512px CLI warmup was verified. Other prompts or image prefixes can fall back to eager.", - description: "Captures a 512px warmup. Text buckets do not pad condition KV; this is not a guaranteed replay recipe.", + flags: (s) => ["--enable-breakable-cuda-graph true", `--warmup-resolutions ${s.resolution || "1024"}x${s.resolution || "1024"}`, "--bcg-text-buckets 64"], + soft: true, softReason: "A 1024px H200 server captured its warmup graph, but tested requests fell back to eager because condition-prefix shapes differed.", + description: "Captures the selected resolution. Condition-prefix shapes must also match warmup; text buckets alone do not ensure replay.", + }, + ], + }, + { + id: "batching", + title: "Request batching", + scope: "serve", + description: "Merge compatible text-to-image requests. Image edits run separately; Outputs controls multiple images within one request. Batching preserves native precision but can change floating-point rounding and output pixels.", + learnMore: "#batching", + default: "off", + options: [ + { id: "off", label: "Off", recommended: true, flags: ["--batching-max-size 1"], description: "Recommended for interactive latency. Resident H200, B200, and RTX PRO 6000 batching did not materially improve throughput in the measured workload." }, + { + id: "2", label: "Up to 2 images", + flags: ["--batching-max-size 2", "--batching-delay-ms 20"], + description: "Wait up to 20 ms to merge compatible queued requests. The tested RTX 4090 offload recipe benefits under concurrent load; each response takes longer.", + }, + { + id: "4", label: "Up to 4 images", + flags: ["--batching-max-size 4", "--batching-delay-ms 20"], + soft: (s) => ["rtx4090", "rtx5090"].includes(s.hw), + softReason: "Four-image batches have not been verified on this GPU and can exceed its memory.", + description: "Larger batches increase activation memory and individual request latency.", }, ], }, @@ -257,12 +283,14 @@ const config = { resource: { limits: { nodes: { min: 1, max: 1 }, gpus_per_node: { min: 1, max: 4 } }, verifiedRecipes: [ - { id: "h200-1-resident", hw: "h200", nodes: 1, gpus_per_node: 1, placement: "resident", tp_size: 1, ulysses_degree: 1, ring_degree: 1, encoder: "auto", attentions: ["fa"], default: true }, - { id: "b200-1-resident", hw: "b200", nodes: 1, gpus_per_node: 1, placement: "resident", tp_size: 1, ulysses_degree: 1, ring_degree: 1, encoder: "auto", attentions: ["fa", "sdpa"], default: true }, - { id: "rtxpro6000-1-resident", hw: "rtxpro6000", nodes: 1, gpus_per_node: 1, placement: "resident", tp_size: 1, ulysses_degree: 1, ring_degree: 1, encoder: "auto", attentions: ["sdpa"], default: true }, + { id: "h200-1-resident", hw: "h200", nodes: 1, gpus_per_node: 1, placement: "resident", tp_size: 1, ulysses_degree: 1, ring_degree: 1, encoder: "auto", attentions: ["fa", "sdpa"], batchSizes: [1, 2, 4], batchAttentions: ["fa"], default: true }, + { id: "b200-1-resident", hw: "b200", nodes: 1, gpus_per_node: 1, placement: "resident", tp_size: 1, ulysses_degree: 1, ring_degree: 1, encoder: "auto", attentions: ["fa", "sdpa"], batchSizes: [1, 2, 4], batchAttentions: ["fa"], default: true }, + { id: "b200-2-tp", hw: "b200", nodes: 1, gpus_per_node: 2, placement: "resident", tp_size: 2, ulysses_degree: 1, ring_degree: 1, encoder: "auto", attentions: ["fa"], batchSizes: [1, 2] }, + { id: "b200-2-ulysses", hw: "b200", nodes: 1, gpus_per_node: 2, placement: "resident", tp_size: 1, ulysses_degree: 2, ring_degree: 1, encoder: "auto", attentions: ["fa"], batchSizes: [1, 2] }, + { id: "rtxpro6000-1-resident", hw: "rtxpro6000", nodes: 1, gpus_per_node: 1, placement: "resident", tp_size: 1, ulysses_degree: 1, ring_degree: 1, encoder: "auto", attentions: ["sdpa"], batchSizes: [1, 2, 4], default: true }, { id: "rtxpro6000-1-offload", hw: "rtxpro6000", nodes: 1, gpus_per_node: 1, placement: "offload", tp_size: 1, ulysses_degree: 1, ring_degree: 1, encoder: "auto", attentions: ["sdpa"] }, { id: "rtx5090-1-offload", hw: "rtx5090", nodes: 1, gpus_per_node: 1, placement: "offload", tp_size: 1, ulysses_degree: 1, ring_degree: 1, encoder: "auto", attentions: ["sdpa"], default: true }, - { id: "rtx4090-1-offload", hw: "rtx4090", nodes: 1, gpus_per_node: 1, placement: "offload", tp_size: 1, ulysses_degree: 1, ring_degree: 1, encoder: "auto", attentions: ["fa"], default: true }, + { id: "rtx4090-1-offload", hw: "rtx4090", nodes: 1, gpus_per_node: 1, placement: "offload", tp_size: 1, ulysses_degree: 1, ring_degree: 1, encoder: "auto", attentions: ["fa", "sdpa"], batchSizes: [1, 2], batchAttentions: ["fa"], default: true }, ], autoTopology: (s) => ({ tp_size: 1, ulysses_degree: Number(s.gpus_per_node), ring_degree: 1 }), validateTopology: (s, topology) => { @@ -293,10 +321,14 @@ const config = { && entry.ulysses_degree === topology.ulysses_degree && entry.ring_degree === topology.ring_degree); const serveVerified = !!recipe && errors.length === 0 && s.encoder === "auto" && recipe.attentions.includes(effectiveAttention(s)) && s.precision === "native" - && s.execution === "eager" && s.vae === "full"; + && s.execution === "eager" && s.vae === "full" + && (!s.batching || s.batching === "off" || ((recipe.batchSizes || [1]).includes(Number(s.batching)) + && (recipe.batchAttentions || recipe.attentions).includes(effectiveAttention(s)))); // Exact HTTP workloads from the validation matrix, not blanket quality coverage. const requestVerified = serveVerified - && ((["text", "edit"].includes(s.mode) && s.resolution === "1024" && Number(s.steps) === 40 && Number(s.outputs) === 1 + && ((recipe.batchSizes && (Number(s.outputs) === 1 || (recipe.batchAttentions || recipe.attentions).includes(effectiveAttention(s))) + && ["text", "edit"].includes(s.mode) && s.resolution === "1024" && Number(s.steps) === 40 && recipe.batchSizes.includes(Number(s.outputs))) + || (["text", "edit"].includes(s.mode) && s.resolution === "1024" && Number(s.steps) === 40 && Number(s.outputs) === 1 && (["h200", "rtxpro6000"].includes(s.hw) || s.mode === "text" || s.background === "scene")) || (s.hw === "h200" && s.background === "scene" && s.mode === "text" && s.resolution === "512" && Number(s.steps) === 4 && Number(s.outputs) === 2) || (s.hw === "h200" && s.background === "scene" && s.mode === "multi" && s.resolution === "512" && Number(s.steps) === 4 && Number(s.outputs) === 1)); @@ -307,6 +339,8 @@ const config = { if (topology.ring_degree > 1) flags.push(`--ring-degree ${topology.ring_degree}`); flags.push("--host {{HOST_IP}}", "--port {{PORT}}"); const warnings = []; + if (s.hw === "rtx4090" && (Number(s.outputs) > 1 || (s.batching && s.batching !== "off"))) warnings.push("This recipe streams all DiT layers. Use the updated Server command if switching from the single-output recipe with 8 resident layers."); + if (s.batching && s.batching !== "off" && s.mode !== "text") warnings.push("Cross-request batching applies to text-to-image requests. Image edits run separately; use Outputs for multiple images in one edit request."); if (!serveVerified && !errors.length) warnings.push("This server combination has not completed an exact HTTP verification run."); if (!requestVerified && !errors.length) warnings.push("This request shape is outside the verified HTTP matrix."); return { diff --git a/python/sglang/multimodal_gen/README.md b/python/sglang/multimodal_gen/README.md index 5efb31b46..868edc17a 100644 --- a/python/sglang/multimodal_gen/README.md +++ b/python/sglang/multimodal_gen/README.md @@ -93,7 +93,11 @@ sglang generate --model-path /models/qwen-image-2.1 --model-id Qwen-Image-2.1 \ Add `--image-path /path/to/input.png` for editing. Dimensions must be multiples of 32. Full-checkpoint generation and editing have been tested on H200; see the [model cookbook](../../../docs/cookbook/diffusion/Qwen-Image/Qwen-Image-2.1.mdx) -for component requirements and optimization boundaries. +for component requirements and optimization boundaries. Compatible text-to-image +requests support opt-in dynamic batching with `--batching-max-size 2` when +serving. Image edits are not merged across requests; use `n` for multiple +outputs within an edit request. Batching can improve offload throughput, but +changes floating-point rounding and is not always faster with resident weights. ### Component residency diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image21.py b/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image21.py index 9e053eca9..b06ede3b4 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image21.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image21.py @@ -27,6 +27,10 @@ class QwenImage21PipelineConfig(ImagePipelineConfig): text_encoder_configs: tuple = field(default_factory=lambda: (Qwen3VLConfig(),)) text_encoder_precisions: tuple[str, ...] = ("bf16",) + def supports_dynamic_batching(self): + # the scheduler excludes reference-image requests from cross-request merging + return True + def prepare_sigmas(self, sigmas, num_inference_steps): return self._prepare_sigmas(sigmas, num_inference_steps) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image21.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image21.py index 22735163b..297bcee9f 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image21.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image21.py @@ -102,6 +102,7 @@ def build_layout(image_slots, image_shapes, axes_dims, device): ) rope = torch.polar(torch.ones_like(angles), angles) return dict( + encoder_seq_len=len(image_slots), text_indices=torch.tensor(indices, device=device, dtype=torch.long), image_indices=torch.tensor(image_indices, device=device, dtype=torch.long), prefix_rope=rope[:prefix_len], @@ -307,7 +308,7 @@ class QwenImage21Attention(nn.Module): v, ) - def forward(self, x, rope, prefix, prefix_rope, segments, cache): + def attend_sample(self, q, k, v, rope, prefix, prefix_rope, segments, cache): if cache: kp, vp = cache["key"], cache["value"] prefix_output = None @@ -319,8 +320,8 @@ class QwenImage21Attention(nn.Module): mask = None if not is_image: mask = ( - torch.arange(end, device=x.device)[None, :] - <= torch.arange(start, end, device=x.device)[:, None] + torch.arange(end, device=q.device)[None, :] + <= torch.arange(start, end, device=q.device)[:, None] ) mask = mask[None, None] outputs.append( @@ -331,7 +332,6 @@ class QwenImage21Attention(nn.Module): prefix_output = self.to_out[0](torch.cat(outputs, dim=1).flatten(2))[0] if cache is not None: cache.update(key=kp, value=vp) - q, k, v = self.project_qkv(x) q = apply_qk_norm_rope(q, self.norm_q, rope) packed = None if ( @@ -358,7 +358,26 @@ class QwenImage21Attention(nn.Module): else: k = apply_qk_norm_rope(k, self.norm_k, rope) out = self.target_attn.forward_with_replicated_kv_prefix(q, kp, vp, k, v) - return self.to_out[0](out.flatten(2))[0], prefix_output + return out, prefix_output + + def forward(self, x, ropes, prefixes, layouts, caches): + # batch target projections while retaining each sample's unpadded prefix + q, k, v = self.project_qkv(x) + outputs, prefix_outputs = [], [] + for sample, layout in enumerate(layouts): + out, prefix_out = self.attend_sample( + q[sample : sample + 1], + k[sample : sample + 1], + v[sample : sample + 1], + ropes[sample], + prefixes[sample], + layout["prefix_rope"], + layout["segments"], + caches[sample], + ) + outputs.append(out) + prefix_outputs.append(prefix_out) + return self.to_out[0](torch.cat(outputs).flatten(2))[0], prefix_outputs class QwenImage21TransformerBlock(nn.Module): @@ -379,25 +398,27 @@ class QwenImage21TransformerBlock(nn.Module): self, hidden_states, modulation, - prefix_state, + prefix_states, prefix_modulation, - layout, - rope, - cache, + layouts, + ropes, + caches, ): - prefix = prefix_state.get("hidden_states") scale1, gate1, scale2, gate2 = modulation - p = None - if not cache: - ps1, pg1, ps2, pg2 = prefix_modulation - p = apply_modulation(prefix, self.img_norm1, ps1) - attention, prefix_attention = self.attn( + prefixes = [ + apply_modulation( + state["hidden_states"], self.img_norm1, prefix_modulation[0] + ) + if not cache + else None + for state, cache in zip(prefix_states, caches, strict=True) + ] + attention, prefix_attentions = self.attn( apply_modulation(hidden_states, self.img_norm1, scale1), - rope, - p, - layout["prefix_rope"], - layout["segments"], - cache, + ropes, + prefixes, + layouts, + caches, ) hidden_states = residual_gate_add(hidden_states, attention, gate1) hidden_states = residual_gate_add( @@ -405,14 +426,15 @@ class QwenImage21TransformerBlock(nn.Module): self.img_mlp(apply_modulation(hidden_states, self.img_norm2, scale2)), gate2, ) - if prefix_attention is not None: - prefix = residual_gate_add(prefix, prefix_attention, pg1) - prefix = residual_gate_add( - prefix, - self.img_mlp(apply_modulation(prefix, self.img_norm2, ps2)), - pg2, - ) - prefix_state["hidden_states"] = prefix + for state, attention in zip(prefix_states, prefix_attentions, strict=True): + if attention is not None: + _, pg1, ps2, pg2 = prefix_modulation + prefix = residual_gate_add(state["hidden_states"], attention, pg1) + state["hidden_states"] = residual_gate_add( + prefix, + self.img_mlp(apply_modulation(prefix, self.img_norm2, ps2)), + pg2, + ) return hidden_states @@ -505,39 +527,35 @@ class QwenImage21Transformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin timestep.new_zeros(1).to(images.dtype), images.dtype ) prefix_modulation = self.prepare_modulation(zero_temb) - outputs = [] + if prefix_caches is None: + prefix_caches = [[None] * len(self.transformer_blocks) for _ in layouts] + prefix_states, ropes = [], [] for sample, layout in enumerate(layouts): - caches = ( - prefix_caches[sample] - if prefix_caches is not None - else [None] * len(self.transformer_blocks) - ) prefix = None - if not caches[0]: + if not prefix_caches[sample][0]: prefix = self.txt_in( - encoder_hidden_states[sample : sample + 1] + encoder_hidden_states[ + sample : sample + 1, : layout["encoder_seq_len"] + ] ).index_select(1, layout["text_indices"]) if condition_latents is not None: prefix[:, layout["image_indices"]] = self.img_in( condition_latents[sample : sample + 1] ) - prefix_state = {"hidden_states": prefix} - x = images[sample : sample + 1] - sample_modulation = tuple( - value[sample : sample + 1] for value in modulation + prefix_states.append({"hidden_states": prefix}) + ropes.append(layout["target_rope"][start:end]) + # visit each block once so layerwise offload transfers weights once per batch + for i, block in enumerate(self.transformer_blocks): + images = block( + images, + modulation, + prefix_states, + prefix_modulation, + layouts, + ropes, + [cache[i] for cache in prefix_caches], ) - for i, block in enumerate(self.transformer_blocks): - x = block( - x, - sample_modulation, - prefix_state, - prefix_modulation, - layout, - layout["target_rope"][start:end], - caches[i], - ) - outputs.append(self.proj_out(self.norm_out(x, temb[sample : sample + 1]))) - output = torch.cat(outputs) + output = self.proj_out(self.norm_out(images, temb)) if sp > 1: output = sequence_model_parallel_all_gather(output, dim=1) return output diff --git a/python/sglang/multimodal_gen/test/unit/test_qwen_image21.py b/python/sglang/multimodal_gen/test/unit/test_qwen_image21.py index 8b4ec9625..d38737612 100644 --- a/python/sglang/multimodal_gen/test/unit/test_qwen_image21.py +++ b/python/sglang/multimodal_gen/test/unit/test_qwen_image21.py @@ -20,6 +20,7 @@ from sglang.multimodal_gen.configs.models.vaes.qwenimage21 import ( from sglang.multimodal_gen.configs.pipeline_configs.qwen_image21 import ( QwenImage21PipelineConfig, ) +from sglang.multimodal_gen.configs.sample.qwenimage21 import QwenImage21SamplingParams from sglang.multimodal_gen.registry import _get_config_info from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import ( ResidencyState, @@ -27,6 +28,7 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager im from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency_strategies import ( ComponentOffloadStrategy, ) +from sglang.multimodal_gen.runtime.managers.scheduler import Scheduler from sglang.multimodal_gen.runtime.models.dits.qwen_image21 import build_layout from sglang.multimodal_gen.runtime.models.encoders.qwen3vl_vision import ( Qwen3VLVisionRotaryEmbedding, @@ -37,6 +39,7 @@ from sglang.multimodal_gen.runtime.models.vaes.autoencoder_kl_qwenimage21 import _patchify, _unpatchify, ) +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req from sglang.multimodal_gen.runtime.pipelines_core.stages.input_validation import ( InputValidationStage, ) @@ -184,6 +187,31 @@ def test_latent_pack_decode_contract(): ) +@pytest.mark.parametrize("outputs", [1, 2]) +def test_dynamic_batching_preserves_output_order_and_seeds(outputs): + scheduler = object.__new__(Scheduler) + config = QwenImage21PipelineConfig() + scheduler.server_args = SimpleNamespace(pipeline_config=config) + scheduler._batch_admission = SimpleNamespace(enabled=True) + assert scheduler._dynamic_batching_enabled() + requests = [] + for i, prompt in enumerate(["short", "a longer prompt"]): + params = QwenImage21SamplingParams( + prompt=prompt, seed=7 + i * 10, num_outputs_per_prompt=outputs + ) + requests.append(Req(request_id=f"request-{i}", sampling_params=params)) + merged = scheduler._try_merge_generation_reqs(requests) + assert merged.prompt == ["short", "a longer prompt"] + assert merged.extra["dynamic_batch_seeds"] == [7, 17] + result = torch.arange(outputs * 2).reshape(-1, 1) + split = scheduler._split_batched_output(OutputBatch(output=result), requests) + assert len(split) == 2 + torch.testing.assert_close(split[0].output, result[:outputs]) + torch.testing.assert_close(split[1].output, result[outputs:]) + requests[1].image_path = "reference.png" + assert scheduler._try_merge_generation_reqs(requests) is None + + @pytest.mark.parametrize("channels", [3, 4]) @pytest.mark.parametrize("tiling", [False, True]) def test_native_vae_roundtrip_shapes_and_checkpoint_names(channels, tiling): diff --git a/python/sglang/multimodal_gen/test/unit/test_qwen_image21_cuda.py b/python/sglang/multimodal_gen/test/unit/test_qwen_image21_cuda.py index 66e0114d7..4b6e23476 100644 --- a/python/sglang/multimodal_gen/test/unit/test_qwen_image21_cuda.py +++ b/python/sglang/multimodal_gen/test/unit/test_qwen_image21_cuda.py @@ -101,6 +101,81 @@ def inputs(seed, edit): ) +def batched_inputs(samples): + max_length = max(sample["encoder_hidden_states"].shape[1] for sample in samples) + return dict( + hidden_states=torch.cat([sample["hidden_states"] for sample in samples]), + encoder_hidden_states=torch.cat( + [ + torch.nn.functional.pad( + sample["encoder_hidden_states"], + (0, 0, 0, max_length - sample["encoder_hidden_states"].shape[1]), + ) + for sample in samples + ] + ), + condition_latents=( + torch.cat([sample["condition_latents"] for sample in samples]) + if samples[0]["condition_latents"] is not None + else None + ), + layouts=[sample["layouts"][0] for sample in samples], + prefix_caches=[sample["prefix_caches"][0] for sample in samples], + timestep=torch.cat([sample["timestep"] for sample in samples]), + ) + + +@pytest.mark.parametrize("edit", [False, True]) +@torch.no_grad() +def test_batched_targets_preserve_ragged_prefixes_and_cache_ownership(model, edit): + samples = [inputs(seed, edit) for seed in (5, 9)] + slots = [False] * 5 + ([True, False, False] if edit else []) + samples[1]["encoder_hidden_states"] = torch.randn(1, len(slots), 16, device="cuda") + samples[1]["layouts"] = [ + build_layout( + slots, ([(1, 2, 4)] if edit else []) + [(1, 4, 4)], (8, 12, 12), "cuda" + ) + ] + batch = batched_inputs(deepcopy(samples)) + for timestep in (700, 300, 10): + for sample in samples: + sample["timestep"].fill_(timestep) + batch["timestep"].fill_(timestep) + with set_forward_context(None, None): + expected = torch.cat([model(**sample) for sample in samples]) + calls = [] + handles = [ + block.register_forward_pre_hook( + lambda module, args: calls.append(args[0].shape) + ) + for block in model.transformer_blocks + ] + try: + with set_forward_context(None, None): + actual = model(**batch) + finally: + for handle in handles: + handle.remove() + assert calls == [torch.Size([2, 16, model.hidden_size])] * len( + model.transformer_blocks + ) + torch.testing.assert_close(actual, expected, atol=2e-6, rtol=1e-5) + for sample, caches in zip(samples, batch["prefix_caches"], strict=True): + for cache, reference in zip( + caches, sample["prefix_caches"][0], strict=True + ): + torch.testing.assert_close( + cache["key"], reference["key"], atol=0, rtol=0 + ) + torch.testing.assert_close( + cache["value"], reference["value"], atol=0, rtol=0 + ) + assert ( + batch["prefix_caches"][0][0]["key"].data_ptr() + != batch["prefix_caches"][1][0]["key"].data_ptr() + ) + + def test_bf16_qk_norm_matches_reference(model): norm = deepcopy(model.transformer_blocks[0].attn.norm_q).bfloat16() reference = ReferenceRMSNorm(32, eps=1e-6).cuda().bfloat16() @@ -182,8 +257,10 @@ def test_cached_prefix_matches_full_recomputation(model, edit): @pytest.mark.parametrize("edit", [False, True]) -def test_graph_replay_uses_new_request_prefix(model, edit): - first, second = inputs(5, edit), inputs(9, edit) +@pytest.mark.parametrize("sample_count", [1, 2]) +def test_graph_replay_uses_new_request_prefix(model, edit, sample_count): + first = batched_inputs([inputs(5 + i, edit) for i in range(sample_count)]) + second = batched_inputs([inputs(9 + i, edit) for i in range(sample_count)]) runner = DiffusionBreakableCudaGraphRunner(model, torch.device("cuda")) try: with torch.no_grad(), set_forward_context(None, None):