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

2026-09-08 19:05:02 +08:00
parent 88a9bfd1ff
commit 554f817948
7 changed files with 167 additions and 2 deletions
@@ -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(
@@ -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",
@@ -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)
@@ -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:
@@ -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,
@@ -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

@@ -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",
    [