Skip to content

BACKEND/UCX: Multi-device support: Allow setting CUDA context to match device of the memory #857

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
Merged
merged 5 commits into from
Oct 14, 2025
Merged

BACKEND/UCX: Multi-device support: Allow setting CUDA context to match device of the memory #857

merged 5 commits into from
Oct 14, 2025

Conversation

tvegas1
Copy link
Contributor

@tvegas1 tvegas1 commented Oct 3, 2025 •
edited
Loading

What?

Allow changing the CUDA context to match the address.

Why?

Allow multi-device support.

How?

$ nixlbench --backend UCX --initiator_seg_type VRAM --target_seg_type VRAM \
    --num_target_dev 8 --num_initiator_dev 8 --mode MG -op_type WRITE
...
Block Size (B)      Batch Size     B/W (GB/Sec)   Avg Lat. (us)  Avg Prep (us)  P99 Prep (us)  Avg Post (us)  P99 Post (us)  Avg Tx (us)    P99 Tx (us)
----------------------------------------------------------------------------------------------------------------------------------------------------------------
16777216            32             357.716629     46.9           220.0          220.0          914.0          942.0          11089.0        11281.0

@github-actions GitHub Actions
Copy link

github-actions bot commented Oct 3, 2025

👋 Hi tvegas1! Thank you for contributing to ai-dynamo/nixl.

Your PR reviewers will review your contribution then trigger the CI to test your changes.

🚀

@tvegas1
Copy link
Contributor Author

tvegas1 commented Oct 3, 2025

Any input on testing it? We might even be able to remove it altogether since UCX will set it internally if needed?

brminich
brminich reviewed Oct 7, 2025
View reviewed changes
@brminich
Copy link
Contributor

brminich commented Oct 7, 2025

/build

@ovidiusm
Copy link
Contributor

/build

@tvegas1
Copy link
Contributor Author

tvegas1 commented Oct 13, 2025

/build

1 similar comment
@ovidiusm
Copy link
Contributor

/build

@ovidiusm ovidiusm changed the title BACKEND/UCX: Allow setting CUDA context to match device of the memory. BACKEND/UCX: Allow setting CUDA context to match device of the memory (multi-device support) Oct 13, 2025
@ovidiusm ovidiusm changed the title BACKEND/UCX: Allow setting CUDA context to match device of the memory (multi-device support) BACKEND/UCX: Multi-device support: Allow setting CUDA context to match device of the memory Oct 13, 2025
@ovidiusm
Copy link
Contributor

/build

@ovidiusm ovidiusm merged commit 2579ec3 into ai-dynamo:main Oct 14, 2025
21 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@brminich brminich brminich approved these changes

@ovidiusm ovidiusm ovidiusm approved these changes

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

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

Assignees

No one assigned

Labels

external-contribution size/S

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@tvegas1 @brminich @ovidiusm