diff --git a/python/sglang/jit_kernel/csrc/diffusion/ltx2_qknorm_split_rope.cuh b/python/sglang/jit_kernel/csrc/diffusion/ltx2_qknorm_split_rope.cuh new file mode 100644 index 000000000000..7940071973f1 --- /dev/null +++ b/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 // For TensorMatcher, SymbolicSize, SymbolicDevice +#include // For RuntimeCheck + +#include // For LaunchKernel and CUDA dtype aliases + +#include +#include + +namespace sglang_ltx2_qknorm_split_rope { + +namespace { + +constexpr int kThreads = 128; + +inline const char* data_ptr(const tvm::ffi::TensorView& t) { + return static_cast(t.data_ptr()) + t.byte_offset(); +} + +inline char* mutable_data_ptr(const tvm::ffi::TensorView& t) { + return static_cast(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(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(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(UINT32_MAX), "LTX2 QKNorm split-RoPE grid is too large"); + host::LaunchKernel(dim3(static_cast(num_rows)), dim3(32, 4), device)( + ltx2_qknorm_split_rope_kernel, + reinterpret_cast(data_ptr(x)), + reinterpret_cast(data_ptr(cos)), + reinterpret_cast(data_ptr(sin)), + reinterpret_cast(data_ptr(weight)), + reinterpret_cast(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(); + + TensorMatcher({batch, q_seq_len, hidden_size}).with_dtype().with_device(device).verify(q).verify(q_out); + TensorMatcher({batch, k_seq_len, hidden_size}).with_dtype().with_device(device).verify(k).verify(k_out); + TensorMatcher({hidden_size}).with_dtype().with_device(device).verify(q_weight); + TensorMatcher({hidden_size}).with_dtype().with_device(device).verify(k_weight); + TensorMatcher({batch, heads, q_seq_len, half_dim}) + .with_strides({-1, -1, -1, 1}) + .with_dtype() + .with_device(device) + .verify(q_cos); + TensorMatcher({batch, heads, q_seq_len, half_dim}) + .with_strides({-1, -1, -1, 1}) + .with_dtype() + .with_device(device) + .verify(q_sin); + TensorMatcher({batch, heads, k_seq_len, half_dim}) + .with_strides({-1, -1, -1, 1}) + .with_dtype() + .with_device(device) + .verify(k_cos); + TensorMatcher({batch, heads, k_seq_len, half_dim}) + .with_strides({-1, -1, -1, 1}) + .with_dtype() + .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(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(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 diff --git a/python/sglang/jit_kernel/diffusion/ltx2_qknorm_split_rope.py b/python/sglang/jit_kernel/diffusion/ltx2_qknorm_split_rope.py new file mode 100644 index 000000000000..45277e2d2cee --- /dev/null +++ b/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), + ) 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 e85c2da447eb..99e7bcfbb84f 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py @@ -10,6 +10,10 @@ 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 @@ 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 @@ def forward( q, _ = self.to_q(x) k, _ = self.to_k(context_) - 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) - + fused_qk = None 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 @@ def forward( 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)) diff --git a/test/registered/jit/benchmark/diffusion/bench_ltx2_qknorm_split_rope.py b/test/registered/jit/benchmark/diffusion/bench_ltx2_qknorm_split_rope.py new file mode 100644 index 000000000000..b92b1a8b2b70 --- /dev/null +++ b/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) diff --git a/test/registered/jit/diffusion/test_ltx2_qknorm_split_rope.py b/test/registered/jit/diffusion/test_ltx2_qknorm_split_rope.py new file mode 100644 index 000000000000..a52991148958 --- /dev/null +++ b/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"]))