@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e9b92dcd89
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

[Kernels] Overlap shared experts with send/recv (#23273)

Signed-off-by: Bill Nell <bnell@redhat.com>

bnellnm committed Sep 3, 2025 at 12:35 UTC e9b92dcd89e8d05f162a8fdaa3d5d60012615514
32 files changed +884 -226
docs/design/fused_moe_modular_kernel.md
+6 -2
@@ -54,8 +54,8 @@ The `FusedMoEModularKernel` acts as a bridge between the `FusedMoEPermuteExperts
54
55 ### FusedMoEPrepareAndFinalize
56
57 -The `FusedMoEPrepareAndFinalize` abstract class exposes `prepare` and `finalize` functions.
58 -The `prepare` function is responsible for input activation Quantization and All2All Dispatch. The `finalize` function is responsible for invoking the All2All Combine. Additionally the `finalize` function may or may not do the TopK weight application and reduction (Please refer to the TopKWeightAndReduce section)
57 +The `FusedMoEPrepareAndFinalize` abstract class exposes `prepare`, `prepare_no_receive` and `finalize` functions.
58 +The `prepare` function is responsible for input activation Quantization and All2All Dispatch. If implemented, The `prepare_no_receive` is like `prepare` except it does not wait to receive results from other workers. Instead it returns a "receiver" callback that must be invoked to wait for the final results of worker. It is not required that this method is supported by all `FusedMoEPrepareAndFinalize` classes, but if it is available, it can be used to interleave work with the initial all to all communication, e.g. interleaving shared experts with fused experts. The `finalize` function is responsible for invoking the All2All Combine. Additionally the `finalize` function may or may not do the TopK weight application and reduction (Please refer to the TopKWeightAndReduce section)
59
60 ![](../assets/design/fused_moe_modular_kernel/prepare_and_finalize_blocks.png "FusedMoEPrepareAndFinalize Blocks")
61
@@ -146,6 +146,10 @@ This section describes the significance of the various functions exposed by the
146
147 `FusedMoEPrepareAndFinalize::prepare()`: The prepare method implements the Quantization and All2All Dispatch. Typically the Dispatch function from the relevant All2All Manager is invoked.
148
149 +`FusedMoEPrepareAndFinalize::has_prepare_no_receive()`: Indicates whether or not this subclass implements `prepare_no_receive`. Defaults to False.
150 +
151 +`FusedMoEPrepareAndFinalize::prepare_no_receive()`: The prepare_no_receive method implements the Quantization and All2All Dispatch. It does not wait for the result of the dispatch operation but instead returns a thunk that can be invoked to wait for the final results. Typically the Dispatch function from the relevant All2All Manager is invoked.
152 +
153 `FusedMoEPrepareAndFinalize::finalize()`: Maybe perform TopK Weight Application and Reduction and All2All Combine. Typically the Combine function from the relevant All2AllManager is invoked.
154
155 `FusedMoEPrepareAndFinalize::activation_format()`: Return `FusedMoEActivationFormat.BatchedExperts` if the output of the prepare method (i.e. the All2All dispatch) is Batched. Return `FusedMoEActivationFormat.Standard` otherwise.
examples/offline_inference/data_parallel.py
+8
@@ -87,6 +87,11 @@ def parse_args():
87 default=0.8,
88 help=("Fraction of GPU memory vLLM is allowed to allocate (0.0, 1.0]."),
89 )
90 + parser.add_argument(
91 + "--compilation-config",
92 + type=int,
93 + help=("Compilation optimization (O) level 0-3."),
94 + )
95 parser.add_argument(
96 "--quantization",
97 type=str,
@@ -106,6 +111,7 @@ def main(
111 trust_remote_code,
112 max_num_seqs,
113 max_model_len,
114 + compilation_config,
115 gpu_memory_utilization,
116 quantization,
117 ):
@@ -162,6 +168,7 @@ def main(
168 max_model_len=max_model_len,
169 gpu_memory_utilization=gpu_memory_utilization,
170 quantization=quantization,
171 + compilation_config=compilation_config,
172 )
173 outputs = llm.generate(prompts, sampling_params)
174 # Print the outputs.
@@ -218,6 +225,7 @@ if __name__ == "__main__":
225 args.trust_remote_code,
226 args.max_num_seqs,
227 args.max_model_len,
228 + args.compilation_config,
229 args.gpu_memory_utilization,
230 args.quantization,
231 ),
tests/kernels/moe/test_pplx_moe.py
+74 -12
@@ -4,10 +4,11 @@
4
5 Run `pytest tests/kernels/test_pplx_moe.py`.
6 """
7 +import copy
8 import itertools
9 import textwrap
10 import traceback
10 -from typing import Callable, Optional
11 +from typing import Callable, Optional, Union
12
13 import pytest
14 import torch
@@ -21,7 +22,10 @@ try:
22 except ImportError:
23 has_pplx = False
24
24 -from tests.kernels.moe.utils import make_test_weights, naive_batched_moe
25 +from tests.kernels.moe.modular_kernel_tools.parallel_utils import (
26 + _set_vllm_config)
27 +from tests.kernels.moe.utils import (make_shared_experts, make_test_weights,
28 + naive_batched_moe)
29 from tests.kernels.quant_utils import dequant
30 from tests.kernels.utils import torch_experts
31 from vllm.config import VllmConfig, set_current_vllm_config
@@ -511,7 +515,8 @@ def pplx_moe(
515 block_shape: Optional[list[int]] = None,
516 use_compile: bool = False,
517 use_cudagraphs: bool = True,
514 -) -> torch.Tensor:
518 + shared_experts: Optional[torch.nn.Module] = None,
519 +) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
520
521 num_tokens, hidden_dim = a.shape
522 num_experts = w1.shape[0]
@@ -546,6 +551,7 @@ def pplx_moe(
551 fused_experts = FusedMoEModularKernel(
552 prepare_finalize,
553 experts,
554 + shared_experts,
555 )
556
557 # Note: workers with the same dp_rank must use the exact same inputs.
@@ -586,7 +592,11 @@ def pplx_moe(
592 global_num_experts=num_experts)
593
594 if use_cudagraphs:
589 - out.fill_(0)
595 + if isinstance(out, tuple):
596 + out[0].fill_(0)
597 + out[1].fill_(0)
598 + else:
599 + out.fill_(0)
600 stream = torch.cuda.Stream()
601 graph = torch.cuda.CUDAGraph()
602 with torch.cuda.graph(graph, stream=stream):
@@ -626,6 +636,7 @@ def _pplx_moe(
636 per_act_token_quant: bool = False,
637 block_shape: Optional[list[int]] = None,
638 use_internode: bool = False,
639 + shared_experts: Optional[torch.nn.Module] = None,
640 ):
641 try:
642 if use_internode:
@@ -666,6 +677,11 @@ def _pplx_moe(
677 with set_current_vllm_config(vllm_config), override_config(moe_config):
678 topk_weight, topk_ids, _ = fused_topk(a, score, topk, False)
679
680 + if shared_experts is not None:
681 + shared_output = shared_experts(a)
682 + else:
683 + shared_output = None
684 +
685 torch_output = torch_experts(
686 a,
687 w1,
@@ -696,7 +712,7 @@ def _pplx_moe(
712 block_shape=block_shape,
713 )
714
699 - pplx_output = pplx_moe(
715 + pplx_outputs = pplx_moe(
716 group_name,
717 rank,
718 world_size,
@@ -713,8 +729,24 @@ def _pplx_moe(
729 quant_dtype=quant_dtype,
730 per_act_token_quant=per_act_token_quant,
731 block_shape=block_shape,
732 + shared_experts=shared_experts,
733 )
734
735 + if shared_experts is None:
736 + pplx_shared_output = None
737 + pplx_output = pplx_outputs
738 + assert isinstance(pplx_output, torch.Tensor)
739 + else:
740 + pplx_shared_output, pplx_output = pplx_outputs
741 +
742 + if shared_output is not None:
743 + assert pplx_shared_output is not None
744 + chunked_shared_output = chunk_by_rank(
745 + shared_output, pgi.rank,
746 + pgi.world_size).to(pplx_shared_output.device)
747 + else:
748 + chunked_shared_output = None
749 +
750 chunked_batch_output = chunk_by_rank(
751 batched_output, pgi.rank, pgi.world_size).to(pplx_output.device)
752
@@ -727,6 +759,15 @@ def _pplx_moe(
759 chunked_batch_output,
760 atol=3e-2,
761 rtol=3e-2)
762 +
763 + if shared_experts is not None:
764 + assert chunked_shared_output is not None
765 + assert pplx_shared_output is not None
766 + torch.testing.assert_close(pplx_shared_output,
767 + chunked_shared_output,
768 + atol=3e-2,
769 + rtol=3e-2)
770 +
771 finally:
772 if use_internode:
773 nvshmem_finalize()
@@ -788,7 +829,8 @@ def test_pplx_moe_slow(
829
830
831 def _pplx_test_loop(pgi: ProcessGroupInfo, dp_size: int, use_internode: bool,
791 - make_weights: bool, test_fn: Callable):
832 + use_shared_experts: bool, make_weights: bool,
833 + test_fn: Callable):
834
835 def format_result(msg, ex=None):
836 if ex is not None:
@@ -803,6 +845,14 @@ def _pplx_test_loop(pgi: ProcessGroupInfo, dp_size: int, use_internode: bool,
845 else:
846 print(f"PASSED {msg}")
847
848 + if use_shared_experts:
849 + # Note: this config is only needed for the non-naive shared experts.
850 + new_vllm_config = copy.deepcopy(vllm_config)
851 + new_vllm_config.parallel_config.data_parallel_size = pgi.world_size
852 + new_vllm_config.parallel_config.enable_expert_parallel = True
853 + _set_vllm_config(new_vllm_config, pgi.world_size, pgi.rank,
854 + pgi.local_rank)
855 +
856 current_platform.seed_everything(7)
857 combos = itertools.product(PPLX_COMBOS, NUM_EXPERTS, TOP_KS, DTYPES,
858 [False, True], [None, [128, 128]])
@@ -819,9 +869,11 @@ def _pplx_test_loop(pgi: ProcessGroupInfo, dp_size: int, use_internode: bool,
869 use_fp8_w8a8 = False
870 quant_dtype = None
871
822 - test_desc = (f"test_pplx_moe[mnk={mnk}, e={e}, topk={topk}, "
823 - f"dtype={dtype}, per_act_token={per_act_token_quant}, "
824 - f"block_shape={block_shape}")
872 + test_desc = (
873 + f"test_pplx_moe[mnk={mnk}, e={e}, topk={topk}, "
874 + f"dtype={dtype}, per_act_token={per_act_token_quant}, "
875 + f"block_shape={block_shape}, use_internode={use_internode}, "
876 + f"use_shared_experts={use_shared_experts}")
877
878 if not use_fp8_w8a8 and (per_act_token_quant
879 or block_shape is not None):
@@ -852,6 +904,14 @@ def _pplx_test_loop(pgi: ProcessGroupInfo, dp_size: int, use_internode: bool,
904 args["w1_s"] = w1_s
905 args["w2_s"] = w2_s
906
907 + if use_shared_experts:
908 + args["shared_experts"] = make_shared_experts(
909 + n,
910 + k,
911 + in_dtype=a.dtype,
912 + quant_dtype=quant_dtype,
913 + )
914 +
915 try:
916 test_fn(
917 pgi=pgi,
@@ -891,18 +951,20 @@ def test_pplx_prepare_finalize(
951 current_platform.seed_everything(7)
952 world_size, dp_size = world_dp_size
953 parallel_launch(world_size * dp_size, _pplx_test_loop, dp_size,
894 - use_internode, False, _pplx_prepare_finalize)
954 + use_internode, False, False, _pplx_prepare_finalize)
955
956
957 @pytest.mark.parametrize("world_dp_size", [[2, 1]])
958 @pytest.mark.parametrize("use_internode", [False])
959 +@pytest.mark.parametrize("use_shared_experts", [False, True])
960 @requires_pplx
961 @multi_gpu_test(num_gpus=2)
962 def test_pplx_moe(
963 world_dp_size: tuple[int, int],
964 use_internode: bool,
965 + use_shared_experts: bool,
966 ):
967 current_platform.seed_everything(7)
968 world_size, dp_size = world_dp_size
907 - parallel_launch(world_size, _pplx_test_loop, dp_size, use_internode, True,
908 - _pplx_moe)
969 + parallel_launch(world_size, _pplx_test_loop, dp_size, use_internode,
970 + use_shared_experts, True, _pplx_moe)
tests/kernels/moe/utils.py
+149
@@ -8,6 +8,7 @@ import vllm._custom_ops as ops
8 from tests.kernels.quant_utils import per_block_cast_to_int8
9 from tests.kernels.quantization.nvfp4_utils import (FLOAT4_E2M1_MAX,
10 FLOAT8_E4M3_MAX)
11 +from vllm.model_executor.layers.activation import SiluAndMul
12 from vllm.model_executor.layers.fused_moe import fused_experts
13 from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
14 BatchedPrepareAndFinalize, BatchedTritonExperts, NaiveBatchedExperts)
@@ -282,3 +283,151 @@ def per_token_cast_to_fp8(
283 x_amax = x_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4)
284 fp8_data = (x_view * (448.0 / x_amax.unsqueeze(2))).to(torch.float8_e4m3fn)
285 return fp8_data.view(m, n + pad_size)[:, :n], (x_amax / 448.0).view(m, -1)
286 +
287 +
288 +# CustomOp?
289 +class BaselineMM(torch.nn.Module):
290 +
291 + def __init__(
292 + self,
293 + b: torch.Tensor,
294 + out_dtype: torch.dtype,
295 + ):
296 + super().__init__()
297 + self.b = b.to(dtype=torch.float32)
298 + self.out_dtype = out_dtype
299 +
300 + def forward(
301 + self,
302 + a: torch.Tensor) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
303 + return torch.mm(a.to(dtype=torch.float32),
304 + self.b).to(self.out_dtype), None
305 +
306 +
307 +class TestMLP(torch.nn.Module):
308 +
309 + def __init__(
310 + self,
311 + w1: torch.Tensor,
312 + w2: torch.Tensor,
313 + out_dtype: torch.dtype,
314 + ):
315 + super().__init__()
316 + self.gate_up_proj = BaselineMM(w1, out_dtype)
317 + self.down_proj = BaselineMM(w2, out_dtype)
318 + self.act_fn = SiluAndMul()
319 +
320 + def forward(self, x):
321 + x, _ = self.gate_up_proj(x)
322 + x = self.act_fn(x)
323 + x, _ = self.down_proj(x)
324 + return x
325 +
326 +
327 +def make_naive_shared_experts(
328 + N: int,
329 + K: int,
330 + in_dtype: torch.dtype = torch.bfloat16,
331 +) -> torch.nn.Module:
332 + w1 = torch.randn((K, N * 2), device="cuda", dtype=in_dtype) / 15
333 + w2 = torch.randn((N, K), device="cuda", dtype=in_dtype) / 15
334 + return TestMLP(w1, w2, out_dtype=in_dtype)
335 +
336 +
337 +class RealMLP(torch.nn.Module):
338 +
339 + def __init__(
340 + self,
341 + hidden_size: int,
342 + intermediate_size: int,
343 + w1: torch.Tensor,
344 + w2: torch.Tensor,
345 + hidden_act: str = "silu",
346 + quant_config=None,
347 + reduce_results: bool = True,
348 + prefix: str = "",
349 + w1_s: Optional[torch.Tensor] = None,
350 + w2_s: Optional[torch.Tensor] = None,
351 + ) -> None:
352 + from vllm.model_executor.layers.linear import (
353 + MergedColumnParallelLinear, RowParallelLinear)
354 +
355 + super().__init__()
356 + self.gate_up_proj = MergedColumnParallelLinear(
357 + hidden_size, [intermediate_size] * 2,
358 + bias=False,
359 + quant_config=quant_config,
360 + prefix=f"{prefix}.gate_up_proj")
361 + self.gate_up_proj.register_parameter(
362 + "weight", torch.nn.Parameter(w1, requires_grad=False))
363 + self.gate_up_proj.register_parameter(
364 + "weight_scale", torch.nn.Parameter(w1_s, requires_grad=False))
365 + self.gate_up_proj.register_parameter(
366 + "input_scale",
367 + None) #torch.nn.Parameter(None, requires_grad=False))
368 + self.down_proj = RowParallelLinear(intermediate_size,
369 + hidden_size,
370 + bias=False,
371 + quant_config=quant_config,
372 + reduce_results=reduce_results,
373 + prefix=f"{prefix}.down_proj")
374 + self.down_proj.register_parameter(
375 + "weight", torch.nn.Parameter(w2, requires_grad=False))
376 + self.down_proj.register_parameter(
377 + "weight_scale", torch.nn.Parameter(w2_s, requires_grad=False))
378 + self.down_proj.register_parameter(
379 + "input_scale",
380 + None) #torch.nn.Parameter(None, requires_grad=False))
381 + if hidden_act != "silu":
382 + raise ValueError(f"Unsupported activation: {hidden_act}. "
383 + "Only silu is supported for now.")
384 + self.act_fn = SiluAndMul()
385 +
386 + def forward(self, x):
387 + gate_up, _ = self.gate_up_proj(x)
388 + x = self.act_fn(gate_up)
389 + x, _ = self.down_proj(x)
390 + return x
391 +
392 +
393 +def make_shared_experts(
394 + N: int,
395 + K: int,
396 + in_dtype: torch.dtype = torch.bfloat16,
397 + quant_dtype: Union[torch.dtype, str, None] = None,
398 +) -> torch.nn.Module:
399 + from vllm.model_executor.layers.quantization.fp8 import Fp8Config
400 +
401 + (_, w1, w1_s, _), (_, w2, w2_s, _) = make_test_weights(
402 + 1,
403 + N,
404 + K,
405 + in_dtype=in_dtype,
406 + quant_dtype=quant_dtype,
407 + )
408 + old_dtype = torch.get_default_dtype()
409 + try:
410 + torch.set_default_dtype(in_dtype)
411 + if quant_dtype == torch.float8_e4m3fn:
412 + w1 = w1[0].transpose(0, 1)
413 + w2 = w2[0].transpose(0, 1)
414 + w1_s = w1_s[0].transpose(0, 1) if w1_s is not None else None
415 + w2_s = w2_s[0].transpose(0, 1) if w2_s is not None else None
416 + quant_config = Fp8Config(True)
417 + else:
418 + w1 = w1[0]
419 + w2 = w2[0]
420 + w1_s = None
421 + w2_s = None
422 + quant_config = None
423 +
424 + return RealMLP(K,
425 + N,
426 + w1,
427 + w2,
428 + "silu",
429 + quant_config,
430 + w1_s=w1_s,
431 + w2_s=w2_s)
432 + finally:
433 + torch.set_default_dtype(old_dtype)
vllm/distributed/device_communicators/all2all.py
+1 -6
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from typing import TYPE_CHECKING, Any
3 +from typing import Any
4
5 import torch
6 import torch.distributed as dist
@@ -13,11 +13,6 @@ from .base_device_communicator import All2AllManagerBase, Cache
13
14 logger = init_logger(__name__)
15
16 -if TYPE_CHECKING:
17 - from vllm.model_executor.layers.fused_moe.layer import FusedMoE
18 -else:
19 - FusedMoE = None
20 -
16
17 class NaiveAll2AllManager(All2AllManagerBase):
18 """
vllm/distributed/device_communicators/base_device_communicator.py
+4 -1
@@ -252,7 +252,10 @@ class DeviceCommunicatorBase:
252
253 moe_modules = [
254 module for module in model.modules()
255 - if module.__class__.__name__ == "FusedMoE"
255 + # TODO(bnell): Should use isinstance but can't. Maybe search for
256 + # presence of quant_method.init_prepare_finalize?
257 + if (module.__class__.__name__ == "FusedMoE"
258 + or module.__class__.__name__ == "SharedFusedMoE")
259 ]
260 for module in moe_modules:
261 module.quant_method.init_prepare_finalize(module)
vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py
+91 -39
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from typing import Optional
3 +from typing import Callable, Optional, Union
4
5 import deep_ep
6 import torch
@@ -25,6 +25,8 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
25 self.num_dispatchers_ = num_dispatchers
26 self.dp_size = dp_size
27 self.rank_expert_offset = rank_expert_offset
28 + self.async_prepare = True
29 +
30 # The dispatch function returns a handle that the combine function
31 # requires. We store the handle here so it is available to the
32 # combine function.
@@ -56,10 +58,16 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
58 return None
59 return deep_ep.Buffer.get_combine_config(self.dp_size)
60
59 - def _do_dispatch(self, tokens: torch.Tensor,
60 - token_scales: Optional[torch.Tensor],
61 - rank_topk_ids: torch.Tensor,
62 - rank_topk_weights: torch.Tensor, num_experts: int):
61 + def _do_dispatch(
62 + self,
63 + tokens: torch.Tensor,
64 + token_scales: Optional[torch.Tensor],
65 + rank_topk_ids: torch.Tensor,
66 + rank_topk_weights: torch.Tensor,
67 + num_experts: int,
68 + a1_scale: Optional[torch.Tensor],
69 + quant_config: FusedMoEQuantConfig,
70 + ) -> Callable:
71
72 has_scales = token_scales is not None
73
@@ -93,9 +101,36 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
101 expert_alignment=1,
102 config=self._get_dispatch_config(),
103 previous_event=None,
96 - async_finish=False,
104 + async_finish=self.async_prepare,
105 allocate_on_comm_stream=False)
106
107 + return lambda: self._receiver(
108 + event,
109 + has_scales,
110 + token_data,
111 + expert_topk_ids,
112 + num_experts,
113 + expert_num_tokens_per_expert_list,
114 + expert_topk_weights,
115 + a1_scale,
116 + quant_config,
117 + )
118 +
119 + def _receiver(
120 + self,
121 + event: deep_ep.EventOverlap,
122 + has_scales: bool,
123 + token_data: Union[tuple[torch.Tensor, torch.Tensor], torch.Tensor],
124 + expert_topk_ids: Optional[torch.Tensor],
125 + num_experts: int,
126 + expert_num_tokens_per_expert_list: list[int],
127 + expert_topk_weights: Optional[torch.Tensor],
128 + a1_scale: Optional[torch.Tensor],
129 + quant_config: FusedMoEQuantConfig,
130 + ) -> mk.PrepareResultType:
131 + if self.async_prepare:
132 + event.current_stream_wait()
133 +
134 if has_scales:
135 expert_x, expert_x_scale = token_data
136 else:
@@ -112,6 +147,7 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
147 # DeepEP's topk_ids output refers to the local experts directly. Offset
148 # the topk_ids to move it back to the global experts space so it aligns
149 # with existing vLLM interfaces.
150 + assert expert_topk_ids is not None
151 expert_topk_ids = torch.where(
152 expert_topk_ids == -1,
153 num_experts - 1 if self.rank_expert_offset == 0 else 0,
@@ -123,10 +159,28 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
159 expert_tokens_meta = mk.ExpertTokensMetadata.make_from_list(
160 expert_num_tokens_per_expert_list, device=expert_x.device)
161
162 + # Dispatch and Quant
163 + # DeepEP kernels only support dispatching block-quantized
164 + # activation scales.
165 + # Dispatch in bfloat16 and quantize afterwards
166 + if not quant_config.is_block_quantized:
167 + # Quantize after dispatch.
168 + expert_x_scale = None
169 + if expert_x.numel() != 0:
170 + expert_x, expert_x_scale = moe_kernel_quantize_input(
171 + expert_x,
172 + a1_scale,
173 + quant_dtype=quant_config.quant_dtype,
174 + per_act_token_quant=False,
175 + block_shape=quant_config.block_shape)
176 +
177 return (expert_x, expert_x_scale, expert_tokens_meta, expert_topk_ids,
178 expert_topk_weights)
179
129 - def prepare(
180 + def supports_async(self) -> bool:
181 + return True
182 +
183 + def prepare_async(
184 self,
185 a1: torch.Tensor,
186 a1_scale: Optional[torch.Tensor],
@@ -137,9 +191,7 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
191 expert_map: Optional[torch.Tensor],
192 apply_router_weight_on_input: bool,
193 quant_config: FusedMoEQuantConfig,
140 - ) -> tuple[torch.Tensor, Optional[torch.Tensor],
141 - Optional[mk.ExpertTokensMetadata], Optional[torch.Tensor],
142 - Optional[torch.Tensor]]:
194 + ) -> Callable:
195
196 if apply_router_weight_on_input:
197 topk = topk_ids.size(1)
@@ -159,37 +211,37 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
211 )
212 if a1q_scale is not None and a1q_scale.numel() == 1:
213 a1q_scale = a1q_scale.view(1, 1)
162 - (expert_x, expert_x_scale, expert_tokens_meta, expert_topk_ids,
163 - expert_topk_weights) = self._do_dispatch(
164 - tokens=a1q,
165 - token_scales=a1q_scale,
166 - rank_topk_ids=topk_ids,
167 - rank_topk_weights=topk_weights,
168 - num_experts=num_experts)
214 + a1_post_scale = None
215 else:
170 - # Dispatch and Quant
171 - # DeepEP kernels only support dispatching block-quantized
172 - # activation scales.
173 - # Dispatch in bfloat16
174 - (expert_x, _, expert_tokens_meta, expert_topk_ids,
175 - expert_topk_weights) = self._do_dispatch(
176 - tokens=a1,
177 - token_scales=None,
178 - rank_topk_ids=topk_ids,
179 - rank_topk_weights=topk_weights,
180 - num_experts=num_experts)
181 - # Quantize after dispatch.
182 - expert_x_scale = None
183 - if expert_x.numel() != 0:
184 - expert_x, expert_x_scale = moe_kernel_quantize_input(
185 - expert_x,
186 - a1_scale,
187 - quant_dtype=quant_config.quant_dtype,
188 - per_act_token_quant=False,
189 - block_shape=quant_config.block_shape)
216 + a1q = a1
217 + a1q_scale = None
218 + a1_post_scale = a1_scale
219
191 - return (expert_x, expert_x_scale, expert_tokens_meta, expert_topk_ids,
192 - expert_topk_weights)
220 + return self._do_dispatch(tokens=a1q,
221 + token_scales=a1q_scale,
222 + rank_topk_ids=topk_ids,
223 + rank_topk_weights=topk_weights,
224 + num_experts=num_experts,
225 + a1_scale=a1_post_scale,
226 + quant_config=quant_config)
227 +
228 + def prepare(
229 + self,
230 + a1: torch.Tensor,
231 + a1_scale: Optional[torch.Tensor],
232 + a2_scale: Optional[torch.Tensor],
233 + topk_weights: torch.Tensor,
234 + topk_ids: torch.Tensor,
235 + num_experts: int,
236 + expert_map: Optional[torch.Tensor],
237 + apply_router_weight_on_input: bool,
238 + quant_config: FusedMoEQuantConfig,
239 + ) -> mk.PrepareResultType:
240 + receiver = self.prepare_async(a1, a1_scale, a2_scale, topk_weights,
241 + topk_ids, num_experts, expert_map,
242 + apply_router_weight_on_input,
243 + quant_config)
244 + return receiver()
245
246 def finalize(
247 self,
vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py
+41 -9
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from typing import Optional, Union
3 +from typing import Callable, Optional, Union
4
5 import deep_ep
6 import torch
@@ -75,7 +75,6 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
75 self,
76 x: Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]],
77 a1_scale: Optional[torch.Tensor],
78 - a2_scale: Optional[torch.Tensor],
78 a1_dtype: torch.dtype,
79 quant_dtype: Union[torch.dtype, str, None],
80 per_act_token_quant: bool,
@@ -110,7 +109,10 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
109
110 return x, x_scales
111
113 - def prepare(
112 + def supports_async(self) -> bool:
113 + return True
114 +
115 + def prepare_async(
116 self,
117 a1: torch.Tensor,
118 a1_scale: Optional[torch.Tensor],
@@ -121,9 +123,7 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
123 expert_map: Optional[torch.Tensor],
124 apply_router_weight_on_input: bool,
125 quant_config: FusedMoEQuantConfig,
124 - ) -> tuple[torch.Tensor, Optional[torch.Tensor],
125 - Optional[mk.ExpertTokensMetadata], Optional[torch.Tensor],
126 - Optional[torch.Tensor]]:
126 + ) -> mk.ReceiverType:
127
128 hidden_size = a1.size(1)
129 assert hidden_size in self.SUPPORTED_HIDDEN_SIZES, \
@@ -155,16 +155,48 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
155 num_experts,
156 use_fp8=self.use_fp8_dispatch,
157 async_finish=False,
158 - return_recv_hook=False)
158 + return_recv_hook=True)
159 +
160 + return lambda: self._receiver(hook, expert_x, expert_num_tokens,
161 + a1_scale, a1.dtype, quant_config)
162 +
163 + def _receiver(
164 + self,
165 + hook: Callable,
166 + expert_x: Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]],
167 + expert_num_tokens: torch.Tensor,
168 + a1_scale,
169 + a1_dtype,
170 + quant_config: FusedMoEQuantConfig,
171 + ) -> mk.PrepareResultType:
172 + hook()
173
174 expert_x, expert_x_scale = self._do_quant(
161 - expert_x, a1_scale, a2_scale, a1.dtype, quant_config.quant_dtype,
175 + expert_x, a1_scale, a1_dtype, quant_config.quant_dtype,
176 quant_config.per_act_token_quant, quant_config.block_shape)
177
178 expert_tokens_meta = mk.ExpertTokensMetadata(
179 expert_num_tokens=expert_num_tokens, expert_num_tokens_cpu=None)
180
167 - return (expert_x, expert_x_scale, expert_tokens_meta, None, None)
181 + return expert_x, expert_x_scale, expert_tokens_meta, None, None
182 +
183 + def prepare(
184 + self,
185 + a1: torch.Tensor,
186 + a1_scale: Optional[torch.Tensor],
187 + a2_scale: Optional[torch.Tensor],
188 + topk_weights: torch.Tensor,
189 + topk_ids: torch.Tensor,
190 + num_experts: int,
191 + expert_map: Optional[torch.Tensor],
192 + apply_router_weight_on_input: bool,
193 + quant_config: FusedMoEQuantConfig,
194 + ) -> mk.PrepareResultType:
195 + receiver = self.prepare_async(a1, a1_scale, a2_scale, topk_weights,
196 + topk_ids, num_experts, expert_map,
197 + apply_router_weight_on_input,
198 + quant_config)
199 + return receiver()
200
201 def finalize(
202 self,
vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py
+1 -3
@@ -56,9 +56,7 @@ class FlashInferCutlassMoEPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
56 apply_router_weight_on_input: bool,
57 # TODO(bnell): use quant_config + scales instead of ctor args
58 quant_config: FusedMoEQuantConfig,
59 - ) -> tuple[torch.Tensor, Optional[torch.Tensor],
60 - Optional[mk.ExpertTokensMetadata], Optional[torch.Tensor],
61 - Optional[torch.Tensor]]:
59 + ) -> mk.PrepareResultType:
60
61 if apply_router_weight_on_input:
62 topk = topk_ids.size(1)
vllm/model_executor/layers/fused_moe/fused_batched_moe.py
+1 -3
@@ -506,9 +506,7 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
506 expert_map: Optional[torch.Tensor],
507 apply_router_weight_on_input: bool,
508 quant_config: FusedMoEQuantConfig,
509 - ) -> tuple[torch.Tensor, Optional[torch.Tensor],
510 - Optional[mk.ExpertTokensMetadata], Optional[torch.Tensor],
511 - Optional[torch.Tensor]]:
509 + ) -> mk.PrepareResultType:
510 assert a1.dim() == 2
511 assert topk_ids.dim() == 2
512 assert topk_ids.size(0) == a1.size(0)
vllm/model_executor/layers/fused_moe/layer.py
+146 -37
@@ -4,7 +4,7 @@
4 from abc import abstractmethod
5 from collections.abc import Iterable
6 from enum import Enum
7 -from typing import Callable, Literal, Optional, overload
7 +from typing import Callable, Literal, Optional, Union, overload
8
9 import torch
10 import torch.nn.functional as F
@@ -215,6 +215,7 @@ class FusedMoEMethodBase(QuantizeMethodBase):
215 self.fused_experts = FusedMoEModularKernel(
216 prepare_finalize,
217 experts,
218 + layer.shared_experts,
219 )
220
221 def select_gemm_impl(
@@ -252,7 +253,7 @@ class FusedMoEMethodBase(QuantizeMethodBase):
253 expert_load_view: Optional[torch.Tensor] = None,
254 logical_to_physical_map: Optional[torch.Tensor] = None,
255 logical_replica_count: Optional[torch.Tensor] = None,
255 - ) -> torch.Tensor:
256 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
257 raise NotImplementedError
258
259
@@ -409,7 +410,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
410 expert_load_view: Optional[torch.Tensor] = None,
411 logical_to_physical_map: Optional[torch.Tensor] = None,
412 logical_replica_count: Optional[torch.Tensor] = None,
412 - ) -> torch.Tensor:
413 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
414 if enable_eplb:
415 assert expert_load_view is not None
416 assert logical_to_physical_map is not None
@@ -461,7 +462,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
462 expert_load_view: Optional[torch.Tensor] = None,
463 logical_to_physical_map: Optional[torch.Tensor] = None,
464 logical_replica_count: Optional[torch.Tensor] = None,
464 - ) -> torch.Tensor:
465 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
466
467 topk_weights, topk_ids = FusedMoE.select_experts(
468 hidden_states=x,
@@ -547,7 +548,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
548 expert_load_view: Optional[torch.Tensor] = None,
549 logical_to_physical_map: Optional[torch.Tensor] = None,
550 logical_replica_count: Optional[torch.Tensor] = None,
550 - ):
551 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
552 if enable_eplb is not False or expert_load_view is not None or \
553 logical_to_physical_map is not None or \
554 logical_replica_count is not None:
@@ -594,7 +595,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
595 expert_load_view: Optional[torch.Tensor] = None,
596 logical_to_physical_map: Optional[torch.Tensor] = None,
597 logical_replica_count: Optional[torch.Tensor] = None,
597 - ):
598 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
599 if enable_eplb is not False or expert_load_view is not None or \
600 logical_to_physical_map is not None or \
601 logical_replica_count is not None:
@@ -633,7 +634,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
634 expert_load_view: Optional[torch.Tensor] = None,
635 logical_to_physical_map: Optional[torch.Tensor] = None,
636 logical_replica_count: Optional[torch.Tensor] = None,
636 - ) -> torch.Tensor:
637 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
638 assert not use_grouped_topk
639 assert num_expert_group is None
640 assert topk_group is None
@@ -948,6 +949,10 @@ class FusedMoE(CustomOp):
949 dtype=moe.in_dtype,
950 device=torch.cuda.current_device())
951
952 + @property
953 + def shared_experts(self) -> Optional[torch.nn.Module]:
954 + return None
955 +
956 @property
957 def tp_size(self):
958 return self.moe_parallel_config.tp_size
@@ -1400,6 +1405,7 @@ class FusedMoE(CustomOp):
1405 return [
1406 weight.view(self.local_num_experts, -1) for name, weight in weights
1407 if name not in NON_EXPERT_WEIGHTS
1408 + and not name.startswith("_shared_experts.")
1409 ]
1410
1411 def set_eplb_state(
@@ -1582,25 +1588,45 @@ class FusedMoE(CustomOp):
1588 else:
1589 return tensor_model_parallel_all_reduce(final_hidden_states)
1590
1585 - def forward(self, hidden_states: torch.Tensor,
1586 - router_logits: torch.Tensor):
1591 + def forward(
1592 + self,
1593 + hidden_states: torch.Tensor,
1594 + router_logits: torch.Tensor,
1595 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
1596 og_hidden_states = hidden_states.shape[-1]
1597 if self.hidden_size != og_hidden_states:
1598 hidden_states = F.pad(hidden_states,
1599 (0, self.hidden_size - og_hidden_states),
1600 mode='constant',
1601 value=0.0)
1593 - # TODO: Once the OOM issue for the TPU backend is resolved, we will
1594 - # switch to using the moe_forward custom op.
1595 - if current_platform.is_tpu():
1596 - return self.forward_impl(hidden_states, router_logits)
1602 +
1603 + if self.shared_experts is None:
1604 + if current_platform.is_tpu():
1605 + # TODO: Once the OOM issue for the TPU backend is resolved, we
1606 + # will switch to using the moe_forward custom op.
1607 + fused_output = self.forward_impl(hidden_states, router_logits)
1608 + assert not isinstance(fused_output, tuple)
1609 + else:
1610 + fused_output = torch.ops.vllm.moe_forward(
1611 + hidden_states, router_logits, self.layer_name)
1612 + return fused_output[..., :og_hidden_states]
1613 else:
1598 - return torch.ops.vllm.moe_forward(
1599 - hidden_states, router_logits,
1600 - self.layer_name)[..., :og_hidden_states]
1614 + if current_platform.is_tpu():
1615 + # TODO: Once the OOM issue for the TPU backend is resolved, we
1616 + # will switch to using the moe_forward custom op.
1617 + shared_output, fused_output = self.forward_impl(
1618 + hidden_states, router_logits)
1619 + else:
1620 + shared_output, fused_output = torch.ops.vllm.moe_forward_shared(
1621 + hidden_states, router_logits, self.layer_name)
1622 + return (shared_output[..., :og_hidden_states],
1623 + fused_output[..., :og_hidden_states])
1624
1602 - def forward_impl_chunked(self, full_hidden_states: torch.Tensor,
1603 - full_router_logits: torch.Tensor):
1625 + def forward_impl_chunked(
1626 + self,
1627 + full_hidden_states: torch.Tensor,
1628 + full_router_logits: torch.Tensor,
1629 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
1630 assert self.batched_hidden_states is not None
1631 assert self.batched_router_logits is not None
1632 assert self.batched_hidden_states.dtype == full_hidden_states.dtype
@@ -1611,7 +1637,10 @@ class FusedMoE(CustomOp):
1637 assert (
1638 self.batched_router_logits.size(-1) == full_router_logits.size(-1))
1639
1614 - full_final_hidden_states = torch.empty_like(full_hidden_states)
1640 + full_fused_final_hidden_states = torch.empty_like(full_hidden_states)
1641 + if self.shared_experts is not None:
1642 + full_shared_final_hidden_states = torch.empty_like(
1643 + full_hidden_states)
1644
1645 def process_chunk(chunk_start, chunk_end, skip_result_store=False):
1646 chunk_size = chunk_end - chunk_start
@@ -1652,9 +1681,21 @@ class FusedMoE(CustomOp):
1681 logical_replica_count=self.logical_replica_count,
1682 )
1683
1684 + assert self.shared_experts is None or isinstance(
1685 + final_hidden_states, tuple)
1686 +
1687 if not skip_result_store:
1656 - full_final_hidden_states[chunk_start:chunk_end, :].copy_(
1657 - final_hidden_states, non_blocking=True)
1688 + if self.shared_experts is None:
1689 + full_fused_final_hidden_states[
1690 + chunk_start:chunk_end, :].copy_(final_hidden_states,
1691 + non_blocking=True)
1692 + else:
1693 + full_shared_final_hidden_states[
1694 + chunk_start:chunk_end, :].copy_(final_hidden_states[0],
1695 + non_blocking=True)
1696 + full_fused_final_hidden_states[
1697 + chunk_start:chunk_end, :].copy_(final_hidden_states[1],
1698 + non_blocking=True)
1699
1700 ctx = get_forward_context()
1701 # flashinfer_cutlass_kernels can handle: optional DP + TP/EP
@@ -1675,10 +1716,17 @@ class FusedMoE(CustomOp):
1716 chunk_end,
1717 skip_result_store=chunk_start_ >= num_tokens)
1718
1678 - return full_final_hidden_states
1719 + if self.shared_experts is None:
1720 + return full_fused_final_hidden_states
1721 + else:
1722 + return (full_shared_final_hidden_states,
1723 + full_fused_final_hidden_states)
1724
1680 - def forward_impl(self, hidden_states: torch.Tensor,
1681 - router_logits: torch.Tensor):
1725 + def forward_impl(
1726 + self,
1727 + hidden_states: torch.Tensor,
1728 + router_logits: torch.Tensor,
1729 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
1730 assert self.quant_method is not None
1731 # Route to the chunked forward path using the FlashInfer Cutlass kernel
1732 # only when data parallelism (DP) is enabled.
@@ -1698,6 +1746,15 @@ class FusedMoE(CustomOp):
1746 hidden_states, router_logits = get_ep_group().dispatch(
1747 hidden_states, router_logits)
1748
1749 + # If there are shared experts but we are not using a modular kernel, the
1750 + # shared experts must be called here
1751 + if (not isinstance(self.quant_method.fused_experts,
1752 + FusedMoEModularKernel)
1753 + and self.shared_experts is not None):
1754 + shared_output = self.shared_experts(hidden_states)
1755 + else:
1756 + shared_output = None
1757 +
1758 # Matrix multiply.
1759 final_hidden_states = self.quant_method.apply(
1760 layer=self,
@@ -1722,14 +1779,30 @@ class FusedMoE(CustomOp):
1779 logical_replica_count=self.logical_replica_count,
1780 )
1781
1725 - if do_naive_dispatch_combine:
1726 - final_hidden_states = get_ep_group().combine(final_hidden_states)
1727 - if self.reduce_results and (self.tp_size > 1 or self.ep_size > 1):
1728 - # Default set to False. (May have to add shared expert outputs.
1729 - final_hidden_states = self.maybe_all_reduce_tensor_model_parallel(
1730 - final_hidden_states)
1782 + if shared_output is not None:
1783 + assert not isinstance(final_hidden_states, tuple)
1784 + assert self.shared_experts is not None
1785 + final_hidden_states = (
1786 + shared_output,
1787 + final_hidden_states,
1788 + )
1789
1732 - return final_hidden_states
1790 + def reduce_output(states: torch.Tensor) -> torch.Tensor:
1791 + if do_naive_dispatch_combine:
1792 + states = get_ep_group().combine(states)
1793 +
1794 + if self.reduce_results and (self.tp_size > 1 or self.ep_size > 1):
1795 + states = self.maybe_all_reduce_tensor_model_parallel(states)
1796 +
1797 + return states
1798 +
1799 + if self.shared_experts is None:
1800 + return reduce_output(final_hidden_states)
1801 + else:
1802 + return (
1803 + reduce_output(final_hidden_states[0]),
1804 + reduce_output(final_hidden_states[1]),
1805 + )
1806
1807 @classmethod
1808 def make_expert_params_mapping(
@@ -1784,17 +1857,22 @@ class FusedMoE(CustomOp):
1857 return s
1858
1859
1787 -def moe_forward(hidden_states: torch.Tensor, router_logits: torch.Tensor,
1788 - layer_name: str) -> torch.Tensor:
1860 +def moe_forward(
1861 + hidden_states: torch.Tensor,
1862 + router_logits: torch.Tensor,
1863 + layer_name: str,
1864 +) -> torch.Tensor:
1865 forward_context: ForwardContext = get_forward_context()
1866 self = forward_context.no_compile_layers[layer_name]
1791 - assert self.quant_method is not None
1792 -
1867 + assert self.shared_experts is None
1868 return self.forward_impl(hidden_states, router_logits)
1869
1870
1796 -def moe_forward_fake(hidden_states: torch.Tensor, router_logits: torch.Tensor,
1797 - layer_name: str) -> torch.Tensor:
1871 +def moe_forward_fake(
1872 + hidden_states: torch.Tensor,
1873 + router_logits: torch.Tensor,
1874 + layer_name: str,
1875 +) -> torch.Tensor:
1876 return torch.empty_like(hidden_states)
1877
1878
@@ -1807,6 +1885,37 @@ direct_register_custom_op(
1885 tags=(torch.Tag.needs_fixed_stride_order, ),
1886 )
1887
1888 +
1889 +def moe_forward_shared(
1890 + hidden_states: torch.Tensor,
1891 + router_logits: torch.Tensor,
1892 + layer_name: str,
1893 +) -> tuple[torch.Tensor, torch.Tensor]:
1894 + forward_context: ForwardContext = get_forward_context()
1895 + self = forward_context.no_compile_layers[layer_name]
1896 + assert self.shared_experts is not None
1897 + return self.forward_impl(hidden_states, router_logits)
1898 +
1899 +
1900 +def moe_forward_shared_fake(
1901 + hidden_states: torch.Tensor,
1902 + router_logits: torch.Tensor,
1903 + layer_name: str,
1904 +) -> tuple[torch.Tensor, torch.Tensor]:
1905 + shared_out = torch.empty_like(hidden_states)
1906 + fused_out = torch.empty_like(hidden_states)
1907 + return shared_out, fused_out
1908 +
1909 +
1910 +direct_register_custom_op(
1911 + op_name="moe_forward_shared",
1912 + op_func=moe_forward_shared,
1913 + mutates_args=["hidden_states"],
1914 + fake_impl=moe_forward_shared_fake,
1915 + dispatch_key=current_platform.dispatch_key,
1916 + tags=(torch.Tag.needs_fixed_stride_order, ),
1917 +)
1918 +
1919 # Mark the FusedMoE weight_loader as supporting MoE-specific parameters
1920 # to avoid expensive runtime reflection in model loading code
1921 FusedMoE.weight_loader.supports_moe_loading = True # type: ignore[attr-defined]
vllm/model_executor/layers/fused_moe/modular_kernel.py
+118 -24
@@ -4,7 +4,7 @@ from abc import ABC, abstractmethod
4 from dataclasses import dataclass
5 from enum import Enum
6 from math import prod
7 -from typing import Optional, final
7 +from typing import Callable, Optional, Union, final
8
9 import torch
10
@@ -141,6 +141,29 @@ class TopKWeightAndReduce(ABC):
141 raise NotImplementedError
142
143
144 +#
145 +# PrepareResultType is a tuple of:
146 +# - quantized + dispatched a.
147 +# - quantized + dispatched a1_scales.
148 +# - Optional ExpertTokensMetadata containing gpu/cpu tensors
149 +# as big as the number of local experts with the information about the
150 +# number of tokens assigned to each local expert.
151 +# - Optional dispatched expert topk IDs
152 +# - Optional dispatched expert topk weight
153 +#
154 +# See `prepare` method below.
155 +#
156 +PrepareResultType = tuple[
157 + torch.Tensor,
158 + Optional[torch.Tensor],
159 + Optional[ExpertTokensMetadata],
160 + Optional[torch.Tensor],
161 + Optional[torch.Tensor],
162 +]
163 +
164 +ReceiverType = Callable[[], PrepareResultType]
165 +
166 +
167 # TODO: pass FusedMoEParallelConfig in as ctor parameter?
168 class FusedMoEPrepareAndFinalize(ABC):
169 """
@@ -160,16 +183,9 @@ class FusedMoEPrepareAndFinalize(ABC):
183 expert_map: Optional[torch.Tensor],
184 apply_router_weight_on_input: bool,
185 quant_config: FusedMoEQuantConfig,
163 - ) -> tuple[
164 - torch.Tensor,
165 - Optional[torch.Tensor],
166 - Optional[ExpertTokensMetadata],
167 - Optional[torch.Tensor],
168 - Optional[torch.Tensor],
169 - ]:
186 + ) -> PrepareResultType:
187 """
171 - Perform any quantization (and/or) dispatching needed
172 - for this kernel.
188 + Perform any quantization (and/or) dispatching needed for this kernel.
189 - a1: The (unquantized) input to the MoE layer.
190 - a1_scale: Optional scales for a1
191 - a2_scale: Optional scales for the second MoE gemm. Required to make
@@ -193,6 +209,51 @@ class FusedMoEPrepareAndFinalize(ABC):
209 """
210 raise NotImplementedError
211
212 + def supports_async(self) -> bool:
213 + """
214 + Indicates whether or not this class implements prepare_async.
215 + """
216 + return False
217 +
218 + def prepare_async(
219 + self,
220 + a1: torch.Tensor,
221 + a1_scale: Optional[torch.Tensor],
222 + a2_scale: Optional[torch.Tensor],
223 + topk_weights: torch.Tensor,
224 + topk_ids: torch.Tensor,
225 + num_experts: int,
226 + expert_map: Optional[torch.Tensor],
227 + apply_router_weight_on_input: bool,
228 + quant_config: FusedMoEQuantConfig,
229 + ) -> ReceiverType:
230 + """
231 + Perform any quantization (and/or) dispatching needed for this kernel
232 + but do not wait for results from other workers.
233 + - a1: The (unquantized) input to the MoE layer.
234 + - a1_scale: Optional scales for a1
235 + - a2_scale: Optional scales for the second MoE gemm. Required to make
236 + sure the quantization is consistent for both gemms.
237 + - topk_ids: The topk ids.
238 + - topk_weights: The topk weights.
239 + - num_experts: The total number of experts in the global expert space.
240 + - expert_map: A tensor mapping expert indices from the global expert
241 + space to the local expert space of the expert parallel shard.
242 + - apply_router_weight_on_input: When True, apply the weights to the
243 + activations, before quantization + dispatching.
244 +
245 + Returns a callback that when invoked waits for results from other
246 + workers and has the same return signature as `prepare`, e.g.
247 +
248 + receiver = obj.prepare_async(...)
249 + a, a_scales, expert_meta, topk_ids, topk_weights = receiver()
250 +
251 + is equivalent to:
252 +
253 + a, a_scales, expert_meta, topk_ids, topk_weights = obj.prepare(...)
254 + """
255 + raise NotImplementedError
256 +
257 @abstractmethod
258 def finalize(
259 self,
@@ -453,10 +514,12 @@ class FusedMoEModularKernel(torch.nn.Module):
514 self,
515 prepare_finalize: FusedMoEPrepareAndFinalize,
516 fused_experts: FusedMoEPermuteExpertsUnpermute,
517 + shared_experts: Optional[torch.nn.Module] = None,
518 ):
519 super().__init__()
520 self.prepare_finalize = prepare_finalize
521 self.fused_experts = fused_experts
522 + self.shared_experts = shared_experts
523 assert prepare_finalize.activation_format == \
524 fused_experts.activation_formats[0], (
525 f"{prepare_finalize.__class__.__name__}."
@@ -692,7 +755,7 @@ class FusedMoEModularKernel(torch.nn.Module):
755 a1_scale: Optional[torch.Tensor] = None,
756 a2_scale: Optional[torch.Tensor] = None,
757 apply_router_weight_on_input: bool = False,
695 - ) -> torch.Tensor:
758 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
759 """
760 This function computes a Mixture of Experts (MoE) layer using two sets
761 of weights, w1 and w2, and top-k gating mechanism.
@@ -736,18 +799,46 @@ class FusedMoEModularKernel(torch.nn.Module):
799 if global_num_experts == -1:
800 global_num_experts = local_num_experts
801
739 - (a1q, a1q_scale, expert_tokens_meta, _expert_topk_ids,
740 - _expert_topk_weights) = self.prepare_finalize.prepare(
741 - a1,
742 - a1_scale,
743 - a2_scale,
744 - topk_weights,
745 - topk_ids,
746 - global_num_experts,
747 - expert_map,
748 - apply_router_weight_on_input,
749 - self.fused_experts.quant_config,
750 - )
802 + shared_output: torch.Tensor
803 +
804 + if (not self.prepare_finalize.supports_async()
805 + or self.shared_experts is None):
806 +
807 + # Run shared experts serially with dispatch.
808 + if self.shared_experts is not None:
809 + shared_output = self.shared_experts(a1)
810 +
811 + (a1q, a1q_scale, expert_tokens_meta, _expert_topk_ids,
812 + _expert_topk_weights) = self.prepare_finalize.prepare(
813 + a1,
814 + a1_scale,
815 + a2_scale,
816 + topk_weights,
817 + topk_ids,
818 + global_num_experts,
819 + expert_map,
820 + apply_router_weight_on_input,
821 + self.fused_experts.quant_config,
822 + )
823 + else:
824 + # Overlap shared expert compute with all2all dispatch.
825 + receiver = self.prepare_finalize.prepare_async(
826 + a1,
827 + a1_scale,
828 + a2_scale,
829 + topk_weights,
830 + topk_ids,
831 + global_num_experts,
832 + expert_map,
833 + apply_router_weight_on_input,
834 + self.fused_experts.quant_config,
835 + )
836 +
837 + assert self.shared_experts is not None
838 + shared_output = self.shared_experts(a1)
839 +
840 + (a1q, a1q_scale, expert_tokens_meta, _expert_topk_ids,
841 + _expert_topk_weights) = receiver()
842
843 # Maybe prepare gathered topk_ids and topk_weights from other EP ranks.
844 topk_ids = topk_ids if _expert_topk_ids is None else _expert_topk_ids
@@ -795,4 +886,7 @@ class FusedMoEModularKernel(torch.nn.Module):
886 self.fused_experts.finalize_weight_and_reduce_impl(),
887 )
888
798 - return output
889 + if self.shared_experts is None:
890 + return output
891 + else:
892 + return shared_output, output
vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py
+71 -6
@@ -84,12 +84,15 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
84 return self.max_num_tokens
85
86 def topk_indices_dtype(self) -> Optional[torch.dtype]:
87 - return torch.int32
87 + return torch.uint32
88
89 def num_dispatchers(self) -> int:
90 return self.num_dispatchers_
91
92 - def prepare(
92 + def supports_async(self) -> bool:
93 + return True
94 +
95 + def prepare_async(
96 self,
97 a1: torch.Tensor,
98 a1_scale: Optional[torch.Tensor],
@@ -100,9 +103,7 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
103 expert_map: Optional[torch.Tensor],
104 apply_router_weight_on_input: bool,
105 quant_config: FusedMoEQuantConfig,
103 - ) -> tuple[torch.Tensor, Optional[torch.Tensor],
104 - Optional[mk.ExpertTokensMetadata], Optional[torch.Tensor],
105 - Optional[torch.Tensor]]:
106 + ) -> mk.ReceiverType:
107 num_tokens = a1.size(0) # M
108 hidden_dim = a1.size(-1) # K
109
@@ -138,6 +139,8 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
139 _validate_scale_shape(a1q, a1q_scale, quant_config.per_act_token_quant,
140 quant_config.block_shape)
141
142 + orig_a_scale_block_shape: Optional[int] = None
143 +
144 if a1q_scale is not None:
145 scalar_scales = a1q_scale.numel() == 1
146
@@ -205,8 +208,45 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
208 out_expert_x_scale=expert_x_scale,
209 dp_x=a1q,
210 dp_x_scale=a1q_scale,
208 - indices=topk_ids.view(dtype=torch.uint32),
211 + indices=topk_ids,
212 + bound_m=bound_m,
213 + do_send=True,
214 + do_recv=False,
215 + )
216 +
217 + return lambda: self._receiver(
218 + expert_num_tokens,
219 + expert_x,
220 + expert_x_scale,
221 + a1q,
222 + a1q_scale,
223 + topk_ids,
224 + bound_m,
225 + orig_a_scale_block_shape,
226 + )
227 +
228 + def _receiver(
229 + self,
230 + expert_num_tokens: torch.Tensor,
231 + expert_x: torch.Tensor,
232 + expert_x_scale: Optional[torch.Tensor],
233 + a1q: torch.Tensor,
234 + a1q_scale: Optional[torch.Tensor],
235 + topk_ids: torch.Tensor,
236 + bound_m: Optional[torch.Tensor],
237 + orig_a_scale_block_shape: Optional[int],
238 + ) -> mk.PrepareResultType:
239 +
240 + self.a2a.dispatch(
241 + out_expert_num_tokens=expert_num_tokens,
242 + out_expert_x=expert_x,
243 + out_expert_x_scale=expert_x_scale,
244 + dp_x=a1q,
245 + dp_x_scale=a1q_scale,
246 + indices=topk_ids,
247 bound_m=bound_m,
248 + do_send=False,
249 + do_recv=True,
250 )
251
252 if expert_x_scale is not None:
@@ -218,6 +258,31 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
258
259 return expert_x, expert_x_scale, expert_tokens_meta, None, None
260
261 + def prepare(
262 + self,
263 + a1: torch.Tensor,
264 + a1_scale: Optional[torch.Tensor],
265 + a2_scale: Optional[torch.Tensor],
266 + topk_weights: torch.Tensor,
267 + topk_ids: torch.Tensor,
268 + num_experts: int,
269 + expert_map: Optional[torch.Tensor],
270 + apply_router_weight_on_input: bool,
271 + quant_config: FusedMoEQuantConfig,
272 + ) -> mk.PrepareResultType:
273 + receiver = self.prepare_async(
274 + a1,
275 + a1_scale,
276 + a2_scale,
277 + topk_weights,
278 + topk_ids,
279 + num_experts,
280 + expert_map,
281 + apply_router_weight_on_input,
282 + quant_config,
283 + )
284 + return receiver()
285 +
286 def finalize(
287 self,
288 output: torch.Tensor,
vllm/model_executor/layers/fused_moe/prepare_finalize.py
+1 -3
@@ -38,9 +38,7 @@ class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):
38 expert_map: Optional[torch.Tensor],
39 apply_router_weight_on_input: bool,
40 quant_config: FusedMoEQuantConfig,
41 - ) -> tuple[torch.Tensor, Optional[torch.Tensor],
42 - Optional[mk.ExpertTokensMetadata], Optional[torch.Tensor],
43 - Optional[torch.Tensor]]:
41 + ) -> mk.PrepareResultType:
42
43 if apply_router_weight_on_input:
44 topk = topk_ids.size(1)
vllm/model_executor/layers/quantization/awq_marlin.py
+2 -2
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from typing import Any, Callable, Optional
4 +from typing import Any, Callable, Optional, Union
5
6 import torch
7 from torch.nn import Parameter
@@ -505,7 +505,7 @@ class AWQMoEMethod(FusedMoEMethodBase):
505 expert_load_view: Optional[torch.Tensor] = None,
506 logical_to_physical_map: Optional[torch.Tensor] = None,
507 logical_replica_count: Optional[torch.Tensor] = None,
508 - ) -> torch.Tensor:
508 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
509 assert self.fused_experts is None
510
511 if enable_eplb:
vllm/model_executor/layers/quantization/bitsandbytes.py
+1 -1
@@ -474,7 +474,7 @@ class BitsAndBytesMoEMethod(FusedMoEMethodBase):
474 expert_load_view: Optional[torch.Tensor] = None,
475 logical_to_physical_map: Optional[torch.Tensor] = None,
476 logical_replica_count: Optional[torch.Tensor] = None,
477 - ) -> torch.Tensor:
477 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
478 from vllm.model_executor.layers.fused_moe import fused_experts
479 assert self.fused_experts is None
480
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+6 -6
@@ -3,7 +3,7 @@
3
4 import enum
5 from enum import Enum
6 -from typing import Callable, Optional
6 +from typing import Callable, Optional, Union
7
8 import torch
9 from compressed_tensors import CompressionFormat
@@ -358,7 +358,7 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
358 expert_load_view: Optional[torch.Tensor] = None,
359 logical_to_physical_map: Optional[torch.Tensor] = None,
360 logical_replica_count: Optional[torch.Tensor] = None,
361 - ) -> torch.Tensor:
361 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
362 assert self.fused_experts is None
363
364 if enable_eplb:
@@ -819,7 +819,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
819 expert_load_view: Optional[torch.Tensor] = None,
820 logical_to_physical_map: Optional[torch.Tensor] = None,
821 logical_replica_count: Optional[torch.Tensor] = None,
822 - ) -> torch.Tensor:
822 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
823 if enable_eplb:
824 raise NotImplementedError(
825 "EPLB not supported for "
@@ -1069,7 +1069,7 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
1069 expert_load_view: Optional[torch.Tensor] = None,
1070 logical_to_physical_map: Optional[torch.Tensor] = None,
1071 logical_replica_count: Optional[torch.Tensor] = None,
1072 - ) -> torch.Tensor:
1072 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
1073 assert self.fused_experts is None
1074
1075 if enable_eplb:
@@ -1375,7 +1375,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1375 expert_load_view: Optional[torch.Tensor] = None,
1376 logical_to_physical_map: Optional[torch.Tensor] = None,
1377 logical_replica_count: Optional[torch.Tensor] = None,
1378 - ) -> torch.Tensor:
1378 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
1379 assert self.fused_experts is None
1380
1381 if enable_eplb:
@@ -1608,7 +1608,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
1608 expert_load_view: Optional[torch.Tensor] = None,
1609 logical_to_physical_map: Optional[torch.Tensor] = None,
1610 logical_replica_count: Optional[torch.Tensor] = None,
1611 - ) -> torch.Tensor:
1611 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
1612 assert self.fused_experts is None
1613
1614 if enable_eplb:
vllm/model_executor/layers/quantization/experts_int8.py
+2 -2
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from typing import Any, Callable, Optional
4 +from typing import Any, Callable, Optional, Union
5
6 import torch
7
@@ -128,7 +128,7 @@ class ExpertsInt8MoEMethod(FusedMoEMethodBase):
128 expert_load_view: Optional[torch.Tensor] = None,
129 logical_to_physical_map: Optional[torch.Tensor] = None,
130 logical_replica_count: Optional[torch.Tensor] = None,
131 - ) -> torch.Tensor:
131 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
132 assert self.fused_experts is None
133
134 if enable_eplb:
vllm/model_executor/layers/quantization/fp8.py
+2 -2
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from typing import TYPE_CHECKING, Any, Callable, Optional
4 +from typing import TYPE_CHECKING, Any, Callable, Optional, Union
5
6 import torch
7 import torch.nn.functional as F
@@ -988,7 +988,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
988 expert_load_view: Optional[torch.Tensor] = None,
989 logical_to_physical_map: Optional[torch.Tensor] = None,
990 logical_replica_count: Optional[torch.Tensor] = None,
991 - ) -> torch.Tensor:
991 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
992 if enable_eplb:
993 assert expert_load_view is not None
994 assert logical_to_physical_map is not None
vllm/model_executor/layers/quantization/gguf.py
+2 -2
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from typing import Any, Callable, Optional
4 +from typing import Any, Callable, Optional, Union
5
6 import gguf
7 import torch
@@ -540,7 +540,7 @@ class GGUFMoEMethod(FusedMoEMethodBase):
540 expert_load_view: Optional[torch.Tensor] = None,
541 logical_to_physical_map: Optional[torch.Tensor] = None,
542 logical_replica_count: Optional[torch.Tensor] = None,
543 - ):
543 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
544 assert self.fused_experts is None
545
546 if enable_eplb:
vllm/model_executor/layers/quantization/gptq_marlin.py
+1 -1
@@ -654,7 +654,7 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
654 expert_load_view: Optional[torch.Tensor] = None,
655 logical_to_physical_map: Optional[torch.Tensor] = None,
656 logical_replica_count: Optional[torch.Tensor] = None,
657 - ) -> torch.Tensor:
657 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
658 assert self.fused_experts is None
659
660 if enable_eplb:
vllm/model_executor/layers/quantization/modelopt.py
+2 -2
@@ -491,7 +491,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
491 expert_load_view: Optional[torch.Tensor] = None,
492 logical_to_physical_map: Optional[torch.Tensor] = None,
493 logical_replica_count: Optional[torch.Tensor] = None,
494 - ) -> torch.Tensor:
494 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
495 if enable_eplb:
496 raise NotImplementedError(
497 "EPLB not supported for `ModelOptFp8MoEMethod` yet.")
@@ -1366,7 +1366,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1366 expert_load_view: Optional[torch.Tensor] = None,
1367 logical_to_physical_map: Optional[torch.Tensor] = None,
1368 logical_replica_count: Optional[torch.Tensor] = None,
1369 - ):
1369 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
1370 if enable_eplb:
1371 raise NotImplementedError(
1372 "EPLB not supported for `ModelOptNvFp4FusedMoE` yet.")
vllm/model_executor/layers/quantization/moe_wna16.py
+2 -2
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from typing import Any, Callable, Optional
4 +from typing import Any, Callable, Optional, Union
5
6 import torch
7
@@ -305,7 +305,7 @@ class MoeWNA16Method(FusedMoEMethodBase):
305 expert_load_view: Optional[torch.Tensor] = None,
306 logical_to_physical_map: Optional[torch.Tensor] = None,
307 logical_replica_count: Optional[torch.Tensor] = None,
308 - ) -> torch.Tensor:
308 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
309 assert self.fused_experts is None
310 if enable_eplb:
311 raise NotImplementedError(
vllm/model_executor/layers/quantization/mxfp4.py
+2 -2
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from typing import Callable, Optional
3 +from typing import Callable, Optional, Union
4
5 import torch
6 from torch.nn.parameter import Parameter
@@ -554,7 +554,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
554 expert_load_view: Optional[torch.Tensor] = None,
555 logical_to_physical_map: Optional[torch.Tensor] = None,
556 logical_replica_count: Optional[torch.Tensor] = None,
557 - ) -> torch.Tensor:
557 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
558
559 if enable_eplb:
560 raise NotImplementedError("EPLB is not supported for mxfp4")
vllm/model_executor/layers/quantization/quark/quark_moe.py
+3 -3
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from typing import Any, Callable, Optional
4 +from typing import Any, Callable, Optional, Union
5
6 import torch
7
@@ -226,7 +226,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
226 expert_load_view: Optional[torch.Tensor] = None,
227 logical_to_physical_map: Optional[torch.Tensor] = None,
228 logical_replica_count: Optional[torch.Tensor] = None,
229 - ) -> torch.Tensor:
229 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
230 assert self.fused_experts is None
231
232 if enable_eplb:
@@ -390,7 +390,7 @@ class QuarkW4A4MXFp4MoEMethod(QuarkMoEMethod):
390 expert_load_view: Optional[torch.Tensor] = None,
391 logical_to_physical_map: Optional[torch.Tensor] = None,
392 logical_replica_count: Optional[torch.Tensor] = None,
393 - ) -> torch.Tensor:
393 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
394 assert self.fused_experts is None
395
396 if enable_eplb:
vllm/model_executor/layers/quantization/rtn.py
+2 -2
@@ -3,7 +3,7 @@
3 # Copyright © 2025, Oracle and/or its affiliates.
4
5 import os
6 -from typing import Any, Callable, Optional
6 +from typing import Any, Callable, Optional, Union
7
8 import torch
9 import torch.nn.functional as F
@@ -291,7 +291,7 @@ class RTNMoEMethod(FusedMoEMethodBase):
291 expert_load_view: Optional[torch.Tensor] = None,
292 logical_to_physical_map: Optional[torch.Tensor] = None,
293 logical_replica_count: Optional[torch.Tensor] = None,
294 - ) -> torch.Tensor:
294 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
295 assert self.fused_experts is None
296
297 if enable_eplb:
vllm/model_executor/layers/shared_fused_moe/__init__.py new
+6
@@ -0,0 +1,6 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from vllm.model_executor.layers.shared_fused_moe.shared_fused_moe import (
4 + SharedFusedMoE)
5 +
6 +__all__ = ["SharedFusedMoE"]
vllm/model_executor/layers/shared_fused_moe/shared_fused_moe.py new
+56
@@ -0,0 +1,56 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import Optional
4 +
5 +import torch
6 +
7 +from vllm.distributed import tensor_model_parallel_all_reduce
8 +from vllm.model_executor.layers.fused_moe.layer import FusedMoE
9 +
10 +
11 +# TODO(bnell): Add shared + fused combo function? e.g. +
12 +class SharedFusedMoE(FusedMoE):
13 + """
14 + A FusedMoE operation that also computes the results of shared experts.
15 + If an all2all communicator is being used the shared expert computation
16 + can be interleaved with the fused all2all dispatch communication step.
17 + """
18 +
19 + def __init__(
20 + self,
21 + shared_experts: torch.nn.Module,
22 + use_overlapped: bool = True,
23 + **kwargs,
24 + ):
25 + super().__init__(**kwargs)
26 + self._shared_experts = shared_experts
27 + self.use_overlapped = use_overlapped
28 +
29 + @property
30 + def shared_experts(self) -> Optional[torch.nn.Module]:
31 + return self._shared_experts if self.use_overlapped else None
32 +
33 + def forward(
34 + self,
35 + hidden_states: torch.Tensor,
36 + router_logits: torch.Tensor,
37 + ) -> tuple[torch.Tensor, torch.Tensor]:
38 + if not self.use_overlapped:
39 + shared_out = self._shared_experts(hidden_states)
40 +
41 + # Reduce outputs if necessary, since the MLP should
42 + # have been created with reduce_results=False.
43 + if (self.reduce_results and self.tp_size > 1
44 + and self.must_reduce_shared_expert_outputs()):
45 + shared_out = tensor_model_parallel_all_reduce(shared_out)
46 +
47 + fused_out = super().forward(
48 + hidden_states=hidden_states,
49 + router_logits=router_logits,
50 + )
51 + else:
52 + shared_out, fused_out = super().forward(
53 + hidden_states=hidden_states,
54 + router_logits=router_logits,
55 + )
56 + return shared_out, fused_out
vllm/model_executor/models/deepseek_v2.py
+63 -40
@@ -49,6 +49,7 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,
49 from vllm.model_executor.layers.logits_processor import LogitsProcessor
50 from vllm.model_executor.layers.quantization import QuantizationConfig
51 from vllm.model_executor.layers.rotary_embedding import get_rope
52 +from vllm.model_executor.layers.shared_fused_moe import SharedFusedMoE
53 from vllm.model_executor.layers.vocab_parallel_embedding import (
54 ParallelLMHead, VocabParallelEmbedding)
55 from vllm.model_executor.model_loader.weight_utils import (
@@ -147,63 +148,85 @@ class DeepseekV2MoE(nn.Module):
148 self.physical_expert_end = (self.physical_expert_start +
149 self.n_local_physical_experts)
150
150 - self.experts = FusedMoE(
151 - num_experts=config.n_routed_experts,
152 - top_k=config.num_experts_per_tok,
153 - hidden_size=config.hidden_size,
154 - intermediate_size=config.moe_intermediate_size,
155 - reduce_results=False,
156 - renormalize=config.norm_topk_prob,
157 - quant_config=quant_config,
158 - use_grouped_topk=True,
159 - num_expert_group=config.n_group,
160 - topk_group=config.topk_group,
161 - prefix=f"{prefix}.experts",
162 - scoring_func=config.scoring_func,
163 - # we do scaling outside, set factor to 1.0 to avoid double mul
164 - routed_scaling_factor=1.0,
165 - e_score_correction_bias=self.gate.e_score_correction_bias,
166 - enable_eplb=self.enable_eplb,
167 - num_redundant_experts=self.n_redundant_experts)
168 -
169 - if config.n_shared_experts is not None:
151 + if config.n_shared_experts is None:
152 + self.experts = FusedMoE(
153 + num_experts=config.n_routed_experts,
154 + top_k=config.num_experts_per_tok,
155 + hidden_size=config.hidden_size,
156 + intermediate_size=config.moe_intermediate_size,
157 + reduce_results=False,
158 + renormalize=config.norm_topk_prob,
159 + quant_config=quant_config,
160 + use_grouped_topk=True,
161 + num_expert_group=config.n_group,
162 + topk_group=config.topk_group,
163 + prefix=f"{prefix}.experts",
164 + scoring_func=config.scoring_func,
165 + # we do scaling outside, set factor to 1.0 to avoid double mul
166 + routed_scaling_factor=1.0,
167 + e_score_correction_bias=self.gate.e_score_correction_bias,
168 + enable_eplb=self.enable_eplb,
169 + num_redundant_experts=self.n_redundant_experts)
170 + self.shared_experts = None
171 + else:
172 intermediate_size = (config.moe_intermediate_size *
173 config.n_shared_experts)
174 +
175 self.shared_experts = DeepseekV2MLP(
176 hidden_size=config.hidden_size,
177 intermediate_size=intermediate_size,
178 hidden_act=config.hidden_act,
179 quant_config=quant_config,
177 - reduce_results=self.experts.must_reduce_shared_expert_outputs(
178 - ),
180 + reduce_results=False,
181 prefix=f"{prefix}.shared_experts",
182 )
183
184 + self.experts = SharedFusedMoE(
185 + shared_experts=self.shared_experts,
186 + num_experts=config.n_routed_experts,
187 + top_k=config.num_experts_per_tok,
188 + hidden_size=config.hidden_size,
189 + intermediate_size=config.moe_intermediate_size,
190 + reduce_results=False,
191 + renormalize=config.norm_topk_prob,
192 + quant_config=quant_config,
193 + use_grouped_topk=True,
194 + num_expert_group=config.n_group,
195 + topk_group=config.topk_group,
196 + prefix=f"{prefix}.experts",
197 + scoring_func=config.scoring_func,
198 + # we do scaling outside, set factor to 1.0 to avoid double mul
199 + routed_scaling_factor=1.0,
200 + e_score_correction_bias=self.gate.e_score_correction_bias,
201 + enable_eplb=self.enable_eplb,
202 + num_redundant_experts=self.n_redundant_experts)
203 +
204 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
205 num_tokens, hidden_dim = hidden_states.shape
206 hidden_states = hidden_states.view(-1, hidden_dim)
185 - if self.n_shared_experts is not None:
186 - shared_output = self.shared_experts(hidden_states)
207 # router_logits: (num_tokens, n_experts)
208 router_logits, _ = self.gate(hidden_states)
209
190 - if hidden_states.dtype != torch.float16:
191 - final_hidden_states = self.experts(
192 - hidden_states=hidden_states,
193 - router_logits=router_logits) * self.routed_scaling_factor
210 + fused_moe_out = self.experts(hidden_states=hidden_states,
211 + router_logits=router_logits)
212 +
213 + if self.shared_experts is not None:
214 + shared_output, final_hidden_states = fused_moe_out
215 else:
195 - # Fix FP16 overflow
196 - # See DeepseekV2DecoderLayer for more details.
197 - final_hidden_states = self.experts(hidden_states=hidden_states,
198 - router_logits=router_logits)
199 - if shared_output is not None:
200 - if hidden_states.dtype != torch.float16:
201 - final_hidden_states = final_hidden_states + shared_output
202 - else:
203 - # Fix FP16 overflow
204 - # See DeepseekV2DecoderLayer for more details.
205 - final_hidden_states = final_hidden_states + shared_output \
206 - * (1. / self.routed_scaling_factor)
216 + shared_output = None
217 + final_hidden_states = fused_moe_out
218 +
219 + # Fix FP16 overflow
220 + # See DeepseekV2DecoderLayer for more details.
221 + if hidden_states.dtype != torch.float16:
222 + final_hidden_states *= self.routed_scaling_factor
223 + elif self.shared_experts is not None:
224 + assert shared_output is not None
225 + shared_output *= (1. / self.routed_scaling_factor)
226 +
227 + if self.shared_experts is not None:
228 + assert shared_output is not None
229 + final_hidden_states += shared_output
230
231 if self.tp_size > 1:
232 final_hidden_states = (
vllm/model_executor/models/glm4_moe.py
+2
@@ -184,6 +184,8 @@ class Glm4MoE(nn.Module):
184
185 if self.n_shared_experts is not None:
186 shared_output = self.shared_experts(hidden_states)
187 + else:
188 + shared_output = None
189 router_logits = self.gate(hidden_states.to(dtype=torch.float32))
190 final_hidden_states = self.experts(
191 hidden_states=hidden_states,
vllm/model_executor/models/llama4.py
+16 -13
@@ -36,6 +36,7 @@ from vllm.model_executor.layers.linear import (QKVParallelLinear,
36 RowParallelLinear)
37 from vllm.model_executor.layers.quantization import QuantizationConfig
38 from vllm.model_executor.layers.rotary_embedding import get_rope
39 +from vllm.model_executor.layers.shared_fused_moe import SharedFusedMoE
40 from vllm.model_executor.model_loader.weight_utils import (
41 default_weight_loader, maybe_remap_kv_scale_name)
42
@@ -73,7 +74,18 @@ class Llama4MoE(nn.Module):
74 quant_config=None,
75 prefix=f"{prefix}.router")
76
76 - self.experts = FusedMoE(
77 + self.shared_expert = LlamaMLP(
78 + hidden_size=config.hidden_size,
79 + intermediate_size=intermediate_size_moe,
80 + hidden_act="silu",
81 + quant_config=quant_config,
82 + bias=False,
83 + prefix=f"{prefix}.shared_expert",
84 + reduce_results=False,
85 + )
86 +
87 + self.experts = SharedFusedMoE(
88 + shared_experts=self.shared_expert,
89 num_experts=config.num_local_experts,
90 top_k=config.num_experts_per_tok,
91 hidden_size=config.hidden_size,
@@ -83,22 +95,13 @@ class Llama4MoE(nn.Module):
95 reduce_results=False,
96 renormalize=False,
97 quant_config=quant_config,
86 - prefix=f"{prefix}.experts")
87 -
88 - self.shared_expert = LlamaMLP(
89 - hidden_size=config.hidden_size,
90 - intermediate_size=intermediate_size_moe,
91 - hidden_act="silu",
92 - quant_config=quant_config,
93 - bias=False,
94 - prefix=f"{prefix}.shared_expert",
95 - reduce_results=self.experts.must_reduce_shared_expert_outputs(),
98 + prefix=f"{prefix}.experts",
99 )
100
101 def forward(self, hidden_states):
102 router_logits, _ = self.router(hidden_states)
100 - shared_out = self.shared_expert(hidden_states)
101 - routed_out = self.experts(
103 +
104 + shared_out, routed_out = self.experts(
105 hidden_states=hidden_states,
106 router_logits=router_logits,
107 )
vllm/v1/worker/gpu_worker.py
+2 -1
@@ -500,7 +500,8 @@ class Worker(WorkerBase):
500 parallel_config = self.vllm_config.parallel_config
501 moe_modules = [
502 module for module in self.model_runner.model.modules()
503 - if module.__class__.__name__ == "FusedMoE"
503 + if (module.__class__.__name__ == "FusedMoE"
504 + or module.__class__.__name__ == "SharedFusedMoE")
505 ]
506 num_local_experts = moe_modules[0].moe_config.num_local_experts
507 assert all(module.moe_config.num_local_experts == num_local_experts