[None][feat] KVCacheManagerV2 C++ translation by lowsfer · Pull Request #14047 · NVIDIA/TensorRT-LLM

lowsfer · 2026-05-12T11:44:10Z

Dev Engineer Review

Review the C++ KVCacheManagerV2 translation, including storage tiers, eviction, radix-tree reuse, CUDA virtual memory, copy/event handling, statistics, lifecycle management, and Python bindings for correctness, performance, API consistency, error handling, and regressions.
Review SHA-256 vendoring, architecture-specific dispatch, CMake integration, pre-commit exclusions, backend selection, and public API compatibility.
Validate configuration checks, memory/resource cleanup, cross-backend behavior, and the MPI shutdown-safety change.
CI follow-up is required: recent pipeline runs failed despite successful helper jobs; rerun with fail-fast disabled, including PerfSanity.

QA Engineer Review

Test code was added or updated in:
- kvCacheManagerV2HostMemTest.cpp
- kvCacheManagerV2StatsTest.cpp
- kvCacheManagerV2TypedIndexTest.cpp
- KV-cache event manager, manager behavior, salting, and stats API Python tests.
Coverage includes host-memory behavior, typed-index contracts, statistics and pending accounting, event parity, reuse introspection, salting, native C++ bindings, and request-ID handling.
No tests/integration/test_lists/, test-db/, or qa/ coverage changes are listed. Test functions should be mapped to CI or manual-QA entries.
Verdict: needs follow-up.

hyukn · 2026-05-15T09:51:10Z

/bot run --disable-fail-fast --stage-list "PerfSanity"

ReuseScope.salt (low 8 bytes of a sha256 digest) and lora_id can occupy the full unsigned 64-bit range, but the C++ ReuseScope stored them as int64_t and the nanobind casters used std::optional<int64_t>. Passing a salt >= 2**63 from Python raised TypeError at ReuseScope(salt=...) on the C++ backend. Type them as uint64 end-to-end, mirroring the RequestIdType approach already used for request ids (rather than reinterpret-casting at the nanobind boundary): - ReuseScope::loraId -> std::optional<LoraTaskIdType>, salt -> std::optional<std::uint64_t>; add a LoraTaskIdType alias in the V2 common.h next to RequestIdType. - Propagate uint64 through the V1 event-hash path (V1RootAttrs, hashV1BlockKey) so no implicit narrowing occurs. - Update the nanobind ReuseScope binding, optionalIntToObject, castOptionalIntAttr, and the int lora_task_id fallback to uint64. Block-key and V1 event hashes are unchanged: toBytes/hash mixing serialize the raw little-endian 8 bytes, which are bit-identical for a given value whether stored as int64 or uint64, and the Python side already serializes with signed=False. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

…key-gen helpers as public API Two related cleanups to KVCacheManagerV2's reuse-key handling. C++ hashing: - Add Hasher::update(ReuseScope) that serializes the scope straight into the SHA-256 state via a shared emitReuseScopeBytes() emitter, removing the std::vector allocation the Hasher(ReuseScope) constructor incurred on the key-generation hot path. - Drop ReuseScope::toBytes() and its nanobind `to_bytes` binding. Its only consumer was a unit test; production hashing never used it (a C++ ReuseScope is hashed internally and never round-tripped through Python). Public key-gen API: - Move the pure-Python reference helpers (gen_multimodal_cache_key_tokens, Hasher, reuse_scope_to_bytes, sequence_to_blockchain_keys) into a new backend-neutral module _cache_key.py and export gen_multimodal_cache_key_tokens and sequence_to_blockchain_keys from the package on both backends (added to __all__ and the .pyi stub). - _block_radix_tree.py re-exports from _cache_key; ReuseScope.to_bytes() delegates to reuse_scope_to_bytes(). Removes the cpp dispatcher branch's duplicate gen_multimodal_cache_key_tokens. - reuse_scope_to_bytes() serializes from the scope's fields, so it is backend-neutral (no dependency on a .to_bytes() method) and byte-identical to the C++ emitReuseScopeBytes() and the Python signed=False layout. Tests: - Import sequence_to_blockchain_keys from the public API in test_router and test_kv_cache_salting. - Rewrite the salting test to assert distinct scopes seed distinct blockchain keys (digests) instead of asserting on raw to_bytes output. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

