gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

.

Browse Source
This commit is contained in:
Philippe Tillet
2022-12-21 14:02:10 -08:00
parent 88e572e54d
commit 033e82060d
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/tutorials/06-fused-attention.py
View File

@@ -222,8 +222,8 @@ class _attention(torch.autograd.Function):
o.stride(0), o.stride(1), o.stride(2), o.stride(3), o.stride(0), o.stride(1), o.stride(2), o.stride(3),
q.shape[0], q.shape[1], q.shape[2], q.shape[0], q.shape[1], q.shape[2],
BLOCK_M=BLOCK, BLOCK_N=BLOCK, BLOCK_M=BLOCK, BLOCK_N=BLOCK,
BLOCK_DMODEL=Lk, num_warps=num_warps, BLOCK_DMODEL=Lk, num_warps=4,
num_stages=1, num_stages=2,
) )
ctx.save_for_backward(q, k, v, o, L, m) ctx.save_for_backward(q, k, v, o, L, m)