feat: Using NIXL for KV cache transfer when using disaggregated serving in TRTLLM #1591
Conversation
tanmayv25
requested review from
GuanLuo,
richardhuo-nv,
rmccorm4,
ptarasiewiczNV,
ishandhanani,
alec-flowers,
nnshah1 and
a team
as code owners
WalkthroughThe changes introduce an experimental NIXL-based KV cache transfer method for TensorRT-LLM, affecting the Dockerfile, build scripts, and documentation. New build arguments, environment variable management, and command-line flags were added to enable or disable this feature. The documentation now details how to use and switch between UCX and NIXL KV cache transfer methods. Changes
Poem
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
Documentation and Community
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
examples/tensorrt_llm/README.md (2)
309-311: Fix grammar in Note about production suitability
The phrase “is not be suitable” is incorrect. Please update to “is not suitable” or “may not be suitable.”
335-336: Add missing article for clarity
Change “built in previous step” to “built in the previous step.”container/build.sh (1)
523-526: Allow overriding NIXL commit inbuild.sh
Currently,build.shalways passes its internalNIXL_COMMITto the wheel builder. To support custom NIXL commits, consider adding a CLI option (e.g.--nixl-commit) that updatesNIXL_COMMITbefore invokingbuild_trtllm_wheel.sh.
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
container/Dockerfile.tensorrt_llm(1 hunks)container/build.sh(5 hunks)container/build_trtllm_wheel.sh(3 hunks)examples/tensorrt_llm/README.md(1 hunks)
🧰 Additional context used
🪛 LanguageTool
examples/tensorrt_llm/README.md
[grammar] ~311-~311: It seems that the verb “be” is not needed here.
Context: ...NIXL support is experimental and is not be suitable for production environments. #### Usin...
(BE_NOT_BE_JJ)
[uncategorized] ~335-~335: You might be missing the article “the” here.
Context: ...w to start the container image built in previous step. 3. **Start the disaggregated ser...
(AI_EN_LECTOR_MISSING_DETERMINER_THE)
⏰ Context from checks skipped due to timeout of 90000ms (1)
- GitHub Check: Build and Test - vllm
🔇 Additional comments (13)
container/Dockerfile.tensorrt_llm (4)
327-329: Ensure correct NIXL library path injection
TheARG ARCH_ALTand the prepend toLD_LIBRARY_PATHproperly direct the runtime linker to the NIXL libraries based on CPU architecture.
331-337: Conditional environment variable for KV cache transfer
TheRUNblock correctly appends eitherTRTLLM_USE_NIXL_KVCACHE=1orTRTLLM_USE_UCX_KVCACHE=1to/etc/environmentbased on the experimental flag.
338-346: Generate runtime script for KV cache mode
Theset_trtllm_env.shscript creation and permission setting correctly exports the chosen variable at container startup.
347-349: Source the KV cache script in shell initialization
Appending thesourcecommand to.bashrcensures the environment variable is loaded for interactive sessions.examples/tensorrt_llm/README.md (1)
301-302: Document NIXL experimental section header
The new section header clearly identifies the experimental NIXL KV cache transfer instructions.container/build_trtllm_wheel.sh (4)
21-26: Add-noption for specifying NIXL commit
The newgetoptsblock correctly parses-nintoNIXL_COMMIT.
27-32: Update usage message to include-nflag
The usage instructions now properly document the-n: NIXL commitoption.
41-43: Capture script directory for relative paths
DefiningMAIN_DIRensures subsequent file operations use a stable absolute path.
93-96: Split build into Triton devel and wheel stages
The two-stepmake tritondevel_buildthenmake wheel_buildwithNIXL_ROOTcorrectly integrates NIXL during wheel compilation.container/build.sh (4)
93-93: Declare default experimental flag
SettingTRTLLM_USE_NIXL_KVCACHE_EXPERIMENTAL="0"provides a clear default for the build argument.
170-176: Handle--trtllm-use-nixl-kvcache-experimentalCLI option
The new case branch properly validates no argument is passed and sets the internal flag.
375-376: Document the new experimental flag in help
The help text now includes--trtllm-use-nixl-kvcache-experimental.
504-507: Pass NIXL flag as Docker build argument
Adding--build-arg TRTLLM_USE_NIXL_KVCACHE_EXPERIMENTALensures the Dockerfile sees the correct value.
Co-authored-by: Ryan McCormick <rmccormick@nvidia.com> Signed-off-by: Tanmay Verma <tanmay2592@gmail.com>
Co-authored-by: Ryan McCormick <rmccormick@nvidia.com> Signed-off-by: Tanmay Verma <tanmay2592@gmail.com>
Co-authored-by: Ryan McCormick <rmccormick@nvidia.com> Signed-off-by: Tanmay Verma <tanmay2592@gmail.com>
| unset TRTLLM_USE_NIXL_KVCACHE | ||
| export TRTLLM_USE_UCX_KVCACHE=1 |
There was a problem hiding this comment.
Is it worth adding all this extra build logic, flags, documentation, etc. if we can just boil down the opt-in steps to something like this?
NIXL (Experimental)
To use NIXL for KV Cache transfer instead of UCX, set the relevant environment variables:
export TRTLLM_USE_NIXL_KVCACHE=1
export TRTLLM_USE_UCX_KVCACHE=0There was a problem hiding this comment.
We are currently using the TensorRT-LLM wheel from pyPI by default. This has reduced the container build time considerably. However, these public wheels are not build with NIXL support.
Currently, there is no way around rebuilding the TRTLLM wheel from scratch with some extra flags if someone wants to use NIXL. We could have enabled NIXL by default but it would have undermined our build time optimization from using pyPI TRTLLM wheels. Once NIXL in TRTLLM is mature enough and available in TRTLLM wheels on pypi, we can simplify instructions to what you have mentioned above.
In this section, I just wanted to point out that TRTLLM wheels built with NIXL support can also be used to use UCX instead.
Updates build scripts to allow users to conveniently use NIXL for KV cache transfer when using dynamo with TRTLLM.
Summary by CodeRabbit
New Features
Documentation
Chores