[Bugfix] Add init_workspace_manager to moe kernel benchmarks (#31042)

Signed-off-by: mgoin <mgoin64@gmail.com>
This commit is contained in:
Michael Goin
2026-01-05 22:14:33 -05:00
committed by GitHub
parent a68e703c32
commit e1cd7a5faf
3 changed files with 15 additions and 0 deletions

View File

@@ -14,6 +14,7 @@ from vllm.model_executor.layers.fused_moe.fused_moe import (
fused_topk,
)
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm.v1.worker.workspace import init_workspace_manager
DEFAULT_MODELS = [
"mistralai/Mixtral-8x7B-Instruct-v0.1",
@@ -364,6 +365,10 @@ def bench_run(
def main(args):
# Initialize workspace manager (required for CUTLASS MoE kernels)
device = torch.device("cuda:0")
init_workspace_manager(device)
print("Benchmarking models:")
for i, model in enumerate(args.models):
print(f"[{i}] {model}")