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 14, 05:30 UTC
[Log] DeepGEMM Update Log for Unaligned Problem Size (#22208)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Aug 4, 2025 at 22:13 UTC
d7b28f34153a5116174383d97e41a1279b51e5cb
3 files changed
+23
-8
vllm/model_executor/layers/fused_moe/deep_gemm_moe.py
+19
-2
@@ -33,7 +33,7 @@ def deep_gemm_block_shape() -> list[int]:
33
return [block, block]
34
35
36
-def _valid_deep_gemm_shape(M: int, N: int, K: int):
36
+def _valid_deep_gemm_shape(M: int, N: int, K: int) -> bool:
37
align = deep_gemm_block_shape()[0]
38
return align <= M and N % align == 0 and K % align == 0
39
@@ -51,9 +51,26 @@ def _valid_deep_gemm(hidden_states: torch.Tensor, w1: torch.Tensor,
51
52
M = hidden_states.size(0)
53
_, K, N = w2.size()
54
+
55
+ align = deep_gemm_block_shape()[0]
56
+
57
if not _valid_deep_gemm_shape(M, N, K):
58
logger.debug_once(
56
- "DeepGemm disabled: unaligned problem size. M: %s, N: %s, K: %s",
59
+ "DeepGemm disabled due to unaligned problem size. "
60
+ "M: %s, N: %s, K: %s. M should >= align size "
61
+ "and N and K must be multiples of %s."
62
+ "This is not an error and we will fall back to triton.",
63
+ M,
64
+ N,
65
+ K,
66
+ align,
67
+ )
68
+ return False
69
+ elif N <= 512:
70
+ logger.debug_once(
71
+ "DeepGemm disabled for N <= 512. M: %s, N: %s, K: %s. "
72
+ "This means we will fallback to triton "
73
+ "for this specific shape for further speed up.",
74
M,
75
N,
76
K,
vllm/model_executor/layers/fused_moe/fused_moe.py
+2
-4
@@ -1360,10 +1360,8 @@ def fused_experts(
1360
# E8M0 scale, which means we requantize the weight and input to the specific
1361
# scale. Fallen back to cutlass or triton for some cases would cause
1362
# accuracy issue.
1363
- N = w1.size(1)
1364
- should_use_deep_gemm = ((N > 512
1365
- and _valid_deep_gemm(hidden_states, w1, w2))
1366
- or is_blackwell_deep_gemm_used())
1363
+ should_use_deep_gemm = is_blackwell_deep_gemm_used() or _valid_deep_gemm(
1364
+ hidden_states, w1, w2)
1365
if (allow_deep_gemm and use_fp8_w8a8 and should_use_deep_gemm):
1366
assert apply_router_weight_on_input is False
1367
assert is_act_and_mul, (
vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py
+2
-2
@@ -107,8 +107,8 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
107
# Note: the deep gemm workspaces are strictly larger than the triton
108
# workspaces so we can be pessimistic here and allocate for DeepGemm
109
# even if we fall back to triton later, e.g. if expert maps are set.
110
- if self.allow_deep_gemm and (_valid_deep_gemm_shape(M, N, K)
111
- or is_blackwell_deep_gemm_used()):
110
+ if self.allow_deep_gemm and (is_blackwell_deep_gemm_used()
111
+ or _valid_deep_gemm_shape(M, N, K)):
112
assert self.deep_gemm_expert is not None
113
return self.deep_gemm_expert.workspace_shapes(
114
a, aq, M, N, K, topk, global_num_experts, local_num_experts,