16f3250527 · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -1,6 +1,7 @@

1 # Fused MoE Modular Kernel

2

3 ## Introduction

4 +

5 FusedMoEModularKernel is implemented [here](gh-file:/vllm/model_executor/layers/fused_moe/modular_kernel.py)

6

7 Based on the format of the input activations, FusedMoE implementations are broadly classified into 2 types.

@@ -31,7 +32,8 @@ As can be seen from the diagrams, there are a lot of operations and there can be

32

33 The rest of the document will focus on the Contiguous / Non-Batched case. Extrapolating to the Batched case should be straight-forward.

34

34 -## ModularKernel Components:

35 +## ModularKernel Components

36 +

37 FusedMoEModularKernel splits the FusedMoE operation into 3 parts,

38

39 1. TopKWeightAndReduce

@@ -39,6 +41,7 @@ FusedMoEModularKernel splits the FusedMoE operation into 3 parts,

41 3. FusedMoEPermuteExpertsUnpermute

42

43 ### TopKWeightAndReduce

44 +

45 The TopK Weight Application and Reduction components happen right after the Unpermute operation and before the All2All Combine. Note that the `FusedMoEPermuteExpertsUnpermute` is responsible for the Unpermute and `FusedMoEPrepareAndFinalize` is responsible for the All2All Combine. There is value in doing the TopK Weight Application and Reduction in the `FusedMoEPermuteExpertsUnpermute`. But some implementations choose to do it `FusedMoEPrepareAndFinalize`. In order to enable this flexibility, we have a TopKWeightAndReduce abstract class.

46

47 Please find the implementations of TopKWeightAndReduce [here](gh-file:vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py).

@@ -50,12 +53,14 @@ The `FusedMoEModularKernel` acts as a bridge between the `FusedMoEPermuteExperts

53 * `FusedMoEPermuteExpertsUnpermute::finalize_weight_and_reduce_impl` method returns `TopKWeightAndReduceContiguous` / `TopKWeightAndReduceNaiveBatched` / `TopKWeightAndReduceDelegate` if the `FusedMoEPermuteExpertsUnpermute` implementation needs the `FusedMoEPrepareAndFinalize::finalize()` to do the weight application and reduction.

54

55 ### FusedMoEPrepareAndFinalize

56 +

57 The `FusedMoEPrepareAndFinalize` abstract class exposes `prepare` and `finalize` functions.

58 The `prepare` function is responsible for input activation Quantization and All2All Dispatch. The `finalize` function is responsible for invoking the All2All Combine. Additionally the `finalize` function may or may not do the TopK weight application and reduction (Please refer to the TopKWeightAndReduce section)

59

60 ![](../assets/design/fused_moe_modular_kernel/prepare_and_finalize_blocks.png "FusedMoEPrepareAndFinalize Blocks")

61

62 ### FusedMoEPermuteExpertsUnpermute

63 +

64 The `FusedMoEPermuteExpertsUnpermute` class is where the crux of the MoE operations happen. The `FusedMoEPermuteExpertsUnpermute` abstract class exposes a few important functions,

65

66 * apply()

@@ -63,6 +68,7 @@ The `FusedMoEPermuteExpertsUnpermute` class is where the crux of the MoE operati

68 * finalize_weight_and_reduce_impl()

69

70 #### apply()

71 +

72 The `apply` method is where the implementations perform

73

74 * Permute

@@ -74,50 +80,56 @@ The `apply` method is where the implementations perform

80 * Maybe TopK Weight Application + Reduction

81

82 #### workspace_shapes()

83 +

84 The core FusedMoE implementation performs a series of operations. It would be inefficient to create output memory for each of these operations separately. To that effect, implementations are required to declare 2 workspace shapes, the workspace datatype and the FusedMoE output shape as outputs of the workspace_shapes() method. This information is used to allocate the workspace tensors and the output tensor in `FusedMoEModularKernel::forward()` and passed on to the `FusedMoEPermuteExpertsUnpermute::apply()` method. The workspaces could then be used as intermediate buffers in the FusedMoE implementation.

85

86 #### finalize_weight_and_reduce_impl()

87 +

88 It is sometimes efficient to perform TopK weight application and Reduction inside the `FusedMoEPermuteExpertsUnpermute::apply()`. Find an example [here](https://github.com/vllm-project/vllm/pull/20228). We have a `TopKWeightAndReduce` abstract class to facilitate such implementations. Please refer to the TopKWeightAndReduce section.

89 `FusedMoEPermuteExpertsUnpermute::finalize_weight_and_reduce_impl()` returns the `TopKWeightAndReduce` object that the implementation wants the `FusedMoEPrepareAndFinalize::finalize()` to use.

90

91 ![](../assets/design/fused_moe_modular_kernel/fused_experts_blocks.png "FusedMoEPermuteExpertsUnpermute Blocks")

92

93 ### FusedMoEModularKernel

94 +

95 `FusedMoEModularKernel` is composed of the `FusedMoEPrepareAndFinalize` and `FusedMoEPermuteExpertsUnpermute` objects.

96 `FusedMoEModularKernel` pseudocode/sketch,

97

89 -```

