A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 06:15 UTC
Fix DeepGEMM ep_scatter output address overflow (#39213)
Signed-off-by: S1ro1 <matej.sirovatka@gmail.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com>
Matej Sirovatka committed
May 5, 2026 at 11:56 UTC
48954de2378d19bc0da771930c1d7a202b5e58e0
1 file changed
+4
-1
vllm/model_executor/layers/fused_moe/deep_gemm_utils.py
+4
-1
@@ -140,6 +140,8 @@ def _fwd_kernel_ep_scatter_2(
140
offset_in_s = tl.arange(0, SCALE_HIDDEN_SIZE_PAD)
141
mask_s = offset_in_s < SCALE_HIDDEN_SIZE
142
143
+ output_tensor_stride0 = output_tensor_stride0.to(tl.int64)
144
+
145
for token_id in range(start_token_id, total_token_num, grid_num):
146
to_copy = tl.load(recv_x + token_id * recv_x_stride0 + offset_in, mask=mask)
147
to_copy_s = tl.load(
@@ -154,12 +156,13 @@ def _fwd_kernel_ep_scatter_2(
156
157
if expert_id >= 0:
158
dest_token_index = tl.atomic_add(expert_start_loc + expert_id, 1)
159
+ dest_token_index_i64 = dest_token_index.to(tl.int64)
160
tl.store(
161
output_index + token_id * output_index_stride0 + topk_index,
162
dest_token_index,
163
)
164
output_tensor_ptr = (
162
- output_tensor + dest_token_index * output_tensor_stride0
165
+ output_tensor + dest_token_index_i64 * output_tensor_stride0
166
)
167
output_tensor_scale_ptr = (
168
output_tensor_scale + dest_token_index * output_tensor_scale_stride0