| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from inspect import signature |
| 5 | from typing import TYPE_CHECKING, Any |
| 6 | |
| 7 | import vllm.model_executor.layers.fused_moe.modular_kernel as mk |
| 8 | from vllm.model_executor.layers.fused_moe.all2all_utils import ( |
| 9 | maybe_make_prepare_finalize, |
| 10 | ) |
| 11 | from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig |
| 12 | from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( |
| 13 | FusedMoEMethodBase, |
| 14 | ) |
| 15 | from vllm.model_executor.layers.fused_moe.fused_moe_modular_method import ( |
| 16 | FusedMoEModularMethod, |
| 17 | ) |
| 18 | from vllm.model_executor.layers.fused_moe.modular_kernel import ( |
| 19 | FusedMoEExpertsModular, |
| 20 | FusedMoEPrepareAndFinalizeModular, |
| 21 | ) |
| 22 | |
| 23 | if TYPE_CHECKING: |
| 24 | from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner |
| 25 | |
| 26 | |
| 27 | def _make_eep_experts( |
| 28 | quant_method: FusedMoEMethodBase, |
| 29 | source_experts: FusedMoEExpertsModular, |
| 30 | prepare_finalize: FusedMoEPrepareAndFinalizeModular, |
| 31 | moe_config: FusedMoEConfig, |
| 32 | ) -> FusedMoEExpertsModular: |
| 33 | experts_cls = source_experts.__class__ |
| 34 | assert quant_method.moe_quant_config is not None |
| 35 | experts_kwargs: dict[str, Any] = { |
| 36 | "moe_config": moe_config, |
| 37 | "quant_config": quant_method.moe_quant_config, |
| 38 | } |
| 39 | if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts: |
| 40 | max_num_tokens = prepare_finalize.max_num_tokens_per_rank() |
| 41 | assert max_num_tokens is not None |
| 42 | experts_kwargs.update( |
| 43 | max_num_tokens=max_num_tokens, |
| 44 | num_dispatchers=prepare_finalize.num_dispatchers(), |
| 45 | ) |
| 46 | |
| 47 | # Expert kernels with extra init params need explicit EEP support. |
| 48 | generic_arg_names = set(signature(mk.FusedMoEExperts.__init__).parameters) |
| 49 | ctor_arg_names = set(signature(experts_cls.__init__).parameters) |
| 50 | unsupported_args = ctor_arg_names - generic_arg_names |
| 51 | missing_args = set(experts_kwargs) - ctor_arg_names |
| 52 | if unsupported_args or missing_args: |
| 53 | raise NotImplementedError( |
| 54 | f"{experts_cls.__name__} experts do not support Elastic EP." |
| 55 | ) |
| 56 | |
| 57 | return experts_cls(**experts_kwargs) |
| 58 | |
| 59 | |
| 60 | def make_eep_staged_quant_method( |
| 61 | module: "MoERunner", |
| 62 | moe_config: FusedMoEConfig, |
| 63 | *, |
| 64 | all2all_manager: Any, |
| 65 | ) -> FusedMoEMethodBase | None: |
| 66 | quant_method = module._quant_method |
| 67 | if not quant_method.supports_internal_mk: |
| 68 | return None |
| 69 | if getattr(quant_method, "wraps_legacy_quant_method", False): |
| 70 | return None |
| 71 | |
| 72 | old_batched_format = ( |
| 73 | module.moe_config.moe_parallel_config.use_batched_activation_format |
| 74 | ) |
| 75 | new_batched_format = moe_config.moe_parallel_config.use_batched_activation_format |
| 76 | assert old_batched_format == new_batched_format |
| 77 | |
| 78 | moe_kernel = quant_method.moe_kernel |
| 79 | if moe_kernel is None: |
| 80 | return None |
| 81 | if moe_kernel.is_monolithic: |
| 82 | raise NotImplementedError( |
| 83 | "Elastic EP full modular-kernel staging is not supported for " |
| 84 | "monolithic fused MoE kernels." |
| 85 | ) |
| 86 | if quant_method.moe_quant_config is None: |
| 87 | raise ValueError( |
| 88 | "Elastic EP full modular-kernel staging requires initialized " |
| 89 | "MoE quant config." |
| 90 | ) |
| 91 | |
| 92 | prepare_finalize = maybe_make_prepare_finalize( |
| 93 | moe_config, |
| 94 | quant_method.moe_quant_config, |
| 95 | routing_tables=None, |
| 96 | allow_new_interface=True, |
| 97 | use_monolithic=quant_method.is_monolithic, |
| 98 | all2all_manager=all2all_manager, |
| 99 | ) |
| 100 | assert prepare_finalize is not None |
| 101 | assert isinstance(prepare_finalize, FusedMoEPrepareAndFinalizeModular) |
| 102 | |
| 103 | source_experts = moe_kernel.fused_experts |
| 104 | assert isinstance(source_experts, FusedMoEExpertsModular) |
| 105 | |
| 106 | experts = _make_eep_experts( |
| 107 | quant_method, |
| 108 | source_experts, |
| 109 | prepare_finalize, |
| 110 | moe_config, |
| 111 | ) |
| 112 | |
| 113 | if isinstance(quant_method, FusedMoEModularMethod): |
| 114 | base_quant_method = quant_method.old_quant_method |
| 115 | else: |
| 116 | base_quant_method = quant_method |
| 117 | |
| 118 | return FusedMoEModularMethod( |
| 119 | base_quant_method, |
| 120 | mk.FusedMoEKernel( |
| 121 | prepare_finalize, |
| 122 | experts, |
| 123 | ), |
| 124 | ) |