From bd2d7ae7c197ba995800bb03a23b4d79c38627a4 Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Thu, 17 Sep 2026 03:01:19 -0400 Subject: [PATCH] [Perf] Reuse MoE workspace for DeepGEMM warmup (#57268) Signed-off-by: Lucas Wilkinson --- vllm/model_executor/warmup/deep_gemm_warmup.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/vllm/model_executor/warmup/deep_gemm_warmup.py b/vllm/model_executor/warmup/deep_gemm_warmup.py index 6aa5d391e0..1266f4ccb2 100644 --- a/vllm/model_executor/warmup/deep_gemm_warmup.py +++ b/vllm/model_executor/warmup/deep_gemm_warmup.py @@ -30,6 +30,7 @@ from vllm.utils.deep_gemm import ( ) from vllm.utils.math_utils import cdiv from vllm.utils.platform_utils import num_compute_units +from vllm.v1.worker.workspace import current_workspace_manager def _generate_optimal_warmup_m_values( @@ -275,15 +276,15 @@ def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup( ) if not warmup_cases: return - device = w1.device def _warmup(w: torch.Tensor, w_scale: torch.Tensor): _, n, k = w.size() - a1q = torch.empty((MAX_M, k), device=device, dtype=torch.float8_e4m3fn) - a1q_scales = torch.zeros( - (MAX_M, k // block_m), device=device, dtype=torch.float32 + a1q, a1q_scales, out = current_workspace_manager().get_simultaneous( + ((MAX_M, k), torch.float8_e4m3fn), + ((MAX_M, k // block_m), torch.float32), + ((MAX_M, n), torch.bfloat16), ) - out = torch.empty((MAX_M, n), device=device, dtype=torch.bfloat16) + a1q_scales.zero_() for num_tokens, align_used, expert_ids in warmup_cases: with mk_alignment_scope(align_used):