@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 5df9999fcf

[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>

Tyler Michael Smith committed Aug 3, 2026 at 11:41 UTC 5df9999fcfaa72d9eb61348a789058fff805f142

3 files changed +154 -3

tests/models/kimi_k3/test_sequence_parallel.py

+79

@@ -10,6 +10,7 @@ import torch

10 from torch import nn

12 from vllm.config import ParallelConfig

13 +from vllm.model_executor.layers.activation import SiluAndMul

14 from vllm.models.common.ops import sequence_parallel as sp_ops

15 from vllm.models.kimi_k3.nvidia import model as kimi_model

16 from vllm.models.kimi_k3.nvidia import mtp as kimi_mtp

@@ -265,6 +266,84 @@ def test_kimi_mtp_restores_sequence_parallel_output(monkeypatch):

266 torch.testing.assert_close(final_norm.call_args.args[0], expected_hidden_states)

267

268

269 +@pytest.mark.parametrize(

270 + ("enabled", "use_sequence_parallel", "eligible", "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()

278 + ],

279 +)

280 +def test_shard_sequence_parallel_mlp_gating(

281 + monkeypatch,

282 + enabled: bool,

283 + use_sequence_parallel: bool,

284 + eligible: bool,

285 + tp_size: int,

286 + expected: bool,

287 +):

288 + monkeypatch.setattr(kimi_model.envs, "VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT", enabled)

289 + monkeypatch.setattr(

290 + kimi_model, "get_tensor_model_parallel_world_size", lambda: tp_size

291 + )

292 +

293 + assert (

294 + kimi_model.shard_sequence_parallel_mlp(

295 + hidden_size=7168,

296 + intermediate_size=6144,

297 + use_sequence_parallel=use_sequence_parallel,

298 + eligible=eligible,

299 + )

300 + is expected

301 + )

302 +

303 +

304 +def test_sharded_sequence_parallel_mlp_matches_replicated(default_vllm_config):

305 + """Sharded SP MLP must reproduce the replicated result for every token.

306 +

307 + Each rank owns a *disjoint* token shard, so a weight shard alone cannot

308 + finish a rank's own tokens: the ranks must gather the full token set,

309 + compute partial sums over their intermediate shard, and reduce-scatter.

310 + Splicing per-rank feature slices together instead silently mixes different

311 + tokens and produces plausible-looking garbage.

312 + """

313 + tp_size, hidden, intermediate, tokens_per_rank = 4, 16, 12, 3

314 + torch.manual_seed(0)

315 + num_tokens = tp_size * tokens_per_rank

316 + x = torch.randn(num_tokens, hidden)

317 + gate_weight = torch.randn(intermediate, hidden)

318 + up_weight = torch.randn(intermediate, hidden)

319 + down_weight = torch.randn(hidden, intermediate)

320 + act_fn = SiluAndMul()

321 +

322 + replicated = act_fn(x @ torch.cat([gate_weight, up_weight]).T) @ down_weight.T

323 +

324 + shard = intermediate // tp_size

325 + # Every rank all-gathers the full token set, then computes its partial.

326 + partials = [

327 + act_fn(

328 + x

329 + @ torch.cat(

330 + [

331 + gate_weight[r * shard : (r + 1) * shard],

332 + up_weight[r * shard : (r + 1) * shard],

333 + ]

334 + ).T

335 + )

336 + @ down_weight[:, r * shard : (r + 1) * shard].T

337 + for r in range(tp_size)

338 + ]

339 + reduced = torch.stack(partials).sum(0)

340 + # Reduce-scatter: rank r keeps only its own token shard.

341 + for r in range(tp_size):

342 + mine = reduced[r * tokens_per_rank : (r + 1) * tokens_per_rank]

343 + expected = replicated[r * tokens_per_rank : (r + 1) * tokens_per_rank]

344 + torch.testing.assert_close(mine, expected, atol=1e-5, rtol=1e-5)

345 +

346 +

347 def test_sp_all_gather_uses_custom_kernel(monkeypatch):

348 hidden_states = torch.arange(4, dtype=torch.float32).view(2, 2)

349 expected = torch.cat([hidden_states, hidden_states])

vllm/envs.py

+13

@@ -195,6 +195,7 @@ if TYPE_CHECKING:

195 ] = "relax"

196 VLLM_USE_FUSED_MOE_GROUPED_TOPK: bool = True

197 VLLM_MOE_SKIP_PADDING: bool = True

198 + VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT: bool = False

199 VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER: bool = True

200 VLLM_USE_FLASHINFER_MOE_INT4: bool = False

201 VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR: str | None = None

