@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a5354b3ed2
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 14, 05:30 UTC

[Bugfix][WideEP] Apply TP Attn + EP MoE fix to other models (#24982)

Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>

Tyler Michael Smith committed Sep 27, 2025 at 10:22 UTC a5354b3ed24723ac6e351896cc11e16dcee0b701
23 files changed +541 -376
vllm/config/parallel.py
+18
@@ -279,6 +279,24 @@ class ParallelConfig:
279 assert last_exc is not None
280 raise last_exc
281
282 + # The all_reduce at the end of attention (during o_proj) means that
283 + # inputs are replicated across each rank of the tensor parallel group.
284 + # If using expert-parallelism with DeepEP All2All ops, replicated
285 + # tokens results in useless duplicate computation and communication.
286 + #
287 + # In this case, ensure the input to the experts is sequence parallel
288 + # to avoid the excess work.
289 + #
290 + # Not needed for pplx-kernels as it can handle duplicate input tokens.
291 + @property
292 + def use_sequence_parallel_moe(self) -> bool:
293 + return (envs.VLLM_ALL2ALL_BACKEND
294 + in ("allgather_reducescatter", "naive",
295 + "deepep_high_throughput", "deepep_low_latency")
296 + and self.enable_expert_parallel
297 + and self.tensor_parallel_size > 1
298 + and self.data_parallel_size > 1)
299 +
300 @staticmethod
301 def has_unfinished_dp(dp_group: ProcessGroup,
302 has_unfinished: bool) -> bool:
vllm/distributed/device_communicators/all2all.py
+82 -40
@@ -6,7 +6,7 @@ import torch
6 import torch.distributed as dist
7
8 import vllm.envs as envs
9 -from vllm.distributed import get_dp_group
9 +from vllm.distributed import get_dp_group, get_ep_group
10 from vllm.forward_context import get_forward_context
11 from vllm.logger import init_logger
12 from vllm.utils import has_deep_ep, has_pplx
@@ -34,41 +34,60 @@ class NaiveAll2AllManager(All2AllManagerBase):
34 super().__init__(cpu_group)
35
36 def naive_multicast(self, x: torch.Tensor,
37 - cu_tokens_across_dp_cpu: torch.Tensor):
37 + cu_tokens_across_sp_cpu: torch.Tensor,
38 + is_sequence_parallel: bool) -> torch.Tensor:
39 assert (len(x.shape) == 2)
39 - buffer = torch.empty((cu_tokens_across_dp_cpu[-1], x.size(1)),
40 + buffer = torch.empty((cu_tokens_across_sp_cpu[-1], x.size(1)),
41 device=x.device,
42 dtype=x.dtype)
43
43 - start = 0 if self.dp_rank == 0 else cu_tokens_across_dp_cpu[
44 - self.dp_rank - 1]
45 - end = cu_tokens_across_dp_cpu[self.dp_rank]
44 + rank = self.rank if is_sequence_parallel else self.dp_rank
45 + world_size = (self.world_size
46 + if is_sequence_parallel else self.dp_world_size)
47 +
48 + start = 0 if rank == 0 else cu_tokens_across_sp_cpu[rank - 1]
49 + end = cu_tokens_across_sp_cpu[rank]
50 buffer[start:end, :].copy_(x)
47 - for idx in range(self.dp_world_size):
48 - start = 0 if idx == 0 else cu_tokens_across_dp_cpu[idx - 1]
49 - end = cu_tokens_across_dp_cpu[idx]
50 - self.dp_group.broadcast(buffer[start:end, :], idx)
51 + for idx in range(world_size):
52 + start = 0 if idx == 0 else cu_tokens_across_sp_cpu[idx - 1]
53 + end = cu_tokens_across_sp_cpu[idx]
54 + get_ep_group().broadcast(buffer[start:end, :], idx)
55
56 return buffer
57
54 - def dispatch(self, hidden_states: torch.Tensor,
55 - router_logits: torch.Tensor):
56 - sizes = get_forward_context(
57 - ).dp_metadata.get_chunk_sizes_across_dp_rank()
58 - hidden_states, router_logits = get_dp_group().all_gatherv(
59 - [hidden_states, router_logits],
60 - dim=0,
61 - sizes=sizes,
62 - )
63 -
58 + def dispatch(
59 + self,
60 + hidden_states: torch.Tensor,
61 + router_logits: torch.Tensor,
62 + is_sequence_parallel: bool = False
63 + ) -> tuple[torch.Tensor, torch.Tensor]:
64 + sp_size = self.tp_group.world_size if is_sequence_parallel else 1
65 + dp_metadata = get_forward_context().dp_metadata
66 + cu_tokens_across_sp_cpu = dp_metadata.cu_tokens_across_sp(sp_size)
67 +
68 + hidden_states = self.naive_multicast(hidden_states,
69 + cu_tokens_across_sp_cpu,
70 + is_sequence_parallel)
71 + router_logits = self.naive_multicast(router_logits,
72 + cu_tokens_across_sp_cpu,
73 + is_sequence_parallel)
74 return hidden_states, router_logits
75
66 - def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
67 - sizes = get_forward_context(
68 - ).dp_metadata.get_chunk_sizes_across_dp_rank()
69 - hidden_states = get_dp_group().reduce_scatterv(hidden_states,
70 - dim=0,
71 - sizes=sizes)
76 + def combine(self,
77 + hidden_states: torch.Tensor,
78 + is_sequence_parallel: bool = False) -> torch.Tensor:
79 +
80 + ep_rank = self.rank if is_sequence_parallel else self.dp_rank
81 +
82 + dp_metadata = get_forward_context().dp_metadata
83 + sp_size = self.tp_group.world_size if is_sequence_parallel else 1
84 + cu_tokens_across_sp_cpu = dp_metadata.cu_tokens_across_sp(sp_size)
85 +
86 + start = 0 if ep_rank == 0 else cu_tokens_across_sp_cpu[ep_rank - 1]
87 + end = cu_tokens_across_sp_cpu[ep_rank]
88 +
89 + all_hidden_states = get_ep_group().all_reduce(hidden_states)
90 + hidden_states = all_hidden_states[start:end, :]
91 return hidden_states
92
93 def destroy(self):
@@ -84,29 +103,40 @@ class AgRsAll2AllManager(All2AllManagerBase):
103 def __init__(self, cpu_group):
104 super().__init__(cpu_group)
105
87 - def dispatch(self, hidden_states: torch.Tensor,
88 - router_logits: torch.Tensor):
106 + def dispatch(
107 + self,
108 + hidden_states: torch.Tensor,
109 + router_logits: torch.Tensor,
110 + is_sequence_parallel: bool = False
111 + ) -> tuple[torch.Tensor, torch.Tensor]:
112 """
113 Gather hidden_states and router_logits from all dp ranks.
114 """
115 sizes = get_forward_context(
116 ).dp_metadata.get_chunk_sizes_across_dp_rank()
94 - hidden_states, router_logits = get_dp_group().all_gatherv(
117 +
118 + dist_group = get_ep_group() if is_sequence_parallel else get_dp_group()
119 + assert sizes[dist_group.rank_in_group] == hidden_states.shape[0]
120 + hidden_states, router_logits = dist_group.all_gatherv(
121 [hidden_states, router_logits],
122 dim=0,
123 sizes=sizes,
124 )
125 return hidden_states, router_logits
126
101 - def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
127 + def combine(self,
128 + hidden_states: torch.Tensor,
129 + is_sequence_parallel: bool = False) -> torch.Tensor:
130 """
131 Reduce-scatter hidden_states across all dp ranks.
132 """
133 sizes = get_forward_context(
134 ).dp_metadata.get_chunk_sizes_across_dp_rank()
107 - hidden_states = get_dp_group().reduce_scatterv(hidden_states,
108 - dim=0,
109 - sizes=sizes)
135 +
136 + dist_group = get_ep_group() if is_sequence_parallel else get_dp_group()
137 + hidden_states = dist_group.reduce_scatterv(hidden_states,
138 + dim=0,
139 + sizes=sizes)
140 return hidden_states
141
142 def destroy(self):
@@ -148,11 +178,17 @@ class PPLXAll2AllManager(All2AllManagerBase):
178 kwargs, pplx.AllToAll.internode
179 if self.internode else pplx.AllToAll.intranode)
180
151 - def dispatch(self, hidden_states: torch.Tensor,
152 - router_logits: torch.Tensor):
181 + def dispatch(
182 + self,
183 + hidden_states: torch.Tensor,
184 + router_logits: torch.Tensor,
185 + is_sequence_parallel: bool = False
186 + ) -> tuple[torch.Tensor, torch.Tensor]:
187 raise NotImplementedError
188
155 - def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
189 + def combine(self,
190 + hidden_states: torch.Tensor,
191 + is_sequence_parallel: bool = False) -> torch.Tensor:
192 raise NotImplementedError
193
194 def destroy(self):
@@ -184,11 +220,17 @@ class DeepEPAll2AllManagerBase(All2AllManagerBase):
220 def get_handle(self, kwargs):
221 raise NotImplementedError
222
187 - def dispatch(self, hidden_states: torch.Tensor,
188 - router_logits: torch.Tensor):
223 + def dispatch(
224 + self,
225 + hidden_states: torch.Tensor,
226 + router_logits: torch.Tensor,
227 + is_sequence_parallel: bool = False
228 + ) -> tuple[torch.Tensor, torch.Tensor]:
229 raise NotImplementedError
230
191 - def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
231 + def combine(self,
232 + hidden_states: torch.Tensor,
233 + is_sequence_parallel: bool = False) -> torch.Tensor:
234 raise NotImplementedError
235
236 def destroy(self):
@@ -395,4 +437,4 @@ class FlashInferAllToAllManager(All2AllManagerBase):
437 self.workspace_tensor = None
438 self.prepare_workspace_tensor = None
439 self.mapping = None
398 - self.initialized = False
\ No newline at end of file
440 + self.initialized = False
vllm/distributed/device_communicators/base_device_communicator.py
+20 -8
@@ -28,6 +28,8 @@ class Cache:
28
29
30 class All2AllManagerBase:
31 + rank: int
32 + world_size: int
33
34 def __init__(self, cpu_group):
35 self.cpu_group = cpu_group
@@ -40,6 +42,7 @@ class All2AllManagerBase:
42 # all2all lives in ep group, which is merged from dp and tp group
43 self.dp_group = get_dp_group()
44 self.tp_group = get_tp_group()
45 +
46 # no self.ep_group since self.ep_group is still in construction
47 # when we create this object
48 self.dp_rank = self.dp_group.rank_in_group
@@ -60,17 +63,21 @@ class All2AllManagerBase:
63 # and reuse it for the same config.
64 raise NotImplementedError
65
66 + def dispatch(self,
67 + hidden_states: torch.Tensor,
68 + router_logits: torch.Tensor,
69 + is_sequence_parallel: bool = False):
70 + raise NotImplementedError
71 +
72 def set_num_sms(self, num_sms: int):
73 pass
74
75 def max_sms_used(self) -> Optional[int]:
76 return None # None means it could use the whole GPU
77
69 - def dispatch(self, hidden_states: torch.Tensor,
70 - router_logits: torch.Tensor):
71 - raise NotImplementedError
72 -
73 - def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
78 + def combine(self,
79 + hidden_states: torch.Tensor,
80 + is_sequence_parallel: bool = False):
81 raise NotImplementedError
82
83 def destroy(self):
@@ -267,15 +274,20 @@ class DeviceCommunicatorBase:
274 module.quant_method.init_prepare_finalize(module)
275
276 def dispatch(
270 - self, hidden_states: torch.Tensor,
271 - router_logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
277 + self,
278 + hidden_states: torch.Tensor,
279 + router_logits: torch.Tensor,
280 + is_sequence_parallel: bool = False
281 + ) -> tuple[torch.Tensor, torch.Tensor]:
282 """
283 Dispatch the hidden states and router logits to the appropriate device.
284 This is a no-op in the base class.
285 """
286 return hidden_states, router_logits
287
278 - def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
288 + def combine(self,
289 + hidden_states: torch.Tensor,
290 + is_sequence_parallel: bool = False) -> torch.Tensor:
291 """
292 Combine the hidden states and router logits from the appropriate device.
293 This is a no-op in the base class.
vllm/distributed/device_communicators/cuda_communicator.py
+12 -10
@@ -39,10 +39,6 @@ class CudaCommunicator(DeviceCommunicatorBase):
39 use_custom_allreduce = _ENABLE_CUSTOM_ALL_REDUCE
40 use_torch_symm_mem = envs.VLLM_ALLREDUCE_USE_SYMM_MEM
41
42 - # ep does not use pynccl
43 - use_pynccl = "ep" not in unique_name
44 -
45 - self.use_pynccl = use_pynccl
42 self.use_custom_allreduce = use_custom_allreduce
43 self.use_torch_symm_mem = use_torch_symm_mem
44
@@ -57,7 +53,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
53 SymmMemCommunicator)
54
55 self.pynccl_comm: Optional[PyNcclCommunicator] = None
60 - if use_pynccl and self.world_size > 1:
56 + if self.world_size > 1:
57 self.pynccl_comm = PyNcclCommunicator(
58 group=self.cpu_group,
59 device=self.device,
@@ -308,14 +304,20 @@ class CudaCommunicator(DeviceCommunicatorBase):
304 return output_list
305
306 def dispatch(
311 - self, hidden_states: torch.Tensor,
312 - router_logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
307 + self,
308 + hidden_states: torch.Tensor,
309 + router_logits: torch.Tensor,
310 + is_sequence_parallel: bool = False
311 + ) -> tuple[torch.Tensor, torch.Tensor]:
312 assert self.all2all_manager is not None
313 hidden_states, router_logits = self.all2all_manager.dispatch(
315 - hidden_states, router_logits)
314 + hidden_states, router_logits, is_sequence_parallel)
315 return hidden_states, router_logits
316
318 - def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
317 + def combine(self,
318 + hidden_states: torch.Tensor,
319 + is_sequence_parallel: bool = False) -> torch.Tensor:
320 assert self.all2all_manager is not None
320 - hidden_states = self.all2all_manager.combine(hidden_states)
321 + hidden_states = self.all2all_manager.combine(hidden_states,
322 + is_sequence_parallel)
323 return hidden_states
vllm/distributed/device_communicators/xpu_communicator.py
+11 -5
@@ -75,14 +75,20 @@ class XpuCommunicator(DeviceCommunicatorBase):
75 dist.broadcast(input_, src=src, group=self.device_group)
76
77 def dispatch(
78 - self, hidden_states: torch.Tensor,
79 - router_logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
78 + self,
79 + hidden_states: torch.Tensor,
80 + router_logits: torch.Tensor,
81 + is_sequence_parallel: bool = False
82 + ) -> tuple[torch.Tensor, torch.Tensor]:
83 assert self.all2all_manager is not None
84 hidden_states, router_logits = self.all2all_manager.dispatch(
82 - hidden_states, router_logits)
85 + hidden_states, router_logits, is_sequence_parallel)
86 return hidden_states, router_logits
87
85 - def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
88 + def combine(self,
89 + hidden_states: torch.Tensor,
90 + is_sequence_parallel: bool = False) -> torch.Tensor:
91 assert self.all2all_manager is not None
87 - hidden_states = self.all2all_manager.combine(hidden_states)
92 + hidden_states = self.all2all_manager.combine(hidden_states,
93 + is_sequence_parallel)
94 return hidden_states
vllm/distributed/parallel_state.py
+12 -5
@@ -871,17 +871,24 @@ class GroupCoordinator:
871 model)
872
873 def dispatch(
874 - self, hidden_states: torch.Tensor,
875 - router_logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
874 + self,
875 + hidden_states: torch.Tensor,
876 + router_logits: torch.Tensor,
877 + is_sequence_parallel: bool = False
878 + ) -> tuple[torch.Tensor, torch.Tensor]:
879 if self.device_communicator is not None:
880 return self.device_communicator.dispatch(hidden_states,
878 - router_logits)
881 + router_logits,
882 + is_sequence_parallel)
883 else:
884 return hidden_states, router_logits
885
882 - def combine(self, hidden_states) -> torch.Tensor:
886 + def combine(self,
887 + hidden_states,
888 + is_sequence_parallel: bool = False) -> torch.Tensor:
889 if self.device_communicator is not None:
884 - return self.device_communicator.combine(hidden_states)
890 + return self.device_communicator.combine(hidden_states,
891 + is_sequence_parallel)
892 else:
893 return hidden_states
894
vllm/forward_context.py
+66 -31
@@ -49,16 +49,29 @@ class BatchDescriptor(NamedTuple):
49 return BatchDescriptor(self.num_tokens, uniform_decode=False)
50
51
52 -def _compute_chunked_local_num_tokens(num_tokens_across_dp_cpu: list[int],
52 +def _compute_sp_num_tokens(num_tokens_across_dp_cpu: torch.Tensor,
53 + sequence_parallel_size: int) -> list[int]:
54 + sp_tokens = ((num_tokens_across_dp_cpu + sequence_parallel_size - 1) //
55 + sequence_parallel_size)
56 +
57 + sp_tokens = sp_tokens.repeat_interleave(sequence_parallel_size)
58 + return sp_tokens.tolist()
59 +
60 +
61 +def _compute_chunked_local_num_tokens(num_tokens_across_dp_cpu: torch.Tensor,
62 + sequence_parallel_size: int,
63 max_num_tokens: int,
64 chunk_idx: int) -> list[int]:
55 - dp_size = len(num_tokens_across_dp_cpu)
65
57 - local_size = [-1] * dp_size
58 - for i in range(dp_size):
59 - dp_tokens = num_tokens_across_dp_cpu[i]
66 + sp_tokens = _compute_sp_num_tokens(num_tokens_across_dp_cpu,
67 + sequence_parallel_size)
68 + sp_size = len(sp_tokens)
69 +
70 + local_size = [-1] * sp_size
71 + for i in range(sp_size):
72 + # Take into account sharding if MoE activation is sequence parallel.
73 local_size[i] = min(max_num_tokens,
61 - dp_tokens - (max_num_tokens * chunk_idx))
74 + sp_tokens[i] - (max_num_tokens * chunk_idx))
75 if local_size[i] <= 0:
76 local_size[i] = 1 # ensure lockstep even if done
77 return local_size
@@ -67,7 +80,9 @@ def _compute_chunked_local_num_tokens(num_tokens_across_dp_cpu: list[int],
80 @dataclass
81 class DPMetadata:
82 max_tokens_across_dp_cpu: torch.Tensor
70 - cu_tokens_across_dp_cpu: torch.Tensor
83 + num_tokens_across_dp_cpu: torch.Tensor
84 +
85 + # NOTE: local_sizes should only be set by the chunked_sizes context manager
86 local_sizes: Optional[list[int]] = None
87
88 @staticmethod
@@ -98,6 +113,17 @@ class DPMetadata:
113 dist.all_reduce(num_tokens_tensor, group=group)
114 return num_tokens_tensor.cpu()
115
116 + # Get the cumulative tokens across sequence parallel ranks.
117 + # In this case the input to the MoEs will be distributed w.r.t both
118 + # DP and TP rank.
119 + # When sp_size==1, this is just the cummulative num tokens across DP.
120 + def cu_tokens_across_sp(self, sp_size: int) -> torch.Tensor:
121 + num_tokens_across_sp_cpu = (
122 + (self.num_tokens_across_dp_cpu - 1 + sp_size) // sp_size)
123 + num_tokens_across_sp_cpu = (
124 + num_tokens_across_sp_cpu.repeat_interleave(sp_size))
125 + return torch.cumsum(num_tokens_across_sp_cpu, dim=0)
126 +
127 @staticmethod
128 def should_ubatch_across_dp(
129 should_ubatch: bool, orig_num_tokens_per_ubatch: int,
@@ -147,10 +173,10 @@ class DPMetadata:
173
174 @staticmethod
175 def make(
150 - parallel_config: ParallelConfig,
151 - attn_metadata: Any,
152 - num_tokens: int,
153 - num_tokens_across_dp: Optional[torch.Tensor] = None
176 + parallel_config: ParallelConfig,
177 + attn_metadata: Any,
178 + num_tokens: int,
179 + num_tokens_across_dp_cpu: Optional[torch.Tensor] = None
180 ) -> "DPMetadata":
181
182 assert parallel_config.data_parallel_size > 1
@@ -167,18 +193,18 @@ class DPMetadata:
193
194 # If num_tokens_across_dp is None, it will be computed by all_reduce
195 # Otherwise, num_tokens_across_dp[dp_rank] should be equal to batchsize
170 - assert (num_tokens_across_dp is None or num_tokens_across_dp[dp_rank]
171 - == batchsize), f"{num_tokens_across_dp[dp_rank]} {batchsize}"
172 - if num_tokens_across_dp is None:
173 - num_tokens_across_dp = DPMetadata.num_tokens_across_dp(
196 + assert (num_tokens_across_dp_cpu is None
197 + or num_tokens_across_dp_cpu[dp_rank] == batchsize
198 + ), f"{num_tokens_across_dp_cpu[dp_rank]} {batchsize}"
199 + if num_tokens_across_dp_cpu is None:
200 + num_tokens_across_dp_cpu = DPMetadata.num_tokens_across_dp(
201 batchsize, dp_size, dp_rank)
175 - max_tokens_across_dp_cpu = torch.max(num_tokens_across_dp)
176 - cu_tokens_across_dp_cpu = torch.cumsum(num_tokens_across_dp, dim=0)
177 - return DPMetadata(max_tokens_across_dp_cpu, cu_tokens_across_dp_cpu,
178 - num_tokens_across_dp)
202 + max_tokens_across_dp_cpu = torch.max(num_tokens_across_dp_cpu)
203 + return DPMetadata(max_tokens_across_dp_cpu, num_tokens_across_dp_cpu)
204
205 @contextmanager
181 - def chunked_sizes(self, max_chunk_size_per_rank: int, chunk_idx: int):
206 + def chunked_sizes(self, sequence_parallel_size: int,
207 + max_chunk_size_per_rank: int, chunk_idx: int):
208 """
209 Context manager to compute and temporarily set the per-rank local token
210 sizes for a specific chunk during chunked forward execution.
@@ -192,31 +218,40 @@ class DPMetadata:
218 `chunk_idx`, this context manager sets `self.local_sizes` to the number
219 of tokens to process in that chunk on each rank.
220
195 - It uses cumulative sizes (`cu_tokens_across_dp_cpu`) to derive the
196 - number of tokens per rank, and calls `_compute_chunked_local_num_tokens`
197 - to determine the chunk-wise split.
198 -
221 `self.local_sizes` is only valid inside the context.
222
223 Args:
224 + sequence_parallel_size: When Attn is TP and MoE layers are EP,
225 + we use SP between the layers to avoid
226 + redundant ops. We need this value to
227 + compute the chunked sizes.
228 max_chunk_size_per_rank: The max number of tokens each rank is
229 allowed to process in this chunk.
230 chunk_idx: The index of the chunk to compute sizes for.
231 """
206 - cu_sizes = self.cu_tokens_across_dp_cpu
207 - num_tokens_across_dp_cpu = [
208 - (cu_sizes[i] -
209 - cu_sizes[i - 1]).item() if i > 0 else cu_sizes[0].item()
210 - for i in range(len(cu_sizes))
211 - ]
232 self.local_sizes = _compute_chunked_local_num_tokens(
213 - num_tokens_across_dp_cpu, max_chunk_size_per_rank, chunk_idx)
233 + self.num_tokens_across_dp_cpu, sequence_parallel_size,
234 + max_chunk_size_per_rank, chunk_idx)
235 + try:
236 + yield self.local_sizes
237 + finally:
238 + self.local_sizes = None
239 +
240 + @contextmanager
241 + def sp_local_sizes(self, sequence_parallel_size: int):
242 + """
243 + Context mamager for setting self.local_sizes. Same as self.chunked_sizes
244 + but without any chunking.
245 + """
246 + self.local_sizes = _compute_sp_num_tokens(
247 + self.num_tokens_across_dp_cpu, sequence_parallel_size)
248 try:
249 yield self.local_sizes
250 finally:
251 self.local_sizes = None
252
253 def get_chunk_sizes_across_dp_rank(self) -> Optional[list[int]]:
254 + assert self.local_sizes is not None
255 return self.local_sizes
256
257
vllm/model_executor/layers/fused_moe/layer.py
+64 -55
@@ -3,6 +3,7 @@
3
4 from abc import abstractmethod
5 from collections.abc import Iterable
6 +from contextlib import nullcontext
7 from enum import Enum
8 from typing import Callable, Literal, Optional, Union, get_args, overload
9
@@ -983,8 +984,7 @@ class FusedMoE(CustomOp):
984 if dp_size is not None else get_dp_group().world_size)
985
986 self.is_sequence_parallel = is_sequence_parallel
986 - if self.is_sequence_parallel:
987 - self.sp_size = tp_size_
987 + self.sp_size = tp_size_ if is_sequence_parallel else 1
988
989 self.moe_parallel_config: FusedMoEParallelConfig = (
990 FusedMoEParallelConfig.make(
@@ -1966,7 +1966,8 @@ class FusedMoE(CustomOp):
1966 # clamp start and end
1967 chunk_start = min(chunk_start, num_tokens - 1)
1968 chunk_end = min(chunk_end, num_tokens)
1969 - with ctx.dp_metadata.chunked_sizes(moe_dp_chunk_size_per_rank,
1969 + with ctx.dp_metadata.chunked_sizes(self.sp_size,
1970 + moe_dp_chunk_size_per_rank,
1971 chunk_idx):
1972 process_chunk(chunk_start,
1973 chunk_end,
@@ -2011,65 +2012,73 @@ class FusedMoE(CustomOp):
2012 else:
2013 shared_output = None
2014
2014 - if do_naive_dispatch_combine:
2015 - hidden_states, router_logits = get_ep_group().dispatch(
2016 - hidden_states, router_logits)
2015 + ctx = get_forward_context()
2016 + sp_ctx = ctx.dp_metadata.sp_local_sizes(
2017 + self.sp_size) if ctx.dp_metadata else nullcontext()
2018
2018 - # Matrix multiply.
2019 - final_hidden_states = self.quant_method.apply(
2020 - layer=self,
2021 - x=hidden_states,
2022 - router_logits=router_logits,
2023 - top_k=self.top_k,
2024 - renormalize=self.renormalize,
2025 - use_grouped_topk=self.use_grouped_topk,
2026 - global_num_experts=self.global_num_experts,
2027 - expert_map=self.expert_map,
2028 - topk_group=self.topk_group,
2029 - num_expert_group=self.num_expert_group,
2030 - custom_routing_function=self.custom_routing_function,
2031 - scoring_func=self.scoring_func,
2032 - routed_scaling_factor=self.routed_scaling_factor,
2033 - e_score_correction_bias=self.e_score_correction_bias,
2034 - activation=self.activation,
2035 - apply_router_weight_on_input=self.apply_router_weight_on_input,
2036 - enable_eplb=self.enable_eplb,
2037 - expert_load_view=self.expert_load_view,
2038 - logical_to_physical_map=self.logical_to_physical_map,
2039 - logical_replica_count=self.logical_replica_count,
2040 - )
2019 + with sp_ctx:
2020 + if do_naive_dispatch_combine:
2021 + hidden_states, router_logits = get_ep_group().dispatch(
2022 + hidden_states, router_logits, self.is_sequence_parallel)
2023
2042 - if shared_output is not None:
2043 - assert not isinstance(final_hidden_states, tuple)
2044 - assert self.shared_experts is not None
2045 - final_hidden_states = (
2046 - shared_output,
2047 - final_hidden_states,
2024 + # Matrix multiply.
2025 + final_hidden_states = self.quant_method.apply(
2026 + layer=self,
2027 + x=hidden_states,
2028 + router_logits=router_logits,
2029 + top_k=self.top_k,
2030 + renormalize=self.renormalize,
2031 + use_grouped_topk=self.use_grouped_topk,
2032 + global_num_experts=self.global_num_experts,
2033 + expert_map=self.expert_map,
2034 + topk_group=self.topk_group,
2035 + num_expert_group=self.num_expert_group,
2036 + custom_routing_function=self.custom_routing_function,
2037 + scoring_func=self.scoring_func,
2038 + routed_scaling_factor=self.routed_scaling_factor,
2039 + e_score_correction_bias=self.e_score_correction_bias,
2040 + activation=self.activation,
2041 + apply_router_weight_on_input=self.apply_router_weight_on_input,
2042 + enable_eplb=self.enable_eplb,
2043 + expert_load_view=self.expert_load_view,
2044 + logical_to_physical_map=self.logical_to_physical_map,
2045 + logical_replica_count=self.logical_replica_count,
2046 )
2049 - elif self.zero_expert_num is not None and self.zero_expert_num > 0:
2050 - assert isinstance(final_hidden_states, tuple)
2051 - final_hidden_states, zero_expert_result = final_hidden_states
2047
2053 - def reduce_output(states: torch.Tensor,
2054 - do_combine: bool = True) -> torch.Tensor:
2055 - if do_naive_dispatch_combine and do_combine:
2056 - states = get_ep_group().combine(states)
2048 + if shared_output is not None:
2049 + assert not isinstance(final_hidden_states, tuple)
2050 + assert self.shared_experts is not None
2051 + final_hidden_states = (
2052 + shared_output,
2053 + final_hidden_states,
2054 + )
2055 + elif self.zero_expert_num is not None and self.zero_expert_num > 0:
2056 + assert isinstance(final_hidden_states, tuple)
2057 + final_hidden_states, zero_expert_result = final_hidden_states
2058 +
2059 + def reduce_output(states: torch.Tensor,
2060 + do_combine: bool = True) -> torch.Tensor:
2061 + if do_naive_dispatch_combine and do_combine:
2062 + states = get_ep_group().combine(states,
2063 + self.is_sequence_parallel)
2064
2058 - if self.reduce_results and (self.tp_size > 1 or self.ep_size > 1):
2059 - states = self.maybe_all_reduce_tensor_model_parallel(states)
2065 + if (not self.is_sequence_parallel and self.reduce_results
2066 + and (self.tp_size > 1 or self.ep_size > 1)):
2067 + states = self.maybe_all_reduce_tensor_model_parallel(
2068 + states)
2069
2061 - return states
2070 + return states
2071
2063 - if self.shared_experts is not None:
2064 - return (
2065 - reduce_output(final_hidden_states[0], do_combine=False),
2066 - reduce_output(final_hidden_states[1]),
2067 - )
2068 - elif self.zero_expert_num is not None and self.zero_expert_num > 0:
2069 - assert isinstance(final_hidden_states, torch.Tensor)
2070 - return reduce_output(final_hidden_states) + zero_expert_result
2071 - else:
2072 - return reduce_output(final_hidden_states)
2072 + if self.shared_experts is not None:
2073 + return (
2074 + reduce_output(final_hidden_states[0], do_combine=False),
2075 + reduce_output(final_hidden_states[1]),
2076 + )
2077 + elif self.zero_expert_num is not None and self.zero_expert_num > 0:
2078 + assert isinstance(final_hidden_states, torch.Tensor)
2079 + return reduce_output(final_hidden_states) + zero_expert_result
2080 + else:
2081 + return reduce_output(final_hidden_states)
2082
2083 @classmethod
2084 def make_expert_params_mapping(
vllm/model_executor/models/aria.py
+7 -9
@@ -9,7 +9,7 @@ from transformers import AriaConfig, AriaTextConfig, BatchFeature
9 from transformers.models.aria.modeling_aria import AriaCrossAttention
10 from transformers.models.aria.processing_aria import AriaProcessor
11
12 -from vllm.config import CacheConfig, QuantizationConfig, VllmConfig
12 +from vllm.config import QuantizationConfig, VllmConfig
13 from vllm.distributed import get_tensor_model_parallel_rank
14 from vllm.model_executor.layers.activation import get_act_fn
15 from vllm.model_executor.layers.fused_moe import FusedMoE
@@ -297,14 +297,12 @@ class AriaTextDecoderLayer(LlamaDecoderLayer):
297 Experts (MoE) Layer.
298 """
299
300 - def __init__(
301 - self,
302 - config: AriaTextConfig,
303 - cache_config: Optional[CacheConfig] = None,
304 - quant_config: Optional[QuantizationConfig] = None,
305 - prefix: str = "",
306 - ) -> None:
307 - super().__init__(config, cache_config, quant_config, prefix)
300 + def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:
301 + super().__init__(vllm_config, prefix)
302 +
303 + config = vllm_config.model_config.hf_config
304 + quant_config = vllm_config.quant_config
305 +
306 self.mlp = AriaTextMoELayer(config,
307 quant_config=quant_config,
308 prefix=f"{prefix}.mlp")
vllm/model_executor/models/deepseek_v2.py
+3 -55
@@ -32,7 +32,6 @@ import torch
32 from torch import nn
33 from transformers import DeepseekV2Config, DeepseekV3Config
34
35 -import vllm.envs as envs
35 from vllm.attention import Attention
36 from vllm.compilation.decorators import support_torch_compile
37 from vllm.config import CacheConfig, ParallelConfig, VllmConfig
@@ -56,8 +55,8 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
55 ParallelLMHead, VocabParallelEmbedding)
56 from vllm.model_executor.model_loader.weight_utils import (
57 default_weight_loader, maybe_remap_kv_scale_name)
58 +from vllm.model_executor.models.utils import sequence_parallel_chunk
59 from vllm.sequence import IntermediateTensors
60 -from vllm.utils import cdiv, direct_register_custom_op
60
61 from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP
62 from .utils import (PPMissingLayer, is_pp_missing_parameter,
@@ -108,43 +107,6 @@ class DeepseekV2MLP(nn.Module):
107 return x
108
109
111 -# Chunk x along the num_tokens axis for sequence parallelism
112 -# NOTE: This is wrapped in a torch custom op to work around the following issue:
113 -# The output tensor can have a sequence length 0 at small input sequence lengths
114 -# even though we explicitly pad to avoid this.
115 -def sequence_parallel_chunk(x: torch.Tensor) -> torch.Tensor:
116 - tp_size = get_tensor_model_parallel_world_size()
117 - tp_rank = get_tensor_model_parallel_rank()
118 -
119 - # all_gather needs the sequence length to be divisible by tp_size
120 - seq_len = x.size(0)
121 - remainder = seq_len % tp_size
122 - if remainder != 0:
123 - pad_len = tp_size - remainder
124 - x = nn.functional.pad(x, (0, 0, 0, pad_len))
125 -
126 - chunk = x.shape[0] // tp_size
127 - start = tp_rank * chunk
128 - return torch.narrow(x, 0, start, chunk)
129 -
130 -
131 -def sequence_parallel_chunk_fake(x: torch.Tensor) -> torch.Tensor:
132 - tp_size = get_tensor_model_parallel_world_size()
133 - seq_len = cdiv(x.size(0), tp_size)
134 - shape = list(x.shape)
135 - shape[0] = seq_len
136 - out = torch.empty(shape, dtype=x.dtype, device=x.device)
137 - return out
138 -
139 -
140 -direct_register_custom_op(
141 - op_name="sequence_parallel_chunk",
142 - op_func=sequence_parallel_chunk,
143 - fake_impl=sequence_parallel_chunk_fake,
144 - tags=(torch.Tag.needs_fixed_stride_order, ),
145 -)
146 -
147 -
110 class DeepseekV2MoE(nn.Module):
111
112 def __init__(
@@ -166,20 +128,7 @@ class DeepseekV2MoE(nn.Module):
128 self.n_routed_experts: int = config.n_routed_experts
129 self.n_shared_experts: int = config.n_shared_experts
130
169 - # The all_reduce at the end of attention (during o_proj) means that
170 - # inputs are replicated across each rank of the tensor parallel group.
171 - # If using expert-parallelism with DeepEP All2All ops, replicated
172 - # tokens results in useless duplicate computation and communication.
173 - #
174 - # In this case, ensure the input to the experts is sequence parallel
175 - # to avoid the excess work.
176 - #
177 - # Not needed for pplx-kernels as it can handle duplicate input tokens.
178 - self.is_sequence_parallel = (envs.VLLM_ALL2ALL_BACKEND
179 - in ("deepep_high_throughput",
180 - "deepep_low_latency")
181 - and parallel_config.enable_expert_parallel
182 - and self.tp_size > 1)
131 + self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
132
133 if config.hidden_act != "silu":
134 raise ValueError(f"Unsupported activation: {config.hidden_act}. "
@@ -278,8 +227,7 @@ class DeepseekV2MoE(nn.Module):
227 # TODO: We can replace the all_reduce at the end of attn with a
228 # reduce_scatter instead of chunking here.
229 if self.is_sequence_parallel:
281 - hidden_states = torch.ops.vllm.sequence_parallel_chunk(
282 - hidden_states)
230 + hidden_states = sequence_parallel_chunk(hidden_states)
231
232 # router_logits: (num_tokens, n_experts)
233 router_logits, _ = self.gate(hidden_states)
vllm/model_executor/models/ernie_mtp.py
+5 -11
@@ -29,10 +29,9 @@ import torch
29 import torch.nn as nn
30 from transformers import PretrainedConfig
31
32 -from vllm.config import CacheConfig, ModelConfig, VllmConfig
32 +from vllm.config import VllmConfig
33 from vllm.model_executor.layers.layernorm import RMSNorm
34 from vllm.model_executor.layers.logits_processor import LogitsProcessor
35 -from vllm.model_executor.layers.quantization import QuantizationConfig
35 from vllm.model_executor.layers.vocab_parallel_embedding import (
36 ParallelLMHead, VocabParallelEmbedding)
37 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
@@ -47,13 +46,11 @@ class ErnieMultiTokenPredictorLayer(nn.Module):
46
47 def __init__(
48 self,
50 - config: PretrainedConfig,
49 + vllm_config: VllmConfig,
50 prefix: str,
52 - model_config: ModelConfig,
53 - cache_config: Optional[CacheConfig] = None,
54 - quant_config: Optional[QuantizationConfig] = None,
51 ) -> None:
52 super().__init__()
53 + config = vllm_config.model_config.hf_config
54
55 self.mtp_emb_norm = RMSNorm(config.hidden_size,
56 eps=config.rms_norm_eps)
@@ -62,8 +59,7 @@ class ErnieMultiTokenPredictorLayer(nn.Module):
59 self.mtp_linear_proj = nn.Linear(config.hidden_size * 2,
60 config.hidden_size,
61 bias=False)
65 - self.mtp_block = LlamaDecoderLayer(config, cache_config, quant_config,
66 - prefix)
62 + self.mtp_block = LlamaDecoderLayer(vllm_config, prefix)
63
64 def forward(
65 self,
@@ -102,10 +98,8 @@ class ErnieMultiTokenPredictor(nn.Module):
98 self.layers = torch.nn.ModuleDict({
99 str(idx):
100 ErnieMultiTokenPredictorLayer(
105 - config,
101 + vllm_config,
102 f"{prefix}.layers.{idx}",
107 - model_config=vllm_config.model_config,
108 - cache_config=vllm_config.cache_config,
103 )
104 for idx in range(self.mtp_start_layer_idx,
105 self.mtp_start_layer_idx + self.num_mtp_layers)
vllm/model_executor/models/glm4.py
+9 -7
@@ -136,14 +136,16 @@ class Glm4Attention(nn.Module):
136
137 class Glm4DecoderLayer(nn.Module):
138
139 - def __init__(
140 - self,
141 - config: Glm4Config,
142 - cache_config: Optional[CacheConfig] = None,
143 - quant_config: Optional[QuantizationConfig] = None,
144 - prefix: str = "",
145 - ) -> None:
139 + def __init__(self,
140 + vllm_config: VllmConfig,
141 + prefix: str = "",
142 + config: Optional[Glm4Config] = None) -> None:
143 super().__init__()
144 +
145 + config = config or vllm_config.model_config.hf_config
146 + cache_config = vllm_config.cache_config
147 + quant_config = vllm_config.quant_config
148 +
149 self.hidden_size = config.hidden_size
150 rope_theta = getattr(config, "rope_theta", 1000000)
151 rope_scaling = getattr(config, "rope_scaling", None)
vllm/model_executor/models/gpt_oss.py
+28 -14
@@ -13,7 +13,8 @@ from vllm.compilation.decorators import support_torch_compile
13 from vllm.config import CacheConfig, VllmConfig
14 from vllm.distributed import (get_ep_group, get_pp_group,
15 get_tensor_model_parallel_rank,
16 - get_tensor_model_parallel_world_size)
16 + get_tensor_model_parallel_world_size,
17 + tensor_model_parallel_all_gather)
18 from vllm.model_executor.layers.fused_moe import FusedMoE
19 from vllm.model_executor.layers.layernorm import RMSNorm
20 from vllm.model_executor.layers.linear import (QKVParallelLinear,
@@ -24,6 +25,7 @@ from vllm.model_executor.layers.rotary_embedding import get_rope
25 from vllm.model_executor.layers.vocab_parallel_embedding import (
26 ParallelLMHead, VocabParallelEmbedding)
27 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
28 +from vllm.model_executor.models.utils import sequence_parallel_chunk
29 from vllm.sequence import IntermediateTensors
30 from vllm.utils import cdiv
31
@@ -132,12 +134,18 @@ class MLPBlock(torch.nn.Module):
134
135 def __init__(
136 self,
135 - config: GptOssConfig,
137 + vllm_config: VllmConfig,
138 layer_idx: int,
137 - quant_config: QuantizationConfig,
139 prefix: str = "",
140 ):
141 super().__init__()
142 +
143 + config = vllm_config.model_config.hf_config
144 + quant_config = vllm_config.quant_config
145 + parallel_config = vllm_config.parallel_config
146 +
147 + self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
148 +
149 self.layer_idx = layer_idx
150 self.num_experts = config.num_local_experts
151 self.experts_per_token = config.num_experts_per_tok
@@ -155,11 +163,20 @@ class MLPBlock(torch.nn.Module):
163 prefix=f"{prefix}.experts",
164 apply_router_weight_on_input=False,
165 has_bias=True,
158 - activation="swigluoai")
166 + activation="swigluoai",
167 + is_sequence_parallel=self.is_sequence_parallel)
168
169 def forward(self, x: torch.Tensor) -> torch.Tensor:
170 + num_tokens = x.shape[0]
171 + if self.is_sequence_parallel:
172 + x = sequence_parallel_chunk(x)
173 +
174 g = self.router(x)
175 x = self.experts(hidden_states=x, router_logits=g)
176 +
177 + if self.is_sequence_parallel:
178 + x = tensor_model_parallel_all_gather(x.contiguous(), 0)
179 + x = x[:num_tokens]
180 return x
181
182
@@ -167,19 +184,20 @@ class TransformerBlock(torch.nn.Module):
184
185 def __init__(
186 self,
170 - config: GptOssConfig,
171 - cache_config: CacheConfig,
172 - quant_config: QuantizationConfig,
187 + vllm_config: VllmConfig,
188 prefix: str = "",
189 ):
190 super().__init__()
191 +
192 + config = vllm_config.model_config.hf_config
193 + cache_config = vllm_config.cache_config
194 +
195 self.layer_idx = extract_layer_index(prefix)
196 self.attn = OAIAttention(config,
197 prefix=f"{prefix}.attn",
198 cache_config=cache_config)
180 - self.mlp = MLPBlock(config,
199 + self.mlp = MLPBlock(vllm_config,
200 self.layer_idx,
182 - quant_config=quant_config,
201 prefix=f"{prefix}.mlp")
202 self.input_layernorm = RMSNorm(config.hidden_size, eps=1e-5)
203 self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=1e-5)
@@ -216,8 +234,6 @@ class GptOssModel(nn.Module):
234 ):
235 super().__init__()
236 self.config = vllm_config.model_config.hf_config
219 - self.cache_config = vllm_config.cache_config
220 - self.quant_config = vllm_config.quant_config
237 self.parallel_config = vllm_config.parallel_config
238 self.config.hidden_size = self.config.hidden_size
239 self.embedding = VocabParallelEmbedding(
@@ -227,9 +243,7 @@ class GptOssModel(nn.Module):
243 self.start_layer, self.end_layer, self.layers = make_layers(
244 self.config.num_hidden_layers,
245 lambda prefix: TransformerBlock(
230 - self.config,
231 - cache_config=self.cache_config,
232 - quant_config=self.quant_config,
246 + vllm_config,
247 prefix=prefix,
248 ),
249 prefix=f"{prefix}.layers",
vllm/model_executor/models/granitemoe.py
+28 -10
@@ -29,12 +29,13 @@ from typing import Any, Optional
29
30 import torch
31 from torch import nn
32 -from transformers.models.granitemoe import GraniteMoeConfig
32
33 from vllm.attention import Attention
34 from vllm.compilation.decorators import support_torch_compile
35 from vllm.config import CacheConfig, VllmConfig
37 -from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
36 +from vllm.distributed import (get_pp_group,
37 + get_tensor_model_parallel_world_size,
38 + tensor_model_parallel_all_gather)
39 from vllm.model_executor.layers.fused_moe import FusedMoE
40 from vllm.model_executor.layers.layernorm import RMSNorm
41 from vllm.model_executor.layers.linear import (QKVParallelLinear,
@@ -48,6 +49,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
49 DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
50 from vllm.model_executor.model_loader.weight_utils import (
51 default_weight_loader, maybe_remap_kv_scale_name)
52 +from vllm.model_executor.models.utils import sequence_parallel_chunk
53 from vllm.sequence import IntermediateTensors
54
55 from .interfaces import SupportsLoRA, SupportsPP
@@ -71,9 +73,11 @@ class GraniteMoeMoE(nn.Module):
73 params_dtype: Optional[torch.dtype] = None,
74 quant_config: Optional[QuantizationConfig] = None,
75 tp_size: Optional[int] = None,
76 + is_sequence_parallel=False,
77 prefix: str = ""):
78 super().__init__()
79 self.hidden_size = hidden_size
80 + self.is_sequence_parallel = is_sequence_parallel
81
82 # Gate always runs at half / full precision for now.
83 self.gate = ReplicatedLinear(hidden_size,
@@ -92,15 +96,27 @@ class GraniteMoeMoE(nn.Module):
96 renormalize=True,
97 quant_config=quant_config,
98 tp_size=tp_size,
95 - prefix=f"{prefix}.experts")
99 + prefix=f"{prefix}.experts",
100 + is_sequence_parallel=self.is_sequence_parallel)
101
102 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
103 # NOTE: hidden_states can have either 1D or 2D shape.
104 orig_shape = hidden_states.shape
105 hidden_states = hidden_states.view(-1, self.hidden_size)
106 +
107 + if self.is_sequence_parallel:
108 + hidden_states = sequence_parallel_chunk(hidden_states)
109 +
110 # router_logits: (num_tokens, n_experts)
111 router_logits, _ = self.gate(hidden_states)
112 final_hidden_states = self.experts(hidden_states, router_logits)
113 +
114 + if self.is_sequence_parallel:
115 + final_hidden_states = tensor_model_parallel_all_gather(
116 + final_hidden_states, 0)
117 + num_tokens = orig_shape[0]
118 + final_hidden_states = final_hidden_states[:num_tokens]
119 +
120 return final_hidden_states.view(orig_shape)
121
122
@@ -191,12 +207,16 @@ class GraniteMoeDecoderLayer(nn.Module):
207
208 def __init__(
209 self,
194 - config: GraniteMoeConfig,
195 - cache_config: Optional[CacheConfig] = None,
196 - quant_config: Optional[QuantizationConfig] = None,
210 + vllm_config: VllmConfig,
211 prefix: str = "",
212 ) -> None:
213 super().__init__()
214 +
215 + config = vllm_config.model_config.hf_config
216 + cache_config = vllm_config.cache_config
217 + quant_config = vllm_config.quant_config
218 + parallel_config = vllm_config.parallel_config
219 +
220 self.hidden_size = config.hidden_size
221 # Requires transformers > 4.32.0
222 rope_theta = getattr(config, "rope_theta", 10000)
@@ -218,6 +238,7 @@ class GraniteMoeDecoderLayer(nn.Module):
238 hidden_size=config.hidden_size,
239 intermediate_size=config.intermediate_size,
240 quant_config=quant_config,
241 + is_sequence_parallel=parallel_config.use_sequence_parallel_moe,
242 prefix=f"{prefix}.block_sparse_moe")
243
244 self.input_layernorm = RMSNorm(config.hidden_size,
@@ -255,7 +276,6 @@ class GraniteMoeModel(nn.Module):
276 super().__init__()
277
278 config = vllm_config.model_config.hf_config
258 - cache_config = vllm_config.cache_config
279 quant_config = vllm_config.quant_config
280 lora_config = vllm_config.lora_config
281
@@ -275,9 +295,7 @@ class GraniteMoeModel(nn.Module):
295
296 self.start_layer, self.end_layer, self.layers = make_layers(
297 config.num_hidden_layers,
278 - lambda prefix: GraniteMoeDecoderLayer(
279 - config, cache_config, quant_config=quant_config, prefix=prefix
280 - ),
298 + lambda prefix: GraniteMoeDecoderLayer(vllm_config, prefix=prefix),
299 prefix=f"{prefix}.layers")
300
301 self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
vllm/model_executor/models/llama.py
+13 -12
@@ -68,6 +68,7 @@ class LlamaMLP(nn.Module):
68 bias: bool = False,
69 prefix: str = "",
70 reduce_results: bool = True,
71 + disable_tp: bool = False,
72 ) -> None:
73 super().__init__()
74 self.gate_up_proj = MergedColumnParallelLinear(
@@ -75,6 +76,7 @@ class LlamaMLP(nn.Module):
76 output_sizes=[intermediate_size] * 2,
77 bias=bias,
78 quant_config=quant_config,
79 + disable_tp=disable_tp,
80 prefix=f"{prefix}.gate_up_proj",
81 )
82 self.down_proj = RowParallelLinear(
@@ -83,6 +85,7 @@ class LlamaMLP(nn.Module):
85 bias=bias,
86 quant_config=quant_config,
87 reduce_results=reduce_results,
88 + disable_tp=disable_tp,
89 prefix=f"{prefix}.down_proj",
90 )
91 if hidden_act != "silu":
@@ -237,14 +240,16 @@ class LlamaAttention(nn.Module):
240
241 class LlamaDecoderLayer(nn.Module):
242
240 - def __init__(
241 - self,
242 - config: LlamaConfig,
243 - cache_config: Optional[CacheConfig] = None,
244 - quant_config: Optional[QuantizationConfig] = None,
245 - prefix: str = "",
246 - ) -> None:
243 + def __init__(self,
244 + vllm_config: VllmConfig,
245 + prefix: str = "",
246 + config: Optional[LlamaConfig] = None) -> None:
247 super().__init__()
248 +
249 + config = config or vllm_config.model_config.hf_config
250 + cache_config = vllm_config.cache_config
251 + quant_config = vllm_config.quant_config
252 +
253 self.hidden_size = config.hidden_size
254 rope_theta = getattr(config, "rope_theta", 10000)
255 rope_scaling = getattr(config, "rope_scaling", None)
@@ -335,7 +340,6 @@ class LlamaModel(nn.Module):
340 super().__init__()
341
342 config = vllm_config.model_config.hf_config
338 - cache_config = vllm_config.cache_config
343 quant_config = vllm_config.quant_config
344 lora_config = vllm_config.lora_config
345
@@ -357,10 +361,7 @@ class LlamaModel(nn.Module):
361 self.embed_tokens = PPMissingLayer()
362 self.start_layer, self.end_layer, self.layers = make_layers(
363 config.num_hidden_layers,
360 - lambda prefix: layer_type(config=config,
361 - cache_config=cache_config,
362 - quant_config=quant_config,
363 - prefix=prefix),
364 + lambda prefix: layer_type(vllm_config=vllm_config, prefix=prefix),
365 prefix=f"{prefix}.layers",
366 )
367 if get_pp_group().is_last_rank:
vllm/model_executor/models/llama4.py
+29 -15
@@ -28,7 +28,8 @@ from vllm.attention import Attention
28 from vllm.attention.layers.chunked_local_attention import ChunkedLocalAttention
29 from vllm.compilation.decorators import support_torch_compile
30 from vllm.config import CacheConfig, VllmConfig
31 -from vllm.distributed import get_tensor_model_parallel_world_size
31 +from vllm.distributed import (get_tensor_model_parallel_world_size,
32 + tensor_model_parallel_all_gather)
33 from vllm.model_executor.layers.fused_moe import FusedMoE
34 from vllm.model_executor.layers.layernorm import RMSNorm
35 from vllm.model_executor.layers.linear import (QKVParallelLinear,
@@ -39,6 +40,7 @@ from vllm.model_executor.layers.rotary_embedding import get_rope
40 from vllm.model_executor.layers.shared_fused_moe import SharedFusedMoE
41 from vllm.model_executor.model_loader.weight_utils import (
42 default_weight_loader, maybe_remap_kv_scale_name)
43 +from vllm.model_executor.models.utils import sequence_parallel_chunk
44
45 from .llama import LlamaForCausalLM, LlamaMLP, LlamaModel
46 from .utils import (AutoWeightsLoader, extract_layer_index, fast_topk,
@@ -59,13 +61,16 @@ class Llama4MoE(nn.Module):
61 router_scores = torch.sigmoid(router_scores.float())
62 return (router_scores, router_indices.to(torch.int32))
63
62 - def __init__(self,
63 - config: Llama4TextConfig,
64 - quant_config: Optional[QuantizationConfig] = None,
65 - prefix: str = ""):
64 + def __init__(self, vllm_config: VllmConfig, prefix: str = ""):
65 super().__init__()
66 +
67 + config = vllm_config.model_config.hf_config
68 + parallel_config = vllm_config.parallel_config
69 + quant_config = vllm_config.quant_config
70 +
71 self.tp_size = get_tensor_model_parallel_world_size()
72 self.top_k = config.num_experts_per_tok
73 + self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
74
75 intermediate_size_moe = config.intermediate_size
76 self.router = ReplicatedLinear(config.hidden_size,
@@ -82,6 +87,7 @@ class Llama4MoE(nn.Module):
87 bias=False,
88 prefix=f"{prefix}.shared_expert",
89 reduce_results=False,
90 + disable_tp=self.is_sequence_parallel,
91 )
92
93 self.experts = SharedFusedMoE(
@@ -96,9 +102,14 @@ class Llama4MoE(nn.Module):
102 renormalize=False,
103 quant_config=quant_config,
104 prefix=f"{prefix}.experts",
105 + is_sequence_parallel=self.is_sequence_parallel,
106 )
107
108 def forward(self, hidden_states):
109 + num_tokens = hidden_states.shape[0]
110 + if self.is_sequence_parallel:
111 + hidden_states = sequence_parallel_chunk(hidden_states)
112 +
113 router_logits, _ = self.router(hidden_states)
114
115 shared_out, routed_out = self.experts(
@@ -107,7 +118,10 @@ class Llama4MoE(nn.Module):
118 )
119 experts_out = routed_out + shared_out
120
110 - if self.tp_size > 1:
121 + if self.is_sequence_parallel:
122 + experts_out = tensor_model_parallel_all_gather(experts_out, 0)
123 + experts_out = experts_out[:num_tokens]
124 + elif self.tp_size > 1:
125 experts_out = self.experts.maybe_all_reduce_tensor_model_parallel(
126 experts_out)
127
@@ -257,15 +271,16 @@ class Llama4Attention(nn.Module):
271
272 class Llama4DecoderLayer(nn.Module):
273
260 - def __init__(
261 - self,
262 - config: Llama4TextConfig,
263 - cache_config: Optional[CacheConfig] = None,
264 - quant_config: Optional[QuantizationConfig] = None,
265 - prefix: str = "",
266 - ) -> None:
274 + def __init__(self,
275 + vllm_config: VllmConfig,
276 + prefix: str = "",
277 + config: Optional[Llama4TextConfig] = None) -> None:
278 super().__init__()
279
280 + config = config or vllm_config.model_config.hf_config
281 + cache_config = vllm_config.cache_config
282 + quant_config = vllm_config.quant_config
283 +
284 self.layer_idx = extract_layer_index(prefix)
285 self.global_layer = config.no_rope_layers[self.layer_idx] == 0
286 self.hidden_size = config.hidden_size
@@ -291,8 +306,7 @@ class Llama4DecoderLayer(nn.Module):
306 self.layer_idx + 1) % config.interleave_moe_layer_step == 0
307 if is_moe_layer:
308 self.feed_forward = Llama4MoE(
294 - config=config,
295 - quant_config=quant_config,
309 + vllm_config=vllm_config,
310 prefix=f"{prefix}.feed_forward",
311 )
312 else:
vllm/model_executor/models/llama4_eagle.py
+2 -2
@@ -68,9 +68,9 @@ class LlamaModel(nn.Module):
68
69 self.layers = nn.ModuleList([
70 Llama4DecoderLayer(
71 - self.config,
72 - quant_config=quant_config,
71 + vllm_config=vllm_config,
72 prefix=maybe_prefix(prefix, f"layers.{i + start_layer_id}"),
73 + config=self.config,
74 ) for i in range(self.config.num_hidden_layers)
75 ])
76 self.fc = torch.nn.Linear(self.config.hidden_size * 2,
vllm/model_executor/models/llama_eagle.py
+5 -3
@@ -28,11 +28,12 @@ class LlamaDecoderLayer(LlamaDecoderLayer):
28
29 def __init__(
30 self,
31 - config: LlamaConfig,
31 + vllm_config: VllmConfig,
32 disable_input_layernorm: bool,
33 prefix: str = "",
34 + config: Optional[LlamaConfig] = None,
35 ) -> None:
35 - super().__init__(config, prefix=prefix)
36 + super().__init__(vllm_config, prefix=prefix, config=config)
37
38 # Skip the input_layernorm
39 # https://github.com/SafeAILab/EAGLE/blob/35c78f6cdc19a73e05cf5c330b4c358dad970c6a/eagle/model/cnets.py#L427
@@ -64,9 +65,10 @@ class LlamaModel(nn.Module):
65
66 self.layers = nn.ModuleList([
67 LlamaDecoderLayer(
67 - self.config,
68 + vllm_config,
69 i == 0,
70 prefix=maybe_prefix(prefix, f"layers.{i + start_layer_id}"),
71 + config=self.config,
72 ) for i in range(self.config.num_hidden_layers)
73 ])
74 self.fc = torch.nn.Linear(self.config.hidden_size * 2,
vllm/model_executor/models/llama_eagle3.py
+11 -16
@@ -8,13 +8,11 @@ import torch
8 import torch.nn as nn
9 from transformers import LlamaConfig
10
11 -from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
11 +from vllm.config import VllmConfig, get_current_vllm_config
12 from vllm.logger import init_logger
13 from vllm.model_executor.layers.layernorm import RMSNorm
14 from vllm.model_executor.layers.linear import QKVParallelLinear
15 from vllm.model_executor.layers.logits_processor import LogitsProcessor
16 -from vllm.model_executor.layers.quantization.base_config import (
17 - QuantizationConfig)
16 from vllm.model_executor.layers.vocab_parallel_embedding import (
17 DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
18 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
@@ -28,17 +26,14 @@ logger = init_logger(__name__)
26
27 class LlamaDecoderLayer(LlamaDecoderLayer):
28
31 - def __init__(
32 - self,
33 - config: LlamaConfig,
34 - cache_config: Optional[CacheConfig] = None,
35 - quant_config: Optional[QuantizationConfig] = None,
36 - prefix: str = "",
37 - ) -> None:
38 - super().__init__(config,
39 - cache_config=cache_config,
40 - quant_config=quant_config,
41 - prefix=prefix)
29 + def __init__(self,
30 + vllm_config: VllmConfig,
31 + prefix: str = "",
32 + config: Optional[LlamaConfig] = None) -> None:
33 + super().__init__(vllm_config, prefix=prefix, config=config)
34 +
35 + config = config or vllm_config.model_config.hf_config
36 + quant_config = vllm_config.quant_config
37
38 # override qkv
39 self.self_attn.qkv_proj = QKVParallelLinear(
@@ -125,9 +120,9 @@ class LlamaModel(nn.Module):
120
121 self.layers = nn.ModuleList([
122 LlamaDecoderLayer(
128 - config=self.config,
129 - cache_config=current_vllm_config.cache_config,
123 + current_vllm_config,
124 prefix=maybe_prefix(prefix, f"layers.{start_layer_id}"),
125 + config=self.config,
126 )
127 ])
128 if hasattr(self.config, "target_hidden_size"):
vllm/model_executor/models/qwen3_moe.py
+33 -27
@@ -29,13 +29,13 @@ from typing import Any, Optional, Union
29
30 import torch
31 from torch import nn
32 -from transformers import Qwen3MoeConfig
32
33 from vllm.attention import Attention
34 from vllm.compilation.decorators import support_torch_compile
35 from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
36 from vllm.distributed import (get_ep_group, get_pp_group,
38 - get_tensor_model_parallel_world_size)
37 + get_tensor_model_parallel_world_size,
38 + tensor_model_parallel_all_gather)
39 from vllm.logger import init_logger
40 from vllm.model_executor.layers.activation import SiluAndMul
41 from vllm.model_executor.layers.fused_moe import FusedMoE
@@ -51,6 +51,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
51 ParallelLMHead, VocabParallelEmbedding)
52 from vllm.model_executor.model_loader.weight_utils import (
53 default_weight_loader, maybe_remap_kv_scale_name)
54 +from vllm.model_executor.models.utils import sequence_parallel_chunk
55 from vllm.sequence import IntermediateTensors
56
57 from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP
@@ -101,12 +102,15 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
102
103 def __init__(
104 self,
104 - config: Qwen3MoeConfig,
105 - quant_config: Optional[QuantizationConfig] = None,
105 + vllm_config: VllmConfig,
106 prefix: str = "",
107 - enable_eplb: bool = False,
107 ):
108 super().__init__()
109 +
110 + config = vllm_config.model_config.hf_config
111 + parallel_config = vllm_config.parallel_config
112 + quant_config = vllm_config.quant_config
113 +
114 self.tp_size = get_tensor_model_parallel_world_size()
115
116 self.ep_group = get_ep_group().device_group
@@ -114,6 +118,8 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
118 self.ep_size = self.ep_group.size()
119 self.n_routed_experts = config.num_experts
120
121 + self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
122 +
123 if self.tp_size > config.num_experts:
124 raise ValueError(
125 f"Tensor parallel size {self.tp_size} is greater than "
@@ -122,7 +128,7 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
128 # Load balancing settings.
129 vllm_config = get_current_vllm_config()
130 eplb_config = vllm_config.parallel_config.eplb_config
125 - self.enable_eplb = enable_eplb
131 + self.enable_eplb = parallel_config.enable_eplb
132
133 self.n_logical_experts = self.n_routed_experts
134 self.n_redundant_experts = eplb_config.num_redundant_experts
@@ -144,7 +150,8 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
150 quant_config=quant_config,
151 prefix=f"{prefix}.experts",
152 enable_eplb=self.enable_eplb,
147 - num_redundant_experts=self.n_redundant_experts)
153 + num_redundant_experts=self.n_redundant_experts,
154 + is_sequence_parallel=self.is_sequence_parallel)
155
156 self.gate = ReplicatedLinear(config.hidden_size,
157 config.num_experts,
@@ -156,14 +163,22 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
163 assert hidden_states.dim(
164 ) <= 2, "Qwen3MoeSparseMoeBlock only supports 1D or 2D inputs"
165 is_input_1d = hidden_states.dim() == 1
159 - hidden_dim = hidden_states.shape[-1]
166 + num_tokens, hidden_dim = hidden_states.shape
167 hidden_states = hidden_states.view(-1, hidden_dim)
168
169 + if self.is_sequence_parallel:
170 + hidden_states = sequence_parallel_chunk(hidden_states)
171 +
172 # router_logits: (num_tokens, n_experts)
173 router_logits, _ = self.gate(hidden_states)
174 final_hidden_states = self.experts(hidden_states=hidden_states,
175 router_logits=router_logits)
176
177 + if self.is_sequence_parallel:
178 + final_hidden_states = tensor_model_parallel_all_gather(
179 + final_hidden_states, 0)
180 + final_hidden_states = final_hidden_states[:num_tokens]
181 +
182 # return to 1d if input is 1d
183 return final_hidden_states.squeeze(0) if is_input_1d else \
184 final_hidden_states
@@ -275,15 +290,13 @@ class Qwen3MoeAttention(nn.Module):
290
291 class Qwen3MoeDecoderLayer(nn.Module):
292
278 - def __init__(
279 - self,
280 - config: Qwen3MoeConfig,
281 - cache_config: Optional[CacheConfig] = None,
282 - quant_config: Optional[QuantizationConfig] = None,
283 - prefix: str = "",
284 - enable_eplb: bool = False,
285 - ) -> None:
293 + def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:
294 super().__init__()
295 +
296 + config = vllm_config.model_config.hf_config
297 + cache_config = vllm_config.cache_config
298 + quant_config = vllm_config.quant_config
299 +
300 self.hidden_size = config.hidden_size
301 rope_theta = getattr(config, "rope_theta", 10000)
302 rope_scaling = getattr(config, "rope_scaling", None)
@@ -315,10 +328,8 @@ class Qwen3MoeDecoderLayer(nn.Module):
328 if (layer_idx not in mlp_only_layers) and (
329 config.num_experts > 0 and
330 (layer_idx + 1) % config.decoder_sparse_step == 0):
318 - self.mlp = Qwen3MoeSparseMoeBlock(config=config,
319 - quant_config=quant_config,
320 - prefix=f"{prefix}.mlp",
321 - enable_eplb=enable_eplb)
331 + self.mlp = Qwen3MoeSparseMoeBlock(vllm_config=vllm_config,
332 + prefix=f"{prefix}.mlp")
333 else:
334 self.mlp = Qwen3MoeMLP(hidden_size=config.hidden_size,
335 intermediate_size=config.intermediate_size,
@@ -362,10 +373,8 @@ class Qwen3MoeModel(nn.Module):
373 super().__init__()
374
375 config = vllm_config.model_config.hf_config.get_text_config()
365 - cache_config = vllm_config.cache_config
376 quant_config = vllm_config.quant_config
377 parallel_config = vllm_config.parallel_config
368 - enable_eplb = parallel_config.enable_eplb
378 eplb_config = parallel_config.eplb_config
379 self.num_redundant_experts = eplb_config.num_redundant_experts
380
@@ -379,11 +388,8 @@ class Qwen3MoeModel(nn.Module):
388 prefix=f"{prefix}.embed_tokens")
389 self.start_layer, self.end_layer, self.layers = make_layers(
390 config.num_hidden_layers,
382 - lambda prefix: Qwen3MoeDecoderLayer(config=config,
383 - cache_config=cache_config,
384 - quant_config=quant_config,
385 - prefix=prefix,
386 - enable_eplb=enable_eplb),
391 + lambda prefix: Qwen3MoeDecoderLayer(vllm_config=vllm_config,
392 + prefix=prefix),
393 prefix=f"{prefix}.layers",
394 )
395 self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
vllm/model_executor/models/qwen3_next.py
+35 -35
@@ -17,7 +17,8 @@ from vllm.config import (CacheConfig, ModelConfig, SpeculativeConfig,
17 VllmConfig, get_current_vllm_config)
18 from vllm.distributed import (divide, get_ep_group, get_pp_group,
19 get_tensor_model_parallel_rank,
20 - get_tensor_model_parallel_world_size)
20 + get_tensor_model_parallel_world_size,
21 + tensor_model_parallel_all_gather)
22 from vllm.forward_context import ForwardContext, get_forward_context
23 from vllm.logger import init_logger
24 from vllm.model_executor.layers.fla.ops import (
@@ -47,6 +48,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
48 from vllm.model_executor.model_loader.weight_utils import (
49 default_weight_loader, sharded_weight_loader)
50 from vllm.model_executor.models.qwen2_moe import Qwen2MoeMLP as Qwen3NextMLP
51 +from vllm.model_executor.models.utils import sequence_parallel_chunk
52 from vllm.model_executor.utils import set_weight_attrs
53 from vllm.platforms import current_platform
54 from vllm.sequence import IntermediateTensors
@@ -69,14 +71,13 @@ KVCache = tuple[torch.Tensor, torch.Tensor]
71
72 class Qwen3NextSparseMoeBlock(nn.Module):
73
72 - def __init__(
73 - self,
74 - config: Qwen3NextConfig,
75 - quant_config: Optional[QuantizationConfig] = None,
76 - prefix: str = "",
77 - enable_eplb: bool = False,
78 - ):
74 + def __init__(self, vllm_config: VllmConfig, prefix: str = ""):
75 super().__init__()
76 +
77 + config = vllm_config.model_config.hf_config
78 + parallel_config = vllm_config.parallel_config
79 + quant_config = vllm_config.quant_config
80 +
81 self.tp_size = get_tensor_model_parallel_world_size()
82
83 self.ep_group = get_ep_group().device_group
@@ -84,6 +85,8 @@ class Qwen3NextSparseMoeBlock(nn.Module):
85 self.ep_size = self.ep_group.size()
86 self.n_routed_experts = config.num_experts
87
88 + self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
89 +
90 if self.tp_size > config.num_experts:
91 raise ValueError(
92 f"Tensor parallel size {self.tp_size} is greater than "
@@ -92,7 +95,7 @@ class Qwen3NextSparseMoeBlock(nn.Module):
95 # Load balancing settings.
96 vllm_config = get_current_vllm_config()
97 eplb_config = vllm_config.parallel_config.eplb_config
95 - self.enable_eplb = enable_eplb
98 + self.enable_eplb = parallel_config.enable_eplb
99
100 self.n_logical_experts = self.n_routed_experts
101 self.n_redundant_experts = eplb_config.num_redundant_experts
@@ -114,7 +117,8 @@ class Qwen3NextSparseMoeBlock(nn.Module):
117 quant_config=quant_config,
118 prefix=f"{prefix}.experts",
119 enable_eplb=self.enable_eplb,
117 - num_redundant_experts=self.n_redundant_experts)
120 + num_redundant_experts=self.n_redundant_experts,
121 + is_sequence_parallel=self.is_sequence_parallel)
122
123 self.gate = ReplicatedLinear(config.hidden_size,
124 config.num_experts,
@@ -141,9 +145,12 @@ class Qwen3NextSparseMoeBlock(nn.Module):
145 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
146 # NOTE: hidden_states can have either 1D or 2D shape.
147 orig_shape = hidden_states.shape
144 - hidden_dim = hidden_states.shape[-1]
148 + num_tokens, hidden_dim = hidden_states.shape
149 hidden_states = hidden_states.view(-1, hidden_dim)
150
151 + if self.is_sequence_parallel:
152 + hidden_states = sequence_parallel_chunk(hidden_states)
153 +
154 shared_output = None
155 if self.shared_expert is not None:
156 shared_output = self.shared_expert(hidden_states)
@@ -158,7 +165,12 @@ class Qwen3NextSparseMoeBlock(nn.Module):
165
166 if shared_output is not None:
167 final_hidden_states = final_hidden_states + shared_output
161 - if self.tp_size > 1:
168 +
169 + if self.is_sequence_parallel:
170 + final_hidden_states = tensor_model_parallel_all_gather(
171 + final_hidden_states, 0)
172 + final_hidden_states = final_hidden_states[:num_tokens]
173 + elif self.tp_size > 1:
174 final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( # noqa E501
175 final_hidden_states)
176
@@ -719,17 +731,17 @@ class Qwen3NextDecoderLayer(nn.Module):
731
732 def __init__(
733 self,
722 - config: Qwen3NextConfig,
734 + vllm_config: VllmConfig,
735 layer_type: str,
724 - model_config: Optional[ModelConfig] = None,
725 - cache_config: Optional[CacheConfig] = None,
726 - quant_config: Optional[QuantizationConfig] = None,
727 - speculative_config: Optional[SpeculativeConfig] = None,
736 prefix: str = "",
729 - enable_eplb: bool = False,
737 ) -> None:
738 super().__init__()
732 - self.config = config
739 +
740 + config = vllm_config.model_config.hf_config
741 + model_config = vllm_config.model_config
742 + cache_config = vllm_config.cache_config
743 + quant_config = vllm_config.quant_config
744 + speculative_config = vllm_config.speculative_config
745
746 self.layer_type = layer_type
747 self.layer_idx = extract_layer_index(prefix)
@@ -759,10 +771,8 @@ class Qwen3NextDecoderLayer(nn.Module):
771 config.num_experts > 0 and
772 (self.layer_idx + 1) % config.decoder_sparse_step == 0):
773 self.mlp = Qwen3NextSparseMoeBlock(
762 - config=config,
763 - quant_config=quant_config,
774 + vllm_config=vllm_config,
775 prefix=f"{prefix}.mlp",
765 - enable_eplb=enable_eplb,
776 )
777 else:
778 self.mlp = Qwen3NextMLP(
@@ -783,14 +793,14 @@ class Qwen3NextDecoderLayer(nn.Module):
793 torch.zeros(
794 1,
795 1,
786 - self.config.hidden_size,
796 + config.hidden_size,
797 dtype=config.torch_dtype,
798 ), )
799 self.ffn_layer_scale = torch.nn.Parameter(
800 torch.zeros(
801 1,
802 1,
793 - self.config.hidden_size,
803 + config.hidden_size,
804 dtype=config.torch_dtype,
805 ), )
806
@@ -858,13 +868,8 @@ class Qwen3NextModel(nn.Module):
868 super().__init__()
869
870 config: Qwen3NextConfig = vllm_config.model_config.hf_config
861 - model_config = vllm_config.model_config
862 - cache_config = vllm_config.cache_config
863 - quant_config = vllm_config.quant_config
871 parallel_config = vllm_config.parallel_config
872 lora_config = vllm_config.lora_config
866 - speculative_config = vllm_config.speculative_config
867 - enable_eplb = parallel_config.enable_eplb
873 eplb_config = parallel_config.eplb_config
874 self.num_redundant_experts = eplb_config.num_redundant_experts
875
@@ -881,14 +886,9 @@ class Qwen3NextModel(nn.Module):
886
887 def get_layer(prefix: str):
888 return Qwen3NextDecoderLayer(
884 - config,
889 + vllm_config,
890 layer_type=config.layer_types[extract_layer_index(prefix)],
886 - model_config=model_config,
887 - cache_config=cache_config,
888 - quant_config=quant_config,
889 - speculative_config=speculative_config,
891 prefix=prefix,
891 - enable_eplb=enable_eplb,
892 )
893
894 self.start_layer, self.end_layer, self.layers = make_layers(
vllm/model_executor/models/qwen3_next_mtp.py
+1 -5
@@ -38,7 +38,6 @@ class Qwen3NextMultiTokenPredictor(nn.Module):
38 super().__init__()
39
40 model_config = vllm_config.model_config
41 - cache_config = vllm_config.cache_config
41 quant_config = vllm_config.quant_config
42 lora_config = vllm_config.lora_config
43 config: Qwen3NextConfig = model_config.hf_config
@@ -68,11 +67,8 @@ class Qwen3NextMultiTokenPredictor(nn.Module):
67
68 self.layers = torch.nn.ModuleList(
69 Qwen3NextDecoderLayer(
71 - config,
70 + vllm_config,
71 layer_type="full_attention",
73 - model_config=model_config,
74 - cache_config=cache_config,
75 - quant_config=quant_config,
72 prefix=f'{prefix}.layers.{idx}',
73 ) for idx in range(self.num_mtp_layers))
74
vllm/model_executor/models/utils.py
+47 -1
@@ -13,11 +13,14 @@ from transformers import PretrainedConfig
13
14 import vllm.envs as envs
15 from vllm.config import VllmConfig
16 +from vllm.distributed import (get_tensor_model_parallel_rank,
17 + get_tensor_model_parallel_world_size)
18 from vllm.logger import init_logger
19 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
20 from vllm.multimodal import NestedTensors
21 from vllm.sequence import IntermediateTensors
20 -from vllm.utils import (get_cuda_view_from_cpu_tensor, is_pin_memory_available,
22 +from vllm.utils import (cdiv, direct_register_custom_op,
23 + get_cuda_view_from_cpu_tensor, is_pin_memory_available,
24 is_uva_available)
25
26 logger = init_logger(__name__)
@@ -743,3 +746,46 @@ def get_model_hidden_size(hf_config: PretrainedConfig) -> int:
746 return hf_config.hidden_size
747 text_config = hf_config.get_text_config()
748 return text_config.hidden_size
749 +
750 +
751 +# Chunk x along the num_tokens axis for sequence parallelism
752 +# NOTE: This is wrapped in a torch custom op to work around the following issue:
753 +# The output tensor can have a sequence length 0 at small input sequence lengths
754 +# even though we explicitly pad to avoid this.
755 +def sequence_parallel_chunk(x: torch.Tensor) -> torch.Tensor:
756 + return torch.ops.vllm.sequence_parallel_chunk_impl(x)
757 +
758 +
759 +def sequence_parallel_chunk_impl(x: torch.Tensor) -> torch.Tensor:
760 + tp_size = get_tensor_model_parallel_world_size()
761 + tp_rank = get_tensor_model_parallel_rank()
762 +
763 + # all_gather needs the sequence length to be divisible by tp_size
764 + seq_len = x.size(0)
765 + remainder = seq_len % tp_size
766 + if remainder != 0:
767 + pad_len = tp_size - remainder
768 + y = nn.functional.pad(x, (0, 0, 0, pad_len))
769 + else:
770 + y = x
771 +
772 + chunk = y.shape[0] // tp_size
773 + start = tp_rank * chunk
774 + return torch.narrow(y, 0, start, chunk)
775 +
776 +
777 +def sequence_parallel_chunk_impl_fake(x: torch.Tensor) -> torch.Tensor:
778 + tp_size = get_tensor_model_parallel_world_size()
779 + seq_len = cdiv(x.size(0), tp_size)
780 + shape = list(x.shape)
781 + shape[0] = seq_len
782 + out = torch.empty(shape, dtype=x.dtype, device=x.device)
783 + return out
784 +
785 +
786 +direct_register_custom_op(
787 + op_name="sequence_parallel_chunk_impl",
788 + op_func=sequence_parallel_chunk_impl,
789 + fake_impl=sequence_parallel_chunk_impl_fake,
790 + tags=(torch.Tag.needs_fixed_stride_order, ),
791 +)