Skip to content

[Diffusion] Optimize LTX-2 QKNorm and split RoPE on Hopper - #38396

Merged
BBuf merged 1 commit into
sgl-project:mainfrom
BBuf:codex/ltx2-hopper-qknorm-rope
Sep 8, 2026
Merged

BBuf merged 1 commit into
sgl-project:mainfrom
BBuf:codex/ltx2-hopper-qknorm-rope

Conversation

@BBuf

@BBuf BBuf commented Sep 8, 2026 •

Copy link
Copy Markdown
Collaborator

Summary

  • enable the existing fused LTX-2 Q/K RMSNorm + split-RoPE CUDA kernel on exact SM90 when the request opts into quality="high" or "extra-high"
  • keep Hopper quality="lossless" on the current eager chain; keep the existing SM100+ default unchanged
  • route the Hopper opt-in through the existing request-scoped quality-fusion mount/unmount policy
  • add SM90 guard, numeric-tolerance, and mount-policy coverage

Profile evidence

Workload: Lightricks/LTX-2, 768x512x121, 40 first-stage + 3 refinement steps, CFG parallel on 2x H200, eager mode.

The source all-stage trace recorded 41,316 BF16 LayerNorm launches (271.51 ms) and 16,512 split-RoPE launches (85.37 ms). A two-step candidate trace records 1,152 fused kernel launches, zero split-RoPE launches, and removes a net 192 launches per step.

End-to-end result

Fixed H200 GPUs, two fresh-process repeats per side:

Path Mean denoise Mean E2E Peak reserved / rank
source, quality=high 8.5027 s 9.7970 s 64.56 GiB
candidate, quality=high 7.6832 s 9.0633 s 64.56 GiB
improvement 9.64% 7.49% unchanged

The candidate repeats are byte deterministic. The fused operator differs from the Hopper eager BF16 chain by at most 0.0625 for Q and 0.125 for K, which is why this path is quality-gated instead of enabled for lossless requests. Candidate versus source high-mode video measures PSNR 25.5355 dB / FFmpeg SSIM 0.897576; visual review preserves the scene, composition, and broad motion. Source high mode itself differs more from lossless by frame metrics.

Validation

  • test_rope_ltx2.py on H200: 2 passed, 5 Blackwell-only cases skipped
  • test_model_fast_paths.py -k ltx2 on H200: 6 passed
  • pre-commit on all changed files: passed

CI States

Latest PR Test (Base): ⏳ Run #34181428105
Latest PR Test (Extra): ❌ Run #34181428049
Latest PR Test (AMD ROCm 7.2): ❌ Run #34181428196

@mickqian

mickqian commented Sep 8, 2026

Copy link
Copy Markdown
Collaborator

/tag-and-rerun-ci

@github-actions github-actions Bot added the run-ci CI: run the baseline test suite on this PR label Sep 8, 2026
@BBuf

BBuf commented Sep 8, 2026

Copy link
Copy Markdown
Collaborator Author

@BBuf
BBuf merged commit 554f817 into sgl-project:main Sep 8, 2026
266 of 309 checks passed
@BBuf
BBuf deleted the codex/ltx2-hopper-qknorm-rope branch September 8, 2026 11:05
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

diffusion SGLang Diffusion jit-kernel run-ci CI: run the baseline test suite on this PR

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants