"""
Author: Eric Lin (xihlin)
"""
"""
... note(bapatra)::
This is written as one big file, instead of splitting into logical components because I was running into issues with transformers auto module
imports when splitting into different files. I've tried keeping the logical partitions demarkated with comment blocks, but it is not ideal.
In the future, would be really good to revisit this and refactor into a more readable file structure.
"""
from typing import TypeVar
from functools import lru_cache
import math
import pytest
import torch
import numpy as np
import triton
import triton.language as tl
import os
import dataclasses
Phi3SmallConfig = TypeVar('Phi3SmallConfig')
# triton 2.0.0: fail at backward on A100, for the examples, if h_dim=128.
# Done
# 1. strided of qkv
# 2. seq len not power of 2
# 3. bf16 with Triton May, 2023
# TODO:
# 1. wip: support non-contiguous backward, also help reduce memory allocation in training (q, k, v split)
# 2. block sparse with different BLOCK_M, BLOCK_N?
# 3. for Lq not divided by BLOCK_M, BLOCK_N, only apply mask to K/V on last batch, still need to apply mask on Q.
# Attempt, fail to compile
# 4. For 2nd iter of inference, BLOCK_M=1, how to make things work? K/V maynot divided by BLOCK_N.
# 5. The inner loop can also be paralled via bigger num_stage(better) or on different thread-block (via m/L and atomic update, but this no-comm/sync between blocks)
###########################################################
################### Kernel Parameters #####################
###########################################################
@dataclasses.dataclass
class BlockSparseParams(object):
block_size: int
kernel_block_size: int
num_local_blocks: int
vert_stride: int
homo_head_pattern: bool = False
@classmethod
def from_config(cls, config: Phi3SmallConfig) -> "BlockSparseParams":
return cls(
block_size=config.blocksparse_block_size,
kernel_block_size=config.blocksparse_triton_kernel_block_size,
num_local_blocks=config.blocksparse_num_local_blocks,
vert_stride=config.blocksparse_vert_stride,
homo_head_pattern=config.blocksparse_homo_head_pattern,
)
###########################################################
###########################################################
###########################################################
################### Utility Functions #####################
###########################################################
# helper functions for 3D sparse pattern
# these function are not optimized and very inefficient. Avoid calling them too frequent.
# currently, it is only called within `get_local_strided_sparse_attention_op`, which is cached.
def dense_to_crow_col(x):
''' Turning a 2D/3D torch tensor (x) to CSR rows/cols indexing.
param:
TODO:
1. improve efficiency, is it faster if done in CPU, or customize a cuda kernel for it?
NOTE: col_indices padded -1
'''
pad = -1
dim = x.dim()
assert x.dim() in (2, 3)
if x.dim() == 2:
x = x[None]
x = [xi.to_sparse_csr() for xi in x]
crows = torch.vstack([xi.crow_indices() for xi in x])
cols = [xi.col_indices() for xi in x]
max_cols = max(len(xi) for xi in cols)
cols = [torch.cat([xi, pad + xi.new_zeros(max_cols - xi.shape[0])]) for xi in cols]
cols = torch.vstack(cols)
if dim == 2:
crows = crows[0]
cols = cols[0]
return crows, cols
def crow_col_to_dense(crows, cols, dtype=torch.float16):
dim = crows.dim()
if dim == 1:
crows = crows[None]
cols = cols[None]
device = crows.device
crows, cols = crows.cpu(), cols.cpu() # faster in cpu
shape = (crows.shape[0], crows.shape[1] - 1, cols.max() + 1)
x = torch.zeros(shape, dtype=dtype)
for i in range(shape[0]):
for j in range(shape[1]):
x[i, j, cols[i, crows[i, j]:crows[i, j+1]]] = 1
if dim == 1:
x = x[0]
return x.to(device)
def dense_to_ccol_row(x):
'''Similar, but to CSC format
'''
x = x.transpose(-2, -1)
return dense_to_crow_col(x)
def ccol_row_to_dense(ccol, rows, dtype=torch.float16):
return crow_col_to_dense(ccol, rows, dtype).permute(0, 2, 1).contiguous()
def _get_sparse_attn_mask_homo_head(q_len, N_CTX, dtype, device, BLOCK=128, local_blocks=4, vert_stride=4, return_dense=False):
'''
:return: a tuple of 3:
- tuple of crow_indices, col_indices representation of CSR format.
- block dense mask
- all token dense mask (be aware that it can be OOM if it is too big) if `return_dense==True`, otherwise, None
'''
with torch.no_grad():
N_BLOCK = triton.cdiv(N_CTX, BLOCK)
q_pos = torch.arange(N_BLOCK)[:, None]
k_pos = torch.arange(N_BLOCK)[None]
mask_vert_strided = (torch.arange(N_BLOCK) + 1) % vert_stride == 0
block_mask_dense = ((q_pos >= k_pos) & ((q_pos - k_pos < local_blocks) | mask_vert_strided)).to(device).to(dtype)
N_BLOCK_Q = triton.cdiv(q_len, BLOCK)
block_mask_dense_output = block_mask_dense[-N_BLOCK_Q:].contiguous().to_sparse_csr()
if return_dense:
mask_dense = torch.kron(block_mask_dense, block_mask_dense.new_ones((BLOCK, BLOCK)))
causal_mask = torch.tril(torch.ones(N_CTX, N_CTX)).type_as(mask_dense)[-q_len:]
mask_dense = mask_dense[-q_len:, :N_CTX] * causal_mask
return (block_mask_dense_output.crow_indices(), block_mask_dense_output.col_indices()), block_mask_dense, mask_dense
else:
return (block_mask_dense_output.crow_indices(), block_mask_dense_output.col_indices()), block_mask_dense, None
def _get_sparse_attn_mask(n_heads, q_len, N_CTX, dtype, device, BLOCK=128, local_blocks=4, vert_stride=4, homo_head=True, return_dense=False):
'''
:return: a tuple of 3:
- tuple of crow_indices, col_indices representation of CSR format.
- block dense mask
- all token dense mask (be aware that it can be OOM if it is too big) if `return_dense==True`, otherwise, None
'''
if homo_head:
with torch.no_grad():
(crow, col), block_mask_dense, mask_dense = _get_sparse_attn_mask_homo_head(q_len, N_CTX, dtype, device, BLOCK, local_blocks, vert_stride, return_dense)
crow = crow[None].expand(n_heads, crow.shape[0])
col = col[None].expand(n_heads, col.shape[0])
if return_dense:
mask_dense = mask_dense[None].expand(n_heads, *mask_dense.shape)
return (crow, col), block_mask_dense, mask_dense
with torch.no_grad():
N_BLOCK = triton.cdiv(N_CTX, BLOCK)
q_pos = torch.arange(N_BLOCK)[None, :, None]
k_pos = torch.arange(N_BLOCK)[None, None]
head_sliding_step = max(1, int(vert_stride / n_heads)) # if vert_stride <= n_heads, rotating the heads
mask_vert_strided = [(torch.arange(N_BLOCK) + h * head_sliding_step + 1) % vert_stride == 0 for h in range(n_heads)]
mask_vert_strided = torch.vstack(mask_vert_strided).unsqueeze(1)
block_mask_dense = ((q_pos >= k_pos) & ((q_pos - k_pos < local_blocks) | mask_vert_strided)).to(device).to(dtype)
N_BLOCK_Q = triton.cdiv(q_len, BLOCK)
block_mask_dense_output = block_mask_dense[:, -N_BLOCK_Q:]
if return_dense:
mask_dense = torch.kron(block_mask_dense, block_mask_dense.new_ones((BLOCK, BLOCK)))
causal_mask = torch.tril(torch.ones(N_CTX, N_CTX)).type_as(mask_dense)[-q_len:]
mask_dense = mask_dense[..., -q_len:, :N_CTX] * causal_mask[None]
return dense_to_crow_col(block_mask_dense_output), block_mask_dense, mask_dense
else:
return dense_to_crow_col(block_mask_dense_output), block_mask_dense, None
def get_sparse_attn_mask(q, N_CTX, *args, **kwargs):
return _get_sparse_attn_mask(q.size(1), q.size(2), N_CTX, q.dtype, q.device, *args, **kwargs)
###########################################################
###########################################################
###########################################################
###################### Training Kernels ###################
###########################################################
# TODO: only apply loading/saving mask on the last iteration for EVEN_N_BLOCK, useful for 1st iteration of inference.
# Experiment failed inside loop.
# Another idea: only on saving? load even out of boundary(will it causes illegal access error)?
@triton.jit
def _fwd_kernel(
Q, K, V, sm_scale,
layout_crow_ptr,
layout_col_ptr,
layout_crow_stride_h, layout_crow_stride_m,
layout_col_stride_h, layout_col_stride_m,
TMP, L, M, # NOTE: TMP is a scratchpad buffer to workaround a compiler bug. TMP, L, M are assumed to have contiguous layouts
Out,
stride_qz, stride_qh, stride_qm, stride_qd,
stride_kz, stride_kh, stride_kn, stride_kd,
stride_vz, stride_vh, stride_vn, stride_vd,
stride_oz, stride_oh, stride_om, stride_od,
Z, H, N_CTX,
PAST_LEN,
Q_ROUNDED_LEN,
BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr,
BLOCK_N: tl.constexpr,
EVEN_M_BLOCK: tl.constexpr,
EVEN_N_BLOCK: tl.constexpr,
INFERENCE: tl.constexpr,
NUM_DBLOCKS: tl.constexpr,
):
Q_LEN = N_CTX - PAST_LEN
start_m = tl.program_id(0)
off_hz = tl.program_id(1)
off_h = off_hz % H
off_z = off_hz // H
Q += off_z * stride_qz + off_h * stride_qh
K += off_z * stride_kz + off_h * stride_kh
V += off_z * stride_vz + off_h * stride_vh
# initialize offsets
offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
offs_n = tl.arange(0, BLOCK_N)
offs_d = tl.arange(0, BLOCK_DMODEL)
off_q = offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qd
# off_k = offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kd
off_k = offs_n[None, :] * stride_kn + offs_d[:, None] * stride_kd
off_v = offs_n[:, None] * stride_vn + offs_d[None, :] * stride_vd
# Initialize pointers to Q, K, V
q_ptrs = Q + off_q
k_ptrs = K + off_k
v_ptrs = V + off_v
# initialize pointer to m and l
t_ptrs = TMP + off_hz * Q_ROUNDED_LEN + offs_m
m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float('inf')
l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
if NUM_DBLOCKS >= 2:
acc2 = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
# load q: it will stay in SRAM throughout
if EVEN_M_BLOCK:
q = tl.load(q_ptrs)
if NUM_DBLOCKS >= 2:
q2 = tl.load(q_ptrs + BLOCK_DMODEL * stride_qd)
else:
q = tl.load(q_ptrs, mask=offs_m[:, None] < Q_LEN)
if NUM_DBLOCKS >= 2:
q2 = tl.load(q_ptrs + BLOCK_DMODEL * stride_qd, mask=offs_m[:, None] < Q_LEN)
layout_ptr = layout_crow_ptr + off_h * layout_crow_stride_h + start_m * layout_crow_stride_m
start_l = tl.load(layout_ptr).to(tl.int32)
end_l = tl.load(layout_ptr + layout_crow_stride_m).to(tl.int32)
# loop over k, v and update accumulator
for col_idx_idx in range(start_l, end_l):
col_idx = tl.load(layout_col_ptr + off_h * layout_col_stride_h + col_idx_idx * layout_col_stride_m).to(tl.int32)
start_n = col_idx * BLOCK_N
# -- compute qk ----
if EVEN_N_BLOCK:
k = tl.load(k_ptrs + start_n * stride_kn)
else:
k = tl.load(k_ptrs + start_n * stride_kn, mask=offs_n[None, :] + start_n < N_CTX)
qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
qk += tl.dot(q, k)
if NUM_DBLOCKS >= 2:
if EVEN_N_BLOCK:
k = tl.load(k_ptrs + start_n * stride_kn + BLOCK_DMODEL * stride_kd)
else:
k = tl.load(k_ptrs + start_n * stride_kn + BLOCK_DMODEL * stride_kd, mask=offs_n[None, :] + start_n < N_CTX)
qk += tl.dot(q2, k)
qk *= sm_scale
qk += tl.where(offs_m[:, None] + PAST_LEN >= (start_n + offs_n[None, :]), 0, float('-inf'))
# -- compute m_ij, p, l_ij
m_ij = tl.max(qk, 1)
p = tl.exp(qk - m_ij[:, None])
l_ij = tl.sum(p, 1)
# -- update m_i and l_i
m_i_new = tl.maximum(m_i, m_ij)
alpha = tl.exp(m_i - m_i_new)
beta = tl.exp(m_ij - m_i_new)
l_i_new = alpha * l_i + beta * l_ij
# -- update output accumulator --
# scale p
p_scale = beta / l_i_new
p = p * p_scale[:, None]
# scale acc
acc_scale = l_i / l_i_new * alpha
# tl.store(t_ptrs, acc_scale)
# acc_scale = tl.load(t_ptrs) # BUG: have to store and immediately load
acc = acc * acc_scale[:, None]
if NUM_DBLOCKS >= 2:
acc2 = acc2 * acc_scale[:, None]
p = p.to(Q.dtype.element_ty)
# update acc
if EVEN_N_BLOCK:
v = tl.load(v_ptrs + start_n * stride_vn)
else:
v = tl.load(v_ptrs + start_n * stride_vn, mask=offs_n[:, None] + start_n < N_CTX)
acc += tl.dot(p, v)
if NUM_DBLOCKS >= 2:
if EVEN_N_BLOCK:
v = tl.load(v_ptrs + start_n * stride_vn + BLOCK_DMODEL * stride_vd)
else:
v = tl.load(v_ptrs + start_n * stride_vn + BLOCK_DMODEL * stride_vd, mask=offs_n[:, None] + start_n < N_CTX)
acc2 += tl.dot(p, v)
# update m_i and l_i
l_i = l_i_new
m_i = m_i_new
# rematerialize offsets to save registers
# start_m = tl.program_id(0)
# offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
# write back l and m
if not INFERENCE:
l_ptrs = L + off_hz * N_CTX + offs_m
m_ptrs = M + off_hz * N_CTX + offs_m
if EVEN_M_BLOCK:
tl.store(l_ptrs, l_i)
tl.store(m_ptrs, m_i)
else:
tl.store(l_ptrs, l_i, mask=offs_m < Q_LEN)
tl.store(m_ptrs, m_i, mask=offs_m < Q_LEN)
# initialize pointers to output
# offs_n = tl.arange(0, BLOCK_DMODEL)
off_o = off_z * stride_oz + off_h * stride_oh + offs_m[:, None] * stride_om + offs_d[None, :] * stride_od
out_ptrs = Out + off_o
tl.store(out_ptrs, acc, mask=offs_m[:, None] < Q_LEN)
if NUM_DBLOCKS >= 2:
tl.store(out_ptrs + BLOCK_DMODEL * stride_od, acc2, mask=offs_m[:, None] < Q_LEN)
## backward
@triton.heuristics(
{
'EVEN_M_BLOCK': lambda kwargs: kwargs['N_CTX'] % kwargs['BLOCK_M'] == 0,
}
)
@triton.jit
def _bwd_preprocess(
Out, DO, L, # assume contiguous for Out, DO, L, NewDO, Delta layout.
NewDO, Delta,
N_CTX,
BLOCK_M: tl.constexpr, D_HEAD: tl.constexpr,
EVEN_M_BLOCK: tl.constexpr,
):
off_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
off_d = tl.arange(0, D_HEAD)
# load
if EVEN_M_BLOCK:
o = tl.load(Out + off_m[:, None] * D_HEAD + off_d[None, :]).to(tl.float32)
do = tl.load(DO + off_m[:, None] * D_HEAD + off_d[None, :]).to(tl.float32)
else:
o = tl.load(Out + off_m[:, None] * D_HEAD + off_d[None, :], mask=off_m[:, None] < N_CTX).to(tl.float32)
do = tl.load(DO + off_m[:, None] * D_HEAD + off_d[None, :], mask=off_m[:, None] < N_CTX).to(tl.float32)
denom = tl.load(L + off_m).to(tl.float32)
# compute
do = do / denom[:, None]
delta = tl.sum(o * do, axis=1)
# write-back
if EVEN_M_BLOCK:
tl.store(NewDO + off_m[:, None] * D_HEAD + off_d[None, :], do)
else:
tl.store(NewDO + off_m[:, None] * D_HEAD + off_d[None, :], do, mask=off_m[:, None] < N_CTX)
tl.store(Delta + off_m, delta)
# Does not suuport unequal seqlen(q) and seqlen(k)
@triton.heuristics(
{
'EVEN_M_BLOCK': lambda kwargs: kwargs['N_CTX'] % kwargs['BLOCK_M'] == 0,
'EVEN_N_BLOCK': lambda kwargs: kwargs['N_CTX'] % kwargs['BLOCK_N'] == 0,
}
)
@triton.jit
def _bwd_kernel(
Q, K, V, sm_scale,
layout_ccol_ptr,
layout_row_ptr,
layout_ccol_stride_h, layout_ccol_stride_m,
layout_row_stride_h, layout_row_stride_m,
Out, DO, # assume contigous: Out, Do, DQ, DK, DV, L, M, D, seq(q) == seq(k), with stride_oz, stride_oh, stride_om, stride_od,
DQ, DK, DV,
L, M,
D,
stride_qz, stride_qh, stride_qm, stride_qd,
stride_kz, stride_kh, stride_kn, stride_kd,
stride_vz, stride_vh, stride_vn, stride_vd,
stride_oz, stride_oh, stride_om, stride_od,
# stride_dz, stride_dh, stride_dm, stride_dd,
Z, H, N_CTX,
num_block,
BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr,
BLOCK_N: tl.constexpr,
EVEN_M_BLOCK: tl.constexpr,
EVEN_N_BLOCK: tl.constexpr,
NUM_DBLOCKS: tl.constexpr,
):
start_n = tl.program_id(0)
off_hz = tl.program_id(1)
off_z = off_hz // H
off_h = off_hz % H
# offset pointers for batch/head
Q += off_z * stride_qz + off_h * stride_qh
K += off_z * stride_kz + off_h * stride_kh
V += off_z * stride_vz + off_h * stride_vh
DO += off_z * stride_oz + off_h * stride_oh
DQ += off_z * stride_oz + off_h * stride_oh
DK += off_z * stride_oz + off_h * stride_oh
DV += off_z * stride_oz + off_h * stride_oh
# Look like this loop can be parallelled
# for start_n in range(0, num_block):
offs_n = start_n * BLOCK_N + tl.arange(0, BLOCK_N)
offs_m = tl.arange(0, BLOCK_M)
offs_d = tl.arange(0, BLOCK_DMODEL)
# initialize pointers to value-like data
k_ptrs = K + (offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kd)
v_ptrs = V + (offs_n[:, None] * stride_vn + offs_d[None, :] * stride_vd)
# pointer to row-wise quantities in value-like data
D_ptrs = D + off_hz * N_CTX
m_ptrs = M + off_hz * N_CTX
# initialize dv amd dk
dv = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32)
dk = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32)
# k and v stay in SRAM throughout
if EVEN_N_BLOCK:
k = tl.load(k_ptrs)
v = tl.load(v_ptrs)
else:
k = tl.load(k_ptrs, mask=offs_n[:, None] < N_CTX)
v = tl.load(v_ptrs, mask=offs_n[:, None] < N_CTX)
if NUM_DBLOCKS >= 2:
dv2 = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32)
dk2 = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32)
if EVEN_N_BLOCK:
k2 = tl.load(k_ptrs + BLOCK_DMODEL * stride_kd)
v2 = tl.load(v_ptrs + BLOCK_DMODEL * stride_vd)
else:
k2 = tl.load(k_ptrs + BLOCK_DMODEL * stride_kd, mask=offs_n[:, None] < N_CTX)
v2 = tl.load(v_ptrs + BLOCK_DMODEL * stride_vd, mask=offs_n[:, None] < N_CTX)
# loop over rows
layout_ptr = layout_ccol_ptr + off_h * layout_ccol_stride_h + start_n * layout_ccol_stride_m
start_l = tl.load(layout_ptr).to(tl.int32)
end_l = tl.load(layout_ptr + layout_ccol_stride_m).to(tl.int32)
for row_idx_idx in range(start_l, end_l):
row_idx = tl.load(layout_row_ptr + off_h * layout_row_stride_h + row_idx_idx * layout_row_stride_m).to(tl.int32)
start_m = row_idx * BLOCK_M
# offs_qm = start_m + tl.arange(0, BLOCK_M)
offs_m_curr = start_m + offs_m
q_ptrs = Q + (offs_m_curr[:, None] * stride_qm + offs_d[None, :] * stride_qd)
do_ptrs = DO + (offs_m_curr[:, None] * stride_om + offs_d[None, :] * stride_od)
dq_ptrs = DQ + (offs_m_curr[:, None] * stride_om + offs_d[None, :] * stride_od)
# load q, k, v, do on-chip
if EVEN_M_BLOCK:
q = tl.load(q_ptrs)
else:
q = tl.load(q_ptrs, mask=offs_m_curr[:, None] < N_CTX)
# re-compute p = softmax(qk, dim=-1).T
# NOTE: `do` is pre-divided by `l`; no normalization here
qk = tl.dot(q, tl.trans(k))
if NUM_DBLOCKS >= 2:
if EVEN_M_BLOCK:
q2 = tl.load(q_ptrs + BLOCK_DMODEL * stride_qd)
else:
q2 = tl.load(q_ptrs + BLOCK_DMODEL * stride_qd, mask=offs_m_curr[:, None] < N_CTX)
qk += tl.dot(q2, tl.trans(k2))
qk += tl.where(offs_m_curr[:, None] >= (offs_n[None, :]), 0, float('-inf'))
if EVEN_M_BLOCK:
m = tl.load(m_ptrs + offs_m_curr)
else:
m = tl.load(m_ptrs + offs_m_curr, mask=offs_m_curr < N_CTX)
p = tl.exp(qk * sm_scale - m[:, None])
# compute dv
if EVEN_M_BLOCK:
do = tl.load(do_ptrs)
else:
do = tl.load(do_ptrs, mask=offs_m_curr[:, None] < N_CTX)
if NUM_DBLOCKS >= 2:
if EVEN_M_BLOCK:
do2 = tl.load(do_ptrs + BLOCK_DMODEL * stride_od)
else:
do2 = tl.load(do_ptrs + BLOCK_DMODEL * stride_od, mask=offs_m_curr[:, None] < N_CTX)
dv += tl.dot(tl.trans(p.to(Q.dtype.element_ty)), do)
if NUM_DBLOCKS >= 2:
dv2 += tl.dot(tl.trans(p.to(Q.dtype.element_ty)), do2)
# compute dp = dot(v, do)
if EVEN_M_BLOCK:
Di = tl.load(D_ptrs + offs_m_curr)
else:
Di = tl.load(D_ptrs + offs_m_curr, mask=offs_m_curr < N_CTX)
dp = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) - Di[:, None]
dp += tl.dot(do, tl.trans(v))
if NUM_DBLOCKS >= 2:
dp += tl.dot(do2, tl.trans(v2))
# compute ds = p * (dp - delta[:, None])
ds = p * dp * sm_scale
# compute dk = dot(ds.T, q)
dk += tl.dot(tl.trans(ds.to(Q.dtype.element_ty)), q)
if NUM_DBLOCKS >= 2:
dk2 += tl.dot(tl.trans(ds.to(Q.dtype.element_ty)), q2)
# # compute dq
dq = tl.dot(ds.to(Q.dtype.element_ty), k)
if EVEN_M_BLOCK:
tl.atomic_add(dq_ptrs, dq)
else:
tl.atomic_add(dq_ptrs, dq, mask=offs_m_curr[:, None] < N_CTX)
if NUM_DBLOCKS >= 2:
dq2 = tl.dot(ds.to(Q.dtype.element_ty), k2)
dq_ptrs2 = dq_ptrs + BLOCK_DMODEL * stride_od
if EVEN_M_BLOCK:
tl.atomic_add(dq_ptrs2, dq2)
else:
tl.atomic_add(dq_ptrs2, dq2, mask=offs_m_curr[:, None] < N_CTX)
# write-back
dv_ptrs = DV + (offs_n[:, None] * stride_om + offs_d[None, :] * stride_od)
dk_ptrs = DK + (offs_n[:, None] * stride_om + offs_d[None, :] * stride_od)
if EVEN_N_BLOCK:
tl.store(dv_ptrs, dv)
tl.store(dk_ptrs, dk)
else:
tl.store(dv_ptrs, dv, mask=offs_n[:, None] < N_CTX)
tl.store(dk_ptrs, dk, mask=offs_n[:, None] < N_CTX)
if NUM_DBLOCKS >= 2:
dv_ptrs2 = dv_ptrs + BLOCK_DMODEL * stride_od
dk_ptrs2 = dk_ptrs + BLOCK_DMODEL * stride_od
if EVEN_N_BLOCK:
tl.store(dv_ptrs2, dv2)
tl.store(dk_ptrs2, dk2)
else:
tl.store(dv_ptrs2, dv2, mask=offs_n[:, None] < N_CTX)
tl.store(dk_ptrs2, dk2, mask=offs_n[:, None] < N_CTX)
def _forward(ctx, q, k, v, layout_crow_indices, layout_col_indices, sm_scale, BLOCK_M, BLOCK_N, num_warps=None, num_stages=1, inference=None, out=None):
'''
:param q, k, v: [batch, n_heads, seq_len, model_dim]. len of q is allowed to be different than k/v.
:param layout_crow_indices, layout_col_indices: same as CSR.crow_indices, and CSR.col_indices used to preresent a sparse tensor.
Each element represent a block, i.e, all elements in a block to be attentdd, or not attended at all..
'''
assert q.shape[-1] == k.shape[-1] == v.shape[-1]
assert k.shape[2] == v.shape[2]
o = out if out is not None else torch.empty_like(q).contiguous()
grid = (triton.cdiv(q.shape[2], BLOCK_M), q.shape[0] * q.shape[1])
q_rounded_len = grid[0] * BLOCK_M
tmp = torch.empty((q.shape[0] * q.shape[1], q_rounded_len), device=q.device, dtype=torch.float32)
if inference is None:
inference = (not q.requires_grad) and (not k.requires_grad) and (not v.requires_grad)
if inference:
L, m = tmp, tmp # no need to use create new tensor
else:
L = torch.empty((q.shape[0] * q.shape[1], q_rounded_len), device=q.device, dtype=torch.float32)
m = torch.empty((q.shape[0] * q.shape[1], q_rounded_len), device=q.device, dtype=torch.float32)
if layout_col_indices.dim() == 1:
layout_crow_indices = layout_crow_indices[None].expand(q.shape[1] , -1)
layout_col_indices = layout_col_indices[None].expand(q.shape[1] , -1)
assert q.shape[-1] in [64, 128]
BLOCK_DMODEL = 64
if num_warps is None:
MIN_D = min(BLOCK_M, BLOCK_N, BLOCK_DMODEL)
num_warps = max(1, 2 ** int(math.log2(MIN_D / 16)))
# print(f'> {BLOCK_M=}, {BLOCK_N=}, {BLOCK_DMODEL=}, {num_warps=}, {num_stages=}')
else:
assert math.log2(num_warps) % 1 == 0, f'''"num_warps" should be power of 2, but got {num_warps}.'''
## For debugging:
# print(f'>> {q.shape=}, {k.shape=}, {BLOCK_M=}, {BLOCK_N=}, {num_warps=}, {BLOCK_DMODEL=}, {q.stride()=}, {k.stride()=}')
# print(f'>> {layout_crow_indices=}\n{layout_col_indices=}\n {layout_crow_indices.stride()=}, {layout_crow_indices.stride()=}')
# print(f'> {q.shape=}, {k.shape=}, {layout_crow_indices.shape}, {layout_col_indices.shape}, {layout_crow_indices.stride()}, \
# {layout_col_indices.stride()}, {layout_crow_indices=}, {layout_col_indices=}')
with torch.cuda.device(q.device.index):
_fwd_kernel[grid](
q, k, v, sm_scale,
layout_crow_indices,
layout_col_indices,
layout_crow_indices.stride(0), layout_crow_indices.stride(1),
layout_col_indices.stride(0), layout_col_indices.stride(1),
tmp, L, 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),
q.shape[0], q.shape[1], k.shape[2],
k.shape[2] - q.shape[2],
q_rounded_len,
BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N,
BLOCK_DMODEL=BLOCK_DMODEL,
EVEN_M_BLOCK=q.shape[2] % BLOCK_M == 0,
EVEN_N_BLOCK=k.shape[2] % BLOCK_N == 0 ,
INFERENCE=inference,
NUM_DBLOCKS=q.shape[-1] // BLOCK_DMODEL,
num_warps=num_warps,
num_stages=num_stages,
)
if inference:
L, m = None, None
ctx.save_for_backward(q, k, v, o, L, m, layout_crow_indices, layout_col_indices)
ctx.BLOCK_M = BLOCK_M
ctx.BLOCK_N = BLOCK_N
ctx.BLOCK_DMODEL = BLOCK_DMODEL
# ctx.BLOCK = BLOCK
ctx.grid = grid
ctx.sm_scale = sm_scale
ctx.num_warps = num_warps
ctx.num_stages = num_stages
return o
def _backward(ctx, do, layout_ccol_indices, layout_row_indices, dq=None, dk=None, dv=None):
# q, k, v, o, l, m = ctx.saved_tensors
q, k, v, o, l, m, layout_crow_indices, layout_col_indices = ctx.saved_tensors
## this following too slow to do online, so get it from inputs, which is cached.
# layout_ccol_indices, layout_row_indices = dense_to_ccol_row(crow_col_to_dense(ctx.layout_crow_indices, ctx.layout_col_indices))
# layout_ccol_indices, layout_row_indices = dense_to_ccol_row(crow_col_to_dense(layout_crow_indices, layout_col_indices))
if not do.is_contiguous():
do = do.contiguous()
## for debugging
# print(f'----> do is not contiguous: {do.stride()=}')
# raise ValueError(f'>>>> output grad is not contiguous: {do.stride()=}')
if not o.is_contiguous():
# TODO: currently only work with contiguous q/k/v.
raise ValueError(f'--> output is not contiguous: {o.stride()=}. This is maybe caused by q/k/v not being contiguous.')
if layout_ccol_indices.dim() == 1:
layout_ccol_indices = layout_ccol_indices[None].expand(q.shape[1], -1)
layout_row_indices = layout_row_indices[None].expand(q.shape[1], -1)
# do = do.contiguous()
dq = dq if dq is not None else torch.zeros_like(q, dtype=torch.float32)
dk = dk if dk is not None else torch.empty_like(k)
dv =dv if dv is not None else torch.empty_like(v)
do_scaled = torch.empty_like(do)
delta = torch.empty_like(l)
assert o.stride() == dq.stride() == dk.stride() == dv.stride() == do_scaled.stride()
_bwd_preprocess[(ctx.grid[0] * ctx.grid[1], )](
o, do, l,
do_scaled, delta,
k.shape[2],
BLOCK_M=ctx.BLOCK_M, D_HEAD=q.shape[-1],
)
grid = (triton.cdiv(q.shape[2], ctx.BLOCK_N), ctx.grid[1])
_bwd_kernel[grid](
q, k, v, ctx.sm_scale,
layout_ccol_indices,
layout_row_indices,
layout_ccol_indices.stride(0), layout_ccol_indices.stride(1),
layout_row_indices.stride(0), layout_row_indices.stride(1),
o, do_scaled,
dq, dk, dv,
l, m,
delta,
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),
q.shape[0], q.shape[1], q.shape[2],
ctx.grid[0],
BLOCK_M=ctx.BLOCK_M,
BLOCK_N=ctx.BLOCK_N,
BLOCK_DMODEL=ctx.BLOCK_DMODEL,
NUM_DBLOCKS=q.shape[-1] // ctx.BLOCK_DMODEL,
num_warps=ctx.num_warps,
num_stages=1,
)
return dq, dk, dv, None, None, None
class _sparse_attention(torch.autograd.Function):
@staticmethod
def forward(ctx, q, k, v, layout_crow_indices, layout_col_indices, sm_scale):
BLOCK = 128
# shape constraints
return _forward(ctx, q, k, v, layout_crow_indices, layout_col_indices, sm_scale, BLOCK, BLOCK)
@staticmethod
def backward(ctx, do):
# q, k, v, o, l, m = ctx.saved_tensors
q, k, v, o, l, m, layout_crow_indices, layout_col_indices = ctx.saved_tensors
# TODO: the following is very inefficient.
# layout_ccol_indices, layout_row_indices = dense_to_ccol_row(crow_col_to_dense(ctx.layout_crow_indices, ctx.layout_col_indices))
layout_ccol_indices, layout_row_indices = dense_to_ccol_row(crow_col_to_dense(layout_crow_indices, layout_col_indices))
return _backward(ctx, do, layout_ccol_indices, layout_row_indices)
# suppressed
class _sparse_attention_inference(_sparse_attention):
# TODO: does not work now, as BLOCK_M cannot be <1, as shape for tl.dot cannot be smaller than 16.
@staticmethod
def forward(ctx, q, k, v, layout_crow_indices, layout_col_indices, sm_scale):
BLOCK = 128
return _forward(ctx, q, k, v, layout_crow_indices, layout_col_indices, sm_scale, 1, BLOCK)
def sparse_attention_factory(BLOCK_M=128, BLOCK_N=128, **kwargs):
class _sparse_attention_config(_sparse_attention):
@staticmethod
def forward(ctx, q, k, v, layout_crow_indices, layout_col_indices, sm_scale):
# shape constraints
return _forward(ctx, q, k, v, layout_crow_indices, layout_col_indices, sm_scale, BLOCK_M, BLOCK_N,
**kwargs
)
return _sparse_attention_config.apply
@lru_cache(maxsize=8)
def get_local_strided_sparse_attention_op(
n_heads: int,
max_seq_len:int,
sparse_block_size: int=128,
local_blocks: int=4,
vert_stride: int=4,
homo_head: bool=False,
dtype=torch.bfloat16,
device='cuda',
active_head_range=None,
verbose=True,
**kwargs):
'''
:param n_heads: total number of attention heads (regardless of tensor/model parallel)
:param max_seq_len: max sequence length. Need to be bigger or equal to the length of sequences.
:param sparse_block_size: sparse block size. Default to 128
:param local_blocks: number of nearest block to attend to. Default to 4, i.e., attention to previous 4xblock_size tokens.
:param vert_stride: Default to 4. Meaning
:param homo_head: if all head shared the same pattern.
:param active_head_range: tuple of start & end of the heads, e..g, (8, 16). Default to use all heads.
Mainly for tensor/model parallelization where heads are splitted to different GPUs.
'''
if verbose:
print((f'> new block_sparse_attn op constructed with config: '
f'{n_heads=}, {max_seq_len=}, {sparse_block_size=}, {local_blocks=}, '
f'{vert_stride=}, {homo_head=}, {active_head_range=}, {kwargs=}'))
# assert math.log2(max_seq_len) % 2 == 0, f"max_seq_len should be power of 2 to be more efficient"
_, block_sparse_pattern, _ = _get_sparse_attn_mask(n_heads, max_seq_len, max_seq_len, dtype, device,
BLOCK=sparse_block_size, local_blocks=local_blocks,
vert_stride=vert_stride, homo_head=homo_head,
return_dense=False)
if (not homo_head) and (active_head_range is not None):
assert isinstance(active_head_range, tuple)
assert len(active_head_range) == 2, '"active_head_range" should be a tuple of start/end index of the heads.'
h_start, h_end = active_head_range
block_sparse_pattern = block_sparse_pattern[h_start:h_end]
# print(block_sparse_pattern)
return get_sparse_attn_op(block_sparse_pattern, sparse_block_size, **kwargs)
def get_sparse_attn_op(
sparse_pattern: torch.tensor,
sparse_block_size: int=128,
kernel_block_size=128,
qkv_format='q,k,v',
**kwargs):
'''
Ccreate a block-sparse op with fixed layout. This is to avoid the need to of create CSR layout and convert it to CSC layout everytime,
which is very inefficient (use python loops on CPU. PyTorch 1.13 supports CSR->CSC, may help.)
:param sparse_pattern: sparse pattern of the blocks. Should be `num_blocks(q) x num_blocks(k)` or `n_heads x num_blocks x num_blocks`.
This tensor should have lower-triangular matrices on the last 2 dimensions for causal attention
:param sparse_block_size: sparse block size. Default to 128
:param kernel_block_size: the tile/block size to launch a triton instance. Default to None, i.e., same as `sparse_block_size`
:param qkv_format: Choices=['q,k,v', 'q, kv', 'qkv'], i.e., separated q,k,v, or kv packed, or qkv packed. Currently, only 'q,k,v' is supported.
:param kwargs: keyward arguments passed to `_forward`
'''
# assert qkv_format in ('q,k,v', 'q, kv', 'qkv') # to save from running `concat` at forward/backward
assert qkv_format == 'q,k,v'
if kernel_block_size is None:
kernel_block_size = sparse_block_size
else:
assert sparse_block_size % kernel_block_size == 0, f"The sparse block size must be a multiple of {kernel_block_size}."
assert kernel_block_size >=16 and math.log2(kernel_block_size) % 1 == 0, f"block_size must be power of 2 and at least 16, but {kernel_block_size} is given"
# print(f'>> {sparse_pattern.shape=}')
# print(f'{sparse_pattern=}')
if sparse_block_size // kernel_block_size > 1:
_mul = sparse_block_size // kernel_block_size
# need to consider if block_m and block_n are different
sparse_pattern = torch.kron(sparse_pattern, sparse_pattern.new_ones(_mul, _mul))
num_sparse_blocks = sparse_pattern.size(-1)
block_causal_mask = torch.arange(0, num_sparse_blocks)[:, None] >= torch.arange(0, num_sparse_blocks)[None]
sparse_pattern *= block_causal_mask.type_as(sparse_pattern)
# print(f'>> after: {sparse_pattern.shape=}')
# print(f'{sparse_pattern=}')
BLOCK_N = kernel_block_size
NUM_BLOCK = sparse_pattern.size(-1)
MAX_SEQ_LEN = kernel_block_size * NUM_BLOCK
grand_layout_crow_indices, grand_layout_col_indices = dense_to_crow_col(sparse_pattern)
# sparse csc layout for backward
grand_layout_ccol_indices, grand_layout_row_indices = dense_to_ccol_row(sparse_pattern)
# cache GPU backward layout. limit the size to avoid OOM as time goes.
# For inference, one only needs to cache one block as sequence length always increases
# Therefore, this cache needs to be reconstructed per every `block_size`-steps.
# For training/finetune, set to 8 to increase cache hit.
# Given an input, the block_len will be the same for all layers, so cache is very helpful.
max_cache_size = 1 if kwargs.get('inference', False) else 8
@lru_cache(maxsize=max_cache_size)
def get_backward_layout_by_block_len(block_len):
assert block_len <= NUM_BLOCK
if block_len == NUM_BLOCK:
return (grand_layout_ccol_indices, grand_layout_row_indices)
return dense_to_ccol_row(sparse_pattern[..., :block_len, :block_len])
# for debugging
# if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0:
# print(f'> {sparse_pattern.cpu().tolist()=}')
# print('----')
# print(f'> {grand_layout_crow_indices.cpu().tolist()=}\n{grand_layout_col_indices.cpu().tolist()=}')
# q, k, v separated
class _q_k_v_sparse_attention(torch.autograd.Function):
@staticmethod
def forward(ctx, q, k, v, sm_scale):
# assert q.shape[2] == 1 or q.shape[2] == k.shape[2]
# shape constraints
MIN_BLOCK_SIZE = 16
assert BLOCK_N >= MIN_BLOCK_SIZE
BLOCK_M = 16 if q.shape[2] <= 16 else BLOCK_N # BLOCK_M has to be power of 2
# this following code only works for causal attention
K_BLOCKS = triton.cdiv(k.shape[2], kernel_block_size)
# Q_START_BLOCKS = K_BLOCKS - 1 if q.shape[2] == 1 else 0
Q_START_BLOCKS = K_BLOCKS - triton.cdiv(q.shape[2], BLOCK_N)
# print(Q_START_BLOCKS, K_BLOCKS)
layout_crow_indices = grand_layout_crow_indices[..., Q_START_BLOCKS:K_BLOCKS+1]
layout_col_indices = grand_layout_col_indices
# print(BLOCK_M, BLOCK_N, Q_START_BLOCKS, K_BLOCKS+1, layout_crow_indices, layout_col_indices)
return _forward(ctx, q, k, v, layout_crow_indices, layout_col_indices, sm_scale, BLOCK_M, BLOCK_N,
**kwargs
)
@staticmethod
def backward(ctx, do):
q, k = ctx.saved_tensors[:2]
assert q.shape[2] == k.shape[2], '> currently backward can only be done if q, k have same length. Contact @EricLin if you need it.'
# assume q, k have same length
block_len = triton.cdiv(do.shape[2], kernel_block_size)
backward_layout = get_backward_layout_by_block_len(block_len)
return _backward(ctx, do, *backward_layout)[:4]
def _q_k_v_sparse_attention_fn(*args):
return _q_k_v_sparse_attention.apply(*args)
_q_k_v_sparse_attention_fn.sparse_pattern = sparse_pattern
_q_k_v_sparse_attention_fn.grand_layout_crow_indices = grand_layout_crow_indices
_q_k_v_sparse_attention_fn.grand_layout_col_indices = grand_layout_col_indices
_q_k_v_sparse_attention_fn.grand_layout_ccol_indices = grand_layout_ccol_indices
_q_k_v_sparse_attention_fn.grand_layout_row_indices = grand_layout_row_indices
return _q_k_v_sparse_attention_fn
###########################################################
###########################################################
###########################################################
################ Inference Kernels ########################
###########################################################
def blocksparse_flash_attn_padded_fwd(
q, k, v, # (batch, tokens, n_heads, head_size)
sm_scale,
sparse_layout,
*,
left_paddings = None,
seqlens = None,
block_size = 64,
max_seqlen = None
):
'''
q, k, v: (batch, tokens, n_heads/n_kv_heads, head_size)
left_paddings: (batch, ), number of left paddings for each sample.
seqlens: can be used to specify right padding. No need to specify if left_paddings is used.
'''
batches, q_len, n_heads, head_size = q.shape
_, k_len, n_kv_heads, _ = k.shape
assert q.dim() == k.dim() == v.dim() == 4
assert q.size(2) % k.size(2) == 0
assert q.size(0) == k.size(0) and q.size(3) == k.size(3)
assert k.shape == v.shape # TODO: allow diff head_size for k, v
assert q_len == 1 or q_len == k_len, \
f'q length can only 1 for decoding for same as k length for prefilling.'
q_k_ratio = q.size(2) // k.size(2)
if max_seqlen:
assert k.size(1) <= max_seqlen, f'k has seqlen {k.size(1)} while max sequence length is set to {max_seqlen}.'
# paddings always has zero output, a little slower than using empty
out = q.new_zeros(q.shape)
layout_crow_indices, layout_col_indices = sparse_layout
block_d = triton.next_power_of_2(head_size)
if left_paddings is not None:
assert left_paddings.shape == (batches,)
k_batch_starts = left_paddings.to(q.device, dtype=torch.int32).contiguous()
else:
k_batch_starts = torch.zeros((batches,), dtype=torch.int32, device=q.device)
if seqlens is not None:
k_batch_ends = k_batch_starts + seqlens.type_as(k_batch_starts)
assert k_batch_ends.max() <= k_len, f'seqlens (+left_paddings if any) exceeds seqlen.'
else:
k_batch_ends = torch.zeros_like(k_batch_starts) + k_len
if q_len == 1:
q_batch_starts = torch.zeros_like(k_batch_starts)
q_batch_ends = q_batch_starts + 1
else:
q_batch_starts = k_batch_starts
q_batch_ends = k_batch_ends
# switch to use cpu to avoid too many kernel lauch when iterate over
q_lens = (q_batch_ends - q_batch_starts).cpu()
n_blocks = (q_lens + block_size - 1) // block_size
q_batch_ids = torch.tensor([i for i, n in enumerate(n_blocks) for _ in range(n)],
dtype=q_batch_starts.dtype,
device=q_batch_starts.device)
q_start_sids = torch.tensor([i * block_size for n in n_blocks for i in range(n)],
dtype=q_batch_starts.dtype,
device=q_batch_starts.device)
grid = (len(q_start_sids), n_heads)
with torch.cuda.device(q.device.index):
_fwd_kernel_batch_inference[grid](
q, k, v, out,
sm_scale,
q_batch_starts,
q_batch_ends,
k_batch_starts,
k_batch_ends,
q_batch_ids,
q_start_sids,
*q.stride(),
*k.stride(),
*v.stride(),
*out.stride(),
layout_crow_indices,
layout_col_indices,
*layout_crow_indices.stride(),
*layout_col_indices.stride(),
q_k_ratio,
HAS_BATCH_DIM = True,
D_HEAD = head_size,
BLOCK_M = block_size,
BLOCK_N = block_size,
BLOCK_D = block_d,
BLOCK_M_LOADING = 16 if q_len == 1 else block_size, # smaller for decoding
EVEN_D = block_d == head_size,
num_warps = 1 if q_len == 1 else 4,
num_stages = 1
)
return out
def blocksparse_flash_attn_varlen_fwd(
q, k, v, # (#tokens, n_heads, head_size)
cu_seqlens_k,
cu_seqlens_q,
sm_scale,
sparse_layout,
*,
block_size=64,
max_seqlen = None
):
# split q to blocks
_, n_heads, head_size = q.shape
batch_size = cu_seqlens_k.size(0) - 1
# print(f'> {q.shape=}, {k.shape=}')
assert q.dim() == k.dim() == v.dim() == 3
assert q.size(1) % k.size(1) == 0
assert q.size(2) == k.size(2)
assert k.shape == v.shape # TODO: allow diff head_size for k, v
assert cu_seqlens_k.dim() == 1
q_k_ratio = q.size(1) // k.size(1)
if cu_seqlens_q is None:
if q.size(0) == batch_size: # decoding only
cu_seqlens_q = torch.arange(0, batch_size + 1,
dtype=cu_seqlens_k.dtype,
device=cu_seqlens_k.device)
elif q.size(0) == k.size(0):
cu_seqlens_q = cu_seqlens_k
else:
raise ValueError('cu_seqlens_q must be specified if it is mix of prefilling and decoding.')
else:
assert cu_seqlens_k.size(0) == cu_seqlens_q.size(0)
# switch to use cpu to avoid too many kernel lauch when iterate over
q_lens = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).cpu()
k_lens = (cu_seqlens_k[1:] - cu_seqlens_k[:-1]).cpu()
assert torch.logical_or(q_lens == 1, k_lens == q_lens).all(), \
'length of q should either be 1 (decoding) or same as k (prefilling).'
if max_seqlen:
assert k_lens.max() <= max_seqlen
n_blocks = (q_lens + block_size - 1) // block_size
q_batch_ids = torch.tensor([i for i, n in enumerate(n_blocks) for _ in range(n)],
dtype=cu_seqlens_q.dtype,
device=cu_seqlens_q.device)
q_start_sids = torch.tensor([i * block_size for n in n_blocks for i in range(n)],
dtype=cu_seqlens_q.dtype,
device=cu_seqlens_q.device)
out = q.new_empty(q.shape)
cu_seqlens_q = cu_seqlens_q.contiguous()
cu_seqlens_k = cu_seqlens_k.contiguous()
layout_crow_indices, layout_col_indices = sparse_layout
block_d = triton.next_power_of_2(head_size)
decoding_only = (q_lens == 1).all()
grid = (len(q_start_sids), n_heads)
with torch.cuda.device(q.device.index):
_fwd_kernel_batch_inference[grid](
q, k, v, out,
sm_scale,
cu_seqlens_q[:-1],
cu_seqlens_q[1:],
cu_seqlens_k[:-1],
cu_seqlens_k[1:],
q_batch_ids,
q_start_sids,
0, *q.stride(),
0, *k.stride(),
0, *v.stride(),
0, *out.stride(),
layout_crow_indices,
layout_col_indices,
*layout_crow_indices.stride(),
*layout_col_indices.stride(),
q_k_ratio,
HAS_BATCH_DIM = False,
D_HEAD = head_size,
BLOCK_M = block_size,
BLOCK_N = block_size,
BLOCK_D = block_d,
BLOCK_M_LOADING = 16 if decoding_only else block_size, # smaller for decoding
EVEN_D = block_d == head_size,
num_warps = 1 if decoding_only else 4,
num_stages = 3
)
return out
@triton.jit
def _fwd_kernel_inner(
acc, l_i, m_i,
q, Q,
k_block_col_idx,
layout_col_ptr,
layout_col_stride_h, layout_col_stride_m,
k_ptrs,
v_ptrs,
off_h, offs_m, offs_n, offs_d,
stride_kt, stride_vt,
sm_scale,
k_seqlen,
past_len,
LAST_K_BLOCK: tl.constexpr,
BLOCK_M_LOADING: tl.constexpr,
BLOCK_N: tl.constexpr,
D_HEAD: tl.constexpr,
EVEN_D: tl.constexpr,
M_LT_N: tl.constexpr
):
k_block_id = tl.load(layout_col_ptr + off_h * layout_col_stride_h + k_block_col_idx * layout_col_stride_m).to(tl.int32)
start_n = k_block_id * BLOCK_N
# -- compute qk ----
if LAST_K_BLOCK:
if EVEN_D:
k = tl.load(k_ptrs + start_n * stride_kt,
mask=offs_n[None, :] + start_n < k_seqlen)
else:
# mask = mask & (offs_d[:, ])
k = tl.load(k_ptrs + start_n * stride_kt,
mask=(offs_n[None, :] + start_n < k_seqlen) & (offs_d[:, None] < D_HEAD))
else:
if EVEN_D:
k = tl.load(k_ptrs + start_n * stride_kt)
else:
k = tl.load(k_ptrs + start_n * stride_kt,
mask=offs_d[:, None] < D_HEAD)
qk = tl.zeros([BLOCK_M_LOADING, BLOCK_N], dtype=tl.float32)
qk += tl.dot(q, k)
qk *= sm_scale
# the following is needed only when LAST_K_BLOCK or BLOCK_M < BLOCK_N
if LAST_K_BLOCK | M_LT_N:
qk += tl.where(offs_m[:, None] + past_len >= (start_n + offs_n[None, :]), 0, float('-inf'))
# -- compute m_ij, p, l_ij
m_ij = tl.max(qk, 1)
p = tl.exp(qk - m_ij[:, None])
l_ij = tl.sum(p, 1)
# -- update m_i and l_i
m_i_new = tl.maximum(m_i, m_ij)
alpha = tl.exp(m_i - m_i_new)
beta = tl.exp(m_ij - m_i_new)
l_i_new = alpha * l_i + beta * l_ij
# -- update output accumulator --
# scale p
p_scale = beta / l_i_new
p = p * p_scale[:, None]
# scale acc
acc_scale = l_i / l_i_new * alpha
acc = acc * acc_scale[:, None]
p = p.to(Q.dtype.element_ty)
# update acc
if LAST_K_BLOCK:
if EVEN_D:
v = tl.load(v_ptrs + start_n * stride_vt,
mask=offs_n[:, None] + start_n < k_seqlen)
else:
v = tl.load(v_ptrs + start_n * stride_vt,
mask=(offs_n[:, None] + start_n < k_seqlen) & (offs_d[None, :] < D_HEAD))
else:
if EVEN_D:
v = tl.load(v_ptrs + start_n * stride_vt)
else:
v = tl.load(v_ptrs + start_n * stride_vt,
mask=offs_d[None, :] < D_HEAD)
acc += tl.dot(p, v)
# update m_i and l_i
l_i = l_i_new
m_i = m_i_new
return acc, l_i, m_i
@triton.heuristics(
{
'M_LT_N': lambda kwargs: kwargs['BLOCK_M'] < kwargs['BLOCK_N'],
}
)
@triton.jit
def _fwd_kernel_batch_inference(
Q, K, V, Out,
sm_scale,
q_batch_starts,
q_batch_ends,
k_batch_starts,
k_batch_ends,
q_batch_ids,
q_start_sids,
stride_qb, stride_qt, stride_qh, stride_qd,
stride_kb, stride_kt, stride_kh, stride_kd,
stride_vb, stride_vt, stride_vh, stride_vd,
stride_ob, stride_ot, stride_oh, stride_od,
layout_crow_ptr,
layout_col_ptr,
layout_crow_stride_h, layout_crow_stride_m,
layout_col_stride_h, layout_col_stride_m,
q_k_ratio,
HAS_BATCH_DIM: tl.constexpr,
D_HEAD: tl.constexpr,
BLOCK_M: tl.constexpr,
BLOCK_N: tl.constexpr,
BLOCK_D: tl.constexpr,
BLOCK_M_LOADING: tl.constexpr,
EVEN_D: tl.constexpr,
M_LT_N: tl.constexpr
):
'''
NOTATION:
pid: position id
sid: storage id
sbid: storage block id
pbid: position block id
offs_m, offs_n: storage offsets of m-dim(q, row) and n-dim(k, col)
q and blocks in KV needs to be contiguous
Arguments:
kv_seq_lens: for compute past_len
kv_storage_offsets: similar to block_tables in vllm, except it is dynamic.
TODO: fix this
TODO:
Optimize grouped-attn
CUDA graph support issue
1. grid is dynamic: vllm set up multiple cuda graph in decoding phase, with diff max token size (16, 32, ...)
since we mix prompt and decoing phase here, it can be more complex.
need to set up diff cuda-graph for diff (off_zm, off_z)
# indeed, q_batch_ids can be padded to maximum number of grid[0], i.e., assume all decoding
therefore, cu_seqlens_q, kv_seq_lens
'''
off_zm = tl.program_id(0)
off_h = tl.program_id(1)
off_h_for_kv = off_h // q_k_ratio
off_z = tl.load(q_batch_ids + off_zm).to(tl.int32) # [0, 0, 0, 1]
q_start_sid = tl.load(q_start_sids + off_zm)
start_m = q_start_sid // BLOCK_M
if HAS_BATCH_DIM:
Q += off_z * stride_qb
K += off_z * stride_kb
V += off_z * stride_vb
Out += off_z * stride_ob
offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M_LOADING)
offs_n = tl.arange(0, BLOCK_N)
offs_d = tl.arange(0, BLOCK_D)
q_cu_start = tl.load(q_batch_starts + off_z).to(tl.int32)
q_seqlen = tl.load(q_batch_ends + off_z).to(tl.int32) - q_cu_start
k_cu_start = tl.load(k_batch_starts + off_z).to(tl.int32)
k_seqlen = tl.load(k_batch_ends + off_z).to(tl.int32) - k_cu_start
past_len = k_seqlen - q_seqlen
Q += q_cu_start * stride_qt + off_h * stride_qh
K += k_cu_start * stride_kt + off_h_for_kv * stride_kh
V += k_cu_start * stride_vt + off_h_for_kv * stride_vh
Out += q_cu_start * stride_ot + off_h * stride_oh
q_pbid = (past_len + q_start_sid) // BLOCK_M
if EVEN_D:
q = tl.load(Q + offs_m[:, None] * stride_qt + offs_d[None, :] * stride_qd,
mask=offs_m[:, None] < q_seqlen)
else:
q = tl.load(Q + offs_m[:, None] * stride_qt + offs_d[None, :] * stride_qd,
mask=(offs_m[:, None] < q_seqlen) & (offs_d[None, :] < D_HEAD),
other=0)
sparse_crow_ptr = layout_crow_ptr + off_h * layout_crow_stride_h + q_pbid * layout_crow_stride_m
# TODO: load at once, supported in new Triton
k_block_start = tl.load(sparse_crow_ptr).to(tl.int32)
k_block_end = tl.load(sparse_crow_ptr + 1).to(tl.int32)
m_i = tl.zeros([BLOCK_M_LOADING], dtype=tl.float32) - float('inf')
l_i = tl.zeros([BLOCK_M_LOADING], dtype=tl.float32)
acc = tl.zeros([BLOCK_M_LOADING, BLOCK_D], dtype=tl.float32)
k_ptrs = K + offs_n[None, :] * stride_kt + offs_d[:, None] * stride_kd
v_ptrs = V + offs_n[:, None] * stride_vt + offs_d[None, :] * stride_vd
for k_block_col_idx in range(k_block_start, k_block_end - 1):
acc, l_i, m_i = _fwd_kernel_inner(
acc, l_i, m_i,
q, Q,
k_block_col_idx,
layout_col_ptr,
layout_col_stride_h, layout_col_stride_m,
k_ptrs,
v_ptrs,
off_h, offs_m, offs_n, offs_d,
stride_kt, stride_vt,
sm_scale,
k_seqlen,
past_len,
False,
BLOCK_M_LOADING,
BLOCK_N,
D_HEAD,
EVEN_D,
M_LT_N
)
acc, l_i, m_i = _fwd_kernel_inner(
acc, l_i, m_i,
q, Q,
k_block_end - 1,
layout_col_ptr,
layout_col_stride_h, layout_col_stride_m,
k_ptrs,
v_ptrs,
off_h, offs_m, offs_n, offs_d,
stride_kt, stride_vt,
sm_scale,
k_seqlen,
past_len,
True,
BLOCK_M_LOADING,
BLOCK_N,
D_HEAD,
EVEN_D,
M_LT_N
)
# write output
if EVEN_D:
tl.store(Out + offs_m[:, None] * stride_ot + offs_d[None, :] * stride_od, acc,
mask=offs_m[:, None] < q_seqlen)
else:
tl.store(Out + offs_m[:, None] * stride_ot + offs_d[None, :] * stride_od, acc,
mask=(offs_m[:, None] < q_seqlen) & (offs_d[None, :] < D_HEAD))
###########################################################
###########################################################
###########################################################
################## Testing Utilities ######################
###########################################################
def torch_attention(q, k, v, attn_mask=None, sm_scale=None, block_attn_mask=None, block_size=128, do=None):
'''
q, k, v: shape=(batch, n_heads, seq, dim)
'''
# for verification
if sm_scale is None:
sm_scale = math.sqrt(float(q.size(-1)))
if block_attn_mask is not None:
assert attn_mask is None
outs = []
for s in range(0, q.size(2), block_size):
e = min(s + block_size, q.size(2))
q_block = q[:, :, s:e]
attn = torch.einsum('bhmd,bhnd->bhmn', q_block, k[:, :, :e]).float() * sm_scale
mask = block_attn_mask[..., s // block_size, : (s // block_size + 1)]
mask = torch.kron(mask, torch.ones(block_size, block_size, device=mask.device))
mask[..., :, s:].masked_fill_(torch.arange(0, block_size)[:, None] <= torch.arange(0, block_size)[None, :], 0)
attn = attn.masked_fill((1 - mask).bool(), float('-inf'))
attn = attn.softmax(-1)
out = torch.einsum('bhmn,bhnd->bhmd', attn.type_as(v), v[:, :, :e])
outs.append(out)
torch_output = torch.cat(outs, dim=2)
else:
attn = torch.einsum('bhmd,bhnd->bhmn', q, k).float() * sm_scale
# import ipdb; ipdb.set_trace()
if attn_mask is not None:
attn = attn.masked_fill((1 - attn_mask).bool(), float('-inf'))
# print(f'> torch attn: {attn.exp().sum(-1)=}')
attn = attn.softmax(-1)
if do is not None:
dv = torch.einsum('bhqk,bhqd->bhkd', attn.type_as(do), do)
print(f'> torch_attn computed dv: {dv=}')
torch_output = torch.einsum('bhmn,bhnd->bhmd', attn.type_as(v), v)
return torch_output
###########################################################
###########################################################
###########################################################
#################### Unit Tests ###########################
###########################################################
@pytest.mark.parametrize('Z, H, N_CTX, D_HEAD', [(2, 8, 2048, 128), (1, 4, 4096, 64)])
def test_op(Z, H, N_CTX, D_HEAD, Q_LEN=None, dtype=torch.bfloat16, homo_head=True, kernel_block_size=None, sparse_block_size=128, backward=True,
sparse_attention_fn=None, local_blocks=4, vert_stride=4, sm_scale=None, max_length=None):
Q_LEN = Q_LEN or N_CTX
torch.manual_seed(20)
q = torch.empty((Z, H, Q_LEN, D_HEAD), dtype=dtype, device='cuda').normal_(mean=0, std=.5) # .requires_grad_()
k = torch.empty((Z, H, N_CTX, D_HEAD), dtype=dtype, device='cuda').normal_(mean=0, std=.5) # .requires_grad_()
v = torch.empty((Z, H, N_CTX, D_HEAD), dtype=dtype, device='cuda').normal_(mean=0, std=.5) # .requires_grad_()
if sm_scale is None:
sm_scale = 1. / math.sqrt(D_HEAD)
# for debugging
# print(f'>> {q.shape=}, {k.shape=}, {v.shape=}, {homo_head=}, {kernel_block_size=}, {sparse_block_size=}, {local_blocks=}, {vert_stride=}')
sm_scale = 0.0078125
if backward:
q.requires_grad_(), k.requires_grad_(), v.requires_grad_()
# qkv = torch.empty((Z, N_CTX, 3*H*D_HEAD), dtype=dtype, device='cuda').normal_(mean=0, std=.5)
# q = qkv[..., :H*D_HEAD]
# k = qkv[..., H*D_HEAD:2*H*D_HEAD]
# v = qkv[..., 2*H*D_HEAD:]
# q = q.view(Z, N_CTX, H, -1).permute(0, 2, 1, 3)
# k = k.view(Z, N_CTX, H, -1).permute(0, 2, 1, 3)
# v = v.view(Z, N_CTX, H, -1).permute(0, 2, 1, 3)
# if Q_LEN and Q_LEN < N_CTX:
# q = q[:, :, -Q_LEN:] # .contiguous()
# q = q.requires_grad_()
# k = k.requires_grad_()
# v = v.requires_grad_()
dout = torch.randn_like(q).contiguous()
# dout = torch.eye(N_CTX)[:, :D_HEAD][None, None].expand_as(q).type_as(q).contiguous()
# print(dout)
mask_csr, _, mask_dense = get_sparse_attn_mask(q, N_CTX, BLOCK=sparse_block_size,
local_blocks=local_blocks, vert_stride=vert_stride, homo_head=homo_head, return_dense=True)
if sparse_attention_fn is None:
sparse_attention_fn = get_local_strided_sparse_attention_op(H, N_CTX,
sparse_block_size=sparse_block_size,
local_blocks=local_blocks,
vert_stride=vert_stride,
homo_head=homo_head,
device=q.device,
dtype=q.dtype,
kernel_block_size=kernel_block_size)
# reference implementation
ref_out = torch_attention(q, k, v, mask_dense, sm_scale)
# lengths = torch.full((Z,), fill_value=N_CTX, device='cuda')
# cu_seqlens = torch.zeros((Z + 1,), device='cuda', dtype=torch.int32)
# cu_seqlens[1:] = lengths.cumsum(0)
# # qkv = torch.randn((Z * N_CTX, 3, H, D_HEAD), dtype=dtype, device='cuda', requires_grad=True)
# qkv_list = list(map(lambda x: x.permute(0, 2, 1, 3).contiguous().view(Z * N_CTX, 1, H, D_HEAD), [q, k, v]))
# qkv = torch.cat(qkv_list, dim=1)
# ref_out0 = flash_attn_func(qkv, cu_seqlens, dropout_p=0, max_s=N_CTX, softmax_scale=sm_scale, causal=True)
# ref_out = ref_out0.view(Z, N_CTX, H, D_HEAD).permute(0, 2, 1, 3).contiguous()
if backward:
ref_out.backward(dout)
ref_dv, v.grad = v.grad.clone(), None
ref_dk, k.grad = k.grad.clone(), None
ref_dq, q.grad = q.grad.clone(), None
tri_out = sparse_attention_fn(q, k, v, sm_scale)
decimal = 1 if dtype == torch.bfloat16 else 2
assert torch.allclose(ref_out.cpu(), tri_out.cpu(), atol=1e-2, rtol=0), f'>> {ref_out[0, 0, :, 0].tolist()=}\n\n{tri_out[0, 0, :, 0].tolist()=}'
if backward:
tri_out.backward(dout)
tri_dv, v.grad = v.grad.clone(), None
tri_dk, k.grad = k.grad.clone(), None
tri_dq, q.grad = q.grad.clone(), None
if backward:
assert torch.allclose(ref_dv, tri_dv, atol=1e-2, rtol=1e-2)
assert torch.allclose(ref_dk, tri_dk, atol=1e-2, rtol=0)
assert torch.allclose(ref_dq, tri_dq, atol=1e-2, rtol=0)
print(f'> test passed: {Z=}, {H=}, {N_CTX=}, {D_HEAD=}, {Q_LEN=}, {dtype=}, {homo_head=}, {sparse_block_size=}')
###########################################################
if __name__ == '__main__':
GPU_TYPE = os.popen('nvidia-smi --query-gpu=name --format=csv | tail -n 1').read().strip()
# print(GPU_TYPE)
support_backward = True # 'A100' in GPU_TYPE. Wasn't supportted in consumer A1000.
###############
# benchmarking
HAS_DENSE_TRITON_FLASH = False
# try:
# from triton.ops.flash_attention import attention as triton_attention
# HAS_DENSE_TRITON_FLASH = True
# except:
# HAS_DENSE_TRITON_FLASH = False
# print('> cannot import Trition flash attn')
try:
from flash_attn.flash_attn_interface import flash_attn_func, flash_attn_unpadded_func
HAS_FLASH = True
except BaseException:
HAS_FLASH = False
print('> cannot import flash_attn')
# BATCH, N_HEADS, N_CTX, D_HEAD = 4, 48, 4096, 64
BATCH, N_HEADS, N_CTX, D_HEAD = 4, 32, 4096, 128 # 6.7B model, with 4k len
# BATCH, N_HEADS, N_CTX, D_HEAD = 4, 16, 4096, 128 # 204m model
BLOCK_SIZE = 64
LOCAl_BLOCKS = 8 # 4
VERT_STRIDE = 1 # 16 # 8
HOMO_HEAD = False
sparse_type = 'home' if HOMO_HEAD else 'hetero'
dtype = torch.bfloat16
modes = ['fwd', 'bwd'] if support_backward else ['fwd']
configs = [triton.testing.Benchmark(
x_names=['SEQ_LEN'],
x_vals=[2**i for i in range(8, 16)],
line_arg='provider',
line_vals=(['triton'] if HAS_DENSE_TRITON_FLASH else []) + (['flash'] if HAS_FLASH else []) + ['triton_sparse'],
line_names=(['Triton-Dense'] if HAS_DENSE_TRITON_FLASH else []) + (['Flash-Dense'] if HAS_FLASH else []) + ['Triton-Sparse'],
styles=[('red', '-'), ('blue', '-'), ('green', '-')],
ylabel='ms',
plot_name=f'fused-attention-batch{BATCH}-head{N_HEADS}-d{D_HEAD}-sparse-local{LOCAl_BLOCKS}-vert{VERT_STRIDE}-{sparse_type}-{dtype}-{mode}',
args={'H': N_HEADS, 'BATCH': BATCH, 'D_HEAD': D_HEAD, 'dtype': dtype, 'mode': mode}
) for mode in modes]
@triton.testing.perf_report(configs)
def bench_flash_attention(BATCH, H, SEQ_LEN, D_HEAD, mode, provider, dtype=torch.bfloat16, device='cuda', sparse_attention_fn=None):
assert mode in ['fwd', 'bwd']
warmup = 25
rep = 100
N_CTX = SEQ_LEN
if provider == 'triton':
q = torch.randn((BATCH, H, N_CTX, D_HEAD), dtype=dtype, device='cuda', requires_grad=True)
k = torch.randn((BATCH, H, N_CTX, D_HEAD), dtype=dtype, device='cuda', requires_grad=True)
v = torch.randn((BATCH, H, N_CTX, D_HEAD), dtype=dtype, device='cuda', requires_grad=True)
sm_scale = 1.3
fn = lambda: triton_attention(q, k, v, sm_scale)
if mode == 'bwd':
o = fn()
do = torch.randn_like(o)
fn = lambda: o.backward(do, retain_graph=True)
ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)
return ms
if provider == 'triton_sparse':
q = torch.randn((BATCH, H, N_CTX, D_HEAD), dtype=dtype, device='cuda', requires_grad=True)
k = torch.randn((BATCH, H, N_CTX, D_HEAD), dtype=dtype, device='cuda', requires_grad=True)
v = torch.randn((BATCH, H, N_CTX, D_HEAD), dtype=dtype, device='cuda', requires_grad=True)
sm_scale = 1.3
# q_pos = torch.arange(N_CTX // BLOCK, device='cuda')[:, None]
# k_pos = torch.arange(N_CTX // BLOCK, device='cuda')[None]
# local_blocks = 4 # num_block per attn, block_size is tied to BLOCK
# vert_stride =N_CTX + 1 # 4
# mask_vert_strided = torch.arange(N_CTX // BLOCK, device='cuda') % vert_stride == vert_stride - 1
# mask_dense = ((q_pos >= k_pos) & ((q_pos - k_pos < local_blocks) | mask_vert_strided)).type_as(q)
# mask = mask_dense.to_sparse_csr()
# mask_csr, _ = get_sparse_attn_mask(q, N_CTX, BLOCK=BLOCK, local_blocks=LOCAl_BLOCKS, vert_stride=VERT_STRIDE, homo_head=HOMO_HEAD)
if sparse_attention_fn is None:
# sparse_attention_fn = sparse_attention
sparse_attention_fn = get_local_strided_sparse_attention_op(H, SEQ_LEN,
local_blocks=LOCAl_BLOCKS,
vert_stride=VERT_STRIDE,
homo_head=HOMO_HEAD,
sparse_block_size=BLOCK_SIZE,
kernel_block_size=BLOCK_SIZE,
device=q.device)
# sparse_attention_fn = sparse_attention_factory(128, 128, num_warps=8)
# fn = lambda: sparse_attention_fn(q, k, v, mask_csr[0], mask_csr[1], sm_scale)
fn = lambda: sparse_attention_fn(q, k, v, sm_scale)
if mode == 'bwd':
o = fn()
do = torch.randn_like(o)
fn = lambda: o.backward(do, retain_graph=True)
ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)
return ms
if provider == 'flash':
lengths = torch.full((BATCH,), fill_value=N_CTX, device=device)
cu_seqlens = torch.zeros((BATCH + 1,), device=device, dtype=torch.int32)
cu_seqlens[1:] = lengths.cumsum(0)
qkv = torch.randn((BATCH * N_CTX, 3, H, D_HEAD), dtype=dtype, device=device, requires_grad=True)
fn = lambda: flash_attn_func(qkv, cu_seqlens, 0., N_CTX, causal=True)
if mode == 'bwd':
o = fn()
do = torch.randn_like(o)
fn = lambda: o.backward(do, retain_graph=True)
ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)
return ms
# if provider == 'torch':
# q = torch.randn((BATCH, H, N_CTX, D_HEAD), dtype=dtype, device='cuda', requires_grad=True)
# k = torch.randn((BATCH, H, N_CTX, D_HEAD), dtype=dtype, device='cuda', requires_grad=True)
# v = torch.randn((BATCH, H, N_CTX, D_HEAD), dtype=dtype, device='cuda', requires_grad=True)
# sm_scale = 1.3
# causal_mask = torch.tril(torch.ones(N_CTX, N_CTX)).type_as(q)
# fn = lambda: torch_attention(q, k, v, causal_mask, sm_scale)
# ms = triton.testing.do_bench(fn, percentiles=None, warmup=warmup, rep=rep)
# return ms
BATCH, N_HEADS, N_CTX, D_HEAD, Q_LEN = 4, 32, 4096, 128, 1 # 6.7B model, with 4k len
BLOCK_SIZE = 64
LOCAl_BLOCKS = 8 # 4
VERT_STRIDE = 16 # 8
HOMO_HEAD = False
sparse_type = 'home' if HOMO_HEAD else 'hetero'
dtype = torch.bfloat16
MAX_N_CTX = 8192
configs = [triton.testing.Benchmark(
x_names=['PAST_LEN'],
x_vals=[2**i - 1 for i in range(8, 14)],
line_arg='provider',
line_vals=['torch'] + (['flash'] if HAS_FLASH else []) + ['triton_sparse', 'triton_dense'],
line_names=['Torch'] + (['Flash-Dense'] if HAS_FLASH else []) + ['Triton-Sparse', 'Triton-Dense'],
styles=[('red', '-'), ('blue', '-'), ('green', '-'), ('cyan', '-')],
ylabel='ms',
plot_name=f'fused-attention-inference-batch{BATCH}-head{N_HEADS}-d{D_HEAD}-sparse-local{LOCAl_BLOCKS}-vert{VERT_STRIDE}-{sparse_type}',
args={'H': N_HEADS, 'BATCH': BATCH, 'D_HEAD': D_HEAD, 'Q_LEN': Q_LEN, 'dtype': torch.float16, 'mode': mode}
) for mode in ['fwd']]
@triton.testing.perf_report(configs)
def bench_flash_attention_inference(BATCH, H, PAST_LEN, D_HEAD, Q_LEN, mode, provider, dtype=torch.bfloat16, device='cuda'):
assert mode in ['fwd']
warmup = 25
rep = 100
N_CTX = PAST_LEN + Q_LEN
if provider == 'torch':
q = torch.randn((BATCH, H, Q_LEN, D_HEAD), dtype=dtype, device='cuda', requires_grad=False)
k = torch.randn((BATCH, H, N_CTX, D_HEAD), dtype=dtype, device='cuda', requires_grad=False)
v = torch.randn((BATCH, H, N_CTX, D_HEAD), dtype=dtype, device='cuda', requires_grad=False)
sm_scale = 1.3
mask_csr, _, mask_dense = get_sparse_attn_mask(q, N_CTX, BLOCK=BLOCK_SIZE,
local_blocks=LOCAl_BLOCKS, vert_stride=VERT_STRIDE, homo_head=VERT_STRIDE, return_dense=True)
fn = lambda: torch_attention(q, k, v, mask_dense, sm_scale=sm_scale, block_size=2048)
ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)
return ms
if provider == 'triton_sparse':
q = torch.randn((BATCH, H, Q_LEN, D_HEAD), dtype=dtype, device='cuda', requires_grad=False)
k = torch.randn((BATCH, H, N_CTX, D_HEAD), dtype=dtype, device='cuda', requires_grad=False)
v = torch.randn((BATCH, H, N_CTX, D_HEAD), dtype=dtype, device='cuda', requires_grad=False)
sm_scale = 1.3
sparse_attention_fn = get_local_strided_sparse_attention_op(H, MAX_N_CTX,
local_blocks=LOCAl_BLOCKS,
vert_stride=VERT_STRIDE,
homo_head=HOMO_HEAD,
sparse_block_size=BLOCK_SIZE,
kernel_block_size=BLOCK_SIZE,
device=q.device,
inference=True)
fn = lambda: sparse_attention_fn(q, k, v, sm_scale)
ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)
return ms
if provider == 'triton_dense':
q = torch.randn((BATCH, H, Q_LEN, D_HEAD), dtype=dtype, device='cuda', requires_grad=False)
k = torch.randn((BATCH, H, N_CTX, D_HEAD), dtype=dtype, device='cuda', requires_grad=False)
v = torch.randn((BATCH, H, N_CTX, D_HEAD), dtype=dtype, device='cuda', requires_grad=False)
sm_scale = 1.3
sparse_attention_fn = get_local_strided_sparse_attention_op(H, MAX_N_CTX,
local_blocks=1,
vert_stride=1,
homo_head=True,
sparse_block_size=BLOCK_SIZE,
kernel_block_size=BLOCK_SIZE,
device=q.device,
inference=True)
fn = lambda: sparse_attention_fn(q, k, v, sm_scale)
ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)
return ms
if provider == 'flash':
assert Q_LEN == 1
lengths = torch.full((BATCH,), fill_value=N_CTX, device=device)
cu_seqlens = torch.zeros((BATCH + 1,), device=device, dtype=torch.int32)
cu_seqlens[1:] = lengths.cumsum(0)
cu_seqlens_q = torch.arange(BATCH + 1, device=device, dtype=torch.int32)
# (total_q, nheads, headdim),
q = torch.randn((BATCH, H, D_HEAD), dtype=dtype, device='cuda', requires_grad=False)
k = torch.randn((BATCH*N_CTX, H, D_HEAD), dtype=dtype, device='cuda', requires_grad=False)
v = torch.randn((BATCH*N_CTX, H, D_HEAD), dtype=dtype, device='cuda', requires_grad=False)
fn = lambda: flash_attn_unpadded_func(q, k, v, cu_seqlens_q, cu_seqlens, 1, N_CTX, dropout_p=0, softmax_scale=1.3, causal=False)
ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)
return ms
test_op(1, 4, 512, 128, dtype=torch.float16, homo_head=False, backward=support_backward)
# bench_flash_attention.run(save_path='.', print_data=True)
bench_flash_attention_inference.run(save_path='.', print_data=True)
exit()
# head_dim=64
test_op(1, 2, 1024, 64, kernel_block_size=64, sparse_block_size=64,
dtype=torch.bfloat16, homo_head=False, backward=support_backward)
# uneven length, bf16
test_op(1, 16, 224, 128, dtype=torch.bfloat16, homo_head=False, backward=False, sparse_block_size=128,
kernel_block_size=64, local_blocks=8, vert_stride=8)
test_op(3, 2, 2047, 128, homo_head=False, backward=False)
# diff kernel/sparse block size
test_op(1, 16, 224, 128, dtype=torch.bfloat16, homo_head=False, backward=False, kernel_block_size=64)
# inference
# test_op(1, 4, 512 + 256, 128, Q_LEN=1, dtype=torch.bfloat16, homo_head=False, backward=support_backward)
# dense flash attn
test_op(1, 2, 1024, 128, kernel_block_size=128, sparse_block_size=128, dtype=torch.bfloat16, homo_head=False,
backward=support_backward, local_blocks=1, vert_stride=1)
# fp16
test_op(1, 4, 512 + 256, 128, dtype=torch.float16, homo_head=False, backward=support_backward)
# longer sequence
test_op(2, 4, 8192, 64, homo_head=False, backward=support_backward)
test_op(2, 4, 8192, 128, dtype=torch.bfloat16, homo_head=False, backward=support_backward)
# homo head
test_op(3, 2, 2048, 64, homo_head=True, dtype=torch.bfloat16, backward=False)
test_op(3, 2, 2048, 64, homo_head=True, backward=support_backward)
# sparse_attention_fn = sparse_attention_factory(16, 128, num_warps=1, INFERENCE=True)
# test_op(8, 1, 2047, 128, 1, backward=False, sparse_attention_fn=None)
# test_op_inference(3, 2, 2048, 128, 2048)
# test_op_inference(3, 2, 2047, 64, 2047)
# test_op_inference(3, 2, 256, 64, 128)
# test_op_inference(3, 2, 2048, 64, 1)
bench_flash_attention.run(save_path='.', print_data=True)
# bench_flash_attention_inference.run(save_path='.', print_data=True)
# ========================
# Some Benchmark Results #
# ========================
# fused-attention-batch4-head48-d64-sparse-local4-vert4-hetero-fwd
# SEQ_LEN Triton-Dense Flash-Dense Triton-Sparse
# 0 256.0 0.057184 0.069646 0.052567
# 1 512.0 0.131688 0.187658 0.110212
# 2 1024.0 0.391844 0.524990 0.247875
# 3 2048.0 1.305190 1.456685 0.596506
# 4 4096.0 4.623019 4.968653 1.600277
# 5 8192.0 17.513062 18.332262 4.802458
# 6 16384.0 68.453377 70.337540 16.052908
# 7 32768.0 270.655487 276.020233 57.938946
# fused-attention-batch4-head48-d64-sparse-local4-vert4-hetero-bwd (num_warp=8):
# SEQ_LEN Triton-Dense Flash-Dense Triton-Sparse
# 0 256.0 0.190120 0.150313 0.181451
# 1 512.0 0.406348 0.391767 0.391177
# 2 1024.0 1.029704 1.182967 0.885741
# 3 2048.0 2.985456 3.843399 2.040469
# 4 4096.0 9.808897 13.073701 5.069609
# 5 8192.0 34.995201 47.863808 13.948782
# 6 16384.0 132.740097 182.579193 42.816513
# 7 32768.0 542.223389 714.820618 147.053574
# fused-attention-inference-batch4-head32-d128-sparse-local4-vert4-hetero:
# PAST_LEN Torch-Dense Flash-Dense Triton-Sparse
# 0 256.0 0.050949 0.032357 0.107513
# 1 512.0 0.073624 0.050651 0.199086
# 2 1024.0 0.107472 0.080379 0.245445
# 3 2048.0 0.178423 0.129448 0.338259
# 4 4096.0 0.327647 0.223106 0.517048
# 5 8192.0 0.588423 0.411263 0.884606
# 6 16384.0 1.098898 0.798941 1.611809
# 7 32768.0 2.094537 1.594726 3.044160
# 6.7B
# fused-attention-batch4-head32-d128-sparse-local4-vert4-hetero-fwd:
# SEQ_LEN Triton-Dense Flash-Dense Triton-Sparse
# 0 256.0 0.069208 0.082156 0.065097
# 1 512.0 0.138271 0.201393 0.144467
# 2 1024.0 0.391521 0.624614 0.322382
# 3 2048.0 1.268443 2.406325 0.784367
# 4 4096.0 4.455703 9.139097 2.100856
# 5 8192.0 16.764315 35.289600 6.328320
# 6 16384.0 65.221634 138.401794 21.069057
# 7 32768.0 257.251343 548.085754 76.111870
# fused-attention-batch4-head32-d128-sparse-local4-vert4-hetero-bwd:
# SEQ_LEN Triton-Dense Flash-Dense Triton-Sparse
# 0 256.0 0.297118 0.266469 0.255255
# 1 512.0 0.672826 0.613685 0.552954
# 2 1024.0 1.718434 1.705066 1.251953
# 3 2048.0 4.936755 5.403875 2.927895
# 4 4096.0 15.911594 18.959362 7.436288
# 5 8192.0 55.357441 70.808578 21.140224
# 6 16384.0 208.188416 273.617920 68.018173
# 7 32768.0 806.037476 1081.453613 218.720261
# fused-attention-inference-batch4-head32-d128-sparse-local4-vert4-hetero:
# PAST_LEN Torch-Dense Flash-Dense Triton-Sparse
# 0 256.0 0.050151 0.032337 0.107593
# 1 512.0 0.073409 0.051737 0.200200
# 2 1024.0 0.107533 0.082099 0.247067
# 3 2048.0 0.177259 0.128891 0.338510
# 4 4096.0 0.325866 0.223621 0.524842
# 5 8192.0 0.586926 0.408913 0.885490
# 6 16384.0 1.100834 0.793277 1.612271
# 7 32768.0 2.098851 1.595831 3.064544
# fused-attention-batch4-head32-d128-sparse-local4-vert8-hetero-fwd:
# SEQ_LEN Triton-Dense Flash-Dense Triton-Sparse
# 0 256.0 0.066673 0.082037 0.065085
# 1 512.0 0.137379 0.201880 0.143473
# 2 1024.0 0.390675 0.624234 0.312046
# 3 2048.0 1.267739 2.406950 0.696045
# 4 4096.0 4.445138 9.136333 1.665788
# 5 8192.0 16.768614 35.265533 4.380486
# 6 16384.0 65.235970 138.393600 12.997633
# 7 32768.0 257.317902 550.442993 42.821121
# fused-attention-batch4-head32-d128-sparse-local4-vert8-hetero-bwd:
# SEQ_LEN Triton-Dense Flash-Dense Triton-Sparse
# 0 256.0 0.296461 0.266581 0.254022
# 1 512.0 0.671427 0.613643 0.551283
# 2 1024.0 1.719918 1.704295 1.229982
# 3 2048.0 4.945305 5.403364 2.721906
# 4 4096.0 15.934293 18.960999 6.259371
# 5 8192.0 55.406593 70.832130 15.676929
# 6 16384.0 208.750595 275.004425 44.837891
# 7 32768.0 808.057861 1080.647705 141.856766
# fused-attention-inference-batch4-head32-d128-sparse-local4-vert8-hetero:
# PAST_LEN Torch-Dense Flash-Dense Triton-Sparse
# 0 256.0 0.050739 0.032886 0.107837
# 1 512.0 0.073507 0.051996 0.200293
# 2 1024.0 0.106394 0.080679 0.240610
# 3 2048.0 0.177659 0.127660 0.287625
# 4 4096.0 0.326326 0.226971 0.377500
# 5 8192.0 0.586339 0.407367 0.559266
# 6 16384.0 1.102279 0.786221 0.920976
# 7 32768.0 2.097370 1.545090 1.644288
################
##### fp16 #####
################
# fused-attention-batch4-head16-d64-sparse-local4-vert8-hetero-fwd:
# SEQ_LEN Triton-Dense Flash-Dense Triton-Sparse
# 0 256.0 0.032518 0.035472 0.029939
# 1 512.0 0.054266 0.087841 0.054320
# 2 1024.0 0.133447 0.263090 0.102045
# 3 2048.0 0.384615 1.023293 0.201763
# 4 4096.0 1.300890 4.023936 0.449555
# 5 8192.0 4.774144 15.816704 1.150854
# 6 16384.0 18.220032 62.771198 3.356001
# 7 32768.0 71.405571 250.273788 10.976142
# fused-attention-batch4-head16-d64-sparse-local4-vert8-hetero-bwd:
# SEQ_LEN Triton-Dense Flash-Dense Triton-Sparse
# 0 256.0 0.083342 0.069742 0.079496
# 1 512.0 0.159894 0.170995 0.151705
# 2 1024.0 0.386071 0.522407 0.331443
# 3 2048.0 1.067715 1.737333 0.715248
# 4 4096.0 3.382731 6.219520 1.597457
# 5 8192.0 11.857793 23.560448 3.879035
# 6 16384.0 44.422142 91.251709 10.626843
# 7 32768.0 175.011841 359.473145 32.340992
################
##### bf16 #####
################
# fused-attention-batch4-head16-d64-sparse-local4-vert8-hetero-fwd:
# SEQ_LEN Triton-Dense Flash-Dense Triton-Sparse
# 0 256.0 0.037636 0.035902 0.031512
# 1 512.0 0.058591 0.087229 0.058125
# 2 1024.0 0.143337 0.263919 0.108443
# 3 2048.0 0.414458 1.025985 0.214114
# 4 4096.0 1.390841 4.020010 0.480550
# 5 8192.0 5.067938 15.808171 1.230874
# 6 16384.0 19.442280 62.765057 3.597274
# 7 32768.0 75.501572 250.443771 11.768959
# fused-attention-batch4-head16-d64-sparse-local4-vert8-hetero-bwd:
# SEQ_LEN Triton-Dense Flash-Dense Triton-Sparse
# 0 256.0 0.084404 0.070663 0.082613
# 1 512.0 0.161510 0.172882 0.157661
# 2 1024.0 0.388954 0.526047 0.339855
# 3 2048.0 1.075814 1.736057 0.732420
# 4 4096.0 3.401622 6.221376 1.636039
# 5 8192.0 11.915136 23.483391 3.968725
# 6 16384.0 44.660225 91.302910 10.857130
# 7 32768.0 175.038467 359.048187 32.778240