Fix device capability

Signed-off-by: ilmarkov <markovilya197@gmail.com>

Author: ilmarkov

vllm/distributed/device_communicators/all_reduce_utils.py

@@ -18,7 +18,6 @@
 import vllm.envs as envs
 from vllm.distributed.device_communicators.cuda_wrapper import CudaRTLibrary
 from vllm.logger import init_logger
-from vllm.platforms import DeviceCapability
 from vllm.utils import (cuda_device_count_stateless,
                         update_environment_variables)
 
@@ -29,13 +28,13 @@
 # Max size for each world size in case symmetric memory is available
 # For different SM architectures
 CUSTOM_ALL_REDUCE_MAX_SIZES = {
-    DeviceCapability(9, 0): {
+    "9.0": {
         2: 64 * MiB,  # 64 MB
         4: 32 * MiB,  # 32 MB
         6: MiB // 2,  # 512 KB
         8: MiB // 4,  # 256 KB
     },
-    DeviceCapability(10, 0): {
+    "10.0": {
         2: 2 * MiB,  # 2 MB
         4: 2 * MiB,  # 2 MB
         6: 2 * MiB,  # 2 MB
@@ -44,13 +43,13 @@
 }
 
 SYMM_MEM_ALL_REDUCE_MAX_SIZES = {
-    DeviceCapability(9, 0): {
+    "9.0": {
         2: 64 * MiB,  # 64 MB
         4: 32 * MiB,  # 32 MB
         6: 64 * MiB,  # 64 MB
         8: 64 * MiB,  # 64 MB
     },
-    DeviceCapability(10, 0): {
+    "10.0": {
         2: 8 * MiB,  # 8 MB
         4: 32 * MiB,  # 32 MB
         6: 128 * MiB,  # 128 MB

vllm/distributed/device_communicators/custom_all_reduce.py

@@ -109,7 +109,8 @@ def __init__(self,
         # now `device` is a `torch.device` object
         assert isinstance(device, torch.device)
         self.device = device
-        device_capability = current_platform.get_device_capability()
+        device_capability = current_platform.get_device_capability(
+        ).as_version_str()
         if (current_platform.is_cuda() and envs.VLLM_ALLREDUCE_USE_SYMM_MEM
                 and device_capability in CUSTOM_ALL_REDUCE_MAX_SIZES):
             max_size = min(

vllm/distributed/device_communicators/symm_mem.py

@@ -43,7 +43,8 @@ def __init__(self, group: ProcessGroup, device: Union[int, str,
         self.device = device
         self.group = group
         self.world_size = dist.get_world_size(self.group)
-        device_capability = current_platform.get_device_capability()
+        device_capability = current_platform.get_device_capability(
+        ).as_version_str()
 
         if device_capability not in SYMM_MEM_ALL_REDUCE_MAX_SIZES:
             logger.warning(