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, 13:30 UTC
Add collective_rpc to llm engine (#16999)
Signed-off-by: Yinghai Lu <yinghai@thinkingmachines.ai>
Yinghai Lu committed
Apr 24, 2025 at 13:16 UTC
fe921763212b881d9629d04c2eaab4496e136fa5
2 files changed
+29
vllm/engine/async_llm_engine.py
+18
@@ -528,6 +528,13 @@ class _AsyncLLMEngine(LLMEngine):
528
async def check_health_async(self) -> None:
529
self.model_executor.check_health()
530
531
+ async def collective_rpc_async(self,
532
+ method: str,
533
+ timeout: Optional[float] = None,
534
+ args: tuple = (),
535
+ kwargs: Optional[dict] = None):
536
+ raise NotImplementedError
537
+
538
539
async def build_guided_decoding_logits_processor_async(
540
sampling_params: SamplingParams, tokenizer: AnyTokenizer,
@@ -1236,6 +1243,17 @@ class AsyncLLMEngine(EngineClient):
1243
async def add_lora(self, lora_request: LoRARequest) -> None:
1244
self.engine.add_lora(lora_request)
1245
1246
+ async def collective_rpc(self,
1247
+ method: str,
1248
+ timeout: Optional[float] = None,
1249
+ args: tuple = (),
1250
+ kwargs: Optional[dict] = None):
1251
+ """
1252
+ Perform a collective RPC call to the given path.
1253
+ """
1254
+ return await self.engine.collective_rpc_async(method, timeout, args,
1255
+ kwargs)
1256
+
1257
1258
# TODO(v1): Remove this class proxy when V1 goes default.
1259
if envs.is_set("VLLM_USE_V1") and envs.VLLM_USE_V1:
vllm/v1/engine/async_llm.py
+11
@@ -492,6 +492,17 @@ class AsyncLLM(EngineClient):
492
"""Prevent an adapter from being evicted."""
493
return await self.engine_core.pin_lora_async(lora_id)
494
495
+ async def collective_rpc(self,
496
+ method: str,
497
+ timeout: Optional[float] = None,
498
+ args: tuple = (),
499
+ kwargs: Optional[dict] = None):
500
+ """
501
+ Perform a collective RPC call to the given path.
502
+ """
503
+ return await self.engine_core.collective_rpc_async(
504
+ method, timeout, args, kwargs)
505
+
506
@property
507
def is_running(self) -> bool:
508
# Is None before the loop is started.