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, 06:00 UTC
[BugFix] : Fix Batched DeepGemm Experts (#19515)
Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com> Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
Varun Sundar Rabindranath committed
Jun 12, 2025 at 22:43 UTC
e3b12667d46f74092d87777ffe9736e4baa116ae
9 files changed
+52
-32
vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py
+11
-8
@@ -47,15 +47,21 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
47
N: int,
48
K: int,
49
topk: int,
50
- num_experts: int,
50
+ global_num_experts: int,
51
+ local_num_experts: int,
52
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
53
assert a.dim() == 2
53
- num_dp = self.world_size // self.dp_size
54
+ # FIXME (varun): We should be able to dispatch only from the leader
55
+ # DP ranks in the case of TP > 1. At the moment, all the Ranks
56
+ # end up sending their tokens. This needs to be fixed.
57
+ num_dispatchers = self.world_size
58
+ num_experts = local_num_experts
59
max_num_tokens = a.size(
60
0) if self.max_num_tokens is None else self.max_num_tokens
56
- workspace13 = (num_experts, max_num_tokens * num_dp, max(K, N))
57
- workspace2 = (num_experts, max_num_tokens * num_dp, (N // 2))
58
- output = (num_experts, max_num_tokens * num_dp, K)
61
+ workspace13 = (num_experts, max_num_tokens * num_dispatchers,
62
+ max(K, N))
63
+ workspace2 = (num_experts, max_num_tokens * num_dispatchers, (N // 2))
64
+ output = (num_experts, max_num_tokens * num_dispatchers, K)
65
return (workspace13, workspace2, output, a.dtype)
66
67
def apply(
@@ -84,9 +90,6 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
90
a1q = hidden_states
91
_, N, K = w1.size()
92
87
- if global_num_experts == -1:
88
- global_num_experts = w1.size(0)
89
-
93
assert w2.size(1) == K
94
95
E, max_num_tokens, N, K, top_k_num = mk._moe_problem_size(
vllm/model_executor/layers/fused_moe/batched_triton_or_deep_gemm_moe.py
+4
-3
@@ -81,18 +81,19 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
81
N: int,
82
K: int,
83
topk: int,
84
- num_experts: int,
84
+ global_num_experts: int,
85
+ local_num_experts: int,
86
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
87
# Note: the deep gemm workspaces are strictly larger than the triton
88
# workspaces so we can be pessimistic here and allocate for DeepGemm
89
# even if we fall back to triton later, e.g. if expert maps are set.
90
if self.allow_deep_gemm and self.batched_deep_gemm_experts is not None:
91
return self.batched_deep_gemm_experts.workspace_shapes(
91
- a, aq, M, N, K, topk, num_experts)
92
+ a, aq, M, N, K, topk, global_num_experts, local_num_experts)
93
else:
94
assert self.batched_triton_experts is not None
95
return self.batched_triton_experts.workspace_shapes(
95
- a, aq, M, N, K, topk, num_experts)
96
+ a, aq, M, N, K, topk, global_num_experts, local_num_experts)
97
98
def apply(
99
self,
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+2
-1
@@ -230,7 +230,8 @@ class CutlassExpertsFp8(mk.FusedMoEPermuteExpertsUnpermute):
230
N: int,
231
K: int,
232
topk: int,
233
- num_experts: int,
233
+ global_num_experts: int,
234
+ local_num_experts: int,
235
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
236
workspace1: tuple[int, ...] = ()
237
workspace2: tuple[int, ...] = ()
vllm/model_executor/layers/fused_moe/deep_gemm_moe.py
+5
-8
@@ -74,15 +74,12 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
74
return True
75
76
def workspace_shapes(
77
- self,
78
- a: torch.Tensor,
79
- aq: torch.Tensor,
80
- M: int,
81
- N: int,
82
- K: int,
83
- topk: int,
84
- num_experts: int,
77
+ self, a: torch.Tensor, aq: torch.Tensor, M: int, N: int, K: int,
78
+ topk: int, global_num_experts: int, local_num_experts: int
79
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
80
+ # We use global_num_experts due to how moe_align_block_size handles
81
+ # expert_maps.
82
+ num_experts = global_num_experts
83
block_m = self.block_shape[0]
84
M_sum = (M * topk) + num_experts * (block_m - 1)
85
M_sum = round_up(M_sum, block_m)
vllm/model_executor/layers/fused_moe/fused_batched_moe.py
+7
-3
@@ -521,10 +521,12 @@ class BatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
521
N: int,
522
K: int,
523
topk: int,
524
- num_experts: int,
524
+ global_num_experts: int,
525
+ local_num_experts: int,
526
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
527
assert a.dim() == 2
527
- num_dp = self.world_size // self.dp_size
528
+ num_dp = self.dp_size
529
+ num_experts = local_num_experts
530
workspace13 = (num_experts, self.max_num_tokens * num_dp, K)
531
workspace2 = (self.max_num_tokens * num_dp, N)
532
return (workspace13, workspace2, workspace13, a.dtype)
@@ -624,10 +626,12 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
626
N: int,
627
K: int,
628
topk: int,
627
- num_experts: int,
629
+ global_num_experts: int,
630
+ local_num_experts: int,
631
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
632
assert a.dim() == 2
633
num_dp = self.world_size // self.dp_size
634
+ num_experts = local_num_experts
635
max_num_tokens = a.size(
636
0) if self.max_num_tokens is None else self.max_num_tokens
637
workspace13 = (num_experts, max_num_tokens * num_dp, max(K, N))
vllm/model_executor/layers/fused_moe/fused_moe.py
+2
-1
@@ -1553,7 +1553,8 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1553
N: int,
1554
K: int,
1555
topk: int,
1556
- num_experts: int,
1556
+ global_num_experts: int,
1557
+ local_num_experts: int,
1558
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
1559
workspace1 = (M, topk, max(N * 2, K))
1560
workspace2 = (M, topk, N)
vllm/model_executor/layers/fused_moe/modular_kernel.py
+10
-5
@@ -194,7 +194,8 @@ class FusedMoEPermuteExpertsUnpermute(ABC):
194
N: int,
195
K: int,
196
topk: int,
197
- num_experts: int,
197
+ global_num_experts: int,
198
+ local_num_experts: int,
199
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
200
"""
201
Compute the shapes for the temporary and final outputs of the two gemms
@@ -372,8 +373,9 @@ class FusedMoEModularKernel(torch.nn.Module):
373
a1 = hidden_states
374
output = a1 if inplace else torch.zeros_like(a1)
375
376
+ local_num_experts = w1.size(0)
377
if global_num_experts == -1:
376
- global_num_experts = w1.size(0)
378
+ global_num_experts = local_num_experts
379
380
(a1q, a1q_scale, expert_num_tokens, _expert_topk_ids,
381
_expert_topk_weights) = self.prepare_finalize.prepare(
@@ -408,16 +410,19 @@ class FusedMoEModularKernel(torch.nn.Module):
410
if num_chunks == 1:
411
(workspace13_shape, workspace2_shape, fused_out_shape,
412
workspace_dtype) = self.fused_experts.workspace_shapes(
411
- a1, a1q, M, N, K, top_k, global_num_experts)
413
+ a1, a1q, M, N, K, top_k, global_num_experts,
414
+ local_num_experts)
415
else:
416
# Use the full M to get the final output shape.
417
_, _, fused_out_shape, _ = (
418
self.fused_experts.workspace_shapes(
416
- a1, a1q, M, N, K, top_k, global_num_experts))
419
+ a1, a1q, M, N, K, top_k, global_num_experts,
420
+ local_num_experts))
421
# Use the CHUNK_SIZE to get the workspace shapes.
422
workspace13_shape, workspace2_shape, _, workspace_dtype = (
423
self.fused_experts.workspace_shapes(
420
- a1, a1q, CHUNK_SIZE, N, K, top_k, global_num_experts))
424
+ a1, a1q, CHUNK_SIZE, N, K, top_k, global_num_experts,
425
+ local_num_experts))
426
427
# We can reuse the memory between cache1 and cache3 because by the
428
# time we need cache3, we're done with cache1.
vllm/model_executor/layers/fused_moe/moe_align_block_size.py
+6
@@ -159,6 +159,12 @@ def moe_align_block_size(
159
Aligns the token distribution across experts to be compatible with block
160
size for matrix multiplication.
161
162
+ Note: In the case of expert_parallel, moe_align_block_size initially
163
+ considers all experts as valid and aligns all tokens appropriately.
164
+ Before the function returns it marks the experts_ids that are not in
165
+ the current GPU rank as -1 so the MoE matmuls could skip those blocks.
166
+ This requires the num_experts input arg to be the num global experts.
167
+
168
Parameters:
169
- topk_ids: A tensor of shape [total_tokens, top_k] representing the
170
top-k expert indices for each token.
vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py
+5
-3
@@ -48,7 +48,8 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
48
N: int,
49
K: int,
50
topk: int,
51
- num_experts: int,
51
+ global_num_experts: int,
52
+ local_num_experts: int,
53
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
54
# Note: the deep gemm workspaces are strictly larger than the triton
55
# workspaces so we can be pessimistic here and allocate for DeepGemm
@@ -56,10 +57,11 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
57
if self.allow_deep_gemm and _valid_deep_gemm_shape(M, N, K):
58
assert self.deep_gemm_expert is not None
59
return self.deep_gemm_expert.workspace_shapes(
59
- a, aq, M, N, K, topk, num_experts)
60
+ a, aq, M, N, K, topk, global_num_experts, local_num_experts)
61
else:
62
return self.triton_expert.workspace_shapes(a, aq, M, N, K, topk,
62
- num_experts)
63
+ global_num_experts,
64
+ local_num_experts)
65
66
def apply(
67
self,