@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 4635cc3e8f

Shard the K3 Latent-MoE up-projection on large batches (#50383)

Jee Jee Li committed Aug 3, 2026 at 15:49 UTC 4635cc3e8f609755ad45f9fc2cdb54fa7552038e

2 files changed +176 -72

vllm/model_executor/layers/fused_moe/runner/latent_moe_runner.py

+174 -58

@@ -1,42 +1,82 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +from enum import IntEnum

4 +

5 import torch

7 import vllm.envs as envs

8 from vllm.config import get_current_vllm_config

7 -from vllm.distributed import tensor_model_parallel_all_reduce

9 +from vllm.distributed import (

10 + get_tensor_model_parallel_rank,

11 + tensor_model_parallel_all_reduce,

12 +)

13 from vllm.logger import init_logger

14 from vllm.model_executor.layers.layernorm import RMSNorm

10 -from vllm.utils.torch_utils import aux_stream, current_stream

15 +from vllm.platforms import current_platform

16 +from vllm.utils.multi_stream_utils import maybe_execute_in_parallel

17 +from vllm.utils.torch_utils import aux_stream

19 from .moe_runner import MoERunner, _unpack

21 logger = init_logger(__name__)

24 +class LatentTailTier(IntEnum):

25 + """Which tail implementation the fused path runs, by token count.

26 +

27 + The tiers share the same replicated up-projection weight, so the choice is

28 + per batch and needs no weight relayout: tiers 0 and 2 read it as a

29 + column-parallel row-shard, tier 1 reads it whole.

30 + """

31 +

32 + # ``_small_batch_tail``. Decode-sized (<= the op's max_num_tokens): one

33 + # CuTeDSL collective fuses the latent reduce, RMSNorm and the shared

34 + # reduce-scatter, then a sharded up-projection multicasts through a Lamport

35 + # copy. Requires SM100, TP 8/16, BF16

36 + TAIL_FUSION = 0

37 +

38 + # ``_overlap_allreduce_tail``. The portable default, up to

39 + # VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD tokens: reduce the latent,

40 + # up-project the full hidden dim from the replicated weight, and add the

41 + # separately reduced shared output, hiding that shared all-reduce behind the

42 + # up-projection GEMM on the aux stream.

43 + ALLREDUCE_OVERLAP = 1

44 +

45 + # ``_shard_up_proj_tail``. Prefill-sized: each rank up-projects only its

46 + # hidden shard and accumulates into the shared partial, so the shared

47 + # all-reduce also stitches the routed shards. Same two all-reduces as tier 1

48 + # at 1/tp of the up-projection FLOPs; it gives up tier 1's aux-stream

49 + # overlap, since the reduce now has to follow the accumulate.

50 + COLUMN_PARALLEL = 2

51 +

52 +

53 class LatentMoERunner(MoERunner):

54 """MoE runner for latent MoE with a replicated routed up-projection.

20 - Fused path (tp>1, un-reduced combine output, shared expert, no SP):

21 - concatenates the un-reduced latent partial (dim d) and the un-reduced

22 - shared partial (dim D) into one contiguous buffer, all-reduces once, then

23 - splits. The latent half is normed and up-projected locally (replicated

24 - up-proj -> full hidden), and the shared add folds into the GEMM epilogue

25 - (``torch.addmm``). One collective total, no post-reduction communication.

56 + The fused path (tp>1, un-reduced combine output, shared expert, no SP)

57 + dispatches over ``LatentTailTier`` by token count; see that enum for what

58 + each tier does and when it applies.

60 Native path: the replicated up-proj produces the full hidden dim on every

28 - rank, so the base runner combines routed + shared correctly at any TP size

29 - (using two collectives instead of the fused path's one).

61 + rank, so the base runner combines routed + shared correctly at any TP size.

62 """

64 def __init__(

65 self,

66 *args,

35 - enable_k3_latent_moe_tail_fusion: bool = False,

67 **kwargs,

68 ) -> None:

69 super().__init__(*args, **kwargs)

39 - self.enable_k3_latent_moe_tail_fusion = enable_k3_latent_moe_tail_fusion

70 +

71 + # The tail-fusion kernels are tcgen05-based, so they require an

72 + # SM100 NVIDIA device; the runner falls back to the default latent

73 + # MoE path everywhere else.

74 + self.enable_k3_latent_moe_tail_fusion = (

75 + current_platform.is_cuda()

76 + and current_platform.is_device_capability_family(100)

77 + )

78 + # Overlap the shared-expert all-reduce with the tier-1 up-projection.

79 + self._shared_ar_events = (torch.cuda.Event(), torch.cuda.Event())

80 use_fused_path = self._use_fused_path()

81 if (

82 self.enable_k3_latent_moe_tail_fusion

@@ -113,6 +153,121 @@ class LatentMoERunner(MoERunner):

153 and not self.moe_config.is_sequence_parallel

154 )

155

156 + def _select_tail_tier(

157 + self,

158 + fused_output: torch.Tensor,

159 + shared_output: torch.Tensor,

160 + ) -> LatentTailTier:

161 + num_tokens = fused_output.shape[0]

162 + # tier 0

163 + if self.enable_k3_latent_moe_tail_fusion and (

164 + 0 < num_tokens <= self._k3_latent_moe_tail_op.contract.max_num_tokens

165 + ):

166 + return LatentTailTier.TAIL_FUSION

167 +

168 + transform = self.routed_output_transform

169 + assert transform is not None

170 + # tier 1

171 + if (

172 + num_tokens <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD

173 + and not envs.VLLM_DISABLE_SHARED_EXPERTS_STREAM

174 + ):

175 + return LatentTailTier.ALLREDUCE_OVERLAP

176 + # tier 2

177 + return LatentTailTier.COLUMN_PARALLEL

178 +

179 + def _small_batch_tail(

180 + self,

181 + fused_output: torch.Tensor,

182 + shared_output: torch.Tensor,

183 + trunc_size: int | None,

184 + ) -> torch.Tensor:

185 + """Tier 0: the CuTeDSL operator fuses the whole tail."""

186 + transform = self.routed_output_transform

187 + assert transform is not None

188 + norm = transform.norm

189 + assert norm is not None

190 +

191 + result = self._k3_latent_moe_tail_op(

192 + fused_output,

193 + shared_output,

194 + norm.weight,

195 + transform.up_proj.weight,

196 + )

197 + # The operator already reduced; this only strips padding.

198 + return self._maybe_reduce_final_output(

199 + result, trunc_size, output_is_reduced=True

200 + )

201 +

202 + def _overlap_allreduce_tail(

203 + self,

204 + fused_output: torch.Tensor,

205 + shared_output: torch.Tensor,

206 + trunc_size: int | None,

207 + ) -> torch.Tensor:

208 + """Tier 1: reduce the latent, up-project the full hidden dim from the

209 + replicated weight, and add the separately reduced shared output.

210 +

211 + Small enough batches hide that shared all-reduce behind the up-projection

212 + GEMM on the aux stream.

213 + """

214 + transform = self.routed_output_transform

215 + assert transform is not None

216 + assert shared_output.size(0) <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD

217 + if transform.norm is not None:

218 + fused_latent = self.allreduce_norm_latent_out(fused_output, transform.norm)

219 + else:

220 + fused_latent = tensor_model_parallel_all_reduce(fused_output)

221 +

222 + # Overlap the shared-expert all-reduce with the up-projection GEMM while

223 + # the batch is small enough for it to pay off.

224 + result, shared_output = maybe_execute_in_parallel(

225 + lambda: torch.mm(fused_latent, transform.up_proj.weight.t()),

226 + lambda: tensor_model_parallel_all_reduce(shared_output),

227 + self._shared_ar_events[0],

228 + self._shared_ar_events[1],

229 + aux_stream(),

230 + )

231 + result.add_(shared_output)

232 +

233 + # Output is already fully reduced; this only strips padding.

234 + return self._maybe_reduce_final_output(

235 + result, trunc_size, output_is_reduced=True

236 + )

237 +

238 + def _shard_up_proj_tail(

239 + self,

240 + fused_output: torch.Tensor,

241 + shared_output: torch.Tensor,

242 + trunc_size: int | None,

243 + ) -> torch.Tensor:

244 + """

245 + Tier 2: column-parallel up-projection folded into the final reduce.

246 + """

247 + transform = self.routed_output_transform

248 + assert transform is not None

249 +

250 + if transform.norm is not None:

251 + latent = self.allreduce_norm_latent_out(fused_output, transform.norm)

252 + else:

253 + latent = tensor_model_parallel_all_reduce(fused_output)

254 +

255 + weight = transform.up_proj.weight

256 + shard_size = weight.shape[0] // self.moe_config.tp_size

257 + shard_start = get_tensor_model_parallel_rank() * shard_size

258 +

259 + # column-parallel

260 + up_proj_shard = weight.narrow(0, shard_start, shard_size)

261 + hidden_shard = shared_output.narrow(-1, shard_start, shard_size)

262 +

263 + # hidden_shard += latent @ up_proj_shard.T, accumulated in the GEMM's

264 + # beta-add epilogue so folding in the shared partial costs no kernel.

265 + hidden_shard.addmm_(latent, up_proj_shard.t())

266 +

267 + return self._maybe_reduce_final_output(

268 + shared_output, trunc_size, output_is_reduced=False

269 + )

270 +

271 def forward(

272 self,

273 hidden_states: torch.Tensor,

@@ -167,54 +322,15 @@ class LatentMoERunner(MoERunner):

322 if og_hidden_dim_pre_xform is not None:

323 fused_output = fused_output[..., :og_hidden_dim_pre_xform]

324

170 - transform = self.routed_output_transform

171 - assert transform is not None

172 -

173 - if self.enable_k3_latent_moe_tail_fusion:

174 - op = self._k3_latent_moe_tail_op

175 - if 0 < fused_output.shape[0] <= op.contract.max_num_tokens:

176 - norm = transform.norm

177 - assert norm is not None

178 - result = op(

179 - fused_output,

180 - shared_output,

181 - norm.weight,

182 - transform.up_proj.weight,

183 - )

184 - result = self._maybe_reduce_final_output(

185 - result, og_hidden_dim_post_xform, output_is_reduced=True

186 - )

187 - return self._maybe_add_zero_expert_output(result)

188 -

189 - fused_latent = None

190 - if transform.norm is not None:

191 - fused_latent = self.allreduce_norm_latent_out(fused_output, transform.norm)

325 + tier = self._select_tail_tier(fused_output, shared_output)

326 + if tier is LatentTailTier.TAIL_FUSION:

327 + latent_tail = self._small_batch_tail

328 + elif tier is LatentTailTier.ALLREDUCE_OVERLAP:

329 + latent_tail = self._overlap_allreduce_tail

330 else:

193 - fused_latent = tensor_model_parallel_all_reduce(fused_output)

331 + latent_tail = self._shard_up_proj_tail

332

195 - shared_expert_stream = (

196 - aux_stream()

197 - if shared_output.size(0) <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD

198 - else None

199 - )

200 - if shared_expert_stream is not None:

201 - # overlap shared expert allreduce with latent up_proj

202 - main = current_stream()

203 - shared_output.record_stream(shared_expert_stream)

204 - shared_expert_stream.wait_stream(main)

205 - with torch.cuda.stream(shared_expert_stream):

206 - shared_output = tensor_model_parallel_all_reduce(shared_output)

207 - result = torch.mm(fused_latent, transform.up_proj.weight.t())

208 - main.wait_stream(shared_expert_stream)

209 - else:

210 - shared_output = tensor_model_parallel_all_reduce(shared_output)

211 - result = torch.mm(fused_latent, transform.up_proj.weight.t())

212 - result.add_(shared_output)

213 -

214 - # Output is already fully reduced; this only strips padding.

215 - result = self._maybe_reduce_final_output(

216 - result, og_hidden_dim_post_xform, output_is_reduced=True

217 - )

333 + result = latent_tail(fused_output, shared_output, og_hidden_dim_post_xform)

334

335 return self._maybe_add_zero_expert_output(result)

336

vllm/models/kimi_k3/nvidia/model.py

+2 -14

@@ -472,8 +472,8 @@ class KimiMoE(nn.Module):

472 if self.num_shared_experts is not None:

473 shared_intermediate_size = moe_intermediate_size * self.num_shared_experts

474 self.shared_experts = KimiMLP(

475 - hidden_size=config.hidden_size,

476 - intermediate_size=shared_intermediate_size,

475 + hidden_size=config.hidden_size, # 7618

476 + intermediate_size=shared_intermediate_size, # 3072*2

477 hidden_act=config.hidden_act,

478 quant_config=quant_config,

479 reduce_results=False,

@@ -553,13 +553,6 @@ class KimiMoE(nn.Module):

553 activation_linear_beta=activation_situ_linear_beta,

554 )

555 else:

556 - # The tail-fusion kernels are tcgen05-based, so they require an

557 - # SM100 NVIDIA device; the runner falls back to the default latent

558 - # MoE path everywhere else.

559 - enable_tail_fusion = (

560 - current_platform.is_cuda()

561 - and current_platform.is_device_capability_family(100)

562 - )

556 self.experts = FusedMoEFactory(

557 shared_experts=self.shared_experts,

558 num_experts=num_experts,

@@ -586,11 +579,6 @@ class KimiMoE(nn.Module):

579 routed_output_transform=self.routed_output_transform,

580 is_sequence_parallel=use_sequence_parallel,

581 runner_cls=LatentMoERunner if self.use_latent_moe else None,

589 - runner_args=(

590 - {"enable_k3_latent_moe_tail_fusion": enable_tail_fusion}

591 - if self.use_latent_moe

592 - else None

593 - ),

582 )

583 if self.padded_moe_intermediate_size != moe_intermediate_size:

584 w13_weight = getattr(self.experts, "w13_weight", None)