lint

Signed-off-by: Barbara Suslova <barbara.suslova@axel-t.com>

csrc/moe/moe_fused_gate.cu

@@ -287,19 +287,18 @@ struct KernelParams {
 
 template <typename T, int VPT, int NUM_EXPERTS, int THREADS_PER_ROW,
           int ROWS_PER_WARP, int ROWS_PER_CTA, int WARPS_PER_CTA>
-__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<VPT, NUM_EXPERTS, THREADS_PER_ROW, ROWS_PER_WARP, ROWS_PER_CTA,
                WARPS_PER_CTA>
       params;
   moe_fused_gate_impl<T>(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<T>(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<at::Tensor> 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<at::Tensor> 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 =

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<at::Tensor> 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<at::Tensor> 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,

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),

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,

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"""
 

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]

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)
 
 

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)

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)

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,

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
 

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"), (

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,

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,

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(

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,

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"

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(

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.")

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).