@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f0c98cae27

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 02:15 UTC

[Misc] MoE ModularKernel : Introduce TopKWeightAndReduce (#20648)

Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com> Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com>

Varun Sundar Rabindranath committed Jul 10, 2025 at 17:40 UTC f0c98cae2758b1a706537aa412c6868bb060c151

14 files changed +297 -59

tests/kernels/moe/test_pplx_moe.py

@@ -32,6 +32,8 @@ from vllm.model_executor.layers.fused_moe.fused_batched_moe import (

32 from vllm.model_executor.layers.fused_moe.fused_moe import get_default_config

33 from vllm.model_executor.layers.fused_moe.modular_kernel import (

34 FusedMoEModularKernel)

35 +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

36 + TopKWeightAndReduceDelegate)

37 from vllm.platforms import current_platform

38 from vllm.utils import round_up

@@ -371,6 +373,7 @@ def pplx_prepare_finalize(

373 chunk_topk_weight,

374 chunk_topk_ids,

375 False,

376 + weight_and_reduce_impl=TopKWeightAndReduceDelegate(),

377 )

378

379 torch.cuda.synchronize()

vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py

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

7 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

8 from vllm.logger import init_logger

9 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig

10 +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

11 + TopKWeightAndReduceDelegate)

12 from vllm.model_executor.layers.fused_moe.utils import _resize_cache

13 from vllm.triton_utils import tl, triton

@@ -217,6 +219,10 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

219 def supports_expert_map(self) -> bool:

220 return False

221

222 + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:

223 + # Let PrepareAndFinalize::finalize() decide the impl.

224 + return TopKWeightAndReduceDelegate()

225 +

226 def workspace_shapes(

227 self,

228 a: torch.Tensor,

vllm/model_executor/layers/fused_moe/batched_triton_or_deep_gemm_moe.py

+19

@@ -88,6 +88,25 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

88 return ((bdge is None or bdge.supports_expert_map())

89 and (bte is None or bte.supports_expert_map()))

91 + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:

92 + bdge = self.batched_deep_gemm_experts

93 + bte = self.batched_triton_experts

94 + bdge_war = bdge.finalize_weight_and_reduce_impl() if bdge else None

95 + bte_war = bte.finalize_weight_and_reduce_impl() if bte else None

96 + is_bdge_war = bdge_war is not None

97 + is_bte_war = bte_war is not None

98 +

99 + if is_bdge_war and is_bte_war:

100 + assert bdge_war == bte_war, (

101 + "Both implementations should agree on WeightAndReduce impls. "

102 + f"Got bdge_war: {bdge_war}, and bte_war: {bte_war}")

103 +

104 + if bdge_war is not None:

105 + return bdge_war

106 +

107 + assert bte_war is not None

108 + return bte_war

109 +

110 def workspace_shapes(

111 self,

112 a: torch.Tensor,

vllm/model_executor/layers/fused_moe/cutlass_moe.py

@@ -11,6 +11,8 @@ from vllm.logger import init_logger

11 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig

12 from vllm.model_executor.layers.fused_moe.prepare_finalize import (

13 MoEPrepareAndFinalizeNoEP)

14 +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

15 + TopKWeightAndReduceDelegate)

16 from vllm.model_executor.layers.fused_moe.utils import (_fp8_perm,

17 _fp8_quantize,

18 _resize_cache)

@@ -255,6 +257,10 @@ class CutlassExpertsFp8(mk.FusedMoEPermuteExpertsUnpermute):

257 def supports_expert_map(self) -> bool:

258 return not self.use_batched_format

259

260 + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:

261 + # Let PrepareAndFinalize::finalize() decide the impl.

262 + return TopKWeightAndReduceDelegate()

263 +

264 def workspace_shapes(

265 self,

266 a: torch.Tensor,

vllm/model_executor/layers/fused_moe/deep_gemm_moe.py

@@ -12,6 +12,8 @@ from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (

12 _moe_permute)

13 from vllm.model_executor.layers.fused_moe.prepare_finalize import (

14 MoEPrepareAndFinalizeNoEP)

15 +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

16 + TopKWeightAndReduceDelegate)

17 from vllm.model_executor.layers.fused_moe.utils import (

18 _resize_cache, per_token_group_quant_fp8)

19 from vllm.utils import has_deep_gemm, round_up

@@ -85,6 +87,10 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

87 def supports_expert_map(self) -> bool:

88 return True

90 + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:

91 + # Let PrepareAndFinalize::finalize() decide the impl.

92 + return TopKWeightAndReduceDelegate()

93 +

94 def workspace_shapes(

95 self, a: torch.Tensor, aq: torch.Tensor, M: int, N: int, K: int,

96 topk: int, global_num_experts: int, local_num_experts: int

vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py

+10 -29

@@ -6,8 +6,9 @@ import deep_ep

6 import torch

8 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

9 -from vllm import _custom_ops as ops

9 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig

10 +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

11 + TopKWeightAndReduceContiguous, TopKWeightAndReduceDelegate)

12 from vllm.model_executor.layers.fused_moe.utils import (

13 moe_kernel_quantize_input)

@@ -187,45 +188,25 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

188 return (expert_x, expert_x_scale, expert_tokens_meta, expert_topk_ids,

189 expert_topk_weights)

190

190 - def _apply_weights_and_reduce(self, num_tokens: int,

191 - fused_expert_output: torch.Tensor,

192 - topk_weights: torch.Tensor,

193 - apply_router_weight_on_input: bool,

194 - output_dtype: torch.dtype):

195 -

196 - hidden_dim = fused_expert_output.size(-1)

197 - if fused_expert_output.ndim == 2:

198 - fused_expert_output = fused_expert_output.view(

199 - num_tokens, -1, hidden_dim)

200 -

201 - if not apply_router_weight_on_input:

202 - # The DeepEP combine kernels don't do the topk weight

203 - # multiplication. We multiply the weights locally.

204 - m_x_topk = fused_expert_output.size(0)

205 - fused_expert_output.mul_(topk_weights.view(m_x_topk, -1, 1))

206 -

207 - out = torch.empty((num_tokens, hidden_dim),

208 - device=fused_expert_output.device,

209 - dtype=output_dtype)

210 - ops.moe_sum(fused_expert_output, out)

211 -

212 - return out

213 -

191 def finalize(self, output: torch.Tensor, fused_expert_output: torch.Tensor,

192 topk_weights: torch.Tensor, topk_ids: torch.Tensor,

216 - apply_router_weight_on_input: bool) -> None:

193 + apply_router_weight_on_input: bool,

194 + weight_and_reduce_impl: mk.TopKWeightAndReduce) -> None:

195

196 assert self.handle is not None

197

198 # fused_expert_output can have 0 tokens - This happens when none of the

199 # tokens from the all2all reach this EP rank.

200 if fused_expert_output.numel() != 0:

223 - fused_expert_output = self._apply_weights_and_reduce(

224 - num_tokens=topk_ids.size(0),

201 + if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate):

202 + weight_and_reduce_impl = TopKWeightAndReduceContiguous()

203 + fused_expert_output = weight_and_reduce_impl.apply(

204 + output=None,

205 fused_expert_output=fused_expert_output,

206 topk_weights=topk_weights,

207 + topk_ids=topk_ids,

208 apply_router_weight_on_input=apply_router_weight_on_input,

228 - output_dtype=output.dtype)

209 + )

210

211 combined_x, _, event = self.buffer.combine(

212 x=fused_expert_output,

vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py

+7 -2

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

8 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

9 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig

10 +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

11 + TopKWeightAndReduceDelegate)

12 from vllm.model_executor.layers.fused_moe.utils import (

13 moe_kernel_quantize_input, normalize_batched_scales_shape)

@@ -166,8 +168,11 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

168

169 def finalize(self, output: torch.Tensor, fused_expert_output: torch.Tensor,

170 topk_weights: torch.Tensor, topk_ids: torch.Tensor,

169 - apply_router_weight_on_input: bool) -> None:

170 -

171 + apply_router_weight_on_input: bool,

172 + weight_and_reduce_impl: mk.TopKWeightAndReduce) -> None:

173 + assert isinstance(

174 + weight_and_reduce_impl, TopKWeightAndReduceDelegate

175 + ), ("Weight application and reduction happens in the combine kernel.")

176 assert self.handle is not None

177

178 combine_topk_weights = topk_weights

vllm/model_executor/layers/fused_moe/fused_batched_moe.py

+20 -18

@@ -11,6 +11,8 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk

11 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig

12 from vllm.model_executor.layers.fused_moe.fused_moe import (

13 get_config_dtype_str, try_get_optimal_moe_config)

