xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-13 21:55:32 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fix benchmark_moe.py tuning for CUDA devices (#14164)

Browse Source
This commit is contained in:
Michael Goin 2025-03-04 00:11:03 -05:00 committed by GitHub
parent 66233af7b6
commit f78c0be80a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
benchmarks/kernels/benchmark_moe.py
View File

@ -2,6 +2,7 @@
import argparse import argparse
import time import time
from contextlib import nullcontext
from datetime import datetime from datetime import datetime
from itertools import product from itertools import product
from typing import Any, TypedDict from typing import Any, TypedDict
@ -412,7 +413,8 @@ class BenchmarkWorker:
hidden_size, search_space, hidden_size, search_space,
is_fp16, topk) is_fp16, topk)
with torch.cuda.device(self.device_id): with torch.cuda.device(self.device_id) if current_platform.is_rocm(
) else nullcontext():
for config in tqdm(search_space): for config in tqdm(search_space):
try: try:
kernel_time = benchmark_config( kernel_time = benchmark_config(