@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / c8dcc15921

Fork 0

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, 07:15 UTC

Allow AsyncLLMEngine.generate to target a specific DP rank (#19102)

Signed-off-by: Jon Swenson <jmswen@gmail.com>

jmswen committed Jun 4, 2025 at 08:26 UTC c8dcc159214a20650451dcd64b226f56671763f1

10 files changed +97 -5

examples/online_serving/multi_instance_data_parallel.py new

+58

@@ -0,0 +1,58 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +import asyncio

3 +from typing import Optional

4 +

5 +from vllm.engine.arg_utils import AsyncEngineArgs

6 +from vllm.engine.async_llm_engine import AsyncLLMEngine

7 +from vllm.outputs import RequestOutput

8 +from vllm.sampling_params import SamplingParams

9 +

10 +"""

11 +To run this example, run the following commands simultaneously with

12 +different CUDA_VISIBLE_DEVICES:

13 + python examples/online_serving/multi_instance_data_parallel.py

14 +

15 + vllm serve ibm-research/PowerMoE-3b -dp 2 -dpr 1 \

16 + --data-parallel-address 127.0.0.1 --data-parallel-rpc-port 62300 \

17 + --data-parallel-size-local 1 --enforce-eager --headless

18 +

19 +Once both instances have completed the handshake, this example will

20 +send a request to the instance with DP rank 1.

21 +"""

22 +

23 +

24 +async def main():

25 + engine_args = AsyncEngineArgs(

26 + model="ibm-research/PowerMoE-3b",

27 + data_parallel_size=2,

28 + dtype="auto",

29 + max_model_len=2048,

30 + data_parallel_address="127.0.0.1",

31 + data_parallel_rpc_port=62300,

32 + data_parallel_size_local=1,

33 + enforce_eager=True,

34 + )

35 +

36 + engine_client = AsyncLLMEngine.from_engine_args(engine_args)

37 +

38 + sampling_params = SamplingParams(

39 + temperature=0.7,

40 + top_p=0.9,

41 + max_tokens=100,

42 + )

43 +

44 + prompt = "Who won the 2004 World Series?"

45 + final_output: Optional[RequestOutput] = None

46 + async for output in engine_client.generate(

47 + prompt=prompt,

48 + sampling_params=sampling_params,

49 + request_id="abcdef",

50 + data_parallel_rank=1,

51 + ):

52 + final_output = output

53 + if final_output:

54 + print(final_output.outputs[0].text)

55 +

56 +

57 +if __name__ == "__main__":

58 + asyncio.run(main())

tests/tokenization/test_detokenize.py

+2 -1

@@ -70,7 +70,8 @@ def _run_incremental_decode(tokenizer,

70 None,

71 0.0,

72 None,

73 - cache_salt=None)

73 + cache_salt=None,

74 + data_parallel_rank=None)

76 if fast is None:

77 detokenizer = IncrementalDetokenizer.from_new_request(

tests/v1/engine/test_engine_core.py

@@ -42,6 +42,7 @@ def make_request() -> EngineCoreRequest:

42 arrival_time=time.time(),

43 lora_request=None,

44 cache_salt=None,

45 + data_parallel_rank=None,

46 )

tests/v1/engine/test_engine_core_client.py

@@ -56,6 +56,7 @@ def make_request(

56 arrival_time=time.time(),

57 lora_request=None,

58 cache_salt=None,

59 + data_parallel_rank=None,

60 )

tests/v1/engine/test_output_processor.py

