ColossalAI/extensions/csrc/cuda
yuehuayingxueluo 600881a8ea
[Inference]Add CUDA KVCache Kernel (#5406)
* add cuda KVCache kernel

* annotation benchmark_kvcache_copy

* add use cuda

* fix import path

* move benchmark scripts to example/

* rm benchmark codes in test_kv_cache_memcpy.py

* rm redundancy codes

* rm redundancy codes

* pr was modified according to the review
2024-02-28 14:36:50 +08:00
..
include [feat] refactored extension module (#5298) 2024-01-25 17:01:48 +08:00
colossal_C_frontend.cpp [feat] refactored extension module (#5298) 2024-01-25 17:01:48 +08:00
colossal_inference_C_frontend.cpp [Inference]Add CUDA KVCache Kernel (#5406) 2024-02-28 14:36:50 +08:00
compat.h [feat] refactored extension module (#5298) 2024-01-25 17:01:48 +08:00
cpu_adam.cpp [feat] refactored extension module (#5298) 2024-01-25 17:01:48 +08:00
cpu_adam.h [feat] refactored extension module (#5298) 2024-01-25 17:01:48 +08:00
decode_kv_cache_memcpy_kernel.cu [Inference]Add CUDA KVCache Kernel (#5406) 2024-02-28 14:36:50 +08:00
layer_norm_cuda.cpp [feat] refactored extension module (#5298) 2024-01-25 17:01:48 +08:00
layer_norm_cuda_kernel.cu [feat] refactored extension module (#5298) 2024-01-25 17:01:48 +08:00
moe_cuda.cpp [feat] refactored extension module (#5298) 2024-01-25 17:01:48 +08:00
moe_cuda_kernel.cu [feat] refactored extension module (#5298) 2024-01-25 17:01:48 +08:00
multi_tensor_adam.cu [feat] refactored extension module (#5298) 2024-01-25 17:01:48 +08:00
multi_tensor_apply.cuh [feat] refactored extension module (#5298) 2024-01-25 17:01:48 +08:00
multi_tensor_l2norm_kernel.cu [feat] refactored extension module (#5298) 2024-01-25 17:01:48 +08:00
multi_tensor_lamb.cu [feat] refactored extension module (#5298) 2024-01-25 17:01:48 +08:00
multi_tensor_scale_kernel.cu [feat] refactored extension module (#5298) 2024-01-25 17:01:48 +08:00
multi_tensor_sgd_kernel.cu [feat] refactored extension module (#5298) 2024-01-25 17:01:48 +08:00
scaled_masked_softmax.cpp [feat] refactored extension module (#5298) 2024-01-25 17:01:48 +08:00
scaled_masked_softmax.h [feat] refactored extension module (#5298) 2024-01-25 17:01:48 +08:00
scaled_masked_softmax_cuda.cu [feat] refactored extension module (#5298) 2024-01-25 17:01:48 +08:00
scaled_upper_triang_masked_softmax.cpp [feat] refactored extension module (#5298) 2024-01-25 17:01:48 +08:00
scaled_upper_triang_masked_softmax.h [feat] refactored extension module (#5298) 2024-01-25 17:01:48 +08:00
scaled_upper_triang_masked_softmax_cuda.cu [feat] refactored extension module (#5298) 2024-01-25 17:01:48 +08:00
type_shim.h [Inference]Add CUDA KVCache Kernel (#5406) 2024-02-28 14:36:50 +08:00