diff --git a/tests/models/kimi_k3/test_amd_latent_moe_tail.py b/tests/models/kimi_k3/test_amd_latent_moe_tail.py new file mode 100644 index 000000000000..b04cc18e3fbd --- /dev/null +++ b/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 diff --git a/vllm/models/kimi_k3/amd/linear.py b/vllm/models/kimi_k3/amd/linear.py index 4e0088e9e192..993b6f0b84eb 100644 --- a/vllm/models/kimi_k3/amd/linear.py +++ b/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)