From 716277d2467eae07d750cc5ae31a81da6ff32f1b Mon Sep 17 00:00:00 2001 From: Jagrit Digani Date: Thu, 21 Nov 2024 13:11:31 -0800 Subject: [PATCH] Enable attn build outside of jit --- mlx/backend/metal/kernels/CMakeLists.txt | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/mlx/backend/metal/kernels/CMakeLists.txt b/mlx/backend/metal/kernels/CMakeLists.txt index e6416581a..9c0856425 100644 --- 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()