[KDA-Pilot] Add LTX2 QKNorm split-RoPE CUDA fast path (#29708)

python/sglang/jit_kernel/csrc/diffusion/ltx2_qknorm_split_rope.cuh

@@ -0,0 +1,277 @@
+// CUDA fast path for LTX2 Q/K RMSNorm + split RoPE.
+//
+// Developed with MIT HAN Lab Kernel Design Agents:
+// https://github.com/mit-han-lab/kernel-design-agents
+//
+// This mirrors the LTX2 eager oracle:
+//   torch.nn.RMSNorm(input) returns fp32 under bf16 autocast, then split RoPE
+//   runs in fp32 and rounds once to bf16 at the final attention input.
+
+#pragma once
+
+#include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize, SymbolicDevice
+#include <sgl_kernel/utils.h>   // For RuntimeCheck
+
+#include <sgl_kernel/utils.cuh>  // For LaunchKernel and CUDA dtype aliases
+
+#include <cstdint>
+#include <cuda_bf16.h>
+
+namespace sglang_ltx2_qknorm_split_rope {
+
+namespace {
+
+constexpr int kThreads = 128;
+
+inline const char* data_ptr(const tvm::ffi::TensorView& t) {
+  return static_cast<const char*>(t.data_ptr()) + t.byte_offset();
+}
+
+inline char* mutable_data_ptr(const tvm::ffi::TensorView& t) {
+  return static_cast<char*>(t.data_ptr()) + t.byte_offset();
+}
+
+SGL_DEVICE float compute_rstd(
+    const bf16_t* __restrict__ xrow,
+    int64_t hidden_size,
+    float eps,
+    int tid,
+    int lane,
+    int warp_id,
+    float* warp_sum,
+    float* s_rstd) {
+  float local = 0.f;
+  const int64_t n_vec = hidden_size >> 2;
+  for (int64_t i = tid; i < n_vec; i += kThreads) {
+    const int64_t base = i << 2;
+    const float v0 = __bfloat162float(xrow[base + 0]);
+    const float v1 = __bfloat162float(xrow[base + 1]);
+    const float v2 = __bfloat162float(xrow[base + 2]);
+    const float v3 = __bfloat162float(xrow[base + 3]);
+    local = fmaf(v0, v0, local);
+    local = fmaf(v1, v1, local);
+    local = fmaf(v2, v2, local);
+    local = fmaf(v3, v3, local);
+  }
+
+#pragma unroll
+  for (int offset = 16; offset > 0; offset >>= 1) {
+    local += __shfl_down_sync(0xffffffffu, local, offset);
+  }
+  if (lane == 0) {
+    warp_sum[warp_id] = local;
+  }
+  __syncthreads();
+
+  if (tid == 0) {
+    const float total = (warp_sum[0] + warp_sum[2]) + (warp_sum[1] + warp_sum[3]);
+    *s_rstd = rsqrtf(total / static_cast<float>(hidden_size) + eps);
+  }
+  __syncthreads();
+  return *s_rstd;
+}
+
+SGL_DEVICE float norm_value(float x, float weight, float rstd) {
+  return weight * (rstd * x);
+}
+
+SGL_DEVICE void rope_pair(float x0, float x1, float cos, float sin, float& y0, float& y1) {
+  const float p0 = x0 * cos;
+  const float p1 = x1 * cos;
+  y0 = fmaf(-sin, x1, p0);
+  y1 = fmaf(sin, x0, p1);
+}
+
+__global__ void ltx2_qknorm_split_rope_kernel(
+    const bf16_t* __restrict__ x,
+    const bf16_t* __restrict__ cos,
+    const bf16_t* __restrict__ sin,
+    const bf16_t* __restrict__ weight,
+    bf16_t* __restrict__ out,
+    float eps,
+    int64_t seq_len,
+    int64_t num_heads,
+    int64_t head_dim,
+    int64_t stride_cos_b,
+    int64_t stride_cos_h,
+    int64_t stride_cos_t,
+    int64_t stride_sin_b,
+    int64_t stride_sin_h,
+    int64_t stride_sin_t) {
+  const int64_t row = static_cast<int64_t>(blockIdx.x);
+  const int64_t batch = row / seq_len;
+  const int64_t token = row - batch * seq_len;
+  const int64_t hidden_size = num_heads * head_dim;
+  const int64_t half_dim = head_dim >> 1;
+  const auto* __restrict__ xrow = x + row * hidden_size;
+  auto* __restrict__ outrow = out + row * hidden_size;
+  const int tid = threadIdx.x + threadIdx.y * 32;
+  const int lane = threadIdx.x;
+  const int warp_id = threadIdx.y;
+
+  __shared__ float warp_sum[4];
+  __shared__ float s_rstd;
+  const float rstd = compute_rstd(xrow, hidden_size, eps, tid, lane, warp_id, warp_sum, &s_rstd);
+
+  const int64_t num_pairs = num_heads * half_dim;
+  for (int64_t pair = tid; pair < num_pairs; pair += kThreads) {
+    const int64_t head = pair / half_dim;
+    const int64_t offset = pair - head * half_dim;
+    const int64_t idx0 = head * head_dim + offset;
+    const int64_t idx1 = idx0 + half_dim;
+    const float n0 = norm_value(__bfloat162float(xrow[idx0]), __bfloat162float(weight[idx0]), rstd);
+    const float n1 = norm_value(__bfloat162float(xrow[idx1]), __bfloat162float(weight[idx1]), rstd);
+    const int64_t cos_offset = batch * stride_cos_b + head * stride_cos_h + token * stride_cos_t + offset;
+    const int64_t sin_offset = batch * stride_sin_b + head * stride_sin_h + token * stride_sin_t + offset;
+
+    float y0;
+    float y1;
+    rope_pair(n0, n1, __bfloat162float(cos[cos_offset]), __bfloat162float(sin[sin_offset]), y0, y1);
+    outrow[idx0] = __float2bfloat16_rn(y0);
+    outrow[idx1] = __float2bfloat16_rn(y1);
+  }
+}
+
+inline void launch_one(
+    const tvm::ffi::TensorView& x,
+    const tvm::ffi::TensorView& cos,
+    const tvm::ffi::TensorView& sin,
+    const tvm::ffi::TensorView& weight,
+    const tvm::ffi::TensorView& out,
+    float eps,
+    int64_t num_rows,
+    int64_t seq_len,
+    int64_t num_heads,
+    int64_t head_dim,
+    int64_t stride_cos_b,
+    int64_t stride_cos_h,
+    int64_t stride_cos_t,
+    int64_t stride_sin_b,
+    int64_t stride_sin_h,
+    int64_t stride_sin_t,
+    DLDevice device) {
+  if (num_rows == 0) {
+    return;
+  }
+  host::RuntimeCheck(num_rows <= static_cast<int64_t>(UINT32_MAX), "LTX2 QKNorm split-RoPE grid is too large");
+  host::LaunchKernel(dim3(static_cast<uint32_t>(num_rows)), dim3(32, 4), device)(
+      ltx2_qknorm_split_rope_kernel,
+      reinterpret_cast<const bf16_t*>(data_ptr(x)),
+      reinterpret_cast<const bf16_t*>(data_ptr(cos)),
+      reinterpret_cast<const bf16_t*>(data_ptr(sin)),
+      reinterpret_cast<const bf16_t*>(data_ptr(weight)),
+      reinterpret_cast<bf16_t*>(mutable_data_ptr(out)),
+      eps,
+      seq_len,
+      num_heads,
+      head_dim,
+      stride_cos_b,
+      stride_cos_h,
+      stride_cos_t,
+      stride_sin_b,
+      stride_sin_h,
+      stride_sin_t);
+}
+
+}  // namespace
+
+struct LTX2QKNormSplitRopeKernel {
+  static void
+  run(tvm::ffi::TensorView q_out,
+      tvm::ffi::TensorView k_out,
+      tvm::ffi::TensorView q,
+      tvm::ffi::TensorView q_cos,
+      tvm::ffi::TensorView q_sin,
+      tvm::ffi::TensorView q_weight,
+      tvm::ffi::TensorView k,
+      tvm::ffi::TensorView k_cos,
+      tvm::ffi::TensorView k_sin,
+      tvm::ffi::TensorView k_weight,
+      double eps,
+      int64_t num_heads,
+      int64_t head_dim) {
+    using namespace host;
+
+    RuntimeCheck(num_heads > 0, "num_heads must be positive");
+    RuntimeCheck(head_dim > 0, "head_dim must be positive");
+    RuntimeCheck(head_dim % 2 == 0, "head_dim must be even");
+    const int64_t hidden_size = num_heads * head_dim;
+    RuntimeCheck(hidden_size % 4 == 0, "hidden size must be divisible by 4");
+
+    auto batch = SymbolicSize{"batch"};
+    auto q_seq_len = SymbolicSize{"q_seq_len"};
+    auto k_seq_len = SymbolicSize{"k_seq_len"};
+    auto heads = SymbolicSize{"num_heads"};
+    auto half_dim = SymbolicSize{"half_dim"};
+    auto device = SymbolicDevice{};
+    heads.set_value(num_heads);
+    half_dim.set_value(head_dim / 2);
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({batch, q_seq_len, hidden_size}).with_dtype<bf16_t>().with_device(device).verify(q).verify(q_out);
+    TensorMatcher({batch, k_seq_len, hidden_size}).with_dtype<bf16_t>().with_device(device).verify(k).verify(k_out);
+    TensorMatcher({hidden_size}).with_dtype<bf16_t>().with_device(device).verify(q_weight);
+    TensorMatcher({hidden_size}).with_dtype<bf16_t>().with_device(device).verify(k_weight);
+    TensorMatcher({batch, heads, q_seq_len, half_dim})
+        .with_strides({-1, -1, -1, 1})
+        .with_dtype<bf16_t>()
+        .with_device(device)
+        .verify(q_cos);
+    TensorMatcher({batch, heads, q_seq_len, half_dim})
+        .with_strides({-1, -1, -1, 1})
+        .with_dtype<bf16_t>()
+        .with_device(device)
+        .verify(q_sin);
+    TensorMatcher({batch, heads, k_seq_len, half_dim})
+        .with_strides({-1, -1, -1, 1})
+        .with_dtype<bf16_t>()
+        .with_device(device)
+        .verify(k_cos);
+    TensorMatcher({batch, heads, k_seq_len, half_dim})
+        .with_strides({-1, -1, -1, 1})
+        .with_dtype<bf16_t>()
+        .with_device(device)
+        .verify(k_sin);
+
+    const int64_t batch_size = batch.unwrap();
+    const DLDevice dl_device = device.unwrap();
+    launch_one(
+        q,
+        q_cos,
+        q_sin,
+        q_weight,
+        q_out,
+        static_cast<float>(eps),
+        batch_size * q_seq_len.unwrap(),
+        q_seq_len.unwrap(),
+        num_heads,
+        head_dim,
+        q_cos.stride(0),
+        q_cos.stride(1),
+        q_cos.stride(2),
+        q_sin.stride(0),
+        q_sin.stride(1),
+        q_sin.stride(2),
+        dl_device);
+    launch_one(
+        k,
+        k_cos,
+        k_sin,
+        k_weight,
+        k_out,
+        static_cast<float>(eps),
+        batch_size * k_seq_len.unwrap(),
+        k_seq_len.unwrap(),
+        num_heads,
+        head_dim,
+        k_cos.stride(0),
+        k_cos.stride(1),
+        k_cos.stride(2),
+        k_sin.stride(0),
+        k_sin.stride(1),
+        k_sin.stride(2),
+        dl_device);
+  }
+};
+
+}  // namespace sglang_ltx2_qknorm_split_rope

python/sglang/jit_kernel/diffusion/ltx2_qknorm_split_rope.py

@@ -0,0 +1,205 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.jit_kernel.utils import cache_once, load_jit
+from sglang.srt.utils.custom_op import register_custom_op
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_ltx2_qknorm_split_rope_module() -> Module:
+    return load_jit(
+        "diffusion_ltx2_qknorm_split_rope",
+        cuda_files=["diffusion/ltx2_qknorm_split_rope.cuh"],
+        cuda_wrappers=[
+            (
+                "ltx2_qknorm_split_rope_pair",
+                "sglang_ltx2_qknorm_split_rope::LTX2QKNormSplitRopeKernel::run",
+            )
+        ],
+    )
+
+
+def _fake_impl(
+    q: torch.Tensor,
+    q_cos: torch.Tensor,
+    q_sin: torch.Tensor,
+    q_weight: torch.Tensor,
+    k: torch.Tensor,
+    k_cos: torch.Tensor,
+    k_sin: torch.Tensor,
+    k_weight: torch.Tensor,
+    eps: float,
+    num_heads: int,
+    head_dim: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    return torch.empty_like(q, dtype=torch.bfloat16), torch.empty_like(
+        k, dtype=torch.bfloat16
+    )
+
+
+@register_custom_op(
+    op_name="diffusion_ltx2_qknorm_split_rope",
+    mutates_args=[],
+    fake_impl=_fake_impl,
+)
+def _ltx2_qknorm_split_rope_custom_op(
+    q: torch.Tensor,
+    q_cos: torch.Tensor,
+    q_sin: torch.Tensor,
+    q_weight: torch.Tensor,
+    k: torch.Tensor,
+    k_cos: torch.Tensor,
+    k_sin: torch.Tensor,
+    k_weight: torch.Tensor,
+    eps: float,
+    num_heads: int,
+    head_dim: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    q_out = torch.empty_like(q, dtype=torch.bfloat16)
+    k_out = torch.empty_like(k, dtype=torch.bfloat16)
+    module = _jit_ltx2_qknorm_split_rope_module()
+    module.ltx2_qknorm_split_rope_pair(
+        q_out,
+        k_out,
+        q,
+        q_cos,
+        q_sin,
+        q_weight,
+        k,
+        k_cos,
+        k_sin,
+        k_weight,
+        float(eps),
+        int(num_heads),
+        int(head_dim),
+    )
+    return q_out, k_out
+
+
+def _supported_side(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    weight: torch.Tensor,
+    *,
+    num_heads: int,
+    head_dim: int,
+) -> bool:
+    return (
+        x.is_cuda
+        and cos.is_cuda
+        and sin.is_cuda
+        and weight.is_cuda
+        and x.device == cos.device == sin.device == weight.device
+        and x.dtype == torch.bfloat16
+        and cos.dtype == torch.bfloat16
+        and sin.dtype == torch.bfloat16
+        and weight.dtype == torch.bfloat16
+        and x.ndim == 3
+        and cos.ndim == 4
+        and sin.ndim == 4
+        and x.is_contiguous()
+        and cos.shape == sin.shape
+        and cos.shape[0] == x.shape[0]
+        and cos.shape[1] == num_heads
+        and cos.shape[2] == x.shape[1]
+        and cos.shape[3] * 2 == head_dim
+        and x.shape[2] == num_heads * head_dim
+        and x.shape[2] == weight.shape[0]
+        and weight.ndim == 1
+        and head_dim % 2 == 0
+        and x.shape[2] % 4 == 0
+        and cos.stride(-1) == 1
+        and sin.stride(-1) == 1
+    )
+
+
+def _is_sm100_or_newer(x: torch.Tensor) -> bool:
+    if not x.is_cuda:
+        return False
+    try:
+        return torch.cuda.get_device_capability(x.device)[0] >= 10
+    except RuntimeError:
+        return False
+
+
+def can_use_ltx2_qknorm_split_rope_cuda(
+    q: torch.Tensor,
+    q_cos: torch.Tensor,
+    q_sin: torch.Tensor,
+    q_weight: torch.Tensor,
+    k: torch.Tensor,
+    k_cos: torch.Tensor,
+    k_sin: torch.Tensor,
+    k_weight: torch.Tensor,
+    *,
+    num_heads: int,
+    head_dim: int,
+) -> bool:
+    return (
+        _is_sm100_or_newer(q)
+        and _supported_side(
+            q,
+            q_cos,
+            q_sin,
+            q_weight,
+            num_heads=num_heads,
+            head_dim=head_dim,
+        )
+        and _supported_side(
+            k,
+            k_cos,
+            k_sin,
+            k_weight,
+            num_heads=num_heads,
+            head_dim=head_dim,
+        )
+    )
+
+
+def ltx2_qknorm_split_rope_cuda(
+    q: torch.Tensor,
+    q_cos: torch.Tensor,
+    q_sin: torch.Tensor,
+    q_weight: torch.Tensor,
+    k: torch.Tensor,
+    k_cos: torch.Tensor,
+    k_sin: torch.Tensor,
+    k_weight: torch.Tensor,
+    *,
+    eps: float,
+    num_heads: int,
+    head_dim: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if 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,
+    ):
+        raise RuntimeError("unsupported input for LTX2 QKNorm split-RoPE CUDA")
+    return _ltx2_qknorm_split_rope_custom_op(
+        q,
+        q_cos,
+        q_sin,
+        q_weight,
+        k,
+        k_cos,
+        k_sin,
+        k_weight,
+        float(eps),
+        int(num_heads),
+        int(head_dim),
+    )

python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py

@@ -10,6 +10,10 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+from sglang.jit_kernel.diffusion.ltx2_qknorm_split_rope import (
+    can_use_ltx2_qknorm_split_rope_cuda,
+    ltx2_qknorm_split_rope_cuda,
+)
 from sglang.jit_kernel.diffusion.residual_gate_add import (
     can_use_residual_gate_add_cuda,
     residual_gate_add_cuda,
@@ -53,6 +57,7 @@ logger = init_logger(__name__)
 ADALN_NUM_BASE_PARAMS = 6
 ADALN_NUM_CROSS_ATTN_PARAMS = 3
 _LTX2_RESIDUAL_GATE_CUDA_DISABLED = False
+_LTX2_QKNORM_SPLIT_ROPE_CUDA_DISABLED = False
 
 
 def _ltx2_residual_gate_add(
@@ -76,6 +81,66 @@ def _ltx2_residual_gate_add(
     return residual + update * gate
 
 
+def _ltx2_try_fused_qknorm_split_rope(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_norm: nn.Module,
+    k_norm: nn.Module,
+    q_cos: torch.Tensor,
+    q_sin: torch.Tensor,
+    k_cos: torch.Tensor,
+    k_sin: torch.Tensor,
+    *,
+    eps: float,
+    num_heads: int,
+    head_dim: int,
+) -> tuple[torch.Tensor, torch.Tensor] | None:
+    global _LTX2_QKNORM_SPLIT_ROPE_CUDA_DISABLED
+
+    if (
+        _LTX2_QKNORM_SPLIT_ROPE_CUDA_DISABLED
+        or get_tp_world_size() != 1
+        or not isinstance(q_norm, nn.RMSNorm)
+        or not isinstance(k_norm, nn.RMSNorm)
+        or float(q_norm.eps) != float(eps)
+        or float(k_norm.eps) != float(eps)
+        or not can_use_ltx2_qknorm_split_rope_cuda(
+            q,
+            q_cos,
+            q_sin,
+            q_norm.weight,
+            k,
+            k_cos,
+            k_sin,
+            k_norm.weight,
+            num_heads=num_heads,
+            head_dim=head_dim,
+        )
+    ):
+        return None
+
+    try:
+        return ltx2_qknorm_split_rope_cuda(
+            q,
+            q_cos,
+            q_sin,
+            q_norm.weight,
+            k,
+            k_cos,
+            k_sin,
+            k_norm.weight,
+            eps=eps,
+            num_heads=num_heads,
+            head_dim=head_dim,
+        )
+    except Exception as exc:
+        if torch.compiler.is_compiling():
+            raise
+        logger.warning_once(f"Disabling LTX2 QKNorm split-RoPE CUDA fast path: {exc}")
+        _LTX2_QKNORM_SPLIT_ROPE_CUDA_DISABLED = True
+        return None
+
+
 _LTX2_FUSED_ADA_VALUES_RUNTIME_DISABLED = False
 
 
@@ -748,11 +813,7 @@ class LTX2Attention(nn.Module):
             q, _ = self.to_q(x)
             k, _ = self.to_k(context_)
 
-            if self.qk_norm:
+            fused_qk = None
-                assert self.q_norm is not None and self.k_norm is not None
-                q = self.q_norm(q)
-                k = self.k_norm(k)
-
             if pe is not None:
                 cos, sin = pe
                 k_cos, k_sin = pe if k_pe is None else k_pe
@@ -765,10 +826,34 @@ class LTX2Attention(nn.Module):
                     k_cos, k_sin = self._slice_rope_for_tp(
                         k_cos, k_sin, tp_rank=tp_rank, tp_size=tp_size
                     )
-                if cos.dim() == 3:
+                if self.qk_norm and cos.dim() != 3:
+                    assert self.q_norm is not None and self.k_norm is not None
+                    fused_qk = _ltx2_try_fused_qknorm_split_rope(
+                        q,
+                        k,
+                        self.q_norm,
+                        self.k_norm,
+                        cos,
+                        sin,
+                        k_cos,
+                        k_sin,
+                        eps=self.norm_eps,
+                        num_heads=self.local_heads,
+                        head_dim=self.dim_head,
+                    )
+
+            if fused_qk is not None:
+                q, k = fused_qk
+            else:
+                if self.qk_norm:
+                    assert self.q_norm is not None and self.k_norm is not None
+                    q = self.q_norm(q)
+                    k = self.k_norm(k)
+
+                if pe is not None and cos.dim() == 3:
                     q = apply_interleaved_rotary_emb(q, (cos, sin))
                     k = apply_interleaved_rotary_emb(k, (k_cos, k_sin))
-                else:
+                elif pe is not None:
                     q = apply_split_rotary_emb(q, (cos, sin))
                     k = apply_split_rotary_emb(k, (k_cos, k_sin))
 

test/registered/jit/benchmark/diffusion/bench_ltx2_qknorm_split_rope.py

@@ -0,0 +1,209 @@
+import random
+import sys
+from dataclasses import dataclass
+
+import torch
+
+from sglang.jit_kernel.diffusion.ltx2_qknorm_split_rope import (
+    ltx2_qknorm_split_rope_cuda,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.utils import is_in_ci
+
+register_cuda_ci(
+    est_time=30,
+    stage="base-b-kernel-benchmark",
+    runner_config="1-gpu-large",
+    disabled="standalone benchmark",
+)
+
+
+@dataclass(frozen=True)
+class Workload:
+    name: str
+    batch: int
+    q_seq: int
+    k_seq: int
+    num_heads: int
+    head_dim: int
+
+
+FULL_WORKLOADS = [
+    Workload("stage1_video_self_q1536_k1536_d4096", 2, 1536, 1536, 32, 128),
+    Workload("stage1_audio_self_q126_k126_d2048", 2, 126, 126, 32, 64),
+    Workload("stage1_audio_to_video_q1536_k126_d2048", 2, 1536, 126, 32, 64),
+    Workload("stage1_video_to_audio_q126_k1536_d2048", 2, 126, 1536, 32, 64),
+    Workload("stage2_video_self_q6144_k6144_d4096", 1, 6144, 6144, 32, 128),
+    Workload("stage2_audio_self_q126_k126_d2048", 1, 126, 126, 32, 64),
+    Workload("stage2_audio_to_video_q6144_k126_d2048", 1, 6144, 126, 32, 64),
+    Workload("stage2_video_to_audio_q126_k6144_d2048", 1, 126, 6144, 32, 64),
+    Workload("hq_stage1_video_self_q8160_k8160_d4096", 1, 8160, 8160, 32, 128),
+    Workload("hq_stage1_audio_to_video_q8160_k126_d2048", 1, 8160, 126, 32, 64),
+    Workload("hq_stage1_video_to_audio_q126_k8160_d2048", 1, 126, 8160, 32, 64),
+    Workload("hq_stage2_video_self_q32640_k32640_d4096", 1, 32640, 32640, 32, 128),
+    Workload("hq_stage2_audio_to_video_q32640_k126_d2048", 1, 32640, 126, 32, 64),
+    Workload("hq_stage2_video_to_audio_q126_k32640_d2048", 1, 126, 32640, 32, 64),
+]
+CI_WORKLOADS = [
+    Workload("stage1_video_self_q16_k16_d4096", 1, 16, 16, 32, 128),
+    Workload("stage1_audio_to_video_q16_k8_d2048", 1, 16, 8, 32, 64),
+]
+
+
+def _make_cos_sin(
+    batch: int, seq_len: int, num_heads: int, head_dim: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    half_dim = head_dim // 2
+    cos = torch.randn(
+        batch, seq_len, num_heads, half_dim, device="cuda", dtype=torch.bfloat16
+    ).transpose(1, 2)
+    sin = torch.randn(
+        batch, seq_len, num_heads, half_dim, device="cuda", dtype=torch.bfloat16
+    ).transpose(1, 2)
+    return cos, sin
+
+
+def _apply_split_rotary_ref(
+    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> torch.Tensor:
+    x_dtype = x.dtype
+    batch = x.shape[0]
+    _, num_heads, seq_len, _ = cos.shape
+    x = x.reshape(batch, seq_len, num_heads, -1).swapaxes(1, 2)
+    last = x.shape[-1]
+    half = last // 2
+    split_x = x.reshape(*x.shape[:-1], 2, half)
+    first_x = split_x[..., :1, :]
+    second_x = split_x[..., 1:, :]
+    cos_u = cos.unsqueeze(-2)
+    sin_u = sin.unsqueeze(-2)
+    out = split_x * cos_u
+    out[..., :1, :].addcmul_(-sin_u, second_x)
+    out[..., 1:, :].addcmul_(sin_u, first_x)
+    out = out.reshape(*out.shape[:-2], last)
+    return out.swapaxes(1, 2).reshape(batch, seq_len, -1).to(dtype=x_dtype)
+
+
+def _reference_pair(inputs):
+    q, k, q_cos, q_sin, k_cos, k_sin, q_norm, k_norm = inputs
+    with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
+        q_out = _apply_split_rotary_ref(q_norm(q), q_cos, q_sin)
+        k_out = _apply_split_rotary_ref(k_norm(k), k_cos, k_sin)
+    return q_out.to(dtype=torch.bfloat16), k_out.to(dtype=torch.bfloat16)
+
+
+def cuda_event_us(fn, warmups: int, repeats: int, rounds: int) -> float:
+    for _ in range(warmups):
+        fn()
+    torch.cuda.synchronize()
+
+    samples = []
+    for _ in range(rounds):
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        for _ in range(repeats):
+            fn()
+        end.record()
+        end.synchronize()
+        samples.append(start.elapsed_time(end) * 1000.0 / repeats)
+    samples.sort()
+    return samples[len(samples) // 2]
+
+
+def benchmark() -> None:
+    if not torch.cuda.is_available():
+        print("CUDA required")
+        return
+
+    torch.manual_seed(20260630)
+    random.seed(20260630)
+    torch.cuda.set_device(0)
+
+    workloads = CI_WORKLOADS if is_in_ci() else FULL_WORKLOADS
+    warmups = 3 if is_in_ci() else 10
+    repeats = 3 if is_in_ci() else 10
+    rounds = 3 if is_in_ci() else 7
+
+    print("| workload | torch us | cuda us | speedup |")
+    print("|---|---:|---:|---:|")
+
+    for workload in workloads:
+        hidden = workload.num_heads * workload.head_dim
+        q = torch.randn(
+            workload.batch,
+            workload.q_seq,
+            hidden,
+            device="cuda",
+            dtype=torch.bfloat16,
+        )
+        k = torch.randn(
+            workload.batch,
+            workload.k_seq,
+            hidden,
+            device="cuda",
+            dtype=torch.bfloat16,
+        )
+        q_cos, q_sin = _make_cos_sin(
+            workload.batch, workload.q_seq, workload.num_heads, workload.head_dim
+        )
+        k_cos, k_sin = _make_cos_sin(
+            workload.batch, workload.k_seq, workload.num_heads, workload.head_dim
+        )
+        q_norm = torch.nn.RMSNorm(hidden, eps=1e-6, device="cuda").to(
+            dtype=torch.bfloat16
+        )
+        k_norm = torch.nn.RMSNorm(hidden, eps=1e-6, device="cuda").to(
+            dtype=torch.bfloat16
+        )
+        inputs = (q, k, q_cos, q_sin, k_cos, k_sin, q_norm, k_norm)
+
+        q_ref, k_ref = _reference_pair(inputs)
+        q_out, k_out = ltx2_qknorm_split_rope_cuda(
+            q,
+            q_cos,
+            q_sin,
+            q_norm.weight,
+            k,
+            k_cos,
+            k_sin,
+            k_norm.weight,
+            eps=1e-6,
+            num_heads=workload.num_heads,
+            head_dim=workload.head_dim,
+        )
+        torch.cuda.synchronize()
+        assert torch.equal(q_ref, q_out)
+        assert torch.equal(k_ref, k_out)
+
+        fns = {
+            "torch": lambda: _reference_pair(inputs),
+            "cuda": lambda: ltx2_qknorm_split_rope_cuda(
+                q,
+                q_cos,
+                q_sin,
+                q_norm.weight,
+                k,
+                k_cos,
+                k_sin,
+                k_norm.weight,
+                eps=1e-6,
+                num_heads=workload.num_heads,
+                head_dim=workload.head_dim,
+            ),
+        }
+        order = ["torch", "cuda"]
+        random.shuffle(order)
+        times = {
+            name: cuda_event_us(fns[name], warmups, repeats, rounds) for name in order
+        }
+        print(
+            f"| {workload.name} | {times['torch']:.2f} | "
+            f"{times['cuda']:.2f} | {times['torch'] / times['cuda']:.3f}x |"
+        )
+        torch.cuda.empty_cache()
+
+
+if __name__ == "__main__":
+    benchmark()
+    sys.exit(0)

test/registered/jit/diffusion/test_ltx2_qknorm_split_rope.py

@@ -0,0 +1,223 @@
+import sys
+
+import pytest
+import torch
+
+from sglang.jit_kernel.diffusion.ltx2_qknorm_split_rope import (
+    can_use_ltx2_qknorm_split_rope_cuda,
+    ltx2_qknorm_split_rope_cuda,
+)
+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")
+
+
+def _require_cuda_b200() -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    if torch.cuda.get_device_capability()[0] < 10:
+        pytest.skip("LTX2 QKNorm split-RoPE CUDA path is validated on B200")
+
+
+@pytest.fixture(autouse=True)
+def cuda_setup():
+    _require_cuda_b200()
+    torch.cuda.manual_seed(20260630)
+
+
+def _make_cos_sin(
+    batch: int, seq_len: int, num_heads: int, head_dim: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    half_dim = head_dim // 2
+    cos = torch.randn(
+        batch, seq_len, num_heads, half_dim, device="cuda", dtype=torch.bfloat16
+    ).transpose(1, 2)
+    sin = torch.randn(
+        batch, seq_len, num_heads, half_dim, device="cuda", dtype=torch.bfloat16
+    ).transpose(1, 2)
+    return cos, sin
+
+
+def _apply_split_rotary_ref(
+    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> torch.Tensor:
+    x_dtype = x.dtype
+    batch = x.shape[0]
+    _, num_heads, seq_len, _ = cos.shape
+    x = x.reshape(batch, seq_len, num_heads, -1).swapaxes(1, 2)
+    last = x.shape[-1]
+    half = last // 2
+
+    split_x = x.reshape(*x.shape[:-1], 2, half)
+    first_x = split_x[..., :1, :]
+    second_x = split_x[..., 1:, :]
+    cos_u = cos.unsqueeze(-2)
+    sin_u = sin.unsqueeze(-2)
+
+    out = split_x * cos_u
+    out[..., :1, :].addcmul_(-sin_u, second_x)
+    out[..., 1:, :].addcmul_(sin_u, first_x)
+    out = out.reshape(*out.shape[:-2], last)
+    return out.swapaxes(1, 2).reshape(batch, seq_len, -1).to(dtype=x_dtype)
+
+
+def _reference(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_cos: torch.Tensor,
+    q_sin: torch.Tensor,
+    k_cos: torch.Tensor,
+    k_sin: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    eps: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    q_norm = torch.nn.RMSNorm(q.shape[-1], eps=eps, device="cuda").to(
+        dtype=torch.bfloat16
+    )
+    k_norm = torch.nn.RMSNorm(k.shape[-1], eps=eps, device="cuda").to(
+        dtype=torch.bfloat16
+    )
+    q_norm.weight.data.copy_(q_weight)
+    k_norm.weight.data.copy_(k_weight)
+    with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
+        q_ref = _apply_split_rotary_ref(q_norm(q), q_cos, q_sin)
+        k_ref = _apply_split_rotary_ref(k_norm(k), k_cos, k_sin)
+    return q_ref.to(dtype=torch.bfloat16), k_ref.to(dtype=torch.bfloat16)
+
+
+@pytest.mark.parametrize(
+    "batch,q_seq,k_seq,num_heads,head_dim",
+    [
+        (1, 3, 3, 32, 128),
+        (1, 5, 2, 32, 64),
+        (2, 4, 3, 32, 64),
+    ],
+)
+def test_ltx2_qknorm_split_rope_matches_torch_exactly(
+    batch: int, q_seq: int, k_seq: int, num_heads: int, head_dim: int
+) -> None:
+    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 = _make_cos_sin(batch, q_seq, num_heads, head_dim)
+    k_cos, k_sin = _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 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,
+    )
+
+    q_ref, k_ref = _reference(q, k, q_cos, q_sin, k_cos, k_sin, q_weight, k_weight, eps)
+    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,
+    )
+    torch.cuda.synchronize()
+
+    assert torch.equal(q_ref, q_out)
+    assert torch.equal(k_ref, k_out)
+
+
+def test_ltx2_qknorm_split_rope_rejects_unsupported_inputs() -> None:
+    q = torch.randn((1, 3, 4096), device="cuda", dtype=torch.bfloat16)
+    k = torch.randn_like(q)
+    q_cos, q_sin = _make_cos_sin(1, 3, 32, 128)
+    q_weight = torch.randn(4096, device="cuda", dtype=torch.bfloat16)
+    k_weight = torch.randn(4096, device="cuda", dtype=torch.bfloat16)
+
+    assert can_use_ltx2_qknorm_split_rope_cuda(
+        q,
+        q_cos,
+        q_sin,
+        q_weight,
+        k,
+        q_cos,
+        q_sin,
+        k_weight,
+        num_heads=32,
+        head_dim=128,
+    )
+    assert not can_use_ltx2_qknorm_split_rope_cuda(
+        q.float(),
+        q_cos,
+        q_sin,
+        q_weight,
+        k,
+        q_cos,
+        q_sin,
+        k_weight,
+        num_heads=32,
+        head_dim=128,
+    )
+    assert not can_use_ltx2_qknorm_split_rope_cuda(
+        q,
+        q_cos,
+        q_sin,
+        q_weight,
+        k,
+        q_cos.transpose(-1, -2),
+        q_sin,
+        k_weight,
+        num_heads=32,
+        head_dim=128,
+    )
+
+
+def test_ltx2_qknorm_split_rope_custom_op_torch_compile_fullgraph() -> None:
+    batch, q_seq, k_seq, num_heads, head_dim = 1, 3, 2, 32, 64
+    hidden = num_heads * head_dim
+    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 = _make_cos_sin(batch, q_seq, num_heads, head_dim)
+    k_cos, k_sin = _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)
+
+    def fn(q, k, q_cos, q_sin, k_cos, k_sin, q_weight, k_weight):
+        return ltx2_qknorm_split_rope_cuda(
+            q,
+            q_cos,
+            q_sin,
+            q_weight,
+            k,
+            k_cos,
+            k_sin,
+            k_weight,
+            eps=1e-6,
+            num_heads=num_heads,
+            head_dim=head_dim,
+        )
+
+    compiled = torch.compile(fn, fullgraph=True)
+    q_out, k_out = compiled(q, k, q_cos, q_sin, k_cos, k_sin, q_weight, k_weight)
+    q_ref, k_ref = _reference(
+        q, k, q_cos, q_sin, k_cos, k_sin, q_weight, k_weight, 1e-6
+    )
+    torch.cuda.synchronize()
+    assert torch.equal(q_ref, q_out)
+    assert torch.equal(k_ref, k_out)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v", "-s"]))