Skip to content

GitLab

Unverified Commit e952eee6 authored by Zhewen Li's avatar Zhewen Li Committed by GitHub
Browse files

[Bugfix] Fix `__syncwarp` on ROCM (#25996)

parent 66bca9b8
Hide whitespace changes
Inline Side-by-side
Showing with 5 additions and 1 deletion
csrc/cache_kernels.cu
View file @ e952eee6
...@@ -536,7 +536,9 @@ __global__ void indexer_k_quant_and_cache_kernel( ...@@ -536,7 +536,9 @@ __global__ void indexer_k_quant_and_cache_kernel(
for (int i = 0; i < VEC_SIZE; i++) { for (int i = 0; i < VEC_SIZE; i++) {
amax = fmaxf(amax, fabsf(float(k_val_ptr[i]))); amax = fmaxf(amax, fabsf(float(k_val_ptr[i])));
} }
#ifndef USE_ROCM
__syncwarp(); __syncwarp();
#endif
// Reduced amax // Reduced amax
for (int mask = 16; mask > 0; mask /= 2) { for (int mask = 16; mask > 0; mask /= 2) {
...@@ -546,7 +548,9 @@ __global__ void indexer_k_quant_and_cache_kernel( ...@@ -546,7 +548,9 @@ __global__ void indexer_k_quant_and_cache_kernel(
amax = fmaxf(amax, __shfl_xor_sync(unsigned(-1), amax, mask)); amax = fmaxf(amax, __shfl_xor_sync(unsigned(-1), amax, mask));
#endif #endif
} }
#ifndef USE_ROCM
__syncwarp(); __syncwarp();
#endif
float scale = fmaxf(amax, 1e-4) / 448.0f; float scale = fmaxf(amax, 1e-4) / 448.0f;
if (use_ue8m0) { if (use_ue8m0) {
scale = exp2f(ceilf(log2f(scale))); scale = exp2f(ceilf(log2f(scale)));
...@@ -1167,4 +1171,4 @@ void indexer_k_quant_and_cache( ...@@ -1167,4 +1171,4 @@ void indexer_k_quant_and_cache(
DISPATCH_BY_KV_CACHE_DTYPE(k.dtype(), "fp8_e4m3", DISPATCH_BY_KV_CACHE_DTYPE(k.dtype(), "fp8_e4m3",
CALL_INDEXER_K_QUANT_AND_CACHE); CALL_INDEXER_K_QUANT_AND_CACHE);
} }
\ No newline at end of file
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