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

Add cache to cuda get_device_capability (#19436)

Browse Source 
Signed-off-by: mgoin <mgoin64@gmail.com>
This commit is contained in:
Michael Goin 2025-06-11 05:37:05 -04:00 committed by GitHub
parent a2142f0196
commit 7484e1fce2
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
vllm/platforms/cuda.py
View File

@ -6,7 +6,7 @@ pynvml. However, it should not initialize cuda context.
import os import os
from datetime import timedelta from datetime import timedelta
from functools import wraps from functools import cache, wraps
from typing import TYPE_CHECKING, Callable, Optional, TypeVar, Union from typing import TYPE_CHECKING, Callable, Optional, TypeVar, Union
import torch import torch
@ -389,6 +389,7 @@ class CudaPlatformBase(Platform):
class NvmlCudaPlatform(CudaPlatformBase): class NvmlCudaPlatform(CudaPlatformBase):
@classmethod @classmethod
@cache
@with_nvml_context @with_nvml_context
def get_device_capability(cls, def get_device_capability(cls,
device_id: int = 0 device_id: int = 0
@ -486,6 +487,7 @@ class NvmlCudaPlatform(CudaPlatformBase):
class NonNvmlCudaPlatform(CudaPlatformBase): class NonNvmlCudaPlatform(CudaPlatformBase):
@classmethod @classmethod
@cache
def get_device_capability(cls, device_id: int = 0) -> DeviceCapability: def get_device_capability(cls, device_id: int = 0) -> DeviceCapability:
major, minor = torch.cuda.get_device_capability(device_id) major, minor = torch.cuda.get_device_capability(device_id)
return DeviceCapability(major=major, minor=minor) return DeviceCapability(major=major, minor=minor)