[#15793][perf] Optimize KV cache for unified memory systems (DGX Spark) - #12301
[#15793][perf] Optimize KV cache for unified memory systems (DGX Spark)#12301mihai-chiorean wants to merge 12 commits into
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
WalkthroughThe changes detect integrated CUDA devices and adapt KV cache configuration and memory management for unified CPU-GPU memory, including secondary-pool allocation, transfer handling, and memory-budget accounting. ChangesUnified memory KV cache handling
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant App as Application
participant API as KvCacheConfig
participant CUDA as CUDA runtime
participant KVMgr as KV cache manager
App->>API: Create cache configuration
API->>CUDA: Detect integrated device
CUDA-->>API: Unified-memory status
API-->>App: Apply unified-memory optimization defaults
App->>KVMgr: Allocate and size KV cache pools
alt Unified-memory system
KVMgr->>KVMgr: Use GPU-accessible secondary memory
KVMgr->>KVMgr: Skip onboard/offload transfers
KVMgr->>KVMgr: Fold secondary budget into primary pool
else Discrete GPU system
KVMgr->>KVMgr: Preserve host-pool and transfer behavior
end
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ 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.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@cpp/include/tensorrt_llm/common/cudaUtils.h`:
- Around line 320-331: The function isUnifiedMemorySystem caches a single global
result in sIsUnified but queries the current device (cudaGetDevice), so if the
active CUDA device changes later the cached value can be stale; update
isUnifiedMemorySystem to either query the device attributes on every call
(remove the static sIsUnified lambda) or maintain a per-device cache keyed by
the device id (call cudaGetDevice, then consult/populate a map< int, bool >
using cudaDeviceGetAttribute and check_cuda_error) so the returned
pageable-memory flag correctly reflects the current device.
In `@tensorrt_llm/llmapi/llm_args.py`:
- Around line 2205-2208: The try/except around the call to
is_device_integrated() should stop catching broad Exception; replace the broad
catch with a specific catch for RuntimeError so CUDA-related failures are
handled while other exceptions propagate. Locate the block where unified =
is_device_integrated() is called and change the except Exception to except
RuntimeError, leaving unified=False on that branch and not swallowing other
errors.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: bdc552fe-e226-42ae-866b-31769cdf4f1c
📒 Files selected for processing (3)
cpp/include/tensorrt_llm/common/cudaUtils.hcpp/tensorrt_llm/batch_manager/kvCacheManager.cpptensorrt_llm/llmapi/llm_args.py
| inline bool isUnifiedMemorySystem() | ||
| { | ||
| static bool const sIsUnified = []() | ||
| { | ||
| int device{-1}; | ||
| check_cuda_error(cudaGetDevice(&device)); | ||
| int pageableMemAccess{0}; | ||
| check_cuda_error(cudaDeviceGetAttribute(&pageableMemAccess, cudaDevAttrPageableMemoryAccess, device)); | ||
| return pageableMemAccess != 0; | ||
| }(); | ||
| return sIsUnified; | ||
| } |
There was a problem hiding this comment.
Device-specific detection is cached globally
Line 322 caches the result once for the process, but the query is for the current CUDA device. If the active device changes later, this can return stale data and apply the wrong KV-cache path.
💡 Proposed fix (evaluate per call to avoid stale cross-device cache)
inline bool isUnifiedMemorySystem()
{
- static bool const sIsUnified = []()
- {
- int device{-1};
- check_cuda_error(cudaGetDevice(&device));
- int pageableMemAccess{0};
- check_cuda_error(cudaDeviceGetAttribute(&pageableMemAccess, cudaDevAttrPageableMemoryAccess, device));
- return pageableMemAccess != 0;
- }();
- return sIsUnified;
+ int device{-1};
+ check_cuda_error(cudaGetDevice(&device));
+ int pageableMemAccess{0};
+ check_cuda_error(cudaDeviceGetAttribute(&pageableMemAccess, cudaDevAttrPageableMemoryAccess, device));
+ return pageableMemAccess != 0;
}🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@cpp/include/tensorrt_llm/common/cudaUtils.h` around lines 320 - 331, The
function isUnifiedMemorySystem caches a single global result in sIsUnified but
queries the current device (cudaGetDevice), so if the active CUDA device changes
later the cached value can be stale; update isUnifiedMemorySystem to either
query the device attributes on every call (remove the static sIsUnified lambda)
or maintain a per-device cache keyed by the device id (call cudaGetDevice, then
consult/populate a map< int, bool > using cudaDeviceGetAttribute and
check_cuda_error) so the returned pageable-memory flag correctly reflects the
current device.
| try: | ||
| unified = is_device_integrated() | ||
| except Exception: | ||
| unified = False |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Verify broad exception handling remains removed from this validator.
rg -nP --type=py -C2 'except\s+Exception\s*:' tensorrt_llm/llmapi/llm_args.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 290
🏁 Script executed:
# Find and examine the is_device_integrated() function definition
rg -n 'def is_device_integrated' --type=py -A5Repository: NVIDIA/TensorRT-LLM
Length of output: 525
🏁 Script executed:
# Check for all broad exception handlers in llm_args.py
rg -nP --type=py 'except\s+(Exception|BaseException|:)' tensorrt_llm/llmapi/llm_args.py | head -20Repository: NVIDIA/TensorRT-LLM
Length of output: 249
🏁 Script executed:
# Get broader context around lines 2205-2208
sed -n '2200,2215p' tensorrt_llm/llmapi/llm_args.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 840
🏁 Script executed:
# Get the full implementation of is_device_integrated()
sed -n '1362,1400p' tensorrt_llm/_utils.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 1473
🏁 Script executed:
# Check if is_device_integrated() can raise specific exceptions
rg -n 'is_device_integrated' tensorrt_llm/_utils.py -A20Repository: NVIDIA/TensorRT-LLM
Length of output: 944
Catch specific exceptions instead of broad Exception
Line 2207 catches Exception, which masks unrelated errors and complicates debugging. The is_device_integrated() function raises RuntimeError on CUDA-related failures, so catch that specifically. Avoid AssertionError in the proposed fix—this function does not raise it.
Suggested fix
if self.enable_unified_memory_optimization is None:
try:
unified = is_device_integrated()
- except Exception:
+ except RuntimeError:
+ logger.debug(
+ "Unified-memory auto-detection failed; defaulting to disabled")
unified = False
object.__setattr__(self, 'enable_unified_memory_optimization',
unified)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| try: | |
| unified = is_device_integrated() | |
| except Exception: | |
| unified = False | |
| if self.enable_unified_memory_optimization is None: | |
| try: | |
| unified = is_device_integrated() | |
| except RuntimeError: | |
| logger.debug( | |
| "Unified-memory auto-detection failed; defaulting to disabled") | |
| unified = False | |
| object.__setattr__(self, 'enable_unified_memory_optimization', | |
| unified) |
🧰 Tools
🪛 Ruff (0.15.6)
[warning] 2207-2207: Do not catch blind exception: Exception
(BLE001)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tensorrt_llm/llmapi/llm_args.py` around lines 2205 - 2208, The try/except
around the call to is_device_integrated() should stop catching broad Exception;
replace the broad catch with a specific catch for RuntimeError so CUDA-related
failures are handled while other exceptions propagate. Locate the block where
unified = is_device_integrated() is called and change the except Exception to
except RuntimeError, leaving unified=False on that branch and not swallowing
other errors.
svc-trtllm-gh-bot
added
the
Community want to contribute
|
@eopXD could you review this? thanks |
|
I wonder why we need a secondary pool at all on DGX Spark. Also do you have any test result or evidence on memory usage reduction? Thanks! |
|
@pengbowang-nv We don't need a secondary pool on DGX Spark — hence Tested on DGX Spark GB10 (SM121, 128GB unified memory):
Rebased on latest main, no conflicts. |
|
I think we shouldn't allow a secondary pool with non-zero size on DGX Spark, it meaningless since we want secondary pool to hold what first pool cannot hold. |
c86f04c to
91dd9a5
Compare
|
@pengbowang-nv Good point — updated. |
|
@pengbowang-nv what do you think? I've addressed your feedback. |
|
Thanks! I'm thinking of complete remove these code guarding operation under unified memory situation, which means we assume that we won't be using secondary pool on unified memory at all. So just guard on python side would be enough. I'm also waiting for @eopXD to give some advice as I'm not sure of impact of such change. |
yeaa that's a fair call.happy to give that a try if that's the preference. |
|
@pengbowang-nv I simplified it. C++ footprint goes from 5 guard sites to 1. Python-side |
a655b61 to
0e3f2b0
Compare
0e3f2b0 to
89babd4
Compare
|
@eopXD Could you review this? Thanks |
|
@hchings — friendly bump. This is a small, scoped change for unified-memory systems (DGX Spark / SM121):
@nvpohanh originally suggested @eopXD might be a good reviewer here. If neither of you is the right person, please feel free to redirect to the kv-cache-manager-devs group — I can't add team reviewers as an external contributor. |
… Spark) On Grace Blackwell systems (e.g. DGX Spark / GB10), CPU and GPU share the same 128 GB LPDDR5x physical memory pool. The KV cache block manager treats GPU and CPU as separate tiers, which forces a non-zero host_cache_size to allocate a meaningless secondary pool and triggers unnecessary memcpy for offload/onload that are physically no-ops. Changes: - Add isUnifiedMemorySystem() to cudaUtils.h, using cudaDevAttrIntegrated (not cudaDevAttrPageableMemoryAccess, which returns true on any HMM-enabled discrete GPU). Queries each call -- the attribute lookup is cheap and supports mixed-device processes. - Add unified_memory_detected field to KvCacheConfig in Python, with auto-detection via is_device_integrated(). When True, host_cache_size is forced to 0 in the validator. This is the primary guard. - Add a single defensive guard in calculateFreeMemBytes() that forces freeSecondaryMemBytes=0 on unified memory if host_cache_size is still non-zero (e.g. C++ executor API users bypassing the Python validator). All changes are no-ops on discrete GPU systems. Refs: NVIDIA#3130 Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
b370d78 to
893d3f4
Compare
Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
|
@hchings @pengbowang-nv @nvpohanh can I get a CI run..? this seems good to go |
|
/bot run |
mihai-chiorean
changed the title
Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com> # Conflicts: # tensorrt_llm/llmapi/llm_args.py
Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
7 participants
Summary
Fixes #15793
On unified-memory systems like DGX Spark / GB10, CPU and GPU share the same physical LPDDR memory pool. The current KV cache sizing path can still treat GPU memory and host memory as separate tiers when
host_cache_sizeis nonzero, but there is no separate host DRAM tier to offload to on these systems.cudaMemGetInfoalready reports the unified pool, so the secondary host-cache budget should be folded away instead of tracked as a separate pool.This patch adds integrated-device detection and folds the host KV cache tier on unified-memory systems:
Detect true unified-memory devices: Add
isUnifiedMemorySystem()in C++ usingcudaDevAttrIntegrated. This intentionally does not usecudaDevAttrPageableMemoryAccess, which can also be true on HMM-enabled discrete GPUs.Fold Python
host_cache_sizedefaults: Addunified_memory_detectedtoKvCacheConfig, auto-detect viais_device_integrated(), and forcehost_cache_sizeto0when unified memory is detected.Defensive C++ budget guard: In
calculateFreeMemBytes(), forcefreeSecondaryMemBytes = 0on unified-memory systems if C++ executor users bypass the Python validator and still provide a nonzerohost_cache_size.All changes are gated on integrated-device detection and preserve existing behavior on discrete GPU systems.
Test Coverage
is_device_integrated()returns true on DGX Spark GB10 / SM121 with about 121.7 GB visible memory.KvCacheConfigauto-detects the unified-memory device path.LLM_MODELS_ROOTand full model data, so they are deferred to NVIDIA CI.Test plan
KvCacheConfigauto-detects the unified-memory device path.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.
Dev Engineer Review
common::isUnifiedMemorySystem()usingcudaDevAttrIntegrated, correctly avoidingcudaDevAttrPageableMemoryAccessand preserving discrete-GPU behavior.QA Engineer Review
No test changes.