Fix tensor core instruction shape (#19)

* fix tensor core instruction shape * remove no-op constraint
modular · Mar 9, 2025 · 5039ede · 5039ede
1 parent 5732f20
commit 5039ede
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/custom-ops-matrix-multiplication/benchmarks.mojo b/custom-ops-matrix-multiplication/benchmarks.mojo
@@ -167,6 +167,7 @@ def matmul():
         bench_matmul_kernel["tiled_register"]()
         bench_matmul_kernel["block_tiled"]()
         bench_matmul_kernel["block_tiled_vectorized"]()
+        bench_matmul_kernel["tensor_core"]()
 
     bench.config.verbose_metric_names = False
     print(bench)

diff --git a/custom-ops-matrix-multiplication/operations/matrix_multiplication.mojo b/custom-ops-matrix-multiplication/operations/matrix_multiplication.mojo
@@ -975,7 +975,7 @@ struct MatrixMultiplication[algorithm: StringLiteral]:
                 alias WN = 32
                 alias MMA_M = 16
                 alias MMA_N = 8
-                alias MMA_K = 8
+                alias MMA_K = 4
                 alias NUM_WARPS = (BM // WM) * (BN // WN)
                 gpu_ctx.enqueue_function[
                     tensor_core_matrix_multiplication[