Skip to content

GitLab

Unverified Commit e9a608e2 authored by Wenhao Xie's avatar Wenhao Xie Committed by GitHub
Browse files

[Bugfix][CI] Bug fixing and migrate CI from ada to hopper (#652) 

* fix CI bugs in hopper

* lint fix

* Update bulk_copy.cc

* Refactor bulk copy logic in LowerBulkCopy function

- Removed unnecessary blank lines for improved code readability.
- Enhanced stride validation by checking for null pointers in global stride calculations, ensuring robustness against symbolic strides.
- Updated pass configuration handling in dynamic tile language tests to streamline dynamic alignment and TMA lower pass settings.

* test fix

* ci fix

* Update flash-attention dependencies and clean up example code

- Downgraded `flash-attn` dependency version in `requirements-test.txt` to `<=2.2.0`.
- Removed unused imports and commented-out code in various example files to enhance readability and maintainability.
- Updated the `flashattn` function signature to include default parameters for `block_M`, `block_N`, `num_stages`, and `threads`.
- Cleaned up the `example_mha_fwd_varlen.py` and `example_mha_bwd_wgmma_pipelined.py` fil...
parent 5bd3f942
Show whitespace changes
Inline Side-by-side
Showing with 230 additions and 525 deletions
.github/workflows/ci.yml
View file @ e9a608e2
......@@ -23,8 +23,8 @@ jobs:
- name: Activate virtual environment and install dependencies
run: |
source tilelang_ci/bin/activate
python -m pip install --upgrade pip
if [ -f requirements-dev.txt ]; then python -m pip install -r requirements-dev.txt; fi
python -m pip install --upgrade pip --no-user
if [ -f requirements-dev.txt ]; then python -m pip install -r requirements-dev.txt --no-user; fi
- name: Update submodules recursively
run: git submodule update --init --recursive
......@@ -55,22 +55,24 @@ jobs:
- name: Activate virtual environment and install dependencies
run: |
source tilelang_ci/bin/activate
python -m pip install --upgrade pip
if [ -f requirements-test.txt ]; then PIP_NO_BUILD_ISOLATION=1 python -m pip install -r requirements-test.txt; fi
python -m pip install --upgrade pip --no-user
if [ -f requirements-test.txt ]; then PIP_NO_BUILD_ISOLATION=1 python -m pip install -r requirements-test.txt --no-user; fi
- name: Install project in wheel mode
run: |
source tilelang_ci/bin/activate
python -m pip install .
python -m pip install . --no-user
- name: Run examples
run: |
source tilelang_ci/bin/activate
cd examples
unset PYTHONPATH
python -m pytest **/test*.py
- name: Run tests
run: |
source tilelang_ci/bin/activate
cd testing/python
unset PYTHONPATH
python -m pytest
tvm @ 979c8e7f
Compare db50d4e1...979c8e7f
Subproject commit db50d4e19e8b04677fff3c32dc7fa4c42799f39a
Subproject commit 979c8e7f94473db7d71a41b26ccf51db7e17a734
examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_varlen_indice.py
View file @ e9a608e2
import torch
import torch.nn.functional as F
import tilelang
from tilelang.autotuner import *
import tilelang.language as T
from einops import rearrange, einsum
import argparse
......@@ -71,7 +70,7 @@ def flashattn(batch, heads, heads_kv, dim, dim_v):
loop_range = (blocks_per_split + T.if_then_else(sid < remaining_blocks, 1, 0))
start = blocks_per_split * sid + T.min(sid, remaining_blocks)
has_valid_block = False
# if (start < num_blocks):
for k in T.Pipelined(loop_range, num_stages=num_stages):
i_s = block_indices[bid, cur_kv_head, start + k]
if i_s >= 0:
......@@ -238,23 +237,12 @@ class SparseFlashAttn(torch.nn.Module):
size_one_kv_head,
is_causal_or_local=True,
max_splits=128)
# print("num_split: ", num_split)
# Function to compile
# def compute_actual_num_blocks(block_indices):
# actual_num_blocks = torch.sum(block_indices != -1, dim=-1).to(torch.int32)
# actual_num_blocks = actual_num_blocks[:, 0] # [batch]
# return actual_num_blocks
# compiled_fn = torch.compile(compute_actual_num_blocks)
# actual_num_blocks = compiled_fn(block_indices)
glse = torch.empty((batch, heads, num_split), dtype=torch.float32, device='cuda')
output_partial = torch.empty((batch, heads, num_split, dim_v),
dtype=torch.float32,
device='cuda')
# output = self.kernel(
# query, key, value, block_indices, cache_seqlens,
# actual_num_blocks, glse, output_partial
# )
output = self.kernel(query, key, value, block_indices, cache_seqlens, glse, output_partial)
return output
......@@ -377,8 +365,6 @@ def debug(name, expect, actual, atol=1e-3, rtol=1e-3):
all_close = torch.allclose(expect, actual, atol=atol, rtol=rtol)
print(name + " all_close={}".format(all_close))
if not all_close:
# print(expect[3, 28])
# print(actual[3, 28])
diff = (expect - actual).abs()
print("all_close={}, max={}, min={}, mean={}".format(all_close,
diff.max().item(),
......
examples/convolution/example_convolution.py
View file @ e9a608e2
......@@ -116,9 +116,8 @@ def main(argv=None):
block_k = 32
num_stages = 3
threads = 256
kernel = tilelang.compile(
convolution(N, C, H, W, F, K, S, D, P, block_m, block_n, block_k, num_stages, threads), out_idx=[2])
program = convolution(N, C, H, W, F, K, S, D, P, block_m, block_n, block_k, num_stages, threads)
kernel = tilelang.compile(program, out_idx=[2])
out_c = kernel(a, b)
ref_c = ref_program(S, P, D)(a, b)
......
examples/convolution/test_example_convolution.py
View file @ e9a608e2
......@@ -4,10 +4,15 @@ import example_convolution
import example_convolution_autotune
# TODO(@cy): TMA with convolution must be fixed in future.
@tilelang.testing.requires_cuda
@tilelang.testing.requires_cuda_compute_version_le(8, 9)
def test_example_convolution():
example_convolution.main([])
@tilelang.testing.requires_cuda
@tilelang.testing.requires_cuda_compute_version_le(8, 9)
def test_example_convolution_autotune():
example_convolution_autotune.main()
......
examples/deepseek_deepgemm/example_deepgemm_fp8_2xAcc.py
View file @ e9a608e2
......@@ -9,7 +9,7 @@ from tilelang.utils.tensor import map_torch_type
tilelang.testing.set_random_seed(42)
@tilelang.jit(out_idx=[2])
@tilelang.jit
def tl_gemm(
M,
N,
......
examples/flash_attention/example_mha_bwd_wgmma_pipelined.py
View file @ e9a608e2
......@@ -23,7 +23,6 @@ def flashattn_fwd(batch, heads, seq_len, dim, is_causal, block_M, block_N):
):
with T.Kernel(T.ceildiv(seq_len, block_M), heads, batch, threads=128) as (bx, by, bz):
Q_shared = T.alloc_shared([block_M, dim], dtype)
# Q_local = T.alloc_fragment([block_M, dim], dtype)
K_shared = T.alloc_shared([block_N, dim], dtype)
V_shared = T.alloc_shared([block_N, dim], dtype)
acc_s = T.alloc_fragment([block_M, block_N], accum_dtype)
......@@ -40,9 +39,7 @@ def flashattn_fwd(batch, heads, seq_len, dim, is_causal, block_M, block_N):
T.fill(acc_o, 0)
T.fill(logsum, 0)
T.fill(scores_max, -T.infinity(accum_dtype))
# T.copy(Q_shared, Q_local)
# for i, j in T.Parallel(block_M, dim):
# Q_local[i, j] *= scale
loop_range = (
T.ceildiv(
(bx + 1) * block_M, block_N) if is_causal else T.ceildiv(seq_len, block_N))
......@@ -264,8 +261,8 @@ class _attention(torch.autograd.Function):
return x
do, q, k, v, o = [maybe_contiguous(x) for x in (do, q, k, v, o)]
block_M = 64
block_N = 64 if D_HEAD <= 64 else 32
block_M = 128
block_N = 128 if D_HEAD <= 64 else 32
mod_prep = flashattn_bwd_preprocess(BATCH, H, N_CTX, D_HEAD)
mod_post = flashattn_bwd_postprocess(BATCH, H, N_CTX, D_HEAD)
delta = mod_prep(o, do)
......
examples/flash_attention/example_mha_fwd_varlen.py
View file @ e9a608e2
......@@ -46,7 +46,7 @@ def generate_qkv(q,
assert v.shape == (batch_size, seqlen_k, nheads_k, d)
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
)
else:
......@@ -58,8 +58,8 @@ def generate_qkv(q,
output_unpad, "(b s) h d -> b s h d", b=batch_size)
if key_padding_mask is not None:
k_unpad, indices_k, cu_seqlens_k, max_seqlen_k, _ = unpad_input(k, key_padding_mask)
v_unpad, _, _, _, _ = unpad_input(v, 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)
else:
k_unpad = rearrange(k, "b s h d -> (b s) h d")
v_unpad = rearrange(v, "b s h d -> (b s) h d")
......@@ -218,23 +218,26 @@ def attention_ref(
return output.to(dtype=dtype_og), attention.to(dtype=dtype_og)
def flashattn(batch_size, UQ, UKV, heads, dim, is_causal):
@tilelang.jit(out_idx=[6])
def flashattn(batch_size,
UQ,
UKV,
heads,
dim,
is_causal,
block_M=64,
block_N=64,
num_stages=0,
threads=32):
scale = (1.0 / dim)**0.5 * 1.44269504 # log2(e)
q_shape = [UQ, heads, dim]
k_shape = [UKV, heads, dim]
v_shape = [UKV, heads, dim]
o_shape = [UQ, heads, dim]
block_M = 64
block_N = 64
num_stages = 0
threads = 32
dtype = "float16"
accum_dtype = "float"
@tilelang.jit(out_idx=[6])
def kernel_func(block_M, block_N, num_stages, threads):
@T.prim_func
def main(
Q_unpad: T.Tensor(q_shape, dtype),
......@@ -306,12 +309,7 @@ def flashattn(batch_size, UQ, UKV, heads, dim, is_causal):
k * block_N + j >= k_current_seqlen),
-T.infinity(acc_s.dtype), 0)
T.gemm(
Q_shared,
K_shared,
acc_s,
transpose_B=True,
policy=T.GemmWarpPolicy.FullRow)
T.gemm(Q_shared, K_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullRow)
# Softmax
T.copy(scores_max, scores_max_prev)
......@@ -357,8 +355,6 @@ def flashattn(batch_size, UQ, UKV, heads, dim, is_causal):
return main
return kernel_func(block_M, block_N, num_stages, threads)
def main(batch: int = 2, heads: int = 16, seq_len: int = 256, dim: int = 32):
flops_per_matmul = 2.0 * batch * heads * seq_len * seq_len * dim
......@@ -402,7 +398,6 @@ def main(batch: int = 2, heads: int = 16, seq_len: int = 256, dim: int = 32):
UKV = k_unpad.shape[0] # unpadded query key length
kernel = flashattn(batch, UQ, UKV, heads, dim, causal)
print(kernel.get_kernel_source())
out_unpad = kernel(q_unpad, k_unpad, v_unpad, cu_seqlens_q, cu_seqlens_k, max_seqlen_q)
out = output_pad_fn(out_unpad)
......@@ -429,6 +424,7 @@ def main(batch: int = 2, heads: int = 16, seq_len: int = 256, dim: int = 32):
)
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, fla_out, rtol=1e-2, atol=1e-2)
print("Assert Equal Passed")
......
examples/flash_attention/example_mha_inference.py deleted 100644 → 0
View file @ 5bd3f942
import torch
import torch.nn.functional as F
import tilelang
from tilelang.autotuner import *
import tilelang.language as T
from functools import partial
num_split = 4
@tilelang.jit(out_idx=[5])
def flashattn(batch, heads, seqlen_q, seqlen_kv, dim, is_causal, block_M, block_N):
scale = (1.0 / dim)**0.5 * 1.44269504 # log2(e)
shape_q = [batch, seqlen_q, heads, dim]
shape_kv = [batch, seqlen_kv, heads, dim]
part_shape = [batch, seqlen_q, heads, num_split, dim]
dtype = "float16"
accum_dtype = "float"
@T.macro
def MMA0(
K: T.Tensor(shape_kv, dtype),
Q_shared: T.SharedBuffer([block_M, dim], dtype),
K_shared: T.SharedBuffer([block_N, dim], dtype),
acc_s: T.FragmentBuffer([block_M, block_N], accum_dtype),
k: T.int32,
mid: T.int32,
hid: T.int32,
bid: T.int32,
sid: T.int32,
):
T.copy(
K[bid, (seqlen_kv // num_split) * sid + k * block_N:(seqlen_kv // num_split) * sid +
(k + 1) * block_N, hid, :], K_shared)
# TODO: Handle causal split case
if is_causal:
for i, j in T.Parallel(block_M, block_N):
acc_s[i, j] = T.if_then_else(mid * block_M + i >= k * block_N + j, 0,
-T.infinity(acc_s.dtype))
else:
T.clear(acc_s)
T.gemm(Q_shared, K_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullRow)
@T.macro
def MMA1(
V: T.Tensor(shape_kv, dtype),
V_shared: T.SharedBuffer([block_M, dim], dtype),
acc_s_cast: T.FragmentBuffer([block_M, block_N], dtype),
acc_o: T.FragmentBuffer([block_M, dim], accum_dtype),
k: T.int32,
hid: T.int32,
bid: T.int32,
sid: T.int32,
):
T.copy(
V[bid, (seqlen_kv // num_split) * sid + k * block_N:(seqlen_kv // num_split) * sid +
(k + 1) * block_N, hid, :], V_shared)
T.gemm(acc_s_cast, V_shared, acc_o, policy=T.GemmWarpPolicy.FullRow)
@T.macro
def Softmax(
acc_s: T.FragmentBuffer([block_M, block_N], accum_dtype),
acc_s_cast: T.FragmentBuffer([block_M, block_N], dtype),
scores_max: T.FragmentBuffer([block_M], accum_dtype),
scores_max_prev: T.FragmentBuffer([block_M], accum_dtype),
scores_scale: T.FragmentBuffer([block_M], accum_dtype),
scores_sum: T.FragmentBuffer([block_M], accum_dtype),
logsum: T.FragmentBuffer([block_M], accum_dtype),
):
T.copy(scores_max, scores_max_prev)
T.fill(scores_max, -T.infinity(accum_dtype))
T.reduce_max(acc_s, scores_max, dim=1, clear=False)
# To do causal softmax, we need to set the scores_max to 0 if it is -inf
# This process is called Check_inf in FlashAttention3 code, and it only need to be done
# in the first ceil_div(kBlockM, kBlockN) steps.
# for i in T.Parallel(block_M):
# scores_max[i] = T.if_then_else(scores_max[i] == -T.infinity(accum_dtype), 0, scores_max[i])
for i in T.Parallel(block_M):
scores_scale[i] = T.exp2(scores_max_prev[i] * scale - scores_max[i] * scale)
for i, j in T.Parallel(block_M, block_N):
# Instead of computing exp(x - max), we compute exp2(x * log_2(e) -
# max * log_2(e)) This allows the compiler to use the ffma
# instruction instead of fadd and fmul separately.
acc_s[i, j] = T.exp2(acc_s[i, j] * scale - scores_max[i] * scale)
T.reduce_sum(acc_s, scores_sum, dim=1)
for i in T.Parallel(block_M):
logsum[i] = logsum[i] * scores_scale[i] + scores_sum[i]
T.copy(acc_s, acc_s_cast)
@T.macro
def Rescale(
acc_o: T.FragmentBuffer([block_M, dim], accum_dtype),
scores_scale: T.FragmentBuffer([block_M], accum_dtype),
):
for i, j in T.Parallel(block_M, dim):
acc_o[i, j] *= scores_scale[i]
@T.macro
def flash_attn_split(
Q: T.Tensor(shape_q, dtype),
K: T.Tensor(shape_kv, dtype),
V: T.Tensor(shape_kv, dtype),
glse: T.Tensor([batch, heads, num_split, seqlen_q], dtype),
Output_partial: T.Tensor(part_shape, dtype),
):
with T.Kernel(
T.ceildiv(seqlen_q, block_M), heads * batch, num_split,
threads=128) as (bx, by, bz):
Q_shared = T.alloc_shared([block_M, dim], dtype)
K_shared = T.alloc_shared([block_N, dim], dtype)
V_shared = T.alloc_shared([block_N, dim], dtype)
O_shared = T.alloc_shared([block_M, dim], dtype)
acc_s = T.alloc_fragment([block_M, block_N], accum_dtype)
acc_s_cast = T.alloc_fragment([block_M, block_N], dtype)
acc_o = T.alloc_fragment([block_M, dim], accum_dtype)
scores_max = T.alloc_fragment([block_M], accum_dtype)
scores_max_prev = T.alloc_fragment([block_M], accum_dtype)
scores_scale = T.alloc_fragment([block_M], accum_dtype)
scores_sum = T.alloc_fragment([block_M], accum_dtype)
logsum = T.alloc_fragment([block_M], accum_dtype)
mid = bx
hid = by % heads
bid = by // heads
sid = bz
T.annotate_layout({Q_shared: tilelang.layout.make_swizzled_layout(Q_shared)})
T.copy(Q[bid, mid * block_M:(mid + 1) * block_M, hid, :], Q_shared)
T.fill(acc_o, 0)
T.fill(logsum, 0)
T.fill(scores_max, -T.infinity(accum_dtype))
# TODO: Handle causal split case
loop_range = (
T.min(T.ceildiv(seqlen_kv, block_N), T.ceildiv(
(mid + 1) * block_M, block_N)) if is_causal else T.ceildiv(
(seqlen_kv // num_split), block_N))
for k in T.Pipelined(loop_range, num_stages=2):
MMA0(K, Q_shared, K_shared, acc_s, k, mid, hid, bid, sid)
Softmax(acc_s, acc_s_cast, scores_max, scores_max_prev, scores_scale, scores_sum,
logsum)
Rescale(acc_o, scores_scale)
MMA1(V, V_shared, acc_s_cast, acc_o, k, hid, bid, sid)
for i, j in T.Parallel(block_M, dim):
acc_o[i, j] /= logsum[i]
for i in T.Parallel(block_M):
logsum[i] = T.log2(logsum[i]) + scores_max[i] * scale
T.copy(logsum, glse[bid, hid, sid, mid * block_M:(mid + 1) * block_M])
T.copy(acc_o, O_shared)
T.copy(O_shared, Output_partial[bid, mid * block_M:(mid + 1) * block_M, hid, sid, :])
@T.macro
def combine(
glse: T.Tensor([batch, heads, num_split, seqlen_q], dtype),
Output_partial: T.Tensor(part_shape, dtype),
Output: T.Tensor(shape_q, dtype),
):
with T.Kernel(T.ceildiv(seqlen_q, block_M), heads, batch, threads=128) as (bx, by, bz):
po_local = T.alloc_fragment([block_M, dim], dtype)
po_shared = T.alloc_shared([block_M, dim], dtype)
o_accum_local = T.alloc_fragment([block_M, dim], accum_dtype)
o_shared = T.alloc_shared([block_M, dim], dtype)
lse_local = T.alloc_fragment([num_split, block_M], dtype)
lse_local_split = T.alloc_fragment([block_M], accum_dtype)
lse_logsum_local = T.alloc_fragment([block_M], accum_dtype)
lse_max_local = T.alloc_fragment([block_M], accum_dtype)
scale_local = T.alloc_fragment([block_M], accum_dtype)
T.annotate_layout({
o_accum_local: T.Fragment(o_accum_local.shape, forward_thread_fn=lambda i, j: i),
lse_local_split: T.Fragment(lse_local_split.shape, forward_thread_fn=lambda i: i),
o_shared: tilelang.layout.make_swizzled_layout(o_shared),
po_shared: tilelang.layout.make_swizzled_layout(po_shared),
})
T.clear(lse_logsum_local)
T.clear(o_accum_local)
T.copy(glse[
bz,
by,
:,
bx * block_M:(bx + 1) * block_M,
], lse_local)
T.reduce_max(lse_local, lse_max_local, dim=0, clear=False)
for k in T.Pipelined(num_split):
T.copy(lse_local[k, :], lse_local_split)
for i in T.Parallel(block_M):
lse_logsum_local[i] += T.exp2(lse_local_split[i] - lse_max_local[i])
for i in T.Parallel(block_M):
lse_logsum_local[i] = T.log2(lse_logsum_local[i]) + lse_max_local[i]
for k in T.Pipelined(num_split, num_stages=2):
T.copy(Output_partial[bz, bx * block_M:(bx + 1) * block_M, by, k, :], po_shared)
T.copy(po_shared, po_local)
T.copy(lse_local[k, :], lse_local_split)
for i in T.Parallel(block_M):
scale_local[i] = T.exp2(lse_local_split[i] - lse_logsum_local[i])
for i, j in T.Parallel(block_M, dim):
o_accum_local[i, j] += po_local[i, j] * scale_local[i]
T.copy(o_accum_local, o_shared)
T.copy(o_shared, Output[bz, bx * block_M:(bx + 1) * block_M, by, :])
@T.prim_func
def main(
Q: T.Tensor(shape_q, dtype),
K: T.Tensor(shape_kv, dtype),
V: T.Tensor(shape_kv, dtype),
glse: T.Tensor([batch, heads, num_split, seqlen_q], dtype),
Output_partial: T.Tensor(part_shape, dtype), # [batch, seqlen_q, heads, num_split, dim]
Output: T.Tensor(shape_q, dtype),
):
flash_attn_split(Q, K, V, glse, Output_partial)
combine(glse, Output_partial, Output)
return main
def ref_program(Q, K, V, glse, Output_partial, causal):
assert causal is False
dim = Q.size(-1)
scores = torch.einsum('bqhd,bkhd->bhqk', Q, K)
scores = scores / torch.sqrt(torch.tensor(dim, dtype=scores.dtype))
attention_weights = F.softmax(scores, dim=-1)
output = torch.einsum('bhqk,bkhd->bqhd', attention_weights, V)
return output
def reduce_ref(Q, K, V, glse, Output_partial, causal):
o = torch.empty_like(Output_partial[:, :, :, 0, :]).fill_(0)
lse_logsum = torch.empty_like(glse[:, :, 0, :]).fill_(0) # [batch, seqlen_q, heads]
lse_max = glse.max(dim=2, keepdim=False).values
for ks in range(num_split):
lse = glse[:, :, ks, :]
lse_logsum += torch.exp2(lse - lse_max)
lse_logsum = torch.log2(lse_logsum) + lse_max
for ks in range(num_split):
lse = glse[:, :, ks, :]
scale = torch.exp2(lse - lse_logsum) # [batch, heads, seqlen_q]
o += Output_partial[:, :, :, ks, :] * scale[:, :, :, None].transpose(1, 2)
return o.to(torch.float16)
def flash_split_ref(Q, K, V, causal):
# [batch, seqlen_q, heads, dim]
batch = Q.size(0)
block_M = Q.size(1)
nheads = Q.size(2)
dim = Q.size(3)
block_N = 128
seqlen_kv = K.size(1)
scale = (1.0 / dim)**0.5 * 1.44269504 # log2(e)
acc_s = torch.empty((batch, nheads, block_M, block_N), device="cuda", dtype=torch.float)
acc_s_cast = torch.empty((batch, nheads, block_M, block_N), device="cuda", dtype=torch.float16)
acc_o = torch.empty((batch, block_M, nheads, dim), device="cuda", dtype=torch.float)
scores_max = torch.empty((batch, nheads, block_M), device="cuda", dtype=torch.float)
scores_max_prev = torch.empty((batch, nheads, block_M), device="cuda", dtype=torch.float)
scores_scale = torch.empty((batch, nheads, block_M), device="cuda", dtype=torch.float)
scores_sum = torch.empty((batch, nheads, block_M), device="cuda", dtype=torch.float)
logsum = torch.empty((batch, nheads, block_M), device="cuda", dtype=torch.float)
gacc_o = torch.empty((num_split, batch, block_M, nheads, dim), device="cuda", dtype=torch.float)
glogsum = torch.empty((num_split, batch, nheads, block_M), device="cuda", dtype=torch.float)
Q_ = Q * scale
for ks in range(num_split):
acc_o.fill_(0)
logsum.fill_(0)
scores_max.fill_(float('-inf'))
scores_max_prev.fill_(float('-inf'))
for i in range(int((seqlen_kv // num_split) / block_N)):
acc_s.fill_(0)
acc_s = torch.einsum('bqhd,bkhd->bhqk', Q_,
K[:, (seqlen_kv // num_split) * ks +
i * block_N:(seqlen_kv // num_split) * ks +
(i + 1) * block_N, :, :]) # [batch, seqlen, nheads, block_N]
scores_max_prev = scores_max
scores_max = acc_s.max(dim=-1, keepdim=False).values # [blockM]
scores_scale = torch.exp2(scores_max_prev - scores_max)
acc_o *= scores_scale[:, :, :, None].transpose(1, 2)
acc_s = torch.exp2(acc_s - scores_max[:, :, :, None])
acc_s_cast = acc_s.to(torch.float16)
acc_o += torch.einsum(
'bhqk,bkhd->bqhd', acc_s_cast,
V[:, (seqlen_kv // num_split) * ks + i * block_N:(seqlen_kv // num_split) * ks +
(i + 1) * block_N, :, :])
scores_sum = acc_s.sum(dim=-1, keepdim=False)
logsum = logsum * scores_scale + scores_sum
acc_o /= logsum[:, :, :, None].transpose(1, 2)
logsum = torch.log2(logsum) + scores_max
gacc_o[ks, :, :, :, :] = acc_o
glogsum[ks, :, :, :] = logsum
return glogsum.to(torch.float16).permute(1, 2, 0,
3), gacc_o.to(torch.float16).permute(1, 2, 3, 0, 4)
def main():
BATCH, H, Q_CTX, KV_CTX, D_HEAD = 1, 32, 128, 8192, 128
causal = False
flops_per_matmul = 2.0 * BATCH * H * Q_CTX * KV_CTX * D_HEAD
total_flops = 2 * flops_per_matmul
if causal:
total_flops *= 0.5
BLOCK_M = 128
BLOCK_N = 64 # if D_HEAD <= 128 else 32
kernel = flashattn(BATCH, H, Q_CTX, KV_CTX, D_HEAD, causal, BLOCK_M, BLOCK_N)
ref_program_processed = partial(ref_program, causal=causal)
profiler = kernel.get_profiler(tilelang.TensorSupplyType.Normal)
profiler.assert_allclose(ref_program_processed, rtol=0.01, atol=0.01)
print("All checks passed!")
latency = profiler.do_bench(ref_program_processed, warmup=500)
print("{:.2f} ms".format(latency))
print("{:.2f} TFlops".format(total_flops / latency * 1e-9))
latency = profiler.do_bench(n_warmup=10, n_repeat=10)
print("{:.4f} ms".format(latency))
print("{:.2f} TFlops".format(total_flops / latency * 1e-9))
if __name__ == "__main__":
main()
examples/flash_attention/test_example_flash_attention.py
View file @ e9a608e2
......@@ -9,7 +9,6 @@ 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
......@@ -64,10 +63,5 @@ 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()
examples/flash_decoding/example_gqa_decode.py
View file @ e9a608e2
......@@ -46,12 +46,12 @@ def get_heuristic_config() -> Tuple[Dict, int]:
return cfg, sm_version
# TODO(lei): fix warp specialized and tma lower pass
def get_pass_configs():
_, sm_version = get_heuristic_config()
if sm_version == 80:
return {tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True}
else:
return {}
return {
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True
}
@autotune(configs=get_configs(), warmup=10, rep=10)
......@@ -465,13 +465,12 @@ def main(batch: int = 1,
o_ref = ref_program(q, k, v, mask, glse, Output_partial)
o_ref_split = ref_split_program(q, k, v, mask, glse, Output_partial)
assert_similar(o, o_ref)
assert_similar(o_ref_split, o_ref)
torch.testing.assert_close(o, o_ref, rtol=0.01, atol=0.01)
torch.testing.assert_close(o_ref_split, o_ref, rtol=0.01, atol=0.01)
print(o)
print(o_ref)
assert_similar(o, o_ref, name="o_ref")
assert_similar(o_ref_split, o_ref, name="o_ref_split")
profiler.assert_allclose(ref_program, rtol=0.01, atol=0.01)
profiler.assert_allclose(ref_split_program, rtol=0.01, atol=0.01)
print("All checks pass.")
latency = profiler.do_bench(ref_program, warmup=500)
print("Ref: {:.2f} ms".format(latency))
......
examples/flash_decoding/example_mha_inference.py
View file @ e9a608e2
......@@ -305,7 +305,6 @@ def main():
BLOCK_N = 64 # if D_HEAD <= 128 else 32
kernel = flashattn(BATCH, H, Q_CTX, KV_CTX, D_HEAD, causal, BLOCK_M, BLOCK_N)
ref_fn = partial(ref_program, causal=causal)
print(kernel.get_kernel_source())
profiler = kernel.get_profiler(tensor_supply_type=tilelang.TensorSupplyType.Normal)
profiler.assert_allclose(ref_fn, rtol=0.01, atol=0.01)
print("All checks passed!")
......
examples/flash_decoding/test_example_flash_decoding.py
View file @ e9a608e2
......@@ -4,6 +4,9 @@ import example_gqa_decode
import example_mha_inference
# TODO(lei): fix the correctness of gqa decode on sm90
@tilelang.testing.requires_cuda
@tilelang.testing.requires_cuda_compute_version_le(8, 9)
def test_example_example_gqa_decode():
example_gqa_decode.main()
......
examples/gemm_streamk/test_example_tilelang_gemm_splitk.py
View file @ e9a608e2
......@@ -3,6 +3,9 @@ import tilelang.testing
from example_tilelang_gemm_streamk import main
# not fully supported on sm90
@tilelang.testing.requires_cuda
@tilelang.testing.requires_cuda_compute_version_le(8, 9)
def test_example_tilelang_gemm_streamk():
main()
......
examples/gemv/example_gemv.py
View file @ e9a608e2
......@@ -291,9 +291,9 @@ def check_correctness_and_bench(kernel, N, K, bench_ref=True):
profiler = kernel.get_profiler()
profiler.assert_allclose(lambda x, y: x @ y.T, atol=1e-2, rtol=1e-2)
if bench_ref:
latency = profiler.do_bench(lambda x, y: x @ y.T, warmup=500)
latency = profiler.do_bench(lambda x, y: x @ y.T, warmup=50)
print(f"Torch Latency: {latency} ms")
latency = profiler.do_bench(kernel, warmup=500)
latency = profiler.do_bench(kernel, warmup=50)
print(f"TileLang Latency: {latency} ms\n")
......
examples/norm/test_rms_norm.py
View file @ e9a608e2
import torch
import tilelang
import tilelang.testing
import tilelang.language as T
......@@ -73,3 +74,7 @@ def test_rms_norm():
pass_configs={"tl.disable_tma_lower": True})
profiler = kernel.get_profiler()
profiler.assert_allclose(ref_program, rtol=0.01, atol=0.01)
if __name__ == "__main__":
tilelang.testing.main()
examples/pytest.ini 0 → 100644
View file @ e9a608e2
[pytest]
norecursedirs = bitnet-1.58b
examples/warp_specialize/example_warp_specialize_flashmla.py
View file @ e9a608e2
......@@ -49,8 +49,12 @@ def flashattn(batch, heads, kv_head_num, seqlen_kv, dim, pe_dim, block_N, block_
scores_max_0 = T.alloc_fragment([block_H], accum_dtype)
scores_max_1 = T.alloc_fragment([block_H], accum_dtype)
scores_max = T.alloc_shared([block_H], accum_dtype)
scores_max_prev_0 = T.alloc_fragment([block_H], accum_dtype)
scores_max_prev_1 = T.alloc_fragment([block_H], accum_dtype)
# TODO(lei): this is a workaround for the bug of replicate if stmt.
# have to be optimized in future with index aware sync thread pass injection.
# scores_max_prev_0 and scores_max_prev_1 should be allocated in fragment.
scores_max_prev_0 = T.alloc_shared([block_H], accum_dtype)
scores_max_prev_1 = T.alloc_shared([block_H], accum_dtype)
scores_scale_0 = T.alloc_shared([block_H], accum_dtype)
scores_scale_1 = T.alloc_shared([block_H], accum_dtype)
scores_sum_0 = T.alloc_fragment([block_H], accum_dtype)
......@@ -391,16 +395,7 @@ def ref_program(q, q_pe, kv, k_pe, glse, Output_partial):
return out
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--batch', type=int, default=132, help='batch size')
parser.add_argument('--heads', type=int, default=128, help='q heads number')
parser.add_argument('--kv_heads', type=int, default=1, help='kv heads number')
parser.add_argument('--kv_ctx', type=int, default=8192, help='kv context length')
parser.add_argument('--dim', type=int, default=512, help='head dim')
parser.add_argument('--pe_dim', type=int, default=64, help='pe head dim')
args = parser.parse_args()
batch, heads, kv_heads, kv_ctx, dim, pe_dim = args.batch, args.heads, args.kv_heads, args.kv_ctx, args.dim, args.pe_dim
def main(batch=132, heads=128, kv_heads=1, kv_ctx=8192, dim=512, pe_dim=64):
qk_flops = 2 * batch * heads * kv_ctx * (dim + pe_dim)
pv_flops = 2 * batch * heads * kv_ctx * dim
total_flops = qk_flops + pv_flops
......@@ -418,4 +413,13 @@ def main():
if __name__ == "__main__":
main()
parser = argparse.ArgumentParser()
parser.add_argument('--batch', type=int, default=132, help='batch size')
parser.add_argument('--heads', type=int, default=128, help='q heads number')
parser.add_argument('--kv_heads', type=int, default=1, help='kv heads number')
parser.add_argument('--kv_ctx', type=int, default=8192, help='kv context length')
parser.add_argument('--dim', type=int, default=512, help='head dim')
parser.add_argument('--pe_dim', type=int, default=64, help='pe head dim')
args = parser.parse_args()
batch, heads, kv_heads, kv_ctx, dim, pe_dim = args.batch, args.heads, args.kv_heads, args.kv_ctx, args.dim, args.pe_dim
main(batch, heads, kv_heads, kv_ctx, dim, pe_dim)
requirements-test.txt
View file @ e9a608e2
......@@ -27,6 +27,6 @@ setuptools
einops
attrs
decorator
flash-attn<=2.8.0
flash-attn<=2.2.0
scipy
tornado
\ No newline at end of file
src/op/bulk_copy.cc
View file @ e9a608e2
......@@ -94,7 +94,6 @@ Stmt Copy::LowerBulkCopy(const LowerArgs &T, arith::Analyzer *analyzer) const {
bool is_load;
if (src.scope() == "global" &&
(dst.scope() == "shared.dyn" || dst.scope() == "shared")) {
// Use the Hopper TMA bulk copy instructions
is_load = true;
} else if (dst.scope() == "global" &&
(src.scope() == "shared.dyn" || src.scope() == "shared")) {
......@@ -106,7 +105,6 @@ Stmt Copy::LowerBulkCopy(const LowerArgs &T, arith::Analyzer *analyzer) const {
Buffer shared_tensor = is_load ? dst : src;
Array<Range> global_range = is_load ? src_range : dst_range;
Array<Range> shared_range = is_load ? dst_range : src_range;
if (T.layout_map.count(global_tensor)) {
LOG(WARNING) << "TMA bulk copy cannot support a non-swizzled global "
"layout, fallback to normal copy.";
......@@ -116,7 +114,6 @@ Stmt Copy::LowerBulkCopy(const LowerArgs &T, arith::Analyzer *analyzer) const {
Array<PrimExpr> indices;
for (auto r : shared_range)
indices.push_back(r->min);
std::vector<PrimExpr> strides;
PrimExpr stride = 1;
for (size_t i = 0; i < shared_tensor->shape.size(); i++) {
......@@ -132,7 +129,6 @@ Stmt Copy::LowerBulkCopy(const LowerArgs &T, arith::Analyzer *analyzer) const {
for (size_t i = 0; i < indices.size(); i++) {
offset += indices[i] * strides[i];
}
Layout shared_layout;
if (T.layout_map.count(shared_tensor)) {
shared_layout = T.layout_map[shared_tensor];
......@@ -140,7 +136,6 @@ Stmt Copy::LowerBulkCopy(const LowerArgs &T, arith::Analyzer *analyzer) const {
}
TMADesc desc;
// Verify copy rank
desc.rank = global_tensor->shape.size();
ICHECK(desc.rank >= 1 && desc.rank <= 5) << desc.rank;
......@@ -175,6 +170,18 @@ Stmt Copy::LowerBulkCopy(const LowerArgs &T, arith::Analyzer *analyzer) const {
desc.global_stride = desc.global_stride.Map([&](PrimExpr e) {
return cast(DataType::Int(64), e) * global_tensor->dtype.bytes();
});
for (size_t i{1}; i < desc.global_stride.size(); i++) {
auto stride = desc.global_stride[i].as<IntImmNode>();
if (stride != nullptr) {
// otherwise, the stride is symbolic, we need to check in future with
// assumptions
if (stride->value % 16 != 0 || stride->value >= (1ULL << 40)) {
LOG(WARNING) << "TMA bulk copy cannot support a global stride of "
<< desc.global_stride[i] << ", fallback to normal copy.";
return Stmt();
}
}
}
// Smem Box
// check smem range and global range is legal
......@@ -184,19 +191,30 @@ Stmt Copy::LowerBulkCopy(const LowerArgs &T, arith::Analyzer *analyzer) const {
if (is_one(g_range->extent)) {
continue;
}
auto s_range = shared_range[s_range_idx++];
// skip one range if it is 1
// in case of global range is [128, 64], while shared range is [1, 128, 64]
// A_shared[0, :, :].
while (is_one(shared_range[s_range_idx]->extent) &&
s_range_idx < shared_range.size()) {
s_range_idx++;
}
if (s_range_idx >= shared_range.size()) {
LOG(FATAL) << "TMA bulk copy cannot support a global range of "
<< global_range << ", shared_range " << shared_range;
}
auto s_range = shared_range[s_range_idx];
s_range_idx++;
ICHECK(StructuralEqual()(g_range->extent, s_range->extent))
<< global_tensor->name << "[" << i << "] is illegal, "
<< global_tensor->name << "[" << i << "] = " << g_range->extent << ", "
<< shared_tensor->name << "[" << s_range_idx
<< "] = " << s_range->extent;
}
desc.smem_box =
ReverseArray(global_range.Map([](Range r) { return r->extent; }));
desc.smem_stride = Array<PrimExpr>(desc.rank, PrimExpr(1));
// L2 & OOB
desc.l2_promotion = static_cast<int>(CU_TENSOR_MAP_L2_PROMOTION_L2_128B);
desc.oob_fill = static_cast<int>(CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
......@@ -230,7 +248,7 @@ Stmt Copy::LowerBulkCopy(const LowerArgs &T, arith::Analyzer *analyzer) const {
shared_tensor->dtype.bits()))) {
desc.swizzle = static_cast<int>(CU_TENSOR_MAP_SWIZZLE_128B);
} else {
ICHECK(0) << "Cannot detect TMA layout.";
return Stmt();
}
}
......@@ -252,6 +270,21 @@ Stmt Copy::LowerBulkCopy(const LowerArgs &T, arith::Analyzer *analyzer) const {
ICHECK((*inner_box_dim) % instruction_dim == 0);
desc.smem_box.Set(0, PrimExpr(instruction_dim));
int inner_box_dim_ = instruction_dim * shared_tensor->dtype.bytes();
if (desc.swizzle == static_cast<int>(CU_TENSOR_MAP_SWIZZLE_NONE) &&
inner_box_dim_ % 256 != 0)
return Stmt();
#define CHECK_INNER_BOX_DIM(N) \
if (desc.swizzle == static_cast<int>(CU_TENSOR_MAP_SWIZZLE_##N##B) && \
inner_box_dim_ > N) \
return Stmt();
CHECK_INNER_BOX_DIM(32);
CHECK_INNER_BOX_DIM(64);
CHECK_INNER_BOX_DIM(128);
#undef CHECK_INNER_BOX_DIM
Call create_descriptor =
Call(DataType::Handle(), create_tma_descriptor(), desc.EncodeCallArgs());
......
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