@@ -59,6 +59,7 @@ def test_incremental_detokenization(request_output_kind: RequestOutputKind,

59 eos_token_id=None,

60 lora_request=None,

61 cache_salt=None,

62 + data_parallel_rank=None,

63 sampling_params=SamplingParams(

64 skip_special_tokens=False,

65 spaces_between_special_tokens=False,

@@ -406,6 +407,7 @@ def test_logprobs_processor(request_output_kind: RequestOutputKind,

407 eos_token_id=None,

408 lora_request=None,

409 cache_salt=None,

410 + data_parallel_rank=None,

411 sampling_params=SamplingParams(

412 skip_special_tokens=False,

413 spaces_between_special_tokens=False,

@@ -569,6 +571,7 @@ def test_stop_token(include_stop_str_in_output: bool,

571 eos_token_id=eos_token_id,

572 lora_request=None,

573 cache_salt=None,

574 + data_parallel_rank=None,

575 sampling_params=SamplingParams(

576 skip_special_tokens=False,

577 spaces_between_special_tokens=False,

@@ -666,6 +669,7 @@ def test_stop_string(include_stop_str_in_output: bool,

669 eos_token_id=None,

670 lora_request=None,

671 cache_salt=None,

672 + data_parallel_rank=None,

673 sampling_params=SamplingParams(

674 skip_special_tokens=False,

675 spaces_between_special_tokens=False,

@@ -780,6 +784,7 @@ def test_iteration_stats(dummy_test_vectors):

784 eos_token_id=None,

785 lora_request=None,

786 cache_salt=None,

787 + data_parallel_rank=None,

788 sampling_params=SamplingParams(),

789 ) for idx, prompt_tokens in enumerate(dummy_test_vectors.prompt_tokens)

790 ]

vllm/engine/async_llm_engine.py

+11 -1

@@ -442,6 +442,7 @@ class _AsyncLLMEngine(LLMEngine):

442 trace_headers: Optional[Mapping[str, str]] = None,

443 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

444 priority: int = 0,

445 + data_parallel_rank: Optional[int] = None,

446 ) -> None:

447 ...

448

@@ -456,6 +457,7 @@ class _AsyncLLMEngine(LLMEngine):

457 trace_headers: Optional[Mapping[str, str]] = None,

458 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

459 priority: int = 0,

460 + data_parallel_rank: Optional[int] = None,

461 ) -> None:

462 ...

463

@@ -473,6 +475,7 @@ class _AsyncLLMEngine(LLMEngine):

475 trace_headers: Optional[Mapping[str, str]] = None,

476 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

477 priority: int = 0,

478 + data_parallel_rank: Optional[int] = None,

479 *,

480 inputs: Optional[PromptType] = None, # DEPRECATED

481 ) -> None:

@@ -902,6 +905,7 @@ class AsyncLLMEngine(EngineClient):

905 trace_headers: Optional[Mapping[str, str]] = None,

906 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

907 priority: int = 0,

908 + data_parallel_rank: Optional[int] = None,

909 ) -> Coroutine[None, None, AsyncGenerator[Union[

910 RequestOutput, PoolingRequestOutput], None]]:

911 ...

@@ -917,6 +921,7 @@ class AsyncLLMEngine(EngineClient):

921 trace_headers: Optional[Mapping[str, str]] = None,

922 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

923 priority: int = 0,

924 + data_parallel_rank: Optional[int] = None,

925 ) -> Coroutine[None, None, AsyncGenerator[Union[

926 RequestOutput, PoolingRequestOutput], None]]:

927 ...

@@ -935,6 +940,7 @@ class AsyncLLMEngine(EngineClient):

940 trace_headers: Optional[Mapping[str, str]] = None,

941 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

942 priority: int = 0,

943 + data_parallel_rank: Optional[int] = None,

944 *,

945 inputs: Optional[PromptType] = None, # DEPRECATED

946 ) -> AsyncGenerator[Union[RequestOutput, PoolingRequestOutput], None]:

@@ -967,6 +973,7 @@ class AsyncLLMEngine(EngineClient):

973 trace_headers=trace_headers,

974 prompt_adapter_request=prompt_adapter_request,

975 priority=priority,

976 + data_parallel_rank=data_parallel_rank,

977 )

978

979 return stream.generator()

@@ -980,6 +987,7 @@ class AsyncLLMEngine(EngineClient):

987 trace_headers: Optional[Mapping[str, str]] = None,

988 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

989 priority: int = 0,

990 + data_parallel_rank: Optional[int] = None,

991 ) -> AsyncGenerator[RequestOutput, None]:

992 """Generate outputs for a request.

993

@@ -999,7 +1007,8 @@ class AsyncLLMEngine(EngineClient):

1007 for generation, if any.

1008 priority: The priority of the request.

1009 Only applicable with priority scheduling.

1002 -

1010 + data_parallel_rank: The (global) data parallel rank that must

1011 + handle this request. Only applicable if DP is enabled.

1012 Yields:

1013 The output `RequestOutput` objects from the LLMEngine

1014 for the request.

@@ -1057,6 +1066,7 @@ class AsyncLLMEngine(EngineClient):

1066 trace_headers=trace_headers,

1067 prompt_adapter_request=prompt_adapter_request,

1068 priority=priority,

1069 + data_parallel_rank=data_parallel_rank,

1070 ):

1071 yield LLMEngine.validate_output(output, RequestOutput)

1072 except asyncio.CancelledError:

vllm/v1/engine/__init__.py

@@ -55,6 +55,7 @@ class EngineCoreRequest(

55 arrival_time: float

56 lora_request: Optional[LoRARequest]

57 cache_salt: Optional[str]

58 + data_parallel_rank: Optional[int]

60 # Index of the client, used to ensure outputs are sent back to the same

61 # client for this request when scaling out the front-end.

vllm/v1/engine/async_llm.py

+4 -1

@@ -229,6 +229,7 @@ class AsyncLLM(EngineClient):

229 trace_headers: Optional[Mapping[str, str]] = None,

230 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

231 priority: int = 0,

232 + data_parallel_rank: Optional[int] = None,

233 ) -> RequestOutputCollector:

234 """Add new request to the AsyncLLM."""

235

@@ -245,7 +246,7 @@ class AsyncLLM(EngineClient):

246 prompt_str, request = self.processor.process_inputs(

247 request_id, prompt, params, arrival_time, lora_request,

248 tokenization_kwargs, trace_headers, prompt_adapter_request,

248 - priority)

249 + priority, data_parallel_rank)

250

251 if params.n == 1:

252 await self._add_request(request, prompt_str, None, 0, queue)

@@ -291,6 +292,7 @@ class AsyncLLM(EngineClient):

292 trace_headers: Optional[Mapping[str, str]] = None,

293 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

294 priority: int = 0,

295 + data_parallel_rank: Optional[int] = None,

296 ) -> AsyncGenerator[RequestOutput, None]:

297 """

298 Main function called by the API server to kick off a request

@@ -321,6 +323,7 @@ class AsyncLLM(EngineClient):

323 trace_headers=trace_headers,

324 prompt_adapter_request=prompt_adapter_request,

325 priority=priority,

326 + data_parallel_rank=data_parallel_rank,

327 )

328

329 # The output_handler task pushes items into the queue.

vllm/v1/engine/core_client.py

+12 -2

@@ -982,7 +982,16 @@ class DPAsyncMPClient(AsyncMPClient):

982 resources.stats_update_task = asyncio.create_task(

983 run_engine_stats_update_task())

984

985 - def get_core_engine_for_request(self) -> CoreEngine:

985 + def get_core_engine_for_request(self,

986 + dp_rank: Optional[int] = None

987 + ) -> CoreEngine:

988 + if dp_rank is not None:

989 + # engines are already in rank order

990 + if dp_rank < 0 or dp_rank >= len(self.core_engines):

991 + raise ValueError(f"Requested DP rank {dp_rank} is out of "

992 + f"range [0, {len(self.core_engines)})")

993 + return self.core_engines[dp_rank]

994 +

995 if not self.lb_engines:

996 return self.core_engines[0]

997 # TODO use P2C alg for larger DP sizes

@@ -1018,7 +1027,8 @@ class DPAsyncMPClient(AsyncMPClient):

1027 request.current_wave = self.current_wave

1028 request.client_index = self.client_index

1029

1021 - chosen_engine = self.get_core_engine_for_request()

1030 + chosen_engine = self.get_core_engine_for_request(

1031 + request.data_parallel_rank)

1032 self.reqs_in_flight[request.request_id] = chosen_engine

1033

1034 to_await = self._send_input(EngineCoreRequestType.ADD, request,

vllm/v1/engine/processor.py

@@ -212,6 +212,7 @@ class Processor:

212 trace_headers: Optional[Mapping[str, str]] = None,

213 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

214 priority: int = 0,

215 + data_parallel_rank: Optional[int] = None,

216 ) -> tuple[Optional[str], EngineCoreRequest]:

217

218 # TODO(woosuk): Support pooling models.

@@ -328,6 +329,7 @@ class Processor:

329 arrival_time=arrival_time,

330 lora_request=lora_request,

331 cache_salt=decoder_inputs.get("cache_salt"),

332 + data_parallel_rank=data_parallel_rank,

333 )

334

335 def _validate_model_inputs(self,