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
[UX] Speedup DeepGEMM warmup with heuristics (#25619)
Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: Michael Goin <mgoin64@gmail.com> Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
Michael Goin committed
Oct 13, 2025 at 10:59 UTC
0d21b9b51eccabfa1f8114eab2df61d75459bee7
3 files changed
+95
-14
vllm/envs.py
+20
-4
@@ -146,7 +146,11 @@ if TYPE_CHECKING:
146
VLLM_TPU_USING_PATHWAYS: bool = False
147
VLLM_USE_DEEP_GEMM: bool = True
148
VLLM_USE_DEEP_GEMM_E8M0: bool = True
149
- VLLM_SKIP_DEEP_GEMM_WARMUP: bool = False
149
+ VLLM_DEEP_GEMM_WARMUP: Literal[
150
+ "skip",
151
+ "full",
152
+ "relax",
153
+ ] = "relax"
154
VLLM_USE_FUSED_MOE_GROUPED_TOPK: bool = True
155
VLLM_USE_FLASHINFER_MOE_FP16: bool = False
156
VLLM_USE_FLASHINFER_MOE_FP8: bool = False
@@ -1088,9 +1092,21 @@ environment_variables: dict[str, Callable[[], Any]] = {
1092
# JIT all the required kernels before model execution so there is no
1093
# JIT'ing in the hot-path. However, this warmup increases the engine
1094
# startup time by a couple of minutes.
1091
- # Set `VLLM_SKIP_DEEP_GEMM_WARMUP` to disable the warmup.
1092
- "VLLM_SKIP_DEEP_GEMM_WARMUP": lambda: bool(
1093
- int(os.getenv("VLLM_SKIP_DEEP_GEMM_WARMUP", "0"))
1095
+ # Available options:
1096
+ # - "skip" : Skip warmup.
1097
+ # - "full" : Warmup deepgemm by running all possible gemm shapes the
1098
+ # engine could encounter.
1099
+ # - "relax" : Select gemm shapes to run based on some heuristics. The
1100
+ # heuristic aims to have the same effect as running all possible gemm
1101
+ # shapes, but provides no guarantees.
1102
+ "VLLM_DEEP_GEMM_WARMUP": env_with_choices(
1103
+ "VLLM_DEEP_GEMM_WARMUP",
1104
+ "relax",
1105
+ [
1106
+ "skip",
1107
+ "full",
1108
+ "relax",
1109
+ ],
1110
),
1111
# Whether to use fused grouped_topk used for MoE expert selection.
1112
"VLLM_USE_FUSED_MOE_GROUPED_TOPK": lambda: bool(
vllm/model_executor/warmup/deep_gemm_warmup.py
+74
-9
@@ -26,6 +26,55 @@ from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
26
from vllm.utils.deep_gemm import fp8_gemm_nt, m_grouped_fp8_gemm_nt_contiguous
27
28
29
+def _generate_optimal_warmup_m_values(
30
+ max_tokens: int, n: int, device: torch.device
31
+) -> list[int]:
32
+ """
33
+ Generate M values that cover all possible DeepGEMM kernel configurations.
34
+ Reference: https://github.com/deepseek-ai/DeepGEMM/blob/79f48ee15a82dd5fad5cd9beaa393c1f755e6b55/csrc/jit_kernels/heuristics/common.hpp
35
+
36
+ Args:
37
+ max_tokens: Maximum number of tokens to warmup for
38
+ n: The actual N dimension from the weight tensor
39
+ device: The torch device to get properties from.
40
+ """
41
+
42
+ def ceil_div(a: int, b: int) -> int:
43
+ return (a + b - 1) // b
44
+
45
+ # DeepGEMM's possible block sizes
46
+ block_ms = [64, 128, 256]
47
+ block_ns = list(range(16, min(257, n + 1), 16))
48
+ num_sms = torch.cuda.get_device_properties(device).multi_processor_count
49
+
50
+ m_values = set()
51
+
52
+ # Always include small cases
53
+ m_values.update([1, 2, 4] + [i for i in range(8, 65, 8)])
54
+
55
+ # Collect M values where different wave patterns occur
56
+ for block_m in block_ms:
57
+ for block_n in block_ns:
58
+ if block_n > n:
59
+ continue
60
+
61
+ # Add key M boundaries for this block combination
62
+ for wave in range(1, 11): # Up to 10 waves
63
+ # M where this block config transitions to next wave
64
+ target_blocks = wave * num_sms
65
+ m = target_blocks * block_m // ceil_div(n, block_n)
66
+ if 1 <= m <= max_tokens:
67
+ m_values.add(m)
68
+
69
+ # Add block_m boundaries
70
+ for multiple in range(1, max_tokens // block_m + 1):
71
+ m = multiple * block_m
72
+ if m <= max_tokens:
73
+ m_values.add(m)
74
+
75
+ return sorted(m_values)
76
+
77
+
78
def _extract_data_from_linear_base_module(
79
m: torch.nn.Module,
80
) -> tuple[torch.Tensor, torch.Tensor, list[int]]:
@@ -136,14 +185,27 @@ def _deepgemm_fp8_gemm_nt_warmup(w: torch.Tensor, ws: torch.Tensor, max_tokens:
185
)
186
out = torch.empty((max_tokens, n), device=device, dtype=torch.bfloat16)
187
139
- pbar = tqdm(total=max_tokens, desc=f"DeepGemm(fp8_gemm_nt) warmup (W={w.size()})")
140
- num_tokens = max_tokens
141
- while num_tokens > 0:
188
+ # Use optimal M values only if VLLM_DEEP_GEMM_WARMUP is set to "relax".
189
+ # Otherwise warmup all token sizes to avoid JIT compilation in hotpath
190
+ if envs.VLLM_DEEP_GEMM_WARMUP == "relax":
191
+ m_values = _generate_optimal_warmup_m_values(max_tokens, n, device)
192
+ desc = f"DeepGemm(fp8_gemm_nt) warmup (W={w.size()}) [relaxed]"
193
+ else:
194
+ assert envs.VLLM_DEEP_GEMM_WARMUP == "full", (
195
+ "Expected "
196
+ 'VLLM_DEEP_GEMM_WARMUP env to be set to "full" but got '
197
+ f"{envs.VLLM_DEEP_GEMM_WARMUP}"
198
+ )
199
+ m_values = list(range(1, max_tokens + 1))
200
+ desc = f"DeepGemm(fp8_gemm_nt) warmup (W={w.size()}) [all tokens]"
201
+
202
+ pbar = tqdm(total=len(m_values), desc=desc)
203
+
204
+ for num_tokens in m_values:
205
fp8_gemm_nt(
206
(a1q[:num_tokens], a1q_scales[:num_tokens]), (w, ws), out[:num_tokens]
207
)
208
pbar.update(1)
146
- num_tokens -= 1
209
210
FP8_GEMM_NT_WARMUP_CACHE.add(w.size())
211
@@ -195,12 +257,16 @@ def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(
257
)
258
out = torch.empty((MAX_M, n), device=device, dtype=torch.bfloat16)
259
260
+ # Generate M values in block_m increments (already optimized for MoE)
261
+ m_values = list(range(block_m, MAX_M + 1, block_m))
262
+
263
pbar = tqdm(
199
- total=MAX_BLOCKS,
200
- desc=f"DeepGemm(m_grouped_fp8_gemm_nt_contiguous) warmup (W={w.size()})",
264
+ total=len(m_values),
265
+ desc=f"DeepGemm(m_grouped_fp8_gemm_nt_contiguous) warmup (W={w.size()}) "
266
+ f"[{len(m_values)} values, block_m={block_m}]",
267
)
202
- num_tokens = MAX_M
203
- while num_tokens > 0:
268
+
269
+ for num_tokens in m_values:
270
m_grouped_fp8_gemm_nt_contiguous(
271
(a1q[:num_tokens], a1q_scales[:num_tokens]),
272
(w, w_scale),
@@ -208,7 +274,6 @@ def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(
274
expert_ids[:num_tokens],
275
)
276
pbar.update(1)
211
- num_tokens = num_tokens - block_m
277
278
for w, ws in [(w1, w1_scale), (w2, w2_scale)]:
279
if w.size() not in GROUPED_FP8_GEMM_NT_CONTIGUOUS_WARMUP_CACHE:
vllm/model_executor/warmup/kernel_warmup.py
+1
-1
@@ -29,7 +29,7 @@ def kernel_warmup(worker: "Worker"):
29
do_deep_gemm_warmup = (
30
envs.VLLM_USE_DEEP_GEMM
31
and is_deep_gemm_supported()
32
- and not envs.VLLM_SKIP_DEEP_GEMM_WARMUP
32
+ and envs.VLLM_DEEP_GEMM_WARMUP != "skip"
33
)
34
if do_deep_gemm_warmup:
35
model = worker.get_model()