[lucifer1004]

Description

Fixes #20363.

BatchedPageCopier (#18541) moves KVCacheManagerV2 pages between GPU and the host tier with an SM copy kernel on coherent (NVLink-C2C) links and with cuMemcpyBatchAsync on PCIe. This PR makes the copy engine the default on coherent links as well; the kernel stays available through TLLM_KV_CACHE_MANAGER_V2_COHERENT_COPY_KERNEL=1, and chunked host registration still forces the kernel as before.

Why (details and data in #20363):

• After the SM kernel has accessed registered host memory a few times, every later kernel boundary in the process (including CUDA graph replays) gets slower on our SM107 + NVLink-C2C system: +1.2 ms (~7%) per decode step for Kimi K3 on 8 GPUs, persistent for the rest of the process. The copy engine path in the same binary removes it with token-identical outputs; forcing the kernel back on in that binary brings it back.
• [None][perf] Use the copy engines for large host<->GPU page copies in KVCacheManagerV2 #20311 measured on NVLink-C2C that the copy engines match or slightly beat the kernel on an idle GPU and are 1.2-4x faster when the SMs are busy (the kernel's CTAs wait for resident blocks), and listed this switch as the natural follow-up.
• Cost measured: a fully host-resident 393K-token prefix reload (~13 GB per rank) takes ~0.19 s (~9%) longer TTFT.

Changes:

• batchedPageCopy.cu: computeConfigs() selects the path through detail::selectCopyKernel(coherentLink, chunkedRegistration, kernelOnCoherentLink); the opt-in is read from the environment when the copier is constructed. The selected path is logged at DEBUG.
• batchedPageCopy.h: declares detail::kCoherentCopyKernelEnv and detail::selectCopyKernel(), and updates the class comment.

No behaviour change on PCIe or with chunked host registration. With the copy engine selected, pageIndexLocation() reports kHost, which the cold-page codec and the storage manager already handle (it is the PCIe path today).

Test Coverage

New cases in cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ColdPageCopyTest.cu:

• CopyEngineIsDefaultOnEveryLink, CopyKernelOptInAppliesOnlyToCoherentLinks, ChunkedRegistrationAlwaysSelectsCopyKernel: the selection truth table.
• CoherentCopyKernelEnv/BatchedPageCopierPathTest.SelectsPathAndRoundTripsPages/{Unset,Value0,Value1}: constructs a copier on GPU 0 with the variable unset / 0 / 1, checks kernelEnabled(), both KernelConfig::useKernel and pageIndexLocation() against the expected selection for the detected topology, then offloads 8 pages to mapped pinned host memory and onboards them back through two permutations and checks the bytes. On a coherent system this exercises both paths; on PCIe both instances use the copy engine.

Validation:

• End-to-end (an equivalent patch on our build): SM107 with NVLink-C2C, Kimi K3, 8 GPUs / 2 nodes, KVCM2 C++ backend with host tier: decode step without the persistent +1.2 ms, 99/99 outputs identical; forcing the kernel through the variable reproduces the slowdown.
• The new gtest file compiles with nvcc 13.4 (-std=c++20). On SM107 (coherent link, branch diff applied on a build from an earlier main) it passed 6/6; not run on a PCIe machine yet.
• Not validated on GB200/GB300 (SM100/SM103) or GH200; the change only affects coherent links.
• pre-commit on the changed files: clean.

PR Checklist

• PR description clearly explains what and why.
• PR follows the TRT-LLM coding guidelines to the best of my knowledge.
• Test cases are provided for new code paths.
• No API changes, no new dependencies, no CODEOWNERS changes.
• Please check this after reviewing the above items as appropriate for this PR.

Dev Engineer Review

• BatchedPageCopier now uses the copy engine by default. It uses the kernel when host registration is chunked, or when TLLM_KV_CACHE_MANAGER_V2_COHERENT_COPY_KERNEL is exactly "1" on a coherent link.
• This changes the default host-tier copy path on NVLink-C2C. The kernel remains available as an opt-in. The chunked-registration fallback remains in place.
• Reported validation found identical outputs in the author’s SM107 comparison. GPU execution of the new unit test is pending, and GB200/GB300 and GH200 were not validated. The reported TTFT tradeoff for reloading a fully host-resident prefix is about 0.19 seconds (9%).
• Current review findings and severity counts were not supplied.

QA Engineer Review

• Added kvCacheManagerV2ColdPageCopyTest.cu, a GPU unit test for copy-path selection and page-copy round trips under unset, "0", and "1" environment values. It also tests chunked-registration selection and, for CUDA versions below 12.8, the kernel fallback for multiple index segments.
• CMake registers the test as kvCacheManagerV2ColdPageCopyTest. No integration test-list entry applies to this unit test.
• Coverage verdict: needs follow-up. The new test compiled with nvcc 13.4 (-std=c++20), but GPU execution is pending.

Per-File QA Perspective

• cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/batchedPageCopy.cu — Verify default copy-engine selection, exact "1" opt-in behavior on coherent links, and forced kernel selection with chunked registration. Also verify copies on PCIe and coherent systems.
• cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/batchedPageCopy.h — Verify that pageIndexLocation() matches the selected copy path, since the kernel requires device-resident indices and the copy engine requires host-resident indices.
• cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ColdPageCopyTest.cu — Tests path selection and D2H/H2D page transfers with permuted host slots. It is a unit test, registered through CMake; it is not listed in the integration CI or manual-QA lists.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/TensorRT-LLM/.coderabbit.yaml
• Review profile: CHILL
• Plan: Enterprise
• Run ID: e694cb84-8d5e-4545-b7e8-11648e06b47b

📥 Commits

Reviewing files that changed from the base of the PR and between 15f060c and e6bd742.

📒 Files selected for processing: 3

• cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/batchedPageCopy.cu
• cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/batchedPageCopy.h
• cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ColdPageCopyTest.cu

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

Walkthrough

The copier now uses the copy engine by default. Chunked host registration selects the copy kernel on any link. On coherent links, the kernel is selected only when TLLM_HDC_COHERENT_COPY_KERNEL is set to "1". Tests cover selection and page transfers.

Changes

KV cache page-copy selection

Layer / File(s)	Summary
Copy path selection cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/batchedPageCopy.h, cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/batchedPageCopy.cu	The selection helper chooses the kernel for chunked registration or when the coherent-link override is enabled. Otherwise, the copier selects the copy engine.
Selection and page-copy tests cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ColdPageCopyTest.cu	Tests check selection for coherent and non-coherent links, then exercise D2H and H2D transfers for unset, "0", and "1" environment values.

Priority: ⬆️ High

Estimated code review effort: 3 (Moderate) | ~20 minutes

Change: Bug fix · Severity of issue fixed: Medium

Suggested reviewers: juney-nvidia

Merge Risk: ⚪ Minimal · up to e6bd7

No actionable issue remains from this review. Merge after normal checks; GPU execution of the new tests is still pending.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 15.79% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 19 functions across 3 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Linked Issues check		Issue #20363 requires the copy engine by default on coherent NVLink-C2C links, an environment-variable opt-in for the kernel, and mandatory kernel use with chunked host registration. `computeConfigs()…
Out of Scope Changes check		The changes stay within issue #20363. They modify BatchedPageCopier path selection and documentation, add the required selection helper and environment declaration, and add focused unit tests for se…
Title check		The title clearly identifies the main change: preferring the copy engine for KVCacheManagerV2 host-tier copies on coherent links. It uses the issue identifier and a concise change summary.
Description check		The description includes the required Description, Test Coverage, and PR Checklist sections. It explains the motivation, implementation, performance impact, test cases, validation status, and known pl…

• Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 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.}

[lucifer1004]

Test report (branch diff applied on a build from an earlier main, SM107, single node):

• kvCacheManagerV2ColdPageCopyTest built from this branch: 6/6 passed (3 TEST + 3 TEST_P instances).
• The debug log shows coherentLink=1 on this machine; with TLLM_KV_CACHE_MANAGER_V2_COHERENT_COPY_KERNEL unset or 0 the copier picks the copy engine, with 1 the copy kernel.
• Not run on a PCIe machine.

[lucifer1004]

Cross-referencing two PRs on the same selection in BatchedPageCopier::computeConfigs(): #20426 (copy engines on coherent links as an opt-in, default unchanged, with GB300 measurements) and #20383 (copy engines by default on coherent links, like this PR). The three should probably converge on one; I'd like the maintainers to pick the direction.

What we measured on SM107 with a coherent NVLink-C2C link (Kimi K3, 8 GPUs over 2 nodes, KVCacheManagerV2 with the host tier on, details in #20363):

• With the SM copy kernel, once pages start going to the host tier every decode step becomes slower by a fixed amount for the rest of the process: +7.0% at a 2K prompt, about +4% at 393K, and +13.5% in a configuration with more kernels per step. In the profile no kernel gets slower; the idle gaps between kernels inside the CUDA graph replay grow. With the copy engine there is no such step (noise level), and outputs are token-identical. Forcing the kernel back on in the copy-engine build brings the step back.
• The cost of the copy engine we saw: reloading a fully host-resident 393K-token prefix takes about 9% longer TTFT than with the kernel.

The trade-off as I see it:

• Default on (this PR, perf(kvcache-v2): default to copy engine on coherent NVLink-C2C, adds opt-in SM kernel fallback #20383): coherent deployments that offload to the host tier do not get the decode slowdown. It is silent and persistent, so users are unlikely to find the knob. The price is the slower bulk reload above, and possibly the higher kernel ceiling on wide links (GH200, [None][perf] Speed up burst KVCM2 resize for very long sequences #18541).
• Opt-in ([None][feat] Let BatchedPageCopier prefer the copy engines on coherent links #20426): keeps today's behavior and the wide-link kernel bandwidth, and needs no new GH200 data. But it leaves the decode regression in place by default on systems like ours. We have not identified its mechanism, so we cannot say whether GB300 or GH200 show it.

If you prefer #20426, I'm happy to close this PR and add our SM107 numbers (and a run with its variable set) there. If the default should change, this PR is ready for review. It has been tested on SM107: the new gtest passed 6/6, and both the default path and the opt-in kernel path were exercised.

cc @lowsfer @tzulingk

[nvpohanh]

[by Claude Code] @yizhang-nv Could you review this PR? Thanks!

[nvpohanh]

[by Claude Code] @yizhang-nv Could you review this PR? Thanks!