diff --git a/tests/model_executor/test_moe_runner_fused_ar_rmsnorm.py b/tests/model_executor/test_moe_runner_fused_ar_rmsnorm.py new file mode 100644 index 000000000000..9d85c3fa2cf7 --- /dev/null +++ b/tests/model_executor/test_moe_runner_fused_ar_rmsnorm.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Eligibility for the ROCm fused all-reduce + RMSNorm latent-MoE path. + +``MoERunner._can_fuse_ar_rmsnorm`` decides whether the latent-MoE all-reduce +and its routed-output RMSNorm collapse into one aiter kernel. Getting it wrong +is silent: a false positive would drop a post-norm op (e.g. an unrecognised +transform) or run the kernel on an unsupported layout. These pin the predicate +device-free by mocking the platform and feeding a tensor stand-in, so the +boolean logic is tested without a GPU or an initialised aiter all-reduce. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from tests.utils import ensure_current_vllm_config +from vllm.model_executor.layers.fused_moe.runner import moe_runner +from vllm.model_executor.layers.layernorm import RMSNorm + +pytestmark = pytest.mark.cpu_test + +HIDDEN = 64 +LATENT = 32 + + +class _FakeTensor: + """Stand-in exposing only what the predicate reads (no GPU needed).""" + + def __init__( + self, + *, + is_cuda: bool = True, + dim: int = 2, + contiguous: bool = True, + dtype: torch.dtype = torch.bfloat16, + ) -> None: + self.is_cuda = is_cuda + self._dim = dim + self._contiguous = contiguous + self.dtype = dtype + + def dim(self) -> int: + return self._dim + + def is_contiguous(self) -> bool: + return self._contiguous + + +@pytest.fixture +def norm(): + # RMSNorm construction needs a current vLLM config for its custom-op setup. + with ensure_current_vllm_config(): + yield RMSNorm(LATENT, eps=1e-5) + + +def _runner(norm, **overrides): + """Bare MoERunner carrying only the attributes the predicate reads.""" + transform: SimpleNamespace | None = SimpleNamespace( + norm=overrides.pop("transform_norm", norm), + up_proj=overrides.pop("transform_up_proj", lambda x: (x, None)), + ) + if overrides.pop("no_transform", False): + transform = None + + runner = object.__new__(moe_runner.MoERunner) + runner.routed_output_transform = transform + runner.routed_scaling_factor = overrides.pop("routed_scaling_factor", 1.0) + runner.moe_config = SimpleNamespace( + tp_size=overrides.pop("tp_size", 8), + ep_size=overrides.pop("ep_size", 1), + is_sequence_parallel=overrides.pop("is_sequence_parallel", False), + ) + assert not overrides, f"unexpected overrides: {overrides}" + return runner + + +@pytest.fixture(autouse=True) +def _force_rocm_with_aiter(monkeypatch): + """Default the environment gates to eligible; each test perturbs one.""" + monkeypatch.setattr(moe_runner.current_platform, "is_rocm", lambda: True) + monkeypatch.setattr(moe_runner, "_aiter_fused_ar_rmsnorm", object()) + + +def test_eligible_under_the_kimi_k3_serving_shape(norm): + runner = _runner(norm) + assert runner._can_fuse_ar_rmsnorm(_FakeTensor(), False) + + +def test_disabled_when_aiter_op_missing(norm, monkeypatch): + monkeypatch.setattr(moe_runner, "_aiter_fused_ar_rmsnorm", None) + runner = _runner(norm) + assert not runner._can_fuse_ar_rmsnorm(_FakeTensor(), False) + + +def test_disabled_off_rocm(norm, monkeypatch): + monkeypatch.setattr(moe_runner.current_platform, "is_rocm", lambda: False) + runner = _runner(norm) + assert not runner._can_fuse_ar_rmsnorm(_FakeTensor(), False) + + +def test_disabled_without_transform(norm): + runner = _runner(norm, no_transform=True) + assert not runner._can_fuse_ar_rmsnorm(_FakeTensor(), False) + + +def test_disabled_when_norm_is_not_rmsnorm(norm): + # A transform whose ``.norm`` is some other module must not be fused, or + # its real normalization would be silently replaced. + runner = _runner(norm, transform_norm=torch.nn.LayerNorm(LATENT)) + assert not runner._can_fuse_ar_rmsnorm(_FakeTensor(), False) + + +def test_disabled_when_up_proj_not_callable(norm): + runner = _runner(norm, transform_up_proj=object()) + assert not runner._can_fuse_ar_rmsnorm(_FakeTensor(), False) + + +def test_disabled_with_routed_scaling_factor(norm): + # The fused path assumes the routed-scale step stays a no-op. + runner = _runner(norm, routed_scaling_factor=2.0) + assert not runner._can_fuse_ar_rmsnorm(_FakeTensor(), False) + + +def test_eligible_under_expert_parallelism_without_tp(norm): + # With EP the MoE reports tp_size == 1 but still all-reduces the routed + # output across the EP ranks, so the fused path must engage on ep_size too. + runner = _runner(norm, tp_size=1, ep_size=8) + assert runner._can_fuse_ar_rmsnorm(_FakeTensor(), False) + + +def test_disabled_without_tp_or_ep(norm): + runner = _runner(norm, tp_size=1, ep_size=1) + assert not runner._can_fuse_ar_rmsnorm(_FakeTensor(), False) + + +def test_disabled_under_sequence_parallelism(norm): + runner = _runner(norm, is_sequence_parallel=True) + assert not runner._can_fuse_ar_rmsnorm(_FakeTensor(), False) + + +def test_disabled_when_already_reduced(norm): + runner = _runner(norm) + assert not runner._can_fuse_ar_rmsnorm(_FakeTensor(), True) + + +def test_disabled_for_non_cuda_tensor(norm): + runner = _runner(norm) + assert not runner._can_fuse_ar_rmsnorm(_FakeTensor(is_cuda=False), False) + + +def test_disabled_for_non_2d_tensor(norm): + runner = _runner(norm) + assert not runner._can_fuse_ar_rmsnorm(_FakeTensor(dim=3), False) + + +def test_disabled_for_non_contiguous_tensor(norm): + runner = _runner(norm) + assert not runner._can_fuse_ar_rmsnorm(_FakeTensor(contiguous=False), False) + + +def test_disabled_for_unsupported_dtype(norm): + runner = _runner(norm) + assert not runner._can_fuse_ar_rmsnorm(_FakeTensor(dtype=torch.float32), False) diff --git a/tests/models/kimi_k3/test_amd_latent_moe_runner.py b/tests/models/kimi_k3/test_amd_latent_moe_runner.py index ec270c04fe6c..f19fddd05b18 100644 --- a/tests/models/kimi_k3/test_amd_latent_moe_runner.py +++ b/tests/models/kimi_k3/test_amd_latent_moe_runner.py @@ -1,9 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""The ROCm latent-MoE tail must equal the replicated up-projection. +"""The ROCm latent-MoE tails must equal the replicated up-projection. -A wrong shard offset or a dropped accumulation still runs and still produces -plausible text, so these pin the arithmetic rather than the behaviour. +A wrong shard offset, a dropped accumulation, or an overlapped all-reduce that +races still runs and still produces plausible text, so these pin the arithmetic +rather than the behaviour. The unit tests cover tier selection, which decides +which tail arithmetic runs. """ from types import SimpleNamespace @@ -23,7 +25,11 @@ from vllm.model_executor.layers.fused_moe.runner import moe_runner from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ReplicatedLinear -from vllm.models.kimi_k3.amd.latent_moe_runner import ROCmLatentMoERunner +from vllm.models.kimi_k3.amd import latent_moe_runner +from vllm.models.kimi_k3.amd.latent_moe_runner import ( + ROCmLatentMoERunner, + ROCmLatentTailTier, +) from vllm.models.kimi_k3.amd.linear import KimiRoutedOutputTransform from vllm.platforms import current_platform from vllm.utils.network_utils import get_open_port @@ -58,7 +64,7 @@ def _build_transform(device: torch.device) -> KimiRoutedOutputTransform: def _tail_runner( transform: KimiRoutedOutputTransform, tp_size: int ) -> ROCmLatentMoERunner: - """A runner carrying only what the tail reads, so no engine is needed. + """A runner carrying only what the tails read, so no engine is needed. ``_maybe_reduce_final_output`` is left as the real base-class method, so the all-reduce that stitches the shards is the real collective. @@ -66,8 +72,9 @@ def _tail_runner( runner = object.__new__(ROCmLatentMoERunner) attrs = { "routed_output_transform": transform, - "_up_proj_shard_size": HIDDEN_SIZE // tp_size, - "_logged_sharded_tail": False, + "_logged_column_parallel": False, + "_logged_overlap_fallback": False, + "_shared_ar_events": (torch.cuda.Event(), torch.cuda.Event()), "moe_config": SimpleNamespace( tp_size=tp_size, ep_size=1, @@ -94,7 +101,28 @@ def _rank_partials( return routed.mul_(0.01), shared -def _check_matches_replicated(device: torch.device, tp_size: int, rank: int) -> None: +def _replicated_reference( + routed_output: torch.Tensor, + shared_output: torch.Tensor, + transform: KimiRoutedOutputTransform, + group, +) -> torch.Tensor: + """latent-space all-reduce, RMSNorm, replicated up-proj, plus shared.""" + expected = F.linear( + F.rms_norm( + _all_reduced(routed_output, group), + (LATENT_SIZE,), + transform.norm.weight, + EPS, + ), + transform.up_proj.weight, + ) + return expected.add_(_all_reduced(shared_output, group)) + + +def _check_shard_matches_replicated( + device: torch.device, tp_size: int, rank: int +) -> None: transform = _build_transform(device) runner = _tail_runner(transform, tp_size) group = get_tp_group().device_group @@ -103,22 +131,29 @@ def _check_matches_replicated(device: torch.device, tp_size: int, rank: int) -> torch.manual_seed(100 * iteration + rank + 1) routed_output, shared_output = _rank_partials(num_tokens, device) - expected = F.linear( - F.rms_norm( - _all_reduced(routed_output, group), - (LATENT_SIZE,), - transform.norm.weight, - EPS, - ), - transform.up_proj.weight, - ) - expected.add_(_all_reduced(shared_output, group)) - + expected = _replicated_reference(routed_output, shared_output, transform, group) actual = runner._shard_up_proj_tail(routed_output, shared_output, None) torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) +def _check_overlap_matches_replicated( + device: torch.device, tp_size: int, rank: int +) -> None: + transform = _build_transform(device) + runner = _tail_runner(transform, tp_size) + group = get_tp_group().device_group + + for iteration, num_tokens in enumerate((1, 5, 8, 16, 5)): + torch.manual_seed(100 * iteration + rank + 1) + routed_output, shared_output = _rank_partials(num_tokens, device) + + expected = _replicated_reference(routed_output, shared_output, transform, group) + actual = runner._overlap_allreduce_tail(routed_output, shared_output, None) + + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) + + def _check_writes_only_its_own_shard( device: torch.device, tp_size: int, rank: int ) -> None: @@ -162,7 +197,8 @@ def _capture(states, trunc_size, output_is_reduced=None): _CHECKS = { - "matches_replicated": _check_matches_replicated, + "shard_matches_replicated": _check_shard_matches_replicated, + "overlap_matches_replicated": _check_overlap_matches_replicated, "own_shard_only": _check_writes_only_its_own_shard, } @@ -187,111 +223,194 @@ def _run_ranks(check: str, tp_size: int) -> None: @multi_gpu_test(num_gpus=4) -def test_sharded_tail_tp4_matches_replicated_projection() -> None: - _run_ranks("matches_replicated", 4) +def test_shard_tail_tp4_matches_replicated_projection() -> None: + _run_ranks("shard_matches_replicated", 4) @multi_gpu_test(num_gpus=8) -def test_sharded_tail_tp8_matches_replicated_projection() -> None: - _run_ranks("matches_replicated", 8) +def test_shard_tail_tp8_matches_replicated_projection() -> None: + _run_ranks("shard_matches_replicated", 8) @multi_gpu_test(num_gpus=4) -def test_sharded_tail_tp4_writes_only_its_own_shard() -> None: +def test_overlap_tail_tp4_matches_replicated_projection() -> None: + _run_ranks("overlap_matches_replicated", 4) + + +@multi_gpu_test(num_gpus=8) +def test_overlap_tail_tp8_matches_replicated_projection() -> None: + _run_ranks("overlap_matches_replicated", 8) + + +@multi_gpu_test(num_gpus=4) +def test_shard_tail_tp4_writes_only_its_own_shard() -> None: _run_ranks("own_shard_only", 4) -def _runner(**attrs) -> ROCmLatentMoERunner: +def _logic_runner( + *, + tp_size: int = 8, + hidden: int = HIDDEN_SIZE, + is_sequence_parallel: bool = False, + has_up_proj: bool = True, + has_shared_experts: bool = True, + routed_scaling_factor: float = 1.0, + fused_output_is_reduced: bool = False, +) -> ROCmLatentMoERunner: + """Build a runner with only the fields the pure-logic gates read. + + Bypasses ``__init__`` so no CUDA events are allocated; the tier gates never + touch the device. + """ runner = object.__new__(ROCmLatentMoERunner) + moe_kernel = ( + SimpleNamespace(output_is_reduced=lambda: True) + if fused_output_is_reduced + else None + ) + # ``_quant_method`` is a read-only property reading ``routed_experts``; + # ``_fused_output_is_reduced`` reaches through it to ``moe_kernel``. + attrs = { + "moe_config": SimpleNamespace( + tp_size=tp_size, is_sequence_parallel=is_sequence_parallel + ), + "_shared_experts": object() if has_shared_experts else None, + "routed_scaling_factor": routed_scaling_factor, + "routed_experts": SimpleNamespace( + quant_method=SimpleNamespace(moe_kernel=moe_kernel) + ), + "routed_output_transform": SimpleNamespace( + norm=None, + up_proj=SimpleNamespace(weight=torch.zeros(hidden, LATENT_SIZE)) + if has_up_proj + else None, + ), + "_logged_overlap_fallback": False, + } for name, value in attrs.items(): object.__setattr__(runner, name, value) return runner +def test_fused_path_under_the_kimi_k3_serving_config() -> None: + runner = _logic_runner() + + assert runner._use_fused_path() + assert runner._column_parallel_shardable() + + @pytest.mark.parametrize( - "shardable,pre_reduced,expect_sharded", + "override", [ - (True, False, True), - (True, True, False), - (False, False, False), - (False, True, False), + pytest.param({"tp_size": 1}, id="no-tp"), + pytest.param({"has_shared_experts": False}, id="no-shared-partial"), + pytest.param({"is_sequence_parallel": True}, id="sequence-parallel"), + pytest.param({"routed_scaling_factor": 2.0}, id="routed-scaling-factor"), + pytest.param({"fused_output_is_reduced": True}, id="already-reduced"), ], ) -def test_forward_shards_only_when_the_tail_is_valid( - shardable: bool, - pre_reduced: bool, - expect_sharded: bool, +def test_falls_back_to_native_path_when_fusion_is_unsafe(override: dict) -> None: + """Each of these makes the fused tail wrong, not merely slower, so the + runner must defer to the base combine.""" + runner = _logic_runner(**override) + + assert not runner._use_fused_path() + + +@pytest.mark.parametrize( + "override", + [ + pytest.param({"hidden": HIDDEN_SIZE + 2}, id="hidden-not-divisible-by-tp"), + pytest.param({"has_up_proj": False}, id="no-up-proj"), + ], +) +def test_not_column_parallel_shardable_but_still_fused(override: dict) -> None: + """These break only the shard tail; the fused overlap tail is still correct, + so the fused path stays on and selection avoids the column-parallel tail.""" + runner = _logic_runner(**override) + + assert runner._use_fused_path() + assert not runner._column_parallel_shardable() + tier = runner._select_tail_tier(torch.zeros(4096, 1)) + assert tier is ROCmLatentTailTier.ALLREDUCE_OVERLAP + + +def test_small_batches_pick_the_overlap_tail( monkeypatch: pytest.MonkeyPatch, ) -> None: - """The shard needs an un-reduced shared partial to accumulate into.""" - base_output, sharded_output = torch.zeros(1), torch.ones(1) - monkeypatch.setattr(moe_runner.MoERunner, "forward", lambda *a, **k: base_output) monkeypatch.setattr( - ROCmLatentMoERunner, "_fused_forward", lambda *a, **k: sharded_output + latent_moe_runner.envs, + "VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD", + 16, + raising=False, ) monkeypatch.setattr( - ROCmLatentMoERunner, "_fused_output_is_reduced", property(lambda _: pre_reduced) + latent_moe_runner.envs, + "VLLM_DISABLE_SHARED_EXPERTS_STREAM", + False, + raising=False, ) - runner = _runner(_tail_shardable=shardable) - - result = runner.forward(torch.zeros(1), torch.zeros(1)) - - assert result is (sharded_output if expect_sharded else base_output) + runner = _logic_runner() + # Small and shardable: overlap still wins under the threshold. + assert ( + runner._select_tail_tier(torch.zeros(8, 1)) + is ROCmLatentTailTier.ALLREDUCE_OVERLAP + ) + # Large and shardable: fold the up-projection into the reduce instead. + assert ( + runner._select_tail_tier(torch.zeros(64, 1)) + is ROCmLatentTailTier.COLUMN_PARALLEL + ) -@pytest.fixture -def build_runner(monkeypatch: pytest.MonkeyPatch): - """Construct through the real subclass ``__init__``, stubbing only the base.""" - def _base_init( - self, - *, - tp_size: int = 8, - hidden: int = HIDDEN_SIZE, - is_sequence_parallel: bool = False, - has_up_proj: bool = True, - has_shared_experts: bool = True, - routed_scaling_factor: float = 1.0, - ) -> None: - self.moe_config = SimpleNamespace( - tp_size=tp_size, is_sequence_parallel=is_sequence_parallel - ) - self.routed_output_transform = SimpleNamespace( - norm=None, - up_proj=SimpleNamespace(weight=torch.zeros(hidden, LATENT_SIZE)) - if has_up_proj - else None, - ) - self.routed_scaling_factor = routed_scaling_factor - self._shared_experts = object() if has_shared_experts else None - - monkeypatch.setattr(moe_runner.MoERunner, "__init__", _base_init) - return ROCmLatentMoERunner +def test_disabling_the_stream_forces_the_shard_tail( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + latent_moe_runner.envs, + "VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD", + 16, + raising=False, + ) + monkeypatch.setattr( + latent_moe_runner.envs, + "VLLM_DISABLE_SHARED_EXPERTS_STREAM", + True, + raising=False, + ) + runner = _logic_runner() + # Even a single-token batch takes the column-parallel tail when the overlap + # stream is disabled, as long as the up-projection is shardable. + assert ( + runner._select_tail_tier(torch.zeros(1, 1)) + is ROCmLatentTailTier.COLUMN_PARALLEL + ) -def test_shards_under_the_kimi_k3_serving_config(build_runner) -> None: - runner = build_runner() - assert runner._tail_shardable - assert runner._up_proj_shard_size == HIDDEN_SIZE // 8 +def _runner(**attrs) -> ROCmLatentMoERunner: + runner = object.__new__(ROCmLatentMoERunner) + for name, value in attrs.items(): + object.__setattr__(runner, name, value) + return runner -@pytest.mark.parametrize( - "override", - [ - pytest.param({"tp_size": 1}, id="no-tp"), - pytest.param({"hidden": HIDDEN_SIZE + 2}, id="hidden-not-divisible-by-tp"), - pytest.param({"has_up_proj": False}, id="no-up-proj"), - pytest.param({"has_shared_experts": False}, id="no-shared-partial"), - pytest.param({"is_sequence_parallel": True}, id="sequence-parallel"), - pytest.param({"routed_scaling_factor": 2.0}, id="routed-scaling-factor"), - ], -) -def test_falls_back_when_the_config_breaks_the_shard( - build_runner, override: dict +@pytest.mark.parametrize("fused_path,expect_fused", [(True, True), (False, False)]) +def test_forward_dispatches_to_fused_only_on_the_fused_path( + fused_path: bool, + expect_fused: bool, + monkeypatch: pytest.MonkeyPatch, ) -> None: - """Each of these makes the sharded tail wrong, not merely slower.""" - runner = build_runner(**override) + """forward() routes to the fused tail exactly when _use_fused_path holds.""" + base_output, fused_output = torch.zeros(1), torch.ones(1) + monkeypatch.setattr(moe_runner.MoERunner, "forward", lambda *a, **k: base_output) + monkeypatch.setattr( + ROCmLatentMoERunner, "_fused_forward", lambda *a, **k: fused_output + ) + monkeypatch.setattr(ROCmLatentMoERunner, "_use_fused_path", lambda self: fused_path) + runner = _runner() + + result = runner.forward(torch.zeros(1), torch.zeros(1)) - assert not runner._tail_shardable - assert runner._up_proj_shard_size == 0 + assert result is (fused_output if expect_fused else base_output) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index 1c34da3b6759..0b7e599b30b1 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -44,6 +44,7 @@ SharedExperts, SharedExpertsOrder, ) +from vllm.model_executor.layers.layernorm import RMSNorm from vllm.platforms import current_platform from vllm.utils.torch_utils import ( _USE_LAYERNAME, @@ -53,6 +54,16 @@ logger = init_logger(__name__) +# ROCm fused all-reduce + RMSNorm op, resolved once at import. Used to collapse +# the latent-MoE all-reduce and its subsequent RMSNorm (the routed output +# transform) into a single aiter kernel. None when aiter is unavailable. +try: + from vllm._aiter_ops import rocm_aiter_ops as _aiter_ops + + _aiter_fused_ar_rmsnorm = _aiter_ops.get_fused_allreduce_rmsnorm_op() +except Exception: + _aiter_fused_ar_rmsnorm = None + def register_layer_for_moe_forward_op( vllm_config: VllmConfig, @@ -431,6 +442,56 @@ def _maybe_reduce_shared_expert_output( shared_output = tensor_model_parallel_all_reduce(shared_output) return shared_output + def _get_zero_residual(self, x: torch.Tensor) -> torch.Tensor: + """Read-only zero residual buffer for the fused AR+RMSNorm kernel. + + The aiter op takes a residual even when there is nothing to add; a + zero residual makes it compute plain ``rmsnorm(all_reduce(x))``. The + buffer is cached and reused across calls, regrown only when the current + input needs more space or a different dtype/device. + """ + buf = getattr(self, "_zero_residual", None) + if ( + buf is None + or buf.numel() < x.numel() + or buf.dtype != x.dtype + or buf.device != x.device + ): + buf = torch.zeros(x.numel(), dtype=x.dtype, device=x.device) + self._zero_residual = buf + return buf[: x.numel()].view_as(x) + + def _can_fuse_ar_rmsnorm( + self, + fused_output: torch.Tensor, + fused_output_is_reduced: bool, + ) -> bool: + """Whether the AR + routed-output RMSNorm pair can use the aiter fused op. + + Only fuses transforms whose post-norm structure is known exactly: an + RMSNorm ``.norm`` followed by a callable ``.up_proj`` (the + ``KimiRoutedOutputTransform`` shape). Any other transform falls back to + the unfused all-reduce + transform path so no post-norm op is dropped. + """ + transform = self.routed_output_transform + norm = getattr(transform, "norm", None) + up_proj = getattr(transform, "up_proj", None) + return ( + _aiter_fused_ar_rmsnorm is not None + and current_platform.is_rocm() + and transform is not None + and isinstance(norm, RMSNorm) + and callable(up_proj) + and self.routed_scaling_factor == 1.0 + and not self.moe_config.is_sequence_parallel + and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1) + and not fused_output_is_reduced + and fused_output.is_cuda + and fused_output.dim() == 2 + and fused_output.is_contiguous() + and fused_output.dtype in (torch.bfloat16, torch.float16) + ) + def _maybe_reduce_routed_output_before_transform( self, fused_output: torch.Tensor, @@ -732,13 +793,31 @@ def forward( fused_output_is_reduced = self._fused_output_is_reduced # Latent routed output has to be reduced before output transform, - # because the transform may include non-linear normalization. - fused_output, fused_output_is_reduced = ( - self._maybe_reduce_routed_output_before_transform( - fused_output, - fused_output_is_reduced, + # because the transform may include non-linear normalization. On ROCm, + # when the transform is a plain RMSNorm + up_proj, the all-reduce and + # the RMSNorm collapse into a single aiter fused kernel; the up_proj + # then runs on the fused result. Guarded on routed_scaling_factor == 1.0 + # so the routed-scale step below stays a no-op and ordering is preserved. + fused_via_aiter = False + fused_transform = self.routed_output_transform + if self._can_fuse_ar_rmsnorm(fused_output, fused_output_is_reduced): + assert fused_transform is not None + transform = fused_transform + normed, _ = _aiter_fused_ar_rmsnorm( + input_=fused_output, + residual=self._get_zero_residual(fused_output), + weight=transform.norm.weight.to(fused_output.dtype), + epsilon=transform.norm.variance_epsilon, + ) + fused_output_is_reduced = True + fused_via_aiter = True + else: + fused_output, fused_output_is_reduced = ( + self._maybe_reduce_routed_output_before_transform( + fused_output, + fused_output_is_reduced, + ) ) - ) # If routed output is already reduced, reduce shared to match. # See note above re: the two all-reduce points. @@ -751,7 +830,11 @@ def forward( ) # Apply output transform (e.g. latent -> full dim) - fused_output = self.apply_routed_output_transform(fused_output) + if fused_via_aiter: + assert fused_transform is not None + fused_output, _ = fused_transform.up_proj(normed) + else: + fused_output = self.apply_routed_output_transform(fused_output) if shared_output is not None: result = shared_output + fused_output diff --git a/vllm/models/kimi_k3/amd/latent_moe_runner.py b/vllm/models/kimi_k3/amd/latent_moe_runner.py index db1c078b5d80..a3e0eb245bfe 100644 --- a/vllm/models/kimi_k3/amd/latent_moe_runner.py +++ b/vllm/models/kimi_k3/amd/latent_moe_runner.py @@ -1,22 +1,58 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from enum import IntEnum + import torch +import vllm.envs as envs from vllm.distributed import ( get_tensor_model_parallel_rank, tensor_model_parallel_all_reduce, ) from vllm.logger import init_logger -from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner, _unpack +from vllm.model_executor.layers.fused_moe.runner.moe_runner import ( + MoERunner, + _aiter_fused_ar_rmsnorm, + _unpack, +) +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.utils.multi_stream_utils import maybe_execute_in_parallel +from vllm.utils.torch_utils import aux_stream logger = init_logger(__name__) +class ROCmLatentTailTier(IntEnum): + """Which portable tail implementation the fused path runs, by token count. + + Mirrors CUDA's ``LatentTailTier`` minus its SM100-only tail-fusion tier, + which relies on tcgen05 kernels with no ROCm equivalent. Both tiers here + share the same replicated up-projection weight, so the choice is per batch + and needs no weight relayout. + """ + + # ``_overlap_allreduce_tail``. The portable default, up to + # VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD tokens: reduce the latent, + # up-project the full hidden dim from the replicated weight, and add the + # separately reduced shared output, hiding that shared all-reduce behind the + # up-projection GEMM on the aux stream. + ALLREDUCE_OVERLAP = 0 + + # ``_shard_up_proj_tail``. Prefill-sized: each rank up-projects only its + # hidden shard and accumulates into the shared partial, so the shared + # all-reduce also stitches the routed shards. Same two all-reduces as the + # overlap tier at 1/tp of the up-projection FLOPs; it gives up that tier's + # aux-stream overlap, since the reduce now has to follow the accumulate. + COLUMN_PARALLEL = 1 + + class ROCmLatentMoERunner(MoERunner): """MoE runner for latent MoE with a replicated routed up-projection. - Mirrors CUDA's LatentMoERunner, but currently only the up projection - -sharded path is implemented. (Tier 2) + Mirrors CUDA's ``LatentMoERunner`` for the two portable tail tiers. The + fused path (tp>1, un-reduced combine output, shared expert, no SP, unit + routed scale) dispatches over ``ROCmLatentTailTier`` by token count; see + that enum for what each tier does and when it applies. Native path: the replicated up-proj produces the full hidden dim on every rank, so the base runner combines routed + shared correctly at any TP size. @@ -29,29 +65,123 @@ def __init__( ) -> None: super().__init__(*args, **kwargs) - transform = self.routed_output_transform - up_proj = getattr(transform, "up_proj", None) - tp_size = self.moe_config.tp_size + # Overlap the shared-expert all-reduce with the tier-1 up-projection. + self._shared_ar_events = (torch.cuda.Event(), torch.cuda.Event()) + self._logged_column_parallel = False + self._logged_overlap_fallback = False - self._up_proj_shard_size = 0 - self._tail_shardable = ( - up_proj is not None - and tp_size > 1 - and up_proj.weight.shape[0] % tp_size == 0 + def _use_fused_path(self) -> bool: + # The fused path merges the latent and shared reductions into one + # all-reduce, so it needs actual TP parallelism, a shared expert (to + # concat), an un-reduced combine output, and no sequence parallelism. + # It also assumes a unit routed scale: the tiers do not apply + # routed_scaling_factor, so non-unit scales fall back to the base path. + return ( + self.moe_config.tp_size > 1 and self._shared_experts is not None + and not self._fused_output_is_reduced and not self.moe_config.is_sequence_parallel and self.routed_scaling_factor == 1.0 ) - if self._tail_shardable: - assert up_proj is not None - self._up_proj_shard_size = up_proj.weight.shape[0] // tp_size - else: + + def _column_parallel_shardable(self) -> bool: + transform = self.routed_output_transform + up_proj = getattr(transform, "up_proj", None) + return ( + up_proj is not None + and up_proj.weight.shape[0] % self.moe_config.tp_size == 0 + ) + + def _select_tail_tier( + self, + fused_output: torch.Tensor, + ) -> ROCmLatentTailTier: + num_tokens = fused_output.shape[0] + # tier 1 + if ( + num_tokens <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD + and not envs.VLLM_DISABLE_SHARED_EXPERTS_STREAM + ): + return ROCmLatentTailTier.ALLREDUCE_OVERLAP + # tier 2, when the up-projection rows divide evenly across ranks; + # otherwise the overlap tier is correct at any size. + if self._column_parallel_shardable(): + return ROCmLatentTailTier.COLUMN_PARALLEL + if not self._logged_overlap_fallback: + self._logged_overlap_fallback = True logger.warning_once( "K3 latent-MoE tail is not shardable under this config, " "falling back to the replicated up-projection.", scope="global", ) - self._logged_sharded_tail = False + return ROCmLatentTailTier.ALLREDUCE_OVERLAP + + def _allreduce_norm_latent_out( + self, + fused_output: torch.Tensor, + norm: RMSNorm, + ) -> torch.Tensor: + """All-reduce the latent routed output and RMSNorm it. + + On ROCm the pair collapses into a single aiter fused kernel when the + input is eligible; otherwise it falls back to a plain all-reduce + followed by the RMSNorm. The zero residual makes the fused kernel + compute ``rmsnorm(all_reduce(x))`` with nothing to add. + """ + if self.moe_config.tp_size == 1: + return norm(fused_output) + + if ( + _aiter_fused_ar_rmsnorm is not None + and fused_output.is_cuda + and fused_output.dim() == 2 + and fused_output.is_contiguous() + and fused_output.dtype in (torch.bfloat16, torch.float16) + ): + normed, _ = _aiter_fused_ar_rmsnorm( + input_=fused_output, + residual=self._get_zero_residual(fused_output), + weight=norm.weight.to(fused_output.dtype), + epsilon=norm.variance_epsilon, + ) + return normed + + return norm(tensor_model_parallel_all_reduce(fused_output)) + + def _overlap_allreduce_tail( + self, + fused_output: torch.Tensor, + shared_output: torch.Tensor, + trunc_size: int | None, + ) -> torch.Tensor: + """Tier 1: reduce the latent, up-project the full hidden dim from the + replicated weight, and add the separately reduced shared output. + + Small enough batches hide that shared all-reduce behind the up-projection + GEMM on the aux stream. + """ + transform = self.routed_output_transform + assert transform is not None + if transform.norm is not None: + fused_latent = self._allreduce_norm_latent_out(fused_output, transform.norm) + else: + fused_latent = tensor_model_parallel_all_reduce(fused_output) + + # Overlap the shared-expert all-reduce with the up-projection GEMM while + # the batch is small enough for it to pay off. + result, shared_output = maybe_execute_in_parallel( + lambda: torch.mm(fused_latent, transform.up_proj.weight.t()), + lambda: tensor_model_parallel_all_reduce(shared_output), + self._shared_ar_events[0], + self._shared_ar_events[1], + aux_stream(), + ) + result.add_(shared_output) + + # Output is already fully reduced; this only strips padding. + return self._maybe_reduce_final_output( + result, trunc_size, output_is_reduced=True + ) def _shard_up_proj_tail( self, @@ -62,8 +192,8 @@ def _shard_up_proj_tail( """ Tier 2: column-parallel up-projection folded into the final reduce. """ - if not self._logged_sharded_tail: - self._logged_sharded_tail = True + if not self._logged_column_parallel: + self._logged_column_parallel = True logger.info_once( "Kimi-K3 latent-MoE tail: up-projecting only this rank's " "hidden shard into the shared output.", @@ -73,13 +203,17 @@ def _shard_up_proj_tail( transform = self.routed_output_transform assert transform is not None - latent = tensor_model_parallel_all_reduce(fused_output) if transform.norm is not None: - latent = transform.norm(latent) + latent = self._allreduce_norm_latent_out(fused_output, transform.norm) + else: + latent = tensor_model_parallel_all_reduce(fused_output) - shard_size = self._up_proj_shard_size + weight = transform.up_proj.weight + shard_size = weight.shape[0] // self.moe_config.tp_size shard_start = get_tensor_model_parallel_rank() * shard_size - up_proj_shard = transform.up_proj.weight.narrow(0, shard_start, shard_size) + + # column-parallel + up_proj_shard = weight.narrow(0, shard_start, shard_size) hidden_shard = shared_output.narrow(-1, shard_start, shard_size) # hidden_shard += latent @ up_proj_shard.T, accumulated in the GEMM's @@ -97,7 +231,7 @@ def forward( input_ids: torch.Tensor | None = None, shared_experts_input: torch.Tensor | None = None, ) -> torch.Tensor: - if self._tail_shardable and not self._fused_output_is_reduced: + if self._use_fused_path(): return self._fused_forward( hidden_states, router_logits, input_ids, shared_experts_input ) @@ -145,8 +279,12 @@ def _fused_forward( if og_hidden_dim_pre_xform is not None: fused_output = fused_output[..., :og_hidden_dim_pre_xform] - result = self._shard_up_proj_tail( - fused_output, shared_output, og_hidden_dim_post_xform - ) + tier = self._select_tail_tier(fused_output) + if tier is ROCmLatentTailTier.ALLREDUCE_OVERLAP: + latent_tail = self._overlap_allreduce_tail + else: + latent_tail = self._shard_up_proj_tail + + result = latent_tail(fused_output, shared_output, og_hidden_dim_post_xform) return self._maybe_add_zero_expert_output(result)