[None][perf] KVv2: zero-materialization token path + SHA-256 hash batching - #3
Draft
lancelly wants to merge 22 commits into
Draft
[None][perf] KVv2: zero-materialization token path + SHA-256 hash batching#3lancelly wants to merge 22 commits into
lancelly wants to merge 22 commits into
Conversation
Add the C++ translation of KVCacheManagerV2 under cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/ with nanobind bindings, vendored BLAKE3, build integration, and unit tests. Turn tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py into a backend dispatcher (TLLM_KV_CACHE_MANAGER_V2_BACKEND, default: cpp) and add the C++-backend introspection paths in pyexecutor and tests. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com> Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com> Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com> Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
Port of main commit 7d600ec (DSv4 PR-6): a null beam_block page may now also mean a skipped stale block, not only an SWA scratch block, so drop the two TLLM_CHECK_DEBUG(mEnableSwaScratchReuse) checks in the page unlock/release paths and update the comments to match the Python reference implementation. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
… KVCacheManagerV2 Translate two main-branch Python changes in runtime/kv_cache_manager_v2/ into the C++ backend under cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/: Translated from: - 17190cf [feat] add commit min snapshot for SSM reuse (NVIDIA#15752) config: drop HelixConfig/ssm_reuse_interval, add commit_min_snapshot; page: add SsmCommittedPage + UncommittedPage::convertToSsmCommitted; blockRadixTree: unlinkPage(expectedPage) guard + snapshot-length-aware SSM match truncation; kvCache: commit(is_end), _commitBlock(commitSsm,moveSsm), rewritten _snapshotSsmToTreeBlock, new _copyPageToTreeBlock and _snapshotPartialBlockToTree, plus commit_min_snapshot assert relaxations; nanobind: commit(is_end) and config/manager surface updates. - 52fc8f4 [fix] Reserve worst-case SWA slots to avoid single-request deadlock (NVIDIA#15588, nvbugs/6330273) storageManager::computeMinSlotsFromConstraints SWA worst-case floor; lifeCycleRegistry getStaleRange comment. Also in this change: - Mirror the Python Block._release_pages() eager reclaim into C++ Block::releasePages()/removeSubtree so page reclamation does not depend on ~Block() destruction timing. - Add kv::AssertionError + a nanobind translator to Python AssertionError so config validation and the commit history-alignment check match the Python backend's exception type. - Add white-box reuse-tree introspection helpers (attention/swa/ssm life-cycle ids, reuse_match_pages) to the C++ _introspection submodule and _introspection.py, and route the shared tests through them instead of Python-only internals. - Adapt the SWA event-manager tests (small window; deeper load for the migration test to force both pool groups to migrate under the new slot floor). - _core/_kv_cache.py: add the same VIRTUAL_STOP loop-break/partial-snapshot guard the C++ commit() uses, keeping the two backends behaviourally aligned. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
Hash token sequences into radix-tree block keys with SHA-256 in the C++ KVCacheManagerV2, matching the Python backend's hashlib.sha256 so both backends produce byte-identical block keys and KV cache event hashes. Vendor Bitcoin Core's MIT-licensed SHA-256 under 3rdparty/sha256/, reduced to the single-block CSHA256 hasher: a portable scalar core plus runtime dispatch to x86 SHA-NI or ARMv8 crypto extensions, built as plain C++17. - blockRadixTree: incremental digests via CSHA256; SHA256AutoDetect() selects the hardware transform once per process. - eventManager: emit the v2_sha256 / v2_sha256_64 event hash algorithms. - CMake: per-architecture -msha / -march=armv8-a+crypto wiring on the vendored translation units; include root 3rdparty/sha256. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com> (cherry picked from commit f0c86ab)
hostPrefaultChunk() referenced MADV_POPULATE_WRITE unconditionally, which is only defined with glibc >= 2.34 / Linux >= 5.14 headers. On the Rocky8 package-sanity CI images the macro is undeclared, breaking the tensorrt_llm_batch_manager_static compile and cascading into all downstream C++/PyTorch test stages. Guard the fast madvise path with #ifdef MADV_POPULATE_WRITE and fall back to explicitly touching the pages via memset when it is unavailable, matching the existing EINVAL/ENOSYS runtime fallback. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
StorageManager::newGpuSlots and newSlotsForPoolGroup gained defaulted MigrationRecorder/DropRecorder parameters (commit 741ab91, stats API migration), but the static_assert member-function-pointer typedefs in kvCacheManagerV2TypedIndexTest.cpp were not updated. A default argument does not change a function-pointer type, so std::is_same failed and the static_assert hard-failed compilation of the batch_manager C++ unit tests. Append the two recorder parameters to StorageManagerNewGpuSlots and StorageManagerNewSlotsForPoolGroup so the typedefs match the current API. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
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>
…resume on both backends Ports the DSv4 warmup-starvation fix (PR NVIDIA#16269) to KVCacheManagerV2's constraint-based min-slot computation, on both the C++ and pure-Python backends so they stay in parity. Two behavior changes in computeMinSlotsFromConstraints / _compute_min_slots_from_constraints: - Constraints stay feasibility floors even under an explicit initial_pool_ratio (previously they were dropped for min-slots), so a declared batch that needs more than its target pool-group share clamps that share up instead of starving during warmup. - The per-constraint slot floors are scaled by 1/max_util_for_resume (clamped to the binding (0, 1) range) because KvCache::resume rejects any pool group above that utilization, so the floor must leave headroom. max_util_for_resume is threaded from KVCacheManagerConfig through the StorageManager constructor on both backends. The C++ side is a rewrite of Lance Liao's commit 15a90e1 (the constraint-floor half): this branch's computeMinSlotsFromConstraints was reworked with structural SWA/SSM floors after his fork point, and neither his branch's nor this branch's Python _storage_manager.py had the fix, so a mechanical cherry-pick did not apply and the Python backend needed the same change. Update the init-ratio test accordingly: with the fix an infeasible initial_pool_ratio is overridden by the constraint's feasibility floor, so test_constraint_floor_overrides_infeasible_initial_pool_ratio asserts the clamped-up share (both backends produce an identical ratio). 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>
…l init MpiPoolSession.shutdown() is reachable from __del__ (via shutdown_abort) and may run on a session that never completed __init__ -- e.g. after release_exit_joins() marks the pool dead, or in the released-session path exercised by test_proxy_fast_death.py:: test_pool_session_shutdown_never_blocks_after_release, which builds the object with __new__. shutdown() read self.n_workers (log line) and self._wait_shutdown unguarded, raising AttributeError, while the _pool_dead check right above already used getattr defensively. Guard both accesses with getattr, matching the existing _pool_dead style, so a released/partially-constructed session shuts down without crashing. This unbreaks a cross-PR merge-skew regression on main: the unguarded self.n_workers log came from NVIDIA#16456 while the __new__-based test came from NVIDIA#16312; each was green alone but the combination fails deterministically. Also add the missing D205 docstring blank lines the ruff-legacy gate flagged on the touched file. Verified: the whole test_proxy_fast_death.py file now passes (23 passed, was 1 failed). Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
… create_kv_cache
Admitting a long-context request previously round-tripped every token id
through a Python object: get_tokens() materializes a list of PyLongs from
the C++ token vector, Python slices it, and castTokenIterable unboxes each
element back to int64 (~24ns/token each way; several ms per 200k-token
admission, the dominant _schedule cost after the C++ manager landed).
Bypass the materialization entirely for text-only prompts:
- GenericLlmRequest.get_tokens_bytes(beam): export the int32 token vector
as raw bytes, one memcpy
- prepare_context: wrap it in memoryview(...).cast('i'); slicing the view
for the reuse lookup is zero-copy; multimodal requests keep the list path
- castTokenIterable: accept contiguous int32 ('i') and int64 ('q'/'l')
buffers via the buffer protocol and bulk-fill the TokenIdExt vector;
reserve() + PyLong_CheckExact fast path on the generic iterable fallback
- create_kv_cache python wrapper: pack plain-int lists into array('q') so
even the fallback path crosses the binding as a buffer
Measured on GB300 4-ctx/1-gen disagg (DeepSeek-V4-Pro, conc=1000, agentx
workload, ~110k requests per arm, controlled A/B on the same binary):
ctx _schedule wall 8.95 -> 5.07 ms/iter (-43%), per-request token prep
1.20 -> 0.05 ms; e2e +2.1% req/s, TTFT p50 -5.2%, ITL p50 -4.0%.
Behavior equivalence validated over 230k requests (identical reuse/commit
counters vs the list path).
Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
The bulk update() overload wrote each token id as a separate 8-byte CSHA256::Write. Once the transform itself runs on the ARMv8 SHA extensions, the per-Write incremental buffering/bookkeeping dominates the hash cost. Pack plain token ids into a 512-token stack buffer and flush with one Write per 4KB (falling back to per-element update() for multimodal digest tokens). The byte stream is identical, so digests are unchanged - verified with a 1000-trial randomized equivalence test on Grace covering buffer rollover and DigestToken boundaries. The old comment claiming batching only helps Python predates the hw-SHA switch and is dropped. Measured on GB300 4-ctx/1-gen disagg (DeepSeek-V4-Pro, conc=1000, 15.6B hashed tokens per arm): hash cost 9.5 -> 6.5 ns/token (1.47x). Hashing is ~11% of ctx _schedule after the zero-materialization path, so e2e movement is within noise (+0.2% req/s); the win is CPU cycles and a shorter admission tail, not throughput. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
lowsfer
force-pushed
the
kvCacheManagerV2-cpp
branch
7 times, most recently
from
July 21, 2026 15:10
3b6d8ed to
cf45113
Compare
lowsfer
force-pushed
the
kvCacheManagerV2-cpp
branch
22 times, most recently
from
July 28, 2026 07:11
0d02e17 to
4fb042f
Compare
lowsfer
force-pushed
the
kvCacheManagerV2-cpp
branch
4 times, most recently
from
July 29, 2026 05:29
cef027a to
f0852d3
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Two admission-path optimizations for the C++ KVCacheManagerV2, ported from our GB300 disagg production study (DeepSeek-V4-Pro, 4×TP8 ctx + 1×TP16 gen, conc=1000, agentx workload, ~110k requests per arm, controlled same-binary A/B).
Commit 1: zero-materialization token path (+2.1% e2e req/s)
Admitting a long-context request round-trips every token id through a Python object:
get_tokens()materializes a list of PyLongs, Python slices it, andcastTokenIterableunboxes each element back (~24 ns/token each way). At 200k-token prompts this is the dominant_schedulecost — measured 43% of ctx_schedulewall on this workload.GenericLlmRequest.get_tokens_bytes(beam): raw int32 token buffer, one memcpyprepare_context:memoryview(...).cast('i')+ zero-copy slice for the reuse lookup (text-only prompts; multimodal keeps the list path)castTokenIterable: buffer-protocol fast path for contiguous int32/int64 buffers;reserve()+PyLong_CheckExactfast loop on the iterable fallbackarray('q')so the fallback also crosses as a bufferMeasured: ctx
_schedule8.95 → 5.07 ms/iter (−43%), per-request token prep 1.20 → 0.05 ms, admission of >20k-token requests 9.4 → 5.9 ms; e2e +2.1% req/s, TTFT p50 −5.2%, ITL p50 −4.0%. Behavior equivalence validated over 230k requests (identical reuse/commit counters vs the list path).Commit 2: batch token ids into 4KB SHA-256 writes (CPU-side win, e2e-neutral)
The bulk
Hasher::updateoverload writes each token as a separate 8-byteCSHA256::Write; with the transform on ARMv8 SHA extensions, the per-Write bookkeeping dominates. Packing ids into a 512-token stack buffer (oneWriteper 4KB) gives 9.5 → 6.5 ns/token (1.47×) on Grace. Byte stream is identical → digests unchanged (1000-trial randomized equivalence incl. buffer rollover and DigestToken boundaries), so the security invariant documented onBlockKeyis untouched.Full disclosure: hashing is only ~11% of ctx
_scheduleafter commit 1, so e2e movement was within noise (+0.2% req/s) — the win is CPU cycles and a shorter admission tail. Take it or leave it independently of commit 1; the stale "batching only helps Python" comment it removes predates the hw-SHA switch either way.