From 21d930aae33356aa8bcc66bcfb18eea64cf5a9ac Mon Sep 17 00:00:00 2001 From: "Wang, FangYuan" <39615225+At1a8@users.noreply.github.com> Date: Mon, 3 Aug 2026 14:35:29 +0800 Subject: [PATCH] [AMD] Fix JIT compile failure in sgl_kernel/warp.cuh (#33195) Co-authored-by: Bingxu Chen --- python/sglang/kernels/jit/include/sgl_kernel/warp.cuh | 4 ++++ 1 file changed, 4 insertions(+) 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;