A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Bugfix] Fix `__syncwarp` on ROCM (#25996)
Zhewen Li committed
Sep 30, 2025 at 21:15 UTC
e952eee698fb8a497baaa864b93387fd180f27e7
1 file changed
+5
-1
csrc/cache_kernels.cu
+5
-1
@@ -536,7 +536,9 @@ __global__ void indexer_k_quant_and_cache_kernel(
536
for (int i = 0; i < VEC_SIZE; i++) {
537
amax = fmaxf(amax, fabsf(float(k_val_ptr[i])));
538
}
539
+#ifndef USE_ROCM
540
__syncwarp();
541
+#endif
542
543
// Reduced amax
544
for (int mask = 16; mask > 0; mask /= 2) {
@@ -546,7 +548,9 @@ __global__ void indexer_k_quant_and_cache_kernel(
548
amax = fmaxf(amax, __shfl_xor_sync(unsigned(-1), amax, mask));
549
#endif
550
}
551
+#ifndef USE_ROCM
552
__syncwarp();
553
+#endif
554
float scale = fmaxf(amax, 1e-4) / 448.0f;
555
if (use_ue8m0) {
556
scale = exp2f(ceilf(log2f(scale)));
@@ -1167,4 +1171,4 @@ void indexer_k_quant_and_cache(
1171
1172
DISPATCH_BY_KV_CACHE_DTYPE(k.dtype(), "fp8_e4m3",
1173
CALL_INDEXER_K_QUANT_AND_CACHE);
1170
-}
\ No newline at end of file
1174
+}