[mickqian]

Motivation

Every LTX-2 denoising step pays two host-side costs.

1. Sigma reads sync the host. Each step moved scheduler.sigmas[i] and [i+1] to the device and called .item() on both. The stage-1 guider path called .item() on sigma again after the forward. Each read blocks the host until the stream drains, so the host cannot enqueue the next step while the GPU runs the current one. Stage 2 also paid a pageable H2D copy per sigma, because its distilled schedule lives on the CPU. Once the forward is graph-replayed (breakable CUDA graphs), this sync is the bubble that remains. Removing it lets the host run ahead.
2. The prompt mask is rebuilt in every block. The 0/1 prompt key mask goes to every block. Each block's two prompt cross-attentions (LocalAttention) turned it into the additive SDPA mask again: a cast, a subtract and a multiply per call. That is 288 small launches per forward over 48 blocks.

Modifications

• ltx_2/denoising.py: _prepare_denoising_loop takes fp32 host and device copies of the sigma schedule once per loop (LTX2DenoisingContext.sigmas_host / sigmas_device). Steps index the device copy for tensor math and the host copy for Python scalars. The guider's x0 conversion takes the device sigma instead of its Python float. _ltx2_velocity_to_x0 scales in fp32 either way.
• layers/attention/layer.py: _prepare_sdpa_mask is renamed prepare_sdpa_mask. LocalAttention now calls it instead of keeping its own inline copy of the same conversion. The values are the same for float and 0/1 masks.
• models/dits/ltx_2.py: the model forward converts the video and audio prompt masks once with prepare_sdpa_mask, and the blocks pass the float mask through. LocalAttention only casts and reshapes a float mask. Prompt cross-attention is always LocalAttention, and the prompt mask never goes through the SP K/V gather.

Tests:

• test_ltx2_bcg_coords.py, test_steps_never_read_device_sigmas_on_host: runs one denoising step in each guidance mode (official CFG, batched guider, TI2V, HQ, legacy, legacy batched). The sigma schedule's item() / __float__ / __bool__ raise, and the outputs must equal a run with a plain tensor. It fails on main in every mode.
• test_ltx2_bcg_coords.py, test_tensor_sigma_x0_matches_python_float: the fp32 x0 conversion gives bitwise the same result for a tensor sigma and a float sigma.
• test_ltx2_prompt_mask.py: LocalAttention gives bitwise the same output for a prebuilt additive mask and for the raw 0/1 mask. A 2-block LTX-2 model converts the masks exactly twice per forward (video and audio), and its output matches the per-block conversion bitwise.

No files are shared with #45942, #45943, #45944 or #45948, or with "Shard LTX-2 stage 2 over the CFG ranks as Ulysses". This PR merges cleanly with all of them.

Accuracy Tests

2xH200, LTX-2 two-stage, default profile. Decoded frames and audio are bitwise identical across base, the sigma commit and the prompt-mask commit, in every round.

Speed Tests and Profiling

2xH200, LTX-2 two-stage, default (eager) profile. Three interleaved rounds, e2e median per run:

round	base	+ sigma reads without sync	+ prompt mask once per forward
1	7.571 s	7.392 s	7.161 s
2	7.507 s	7.487 s	7.248 s
3	7.395 s	7.422 s	7.260 s

• Prompt-mask commit vs sigma commit: −0.231 / −0.239 / −0.162 s (−2.9%), consistent in every round. The stage-1 step went from 0.130 s to 0.123 s.
• The sigma commit alone cannot be separated from noise in eager mode: its sign differs between rounds. It pays off under graph replay, where the per-step sync is what stops the host from running ahead.
• The base for these numbers was main plus the fixes in [Fix] diffusion: Keep the linear-GELU verdict cache out of torch.compile traces #45944 and [Fix] diffusion: Keep dynamic LoRA deltas on fused linear-GELU sites #45943 (the latter in an earlier revision). This PR does not depend on either.

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 #39908805023
Latest PR Test (Extra): ❌ Run #39908805019
Latest PR Test (AMD ROCm 10): ❌ Run #39908805036

[mickqian]

/tag-and-rerun-ci