[Feature][Frontend] Add KVTransferParams for disaggregated prefill feature

[ShangmingCai]

This is the initial PR to implement XpYd as described in the design doc, which has gone through several rounds of discussions in the #feat-prefill-disaggregation slack channel and reached a preliminary consensus.

This PR mainly includes the changes related to KVTransferParams. In the implementation of XpYd, we will use this parameter to coordinate metadata among all vllm instances (prefill nodes and decode nodes), the KVStore Server, and the global proxy. With this param, the global proxy can attach the disaggregated prefill metadata info in the CompletionRequest or ChatCompletionRequest, and all kv-transfer-related modules can access it through the model_input in function model_runner.execute_model and use these params to interact with the KVStore Server at the right place under current 1P1D implementation.

Based on this initial PR, the Mooncake team will implement the KVStoreConnector in the next PR, which utilizes KVTransferParams and supports several third-party key-value stores to act as the KVLookupBuffer such as MooncakeStore, Valkey, LMCache, etc. With KVStoreConnector, we can decouple the connections between the Prefill nodes and the Decode nodes, making the entire XpYd implementation simpler and more stable.

The metadata info required in KVTransferParams will be provided by a global proxy (or we can call it the XpYd orchestration mechanism). Therefore, a global proxy interface will be added in the next PR too. Mooncake team will implement a simple version based on round-robin. In the production stage, service providers and users can also implement corresponding global proxy strategies according to their needs.

CC list: @KuntaiDu, @youkaichao, @james0zan, @alogfans, @stmatengss, @zeroorhero, @pizhenwei

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[ShangmingCai]

@KuntaiDu Hello, this is the initial PR for the XpYd design as we discussed before the Chinese New Year. I think it is ready for review :)

[ShangmingCai]

@youkaichao I have verified that the failure of distributed-tests-4-gpus ci can be solved by #12959. However, the image created by docker-build-image does not include the latest fix. So re-triggering ci without rebase does not seem to solve the ci failure that happens to exist when this PR is created but fixed later. And ever since DCO was introduced, I found that rebase has become difficult. I wonder if there is a better solution to rebase or merge without triggering DCO, so that we can fix the failed CI that has been introduced by the stale code of the main branch that happens to exist when this PR is created.

[KuntaiDu]

Since it is List[int], It seems to require the orchestrator to run per-request tokenization. Will that limits the scalability of the orchestrator?

[ShangmingCai]

In practice, if we were only implementing the XpYd scheme, the orchestrator would only need to handle request routing between p-nodes and d-nodes, along with managing the keys for KV cache transmission. However, to ensure compatibility with prefix caching features in the future and enable the proxy to maintain a prefix tree constructed from token ids for more content-aware disaggregated prefill scheduling, this parameter is introduced to support broader KVCache transfer options and scheduling scenarios. In fact, it could be absent in the first XpYd implementation since we assume prefix_prompt_ids should essentially match the prompt_token_ids in disaggregated prefill scenarios.

[KuntaiDu]

Let me discuss with other people and see how to deal with KVTransferParams in v1.

[KuntaiDu]

Just want to learn more on the rationales behind isolating the KVTransferParams in an isolated file instead of stuffing it into one existing file.

[ShangmingCai]

Good question. Initially, I planned to add it in sampling_params.py, but since the kv transfer process is unrelated to sampling, combining them might make the code harder to understand. I referenced the pooling_params.py implementation and used a separate file for this class to improve readability. Regarding the location, while placing it in the root directory (like other params) works, moving it to the vllm/distributed/kv_transfer directory is also an option. Let me know which approach you prefer!

[mergify]

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

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