@@ -1524,6 +1525,18 @@ environment_variables: dict[str, Callable[[], Any]] = {

1525 # ids to -1 so the dispatch and experts drop them. Requires a MoE kernel that

1526 # treats topk_id == -1 as a skip sentinel

1527 "VLLM_MOE_SKIP_PADDING": lambda: bool(int(os.getenv("VLLM_MOE_SKIP_PADDING", "1"))),

1528 + # Kimi-K3 only. Under sequence-parallel MoE the dense and shared-expert MLPs

1529 + # are replicated on every rank, so each rank streams the whole weight to

1530 + # serve its own token shard. Shard them across TP instead: the MLP then

1531 + # all-gathers the full token set, computes this rank's partial, and

1532 + # reduce-scatters (which both sums across TP and restores the sequence

1533 + # sharding). Trades weight bandwidth and resident memory for two collectives

1534 + # per layer, so it only wins at low token counts: intended for decode

1535 + # instances in a P/D disaggregated deployment, not for prefill or unified

1536 + # serving.

1537 + "VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT": lambda: bool(

1538 + int(os.getenv("VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT", "0"))

1539 + ),

1540 # Allow use of FlashInfer FP8 block-scale GEMM for linear layers.

1541 # This uses TensorRT-LLM kernels and requires SM90+ (Hopper).

1542 "VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER": lambda: bool(

vllm/models/kimi_k3/nvidia/model.py

+62 -3

@@ -129,7 +129,42 @@ logger = init_logger(__name__)

129 _ROUTED_DOWN_PROJ_STREAM_TOKEN_THRESHOLD = 256

130

131

132 +def shard_sequence_parallel_mlp(

133 + hidden_size: int,

134 + intermediate_size: int,

135 + use_sequence_parallel: bool,

136 + eligible: bool,

137 +) -> bool:

138 + """Whether to TP-shard a sequence-parallel MLP instead of replicating it.

139 +

140 + Opt-in via ``VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT``; see :class:`KimiMLP` for

141 + the trade-off and :mod:`vllm.envs` for when it is worth enabling.

142 + """

143 + enabled = envs.VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT

144 + if not (use_sequence_parallel and eligible and enabled):

145 + return False

146 + tp_size = get_tensor_model_parallel_world_size()

147 + return (

148 + tp_size > 1 and intermediate_size % tp_size == 0 and hidden_size % tp_size == 0

149 + )

150 +

151 +

152 class KimiMLP(nn.Module):

153 + """Dense / shared-expert MLP, optionally TP-sharded under sequence parallel.

154 +

155 + Under sequence parallelism each rank owns a distinct slice of the tokens, so

156 + by default both projections are replicated (``disable_tp``) and the block

157 + needs no collective. That makes every rank stream the entire weight to serve

158 + its own token shard.

159 +

160 + With ``VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT`` the weights are TP-sharded

161 + instead. A rank then holds only a slice of the intermediate dim, so it

162 + cannot finish its own tokens alone: ``forward`` all-gathers the full token

163 + set, computes this rank's partial, and reduce-scatters. The reduce-scatter

164 + sums across TP and restores the sequence sharding in one collective, so the

165 + block still ends with one collective per direction.

166 + """

167 +

168 def __init__(

169 self,

170 hidden_size: int,

@@ -138,18 +173,27 @@ class KimiMLP(nn.Module):

173 quant_config: QuantizationConfig | None = None,

174 reduce_results: bool = True,

175 use_sequence_parallel: bool = False,

176 + can_shard_sequence_parallel: bool = False,

177 prefix: str = "",

178 activation_situ_beta: float | None = None,

179 activation_situ_linear_beta: float | None = None,

180 ) -> None:

181 super().__init__()

182

183 + self.shard_sequence_parallel = shard_sequence_parallel_mlp(

184 + hidden_size,

185 + intermediate_size,

186 + use_sequence_parallel,

187 + can_shard_sequence_parallel,

188 + )

189 + replicate = use_sequence_parallel and not self.shard_sequence_parallel

190 +

191 self.gate_up_proj = MergedColumnParallelLinear(

192 hidden_size,

193 [intermediate_size] * 2,

194 bias=False,

195 quant_config=quant_config,

152 - disable_tp=use_sequence_parallel,

196 + disable_tp=replicate,

197 prefix=f"{prefix}.gate_up_proj",

198 )

199 self.down_proj = RowParallelLinear(

@@ -157,8 +201,10 @@ class KimiMLP(nn.Module):

201 hidden_size,

202 bias=False,

203 quant_config=quant_config,

160 - reduce_results=reduce_results,

161 - disable_tp=use_sequence_parallel,

204 + # Sharded sequence parallel reduces via the reduce-scatter in

205 + # forward(), which also restores the sequence sharding.

206 + reduce_results=False if self.shard_sequence_parallel else reduce_results,

207 + disable_tp=replicate,

208 prefix=f"{prefix}.down_proj",

209 )

210 if hidden_act == "silu":

@@ -175,9 +221,17 @@ class KimiMLP(nn.Module):

221 )

222

223 def forward(self, x):

224 + if self.shard_sequence_parallel:

225 + # Each rank holds a weight shard but only its own tokens, so it

226 + # cannot finish those tokens alone: gather the full token set,

227 + # compute this rank's partial for all of them, then reduce-scatter,

228 + # which sums across TP and restores the sequence sharding.

229 + x = sp_all_gather(x)

230 gate_up, _ = self.gate_up_proj(x)

231 x = self.act_fn(gate_up)

232 x, _ = self.down_proj(x)

233 + if self.shard_sequence_parallel:

234 + x = sp_reduce_scatter(x)

235 return x

236

237

@@ -478,6 +532,10 @@ class KimiMoE(nn.Module):

532 quant_config=quant_config,

533 reduce_results=False,

534 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,

539 prefix=f"{prefix}.shared_experts",

540 activation_situ_beta=activation_situ_beta,

541 activation_situ_linear_beta=activation_situ_linear_beta,

@@ -788,6 +846,7 @@ class KimiDecoderLayer(nn.Module):

846 quant_config=quant_config,

847 prefix=f"{prefix}.mlp",

848 use_sequence_parallel=self.use_sequence_parallel,

849 + can_shard_sequence_parallel=True,

850 activation_situ_beta=config.activation_situ_beta,

851 activation_situ_linear_beta=config.activation_situ_linear_beta,

852 )