@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e1e318af01

Fork 0

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 5, 04:45 UTC

[MoE Refactor] Remove MoE DP chunking (#39107)

Signed-off-by: Bill Nell <bnell@redhat.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>

bnellnm committed Apr 14, 2026 at 09:48 UTC e1e318af010b4f92d39324d03231cfd409766bf9

20 files changed +76 -389

.buildkite/test_areas/kernels.yaml

+16 -2

@@ -200,7 +200,14 @@ steps:

200 timeout_in_minutes: 90

201 device: h100

202 num_devices: 2

203 - optional: true

203 + source_file_dependencies:

204 + - csrc/quantization/cutlass_w8a8/moe/

205 + - csrc/moe/

206 + - tests/kernels/moe

207 + - vllm/model_executor/layers/fused_moe/

208 + - vllm/model_executor/layers/quantization/

209 + - vllm/distributed/device_communicators/

210 + - vllm/config

211 commands:

212 - pytest -v -s kernels/moe/test_moe_layer.py

213

@@ -209,6 +216,13 @@ steps:

216 timeout_in_minutes: 90

217 device: b200

218 num_devices: 2

212 - optional: true

219 + source_file_dependencies:

220 + - csrc/quantization/cutlass_w8a8/moe/

221 + - csrc/moe/

222 + - tests/kernels/moe

223 + - vllm/model_executor/layers/fused_moe/

224 + - vllm/model_executor/layers/quantization/

225 + - vllm/distributed/device_communicators/

226 + - vllm/config

227 commands:

228 - pytest -v -s kernels/moe/test_moe_layer.py

tests/kernels/moe/modular_kernel_tools/common.py

+1 -7

@@ -46,6 +46,7 @@ from vllm.utils.import_utils import (

46 has_deep_gemm,

47 has_mori,

48 )

49 +from vllm.utils.math_utils import next_power_of_2

51 from .mk_objects import (

52 TestMoEQuantConfig,

@@ -604,13 +605,6 @@ def make_modular_kernel(

605 vllm_config: VllmConfig,

606 quant_config: FusedMoEQuantConfig,

607 ) -> mk.FusedMoEKernel:

607 - def next_power_of_2(x):

608 - import math

609 -

610 - if x == 0:

611 - return 1

612 - return 2 ** math.ceil(math.log2(x))

613 -

608 # make moe config

609 moe_parallel_config: FusedMoEParallelConfig = FusedMoEParallelConfig.make(

610 tp_size_=get_tensor_model_parallel_world_size(),

tests/kernels/moe/modular_kernel_tools/parallel_utils.py

+1 -1

@@ -126,7 +126,7 @@ def parallel_launch_with_config(

126 world_size: int,

127 worker: Callable[Concatenate[ProcessGroupInfo, VllmConfig, Any, P], None],

128 vllm_config: VllmConfig,

129 - env_dict: dict[Any, Any],

129 + env_dict: dict[Any, Any] | None,

130 *args: P.args,

131 **kwargs: P.kwargs,

132 ) -> None:

tests/kernels/moe/test_deepep_deepgemm_moe.py

+1 -8

@@ -29,6 +29,7 @@ from vllm.utils.deep_gemm import (

29 is_deep_gemm_supported,

30 )

31 from vllm.utils.import_utils import has_deep_ep, has_deep_gemm

32 +from vllm.utils.math_utils import next_power_of_2

33 from vllm.utils.torch_utils import set_random_seed

34 from vllm.v1.worker.workspace import init_workspace_manager

@@ -84,14 +85,6 @@ def with_dp_metadata(M: int, world_size: int):

85 yield

87 -def next_power_of_2(x):

88 - import math

89 -

90 - if x == 0:

91 - return 1

92 - return 2 ** math.ceil(math.log2(x))

93 -

94 -

88 def make_block_quant_fp8_weights(

89 e: int,

90 n: int,

tests/kernels/moe/test_flashinfer.py

@@ -32,6 +32,7 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

32 from vllm.model_executor.layers.quantization.utils.fp8_utils import input_to_float8

33 from vllm.model_executor.models.llama4 import Llama4MoE

34 from vllm.platforms import current_platform

35 +from vllm.utils.math_utils import next_power_of_2

36 from vllm.utils.torch_utils import set_random_seed

38 try:

@@ -174,6 +175,7 @@ class TestData:

175 routing_method=layer.routing_method_type,

176 activation=activation,

177 device=w13_quantized.device,

178 + max_num_tokens=next_power_of_2(m),

179 )

180

181 return TestData(

@@ -348,6 +350,7 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(

350 in_dtype=torch.bfloat16,

351 is_act_and_mul=activation.is_gated,

352 routing_method=RoutingMethodType.TopK,

353 + max_num_tokens=next_power_of_2(m),

354 )

355

356 kernel = mk.FusedMoEKernel(

tests/kernels/moe/test_flashinfer_moe.py

@@ -29,6 +29,7 @@ from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (

29 from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel

30 from vllm.platforms import current_platform

31 from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe

32 +from vllm.utils.math_utils import next_power_of_2

33 from vllm.utils.torch_utils import set_random_seed

35 if not has_flashinfer_cutlass_fused_moe() or not current_platform.has_device_capability(

@@ -105,6 +106,7 @@ def test_flashinfer_fp4_moe_no_graph(

106 in_dtype=dtype,

107 is_act_and_mul=is_gated_act,

108 routing_method=RoutingMethodType.TopK,

109 + max_num_tokens=next_power_of_2(m),

110 )

111

112 flashinfer_experts = FusedMoEKernel(

tests/kernels/moe/test_moe.py

+2 -1

@@ -59,6 +59,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import quantize_w

59 from vllm.model_executor.models.mixtral import MixtralMoE

60 from vllm.platforms import current_platform

61 from vllm.scalar_type import ScalarType, scalar_types

62 +from vllm.utils.math_utils import next_power_of_2

63 from vllm.utils.torch_utils import set_random_seed

64 from vllm.v1.worker.workspace import init_workspace_manager

@@ -1676,7 +1677,7 @@ def test_unquantized_bf16_flashinfer_trtllm_backend(

1677 in_dtype=dtype,

1678 is_act_and_mul=True,

1679 routing_method=RoutingMethodType.Renormalize,

1679 - max_num_tokens=m,

1680 + max_num_tokens=next_power_of_2(m),

1681 )

1682

1683 with set_current_vllm_config(vllm_config):

tests/kernels/moe/test_moe_layer.py

+13 -10

@@ -26,6 +26,7 @@ from tests.kernels.moe.utils import TestMLP, make_test_weights, moe_quantize_wei

26 from vllm.config import (

27 CompilationConfig,

28 ParallelConfig,

29 + SchedulerConfig,

30 VllmConfig,

31 set_current_vllm_config,

32 )

@@ -53,7 +54,7 @@ from vllm.utils.flashinfer import (

54 has_flashinfer_nvlink_two_sided,

55 )

56 from vllm.utils.import_utils import has_deep_ep, has_mori, has_nixl_ep

56 -from vllm.utils.math_utils import cdiv

57 +from vllm.utils.math_utils import cdiv, next_power_of_2

58 from vllm.utils.torch_utils import set_random_seed

59 from vllm.v1.worker.workspace import (

60 init_workspace_manager,

@@ -65,8 +66,9 @@ fp8_dtype = torch.float8_e4m3fn # current_platform.fp8_dtype

66 SHAPE_COMBOS = [

67 (1, 128, 256),

68 (32, 1024, 512),

68 - (222, 2048, 2048), # should be big enough to exercise DP chunking

69 + (222, 2048, 2048),

70 ]

71 +MAX_M = max([x[0] for x in SHAPE_COMBOS])

73 NUM_EXPERTS = [8, 64]

74 TOP_KS = [2, 6]

@@ -112,7 +114,7 @@ BACKEND_SUPPORTED_QUANTS: dict[str, set[str | None]] = {

114 "mori": {None, "fp8", "modelopt_fp8"},

115 "flashinfer_nvlink_two_sided": {None, "modelopt_fp8", "modelopt_fp4"},

116 "flashinfer_nvlink_one_sided": {None, "modelopt_fp8", "modelopt_fp4"},

115 - "deepep_low_latency": {None, "fp8", "modelopt_fp8", "modelopt_fp4"},

117 + "deepep_low_latency": {None, "modelopt_fp8", "modelopt_fp4"},

118 "deepep_high_throughput": {None, "fp8", "modelopt_fp8", "modelopt_fp4"},

119 "nixl_ep": {None, "fp8", "modelopt_fp8"},

120 }

@@ -363,9 +365,9 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:

365 )

366

367 # routed_input_transform + quantization + high hidden dimensions

366 - # TODO: Disable >= 2048 w/fp8 + deepep LL for now due to insane errors.

368 + # TODO: Disable >= 2048 for now due to insane errors.

369 if (

368 - (config.use_routed_input_transform or config.backend == "deepep_low_latency")

370 + config.use_routed_input_transform

371 and config.quantization is not None

372 and config.k >= 2048

373 ):

@@ -1663,9 +1665,6 @@ def test_moe_layer(

1665

1666 verbosity = pytestconfig.getoption("verbose")

1667

1666 - test_env = dict()

1667 - test_env["VLLM_MOE_DP_CHUNK_SIZE"] = "128"

1668 - monkeypatch.setenv("VLLM_MOE_DP_CHUNK_SIZE", "128")

1668 if os.environ.get("VLLM_LOGGING_LEVEL") is None:

1669 monkeypatch.setenv("VLLM_LOGGING_LEVEL", "ERROR")

1670

@@ -1690,7 +1689,11 @@ def test_moe_layer(

1689 compilation_config.pass_config.fuse_allreduce_rms = False # for now

1690

1691 vllm_config = VllmConfig(

1693 - parallel_config=parallel_config, compilation_config=compilation_config

1692 + parallel_config=parallel_config,

1693 + compilation_config=compilation_config,

1694 + scheduler_config=SchedulerConfig.default_factory(

1695 + max_num_batched_tokens=next_power_of_2(MAX_M)

1696 + ),

1697 )

1698

1699 test_configs = generate_valid_test_configs(

@@ -1718,7 +1721,7 @@ def test_moe_layer(

1721 world_size,

1722 _parallel_worker,

1723 vllm_config,

1721 - test_env,

1724 + None,

1725 test_configs,

1726 verbosity,

1727 )

tests/kernels/moe/utils.py

@@ -69,6 +69,7 @@ def make_dummy_moe_config(

69 in_dtype=in_dtype,

70 device="cuda",

71 routing_method=RoutingMethodType.TopK,

72 + max_num_tokens=512,

73 )

vllm/config/parallel.py

+12

@@ -622,6 +622,18 @@ class ParallelConfig:

622 and self.data_parallel_size > 1

623 )

624

625 + @property

626 + def use_batched_dp_moe(self) -> bool:

627 + return (

628 + self.all2all_backend

629 + in (

630 + "deepep_low_latency",

631 + "nixl_ep",

632 + )

633 + and self.enable_expert_parallel

634 + and self.data_parallel_size > 1

635 + )

636 +

637 @property

638 def node_rank_within_dp(self) -> int:

639 return self.node_rank % self.nnodes_within_dp

vllm/config/scheduler.py

@@ -40,6 +40,7 @@ class SchedulerConfig:

40 """

42 DEFAULT_MAX_NUM_BATCHED_TOKENS: ClassVar[int] = 2048

43 + DEFAULT_MAX_NUM_BATCHED_TOKENS_FOR_BATCHED_DP: ClassVar[int] = 256

44 DEFAULT_MAX_NUM_SEQS: ClassVar[int] = 128

46 runner_type: RunnerType = "generate"

vllm/engine/arg_utils.py

+16 -7

@@ -1588,9 +1588,6 @@ class EngineArgs:

1588

1589 self._check_feature_supported()

1590 self._set_default_chunked_prefill_and_prefix_caching_args(model_config)

1591 - self._set_default_max_num_seqs_and_batched_tokens_args(

1592 - usage_context, model_config

1593 - )

1591 self._set_default_reasoning_config_args()

1592 sliding_window: int | None = None

1593 if not is_interleaved(model_config.hf_text_config):

@@ -1846,6 +1843,12 @@ class EngineArgs:

1843 target_parallel_config=parallel_config,

1844 )

1845

1846 + self._set_default_max_num_seqs_and_batched_tokens_args(

1847 + usage_context,

1848 + model_config,

1849 + parallel_config,

1850 + )

1851 +

1852 assert self.max_num_batched_tokens is not None, (

1853 "max_num_batched_tokens must be set by this point"

1854 )

@@ -2244,6 +2247,7 @@ class EngineArgs:

2247 self,

2248 usage_context: UsageContext | None,

2249 model_config: ModelConfig,

2250 + parallel_config: ParallelConfig,

2251 ):

2252 world_size = self.pipeline_parallel_size * self.tensor_parallel_size

2253 (

@@ -2255,10 +2259,15 @@ class EngineArgs:

2259 orig_max_num_seqs = self.max_num_seqs

2260

2261 if self.max_num_batched_tokens is None:

2258 - self.max_num_batched_tokens = default_max_num_batched_tokens.get(

2259 - usage_context,

2260 - SchedulerConfig.DEFAULT_MAX_NUM_BATCHED_TOKENS,

2261 - )

2262 + if parallel_config.use_batched_dp_moe:

2263 + self.max_num_batched_tokens = (

2264 + SchedulerConfig.DEFAULT_MAX_NUM_BATCHED_TOKENS_FOR_BATCHED_DP

2265 + )

2266 + else:

2267 + self.max_num_batched_tokens = default_max_num_batched_tokens.get(

2268 + usage_context,

2269 + SchedulerConfig.DEFAULT_MAX_NUM_BATCHED_TOKENS,

2270 + )

2271

2272 if self.max_num_seqs is None:

2273 self.max_num_seqs = default_max_num_seqs.get(

vllm/envs.py

-11

@@ -146,8 +146,6 @@ if TYPE_CHECKING:

146 VLLM_ENABLE_PREGRAD_PASSES: bool = False

147 VLLM_DP_MASTER_IP: str = ""

148 VLLM_DP_MASTER_PORT: int = 0

149 - VLLM_MOE_DP_CHUNK_SIZE: int = 256

150 - VLLM_ENABLE_MOE_DP_CHUNK: bool = True

149 VLLM_RANDOMIZE_DP_DUMMY_INPUTS: bool = False

150 VLLM_RAY_DP_PACK_STRATEGY: Literal["strict", "fill", "span"] = "strict"

151 VLLM_RAY_EXTRA_ENV_VAR_PREFIXES_TO_COPY: str = ""

@@ -1140,15 +1138,6 @@ environment_variables: dict[str, Callable[[], Any]] = {

1138 "VLLM_DP_MASTER_IP": lambda: os.getenv("VLLM_DP_MASTER_IP", "127.0.0.1"),

1139 # Port of the master node in the data parallel setting

1140 "VLLM_DP_MASTER_PORT": lambda: int(os.getenv("VLLM_DP_MASTER_PORT", "0")),

1143 - # In the context of executing MoE models with Data-Parallel, Expert-Parallel

1144 - # and Batched All-to-All dispatch/combine kernels, VLLM_MOE_DP_CHUNK_SIZE

1145 - # dictates the quantum of tokens that can be dispatched from a DP

1146 - # rank. All DP ranks process the activations in VLLM_MOE_DP_CHUNK_SIZE

1147 - # units.

1148 - "VLLM_MOE_DP_CHUNK_SIZE": lambda: int(os.getenv("VLLM_MOE_DP_CHUNK_SIZE", "256")),

1149 - "VLLM_ENABLE_MOE_DP_CHUNK": lambda: bool(

1150 - int(os.getenv("VLLM_ENABLE_MOE_DP_CHUNK", "1"))

1151 - ),

1141 # Randomize inputs during dummy runs when using Data Parallel

1142 "VLLM_RANDOMIZE_DP_DUMMY_INPUTS": lambda: (

1143 os.environ.get("VLLM_RANDOMIZE_DP_DUMMY_INPUTS", "0") == "1"

vllm/forward_context.py

+1 -60

@@ -70,27 +70,8 @@ def _compute_sp_num_tokens(

70 return sp_tokens.tolist()

73 -def _compute_chunked_local_num_tokens(

74 - num_tokens_across_dp_cpu: torch.Tensor,

75 - sequence_parallel_size: int,

76 - max_num_tokens: int,

77 - chunk_idx: int,

78 -) -> list[int]:

79 - sp_tokens = _compute_sp_num_tokens(num_tokens_across_dp_cpu, sequence_parallel_size)

80 - sp_size = len(sp_tokens)

81 -

82 - local_size = [-1] * sp_size

83 - for i in range(sp_size):

84 - # Take into account sharding if MoE activation is sequence parallel.

85 - local_size[i] = min(max_num_tokens, sp_tokens[i] - (max_num_tokens * chunk_idx))

86 - if local_size[i] <= 0:

87 - local_size[i] = 1 # ensure lockstep even if done

88 - return local_size

89 -

90 -

73 @dataclass

74 class DPMetadata:

93 - max_tokens_across_dp_cpu: torch.Tensor

75 num_tokens_across_dp_cpu: torch.Tensor

77 # NOTE: local_sizes should only be set by the chunked_sizes context manager

@@ -113,47 +94,7 @@ class DPMetadata:

94 assert num_tokens_across_dp_cpu[dp_rank] == batchsize, (

95 f"{num_tokens_across_dp_cpu[dp_rank]} {batchsize}"

96 )

116 - max_tokens_across_dp_cpu = torch.max(num_tokens_across_dp_cpu)

117 - return DPMetadata(max_tokens_across_dp_cpu, num_tokens_across_dp_cpu)

118 -

119 - @contextmanager

120 - def chunked_sizes(

121 - self, sequence_parallel_size: int, max_chunk_size_per_rank: int, chunk_idx: int

122 - ):

123 - """

124 - Context manager to compute and temporarily set the per-rank local token

125 - sizes for a specific chunk during chunked forward execution.

126 -

127 - This is necessary to ensure each DP (data parallel) rank processes its

128 - designated portion of tokens in lockstep with others, even when the

129 - token counts are uneven or some ranks have completed their input early.

130 -

131 - For chunked execution, we break up the total tokens on each rank into

132 - multiple chunks (of at most `max_chunk_size_per_rank`), and for a given

133 - `chunk_idx`, this context manager sets `self.local_sizes` to the number

134 - of tokens to process in that chunk on each rank.

135 -

136 - `self.local_sizes` is only valid inside the context.

137 -

138 - Args:

139 - sequence_parallel_size: When Attn is TP and MoE layers are EP,

140 - we use SP between the layers to avoid

141 - redundant ops. We need this value to

142 - compute the chunked sizes.

143 - max_chunk_size_per_rank: The max number of tokens each rank is

144 - allowed to process in this chunk.

145 - chunk_idx: The index of the chunk to compute sizes for.

146 - """

147 - self.local_sizes = _compute_chunked_local_num_tokens(

148 - self.num_tokens_across_dp_cpu,

149 - sequence_parallel_size,

150 - max_chunk_size_per_rank,

151 - chunk_idx,

152 - )

153 - try:

154 - yield self.local_sizes

155 - finally:

156 - self.local_sizes = None

97 + return DPMetadata(num_tokens_across_dp_cpu)

99 @contextmanager

100 def sp_local_sizes(self, sequence_parallel_size: int):

vllm/model_executor/layers/fused_moe/config.py

+2 -12

@@ -6,8 +6,7 @@ from typing import Union

7 import torch

9 -import vllm.envs as envs

10 -from vllm.config import ParallelConfig

9 +from vllm.config import ParallelConfig, SchedulerConfig

10 from vllm.distributed import get_dp_group, get_pcp_group, get_tensor_model_parallel_rank

11 from vllm.logger import init_logger

12 from vllm.model_executor.layers.fused_moe.activation import MoEActivation

@@ -937,15 +936,6 @@ class FusedMoEParallelConfig:

936 all2all_backend: str # all2all backend for MoE communication

937 enable_eplb: bool # whether to enable expert load balancing

938

940 - @property

941 - def use_dp_chunking(self) -> bool:

942 - return (

943 - self.use_deepep_ll_kernels

944 - or self.use_mori_kernels

945 - or self.use_fi_nvl_two_sided_kernels

946 - or self.use_nixl_ep_kernels

947 - ) and envs.VLLM_ENABLE_MOE_DP_CHUNK

948 -

939 @property

940 def is_sequence_parallel(self) -> bool:

941 return self.sp_size > 1

@@ -1184,7 +1174,7 @@ class FusedMoEConfig:

1174 intermediate_size_per_partition_unpadded: int | None = None

1175

1176 moe_backend: str = "auto"

1187 - max_num_tokens: int = envs.VLLM_MOE_DP_CHUNK_SIZE

1177 + max_num_tokens: int = SchedulerConfig.DEFAULT_MAX_NUM_BATCHED_TOKENS_FOR_BATCHED_DP

1178 has_bias: bool = False

1179 is_act_and_mul: bool = True

1180 is_lora_enabled: bool = False

vllm/model_executor/layers/fused_moe/layer.py

+1 -2

@@ -8,7 +8,6 @@ from typing import Literal, cast, get_args, overload

8 import torch

9 from torch.nn.parameter import UninitializedParameter

11 -import vllm.envs as envs

11 from vllm._aiter_ops import rocm_aiter_ops

12 from vllm.config import VllmConfig, get_current_vllm_config

13 from vllm.config.parallel import ExpertPlacementStrategy

@@ -479,7 +478,7 @@ class FusedMoE(CustomOp):

478 in_dtype=moe_in_dtype,

479 moe_backend=vllm_config.kernel_config.moe_backend,

480 router_logits_dtype=router_logits_dtype,

482 - max_num_tokens=envs.VLLM_MOE_DP_CHUNK_SIZE,

481 + max_num_tokens=vllm_config.scheduler_config.max_num_batched_tokens,

482 has_bias=has_bias,

483 is_act_and_mul=is_act_and_mul,

484 is_lora_enabled=vllm_config.lora_config is not None,

vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py deleted

-243

@@ -1,243 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -

4 -import torch

5 -

6 -from vllm.forward_context import (

7 - get_forward_context,

8 -)

9 -from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (

10 - FusedMoEMethodBase,

11 -)

12 -from vllm.model_executor.layers.fused_moe.runner.moe_runner_base import MoERunnerBase

13 -from vllm.model_executor.layers.fused_moe.runner.shared_experts import (

14 - SharedExperts,

15 -)

16 -from vllm.utils.math_utils import cdiv

17 -from vllm.v1.worker.ubatching import dbo_current_ubatch_id

18 -from vllm.v1.worker.workspace import current_workspace_manager

19 -

20 -

21 -class ChunkingMoERunner(MoERunnerBase):

22 - """

23 - MoE runner wrapper that adds chunked processing to any MoERunnerBase.

24 -

25 - This runner wraps an inner MoERunnerBase and overrides _forward_impl to

26 - process large batches by breaking them into smaller chunks. Each chunk

27 - is delegated to the inner runner's _forward_impl, making chunking

28 - composable with any runner implementation.

29 -

30 - All MoERunnerBase state (moe_config, router, quant_method, etc.) is

31 - transparently delegated to the inner runner via __getattr__.

32 - ChunkingMoERunner only owns chunking-specific state: the pre-allocated

33 - workspace buffers and the reduce_results override.

34 -

35 - Key behaviors:

36 - - Pre-allocates workspace tensors for CUDA graph compatibility

37 - - Processes chunks via inner._forward_impl per chunk

38 - - Never reduces results (reduce_results always returns False)

39 - """

40 -

41 - def __init__(self, inner: MoERunnerBase):

42 - # Assert that _maybe_dispatch/_maybe_combine will be nops.

43 - assert inner.moe_config.pcp_size == 1

44 -

45 - # Skip MoERunnerBase.__init__ — all state is delegated to inner

46 - # via __getattr__. Only chunking-specific state lives here.

47 - self._inner = inner

48 -

49 - # Pre-allocated staging buffers. These need to exist ahead of time

50 - # due to CUDA graph construction needing fixed buffer addresses.

51 - self.batched_hidden_states, self.batched_router_logits = (

52 - self._init_dp_chunking()

53 - )

54 -

55 - def __getattr__(self, name):

56 - # Delegate attribute access to the inner runner. This is only

57 - # called when normal lookup (instance __dict__, class MRO) fails,

58 - # so ChunkingMoERunner's own attributes and methods take priority.

59 - return getattr(self._inner, name)

60 -

61 - @property

62 - def shared_experts(self) -> SharedExperts | None:

63 - return self._inner.shared_experts

64 -

65 - # TODO(bnell): temporary hack, do not call this method.

66 - def _replace_quant_method(self, quant_method: FusedMoEMethodBase):

67 - self._inner._replace_quant_method(quant_method)

68 - self.quant_method = quant_method

69 -

70 - def is_internal_router(self) -> bool:

71 - return self._inner.gate is not None

72 -

73 - # Reducing results when chunking is handled by the MK finalize operations

74 - # when DP chunking is enabled..

75 - # This will be removed by #35949

76 - @property

77 - def reduce_results(self) -> bool:

78 - return False

79 -

80 - def _init_dp_chunking(self) -> list[torch.Tensor]:

81 - states_shape: tuple[int, ...]

82 - logits_shape: tuple[int, ...]

83 -

84 - moe = self.moe_config

85 -

86 - if self.enable_dbo:

87 - states_shape = (2, moe.max_num_tokens, self.moe_config.hidden_dim)

88 - logits_shape = (2, moe.max_num_tokens, self.moe_config.num_logical_experts)

89 - else:

90 - states_shape = (moe.max_num_tokens, self.moe_config.hidden_dim)

91 - logits_shape = (moe.max_num_tokens, self.moe_config.num_logical_experts)

92 -

93 - # Does this need some kind of profiling run check like modular_kernel.py?

94 - return current_workspace_manager().get_simultaneous(

95 - (states_shape, moe.in_dtype),

96 - (logits_shape, moe.router_logits_dtype),

97 - )

98 -

99 - def _allocate_dp_chunking_outputs(

100 - self,

101 - hidden_states: torch.Tensor,

102 - router_logits: torch.Tensor,

103 - shared_experts_input: torch.Tensor | None,

104 - ) -> tuple[torch.Tensor | None, torch.Tensor]:

105 - # Assert the inputs are of the proper type and shape.

106 - assert self.batched_hidden_states is not None

107 - assert self.batched_router_logits is not None

108 -

109 - assert self.batched_hidden_states.dtype == hidden_states.dtype, (

110 - f"{self.batched_hidden_states.dtype} == {hidden_states.dtype}"

111 - )

112 - assert self.batched_router_logits.dtype == router_logits.dtype, (

113 - f"{self.batched_router_logits.dtype} == {router_logits.dtype}"

114 - )

115 -

116 - # Check size compatibility.

117 - assert self.batched_hidden_states.size(-1) == hidden_states.size(-1)

118 - assert self.batched_router_logits.size(-1) == router_logits.size(-1)

119 -

120 - final_fused_hidden_states = torch.empty_like(hidden_states)

121 - if self.shared_experts is not None:

122 - if shared_experts_input is not None:

123 - final_shared_hidden_states = torch.empty_like(shared_experts_input)

124 - else:

125 - final_shared_hidden_states = torch.empty_like(hidden_states)

126 - else:

127 - final_shared_hidden_states = None

128 -

129 - return final_shared_hidden_states, final_fused_hidden_states

130 -

131 - def _slice_and_copy_input(

132 - self,

133 - out_slice: torch.Tensor,

134 - orig: torch.Tensor | None,

135 - start: int,

136 - end: int,

137 - ) -> torch.Tensor:

138 - assert orig is not None

139 - slice_size = end - start

140 - orig_slice = orig[start:end, :]

141 - if self.enable_dbo:

142 - assert out_slice.dim() == 3

143 - batch_buffer_idx = dbo_current_ubatch_id()

144 - out_slice = out_slice[batch_buffer_idx, :]

145 -

146 - assert out_slice.size(0) >= slice_size

147 - out_slice = out_slice[:slice_size, :]

148 - out_slice.copy_(orig_slice, non_blocking=True)

149 - return out_slice

150 -

151 - def _forward_impl(

152 - self,

153 - layer: torch.nn.Module,

154 - hidden_states: torch.Tensor,

155 - router_logits: torch.Tensor,

156 - shared_experts_input: torch.Tensor | None,

157 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

158 - final_shared_hidden_states, final_fused_hidden_states = (

159 - self._allocate_dp_chunking_outputs(

160 - hidden_states, router_logits, shared_experts_input

161 - )

162 - )

163 -

164 - ctx = get_forward_context()

165 - # flashinfer_cutlass_kernels can handle: optional DP + TP/EP

166 - max_tokens_across_dispatchers = ctx.dp_metadata.max_tokens_across_dp_cpu

167 - moe_dp_chunk_size_per_rank = self.moe_config.max_num_tokens

168 -

169 - # If the input to the MoE is sequence parallel then divide by sp_size

170 - # to find the maximum number of tokens for any individual dispatcher.

171 - if self.moe_config.is_sequence_parallel:

172 - max_tokens_across_dispatchers = cdiv(

173 - max_tokens_across_dispatchers, self.moe_config.sp_size

174 - )

175 -

176 - num_tokens = hidden_states.size(0)

177 - for chunk_idx, chunk_start_ in enumerate(

178 - range(0, max_tokens_across_dispatchers, moe_dp_chunk_size_per_rank)

179 - ):

180 - chunk_start = chunk_start_

181 - chunk_end = min(

182 - chunk_start + moe_dp_chunk_size_per_rank, max_tokens_across_dispatchers

183 - )

184 - # clamp start and end

185 - chunk_start = min(chunk_start, num_tokens - 1)

186 - chunk_end = min(chunk_end, num_tokens)

187 - chunk_sizes = ctx.dp_metadata.chunked_sizes(

188 - self.moe_config.sp_size, moe_dp_chunk_size_per_rank, chunk_idx

189 - )

190 - with chunk_sizes:

191 - hidden_states_chunk = self._slice_and_copy_input(

192 - self.batched_hidden_states,

193 - hidden_states,

194 - chunk_start,

195 - chunk_end,

196 - )

197 -

198 - router_logits_chunk = self._slice_and_copy_input(

199 - self.batched_router_logits,

200 - router_logits,

201 - chunk_start,

202 - chunk_end,

203 - )

204 -

205 - shared_experts_input_chunk = (

206 - shared_experts_input[chunk_start:chunk_end, :]

207 - if shared_experts_input is not None

208 - else None

209 - )

210 -

211 - # Delegate per-chunk computation to the inner runner.

212 - chunk_result = self._inner._forward_impl(

213 - layer=layer,

214 - hidden_states=hidden_states_chunk,

215 - router_logits=router_logits_chunk,

216 - shared_experts_input=shared_experts_input_chunk,

217 - )

218 -

219 - # Store outputs

220 - # TODO(bnell): document when chunk_start >= num_tokens

221 - if chunk_start < num_tokens:

222 - if self.shared_experts is not None:

223 - assert isinstance(chunk_result, tuple)

224 - shared_output_chunk, hidden_states_chunk = chunk_result

225 - final_fused_hidden_states[chunk_start:chunk_end, :].copy_(

226 - hidden_states_chunk, non_blocking=True

227 - )

228 - assert shared_output_chunk is not None

229 - assert final_shared_hidden_states is not None

230 - final_shared_hidden_states[chunk_start:chunk_end, :].copy_(

231 - shared_output_chunk, non_blocking=True

232 - )

233 - else:

234 - assert isinstance(chunk_result, torch.Tensor)

235 - final_fused_hidden_states[chunk_start:chunk_end, :].copy_(

236 - chunk_result, non_blocking=True

237 - )

238 -

239 - if self.shared_experts is None:

240 - return final_fused_hidden_states

241 - else:

242 - assert final_shared_hidden_states is not None

243 - return (final_shared_hidden_states, final_fused_hidden_states)

vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py

+1 -7

@@ -12,9 +12,6 @@ from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (

12 from vllm.model_executor.layers.fused_moe.router.fused_moe_router import (

13 FusedMoERouter,

14 )

15 -from vllm.model_executor.layers.fused_moe.runner.chunking_moe_runner import (

16 - ChunkingMoERunner,

17 -)

15 from vllm.model_executor.layers.fused_moe.runner.default_moe_runner import (

16 DefaultMoERunner,

17 )

@@ -35,7 +32,7 @@ def create_moe_runner(

32 reduce_results: bool,

33 enable_dbo: bool,

34 ) -> MoERunner:

38 - runner = DefaultMoERunner(

35 + return DefaultMoERunner(

36 layer_name,

37 moe_config,

38 router,

@@ -46,6 +43,3 @@ def create_moe_runner(

43 reduce_results,

44 enable_dbo,

45 )

49 - if moe_config.moe_parallel_config.use_dp_chunking:

50 - return ChunkingMoERunner(runner)

51 - return runner

vllm/model_executor/layers/fused_moe/runner/shared_experts.py

-16

@@ -69,7 +69,6 @@ class SharedExperts:

69 self._moe_config = moe_config

70 self._quant_method = quant_method

71 self._reduce_results = reduce_results

72 - self._use_dp_chunking = moe_config.moe_parallel_config.use_dp_chunking

73 # Allow disabling of the separate shared experts stream for

74 # debug purposes.

@@ -87,20 +86,6 @@ class SharedExperts:

86 "Enabled separate cuda stream for MoE shared_experts", scope="local"

87 )

90 - @property

91 - def _use_external_experts(self) -> bool:

92 - if self._use_dp_chunking:

93 - return False

94 -

95 - # Disable shared expert overlap if:

96 - # - we are using eplb with non-default backend, because of correctness issues

97 - # - we are using flashinfer with DP, since there nothing to gain

98 - backend = self._moe_config.moe_parallel_config.all2all_backend

99 - return (

100 - self._moe_config.moe_parallel_config.enable_eplb

101 - and backend != "allgather_reducescatter"

102 - ) or self._moe_config.moe_parallel_config.use_fi_nvl_two_sided_kernels

103 -

89 def _determine_shared_experts_order(

90 self,

91 hidden_states: torch.Tensor,

@@ -110,7 +95,6 @@ class SharedExperts:

96 should_run_shared_in_aux_stream = (

97 current_platform.is_cuda()

113 - and not self._use_dp_chunking

98 and self._stream is not None

99 and hidden_states.shape[0]

100 <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD

vllm/model_executor/layers/quantization/quark/quark_moe.py

+2 -2

@@ -1502,9 +1502,9 @@ class QuarkOCP_MX_MoEMethod_OSS(QuarkOCP_MX_MoEMethod):

1502 layer.w2_bias = torch.nn.Parameter(w2_bias, requires_grad=False)

1503

1504 # FIXME warp need to be adjusted based on batch size

1505 - # only apply to batched mode

1505 + # only apply to batched mode

1506 if self.moe.use_ep:

1507 - num_warps = 4 if envs.VLLM_MOE_DP_CHUNK_SIZE <= 512 else 8

1507 + num_warps = 4 if self.moe.max_num_tokens <= 512 else 8

1508 else:

1509 num_warps = 8

1510