90 -FusedMoEModularKernel::__init__(self,

91 - prepare_finalize: FusedMoEPrepareAndFinalize,

92 - fused_experts: FusedMoEPermuteExpertsUnpermute):

98 +```py

99 +class FusedMoEModularKernel:

100 + def __init__(self,

101 + prepare_finalize: FusedMoEPrepareAndFinalize,

102 + fused_experts: FusedMoEPermuteExpertsUnpermute):

103 +

104 + self.prepare_finalize = prepare_finalize

105 + self.fused_experts = fused_experts

106

94 - self.prepare_finalize = prepare_finalize

95 - self.fused_experts = fused_experts

107 + def forward(self, DP_A):

108

97 -FusedMoEModularKernel::forward(self, DP_A):

109 + Aq, A_scale, _, _, _ = self.prepare_finalize.prepare(DP_A, ...)

110

99 - Aq, A_scale, _, _, _ = self.prepare_finalize.prepare(DP_A, ...)

111 + workspace13_shape, workspace2_shape, _, _ = self.fused_experts.workspace_shapes(...)

112

101 - workspace13_shape, workspace2_shape, _, _ = self.fused_experts.workspace_shapes(...)

113 + # allocate workspaces

114 + workspace_13 = torch.empty(workspace13_shape, ...)

115 + workspace_2 = torch.empty(workspace2_shape, ...)

116

103 - # allocate workspaces

104 - workspace_13 = torch.empty(workspace13_shape, ...)

105 - workspace_2 = torch.empty(workspace2_shape, ...)

117 + # execute fused_experts

118 + fe_out = self.fused_experts.apply(Aq, A_scale, workspace13, workspace2, ...)

119

107 - # execute fused_experts

108 - fe_out = self.fused_experts.apply(Aq, A_scale, workspace13, workspace2, ...)

120 + # war_impl is an object of type TopKWeightAndReduceNoOp if the fused_experts implementations

121 + # performs the TopK Weight Application and Reduction.

122 + war_impl = self.fused_experts.finalize_weight_and_reduce_impl()

123

110 - # war_impl is an object of type TopKWeightAndReduceNoOp if the fused_experts implementations performs the TopK Weight Application and Reduction.

111 - war_impl = self.fused_experts.finalize_weight_and_reduce_impl()

124 + output = self.prepare_finalize.finalize(fe_out, war_impl,...)

125

113 - output = self.prepare_finalize.finalize(fe_out, war_impl,...)

114 -

115 - return output

126 + return output

127 ```

128

129 ## How-To

130

131 ### How To Add a FusedMoEPrepareAndFinalize Type

132 +

133 Typically a FusedMoEPrepareAndFinalize type is backed by an All2All Dispatch & Combine implementation / kernel. For example,

134

135 * PplxPrepareAndFinalize type is backed by Pplx All2All kernels,

@@ -125,9 +137,11 @@ Typically a FusedMoEPrepareAndFinalize type is backed by an All2All Dispatch & C

137 * DeepEPLLPrepareAndFinalize type is backed by DeepEP Low-Latency All2All kernels.

138

139 #### Step 1: Add an All2All manager

140 +

141 The purpose of the All2All Manager is to setup the All2All kernel implementations. The `FusedMoEPrepareAndFinalize` implementations typically fetch a kernel-implementation "handle" from the All2All Manager to invoke the Dispatch and Combine functions. Please look at the All2All Manager implementations [here](gh-file:vllm/distributed/device_communicators/all2all.py).

142

143 #### Step 2: Add a FusedMoEPrepareAndFinalize Type

144 +

145 This section describes the significance of the various functions exposed by the `FusedMoEPrepareAndFinalize` abstract class.

146

147 `FusedMoEPrepareAndFinalize::prepare()`: The prepare method implements the Quantization and All2All Dispatch. Typically the Dispatch function from the relevant All2All Manager is invoked.

@@ -145,6 +159,7 @@ This section describes the significance of the various functions exposed by the

159 We suggest picking an already existing `FusedMoEPrepareAndFinalize` implementation that matches your All2All implementation closely and using it as a reference.

160

161 ### How To Add a FusedMoEPermuteExpertsUnpermute Type

162 +

