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 15, 21:00 UTC
Fix benchmark_moe.py tuning for CUDA devices (#14164)
Michael Goin committed
Mar 4, 2025 at 00:11 UTC
f78c0be80a8341167a5ebf20ce4eb62421a351a6
1 file changed
+3
-1
benchmarks/kernels/benchmark_moe.py
+3
-1
@@ -2,6 +2,7 @@
2
3
import argparse
4
import time
5
+from contextlib import nullcontext
6
from datetime import datetime
7
from itertools import product
8
from typing import Any, TypedDict
@@ -412,7 +413,8 @@ class BenchmarkWorker:
413
hidden_size, search_space,
414
is_fp16, topk)
415
415
- with torch.cuda.device(self.device_id):
416
+ with torch.cuda.device(self.device_id) if current_platform.is_rocm(
417
+ ) else nullcontext():
418
for config in tqdm(search_space):
419
try:
420
kernel_time = benchmark_config(