Skip to content

GitLab

Commit 9fed1f5d authored by zhuwenwen's avatar zhuwenwen
Browse files

add bf16

parent 3f1166ab
Hide whitespace changes
Inline Side-by-side
Showing with 9 additions and 9 deletions
csrc/attention/attention_dtypes.h
View file @ 9fed1f5d
......@@ -3,5 +3,5 @@
#include "attention_generic.cuh"
#include "dtype_float16.cuh"
#include "dtype_float32.cuh"
// #include "dtype_bfloat16.cuh"
#include "dtype_bfloat16.cuh"
// #include "dtype_fp8_e5m2.cuh"
csrc/attention/attention_kernels.cu
View file @ 9fed1f5d
......@@ -734,8 +734,8 @@ void paged_attention_v1(
CALL_V1_LAUNCHER_BLOCK_SIZE(float, float, false);
} else if (query.dtype() == at::ScalarType::Half) {
CALL_V1_LAUNCHER_BLOCK_SIZE(uint16_t, uint16_t, false);
// } else if (query.dtype() == at::ScalarType::BFloat16) {
// CALL_V1_LAUNCHER_BLOCK_SIZE(__nv_bfloat16, __nv_bfloat16, false);
} else if (query.dtype() == at::ScalarType::BFloat16) {
CALL_V1_LAUNCHER_BLOCK_SIZE(__nv_bfloat16, __nv_bfloat16, false);
} else {
TORCH_CHECK(false, "Unsupported data type: ", query.dtype());
}
......@@ -927,8 +927,8 @@ void paged_attention_v2(
CALL_V2_LAUNCHER_BLOCK_SIZE(float, float, false);
} else if (query.dtype() == at::ScalarType::Half) {
CALL_V2_LAUNCHER_BLOCK_SIZE(uint16_t, uint16_t, false);
// } else if (query.dtype() == at::ScalarType::BFloat16) {
// CALL_V2_LAUNCHER_BLOCK_SIZE(__nv_bfloat16, __nv_bfloat16, false);
} else if (query.dtype() == at::ScalarType::BFloat16) {
CALL_V2_LAUNCHER_BLOCK_SIZE(__nv_bfloat16, __nv_bfloat16, false);
} else {
TORCH_CHECK(false, "Unsupported data type: ", query.dtype());
}
......
csrc/cache_kernels.cu
View file @ 9fed1f5d
......@@ -13,10 +13,10 @@
#include <map>
#include <vector>
// #ifdef USE_ROCM
// #include <hip/hip_bf16.h>
// typedef __hip_bfloat16 __nv_bfloat16;
// #endif
#ifdef USE_ROCM
#include <hip/hip_bf16.h>
typedef __hip_bfloat16 __nv_bfloat16;
#endif
void swap_blocks(
torch::Tensor& src,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment