add def of flash_attention

Signed-off-by: Alex Chi Z <[email protected]>

Author: skyzh

book/src/week2-overview.md

@@ -11,3 +11,7 @@ speculative decoding
 prefill and decode separation
 quantized kv cache
 Assert return data type
+
+https://github.com/ml-explore/mlx/blob/main/mlx/backend/metal/kernels/sdpa_vector.h
+https://github.com/philipturner/metal-flash-attention
+https://github.com/ml-explore/mlx/blob/main/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h

build_ext.sh

@@ -3,4 +3,3 @@
 set -e
 pdm run build-ext-ref
 cp src/extensions_ref/build/lib/tiny_llm_ext_ref/tiny_llm_ext_ref.metallib .venv/lib/python3.12/site-packages/mlx/lib/
-pdm run test-week2-ref -k 'week_2_day_2'

pyproject.toml

@@ -29,6 +29,7 @@ test.cmd = "pytest tests"
 test-week1-ref.cmd = "pytest tests_ref_impl_week1"
 test-week2-ref.cmd = "pytest tests_ref_impl_week2"
 format = "ruff format"
+format-cpp.shell = "find src/extensions_ref -type file \\( -name '*.h' -or -name '*.cpp' \\) | xargs -n1 clang-format -i"
 
 [tool.pytest.ini_options]
 addopts = [

src/extensions_ref/CMakeLists.txt

@@ -35,8 +35,8 @@ add_library(tiny_llm_ext_ref)
 target_sources(
   tiny_llm_ext_ref
   PUBLIC
-  ${CMAKE_CURRENT_LIST_DIR}/axpby/axpby.cpp
   ${CMAKE_CURRENT_LIST_DIR}/src/quantized_matmul.cpp
+  ${CMAKE_CURRENT_LIST_DIR}/src/flash_attention.cpp
 )
 
 # Add include headers
@@ -57,8 +57,8 @@ if(MLX_BUILD_METAL)
     TITLE
     tiny_llm_ext_ref
     SOURCES
-    ${CMAKE_CURRENT_LIST_DIR}/axpby/axpby.metal
     ${CMAKE_CURRENT_LIST_DIR}/src/quantized_matmul.metal
+    ${CMAKE_CURRENT_LIST_DIR}/src/flash_attention.metal
     INCLUDE_DIRS
     ${PROJECT_SOURCE_DIR}
     ${MLX_INCLUDE_DIRS}