[Kimi-K3] Add option to shard the shared expert instead of replicating (#50656)

Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>

tests/models/kimi_k3/test_sequence_parallel.py

@@ -10,6 +10,7 @@ import torch
 from torch import nn
 
 from vllm.config import ParallelConfig
+from vllm.model_executor.layers.activation import SiluAndMul
 from vllm.models.common.ops import sequence_parallel as sp_ops
 from vllm.models.kimi_k3.nvidia import model as kimi_model
 from vllm.models.kimi_k3.nvidia import mtp as kimi_mtp
@@ -265,6 +266,84 @@ def test_kimi_mtp_restores_sequence_parallel_output(monkeypatch):
     torch.testing.assert_close(final_norm.call_args.args[0], expected_hidden_states)
 
 
+@pytest.mark.parametrize(
+    ("enabled", "use_sequence_parallel", "eligible", "tp_size", "expected"),
+    [
+        (True, True, True, 8, True),
+        (False, True, True, 8, False),  # opt-in only
+        (True, False, True, 8, False),  # replication only exists under SP
+        (True, True, False, 8, False),  # FusedMoE path owns the reduction
+        (True, True, True, 1, False),  # nothing to shard
+        (True, True, True, 5, False),  # 6144 % 5 -- would fail divide()
+    ],
+)
+def test_shard_sequence_parallel_mlp_gating(
+    monkeypatch,
+    enabled: bool,
+    use_sequence_parallel: bool,
+    eligible: bool,
+    tp_size: int,
+    expected: bool,
+):
+    monkeypatch.setattr(kimi_model.envs, "VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT", enabled)
+    monkeypatch.setattr(
+        kimi_model, "get_tensor_model_parallel_world_size", lambda: tp_size
+    )
+
+    assert (
+        kimi_model.shard_sequence_parallel_mlp(
+            hidden_size=7168,
+            intermediate_size=6144,
+            use_sequence_parallel=use_sequence_parallel,
+            eligible=eligible,
+        )
+        is expected
+    )
+
+
+def test_sharded_sequence_parallel_mlp_matches_replicated(default_vllm_config):
+    """Sharded SP MLP must reproduce the replicated result for every token.
+
+    Each rank owns a *disjoint* token shard, so a weight shard alone cannot
+    finish a rank's own tokens: the ranks must gather the full token set,
+    compute partial sums over their intermediate shard, and reduce-scatter.
+    Splicing per-rank feature slices together instead silently mixes different
+    tokens and produces plausible-looking garbage.
+    """
+    tp_size, hidden, intermediate, tokens_per_rank = 4, 16, 12, 3
+    torch.manual_seed(0)
+    num_tokens = tp_size * tokens_per_rank
+    x = torch.randn(num_tokens, hidden)
+    gate_weight = torch.randn(intermediate, hidden)
+    up_weight = torch.randn(intermediate, hidden)
+    down_weight = torch.randn(hidden, intermediate)
+    act_fn = SiluAndMul()
+
+    replicated = act_fn(x @ torch.cat([gate_weight, up_weight]).T) @ down_weight.T
+
+    shard = intermediate // tp_size
+    # Every rank all-gathers the full token set, then computes its partial.
+    partials = [
+        act_fn(
+            x
+            @ torch.cat(
+                [
+                    gate_weight[r * shard : (r + 1) * shard],
+                    up_weight[r * shard : (r + 1) * shard],
+                ]
+            ).T
+        )
+        @ down_weight[:, r * shard : (r + 1) * shard].T
+        for r in range(tp_size)
+    ]
+    reduced = torch.stack(partials).sum(0)
+    # Reduce-scatter: rank r keeps only its own token shard.
+    for r in range(tp_size):
+        mine = reduced[r * tokens_per_rank : (r + 1) * tokens_per_rank]
+        expected = replicated[r * tokens_per_rank : (r + 1) * tokens_per_rank]
+        torch.testing.assert_close(mine, expected, atol=1e-5, rtol=1e-5)
+
+
 def test_sp_all_gather_uses_custom_kernel(monkeypatch):
     hidden_states = torch.arange(4, dtype=torch.float32).view(2, 2)
     expected = torch.cat([hidden_states, hidden_states])

vllm/envs.py

@@ -195,6 +195,7 @@ if TYPE_CHECKING:
     ] = "relax"
     VLLM_USE_FUSED_MOE_GROUPED_TOPK: bool = True
     VLLM_MOE_SKIP_PADDING: bool = True
+    VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT: bool = False
     VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER: bool = True
     VLLM_USE_FLASHINFER_MOE_INT4: bool = False
     VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR: str | None = None