The nanobind module registers its own OutOfPagesError exception class, so 'except OutOfPagesError' clauses importing the pure-Python class from kv_cache_manager_v2._exceptions never match on the C++ backend and the guarded allocation paths crash instead of backing off. The package __init__ already re-exports the backend-appropriate class; import that. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> (cherry picked from commit fdbe4f9)

release_index_slot passes None to clear the DSA index buffer, but the buf argument rejected None at the nanobind layer (missing .none() annotation) even though the lambda body already handles is_none(). On the C++ backend every context worker died with 'incompatible function arguments ... NoneType' at the first index-slot release. Split from Lance Liao's commit 15a90e1 (the constraint-floor half of that change lands in a separate commit). Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

Port the Python-side fix from PR NVIDIA#16484 (nvbugs/6438658) to the C++ StorageManager so both backends behave identically: thread max_util_for_resume from KVCacheManagerConfig through the StorageManager constructor, and scale each constraint-derived min-slot floor by 1/max_util_for_resume in computeMinSlotsFromConstraints so a declared batch stays below the utilization gate checked by KvCache::resume. This mirrors main's behavior exactly: an explicit initial_pool_ratio still overrides constraints for min-slots (the gating is kept), and the scale factor is applied unconditionally (no clamping). Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

REVIEW NOTE: this is the cpp-branch behavior delta on top of the previous commit (which mirrors main's PR NVIDIA#16484). It is isolated so it can be reviewed or dropped independently. If we decide to match main exactly, revert just this commit. Diverges from main's resume-utilization fix in two ways, on both the Python and C++ backends: - Constraints stay feasibility floors even under an explicit initial_pool_ratio (main drops them): a declared batch that needs more than its target pool-group share clamps that share up instead of starving during warmup. Follows Lanyu's C++ commit 15a90e1 (which cited PR NVIDIA#16269). - Clamp max_util_for_resume to the binding (0, 1) range before scaling, so a degenerate 0 can't divide-by-zero and a value >= 1 can't shrink floors. Renames test_initial_pool_ratio_overrides_typical_step_and_constraints to test_constraint_floor_overrides_infeasible_initial_pool_ratio to assert the new behavior (main's test_constraint_reserves_resume_headroom is kept as-is). Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

The license_checker rejected pendingStats.h and stats.h because their boilerplate wrapped as 'permissions and limitations under / the License' instead of the canonical 'permissions and / limitations under the License' used by every passing C++ source file. Fix the wrap in the two flagged headers plus the sibling kvCacheManagerV2StatsTest.cpp (same style). Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

test_kv_pool_rebalance imported OutOfPagesError from the ._exceptions submodule (pure-python class), but py_executor catches the dispatcher package's OutOfPagesError, which is _cpp.OutOfPagesError under the C++ backend. The test raised the python class while the code caught the C++ class, so the expected-failure path was not swallowed and the test failed under the cpp backend (it passed under the python backend where both are the same class). Import the backend-aware symbol so the class matches the one py_executor catches, mirroring the py_executor import. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

…for unregistered role get_index_k_buffer returns None when a layer has no INDEX_KEY buffer registered (dense layers, base V2 manager). It detected this by catching KeyError from the role probe. That holds on the python backend (dict miss -> KeyError), but the C++ backend's getBufferAttr throws std::out_of_range, which nanobind maps to IndexError, so the guard missed it and the exception escaped -> the None-path contract broke and TestIndexKeyBufferAccessor::test_accessor_returns_none_on_dense_layer and ::test_accessor_returns_none_on_base_v2_default_manager failed under the C++ backend on H100/B300. Catch (KeyError, IndexError) so 'role not registered -> None' holds on either backend. Verified on H100 with TLLM_KV_CACHE_MANAGER_V2_BACKEND=cpp: the full TestIndexKeyBufferAccessor class passes (was 2 failing). Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

