From 441910f9265612899192ba03624a109e40678618 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Mon, 10 Aug 2026 09:46:17 +0800 Subject: [PATCH] [diffusion] LTX-2 quality=high fused RMSNorm+modulate + FFN GELU epilogue (H200 ltx23-one-stage denoise 45.85->43.24 s, ~matches torch.compile) (#34172) Co-authored-by: Claude Opus 4.8 --- .../ops/diffusion/ltx2_rmsnorm_modulate.py | 80 +++++++++++++++++ .../runtime/models/dits/ltx_2.py | 86 +++++++++++++++---- .../pipelines_core/stages/denoising.py | 9 ++ .../diffusion/test_ltx2_rms_norm_modulate.py | 71 +++++++++++++++ 4 files changed, 229 insertions(+), 17 deletions(-) create mode 100644 python/sglang/kernels/ops/diffusion/ltx2_rmsnorm_modulate.py create mode 100644 test/registered/kernels/ops/diffusion/test_ltx2_rms_norm_modulate.py diff --git a/python/sglang/kernels/ops/diffusion/ltx2_rmsnorm_modulate.py b/python/sglang/kernels/ops/diffusion/ltx2_rmsnorm_modulate.py new file mode 100644 index 000000000..f2d36ec7e --- /dev/null +++ b/python/sglang/kernels/ops/diffusion/ltx2_rmsnorm_modulate.py @@ -0,0 +1,80 @@ +"""Weightless RMSNorm + adaLN modulate folded into one kernel for LTX-2. + +``rms_norm(x) * (1 + scale) + shift`` at the LTX-2 transformer-block adaLN +sites is otherwise an aten ``F.rms_norm`` plus a separate ``mul``/``add`` +modulate (one reduction kernel plus several pointwise passes per site). This +folds the whole chain into a single ``fused_rmsnorm_scale_shift_bitexact`` +launch. + +The fused kernel reproduces the RMSNorm math via ``rsqrt.approx`` (the +flashinfer CuTe form), which differs from aten's refined ``rsqrtf`` by at +most one bf16 ULP on a small fraction of elements. It is therefore *not* +bit-exact vs the eager reference, so the fold is opt-in per batch: model code +marks its adaLN sites with :func:`mark_ltx2_rms_norm_modulate_site` (default +off, reference path) and the denoising stage calls +:func:`mount_ltx2_rms_norm_modulate` / :func:`unmount_ltx2_rms_norm_modulate` +at batch boundaries for ``quality="high"`` requests. +""" + +from __future__ import annotations + +import torch +from torch import nn + +from sglang.kernels.ops.diffusion.quality_gate import QualityGatedFusion +from sglang.kernels.ops.diffusion.triton.rmsnorm_scale_shift_bitexact import ( + can_use_fused_rmsnorm_scale_shift, + fused_rmsnorm_scale_shift_bitexact, +) + +_SITE_MARKER_ATTR = "_sgl_ltx2_rms_norm_modulate_site" +_SITE_ENABLED_ATTR = "_sgl_ltx2_rms_norm_modulate_enabled" +_FUSION = QualityGatedFusion( + name="LTX-2 RMSNorm+modulate", + marker_attr=_SITE_MARKER_ATTR, + enabled_attr=_SITE_ENABLED_ATTR, +) + +# ``RMSNormNoWeight`` applies no scale, so a ones weight reproduces it exactly. +_ONES_WEIGHT_CACHE: dict[tuple[torch.device, int], torch.Tensor] = {} + + +def mark_ltx2_rms_norm_modulate_site(module: nn.Module) -> None: + """Mark ``module`` as an LTX-2 RMSNorm+modulate fusion site (mounted off).""" + _FUSION.mark(module) + + +def ltx2_rms_norm_modulate_active(module: nn.Module) -> bool: + return _FUSION.is_enabled(module) + + +def mount_ltx2_rms_norm_modulate(root: nn.Module) -> bool: + return _FUSION.mount(root) + + +def unmount_ltx2_rms_norm_modulate(root: nn.Module) -> None: + _FUSION.unmount(root) + + +def _ones_weight(x: torch.Tensor) -> torch.Tensor: + key = (x.device, int(x.shape[-1])) + w = _ONES_WEIGHT_CACHE.get(key) + if w is None: + w = torch.ones(x.shape[-1], device=x.device, dtype=torch.bfloat16) + _ONES_WEIGHT_CACHE[key] = w + return w + + +def can_fuse_ltx2_rms_norm_modulate( + x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor +) -> bool: + if x.dtype is not torch.bfloat16 or not x.is_cuda: + return False + return can_use_fused_rmsnorm_scale_shift(x, _ones_weight(x), scale, shift) + + +def fused_ltx2_rms_norm_modulate( + x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor, eps: float +) -> torch.Tensor: + """``rms_norm(x) * (1 + scale) + shift`` as one kernel (weightless RMSNorm).""" + return fused_rmsnorm_scale_shift_bitexact(x, _ones_weight(x), scale, shift, eps) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py index ce878e904..77063c6c3 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py @@ -10,10 +10,22 @@ import torch import torch.nn as nn import torch.nn.functional as F +from sglang.kernels.ops.diffusion.fused_linear_gelu import ( + can_fuse_linear_gelu, + fused_gelu_active, + fused_linear_gelu_tanh, + mark_fused_gelu_site, +) from sglang.kernels.ops.diffusion.ltx2_qknorm_split_rope import ( can_use_ltx2_qknorm_split_rope_cuda, ltx2_qknorm_split_rope_cuda, ) +from sglang.kernels.ops.diffusion.ltx2_rmsnorm_modulate import ( + can_fuse_ltx2_rms_norm_modulate, + fused_ltx2_rms_norm_modulate, + ltx2_rms_norm_modulate_active, + mark_ltx2_rms_norm_modulate_site, +) from sglang.kernels.ops.diffusion.residual_gate_add import residual_gate_add from sglang.multimodal_gen.configs.models.dits.ltx_2 import LTX2ArchConfig, LTX2Config from sglang.multimodal_gen.runtime.distributed import ( @@ -125,6 +137,29 @@ def adaln_embedding_coefficient(cross_attention_adaln: bool) -> int: ) +def _ltx2_rms_norm_modulate( + block: nn.Module, + rms_norm: nn.Module, + x: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + eps: float, +) -> torch.Tensor: + """``rms_norm(x) * (1 + scale) + shift`` for the LTX-2 adaLN sites. + + Folds the weightless RMSNorm and the modulate into one kernel when the + ``quality="high"`` fusion is mounted on ``block`` and the per-call guard + passes; otherwise the verbatim eager reference chain (the ``lossless`` + default). The fused kernel is not bit-exact (<=1 bf16 ULP) so it is gated + on the request-scoped mount rather than a runtime self-check. + """ + if ltx2_rms_norm_modulate_active(block) and can_fuse_ltx2_rms_norm_modulate( + x, scale, shift + ): + return fused_ltx2_rms_norm_modulate(x, scale, shift, eps) + return rms_norm(x, eps) * (1 + scale) + shift + + def _ltx2_disable_fused_ada_values(exc: Exception) -> None: global _LTX2_FUSED_ADA_VALUES_RUNTIME_DISABLED _LTX2_FUSED_ADA_VALUES_RUNTIME_DISABLED = True @@ -977,10 +1012,14 @@ class LTX2FeedForward(nn.Module): input_is_parallel=True, quant_config=quant_config, ) + mark_fused_gelu_site(self, "proj_in") def forward(self, x: torch.Tensor) -> torch.Tensor: - x, _ = self.proj_in(x) - x = self.act(x) + if fused_gelu_active(self) and can_fuse_linear_gelu(self.proj_in, x): + x = fused_linear_gelu_tanh(x, self.proj_in.weight, self.proj_in.bias) + else: + x, _ = self.proj_in(x) + x = self.act(x) x, _ = self.proj_out(x) return x @@ -1108,6 +1147,7 @@ class LTX2TransformerBlock(nn.Module): # 4. Feedforward layers self.ff = LTX2FeedForward(dim, dim_out=dim, quant_config=quant_config) + mark_ltx2_rms_norm_modulate_site(self) self.audio_ff = LTX2FeedForward( audio_dim, dim_out=audio_dim, quant_config=quant_config ) @@ -1200,8 +1240,8 @@ class LTX2TransformerBlock(nn.Module): ) else: vshift_msa, vscale_msa, vgate_msa = video_ada_values[0:3] - norm_hidden_states = ( - self.rms_norm(hidden_states, self.norm_eps) * (1 + vscale_msa) + vshift_msa + norm_hidden_states = _ltx2_rms_norm_modulate( + self, self.rms_norm, hidden_states, vscale_msa, vshift_msa, self.norm_eps ) attn_hidden_states = self.attn1( norm_hidden_states, @@ -1220,9 +1260,13 @@ class LTX2TransformerBlock(nn.Module): ) else: ashift_msa, ascale_msa, agate_msa = audio_ada_values[0:3] - norm_audio_hidden_states = ( - self.rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_msa) - + ashift_msa + norm_audio_hidden_states = _ltx2_rms_norm_modulate( + self, + self.rms_norm, + audio_hidden_states, + ascale_msa, + ashift_msa, + self.norm_eps, ) attn_audio_hidden_states = self.audio_attn1( norm_audio_hidden_states, @@ -1251,8 +1295,8 @@ class LTX2TransformerBlock(nn.Module): v_prompt_shift, v_prompt_scale = self.get_ada_values( self.prompt_scale_shift_table, batch_size, temb_prompt, slice(None) ) - norm_hidden_states = ( - self.rms_norm(hidden_states, self.norm_eps) * (1 + vscale_q) + vshift_q + norm_hidden_states = _ltx2_rms_norm_modulate( + self, self.rms_norm, hidden_states, vscale_q, vshift_q, self.norm_eps ) mod_encoder_hidden_states = ( encoder_hidden_states * (1 + v_prompt_scale) + v_prompt_shift @@ -1278,9 +1322,13 @@ class LTX2TransformerBlock(nn.Module): temb_audio_prompt, slice(None), ) - norm_audio_hidden_states = ( - self.rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_q) - + ashift_q + norm_audio_hidden_states = _ltx2_rms_norm_modulate( + self, + self.rms_norm, + audio_hidden_states, + ascale_q, + ashift_q, + self.norm_eps, ) mod_audio_encoder_hidden_states = ( audio_encoder_hidden_states * (1 + a_prompt_scale) + a_prompt_shift @@ -1428,8 +1476,8 @@ class LTX2TransformerBlock(nn.Module): ) else: vshift_mlp, vscale_mlp, vgate_mlp = video_ada_values[3:6] - norm_hidden_states = ( - self.rms_norm(hidden_states, self.norm_eps) * (1 + vscale_mlp) + vshift_mlp + norm_hidden_states = _ltx2_rms_norm_modulate( + self, self.rms_norm, hidden_states, vscale_mlp, vshift_mlp, self.norm_eps ) ff_output = self.ff(norm_hidden_states) hidden_states = residual_gate_add(hidden_states, ff_output, vgate_mlp) @@ -1440,9 +1488,13 @@ class LTX2TransformerBlock(nn.Module): ) else: ashift_mlp, ascale_mlp, agate_mlp = audio_ada_values[3:6] - norm_audio_hidden_states = ( - self.rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_mlp) - + ashift_mlp + norm_audio_hidden_states = _ltx2_rms_norm_modulate( + self, + self.rms_norm, + audio_hidden_states, + ascale_mlp, + ashift_mlp, + self.norm_eps, ) audio_ff_output = self.audio_ff(norm_audio_hidden_states) audio_hidden_states = residual_gate_add( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 8f3e91aad..9e59fe530 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -32,6 +32,10 @@ from sglang.kernels.ops.diffusion.fused_ln_modulate import ( mount_fused_ln_modulate, unmount_fused_ln_modulate, ) +from sglang.kernels.ops.diffusion.ltx2_rmsnorm_modulate import ( + mount_ltx2_rms_norm_modulate, + unmount_ltx2_rms_norm_modulate, +) from sglang.multimodal_gen import envs from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType, STA_Mode from sglang.multimodal_gen.configs.pipeline_configs.flux import ( @@ -156,6 +160,11 @@ _QUALITY_FUSION_HANDLERS: tuple[ mount_fused_ln_modulate, unmount_fused_ln_modulate, ), + ( + "LTX-2 fused RMSNorm+modulate", + mount_ltx2_rms_norm_modulate, + unmount_ltx2_rms_norm_modulate, + ), ( "fused gate RMSNorm (BF16-native Triton)", mount_fused_gate_rmsnorm, diff --git a/test/registered/kernels/ops/diffusion/test_ltx2_rms_norm_modulate.py b/test/registered/kernels/ops/diffusion/test_ltx2_rms_norm_modulate.py new file mode 100644 index 000000000..292452469 --- /dev/null +++ b/test/registered/kernels/ops/diffusion/test_ltx2_rms_norm_modulate.py @@ -0,0 +1,71 @@ +"""LTX-2 quality=high RMSNorm+modulate fusion: gated, close to eager.""" + +import sys + +import pytest +import torch +from torch import nn + +from sglang.kernels.ops.diffusion.ltx2_rmsnorm_modulate import ( + fused_ltx2_rms_norm_modulate, + mark_ltx2_rms_norm_modulate_site, + mount_ltx2_rms_norm_modulate, + unmount_ltx2_rms_norm_modulate, +) +from sglang.multimodal_gen.runtime.layers.layernorm import RMSNormNoWeight +from sglang.multimodal_gen.runtime.models.dits.ltx_2 import _ltx2_rms_norm_modulate +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci + +register_cuda_ci(est_time=8, stage="base-b-kernel-unit", runner_config="1-gpu-large") +register_amd_ci(est_time=8, suite="nightly-amd-kernel-1-gpu", nightly=True) + + +@pytest.fixture(autouse=True) +def _setup(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + torch.cuda.manual_seed(0) + + +def _eager(rms, x, scale, shift, eps): + return rms(x, eps) * (1 + scale) + shift + + +def _inputs(hidden, batch=1, seq=4096): + rms = RMSNormNoWeight() + x = torch.randn(batch, seq, hidden, device="cuda", dtype=torch.bfloat16) + scale = torch.randn(batch, 1, hidden, device="cuda", dtype=torch.bfloat16) * 0.1 + shift = torch.randn(batch, 1, hidden, device="cuda", dtype=torch.bfloat16) * 0.1 + return rms, x, scale, shift + + +# hidden 4096 = LTX-2 video stream, 2048 = audio stream. +@pytest.mark.parametrize("hidden", [4096, 2048]) +def test_lossless_default_is_bitexact(hidden): + # A marked-but-unmounted site (the lossless default) runs verbatim eager. + block = nn.Module() + mark_ltx2_rms_norm_modulate_site(block) + rms, x, scale, shift = _inputs(hidden) + out = _ltx2_rms_norm_modulate(block, rms, x, scale, shift, 1e-6) + assert torch.equal(out, _eager(rms, x, scale, shift, 1e-6)) + + +@pytest.mark.parametrize("hidden", [4096, 2048]) +def test_mounted_high_uses_fused_kernel(hidden): + block = nn.Module() + mark_ltx2_rms_norm_modulate_site(block) + assert mount_ltx2_rms_norm_modulate(block) + try: + rms, x, scale, shift = _inputs(hidden) + out = _ltx2_rms_norm_modulate(block, rms, x, scale, shift, 1e-6) + # The mounted path routes through the fused kernel exactly. + assert torch.equal(out, fused_ltx2_rms_norm_modulate(x, scale, shift, 1e-6)) + # And stays within half-precision rounding of the eager reference. + ref = _eager(rms, x, scale, shift, 1e-6) + assert torch.allclose(out.float(), ref.float(), atol=3e-2, rtol=1e-2) + finally: + unmount_ltx2_rms_norm_modulate(block) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v", "-s"]))