@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 569 lines 18.8 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Taken from https://github.com/ModelTC/LightLLM/blob/8ed97c74c18f11505b048b1ba00ba5c0cef8bff6/lightllm/common/fused_moe/deepep_scatter_gather.py
4 and updated to fit vllm needs and terminology.
5 """
6
7 import math
8
9 import torch
10
11 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
12 from vllm.model_executor.layers.fused_moe.utils import count_expert_num_tokens
13 from vllm.triton_utils import tl, triton
14 from vllm.utils.deep_gemm import get_mk_alignment_for_contiguous_layout
15 from vllm.utils.math_utils import round_up
16
17
18 def expert_num_tokens_round_up_and_sum(
19 expert_num_tokens: torch.Tensor, alignment: int
20 ) -> int:
21 # Round up each element in expert_num_tokens to the nearest multiple of
22 # alignment.
23 ent = (expert_num_tokens.to(torch.int64) + (alignment - 1)) // alignment * alignment
24 return torch.sum(ent).item()
25
26
27 def compute_aligned_M_and_alignment(
28 M: int,
29 num_topk: int,
30 local_num_experts: int,
31 alignment: int,
32 expert_tokens_meta: mk.ExpertTokensMetadata | None,
33 ) -> tuple[int, int]:
34 """Return (M_sum, alignment_used).
35
36 `alignment_used` may be smaller than the caller-supplied `alignment` on
37 SM100/SM120 when DeepGEMM can JIT a smaller BLOCK_M for the per-call
38 expected_m. Callers that index by block size (e.g. ``M_sum // block_m``)
39 or assert workspace alignment must use the returned `alignment_used`,
40 not their original `alignment` argument.
41
42 Prefer this over the int-returning :func:`compute_aligned_M` when the
43 GEMM call site needs to wrap itself in ``mk_alignment_scope`` or
44 otherwise reason about the actual per-expert padding.
45 """
46 if (expert_tokens_meta is not None) and (
47 expert_tokens_meta.expert_num_tokens_cpu is not None
48 ):
49 return (
50 expert_num_tokens_round_up_and_sum(
51 expert_tokens_meta.expert_num_tokens_cpu, alignment=alignment
52 ),
53 alignment,
54 )
55
56 # expert_num_tokens not on cpu. Cap padding by min(M*num_topk,
57 # local_num_experts) — at batch=1 decode only `num_topk` experts can be
58 # active, so the worst-case `local_num_experts*(align-1)` is too loose.
59 # Also shrink `alignment` to DeepGEMM's per-call theoretical BLOCK_M on
60 # SM100/SM120 when smaller.
61 expected_m = M * num_topk
62 try:
63 from vllm.utils.deep_gemm import (
64 get_theoretical_mk_alignment_for_contiguous_layout,
65 )
66
67 # num_groups=local_num_experts so the helper recovers per-expert em;
68 # omitting it over-picks BLOCK_M on SM120 (heuristic assumes em is
69 # already per-expert).
70 per_call_align = get_theoretical_mk_alignment_for_contiguous_layout(
71 expected_m=expected_m,
72 num_groups=local_num_experts,
73 )
74 if per_call_align and per_call_align <= alignment:
75 alignment = per_call_align
76 except Exception:
77 pass
78
79 max_active_experts = min(M * num_topk, local_num_experts)
80 M_sum = (M * num_topk) + max_active_experts * (alignment - 1)
81 M_sum = round_up(M_sum, alignment)
82 return M_sum, alignment
83
84
85 def compute_aligned_M(
86 M: int,
87 num_topk: int,
88 local_num_experts: int,
89 alignment: int,
90 expert_tokens_meta: mk.ExpertTokensMetadata | None,
91 ) -> int:
92 """Return ``M_sum`` only (backward-compat wrapper).
93
94 Equivalent to :func:`compute_aligned_M_and_alignment`'s first return
95 value. Existing downstream callers and the warmup path that only size
96 a workspace use this. Call sites that need the actual per-expert
97 alignment (to wrap GEMMs in ``mk_alignment_scope``) should use
98 :func:`compute_aligned_M_and_alignment` instead.
99 """
100 M_sum, _ = compute_aligned_M_and_alignment(
101 M, num_topk, local_num_experts, alignment, expert_tokens_meta
102 )
103 return M_sum
104
105
106 @triton.jit
107 def apply_expert_map(expert_id, expert_map):
108 if expert_id != -1:
109 expert_id = tl.load(expert_map + expert_id).to(expert_id.dtype)
110 return expert_id
111
112
113 @triton.jit
114 def _fwd_kernel_ep_scatter_1(
115 num_recv_tokens_per_expert,
116 expert_start_loc,
117 m_indices,
118 num_experts: tl.constexpr,
119 BLOCK_E: tl.constexpr,
120 BLOCK_EXPERT_NUM: tl.constexpr,
121 ALIGN_M: tl.constexpr,
122 ):
123 cur_expert = tl.program_id(0)
124
125 offset_cumsum = tl.arange(0, BLOCK_EXPERT_NUM)
126 tokens_per_expert = tl.load(
127 num_recv_tokens_per_expert + offset_cumsum,
128 mask=offset_cumsum < num_experts,
129 other=0,
130 )
131 # Round up to ALIGN_M so cumsum matches the workspace's per-expert slices.
132 tokens_per_expert = ((tokens_per_expert + ALIGN_M - 1) // ALIGN_M) * ALIGN_M
133 cumsum = tl.cumsum(tokens_per_expert) - tokens_per_expert
134
135 # Extract this block's offset from the register vector (warp shuffle,
136 # no global memory round-trip) then write it once to expert_start_loc.
137 cur_expert_start = tl.sum(
138 tl.where(offset_cumsum == cur_expert, cumsum, tl.zeros_like(cumsum))
139 )
140 tl.store(expert_start_loc + cur_expert, cur_expert_start)
141 cur_expert_token_num = tl.load(num_recv_tokens_per_expert + cur_expert)
142
143 m_indices_start_ptr = m_indices + cur_expert_start
144 off_expert = tl.arange(0, BLOCK_E)
145
146 # any rows in the per-expert aligned region that do not correspond to
147 # real tokens are left untouched here and should remain initialized to
148 # -1 so DeepGEMM can skip them
149 for start_m in tl.range(0, cur_expert_token_num, BLOCK_E):
150 offs = start_m + off_expert
151 mask = offs < cur_expert_token_num
152 tl.store(
153 m_indices_start_ptr + offs,
154 cur_expert,
155 mask=mask,
156 )
157
158
159 @triton.jit
160 def _fwd_kernel_ep_scatter_2(
161 total_token_num,
162 expert_start_loc,
163 recv_x,
164 recv_x_stride0,
165 recv_x_stride1,
166 recv_x_scale,
167 recv_x_scale_stride0,
168 recv_x_scale_stride1,
169 recv_topk,
170 recv_topk_stride0,
171 recv_topk_stride1,
172 output_tensor,
173 output_tensor_stride0,
174 output_tensor_stride1,
175 output_tensor_scale,
176 output_tensor_scale_stride0,
177 output_tensor_scale_stride1,
178 output_index,
179 output_index_stride0,
180 output_index_stride1,
181 topk_num: tl.constexpr,
182 expert_map,
183 HAS_EXPERT_MAP: tl.constexpr,
184 HIDDEN_SIZE: tl.constexpr,
185 HIDDEN_SIZE_PAD: tl.constexpr,
186 SCALE_HIDDEN_SIZE: tl.constexpr,
187 SCALE_HIDDEN_SIZE_PAD: tl.constexpr,
188 PACK_UE8M0: tl.constexpr,
189 SCALE_PACKED_SIZE: tl.constexpr,
190 SCALE_PACKED_SIZE_PAD: tl.constexpr,
191 ):
192 start_token_id = tl.program_id(0)
193 grid_num = tl.num_programs(0)
194
195 offset_in = tl.arange(0, HIDDEN_SIZE_PAD)
196 mask = offset_in < HIDDEN_SIZE
197
198 output_tensor_stride0 = output_tensor_stride0.to(tl.int64)
199
200 if PACK_UE8M0:
201 # One int32 per 4 consecutive 32-wide UE8M0 groups, stored MN-major.
202 offs_pk = tl.arange(0, SCALE_PACKED_SIZE_PAD)
203 mask_pk = offs_pk < SCALE_PACKED_SIZE
204 else:
205 offset_in_s = tl.arange(0, SCALE_HIDDEN_SIZE_PAD)
206 mask_s = offset_in_s < SCALE_HIDDEN_SIZE
207
208 for token_id in range(start_token_id, total_token_num, grid_num):
209 to_copy = tl.load(recv_x + token_id * recv_x_stride0 + offset_in, mask=mask)
210
211 if PACK_UE8M0:
212 # Pack 4 UE8M0 bytes into one int32 (byte j = group 4*pk+j).
213 base_s = recv_x_scale + token_id * recv_x_scale_stride0
214 g0, g1 = offs_pk * 4, offs_pk * 4 + 1
215 g2, g3 = offs_pk * 4 + 2, offs_pk * 4 + 3
216 b0 = tl.load(
217 base_s + g0 * recv_x_scale_stride1, mask=g0 < SCALE_HIDDEN_SIZE
218 )
219 b1 = tl.load(
220 base_s + g1 * recv_x_scale_stride1, mask=g1 < SCALE_HIDDEN_SIZE
221 )
222 b2 = tl.load(
223 base_s + g2 * recv_x_scale_stride1, mask=g2 < SCALE_HIDDEN_SIZE
224 )
225 b3 = tl.load(
226 base_s + g3 * recv_x_scale_stride1, mask=g3 < SCALE_HIDDEN_SIZE
227 )
228 packed_s = (
229 b0.to(tl.int32)
230 | (b1.to(tl.int32) << 8)
231 | (b2.to(tl.int32) << 16)
232 | (b3.to(tl.int32) << 24)
233 )
234 else:
235 to_copy_s = tl.load(
236 recv_x_scale + token_id * recv_x_scale_stride0 + offset_in_s,
237 mask=mask_s,
238 )
239
240 for topk_index in tl.range(0, topk_num, 1, num_stages=4):
241 expert_id = tl.load(recv_topk + token_id * recv_topk_stride0 + topk_index)
242
243 if HAS_EXPERT_MAP:
244 expert_id = apply_expert_map(expert_id, expert_map)
245
246 if expert_id >= 0:
247 dest_token_index = tl.atomic_add(expert_start_loc + expert_id, 1)
248 dest_token_index_i64 = dest_token_index.to(tl.int64)
249 tl.store(
250 output_index + token_id * output_index_stride0 + topk_index,
251 dest_token_index,
252 )
253 output_tensor_ptr = (
254 output_tensor + dest_token_index_i64 * output_tensor_stride0
255 )
256 tl.store(output_tensor_ptr + offset_in, to_copy, mask=mask)
257
258 output_tensor_scale_ptr = (
259 output_tensor_scale + dest_token_index * output_tensor_scale_stride0
260 )
261 if PACK_UE8M0:
262 tl.store(
263 output_tensor_scale_ptr + offs_pk * output_tensor_scale_stride1,
264 packed_s,
265 mask=mask_pk,
266 )
267 else:
268 tl.store(
269 output_tensor_scale_ptr + offset_in_s, to_copy_s, mask=mask_s
270 )
271
272
273 @torch.no_grad()
274 def ep_scatter(
275 recv_x: torch.Tensor,
276 recv_x_scale: torch.Tensor,
277 recv_topk: torch.Tensor,
278 num_recv_tokens_per_expert: torch.Tensor,
279 expert_map: torch.Tensor | None,
280 expert_start_loc: torch.Tensor,
281 output_tensor: torch.Tensor,
282 output_tensor_scale: torch.Tensor,
283 m_indices: torch.Tensor,
284 output_index: torch.Tensor,
285 align_m: int = 128,
286 block_size: int = 128,
287 pack_ue8m0: bool = False,
288 ):
289 # BLOCK_E is the m_indices fill-loop tile (masked), independent of align_m.
290 BLOCK_E = 128
291 BLOCK_D = block_size # block size of activation-scale quantization
292 num_warps = 8
293 num_experts = num_recv_tokens_per_expert.shape[0]
294 hidden_size = recv_x.shape[1]
295 # grid = (triton.cdiv(hidden_size, BLOCK_D), num_experts)
296 grid = num_experts
297
298 assert m_indices.shape[0] % align_m == 0
299 assert expert_start_loc.shape[0] == num_experts
300
301 # pack_ue8m0: scatter packs 4 UE8M0 bytes per int32; else copies scales as-is.
302 scale_hidden_size = hidden_size // BLOCK_D
303 scale_packed_size = (scale_hidden_size + 3) // 4 if pack_ue8m0 else 1
304
305 _fwd_kernel_ep_scatter_1[(grid,)](
306 num_recv_tokens_per_expert,
307 expert_start_loc,
308 m_indices,
309 num_experts=num_experts,
310 num_warps=num_warps,
311 BLOCK_E=BLOCK_E,
312 BLOCK_EXPERT_NUM=triton.next_power_of_2(num_experts),
313 ALIGN_M=align_m,
314 )
315
316 grid = min(recv_topk.shape[0], 1024 * 8)
317
318 _fwd_kernel_ep_scatter_2[(grid,)](
319 recv_topk.shape[0],
320 expert_start_loc,
321 recv_x,
322 recv_x.stride(0),
323 recv_x.stride(1),
324 recv_x_scale,
325 recv_x_scale.stride(0),
326 recv_x_scale.stride(1),
327 recv_topk,
328 recv_topk.stride(0),
329 recv_topk.stride(1),
330 output_tensor,
331 output_tensor.stride(0),
332 output_tensor.stride(1),
333 output_tensor_scale,
334 output_tensor_scale.stride(0),
335 output_tensor_scale.stride(1),
336 output_index,
337 output_index.stride(0),
338 output_index.stride(1),
339 topk_num=recv_topk.shape[1],
340 expert_map=expert_map,
341 HAS_EXPERT_MAP=expert_map is not None,
342 num_warps=num_warps,
343 HIDDEN_SIZE=hidden_size,
344 HIDDEN_SIZE_PAD=triton.next_power_of_2(hidden_size),
345 SCALE_HIDDEN_SIZE=scale_hidden_size,
346 SCALE_HIDDEN_SIZE_PAD=triton.next_power_of_2(scale_hidden_size),
347 PACK_UE8M0=pack_ue8m0,
348 SCALE_PACKED_SIZE=scale_packed_size,
349 SCALE_PACKED_SIZE_PAD=triton.next_power_of_2(scale_packed_size),
350 )
351 return
352
353
354 @triton.jit
355 def _fwd_kernel_ep_gather(
356 total_token_num,
357 input_tensor,
358 input_tensor_stride0,
359 input_tensor_stride1,
360 recv_topk_ids,
361 recv_topk_ids_stride0,
362 recv_topk_ids_stride1,
363 recv_topk_weight,
364 recv_topk_weight_stride0,
365 recv_topk_weight_stride1,
366 input_index,
367 input_index_stride0,
368 input_index_stride1,
369 output_tensor,
370 output_tensor_stride0,
371 output_tensor_stride1,
372 topk_num: tl.constexpr,
373 expert_map,
374 HAS_EXPERT_MAP: tl.constexpr,
375 BLOCK_D: tl.constexpr,
376 ):
377 cur_block = tl.program_id(0)
378 start_cur_token = tl.program_id(1)
379 grid_num = tl.num_programs(1)
380
381 for cur_token in range(start_cur_token, total_token_num, grid_num):
382 off_d = tl.arange(0, BLOCK_D)
383 accumulator = tl.zeros([BLOCK_D], dtype=tl.float32)
384 for topk_index in range(0, topk_num):
385 expert_id = tl.load(
386 recv_topk_ids + cur_token * recv_topk_ids_stride0 + topk_index
387 )
388
389 if HAS_EXPERT_MAP:
390 expert_id = apply_expert_map(expert_id, expert_map)
391
392 if expert_id >= 0:
393 source_token_index = tl.load(
394 input_index + cur_token * input_index_stride0 + topk_index
395 )
396 source_token_index_i64 = source_token_index.to(tl.int64)
397 acc_weight = tl.load(
398 recv_topk_weight + cur_token * recv_topk_weight_stride0 + topk_index
399 )
400 tmp = tl.load(
401 input_tensor
402 + source_token_index_i64 * input_tensor_stride0
403 + cur_block * BLOCK_D
404 + off_d
405 )
406 accumulator += tmp.to(tl.float32) * acc_weight
407
408 tl.store(
409 output_tensor
410 + cur_token * output_tensor_stride0
411 + cur_block * BLOCK_D
412 + off_d,
413 accumulator.to(output_tensor.dtype.element_ty),
414 )
415
416
417 @torch.no_grad()
418 def ep_gather(
419 input_tensor: torch.Tensor,
420 recv_topk_ids: torch.Tensor,
421 recv_topk_weight: torch.Tensor,
422 input_index: torch.Tensor,
423 expert_map: torch.Tensor | None,
424 output_tensor: torch.Tensor,
425 ):
426 num_warps = 2
427 num_tokens = output_tensor.shape[0]
428 hidden_size = input_tensor.shape[1]
429 BLOCK_D = math.gcd(hidden_size, 1024)
430 assert hidden_size % BLOCK_D == 0
431 grid = (triton.cdiv(hidden_size, BLOCK_D), min(num_tokens, 1024))
432
433 _fwd_kernel_ep_gather[grid](
434 num_tokens,
435 input_tensor,
436 input_tensor.stride(0),
437 input_tensor.stride(1),
438 recv_topk_ids,
439 recv_topk_ids.stride(0),
440 recv_topk_ids.stride(1),
441 recv_topk_weight,
442 recv_topk_weight.stride(0),
443 recv_topk_weight.stride(1),
444 input_index,
445 input_index.stride(0),
446 input_index.stride(1),
447 output_tensor,
448 output_tensor.stride(0),
449 output_tensor.stride(1),
450 topk_num=recv_topk_ids.shape[1],
451 expert_map=expert_map,
452 HAS_EXPERT_MAP=expert_map is not None,
453 num_warps=num_warps,
454 BLOCK_D=BLOCK_D,
455 )
456 return
457
458
459 def deepgemm_moe_permute(
460 aq: torch.Tensor,
461 aq_scale: torch.Tensor,
462 topk_ids: torch.Tensor,
463 local_num_experts: int,
464 expert_map: torch.Tensor | None,
465 expert_tokens_meta: mk.ExpertTokensMetadata | None,
466 aq_out: torch.Tensor | None = None,
467 block_size: int | None = None,
468 ):
469 assert aq.ndim == 2
470 assert topk_ids.dtype.is_signed, "The kernel uses -1 to represent invalid topk_ids"
471 H = aq.size(1)
472 device = aq.device
473
474 block_m, block_k = get_mk_alignment_for_contiguous_layout()
475 # The activation-scale group size may differ from the M/K tile alignment
476 # (e.g. MXFP8 uses a 32-element scale group while block_k stays 128).
477 if block_size is not None:
478 block_k = block_size
479
480 M_sum, align_used = compute_aligned_M_and_alignment(
481 M=topk_ids.size(0),
482 num_topk=topk_ids.size(1),
483 local_num_experts=local_num_experts,
484 alignment=block_m,
485 expert_tokens_meta=expert_tokens_meta,
486 )
487
488 expert_start_loc = torch.empty(
489 (local_num_experts), device=device, dtype=torch.int32
490 )
491
492 assert aq_out is None or aq_out.shape == (M_sum, H)
493 if aq_out is None:
494 aq_out = torch.empty((M_sum, H), device=device, dtype=aq.dtype)
495
496 # uint8 UE8M0 (MXFP8) -> scatter packs into DeepGEMM's int32 MN-major
497 # TMA-aligned layout; float32 (FP8/FP4) scattered row-major as-is.
498 pack_ue8m0 = aq_scale.dtype == torch.uint8
499 sf_k = H // block_k
500 if pack_ue8m0:
501 packed_sf_k = (sf_k + 3) // 4
502 tma_aligned_mn = round_up(M_sum, 4)
503 aq_scale_out = torch.empty_strided(
504 (M_sum, packed_sf_k),
505 (1, tma_aligned_mn),
506 device=device,
507 dtype=torch.int32,
508 )
509 else:
510 aq_scale_out = torch.zeros((M_sum, sf_k), device=device, dtype=torch.float32)
511
512 # DeepGEMM uses negative values in m_indices (here expert_ids) to mark
513 # completely invalid / padded blocks that should be skipped. We always
514 # initialize expert_ids to -1 so any row that is not explicitly written
515 # by the scatter kernel will be treated as invalid and skipped by
516 # DeepGEMM's scheduler.
517 expert_ids = torch.full(
518 (M_sum,),
519 fill_value=-1,
520 device=device,
521 dtype=torch.int32,
522 )
523 inv_perm = torch.empty(topk_ids.shape, device=device, dtype=torch.int32)
524
525 expert_num_tokens = None
526 if expert_tokens_meta is not None:
527 expert_num_tokens = expert_tokens_meta.expert_num_tokens
528 if expert_num_tokens is None:
529 # A cudagraph-decode carrier may hold only psum_recv_per_rank and leave
530 # expert_num_tokens None; count locally as in the meta-absent case.
531 expert_num_tokens = count_expert_num_tokens(
532 topk_ids, local_num_experts, expert_map
533 )
534
535 ep_scatter(
536 recv_x=aq,
537 recv_x_scale=aq_scale,
538 recv_topk=topk_ids,
539 num_recv_tokens_per_expert=expert_num_tokens,
540 expert_start_loc=expert_start_loc,
541 expert_map=expert_map,
542 output_tensor=aq_out,
543 output_tensor_scale=aq_scale_out,
544 m_indices=expert_ids,
545 output_index=inv_perm,
546 align_m=align_used,
547 block_size=block_k,
548 pack_ue8m0=pack_ue8m0,
549 )
550
551 return aq_out, aq_scale_out, expert_ids, inv_perm, align_used
552
553
554 def deepgemm_unpermute_and_reduce(
555 a: torch.Tensor, # Grouped gemm output
556 topk_ids: torch.Tensor,
557 topk_weights: torch.Tensor,
558 inv_perm: torch.Tensor,
559 expert_map: torch.Tensor | None,
560 output: torch.Tensor,
561 ):
562 return ep_gather(
563 input_tensor=a,
564 recv_topk_ids=topk_ids,
565 recv_topk_weight=topk_weights,
566 input_index=inv_perm,
567 expert_map=expert_map,
568 output_tensor=output,
569 )