[ROCm][MoE] Accumulate the shared-expert output in AITER's stage 2; fold routed scaling into top-k

vllm/_aiter_ops.py

@@ -325,7 +325,7 @@ def _validate_rocm_aiter_fused_moe_shared_expert_args(
     return has_shared_expert
 
 
-def _rocm_aiter_fused_moe_impl(
+def _aiter_fused_moe(
     hidden_states: torch.Tensor,
     w1: torch.Tensor,
     w2: torch.Tensor,
@@ -356,6 +356,7 @@ def _rocm_aiter_fused_moe_impl(
     shared_w2_scale: torch.Tensor | None = None,
     shared_expert_id: int = -1,
     q_dtype_a: torch.dtype | None = None,
+    **fused_moe_kwargs,
 ) -> torch.Tensor:
     has_shared_expert = _validate_rocm_aiter_fused_moe_shared_expert_args(
         shared_w1,
@@ -395,6 +396,7 @@ def _rocm_aiter_fused_moe_impl(
         # is checked at config time, so this is only reached on an AITER
         # build new enough to accept it.
         extra_kwargs["quant_dtype_a"] = q_dtype_a
+    extra_kwargs.update(fused_moe_kwargs)
 
     return fused_moe(
         hidden_states,
@@ -422,6 +424,41 @@ def _rocm_aiter_fused_moe_impl(
     )
 
 
+def _rocm_aiter_fused_moe_impl(
+    hidden_states: torch.Tensor,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    topk_weight: torch.Tensor,
+    topk_ids: torch.Tensor,
+    expert_mask: torch.Tensor | None = None,
+    activation_method: int = 0,
+    quant_method: int = 0,
+    doweight_stage1: bool = False,
+    w1_scale: torch.Tensor | None = None,
+    w2_scale: torch.Tensor | None = None,
+    a1_scale: torch.Tensor | None = None,
+    a2_scale: torch.Tensor | None = None,
+    num_local_tokens: torch.Tensor | None = None,
+    output_dtype: torch.dtype | None = None,
+    hidden_pad: int = 0,
+    intermediate_pad: int = 0,
+    gate_mode: str = "",
+    bias1: torch.Tensor | None = None,
+    bias2: torch.Tensor | None = None,
+    moe_sorting_dispatch_policy: int = 0,
+    swiglu_limit: float = 0.0,
+    beta: float | None = None,
+    linear_beta: float | None = None,
+    shared_w1: torch.Tensor | None = None,
+    shared_w2: torch.Tensor | None = None,
+    shared_w1_scale: torch.Tensor | None = None,
+    shared_w2_scale: torch.Tensor | None = None,
+    shared_expert_id: int = -1,
+    q_dtype_a: torch.dtype | None = None,
+) -> torch.Tensor:
+    return _aiter_fused_moe(**locals())
+
+
 def _rocm_aiter_fused_moe_fake(
     hidden_states: torch.Tensor,
     w1: torch.Tensor,
@@ -2883,6 +2920,17 @@ def fused_moe_supports_quant_dtype_a(cls) -> bool:
 
         return "quant_dtype_a" in inspect.signature(fused_moe).parameters
 
+    @classmethod
+    @if_aiter_supported
+    @functools.cache
+    def fused_moe_supports_accumulate_output(cls) -> bool:
+        """Probe whether the installed aiter.fused_moe accepts `accumulate_output`."""
+        import inspect
+
+        from aiter.fused_moe import fused_moe
+
+        return "output_ready_event" in inspect.signature(fused_moe).parameters
+
     @staticmethod
     def _probe_dsv4_i384_fhmoe_capability(num_tokens: int) -> bool:
         """Probe AITER's CSV-backed DSV4 native-I384 FHMoE contract."""
@@ -3575,7 +3623,48 @@ def fused_moe(
         shared_w2_scale: torch.Tensor | None = None,
         shared_expert_id: int = -1,
         q_dtype_a: torch.dtype | None = None,
+        addend: torch.Tensor | None = None,
+        addend_ready: torch.cuda.Event | None = None,
     ) -> torch.Tensor:
+        """With `addend`, accumulates into it and returns it, waiting on
+        `addend_ready` only before stage 2. Eager only: the event can't cross
+        the custom-op schema."""
+        if addend is not None:
+            return _aiter_fused_moe(
+                hidden_states,
+                w1,
+                w2,
+                topk_weight,
+                topk_ids,
+                expert_mask,
+                activation_method,
+                quant_method,
+                doweight_stage1,
+                w1_scale,
+                w2_scale,
+                a1_scale,
+                a2_scale,
+                num_local_tokens,
+                output_dtype,
+                hidden_pad,
+                intermediate_pad,
+                gate_mode,
+                bias1,
+                bias2,
+                moe_sorting_dispatch_policy,
+                swiglu_limit,
+                beta,
+                linear_beta,
+                shared_w1,
+                shared_w2,
+                shared_w1_scale,
+                shared_w2_scale,
+                shared_expert_id,
+                q_dtype_a,
+                output=addend,
+                accumulate_output=True,
+                output_ready_event=addend_ready,
+            )
         return torch.ops.vllm.rocm_aiter_fused_moe(
             hidden_states,
             w1,

vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py

@@ -291,6 +291,8 @@ def rocm_aiter_fused_experts(
     shared_w1_scale: torch.Tensor | None = None,
     shared_w2_scale: torch.Tensor | None = None,
     shared_expert_id: int = -1,
+    addend: torch.Tensor | None = None,
+    addend_ready: torch.cuda.Event | None = None,
 ) -> torch.Tensor:
     """ROCm AITER fused MoE expert computation."""
     if quant_config is None:
@@ -338,7 +340,7 @@ def rocm_aiter_fused_experts(
             "AITER tkw1 kernel does not support `num_local_tokens`"
         )
 
-        return rocm_aiter_ops.asm_moe_tkw1(
+        result = rocm_aiter_ops.asm_moe_tkw1(
             hidden_states,
             w1,
             w2,
@@ -353,6 +355,11 @@ def rocm_aiter_fused_experts(
             expert_mask=expert_mask,
             activation_method=activation_method,
         )
+        if addend is None:
+            return result
+        if addend_ready is not None:
+            torch.cuda.current_stream().wait_event(addend_ready)
+        return addend.add_(result)
 
     else:
         quant_method = QuantMethod.NO.value
@@ -475,6 +482,8 @@ def rocm_aiter_fused_experts(
             shared_w2_scale=shared_w2_scale,
             shared_expert_id=shared_expert_id,
             q_dtype_a=q_dtype_a,
+            addend=addend,
+            addend_ready=addend_ready,
         )
 
 
@@ -537,6 +546,9 @@ def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bo
     def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
         return TopKWeightAndReduceNoOP()
 
+    def supports_output_addend(self) -> bool:
+        return rocm_aiter_ops.fused_moe_supports_accumulate_output()
+
     def workspace_shapes(
         self,
         M: int,
@@ -576,6 +588,8 @@ def apply(
         workspace2: torch.Tensor,
         expert_tokens_meta: mk.ExpertTokensMetadata | None,
         apply_router_weight_on_input: bool,
+        addend: torch.Tensor | None = None,
+        addend_ready: torch.cuda.Event | None = None,
     ):
         # TODO(rob): rocm_aiter_fused_experts uses self.quant_config's
         # a_scales for static quantization. Update this to fit better
@@ -601,6 +615,8 @@ def apply(
             num_local_tokens=num_local_tokens,
             output_dtype=output.dtype,
             moe_sorting_dispatch_policy=rocm_aiter_ops.get_moe_dispatch_policy(),
+            addend=addend,
+            addend_ready=addend_ready,
         )
         # avoid redundant copy when output is a view of the result
         if (

vllm/model_executor/layers/fused_moe/modular_kernel.py

@@ -954,6 +954,10 @@ def activation(
     def finalize_weight_and_reduce_impl(self) -> TopKWeightAndReduce:
         raise NotImplementedError
 
+    def supports_output_addend(self) -> bool:
+        """Whether `apply` can accumulate into `addend` after `addend_ready`."""
+        return False
+
     @abstractmethod
     def apply(
         self,
@@ -1336,6 +1340,7 @@ def _fused_experts(
         apply_router_weight_on_input: bool,
         expert_tokens_meta: ExpertTokensMetadata | None,
         output_alias: torch.Tensor | None = None,
+        shared_experts: SharedExperts | None = None,
     ) -> torch.Tensor | UnfinalizedMoEOutput:
         _, M_full, N, K, top_k = self.fused_experts.moe_problem_size(
             a1q, w1, w2, topk_ids
@@ -1388,6 +1393,16 @@ def _fused_experts(
         elif use_output_alias:
             fused_out = output_alias
 
+        addend_kwargs: dict[str, Any] = {}
+        if (
+            shared_experts is not None
+            and not self.prepare_finalize.supports_async()
+            and self.fused_experts.supports_output_addend()
+        ):
+            claimed = shared_experts.claim_output(fused_out)
+            if claimed is not None:
+                addend_kwargs = dict(addend=claimed[0], addend_ready=claimed[1])
+
         unfinalized = self.fused_experts.apply(
             output=fused_out,
             hidden_states=a1q,
@@ -1404,6 +1419,7 @@ def _fused_experts(
             workspace2=workspace2,
             expert_tokens_meta=expert_tokens_meta,
             apply_router_weight_on_input=apply_router_weight_on_input,
+            **addend_kwargs,
         )
 
         # Experts that stopped after GEMM2 wrote nothing into `fused_out`; the
@@ -1575,6 +1591,7 @@ def apply(
             apply_router_weight_on_input=apply_router_weight_on_input,
             expert_tokens_meta=expert_tokens_meta,
             output_alias=output,
+            shared_experts=shared_experts,
         )
 
         if lora_ctx is not None:
@@ -1727,6 +1744,13 @@ def __init__(
 
         self._post_init_setup()
 
+        # Resolved eagerly: the runner queries this while being traced.
+        self._supports_output_addend = (
+            isinstance(self.impl, FusedMoEKernelModularImpl)
+            and not self.prepare_finalize.supports_async()
+            and self.impl.fused_experts.supports_output_addend()
+        )
+
     @property
     def can_overlap_shared_experts(self) -> bool:
         if isinstance(self.impl, FusedMoEKernelModularImpl):
@@ -1772,6 +1796,9 @@ def output_is_reduced(self) -> bool:
     def supports_deferred_moe_finalize(self) -> bool:
         return self.prepare_finalize.supports_deferred_moe_finalize()
 
+    def supports_output_addend(self) -> bool:
+        return self._supports_output_addend
+
     def apply_monolithic(
         self,
         hidden_states: torch.Tensor,

vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py

@@ -272,9 +272,8 @@ def fused_topk_bias(
                 num_expert_group=num_expert_group,
                 topk_group=num_expert_group,
                 need_renorm=renormalize,
+                routed_scaling_factor=routed_scaling_factor,
             )
-            if routed_scaling_factor != 1.0:
-                topk_weights *= routed_scaling_factor
             return topk_weights, topk_ids
 
     if scoring_func == "sqrtsoftplus":

vllm/model_executor/layers/fused_moe/runner/moe_runner.py

@@ -426,6 +426,25 @@ def _fused_output_is_reduced(self) -> bool:
             and self._quant_method.moe_kernel.output_is_reduced()
         )
 
+    @property
+    def _fold_shared_output(self) -> bool:
+        """Whether the MoE op returns the shared + routed sum."""
+        moe_kernel = self._quant_method.moe_kernel
+        return (
+            self._shared_experts is not None
+            and moe_kernel is not None
+            and moe_kernel.supports_output_addend()
+            and self.routed_input_transform is None
+            and self.routed_output_transform is None
+            and self.routed_scaling_factor == 1.0
+            and not self.do_naive_dispatch_combine
+            and self.moe_config.pcp_size == 1
+            and not self._fused_output_is_reduced
+            and not self._quant_method.has_unpadded_output
+            and self.moe_config.hidden_dim_unpadded
+            in (None, self.moe_config.hidden_dim)
+        )
+
     def _maybe_reduce_shared_expert_output(
         self,
         shared_output: torch.Tensor | None,
@@ -646,10 +665,16 @@ def _apply_quant_method(
             assert self._shared_experts is not None
             self._shared_experts.wait()
 
-        return (
-            self._shared_experts.output if self._shared_experts is not None else None,
-            fused_out,
-        )
+        if self._shared_experts is None:
+            return None, fused_out
+
+        claimed = self._shared_experts.output_claimed
+        shared_output = self._shared_experts.output
+        if not self._shared_experts.fold_output:
+            return shared_output, fused_out
+        if not claimed:
+            fused_out = shared_output + cast(torch.Tensor, fused_out)
+        return None, fused_out
 
     def _sequence_parallel_context(self):
         """Return a context manager for sequence-parallel token
@@ -735,7 +760,12 @@ def forward(
             )
         )
 
-        result = self._forward_entry(
+        forward_entry = (
+            torch.ops.vllm.moe_forward
+            if self._fold_shared_output
+            else self._forward_entry
+        )
+        result = forward_entry(
             hidden_states,
             router_logits,
             shared_experts_input,
@@ -860,8 +890,7 @@ def _maybe_combine(
                 )
             hidden_states = get_pcp_group().reduce_scatter(hidden_states, dim=0)
 
-        if self.shared_experts is not None:
-            assert shared_output is not None
+        if shared_output is not None:
             return shared_output, hidden_states
         else:
             return hidden_states
@@ -897,6 +926,7 @@ def _forward_impl(
         # before routed expert dispatch.
         shared_experts_overlapping = False
         if self._shared_experts is not None:
+            self._shared_experts.fold_output = self._fold_shared_output
             shared_experts_overlapping = self._shared_experts.maybe_forward_async(
                 shared_experts_input
             )