[Perf] Reuse MoE workspace for DeepGEMM warmup (#57268)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

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):