[Bugfix] Machete garbage results for some models (large K dim) (#9212)

2024-10-10 02:16:17 -04:00
parent ce00231a8b
commit a64e7b9407
2 changed files with 16 additions and 12 deletions
--- a/tests/kernels/test_machete_gemm.py
+++ b/tests/kernels/test_machete_gemm.py
@@ -24,13 +24,14 @@ MNK_SHAPES = [
    (1, 128, 128),
    (1, 512, 1024),
    (1, 4096, 4096),
+    (1, 8192, 28672),
    (13, 8192, 4096),
    (26, 4096, 8192),
-    (1, 4096, 4096),
+    (64, 4096, 4096),
+    (64, 8192, 28672),
    (257, 128, 4096),
    (257, 4224, 4160),
    (257, 4096, 4096),
-    (64, 4096, 4096),
    (1024, 4096, 8192),
    (1024, 8192, 4096),
 ]