The radix tree is a globally shared cross-request/cross-tenant cache index whose prefix matching is decided purely by digest equality (no token re-verification), and its hashed input (tokens, user-supplied cache_salt, multimodal bytes) is attacker-influenceable. Document, on both the C++ (blockRadixTree Hasher/BlockKey) and Python (_cache_key Hasher) sides plus the kvcache doc, that the block-key hash must stay cryptographic, collision-resistant, and >=256-bit (SHA-256), and must not be swapped for a non-cryptographic hash or truncated without adding a token-content equality check on match. Comment/doc-only. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

Port PlannedDropHandle and plan_committed_block_drop() to the C++ backend, mirroring the Python implementation added in fd9c32f. - CommittedPage gains a planned_drop_count field. - PlannedDropHandle holds weak references to committed pages, bumps their planned-drop count on construction, and on drop() decrements each live page's count, removing an already-droppable/queued page from eviction tracking on its final plan. Double drop raises ValueError. - KvCache::planCommittedBlockDrop() collects the SWA in-window committed pages for each attention life cycle (excluding full-attention and sink blocks), returning nullptr if any required page is unavailable. - nanobind exposes PlannedDropHandle and _KVCache.plan_committed_block_drop; the dispatcher now resolves PlannedDropHandle from the C++ module. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

…ager-devs The new implementation directory cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/ is not matched by the existing kvCache*/blockKey* globs (a CODEOWNERS '*' does not cross '/'), so it fell through to trt-llm-runtime-devs. Add an explicit directory entry so the KV cache manager team owns the C++ v2 core. The nanobind binding, *Utils, and unit tests are already covered by existing globs. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

Fixes verified-valid issues from PR NVIDIA#14047 review: C++ (cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2): - config.cpp: guard tokensPerBlockOverride against 0 (avoid % by zero UB) - storage/core.cpp: assert non-empty slotSizeList before max_element; check cuCtxGetDevice/cuDeviceTotalMem via cuCheck; guard log2(0) for quotas below 1 GiB - copyEngine.cpp: short-circuit empty/zero-byte transfers (two-hop div-by-zero) - utils/math.h: fix overlap() doc comment; fix DynamicBitset::resize to keep numSetBits accurate and clear stale bits on shrink - utils/cudaEvent.h: increment SimplePool outstanding count only after a successful create (no leak if create throws) - CMakeLists.txt: bump copyright year Python: - _utils.py: mirror the DynamicBitset.resize fix - _storage_manager.py: validate max_util_for_resume via ValueError not assert - pyexecutor/kv_cache_manager_v2.py: narrow the INDEX_KEY try block; use {role!s} Docs: - kvcache.md: distinguish SHA-256 digest size from collision resistance Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

…es.replace-compatible The DeepSeek-V4 cache manager (_build_cache_config) and our own host-tier fallback call dataclasses.replace() on the KVCacheManagerConfig. On the Python backend this was a @DataClass; the C++ translation exposes a nanobind class instead, so dataclasses.replace() raised 'TypeError: replace() should be called on dataclass instances', failing all DeepSeek-V4 cache-manager unit tests (compressor_module, cache_manager, indices_transform, sparse_mla, kv_transfer). dataclasses.replace() is a free function keyed on __dataclass_fields__: it reads each field via getattr and rebuilds via cls(**fields). The C++ binding already has a full keyword __init__ and readable fields, so advertise the dataclass field set on the binding. The read-only enable_swa_scratch_reuse property is excluded (not a constructor field), matching the Python dataclass. Python-only change; no C++ rebuild required. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

…c bindings The Python KVCacheDesc/BatchDesc were @DataClass(slots=True, frozen=True), so they compared by value. The C++ translation replaced them with plain structs whose nanobind bindings exposed no __eq__, so Python equality fell back to identity and always returned False for distinct objects. This broke tests/unittest/_torch/executor/test_kv_cache_manager_v2.py (test_builds_warmup_constraints, test_avg_seq_len_updates_typical_step, test_extra_tokens_are_in_context_capacity), which assert config.typical_step / config.constraints against expected BatchDesc values. Add value operator== (and !=) to the C++ structs and bind __eq__ + __repr__ so the bindings match the original dataclass semantics; __repr__ also gives readable assertion diffs. Nothing hashes these types, so __hash__ is not added (BatchDesc holds a list and was not reliably hashable in Python). Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

