mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-08 05:23:41 +08:00
Daniel Koepping
47 lines
1.9 KiB
Python
47 lines
1.9 KiB
Python
import unittest
|
|
from tinygrad import Tensor, dtypes, TinyJit, UOp
|
|
from tinygrad.llm.kernels.amd import Linear
|
|
from tinygrad.llm.gguf import ggml_data_to_tensor
|
|
from tinygrad.llm.model import apply_rope as apply_rope_new, precompute_freqs_cis
|
|
from test.helpers import assert_jit_cache_len, check_schedule
|
|
|
|
def apply_rope(x:Tensor, start_pos:int):
|
|
B, H, T, Hd = x.shape
|
|
precompute_freqs_cis.cache_clear()
|
|
freqs_cis = precompute_freqs_cis(Hd, start_pos+T)[start_pos:start_pos+T]
|
|
return apply_rope_new(x, freqs_cis)
|
|
|
|
class TestAttention(unittest.TestCase):
|
|
def test_half_qkv_buffers(self):
|
|
BS, seqlen, dim = 10, 4, 100
|
|
q = Tensor.ones(BS, seqlen, dim, dtype=dtypes.half).contiguous().realize()
|
|
k = Tensor.ones(BS, seqlen, dim, dtype=dtypes.half).contiguous().realize()
|
|
v = Tensor.ones(BS, seqlen, dim, dtype=dtypes.half).contiguous().realize()
|
|
attn = q.scaled_dot_product_attention(k, v)
|
|
# attention has 4 kernels now
|
|
check_schedule(attn, 4)
|
|
|
|
def test_apply_rope_jit_prune(self):
|
|
def rope_fn(x_in, pos): return apply_rope(x_in, pos)
|
|
rope_noprune = TinyJit(rope_fn)
|
|
rope_prune = TinyJit(rope_fn, prune=True)
|
|
|
|
v_pos = UOp.variable("start_pos", 0, 100)
|
|
for _ in range(3):
|
|
rope_noprune(Tensor.randn(1, 2, 4, 8, dtype=dtypes.float32), v_pos.bind(1))
|
|
rope_prune(Tensor.randn(1, 2, 4, 8, dtype=dtypes.float32), v_pos.bind(1))
|
|
assert_jit_cache_len(rope_prune, 1)
|
|
assert_jit_cache_len(rope_noprune, 3)
|
|
|
|
class TestLinear(unittest.TestCase):
|
|
def test_recovers_packed_ggml_weight(self):
|
|
for ggml_type,packed_size in ((13, 176), (14, 210), (23, 136)):
|
|
packed = Tensor.empty(packed_size+4, dtype=dtypes.uint8)[4:]
|
|
decoded = ggml_data_to_tensor(packed, 256, ggml_type).reshape(1, 256)
|
|
linear = Linear(256, 1, bias=False)
|
|
linear.set_quantized(decoded)
|
|
self.assertEqual((linear.ggml_type, linear.weight.nbytes()), (ggml_type, packed_size))
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|