56 lines
2.3 KiB
Python
56 lines
2.3 KiB
Python
import unittest
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
from prain import gemv
|
|
|
|
GPU = torch.cuda.is_available() and gemv.fast_device(0)
|
|
|
|
|
|
def rel(a, b):
|
|
return ((a.float() - b.float()).norm() / b.float().norm()).item()
|
|
|
|
|
|
@unittest.skipUnless(GPU, "needs a gfx10 GPU")
|
|
class TestGemv(unittest.TestCase):
|
|
def test_up_to_four_rows_match_a_float32_linear(self):
|
|
torch.manual_seed(0)
|
|
for n, k in ((1000, 2816), (2816, 2112), (130, 70)):
|
|
w = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") * 0.02
|
|
for rows in (1, 2, 3, 4):
|
|
x = torch.randn(rows, k, dtype=torch.bfloat16, device="cuda")
|
|
y = gemv.linear(x, w)
|
|
self.assertEqual((y.shape, y.dtype), ((rows, n), torch.bfloat16))
|
|
self.assertLess(rel(y, F.linear(x.float(), w.float())), 1e-2, (n, k, rows))
|
|
x = torch.randn(2, 1, 70, dtype=torch.float16, device="cuda") # leading dims, fp16
|
|
w = torch.randn(130, 70, dtype=torch.float16, device="cuda")
|
|
self.assertEqual(gemv.linear(x, w).shape, (2, 1, 130))
|
|
self.assertLess(rel(gemv.linear(x, w), F.linear(x.float(), w.float())), 1e-2)
|
|
|
|
def test_more_rows_bias_or_strided_weights_use_f_linear(self):
|
|
w = torch.randn(64, 96, dtype=torch.bfloat16, device="cuda")
|
|
bias = torch.randn(64, dtype=torch.bfloat16, device="cuda")
|
|
x = torch.randn(5, 96, dtype=torch.bfloat16, device="cuda")
|
|
self.assertTrue(torch.equal(gemv.linear(x, w), F.linear(x, w)))
|
|
self.assertTrue(torch.equal(gemv.linear(x[:1], w, bias), F.linear(x[:1], w, bias)))
|
|
strided = torch.randn(96, 64, dtype=torch.bfloat16, device="cuda").t()
|
|
self.assertTrue(torch.equal(gemv.linear(x[:1], strided), F.linear(x[:1], strided)))
|
|
|
|
def test_wrapped_dispatch_sends_only_decode_shapes_to_the_kernel(self):
|
|
calls = []
|
|
|
|
def fallback(layer, x, weight, bias=None):
|
|
calls.append(x.shape[0])
|
|
return F.linear(x, weight, bias)
|
|
|
|
gemm = gemv.wrap(lambda *args, **kwargs: fallback)("auto")
|
|
w = torch.randn(64, 96, dtype=torch.bfloat16, device="cuda")
|
|
gemm(None, torch.randn(1, 96, dtype=torch.bfloat16, device="cuda"), w)
|
|
gemm(None, torch.randn(9, 96, dtype=torch.bfloat16, device="cuda"), w)
|
|
self.assertEqual(calls, [9])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|