@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 0a3e2dbc09

[Optimization] Skip DP padding tokens in MoE (#46428)

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Woosuk Kwon committed Jun 23, 2026 at 14:54 UTC 0a3e2dbc09c8a70dfd18f728bb829bf29ffa7da6

9 files changed +154 -1

tests/models/test_deepseek_v4_mega_moe.py

+80 -1

@@ -46,7 +46,8 @@ def test_deepseek_v4_mega_moe_ue8m0_uint8_to_float():

47 def test_deepseek_v4_mega_moe_weight_loader_uses_ep_expert_ownership():

48 vllm_config = SimpleNamespace(

49 - scheduler_config=SimpleNamespace(max_num_batched_tokens=4)

49 + scheduler_config=SimpleNamespace(max_num_batched_tokens=4),

50 + compilation_config=SimpleNamespace(static_forward_context={}),

51 )

52 experts = DeepseekV4MegaMoEExperts(

53 vllm_config,

@@ -182,3 +183,81 @@ def test_deepseek_v4_mega_moe_fused_input_staging_is_bitwise_exact():

183 fused_topk_weights.view(torch.uint8),

184 ref_topk_weights.view(torch.uint8),

185 )

186 +

187 +

188 +@pytest.mark.skipif(

189 + not torch.cuda.is_available(),

190 + reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.",

191 +)

192 +def test_deepseek_v4_mega_moe_fused_input_staging_masks_padding():

193 + from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8

194 +

195 + device = torch.device("cuda")

196 + num_tokens = 7

197 + hidden_size = 256

198 + top_k = 8

199 +

200 + generator = torch.Generator(device=device)

201 + generator.manual_seed(1)

202 + hidden_states = torch.randn(

203 + num_tokens,

204 + hidden_size,

205 + device=device,

206 + dtype=torch.bfloat16,

207 + generator=generator,

208 + )

209 + topk_ids = torch.randint(

210 + 0,

211 + 256,

212 + (num_tokens, top_k),

213 + device=device,

214 + dtype=torch.int32,

215 + generator=generator,

216 + )

217 + topk_weights = torch.randn(

218 + num_tokens,

219 + top_k,

220 + device=device,

221 + dtype=torch.float32,

222 + generator=generator,

223 + )

224 + is_padding = torch.tensor(

225 + [False, True, False, False, True, False, True],

226 + device=device,

227 + )

228 +

229 + ref_x, ref_x_sf = per_token_cast_to_fp8(

230 + hidden_states,

231 + use_ue8m0=True,

232 + gran_k=32,

233 + use_packed_ue8m0=True,

234 + )

235 + ref_topk_idx = topk_ids.to(torch.int64)

236 + ref_topk_idx[is_padding] = -1

237 + ref_topk_weights = topk_weights.clone()

238 + ref_topk_weights[is_padding] = 0.0

239 +

240 + fused_x = torch.empty_like(ref_x)

241 + fused_x_sf = torch.empty_like(ref_x_sf)

242 + fused_topk_idx = torch.empty_like(ref_topk_idx)

243 + fused_topk_weights = torch.empty_like(ref_topk_weights)

244 +

245 + prepare_megamoe_inputs(

246 + hidden_states,

247 + topk_weights,

248 + topk_ids,

249 + fused_x,

250 + fused_x_sf,

251 + fused_topk_idx,

252 + fused_topk_weights,

253 + is_padding=is_padding,

254 + )

255 + torch.accelerator.synchronize()

256 +

257 + assert torch.equal(fused_x.view(torch.uint8), ref_x.view(torch.uint8))

258 + assert torch.equal(fused_x_sf, ref_x_sf)

259 + assert torch.equal(fused_topk_idx, ref_topk_idx)

260 + assert torch.equal(

261 + fused_topk_weights.view(torch.uint8),

262 + ref_topk_weights.view(torch.uint8),

263 + )

vllm/envs.py

@@ -185,6 +185,7 @@ if TYPE_CHECKING:

185 "relax",

186 ] = "relax"

187 VLLM_USE_FUSED_MOE_GROUPED_TOPK: bool = True

188 + VLLM_MOE_SKIP_PADDING: bool = False

189 VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER: bool = True

190 VLLM_USE_FLASHINFER_MOE_INT4: bool = False

191 VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR: str | None = None

