[KDA-Pilot] Add LTX2 QKNorm split-RoPE CUDA fast path (#29708)
This commit is contained in:
Xiaoyu Zhang
@@ -0,0 +1,277 @@
|
||||
// CUDA fast path for LTX2 Q/K RMSNorm + split RoPE.
|
||||
//
|
||||
// Developed with MIT HAN Lab Kernel Design Agents:
|
||||
// https://github.com/mit-han-lab/kernel-design-agents
|
||||
//
|
||||
// This mirrors the LTX2 eager oracle:
|
||||
// torch.nn.RMSNorm(input) returns fp32 under bf16 autocast, then split RoPE
|
||||
// runs in fp32 and rounds once to bf16 at the final attention input.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <sgl_kernel/tensor.h> // For TensorMatcher, SymbolicSize, SymbolicDevice
|
||||
#include <sgl_kernel/utils.h> // For RuntimeCheck
|
||||
|
||||
#include <sgl_kernel/utils.cuh> // For LaunchKernel and CUDA dtype aliases
|
||||
|
||||
#include <cstdint>
|
||||
#include <cuda_bf16.h>
|
||||
|
||||
namespace sglang_ltx2_qknorm_split_rope {
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr int kThreads = 128;
|
||||
|
||||
inline const char* data_ptr(const tvm::ffi::TensorView& t) {
|
||||
return static_cast<const char*>(t.data_ptr()) + t.byte_offset();
|
||||
}
|
||||
|
||||
inline char* mutable_data_ptr(const tvm::ffi::TensorView& t) {
|
||||
return static_cast<char*>(t.data_ptr()) + t.byte_offset();
|
||||
}
|
||||
|
||||
SGL_DEVICE float compute_rstd(
|
||||
const bf16_t* __restrict__ xrow,
|
||||
int64_t hidden_size,
|
||||
float eps,
|
||||
int tid,
|
||||
int lane,
|
||||
int warp_id,
|
||||
float* warp_sum,
|
||||
float* s_rstd) {
|
||||
float local = 0.f;
|
||||
const int64_t n_vec = hidden_size >> 2;
|
||||
for (int64_t i = tid; i < n_vec; i += kThreads) {
|
||||
const int64_t base = i << 2;
|
||||
const float v0 = __bfloat162float(xrow[base + 0]);
|
||||
const float v1 = __bfloat162float(xrow[base + 1]);
|
||||
const float v2 = __bfloat162float(xrow[base + 2]);
|
||||
const float v3 = __bfloat162float(xrow[base + 3]);
|
||||
local = fmaf(v0, v0, local);
|
||||
local = fmaf(v1, v1, local);
|
||||
local = fmaf(v2, v2, local);
|
||||
local = fmaf(v3, v3, local);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int offset = 16; offset > 0; offset >>= 1) {
|
||||
local += __shfl_down_sync(0xffffffffu, local, offset);
|
||||
}
|
||||
if (lane == 0) {
|
||||
warp_sum[warp_id] = local;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (tid == 0) {
|
||||
const float total = (warp_sum[0] + warp_sum[2]) + (warp_sum[1] + warp_sum[3]);
|
||||
*s_rstd = rsqrtf(total / static_cast<float>(hidden_size) + eps);
|
||||
}
|
||||
__syncthreads();
|
||||
return *s_rstd;
|
||||
}
|
||||
|
||||
SGL_DEVICE float norm_value(float x, float weight, float rstd) {
|
||||
return weight * (rstd * x);
|
||||
}
|
||||
|
||||
SGL_DEVICE void rope_pair(float x0, float x1, float cos, float sin, float& y0, float& y1) {
|
||||
const float p0 = x0 * cos;
|
||||
const float p1 = x1 * cos;
|
||||
y0 = fmaf(-sin, x1, p0);
|
||||
y1 = fmaf(sin, x0, p1);
|
||||
}
|
||||
|
||||
__global__ void ltx2_qknorm_split_rope_kernel(
|
||||
const bf16_t* __restrict__ x,
|
||||
const bf16_t* __restrict__ cos,
|
||||
const bf16_t* __restrict__ sin,
|
||||
const bf16_t* __restrict__ weight,
|
||||
bf16_t* __restrict__ out,
|
||||
float eps,
|
||||
int64_t seq_len,
|
||||
int64_t num_heads,
|
||||
int64_t head_dim,
|
||||
int64_t stride_cos_b,
|
||||
int64_t stride_cos_h,
|
||||
int64_t stride_cos_t,
|
||||
int64_t stride_sin_b,
|
||||
int64_t stride_sin_h,
|
||||
int64_t stride_sin_t) {
|
||||
const int64_t row = static_cast<int64_t>(blockIdx.x);
|
||||
const int64_t batch = row / seq_len;
|
||||
const int64_t token = row - batch * seq_len;
|
||||
const int64_t hidden_size = num_heads * head_dim;
|
||||
const int64_t half_dim = head_dim >> 1;
|
||||
const auto* __restrict__ xrow = x + row * hidden_size;
|
||||
auto* __restrict__ outrow = out + row * hidden_size;
|
||||
const int tid = threadIdx.x + threadIdx.y * 32;
|
||||
const int lane = threadIdx.x;
|
||||
const int warp_id = threadIdx.y;
|
||||
|
||||
__shared__ float warp_sum[4];
|
||||
__shared__ float s_rstd;
|
||||
const float rstd = compute_rstd(xrow, hidden_size, eps, tid, lane, warp_id, warp_sum, &s_rstd);
|
||||
|
||||
const int64_t num_pairs = num_heads * half_dim;
|
||||
for (int64_t pair = tid; pair < num_pairs; pair += kThreads) {
|
||||
const int64_t head = pair / half_dim;
|
||||
const int64_t offset = pair - head * half_dim;
|
||||
const int64_t idx0 = head * head_dim + offset;
|
||||
const int64_t idx1 = idx0 + half_dim;
|
||||
const float n0 = norm_value(__bfloat162float(xrow[idx0]), __bfloat162float(weight[idx0]), rstd);
|
||||
const float n1 = norm_value(__bfloat162float(xrow[idx1]), __bfloat162float(weight[idx1]), rstd);
|
||||
const int64_t cos_offset = batch * stride_cos_b + head * stride_cos_h + token * stride_cos_t + offset;
|
||||
const int64_t sin_offset = batch * stride_sin_b + head * stride_sin_h + token * stride_sin_t + offset;
|
||||
|
||||
float y0;
|
||||
float y1;
|
||||
rope_pair(n0, n1, __bfloat162float(cos[cos_offset]), __bfloat162float(sin[sin_offset]), y0, y1);
|
||||
outrow[idx0] = __float2bfloat16_rn(y0);
|
||||
outrow[idx1] = __float2bfloat16_rn(y1);
|
||||
}
|
||||
}
|
||||
|
||||
inline void launch_one(
|
||||
const tvm::ffi::TensorView& x,
|
||||
const tvm::ffi::TensorView& cos,
|
||||
const tvm::ffi::TensorView& sin,
|
||||
const tvm::ffi::TensorView& weight,
|
||||
const tvm::ffi::TensorView& out,
|
||||
float eps,
|
||||
int64_t num_rows,
|
||||
int64_t seq_len,
|
||||
int64_t num_heads,
|
||||
int64_t head_dim,
|
||||
int64_t stride_cos_b,
|
||||
int64_t stride_cos_h,
|
||||
int64_t stride_cos_t,
|
||||
int64_t stride_sin_b,
|
||||
int64_t stride_sin_h,
|
||||
int64_t stride_sin_t,
|
||||
DLDevice device) {
|
||||
if (num_rows == 0) {
|
||||
return;
|
||||
}
|
||||
host::RuntimeCheck(num_rows <= static_cast<int64_t>(UINT32_MAX), "LTX2 QKNorm split-RoPE grid is too large");
|
||||
host::LaunchKernel(dim3(static_cast<uint32_t>(num_rows)), dim3(32, 4), device)(
|
||||
ltx2_qknorm_split_rope_kernel,
|
||||
reinterpret_cast<const bf16_t*>(data_ptr(x)),
|
||||
reinterpret_cast<const bf16_t*>(data_ptr(cos)),
|
||||
reinterpret_cast<const bf16_t*>(data_ptr(sin)),
|
||||
reinterpret_cast<const bf16_t*>(data_ptr(weight)),
|
||||
reinterpret_cast<bf16_t*>(mutable_data_ptr(out)),
|
||||
eps,
|
||||
seq_len,
|
||||
num_heads,
|
||||
head_dim,
|
||||
stride_cos_b,
|
||||
stride_cos_h,
|
||||
stride_cos_t,
|
||||
stride_sin_b,
|
||||
stride_sin_h,
|
||||
stride_sin_t);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
struct LTX2QKNormSplitRopeKernel {
|
||||
static void
|
||||
run(tvm::ffi::TensorView q_out,
|
||||
tvm::ffi::TensorView k_out,
|
||||
tvm::ffi::TensorView q,
|
||||
tvm::ffi::TensorView q_cos,
|
||||
tvm::ffi::TensorView q_sin,
|
||||
tvm::ffi::TensorView q_weight,
|
||||
tvm::ffi::TensorView k,
|
||||
tvm::ffi::TensorView k_cos,
|
||||
tvm::ffi::TensorView k_sin,
|
||||
tvm::ffi::TensorView k_weight,
|
||||
double eps,
|
||||
int64_t num_heads,
|
||||
int64_t head_dim) {
|
||||
using namespace host;
|
||||
|
||||
RuntimeCheck(num_heads > 0, "num_heads must be positive");
|
||||
RuntimeCheck(head_dim > 0, "head_dim must be positive");
|
||||
RuntimeCheck(head_dim % 2 == 0, "head_dim must be even");
|
||||
const int64_t hidden_size = num_heads * head_dim;
|
||||
RuntimeCheck(hidden_size % 4 == 0, "hidden size must be divisible by 4");
|
||||
|
||||
auto batch = SymbolicSize{"batch"};
|
||||
auto q_seq_len = SymbolicSize{"q_seq_len"};
|
||||
auto k_seq_len = SymbolicSize{"k_seq_len"};
|
||||
auto heads = SymbolicSize{"num_heads"};
|
||||
auto half_dim = SymbolicSize{"half_dim"};
|
||||
auto device = SymbolicDevice{};
|
||||
heads.set_value(num_heads);
|
||||
half_dim.set_value(head_dim / 2);
|
||||
device.set_options<kDLCUDA>();
|
||||
|
||||
TensorMatcher({batch, q_seq_len, hidden_size}).with_dtype<bf16_t>().with_device(device).verify(q).verify(q_out);
|
||||
TensorMatcher({batch, k_seq_len, hidden_size}).with_dtype<bf16_t>().with_device(device).verify(k).verify(k_out);
|
||||
TensorMatcher({hidden_size}).with_dtype<bf16_t>().with_device(device).verify(q_weight);
|
||||
TensorMatcher({hidden_size}).with_dtype<bf16_t>().with_device(device).verify(k_weight);
|
||||
TensorMatcher({batch, heads, q_seq_len, half_dim})
|
||||
.with_strides({-1, -1, -1, 1})
|
||||
.with_dtype<bf16_t>()
|
||||
.with_device(device)
|
||||
.verify(q_cos);
|
||||
TensorMatcher({batch, heads, q_seq_len, half_dim})
|
||||
.with_strides({-1, -1, -1, 1})
|
||||
.with_dtype<bf16_t>()
|
||||
.with_device(device)
|
||||
.verify(q_sin);
|
||||
TensorMatcher({batch, heads, k_seq_len, half_dim})
|
||||
.with_strides({-1, -1, -1, 1})
|
||||
.with_dtype<bf16_t>()
|
||||
.with_device(device)
|
||||
.verify(k_cos);
|
||||
TensorMatcher({batch, heads, k_seq_len, half_dim})
|
||||
.with_strides({-1, -1, -1, 1})
|
||||
.with_dtype<bf16_t>()
|
||||
.with_device(device)
|
||||
.verify(k_sin);
|
||||
|
||||
const int64_t batch_size = batch.unwrap();
|
||||
const DLDevice dl_device = device.unwrap();
|
||||
launch_one(
|
||||
q,
|
||||
q_cos,
|
||||
q_sin,
|
||||
q_weight,
|
||||
q_out,
|
||||
static_cast<float>(eps),
|
||||
batch_size * q_seq_len.unwrap(),
|
||||
q_seq_len.unwrap(),
|
||||
num_heads,
|
||||
head_dim,
|
||||
q_cos.stride(0),
|
||||
q_cos.stride(1),
|
||||
q_cos.stride(2),
|
||||
q_sin.stride(0),
|
||||
q_sin.stride(1),
|
||||
q_sin.stride(2),
|
||||
dl_device);
|
||||
launch_one(
|
||||
k,
|
||||
k_cos,
|
||||
k_sin,
|
||||
k_weight,
|
||||
k_out,
|
||||
static_cast<float>(eps),
|
||||
batch_size * k_seq_len.unwrap(),
|
||||
k_seq_len.unwrap(),
|
||||
num_heads,
|
||||
head_dim,
|
||||
k_cos.stride(0),
|
||||
k_cos.stride(1),
|
||||
k_cos.stride(2),
|
||||
k_sin.stride(0),
|
||||
k_sin.stride(1),
|
||||
k_sin.stride(2),
|
||||
dl_device);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace sglang_ltx2_qknorm_split_rope
|
||||
@@ -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),
|
||||
)
|
||||
@@ -10,6 +10,10 @@ import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from sglang.jit_kernel.diffusion.ltx2_qknorm_split_rope import (
|
||||
can_use_ltx2_qknorm_split_rope_cuda,
|
||||
ltx2_qknorm_split_rope_cuda,
|
||||
)
|
||||
from sglang.jit_kernel.diffusion.residual_gate_add import (
|
||||
can_use_residual_gate_add_cuda,
|
||||
residual_gate_add_cuda,
|
||||
@@ -53,6 +57,7 @@ logger = init_logger(__name__)
|
||||
ADALN_NUM_BASE_PARAMS = 6
|
||||
ADALN_NUM_CROSS_ATTN_PARAMS = 3
|
||||
_LTX2_RESIDUAL_GATE_CUDA_DISABLED = False
|
||||
_LTX2_QKNORM_SPLIT_ROPE_CUDA_DISABLED = False
|
||||
|
||||
|
||||
def _ltx2_residual_gate_add(
|
||||
@@ -76,6 +81,66 @@ def _ltx2_residual_gate_add(
|
||||
return residual + update * gate
|
||||
|
||||
|
||||
def _ltx2_try_fused_qknorm_split_rope(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
q_norm: nn.Module,
|
||||
k_norm: nn.Module,
|
||||
q_cos: torch.Tensor,
|
||||
q_sin: torch.Tensor,
|
||||
k_cos: torch.Tensor,
|
||||
k_sin: torch.Tensor,
|
||||
*,
|
||||
eps: float,
|
||||
num_heads: int,
|
||||
head_dim: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor] | None:
|
||||
global _LTX2_QKNORM_SPLIT_ROPE_CUDA_DISABLED
|
||||
|
||||
if (
|
||||
_LTX2_QKNORM_SPLIT_ROPE_CUDA_DISABLED
|
||||
or get_tp_world_size() != 1
|
||||
or not isinstance(q_norm, nn.RMSNorm)
|
||||
or not isinstance(k_norm, nn.RMSNorm)
|
||||
or float(q_norm.eps) != float(eps)
|
||||
or float(k_norm.eps) != float(eps)
|
||||
or not can_use_ltx2_qknorm_split_rope_cuda(
|
||||
q,
|
||||
q_cos,
|
||||
q_sin,
|
||||
q_norm.weight,
|
||||
k,
|
||||
k_cos,
|
||||
k_sin,
|
||||
k_norm.weight,
|
||||
num_heads=num_heads,
|
||||
head_dim=head_dim,
|
||||
)
|
||||
):
|
||||
return None
|
||||
|
||||
try:
|
||||
return ltx2_qknorm_split_rope_cuda(
|
||||
q,
|
||||
q_cos,
|
||||
q_sin,
|
||||
q_norm.weight,
|
||||
k,
|
||||
k_cos,
|
||||
k_sin,
|
||||
k_norm.weight,
|
||||
eps=eps,
|
||||
num_heads=num_heads,
|
||||
head_dim=head_dim,
|
||||
)
|
||||
except Exception as exc:
|
||||
if torch.compiler.is_compiling():
|
||||
raise
|
||||
logger.warning_once(f"Disabling LTX2 QKNorm split-RoPE CUDA fast path: {exc}")
|
||||
_LTX2_QKNORM_SPLIT_ROPE_CUDA_DISABLED = True
|
||||
return None
|
||||
|
||||
|
||||
_LTX2_FUSED_ADA_VALUES_RUNTIME_DISABLED = False
|
||||
|
||||
|
||||
@@ -748,11 +813,7 @@ class LTX2Attention(nn.Module):
|
||||
q, _ = self.to_q(x)
|
||||
k, _ = self.to_k(context_)
|
||||
|
||||
if self.qk_norm:
|
||||
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 @@ class LTX2Attention(nn.Module):
|
||||
k_cos, k_sin = self._slice_rope_for_tp(
|
||||
k_cos, k_sin, tp_rank=tp_rank, tp_size=tp_size
|
||||
)
|
||||
if cos.dim() == 3:
|
||||
if self.qk_norm and cos.dim() != 3:
|
||||
assert self.q_norm is not None and self.k_norm is not None
|
||||
fused_qk = _ltx2_try_fused_qknorm_split_rope(
|
||||
q,
|
||||
k,
|
||||
self.q_norm,
|
||||
self.k_norm,
|
||||
cos,
|
||||
sin,
|
||||
k_cos,
|
||||
k_sin,
|
||||
eps=self.norm_eps,
|
||||
num_heads=self.local_heads,
|
||||
head_dim=self.dim_head,
|
||||
)
|
||||
|
||||
if fused_qk is not None:
|
||||
q, k = fused_qk
|
||||
else:
|
||||
if self.qk_norm:
|
||||
assert self.q_norm is not None and self.k_norm is not None
|
||||
q = self.q_norm(q)
|
||||
k = self.k_norm(k)
|
||||
|
||||
if pe is not None and cos.dim() == 3:
|
||||
q = apply_interleaved_rotary_emb(q, (cos, sin))
|
||||
k = apply_interleaved_rotary_emb(k, (k_cos, k_sin))
|
||||
else:
|
||||
elif pe is not None:
|
||||
q = apply_split_rotary_emb(q, (cos, sin))
|
||||
k = apply_split_rotary_emb(k, (k_cos, k_sin))
|
||||
|
||||
|
||||
Reference in New Issue
Block a user