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, 06:00 UTC
Add cache to cuda get_device_capability (#19436)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Jun 11, 2025 at 05:37 UTC
7484e1fce2fc2d6b4851d70473b5ffbf8477a57a
1 file changed
+3
-1
vllm/platforms/cuda.py
+3
-1
@@ -6,7 +6,7 @@ pynvml. However, it should not initialize cuda context.
6
7
import os
8
from datetime import timedelta
9
-from functools import wraps
9
+from functools import cache, wraps
10
from typing import TYPE_CHECKING, Callable, Optional, TypeVar, Union
11
12
import torch
@@ -389,6 +389,7 @@ class CudaPlatformBase(Platform):
389
class NvmlCudaPlatform(CudaPlatformBase):
390
391
@classmethod
392
+ @cache
393
@with_nvml_context
394
def get_device_capability(cls,
395
device_id: int = 0
@@ -486,6 +487,7 @@ class NvmlCudaPlatform(CudaPlatformBase):
487
class NonNvmlCudaPlatform(CudaPlatformBase):
488
489
@classmethod
490
+ @cache
491
def get_device_capability(cls, device_id: int = 0) -> DeviceCapability:
492
major, minor = torch.cuda.get_device_capability(device_id)
493
return DeviceCapability(major=major, minor=minor)