Repository navigation
[#15793][perf] Optimize KV cache for unified memory systems (DGX Spark) - #12301
mihai-chiorean wants to merge 6 commits into
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
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 (4)
🚧 Files skipped from review as they are similar to previous changes (4)
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review. WalkthroughThe changes detect integrated CUDA devices and prevent a separate host KV cache budget on unified-memory systems. Python configuration and C++ memory accounting set the host or secondary cache budget to zero when the device is integrated. ChangesUnified memory KV cache handling
Estimated code review effort: 3 (Moderate) | ~25 minutes Merge Risk: ⚪ Minimal · up to This change folds the secondary KV-cache budget into the unified memory pool only on integrated-memory systems while preserving discrete-GPU behavior; no actionable merge-blocking risk remains beyond normal checks and review. Sequence Diagram(s)sequenceDiagram
participant KvCacheConfig
participant CUDA as CUDA runtime
participant KVCacheManager
participant Tests
KvCacheConfig->>CUDA: Query integrated-device status
CUDA-->>KvCacheConfig: Return unified-memory status
KvCacheConfig->>KvCacheConfig: Set host_cache_size to zero
KVCacheManager->>CUDA: Query integrated-device status
CUDA-->>KVCacheManager: Return unified-memory status
KVCacheManager->>KVCacheManager: Set secondary free-memory bytes to zero
Tests->>KvCacheConfig: Validate integrated, discrete, and failure cases
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 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
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. |
b370d78 to
893d3f4
Compare
|
@hchings @pengbowang-nv @nvpohanh can I get a CI run..? this seems good to go |
|
/bot run |
mihai-chiorean
changed the title
… 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>
Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
cb863d9 to
38ab858
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
🧹 Nitpick comments (3)
tests/unittest/llmapi/test_llm_args.py (3)
938-943: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMark the integrated-memory test as CPU-only.
This test replaces
is_device_integratedwith a lambda and performs no CUDA operation. Add@pytest.mark.cpu_only, matching the discrete-memory test.As per path instructions, review test changes as QA coverage.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/llmapi/test_llm_args.py` around lines 938 - 943, Add the `@pytest.mark.cpu_only` decorator to test_KvCacheConfig_unified_memory_forces_host_cache_size_zero, matching the corresponding discrete-memory test while leaving its existing setup and assertions unchanged.Source: Path instructions
938-948: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRename the two added test functions to snake_case.
Rename them to:
test_kv_cache_config_unified_memory_forces_host_cache_size_zerotest_kv_cache_config_discrete_memory_preserves_host_cache_sizeAs per coding guidelines, use snake_case for files, functions, methods, locals, and mutable globals.
Proposed rename
-def test_KvCacheConfig_unified_memory_forces_host_cache_size_zero(monkeypatch): +def test_kv_cache_config_unified_memory_forces_host_cache_size_zero(monkeypatch): -def test_KvCacheConfig_discrete_memory_preserves_host_cache_size(monkeypatch): +def test_kv_cache_config_discrete_memory_preserves_host_cache_size(monkeypatch):🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/llmapi/test_llm_args.py` around lines 938 - 948, Rename the two test functions KvCacheConfig tests to snake_case: test_kv_cache_config_unified_memory_forces_host_cache_size_zero and test_kv_cache_config_discrete_memory_preserves_host_cache_size. Keep their test logic unchanged.Source: Coding guidelines
684-686: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftAdd coverage for the fallback and C++ safeguard.
Modified:
test_KvCacheConfig_declaration. Added:test_KvCacheConfig_unified_memory_forces_host_cache_size_zeroandtest_KvCacheConfig_discrete_memory_preserves_host_cache_size. The file is listed inl0_cpu.ymlandl0_a10.yml. Add coverage for theRuntimeErrorfallback and C++ unified-memory bypass. Runpytest tests/unittest/.Test coverage verdict: needs follow-up.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/llmapi/test_llm_args.py` around lines 684 - 686, Extend the KvCacheConfig tests around test_KvCacheConfig_declaration to cover the RuntimeError fallback path and the C++ unified-memory safeguard, including the expected host cache size behavior. Add focused cases for both scenarios and run pytest tests/unittest/.Sources: Coding guidelines, Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/unittest/llmapi/test_llm_args.py`:
- Around line 938-943: Add the `@pytest.mark.cpu_only` decorator to
test_KvCacheConfig_unified_memory_forces_host_cache_size_zero, matching the
corresponding discrete-memory test while leaving its existing setup and
assertions unchanged.
- Around line 938-948: Rename the two test functions KvCacheConfig tests to
snake_case: test_kv_cache_config_unified_memory_forces_host_cache_size_zero and
test_kv_cache_config_discrete_memory_preserves_host_cache_size. Keep their test
logic unchanged.
- Around line 684-686: Extend the KvCacheConfig tests around
test_KvCacheConfig_declaration to cover the RuntimeError fallback path and the
C++ unified-memory safeguard, including the expected host cache size behavior.
Add focused cases for both scenarios and run pytest tests/unittest/.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d384c951-e5f0-411a-93a6-afa466cb8751
📒 Files selected for processing (4)
cpp/include/tensorrt_llm/common/cudaUtils.hcpp/tensorrt_llm/batch_manager/kvCacheManager.cpptensorrt_llm/llmapi/llm_args.pytests/unittest/llmapi/test_llm_args.py
🚧 Files skipped from review as they are similar to previous changes (1)
- cpp/include/tensorrt_llm/common/cudaUtils.h
Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
|
Addressed the test nits in |
mikeiovine
left a comment
mikeiovine
left a comment
There was a problem hiding this comment.
Stamp on behalf of runtime devs, @NVIDIA/trt-llm-kv-cache-manager-devs to give final approval
…0260806 Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com> # Conflicts: # tests/unittest/llmapi/test_llm_args.py
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
|
Rebased onto current main in 00be09f. The only conflict was in test_llm_args.py; I preserved upstream’s new KV-manager tests and this PR’s unified-memory monkeypatch coverage. Validation: scoped pre-commit passed, and 15 focused KvCacheConfig tests passed on Spark SM121 against the current ARM64 release image. The refreshed API-label check now requires api-compatible because llm_args.py is in the protected path; this changes no public field or signature, and I cannot apply repository labels as an external contributor. @mikeiovine, could you please add api-compatible and re-stamp the main-only merge when convenient? |
|
Current head |
| unified_memory_detected = False | ||
|
|
||
| if unified_memory_detected: | ||
| if self.host_cache_size and self.host_cache_size > 0: |
There was a problem hiding this comment.
On current main this only catches an explicit value: host_cache_size=None still gets the V2 auto host tier, which #18926 made integrated-aware. So on Spark the default keeps a host tier while an explicit size gets none (MemFree 100 GiB, 60 GiB device quota: None gives 20 GiB, 4 GiB gives 0). Could the explicit case be handled next to #18926 in KVCacheManagerV2 (and V1's resource_manager.py sizing) instead of rewriting the config at parse time? The branch no longer merges either: #18538 removed calculateFreeMemBytes, so the C++ guard has nothing left to patch.
|
|
||
| if unified_memory_detected: | ||
| if self.host_cache_size and self.host_cache_size > 0: | ||
| logger.info( |
There was a problem hiding this comment.
This drops a value the user set explicitly, so logger.warning (as disable_periodic_mamba_snapshots_for_conversations does in this class) would make it visible.
|
Thanks @JunyiXu-nv, agreed on both points. Closing this as superseded.
The remaining Spark sizing problem, reclaimable page cache not being counted as available memory, is what #20951 addresses, so I'll follow up there. One gap it doesn't cover: an explicit V2 |
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()withcudaDevAttrIntegrated.host_cache_sizeto zero.QA Engineer Review
test_KvCacheConfig_declaration.test_kv_cache_config_unified_memory_forces_host_cache_size_zerotest_kv_cache_config_discrete_memory_preserves_host_cache_sizetest_kv_cache_config_preserves_host_cache_size_when_detection_failstests/integration/test_lists/.