cleanup

Signed-off-by: Bill Nell <bnell@redhat.com>

Author: bnellnm

vllm/model_executor/layers/fused_moe/layer.py

@@ -237,8 +237,6 @@ def maybe_make_prepare_finalize(self) -> FusedMoEPrepareAndFinalize | None:
         else:
             return None
 
-    # Note: init_prepare_finalize should only be called by
-    # prepare_communication_buffer_for_model.
     def init_prepare_finalize(
         self, layer: torch.nn.Module
     ) -> FusedMoEModularKernel | None:
@@ -266,8 +264,8 @@ def init_prepare_finalize(
                 experts,
                 layer.shared_experts,
             )
-
-        return None
+        else:
+            return None
 
     def select_gemm_impl(
         self,
@@ -321,7 +319,9 @@ def apply(
 @CustomOp.register("modular_fused_moe")
 class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
     def __init__(
-        self, old_moe_method: FusedMoEMethodBase, fused_experts: FusedMoEModularKernel
+        self,
+        old_moe_method: FusedMoEMethodBase,
+        fused_experts: FusedMoEModularKernel,
     ):
         super().__init__(old_moe_method.moe)
         # Find better way to copy attributes
@@ -374,6 +374,8 @@ def apply(
         logical_to_physical_map: torch.Tensor | None = None,
         logical_replica_count: torch.Tensor | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        assert self.fused_experts is not None
+
         # Is getattr needed?
         zero_expert_num = getattr(layer, "zero_expert_num", 0)
         zero_expert_type = getattr(layer, "zero_expert_type", None)
@@ -1372,6 +1374,8 @@ def __init__(
                 logits_shape, dtype=moe.in_dtype, device=torch.cuda.current_device()
             )
 
+    # Note: init_prepare_finalize should only be called by
+    # prepare_communication_buffer_for_model.
     def init_prepare_finalize(self) -> None:
         mk = self.quant_method.init_prepare_finalize(self)
         if mk is not None: