[Kimi K3 Perf] option to shard the shared expert for non mega case, 16.98 GiB memory/GPU saved (#50912)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Aug 4, 2026 at 14:00 UTC
d31de3c421c0281a959ac1a3cbe1a7f354bae179
3 files changed
+10
-18
tests/models/kimi_k3/test_sequence_parallel.py
+6
-9
@@ -267,21 +267,19 @@ def test_kimi_mtp_restores_sequence_parallel_output(monkeypatch):
267
268
269
@pytest.mark.parametrize(
270
- ("enabled", "use_sequence_parallel", "eligible", "tp_size", "expected"),
270
+ ("enabled", "use_sequence_parallel", "tp_size", "expected"),
271
[
272
- (True, True, True, 8, True),
273
- (False, True, True, 8, False), # opt-in only
274
- (True, False, True, 8, False), # replication only exists under SP
275
- (True, True, False, 8, False), # FusedMoE path owns the reduction
276
- (True, True, True, 1, False), # nothing to shard
277
- (True, True, True, 5, False), # 6144 % 5 -- would fail divide()
272
+ (True, True, 8, True),
273
+ (False, True, 8, False), # opt-in only
274
+ (True, False, 8, False), # replication only exists under SP
275
+ (True, True, 1, False), # nothing to shard
276
+ (True, True, 5, False), # 6144 % 5 -- would fail divide()
277
],
278
)
279
def test_shard_sequence_parallel_mlp_gating(
280
monkeypatch,
281
enabled: bool,
282
use_sequence_parallel: bool,
284
- eligible: bool,
283
tp_size: int,
284
expected: bool,
285
):
@@ -295,7 +293,6 @@ def test_shard_sequence_parallel_mlp_gating(
293
hidden_size=7168,
294
intermediate_size=6144,
295
use_sequence_parallel=use_sequence_parallel,
298
- eligible=eligible,
296
)
297
is expected
298
)
vllm/model_executor/layers/fused_moe/runner/shared_experts.py
+3
@@ -81,6 +81,9 @@ class SharedExperts(torch.nn.Module):
81
# - we are using eplb with non-default backend, because of correctness issues
82
# - we are using flashinfer with DP, since there nothing to gain
83
parallel_config = self._moe_config.moe_parallel_config
84
+ if getattr(self._layer, "shard_sequence_parallel", False):
85
+ # TODO: we may enable this to optimize further
86
+ return True
87
return (
88
parallel_config.enable_eplb
89
and parallel_config.all2all_backend != "allgather_reducescatter"
vllm/models/kimi_k3/nvidia/model.py
+1
-9
@@ -133,7 +133,6 @@ def shard_sequence_parallel_mlp(
133
hidden_size: int,
134
intermediate_size: int,
135
use_sequence_parallel: bool,
136
- eligible: bool,
136
) -> bool:
137
"""Whether to TP-shard a sequence-parallel MLP instead of replicating it.
138
@@ -141,7 +140,7 @@ def shard_sequence_parallel_mlp(
140
the trade-off and :mod:`vllm.envs` for when it is worth enabling.
141
"""
142
enabled = envs.VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT
144
- if not (use_sequence_parallel and eligible and enabled):
143
+ if not (use_sequence_parallel and enabled):
144
return False
145
tp_size = get_tensor_model_parallel_world_size()
146
return (
@@ -173,7 +172,6 @@ class KimiMLP(nn.Module):
172
quant_config: QuantizationConfig | None = None,
173
reduce_results: bool = True,
174
use_sequence_parallel: bool = False,
176
- can_shard_sequence_parallel: bool = False,
175
prefix: str = "",
176
activation_situ_beta: float | None = None,
177
activation_situ_linear_beta: float | None = None,
@@ -184,7 +182,6 @@ class KimiMLP(nn.Module):
182
hidden_size,
183
intermediate_size,
184
use_sequence_parallel,
187
- can_shard_sequence_parallel,
185
)
186
replicate = use_sequence_parallel and not self.shard_sequence_parallel
187
@@ -532,10 +529,6 @@ class KimiMoE(nn.Module):
529
quant_config=quant_config,
530
reduce_results=False,
531
use_sequence_parallel=use_sequence_parallel,
535
- # Only the MegaMoE path calls the shared experts directly; the
536
- # FusedMoE path below hands them to the runner, which fuses
537
- # their reduction and assumes the replicated layout.
538
- can_shard_sequence_parallel=self.use_mega_moe,
532
prefix=f"{prefix}.shared_experts",
533
activation_situ_beta=activation_situ_beta,
534
activation_situ_linear_beta=activation_situ_linear_beta,
@@ -846,7 +839,6 @@ class KimiDecoderLayer(nn.Module):
839
quant_config=quant_config,
840
prefix=f"{prefix}.mlp",
841
use_sequence_parallel=self.use_sequence_parallel,
849
- can_shard_sequence_parallel=True,
842
activation_situ_beta=config.activation_situ_beta,
843
activation_situ_linear_beta=config.activation_situ_linear_beta,
844
)