Skip to content

GitLab

Commit 51f194f2 authored by zhuwenwen's avatar zhuwenwen
Browse files

update indexer_k_cache_kernel

parent c2ef7fdd
Show whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
csrc/cache_kernels.cu
View file @ 51f194f2
...@@ -806,7 +806,8 @@ __global__ void indexer_k_cache_kernel( ...@@ -806,7 +806,8 @@ __global__ void indexer_k_cache_kernel(
kv_cache[dst_offset + i] = __float2half(val); kv_cache[dst_offset + i] = __float2half(val);
} else if constexpr (std::is_same<cache_t, at::BFloat16>::value || } else if constexpr (std::is_same<cache_t, at::BFloat16>::value ||
std::is_same<cache_t, __nv_bfloat16>::value) { std::is_same<cache_t, __nv_bfloat16>::value) {
kv_cache[dst_offset + i] = __float2bfloat16(val); __hip_bfloat16 bf16_val = __float2bfloat16(val);
kv_cache[dst_offset + i] = *reinterpret_cast<at::BFloat16*>(&bf16_val);
} else if constexpr (std::is_same<cache_t, float>::value) { } else if constexpr (std::is_same<cache_t, float>::value) {
kv_cache[dst_offset + i] = val; kv_cache[dst_offset + i] = val;
} else { } else {
......
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