14 +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

15 + TopKWeightAndReduceDelegate, TopKWeightAndReduceNaiveBatched)

16 from vllm.model_executor.layers.fused_moe.utils import (

17 _resize_cache, moe_kernel_quantize_input, normalize_batched_scales_shape,

18 normalize_scales_shape)

@@ -600,25 +602,17 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

602 topk_weights: torch.Tensor,

603 topk_ids: torch.Tensor,

604 apply_router_weight_on_input: bool,

605 + weight_and_reduce_impl: mk.TopKWeightAndReduce,

606 ) -> None:

604 - num_tokens = topk_ids.size(0)

605 - num_local_experts = fused_expert_output.size(0)

606 - K = fused_expert_output.size(-1)

607 - assert output.size(0) == num_tokens and output.size(1) == K

608 -

609 - output.fill_(0)

610 -

611 - first_expert = num_local_experts * self.rank

612 - last_expert = first_expert + num_local_experts

613 -

614 - for expert_id in range(first_expert, last_expert):

615 - matching_tokens = topk_ids == expert_id

616 - topks = torch.any(matching_tokens, dim=1).flatten()

617 - rows = torch.count_nonzero(topks)

618 - rhs = fused_expert_output[expert_id - first_expert, :rows, :]

619 - if not apply_router_weight_on_input:

620 - rhs.mul_(topk_weights[matching_tokens].view(rhs.size(0), 1))

621 - output[topks] = output[topks] + rhs

607 + if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate):

608 + weight_and_reduce_impl = TopKWeightAndReduceNaiveBatched(self.rank)

609 + weight_and_reduce_impl.apply(

610 + output=output,

611 + fused_expert_output=fused_expert_output,

612 + topk_weights=topk_weights,

613 + topk_ids=topk_ids,

614 + apply_router_weight_on_input=apply_router_weight_on_input,

615 + )

616

617

618 class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):

@@ -670,6 +664,10 @@ class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):

664 def supports_expert_map(self) -> bool:

665 return False

666

667 + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:

668 + # Let PrepareAndFinalize::finalize() decide the impl.

669 + return TopKWeightAndReduceDelegate()

670 +

