[Bugfix] Fix int32 destination offset overflow in CUTLASS MoE pre-reorder (#35188)

This commit is contained in:
Benjamin Truong
2026-08-25 20:56:42 +08:00
committed by GitHub
parent 829138a31e
commit a1f9508dd4
@@ -241,7 +241,7 @@ def pre_reorder_triton_kernel_for_cutlass_moe(
for idx in range(topk):
expert_id = tl.load(token_topk_ids_ptr + idx)
if expert_id != num_local_experts:
dst_idx = tl.load(token_src2dst_ptr + idx)
dst_idx = tl.load(token_src2dst_ptr + idx).to(tl.int64)
tl.store(dst_ptr_offs + dst_idx * hidden_size, out_data, mask=mask)