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

bugfix

Browse Source
This commit is contained in:
geohot
2026-08-04 19:41:01 +00:00
parent 330543c9af
commit 80e8ffc7c2
2 changed files with 21 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/unit/test_llm_amd.py
+19 -1
View File
@@ -1,6 +1,6 @@
import unittest
import numpy as np
from tinygrad import Tensor, dtypes, nn
from tinygrad import Tensor, UOp, dtypes, nn
from tinygrad.llm.kernels import Linear
from tinygrad.llm.kernels.amd import q8_quantize, quantized_attention
from tinygrad.llm.gguf import ggml_data_to_tensor
@@ -34,4 +34,22 @@ class TestQ8Quantize(unittest.TestCase):
out = quantized_attention(q, Tensor.stack(k, v), cache, scale, 0).realize()
np.testing.assert_allclose(out.numpy(), v.expand(1, 2, 1, 32).numpy(), rtol=2e-2, atol=2e-2)
def test_prefill_attention_unaligned_start(self):
if not str(Tensor.empty(1).device).startswith("AMD"): self.skipTest("AMD required")
rng = np.random.default_rng(42)
start_pos = 31
q = Tensor.zeros(1, 8, 32, 128)
old_kv = rng.normal(size=(2, 1, 1, start_pos, 128)).astype(np.float32)
new_kv = rng.normal(size=(2, 1, 1, 32, 128)).astype(np.float32)
cache = Tensor.empty(2, 1, 1, 256, 128, dtype=dtypes.int8).contiguous()
scale = Tensor.empty(2, 1, 1, 256, dtype=dtypes.float16).contiguous()
old_scale = np.maximum(np.max(np.abs(old_kv), axis=-1, keepdims=True) / 127, 1e-8).astype(np.float16)
Tensor.realize(cache[:, :, :, :start_pos].assign(Tensor(np.rint(old_kv / old_scale).astype(np.int8))),
scale[:, :, :, :start_pos].assign(Tensor(old_scale.squeeze(-1))))
out = quantized_attention(q, Tensor(new_kv), cache, scale, UOp.variable("start_pos", 0, 255).bind(start_pos)).realize()
values = cache[1, 0, 0, :start_pos+32].numpy().astype(np.float32) * \
scale[1, 0, 0, :start_pos+32].numpy().astype(np.float32)[:, None]
expected = np.stack([values[:start_pos+i+1].mean(0) for i in range(32)])[None, None].repeat(8, axis=1)
np.testing.assert_allclose(out.numpy(), expected, rtol=2e-3, atol=2e-3)
if __name__ == "__main__": unittest.main()
tinygrad/llm/kernels/amd.py
+2 -1
View File
@@ -134,7 +134,8 @@ def _amd_flash_attention(o:UOp, q:UOp, cache:UOp, kv_scale:UOp, valid_kv_len:int
m_ij = m_ij.after(m_ij.store(m_ij.after(rm2).maximum(S_reg[:, rm2])).end(rm2))
ri_w = UOp.range(TM, 270)
m_ij = m_ij.after(m_ij[ri_w].store(warp_reduce(m_ij[ri_w], lane, maximum=True)).end(ri_w))
S_reg = S_reg.after(S_reg.store(((S_reg - m_ij.reshape(TM, 1).expand(TM, TN)) * LOG2E).exp2()))
tile_max = m_ij.reshape(TM, 1).expand(TM, TN).maximum(-1e30)
S_reg = S_reg.after(S_reg.store(((S_reg - tile_max) * LOG2E).exp2()))
p_local, ri_ws = _reg((TM,), 8, 0, n_tile), UOp.range(TM, 295, AxisType.WEAK)
p_sum = p_local.after(p_local[ri_ws].store(sum((warp_reduce(S_reg[ri_ws, rn], lane) for rn in range(TN)), S_reg.const_like(0))).end(ri_ws))
P_lds = QP_lds.flatten()[:WAVES_N * BLOCK_M * BLOCK_N].reshape(WAVES_N, BLOCK_M, BLOCK_N)