[Bugfix] Add init_workspace_manager to moe kernel benchmarks (#31042)
Signed-off-by: mgoin <mgoin64@gmail.com>
This commit is contained in:
@@ -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}")
|
||||
|
||||
Reference in New Issue
Block a user