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 10, 16:30 UTC
[Bugfix] Add init_workspace_manager to moe kernel benchmarks (#31042)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Jan 5, 2026 at 22:14 UTC
e1cd7a5faffd188cd204f7b54eea6cb35f787ee9
3 files changed
+15
benchmarks/kernels/benchmark_cutlass_moe_fp8.py
+5
@@ -15,6 +15,7 @@ from vllm.model_executor.layers.fused_moe.cutlass_moe import cutlass_moe_fp8
15
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk
16
from vllm.platforms import current_platform
17
from vllm.utils.argparse_utils import FlexibleArgumentParser
18
+from vllm.v1.worker.workspace import init_workspace_manager
19
20
# Weight shapes for different models: [num_experts, topk, hidden_size,
21
# intermediate_size]
@@ -297,6 +298,10 @@ def bench_run(
298
299
300
def main(args):
301
+ # Initialize workspace manager (required for CUTLASS MoE kernels)
302
+ device = torch.device("cuda:0")
303
+ init_workspace_manager(device)
304
+
305
print("Benchmarking models:")
306
for i, model in enumerate(args.models):
307
print(f"[{i}] {model}")
benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py
renamed
+5
@@ -21,6 +21,7 @@ from vllm.model_executor.layers.fused_moe.cutlass_moe import cutlass_moe_fp4
21
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk
22
from vllm.scalar_type import scalar_types
23
from vllm.utils.argparse_utils import FlexibleArgumentParser
24
+from vllm.v1.worker.workspace import init_workspace_manager
25
26
WEIGHT_SHAPES_MOE = {
27
"nvidia/DeepSeek-R1-FP4": [
@@ -441,6 +442,10 @@ def bench_run(
442
443
444
def main(args):
445
+ # Initialize workspace manager (required for CUTLASS MoE kernels)
446
+ device = torch.device("cuda:0")
447
+ init_workspace_manager(device)
448
+
449
print("Benchmarking models:")
450
for i, model in enumerate(args.models):
451
print(f"[{i}] {model}")
benchmarks/kernels/benchmark_grouped_gemm_cutlass.py
+5
@@ -14,6 +14,7 @@ from vllm.model_executor.layers.fused_moe.fused_moe import (
14
fused_topk,
15
)
16
from vllm.utils.argparse_utils import FlexibleArgumentParser
17
+from vllm.v1.worker.workspace import init_workspace_manager
18
19
DEFAULT_MODELS = [
20
"mistralai/Mixtral-8x7B-Instruct-v0.1",
@@ -364,6 +365,10 @@ def bench_run(
365
366
367
def main(args):
368
+ # Initialize workspace manager (required for CUTLASS MoE kernels)
369
+ device = torch.device("cuda:0")
370
+ init_workspace_manager(device)
371
+
372
print("Benchmarking models:")
373
for i, model in enumerate(args.models):
374
print(f"[{i}] {model}")