@@ -1524,6 +1525,18 @@ environment_variables: dict[str, Callable[[], Any]] = {
     # ids to -1 so the dispatch and experts drop them. Requires a MoE kernel that
     # treats topk_id == -1 as a skip sentinel
     "VLLM_MOE_SKIP_PADDING": lambda: bool(int(os.getenv("VLLM_MOE_SKIP_PADDING", "1"))),
+    # Kimi-K3 only. Under sequence-parallel MoE the dense and shared-expert MLPs
+    # are replicated on every rank, so each rank streams the whole weight to
+    # serve its own token shard. Shard them across TP instead: the MLP then
+    # all-gathers the full token set, computes this rank's partial, and
+    # reduce-scatters (which both sums across TP and restores the sequence
+    # sharding). Trades weight bandwidth and resident memory for two collectives
+    # per layer, so it only wins at low token counts: intended for decode
+    # instances in a P/D disaggregated deployment, not for prefill or unified
+    # serving.
+    "VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT": lambda: bool(
+        int(os.getenv("VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT", "0"))
+    ),
     # Allow use of FlashInfer FP8 block-scale GEMM for linear layers.
     # This uses TensorRT-LLM kernels and requires SM90+ (Hopper).
     "VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER": lambda: bool(

vllm/models/kimi_k3/nvidia/model.py

@@ -129,7 +129,42 @@ logger = init_logger(__name__)
 _ROUTED_DOWN_PROJ_STREAM_TOKEN_THRESHOLD = 256
 
 
+def shard_sequence_parallel_mlp(
+    hidden_size: int,
+    intermediate_size: int,
+    use_sequence_parallel: bool,
+    eligible: bool,
+) -> bool:
+    """Whether to TP-shard a sequence-parallel MLP instead of replicating it.
+
+    Opt-in via ``VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT``; see :class:`KimiMLP` for
+    the trade-off and :mod:`vllm.envs` for when it is worth enabling.
+    """
+    enabled = envs.VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT
+    if not (use_sequence_parallel and eligible and enabled):
+        return False
+    tp_size = get_tensor_model_parallel_world_size()
+    return (
+        tp_size > 1 and intermediate_size % tp_size == 0 and hidden_size % tp_size == 0
+    )
+
+
 class KimiMLP(nn.Module):
+    """Dense / shared-expert MLP, optionally TP-sharded under sequence parallel.
+
+    Under sequence parallelism each rank owns a distinct slice of the tokens, so
+    by default both projections are replicated (``disable_tp``) and the block
+    needs no collective. That makes every rank stream the entire weight to serve
+    its own token shard.
+
+    With ``VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT`` the weights are TP-sharded
+    instead. A rank then holds only a slice of the intermediate dim, so it
+    cannot finish its own tokens alone: ``forward`` all-gathers the full token
+    set, computes this rank's partial, and reduce-scatters. The reduce-scatter
+    sums across TP and restores the sequence sharding in one collective, so the
+    block still ends with one collective per direction.
+    """
+
     def __init__(
         self,
         hidden_size: int,
@@ -138,18 +173,27 @@ class KimiMLP(nn.Module):
         quant_config: QuantizationConfig | None = None,
         reduce_results: bool = True,
         use_sequence_parallel: bool = False,
+        can_shard_sequence_parallel: bool = False,
         prefix: str = "",
         activation_situ_beta: float | None = None,
         activation_situ_linear_beta: float | None = None,
     ) -> None:
         super().__init__()
 
+        self.shard_sequence_parallel = shard_sequence_parallel_mlp(
+            hidden_size,
+            intermediate_size,
+            use_sequence_parallel,
+            can_shard_sequence_parallel,
+        )
+        replicate = use_sequence_parallel and not self.shard_sequence_parallel
+
         self.gate_up_proj = MergedColumnParallelLinear(
             hidden_size,
             [intermediate_size] * 2,
             bias=False,
             quant_config=quant_config,
-            disable_tp=use_sequence_parallel,
+            disable_tp=replicate,
             prefix=f"{prefix}.gate_up_proj",
         )
         self.down_proj = RowParallelLinear(
@@ -157,8 +201,10 @@ class KimiMLP(nn.Module):
             hidden_size,
             bias=False,
             quant_config=quant_config,
-            reduce_results=reduce_results,
-            disable_tp=use_sequence_parallel,
+            # Sharded sequence parallel reduces via the reduce-scatter in
+            # forward(), which also restores the sequence sharding.
+            reduce_results=False if self.shard_sequence_parallel else reduce_results,
+            disable_tp=replicate,
             prefix=f"{prefix}.down_proj",
         )
         if hidden_act == "silu":
@@ -175,9 +221,17 @@ class KimiMLP(nn.Module):
             )
 
     def forward(self, x):
+        if self.shard_sequence_parallel:
+            # Each rank holds a weight shard but only its own tokens, so it
+            # cannot finish those tokens alone: gather the full token set,
+            # compute this rank's partial for all of them, then reduce-scatter,
+            # which sums across TP and restores the sequence sharding.
+            x = sp_all_gather(x)
         gate_up, _ = self.gate_up_proj(x)
         x = self.act_fn(gate_up)
         x, _ = self.down_proj(x)
+        if self.shard_sequence_parallel:
+            x = sp_reduce_scatter(x)
         return x
 
 
@@ -478,6 +532,10 @@ class KimiMoE(nn.Module):
                 quant_config=quant_config,
                 reduce_results=False,
                 use_sequence_parallel=use_sequence_parallel,
+                # Only the MegaMoE path calls the shared experts directly; the
+                # FusedMoE path below hands them to the runner, which fuses
+                # their reduction and assumes the replicated layout.
+                can_shard_sequence_parallel=self.use_mega_moe,
                 prefix=f"{prefix}.shared_experts",
                 activation_situ_beta=activation_situ_beta,
                 activation_situ_linear_beta=activation_situ_linear_beta,
@@ -788,6 +846,7 @@ class KimiDecoderLayer(nn.Module):
                 quant_config=quant_config,
                 prefix=f"{prefix}.mlp",
                 use_sequence_parallel=self.use_sequence_parallel,
+                can_shard_sequence_parallel=True,
                 activation_situ_beta=config.activation_situ_beta,
                 activation_situ_linear_beta=config.activation_situ_linear_beta,
             )