@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 11dfdf21bf
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, 01:00 UTC

[Kernel] DeepGemm MoE : Integrate triton permute / unpermute kernels (#20903)

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

Varun Sundar Rabindranath committed Jul 17, 2025 at 13:40 UTC 11dfdf21bf3ab057986a7f0053704dfcfee27a4e
10 files changed +491 -59
tests/kernels/moe/modular_kernel_tools/cli_args.py
-1
@@ -85,7 +85,6 @@ def make_config_arg_parser(description: str):
85 help="num topk")
86 parser.add_argument(
87 "--fused-moe-chunk-size",
88 - nargs="+",
88 type=int,
89 help="Fused moe chunk size used for the non-batched fused experts impl."
90 )
vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py
+1
@@ -239,6 +239,7 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
239 topk: int,
240 global_num_experts: int,
241 local_num_experts: int,
242 + expert_tokens_metadata: Optional[mk.ExpertTokensMetadata],
243 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
244 assert a.dim() == 2
245 # FIXME (varun): We should be able to dispatch only from the leader
vllm/model_executor/layers/fused_moe/batched_triton_or_deep_gemm_moe.py
+5 -2
@@ -116,6 +116,7 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
116 topk: int,
117 global_num_experts: int,
118 local_num_experts: int,
119 + expert_tokens_metadata: Optional[mk.ExpertTokensMetadata],
120 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
121 # Note: the deep gemm workspaces are strictly larger than the triton
122 # workspaces so we can be pessimistic here and allocate for DeepGemm
@@ -123,11 +124,13 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
124 if self.allow_deep_gemm:
125 assert self.batched_deep_gemm_experts is not None
126 return self.batched_deep_gemm_experts.workspace_shapes(
126 - a, aq, M, N, K, topk, global_num_experts, local_num_experts)
127 + a, aq, M, N, K, topk, global_num_experts, local_num_experts,
128 + expert_tokens_metadata)
129 else:
130 assert self.batched_triton_experts is not None
131 return self.batched_triton_experts.workspace_shapes(
130 - a, aq, M, N, K, topk, global_num_experts, local_num_experts)
132 + a, aq, M, N, K, topk, global_num_experts, local_num_experts,
133 + expert_tokens_metadata)
134
135 def apply(self, output: torch.Tensor, hidden_states: torch.Tensor,
136 w1: torch.Tensor, w2: torch.Tensor, topk_weights: torch.Tensor,
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+1
@@ -271,6 +271,7 @@ class CutlassExpertsFp8(mk.FusedMoEPermuteExpertsUnpermute):
271 topk: int,
272 global_num_experts: int,
273 local_num_experts: int,
274 + expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
275 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
276 workspace1: tuple[int, ...] = ()
277 workspace2: tuple[int, ...] = ()
vllm/model_executor/layers/fused_moe/deep_gemm_moe.py
+53 -48
@@ -8,16 +8,16 @@ 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.moe_permute_unpermute import (
12 - _moe_permute)
11 +from vllm.model_executor.layers.fused_moe.deep_gemm_utils import (
12 + compute_aligned_M, deepgemm_moe_permute, deepgemm_unpermute_and_reduce)
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 - TopKWeightAndReduceContiguous, TopKWeightAndReduceNoOP)
16 + TopKWeightAndReduceNoOP)
17 from vllm.model_executor.layers.fused_moe.utils import _resize_cache
18 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
19 per_token_group_quant_fp8)
20 -from vllm.utils import has_deep_gemm, round_up
20 +from vllm.utils import has_deep_gemm
21 from vllm.utils.deep_gemm import m_grouped_fp8_gemm_nt_contiguous
22
23 logger = init_logger(__name__)
@@ -93,18 +93,25 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
93 return TopKWeightAndReduceNoOP()
94
95 def workspace_shapes(
96 - self, a: torch.Tensor, aq: torch.Tensor, M: int, N: int, K: int,
97 - topk: int, global_num_experts: int, local_num_experts: int
96 + self,
97 + a: torch.Tensor,
98 + aq: torch.Tensor,
99 + M: int,
100 + N: int,
101 + K: int,
102 + topk: int,
103 + global_num_experts: int,
104 + local_num_experts: int,
105 + expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
106 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
107 assert self.block_shape is not None
100 - # We use global_num_experts due to how moe_align_block_size handles
101 - # expert_maps.
102 - num_experts = global_num_experts
108 block_m = self.block_shape[0]
104 - M_sum = (M * topk) + num_experts * (block_m - 1)
105 - M_sum = round_up(M_sum, block_m)
106 - workspace1 = (M_sum, max(N // 2, K))
107 - workspace2 = (M_sum, max(N, K))
109 + M_sum = compute_aligned_M(M, topk, local_num_experts, block_m,
110 + expert_tokens_meta)
111 + assert M_sum % block_m == 0
112 +
113 + workspace1 = (M_sum, max(N, K))
114 + workspace2 = (M_sum, max(N // 2, K))
115 output = (M, K)
116 return (workspace1, workspace2, output, a.dtype)
117
@@ -131,43 +138,40 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
138 apply_router_weight_on_input: bool,
139 ):
140 assert self.block_shape is not None
141 + assert a1q_scale is not None
142
143 a1q = hidden_states
144 _, N, K = w1.size()
137 - M, _ = output.size()
138 - num_topk = topk_ids.size(1)
145
146 + local_num_experts = w1.size(0)
147 if global_num_experts == -1:
141 - global_num_experts = w1.size(0)
148 + global_num_experts = local_num_experts
149
150 assert w2.size(1) == K
151
145 - a1q, a1q_scale, _, expert_ids, inv_perm = _moe_permute(
146 - a1q,
147 - a1q_scale,
148 - topk_ids,
149 - global_num_experts,
150 - expert_map,
151 - self.block_shape[0],
152 - )
153 -
154 - if expert_map is not None:
155 - # DeepGemm (Grouped Contiguous) kernel needs a valid B index
156 - # for all rows of A. To that effect, simply compute with
157 - # the 0th weight matrix.
158 - # Note that this relies on the fact that corresponding topk
159 - # weights would be 0 during weight multiplication.
160 - expert_ids = torch.where(expert_ids == -1, 0, expert_ids)
161 -
162 - # Note: M_sum is different than the pre-permuted shape of a1q.
163 - M_sum = a1q.size(0)
164 -
165 - mm1_out = _resize_cache(workspace2, (M_sum, N))
166 - act_out = _resize_cache(workspace13, (M_sum, N // 2))
167 - quant_out = _resize_cache(workspace2.view(dtype=torch.float8_e4m3fn),
152 + M_sum = compute_aligned_M(M=topk_ids.size(0),
153 + num_topk=topk_ids.size(1),
154 + local_num_experts=local_num_experts,
155 + alignment=deep_gemm_block_shape()[0],
156 + expert_tokens_meta=expert_tokens_meta)
157 +
158 + a1q_perm = _resize_cache(workspace2.view(dtype=torch.float8_e4m3fn),
159 + (M_sum, K))
160 + mm1_out = _resize_cache(workspace13, (M_sum, N))
161 + act_out = _resize_cache(workspace2, (M_sum, N // 2))
162 + quant_out = _resize_cache(workspace13.view(dtype=torch.float8_e4m3fn),
163 (M_sum, N // 2))
169 - mm2_out = _resize_cache(workspace13, (M_sum, K))
170 - perm_out = _resize_cache(workspace2, (M * num_topk, K))
164 + mm2_out = _resize_cache(workspace2, (M_sum, K))
165 +
166 + a1q, a1q_scale, expert_ids, inv_perm = deepgemm_moe_permute(
167 + aq=a1q,
168 + aq_scale=a1q_scale,
169 + topk_ids=topk_ids,
170 + local_num_experts=local_num_experts,
171 + expert_map=expert_map,
172 + expert_tokens_meta=expert_tokens_meta,
173 + aq_out=a1q_perm)
174 + assert a1q.size(0) == M_sum
175
176 m_grouped_fp8_gemm_nt_contiguous((a1q, a1q_scale), (w1, w1_scale),
177 mm1_out, expert_ids)
@@ -183,14 +187,15 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
187 m_grouped_fp8_gemm_nt_contiguous((a2q, a2q_scale), (w2, w2_scale),
188 mm2_out, expert_ids)
189
186 - torch.index_select(mm2_out, 0, inv_perm, out=perm_out)
190 + if apply_router_weight_on_input:
191 + topk_weights = torch.ones_like(topk_weights)
192
188 - TopKWeightAndReduceContiguous().apply(
189 - output=output,
190 - fused_expert_output=perm_out,
191 - topk_weights=topk_weights,
192 - topk_ids=topk_ids,
193 - apply_router_weight_on_input=apply_router_weight_on_input)
193 + deepgemm_unpermute_and_reduce(a=mm2_out,
194 + topk_ids=topk_ids,
195 + topk_weights=topk_weights,
196 + inv_perm=inv_perm,
197 + expert_map=expert_map,
198 + output=output)
199
200
201 def deep_gemm_moe_fp8(
vllm/model_executor/layers/fused_moe/deep_gemm_utils.py new
+413
@@ -0,0 +1,413 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Taken from https://github.com/ModelTC/LightLLM/blob/8ed97c74c18f11505b048b1ba00ba5c0cef8bff6/lightllm/common/fused_moe/deepep_scatter_gather.py
5 +and updated to fit vllm needs and terminology.
6 +"""
7 +
8 +import functools
9 +from typing import Optional
10 +
11 +import torch
12 +
13 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
14 +from vllm.model_executor.layers.fused_moe.utils import count_expert_num_tokens
15 +from vllm.triton_utils import tl, triton
16 +from vllm.utils import round_up
17 +
18 +
19 +@functools.cache
20 +def deep_gemm_block_shape() -> list[int]:
21 + # Lazy import to avoid CUDA initialization problems.
22 + import deep_gemm as dg
23 + block = dg.get_m_alignment_for_contiguous_layout()
24 + return [block, block]
25 +
26 +
27 +def expert_num_tokens_round_up_and_sum(expert_num_tokens: torch.Tensor,
28 + alignment: int) -> int:
29 + # Round up each element in expert_num_tokens to the nearest multiple of
30 + # alignment.
31 + ent = (expert_num_tokens.to(torch.int64) +
32 + (alignment - 1)) // alignment * alignment
33 + return torch.sum(ent).item()
34 +
35 +
36 +def compute_aligned_M(M: int, num_topk: int, local_num_experts: int,
37 + alignment: int,
38 + expert_tokens_meta: Optional[mk.ExpertTokensMetadata]):
39 +
40 + if ((expert_tokens_meta is not None)
41 + and (expert_tokens_meta.expert_num_tokens_cpu is not None)):
42 + return expert_num_tokens_round_up_and_sum(
43 + expert_tokens_meta.expert_num_tokens_cpu, alignment=alignment)
44 +
45 + # expert_num_tokens information is not available on the cpu.
46 + # compute the max required size.
47 + M_sum = (M * num_topk) + local_num_experts * (alignment - 1)
48 + M_sum = round_up(M_sum, alignment)
49 + return M_sum
50 +
51 +
52 +@triton.jit
53 +def apply_expert_map(expert_id, expert_map):
54 + if expert_id != -1:
55 + expert_id = tl.load(expert_map + expert_id).to(tl.int64)
56 + return expert_id
57 +
58 +
59 +@triton.jit
60 +def round_up_128(x: int) -> int:
61 + y = 128
62 + return ((x + y - 1) // y) * y
63 +
64 +
65 +@triton.jit
66 +def _fwd_kernel_ep_scatter_1(
67 + num_recv_tokens_per_expert,
68 + expert_start_loc,
69 + m_indices,
70 + num_experts: tl.constexpr,
71 + BLOCK_E: tl.constexpr,
72 + BLOCK_EXPERT_NUM: tl.constexpr,
73 +):
74 + cur_expert = tl.program_id(0)
75 +
76 + offset_cumsum = tl.arange(0, BLOCK_EXPERT_NUM)
77 + tokens_per_expert = tl.load(num_recv_tokens_per_expert + offset_cumsum,
78 + mask=offset_cumsum < num_experts,
79 + other=0)
80 + tokens_per_expert = round_up_128(tokens_per_expert)
81 + cumsum = tl.cumsum(tokens_per_expert) - tokens_per_expert
82 + tl.store(expert_start_loc + offset_cumsum,
83 + cumsum,
84 + mask=offset_cumsum < num_experts)
85 +
86 + cur_expert_start = tl.load(expert_start_loc + cur_expert)
87 + cur_expert_token_num = tl.load(num_recv_tokens_per_expert + cur_expert)
88 +
89 + m_indices_start_ptr = m_indices + cur_expert_start
90 + off_expert = tl.arange(0, BLOCK_E)
91 +
92 + for start_m in tl.range(0, cur_expert_token_num, BLOCK_E, num_stages=4):
93 + tl.store(
94 + m_indices_start_ptr + start_m + off_expert,
95 + cur_expert,
96 + )
97 +
98 +
99 +@triton.jit
100 +def _fwd_kernel_ep_scatter_2(
101 + total_token_num,
102 + expert_start_loc,
103 + recv_x,
104 + recv_x_stride0,
105 + recv_x_stride1,
106 + recv_x_scale,
107 + recv_x_scale_stride0,
108 + recv_x_scale_stride1,
109 + recv_topk,
110 + recv_topk_stride0,
111 + recv_topk_stride1,
112 + output_tensor,
113 + output_tensor_stride0,
114 + output_tensor_stride1,
115 + output_tensor_scale,
116 + output_tensor_scale_stride0,
117 + output_tensor_scale_stride1,
118 + output_index,
119 + output_index_stride0,
120 + output_index_stride1,
121 + topk_num: tl.constexpr,
122 + expert_map,
123 + HAS_EXPERT_MAP: tl.constexpr,
124 + HIDDEN_SIZE: tl.constexpr,
125 + HIDDEN_SIZE_PAD: tl.constexpr,
126 + SCALE_HIDDEN_SIZE: tl.constexpr,
127 + SCALE_HIDDEN_SIZE_PAD: tl.constexpr,
128 +):
129 + start_token_id = tl.program_id(0)
130 + grid_num = tl.num_programs(0)
131 +
132 + offset_in = tl.arange(0, HIDDEN_SIZE_PAD)
133 + mask = offset_in < HIDDEN_SIZE
134 +
135 + offset_in_s = tl.arange(0, SCALE_HIDDEN_SIZE_PAD)
136 + mask_s = offset_in_s < SCALE_HIDDEN_SIZE
137 +
138 + for token_id in range(start_token_id, total_token_num, grid_num):
139 + to_copy = tl.load(recv_x + token_id * recv_x_stride0 + offset_in,
140 + mask=mask)
141 + to_copy_s = tl.load(recv_x_scale + token_id * recv_x_scale_stride0 +
142 + offset_in_s,
143 + mask=mask_s)
144 +
145 + for topk_index in tl.range(0, topk_num, 1, num_stages=4):
146 + expert_id = tl.load(recv_topk + token_id * recv_topk_stride0 +
147 + topk_index)
148 +
149 + if HAS_EXPERT_MAP:
150 + expert_id = apply_expert_map(expert_id, expert_map)
151 +
152 + if expert_id >= 0:
153 + dest_token_index = tl.atomic_add(expert_start_loc + expert_id,
154 + 1)
155 + tl.store(
156 + output_index + token_id * output_index_stride0 +
157 + topk_index, dest_token_index)
158 + output_tensor_ptr = (output_tensor +
159 + dest_token_index * output_tensor_stride0)
160 + output_tensor_scale_ptr = (
161 + output_tensor_scale +
162 + dest_token_index * output_tensor_scale_stride0)
163 + tl.store(output_tensor_ptr + offset_in, to_copy, mask=mask)
164 + tl.store(output_tensor_scale_ptr + offset_in_s,
165 + to_copy_s,
166 + mask=mask_s)
167 +
168 +
169 +@torch.no_grad()
170 +def ep_scatter(
171 + recv_x: torch.Tensor,
172 + recv_x_scale: torch.Tensor,
173 + recv_topk: torch.Tensor,
174 + num_recv_tokens_per_expert: torch.Tensor,
175 + expert_map: Optional[torch.Tensor],
176 + expert_start_loc: torch.Tensor,
177 + output_tensor: torch.Tensor,
178 + output_tensor_scale: torch.Tensor,
179 + m_indices: torch.Tensor,
180 + output_index: torch.Tensor,
181 +):
182 + BLOCK_E = 128 # token num of per expert is aligned to 128
183 + BLOCK_D = 128 # block size of quantization
184 + num_warps = 8
185 + num_experts = num_recv_tokens_per_expert.shape[0]
186 + hidden_size = recv_x.shape[1]
187 + # grid = (triton.cdiv(hidden_size, BLOCK_D), num_experts)
188 + grid = num_experts
189 +
190 + assert m_indices.shape[0] % BLOCK_E == 0
191 +
192 + _fwd_kernel_ep_scatter_1[(grid, )](
193 + num_recv_tokens_per_expert,
194 + expert_start_loc,
195 + m_indices,
196 + num_experts=num_experts,
197 + num_warps=num_warps,
198 + BLOCK_E=BLOCK_E,
199 + BLOCK_EXPERT_NUM=triton.next_power_of_2(num_experts),
200 + )
201 +
202 + grid = min(recv_topk.shape[0], 1024 * 8)
203 +
204 + _fwd_kernel_ep_scatter_2[(grid, )](
205 + recv_topk.shape[0],
206 + expert_start_loc,
207 + recv_x,
208 + recv_x.stride(0),
209 + recv_x.stride(1),
210 + recv_x_scale,
211 + recv_x_scale.stride(0),
212 + recv_x_scale.stride(1),
213 + recv_topk,
214 + recv_topk.stride(0),
215 + recv_topk.stride(1),
216 + output_tensor,
217 + output_tensor.stride(0),
218 + output_tensor.stride(1),
219 + output_tensor_scale,
220 + output_tensor_scale.stride(0),
221 + output_tensor_scale.stride(1),
222 + output_index,
223 + output_index.stride(0),
224 + output_index.stride(1),
225 + topk_num=recv_topk.shape[1],
226 + expert_map=expert_map,
227 + HAS_EXPERT_MAP=expert_map is not None,
228 + num_warps=num_warps,
229 + HIDDEN_SIZE=hidden_size,
230 + HIDDEN_SIZE_PAD=triton.next_power_of_2(hidden_size),
231 + SCALE_HIDDEN_SIZE=hidden_size // BLOCK_D,
232 + SCALE_HIDDEN_SIZE_PAD=triton.next_power_of_2(hidden_size // BLOCK_D),
233 + )
234 + return
235 +
236 +
237 +@triton.jit
238 +def _fwd_kernel_ep_gather(
239 + total_token_num,
240 + input_tensor,
241 + input_tensor_stride0,
242 + input_tensor_stride1,
243 + recv_topk_ids,
244 + recv_topk_ids_stride0,
245 + recv_topk_ids_stride1,
246 + recv_topk_weight,
247 + recv_topk_weight_stride0,
248 + recv_topk_weight_stride1,
249 + input_index,
250 + input_index_stride0,
251 + input_index_stride1,
252 + output_tensor,
253 + output_tensor_stride0,
254 + output_tensor_stride1,
255 + topk_num: tl.constexpr,
256 + expert_map,
257 + HAS_EXPERT_MAP: tl.constexpr,
258 + BLOCK_D: tl.constexpr,
259 +):
260 + cur_block = tl.program_id(0)
261 + start_cur_token = tl.program_id(1)
262 + grid_num = tl.num_programs(1)
263 +
264 + for cur_token in range(start_cur_token, total_token_num, grid_num):
265 + off_d = tl.arange(0, BLOCK_D)
266 + accumulator = tl.zeros([BLOCK_D], dtype=tl.float32)
267 + for topk_index in range(0, topk_num):
268 + expert_id = tl.load(recv_topk_ids +
269 + cur_token * recv_topk_ids_stride0 + topk_index)
270 +
271 + if HAS_EXPERT_MAP:
272 + expert_id = apply_expert_map(expert_id, expert_map)
273 +
274 + if expert_id >= 0:
275 + source_token_index = tl.load(input_index +
276 + cur_token * input_index_stride0 +
277 + topk_index)
278 + acc_weight = tl.load(recv_topk_weight +
279 + cur_token * recv_topk_weight_stride0 +
280 + topk_index)
281 + tmp = tl.load(input_tensor +
282 + source_token_index * input_tensor_stride0 +
283 + cur_block * BLOCK_D + off_d)
284 + accumulator += tmp.to(tl.float32) * acc_weight
285 +
286 + tl.store(
287 + output_tensor + cur_token * output_tensor_stride0 +
288 + cur_block * BLOCK_D + off_d,
289 + accumulator.to(output_tensor.dtype.element_ty),
290 + )
291 +
292 +
293 +@torch.no_grad()
294 +def ep_gather(
295 + input_tensor: torch.Tensor,
296 + recv_topk_ids: torch.Tensor,
297 + recv_topk_weight: torch.Tensor,
298 + input_index: torch.Tensor,
299 + expert_map: Optional[torch.Tensor],
300 + output_tensor: torch.Tensor,
301 +):
302 + num_warps = 2
303 + num_tokens = output_tensor.shape[0]
304 + hidden_size = input_tensor.shape[1]
305 + BLOCK_D = min(hidden_size, 1024)
306 + assert hidden_size % BLOCK_D == 0
307 + grid = (triton.cdiv(hidden_size, BLOCK_D), min(num_tokens, 1024))
308 +
309 + _fwd_kernel_ep_gather[grid](
310 + num_tokens,
311 + input_tensor,
312 + input_tensor.stride(0),
313 + input_tensor.stride(1),
314 + recv_topk_ids,
315 + recv_topk_ids.stride(0),
316 + recv_topk_ids.stride(1),
317 + recv_topk_weight,
318 + recv_topk_weight.stride(0),
319 + recv_topk_weight.stride(1),
320 + input_index,
321 + input_index.stride(0),
322 + input_index.stride(1),
323 + output_tensor,
324 + output_tensor.stride(0),
325 + output_tensor.stride(1),
326 + topk_num=recv_topk_ids.shape[1],
327 + expert_map=expert_map,
328 + HAS_EXPERT_MAP=expert_map is not None,
329 + num_warps=num_warps,
330 + BLOCK_D=BLOCK_D,
331 + )
332 + return
333 +
334 +
335 +def deepgemm_moe_permute(aq: torch.Tensor,
336 + aq_scale: torch.Tensor,
337 + topk_ids: torch.Tensor,
338 + local_num_experts: int,
339 + expert_map: Optional[torch.Tensor],
340 + expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
341 + aq_out: Optional[torch.Tensor] = None):
342 +
343 + assert aq.ndim == 2
344 + assert topk_ids.dtype.is_signed, (
345 + "The kernel uses -1 to represent invalid topk_ids")
346 + H = aq.size(1)
347 + device = aq.device
348 +
349 + block_m = deep_gemm_block_shape()[0]
350 + block_k = deep_gemm_block_shape()[1]
351 +
352 + M_sum = compute_aligned_M(M=topk_ids.size(0),
353 + num_topk=topk_ids.size(1),
354 + local_num_experts=local_num_experts,
355 + alignment=block_m,
356 + expert_tokens_meta=expert_tokens_meta)
357 +
358 + expert_start_loc = torch.empty((local_num_experts),
359 + device=device,
360 + dtype=torch.int32)
361 +
362 + assert aq_out is None or aq_out.shape == (M_sum, H)
363 + if aq_out is None:
364 + aq_out = torch.empty((M_sum, H), device=device, dtype=aq.dtype)
365 +
366 + aq_scale_out = torch.empty((M_sum, H // block_k),
367 + device=device,
368 + dtype=torch.float32)
369 +
370 + maybe_has_empty_blocks = ((expert_tokens_meta is None)
371 + or (expert_tokens_meta.expert_num_tokens_cpu
372 + is None))
373 + expert_ids_init = torch.zeros if maybe_has_empty_blocks else torch.empty
374 +
375 + expert_ids = expert_ids_init((M_sum), device=device, dtype=torch.int32)
376 + inv_perm = torch.empty(topk_ids.shape, device=device, dtype=torch.int32)
377 +
378 + expert_num_tokens = None
379 + if expert_tokens_meta is not None:
380 + expert_num_tokens = expert_tokens_meta.expert_num_tokens
381 + else:
382 + expert_num_tokens = count_expert_num_tokens(topk_ids,
383 + local_num_experts,
384 + expert_map)
385 +
386 + ep_scatter(recv_x=aq,
387 + recv_x_scale=aq_scale,
388 + recv_topk=topk_ids,
389 + num_recv_tokens_per_expert=expert_num_tokens,
390 + expert_start_loc=expert_start_loc,
391 + expert_map=expert_map,
392 + output_tensor=aq_out,
393 + output_tensor_scale=aq_scale_out,
394 + m_indices=expert_ids,
395 + output_index=inv_perm)
396 +
397 + return aq_out, aq_scale_out, expert_ids, inv_perm
398 +
399 +
400 +def deepgemm_unpermute_and_reduce(
401 + a: torch.Tensor, # Grouped gemm output
402 + topk_ids: torch.Tensor,
403 + topk_weights: torch.Tensor,
404 + inv_perm: torch.Tensor,
405 + expert_map: Optional[torch.Tensor],
406 + output: torch.Tensor):
407 +
408 + return ep_gather(input_tensor=a,
409 + recv_topk_ids=topk_ids,
410 + recv_topk_weight=topk_weights,
411 + input_index=inv_perm,
412 + expert_map=expert_map,
413 + output_tensor=output)
vllm/model_executor/layers/fused_moe/fused_batched_moe.py
+2
@@ -677,6 +677,7 @@ class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
677 topk: int,
678 global_num_experts: int,
679 local_num_experts: int,
680 + expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
681 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
682 assert a.dim() == 2
683 num_dp = self.num_dispatchers
@@ -889,6 +890,7 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
890 topk: int,
891 global_num_experts: int,
892 local_num_experts: int,
893 + expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
894 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
895 assert a.dim() == 2
896 num_dp = self.num_dispatchers
vllm/model_executor/layers/fused_moe/fused_moe.py
+1
@@ -1618,6 +1618,7 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1618 topk: int,
1619 global_num_experts: int,
1620 local_num_experts: int,
1621 + expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
1622 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
1623 workspace1 = (M, topk, max(N // 2, K))
1624 workspace2 = (M, topk, max(N, K))
vllm/model_executor/layers/fused_moe/modular_kernel.py
+10 -6
@@ -317,6 +317,7 @@ class FusedMoEPermuteExpertsUnpermute(ABC):
317 topk: int,
318 global_num_experts: int,
319 local_num_experts: int,
320 + expert_tokens_meta: Optional[ExpertTokensMetadata],
321 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
322 """
323 Compute the shapes for the temporary and final outputs of the two gemms
@@ -479,7 +480,8 @@ class FusedMoEModularKernel(torch.nn.Module):
480
481 (workspace13_shape, workspace2_shape, fused_out_shape,
482 workspace_dtype) = self.fused_experts.workspace_shapes(
482 - a1, a1q, M, N, K, top_k, global_num_experts, local_num_experts)
483 + a1, a1q, M, N, K, top_k, global_num_experts, local_num_experts,
484 + expert_tokens_meta)
485
486 # We can reuse the memory between cache1 and cache3 because by the
487 # time we need cache3, we're done with cache1.
@@ -572,10 +574,9 @@ class FusedMoEModularKernel(torch.nn.Module):
574 assert num_chunks > 1
575
576 # Construct the entire output that can then be processed in chunks.
575 - (_, _, fused_out_shape,
576 - _) = self.fused_experts.workspace_shapes(a1, a1q, M, N, K, top_k,
577 - global_num_experts,
578 - local_num_experts)
577 + (_, _, fused_out_shape, _) = self.fused_experts.workspace_shapes(
578 + a1, a1q, M, N, K, top_k, global_num_experts, local_num_experts,
579 + expert_tokens_meta)
580 fused_out = torch.empty(fused_out_shape,
581 device=a1q.device,
582 dtype=a1.dtype)
@@ -613,8 +614,11 @@ class FusedMoEModularKernel(torch.nn.Module):
614 need_expert_num_tokens_cpu = (
615 full_expert_tokens_meta.expert_num_tokens_cpu is not None)
616 if need_expert_num_tokens_cpu:
617 + # This is blocking as some implementations need the count
618 + # on the CPU to determine appropriate input/out fused-moe
619 + # buffers
620 c_expert_num_tokens_cpu = c_expert_num_tokens.to(
617 - "cpu", non_blocking=True)
621 + "cpu", non_blocking=False)
622
623 return ExpertTokensMetadata(
624 expert_num_tokens=c_expert_num_tokens,
vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py
+5 -2
@@ -102,6 +102,7 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
102 topk: int,
103 global_num_experts: int,
104 local_num_experts: int,
105 + expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
106 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
107 # Note: the deep gemm workspaces are strictly larger than the triton
108 # workspaces so we can be pessimistic here and allocate for DeepGemm
@@ -110,11 +111,13 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
111 or is_blackwell_deep_gemm_used()):
112 assert self.deep_gemm_expert is not None
113 return self.deep_gemm_expert.workspace_shapes(
113 - a, aq, M, N, K, topk, global_num_experts, local_num_experts)
114 + a, aq, M, N, K, topk, global_num_experts, local_num_experts,
115 + expert_tokens_meta)
116 else:
117 return self.triton_expert.workspace_shapes(a, aq, M, N, K, topk,
118 global_num_experts,
117 - local_num_experts)
119 + local_num_experts,
120 + expert_tokens_meta)
121
122 def apply(
123 self,