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 6, 19:15 UTC
[Bugfix] ep_scatter kernel store-load race condition (#34991)
Signed-off-by: Yifan Qiao <yifanqiao@berkeley.edu>
Yifan Qiao committed
Mar 12, 2026 at 18:07 UTC
55d8073d066dee229ea218579f3884305f27c327
1 file changed
+8
-3
vllm/model_executor/layers/fused_moe/deep_gemm_utils.py
+8
-3
@@ -76,9 +76,13 @@ def _fwd_kernel_ep_scatter_1(
76
)
77
tokens_per_expert = round_up_128(tokens_per_expert)
78
cumsum = tl.cumsum(tokens_per_expert) - tokens_per_expert
79
- tl.store(expert_start_loc + offset_cumsum, cumsum, mask=offset_cumsum < num_experts)
79
81
- cur_expert_start = tl.load(expert_start_loc + cur_expert)
80
+ # Extract this block's offset from the register vector (warp shuffle,
81
+ # no global memory round-trip) then write it once to expert_start_loc.
82
+ cur_expert_start = tl.sum(
83
+ tl.where(offset_cumsum == cur_expert, cumsum, tl.zeros_like(cumsum))
84
+ )
85
+ tl.store(expert_start_loc + cur_expert, cur_expert_start)
86
cur_expert_token_num = tl.load(num_recv_tokens_per_expert + cur_expert)
87
88
m_indices_start_ptr = m_indices + cur_expert_start
@@ -87,7 +91,7 @@ def _fwd_kernel_ep_scatter_1(
91
# any rows in the per-expert aligned region that do not correspond to
92
# real tokens are left untouched here and should remain initialized to
93
# -1 so DeepGEMM can skip them
90
- for start_m in tl.range(0, cur_expert_token_num, BLOCK_E, num_stages=4):
94
+ for start_m in tl.range(0, cur_expert_token_num, BLOCK_E):
95
offs = start_m + off_expert
96
mask = offs < cur_expert_token_num
97
tl.store(
@@ -186,6 +190,7 @@ def ep_scatter(
190
grid = num_experts
191
192
assert m_indices.shape[0] % BLOCK_E == 0
193
+ assert expert_start_loc.shape[0] == num_experts
194
195
_fwd_kernel_ep_scatter_1[(grid,)](
196
num_recv_tokens_per_expert,