tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 21:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

stub attention (#13196)

Browse Source 
* stub attention

* name the kernels
This commit is contained in:
George Hotz
2025-11-10 13:48:38 -08:00
committed by GitHub
parent 58c30fc7ce
commit 0c978d45e6
2 changed files with 15 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/models/llama.py
+14
View File
@@ -89,6 +89,20 @@ class Attention:
keys, values = repeat_kv(keys, self.n_rep), repeat_kv(values, self.n_rep)
xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
attn = xq.scaled_dot_product_attention(keys, values, mask).transpose(1, 2)
if getenv("STUB_ATTENTION"):
# TODO: do we need mask?
from tinygrad.uop.ops import UOp, KernelInfo
def fa_custom_forward(attn:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
return UOp.sink(arg=KernelInfo(name="fa_custom_forward"))
def fa_custom_backward(out_q:UOp, out_k:UOp, out_v:UOp, grad:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
return UOp.sink(arg=KernelInfo(name="fa_custom_backward"))
def fa_backward(grad:UOp, kernel:UOp) -> tuple[None, UOp, UOp, UOp]:
grad_q = Tensor.empty_like(q:=Tensor(kernel.src[1]))
grad_k = Tensor.empty_like(k:=Tensor(kernel.src[2]))
grad_v = Tensor.empty_like(v:=Tensor(kernel.src[3]))
ck = Tensor.custom_kernel(grad_q, grad_k, grad_v, Tensor(grad), q, k, v, fxn=fa_custom_backward)[:3]
return (None, ck[0].uop, ck[1].uop, ck[2].uop)
attn = Tensor.empty_like(attn).custom_kernel(xq, keys, values, fxn=fa_custom_forward, grad_fxn=fa_backward)[0]
attn = attn.reshape(bsz, seqlen, -1)
return self.wo(attn)
tinygrad/schedule/rangeify.py
+1
View File
@@ -538,6 +538,7 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
tsink = graph_rewrite(tsink, symbolic+pm_reduce_simplify+pm_const_buffer_folding, name="symbolic+reduce_collapse") # this does const folding
tsink = graph_rewrite(tsink, pm_remove_bufferize, bottom_up=True, name="remove bufferize with cost function")
tsink = graph_rewrite(tsink, symbolic+pm_reduce_simplify+pm_const_buffer_folding, name="symbolic+reduce_collapse pt 2")
tsink = graph_rewrite(tsink, pm_limit_bufs, ctx=rctx, name="limit buffers")
# rebuild the sink with all the BUFFERIZEs with tags, this is what's ending up in the tensor graph