[Perf] diffusion: Remove per-step host syncs and uploads in image models by mickqian · Pull Request #46046 · sgl-project/sglang

mickqian · 2026-10-23T11:56:53Z

Depends on #46045 (host-sync guards in test/unit/host_sync_guard.py); this branch is stacked on it.

Motivation

Several image pipelines read device values on the host or upload constants with pageable copies inside the denoising loop, or redo request-constant work in every forward. Each such read blocks the host until the stream drains, so it cannot queue the next step while the GPU runs this one. This PR removes them model by model. All values are unchanged.

Modifications

Qwen-Image (dits/qwen_image.py, pipeline_configs/qwen_image.py, qwen_image21.py, vaes/base.py, qwen_image_layered.py):
- Layered built additional_t_cond with a pageable torch.tensor([0]) per CFG branch. The branches got different objects, so the DiT's modulation cache never hit on the negative branch and every step recomputed 60 blocks x 2 modulation GEMMs. Both branches now share one device tensor.
- Batched prompts of different lengths rebuilt their varlen metadata every forward, uploading the lengths from pageable memory (a sync per step and branch). The metadata is cached by its lengths; each forward gets a fresh dict, since attention may memoize into it.
- Layered's condition-image encode, the channel-wise VAE stats helper (Qwen-Image edit encode/decode, Krea-2, JoyImage) and Qwen-Image-2.1 decode uploaded latent mean/std per request; they are uploaded once per (device, dtype).
FLUX.3 Action (flux3_action.py): cosmos_unipc called float(tick) on each device tick, a sync per step. The ticks are copied to the host once per request; each step still gets the same device tick.
Ideogram-4 (dits/ideogram.py, model_specific_stages/ideogram.py, progressive_resolution/ideogram.py): every forward cast the (B, T+I, 53248) fp32 LLM features to bf16, masked them and ran the RMSNorm and the 53248 -> 4608 projection (plus its TP all-gather) on inputs that never change within a request. The stage projects each branch's conditioning on its first step, inside that step's forward context, and passes it to later steps as llm_cond. BCG replays still pass the raw features, which their prompt padder pads; the progressive stage drops the projection when the sequence grows. Decode uploads its latent scale/shift once per (device, dtype).
GLM-Image (dits/glm_image.py, pipeline_configs/glm_image.py): the eager forward dropped the unconditional branch's prior tokens with prior_embedding[prior_token_drop] *= 0.0, two nonzero syncs per forward. It now selects prior_embedding * 0.0 with torch.where (same values, signed zeros included). Decode stats are uploaded once.
Krea-2 (pipeline_configs/krea2.py): without SP the DiT always got a key mask and ran bool(mask.all()) every forward. The config decides once per request and passes mask=None for unpadded prompts, the maskless path the DiT took before.
HunyuanImage-3 (hunyuan_image3_utils.py, hunyuan_image3/ar_stage.py): per step, a pageable timestep upload, host-built timestep_embedding frequencies uploaded on every call (five calls per CFG step), masked_select on the request-constant image mask for the scatter and the gather, and per-step condition-timestep uploads on the edit path. Positions, device timesteps and condition timesteps are taken once per request, predictions are gathered by position, and the frequencies are uploaded once per device.
SenseNova-U1 (modeling_neo_chat.py, modeling_fm_modules.py): the CFG-window check compared the device timestep in Python (two syncs per step); the window is computed for the whole schedule with the same fp32 comparisons and read once. The timestep embedder's frequencies are uploaded once per device.
SANA (dits/sana.py): every block cast the int64 prompt mask to bool; the forward casts it once.

Tests:

test_qwen_image_rope.py: batched prompts of different lengths build their varlen metadata once, and the second forward matches the first bitwise.
test_qwen_image_layered.py: both CFG branches share additional_t_cond; VAE and condition-latent stats are uploaded once.
test_qwen_image21.py, test_glm_image_dit.py, test_ideogram4.py: decode stats/constants are uploaded once and reused bitwise.
test_flux3_action.py: cosmos_unipc with SyncOnRead ticks matches the plain run bitwise.
test_ideogram4.py: the LLM projection runs once per branch and every step matches the unprojected forward bitwise; BCG replays get the raw features. progressive_resolution/test_progressive.py: a resolution change drops the projection.
test_glm_image_dit.py: the prior-token drop runs under forbid_host_syncs and equals the graph path's masked_fill.
test_krea2_pipeline_config.py: only ragged prompts carry a key mask.
test_hunyuan_image3_denoise_step.py: timestep frequencies are built once; the image-token scatter/gather runs under forbid_host_syncs through the precomputed positions.
test_sensenova_u1.py: timestep frequencies are built once per device.
test_sana_dit.py: the prompt mask is cast once per forward, output unchanged.

