[EP/DP][API Server] Enable DP-aware routing in OpenAI API requests

[Prowindy]

Signed-off-by: Cong Chen [email protected]

Purpose

This PR adds support for the data_parallel_rank parameter in vLLM's OpenAI API, enabling external routers to specify which data parallel rank
should handle a request. This enhancement improves load distribution and cache locality in multi-GPU deployments by allowing DP-aware routing
from router to engine core.

Key functionality:

• External routers can inject data_parallel_rank into API requests
• vLLM server routes requests to specific GPU ranks based on the parameter
• Maintains backward compatibility when parameter is not provided
• Enables intelligent load balancing and cache-aware routing strategies

Test Plan

1. End-to-End DP-Aware Routing Test:

   # Start vLLM server with 8 DP ranks
   vllm serve meta-llama/Meta-Llama-3.1-8B-Instruct \
     --data-parallel-size 8 --tensor-parallel-size 1 --port 8000
   
   # Start vLLM-router with DP-aware routing (selective ranks 0,2,4,6)
   vllm-router --worker-urls http://0.0.0.0:8000 \
     --dp-aware --policy round_robin --port 30000
   
   # Generate test traffic
   vllm bench serve --num-prompts 100 --port 30000 \
     --endpoint /v1/completions --max-concurrency 32
   

2. Validation Methods:
   - Added comprehensive logging throughout the request pipeline
   - Implemented rank counters to track request distribution
   - Verified GPU device assignment for each request
   - Tested with multiple routing policies (round_robin, cache_aware)

Test Result

✅ Successful DP-Aware Routing Validation:

1. Router Discovery: Successfully identified 4 DP-aware workers
   workers: ["http://0.0.0.0:8000@0", "http://0.0.0.0:8000@2", "http://0.0.0.0:8000@4", "http://0.0.0.0:8000@6"]
2. Request Distribution: Balanced traffic across selected ranks
   Current counters: {0: 2, 2: 3, 4: 3, 6: 2}
3. End-to-End Pipeline Verification:
   API_SERVER_DEBUG: Received completion request with data_parallel_rank=4
   ROUTING_DEBUG: Using data_parallel_rank=4 directly
   ENGINE_DISPATCH DEBUG: Request dispatched to engine identity 0400 (rank 4)
   GPU_DEVICE_DEBUG: Request received on GPU device, parallel_config.rank=4
4. Multi-Process Validation: Confirmed different engine processes correctly received assigned requests:
   - EngineCore_0 pid=296341 → Rank 0 requests
   - EngineCore_2 pid=296343 → Rank 2 requests
   - EngineCore_4 pid=296345 → Rank 4 requests
   - EngineCore_6 pid=296347 → Rank 6 requests

[gemini-code-assist]

Code Review

This pull request effectively enables DP-aware routing by introducing the data_parallel_rank parameter in OpenAI API requests. The changes are well-implemented across the protocol, serving layer, and client logic, and include a new test for validation. The core functionality appears solid. My review identified one critical issue in a new Kubernetes manifest file which seems to be a copy-paste error, making the file invalid.

[gemini-code-assist]

The content of this YAML file is duplicated from line 98 onwards, which makes the file invalid and will cause parsing errors. This appears to be a copy-paste error that needs to be fixed by removing the duplicated section.

  restartPolicy: Never

[robertgshaw2-redhat]

I think this should be implemented as a header rather than in the body similar to how external processes can add request id (

vllm/vllm/entrypoints/openai/serving_engine.py

Line 962 in 4e5affe

return raw_request.headers.get("X-Request-Id", default)

)

So this could be something like X-data-parallel-rank

This is more canonical with how extra arguments are added to a standard api

[Prowindy]

@robertgshaw2-redhat The current approach is consistent with how SGLang handles DP-aware requests. With this change to the vLLM server, the sgl-router fork can send requests directly to the vLLM server without requiring any API modifications.

That said, your suggestion is valid—P/D disaggregation also embeds the prefill/decode address in the X-Request-Id. I’ll make a few adjustments to incorporate your recommendation.

[robertgshaw2-redhat]

Is this needed? I think its usually not great to allow external users to see this type of server info

I think that instead we should make sure the server is in DEV mode to expose this

[Prowindy]

This design is inherited from sglang API which allows its router to fetch entire server config.

We have a couple of alternative impl options:

1. Pass data_parallel_size as an input flag to the router. This would require all vllm servers to use the same DP size (or we could further separate it into prefill and decode configs).
2. Create a /get_data_parallel_size endpoint, allowing us to expose only this specific configuration to callers.

@robertgshaw2-redhat Which of these approaches would you recommend?

[njhill]

I agree it's good to keep the discovery part separate (can figure out in future PR). The header changes in this one look good to me now.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @Prowindy.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

Documentation preview: https://vllm--24945.org.readthedocs.build/en/24945/

[njhill]

Thanks @Prowindy! LGTM

Looks like there is a commit which needs DCO sign-off, see https://github.com/vllm-project/vllm/pull/24945/checks?check_run_id=53896850372, and some pre-commit errors.

Also I wonder if we could add something to the docs to cover this. Perhaps add a note to the external router section of the data parallel page: https://docs.vllm.ai/en/latest/serving/data_parallel_deployment.html#external-load-balancing

[robertgshaw2-redhat]

an attack vector:

• user specifies X-data-parallel-rank that is > DP size or is negative
• user specifies X-data-parallel-rank when the server is not using DP

What happens? We need to be careful as this is untrusted user input that could case DOS issues.

Perhaps we should only parse this field if the server operator ops in and log a warning about it. We also need to santize the inputs

[Prowindy]

The engine performs a sanity check on the input DP rank range before utilizing it:

vllm/vllm/v1/engine/processor.py

Lines 356 to 363 in 141e6a0

	data_parallel_size = self.vllm_config.parallel_config.data_parallel_size
	if data_parallel_rank is not None and not (
	0 <= data_parallel_rank < data_parallel_size
	):
	raise ValueError(
	f"data_parallel_rank {data_parallel_rank} "
	f"is out of range [0, {data_parallel_size})."
	)

Do we want to enhance safety by clamping, or raising exceptions during parsing?

[Prowindy]

Invalid-dp-tests.txt

@robertgshaw2-redhat I tested the vLLM patch by requesting with invalid DP ranks, including negative or large numbers, and non-integer characters. In these cases, the system fallback to a regular request without DP rank info. Would this be our expected behavior?