| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from collections.abc import Callable |
| 5 | from typing import Any |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | from vllm._aiter_ops import rocm_aiter_ops |
| 10 | from vllm.config import ParallelConfig, get_current_vllm_config |
| 11 | from vllm.distributed import ( |
| 12 | get_dp_group, |
| 13 | get_pcp_group, |
| 14 | get_tensor_model_parallel_world_size, |
| 15 | ) |
| 16 | from vllm.distributed.eplb.eplb_state import EplbLayerState |
| 17 | from vllm.logger import init_logger |
| 18 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 19 | from vllm.model_executor.layers.fused_moe.config import ( |
| 20 | FusedMoEConfig, |
| 21 | FusedMoEParallelConfig, |
| 22 | ) |
| 23 | from vllm.model_executor.layers.fused_moe.expert_map_manager import ( |
| 24 | ExpertMapManager, |
| 25 | ) |
| 26 | from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts |
| 27 | from vllm.model_executor.layers.fused_moe.router.fused_moe_router import ( |
| 28 | FusedMoERouter, |
| 29 | ) |
| 30 | from vllm.model_executor.layers.fused_moe.router.router_factory import ( |
| 31 | create_fused_moe_router, |
| 32 | ) |
| 33 | from vllm.model_executor.layers.fused_moe.runner.moe_runner import ( |
| 34 | MoERunner, |
| 35 | ) |
| 36 | from vllm.model_executor.layers.quantization.base_config import ( |
| 37 | QuantizationConfig, |
| 38 | ) |
| 39 | |
| 40 | logger = init_logger(__name__) |
| 41 | |
| 42 | |
| 43 | def make_parallel_config( |
| 44 | tp_size: int | None, |
| 45 | dp_size: int | None, |
| 46 | pcp_size: int | None, |
| 47 | is_sequence_parallel: bool, |
| 48 | parallel_config: ParallelConfig, |
| 49 | ) -> FusedMoEParallelConfig: |
| 50 | tp_size_ = ( |
| 51 | tp_size if tp_size is not None else get_tensor_model_parallel_world_size() |
| 52 | ) |
| 53 | dp_size_ = dp_size if dp_size is not None else get_dp_group().world_size |
| 54 | pcp_size_ = pcp_size if pcp_size is not None else get_pcp_group().world_size |
| 55 | sp_size = tp_size_ if is_sequence_parallel else 1 |
| 56 | |
| 57 | moe_parallel_config = FusedMoEParallelConfig.make( |
| 58 | tp_size_=tp_size_, |
| 59 | pcp_size_=pcp_size_, |
| 60 | dp_size_=dp_size_, |
| 61 | sp_size_=sp_size, |
| 62 | vllm_parallel_config=parallel_config, |
| 63 | ) |
| 64 | |
| 65 | assert moe_parallel_config.is_sequence_parallel == is_sequence_parallel |
| 66 | |
| 67 | logger.debug("FusedMoEParallelConfig = %s", str(moe_parallel_config)) |
| 68 | |
| 69 | return moe_parallel_config |
| 70 | |
| 71 | |
| 72 | def determine_expert_counts( |
| 73 | num_experts: int, |
| 74 | num_redundant_experts: int, |
| 75 | n_shared_experts: int | None, |
| 76 | fuse_shared_experts: bool, |
| 77 | ) -> tuple[int, int, int]: |
| 78 | global_num_experts = num_experts + num_redundant_experts |
| 79 | logical_num_experts = num_experts |
| 80 | |
| 81 | num_fused_shared_experts = ( |
| 82 | n_shared_experts if n_shared_experts is not None and fuse_shared_experts else 0 |
| 83 | ) |
| 84 | |
| 85 | return global_num_experts, logical_num_experts, num_fused_shared_experts |
| 86 | |
| 87 | |
| 88 | def FusedMoEFactory( |
| 89 | num_experts: int, # Global number of experts |
| 90 | top_k: int, |
| 91 | hidden_size: int, |
| 92 | intermediate_size: int, |
| 93 | intermediate_pad: int | None = None, |
| 94 | params_dtype: torch.dtype | None = None, |
| 95 | renormalize: bool = True, |
| 96 | use_grouped_topk: bool = False, |
| 97 | num_expert_group: int | None = None, |
| 98 | topk_group: int | None = None, |
| 99 | quant_config: QuantizationConfig | None = None, |
| 100 | tp_size: int | None = None, |
| 101 | dp_size: int | None = None, |
| 102 | pcp_size: int | None = None, |
| 103 | prefix: str = "", |
| 104 | custom_routing_function: Callable | None = None, |
| 105 | router: FusedMoERouter | None = None, |
| 106 | scoring_func: str = "softmax", |
| 107 | routed_scaling_factor: float = 1.0, |
| 108 | swiglu_limit: float | None = None, |
| 109 | swiglu_alpha: float | None = None, |
| 110 | swiglu_beta: float | None = None, |
| 111 | activation_situ_beta: float | None = None, |
| 112 | activation_situ_linear_beta: float | None = None, |
| 113 | e_score_correction_bias: torch.Tensor | None = None, |
| 114 | apply_router_weight_on_input: bool = False, |
| 115 | activation: str = "silu", |
| 116 | enable_eplb: bool = False, |
| 117 | num_redundant_experts: int = 0, |
| 118 | has_bias: bool = False, |
| 119 | is_sequence_parallel: bool = False, |
| 120 | reduce_results: bool = True, |
| 121 | ckpt_names: tuple[str, str, str] = ("gate_proj", "down_proj", "up_proj"), |
| 122 | is_fused_checkpoint_transposed: bool = False, |
| 123 | n_shared_experts: int | None = None, |
| 124 | fuse_shared_experts: bool = False, |
| 125 | router_logits_dtype: torch.dtype | None = None, |
| 126 | gate: torch.nn.Module | None = None, |
| 127 | shared_experts: torch.nn.Module | None = None, |
| 128 | shared_expert_gate: torch.nn.Module | None = None, |
| 129 | routed_input_transform: torch.nn.Module | None = None, |
| 130 | routed_output_transform: torch.nn.Module | None = None, |
| 131 | apply_routed_scale_to_output: bool = False, |
| 132 | zero_expert_type: str | None = None, |
| 133 | hash_indices_table: torch.Tensor | None = None, |
| 134 | bias_vl: torch.Tensor | None = None, |
| 135 | image_sentinel_lo: int = 0, |
| 136 | runner_cls: type[MoERunner] | None = None, |
| 137 | runner_args: dict[str, Any] | None = None, |
| 138 | routed_experts_cls: type[RoutedExperts] | None = None, |
| 139 | routed_experts_args: dict[str, Any] | None = None, |
| 140 | skip_padding: bool = False, |
| 141 | ) -> MoERunner: |
| 142 | """Factory function for creating MoE execution pipeline. |
| 143 | |
| 144 | Creates and configures a complete MoE execution pipeline including: |
| 145 | - Router (for token-to-expert assignment) |
| 146 | - RoutedExperts (containing expert weight parameters) |
| 147 | - MoERunner (orchestrates the complete forward pass) |
| 148 | |
| 149 | The experts contain both MergedColumnParallel weights (gate_up_proj/w13) |
| 150 | and RowParallelLinear weights (down_proj/w2). |
| 151 | |
| 152 | Note: Mixtral uses w1, w2, and w3 for gate, up, and down_proj. We |
| 153 | copy that naming convention here and handle any remapping in the |
| 154 | load_weights function in each model implementation. |
| 155 | |
| 156 | Args: |
| 157 | intermediate_pad: Padding added to the intermediate size, if any. |
| 158 | swiglu_alpha: Optional alpha parameter for the SwiGLU activation. |
| 159 | swiglu_beta: Optional beta parameter for the SwiGLU activation. |
| 160 | |
| 161 | Args: |
| 162 | num_experts: Number of experts in the model (global count) |
| 163 | top_k: Number of experts selected for each token |
| 164 | hidden_size: Input hidden state size of the transformer |
| 165 | intermediate_size: Intermediate size of the experts |
| 166 | params_dtype: Data type for the parameters |
| 167 | renormalize: Whether to renormalize the logits in the router |
| 168 | use_grouped_topk: Whether to use grouped top-k routing |
| 169 | num_expert_group: Number of expert groups for grouped top-k |
| 170 | topk_group: Top-k value per group for grouped top-k |
| 171 | quant_config: Quantization configuration |
| 172 | tp_size: Tensor parallelism size (None = use global default) |
| 173 | dp_size: Data parallelism size (None = use global default) |
| 174 | pcp_size: Pipeline context parallelism size (None = use global default) |
| 175 | prefix: Layer name prefix for weight loading |
| 176 | custom_routing_function: Custom routing function override |
| 177 | router: Pre-configured router instance (None = create default) |
| 178 | scoring_func: Scoring function for routing ("softmax" or others) |
| 179 | routed_scaling_factor: Scaling factor applied to topk_weights or output |
| 180 | swiglu_limit: SwiGLU activation limit |
| 181 | activation_situ_beta: SituGLU activation beta |
| 182 | activation_situ_linear_beta: SituGLU linear beta |
| 183 | e_score_correction_bias: Expert score correction bias tensor |
| 184 | apply_router_weight_on_input: Whether to apply router weights on input |
| 185 | activation: Activation function name ("silu", "gelu", etc.) |
| 186 | enable_eplb: Whether to enable expert parallelism load balancer |
| 187 | num_redundant_experts: Number of redundant experts for EPLB |
| 188 | has_bias: Whether expert layers have bias terms |
| 189 | is_sequence_parallel: Whether sequence parallelism is enabled |
| 190 | reduce_results: Whether to all-reduce the final output. Setting this |
| 191 | to False (to fuse the all-reduce downstream) is only honored on |
| 192 | the late-AR path. |
| 193 | ckpt_names: Checkpoint parameter name tuple (gate_proj, down_proj, |
| 194 | up_proj) used for weight loading |
| 195 | is_fused_checkpoint_transposed: Whether fused checkpoint weights and |
| 196 | block scales use transposed storage. |
| 197 | n_shared_experts: Number of shared experts to fuse into the routed |
| 198 | grouped GEMM (ROCm; requires aiter FSE or the router-append path) |
| 199 | fuse_shared_experts: Whether to enable shared-expert fusion. |
| 200 | router_logits_dtype: Data type for router logits buffers |
| 201 | gate: Pre-configured gate module |
| 202 | shared_experts: Pre-configured shared experts module |
| 203 | shared_expert_gate: Pre-configured shared expert gate module |
| 204 | routed_input_transform: Input transformation module |
| 205 | routed_output_transform: Output transformation module |
| 206 | apply_routed_scale_to_output: Whether to apply routed_scaling_factor to |
| 207 | output instead of topk_weights |
| 208 | zero_expert_type: Type of zero expert handling |
| 209 | hash_indices_table: Hash table for expert indices |
| 210 | bias_vl: Vision routing bias for image tokens (Deepseek V4) |
| 211 | image_sentinel_lo: First of five consecutive in-vocab image sentinel |
| 212 | ids (0 = vision routing disabled) |
| 213 | runner_cls: Custom MoERunner class (None = use default MoERunner) |
| 214 | runner_args: Additional arguments for runner constructor |
| 215 | routed_experts_cls: Custom RoutedExperts class (None = use default) |
| 216 | routed_experts_args: Additional arguments for routed_experts constructor |
| 217 | skip_padding: Whether grouped routing should invalidate padding rows. |
| 218 | |
| 219 | Returns: |
| 220 | MoERunner: Configured MoE execution pipeline ready for forward passes |
| 221 | |
| 222 | """ |
| 223 | vllm_config = get_current_vllm_config() |
| 224 | |
| 225 | layer_name = prefix |
| 226 | |
| 227 | moe_activation = MoEActivation.from_str(activation) |
| 228 | is_act_and_mul = moe_activation.is_gated |
| 229 | |
| 230 | moe_parallel_config = make_parallel_config( |
| 231 | tp_size=tp_size, |
| 232 | dp_size=dp_size, |
| 233 | pcp_size=pcp_size, |
| 234 | is_sequence_parallel=is_sequence_parallel, |
| 235 | parallel_config=vllm_config.parallel_config, |
| 236 | ) |
| 237 | |
| 238 | # Resolve the deferred all-reduce request against the parallel config. |
| 239 | skip_final_all_reduce = ( |
| 240 | not reduce_results |
| 241 | and not moe_parallel_config.use_all2all_kernels |
| 242 | and not moe_parallel_config.is_sequence_parallel |
| 243 | and zero_expert_type is None |
| 244 | ) |
| 245 | |
| 246 | global_num_experts, logical_num_experts, num_fused_shared_experts = ( |
| 247 | determine_expert_counts( |
| 248 | num_experts, |
| 249 | num_redundant_experts, |
| 250 | n_shared_experts, |
| 251 | fuse_shared_experts, |
| 252 | ) |
| 253 | ) |
| 254 | |
| 255 | # Initialize EPLB manager (or None?) |
| 256 | eplb_state: EplbLayerState | None = None |
| 257 | if enable_eplb: |
| 258 | use_ep = moe_parallel_config.use_ep |
| 259 | ep_size = moe_parallel_config.ep_size |
| 260 | if use_ep and global_num_experts % ep_size != 0: |
| 261 | raise ValueError( |
| 262 | f"EPLB currently only supports even distribution of " |
| 263 | f"experts across ranks. Got {global_num_experts} experts " |
| 264 | f"and {ep_size} EP ranks." |
| 265 | ) |
| 266 | eplb_state = EplbLayerState() |
| 267 | else: |
| 268 | assert num_redundant_experts == 0, ( |
| 269 | "Redundant experts are only supported with EPLB." |
| 270 | ) |
| 271 | |
| 272 | max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens |
| 273 | |
| 274 | # Create ExpertMapManager to handle expert mapping and placement for EP. |
| 275 | # See ExpertMapManager for a detailed description of what it does and when |
| 276 | # it is required. |
| 277 | expert_map_manager = ExpertMapManager( |
| 278 | max_num_batched_tokens=max_num_batched_tokens, |
| 279 | top_k=top_k, |
| 280 | global_num_experts=global_num_experts, |
| 281 | num_redundant_experts=num_redundant_experts, |
| 282 | num_expert_group=num_expert_group, |
| 283 | moe_parallel_config=moe_parallel_config, |
| 284 | placement_strategy=vllm_config.parallel_config.expert_placement_strategy, |
| 285 | enable_eplb=eplb_state is not None, |
| 286 | num_fused_shared_experts=num_fused_shared_experts, |
| 287 | rocm_aiter_enabled=rocm_aiter_ops.is_fused_moe_enabled() and is_act_and_mul, |
| 288 | ) |
| 289 | |
| 290 | # TODO(bnell): we should not have to create a router if the kernel is |
| 291 | # monolithic. |
| 292 | if router is None: |
| 293 | router = create_fused_moe_router( |
| 294 | top_k=top_k, |
| 295 | global_num_experts=global_num_experts, |
| 296 | eplb_state=eplb_state, |
| 297 | renormalize=renormalize, |
| 298 | use_grouped_topk=use_grouped_topk, |
| 299 | num_expert_group=num_expert_group, |
| 300 | topk_group=topk_group, |
| 301 | custom_routing_function=custom_routing_function, |
| 302 | scoring_func=scoring_func, |
| 303 | # When apply_routed_scale_to_output is True, we set the scaling factor |
| 304 | # to 1.0 so it ends up being a nop. Applying the scale will be handled |
| 305 | # by the runner in this case. |
| 306 | # The member variable must be set in the same way as the router since |
| 307 | # some quantization methods can access it. |
| 308 | routed_scaling_factor=routed_scaling_factor |
| 309 | if not apply_routed_scale_to_output |
| 310 | else 1.0, |
| 311 | e_score_correction_bias=e_score_correction_bias, |
| 312 | num_fused_shared_experts=num_fused_shared_experts, |
| 313 | skip_padding=(skip_padding and moe_parallel_config.use_deepep_v2_kernels), |
| 314 | # Fused shared-expert slot weight. With apply_routed_scale_to_output |
| 315 | # the runner scales the combined output by routed_scaling_factor, so |
| 316 | # the shared slot weight must be 1/routed_scaling_factor for its net |
| 317 | # contribution to be 1.0 (matching the un-scaled separate-MLP add). |
| 318 | shared_expert_weight=( |
| 319 | (1.0 / routed_scaling_factor) |
| 320 | if ( |
| 321 | apply_routed_scale_to_output |
| 322 | and num_fused_shared_experts > 0 |
| 323 | and routed_scaling_factor |
| 324 | ) |
| 325 | else 1.0 |
| 326 | ), |
| 327 | zero_expert_type=zero_expert_type, |
| 328 | num_logical_experts=logical_num_experts, |
| 329 | hash_indices_table=hash_indices_table, |
| 330 | bias_vl=bias_vl, |
| 331 | image_sentinel_lo=image_sentinel_lo, |
| 332 | ) |
| 333 | |
| 334 | if params_dtype is None: |
| 335 | params_dtype = torch.get_default_dtype() |
| 336 | |
| 337 | # FIXME (varun): We should have a better way of inferring the activation |
| 338 | # datatype. This works for now as the tensor datatype entering the MoE |
| 339 | # operation is typically unquantized (i.e. float16/bfloat16). |
| 340 | if vllm_config.model_config is not None: |
| 341 | moe_in_dtype = vllm_config.model_config.dtype |
| 342 | else: |
| 343 | # TODO (bnell): This is a hack to get test_mixtral_moe to work |
| 344 | # since model_config is not set in the pytest test. |
| 345 | moe_in_dtype = params_dtype |
| 346 | |
| 347 | moe_config = FusedMoEConfig( |
| 348 | num_experts=global_num_experts, |
| 349 | experts_per_token=top_k, |
| 350 | hidden_dim=hidden_size, |
| 351 | intermediate_size=intermediate_size, |
| 352 | intermediate_pad=intermediate_pad, |
| 353 | num_local_experts=expert_map_manager.local_num_experts, |
| 354 | num_logical_experts=logical_num_experts, |
| 355 | moe_parallel_config=moe_parallel_config, |
| 356 | in_dtype=moe_in_dtype, |
| 357 | moe_backend=vllm_config.kernel_config.moe_backend, |
| 358 | router_logits_dtype=router_logits_dtype, |
| 359 | max_num_tokens=max_num_batched_tokens, |
| 360 | elastic_ep_max_dp_size=vllm_config.parallel_config.elastic_ep_max_dp_size, |
| 361 | has_bias=has_bias, |
| 362 | is_lora_enabled=vllm_config.lora_config is not None, |
| 363 | activation=moe_activation, |
| 364 | device=vllm_config.device_config.device, |
| 365 | routing_method=router.routing_method_type, # Not ideal |
| 366 | has_hash_routing=hash_indices_table is not None, |
| 367 | swiglu_limit=swiglu_limit, |
| 368 | swiglu_alpha=swiglu_alpha, |
| 369 | swiglu_beta=swiglu_beta, |
| 370 | activation_situ_beta=activation_situ_beta, |
| 371 | activation_situ_linear_beta=activation_situ_linear_beta, |
| 372 | max_capture_size=vllm_config.compilation_config.max_cudagraph_capture_size, |
| 373 | skip_final_all_reduce=skip_final_all_reduce, |
| 374 | ) |
| 375 | |
| 376 | logger.debug("FusedMoEConfig = %s", moe_config) |
| 377 | |
| 378 | # Create RoutedExperts instance BEFORE create_weights() |
| 379 | # This will hold all expert weight parameters |
| 380 | if routed_experts_cls is None: |
| 381 | routed_experts_cls = RoutedExperts |
| 382 | |
| 383 | assert params_dtype is not None |
| 384 | routed_experts = routed_experts_cls( |
| 385 | layer_name, |
| 386 | params_dtype, |
| 387 | moe_config, |
| 388 | quant_config, |
| 389 | expert_map_manager=expert_map_manager, |
| 390 | ckpt_gate_proj_name=ckpt_names[0], |
| 391 | ckpt_down_proj_name=ckpt_names[1], |
| 392 | ckpt_up_proj_name=ckpt_names[2], |
| 393 | is_fused_checkpoint_transposed=is_fused_checkpoint_transposed, |
| 394 | # Extra params that are needed by quant_methods, pass along for now |
| 395 | # Prefer getting these from other sources, e.g. moe_config or |
| 396 | # router object |
| 397 | renormalize=renormalize, |
| 398 | use_grouped_topk=use_grouped_topk, |
| 399 | num_expert_group=num_expert_group, |
| 400 | topk_group=topk_group, |
| 401 | custom_routing_function=custom_routing_function, |
| 402 | scoring_func=scoring_func, |
| 403 | routed_scaling_factor=routed_scaling_factor |
| 404 | if not apply_routed_scale_to_output |
| 405 | else 1.0, |
| 406 | swiglu_limit=swiglu_limit, |
| 407 | swiglu_alpha=swiglu_alpha, |
| 408 | swiglu_beta=swiglu_beta, |
| 409 | # TODO get from router? needs to be truncated? |
| 410 | e_score_correction_bias=e_score_correction_bias, |
| 411 | apply_router_weight_on_input=apply_router_weight_on_input, |
| 412 | **routed_experts_args if routed_experts_args is not None else {}, |
| 413 | ) |
| 414 | |
| 415 | if runner_cls is None: |
| 416 | runner_cls = MoERunner |
| 417 | |
| 418 | runner = runner_cls( |
| 419 | layer_name=layer_name, |
| 420 | moe_config=moe_config, |
| 421 | router=router, |
| 422 | routed_experts=routed_experts, |
| 423 | enable_dbo=vllm_config.parallel_config.enable_dbo, |
| 424 | gate=gate, |
| 425 | shared_expert_gate=shared_expert_gate, |
| 426 | shared_experts=shared_experts, |
| 427 | routed_input_transform=routed_input_transform, |
| 428 | routed_output_transform=routed_output_transform, |
| 429 | # When apply_routed_scale_to_output is True, we allow |
| 430 | # the scaling factor to be passed to the runner, otherwise |
| 431 | # we pass 1.0 so it ends up being a nop. |
| 432 | routed_scaling_factor=routed_scaling_factor |
| 433 | if apply_routed_scale_to_output |
| 434 | else 1.0, |
| 435 | **runner_args if runner_args is not None else {}, |
| 436 | ) |
| 437 | |
| 438 | return runner |
| 439 | |
| 440 | |
| 441 | def fused_moe_make_expert_params_mapping( |
| 442 | model: torch.nn.Module, |
| 443 | ckpt_gate_proj_name: str, |
| 444 | ckpt_down_proj_name: str, |
| 445 | ckpt_up_proj_name: str, |
| 446 | num_experts: int, |
| 447 | num_redundant_experts: int = 0, |
| 448 | routed_experts_prefix: str = "routed_experts", |
| 449 | ) -> list[tuple[str, str, int, str]]: |
| 450 | """Delegate to EPLB manager.""" |
| 451 | return RoutedExperts.make_expert_params_mapping( |
| 452 | model, |
| 453 | ckpt_gate_proj_name, |
| 454 | ckpt_down_proj_name, |
| 455 | ckpt_up_proj_name, |
| 456 | num_experts, |
| 457 | num_redundant_experts, |
| 458 | routed_experts_prefix, |
| 459 | ) |