The model tests run their tiny DiTs on CUDA when the runner has one (the diffusion unit suite runs on an H100).

No hunk overlaps #45942, #45943, #45944, #45948, #45965, #45966, #45967, #45977, #45978, #45979 or #45987; each of them still applies on top of this branch. #45967 also uploads Qwen-Image-2.1 latent stats once, in the encoding stage; this PR covers decode, a different function.

Accuracy Tests

2xH200 (1 GPU for Qwen-Image-2.1), eager profiles, two requests per model with different inputs. Decoded pixels are bitwise identical to main:

model (request)	GPUs	outputs compared	main vs PRs
Qwen-Image-2.1 T2I 1024x1024, 40 steps	1	pixels	identical (2/2)
Qwen-Image-2.1 edit 1024x1024, 40 steps	1	pixels	identical (2/2)
Z-Image-Turbo 1024x1024, 9 steps, TP2	2	pixels	identical (2/2)
FLUX.3 Klein Base 1024x1024, 50 steps, CFG parallel	2	pixels	identical (2/2)
Ideogram-4 FP8 1024x1024, 20 steps, TP2	2	pixels	identical (2/2)
Cosmos3-Nano T2I 1280x720, 35 steps, CFG parallel	2	pixels	identical (2/2)
Cosmos3-Nano T2V 1280x720x61, 35 steps, CFG parallel	2	frames	identical (2/2)
HunyuanVideo 848x480x61, 30 steps, Ulysses 2	2	frames	identical (2/2)
LTX-2 two-stage 768x512x121, 40 steps, CFG parallel	2	frames + audio	identical (2/2)
Wan2.2-T2V-A14B 1280x720x33, 40 steps, CFG parallel	2	frames	identical (2/2)
MiniMax-H3 T2VA 1344x768 4 s, 12 steps, Ulysses 2	2	frames + audio	identical (2/2)
MOVA-360p I2VA 640x352x97, 20 steps, TP2	2	frames + audio	identical (2/2)

Speed Tests and Profiling

Host syncs inside the denoising stage per step (rank 0, torch.cuda.set_sync_debug_mode("warn") around the stage, real server requests):

