diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 0e55e827c20..f26e477ea2e 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -10,6 +10,7 @@ from typing import final import torch import vllm.envs as envs +from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.activation import ( MoEActivation, @@ -1132,6 +1133,17 @@ class FusedMoEKernelModularImpl: The _prepare method is a wrapper around self.prepare_finalize.prepare that handles DBO and async. """ + # Skip cudagraph/DP padding tokens uniformly across all a2a backends: + # forcing padded rows' expert ids to -1 makes every prepare_finalize drop + # them (not dispatched / not computed by the experts). The V2 model runner + # marks them in forward_context.is_padding; it is None for runners that do + # not populate it, leaving topk_ids unchanged. + is_padding = get_forward_context().is_padding + if is_padding is not None: + n = topk_ids.shape[0] + # TODO: Properly support DBO (padding lives at the batch tail). + topk_ids = torch.where(is_padding[:n].unsqueeze(1), -1, topk_ids) + if not self.prepare_finalize.supports_async(): # We shouldn't be running an a2a kernel that doesn't # support async prepare/finalize diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py index d0d487be692..129e3b5d5c2 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py @@ -289,12 +289,6 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): ) a1 = a1 * topk_weights.to(a1.dtype) - is_padding = get_forward_context().is_padding - if is_padding is not None: - n = topk_ids.shape[0] - # TODO: Properly support DBO. - topk_ids = torch.where(is_padding[:n].unsqueeze(1), -1, topk_ids) - if quant_config.is_block_quantized and not defer_input_quant: a1q, a1q_scale = moe_kernel_quantize_input( a1,