Skip to content

[P/D] Heterogeneous TP #18079

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Closed
wants to merge 11 commits into from
Closed

[P/D] Heterogeneous TP #18079

wants to merge 11 commits into from

Conversation

@NickLucche
Copy link
Collaborator

@NickLucche NickLucche commented May 13, 2025 •
edited by github-actions bot
Loading

What this PR does:

  • Adds support for heterogenous TP sizes. So far tested with P_TP=1 and D_TP=2/4.
  • It does so by splitting the remote kv cache among D_TP/P_TP D workers along the kv_head dim. block_size and kv precision must match.
  • (JIT) Discovery is done by first querying the rank0 of the dst_engine_id to get the TP size of destination group. After that, every D TP worker will only pull from a single remote P TP worker, hence the setup only requires two exchanges on the side channel (rank0 and target rank_j).

How to test

DECODER_TP_SIZE=2 PREFILLER_TP_SIZE=2 NUM_DECODE_INSTANCES=1 bash tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh

DP TP worker rank and kv split assignment:

image

@NickLucche NickLucche changed the title [P/D] Heterogenous tp [P/D] Heterogenous TP May 13, 2025
@mergify mergify bot added the v1 label May 13, 2025
NickLucche
NickLucche commented May 13, 2025
View reviewed changes
tests/v1/kv_connector/nixl_integration/toy_proxy_server.py
async def send_request_to_service(client_info: dict, endpoint: str,
req_data: dict, request_id: str):
"""
Send a request to a service using a client from the pool.
Copy link
Collaborator Author

@NickLucche NickLucche May 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I pulled this from the old nm/disagg_pd_dev branch. The one on was triggering some TP communication issues between workers.

Copy link
Collaborator Author

@NickLucche NickLucche May 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

let's hold until I figure this out

NickLucche
NickLucche commented May 13, 2025
View reviewed changes
tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh Outdated
}

set_cli_args() {
PREFILLER_TP_SIZE=1
Copy link
Collaborator Author

@NickLucche NickLucche May 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I suppose this could be a simpler pair of env vars

@NickLucche NickLucche mentioned this pull request May 14, 2025
Closed
@mergify
Copy link

mergify bot commented May 15, 2025

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

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

@mergify mergify bot added the needs-rebase label May 15, 2025
NickLucche
NickLucche commented May 16, 2025
View reviewed changes
tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh
PORT=$((8200 + i))
# Calculate side channel port
SIDE_CHANNEL_PORT=$((5659 + i))
SIDE_CHANNEL_PORT=$((5659 + i * $DECODER_TP_SIZE))
Copy link
Collaborator Author

@NickLucche NickLucche May 16, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is bugged only works with 1P 1D btw

NickLucche added 10 commits May 20, 2025 07:45
@NickLucche
split kv_cache along head dim
20ca3aa 
fix descr indexing

change remote worker selection indexing; test ptp2-dtp4

Signed-off-by: nicklucche <[email protected]>
@NickLucche
address race condition
9ed29aa 
Signed-off-by: nicklucche <[email protected]>
@NickLucche
integration test
c59ca52 
Signed-off-by: nicklucche <[email protected]>
@NickLucche
address review: move src_xfer_side_handle init out of hotpath and docs
0e8f82a 
Signed-off-by: nicklucche <[email protected]>
@NickLucche
address review: clarity
3e11240 
Signed-off-by: nicklucche <[email protected]>
@NickLucche
fix req_finished sync bug
e152305 
Signed-off-by: nicklucche <[email protected]>
@NickLucche
simplify script
f2c2180 
Signed-off-by: nicklucche <[email protected]>
@NickLucche
optimize req state checking loop
a9cb4a5 
Signed-off-by: nicklucche <[email protected]>
@NickLucche
release_xfer_handle on status DONE
ef200c0 
Signed-off-by: nicklucche <[email protected]>
@NickLucche
send notif to agent_name
355c2f1 
Signed-off-by: nicklucche <[email protected]>
@mergify mergify bot removed the needs-rebase label May 20, 2025
@NickLucche
Copy link
Collaborator Author

NickLucche commented May 21, 2025

Update:
this PR was tested on lm-eval (run_accuracy_test.sh ) on nixl 0.2.1 (4f37f07) .

Unfortunately it does not work with nixl 0.1.1, as during lm-eval the decoder will hang, triggering TimeoutError.
I've spent some time trying to nail this down, but I am still unsure about the root cause.

@NickLucche
fix req abort
9d65ae3 
Signed-off-by: nicklucche <[email protected]>
@mergify
Copy link

mergify bot commented May 21, 2025

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

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

@mergify mergify bot added the needs-rebase label May 21, 2025
@robertgshaw2-redhat robertgshaw2-redhat changed the title [P/D] Heterogenous TP [P/D] Heterogeneous TP May 25, 2025
@NickLucche NickLucche mentioned this pull request May 28, 2025
Merged
@NickLucche NickLucche mentioned this pull request Jun 4, 2025
Open
@NickLucche
Copy link
Collaborator Author

NickLucche commented Jun 4, 2025

Closing in favor of #18833.
The higher number of (small-size) descriptors this PR proposed didn't work with nixl, causing significant slowdowns.

@NickLucche NickLucche closed this Jun 4, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@WoosukKwon WoosukKwon Awaiting requested review from WoosukKwon WoosukKwon is a code owner

@robertgshaw2-redhat robertgshaw2-redhat Awaiting requested review from robertgshaw2-redhat robertgshaw2-redhat is a code owner

@njhill njhill Awaiting requested review from njhill njhill is a code owner

@ywang96 ywang96 Awaiting requested review from ywang96 ywang96 is a code owner

@comaniac comaniac Awaiting requested review from comaniac comaniac is a code owner

@alexm-redhat alexm-redhat Awaiting requested review from alexm-redhat alexm-redhat is a code owner

Assignees

No one assigned

Labels

needs-rebase v1

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@NickLucche