prain: Triton GEMV for decode GEMMs on gfx10 (rocBLAS ran bf16 GEMVs at 6% of the bandwidth; agent setup 1.58 -> 2.77 tok/s, same answers)

tests/test_gemv.py

@@ -0,0 +1,55 @@
+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()