163 FusedMoEPermuteExpertsUnpermute performs the core of the FusedMoE operations. The various functions exposed by the abstract class and their significance is as follows,

164

165 `FusedMoEPermuteExpertsUnpermute::activation_formats()`: Return the supported Input and Output activation formats. i.e. Contiguous / Batched format.

@@ -159,12 +174,14 @@ implementations that input `FusedMoEActivationFormat.Standard` support chunking

174 `FusedMoEPermuteExpertsUnpermute::apply`: Refer to `FusedMoEPermuteExpertsUnpermute` section above.

175

176 ### FusedMoEModularKernel Initialization

177 +

178 `FusedMoEMethodBase` class has 2 methods that are collectively responsible in creating the `FusedMoEModularKernel` object. They are,

179

180 * select_gemm_impl, and

181 * init_prepare_finalize

182

183 #### select_gemm_impl

184 +

185 The `select_gemm_impl` method is undefined in the base class. It is the responsibility of the derived class to implement a method that constructs a valid/appropriate `FusedMoEPermuteExpertsUnpermute` object.

186 Please refer to the implementations in,

187

@@ -176,12 +193,14 @@ Please refer to the implementations in,

193 dervied classes.

194

195 #### init_prepare_finalize

196 +

197 Based on the input and env settings, the `init_prepare_finalize` method creates the appropriate `FusedMoEPrepareAndFinalize` object. The method then queries `select_gemm_impl` for the appropriate `FusedMoEPermuteExpertsUnpermute` object and builds the `FusedMoEModularKernel` object

198

199 Please take a look at [init_prepare_finalize](https://github.com/vllm-project/vllm/blob/1cbf951ba272c230823b947631065b826409fa62/vllm/model_executor/layers/fused_moe/layer.py#L188).

200 **Important**: The `FusedMoEMethodBase` derived classes use the `FusedMoEMethodBase::fused_experts` object in their `apply` methods. When settings permit the construction of a valid `FusedMoEModularKernel` object, we override `FusedMoEMethodBase::fused_experts` with it. This essentially makes the derived classes agnostic to what FusedMoE implementation is used.

201

202 ### How To Unit Test

203 +

204 We have `FusedMoEModularKernel` unit tests at [test_modular_kernel_combinations.py](gh-file:tests/kernels/moe/test_modular_kernel_combinations.py).

205

206 The unit test iterates through all combinations of `FusedMoEPrepareAndFinalize` and `FusedMoEPremuteExpertsUnpermute` types and if they are

@@ -196,18 +215,21 @@ If you are adding some `FusedMoEPrepareAndFinalize` / `FusedMoEPermuteExpertsUnp

215 Doing this will add the new implementation to the test suite.

216

217 ### How To Check `FusedMoEPrepareAndFinalize` & `FusedMoEPermuteExpertsUnpermute` Compatibility

218 +

219 The unit test file [test_modular_kernel_combinations.py](gh-file:tests/kernels/moe/test_modular_kernel_combinations.py) can also be executed as a standalone script.

220 Example: `python3 -m tests.kernels.moe.test_modular_kernel_combinations --pf-type PplxPrepareAndFinalize --experts-type BatchedTritonExperts`

221 As a side-effect, this script can be used to test `FusedMoEPrepareAndFinalize` & `FusedMoEPermuteExpertsUnpermute` compatibility. When invoked

222 with incompatible types, the script will error.

223

224 ### How To Profile

225 +

226 Please take a look at [profile_modular_kernel.py](gh-file:tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py)

227 The script can be used to generate Torch traces for a single `FusedMoEModularKernel::forward()` call for any compatible

228 `FusedMoEPrepareAndFinalize` and `FusedMoEPermuteExpertsUnpermute` types.

229 Example: `python3 -m tests.kernels.moe.modular_kernel_tools.profile_modular_kernel --pf-type PplxPrepareAndFinalize --experts-type BatchedTritonExperts`

230

231 ## FusedMoEPrepareAndFinalize Implementations

232 +

233 The following table lists the `FusedMoEPrepareAndFinalize` implementations at the time of writing,

234

235 | Implementation | Type | Comments |

@@ -220,6 +242,7 @@ The following table lists the `FusedMoEPrepareAndFinalize` implementations at th

242 | BatchedPrepareAndFinalize | Batched | A reference prepare/finalize class that reorganizes the tokens into expert batched format, i.e. E x max_num_tokens x K. (Doesn’t use any all2all kernels. This is primarily used in unit testing) |

243

244 ## FusedMoEPermuteExpertsUnpermute

245 +

246 The following table lists the `FusedMoEPermuteExpertsUnpermute` implementations at the time of writing,

247

248 | Implementation | Type | Comment |

[CI/Build] Fix pre-commit failure in docs (#21897)