Add kv cache related ops (#2968)

Summary: Pull Request resolved: #2968 X-link: facebookresearch/FBGEMM#65 Reviewed By: Aya-ZIbra Differential Revision: D60951847 fbshipit-source-id: 9d3fa828fcdb3c905c8cf8ca6fbce2066252724b
pytorch · Aug 15, 2024 · 7105d5f · 7105d5f
1 parent 7143818
commit 7105d5f
Show file tree

Hide file tree

Showing 10 changed files with 3,182 additions and 618 deletions.
diff --git a/fbgemm_gpu/experimental/gen_ai/CMakeLists.txt b/fbgemm_gpu/experimental/gen_ai/CMakeLists.txt
@@ -56,10 +56,15 @@ set(comm_ops_sources
     src/comm/car.cu
     src/comm/car.cpp)
 
+set(kv_cache_ops_sources
+    src/kv_cache/kv_cache.cu
+    src/kv_cache/kv_cache.cpp)
+
 set(experimental_gen_ai_cpp_source_files
     ${attention_ops_sources}
     ${quantize_ops_sources}
-    ${comm_ops_sources})
+    ${comm_ops_sources}
+    ${kv_cache_ops_sources})
 
 # Set the source file for FB only CPP
 if(USE_FB_ONLY)