[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)