@@ -1469,6 +1470,11 @@ environment_variables: dict[str, Callable[[], Any]] = {

1470 "VLLM_USE_FUSED_MOE_GROUPED_TOPK": lambda: bool(

1471 int(os.getenv("VLLM_USE_FUSED_MOE_GROUPED_TOPK", "1"))

1472 ),

1473 + # Skip cudagraph/DP padding tokens in the MoE path by forcing their expert

1474 + # ids to -1 so the dispatch and experts drop them. Requires a MoE kernel that

1475 + # treats topk_id == -1 as a skip sentinel; off by default because not all

1476 + # kernels support it yet.

1477 + "VLLM_MOE_SKIP_PADDING": lambda: bool(int(os.getenv("VLLM_MOE_SKIP_PADDING", "0"))),

1478 # Allow use of FlashInfer FP8 block-scale GEMM for linear layers.

1479 # This uses TensorRT-LLM kernels and requires SM90+ (Hopper).

1480 "VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER": lambda: bool(

vllm/forward_context.py

@@ -147,6 +147,11 @@ class ForwardContext:

147

148 ubatch_slices: UBatchSlices | None = None

149

150 + # Boolean mask over the token axis: True for padding rows that are not real

151 + # tokens. Consumers can use it to skip work for padded tokens. None when

152 + # the producer does not set it.

153 + is_padding: torch.Tensor | None = None

154 +

155 # If True, bypass the compiled model call, e.g. by using .forward() directly

156 skip_compiled: bool = False

157

@@ -211,6 +216,7 @@ def create_forward_context(

216 slot_mapping: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None,

217 additional_kwargs: dict[str, Any] | None = None,

218 skip_compiled: bool = False,

219 + is_padding: torch.Tensor | None = None,

220 ):

221 if vllm_config.compilation_config.fast_moe_cold_start:

222 all_moe_layers = vllm_config.compilation_config.static_all_moe_layers

@@ -228,6 +234,7 @@ def create_forward_context(

234 ubatch_slices=ubatch_slices,

235 skip_compiled=skip_compiled,

236 additional_kwargs=additional_kwargs or {},

237 + is_padding=is_padding,

238 )

239

240

@@ -257,6 +264,7 @@ def set_forward_context(

264 ubatch_slices: UBatchSlices | None = None,

265 slot_mapping: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None,

266 skip_compiled: bool = False,

267 + is_padding: torch.Tensor | None = None,

268 ):

269 """A context manager that stores the current forward context,

270 can be attention metadata, etc.

@@ -316,6 +324,7 @@ def set_forward_context(

324 slot_mapping,

325 additional_kwargs,

326 skip_compiled,

327 + is_padding=is_padding,

328 )

329

330 try:

vllm/model_executor/layers/fused_moe/modular_kernel.py

+17

@@ -10,6 +10,7 @@ from typing import final

10 import torch

12 import vllm.envs as envs

13 +from vllm.forward_context import get_forward_context, is_forward_context_available

14 from vllm.logger import init_logger

15 from vllm.model_executor.layers.fused_moe.activation import (

16 MoEActivation,

@@ -1133,6 +1134,22 @@ class FusedMoEKernelModularImpl:

1134 The _prepare method is a wrapper around self.prepare_finalize.prepare

1135 that handles DBO and async.

1136 """

1137 + # Skip cudagraph/DP padding tokens uniformly across all a2a backends:

1138 + # forcing padded rows' expert ids to -1 makes every prepare_finalize drop

1139 + # them (not dispatched / not computed by the experts). The V2 model runner

1140 + # marks them in forward_context.is_padding; it is None for runners that do

1141 + # not populate it, leaving topk_ids unchanged.

1142 + # Gated by VLLM_MOE_SKIP_PADDING (off by default) because this requires the

1143 + # experts kernel to treat topk_id == -1 as a skip sentinel, which not all

1144 + # MoE backends support yet.

1145 + is_padding = None

1146 + if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available():

1147 + is_padding = get_forward_context().is_padding

1148 + if is_padding is not None:

1149 + n = topk_ids.shape[0]

1150 + # TODO: Properly support DBO (padding lives at the batch tail).

1151 + topk_ids = torch.where(is_padding[:n].unsqueeze(1), -1, topk_ids)

1152 +

1153 if not self.prepare_finalize.supports_async():

1154 # We shouldn't be running an a2a kernel that doesn't

1155 # support async prepare/finalize

vllm/models/deepseek_v4/nvidia/model.py

+10

@@ -8,6 +8,7 @@ import regex as re

8 import torch

9 import torch.nn as nn

11 +import vllm.envs as envs

12 from vllm.config import VllmConfig

13 from vllm.distributed import (

14 get_ep_group,

@@ -16,6 +17,7 @@ from vllm.distributed import (

17 get_tensor_model_parallel_world_size,

18 )

19 from vllm.distributed.eplb.eplb_state import EplbLayerState

20 +from vllm.forward_context import get_forward_context, is_forward_context_available

21 from vllm.model_executor.kernels.mhc.tilelang import (

22 hc_head_fused_kernel_tilelang,

23 mhc_fused_post_pre_tilelang,

@@ -442,6 +444,11 @@ class DeepseekV4MegaMoEExperts(nn.Module):

444

445 symm_buffer = self.get_symm_buffer()

446 num_tokens = hidden_states.shape[0]

447 + is_padding = None

448 + if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available():

449 + is_padding = get_forward_context().is_padding

450 + if is_padding is not None:

451 + is_padding = is_padding[:num_tokens]

452

453 # EPLB: map logical expert IDs to physical replicas and record load.

454 eplb_state = self.eplb_state

@@ -449,6 +456,8 @@ class DeepseekV4MegaMoEExperts(nn.Module):

456 assert eplb_state.expert_load_view is not None

457 assert eplb_state.logical_replica_count is not None

458 assert eplb_state.should_record_tensor is not None

459 + if is_padding is not None:

460 + topk_ids = torch.where(is_padding.unsqueeze(1), -1, topk_ids)

461 topk_ids = eplb_map_to_physical_and_record(

462 topk_ids=topk_ids,

463 expert_load_view=eplb_state.expert_load_view,

@@ -465,6 +474,7 @@ class DeepseekV4MegaMoEExperts(nn.Module):

474 symm_buffer.x_sf[:num_tokens],

475 symm_buffer.topk_idx[:num_tokens],

476 symm_buffer.topk_weights[:num_tokens],

477 + is_padding=is_padding,

478 )

479

480 # This method must have been already called during the weight loading phase.

vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py

+11

@@ -19,6 +19,7 @@ def _prepare_megamoe_inputs_kernel(

19 x_sf,

20 topk_ids,

21 topk_weights,

22 + is_padding,

23 topk_idx_out,

24 topk_weights_out,

25 hidden_stride_m: tl.constexpr,

@@ -31,6 +32,7 @@ def _prepare_megamoe_inputs_kernel(

32 topk_ids_stride_k: tl.constexpr,

33 topk_weights_stride_m: tl.constexpr,

34 topk_weights_stride_k: tl.constexpr,

35 + is_padding_stride_m: tl.constexpr,

36 topk_idx_stride_m: tl.constexpr,

37 topk_idx_stride_k: tl.constexpr,

38 topk_weights_out_stride_m: tl.constexpr,

@@ -85,12 +87,16 @@ def _prepare_megamoe_inputs_kernel(

87 if k_block_id == 0:

88 topk_offsets = tl.arange(0, BLOCK_TOPK)

89 topk_mask = topk_offsets < top_k

90 + token_is_padding = False

91 + if is_padding is not None:

92 + token_is_padding = tl.load(is_padding + token_id * is_padding_stride_m)

94 ids = tl.load(

95 topk_ids + token_id * topk_ids_stride_m + topk_offsets * topk_ids_stride_k,

96 mask=topk_mask,

97 other=0,

98 ).to(tl.int64)

99 + ids = tl.where(token_is_padding, -1, ids)

100 tl.store(

101 topk_idx_out

102 + token_id * topk_idx_stride_m

@@ -106,6 +112,7 @@ def _prepare_megamoe_inputs_kernel(

112 mask=topk_mask,

113 other=0.0,

114 )

115 + weights = tl.where(token_is_padding, 0.0, weights)

116 tl.store(

117 topk_weights_out

118 + token_id * topk_weights_out_stride_m

@@ -123,6 +130,7 @@ def prepare_megamoe_inputs(

130 x_sf: torch.Tensor,

131 topk_idx_out: torch.Tensor,

132 topk_weights_out: torch.Tensor,

133 + is_padding: torch.Tensor | None = None,

134 ) -> None:

135 num_tokens, hidden_size = hidden_states.shape

136 if num_tokens == 0:

@@ -142,12 +150,14 @@ def prepare_megamoe_inputs(

150 block_k = 128

151 grid = (num_tokens, triton.cdiv(hidden_size, block_k))

152 block_topk = triton.next_power_of_2(top_k)

153 + padding_stride_m = is_padding.stride(0) if is_padding is not None else 0

154 _prepare_megamoe_inputs_kernel[grid](

155 hidden_states,

156 x_fp8,

157 x_sf,

158 topk_ids,

159 topk_weights,

160 + is_padding,

161 topk_idx_out,

162 topk_weights_out,

163 hidden_states.stride(0),

@@ -160,6 +170,7 @@ def prepare_megamoe_inputs(

170 topk_ids.stride(1),

171 topk_weights.stride(0),

172 topk_weights.stride(1),

173 + padding_stride_m,

174 topk_idx_out.stride(0),

175 topk_idx_out.stride(1),

176 topk_weights_out.stride(0),

vllm/v1/worker/gpu/cudagraph_utils.py

@@ -471,6 +471,9 @@ class ModelCudaGraphManager(CudaGraphManager):

471 skip_attn=(desc.cg_mode == CUDAGraphMode.PIECEWISE),

472 )

473

474 + # Capture with dummy rows marked as padding.

475 + input_buffers.is_padding.fill_(True)

476 +

477 def forward_fn(cg_mode: CUDAGraphMode) -> None:

478 batch_descriptor = None

479 if cg_mode == CUDAGraphMode.PIECEWISE:

@@ -488,6 +491,7 @@ class ModelCudaGraphManager(CudaGraphManager):

491 num_tokens_across_dp=num_tokens_across_dp,

492 slot_mapping=slot_mappings,

493 batch_descriptor=batch_descriptor,

494 + is_padding=input_buffers.is_padding[:num_tokens],

495 ):

496 if cg_mode == CUDAGraphMode.PIECEWISE:

497 # PIECEWISE graph (compiled PW or breakable, chosen inside

vllm/v1/worker/gpu/input_batch.py

@@ -22,6 +22,7 @@ class InputBuffers:

23 self.input_ids = torch.zeros(max_num_tokens, dtype=torch.int32, device=device)

24 self.positions = torch.zeros(max_num_tokens, dtype=torch.int64, device=device)

25 + self.is_padding = torch.zeros(max_num_tokens, dtype=torch.bool, device=device)

26 self.query_start_loc = torch.zeros(

27 max_num_reqs + 1, dtype=torch.int32, device=device

28 )

@@ -83,6 +84,8 @@ class InputBatch:

84 input_ids: torch.Tensor

85 # [num_tokens_after_padding]

86 positions: torch.Tensor

87 + # [num_tokens_after_padding]

88 + is_padding: torch.Tensor

90 # [total_num_logits]

91 logits_indices: torch.Tensor

@@ -134,6 +137,9 @@ class InputBatch:

137 input_ids = input_buffers.input_ids[:num_tokens].zero_()

138 positions = input_buffers.positions[:num_tokens].zero_()

139

140 + input_buffers.is_padding[:num_tokens].fill_(True)

141 + is_padding = input_buffers.is_padding[:num_tokens]

142 +

143 logits_indices = query_start_loc[1:] - 1

144 cu_num_logits = torch.arange(num_reqs + 1, device=device, dtype=torch.int32)

145 cu_num_logits_np = np.arange(num_reqs + 1, dtype=np.int32)

@@ -164,6 +170,7 @@ class InputBatch:

170 max_seq_len_np=None,

171 input_ids=input_ids,

172 positions=positions,

173 + is_padding=is_padding,

174 logits_indices=logits_indices,

175 cu_num_logits=cu_num_logits,

176 cu_num_logits_np=cu_num_logits_np,

vllm/v1/worker/gpu/model_runner.py

+10

@@ -27,6 +27,7 @@ import numpy as np

27 import torch

28 import torch.nn as nn

30 +import vllm.envs as envs

31 from vllm.compilation.counter import compilation_counter

32 from vllm.config import VllmConfig

33 from vllm.config.compilation import CUDAGraphMode

@@ -847,6 +848,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):

848 num_tokens = scheduler_output.total_num_scheduled_tokens

849 num_tokens_after_padding = batch_desc.num_tokens

850 assert num_tokens > 0

851 + if envs.VLLM_MOE_SKIP_PADDING:

852 + # Mark trailing cudagraph-padding rows so kernels can skip work for

853 + # them when supported.

854 + self.input_buffers.is_padding[:num_tokens].fill_(False)

855 + self.input_buffers.is_padding[num_tokens:num_tokens_after_padding].fill_(

856 + True

857 + )

858 num_tokens_per_req = scheduler_output.num_scheduled_tokens

859 num_reqs = len(num_tokens_per_req)

860

@@ -1001,6 +1009,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1009 max_seq_len_np=max_seq_len_np,

1010 input_ids=self.input_buffers.input_ids[:num_tokens_after_padding],

1011 positions=self.input_buffers.positions[:num_tokens_after_padding],

1012 + is_padding=self.input_buffers.is_padding[:num_tokens_after_padding],

1013 logits_indices=logits_indices,

1014 cu_num_logits=cu_num_logits,

1015 cu_num_logits_np=cu_num_logits_np,

@@ -1277,6 +1286,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1286 batch_descriptor=batch_descriptor,

1287 slot_mapping=slot_mappings_by_layer,

1288 skip_compiled=skip_compiled,

1289 + is_padding=input_batch.is_padding,

1290 ):

1291 self.kv_connector.pre_forward(scheduler_output)

1292 if batch_desc.cg_mode == CUDAGraphMode.PIECEWISE: