Skip to content

Commit

Permalink
A working version but load Q at every loop
Browse files Browse the repository at this point in the history
  • Loading branch information
@zhanglx13
zhanglx13 committed Nov 7, 2023
1 parent 56ad004 commit a79e6f2
Show file tree
Hide file tree
Showing 2 changed files with 51 additions and 27 deletions.
3 changes: 2 additions & 1 deletion lib/Dialect/TritonGPU/Transforms/AccelerateAMDMatmul.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ using ttg::SliceEncodingAttr;

SmallVector<unsigned, 2>
warpsPerTileMFMA(tt::DotOp dotOp, const ArrayRef<int64_t> shape, int numWarps) {
return {(unsigned)numWarps, 1};
// TODO: needs to be updated with appropriate shapePerWarp etc.
auto filter = [&dotOp](Operation *op) {
return op->getParentRegion() == dotOp->getParentRegion();
Expand Down Expand Up @@ -186,7 +187,7 @@ class BlockedToMFMA : public mlir::RewritePattern {

bool isTransposed = isChainDot(dotOp);
mfmaEnc = ttg::MfmaEncodingAttr::get(oldRetType.getContext(), nonKDim,
warpsPerTile, isTransposed, CTALayout);
warpsPerTile, true, CTALayout);

auto newRetType =
RankedTensorType::get(retShape, oldRetType.getElementType(), mfmaEnc);
Expand Down
75 changes: 49 additions & 26 deletions python/perf-kernels/06-fused-attention-sliceQKV.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,11 +25,11 @@ def max_fn(x, y):

@triton.autotune(
configs=[
triton.Config({'BLOCK_M': 256, 'BLOCK_N': 64, 'waves_per_eu': 2, 'pre_load_v': False}, num_stages=1, num_warps=8),
#triton.Config({'BLOCK_M': 256, 'BLOCK_N': 64, 'waves_per_eu': 2, 'pre_load_v': False}, num_stages=1, num_warps=8),
triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'waves_per_eu': 2, 'pre_load_v': False}, num_stages=1, num_warps=4),
triton.Config({'BLOCK_M': 256, 'BLOCK_N': 128, 'waves_per_eu': 2, 'pre_load_v': False}, num_stages=1, num_warps=8),
triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'waves_per_eu': 3, 'pre_load_v': True}, num_stages=1, num_warps=4),
triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'waves_per_eu': 3, 'pre_load_v': False}, num_stages=1, num_warps=4),
#triton.Config({'BLOCK_M': 256, 'BLOCK_N': 128, 'waves_per_eu': 2, 'pre_load_v': False}, num_stages=1, num_warps=8),
#triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'waves_per_eu': 3, 'pre_load_v': True}, num_stages=1, num_warps=4),
#triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'waves_per_eu': 3, 'pre_load_v': False}, num_stages=1, num_warps=4),
],
key=['Z', 'H', 'N_CTX', 'STAGE', 'BLOCK_DMODEL'],
)
Expand Down Expand Up @@ -58,15 +58,15 @@ def _attn_fwd(
shape=(N_CTX, BLOCK_DMODEL),
strides=(stride_qm, stride_qk),
offsets=(start_m * BLOCK_M, 0),
block_shape=(BLOCK_M, BLOCK_DMODEL),
block_shape=(BLOCK_M, 32),
order=(1, 0)
)
K_block_ptr = tl.make_block_ptr(
base=K + qkv_offset,
shape=(BLOCK_DMODEL, N_CTX),
strides=(stride_kk, stride_kn),
offsets=(0, 0),
block_shape=(BLOCK_DMODEL, BLOCK_N),
block_shape=(32, BLOCK_N),
order=(0, 1)
)
V_block_ptr = tl.make_block_ptr(
Expand All @@ -88,21 +88,27 @@ def _attn_fwd(
# 2^x instead of exp in the loop because CSE and LICM
# don't work as expected with `exp` in the loop
qk_scale = sm_scale * 1.44269504
q = tl.load(Q_block_ptr)
q = (q * qk_scale).to(tl.float16)
#q = tl.load(Q_block_ptr)
#q = (q * qk_scale).to(tl.float16)
lo, hi = 0, N_CTX
# loop over k, v and update accumulator
for start_n in range(lo, hi, BLOCK_N):
start_n = tl.multiple_of(start_n, BLOCK_N)
# -- compute qk ----
k = tl.load(K_block_ptr)
if pre_load_v:
v = tl.load(V_block_ptr)
#k = tl.load(K_block_ptr)
qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
#if STAGE == 2:
# mask = offs_m[:, None] >= (start_n + offs_n[None, :])
# qk = tl.where(mask, qk, float("-inf"))
qk += tl.dot(q, k)
for i in range (0,4):
q = tl.load(Q_block_ptr)
q = (q * qk_scale).to(tl.float16)
k = tl.load(K_block_ptr)
qk += tl.dot(q, k)
Q_block_ptr = tl.advance(Q_block_ptr, (0, 32))
K_block_ptr = tl.advance(K_block_ptr, (32, 0))
if pre_load_v:
v = tl.load(V_block_ptr)
m_ij = tl.maximum(m_i, tl.max(qk, 1))
qk = qk - m_ij[:, None]
p = tl.math.exp2(qk)
Expand All @@ -118,7 +124,24 @@ def _attn_fwd(
# update m_i and l_i
m_i = m_ij
V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))
K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))
#K_block_ptr = tl.advance(K_block_ptr, (-128, BLOCK_N))
K_block_ptr = tl.make_block_ptr(
base=K + qkv_offset,
shape=(BLOCK_DMODEL, N_CTX),
strides=(stride_kk, stride_kn),
offsets=(0, start_n+BLOCK_N),
block_shape=(32, BLOCK_N),
order=(0, 1)
)
#Q_block_ptr = tl.advance(Q_block_ptr, (0, -128))
Q_block_ptr = tl.make_block_ptr(
base=Q + qkv_offset,
shape=(N_CTX, BLOCK_DMODEL),
strides=(stride_qm, stride_qk),
offsets=(start_m * BLOCK_M, 0),
block_shape=(BLOCK_M, 32),
order=(1, 0)
)
acc = acc / l_i[:, None]
# write back O
O_block_ptr = tl.make_block_ptr(
Expand Down Expand Up @@ -188,16 +211,16 @@ def forward(ctx, q, k, v, causal, sm_scale, split_kernel=False):


@pytest.mark.parametrize('Z, H, N_CTX, D_HEAD',
[(4, 48, 1024, 64),
(4, 48, 2048, 64),
(4, 48, 4096, 64),
[#(4, 48, 1024, 64),
#(4, 48, 2048, 64),
#(4, 48, 4096, 64),
(4, 48, 1024, 128),
(4, 48, 2048, 128),
(4, 48, 4096, 128),
#(4, 48, 8192, 64),
#(4, 48, 16384, 64)
])
@pytest.mark.parametrize('causal', [False, True])
@pytest.mark.parametrize('causal', [False])
def test_op_fwd(Z, H, N_CTX, D_HEAD, causal, dtype=torch.float16):
torch.manual_seed(20)
q = (
Expand Down Expand Up @@ -249,16 +272,16 @@ def test_op_fwd(Z, H, N_CTX, D_HEAD, causal, dtype=torch.float16):
for causal in [False]:
configs.append(triton.testing.Benchmark(
x_names=['BATCH', 'H','N_CTX'],
x_vals=[(16, 16, 1024),
(8, 16, 2048),
(4, 16, 4096),
(2, 16, 8192),
(1, 16, 16384),
(4, 48, 1024),
(4, 48, 2048),
x_vals=[#(16, 16, 1024),
#(8, 16, 2048),
#(4, 16, 4096),
#(2, 16, 8192),
#(1, 16, 16384),
#(4, 48, 1024),
#(4, 48, 2048),
(4, 48, 4096),
(4, 48, 8192),
(4, 48, 16384),
#(4, 48, 8192),
#(4, 48, 16384),
],
line_arg='provider',
line_vals=['triton'] + (['flash'] if HAS_FLASH else []),
Expand Down

0 comments on commit a79e6f2

Please sign in to comment.