Repository navigation
[None][perf] BREAKING: Enable zero-copy mode for Mooncake Store - #20108
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
WalkthroughMooncake KV transfers now use direct GPU-pool registration by default. The connector splits registration and transfer ranges at allocation boundaries and reports GPUDirect RDMA registration facts. Host staging remains an optional alternative. Zero-copy loads wait for KV-page allocation readiness. ChangesMooncake KV Pool Registration
Priority: ➖ Normal Estimated code review effort: 4 (Complex) | ~45 minutes Change: Bug fix Sequence Diagram(s)sequenceDiagram
participant KVCacheManagerV2
participant KvCacheConnectorManager
participant KvCacheConnectorWorker
participant MooncakeStoreConnectorWorker
KVCacheManagerV2->>KvCacheConnectorManager: Provide wait_allocation_fences callback
KvCacheConnectorManager->>KvCacheConnectorWorker: Bind callback during metadata handling
MooncakeStoreConnectorWorker->>KvCacheConnectorWorker: Retrieve callback
MooncakeStoreConnectorWorker->>KVCacheManagerV2: Wait for request allocation fences
MooncakeStoreConnectorWorker->>MooncakeStoreConnectorWorker: Continue zero-copy load after wait
Merge Risk: ⚪ Minimal · up to
|
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 56.71% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 164 functions across 20 files. (1 skipped… | Write docstrings for the functions missing them to satisfy the coverage threshold. |
-
Check name Status Explanation Title check The title is concise, specific, and accurately identifies the primary change: enabling zero-copy mode for Mooncake Store. It also includes the required ticket, type, and BREAKING markers. Description check The description explains the problem, solution, registration requirements, fallback behavior, and relevant test coverage. It includes the required sections and acknowledges the breaking change. Test o… Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request.
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 56.71% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 164 functions across 20 files. (1 skipped: 1 unsupported.)
✨ Finishing Touches: 💡 1
-
⚔️ Resolve merge conflicts: 💡
-
- Resolve merge conflict in branch
user/brb/use-dma-buff-main
- Resolve merge conflict in branch
🧪 Generate unit tests (beta)
-
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
|
/bot run --disable-fail-fast |
|
PR_Github #77163 [ run ] triggered by Bot. Commit: |
71127af to
b529dac
Compare
brb-nv
added
the
api-breaking
brb-nv
changed the title
brb-nv
changed the title
brb-nv
changed the title
brb-nv
changed the title
brb-nv
changed the title
|
PR_Github #77163 [ run ] completed with state
|
thorjohnsen
left a comment
thorjohnsen
left a comment
There was a problem hiding this comment.
Mostly looks good, but I have concerns about registering pages interfering with rebalancing.
chuangz0
left a comment
chuangz0
left a comment
There was a problem hiding this comment.
Two things worth fixing before merge (inline).
|
/bot run --disable-fail-fast |
|
PR_Github #78999 [ run ] triggered by Bot. Commit: |
…tion for GPUDirect RDMA A KV cache V2 pool is many cuMemCreate allocations mapped back to back in one virtual range, and GPUDirect RDMA cannot register a range that crosses two of them (EFAULT under nvidia_peermem, EINVAL under dma-buf). - Split each pool at the allocation boundaries cuMemGetAddressRange reports and register it piece by piece; transfer buffers are split the same way. - Unregister the pieces already registered when one fails, and every piece at shutdown. - Restore stage_through_host=False as the default, with the registered-pool path. - Select dma-buf (WITH_NVIDIA_PEERMEM=0) before the first Mooncake client in the process starts when no peer-memory module is loaded and every visible GPU supports dma-buf; refuse later clients that would need a different path. - Add KvCache::waitAllocationFence so zero-copy loads wait only for the busy recycled pages of the requests being loaded, instead of the whole stream. - State the WITH_NVIDIA_PEERMEM and gpuDirectRDMACapable preconditions in the registration error and docs, and log the per-range driver report at debug level. - Warn from PooledPhysMemAllocator when it falls back to memory without gpuDirectRDMACapable. - Add unit tests for the allocation split, rollback, shutdown unregistration, registration-path selection and the allocation fence, and an RDMA round trip over a real KVCacheManagerV2 pool. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
7ca92fe to
029c8a0
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #79023 [ run ] triggered by Bot. Commit: |
|
PR_Github #78999 [ run ] completed with state |
|
PR_Github #79023 [ run ] completed with state
|
… fence Remove KvCache::waitAllocationFence. Connector-served context caches are resumed on a side stream until their batch is reported, then adopted by the execution stream through the cuda_stream setter. Zero-copy loads synchronize the side stream, which holds only the allocation's waits and copies. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
029c8a0 to
74ea7be
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #79086 [ run ] triggered by Bot. Commit: |
|
PR_Github #79086 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #79367 [ run ] triggered by Bot. Commit: |
|
PR_Github #79367 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #79394 [ run ] triggered by Bot. Commit: |
|
PR_Github #79394 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #79522 [ run ] triggered by Bot. Commit: |
|
PR_Github #79522 [ run ] completed with state |
Description
Zero-copy mode (
stage_through_host=False) failed to register KV cache V2 pools with Mooncake over RDMA.This PR fixes the registration and restores zero-copy as the default.
Why
cuMemCreateallocations mapped back to back in one virtual range.EFAULTunder nvidia_peermem andEINVALunder dma-buf (reported by chuangz0 in [TRTLLM-12891][feat] Integrate Mooncake connector into KVCMv2 #19171).What changed
cuMemGetAddressRange, and registered piece by piece.stage_through_hostdefaults toFalseagain, with the registered-pool path restored.WITH_NVIDIA_PEERMEM=0: dma-buf, which needs neither of the below.gpuDirectRDMACapable=1.gpuDirectRDMACapable.TRTLLM_MOONCAKE_STORE_DEBUG_REGISTRATION=1to log what the driver reports for each range.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)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.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.Dev Engineer Review
Mooncake now defaults to zero-copy transfers. The worker splits KV cache registrations at CUDA allocation boundaries, rolls back partial registrations on failure, and unregisters each range at shutdown. Host staging remains available with
stage_through_host=True.The connector selects the registration path before Mooncake’s first client starts. It preserves an explicit
WITH_NVIDIA_PEERMEMsetting and rejects a later RDMA client that requires a different path. Zero-copy loads require an allocation fence and wait for the requested pages before the NIC writes to them. The KV cache manager tracks work on pages handed out while busy.Verify both RDMA registration paths, allocation-fence ordering, cleanup after failures, and staging fallback on supported hardware. The supplied CI reports show
PR_Githubsuccess andL0_MergeRequest_PRfailure for two commits. The failing tests are not identified, so test and merge status remain unresolved.QA Engineer Review
Changed unit tests cover staging configuration and buffer behavior; allocation-boundary range and page splitting; GPU registration-path selection and diagnosis; registration rollback and shutdown cleanup; zero-copy and staged transfer ordering; allocation fences; and real-pool TCP and RDMA round trips. The integration tests cover Mooncake prefix reuse across engines and tensor-parallel ranks, and chunked prefill with a small cache.
The changed integration test functions are listed in
tests/integration/test_lists/test-db/l0_a10.ymlandtests/integration/test_lists/test-db/l0_dgx_h100.yml. No corresponding Mooncake entry appears in the manual-QA lists. The objectives report twoPR_Githubsuccesses and twoL0_MergeRequest_PRfailures, but do not identify the failed tests. Coverage verdict: needs follow-up until the CI failures and hardware-specific RDMA paths are verified.Per-File QA Perspective
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/cudaVirtMem.cpp— Verify the warning identifies allocations withoutgpuDirectRDMACapableand describes registration implications accurately.docs/source/features/kv-cache-connector.md— Verify the registration requirements, defaults, diagnostics, and staging fallback match runtime behavior.tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/addressing.py— Verify registration ranges and page buffers split at allocation boundaries, including unresolved lookup results.tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py— Verify thestage_through_host=Falsedefault and explicit staging configuration.tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/gpudirect.py— Verify registration-path selection and diagnostics for available and unavailable driver information.tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py— Verify the documented fallback conditions match worker behavior.tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py— Verify staging memory counts toward the budget only when staging is enabled.tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py— Verify per-range registration, failure rollback, shutdown cleanup, and allocation-fence waits for zero-copy loads.tensorrt_llm/llmapi/llm_args.py— Verify the public staging default and capacity-role behavior match connector behavior.tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py— Verify the manager’s allocation-fence callback is bound to the worker and cleared after metadata handling.tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py— Verify allocation-fence waits cover the requested cache entries and order earlier page fills before external writes.cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp— Verify allocation-fence events cover pages handed out while busy and are synchronized before return.cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h— Verify fence event tracking and pruning remain protected by the API lock.cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp— Verify the Python binding releases the GIL while waiting on the allocation fence.tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi— Verify the stub documents the bound allocation-fence API and its wait behavior.tests/integration/defs/llmapi/test_llm_api_connector_mooncake.py— Covers cross-engine and tensor-parallel prefix reuse and chunked prefill. The changed test functions are listed in the CI test lists above; no manual-QA listing was found.tests/unittest/_torch/executor/test_mooncake_store_common.py— Covers configuration defaults and staging-pool registration, bounds, and data movement. Unit test; integration test-list membership does not apply.tests/unittest/_torch/executor/test_mooncake_store_connector.py— Covers allocation splitting, path diagnosis and selection, registration cleanup, debug logging, and staged and zero-copy ordering. Unit test; integration test-list membership does not apply.tests/unittest/_torch/executor/test_mooncake_store_real_pool.py— Covers TCP and RDMA pool round trips, staged and zero-copy transfers, allocation boundaries, and shutdown cleanup. Unit test; integration test-list membership does not apply.tests/unittest/_torch/executor/test_kv_connector_reservations.py— Covers allocation-fence propagation from the cache manager to the worker. Unit test; integration test-list membership does not apply.tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py— Covers a fence wait for a recycled slot with unfinished work from its prior owner. Unit test; integration test-list membership does not apply.