[Perf][ROCm][Kimi-K3] Fuse all-reduce with RMSNorm in latent-MoE

tests/models/common/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

tests/models/common/test_amd_fused_allreduce_rms_norm.py

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""AITER fused AR+RMSNorm must equal RMSNorm(all_reduce(partial)).
+
+One TP2 spawn covers the fused kernel (4 tokens) and the unfused fallback
+(256 tokens). Hidden dim 3584 is outside AITER's old static-template set.
+"""
+
+import pytest
+import torch
+from torch.multiprocessing import spawn
+
+from tests.utils import ensure_current_vllm_config, init_test_distributed_environment
+from vllm.distributed import cleanup_dist_env_and_memory
+from vllm.distributed.communication_op import tensor_model_parallel_all_reduce
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.models.common.amd.ops.fused_allreduce_rms_norm import (
+    can_fuse_allreduce_rms_norm,
+    fused_allreduce_rms_norm_out,
+)
+from vllm.platforms import current_platform
+from vllm.utils.network_utils import get_open_port
+from vllm.utils.torch_utils import set_random_seed
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_rocm(),
+    reason="AITER fused AR+RMSNorm is only wired up on ROCm",
+)
+
+HIDDEN_SIZE = 3584
+EPS = 1e-5
+DTYPE = torch.bfloat16
+WORLD_SIZE = 2
+
+
+def _check_matches_unfused(norm: RMSNorm, partial: torch.Tensor) -> None:
+    ref = norm(tensor_model_parallel_all_reduce(partial.clone()))
+    out = fused_allreduce_rms_norm_out(partial.clone(), norm)
+    torch.accelerator.synchronize()
+    torch.testing.assert_close(out, ref, atol=2e-2, rtol=2e-2)
+
+
+@ensure_current_vllm_config()
+def _worker_fused_ar_rms(local_rank, world_size, port, seed):
+    device = torch.device(f"cuda:{local_rank}")
+    torch.accelerator.set_device_index(device)
+    init_test_distributed_environment(
+        world_size, 1, local_rank, port, local_rank=local_rank
+    )
+
+    set_random_seed(seed)
+    norm = RMSNorm(HIDDEN_SIZE, eps=EPS).to(device=device, dtype=DTYPE)
+    with torch.no_grad():
+        norm.weight.normal_(mean=1.0, std=0.1)
+
+    torch.manual_seed(seed + local_rank)
+    fused = torch.randn(4, HIDDEN_SIZE, dtype=DTYPE, device=device)
+    fallback = torch.randn(256, HIDDEN_SIZE, dtype=DTYPE, device=device)
+
+    _check_matches_unfused(norm, fused)
+    assert not can_fuse_allreduce_rms_norm(fallback)
+    _check_matches_unfused(norm, fallback)
+
+    cleanup_dist_env_and_memory()
+
+
+def test_fused_allreduce_rms_norm_out():
+    if current_platform.device_count() < WORLD_SIZE:
+        pytest.skip(f"Need >= {WORLD_SIZE} GPUs")
+    spawn(
+        _worker_fused_ar_rms,
+        args=(WORLD_SIZE, str(get_open_port()), 42),
+        nprocs=WORLD_SIZE,
+        join=True,
+    )
+
+
+def test_gate_is_closed_without_aiter_custom_ar(monkeypatch):
+    from vllm import _aiter_ops
+
+    monkeypatch.setattr(
+        _aiter_ops.rocm_aiter_ops, "is_custom_all_reduce_enabled", lambda: False
+    )
+    hidden = torch.empty(4, HIDDEN_SIZE, dtype=DTYPE)
+    assert not can_fuse_allreduce_rms_norm(hidden)

vllm/models/common/amd/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/models/common/amd/ops/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/models/common/amd/ops/fused_allreduce_rms_norm.py

@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""AITER fused all-reduce + RMSNorm (no residual)."""
+
+import torch
+
+from vllm._aiter_ops import rocm_aiter_ops
+from vllm.distributed import (
+    get_tensor_model_parallel_world_size,
+    tensor_model_parallel_all_reduce,
+)
+from vllm.model_executor.layers.layernorm import RMSNorm
+
+# Zero residual, cached per shape so CUDA-graph captures keep a stable pointer.
+_zero_residuals: dict[tuple[torch.Size, torch.dtype, torch.device], torch.Tensor] = {}
+
+
+def _zero_residual(hidden_states: torch.Tensor) -> torch.Tensor:
+    key = (hidden_states.shape, hidden_states.dtype, hidden_states.device)
+    buf = _zero_residuals.get(key)
+    if buf is None:
+        buf = torch.zeros_like(hidden_states)
+        _zero_residuals[key] = buf
+    return buf
+
+
+def can_fuse_allreduce_rms_norm(hidden_states: torch.Tensor) -> bool:
+    """True when AITER's one-stage fused AR+RMSNorm covers this tensor."""
+    if not rocm_aiter_ops.is_custom_all_reduce_enabled():
+        return False
+    # gfx1250 custom AR has no fused_allreduce_rmsnorm:
+    # https://github.com/ROCm/aiter/blob/v0.1.22/csrc/include/custom_all_reduce_gfx1250.h
+    # https://github.com/ROCm/aiter/blob/v0.1.22/csrc/include/custom_all_reduce.h#L69
+    from vllm.platforms.rocm import on_gfx1250
+
+    if on_gfx1250():
+        return False
+    if hidden_states.dim() != 2 or not hidden_states.is_contiguous():
+        return False
+    if hidden_states.dtype not in (torch.bfloat16, torch.float16):
+        return False
+
+    aiter_ar = rocm_aiter_ops.get_aiter_allreduce()
+    if aiter_ar is None or aiter_ar.disabled:
+        return False
+    if getattr(aiter_ar.aiter_ca, "_is_gfx1250", False):
+        return False
+    if not aiter_ar.supports_dynamic_hidden_dim:
+        return False
+    if not aiter_ar.should_custom_ar(hidden_states):
+        return False
+    return aiter_ar.use_1stage_fused_ar_rms(hidden_states)
+
+
+def fused_allreduce_rms_norm_out(
+    hidden_states: torch.Tensor,
+    norm: RMSNorm,
+) -> torch.Tensor:
+    """Equivalent to ``norm(all_reduce(hidden_states))``."""
+    if get_tensor_model_parallel_world_size() == 1:
+        return norm(hidden_states)
+
+    if not can_fuse_allreduce_rms_norm(hidden_states):
+        return norm(tensor_model_parallel_all_reduce(hidden_states))
+
+    norm_out, _ = rocm_aiter_ops.get_fused_allreduce_rmsnorm_op()(
+        input_=hidden_states,
+        residual=_zero_residual(hidden_states),
+        weight=norm.weight,
+        epsilon=norm.variance_epsilon,
+        gemma_norm=False,
+    )
+    return norm_out

vllm/models/kimi_k3/amd/latent_moe_runner.py

@@ -11,6 +11,9 @@
 )
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner
+from vllm.models.common.amd.ops.fused_allreduce_rms_norm import (
+    fused_allreduce_rms_norm_out,
+)
 
 logger = init_logger(__name__)
 
@@ -23,6 +26,11 @@ class ROCmLatentMoERunner(MoERunner):
 
     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.
+
+    The latent all-reduce is fused with the following RMSNorm via AITER's
+    one-stage custom AR when that kernel's gate admits the tensor; see
+    ``vllm.models.common.amd.ops.fused_allreduce_rms_norm``. Everything
+    else keeps the unfused all-reduce.
     """
 
     def __init__(
@@ -74,9 +82,10 @@ 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 = fused_allreduce_rms_norm_out(fused_output, transform.norm)
+        else:
+            latent = tensor_model_parallel_all_reduce(fused_output)
 
         shard_size = self._up_proj_shard_size
         shard_start = get_tensor_model_parallel_rank() * shard_size