[MoE Refactor] Remove MoE legacy code (#51078)
Signed-off-by: Bill Nell <bnell@redhat.com>
bnellnm committed
Aug 5, 2026 at 17:35 UTC
373fe8b83ee772d27b1a5b2b908618a1a414bb11
22 files changed
+3
-295
tests/v1/worker/test_gpu_model_runner_v2_eplb.py
-6
@@ -93,7 +93,6 @@ def _make_runner(**overrides: Any) -> Any:
93
def test_v2_load_model_registers_moe_with_eplb(monkeypatch):
94
FakeEplbState.instances.clear()
95
model = SimpleNamespace(is_moe=True)
96
- prepared: list[object] = []
96
97
monkeypatch.setattr(mrv2, "DeviceMemoryProfiler", FakeMemoryProfiler)
98
monkeypatch.setattr(eplb, "EplbState", FakeEplbState)
@@ -102,7 +101,6 @@ def test_v2_load_model_registers_moe_with_eplb(monkeypatch):
101
"get_model_loader",
102
lambda load_config: SimpleNamespace(load_model=lambda **_: model),
103
)
105
- monkeypatch.setattr(mrv2, "prepare_communication_buffer_for_model", prepared.append)
104
monkeypatch.setattr(
105
mrv2,
106
"init_model_state",
@@ -119,7 +117,6 @@ def test_v2_load_model_registers_moe_with_eplb(monkeypatch):
117
118
assert runner.model is model
119
assert runner.model_state is not None
122
- assert prepared == [model]
120
assert runner.eplb_state is not None
121
assert runner.eplb_state.add_model_calls == [(model, runner.model_config)]
122
assert runner.eplb_state.async_started is True
@@ -128,7 +125,6 @@ def test_v2_load_model_registers_moe_with_eplb(monkeypatch):
125
def test_v2_load_model_with_dummy_weights_skips_eplb_registration(monkeypatch):
126
FakeEplbState.instances.clear()
127
model = SimpleNamespace(is_moe=True)
131
- prepared: list[object] = []
128
129
monkeypatch.setattr(mrv2, "DeviceMemoryProfiler", FakeMemoryProfiler)
130
monkeypatch.setattr(eplb, "EplbState", FakeEplbState)
@@ -137,7 +133,6 @@ def test_v2_load_model_with_dummy_weights_skips_eplb_registration(monkeypatch):
133
"get_model_loader",
134
lambda load_config: SimpleNamespace(load_model=lambda **_: model),
135
)
140
- monkeypatch.setattr(mrv2, "prepare_communication_buffer_for_model", prepared.append)
136
monkeypatch.setattr(
137
mrv2,
138
"init_model_state",
@@ -149,7 +144,6 @@ def test_v2_load_model_with_dummy_weights_skips_eplb_registration(monkeypatch):
144
mrv2.GPUModelRunner.load_model(runner, load_dummy_weights=True)
145
146
assert runner.load_config.load_format == "dummy"
152
- assert prepared == []
147
assert runner.eplb_state is not None
148
assert runner.eplb_state.add_model_calls == []
149
assert runner.eplb_state.async_started is False
vllm/distributed/device_communicators/base_device_communicator.py
-12
@@ -8,7 +8,6 @@ import torch.distributed as dist
8
from torch.distributed import ProcessGroup
9
10
from vllm.logger import init_logger
11
-from vllm.utils import is_moe_layer
11
12
logger = init_logger(__name__)
13
@@ -358,17 +357,6 @@ class DeviceCommunicatorBase:
357
def destroy(self):
358
pass
359
361
- def prepare_communication_buffer_for_model(self, model: torch.nn.Module) -> None:
362
- """
363
- Prepare the communication buffer for the model.
364
- """
365
- if not self.is_ep_communicator:
366
- return
367
-
368
- moe_modules = [module for module in model.modules() if is_moe_layer(module)]
369
- for module in moe_modules:
370
- module.maybe_init_modular_kernel()
371
-
360
def dispatch_router_logits(
361
self,
362
hidden_states: torch.Tensor,
vllm/distributed/elastic_ep/elastic_execute.py
-6
@@ -35,7 +35,6 @@ from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator
35
from vllm.distributed.parallel_state import (
36
_replace_active_groups,
37
get_eplb_group,
38
- prepare_communication_buffer_for_model,
38
)
39
from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator
40
from vllm.logger import init_logger
@@ -485,7 +484,6 @@ class ElasticEPScalingExecutor:
484
for module in moe_modules:
485
if getattr(module._quant_method, "wraps_legacy_quant_method", False):
486
module._replace_quant_method(module._quant_method.old_quant_method)
488
- prepare_communication_buffer_for_model(self.worker.model_runner.model)
487
488
eplb_model_state.expert_buffer = [
489
torch.empty_like(w) for w in model.expert_weights[0]
@@ -652,10 +650,6 @@ class ElasticEPScalingExecutor:
650
old_num_physical_experts,
651
)
652
655
- def prepare_new_worker(self) -> None:
656
- with set_current_vllm_config(self.worker.vllm_config):
657
- prepare_communication_buffer_for_model(self.worker.model_runner.get_model())
658
-
653
def warmup_local_kernels(self) -> None:
654
with set_current_vllm_config(self.worker.vllm_config):
655
kernel_warmup(self.worker, process_local_only=True)
vllm/distributed/elastic_ep/elastic_state.py
-1
@@ -214,7 +214,6 @@ class ElasticEPScalingState:
214
self.engine_core.available_gpu_memory_for_kv_cache = (
215
ParallelConfig.sync_kv_cache_memory_size(self.new_dp_group, -1)
216
)
217
- self._collective_rpc("elastic_ep_execute", args=("prepare_new_worker",))
217
self.state = ScaleUpNewEngineState.PREPARE
218
return True
219
vllm/distributed/parallel_state.py
-25
@@ -1234,10 +1234,6 @@ class GroupCoordinator:
1234
if self.mq_broadcaster is not None:
1235
self.mq_broadcaster = None
1236
1237
- def prepare_communication_buffer_for_model(self, model: torch.nn.Module):
1238
- if self.device_communicator is not None:
1239
- self.device_communicator.prepare_communication_buffer_for_model(model)
1240
-
1237
def dispatch_router_logits(
1238
self,
1239
hidden_states: torch.Tensor,
@@ -2041,27 +2037,6 @@ def ensure_model_parallel_initialized(
2037
)
2038
2039
2044
-def prepare_communication_buffer_for_model(model: torch.nn.Module):
2045
- """Prepare the communication buffer for the model.
2046
- Traditional communication libraries like NCCL are almost
2047
- model agnostic. However, emerging new communication libraries like
2048
- MoE all2all (DeepEP) usually allocate the communication buffer
2049
- based on the model shape for optimal performance.
2050
- """
2051
- if _TP is not None:
2052
- _TP.prepare_communication_buffer_for_model(model)
2053
- if _PCP is not None:
2054
- _PCP.prepare_communication_buffer_for_model(model)
2055
- if _PP is not None:
2056
- _PP.prepare_communication_buffer_for_model(model)
2057
- if _DP is not None:
2058
- _DP.prepare_communication_buffer_for_model(model)
2059
- if _EP is not None:
2060
- _EP.prepare_communication_buffer_for_model(model)
2061
- if _EPLB is not None:
2062
- _EPLB.prepare_communication_buffer_for_model(model)
2063
-
2064
-
2040
def checkpoint_prepare_distributed_state() -> None:
2041
"""Prepare every device communicator for a process checkpoint."""
2042
torch.accelerator.synchronize()
vllm/lora/layers/fused_moe.py
+2
-11
@@ -19,12 +19,8 @@ from vllm.model_executor.layers.fused_moe.fused_moe_modular_method import (
19
FusedMoEModularMethod,
20
)
21
from vllm.model_executor.layers.fused_moe.modular_kernel import (
22
- FusedMoEKernel,
22
FusedMoEKernelModularImpl,
23
)
25
-from vllm.model_executor.layers.fused_moe.prepare_finalize import (
26
- MoEPrepareAndFinalizeNoDPEPModular,
27
-)
24
from vllm.platforms import current_platform
25
26
from .utils import _get_lora_aux_cuda_stream, _get_lora_device
@@ -76,13 +72,8 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
72
assert isinstance(moe_kernel.impl, FusedMoEKernelModularImpl)
73
moe_kernel.impl.shared_experts = None
74
else:
79
- prepare_finalize = MoEPrepareAndFinalizeNoDPEPModular()
80
- moe_kernel = FusedMoEKernel(
81
- prepare_finalize,
82
- routed_experts.quant_method.select_gemm_impl(
83
- prepare_finalize, routed_experts
84
- ),
85
- )
75
+ raise RuntimeError("Old-style modular kernel not supported.")
76
+
77
assert moe_kernel.supports_lora(), (
78
f"{type(moe_kernel.fused_experts).__name__} does not support LoRA. "
79
"For unquantized MoE, set moe_backend='triton' or moe_backend='auto' "
vllm/model_executor/layers/fused_moe/all2all_utils.py
-14
@@ -123,20 +123,6 @@ def maybe_make_prepare_finalize(
123
use_monolithic: bool = False,
124
eep_stage: bool = False,
125
) -> FusedMoEPrepareAndFinalize | None:
126
- # NOTE(rob): we are migrating each quant_method to hold the MK
127
- # in all cases. The allow_new_interface=False flag allow us to fall
128
- # back to the old method for methods that have not yet been migrated.
129
- #
130
- # In old method:
131
- # * maybe_init_modular_kernel() calls this function. If we are
132
- # using no Dp/Ep or naive all2all, we return None this function
133
- # returns None and no ModularKernelMethod is created. If non-naive
134
- # all2all is used, this returns a PrepareAndFinalize object and
135
- # a ModularKernelMethod is created.
136
- # In new method:
137
- # * maybe_make_prepare_finalize() is called from the oracle. We
138
- # always return a PrepareAndFinalize object and the quant method
139
- # holds the ModularKernel.
126
if not moe.moe_parallel_config.use_all2all_kernels:
127
if not allow_new_interface:
128
return None
vllm/model_executor/layers/fused_moe/fused_moe_method_base.py
-28
@@ -13,10 +13,6 @@ from vllm.model_executor.layers.fused_moe.config import (
13
FusedMoEParallelConfig,
14
FusedMoEQuantConfig,
15
)
16
-from vllm.model_executor.layers.fused_moe.modular_kernel import (
17
- FusedMoEExpertsModular,
18
- FusedMoEPrepareAndFinalizeModular,
19
-)
16
from vllm.model_executor.layers.quantization.base_config import (
17
QuantizeMethodBase,
18
)
@@ -103,30 +99,6 @@ class FusedMoEMethodBase(QuantizeMethodBase):
99
hidden_size, act_dtype, moe_parallel_config
100
), intermediate_size_per_partition
101
106
- def maybe_make_prepare_finalize(
107
- self,
108
- routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
109
- ) -> FusedMoEPrepareAndFinalizeModular | None:
110
- from .all2all_utils import maybe_make_prepare_finalize
111
-
112
- pf = maybe_make_prepare_finalize(
113
- self.moe, self.moe_quant_config, routing_tables
114
- )
115
- assert pf is None or isinstance(pf, FusedMoEPrepareAndFinalizeModular)
116
- return pf
117
-
118
- def select_gemm_impl(
119
- self,
120
- prepare_finalize: FusedMoEPrepareAndFinalizeModular,
121
- layer: "RoutedExperts",
122
- ) -> FusedMoEExpertsModular:
123
- # based on the all2all implementation, select the appropriate
124
- # gemm implementation
125
- raise ValueError(
126
- f"{self.__class__.__name__} uses the new modular kernel initialization "
127
- "logic. This function should not be called."
128
- )
129
-
102
@abstractmethod
103
def get_fused_moe_quant_config(
104
self, layer: "RoutedExperts"
vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py
-15
@@ -15,7 +15,6 @@ from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
15
)
16
from vllm.model_executor.layers.fused_moe.modular_kernel import (
17
FusedMoEKernel,
18
- FusedMoEPrepareAndFinalizeModular,
18
)
19
from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
20
SharedExperts,
@@ -47,20 +46,6 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
46
def wraps_legacy_quant_method(self) -> bool:
47
return not self.old_quant_method.supports_internal_mk
48
50
- @staticmethod
51
- def make(
52
- routed_experts: "RoutedExperts",
53
- old_quant_method: FusedMoEMethodBase,
54
- prepare_finalize: FusedMoEPrepareAndFinalizeModular,
55
- ) -> "FusedMoEModularMethod":
56
- return FusedMoEModularMethod(
57
- old_quant_method,
58
- FusedMoEKernel(
59
- prepare_finalize,
60
- old_quant_method.select_gemm_impl(prepare_finalize, routed_experts),
61
- ),
62
- )
63
-
49
@property
50
def skip_forward_padding(self) -> bool:
51
return self.old_quant_method.skip_forward_padding
vllm/model_executor/layers/fused_moe/runner/moe_runner.py
-41
@@ -28,9 +28,6 @@ from vllm.model_executor.layers.fused_moe.config import (
28
from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
29
FusedMoEMethodBase,
30
)
31
-from vllm.model_executor.layers.fused_moe.fused_moe_modular_method import (
32
- FusedMoEModularMethod,
33
-)
31
from vllm.model_executor.layers.fused_moe.routed_experts import (
32
RoutedExperts,
33
)
@@ -896,44 +893,6 @@ class MoERunner(MoERunnerInterface):
893
#
894
#########################################################
895
899
- # Note: maybe_init_modular_kernel should only be called by
900
- # prepare_communication_buffer_for_model.
901
- # This is called after all weight loading and post-processing, so it
902
- # should be safe to swap out the quant_method.
903
- def maybe_init_modular_kernel(self) -> None:
904
- # NOTE(rob): WIP refactor. For quant methods that own the MK
905
- # we create the MK during process_weights_after_loading.
906
- if (
907
- self.routed_experts.quant_method.supports_internal_mk
908
- or self.routed_experts.quant_method.is_monolithic
909
- ):
910
- return None
911
-
912
- self.routed_experts._ensure_moe_quant_config_init()
913
- # routing_tables only needed for round-robin expert placement with
914
- # DeepEP all2all backend.
915
- routing_tables = self._expert_routing_tables()
916
-
917
- if isinstance(self.routed_experts.quant_method, FusedMoEModularMethod):
918
- base_quant_method = self.routed_experts.quant_method.old_quant_method
919
- else:
920
- base_quant_method = self.routed_experts.quant_method
921
-
922
- prepare_finalize = base_quant_method.maybe_make_prepare_finalize(
923
- routing_tables=routing_tables
924
- )
925
- if prepare_finalize is not None:
926
- logger.debug(
927
- "%s for %s(%s)", prepare_finalize.__class__.__name__, self, id(self)
928
- )
929
- self._replace_quant_method(
930
- FusedMoEModularMethod.make(
931
- self.routed_experts,
932
- base_quant_method,
933
- prepare_finalize,
934
- )
935
- )
936
-
896
#
897
# Properties
898
#
vllm/model_executor/layers/fused_moe/runner/moe_runner_interface.py
-4
@@ -66,10 +66,6 @@ class MoERunnerInterface(PluggableLayer, ABC):
66
#
67
########################################################################
68
69
- @abstractmethod
70
- def maybe_init_modular_kernel(self) -> None:
71
- raise NotImplementedError
72
-
69
@property
70
@abstractmethod
71
def layer_id(self):
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
-23
@@ -17,10 +17,6 @@ from vllm.model_executor.layers.fused_moe.config import (
17
from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
18
FusedMoEMethodBase,
19
)
20
-from vllm.model_executor.layers.fused_moe.modular_kernel import (
21
- FusedMoEExpertsModular,
22
- FusedMoEPrepareAndFinalizeModular,
23
-)
20
from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
21
UnquantizedMoeBackend,
22
convert_to_unquantized_kernel_format,
@@ -56,25 +52,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
52
def supports_eplb(self) -> bool:
53
return True
54
59
- def maybe_make_prepare_finalize(
60
- self,
61
- routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
62
- ):
63
- raise ValueError(
64
- f"{self.__class__.__name__} uses the new modular kernel initialization "
65
- "logic. So this function should not be called."
66
- )
67
-
68
- def select_gemm_impl(
69
- self,
70
- prepare_finalize: FusedMoEPrepareAndFinalizeModular,
71
- layer: "RoutedExperts",
72
- ) -> FusedMoEExpertsModular:
73
- raise ValueError(
74
- f"{self.__class__.__name__} uses the new modular kernel initialization "
75
- "logic. This function should not be called."
76
- )
77
-
55
def create_weights(
56
self,
57
layer: "RoutedExperts",
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a16_flydsl.py
-8
@@ -8,7 +8,6 @@ from compressed_tensors.quantization import (
8
QuantizationArgs,
9
)
10
11
-import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11
from vllm.logger import init_logger
12
from vllm.model_executor.layers.fused_moe import (
13
RoutedExperts,
@@ -310,13 +309,6 @@ class CompressedTensorsW4A16FlydslMoEMethod(CompressedTensorsMoEMethod):
309
block_shape=[0, self.group_size],
310
)
311
313
- def select_gemm_impl(
314
- self,
315
- prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular,
316
- layer: torch.nn.Module,
317
- ) -> mk.FusedMoEExpertsModular:
318
- raise NotImplementedError
319
-
312
def apply(
313
self,
314
layer: RoutedExperts,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py
-10
@@ -4,7 +4,6 @@
4
5
import torch
6
7
-import vllm.model_executor.layers.fused_moe.modular_kernel as mk
7
from vllm.logger import init_logger
8
from vllm.model_executor.layers.fused_moe import (
9
FusedMoeWeightScaleSupported,
@@ -242,15 +241,6 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
241
)
242
self.moe_kernel.fused_experts.process_weights_after_loading(layer)
243
245
- def maybe_make_prepare_finalize(
246
- self,
247
- routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
248
- ) -> mk.FusedMoEPrepareAndFinalizeModular | None:
249
- raise ValueError(
250
- f"{self.__class__.__name__} uses the new modular kernel initialization "
251
- "logic. This function should not be called."
252
- )
253
-
244
def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
245
return make_nvfp4_moe_quant_config(
246
backend=self.nvfp4_backend,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py
-10
@@ -8,7 +8,6 @@ from compressed_tensors.quantization import (
8
QuantizationStrategy,
9
)
10
11
-import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11
from vllm.distributed import get_tensor_model_parallel_world_size
12
from vllm.logger import init_logger
13
from vllm.model_executor.layers.fused_moe import (
@@ -341,15 +340,6 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
340
layer=layer,
341
)
342
344
- def maybe_make_prepare_finalize(
345
- self,
346
- routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
347
- ) -> mk.FusedMoEPrepareAndFinalizeModular | None:
348
- raise ValueError(
349
- f"{self.__class__.__name__} uses the new modular kernel initialization "
350
- "logic. This function should not be called."
351
- )
352
-
343
def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
344
is_per_token = self.input_quant.strategy == QuantizationStrategy.TOKEN
345
return make_fp8_moe_quant_config(
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_int8.py
-10
@@ -7,7 +7,6 @@ from compressed_tensors.quantization import (
7
QuantizationStrategy,
8
)
9
10
-import vllm.model_executor.layers.fused_moe.modular_kernel as mk
10
from vllm.logger import init_logger
11
from vllm.model_executor.layers.fused_moe import (
12
FusedMoeWeightScaleSupported,
@@ -164,15 +163,6 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
163
)
164
self.moe_kernel.fused_experts.process_weights_after_loading(layer)
165
167
- def maybe_make_prepare_finalize(
168
- self,
169
- routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
170
- ) -> mk.FusedMoEPrepareAndFinalizeModular | None:
171
- raise ValueError(
172
- f"{self.__class__.__name__} uses the new modular kernel initialization "
173
- "logic. This function should not be called."
174
- )
175
-
166
def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
167
return make_int8_moe_quant_config(
168
int8_backend=self.int8_backend,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py
-10
@@ -3,7 +3,6 @@
3
4
import torch
5
6
-import vllm.model_executor.layers.fused_moe.modular_kernel as mk
6
from vllm.model_executor.layers.fused_moe import (
7
FusedMoeWeightScaleSupported,
8
RoutedExperts,
@@ -159,15 +158,6 @@ class CompressedTensorsW8A8Mxfp8MoEMethod(CompressedTensorsMoEMethod):
158
layer=layer,
159
)
160
162
- def maybe_make_prepare_finalize(
163
- self,
164
- routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
165
- ) -> mk.FusedMoEPrepareAndFinalizeModular | None:
166
- raise ValueError(
167
- f"{self.__class__.__name__} uses the new modular kernel "
168
- "initialization logic. This function should not be called."
169
- )
170
-
161
def apply_monolithic(
162
self,
163
layer: RoutedExperts,
vllm/model_executor/layers/quantization/fp8.py
-10
@@ -6,7 +6,6 @@ from typing import TYPE_CHECKING, Any
6
import torch
7
8
import vllm.envs as envs
9
-import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9
from vllm.config import get_current_vllm_config
10
from vllm.distributed import get_tensor_model_parallel_world_size
11
from vllm.logger import init_logger
@@ -731,15 +730,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):
730
layer, w13, w2, w13_scale, w2_scale, w13_input_scale, w2_input_scale
731
)
732
734
- def maybe_make_prepare_finalize(
735
- self,
736
- routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
737
- ) -> mk.FusedMoEPrepareAndFinalizeModular | None:
738
- raise ValueError(
739
- f"{self.__class__.__name__} uses the new modular kernel initialization "
740
- "logic. This function should not be called."
741
- )
742
-
733
def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig:
734
w1_scale = getattr(layer, f"w13_{self.weight_scale_name}")
735
w2_scale = getattr(layer, f"w2_{self.weight_scale_name}")
vllm/model_executor/layers/quantization/modelopt.py
-28
@@ -8,7 +8,6 @@ import torch
8
from torch.nn.parameter import Parameter
9
10
import vllm.envs as envs
11
-import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11
from vllm.config import get_current_vllm_config
12
from vllm.logger import init_logger
13
from vllm.model_executor.kernels.linear import (
@@ -771,15 +770,6 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
770
activation_key=kFp8StaticTensorSym,
771
)
772
774
- def maybe_make_prepare_finalize(
775
- self,
776
- routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
777
- ) -> mk.FusedMoEPrepareAndFinalizeModular | None:
778
- raise ValueError(
779
- f"{self.__class__.__name__} uses the new modular kernel initialization "
780
- "logic. This function should not be called."
781
- )
782
-
773
def create_weights(
774
self,
775
layer: RoutedExperts,
@@ -1409,15 +1399,6 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1399
self.nvfp4_backend
1400
)
1401
1412
- def maybe_make_prepare_finalize(
1413
- self,
1414
- routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
1415
- ) -> mk.FusedMoEPrepareAndFinalizeModular | None:
1416
- raise ValueError(
1417
- f"{self.__class__.__name__} uses the new modular kernel initialization "
1418
- "logic. This function should not be called."
1419
- )
1420
-
1402
def uses_weight_scale_2_pattern(self) -> bool:
1403
"""
1404
FP4 variants use 'weight_scale_2' pattern for per-tensor weight scales.
@@ -2087,15 +2068,6 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
2068
):
2069
self._dequant_mxfp8_weights_to_bf16(layer)
2070
2090
- def maybe_make_prepare_finalize(
2091
- self,
2092
- routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
2093
- ) -> mk.FusedMoEPrepareAndFinalizeModular | None:
2094
- raise ValueError(
2095
- f"{self.__class__.__name__} uses the new modular kernel initialization "
2096
- "logic. This function should not be called."
2097
- )
2098
-
2071
def get_fused_moe_quant_config(
2072
self, layer: RoutedExperts
2073
) -> FusedMoEQuantConfig | None:
vllm/model_executor/layers/quantization/online/moe_base.py
-10
@@ -5,7 +5,6 @@ from abc import abstractmethod
5
6
import torch
7
8
-import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8
from vllm.model_executor.layers.fused_moe import (
9
FusedMoEMethodBase,
10
RoutedExperts,
@@ -100,15 +99,6 @@ class OnlineMoEMethodBase(FusedMoEMethodBase):
99
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
100
pass
101
103
- def maybe_make_prepare_finalize(
104
- self,
105
- routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
106
- ) -> mk.FusedMoEPrepareAndFinalizeModular | None:
107
- raise ValueError(
108
- f"{self.__class__.__name__} uses the new modular kernel "
109
- "initialization logic. This function should not be called."
110
- )
111
-
102
@property
103
def supports_eplb(self) -> bool:
104
return True
vllm/v1/worker/gpu/model_runner.py
-6
@@ -34,7 +34,6 @@ from vllm.config.compilation import CUDAGraphMode
34
from vllm.distributed.parallel_state import (
35
get_dcp_group,
36
get_pp_group,
37
- prepare_communication_buffer_for_model,
37
)
38
from vllm.forward_context import BatchDescriptor, set_forward_context
39
from vllm.logger import init_logger
@@ -354,11 +353,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
353
time_after_load - time_before_load,
354
)
355
357
- if not load_dummy_weights:
358
- prepare_communication_buffer_for_model(self.model)
359
- if self.speculator is not None:
360
- prepare_communication_buffer_for_model(self.speculator.model)
361
-
356
# Initialize the components that require the model.
357
self.model_state = init_model_state(
358
self.vllm_config, self.model, self.encoder_cache, self.device
vllm/v1/worker/gpu_model_runner.py
+1
-7
@@ -50,7 +50,6 @@ from vllm.distributed.parallel_state import (
50
get_tp_group,
51
graph_capture,
52
is_global_first_rank,
53
- prepare_communication_buffer_for_model,
53
)
54
from vllm.forward_context import (
55
BatchDescriptor,
@@ -5387,12 +5386,7 @@ class GPUModelRunner(
5386
format_gib(self.model_memory_usage),
5387
time_after_load - time_before_load,
5388
)
5390
- if not load_dummy_weights:
5391
- prepare_communication_buffer_for_model(self.model)
5392
- if (drafter := getattr(self, "drafter", None)) and (
5393
- drafter_model := getattr(drafter, "model", None)
5394
- ):
5395
- prepare_communication_buffer_for_model(drafter_model)
5389
+
5390
mm_config = self.model_config.multimodal_config
5391
self.is_multimodal_pruning_enabled = (
5392
supports_multimodal_pruning(self.get_model())