…offset The NVFP4 block-scale pool pointer in _prepare_page_table_tensor must share the SAME origin convention as the KEY pool pointer, so that the 'block_scale_offset == offset' invariant holds on both backends. The KEY pointer is a (representative-layer, 0-offset) origin against which each layer's kv_cache_pool_mapping offset is resolved by _kv_pool_mapping_offset(). That offset function is overridable: the base manager makes it relative to the representative layer (offset(rep) == 0), while address-ranked sparse subclasses (MiniMax-M3 / DeepSeek-V4) make it the layer's rank in physical-address order (offset(rep) may be non-zero, since layer_grouping iteration order is NOT a V2 API contract and the C++ unordered_map need not return the pool-base layer first). Compute the block-scale origin the same way: mirror the KEY base for the representative layer, then shift it back by that layer's offset(rep) * scale_stride. For the base manager this is just the rep layer's scale base; for address-ranked subclasses the shift lands the origin on the pool's slot-0 scale address. Either way block_scale_offset reproduces the same per-layer offset() as KEY, with no dependency on layer_grouping order. This supersedes the earlier pool_group_descs approach, which fixed the address-ranked subclasses but broke the base manager: it pinned the scale base to the pool's slot-0 address while KEY stayed relative to a non-slot-0 representative layer, so the two origins disagreed and test_kv_cache_v2_extra_buffers::test_default_hook_keeps_nvfp4_scale_buffers failed. The now-unused _pool_base_addr_by_group_role() helper is removed. Verified on 1x H100: test_kv_cache_v2_extra_buffers (15), test_per_layer_head_dim (8), test_minimax_m3_kv_transfer (55) and the full executor KV-cache suite (342) all pass. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

Signed-off-by: Tyler Burt <195370667+tburt-nv@users.noreply.github.com>

…ude() kv_cache_manager_v2/CMakeLists.txt is include()'d from batch_manager/ CMakeLists.txt, so CMAKE_CURRENT_SOURCE_DIR is batch_manager/, not this file's own directory. SHA256_DIR used ../../common/sha256, which resolved to cpp/common/sha256 (nonexistent) and failed CMake generation on both x86_64 and SBSA builds. Use ../common/sha256 (one level up from batch_manager/) so it resolves to cpp/tensorrt_llm/common/sha256, consistent with the file's existing CMAKE_CURRENT_SOURCE_DIR=batch_manager convention (SRCS listed as kv_cache_manager_v2/*.cpp; KV_CACHE_MANAGER_V2_INCLUDE_DIR). Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

The Release-Check stage runs `license_checker -config jenkins/license_cpp.json`, which requires an Apache header for all files under tensorrt_llm/. The newly added SHA-256 implementation under cpp/tensorrt_llm/common/sha256/ is vendored from Bitcoin Core and carries its upstream MIT license header, so the checker reports "license not exist" (errcount: 6) and fails the pipeline (fail-fast then aborts the downstream GPU test stages). Add the six vendored files to the skip map, matching the existing convention for external / dual-licensed third-party sources (selectiveScan.h, flashMLA, deep_gemm, causalConv1d, vec_dtypes). Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

longlee0622 · 2026-07-28T07:09:26Z

/bot run --disable-fail-fast

lowsfer · 2026-07-28T07:15:25Z

/bot run

tensorrt-cicd · 2026-07-28T07:15:51Z

PR_Github #62144 [ run ] triggered by Bot. Commit: 4fb042f Link to invocation

tensorrt-cicd · 2026-07-28T07:21:06Z

PR_Github #62146 [ run ] triggered by Bot. Commit: 4fb042f Link to invocation

tensorrt-cicd · 2026-07-28T07:25:05Z

PR_Github #62144 [ run ] completed with state ABORTED. Commit: 4fb042f

Link to invocation

tongyuantongyu

LGTM for the runtime modifications that's not in KV Cache manager source files. I'm trusting trt-llm-kv-cache-manager-devs for them.

Shixiaowei02

LGTM for the dis-agg related changes.

tensorrt-cicd · 2026-07-28T08:36:31Z

