@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 6d30655b13
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 2, 23:15 UTC

elastic_ep: stage/commit MoE quant method on reconfigure (#40881)

Signed-off-by: Itay Alroy <ialroy@nvidia.com>

Itay Alroy committed May 23, 2026 at 01:57 UTC 6d30655b136bec37e7e7c7d045f1bcedc8cf3e2b
7 files changed +351 -59
vllm/distributed/device_communicators/all2all.py
+98 -35
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import threading
4 +from dataclasses import dataclass
5 from typing import Any
6
7 import torch
@@ -324,15 +325,21 @@ class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase):
325 return 0
326
327
328 +@dataclass
329 +class _NixlEPBufferState:
330 + buffer: Any
331 + connected_ep_size: int
332 + active_ep_size: int
333 +
334 +
335 class NixlEPAll2AllManager(All2AllManagerBase):
336 """
337 All2All communication based on NIXL EP kernels.
338 This backend supports elastic EP with dynamic rank connection/disconnection.
339 """
340
333 - # (nixl_ep_buffer, ep_size)
334 - _buffer: tuple[Any, int] | None = None
335 - _lock = threading.Lock()
341 + _buffer: _NixlEPBufferState | None = None
342 + _lock = threading.RLock()
343
344 def __init__(self, cpu_group, tcp_store_group=None):
345 assert tcp_store_group is not None
@@ -367,47 +374,103 @@ class NixlEPAll2AllManager(All2AllManagerBase):
374 num_experts_per_rank=num_experts_per_rank,
375 num_rdma_bytes=num_rdma_bytes,
376 )
370 - ranks_to_connect = list(range(self.cpu_group.size()))
377 + ranks_to_connect = list(range(self.world_size))
378 buffer.connect_ranks(ranks_to_connect)
372 - NixlEPAll2AllManager._buffer = (buffer, self.cpu_group.size())
379 + NixlEPAll2AllManager._buffer = _NixlEPBufferState(
380 + buffer=buffer,
381 + connected_ep_size=self.world_size,
382 + active_ep_size=self.world_size,
383 + )
384 +
385 + def _connect_to_ep_size(self, ep_size: int, *, make_active: bool) -> None:
386 + assert NixlEPAll2AllManager._buffer is not None
387 + state = NixlEPAll2AllManager._buffer
388 + if ep_size <= state.connected_ep_size:
389 + return
390 +
391 + state.buffer.set_tcp_store_group(self.tcp_store_group.store)
392 + ranks_to_connect = list(range(state.connected_ep_size, ep_size))
393 + state.buffer.connect_ranks(ranks_to_connect, activate=make_active)
394 + state.connected_ep_size = ep_size
395 + if make_active:
396 + state.active_ep_size = ep_size
397 +
398 + def _disconnect_to_ep_size(self, ep_size: int) -> None:
399 + assert NixlEPAll2AllManager._buffer is not None
400 + state = NixlEPAll2AllManager._buffer
401 + if ep_size >= state.connected_ep_size:
402 + return
403 +
404 + state.buffer.set_tcp_store_group(self.tcp_store_group.store)
405 + ranks_to_disconnect = list(range(ep_size, state.connected_ep_size))
406 + state.buffer.disconnect_ranks(ranks_to_disconnect)
407 + state.connected_ep_size = ep_size
408 + state.active_ep_size = min(state.active_ep_size, ep_size)
409 +
410 + def _unmask_connected_ranks(self, target_ep_size: int) -> None:
411 + assert NixlEPAll2AllManager._buffer is not None
412 + state = NixlEPAll2AllManager._buffer
413 + state.buffer.set_tcp_store_group(self.tcp_store_group.store)
414 + if target_ep_size <= state.active_ep_size:
415 + return
416 + assert state.connected_ep_size >= target_ep_size
417 +
418 + for rank in range(state.active_ep_size, target_ep_size):
419 + state.buffer.update_mask_buffer(rank, mask=False)
420 + state.active_ep_size = target_ep_size
421
374 - def _update_buffer(self):
422 + def _stage_ep_size(self) -> None:
423 assert NixlEPAll2AllManager._buffer is not None
376 - buffer, current_ep_size = NixlEPAll2AllManager._buffer
377 - current_ranks = list(range(current_ep_size))
378 - new_ep_size = self.cpu_group.size()
379 - buffer.set_tcp_store_group(self.tcp_store_group.store)
380 - if new_ep_size > len(current_ranks):
381 - ranks_to_connect = list(range(len(current_ranks), new_ep_size))
382 - buffer.connect_ranks(ranks_to_connect)
424 + state = NixlEPAll2AllManager._buffer
425 + target_ep_size = self.world_size
426 +
427 + # Scale-up can safely connect standby ranks while leaving them masked.
428 + # Scale-down must not disconnect active ranks until commit.
429 + if target_ep_size > state.connected_ep_size:
430 + self._connect_to_ep_size(target_ep_size, make_active=False)
431 +
432 + def commit_staged_state(self) -> None:
433 + """Commit staged NIXL EP state to the active communication set."""
434 + with NixlEPAll2AllManager._lock:
435 + assert NixlEPAll2AllManager._buffer is not None
436 + state = NixlEPAll2AllManager._buffer
437 + target_ep_size = self.world_size
438 +
439 + if target_ep_size < state.connected_ep_size:
440 + self._disconnect_to_ep_size(target_ep_size)
441 + elif target_ep_size > state.connected_ep_size:
442 + self._connect_to_ep_size(target_ep_size, make_active=True)
443 +
444 + self._unmask_connected_ranks(target_ep_size)
445 +
446 + def _ensure_ep_size(self, *, stage: bool) -> None:
447 + if stage:
448 + self._stage_ep_size()
449 else:
384 - ranks_to_disconnect = current_ranks[new_ep_size:]
385 - buffer.disconnect_ranks(ranks_to_disconnect)
386 - NixlEPAll2AllManager._buffer = (buffer, new_ep_size)
450 + self.commit_staged_state()
451
452 def get_handle(self, kwargs):
453 with NixlEPAll2AllManager._lock:
390 - if (
391 - NixlEPAll2AllManager._buffer is not None
392 - and NixlEPAll2AllManager._buffer[1] == self.cpu_group.size()
393 - ):
394 - return NixlEPAll2AllManager._buffer[0]
395 -
396 - num_experts_per_rank = (
397 - kwargs["num_global_experts"] // kwargs["num_ep_ranks"]
398 - )
399 - nixl_kwargs = dict(
400 - max_num_tokens_per_dp_rank=kwargs["max_num_tokens_per_dp_rank"],
401 - token_hidden_size=kwargs["token_hidden_size"],
402 - num_experts_per_rank=num_experts_per_rank,
403 - )
404 - if NixlEPAll2AllManager._buffer is None:
405 - self._init_buffer(**nixl_kwargs)
454 + stage = bool(kwargs.get("stage", False))
455 + state = NixlEPAll2AllManager._buffer
456 + if state is None:
457 + assert not stage, (
458 + "NIXL EP staged initialization requires an existing buffer"
459 + )
460 + max_num_tokens_per_dp_rank = kwargs["max_num_tokens_per_dp_rank"]
461 + num_experts_per_rank = (
462 + kwargs["num_global_experts"] // kwargs["num_ep_ranks"]
463 + )
464 + self._init_buffer(
465 + max_num_tokens_per_dp_rank=max_num_tokens_per_dp_rank,
466 + token_hidden_size=kwargs["token_hidden_size"],
467 + num_experts_per_rank=num_experts_per_rank,
468 + )
469 else:
407 - self._update_buffer()
470 + self._ensure_ep_size(stage=stage)
471
472 assert NixlEPAll2AllManager._buffer is not None
410 - handle = NixlEPAll2AllManager._buffer[0]
473 + handle = NixlEPAll2AllManager._buffer.buffer
474 return handle
475
476 def dispatch(
@@ -432,7 +495,7 @@ class NixlEPAll2AllManager(All2AllManagerBase):
495 # NOTE(yongji): NIXLEPAll2AllManager instance is recreated during
496 # scale-up/down, so we cannot destroy the persistent buffer here.
497 assert NixlEPAll2AllManager._buffer is not None
435 - buffer = NixlEPAll2AllManager._buffer[0]
498 + buffer = NixlEPAll2AllManager._buffer.buffer
499 buffer.set_tcp_store_group(None)
500
501 # NIXL EP uses RDMA so no SMs are used for communication
vllm/distributed/elastic_ep/elastic_execute.py
+74 -15
@@ -4,6 +4,8 @@ import copy
4 import gc
5 import weakref
6 from collections.abc import Iterable, Sequence
7 +from dataclasses import replace
8 +from typing import TYPE_CHECKING
9
10 import torch
11 import torch.nn as nn
@@ -37,7 +39,10 @@ from vllm.distributed.parallel_state import (
39 )
40 from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator
41 from vllm.logger import init_logger
40 -from vllm.model_executor.layers.fused_moe.layer import FusedMoEParallelConfig
42 +from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig
43 +from vllm.model_executor.layers.fused_moe.eep_reconfigure import (
44 + make_eep_staged_quant_method,
45 +)
46 from vllm.utils import is_moe_layer
47 from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
48 from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper
@@ -45,6 +50,11 @@ from vllm.v1.worker.workspace import lock_workspace, unlock_workspace
50
51 logger = init_logger(__name__)
52
53 +if TYPE_CHECKING:
54 + from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
55 + FusedMoEMethodBase,
56 + )
57 +
58
59 def batch_transfer_weights(
60 model: nn.Module,
@@ -134,6 +144,7 @@ class ElasticEPScalingExecutor:
144 def __init__(self, worker):
145 self.worker_ref = weakref.ref(worker)
146 self.reconfig_request = None
147 + self._staged_moe_quant_methods: dict[nn.Module, FusedMoEMethodBase] = {}
148
149 @property
150 def worker(self):
@@ -196,6 +207,8 @@ class ElasticEPScalingExecutor:
207 )
208 if new_dp_size > old_dp_size:
209 self._set_eplb_suppressed(True)
210 + elif new_dp_size < old_dp_size:
211 + self._stage_standby_moe_quant_methods()
212
213 def transfer_weights(self, old_dp_size: int, new_dp_size: int) -> None:
214 standby_dp_group = get_standby_dp_group()
@@ -262,6 +275,58 @@ class ElasticEPScalingExecutor:
275 src_rank=0,
276 device=self.worker.device,
277 )
278 + # New workers enter load_model after receiving the expert mapping.
279 + # Stage replacement MoE kernels before returning to the state machine
280 + # so existing ranks can participate in collective EP comm creation.
281 + self._stage_standby_moe_quant_methods()
282 +
283 + def _make_eep_moe_config(self, module, dp_group, ep_group):
284 + parallel_config = self.worker.vllm_config.parallel_config
285 + tp_size = get_tp_group().world_size
286 + sp_size = tp_size if parallel_config.use_sequence_parallel_moe else 1
287 + moe_parallel_config = FusedMoEParallelConfig.make(
288 + tp_size_=tp_size,
289 + pcp_size_=get_pcp_group().world_size,
290 + dp_size_=dp_group.world_size,
291 + sp_size_=sp_size,
292 + vllm_parallel_config=parallel_config,
293 + )
294 + return replace(
295 + module.moe_config,
296 + num_experts=module.moe_config.num_local_experts * ep_group.world_size,
297 + moe_parallel_config=moe_parallel_config,
298 + )
299 +
300 + def _stage_standby_moe_quant_methods(self) -> None:
301 + standby_dp_group = get_standby_dp_group()
302 + standby_ep_group = get_standby_ep_group()
303 + model = self.worker.model_runner.get_model()
304 + moe_modules = [module for module in model.modules() if is_moe_layer(module)]
305 + self._staged_moe_quant_methods.clear()
306 + with set_current_vllm_config(self.worker.vllm_config):
307 + for module in moe_modules:
308 + staged_quant_method = make_eep_staged_quant_method(
309 + module,
310 + self._make_eep_moe_config(
311 + module,
312 + standby_dp_group,
313 + standby_ep_group,
314 + ),
315 + )
316 + if staged_quant_method is not None:
317 + self._staged_moe_quant_methods[module] = staged_quant_method
318 +
319 + def _commit_staged_moe_quant_methods(self) -> None:
320 + model = self.worker.model_runner.get_model()
321 + moe_modules = [module for module in model.modules() if is_moe_layer(module)]
322 + for module in moe_modules:
323 + staged_quant_method = self._staged_moe_quant_methods.pop(module, None)
324 + if staged_quant_method is None:
325 + continue
326 + assert staged_quant_method.moe_kernel is not None
327 + module._replace_quant_method(staged_quant_method)
328 + staged_quant_method.moe_kernel.prepare_finalize.on_commit()
329 + self._staged_moe_quant_methods.clear()
330
331 def _release_cuda_graphs(self) -> None:
332 if isinstance(self.worker.model_runner.model, CUDAGraphWrapper):
@@ -327,19 +392,13 @@ class ElasticEPScalingExecutor:
392 module.moe_config.num_local_experts == num_local_experts
393 for module in moe_modules
394 ), "All MoE modules must have the same number of experts"
395 + dp_group = get_dp_group()
396 + ep_group = get_ep_group()
397 for module in moe_modules:
331 - module.moe_config.num_experts = num_local_experts * new_ep_size
398 + new_moe_config = self._make_eep_moe_config(module, dp_group, ep_group)
399 + module.moe_config.num_experts = new_moe_config.num_experts
400 module.global_num_experts = module.moe_config.num_experts
333 - tp_size = get_tp_group().world_size
334 - is_sequence_parallel = parallel_config.use_sequence_parallel_moe
335 - sp_size = tp_size if is_sequence_parallel else 1
336 - module.moe_parallel_config = FusedMoEParallelConfig.make(
337 - tp_size_=tp_size,
338 - pcp_size_=get_pcp_group().world_size,
339 - dp_size_=get_dp_group().world_size,
340 - sp_size_=sp_size,
341 - vllm_parallel_config=parallel_config,
342 - )
401 + module.moe_parallel_config = new_moe_config.moe_parallel_config
402 module.moe_config.moe_parallel_config = module.moe_parallel_config
403
404 # Update EPLB state
@@ -404,10 +463,10 @@ class ElasticEPScalingExecutor:
463 num_physical_experts=num_physical_experts,
464 num_local_physical_experts=num_local_experts,
465 )
407 - # Force re-creation of the modular kernel (and all2all manager)
408 - # for the new EP size by resetting quant_method to base
466 + self._commit_staged_moe_quant_methods()
467 + # Legacy modular methods need to be recreated for the new EP size.
468 for module in moe_modules:
410 - if hasattr(module.quant_method, "old_quant_method"):
469 + if getattr(module.quant_method, "wraps_legacy_quant_method", False):
470 module._replace_quant_method(module.quant_method.old_quant_method)
471 prepare_communication_buffer_for_model(self.worker.model_runner.model)
472
vllm/model_executor/layers/fused_moe/all2all_utils.py
+21 -8
@@ -49,6 +49,22 @@ if current_platform.is_cuda_alike():
49 )
50
51
52 +def _get_ep_all2all_manager(eep_stage: bool = False) -> Any:
53 + if eep_stage:
54 + from vllm.distributed.elastic_ep.standby_state import get_standby_ep_group
55 +
56 + ep_group = get_standby_ep_group()
57 + assert ep_group is not None
58 + device_communicator = ep_group.device_communicator
59 + else:
60 + device_communicator = get_ep_group().device_communicator
61 +
62 + assert device_communicator is not None
63 + all2all_manager = device_communicator.all2all_manager
64 + assert all2all_manager is not None
65 + return all2all_manager
66 +
67 +
68 def maybe_roundup_layer_hidden_size(
69 hidden_size: int,
70 act_dtype: torch.dtype,
@@ -92,6 +108,7 @@ def maybe_make_prepare_finalize(
108 routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
109 allow_new_interface: bool = False,
110 use_monolithic: bool = False,
111 + eep_stage: bool = False,
112 ) -> FusedMoEPrepareAndFinalize | None:
113 # NOTE(rob): we are migrating each quant_method to hold the MK
114 # in all cases. The allow_new_interface=False flag allow us to fall
@@ -117,21 +134,16 @@ def maybe_make_prepare_finalize(
134 "Detected DP deployment with no --enable-expert-parallel. "
135 "Falling back to AllGather+ReduceScatter dispatch/combine."
136 )
120 - device_communicator = get_ep_group().device_communicator
121 - assert device_communicator is not None
122 - assert device_communicator.all2all_manager is not None
137 + all2all_manager = _get_ep_all2all_manager(eep_stage)
138 return make_moe_prepare_and_finalize_naive_dp_ep(
139 is_sequence_parallel=moe.moe_parallel_config.is_sequence_parallel,
125 - num_dispatchers=(device_communicator.all2all_manager.world_size),
140 + num_dispatchers=all2all_manager.world_size,
141 use_monolithic=use_monolithic,
142 )
143 else:
144 return make_moe_prepare_and_finalize_no_dp_ep(use_monolithic)
145
131 - device_communicator = get_ep_group().device_communicator
132 - assert device_communicator is not None
133 - all2all_manager = device_communicator.all2all_manager
134 - assert all2all_manager is not None
146 + all2all_manager = _get_ep_all2all_manager(eep_stage)
147
148 prepare_finalize: FusedMoEPrepareAndFinalize | None = None
149
@@ -283,6 +295,7 @@ def maybe_make_prepare_finalize(
295 num_ep_ranks=all2all_manager.world_size,
296 num_global_experts=moe.num_experts,
297 num_local_experts=moe.num_experts // all2all_manager.world_size,
298 + stage=eep_stage,
299 )
300 handle = all2all_manager.get_handle(all_to_all_args)
301
vllm/model_executor/layers/fused_moe/eep_reconfigure.py new
+123
@@ -0,0 +1,123 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from inspect import signature
5 +from typing import TYPE_CHECKING, Any
6 +
7 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8 +from vllm.model_executor.layers.fused_moe.all2all_utils import (
9 + maybe_make_prepare_finalize,
10 +)
11 +from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
12 +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
13 + FusedMoEMethodBase,
14 +)
15 +from vllm.model_executor.layers.fused_moe.fused_moe_modular_method import (
16 + FusedMoEModularMethod,
17 +)
18 +from vllm.model_executor.layers.fused_moe.modular_kernel import (
19 + FusedMoEExpertsModular,
20 + FusedMoEPrepareAndFinalizeModular,
21 +)
22 +
23 +if TYPE_CHECKING:
24 + from vllm.model_executor.layers.fused_moe.layer import FusedMoE
25 +
26 +
27 +def _make_eep_experts(
28 + quant_method: FusedMoEMethodBase,
29 + source_experts: FusedMoEExpertsModular,
30 + prepare_finalize: FusedMoEPrepareAndFinalizeModular,
31 + moe_config: FusedMoEConfig,
32 +) -> FusedMoEExpertsModular:
33 + experts_cls = source_experts.__class__
34 + assert quant_method.moe_quant_config is not None
35 + experts_kwargs: dict[str, Any] = {
36 + "moe_config": moe_config,
37 + "quant_config": quant_method.moe_quant_config,
38 + }
39 + if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts:
40 + max_num_tokens = prepare_finalize.max_num_tokens_per_rank()
41 + assert max_num_tokens is not None
42 + experts_kwargs.update(
43 + max_num_tokens=max_num_tokens,
44 + num_dispatchers=prepare_finalize.num_dispatchers(),
45 + )
46 +
47 + # Expert kernels with extra init params need explicit EEP support.
48 + generic_arg_names = set(signature(mk.FusedMoEExperts.__init__).parameters)
49 + ctor_arg_names = set(signature(experts_cls.__init__).parameters)
50 + unsupported_args = ctor_arg_names - generic_arg_names
51 + missing_args = set(experts_kwargs) - ctor_arg_names
52 + if unsupported_args or missing_args:
53 + raise NotImplementedError(
54 + f"{experts_cls.__name__} experts do not support Elastic EP."
55 + )
56 +
57 + return experts_cls(**experts_kwargs)
58 +
59 +
60 +def make_eep_staged_quant_method(
61 + module: "FusedMoE",
62 + moe_config: FusedMoEConfig,
63 +) -> FusedMoEMethodBase | None:
64 + quant_method = module.quant_method
65 + if not quant_method.supports_internal_mk:
66 + return None
67 + if getattr(quant_method, "wraps_legacy_quant_method", False):
68 + return None
69 +
70 + old_batched_format = (
71 + module.moe_config.moe_parallel_config.use_batched_activation_format
72 + )
73 + new_batched_format = moe_config.moe_parallel_config.use_batched_activation_format
74 + assert old_batched_format == new_batched_format
75 +
76 + moe_kernel = quant_method.moe_kernel
77 + if moe_kernel is None:
78 + return None
79 + if moe_kernel.is_monolithic:
80 + raise NotImplementedError(
81 + "Elastic EP full modular-kernel staging is not supported for "
82 + "monolithic fused MoE kernels."
83 + )
84 + if quant_method.moe_quant_config is None:
85 + raise ValueError(
86 + "Elastic EP full modular-kernel staging requires initialized "
87 + "MoE quant config."
88 + )
89 +
90 + prepare_finalize = maybe_make_prepare_finalize(
91 + moe_config,
92 + quant_method.moe_quant_config,
93 + routing_tables=None,
94 + allow_new_interface=True,
95 + use_monolithic=quant_method.is_monolithic,
96 + eep_stage=True,
97 + )
98 + assert prepare_finalize is not None
99 + assert isinstance(prepare_finalize, FusedMoEPrepareAndFinalizeModular)
100 +
101 + source_experts = moe_kernel.fused_experts
102 + assert isinstance(source_experts, FusedMoEExpertsModular)
103 +
104 + experts = _make_eep_experts(
105 + quant_method,
106 + source_experts,
107 + prepare_finalize,
108 + moe_config,
109 + )
110 +
111 + if isinstance(quant_method, FusedMoEModularMethod):
112 + base_quant_method = quant_method.old_quant_method
113 + else:
114 + base_quant_method = quant_method
115 +
116 + return FusedMoEModularMethod(
117 + base_quant_method,
118 + mk.FusedMoEKernel(
119 + prepare_finalize,
120 + experts,
121 + inplace=moe_kernel.inplace,
122 + ),
123 + )
vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py
+5 -1
@@ -31,7 +31,7 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
31 def __init__(
32 self, old_quant_method: FusedMoEMethodBase, moe_kernel: FusedMoEKernel
33 ):
34 - super().__init__(old_quant_method.moe)
34 + super().__init__(moe_kernel.moe_config)
35 self.moe_quant_config = old_quant_method.moe_quant_config
36 self.moe_kernel = moe_kernel
37 self.disable_expert_map = getattr(
@@ -42,6 +42,10 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
42 self.old_quant_method = old_quant_method
43 logger.debug("Swapping out %s", self.old_quant_method.__class__.__name__)
44
45 + @property
46 + def wraps_legacy_quant_method(self) -> bool:
47 + return not self.old_quant_method.supports_internal_mk
48 +
49 @staticmethod
50 def make(
51 moe_layer: torch.nn.Module,
vllm/model_executor/layers/fused_moe/modular_kernel.py
+17
@@ -246,6 +246,13 @@ class FusedMoEPrepareAndFinalize(ABC):
246 """
247 return False
248
249 + def on_commit(self) -> None:
250 + """
251 + Runs after this prepare/finalize has been committed to the active
252 + MoE kernel.
253 + """
254 + return
255 +
256
257 # TODO: pass FusedMoEParallelConfig in as ctor parameter?
258 class FusedMoEPrepareAndFinalizeModular(FusedMoEPrepareAndFinalize):
@@ -1542,6 +1549,12 @@ class FusedMoEKernel:
1549 else:
1550 return False
1551
1552 + @property
1553 + def inplace(self) -> bool:
1554 + if isinstance(self.impl, FusedMoEKernelModularImpl):
1555 + return self.impl.inplace
1556 + return False
1557 +
1558 @property
1559 def is_monolithic(self) -> bool:
1560 return isinstance(self.impl, FusedMoEKernelMonolithicImpl)
@@ -1554,6 +1567,10 @@ class FusedMoEKernel:
1567 def fused_experts(self) -> FusedMoEExperts:
1568 return self.impl.fused_experts
1569
1570 + @property
1571 + def moe_config(self) -> FusedMoEConfig:
1572 + return self.fused_experts.moe_config
1573 +
1574 def supports_lora(self) -> bool:
1575 return self.fused_experts.supports_lora()
1576
vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py
+13
@@ -7,6 +7,8 @@ import torch
7
8 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9 from vllm import envs
10 +from vllm.distributed import get_ep_group
11 +from vllm.distributed.device_communicators.all2all import NixlEPAll2AllManager
12 from vllm.logger import init_logger
13 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
14 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
@@ -138,6 +140,17 @@ class NixlEPPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
140 def max_num_tokens_per_rank(self) -> int | None:
141 return self.max_tokens_per_rank
142
143 + def on_commit(self) -> None:
144 + device_communicator = get_ep_group().device_communicator
145 + assert device_communicator is not None
146 + all2all_manager = device_communicator.all2all_manager
147 + assert isinstance(all2all_manager, NixlEPAll2AllManager)
148 + # maybe_make_prepare_finalize(..., eep_stage=True) initializes self.buffer
149 + # with get_handle(..., stage=True), which stages global NIXL state for the
150 + # new config but leaves it inactive while the old config remains active.
151 + # When EEP commit switches to this P/F, this P/F needs to commit that state.
152 + all2all_manager.commit_staged_state()
153 +
154 def topk_indices_dtype(self) -> torch.dtype | None:
155 return torch.int64
156