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 14, 00:30 UTC
[EP/DP][API Server] Enable DP-aware routing in OpenAI API requests (#24945)
Co-authored-by: Cong Chen <prowindy@gmail.com>
cong-meta committed
Oct 30, 2025 at 12:10 UTC
a2981c42720a34b5abf59c4c14df701f8105d4cd
4 files changed
+99
tests/entrypoints/openai/test_serving_chat.py
+76
@@ -651,3 +651,79 @@ async def test_serving_chat_did_set_correct_cache_salt(model_type):
651
await serving_chat.create_chat_completion(req)
652
engine_prompt = serving_chat._process_inputs.await_args_list[1].args[1]
653
assert engine_prompt.get("cache_salt") == "test_salt"
654
+
655
+
656
+@pytest.mark.asyncio
657
+async def test_serving_chat_data_parallel_rank_extraction():
658
+ """Test that data_parallel_rank is properly extracted from header and
659
+ passed to engine."""
660
+ mock_engine = MagicMock(spec=AsyncLLM)
661
+ mock_engine.get_tokenizer.return_value = get_tokenizer(MODEL_NAME)
662
+ mock_engine.errored = False
663
+ mock_engine.model_config = MockModelConfig()
664
+ mock_engine.processor = MagicMock()
665
+ mock_engine.io_processor = MagicMock()
666
+
667
+ # Mock the generate method to return an async generator
668
+ async def mock_generate(*args, **kwargs):
669
+ # Yield a fake RequestOutput
670
+ from vllm.outputs import CompletionOutput, RequestOutput
671
+
672
+ yield RequestOutput(
673
+ request_id="test-request",
674
+ prompt="test prompt",
675
+ prompt_token_ids=[1, 2, 3],
676
+ prompt_logprobs=None,
677
+ outputs=[
678
+ CompletionOutput(
679
+ index=0,
680
+ text="test response",
681
+ token_ids=[4, 5, 6],
682
+ cumulative_logprob=0.0,
683
+ logprobs=None,
684
+ finish_reason="stop",
685
+ stop_reason=None,
686
+ )
687
+ ],
688
+ finished=True,
689
+ )
690
+
691
+ mock_engine.generate = AsyncMock(side_effect=mock_generate)
692
+
693
+ serving_chat = _build_serving_chat(mock_engine)
694
+
695
+ # Test when data_parallel_rank is present in header
696
+ req = ChatCompletionRequest(
697
+ model=MODEL_NAME,
698
+ messages=[{"role": "user", "content": "what is 1+1?"}],
699
+ )
700
+
701
+ # Mock request with X-data-parallel-rank header
702
+ mock_raw_request = MagicMock()
703
+ mock_raw_request.headers = {"X-data-parallel-rank": "2"}
704
+ mock_raw_request.state = MagicMock()
705
+
706
+ with suppress(Exception):
707
+ await serving_chat.create_chat_completion(req, mock_raw_request)
708
+
709
+ # Verify that data_parallel_rank was passed to engine.generate
710
+ assert "data_parallel_rank" in mock_engine.generate.call_args.kwargs
711
+ assert mock_engine.generate.call_args.kwargs["data_parallel_rank"] == 2
712
+
713
+ # Test when data_parallel_rank is not present (defaults to None)
714
+ req_no_dp = ChatCompletionRequest(
715
+ model=MODEL_NAME,
716
+ messages=[{"role": "user", "content": "what is 2+2?"}],
717
+ )
718
+
719
+ # Mock request with no header
720
+ mock_raw_request_no_dp = MagicMock()
721
+ mock_raw_request_no_dp.headers = {}
722
+ mock_raw_request_no_dp.state = MagicMock()
723
+
724
+ with suppress(Exception):
725
+ await serving_chat.create_chat_completion(req_no_dp, mock_raw_request_no_dp)
726
+
727
+ # Verify that data_parallel_rank defaults to None
728
+ assert "data_parallel_rank" in mock_engine.generate.call_args.kwargs
729
+ assert mock_engine.generate.call_args.kwargs["data_parallel_rank"] is None
vllm/entrypoints/openai/serving_chat.py
+4
@@ -264,6 +264,9 @@ class OpenAIServingChat(OpenAIServing):
264
if raw_request:
265
raw_request.state.request_metadata = request_metadata
266
267
+ # Extract data_parallel_rank from header (router can inject it)
268
+ data_parallel_rank = self._get_data_parallel_rank(raw_request)
269
+
270
# Schedule the request and get the result generator.
271
generators: list[AsyncGenerator[RequestOutput, None]] = []
272
try:
@@ -331,6 +334,7 @@ class OpenAIServingChat(OpenAIServing):
334
priority=request.priority,
335
prompt_text=prompt_text,
336
tokenization_kwargs=tokenization_kwargs,
337
+ data_parallel_rank=data_parallel_rank,
338
)
339
340
generators.append(generator)
vllm/entrypoints/openai/serving_completion.py
+4
@@ -141,6 +141,9 @@ class OpenAIServingCompletion(OpenAIServing):
141
logger.exception("Error in preprocessing prompt inputs")
142
return self.create_error_response(str(e))
143
144
+ # Extract data_parallel_rank from header (router can inject it)
145
+ data_parallel_rank = self._get_data_parallel_rank(raw_request)
146
+
147
# Schedule the request and get the result generator.
148
generators: list[AsyncGenerator[RequestOutput, None]] = []
149
try:
@@ -224,6 +227,7 @@ class OpenAIServingCompletion(OpenAIServing):
227
priority=request.priority,
228
prompt_text=prompt_text,
229
tokenization_kwargs=tokenization_kwargs,
230
+ data_parallel_rank=data_parallel_rank,
231
)
232
233
generators.append(generator)
vllm/entrypoints/openai/serving_engine.py
+15
@@ -1298,6 +1298,21 @@ class OpenAIServing:
1298
1299
return raw_request.headers.get("X-Request-Id", default)
1300
1301
+ @staticmethod
1302
+ def _get_data_parallel_rank(raw_request: Request | None) -> int | None:
1303
+ """Pulls the data parallel rank from a header, if provided"""
1304
+ if raw_request is None:
1305
+ return None
1306
+
1307
+ rank_str = raw_request.headers.get("X-data-parallel-rank")
1308
+ if rank_str is None:
1309
+ return None
1310
+
1311
+ try:
1312
+ return int(rank_str)
1313
+ except ValueError:
1314
+ return None
1315
+
1316
@staticmethod
1317
def _get_decoded_token(
1318
logprob: Logprob,