model (request)	main: per step in the loop / per request	this PR: per step in the loop / per request
Qwen-Image-2.1 T2I / edit 1024x1024, 40 steps, 1 GPU	0 / 1	0 / 1
Z-Image-Turbo 1024x1024, 9 steps, TP2	0 / 2	0 / 2
FLUX.3 Klein Base 1024x1024, 50 steps, CFG parallel	0 / 1	0 / 1
Ideogram-4 FP8 1024x1024, 20 steps, TP2	0 / 5	0 / 5
Cosmos3-Nano T2I 1280x720, 35 steps, CFG parallel	1 / 42	1 / 42 (the CFG-window read #45967 removes)

Client latency, interleaved A/B (main, PR, PR, main), fresh server per run, median of 5 measured requests per run:

cell	round	main	this PR (with #46045)
Ideogram-4 FP8 1024x1024, 20 steps, TP2	1	3.671 s	3.588 s
	2	3.676 s	3.589 s
	median	3.673 s	3.588 s (-2.3%)
Qwen-Image-2.1 edit 1024x1024, 40 steps, 1 GPU	1	5.288 s	5.294 s
	2	5.312 s	5.296 s
	median	5.300 s	5.295 s (-0.1%, noise)

Ideogram-4: the denoising stage went from 3.38 s to 3.29 s; the gain is the LLM projection (and its TP all-gather) no longer running in each of the 40 forwards (20 steps x 2 CFG branches). Each run's 5 requests spread under 0.5%.
Qwen-Image-2.1 had no per-step syncs on main, so only the per-request decode-stats upload goes away; the change is within noise.
The per-step sync removals for GLM-Image, HunyuanImage-3, Krea-2, SenseNova-U1, Qwen-Image-Layered/batched prompts, FLUX.3 Action and SANA were not benchmarked on GPU here; their unit tests cover them.

Checklist

Format your code according to the Format code with pre-commit.
Add unit tests according to the Run and add unit tests.
Update documentation according to Write documentations.
Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
Follow the SGLang code style guidance.

🤖 Generated with Claude Code

CI States

Latest PR Test (Base): ✅ Run #40059865618
Latest PR Test (Extra): ❌ Run #40059865542
Latest PR Test (AMD ROCm 10): ❌ Run #40059865627

A denoising step that reads a device value on the host (item(), float(t), if t:, nonzero) or uploads a constant with a pageable copy blocks the host until the stream drains, so it cannot queue the next step while the GPU runs this one. Regressions of this kind pass every numeric test. test/unit/host_sync_guard.py gives unit tests three guards: - no_cuda_sync(allow=...) / assert_no_cuda_sync(fn, ...): runs the code under torch.cuda.set_sync_debug_mode plus explicit synchronize() calls, and fails with the call site of every sync whose issuing function is not allow-listed. Skips without CUDA; restores the debug mode. - SyncOnRead: a CPU stand-in for device tensors whose host reads raise, so the same check runs in CPU-only tests. - forbid_host_syncs: a dispatch mode that fails on the ops that sync on CUDA (scalar reads, nonzero, boolean-mask indexing, data-dependent shapes), for CPU tests of code that builds its own tensors. test_denoising_host_sync.py applies both to one CFG request through the shared DenoisingStage loop with the flow-match Euler scheduler. The contributing guide documents the pattern for model tests. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

1. NDRotaryEmbedding: each axis's 1D generator keyed an LRU cache on tuple(pos.tolist()). That reads the positions back to the host (a sync per axis), hashes thousands of floats, and on a miss rebuilds the tensor with a pageable upload. FLUX.1/Kontext/FLUX.2 pay it per image, text and condition axis on every request, as do JoyImage, LingBot and the Wan sequence-shard path. OneDRotaryEmbedding.forward now computes from the device positions directly, with the same ops on the same values; the outputs are bitwise identical (dtype and strides too). 2. apply_qk_norm_rope, apply_flashinfer_rope_qk_inplace and both RotaryEmbedding paths rebuilt the [offset, offset + seq) positions with an arange (and a repeat for batch > 1) in every block when the caller passes none: FLUX.2, Z-Image, Qwen-Image, Wan and others, 64-120 launches per step in eager mode. rope_positions() shares one read-only tensor per shape; compile and graph capture still build their own, so a capture never fills the cache with an unexecuted arange. Tests (test_rope_host_sync.py, fail on main): FluxPosEmbed with SyncOnRead positions matches the plain run bitwise; positions are built once per shape, rebuilt under compile and capture, and apply_qk_norm_rope takes them from rope_positions with the same output as explicit positions. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…edulers 1. PiflowScheduler (Kandinsky-6 distilled): the policy rollout looped range(substeps.max().item()), a device read right after the DiT forward on every step, for the video and the audio scheduler. The host then launched ~1k substep ops with the GPU idle. set_timesteps now derives each step's substep count from a host copy of the same fp32 math; the device substep sizes and masks are unchanged. 2. HeliosScheduler UniPC: torch.tensor() over lists of 0-dim device tensors read every rk and b back to the host, and the 0.5 constants were pageable uploads, about a dozen syncs per step. They are now stacked and filled on the device (same dtypes and values), and linalg.solve becomes linalg.solve_ex, which skips the host-side singularity check over the same kernels. 3. HeliosScheduler DMD: torch.full() with a device timestep read it on the host twice per step; the timestep is now cast and expanded on the device (same truncation to int64). Outputs are bitwise identical to main on CPU, NaN bits included. test_scheduler_host_sync.py runs both schedulers with SyncOnRead schedules and latents against plain runs; it fails on main. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

DmdDenoisingStage (Wan DMD / FastWan) synced twice per step: 1. t_int = int(t.item()) on the device timestep, although only the two-expert branch uses it. It now comes from a host copy of the same dmd_denoising_steps list. 2. assert not torch.isnan(latent_model_input).any() read a device flag every step. The flags are OR-ed on the device and checked once after the loop, with the same message; outputs are unchanged, a NaN request still fails. test_denoising_host_sync.py runs one DMD request with SyncOnRead device timesteps against a plain run (bitwise equal; fails on main) and checks that NaN latents still raise. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

FlashAttentionImpl.forward_ring_kv_chunk built cu_seqlens_q/k with torch.tensor(..., device=...) on every ring hop of every block. Each is a pageable host-to-device copy that blocks the host until the stream drains, which also stalls the hop's P2P/compute overlap. The ring tail-pad path (ring > 1 on the FA backend) hit it twice per hop. The bounds depend only on the chunk lengths, so the hop now takes them from the cached cu_seqlens_on_device, moved from layer.py into the FA backend so both the local varlen path and the ring path share it. The values and dtype are unchanged. test_attention_host_sync.py checks two ring hops get [0, n] int32 bounds and reuse the same device tensors; it fails on main. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…in Qwen-Image - Layered built `additional_t_cond` with a pageable `torch.tensor([0])` upload per branch. The two branches got different objects, so the DiT's modulation cache never hit on the negative branch: every step recomputed 60 blocks x 2 modulation GEMMs. Both branches now share one on-device zero tensor. - Batched prompts of different lengths rebuilt their varlen metadata every forward, uploading the lengths from pageable memory: a host sync per step and CFG branch. The metadata is cached by its lengths; each forward gets a fresh dict because attention may memoize into it. - Layered's condition-image encode, the shared channel-wise VAE stats helper (Qwen-Image edit encode/decode, Krea2, JoyImage) and Qwen-Image-2.1 decode uploaded latent mean/std per request. Each pageable upload stalls the host until the queued VAE encode or denoising loop drains. They are now uploaded once per (device, dtype). All values are unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…ost once `cosmos_unipc` iterated the scheduler's timesteps, which live on the device, and called `float(tick)` on each one. Every step therefore blocked the host until the previous step's kernels drained. The loop now copies the ticks to the host once per request and still hands each step the same device tick. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

… per request - Every DiT forward cast the (B, T+I, 53248) fp32 LLM features to bf16, masked them, ran the RMSNorm and the 53248 -> 4608 projection (plus its TP all-gather), all on inputs that never change within a request. The denoising stage now projects each branch's conditioning on its first step, inside that step's forward context, and passes it to later steps as `llm_cond`. BCG replays keep passing the raw features, which their prompt padder pads, and the progressive stage drops the projection when the sequence grows. - Decode uploaded the latent scale/shift per request; that pageable upload stalled the host until the denoising loop drained. They are uploaded once per (device, dtype). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…t sync - The eager forward dropped the unconditional branch's prior tokens with `prior_embedding[prior_token_drop] *= 0.0`. The boolean-mask read and write each run a nonzero, so every forward (every step x CFG branch) blocked the host twice. It now selects `prior_embedding * 0.0` with `torch.where`, which keeps the exact values, signed zeros included. - Decode uploaded the latent mean/std per request, stalling the host until the denoising loop drained; they are uploaded once per (device, dtype). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Without sequence parallelism the pipeline config always handed the DiT a key mask, even for a single prompt with no padding. Every forward then ran `bool(mask.all())` to find out the mask was all valid, which blocked the host once per step and CFG branch. The config now makes that decision once per request and passes `mask=None` for unpadded prompts, the same maskless path the DiT took before. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…off the host Each step blocked the host on the device several times: - `t.repeat(bs).to(device)` uploaded the CPU timestep from pageable memory. - `timestep_embedding` built its frequencies on the host and uploaded them on every call, and a CFG step embeds timesteps five times. - The image-token scatter and the prediction gather ran `masked_select` on the request-constant image mask (and the legacy boolean timestep mask), whose output size is only known on the host. - The edit path uploaded the condition timesteps every step. The loop now takes the mask positions, the device timesteps and the condition timesteps once per request, gathers the predictions with those positions, and uploads the host-built frequencies once per device. Values are unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…p syncs - `t2i_generate` checked the CFG window with `t >= lo and t <= hi` on the device timestep, two host syncs per step. The window is now computed for the whole schedule with the same float32 comparisons and read once per request. - The timestep embedder built its frequencies on the host and uploaded them on every call, which blocked the host each step; they are uploaded once per device. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Every block's cross-attention cast the int64 tokenizer mask to bool, one launch per block for a mask that is the same for the whole forward (20 blocks x 2 CFG branches per step on SANA 1.6B, which its own docstring notes is launch-bound in eager mode). The forward now casts it once; the blocks' cast becomes a no-op on the bool mask. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

mintlify · 2026-10-23T11:56:57Z

Preview deployment for your docs. Learn more about Mintlify Previews.

Project	Status	Preview	Updated
sglang-doc	🟢 Ready	View Preview	Oct 24, 2026, 3:49 AM

💡 Tip: Enable Automations to automatically generate PRs for you.

mickqian · 2026-10-23T11:57:19Z

/tag-and-rerun-ci

Fold the guard's own tests and the shared denoising, scheduler, RoPE and attention host-sync tests into test_host_sync.py, and compare outputs with torch.testing.assert_close(rtol=0, atol=0). Add the helpers the model host-overhead tests share to host_sync_guard.py: RecordOps (aten op recorder counting casts and host-data uploads), CountedUploads (host tensor counting its copies to the device), and tiny_dit_runtime / reference_custom_ops (SDPA attention and reference CustomOps for tiny single-process DiTs). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…ffusion-image-models-host-overhead Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

… model files GLM-Image, HunyuanImage-3, Krea-2 and SANA tests join the models' existing test files (test_glm_image_multi_output.py, test_hunyuan_image3_tp_broadcast.py, test_krea2_fp8.py and the SANA family's test_sana_video.py), and SANA counts its bool casts with the shared RecordOps. No test is dropped. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

mickqian and others added 13 commits October 23, 2026 16:17

mickqian requested review from AgainstEntropy, BBuf, HaiShaw, JustinTong0323, kevin-mii, niehen6174, ping1jing2, sogalin, wisclmy0611, yichiche and zijiexia as code owners October 23, 2026 11:56

github-actions Bot added documentation Improvements or additions to documentation diffusion SGLang Diffusion labels Oct 23, 2026

github-actions Bot added the run-ci CI: run the baseline test suite on this PR label Oct 23, 2026

mintlify Bot deployed to staging - docs October 23, 2026 12:02

View deployment

mickqian and others added 3 commits October 24, 2026 11:26

Merge branch 'mick/diffusion-host-sync-guard-and-shared' into mick/di…

e9fa058

…ffusion-image-models-host-overhead Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

mintlify Bot deployed to staging - docs October 24, 2026 03:49

View deployment

mickqian merged commit c9f4055 into main Oct 24, 2026
145 of 153 checks passed

mickqian deleted the mick/diffusion-image-models-host-overhead branch October 24, 2026 13:33

This was referenced Oct 24, 2026

[CI Monitor] Daily Report - 2026-10-24 bingxche/sglang-ci-bot#243

Open

[CI Monitor] Daily Report - 2026-10-25 bingxche/sglang-ci-bot#244

Open

[CI Monitor] Daily Report - 2026-10-27 bingxche/sglang-ci-bot#249

Open

javey-q mentioned this pull request Nov 3, 2026

[Feat] SenseNova-U1: CFG parallelism on two GPUs (1.95x, bit-identical) #45161

Open

5 tasks done

Navigation Menu

Repository navigation

[Perf] diffusion: Remove per-step host syncs and uploads in image models - #46046

[Perf] diffusion: Remove per-step host syncs and uploads in image models#46046
mickqian merged 16 commits into
mainfrom
mick/diffusion-image-models-host-overhead

mickqian commented Oct 23, 2026 •

edited by github-actions Bot

Loading

Uh oh!

mintlify Bot commented Oct 23, 2026 •

edited

Loading

Uh oh!

mickqian commented Oct 23, 2026

Uh oh!

Uh oh!

This branch was successfully deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

mickqian commented Oct 23, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy Tests

Speed Tests and Profiling

Checklist

CI States

Uh oh!

mintlify Bot commented Oct 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mickqian commented Oct 23, 2026

Uh oh!

Uh oh!

This branch was successfully deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

mickqian commented Oct 23, 2026 •

edited by github-actions Bot

Loading

mintlify Bot commented Oct 23, 2026 •

edited

Loading