Skip to content

GitLab

Commit 24403aea authored by xs-keju's avatar xs-keju Committed by LeiWang1999
Browse files

[CI] Add CI test for flash_attention examples (#558) 



* [CI] Add CI test for flash_attention examples

* Update example_gqa_fwd_bshd.py

* Update example_mha_fwd_bshd_wgmma_pipelined.py

* [CI] Added conditional annotations for tests in flash_attention

* [CI] Added conditional annotations for tests in flash_attention

---------
Co-authored-by: default avatarLei Wang <34334180+LeiWang1999@users.noreply.github.com>
parent 17f7394f
Show whitespace changes
Inline Side-by-side
Showing with 417 additions and 226 deletions
examples/flash_attention/example_gqa_bwd.py
View file @ 24403aea
import torch import torch
import torch.nn.functional as F import torch.nn.functional as F
import tilelang import tilelang
from tilelang import cached
from tilelang.autotuner import * from tilelang.autotuner import *
import tilelang.language as T import tilelang.language as T
import argparse import argparse
def flashattn_fwd(batch, heads, seq_len, dim_qk, dim_v, is_casual, block_M, block_N, groups=1): @tilelang.jit(out_idx=[3, 4])
def flashattn_fwd(batch, heads, seq_len, dim_qk, dim_v, is_causal, block_M, block_N, groups=1):
scale = (1.0 / dim_qk)**0.5 * 1.44269504 # log2(e) scale = (1.0 / dim_qk)**0.5 * 1.44269504 # log2(e)
head_kv = heads // groups head_kv = heads // groups
q_shape = [batch, seq_len, heads, dim_qk] q_shape = [batch, seq_len, heads, dim_qk]
...@@ -24,7 +24,7 @@ def flashattn_fwd(batch, heads, seq_len, dim_qk, dim_v, is_casual, block_M, bloc ...@@ -24,7 +24,7 @@ def flashattn_fwd(batch, heads, seq_len, dim_qk, dim_v, is_casual, block_M, bloc
Output: T.Tensor([batch, seq_len, heads, dim_v], dtype), # type: ignore Output: T.Tensor([batch, seq_len, heads, dim_v], dtype), # type: ignore
lse: T.Tensor([batch, heads, seq_len], accum_dtype), # type: ignore lse: T.Tensor([batch, heads, seq_len], accum_dtype), # type: ignore
): ):
with T.Kernel(T.ceildiv(seq_len, block_M), heads, batch, threads=128) as (bx, by, bz): with T.Kernel(T.ceildiv(seq_len, block_M), heads, batch, threads=256) as (bx, by, bz):
Q_shared = T.alloc_shared([block_M, dim_qk], dtype) Q_shared = T.alloc_shared([block_M, dim_qk], dtype)
K_shared = T.alloc_shared([block_N, dim_qk], dtype) K_shared = T.alloc_shared([block_N, dim_qk], dtype)
V_shared = T.alloc_shared([block_N, dim_v], dtype) V_shared = T.alloc_shared([block_N, dim_v], dtype)
...@@ -44,10 +44,10 @@ def flashattn_fwd(batch, heads, seq_len, dim_qk, dim_v, is_casual, block_M, bloc ...@@ -44,10 +44,10 @@ def flashattn_fwd(batch, heads, seq_len, dim_qk, dim_v, is_casual, block_M, bloc
T.fill(scores_max, -T.infinity(accum_dtype)) T.fill(scores_max, -T.infinity(accum_dtype))
loop_range = ( loop_range = (
T.ceildiv( T.ceildiv(
(bx + 1) * block_M, block_N) if is_casual else T.ceildiv(seq_len, block_N)) (bx + 1) * block_M, block_N) if is_causal else T.ceildiv(seq_len, block_N))
for k in T.Pipelined(loop_range, num_stages=1): for k in T.Pipelined(loop_range, num_stages=1):
T.copy(K[bz, k * block_N:(k + 1) * block_N, by // groups, :], K_shared) T.copy(K[bz, k * block_N:(k + 1) * block_N, by // groups, :], K_shared)
if is_casual: if is_causal:
for i, j in T.Parallel(block_M, block_N): for i, j in T.Parallel(block_M, block_N):
acc_s[i, j] = T.if_then_else(bx * block_M + i >= k * block_N + j, 0, acc_s[i, j] = T.if_then_else(bx * block_M + i >= k * block_N + j, 0,
-T.infinity(acc_s.dtype)) -T.infinity(acc_s.dtype))
...@@ -78,6 +78,7 @@ def flashattn_fwd(batch, heads, seq_len, dim_qk, dim_v, is_casual, block_M, bloc ...@@ -78,6 +78,7 @@ def flashattn_fwd(batch, heads, seq_len, dim_qk, dim_v, is_casual, block_M, bloc
return flash_fwd return flash_fwd
@tilelang.jit(out_idx=[2])
def flashattn_bwd_preprocess(batch, heads, seq_len, dim_v): def flashattn_bwd_preprocess(batch, heads, seq_len, dim_v):
dtype = "float16" dtype = "float16"
accum_dtype = "float" accum_dtype = "float"
...@@ -113,6 +114,7 @@ def make_dq_layout(dQ): ...@@ -113,6 +114,7 @@ def make_dq_layout(dQ):
lambda b, l, h, d: [b, l // 8, h, d // 8, (d % 2), 4 * (l % 8) + (d % 8) // 2]) lambda b, l, h, d: [b, l // 8, h, d // 8, (d % 2), 4 * (l % 8) + (d % 8) // 2])
@tilelang.jit(out_idx=[1])
def flashattn_bwd_postprocess(batch, heads, seq_len, dim_qk): def flashattn_bwd_postprocess(batch, heads, seq_len, dim_qk):
dtype = "float16" dtype = "float16"
accum_dtype = "float" accum_dtype = "float"
...@@ -134,7 +136,8 @@ def flashattn_bwd_postprocess(batch, heads, seq_len, dim_qk): ...@@ -134,7 +136,8 @@ def flashattn_bwd_postprocess(batch, heads, seq_len, dim_qk):
return flash_bwd_post return flash_bwd_post
def flashattn_bwd(batch, heads, seq_len, dim_qk, dim_v, is_casual, block_M, block_N, groups=1): @tilelang.jit
def flashattn_bwd(batch, heads, seq_len, dim_qk, dim_v, is_causal, block_M, block_N, groups=1):
sm_scale = (1.0 / dim_qk)**0.5 sm_scale = (1.0 / dim_qk)**0.5
scale = (1.0 / dim_qk)**0.5 * 1.44269504 # log2(e) scale = (1.0 / dim_qk)**0.5 * 1.44269504 # log2(e)
head_kv = heads // groups head_kv = heads // groups
...@@ -156,7 +159,7 @@ def flashattn_bwd(batch, heads, seq_len, dim_qk, dim_v, is_casual, block_M, bloc ...@@ -156,7 +159,7 @@ def flashattn_bwd(batch, heads, seq_len, dim_qk, dim_v, is_casual, block_M, bloc
dK: T.Tensor(k_shape, dtype), # type: ignore dK: T.Tensor(k_shape, dtype), # type: ignore
dV: T.Tensor(v_shape, dtype), # type: ignore dV: T.Tensor(v_shape, dtype), # type: ignore
): ):
with T.Kernel(heads, T.ceildiv(seq_len, block_M), batch, threads=256) as (bx, by, bz): with T.Kernel(heads, T.ceildiv(seq_len, block_M), batch, threads=128) as (bx, by, bz):
K_shared = T.alloc_shared([block_M, dim_qk], dtype) K_shared = T.alloc_shared([block_M, dim_qk], dtype)
dsT_shared = T.alloc_shared([block_M, block_N], dtype) dsT_shared = T.alloc_shared([block_M, block_N], dtype)
q = T.alloc_shared([block_N, dim_qk], dtype) q = T.alloc_shared([block_N, dim_qk], dtype)
...@@ -185,7 +188,7 @@ def flashattn_bwd(batch, heads, seq_len, dim_qk, dim_v, is_casual, block_M, bloc ...@@ -185,7 +188,7 @@ def flashattn_bwd(batch, heads, seq_len, dim_qk, dim_v, is_casual, block_M, bloc
T.copy(V[bz, by * block_M:(by + 1) * block_M, bx // groups, :], V_shared) T.copy(V[bz, by * block_M:(by + 1) * block_M, bx // groups, :], V_shared)
T.clear(dv) T.clear(dv)
T.clear(dk) T.clear(dk)
loop_st = T.floordiv(by * block_M, block_N) if is_casual else 0 loop_st = T.floordiv(by * block_M, block_N) if is_causal else 0
loop_ed = T.ceildiv(seq_len, block_N) loop_ed = T.ceildiv(seq_len, block_N)
for k in T.Pipelined(loop_st, loop_ed, num_stages=1): for k in T.Pipelined(loop_st, loop_ed, num_stages=1):
T.copy(Q[bz, k * block_N:(k + 1) * block_N, bx, :], q) T.copy(Q[bz, k * block_N:(k + 1) * block_N, bx, :], q)
...@@ -194,7 +197,7 @@ def flashattn_bwd(batch, heads, seq_len, dim_qk, dim_v, is_casual, block_M, bloc ...@@ -194,7 +197,7 @@ def flashattn_bwd(batch, heads, seq_len, dim_qk, dim_v, is_casual, block_M, bloc
T.copy(lse[bz, bx, k * block_N:(k + 1) * block_N], lse_shared) T.copy(lse[bz, bx, k * block_N:(k + 1) * block_N], lse_shared)
for i, j in T.Parallel(block_M, block_N): for i, j in T.Parallel(block_M, block_N):
qkT[i, j] = T.exp2(qkT[i, j] * scale - lse_shared[j]) qkT[i, j] = T.exp2(qkT[i, j] * scale - lse_shared[j])
if is_casual: if is_causal:
for i, j in T.Parallel(block_M, block_N): for i, j in T.Parallel(block_M, block_N):
qkT[i, j] = T.if_then_else(by * block_M + i <= k * block_N + j, qkT[i, j], qkT[i, j] = T.if_then_else(by * block_M + i <= k * block_N + j, qkT[i, j],
0) 0)
...@@ -225,16 +228,16 @@ def flashattn_bwd(batch, heads, seq_len, dim_qk, dim_v, is_casual, block_M, bloc ...@@ -225,16 +228,16 @@ def flashattn_bwd(batch, heads, seq_len, dim_qk, dim_v, is_casual, block_M, bloc
return flash_bwd return flash_bwd
@torch.compile
class _attention(torch.autograd.Function): class _attention(torch.autograd.Function):
@staticmethod @staticmethod
def forward(ctx, q, k, v, causal, groups=1): def forward(ctx, q, k, v, causal, groups=1):
BATCH, N_CTX, H, D_HEAD_QK = q.shape BATCH, N_CTX, H, D_HEAD_QK = q.shape
D_HEAD_V = v.shape[-1] D_HEAD_V = v.shape[-1]
block_M = 64 block_M = 128
block_N = 64 block_N = 64
mod = cached(flashattn_fwd, [3, 4], BATCH, H, N_CTX, D_HEAD_QK, D_HEAD_V, causal, block_M, mod = flashattn_fwd(BATCH, H, N_CTX, D_HEAD_QK, D_HEAD_V, causal, block_M, block_N, groups)
block_N, groups)
o, lse = mod(q, k, v) o, lse = mod(q, k, v)
ctx.save_for_backward(q, k, v, o, lse) ctx.save_for_backward(q, k, v, o, lse)
ctx.causal = causal ctx.causal = causal
...@@ -243,6 +246,9 @@ class _attention(torch.autograd.Function): ...@@ -243,6 +246,9 @@ class _attention(torch.autograd.Function):
@staticmethod @staticmethod
def backward(ctx, do): def backward(ctx, do):
q, k, v, o, lse = ctx.saved_tensors q, k, v, o, lse = ctx.saved_tensors
BATCH, N_CTX, H, D_HEAD_QK = q.shape
HEAD_KV, D_HEAD_V, = v.shape[-2], v.shape[-1]
groups = H // HEAD_KV
def maybe_contiguous(x): def maybe_contiguous(x):
if x.stride(-1) != 1: if x.stride(-1) != 1:
...@@ -250,14 +256,20 @@ class _attention(torch.autograd.Function): ...@@ -250,14 +256,20 @@ class _attention(torch.autograd.Function):
return x return x
do, q, k, v, o = [maybe_contiguous(x) for x in (do, q, k, v, o)] do, q, k, v, o = [maybe_contiguous(x) for x in (do, q, k, v, o)]
block_M = 128 block_M = 64
block_N = 128 block_N = 32
mod_prep = cached(flashattn_bwd_preprocess, [2], BATCH, H, N_CTX, D_HEAD_V) mod_prep = flashattn_bwd_preprocess(BATCH, H, N_CTX, D_HEAD_V)
mod_post = cached(flashattn_bwd_postprocess, [1], BATCH, H, N_CTX, D_HEAD_QK) mod_post = flashattn_bwd_postprocess(BATCH, H, N_CTX, D_HEAD_QK)
delta = mod_prep(o, do) delta = mod_prep(o, do)
mod = cached(flashattn_bwd, [6, 7, 8], BATCH, H, N_CTX, D_HEAD_QK, D_HEAD_V, ctx.causal, kernel = flashattn_bwd(BATCH, H, N_CTX, D_HEAD_QK, D_HEAD_V, ctx.causal, block_M, block_N,
block_M, block_N, groups) groups)
dq, dk, dv = mod(q, k, v, do, lse, delta) shape_q = [BATCH, N_CTX, H, D_HEAD_QK]
shape_k = [BATCH, N_CTX, HEAD_KV, D_HEAD_QK]
shape_v = [BATCH, N_CTX, HEAD_KV, D_HEAD_V]
dq = torch.zeros(shape_q, dtype=torch.float32, device=q.device)
dk = torch.zeros(shape_k, dtype=torch.float16, device=q.device)
dv = torch.zeros(shape_v, dtype=torch.float16, device=q.device)
kernel(q, k, v, do, lse, delta, dq, dk, dv)
dq = mod_post(dq) dq = mod_post(dq)
return dq, dk, dv, None, None return dq, dk, dv, None, None
...@@ -290,22 +302,17 @@ def ref_program(Q, K, V, is_causal, groups=1): ...@@ -290,22 +302,17 @@ def ref_program(Q, K, V, is_causal, groups=1):
return output return output
if __name__ == "__main__": def main(BATCH: int = 8,
parser = argparse.ArgumentParser() H: int = 32,
parser.add_argument('--batch', type=int, default=8, help='Batch size') N_CTX: int = 1024,
parser.add_argument('--h', type=int, default=32, help='Number of heads') D_HEAD_QK: int = 192,
parser.add_argument('--n_ctx', type=int, default=1024, help='Context size') D_HEAD_V: int = 128,
parser.add_argument('--d_head_qk', type=int, default=192, help='Head dimension for Q/K') groups: int = 16,
parser.add_argument('--d_head_v', type=int, default=128, help='Head dimension for V') causal: bool = False):
parser.add_argument('--casual', type=bool, default=False, help='Casual flag')
parser.add_argument('--groups', type=int, default=16, help='groups')
args = parser.parse_args()
BATCH, H, N_CTX, D_HEAD_QK, D_HEAD_V, groups = args.batch, args.h, args.n_ctx, args.d_head_qk, args.d_head_v, args.groups
casual = args.casual
flops_per_qk = 2.0 * BATCH * H * N_CTX * N_CTX * D_HEAD_QK flops_per_qk = 2.0 * BATCH * H * N_CTX * N_CTX * D_HEAD_QK
flops_per_v = 2.0 * BATCH * H * N_CTX * N_CTX * D_HEAD_V flops_per_v = 2.0 * BATCH * H * N_CTX * N_CTX * D_HEAD_V
total_flops = 3 * flops_per_qk + 2 * flops_per_v total_flops = 3 * flops_per_qk + 2 * flops_per_v
if casual: if causal:
total_flops *= 0.5 total_flops *= 0.5
Q = ( Q = (
torch.empty(BATCH, N_CTX, H, D_HEAD_QK, dtype=torch.half, torch.empty(BATCH, N_CTX, H, D_HEAD_QK, dtype=torch.half,
...@@ -321,20 +328,21 @@ if __name__ == "__main__": ...@@ -321,20 +328,21 @@ if __name__ == "__main__":
dO = ( dO = (
torch.empty(BATCH, N_CTX, H, D_HEAD_V, dtype=torch.half, torch.empty(BATCH, N_CTX, H, D_HEAD_V, dtype=torch.half,
device="cuda").normal_().requires_grad_()) device="cuda").normal_().requires_grad_())
O = attention(Q, K, V, casual, groups) O = attention(Q, K, V, causal, groups)
O.backward(dO, retain_graph=True) O.backward(dO, retain_graph=True)
dQ, Q.grad = Q.grad.clone(), None dQ, Q.grad = Q.grad.clone(), None
dK, K.grad = K.grad.clone(), None dK, K.grad = K.grad.clone(), None
dV, V.grad = V.grad.clone(), None dV, V.grad = V.grad.clone(), None
O_ref = ref_program(Q, K, V, casual, groups) O_ref = ref_program(Q, K, V, causal, groups)
O_ref.backward(dO, retain_graph=True) O_ref.backward(dO, retain_graph=True)
dQ_ref, Q.grad = Q.grad.clone(), None dQ_ref, Q.grad = Q.grad.clone(), None
dK_ref, K.grad = K.grad.clone(), None dK_ref, K.grad = K.grad.clone(), None
dV_ref, V.grad = V.grad.clone(), None dV_ref, V.grad = V.grad.clone(), None
assert torch.allclose(O, O_ref, rtol=1e-2, atol=1e-2) assert torch.allclose(O, O_ref, rtol=1e-2, atol=1e-2)
assert torch.allclose(dV, dV_ref, rtol=1e-2, atol=1e-2) torch.testing.assert_close(dV, dV_ref, rtol=1e-2, atol=1e-2)
# assert torch.allclose(dV, dV_ref, rtol=1e-2, atol=1e-2)
assert torch.allclose(dK, dK_ref, rtol=1e-2, atol=1e-2) assert torch.allclose(dK, dK_ref, rtol=1e-2, atol=1e-2)
assert torch.allclose(dQ, dQ_ref, rtol=1e-2, atol=1e-2) assert torch.allclose(dQ, dQ_ref, rtol=1e-2, atol=1e-2)
...@@ -352,3 +360,16 @@ if __name__ == "__main__": ...@@ -352,3 +360,16 @@ if __name__ == "__main__":
latency = do_bench(run1, warmup=500) latency = do_bench(run1, warmup=500)
print("tilelang: {:.2f} ms".format(latency)) print("tilelang: {:.2f} ms".format(latency))
print("tilelang: {:.2f} TFlops".format(total_flops / latency * 1e-9)) print("tilelang: {:.2f} TFlops".format(total_flops / latency * 1e-9))
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--batch', type=int, default=8, help='Batch size')
parser.add_argument('--h', type=int, default=32, help='Number of heads')
parser.add_argument('--n_ctx', type=int, default=1024, help='Context size')
parser.add_argument('--d_head_qk', type=int, default=192, help='Head dimension for Q/K')
parser.add_argument('--d_head_v', type=int, default=128, help='Head dimension for V')
parser.add_argument('--causal', type=bool, default=False, help='Causal flag')
parser.add_argument('--groups', type=int, default=16, help='groups')
args = parser.parse_args()
main(args.batch, args.h, args.n_ctx, args.d_head_qk, args.d_head_v, args.groups, args.causal)
examples/flash_attention/example_gqa_fwd_bshd.py
View file @ 24403aea
...@@ -227,42 +227,52 @@ def ref_program(Q, K, V, is_causal, groups=1): ...@@ -227,42 +227,52 @@ def ref_program(Q, K, V, is_causal, groups=1):
return output return output
if __name__ == "__main__": def main(batch: int = 1,
parser = argparse.ArgumentParser() heads: int = 64,
parser.add_argument('--batch', type=int, default=1, help='batch size') seq_len: int = 4096,
parser.add_argument('--heads', type=int, default=64, help='heads') dim: int = 128,
parser.add_argument('--seq_len', type=int, default=8192, help='sequence length') is_causal: bool = False,
parser.add_argument('--dim', type=int, default=128, help='dim') groups: int = 16,
parser.add_argument('--is_causal', action='store_true', help='causal') tune: bool = False):
parser.add_argument('--tune', action='store_true', help='tune configs')
parser.add_argument('--groups', type=int, default=16, help='groups')
args = parser.parse_args()
batch, heads, seq_len, dim, is_causal, groups = args.batch, args.heads, args.seq_len, args.dim, args.is_causal, args.groups
flops_per_matmul = 2.0 * batch * heads * seq_len * seq_len * dim flops_per_matmul = 2.0 * batch * heads * seq_len * seq_len * dim
total_flops = 2 * flops_per_matmul total_flops = 2 * flops_per_matmul
if is_causal: if is_causal:
total_flops *= 0.5 total_flops *= 0.5
if (not args.tune): if (not tune):
program = flashattn( program = flashattn(
batch, heads, seq_len, dim, is_causal, tune=args.tune, groups=groups)( batch, heads, seq_len, dim, is_causal, tune=tune, groups=groups)(
block_M=128, block_N=128, num_stages=2, threads=128) block_M=64, block_N=64, num_stages=2, threads=128)
ref_program = partial(ref_program, is_causal=is_causal, groups=groups) ref_program_processed = partial(ref_program, is_causal=is_causal, groups=groups)
kernel = tilelang.compile(program, out_idx=[3]) kernel = tilelang.compile(program, out_idx=[3])
profiler = kernel.get_profiler(tensor_supply_type=tilelang.TensorSupplyType.Normal) profiler = kernel.get_profiler(tensor_supply_type=tilelang.TensorSupplyType.Normal)
profiler.assert_allclose(ref_program, rtol=0.01, atol=0.01) profiler.assert_allclose(ref_program_processed, rtol=0.01, atol=0.01)
print("All checks pass.") print("All checks pass.")
latency = profiler.do_bench(ref_program, warmup=500) latency = profiler.do_bench(ref_program_processed, warmup=500)
print("Ref: {:.2f} ms".format(latency)) print("Ref: {:.2f} ms".format(latency))
print("Ref: {:.2f} TFlops".format(total_flops / latency * 1e-9)) print("Ref: {:.2f} TFlops".format(total_flops / latency * 1e-9))
latency = profiler.do_bench(warmup=500) latency = profiler.do_bench(warmup=500)
print("Tile-lang: {:.2f} ms".format(latency)) print("Tile-lang: {:.2f} ms".format(latency))
print("Tile-lang: {:.2f} TFlops".format(total_flops / latency * 1e-9)) print("Tile-lang: {:.2f} TFlops".format(total_flops / latency * 1e-9))
else: else:
best_result = flashattn(batch, heads, seq_len, dim, is_causal, tune=args.tune) best_result = flashattn(batch, heads, seq_len, dim, is_causal, tune=tune)
best_latency = best_result.latency best_latency = best_result.latency
best_config = best_result.config best_config = best_result.config
ref_latency = best_result.ref_latency ref_latency = best_result.ref_latency
print(f"Best latency: {best_latency}") print(f"Best latency: {best_latency}")
print(f"Best TFlops: {total_flops / best_latency * 1e-9}") print(f"Best TFlops: {total_flops / best_latency * 1e-9}")
print(f"Best config: {best_config}") print(f"Best config: {best_config}")
print(f"Ref latency: {ref_latency}")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--batch', type=int, default=1, help='batch size')
parser.add_argument('--heads', type=int, default=64, help='heads')
parser.add_argument('--seq_len', type=int, default=4096, help='sequence length')
parser.add_argument('--dim', type=int, default=128, help='dim')
parser.add_argument('--is_causal', action='store_true', help='causal')
parser.add_argument('--tune', action='store_true', help='tune configs')
parser.add_argument('--groups', type=int, default=16, help='groups')
args = parser.parse_args()
main(args.batch, args.heads, args.seq_len, args.dim, args.is_causal, args.groups, args.tune)
examples/flash_attention/example_gqa_fwd_bshd_wgmma_pipelined.py
View file @ 24403aea
...@@ -199,42 +199,54 @@ def ref_program(Q, K, V, is_causal, groups=1): ...@@ -199,42 +199,54 @@ def ref_program(Q, K, V, is_causal, groups=1):
return output return output
if __name__ == "__main__": def main(
parser = argparse.ArgumentParser() batch: int = 1,
parser.add_argument('--batch', type=int, default=1, help='batch size') heads: int = 64,
parser.add_argument('--heads', type=int, default=64, help='heads') seq_len: int = 4096,
parser.add_argument('--seq_len', type=int, default=8192, help='sequence length') dim: int = 128,
parser.add_argument('--dim', type=int, default=128, help='dim') is_causal: bool = False,
parser.add_argument('--is_causal', action='store_true', help='causal') groups: int = 16,
parser.add_argument('--tune', action='store_true', help='tune configs') tune: bool = False,
parser.add_argument('--groups', type=int, default=16, help='groups') ):
args = parser.parse_args()
batch, heads, seq_len, dim, is_causal, groups = args.batch, args.heads, args.seq_len, args.dim, args.is_causal, args.groups
flops_per_matmul = 2.0 * batch * heads * seq_len * seq_len * dim flops_per_matmul = 2.0 * batch * heads * seq_len * seq_len * dim
total_flops = 2 * flops_per_matmul total_flops = 2 * flops_per_matmul
if is_causal: if is_causal:
total_flops *= 0.5 total_flops *= 0.5
if (not args.tune): if (not tune):
program = flashattn( program = flashattn(
batch, heads, seq_len, dim, is_causal, tune=args.tune, groups=groups)( batch, heads, seq_len, dim, is_causal, tune=tune, groups=groups)(
block_M=128, block_N=128, num_stages=2, threads=256) block_M=128, block_N=128, num_stages=2, threads=256)
ref_program = partial(ref_program, is_causal=is_causal, groups=groups) ref_program_processed = partial(ref_program, is_causal=is_causal, groups=groups)
kernel = tilelang.compile(program, out_idx=[3]) kernel = tilelang.compile(program, out_idx=[3])
profiler = kernel.get_profiler(tensor_supply_type=tilelang.TensorSupplyType.Normal) profiler = kernel.get_profiler(tensor_supply_type=tilelang.TensorSupplyType.Normal)
profiler.assert_allclose(ref_program, rtol=0.01, atol=0.01) profiler.assert_allclose(ref_program_processed, rtol=0.01, atol=0.01)
print("All checks pass.") print("All checks pass.")
latency = profiler.do_bench(ref_program, warmup=500) latency = profiler.do_bench(ref_program_processed, warmup=500)
print("Ref: {:.2f} ms".format(latency)) print("Ref: {:.2f} ms".format(latency))
print("Ref: {:.2f} TFlops".format(total_flops / latency * 1e-9)) print("Ref: {:.2f} TFlops".format(total_flops / latency * 1e-9))
latency = profiler.do_bench(warmup=500) latency = profiler.do_bench(warmup=500)
print("Tile-lang: {:.2f} ms".format(latency)) print("Tile-lang: {:.2f} ms".format(latency))
print("Tile-lang: {:.2f} TFlops".format(total_flops / latency * 1e-9)) print("Tile-lang: {:.2f} TFlops".format(total_flops / latency * 1e-9))
else: else:
best_result = flashattn(batch, heads, seq_len, dim, is_causal, tune=args.tune) best_result = flashattn(batch, heads, seq_len, dim, is_causal, tune=tune)
best_latency = best_result.latency best_latency = best_result.latency
best_config = best_result.config best_config = best_result.config
ref_latency = best_result.ref_latency ref_latency = best_result.ref_latency
print(f"Best latency: {best_latency}") print(f"Best latency: {best_latency}")
print(f"Best TFlops: {total_flops / best_latency * 1e-9}") print(f"Best TFlops: {total_flops / best_latency * 1e-9}")
print(f"Best config: {best_config}") print(f"Best config: {best_config}")
print(f"Ref latency: {ref_latency}")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--batch', type=int, default=1, help='batch size')
parser.add_argument('--heads', type=int, default=64, help='heads')
parser.add_argument('--seq_len', type=int, default=4096, help='sequence length')
parser.add_argument('--dim', type=int, default=128, help='dim')
parser.add_argument('--is_causal', action='store_true', help='causal')
parser.add_argument('--tune', action='store_true', help='tune configs')
parser.add_argument('--groups', type=int, default=16, help='groups')
args = parser.parse_args()
main(args.batch, args.heads, args.seq_len, args.dim, args.is_causal, args.groups, args.tune)
examples/flash_attention/example_mha_bwd.py
View file @ 24403aea
...@@ -6,7 +6,8 @@ import tilelang.language as T ...@@ -6,7 +6,8 @@ import tilelang.language as T
import argparse import argparse
def flashattn_fwd(batch, heads, seq_len, dim, is_casual, block_M, block_N): @tilelang.jit(out_idx=[3, 4])
def flashattn_fwd(batch, heads, seq_len, dim, is_causal, block_M, block_N):
scale = (1.0 / dim)**0.5 * 1.44269504 # log2(e) scale = (1.0 / dim)**0.5 * 1.44269504 # log2(e)
shape = [batch, seq_len, heads, dim] shape = [batch, seq_len, heads, dim]
dtype = "float16" dtype = "float16"
...@@ -44,10 +45,10 @@ def flashattn_fwd(batch, heads, seq_len, dim, is_casual, block_M, block_N): ...@@ -44,10 +45,10 @@ def flashattn_fwd(batch, heads, seq_len, dim, is_casual, block_M, block_N):
# Q_local[i, j] *= scale # Q_local[i, j] *= scale
loop_range = ( loop_range = (
T.ceildiv( T.ceildiv(
(bx + 1) * block_M, block_N) if is_casual else T.ceildiv(seq_len, block_N)) (bx + 1) * block_M, block_N) if is_causal else T.ceildiv(seq_len, block_N))
for k in T.Pipelined(loop_range, num_stages=1): for k in T.Pipelined(loop_range, num_stages=1):
T.copy(K[bz, k * block_N:(k + 1) * block_N, by, :], K_shared) T.copy(K[bz, k * block_N:(k + 1) * block_N, by, :], K_shared)
if is_casual: if is_causal:
for i, j in T.Parallel(block_M, block_N): for i, j in T.Parallel(block_M, block_N):
acc_s[i, j] = T.if_then_else(bx * block_M + i >= k * block_N + j, 0, acc_s[i, j] = T.if_then_else(bx * block_M + i >= k * block_N + j, 0,
-T.infinity(acc_s.dtype)) -T.infinity(acc_s.dtype))
...@@ -78,6 +79,7 @@ def flashattn_fwd(batch, heads, seq_len, dim, is_casual, block_M, block_N): ...@@ -78,6 +79,7 @@ def flashattn_fwd(batch, heads, seq_len, dim, is_casual, block_M, block_N):
return flash_fwd return flash_fwd
@tilelang.jit(out_idx=[2])
def flashattn_bwd_preprocess(batch, heads, seq_len, dim): def flashattn_bwd_preprocess(batch, heads, seq_len, dim):
dtype = "float16" dtype = "float16"
accum_dtype = "float" accum_dtype = "float"
...@@ -113,6 +115,7 @@ def make_dq_layout(dQ): ...@@ -113,6 +115,7 @@ def make_dq_layout(dQ):
lambda b, l, h, d: [b, l // 8, h, d // 8, (d % 2), 4 * (l % 8) + (d % 8) // 2]) lambda b, l, h, d: [b, l // 8, h, d // 8, (d % 2), 4 * (l % 8) + (d % 8) // 2])
@tilelang.jit(out_idx=[1])
def flashattn_bwd_postprocess(batch, heads, seq_len, dim): def flashattn_bwd_postprocess(batch, heads, seq_len, dim):
dtype = "float16" dtype = "float16"
accum_dtype = "float" accum_dtype = "float"
...@@ -134,7 +137,8 @@ def flashattn_bwd_postprocess(batch, heads, seq_len, dim): ...@@ -134,7 +137,8 @@ def flashattn_bwd_postprocess(batch, heads, seq_len, dim):
return flash_bwd_post return flash_bwd_post
def flashattn_bwd(batch, heads, seq_len, dim, is_casual, block_M, block_N): @tilelang.jit
def flashattn_bwd(batch, heads, seq_len, dim, is_causal, block_M, block_N):
sm_scale = (1.0 / dim)**0.5 sm_scale = (1.0 / dim)**0.5
scale = (1.0 / dim)**0.5 * 1.44269504 # log2(e) scale = (1.0 / dim)**0.5 * 1.44269504 # log2(e)
shape = [batch, seq_len, heads, dim] shape = [batch, seq_len, heads, dim]
...@@ -153,7 +157,7 @@ def flashattn_bwd(batch, heads, seq_len, dim, is_casual, block_M, block_N): ...@@ -153,7 +157,7 @@ def flashattn_bwd(batch, heads, seq_len, dim, is_casual, block_M, block_N):
dK: T.Tensor(shape, dtype), # type: ignore dK: T.Tensor(shape, dtype), # type: ignore
dV: T.Tensor(shape, dtype), # type: ignore dV: T.Tensor(shape, dtype), # type: ignore
): ):
with T.Kernel(heads, T.ceildiv(seq_len, block_M), batch, threads=256) as (bx, by, bz): with T.Kernel(heads, T.ceildiv(seq_len, block_M), batch, threads=128) as (bx, by, bz):
K_shared = T.alloc_shared([block_M, dim], dtype) K_shared = T.alloc_shared([block_M, dim], dtype)
dsT_shared = T.alloc_shared([block_M, block_N], dtype) dsT_shared = T.alloc_shared([block_M, block_N], dtype)
# should not store K to local if dim is large # should not store K to local if dim is large
...@@ -181,12 +185,11 @@ def flashattn_bwd(batch, heads, seq_len, dim, is_casual, block_M, block_N): ...@@ -181,12 +185,11 @@ def flashattn_bwd(batch, heads, seq_len, dim, is_casual, block_M, block_N):
dv_shared: tilelang.layout.make_swizzled_layout(dv_shared), dv_shared: tilelang.layout.make_swizzled_layout(dv_shared),
dk_shared: tilelang.layout.make_swizzled_layout(dk_shared), dk_shared: tilelang.layout.make_swizzled_layout(dk_shared),
}) })
T.copy(K[bz, by * block_M:(by + 1) * block_M, bx, :], K_shared) T.copy(K[bz, by * block_M:(by + 1) * block_M, bx, :], K_shared)
T.copy(V[bz, by * block_M:(by + 1) * block_M, bx, :], V_shared) T.copy(V[bz, by * block_M:(by + 1) * block_M, bx, :], V_shared)
T.clear(dv) T.clear(dv)
T.clear(dk) T.clear(dk)
loop_st = T.floordiv(by * block_M, block_N) if is_casual else 0 loop_st = T.floordiv(by * block_M, block_N) if is_causal else 0
loop_ed = T.ceildiv(seq_len, block_N) loop_ed = T.ceildiv(seq_len, block_N)
for k in T.Pipelined(loop_st, loop_ed, num_stages=2): for k in T.Pipelined(loop_st, loop_ed, num_stages=2):
T.copy(Q[bz, k * block_N:(k + 1) * block_N, bx, :], q) T.copy(Q[bz, k * block_N:(k + 1) * block_N, bx, :], q)
...@@ -195,7 +198,7 @@ def flashattn_bwd(batch, heads, seq_len, dim, is_casual, block_M, block_N): ...@@ -195,7 +198,7 @@ def flashattn_bwd(batch, heads, seq_len, dim, is_casual, block_M, block_N):
T.copy(lse[bz, bx, k * block_N:(k + 1) * block_N], lse_shared) T.copy(lse[bz, bx, k * block_N:(k + 1) * block_N], lse_shared)
for i, j in T.Parallel(block_M, block_N): for i, j in T.Parallel(block_M, block_N):
qkT[i, j] = T.exp2(qkT[i, j] * scale - lse_shared[j]) qkT[i, j] = T.exp2(qkT[i, j] * scale - lse_shared[j])
if is_casual: if is_causal:
for i, j in T.Parallel(block_M, block_N): for i, j in T.Parallel(block_M, block_N):
qkT[i, j] = T.if_then_else(by * block_M + i <= k * block_N + j, qkT[i, j], qkT[i, j] = T.if_then_else(by * block_M + i <= k * block_N + j, qkT[i, j],
0) 0)
...@@ -232,12 +235,7 @@ class _attention(torch.autograd.Function): ...@@ -232,12 +235,7 @@ class _attention(torch.autograd.Function):
BATCH, N_CTX, H, D_HEAD = q.shape BATCH, N_CTX, H, D_HEAD = q.shape
block_M = 64 block_M = 64
block_N = 64 if D_HEAD <= 128 else 32 block_N = 64 if D_HEAD <= 128 else 32
kernel = tilelang.compile( o, lse = flashattn_fwd(BATCH, H, N_CTX, D_HEAD, causal, block_M, block_N)(q, k, v)
flashattn_fwd(BATCH, H, N_CTX, D_HEAD, causal, block_M, block_N),
out_idx=[3, 4],
target="cuda",
execution_backend="cython")
o, lse = kernel(q, k, v)
ctx.save_for_backward(q, k, v, o, lse) ctx.save_for_backward(q, k, v, o, lse)
ctx.causal = causal ctx.causal = causal
return o return o
...@@ -245,6 +243,7 @@ class _attention(torch.autograd.Function): ...@@ -245,6 +243,7 @@ class _attention(torch.autograd.Function):
@staticmethod @staticmethod
def backward(ctx, do): def backward(ctx, do):
q, k, v, o, lse = ctx.saved_tensors q, k, v, o, lse = ctx.saved_tensors
BATCH, N_CTX, H, D_HEAD = q.shape
def maybe_contiguous(x): def maybe_contiguous(x):
if x.stride(-1) != 1: if x.stride(-1) != 1:
...@@ -252,25 +251,17 @@ class _attention(torch.autograd.Function): ...@@ -252,25 +251,17 @@ class _attention(torch.autograd.Function):
return x return x
do, q, k, v, o = [maybe_contiguous(x) for x in (do, q, k, v, o)] do, q, k, v, o = [maybe_contiguous(x) for x in (do, q, k, v, o)]
block_M = 128 block_M = 64
block_N = 128 if D_HEAD <= 64 else 32 block_N = 64 if D_HEAD <= 64 else 32
kernel_prep = tilelang.compile( kernel_prep = flashattn_bwd_preprocess(BATCH, H, N_CTX, D_HEAD)
flashattn_bwd_preprocess(BATCH, H, N_CTX, D_HEAD), kernel_post = flashattn_bwd_postprocess(BATCH, H, N_CTX, D_HEAD)
out_idx=[2],
target="cuda",
execution_backend="cython")
kernel_post = tilelang.compile(
flashattn_bwd_postprocess(BATCH, H, N_CTX, D_HEAD),
out_idx=[1],
target="cuda",
execution_backend="cython")
delta = kernel_prep(o, do) delta = kernel_prep(o, do)
kernel = tilelang.compile( kernel = flashattn_bwd(BATCH, H, N_CTX, D_HEAD, ctx.causal, block_M, block_N)
flashattn_bwd(BATCH, H, N_CTX, D_HEAD, ctx.causal, block_M, block_N), shape = [BATCH, N_CTX, H, D_HEAD]
out_idx=[6, 7, 8], dq = torch.zeros(shape, dtype=torch.float32, device=q.device)
target="cuda", dk = torch.empty(shape, dtype=torch.float16, device=q.device)
execution_backend="cython") dv = torch.empty(shape, dtype=torch.float16, device=q.device)
dq, dk, dv = kernel(q, k, v, do, lse, delta) kernel(q, k, v, do, lse, delta, dq, dk, dv)
dq = kernel_post(dq) dq = kernel_post(dq)
return dq, dk, dv, None return dq, dk, dv, None
...@@ -292,19 +283,16 @@ def ref_program(Q, K, V, is_causal): ...@@ -292,19 +283,16 @@ def ref_program(Q, K, V, is_causal):
return output return output
if __name__ == "__main__": def main(
parser = argparse.ArgumentParser() BATCH: int = 8,
parser.add_argument('--batch', type=int, default=8, help='Batch size') H: int = 32,
parser.add_argument('--h', type=int, default=32, help='Number of heads') N_CTX: int = 1024,
parser.add_argument('--n_ctx', type=int, default=1024, help='Context size') D_HEAD: int = 64,
parser.add_argument('--d_head', type=int, default=64, help='Head dimension') causal: bool = False,
parser.add_argument('--casual', type=bool, default=False, help='Casual flag') ):
args = parser.parse_args()
BATCH, H, N_CTX, D_HEAD = args.batch, args.h, args.n_ctx, args.d_head
casual = args.casual
flops_per_matmul = 2.0 * BATCH * H * N_CTX * N_CTX * D_HEAD flops_per_matmul = 2.0 * BATCH * H * N_CTX * N_CTX * D_HEAD
total_flops = 5 * flops_per_matmul total_flops = 5 * flops_per_matmul
if casual: if causal:
total_flops *= 0.5 total_flops *= 0.5
Q = ( Q = (
torch.empty(BATCH, N_CTX, H, D_HEAD, dtype=torch.half, torch.empty(BATCH, N_CTX, H, D_HEAD, dtype=torch.half,
...@@ -312,13 +300,13 @@ if __name__ == "__main__": ...@@ -312,13 +300,13 @@ if __name__ == "__main__":
K = torch.empty_like(Q).normal_().requires_grad_() K = torch.empty_like(Q).normal_().requires_grad_()
V = torch.empty_like(Q).normal_().requires_grad_() V = torch.empty_like(Q).normal_().requires_grad_()
dO = torch.randn_like(Q) dO = torch.randn_like(Q)
O = attention(Q, K, V, casual) O = attention(Q, K, V, causal)
O.backward(dO, retain_graph=True) O.backward(dO, retain_graph=True)
dQ, Q.grad = Q.grad.clone(), None dQ, Q.grad = Q.grad.clone(), None
dK, K.grad = K.grad.clone(), None dK, K.grad = K.grad.clone(), None
dV, V.grad = V.grad.clone(), None dV, V.grad = V.grad.clone(), None
O_ref = ref_program(Q, K, V, casual) O_ref = ref_program(Q, K, V, causal)
O_ref.backward(dO, retain_graph=True) O_ref.backward(dO, retain_graph=True)
dQ_ref, Q.grad = Q.grad.clone(), None dQ_ref, Q.grad = Q.grad.clone(), None
dK_ref, K.grad = K.grad.clone(), None dK_ref, K.grad = K.grad.clone(), None
...@@ -343,3 +331,14 @@ if __name__ == "__main__": ...@@ -343,3 +331,14 @@ if __name__ == "__main__":
latency = do_bench(run1, warmup=500) latency = do_bench(run1, warmup=500)
print("tilelang: {:.2f} ms".format(latency)) print("tilelang: {:.2f} ms".format(latency))
print("tilelang: {:.2f} TFlops".format(total_flops / latency * 1e-9)) print("tilelang: {:.2f} TFlops".format(total_flops / latency * 1e-9))
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--batch', type=int, default=8, help='Batch size')
parser.add_argument('--h', type=int, default=32, help='Number of heads')
parser.add_argument('--n_ctx', type=int, default=1024, help='Context size')
parser.add_argument('--d_head', type=int, default=64, help='Head dimension')
parser.add_argument('--causal', type=bool, default=False, help='Causal flag')
args = parser.parse_args()
main(args.batch, args.h, args.n_ctx, args.d_head, args.causal)
examples/flash_attention/example_mha_bwd_wgmma_pipelined.py
View file @ 24403aea
import torch import torch
import torch.nn.functional as F import torch.nn.functional as F
import tilelang import tilelang
from tilelang import cached
from tilelang.autotuner import * from tilelang.autotuner import *
import tilelang.language as T import tilelang.language as T
import argparse import argparse
def flashattn_fwd(batch, heads, seq_len, dim, is_casual, block_M, block_N): @tilelang.jit(out_idx=[3, 4])
def flashattn_fwd(batch, heads, seq_len, dim, is_causal, block_M, block_N):
scale = (1.0 / dim)**0.5 * 1.44269504 # log2(e) scale = (1.0 / dim)**0.5 * 1.44269504 # log2(e)
shape = [batch, seq_len, heads, dim] shape = [batch, seq_len, heads, dim]
dtype = "float16" dtype = "float16"
...@@ -45,10 +45,10 @@ def flashattn_fwd(batch, heads, seq_len, dim, is_casual, block_M, block_N): ...@@ -45,10 +45,10 @@ def flashattn_fwd(batch, heads, seq_len, dim, is_casual, block_M, block_N):
# Q_local[i, j] *= scale # Q_local[i, j] *= scale
loop_range = ( loop_range = (
T.ceildiv( T.ceildiv(
(bx + 1) * block_M, block_N) if is_casual else T.ceildiv(seq_len, block_N)) (bx + 1) * block_M, block_N) if is_causal else T.ceildiv(seq_len, block_N))
for k in T.Pipelined(loop_range, num_stages=1): for k in T.Pipelined(loop_range, num_stages=1):
T.copy(K[bz, k * block_N:(k + 1) * block_N, by, :], K_shared) T.copy(K[bz, k * block_N:(k + 1) * block_N, by, :], K_shared)
if is_casual: if is_causal:
for i, j in T.Parallel(block_M, block_N): for i, j in T.Parallel(block_M, block_N):
acc_s[i, j] = T.if_then_else(bx * block_M + i >= k * block_N + j, 0, acc_s[i, j] = T.if_then_else(bx * block_M + i >= k * block_N + j, 0,
-T.infinity(acc_s.dtype)) -T.infinity(acc_s.dtype))
...@@ -79,6 +79,7 @@ def flashattn_fwd(batch, heads, seq_len, dim, is_casual, block_M, block_N): ...@@ -79,6 +79,7 @@ def flashattn_fwd(batch, heads, seq_len, dim, is_casual, block_M, block_N):
return flash_fwd return flash_fwd
@tilelang.jit(out_idx=[2])
def flashattn_bwd_preprocess(batch, heads, seq_len, dim): def flashattn_bwd_preprocess(batch, heads, seq_len, dim):
dtype = "float16" dtype = "float16"
accum_dtype = "float" accum_dtype = "float"
...@@ -114,6 +115,7 @@ def make_dq_layout(dQ): ...@@ -114,6 +115,7 @@ def make_dq_layout(dQ):
lambda b, l, h, d: [b, l // 8, h, d // 8, (d % 2), 4 * (l % 8) + (d % 8) // 2]) lambda b, l, h, d: [b, l // 8, h, d // 8, (d % 2), 4 * (l % 8) + (d % 8) // 2])
@tilelang.jit(out_idx=[1])
def flashattn_bwd_postprocess(batch, heads, seq_len, dim): def flashattn_bwd_postprocess(batch, heads, seq_len, dim):
dtype = "float16" dtype = "float16"
accum_dtype = "float" accum_dtype = "float"
...@@ -135,6 +137,7 @@ def flashattn_bwd_postprocess(batch, heads, seq_len, dim): ...@@ -135,6 +137,7 @@ def flashattn_bwd_postprocess(batch, heads, seq_len, dim):
return flash_bwd_post return flash_bwd_post
@tilelang.jit
def flashattn_bwd(batch, heads, seq_len, dim, is_casual, block_M, block_N): def flashattn_bwd(batch, heads, seq_len, dim, is_casual, block_M, block_N):
sm_scale = (1.0 / dim)**0.5 sm_scale = (1.0 / dim)**0.5
scale = (1.0 / dim)**0.5 * 1.44269504 # log2(e) scale = (1.0 / dim)**0.5 * 1.44269504 # log2(e)
...@@ -244,7 +247,7 @@ class _attention(torch.autograd.Function): ...@@ -244,7 +247,7 @@ class _attention(torch.autograd.Function):
BATCH, N_CTX, H, D_HEAD = q.shape BATCH, N_CTX, H, D_HEAD = q.shape
block_M = 64 block_M = 64
block_N = 64 if D_HEAD <= 128 else 32 block_N = 64 if D_HEAD <= 128 else 32
mod = cached(flashattn_fwd, [3, 4], BATCH, H, N_CTX, D_HEAD, causal, block_M, block_N) mod = flashattn_fwd(BATCH, H, N_CTX, D_HEAD, causal, block_M, block_N)
o, lse = mod(q, k, v) o, lse = mod(q, k, v)
ctx.save_for_backward(q, k, v, o, lse) ctx.save_for_backward(q, k, v, o, lse)
ctx.causal = causal ctx.causal = causal
...@@ -253,6 +256,7 @@ class _attention(torch.autograd.Function): ...@@ -253,6 +256,7 @@ class _attention(torch.autograd.Function):
@staticmethod @staticmethod
def backward(ctx, do): def backward(ctx, do):
q, k, v, o, lse = ctx.saved_tensors q, k, v, o, lse = ctx.saved_tensors
BATCH, N_CTX, H, D_HEAD = q.shape
def maybe_contiguous(x): def maybe_contiguous(x):
if x.stride(-1) != 1: if x.stride(-1) != 1:
...@@ -260,14 +264,17 @@ class _attention(torch.autograd.Function): ...@@ -260,14 +264,17 @@ class _attention(torch.autograd.Function):
return x return x
do, q, k, v, o = [maybe_contiguous(x) for x in (do, q, k, v, o)] do, q, k, v, o = [maybe_contiguous(x) for x in (do, q, k, v, o)]
block_M = 128 block_M = 64
block_N = 128 if D_HEAD <= 64 else 32 block_N = 64 if D_HEAD <= 64 else 32
mod_prep = cached(flashattn_bwd_preprocess, [2], BATCH, H, N_CTX, D_HEAD) mod_prep = flashattn_bwd_preprocess(BATCH, H, N_CTX, D_HEAD)
mod_post = cached(flashattn_bwd_postprocess, [1], BATCH, H, N_CTX, D_HEAD) mod_post = flashattn_bwd_postprocess(BATCH, H, N_CTX, D_HEAD)
delta = mod_prep(o, do) delta = mod_prep(o, do)
mod = cached(flashattn_bwd, [6, 7, 8], BATCH, H, N_CTX, D_HEAD, ctx.causal, block_M, mod = flashattn_bwd(BATCH, H, N_CTX, D_HEAD, ctx.causal, block_M, block_N)
block_N) shape = [BATCH, N_CTX, H, D_HEAD]
dq, dk, dv = mod(q, k, v, do, lse, delta) dq = torch.zeros(shape, dtype=torch.float32, device=q.device)
dk = torch.empty(shape, dtype=torch.float16, device=q.device)
dv = torch.empty(shape, dtype=torch.float16, device=q.device)
mod(q, k, v, do, lse, delta, dq, dk, dv)
dq = mod_post(dq) dq = mod_post(dq)
return dq, dk, dv, None return dq, dk, dv, None
...@@ -289,19 +296,16 @@ def ref_program(Q, K, V, is_causal): ...@@ -289,19 +296,16 @@ def ref_program(Q, K, V, is_causal):
return output return output
if __name__ == "__main__": def main(
parser = argparse.ArgumentParser() BATCH: int = 8,
parser.add_argument('--batch', type=int, default=8, help='Batch size') H: int = 32,
parser.add_argument('--h', type=int, default=32, help='Number of heads') N_CTX: int = 1024,
parser.add_argument('--n_ctx', type=int, default=1024, help='Context size') D_HEAD: int = 64,
parser.add_argument('--d_head', type=int, default=64, help='Head dimension') causal: bool = False,
parser.add_argument('--casual', type=bool, default=False, help='Casual flag') ):
args = parser.parse_args()
BATCH, H, N_CTX, D_HEAD = args.batch, args.h, args.n_ctx, args.d_head
casual = args.casual
flops_per_matmul = 2.0 * BATCH * H * N_CTX * N_CTX * D_HEAD flops_per_matmul = 2.0 * BATCH * H * N_CTX * N_CTX * D_HEAD
total_flops = 5 * flops_per_matmul total_flops = 5 * flops_per_matmul
if casual: if causal:
total_flops *= 0.5 total_flops *= 0.5
Q = ( Q = (
torch.empty(BATCH, N_CTX, H, D_HEAD, dtype=torch.half, torch.empty(BATCH, N_CTX, H, D_HEAD, dtype=torch.half,
...@@ -309,13 +313,13 @@ if __name__ == "__main__": ...@@ -309,13 +313,13 @@ if __name__ == "__main__":
K = torch.empty_like(Q).normal_().requires_grad_() K = torch.empty_like(Q).normal_().requires_grad_()
V = torch.empty_like(Q).normal_().requires_grad_() V = torch.empty_like(Q).normal_().requires_grad_()
dO = torch.randn_like(Q) dO = torch.randn_like(Q)
O = attention(Q, K, V, casual) O = attention(Q, K, V, causal)
O.backward(dO, retain_graph=True) O.backward(dO, retain_graph=True)
dQ, Q.grad = Q.grad.clone(), None dQ, Q.grad = Q.grad.clone(), None
dK, K.grad = K.grad.clone(), None dK, K.grad = K.grad.clone(), None
dV, V.grad = V.grad.clone(), None dV, V.grad = V.grad.clone(), None
O_ref = ref_program(Q, K, V, casual) O_ref = ref_program(Q, K, V, causal)
O_ref.backward(dO, retain_graph=True) O_ref.backward(dO, retain_graph=True)
dQ_ref, Q.grad = Q.grad.clone(), None dQ_ref, Q.grad = Q.grad.clone(), None
dK_ref, K.grad = K.grad.clone(), None dK_ref, K.grad = K.grad.clone(), None
...@@ -340,3 +344,14 @@ if __name__ == "__main__": ...@@ -340,3 +344,14 @@ if __name__ == "__main__":
latency = do_bench(run1, warmup=500) latency = do_bench(run1, warmup=500)
print("tilelang: {:.2f} ms".format(latency)) print("tilelang: {:.2f} ms".format(latency))
print("tilelang: {:.2f} TFlops".format(total_flops / latency * 1e-9)) print("tilelang: {:.2f} TFlops".format(total_flops / latency * 1e-9))
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--batch', type=int, default=8, help='Batch size')
parser.add_argument('--h', type=int, default=32, help='Number of heads')
parser.add_argument('--n_ctx', type=int, default=1024, help='Context size')
parser.add_argument('--d_head', type=int, default=64, help='Head dimension')
parser.add_argument('--causal', type=bool, default=False, help='Causal flag')
args = parser.parse_args()
main(args.batch, args.h, args.n_ctx, args.d_head, args.causal)
examples/flash_attention/example_mha_fwd_bhsd.py
View file @ 24403aea
...@@ -182,43 +182,55 @@ def ref_program(Q, K, V, is_causal): ...@@ -182,43 +182,55 @@ def ref_program(Q, K, V, is_causal):
return output return output
if __name__ == "__main__": def main(
parser = argparse.ArgumentParser() batch: int = 1,
parser.add_argument('--batch', type=int, default=1, help='batch size') heads: int = 1,
parser.add_argument('--heads', type=int, default=1, help='heads') seq_q: int = 256,
parser.add_argument('--seq_q', type=int, default=256, help='query sequence length') seq_kv: int = 256,
parser.add_argument('--seq_kv', type=int, default=256, help='key/value sequence length') dim: int = 64,
parser.add_argument('--dim', type=int, default=64, help='dim') is_causal: bool = False,
parser.add_argument('--is_causal', action='store_true', help='causal') tune: bool = False,
parser.add_argument('--tune', action='store_true', help='tune configs') ):
args = parser.parse_args()
batch, heads, seq_q, seq_kv, dim, is_causal = args.batch, args.heads, args.seq_q, args.seq_kv, args.dim, args.is_causal
flops_per_matmul = 2.0 * batch * heads * seq_q * seq_kv * dim flops_per_matmul = 2.0 * batch * heads * seq_q * seq_kv * dim
total_flops = 2 * flops_per_matmul total_flops = 2 * flops_per_matmul
if is_causal: if is_causal:
total_flops *= 0.5 total_flops *= 0.5
if (not args.tune): if (not tune):
program = flashattn( program = flashattn(
batch, heads, seq_q, seq_kv, dim, is_causal, tune=args.tune)( batch, heads, seq_q, seq_kv, dim, is_causal, tune=tune)(
block_M=64, block_N=64, num_stages=1, threads=128) block_M=64, block_N=64, num_stages=1, threads=128)
ref_program = partial(ref_program, is_causal=is_causal) ref_program_processed = partial(ref_program, is_causal=is_causal)
kernel = tilelang.compile(program, out_idx=[3]) kernel = tilelang.compile(program, out_idx=[3])
profiler = kernel.get_profiler() profiler = kernel.get_profiler()
profiler.assert_allclose(ref_program, rtol=0.01, atol=0.01) profiler.assert_allclose(ref_program_processed, rtol=0.01, atol=0.01)
print("All checks pass.") print("All checks pass.")
latency = profiler.do_bench(ref_program, warmup=500) latency = profiler.do_bench(ref_program_processed, warmup=500)
print("Ref: {:.2f} ms".format(latency)) print("Ref: {:.2f} ms".format(latency))
print("Ref: {:.2f} TFlops".format(total_flops / latency * 1e-9)) print("Ref: {:.2f} TFlops".format(total_flops / latency * 1e-9))
latency = profiler.do_bench(warmup=500) latency = profiler.do_bench(warmup=500)
print("Tile-lang: {:.2f} ms".format(latency)) print("Tile-lang: {:.2f} ms".format(latency))
print("Tile-lang: {:.2f} TFlops".format(total_flops / latency * 1e-9)) print("Tile-lang: {:.2f} TFlops".format(total_flops / latency * 1e-9))
else: else:
best_result = flashattn(batch, heads, seq_q, seq_kv, dim, is_causal, tune=args.tune) best_result = flashattn(batch, heads, seq_q, seq_kv, dim, is_causal, tune=tune)
best_latency = best_result.latency best_latency = best_result.latency
best_config = best_result.config best_config = best_result.config
ref_latency = best_result.ref_latency ref_latency = best_result.ref_latency
print(f"Best latency: {best_latency}") print(f"Best latency: {best_latency}")
print(f"Best TFlops: {total_flops / best_latency * 1e-9}") print(f"Best TFlops: {total_flops / best_latency * 1e-9}")
print(f"Best config: {best_config}") print(f"Best config: {best_config}")
print(f"Ref latency: {ref_latency}")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--batch', type=int, default=1, help='batch size')
parser.add_argument('--heads', type=int, default=1, help='heads')
parser.add_argument('--seq_q', type=int, default=256, help='query sequence length')
parser.add_argument('--seq_kv', type=int, default=256, help='key/value sequence length')
parser.add_argument('--dim', type=int, default=64, help='dim')
parser.add_argument('--is_causal', action='store_true', help='causal')
parser.add_argument('--tune', action='store_true', help='tune configs')
args = parser.parse_args()
main(args.batch, args.heads, args.seq_q, args.seq_kv, args.dim, args.is_causal, args.tune)
examples/flash_attention/example_mha_fwd_bhsd_wgmma_pipelined.py
View file @ 24403aea
...@@ -187,43 +187,55 @@ def ref_program(Q, K, V, is_causal): ...@@ -187,43 +187,55 @@ def ref_program(Q, K, V, is_causal):
return output return output
if __name__ == "__main__": def main(
parser = argparse.ArgumentParser() batch: int = 8,
parser.add_argument('--batch', type=int, default=8, help='batch size') heads: int = 32,
parser.add_argument('--heads', type=int, default=32, help='heads') seq_q: int = 4096,
parser.add_argument('--seq_q', type=int, default=4096, help='query sequence length') seq_kv: int = 4096,
parser.add_argument('--seq_kv', type=int, default=4096, help='key/value sequence length') dim: int = 128,
parser.add_argument('--dim', type=int, default=128, help='dim') is_causal: bool = False,
parser.add_argument('--is_causal', action='store_true', help='causal') tune: bool = False,
parser.add_argument('--tune', action='store_true', help='tune configs') ):
args = parser.parse_args()
batch, heads, seq_q, seq_kv, dim, is_causal = args.batch, args.heads, args.seq_q, args.seq_kv, args.dim, args.is_causal
flops_per_matmul = 2.0 * batch * heads * seq_q * seq_kv * dim flops_per_matmul = 2.0 * batch * heads * seq_q * seq_kv * dim
total_flops = 2 * flops_per_matmul total_flops = 2 * flops_per_matmul
if is_causal: if is_causal:
total_flops *= 0.5 total_flops *= 0.5
if (not args.tune): if (not tune):
program = flashattn( program = flashattn(
batch, heads, seq_q, seq_kv, dim, is_causal, tune=args.tune)( batch, heads, seq_q, seq_kv, dim, is_causal, tune=tune)(
block_M=128, block_N=128, num_stages=2, threads=256) block_M=128, block_N=128, num_stages=2, threads=256)
ref_program = partial(ref_program, is_causal=is_causal) ref_program_processed = partial(ref_program, is_causal=is_causal)
kernel = tilelang.compile(program, out_idx=[3]) kernel = tilelang.compile(program, out_idx=[3])
profiler = kernel.get_profiler() profiler = kernel.get_profiler()
profiler.assert_allclose(ref_program, rtol=0.01, atol=0.01) profiler.assert_allclose(ref_program_processed, rtol=0.01, atol=0.01)
print("All checks pass.") print("All checks pass.")
latency = profiler.do_bench(ref_program, warmup=500) latency = profiler.do_bench(ref_program_processed, warmup=500)
print("Ref: {:.2f} ms".format(latency)) print("Ref: {:.2f} ms".format(latency))
print("Ref: {:.2f} TFlops".format(total_flops / latency * 1e-9)) print("Ref: {:.2f} TFlops".format(total_flops / latency * 1e-9))
latency = profiler.do_bench(warmup=500) latency = profiler.do_bench(warmup=500)
print("Tile-lang: {:.2f} ms".format(latency)) print("Tile-lang: {:.2f} ms".format(latency))
print("Tile-lang: {:.2f} TFlops".format(total_flops / latency * 1e-9)) print("Tile-lang: {:.2f} TFlops".format(total_flops / latency * 1e-9))
else: else:
best_result = flashattn(batch, heads, seq_q, seq_kv, dim, is_causal, tune=args.tune) best_result = flashattn(batch, heads, seq_q, seq_kv, dim, is_causal, tune=tune)
best_latency = best_result.latency best_latency = best_result.latency
best_config = best_result.config best_config = best_result.config
ref_latency = best_result.ref_latency ref_latency = best_result.ref_latency
print(f"Best latency: {best_latency}") print(f"Best latency: {best_latency}")
print(f"Best TFlops: {total_flops / best_latency * 1e-9}") print(f"Best TFlops: {total_flops / best_latency * 1e-9}")
print(f"Best config: {best_config}") print(f"Best config: {best_config}")
print(f"Ref latency: {ref_latency}")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--batch', type=int, default=8, help='batch size')
parser.add_argument('--heads', type=int, default=32, help='heads')
parser.add_argument('--seq_q', type=int, default=4096, help='query sequence length')
parser.add_argument('--seq_kv', type=int, default=4096, help='key/value sequence length')
parser.add_argument('--dim', type=int, default=128, help='dim')
parser.add_argument('--is_causal', action='store_true', help='causal')
parser.add_argument('--tune', action='store_true', help='tune configs')
args = parser.parse_args()
main(args.batch, args.heads, args.seq_q, args.seq_kv, args.dim, args.is_causal, args.tune)
examples/flash_attention/example_mha_fwd_bshd.py
View file @ 24403aea
...@@ -177,41 +177,52 @@ def ref_program(Q, K, V, is_causal): ...@@ -177,41 +177,52 @@ def ref_program(Q, K, V, is_causal):
return output return output
if __name__ == "__main__": def main(
parser = argparse.ArgumentParser() batch: int = 8,
parser.add_argument('--batch', type=int, default=8, help='batch size') heads: int = 32,
parser.add_argument('--heads', type=int, default=32, help='heads') seq_len: int = 4096,
parser.add_argument('--seq_len', type=int, default=4096, help='sequence length') dim: int = 128,
parser.add_argument('--dim', type=int, default=128, help='dim') is_causal: bool = False,
parser.add_argument('--is_causal', action='store_true', help='causal') tune: bool = False,
parser.add_argument('--tune', action='store_true', help='tune configs') ):
args = parser.parse_args()
batch, heads, seq_len, dim, is_causal = args.batch, args.heads, args.seq_len, args.dim, args.is_causal
flops_per_matmul = 2.0 * batch * heads * seq_len * seq_len * dim flops_per_matmul = 2.0 * batch * heads * seq_len * seq_len * dim
total_flops = 2 * flops_per_matmul total_flops = 2 * flops_per_matmul
if is_causal: if is_causal:
total_flops *= 0.5 total_flops *= 0.5
if (not args.tune): if (not tune):
program = flashattn( program = flashattn(
batch, heads, seq_len, dim, is_causal, tune=args.tune)( batch, heads, seq_len, dim, is_causal, tune=tune)(
block_M=128, block_N=128, num_stages=1, threads=128) block_M=128, block_N=128, num_stages=1, threads=128)
ref_program = partial(ref_program, is_causal=is_causal) ref_program_processed = partial(ref_program, is_causal=is_causal)
kernel = tilelang.compile(program, out_idx=[3]) kernel = tilelang.compile(program, out_idx=[3])
profiler = kernel.get_profiler() profiler = kernel.get_profiler()
profiler.assert_allclose(ref_program, rtol=0.01, atol=0.01) profiler.assert_allclose(ref_program_processed, rtol=0.01, atol=0.01)
print("All checks pass.") print("All checks pass.")
latency = profiler.do_bench(ref_program, warmup=500) latency = profiler.do_bench(ref_program_processed, warmup=500)
print("Ref: {:.2f} ms".format(latency)) print("Ref: {:.2f} ms".format(latency))
print("Ref: {:.2f} TFlops".format(total_flops / latency * 1e-9)) print("Ref: {:.2f} TFlops".format(total_flops / latency * 1e-9))
latency = profiler.do_bench(warmup=500) latency = profiler.do_bench(warmup=500)
print("Tile-lang: {:.2f} ms".format(latency)) print("Tile-lang: {:.2f} ms".format(latency))
print("Tile-lang: {:.2f} TFlops".format(total_flops / latency * 1e-9)) print("Tile-lang: {:.2f} TFlops".format(total_flops / latency * 1e-9))
else: else:
best_result = flashattn(batch, heads, seq_len, dim, is_causal, tune=args.tune) best_result = flashattn(batch, heads, seq_len, dim, is_causal, tune=tune)
best_latency = best_result.latency best_latency = best_result.latency
best_config = best_result.config best_config = best_result.config
ref_latency = best_result.ref_latency ref_latency = best_result.ref_latency
print(f"Best latency: {best_latency}") print(f"Best latency: {best_latency}")
print(f"Best TFlops: {total_flops / best_latency * 1e-9}") print(f"Best TFlops: {total_flops / best_latency * 1e-9}")
print(f"Best config: {best_config}") print(f"Best config: {best_config}")
print(f"Ref latency: {ref_latency}")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--batch', type=int, default=8, help='batch size')
parser.add_argument('--heads', type=int, default=32, help='heads')
parser.add_argument('--seq_len', type=int, default=4096, help='sequence length')
parser.add_argument('--dim', type=int, default=128, help='dim')
parser.add_argument('--is_causal', action='store_true', help='causal')
parser.add_argument('--tune', action='store_true', help='tune configs')
args = parser.parse_args()
main(args.batch, args.heads, args.seq_len, args.dim, args.is_causal, args.tune)
examples/flash_attention/example_mha_fwd_bshd_wgmma_pipelined.py
View file @ 24403aea
...@@ -182,41 +182,52 @@ def ref_program(Q, K, V, is_causal): ...@@ -182,41 +182,52 @@ def ref_program(Q, K, V, is_causal):
return output return output
if __name__ == "__main__": def main(
parser = argparse.ArgumentParser() batch: int = 8,
parser.add_argument('--batch', type=int, default=8, help='batch size') heads: int = 32,
parser.add_argument('--heads', type=int, default=32, help='heads') seq_len: int = 4096,
parser.add_argument('--seq_len', type=int, default=4096, help='sequence length') dim: int = 128,
parser.add_argument('--dim', type=int, default=128, help='dim') is_causal: bool = False,
parser.add_argument('--is_causal', action='store_true', help='causal') tune: bool = False,
parser.add_argument('--tune', action='store_true', help='tune configs') ):
args = parser.parse_args()
batch, heads, seq_len, dim, is_causal = args.batch, args.heads, args.seq_len, args.dim, args.is_causal
flops_per_matmul = 2.0 * batch * heads * seq_len * seq_len * dim flops_per_matmul = 2.0 * batch * heads * seq_len * seq_len * dim
total_flops = 2 * flops_per_matmul total_flops = 2 * flops_per_matmul
if is_causal: if is_causal:
total_flops *= 0.5 total_flops *= 0.5
if (not args.tune): if (not tune):
program = flashattn( program = flashattn(
batch, heads, seq_len, dim, is_causal, tune=args.tune)( batch, heads, seq_len, dim, is_causal, tune=tune)(
block_M=128, block_N=128, num_stages=2, threads=256) block_M=128, block_N=128, num_stages=2, threads=256)
ref_program = partial(ref_program, is_causal=is_causal) ref_program_processed = partial(ref_program, is_causal=is_causal)
kernel = tilelang.compile(program, out_idx=[3]) kernel = tilelang.compile(program, out_idx=[3])
profiler = kernel.get_profiler(tensor_supply_type=tilelang.TensorSupplyType.Normal) profiler = kernel.get_profiler(tensor_supply_type=tilelang.TensorSupplyType.Normal)
profiler.assert_allclose(ref_program, rtol=0.01, atol=0.01) profiler.assert_allclose(ref_program_processed, rtol=0.01, atol=0.01)
print("All checks pass.") print("All checks pass.")
latency = profiler.do_bench(ref_program, warmup=500) latency = profiler.do_bench(ref_program_processed, warmup=500)
print("Ref: {:.2f} ms".format(latency)) print("Ref: {:.2f} ms".format(latency))
print("Ref: {:.2f} TFlops".format(total_flops / latency * 1e-9)) print("Ref: {:.2f} TFlops".format(total_flops / latency * 1e-9))
latency = profiler.do_bench(warmup=500) latency = profiler.do_bench(warmup=500)
print("Tile-lang: {:.2f} ms".format(latency)) print("Tile-lang: {:.2f} ms".format(latency))
print("Tile-lang: {:.2f} TFlops".format(total_flops / latency * 1e-9)) print("Tile-lang: {:.2f} TFlops".format(total_flops / latency * 1e-9))
else: else:
best_result = flashattn(batch, heads, seq_len, dim, is_causal, tune=args.tune) best_result = flashattn(batch, heads, seq_len, dim, is_causal, tune=tune)
best_latency = best_result.latency best_latency = best_result.latency
best_config = best_result.config best_config = best_result.config
ref_latency = best_result.ref_latency ref_latency = best_result.ref_latency
print(f"Best latency: {best_latency}") print(f"Best latency: {best_latency}")
print(f"Best TFlops: {total_flops / best_latency * 1e-9}") print(f"Best TFlops: {total_flops / best_latency * 1e-9}")
print(f"Best config: {best_config}") print(f"Best config: {best_config}")
print(f"Ref latency: {ref_latency}")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--batch', type=int, default=8, help='batch size')
parser.add_argument('--heads', type=int, default=32, help='heads')
parser.add_argument('--seq_len', type=int, default=4096, help='sequence length')
parser.add_argument('--dim', type=int, default=128, help='dim')
parser.add_argument('--is_causal', action='store_true', help='causal')
parser.add_argument('--tune', action='store_true', help='tune configs')
args = parser.parse_args()
main(args.batch, args.heads, args.seq_len, args.dim, args.is_causal, args.tune)
examples/flash_attention/example_mha_fwd_varlen.py
View file @ 24403aea
...@@ -46,7 +46,7 @@ def generate_qkv(q, ...@@ -46,7 +46,7 @@ def generate_qkv(q,
assert v.shape == (batch_size, seqlen_k, nheads_k, d) assert v.shape == (batch_size, seqlen_k, nheads_k, d)
if query_padding_mask is not None: if query_padding_mask is not None:
q_unpad, indices_q, cu_seqlens_q, max_seqlen_q = unpad_input(q, query_padding_mask) q_unpad, indices_q, cu_seqlens_q, max_seqlen_q, _ = unpad_input(q, query_padding_mask)
output_pad_fn = lambda output_unpad: pad_input(output_unpad, indices_q, batch_size, seqlen_q output_pad_fn = lambda output_unpad: pad_input(output_unpad, indices_q, batch_size, seqlen_q
) )
else: else:
...@@ -58,8 +58,8 @@ def generate_qkv(q, ...@@ -58,8 +58,8 @@ def generate_qkv(q,
output_unpad, "(b s) h d -> b s h d", b=batch_size) output_unpad, "(b s) h d -> b s h d", b=batch_size)
if key_padding_mask is not None: if key_padding_mask is not None:
k_unpad, indices_k, cu_seqlens_k, max_seqlen_k = unpad_input(k, key_padding_mask) k_unpad, indices_k, cu_seqlens_k, max_seqlen_k, _ = unpad_input(k, key_padding_mask)
v_unpad, _, _, _ = unpad_input(v, key_padding_mask) v_unpad, _, _, _, _ = unpad_input(v, key_padding_mask)
else: else:
k_unpad = rearrange(k, "b s h d -> (b s) h d") k_unpad = rearrange(k, "b s h d -> (b s) h d")
v_unpad = rearrange(v, "b s h d -> (b s) h d") v_unpad = rearrange(v, "b s h d -> (b s) h d")
...@@ -197,6 +197,7 @@ def attention_ref( ...@@ -197,6 +197,7 @@ def attention_ref(
dtype_og = q.dtype dtype_og = q.dtype
if upcast: if upcast:
q, k, v = q.float(), k.float(), v.float() q, k, v = q.float(), k.float(), v.float()
dim = q.shape[-1]
scale = (1.0 / dim)**0.5 # log2(e) scale = (1.0 / dim)**0.5 # log2(e)
k = repeat(k, "b s h d -> b s (h g) d", g=q.shape[2] // k.shape[2]) k = repeat(k, "b s h d -> b s (h g) d", g=q.shape[2] // k.shape[2])
v = repeat(v, "b s h d -> b s (h g) d", g=q.shape[2] // v.shape[2]) v = repeat(v, "b s h d -> b s (h g) d", g=q.shape[2] // v.shape[2])
...@@ -358,15 +359,7 @@ def flashattn(batch_size, UQ, UKV, heads, dim, is_causal): ...@@ -358,15 +359,7 @@ def flashattn(batch_size, UQ, UKV, heads, dim, is_causal):
return kernel_func(block_M, block_N, num_stages, threads) return kernel_func(block_M, block_N, num_stages, threads)
if __name__ == "__main__": def main(batch: int = 2, heads: int = 16, seq_len: int = 256, dim: int = 32):
parser = argparse.ArgumentParser()
parser.add_argument('--batch', type=int, default=2, help='batch size')
parser.add_argument('--heads', type=int, default=16, help='heads')
parser.add_argument('--seq_len', type=int, default=256, help='sequence length')
parser.add_argument('--dim', type=int, default=32, help='dim')
args = parser.parse_args()
batch, heads, seq_len, dim = args.batch, args.heads, args.seq_len, args.dim
flops_per_matmul = 2.0 * batch * heads * seq_len * seq_len * dim flops_per_matmul = 2.0 * batch * heads * seq_len * seq_len * dim
total_flops = 2 * flops_per_matmul total_flops = 2 * flops_per_matmul
...@@ -408,7 +401,7 @@ if __name__ == "__main__": ...@@ -408,7 +401,7 @@ if __name__ == "__main__":
UKV = k_unpad.shape[0] # unpadded query key length UKV = k_unpad.shape[0] # unpadded query key length
program = flashattn(batch, UQ, UKV, heads, dim, causal) program = flashattn(batch, UQ, UKV, heads, dim, causal)
kernel = tilelang.compile(program, out_idx=-1, execution_backend="cython") kernel = tilelang.compile(program, [6])
print(kernel.get_kernel_source()) print(kernel.get_kernel_source())
out_unpad = kernel(q_unpad, k_unpad, v_unpad, cu_seqlens_q, cu_seqlens_k, max_seqlen_q) out_unpad = kernel(q_unpad, k_unpad, v_unpad, cu_seqlens_q, cu_seqlens_k, max_seqlen_q)
...@@ -437,3 +430,14 @@ if __name__ == "__main__": ...@@ -437,3 +430,14 @@ if __name__ == "__main__":
fla_out = output_pad_fn(fla_out_unpad) fla_out = output_pad_fn(fla_out_unpad)
torch.testing.assert_close(out, out_ref, rtol=1e-2, atol=1e-2) torch.testing.assert_close(out, out_ref, rtol=1e-2, atol=1e-2)
print("Assert Equal Passed") print("Assert Equal Passed")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--batch', type=int, default=2, help='batch size')
parser.add_argument('--heads', type=int, default=16, help='heads')
parser.add_argument('--seq_len', type=int, default=256, help='sequence length')
parser.add_argument('--dim', type=int, default=32, help='dim')
args = parser.parse_args()
main(args.batch, args.heads, args.seq_len, args.dim)
examples/flash_attention/example_mha_inference.py
View file @ 24403aea
...@@ -294,7 +294,7 @@ def flash_split_ref(Q, K, V, causal): ...@@ -294,7 +294,7 @@ def flash_split_ref(Q, K, V, causal):
3), gacc_o.to(torch.float16).permute(1, 2, 3, 0, 4) 3), gacc_o.to(torch.float16).permute(1, 2, 3, 0, 4)
if __name__ == "__main__": def main():
BATCH, H, Q_CTX, KV_CTX, D_HEAD = 1, 32, 128, 8192, 128 BATCH, H, Q_CTX, KV_CTX, D_HEAD = 1, 32, 128, 8192, 128
causal = False causal = False
flops_per_matmul = 2.0 * BATCH * H * Q_CTX * KV_CTX * D_HEAD flops_per_matmul = 2.0 * BATCH * H * Q_CTX * KV_CTX * D_HEAD
...@@ -304,15 +304,19 @@ if __name__ == "__main__": ...@@ -304,15 +304,19 @@ if __name__ == "__main__":
BLOCK_M = 128 BLOCK_M = 128
BLOCK_N = 64 # if D_HEAD <= 128 else 32 BLOCK_N = 64 # if D_HEAD <= 128 else 32
program = flashattn(BATCH, H, Q_CTX, KV_CTX, D_HEAD, causal, BLOCK_M, BLOCK_N) program = flashattn(BATCH, H, Q_CTX, KV_CTX, D_HEAD, causal, BLOCK_M, BLOCK_N)
ref_program = partial(ref_program, causal=causal) ref_program_processed = partial(ref_program, causal=causal)
kernel = tilelang.compile(program, out_idx=[5]) kernel = tilelang.compile(program, out_idx=[5])
profiler = kernel.get_profiler(tilelang.TensorSupplyType.Normal) profiler = kernel.get_profiler(tilelang.TensorSupplyType.Normal)
profiler.assert_allclose(ref_program, rtol=0.01, atol=0.01) profiler.assert_allclose(ref_program_processed, rtol=0.01, atol=0.01)
print("All checks passed!") print("All checks passed!")
latency = profiler.do_bench(ref_program, warmup=500) latency = profiler.do_bench(ref_program_processed, warmup=500)
print("{:.2f} ms".format(latency)) print("{:.2f} ms".format(latency))
print("{:.2f} TFlops".format(total_flops / latency * 1e-9)) print("{:.2f} TFlops".format(total_flops / latency * 1e-9))
latency = profiler.do_bench(n_warmup=10, n_repeat=10, profiler="tvm") latency = profiler.do_bench(n_warmup=10, n_repeat=10)
print("{:.4f} ms".format(latency)) print("{:.4f} ms".format(latency))
print("{:.2f} TFlops".format(total_flops / latency * 1e-9)) print("{:.2f} TFlops".format(total_flops / latency * 1e-9))
if __name__ == "__main__":
main()
examples/flash_attention/test_example_flash_attention.py 0 → 100644
View file @ 24403aea
import tilelang.testing
import example_gqa_bwd
import example_mha_bwd
import example_mha_fwd_bhsd_wgmma_pipelined
import example_gqa_fwd_bshd
import example_mha_fwd_bshd
import example_gqa_fwd_bshd_wgmma_pipelined
import example_mha_fwd_bshd_wgmma_pipelined
import example_mha_fwd_varlen
import example_mha_bwd_wgmma_pipelined
import example_mha_inference
import example_mha_fwd_bhsd
@tilelang.testing.requires_cuda
def test_example_gqa_bwd():
example_gqa_bwd.main()
@tilelang.testing.requires_cuda
def test_example_mha_bwd():
example_mha_bwd.main()
@tilelang.testing.requires_cuda
@tilelang.testing.requires_cuda_compute_version_ge(9, 0)
def test_example_mha_bwd_wgmma_pipelined():
example_mha_bwd_wgmma_pipelined.main()
@tilelang.testing.requires_cuda
def test_example_gqa_fwd_bshd_wgmma_pipelined():
example_gqa_fwd_bshd_wgmma_pipelined.main()
@tilelang.testing.requires_cuda
def test_example_gqa_fwd_bshd():
example_gqa_fwd_bshd.main()
@tilelang.testing.requires_cuda
def test_example_mha_fwd_bhsd_wgmma_pipelined():
example_mha_fwd_bhsd_wgmma_pipelined.main()
@tilelang.testing.requires_cuda
def test_example_mha_fwd_bhsd():
example_mha_fwd_bhsd.main()
@tilelang.testing.requires_cuda
def test_example_mha_fwd_bshd_wgmma_pipelined():
example_mha_fwd_bshd_wgmma_pipelined.main()
@tilelang.testing.requires_cuda
def test_example_mha_fwd_bshd():
example_mha_fwd_bshd.main()
@tilelang.testing.requires_cuda
def test_example_mha_fwd_varlen():
example_mha_fwd_varlen.main()
@tilelang.testing.requires_cuda
def test_example_mha_inference():
example_mha_inference.main()
if __name__ == "__main__":
tilelang.testing.main()
tilelang/engine/phase.py
View file @ 24403aea
...@@ -46,6 +46,13 @@ def allow_vectorize(pass_ctx: Optional[PassContext] = None) -> bool: ...@@ -46,6 +46,13 @@ def allow_vectorize(pass_ctx: Optional[PassContext] = None) -> bool:
return not disable_vectorize return not disable_vectorize
def allow_global_thread_synchronization(pass_ctx: Optional[PassContext] = None) -> bool:
if pass_ctx is None:
pass_ctx = tilelang.transform.get_pass_context()
enable_global_thread_sync = pass_ctx.config.get("tir.detect_global_barrier", False)
return enable_global_thread_sync
def LowerAndLegalize(mod: IRModule, target: Target) -> IRModule: def LowerAndLegalize(mod: IRModule, target: Target) -> IRModule:
# Bind the target device information to the module # Bind the target device information to the module
mod = tir.transform.BindTarget(target)(mod) mod = tir.transform.BindTarget(target)(mod)
...@@ -133,10 +140,10 @@ def OptimizeForTarget(mod: IRModule, target: Target) -> IRModule: ...@@ -133,10 +140,10 @@ def OptimizeForTarget(mod: IRModule, target: Target) -> IRModule:
mod = tir.transform.InferFragment()(mod) mod = tir.transform.InferFragment()(mod)
mod = tir.transform.LowerThreadAllreduce()(mod) mod = tir.transform.LowerThreadAllreduce()(mod)
mod = tilelang.transform.LowerHopperIntrin()(mod) mod = tilelang.transform.LowerHopperIntrin()(mod)
mod = tilelang.transform.ThreadSync("global")(mod)
# Global Barrier Synchronization must be applied before # Global Barrier Synchronization must be applied before
# SplitHostDevice pass, as the global barrier # SplitHostDevice pass, as the global barrier
if allow_global_thread_synchronization():
mod = tilelang.transform.ThreadSync("global")(mod) mod = tilelang.transform.ThreadSync("global")(mod)
mod = tilelang.transform.AnnotateDeviceRegions()(mod) mod = tilelang.transform.AnnotateDeviceRegions()(mod)
mod = tir.transform.SplitHostDevice()(mod) mod = tir.transform.SplitHostDevice()(mod)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment