diff --git a/python/sglang/kernels/jit/include/sgl_kernel/warp.cuh b/python/sglang/kernels/jit/include/sgl_kernel/warp.cuh
index 6631a8cc4..2a787bdcf 100644
--- a/python/sglang/kernels/jit/include/sgl_kernel/warp.cuh
+++ b/python/sglang/kernels/jit/include/sgl_kernel/warp.cuh
@@ -181,7 +181,11 @@ SGL_DEVICE uint32_t inclusive_sum(uint32_t lane_id, uint32_t val) {
   static_assert(kWarpThreads == 32);
 #pragma unroll
   for (uint32_t offset = 1; offset < 32; offset *= 2) {
+#ifndef USE_ROCM
     uint32_t n = __shfl_up_sync(0xFFFFFFFF, val, offset);
+#else
+    uint32_t n = __shfl_up_sync(kFullMask, val, offset, kWarpThreads);
+#endif
     if (lane_id >= offset) val += n;
   }
   return val;