diff --git a/python/sglang/kernels/kda_kernels/ltx2_qknorm_split_rope_jit.py b/python/sglang/kernels/kda_kernels/ltx2_qknorm_split_rope_jit.py index abd50fb54..33a99cfeb 100644 --- a/python/sglang/kernels/kda_kernels/ltx2_qknorm_split_rope_jit.py +++ b/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( diff --git a/python/sglang/kernels/ops/diffusion/__init__.py b/python/sglang/kernels/ops/diffusion/__init__.py index 468d11bbd..8761786c7 100644 --- a/python/sglang/kernels/ops/diffusion/__init__.py +++ b/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", diff --git a/python/sglang/kernels/ops/diffusion/sites/ltx2_qknorm_split_rope_site.py b/python/sglang/kernels/ops/diffusion/sites/ltx2_qknorm_split_rope_site.py new file mode 100644 index 000000000..1f95edbbc --- /dev/null +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py index 9471d99c7..63e93cda3 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 277d59718..d85370940 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/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, diff --git a/test/registered/kernels/ops/diffusion/test_model_fast_paths.py b/test/registered/kernels/ops/diffusion/test_model_fast_paths.py index 238fde2e5..811fdc9eb 100644 --- a/test/registered/kernels/ops/diffusion/test_model_fast_paths.py +++ b/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 diff --git a/test/registered/kernels/ops/diffusion/test_rope_ltx2.py b/test/registered/kernels/ops/diffusion/test_rope_ltx2.py index e6e1e4bc7..9876a89b4 100644 --- a/test/registered/kernels/ops/diffusion/test_rope_ltx2.py +++ b/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", [