671 def workspace_shapes(

672 self,

673 a: torch.Tensor,

@@ -877,6 +875,10 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

875 def supports_expert_map(self) -> bool:

876 return False

877

878 + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:

879 + # Let PrepareAndFinalize::finalize() decide the impl.

880 + return TopKWeightAndReduceDelegate()

881 +

882 def workspace_shapes(

883 self,

884 a: torch.Tensor,

vllm/model_executor/layers/fused_moe/fused_moe.py

@@ -25,6 +25,8 @@ from vllm.model_executor.layers.fused_moe.moe_align_block_size import (

25 moe_align_block_size)

26 from vllm.model_executor.layers.fused_moe.prepare_finalize import (

27 MoEPrepareAndFinalizeNoEP)

28 +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

29 + TopKWeightAndReduceDelegate)

30 from vllm.model_executor.layers.fused_moe.utils import (

31 _resize_cache, moe_kernel_quantize_input)

32 from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (

@@ -1596,6 +1598,10 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

1598 def supports_expert_map(self) -> bool:

1599 return True

1600

1601 + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:

1602 + # Let PrepareAndFinalize::finalize() decide the impl.

1603 + return TopKWeightAndReduceDelegate()

1604 +

1605 def workspace_shapes(

1606 self,

1607 a: torch.Tensor,

vllm/model_executor/layers/fused_moe/modular_kernel.py

+38 -6

@@ -23,7 +23,7 @@ from vllm.utils import cdiv

23 #

24 # [Router] → [Quantize-Dispatch] → [Permute-Experts-Unpermute] → [Combine]

25 #

26 -# Each component will be independent of the others except for

26 +# Each component will be independent of (but may inform) the others except for

27 # [Quantize-Dispatch] and `[Combine] (see below). The components can then be

28 # mixed and matched with so that DP+EP can be supported easily for multiple

29 # MoE kernel implementations.

@@ -32,13 +32,19 @@ from vllm.utils import cdiv

32 # * FusedMoEPrepareAndFinalize - an abstract base class for preparation of MoE

33 # inputs (e.g. quantization, distribution) and finalization of Moe outputs.

34 # The prepare method must take care of any needed quantization and the

35 -# finalize method must apply weights and do the final reduction of the output.

35 +# finalize method, informed by the FusedMoEPermuteExpertsUnpermute method,

36 +# may apply weights and/or do the final reduction of the output.

37 # * FusedMoEPermuteExpertsUnpermute - an abstract base class for the main fused

37 -# MoE operation. One important feature to note is that this class does not

38 -# apply topk weights or reduce the final output.

38 +# MoE operation, i.e matmul + act_mul + optionally quant + matmul.

39 +# Some FusedMoEPermuteExpertsUnpermute implementations may choose to do

40 +# the weight application and/or reduction. The class communicates this

41 +# to [Finalize] via a TopKWeightAndReduce object.

42 # * FusedMoEModularKernel - an interface class that combines a

43 # FusedMoEPrepareAndFinalize and a FusedMoEPermuteExpertsUnpermute to

44 # provide the standard fused MoE kernel interface.

45 +# * TopKWeightAndReduce - A TopKWeightAndReduce implementation chosen

46 +# by the FusedMoEPermuteExpertsUnpermute implementation that is passed

47 +# on to [Finalize].

48 #

49 # [Quantize-Prepare] and [Finalize] functionality are bundled into a single

50 # class `FusedMoEPrepareAndFinalize` since they could use collective

@@ -117,6 +123,24 @@ class ExpertTokensMetadata:

123 expert_num_tokens_cpu=expert_num_tokens_cpu)

124

125

126 +class TopKWeightAndReduce(ABC):

127 + """

128 + An abstract base class for weight application and reduction implementations.

129 + """

130 +

131 + @abstractmethod

132 + def apply(self, output: Optional[torch.Tensor],

133 + fused_expert_output: torch.Tensor, topk_weights: torch.Tensor,

134 + topk_ids: torch.Tensor,

135 + apply_router_weight_on_input: bool) -> torch.Tensor:

136 + """

137 + Apply topk_weights to the fused_experts_outputs and/or reduce.

138 + If an output tensor is not passed, it will be created in the

139 + function.

140 + """

141 + raise NotImplementedError

142 +

143 +

144 # TODO: pass FusedMoEParallelConfig in as ctor parameter?

145 class FusedMoEPrepareAndFinalize(ABC):

146 """

@@ -173,6 +197,7 @@ class FusedMoEPrepareAndFinalize(ABC):

197 topk_weights: torch.Tensor,

198 topk_ids: torch.Tensor,

199 apply_router_weight_on_input: bool,

200 + weight_and_reduce_impl: TopKWeightAndReduce,

201 ) -> None:

202 """

203 Perform any combine plus apply weights and perform a reduction on the

@@ -184,6 +209,8 @@ class FusedMoEPrepareAndFinalize(ABC):

209 - topk_ids: The topk_ids.

210 - apply_router_weight_on_input: When False, apply the weights to

211 fused_expert_output.

212 + - weight_and_reduce_impl: An optional TopKWeightAndReduce

213 + implementation.

214 """

215 raise NotImplementedError

216

@@ -323,6 +350,9 @@ class FusedMoEPermuteExpertsUnpermute(ABC):

350 return envs.VLLM_ENABLE_FUSED_MOE_ACTIVATION_CHUNKING and \

351 self.supports_chunking()

352

353 + def finalize_weight_and_reduce_impl(self) -> TopKWeightAndReduce:

354 + raise NotImplementedError

355 +

356 @abstractmethod

357 def apply(

358 self,

@@ -702,7 +732,9 @@ class FusedMoEModularKernel(torch.nn.Module):

732 a2_scale=a2_scale,

733 expert_tokens_meta=expert_tokens_meta)

734

705 - self.prepare_finalize.finalize(output, fused_out, topk_weights,

706 - topk_ids, apply_router_weight_on_input)

735 + self.prepare_finalize.finalize(

736 + output, fused_out, topk_weights, topk_ids,

737 + apply_router_weight_on_input,

738 + self.fused_experts.finalize_weight_and_reduce_impl())

739

740 return output

vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py

@@ -8,6 +8,8 @@ import torch

8 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

9 from vllm.logger import init_logger

10 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig

11 +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

12 + TopKWeightAndReduceDelegate)

13 from vllm.model_executor.layers.fused_moe.utils import (

14 _validate_scale_shape, moe_kernel_quantize_input)

15 from vllm.utils import cdiv, round_up

@@ -222,7 +224,12 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

224 topk_weights: torch.Tensor,

225 topk_ids: torch.Tensor,

226 apply_router_weight_on_input: bool,

227 + weight_and_reduce_impl: mk.TopKWeightAndReduce,

228 ) -> None:

229 + assert isinstance(

230 + weight_and_reduce_impl, TopKWeightAndReduceDelegate

231 + ), ("Weight application and reduction happens in the combine kernel.")

232 +

233 # This argument is optional

234 # There's not much point setting this unless it is != topk_ids.size(0)

235 bound_m: Optional[torch.Tensor] = None

vllm/model_executor/layers/fused_moe/prepare_finalize.py

+11 -4

@@ -6,8 +6,8 @@ import torch

7 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

8 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig

9 -from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (

10 - _moe_unpermute_and_reduce)

9 +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

10 + TopKWeightAndReduceContiguous, TopKWeightAndReduceDelegate)

11 from vllm.model_executor.layers.fused_moe.utils import (

12 moe_kernel_quantize_input)

@@ -62,6 +62,13 @@ class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):

62 topk_weights: torch.Tensor,

63 topk_ids: torch.Tensor,

64 apply_router_weight_on_input: bool,

65 + weight_and_reduce_impl: mk.TopKWeightAndReduce,

66 ) -> None:

66 - _moe_unpermute_and_reduce(output, fused_expert_output, None,

67 - topk_weights, apply_router_weight_on_input)

67 + if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate):

68 + weight_and_reduce_impl = TopKWeightAndReduceContiguous()

69 + weight_and_reduce_impl.apply(

70 + output=output,

71 + fused_expert_output=fused_expert_output,

72 + topk_weights=topk_weights,

73 + topk_ids=topk_ids,

74 + apply_router_weight_on_input=apply_router_weight_on_input)

vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py new

+139

@@ -0,0 +1,139 @@

1 +# SPDX-License-Identifier: Apache-2.0

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

3 +

4 +from typing import Optional

5 +

6 +import torch

7 +

8 +import vllm._custom_ops as ops

9 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk

10 +

11 +

12 +class TopKWeightAndReduceDelegate(mk.TopKWeightAndReduce):

13 + """

14 + Useful in the case when some FusedMoEPermuteExpertsUnpermute

15 + implementation does not perform weight application and reduction

16 + but cannot address the needs of all the compatible PrepareAndFinalize

17 + implementations.

18 + For example, BatchedTritonExperts is compatible with both

19 + PplxPrepareAndFinalize and BatchedPrepareAndFinalize. PplxPrepareAndFinalize

20 + does the weight-application + reduction as part of the pplx combine kernel.

21 + But the BatchedPrepareAndFinalize needs an implementation. To facilitate

22 + this case, the BatchedTritonExperts could use TopKWeightAndReduceDelegate

23 + so the PrepareAndFinalize implementations could choose how to

24 + weight + reduce.

25 + """

26 +

27 + def __eq__(self, other):

28 + return isinstance(other, TopKWeightAndReduceDelegate)

29 +

30 + def apply(self, output: Optional[torch.Tensor],

31 + fused_expert_output: torch.Tensor, topk_weights: torch.Tensor,

32 + topk_ids: torch.Tensor,

33 + apply_router_weight_on_input: bool) -> torch.Tensor:

34 + raise RuntimeError("The caller is expected to choose an appropriate "

35 + "TopKWeightAndReduce implementation.")

36 +

37 +

38 +class TopKWeightAndReduceNoOP(mk.TopKWeightAndReduce):

39 + """

40 + The fused_experts outputs have already been weight applied and reduced.

41 + This implementation is a no-op.

42 + """

43 +

44 + def __eq__(self, other):

45 + return isinstance(other, TopKWeightAndReduceNoOP)

46 +

47 + def apply(self, output: Optional[torch.Tensor],

48 + fused_expert_output: torch.Tensor, topk_weights: torch.Tensor,

49 + topk_ids: torch.Tensor,

50 + apply_router_weight_on_input: bool) -> torch.Tensor:

51 + # Relax this if an explicit copy is necessary. Note that,

52 + # if a copy is employed we have to make sure that the

53 + # tensors don't overlap

54 + assert output is None

55 + return fused_expert_output

56 +

57 +

58 +class TopKWeightAndReduceContiguous(mk.TopKWeightAndReduce):

59 + """

60 + TopKWeightAndReduce implementation for a fused_experts output

61 + of shape (m, topk, K)

62 + """

63 +

64 + def __eq__(self, other):

65 + return isinstance(other, TopKWeightAndReduceContiguous)

66 +

67 + def apply(self, output: Optional[torch.Tensor],

68 + fused_expert_output: torch.Tensor, topk_weights: torch.Tensor,

69 + topk_ids: torch.Tensor,

70 + apply_router_weight_on_input: bool) -> torch.Tensor:

71 +

72 + m, num_topk = topk_ids.size()

73 + k = fused_expert_output.size(-1)

74 + if fused_expert_output.ndim == 2:

75 + fused_expert_output = fused_expert_output.view(m, num_topk, k)

76 +

77 + assert fused_expert_output.size() == (m, num_topk, k), (

78 + f"Expected fused_expert_output size {(m, num_topk, k)}. But got "

79 + f"{fused_expert_output.size()}")

80 +

81 + if not apply_router_weight_on_input:

82 + fused_expert_output.mul_(topk_weights.view(m, -1, 1))

83 +

84 + if output is None:

85 + output = torch.empty((m, k),

86 + device=fused_expert_output.device,

87 + dtype=fused_expert_output.dtype)

88 + assert output.size() == (m, k), (

89 + f"Expected output size {(m, k)}. But got {output.size()}")

90 +

91 + ops.moe_sum(fused_expert_output, output)

92 + return output

93 +

94 +

95 +class TopKWeightAndReduceNaiveBatched(mk.TopKWeightAndReduce):

96 + """

97 + TopKWeightAndReduce implementation for a fused_experts output

98 + of shape (num_experts, batch_size, K)

99 + """

100 +

101 + def __init__(self, rank: int):

102 + self.rank = rank

103 +

104 + def __eq__(self, other):

105 + return (isinstance(other, TopKWeightAndReduceNaiveBatched)

106 + and (other.rank == self.rank))

107 +

108 + def apply(self, output: Optional[torch.Tensor],

109 + fused_expert_output: torch.Tensor, topk_weights: torch.Tensor,

110 + topk_ids: torch.Tensor,

111 + apply_router_weight_on_input: bool) -> torch.Tensor:

112 + assert fused_expert_output.ndim == 3

113 + num_tokens = topk_ids.size(0)

114 + num_local_experts = fused_expert_output.size(0)

115 + K = fused_expert_output.size(-1)

116 +

117 + if output is None:

118 + output = torch.zeros((num_tokens, K),

119 + device=fused_expert_output.device,

120 + dtype=fused_expert_output.dtype)

121 + else:

122 + output.fill_(0)

123 +

124 + assert output.size() == (num_tokens, K), (

125 + f"Expected output size {(num_tokens, K)}, but got {output.size()}")

126 +

127 + first_expert = num_local_experts * self.rank

128 + last_expert = first_expert + num_local_experts

129 +

130 + for expert_id in range(first_expert, last_expert):

131 + matching_tokens = topk_ids == expert_id

132 + topks = torch.any(matching_tokens, dim=1).flatten()

133 + rows = torch.count_nonzero(topks)

134 + rhs = fused_expert_output[expert_id - first_expert, :rows, :]

135 + if not apply_router_weight_on_input:

136 + rhs.mul_(topk_weights[matching_tokens].view(rhs.size(0), 1))

137 + output[topks] = output[topks] + rhs

138 +

139 + return output

vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py

+19

@@ -69,6 +69,25 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

69 return ((dge is None or dge.supports_expert_map())

70 and (te is None or te.supports_expert_map()))

72 + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:

73 + dge = self.deep_gemm_expert

74 + te = self.triton_expert

75 + dge_war = dge.finalize_weight_and_reduce_impl() if dge else None

76 + te_war = te.finalize_weight_and_reduce_impl() if te else None

77 + is_dge_war = dge_war is not None

78 + is_te_war = te_war is not None

79 +

80 + if is_dge_war and is_te_war:

81 + assert dge_war == te_war, (

82 + "Both implementations should agree on WeightAndReduce impls. "

83 + f"Got dge_war: {dge_war}, and te_war: {te_war}")

84 +

85 + if dge_war is not None:

86 + return dge_war

87 +

88 + assert te_war is not None

89 + return te_war

90 +

91 def workspace_shapes(

92 self,

93 a: torch.Tensor,