sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions 8 Activity

Revert "move attention upcast (#7830)" (#7903)

Browse Source 
This reverts commit c07daf40e7.
This commit is contained in:
chenyu
2024-11-25 18:59:51 -05:00
committed by GitHub
parent 04bee97d2a
commit ff3f2a9c1a
2 changed files with 1 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/unit/test_attention.py
-17
View File
@@ -1,17 +0,0 @@
#!/usr/bin/env python
import unittest
from tinygrad import Tensor, dtypes
from tinygrad.engine.schedule import create_schedule
class TestAttention(unittest.TestCase):
def test_half_intermediate_dtypes(self):
q = Tensor.empty(1, 64, 128, dtype=dtypes.half).realize()
k = Tensor.empty(1, 64, 128, dtype=dtypes.half).realize()
v = Tensor.empty(1, 64, 128, dtype=dtypes.half).realize()
attn = q.scaled_dot_product_attention(k, v)
sched = create_schedule(attn.lazydata.lbs)
# TODO: make attention 1 kernel
self.assertEqual(len(sched), 5)
# store in half after after matmul
for buf in sched[0].outputs: self.assertEqual(buf.dtype, dtypes.half)
tinygrad/tensor.py
+1 -1
View File
@@ -3356,7 +3356,7 @@ class Tensor(SimpleMathTrait):
assert all_int(self.shape), f"does not support symbolic shape {self.shape}"
if is_causal: attn_mask = Tensor.ones(self.shape[-2], key.shape[-2], requires_grad=False, device=self.device).tril(0).cast(dtypes.bool)
if attn_mask is not None and attn_mask.dtype == dtypes.bool: attn_mask = (attn_mask == 0).where(-float("inf"), 0)
qk = (self.matmul(key.transpose(-2,-1)) / math.sqrt(self.shape[-1])).cast(least_upper_dtype(self.dtype, key.dtype, dtypes.float32))
qk = self.matmul(key.transpose(-2,-1), acc_dtype=least_upper_dtype(self.dtype, key.dtype, dtypes.float32)) / math.sqrt(self.shape[-1])
return ((qk+attn_mask) if attn_mask is not None else qk).softmax(-1).cast(self.dtype).dropout(dropout_p) @ value
def _do_reduction(self, reduction:ReductionStr="mean") -> Tensor: