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

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