376 lines
13 KiB
Python
376 lines
13 KiB
Python
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 is_cdna3_cdna4():
|
|
target = triton.runtime.driver.active.get_current_target()
|
|
return (target.arch == 'gfx1201' or target.arch == 'gfx1101' or target.arch == 'gfx1100'
|
|
or target.arch == 'gfx1030')
|
|
|
|
|
|
def get_common_autotune_config():
|
|
# cdna arch does not support a 4-stage software pipeline, see https://github.com/ROCm/triton/issues/916
|
|
supported_num_staged = [1, 2] if is_cdna3_cdna4() else [1, 2, 3, 4]
|
|
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 supported_num_staged\
|
|
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
|