@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a1448b4b69
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 13, 12:00 UTC

[Kernels] Split up fused_moe/layer.py, isolate more modular kernel code (#28064)

bnellnm committed Nov 11, 2025 at 09:29 UTC a1448b4b69b15c33b4fbc9a883c4f3b9559ee7db
10 files changed +1064 -948
tests/kernels/moe/modular_kernel_tools/mk_objects.py
+5 -4
@@ -6,6 +6,10 @@ import torch
6
7 # Fused experts and PrepareFinalize imports
8 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9 +from vllm.model_executor.layers.fused_moe import TritonExperts
10 +from vllm.model_executor.layers.fused_moe.all2all_utils import (
11 + maybe_make_prepare_finalize,
12 +)
13 from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
14 BatchedDeepGemmExperts,
15 )
@@ -21,7 +25,6 @@ from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
25 BatchedTritonExperts,
26 NaiveBatchedExperts,
27 )
24 -from vllm.model_executor.layers.fused_moe.layer import FusedMoEMethodBase, TritonExperts
28 from vllm.model_executor.layers.fused_moe.prepare_finalize import (
29 MoEPrepareAndFinalizeNoEP,
30 )
@@ -399,9 +402,7 @@ def make_prepare_finalize(
402 quant_config: FusedMoEQuantConfig,
403 ) -> mk.FusedMoEPrepareAndFinalize:
404 if backend != "naive" and backend is not None:
402 - prepare_finalize = FusedMoEMethodBase._maybe_make_prepare_finalize(
403 - moe, quant_config
404 - )
405 + prepare_finalize = maybe_make_prepare_finalize(moe, quant_config)
406 assert prepare_finalize is not None
407 return prepare_finalize
408 elif prepare_finalize_type == FlashInferCutlassMoEPrepareAndFinalize:
vllm/lora/layers/fused_moe.py
+3 -1
@@ -25,7 +25,9 @@ from vllm.model_executor.layers.fused_moe.fused_moe import (
25 modular_triton_fused_moe,
26 try_get_optimal_moe_config,
27 )
28 -from vllm.model_executor.layers.fused_moe.layer import FusedMoEModularMethod
28 +from vllm.model_executor.layers.fused_moe.fused_moe_modular_method import (
29 + FusedMoEModularMethod,
30 +)
31
32
33 class FusedMoEWithLoRA(BaseLayerWithLoRA):
vllm/model_executor/layers/fused_moe/__init__.py
+3 -1
@@ -5,9 +5,11 @@ from contextlib import contextmanager
5 from typing import Any
6
7 from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
8 +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
9 + FusedMoEMethodBase,
10 +)
11 from vllm.model_executor.layers.fused_moe.layer import (
12 FusedMoE,
10 - FusedMoEMethodBase,
13 FusedMoeWeightScaleSupported,
14 )
15 from vllm.model_executor.layers.fused_moe.modular_kernel import (
vllm/model_executor/layers/fused_moe/all2all_utils.py new
+160
@@ -0,0 +1,160 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +
5 +import torch
6 +
7 +from vllm.distributed import (
8 + get_ep_group,
9 +)
10 +from vllm.model_executor.layers.fused_moe.config import (
11 + FusedMoEConfig,
12 + FusedMoEParallelConfig,
13 + FusedMoEQuantConfig,
14 +)
15 +from vllm.model_executor.layers.fused_moe.modular_kernel import (
16 + FusedMoEPrepareAndFinalize,
17 +)
18 +from vllm.platforms import current_platform
19 +from vllm.utils.import_utils import has_deep_ep, has_pplx
20 +
21 +if current_platform.is_cuda_alike():
22 + if has_pplx():
23 + from .pplx_prepare_finalize import (
24 + PplxPrepareAndFinalize,
25 + pplx_hidden_dim_scale_bytes,
26 + )
27 + if has_deep_ep():
28 + from .deepep_ht_prepare_finalize import DeepEPHTPrepareAndFinalize
29 + from .deepep_ll_prepare_finalize import (
30 + DEEPEP_QUANT_BLOCK_SHAPE,
31 + DeepEPLLPrepareAndFinalize,
32 + )
33 +
34 +
35 +def maybe_roundup_layer_hidden_size(
36 + hidden_size: int,
37 + act_dtype: torch.dtype,
38 + moe_parallel_config: FusedMoEParallelConfig,
39 +) -> int:
40 + """
41 + Given layer hidden size and MoE configurations, round up hidden_size
42 + if necessary.
43 +
44 + Args:
45 + hidden_size: Layer hidden-size
46 + act_dtype: Data type of the layer activations.
47 + moe_parallel_config: Fused MoE parallelization strategy configuration.
48 +
49 + Return:
50 + Rounded up hidden_size if rounding up is required based on the configs
51 + and all2all backend.
52 + Original hidden size otherwise.
53 + """
54 + if moe_parallel_config.use_deepep_ht_kernels:
55 + hidden_size = DeepEPHTPrepareAndFinalize.maybe_roundup_layer_hidden_size(
56 + hidden_size, act_dtype
57 + )
58 +
59 + if moe_parallel_config.use_deepep_ll_kernels:
60 + hidden_size = DeepEPLLPrepareAndFinalize.maybe_roundup_layer_hidden_size(
61 + hidden_size
62 + )
63 +
64 + return hidden_size
65 +
66 +
67 +def maybe_make_prepare_finalize(
68 + moe: FusedMoEConfig,
69 + quant_config: FusedMoEQuantConfig | None,
70 +) -> FusedMoEPrepareAndFinalize | None:
71 + if not moe.moe_parallel_config.use_all2all_kernels:
72 + return None
73 +
74 + all2all_manager = get_ep_group().device_communicator.all2all_manager
75 + assert all2all_manager is not None
76 +
77 + prepare_finalize: FusedMoEPrepareAndFinalize | None = None
78 +
79 + # TODO: could allow this now
80 + assert not moe.use_flashinfer_cutlass_kernels, "Must be created in modelopt.py"
81 +
82 + if moe.use_pplx_kernels:
83 + assert quant_config is not None
84 +
85 + hidden_dim_bytes, hidden_scale_bytes = pplx_hidden_dim_scale_bytes(
86 + moe.max_num_tokens,
87 + moe.hidden_dim,
88 + moe.in_dtype,
89 + quant_config.quant_dtype,
90 + per_act_token_quant=quant_config.per_act_token_quant,
91 + block_shape=quant_config.block_shape,
92 + )
93 +
94 + all_to_all_args = dict(
95 + max_num_tokens=moe.max_num_tokens,
96 + num_experts=moe.num_experts,
97 + experts_per_token=moe.experts_per_token, # topk
98 + rank=all2all_manager.rank,
99 + world_size=all2all_manager.world_size,
100 + # dp_size actually means tp_size, bug in pplx kernels
101 + dp_size=all2all_manager.tp_group.world_size,
102 + hidden_dim=moe.hidden_dim,
103 + hidden_dim_bytes=hidden_dim_bytes,
104 + hidden_dim_scale_bytes=hidden_scale_bytes,
105 + )
106 +
107 + num_dispatchers = (
108 + all2all_manager.world_size // all2all_manager.tp_group.world_size
109 + )
110 +
111 + # Intranode pplx a2a takes a group name while internode does not.
112 + if not all2all_manager.internode:
113 + all_to_all_args["group_name"] = all2all_manager.cpu_group.group_name
114 +
115 + handle = all2all_manager.get_handle(all_to_all_args)
116 +
117 + prepare_finalize = PplxPrepareAndFinalize(
118 + handle,
119 + max_num_tokens=moe.max_num_tokens,
120 + num_local_experts=moe.num_local_experts,
121 + num_dispatchers=num_dispatchers,
122 + )
123 + elif moe.use_deepep_ht_kernels:
124 + assert moe.dp_size == all2all_manager.dp_world_size
125 +
126 + all_to_all_args = dict()
127 + handle = all2all_manager.get_handle(all_to_all_args)
128 + prepare_finalize = DeepEPHTPrepareAndFinalize(
129 + handle,
130 + num_dispatchers=all2all_manager.world_size,
131 + dp_size=all2all_manager.dp_world_size,
132 + rank_expert_offset=all2all_manager.rank * moe.num_local_experts,
133 + )
134 +
135 + elif moe.use_deepep_ll_kernels:
136 + assert quant_config is not None
137 + all_to_all_args = dict(
138 + max_num_tokens_per_dp_rank=moe.max_num_tokens,
139 + token_hidden_size=moe.hidden_dim,
140 + num_ep_ranks=all2all_manager.world_size,
141 + num_global_experts=moe.num_experts,
142 + num_local_experts=moe.num_experts // all2all_manager.world_size,
143 + )
144 + handle = all2all_manager.get_handle(all_to_all_args)
145 +
146 + # Note: We may want to use FP8 dispatch just to reduce
147 + # data movement.
148 + use_fp8_dispatch = (
149 + quant_config.quant_dtype == current_platform.fp8_dtype()
150 + and quant_config.block_shape == DEEPEP_QUANT_BLOCK_SHAPE
151 + )
152 +
153 + prepare_finalize = DeepEPLLPrepareAndFinalize(
154 + handle,
155 + max_tokens_per_rank=moe.max_num_tokens,
156 + num_dispatchers=all2all_manager.world_size,
157 + use_fp8_dispatch=use_fp8_dispatch,
158 + )
159 +
160 + return prepare_finalize
vllm/model_executor/layers/fused_moe/fused_moe_method_base.py new
+112
@@ -0,0 +1,112 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from abc import abstractmethod
5 +from collections.abc import Callable
6 +
7 +import torch
8 +
9 +from vllm.logger import init_logger
10 +from vllm.model_executor.layers.fused_moe.config import (
11 + FusedMoEConfig,
12 + FusedMoEQuantConfig,
13 +)
14 +from vllm.model_executor.layers.fused_moe.modular_kernel import (
15 + FusedMoEPermuteExpertsUnpermute,
16 + FusedMoEPrepareAndFinalize,
17 +)
18 +from vllm.model_executor.layers.quantization.base_config import (
19 + QuantizeMethodBase,
20 +)
21 +
22 +logger = init_logger(__name__)
23 +
24 +
25 +class FusedMoEMethodBase(QuantizeMethodBase):
26 + def __init__(self, moe: FusedMoEConfig):
27 + super().__init__()
28 + self.moe: FusedMoEConfig = moe
29 + self.moe_quant_config: FusedMoEQuantConfig | None = None
30 +
31 + @abstractmethod
32 + def create_weights(
33 + self,
34 + layer: torch.nn.Module,
35 + num_experts: int,
36 + hidden_size: int,
37 + intermediate_size_per_partition: int,
38 + params_dtype: torch.dtype,
39 + **extra_weight_attrs,
40 + ):
41 + raise NotImplementedError
42 +
43 + def uses_weight_scale_2_pattern(self) -> bool:
44 + """
45 + Returns True if this quantization method uses 'weight_scale_2' pattern
46 + for per-tensor weight scales (e.g., FP4 variants), False otherwise.
47 +
48 + This method should be overridden by subclasses that use the
49 + 'weight_scale_2' pattern instead of the standard 'weight_scale' pattern.
50 + """
51 + return False
52 +
53 + def maybe_make_prepare_finalize(self) -> FusedMoEPrepareAndFinalize | None:
54 + from .all2all_utils import maybe_make_prepare_finalize
55 +
56 + return maybe_make_prepare_finalize(self.moe, self.moe_quant_config)
57 +
58 + def select_gemm_impl(
59 + self,
60 + prepare_finalize: FusedMoEPrepareAndFinalize,
61 + layer: torch.nn.Module,
62 + ) -> FusedMoEPermuteExpertsUnpermute:
63 + # based on the all2all implementation, select the appropriate
64 + # gemm implementation
65 + raise NotImplementedError(
66 + f"{self.__class__.__name__} must select appropriate gemm "
67 + "implementation based on the prepare_finalize"
68 + )
69 +
70 + @abstractmethod
71 + def get_fused_moe_quant_config(
72 + self, layer: torch.nn.Module
73 + ) -> FusedMoEQuantConfig | None:
74 + raise NotImplementedError
75 +
76 + @property
77 + def topk_indices_dtype(self) -> torch.dtype | None:
78 + return None
79 +
80 + @property
81 + def supports_eplb(self) -> bool:
82 + return False
83 +
84 + @property
85 + def allow_inplace(self) -> bool:
86 + return False
87 +
88 + @abstractmethod
89 + def apply(
90 + self,
91 + layer: torch.nn.Module,
92 + x: torch.Tensor,
93 + router_logits: torch.Tensor,
94 + top_k: int,
95 + renormalize: bool,
96 + use_grouped_topk: bool = False,
97 + topk_group: int | None = None,
98 + num_expert_group: int | None = None,
99 + global_num_experts: int = -1,
100 + expert_map: torch.Tensor | None = None,
101 + custom_routing_function: Callable | None = None,
102 + scoring_func: str = "softmax",
103 + routed_scaling_factor: float = 1.0,
104 + e_score_correction_bias: torch.Tensor | None = None,
105 + apply_router_weight_on_input: bool = False,
106 + activation: str = "silu",
107 + enable_eplb: bool = False,
108 + expert_load_view: torch.Tensor | None = None,
109 + logical_to_physical_map: torch.Tensor | None = None,
110 + logical_replica_count: torch.Tensor | None = None,
111 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
112 + raise NotImplementedError
vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py new
+164
@@ -0,0 +1,164 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from collections.abc import Callable
5 +
6 +import torch
7 +
8 +from vllm.logger import init_logger
9 +from vllm.model_executor.custom_op import CustomOp
10 +from vllm.model_executor.layers.fused_moe.config import (
11 + FusedMoEQuantConfig,
12 +)
13 +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
14 + FusedMoEMethodBase,
15 +)
16 +from vllm.model_executor.layers.fused_moe.modular_kernel import (
17 + FusedMoEModularKernel,
18 + FusedMoEPrepareAndFinalize,
19 +)
20 +
21 +logger = init_logger(__name__)
22 +
23 +
24 +@CustomOp.register("modular_fused_moe")
25 +class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
26 + def __init__(
27 + self, old_quant_method: FusedMoEMethodBase, experts: FusedMoEModularKernel
28 + ):
29 + super().__init__(old_quant_method.moe)
30 + self.moe_quant_config = old_quant_method.moe_quant_config
31 + self.fused_experts = experts
32 + self.disable_expert_map = getattr(
33 + old_quant_method,
34 + "disable_expert_map",
35 + not self.fused_experts.supports_expert_map(),
36 + )
37 + self.old_quant_method = old_quant_method
38 + logger.debug("Swapping out %s", self.old_quant_method.__class__.__name__)
39 +
40 + @staticmethod
41 + def make(
42 + moe_layer: torch.nn.Module,
43 + old_quant_method: FusedMoEMethodBase,
44 + prepare_finalize: FusedMoEPrepareAndFinalize,
45 + shared_experts: torch.nn.Module | None,
46 + ) -> "FusedMoEModularMethod":
47 + return FusedMoEModularMethod(
48 + old_quant_method,
49 + FusedMoEModularKernel(
50 + prepare_finalize,
51 + old_quant_method.select_gemm_impl(prepare_finalize, moe_layer),
52 + shared_experts,
53 + ),
54 + )
55 +
56 + @property
57 + def topk_indices_dtype(self) -> torch.dtype | None:
58 + return self.fused_experts.prepare_finalize.topk_indices_dtype()
59 +
60 + @property
61 + def supports_eplb(self) -> bool:
62 + return self.old_quant_method.supports_eplb
63 +
64 + @property
65 + def allow_inplace(self) -> bool:
66 + return self.old_quant_method.allow_inplace
67 +
68 + def create_weights(
69 + self,
70 + layer: torch.nn.Module,
71 + num_experts: int,
72 + hidden_size: int,
73 + intermediate_size_per_partition: int,
74 + params_dtype: torch.dtype,
75 + **extra_weight_attrs,
76 + ):
77 + raise NotImplementedError
78 +
79 + def get_fused_moe_quant_config(
80 + self, layer: torch.nn.Module
81 + ) -> FusedMoEQuantConfig | None:
82 + return self.moe_quant_config
83 +
84 + def apply(
85 + self,
86 + layer: torch.nn.Module,
87 + x: torch.Tensor,
88 + router_logits: torch.Tensor,
89 + top_k: int,
90 + renormalize: bool,
91 + use_grouped_topk: bool = False,
92 + topk_group: int | None = None,
93 + num_expert_group: int | None = None,
94 + global_num_experts: int = -1,
95 + expert_map: torch.Tensor | None = None,
96 + custom_routing_function: Callable | None = None,
97 + scoring_func: str = "softmax",
98 + routed_scaling_factor: float = 1.0,
99 + e_score_correction_bias: torch.Tensor | None = None,
100 + apply_router_weight_on_input: bool = False,
101 + activation: str = "silu",
102 + enable_eplb: bool = False,
103 + expert_load_view: torch.Tensor | None = None,
104 + logical_to_physical_map: torch.Tensor | None = None,
105 + logical_replica_count: torch.Tensor | None = None,
106 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
107 + # Is getattr needed?
108 + zero_expert_num = getattr(layer, "zero_expert_num", 0)
109 + zero_expert_type = getattr(layer, "zero_expert_type", None)
110 +
111 + if enable_eplb:
112 + if self.supports_eplb:
113 + assert expert_load_view is not None
114 + assert logical_to_physical_map is not None
115 + assert logical_replica_count is not None
116 + else:
117 + raise NotImplementedError(
118 + "EPLB is not supported for "
119 + f"{self.old_quant_method.__class__.__name__}."
120 + )
121 +
122 + topk_weights, topk_ids, zero_expert_result = layer.select_experts(
123 + hidden_states=x,
124 + router_logits=router_logits,
125 + use_grouped_topk=use_grouped_topk,
126 + top_k=top_k,
127 + renormalize=renormalize,
128 + topk_group=topk_group,
129 + num_expert_group=num_expert_group,
130 + custom_routing_function=custom_routing_function,
131 + scoring_func=scoring_func,
132 + routed_scaling_factor=routed_scaling_factor,
133 + e_score_correction_bias=e_score_correction_bias,
134 + indices_type=self.topk_indices_dtype,
135 + enable_eplb=enable_eplb,
136 + expert_map=expert_map,
137 + expert_load_view=expert_load_view,
138 + logical_to_physical_map=logical_to_physical_map,
139 + logical_replica_count=logical_replica_count,
140 + global_num_experts=global_num_experts,
141 + zero_expert_num=zero_expert_num,
142 + zero_expert_type=zero_expert_type,
143 + )
144 +
145 + result = self.fused_experts(
146 + hidden_states=x,
147 + w1=layer.w13_weight,
148 + w2=layer.w2_weight,
149 + topk_weights=topk_weights,
150 + topk_ids=topk_ids,
151 + inplace=self.allow_inplace,
152 + activation=activation,
153 + global_num_experts=global_num_experts,
154 + apply_router_weight_on_input=apply_router_weight_on_input,
155 + expert_map=None if self.disable_expert_map else expert_map,
156 + )
157 +
158 + if zero_expert_num != 0 and zero_expert_type is not None:
159 + assert not isinstance(result, tuple), (
160 + "Shared + zero experts are mutually exclusive not yet supported"
161 + )
162 + return result, zero_expert_result
163 + else:
164 + return result
vllm/model_executor/layers/fused_moe/layer.py
+32 -918
@@ -1,7 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from abc import abstractmethod
4 from collections.abc import Callable, Iterable
5 from contextlib import nullcontext
6 from enum import Enum
@@ -27,17 +26,13 @@ from vllm.forward_context import ForwardContext, get_forward_context
26 from vllm.logger import init_logger
27 from vllm.model_executor.custom_op import CustomOp
28 from vllm.model_executor.layers.fused_moe.config import (
30 - FUSED_MOE_UNQUANTIZED_CONFIG,
29 FusedMoEConfig,
30 FusedMoEParallelConfig,
31 FusedMoEQuantConfig,
32 RoutingMethodType,
35 - biased_moe_quant_config,
33 )
34 from vllm.model_executor.layers.fused_moe.fused_moe import zero_experts_compute_triton
35 from vllm.model_executor.layers.fused_moe.modular_kernel import (
39 - FusedMoEActivationFormat,
40 - FusedMoEModularKernel,
36 FusedMoEPermuteExpertsUnpermute,
37 FusedMoEPrepareAndFinalize,
38 )
@@ -47,35 +42,17 @@ from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
42 from vllm.model_executor.layers.fused_moe.routing_simulator import RoutingSimulator
43 from vllm.model_executor.layers.quantization.base_config import (
44 QuantizationConfig,
50 - QuantizeMethodBase,
45 )
46 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
47 is_flashinfer_supporting_global_sf,
48 )
55 -from vllm.model_executor.utils import set_weight_attrs
49 from vllm.platforms import current_platform
57 -from vllm.platforms.interface import CpuArchEnum
58 -from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
59 -from vllm.utils.import_utils import has_deep_ep, has_pplx
50 from vllm.utils.math_utils import cdiv, round_up
51 from vllm.utils.torch_utils import current_stream, direct_register_custom_op
52 from vllm.v1.worker.ubatching import dbo_current_ubatch_id
53
54 if current_platform.is_cuda_alike():
65 - from .fused_batched_moe import BatchedTritonExperts
66 - from .fused_moe import TritonExperts, eplb_map_to_physical_and_record, fused_experts
67 -
68 - if has_pplx():
69 - from .pplx_prepare_finalize import (
70 - PplxPrepareAndFinalize,
71 - pplx_hidden_dim_scale_bytes,
72 - )
73 - if has_deep_ep():
74 - from .deepep_ht_prepare_finalize import DeepEPHTPrepareAndFinalize
75 - from .deepep_ll_prepare_finalize import (
76 - DEEPEP_QUANT_BLOCK_SHAPE,
77 - DeepEPLLPrepareAndFinalize,
78 - )
55 + from .fused_moe import eplb_map_to_physical_and_record, fused_experts
56 else:
57 fused_experts = None # type: ignore
58 FusedMoEPermuteExpertsUnpermute = object # type: ignore
@@ -102,6 +79,16 @@ if current_platform.is_tpu():
79 else:
80 fused_moe_pallas = None # type: ignore
81
82 +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
83 + FusedMoEMethodBase,
84 +)
85 +from vllm.model_executor.layers.fused_moe.fused_moe_modular_method import (
86 + FusedMoEModularMethod,
87 +)
88 +from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
89 + UnquantizedFusedMoEMethod,
90 +)
91 +
92 logger = init_logger(__name__)
93
94
@@ -112,885 +99,6 @@ class FusedMoeWeightScaleSupported(Enum):
99 BLOCK = "block"
100
101
115 -class FusedMoEMethodBase(QuantizeMethodBase):
116 - def __init__(self, moe: FusedMoEConfig):
117 - super().__init__()
118 - self.moe: FusedMoEConfig = moe
119 - self.moe_quant_config: FusedMoEQuantConfig | None = None
120 -
121 - @abstractmethod
122 - def create_weights(
123 - self,
124 - layer: torch.nn.Module,
125 - num_experts: int,
126 - hidden_size: int,
127 - intermediate_size_per_partition: int,
128 - params_dtype: torch.dtype,
129 - **extra_weight_attrs,
130 - ):
131 - raise NotImplementedError
132 -
133 - def uses_weight_scale_2_pattern(self) -> bool:
134 - """
135 - Returns True if this quantization method uses 'weight_scale_2' pattern
136 - for per-tensor weight scales (e.g., FP4 variants), False otherwise.
137 -
138 - This method should be overridden by subclasses that use the
139 - 'weight_scale_2' pattern instead of the standard 'weight_scale' pattern.
140 - """
141 - return False
142 -
143 - @staticmethod
144 - def _maybe_make_prepare_finalize(
145 - moe: FusedMoEConfig,
146 - quant_config: FusedMoEQuantConfig | None,
147 - ) -> FusedMoEPrepareAndFinalize | None:
148 - all2all_manager = get_ep_group().device_communicator.all2all_manager
149 - assert all2all_manager is not None
150 -
151 - prepare_finalize: FusedMoEPrepareAndFinalize | None = None
152 -
153 - # TODO: could allow this now
154 - assert not moe.use_flashinfer_cutlass_kernels, "Must be created in modelopt.py"
155 -
156 - if moe.use_pplx_kernels:
157 - assert quant_config is not None
158 -
159 - hidden_dim_bytes, hidden_scale_bytes = pplx_hidden_dim_scale_bytes(
160 - moe.max_num_tokens,
161 - moe.hidden_dim,
162 - moe.in_dtype,
163 - quant_config.quant_dtype,
164 - per_act_token_quant=quant_config.per_act_token_quant,
165 - block_shape=quant_config.block_shape,
166 - )
167 -
168 - all_to_all_args = dict(
169 - max_num_tokens=moe.max_num_tokens,
170 - num_experts=moe.num_experts,
171 - experts_per_token=moe.experts_per_token, # topk
172 - rank=all2all_manager.rank,
173 - world_size=all2all_manager.world_size,
174 - # dp_size actually means tp_size, bug in pplx kernels
175 - dp_size=all2all_manager.tp_group.world_size,
176 - hidden_dim=moe.hidden_dim,
177 - hidden_dim_bytes=hidden_dim_bytes,
178 - hidden_dim_scale_bytes=hidden_scale_bytes,
179 - )
180 -
181 - num_dispatchers = (
182 - all2all_manager.world_size // all2all_manager.tp_group.world_size
183 - )
184 -
185 - # Intranode pplx a2a takes a group name while internode does not.
186 - if not all2all_manager.internode:
187 - all_to_all_args["group_name"] = all2all_manager.cpu_group.group_name
188 -
189 - handle = all2all_manager.get_handle(all_to_all_args)
190 -
191 - prepare_finalize = PplxPrepareAndFinalize(
192 - handle,
193 - max_num_tokens=moe.max_num_tokens,
194 - num_local_experts=moe.num_local_experts,
195 - num_dispatchers=num_dispatchers,
196 - )
197 - elif moe.use_deepep_ht_kernels:
198 - assert moe.dp_size == all2all_manager.dp_world_size
199 -
200 - all_to_all_args = dict()
201 - handle = all2all_manager.get_handle(all_to_all_args)
202 - prepare_finalize = DeepEPHTPrepareAndFinalize(
203 - handle,
204 - num_dispatchers=all2all_manager.world_size,
205 - dp_size=all2all_manager.dp_world_size,
206 - rank_expert_offset=all2all_manager.rank * moe.num_local_experts,
207 - )
208 -
209 - elif moe.use_deepep_ll_kernels:
210 - assert quant_config is not None
211 - all_to_all_args = dict(
212 - max_num_tokens_per_dp_rank=moe.max_num_tokens,
213 - token_hidden_size=moe.hidden_dim,
214 - num_ep_ranks=all2all_manager.world_size,
215 - num_global_experts=moe.num_experts,
216 - num_local_experts=moe.num_experts // all2all_manager.world_size,
217 - )
218 - handle = all2all_manager.get_handle(all_to_all_args)
219 -
220 - # Note: We may want to use FP8 dispatch just to reduce
221 - # data movement.
222 - use_fp8_dispatch = (
223 - quant_config.quant_dtype == current_platform.fp8_dtype()
224 - and quant_config.block_shape == DEEPEP_QUANT_BLOCK_SHAPE
225 - )
226 -
227 - prepare_finalize = DeepEPLLPrepareAndFinalize(
228 - handle,
229 - max_tokens_per_rank=moe.max_num_tokens,
230 - num_dispatchers=all2all_manager.world_size,
231 - use_fp8_dispatch=use_fp8_dispatch,
232 - )
233 -
234 - return prepare_finalize
235 -
236 - def maybe_make_prepare_finalize(self) -> FusedMoEPrepareAndFinalize | None:
237 - if self.moe.moe_parallel_config.use_all2all_kernels:
238 - return FusedMoEMethodBase._maybe_make_prepare_finalize(
239 - self.moe, self.moe_quant_config
240 - )
241 - else:
242 - return None
243 -
244 - def maybe_init_modular_kernel(
245 - self, layer: torch.nn.Module
246 - ) -> FusedMoEModularKernel | None:
247 - assert self.moe is not None
248 -
249 - # We must get the quant config here so that the layer is
250 - # completely initialized, i.e. all weights loaded and post
251 - # processed.
252 - self.moe_quant_config = self.get_fused_moe_quant_config(layer)
253 -
254 - prepare_finalize = self.maybe_make_prepare_finalize()
255 -
256 - if prepare_finalize is not None:
257 - logger.debug(
258 - "%s for %s(%s)", prepare_finalize.__class__.__name__, self, id(self)
259 - )
260 - experts = self.select_gemm_impl(prepare_finalize, layer)
261 - return FusedMoEModularKernel(
262 - prepare_finalize,
263 - experts,
264 - layer.shared_experts,
265 - )
266 - else:
267 - return None
268 -
269 - def select_gemm_impl(
270 - self,
271 - prepare_finalize: FusedMoEPrepareAndFinalize,
272 - layer: torch.nn.Module,
273 - ) -> FusedMoEPermuteExpertsUnpermute:
274 - # based on the all2all implementation, select the appropriate
275 - # gemm implementation
276 - raise NotImplementedError(
277 - f"{self.__class__.__name__} must select appropriate gemm "
278 - "implementation based on the prepare_finalize"
279 - )
280 -
281 - @abstractmethod
282 - def get_fused_moe_quant_config(
283 - self, layer: torch.nn.Module
284 - ) -> FusedMoEQuantConfig | None:
285 - raise NotImplementedError
286 -
287 - @property
288 - def topk_indices_dtype(self) -> torch.dtype | None:
289 - return None
290 -
291 - @property
292 - def supports_eplb(self) -> bool:
293 - return False
294 -
295 - @property
296 - def allow_inplace(self) -> bool:
297 - return False
298 -
299 - @abstractmethod
300 - def apply(
301 - self,
302 - layer: torch.nn.Module,
303 - x: torch.Tensor,
304 - router_logits: torch.Tensor,
305 - top_k: int,
306 - renormalize: bool,
307 - use_grouped_topk: bool = False,
308 - topk_group: int | None = None,
309 - num_expert_group: int | None = None,
310 - global_num_experts: int = -1,
311 - expert_map: torch.Tensor | None = None,
312 - custom_routing_function: Callable | None = None,
313 - scoring_func: str = "softmax",
314 - routed_scaling_factor: float = 1.0,
315 - e_score_correction_bias: torch.Tensor | None = None,
316 - apply_router_weight_on_input: bool = False,
317 - activation: str = "silu",
318 - enable_eplb: bool = False,
319 - expert_load_view: torch.Tensor | None = None,
320 - logical_to_physical_map: torch.Tensor | None = None,
321 - logical_replica_count: torch.Tensor | None = None,
322 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
323 - raise NotImplementedError
324 -
325 -
326 -@CustomOp.register("modular_fused_moe")
327 -class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
328 - def __init__(
329 - self,
330 - old_quant_method: FusedMoEMethodBase,
331 - fused_experts: FusedMoEModularKernel,
332 - ):
333 - super().__init__(old_quant_method.moe)
334 - # Find better way to copy attributes? Should we even copy attributes?
335 - # self.__dict__.update(old_quant_method.__dict__)
336 - self.moe_quant_config = old_quant_method.moe_quant_config
337 - self.fused_experts = fused_experts
338 - self.disable_expert_map = getattr(
339 - old_quant_method,
340 - "disable_expert_map",
341 - not fused_experts.supports_expert_map(),
342 - )
343 - self.old_quant_method = old_quant_method
344 - logger.debug("Swapping out %s", self.old_quant_method.__class__.__name__)
345 -
346 - @property
347 - def topk_indices_dtype(self) -> torch.dtype | None:
348 - return self.fused_experts.prepare_finalize.topk_indices_dtype()
349 -
350 - @property
351 - def supports_eplb(self) -> bool:
352 - return self.old_quant_method.supports_eplb
353 -
354 - @property
355 - def allow_inplace(self) -> bool:
356 - return self.old_quant_method.allow_inplace
357 -
358 - def create_weights(
359 - self,
360 - layer: torch.nn.Module,
361 - num_experts: int,
362 - hidden_size: int,
363 - intermediate_size_per_partition: int,
364 - params_dtype: torch.dtype,
365 - **extra_weight_attrs,
366 - ):
367 - raise NotImplementedError
368 -
369 - def get_fused_moe_quant_config(
370 - self, layer: torch.nn.Module
371 - ) -> FusedMoEQuantConfig | None:
372 - return self.moe_quant_config
373 -
374 - def apply(
375 - self,
376 - layer: torch.nn.Module,
377 - x: torch.Tensor,
378 - router_logits: torch.Tensor,
379 - top_k: int,
380 - renormalize: bool,
381 - use_grouped_topk: bool = False,
382 - topk_group: int | None = None,
383 - num_expert_group: int | None = None,
384 - global_num_experts: int = -1,
385 - expert_map: torch.Tensor | None = None,
386 - custom_routing_function: Callable | None = None,
387 - scoring_func: str = "softmax",
388 - routed_scaling_factor: float = 1.0,
389 - e_score_correction_bias: torch.Tensor | None = None,
390 - apply_router_weight_on_input: bool = False,
391 - activation: str = "silu",
392 - enable_eplb: bool = False,
393 - expert_load_view: torch.Tensor | None = None,
394 - logical_to_physical_map: torch.Tensor | None = None,
395 - logical_replica_count: torch.Tensor | None = None,
396 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
397 - # Is getattr needed?
398 - zero_expert_num = getattr(layer, "zero_expert_num", 0)
399 - zero_expert_type = getattr(layer, "zero_expert_type", None)
400 -
401 - if enable_eplb:
402 - if self.supports_eplb:
403 - assert expert_load_view is not None
404 - assert logical_to_physical_map is not None
405 - assert logical_replica_count is not None
406 - assert isinstance(layer, FusedMoE)
407 - else:
408 - raise NotImplementedError(
409 - "EPLB is not supported for "
410 - f"{self.old_quant_method.__class__.__name__}."
411 - )
412 -
413 - topk_weights, topk_ids, zero_expert_result = FusedMoE.select_experts(
414 - hidden_states=x,
415 - router_logits=router_logits,
416 - use_grouped_topk=use_grouped_topk,
417 - top_k=top_k,
418 - renormalize=renormalize,
419 - topk_group=topk_group,
420 - num_expert_group=num_expert_group,
421 - custom_routing_function=custom_routing_function,
422 - scoring_func=scoring_func,
423 - routed_scaling_factor=routed_scaling_factor,
424 - e_score_correction_bias=e_score_correction_bias,
425 - indices_type=self.topk_indices_dtype,
426 - enable_eplb=enable_eplb,
427 - expert_map=expert_map,
428 - expert_load_view=expert_load_view,
429 - logical_to_physical_map=logical_to_physical_map,
430 - logical_replica_count=logical_replica_count,
431 - global_num_experts=global_num_experts,
432 - zero_expert_num=zero_expert_num,
433 - zero_expert_type=zero_expert_type,
434 - )
435 -
436 - result = self.fused_experts(
437 - hidden_states=x,
438 - w1=layer.w13_weight,
439 - w2=layer.w2_weight,
440 - topk_weights=topk_weights,
441 - topk_ids=topk_ids,
442 - inplace=self.allow_inplace,
443 - activation=activation,
444 - global_num_experts=global_num_experts,
445 - apply_router_weight_on_input=apply_router_weight_on_input,
446 - expert_map=None if self.disable_expert_map else expert_map,
447 - )
448 -
449 - if zero_expert_num != 0 and zero_expert_type is not None:
450 - assert not isinstance(result, tuple), (
451 - "Shared + zero experts are mutually exclusive not yet supported"
452 - )
453 - return result, zero_expert_result
454 - else:
455 - return result
456 -
457 -
458 -@CustomOp.register("unquantized_fused_moe")
459 -class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
460 - """MoE method without quantization."""
461 -
462 - def __init__(self, moe: FusedMoEConfig):
463 - super().__init__(moe)
464 -
465 - self.rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
466 - if self.rocm_aiter_moe_enabled:
467 - from .rocm_aiter_fused_moe import rocm_aiter_fused_experts
468 -
469 - self.rocm_aiter_fused_experts = rocm_aiter_fused_experts
470 - else:
471 - self.rocm_aiter_fused_experts = None # type: ignore
472 -
473 - # FlashInfer CUTLASS MoE is only supported on Hopper and later GPUS
474 - self.flashinfer_cutlass_moe_enabled = (
475 - has_flashinfer_cutlass_fused_moe()
476 - and envs.VLLM_USE_FLASHINFER_MOE_FP16
477 - and self.moe.moe_parallel_config.use_ep
478 - and self.moe.moe_parallel_config.dp_size == 1
479 - and current_platform.get_device_capability()[0] >= 9
480 - )
481 - if self.flashinfer_cutlass_moe_enabled:
482 - logger.info_once(
483 - "Enabling FlashInfer CUTLASS MoE for UnquantizedFusedMoEMethod"
484 - )
485 - from functools import partial
486 -
487 - from .flashinfer_cutlass_moe import flashinfer_cutlass_moe
488 -
489 - self.flashinfer_cutlass_moe = partial(
490 - flashinfer_cutlass_moe,
491 - quant_config=FUSED_MOE_UNQUANTIZED_CONFIG,
492 - tp_rank=self.moe.moe_parallel_config.tp_rank,
493 - tp_size=self.moe.moe_parallel_config.tp_size,
494 - ep_rank=self.moe.moe_parallel_config.ep_rank,
495 - ep_size=self.moe.moe_parallel_config.ep_size,
496 - )
497 - else:
498 - if (
499 - self.moe.moe_parallel_config.use_ep
500 - and self.moe.moe_parallel_config.dp_size == 1
501 - ):
502 - logger.info_once(
503 - "FlashInfer CUTLASS MoE is available for EP"
504 - " but not enabled, consider setting"
505 - " VLLM_USE_FLASHINFER_MOE_FP16=1 to enable it.",
506 - scope="local",
507 - )
508 - elif self.moe.moe_parallel_config.dp_size > 1:
509 - logger.info_once(
510 - "FlashInfer CUTLASS MoE is currently not available for DP.",
511 - scope="local",
512 - )
513 - self.flashinfer_cutlass_moe = None # type: ignore
514 -
515 - @property
516 - def supports_eplb(self) -> bool:
517 - return True
518 -
519 - @property
520 - def allow_inplace(self) -> bool:
521 - return True
522 -
523 - def maybe_make_prepare_finalize(self) -> FusedMoEPrepareAndFinalize | None:
524 - if self.rocm_aiter_moe_enabled:
525 - return None
526 - else:
527 - return super().maybe_make_prepare_finalize()
528 -
529 - def select_gemm_impl(
530 - self,
531 - prepare_finalize: FusedMoEPrepareAndFinalize,
532 - layer: torch.nn.Module,
533 - ) -> FusedMoEPermuteExpertsUnpermute:
534 - assert self.moe_quant_config is not None
535 - if (
536 - prepare_finalize.activation_format
537 - == FusedMoEActivationFormat.BatchedExperts
538 - ):
539 - logger.debug("BatchedTritonExperts %s", self.moe)
540 - return BatchedTritonExperts(
541 - max_num_tokens=self.moe.max_num_tokens,
542 - num_dispatchers=prepare_finalize.num_dispatchers(),
543 - quant_config=self.moe_quant_config,
544 - )
545 - else:
546 - logger.debug("TritonExperts %s", self.moe)
547 - return TritonExperts(self.moe_quant_config)
548 -
549 - def create_weights(
550 - self,
551 - layer: torch.nn.Module,
552 - num_experts: int,
553 - hidden_size: int,
554 - intermediate_size_per_partition: int,
555 - params_dtype: torch.dtype,
556 - **extra_weight_attrs,
557 - ):
558 - if self.moe.is_act_and_mul:
559 - w13_up_dim = 2 * intermediate_size_per_partition
560 - else:
561 - w13_up_dim = intermediate_size_per_partition
562 - # Fused gate_up_proj (column parallel)
563 - w13_weight = torch.nn.Parameter(
564 - torch.empty(
565 - num_experts,
566 - w13_up_dim,
567 - hidden_size,
568 - dtype=params_dtype,
569 - ),
570 - requires_grad=False,
571 - )
572 - layer.register_parameter("w13_weight", w13_weight)
573 - set_weight_attrs(w13_weight, extra_weight_attrs)
574 - if self.moe.has_bias:
575 - w13_bias = torch.nn.Parameter(
576 - torch.zeros(num_experts, w13_up_dim, dtype=params_dtype),
577 - requires_grad=False,
578 - )
579 - layer.register_parameter("w13_bias", w13_bias)
580 - set_weight_attrs(w13_bias, extra_weight_attrs)
581 - # down_proj (row parallel)
582 - w2_weight = torch.nn.Parameter(
583 - torch.empty(
584 - num_experts,
585 - hidden_size,
586 - intermediate_size_per_partition,
587 - dtype=params_dtype,
588 - ),
589 - requires_grad=False,
590 - )
591 - layer.register_parameter("w2_weight", w2_weight)
592 - set_weight_attrs(w2_weight, extra_weight_attrs)
593 - if self.moe.has_bias:
594 - w2_bias = torch.nn.Parameter(
595 - torch.zeros(num_experts, hidden_size, dtype=params_dtype),
596 - requires_grad=False,
597 - )
598 - layer.register_parameter("w2_bias", w2_bias)
599 - set_weight_attrs(w2_bias, extra_weight_attrs)
600 -
601 - def _maybe_pad_weight(self, weight: torch.Tensor) -> torch.Tensor:
602 - # Pad the weight tensor. This is an optimization on ROCm platform, which
603 - # can benefit from tensors located far enough from one another in memory
604 - if (
605 - envs.VLLM_ROCM_MOE_PADDING
606 - and current_platform.is_rocm()
607 - and weight.stride(-1) == 1
608 - and (weight.stride(-2) * weight.element_size()) % 512 == 0
609 - ):
610 - num_pad = 256 // weight.element_size()
611 - weight = F.pad(weight, (0, num_pad), "constant", 0)[..., :-num_pad]
612 - torch.cuda.empty_cache()
613 -
614 - return weight
615 -
616 - def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
617 - super().process_weights_after_loading(layer)
618 -
619 - # Padding the weight for better performance on ROCm
620 - layer.w13_weight.data = self._maybe_pad_weight(layer.w13_weight.data)
621 - layer.w2_weight.data = self._maybe_pad_weight(layer.w2_weight.data)
622 -
623 - if self.rocm_aiter_moe_enabled:
624 - shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights(
625 - layer.w13_weight.data, layer.w2_weight.data
626 - )
627 -
628 - layer.w13_weight.data = shuffled_w13
629 - layer.w2_weight.data = shuffled_w2
630 -
631 - if self.flashinfer_cutlass_moe_enabled:
632 - # Swap halves to arrange as [w3; w1] (kernel expectation)
633 - w1_w, w3_w = torch.chunk(layer.w13_weight.data, 2, dim=1)
634 - w13_weight_swapped = torch.cat([w3_w, w1_w], dim=1)
635 - layer.w13_weight.data = w13_weight_swapped.contiguous()
636 -
637 - if current_platform.is_xpu():
638 - import intel_extension_for_pytorch as ipex
639 -
640 - ep_rank_start = self.moe.ep_rank * self.moe.num_local_experts
641 - layer.ipex_fusion = ipex.llm.modules.GatedMLPMOE(
642 - layer.w13_weight,
643 - layer.w2_weight,
644 - use_prepack=True,
645 - experts_start_id=ep_rank_start,
646 - )
647 - elif current_platform.is_cpu():
648 - from vllm.model_executor.layers.fused_moe import cpu_fused_moe
649 -
650 - if current_platform.get_cpu_architecture() == CpuArchEnum.X86:
651 - from vllm.model_executor.layers.utils import check_cpu_sgl_kernel
652 -
653 - dtype_w13 = layer.w13_weight.dtype
654 - _, n_w13, k_w13 = layer.w13_weight.size()
655 - dtype_w2 = layer.w2_weight.dtype
656 - _, n_w2, k_w2 = layer.w2_weight.size()
657 - if (
658 - envs.VLLM_CPU_SGL_KERNEL
659 - and check_cpu_sgl_kernel(n_w13, k_w13, dtype_w13)
660 - and check_cpu_sgl_kernel(n_w2, k_w2, dtype_w2)
661 - ):
662 - packed_w13_weight = torch.ops._C.convert_weight_packed(
663 - layer.w13_weight
664 - )
665 - assert packed_w13_weight.size() == layer.w13_weight.size()
666 - layer.w13_weight.copy_(packed_w13_weight)
667 - del packed_w13_weight
668 - packed_w2_weight = torch.ops._C.convert_weight_packed(
669 - layer.w2_weight
670 - )
671 - assert packed_w2_weight.size() == layer.w2_weight.size()
672 - layer.w2_weight.copy_(packed_w2_weight)
673 - layer.cpu_fused_moe = cpu_fused_moe.SGLFusedMOE(layer)
674 - else:
675 - layer.cpu_fused_moe = cpu_fused_moe.IPEXFusedMOE(layer)
676 - else:
677 - layer.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer)
678 -
679 - def apply(
680 - self,
681 - layer: torch.nn.Module,
682 - x: torch.Tensor,
683 - router_logits: torch.Tensor,
684 - top_k: int,
685 - renormalize: bool,
686 - use_grouped_topk: bool = False,
687 - topk_group: int | None = None,
688 - num_expert_group: int | None = None,
689 - global_num_experts: int = -1,
690 - expert_map: torch.Tensor | None = None,
691 - custom_routing_function: Callable | None = None,
692 - scoring_func: str = "softmax",
693 - routed_scaling_factor: float = 1.0,
694 - e_score_correction_bias: torch.Tensor | None = None,
695 - apply_router_weight_on_input: bool = False,
696 - activation: str = "silu",
697 - enable_eplb: bool = False,
698 - expert_load_view: torch.Tensor | None = None,
699 - logical_to_physical_map: torch.Tensor | None = None,
700 - logical_replica_count: torch.Tensor | None = None,
701 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
702 - if enable_eplb:
703 - assert expert_load_view is not None
704 - assert logical_to_physical_map is not None
705 - assert logical_replica_count is not None
706 - assert isinstance(layer, FusedMoE)
707 -
708 - return self.forward(
709 - x=x,
710 - layer=layer,
711 - router_logits=router_logits,
712 - top_k=top_k,
713 - renormalize=renormalize,
714 - use_grouped_topk=use_grouped_topk,
715 - topk_group=topk_group,
716 - num_expert_group=num_expert_group,
717 - global_num_experts=global_num_experts,
718 - expert_map=expert_map,
719 - custom_routing_function=custom_routing_function,
720 - scoring_func=scoring_func,
721 - routed_scaling_factor=routed_scaling_factor,
722 - e_score_correction_bias=e_score_correction_bias,
723 - activation=activation,
724 - apply_router_weight_on_input=apply_router_weight_on_input,
725 - enable_eplb=enable_eplb,
726 - expert_load_view=expert_load_view,
727 - logical_to_physical_map=logical_to_physical_map,
728 - logical_replica_count=logical_replica_count,
729 - )
730 -
731 - def get_fused_moe_quant_config(
732 - self, layer: torch.nn.Module
733 - ) -> FusedMoEQuantConfig | None:
734 - if self.moe.has_bias:
735 - return biased_moe_quant_config(
736 - layer.w13_bias,
737 - layer.w2_bias,
738 - )
739 - else:
740 - return FUSED_MOE_UNQUANTIZED_CONFIG
741 -
742 - def forward_cuda(
743 - self,
744 - layer: torch.nn.Module,
745 - x: torch.Tensor,
746 - use_grouped_topk: bool,
747 - top_k: int,
748 - router_logits: torch.Tensor,
749 - renormalize: bool,
750 - topk_group: int | None = None,
751 - num_expert_group: int | None = None,
752 - global_num_experts: int = -1,
753 - expert_map: torch.Tensor | None = None,
754 - custom_routing_function: Callable | None = None,
755 - scoring_func: str = "softmax",
756 - routed_scaling_factor: float = 1.0,
757 - e_score_correction_bias: torch.Tensor | None = None,
758 - apply_router_weight_on_input: bool = False,
759 - activation: str = "silu",
760 - enable_eplb: bool = False,
761 - expert_load_view: torch.Tensor | None = None,
762 - logical_to_physical_map: torch.Tensor | None = None,
763 - logical_replica_count: torch.Tensor | None = None,
764 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
765 - zero_expert_num = getattr(layer, "zero_expert_num", 0)
766 - zero_expert_type = getattr(layer, "zero_expert_type", None)
767 -
768 - topk_weights, topk_ids, zero_expert_result = FusedMoE.select_experts(
769 - hidden_states=x,
770 - router_logits=router_logits,
771 - use_grouped_topk=use_grouped_topk,
772 - top_k=top_k,
773 - renormalize=renormalize,
774 - topk_group=topk_group,
775 - num_expert_group=num_expert_group,
776 - custom_routing_function=custom_routing_function,
777 - scoring_func=scoring_func,
778 - routed_scaling_factor=routed_scaling_factor,
779 - e_score_correction_bias=e_score_correction_bias,
780 - indices_type=self.topk_indices_dtype,
781 - enable_eplb=enable_eplb,
782 - expert_map=expert_map,
783 - expert_load_view=expert_load_view,
784 - logical_to_physical_map=logical_to_physical_map,
785 - logical_replica_count=logical_replica_count,
786 - global_num_experts=global_num_experts,
787 - zero_expert_num=zero_expert_num,
788 - zero_expert_type=zero_expert_type,
789 - num_fused_shared_experts=layer.num_fused_shared_experts,
790 - )
791 -
792 - if self.rocm_aiter_moe_enabled:
793 - result = self.rocm_aiter_fused_experts(
794 - hidden_states=x,
795 - w1=layer.w13_weight,
796 - w2=layer.w2_weight,
797 - topk_weights=topk_weights,
798 - topk_ids=topk_ids,
799 - expert_map=expert_map,
800 - activation=activation,
801 - apply_router_weight_on_input=apply_router_weight_on_input,
802 - )
803 - elif self.flashinfer_cutlass_moe_enabled:
804 - return self.flashinfer_cutlass_moe(
805 - hidden_states=x,
806 - w1=layer.w13_weight,
807 - w2=layer.w2_weight,
808 - topk_weights=topk_weights,
809 - topk_ids=topk_ids,
810 - activation=activation,
811 - apply_router_weight_on_input=apply_router_weight_on_input,
812 - )
813 - else:
814 - result = fused_experts(
815 - hidden_states=x,
816 - w1=layer.w13_weight,
817 - w2=layer.w2_weight,
818 - topk_weights=topk_weights,
819 - topk_ids=topk_ids,
820 - inplace=True,
821 - activation=activation,
822 - quant_config=self.moe_quant_config,
823 - apply_router_weight_on_input=apply_router_weight_on_input,
824 - global_num_experts=global_num_experts,
825 - expert_map=expert_map,
826 - )
827 -
828 - if zero_expert_num != 0 and zero_expert_type is not None:
829 - assert not isinstance(result, tuple), (
830 - "Shared + zero experts are mutually exclusive not yet supported"
831 - )
832 - return result, zero_expert_result
833 - else:
834 - return result
835 -
836 - def forward_cpu(
837 - self,
838 - layer: torch.nn.Module,
839 - x: torch.Tensor,
840 - use_grouped_topk: bool,
841 - top_k: int,
842 - router_logits: torch.Tensor,
843 - renormalize: bool,
844 - topk_group: int | None = None,
845 - num_expert_group: int | None = None,
846 - global_num_experts: int = -1,
847 - expert_map: torch.Tensor | None = None,
848 - custom_routing_function: Callable | None = None,
849 - scoring_func: str = "softmax",
850 - routed_scaling_factor: float = 1.0,
851 - e_score_correction_bias: torch.Tensor | None = None,
852 - apply_router_weight_on_input: bool = False,
853 - activation: str = "silu",
854 - enable_eplb: bool = False,
855 - expert_load_view: torch.Tensor | None = None,
856 - logical_to_physical_map: torch.Tensor | None = None,
857 - logical_replica_count: torch.Tensor | None = None,
858 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
859 - if (
860 - enable_eplb is not False
861 - or expert_load_view is not None
862 - or logical_to_physical_map is not None
863 - or logical_replica_count is not None
864 - ):
865 - raise NotImplementedError("Expert load balancing is not supported for CPU.")
866 - return layer.cpu_fused_moe(
867 - layer,
868 - x,
869 - use_grouped_topk,
870 - top_k,
871 - router_logits,
872 - renormalize,
873 - topk_group,
874 - num_expert_group,
875 - global_num_experts,
876 - expert_map,
877 - custom_routing_function,
878 - scoring_func,
879 - routed_scaling_factor,
880 - e_score_correction_bias,
881 - apply_router_weight_on_input,
882 - activation,
883 - )
884 -
885 - def forward_xpu(
886 - self,
887 - layer: torch.nn.Module,
888 - x: torch.Tensor,
889 - use_grouped_topk: bool,
890 - top_k: int,
891 - router_logits: torch.Tensor,
892 - renormalize: bool,
893 - topk_group: int | None = None,
894 - num_expert_group: int | None = None,
895 - global_num_experts: int = -1,
896 - expert_map: torch.Tensor | None = None,
897 - custom_routing_function: Callable | None = None,
898 - scoring_func: str = "softmax",
899 - routed_scaling_factor: float = 1.0,
900 - e_score_correction_bias: torch.Tensor | None = None,
901 - apply_router_weight_on_input: bool = False,
902 - activation: str = "silu",
903 - enable_eplb: bool = False,
904 - expert_load_view: torch.Tensor | None = None,
905 - logical_to_physical_map: torch.Tensor | None = None,
906 - logical_replica_count: torch.Tensor | None = None,
907 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
908 - if (
909 - enable_eplb is not False
910 - or expert_load_view is not None
911 - or logical_to_physical_map is not None
912 - or logical_replica_count is not None
913 - ):
914 - raise NotImplementedError("Expert load balancing is not supported for XPU.")
915 - return layer.ipex_fusion(
916 - x,
917 - use_grouped_topk,
918 - top_k,
919 - router_logits,
920 - renormalize,
921 - topk_group,
922 - num_expert_group,
923 - custom_routing_function=custom_routing_function,
924 - )
925 -
926 - def forward_tpu(
927 - self,
928 - layer: torch.nn.Module,
929 - x: torch.Tensor,
930 - use_grouped_topk: bool,
931 - top_k: int,
932 - router_logits: torch.Tensor,
933 - renormalize: bool,
934 - topk_group: int | None = None,
935 - num_expert_group: int | None = None,
936 - global_num_experts: int = -1,
937 - expert_map: torch.Tensor | None = None,
938 - custom_routing_function: Callable | None = None,
939 - scoring_func: str = "softmax",
940 - routed_scaling_factor: float = 1.0,
941 - e_score_correction_bias: torch.Tensor | None = None,
942 - apply_router_weight_on_input: bool = False,
943 - activation: str = "silu",
944 - enable_eplb: bool = False,
945 - expert_load_view: torch.Tensor | None = None,
946 - logical_to_physical_map: torch.Tensor | None = None,
947 - logical_replica_count: torch.Tensor | None = None,
948 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
949 - assert not use_grouped_topk
950 - assert num_expert_group is None
951 - assert topk_group is None
952 - assert custom_routing_function is None
953 - assert apply_router_weight_on_input is False
954 - if scoring_func != "softmax":
955 - raise NotImplementedError(
956 - "Only softmax scoring function is supported for TPU."
957 - )
958 - if e_score_correction_bias is not None:
959 - raise NotImplementedError(
960 - "Expert score correction bias is not supported for TPU."
961 - )
962 - assert activation == "silu", f"{activation} is not supported for TPU."
963 - assert routed_scaling_factor == 1.0, (
964 - f"routed_scaling_factor {routed_scaling_factor} is not supported for TPU."
965 - )
966 - if (
967 - enable_eplb is not False
968 - or expert_load_view is not None
969 - or logical_to_physical_map is not None
970 - or logical_replica_count is not None
971 - ):
972 - raise NotImplementedError("Expert load balancing is not supported for TPU.")
973 - return fused_moe_pallas(
974 - hidden_states=x,
975 - w1=layer.w13_weight,
976 - w2=layer.w2_weight,
977 - topk=top_k,
978 - gating_output=router_logits,
979 - global_num_experts=global_num_experts,
980 - expert_map=expert_map,
981 - renormalize=renormalize,
982 - )
983 -
984 - if current_platform.is_tpu():
985 - forward_native = forward_tpu
986 - elif current_platform.is_cpu():
987 - forward_native = forward_cpu
988 - elif current_platform.is_xpu():
989 - forward_native = forward_xpu
990 - else:
991 - forward_native = forward_cuda
992 -
993 -
102 def determine_expert_map(
103 ep_size: int,
104 ep_rank: int,
@@ -1125,16 +233,13 @@ def maybe_roundup_hidden_size(
233 Rounded up hidden_size if rounding up is required based on the configs.
234 Original hidden size otherwise.
235 """
236 + from vllm.model_executor.layers.fused_moe.all2all_utils import (
237 + maybe_roundup_layer_hidden_size,
238 + )
239
1129 - if moe_parallel_config.use_deepep_ht_kernels:
1130 - hidden_size = DeepEPHTPrepareAndFinalize.maybe_roundup_layer_hidden_size(
1131 - hidden_size, act_dtype
1132 - )
1133 -
1134 - if moe_parallel_config.use_deepep_ll_kernels:
1135 - hidden_size = DeepEPLLPrepareAndFinalize.maybe_roundup_layer_hidden_size(
1136 - hidden_size
1137 - )
240 + hidden_size = maybe_roundup_layer_hidden_size(
241 + hidden_size, act_dtype, moe_parallel_config
242 + )
243
244 # we are padding globally so EP buffer allocation works
245 if quant_config and quant_config.get_name() == "mxfp4":
@@ -1430,7 +535,6 @@ class FusedMoE(CustomOp):
535 is_lora_enabled=vllm_config.lora_config is not None,
536 )
537
1433 - self.moe_quant_config: FusedMoEQuantConfig | None = None
538 self.quant_config = quant_config
539
540 def _get_quant_method() -> FusedMoEMethodBase:
@@ -1508,9 +612,15 @@ class FusedMoE(CustomOp):
612 # This is called after all weight loading and post-processing, so it
613 # should be safe to swap out the quant_method.
614 def maybe_init_modular_kernel(self) -> None:
1511 - mk = self.quant_method.maybe_init_modular_kernel(self)
1512 - if mk is not None:
1513 - self.quant_method = FusedMoEModularMethod(self.quant_method, mk)
615 + self.ensure_moe_quant_config_init()
616 + prepare_finalize = self.quant_method.maybe_make_prepare_finalize()
617 + if prepare_finalize is not None:
618 + logger.debug(
619 + "%s for %s(%s)", prepare_finalize.__class__.__name__, self, id(self)
620 + )
621 + self.quant_method = FusedMoEModularMethod.make(
622 + self, self.quant_method, prepare_finalize, self.shared_experts
623 + )
624
625 @property
626 def shared_experts(self) -> torch.nn.Module | None:
@@ -2142,12 +1252,16 @@ class FusedMoE(CustomOp):
1252
1253 def ensure_moe_quant_config_init(self):
1254 if self.quant_method.moe_quant_config is None:
1255 + # Note: the moe_quant_config can't be constructed until after
1256 + # weight loading post processing.
1257 self.quant_method.moe_quant_config = (
1258 self.quant_method.get_fused_moe_quant_config(self)
1259 )
1260
2149 - if self.moe_quant_config is None:
2150 - self.moe_quant_config = self.quant_method.moe_quant_config
1261 + @property
1262 + def moe_quant_config(self) -> FusedMoEQuantConfig | None:
1263 + self.ensure_moe_quant_config_init()
1264 + return self.quant_method.moe_quant_config
1265
1266 def ensure_dp_chunking_init(self):
1267 if not self.use_dp_chunking or self.batched_hidden_states is not None:
vllm/model_executor/layers/fused_moe/shared_fused_moe.py
+1 -1
@@ -38,7 +38,7 @@ class SharedFusedMoE(FusedMoE):
38 and not (
39 # TODO(wentao): find the root cause and remove this condition
40 self.enable_eplb
41 - or (self.use_flashinfer_cutlass_kernels and self.dp_size > 1)
41 + or (self.moe_config.use_flashinfer_cutlass_kernels and self.dp_size > 1)
42 )
43 and self._shared_experts is not None
44 )
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py new
+578
@@ -0,0 +1,578 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from collections.abc import Callable
5 +
6 +import torch
7 +import torch.nn.functional as F
8 +
9 +import vllm.envs as envs
10 +from vllm._aiter_ops import rocm_aiter_ops
11 +from vllm.logger import init_logger
12 +from vllm.model_executor.custom_op import CustomOp
13 +from vllm.model_executor.layers.fused_moe.config import (
14 + FUSED_MOE_UNQUANTIZED_CONFIG,
15 + FusedMoEConfig,
16 + FusedMoEQuantConfig,
17 + biased_moe_quant_config,
18 +)
19 +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
20 + FusedMoEMethodBase,
21 +)
22 +from vllm.model_executor.layers.fused_moe.modular_kernel import (
23 + FusedMoEActivationFormat,
24 + FusedMoEPermuteExpertsUnpermute,
25 + FusedMoEPrepareAndFinalize,
26 +)
27 +from vllm.model_executor.utils import set_weight_attrs
28 +from vllm.platforms import current_platform
29 +from vllm.platforms.interface import CpuArchEnum
30 +from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
31 +
32 +if current_platform.is_cuda_alike():
33 + from .fused_batched_moe import BatchedTritonExperts
34 + from .fused_moe import TritonExperts, fused_experts
35 +else:
36 + fused_experts = None # type: ignore
37 +
38 +if current_platform.is_tpu():
39 + from .moe_pallas import fused_moe as fused_moe_pallas
40 +else:
41 + fused_moe_pallas = None # type: ignore
42 +
43 +logger = init_logger(__name__)
44 +
45 +
46 +@CustomOp.register("unquantized_fused_moe")
47 +class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
48 + """MoE method without quantization."""
49 +
50 + def __init__(self, moe: FusedMoEConfig):
51 + super().__init__(moe)
52 +
53 + self.rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
54 + if self.rocm_aiter_moe_enabled:
55 + from .rocm_aiter_fused_moe import rocm_aiter_fused_experts
56 +
57 + self.rocm_aiter_fused_experts = rocm_aiter_fused_experts
58 + else:
59 + self.rocm_aiter_fused_experts = None # type: ignore
60 +
61 + # FlashInfer CUTLASS MoE is only supported on Hopper and later GPUS
62 + self.flashinfer_cutlass_moe_enabled = (
63 + has_flashinfer_cutlass_fused_moe()
64 + and envs.VLLM_USE_FLASHINFER_MOE_FP16
65 + and self.moe.moe_parallel_config.use_ep
66 + and self.moe.moe_parallel_config.dp_size == 1
67 + and current_platform.get_device_capability()[0] >= 9
68 + )
69 + if self.flashinfer_cutlass_moe_enabled:
70 + logger.info_once(
71 + "Enabling FlashInfer CUTLASS MoE for UnquantizedFusedMoEMethod"
72 + )
73 + from functools import partial
74 +
75 + from .flashinfer_cutlass_moe import flashinfer_cutlass_moe
76 +
77 + self.flashinfer_cutlass_moe = partial(
78 + flashinfer_cutlass_moe,
79 + quant_config=FUSED_MOE_UNQUANTIZED_CONFIG,
80 + tp_rank=self.moe.moe_parallel_config.tp_rank,
81 + tp_size=self.moe.moe_parallel_config.tp_size,
82 + ep_rank=self.moe.moe_parallel_config.ep_rank,
83 + ep_size=self.moe.moe_parallel_config.ep_size,
84 + )
85 + else:
86 + if (
87 + self.moe.moe_parallel_config.use_ep
88 + and self.moe.moe_parallel_config.dp_size == 1
89 + ):
90 + logger.info_once(
91 + "FlashInfer CUTLASS MoE is available for EP"
92 + " but not enabled, consider setting"
93 + " VLLM_USE_FLASHINFER_MOE_FP16=1 to enable it.",
94 + scope="local",
95 + )
96 + elif self.moe.moe_parallel_config.dp_size > 1:
97 + logger.info_once(
98 + "FlashInfer CUTLASS MoE is currently not available for DP.",
99 + scope="local",
100 + )
101 + self.flashinfer_cutlass_moe = None # type: ignore
102 +
103 + @property
104 + def supports_eplb(self) -> bool:
105 + return True
106 +
107 + @property
108 + def allow_inplace(self) -> bool:
109 + return True
110 +
111 + def maybe_make_prepare_finalize(self) -> FusedMoEPrepareAndFinalize | None:
112 + if self.rocm_aiter_moe_enabled:
113 + return None
114 + else:
115 + return super().maybe_make_prepare_finalize()
116 +
117 + def select_gemm_impl(
118 + self,
119 + prepare_finalize: FusedMoEPrepareAndFinalize,
120 + layer: torch.nn.Module,
121 + ) -> FusedMoEPermuteExpertsUnpermute:
122 + assert self.moe_quant_config is not None
123 + if (
124 + prepare_finalize.activation_format
125 + == FusedMoEActivationFormat.BatchedExperts
126 + ):
127 + logger.debug("BatchedTritonExperts %s", self.moe)
128 + return BatchedTritonExperts(
129 + max_num_tokens=self.moe.max_num_tokens,
130 + num_dispatchers=prepare_finalize.num_dispatchers(),
131 + quant_config=self.moe_quant_config,
132 + )
133 + else:
134 + logger.debug("TritonExperts %s", self.moe)
135 + return TritonExperts(self.moe_quant_config)
136 +
137 + def create_weights(
138 + self,
139 + layer: torch.nn.Module,
140 + num_experts: int,
141 + hidden_size: int,
142 + intermediate_size_per_partition: int,
143 + params_dtype: torch.dtype,
144 + **extra_weight_attrs,
145 + ):
146 + if self.moe.is_act_and_mul:
147 + w13_up_dim = 2 * intermediate_size_per_partition
148 + else:
149 + w13_up_dim = intermediate_size_per_partition
150 + # Fused gate_up_proj (column parallel)
151 + w13_weight = torch.nn.Parameter(
152 + torch.empty(
153 + num_experts,
154 + w13_up_dim,
155 + hidden_size,
156 + dtype=params_dtype,
157 + ),
158 + requires_grad=False,
159 + )
160 + layer.register_parameter("w13_weight", w13_weight)
161 + set_weight_attrs(w13_weight, extra_weight_attrs)
162 + if self.moe.has_bias:
163 + w13_bias = torch.nn.Parameter(
164 + torch.zeros(num_experts, w13_up_dim, dtype=params_dtype),
165 + requires_grad=False,
166 + )
167 + layer.register_parameter("w13_bias", w13_bias)
168 + set_weight_attrs(w13_bias, extra_weight_attrs)
169 + # down_proj (row parallel)
170 + w2_weight = torch.nn.Parameter(
171 + torch.empty(
172 + num_experts,
173 + hidden_size,
174 + intermediate_size_per_partition,
175 + dtype=params_dtype,
176 + ),
177 + requires_grad=False,
178 + )
179 + layer.register_parameter("w2_weight", w2_weight)
180 + set_weight_attrs(w2_weight, extra_weight_attrs)
181 + if self.moe.has_bias:
182 + w2_bias = torch.nn.Parameter(
183 + torch.zeros(num_experts, hidden_size, dtype=params_dtype),
184 + requires_grad=False,
185 + )
186 + layer.register_parameter("w2_bias", w2_bias)
187 + set_weight_attrs(w2_bias, extra_weight_attrs)
188 +
189 + def _maybe_pad_weight(self, weight: torch.Tensor) -> torch.Tensor:
190 + # Pad the weight tensor. This is an optimization on ROCm platform, which
191 + # can benefit from tensors located far enough from one another in memory
192 + if (
193 + envs.VLLM_ROCM_MOE_PADDING
194 + and current_platform.is_rocm()
195 + and weight.stride(-1) == 1
196 + and (weight.stride(-2) * weight.element_size()) % 512 == 0
197 + ):
198 + num_pad = 256 // weight.element_size()
199 + weight = F.pad(weight, (0, num_pad), "constant", 0)[..., :-num_pad]
200 + torch.cuda.empty_cache()
201 +
202 + return weight
203 +
204 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
205 + super().process_weights_after_loading(layer)
206 +
207 + # Padding the weight for better performance on ROCm
208 + layer.w13_weight.data = self._maybe_pad_weight(layer.w13_weight.data)
209 + layer.w2_weight.data = self._maybe_pad_weight(layer.w2_weight.data)
210 +
211 + if self.rocm_aiter_moe_enabled:
212 + shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights(
213 + layer.w13_weight.data, layer.w2_weight.data
214 + )
215 +
216 + layer.w13_weight.data = shuffled_w13
217 + layer.w2_weight.data = shuffled_w2
218 +
219 + if self.flashinfer_cutlass_moe_enabled:
220 + # Swap halves to arrange as [w3; w1] (kernel expectation)
221 + w1_w, w3_w = torch.chunk(layer.w13_weight.data, 2, dim=1)
222 + w13_weight_swapped = torch.cat([w3_w, w1_w], dim=1)
223 + layer.w13_weight.data = w13_weight_swapped.contiguous()
224 +
225 + if current_platform.is_xpu():
226 + import intel_extension_for_pytorch as ipex
227 +
228 + ep_rank_start = self.moe.ep_rank * self.moe.num_local_experts
229 + layer.ipex_fusion = ipex.llm.modules.GatedMLPMOE(
230 + layer.w13_weight,
231 + layer.w2_weight,
232 + use_prepack=True,
233 + experts_start_id=ep_rank_start,
234 + )
235 + elif current_platform.is_cpu():
236 + from vllm.model_executor.layers.fused_moe import cpu_fused_moe
237 +
238 + if current_platform.get_cpu_architecture() == CpuArchEnum.X86:
239 + from vllm.model_executor.layers.utils import check_cpu_sgl_kernel
240 +
241 + dtype_w13 = layer.w13_weight.dtype
242 + _, n_w13, k_w13 = layer.w13_weight.size()
243 + dtype_w2 = layer.w2_weight.dtype
244 + _, n_w2, k_w2 = layer.w2_weight.size()
245 + if (
246 + envs.VLLM_CPU_SGL_KERNEL
247 + and check_cpu_sgl_kernel(n_w13, k_w13, dtype_w13)
248 + and check_cpu_sgl_kernel(n_w2, k_w2, dtype_w2)
249 + ):
250 + packed_w13_weight = torch.ops._C.convert_weight_packed(
251 + layer.w13_weight
252 + )
253 + assert packed_w13_weight.size() == layer.w13_weight.size()
254 + layer.w13_weight.copy_(packed_w13_weight)
255 + del packed_w13_weight
256 + packed_w2_weight = torch.ops._C.convert_weight_packed(
257 + layer.w2_weight
258 + )
259 + assert packed_w2_weight.size() == layer.w2_weight.size()
260 + layer.w2_weight.copy_(packed_w2_weight)
261 + layer.cpu_fused_moe = cpu_fused_moe.SGLFusedMOE(layer)
262 + else:
263 + layer.cpu_fused_moe = cpu_fused_moe.IPEXFusedMOE(layer)
264 + else:
265 + layer.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer)
266 +
267 + def apply(
268 + self,
269 + layer: torch.nn.Module,
270 + x: torch.Tensor,
271 + router_logits: torch.Tensor,
272 + top_k: int,
273 + renormalize: bool,
274 + use_grouped_topk: bool = False,
275 + topk_group: int | None = None,
276 + num_expert_group: int | None = None,
277 + global_num_experts: int = -1,
278 + expert_map: torch.Tensor | None = None,
279 + custom_routing_function: Callable | None = None,
280 + scoring_func: str = "softmax",
281 + routed_scaling_factor: float = 1.0,
282 + e_score_correction_bias: torch.Tensor | None = None,
283 + apply_router_weight_on_input: bool = False,
284 + activation: str = "silu",
285 + enable_eplb: bool = False,
286 + expert_load_view: torch.Tensor | None = None,
287 + logical_to_physical_map: torch.Tensor | None = None,
288 + logical_replica_count: torch.Tensor | None = None,
289 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
290 + if enable_eplb:
291 + assert expert_load_view is not None
292 + assert logical_to_physical_map is not None
293 + assert logical_replica_count is not None
294 +
295 + return self.forward(
296 + x=x,
297 + layer=layer,
298 + router_logits=router_logits,
299 + top_k=top_k,
300 + renormalize=renormalize,
301 + use_grouped_topk=use_grouped_topk,
302 + topk_group=topk_group,
303 + num_expert_group=num_expert_group,
304 + global_num_experts=global_num_experts,
305 + expert_map=expert_map,
306 + custom_routing_function=custom_routing_function,
307 + scoring_func=scoring_func,
308 + routed_scaling_factor=routed_scaling_factor,
309 + e_score_correction_bias=e_score_correction_bias,
310 + activation=activation,
311 + apply_router_weight_on_input=apply_router_weight_on_input,
312 + enable_eplb=enable_eplb,
313 + expert_load_view=expert_load_view,
314 + logical_to_physical_map=logical_to_physical_map,
315 + logical_replica_count=logical_replica_count,
316 + )
317 +
318 + def get_fused_moe_quant_config(
319 + self, layer: torch.nn.Module
320 + ) -> FusedMoEQuantConfig | None:
321 + if self.moe.has_bias:
322 + return biased_moe_quant_config(
323 + layer.w13_bias,
324 + layer.w2_bias,
325 + )
326 + else:
327 + return FUSED_MOE_UNQUANTIZED_CONFIG
328 +
329 + def forward_cuda(
330 + self,
331 + layer: torch.nn.Module,
332 + x: torch.Tensor,
333 + use_grouped_topk: bool,
334 + top_k: int,
335 + router_logits: torch.Tensor,
336 + renormalize: bool,
337 + topk_group: int | None = None,
338 + num_expert_group: int | None = None,
339 + global_num_experts: int = -1,
340 + expert_map: torch.Tensor | None = None,
341 + custom_routing_function: Callable | None = None,
342 + scoring_func: str = "softmax",
343 + routed_scaling_factor: float = 1.0,
344 + e_score_correction_bias: torch.Tensor | None = None,
345 + apply_router_weight_on_input: bool = False,
346 + activation: str = "silu",
347 + enable_eplb: bool = False,
348 + expert_load_view: torch.Tensor | None = None,
349 + logical_to_physical_map: torch.Tensor | None = None,
350 + logical_replica_count: torch.Tensor | None = None,
351 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
352 + zero_expert_num = getattr(layer, "zero_expert_num", 0)
353 + zero_expert_type = getattr(layer, "zero_expert_type", None)
354 +
355 + topk_weights, topk_ids, zero_expert_result = layer.select_experts(
356 + hidden_states=x,
357 + router_logits=router_logits,
358 + use_grouped_topk=use_grouped_topk,
359 + top_k=top_k,
360 + renormalize=renormalize,
361 + topk_group=topk_group,
362 + num_expert_group=num_expert_group,
363 + custom_routing_function=custom_routing_function,
364 + scoring_func=scoring_func,
365 + routed_scaling_factor=routed_scaling_factor,
366 + e_score_correction_bias=e_score_correction_bias,
367 + indices_type=self.topk_indices_dtype,
368 + enable_eplb=enable_eplb,
369 + expert_map=expert_map,
370 + expert_load_view=expert_load_view,
371 + logical_to_physical_map=logical_to_physical_map,
372 + logical_replica_count=logical_replica_count,
373 + global_num_experts=global_num_experts,
374 + zero_expert_num=zero_expert_num,
375 + zero_expert_type=zero_expert_type,
376 + num_fused_shared_experts=layer.num_fused_shared_experts,
377 + )
378 +
379 + if self.rocm_aiter_moe_enabled:
380 + result = self.rocm_aiter_fused_experts(
381 + hidden_states=x,
382 + w1=layer.w13_weight,
383 + w2=layer.w2_weight,
384 + topk_weights=topk_weights,
385 + topk_ids=topk_ids,
386 + expert_map=expert_map,
387 + activation=activation,
388 + apply_router_weight_on_input=apply_router_weight_on_input,
389 + )
390 + elif self.flashinfer_cutlass_moe_enabled:
391 + return self.flashinfer_cutlass_moe(
392 + hidden_states=x,
393 + w1=layer.w13_weight,
394 + w2=layer.w2_weight,
395 + topk_weights=topk_weights,
396 + topk_ids=topk_ids,
397 + activation=activation,
398 + apply_router_weight_on_input=apply_router_weight_on_input,
399 + )
400 + else:
401 + result = fused_experts(
402 + hidden_states=x,
403 + w1=layer.w13_weight,
404 + w2=layer.w2_weight,
405 + topk_weights=topk_weights,
406 + topk_ids=topk_ids,
407 + inplace=True,
408 + activation=activation,
409 + quant_config=self.moe_quant_config,
410 + apply_router_weight_on_input=apply_router_weight_on_input,
411 + global_num_experts=global_num_experts,
412 + expert_map=expert_map,
413 + )
414 +
415 + if zero_expert_num != 0 and zero_expert_type is not None:
416 + assert not isinstance(result, tuple), (
417 + "Shared + zero experts are mutually exclusive not yet supported"
418 + )
419 + return result, zero_expert_result
420 + else:
421 + return result
422 +
423 + def forward_cpu(
424 + self,
425 + layer: torch.nn.Module,
426 + x: torch.Tensor,
427 + use_grouped_topk: bool,
428 + top_k: int,
429 + router_logits: torch.Tensor,
430 + renormalize: bool,
431 + topk_group: int | None = None,
432 + num_expert_group: int | None = None,
433 + global_num_experts: int = -1,
434 + expert_map: torch.Tensor | None = None,
435 + custom_routing_function: Callable | None = None,
436 + scoring_func: str = "softmax",
437 + routed_scaling_factor: float = 1.0,
438 + e_score_correction_bias: torch.Tensor | None = None,
439 + apply_router_weight_on_input: bool = False,
440 + activation: str = "silu",
441 + enable_eplb: bool = False,
442 + expert_load_view: torch.Tensor | None = None,
443 + logical_to_physical_map: torch.Tensor | None = None,
444 + logical_replica_count: torch.Tensor | None = None,
445 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
446 + if (
447 + enable_eplb is not False
448 + or expert_load_view is not None
449 + or logical_to_physical_map is not None
450 + or logical_replica_count is not None
451 + ):
452 + raise NotImplementedError("Expert load balancing is not supported for CPU.")
453 + return layer.cpu_fused_moe(
454 + layer,
455 + x,
456 + use_grouped_topk,
457 + top_k,
458 + router_logits,
459 + renormalize,
460 + topk_group,
461 + num_expert_group,
462 + global_num_experts,
463 + expert_map,
464 + custom_routing_function,
465 + scoring_func,
466 + routed_scaling_factor,
467 + e_score_correction_bias,
468 + apply_router_weight_on_input,
469 + activation,
470 + )
471 +
472 + def forward_xpu(
473 + self,
474 + layer: torch.nn.Module,
475 + x: torch.Tensor,
476 + use_grouped_topk: bool,
477 + top_k: int,
478 + router_logits: torch.Tensor,
479 + renormalize: bool,
480 + topk_group: int | None = None,
481 + num_expert_group: int | None = None,
482 + global_num_experts: int = -1,
483 + expert_map: torch.Tensor | None = None,
484 + custom_routing_function: Callable | None = None,
485 + scoring_func: str = "softmax",
486 + routed_scaling_factor: float = 1.0,
487 + e_score_correction_bias: torch.Tensor | None = None,
488 + apply_router_weight_on_input: bool = False,
489 + activation: str = "silu",
490 + enable_eplb: bool = False,
491 + expert_load_view: torch.Tensor | None = None,
492 + logical_to_physical_map: torch.Tensor | None = None,
493 + logical_replica_count: torch.Tensor | None = None,
494 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
495 + if (
496 + enable_eplb is not False
497 + or expert_load_view is not None
498 + or logical_to_physical_map is not None
499 + or logical_replica_count is not None
500 + ):
501 + raise NotImplementedError("Expert load balancing is not supported for XPU.")
502 + return layer.ipex_fusion(
503 + x,
504 + use_grouped_topk,
505 + top_k,
506 + router_logits,
507 + renormalize,
508 + topk_group,
509 + num_expert_group,
510 + custom_routing_function=custom_routing_function,
511 + )
512 +
513 + def forward_tpu(
514 + self,
515 + layer: torch.nn.Module,
516 + x: torch.Tensor,
517 + use_grouped_topk: bool,
518 + top_k: int,
519 + router_logits: torch.Tensor,
520 + renormalize: bool,
521 + topk_group: int | None = None,
522 + num_expert_group: int | None = None,
523 + global_num_experts: int = -1,
524 + expert_map: torch.Tensor | None = None,
525 + custom_routing_function: Callable | None = None,
526 + scoring_func: str = "softmax",
527 + routed_scaling_factor: float = 1.0,
528 + e_score_correction_bias: torch.Tensor | None = None,
529 + apply_router_weight_on_input: bool = False,
530 + activation: str = "silu",
531 + enable_eplb: bool = False,
532 + expert_load_view: torch.Tensor | None = None,
533 + logical_to_physical_map: torch.Tensor | None = None,
534 + logical_replica_count: torch.Tensor | None = None,
535 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
536 + assert not use_grouped_topk
537 + assert num_expert_group is None
538 + assert topk_group is None
539 + assert custom_routing_function is None
540 + assert apply_router_weight_on_input is False
541 + if scoring_func != "softmax":
542 + raise NotImplementedError(
543 + "Only softmax scoring function is supported for TPU."
544 + )
545 + if e_score_correction_bias is not None:
546 + raise NotImplementedError(
547 + "Expert score correction bias is not supported for TPU."
548 + )
549 + assert activation == "silu", f"{activation} is not supported for TPU."
550 + assert routed_scaling_factor == 1.0, (
551 + f"routed_scaling_factor {routed_scaling_factor} is not supported for TPU."
552 + )
553 + if (
554 + enable_eplb is not False
555 + or expert_load_view is not None
556 + or logical_to_physical_map is not None
557 + or logical_replica_count is not None
558 + ):
559 + raise NotImplementedError("Expert load balancing is not supported for TPU.")
560 + return fused_moe_pallas(
561 + hidden_states=x,
562 + w1=layer.w13_weight,
563 + w2=layer.w2_weight,
564 + topk=top_k,
565 + gating_output=router_logits,
566 + global_num_experts=global_num_experts,
567 + expert_map=expert_map,
568 + renormalize=renormalize,
569 + )
570 +
571 + if current_platform.is_tpu():
572 + forward_native = forward_tpu
573 + elif current_platform.is_cpu():
574 + forward_native = forward_cpu
575 + elif current_platform.is_xpu():
576 + forward_native = forward_xpu
577 + else:
578 + forward_native = forward_cuda
vllm/model_executor/layers/quantization/mxfp4.py
+6 -23
@@ -741,15 +741,10 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
741 weight_scale=w2_scale, flex_ctx=FlexCtx(rhs_data=w2_flex)
742 )
743
744 - self.w13_weight_triton_tensor = w13_weight
745 - self.w2_weight_triton_tensor = w2_weight
746 -
747 - # need to delete the original weights to save memory on single GPU
748 - del layer.w13_weight
749 - del layer.w2_weight
750 - layer.w13_weight = None
751 - layer.w2_weight = None
752 - torch.cuda.empty_cache()
744 + self.w13_weight = w13_weight
745 + self.w2_weight = w2_weight
746 + layer.w13_weight = w13_weight
747 + layer.w2_weight = w2_weight
748 else:
749 raise ValueError(f"Unsupported backend: {self.mxfp4_backend}")
750
@@ -824,18 +819,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
819 "EP batched experts format"
820 )
821 else:
827 - layer.w13_weight = (
828 - self.w13_weight_triton_tensor
829 - if layer.w13_weight is None
830 - else layer.w13_weight
831 - )
832 - layer.w2_weight = (
833 - self.w2_weight_triton_tensor
834 - if layer.w2_weight is None
835 - else layer.w2_weight
836 - )
837 - assert all([w is not None for w in [layer.w13_weight, layer.w2_weight]])
838 -
822 assert self.moe_quant_config is not None
823 if (
824 self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM
@@ -1070,8 +1053,8 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
1053
1054 return triton_kernel_moe_forward(
1055 hidden_states=x,
1073 - w1=self.w13_weight_triton_tensor,
1074 - w2=self.w2_weight_triton_tensor,
1056 + w1=self.w13_weight,
1057 + w2=self.w2_weight,
1058 gating_output=router_logits,
1059 topk=top_k,
1060 renormalize=renormalize,