[DeepEP V2] Fill invalid recv_topk_idx with -1 (#46432)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Jun 22, 2026 at 21:45 UTC
04c2a8deac44fdb1ca3e2b5ec3e6bf16f3f6a914
1 file changed
+64
-10
vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py
+64
-10
@@ -13,6 +13,7 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
13
TopKWeightAndReduceDelegate,
14
)
15
from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
16
+from vllm.triton_utils import tl, triton
17
from vllm.utils.math_utils import round_up
18
from vllm.v1.worker.ubatching import (
19
dbo_current_ubatch_id,
@@ -220,17 +221,22 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
221
)
222
recv_topk_idx = recv_topk_idx.unsqueeze(1)
223
else:
223
- # do_expand=False (decode/cudagraph mode): recv_topk_idx has
224
- # LOCAL expert IDs (-1 for non-local and padding rows).
225
- # Convert valid local IDs to global. Rows with -1 are
226
- # skipped by expert kernels (TrtLLM tile-level skipping,
227
- # DeepGemm is_computation_valid), so no need to zero
228
- # hidden states, scales, or weights for padding rows.
229
- valid_mask = recv_topk_idx >= 0
230
- recv_topk_idx = torch.where(
231
- valid_mask,
232
- recv_topk_idx + self.rank_expert_offset,
224
+ # do_expand=False (decode/cudagraph mode): the dispatch only writes
225
+ # rows [0, num_recv_tokens); the rest of the worst-case-allocated
226
+ # buffer is left UNINITIALIZED. For valid rows, recv_topk_idx holds
227
+ # LOCAL expert IDs (-1 for non-local slots). Convert valid local IDs
228
+ # to global and force everything else to -1:
229
+ # * non-local / out-of-range expert slots, and
230
+ # * every row >= num_recv_tokens (uninitialized padding): its
231
+ # stale contents can alias valid expert IDs and would otherwise
232
+ # be treated as real routed tokens by experts that build routing
233
+ # over *all* rows (e.g. triton MoE backend's make_routing_data),
234
+ # polluting the per-expert token lists and corrupting real tokens.
235
+ recv_topk_idx = _globalize_recv_topk_idx(
236
recv_topk_idx,
237
+ psum_recv_per_rank,
238
+ self.rank_expert_offset,
239
+ self.num_experts,
240
)
241
242
# Reshape recv_topk_weights to match recv_topk_idx shape [N, 1]
@@ -416,3 +422,51 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
422
weight_and_reduce_impl,
423
False,
424
)
425
+
426
+
427
+@triton.jit
428
+def _globalize_recv_topk_idx_kernel(
429
+ topk_idx_ptr, # [N*topk] local expert IDs (-1 = non-local), modified in place
430
+ psum_ptr, # [P] per-scaleup-rank recv prefix sum; num_recv = psum[P-1]
431
+ P,
432
+ rank_expert_offset,
433
+ num_experts,
434
+ n_elements, # N * topk
435
+ topk: tl.constexpr,
436
+ BLOCK: tl.constexpr,
437
+):
438
+ pid = tl.program_id(0)
439
+ offs = pid * BLOCK + tl.arange(0, BLOCK)
440
+ mask = offs < n_elements
441
+ # num_recv_tokens read on-device (no host sync) -> cudagraph-safe.
442
+ num_recv = tl.load(psum_ptr + P - 1)
443
+ val = tl.load(topk_idx_ptr + offs, mask=mask, other=-1)
444
+ g = val + rank_expert_offset
445
+ row = offs // topk
446
+ # Keep a slot iff: it is a local expert (val >= 0), its global id is in
447
+ # range, and its row is a real received token (< num_recv). Otherwise -1.
448
+ valid = (val >= 0) & (g < num_experts) & (row < num_recv)
449
+ tl.store(topk_idx_ptr + offs, tl.where(valid, g, -1), mask=mask)
450
+
451
+
452
+def _globalize_recv_topk_idx(
453
+ recv_topk_idx: torch.Tensor, # [N, topk] local expert IDs, -1 = non-local
454
+ psum_recv_per_rank: torch.Tensor,
455
+ rank_expert_offset: int,
456
+ num_experts: int,
457
+) -> torch.Tensor:
458
+ N, topk = recv_topk_idx.shape
459
+ n = N * topk
460
+ BLOCK = 1024
461
+ grid = (triton.cdiv(n, BLOCK),)
462
+ _globalize_recv_topk_idx_kernel[grid](
463
+ recv_topk_idx,
464
+ psum_recv_per_rank,
465
+ psum_recv_per_rank.shape[0],
466
+ rank_expert_offset,
467
+ num_experts,
468
+ n,
469
+ topk=topk,
470
+ BLOCK=BLOCK,
471
+ )
472
+ return recv_topk_idx