Skip to content

Commit 5f6756b

Browse files
authored
[BugFix] fix pre_reorder_triton_kernel default int32 issue (sgl-project#7814)
1 parent 98aa836 commit 5f6756b

File tree

1 file changed

+4
-2
lines changed

1 file changed

+4
-2
lines changed

python/sglang/srt/layers/moe/ep_moe/kernels.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -236,7 +236,8 @@ def pre_reorder_triton_kernel(
236236
):
237237
OutDtype = gateup_input_ptr.dtype.element_ty
238238

239-
src_idx = tl.program_id(0)
239+
src_idx_int32 = tl.program_id(0)
240+
src_idx = src_idx_int32.to(tl.int64)
240241
src2dst_ptr = src2dst_ptr + src_idx * topk
241242
topk_ids_ptr = topk_ids_ptr + src_idx * topk
242243
src_ptr = input_ptr + src_idx * hidden_size
@@ -255,7 +256,8 @@ def pre_reorder_triton_kernel(
255256
else:
256257
scale = 1.0
257258

258-
dst_idx = tl.load(src2dst_ptr + idx)
259+
dst_idx_int32 = tl.load(src2dst_ptr + idx)
260+
dst_idx = dst_idx_int32.to(tl.int64)
259261
dst_ptr = gateup_input_ptr + dst_idx * hidden_size
260262
for start_offset in tl.range(0, hidden_size, BLOCK_SIZE):
261263
offset = start_offset + vec

0 commit comments

Comments
 (0)