diff --git a/tests/models/kimi_k3/test_amd_prefill_sp.py b/tests/models/kimi_k3/test_amd_prefill_sp.py new file mode 100644 index 000000000000..17ff5d7c90e1 --- /dev/null +++ b/tests/models/kimi_k3/test_amd_prefill_sp.py @@ -0,0 +1,488 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Token-sharded Kimi-K3 prefill on ROCm must equal the replicated path. + +Each rank ends up holding only its own token rows, so the multi-GPU checks +compare every rank's shard with the matching rows of the all-reduce result. +""" + +from collections.abc import Callable +from functools import partial +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.nn.functional as F +from torch.multiprocessing import spawn + +from tests.utils import ( + ensure_current_vllm_config, + init_test_distributed_environment, + multi_gpu_test, +) +from vllm.config import VllmConfig, set_current_vllm_config +from vllm.distributed import get_tp_group +from vllm.forward_context import set_forward_context +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ColumnParallelLinear, ReplicatedLinear +from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8DynamicTokenSym, +) +from vllm.models.kimi_k3.amd import sp +from vllm.models.kimi_k3.amd.latent_moe_runner import ROCmLatentMoERunner +from vllm.models.kimi_k3.amd.linear import ( + KimiDecoderLayer, + KimiMLP, + KimiMoE, + KimiRoutedOutputTransform, +) +from vllm.platforms import current_platform +from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig +from vllm.utils.network_utils import get_open_port +from vllm.v1.worker.workspace import init_workspace_manager + +pytestmark = pytest.mark.skipif( + not current_platform.is_rocm(), + reason="Token-sharded Kimi-K3 prefill is only wired up on ROCm", +) + +HIDDEN_SIZE = 7168 +LATENT_SIZE = 3584 +TOP_K = 16 +NUM_EXPERTS = 896 +EPS = 1e-5 +DTYPE = torch.bfloat16 + + +def _build_transform( + device: torch.device, tp_world: int, rank: int, row_sharded: bool +) -> tuple[KimiRoutedOutputTransform, torch.Tensor]: + """Returns the transform and the full up-projection weight.""" + norm = RMSNorm(LATENT_SIZE, eps=EPS).to(device=device, dtype=DTYPE) + torch.manual_seed(0) + norm.weight.data.copy_(1 + 0.1 * torch.randn_like(norm.weight)) + weight = torch.randn(HIDDEN_SIZE, LATENT_SIZE, device=device, dtype=DTYPE) + weight /= LATENT_SIZE**0.5 + + up_proj: ColumnParallelLinear | ReplicatedLinear + if row_sharded: + # The layout KimiMoE uses under TP: each rank stores its output rows. + up_proj = ColumnParallelLinear( + LATENT_SIZE, + HIDDEN_SIZE, + bias=False, + gather_output=True, + params_dtype=DTYPE, + prefix="routed_expert_up_proj", + ).to(device=device) + rows = HIDDEN_SIZE // tp_world + up_proj.weight.data.copy_(weight[rank * rows : (rank + 1) * rows]) + else: + up_proj = ReplicatedLinear( + LATENT_SIZE, + HIDDEN_SIZE, + bias=False, + params_dtype=DTYPE, + prefix="routed_expert_up_proj", + ).to(device=device) + up_proj.weight.data.copy_(weight) + return KimiRoutedOutputTransform(norm, up_proj, row_sharded=row_sharded), weight + + +def _tail_runner( + transform: KimiRoutedOutputTransform, tp_world: int +) -> ROCmLatentMoERunner: + runner = object.__new__(ROCmLatentMoERunner) + attrs = { + "routed_output_transform": transform, + "_up_proj_preshard": transform.row_sharded, + "_up_proj_shard_size": HIDDEN_SIZE // tp_world, + "_logged_sharded_tail": False, + "moe_config": SimpleNamespace( + tp_size=tp_world, + ep_size=1, + is_sequence_parallel=False, + skip_final_all_reduce=False, + ), + } + for name, value in attrs.items(): + object.__setattr__(runner, name, value) + return runner + + +def _all_reduced(tensor: torch.Tensor, group) -> torch.Tensor: + reduced = tensor.clone() + dist.all_reduce(reduced, group=group) + return reduced + + +def _check_sharded_tail(device: torch.device, tp_world: int, rank: int) -> None: + """For both up-projection layouts, the column-parallel tail matches the + replicated result, and its token-sharded variant returns this rank's rows.""" + for row_sharded in (True, False): + _check_tail_layout(device, tp_world, rank, row_sharded) + + +def _check_tail_layout( + device: torch.device, tp_world: int, rank: int, row_sharded: bool +) -> None: + transform, weight = _build_transform(device, tp_world, rank, row_sharded) + assert transform.norm is not None + norm_weight = transform.norm.weight + runner = _tail_runner(transform, tp_world) + group = get_tp_group().device_group + + for iteration, num_tokens in enumerate((tp_world, 2 * tp_world, 64)): + torch.manual_seed(100 * iteration + rank + 1) + routed = 0.01 * torch.randn(num_tokens, LATENT_SIZE, device=device, dtype=DTYPE) + shared = torch.randn(num_tokens, HIDDEN_SIZE, device=device, dtype=DTYPE) + + expected = F.linear( + F.rms_norm( + _all_reduced(routed, group), + (LATENT_SIZE,), + norm_weight, + EPS, + ), + weight, + ) + expected.add_(_all_reduced(shared, group)) + rows = num_tokens // tp_world + + full = runner._shard_up_proj_tail(routed.clone(), shared.clone(), None) + actual = runner._shard_up_proj_tail(routed, shared, None, token_sharded=True) + + torch.testing.assert_close(full, expected, atol=8e-2, rtol=3e-2) + assert actual.shape == (rows, HIDDEN_SIZE) + torch.testing.assert_close( + actual, expected[rank * rows : (rank + 1) * rows], atol=8e-2, rtol=3e-2 + ) + + +def _check_routing_gather(device: torch.device, tp_world: int, rank: int) -> None: + """Packed routing survives the gather bit-exactly and in rank order.""" + rows = 5 + gen = torch.Generator(device=device).manual_seed(7) + weights = torch.rand(tp_world * rows, TOP_K, device=device, generator=gen) + ids = torch.randint( + NUM_EXPERTS, (tp_world * rows, TOP_K), device=device, generator=gen + ).to(torch.int32) + shard = slice(rank * rows, (rank + 1) * rows) + + got_weights, got_ids = sp.all_gather_routing(weights[shard], ids[shard]) + + assert got_weights.dtype == weights.dtype and got_ids.dtype == ids.dtype + torch.testing.assert_close(got_weights, weights, atol=0, rtol=0) + torch.testing.assert_close(got_ids, ids, atol=0, rtol=0) + + +def _check_padded_gather_scatter( + device: torch.device, tp_world: int, rank: int +) -> None: + """Token counts that do not divide TP round-trip through the padded shards.""" + group = get_tp_group().device_group + for num_tokens in (tp_world * 8 - 3, tp_world * 8 + 1, tp_world * 8): + padded = -(-num_tokens // tp_world) * tp_world + rows = padded // tp_world + gen = torch.Generator(device=device).manual_seed(num_tokens) + full = torch.randn(padded, 64, device=device, generator=gen) + full[num_tokens:] = 0 + + gathered = sp.gather_tokens(full[rank * rows : (rank + 1) * rows], num_tokens) + torch.manual_seed(rank + 1) + partial = torch.randn(num_tokens, 64, device=device) + scattered = sp.scatter_tokens(partial) + + torch.testing.assert_close(gathered, full[:num_tokens], atol=0, rtol=0) + expected = F.pad(_all_reduced(partial, group), (0, 0, 0, padded - num_tokens)) + torch.testing.assert_close( + scattered, expected[rank * rows : (rank + 1) * rows], atol=1e-5, rtol=1e-5 + ) + + data = full.to(current_platform.fp8_dtype()) + scale = torch.rand(padded, 1, device=device, generator=gen) + mine = slice(rank * rows, (rank + 1) * rows) + quantized = QuantizedActivation( + data[mine], scale[mine], DTYPE, torch.Size((rows, 64)), kFp8DynamicTokenSym + ) + for keep, expected_rows in ((num_tokens, num_tokens), (None, padded)): + got = sp.gather_activation(quantized, keep) + assert isinstance(got, QuantizedActivation) + assert got.orig_shape == (expected_rows, 64) + assert got.orig_dtype == DTYPE and got.quant_key == kFp8DynamicTokenSym + torch.testing.assert_close( + got.data.view(torch.uint8), + data[:expected_rows].view(torch.uint8), + atol=0, + rtol=0, + ) + torch.testing.assert_close(got.scale, scale[:expected_rows], atol=0, rtol=0) + + +MLP_HIDDEN_SIZE = 512 +# At least the 256 per rank that KimiMoE would otherwise pad up to. +MLP_INTERMEDIATE_PER_RANK = 256 + + +def _build_mlp( + device: torch.device, tp_world: int, kind: str +) -> tuple[KimiMLP | KimiMoE, VllmConfig]: + """A small dense KimiMLP or latent KimiMoE with the same weights on every + rank's view of the full layer. + + ``kind`` is "dense", "moe" (token-shardable tail) or "moe_no_shared" + (without shared experts the tail is not token-shardable, so the layer + takes the all-gather fallback). + """ + init_workspace_manager(device) + vllm_config = VllmConfig() + vllm_config.compilation_config.static_forward_context = {} + torch.set_default_dtype(DTYPE) + mlp: KimiMLP | KimiMoE + with set_current_vllm_config(vllm_config), torch.device(device): + if kind == "dense": + mlp = KimiMLP( + MLP_HIDDEN_SIZE, + MLP_INTERMEDIATE_PER_RANK * tp_world, + "silu", + prefix="model.layers.0.mlp", + ) + else: + config = KimiLinearConfig( + hidden_size=MLP_HIDDEN_SIZE, + hidden_act="silu", + rms_norm_eps=EPS, + moe_intermediate_size=MLP_INTERMEDIATE_PER_RANK * tp_world, + num_experts=32, + num_experts_per_token=4, + num_shared_experts=None if kind == "moe_no_shared" else 1, + routed_expert_hidden_size=256, + latent_moe_use_norm=True, + use_grouped_topk=True, + num_expert_group=1, + topk_group=1, + ) + mlp = KimiMoE(config, prefix="model.layers.1.mlp", layer_idx=1) + assert mlp.experts.token_shardable is (kind == "moe") + + # Same seed on every rank: replicated weights match across ranks, and each + # rank's TP shard is just a slice of some consistent full weight. + torch.manual_seed(0) + with torch.no_grad(): + for param in mlp.parameters(): + if param.dim() == 1: + param.copy_(1 + 0.1 * torch.randn_like(param)) + else: + param.copy_(torch.randn_like(param) / param.size(-1) ** 0.5) + for module in mlp.modules(): + quant_method = getattr(module, "quant_method", None) + if isinstance(quant_method, QuantizeMethodBase): + quant_method.process_weights_after_loading(module) + return mlp, vllm_config + + +def _check_mlp_rows( + device: torch.device, + tp_world: int, + rank: int, + kind: str, + run_sharded: Callable[[SimpleNamespace, torch.Tensor], torch.Tensor], +) -> None: + """``run_sharded`` on this rank's token shard returns this rank's rows of + the unsharded layer. The token count does not divide TP, so the last + rank's shard ends in zero padding rows.""" + mlp, vllm_config = _build_mlp(device, tp_world, kind) + rows = 6 + num_tokens = tp_world * rows - 3 + torch.manual_seed(1) + hidden = torch.randn(num_tokens, MLP_HIDDEN_SIZE, device=device, dtype=DTYPE) + padded = F.pad(hidden, (0, 0, 0, tp_world * rows - num_tokens)) + mine = slice(rank * rows, (rank + 1) * rows) + real = max(0, min(rows, num_tokens - rank * rows)) + layer = SimpleNamespace( + mlp=mlp, _output_unreduced=KimiDecoderLayer._output_unreduced + ) + with set_forward_context(None, vllm_config, num_tokens=num_tokens): + expected = mlp(hidden) + actual = run_sharded(layer, padded[mine]) + + assert actual.shape == (rows, MLP_HIDDEN_SIZE) + torch.testing.assert_close( + actual[:real], expected[mine][:real], atol=8e-2, rtol=3e-2 + ) + + +def _check_shard_routed_moe(device: torch.device, tp_world: int, rank: int) -> None: + """A latent KimiMoE routed on the token shard, with the gathered top-k + replacing routing.""" + _check_mlp_rows( + device, tp_world, rank, "moe", KimiDecoderLayer._run_moe_sp_shard_routed + ) + + +def _check_mlp_sp(device: torch.device, tp_world: int, rank: int, kind: str) -> None: + """The decoder layer's sharded MLP step: reduce-scatter of the unreduced + down_proj for a dense MLP, all-gather then this rank's rows for an MoE + whose tail is not token-shardable.""" + _check_mlp_rows(device, tp_world, rank, kind, KimiDecoderLayer._run_mlp_sp) + + +_CHECKS: dict[str, Callable[[torch.device, int, int], None]] = { + "sharded_tail": _check_sharded_tail, + "routing_gather": _check_routing_gather, + "padded_gather_scatter": _check_padded_gather_scatter, + "shard_routed_moe": _check_shard_routed_moe, + "mlp_sp_dense": partial(_check_mlp_sp, kind="dense"), + "mlp_sp_moe_no_shared": partial(_check_mlp_sp, kind="moe_no_shared"), +} + + +def _worker(local_rank: int, world_size: int, port: str, check: str) -> None: + device = torch.device(f"cuda:{local_rank}") + torch.accelerator.set_device_index(device) + with ensure_current_vllm_config(): + init_test_distributed_environment( + world_size, 1, local_rank, port, local_rank=local_rank + ) + _CHECKS[check](device, world_size, local_rank) + + +def _run_ranks(check: str, tp_size: int) -> None: + spawn( + _worker, + args=(tp_size, str(get_open_port()), check), + nprocs=tp_size, + join=True, + ) + + +@multi_gpu_test(num_gpus=4) +def test_sharded_tail_tp4_matches_replicated_rows() -> None: + _run_ranks("sharded_tail", 4) + + +@multi_gpu_test(num_gpus=8) +def test_sharded_tail_tp8_matches_replicated_rows() -> None: + _run_ranks("sharded_tail", 8) + + +@multi_gpu_test(num_gpus=4) +def test_routing_gather_tp4_is_exact() -> None: + _run_ranks("routing_gather", 4) + + +@multi_gpu_test(num_gpus=4) +def test_padded_gather_scatter_tp4() -> None: + _run_ranks("padded_gather_scatter", 4) + + +@multi_gpu_test(num_gpus=8) +def test_padded_gather_scatter_tp8() -> None: + _run_ranks("padded_gather_scatter", 8) + + +def _run_shard_routed_moe( + tp_size: int, use_aiter: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("VLLM_ROCM_USE_AITER", str(int(use_aiter))) + monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", str(int(use_aiter))) + _run_ranks("shard_routed_moe", tp_size) + + +@multi_gpu_test(num_gpus=4) +@pytest.mark.parametrize("use_aiter", [False, True]) +def test_shard_routed_moe_tp4_matches_unsharded_rows( + use_aiter: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + _run_shard_routed_moe(4, use_aiter, monkeypatch) + + +@multi_gpu_test(num_gpus=8) +@pytest.mark.parametrize("use_aiter", [False, True]) +def test_shard_routed_moe_tp8_matches_unsharded_rows( + use_aiter: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + _run_shard_routed_moe(8, use_aiter, monkeypatch) + + +@multi_gpu_test(num_gpus=4) +@pytest.mark.parametrize("kind", ["dense", "moe_no_shared"]) +def test_mlp_sp_tp4_matches_unsharded_rows(kind: str) -> None: + _run_ranks(f"mlp_sp_{kind}", 4) + + +class _Router(torch.nn.Module): + def select_experts(self, hidden_states, router_logits, **kwargs): + return "routed" + + +def test_pre_routed_overrides_then_restores() -> None: + router = _Router() + topk = (torch.ones(2, TOP_K), torch.zeros(2, TOP_K, dtype=torch.int32)) + + with sp.pre_routed(router, topk): + assert router.select_experts(hidden_states=None, router_logits=None) is topk + assert router.select_experts(None, None) == "routed" + assert "select_experts" not in router.__dict__ + + with pytest.raises(RuntimeError), sp.pre_routed(router, topk): + raise RuntimeError + assert router.select_experts(None, None) == "routed" + + +@pytest.mark.parametrize("calls", [0, 2]) +def test_pre_routed_requires_exactly_one_route(calls: int) -> None: + router = _Router() + with pytest.raises(AssertionError), sp.pre_routed(router, ("w", "i")): + for _ in range(calls): + router.select_experts(None, None) + assert router.select_experts(None, None) == "routed" + + +def test_pre_routed_keeps_an_instance_override() -> None: + router = _Router() + router.__dict__["select_experts"] = lambda *args, **kwargs: "patched" + + with sp.pre_routed(router, ("w", "i")): + assert router.select_experts(None, None) == ("w", "i") + assert router.select_experts(None, None) == "patched" + + +@pytest.mark.parametrize( + "min_tokens,tp_size,num_tokens,has_residual,has_aux,last_rank,dp_size,observed,expected", + [ + (1024, 8, 16384, False, False, True, 1, False, True), + (1024, 8, 1024, False, False, True, 1, False, True), + (0, 8, 16384, False, False, True, 1, False, False), + (1024, 1, 16384, False, False, True, 1, False, False), + (1024, 8, 1016, False, False, True, 1, False, False), + (1024, 8, 16381, False, False, True, 1, False, True), + (1024, 8, 16384, True, False, True, 1, False, False), + (1024, 8, 16384, False, True, True, 1, False, False), + (1024, 8, 16384, False, False, False, 1, False, False), + (1024, 8, 16384, False, False, True, 2, False, False), + (1024, 8, 16384, False, False, True, 1, True, False), + ], +) +def test_should_shard( + monkeypatch: pytest.MonkeyPatch, + min_tokens: int, + tp_size: int, + num_tokens: int, + has_residual: bool, + has_aux: bool, + last_rank: bool, + dp_size: int, + observed: bool, + expected: bool, +) -> None: + monkeypatch.setenv("VLLM_KIMI_K3_AMD_PREFILL_SP_MIN_TOKENS", str(min_tokens)) + monkeypatch.setattr(sp, "get_tensor_model_parallel_world_size", lambda: tp_size) + monkeypatch.setattr( + sp, "get_pp_group", lambda: SimpleNamespace(is_last_rank=last_rank) + ) + monkeypatch.setattr(sp, "get_dp_group", lambda: SimpleNamespace(world_size=dp_size)) + assert sp.should_shard(num_tokens, has_residual, has_aux, observed) is expected diff --git a/vllm/envs.py b/vllm/envs.py index 857085ee1b26..d558096aa643 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -217,6 +217,7 @@ VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT: bool = False VLLM_KIMI_K3_AUX_ATTN_RES_STREAM: bool = False VLLM_KIMI_K3_GEMM_AR: bool = True + VLLM_KIMI_K3_AMD_PREFILL_SP_MIN_TOKENS: int = 0 VLLM_ENABLE_GEMM_RS: bool = False VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER: bool = True VLLM_USE_FLASHINFER_MOE_INT4: bool = False @@ -1693,6 +1694,16 @@ def _resolve_rust_cli_path() -> str | None: # Use the SM100 BF16 GEMM-AR kernel for eligible Kimi-K3 row-parallel # attention projections. All TP ranks must belong to one NVLink domain. "VLLM_KIMI_K3_GEMM_AR": lambda: bool(int(os.getenv("VLLM_KIMI_K3_GEMM_AR", "1"))), + # Kimi-K3 on ROCm only (TP, no PP). Eager steps with at least this many + # tokens run the decoder with a token-sharded residual stream (padded to a + # multiple of the TP size): each attention / MLP output all-reduce + # becomes a reduce-scatter plus an all-gather before the next block, the + # attention-residual mixing runs on the local shard, and latent-MoE routing + # runs once per shard instead of on every rank. Speeds up long prefills; + # CUDA-graph decode is unaffected. 0 disables it; 1024 is a good value. + "VLLM_KIMI_K3_AMD_PREFILL_SP_MIN_TOKENS": lambda: int( + os.getenv("VLLM_KIMI_K3_AMD_PREFILL_SP_MIN_TOKENS", "0") + ), # Fuse eligible sequence-parallel row-parallel projections with their TP # reduce-scatter using the SM100 BF16/MXFP8 GEMM-RS kernel (Kimi-K3 # attention/shared-expert projections, DeepSeek-V4.1 ``wo_b``). All TP diff --git a/vllm/models/kimi_k3/amd/latent_moe_runner.py b/vllm/models/kimi_k3/amd/latent_moe_runner.py index f6508d6b559f..6455b39a2637 100644 --- a/vllm/models/kimi_k3/amd/latent_moe_runner.py +++ b/vllm/models/kimi_k3/amd/latent_moe_runner.py @@ -8,8 +8,12 @@ get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce, + tensor_model_parallel_reduce_scatter, ) from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.router.zero_expert_router import ( + ZeroExpertRouter, +) from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner logger = init_logger(__name__) @@ -61,11 +65,37 @@ def __init__( ) self._logged_sharded_tail = False + @property + def token_shardable(self) -> bool: + """Whether ``forward`` supports ``token_sharded=True``. Monolithic + kernels route from ``router_logits`` themselves, so they cannot take the + all-gathered top-k.""" + return ( + self._tail_shardable + and not self.is_monolithic + and not self._fused_output_is_reduced + and not self.moe_config.skip_final_all_reduce + and not isinstance(self.router, ZeroExpertRouter) + ) + + def sp_route_shard( + self, hidden_shard: torch.Tensor, router_logits: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Latent projection and top-k for this rank's token shard only.""" + latent, _ = self.apply_routed_input_transform(hidden_shard) + topk_weights, topk_ids = self.router.select_experts( + hidden_states=latent, + router_logits=router_logits, + topk_indices_dtype=self._quant_method.topk_indices_dtype, + ) + return latent, topk_weights, topk_ids + def _shard_up_proj_tail( self, fused_output: torch.Tensor, shared_output: torch.Tensor, trunc_size: int | None, + token_sharded: bool = False, ) -> torch.Tensor: """Tier 2: column-parallel up-projection folded into the final reduce.""" if not self._logged_sharded_tail: @@ -97,6 +127,10 @@ def _shard_up_proj_tail( # beta-add epilogue so folding in the shared partial costs no kernel. hidden_shard.addmm_(latent, up_proj_shard.t()) + if token_sharded: + # Token-sharded residual stream: keep only this rank's tokens. + out = tensor_model_parallel_reduce_scatter(shared_output, dim=0) + return out[..., :trunc_size] if trunc_size is not None else out return self._maybe_reduce_final_output( shared_output, trunc_size, output_is_reduced=False ) @@ -107,10 +141,18 @@ def forward( router_logits: torch.Tensor, input_ids: torch.Tensor | None = None, shared_experts_input: torch.Tensor | None = None, + token_sharded: bool = False, ) -> torch.Tensor: + """``token_sharded`` reduce-scatters the output over tokens instead of + all-reducing it; it requires ``token_shardable``.""" + assert not token_sharded or self.token_shardable if self._tail_shardable and not self._fused_output_is_reduced: return self._fused_forward( - hidden_states, router_logits, input_ids, shared_experts_input + hidden_states, + router_logits, + input_ids, + shared_experts_input, + token_sharded, ) return super().forward( hidden_states, router_logits, input_ids, shared_experts_input @@ -122,6 +164,7 @@ def _fused_forward( router_logits: torch.Tensor, input_ids: torch.Tensor | None, shared_experts_input: torch.Tensor | None, + token_sharded: bool, ) -> torch.Tensor: # When the caller pre-applies the routed input transform outside the # runner (e.g. to overlap it on a separate stream), it passes the @@ -156,7 +199,7 @@ def _fused_forward( fused_output = fused_output[..., :og_hidden_dim_pre_xform] result = self._shard_up_proj_tail( - fused_output, shared_output, og_hidden_dim_post_xform + fused_output, shared_output, og_hidden_dim_post_xform, token_sharded ) return self._maybe_add_zero_expert_output(result) diff --git a/vllm/models/kimi_k3/amd/linear.py b/vllm/models/kimi_k3/amd/linear.py index 0671a69d293e..d6e146fbf65c 100644 --- a/vllm/models/kimi_k3/amd/linear.py +++ b/vllm/models/kimi_k3/amd/linear.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Iterable +import contextlib +from collections.abc import Iterable, Iterator from typing import Any, cast import torch @@ -11,7 +12,10 @@ from vllm.config import CacheConfig, VllmConfig from vllm.distributed import ( get_pp_group, + get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, + tensor_model_parallel_all_gather, + tensor_model_parallel_reduce_scatter, ) from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMul, SituAndMul @@ -71,6 +75,8 @@ make_layers, maybe_prefix, ) +from vllm.models.common.ops.sequence_parallel import sp_shard +from vllm.models.kimi_k3.amd import sp from vllm.models.kimi_k3.amd.kda import KimiK3DeltaAttention from vllm.models.kimi_k3.amd.latent_moe_runner import ROCmLatentMoERunner from vllm.models.kimi_k3.amd.mla import KimiK3MultiHeadLatentAttentionWrapper @@ -654,12 +660,73 @@ def _run_self_attn( positions=positions, ) + @staticmethod + @contextlib.contextmanager + def _output_unreduced(proj: RowParallelLinear) -> Iterator[None]: + # Row-parallel outputs stay partial so the caller can reduce-scatter. + assert isinstance(proj, RowParallelLinear) and proj.reduce_results + proj.reduce_results = False + try: + yield + finally: + proj.reduce_results = True + + def _run_self_attn_sp( + self, + positions: torch.Tensor, + hidden_shard: torch.Tensor | QuantizedActivation, + num_tokens: int, + ) -> torch.Tensor: + hidden_states = sp.gather_activation(hidden_shard, num_tokens) + with self._output_unreduced(self.self_attn.o_proj): + out = self._run_self_attn(positions, hidden_states) + return sp.scatter_tokens(out) + + def _run_moe_sp_shard_routed(self, hidden_shard: torch.Tensor) -> torch.Tensor: + runner = self.mlp.experts + assert isinstance(runner, ROCmLatentMoERunner) + router_logits, _ = self.mlp.gate(hidden_shard) + latent, topk_weights, topk_ids = runner.sp_route_shard( + hidden_shard, router_logits + ) + hidden_states = tensor_model_parallel_all_gather(hidden_shard, dim=0) + latent = tensor_model_parallel_all_gather(latent.contiguous(), dim=0) + topk = sp.all_gather_routing(topk_weights, topk_ids) + # Only read by select_experts, which returns the gathered top-k instead. + logits = router_logits.new_empty((latent.size(0), router_logits.size(1))) + with sp.pre_routed(runner.router, topk): + return runner( + hidden_states=latent, + router_logits=logits, + shared_experts_input=hidden_states, + token_sharded=True, + ) + + def _run_mlp_sp( + self, hidden_shard: torch.Tensor | QuantizedActivation + ) -> torch.Tensor: + # MLPs are per-token, so they run on the TP-padded token set as is. + if isinstance(self.mlp, KimiMoE): + assert isinstance(hidden_shard, torch.Tensor) + if self.mlp.use_latent_moe and self.mlp.experts.token_shardable: + return self._run_moe_sp_shard_routed(hidden_shard) + hidden_states = tensor_model_parallel_all_gather(hidden_shard, dim=0) + out = self.mlp(hidden_states) + # The MoE output is fully reduced; keep this rank's tokens. + shard = hidden_shard.size(0) + return out.narrow(0, get_tensor_model_parallel_rank() * shard, shard) + hidden_states = sp.gather_activation(hidden_shard) + with self._output_unreduced(self.mlp.down_proj): + out = self.mlp(hidden_states) + return tensor_model_parallel_reduce_scatter(out, dim=0) + def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, residual: torch.Tensor | None, prefix_delta: torch.Tensor | None = None, + sp_num_tokens: int | None = None, **kwargs, ) -> ( tuple[torch.Tensor, torch.Tensor] @@ -668,8 +735,9 @@ def forward( if self.use_attn_residuals: assert residual is not None return self.forward_attn_residual( - positions, hidden_states, residual, prefix_delta + positions, hidden_states, residual, prefix_delta, sp_num_tokens ) + assert sp_num_tokens is None assert prefix_delta is None # Self Attention @@ -692,7 +760,10 @@ def forward_attn_residual( hidden_states: torch.Tensor, block_residual: torch.Tensor, prefix_delta: torch.Tensor | None, + sp_num_tokens: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """``sp_num_tokens`` is the unpadded token count when the residual + stream is token-sharded (see :mod:`.sp`), else None.""" prefix_sum = hidden_states attention_quant_key = ( self.self_attn.get_input_quant_key() @@ -714,7 +785,12 @@ def forward_attn_residual( if self.is_block_write_layer: prefix_sum = None - hidden_states = self._run_self_attn(positions, hidden_states) + if sp_num_tokens is not None: + hidden_states = self._run_self_attn_sp( + positions, hidden_states, sp_num_tokens + ) + else: + hidden_states = self._run_self_attn(positions, hidden_states) if prefix_sum is None: prefix_sum = hidden_states @@ -740,7 +816,10 @@ def forward_attn_residual( ), ) - hidden_states = self.mlp(hidden_states) + if sp_num_tokens is not None: + hidden_states = self._run_mlp_sp(hidden_states) + else: + hidden_states = self.mlp(hidden_states) return prefix_sum, block_residual, hidden_states @@ -750,6 +829,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): config = vllm_config.model_config.hf_text_config self.config = config + self.routing_observed = ( + vllm_config.aux_output_config.enable_return_routed_experts + or vllm_config.parallel_config.enable_eplb + ) self.vocab_size = config.vocab_size @@ -891,6 +974,16 @@ def forward( return hidden_states attn_res_block_num = cdiv(self.end_layer, self.config.attn_res_block_size) + num_tokens = hidden_states.size(0) + # Specialized decode kernels take the full token batch. + use_sp = step is None and sp.should_shard( + num_tokens, + has_residual=residual is not None, + has_aux_layers=bool(self.aux_hidden_state_layers), + routing_observed=self.routing_observed, + ) + if use_sp: + hidden_states = sp_shard(hidden_states) block_residual = hidden_states.new_empty( hidden_states.size(0), attn_res_block_num, hidden_states.size(1) ) @@ -904,9 +997,18 @@ def forward( start=self.start_layer, ): kernel = step.layer(layer_idx) if step is not None else None - hidden_states, residual, prefix_delta = (kernel or layer)( - positions, hidden_states, residual, prefix_delta - ) + if kernel is not None: + hidden_states, residual, prefix_delta = kernel( + positions, hidden_states, residual, prefix_delta + ) + else: + hidden_states, residual, prefix_delta = layer( + positions=positions, + hidden_states=hidden_states, + residual=residual, + prefix_delta=prefix_delta, + sp_num_tokens=num_tokens if use_sp else None, + ) if (layer_idx + 1) in self.aux_hidden_state_layers: self._maybe_add_hidden_state( aux_hidden_states, @@ -929,6 +1031,8 @@ def forward( attn_res_block_num, delta=prefix_delta, ) + if use_sp: + hidden_states = sp.gather_tokens(hidden_states, num_tokens) # NOTE: the final norm is applied in compute_logits instead of here, so # the MTP draft model receives the pre-norm hidden states. if aux_hidden_states: diff --git a/vllm/models/kimi_k3/amd/sp.py b/vllm/models/kimi_k3/amd/sp.py new file mode 100644 index 000000000000..7bbc4b29394f --- /dev/null +++ b/vllm/models/kimi_k3/amd/sp.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Token-sharded residual stream for Kimi-K3 prefill on ROCm. + +With TP, every attention / MLP output all-reduce is replaced by a +reduce-scatter over tokens, and the (per-token) attention-residual mixing runs +on the local token shard only; the shard is all-gathered right before each +attention / MLP input. For attention and dense MLPs the communication volume +is unchanged (AR = RS + AG), while attn_res traffic and the block-residual bank +shrink by the TP size. + +In latent-MoE layers the router gate, latent down-projection and top-k also run +on the token shard (their outputs are all-gathered instead of being recomputed +on every rank), and the column-parallel up-projection tail reduce-scatters +its output instead of all-reducing it. The extra all-gathers of the latent and +the top-k add about 17% bytes per MoE layer at K3's sizes. An MoE whose tail is +not token-shardable all-gathers its input and all-reduces its output in full. + +Enabled by ``VLLM_KIMI_K3_AMD_PREFILL_SP_MIN_TOKENS``; see :mod:`vllm.envs`. +""" + +import contextlib +from collections.abc import Iterator +from dataclasses import replace +from typing import Any + +import torch +import torch.nn.functional as F + +import vllm.envs as envs +from vllm.distributed import ( + get_dp_group, + get_pp_group, + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_gather, + tensor_model_parallel_reduce_scatter, +) +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation +from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape + + +def should_shard( + num_tokens: int, + has_residual: bool, + has_aux_layers: bool, + routing_observed: bool, +) -> bool: + """Whether this forward runs with a token-sharded residual stream. + + The residual stream is padded to a multiple of the TP size, so every eager + step at or above the threshold shards; otherwise the unsharded path's + larger activations would exceed the profiled peak. CUDA-graph capture and + specialized-decode steps (checked by the caller) never shard; both are far + smaller than the profile run. + + ``routing_observed`` (routed-experts capture or EPLB load recording) keeps + it off: routing runs on each rank's token shard only, so those would see + just this rank's tokens. + """ + min_tokens = envs.VLLM_KIMI_K3_AMD_PREFILL_SP_MIN_TOKENS + return ( + min_tokens > 0 + and get_tensor_model_parallel_world_size() > 1 + and num_tokens >= min_tokens + and not has_residual + and not has_aux_layers + and not routing_observed + and get_pp_group().is_last_rank + and get_dp_group().world_size == 1 + and not torch.cuda.is_current_stream_capturing() + ) + + +def gather_tokens(shard: torch.Tensor, num_tokens: int | None = None) -> torch.Tensor: + """All-gather the token shards, dropping the TP padding rows if + ``num_tokens`` is given.""" + out = tensor_model_parallel_all_gather(shard, dim=0) + return out if num_tokens is None else out[:num_tokens] + + +def gather_activation( + shard: torch.Tensor | QuantizedActivation, num_tokens: int | None = None +) -> torch.Tensor | QuantizedActivation: + """:func:`gather_tokens` that also accepts a per-token quantized + activation, gathering its data and scale and rebuilding it.""" + if isinstance(shard, torch.Tensor): + return gather_tokens(shard, num_tokens) + key = shard.quant_key + assert key.scale.group_shape == GroupShape.PER_TOKEN and key.scale2 is None + data = gather_tokens(shard.data, num_tokens) + return replace( + shard, + data=data, + scale=gather_tokens(shard.scale, num_tokens), + orig_shape=torch.Size((data.size(0), *shard.orig_shape[1:])), + ) + + +def scatter_tokens(partial: torch.Tensor) -> torch.Tensor: + """Pad unpadded per-token partial sums back and reduce-scatter them.""" + pad = (-partial.size(0)) % get_tensor_model_parallel_world_size() + if pad: + partial = F.pad(partial, (0, 0, 0, pad)) + return tensor_model_parallel_reduce_scatter(partial, dim=0) + + +def all_gather_routing( + topk_weights: torch.Tensor, topk_ids: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """All-gather this rank's top-k over tokens in one collective.""" + w_bytes = topk_weights.contiguous().view(torch.uint8) + packed = torch.cat([w_bytes, topk_ids.contiguous().view(torch.uint8)], dim=1) + packed = tensor_model_parallel_all_gather(packed, dim=0) + nw = w_bytes.size(1) + return ( + packed[:, :nw].contiguous().view(topk_weights.dtype), + packed[:, nw:].contiguous().view(topk_ids.dtype), + ) + + +@contextlib.contextmanager +def pre_routed(router: Any, topk: tuple[torch.Tensor, torch.Tensor]) -> Iterator[None]: + """Make ``router.select_experts`` return ``topk`` for the duration. The + body must route exactly once, so the caller's placeholder router logits + are never read.""" + calls = 0 + + def select_experts(*args, **kwargs) -> tuple[torch.Tensor, torch.Tensor]: + nonlocal calls + calls += 1 + return topk + + shadowed = router.__dict__.get("select_experts") + router.select_experts = select_experts + try: + yield + finally: + if shadowed is None: + del router.select_experts + else: + router.select_experts = shadowed + assert calls == 1, f"expected one select_experts call, got {calls}"