[AMD] Fix JIT compile failure in sgl_kernel/warp.cuh (#33195)

Co-authored-by: Bingxu Chen <bingxche@amd.com>
2026-08-02 23:35:29 -07:00
parent e9366d7f79
commit 21d930aae3
1 changed files with 4 additions and 0 deletions
@@ -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;