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, 16:00 UTC
[Distributed] [ROCM] Fix custom allreduce enable checks (#16010)
Signed-off-by: ilmarkov <imarkov@redhat.com> Co-authored-by: ilmarkov <imarkov@redhat.com>
Ilya Markov committed
Apr 4, 2025 at 18:39 UTC
ef608c37a7cf106fcd48aee24ec44248d2a6665b
4 files changed
+21
-4
vllm/config.py
+3
-4
@@ -1619,13 +1619,12 @@ class ParallelConfig:
1619
if self.use_ray:
1620
from vllm.executor import ray_utils
1621
ray_utils.assert_ray_available()
1622
- device_capability = current_platform.get_device_capability()
1623
- if (current_platform.is_rocm() and device_capability is not None
1624
- and device_capability < (9, 4)):
1622
+
1623
+ if not current_platform.use_custom_allreduce():
1624
self.disable_custom_all_reduce = True
1625
logger.info(
1626
"Disabled the custom all-reduce kernel because it is not "
1628
- "supported on AMD GPUs older than MI300X.")
1627
+ "supported on current platform.")
1628
if self.ray_workers_use_nsight and not self.use_ray:
1629
raise ValueError("Unable to use nsight profiling unless workers "
1630
"run with Ray.")
vllm/platforms/cuda.py
+4
@@ -308,6 +308,10 @@ class CudaPlatformBase(Platform):
308
def supports_v1(cls, model_config: ModelConfig) -> bool:
309
return True
310
311
+ @classmethod
312
+ def use_custom_allreduce(cls) -> bool:
313
+ return True
314
+
315
316
# NVML utils
317
# Note that NVML is not affected by `CUDA_VISIBLE_DEVICES`,
vllm/platforms/interface.py
+7
@@ -379,6 +379,13 @@ class Platform:
379
"""
380
return False
381
382
+ @classmethod
383
+ def use_custom_allreduce(cls) -> bool:
384
+ """
385
+ Returns if custom allreduce is supported on the current platform
386
+ """
387
+ return False
388
+
389
390
class UnspecifiedPlatform(Platform):
391
_enum = PlatformEnum.UNSPECIFIED
vllm/platforms/rocm.py
+7
@@ -302,3 +302,10 @@ class RocmPlatform(Platform):
302
def supports_v1(cls, model_config: ModelConfig) -> bool:
303
# V1 support on AMD gpus is experimental
304
return True
305
+
306
+ @classmethod
307
+ def use_custom_allreduce(cls) -> bool:
308
+ # We only enable custom allreduce for MI300 series
309
+ gcn_arch = torch.cuda.get_device_properties(0).gcnArchName
310
+ supported_archs = ['gfx94']
311
+ return any(gfx in gcn_arch for gfx in supported_archs)