diff --git a/python/sglang/jit_kernel/diffusion/triton/ltx2_rotary.py b/python/sglang/jit_kernel/diffusion/triton/ltx2_rotary.py new file mode 100644 index 000000000..6bf8791ce --- /dev/null +++ b/python/sglang/jit_kernel/diffusion/triton/ltx2_rotary.py @@ -0,0 +1,90 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _ltx2_split_rotary_kernel( + out_ptr, + x_ptr, + cos_ptr, + sin_ptr, + seq_len: tl.constexpr, + num_heads: tl.constexpr, + head_dim: tl.constexpr, + half_dim: tl.constexpr, + stride_cos_b: tl.constexpr, + stride_cos_h: tl.constexpr, + stride_cos_t: tl.constexpr, + stride_sin_b: tl.constexpr, + stride_sin_h: tl.constexpr, + stride_sin_t: tl.constexpr, + BLOCK_HALF: tl.constexpr, +): + pid_bt = tl.program_id(0) + head = tl.program_id(1) + batch = pid_bt // seq_len + token = pid_bt - batch * seq_len + offsets = tl.arange(0, BLOCK_HALF) + mask = offsets < half_dim + + x_base = ((batch * seq_len + token) * num_heads + head) * head_dim + cos_base = batch * stride_cos_b + head * stride_cos_h + token * stride_cos_t + sin_base = batch * stride_sin_b + head * stride_sin_h + token * stride_sin_t + + x_first = tl.load(x_ptr + x_base + offsets, mask=mask, other=0.0) + x_second = tl.load(x_ptr + x_base + half_dim + offsets, mask=mask, other=0.0) + cos = tl.load(cos_ptr + cos_base + offsets, mask=mask, other=0.0) + sin = tl.load(sin_ptr + sin_base + offsets, mask=mask, other=0.0) + + # Match the original PyTorch order: x * cos is written as BF16 first, then + # addcmul_ computes the sine product in FP32 before the final BF16 store. + out_first = (x_first * cos).to(tl.bfloat16).to(tl.float32) + ( + -x_second.to(tl.float32) * sin.to(tl.float32) + ) + out_second = (x_second * cos).to(tl.bfloat16).to(tl.float32) + ( + x_first.to(tl.float32) * sin.to(tl.float32) + ) + + tl.store(out_ptr + x_base + offsets, out_first, mask=mask) + tl.store(out_ptr + x_base + half_dim + offsets, out_second, mask=mask) + + +def apply_ltx2_split_rotary_emb( + x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor +) -> torch.Tensor: + batch, seq_len, inner_dim = x.shape + cos_batch, num_heads, cos_seq_len, half_dim = cos.shape + head_dim = half_dim * 2 + if ( + cos_batch != batch + or cos_seq_len != seq_len + or inner_dim != num_heads * head_dim + or sin.shape != cos.shape + ): + raise ValueError( + "LTX2 split RoPE shape mismatch: " + f"x={tuple(x.shape)}, cos={tuple(cos.shape)}, sin={tuple(sin.shape)}" + ) + + out = torch.empty_like(x) + block_half = triton.next_power_of_2(half_dim) + _ltx2_split_rotary_kernel[(batch * seq_len, num_heads)]( + out, + x, + cos, + sin, + seq_len, + num_heads, + head_dim, + half_dim, + cos.stride(0), + cos.stride(1), + cos.stride(2), + sin.stride(0), + sin.stride(1), + sin.stride(2), + BLOCK_HALF=block_half, + num_warps=1, + ) + return out 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 e3bc1c8fd..7a4c649ae 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py @@ -102,6 +102,24 @@ def apply_split_rotary_emb( x: torch.Tensor, freqs: Tuple[torch.Tensor, torch.Tensor] ) -> torch.Tensor: cos, sin = freqs + if ( + x.ndim == 3 + and cos.ndim == 4 + and sin.ndim == 4 + and x.dtype == torch.bfloat16 + and cos.dtype == torch.bfloat16 + and sin.dtype == torch.bfloat16 + and x.is_cuda + and x.is_contiguous() + and cos.is_cuda + and sin.is_cuda + ): + from sglang.jit_kernel.diffusion.triton.ltx2_rotary import ( + apply_ltx2_split_rotary_emb, + ) + + return apply_ltx2_split_rotary_emb(x, cos, sin) + x_dtype = x.dtype needs_reshape = False if x.ndim != 4 and cos.ndim == 4: