Repository navigation
[Perf] diffusion: Remove per-step host syncs and uploads in video models - #46047
Merged
Merged
Conversation
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>
- The loop uploaded both timesteps to the device every step (blocking copies), then passed the device copies to step_from_to, whose float() read them back: four stream syncs per step. The paired schedule is now uploaded once per request; the scheduler steps on the host values. - Every forward re-uploaded the CPU complex128 RoPE tables of both towers (six blocking copies, ~17 MB for audio) and redid the full-sequence concats; every bridge layer rebuilt its fp32 cos/sin caches (30 layers x 2 directions). The tables and the two bridge caches are now built once per request and freed when the loop ends. All values are unchanged: the same copies, casts and concats run once. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…and AdaLN cache - FastH3 PDD (release pdd_fused_heads): every step copied that step's video/audio head weight and bias from the host stacks with blocking copies, four stream syncs per step. The stacks are now uploaded once per (device, dtype) and indexed by step. - AdaLN online cache: block_all/final indexed the slabs with the 0-dim device slot, which reads the slot back to the host, two syncs per step. They now gather with index_select and keep the old strides. Values are unchanged: the same weights and slab rows are read. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The forward built the RoPE cos/sin on the host (cached) and copied both to the device on every forward: two blocking copies, ~59 MB per step at 720p x 125 frames, each stalling the host on the previous step's kernels. The uploaded tables are now kept per request (in an underscore Req.extra key, which the disaggregation transport skips). The host computation is unchanged, so the values are identical. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ck modulation casts - The stage built the frame indices on the host. Every forward embeds them four times (current, short, mid, long RoPE), and each embedding copied them to the device with a blocking copy: eight stream syncs per CFG step. They are now built on the latent device; the small integers are exact. - Every block cast the full-sequence [B, L, 6, D] modulation to fp32 before adding its table (~1.4 GB per block at 384x640). The compact per-timestep rows are now cast once before the expand/concat, so the blocks' .float() is a no-op; casting commutes with expand and cat, so the values are identical. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…block modulation casts - Causal Wan (LongLive2): every forward copied the host-built RoPE tables to the device with blocking copies, two stream syncs per forward. They are now uploaded once per request and chunk position, in an underscore Req.extra key (the disaggregation transport skips those). - Wan and causal Wan blocks each cast the batch-wide modulation to fp32 before adding their table; the forward now casts it once (the per-token TI2V table keeps its fused bf16 path). Outputs are bit-identical: the same values reach the same ops. The DMD stage's per-step syncs are left to the shared-stage change. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…syncs
Paths not covered by the open LTX-2 step-overhead PR; the hunks are in
separate functions from it.
- prepare_video_coords uploaded the patch-size delta and scale factors
with two blocking copies per call. It runs every forward on the LTX-2.3
one-stage path and per step for JoyEcho memory coords; the constants are
now uploaded once per device.
- HQ res2s (opt-in two-stage HQ pipeline):
- _ltx2_get_sde_coeff replaced NaNs through a 0-dim boolean index, a
nonzero sync on get and on set, four calls per res2s step. torch.where
selects the same values.
- The 100-iteration bongmath refinement re-cast the denoised estimates
to fp64 and recomputed h * a21 every iteration (~400 extra launches and
200 full-latent casts per refined step). Both are hoisted;
h * a21 * eps already evaluated as (h * a21) * eps.
- The stage-1 HQ step size read both device sigmas back with .cpu(); it
now starts from the step's host sigmas, the same fp32 values.
Values are bit-identical.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…t rebuilds - Every DMD step read sigma and sigma_next back from the device schedule with .item(), two stream syncs per step. The loop now takes fp32 host and device copies of the schedule once (in ctx.extra, without touching the LTX-2 context) and steps index them. - With the memory bank (shot >= 2), every step rebuilt the shot-constant memory inputs: the a2v/v2a cross masks, the audio block mask and the memory video/audio RoPE coords. They are now built once per loop; each step only concatenates the latents and prefixes the timesteps. Values are bit-identical. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ost syncs SANA-WM streaming / realtime: - The cache-type flag was uploaded to the device in every block's clean pass and read back with .item() by the next chunk's cache assembly (20 uploads + 20 reads per chunk). It only steers Python control flow; it is now a host tensor. - Every self-forcing step read its device timestep with int(t.item()) for a forward context the realtime path ignores, and the first step looked its index up with nonzero().item(). The batch path now reads the schedule once per chunk; the scheduler starts at index 0 (the per-token update reads no step index). - forward_long re-sliced the camera / Plucker tensors on every call, so the UCPE and Plucker caches never hit (full UCPE rebuild per step), and rebuilt the frame_index RoPE. The windowed inputs are now built once per chunk. - UCPE: `if bool(invalid.any()): raymats[invalid] = eye` synced per build; torch.where selects the same values. compute_chunk_plucker read its frame indices back with .tolist(); they are now built on the host. - Chunked refiner: each block read its sigmas with 7 .item() calls and built the same RoPE four times with three uploads each. The runner keeps a host copy of the schedule, the builder makes positions, patch ends and scales on the device, and a block's levels and capture share one table. SANA-Video2: every forward uploaded the host complex128 RoPE frequencies twice and rebuilt both tables; they are now built once per latent shape. Outputs are bit-identical. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The default Cosmos3 T2V/I2V loop has no per-step syncs; these are in its world-model variants. - Dreams: every GEN layer zeroed the null-action V slots by indexing with a Python position list, a blocking upload of the list in each of 36 layers per forward (frame 0 always; every frame without actions). The slots are now given as per-frame slices. - Sim-Auto stage: each step's scatter/gather built the generated-frame indexes from host lists (two torch.tensor uploads and two list indexes). _FlatLayout now keeps them as device indexes built once per chunk; camera_major_slice/write take a device index. - Sim-Auto transformer: every step re-uploaded each item's frame ids for the mRoPE ids and each target's generated-frame mask. Both are now built once per chunk, keyed by the chunk's items. Outputs are bit-identical: the same slots are zeroed and the same rows copied. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
mickqian
requested review from
AgainstEntropy,
BBuf,
HaiShaw,
JustinTong0323,
kevin-mii,
niehen6174,
ping1jing2,
sogalin,
wisclmy0611,
yichiche and
zijiexia
as code owners
October 23, 2026 11:57
Contributor
|
Preview deployment for your docs. Learn more about Mintlify Previews.
💡 Tip: Enable Automations to automatically generate PRs for you. |
Collaborator
Author
|
/tag-and-rerun-ci |
4 of 5 tasks
The loop test checks each forward's device timesteps against the scheduler's own pairs instead of a second reference run. The RoPE test counts the bridge table builds and cos/sin cache conversions in one run, drops the counting tensor subclass on the DiT frequencies (every table is built in _rope_freqs), and compares against per-step rebuilds with a bitwise assert_close. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
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-video-models-host-overhead Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Brings in test_mova_denoising_loop.py, the MOVA tests' home. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… model files MOVA, Helios, Wan/causal Wan, LTX-2, SANA-WM, Cosmos3 Dreams and Sim-Auto tests join the models' existing test files; HunyuanVideo and JoyEcho, which have none, share test_video_models_host_sync.py. The tests use the shared CountedUploads, RecordOps and tiny_dit_runtime helpers and the single_process_model_parallel fixture instead of per-file copies, and compare outputs with assert_close(rtol=0, atol=0). No test is dropped. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This was referenced Oct 25, 2026
This was referenced Nov 4, 2026
niehen6174
added a commit
to niehen6174/sglang
that referenced
this pull request
Nov 5, 2026
Reviewed the six upstream minimax_h3.py changes since the last pin (sgl-project#46047, sgl-project#46709, sgl-project#47567, sgl-project#47569, sgl-project#47566 and the BCG output buffers). None touches what h3_comfyui_v1 relies on in its single-GPU eager torch_sdpa profile: DiTBlock.forward is unchanged, _attend still calls fused_inplace_qknorm_rope on the _use_fused_qknorm_rope branch with the same arguments (the new pipelined branch needs Ulysses, which the profile rejects), and _embed still returns (embeddings, t_emb) in eager mode.
This branch was successfully deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Depends on #46045 (host-sync guards in
test/unit/host_sync_guard.py); this branch is stacked on it.Motivation
Video pipelines run long denoising loops where every blocking host read or pageable upload stalls the host on the previous step's kernels, and several of them rebuilt request-constant RoPE tables or masks on every forward. This PR removes those per-step costs model by model. Outputs are bitwise unchanged.
Modifications
model_specific_stages/mova.py,bridges/mova_dual_tower.py): the loop uploaded both timesteps every step andstep_from_toread them back withfloat(), four syncs per step. The paired schedule is uploaded once and the scheduler steps on the host values. Every forward re-uploaded both towers' CPU complex128 RoPE tables (six blocking copies, ~17 MB for audio) and redid the full-sequence concats, and every bridge layer rebuilt its fp32 cos/sin caches (30 layers x 2 directions); tables and bridge caches are built once per request and freed when the loop ends.dits/minimax_h3.py,minimax_h3_adaln_cache.py): FastH3 PDD heads copied that step's weight and bias from the host stacks, four blocking copies per step; the stacks are uploaded once per (device, dtype) and indexed by step. The AdaLN online cache indexed its slabs with a 0-dim device slot (a host read, two per step); it now gathers withindex_selectand keeps the old strides.dits/hunyuanvideo.py): the forward built the RoPE cos/sin on the host and copied both to the device every forward (~59 MB per step at 720p x 125 frames). The uploaded tables are kept per request in an underscoreReq.extrakey, which the disaggregation transport skips.helios_denoising.py,dits/helios.py): frame indices were built on the host and uploaded by each of four RoPE embeddings per forward (eight syncs per CFG step); they are built on the latent device. Every block cast the full-sequence [B, L, 6, D] modulation to fp32 (~1.4 GB per block at 384x640); the compact rows are cast once before the expand/concat.dits/causal_wanvideo.py,dits/wanvideo.py): causal Wan (LongLive2) copied its host RoPE tables to the device every forward; they are uploaded once per request and chunk position. Wan and causal Wan blocks each cast the batch-wide modulation to fp32; the forward casts it once (the per-token TI2V table keeps its fused bf16 path).dits/ltx_2.py,ltx_2/denoising.py; paths not touched by [Perf] diffusion: Cut LTX-2 denoising-step host overhead #45966):prepare_video_coordsuploaded two constant tensors per call (every forward on the LTX-2.3 one-stage path, every step for JoyEcho memory coords); they are uploaded once per device. HQ res2s: the SDE coefficient NaN replacement went through a 0-dim boolean index (nonzero on get and set, four calls per step) and is now atorch.where; the 100-iteration bongmath refinement hoists its fp64 casts andh * a21; the stage-1 step size starts from the step's host sigmas instead of.cpu()reads.joy_echo/denoising.py): every DMD step read sigma and sigma_next with.item(); the loop takes fp32 host and device copies of the schedule once. With the memory bank (shot >= 2) the shot-constant cross masks, audio block mask and memory RoPE coords are built once per loop instead of every step.sana_wm*.py,sana_wm/streaming*.py,sana_video2.py): the cache-type flag was uploaded in every block's clean pass and read back by the next chunk (20 uploads + 20 reads per chunk) and is now a host tensor; self-forcing steps read their timestep once per chunk;forward_longwindows camera/Plucker inputs once per chunk so the UCPE and Plucker caches hit; the UCPE invalid-ray fix is atorch.where; the chunked refiner reads its sigmas from a host copy and builds each block's RoPE once on the device. SANA-Video2 builds its RoPE tables once per latent shape instead of uploading complex128 tables twice per forward.dits/cosmos_dreams.py,dits/cosmos3_sim_auto.py,cosmos3_simulation/sim_auto.py): Dreams zeroed null-action V slots by indexing with a Python list in each of 36 layers (a blocking upload each); the slots are per-frame slices. Sim-Auto's scatter/gather indexes, mRoPE frame ids and generated-frame masks are built once per chunk on the device. The default Cosmos3 T2V/I2V loop is untouched ([Perf] diffusion: Remove per-step and per-request host syncs in Cosmos3 and Qwen-Image-2.1 #45967 covers its CFG window).Tests:
test_mova_denoising.py: steps never upload or read device timesteps (SyncOnReadschedule, one upload); RoPE tables built once per request; outputs bitwise equal.test_minimax_h3_pdd_offline.py,test_minimax_h3_adaln_cache.py: heads uploaded once; slab gathers run underforbid_host_syncswith the old values and strides.test_hunyuanvideo_rope_upload.py,test_wan_step_host_overhead.py: RoPE tables uploaded once per request (per chunk position for causal Wan), identical tables; blocks receive fp32 modulation.test_helios_step_host_overhead.py: frame indices live on the latent device; blocks receive fp32 modulation.test_ltx2_step_host_overhead.py: video-coord constants uploaded once; SDE coefficients underforbid_host_syncs; the refinement casts once.test_joy_echo_denoising_host_overhead.py: host/device sigma copies, steps never read device sigmas, the memory prefix is built once.sana_wm/test_host_overhead.py,test_sana_video2.py,test_cosmos3_simulation_host_overhead.py,test_cosmos_dreams.py,test_cosmos3_sim_auto.py: the SANA-WM, SANA-Video2 and Cosmos3 world-model paths above.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.
Accuracy Tests
2xH200, eager profiles, two requests per model with different inputs. Decoded frames (and audio where the model produces it) are bitwise identical to main:
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):Client latency, interleaved A/B (main, PR, PR, main), fresh server per run, median of 2 measured requests per run:
Checklist
🤖 Generated with Claude Code
CI States
Latest PR Test (Base): ✅ Run #40059879645
Latest PR Test (Extra): ❌ Run #40059879638
Latest PR Test (AMD ROCm 10): ❌ Run #40059879708