applied few suggestions from code-assistant

Author: whitememory

vllm/distributed/device_communicators/all2all.py

@@ -510,8 +510,10 @@ def _ensure_shmem_initialized(self):
                 logger.debug(
                     "[rank %s] torch process group shmem init failed: %s",
                     self.rank, torch_error)
-
-            self._shmem_initialized = True
+                self._shmem_initialized = True
+                logger.warning(
+                    "[rank %s] Continuing without mori shmem optimization",
+                    self.rank)
 
         except Exception as e:
             logger.error("[rank %s] mori shmem initialization failed: %s",

vllm/model_executor/layers/fused_moe/mori_prepare_finalize.py

@@ -88,9 +88,9 @@ def prepare(
 
         Args:
             a1: Input hidden states [num_tokens, hidden_dim]
-            a1_scale: Input activation scales
             topk_weights: Top-k routing weights [num_experts, experts_per_token]
             topk_ids: Top-k expert indices [num_experts, experts_per_token]
+            apply_router_weight_on_input: Whether to apply router weight
             quant_config: Quantization config
 
         Returns: