From 8156003380e5e6bcc7881d8257fad867d7901cbb Mon Sep 17 00:00:00 2001
From: Vansh Bhatia <vansh.bhatia@amd.com>
Date: Thu, 1 Oct 2026 21:30:59 +0000
Subject: [PATCH 1/7] [ROCm][Perf] Kimi-K3: token-sharded residual stream for
 long prefills

Signed-off-by: Vansh Bhatia <vansh.bhatia@amd.com>
---
 tests/models/kimi_k3/test_amd_prefill_sp.py  | 285 +++++++++++++++++++
 vllm/envs.py                                 |  11 +
 vllm/models/kimi_k3/amd/latent_moe_runner.py |  40 ++-
 vllm/models/kimi_k3/amd/linear.py            |  90 +++++-
 vllm/models/kimi_k3/amd/sp.py                |  93 ++++++
 5 files changed, 515 insertions(+), 4 deletions(-)
 create mode 100644 tests/models/kimi_k3/test_amd_prefill_sp.py
 create mode 100644 vllm/models/kimi_k3/amd/sp.py

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..b004123eadda
--- /dev/null
+++ b/tests/models/kimi_k3/test_amd_prefill_sp.py
@@ -0,0 +1,285 @@
+# 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 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.distributed import get_tp_group
+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 import latent_moe_runner, sp
+from vllm.models.kimi_k3.amd.latent_moe_runner import ROCmLatentMoERunner
+from vllm.models.kimi_k3.amd.linear import KimiRoutedOutputTransform
+from vllm.platforms import current_platform
+from vllm.utils.network_utils import get_open_port
+
+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) -> KimiRoutedOutputTransform:
+    norm = RMSNorm(LATENT_SIZE, eps=EPS).to(device=device, dtype=DTYPE)
+    up_proj = ReplicatedLinear(
+        LATENT_SIZE,
+        HIDDEN_SIZE,
+        bias=False,
+        params_dtype=DTYPE,
+        prefix="routed_expert_up_proj",
+    ).to(device=device)
+
+    torch.manual_seed(0)
+    norm.weight.data.copy_(1 + 0.1 * torch.randn_like(norm.weight))
+    up_proj.weight.data.copy_(torch.randn_like(up_proj.weight) / LATENT_SIZE**0.5)
+    return KimiRoutedOutputTransform(norm, up_proj)
+
+
+def _tail_runner(
+    transform: KimiRoutedOutputTransform, tp_world: int
+) -> ROCmLatentMoERunner:
+    runner = object.__new__(ROCmLatentMoERunner)
+    attrs = {
+        "routed_output_transform": transform,
+        "_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:
+    """The sharded tail returns this rank's rows of the replicated tail."""
+    transform = _build_transform(device)
+    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,),
+                transform.norm.weight,
+                EPS,
+            ),
+            transform.up_proj.weight,
+        )
+        expected.add_(_all_reduced(shared, group))
+        rows = num_tokens // tp_world
+
+        sp.ACTIVE = True
+        try:
+            actual = runner._shard_up_proj_tail(routed, shared, None)
+        finally:
+            sp.ACTIVE = False
+
+        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
+        )
+
+
+_CHECKS = {
+    "sharded_tail": _check_sharded_tail,
+    "routing_gather": _check_routing_gather,
+    "padded_gather_scatter": _check_padded_gather_scatter,
+}
+
+
+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=8)
+def test_padded_gather_scatter_tp8() -> None:
+    _run_ranks("padded_gather_scatter", 8)
+
+
+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"
+
+
+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,expected",
+    [
+        (1024, 8, 16384, False, False, True, True),
+        (1024, 8, 1024, False, False, True, True),
+        (0, 8, 16384, False, False, True, False),
+        (1024, 1, 16384, False, False, True, False),
+        (1024, 8, 1016, False, False, True, False),
+        (1024, 8, 16381, False, False, True, True),
+        (1024, 8, 16384, True, False, True, False),
+        (1024, 8, 16384, False, True, True, False),
+        (1024, 8, 16384, False, False, False, 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,
+    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)
+    )
+    assert sp.should_shard(num_tokens, has_residual, has_aux) is expected
+
+
+def test_unshardable_tail_narrows_to_own_rows(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Without the sharded tail the runner still returns only this rank's rows."""
+    full = torch.arange(8 * 3, dtype=torch.float32).view(8, 3)
+    monkeypatch.setattr(moe_runner.MoERunner, "forward", lambda *a, **k: full)
+    monkeypatch.setattr(
+        latent_moe_runner, "get_tensor_model_parallel_world_size", lambda: 4
+    )
+    monkeypatch.setattr(latent_moe_runner, "get_tensor_model_parallel_rank", lambda: 2)
+    runner = object.__new__(ROCmLatentMoERunner)
+    object.__setattr__(runner, "_tail_shardable", False)
+
+    monkeypatch.setattr(sp, "ACTIVE", True)
+    out = runner.forward(torch.empty(8, 3), torch.empty(8, 2))
+    torch.testing.assert_close(out, full[4:6])
+
+    monkeypatch.setattr(sp, "ACTIVE", False)
+    assert runner.forward(torch.empty(8, 3), torch.empty(8, 2)) is full
diff --git a/vllm/envs.py b/vllm/envs.py
index e63ed5cc02c5..79ff48c6f346 100755
--- a/vllm/envs.py
+++ b/vllm/envs.py
@@ -219,6 +219,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
@@ -1670,6 +1671,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 dbebfd418203..74660b5d32be 100644
--- a/vllm/models/kimi_k3/amd/latent_moe_runner.py
+++ b/vllm/models/kimi_k3/amd/latent_moe_runner.py
@@ -8,9 +8,11 @@
     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.runner.moe_runner import MoERunner
+from vllm.models.kimi_k3.amd import sp
 
 logger = init_logger(__name__)
 
@@ -56,6 +58,18 @@ def __init__(
             )
         self._logged_sharded_tail = False
 
+    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,
@@ -74,6 +88,9 @@ def _shard_up_proj_tail(
         transform = self.routed_output_transform
         assert transform is not None
 
+        if sp.ACTIVE:
+            return self._sp_tail(fused_output, shared_output, trunc_size)
+
         latent = tensor_model_parallel_all_reduce(fused_output)
         if transform.norm is not None:
             latent = transform.norm(latent)
@@ -91,6 +108,23 @@ def _shard_up_proj_tail(
             shared_output, trunc_size, output_is_reduced=False
         )
 
+    def _sp_tail(
+        self,
+        fused_output: torch.Tensor,
+        shared_output: torch.Tensor,
+        trunc_size: int | None,
+    ) -> torch.Tensor:
+        """Token-sharded tail: reduce-scatter the latent (half the bytes of an
+        all-reduce) and up-project this rank's tokens with the full weight."""
+        transform = self.routed_output_transform
+        assert transform is not None
+        latent = tensor_model_parallel_reduce_scatter(fused_output.contiguous(), dim=0)
+        if transform.norm is not None:
+            latent = transform.norm(latent)
+        out = tensor_model_parallel_reduce_scatter(shared_output, dim=0)
+        out.addmm_(latent.to(out.dtype), transform.up_proj.weight.t())
+        return out[..., :trunc_size] if trunc_size is not None else out
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -102,9 +136,13 @@ def forward(
             return self._fused_forward(
                 hidden_states, router_logits, input_ids, shared_experts_input
             )
-        return super().forward(
+        out = super().forward(
             hidden_states, router_logits, input_ids, shared_experts_input
         )
+        if sp.ACTIVE:
+            shard = out.size(0) // get_tensor_model_parallel_world_size()
+            out = out.narrow(0, get_tensor_model_parallel_rank() * shard, shard)
+        return out
 
     def _fused_forward(
         self,
diff --git a/vllm/models/kimi_k3/amd/linear.py b/vllm/models/kimi_k3/amd/linear.py
index 65dad661ab12..d4546fdddd6a 100644
--- a/vllm/models/kimi_k3/amd/linear.py
+++ b/vllm/models/kimi_k3/amd/linear.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import contextlib
 from collections.abc import Iterable
 from typing import Any
 
@@ -10,7 +11,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
@@ -61,6 +65,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
@@ -313,7 +319,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         final_hidden_states = self.experts(
             hidden_states=hidden_states, router_logits=router_logits
         )
-        return final_hidden_states.view(num_tokens, hidden_size)
+        return final_hidden_states.view(-1, hidden_size)
 
 
 class KimiMLAAttention(nn.Module):
@@ -586,6 +592,66 @@ def _run_self_attn(
             positions=positions,
         )
 
+    @contextlib.contextmanager
+    def _output_unreduced(self, module: nn.Module):
+        # Row-parallel outputs stay partial so the caller can reduce-scatter.
+        rows = [
+            m
+            for m in module.modules()
+            if isinstance(m, RowParallelLinear) and m.reduce_results
+        ]
+        for m in rows:
+            m.reduce_results = False
+        try:
+            yield
+        finally:
+            for m in rows:
+                m.reduce_results = True
+
+    def _run_self_attn_sp(
+        self, positions: torch.Tensor, hidden_shard: torch.Tensor
+    ) -> torch.Tensor:
+        hidden_states = sp.gather_tokens(hidden_shard, sp.NUM_TOKENS)
+        with self._output_unreduced(self.self_attn):
+            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
+        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,
+            )
+
+    def _run_mlp_sp(self, hidden_shard: torch.Tensor) -> torch.Tensor:
+        # MLPs are per-token, so they run on the TP-padded token set as is.
+        if isinstance(self.mlp, KimiMoE):
+            if self.mlp.use_latent_moe and self.mlp.experts._tail_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)
+            if self.mlp.use_latent_moe:
+                # ROCmLatentMoERunner already narrows to this rank's tokens.
+                return out
+            # 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 = tensor_model_parallel_all_gather(hidden_shard, dim=0)
+        with self._output_unreduced(self.mlp):
+            out = self.mlp(hidden_states)
+        return tensor_model_parallel_reduce_scatter(out, dim=0)
+
     def forward(
         self,
         positions: torch.Tensor,
@@ -640,7 +706,10 @@ def forward_attn_residual(
         if self.is_block_write_layer:
             prefix_sum = None
 
-        hidden_states = self._run_self_attn(positions, hidden_states)
+        if sp.ACTIVE:
+            hidden_states = self._run_self_attn_sp(positions, hidden_states)
+        else:
+            hidden_states = self._run_self_attn(positions, hidden_states)
 
         if prefix_sum is None:
             prefix_sum = hidden_states
@@ -661,7 +730,10 @@ def forward_attn_residual(
             output_norm=self.post_attention_layernorm,
         )
 
-        hidden_states = self.mlp(hidden_states)
+        if sp.ACTIVE:
+            hidden_states = self._run_mlp_sp(hidden_states)
+        else:
+            hidden_states = self.mlp(hidden_states)
         return prefix_sum, block_residual, hidden_states
 
 
@@ -811,6 +883,14 @@ 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)
+        use_sp = sp.should_shard(
+            num_tokens,
+            has_residual=residual is not None,
+            has_aux_layers=bool(self.aux_hidden_state_layers),
+        )
+        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)
         )
@@ -819,6 +899,7 @@ def forward(
         residual = block_residual
         prefix_delta = None
 
+        sp.ACTIVE, sp.NUM_TOKENS = use_sp, num_tokens
         for layer_idx, layer in enumerate(
             self.layers[self.start_layer : self.end_layer],
             start=self.start_layer,
@@ -836,6 +917,7 @@ def forward(
                     hidden_states + prefix_delta,
                     residual,
                 )
+        sp.ACTIVE = False
 
         if not get_pp_group().is_last_rank:
             hidden_states = hidden_states + prefix_delta
@@ -851,6 +933,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..7be60c689ce5
--- /dev/null
+++ b/vllm/models/kimi_k3/amd/sp.py
@@ -0,0 +1,93 @@
+# 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. 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 latent is reduce-scattered instead of all-reduced.
+
+Enabled by ``VLLM_KIMI_K3_AMD_PREFILL_SP_MIN_TOKENS``; see :mod:`vllm.envs`.
+"""
+
+import contextlib
+from collections.abc import Iterator
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+
+import vllm.envs as envs
+from vllm.distributed import (
+    get_pp_group,
+    get_tensor_model_parallel_world_size,
+    tensor_model_parallel_all_gather,
+    tensor_model_parallel_reduce_scatter,
+)
+
+# True while the decoder stack runs token-sharded.
+ACTIVE = False
+# Unpadded token count of the sharded step. The residual stream is padded to a
+# multiple of the TP size, so every step above the threshold shards; otherwise
+# the unsharded path's larger activations would exceed the profiled peak.
+NUM_TOKENS = 0
+
+
+def should_shard(num_tokens: int, has_residual: bool, has_aux_layers: bool) -> bool:
+    """Whether this forward runs with a token-sharded residual stream."""
+    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 get_pp_group().is_last_rank
+        and not torch.cuda.is_current_stream_capturing()
+    )
+
+
+def gather_tokens(shard: torch.Tensor, num_tokens: int) -> torch.Tensor:
+    """All-gather the token shards and drop the TP padding rows."""
+    return tensor_model_parallel_all_gather(shard, dim=0)[:num_tokens]
+
+
+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."""
+    shadowed = router.__dict__.get("select_experts")
+    router.select_experts = lambda *args, **kwargs: topk
+    try:
+        yield
+    finally:
+        if shadowed is None:
+            del router.select_experts
+        else:
+            router.select_experts = shadowed

From b4e214f056bf3c5694417ebd7848c0cf1a396123 Mon Sep 17 00:00:00 2001
From: Vansh Bhatia <vansh.bhatia@amd.com>
Date: Wed, 7 Oct 2026 19:35:40 +0000
Subject: [PATCH 2/7] [ROCm] Kimi-K3 latent-MoE tail: avoid hipBLASLt addmm_
 fault

Both up-projection tails folded the shared output in with addmm_. On ROCm,
hipBLASLt's C-accumulating bf16 GEMM faults at some row counts for these
shapes: 2914-2925 rows per rank at 7168x3584 in the token-sharded tail (a
~23.4K-token prefill step at TP8), and bands from ~16.5K tokens at 896x3584
in the existing tail. Use a plain GEMM plus an add, which costs about 10 us
per MoE layer at 3K rows.

Signed-off-by: Vansh Bhatia <vansh.bhatia@amd.com>
---
 vllm/models/kimi_k3/amd/latent_moe_runner.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/vllm/models/kimi_k3/amd/latent_moe_runner.py b/vllm/models/kimi_k3/amd/latent_moe_runner.py
index 74660b5d32be..056242fcd9ac 100644
--- a/vllm/models/kimi_k3/amd/latent_moe_runner.py
+++ b/vllm/models/kimi_k3/amd/latent_moe_runner.py
@@ -100,9 +100,9 @@ def _shard_up_proj_tail(
         up_proj_shard = transform.up_proj.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
-        # beta-add epilogue so folding in the shared partial costs no kernel.
-        hidden_shard.addmm_(latent, up_proj_shard.t())
+        # Not addmm_: hipBLASLt's C-accumulating bf16 GEMM faults at some row
+        # counts for this shape (e.g. 23393-23405 rows at 896x3584).
+        hidden_shard += latent @ up_proj_shard.t()
 
         return self._maybe_reduce_final_output(
             shared_output, trunc_size, output_is_reduced=False
@@ -122,7 +122,9 @@ def _sp_tail(
         if transform.norm is not None:
             latent = transform.norm(latent)
         out = tensor_model_parallel_reduce_scatter(shared_output, dim=0)
-        out.addmm_(latent.to(out.dtype), transform.up_proj.weight.t())
+        # Not addmm_: hipBLASLt's C-accumulating bf16 GEMM faults at some row
+        # counts for this shape (e.g. 2914-2925 rows at 7168x3584).
+        out += latent.to(out.dtype) @ transform.up_proj.weight.t()
         return out[..., :trunc_size] if trunc_size is not None else out
 
     def forward(

From 239cee6640037fa899d6b3ec832b46183fe6ee60 Mon Sep 17 00:00:00 2001
From: Vansh Bhatia <vansh.bhatia@amd.com>
Date: Wed, 14 Oct 2026 16:27:36 +0000
Subject: [PATCH 3/7] [ROCm][Kimi-K3] Address review on token-sharded prefill

- Reuse the column-parallel up-projection tail and reduce-scatter its
  output over tokens, so the path works with the row-sharded up_proj.
- Gather pre-quantized AttnRes outputs as data plus per-token scale and
  rebuild them, keeping the FP8 input fusion.
- Pass the sharded token count to the decoder layers and the MoE runner
  as arguments instead of module globals.
- Only shard with DP == 1.
- Drop the addmm_ workaround; the hipBLASLt fault it avoided only
  reproduces on the ROCm 10 preview image.
- Test the ColumnParallelLinear up_proj layout and the quantized gather.

Signed-off-by: Vansh Bhatia <vansh.bhatia@amd.com>
---
 tests/models/kimi_k3/test_amd_prefill_sp.py  | 134 ++++++++++++-------
 vllm/models/kimi_k3/amd/latent_moe_runner.py |  66 ++++-----
 vllm/models/kimi_k3/amd/linear.py            |  41 +++---
 vllm/models/kimi_k3/amd/sp.py                |  48 +++++--
 4 files changed, 178 insertions(+), 111 deletions(-)

diff --git a/tests/models/kimi_k3/test_amd_prefill_sp.py b/tests/models/kimi_k3/test_amd_prefill_sp.py
index b004123eadda..c0279874af97 100644
--- a/tests/models/kimi_k3/test_amd_prefill_sp.py
+++ b/tests/models/kimi_k3/test_amd_prefill_sp.py
@@ -20,10 +20,13 @@
     multi_gpu_test,
 )
 from vllm.distributed import get_tp_group
-from vllm.model_executor.layers.fused_moe.runner import moe_runner
+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 ReplicatedLinear
-from vllm.models.kimi_k3.amd import latent_moe_runner, sp
+from vllm.model_executor.layers.linear import ColumnParallelLinear, ReplicatedLinear
+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 KimiRoutedOutputTransform
 from vllm.platforms import current_platform
@@ -42,20 +45,39 @@
 DTYPE = torch.bfloat16
 
 
-def _build_transform(device: torch.device) -> KimiRoutedOutputTransform:
+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)
-    up_proj = ReplicatedLinear(
-        LATENT_SIZE,
-        HIDDEN_SIZE,
-        bias=False,
-        params_dtype=DTYPE,
-        prefix="routed_expert_up_proj",
-    ).to(device=device)
-
     torch.manual_seed(0)
     norm.weight.data.copy_(1 + 0.1 * torch.randn_like(norm.weight))
-    up_proj.weight.data.copy_(torch.randn_like(up_proj.weight) / LATENT_SIZE**0.5)
-    return KimiRoutedOutputTransform(norm, up_proj)
+    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(
@@ -64,6 +86,7 @@ def _tail_runner(
     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(
@@ -85,8 +108,16 @@ def _all_reduced(tensor: torch.Tensor, group) -> torch.Tensor:
 
 
 def _check_sharded_tail(device: torch.device, tp_world: int, rank: int) -> None:
-    """The sharded tail returns this rank's rows of the replicated tail."""
-    transform = _build_transform(device)
+    """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)
     runner = _tail_runner(transform, tp_world)
     group = get_tp_group().device_group
 
@@ -102,17 +133,15 @@ def _check_sharded_tail(device: torch.device, tp_world: int, rank: int) -> None:
                 transform.norm.weight,
                 EPS,
             ),
-            transform.up_proj.weight,
+            weight,
         )
         expected.add_(_all_reduced(shared, group))
         rows = num_tokens // tp_world
 
-        sp.ACTIVE = True
-        try:
-            actual = runner._shard_up_proj_tail(routed, shared, None)
-        finally:
-            sp.ACTIVE = False
+        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
@@ -159,6 +188,25 @@ def _check_padded_gather_scatter(
             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)
+
 
 _CHECKS = {
     "sharded_tail": _check_sharded_tail,
@@ -235,17 +283,18 @@ def test_pre_routed_keeps_an_instance_override() -> None:
 
 
 @pytest.mark.parametrize(
-    "min_tokens,tp_size,num_tokens,has_residual,has_aux,last_rank,expected",
+    "min_tokens,tp_size,num_tokens,has_residual,has_aux,last_rank,dp_size,expected",
     [
-        (1024, 8, 16384, False, False, True, True),
-        (1024, 8, 1024, False, False, True, True),
-        (0, 8, 16384, False, False, True, False),
-        (1024, 1, 16384, False, False, True, False),
-        (1024, 8, 1016, False, False, True, False),
-        (1024, 8, 16381, False, False, True, True),
-        (1024, 8, 16384, True, False, True, False),
-        (1024, 8, 16384, False, True, True, False),
-        (1024, 8, 16384, False, False, False, False),
+        (1024, 8, 16384, False, False, True, 1, True),
+        (1024, 8, 1024, False, False, True, 1, True),
+        (0, 8, 16384, False, False, True, 1, False),
+        (1024, 1, 16384, False, False, True, 1, False),
+        (1024, 8, 1016, False, False, True, 1, False),
+        (1024, 8, 16381, False, False, True, 1, True),
+        (1024, 8, 16384, True, False, True, 1, False),
+        (1024, 8, 16384, False, True, True, 1, False),
+        (1024, 8, 16384, False, False, False, 1, False),
+        (1024, 8, 16384, False, False, True, 2, False),
     ],
 )
 def test_should_shard(
@@ -256,6 +305,7 @@ def test_should_shard(
     has_residual: bool,
     has_aux: bool,
     last_rank: bool,
+    dp_size: int,
     expected: bool,
 ) -> None:
     monkeypatch.setenv("VLLM_KIMI_K3_AMD_PREFILL_SP_MIN_TOKENS", str(min_tokens))
@@ -263,23 +313,5 @@ def test_should_shard(
     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) is expected
-
-
-def test_unshardable_tail_narrows_to_own_rows(monkeypatch: pytest.MonkeyPatch) -> None:
-    """Without the sharded tail the runner still returns only this rank's rows."""
-    full = torch.arange(8 * 3, dtype=torch.float32).view(8, 3)
-    monkeypatch.setattr(moe_runner.MoERunner, "forward", lambda *a, **k: full)
-    monkeypatch.setattr(
-        latent_moe_runner, "get_tensor_model_parallel_world_size", lambda: 4
-    )
-    monkeypatch.setattr(latent_moe_runner, "get_tensor_model_parallel_rank", lambda: 2)
-    runner = object.__new__(ROCmLatentMoERunner)
-    object.__setattr__(runner, "_tail_shardable", False)
-
-    monkeypatch.setattr(sp, "ACTIVE", True)
-    out = runner.forward(torch.empty(8, 3), torch.empty(8, 2))
-    torch.testing.assert_close(out, full[4:6])
-
-    monkeypatch.setattr(sp, "ACTIVE", False)
-    assert runner.forward(torch.empty(8, 3), torch.empty(8, 2)) is full
diff --git a/vllm/models/kimi_k3/amd/latent_moe_runner.py b/vllm/models/kimi_k3/amd/latent_moe_runner.py
index 29814a40870b..6b758c1772b3 100644
--- a/vllm/models/kimi_k3/amd/latent_moe_runner.py
+++ b/vllm/models/kimi_k3/amd/latent_moe_runner.py
@@ -11,8 +11,10 @@
     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
-from vllm.models.kimi_k3.amd import sp
 
 logger = init_logger(__name__)
 
@@ -63,6 +65,16 @@ def __init__(
             )
         self._logged_sharded_tail = False
 
+    @property
+    def token_shardable(self) -> bool:
+        """Whether ``forward`` supports ``token_sharded=True``."""
+        return (
+            self._tail_shardable
+            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]:
@@ -80,6 +92,7 @@ def _shard_up_proj_tail(
         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:
@@ -93,9 +106,6 @@ def _shard_up_proj_tail(
         transform = self.routed_output_transform
         assert transform is not None
 
-        if sp.ACTIVE:
-            return self._sp_tail(fused_output, shared_output, trunc_size)
-
         latent = tensor_model_parallel_all_reduce(fused_output)
         if transform.norm is not None:
             latent = transform.norm(latent)
@@ -110,51 +120,40 @@ def _shard_up_proj_tail(
         )
         hidden_shard = shared_output.narrow(-1, shard_start, shard_size)
 
-        # Not addmm_: hipBLASLt's C-accumulating bf16 GEMM faults at some row
-        # counts for this shape (e.g. 23393-23405 rows at 896x3584).
-        hidden_shard += latent @ up_proj_shard.t()
+        # hidden_shard += latent @ up_proj_shard.T, accumulated in the GEMM's
+        # 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
         )
 
-    def _sp_tail(
-        self,
-        fused_output: torch.Tensor,
-        shared_output: torch.Tensor,
-        trunc_size: int | None,
-    ) -> torch.Tensor:
-        """Token-sharded tail: reduce-scatter the latent (half the bytes of an
-        all-reduce) and up-project this rank's tokens with the full weight."""
-        transform = self.routed_output_transform
-        assert transform is not None
-        latent = tensor_model_parallel_reduce_scatter(fused_output.contiguous(), dim=0)
-        if transform.norm is not None:
-            latent = transform.norm(latent)
-        out = tensor_model_parallel_reduce_scatter(shared_output, dim=0)
-        # Not addmm_: hipBLASLt's C-accumulating bf16 GEMM faults at some row
-        # counts for this shape (e.g. 2914-2925 rows at 7168x3584).
-        out += latent.to(out.dtype) @ transform.up_proj.weight.t()
-        return out[..., :trunc_size] if trunc_size is not None else out
-
     def forward(
         self,
         hidden_states: torch.Tensor,
         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,
             )
-        out = super().forward(
+        return super().forward(
             hidden_states, router_logits, input_ids, shared_experts_input
         )
-        if sp.ACTIVE:
-            shard = out.size(0) // get_tensor_model_parallel_world_size()
-            out = out.narrow(0, get_tensor_model_parallel_rank() * shard, shard)
-        return out
 
     def _fused_forward(
         self,
@@ -162,6 +161,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
@@ -196,7 +196,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 beb89ab0cd0d..18e177320c25 100644
--- a/vllm/models/kimi_k3/amd/linear.py
+++ b/vllm/models/kimi_k3/amd/linear.py
@@ -367,7 +367,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         final_hidden_states = self.experts(
             hidden_states=hidden_states, router_logits=router_logits
         )
-        return final_hidden_states.view(-1, hidden_size)
+        return final_hidden_states.view(num_tokens, hidden_size)
 
 
 class KimiMLAAttention(nn.Module):
@@ -660,15 +660,19 @@ def _output_unreduced(self, module: nn.Module):
                 m.reduce_results = True
 
     def _run_self_attn_sp(
-        self, positions: torch.Tensor, hidden_shard: torch.Tensor
+        self,
+        positions: torch.Tensor,
+        hidden_shard: torch.Tensor | QuantizedActivation,
+        num_tokens: int,
     ) -> torch.Tensor:
-        hidden_states = sp.gather_tokens(hidden_shard, sp.NUM_TOKENS)
+        hidden_states = sp.gather_activation(hidden_shard, num_tokens)
         with self._output_unreduced(self.self_attn):
             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
@@ -683,22 +687,23 @@ def _run_moe_sp_shard_routed(self, hidden_shard: torch.Tensor) -> torch.Tensor:
                 hidden_states=latent,
                 router_logits=logits,
                 shared_experts_input=hidden_states,
+                token_sharded=True,
             )
 
-    def _run_mlp_sp(self, hidden_shard: torch.Tensor) -> torch.Tensor:
+    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):
-            if self.mlp.use_latent_moe and self.mlp.experts._tail_shardable:
+            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)
-            if self.mlp.use_latent_moe:
-                # ROCmLatentMoERunner already narrows to this rank's tokens.
-                return out
             # 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 = tensor_model_parallel_all_gather(hidden_shard, dim=0)
+        hidden_states = sp.gather_activation(hidden_shard)
         with self._output_unreduced(self.mlp):
             out = self.mlp(hidden_states)
         return tensor_model_parallel_reduce_scatter(out, dim=0)
@@ -709,6 +714,7 @@ def forward(
         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]
@@ -717,8 +723,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
@@ -741,7 +748,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()
@@ -763,8 +773,10 @@ def forward_attn_residual(
         if self.is_block_write_layer:
             prefix_sum = None
 
-        if sp.ACTIVE:
-            hidden_states = self._run_self_attn_sp(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)
 
@@ -792,7 +804,7 @@ def forward_attn_residual(
             ),
         )
 
-        if sp.ACTIVE:
+        if sp_num_tokens is not None:
             hidden_states = self._run_mlp_sp(hidden_states)
         else:
             hidden_states = self.mlp(hidden_states)
@@ -961,7 +973,6 @@ def forward(
         residual = block_residual
         prefix_delta = None
 
-        sp.ACTIVE, sp.NUM_TOKENS = use_sp, num_tokens
         for layer_idx, layer in enumerate(
             self.layers[self.start_layer : self.end_layer],
             start=self.start_layer,
@@ -971,6 +982,7 @@ def forward(
                 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(
@@ -979,7 +991,6 @@ def forward(
                     hidden_states + prefix_delta,
                     residual,
                 )
-        sp.ACTIVE = False
 
         if not get_pp_group().is_last_rank:
             hidden_states = hidden_states + prefix_delta
diff --git a/vllm/models/kimi_k3/amd/sp.py b/vllm/models/kimi_k3/amd/sp.py
index 7be60c689ce5..25f7ea1d4745 100644
--- a/vllm/models/kimi_k3/amd/sp.py
+++ b/vllm/models/kimi_k3/amd/sp.py
@@ -10,13 +10,15 @@
 
 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 latent is reduce-scattered instead of all-reduced.
+on every rank), and the column-parallel up-projection tail reduce-scatters
+its output instead of all-reducing it.
 
 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
@@ -24,22 +26,23 @@
 
 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,
 )
-
-# True while the decoder stack runs token-sharded.
-ACTIVE = False
-# Unpadded token count of the sharded step. The residual stream is padded to a
-# multiple of the TP size, so every step above the threshold shards; otherwise
-# the unsharded path's larger activations would exceed the profiled peak.
-NUM_TOKENS = 0
+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) -> bool:
-    """Whether this forward runs with a token-sharded residual stream."""
+    """Whether this forward runs with a token-sharded residual stream.
+
+    The residual stream is padded to a multiple of the TP size, so every step
+    above the threshold shards; otherwise the unsharded path's larger
+    activations would exceed the profiled peak.
+    """
     min_tokens = envs.VLLM_KIMI_K3_AMD_PREFILL_SP_MIN_TOKENS
     return (
         min_tokens > 0
@@ -48,13 +51,34 @@ def should_shard(num_tokens: int, has_residual: bool, has_aux_layers: bool) -> b
         and not has_residual
         and not has_aux_layers
         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) -> torch.Tensor:
-    """All-gather the token shards and drop the TP padding rows."""
-    return tensor_model_parallel_all_gather(shard, dim=0)[:num_tokens]
+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:

From adb239241bae74c6937da67074ca39ee6147cfa5 Mon Sep 17 00:00:00 2001
From: Vansh Bhatia <vansh.bhatia@amd.com>
Date: Mon, 19 Oct 2026 16:14:16 +0000
Subject: [PATCH 4/7] [ROCm][Kimi-K3] Check shard routing on a real MoE; target
 unreduced projs

- pre_routed now asserts the MoE forward routes exactly once, so the
  placeholder router logits can never be read.
- Test a small latent KimiMoE: routing on the token shard returns this
  rank's rows of the unsharded layer at TP4/TP8, on Triton and AITER.
- _output_unreduced takes the one RowParallelLinear to leave partial
  (attention o_proj, MLP down_proj) instead of every one in the module.

Signed-off-by: Vansh Bhatia <vansh.bhatia@amd.com>
---
 tests/models/kimi_k3/test_amd_prefill_sp.py | 99 ++++++++++++++++++++-
 vllm/models/kimi_k3/amd/linear.py           | 21 ++---
 vllm/models/kimi_k3/amd/sp.py               | 14 ++-
 3 files changed, 118 insertions(+), 16 deletions(-)

diff --git a/tests/models/kimi_k3/test_amd_prefill_sp.py b/tests/models/kimi_k3/test_amd_prefill_sp.py
index c0279874af97..706047a09da5 100644
--- a/tests/models/kimi_k3/test_amd_prefill_sp.py
+++ b/tests/models/kimi_k3/test_amd_prefill_sp.py
@@ -19,18 +19,27 @@
     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 KimiRoutedOutputTransform
+from vllm.models.kimi_k3.amd.linear import (
+    KimiDecoderLayer,
+    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(),
@@ -208,10 +217,65 @@ def _check_padded_gather_scatter(
             torch.testing.assert_close(got.scale, scale[:expected_rows], atol=0, rtol=0)
 
 
+def _check_shard_routed_moe(device: torch.device, tp_world: int, rank: int) -> None:
+    """A small latent KimiMoE routed on the token shard returns this rank's
+    rows of the unsharded layer, with the gathered top-k replacing routing."""
+    init_workspace_manager(device)
+    config = KimiLinearConfig(
+        hidden_size=512,
+        hidden_act="silu",
+        rms_norm_eps=EPS,
+        # At least the 256 per rank that KimiMoE would otherwise pad up to.
+        moe_intermediate_size=256 * tp_world,
+        num_experts=32,
+        num_experts_per_token=4,
+        num_shared_experts=1,
+        routed_expert_hidden_size=256,
+        latent_moe_use_norm=True,
+        use_grouped_topk=True,
+        num_expert_group=1,
+        topk_group=1,
+    )
+    vllm_config = VllmConfig()
+    vllm_config.compilation_config.static_forward_context = {}
+    torch.set_default_dtype(DTYPE)
+    with set_current_vllm_config(vllm_config), torch.device(device):
+        moe = KimiMoE(config, prefix="model.layers.1.mlp", layer_idx=1)
+    assert moe.experts.token_shardable
+
+    # 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 moe.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 moe.modules():
+        quant_method = getattr(module, "quant_method", None)
+        if isinstance(quant_method, QuantizeMethodBase):
+            quant_method.process_weights_after_loading(module)
+
+    rows = 6
+    num_tokens = tp_world * rows
+    torch.manual_seed(1)
+    hidden = torch.randn(num_tokens, config.hidden_size, device=device, dtype=DTYPE)
+    mine = slice(rank * rows, (rank + 1) * rows)
+    layer = SimpleNamespace(mlp=moe)
+    with set_forward_context(None, vllm_config, num_tokens=num_tokens):
+        expected = moe(hidden)
+        actual = KimiDecoderLayer._run_moe_sp_shard_routed(layer, hidden[mine])
+
+    assert actual.shape == (rows, config.hidden_size)
+    torch.testing.assert_close(actual, expected[mine], atol=8e-2, rtol=3e-2)
+
+
 _CHECKS = {
     "sharded_tail": _check_sharded_tail,
     "routing_gather": _check_routing_gather,
     "padded_gather_scatter": _check_padded_gather_scatter,
+    "shard_routed_moe": _check_shard_routed_moe,
 }
 
 
@@ -254,6 +318,30 @@ 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)
+
+
 class _Router(torch.nn.Module):
     def select_experts(self, hidden_states, router_logits, **kwargs):
         return "routed"
@@ -273,6 +361,15 @@ def test_pre_routed_overrides_then_restores() -> None:
     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"
diff --git a/vllm/models/kimi_k3/amd/linear.py b/vllm/models/kimi_k3/amd/linear.py
index 18e177320c25..89f0ed4a4121 100644
--- a/vllm/models/kimi_k3/amd/linear.py
+++ b/vllm/models/kimi_k3/amd/linear.py
@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import contextlib
-from collections.abc import Iterable
+from collections.abc import Iterable, Iterator
 from typing import Any, cast
 
 import torch
@@ -643,21 +643,16 @@ def _run_self_attn(
             positions=positions,
         )
 
+    @staticmethod
     @contextlib.contextmanager
-    def _output_unreduced(self, module: nn.Module):
+    def _output_unreduced(proj: RowParallelLinear) -> Iterator[None]:
         # Row-parallel outputs stay partial so the caller can reduce-scatter.
-        rows = [
-            m
-            for m in module.modules()
-            if isinstance(m, RowParallelLinear) and m.reduce_results
-        ]
-        for m in rows:
-            m.reduce_results = False
+        assert isinstance(proj, RowParallelLinear) and proj.reduce_results
+        proj.reduce_results = False
         try:
             yield
         finally:
-            for m in rows:
-                m.reduce_results = True
+            proj.reduce_results = True
 
     def _run_self_attn_sp(
         self,
@@ -666,7 +661,7 @@ def _run_self_attn_sp(
         num_tokens: int,
     ) -> torch.Tensor:
         hidden_states = sp.gather_activation(hidden_shard, num_tokens)
-        with self._output_unreduced(self.self_attn):
+        with self._output_unreduced(self.self_attn.o_proj):
             out = self._run_self_attn(positions, hidden_states)
         return sp.scatter_tokens(out)
 
@@ -704,7 +699,7 @@ def _run_mlp_sp(
             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):
+        with self._output_unreduced(self.mlp.down_proj):
             out = self.mlp(hidden_states)
         return tensor_model_parallel_reduce_scatter(out, dim=0)
 
diff --git a/vllm/models/kimi_k3/amd/sp.py b/vllm/models/kimi_k3/amd/sp.py
index 25f7ea1d4745..ae6125d902f9 100644
--- a/vllm/models/kimi_k3/amd/sp.py
+++ b/vllm/models/kimi_k3/amd/sp.py
@@ -105,9 +105,18 @@ def all_gather_routing(
 
 @contextlib.contextmanager
 def pre_routed(router: Any, topk: tuple[torch.Tensor, torch.Tensor]) -> Iterator[None]:
-    """Make ``router.select_experts`` return ``topk`` for the duration."""
+    """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 = lambda *args, **kwargs: topk
+    router.select_experts = select_experts
     try:
         yield
     finally:
@@ -115,3 +124,4 @@ def pre_routed(router: Any, topk: tuple[torch.Tensor, torch.Tensor]) -> Iterator
             del router.select_experts
         else:
             router.select_experts = shadowed
+    assert calls == 1, f"expected one select_experts call, got {calls}"

From f0c6a4bb0b805b408a3631083efbe4f63c51c6d2 Mon Sep 17 00:00:00 2001
From: Vansh Bhatia <vansh.bhatia@amd.com>
Date: Tue, 20 Oct 2026 18:49:45 +0000
Subject: [PATCH 5/7] [ROCm][Kimi-K3] Keep prefill SP off when routing is
 observed; test the sharded MLP step

Routed-experts capture and EPLB load recording see routing from
select_experts, which under SP only runs on this rank's token shard, so
should_shard now also requires neither to be enabled.

Tests: TP4 padded gather/scatter, a token count that does not divide TP in
the shard-routed MoE check, and _run_mlp_sp against the unsharded rows for a
dense KimiMLP and a non-token-shardable KimiMoE. The sp.py docstring now
notes the extra latent and top-k all-gathers in latent-MoE layers.

Signed-off-by: Vansh Bhatia <vansh.bhatia@amd.com>
---
 tests/models/kimi_k3/test_amd_prefill_sp.py | 156 ++++++++++++++------
 vllm/models/kimi_k3/amd/linear.py           |   5 +
 vllm/models/kimi_k3/amd/sp.py               |  21 ++-
 3 files changed, 136 insertions(+), 46 deletions(-)

diff --git a/tests/models/kimi_k3/test_amd_prefill_sp.py b/tests/models/kimi_k3/test_amd_prefill_sp.py
index 706047a09da5..7b856de45fd5 100644
--- a/tests/models/kimi_k3/test_amd_prefill_sp.py
+++ b/tests/models/kimi_k3/test_amd_prefill_sp.py
@@ -6,6 +6,8 @@
 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
@@ -33,6 +35,7 @@
 from vllm.models.kimi_k3.amd.latent_moe_runner import ROCmLatentMoERunner
 from vllm.models.kimi_k3.amd.linear import (
     KimiDecoderLayer,
+    KimiMLP,
     KimiMoE,
     KimiRoutedOutputTransform,
 )
@@ -217,65 +220,121 @@ def _check_padded_gather_scatter(
             torch.testing.assert_close(got.scale, scale[:expected_rows], atol=0, rtol=0)
 
 
-def _check_shard_routed_moe(device: torch.device, tp_world: int, rank: int) -> None:
-    """A small latent KimiMoE routed on the token shard returns this rank's
-    rows of the unsharded layer, with the gathered top-k replacing routing."""
+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)
-    config = KimiLinearConfig(
-        hidden_size=512,
-        hidden_act="silu",
-        rms_norm_eps=EPS,
-        # At least the 256 per rank that KimiMoE would otherwise pad up to.
-        moe_intermediate_size=256 * tp_world,
-        num_experts=32,
-        num_experts_per_token=4,
-        num_shared_experts=1,
-        routed_expert_hidden_size=256,
-        latent_moe_use_norm=True,
-        use_grouped_topk=True,
-        num_expert_group=1,
-        topk_group=1,
-    )
     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):
-        moe = KimiMoE(config, prefix="model.layers.1.mlp", layer_idx=1)
-    assert moe.experts.token_shardable
+        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 moe.parameters():
+        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 moe.modules():
+    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
+    num_tokens = tp_world * rows - 3
     torch.manual_seed(1)
-    hidden = torch.randn(num_tokens, config.hidden_size, device=device, dtype=DTYPE)
+    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)
-    layer = SimpleNamespace(mlp=moe)
+    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 = moe(hidden)
-        actual = KimiDecoderLayer._run_moe_sp_shard_routed(layer, hidden[mine])
+        expected = mlp(hidden)
+        actual = run_sharded(layer, padded[mine])
 
-    assert actual.shape == (rows, config.hidden_size)
-    torch.testing.assert_close(actual, expected[mine], atol=8e-2, rtol=3e-2)
+    assert actual.shape == (rows, MLP_HIDDEN_SIZE)
+    torch.testing.assert_close(
+        actual[:real], expected[mine][:real], atol=8e-2, rtol=3e-2
+    )
 
 
-_CHECKS = {
+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"),
 }
 
 
@@ -313,6 +372,11 @@ 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)
@@ -342,6 +406,12 @@ def test_shard_routed_moe_tp8_matches_unsharded_rows(
     _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"
@@ -380,18 +450,19 @@ def test_pre_routed_keeps_an_instance_override() -> None:
 
 
 @pytest.mark.parametrize(
-    "min_tokens,tp_size,num_tokens,has_residual,has_aux,last_rank,dp_size,expected",
+    "min_tokens,tp_size,num_tokens,has_residual,has_aux,last_rank,dp_size,observed,expected",
     [
-        (1024, 8, 16384, False, False, True, 1, True),
-        (1024, 8, 1024, False, False, True, 1, True),
-        (0, 8, 16384, False, False, True, 1, False),
-        (1024, 1, 16384, False, False, True, 1, False),
-        (1024, 8, 1016, False, False, True, 1, False),
-        (1024, 8, 16381, False, False, True, 1, True),
-        (1024, 8, 16384, True, False, True, 1, False),
-        (1024, 8, 16384, False, True, True, 1, False),
-        (1024, 8, 16384, False, False, False, 1, False),
-        (1024, 8, 16384, False, False, True, 2, False),
+        (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(
@@ -403,6 +474,7 @@ def test_should_shard(
     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))
@@ -411,4 +483,4 @@ def test_should_shard(
         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) is expected
+    assert sp.should_shard(num_tokens, has_residual, has_aux, observed) is expected
diff --git a/vllm/models/kimi_k3/amd/linear.py b/vllm/models/kimi_k3/amd/linear.py
index 89f0ed4a4121..3596d7ed30d0 100644
--- a/vllm/models/kimi_k3/amd/linear.py
+++ b/vllm/models/kimi_k3/amd/linear.py
@@ -812,6 +812,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
 
@@ -957,6 +961,7 @@ def forward(
             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)
diff --git a/vllm/models/kimi_k3/amd/sp.py b/vllm/models/kimi_k3/amd/sp.py
index ae6125d902f9..db640f323bac 100644
--- a/vllm/models/kimi_k3/amd/sp.py
+++ b/vllm/models/kimi_k3/amd/sp.py
@@ -5,13 +5,16 @@
 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. Communication volume is unchanged (AR = RS + AG), while
-attn_res traffic and the block-residual bank shrink by the TP size.
+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.
+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`.
 """
@@ -36,12 +39,21 @@
 from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
 
 
-def should_shard(num_tokens: int, has_residual: bool, has_aux_layers: bool) -> bool:
+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 step
     above the threshold shards; otherwise the unsharded path's larger
     activations would exceed the profiled peak.
+
+    ``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 (
@@ -50,6 +62,7 @@ def should_shard(num_tokens: int, has_residual: bool, has_aux_layers: bool) -> b
         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()

From c7ff67b8ce98ced96469c0198b5be33ffe648ca1 Mon Sep 17 00:00:00 2001
From: Vansh Bhatia <vansh.bhatia@amd.com>
Date: Fri, 30 Oct 2026 17:01:09 +0000
Subject: [PATCH 6/7] [ROCm][Kimi-K3] Narrow the optional latent norm in the
 prefill SP test

Signed-off-by: Vansh Bhatia <vansh.bhatia@amd.com>
---
 tests/models/kimi_k3/test_amd_prefill_sp.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tests/models/kimi_k3/test_amd_prefill_sp.py b/tests/models/kimi_k3/test_amd_prefill_sp.py
index 7b856de45fd5..17ff5d7c90e1 100644
--- a/tests/models/kimi_k3/test_amd_prefill_sp.py
+++ b/tests/models/kimi_k3/test_amd_prefill_sp.py
@@ -130,6 +130,8 @@ 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
 
@@ -142,7 +144,7 @@ def _check_tail_layout(
             F.rms_norm(
                 _all_reduced(routed, group),
                 (LATENT_SIZE,),
-                transform.norm.weight,
+                norm_weight,
                 EPS,
             ),
             weight,

From 6214ff8b3f8a3f5c0035d6e1d887c3e1d1c765cf Mon Sep 17 00:00:00 2001
From: Vansh Bhatia <vansh.bhatia@amd.com>
Date: Mon, 2 Nov 2026 00:45:34 +0000
Subject: [PATCH 7/7] [ROCm][Kimi-K3] Keep monolithic MoE kernels off the
 shard-routed path

Signed-off-by: Vansh Bhatia <vansh.bhatia@amd.com>
---
 vllm/models/kimi_k3/amd/latent_moe_runner.py | 5 ++++-
 vllm/models/kimi_k3/amd/sp.py                | 8 +++++---
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/vllm/models/kimi_k3/amd/latent_moe_runner.py b/vllm/models/kimi_k3/amd/latent_moe_runner.py
index 6b758c1772b3..6455b39a2637 100644
--- a/vllm/models/kimi_k3/amd/latent_moe_runner.py
+++ b/vllm/models/kimi_k3/amd/latent_moe_runner.py
@@ -67,9 +67,12 @@ def __init__(
 
     @property
     def token_shardable(self) -> bool:
-        """Whether ``forward`` supports ``token_sharded=True``."""
+        """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)
diff --git a/vllm/models/kimi_k3/amd/sp.py b/vllm/models/kimi_k3/amd/sp.py
index db640f323bac..7bbc4b29394f 100644
--- a/vllm/models/kimi_k3/amd/sp.py
+++ b/vllm/models/kimi_k3/amd/sp.py
@@ -47,9 +47,11 @@ def should_shard(
 ) -> 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 step
-    above the threshold shards; otherwise the unsharded path's larger
-    activations would exceed the profiled peak.
+    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