Masked gemv (#1211)

ml-explore · Jun 14, 2024 · 2d6cd47 · 2d6cd47
1 parent fe3167d
commit 2d6cd47
Show file tree

Hide file tree

Showing 6 changed files with 1,627 additions and 319 deletions.
diff --git a/mlx/backend/metal/kernels/CMakeLists.txt b/mlx/backend/metal/kernels/CMakeLists.txt
@@ -37,6 +37,7 @@ endfunction(build_kernel)
 build_kernel(arg_reduce)
 build_kernel(conv steel/conv/params.h)
 build_kernel(gemv steel/utils.h)
+build_kernel(gemv_masked steel/utils.h)
 build_kernel(layer_norm)
 build_kernel(random)
 build_kernel(rms_norm)