perf(rocm): fuse Kimi-K3 B1 latent MoE tail

tests/models/kimi_k3/test_amd_latent_moe_tail.py

@@ -0,0 +1,125 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import importlib
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from vllm._aiter_ops import rocm_aiter_ops
+from vllm.models.kimi_k3.amd.linear import (
+    KimiAMDLatentMoERunner,
+    KimiRoutedOutputTransform,
+)
+from vllm.platforms import current_platform
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_rocm(),
+    reason="Kimi-K3 AITER latent-MoE tail requires ROCm",
+)
+
+
+def _transform() -> KimiRoutedOutputTransform:
+    transform = object.__new__(KimiRoutedOutputTransform)
+    nn.Module.__init__(transform)
+    transform.norm = SimpleNamespace(
+        weight=torch.empty(3584, device="meta"),
+        variance_epsilon=1.0e-6,
+    )
+    transform.up_proj = SimpleNamespace(weight=torch.empty(7168, 3584, device="meta"))
+    return transform
+
+
+def test_forward_with_shared_delegates_to_supported_aiter_kernel(monkeypatch):
+    latent_moe_tail_module = importlib.import_module("aiter.ops.flydsl.latent_moe_tail")
+
+    transform = _transform()
+    routed = torch.empty(1, 3584)
+    shared = torch.empty(1, 7168)
+    expected = torch.empty_like(shared)
+    calls = []
+
+    monkeypatch.setattr(rocm_aiter_ops, "is_enabled", lambda: True)
+    monkeypatch.setattr(
+        latent_moe_tail_module,
+        "supports_latent_moe_tail",
+        lambda *args: True,
+    )
+
+    def fused_tail(*args):
+        calls.append(args)
+        return expected
+
+    monkeypatch.setattr(latent_moe_tail_module, "latent_moe_tail", fused_tail)
+
+    assert transform.forward_with_shared(routed, shared) is expected
+    assert len(calls) == 1
+    assert calls[0][0] is routed
+    assert calls[0][1] is shared
+    assert calls[0][2] is transform.norm.weight
+    assert calls[0][3] is transform.up_proj.weight
+    assert calls[0][4] == transform.norm.variance_epsilon
+
+
+def test_forward_with_shared_preserves_fallbacks(monkeypatch):
+    transform = _transform()
+    routed = torch.empty(8, 3584)
+    shared = torch.empty(8, 7168)
+
+    monkeypatch.setattr(rocm_aiter_ops, "is_enabled", lambda: False)
+    assert transform.forward_with_shared(routed, shared) is None
+
+    monkeypatch.setattr(rocm_aiter_ops, "is_enabled", lambda: True)
+    latent_moe_tail_module = importlib.import_module("aiter.ops.flydsl.latent_moe_tail")
+
+    monkeypatch.setattr(
+        latent_moe_tail_module,
+        "supports_latent_moe_tail",
+        lambda *args: False,
+    )
+    monkeypatch.setattr(
+        latent_moe_tail_module,
+        "latent_moe_tail",
+        lambda *args: pytest.fail("unsupported inputs must use the fallback"),
+    )
+    assert transform.forward_with_shared(routed, shared) is None
+
+
+def test_runner_consumes_fused_tail_once_without_leaking_state(monkeypatch):
+    runner = object.__new__(KimiAMDLatentMoERunner)
+    nn.Module.__init__(runner)
+    runner.routed_scaling_factor = 1.0
+    transform = _transform()
+    runner.routed_output_transform = transform
+
+    routed = torch.tensor([[1.0]])
+    shared = torch.tensor([[2.0]])
+    fused_result = torch.tensor([[3.0]])
+
+    monkeypatch.setattr(
+        transform,
+        "forward_with_shared",
+        lambda routed, shared: fused_result,
+    )
+    remaining_shared, result = runner._maybe_apply_routed_scale_to_output(
+        shared, routed
+    )
+    assert remaining_shared is None
+    assert result is fused_result
+    assert runner.apply_routed_output_transform(result) is fused_result
+
+    fallback_result = torch.tensor([[4.0]])
+    monkeypatch.setattr(
+        transform,
+        "forward_with_shared",
+        lambda routed, shared: None,
+    )
+    monkeypatch.setattr(transform, "forward", lambda routed: fallback_result)
+    remaining_shared, result = runner._maybe_apply_routed_scale_to_output(
+        shared, routed
+    )
+    assert remaining_shared is shared
+    assert result is routed
+    assert runner.apply_routed_output_transform(result) is fallback_result

