Enable attn build outside of jit

ml-explore · Nov 21, 2024 · 9702b37 · 9702b37
1 parent 292c6cd
commit 9702b37
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/mlx/backend/metal/kernels/CMakeLists.txt b/mlx/backend/metal/kernels/CMakeLists.txt
@@ -82,6 +82,8 @@ set(STEEL_ATTN_HEADERS
     steel/attn/transforms.h
     steel/attn/kernels/steel_attention.h)
 
+build_kernel(steel/attn/kernels/steel_attention ${STEEL_ATTN_HEADERS})
+
 if(NOT MLX_METAL_JIT)
   build_kernel(arange arange.h)
   build_kernel(binary binary.h binary_ops.h)
@@ -107,7 +109,6 @@ if(NOT MLX_METAL_JIT)
   build_kernel(steel/gemm/kernels/steel_gemm_fused ${STEEL_HEADERS})
   build_kernel(steel/gemm/kernels/steel_gemm_masked ${STEEL_HEADERS})
   build_kernel(steel/gemm/kernels/steel_gemm_splitk ${STEEL_HEADERS})
-  build_kernel(steel/attn/kernels/steel_attention ${STEEL_ATTN_HEADERS})
   build_kernel(gemv_masked steel/utils.h)
 endif()