From ff3f2a9c1a497d90df96afdccd7f48d61b77dfa9 Mon Sep 17 00:00:00 2001 From: chenyu Date: Mon, 25 Nov 2024 18:59:51 -0500 Subject: [PATCH] Revert "move attention upcast (#7830)" (#7903) This reverts commit c07daf40e79da6ca76c0ed9b888bb2e488d197e9. --- test/unit/test_attention.py | 17 ----------------- tinygrad/tensor.py | 2 +- 2 files changed, 1 insertion(+), 18 deletions(-) delete mode 100644 test/unit/test_attention.py diff --git a/test/unit/test_attention.py b/test/unit/test_attention.py deleted file mode 100644 index 2c1060d6d8..0000000000 --- a/test/unit/test_attention.py +++ /dev/null @@ -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) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index e204c25cb4..d84e6e52c0 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -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: