File tree Expand file tree Collapse file tree 1 file changed +4
-2
lines changed
python/sglang/srt/layers/moe/ep_moe Expand file tree Collapse file tree 1 file changed +4
-2
lines changed Original file line number Diff line number Diff line change @@ -236,7 +236,8 @@ def pre_reorder_triton_kernel(
236
236
):
237
237
OutDtype = gateup_input_ptr .dtype .element_ty
238
238
239
- src_idx = tl .program_id (0 )
239
+ src_idx_int32 = tl .program_id (0 )
240
+ src_idx = src_idx_int32 .to (tl .int64 )
240
241
src2dst_ptr = src2dst_ptr + src_idx * topk
241
242
topk_ids_ptr = topk_ids_ptr + src_idx * topk
242
243
src_ptr = input_ptr + src_idx * hidden_size
@@ -255,7 +256,8 @@ def pre_reorder_triton_kernel(
255
256
else :
256
257
scale = 1.0
257
258
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 )
259
261
dst_ptr = gateup_input_ptr + dst_idx * hidden_size
260
262
for start_offset in tl .range (0 , hidden_size , BLOCK_SIZE ):
261
263
offset = start_offset + vec
You can’t perform that action at this time.
0 commit comments