PR_Github #62146 [ run ] completed with state SUCCESS. Commit: 4fb042f
/LLM/main/L0_MergeRequest_PR pipeline #50321 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

lowsfer · 2026-07-28T08:41:12Z

/bot run

github-actions Bot assigned lowsfer May 12, 2026

lowsfer force-pushed the kvCacheManagerV2-cpp branch 2 times, most recently from 5bd984d to 456cfa8 Compare May 26, 2026 08:44

lowsfer force-pushed the kvCacheManagerV2-cpp branch 3 times, most recently from 4f4dd8d to 0802c6c Compare June 15, 2026 10:25

lowsfer force-pushed the kvCacheManagerV2-cpp branch 6 times, most recently from 4c24d8a to 44c797c Compare June 18, 2026 09:30

lowsfer changed the title ~~Draft: Kv cache manager v2 cpp~~ [None][feat] KVCacheManagerV2 C++ translation Jun 19, 2026

lowsfer force-pushed the kvCacheManagerV2-cpp branch 5 times, most recently from b4592c1 to 5a1fa77 Compare July 10, 2026 10:10

lowsfer added the api-compatible Accepted LLM API contract change that is backwards-compatible label Jul 10, 2026

juney-nvidia requested a review from thorjohnsen July 14, 2026 06:37

lowsfer force-pushed the kvCacheManagerV2-cpp branch 3 times, most recently from fd7d167 to 99bfc50 Compare July 20, 2026 03:03

lowsfer marked this pull request as ready for review July 20, 2026 05:16

lowsfer requested review from a team as code owners July 20, 2026 05:16

lowsfer and others added 19 commits July 28, 2026 07:07

tweak copyright headers

b08efc4

Signed-off-by: Tyler Burt <195370667+tburt-nv@users.noreply.github.com>

tongyuantongyu approved these changes Jul 28, 2026

View reviewed changes

Shixiaowei02 approved these changes Jul 28, 2026

View reviewed changes

eopXD mentioned this pull request Jul 28, 2026

[None][feat] Opt GPT-OSS in to KV cache manager V2 by default lowsfer/TensorRT-LLM#5

Closed

1 task

eopXD mentioned this pull request Jul 28, 2026

[None][feat] Opt GPT-OSS in to KV cache manager V2 by default #16942

Draft

1 task

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[None][feat] KVCacheManagerV2 C++ translation - #14047

[None][feat] KVCacheManagerV2 C++ translation#14047
lowsfer wants to merge 30 commits into
NVIDIA:mainfrom
lowsfer:kvCacheManagerV2-cpp

lowsfer commented May 12, 2026 •

edited by coderabbitai Bot

Loading

Uh oh!

hyukn commented May 15, 2026

Uh oh!

longlee0622 commented Jul 28, 2026

Uh oh!

lowsfer commented Jul 28, 2026

Uh oh!

tensorrt-cicd commented Jul 28, 2026

Uh oh!

tensorrt-cicd commented Jul 28, 2026

Uh oh!

tensorrt-cicd commented Jul 28, 2026

Uh oh!

tongyuantongyu left a comment

Uh oh!

Shixiaowei02 left a comment

Uh oh!

tensorrt-cicd commented Jul 28, 2026

Uh oh!

lowsfer commented Jul 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

11 participants

Uh oh!

Conversation

lowsfer commented May 12, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Dev Engineer Review

QA Engineer Review

Uh oh!

hyukn commented May 15, 2026

Uh oh!

longlee0622 commented Jul 28, 2026

Uh oh!

lowsfer commented Jul 28, 2026

Uh oh!

tensorrt-cicd commented Jul 28, 2026

Uh oh!

tensorrt-cicd commented Jul 28, 2026

Uh oh!

tensorrt-cicd commented Jul 28, 2026

Uh oh!

tongyuantongyu left a comment

Choose a reason for hiding this comment

Uh oh!

Shixiaowei02 left a comment

Choose a reason for hiding this comment

Uh oh!

tensorrt-cicd commented Jul 28, 2026

Uh oh!

lowsfer commented Jul 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

11 participants

lowsfer commented May 12, 2026 •

edited by coderabbitai Bot

Loading