geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions 7 Activity

fa: paas through window (#17517)

Browse Source
This commit is contained in:
wozeparrot
2026-08-13 14:45:57 +08:00
committed by GitHub
parent 39d144546e
commit 1b7f040984
1 changed files with 55 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/thunder/amd/fa.py
+55 -9
View File
@@ -110,7 +110,49 @@ def _sharded_empty_like(ref:Tensor, axis:int|None=None) -> Tensor:
return _sharded_empty(ref.shape, ref, axis)
@functools.cache
def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink):
def _windowed_lse(xq:Tensor, xk:Tensor, sinks, W:int) -> Tensor:
B, N, H, hd = xq.shape
H_KV = xk.shape[2]; R = H // H_KV; nb = N // W; sm = hd ** -0.5
q = xq.reshape(B, N, H_KV, R, hd).permute(0, 2, 3, 1, 4).reshape(B, H_KV, R, nb, W, hd).float()
k = xk.permute(0, 2, 1, 3).reshape(B, H_KV, 1, nb, W, hd).float()
k_prev = k.pad((None, None, None, (1, 0), None, None))[:, :, :, :nb]
sc_d = (q @ k.transpose(-1, -2)) * sm
sc_p = (q @ k_prev.transpose(-1, -2)) * sm
li, lj = Tensor.arange(W).reshape(W, 1), Tensor.arange(W).reshape(1, W)
pv = (Tensor.arange(nb).reshape(nb, 1, 1) >= 1)
sc_d = (lj <= li).where(sc_d, -float("inf"))
sc_p = ((li < lj) & pv).where(sc_p, -float("inf"))
m = sc_d.max(-1, keepdim=True).maximum(sc_p.max(-1, keepdim=True))
if sinks is not None: m = m.maximum(sinks.reshape(1, H_KV, R, 1, 1, 1).float())
denom = (sc_d - m).exp().sum(-1, keepdim=True) + (sc_p - m).exp().sum(-1, keepdim=True)
if sinks is not None: denom = denom + (sinks.reshape(1, H_KV, R, 1, 1, 1).float() - m).exp()
return (m + denom.log()).reshape(B, H, N).unsqueeze(2) # (B, H, 1, N), matches saved l_vec
def _windowed_delta(xq:Tensor, xk:Tensor, xv:Tensor, do:Tensor, sinks, W:int) -> Tensor:
B, N, H, hd = xq.shape
H_KV = xk.shape[2]; R = H // H_KV; nb = N // W; sm = hd ** -0.5
q = xq.reshape(B, N, H_KV, R, hd).permute(0, 2, 3, 1, 4).reshape(B, H_KV, R, nb, W, hd).float()
k = xk.permute(0, 2, 1, 3).reshape(B, H_KV, 1, nb, W, hd).float()
v = xv.permute(0, 2, 1, 3).reshape(B, H_KV, 1, nb, W, hd).float()
dob = do.reshape(B, N, H_KV, R, hd).permute(0, 2, 3, 1, 4).reshape(B, H_KV, R, nb, W, hd).float()
k_prev = k.pad((None, None, None, (1, 0), None, None))[:, :, :, :nb]
v_prev = v.pad((None, None, None, (1, 0), None, None))[:, :, :, :nb]
sc_d = (q @ k.transpose(-1, -2)) * sm
sc_p = (q @ k_prev.transpose(-1, -2)) * sm
li, lj = Tensor.arange(W).reshape(W, 1), Tensor.arange(W).reshape(1, W)
pv = (Tensor.arange(nb).reshape(nb, 1, 1) >= 1)
sc_d = (lj <= li).where(sc_d, -float("inf"))
sc_p = ((li < lj) & pv).where(sc_p, -float("inf"))
m = sc_d.max(-1, keepdim=True).maximum(sc_p.max(-1, keepdim=True))
if sinks is not None: m = m.maximum(sinks.reshape(1, H_KV, R, 1, 1, 1).float())
e_d, e_p = (sc_d - m).exp(), (sc_p - m).exp()
denom = e_d.sum(-1, keepdim=True) + e_p.sum(-1, keepdim=True)
if sinks is not None: denom = denom + (sinks.reshape(1, H_KV, R, 1, 1, 1).float() - m).exp()
o = ((e_d / denom) @ v) + ((e_p / denom) @ v_prev)
delta = (dob * o).sum(-1)
return delta.reshape(B, H, N).unsqueeze(2)
def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink, window=0):
def grad(dou:UOp, ker:UOp) -> tuple:
do = Tensor(dou, device=dou.device)
attn = Tensor(ker.src[1].after(ker), device=ker.src[1].device)
@@ -118,6 +160,8 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
xq = Tensor(ker.src[3], device=ker.src[3].device)
xk = Tensor(ker.src[4], device=ker.src[4].device)
xv = Tensor(ker.src[5], device=ker.src[5].device)
if window:
l_vec = _windowed_lse(xq, xk, Tensor(ker.src[6], device=ker.src[6].device) if has_sink else None, window)
dq = _sharded_empty((B, H, N, D), xq, axis=shard_axis_t)
GROUP_SIZE = H_local // H_KV_local
@@ -128,8 +172,10 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
# delta_vec = (do * attn).sum(-1, dtype=dtypes.float32).transpose(1, 2).unsqueeze(-2).detach()
delta_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)
delta_vec, dq = Tensor.custom_kernel(delta_vec, dq, attn, do, fxn=functools.partial(custom_fa_backward_pre, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:2]
if window:
delta_vec = _windowed_delta(xq, xk, xv, do, Tensor(ker.src[6], device=ker.src[6].device) if has_sink else None, window)
dq, dk_partial, dv_partial = Tensor.custom_kernel(dq, dk_partial, dv_partial, do, xq, xk, xv, l_vec, delta_vec, fxn=functools.partial(custom_fa_backward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:3]
dq, dk_partial, dv_partial = Tensor.custom_kernel(dq, dk_partial, dv_partial, do, xq, xk, xv, l_vec, delta_vec, fxn=functools.partial(custom_fa_backward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D, window=window))[:3]
if D == 64:
dq = dq.reshape(B, H, N//16, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2).permute(0, 1, 2, 8, 9, 10, 11, 3, 4, 6, 7, 5, 12).reshape(B, H, N, D).transpose(1, 2)
@@ -149,7 +195,7 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
return grad
# TODO: remove write_flat once scheduler can remove reshapes between custom_kernel. TestCustomKernel.test_simple_reshape
def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False, write_flat:bool=False, sinks:Tensor|None=None):
def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False, write_flat:bool=False, sinks:Tensor|None=None, window:int=0):
assert attn_mask is None, "attn_mask not supported"
assert is_causal, "only causal attention supported"
@@ -176,18 +222,18 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
attn = _sharded_empty((B, N, H * D), xq, axis=shard_axis) if write_flat else _sharded_empty_like(xq, axis=shard_axis)
l_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)
grad = _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink)
grad = _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink, window=window)
fwd_inputs = (attn, l_vec, xq, xk, xv) + ((sinks,) if has_sink else ())
attn, l_vec = Tensor.custom_kernel(*fwd_inputs, fxn=functools.partial(custom_fa_forward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D, has_sink=has_sink), grad_fxn=grad)[:2]
attn, l_vec = Tensor.custom_kernel(*fwd_inputs, fxn=functools.partial(custom_fa_forward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D, has_sink=has_sink, window=window), grad_fxn=grad)[:2]
return attn, attn, l_vec
@functools.cache
def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, sinks:UOp|None=None, *, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int, has_sink:bool=True):
def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, sinks:UOp|None=None, *, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int, has_sink:bool=True, window:int=0):
code = (pathlib.Path(__file__).parent / "fa_fwd_causal.cpp").read_text()
compile_args = [f"-I{(pathlib.Path(__file__).parent / 'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-DHIP_ENABLE_WARP_SYNC_BUILTINS", "-ffast-math",
f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DATTN_SINK={int(has_sink)}"]
f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DATTN_SINK={int(has_sink)}", f"-DWINDOW={window}"]
Q_BLOCK_SIZE = 32
NUM_WARPS = 8
@@ -247,10 +293,10 @@ def custom_fa_backward_pre(delta_vec:UOp, dq:UOp, o:UOp, do:UOp, device:str, arc
src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
@functools.cache
def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_vec:UOp, delta_vec:UOp, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int):
def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_vec:UOp, delta_vec:UOp, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int, window:int=0):
code = (pathlib.Path(__file__).parent / "fa_bwd_causal.cpp").read_text()
compile_args = [f"-I{(pathlib.Path(__file__).parent / 'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-DHIP_ENABLE_WARP_SYNC_BUILTINS", "-ffast-math",
f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}"]
f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DWINDOW={window}"]
BLOCK_SIZE_KV = 256
GROUP_SIZE = H // H_KV