nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

asm gemm test cleanups try 2 (#16964)

Browse Source 
* fix hipcc requirement

* move hipcc

* fp8 in llama

* update some of those skips

* update dtype

* change
This commit is contained in:
qazal
2026-07-10 16:00:24 +09:00
committed by GitHub
parent 34452efa22
commit bc3bf1988a
1 changed files with 26 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/backend/test_asm_gemm.py
+26 -29
View File
@@ -9,6 +9,11 @@ from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8, FP8_MAX
# Use DEV=NULL:HIP:gfx950 to also test the assembly
def is_cdna4(): return Device[Device.DEFAULT].renderer.target.arch.startswith("gfx950")
def has_hipcc():
try: system("hipcc --version")
except Exception: return False
return True
def run_asm_gemm(a_shape, b_shape, dtype=dtypes.bfloat16, a_shard=None, b_shard=None, gpus:int=1) -> None:
Tensor.manual_seed(0)
input_dtype = dtypes.bfloat16 if dtype == FP8_DTYPE else dtype
@@ -141,11 +146,13 @@ class TestAsmGEMM(unittest.TestCase):
verify_asm_gemm(1, 256, 1000, 256)
# test the Asm GEMM with Llama shapes, only run on the real machine for speed
@unittest.skipUnless(has_hipcc(), "requires hipcc to compile")
class TestGemmLlama(unittest.TestCase):
dtype = dtypes.bfloat16
dtype = FP8_DTYPE
def setUp(self):
if not is_cdna4() or DEV.interface.startswith("MOCK") or not has_hipcc():
if not is_cdna4() or DEV.interface.startswith("MOCK"):
self.skipTest("very slow on non mi350x")
def test_empty(self): asm_gemm(Tensor.empty(N:=getenv("N", 4096), N, dtype=self.dtype), Tensor.empty(N, N, dtype=self.dtype)).realize()
@@ -172,18 +179,16 @@ class TestGemmLlama(unittest.TestCase):
def test_gemm_batched(self): verify_asm_gemm(2, 8192, 4096, 4096, dtype=self.dtype)
def test_gemm1(self): verify_asm_gemm(8, 8192, 4096, 14336, dtype=self.dtype, gpus=8)
@unittest.skip("disabled, asm in this shape is slower than tinygrad")
def test_gemm2(self): verify_asm_gemm(8, 8192, 128256, 4096, dtype=self.dtype, gpus=8)
def test_gemm3(self): verify_asm_gemm(8, 8192, 14336, 4096, dtype=self.dtype, gpus=8)
def test_gemm4(self): verify_asm_gemm(8, 4096, 14336, 4096, dtype=self.dtype, gpus=8)
def test_gemm5(self): verify_asm_gemm(8, 4096, 4096, 14336, dtype=self.dtype, gpus=8)
def test_gemm6(self): verify_asm_gemm(16, 4096, 4096, 14336, dtype=self.dtype, gpus=8)
@unittest.skip("disabled, asm in this shape is slower than tinygrad")
def test_gemm7(self): verify_asm_gemm(1, 8192, 128256, 4096, dtype=self.dtype)
def test_gemm8(self): verify_asm_gemm(1, 4096, 14336, 8192, dtype=self.dtype)
def test_gemm9(self): verify_asm_gemm(8, 4096, 14336, 8192, dtype=self.dtype, gpus=8)
def test_gemm10(self): verify_asm_gemm(1, 4096, 8192, 4096, dtype=self.dtype)
def test_gemm_previously_unsupported(self): verify_asm_gemm(8, 1024, 1024, 4096, gpus=8)
def test_gemm11(self): verify_asm_gemm(8, 1024, 1024, 4096, dtype=self.dtype, gpus=8)
def test_k_sharded_1(self): verify_asm_gemm_k_sharded(14336, 4096, 8*8192, dtype=self.dtype, gpus=8)
def test_k_sharded_2(self): verify_asm_gemm_k_sharded(4096, 14336, 8*8192, dtype=self.dtype, gpus=8)
def test_k_sharded_3(self): verify_asm_gemm_k_sharded(4096, 4096, 8*8192, dtype=self.dtype, gpus=8)
@@ -203,33 +208,25 @@ class TestGemmLlama(unittest.TestCase):
def test_tp_k_sharded_w2(self): verify_asm_gemm_k_sharded_3d(1, 8192, 4096, 14336, dtype=self.dtype, gpus=8)
# more shapes: vary M, N, K independently
def test_shape_small_square(self): verify_asm_gemm(1, 256, 256, 256)
def test_shape_small_rect_m(self): verify_asm_gemm(1, 512, 256, 256)
def test_shape_small_rect_n(self): verify_asm_gemm(1, 256, 512, 256)
def test_shape_small_rect_k(self): verify_asm_gemm(1, 256, 256, 512)
def test_shape_tall(self): verify_asm_gemm(1, 2048, 256, 256)
def test_shape_wide(self): verify_asm_gemm(1, 256, 2048, 256)
def test_shape_deep(self): verify_asm_gemm(1, 256, 256, 4096)
def test_shape_non_square(self): verify_asm_gemm(1, 1024, 2048, 512)
def test_shape_batched_small(self): verify_asm_gemm(2, 256, 256, 256)
def test_shape_batched_rect(self): verify_asm_gemm(2, 512, 1024, 256)
# K edge cases: iters=1,2,3 exercise different loop paths
def test_shape_k64(self): verify_asm_gemm(1, 256, 256, 64)
def test_shape_k128(self): verify_asm_gemm(1, 256, 256, 128)
def test_shape_k192(self): verify_asm_gemm(1, 256, 256, 192)
def test_shape_small_square(self): verify_asm_gemm(1, 256, 256, 256, dtype=self.dtype)
def test_shape_small_rect_m(self): verify_asm_gemm(1, 512, 256, 256, dtype=self.dtype)
def test_shape_small_rect_n(self): verify_asm_gemm(1, 256, 512, 256, dtype=self.dtype)
def test_shape_small_rect_k(self): verify_asm_gemm(1, 256, 256, 512, dtype=self.dtype)
def test_shape_tall(self): verify_asm_gemm(1, 2048, 256, 256, dtype=self.dtype)
def test_shape_wide(self): verify_asm_gemm(1, 256, 2048, 256, dtype=self.dtype)
def test_shape_deep(self): verify_asm_gemm(1, 256, 256, 4096, dtype=self.dtype)
def test_shape_non_square(self): verify_asm_gemm(1, 1024, 2048, 512, dtype=self.dtype)
def test_shape_batched_small(self): verify_asm_gemm(2, 256, 256, 256, dtype=self.dtype)
def test_shape_batched_rect(self): verify_asm_gemm(2, 512, 1024, 256, dtype=self.dtype)
# K edge cases: iters=1,2,3 exercise different loop path
def test_shape_k64(self): verify_asm_gemm(1, 256, 256, 64, dtype=self.dtype)
def test_shape_k128(self): verify_asm_gemm(1, 256, 256, 128, dtype=self.dtype)
def test_shape_k192(self): verify_asm_gemm(1, 256, 256, 192, dtype=self.dtype)
def test_llama3_out1(self): verify_asm_gemm(1, 8192, 128256, 4096, dtype=self.dtype)
def test_llama3_out2(self): verify_asm_gemm(1, 8192, 4096, 128256, dtype=self.dtype)
def test_llama3_out3(self): verify_asm_gemm(1, 4096, 128256, 8192, dtype=self.dtype)
def has_hipcc():
try: system("hipcc --version")
except Exception: return False
return True
@unittest.skipUnless(has_hipcc(), "FP8 gemm requires hipcc to compile")
class TestGemmLlamaFP8(TestGemmLlama): dtype = FP8_DTYPE
# mxfp8: 1x32 block scaling along K, e8m0 scales packed iteration-major (K/128, dim) uint32
def quantize_mxfp8(x:Tensor) -> tuple[Tensor, Tensor, Tensor]:
rows, K = x.shape
@@ -327,7 +324,7 @@ def run_mx_prequant(M:int, N:int, K:int) -> None:
err = ((t.float() - r.float()).abs().mean() / (r.float().abs().mean() + 1e-8)).item()
assert err < 6e-2, f"{name} prequant vs analytic rel err {err}"
@unittest.skipUnless(has_hipcc(), "MXFP8 gemm requires hipcc to compile")
@unittest.skipUnless(has_hipcc(), "requires hipcc to compile")
class TestGemmMXFP8(unittest.TestCase):
def setUp(self):
if not is_cdna4() or DEV.interface.startswith("MOCK"): self.skipTest("mxfp8 gemm is only for cdna4")
@@ -368,7 +365,7 @@ def run_atb_gemm(rows, M, N, a_shard=None, b_shard=None, gpus=1, atol=1.0, rtol=
out = hk_bf16_atb_gemm(a, b)
np.testing.assert_allclose(out.float().numpy(), ref.numpy(), atol=atol, rtol=rtol)
@unittest.skipUnless(has_hipcc(), "MXFP8 gemm requires hipcc to compile")
@unittest.skipUnless(has_hipcc(), "requires hipcc to compile")
class TestHkBf16AtbGemm(unittest.TestCase):
def setUp(self):
if not is_cdna4(): self.skipTest("hk bf16 atb gemm is cdna4 only")