vllm-project / vllm-project/aibrix

[Feat] Support TensorRT-LLM generation-first (parallel) P/D scheduling

Open
#2,018 0 comments 0 reactions 0 assignees View on GitHub
area/distributed kind/feature priority/important-longterm
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.