[KV Connector]: Support KV push from Prefill to Decode node using Nixl KV Connector

[snadampal]

RFC #36923

Implemented KV push feature where Prefill node pushes
its KV blocks to Decode node as soon as the model executor
completes the forward pass and finishes request.
The implementation supports heterogeneous TP and
heterogeneous block sizes between P and D nodes

And it covers both the scenarios:
Scenario 1: D registers blocks with P before P finishes generating KV
Scenario 2: P has the KV ready before D registers

Purpose

To improve TTFT for P-D disaggregated inference deployment.

Test Plan

Manually tested Inference on P-D disaggregated setup

Test Result

Tested 1p1d configuration for pushing KV from prefill to decode node and validated accuracy of the results.

Benchmarking results

I ran vllm bench serve with sonnet dataset for different input and output token lengths. KV push mode (this PR) showed 1.2x - 3.0x improvements over pull mode across different input and out lengths. Following are the performance numbers measured on AWS P5en instance and similar improvements are observed on AWS Trn2 instances as well where the feature was originally developed.

Mode	Prefill/Decode TP	Input Len	Output Len	QPS	Mean TTFT (ms)	Median TTFT (ms)	P99 TTFT (ms)	Mean TPOT (ms)	Req/s
pull	4	512	64	4	110.94	85.07	740.03	12.51	3.86
push	4	512	64	4	74.46	77.13	100.24	12.50	3.87
pull	4	512	128	4	301.92	89.22	1125.98	12.66	3.63
push	4	512	128	4	93.77	79.95	302.37	12.59	3.75
pull	4	512	128	8	2948.11	2682.32	7258.96	12.70	4.65
push	4	512	128	8	1212.48	1275.05	2826.72	12.67	5.92
pull	4	1024	128	4	330.94	91.26	1163.70	12.74	3.62
push	4	1024	128	4	273.83	84.10	1096.88	12.79	3.67
pull	4	2048	128	4	352.22	98.53	1225.20	12.88	3.61
push	4	2048	128	4	239.65	88.23	1042.32	12.85	3.73
pull	8	512	64	8	85.77	65.69	257.58	8.95	7.62
push	8	512	64	8	59.41	60.32	71.88	8.84	7.63
pull	8	512	128	16	3672.89	3628.99	7880.98	9.07	6.45
push	8	512	128	16	1653.70	1791.73	3681.04	9.06	9.04
pull	8	1024	128	8	986.94	776.67	2764.92	9.10	6.03
push	8	1024	128	8	674.82	551.53	2082.18	9.09	6.37
pull	8	2048	128	8	1048.93	828.82	2905.03	9.20	5.99
push	8	2048	128	8	648.08	371.25	1859.21	9.19	6.46

TODO:

1. Run unit tests

Phase2:

1. Add layer-wise KV push to bring out the actual performance gains from kv push model

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request introduces a significant feature to support KV cache push from a Prefill node to a Decode node using the Nixl KV Connector, aiming to improve TTFT in disaggregated inference setups. The implementation correctly handles two synchronization scenarios between the nodes and includes necessary changes across the connector, engine, and worker components. The logic for handling block registration and the two push scenarios seems well-thought-out. However, I've identified a critical bug in the _write_blocks method that could lead to an UnboundLocalError for models with local attention, which needs to be addressed.

[mergify]

Hi @snadampal, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[yewentao256]

Please solve the comments, pre-commit issue.

Also, please reduce diff to make it easier for review.

Adding lm_eval metrics to make sure we have the correct acc

[snadampal]

@yewentao256 , thanks for the review, will take care of the comments. btw, I'm implementing heterogeneous TP and heterogeneous block_size between P and D nodes (the current version supports only homogeneous) and will update the PR soon.
I'm also unifying the pull and push mode logic for example, combing _read_blocks and _write_blocks to remove redundant code which should bring down the size a bit.

[mergify]

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

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

[snadampal]

Hi @yewentao256 , I have updated the PR with the features I mentioned above and also I've unified the push and pull modes as much as possible to avoid any redundant code. Still it's a large PR, but this is the best I could condense it to. Please review and let me know your comments.

I have started this on yesterday's commit and today it already has conflicts :(
I will rebase again and resolve them after the review is done.

[snadampal]

I will write a design doc (in RFC format) to help with the reviews.

[snadampal]

I have captured design and implementation details in the following design/RDC doc. I hope this helps with reviewing the PR.
#36923 - [RFC]: [KV Connector]: Support KV push from Prefill to Decode node using Nixl Connector

cc: @yewentao256 , @NickLucche

[snadampal]

I have updated the PR to split the implementation into 8 commits; the first 4 commits being just the existing code refactor with zero functionality change, and the last 4 commits are for
(i) P allows block registration from remote
(ii) D registers blocks with P
(iii) P pushes its local blocks into remote blocks, and
(iv) D uses push mode

so that it will be much easier for review and also we can take commit by commit into mainline without disturbing anything. I have rebased it to the latest ( so it's now with hybrid kv feature) and run the benchmarks and accuracy tests, the performance numbers are same as those I had posted above.

As per Nicolo's suggestion, I have also hosted it as an out-of-tree connector as a plugin for people to try out quickly on top of vllm.
https://github.com/snadampal/vllm-nixl-pushpull-connector

cc:@NickLucche, @yewentao256

[mergify]

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

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

[NickLucche]

thanks for the work @snadampal !
Will try to run the numbers on my side too and then we can iron out a few details with this PR

[snadampal]

Hi @NickLucche, I have done some more deep dive into Push mode performance gains and got some interesting findings:

1. Push mode gives 3x or more gains than Pull mode for TTFT for the scenario where there is a 100% prefix cache hit on D node. In this case,

• Pull mode incurs the latency of Prefill node request handling + proxy round-trip. If the prefill node doesn't hit cache, it includes complete prefill forward pass latency which is of the order of several 100s of milliseconds
• Where as in Push mode, the request can start decoding right away and there is no serialization with Prefill or proxy.

2. Push mode TTFT is around 1.3x - 2x faster and TPOT is 30% faster (and hence higher decode throughput) than Pull mode when KV cache memory is fragmented. When the memory is fragmented, the KV blocks are non-contiguous and hence Nixl can't coalesce the descriptors into fewer. For llama-3.3-70b-instruct model I see the descriptor count grew from 118 to 18880 which caused the nixl postXfer() duration to increase by 100x (460usec to 46ms).

• In Pull mode, the start_load_kv -> nixl.transfer() -> postXfer() comes in decode critical path, just before the forward pass. Hence it directly affects TPOT .
• Where as in Push mode, the start_push_kv -> nixl.transfer() -> postXfer() happens in a separate thread. it doesn't affect the worker forward pass. And moreover, this happens on Prefill node and it doesn't effect the Decode latencies.
  Note: Even when Prefill node has to pull KV (my other PR for multi-turn optimization), the postXfer() affect will not be as bad as in Decode node, because the engine loop steps are much longer, order of 100s of ms, so, the additional overhead from postXfer() will not be magnified.

3. Push mode is around 10% faster than Pull mode when there is no prefix cache on Decode node and KV cache is not fragmented. But this scenario happens only a fresh instance start.

Please let me know your observations.

[mergify]

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

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

[NickLucche]

Mega job here @snadampal !
Code is clean and the docs diagram goes a long way into explaining the flow.
I also feel pretty good about the pull separation in terms of maintainability.

I left some comments, none of them is major, so feel free to address them as you see fit, we can even defer some of these ideas to a follow up PR (especially the one around thread state separation, if at all).

[snadampal]

Hi @NickLucche , thanks for the great feedback, as always! I have incorporated everything except making a new Writer class.
I'm deferring the Writer change for the next PR because, based on perf improvements in Push mode, I'm adding the similar threading logic for Pull mode as well, and as part of that I was planning to move the threading logic to BaseConnector if possible.

The PR is now has been rebased, all the feedback incorporated, and tested for unit tests, proxy, and benchmarking script.

[mergify]

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

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

[mergify]

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

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