vllm/models/kimi_k3/amd/linear.py

@@ -7,6 +7,7 @@
 import torch
 from torch import nn
 
+from vllm._aiter_ops import rocm_aiter_ops
 from vllm.config import CacheConfig, VllmConfig
 from vllm.distributed import (
     get_pp_group,
@@ -19,6 +20,7 @@
     fused_moe_make_expert_params_mapping,
 )
 from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
+from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
@@ -134,6 +136,78 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         hidden_states, _ = self.up_proj(hidden_states)
         return hidden_states
 
+    def forward_with_shared(
+        self,
+        hidden_states: torch.Tensor,
+        shared_output: torch.Tensor,
+    ) -> torch.Tensor | None:
+        """Fuse the supported local tail, or return ``None`` for fallback."""
+
+        if self.norm is None or not rocm_aiter_ops.is_enabled():
+            return None
+        try:
+            from aiter.ops.flydsl.latent_moe_tail import (
+                latent_moe_tail,
+                supports_latent_moe_tail,
+            )
+        except (ImportError, ModuleNotFoundError):
+            return None
+
+        up_weight = self.up_proj.weight
+        if not supports_latent_moe_tail(
+            hidden_states,
+            shared_output,
+            self.norm.weight,
+            up_weight,
+            self.norm.variance_epsilon,
+        ):
+            return None
+        return latent_moe_tail(
+            hidden_states,
+            shared_output,
+            self.norm.weight,
+            up_weight,
+            self.norm.variance_epsilon,
+        )
+
+
+class KimiAMDLatentMoERunner(MoERunner):
+    """Use the AMD local-tail primitive after routed/shared reductions."""
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self._skip_next_routed_output_transform = False
+
+    def _maybe_apply_routed_scale_to_output(
+        self,
+        shared_output: torch.Tensor | None,
+        fused_output: torch.Tensor,
+    ) -> tuple[torch.Tensor | None, torch.Tensor]:
+        shared_output, fused_output = super()._maybe_apply_routed_scale_to_output(
+            shared_output, fused_output
+        )
+        self._skip_next_routed_output_transform = False
+        transform = self.routed_output_transform
+        if shared_output is not None and isinstance(
+            transform, KimiRoutedOutputTransform
+        ):
+            result = transform.forward_with_shared(fused_output, shared_output)
+            if result is not None:
+                # MoERunner applies the routed transform in the next synchronous
+                # pipeline step. The fused primitive has already performed it.
+                self._skip_next_routed_output_transform = True
+                return None, result
+        return shared_output, fused_output
+
+    def apply_routed_output_transform(
+        self,
+        fused_output: torch.Tensor,
+    ) -> torch.Tensor:
+        if self._skip_next_routed_output_transform:
+            self._skip_next_routed_output_transform = False
+            return fused_output
+        return super().apply_routed_output_transform(fused_output)
+
 
 def _apply_attn_res(
     prefix_sum: torch.Tensor,
@@ -280,6 +354,7 @@ def __init__(
             routed_scaling_factor=self.routed_scaling_factor,
             routed_input_transform=self.routed_expert_down_proj,
             routed_output_transform=self.routed_output_transform,
+            runner_cls=KimiAMDLatentMoERunner if self.use_latent_moe else None,
         )
         if self.padded_moe_intermediate_size != moe_intermediate_size:
             w13_weight = getattr(self.experts, "w13_weight", None)