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

python/sglang/kernels/kda_kernels/ltx2_qknorm_split_rope_jit.py

@@ -130,6 +130,15 @@ def _is_sm100_or_newer(x: torch.Tensor) -> bool:
         return False
 
 
+def _is_sm90(x: torch.Tensor) -> bool:
+    if not x.is_cuda or torch.version.hip is not None:
+        return False
+    try:
+        return torch.cuda.get_device_capability(x.device) == (9, 0)
+    except RuntimeError:
+        return False
+
+
 def can_use_ltx2_qknorm_split_rope_cuda(
     q: torch.Tensor,
     q_cos: torch.Tensor,
@@ -142,9 +151,10 @@ def can_use_ltx2_qknorm_split_rope_cuda(
     *,
     num_heads: int,
     head_dim: int,
+    allow_sm90: bool = False,
 ) -> bool:
     return (
-        _is_sm100_or_newer(q)
+        (_is_sm100_or_newer(q) or (allow_sm90 and _is_sm90(q)))
         and _supported_side(
             q,
             q_cos,
@@ -177,6 +187,7 @@ def ltx2_qknorm_split_rope_cuda(
     eps: float,
     num_heads: int,
     head_dim: int,
+    allow_sm90: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     if not can_use_ltx2_qknorm_split_rope_cuda(
         q,
@@ -189,6 +200,7 @@ def ltx2_qknorm_split_rope_cuda(
         k_weight,
         num_heads=num_heads,
         head_dim=head_dim,
+        allow_sm90=allow_sm90,
     ):
         raise RuntimeError("unsupported input for LTX2 QKNorm split-RoPE CUDA")
     return _ltx2_qknorm_split_rope_custom_op(

python/sglang/kernels/ops/diffusion/__init__.py

@@ -587,6 +587,10 @@ _EXPORTS: dict[str, str] = {
     "mark_ltx2_rms_norm_modulate_site": "sites.ltx2_rmsnorm_modulate_site",
     "mount_ltx2_rms_norm_modulate": "sites.ltx2_rmsnorm_modulate_site",
     "unmount_ltx2_rms_norm_modulate": "sites.ltx2_rmsnorm_modulate_site",
+    "ltx2_qknorm_split_rope_active": "sites.ltx2_qknorm_split_rope_site",
+    "mark_ltx2_qknorm_split_rope_site": "sites.ltx2_qknorm_split_rope_site",
+    "mount_ltx2_qknorm_split_rope": "sites.ltx2_qknorm_split_rope_site",
+    "unmount_ltx2_qknorm_split_rope": "sites.ltx2_qknorm_split_rope_site",
     "lingbot_video_rmsnorm_active": "sites.lingbot_video_rmsnorm_site",
     "mark_lingbot_video_rmsnorm_site": "sites.lingbot_video_rmsnorm_site",
     "mount_lingbot_video_rmsnorm": "sites.lingbot_video_rmsnorm_site",

python/sglang/kernels/ops/diffusion/sites/ltx2_qknorm_split_rope_site.py

@@ -0,0 +1,35 @@
+"""Quality gate for the LTX-2 Q/K RMSNorm + split-RoPE Hopper path.
+
+The fused CUDA kernel is already the default on SM100+, but its fused
+rounding differs from the Hopper eager chain.  LTX-2 attention sites therefore
+enable the SM90 path only for requests whose quality policy allows approximate
+kernel fusions.
+"""
+
+from __future__ import annotations
+
+from torch import nn
+
+from sglang.kernels.ops.diffusion.sites.quality_gate import QualityGatedFusion
+
+_FUSION = QualityGatedFusion(
+    name="LTX-2 Hopper QKNorm+split-RoPE",
+    marker_attr="_sgl_ltx2_qknorm_split_rope_site",
+    enabled_attr="_sgl_ltx2_qknorm_split_rope_enabled",
+)
+
+
+def mark_ltx2_qknorm_split_rope_site(module: nn.Module) -> None:
+    _FUSION.mark(module)
+
+
+def ltx2_qknorm_split_rope_active(module: nn.Module) -> bool:
+    return _FUSION.is_enabled(module)
+
+
+def mount_ltx2_qknorm_split_rope(root: nn.Module) -> bool:
+    return _FUSION.mount(root)
+
+
+def unmount_ltx2_qknorm_split_rope(root: nn.Module) -> None:
+    _FUSION.unmount(root)

python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py

@@ -19,9 +19,11 @@ from sglang.kernels.ops.diffusion import (
     fused_gelu_active,
     fused_linear_gelu_tanh,
     fused_ltx2_rms_norm_modulate,
+    ltx2_qknorm_split_rope_active,
     ltx2_qknorm_split_rope_cuda,
     ltx2_rms_norm_modulate_active,
     mark_fused_gelu_site,
+    mark_ltx2_qknorm_split_rope_site,
     mark_ltx2_rms_norm_modulate_site,
     modulate_scale_shift_cuda,
     residual_gate_add,
@@ -83,6 +85,7 @@ def _ltx2_try_fused_qknorm_split_rope(
     eps: float,
     num_heads: int,
     head_dim: int,
+    allow_sm90: bool,
 ) -> tuple[torch.Tensor, torch.Tensor] | None:
     global _LTX2_QKNORM_SPLIT_ROPE_CUDA_DISABLED
 
@@ -104,6 +107,7 @@ def _ltx2_try_fused_qknorm_split_rope(
             k_norm.weight,
             num_heads=num_heads,
             head_dim=head_dim,
+            allow_sm90=allow_sm90,
         )
     ):
         return None
@@ -121,6 +125,7 @@ def _ltx2_try_fused_qknorm_split_rope(
             eps=eps,
             num_heads=num_heads,
             head_dim=head_dim,
+            allow_sm90=allow_sm90,
         )
     except Exception as exc:
         if torch.compiler.is_compiling():
@@ -753,6 +758,7 @@ class LTX2Attention(nn.Module):
         self.apply_gated_attention = bool(apply_gated_attention)
         self.enable_packed_qkv_input_a2a = bool(enable_packed_qkv_input_a2a)
         self.prefix = prefix
+        mark_ltx2_qknorm_split_rope_site(self)
 
         tp_size = get_tp_world_size()
         if tp_size <= 0:
@@ -910,6 +916,7 @@ class LTX2Attention(nn.Module):
                         eps=self.norm_eps,
                         num_heads=self.local_heads,
                         head_dim=self.dim_head,
+                        allow_sm90=ltx2_qknorm_split_rope_active(self),
                     )
 
             if fused_qk is not None:

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py

@@ -29,6 +29,7 @@ from sglang.kernels.ops.diffusion import (
     mount_hunyuan_qknorm,
     mount_lingbot_video_gated_residual,
     mount_lingbot_video_rmsnorm,
+    mount_ltx2_qknorm_split_rope,
     mount_ltx2_rms_norm_modulate,
     mount_nvfp4_bias_gelu,
     mount_qwen_image_added_qkv,
@@ -41,6 +42,7 @@ from sglang.kernels.ops.diffusion import (
     unmount_hunyuan_qknorm,
     unmount_lingbot_video_gated_residual,
     unmount_lingbot_video_rmsnorm,
+    unmount_ltx2_qknorm_split_rope,
     unmount_ltx2_rms_norm_modulate,
     unmount_nvfp4_bias_gelu,
     unmount_qwen_image_added_qkv,
@@ -204,6 +206,11 @@ _QUALITY_FUSION_HANDLERS: tuple[
         mount_fused_ln_modulate,
         unmount_fused_ln_modulate,
     ),
+    (
+        "LTX-2 Hopper QKNorm+split-RoPE",
+        mount_ltx2_qknorm_split_rope,
+        unmount_ltx2_qknorm_split_rope,
+    ),
     (
         "LTX-2 fused RMSNorm+modulate",
         mount_ltx2_rms_norm_modulate,

test/registered/kernels/ops/diffusion/test_model_fast_paths.py

@@ -47,14 +47,17 @@ from sglang.kernels.ops.diffusion import (
     hunyuan_qkv_rope_pack,
     mark_fused_ln_modulate_site,
     mark_hunyuan_qknorm_site,
+    mark_ltx2_qknorm_split_rope_site,
     mark_ltx2_rms_norm_modulate_site,
     mark_qwen_image_added_qkv_site,
     mount_fused_ln_modulate,
     mount_hunyuan_qknorm,
+    mount_ltx2_qknorm_split_rope,
     mount_ltx2_rms_norm_modulate,
     mount_qwen_image_added_qkv,
     try_flux2_token_cat_nvfp4,
     unmount_hunyuan_qknorm,
+    unmount_ltx2_qknorm_split_rope,
     unmount_ltx2_rms_norm_modulate,
     unmount_qwen_image_added_qkv,
     wan_rmsnorm_silu,
@@ -856,6 +859,16 @@ def test_longcat_qknorm_rope_is_bit_exact():
 # -------------------------------------------------------------------------
 
 
+def test_ltx2_qknorm_split_rope_hopper_site_is_quality_gated():
+    attention = nn.Module()
+    mark_ltx2_qknorm_split_rope_site(attention)
+    assert not ltx2_module.ltx2_qknorm_split_rope_active(attention)
+    assert mount_ltx2_qknorm_split_rope(attention)
+    assert ltx2_module.ltx2_qknorm_split_rope_active(attention)
+    unmount_ltx2_qknorm_split_rope(attention)
+    assert not ltx2_module.ltx2_qknorm_split_rope_active(attention)
+
+
 def _ltx2_eager(rms, x, scale, shift, eps):
     return rms(x, eps) * (1 + scale) + shift
 

test/registered/kernels/ops/diffusion/test_rope_ltx2.py

@@ -1,7 +1,8 @@
 """``diffusion.rope``: the LTX-2 QK-norm + split-RoPE CUDA kernel.
 
 Split out of ``test_rope.py`` rather than merged with the other RoPE kernels:
-this one is validated on B200 and registered on that lane alone, while the
+the lossless-default path is validated on B200, while the explicitly
+quality-gated SM90 path is also checked on the large-GPU lane. The
 ``fused_inplace_qknorm_rope`` cases there are held to the *split* baseline,
 whose sgl_kernel / FlashInfer dispatch differs on Blackwell -- their bit-exact
 assertions fail on B200.  One file cannot carry both lane sets.
@@ -13,6 +14,7 @@ import pytest
 import torch
 import torch.nn.functional as F
 
+import sglang.kernels.kda_kernels.ltx2_qknorm_split_rope_jit as ltx2_qknorm_jit
 from sglang.kernels.ops.diffusion import (
     can_use_ltx2_qknorm_split_rope_cuda,
     ltx2_qknorm_split_rope_cuda,
@@ -20,11 +22,27 @@ from sglang.kernels.ops.diffusion import (
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
+register_cuda_ci(est_time=15, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 
 DEVICE = "cuda"
 BF16_FUSED_ATOL = 1.6e-1
 
 
+def test_ltx2_qknorm_hopper_requires_explicit_quality_gate(monkeypatch) -> None:
+    sentinel = object()
+    monkeypatch.setattr(ltx2_qknorm_jit, "_is_sm100_or_newer", lambda _x: False)
+    monkeypatch.setattr(ltx2_qknorm_jit, "_is_sm90", lambda _x: True)
+    monkeypatch.setattr(ltx2_qknorm_jit, "_supported_side", lambda *_a, **_k: True)
+
+    args = (sentinel,) * 8
+    assert not ltx2_qknorm_jit.can_use_ltx2_qknorm_split_rope_cuda(
+        *args, num_heads=32, head_dim=128
+    )
+    assert ltx2_qknorm_jit.can_use_ltx2_qknorm_split_rope_cuda(
+        *args, num_heads=32, head_dim=128, allow_sm90=True
+    )
+
+
 def _require_b200() -> None:
     if not torch.cuda.is_available():
         pytest.skip("CUDA required")
@@ -32,6 +50,13 @@ def _require_b200() -> None:
         pytest.skip("LTX2 QKNorm split-RoPE CUDA path is validated on B200")
 
 
+def _require_sm90() -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    if torch.version.hip is not None or torch.cuda.get_device_capability() != (9, 0):
+        pytest.skip("quality-gated LTX2 Hopper path requires SM90")
+
+
 def _ltx2_make_cos_sin(
     batch: int, seq_len: int, num_heads: int, head_dim: int
 ) -> tuple[torch.Tensor, torch.Tensor]:
@@ -88,6 +113,68 @@ def _ltx2_reference(
     return q_ref.to(dtype=torch.bfloat16), k_ref.to(dtype=torch.bfloat16)
 
 
+def test_ltx2_qknorm_hopper_quality_path_matches_within_bf16() -> None:
+    _require_sm90()
+    torch.cuda.manual_seed(20260908)
+    batch, q_seq, k_seq, num_heads, head_dim = 1, 17, 9, 32, 64
+    hidden = num_heads * head_dim
+    eps = 1e-6
+    q = torch.randn(batch, q_seq, hidden, device="cuda", dtype=torch.bfloat16)
+    k = torch.randn(batch, k_seq, hidden, device="cuda", dtype=torch.bfloat16)
+    q_cos, q_sin = _ltx2_make_cos_sin(batch, q_seq, num_heads, head_dim)
+    k_cos, k_sin = _ltx2_make_cos_sin(batch, k_seq, num_heads, head_dim)
+    q_weight = torch.randn(hidden, device="cuda", dtype=torch.bfloat16)
+    k_weight = torch.randn(hidden, device="cuda", dtype=torch.bfloat16)
+
+    assert not can_use_ltx2_qknorm_split_rope_cuda(
+        q,
+        q_cos,
+        q_sin,
+        q_weight,
+        k,
+        k_cos,
+        k_sin,
+        k_weight,
+        num_heads=num_heads,
+        head_dim=head_dim,
+    )
+    assert can_use_ltx2_qknorm_split_rope_cuda(
+        q,
+        q_cos,
+        q_sin,
+        q_weight,
+        k,
+        k_cos,
+        k_sin,
+        k_weight,
+        num_heads=num_heads,
+        head_dim=head_dim,
+        allow_sm90=True,
+    )
+
+    q_norm = F.rms_norm(q, (hidden,), q_weight, eps)
+    k_norm = F.rms_norm(k, (hidden,), k_weight, eps)
+    q_ref = _apply_split_rotary_ref(q_norm, q_cos, q_sin)
+    k_ref = _apply_split_rotary_ref(k_norm, k_cos, k_sin)
+    q_out, k_out = ltx2_qknorm_split_rope_cuda(
+        q,
+        q_cos,
+        q_sin,
+        q_weight,
+        k,
+        k_cos,
+        k_sin,
+        k_weight,
+        eps=eps,
+        num_heads=num_heads,
+        head_dim=head_dim,
+        allow_sm90=True,
+    )
+    torch.cuda.synchronize()
+    torch.testing.assert_close(q_out, q_ref, rtol=0, atol=BF16_FUSED_ATOL)
+    torch.testing.assert_close(k_out, k_ref, rtol=0, atol=BF16_FUSED_ATOL)
+
+
 @pytest.mark.parametrize(
     "batch,q_seq,k_seq,num_heads,head_dim",
     [