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