Skip to content

GitLab

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

update indexer_k_cache_kernel

parent c2ef7fdd
Hide 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(
kv_cache[dst_offset + i] = __float2half(val);
} else if constexpr (std::is_same<cache_t, at::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) {
kv_cache[dst_offset + i] = val;
} 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