[mickqian]

Motivation

A denoising step that reads a device value on the host (item(), float(t), if t:, nonzero, boolean-mask indexing) or uploads a constant with a pageable copy blocks the host until the stream drains. The host then cannot queue the next step while the GPU runs this one. These regressions pass every numeric test, so they come back. This PR adds unit-test guards for them and removes the ones in shared code: RoPE position setup, two schedulers, the DMD stage and the ring-attention hop.

The model-specific fixes are in two follow-ups that use these guards: image models (#) and video models (#).

Modifications

Guards (test/unit/host_sync_guard.py):

• no_cuda_sync(allow=...) / assert_no_cuda_sync(fn, ...) run the code under torch.cuda.set_sync_debug_mode("warn") plus explicit synchronize() calls, and fail with the call site of every sync whose issuing function is not allow-listed. They skip without CUDA and restore the debug mode.
• SyncOnRead is a CPU stand-in for device tensors (timesteps, sigmas, masks): host reads of its values raise, and .cpu() is the explicit copy that is allowed.
• forbid_host_syncs() is a TorchDispatchMode that fails on the ops that sync on CUDA whatever tensor they read: scalar reads, nonzero, boolean-mask indexing, data-dependent output shapes.
• The contributing guide documents the three.

Shared hot paths:

• rotary_embedding/mrope.py: OneDRotaryEmbedding keyed an LRU cache on tuple(pos.tolist()), a device read per axis plus a hash over thousands of floats, and a pageable upload on a miss. FLUX.1/Kontext/FLUX.2, JoyImage, LingBot and the Wan sequence-shard path paid it per image, text and condition axis on every request. It now computes from the device positions directly (same ops, same values).
• rotary_embedding/{utils,base}.py, layernorm.py: apply_qk_norm_rope, apply_flashinfer_rope_qk_inplace and both RotaryEmbedding paths rebuilt [offset, offset + seq) with an arange (and a repeat for batch > 1) in every block when the caller passes no positions: 64-120 launches per step in eager mode for FLUX.2, Z-Image, Qwen-Image, Wan and others. 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 arange it did not execute.
• schedulers/kandinsky6_piflow.py: the policy rollout looped range(substeps.max().item()), a device read right after the DiT forward on every step (video and audio scheduler). The substep counts now come from a host copy of the same fp32 math.
• schedulers/scheduling_helios.py: UniPC built torch.tensor() over lists of 0-dim device tensors and uploaded 0.5 constants, about a dozen syncs per step; they are stacked and filled on the device, and linalg.solve becomes linalg.solve_ex (same kernels, no host-side singularity check). The DMD path read the timestep twice per step through torch.full(..., tensor); it is now cast and expanded on the device.
• stages/denoising_dmd.py (Wan DMD / FastWan): int(t.item()) per step now comes from a host copy of the same schedule, and the per-step NaN assert is OR-ed on the device and checked once after the loop (same message; a NaN request still fails).
• attention/backends/flash_attn.py, attention/layer.py: the ring hop built cu_seqlens_q/k with torch.tensor(..., device=...) on every hop of every block, twice per hop for the tail-pad path. The hop now uses the cached cu_seqlens_on_device, moved next to the FA backend so the varlen and ring paths share it.

Tests (the fix tests fail on main):

• test_host_sync_guard.py: the three guards catch scalar reads, explicit synchronizes, boolean-mask indexing and data-dependent shapes, let device-side selects through, honor the allow list and restore the debug mode.
• test_denoising_host_sync.py: one CFG request through the shared DenoisingStage loop with SyncOnRead timesteps matches a plain run, and runs under no_cuda_sync on CUDA; one DMD request with SyncOnRead timesteps matches a plain run bitwise, and NaN latents still raise.
• test_rope_host_sync.py: FluxPosEmbed with SyncOnRead positions matches bitwise; positions are built once per shape, rebuilt under compile and capture.
• test_scheduler_host_sync.py: PiFlow and Helios UniPC/DMD steps with SyncOnRead schedules match plain runs bitwise, NaN bits included.
• test_attention_host_sync.py: two ring hops get [0, n] int32 bounds from the same device tensors.

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. Decoded outputs are bitwise identical to main for every model below, two requests each with different inputs (pixels for images, decoded frames and audio for video), with this PR and both follow-ups applied:

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, counted with torch.cuda.set_sync_debug_mode("warn") around the stage's forward on a real server request (rank 0, second request, per denoising step):

None of the twelve models above runs the code this PR changes in a way that syncs per step: their counts move only through the follow-ups (#, #). The paths changed here are FLUX.1/FLUX.2/JoyImage/LingBot RoPE setup, Kandinsky-6 PiFlow, Helios, Wan DMD / FastWan and ring attention on the FA backend. A GPU sync count and A/B for those (FastWan2.2-TI2V DMD, Wan2.1 with ring degree 2, FLUX.2 klein) is still to be run.

The CUDA guard tests (no_cuda_sync) ran on an H200 and pass.

Keep this PR in draft until the A/B above exists; the RoPE-position change removes 64-120 small launches per eager step, which should show on launch-bound models (FLUX.2, Z-Image), but it has not been measured in isolation.

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 #40059851951
Latest PR Test (Extra): ❌ Run #40059851924
Latest PR Test (AMD ROCm 10): ❌ Run #40059851923

[mintlify]

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]

/tag-and-rerun-ci