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

[2/N] Elastic EP Milestone 2: Integrating NIXL-EP (#35627)

Signed-off-by: Itay Alroy <ialroy@nvidia.com> Co-authored-by: Yongji Wu <wuyongji317@gmail.com> Co-authored-by: Ron Tourgeman <rtourgeman@nvidia.com>

Itay Alroy committed Mar 13, 2026 at 15:25 UTC d5af196c183bef2e886c7ec12db63b6161cacfde
14 files changed +635 -11
tests/v1/engine/test_engine_core_client.py
+1
@@ -150,6 +150,7 @@ def test_mp_client_uses_env_timeout(monkeypatch: pytest.MonkeyPatch):
150 data_parallel_hybrid_lb=False,
151 data_parallel_external_lb=False,
152 local_engines_only=False,
153 + enable_elastic_ep=False,
154 )
155 vllm_config = SimpleNamespace(parallel_config=parallel_config)
156
vllm/config/parallel.py
+3
@@ -43,6 +43,7 @@ All2AllBackend = Literal[
43 "deepep_high_throughput",
44 "deepep_low_latency",
45 "mori",
46 + "nixl_ep",
47 "allgather_reducescatter",
48 "flashinfer_all2allv",
49 ]
@@ -156,6 +157,7 @@ class ParallelConfig:
157 - "deepep_high_throughput": Use deepep high-throughput kernels\n
158 - "deepep_low_latency": Use deepep low-latency kernels\n
159 - "mori": Use mori kernels\n
160 + - "nixl_ep": Use nixl-ep kernels\n
161 - "flashinfer_all2allv": Use flashinfer alltoallv kernels for mnnvl"""
162
163 max_parallel_loading_workers: int | None = None
@@ -580,6 +582,7 @@ class ParallelConfig:
582 "deepep_high_throughput",
583 "deepep_low_latency",
584 "mori",
585 + "nixl_ep",
586 )
587 and self.enable_expert_parallel
588 and self.tensor_parallel_size > 1
vllm/distributed/device_communicators/all2all.py
+116
@@ -1,5 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import threading
4 from typing import Any
5
6 import torch
@@ -413,6 +414,121 @@ class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase):
414 return 0
415
416
417 +class NixlEPAll2AllManager(All2AllManagerBase):
418 + """
419 + All2All communication based on NIXL EP kernels.
420 + This backend supports elastic EP with dynamic rank connection/disconnection.
421 + """
422 +
423 + # (nixl_ep_buffer, ep_size)
424 + _buffer: tuple[Any, int] | None = None
425 + _lock = threading.Lock()
426 +
427 + def __init__(self, cpu_group, tcp_store_group=None):
428 + super().__init__(cpu_group, tcp_store_group)
429 +
430 + self.max_num_ep_ranks = envs.VLLM_NIXL_EP_MAX_NUM_RANKS
431 +
432 + def _init_buffer(
433 + self,
434 + max_num_tokens_per_dp_rank: int,
435 + token_hidden_size: int,
436 + num_experts_per_rank: int,
437 + ) -> None:
438 + from nixl_ep import Buffer # type: ignore[import-not-found]
439 +
440 + max_num_global_experts = self.max_num_ep_ranks * num_experts_per_rank
441 + num_rdma_bytes = Buffer.get_rdma_size_hint(
442 + num_max_dispatch_tokens_per_rank=max_num_tokens_per_dp_rank,
443 + hidden=token_hidden_size,
444 + num_ranks=self.max_num_ep_ranks,
445 + num_experts=max_num_global_experts,
446 + )
447 + assert NixlEPAll2AllManager._buffer is None, (
448 + "NIXL EP buffer already initialized"
449 + )
450 + buffer = Buffer(
451 + rank=self.rank,
452 + tcp_store_group=self.tcp_store_group.store,
453 + )
454 + buffer.update_memory_buffers(
455 + num_ranks=self.max_num_ep_ranks,
456 + num_experts_per_rank=num_experts_per_rank,
457 + num_rdma_bytes=num_rdma_bytes,
458 + )
459 + ranks_to_connect = list(range(self.cpu_group.size()))
460 + buffer.connect_ranks(ranks_to_connect)
461 + NixlEPAll2AllManager._buffer = (buffer, self.cpu_group.size())
462 +
463 + def _update_buffer(self):
464 + assert NixlEPAll2AllManager._buffer is not None
465 + buffer, current_ep_size = NixlEPAll2AllManager._buffer
466 + current_ranks = list(range(current_ep_size))
467 + new_ep_size = self.cpu_group.size()
468 + buffer.set_tcp_store_group(self.tcp_store_group.store)
469 + if new_ep_size > len(current_ranks):
470 + ranks_to_connect = list(range(len(current_ranks), new_ep_size))
471 + buffer.connect_ranks(ranks_to_connect)
472 + else:
473 + ranks_to_disconnect = current_ranks[new_ep_size:]
474 + buffer.disconnect_ranks(ranks_to_disconnect)
475 + NixlEPAll2AllManager._buffer = (buffer, new_ep_size)
476 +
477 + def get_handle(self, kwargs):
478 + with NixlEPAll2AllManager._lock:
479 + if (
480 + NixlEPAll2AllManager._buffer is not None
481 + and NixlEPAll2AllManager._buffer[1] == self.cpu_group.size()
482 + ):
483 + return NixlEPAll2AllManager._buffer[0]
484 +
485 + num_experts_per_rank = (
486 + kwargs["num_global_experts"] // kwargs["num_ep_ranks"]
487 + )
488 + nixl_kwargs = dict(
489 + max_num_tokens_per_dp_rank=kwargs["max_num_tokens_per_dp_rank"],
490 + token_hidden_size=kwargs["token_hidden_size"],
491 + num_experts_per_rank=num_experts_per_rank,
492 + )
493 + if NixlEPAll2AllManager._buffer is None:
494 + self._init_buffer(**nixl_kwargs)
495 + else:
496 + self._update_buffer()
497 +
498 + assert NixlEPAll2AllManager._buffer is not None
499 + handle = NixlEPAll2AllManager._buffer[0]
500 + return handle
501 +
502 + def dispatch(
503 + self,
504 + hidden_states: torch.Tensor,
505 + topk_weights: torch.Tensor,
506 + topk_ids: torch.Tensor,
507 + is_sequence_parallel: bool = False,
508 + extra_tensors: list[torch.Tensor] | None = None,
509 + ) -> (
510 + tuple[torch.Tensor, torch.Tensor, torch.Tensor]
511 + | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
512 + ):
513 + raise NotImplementedError
514 +
515 + def combine(
516 + self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
517 + ) -> torch.Tensor:
518 + raise NotImplementedError
519 +
520 + def destroy(self):
521 + # NOTE(yongji): NIXLEPAll2AllManager instance is recreated during
522 + # scale-up/down, so we cannot destroy the persistent buffer here.
523 + assert NixlEPAll2AllManager._buffer is not None
524 + buffer = NixlEPAll2AllManager._buffer[0]
525 + buffer.set_tcp_store_group(None)
526 +
527 + # NIXL EP uses RDMA so no SMs are used for communication
528 + def max_sms_used(self) -> int | None:
529 + return 0
530 +
531 +
532 class FlashInferAllToAllManager(All2AllManagerBase):
533 """
534 All2All communication based on flashinfer kernels.
vllm/distributed/device_communicators/cuda_communicator.py
+6
@@ -143,6 +143,12 @@ class CudaCommunicator(DeviceCommunicatorBase):
143 from .all2all import MoriAll2AllManager
144
145 self.all2all_manager = MoriAll2AllManager(self.cpu_group)
146 + elif self.all2all_backend == "nixl_ep":
147 + from .all2all import NixlEPAll2AllManager
148 +
149 + self.all2all_manager = NixlEPAll2AllManager(
150 + self.cpu_group, tcp_store_group
151 + )
152 elif self.all2all_backend == "flashinfer_all2allv":
153 from .all2all import FlashInferAllToAllManager
154
vllm/envs.py
+5
@@ -244,6 +244,7 @@ if TYPE_CHECKING:
244 VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: bool = False
245 VLLM_ELASTIC_EP_DRAIN_REQUESTS: bool = False
246 VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS: bool = False
247 + VLLM_NIXL_EP_MAX_NUM_RANKS: int = 32
248
249
250 def get_default_cache_root():
@@ -1628,6 +1629,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
1629 "VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS": lambda: bool(
1630 int(os.getenv("VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS", "0"))
1631 ),
1632 + # NIXL EP environment variables
1633 + "VLLM_NIXL_EP_MAX_NUM_RANKS": lambda: int(
1634 + os.getenv("VLLM_NIXL_EP_MAX_NUM_RANKS", "32")
1635 + ),
1636 }
1637
1638
vllm/model_executor/layers/fused_moe/all2all_utils.py
+46 -1
@@ -25,7 +25,7 @@ from vllm.model_executor.layers.fused_moe.prepare_finalize import (
25 make_moe_prepare_and_finalize_no_dp_ep,
26 )
27 from vllm.platforms import current_platform
28 -from vllm.utils.import_utils import has_deep_ep, has_mori
28 +from vllm.utils.import_utils import has_deep_ep, has_mori, has_nixl_ep
29
30 logger = init_logger(__name__)
31
@@ -38,6 +38,11 @@ if current_platform.is_cuda_alike():
38 )
39 if has_mori():
40 from .mori_prepare_finalize import MoriPrepareAndFinalize
41 + if has_nixl_ep():
42 + from .nixl_ep_prepare_finalize import (
43 + NIXL_EP_QUANT_BLOCK_SHAPE,
44 + NixlEPPrepareAndFinalize,
45 + )
46
47
48 def maybe_roundup_layer_hidden_size(
@@ -69,6 +74,11 @@ def maybe_roundup_layer_hidden_size(
74 hidden_size
75 )
76
77 + if moe_parallel_config.use_nixl_ep_kernels:
78 + hidden_size = NixlEPPrepareAndFinalize.maybe_roundup_layer_hidden_size(
79 + hidden_size
80 + )
81 +
82 return hidden_size
83
84
@@ -209,4 +219,39 @@ def maybe_make_prepare_finalize(
219 num_dispatchers=all2all_manager.world_size,
220 )
221
222 + elif moe.use_nixl_ep_kernels:
223 + assert quant_config is not None
224 + global_to_physical = physical_to_global = local_expert_global_ids = None
225 + if routing_tables is not None:
226 + (
227 + global_to_physical,
228 + physical_to_global,
229 + local_expert_global_ids,
230 + ) = routing_tables
231 + all_to_all_args = dict(
232 + max_num_tokens_per_dp_rank=moe.max_num_tokens,
233 + token_hidden_size=moe.hidden_dim,
234 + num_ep_ranks=all2all_manager.world_size,
235 + num_global_experts=moe.num_experts,
236 + num_local_experts=moe.num_experts // all2all_manager.world_size,
237 + )
238 + handle = all2all_manager.get_handle(all_to_all_args)
239 +
240 + # Note: We may want to use FP8 dispatch just to reduce
241 + # data movement.
242 + use_fp8_dispatch = (
243 + quant_config.quant_dtype == current_platform.fp8_dtype()
244 + and quant_config.block_shape == NIXL_EP_QUANT_BLOCK_SHAPE
245 + )
246 +
247 + prepare_finalize = NixlEPPrepareAndFinalize(
248 + handle,
249 + max_tokens_per_rank=moe.max_num_tokens,
250 + num_dispatchers=all2all_manager.world_size,
251 + use_fp8_dispatch=use_fp8_dispatch,
252 + global_to_physical=global_to_physical,
253 + physical_to_global=physical_to_global,
254 + local_expert_global_ids=local_expert_global_ids,
255 + )
256 +
257 return prepare_finalize
vllm/model_executor/layers/fused_moe/config.py
+8
@@ -976,6 +976,10 @@ class FusedMoEParallelConfig:
976 def use_mori_kernels(self):
977 return self.use_all2all_kernels and self.all2all_backend == "mori"
978
979 + @property
980 + def use_nixl_ep_kernels(self):
981 + return self.use_all2all_kernels and self.all2all_backend == "nixl_ep"
982 +
983 @staticmethod
984 def flatten_tp_across_dp_and_pcp(
985 tp_size: int, dp_size: int, dp_rank: int, pcp_size: int, pcp_rank: int
@@ -1242,3 +1246,7 @@ class FusedMoEConfig:
1246 @property
1247 def use_naive_all2all_kernels(self):
1248 return self.moe_parallel_config.use_naive_all2all_kernels
1249 +
1250 + @property
1251 + def use_nixl_ep_kernels(self):
1252 + return self.moe_parallel_config.use_nixl_ep_kernels
vllm/model_executor/layers/fused_moe/layer.py
+6 -5
@@ -177,10 +177,11 @@ def determine_expert_placement_strategy(
177 if (
178 moe_parallel_config.use_all2all_kernels
179 and not moe_parallel_config.use_deepep_ll_kernels
180 + and not moe_parallel_config.use_nixl_ep_kernels
181 ):
182 logger.warning(
183 "Round-robin expert placement currently only supports "
183 - "the DeepEP low-latency backend, but '%s' was configured. "
184 + "the DeepEP low-latency or NIXL EP backend, but '%s' was configured. "
185 "Falling back to linear expert placement.",
186 moe_parallel_config.all2all_backend,
187 )
@@ -745,10 +746,10 @@ class FusedMoE(CustomOp):
746 self,
747 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None:
748 # Currently routing_tables only needed for round-robin expert placement
748 - # with DeepEP-ll all2all backend.
749 - if (
750 - self.expert_placement_strategy != "round_robin"
751 - or not self.moe_parallel_config.use_deepep_ll_kernels
749 + # with DeepEP-ll or NIXL EP all2all backends.
750 + if self.expert_placement_strategy != "round_robin" or (
751 + not self.moe_parallel_config.use_deepep_ll_kernels
752 + and not self.moe_parallel_config.use_nixl_ep_kernels
753 ):
754 return None
755
vllm/model_executor/layers/fused_moe/nixl_ep_prepare_finalize.py new
+406
@@ -0,0 +1,406 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from collections.abc import Callable
4 +
5 +import nixl_ep
6 +import torch
7 +
8 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9 +from vllm import envs
10 +from vllm.logger import init_logger
11 +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
12 +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
13 + TopKWeightAndReduceDelegate,
14 +)
15 +from vllm.model_executor.layers.fused_moe.utils import (
16 + moe_kernel_quantize_input,
17 + normalize_batched_scales_shape,
18 +)
19 +from vllm.v1.worker.ubatching import (
20 + dbo_current_ubatch_id,
21 + dbo_enabled,
22 + dbo_maybe_run_recv_hook,
23 +)
24 +
25 +logger = init_logger(__name__)
26 +
27 +# NIXL EP kernels quantize dispatch inputs in 128 element chunks.
28 +NIXL_EP_QUANT_BLOCK_SIZE = 128
29 +NIXL_EP_QUANT_BLOCK_SHAPE = [NIXL_EP_QUANT_BLOCK_SIZE, NIXL_EP_QUANT_BLOCK_SIZE]
30 +
31 +
32 +def dequant_fp8(
33 + expert_x_fp8: torch.Tensor, expert_x_scales: torch.Tensor
34 +) -> torch.Tensor:
35 + """
36 + Return dequantized tensor in fp32
37 + """
38 + assert expert_x_fp8.is_contiguous()
39 + expert_x_scales = expert_x_scales.contiguous()
40 + num_experts = expert_x_fp8.size(0)
41 +
42 + expert_x_fp32 = expert_x_fp8.to(torch.float32).view(
43 + num_experts, -1, NIXL_EP_QUANT_BLOCK_SIZE
44 + )
45 + expert_x_scales = expert_x_scales.view(num_experts, -1, 1)
46 + return (expert_x_fp32 * expert_x_scales).view(expert_x_fp8.size())
47 +
48 +
49 +class NixlEPPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
50 + """
51 + Prepare/Finalize using NIXL EP kernels.
52 + """
53 +
54 + # NIXL EP kernels are compiled only for certain specific hidden sizes.
55 + # NOTE: Keep this list sorted, maybe_roundup_layer_hidden_size depends
56 + # on it.
57 + SUPPORTED_HIDDEN_SIZES = [2048, 2560, 3072, 4096, 5120, 6144, 7168, 8192]
58 + assert sorted(set(SUPPORTED_HIDDEN_SIZES)) == SUPPORTED_HIDDEN_SIZES
59 +
60 + @staticmethod
61 + def maybe_roundup_layer_hidden_size(hidden_size: int) -> int:
62 + # Round up hidden size to the closest supported hidden size.
63 + _supported_hs = NixlEPPrepareAndFinalize.SUPPORTED_HIDDEN_SIZES
64 +
65 + for x in _supported_hs:
66 + if x >= hidden_size:
67 + return x
68 +
69 + raise ValueError(
70 + f"Hidden Size {hidden_size} is greater than the "
71 + f"maximum supported hidden size {_supported_hs[-1]}"
72 + )
73 +
74 + def __init__(
75 + self,
76 + buffer: nixl_ep.Buffer,
77 + max_tokens_per_rank: int,
78 + num_dispatchers: int,
79 + use_fp8_dispatch: bool = False,
80 + global_to_physical: torch.Tensor | None = None,
81 + physical_to_global: torch.Tensor | None = None,
82 + local_expert_global_ids: torch.Tensor | None = None,
83 + ):
84 + super().__init__()
85 +
86 + self.buffer = buffer
87 + self.max_tokens_per_rank = max_tokens_per_rank
88 + self.use_fp8_dispatch = use_fp8_dispatch
89 + # The dispatch function returns a handle that the combine function
90 + # requires. We store the handle here so it is available to the
91 + # combine function.
92 + self.handles: list[tuple | None] = [None, None]
93 + self.num_dispatchers_ = num_dispatchers
94 +
95 + topk_indices_dtype = self.topk_indices_dtype()
96 +
97 + def _maybe_cast(tensor: torch.Tensor | None) -> torch.Tensor | None:
98 + if tensor is None or topk_indices_dtype is None:
99 + return tensor
100 + return tensor.to(dtype=topk_indices_dtype)
101 +
102 + self.global_to_physical = _maybe_cast(global_to_physical)
103 + self.physical_to_global = _maybe_cast(physical_to_global)
104 + self.local_expert_global_ids = _maybe_cast(local_expert_global_ids)
105 +
106 + # We don't have enough information to determine if we should dispatch
107 + # activation scales in a packed ue8m0 format during object construction
108 + # time. This setting is handled by post_init_setup.
109 + self.use_ue8m0_dispatch = False
110 +
111 + def post_init_setup(self, fused_experts: mk.FusedMoEExperts):
112 + if not fused_experts.supports_packed_ue8m0_act_scales():
113 + # Early exit.
114 + return
115 +
116 + if self.use_fp8_dispatch:
117 + logger.debug_once(
118 + "Update NixlEPPrepareAndFinalize to do packed ue8m0 scales dispatch."
119 + )
120 + self.use_ue8m0_dispatch = True
121 + else:
122 + logger.warning_once(
123 + "NixlEPPrepareAndFinalize is setup to dispatch raw/unquantized "
124 + f"activations despite ({fused_experts.__class__.__name__}) being able "
125 + "to support quantized activations.",
126 + scope="local",
127 + )
128 +
129 + def num_dispatchers(self) -> int:
130 + return self.num_dispatchers_
131 +
132 + def output_is_reduced(self) -> bool:
133 + return True
134 +
135 + @property
136 + def activation_format(self) -> mk.FusedMoEActivationFormat:
137 + return mk.FusedMoEActivationFormat.BatchedExperts
138 +
139 + def max_num_tokens_per_rank(self) -> int | None:
140 + return self.max_tokens_per_rank
141 +
142 + def topk_indices_dtype(self) -> torch.dtype | None:
143 + return torch.int64
144 +
145 + def _map_global_to_physical_ids(self, topk_ids: torch.Tensor) -> torch.Tensor:
146 + if self.global_to_physical is None:
147 + return topk_ids
148 + return self.global_to_physical[topk_ids]
149 +
150 + def _map_local_to_global_ids(self, expert_topk_ids: torch.Tensor) -> torch.Tensor:
151 + if self.local_expert_global_ids is None:
152 + return expert_topk_ids
153 + return self.local_expert_global_ids[expert_topk_ids]
154 +
155 + def _do_quant(
156 + self,
157 + x: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
158 + a1_dtype: torch.dtype,
159 + quant_config: FusedMoEQuantConfig,
160 + ) -> tuple[torch.Tensor, torch.Tensor | None]:
161 + if self.use_fp8_dispatch:
162 + block_k = (
163 + quant_config.block_shape[1]
164 + if quant_config.block_shape is not None
165 + else None
166 + )
167 + if block_k == NIXL_EP_QUANT_BLOCK_SIZE:
168 + # NIXL EP kernels did the quantization for us.
169 + x, x_scales = x
170 + return x, x_scales
171 +
172 + # Dequant to get back the tokens in the datatype we dispatched in.
173 + x_fp8, x_scales = x
174 + x = dequant_fp8(x_fp8, x_scales).to(dtype=a1_dtype)
175 +
176 + assert isinstance(x, torch.Tensor)
177 +
178 + num_experts, max_tokens, hidden_dim = x.size()
179 +
180 + x = x.view((-1, hidden_dim))
181 + q_dtype = quant_config.quant_dtype
182 +
183 + if envs.VLLM_FLASHINFER_MOE_BACKEND == "masked_gemm":
184 + logger.info_once(
185 + "Skip quantization when using FlashInfer CUTEDSL(masked_gemm) "
186 + "for ModelOptNvFp4FusedMoE."
187 + )
188 + q_dtype = None
189 +
190 + x, x_scales = moe_kernel_quantize_input(
191 + x,
192 + quant_config.a1_scale,
193 + q_dtype,
194 + quant_config.per_act_token_quant,
195 + quant_config.block_shape,
196 + )
197 + x = x.view((num_experts, -1, hidden_dim))
198 +
199 + if q_dtype is not None:
200 + assert x_scales is not None
201 + x_scales = normalize_batched_scales_shape(x_scales, num_experts)
202 +
203 + return x, x_scales
204 +
205 + def supports_async(self) -> bool:
206 + return True
207 +
208 + def prepare_async(
209 + self,
210 + a1: torch.Tensor,
211 + topk_weights: torch.Tensor,
212 + topk_ids: torch.Tensor,
213 + num_experts: int,
214 + expert_map: torch.Tensor | None,
215 + apply_router_weight_on_input: bool,
216 + quant_config: FusedMoEQuantConfig,
217 + defer_input_quant: bool = False,
218 + ) -> tuple[Callable, mk.ReceiverType]:
219 + if defer_input_quant:
220 + raise NotImplementedError(
221 + f"{self.__class__.__name__} does not support defer_input_quant=True. "
222 + "Please select an MoE kernel that accepts quantized inputs."
223 + )
224 +
225 + hidden_size = a1.size(1)
226 + assert hidden_size in self.SUPPORTED_HIDDEN_SIZES, (
227 + f"Hidden Size {hidden_size} not in supported list of hidden sizes"
228 + f"{self.SUPPORTED_HIDDEN_SIZES}"
229 + )
230 +
231 + a2a_idx = dbo_current_ubatch_id()
232 +
233 + if self.use_fp8_dispatch:
234 + assert hidden_size % 128 == 0, (
235 + "NIXL EP kernels quantize the inputs in blocks of shape 128"
236 + )
237 +
238 + has_per_token_scales = (
239 + quant_config.a1_scale.numel() != 1
240 + if quant_config.a1_scale is not None
241 + else (
242 + quant_config.a2_scale.numel() != 1
243 + if quant_config.a2_scale is not None
244 + else False
245 + )
246 + )
247 + assert not has_per_token_scales, (
248 + "NIXL EP kernels don't support dispatching per-token scales"
249 + )
250 +
251 + if apply_router_weight_on_input:
252 + topk = topk_ids.size(1)
253 + # TODO: this only works for topK=1, will need to update for topK>1
254 + assert topk == 1, (
255 + "apply_router_weight_on_input is only implemented for topk=1"
256 + )
257 + a1 = a1 * topk_weights.to(a1.dtype)
258 +
259 + # Dispatch
260 + dispatch_topk_ids = self._map_global_to_physical_ids(topk_ids)
261 + expert_x, expert_num_tokens, handle, _, hook = self.buffer.dispatch(
262 + a1,
263 + dispatch_topk_ids,
264 + self.max_tokens_per_rank,
265 + num_experts,
266 + use_fp8=self.use_fp8_dispatch,
267 + # round_scale needs to be set to dispatch in ue8m0
268 + round_scale=self.use_ue8m0_dispatch,
269 + use_ue8m0=self.use_ue8m0_dispatch,
270 + async_finish=False,
271 + return_recv_hook=True,
272 + )
273 + self.handles[a2a_idx] = handle
274 +
275 + return (
276 + hook,
277 + lambda: self._receiver(
278 + expert_x,
279 + expert_num_tokens,
280 + quant_config.a1_scale,
281 + a1.dtype,
282 + quant_config,
283 + ),
284 + )
285 +
286 + def _receiver(
287 + self,
288 + expert_x: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
289 + expert_num_tokens: torch.Tensor,
290 + a1_scale: torch.Tensor | None,
291 + a1_dtype: torch.dtype,
292 + quant_config: FusedMoEQuantConfig,
293 + ) -> mk.PrepareResultType:
294 + expert_x, expert_x_scale = self._do_quant(expert_x, a1_dtype, quant_config)
295 +
296 + expert_tokens_meta = mk.ExpertTokensMetadata(
297 + expert_num_tokens=expert_num_tokens, expert_num_tokens_cpu=None
298 + )
299 +
300 + return expert_x, expert_x_scale, expert_tokens_meta, None, None
301 +
302 + def prepare(
303 + self,
304 + a1: torch.Tensor,
305 + topk_weights: torch.Tensor,
306 + topk_ids: torch.Tensor,
307 + num_experts: int,
308 + expert_map: torch.Tensor | None,
309 + apply_router_weight_on_input: bool,
310 + quant_config: FusedMoEQuantConfig,
311 + defer_input_quant: bool = False,
312 + ) -> mk.PrepareResultType:
313 + if defer_input_quant:
314 + raise NotImplementedError(
315 + f"{self.__class__.__name__} does not support defer_input_quant=True. "
316 + "Please select an MoE kernel that accepts quantized inputs."
317 + )
318 + hook, receiver = self.prepare_async(
319 + a1,
320 + topk_weights,
321 + topk_ids,
322 + num_experts,
323 + expert_map,
324 + apply_router_weight_on_input,
325 + quant_config,
326 + )
327 + hook()
328 + return receiver()
329 +
330 + def _finalize(
331 + self,
332 + output: torch.Tensor,
333 + fused_expert_output: torch.Tensor,
334 + topk_weights: torch.Tensor,
335 + topk_ids: torch.Tensor,
336 + apply_router_weight_on_input: bool,
337 + weight_and_reduce_impl: mk.TopKWeightAndReduce,
338 + do_async: bool,
339 + ) -> tuple[Callable, Callable]:
340 + assert isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate), (
341 + "Weight application and reduction happens in the combine kernel."
342 + )
343 +
344 + a2a_idx = dbo_current_ubatch_id()
345 + do_recv_hook = dbo_enabled() or do_async
346 + handle = self.handles[a2a_idx]
347 + assert handle is not None
348 +
349 + combine_topk_weights = topk_weights
350 + if apply_router_weight_on_input:
351 + # weights have already been applied.
352 + combine_topk_weights = torch.ones_like(topk_weights)
353 +
354 + combine_topk_ids = self._map_global_to_physical_ids(topk_ids)
355 + # TODO (varun) : Enable zero copy mode
356 + dbo_maybe_run_recv_hook()
357 + _, _, recv_hook = self.buffer.combine(
358 + fused_expert_output,
359 + combine_topk_ids,
360 + combine_topk_weights,
361 + handle,
362 + async_finish=False,
363 + zero_copy=False,
364 + return_recv_hook=do_recv_hook,
365 + out=output,
366 + )
367 +
368 + return recv_hook, lambda: None
369 +
370 + def finalize_async(
371 + self,
372 + output: torch.Tensor,
373 + fused_expert_output: torch.Tensor,
374 + topk_weights: torch.Tensor,
375 + topk_ids: torch.Tensor,
376 + apply_router_weight_on_input: bool,
377 + weight_and_reduce_impl: mk.TopKWeightAndReduce,
378 + ) -> tuple[Callable, Callable]:
379 + return self._finalize(
380 + output,
381 + fused_expert_output,
382 + topk_weights,
383 + topk_ids,
384 + apply_router_weight_on_input,
385 + weight_and_reduce_impl,
386 + do_async=True,
387 + )
388 +
389 + def finalize(
390 + self,
391 + output: torch.Tensor,
392 + fused_expert_output: torch.Tensor,
393 + topk_weights: torch.Tensor,
394 + topk_ids: torch.Tensor,
395 + apply_router_weight_on_input: bool,
396 + weight_and_reduce_impl: mk.TopKWeightAndReduce,
397 + ) -> None:
398 + self._finalize(
399 + output,
400 + fused_expert_output,
401 + topk_weights,
402 + topk_ids,
403 + apply_router_weight_on_input,
404 + weight_and_reduce_impl,
405 + do_async=False,
406 + )
vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py
+1
@@ -234,6 +234,7 @@ class DefaultMoERunner(MoERunner):
234 self.moe_config.moe_parallel_config.use_deepep_ll_kernels
235 or self.moe_config.moe_parallel_config.use_mori_kernels
236 or self.moe_config.moe_parallel_config.use_fi_all2allv_kernels
237 + or self.moe_config.moe_parallel_config.use_nixl_ep_kernels
238 ) and envs.VLLM_ENABLE_MOE_DP_CHUNK
239
240 def _maybe_setup_shared_experts_stream(
vllm/model_executor/layers/quantization/mxfp4.py
+3 -1
@@ -896,7 +896,9 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
896 # batched activation format. As self.fused_experts is not
897 # initialized at this point, we resort to checking the MoE config
898 # directly.
899 - is_batched_moe = self.moe.use_deepep_ll_kernels
899 + is_batched_moe = (
900 + self.moe.use_deepep_ll_kernels or self.moe.use_nixl_ep_kernels
901 + )
902 if is_batched_moe:
903 num_warps = 4 if envs.VLLM_MOE_DP_CHUNK_SIZE <= 512 else 8
904 else:
vllm/utils/import_utils.py
+5
@@ -412,6 +412,11 @@ def has_deep_gemm() -> bool:
412 return _has_module("deep_gemm")
413
414
415 +def has_nixl_ep() -> bool:
416 + """Whether the optional `nixl_ep` package is available."""
417 + return _has_module("nixl_ep")
418 +
419 +
420 def has_triton_kernels() -> bool:
421 """Whether the optional `triton_kernels` package is available."""
422 is_available = _has_module("triton_kernels") or _has_module(
vllm/utils/network_utils.py
+14 -1
@@ -288,6 +288,7 @@ def make_zmq_socket(
288 bind: bool | None = None,
289 identity: bytes | None = None,
290 linger: int | None = None,
291 + router_handover: bool = False,
292 ) -> zmq.Socket | zmq.asyncio.Socket: # type: ignore[name-defined]
293 """Make a ZMQ socket with the proper bind/connect semantics."""
294
@@ -314,6 +315,10 @@ def make_zmq_socket(
315 socket.setsockopt(zmq.SNDHWM, 0)
316 socket.setsockopt(zmq.SNDBUF, buf_size)
317
318 + if socket_type == zmq.ROUTER and router_handover:
319 + # Let a new connection take over an identity left behind by a dead one.
320 + socket.setsockopt(zmq.ROUTER_HANDOVER, 1)
321 +
322 if identity is not None:
323 socket.setsockopt(zmq.IDENTITY, identity)
324
@@ -344,12 +349,20 @@ def zmq_socket_ctx(
349 bind: bool | None = None,
350 linger: int = 0,
351 identity: bytes | None = None,
352 + router_handover: bool = False,
353 ) -> Iterator[zmq.Socket]:
354 """Context manager for a ZMQ socket"""
355
356 ctx = zmq.Context() # type: ignore[attr-defined]
357 try:
352 - yield make_zmq_socket(ctx, path, socket_type, bind=bind, identity=identity)
358 + yield make_zmq_socket(
359 + ctx,
360 + path,
361 + socket_type,
362 + bind=bind,
363 + identity=identity,
364 + router_handover=router_handover,
365 + )
366 except KeyboardInterrupt:
367 logger.debug("Got Keyboard Interrupt.")
368
vllm/v1/engine/core_client.py
+15 -3
@@ -544,6 +544,11 @@ class MPClient(EngineCoreClient):
544 try:
545 # State used for data parallel.
546 self.engines_running = False
547 + parallel_config = vllm_config.parallel_config
548 + # Elastic EP can remove a rank and later add it back with the same
549 + # identity. The client input ROUTER needs handover to allow the new
550 + # engine to replace the dead connection.
551 + enable_input_socket_handover = parallel_config.enable_elastic_ep
552
553 self.stats_update_address: str | None = None
554 if client_addresses:
@@ -552,7 +557,11 @@ class MPClient(EngineCoreClient):
557 output_address = client_addresses["output_address"]
558 self.stats_update_address = client_addresses.get("stats_update_address")
559 self.input_socket = self.resources.input_socket = make_zmq_socket(
555 - self.ctx, input_address, zmq.ROUTER, bind=True
560 + self.ctx,
561 + input_address,
562 + zmq.ROUTER,
563 + bind=True,
564 + router_handover=enable_input_socket_handover,
565 )
566 self.resources.output_socket = make_zmq_socket(
567 self.ctx, output_address, zmq.PULL
@@ -561,7 +570,11 @@ class MPClient(EngineCoreClient):
570 # Engines are managed by this client.
571 addresses = get_engine_zmq_addresses(vllm_config)
572 self.input_socket = self.resources.input_socket = make_zmq_socket(
564 - self.ctx, addresses.inputs[0], zmq.ROUTER, bind=True
573 + self.ctx,
574 + addresses.inputs[0],
575 + zmq.ROUTER,
576 + bind=True,
577 + router_handover=enable_input_socket_handover,
578 )
579 self.resources.output_socket = make_zmq_socket(
580 self.ctx, addresses.outputs[0], zmq.PULL
@@ -582,7 +595,6 @@ class MPClient(EngineCoreClient):
595 coordinator.get_stats_publish_address()
596 )
597
585 - parallel_config = vllm_config.parallel_config
598 dp_size = parallel_config.data_parallel_size
599 dp_rank = parallel_config.data_parallel_index
600 dp_local_size = parallel_config.data_parallel_size_local