vllm-project / vllm-project/aibrix
[Feat] Support TensorRT-LLM generation-first (parallel) P/D scheduling
- Dominant language
- Go
- Stars
- 5.1k
- Forks
- 694
- Avg merge
- 1d 19h
- Merged PRs (30d)
- 98
Description
### 🚀 Feature Description and Motivation
low priority since context-first mode has been supported
The pd router currently only supports context-first scheduling (sequential: prefill → wait for response → decode).
TRT-LLM also supports generation-first mode where the proxy sends context_only and generation_only requests in parallel — the GEN worker enters `DISAGG_GENERATION_INIT` state, uses ctx_info_endpoint + disagg_request_id to pull KV cache directly from the CTX worker while it's still computing, reducing end-to-end latency.
Implementation requires: caching `/server_info` from CTX workers at router init (returns ctx_info_endpoint and ctx_dp_rank), dispatching both requests concurrently via goroutines, injecting the cached server info into the GEN request's disaggregated_params, and setting schedule_style:
> Note that ctx_dp_rank is the attention data-parallel rank — relevant when CTX workers use data parallelism (TP x DP). For single-instance deployments it's 0, but for multi-instance setups the router needs to track which DP rank each CTX worker exposes and match correctly when routing.
### Use Case
Support generation-first strategy
### Proposed Solution
_No response_
Contributor guide
Research direction
Start by locating the pd router's context-first request flow and its /server_info handling. Trace how CTX and GEN requests are routed, then define completion as concurrent dispatch, cached ctx_info_endpoint and ctx_dp_rank propagation, correct disaggregated_params, and generation-first scheduling with multi-DP routing support.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- go
- Domain
- ai-infra-agents, distributed-systems
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 30/100