diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index 3daae242d45..25b431429bd 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/tests/models/test_deepseek_v4_mega_moe.py @@ -46,7 +46,8 @@ def test_deepseek_v4_mega_moe_ue8m0_uint8_to_float(): def test_deepseek_v4_mega_moe_weight_loader_uses_ep_expert_ownership(): vllm_config = SimpleNamespace( - scheduler_config=SimpleNamespace(max_num_batched_tokens=4) + scheduler_config=SimpleNamespace(max_num_batched_tokens=4), + compilation_config=SimpleNamespace(static_forward_context={}), ) experts = DeepseekV4MegaMoEExperts( vllm_config, @@ -182,3 +183,81 @@ def test_deepseek_v4_mega_moe_fused_input_staging_is_bitwise_exact(): fused_topk_weights.view(torch.uint8), ref_topk_weights.view(torch.uint8), ) + + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.", +) +def test_deepseek_v4_mega_moe_fused_input_staging_masks_padding(): + from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8 + + device = torch.device("cuda") + num_tokens = 7 + hidden_size = 256 + top_k = 8 + + generator = torch.Generator(device=device) + generator.manual_seed(1) + hidden_states = torch.randn( + num_tokens, + hidden_size, + device=device, + dtype=torch.bfloat16, + generator=generator, + ) + topk_ids = torch.randint( + 0, + 256, + (num_tokens, top_k), + device=device, + dtype=torch.int32, + generator=generator, + ) + topk_weights = torch.randn( + num_tokens, + top_k, + device=device, + dtype=torch.float32, + generator=generator, + ) + is_padding = torch.tensor( + [False, True, False, False, True, False, True], + device=device, + ) + + ref_x, ref_x_sf = per_token_cast_to_fp8( + hidden_states, + use_ue8m0=True, + gran_k=32, + use_packed_ue8m0=True, + ) + ref_topk_idx = topk_ids.to(torch.int64) + ref_topk_idx[is_padding] = -1 + ref_topk_weights = topk_weights.clone() + ref_topk_weights[is_padding] = 0.0 + + fused_x = torch.empty_like(ref_x) + fused_x_sf = torch.empty_like(ref_x_sf) + fused_topk_idx = torch.empty_like(ref_topk_idx) + fused_topk_weights = torch.empty_like(ref_topk_weights) + + prepare_megamoe_inputs( + hidden_states, + topk_weights, + topk_ids, + fused_x, + fused_x_sf, + fused_topk_idx, + fused_topk_weights, + is_padding=is_padding, + ) + torch.accelerator.synchronize() + + assert torch.equal(fused_x.view(torch.uint8), ref_x.view(torch.uint8)) + assert torch.equal(fused_x_sf, ref_x_sf) + assert torch.equal(fused_topk_idx, ref_topk_idx) + assert torch.equal( + fused_topk_weights.view(torch.uint8), + ref_topk_weights.view(torch.uint8), + ) diff --git a/vllm/envs.py b/vllm/envs.py index d38014e468b..9cfd4792e14 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -185,6 +185,7 @@ if TYPE_CHECKING: "relax", ] = "relax" VLLM_USE_FUSED_MOE_GROUPED_TOPK: bool = True + VLLM_MOE_SKIP_PADDING: bool = False VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER: bool = True VLLM_USE_FLASHINFER_MOE_INT4: bool = False VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR: str | None = None @@ -1469,6 +1470,11 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_USE_FUSED_MOE_GROUPED_TOPK": lambda: bool( int(os.getenv("VLLM_USE_FUSED_MOE_GROUPED_TOPK", "1")) ), + # Skip cudagraph/DP padding tokens in the MoE path by forcing their expert + # ids to -1 so the dispatch and experts drop them. Requires a MoE kernel that + # treats topk_id == -1 as a skip sentinel; off by default because not all + # kernels support it yet. + "VLLM_MOE_SKIP_PADDING": lambda: bool(int(os.getenv("VLLM_MOE_SKIP_PADDING", "0"))), # Allow use of FlashInfer FP8 block-scale GEMM for linear layers. # This uses TensorRT-LLM kernels and requires SM90+ (Hopper). "VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER": lambda: bool( diff --git a/vllm/forward_context.py b/vllm/forward_context.py index 5527ec13b06..10f400364ee 100644 --- a/vllm/forward_context.py +++ b/vllm/forward_context.py @@ -147,6 +147,11 @@ class ForwardContext: ubatch_slices: UBatchSlices | None = None + # Boolean mask over the token axis: True for padding rows that are not real + # tokens. Consumers can use it to skip work for padded tokens. None when + # the producer does not set it. + is_padding: torch.Tensor | None = None + # If True, bypass the compiled model call, e.g. by using .forward() directly skip_compiled: bool = False @@ -211,6 +216,7 @@ def create_forward_context( slot_mapping: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None, additional_kwargs: dict[str, Any] | None = None, skip_compiled: bool = False, + is_padding: torch.Tensor | None = None, ): if vllm_config.compilation_config.fast_moe_cold_start: all_moe_layers = vllm_config.compilation_config.static_all_moe_layers @@ -228,6 +234,7 @@ def create_forward_context( ubatch_slices=ubatch_slices, skip_compiled=skip_compiled, additional_kwargs=additional_kwargs or {}, + is_padding=is_padding, ) @@ -257,6 +264,7 @@ def set_forward_context( ubatch_slices: UBatchSlices | None = None, slot_mapping: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None, skip_compiled: bool = False, + is_padding: torch.Tensor | None = None, ): """A context manager that stores the current forward context, can be attention metadata, etc. @@ -316,6 +324,7 @@ def set_forward_context( slot_mapping, additional_kwargs, skip_compiled, + is_padding=is_padding, ) try: diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 9f3ac1fd79d..cca978d884a 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, is_forward_context_available from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.activation import ( MoEActivation, @@ -1133,6 +1134,22 @@ 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. + # Gated by VLLM_MOE_SKIP_PADDING (off by default) because this requires the + # experts kernel to treat topk_id == -1 as a skip sentinel, which not all + # MoE backends support yet. + is_padding = None + if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): + 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/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index aa60ad34ce3..99373361922 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -8,6 +8,7 @@ import regex as re import torch import torch.nn as nn +import vllm.envs as envs from vllm.config import VllmConfig from vllm.distributed import ( get_ep_group, @@ -16,6 +17,7 @@ from vllm.distributed import ( get_tensor_model_parallel_world_size, ) from vllm.distributed.eplb.eplb_state import EplbLayerState +from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.model_executor.kernels.mhc.tilelang import ( hc_head_fused_kernel_tilelang, mhc_fused_post_pre_tilelang, @@ -442,6 +444,11 @@ class DeepseekV4MegaMoEExperts(nn.Module): symm_buffer = self.get_symm_buffer() num_tokens = hidden_states.shape[0] + is_padding = None + if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): + is_padding = get_forward_context().is_padding + if is_padding is not None: + is_padding = is_padding[:num_tokens] # EPLB: map logical expert IDs to physical replicas and record load. eplb_state = self.eplb_state @@ -449,6 +456,8 @@ class DeepseekV4MegaMoEExperts(nn.Module): assert eplb_state.expert_load_view is not None assert eplb_state.logical_replica_count is not None assert eplb_state.should_record_tensor is not None + if is_padding is not None: + topk_ids = torch.where(is_padding.unsqueeze(1), -1, topk_ids) topk_ids = eplb_map_to_physical_and_record( topk_ids=topk_ids, expert_load_view=eplb_state.expert_load_view, @@ -465,6 +474,7 @@ class DeepseekV4MegaMoEExperts(nn.Module): symm_buffer.x_sf[:num_tokens], symm_buffer.topk_idx[:num_tokens], symm_buffer.topk_weights[:num_tokens], + is_padding=is_padding, ) # This method must have been already called during the weight loading phase. diff --git a/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py b/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py index 7cdb39e9b68..dac86be6edb 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py +++ b/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py @@ -19,6 +19,7 @@ def _prepare_megamoe_inputs_kernel( x_sf, topk_ids, topk_weights, + is_padding, topk_idx_out, topk_weights_out, hidden_stride_m: tl.constexpr, @@ -31,6 +32,7 @@ def _prepare_megamoe_inputs_kernel( topk_ids_stride_k: tl.constexpr, topk_weights_stride_m: tl.constexpr, topk_weights_stride_k: tl.constexpr, + is_padding_stride_m: tl.constexpr, topk_idx_stride_m: tl.constexpr, topk_idx_stride_k: tl.constexpr, topk_weights_out_stride_m: tl.constexpr, @@ -85,12 +87,16 @@ def _prepare_megamoe_inputs_kernel( if k_block_id == 0: topk_offsets = tl.arange(0, BLOCK_TOPK) topk_mask = topk_offsets < top_k + token_is_padding = False + if is_padding is not None: + token_is_padding = tl.load(is_padding + token_id * is_padding_stride_m) ids = tl.load( topk_ids + token_id * topk_ids_stride_m + topk_offsets * topk_ids_stride_k, mask=topk_mask, other=0, ).to(tl.int64) + ids = tl.where(token_is_padding, -1, ids) tl.store( topk_idx_out + token_id * topk_idx_stride_m @@ -106,6 +112,7 @@ def _prepare_megamoe_inputs_kernel( mask=topk_mask, other=0.0, ) + weights = tl.where(token_is_padding, 0.0, weights) tl.store( topk_weights_out + token_id * topk_weights_out_stride_m @@ -123,6 +130,7 @@ def prepare_megamoe_inputs( x_sf: torch.Tensor, topk_idx_out: torch.Tensor, topk_weights_out: torch.Tensor, + is_padding: torch.Tensor | None = None, ) -> None: num_tokens, hidden_size = hidden_states.shape if num_tokens == 0: @@ -142,12 +150,14 @@ def prepare_megamoe_inputs( block_k = 128 grid = (num_tokens, triton.cdiv(hidden_size, block_k)) block_topk = triton.next_power_of_2(top_k) + padding_stride_m = is_padding.stride(0) if is_padding is not None else 0 _prepare_megamoe_inputs_kernel[grid]( hidden_states, x_fp8, x_sf, topk_ids, topk_weights, + is_padding, topk_idx_out, topk_weights_out, hidden_states.stride(0), @@ -160,6 +170,7 @@ def prepare_megamoe_inputs( topk_ids.stride(1), topk_weights.stride(0), topk_weights.stride(1), + padding_stride_m, topk_idx_out.stride(0), topk_idx_out.stride(1), topk_weights_out.stride(0), diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index dad1777b47e..aa022f6d99e 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -471,6 +471,9 @@ class ModelCudaGraphManager(CudaGraphManager): skip_attn=(desc.cg_mode == CUDAGraphMode.PIECEWISE), ) + # Capture with dummy rows marked as padding. + input_buffers.is_padding.fill_(True) + def forward_fn(cg_mode: CUDAGraphMode) -> None: batch_descriptor = None if cg_mode == CUDAGraphMode.PIECEWISE: @@ -488,6 +491,7 @@ class ModelCudaGraphManager(CudaGraphManager): num_tokens_across_dp=num_tokens_across_dp, slot_mapping=slot_mappings, batch_descriptor=batch_descriptor, + is_padding=input_buffers.is_padding[:num_tokens], ): if cg_mode == CUDAGraphMode.PIECEWISE: # PIECEWISE graph (compiled PW or breakable, chosen inside diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index 6b750fe7ebf..d745dc6abf9 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -22,6 +22,7 @@ class InputBuffers: self.input_ids = torch.zeros(max_num_tokens, dtype=torch.int32, device=device) self.positions = torch.zeros(max_num_tokens, dtype=torch.int64, device=device) + self.is_padding = torch.zeros(max_num_tokens, dtype=torch.bool, device=device) self.query_start_loc = torch.zeros( max_num_reqs + 1, dtype=torch.int32, device=device ) @@ -83,6 +84,8 @@ class InputBatch: input_ids: torch.Tensor # [num_tokens_after_padding] positions: torch.Tensor + # [num_tokens_after_padding] + is_padding: torch.Tensor # [total_num_logits] logits_indices: torch.Tensor @@ -134,6 +137,9 @@ class InputBatch: input_ids = input_buffers.input_ids[:num_tokens].zero_() positions = input_buffers.positions[:num_tokens].zero_() + input_buffers.is_padding[:num_tokens].fill_(True) + is_padding = input_buffers.is_padding[:num_tokens] + logits_indices = query_start_loc[1:] - 1 cu_num_logits = torch.arange(num_reqs + 1, device=device, dtype=torch.int32) cu_num_logits_np = np.arange(num_reqs + 1, dtype=np.int32) @@ -164,6 +170,7 @@ class InputBatch: max_seq_len_np=None, input_ids=input_ids, positions=positions, + is_padding=is_padding, logits_indices=logits_indices, cu_num_logits=cu_num_logits, cu_num_logits_np=cu_num_logits_np, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 30ca2ddc562..4e1594e8065 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -27,6 +27,7 @@ import numpy as np import torch import torch.nn as nn +import vllm.envs as envs from vllm.compilation.counter import compilation_counter from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode @@ -847,6 +848,13 @@ class GPUModelRunner(LoRAModelRunnerMixin): num_tokens = scheduler_output.total_num_scheduled_tokens num_tokens_after_padding = batch_desc.num_tokens assert num_tokens > 0 + if envs.VLLM_MOE_SKIP_PADDING: + # Mark trailing cudagraph-padding rows so kernels can skip work for + # them when supported. + self.input_buffers.is_padding[:num_tokens].fill_(False) + self.input_buffers.is_padding[num_tokens:num_tokens_after_padding].fill_( + True + ) num_tokens_per_req = scheduler_output.num_scheduled_tokens num_reqs = len(num_tokens_per_req) @@ -1001,6 +1009,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): max_seq_len_np=max_seq_len_np, input_ids=self.input_buffers.input_ids[:num_tokens_after_padding], positions=self.input_buffers.positions[:num_tokens_after_padding], + is_padding=self.input_buffers.is_padding[:num_tokens_after_padding], logits_indices=logits_indices, cu_num_logits=cu_num_logits, cu_num_logits_np=cu_num_logits_np, @@ -1277,6 +1286,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): batch_descriptor=batch_descriptor, slot_mapping=slot_mappings_by_layer, skip_compiled=skip_compiled, + is_padding=input_batch.is_padding, ): self.kv_connector.pre_forward(scheduler_output) if batch_desc.cg_mode == CUDAGraphMode.PIECEWISE: