applied pre-commit results and Read the Docs build results

Signed-off-by: HakJu Kim <[email protected]>

Author: whitememory

vllm/distributed/device_communicators/all2all.py

@@ -9,7 +9,7 @@
 from vllm.distributed import get_dp_group, get_ep_group
 from vllm.forward_context import get_forward_context
 from vllm.logger import init_logger
-from vllm.utils import has_deep_ep, has_pplx, has_mori
+from vllm.utils import has_deep_ep, has_mori, has_pplx
 from vllm.utils.flashinfer import has_flashinfer_all2all
 
 from .base_device_communicator import All2AllManagerBase, Cache
@@ -439,24 +439,26 @@ def cleanup(self):
                 self.mapping = None
                 self.initialized = False
 
+
 class MoriAll2AllManager(All2AllManagerBase):
     """
     All2All communication based on mori kernels.
     """
+
     def __init__(self, cpu_group):
         assert has_mori(
-        ), "mori not found. Please follow https://github.com/ROCm/mori/blob/main/README.md#installation to install mori."  # noqa
+        ), "Please install mori from https://github.com/ROCm/mori."
 
         super().__init__(cpu_group)
         self.handle_cache = Cache()
         self.config = None
         self._op_handles = {}  # Cache for EpDispatchCombineOp instances
         self._shmem_initialized = False
         # Delay mori shmem initialization until first use
-        logger.debug(f"[rank {self.rank}] MoriAll2AllManager created, shmem will be initialized lazily")
+        logger.debug("[rank %s] MoriAll2AllManager created", self.rank)
 
     def _ensure_shmem_initialized(self):
-        """Ensure mori's shared memory system is initialized (lazy initialization)"""
+        """Initialize mori's shared memory system lazily"""
         if self._shmem_initialized:
             return
 
@@ -473,45 +475,60 @@ def _ensure_shmem_initialized(self):
             if backend is None:
                 raise RuntimeError("No valid distributed backend found")
 
-            logger.debug(f"[rank {self.rank}] PyTorch distributed ready with backend: {backend}")
+            logger.debug(
+                "[rank %s] PyTorch distributed ready with backend: %s",
+                self.rank, backend)
 
-            current_group = self.cpu_group if self.cpu_group is not None else dist.group.WORLD
+            current_group = (self.cpu_group if self.cpu_group is not None else
+                             dist.group.WORLD)
 
             # TODO(inhyeok): make group_name more reasonable
             group_name = "default"
             try:
+                import contextlib
+
                 import torch._C._distributed_c10d as c10d
 
                 # Try to unregister first in case it exists
-                try:
+                with contextlib.suppress(RuntimeError):
                     c10d._unregister_process_group(group_name)
-                except:
-                    pass
 
                 # Register the current process group
                 c10d._register_process_group(group_name, current_group)
-                logger.debug(f"[rank {self.rank}] Registered process group '{group_name}'")
+                logger.debug("[rank %s] Registered process group '%s'",
+                             self.rank, group_name)
 
                 # Initialize mori shmem with the registered group
                 mori.shmem.shmem_torch_process_group_init(group_name)
-                logger.debug(f"[rank {self.rank}] Torch process group shmem initialization successful")
+                logger.debug(
+                    "[rank %s] torch process group shmem init success",
+                    self.rank)
                 self._shmem_initialized = True
                 return
 
             except Exception as torch_error:
-                logger.debug(f"[rank {self.rank}] Torch process group shmem init failed: {torch_error}")
+                logger.debug(
+                    "[rank %s] torch process group shmem init failed: %s",
+                    self.rank, torch_error)
 
             self._shmem_initialized = True
 
         except Exception as e:
-            logger.error(f"[rank {self.rank}] mori shmem initialization failed: {e}")
+            logger.error("[rank %s] mori shmem initialization failed: %s",
+                         self.rank, e)
             # Don't fail completely - mark as initialized to avoid retry loops
             self._shmem_initialized = True
-            logger.warning(f"[rank {self.rank}] Continuing without mori shmem optimization")
-
-    def _make_mori_config(self, max_num_tokens: int, num_local_experts: int,
-                          experts_per_token: int, hidden_dim: int,
-                          scale_dim: int, scale_type_size: int,
+            logger.warning(
+                "[rank %s] Continuing without mori shmem optimization",
+                self.rank)
+
+    def _make_mori_config(self,
+                          max_num_tokens: int,
+                          num_local_experts: int,
+                          experts_per_token: int,
+                          hidden_dim: int,
+                          scale_dim: int,
+                          scale_type_size: int,
                           data_type: torch.dtype = torch.bfloat16,
                           quant_dtype: torch.dtype = None):
         """Create mori EpDispatchCombineConfig"""
@@ -546,9 +563,8 @@ def _make_mori_config(self, max_num_tokens: int, num_local_experts: int,
 
             # Determine kernel type based on topology
             kernel_type=(EpDispatchCombineKernelType.InterNode
-                        if self.internode
-                        else EpDispatchCombineKernelType.IntraNode)
-        )
+                         if self.internode else
+                         EpDispatchCombineKernelType.IntraNode))
 
         return config
 
@@ -578,13 +594,16 @@ def get_handle(self, kwargs):
         scale_type_size = kwargs.get('scale_type_size')
 
         # Validate required parameters
-        if any(param is None for param in [max_num_tokens, num_local_experts,
-                                          experts_per_token, hidden_dim]):
-            raise ValueError("Missing required parameters for mori handle creation")
+        if any(
+                param is None for param in
+            [max_num_tokens, num_local_experts, experts_per_token, hidden_dim
+             ]):
+            raise ValueError(
+                "Missing required parameters for mori handle creation")
 
         # Create cache key
         cache_key = (max_num_tokens, num_local_experts, experts_per_token,
-                    hidden_dim, data_type)
+                     hidden_dim, data_type)
 
         # Check cache first
         if cache_key in self._op_handles:
@@ -607,17 +626,22 @@ def get_handle(self, kwargs):
         # Cache the handle
         self._op_handles[cache_key] = op
 
-        logger.debug(f"[rank {self.dp_rank}] Created mori handle with config: "
-                    f"tokens={max_num_tokens}, experts={num_local_experts}, "
-                    f"topk={experts_per_token}, hidden={hidden_dim}")
+        logger.debug(
+            "[rank %s] Created mori handle with config: tokens=%d, experts=%d,"
+            " topk=%d, hidden_dim=%d", self.dp_rank, max_num_tokens,
+            num_local_experts, experts_per_token, hidden_dim)
 
         return op
 
-    def dispatch(self, hidden_states: torch.Tensor,
-                 router_logits: torch.Tensor):
+    def dispatch(self,
+                 hidden_states: torch.Tensor,
+                 router_logits: torch.Tensor,
+                 is_sequence_parallel: bool = False):
         raise NotImplementedError
 
-    def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    def combine(self,
+                hidden_states: torch.Tensor,
+                is_sequence_parallel: bool = False):
         raise NotImplementedError
 
     def destroy(self):
@@ -626,17 +650,23 @@ def destroy(self):
             # Clear operation handle cache
             self._op_handles.clear()
 
-            # Try to finalize mori shared memory if it was successfully initialized
+            # finalize mori shared memory if it was initialized
             if self._shmem_initialized:
                 try:
                     import mori.shmem
+
                     # Check if shmem is actually active before finalizing
                     mori.shmem.shmem_finalize()
-                    logger.debug(f"[rank {self.dp_rank}] mori shmem finalized")
+                    logger.debug("[rank %s] mori shmem finalized",
+                                 self.dp_rank)
                 except Exception as shmem_error:
-                    logger.debug(f"[rank {self.dp_rank}] shmem finalize failed (may not have been active): {shmem_error}")
+                    logger.debug(
+                        "[rank %s] shmem finalize failed "
+                        "(may not have been active): %s", self.dp_rank,
+                        shmem_error)
 
-            logger.debug(f"[rank {self.dp_rank}] mori resources cleaned up")
+            logger.debug("[rank %s] mori resources cleaned up", self.dp_rank)
 
         except Exception as e:
-            logger.warning(f"[rank {self.dp_rank}] Error during mori cleanup: {e}")
+            logger.warning("[rank %s] Error during mori cleanup: %s",
+                           self.dp_rank, e)

vllm/distributed/device_communicators/base_device_communicator.py

@@ -7,7 +7,9 @@
 import torch
 import torch.distributed as dist
 from torch.distributed import ProcessGroup
+
 from vllm.logger import init_logger
+
 logger = init_logger(__name__)
 
 

vllm/model_executor/layers/fused_moe/aiter_experts.py

@@ -1,16 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """
 Aiter-based expert processing for Mori integration.
 """
 
-from typing import Any, Optional
+from typing import Optional
 
 import torch
 
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
 from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
-    rocm_aiter_fused_experts,
-)
+    rocm_aiter_fused_experts)
+from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
+    TopKWeightAndReduceNoOP)
 
 
 class AiterExperts(mk.FusedMoEPermuteExpertsUnpermute):
@@ -24,11 +27,9 @@ class AiterExperts(mk.FusedMoEPermuteExpertsUnpermute):
     def __init__(
         self,
         max_num_tokens: int,
-        quant_config: FusedMoEQuantConfig = None,
+        quant_config: FusedMoEQuantConfig,
     ):
-        super().__init__(
-            quant_config=quant_config,
-        )
+        super().__init__(quant_config=quant_config, )
         self.max_num_tokens = max_num_tokens
 
     @property
@@ -51,10 +52,6 @@ def supports_expert_map(self) -> bool:
 
     def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
         """Aiter handles weight and reduce internally."""
-        from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
-            TopKWeightAndReduceNoOP,
-        )
-
         return TopKWeightAndReduceNoOP()
 
     def workspace_shapes(
@@ -101,6 +98,11 @@ def apply(
         Process expert computation using Aiter kernels.
         Works with pre-dispatched tokens from Mori all2all.
         """
+        if expert_tokens_meta is not None:
+            expert_num_tokens = expert_tokens_meta.expert_num_tokens
+        else:
+            expert_num_tokens = None
+
         # Call Aiter fused MoE expert processing
         result = rocm_aiter_fused_experts(
             hidden_states=hidden_states,
@@ -111,7 +113,7 @@ def apply(
             activation=activation,
             apply_router_weight_on_input=apply_router_weight_on_input,
             expert_map=expert_map,
-            expert_num_tokens=expert_tokens_meta.expert_num_tokens if expert_tokens_meta is not None else None,
+            expert_num_tokens=expert_num_tokens,
             output_dtype=output.dtype,
             quant_config=self.quant_config,
             a1q_scale=a1q_scale,

vllm/model_executor/layers/fused_moe/layer.py

@@ -40,8 +40,8 @@
 from vllm.model_executor.utils import set_weight_attrs
 from vllm.platforms import current_platform
 from vllm.platforms.interface import CpuArchEnum
-from vllm.utils import (cdiv, direct_register_custom_op, has_deep_ep, has_pplx,
-                        has_mori, round_up)
+from vllm.utils import (cdiv, direct_register_custom_op, has_deep_ep, has_mori,
+                        has_pplx, round_up)
 from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
 from vllm.v1.worker.ubatching import dbo_current_ubatch_id
 
@@ -75,9 +75,12 @@ def _eplb_map_to_physical_and_record(
 
 if is_rocm_aiter_moe_enabled():
     from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (  # noqa: E501
-        rocm_aiter_grouped_topk as grouped_topk)
+        rocm_aiter_grouped_topk)
+    grouped_topk_impl = rocm_aiter_grouped_topk
 else:
     from vllm.model_executor.layers.fused_moe.fused_moe import grouped_topk
+    grouped_topk_impl = grouped_topk
+
 if current_platform.is_tpu():
     from .moe_pallas import fused_moe as fused_moe_pallas
 else:
@@ -210,21 +213,20 @@ def _maybe_make_prepare_finalize(
                 use_fp8_dispatch=use_fp8_dispatch,
             )
         elif moe.use_mori_kernels:
-            use_fp8_dispatch = (
-                quant_config is not None
-                and quant_config.quant_dtype == current_platform.fp8_dtype()
-            )
+            use_fp8_dispatch = (quant_config is not None
+                                and quant_config.quant_dtype
+                                == current_platform.fp8_dtype())
             scale_dim = 0
             scale_type_size = 0
             quant_dtype = None
             if use_fp8_dispatch:
+                assert quant_config is not None
                 scale_dim = quant_config.scale_shape(
                     moe.max_num_tokens,
                     moe.hidden_dim,
                 )[-1]
-                scale_type_size = (
-                    torch.float32.itemsize
-                )  # aiter quantization uses float32 scale
+                scale_type_size = (torch.float32.itemsize
+                                   )  # aiter quantization uses float32 scale
                 quant_dtype = quant_config.quant_dtype
 
             all_to_all_args = dict(
@@ -394,7 +396,7 @@ def select_gemm_impl(
                 quant_config=self.moe_quant_config,
             )
         elif (prepare_finalize.activation_format ==
-                FusedMoEActivationFormat.BatchedExperts):
+              FusedMoEActivationFormat.BatchedExperts):
             logger.debug("BatchedTritonExperts %s", self.moe)
             return BatchedTritonExperts(
                 max_num_tokens=self.moe.max_num_tokens,
@@ -1760,7 +1762,7 @@ def select_experts(
         if use_grouped_topk:
             assert topk_group is not None
             assert num_expert_group is not None
-            topk_weights, topk_ids = grouped_topk(
+            topk_weights, topk_ids = grouped_topk_impl(
                 hidden_states=hidden_states,
                 gating_output=router_logits,
                 topk=top_k,