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 11, 06:15 UTC
[UX] Reduce DeepGEMM warmup log output to single progress bar (#30903)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Matthew Bonanni committed
Dec 17, 2025 at 23:21 UTC
4a8412f773c67e0ba1eb1d4992095d6e0204f0ce
1 file changed
+99
-42
vllm/model_executor/warmup/deep_gemm_warmup.py
+99
-42
@@ -10,7 +10,7 @@ import torch
10
from tqdm import tqdm
11
12
import vllm.envs as envs
13
-from vllm.distributed.parallel_state import get_dp_group
13
+from vllm.distributed.parallel_state import get_dp_group, is_global_first_rank
14
from vllm.model_executor.layers.fused_moe.deep_gemm_moe import DeepGemmExperts
15
from vllm.model_executor.layers.fused_moe.deep_gemm_utils import compute_aligned_M
16
from vllm.model_executor.layers.fused_moe.layer import FusedMoE, FusedMoEModularMethod
@@ -175,7 +175,30 @@ def _fused_moe_grouped_gemm_may_use_deep_gemm(module: torch.nn.Module) -> bool:
175
FP8_GEMM_NT_WARMUP_CACHE: set[torch.Size] = set()
176
177
178
-def _deepgemm_fp8_gemm_nt_warmup(w: torch.Tensor, ws: torch.Tensor, max_tokens: int):
178
+def _get_fp8_gemm_nt_m_values(w: torch.Tensor, max_tokens: int) -> list[int]:
179
+ """Get the M values to warmup for a given weight tensor."""
180
+ n, _ = w.size()
181
+ device = w.device
182
+
183
+ # Use optimal M values only if VLLM_DEEP_GEMM_WARMUP is set to "relax".
184
+ # Otherwise warmup all token sizes to avoid JIT compilation in hotpath
185
+ if envs.VLLM_DEEP_GEMM_WARMUP == "relax":
186
+ return _generate_optimal_warmup_m_values(max_tokens, n, device)
187
+ else:
188
+ assert envs.VLLM_DEEP_GEMM_WARMUP == "full", (
189
+ "Expected "
190
+ 'VLLM_DEEP_GEMM_WARMUP env to be set to "full" but got '
191
+ f"{envs.VLLM_DEEP_GEMM_WARMUP}"
192
+ )
193
+ return list(range(1, max_tokens + 1))
194
+
195
+
196
+def _deepgemm_fp8_gemm_nt_warmup(
197
+ w: torch.Tensor,
198
+ ws: torch.Tensor,
199
+ max_tokens: int,
200
+ pbar: tqdm | None = None,
201
+):
202
if w.size() in FP8_GEMM_NT_WARMUP_CACHE:
203
return
204
@@ -189,27 +212,14 @@ def _deepgemm_fp8_gemm_nt_warmup(w: torch.Tensor, ws: torch.Tensor, max_tokens:
212
)
213
out = torch.empty((max_tokens, n), device=device, dtype=torch.bfloat16)
214
192
- # Use optimal M values only if VLLM_DEEP_GEMM_WARMUP is set to "relax".
193
- # Otherwise warmup all token sizes to avoid JIT compilation in hotpath
194
- if envs.VLLM_DEEP_GEMM_WARMUP == "relax":
195
- m_values = _generate_optimal_warmup_m_values(max_tokens, n, device)
196
- desc = f"DeepGemm(fp8_gemm_nt) warmup (W={w.size()}) [relaxed]"
197
- else:
198
- assert envs.VLLM_DEEP_GEMM_WARMUP == "full", (
199
- "Expected "
200
- 'VLLM_DEEP_GEMM_WARMUP env to be set to "full" but got '
201
- f"{envs.VLLM_DEEP_GEMM_WARMUP}"
202
- )
203
- m_values = list(range(1, max_tokens + 1))
204
- desc = f"DeepGemm(fp8_gemm_nt) warmup (W={w.size()}) [all tokens]"
205
-
206
- pbar = tqdm(total=len(m_values), desc=desc)
215
+ m_values = _get_fp8_gemm_nt_m_values(w, max_tokens)
216
217
for num_tokens in m_values:
218
fp8_gemm_nt(
219
(a1q[:num_tokens], a1q_scales[:num_tokens]), (w, ws), out[:num_tokens]
220
)
212
- pbar.update(1)
221
+ if pbar is not None:
222
+ pbar.update(1)
223
224
FP8_GEMM_NT_WARMUP_CACHE.add(w.size())
225
@@ -217,20 +227,12 @@ def _deepgemm_fp8_gemm_nt_warmup(w: torch.Tensor, ws: torch.Tensor, max_tokens:
227
GROUPED_FP8_GEMM_NT_CONTIGUOUS_WARMUP_CACHE: set[torch.Size] = set()
228
229
220
-def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(
230
+def _get_grouped_gemm_params(
231
w1: torch.Tensor,
232
w2: torch.Tensor,
223
- w1_scale: torch.Tensor,
224
- w2_scale: torch.Tensor,
233
num_topk: int,
234
max_tokens: int,
227
-):
228
- if (
229
- w1.size() in GROUPED_FP8_GEMM_NT_CONTIGUOUS_WARMUP_CACHE
230
- and w2.size() in GROUPED_FP8_GEMM_NT_CONTIGUOUS_WARMUP_CACHE
231
- ):
232
- return
233
-
235
+) -> tuple[int, int, torch.Tensor]:
236
assert w1.size(0) == w2.size(0), "w1 and w2 must have the same number of experts"
237
238
block_m = get_mk_alignment_for_contiguous_layout()[0]
@@ -253,6 +255,27 @@ def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(
255
)
256
expert_ids = torch.repeat_interleave(expert_ids_block, block_m, dim=0)
257
258
+ return MAX_M, block_m, expert_ids
259
+
260
+
261
+def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(
262
+ w1: torch.Tensor,
263
+ w2: torch.Tensor,
264
+ w1_scale: torch.Tensor,
265
+ w2_scale: torch.Tensor,
266
+ num_topk: int,
267
+ max_tokens: int,
268
+ pbar: tqdm | None = None,
269
+):
270
+ if (
271
+ w1.size() in GROUPED_FP8_GEMM_NT_CONTIGUOUS_WARMUP_CACHE
272
+ and w2.size() in GROUPED_FP8_GEMM_NT_CONTIGUOUS_WARMUP_CACHE
273
+ ):
274
+ return
275
+
276
+ MAX_M, block_m, expert_ids = _get_grouped_gemm_params(w1, w2, num_topk, max_tokens)
277
+ device = w1.device
278
+
279
def _warmup(w: torch.Tensor, w_scale: torch.Tensor):
280
_, n, k = w.size()
281
a1q = torch.empty((MAX_M, k), device=device, dtype=torch.float8_e4m3fn)
@@ -261,15 +284,8 @@ def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(
284
)
285
out = torch.empty((MAX_M, n), device=device, dtype=torch.bfloat16)
286
264
- # Generate M values in block_m increments (already optimized for MoE)
287
m_values = list(range(block_m, MAX_M + 1, block_m))
288
267
- pbar = tqdm(
268
- total=len(m_values),
269
- desc=f"DeepGemm(m_grouped_fp8_gemm_nt_contiguous) warmup (W={w.size()}) "
270
- f"[{len(m_values)} values, block_m={block_m}]",
271
- )
272
-
289
for num_tokens in m_values:
290
m_grouped_fp8_gemm_nt_contiguous(
291
(a1q[:num_tokens], a1q_scales[:num_tokens]),
@@ -277,7 +293,8 @@ def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(
293
out[:num_tokens],
294
expert_ids[:num_tokens],
295
)
280
- pbar.update(1)
296
+ if pbar is not None:
297
+ pbar.update(1)
298
299
for w, ws in [(w1, w1_scale), (w2, w2_scale)]:
300
if w.size() not in GROUPED_FP8_GEMM_NT_CONTIGUOUS_WARMUP_CACHE:
@@ -285,16 +302,18 @@ def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(
302
GROUPED_FP8_GEMM_NT_CONTIGUOUS_WARMUP_CACHE.add(w.size())
303
304
288
-def deepgemm_fp8_gemm_nt_warmup(model: torch.nn.Module, max_tokens: int):
305
+def deepgemm_fp8_gemm_nt_warmup(
306
+ model: torch.nn.Module, max_tokens: int, pbar: tqdm | None = None
307
+):
308
dg_modules = [m for m in model.modules() if _fp8_linear_may_use_deep_gemm(m)]
309
310
for dgm in dg_modules:
311
w, ws, _ = _extract_data_from_linear_base_module(dgm)
293
- _deepgemm_fp8_gemm_nt_warmup(w=w, ws=ws, max_tokens=max_tokens)
312
+ _deepgemm_fp8_gemm_nt_warmup(w=w, ws=ws, max_tokens=max_tokens, pbar=pbar)
313
314
315
def deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(
297
- model: torch.nn.Module, max_tokens: int
316
+ model: torch.nn.Module, max_tokens: int, pbar: tqdm | None = None
317
):
318
dg_modules = [
319
m for m in model.modules() if _fused_moe_grouped_gemm_may_use_deep_gemm(m)
@@ -305,10 +324,48 @@ def deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(
324
dgm
325
)
326
_deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(
308
- w13, w2, w13_scale, w2_scale, num_topk, max_tokens
327
+ w13, w2, w13_scale, w2_scale, num_topk, max_tokens, pbar=pbar
328
)
329
330
331
+def _count_warmup_iterations(model: torch.nn.Module, max_tokens: int) -> int:
332
+ seen_fp8_sizes: set[torch.Size] = set(FP8_GEMM_NT_WARMUP_CACHE)
333
+ seen_grouped_sizes: set[torch.Size] = set(
334
+ GROUPED_FP8_GEMM_NT_CONTIGUOUS_WARMUP_CACHE
335
+ )
336
+
337
+ total = 0
338
+ for m in model.modules():
339
+ if _fp8_linear_may_use_deep_gemm(m):
340
+ w, _, _ = _extract_data_from_linear_base_module(m)
341
+ if w.size() not in seen_fp8_sizes:
342
+ total += len(_get_fp8_gemm_nt_m_values(w, max_tokens))
343
+ seen_fp8_sizes.add(w.size())
344
+ elif _fused_moe_grouped_gemm_may_use_deep_gemm(m):
345
+ w13, _, w2, _, num_topk = _extract_data_from_fused_moe_module(m)
346
+ if w13.size() in seen_grouped_sizes and w2.size() in seen_grouped_sizes:
347
+ continue
348
+ MAX_M, block_m, _ = _get_grouped_gemm_params(w13, w2, num_topk, max_tokens)
349
+ n_values = (MAX_M - block_m) // block_m + 1
350
+ if w13.size() not in seen_grouped_sizes:
351
+ total += n_values
352
+ seen_grouped_sizes.add(w13.size())
353
+ if w2.size() not in seen_grouped_sizes:
354
+ total += n_values
355
+ seen_grouped_sizes.add(w2.size())
356
+ return total
357
+
358
+
359
def deep_gemm_warmup(model: torch.nn.Module, max_tokens: int):
313
- deepgemm_fp8_gemm_nt_warmup(model, max_tokens)
314
- deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(model, max_tokens)
360
+ total = _count_warmup_iterations(model, max_tokens)
361
+ if total == 0:
362
+ return
363
+
364
+ # Only show progress bar on rank 0 to avoid cluttered output
365
+ if is_global_first_rank():
366
+ with tqdm(total=total, desc="DeepGEMM warmup") as pbar:
367
+ deepgemm_fp8_gemm_nt_warmup(model, max_tokens, pbar)
368
+ deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(model, max_tokens, pbar)
369
+ else:
370
+ deepgemm_fp8_gemm_nt_warmup(model, max_tokens, None)
371
+ deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(model, max_tokens, None)