[Perf] Reuse MoE workspace for DeepGEMM warmup (#57268)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Sep 17, 2026 at 03:01 UTC
bd2d7ae7c197ba995800bb03a23b4d79c38627a4
1 file changed
+6
-5
vllm/model_executor/warmup/deep_gemm_warmup.py
+6
-5
@@ -30,6 +30,7 @@ from vllm.utils.deep_gemm import (
30
)
31
from vllm.utils.math_utils import cdiv
32
from vllm.utils.platform_utils import num_compute_units
33
+from vllm.v1.worker.workspace import current_workspace_manager
34
35
36
def _generate_optimal_warmup_m_values(
@@ -275,15 +276,15 @@ def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(
276
)
277
if not warmup_cases:
278
return
278
- device = w1.device
279
280
def _warmup(w: torch.Tensor, w_scale: torch.Tensor):
281
_, n, k = w.size()
282
- a1q = torch.empty((MAX_M, k), device=device, dtype=torch.float8_e4m3fn)
283
- a1q_scales = torch.zeros(
284
- (MAX_M, k // block_m), device=device, dtype=torch.float32
282
+ a1q, a1q_scales, out = current_workspace_manager().get_simultaneous(
283
+ ((MAX_M, k), torch.float8_e4m3fn),
284
+ ((MAX_M, k // block_m), torch.float32),
285
+ ((MAX_M, n), torch.bfloat16),
286
)
286
- out = torch.empty((MAX_M, n), device=device, dtype=torch.bfloat16)
287
+ a1q_scales.zero_()
288
289
for num_tokens, align_used, expert_ids in warmup_cases:
290
with mk_alignment_scope(align_used):