[None][feat] kv cache manager v1 use fabric memory to enable mnnvl transfer for kv cache transfer in disaggregated serving - #12490
Conversation
|
/bot run |
|
PR_Github #40094 [ run ] triggered by Bot. Commit: |
📝 WalkthroughWalkthroughThis pull request introduces fabric memory support for KV-cache pool allocation in TensorRT-LLM. It adds an environment variable flag and integrates FabricMemory into the WindowBlockManager to optionally allocate KV-cache pools using fabric memory instead of standard GPU or UVM allocation paths. Changes
Sequence DiagramsequenceDiagram
participant Config as Configuration
participant WBM as WindowBlockManager
participant FM as FabricMemory
participant BM as BufferManager
participant Pool as Pool Storage
Config->>Config: Read TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY
WBM->>Config: getEnvKVCachePoolUseFabricMemory()
Config-->>WBM: fabric_enabled (bool)
rect rgba(100, 150, 200, 0.5)
alt Fabric Memory Enabled
WBM->>FM: Allocate pool buffer<br/>(size = tensor_volume × dtype_size)
FM-->>WBM: FabricMemory pointer
WBM->>WBM: Wrap as ITensor
WBM->>Pool: Store in mFabricMemoryPools
else Fallback to Standard
WBM->>BM: Allocate via managed()/gpuSync()
BM-->>WBM: Buffer pointer
end
end
WBM->>Pool: pool.primaryPtr assignment
Pool-->>WBM: Pool allocated
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~22 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
cpp/tensorrt_llm/common/envUtils.cpp (1)
1-2:⚠️ Potential issue | 🟡 MinorUpdate the modified-file copyright year.
This file changed in 2026, so the header should be bumped from
2022-2024to2022-2026.As per coding guidelines, "Add NVIDIA copyright header on ALL new files, and update year on modified files."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cpp/tensorrt_llm/common/envUtils.cpp` around lines 1 - 2, Update the file header comment string that currently reads "SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES." to bump the year range to "2022-2026" in the top-of-file license block (the SPDX header in envUtils.cpp).cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp (1)
939-955:⚠️ Potential issue | 🔴 CriticalSynchronize outstanding work before freeing fabric-backed pools.
mFabricMemoryPools.clear()destroys the external allocation before any stream or transfer-manager fence runs. If a transfer or kernel is still in flight againstpool.primaryPtr, this becomes a use-after-free.💡 Suggested fix
void WindowBlockManager::releasePools() { + mTransferManager->syncTransfers(); + mBufferManager.getStream().synchronize(); for (auto& pool : mPools) { if (pool.primaryPtr) { pool.primaryPtr->release(); @@ { pool.secondaryPtr->release(); } } - // Release fabric memory backing (must happen after ITensor release) + // Release fabric memory backing after all in-flight work has completed. mFabricMemoryPools.clear(); - mBufferManager.getStream().synchronize(); mBufferManager.memoryPoolTrimTo(0); }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp` around lines 939 - 955, In WindowBlockManager::releasePools the code clears mFabricMemoryPools before synchronizing, risking freeing fabric-backed allocations still in use; move the stream synchronization and memoryPoolTrimTo(0) calls to occur before mFabricMemoryPools.clear() (i.e., call mBufferManager.getStream().synchronize(); mBufferManager.memoryPoolTrimTo(0); then clear mFabricMemoryPools), ensuring any in-flight transfers/kernels touching pool.primaryPtr/pool.secondaryPtr complete before the fabric-backed memory is destroyed.cpp/tensorrt_llm/common/envUtils.h (1)
1-2:⚠️ Potential issue | 🟡 MinorUpdate the modified-file copyright year.
This header changed in 2026, so the banner should be bumped from
2022-2024to2022-2026.As per coding guidelines, "Add NVIDIA copyright header on ALL new files, and update year on modified files."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cpp/tensorrt_llm/common/envUtils.h` around lines 1 - 2, Update the copyright banner at the top of cpp/tensorrt_llm/common/envUtils.h: change the SPDX-FileCopyrightText year range string from "2022-2024" to "2022-2026" so the header reflects the 2026 modification (look for the comment block starting with "/*" and the line containing "SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.").
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Outside diff comments:
In `@cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp`:
- Around line 939-955: In WindowBlockManager::releasePools the code clears
mFabricMemoryPools before synchronizing, risking freeing fabric-backed
allocations still in use; move the stream synchronization and
memoryPoolTrimTo(0) calls to occur before mFabricMemoryPools.clear() (i.e., call
mBufferManager.getStream().synchronize(); mBufferManager.memoryPoolTrimTo(0);
then clear mFabricMemoryPools), ensuring any in-flight transfers/kernels
touching pool.primaryPtr/pool.secondaryPtr complete before the fabric-backed
memory is destroyed.
In `@cpp/tensorrt_llm/common/envUtils.cpp`:
- Around line 1-2: Update the file header comment string that currently reads
"SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION &
AFFILIATES." to bump the year range to "2022-2026" in the top-of-file license
block (the SPDX header in envUtils.cpp).
In `@cpp/tensorrt_llm/common/envUtils.h`:
- Around line 1-2: Update the copyright banner at the top of
cpp/tensorrt_llm/common/envUtils.h: change the SPDX-FileCopyrightText year range
string from "2022-2024" to "2022-2026" so the header reflects the 2026
modification (look for the comment block starting with "/*" and the line
containing "SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION &
AFFILIATES. All rights reserved.").
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 7ad5e998-6972-4e13-823d-ffc6442cdd48
📒 Files selected for processing (4)
cpp/include/tensorrt_llm/batch_manager/kvCacheManager.hcpp/tensorrt_llm/batch_manager/kvCacheManager.cppcpp/tensorrt_llm/common/envUtils.cppcpp/tensorrt_llm/common/envUtils.h
|
PR_Github #40094 [ run ] completed with state
|
3411faa to
d45c7de
Compare
|
/bot run |
|
PR_Github #40305 [ run ] triggered by Bot. Commit: |
|
PR_Github #40305 [ run ] completed with state
|
d45c7de to
3cadb90
Compare
|
/bot run |
|
PR_Github #40668 [ run ] triggered by Bot. Commit: |
|
PR_Github #40668 [ run ] completed with state |
|
/bot run |
|
PR_Github #40681 [ run ] triggered by Bot. Commit: |
|
PR_Github #40681 [ run ] completed with state
|
3cadb90 to
fc8a045
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #40897 [ run ] triggered by Bot. Commit: |
|
PR_Github #40897 [ run ] completed with state
|
…rom PR NVIDIA#12490) Enable MNNVL transfer by setting TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY=1. When enabled, KV cache pools are allocated from FabricMemory instead of standard GPU allocation. Signed-off-by: Ethan Kou <ekou@nvidia.com> Signed-off-by: Ethan Kou <ekou@login-lyris01.lyris.clusters.nvidia.com>
added |
716f1b3 to
2c6fbd8
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #41142 [ run ] triggered by Bot. Commit: |
|
PR_Github #41142 [ run ] completed with state
|
644747e to
ff35030
Compare
|
/bot run --stage-list "A30-AutoDeploy-1, DGX_B200-AutoDeploy-1, H100_PCIe-AutoDeploy-1, DGX_B200-4_GPUs-PyTorch-Ray-1, DGX_B200-8_GPUs-PyTorch-1, DGX_H100-4_GPUs-PyTorch-Others-1, GB200-8_GPUs-2_Nodes-PyTorch-1" |
ff35030 to
7861500
Compare
|
/bot run --stage-list "A30-AutoDeploy-1, DGX_B200-AutoDeploy-1, H100_PCIe-AutoDeploy-1, DGX_B200-4_GPUs-PyTorch-Ray-1, DGX_B200-8_GPUs-PyTorch-1, DGX_H100-4_GPUs-PyTorch-Others-1, GB200-8_GPUs-2_Nodes-PyTorch-1" |
|
PR_Github #56260 [ run ] triggered by Bot. Commit: |
|
PR_Github #56260 [ run ] completed with state
|
7861500 to
3121b16
Compare
|
/bot run --stage-list "DGX_B200-AutoDeploy-1, H100_PCIe-AutoDeploy-1, DGX_B200-4_GPUs-PyTorch-Ray-1, DGX_B200-8_GPUs-PyTorch-1, DGX_H100-4_GPUs-PyTorch-Others-1, GB200-8_GPUs-2_Nodes-PyTorch-1" |
|
PR_Github #56333 [ run ] triggered by Bot. Commit: |
|
PR_Github #56333 [ run ] completed with state
|
|
/bot run --stage-list "DGX_B200-AutoDeploy-1, H100_PCIe-AutoDeploy-1, DGX_B200-4_GPUs-PyTorch-Ray-1, DGX_B200-8_GPUs-PyTorch-1, DGX_H100-4_GPUs-PyTorch-Others-1, GB200-8_GPUs-2_Nodes-PyTorch-1" |
|
PR_Github #56356 [ run ] triggered by Bot. Commit: |
|
PR_Github #56356 [ run ] completed with state
|
3121b16 to
0858b0e
Compare
|
/bot run --stage-list "H100_PCIe-AutoDeploy-1, DGX_B200-AutoDeploy-1" |
Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
|
/bot run --stage-list "H100_PCIe-AutoDeploy-1, DGX_B200-AutoDeploy-1" |
|
/bot skip --comment "all related tests have passed" |
|
PR_Github #56556 [ run ] triggered by Bot. Commit: |
|
PR_Github #56557 [ skip ] triggered by Bot. Commit: |
|
PR_Github #56556 [ run ] completed with state |
|
PR_Github #56557 [ skip ] completed with state |
7 participants
Summary by CodeRabbit
Description
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.