Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fused ROPE and reshape cache kernel #229

Open
wants to merge 21 commits into
base: main
Choose a base branch
from
Open

Fused ROPE and reshape cache kernel #229

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
21 commits
Select commit Hold shift + click to select a range
eee3c78
Initial Frame, without kernel itself
Sep 19, 2024
2f19c01
fused kernel
Sep 20, 2024
beba692
Merge branch 'main' into fused_re_and_reshape_cache_kernel
Oct 7, 2024
d7eaac6
compilation passes, begging query
Oct 9, 2024
deaf6aa
query correctly ROPEd
Oct 9, 2024
108fb65
teduced test scope is working
Oct 11, 2024
c5dca86
ordered per slot mapping test passed
Oct 11, 2024
58a93c1
kernel ready except fp8
Oct 11, 2024
ef9a664
+ fp8 for KV cache support
Oct 11, 2024
b30a7f7
vllm and test runs, vllm has corretness issue
Oct 14, 2024
f0648c3
correctness is fixed
Oct 15, 2024
afa466b
Merge branch 'main' into fused_re_and_reshape_cache_kernel
Oct 15, 2024
2548322
honest switch from fused ROPE and original
Oct 15, 2024
3af41ef
timing test
Oct 16, 2024
3806814
Merge branch 'main' into fused_re_and_reshape_cache_kernel
Nov 13, 2024
52a70b3
small refactoring of kernel
Nov 19, 2024
97582cc
moving load cos/sin to common code
maleksan85 Nov 21, 2024
47ff040
everything except values to cache
Dec 3, 2024
81bcf3e
+ value cache filling
Dec 3, 2024
d14c054
union with float4
Dec 4, 2024
f066592
dwordx4 seen
Dec 5, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -496,6 +496,7 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
"csrc/rocm/attention.cu"
"csrc/rocm/custom_kernels.cu"
"csrc/rocm/fused_kernels.cu"
"csrc/rocm/fused_rope_and_reshape_cache.cu"
"csrc/rocm/custom.cu")

define_gpu_extension_target(
Expand Down
46 changes: 23 additions & 23 deletions csrc/rocm/attention.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1266,29 +1266,29 @@ void paged_attention(
const std::string& kv_cache_dtype, double k_scale, double v_scale,
const c10::optional<torch::Tensor>& fp8_out_scale, int64_t partition_size) {
const int head_size = query.size(2);
if (kv_cache_dtype == "auto") {
if (query.dtype() == at::ScalarType::Half) {
CALL_CUSTOM_LAUNCHER_BLK_HEAD(_Float16, _Float16,
vllm::Fp8KVCacheDataType::kAuto);
} else if (query.dtype() == at::ScalarType::BFloat16) {
CALL_CUSTOM_LAUNCHER_BLK_HEAD(__hip_bfloat16, __hip_bfloat16,
vllm::Fp8KVCacheDataType::kAuto);
} else {
TORCH_CHECK(false, "Unsupported data type: ", query.dtype());
}
} else if (kv_cache_dtype == "fp8" || kv_cache_dtype == "fp8_e4m3") {
if (query.dtype() == at::ScalarType::Half) {
CALL_CUSTOM_LAUNCHER_BLK_HEAD(_Float16, uint8_t,
vllm::Fp8KVCacheDataType::kFp8E4M3);
} else if (query.dtype() == at::ScalarType::BFloat16) {
CALL_CUSTOM_LAUNCHER_BLK_HEAD(__hip_bfloat16, uint8_t,
vllm::Fp8KVCacheDataType::kFp8E4M3);
} else {
TORCH_CHECK(false, "Unsupported data type: ", query.dtype());
}
} else {
TORCH_CHECK(false, "Unsupported KV cache dtype: ", kv_cache_dtype);
}
// if (kv_cache_dtype == "auto") {
// if (query.dtype() == at::ScalarType::Half) {
// CALL_CUSTOM_LAUNCHER_BLK_HEAD(_Float16, _Float16,
// vllm::Fp8KVCacheDataType::kAuto);
// } else if (query.dtype() == at::ScalarType::BFloat16) {
// CALL_CUSTOM_LAUNCHER_BLK_HEAD(__hip_bfloat16, __hip_bfloat16,
// vllm::Fp8KVCacheDataType::kAuto);
// } else {
// TORCH_CHECK(false, "Unsupported data type: ", query.dtype());
// }
// } else if (kv_cache_dtype == "fp8" || kv_cache_dtype == "fp8_e4m3") {
// if (query.dtype() == at::ScalarType::Half) {
// CALL_CUSTOM_LAUNCHER_BLK_HEAD(_Float16, uint8_t,
// vllm::Fp8KVCacheDataType::kFp8E4M3);
// } else if (query.dtype() == at::ScalarType::BFloat16) {
// CALL_CUSTOM_LAUNCHER_BLK_HEAD(__hip_bfloat16, uint8_t,
// vllm::Fp8KVCacheDataType::kFp8E4M3);
// } else {
// TORCH_CHECK(false, "Unsupported data type: ", query.dtype());
// }
// } else {
// TORCH_CHECK(false, "Unsupported KV cache dtype: ", kv_cache_dtype);
// }
}

#undef WARP_SIZE
Expand Down
Loading
Loading