biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Implement single_query_cached_kv_attention kernel (#3)

Browse Source
This commit is contained in:
Woosuk Kwon
2023-03-01 15:02:19 -08:00
committed by GitHub
parent cbf8779afa
commit 0deacbce6e
12 changed files with 2140 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
setup.py
View File

@@ -9,15 +9,22 @@ ext_modules = []
# Cache operations.
cache_extension = cpp_extension.CUDAExtension(
name='cacheflow.ops',
name='cacheflow.cache_ops',
sources=['csrc/cache.cpp', 'csrc/cache_kernels.cu'],
extra_compile_args={'cxx': CXX_FLAGS, 'nvcc': NVCC_FLAGS},
)
ext_modules.append(cache_extension)
# Attention kernels.
attention_extension = cpp_extension.CUDAExtension(
name='cacheflow.attention_ops',
sources=['csrc/attention.cpp', 'csrc/attention_kernels.cu'],
extra_compile_args={'cxx': CXX_FLAGS, 'nvcc': NVCC_FLAGS},
)
ext_modules.append(attention_extension)
setuptools.setup(
name='cacheflow',
requires_python='>=3.9',
ext_modules=ext_modules,
cmdclass={'build_ext': cpp_extension.BuildExtension},
)