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

Refactor pplx init logic to make it modular (prepare for deepep) (#18200)

Signed-off-by: youkaichao <youkaichao@gmail.com>

youkaichao committed May 23, 2025 at 23:43 UTC 6a7988c55bb0a54240f479033313eb98344e6431
16 files changed +297 -284
vllm/distributed/device_communicators/all2all.py
+67 -34
@@ -1,44 +1,24 @@
1 # SPDX-License-Identifier: Apache-2.0
2 +import importlib.util
3 +from typing import TYPE_CHECKING
4 +
5 import torch
6 +import torch.distributed as dist
7
8 from vllm.forward_context import get_forward_context
9 +from vllm.logger import init_logger
10
11 +from .base_device_communicator import All2AllManagerBase, Cache
12
7 -class All2AllBase:
8 -
9 - def __init__(self, cpu_group, model):
10 - self.cpu_group = cpu_group
11 -
12 - # compute some common properties
13 - from vllm.distributed.parallel_state import (get_dp_group,
14 - get_ep_group,
15 - get_tp_group,
16 - in_the_same_node_as)
17 -
18 - # all2all lives in ep group, which is merged from dp and tp group
19 - self.dp_group = get_dp_group()
20 - self.tp_group = get_tp_group()
21 - self.ep_group = get_ep_group()
22 - self.dp_rank = self.dp_group.rank_in_group
23 - self.dp_world_size = self.dp_group.world_size
24 -
25 - # all2all communication often has separate implementations for
26 - # intra-node and inter-node communication
27 - self.intranode = in_the_same_node_as(cpu_group, source_rank=0)
28 - self.internode = not self.intranode
29 -
30 - def dispatch(self, hidden_states: torch.Tensor,
31 - router_logits: torch.Tensor):
32 - raise NotImplementedError
33 -
34 - def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
35 - raise NotImplementedError
13 +logger = init_logger(__name__)
14
37 - def destroy(self):
38 - pass
15 +if TYPE_CHECKING:
16 + from vllm.model_executor.layers.fused_moe.layer import FusedMoE
17 +else:
18 + FusedMoE = None
19
20
41 -class NaiveAll2All(All2AllBase):
21 +class NaiveAll2AllManager(All2AllManagerBase):
22 """
23 A naive implementation of all2all communication.
24 It uses all-reduce under the hood, which is not
@@ -46,8 +26,8 @@ class NaiveAll2All(All2AllBase):
26 debugging.
27 """
28
49 - def __init__(self, cpu_group, model):
50 - super().__init__(cpu_group, model)
29 + def __init__(self, cpu_group):
30 + super().__init__(cpu_group)
31
32 def naive_multicast(self, x: torch.Tensor,
33 cu_tokens_across_dp_cpu: torch.Tensor):
@@ -91,3 +71,56 @@ class NaiveAll2All(All2AllBase):
71
72 def destroy(self):
73 pass
74 +
75 +
76 +class PPLXAll2AllManager(All2AllManagerBase):
77 + """
78 + All2All communication based on PPLX kernels.
79 + """
80 +
81 + def __init__(self, cpu_group):
82 + has_pplx = importlib.util.find_spec("pplx_kernels") is not None
83 + assert has_pplx, "pplx_kernels not found. Please follow https://github.com/vllm-project/vllm/blob/main/tools/ep_kernels/README.md to install pplx_kernels." # noqa
84 + super().__init__(cpu_group)
85 +
86 + if self.internode:
87 + # inter-node communication needs nvshmem,
88 + # intra-node communication uses p2p mapping directly
89 + from pplx_kernels.nvshmem import (nvshmem_alloc_empty_unique_id,
90 + nvshmem_get_unique_id,
91 + nvshmem_init)
92 + logger.debug(
93 + "Initialize NVSHMEM for pplx_kernels: "
94 + "rank=%d, world size=%d", self.rank, self.world_size)
95 + uid = nvshmem_get_unique_id(
96 + ) if self.rank == 0 else nvshmem_alloc_empty_unique_id()
97 + dist.broadcast(uid,
98 + src=dist.get_process_group_ranks(self.cpu_group)[0],
99 + group=self.cpu_group)
100 + logger.debug("PPLX NVSHMEM UID = %s", uid)
101 + nvshmem_init(uid, self.rank, self.world_size)
102 +
103 + self.handle_cache = Cache()
104 +
105 + def get_handle(self, kwargs):
106 + import pplx_kernels as pplx
107 + return self.handle_cache.get_or_create(
108 + kwargs, pplx.AllToAll.internode
109 + if self.internode else pplx.AllToAll.intranode)
110 +
111 + def dispatch(self, hidden_states: torch.Tensor,
112 + router_logits: torch.Tensor):
113 + raise NotImplementedError
114 +
115 + def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
116 + raise NotImplementedError
117 +
118 + def destroy(self):
119 + with self.handle_cache._lock:
120 + for _, handle in self.handle_cache._cache.items():
121 + handle.destroy()
122 +
123 + if self.internode:
124 + from pplx_kernels.nvshmem import nvshmem_finalize
125 + logger.debug("PPLX NVSHMEM finalize")
126 + nvshmem_finalize()
vllm/distributed/device_communicators/base_device_communicator.py
+87 -2
@@ -1,11 +1,76 @@
1 # SPDX-License-Identifier: Apache-2.0
2 +import threading
3 from typing import Optional
4 +from weakref import WeakValueDictionary
5
6 import torch
7 import torch.distributed as dist
8 from torch.distributed import ProcessGroup
9
10
11 +class Cache:
12 +
13 + def __init__(self):
14 + self._cache: WeakValueDictionary = WeakValueDictionary()
15 + self._lock = threading.RLock() # Reentrant lock for thread safety
16 +
17 + def get_or_create(self, kwargs, func):
18 + # Create a hashable key from the kwargs
19 + key = tuple(sorted((k, v) for k, v in kwargs.items()))
20 +
21 + with self._lock:
22 + instance = self._cache.get(key)
23 + if instance is None:
24 + instance = func(**kwargs)
25 + self._cache[key] = instance
26 + return instance
27 +
28 +
29 +class All2AllManagerBase:
30 +
31 + def __init__(self, cpu_group):
32 + self.cpu_group = cpu_group
33 +
34 + # compute some common properties
35 + from vllm.distributed.parallel_state import (get_dp_group,
36 + get_tp_group,
37 + in_the_same_node_as)
38 +
39 + # all2all lives in ep group, which is merged from dp and tp group
40 + self.dp_group = get_dp_group()
41 + self.tp_group = get_tp_group()
42 + # no self.ep_group since self.ep_group is still in construction
43 + # when we create this object
44 + self.dp_rank = self.dp_group.rank_in_group
45 + self.dp_world_size = self.dp_group.world_size
46 + self.rank = dist.get_rank(cpu_group)
47 + self.world_size = dist.get_world_size(cpu_group)
48 +
49 + # all2all communication often has separate implementations for
50 + # intra-node and inter-node communication
51 + self.intranode = in_the_same_node_as(cpu_group, source_rank=0)
52 + self.internode = not self.intranode
53 +
54 + def get_handle(self, kwargs):
55 + # get a handle for the all2all communication,
56 + # based on the kwargs.
57 + # different layers can have different configs,
58 + # e.g. one layer has hidden size 1024, another has 2048.
59 + # usually the underlying implementation caches the handle
60 + # and reuse it for the same config.
61 + raise NotImplementedError
62 +
63 + def dispatch(self, hidden_states: torch.Tensor,
64 + router_logits: torch.Tensor):
65 + raise NotImplementedError
66 +
67 + def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
68 + raise NotImplementedError
69 +
70 + def destroy(self):
71 + pass
72 +
73 +
74 class DeviceCommunicatorBase:
75 """
76 Base class for device-specific communicator.
@@ -31,6 +96,18 @@ class DeviceCommunicatorBase:
96 self.rank_in_group = dist.get_group_rank(self.cpu_group,
97 self.global_rank)
98
99 + use_ep = False
100 + from vllm.config import get_current_vllm_config
101 + config = get_current_vllm_config()
102 + if config is not None:
103 + # as long as we use data parallel (coupled data parallel
104 + # where all data parallel ranks execute forward together),
105 + # we initialize the all2all manager used in expert parallel.
106 + use_ep = config.parallel_config.data_parallel_size > 1
107 +
108 + self.use_all2all = "ep" in unique_name and use_ep
109 + self.all2all_manager: Optional[All2AllManagerBase] = None
110 +
111 def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
112 dist.all_reduce(input_, group=self.device_group)
113 return input_
@@ -154,9 +231,17 @@ class DeviceCommunicatorBase:
231 model: torch.nn.Module) -> None:
232 """
233 Prepare the communication buffer for the model.
157 - This is a no-op in the base class.
234 """
159 - pass
235 + if not self.use_all2all:
236 + return
237 +
238 + moe_modules = [
239 + module for module in model.modules()
240 + if module.__class__.__name__ == "FusedMoE"
241 + ]
242 + for module in moe_modules:
243 + module.quant_method.init_prepare_finalize(module.moe_config,
244 + module.quant_config)
245
246 def dispatch(
247 self, hidden_states: torch.Tensor,
vllm/distributed/device_communicators/cuda_communicator.py
+23 -22
@@ -6,10 +6,12 @@ import torch
6 from torch.distributed import ProcessGroup
7
8 import vllm.envs as envs
9 +from vllm.logger import init_logger
10
10 -from .all2all import All2AllBase
11 from .base_device_communicator import DeviceCommunicatorBase
12
13 +logger = init_logger(__name__)
14 +
15
16 class CudaCommunicator(DeviceCommunicatorBase):
17
@@ -31,8 +33,6 @@ class CudaCommunicator(DeviceCommunicatorBase):
33 use_pynccl = "ep" not in unique_name
34
35 self.use_pynccl = use_pynccl
34 - self.use_all2all = "ep" in unique_name
35 - self.all2all_impl: Optional[All2AllBase] = None
36 self.use_custom_allreduce = use_custom_allreduce
37
38 # lazy import to avoid documentation build error
@@ -56,6 +56,19 @@ class CudaCommunicator(DeviceCommunicatorBase):
56 device=self.device,
57 )
58
59 + if self.use_all2all:
60 + all2all_backend = envs.VLLM_ALL2ALL_BACKEND
61 + if all2all_backend == "naive":
62 + from .all2all import NaiveAll2AllManager
63 + self.all2all_manager = NaiveAll2AllManager(self.cpu_group)
64 + logger.info("Using naive all2all manager.")
65 + elif all2all_backend == "pplx":
66 + from .all2all import PPLXAll2AllManager
67 + self.all2all_manager = PPLXAll2AllManager(self.cpu_group)
68 + logger.info("Using PPLX all2all manager.")
69 + else:
70 + raise ValueError(f"Unknown all2all backend: {all2all_backend}")
71 +
72 def all_reduce(self, input_):
73 # always try custom allreduce first,
74 # and then pynccl.
@@ -136,31 +149,19 @@ class CudaCommunicator(DeviceCommunicatorBase):
149 self.pynccl_comm = None
150 if self.ca_comm is not None:
151 self.ca_comm = None
139 - if self.all2all_impl is not None:
140 - self.all2all_impl.destroy()
141 - self.all2all_impl = None
142 -
143 - def prepare_communication_buffer_for_model(self,
144 - model: torch.nn.Module) -> None:
145 - """
146 - Prepare the communication buffer for the model.
147 - """
148 - if not self.use_all2all:
149 - return
150 - all2all_backend = envs.VLLM_ALL2ALL_BACKEND
151 - if all2all_backend == "naive":
152 - from .all2all import NaiveAll2All
153 - self.all2all_impl = NaiveAll2All(self.cpu_group, model)
152 + if self.all2all_manager is not None:
153 + self.all2all_manager.destroy()
154 + self.all2all_manager = None
155
156 def dispatch(
157 self, hidden_states: torch.Tensor,
158 router_logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
158 - assert self.all2all_impl is not None
159 - hidden_states, router_logits = self.all2all_impl.dispatch(
159 + assert self.all2all_manager is not None
160 + hidden_states, router_logits = self.all2all_manager.dispatch(
161 hidden_states, router_logits)
162 return hidden_states, router_logits
163
164 def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
164 - assert self.all2all_impl is not None
165 - hidden_states = self.all2all_impl.combine(hidden_states)
165 + assert self.all2all_manager is not None
166 + hidden_states = self.all2all_manager.combine(hidden_states)
167 return hidden_states
vllm/distributed/parallel_state.py
+6 -50
@@ -23,7 +23,6 @@ If you only need to use the distributed environment without model/pipeline
23 """
24 import contextlib
25 import gc
26 -import importlib.util
26 import pickle
27 import weakref
28 from collections import namedtuple
@@ -43,7 +42,7 @@ from vllm.distributed.device_communicators.base_device_communicator import (
42 from vllm.distributed.utils import StatelessProcessGroup
43 from vllm.logger import init_logger
44 from vllm.utils import (direct_register_custom_op, resolve_obj_by_qualname,
46 - run_once, supports_custom_op)
45 + supports_custom_op)
46
47
48 @dataclass
@@ -791,10 +790,14 @@ class GroupCoordinator:
790 if self.device_communicator is not None:
791 return self.device_communicator.dispatch(hidden_states,
792 router_logits)
793 + else:
794 + return hidden_states, router_logits
795
796 def combine(self, hidden_states) -> torch.Tensor:
797 if self.device_communicator is not None:
798 return self.device_communicator.combine(hidden_states)
799 + else:
800 + return hidden_states
801
802
803 _WORLD: Optional[GroupCoordinator] = None
@@ -959,49 +962,9 @@ def init_distributed_environment(
962 "world group already initialized with a different world size")
963
964
962 -PPLX_DID_INIT: bool = False
963 -
964 -
965 -@run_once
966 -def pplx_init(rank, world_size):
967 - has_pplx = importlib.util.find_spec("pplx_kernels") is not None
968 -
969 - if has_pplx and world_size > 1:
970 - from pplx_kernels.nvshmem import (nvshmem_alloc_empty_unique_id,
971 - nvshmem_get_unique_id, nvshmem_init)
972 - try:
973 - global PPLX_DID_INIT
974 - logger.debug(
975 - "Initialize NVSHMEM for PPLX kernels: rank=%d, "
976 - "world size=%d", rank, world_size)
977 - uid = nvshmem_get_unique_id(
978 - ) if rank == 0 else nvshmem_alloc_empty_unique_id()
979 - uid_gpu = uid.cuda()
980 - get_world_group().broadcast(uid_gpu, src=0)
981 - uid = uid_gpu.to(device='cpu')
982 - logger.debug("PPLX NVSHMEM UID = %s", uid)
983 - nvshmem_init(uid, rank, world_size)
984 - PPLX_DID_INIT = True
985 - except Exception as ex:
986 - logger.error("Failed to initialize NVSHMEM for PPLX: %s", ex)
987 -
988 -
989 -@run_once
990 -def pplx_finalize():
991 - global PPLX_DID_INIT
992 - if PPLX_DID_INIT:
993 - from pplx_kernels.nvshmem import nvshmem_finalize
994 - logger.debug("PPLX NVSHMEM finalize")
995 - from vllm.model_executor.layers.fused_moe.layer import (
996 - _all_to_all_cache)
997 - _all_to_all_cache.destroy()
998 - nvshmem_finalize()
999 -
1000 -
965 def initialize_model_parallel(
966 tensor_model_parallel_size: int = 1,
967 pipeline_model_parallel_size: int = 1,
1004 - enable_expert_parallel: bool = False,
968 backend: Optional[str] = None,
969 ) -> None:
970 """
@@ -1104,14 +1067,10 @@ def initialize_model_parallel(
1067 _DP.rank_in_group, _PP.rank_in_group, _TP.rank_in_group,
1068 _EP.rank_in_group)
1069
1107 - if enable_expert_parallel:
1108 - pplx_init(rank, world_size)
1109 -
1070
1071 def ensure_model_parallel_initialized(
1072 tensor_model_parallel_size: int,
1073 pipeline_model_parallel_size: int,
1114 - enable_expert_parallel: bool = False,
1074 backend: Optional[str] = None,
1075 ) -> None:
1076 """Helper to initialize model parallel groups if they are not initialized,
@@ -1122,8 +1081,7 @@ def ensure_model_parallel_initialized(
1081 get_world_group().device_group)
1082 if not model_parallel_is_initialized():
1083 initialize_model_parallel(tensor_model_parallel_size,
1125 - pipeline_model_parallel_size,
1126 - enable_expert_parallel, backend)
1084 + pipeline_model_parallel_size, backend)
1085 return
1086
1087 assert (
@@ -1202,8 +1160,6 @@ def destroy_model_parallel():
1160 """Set the groups to none and destroy them."""
1161 global _TP
1162
1205 - pplx_finalize()
1206 -
1163 if _TP:
1164 _TP.destroy()
1165 _TP = None
vllm/envs.py
+3
@@ -809,6 +809,9 @@ environment_variables: dict[str, Callable[[], Any]] = {
809 lambda: int(os.getenv("VLLM_NIXL_SIDE_CHANNEL_PORT", "5557")),
810
811 # all2all backend for vllm's expert parallel communication
812 + # Available options:
813 + # - "naive": naive all2all implementation using all-reduce
814 + # - "pplx": use pplx kernels
815 "VLLM_ALL2ALL_BACKEND":
816 lambda: os.getenv("VLLM_ALL2ALL_BACKEND", "naive"),
817 }
vllm/model_executor/layers/fused_moe/layer.py
+97 -143
@@ -1,12 +1,10 @@
1 # SPDX-License-Identifier: Apache-2.0
2
3 import importlib
4 -import threading
4 from abc import abstractmethod
5 from dataclasses import dataclass
6 from enum import Enum
7 from typing import Callable, Optional
9 -from weakref import WeakValueDictionary
8
9 import torch
10 import torch.nn.functional as F
@@ -73,7 +71,8 @@ class FusedMoEParallelConfig:
71
72 @property
73 def use_pplx_kernels(self):
76 - return self.dp_size > 1 and self.use_ep and has_pplx
74 + return self.dp_size > 1 and self.use_ep and \
75 + envs.VLLM_ALL2ALL_BACKEND == "pplx"
76
77 @staticmethod
78 def make(tp_size_: int, dp_size_: int,
@@ -196,6 +195,8 @@ class MoEConfig:
195 # TODO: add more quantization params, blocked, per-token, etc.
196 block_size: int = 128
197
198 + max_num_tokens: int = MOE_DP_CHUNK_SIZE
199 +
200 @property
201 def tp_size(self):
202 return self.moe_parallel_config.tp_size
@@ -244,13 +245,59 @@ class FusedMoEMethodBase(QuantizeMethodBase):
245 params_dtype: torch.dtype, **extra_weight_attrs):
246 raise NotImplementedError
247
247 - def set_prepare_finalize(
248 - self,
249 - dp_size: int,
250 - world_size: int,
251 - prepare_finalize: FusedMoEPrepareAndFinalize,
252 - ) -> bool:
253 - return False
248 + def init_prepare_finalize(self, moe: MoEConfig,
249 + quant_config: Optional[QuantizationConfig]):
250 + all2all_manager = get_ep_group().device_communicator.all2all_manager
251 + assert all2all_manager is not None
252 +
253 + prepare_finalize = None
254 + if moe.use_pplx_kernels:
255 + all_to_all_args = dict(
256 + max_num_tokens=moe.max_num_tokens,
257 + num_experts=moe.num_experts,
258 + experts_per_token=moe.experts_per_token, # topk
259 + rank=all2all_manager.rank,
260 + world_size=all2all_manager.world_size,
261 + # dp_size actually means tp_size, bug in pplx kernels
262 + dp_size=all2all_manager.tp_group.world_size,
263 + hidden_dim=moe.hidden_dim,
264 + hidden_dim_bytes=moe.hidden_dim * moe.in_dtype.itemsize,
265 + # For blocked per token: set to
266 + # ceil_div(hidden_dim, block_size) * sizeof(float32)
267 + # For per-token: set to sizeof(float32)
268 + hidden_dim_scale_bytes=(0 if moe.in_dtype.itemsize != 1 else (
269 + (moe.hidden_dim + moe.block_size - 1) // moe.block_size *
270 + torch.float32.itemsize)),
271 + group_name=all2all_manager.cpu_group.group_name,
272 + )
273 +
274 + handle = all2all_manager.get_handle(all_to_all_args)
275 +
276 + prepare_finalize = PplxPrepareAndFinalize(
277 + handle,
278 + max_num_tokens=moe.max_num_tokens,
279 + world_size=all2all_manager.world_size,
280 + rank=all2all_manager.rank,
281 + # dp_size actually means tp_size, bug in pplx kernels
282 + dp_size=all2all_manager.tp_group.world_size,
283 + quant_dtype=moe.in_dtype,
284 + )
285 +
286 + if prepare_finalize is not None:
287 + experts = self.select_gemm_impl(prepare_finalize)
288 + self.fused_experts = FusedMoEModularKernel(
289 + prepare_finalize,
290 + experts,
291 + )
292 +
293 + def select_gemm_impl(
294 + self, prepare_finalize: Optional[FusedMoEPrepareAndFinalize]
295 + ) -> FusedMoEPermuteExpertsUnpermute:
296 + # based on the all2all implementation, select the appropriate
297 + # gemm implementation
298 + raise NotImplementedError(
299 + "Subclass must select appropriate gemm implementation"
300 + " based on the prepare_finalize")
301
302 @abstractmethod
303 def apply(
@@ -274,53 +321,13 @@ class FusedMoEMethodBase(QuantizeMethodBase):
321 raise NotImplementedError
322
323
277 -class AllToAllCache:
278 -
279 - def __init__(self):
280 - self._cache: WeakValueDictionary = WeakValueDictionary()
281 - self._lock = threading.RLock() # Reentrant lock for thread safety
282 -
283 - def destroy(self):
284 - with self._lock:
285 - # TODO: can we do del self._cache?
286 - for _, a2a in self._cache.items():
287 - a2a.destroy()
288 -
289 - def get_or_create(self, **kwargs):
290 - assert has_pplx
291 - import pplx_kernels as pplx
292 -
293 - # Create a hashable key from the kwargs
294 - key = tuple(sorted((k, v) for k, v in kwargs.items()))
295 -
296 - with self._lock:
297 - instance = self._cache.get(key)
298 - if instance is None:
299 - # TODO (varun): Add support to switch to intranode
300 - # when all communications are within the same
301 - # node.
302 - logger.debug("Create AllToAll %s", kwargs)
303 - instance = pplx.AllToAll.internode(**kwargs)
304 - self._cache[key] = instance
305 - return instance
306 -
307 -
308 -# Global singleton
309 -_all_to_all_cache = AllToAllCache()
310 -
311 -
312 -# Factory function as a cleaner interface
313 -def get_all_to_all(**kwargs):
314 - return _all_to_all_cache.get_or_create(**kwargs)
315 -
316 -
324 @CustomOp.register("unquantized_fused_moe")
325 class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
326 """MoE method without quantization."""
327
328 def __init__(self, moe: MoEConfig):
329 super().__init__()
323 - self.fused_experts = fused_experts
330 + self.fused_experts = fused_experts # type: ignore
331 self.moe = moe
332
333 self.rocm_aiter_moe_enabled = is_rocm_aiter_moe_enabled()
@@ -330,6 +337,42 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
337 else:
338 self.rocm_aiter_fused_experts = None # type: ignore
339
340 + def select_gemm_impl(
341 + self, prepare_finalize: Optional[FusedMoEPrepareAndFinalize]):
342 +
343 + assert self.fused_experts == fused_experts
344 +
345 + all2all_manager = get_ep_group().device_communicator.all2all_manager
346 + assert all2all_manager is not None
347 +
348 + experts: Optional[FusedMoEPermuteExpertsUnpermute] = None
349 +
350 + if isinstance(prepare_finalize,
351 + (BatchedPrepareAndFinalize, PplxPrepareAndFinalize)):
352 + logger.debug("BatchedTritonExperts %s", self.moe)
353 + experts = BatchedTritonExperts(
354 + max_num_tokens=MOE_DP_CHUNK_SIZE,
355 + world_size=all2all_manager.world_size,
356 + # dp_size actually means tp_size, bug in pplx kernels
357 + dp_size=all2all_manager.tp_group.world_size,
358 + use_fp8_w8a8=False,
359 + use_int8_w8a8=False,
360 + use_int8_w8a16=False,
361 + use_int4_w4a16=False,
362 + block_shape=None,
363 + )
364 + else:
365 + logger.debug("TritonExperts %s", self.moe)
366 + experts = TritonExperts(
367 + use_fp8_w8a8=False,
368 + use_int8_w8a8=False,
369 + use_int8_w8a16=False,
370 + use_int4_w4a16=False,
371 + block_shape=None,
372 + per_channel_quant=False,
373 + )
374 + return experts
375 +
376 def create_weights(self, layer: torch.nn.Module, num_experts: int,
377 hidden_size: int, intermediate_size_per_partition: int,
378 params_dtype: torch.dtype, **extra_weight_attrs):
@@ -429,47 +472,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
472 activation=activation,
473 apply_router_weight_on_input=apply_router_weight_on_input)
474
432 - def set_prepare_finalize(
433 - self,
434 - dp_size: int,
435 - world_size: int,
436 - prepare_finalize: FusedMoEPrepareAndFinalize,
437 - ) -> bool:
438 - assert self.fused_experts == fused_experts
439 -
440 - experts: Optional[FusedMoEPermuteExpertsUnpermute] = None
441 -
442 - if isinstance(prepare_finalize,
443 - (BatchedPrepareAndFinalize, PplxPrepareAndFinalize)):
444 - logger.debug("BatchedTritonExperts %s", self.moe)
445 - experts = BatchedTritonExperts(
446 - max_num_tokens=MOE_DP_CHUNK_SIZE,
447 - world_size=world_size,
448 - dp_size=dp_size,
449 - use_fp8_w8a8=False,
450 - use_int8_w8a8=False,
451 - use_int8_w8a16=False,
452 - use_int4_w4a16=False,
453 - block_shape=None,
454 - )
455 - else:
456 - logger.debug("TritonExperts %s", self.moe)
457 - experts = TritonExperts(
458 - use_fp8_w8a8=False,
459 - use_int8_w8a8=False,
460 - use_int8_w8a16=False,
461 - use_int4_w4a16=False,
462 - block_shape=None,
463 - per_channel_quant=False,
464 - )
465 -
466 - self.fused_experts = FusedMoEModularKernel(
467 - prepare_finalize,
468 - experts,
469 - )
470 -
471 - return True
472 -
475 def forward_cuda(
476 self,
477 layer: torch.nn.Module,
@@ -679,45 +681,6 @@ def determine_expert_map(
681 return (local_num_experts, expert_map)
682
683
682 -def _construct_prepare_finalize(
683 - moe: MoEConfig, quant_config: Optional[QuantizationConfig]
684 -) -> Optional[FusedMoEPrepareAndFinalize]:
685 - max_num_tokens = MOE_DP_CHUNK_SIZE
686 - world_size = moe.ep_size
687 - dp_size = moe.ep_size // moe.dp_size # dp_size actually means TP.
688 - rank = moe.ep_rank
689 -
690 - if moe.use_pplx_kernels:
691 - logger.debug("using PplxPrepareAndFinalize")
692 -
693 - all_to_all = get_all_to_all(
694 - max_num_tokens=max_num_tokens,
695 - num_experts=moe.num_experts,
696 - experts_per_token=moe.experts_per_token, # topk
697 - rank=rank,
698 - world_size=world_size,
699 - dp_size=dp_size,
700 - hidden_dim=moe.hidden_dim,
701 - hidden_dim_bytes=moe.hidden_dim * moe.in_dtype.itemsize,
702 - # For blocked per token: set to
703 - # ceil_div(hidden_dim, block_size) * sizeof(float32)
704 - # For per-token: set to sizeof(float32)
705 - hidden_dim_scale_bytes=(0 if moe.in_dtype.itemsize != 1 else
706 - ((moe.hidden_dim + moe.block_size - 1) //
707 - moe.block_size * torch.float32.itemsize)))
708 -
709 - return PplxPrepareAndFinalize(
710 - all_to_all,
711 - max_num_tokens=max_num_tokens,
712 - world_size=world_size,
713 - rank=rank,
714 - dp_size=dp_size,
715 - quant_dtype=moe.in_dtype,
716 - )
717 -
718 - return None
719 -
720 -
684 class FusedMoE(torch.nn.Module):
685 """FusedMoE layer for MoE models.
686
@@ -831,7 +794,10 @@ class FusedMoE(torch.nn.Module):
794 moe_parallel_config=self.moe_parallel_config,
795 # TODO (bnell): this needs to be fixed for quantized types.
796 in_dtype=params_dtype,
797 + max_num_tokens=MOE_DP_CHUNK_SIZE,
798 )
799 + self.moe_config = moe
800 + self.quant_config = quant_config
801
802 # Note: get_quant_method will look at the layer's local_num_experts
803 # for heuristic purposes, so it must be initialized first.
@@ -839,25 +805,13 @@ class FusedMoE(torch.nn.Module):
805
806 if quant_config is None:
807 quant_method = UnquantizedFusedMoEMethod(moe)
842 - prepare_finalize = _construct_prepare_finalize(moe, quant_config)
808 else:
809 quant_method = quant_config.get_quant_method(self, prefix)
845 - # No pplx for quantized types yet.
846 - prepare_finalize = None
810
811 assert quant_method is not None
812 assert isinstance(quant_method, FusedMoEMethodBase)
813 self.quant_method = quant_method
814
852 - if prepare_finalize is not None:
853 - world_size = moe.ep_size
854 - dp_size = int(moe.ep_size // moe.dp_size)
855 - success = self.quant_method.set_prepare_finalize(
856 - dp_size, world_size, prepare_finalize)
857 - if not success:
858 - logger.warning("DP+EP not supported for %s.",
859 - type(self.quant_method))
860 -
815 moe_quant_params = {
816 "num_experts": self.local_num_experts,
817 "hidden_size": hidden_size,
vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py
-1
@@ -9,7 +9,6 @@ from vllm.model_executor.layers.fused_moe.utils import (
9 moe_kernel_quantize_input)
10
11
12 -# Note use: layer.get_all_to_all() to get an AllToAll instance
12 # The max_num_tokens, world_size and dp_size must be the same
13 # as the ones used to create the AllToAll.
14 class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
vllm/model_executor/layers/quantization/fp8.py
+5 -16
@@ -10,7 +10,6 @@ from torch.nn import Module
10 from torch.nn.parameter import Parameter
11
12 import vllm.envs as envs
13 -import vllm.model_executor.layers.fused_moe.modular_kernel as mk
13 from vllm import _custom_ops as ops
14 from vllm.distributed import get_tensor_model_parallel_world_size
15 from vllm.logger import init_logger
@@ -461,7 +460,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
460 logger.warning_once(
461 "DeepGemm not supported on the current platform.")
462
464 - self.fused_experts = functools.partial(
463 + self.fused_experts = functools.partial( # type: ignore
464 fused_experts,
465 block_shape=self.quant_config.weight_block_size,
466 allow_deep_gemm=self.allow_deep_gemm)
@@ -791,17 +790,12 @@ class Fp8MoEMethod(FusedMoEMethodBase):
790 del layer.w13_input_scale
791 del layer.w2_input_scale
792
794 - def set_prepare_finalize(
795 - self,
796 - dp_size: int,
797 - world_size: int,
798 - prepare_finalize: mk.FusedMoEPrepareAndFinalize,
799 - ) -> bool:
793 + def select_gemm_impl(self, prepare_finalize):
794 from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
795 TritonOrDeepGemmExperts)
796
803 - if self.use_marlin or self.rocm_aiter_moe_enabled:
804 - return False
797 + assert not self.use_marlin and not self.rocm_aiter_moe_enabled, (
798 + "Marlin and ROCm AITER are not supported with all2all yet.")
799
800 experts = TritonOrDeepGemmExperts(
801 use_fp8_w8a8=True,
@@ -809,12 +803,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
803 allow_deep_gemm=self.allow_deep_gemm,
804 )
805
812 - self.fused_experts = mk.FusedMoEModularKernel(
813 - prepare_finalize,
814 - experts,
815 - )
816 -
817 - return True
806 + return experts
807
808 def apply(
809 self,
vllm/platforms/cuda.py
+1
@@ -158,6 +158,7 @@ class CudaPlatformBase(Platform):
158 "currently not supported with CUDA Graphs.")
159 vllm_config.model_config.enforce_eager = True
160 compilation_config.use_cudagraph = False
161 + # FIXME: inductor breaks cudagraph (from @bnell)
162 compilation_config.use_inductor = False
163
164 @classmethod
vllm/v1/worker/gpu_worker.py
+1 -2
@@ -348,8 +348,7 @@ def init_worker_distributed_environment(
348 distributed_init_method, local_rank)
349
350 ensure_model_parallel_initialized(parallel_config.tensor_parallel_size,
351 - parallel_config.pipeline_parallel_size,
352 - parallel_config.enable_expert_parallel)
351 + parallel_config.pipeline_parallel_size)
352
353 ensure_kv_transfer_initialized(vllm_config)
354
vllm/v1/worker/tpu_worker.py
+1 -2
@@ -265,8 +265,7 @@ def init_tpu_worker_distributed_environment(
265 backend="gloo",
266 )
267 ensure_model_parallel_initialized(parallel_config.tensor_parallel_size,
268 - parallel_config.pipeline_parallel_size,
269 - parallel_config.enable_expert_parallel)
268 + parallel_config.pipeline_parallel_size)
269
270
271 try:
vllm/worker/cpu_worker.py
+1 -2
@@ -390,8 +390,7 @@ class CPUWorker(LocalOrDistributedWorkerBase):
390
391 ensure_model_parallel_initialized(
392 parallel_config.tensor_parallel_size,
393 - parallel_config.pipeline_parallel_size,
394 - parallel_config.enable_expert_parallel)
393 + parallel_config.pipeline_parallel_size)
394
395 def get_cache_block_size_bytes(self) -> int:
396 """Return the size in bytes of a single KV cache block.
vllm/worker/hpu_worker.py
+2 -4
@@ -415,8 +415,7 @@ def init_worker_distributed_environment(
415 backend='hccl')
416
417 ensure_model_parallel_initialized(parallel_config.tensor_parallel_size,
418 - parallel_config.pipeline_parallel_size,
419 - parallel_config.enable_expert_parallel)
418 + parallel_config.pipeline_parallel_size)
419
420 if torch.distributed.is_initialized():
421 torch_world_size = torch.distributed.get_world_size()
@@ -442,8 +441,7 @@ def init_worker_distributed_environment(
441 torch.distributed.all_reduce(dummy_tensor_hpu)
442 assert dummy_tensor_hpu.item() == parallel_config.world_size
443 ensure_model_parallel_initialized(parallel_config.tensor_parallel_size,
445 - parallel_config.pipeline_parallel_size,
446 - parallel_config.enable_expert_parallel)
444 + parallel_config.pipeline_parallel_size)
445
446
447 def raise_if_cache_size_invalid(num_gpu_blocks, block_size, max_model_len,
vllm/worker/tpu_worker.py
+1 -2
@@ -76,8 +76,7 @@ class TPUWorker(LoRANotSupportedWorkerBase, LocalOrDistributedWorkerBase):
76 )
77 ensure_model_parallel_initialized(
78 self.parallel_config.tensor_parallel_size,
79 - self.parallel_config.pipeline_parallel_size,
80 - self.parallel_config.enable_expert_parallel)
79 + self.parallel_config.pipeline_parallel_size)
80
81 # Device initialization should happen after initializing the distributed
82 # runtime.
vllm/worker/worker.py
+1 -2
@@ -529,8 +529,7 @@ def init_worker_distributed_environment(
529 init_distributed_environment(parallel_config.world_size, rank,
530 distributed_init_method, local_rank)
531 ensure_model_parallel_initialized(parallel_config.tensor_parallel_size,
532 - parallel_config.pipeline_parallel_size,
533 - parallel_config.enable_expert_parallel)
532 + parallel_config.pipeline_parallel_size)
533
534 ensure_kv_transfer_initialized(vllm_config)
535
vllm/worker/xpu_worker.py
+1 -2
@@ -175,8 +175,7 @@ class XPUWorker(LoRANotSupportedWorkerBase, Worker):
175
176 ensure_model_parallel_initialized(
177 parallel_config.tensor_parallel_size,
178 - parallel_config.pipeline_parallel_size,
179 - parallel_config.enable_expert_parallel)
178 + parallel_config.pipeline_parallel_size)
179 # global all_reduce needed for overall oneccl warm up
180 torch.distributed.all_reduce(torch.zeros(1).xpu())
181