From bc3bf1988a6a55198130d5ee3c441925554ab42a Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Fri, 10 Jul 2026 15:00:24 +0800 Subject: [PATCH] asm gemm test cleanups try 2 (#16964) * fix hipcc requirement * move hipcc * fp8 in llama * update some of those skips * update dtype * change --- test/backend/test_asm_gemm.py | 55 +++++++++++++++++------------------ 1 file changed, 26 insertions(+), 29 deletions(-) diff --git a/test/backend/test_asm_gemm.py b/test/backend/test_asm_gemm.py index 9775e8e611..7c64897d87 100644 --- a/test/backend/test_asm_gemm.py +++ b/test/backend/test_asm_gemm.py @@ -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")