diff --git a/docs/docs/sglang-diffusion/fused_kernels.mdx b/docs/docs/sglang-diffusion/fused_kernels.mdx index 592e5a125..3f1209e50 100644 --- a/docs/docs/sglang-diffusion/fused_kernels.mdx +++ b/docs/docs/sglang-diffusion/fused_kernels.mdx @@ -98,6 +98,7 @@ These fusion families mount under `quality="high"`: | `fused_inplace_qknorm_rope` | JIT CUDA | one bf16 rounding step vs the split baseline; exact with `round_norm_before_rope=True` | separate QK-norm kernel + RoPE | | `rope_rotate_half` | Triton | bit-exact | `chunk` → `cat(-x2, x1)` → two muls + add → `cat(tail)`, about 7 kernels per projection | | `ltx2_qknorm_split_rope` | JIT CUDA | close (validated on B200) | LTX-2 QK-norm + split RoPE | +| `ltx25_decoder_rope` | JIT CUDA | bit-exact | paired LTX-2.5 decoder 3D RoPE from cached compact axis tables | | `hunyuan_qkv_rope_pack` | Triton | bit-exact | QKV pack and RoPE in one pass | ### Activation @@ -142,6 +143,7 @@ Kernels are written against a specific eager chain in a specific model, so cover | Z-Image | BF16-native RMSNorm scale / tanh-residual, per-head QK RMSNorm | | Ideogram 4 | gate RMSNorm, SwiGLU, rotate-half RoPE, modulate, residual-gate add | | LTX-2 | QK-norm + split RoPE, ada-values split, RMSNorm+modulate, modulate, residual-gate add, linear+GELU | +| LTX-2.5 decoder | paired 3D RoPE with shared axis-table cache | | HunyuanVideo | QKV+RoPE pack, strided QK RMSNorm, linear+GELU | | Sana | LN+modulate, GLUMB bias+SiLU / bias+GLU, residual-gate add | | Sana-WM | bidirectional gated delta-net, fused QK inverse-RMS | diff --git a/python/sglang/kernels/jit/csrc/diffusion/ltx25_decoder_rope.cuh b/python/sglang/kernels/jit/csrc/diffusion/ltx25_decoder_rope.cuh new file mode 100644 index 000000000..06e1662db --- /dev/null +++ b/python/sglang/kernels/jit/csrc/diffusion/ltx25_decoder_rope.cuh @@ -0,0 +1,198 @@ +#pragma once + +#include +#include + +#include +#include +#include + +#include +#include +#include + +namespace sglang { + +namespace ltx25_decoder_rope { + +namespace { + +constexpr uint32_t kBlockSize = 256; +constexpr uint32_t kMaxGrid = 65535; + +template +SGL_DEVICE device::AlignedVector rotate_pair(device::AlignedVector input, float cos, float sin) { + static_assert(std::is_same_v); + const float even = static_cast(input[0]); + const float odd = static_cast(input[1]); + const float even_cos = __fmul_rn(even, cos); + const float odd_sin = __fmul_rn(odd, sin); + const float even_sin = __fmul_rn(even, sin); + const float odd_cos = __fmul_rn(odd, cos); + device::AlignedVector output; + output[0] = static_cast(__fsub_rn(even_cos, odd_sin)); + output[1] = static_cast(__fadd_rn(even_sin, odd_cos)); + return output; +} + +template +__global__ void ltx25_decoder_rope_kernel( + T* __restrict__ q_out, + T* __restrict__ k_out, + const T* __restrict__ q, + const T* __restrict__ k, + const float* __restrict__ cos_t, + const float* __restrict__ sin_t, + const float* __restrict__ cos_h, + const float* __restrict__ sin_h, + const float* __restrict__ cos_w, + const float* __restrict__ sin_w, + int64_t num_pairs, + int64_t num_frames, + int64_t height, + int64_t width, + int64_t num_heads, + int64_t pairs_per_head, + int64_t t_pairs, + int64_t h_pairs) { + using Pair = device::AlignedVector; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t pair_index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; pair_index < num_pairs; + pair_index += stride) { + const int64_t pair_in_head = pair_index % pairs_per_head; + const int64_t row = pair_index / (num_heads * pairs_per_head); + const int64_t token = row % (num_frames * height * width); + const int64_t frame = token / (height * width); + const int64_t spatial = token % (height * width); + const int64_t y = spatial / width; + const int64_t x = spatial % width; + + const float* cos_table; + const float* sin_table; + int64_t table_index; + if (pair_in_head < t_pairs) { + cos_table = cos_t; + sin_table = sin_t; + table_index = frame * t_pairs + pair_in_head; + } else if (pair_in_head < t_pairs + h_pairs) { + cos_table = cos_h; + sin_table = sin_h; + table_index = y * h_pairs + pair_in_head - t_pairs; + } else { + const int64_t w_pairs = pairs_per_head - t_pairs - h_pairs; + cos_table = cos_w; + sin_table = sin_w; + table_index = x * w_pairs + pair_in_head - t_pairs - h_pairs; + } + const float cos_value = SGLANG_LDG(cos_table + table_index); + const float sin_value = SGLANG_LDG(sin_table + table_index); + + Pair q_pair; + q_pair.load(q, pair_index); + rotate_pair(q_pair, cos_value, sin_value).store(q_out, pair_index); + + Pair k_pair; + k_pair.load(k, pair_index); + rotate_pair(k_pair, cos_value, sin_value).store(k_out, pair_index); + } +} + +} // namespace + +/** + * \brief Apply LTX-2.5 decoder 3D RoPE to Q and K from compact axis tables. + * + * The separate fp32 multiply and add/subtract roundings mirror the eager + * PyTorch expression used by the decoder. + * + * \tparam T Activation type; currently bf16 only. + */ +template +struct LTX25DecoderRopeKernel { + static_assert(std::is_same_v); + + static void + run(tvm::ffi::TensorView q_out, + tvm::ffi::TensorView k_out, + tvm::ffi::TensorView q, + tvm::ffi::TensorView k, + tvm::ffi::TensorView cos_t, + tvm::ffi::TensorView sin_t, + tvm::ffi::TensorView cos_h, + tvm::ffi::TensorView sin_h, + tvm::ffi::TensorView cos_w, + tvm::ffi::TensorView sin_w, + int64_t batch_size, + int64_t num_frames, + int64_t height, + int64_t width, + int64_t num_heads, + int64_t head_dim, + int64_t dim_t, + int64_t dim_h) { + using namespace host; + using Pair = device::AlignedVector; + + auto N = SymbolicSize{"activation_elements"}; + auto RT = SymbolicSize{"temporal_table_elements"}; + auto RH = SymbolicSize{"height_table_elements"}; + auto RW = SymbolicSize{"width_table_elements"}; + auto device = SymbolicDevice{}; + device.set_options(); + TensorMatcher({N}).with_dtype().with_device(device).verify(q_out).verify(k_out).verify(q).verify(k); + TensorMatcher({RT}).with_dtype().with_device(device).verify(cos_t).verify(sin_t); + TensorMatcher({RH}).with_dtype().with_device(device).verify(cos_h).verify(sin_h); + TensorMatcher({RW}).with_dtype().with_device(device).verify(cos_w).verify(sin_w); + + CHECK_HOST(batch_size > 0 && num_frames > 0 && height > 0 && width > 0 && num_heads > 0 && head_dim > 0) + << "ltx25_decoder_rope dimensions must be positive"; + CHECK_HOST(head_dim % 2 == 0 && dim_t > 0 && dim_h > 0 && dim_t % 2 == 0 && dim_h % 2 == 0) + << "ltx25_decoder_rope dimensions must split into positive rotation pairs"; + const int64_t dim_w = head_dim - dim_t - dim_h; + CHECK_HOST(dim_w > 0 && dim_w % 2 == 0) << "ltx25_decoder_rope width dimension must be positive and even"; + CHECK_HOST(N.unwrap() == batch_size * num_frames * height * width * num_heads * head_dim) + << "ltx25_decoder_rope activation shape does not match dimensions"; + CHECK_HOST(RT.unwrap() == num_frames * dim_t / 2) << "ltx25_decoder_rope temporal table shape mismatch"; + CHECK_HOST(RH.unwrap() == height * dim_h / 2) << "ltx25_decoder_rope height table shape mismatch"; + CHECK_HOST(RW.unwrap() == width * dim_w / 2) << "ltx25_decoder_rope width table shape mismatch"; + CHECK_HOST( + q_out.data_ptr() != k_out.data_ptr() && q_out.data_ptr() != q.data_ptr() && q_out.data_ptr() != k.data_ptr() && + k_out.data_ptr() != q.data_ptr() && k_out.data_ptr() != k.data_ptr()) + << "ltx25_decoder_rope outputs must not alias inputs"; + CHECK_HOST( + reinterpret_cast(q_out.data_ptr()) % alignof(Pair) == 0 && + reinterpret_cast(k_out.data_ptr()) % alignof(Pair) == 0 && + reinterpret_cast(q.data_ptr()) % alignof(Pair) == 0 && + reinterpret_cast(k.data_ptr()) % alignof(Pair) == 0) + << "ltx25_decoder_rope activations must be aligned to rotation pairs"; + + const int64_t num_pairs = N.unwrap() / 2; + const int64_t pairs_per_head = head_dim / 2; + const auto blocks = + static_cast(std::min(div_ceil(num_pairs, static_cast(kBlockSize)), kMaxGrid)); + LaunchKernel(blocks, kBlockSize, device.unwrap())( + ltx25_decoder_rope_kernel, + static_cast(q_out.data_ptr()), + static_cast(k_out.data_ptr()), + static_cast(q.data_ptr()), + static_cast(k.data_ptr()), + static_cast(cos_t.data_ptr()), + static_cast(sin_t.data_ptr()), + static_cast(cos_h.data_ptr()), + static_cast(sin_h.data_ptr()), + static_cast(cos_w.data_ptr()), + static_cast(sin_w.data_ptr()), + num_pairs, + num_frames, + height, + width, + num_heads, + pairs_per_head, + dim_t / 2, + dim_h / 2); + } +}; + +} // namespace ltx25_decoder_rope + +} // namespace sglang diff --git a/python/sglang/kernels/ops/diffusion/README.md b/python/sglang/kernels/ops/diffusion/README.md index c4815be15..6de12b39b 100644 --- a/python/sglang/kernels/ops/diffusion/README.md +++ b/python/sglang/kernels/ops/diffusion/README.md @@ -114,6 +114,7 @@ Several norms look interchangeable and are not. Start here. | `fused_qknorm_rope_pack_kv` | JIT CUDA | as above, also packs prefix K/V | | `fused_rope_rotate_half_bitexact` | Triton | bit-exact (elementwise only) | | `ltx2_qknorm_split_rope_cuda` | JIT CUDA | close; **validated on B200** | +| `fused_ltx25_decoder_rope` | JIT CUDA | bit-exact paired 3D RoPE from cached compact axis tables | | `apply_rotary_embedding` | Triton (+fallbacks) | close; the generic entry point | | `hunyuan_qkv_rope_pack` | Triton | bit-exact; packs QKV and applies RoPE in one pass | diff --git a/python/sglang/kernels/ops/diffusion/__init__.py b/python/sglang/kernels/ops/diffusion/__init__.py index 2759c5d25..3310fad5f 100644 --- a/python/sglang/kernels/ops/diffusion/__init__.py +++ b/python/sglang/kernels/ops/diffusion/__init__.py @@ -200,6 +200,13 @@ _SPECS: tuple[tuple[str, KernelBackend, str, frozenset, str], ...] = ( _CUDA, "LTX-2 QK-norm + split RoPE.", ), + ( + "diffusion.ltx25_decoder_rope", + KernelBackend.JIT, + "rope.ltx25_decoder_rope_jit:fused_ltx25_decoder_rope", + _CUDA, + "Paired LTX-2.5 decoder 3D RoPE.", + ), ( "diffusion.rope_rotate_half", KernelBackend.TRITON, @@ -386,6 +393,8 @@ _EXPORTS: dict[str, str] = { "can_use_ltx2_qknorm_split_rope_cuda": "rope.ltx2_qknorm_split_rope_jit", "ltx2_qknorm_split_rope_cuda": "rope.ltx2_qknorm_split_rope_jit", "apply_ltx2_split_rotary_emb": "rope.ltx2_rotary_triton", + "can_use_ltx25_decoder_rope": "rope.ltx25_decoder_rope_jit", + "fused_ltx25_decoder_rope": "rope.ltx25_decoder_rope_jit", "can_use_fused_inplace_qknorm_rope": "rope.qknorm_rope_jit", "fused_inplace_qknorm_rope": "rope.qknorm_rope_jit", "fused_qknorm_rope_pack_kv": "rope.qknorm_rope_jit", diff --git a/python/sglang/kernels/ops/diffusion/rope/ltx25_decoder_rope_jit.py b/python/sglang/kernels/ops/diffusion/rope/ltx25_decoder_rope_jit.py new file mode 100644 index 000000000..f2a714d57 --- /dev/null +++ b/python/sglang/kernels/ops/diffusion/rope/ltx25_decoder_rope_jit.py @@ -0,0 +1,140 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args +from sglang.srt.utils.custom_op import register_custom_op + +if TYPE_CHECKING: + from tvm_ffi.module import Module + + +@cache_once +def _jit_ltx25_decoder_rope_module(dtype: torch.dtype) -> Module: + if dtype is not torch.bfloat16: + raise RuntimeError(f"Unsupported ltx25_decoder_rope dtype: {dtype}") + args = make_cpp_args(dtype) + return load_jit( + "diffusion_ltx25_decoder_rope", + *args, + cuda_files=["diffusion/ltx25_decoder_rope.cuh"], + cuda_wrappers=[ + ( + "ltx25_decoder_rope", + f"ltx25_decoder_rope::LTX25DecoderRopeKernel<{args}>::run", + ), + ], + ) + + +def _fake_impl( + q: torch.Tensor, + k: torch.Tensor, + cos_t: torch.Tensor, + sin_t: torch.Tensor, + cos_h: torch.Tensor, + sin_h: torch.Tensor, + cos_w: torch.Tensor, + sin_w: torch.Tensor, + dim_t: int, + dim_h: int, +) -> tuple[torch.Tensor, torch.Tensor]: + del cos_t, sin_t, cos_h, sin_h, cos_w, sin_w, dim_t, dim_h + return torch.empty_like(q), torch.empty_like(k) + + +@register_custom_op( + op_name="diffusion_ltx25_decoder_rope", + mutates_args=[], + fake_impl=_fake_impl, +) +def fused_ltx25_decoder_rope( + q: torch.Tensor, + k: torch.Tensor, + cos_t: torch.Tensor, + sin_t: torch.Tensor, + cos_h: torch.Tensor, + sin_h: torch.Tensor, + cos_w: torch.Tensor, + sin_w: torch.Tensor, + dim_t: int, + dim_h: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Apply paired LTX-2.5 decoder RoPE from compact 3D tables.""" + q_out = torch.empty_like(q) + k_out = torch.empty_like(k) + module = _jit_ltx25_decoder_rope_module(q.dtype) + module.ltx25_decoder_rope( + q_out.view(-1), + k_out.view(-1), + q.view(-1), + k.view(-1), + cos_t.view(-1), + sin_t.view(-1), + cos_h.view(-1), + sin_h.view(-1), + cos_w.view(-1), + sin_w.view(-1), + q.shape[0], + q.shape[1], + q.shape[2], + q.shape[3], + q.shape[4], + q.shape[5], + dim_t, + dim_h, + ) + return q_out, k_out + + +def can_use_ltx25_decoder_rope( + q: torch.Tensor, + k: torch.Tensor, + tables: tuple[tuple[torch.Tensor, torch.Tensor], ...], + dim_split: tuple[int, int, int], +) -> bool: + if ( + q.dim() != 6 + or len(tables) != 3 + or any(len(pair) != 2 for pair in tables) + or len(dim_split) != 3 + ): + return False + dim_t, dim_h, dim_w = dim_split + expected_shapes = ( + (q.shape[1], dim_t // 2), + (q.shape[2], dim_h // 2), + (q.shape[3], dim_w // 2), + ) + flat_tables = tuple(table for pair in tables for table in pair) + return ( + q.dtype is torch.bfloat16 + and k.dtype is q.dtype + and q.is_cuda + and k.is_cuda + and q.device == k.device + and k.shape == q.shape + and all(size > 0 for size in q.shape) + and q.shape[-1] == sum(dim_split) + and all(dim > 0 and dim % 2 == 0 for dim in dim_split) + and q.is_contiguous() + and k.is_contiguous() + and q.data_ptr() % 4 == 0 + and k.data_ptr() % 4 == 0 + and all(table.device == q.device for table in flat_tables) + and all(table.dtype is torch.float32 for table in flat_tables) + and all( + cos.shape == sin.shape == expected_shape + and cos.is_contiguous() + and sin.is_contiguous() + for (cos, sin), expected_shape in zip(tables, expected_shapes, strict=True) + ) + ) + + +__all__ = [ + "can_use_ltx25_decoder_rope", + "fused_ltx25_decoder_rope", +] diff --git a/python/sglang/multimodal_gen/runtime/models/decoders/ltx_2_5_diffusion_decoder.py b/python/sglang/multimodal_gen/runtime/models/decoders/ltx_2_5_diffusion_decoder.py index 8ec808290..1b336f6dd 100644 --- a/python/sglang/multimodal_gen/runtime/models/decoders/ltx_2_5_diffusion_decoder.py +++ b/python/sglang/multimodal_gen/runtime/models/decoders/ltx_2_5_diffusion_decoder.py @@ -17,6 +17,12 @@ import torch import torch.nn.functional as F from torch import nn +from sglang.kernels.ops.diffusion import ( + BitExactFusionGate, + can_use_ltx25_decoder_rope, + fused_ltx25_decoder_rope, + tensors_equal, +) from sglang.multimodal_gen.configs.models.decoders.ltx_2_5_diffusion_decoder import ( LTX25DiffusionDecoderConfig, ) @@ -152,6 +158,12 @@ def _unpatchify(x: torch.Tensor, patch_size: int) -> torch.Tensor: _BLOCK_MASK_CACHE: dict = {} _BLOCK_MASK_CACHE_MAX = 16 +# These tables are tiny and shared by every attention block at the same decoder +# grid. Without this cache, each Q and K rotation rebuilds them independently. +_ROPE_TABLE_CACHE: dict[tuple, tuple[tuple[torch.Tensor, torch.Tensor], ...]] = {} +_ROPE_TABLE_CACHE_MAX = 16 +_LTX25_DECODER_ROPE = BitExactFusionGate("LTX-2.5 decoder fused RoPE") + def _neighborhood_block_mask( num_frames: int, @@ -232,15 +244,47 @@ class LTX2VideoVaeRotaryPosEmbed3D(nn.Module): self.rope_dim_split = (dim_t, dim_hw, dim_hw) self.base = base - def _inv_freqs(self, dim: int, device: torch.device) -> torch.Tensor: + def _axis_tables( + self, length: int, dim: int, device: torch.device + ) -> tuple[torch.Tensor, torch.Tensor]: exponents = torch.arange(0, dim, 2, dtype=torch.float64, device=device) / dim - return (1.0 / self.base**exponents).to(torch.float32) + inv_freqs = (1.0 / self.base**exponents).to(torch.float32) + positions = torch.arange(length, dtype=torch.float32, device=device) + angles = positions[:, None] * inv_freqs[None, :] + return angles.cos(), angles.sin() + + def _tables( + self, hidden_states: torch.Tensor + ) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]: + num_frames, height, width = hidden_states.shape[1:4] + cache_key = ( + num_frames, + height, + width, + self.rope_dim_split, + self.base, + hidden_states.device, + ) + cached = _ROPE_TABLE_CACHE.get(cache_key) + if cached is not None: + return cached + + tables = tuple( + self._axis_tables(length, dim, hidden_states.device) + for length, dim in zip( + (num_frames, height, width), self.rope_dim_split, strict=True + ) + ) + if len(_ROPE_TABLE_CACHE) >= _ROPE_TABLE_CACHE_MAX: + _ROPE_TABLE_CACHE.pop(next(iter(_ROPE_TABLE_CACHE))) + _ROPE_TABLE_CACHE[cache_key] = tables + return tables def _rotate_axis( self, x: torch.Tensor, - positions: torch.Tensor, - inv_freqs: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, axis: int, ) -> torch.Tensor: out_dtype = x.dtype @@ -248,36 +292,72 @@ class LTX2VideoVaeRotaryPosEmbed3D(nn.Module): even = pairs[..., 0].float() odd = pairs[..., 1].float() # Broadcast over (B, T, H, W, heads, dim // 2), varying only along `axis`. - shape = [1, 1, 1, 1, 1, inv_freqs.shape[0]] - shape[axis] = positions.shape[0] - angles = (positions[:, None] * inv_freqs[None, :]).reshape(shape) - cos, sin = angles.cos(), angles.sin() + shape = [1, 1, 1, 1, 1, cos.shape[1]] + shape[axis] = cos.shape[0] + cos = cos.reshape(shape) + sin = sin.reshape(shape) rotated = torch.stack([even * cos - odd * sin, even * sin + odd * cos], dim=-1) return rotated.reshape(x.shape).to(out_dtype) - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - """`hidden_states`: `(B, T, H, W, heads, head_dim)`.""" + def _apply_rope( + self, + hidden_states: torch.Tensor, + tables: tuple[tuple[torch.Tensor, torch.Tensor], ...], + ) -> torch.Tensor: dim_t, dim_h, _ = self.rope_dim_split - num_frames, height, width = hidden_states.shape[1:4] - device = hidden_states.device - inv_t, inv_h, inv_w = ( - self._inv_freqs(dim, device) for dim in self.rope_dim_split - ) - - positions_t = torch.arange(num_frames, dtype=torch.float32, device=device) - positions_h = torch.arange(height, dtype=torch.float32, device=device) - positions_w = torch.arange(width, dtype=torch.float32, device=device) - rotated_t = self._rotate_axis( - hidden_states[..., :dim_t], positions_t, inv_t, axis=1 - ) + rotated_t = self._rotate_axis(hidden_states[..., :dim_t], *tables[0], axis=1) rotated_h = self._rotate_axis( - hidden_states[..., dim_t : dim_t + dim_h], positions_h, inv_h, axis=2 + hidden_states[..., dim_t : dim_t + dim_h], *tables[1], axis=2 ) rotated_w = self._rotate_axis( - hidden_states[..., dim_t + dim_h :], positions_w, inv_w, axis=3 + hidden_states[..., dim_t + dim_h :], *tables[2], axis=3 ) return torch.cat([rotated_t, rotated_h, rotated_w], dim=-1) + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """`hidden_states`: `(B, T, H, W, heads, head_dim)`.""" + return self._apply_rope(hidden_states, self._tables(hidden_states)) + + def forward_pair( + self, query: torch.Tensor, key: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Rotate Q/K together, with a bit-exact eager fallback.""" + tables = self._tables(query) + verified = _LTX25_DECODER_ROPE.verified + if ( + not _LTX25_DECODER_ROPE.disabled + and can_use_ltx25_decoder_rope(query, key, tables, self.rope_dim_split) + and (verified or _LTX25_DECODER_ROPE.can_attempt_once()) + ): + dim_t, dim_h, _ = self.rope_dim_split + try: + fused = fused_ltx25_decoder_rope( + query, + key, + *tables[0], + *tables[1], + *tables[2], + dim_t, + dim_h, + ) + except Exception as exc: + _LTX25_DECODER_ROPE.on_exception(exc, logger=logger) + else: + if verified: + return fused + eager = self._apply_rope(query, tables), self._apply_rope(key, tables) + return _LTX25_DECODER_ROPE.accept_or_fallback( + fused, + eager, + equal=tensors_equal, + logger=logger, + mismatch_msg=( + "LTX-2.5 decoder fused RoPE is not bit-exact on this " + "platform; falling back to eager" + ), + ) + return self._apply_rope(query, tables), self._apply_rope(key, tables) + class LTX2VideoVaeNeighborhoodAttention(nn.Module): """3D neighborhood attention over a channels-last `(B, T, H, W, C)` volume.""" @@ -322,7 +402,8 @@ class LTX2VideoVaeNeighborhoodAttention(nn.Module): query = self.norm_q(query) key = self.norm_k(key) query = query * self.scale - return self.rope(query), self.rope(key), value + query, key = self.rope.forward_pair(query, key) + return query, key, value def build_block_mask(self, hidden_states: torch.Tensor): """The window mask for this grid, or `None` when NATTEN handles it. diff --git a/python/sglang/multimodal_gen/test/unit/test_ltx2_5_config.py b/python/sglang/multimodal_gen/test/unit/test_ltx2_5_config.py index ee7f0565b..279e582c2 100644 --- a/python/sglang/multimodal_gen/test/unit/test_ltx2_5_config.py +++ b/python/sglang/multimodal_gen/test/unit/test_ltx2_5_config.py @@ -478,6 +478,66 @@ class TestLTX25DiffusionDecoder(unittest.TestCase): cls, _ = ModelRegistry.resolve_model_cls("LTX2VideoDiffusionDecoderModel") self.assertEqual(cls.__name__, "LTX2VideoDiffusionDecoderModel") + def test_rotary_pair_cpu_fallback_matches_original_expression(self): + import torch + + from sglang.multimodal_gen.runtime.models.decoders.ltx_2_5_diffusion_decoder import ( + LTX2VideoVaeRotaryPosEmbed3D, + ) + + def reference(module, hidden_states): + outputs = [] + offset = 0 + for axis, (length, dim) in enumerate( + zip(hidden_states.shape[1:4], module.rope_dim_split, strict=True), + 1, + ): + chunk = hidden_states[..., offset : offset + dim] + pairs = chunk.reshape(*chunk.shape[:-1], dim // 2, 2) + even = pairs[..., 0].float() + odd = pairs[..., 1].float() + exponents = torch.arange(0, dim, 2, dtype=torch.float64) / dim + inv_freqs = (1.0 / module.base**exponents).to(torch.float32) + positions = torch.arange(length, dtype=torch.float32) + angles = positions[:, None] * inv_freqs[None, :] + shape = [1, 1, 1, 1, 1, dim // 2] + shape[axis] = length + cos = angles.cos().reshape(shape) + sin = angles.sin().reshape(shape) + rotated = torch.stack( + [even * cos - odd * sin, even * sin + odd * cos], dim=-1 + ) + outputs.append(rotated.reshape(chunk.shape).to(hidden_states.dtype)) + offset += dim + return torch.cat(outputs, dim=-1) + + torch.manual_seed(42) + rope = LTX2VideoVaeRotaryPosEmbed3D(64) + query = torch.randn(1, 3, 7, 7, 2, 64, dtype=torch.bfloat16) + key = torch.randn_like(query) + + query_out, key_out = rope.forward_pair(query, key) + + self.assertTrue(torch.equal(query_out, reference(rope, query))) + self.assertTrue(torch.equal(key_out, reference(rope, key))) + + def test_rotary_tables_are_shared_across_decoder_blocks(self): + import torch + + from sglang.multimodal_gen.runtime.models.decoders.ltx_2_5_diffusion_decoder import ( + _ROPE_TABLE_CACHE, + LTX2VideoVaeRotaryPosEmbed3D, + ) + + _ROPE_TABLE_CACHE.clear() + hidden_states = torch.empty(1, 3, 7, 7, 2, 64, dtype=torch.bfloat16) + first = LTX2VideoVaeRotaryPosEmbed3D(64)._tables(hidden_states) + second = LTX2VideoVaeRotaryPosEmbed3D(64)._tables(hidden_states) + + self.assertIs(first, second) + self.assertEqual(len(_ROPE_TABLE_CACHE), 1) + _ROPE_TABLE_CACHE.clear() + class TestLTX25OptionalDecoderLoading(unittest.TestCase): @staticmethod diff --git a/test/registered/kernels/benchmark/diffusion/bench_ltx25_decoder_rope.py b/test/registered/kernels/benchmark/diffusion/bench_ltx25_decoder_rope.py new file mode 100644 index 000000000..308e15278 --- /dev/null +++ b/test/registered/kernels/benchmark/diffusion/bench_ltx25_decoder_rope.py @@ -0,0 +1,118 @@ +from dataclasses import dataclass + +import torch + +from sglang.kernels.jit.benchmark import marker +from sglang.kernels.ops.diffusion import fused_ltx25_decoder_rope +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci( + est_time=8, stage="base-b-kernel-benchmark", runner_config="1-gpu-large" +) + + +@dataclass(frozen=True) +class Case: + name: str + frames: int + height: int + width: int + heads: int + + +CASES = { + case.name: case + for case in ( + Case("stage0", 18, 17, 30, 32), + Case("stage4_tile", 16, 68, 96, 8), + Case("stage5_tile", 31, 136, 192, 4), + ) +} +DIM_SPLIT = (16, 24, 24) + + +def make_tables(case: Case) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]: + tables = [] + for length, dim in zip( + (case.frames, case.height, case.width), DIM_SPLIT, strict=True + ): + exponents = torch.arange(0, dim, 2, dtype=torch.float64, device="cuda") / dim + inv_freqs = (1.0 / 10000.0**exponents).to(torch.float32) + positions = torch.arange(length, dtype=torch.float32, device="cuda") + angles = positions[:, None] * inv_freqs[None, :] + tables.append((angles.cos(), angles.sin())) + return tuple(tables) + + +def eager_rope( + q: torch.Tensor, + k: torch.Tensor, + tables: tuple[tuple[torch.Tensor, torch.Tensor], ...], +) -> tuple[torch.Tensor, torch.Tensor]: + def apply(hidden_states: torch.Tensor) -> torch.Tensor: + outputs = [] + offset = 0 + for axis, (dim, (cos, sin)) in enumerate( + zip(DIM_SPLIT, tables, strict=True), 1 + ): + chunk = hidden_states[..., offset : offset + dim] + pairs = chunk.reshape(*chunk.shape[:-1], dim // 2, 2) + even = pairs[..., 0].float() + odd = pairs[..., 1].float() + shape = [1, 1, 1, 1, 1, dim // 2] + shape[axis] = cos.shape[0] + cos_view = cos.reshape(shape) + sin_view = sin.reshape(shape) + rotated = torch.stack( + [ + even * cos_view - odd * sin_view, + even * sin_view + odd * cos_view, + ], + dim=-1, + ) + outputs.append(rotated.reshape(chunk.shape).to(hidden_states.dtype)) + offset += dim + return torch.cat(outputs, dim=-1) + + return apply(q), apply(k) + + +@marker.parametrize("case_name", list(CASES), ci_vals=["stage5_tile"]) +@marker.benchmark("impl", ["eager", "jit"], unit="ms") +def benchmark(case_name: str, impl: str) -> marker.BenchResult: + case = CASES[case_name] + generator = torch.Generator(device="cuda").manual_seed(42) + q = torch.randn( + 1, + case.frames, + case.height, + case.width, + case.heads, + 64, + dtype=torch.bfloat16, + device="cuda", + generator=generator, + ) + k = torch.randn_like(q) + tables = make_tables(case) + if impl == "eager": + fn = eager_rope + else: + fn = lambda q, k, tables: fused_ltx25_decoder_rope( + q, + k, + *tables[0], + *tables[1], + *tables[2], + DIM_SPLIT[0], + DIM_SPLIT[1], + ) + return marker.do_bench( + fn, + input_args=(q, k, tables), + disable_log_bandwidth=True, + ) + + +if __name__ == "__main__": + benchmark.run() diff --git a/test/registered/kernels/ops/diffusion/test_ltx25_decoder_rope.py b/test/registered/kernels/ops/diffusion/test_ltx25_decoder_rope.py new file mode 100644 index 000000000..f4ba23a52 --- /dev/null +++ b/test/registered/kernels/ops/diffusion/test_ltx25_decoder_rope.py @@ -0,0 +1,114 @@ +import sys + +import pytest +import torch + +from sglang.kernels.jit.utils import get_ci_test_range +from sglang.kernels.ops.diffusion import ( + can_use_ltx25_decoder_rope, + fused_ltx25_decoder_rope, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=40, stage="base-b-kernel-unit", runner_config="1-gpu-large") + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + +DIM_SPLIT = (16, 24, 24) + + +def make_tables( + frames: int, height: int, width: int +) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]: + tables = [] + for length, dim in zip((frames, height, width), DIM_SPLIT, strict=True): + exponents = torch.arange(0, dim, 2, dtype=torch.float64, device="cuda") / dim + inv_freqs = (1.0 / 10000.0**exponents).to(torch.float32) + positions = torch.arange(length, dtype=torch.float32, device="cuda") + angles = positions[:, None] * inv_freqs[None, :] + tables.append((angles.cos(), angles.sin())) + return tuple(tables) + + +def eager_rope( + hidden_states: torch.Tensor, + tables: tuple[tuple[torch.Tensor, torch.Tensor], ...], +) -> torch.Tensor: + outputs = [] + offset = 0 + for axis, (dim, (cos, sin)) in enumerate(zip(DIM_SPLIT, tables, strict=True), 1): + chunk = hidden_states[..., offset : offset + dim] + pairs = chunk.reshape(*chunk.shape[:-1], dim // 2, 2) + even = pairs[..., 0].float() + odd = pairs[..., 1].float() + shape = [1, 1, 1, 1, 1, dim // 2] + shape[axis] = cos.shape[0] + cos = cos.reshape(shape) + sin = sin.reshape(shape) + rotated = torch.stack([even * cos - odd * sin, even * sin + odd * cos], dim=-1) + outputs.append(rotated.reshape(chunk.shape).to(hidden_states.dtype)) + offset += dim + return torch.cat(outputs, dim=-1) + + +CASES = get_ci_test_range( + [ + (1, 3, 7, 7, 1), + (1, 18, 17, 30, 32), + (1, 16, 68, 96, 8), + (1, 31, 136, 192, 4), + ], + [(1, 3, 7, 7, 1), (1, 31, 136, 192, 4)], +) + + +@pytest.mark.parametrize("batch,frames,height,width,heads", CASES) +def test_ltx25_decoder_rope_is_bit_exact( + batch: int, frames: int, height: int, width: int, heads: int +) -> None: + generator = torch.Generator(device="cuda").manual_seed(42) + q = torch.randn( + batch, + frames, + height, + width, + heads, + 64, + dtype=torch.bfloat16, + device="cuda", + generator=generator, + ) + k = torch.randn_like(q) + tables = make_tables(frames, height, width) + + q_out, k_out = fused_ltx25_decoder_rope( + q, k, *tables[0], *tables[1], *tables[2], DIM_SPLIT[0], DIM_SPLIT[1] + ) + + assert torch.equal(q_out, eager_rope(q, tables)) + assert torch.equal(k_out, eager_rope(k, tables)) + assert q_out.data_ptr() not in (q.data_ptr(), k.data_ptr()) + assert k_out.data_ptr() not in (q.data_ptr(), k.data_ptr()) + + +def test_ltx25_decoder_rope_predicate_rejects_unsupported_inputs() -> None: + q = torch.empty(1, 3, 7, 7, 1, 64, dtype=torch.bfloat16, device="cuda") + k = torch.empty_like(q) + tables = make_tables(3, 7, 7) + + assert can_use_ltx25_decoder_rope(q, k, tables, DIM_SPLIT) + assert not can_use_ltx25_decoder_rope(q.flatten(), k, tables, DIM_SPLIT) + assert not can_use_ltx25_decoder_rope(q.float(), k, tables, DIM_SPLIT) + assert not can_use_ltx25_decoder_rope(q, k[..., ::2], tables, DIM_SPLIT) + assert not can_use_ltx25_decoder_rope(q, k, tables, (16, 16, 16)) + bad_tables = (tables[0], tables[1], (tables[2][0].double(), tables[2][1])) + assert not can_use_ltx25_decoder_rope(q, k, bad_tables, DIM_SPLIT) + unaligned = torch.empty(q.numel() + 1, dtype=torch.bfloat16, device="cuda")[ + 1: + ].view_as(q) + assert unaligned.is_contiguous() + assert not can_use_ltx25_decoder_rope(unaligned, k, tables, DIM_SPLIT) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v", "-s"]))