[diffusion] Fuse LTX-2.5 decoder 3D RoPE (#35698)
This commit is contained in:
Xiaoyu Zhang
@@ -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 |
|
||||
|
||||
@@ -0,0 +1,198 @@
|
||||
#pragma once
|
||||
|
||||
#include <sgl_kernel/tensor.h>
|
||||
#include <sgl_kernel/utils.h>
|
||||
|
||||
#include <sgl_kernel/type.cuh>
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
#include <sgl_kernel/vec.cuh>
|
||||
|
||||
#include <algorithm>
|
||||
#include <cstdint>
|
||||
#include <type_traits>
|
||||
|
||||
namespace sglang {
|
||||
|
||||
namespace ltx25_decoder_rope {
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr uint32_t kBlockSize = 256;
|
||||
constexpr uint32_t kMaxGrid = 65535;
|
||||
|
||||
template <typename T>
|
||||
SGL_DEVICE device::AlignedVector<T, 2> rotate_pair(device::AlignedVector<T, 2> input, float cos, float sin) {
|
||||
static_assert(std::is_same_v<T, bf16_t>);
|
||||
const float even = static_cast<float>(input[0]);
|
||||
const float odd = static_cast<float>(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<T, 2> output;
|
||||
output[0] = static_cast<T>(__fsub_rn(even_cos, odd_sin));
|
||||
output[1] = static_cast<T>(__fadd_rn(even_sin, odd_cos));
|
||||
return output;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__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<T, 2>;
|
||||
const int64_t stride = static_cast<int64_t>(gridDim.x) * blockDim.x;
|
||||
for (int64_t pair_index = static_cast<int64_t>(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 <typename T>
|
||||
struct LTX25DecoderRopeKernel {
|
||||
static_assert(std::is_same_v<T, bf16_t>);
|
||||
|
||||
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<T, 2>;
|
||||
|
||||
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<kDLCUDA>();
|
||||
TensorMatcher({N}).with_dtype<T>().with_device(device).verify(q_out).verify(k_out).verify(q).verify(k);
|
||||
TensorMatcher({RT}).with_dtype<float>().with_device(device).verify(cos_t).verify(sin_t);
|
||||
TensorMatcher({RH}).with_dtype<float>().with_device(device).verify(cos_h).verify(sin_h);
|
||||
TensorMatcher({RW}).with_dtype<float>().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<uintptr_t>(q_out.data_ptr()) % alignof(Pair) == 0 &&
|
||||
reinterpret_cast<uintptr_t>(k_out.data_ptr()) % alignof(Pair) == 0 &&
|
||||
reinterpret_cast<uintptr_t>(q.data_ptr()) % alignof(Pair) == 0 &&
|
||||
reinterpret_cast<uintptr_t>(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<uint32_t>(std::min<int64_t>(div_ceil(num_pairs, static_cast<int64_t>(kBlockSize)), kMaxGrid));
|
||||
LaunchKernel(blocks, kBlockSize, device.unwrap())(
|
||||
ltx25_decoder_rope_kernel<T>,
|
||||
static_cast<T*>(q_out.data_ptr()),
|
||||
static_cast<T*>(k_out.data_ptr()),
|
||||
static_cast<const T*>(q.data_ptr()),
|
||||
static_cast<const T*>(k.data_ptr()),
|
||||
static_cast<const float*>(cos_t.data_ptr()),
|
||||
static_cast<const float*>(sin_t.data_ptr()),
|
||||
static_cast<const float*>(cos_h.data_ptr()),
|
||||
static_cast<const float*>(sin_h.data_ptr()),
|
||||
static_cast<const float*>(cos_w.data_ptr()),
|
||||
static_cast<const float*>(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
|
||||
@@ -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 |
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
+106
-25
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
@@ -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"]))
|
||||
Reference in New Issue
Block a user