[kernel] Reorg and fix fastvideo-kernel (#962)
Co-authored-by: ShreejithSG <shreejithsg@gmail.com>
This commit is contained in:
William Lin
co-authored by
ShreejithSG
ShreejithSG
parent
f8bfc76015
commit
2f6230abcf
@@ -0,0 +1,21 @@
|
||||
__version__ = "0.1.0"
|
||||
|
||||
from fastvideo_kernel.ops import (
|
||||
sliding_tile_attention,
|
||||
video_sparse_attn,
|
||||
)
|
||||
|
||||
from fastvideo_kernel.vmoba import (
|
||||
moba_attn_varlen,
|
||||
process_moba_input,
|
||||
process_moba_output,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"sliding_tile_attention",
|
||||
"video_sparse_attn",
|
||||
"moba_attn_varlen",
|
||||
"process_moba_input",
|
||||
"process_moba_output",
|
||||
"__version__",
|
||||
]
|
||||
@@ -0,0 +1,112 @@
|
||||
import math
|
||||
import torch
|
||||
from .triton_kernels.block_sparse_attn_triton import triton_block_sparse_attn_forward
|
||||
from .triton_kernels.st_attn_triton import sliding_tile_attention_triton
|
||||
from .triton_kernels.index import map_to_index
|
||||
|
||||
# Try to load the C++ extension
|
||||
try:
|
||||
from fastvideo_kernel._C import fastvideo_kernel_ops
|
||||
sta_fwd = getattr(fastvideo_kernel_ops, "sta_fwd", None)
|
||||
block_sparse_fwd = getattr(fastvideo_kernel_ops, "block_sparse_fwd", None)
|
||||
block_sparse_bwd = getattr(fastvideo_kernel_ops, "block_sparse_bwd", None)
|
||||
except ImportError:
|
||||
sta_fwd = None
|
||||
block_sparse_fwd = None
|
||||
block_sparse_bwd = None
|
||||
|
||||
def sliding_tile_attention(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
window_size: list,
|
||||
text_length: int,
|
||||
has_text: bool = True,
|
||||
seq_shape: str = "30x48x80",
|
||||
) -> torch.Tensor:
|
||||
# Check if the specific op is available
|
||||
if sta_fwd is None:
|
||||
return sliding_tile_attention_triton(
|
||||
q, k, v, window_size, text_length, has_text, seq_shape
|
||||
)
|
||||
|
||||
seq_length = q.shape[2]
|
||||
shape_map = {"30x48x80": 1, "36x48x48": 2, "18x48x80": 3}
|
||||
|
||||
if has_text:
|
||||
target_size = math.ceil(seq_length / 384) * 384
|
||||
pad_size = target_size - seq_length
|
||||
if pad_size > 0:
|
||||
q = torch.cat([q, q[:, :, -pad_size:]], dim=2)
|
||||
k = torch.cat([k, k[:, :, -pad_size:]], dim=2)
|
||||
v = torch.cat([v, v[:, :, -pad_size:]], dim=2)
|
||||
|
||||
output = torch.empty_like(q)
|
||||
flag = shape_map[seq_shape]
|
||||
|
||||
for head_idx, (t, h, w) in enumerate(window_size):
|
||||
sta_fwd(
|
||||
q[:, head_idx:head_idx + 1], k[:, head_idx:head_idx + 1],
|
||||
v[:, head_idx:head_idx + 1], output[:, head_idx:head_idx + 1],
|
||||
t, h, w, text_length, False, has_text, flag
|
||||
)
|
||||
|
||||
if has_text:
|
||||
sta_fwd(q, k, v, output, 3, 3, 3, text_length, True, True, flag)
|
||||
|
||||
return output[:, :, :seq_length]
|
||||
|
||||
|
||||
def video_sparse_attn(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
variable_block_sizes: torch.Tensor,
|
||||
topk: int,
|
||||
block_size: int | tuple = 64,
|
||||
compress_attn_weight: torch.Tensor = None,
|
||||
) -> torch.Tensor:
|
||||
if isinstance(block_size, int):
|
||||
block_size = (block_size, block_size, block_size)
|
||||
|
||||
block_elements = block_size[0] * block_size[1] * block_size[2]
|
||||
batch, heads, seq_len, dim = q.shape
|
||||
|
||||
# Compression branch
|
||||
q_c = q.view(batch, heads, seq_len // block_elements, block_elements, dim)
|
||||
k_c = k.view(batch, heads, seq_len // block_elements, block_elements, dim)
|
||||
v_c = v.view(batch, heads, seq_len // block_elements, block_elements, dim)
|
||||
|
||||
q_c = (q_c.float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(
|
||||
q.dtype)
|
||||
k_c = (k_c.float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(
|
||||
k.dtype)
|
||||
v_c = (v_c.float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(
|
||||
v.dtype)
|
||||
|
||||
scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (dim**0.5)
|
||||
attn = torch.softmax(scores, dim=-1)
|
||||
out_c = torch.matmul(attn, v_c)
|
||||
|
||||
out_c = out_c.view(batch, heads, seq_len // block_elements, 1, dim)
|
||||
out_c = out_c.repeat(1, 1, 1, block_elements,
|
||||
1).view(batch, heads, seq_len, dim)
|
||||
|
||||
# Sparse branch
|
||||
topk_idx = torch.topk(scores, topk, dim=-1).indices
|
||||
mask = torch.zeros_like(scores,
|
||||
dtype=torch.bool).scatter_(-1, topk_idx, True)
|
||||
|
||||
idx, num = map_to_index(mask)
|
||||
|
||||
if block_sparse_fwd is not None:
|
||||
out_s = block_sparse_fwd(
|
||||
q, k, v, idx, num, variable_block_sizes.int()
|
||||
)[0] # block_sparse_fwd returns vector<Tensor>
|
||||
else:
|
||||
out_s, _ = triton_block_sparse_attn_forward(q, k, v, idx, num,
|
||||
variable_block_sizes)
|
||||
|
||||
if compress_attn_weight is not None:
|
||||
return out_c * compress_attn_weight + out_s
|
||||
return out_c + out_s
|
||||
@@ -0,0 +1,583 @@
|
||||
"""
|
||||
Fused Attention
|
||||
===============
|
||||
|
||||
This is a Triton implementation of the Flash Attention v2 algorithm from Tri Dao
|
||||
(https://tridao.me/publications/flash2/flash2.pdf)
|
||||
|
||||
Credits: OpenAI kernel team
|
||||
"""
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
# ──────────────────────────── SPARSE ADDITION BEGIN ───────────────────────────
|
||||
import math # small utility needed by the sparse wrapper
|
||||
# ──────────────────────────── SPARSE ADDITION END ─────────────────────────────
|
||||
|
||||
# We don't run auto-tuning every time to keep the tutorial fast. Keeping
|
||||
# the code below and commenting out the equivalent parameters is convenient for
|
||||
# re-tuning.
|
||||
configs = [
|
||||
triton.Config({'BLOCK_M': BM, 'BLOCK_N': BN}, num_stages=s, num_warps=w) \
|
||||
for BM in [64]\
|
||||
for BN in [64]\
|
||||
for s in [3, 4, 7]\
|
||||
for w in [4, 8]\
|
||||
]
|
||||
|
||||
|
||||
# ──────────────────────────── SPARSE ADDITION BEGIN ───────────────────────────
|
||||
@triton.autotune(configs, key=["N_CTX", "HEAD_DIM"])
|
||||
@triton.jit
|
||||
def _attn_fwd_sparse(
|
||||
Q,
|
||||
K,
|
||||
V,
|
||||
sm_scale, #
|
||||
q2k_index,
|
||||
q2k_num,
|
||||
max_kv_blks, #
|
||||
variable_block_sizes,
|
||||
M,
|
||||
Out, #
|
||||
stride_qz,
|
||||
stride_qh,
|
||||
stride_qm,
|
||||
stride_qk,
|
||||
stride_kz,
|
||||
stride_kh,
|
||||
stride_kn,
|
||||
stride_kk,
|
||||
stride_vz,
|
||||
stride_vh,
|
||||
stride_vk,
|
||||
stride_vn,
|
||||
stride_oz,
|
||||
stride_oh,
|
||||
stride_om,
|
||||
stride_on,
|
||||
Z,
|
||||
H,
|
||||
N_CTX, #
|
||||
HEAD_DIM: tl.constexpr, #
|
||||
BLOCK_M: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr,
|
||||
STAGE: tl.constexpr):
|
||||
"""
|
||||
64×64 **block-sparse** forward kernel. Back-prop kernels remain dense
|
||||
(32×64 and 64×32) – memory footprint unchanged.
|
||||
"""
|
||||
|
||||
# ----- program-id mapping -----
|
||||
q_blk = tl.program_id(0) # Q-tile index
|
||||
off_hz = tl.program_id(1) # fused (batch, head)
|
||||
b = off_hz // H
|
||||
h = off_hz % H
|
||||
q_tiles = N_CTX // BLOCK_M
|
||||
meta_base = ((b * H + h) * q_tiles + q_blk)
|
||||
|
||||
kv_blocks = tl.load(q2k_num + meta_base) # int32
|
||||
kv_ptr = q2k_index + meta_base * max_kv_blks # ptr to list
|
||||
|
||||
# ----- base pointers -----
|
||||
qvk_off = (b.to(tl.int64) * stride_qz + h.to(tl.int64) * stride_qh)
|
||||
|
||||
Q_ptr = tl.make_block_ptr(base=Q + qvk_off,
|
||||
shape=(N_CTX, HEAD_DIM),
|
||||
strides=(stride_qm, stride_qk),
|
||||
offsets=(q_blk * BLOCK_M, 0),
|
||||
block_shape=(BLOCK_M, HEAD_DIM),
|
||||
order=(1, 0))
|
||||
|
||||
K_base = tl.make_block_ptr(base=K + qvk_off,
|
||||
shape=(HEAD_DIM, N_CTX),
|
||||
strides=(stride_kk, stride_kn),
|
||||
offsets=(0, 0),
|
||||
block_shape=(HEAD_DIM, BLOCK_N),
|
||||
order=(0, 1))
|
||||
|
||||
v_order: tl.constexpr = (0, 1) if V.dtype.element_ty == tl.float8e5 else (1,
|
||||
0)
|
||||
V_base = tl.make_block_ptr(base=V + qvk_off,
|
||||
shape=(N_CTX, HEAD_DIM),
|
||||
strides=(stride_vk, stride_vn),
|
||||
offsets=(0, 0),
|
||||
block_shape=(BLOCK_N, HEAD_DIM),
|
||||
order=v_order)
|
||||
|
||||
O_ptr = tl.make_block_ptr(base=Out + qvk_off,
|
||||
shape=(N_CTX, HEAD_DIM),
|
||||
strides=(stride_om, stride_on),
|
||||
offsets=(q_blk * BLOCK_M, 0),
|
||||
block_shape=(BLOCK_M, HEAD_DIM),
|
||||
order=(1, 0))
|
||||
|
||||
# ----- accumulators -----
|
||||
offs_m = q_blk * BLOCK_M + tl.arange(0, BLOCK_M)
|
||||
m_i = tl.full([BLOCK_M], -float("inf"), tl.float32)
|
||||
l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0
|
||||
acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)
|
||||
qk_scale = sm_scale * 1.44269504 # 1/ln2
|
||||
q = tl.load(Q_ptr)
|
||||
|
||||
# ----- sparse loop over valid K/V tiles -----
|
||||
for i in range(0, kv_blocks):
|
||||
kv_idx = tl.load(kv_ptr + i).to(tl.int32)
|
||||
block_size = tl.load(variable_block_sizes + kv_idx)
|
||||
K_ptr = tl.advance(K_base, (0, kv_idx * BLOCK_N))
|
||||
V_ptr = tl.advance(V_base, (kv_idx * BLOCK_N, 0))
|
||||
|
||||
k = tl.load(K_ptr)
|
||||
qk = tl.dot(q, k)
|
||||
# mask out invalid columns
|
||||
mask = tl.arange(0, BLOCK_N) < block_size
|
||||
qk = tl.where(mask[None, :], qk, -float("inf"))
|
||||
|
||||
m_ij = tl.maximum(m_i, tl.max(qk, 1) * qk_scale)
|
||||
p = tl.math.exp2(qk * qk_scale - m_ij[:, None])
|
||||
l_ij = tl.sum(p, 1)
|
||||
|
||||
alpha = tl.math.exp2(m_i - m_ij)
|
||||
l_i = l_i * alpha + l_ij
|
||||
acc = acc * alpha[:, None]
|
||||
|
||||
v = tl.load(V_ptr)
|
||||
acc = tl.dot(p.to(tl.bfloat16), v, acc)
|
||||
m_i = m_ij
|
||||
|
||||
# ----- epilogue -----
|
||||
m_i += tl.math.log2(l_i)
|
||||
acc = acc / l_i[:, None]
|
||||
tl.store(M + off_hz * N_CTX + offs_m, m_i)
|
||||
tl.store(O_ptr, acc.to(Out.type.element_ty))
|
||||
|
||||
|
||||
# ──────────────────────────── SPARSE ADDITION END ─────────────────────────────
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _attn_bwd_preprocess(
|
||||
O,
|
||||
DO, #
|
||||
Delta, #
|
||||
Z,
|
||||
H,
|
||||
N_CTX, #
|
||||
BLOCK_M: tl.constexpr,
|
||||
HEAD_DIM: tl.constexpr #
|
||||
):
|
||||
off_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
|
||||
off_hz = tl.program_id(1)
|
||||
off_n = tl.arange(0, HEAD_DIM)
|
||||
# load
|
||||
o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM +
|
||||
off_n[None, :])
|
||||
do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM +
|
||||
off_n[None, :]).to(tl.float32)
|
||||
delta = tl.sum(o * do, axis=1)
|
||||
# write-back
|
||||
tl.store(Delta + off_hz * N_CTX + off_m, delta)
|
||||
|
||||
|
||||
# The main inner-loop logic for computing dK and dV.
|
||||
@triton.jit
|
||||
def _attn_bwd_dkdv(
|
||||
dk,
|
||||
dv, #
|
||||
Q,
|
||||
k,
|
||||
v,
|
||||
sm_scale, #
|
||||
DO, #
|
||||
M,
|
||||
D, #
|
||||
k2q_index,
|
||||
k2q_num,
|
||||
max_q_blks,
|
||||
variable_block_sizes,
|
||||
# shared by Q/K/V/DO.
|
||||
stride_tok,
|
||||
stride_d, #
|
||||
H,
|
||||
N_CTX,
|
||||
BLOCK_M1: tl.constexpr, #
|
||||
BLOCK_N1: tl.constexpr, #
|
||||
HEAD_DIM: tl.constexpr, #
|
||||
# Filled in by the wrapper.
|
||||
start_n,
|
||||
start_m,
|
||||
num_steps):
|
||||
offs_m = start_m + tl.arange(0, BLOCK_M1)
|
||||
offs_n = start_n + tl.arange(0, BLOCK_N1)
|
||||
offs_k = tl.arange(0, HEAD_DIM)
|
||||
qT_ptrs = Q + offs_m[None, :] * stride_tok + offs_k[:, None] * stride_d
|
||||
do_ptrs = DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
|
||||
# BLOCK_N1 must be a multiple of BLOCK_M1, otherwise the code wouldn't work.
|
||||
tl.static_assert(BLOCK_N1 % BLOCK_M1 == 0)
|
||||
step_m = BLOCK_M1
|
||||
kv_blk = tl.program_id(0) # Q-tile index
|
||||
off_hz = tl.program_id(2) # fused (batch, head)
|
||||
b = off_hz // H
|
||||
h = off_hz % H
|
||||
q_tiles = N_CTX // BLOCK_N1
|
||||
meta_base = ((b * H + h) * q_tiles + kv_blk)
|
||||
|
||||
q_blocks = tl.load(k2q_num + meta_base) # int32
|
||||
q_ptr = k2q_index + meta_base * max_q_blks # ptr to list
|
||||
block_size = tl.load(variable_block_sizes + kv_blk)
|
||||
|
||||
for blk_idx in range(q_blocks * 2):
|
||||
block_sparse_offset = (tl.load(q_ptr + blk_idx // 2).to(tl.int32) * 2 +
|
||||
blk_idx % 2) * step_m
|
||||
qT = tl.load(qT_ptrs + block_sparse_offset * stride_tok)
|
||||
# Load m before computing qk to reduce pipeline stall.
|
||||
offs_m = start_m + block_sparse_offset + tl.arange(0, BLOCK_M1)
|
||||
m = tl.load(M + offs_m)
|
||||
qkT = tl.dot(k, qT)
|
||||
pT = tl.math.exp2(qkT - m[None, :])
|
||||
mask = tl.arange(0, BLOCK_N1) < block_size
|
||||
pT = tl.where(mask[:, None], pT, 0.0)
|
||||
|
||||
do = tl.load(do_ptrs + block_sparse_offset * stride_tok)
|
||||
# Compute dV.
|
||||
ppT = pT
|
||||
ppT = ppT.to(tl.bfloat16)
|
||||
dv += tl.dot(ppT, do)
|
||||
# D (= delta) is pre-divided by ds_scale.
|
||||
Di = tl.load(D + offs_m)
|
||||
# Compute dP and dS.
|
||||
dpT = tl.dot(v, tl.trans(do)).to(tl.float32)
|
||||
dsT = pT * (dpT - Di[None, :])
|
||||
dsT = dsT.to(tl.bfloat16)
|
||||
dk += tl.dot(dsT, tl.trans(qT))
|
||||
# Increment pointers.
|
||||
return dk, dv
|
||||
|
||||
|
||||
# the main inner-loop logic for computing dQ
|
||||
@triton.jit
|
||||
def _attn_bwd_dq(
|
||||
dq,
|
||||
q,
|
||||
K,
|
||||
V, #
|
||||
do,
|
||||
m,
|
||||
D,
|
||||
# shared by Q/K/V/DO.
|
||||
q2k_index,
|
||||
q2k_num,
|
||||
max_kv_blks,
|
||||
variable_block_sizes,
|
||||
stride_tok,
|
||||
stride_d, #
|
||||
H,
|
||||
N_CTX, #
|
||||
BLOCK_M2: tl.constexpr, #
|
||||
BLOCK_N2: tl.constexpr, #
|
||||
HEAD_DIM: tl.constexpr,
|
||||
# Filled in by the wrapper.
|
||||
start_m,
|
||||
start_n,
|
||||
num_steps):
|
||||
offs_m = start_m + tl.arange(0, BLOCK_M2)
|
||||
offs_n = start_n + tl.arange(0, BLOCK_N2)
|
||||
offs_k = tl.arange(0, HEAD_DIM)
|
||||
kT_ptrs = K + offs_n[None, :] * stride_tok + offs_k[:, None] * stride_d
|
||||
vT_ptrs = V + offs_n[None, :] * stride_tok + offs_k[:, None] * stride_d
|
||||
# D (= delta) is pre-divided by ds_scale.
|
||||
Di = tl.load(D + offs_m)
|
||||
# BLOCK_M2 must be a multiple of BLOCK_N2, otherwise the code wouldn't work.
|
||||
tl.static_assert(BLOCK_M2 % BLOCK_N2 == 0)
|
||||
step_n = BLOCK_N2
|
||||
|
||||
q_blk = tl.program_id(0) # Q-tile index
|
||||
off_hz = tl.program_id(2) # fused (batch, head)
|
||||
b = off_hz // H
|
||||
h = off_hz % H
|
||||
q_tiles = N_CTX // BLOCK_M2
|
||||
meta_base = ((b * H + h) * q_tiles + q_blk)
|
||||
|
||||
kv_blocks = tl.load(q2k_num + meta_base) # int32
|
||||
kv_ptr = q2k_index + meta_base * max_kv_blks # ptr to list
|
||||
block_size = tl.load(variable_block_sizes + q_blk)
|
||||
|
||||
for blk_idx in range(kv_blocks * 2):
|
||||
block_sparse_offset = (tl.load(kv_ptr + blk_idx // 2).to(tl.int32) * 2 +
|
||||
blk_idx % 2) * step_n * stride_tok
|
||||
kT = tl.load(kT_ptrs + block_sparse_offset)
|
||||
vT = tl.load(vT_ptrs + block_sparse_offset)
|
||||
qk = tl.dot(q, kT)
|
||||
p = tl.math.exp2(qk - m)
|
||||
mask = tl.arange(0, BLOCK_N2) < block_size.to(tl.int32)
|
||||
p = tl.where(mask[None, :], p, 0.0)
|
||||
# Compute dP and dS.
|
||||
dp = tl.dot(do, vT).to(tl.float32)
|
||||
ds = p * (dp - Di[:, None])
|
||||
ds = ds.to(tl.bfloat16)
|
||||
# Compute dQ.
|
||||
# NOTE: We need to de-scale dq in the end, because kT was pre-scaled.
|
||||
dq += tl.dot(ds, tl.trans(kT))
|
||||
# Increment pointers.
|
||||
return dq
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _attn_bwd(
|
||||
Q,
|
||||
K,
|
||||
V,
|
||||
sm_scale, #
|
||||
DO, #
|
||||
DQ,
|
||||
DK,
|
||||
DV, #
|
||||
M,
|
||||
D,
|
||||
q2k_index,
|
||||
q2k_num,
|
||||
max_kv_blks,
|
||||
k2q_index,
|
||||
k2q_num,
|
||||
max_q_blks,
|
||||
variable_block_sizes,
|
||||
# shared by Q/K/V/DO.
|
||||
stride_z,
|
||||
stride_h,
|
||||
stride_tok,
|
||||
stride_d, #
|
||||
H,
|
||||
N_CTX, #
|
||||
BLOCK_M1: tl.constexpr, #
|
||||
BLOCK_N1: tl.constexpr, #
|
||||
BLOCK_M2: tl.constexpr, #
|
||||
BLOCK_N2: tl.constexpr, #
|
||||
HEAD_DIM: tl.constexpr):
|
||||
LN2 = 0.6931471824645996 # = ln(2)
|
||||
|
||||
bhid = tl.program_id(2)
|
||||
off_chz = (bhid * N_CTX).to(tl.int64)
|
||||
adj = (stride_h * (bhid % H) + stride_z * (bhid // H)).to(tl.int64)
|
||||
pid = tl.program_id(0)
|
||||
|
||||
# offset pointers for batch/head
|
||||
Q += adj
|
||||
K += adj
|
||||
V += adj
|
||||
DO += adj
|
||||
DQ += adj
|
||||
DK += adj
|
||||
DV += adj
|
||||
M += off_chz
|
||||
D += off_chz
|
||||
|
||||
# load scales
|
||||
offs_k = tl.arange(0, HEAD_DIM)
|
||||
|
||||
start_n = pid * BLOCK_N1
|
||||
start_m = 0
|
||||
|
||||
offs_n = start_n + tl.arange(0, BLOCK_N1)
|
||||
|
||||
dv = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
|
||||
dk = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
|
||||
|
||||
# load K and V: they stay in SRAM throughout the inner loop.
|
||||
k = tl.load(K + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
|
||||
v = tl.load(V + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
|
||||
|
||||
num_steps = N_CTX // BLOCK_M1
|
||||
|
||||
dk, dv = _attn_bwd_dkdv( #
|
||||
dk,
|
||||
dv, #
|
||||
Q,
|
||||
k,
|
||||
v,
|
||||
sm_scale, #
|
||||
DO, #
|
||||
M,
|
||||
D, #
|
||||
k2q_index,
|
||||
k2q_num,
|
||||
max_q_blks,
|
||||
variable_block_sizes,
|
||||
stride_tok,
|
||||
stride_d, #
|
||||
H,
|
||||
N_CTX, #
|
||||
BLOCK_M1,
|
||||
BLOCK_N1,
|
||||
HEAD_DIM, #
|
||||
start_n,
|
||||
start_m,
|
||||
num_steps #
|
||||
)
|
||||
|
||||
dv_ptrs = DV + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
|
||||
tl.store(dv_ptrs, dv)
|
||||
|
||||
# Write back dK.
|
||||
dk *= sm_scale
|
||||
dk_ptrs = DK + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
|
||||
tl.store(dk_ptrs, dk)
|
||||
|
||||
# THIS BLOCK DOES DQ:
|
||||
start_m = pid * BLOCK_M2
|
||||
end_n = 0
|
||||
|
||||
offs_m = start_m + tl.arange(0, BLOCK_M2)
|
||||
|
||||
q = tl.load(Q + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)
|
||||
dq = tl.zeros([BLOCK_M2, HEAD_DIM], dtype=tl.float32)
|
||||
do = tl.load(DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)
|
||||
|
||||
m = tl.load(M + offs_m)
|
||||
m = m[:, None]
|
||||
|
||||
num_steps = N_CTX // BLOCK_N2
|
||||
dq = _attn_bwd_dq(
|
||||
dq,
|
||||
q,
|
||||
K,
|
||||
V, #
|
||||
do,
|
||||
m,
|
||||
D, #
|
||||
q2k_index,
|
||||
q2k_num,
|
||||
max_kv_blks,
|
||||
variable_block_sizes,
|
||||
stride_tok,
|
||||
stride_d, #
|
||||
H,
|
||||
N_CTX, #
|
||||
BLOCK_M2,
|
||||
BLOCK_N2,
|
||||
HEAD_DIM, #
|
||||
start_m,
|
||||
end_n,
|
||||
num_steps #
|
||||
)
|
||||
# Write back dQ.
|
||||
dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
|
||||
dq *= LN2
|
||||
tl.store(dq_ptrs, dq)
|
||||
|
||||
|
||||
# ──────────────────────────── SPARSE ADDITION BEGIN ───────────────────────────
|
||||
def triton_block_sparse_attn_forward(q, k, v, q2k_index, q2k_num,
|
||||
variable_block_sizes):
|
||||
B, H, T, D = q.shape
|
||||
sm_scale = 1.0 / math.sqrt(D)
|
||||
max_kv_blks = q2k_index.shape[-1]
|
||||
assert T % 64 == 0, f"T must be a multiple of 64, but got {T}"
|
||||
assert q2k_num.shape[
|
||||
-1] == T // 64, f"shape mismatch, T // 64 = {T // 64}, q2k_num.shape[-2] = {q2k_num.shape[-2]}"
|
||||
o = torch.empty_like(q)
|
||||
M = torch.empty((B, H, T), dtype=torch.float32, device=q.device)
|
||||
|
||||
grid = lambda _: (triton.cdiv(T, 64), B * H, 1)
|
||||
_attn_fwd_sparse[grid](q,
|
||||
k,
|
||||
v,
|
||||
sm_scale,
|
||||
q2k_index,
|
||||
q2k_num,
|
||||
max_kv_blks,
|
||||
variable_block_sizes,
|
||||
M,
|
||||
o,
|
||||
q.stride(0),
|
||||
q.stride(1),
|
||||
q.stride(2),
|
||||
q.stride(3),
|
||||
k.stride(0),
|
||||
k.stride(1),
|
||||
k.stride(2),
|
||||
k.stride(3),
|
||||
v.stride(0),
|
||||
v.stride(1),
|
||||
v.stride(2),
|
||||
v.stride(3),
|
||||
o.stride(0),
|
||||
o.stride(1),
|
||||
o.stride(2),
|
||||
o.stride(3),
|
||||
B,
|
||||
H,
|
||||
T,
|
||||
HEAD_DIM=D,
|
||||
STAGE=3)
|
||||
|
||||
return o, M
|
||||
|
||||
|
||||
def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num,
|
||||
k2q_index, k2q_num, variable_block_sizes):
|
||||
assert do.is_contiguous()
|
||||
assert q.stride() == k.stride() == v.stride() == o.stride() == do.stride()
|
||||
|
||||
B, H, T, D = q.shape
|
||||
sm_scale = 1.0 / math.sqrt(D)
|
||||
dq = torch.empty_like(q)
|
||||
dk = torch.empty_like(k)
|
||||
dv = torch.empty_like(v)
|
||||
BATCH, N_HEAD, N_CTX = q.shape[:3]
|
||||
BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2 = 32, 64, 64, 32
|
||||
RCP_LN2 = 1.4426950408889634 # = 1.0 / ln(2)
|
||||
arg_k = k
|
||||
arg_k = arg_k * (sm_scale * RCP_LN2)
|
||||
PRE_BLOCK = 64
|
||||
assert N_CTX % PRE_BLOCK == 0
|
||||
pre_grid = (N_CTX // PRE_BLOCK, BATCH * N_HEAD)
|
||||
delta = torch.empty_like(M)
|
||||
_attn_bwd_preprocess[pre_grid](
|
||||
o,
|
||||
do, #
|
||||
delta, #
|
||||
BATCH,
|
||||
N_HEAD,
|
||||
N_CTX, #
|
||||
BLOCK_M=PRE_BLOCK,
|
||||
HEAD_DIM=D #
|
||||
)
|
||||
|
||||
max_q_blks = k2q_index.shape[-1]
|
||||
max_kv_blks = q2k_index.shape[-1]
|
||||
|
||||
grid = (N_CTX // BLOCK_N1, 1, BATCH * N_HEAD)
|
||||
_attn_bwd[grid](
|
||||
q,
|
||||
arg_k,
|
||||
v,
|
||||
sm_scale,
|
||||
do,
|
||||
dq,
|
||||
dk,
|
||||
dv, #
|
||||
M,
|
||||
delta, #
|
||||
q2k_index,
|
||||
q2k_num,
|
||||
max_kv_blks,
|
||||
k2q_index,
|
||||
k2q_num,
|
||||
max_q_blks,
|
||||
variable_block_sizes,
|
||||
q.stride(0),
|
||||
q.stride(1),
|
||||
q.stride(2),
|
||||
q.stride(3), #
|
||||
N_HEAD,
|
||||
N_CTX, #
|
||||
BLOCK_M1=BLOCK_M1,
|
||||
BLOCK_N1=BLOCK_N1, #
|
||||
BLOCK_M2=BLOCK_M2,
|
||||
BLOCK_N2=BLOCK_N2, #
|
||||
HEAD_DIM=D #
|
||||
)
|
||||
|
||||
return dq, dk, dv
|
||||
@@ -0,0 +1,155 @@
|
||||
## pytorch sdpa version of block sparse ##
|
||||
import triton
|
||||
import triton.language as tl
|
||||
import torch
|
||||
|
||||
|
||||
@triton.jit
|
||||
def topk_index_to_map_kernel(
|
||||
map_ptr,
|
||||
index_ptr,
|
||||
map_bs_stride,
|
||||
map_h_stride,
|
||||
map_q_stride,
|
||||
map_kv_stride,
|
||||
index_bs_stride,
|
||||
index_h_stride,
|
||||
index_q_stride,
|
||||
index_kv_stride,
|
||||
topk,
|
||||
):
|
||||
b, h, q = tl.program_id(0), tl.program_id(1), tl.program_id(2)
|
||||
index_ptr_base = index_ptr + b * index_bs_stride + h * index_h_stride + q * index_q_stride
|
||||
map_ptr_base = map_ptr + b * map_bs_stride + h * map_h_stride + q * map_q_stride
|
||||
|
||||
for i in tl.static_range(topk):
|
||||
index = tl.load(index_ptr_base + i * index_kv_stride)
|
||||
tl.store(map_ptr_base + index * map_kv_stride, 1.0)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def map_to_index_kernel(
|
||||
map_ptr,
|
||||
index_ptr,
|
||||
index_num_ptr,
|
||||
map_bs_stride,
|
||||
map_h_stride,
|
||||
map_q_stride,
|
||||
map_kv_stride,
|
||||
index_bs_stride,
|
||||
index_h_stride,
|
||||
index_q_stride,
|
||||
index_kv_stride,
|
||||
index_num_bs_stride,
|
||||
index_num_h_stride,
|
||||
index_num_q_stride,
|
||||
num_kv_blocks,
|
||||
):
|
||||
b, h, q = tl.program_id(0), tl.program_id(1), tl.program_id(2)
|
||||
index_ptr_base = index_ptr + b * index_bs_stride + h * index_h_stride + q * index_q_stride
|
||||
map_ptr_base = map_ptr + b * map_bs_stride + h * map_h_stride + q * map_q_stride
|
||||
|
||||
num = 0
|
||||
for i in tl.range(num_kv_blocks):
|
||||
map_entry = tl.load(map_ptr_base + i * map_kv_stride)
|
||||
if map_entry:
|
||||
tl.store(index_ptr_base + num * index_kv_stride, i)
|
||||
num += 1
|
||||
|
||||
tl.store(
|
||||
index_num_ptr + b * index_num_bs_stride + h * index_num_h_stride +
|
||||
q * index_num_q_stride, num)
|
||||
|
||||
|
||||
def topk_index_to_map(index: torch.Tensor,
|
||||
num_kv_blocks: int,
|
||||
transpose_map: bool = False):
|
||||
"""
|
||||
Convert topk indices to a map.
|
||||
|
||||
Args:
|
||||
index: [bs, h, num_q_blocks, topk]
|
||||
The topk indices tensor.
|
||||
num_kv_blocks: int
|
||||
The number of key-value blocks in the block_map returned
|
||||
transpose_map: bool
|
||||
If True, the block_map will be transposed on the final two dimensions.
|
||||
|
||||
Returns:
|
||||
block_map: [bs, h, num_q_blocks, num_kv_blocks]
|
||||
A binary map where 1 indicates that the q block attends to the kv block.
|
||||
"""
|
||||
bs, h, num_q_blocks, topk = index.shape
|
||||
|
||||
if transpose_map is False:
|
||||
block_map = torch.zeros((bs, h, num_q_blocks, num_kv_blocks),
|
||||
dtype=torch.bool,
|
||||
device=index.device)
|
||||
else:
|
||||
block_map = torch.zeros((bs, h, num_kv_blocks, num_q_blocks),
|
||||
dtype=torch.bool,
|
||||
device=index.device)
|
||||
block_map = block_map.transpose(2, 3)
|
||||
|
||||
grid = (bs, h, num_q_blocks)
|
||||
topk_index_to_map_kernel[grid](
|
||||
block_map,
|
||||
index,
|
||||
block_map.stride(0),
|
||||
block_map.stride(1),
|
||||
block_map.stride(2),
|
||||
block_map.stride(3),
|
||||
index.stride(0),
|
||||
index.stride(1),
|
||||
index.stride(2),
|
||||
index.stride(3),
|
||||
topk=topk,
|
||||
)
|
||||
|
||||
return block_map
|
||||
|
||||
|
||||
def map_to_index(block_map: torch.Tensor):
|
||||
"""
|
||||
Convert a block map to indices and counts.
|
||||
|
||||
Args:
|
||||
block_map: [bs, h, num_q_blocks, num_kv_blocks]
|
||||
The block map tensor.
|
||||
|
||||
Returns:
|
||||
index: [bs, h, num_q_blocks, num_kv_blocks]
|
||||
The indices of the blocks.
|
||||
index_num: [bs, h, num_q_blocks]
|
||||
The number of blocks for each q block.
|
||||
"""
|
||||
bs, h, num_q_blocks, num_kv_blocks = block_map.shape
|
||||
|
||||
index = torch.full((block_map.shape),
|
||||
-1,
|
||||
dtype=torch.int32,
|
||||
device=block_map.device)
|
||||
index_num = torch.empty((bs, h, num_q_blocks),
|
||||
dtype=torch.int32,
|
||||
device=block_map.device)
|
||||
|
||||
grid = (bs, h, num_q_blocks)
|
||||
map_to_index_kernel[grid](
|
||||
block_map,
|
||||
index,
|
||||
index_num,
|
||||
block_map.stride(0),
|
||||
block_map.stride(1),
|
||||
block_map.stride(2),
|
||||
block_map.stride(3),
|
||||
index.stride(0),
|
||||
index.stride(1),
|
||||
index.stride(2),
|
||||
index.stride(3),
|
||||
index_num.stride(0),
|
||||
index_num.stride(1),
|
||||
index_num.stride(2),
|
||||
num_kv_blocks=num_kv_blocks,
|
||||
)
|
||||
|
||||
return index, index_num
|
||||
@@ -0,0 +1,327 @@
|
||||
import math
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
def is_cuda():
|
||||
return triton.runtime.driver.active.get_current_target().backend == "cuda"
|
||||
|
||||
|
||||
def is_hip():
|
||||
target = triton.runtime.driver.active.get_current_target()
|
||||
return target.backend == 'hip'
|
||||
|
||||
def get_common_autotune_config():
|
||||
configs = [
|
||||
triton.Config({'BLOCK_Q': BLOCK_Q, 'BLOCK_KV': BLOCK_KV}, num_stages=s, num_warps=w) \
|
||||
for BLOCK_Q in [32, 64, 128]\
|
||||
for BLOCK_KV in [32, 64, 128]\
|
||||
for s in [1, 2, 3, 4]\
|
||||
for w in [4, 8]\
|
||||
]
|
||||
return configs
|
||||
|
||||
|
||||
def get_cuda_autotune_config():
|
||||
# cuda and hip can use differnt autotune configs
|
||||
return get_common_autotune_config()
|
||||
|
||||
|
||||
def get_hip_autotune_config():
|
||||
# cuda and hip can use differnt autotune configs
|
||||
return get_common_autotune_config()
|
||||
|
||||
|
||||
def get_autotune_config():
|
||||
if is_cuda():
|
||||
return get_cuda_autotune_config()
|
||||
else:
|
||||
return get_hip_autotune_config()
|
||||
|
||||
|
||||
@triton.jit
|
||||
def clamp_int(value, min_val, max_val):
|
||||
ret = tl.where(value > max_val, max_val, value)
|
||||
ret = tl.where(ret < min_val, min_val, ret)
|
||||
return ret
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _attn_fwd_loop(
|
||||
q, k, v, kv_mask, m, l, acc, sm_scale,
|
||||
MASK_KV: tl.constexpr,
|
||||
):
|
||||
scores = tl.dot(q, k.T) #[BLOCK_Q, BLOCK_KV]
|
||||
scores = scores * sm_scale
|
||||
if MASK_KV:
|
||||
scores = tl.where(kv_mask[None, :], scores, -float('inf'))
|
||||
|
||||
current_m = tl.max(scores, axis=1)
|
||||
new_m = tl.maximum(m, current_m)
|
||||
exp_scores = tl.math.exp2(scores - new_m[:, None])
|
||||
current_l = tl.sum(exp_scores, axis=1)
|
||||
|
||||
# Update L <- L * exp(M - M') + L1, M <- M'
|
||||
alpha = tl.math.exp2(m - new_m)
|
||||
l = l * alpha + current_l
|
||||
m = new_m
|
||||
|
||||
# Update O <- O * exp(M - M') + P @ V
|
||||
acc = (acc * alpha[:, None] + tl.dot(exp_scores.to(v.type.element_ty), v))
|
||||
|
||||
return m, l, acc
|
||||
|
||||
|
||||
@triton.autotune(
|
||||
configs=get_autotune_config(),
|
||||
key=['head_dim'],
|
||||
)
|
||||
@triton.jit
|
||||
def triton_sta_kernel(
|
||||
Q, K, V, output,
|
||||
batch_size: int, num_heads: int, seq_len: int, head_dim: int,
|
||||
img_seq_len: int,
|
||||
text_length: int,
|
||||
canvas_t: int, canvas_h: int, canvas_w: int,
|
||||
kernel_t: int, kernel_h: int, kernel_w: int,
|
||||
tile_t: int, tile_h: int, tile_w: int,
|
||||
scale: float,
|
||||
has_text: tl.constexpr,
|
||||
text_q: tl.constexpr,
|
||||
BLOCK_Q: tl.constexpr,
|
||||
BLOCK_KV: tl.constexpr,
|
||||
BLOCK_DIM: tl.constexpr,
|
||||
):
|
||||
total_tile_size = tile_t * tile_h * tile_w
|
||||
q_block_per_tile = (total_tile_size + BLOCK_Q - 1) // BLOCK_Q
|
||||
|
||||
batch_idx = tl.program_id(0)
|
||||
head_idx = tl.program_id(1)
|
||||
if text_q:
|
||||
q_block_idx = tl.program_id(2)
|
||||
else:
|
||||
q_tile_flat = tl.program_id(2) // q_block_per_tile
|
||||
q_block_idx = tl.program_id(2) % q_block_per_tile
|
||||
|
||||
m = tl.full((BLOCK_Q,), -float('inf'), dtype=tl.float32)
|
||||
l = tl.zeros((BLOCK_Q,), dtype=tl.float32)
|
||||
acc = tl.zeros((BLOCK_Q, BLOCK_DIM), dtype=tl.float32)
|
||||
|
||||
q_offset = (batch_idx * num_heads + head_idx) * seq_len * head_dim
|
||||
if text_q:
|
||||
q_base_idx = img_seq_len + q_block_idx * BLOCK_Q
|
||||
else:
|
||||
q_base_idx = q_tile_flat * total_tile_size + q_block_idx * BLOCK_Q
|
||||
|
||||
q_offset_in_tile = tl.arange(0, BLOCK_Q)
|
||||
q_idx = q_base_idx + q_offset_in_tile
|
||||
q_mask = (q_block_idx * BLOCK_Q + tl.arange(0, BLOCK_Q)) < total_tile_size
|
||||
|
||||
q = tl.load(
|
||||
Q + q_offset + q_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :],
|
||||
mask=q_mask[:, None],
|
||||
other=0.0
|
||||
) # [BLOCK_Q, BLOCK_DIM]
|
||||
|
||||
# Scale sm_scale by log_2(e) and use 2^x instead of exp
|
||||
sm_scale = scale * 1.4426950408889634
|
||||
|
||||
num_tiles_t = canvas_t // tile_t
|
||||
num_tiles_h = canvas_h // tile_h
|
||||
num_tiles_w = canvas_w // tile_w
|
||||
tiles_per_hw = num_tiles_h * num_tiles_w
|
||||
|
||||
if text_q:
|
||||
kv_tile_start_t = 0
|
||||
kv_tile_end_t = num_tiles_t
|
||||
|
||||
kv_tile_start_h = 0
|
||||
kv_tile_end_h = num_tiles_h
|
||||
|
||||
kv_tile_start_w = 0
|
||||
kv_tile_end_w = num_tiles_w
|
||||
|
||||
else:
|
||||
q_tile_t = q_tile_flat // tiles_per_hw
|
||||
remaining = q_tile_flat % tiles_per_hw
|
||||
q_tile_h = remaining // num_tiles_w
|
||||
q_tile_w = remaining % num_tiles_w
|
||||
|
||||
kernel_center_t = clamp_int(q_tile_t, kernel_t // 2, (num_tiles_t - 1) - kernel_t // 2)
|
||||
kernel_center_h = clamp_int(q_tile_h, kernel_h // 2, (num_tiles_h - 1) - kernel_h // 2)
|
||||
kernel_center_w = clamp_int(q_tile_w, kernel_w // 2, (num_tiles_w - 1) - kernel_w // 2)
|
||||
|
||||
kv_tile_start_t = kernel_center_t - kernel_t // 2
|
||||
kv_tile_end_t = kernel_center_t + kernel_t // 2 + 1
|
||||
kv_tile_end_t = tl.where(kv_tile_end_t > num_tiles_t, num_tiles_t, kv_tile_end_t)
|
||||
|
||||
kv_tile_start_h = kernel_center_h - kernel_h // 2
|
||||
kv_tile_end_h = kernel_center_h + kernel_h // 2 + 1
|
||||
kv_tile_end_h = tl.where(kv_tile_end_h > num_tiles_h, num_tiles_h, kv_tile_end_h)
|
||||
|
||||
kv_tile_start_w = kernel_center_w - kernel_w // 2
|
||||
kv_tile_end_w = kernel_center_w + kernel_w // 2 + 1
|
||||
kv_tile_end_w = tl.where(kv_tile_end_w > num_tiles_w, num_tiles_w, kv_tile_end_w)
|
||||
|
||||
# for kv_img
|
||||
for kv_tile_t in tl.range(kv_tile_start_t, kv_tile_end_t):
|
||||
for kv_tile_h in tl.range(kv_tile_start_h, kv_tile_end_h):
|
||||
for kv_tile_w in tl.range(kv_tile_start_w, kv_tile_end_w):
|
||||
kv_base_idx = (kv_tile_t * num_tiles_h * num_tiles_w + kv_tile_h * num_tiles_w + kv_tile_w) * total_tile_size
|
||||
|
||||
for kv_block_idx in tl.range(0, total_tile_size, BLOCK_KV):
|
||||
kv_offset_in_block = tl.arange(0, BLOCK_KV)
|
||||
kv_idx = kv_base_idx + kv_block_idx + kv_offset_in_block
|
||||
kv_mask = (kv_block_idx + tl.arange(0, BLOCK_KV)) < total_tile_size
|
||||
|
||||
kv_offset = (batch_idx * num_heads + head_idx) * seq_len * head_dim
|
||||
|
||||
k = tl.load(
|
||||
K + kv_offset + kv_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :],
|
||||
mask=kv_mask[:, None],
|
||||
other=0.0
|
||||
) # [BLOCK_KV, BLOCK_DIM]
|
||||
v = tl.load(
|
||||
V + kv_offset + kv_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :],
|
||||
mask=kv_mask[:, None],
|
||||
other=0.0
|
||||
) # [BLOCK_KV, BLOCK_DIM]
|
||||
|
||||
m, l, acc = _attn_fwd_loop(q, k, v, kv_mask, m, l, acc, sm_scale, False)
|
||||
|
||||
|
||||
# for kv_text
|
||||
if has_text:
|
||||
kv_base_idx = img_seq_len
|
||||
for kv_block_idx in tl.range(0, total_tile_size, BLOCK_KV):
|
||||
kv_offset_in_block = tl.arange(0, BLOCK_KV)
|
||||
kv_idx = kv_base_idx + kv_block_idx + kv_offset_in_block
|
||||
kv_mask = (kv_block_idx + tl.arange(0, BLOCK_KV)) < text_length
|
||||
|
||||
kv_offset = (batch_idx * num_heads + head_idx) * seq_len * head_dim
|
||||
|
||||
k = tl.load(
|
||||
K + kv_offset + kv_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :],
|
||||
mask=kv_mask[:, None],
|
||||
other=0.0
|
||||
) # [BLOCK_KV, BLOCK_DIM]
|
||||
v = tl.load(
|
||||
V + kv_offset + kv_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :],
|
||||
mask=kv_mask[:, None],
|
||||
other=0.0
|
||||
) # [BLOCK_KV, BLOCK_DIM]
|
||||
|
||||
m, l, acc = _attn_fwd_loop(q, k, v, kv_mask, m, l, acc, sm_scale, True)
|
||||
|
||||
|
||||
output_acc = acc / l[:, None]
|
||||
tl.store(
|
||||
output + q_offset + q_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :],
|
||||
output_acc,
|
||||
mask=q_mask[:, None]
|
||||
) # [BLOCK_Q, BLOCK_DIM]
|
||||
|
||||
|
||||
def sliding_tile_attention_triton(
|
||||
q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
|
||||
window_size, text_length: int,
|
||||
has_text=True, dit_seq_shape='30x48x80') -> torch.Tensor:
|
||||
seq_length = q.shape[2]
|
||||
if has_text:
|
||||
assert q.shape[2] >= 115200 and q.shape[2] <= 115456, f"Unsupported {dit_seq_shape}, current shape is {q.shape}, only support '30x48x80' for HunyuanVideo"
|
||||
target_size = math.ceil(seq_length / 384) * 384
|
||||
pad_size = target_size - seq_length
|
||||
if pad_size > 0:
|
||||
q = torch.cat([q, q[:, :, -pad_size:]], dim=2)
|
||||
k = torch.cat([k, k[:, :, -pad_size:]], dim=2)
|
||||
v = torch.cat([v, v[:, :, -pad_size:]], dim=2)
|
||||
else:
|
||||
if dit_seq_shape == '36x48x48': # Stepvideo
|
||||
assert q.shape[2] == 82944
|
||||
elif dit_seq_shape == '18x48x80': # Wan
|
||||
assert q.shape[2] == 69120
|
||||
else:
|
||||
raise ValueError(f"Unsupported {dit_seq_shape}, current shape is {q.shape}, only support '36x48x48' for Stepvideo and '18x48x80' for Wan")
|
||||
assert q.shape[1] == len(window_size), "Number of heads must match the number of window sizes"
|
||||
|
||||
batch_size, num_heads, seq_len, head_dim = q.shape
|
||||
if dit_seq_shape == '30x48x80': # Hunyuan
|
||||
canvas_t, canvas_h, canvas_w = 30, 48, 80
|
||||
tile_t, tile_h, tile_w = 6, 8, 8
|
||||
elif dit_seq_shape == '36x48x48': # Stepvideo
|
||||
canvas_t, canvas_h, canvas_w = 36, 48, 48
|
||||
tile_t, tile_h, tile_w = 6, 8, 8
|
||||
elif dit_seq_shape == '18x48x80': # Wan
|
||||
canvas_t, canvas_h, canvas_w = 18, 48, 80
|
||||
tile_t, tile_h, tile_w = 6, 8, 8
|
||||
|
||||
img_seq_len = canvas_t * canvas_h * canvas_w
|
||||
|
||||
num_tiles_t = canvas_t // tile_t
|
||||
num_tiles_h = canvas_h // tile_h
|
||||
num_tiles_w = canvas_w // tile_w
|
||||
num_tiles = num_tiles_t * num_tiles_h * num_tiles_w
|
||||
|
||||
total_tile_size = tile_t * tile_h * tile_w
|
||||
|
||||
# BLOCK_Q=128
|
||||
# BLOCK_KV=128
|
||||
BLOCK_DIM = head_dim
|
||||
|
||||
output = torch.empty_like(q)
|
||||
|
||||
# for q_img
|
||||
# kernel_size maybe different for different head
|
||||
# This for loop is ugly. but it is actually quite efficient. The sequence dimension alone can already oversubscribe SMs
|
||||
for head_index, (kernel_t, kernel_h, kernel_w) in enumerate(window_size):
|
||||
for batch in range(batch_size):
|
||||
q_head, k_head, v_head, o_head = (q[batch:batch + 1, head_index:head_index + 1],
|
||||
k[batch:batch + 1, head_index:head_index + 1],
|
||||
v[batch:batch + 1, head_index:head_index + 1],
|
||||
output[batch:batch + 1, head_index:head_index + 1])
|
||||
|
||||
# triton_sta_kernel[(1, 1, num_tiles * triton.cdiv(total_tile_size, BLOCK_Q))](
|
||||
grid = lambda META: (1, 1, num_tiles * triton.cdiv(total_tile_size, META['BLOCK_Q']))
|
||||
triton_sta_kernel[grid](
|
||||
q_head, k_head, v_head, o_head,
|
||||
1, 1, seq_len, head_dim,
|
||||
img_seq_len,
|
||||
text_length,
|
||||
canvas_t, canvas_h, canvas_w,
|
||||
kernel_t, kernel_h, kernel_w,
|
||||
tile_t, tile_h, tile_w,
|
||||
scale=1.0 / (head_dim ** 0.5),
|
||||
has_text=has_text,
|
||||
text_q=False,
|
||||
# BLOCK_Q=BLOCK_Q,
|
||||
# BLOCK_KV=BLOCK_KV,
|
||||
BLOCK_DIM=BLOCK_DIM,
|
||||
)
|
||||
|
||||
# for q_text
|
||||
# kernel_t, kernel_h, kernel_w is not used, set to (3, 3, 3)
|
||||
if has_text:
|
||||
# triton_sta_kernel[(batch_size, num_heads, triton.cdiv(total_tile_size, BLOCK_Q))](
|
||||
grid = lambda META: (batch_size, num_heads, triton.cdiv(total_tile_size, META['BLOCK_Q']))
|
||||
triton_sta_kernel[grid](
|
||||
q, k, v, output,
|
||||
batch_size, num_heads, seq_len, head_dim,
|
||||
img_seq_len,
|
||||
text_length,
|
||||
canvas_t, canvas_h, canvas_w,
|
||||
3, 3, 3,
|
||||
#kernel_t, kernel_h, kernel_w,
|
||||
tile_t, tile_h, tile_w,
|
||||
scale=1.0 / (head_dim ** 0.5),
|
||||
has_text=has_text,
|
||||
text_q=True,
|
||||
# BLOCK_Q=BLOCK_Q,
|
||||
# BLOCK_KV=BLOCK_KV,
|
||||
BLOCK_DIM=BLOCK_DIM,
|
||||
)
|
||||
|
||||
if has_text:
|
||||
if pad_size > 0:
|
||||
output = output[:, :, :seq_length]
|
||||
|
||||
return output
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user