chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[diffusion] fix: fix accuracy bug caused by #14717 (#18296)

Browse Source
This commit is contained in:
陈一涵
2026-02-05 20:36:18 +08:00
committed by GitHub
parent 8f8c1724ae
commit 4aa03d91fd
3 changed files with 6 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/sglang/jit_kernel/diffusion/cutedsl/common/reduce.py
View File

@@ -17,7 +17,7 @@ def cta_reduce_sum(
val: cute.Numeric, num_warps: cutlass.Constexpr, tidx: cutlass.Int32
) -> cute.Numeric:
smem = cutlass.utils.SmemAllocator()
acc = smem.allocate_tensor(cutlass.Float32, num_warps)
acc = smem.allocate_tensor(cutlass.Float32, num_warps + 1)
warp_id = tidx >> 5
lane_id = tidx & 31
if lane_id == 0:
@@ -27,7 +27,7 @@ def cta_reduce_sum(
val = acc[lane_id] if lane_id < num_warps else cutlass.Float32(0)
val = warp_reduce_sum(val)
if lane_id == 0:
acc[0] = val
acc[num_warps] = val
cute.arch.sync_threads()
val = acc[0]
val = acc[num_warps]
return val