diff --git a/tests/models/kimi_k3/test_sequence_parallel.py b/tests/models/kimi_k3/test_sequence_parallel.py index b61bb498cf..9e8faea099 100644 --- a/tests/models/kimi_k3/test_sequence_parallel.py +++ b/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]) diff --git a/vllm/envs.py b/vllm/envs.py index 4c247e6612..da45b92703 100755 --- a/vllm/envs.py +++ b/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( diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 4d97c0a507..a47d4ed08b 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/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, )