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

[diffusion] chore: batch qwen-image 2.1 targets and document measured deployment recipes (#40408)

Browse Source 
Co-authored-by: Mick Qian <mickqian@users.noreply.github.com>
This commit is contained in:
Mick
2026-09-20 13:59:36 +08:00
committed by GitHub
co-authored by Mick Qian
parent 99d53fe0c2
commit 031bff5dd3
10 changed files with 369 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/README.md
+5 -1
View File
@@ -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
python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image21.py
+4
View File
@@ -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)
python/sglang/multimodal_gen/runtime/models/dits/qwen_image21.py
+70 -52
View File
@@ -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
python/sglang/multimodal_gen/test/unit/test_qwen_image21.py
+28
View File
@@ -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):
python/sglang/multimodal_gen/test/unit/test_qwen_image21_cuda.py
+79 -2
View File
@@ -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):