Skip to content

NIXL/UCX threadpool - #573

Merged
brminich merged 20 commits into
ai-dynamo:mainfrom
iyastreb:nixl-ucx-workers-pool
Aug 15, 2025
Merged

brminich merged 20 commits into
ai-dynamo:mainfrom
iyastreb:nixl-ucx-workers-pool

Conversation

@iyastreb

@iyastreb iyastreb commented Jul 15, 2025 •

Copy link
Copy Markdown
Contributor

What?

Enable threadpool in NIXL/UCX to parallelise posting of large requests (in terms of batch size) .

To enable this feature client needs to set "num_threads" in UCX backend params (close to "num_workers").
This essentially specifies the number of dedicated workers and threads (1 worker per thread) that are used to post requests and progress.

Progress thread logic remains the same as before: if configured, a single progress thread is used for all shared workers. There must be at least one shared worker configured in any use case.

Why?

Clients report long posting time with large batch sizes (> 1000).
For example, posting of 32k batch size with cuda_ipc protocol being used takes up to 0.25s.
With this fix we can drastically reduce the time needed to post the request.

How?

This PR depends on refactoring PR: #606 (merged)
It depends also on PR #673, which bumps UCX from 1.18 to 1.19
The performance evaluation

Summary:
With communication over IB devices, using multiple threads for posting+progress makes a huge difference in terms of request posting time, and also improves BW.
E.g. on 64k message size, the posting time is reduced by 8x, and BW is 10% better

With cuda_ipc communication I see not so huge decrease of posting time, it's basically around 30%, but BW is also 30% better.

@github-actions

Copy link
Copy Markdown

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

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

🚀

Comment thread benchmark/kvbench/README.md
Comment thread src/plugins/ucx/meson.build
Comment thread src/plugins/ucx/ucx_backend.cpp Outdated
Comment thread src/plugins/ucx/ucx_backend.cpp Outdated
@copy-pr-bot

copy-pr-bot Bot commented Jul 17, 2025

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@iyastreb

iyastreb commented Aug 8, 2025

Copy link
Copy Markdown
Contributor Author

Tests fail because of using UCX 1.18, should be fixed by #673

Comment thread src/plugins/ucx/ucx_backend.h
Comment thread src/plugins/ucx/ucx_backend.h
Comment thread src/api/python/_api.py Outdated
mkhazraee
mkhazraee previously approved these changes Aug 12, 2025
Comment thread test/gtest/error_handling.cpp
Comment thread src/plugins/ucx/ucx_backend.cpp Outdated
Comment thread src/plugins/ucx/ucx_backend.cpp Outdated
Comment thread src/plugins/ucx/ucx_backend.cpp Outdated
Comment thread src/plugins/ucx/ucx_backend.cpp Outdated
Comment thread src/plugins/ucx/ucx_backend.cpp Outdated
Comment thread src/plugins/ucx/ucx_backend.cpp
Comment thread src/plugins/ucx/ucx_backend.h
Comment thread src/plugins/ucx/ucx_backend.h
Comment thread src/plugins/ucx/ucx_backend.h
Comment thread src/plugins/ucx/ucx_backend.cpp
Comment thread src/plugins/ucx/ucx_backend.h
Comment thread src/plugins/ucx/ucx_backend.cpp Outdated
Comment thread src/plugins/ucx/ucx_backend.cpp Outdated
Comment thread src/plugins/ucx/ucx_backend.cpp
Comment thread src/plugins/ucx/ucx_backend.cpp
brminich
brminich previously approved these changes Aug 15, 2025
Signed-off-by: Ilia Yastrebov <iyastrebov@nvidia.com>
@brminich

Copy link
Copy Markdown
Contributor

/build

@brminich

Copy link
Copy Markdown
Contributor

/build

This branch was previously deployed

2 inactive deployments
GITLAB — 97a9091b Deployed Aug 15, 2025 by copy-pr-bot[bot] via Trigger CI Pipeline #3932
SWX_AWS — 97a9091b Deployed Aug 15, 2025 by copy-pr-bot[bot] via Run AWS S3 Tests #644
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

7 participants