From 4d9f1ad714e44cd2ceab41a91ec82f9b2f251aa4 Mon Sep 17 00:00:00 2001 From: Barbara Suslova Date: Wed, 12 Nov 2025 14:26:38 +0000 Subject: [PATCH] lint Signed-off-by: Barbara Suslova --- csrc/moe/moe_fused_gate.cu | 28 +++++------ csrc/moe/moe_ops.h | 12 ++--- tests/kernels/moe/test_moe_fused_gate.py | 29 +++++------ vllm/_custom_ops.py | 2 + vllm/config/parallel.py | 4 +- vllm/envs.py | 11 +++-- .../layers/fused_moe/fused_moe.py | 30 ++++++------ .../fused_moe/fused_moe_modular_method.py | 1 + vllm/model_executor/layers/fused_moe/layer.py | 49 +++++++++++-------- .../layers/quantization/awq_marlin.py | 2 +- .../layers/quantization/bitsandbytes.py | 1 + .../compressed_tensors_moe.py | 3 +- .../layers/quantization/experts_int8.py | 2 +- .../layers/quantization/ipex_quant.py | 1 + .../layers/quantization/modelopt.py | 3 +- .../layers/quantization/moe_wna16.py | 2 +- .../layers/quantization/mxfp4.py | 2 + .../layers/quantization/quark/quark_moe.py | 3 +- .../model_executor/layers/quantization/rtn.py | 1 + vllm/model_executor/models/deepseek_v2.py | 18 +++---- 20 files changed, 108 insertions(+), 96 deletions(-) diff --git a/csrc/moe/moe_fused_gate.cu b/csrc/moe/moe_fused_gate.cu index 48c1cf4dd17..951248a7594 100644 --- a/csrc/moe/moe_fused_gate.cu +++ b/csrc/moe/moe_fused_gate.cu @@ -287,19 +287,18 @@ struct KernelParams { template -__global__ void moe_fused_gate_kernel(void* input, void* bias, - float* output_ptr, int32_t* indices_ptr, - int64_t num_rows, int64_t topk_group, - int64_t topk, - int64_t num_fused_shared_experts, - double routed_scaling_factor, - bool apply_routed_scaling_factor_on_output) { +__global__ void moe_fused_gate_kernel( + void* input, void* bias, float* output_ptr, int32_t* indices_ptr, + int64_t num_rows, int64_t topk_group, int64_t topk, + int64_t num_fused_shared_experts, double routed_scaling_factor, + bool apply_routed_scaling_factor_on_output) { KernelParams params; moe_fused_gate_impl(input, bias, output_ptr, indices_ptr, num_rows, topk_group, topk, num_fused_shared_experts, - routed_scaling_factor, apply_routed_scaling_factor_on_output, params); + routed_scaling_factor, + apply_routed_scaling_factor_on_output, params); } // Macro to compute compile-time constants and launch the kernel. @@ -352,18 +351,17 @@ __global__ void moe_fused_gate_kernel_dynamic( moe_fused_gate_impl(input, bias, output_ptr, indices_ptr, num_rows, topk_group, topk, num_fused_shared_experts, - routed_scaling_factor, apply_routed_scaling_factor_on_output, params); + routed_scaling_factor, + apply_routed_scaling_factor_on_output, params); } //------------------------------------------------------------------------------ // Host Launcher Function //------------------------------------------------------------------------------ -std::vector moe_fused_gate(at::Tensor& input, at::Tensor& bias, - int64_t num_expert_group, - int64_t topk_group, int64_t topk, - int64_t num_fused_shared_experts, - double routed_scaling_factor, - bool apply_routed_scaling_factor_on_output) { +std::vector moe_fused_gate( + at::Tensor& input, at::Tensor& bias, int64_t num_expert_group, + int64_t topk_group, int64_t topk, int64_t num_fused_shared_experts, + double routed_scaling_factor, bool apply_routed_scaling_factor_on_output) { int64_t num_rows = input.size(0); int32_t num_experts = input.size(1); auto options = diff --git a/csrc/moe/moe_ops.h b/csrc/moe/moe_ops.h index a29606ccd2f..0238b73f85d 100644 --- a/csrc/moe/moe_ops.h +++ b/csrc/moe/moe_ops.h @@ -28,14 +28,10 @@ void moe_lora_align_block_size( torch::Tensor num_tokens_post_pad, torch::Tensor adapter_enabled, torch::Tensor lora_ids); -std::vector moe_fused_gate(torch::Tensor& input, - torch::Tensor& bias, - int64_t num_expert_group, - int64_t topk_group, int64_t topk, - int64_t num_fused_shared_experts, - double routed_scaling_factor, - bool apply_routed_scaling_factor_on_output - ); +std::vector moe_fused_gate( + torch::Tensor& input, torch::Tensor& bias, int64_t num_expert_group, + int64_t topk_group, int64_t topk, int64_t num_fused_shared_experts, + double routed_scaling_factor, bool apply_routed_scaling_factor_on_output); #ifndef USE_ROCM torch::Tensor moe_wna16_gemm(torch::Tensor input, torch::Tensor output, diff --git a/tests/kernels/moe/test_moe_fused_gate.py b/tests/kernels/moe/test_moe_fused_gate.py index fcf14510661..6c2699d5030 100644 --- a/tests/kernels/moe/test_moe_fused_gate.py +++ b/tests/kernels/moe/test_moe_fused_gate.py @@ -10,14 +10,12 @@ from vllm.model_executor.layers.fused_moe.fused_moe import grouped_topk @pytest.mark.parametrize( "seq_length", - list(range(1, 10)) + - [16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536], + list(range(1, 10)) + + [16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536], ) @pytest.mark.parametrize( "dtype", - [ - torch.float32 - ] # torch.float16, torch.bfloat16 - aren't working correctly yet + [torch.float32], # torch.float16, torch.bfloat16 - aren't working correctly yet ) @pytest.mark.parametrize( "params", @@ -36,7 +34,7 @@ from vllm.model_executor.layers.fused_moe.fused_moe import grouped_topk ) def test_moe_fused_gate_combined( seq_length, dtype, params, num_fused_shared_experts, monkeypatch - ): +): num_experts, num_expert_group, topk_group, topk = params topk += 1 if num_fused_shared_experts > 0 else 0 @@ -82,16 +80,19 @@ def test_moe_fused_gate_combined( shared_indices = original_indices[:, -1] shared_ref_indices = original_ref_indices[:, -1] if shared_indices is not None: - assert torch.all((shared_indices >= valid_min) & ( - shared_indices < valid_max)), ( - "Shared expert indices out of range: ", - f"found values outside [{valid_min}, {valid_max})") + assert torch.all( + (shared_indices >= valid_min) & (shared_indices < valid_max) + ), ( + "Shared expert indices out of range: ", + f"found values outside [{valid_min}, {valid_max})", + ) if shared_ref_indices is not None: assert torch.all( - (shared_ref_indices >= valid_min) - & (shared_ref_indices < valid_max)), ( - "Shared expert reference indices out of range: ", - f"found values outside [{valid_min}, {valid_max})") + (shared_ref_indices >= valid_min) & (shared_ref_indices < valid_max) + ), ( + "Shared expert reference indices out of range: ", + f"found values outside [{valid_min}, {valid_max})", + ) vllm_idx_check = torch.allclose( ref_vllm_indices.sort()[0].to(torch.int32), diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index e34a9516777..69b27db7cbe 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -1911,6 +1911,7 @@ def moe_lora_align_block_size( lora_ids, ) + def moe_fused_gate( input_tensor: torch.Tensor, bias: torch.Tensor, @@ -1956,6 +1957,7 @@ if hasattr(torch.ops._moe_C, "moe_fused_gate"): device=input_tensor.device, ) + def moe_wna16_gemm( input: torch.Tensor, output: torch.Tensor, diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 28b28890a6c..2c9f2373b08 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -280,10 +280,10 @@ class ParallelConfig: This is an internal config that is only valid for and should only be set by API server scale-out. """ - + enable_fused_shared_experts: bool = False """Enable the fusion of the shared experts of the model with other experts.""" - + enable_fused_moe_router: bool = False """Use the fused grouped top-k MoE expert selection router""" diff --git a/vllm/envs.py b/vllm/envs.py index 8f776b4186f..7893972c487 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -1531,12 +1531,13 @@ environment_variables: dict[str, Callable[[], Any]] = { # FlatLogprobs. "VLLM_FLAT_LOGPROBS": lambda: bool(int(os.getenv("VLLM_FLAT_LOGPROBS", "0"))), # Enable the fusion of the shared experts of the model with other experts. - "VLLM_USE_CUDA_FUSION_SHARED_EXPERTS": - lambda: bool(int(os.getenv("VLLM_USE_CUDA_FUSION_SHARED_EXPERTS", "0"))), - + "VLLM_USE_CUDA_FUSION_SHARED_EXPERTS": lambda: bool( + int(os.getenv("VLLM_USE_CUDA_FUSION_SHARED_EXPERTS", "0")) + ), # Use the fused grouped top-k MoE expert selection router - "VLLM_USE_FUSED_MOE_ROUTER": - lambda: bool(int(os.getenv("VLLM_USE_FUSED_MOE_ROUTER", "0"))), + "VLLM_USE_FUSED_MOE_ROUTER": lambda: bool( + int(os.getenv("VLLM_USE_FUSED_MOE_ROUTER", "0")) + ), } # --8<-- [end:env-vars-definition] diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py index 1329c4ee38c..3000486b968 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe.py @@ -1173,7 +1173,7 @@ def grouped_topk( logger.info( "Fused MoE grouped topk is enabled with fused shared experts.", "Only one of these options can be used at a time", - "Fused MoE grouped topk is disabled." + "Fused MoE grouped topk is disabled.", ) use_fused_moe_grouped_topk = False @@ -1239,21 +1239,20 @@ def grouped_topk( # Use original unbiased scores for the routing weights topk_weights = original_scores.gather(1, topk_ids) else: - topk_weights, topk_ids = torch.topk(tmp_scores, - k=topk, - dim=-1, - sorted=use_sorted) + topk_weights, topk_ids = torch.topk( + tmp_scores, k=topk, dim=-1, sorted=use_sorted + ) if num_fused_shared_experts > 0: - assert routed_scaling_factor is not None, \ - "With num_fused_shared_experts>0" + assert routed_scaling_factor is not None, "With num_fused_shared_experts>0" ", routed_scaling_factor need to be provided" - topk_ids[:, -1] = torch.randint(low=num_experts, - high=num_experts + - num_fused_shared_experts, - size=(topk_ids.size(0), ), - dtype=topk_ids.dtype, - device=topk_ids.device) + topk_ids[:, -1] = torch.randint( + low=num_experts, + high=num_experts + num_fused_shared_experts, + size=(topk_ids.size(0),), + dtype=topk_ids.dtype, + device=topk_ids.device, + ) topk_weights[:, -1] = topk_weights[:, :-1].sum(dim=-1) / routed_scaling_factor if renormalize: @@ -1263,9 +1262,8 @@ def grouped_topk( topk_weights_sum = topk_weights[:, :-1].sum(dim=-1, keepdim=True) topk_weights = topk_weights / topk_weights_sum - if num_fused_shared_experts == 0: - if routed_scaling_factor != 1.0: - topk_weights = topk_weights * routed_scaling_factor + if num_fused_shared_experts == 0 and routed_scaling_factor != 1.0: + topk_weights = topk_weights * routed_scaling_factor return topk_weights.to(torch.float32), topk_ids.to(torch.int32) diff --git a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py index 43974ba917e..508ebed2cd5 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py @@ -103,6 +103,7 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp): expert_load_view: torch.Tensor | None = None, logical_to_physical_map: torch.Tensor | None = None, logical_replica_count: torch.Tensor | None = None, + enable_fused_moe_router: bool = False, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: # Is getattr needed? zero_expert_num = getattr(layer, "zero_expert_num", 0) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 230ed1748f6..1f3f302930d 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import math -from abc import abstractmethod from collections.abc import Callable, Iterable from contextlib import nullcontext from enum import Enum @@ -59,8 +58,9 @@ from vllm.utils.torch_utils import ( from vllm.v1.worker.ubatching import dbo_current_ubatch_id if current_platform.is_cuda_alike(): - from .fused_moe import eplb_map_to_physical_and_record, fused_experts from vllm._custom_ops import moe_fused_gate + + from .fused_moe import eplb_map_to_physical_and_record, fused_experts else: fused_experts = None # type: ignore FusedMoEPermuteExpertsUnpermute = object # type: ignore @@ -418,9 +418,10 @@ class FusedMoE(CustomOp): dp_size_=dp_size_, vllm_parallel_config=vllm_config.parallel_config, ) - + self.enable_fused_shared_experts = enable_fused_shared_experts if self.enable_fused_shared_experts: + assert n_shared_experts is not None num_experts += n_shared_experts top_k += n_shared_experts @@ -464,10 +465,11 @@ class FusedMoE(CustomOp): self.num_fused_shared_experts = ( n_shared_experts - if ( - n_shared_experts is not None - and self.aiter_fmoe_shared_expert_enabled - ) or self.enable_fused_shared_experts + if n_shared_experts is not None + and ( + self.aiter_fmoe_shared_expert_enabled + or self.enable_fused_shared_experts + ) else 0 ) if ( @@ -526,12 +528,15 @@ class FusedMoE(CustomOp): self.global_num_experts, get_compressed_expert_map(self.expert_map), ) - if (self.num_fused_shared_experts > 0): + if self.num_fused_shared_experts > 0: logger.warning( "With EP enabled and share expert fusion enabled" ", share expert replica should be same as ep_size" "got share expert replica = %d" - "and ep_size = %d", self.num_fused_shared_experts, self.ep_size) + "and ep_size = %d", + self.num_fused_shared_experts, + self.ep_size, + ) else: self.local_num_experts, self.expert_map, self.expert_mask = ( self.global_num_experts, @@ -1514,13 +1519,12 @@ class FusedMoE(CustomOp): assert topk_group is not None assert num_expert_group is not None if hidden_states.shape[0] == 0: - topk_ids = torch.full((0, top_k), - -1, - dtype=torch.int, - device=hidden_states.device) - topk_weights = torch.empty((0, top_k), - dtype=torch.float32, - device=hidden_states.device) + topk_ids = torch.full( + (0, top_k), -1, dtype=torch.int, device=hidden_states.device + ) + topk_weights = torch.empty( + (0, top_k), dtype=torch.float32, device=hidden_states.device + ) elif rocm_aiter_ops.is_fused_moe_enabled(): if not rocm_aiter_ops.is_fusion_moe_shared_experts_enabled(): assert num_fused_shared_experts == 0 @@ -1528,9 +1532,11 @@ class FusedMoE(CustomOp): else: grouped_topk_impl = grouped_topk - if (enable_fused_moe_router - and e_score_correction_bias is not None - and is_power_of_two(e_score_correction_bias.shape[0])): + if ( + enable_fused_moe_router + and e_score_correction_bias is not None + and is_power_of_two(e_score_correction_bias.shape[0]) + ): # The fused kernel can only work with 128/256 experts topk_weights, topk_ids = moe_fused_gate( input_tensor=router_logits.to(dtype=torch.float32), @@ -1540,7 +1546,8 @@ class FusedMoE(CustomOp): topk=top_k, num_fused_shared_experts=num_fused_shared_experts, routed_scaling_factor=routed_scaling_factor - if routed_scaling_factor is not None else 1.0, + if routed_scaling_factor is not None + else 1.0, apply_routed_scaling_factor_on_output=False, ) else: @@ -1554,7 +1561,7 @@ class FusedMoE(CustomOp): scoring_func=scoring_func, routed_scaling_factor=routed_scaling_factor, e_score_correction_bias=e_score_correction_bias, - num_fused_shared_experts=num_fused_shared_experts + num_fused_shared_experts=num_fused_shared_experts, ) if indices_type is not None: topk_ids = topk_ids.to(dtype=indices_type) diff --git a/vllm/model_executor/layers/quantization/awq_marlin.py b/vllm/model_executor/layers/quantization/awq_marlin.py index 46f2aa38bbc..de20edbedda 100644 --- a/vllm/model_executor/layers/quantization/awq_marlin.py +++ b/vllm/model_executor/layers/quantization/awq_marlin.py @@ -639,7 +639,7 @@ class AWQMoEMethod(FusedMoEMethodBase): indices_type=self.topk_indices_dtype, num_fused_shared_experts=layer.num_fused_shared_experts, enable_fused_moe_router=enable_fused_moe_router, - ) + ) return fused_marlin_moe( x, diff --git a/vllm/model_executor/layers/quantization/bitsandbytes.py b/vllm/model_executor/layers/quantization/bitsandbytes.py index e5a741e639a..74a376ee988 100644 --- a/vllm/model_executor/layers/quantization/bitsandbytes.py +++ b/vllm/model_executor/layers/quantization/bitsandbytes.py @@ -515,6 +515,7 @@ class BitsAndBytesMoEMethod(FusedMoEMethodBase): expert_load_view: torch.Tensor | None = None, logical_to_physical_map: torch.Tensor | None = None, logical_replica_count: torch.Tensor | None = None, + enable_fused_moe_router: bool = False, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: from vllm.model_executor.layers.fused_moe import fused_experts diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py index 26d038cde2c..bfd434ba2df 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -1687,7 +1687,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): indices_type=self.topk_indices_dtype, num_fused_shared_experts=layer.num_fused_shared_experts, enable_fused_moe_router=enable_fused_moe_router, - ) + ) return fused_marlin_moe( x, @@ -2275,6 +2275,7 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod): expert_load_view: torch.Tensor | None = None, logical_to_physical_map: torch.Tensor | None = None, logical_replica_count: torch.Tensor | None = None, + enable_fused_moe_router: bool = False, ) -> torch.Tensor: assert not enable_eplb, "EPLB not supported for W4A8-int MoE yet." assert activation in ("silu", "swigluoai", "swiglu"), ( diff --git a/vllm/model_executor/layers/quantization/experts_int8.py b/vllm/model_executor/layers/quantization/experts_int8.py index 1c30aabbbf4..6592a074910 100644 --- a/vllm/model_executor/layers/quantization/experts_int8.py +++ b/vllm/model_executor/layers/quantization/experts_int8.py @@ -181,7 +181,7 @@ class ExpertsInt8MoEMethod(FusedMoEMethodBase): indices_type=self.topk_indices_dtype, num_fused_shared_experts=layer.num_fused_shared_experts, enable_fused_moe_router=enable_fused_moe_router, - ) + ) return fused_experts( x, diff --git a/vllm/model_executor/layers/quantization/ipex_quant.py b/vllm/model_executor/layers/quantization/ipex_quant.py index 22c4bae041a..27e2fa04573 100644 --- a/vllm/model_executor/layers/quantization/ipex_quant.py +++ b/vllm/model_executor/layers/quantization/ipex_quant.py @@ -454,6 +454,7 @@ class XPUFp8MoEMethod(FusedMoEMethodBase): expert_load_view: torch.Tensor | None = None, logical_to_physical_map: torch.Tensor | None = None, logical_replica_count: torch.Tensor | None = None, + enable_fused_moe_router: bool = False, ) -> torch.Tensor: return layer.ipex_fusion( x, diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 9225feaff0b..d5eb0b40c4c 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -715,6 +715,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase): expert_load_view: torch.Tensor | None = None, logical_to_physical_map: torch.Tensor | None = None, logical_replica_count: torch.Tensor | None = None, + enable_fused_moe_router: bool = False, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: if enable_eplb: raise NotImplementedError( @@ -1694,7 +1695,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase): indices_type=self.topk_indices_dtype, num_fused_shared_experts=layer.num_fused_shared_experts, enable_fused_moe_router=enable_fused_moe_router, - ) + ) if self.use_marlin: return fused_marlin_moe( diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py index 24ecbc4bef5..5dc3eb2f96b 100644 --- a/vllm/model_executor/layers/quantization/moe_wna16.py +++ b/vllm/model_executor/layers/quantization/moe_wna16.py @@ -402,7 +402,7 @@ class MoeWNA16Method(FusedMoEMethodBase): indices_type=self.topk_indices_dtype, num_fused_shared_experts=layer.num_fused_shared_experts, enable_fused_moe_router=enable_fused_moe_router, - ) + ) return fused_experts( x, diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index 66ae2e94c60..75c89788a6d 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -882,6 +882,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): expert_load_view: torch.Tensor | None = None, logical_to_physical_map: torch.Tensor | None = None, logical_replica_count: torch.Tensor | None = None, + enable_fused_moe_router: bool = False, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: if enable_eplb: raise NotImplementedError("EPLB is not supported for mxfp4") @@ -1144,6 +1145,7 @@ class IpexMxfp4MoEMethod(Mxfp4MoEMethod): expert_load_view: torch.Tensor | None = None, logical_to_physical_map: torch.Tensor | None = None, logical_replica_count: torch.Tensor | None = None, + enable_fused_moe_router: bool = False, ) -> torch.Tensor: assert activation == "swigluoai", ( "Only swiglu_oai activation is supported for IPEX MXFP4 MoE" diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index 077a6e8cb48..46e9e63a73c 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -376,7 +376,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod): indices_type=self.topk_indices_dtype, num_fused_shared_experts=layer.num_fused_shared_experts, enable_fused_moe_router=enable_fused_moe_router, - ) + ) if self.rocm_aiter_moe_enabled: from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( @@ -632,6 +632,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): expert_load_view: torch.Tensor | None = None, logical_to_physical_map: torch.Tensor | None = None, logical_replica_count: torch.Tensor | None = None, + enable_fused_moe_router: bool = False, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: if enable_eplb: raise NotImplementedError( diff --git a/vllm/model_executor/layers/quantization/rtn.py b/vllm/model_executor/layers/quantization/rtn.py index 52656263a60..2c9069d8dab 100644 --- a/vllm/model_executor/layers/quantization/rtn.py +++ b/vllm/model_executor/layers/quantization/rtn.py @@ -376,6 +376,7 @@ class RTNMoEMethod(FusedMoEMethodBase): expert_load_view: torch.Tensor | None = None, logical_to_physical_map: torch.Tensor | None = None, logical_replica_count: torch.Tensor | None = None, + enable_fused_moe_router: bool = False, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: if enable_eplb: raise NotImplementedError("EPLB not supported for `RTNMoEMethod` yet.") diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 600fdcf91ae..9d39f14af5a 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -1386,7 +1386,9 @@ class DeepseekV2ForCausalLM( self.config.num_hidden_layers - self.config.first_k_dense_replace ) self.set_moe_parameters() - self.enable_fused_shared_experts = vllm_config.parallel_config.enable_fused_shared_experts + self.enable_fused_shared_experts = ( + vllm_config.parallel_config.enable_fused_shared_experts + ) def set_moe_parameters(self): self.expert_weights = [] @@ -1466,8 +1468,10 @@ class DeepseekV2ForCausalLM( stacked_params_mapping.extend(mla_params_mapping) if self.enable_fused_shared_experts: - logger.info("Cloning %s replicas of the shared expert into MoE", - self.num_shared_experts) + logger.info( + "Cloning %s replicas of the shared expert into MoE", + self.num_shared_experts, + ) expert_params_mapping = SharedFusedMoE.make_expert_params_mapping( ckpt_gate_proj_name="gate_proj", @@ -1494,12 +1498,8 @@ class DeepseekV2ForCausalLM( continue # skip spec decode layers for main model is_fuse_shared_experts_layer = ( - ( - self.enable_fused_shared_experts - or rocm_aiter_moe_shared_expert_enabled - ) - and ("mlp.shared_experts" in name) - ) + self.enable_fused_shared_experts or rocm_aiter_moe_shared_expert_enabled + ) and ("mlp.shared_experts" in name) for param_name, weight_name, shard_id in stacked_params_mapping: # Skip non-stacked layers and experts (experts handled below).