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 7, 05:15 UTC
Add platform method to enable custom collective ops registration (#34760)
Signed-off-by: Naina Kuruballi Mahesh <nainakm@meta.com>
nkm-meta committed
Mar 4, 2026 at 16:50 UTC
792cbd64ca1ad7b2b3bc927f1a11cf2532f624da
4 files changed
+22
-3
vllm/distributed/parallel_state.py
+3
-1
@@ -385,8 +385,10 @@ class GroupCoordinator:
385
self.cpu_group, 1 << 22, 6
386
)
387
388
+ # TODO(#35915): Remove is_tpu() check once tpu_inference
389
+ # overrides use_custom_op_collectives() to return True.
390
self.use_custom_op_call = (
389
- current_platform.is_cuda_alike() or current_platform.is_tpu()
391
+ current_platform.is_tpu() or current_platform.use_custom_op_collectives()
392
)
393
394
self.use_cpu_custom_send_recv = current_platform.is_cpu() and hasattr(
vllm/platforms/cuda.py
+5
-1
@@ -574,9 +574,13 @@ class CudaPlatformBase(Platform):
574
return True
575
576
@classmethod
577
- def num_compute_units(cls, device_id=0):
577
+ def num_compute_units(cls, device_id: int = 0) -> int:
578
return torch.cuda.get_device_properties(device_id).multi_processor_count
579
580
+ @classmethod
581
+ def use_custom_op_collectives(cls) -> bool:
582
+ return True
583
+
584
585
# NVML utils
586
# Note that NVML is not affected by `CUDA_VISIBLE_DEVICES`,
vllm/platforms/interface.py
+9
@@ -654,6 +654,15 @@ class Platform:
654
"""
655
return False
656
657
+ @classmethod
658
+ def use_custom_op_collectives(cls) -> bool:
659
+ """
660
+ Whether this platform should use torch.ops.vllm.* custom ops for collectives.
661
+
662
+ Returns False by default - platforms must explicitly opt-in.
663
+ """
664
+ return False
665
+
666
@classmethod
667
def use_sync_weight_loader(cls) -> bool:
668
"""
vllm/platforms/rocm.py
+5
-1
@@ -820,5 +820,9 @@ class RocmPlatform(Platform):
820
return True
821
822
@classmethod
823
- def num_compute_units(cls, device_id=0):
823
+ def num_compute_units(cls, device_id: int = 0) -> int:
824
return torch.cuda.get_device_properties(device_id).multi_processor_count
825
+
826
+ @classmethod
827
+ def use_custom_op_collectives(cls) -> bool:
828
+ return True