diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index d2f479d25b9d..1a01b72572c0 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -210,6 +210,7 @@ /cpp/tensorrt_llm/batch_manager/blockKey* @NVIDIA/trt-llm-kv-cache-manager-devs /cpp/tensorrt_llm/batch_manager/evictionPolicy* @NVIDIA/trt-llm-kv-cache-manager-devs /cpp/tensorrt_llm/batch_manager/kvCache* @NVIDIA/trt-llm-kv-cache-manager-devs +/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2 @NVIDIA/trt-llm-kv-cache-manager-devs /cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager* @NVIDIA/trt-llm-kv-cache-manager-devs /cpp/tests/unit_tests/batch_manager/blockKey* @NVIDIA/trt-llm-kv-cache-manager-devs /cpp/tests/unit_tests/batch_manager/evictionPolicy* @NVIDIA/trt-llm-kv-cache-manager-devs diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 8de0dc5489ab..7c621df88804 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1559,7 +1559,7 @@ static-analysis-files: &static_analysis_files | )$ # Global exclude: vendored code + trtllm-gen FMHA artifacts (cubin pointers, export headers, cuda_ptx) -exclude: '(^triton_kernels/|trtllmGenKernels/fmha/cubin/kernelMetaInfo\.h$|cubin\.cpp$|cubin\.h$|trtllmGenKernels/fmha/trtllmGen_fmha_export/|trtllmGenKernels/fmha/cuda_ptx/)' +exclude: '(^cpp/tensorrt_llm/common/sha256/|^triton_kernels/|trtllmGenKernels/fmha/cubin/kernelMetaInfo\.h$|cubin\.cpp$|cubin\.h$|trtllmGenKernels/fmha/trtllmGen_fmha_export/|trtllmGenKernels/fmha/cuda_ptx/)' default_install_hook_types: [pre-commit, commit-msg] repos: diff --git a/LICENSE b/LICENSE index ed5ca3b261ba..bf9a933a6697 100644 --- a/LICENSE +++ b/LICENSE @@ -1,4 +1,4 @@ -Copyright (c) 2011-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +Copyright (c) 2011-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. This project is licensed under the Apache 2.0 license, whose full license text is available below. @@ -12,6 +12,14 @@ file headers for specific copyright and license information. Below is a list of other projects that have portions contained by this project: +-------------------------------------------------------------------------------- +Bitcoin Core +-------------------------------------------------------------------------------- +Original Source: https://github.com/bitcoin/bitcoin +Copyright (c) 2009-2026 The Bitcoin Core developers +Copyright (c) 2009-2026 Bitcoin Developers +Licensed under the MIT License + -------------------------------------------------------------------------------- causal-conv1d -------------------------------------------------------------------------------- diff --git a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt index c7fea62c023a..f61e58e16b28 100644 --- a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt +++ b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt @@ -76,12 +76,18 @@ else() # Windows set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /W4") endif() +# Include kv_cache_manager_v2 C++ sources. +include(${CMAKE_CURRENT_SOURCE_DIR}/kv_cache_manager_v2/CMakeLists.txt) +list(APPEND SRCS ${KV_CACHE_MANAGER_V2_SRCS}) + add_library(${BATCH_MANAGER_STATIC_TARGET} STATIC ${SRCS}) target_include_directories( ${BATCH_MANAGER_STATIC_TARGET} PUBLIC ${xgrammar_source_dir}/3rdparty/picojson ${xgrammar_source_dir}/3rdparty/dlpack/include - ${xgrammar_source_dir}/include) + ${xgrammar_source_dir}/include + ${KV_CACHE_MANAGER_V2_INCLUDE_DIR} + ${PROJECT_SOURCE_DIR}/tensorrt_llm/common/sha256) set_target_properties( ${BATCH_MANAGER_STATIC_TARGET} diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cu b/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cu index a2778e169154..b1855ab43b64 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cu +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cu @@ -119,7 +119,7 @@ CUresult launchBatchedCopyImpl( } uint32_t const nbSplits = lowBandwidth ? 1 : divUp(nbBytes, grainBytes * ctaSize * 2); void* args[] = {(void*) pTasks, (void*) &nbBytes}; - static CUkernel const kernel = [] -> CUkernel + static CUkernel const kernel = []() -> CUkernel { cudaKernel_t kernel = nullptr; TLLM_CUDA_CHECK(cudaGetKernel(&kernel, reinterpret_cast(&batchedCopy))); diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h b/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h index e32c727d0d81..5fa88a78468e 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h @@ -17,6 +17,7 @@ #pragma once +#include "kv_cache_manager_v2/common.h" #include "tensorrt_llm/batch_manager/llmRequest.h" #include "tensorrt_llm/kernels/kvCacheIndex.h" #include "tensorrt_llm/runtime/iBuffer.h" @@ -34,13 +35,7 @@ using ITensor = tensorrt_llm::runtime::ITensor; namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 { -struct DiskAddress -{ - int fd; - ssize_t pos; -}; - -using MemAddress = std::uintptr_t; +// DiskAddress and MemAddress are defined in kv_cache_manager_v2/common.h (included above). // Please make sure to align with the definition in tensorrt_llm/runtime/kv_cache_manager_v2/_common.py constexpr tk::KVCacheIndex::UnderlyingType BAD_PAGE_INDEX = -1; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/CMakeLists.txt b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/CMakeLists.txt new file mode 100644 index 000000000000..163ad82ffd3f --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/CMakeLists.txt @@ -0,0 +1,65 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may not +# use this file except in compliance with the License. You may obtain a copy of +# the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +# License for the specific language governing permissions and limitations under +# the License. + +# Sources for the C++ KVCacheManagerV2 implementation. These are added to the +# tensorrt_llm_batch_manager_static target by the parent CMakeLists.txt. + +# Bitcoin Core SHA-256 (vendored, MIT; see common/sha256/README.md) — needed by +# blockRadixTree.cpp. Slimmed to the single-block CSHA256 path: portable scalar +# core + runtime dispatch (sha256.cpp) plus a per-architecture hardware +# transform. Plain C++17; the include root common/sha256 is added on the +# batch_manager target so "sha256.h" resolves. +set(SHA256_DIR ${CMAKE_CURRENT_SOURCE_DIR}/../common/sha256) +set(SHA256_SRCS ${SHA256_DIR}/sha256.cpp) +# x86/x86-64: Intel SHA-NI transform. +if(CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|AMD64|i[3-6]86") + set(SHA256_X86_SHANI ${SHA256_DIR}/sha256_x86_shani.cpp) + set_source_files_properties( + ${SHA256_X86_SHANI} PROPERTIES COMPILE_OPTIONS "-msse4.1;-msse4;-msha") + list(APPEND SHA256_SRCS ${SHA256_X86_SHANI}) +endif() +# AArch64: ARMv8 crypto-extension (SHA-256) transform. +if(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64|arm64|ARM64") + set(SHA256_ARM_SHANI ${SHA256_DIR}/sha256_arm_shani.cpp) + set_source_files_properties( + ${SHA256_ARM_SHANI} PROPERTIES COMPILE_OPTIONS "-march=armv8-a+crypto") + list(APPEND SHA256_SRCS ${SHA256_ARM_SHANI}) +endif() + +set(KV_CACHE_MANAGER_V2_SRCS + kv_cache_manager_v2/common.cpp + kv_cache_manager_v2/config.cpp + kv_cache_manager_v2/lifeCycleRegistry.cpp + kv_cache_manager_v2/utils/cudaEvent.cpp + kv_cache_manager_v2/utils/hostMem.cpp + kv_cache_manager_v2/cudaVirtMem.cpp + kv_cache_manager_v2/storage/config.cpp + kv_cache_manager_v2/storage/core.cpp + kv_cache_manager_v2/evictionController.cpp + kv_cache_manager_v2/copyEngine.cpp + kv_cache_manager_v2/blockRadixTree.cpp + kv_cache_manager_v2/eventManager.cpp + kv_cache_manager_v2/page.cpp + kv_cache_manager_v2/storageManager.cpp + kv_cache_manager_v2/introspection.cpp + kv_cache_manager_v2/kvCache.cpp + kv_cache_manager_v2/kvCacheManager.cpp + ${SHA256_SRCS}) + +# Headers and sources are co-located; the include root is batch_manager/ so that +# #include "kv_cache_manager_v2/common.h" resolves correctly. NOTE: this file is +# include()'d (not add_subdirectory()'d), so CMAKE_CURRENT_SOURCE_DIR is already +# batch_manager/ — no PARENT_SCOPE needed. +set(KV_CACHE_MANAGER_V2_INCLUDE_DIR ${CMAKE_CURRENT_SOURCE_DIR}) diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp new file mode 100644 index 000000000000..212161acfda0 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp @@ -0,0 +1,880 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "kv_cache_manager_v2/blockRadixTree.h" +#include "kv_cache_manager_v2/common.h" +#include "kv_cache_manager_v2/page.h" +#include "kv_cache_manager_v2/storageManager.h" +#include "kv_cache_manager_v2/utils/math.h" + +#include "sha256.h" + +#include "tensorrt_llm/common/assert.h" +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// --------------------------------------------------------------------------- +// ReuseScope +// --------------------------------------------------------------------------- + +// Serialized layout of a ReuseScope, consumed by Hasher::update(ReuseScope). +// Must stay byte-identical to the Python ReuseScope.to_bytes(): a mask byte +// followed by one little-endian uint64 per present field (signed=False). +template +static void emitReuseScopeBytes(ReuseScope const& scope, Emit&& emit) +{ + uint8_t mask = 0; + if (scope.loraId.has_value()) + { + mask |= 1U << 0; + } + if (scope.salt.has_value()) + { + mask |= 1U << 1; + } + emit(&mask, sizeof(mask)); + if (scope.loraId.has_value()) + { + std::uint64_t const value = *scope.loraId; + emit(reinterpret_cast(&value), sizeof(value)); + } + if (scope.salt.has_value()) + { + std::uint64_t const value = *scope.salt; + emit(reinterpret_cast(&value), sizeof(value)); + } +} + +// --------------------------------------------------------------------------- +// Hasher +// --------------------------------------------------------------------------- + +namespace +{ +// Select the best available SHA-256 back-end (x86 SHA-NI, ARMv8 crypto, SSE4, +// AVX2) once, falling back to a portable scalar transform. CSHA256 dispatches +// through function pointers that SHA256AutoDetect() installs; the magic-static +// guarantees this runs exactly once and is thread-safe. +void ensureSha256Detected() +{ + static std::string const impl = SHA256AutoDetect(); + (void) impl; +} +} // namespace + +static void hashInt64(CSHA256& h, int64_t v) +{ + unsigned char buf[8]; + auto const unsignedValue = static_cast(v); + for (int i = 0; i < 8; ++i) + { + buf[i] = static_cast((unsignedValue >> (8 * i)) & 0xFFU); + } + h.Write(buf, sizeof(buf)); +} + +Hasher::Hasher() +{ + ensureSha256Detected(); +} + +Hasher::Hasher(ReuseScope const& seed) +{ + ensureSha256Detected(); + update(seed); +} + +Hasher& Hasher::update(ReuseScope const& scope) +{ + // Feed the serialized ReuseScope straight into the hash state without any + // intermediate heap buffer. + emitReuseScopeBytes(scope, [this](uint8_t const* data, size_t count) { mState.Write(data, count); }); + return *this; +} + +Hasher& Hasher::update(TokenId token) +{ + hashInt64(mState, static_cast(token)); + return *this; +} + +Hasher& Hasher::update(BlockKey const& key) +{ + mState.Write(reinterpret_cast(key.data()), key.size()); + return *this; +} + +Hasher& Hasher::update(std::vector const& bytes) +{ + mState.Write(bytes.data(), bytes.size()); + return *this; +} + +Hasher& Hasher::update(TokenIdExt const& tokenExt) +{ + std::visit( + [this](auto const& v) + { + using T = std::decay_t; + if constexpr (std::is_same_v) + hashInt64(mState, static_cast(v)); + else + mState.Write(reinterpret_cast(v.data()), v.size()); + }, + tokenExt); + return *this; +} + +Hasher& Hasher::update(TokenIdExt const* tokens, size_t count) +{ + // Python uses array("Q", data).tobytes() to reduce per-token interpreter + // overhead. In C++ the compiler inlines each update() call, so the loop + // is already optimal; batching would only add a heap allocation. + for (size_t i = 0; i < count; ++i) + update(tokens[i]); + return *this; +} + +BlockKey Hasher::digest() const +{ + // Finalize into a 32-byte key. CSHA256::Finalize consumes the state, so we + // finalize a copy to keep this method const and allow further updates. + BlockKey out; + CSHA256 copy = mState; + copy.Finalize(reinterpret_cast(out.data())); + return out; +} + +// --------------------------------------------------------------------------- +// genMultiModalTokens +// --------------------------------------------------------------------------- + +std::vector genMultiModalTokens( + int idOffset, std::vector const& multiModalDataDigest, int numTokens, int tokenOffset) +{ + TLLM_CHECK_DEBUG(numTokens > 0); + TLLM_CHECK_DEBUG(tokenOffset >= 0); + TLLM_CHECK_DEBUG(multiModalDataDigest.size() == kDIGEST_LEN); + std::vector result; + result.reserve(static_cast(numTokens)); + for (int i = 0; i < numTokens; ++i) + { + if (tokenOffset + i == 0) + { + Digest d; + std::memcpy(d.data(), multiModalDataDigest.data(), kDIGEST_LEN); + result.emplace_back(DigestToken(d)); + } + else + { + result.emplace_back(TokenId(idOffset + tokenOffset + i)); + } + } + return result; +} + +// --------------------------------------------------------------------------- +// makeBlockchainKeyGenerator — lazy key generator. +// Returns a callable that yields one BlockKey per call (nullopt when done). +// First call yields root entry (empty token block). Mirrors Python's generator. +// --------------------------------------------------------------------------- + +static auto makeBlockchainKeyGenerator( + int tokensPerBlock, ReuseScope reuseScope, TokenIdExt const* tokens, size_t numTokens) +{ + // digest carries the running hash from the previous block. + BlockKey digest = Hasher(reuseScope).digest(); + // ordinal = -1: next call yields root (reuseScope digest). + // ordinal >= 0: next call yields key for tokens[ordinal*tpb .. (ordinal+1)*tpb). + int ordinal = -1; + + return [=]() mutable -> std::optional + { + if (ordinal == -1) + { + ordinal++; + return digest; // root key + } + + size_t beg = static_cast(ordinal) * static_cast(tokensPerBlock); + if (beg >= numTokens) + return std::nullopt; + + size_t end = std::min(beg + static_cast(tokensPerBlock), numTokens); + + Hasher h; + h.update(digest); + h.update(tokens + beg, end - beg); + digest = h.digest(); + + ordinal++; + return digest; + }; +} + +// Eager wrapper for callers that need all keys at once. +std::vector sequenceToBlockchainKeys( + int tokensPerBlock, ReuseScope const& reuseScope, std::vector const& tokens) +{ + std::vector result; + auto gen = makeBlockchainKeyGenerator(tokensPerBlock, reuseScope, tokens.data(), tokens.size()); + while (auto key = gen()) + result.push_back(*key); + return result; +} + +// --------------------------------------------------------------------------- +// RootBlock +// --------------------------------------------------------------------------- + +BlockKey RootBlock::makeKey(ReuseScope const& reuseScope) +{ + return Hasher(reuseScope).digest(); +} + +RootBlock::RootBlock(ReuseScope reuseScope_, BlockRadixTree* treePtr) + : NodeBase(makeKey(reuseScope_), treePtr->eventSink().get()) + , reuseScope(std::move(reuseScope_)) + , tree(treePtr) +{ +} + +int RootBlock::tokensPerBlock() const noexcept +{ + return tree->tokensPerBlock(); +} + +// --------------------------------------------------------------------------- +// NodeBase +// --------------------------------------------------------------------------- + +NodeBase::~NodeBase() +{ + // Detach children before next is destroyed (implicit member destruction). + // This ensures that when a child's ~Block() runs, it sees prev == nullptr + // and skips parent cleanup — avoiding virtual calls on a mid-destruction parent. + for (auto& [k, child] : next) + { + child->prev = nullptr; + } +} + +SharedPtr NodeBase::detachNext(BlockKey const& blockKey) +{ + auto it = next.find(blockKey); + if (it == next.end()) + { + return nullptr; + } + + auto block = it->second; + block->prev = nullptr; + next.erase(it); + if (eventSink) + { + eventSink->addRemovedBlock(block->key); + } + if (type() == Type::kROOT_BLOCK && next.empty()) + { + auto* root = static_cast(this); + root->tree->proposeToEraseEmptyRoot(root->key); + } + return block; +} + +// --------------------------------------------------------------------------- +// Block +// --------------------------------------------------------------------------- + +namespace +{ + +static bool isPrefix(std::vector const& prefix, std::vector const& full) +{ + if (prefix.size() > full.size()) + return false; + for (size_t i = 0; i < prefix.size(); ++i) + { + if (prefix[i] != full[i]) + return false; + } + return true; +} + +} // anonymous namespace + +BlockKey Block::makeKey(BlockKey const& prevKey, TokenIdExt const* tokens, size_t count) +{ + Hasher h; + h.update(prevKey); + h.update(tokens, count); + return h.digest(); +} + +Block::Block(BlockKey k, std::vector toks, NodeBase* prevNode, LifeCycleId numLifeCycles) + : NodeBase(k, prevNode->eventSink) + , tokens(std::move(toks)) + , prev(prevNode) + , storage(numLifeCycles, nullptr) + , mOrdinal(prevNode->ordinal() + 1) +{ +} + +int Block::tokensPerBlock() const noexcept +{ + TLLM_CHECK_DEBUG_WITH_INFO(prev, "Block must have a prev"); + // Mirrors Python: prev.tokens_per_block if isinstance(prev, RootBlock) else len(prev.tokens) + if (prev->type() == Type::kROOT_BLOCK) + return prev->tokensPerBlock(); + return static_cast(static_cast(prev)->tokens.size()); +} + +void Block::releasePages() +{ + // Mirrors Python Block._release_pages(): for each stored page, if alive and + // DROPPABLE and scheduled for eviction, exclude from eviction. Also null out + // the page's back-pointer so that CommittedPage::~CommittedPage() doesn't + // attempt cleanup through this Block. Idempotent — storage is empty afterwards. + for (LifeCycleId lcIdx{0}; lcIdx < storage.size(); ++lcIdx) + { + auto const page = storage[lcIdx]; + if (page != nullptr) + { + TLLM_CHECK_DEBUG(page->block == this); + unlinkPage(lcIdx); + if (page->status() == PageStatus::DROPPABLE && page->scheduledForEviction()) + { + page->manager->excludeFromEviction(*page); + } + } + } +} + +Block::~Block() +{ + releasePages(); +} + +bool Block::isOrphan() const noexcept +{ + TLLM_CHECK_DEBUG(prev == nullptr || (prev->next.count(key) == 1 && prev->next.at(key).get() == this)); + return prev == nullptr; +} + +int Block::partialMatchThisNode(TokenIdExt const* otherTokens, size_t otherCount) const +{ + int count = 0; + for (size_t i = 0; i < std::min(tokens.size(), otherCount); ++i) + { + if (tokens[i] != otherTokens[i]) + break; + ++count; + } + return count; +} + +CommittedPage* Block::unlinkPage(LifeCycleId lcIdx, CommittedPage* expectedPage) +{ + auto& slot = storage.at(lcIdx); + CommittedPage* page = slot; + if (page == nullptr) + return nullptr; + if (expectedPage != nullptr && page != expectedPage) + return nullptr; + page->block = nullptr; + slot = nullptr; + return page; +} + +std::vector> Block::clearStaleBlocksAfterPageUnlink( + Block& block, LifeCycleId lcIdx, LifeCycle const& lc) +{ + std::vector> detachedBlocks; + TLLM_CHECK_DEBUG(block.storage.at(lcIdx) == nullptr); + if (block.isOrphan()) + { + return detachedBlocks; + } + + // Reuse cleanup only applies to attention lifecycles. + // SSM lifecycles are allowed in the tree but don't trigger subtree eviction. + auto const* const alc = std::get_if(&lc); + NodeBase* pruneStart = █ + + // If this is a full-attention block or a sink block: evict subtree. + // Mirrors Python: pages = remove_subtree(self) + if (alc && (!alc->windowSize.has_value() || block.ordinal() < BlockOrdinal{alc->numSinkBlocks})) + { + pruneStart = block.prev; + detachedBlocks.push_back(removeSubtree(block)); + } + else if (block.eventSink) + { + block.eventSink->addRemovedLifeCycle(block.key, lcIdx); + } + + // Prune empty tail nodes up the chain. + // Save prev, key, and type before erasing, because the erase may destroy + // curr when its last shared_ptr is dropped. + Block* curr + = pruneStart && pruneStart->type() == NodeBase::Type::kBLOCK ? static_cast(pruneStart) : nullptr; + while (curr && curr->next.empty() && curr->storage.at(lcIdx) == nullptr) + { + NodeBase* prevNode = curr->prev; + BlockKey const currKey = curr->key; + bool const prevIsBlock = prevNode && prevNode->type() == NodeBase::Type::kBLOCK; + if (prevNode) + { + auto detached = prevNode->detachNext(currKey); // may destroy curr + TLLM_CHECK_DEBUG(detached && detached.get() == curr); + detachedBlocks.push_back(std::move(detached)); + } + // Walk up only through Block nodes; stop at RootBlock. + curr = prevIsBlock ? static_cast(prevNode) : nullptr; + } + return detachedBlocks; +} + +// --------------------------------------------------------------------------- +// addOrGetExistingBlock +// --------------------------------------------------------------------------- + +SharedPtr addOrGetExistingBlock( + NodeBase* prev, LifeCycleId numLifeCycles, std::vector tokens, bool* isNew) +{ + TLLM_CHECK_DEBUG_WITH_INFO(prev, "prev must not be null"); + + // Prev must be a full block if it is a Block (mirrors Python: "prev must be a full block"). + if (prev->type() == NodeBase::Type::kBLOCK) + { + TLLM_CHECK_DEBUG_WITH_INFO(static_cast(prev)->isFull(), "prev must be a full block"); + } + + auto& prevNext = prev->next; + int const tpb = prev->tokensPerBlock(); + BlockKey newKey = Block::makeKey(prev->key, tokens.data(), tokens.size()); + + // Exact match: return existing block (not new — mirrors Python's UselessBlockError path). + auto it = prevNext.find(newKey); + if (it != prevNext.end()) + { + if (isNew) + *isNew = false; + return it->second; + } + + // Useless check: is this block's token prefix covered by a sibling? + // Mirrors Python's UselessBlockError — throw with the sibling block. + if (static_cast(tokens.size()) < tpb) + { + for (auto const& [k, sibling] : prevNext) + { + if (sibling->tokens.size() >= tokens.size() && isPrefix(tokens, sibling->tokens)) + throw UselessBlockError(sibling); + } + } + + // Remove siblings whose tokens are a strict prefix of ours. + std::vector toRemove; + for (auto const& [k, sibling] : prevNext) + { + if (sibling->tokens.size() < tokens.size() && isPrefix(sibling->tokens, tokens)) + { + TLLM_CHECK_DEBUG(!sibling->isFull() && sibling->key == k && sibling->next.empty()); + toRemove.push_back(k); + } + } + for (auto const& k : toRemove) + { + auto erasedBlock = prev->detachNext(k); + TLLM_CHECK_DEBUG(erasedBlock); + TLLM_CHECK_DEBUG_WITH_INFO(erasedBlock->isOrphan(), "erased sibling must be orphan after removal"); + (void) erasedBlock; + } + + // Create the new block. ordinal and tokensPerBlock are derived from prev inside the Block ctor. + auto block = makeShared(newKey, std::move(tokens), prev, numLifeCycles); + + prevNext[newKey] = block; + if (isNew) + *isNew = true; + return block; +} + +// --------------------------------------------------------------------------- +// removeSubtree +// --------------------------------------------------------------------------- + +SharedPtr removeSubtree(Block& root) +{ + Block* current = &root; + SharedPtr detachedRoot; + + // Post-order traversal using prev/next links — O(1) extra space. + // Descend to leaves first, remove on the way back up. + // Each block's pages are reclaimed eagerly via releasePages() while the + // StorageManager is still alive, rather than deferring to ~Block(): an external + // reference can keep a Block alive past StorageManager teardown, after which + // page->manager would be dangling. Mirrors Python's remove_subtree(). + while (true) + { + // Descend: if the current block has children, go to the first child. + if (!current->next.empty()) + { + current = current->next.begin()->second.get(); + } + else + { + current->releasePages(); + // Remove this block from its parent's next map. + // Null prev to detach — the block may outlive the tree if held + // externally (e.g., by nanobind/Python shared_ptr). + NodeBase* parent = current->prev; + BlockKey const currentKey = current->key; + auto detached = parent->detachNext(currentKey); + TLLM_CHECK_DEBUG(detached && detached.get() == current); + (void) detached; + + if (current == &root) + { + detachedRoot = std::move(detached); + break; + } + + TLLM_CHECK_DEBUG(parent->type() == NodeBase::Type::kBLOCK); + current = static_cast(parent); + } + } + TLLM_CHECK_DEBUG(detachedRoot); + return detachedRoot; +} + +// --------------------------------------------------------------------------- +// BlockRadixTree +// --------------------------------------------------------------------------- + +BlockRadixTree::BlockRadixTree( + LifeCycleRegistry const& lifeCycles, int tokensPerBlock, std::shared_ptr eventSink) + : mLifeCycles(lifeCycles) + , mTokensPerBlock(tokensPerBlock) + , mEventSink(std::move(eventSink)) +{ +} + +BlockRadixTree::~BlockRadixTree() +{ + // Clear all roots (which will drop all blocks). + mRoots.clear(); +} + +LifeCycleId BlockRadixTree::numLifeCycles() const noexcept +{ + return mLifeCycles.size(); +} + +void BlockRadixTree::drainPendingRootErases() const +{ + if (mPendingRootErases.empty()) + { + return; + } + // Move to local to allow re-entrancy (proposeToEraseEmptyRoot during erase). + std::vector pending; + pending.swap(mPendingRootErases); + auto& roots = const_cast>&>(mRoots); + for (auto const& key : pending) + { + auto it = roots.find(key); + // Only erase if the root exists and is still childless. + if (it != roots.end() && it->second->next.empty()) + { + roots.erase(it); + } + } +} + +RootBlock& BlockRadixTree::addOrGetExisting(ReuseScope const& reuseScope) +{ + drainPendingRootErases(); + + BlockKey key = RootBlock::makeKey(reuseScope); + auto it = mRoots.find(key); + if (it != mRoots.end()) + { + return *it->second; + } + + auto rb = makeShared(reuseScope, this); + auto [newIt, inserted] = mRoots.emplace(key, std::move(rb)); + return *newIt->second; +} + +// Among all child nodes, find the one whose tokens have the longest leading match. +// Returns (block, numMatchedTokens) or (nullptr, 0) if no match. +// Mirrors Python's find_best_partial_match_in_next_nodes(). +std::pair findBestPartialMatchInNextNodes( + std::unordered_map> const& nextMap, TokenIdExt const* tokens, size_t tokenCount) +{ + // Skip heuristic: too many children would be slow to iterate. + if (nextMap.size() >= 32) + return {nullptr, 0}; + Block* best = nullptr; + int bestMatch = 0; + for (auto const& [k, child] : nextMap) + { + int m = child->partialMatchThisNode(tokens, tokenCount); + if (m > bestMatch) + { + bestMatch = m; + best = child.get(); + } + } + return {best, bestMatch}; +} + +namespace +{ + +int numMatchedTokens(std::vector const& matched, int tokensPerBlock) +{ + if (matched.empty()) + { + return 0; + } + return tokensPerBlock * (static_cast(matched.size()) - 1) + matched.back().numMatchedTokens; +} + +bool hasPage(Block const& block, LifeCycleId lcId) +{ + return block.storage.at(lcId) != nullptr; +} + +} // anonymous namespace + +std::vector BlockRadixTree::matchTokenPath( + ReuseScope const& reuseScope, std::vector const& tokens, bool enablePartialMatch) const +{ + drainPendingRootErases(); + + std::vector results; + + // Lazily compute one key per iteration — no wasted hashing on early miss. + auto gen = makeBlockchainKeyGenerator(mTokensPerBlock, reuseScope, tokens.data(), tokens.size()); + + // First key is the root key. + auto rootKey = gen(); + if (!rootKey) + return results; + auto rootIt = mRoots.find(*rootKey); + if (rootIt == mRoots.end()) + return results; + + RootBlock const& root = *rootIt->second; + std::unordered_map> const* currentNext = &root.next; + // ordinal tracks which block we're on (0-based, after root). + BlockOrdinal ordinal{0}; + bool missed = false; + + while (auto key = gen()) + { + auto blockIt = currentNext->find(*key); + if (blockIt == currentNext->end()) + { + missed = true; + break; + } + size_t beg = toSizeT(ordinal) * static_cast(mTokensPerBlock); + int numTokens = static_cast(std::min(static_cast(mTokensPerBlock), tokens.size() - beg)); + Block* block = blockIt->second.get(); + results.push_back({block, numTokens}); + currentNext = &block->next; + ordinal++; + } + + // Partial match in children of current node. + if (missed && enablePartialMatch) + { + size_t beg = toSizeT(ordinal) * static_cast(mTokensPerBlock); + size_t missedCount = std::min(static_cast(mTokensPerBlock), tokens.size() - beg); + auto [best, bestMatch] = findBestPartialMatchInNextNodes(*currentNext, tokens.data() + beg, missedCount); + if (best) + results.push_back({best, bestMatch}); + } + + return results; +} + +std::vector BlockRadixTree::pruneMatch(std::vector matched) const +{ + // All blocks except the last must be fully matched (mirrors Python: matched[:-1]). + TLLM_CHECK_DEBUG(matched.size() <= 1 + || std::all_of(matched.begin(), matched.end() - 1, + [this](auto const& m) { return m.numMatchedTokens == mTokensPerBlock; })); + + auto attnLcs = mLifeCycles.attentionLifeCycles(); + + // Full-attention layers require pages on every matched block. + std::vector fullAttnLcList; + for (auto [lcId, attn] : attnLcs) + { + if (!attn->windowSize.has_value()) + { + fullAttnLcList.push_back(lcId); + } + } + if (!fullAttnLcList.empty()) + { + int n = findIndex(matched.begin(), matched.end(), + [&](auto const& match) + { + return std::any_of(fullAttnLcList.begin(), fullAttnLcList.end(), + [&](LifeCycleId lcId) { return !hasPage(*match.block, lcId); }); + }); + matched.resize(static_cast(n)); + } + + std::vector> swaLcs; + for (auto [lcId, attn] : attnLcs) + { + if (attn->windowSize.has_value()) + { + swaLcs.push_back({lcId, attn}); + } + } + + // SWA sink blocks must all be available. + for (auto [lcId, attn] : swaLcs) + { + int const sinkBlocks = attn->numSinkBlocks; + int const limit = std::min(sinkBlocks, static_cast(matched.size())); + int n = findIndex(matched.begin(), matched.begin() + limit, + [&, lcId = lcId](auto const& match) { return !hasPage(*match.block, lcId); }); + if (n < sinkBlocks) + { + matched.resize(static_cast(n)); + } + } + + auto ssmLcId = mLifeCycles.ssmLifeCycleId(); + while (!matched.empty()) + { + if (ssmLcId.has_value()) + { + // Truncate to the last block whose SSM snapshot is reusable at that + // block's matched-token count, then clamp the tail entry's matched + // token count to the snapshot length (mirrors _block_radix_tree.py). + int ssmTrunc = 0; + int ssmMatchLen = 0; + for (int i = static_cast(matched.size()) - 1; i >= 0; --i) + { + CommittedPage* page = matched[static_cast(i)].block->storage.at(*ssmLcId); + if (page == nullptr) + { + continue; + } + auto* ssmPage = dynamic_cast(page); + TLLM_CHECK_DEBUG(ssmPage != nullptr); + int const snapshotLen = ssmPage->numTokensInBlock; + if (matched[static_cast(i)].numMatchedTokens >= snapshotLen) + { + ssmTrunc = i + 1; + ssmMatchLen = snapshotLen; + break; + } + } + matched.resize(static_cast(ssmTrunc)); + if (matched.empty()) + { + break; + } + matched.back().numMatchedTokens = ssmMatchLen; + } + + int const numTok = numMatchedTokens(matched, mTokensPerBlock); + bool trimmed = false; + for (auto [lcId, attn] : swaLcs) + { + int n = findIndex(matched.rbegin(), matched.rend(), + [&, lcId = lcId](auto const& match) { return hasPage(*match.block, lcId); }); + if (n != 0) + { + matched.resize(matched.size() - static_cast(n)); + trimmed = true; + break; + } + + auto staleRange = attn->getStaleRange(numTok, mTokensPerBlock); + BlockOrdinal const staleEnd = staleRange.end; + if (staleEnd < BlockOrdinal{static_cast(matched.size())}) + { + auto tailBegin = matched.begin() + static_cast(toSizeT(staleEnd)); + int nMissing = findIndex(matched.rbegin(), std::make_reverse_iterator(tailBegin), + [&, lcId = lcId](auto const& match) { return !hasPage(*match.block, lcId); }); + if (BlockOrdinal{static_cast(matched.size()) - nMissing} > staleEnd) + { + matched.resize(matched.size() - static_cast(nMissing) - 1); + trimmed = true; + break; + } + } + } + if (!trimmed) + { + break; + } + } + + return matched; +} + +BlockRadixTree::ReuseMatch BlockRadixTree::match( + ReuseScope const& reuseScope, std::vector const& tokens, bool enablePartialMatch) const +{ + auto const matched = pruneMatch(matchTokenPath(reuseScope, tokens, enablePartialMatch)); + ReuseMatch result{}; + result.numTokens = numMatchedTokens(matched, mTokensPerBlock); + result.numLookupTokens = static_cast(tokens.size()); + result.blocks.reserve(BlockOrdinal{static_cast(matched.size())}); + for (auto const& match : matched) + { + result.blocks.push_back(match.block); + } + return result; +} + +void BlockRadixTree::clear() +{ + // detachNext() may call proposeToEraseEmptyRoot, but won't modify mRoots directly. + for (auto& [rootKey, root] : mRoots) + { + while (!root->next.empty()) + { + removeSubtree(*root->next.begin()->second); + } + } + TLLM_CHECK_DEBUG(mRoots.size() == mPendingRootErases.size()); + mRoots.clear(); + mPendingRootErases.clear(); +} + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h new file mode 100644 index 000000000000..c3ae3035415c --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h @@ -0,0 +1,356 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "kv_cache_manager_v2/common.h" +#include "kv_cache_manager_v2/eventSink.h" +#include "kv_cache_manager_v2/lifeCycleRegistry.h" +#include "kv_cache_manager_v2/utils/sharedPtr.h" + +#include "sha256.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// Forward declarations +class CommittedPage; +class BlockRadixTree; +struct NodeBase; +struct RootBlock; +struct Block; + +// --------------------------------------------------------------------------- +// ReuseScope — per-request namespace for prefix reuse. +// Mirrors Python's ReuseScope(lora_id, salt). +// --------------------------------------------------------------------------- +struct ReuseScope +{ + std::optional loraId; + std::optional salt; + + bool operator==(ReuseScope const& other) const noexcept + { + return loraId == other.loraId && salt == other.salt; + } +}; + +// --------------------------------------------------------------------------- +// BlockKey — SHA-256 digest (32 bytes), used as radix-tree node identifier. +// Matches Python's hashlib.sha256 32-byte digest. +// +// SECURITY INVARIANT: the block hash MUST remain cryptographically +// collision-resistant and >= 256-bit. The radix tree is a globally shared, +// cross-request/cross-tenant cache index, prefix matching is decided purely by +// digest equality with NO re-verification of the underlying tokens, and the +// hashed input (tokens, the user-supplied cache_salt in ReuseScope, multimodal +// content bytes) is attacker-influenceable. A hash collision therefore silently +// reuses another request's KV blocks (cross-request corruption / data leak), +// and tenant isolation via cache_salt relies entirely on this hash's collision +// resistance. Do NOT substitute a non-cryptographic hash (xxHash, HighwayHash, +// City, ...) or truncate below 256 bits without first adding a token-content +// equality check on match. +// --------------------------------------------------------------------------- +using BlockKey = Digest; +static_assert(kDIGEST_LEN == CSHA256::OUTPUT_SIZE); // 32 bytes + +// --------------------------------------------------------------------------- +// Hasher — thin wrapper around SHA-256 (CSHA256) for incremental digests. +// Mirrors Python's Hasher class (hashlib.sha256). See the SECURITY INVARIANT on +// BlockKey above before changing the hash algorithm or digest width. +// --------------------------------------------------------------------------- +class Hasher +{ +public: + Hasher(); + explicit Hasher(ReuseScope const& seed); + + Hasher& update(TokenId token); + Hasher& update(BlockKey const& key); + Hasher& update(ReuseScope const& scope); + Hasher& update(std::vector const& bytes); + Hasher& update(TokenIdExt const& tokenExt); + Hasher& update(TokenIdExt const* tokens, size_t count); + + BlockKey digest() const; + +private: + CSHA256 mState; +}; + +// --------------------------------------------------------------------------- +// Utility: convert a token sequence → list of BlockKeys. +// First key is the root (reuseScope digest), then one per token block. +// Mirrors Python's sequence_to_blockchain_keys(). +// --------------------------------------------------------------------------- +std::vector sequenceToBlockchainKeys( + int tokensPerBlock, ReuseScope const& reuseScope, std::vector const& tokens); + +// Generate multi-modal token IDs (mirrors gen_multi_modal_tokens in Python). +std::vector genMultiModalTokens( + int idOffset, std::vector const& multiModalDataDigest, int numTokens, int tokenOffset = 0); + +// --------------------------------------------------------------------------- +// NodeBase — common base for RootBlock and Block (nodes in the radix tree). +// Holds shared fields: key, next map, ordinal, and tokens-per-block. +// Mirrors Python's common interface between RootBlock and Block. +// --------------------------------------------------------------------------- +struct NodeBase +{ + enum class Type : uint8_t + { + kROOT_BLOCK, + kBLOCK + }; + + BlockKey key; + std::unordered_map> next; + EventSink* eventSink; + + virtual ~NodeBase(); + + virtual Type type() const noexcept = 0; + virtual BlockOrdinal ordinal() const noexcept = 0; + + SharedPtr detachNext(BlockKey const& key); + + /// RootBlock: delegates to tree. Block: len(prev->tokens) or prev->tokensPerBlock(). + virtual int tokensPerBlock() const noexcept = 0; + +protected: + NodeBase(BlockKey k, EventSink* sink) + : key(k) + , eventSink(sink) + { + } +}; + +// --------------------------------------------------------------------------- +// RootBlock — one root per ReuseScope in a BlockRadixTree. +// Holds a map of child Blocks keyed by BlockKey. +// Mirrors Python's RootBlock. +// --------------------------------------------------------------------------- +struct RootBlock : NodeBase +{ + ReuseScope reuseScope; + BlockRadixTree* tree; // back-reference (non-owning) + + RootBlock(ReuseScope reuseScope, BlockRadixTree* tree); + + static BlockKey makeKey(ReuseScope const& reuseScope); + + Type type() const noexcept override + { + return Type::kROOT_BLOCK; + } + + BlockOrdinal ordinal() const noexcept override + { + return kBadBlockOrdinal; + } + + int tokensPerBlock() const noexcept override; +}; + +// --------------------------------------------------------------------------- +// Block — one full (or partial) token block in the radix tree. +// storage[lifeCycleId] = raw observer pointer to CommittedPage (null if not cached). +// Mirrors Python's Block. +// --------------------------------------------------------------------------- +struct Block : NodeBase, EnableSharedFromThis +{ + std::vector tokens; + + // Previous node in the chain (RootBlock or Block). Null after detaching from the tree. + // Raw non-owning pointer: while attached, the prev node's `next` map owns us via shared_ptr. + NodeBase* prev{nullptr}; + + TypedVec storage; + + Block(BlockKey key, std::vector tokens, NodeBase* prev, LifeCycleId numLifeCycles); + ~Block() override; + + static BlockKey makeKey(BlockKey const& prevKey, TokenIdExt const* tokens, size_t count); + + Type type() const noexcept override + { + return Type::kBLOCK; + } + + BlockOrdinal ordinal() const noexcept override + { + return mOrdinal; + } + + int tokensPerBlock() const noexcept override; + + LifeCycleId numLifeCycles() const noexcept + { + return storage.size(); + } + + bool isFull() const noexcept + { + return static_cast(tokens.size()) == tokensPerBlock(); + } + + bool isOrphan() const noexcept; + + // Returns how many leading tokens match `otherTokens`. + int partialMatchThisNode(TokenIdExt const* otherTokens, size_t count) const; + + // Break the bidirectional link to the cached page for a lifecycle. + // Returns the previously-stored CommittedPage* (nullptr if already unlinked). + // If `expectedPage` is non-null and the stored page differs from it, the link + // is left untouched and nullptr is returned (mirrors Python's unset_page + // `expected_page` guard: a newer page may already occupy the slot). + CommittedPage* unlinkPage(LifeCycleId lcIdx, CommittedPage* expectedPage = nullptr); + + // Return the cached page for a lifecycle (nullptr if none). Mirrors Python's Block.get_page(). + CommittedPage* getPage(LifeCycleId lcIdx) const + { + return storage[lcIdx]; + } + + // Clear stale tree nodes after a lifecycle page has been unlinked. + // Returns detached blocks that must stay alive until cleanup completes. + static std::vector> clearStaleBlocksAfterPageUnlink( + Block& block, LifeCycleId lcIdx, LifeCycle const& lc); + + // Reclaim every page held by this block: null each page's back-pointer and, for + // DROPPABLE pages still scheduled for eviction, remove them from the eviction + // controller (releasing their storage slots). Idempotent. Must run during tree + // teardown (removeSubtree) rather than being deferred to ~Block(), so page + // reclamation does not depend on this Block's destruction timing — an external + // reference can keep a Block alive past StorageManager teardown, after which + // page->manager would be dangling. Mirrors Python's Block._release_pages(). + void releasePages(); + +private: + BlockOrdinal mOrdinal; +}; + +// --------------------------------------------------------------------------- +// BlockRadixTree — the global cache index. +// next: reuseScope digest → RootBlock. +// Mirrors Python's BlockRadixTree. +// --------------------------------------------------------------------------- +class BlockRadixTree +{ +public: + BlockRadixTree( + LifeCycleRegistry const& lifeCycles, int tokensPerBlock, std::shared_ptr eventSink = nullptr); + ~BlockRadixTree(); + + // Get (or create) the RootBlock for the given reuse scope. + RootBlock& addOrGetExisting(ReuseScope const& reuseScope); + + // Match tokens against the tree, yielding (block, numMatchedTokens) pairs. + // Partial matching: if enablePartialMatch, also yields blocks with a partial + // leading-token match. + struct MatchResult + { + Block* block; + int numMatchedTokens; + }; + + struct ReuseMatch + { + TypedVec blocks; + int numTokens; + // Total query length passed to match() (== len(tokens)). + int numLookupTokens; + }; + + ReuseMatch match( + ReuseScope const& reuseScope, std::vector const& tokens, bool enablePartialMatch = false) const; + + // Clear all cached pages. ~Block() handles excludeFromEviction for DROPPABLE pages. + void clear(); + + int tokensPerBlock() const noexcept + { + return mTokensPerBlock; + } + + LifeCycleId numLifeCycles() const noexcept; + + LifeCycleRegistry const& lifeCycles() const noexcept + { + return mLifeCycles; + } + + std::shared_ptr const& eventSink() const noexcept + { + return mEventSink; + } + + // Read-only access to the root map (used by nanobind introspection). + std::unordered_map> const& roots() const noexcept + { + return mRoots; + } + + // Propose removal of an empty root block. Deferred to avoid destroying + // objects during destructor chains. Drained at safe points (addOrGetExisting, match). + void proposeToEraseEmptyRoot(BlockKey const& key) + { + mPendingRootErases.push_back(key); + } + +private: + std::vector matchTokenPath( + ReuseScope const& reuseScope, std::vector const& tokens, bool enablePartialMatch) const; + std::vector pruneMatch(std::vector matched) const; + + // Erase any pending empty root blocks from mRoots. + // Const-qualified: deferred cleanup is not a logical mutation. + void drainPendingRootErases() const; + + LifeCycleRegistry const& mLifeCycles; + int mTokensPerBlock; + std::shared_ptr mEventSink; + + std::unordered_map> mRoots; + mutable std::vector mPendingRootErases; +}; + +// --------------------------------------------------------------------------- +// Helpers used by Block and the tree traversal. +// --------------------------------------------------------------------------- + +// Add a block to prev's `next` map, or return the existing one on collision. +// Throws UselessBlockError (with the sibling block) if the block's tokens are a +// prefix of an existing sibling — mirrors Python's UselessBlockError. +// If isNew is non-null, *isNew is set to true if a new block was created, false +// if an existing block was returned. +SharedPtr addOrGetExistingBlock( + NodeBase* prev, LifeCycleId numLifeCycles, std::vector tokens, bool* isNew = nullptr); + +// Post-order traversal: remove a subtree rooted at `root` from its parent's +// next map. ~Block() handles page cleanup. Mirrors Python's remove_subtree(). +SharedPtr removeSubtree(Block& root); + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.cpp new file mode 100644 index 000000000000..bd2928adb2f3 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.cpp @@ -0,0 +1,26 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "kv_cache_manager_v2/common.h" +#include "tensorrt_llm/common/assert.h" + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +bool const gDebug = tensorrt_llm::DebugConfig::isCheckDebugEnabled(); + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h new file mode 100644 index 000000000000..0f53938b95e2 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h @@ -0,0 +1,243 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "kv_cache_manager_v2/utils/typedIndex.h" +#include "tensorrt_llm/batch_manager/common.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// --------------------------------------------------------------------------- +// Debug flag — true when TLLM_DEBUG_MODE=1. +// Delegates to DebugConfig::isCheckDebugEnabled() for consistency with TLLM_CHECK_DEBUG. +// --------------------------------------------------------------------------- +extern bool const gDebug; // true == debug mode (expensive assertions enabled) + +// --------------------------------------------------------------------------- +// Enumerations +// --------------------------------------------------------------------------- + +enum class PageStatus : int +{ + LOCKED = 0, // Required in GPU. Eviction/dropping not allowed. + HELD = 1, // Allow eviction but not dropping. + DROPPABLE = 2, // Allow eviction and dropping. +}; + +enum class CacheTier : int +{ + GPU_MEM = 0, + HOST_MEM = 1, + DISK = 2, +}; + +// PageIndexMode — how converted page indices relate to layers within a layer group. +// Mirrors _common.py::PageIndexMode. +enum class PageIndexMode : int +{ + // Converted index list is shared across layers in the same LayerGroup. + // Base pointer is per-layer (includes attr.offset). + SHARED = 0, + // Converted index list is per-layer. + // Base pointer is shared (pool group base, no attr.offset). + PER_LAYER = 1, +}; + +// --------------------------------------------------------------------------- +// Strongly-typed integer aliases (mirroring Python NewType wrappers). +// --------------------------------------------------------------------------- + +// Index of a cache level (0 = GPU, 1 = host, 2 = disk, ...). +using CacheLevel = StrongIndex; +inline constexpr CacheLevel kGpuLevel{0}; + +// Vocabulary token identifier (normal tokens only). +using TokenId = int64_t; + +// Opaque request identifier shared with the rest of the batch manager. +using RequestIdType = tensorrt_llm::batch_manager::RequestIdType; + +// Opaque LoRA task identifier shared with the rest of the batch manager. +using LoraTaskIdType = tensorrt_llm::runtime::LoraTaskIdType; + +// 32-byte aligned to enable SIMD. +inline constexpr int kDIGEST_LEN = 32; + +struct alignas(kDIGEST_LEN) Digest : std::array +{ + // Custom operator== needed to emit SIMD code + bool operator==(Digest const& o) const noexcept + { + return std::memcmp(this, &o, kDIGEST_LEN) == 0; + } + + bool operator!=(Digest const& o) const noexcept + { + return !(*this == o); + } +}; + +// Heap-allocated digest token for multi-modal tokens. +// Copyable (deep-copies the digest) with value-based equality. +// Digest tokens are rare, so unique_ptr keeps sizeof(TokenIdExt) small. +class DigestToken +{ +public: + explicit DigestToken(Digest const& d) + : mData(std::make_unique(d)) + { + } + + explicit DigestToken(std::unique_ptr d) + : mData(std::move(d)) + { + } + + DigestToken(DigestToken const& o) + : mData(std::make_unique(*o.mData)) + { + } + + DigestToken(DigestToken&&) noexcept = default; + + DigestToken& operator=(DigestToken const& o) + { + if (this != &o) + mData = std::make_unique(*o.mData); + return *this; + } + + DigestToken& operator=(DigestToken&&) noexcept = default; + + bool operator==(DigestToken const& o) const + { + return *mData == *o.mData; + } + + bool operator!=(DigestToken const& o) const + { + return !(*this == o); + } + + std::byte const* data() const noexcept + { + return mData->data(); + } + + size_t size() const noexcept + { + return mData->size(); + } + + Digest const& digest() const noexcept + { + return *mData; + } + +private: + std::unique_ptr mData; +}; + +// Extended token id: normal TokenId or a heap-allocated digest for multi-modal tokens. +using TokenIdExt = std::variant; + +// Ordinal index of a KV cache block (sequence of tokens). +using BlockOrdinal = StrongIndex; +inline constexpr BlockOrdinal kBadBlockOrdinal{-1}; + +// Identifier of an attention layer. +using LayerId = int; + +// Raw CUDA stream handle (CUstream cast to uintptr_t). +using CudaStream = uintptr_t; + +// Index of a beam in beam-search. +using BeamIndex = StrongIndex; +inline constexpr BeamIndex kDefaultBeamIndex{0}; + +// User-defined request/session identifier. +using UserId = int64_t; + +// Host or device memory address (uintptr_t). +using MemAddress = std::uintptr_t; + +// OS file descriptor. +using FileDescriptor = int; +inline constexpr FileDescriptor kBadFileDescriptor = -1; + +// Index into a page table. +using PageIndex = StrongIndex; +inline constexpr PageIndex kBadPageIndex{-1}; + +// Eviction priority (0 = highest priority to evict, 100 = lowest). +using Priority = int; +inline constexpr Priority kPriorityMin = 0; +inline constexpr Priority kPriorityMax = 100; +inline constexpr Priority kPriorityDefault = 35; + +// Optional sliding window size (nullopt = no sliding window). +using SlidingWindowSize = std::optional; + +// --------------------------------------------------------------------------- +// Address types +// --------------------------------------------------------------------------- + +// Disk address: (fd, byte-offset). +struct DiskAddress +{ + int fd = kBadFileDescriptor; + ssize_t pos = 0; + + bool operator==(DiskAddress const& o) const noexcept + { + return fd == o.fd && pos == o.pos; + } +}; + +// Unified address: either a host/device memory pointer or a disk address. +using Address = std::variant; + +// --------------------------------------------------------------------------- +// DataRole — string-typed tag for a buffer inside one attention layer. +// --------------------------------------------------------------------------- +using DataRole = std::string; + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 + +// std::hash specialization for Digest/BlockKey so unordered_map works without a custom hasher. +template <> +struct std::hash +{ + size_t operator()(tensorrt_llm::batch_manager::kv_cache_manager_v2::Digest const& k) const noexcept + { + // First 8 bytes of a SHA-256 digest are already well-distributed. + uint64_t v; + std::memcpy(&v, k.data(), sizeof(v)); + return static_cast(v); + } +}; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.cpp new file mode 100644 index 000000000000..c31ea7be3586 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.cpp @@ -0,0 +1,96 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "kv_cache_manager_v2/config.h" +#include "kv_cache_manager_v2/exceptions.h" + +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +void DiskCacheTierConfig::assertValid() const +{ + if (quota == 0) + { + throw std::invalid_argument("DiskCacheTierConfig: quota must be > 0"); + } + if (!std::filesystem::is_directory(path)) + { + throw std::invalid_argument("DiskCacheTierConfig: path '" + path + "' is not a directory"); + } +} + +void KVCacheManagerConfig::validate() const +{ + if (swaScratchReuse.has_value()) + { + swaScratchReuse->validate(); + } + + // These mirror Python's KVCacheManagerConfig.__post_init__ asserts, so they + // throw AssertionError (translated in the binding layer) rather than ValueError. + if (cacheTiers.empty() || cacheTierOf(cacheTiers[0]) != CacheTier::GPU_MEM) + { + throw AssertionError("KVCacheManagerConfig: first cache tier must be GPU memory"); + } + + // Check for duplicate layer ids. + std::set seenLayerIds; + for (auto const& layer : layers) + { + std::visit( + [&](auto const& cfg) + { + if (!seenLayerIds.insert(cfg.layerId).second) + { + throw AssertionError("KVCacheManagerConfig: duplicate layer id"); + } + for (auto const& buf : cfg.buffers) + { + if (buf.tokensPerBlockOverride.has_value() + && (*buf.tokensPerBlockOverride <= 0 || tokensPerBlock % *buf.tokensPerBlockOverride != 0)) + { + throw AssertionError( + "KVCacheManagerConfig: tokensPerBlockOverride must be a divisor of " + "tokensPerBlock"); + } + } + }, + layer); + } + + // SSM-specific validation. + bool hasSSM = false; + for (auto const& layer : layers) + { + if (std::holds_alternative(layer)) + { + hasSSM = true; + break; + } + } + if (hasSSM) + { + if (!commitMinSnapshot) + throw AssertionError("KVCacheManagerConfig: commit_min_snapshot must be True when SSM layers are present"); + } +} + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h new file mode 100644 index 000000000000..6852ed51f5a1 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h @@ -0,0 +1,304 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "kv_cache_manager_v2/common.h" + +#include "tensorrt_llm/common/assert.h" +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// --------------------------------------------------------------------------- +// Cache tier configuration structs (mirrors _config.py) +// --------------------------------------------------------------------------- + +struct GpuCacheTierConfig +{ + size_t quota = 0; // bytes + + CacheTier tier() const noexcept + { + return CacheTier::GPU_MEM; + } + + void assertValid() const + { + if (quota == 0) + throw std::invalid_argument("GpuCacheTierConfig: quota must be > 0"); + } +}; + +struct HostCacheTierConfig +{ + size_t quota = 0; // bytes + + CacheTier tier() const noexcept + { + return CacheTier::HOST_MEM; + } + + void assertValid() const + { + if (quota == 0) + throw std::invalid_argument("HostCacheTierConfig: quota must be > 0"); + } +}; + +struct DiskCacheTierConfig +{ + size_t quota = 0; // bytes + std::string path; // directory for temp files + + CacheTier tier() const noexcept + { + return CacheTier::DISK; + } + + void assertValid() const; +}; + +// Variant holding any tier config. +using CacheTierConfig = std::variant; + +// Helper to extract tier from a variant. +inline CacheTier cacheTierOf(CacheTierConfig const& cfg) +{ + return std::visit([](auto const& c) { return c.tier(); }, cfg); +} + +inline size_t cacheTierQuota(CacheTierConfig const& cfg) +{ + return std::visit([](auto const& c) { return c.quota; }, cfg); +} + +// --------------------------------------------------------------------------- +// Buffer configuration (one KV buffer inside an attention layer). +// --------------------------------------------------------------------------- + +struct BufferConfig +{ + DataRole role; + size_t size = 0; // bytes per page (without expansion) + + // If set, overrides tokens_per_block for this buffer. + // Must be a divisor of KVCacheManagerConfig::tokensPerBlock. + std::optional tokensPerBlockOverride; +}; + +// --------------------------------------------------------------------------- +// Layer type discriminator. +// --------------------------------------------------------------------------- + +enum class LayerType : int +{ + ATTENTION = 0, + SSM = 1, +}; + +// --------------------------------------------------------------------------- +// Attention layer configuration. +// --------------------------------------------------------------------------- + +namespace detail +{ + +inline void validateNoDuplicateBufferRoles(std::vector const& buffers) +{ + std::unordered_set roles; + for (auto const& buf : buffers) + { + if (!roles.insert(buf.role).second) + throw std::invalid_argument("duplicate buffer role"); + } +} + +} // namespace detail + +struct AttentionLayerConfig +{ + static constexpr LayerType type = LayerType::ATTENTION; + + LayerId layerId = 0; + std::vector buffers; + + // nullopt = no sliding window. + std::optional slidingWindowSize; + + // nullopt or 0 = no sink tokens. + std::optional numSinkTokens; + + std::optional windowSize() const noexcept + { + return slidingWindowSize; + } + + void validate() const + { + detail::validateNoDuplicateBufferRoles(buffers); + } +}; + +// --------------------------------------------------------------------------- +// SSM (State Space Model) layer configuration. +// --------------------------------------------------------------------------- + +struct SsmLayerConfig +{ + static constexpr LayerType type = LayerType::SSM; + + LayerId layerId = 0; + std::vector buffers; + + void validate() const + { + detail::validateNoDuplicateBufferRoles(buffers); + for (auto const& buf : buffers) + { + if (buf.tokensPerBlockOverride.has_value()) + throw std::invalid_argument("tokensPerBlockOverride not supported for SSM layers"); + } + } +}; + +using LayerConfig = std::variant; + +// --------------------------------------------------------------------------- +// KVCacheDesc — describes one KV cache request's capacity and history length. +// Mirrors _config.py::KVCacheDesc. +// --------------------------------------------------------------------------- +struct KVCacheDesc +{ + int capacity = 0; + int historyLength = 0; + + void validate() const + { + TLLM_CHECK_DEBUG(0 <= historyLength && historyLength <= capacity); + } + + // Value equality, mirroring the Python @dataclass(frozen=True) semantics the + // bindings replace. Required so tests can compare descs by value. + bool operator==(KVCacheDesc const& other) const noexcept + { + return capacity == other.capacity && historyLength == other.historyLength; + } + + bool operator!=(KVCacheDesc const& other) const noexcept + { + return !(*this == other); + } +}; + +// --------------------------------------------------------------------------- +// BatchDesc — a batch of requests that the KVCacheManager must support. +// Mirrors _config.py::BatchDesc. +// --------------------------------------------------------------------------- +struct BatchDesc +{ + std::vector kvCaches; + int systemPromptLength = 0; // tokens shared by all requests (0 if no reuse) + + void validate() const + { + TLLM_CHECK_DEBUG(systemPromptLength >= 0); + } + + // Value equality, mirroring the Python @dataclass(frozen=True) semantics the + // bindings replace. Uses KVCacheDesc::operator== elementwise. + bool operator==(BatchDesc const& other) const noexcept + { + return systemPromptLength == other.systemPromptLength && kvCaches == other.kvCaches; + } + + bool operator!=(BatchDesc const& other) const noexcept + { + return !(*this == other); + } +}; + +// --------------------------------------------------------------------------- +// SWA scratch reuse configuration. +// --------------------------------------------------------------------------- +struct SwaScratchReuseConfig +{ + int maxRewindLen = 0; + + void validate() const + { + if (maxRewindLen < 0) + { + throw std::invalid_argument("SwaScratchReuseConfig: max_rewind_len must be non-negative"); + } + } +}; + +// --------------------------------------------------------------------------- +// Top-level KV cache manager configuration (mirrors _config.py::KVCacheManagerConfig). +// --------------------------------------------------------------------------- + +struct KVCacheManagerConfig +{ + int tokensPerBlock = 0; + + // Ordered from warm (GPU) to cold (disk). First must be GPU memory. + std::vector cacheTiers; + + // Layer configs (attention or SSM). Layer IDs must be unique. + std::vector layers; + + // Suspend/resume threshold: if utilization > this, resuming will fail. + float maxUtilForResume = 0.97f; + + // Try to reuse tokens from partially matched blocks. + bool enablePartialReuse = true; + + // Constraint-based memory partitioning. + std::vector constraints; // batches that must always be supportable + std::optional typicalStep; // typical step for initial ratio computation + std::optional> initialPoolRatio; // explicit initial ratio, overrides inferred sizing inputs + + // When set, SWA layers reuse physical pages for out-of-window blocks during prefill. + // Scratch blocks share coalesced slot sub-pages across blocks for the currently executing + // layer, reducing peak memory. Trade-off: KV cache reuse is degraded because scratch blocks + // have no preserved data after the step. + std::optional swaScratchReuse; + + // If true, commit() records only the minimum cache snapshot reusable at the post-call + // numCommittedTokens. Only the minimum amount of pages required for such reuse will be + // preserved. Required when SSM layers are present. + bool commitMinSnapshot = false; + + // Collect V2 KV cache allocation, reuse, and transfer statistics. + bool enableStats = true; + + bool enableSwaScratchReuse() const noexcept + { + return swaScratchReuse.has_value(); + } + + void validate() const; +}; + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/copyEngine.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/copyEngine.cpp new file mode 100644 index 000000000000..d8579005b2f1 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/copyEngine.cpp @@ -0,0 +1,256 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "kv_cache_manager_v2/copyEngine.h" +#include "kv_cache_manager_v2/common.h" +#include "kv_cache_manager_v2/exceptions.h" +#include "kv_cache_manager_v2/utils/math.h" + +// Reuse existing copy implementations (no Python round-trip). +#include "tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h" + +#include "tensorrt_llm/common/assert.h" +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// --------------------------------------------------------------------------- +// dispatchCopy — template that maps CacheTier pair to +// address types and the underlying copy function via if-constexpr. +// Replaces 7 near-identical static dispatcher functions. +// --------------------------------------------------------------------------- + +template +using TierAddr = std::conditional_t; + +template +void dispatchCopy(std::vector const& tasks, size_t numBytes, CUstream stream) +{ + using DstAddr = TierAddr; + using SrcAddr = TierAddr; + + std::vector> t; + t.reserve(tasks.size()); + for (auto const& task : tasks) + t.push_back({std::get(task.dst), std::get(task.src)}); + + constexpr auto G = CacheTier::GPU_MEM; + constexpr auto H = CacheTier::HOST_MEM; + constexpr auto D = CacheTier::DISK; + + // clang-format off + if constexpr (DstTier == G && SrcTier == G) cuCheck(copyDeviceToDevice(std::move(t), static_cast(numBytes), stream)); + else if constexpr (DstTier == G && SrcTier == H) cuCheck(copyHostToDevice(std::move(t), static_cast(numBytes), stream)); + else if constexpr (DstTier == H && SrcTier == G) cuCheck(copyDeviceToHost(std::move(t), static_cast(numBytes), stream)); + else if constexpr (DstTier == H && SrcTier == H) cuCheck(copyHostToHost(std::move(t), static_cast(numBytes), stream)); + else if constexpr (DstTier == D && SrcTier == D) cuCheck(copyDiskToDisk(std::move(t), static_cast(numBytes), stream)); + else if constexpr (DstTier == H && SrcTier == D) cuCheck(copyDiskToHost(std::move(t), static_cast(numBytes), stream)); + else if constexpr (DstTier == D && SrcTier == H) cuCheck(copyHostToDisk(std::move(t), static_cast(numBytes), stream)); + // clang-format on +} + +// --------------------------------------------------------------------------- +// StagingBuffer +// --------------------------------------------------------------------------- + +StagingBuffer::StagingBuffer(StagingBufferManager& manager, size_t minSize, size_t maxSize, CUstream stream) + : mManager(manager) + , mStream(stream) +{ + if (minSize > manager.totalSize()) + { + throw std::invalid_argument("StagingBuffer: minSize exceeds total staging buffer size"); + } + + std::unique_lock lock(mManager.mMutex); + + // Compute how many contiguous grains to use. If the suffix cannot satisfy + // the required minimum, skip it and wrap before allocating. + size_t const minGrains = divUp(minSize, kGranularity); + size_t availableGrains = mManager.suggestNextMaxGrains(); + if (minGrains > availableGrains) + { + mManager.mNext = 0; + availableGrains = mManager.suggestNextMaxGrains(); + } + TLLM_CHECK_DEBUG(minGrains <= availableGrains); + + size_t const available = availableGrains * kGranularity; + size_t actualSize = std::min(maxSize, available); + actualSize = std::max(actualSize, minSize); + mNumGrains = divUp(actualSize, kGranularity); + TLLM_CHECK_DEBUG(mNumGrains <= availableGrains); + mSize = actualSize; + mStartGrain = mManager.mNext; + mManager.mNext += mNumGrains; + TLLM_CHECK_DEBUG(mManager.mNext <= mManager.numGrains()); + if (mManager.mNext == mManager.numGrains()) + { + mManager.mNext = 0; + } + + mAddress = mManager.baseAddress() + mStartGrain * kGranularity; + lock.unlock(); + + // Lock grains and collect their ready events for deduplicated waiting. + // Mirrors Python's stream_wait_events(stream, lock_and_consume_events()) which + // deduplicates via set() — adjacent grains often share the same event. + std::vector readyEvents; + readyEvents.reserve(mNumGrains); + for (size_t i = 0; i < mNumGrains; ++i) + { + GrainMetadata& g = mManager.mGrains[mStartGrain + i]; + g.mutex.lock(); + readyEvents.push_back(&g.readyEvent); + } + streamWaitEvents(reinterpret_cast(mStream), readyEvents); + for (size_t i = 0; i < mNumGrains; ++i) + mManager.mGrains[mStartGrain + i].readyEvent.close(); +} + +StagingBuffer::~StagingBuffer() +{ + // One shared completion event for all grains (all on the same stream → same completion point). + CachedCudaEvent finishEvent(reinterpret_cast(mStream)); + for (int i = static_cast(mNumGrains) - 1; i >= 0; --i) + { + GrainMetadata& g = mManager.mGrains[mStartGrain + static_cast(i)]; + g.readyEvent = finishEvent; + g.mutex.unlock(); + } +} + +// --------------------------------------------------------------------------- +// StagingBufferManager +// --------------------------------------------------------------------------- + +StagingBufferManager::StagingBufferManager(size_t size) + : mBuffer(size) + , mGrains(size / kGranularity) +{ + TLLM_CHECK_DEBUG(size % kGranularity == 0); +} + +StagingBuffer StagingBufferManager::acquire(size_t minSize, size_t maxSize, CUstream stream) +{ + return StagingBuffer(*this, minSize, maxSize, stream); +} + +// --------------------------------------------------------------------------- +// CopyEngine +// --------------------------------------------------------------------------- + +StagingBufferManager& CopyEngine::getStagingManager() +{ + if (!mStagingManager) + mStagingManager = std::make_unique(64u << 20u); // 64 MB + return *mStagingManager; +} + +// Two-hop transfer via host staging buffer (e.g., GPU→Disk or Disk→GPU). +// SrcTier → MidTier (staging) → DstTier. MidTier is the staging tier (HOST_MEM). +template +static void twoHopTransfer( + StagingBufferManager& manager, size_t numBytes, std::vector const& tasks, CUstream stream) +{ + size_t remaining = tasks.size(); + size_t offset = 0; + + while (remaining > 0) + { + StagingBuffer buf = manager.acquire(numBytes, numBytes * remaining, stream); + MemAddress addr = buf.address(); + size_t n = buf.size() / numBytes; + TLLM_CHECK_DEBUG(n > 0 && n <= remaining); + + // First hop: src → staging + { + std::vector hop1; + hop1.reserve(n); + for (size_t i = 0; i < n; ++i) + hop1.push_back({Address{addr + numBytes * i}, tasks[offset + i].src}); + dispatchCopy(hop1, numBytes, buf.stream()); + } + + // Second hop: staging → dst + { + std::vector hop2; + hop2.reserve(n); + for (size_t i = 0; i < n; ++i) + hop2.push_back({tasks[offset + i].dst, Address{addr + numBytes * i}}); + dispatchCopy(hop2, numBytes, buf.stream()); + } + + offset += n; + remaining -= n; + } +} + +void CopyEngine::transfer( + CacheTier dstTier, CacheTier srcTier, size_t numBytes, std::vector const& tasks, CUstream stream) +{ + // Nothing to copy. Also guards the two-hop path against a div-by-zero on + // `buf.size() / numBytes` when numBytes == 0. + if (tasks.empty() || numBytes == 0) + { + return; + } + + constexpr auto G = CacheTier::GPU_MEM; + constexpr auto H = CacheTier::HOST_MEM; + constexpr auto D = CacheTier::DISK; + + if (dstTier == G && srcTier == G) + dispatchCopy(tasks, numBytes, stream); + else if (dstTier == G && srcTier == H) + dispatchCopy(tasks, numBytes, stream); + else if (dstTier == G && srcTier == D) + twoHopTransfer(getStagingManager(), numBytes, tasks, stream); + else if (dstTier == H && srcTier == G) + dispatchCopy(tasks, numBytes, stream); + else if (dstTier == H && srcTier == H) + dispatchCopy(tasks, numBytes, stream); + else if (dstTier == H && srcTier == D) + dispatchCopy(tasks, numBytes, stream); + else if (dstTier == D && srcTier == G) + twoHopTransfer(getStagingManager(), numBytes, tasks, stream); + else if (dstTier == D && srcTier == H) + dispatchCopy(tasks, numBytes, stream); + else if (dstTier == D && srcTier == D) + dispatchCopy(tasks, numBytes, stream); + else + throw std::invalid_argument("CopyEngine::transfer: unsupported tier combination"); +} + +// --------------------------------------------------------------------------- +// Module-level singleton +// --------------------------------------------------------------------------- + +CopyEngine& globalCopyEngine() +{ + static CopyEngine engine; + return engine; +} + +void batchedCopy(CacheTier dstTier, CacheTier srcTier, size_t numBytes, std::vector tasks, CUstream stream) +{ + globalCopyEngine().transfer(dstTier, srcTier, numBytes, std::move(tasks), stream); +} + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/copyEngine.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/copyEngine.h new file mode 100644 index 000000000000..5fca2de91abc --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/copyEngine.h @@ -0,0 +1,199 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "kv_cache_manager_v2/common.h" +#include "kv_cache_manager_v2/utils/cudaEvent.h" +#include "kv_cache_manager_v2/utils/hostMem.h" + +#include "tensorrt_llm/common/assert.h" +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// --------------------------------------------------------------------------- +// CopyTask — source and destination address pair for a bulk copy. +// --------------------------------------------------------------------------- +struct CopyTask +{ + Address dst; + Address src; +}; + +// --------------------------------------------------------------------------- +// GrainMetadata — per-grain synchronization state inside StagingBufferManager. +// --------------------------------------------------------------------------- +struct GrainMetadata +{ + std::mutex mutex; + CachedCudaEvent readyEvent; // event protecting this grain + + GrainMetadata() + : readyEvent(CachedCudaEvent::makeNull()) + { + } + + GrainMetadata(GrainMetadata const&) = delete; + GrainMetadata& operator=(GrainMetadata const&) = delete; + GrainMetadata(GrainMetadata&&) = delete; + GrainMetadata& operator=(GrainMetadata&&) = delete; +}; + +class StagingBufferManager; + +// --------------------------------------------------------------------------- +// StagingBuffer — RAII handle to a slice of the StagingBufferManager's buffer. +// +// On construction, acquires grain locks and waits for readyEvents. +// On destruction, records a new event and releases grain locks. +// --------------------------------------------------------------------------- +class StagingBuffer +{ +public: + static constexpr size_t kGranularity = 1u << 20; // 1 MB grains + + StagingBuffer(StagingBufferManager& manager, size_t minSize, size_t maxSize, CUstream stream); + ~StagingBuffer(); + + StagingBuffer(StagingBuffer const&) = delete; + StagingBuffer& operator=(StagingBuffer const&) = delete; + StagingBuffer(StagingBuffer&&) = delete; + StagingBuffer& operator=(StagingBuffer&&) = delete; + + MemAddress address() const noexcept + { + return mAddress; + } + + size_t size() const noexcept + { + return mSize; + } + + CUstream stream() const noexcept + { + return mStream; + } + +private: + StagingBufferManager& mManager; + size_t mStartGrain{0}; + size_t mNumGrains{0}; + size_t mSize{0}; + MemAddress mAddress{0}; + CUstream mStream; +}; + +// --------------------------------------------------------------------------- +// StagingBufferManager — ring-buffer allocator over a CUDA-registered HostMem. +// Used for two-hop GPU↔Disk transfers. +// --------------------------------------------------------------------------- +class StagingBufferManager +{ +public: + static constexpr size_t kGranularity = StagingBuffer::kGranularity; // 1 MB + + explicit StagingBufferManager(size_t size); + + // Not movable: GrainMetadata has mutexes. + StagingBufferManager(StagingBufferManager const&) = delete; + StagingBufferManager& operator=(StagingBufferManager const&) = delete; + StagingBufferManager(StagingBufferManager&&) = delete; + StagingBufferManager& operator=(StagingBufferManager&&) = delete; + + // Acquire a staging slice. Thread-safe. + // minSize: minimum required bytes. maxSize: best-effort upper bound. + // Returns an RAII StagingBuffer that holds grain locks until destroyed. + StagingBuffer acquire(size_t minSize, size_t maxSize, CUstream stream); + + size_t totalSize() const noexcept + { + TLLM_CHECK_DEBUG_WITH_INFO( + mGrains.size() * kGranularity == mBuffer.size(), "grain count * granularity must equal buffer size"); + return mBuffer.size(); + } + + size_t numGrains() const noexcept + { + return mGrains.size(); + } + + MemAddress baseAddress() const noexcept + { + return mBuffer.address(); + } + +private: + friend class StagingBuffer; + + // Caller must hold mMutex. + size_t suggestNextMaxGrains() const noexcept + { + return numGrains() - mNext; + } + + std::mutex mMutex; + HostMem mBuffer; + std::vector mGrains; + size_t mNext{0}; +}; + +// --------------------------------------------------------------------------- +// CopyEngine — dispatches bulk transfers between cache tiers. +// +// Single-hop pairs call the appropriate copy function from kvCacheManagerV2Utils. +// Two-hop pairs (GPU↔Disk) route through a lazily-allocated StagingBuffer. +// --------------------------------------------------------------------------- +class CopyEngine +{ +public: + CopyEngine() = default; + ~CopyEngine() = default; + + CopyEngine(CopyEngine const&) = delete; + CopyEngine& operator=(CopyEngine const&) = delete; + + // Transfer num_bytes per task. tasks must all share the same (dstTier, srcTier). + // stream: the CUDA stream on which GPU ops are enqueued. + void transfer( + CacheTier dstTier, CacheTier srcTier, size_t numBytes, std::vector const& tasks, CUstream stream); + + void close() noexcept + { + mStagingManager.reset(); + } + +private: + StagingBufferManager& getStagingManager(); + + std::unique_ptr mStagingManager; +}; + +// --------------------------------------------------------------------------- +// Module-level singleton (mirrors Python's _copy_engine global). +// --------------------------------------------------------------------------- +CopyEngine& globalCopyEngine(); + +// Convenience wrapper used by higher layers. +void batchedCopy(CacheTier dstTier, CacheTier srcTier, size_t numBytes, std::vector tasks, CUstream stream); + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/cudaVirtMem.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/cudaVirtMem.cpp new file mode 100644 index 000000000000..240f7d08f1d5 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/cudaVirtMem.cpp @@ -0,0 +1,217 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "kv_cache_manager_v2/cudaVirtMem.h" +#include "kv_cache_manager_v2/exceptions.h" +#include "kv_cache_manager_v2/utils/math.h" +#include "tensorrt_llm/common/assert.h" + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// --------------------------------------------------------------------------- +// PooledPhysMemAllocator +// --------------------------------------------------------------------------- + +static bool isPropSupported(CUmemAllocationProp const& prop) +{ + CUmemGenericAllocationHandle handle; + CUresult err = cuMemCreate(&handle, 2ULL << 20, &prop, 0); + if (err == CUDA_ERROR_NOT_PERMITTED || err == CUDA_ERROR_NOT_SUPPORTED || err == CUDA_ERROR_INVALID_DEVICE + || err == CUDA_ERROR_INVALID_VALUE) + { + return false; + } + if (err == CUDA_SUCCESS) + { + cuMemRelease(handle); + return true; + } + throw CuError(err); +} + +// --------------------------------------------------------------------------- +// PhysMemChunk +// --------------------------------------------------------------------------- + +PhysMemWrapper::PhysMemWrapper(size_t size, CUmemAllocationProp const& prop) +{ + cuCheck(cuMemCreate(&mHandle, size, &prop, 0)); +} + +PhysMemWrapper::~PhysMemWrapper() +{ + cuMemRelease(mHandle); +} + +// --------------------------------------------------------------------------- +// PooledPhysMemAllocator +// --------------------------------------------------------------------------- + +PooledPhysMemAllocator::PooledPhysMemAllocator(size_t physMemSize) + : mPhysMemSize(physMemSize) + , mPool([this]() -> PhysMemWrapper* { return new PhysMemWrapper(mPhysMemSize, mProp); }, + [](PhysMemWrapper* chunk) { delete chunk; }) +{ + // Get current device. + cuCheck(cuCtxGetDevice(&mDeviceId)); + + // Build the best allocation property. + mProp.type = CU_MEM_ALLOCATION_TYPE_PINNED; + mProp.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + mProp.location.id = mDeviceId; + mProp.allocFlags.gpuDirectRDMACapable = 1; + mProp.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC; + + if (!isPropSupported(mProp)) + { + mProp.requestedHandleTypes = CU_MEM_HANDLE_TYPE_NONE; + if (!isPropSupported(mProp)) + { + mProp.allocFlags.gpuDirectRDMACapable = 0; + if (!isPropSupported(mProp)) + { + throw std::runtime_error("PooledPhysMemAllocator: no supported physical memory allocation property"); + } + } + } +} + +PooledPhysMemAllocator::~PooledPhysMemAllocator() +{ + mPool.clear(); +} + +PooledPhysMemAllocator::PooledPhysMem PooledPhysMemAllocator::acquire() +{ + return mPool.get(); +} + +// --------------------------------------------------------------------------- +// VirtMem +// --------------------------------------------------------------------------- + +VirtMem::VirtMem(size_t vmSize, PooledPhysMemAllocator& physMemAllocator, size_t initNumPhysMem) + : mVmSize(vmSize) + , mPhysMemAllocator(physMemAllocator) +{ + TLLM_CHECK_DEBUG(vmSize % physMemAllocator.physMemSize() == 0); + + cuCheck(cuMemAddressReserve(&mAddr, vmSize, 0, 0, 0)); + + mAccessDesc.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + mAccessDesc.location.id = physMemAllocator.deviceId(); + mAccessDesc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; + + extend(initNumPhysMem); +} + +VirtMem::~VirtMem() noexcept +{ + try + { + destroy(); + } + catch (...) + { + // Destructors cannot surface CUDA cleanup failures. Explicit destroy() + // still reports them to match Python VirtMem.destroy(). + } +} + +void VirtMem::push(PooledPhysMemAllocator::PooledPhysMem handle) +{ + size_t physSize = mPhysMemAllocator.physMemSize(); + CUdeviceptr offset = mAddr + physSize * static_cast(mPhysHandles.size()); + TLLM_CHECK_DEBUG(physSize * (mPhysHandles.size() + 1) <= mVmSize); + + cuCheck(cuMemMap(offset, physSize, 0, handle->handle(), 0)); + cuCheck(cuMemSetAccess(offset, physSize, &mAccessDesc, 1)); + mPhysHandles.push_back(std::move(handle)); +} + +void VirtMem::pop() +{ + TLLM_CHECK_DEBUG(!mPhysHandles.empty()); + size_t physSize = mPhysMemAllocator.physMemSize(); + CUdeviceptr offset = mAddr + physSize * (mPhysHandles.size() - 1); + cuCheck(cuMemUnmap(offset, physSize)); + mPhysHandles.pop_back(); // PhysMemHandle destructor returns to pool +} + +void VirtMem::extend(size_t numToAdd) +{ + size_t old = numPhysMem(); + try + { + for (size_t i = 0; i < numToAdd; ++i) + { + push(mPhysMemAllocator.acquire()); + } + } + catch (...) + { + // Rollback: remove any newly-mapped chunks. + while (numPhysMem() > old) + { + pop(); + } + throw; + } +} + +void VirtMem::shrink(size_t numToRemove) +{ + cuCheck(cuCtxSynchronize()); + for (size_t i = 0; i < numToRemove; ++i) + { + pop(); + } +} + +void VirtMem::realloc(size_t numBytes) +{ + size_t physSize = mPhysMemAllocator.physMemSize(); + size_t required = divUp(numBytes, physSize); + size_t current = numPhysMem(); + if (required > current) + { + extend(required - current); + } + else if (required < current) + { + shrink(current - required); + } +} + +void VirtMem::destroy() +{ + if (mVmSize == 0) + { + return; + } + cuCheck(cuCtxSynchronize()); + while (!mPhysHandles.empty()) + { + pop(); + } + cuCheck(cuMemAddressFree(mAddr, mVmSize)); + mAddr = 0; + mVmSize = 0; +} + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/cudaVirtMem.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/cudaVirtMem.h new file mode 100644 index 000000000000..5c560f036594 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/cudaVirtMem.h @@ -0,0 +1,163 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "kv_cache_manager_v2/common.h" +#include "kv_cache_manager_v2/utils/cudaEvent.h" + +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// --------------------------------------------------------------------------- +// PhysMemChunk — RAII wrapper for a CUmemGenericAllocationHandle. +// Constructor calls cuMemCreate, destructor calls cuMemRelease. +// --------------------------------------------------------------------------- +class PhysMemWrapper +{ +public: + PhysMemWrapper(size_t size, CUmemAllocationProp const& prop); + ~PhysMemWrapper(); + + PhysMemWrapper(PhysMemWrapper const&) = delete; + PhysMemWrapper& operator=(PhysMemWrapper const&) = delete; + + [[nodiscard]] CUmemGenericAllocationHandle handle() const noexcept + { + return mHandle; + } + +private: + CUmemGenericAllocationHandle mHandle; +}; + +// --------------------------------------------------------------------------- +// PooledPhysMemAllocator — creates and pools physical GPU memory chunks. +// Mirrors _cuda_virt_mem.py::PooledPhysMemAllocator. +// --------------------------------------------------------------------------- +class PooledPhysMemAllocator +{ +public: + using PooledPhysMem = SimplePool::PoolItem; + + // physMemSize: size of each physical chunk in bytes. + explicit PooledPhysMemAllocator(size_t physMemSize); + ~PooledPhysMemAllocator(); + + PooledPhysMemAllocator(PooledPhysMemAllocator const&) = delete; + PooledPhysMemAllocator& operator=(PooledPhysMemAllocator const&) = delete; + + // Borrow a physical memory handle from the pool (or allocate a new one). + // Dropping the returned PhysMemHandle returns it to the pool. + [[nodiscard]] PooledPhysMem acquire(); + + // Release all cached (unused) physical memory back to the driver. + // Mirrors Python PooledPhysMemAllocator.clear(). + void clear() + { + mPool.clear(); + } + + [[nodiscard]] size_t physMemSize() const noexcept + { + return mPhysMemSize; + } + + int deviceId() const noexcept + { + return mDeviceId; + } + +private: + size_t mPhysMemSize{}; + int mDeviceId{}; + CUmemAllocationProp mProp{}; + SimplePool mPool; +}; + +// --------------------------------------------------------------------------- +// VirtMem — a virtual address range backed by physical GPU memory chunks. +// Physical chunks are mapped/unmapped at the end of the range (stack discipline). +// Mirrors _cuda_virt_mem.py::VirtMem. +// +// Invariant: mappedBytes() == numPhysMem() * physMemSize() +// mappedBytes() <= virtualBytes() +// --------------------------------------------------------------------------- +class VirtMem +{ +public: + // vmSize: total virtual address space reserved (bytes). + // Must be a multiple of physMemSize. + // physMemAllocator: shared allocator for physical chunks. + // initNumPhysMem: number of physical chunks to map immediately. + VirtMem(size_t vmSize, PooledPhysMemAllocator& physMemAllocator, size_t initNumPhysMem = 0); + ~VirtMem() noexcept; + + VirtMem(VirtMem const&) = delete; + VirtMem& operator=(VirtMem const&) = delete; + + // Map numPhysMem additional chunks at the top of the address range. + void extend(size_t numPhysMem); + + // Unmap numPhysMem chunks from the top (synchronizes CUDA first). + void shrink(size_t numPhysMem); + + // Adjust mapped bytes to exactly numBytes (extend or shrink). + void realloc(size_t numBytes); + + void destroy(); + + [[nodiscard]] MemAddress address() const noexcept + { + return static_cast(mAddr); + } + + [[nodiscard]] size_t physMemSize() const noexcept + { + return mPhysMemAllocator.physMemSize(); + } + + [[nodiscard]] size_t mappedBytes() const noexcept + { + return mPhysMemAllocator.physMemSize() * numPhysMem(); + } + + [[nodiscard]] size_t virtualBytes() const noexcept + { + return mVmSize; + } + + [[nodiscard]] size_t numPhysMem() const noexcept + { + return mPhysHandles.size(); + } + +private: + void push(PooledPhysMemAllocator::PooledPhysMem handle); + void pop(); + + CUdeviceptr mAddr = 0; + size_t mVmSize = 0; + PooledPhysMemAllocator& mPhysMemAllocator; + std::vector mPhysHandles; + CUmemAccessDesc mAccessDesc{}; +}; + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cpp new file mode 100644 index 000000000000..73fec6f57cd1 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cpp @@ -0,0 +1,683 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "kv_cache_manager_v2/eventManager.h" + +#include "kv_cache_manager_v2/blockRadixTree.h" +#include "kv_cache_manager_v2/page.h" +#include "tensorrt_llm/common/logger.h" + +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ +namespace +{ + +constexpr uint32_t kUint32HashConst = 0x045D9F3BU; +constexpr uint64_t kUint64HashConst1 = 0xBF58476D1CE4E5B9ULL; +constexpr uint64_t kUint64HashConst2 = 0x94D049BB133111EBULL; +constexpr uint32_t kHashCombineConst = 0x9E3779B9U; +constexpr uint64_t kParentHashConst = 0xBF58476D1CE4E5B9ULL; + +uint64_t hash32Mix(int64_t input, uint64_t seed) +{ + uint32_t value = static_cast(input); + value = ((value >> 16U) ^ value) * kUint32HashConst; + value = ((value >> 16U) ^ value) * kUint32HashConst; + value = (value >> 16U) ^ value; + value += kHashCombineConst; + return seed ^ (static_cast(value) + (seed << 6U) + (seed >> 2U)); +} + +uint64_t hash64Mix(int64_t input, uint64_t seed) +{ + uint64_t value = static_cast(input); + value = (value ^ (value >> 30U)) * kUint64HashConst1; + value = (value ^ (value >> 27U)) * kUint64HashConst2; + value ^= value >> 31U; + return seed ^ (value + static_cast(kHashCombineConst) + (seed << 6U) + (seed >> 2U)); +} + +} // namespace + +EventManager::EventManager(int maxKvEventEntries, int windowSize, std::optional attentionDpRank, + AttentionDpGatherFn attentionDpGather, std::string hashAlgo, std::map windowSizeByLayerGroup) + : mMaxKvEventEntries(maxKvEventEntries) + , mWindowSize(windowSize) + , mWindowSizeByLayerGroup(std::move(windowSizeByLayerGroup)) + , mAttentionDpRank(attentionDpRank) + , mAttentionDpGather(std::move(attentionDpGather)) +{ + std::tie(mHashAlgo, mHashAlgoName) = parseHashAlgorithm(hashAlgo); +} + +std::pair EventManager::parseHashAlgorithm(std::string const& hashAlgo) +{ + if (hashAlgo == "auto" || hashAlgo == "v1_block_key") + { + return {HashAlgorithm::kV1, "v1_block_key"}; + } + if (hashAlgo == "v2_sha256") + { + return {HashAlgorithm::kV2Sha256, "v2_sha256"}; + } + if (hashAlgo == "v2_sha256_64") + { + return {HashAlgorithm::kV2Sha256_64, "v2_sha256_64"}; + } + throw std::invalid_argument("Unsupported V2 KV cache event hash algorithm: " + hashAlgo); +} + +void EventManager::addCreatedEvent( + std::vector numBlocksPerCacheLevel, std::optional> layerGroupIds) +{ + if (mMaxKvEventEntries <= 0) + { + return; + } + std::lock_guard lock(mMutex); + KVCacheCreatedData data{std::move(numBlocksPerCacheLevel)}; + if (!layerGroupIds.has_value()) + { + addEventUnlocked(std::move(data), std::nullopt); + return; + } + for (int layerGroupId : *layerGroupIds) + { + addEventUnlocked(data, layerGroupId); + } +} + +void EventManager::setLayerGroupWindowSizes(std::map windowSizes) +{ + std::lock_guard lock(mMutex); + mWindowSizeByLayerGroup = std::move(windowSizes); +} + +void EventManager::addStoredEvent(KVCacheStoredData data, EventLayerGroupId layerGroupId) +{ + if (data.blocks.empty() || mMaxKvEventEntries <= 0) + { + return; + } + std::lock_guard lock(mMutex); + flushRemovedEventUnlocked(layerGroupId); + addStoredEventUnlocked(std::move(data), layerGroupId); +} + +void EventManager::addRemovedEvent(std::vector blockHashes, EventLayerGroupId layerGroupId) +{ + if (blockHashes.empty() || mMaxKvEventEntries <= 0) + { + return; + } + std::lock_guard lock(mMutex); + enqueueRemovedEventUnlocked(std::move(blockHashes), layerGroupId); +} + +void EventManager::addUpdatedEvent(EventBlockHash blockHash, std::optional cacheLevel, + std::optional priority, EventLayerGroupId layerGroupId) +{ + if ((!cacheLevel.has_value() && !priority.has_value()) || mMaxKvEventEntries <= 0) + { + return; + } + std::lock_guard lock(mMutex); + addEventUnlocked(KVCacheUpdatedData{std::move(blockHash), cacheLevel, priority}, layerGroupId); +} + +void EventManager::addUpdatedEvent(Digest const& blockKey, std::optional cacheLevel, + std::optional priority, EventLayerGroupId layerGroupId) +{ + if ((!cacheLevel.has_value() && !priority.has_value()) || mMaxKvEventEntries <= 0) + { + return; + } + std::lock_guard lock(mMutex); + auto const state = mStoredBlocks.find(blockKey); + if (state == mStoredBlocks.end()) + { + return; + } + addEventUnlocked(KVCacheUpdatedData{state->second.blockHash, cacheLevel, priority}, layerGroupId); +} + +void EventManager::addStoredBlock(Block const& block) +{ + if (mMaxKvEventEntries <= 0) + { + return; + } + std::lock_guard lock(mMutex); + addStoredBlockUnlocked(block); +} + +void EventManager::addStoredBlockUnlocked(Block const& block) +{ + std::set lifeCycleIds; + for (LifeCycleId lifeCycle{0}; lifeCycle < block.storage.size(); ++lifeCycle) + { + if (block.storage[lifeCycle] != nullptr) + { + lifeCycleIds.insert(lifeCycle.value()); + } + } + if (lifeCycleIds.empty()) + { + return; + } + + EventBlockHash blockHash = hashFromBlock(block); + mStoredBlocks.insert_or_assign(block.key, StoredBlockState{blockHash, lifeCycleIds}); + auto parentHash = parentHashFromBlock(block); + for (int lifeCycleId : lifeCycleIds) + { + auto blockData = storedBlockFromBlock(block, std::set{lifeCycleId}); + if (blockData.has_value()) + { + flushRemovedEventUnlocked(lifeCycleId); + addStoredEventUnlocked(KVCacheStoredData{parentHash, {std::move(*blockData)}}, lifeCycleId); + } + } +} + +void EventManager::addStoredLifeCycle(Block const& block, LifeCycleId lifeCycle) +{ + if (mMaxKvEventEntries <= 0) + { + return; + } + std::lock_guard lock(mMutex); + auto state = mStoredBlocks.find(block.key); + if (state == mStoredBlocks.end()) + { + addStoredBlockUnlocked(block); + return; + } + int const lifeCycleId = lifeCycle.value(); + if (state->second.lifeCycleIds.count(lifeCycleId) != 0) + { + return; + } + auto blockData = storedBlockFromBlock(block, std::set{lifeCycleId}); + if (!blockData.has_value()) + { + return; + } + state->second.lifeCycleIds.insert(lifeCycleId); + flushRemovedEventUnlocked(lifeCycleId); + addStoredEventUnlocked(KVCacheStoredData{parentHashFromBlock(block), {std::move(*blockData)}}, lifeCycleId); +} + +void EventManager::addRemovedBlock(Digest const& blockKey) +{ + if (mMaxKvEventEntries <= 0) + { + return; + } + std::lock_guard lock(mMutex); + auto state = mStoredBlocks.find(blockKey); + if (state == mStoredBlocks.end()) + { + return; + } + EventBlockHash blockHash = state->second.blockHash; + auto lifeCycleIds = state->second.lifeCycleIds; + mStoredBlocks.erase(state); + dropHashCache(blockKey); + + if (lifeCycleIds.empty()) + { + enqueueRemovedEventUnlocked({std::move(blockHash)}, std::nullopt); + return; + } + for (int lifeCycleId : lifeCycleIds) + { + enqueueRemovedEventUnlocked({blockHash}, lifeCycleId); + } +} + +void EventManager::addRemovedLifeCycle(Digest const& blockKey, LifeCycleId lifeCycle) +{ + if (mMaxKvEventEntries <= 0) + { + return; + } + std::lock_guard lock(mMutex); + auto state = mStoredBlocks.find(blockKey); + int const lifeCycleId = lifeCycle.value(); + if (state == mStoredBlocks.end() || state->second.lifeCycleIds.erase(lifeCycleId) == 0) + { + return; + } + EventBlockHash blockHash = state->second.blockHash; + if (state->second.lifeCycleIds.empty()) + { + mStoredBlocks.erase(state); + dropHashCache(blockKey); + } + enqueueRemovedEventUnlocked({std::move(blockHash)}, lifeCycleId); +} + +void EventManager::addCacheLevelUpdated( + Digest const& blockKey, CacheLevel oldLevel, CacheLevel newLevel, LifeCycleId lifeCycle) +{ + if (mMaxKvEventEntries <= 0) + { + return; + } + std::lock_guard lock(mMutex); + auto state = mStoredBlocks.find(blockKey); + if (state == mStoredBlocks.end()) + { + return; + } + addEventUnlocked( + KVCacheUpdatedData{state->second.blockHash, KVCacheEventDiff{oldLevel.value(), newLevel.value()}, std::nullopt}, + lifeCycle.value()); +} + +void EventManager::addStoredEventUnlocked(KVCacheStoredData data, EventLayerGroupId layerGroupId) +{ + bool const hasPendingRemovedEvents = !mLatestRemovedBlockHashes.empty(); + auto latest = mLatestStoredEventIds.find(layerGroupId); + if (!hasPendingRemovedEvents && latest != mLatestStoredEventIds.end()) + { + auto pending = std::find_if(mPendingEvents.rbegin(), mPendingEvents.rend(), + [&](KVCacheEvent const& event) { return event.eventId == latest->second; }); + if (pending == mPendingEvents.rend()) + { + throw std::logic_error("Stored event coalescing lost the pending event"); + } + if (auto* stored = std::get_if(&pending->data); stored != nullptr && !stored->blocks.empty() + && data.parentHash.has_value() && stored->blocks.back().blockHash == *data.parentHash) + { + std::move(data.blocks.begin(), data.blocks.end(), std::back_inserter(stored->blocks)); + return; + } + } + + auto& event = addEventUnlocked(std::move(data), layerGroupId); + mLatestStoredEventIds.insert_or_assign(layerGroupId, event.eventId); +} + +void EventManager::enqueueRemovedEventUnlocked(std::vector blockHashes, EventLayerGroupId layerGroupId) +{ + if (blockHashes.empty()) + { + return; + } + auto& pending = mLatestRemovedBlockHashes[layerGroupId]; + std::move(blockHashes.begin(), blockHashes.end(), std::back_inserter(pending)); + mLatestStoredEventIds.erase(layerGroupId); +} + +void EventManager::flushRemovedEventUnlocked(EventLayerGroupId layerGroupId) +{ + auto removed = mLatestRemovedBlockHashes.find(layerGroupId); + if (removed == mLatestRemovedBlockHashes.end() || removed->second.empty()) + { + return; + } + auto blockHashes = std::move(removed->second); + mLatestRemovedBlockHashes.erase(removed); + addEventUnlocked(KVCacheRemovedData{std::move(blockHashes)}, layerGroupId); +} + +void EventManager::flushAllRemovedEventsUnlocked() +{ + while (!mLatestRemovedBlockHashes.empty()) + { + flushRemovedEventUnlocked(mLatestRemovedBlockHashes.begin()->first); + } +} + +KVCacheEvent& EventManager::addEventUnlocked(KVCacheEventData data, EventLayerGroupId layerGroupId) +{ + if (mMaxKvEventEntries <= 0) + { + throw std::logic_error("Cannot add an event when the event queue is disabled"); + } + if (!std::holds_alternative(data)) + { + flushAllRemovedEventsUnlocked(); + } + mPendingEvents.push_back(KVCacheEvent{ + mNextEventId++, std::move(data), getWindowSize(layerGroupId), mHashAlgoName, mAttentionDpRank, layerGroupId}); + if (!std::holds_alternative(mPendingEvents.back().data)) + { + mLatestStoredEventIds.erase(layerGroupId); + } + return mPendingEvents.back(); +} + +std::vector EventManager::drainPendingEventsUnlocked() +{ + flushAllRemovedEventsUnlocked(); + auto events = std::move(mPendingEvents); + mPendingEvents.clear(); + mLatestStoredEventIds.clear(); + return events; +} + +void EventManager::publishEventsUnlocked(std::vector events, std::optional maxKvEventEntries) +{ + if (events.empty()) + { + return; + } + int const capacity = maxKvEventEntries.value_or(mMaxKvEventEntries); + std::move(events.begin(), events.end(), std::back_inserter(mEvents)); + while (static_cast(mEvents.size()) > capacity) + { + mEvents.pop_front(); + } +} + +std::vector EventManager::trimEvents(std::vector events, int maxKvEventEntries) +{ + if (maxKvEventEntries <= 0) + { + return {}; + } + if (static_cast(events.size()) > maxKvEventEntries) + { + events.erase(events.begin(), events.end() - maxKvEventEntries); + } + return events; +} + +void EventManager::flushIterationEvents() +{ + if (mAttentionDpGather) + { + std::vector localEvents; + { + std::lock_guard lock(mMutex); + localEvents = trimEvents(drainPendingEventsUnlocked(), mMaxKvEventEntries); + } + auto gatheredEvents = mAttentionDpGather(localEvents); + if (mAttentionDpRank != std::optional{0}) + { + return; + } + + std::vector events; + for (auto& rankEvents : gatheredEvents) + { + auto trimmed = trimEvents(std::move(rankEvents), mMaxKvEventEntries); + std::move(trimmed.begin(), trimmed.end(), std::back_inserter(events)); + } + { + std::lock_guard lock(mMutex); + publishEventsUnlocked(std::move(events), mMaxKvEventEntries * std::max(1, gatheredEvents.size())); + } + mCondition.notify_all(); + return; + } + + { + std::lock_guard lock(mMutex); + publishEventsUnlocked(drainPendingEventsUnlocked()); + } + mCondition.notify_all(); +} + +std::vector EventManager::getLatestEvents(std::optional timeoutMs) +{ + std::unique_lock lock(mMutex); + if (mEvents.empty() && !timeoutMs.has_value()) + { + mCondition.wait(lock, [&] { return !mEvents.empty(); }); + } + else if (mEvents.empty() && *timeoutMs > 0) + { + mCondition.wait_for( + lock, std::chrono::duration(*timeoutMs), [&] { return !mEvents.empty(); }); + } + std::vector events; + events.reserve(mEvents.size()); + std::move(mEvents.begin(), mEvents.end(), std::back_inserter(events)); + mEvents.clear(); + return events; +} + +int EventManager::getWindowSize(EventLayerGroupId layerGroupId) const +{ + if (!layerGroupId.has_value()) + { + return mWindowSize; + } + auto const windowSize = mWindowSizeByLayerGroup.find(*layerGroupId); + return windowSize == mWindowSizeByLayerGroup.end() ? mWindowSize : windowSize->second; +} + +std::string EventManager::digestToHex(Digest const& digest) +{ + constexpr char kHex[] = "0123456789abcdef"; + std::string result; + result.resize(digest.size() * 2); + for (size_t i = 0; i < digest.size(); ++i) + { + auto const value = std::to_integer(digest[i]); + result[2 * i] = kHex[value >> 4U]; + result[2 * i + 1] = kHex[value & 0x0FU]; + } + return result; +} + +uint64_t EventManager::truncateDigestToInt64(Digest const& digest) +{ + uint64_t result = 0; + for (int i = 0; i < 8; ++i) + { + result = (result << 8U) | std::to_integer(digest[static_cast(i)]); + } + return result; +} + +EventBlockHash EventManager::normalizeDigest(Digest const& digest) const +{ + if (mHashAlgo == HashAlgorithm::kV2Sha256_64) + { + return truncateDigestToInt64(digest); + } + return digestToHex(digest); +} + +EventBlockHash EventManager::hashFromBlock(Block const& block) +{ + if (mHashAlgo == HashAlgorithm::kV1) + { + return v1HashFromBlock(block); + } + return normalizeDigest(block.key); +} + +std::optional EventManager::parentHashFromBlock(Block const& block) +{ + if (block.prev == nullptr) + { + throw std::logic_error("Cannot hash an orphan KV cache block"); + } + if (block.prev->type() == NodeBase::Type::kROOT_BLOCK) + { + return std::nullopt; + } + return hashFromBlock(*static_cast(block.prev)); +} + +std::optional EventManager::storedBlockFromBlock( + Block const& block, std::optional> const& lifeCycleIds) +{ + CacheLevel cacheLevel = kGpuLevel; + Priority priority = kPriorityDefault; + bool foundPage = false; + for (LifeCycleId lifeCycle{0}; lifeCycle < block.storage.size(); ++lifeCycle) + { + if (lifeCycleIds.has_value() && lifeCycleIds->count(lifeCycle.value()) == 0) + { + continue; + } + auto const* page = block.storage[lifeCycle]; + if (page != nullptr) + { + cacheLevel = page->cacheLevel; + priority = page->priority; + foundPage = true; + break; + } + } + if (lifeCycleIds.has_value() && !foundPage) + { + return std::nullopt; + } + + std::vector tokens; + tokens.reserve(block.tokens.size()); + for (auto const& token : block.tokens) + { + if (auto const* tokenId = std::get_if(&token)) + { + UniqueToken uniqueToken; + uniqueToken.tokenId = EventTokenId{std::in_place_index<0>, *tokenId}; + tokens.push_back(std::move(uniqueToken)); + } + else + { + UniqueToken uniqueToken; + uniqueToken.tokenId + = EventTokenId{std::in_place_index<1>, digestToHex(std::get(token).digest())}; + tokens.push_back(std::move(uniqueToken)); + } + } + return KVCacheStoredBlockData{ + hashFromBlock(block), std::move(tokens), cacheLevel.value(), priority, {}, std::nullopt}; +} + +uint64_t EventManager::hashV1BlockKey(std::vector const& tokens, uint64_t parentHash, + std::optional loraTaskId, std::optional cacheSaltId) +{ + uint64_t seed = static_cast(tokens.size()) ^ (parentHash * kParentHashConst); + if (parentHash == 0 && cacheSaltId.has_value()) + { + seed = hash64Mix(*cacheSaltId, seed); + } + for (TokenId token : tokens) + { + seed = hash32Mix(token, seed); + } + if (loraTaskId.has_value()) + { + seed = hash64Mix(*loraTaskId, seed); + } + return seed; +} + +uint64_t EventManager::v1HashFromBlock(Block const& block) +{ + if (auto const cached = mV1HashByBlockKey.find(block.key); cached != mV1HashByBlockKey.end()) + { + return cached->second; + } + + std::vector chain; + NodeBase const* current = █ + uint64_t parentHash = 0; + bool parentIsV1Compatible = true; + V1RootAttrs rootAttrs; + while (current->type() == NodeBase::Type::kBLOCK) + { + auto const* currentBlock = static_cast(current); + if (auto const cached = mV1HashByBlockKey.find(currentBlock->key); cached != mV1HashByBlockKey.end()) + { + parentHash = cached->second; + parentIsV1Compatible = mV1HashCompatibleKeys.count(currentBlock->key) != 0; + rootAttrs = mV1RootAttrsByBlockKey.at(currentBlock->key); + break; + } + chain.push_back(currentBlock); + current = currentBlock->prev; + if (current == nullptr) + { + throw std::logic_error("Cannot hash an orphan KV cache block"); + } + } + if (current->type() == NodeBase::Type::kROOT_BLOCK) + { + auto const& reuseScope = static_cast(current)->reuseScope; + rootAttrs = {reuseScope.loraId, reuseScope.salt}; + } + + for (auto chainIter = chain.rbegin(); chainIter != chain.rend(); ++chainIter) + { + Block const& currentBlock = **chainIter; + std::vector textTokens; + textTokens.reserve(currentBlock.tokens.size()); + if (parentIsV1Compatible) + { + for (auto const& token : currentBlock.tokens) + { + auto const* tokenId = std::get_if(&token); + if (tokenId == nullptr) + { + parentIsV1Compatible = false; + break; + } + textTokens.push_back(*tokenId); + } + } + if (parentIsV1Compatible) + { + parentHash = hashV1BlockKey(textTokens, parentHash, rootAttrs.first, rootAttrs.second); + mV1HashCompatibleKeys.insert(currentBlock.key); + } + else + { + parentHash = fallbackV1Hash(currentBlock.key); + } + mV1HashByBlockKey.insert_or_assign(currentBlock.key, parentHash); + mV1RootAttrsByBlockKey.insert_or_assign(currentBlock.key, rootAttrs); + } + return parentHash; +} + +uint64_t EventManager::fallbackV1Hash(Digest const& blockKey) +{ + if (!mWarnedV1HashFallback) + { + TLLM_LOG_WARNING( + "V2 KV cache event hash algorithm v1_block_key only matches v1 for text-token radix blocks. " + "Falling back to truncated SHA-256 block hash for unsupported blocks."); + mWarnedV1HashFallback = true; + } + return truncateDigestToInt64(blockKey); +} + +void EventManager::dropHashCache(Digest const& blockKey) +{ + mV1HashByBlockKey.erase(blockKey); + mV1HashCompatibleKeys.erase(blockKey); + mV1RootAttrsByBlockKey.erase(blockKey); +} + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.h new file mode 100644 index 000000000000..871a8bf4199a --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.h @@ -0,0 +1,262 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "kv_cache_manager_v2/common.h" +#include "kv_cache_manager_v2/eventSink.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +using EventBlockHash = std::variant; +using EventTokenId = std::variant; +using EventLayerGroupId = std::optional; + +struct UniqueToken +{ + EventTokenId tokenId; + int64_t tokenExtraId = 0; + + bool operator==(UniqueToken const& other) const + { + return tokenId == other.tokenId && tokenExtraId == other.tokenExtraId; + } +}; + +struct KVCacheCreatedData +{ + std::vector numBlocksPerCacheLevel; + + bool operator==(KVCacheCreatedData const& other) const + { + return numBlocksPerCacheLevel == other.numBlocksPerCacheLevel; + } +}; + +struct MmKey +{ + std::string hash; + int startOffset = 0; + std::optional uuid; + bool hasUuidField = false; + + bool operator==(MmKey const& other) const + { + return hash == other.hash && startOffset == other.startOffset && uuid == other.uuid + && hasUuidField == other.hasUuidField; + } +}; + +struct KVCacheStoredBlockData +{ + EventBlockHash blockHash; + std::vector tokens; + int cacheLevel = kGpuLevel.value(); + int priority = kPriorityDefault; + std::vector mmKeys; + std::optional cacheSalt; + + bool operator==(KVCacheStoredBlockData const& other) const + { + return blockHash == other.blockHash && tokens == other.tokens && cacheLevel == other.cacheLevel + && priority == other.priority && mmKeys == other.mmKeys && cacheSalt == other.cacheSalt; + } +}; + +struct KVCacheStoredData +{ + std::optional parentHash; + std::vector blocks; + + bool operator==(KVCacheStoredData const& other) const + { + return parentHash == other.parentHash && blocks == other.blocks; + } +}; + +struct KVCacheRemovedData +{ + std::vector blockHashes; + + bool operator==(KVCacheRemovedData const& other) const + { + return blockHashes == other.blockHashes; + } +}; + +struct KVCacheEventDiff +{ + int oldValue = 0; + int newValue = 0; + + bool operator==(KVCacheEventDiff const& other) const + { + return oldValue == other.oldValue && newValue == other.newValue; + } +}; + +struct KVCacheUpdatedData +{ + EventBlockHash blockHash; + std::optional cacheLevel; + std::optional priority; + + bool operator==(KVCacheUpdatedData const& other) const + { + return blockHash == other.blockHash && cacheLevel == other.cacheLevel && priority == other.priority; + } +}; + +using KVCacheEventData = std::variant; + +struct KVCacheEvent +{ + int64_t eventId = 0; + KVCacheEventData data; + int windowSize = 0; + std::optional hashAlgo; + std::optional attentionDpRank; + EventLayerGroupId layerGroupId; + + bool operator==(KVCacheEvent const& other) const + { + return eventId == other.eventId && data == other.data && windowSize == other.windowSize + && hashAlgo == other.hashAlgo && attentionDpRank == other.attentionDpRank + && layerGroupId == other.layerGroupId; + } +}; + +class EventManager final : public EventSink +{ +public: + using AttentionDpGatherFn = std::function>(std::vector const&)>; + + EventManager(int maxKvEventEntries, int windowSize = 0, std::optional attentionDpRank = std::nullopt, + AttentionDpGatherFn attentionDpGather = {}, std::string hashAlgo = "v2_sha256", + std::map windowSizeByLayerGroup = {}); + + void addCreatedEvent( + std::vector numBlocksPerCacheLevel, std::optional> layerGroupIds = std::nullopt); + void setLayerGroupWindowSizes(std::map windowSizes); + void addStoredEvent(KVCacheStoredData data, EventLayerGroupId layerGroupId = std::nullopt); + void addRemovedEvent(std::vector blockHashes, EventLayerGroupId layerGroupId = std::nullopt); + void addUpdatedEvent(EventBlockHash blockHash, std::optional cacheLevel = std::nullopt, + std::optional priority = std::nullopt, EventLayerGroupId layerGroupId = std::nullopt); + void addUpdatedEvent(Digest const& blockKey, std::optional cacheLevel = std::nullopt, + std::optional priority = std::nullopt, EventLayerGroupId layerGroupId = std::nullopt); + + void flushIterationEvents(); + std::vector getLatestEvents(std::optional timeoutMs = std::nullopt); + + std::string const& hashAlgorithm() const noexcept + { + return mHashAlgoName; + } + + static uint64_t hashV1BlockKey(std::vector const& tokens, uint64_t parentHash = 0, + std::optional loraTaskId = std::nullopt, + std::optional cacheSaltId = std::nullopt); + + void addStoredBlock(Block const& block) override; + void addStoredLifeCycle(Block const& block, LifeCycleId lifeCycle) override; + void addRemovedBlock(Digest const& blockKey) override; + void addRemovedLifeCycle(Digest const& blockKey, LifeCycleId lifeCycle) override; + void addCacheLevelUpdated( + Digest const& blockKey, CacheLevel oldLevel, CacheLevel newLevel, LifeCycleId lifeCycle) override; + +private: + enum class HashAlgorithm + { + kV1, + kV2Sha256, + kV2Sha256_64, + }; + + struct StoredBlockState + { + EventBlockHash blockHash; + std::set lifeCycleIds; + }; + + using V1RootAttrs = std::pair, std::optional>; + + static std::pair parseHashAlgorithm(std::string const& hashAlgo); + static std::string digestToHex(Digest const& digest); + static uint64_t truncateDigestToInt64(Digest const& digest); + static std::vector trimEvents(std::vector events, int maxKvEventEntries); + + EventBlockHash normalizeDigest(Digest const& digest) const; + EventBlockHash hashFromBlock(Block const& block); + uint64_t v1HashFromBlock(Block const& block); + uint64_t fallbackV1Hash(Digest const& blockKey); + std::optional parentHashFromBlock(Block const& block); + std::optional storedBlockFromBlock( + Block const& block, std::optional> const& lifeCycleIds = std::nullopt); + + void addStoredBlockUnlocked(Block const& block); + void addStoredEventUnlocked(KVCacheStoredData data, EventLayerGroupId layerGroupId); + void enqueueRemovedEventUnlocked(std::vector blockHashes, EventLayerGroupId layerGroupId); + void flushRemovedEventUnlocked(EventLayerGroupId layerGroupId); + void flushAllRemovedEventsUnlocked(); + KVCacheEvent& addEventUnlocked(KVCacheEventData data, EventLayerGroupId layerGroupId); + std::vector drainPendingEventsUnlocked(); + void publishEventsUnlocked(std::vector events, std::optional maxKvEventEntries = std::nullopt); + int getWindowSize(EventLayerGroupId layerGroupId) const; + void dropHashCache(Digest const& blockKey); + + int mMaxKvEventEntries; + int mWindowSize; + std::map mWindowSizeByLayerGroup; + std::optional mAttentionDpRank; + AttentionDpGatherFn mAttentionDpGather; + HashAlgorithm mHashAlgo; + std::string mHashAlgoName; + int64_t mNextEventId = 0; + + std::unordered_map mStoredBlocks; + std::map mLatestStoredEventIds; + std::map> mLatestRemovedBlockHashes; + std::vector mPendingEvents; + std::deque mEvents; + + std::unordered_map mV1HashByBlockKey; + std::unordered_set mV1HashCompatibleKeys; + std::unordered_map mV1RootAttrsByBlockKey; + bool mWarnedV1HashFallback = false; + + mutable std::mutex mMutex; + std::condition_variable mCondition; +}; + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventSink.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventSink.h new file mode 100644 index 000000000000..0be2a6d10fa1 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventSink.h @@ -0,0 +1,43 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "kv_cache_manager_v2/common.h" +#include "kv_cache_manager_v2/lifeCycleRegistry.h" + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +struct Block; + +// Boundary between the cache implementation and the native event manager. +class EventSink +{ +public: + virtual ~EventSink() = default; + + virtual void addStoredBlock(Block const& block) = 0; + virtual void addStoredLifeCycle(Block const& block, LifeCycleId lifeCycle) = 0; + virtual void addRemovedBlock(Digest const& blockKey) = 0; + virtual void addRemovedLifeCycle(Digest const& blockKey, LifeCycleId lifeCycle) = 0; + virtual void addCacheLevelUpdated( + Digest const& blockKey, CacheLevel oldLevel, CacheLevel newLevel, LifeCycleId lifeCycle) + = 0; +}; + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/evictionController.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/evictionController.cpp new file mode 100644 index 000000000000..f29a9bc046ff --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/evictionController.cpp @@ -0,0 +1,242 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "kv_cache_manager_v2/evictionController.h" +#include "kv_cache_manager_v2/common.h" +#include "kv_cache_manager_v2/exceptions.h" +#include "kv_cache_manager_v2/page.h" + +#include "tensorrt_llm/common/assert.h" +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// --------------------------------------------------------------------------- +// LRUEvictionPolicy +// --------------------------------------------------------------------------- + +NodeRef LRUEvictionPolicy::push(SharedPtr page, bool evictFirst) +{ + TLLM_CHECK_DEBUG_WITH_INFO(!page->nodeRef.has_value(), "page must not already be scheduled for eviction"); + if (evictFirst) + { + mQueue.push_front(std::move(page)); + return mQueue.begin(); + } + else + { + mQueue.push_back(std::move(page)); + auto it = mQueue.end(); + --it; + return it; + } +} + +SharedPtr LRUEvictionPolicy::pop() +{ + TLLM_CHECK_DEBUG(!mQueue.empty()); + auto page = mQueue.front(); + mQueue.pop_front(); + return page; +} + +SharedPtr LRUEvictionPolicy::remove(NodeRef node) +{ + auto page = *node; + TLLM_CHECK_DEBUG_WITH_INFO( + page->nodeRef.has_value() && page->nodeRef.value() == node, "node's page must reference this node"); + mQueue.erase(node); + return page; +} + +// --------------------------------------------------------------------------- +// PrioritizedEvictionPolicy +// --------------------------------------------------------------------------- + +LRUEvictionPolicy& PrioritizedEvictionPolicy::getOrCreate(Priority p) +{ + return mPolicies[p]; // std::map default-constructs if not present +} + +NodeRef PrioritizedEvictionPolicy::push(SharedPtr page, bool evictFirst) +{ + Priority p = page->priority; + LRUEvictionPolicy& policy = getOrCreate(p); + return policy.push(std::move(page), evictFirst); +} + +SharedPtr PrioritizedEvictionPolicy::pop() +{ + TLLM_CHECK_DEBUG(!mPolicies.empty()); + // Lowest priority key evicted first (std::map iterates in ascending key order) + auto it = mPolicies.begin(); + auto page = it->second.pop(); + if (it->second.empty()) + { + mPolicies.erase(it); + } + return page; +} + +SharedPtr PrioritizedEvictionPolicy::remove(NodeRef node) +{ + auto page = *node; + Priority p = page->priority; + auto it = mPolicies.find(p); + TLLM_CHECK_DEBUG(it != mPolicies.end()); + it->second.remove(node); + if (it->second.empty()) + { + mPolicies.erase(it); + } + return page; +} + +SlotCount PrioritizedEvictionPolicy::size() const noexcept +{ + SlotCount total = 0; + for (auto const& [p, policy] : mPolicies) + { + total += slotCountValueFromSize(policy.size()); + } + return total; +} + +// allPages() removed — use begin()/end() iterator instead. + +// --------------------------------------------------------------------------- +// PerLevelEvictionController +// --------------------------------------------------------------------------- + +PerLevelEvictionController::PerLevelEvictionController( + TypedVec const& lifeCycleGrouping, CacheLevel cacheLevel) + : mCacheLevel(cacheLevel) + , mLifeCycleGrouping(lifeCycleGrouping) +{ + // Compute number of pool groups = max(grouping) + 1 + PoolGroupIndex numPoolGroups{0}; + for (auto g : mLifeCycleGrouping) + { + if (g + 1 > numPoolGroups) + numPoolGroups = g + 1; + } + // Pool group indices must be contiguous 0..N-1. + TLLM_CHECK_DEBUG(numPoolGroups + == PoolGroupIndex{ + static_cast(std::set(mLifeCycleGrouping.begin(), mLifeCycleGrouping.end()).size())}); + mPolicies.resize(numPoolGroups); +} + +PerLevelEvictionController::~PerLevelEvictionController() +{ + TLLM_CHECK_DEBUG_WITH_INFO(std::all_of(mPolicies.begin(), mPolicies.end(), [](auto const& p) { return p.empty(); }), + "Eviction controller is not empty on destruction"); +} + +PrioritizedEvictionPolicy& PerLevelEvictionController::getPolicy(LifeCycleId lcId) +{ + PoolGroupIndex pgIdx = mLifeCycleGrouping.at(lcId); + return mPolicies.at(pgIdx); +} + +void PerLevelEvictionController::scheduleForEviction(Page& page, bool evictFirst) +{ + TLLM_CHECK_DEBUG(page.nodeRef == std::nullopt); + TLLM_CHECK_DEBUG(page.cacheLevel == mCacheLevel); + auto sharedPage = page.sharedFromThis(); + NodeRef ref = getPolicy(page.lifeCycle).push(sharedPage, evictFirst); + page.nodeRef = ref; + TLLM_CHECK_DEBUG_WITH_INFO(*ref == sharedPage, "stored iterator must dereference to this page"); +} + +TypedVec>> PerLevelEvictionController::evict( + TypedVec const& minNumPages) +{ + TLLM_CHECK_DEBUG(minNumPages.size() == numPoolGroups()); + + TypedVec>> ret(numPoolGroups()); + + try + { + for (PoolGroupIndex pgIdx{0}; pgIdx < minNumPages.size(); ++pgIdx) + { + auto& policy = mPolicies.at(pgIdx); + SlotCount const count = minNumPages.at(pgIdx); + if (count < 0) + { + throw LogicError("PerLevelEvictionController::evict: page count must be non-negative"); + } + SlotCount const available = policy.size() + slotCountValueFromSize(ret[pgIdx].size()); + if (available < count) + { + throw OutOfPagesError("Not enough pages to evict in group " + std::to_string(pgIdx.value())); + } + while (slotCountValueFromSize(ret[pgIdx].size()) < count) + { + auto page = policy.pop(); + page->nodeRef = std::nullopt; + ret[pgIdx].push_back(page); + // @TODO: evict dependencies (like Python _evict_dependencies) + } + } + } + catch (...) + { + // Re-queue evicted pages in reverse order (push to front so they are evicted first next time) + for (PoolGroupIndex pgIdx = ret.size(); pgIdx > PoolGroupIndex{0};) + { + --pgIdx; + auto& group = ret[pgIdx]; + while (!group.empty()) + { + auto page = std::move(group.back()); + group.pop_back(); + scheduleForEviction(*page, /*evictFirst=*/true); + } + } + throw; + } + + TLLM_CHECK_DEBUG_WITH_INFO(std::all_of(ret.begin(), ret.end(), + [this](auto const& group) { + return std::all_of(group.begin(), group.end(), + [this](auto const& p) { return p->cacheLevel == mCacheLevel; }); + }), + "Corrupted eviction controller"); + + return ret; +} + +void PerLevelEvictionController::remove(NodeRef node) +{ + auto page = *node; + TLLM_CHECK_DEBUG_WITH_INFO( + page->nodeRef.has_value() && page->nodeRef.value() == node, "page's nodeRef must match the node being removed"); + getPolicy(page->lifeCycle).remove(node); + page->nodeRef = std::nullopt; +} + +SlotCount PerLevelEvictionController::numEvictablePages(PoolGroupIndex pgIdx) const +{ + return mPolicies.at(pgIdx).size(); +} + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/evictionController.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/evictionController.h new file mode 100644 index 000000000000..b7ff84816d7f --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/evictionController.h @@ -0,0 +1,198 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "kv_cache_manager_v2/common.h" +#include "kv_cache_manager_v2/exceptions.h" +#include "kv_cache_manager_v2/storage/config.h" +#include "kv_cache_manager_v2/utils/sharedPtr.h" + +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// Forward declaration: Page is defined in page.h, but evictionController.h +// must not include page.h to avoid circular dependencies. +// We use a forward-declared abstract interface instead. +class Page; + +// --------------------------------------------------------------------------- +// NodeRef — stable iterator into an LRU list. +// Using list>: erasing any element does not invalidate +// other iterators (std::list guarantee). This is the C++ equivalent of the +// Python dllistnode. +// +// IMPORTANT: A NodeRef is always tied to the specific list that issued it. +// Never mix NodeRefs from different LRUEvictionPolicy instances. +// --------------------------------------------------------------------------- +using EvictionList = std::list>; +using NodeRef = EvictionList::iterator; + +// --------------------------------------------------------------------------- +// LRUEvictionPolicy — LRU queue backed by std::list. +// Mirrors _eviction_controller.py::LRUEvictionPolicy. +// push() → O(1), pop() → O(1), remove(NodeRef) → O(1). +// --------------------------------------------------------------------------- +class LRUEvictionPolicy +{ +public: + // Push a page into the eviction queue. + // evictFirst=true puts it at the front (will be evicted first). + // Returns a stable iterator (NodeRef) for later O(1) removal. + NodeRef push(SharedPtr page, bool evictFirst = false); + + // Remove and return the front (least-recently-used) page. + SharedPtr pop(); + + // Remove an arbitrary page via its iterator. O(1). + SharedPtr remove(NodeRef node); + + size_t size() const noexcept + { + return mQueue.size(); + } + + bool empty() const noexcept + { + return mQueue.empty(); + } + + EvictionList::const_iterator cbegin() const + { + return mQueue.cbegin(); + } + + EvictionList::const_iterator cend() const + { + return mQueue.cend(); + } + +private: + EvictionList mQueue; +}; + +// --------------------------------------------------------------------------- +// PrioritizedEvictionPolicy — wraps per-priority LRU sub-queues. +// Mirrors _eviction_controller.py::PrioritizedEvictionPolicy. +// Lower priority key = evicted first. +// --------------------------------------------------------------------------- +class PrioritizedEvictionPolicy +{ +public: + NodeRef push(SharedPtr page, bool evictFirst = false); + SharedPtr pop(); + SharedPtr remove(NodeRef node); + + SlotCount size() const noexcept; + + bool empty() const noexcept + { + return size() == 0; + } + + // Generator: returns a mutable lambda that yields shared_ptr const* + // in eviction order (lowest priority first, LRU within). Returns nullptr when exhausted. + // No extra strong refs — the pointer references the shared_ptr inside the list node. + [[nodiscard]] auto pageGenerator() const + { + using MapIt = std::map::const_iterator; + MapIt mapIt = mPolicies.cbegin(); + MapIt mapEnd = mPolicies.cend(); + EvictionList::const_iterator listIt; + // Advance to first non-empty sub-queue. + while (mapIt != mapEnd && mapIt->second.empty()) + ++mapIt; + if (mapIt != mapEnd) + listIt = mapIt->second.cbegin(); + return [=]() mutable -> SharedPtr const* + { + if (mapIt == mapEnd) + return nullptr; + auto* result = &(*listIt); + ++listIt; + if (listIt == mapIt->second.cend()) + { + ++mapIt; + while (mapIt != mapEnd && mapIt->second.empty()) + ++mapIt; + if (mapIt != mapEnd) + listIt = mapIt->second.cbegin(); + } + return result; + }; + } + +private: + LRUEvictionPolicy& getOrCreate(Priority p); + + // std::map keeps keys sorted. We evict from the lowest-priority key first. + std::map mPolicies; +}; + +// --------------------------------------------------------------------------- +// PerLevelEvictionController — one eviction controller per cache level. +// Holds one PrioritizedEvictionPolicy per pool group. +// Mirrors _eviction_controller.py::PerLevelEvictionController. +// --------------------------------------------------------------------------- +class PerLevelEvictionController +{ +public: + // lifeCycleGrouping: maps LifeCycleId → PoolGroupIndex. + // cacheLevel: the level this controller manages. + PerLevelEvictionController(TypedVec const& lifeCycleGrouping, CacheLevel cacheLevel); + + ~PerLevelEvictionController(); + + void scheduleForEviction(Page& page, bool evictFirst = false); + + // Evict at least minNumPages[pgIdx] pages per pool group. + // Returns evicted pages per pool group. + // On failure, re-queues any already-evicted pages and throws OutOfPagesError. + TypedVec>> evict( + TypedVec const& minNumPages); + + // Remove a page from the queue by its NodeRef. + void remove(NodeRef node); + + [[nodiscard]] SlotCount numEvictablePages(PoolGroupIndex pgIdx) const; + + [[nodiscard]] PoolGroupIndex numPoolGroups() const noexcept + { + return mPolicies.size(); + } + + // All pages in eviction order for a pool group. + [[nodiscard]] auto pageGenerator(PoolGroupIndex pgIdx) const + { + return mPolicies.at(pgIdx).pageGenerator(); + } + +private: + PrioritizedEvictionPolicy& getPolicy(LifeCycleId lcId); + + CacheLevel mCacheLevel; + TypedVec mLifeCycleGrouping; + TypedVec mPolicies; // one per pool group +}; + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/exceptions.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/exceptions.h new file mode 100644 index 000000000000..0b83db95a010 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/exceptions.h @@ -0,0 +1,184 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "kv_cache_manager_v2/utils/sharedPtr.h" + +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// --------------------------------------------------------------------------- +// Exception hierarchy (mirrors _exceptions.py) +// --------------------------------------------------------------------------- + +class OutOfMemoryError : public std::runtime_error +{ +public: + explicit OutOfMemoryError(std::string const& msg = "Out of memory") + : std::runtime_error(msg) + { + } +}; + +class HostOOMError : public OutOfMemoryError +{ +public: + explicit HostOOMError(std::string const& msg = "Host out of memory") + : OutOfMemoryError(msg) + { + } +}; + +class DiskOOMError : public OutOfMemoryError +{ +public: + explicit DiskOOMError(std::string const& msg = "Disk out of memory") + : OutOfMemoryError(msg) + { + } +}; + +class CuOOMError : public OutOfMemoryError +{ +public: + explicit CuOOMError(std::string const& msg = "CUDA out of memory") + : OutOfMemoryError(msg) + { + } +}; + +// Indicates a bug in the KV cache manager code. +class LogicError : public std::logic_error +{ +public: + explicit LogicError(std::string const& msg) + : std::logic_error(msg) + { + } +}; + +// Mirrors a Python `assert` failure: the binding layer translates this to a +// Python AssertionError so shared tests observe the same exception type as the +// pure-Python backend. +class AssertionError : public std::logic_error +{ +public: + explicit AssertionError(std::string const& msg) + : std::logic_error(msg) + { + } +}; + +// Wraps a CUDA driver API error (CUresult). +class CuError : public std::runtime_error +{ +public: + CUresult errorCode; + + explicit CuError(CUresult result) + : std::runtime_error(makeMessage(result)) + , errorCode(result) + { + } + +private: + static std::string makeMessage(CUresult result) + { + char const* errStr = nullptr; + cuGetErrorString(result, &errStr); + std::string msg = "CUDA driver error: "; + msg += errStr ? errStr : ""; + return msg; + } +}; + +// A resource (e.g., a page lock) is still in use. +class ResourceBusyError : public std::runtime_error +{ +public: + explicit ResourceBusyError(std::string const& msg = "Resource is busy") + : std::runtime_error(msg) + { + } +}; + +// Not enough free pages to satisfy an allocation request. +class OutOfPagesError : public std::runtime_error +{ +public: + explicit OutOfPagesError(std::string const& msg = "Out of pages") + : std::runtime_error(msg) + { + } +}; + +// Block creation rejected because its tokens are fully covered by an existing sibling. +// Mirrors Python's UselessBlockError — carries the sibling block. +// TODO: Once Python is removed and C++ becomes the primary development target, +// replace this exception-based flow with a simple if-condition return in +// addOrGetExistingBlock (returning the sibling block directly instead of throwing). +// The exception pattern exists only to maintain parity with the Python code path. +// Forward-declared; Block definition is in blockRadixTree.h. +struct Block; + +class UselessBlockError : public std::runtime_error +{ +public: + SharedPtr block; + + explicit UselessBlockError(SharedPtr blk) + : std::runtime_error("Block is useless — covered by existing sibling") + , block(std::move(blk)) + { + } +}; + +// --------------------------------------------------------------------------- +// Helper: unwrap a weak_ptr, throw LogicError on dangling reference. +// Mirrors Python's unwrap_rawref(_utils.py:163). +// --------------------------------------------------------------------------- +template +SharedPtr unwrap(WeakPtr const& ref) +{ + auto ptr = ref.lock(); + if (!ptr) + throw LogicError("Dereferencing a dangling weak_ptr"); + return ptr; +} + +// --------------------------------------------------------------------------- +// Helper: unwrap CUresult, throw CuError/CuOOMError on failure. +// --------------------------------------------------------------------------- +inline void cuCheck(CUresult result) +{ + if (result == CUDA_SUCCESS) + { + return; + } + if (result == CUDA_ERROR_OUT_OF_MEMORY) + { + throw CuOOMError(); + } + throw CuError(result); +} + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.cpp new file mode 100644 index 000000000000..835c48c9a13b --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.cpp @@ -0,0 +1,131 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "kv_cache_manager_v2/introspection.h" + +#include "kv_cache_manager_v2/blockRadixTree.h" + +#include "kv_cache_manager_v2/kvCache.h" +#include "kv_cache_manager_v2/kvCacheManager.h" +#include "kv_cache_manager_v2/storageManager.h" + +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +namespace +{ + +bool allBlockPagesDroppable(Block const& block) +{ + for (auto const* page : block.storage) + { + if (page != nullptr && page->status() != PageStatus::DROPPABLE) + { + return false; + } + } + + for (auto const& [_, child] : block.next) + { + if (!allBlockPagesDroppable(*child)) + { + return false; + } + } + return true; +} + +} // namespace + +KvCacheIntrospection::ActivePageStats KvCacheIntrospection::activePageStats(KvCache const& kvCache) +{ + auto& storageMgr = kvCache.manager().storage(); + CacheLevel const numTiers = storageMgr.numCacheLevels(); + TypedVec counts(numTiers, 0); + TypedVec unscheduledEvictable(numTiers, 0); + + for (auto const& activePage : kvCache._activePages()) + { + auto page = kvCache._page(activePage.ordinal, activePage.beamIdx, activePage.lcId); + if (!page) + { + continue; + } + + CacheLevel const level = page->cacheLevel; + counts.at(level) += 1; + if (storageMgr.isEvictable(*page) && !page->scheduledForEviction()) + { + unscheduledEvictable.at(level) += 1; + } + } + + return {std::move(counts), std::move(unscheduledEvictable)}; +} + +TypedVec KvCacheIntrospection::storageStatistics( + KvCacheManager& manager, CacheLevel level) +{ + TypedVec result; + PoolGroupIndex const numPoolGroups = manager.storage().numPoolGroups(); + result.reserve(numPoolGroups); + for (PoolGroupIndex pgIdx{0}; pgIdx < numPoolGroups; ++pgIdx) + { + result.push_back(manager.storage().getStatistics(level, pgIdx)); + } + return result; +} + +TypedVec KvCacheIntrospection::computeSlotsForBatch(KvCacheManager& manager, + BatchDesc const& batch, int tokensPerBlock, std::optional const& swaScratchReuse) +{ + return manager.storage().computeSlotsForBatch(batch, tokensPerBlock, swaScratchReuse); +} + +bool KvCacheIntrospection::allTreePagesDroppable(KvCacheManager& manager) +{ + for (auto const& [_, root] : manager.radixTree().roots()) + { + for (auto const& [__, block] : root->next) + { + if (!allBlockPagesDroppable(*block)) + { + return false; + } + } + } + return true; +} + +void KvCacheIntrospection::setNumSampledKvCaches(KvCacheManager& manager, int value) +{ + manager.mNumSampledKvCaches = value; +} + +void KvCacheIntrospection::setLastAdjustmentTime(KvCacheManager& manager, double value) +{ + manager.mLastAdjustmentTime = value; +} + +void KvCacheIntrospection::setTargetRatioListGpu(KvCacheManager& manager, TypedVec value) +{ + manager.mTargetRatioListGpu = std::move(value); +} + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.h new file mode 100644 index 000000000000..07f0dd9acef0 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.h @@ -0,0 +1,53 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "kv_cache_manager_v2/common.h" +#include "kv_cache_manager_v2/kvCache.h" +#include "kv_cache_manager_v2/storageManager.h" + +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +class KvCacheIntrospection +{ +public: + using ActivePageStats = std::tuple, TypedVec>; + + static ActivePageStats activePageStats(KvCache const& kvCache); + static bool allTreePagesDroppable(KvCacheManager& manager); + static TypedVec storageStatistics(KvCacheManager& manager, CacheLevel level); + + // White-box hook: minimum per-pool-group slot counts to support a BatchDesc. + // Reaches StorageManager::computeSlotsForBatch() (private) via friendship. + static TypedVec computeSlotsForBatch(KvCacheManager& manager, BatchDesc const& batch, + int tokensPerBlock, std::optional const& swaScratchReuse); + + // White-box test hooks: mutate auto-tuner state so accuracy tests can force a + // pool rebalance. Reach KvCacheManager's private members via friendship. + static void setNumSampledKvCaches(KvCacheManager& manager, int value); + static void setLastAdjustmentTime(KvCacheManager& manager, double value); + static void setTargetRatioListGpu(KvCacheManager& manager, TypedVec value); +}; + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp new file mode 100644 index 000000000000..1715f44c7463 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp @@ -0,0 +1,2567 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "kv_cache_manager_v2/kvCache.h" +#include "kv_cache_manager_v2/blockRadixTree.h" +#include "kv_cache_manager_v2/common.h" +#include "kv_cache_manager_v2/copyEngine.h" +#include "kv_cache_manager_v2/exceptions.h" +#include "kv_cache_manager_v2/kvCacheManager.h" +#include "kv_cache_manager_v2/storageManager.h" +#include "kv_cache_manager_v2/utils/math.h" + +#include "tensorrt_llm/common/assert.h" +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +namespace +{ + +// Copy one slot's data across all pools in a pool group. +// Used by resume() (GPU→GPU partial block copy) and _snapshotSsmToTreeBlock(). +void copySlotData(StorageManager& storageMgr, CacheLevel dstLevel, CacheLevel srcLevel, PoolGroupIndex pgIdx, + SlotId dstSlotId, SlotId srcSlotId, CUstream stream) +{ + auto slotSizes = storageMgr.slotSize(pgIdx); + CacheTier dstTier = storageMgr.cacheTier(dstLevel); + CacheTier srcTier = storageMgr.cacheTier(srcLevel); + for (PoolIndex poolIdx{0}; poolIdx < slotSizes.size(); ++poolIdx) + { + Address dst = storageMgr.slotAddress(dstLevel, pgIdx, dstSlotId, poolIdx); + Address src = storageMgr.slotAddress(srcLevel, pgIdx, srcSlotId, poolIdx); + batchedCopy(dstTier, srcTier, static_cast(slotSizes[poolIdx]), {{dst, src}}, stream); + } +} + +} // anonymous namespace + +// --------------------------------------------------------------------------- +// KvCache constructor +// --------------------------------------------------------------------------- + +KvCache::KvCache(KvCacheManager& manager, ReuseScope reuseScope, std::optional reuseMatch, + std::optional mId, PriorityCb priorityCb, std::optional expectedPromptLength) + : id(mId) + , mManager(manager.shared_from_this()) + , mReuseScope(std::move(reuseScope)) + , mPriorityCb(priorityCb ? std::move(priorityCb) : [](BlockOrdinal, LifeCycleId) { return kPriorityDefault; }) + , mStatus(Status::SUSPENDED) + , mCommitState(CommitState::ALLOWED) + , mBeamWidth(BeamIndex{1}) + , mCapacity(0) + , mHistoryLength(0) + , mExpectedPromptLength( + expectedPromptLength.has_value() ? std::optional{std::max(*expectedPromptLength, 0)} : std::nullopt) + , mNumCommittedBlocks(0) + , mTokensPerBlock(manager.tokensPerBlock()) +{ + LifeCycleId numLc = manager.storage().numLifeCycles(); + + // Initialise page index buffers: [beamIdx][lcId] = empty vector + mBasePageIndices.resize(mBeamWidth); + for (auto& beamIndices : mBasePageIndices) + { + beamIndices.resize(numLc, PageIndexBuf{std::vector{}}); + } + + // Always initialise mSsmBlocks (matching Python: no longer optional). + mSsmBlocks.resize(mBeamWidth); + for (auto& beam : mSsmBlocks) + { + beam.resize(numLc); // default-constructs to monostate + } + + mEnableSwaScratchReuse = manager.isSwaScratchReuseEnabled(); + mScratchSlots.resize(manager.storage().numLifeCycles()); + + if (reuseMatch.has_value()) + { + _setupForReuse(*reuseMatch); + } + + _refreshGenerationAllocReady(); + + mAvgHistoryLength.update(static_cast(mHistoryLength)); + + mManager->registerKvCache(this); + mManager->updateAvgReusedLength(static_cast(mHistoryLength)); + TLLM_CHECK_DEBUG(_checkSanity()); +} + +KvCache::~KvCache() +{ + try + { + close(); + } + catch (...) + { + // Destructors must not propagate exceptions (implicitly noexcept in C++11). + // close() should not throw in normal usage; if it does, suppress and accept leak. + } +} + +// --------------------------------------------------------------------------- +// State machine +// --------------------------------------------------------------------------- + +CUstream KvCache::cudaStream() const +{ + TLLM_CHECK_DEBUG_WITH_INFO(mCudaStream.has_value(), "No CUDA stream attached"); + return *mCudaStream; +} + +CachedCudaEvent KvCache::finishEvent() const +{ + return mFinishEvent.value(); +} + +Priority KvCache::getPriority(BlockOrdinal ordinal, LifeCycleId lc) const +{ + return mPriorityCb(ordinal, lc); +} + +StorageManager* KvCache::storageManager() const +{ + return &mManager->storage(); +} + +std::vector KvCache::_activePages() const +{ + std::vector result; + auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); + auto const& lcs = mManager->lifeCycles(); + LifeCycleId numLc = mManager->storage().numLifeCycles(); + + for (LifeCycleId lcId{0}; lcId < numLc; ++lcId) + { + // SSM lifecycle → yield from mSsmBlocks (check individual entries). + if (ssmLcId.has_value() && lcId == *ssmLcId) + { + for (BeamIndex bi{0}; bi < mBeamWidth; ++bi) + { + if (!blockPageIsNull(mSsmBlocks[bi][*ssmLcId])) + result.push_back({kBadBlockOrdinal, bi, lcId}); + } + continue; + } + + // Attention lifecycle: yield non-stale blocks (sink + window). + LifeCycle const& lc = lcs.getLifeCycle(lcId); + auto staleRange = _getStaleRange(mHistoryLength, lc); + BlockOrdinal staleBeg = staleRange.beg; + BlockOrdinal staleEnd = staleRange.end; + auto scratchRange = _getScratchRange(lc); + + // Sink blocks: [0, staleBeg) + for (BlockOrdinal ord{0}; ord < staleBeg; ++ord) + for (BeamIndex bi{0}; bi < mBeamWidth; ++bi) + result.push_back({ord, bi, lcId}); + + // Window blocks: [staleEnd, numBlocks) — skip scratch blocks. + for (BlockOrdinal ord{staleEnd}; ord < mBlocks.size(); ++ord) + { + auto& block = mBlocks[ord]; + for (BeamIndex bi{0}; bi < mBeamWidth; ++bi) + { + bool isScratch = scratchRange.contains(ord); + TLLM_CHECK_DEBUG(isScratch == blockPageIsNull(block.pages[bi][lcId])); + if (!isScratch) + result.push_back({ord, bi, lcId}); + } + } + } + return result; +} + +SharedPtr KvCache::_page(BlockOrdinal ordinal, BeamIndex beamIdx, LifeCycleId lcId) const +{ + bool const isSsm = ordinal == kBadBlockOrdinal; + auto const ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); + TLLM_CHECK_DEBUG((ssmLcId.has_value() && lcId == *ssmLcId) == isSsm); + auto const& blockPage = isSsm ? mSsmBlocks.at(beamIdx).at(lcId) : mBlocks.at(ordinal).pages.at(beamIdx).at(lcId); + return blockPageGetPage(blockPage); +} + +void KvCache::activate() +{ + TLLM_CHECK_DEBUG(mStatus == Status::SUSPENDED); + TLLM_CHECK_DEBUG_WITH_INFO(mCudaStream.has_value(), "cuda_stream must be set before activate()"); + + mFinishEvent.reset(); + + // Lock only active (non-stale) pages to GPU — mirrors Python's _active_pages(). + auto activePages = _activePages(); + std::vector targets; + targets.reserve(activePages.size()); + + for (auto const& ap : activePages) + { + BlockPage* bp = nullptr; + if (ap.ordinal == kBadBlockOrdinal) + { + bp = &mSsmBlocks[ap.beamIdx][ap.lcId]; + } + else + { + bp = &mBlocks[ap.ordinal].pages[ap.beamIdx][ap.lcId]; + } + auto& holder = std::get>(*bp); + TLLM_CHECK_DEBUG(holder); + targets.push_back({holder->page, ap.beamIdx, ap.ordinal, ap.lcId}); + } + + { + auto locks = batchedLockToGpu(*this, targets); + size_t idx = 0; + for (auto& t : targets) + { + TLLM_CHECK_DEBUG(t.page == locks[idx].page()); + BeamIndex bi = t.beamIndex; + LifeCycleId lc = t.lifeCycle; + if (t.ordinal == kBadBlockOrdinal) + mSsmBlocks[bi][lc] = std::move(locks[idx++]); + else + mBlocks[t.ordinal].pages[bi][lc] = std::move(locks[idx++]); + } + } +} + +bool KvCache::resume(std::optional stream) +{ + TLLM_CHECK_DEBUG(mStatus == Status::SUSPENDED); + + // Set stream first (mirrors Python: self.cuda_stream = cuda_stream). + if (stream.has_value()) + { + setCudaStream(*stream); + } + TLLM_CHECK_DEBUG_WITH_INFO(mCudaStream.has_value(), "cuda_stream is never set"); + TLLM_CHECK_DEBUG(!mFinishEvent.has_value()); + + // Check utilization against threshold. + auto const utilizations = mManager->storage().getUtilization(kGpuLevel); + float const utilization = utilizations.empty() ? 0.f : *std::max_element(utilizations.begin(), utilizations.end()); + if (utilization > mManager->config().maxUtilForResume) + { + return false; + } + + auto& storageMgr = mManager->storage(); + auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); + LifeCycleId numLc = storageMgr.numLifeCycles(); + + // Pre-allocate GPU slots for deferred copies (partial blocks + SSM) and scratch slots + // before locking, so we never end up in a state where pages are locked but we can't allocate. + TypedVec> deferredSlots(numLc); + + // Compute scratch slot deltas UNCONDITIONALLY (mirrors Python: _take_excess_scratch_slots + // is called outside _never_resumed). + auto [excessScratch, scratchDeltaCounts, scratchRanges] = _takeExcessScratchSlots(mCapacity, mHistoryLength); + TLLM_CHECK_DEBUG(excessScratch.size() == numLc + && std::all_of(excessScratch.begin(), excessScratch.end(), [](auto const& s) { return s.empty(); })); + + TypedVec numSlotsNeeded(numLc, 0); + bool hasPartial = false; + if (mNeverResumed) + { + TLLM_CHECK_DEBUG(mBeamWidth == BeamIndex{1}); + hasPartial = numCommittedTokens() % mTokensPerBlock != 0; + for (LifeCycleId lc{0}; lc < numLc; ++lc) + { + bool isSsm = ssmLcId.has_value() && lc == *ssmLcId; + if (isSsm || hasPartial) + numSlotsNeeded[lc] += 1; + } + } + + // Add scratch slot needs UNCONDITIONALLY (mirrors Python: delta loop is outside _never_resumed). + for (LifeCycleId lc{0}; lc < numLc; ++lc) + numSlotsNeeded[lc] += std::max(0, scratchDeltaCounts[lc]); + + // Only allocate if any slots are needed. + bool anyNeeded = std::any_of(numSlotsNeeded.begin(), numSlotsNeeded.end(), [](SlotCount n) { return n > 0; }); + if (anyNeeded) + { + TypedVec> tmpSlots; + try + { + MigrationRecorder const migrationRecorder + = [this](std::vector> const& pages, std::vector const& slots, CacheLevel srcLevel, + CacheLevel dstLevel) { _recordMigratedSlots(pages, slots, srcLevel, dstLevel); }; + DropRecorder const dropRecorder = [this](std::vector> const& pages, CacheLevel cacheLevel) + { _recordDroppedPages(pages, cacheLevel); }; + tmpSlots = storageMgr.newGpuSlots(numSlotsNeeded, migrationRecorder, dropRecorder); + } + catch (OutOfPagesError const&) + { + return false; + } + + // Separate deferred vs scratch slots, and collect scratch ready events. + std::vector scratchReadyEvents; + for (LifeCycleId lc{0}; lc < numLc; ++lc) + { + if (!tmpSlots[lc].empty()) + { + // Mirrors Python: `if self._never_resumed and (... SsmLifeCycle or has_partial):` + bool needsDeferred = mNeverResumed && ((ssmLcId.has_value() && lc == *ssmLcId) || hasPartial); + if (needsDeferred) + { + // Python uses pop() here: reserve one slot for deferred copy, then treat the rest as scratch. + deferredSlots[lc] = std::move(tmpSlots[lc].back()); + tmpSlots[lc].pop_back(); + } + // Remaining slots are scratch slots. + for (auto& slot : tmpSlots[lc]) + { + mScratchSlots[lc].emplace_back(std::move(slot), *this, lc, + /*skipWait=*/true); + scratchReadyEvents.push_back(&mScratchSlots[lc].back().slot().readyEvent); + } + } + } + + // Wait only for newly-added scratch slots (mirrors Python's + // stream_wait_events for scratch_slots_to_add). + if (!scratchReadyEvents.empty()) + streamWaitEvents(reinterpret_cast(cudaStream()), scratchReadyEvents); + } + + try + { + activate(); + } + catch (OutOfPagesError const&) + { + // Release pre-allocated deferred slots on failure. + for (LifeCycleId lc{0}; lc < numLc; ++lc) + { + if (deferredSlots[lc].has_value()) + storageMgr.releaseSlot(lc, kGpuLevel, std::move(*deferredSlots[lc])); + } + // Scratch slots stay in mScratchSlots — they'll be freed by close() inside + // a recordEventScope, matching Python behavior. + return false; + } + + // Deferred copy: for partial blocks and SSM, copy from now-locked source pages + // to pre-allocated GPU slots, then unlock sources and replace with new pages. + if (mNeverResumed) + { + BeamIndex beamIdx = kDefaultBeamIndex; + auto const lastOrdinal = BlockOrdinal{mBlocks.empty() ? 0 : (numCommittedTokens() - 1) / mTokensPerBlock}; + CUstream cudaStr = cudaStream(); + + // Wait for all new slots to be ready (deduplicated). + { + std::vector slotEvents; + for (auto& optSlot : deferredSlots) + { + if (optSlot.has_value()) + slotEvents.push_back(&optSlot->readyEvent); + } + streamWaitEvents(reinterpret_cast(cudaStr), slotEvents); + } + + // Phase 1: Copy GPU→GPU from locked source pages to pre-allocated slots. + std::vector srcLocks; + for (LifeCycleId lcIdx{0}; lcIdx < numLc; ++lcIdx) + { + if (!deferredSlots[lcIdx].has_value()) + continue; + auto& newSlot = *deferredSlots[lcIdx]; + + BlockPage* sourcePage = nullptr; + if (ssmLcId.has_value() && lcIdx == *ssmLcId) + { + if (numCommittedTokens() == 0) + continue; // fresh SSM — no source to copy from + sourcePage = &mSsmBlocks[beamIdx][lcIdx]; + } + else + { + sourcePage = &mBlocks[lastOrdinal].pages[beamIdx][lcIdx]; + } + auto* lock = std::get_if(sourcePage); + TLLM_CHECK_DEBUG(lock && lock->isValid()); + bool const hasPartialReuseSource = _hasReuseSource(*sourcePage); + srcLocks.push_back(lock); + + PoolGroupIndex pgIdx = storageMgr.getPoolGroupIndex(lcIdx); + copySlotData(storageMgr, kGpuLevel, kGpuLevel, pgIdx, newSlot.slotId(), lock->page()->slotId(), cudaStr); + if ((!ssmLcId.has_value() || lcIdx != *ssmLcId) && _shouldRecordStats()) + { + bool const changed = mPendingStats.recordAllocationRange(lcIdx, lastOrdinal, lastOrdinal + 1, + /*beamWidth=*/1, /*countAsMissed=*/!hasPartialReuseSource); + if (changed) + { + mManager->markStatsDirty(id); + } + } + KVCacheIterationStatsDelta iterationStats; + iterationStats.iterIntraDeviceCopyBlocks = 1; + for (size_t const size : storageMgr.slotSize(pgIdx)) + { + iterationStats.iterIntraDeviceCopyBytes += static_cast(size); + } + _recordDirectIterationStats(lcIdx, iterationStats); + } + + // Unlock source pages — recordEventScope captures all prior CUDA work + // so the original pages know when we're done reading from them. + if (!srcLocks.empty()) + { + auto scope = recordEventScope(); + for (auto* lock : srcLocks) + lock->unlock(); + } + + // Phase 2: Replace with new UncommittedPages (both copied and fresh SSM). + for (LifeCycleId lcIdx{0}; lcIdx < numLc; ++lcIdx) + { + if (!deferredSlots[lcIdx].has_value()) + continue; + auto& newSlot = *deferredSlots[lcIdx]; + + BlockPage* targetBp; + BlockOrdinal blockOrdinal; + if (ssmLcId.has_value() && lcIdx == *ssmLcId) + { + targetBp = &mSsmBlocks[beamIdx][lcIdx]; + blockOrdinal = kBadBlockOrdinal; + } + else + { + targetBp = &mBlocks[lastOrdinal].pages[beamIdx][lcIdx]; + blockOrdinal = lastOrdinal; + } + + auto newPage = makeShared(*this, blockOrdinal, lcIdx, kGpuLevel, beamIdx); + newPage->setSlot(newSlot); + auto newLock = newPage->lock(*this, beamIdx, blockOrdinal, lcIdx, /*skipWait=*/true); + *targetBp = std::move(newLock); + } + + // Clear treeBlock for partial last block (mirrors Python: partial block is uncommitted). + if (numCommittedTokens() % mTokensPerBlock != 0) + mBlocks[lastOrdinal].treeBlock = nullptr; + } + + mNeverResumed = false; + mStatus = Status::ACTIVE; + return true; +} + +bool KvCache::prefetch(CacheLevel target) +{ + TLLM_CHECK_DEBUG(mStatus == Status::SUSPENDED); + auto& storageMgr = mManager->storage(); + CacheLevel const numTiers = storageMgr.numCacheLevels(); + TLLM_CHECK_DEBUG(kGpuLevel <= target && target < numTiers); + + PoolGroupIndex const numPoolGroups = storageMgr.numPoolGroups(); + TypedVec>>> allPages( + numPoolGroups, TypedVec>>(numTiers)); + + for (auto const& activePage : _activePages()) + { + auto page = _page(activePage.ordinal, activePage.beamIdx, activePage.lcId); + if (!page) + { + continue; + } + CacheLevel const level = page->cacheLevel; + if (level < target) + { + continue; + } + auto const pgIdx = storageMgr.getPoolGroupIndex(activePage.lcId); + allPages.at(pgIdx).at(level).push_back(std::move(page)); + } + + try + { + storageMgr.prefetch(target, allPages); + } + catch (OutOfPagesError const&) + { + return false; + } + return true; +} + +void KvCache::suspend() +{ + TLLM_CHECK_DEBUG(mStatus == Status::ACTIVE); + TLLM_CHECK_DEBUG(_checkSanity()); + TLLM_CHECK_DEBUG(!mFinishEvent.has_value()); + + // Copy data from external buffers back to internal vectors (mirrors Python's suspend). + for (BeamIndex bi{0}; bi < mBeamWidth; ++bi) + for (LifeCycleId lcId{0}; lcId < mBasePageIndices[bi].size(); ++lcId) + if (std::holds_alternative>(mBasePageIndices[bi][lcId])) + setBasePageIndexBuf(bi, lcId, nullptr, 0); + + // Record event scope — mirrors Python's `with self._record_event()`. + // SharedPageLock destructors inside the scope use finishEvent() to synchronize. + { + auto scope = recordEventScope(); + auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); + + // Convert SharedPageLocks → PageHolders for active (non-stale) pages only. + // Mirrors Python: for ordinal, beam_idx, lc_idx in self._active_pages() + for (auto const& ap : _activePages()) + { + auto& bp = (ap.lcId != ssmLcId) ? mBlocks[ap.ordinal].pages[ap.beamIdx][ap.lcId] + : mSsmBlocks[ap.beamIdx][ap.lcId]; + // expect_type(_SharedPageLock, beam_block[lc_idx]) → std::get raises on wrong type + auto& lock = std::get(bp); + auto holder = lock.page()->hold(); + bp = std::move(holder); // ~SharedPageLock calls unlock() → notifyFinish(finishEvent()) + } + // Free scratch slots inside scope — unlock() needs finishEvent(). + // Mirrors Python: _free_scratch_slots() inside `with self._record_event()`. + _freeScratchSlots(); + } + mStatus = Status::SUSPENDED; +} + +void KvCache::close() +{ + TLLM_CHECK_DEBUG(_checkSanity()); + if (mStatus == Status::CLOSED) + return; + + discardPendingStats(); + stopCommitting(); + TLLM_CHECK_DEBUG(_checkSanity()); + + if (mCapacity > 0) + { + mAvgCapacity.update(static_cast(mCapacity)); + mManager->updateAvgSqrCapacity(mAvgCapacity.value() * mAvgCapacity.value()); + mManager->updateAvgSqrHistoryLength(mAvgHistoryLength.value() * mAvgHistoryLength.value()); + mManager->incrementNumSampledKvCaches(); + mManager->tryUpdateTargetRatios(); + } + + // Record event scope — mirrors Python's `with self._record_event()`. + // Python always enters _record_event() here; _cuda_stream is valid for both ACTIVE and SUSPENDED. + { + auto scope = recordEventScope(); + _clearBlocks(); + } + mStatus = Status::CLOSED; + mManager->unregisterKvCache(this); +} + +KVCacheStatsDelta KvCache::commitPendingStats() +{ + if (!_shouldRecordStats()) + { + discardPendingStats(); + return {}; + } + + mManager->commitStats(mPendingStats.globalStats(), mPendingStats.iterationStatsByLifeCycle()); + mManager->commitSsmSnapshotIterationStats(mPendingStats.ssmSnapshotIterationStatsByLifeCycle()); + KVCacheStatsDelta const requestStats = mPendingStats.requestStats().copy(); + mPendingStats.clear(); + mManager->clearStatsDirty(id); + return requestStats; +} + +void KvCache::discardPendingStats() +{ + mPendingStats.clear(); + mManager->clearStatsDirty(id); +} + +bool KvCache::_shouldRecordStats() const +{ + return mManager->config().enableStats && !mManager->isStatsExcluded(id); +} + +void KvCache::_refreshStatsDirtyState() +{ + if (!mPendingStats.empty()) + { + mManager->markStatsDirty(id); + } + else + { + mManager->clearStatsDirty(id); + } +} + +void KvCache::_recordDirectIterationStats(LifeCycleId lifeCycle, KVCacheIterationStatsDelta const& iterationStats) +{ + if (!_shouldRecordStats() || iterationStats.empty() + || !std::holds_alternative(mManager->lifeCycles().getLifeCycle(lifeCycle))) + { + return; + } + IterationStatsByLifeCycle iterationStatsByLifeCycle; + iterationStatsByLifeCycle.emplace(lifeCycle, iterationStats.copy()); + mManager->commitStats({}, iterationStatsByLifeCycle); +} + +void KvCache::_recordMigratedSlots( + std::vector> const& pages, std::vector const& slots, CacheLevel srcLevel, CacheLevel dstLevel) +{ + if (!_shouldRecordStats()) + { + return; + } + TLLM_CHECK_DEBUG(pages.size() == slots.size()); + for (auto const& page : pages) + { + LifeCycleId const lifeCycle = page->lifeCycle; + if (!std::holds_alternative(mManager->lifeCycles().getLifeCycle(lifeCycle))) + { + continue; + } + + PoolGroupIndex const poolGroup = mManager->storage().getPoolGroupIndex(lifeCycle); + int64_t pageSize = 0; + for (size_t const size : mManager->storage().slotSize(poolGroup)) + { + pageSize += static_cast(size); + } + + KVCacheStatsDelta stats; + KVCacheIterationStatsDelta iterationStats; + if (srcLevel == kGpuLevel && dstLevel > kGpuLevel) + { + iterationStats.iterOffloadBlocks = 1; + iterationStats.iterOffloadBytes = pageSize; + } + else if (dstLevel == kGpuLevel) + { + stats.allocTotalBlocks = 1; + stats.allocNewBlocks = 1; + iterationStats.iterAllocTotalBlocks = 1; + iterationStats.iterAllocNewBlocks = 1; + if (srcLevel > kGpuLevel) + { + iterationStats.iterOnboardBlocks = 1; + iterationStats.iterOnboardBytes = pageSize; + } + else if (srcLevel == kGpuLevel) + { + iterationStats.iterIntraDeviceCopyBlocks = 1; + iterationStats.iterIntraDeviceCopyBytes = pageSize; + } + } + + if (!stats.empty() || !iterationStats.empty()) + { + IterationStatsByLifeCycle iterationStatsByLifeCycle; + iterationStatsByLifeCycle.emplace(lifeCycle, iterationStats); + mManager->commitStats(stats, iterationStatsByLifeCycle); + } + } +} + +void KvCache::_recordDroppedPages(std::vector> const& pages, CacheLevel cacheLevel) +{ + (void) cacheLevel; + if (!_shouldRecordStats()) + { + return; + } + for (auto const& page : pages) + { + LifeCycleId const lifeCycle = page->lifeCycle; + if (!std::holds_alternative(mManager->lifeCycles().getLifeCycle(lifeCycle))) + { + continue; + } + PoolGroupIndex const poolGroup = mManager->storage().getPoolGroupIndex(lifeCycle); + int64_t pageSize = 0; + for (size_t const size : mManager->storage().slotSize(poolGroup)) + { + pageSize += static_cast(size); + } + KVCacheIterationStatsDelta iterationStats; + iterationStats.iterHostDroppedBlocks = 1; + iterationStats.iterHostDroppedBytes = pageSize; + _recordDirectIterationStats(lifeCycle, iterationStats); + } +} + +void KvCache::_recordResizePendingAllocations(BlockOrdinal blockBegin, BlockOrdinal blockEnd, + TypedVec> const& excludedRanges, bool countAsGeneration) +{ + if (!_shouldRecordStats() || blockBegin >= blockEnd) + { + return; + } + + bool changed = false; + for (auto const& [lifeCycle, unused] : mManager->lifeCycles().attentionLifeCycles()) + { + (void) unused; + auto const& excluded = excludedRanges[lifeCycle]; + BlockOrdinal const firstEnd = std::min(blockEnd, excluded.beg); + if (blockBegin < firstEnd) + { + changed |= mPendingStats.recordAllocationRange( + lifeCycle, blockBegin, firstEnd, mBeamWidth.value(), !countAsGeneration, countAsGeneration); + } + BlockOrdinal const secondBegin = std::max(blockBegin, excluded.end); + if (secondBegin < blockEnd) + { + changed |= mPendingStats.recordAllocationRange( + lifeCycle, secondBegin, blockEnd, mBeamWidth.value(), !countAsGeneration, countAsGeneration); + } + } + if (changed) + { + mManager->markStatsDirty(id); + } +} + +void KvCache::_subtractPendingAllocationRange(BlockOrdinal blockBegin, BlockOrdinal blockEnd) +{ + if (mPendingStats.subtractAllocationRange(blockBegin, blockEnd)) + { + _refreshStatsDirtyState(); + } +} + +bool KvCache::_hasReuseSource(BlockPage const& page) +{ + auto const committedPage = dynamicPointerCast(blockPageGetPage(page)); + return committedPage && committedPage->block != nullptr; +} + +// --------------------------------------------------------------------------- +// _clearBlocks +// --------------------------------------------------------------------------- + +void KvCache::_clearBlocks() +{ + // Drop last block first (mirrors Python: while self._blocks: self._blocks.pop()). + while (!mBlocks.empty()) + mBlocks.pop_back(); + _freeScratchSlots(); + // Clear SSM blocks. + auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); + if (ssmLcId.has_value()) + { + for (auto& beamBlock : mSsmBlocks) + beamBlock[*ssmLcId] = std::monostate{}; + } +} + +// --------------------------------------------------------------------------- +// _copyPageToTreeBlock: copy a page into a new committed (or SSM committed) page +// attached to a radix tree block. Mirrors Python's _copy_page_to_tree_block. +// --------------------------------------------------------------------------- + +void KvCache::_copyPageToTreeBlock(SharedPtr const& treeBlock, LifeCycleId lcIdx, SharedPtr const& srcPage, + std::optional ssmNumTokensInBlock) +{ + if (treeBlock->storage.at(lcIdx) != nullptr) + { + return; // block already holds a page for this lifecycle + } + auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); + bool const isSsm = ssmLcId.has_value() && lcIdx == *ssmLcId; + TLLM_CHECK_DEBUG(isSsm == ssmNumTokensInBlock.has_value()); + + auto& storageMgr = mManager->storage(); + PoolGroupIndex pgIdx = storageMgr.getPoolGroupIndex(lcIdx); + + for (CacheLevel lvl = srcPage->cacheLevel; lvl < storageMgr.numCacheLevels(); ++lvl) + { + Slot newSlot; + try + { + auto slots = storageMgr.newSlotsForPoolGroup(lvl, pgIdx, 1); + newSlot = std::move(slots[0]); + } + catch (OutOfPagesError const&) + { + continue; + } + + CUstream stream = cudaStream(); + newSlot.readyEvent.waitInStream(reinterpret_cast(stream)); + copySlotData(storageMgr, lvl, srcPage->cacheLevel, pgIdx, newSlot.slotId(), srcPage->slotId(), stream); + + CachedCudaEvent readyEv(reinterpret_cast(stream)); + auto tempPage = makeShared(*this, treeBlock->ordinal(), lcIdx, lvl, kDefaultBeamIndex); + tempPage->setSlot(newSlot); + SharedPtr committed; + if (ssmNumTokensInBlock.has_value()) + { + committed = tempPage->convertToSsmCommitted(treeBlock, std::move(readyEv), *ssmNumTokensInBlock); + } + else + { + committed = tempPage->convertToCommitted(treeBlock, std::move(readyEv)); + } + + // Schedule for eviction so eviction controller keeps a strong reference, + // preventing the page from being destroyed. + storageMgr.scheduleForEviction(*committed); + return; // success + } + // No pages available in any level, silently skip snapshot (matches Python). +} + +// --------------------------------------------------------------------------- +// _snapshotSsmToTreeBlock: snapshot live SSM state to a radix tree block for the +// given committed token count. Mirrors Python's _snapshot_ssm_to_tree_block. +// --------------------------------------------------------------------------- + +void KvCache::_snapshotSsmToTreeBlock(SharedPtr const& treeBlock, LifeCycleId ssmLcId, int numTokens, bool move) +{ + int const numTokensInBlock = numTokens - treeBlock->ordinal().value() * mTokensPerBlock; + TLLM_CHECK_DEBUG(0 < numTokensInBlock && numTokensInBlock <= mTokensPerBlock); + + CommittedPage* existingRaw = treeBlock->storage.at(ssmLcId); + int existingNumTokens = 0; + if (existingRaw != nullptr) + { + auto* existingSsm = dynamic_cast(existingRaw); + TLLM_CHECK_DEBUG(existingSsm != nullptr); + existingNumTokens = existingSsm->numTokensInBlock; + } + if (existingNumTokens >= numTokensInBlock) + { + return; + } + if (existingRaw != nullptr) + { + // Detach the smaller snapshot; the CommittedPage dtor's expected-page guard + // makes the later unlink a no-op (the slot is already null). + treeBlock->storage.at(ssmLcId) = nullptr; + if (existingRaw->scheduledForEviction()) + { + existingRaw->manager->excludeFromEviction(*existingRaw); + } + } + + auto& ssmBlock = mSsmBlocks[kDefaultBeamIndex]; + auto* ssmLock = std::get_if(&ssmBlock[ssmLcId]); + TLLM_CHECK_DEBUG(ssmLock && ssmLock->isValid()); + auto srcPage = ssmLock->page(); + if (move) + { + auto unlocked = ssmLock->unlock(); + auto up = dynamicPointerCast(unlocked); + TLLM_CHECK_DEBUG(up != nullptr); + ssmBlock[ssmLcId] = std::monostate{}; + auto committed = up->convertToSsmCommitted(treeBlock, finishEvent(), numTokensInBlock); + mManager->storage().scheduleForEviction(*committed); + return; + } + + _copyPageToTreeBlock(treeBlock, ssmLcId, srcPage, numTokensInBlock); +} + +// --------------------------------------------------------------------------- +// _snapshotPartialBlockToTree: snapshot a partial final block into the radix +// tree. Mirrors Python's _snapshot_partial_block_to_tree. +// --------------------------------------------------------------------------- + +void KvCache::_snapshotPartialBlockToTree(BlockOrdinal ordinal, bool commitSsm) +{ + int const start = ordinal.value() * mTokensPerBlock; + int const end = std::min(start + mTokensPerBlock, static_cast(mCommittedTokens.size())); + std::vector tokens(mCommittedTokens.begin() + start, mCommittedTokens.begin() + end); + int const numTokens = static_cast(tokens.size()); + TLLM_CHECK_DEBUG(0 < numTokens && numTokens < mTokensPerBlock); + + LifeCycleId numLc = mManager->storage().numLifeCycles(); + NodeBase* prevNode = nullptr; + RootBlock& root = mManager->radixTree().addOrGetExisting(mReuseScope); + if (ordinal == BlockOrdinal{0}) + { + prevNode = &root; + } + else + { + prevNode = _getTreeBlock(BlockOrdinal{ordinal.value() - 1}).get(); + } + + bool isNew = false; + SharedPtr treeBlock; + try + { + treeBlock = addOrGetExistingBlock(prevNode, numLc, tokens, &isNew); + } + catch (UselessBlockError const& e) + { + treeBlock = e.block; + isNew = false; + } + TLLM_CHECK_DEBUG(treeBlock); + TLLM_CHECK_DEBUG(isNew || std::equal(tokens.begin(), tokens.end(), treeBlock->tokens.begin())); + + auto& beamBlock = mBlocks.at(ordinal).pages[kDefaultBeamIndex]; + auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); + // Only attach partial attention pages to a tree block whose token span is + // exactly the partial snapshot. A longer existing sibling may already have + // full attention pages; otherwise a partial page would make it look more + // reusable than it is. + if (static_cast(treeBlock->tokens.size()) == numTokens) + { + for (auto const& [lcIdx, attn] : mManager->lifeCycles().attentionLifeCycles()) + { + (void) attn; + auto& bp = beamBlock[lcIdx]; + if (blockPageIsNull(bp) || treeBlock->storage.at(lcIdx) != nullptr) + { + continue; + } + _copyPageToTreeBlock(treeBlock, lcIdx, blockPageGetPage(bp)); + } + } + if (commitSsm) + { + TLLM_CHECK_DEBUG(ssmLcId.has_value()); + _snapshotSsmToTreeBlock(treeBlock, *ssmLcId, start + numTokens); + } + if (isNew && treeBlock->eventSink) + { + treeBlock->eventSink->addStoredBlock(*treeBlock); + } +} + +// --------------------------------------------------------------------------- +// resize +// --------------------------------------------------------------------------- + +bool KvCache::resize(std::optional capacity, std::optional historyLength) +{ + TLLM_CHECK_DEBUG(mStatus == Status::ACTIVE); + TLLM_CHECK_DEBUG(mBlocks.size() == BlockOrdinal{divUp(mCapacity, mTokensPerBlock)}); + + int newCap = capacity.value_or(mCapacity); + int newHist = historyLength.value_or(mHistoryLength); + + if (capacity.has_value()) + mAvgCapacity.update(static_cast(newCap)); + if (historyLength.has_value()) + mAvgHistoryLength.update(static_cast(newHist)); + + if (newHist < mHistoryLength) + throw std::invalid_argument("History length cannot be decreased"); + if (newCap < newHist) + throw std::invalid_argument("History length cannot exceed capacity"); + + // Scratch reuse: enforce constraint. + bool enableScratch = mEnableSwaScratchReuse; + if (TLLM_UNLIKELY(gDebug) && enableScratch && newCap != mCapacity) + { + int const maxRewindLen = _swaScratchMaxRewindLen(); + int const minHistoryLength = std::max(0, mCapacity - maxRewindLen); + bool const validSwaScratchHistory = minHistoryLength <= newHist && newHist <= mCapacity; + TLLM_CHECK_WITH_INFO(validSwaScratchHistory, + "SWA scratch requires old_capacity - max_rewind_len <= history_length <= old_capacity"); + (void) validSwaScratchHistory; + } + + bool const recordGenerationAllocStats = mGenerationAllocReady && newCap > mCapacity; + if (!enableScratch && _shortcutSetCapacity(newCap) && _shortcutSetHistoryLength(newHist)) + { + _refreshGenerationAllocReady(); + return true; + } + + BlockOrdinal oldNumBlocks{divUp(mCapacity, mTokensPerBlock)}; + BlockOrdinal newNumBlocks{divUp(newCap, mTokensPerBlock)}; + LifeCycleId numLc = mManager->storage().numLifeCycles(); + auto const& lcs = mManager->lifeCycles(); + + _checkPageIndexBufferCapacity(newNumBlocks); + + auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); + auto backupHolders = _unlockStaleBlocks(newHist); + + if (newNumBlocks < oldNumBlocks) + { + TLLM_CHECK_DEBUG_WITH_INFO(!hasScratchSlots(), "Cannot shrink while scratch slots exist"); + _subtractPendingAllocationRange(newNumBlocks, oldNumBlocks); + auto scope = recordEventScope(); + _decreaseCapacity(newNumBlocks); + } + + // Compute scratch deltas. + auto [excessScratchSlots, deltaScratchSlots, scratchRanges] = _takeExcessScratchSlots(newCap, newHist); + + if (newNumBlocks >= oldNumBlocks) + { + // Compute new normal slots needed per lifecycle. + TypedVec numNewSlots(numLc, 0); + TypedVec> staleRanges(numLc); + for (LifeCycleId lc{0}; lc < numLc; ++lc) + { + if (ssmLcId.has_value() && lc == *ssmLcId) + continue; + staleRanges[lc] = _getStaleRange(newHist, lcs.getLifeCycle(lc)); + + if (enableScratch) + { + HalfOpenRange const newBlockRange{oldNumBlocks, newNumBlocks}; + int const numNewBlocksUsingScratch = intersect(scratchRanges[lc], newBlockRange).length(); + int const numNewNormalBlocks = newBlockRange.length() - numNewBlocksUsingScratch; + numNewSlots[lc] = static_cast(numNewNormalBlocks) * mBeamWidth.value(); + } + else + { + auto [staleBeg, staleEnd] = staleRanges[lc]; + int numNewBlocksToAdd; + if (oldNumBlocks < staleBeg) + { + TLLM_CHECK_DEBUG(newNumBlocks >= staleEnd); + numNewBlocksToAdd = (staleBeg - oldNumBlocks) + (newNumBlocks - staleEnd); + } + else + { + numNewBlocksToAdd = newNumBlocks - std::max(staleEnd, oldNumBlocks); + } + numNewSlots[lc] = static_cast(numNewBlocksToAdd) * mBeamWidth.value(); + } + } + + // Compute net allocation counts (normal + scratch delta). + TypedVec netAllocCounts(numLc, 0); + for (LifeCycleId lc{0}; lc < numLc; ++lc) + netAllocCounts[lc] = numNewSlots[lc] + static_cast(deltaScratchSlots[lc]); + + // Allocate new slots. + TypedVec> newSlots; + bool anyPositive = std::any_of(netAllocCounts.begin(), netAllocCounts.end(), [](SlotCount c) { return c > 0; }); + if (anyPositive) + { + TypedVec allocCounts(numLc); + for (LifeCycleId lc{0}; lc < numLc; ++lc) + allocCounts[lc] = std::max(SlotCount{0}, netAllocCounts[lc]); + try + { + MigrationRecorder const migrationRecorder + = [this](std::vector> const& pages, std::vector const& slots, + CacheLevel srcLevel, CacheLevel dstLevel) + { _recordMigratedSlots(pages, slots, srcLevel, dstLevel); }; + DropRecorder const dropRecorder + = [this](std::vector> const& pages, CacheLevel cacheLevel) + { _recordDroppedPages(pages, cacheLevel); }; + newSlots = mManager->storage().newGpuSlots(allocCounts, migrationRecorder, dropRecorder); + } + catch (OutOfPagesError const&) + { + _recoverExcessScratchSlots(excessScratchSlots); + _lockHeldBlocks(backupHolders); + return false; + } + } + else + { + newSlots.resize(numLc); + } + + // Wait on newly allocated slots. + { + std::vector readyEvents; + for (auto const& lcSlots : newSlots) + for (auto const& slot : lcSlots) + readyEvents.push_back(&slot.readyEvent); + if (!readyEvents.empty()) + streamWaitEvents(reinterpret_cast(cudaStream()), readyEvents); + } + + // Combine: new slots + excess scratch detached slots. + TypedVec> slots(numLc); + for (LifeCycleId lc{0}; lc < numLc; ++lc) + { + auto& combined = slots[lc]; + combined = std::move(newSlots[lc]); + for (auto& lock : excessScratchSlots[lc]) + combined.push_back(lock.detachSlot()); + excessScratchSlots[lc].clear(); + } + + // Release excess if net is negative. + if (std::any_of(netAllocCounts.begin(), netAllocCounts.end(), [](SlotCount c) { return c < 0; })) + { + auto scope = recordEventScope(); + for (LifeCycleId lc{0}; lc < numLc; ++lc) + { + for (SlotCount i = 0; i < -netAllocCounts[lc]; ++i) + { + auto slot = std::move(slots[lc].back()); + slots[lc].pop_back(); + slot.readyEvent = finishEvent(); + mManager->storage().releaseSlot(lc, kGpuLevel, std::move(slot)); + } + } + } + + // Assert correct combined slot count. + if (TLLM_UNLIKELY(gDebug)) + { + for (LifeCycleId lc{0}; lc < numLc; ++lc) + { + SlotCount const expected + = numNewSlots[lc] + std::max(SlotCount{0}, static_cast(deltaScratchSlots[lc])); + TLLM_CHECK(static_cast(slots[lc].size()) == expected); + } + } + + // Fulfill additional scratch slots (pop from end of combined slots). + for (LifeCycleId lc{0}; lc < numLc; ++lc) + { + for (int i = 0; i < deltaScratchSlots[lc]; ++i) + { + auto slot = std::move(slots[lc].back()); + slots[lc].pop_back(); + mScratchSlots[lc].emplace_back(std::move(slot), *this, lc, /*skipWait=*/true); + } + } + + // Wait for all slots that will back new pages. This includes detached excess scratch slots, + // which are not covered by the earlier wait on newly allocated slots. + { + std::vector readyEvents; + for (auto const& lcSlots : slots) + for (auto const& slot : lcSlots) + readyEvents.push_back(&slot.readyEvent); + if (!readyEvents.empty()) + streamWaitEvents(reinterpret_cast(cudaStream()), readyEvents); + } + + // Scratch and stale blocks do not consume per-request KV pages, so exclude them from allocation stats. + auto const& excludedRanges = enableScratch ? scratchRanges : staleRanges; + _recordResizePendingAllocations(oldNumBlocks, newNumBlocks, excludedRanges, recordGenerationAllocStats); + + // Resize page index buffers. + TLLM_CHECK_DEBUG(std::all_of(mBasePageIndices.begin(), mBasePageIndices.end(), + [oldNumBlocks](auto const& beamIndices) + { + return std::all_of(beamIndices.begin(), beamIndices.end(), + [oldNumBlocks](auto const& buf) + { + auto const* vec = std::get_if>(&buf); + return !vec || vec->size() == toSizeT(oldNumBlocks); + }); + })); + + // Create SeqBlocks for new ordinals (pop slots from end, matching Python). + _resizePageIndexBuffers(newNumBlocks); + for (BlockOrdinal ord = oldNumBlocks; ord < newNumBlocks; ++ord) + { + SeqBlock sb; + sb.pages.resize(mBeamWidth); + for (auto& row : sb.pages) + row.resize(numLc); + for (BeamIndex bi{0}; bi < mBeamWidth; ++bi) + { + for (LifeCycleId lc{0}; lc < numLc; ++lc) + { + if (ssmLcId.has_value() && lc == *ssmLcId) + continue; // SSM pages live in mSsmBlocks, not in mBlocks + if (enableScratch) + { + if (scratchRanges[lc].contains(ord)) + continue; // Scratch block — no per-block page allocation. + } + else + { + auto [staleBeg, staleEnd] = staleRanges[lc]; + if (staleBeg <= ord && ord < staleEnd) + continue; + } + auto slot = std::move(slots[lc].back()); + slots[lc].pop_back(); + auto page = makeShared(*this, ord, lc, kGpuLevel, bi); + page->setSlot(slot); + sb.pages[bi][lc] = page->lock(*this, bi, ord, lc, /*skipWait=*/true); + } + } + mBlocks.push_back(std::move(sb)); + } + TLLM_CHECK_DEBUG(std::all_of(slots.begin(), slots.end(), [](auto const& vec) { return vec.empty(); })); + } + + mCapacity = newCap; + mHistoryLength = newHist; + _refreshGenerationAllocReady(); + _evictOutOfWindowBlocks(newHist); + TLLM_CHECK_DEBUG(_checkSanity()); + return true; +} + +void KvCache::setCapacity(int cap) +{ + if (mEnableSwaScratchReuse) + { + throw std::invalid_argument( + "Cannot use capacity setter when SWA scratch reuse is enabled. " + "Use resize(capacity, history_length) instead."); + } + if (!resize(cap, std::nullopt)) + throw OutOfPagesError("Not enough pages in GPU memory"); +} + +void KvCache::setHistoryLength(int hist) +{ + bool success = resize(std::nullopt, hist); + TLLM_CHECK_DEBUG(success); + (void) success; +} + +bool KvCache::_shortcutSetCapacity(int newCap) +{ + if (newCap == mCapacity) + return true; + // No shortcut if block count changes. + if (divUp(newCap, mTokensPerBlock) != divUp(mCapacity, mTokensPerBlock)) + return false; + mCapacity = newCap; + return true; +} + +bool KvCache::_shortcutSetHistoryLength(int newHist) +{ + if (newHist == mHistoryLength) + return true; + // Check if stale range changes for any lifecycle. + for (auto [lcId, lc] : mManager->lifeCycles()) + { + bool changed = std::visit( + [&](auto const& v) -> bool + { + using T = std::decay_t; + if constexpr (std::is_same_v) + { + // history_length change does not impact blocks at all. + return false; + } + else + { + static_assert(std::is_same_v); + if (!v.windowSize.has_value()) + return false; + return v.getStaleRange(newHist, mTokensPerBlock) + != v.getStaleRange(mHistoryLength, mTokensPerBlock); + } + }, + lc); + if (changed) + return false; + } + mHistoryLength = newHist; + return true; +} + +void KvCache::_refreshGenerationAllocReady() +{ + if (mExpectedPromptLength.has_value() && mHistoryLength >= *mExpectedPromptLength) + { + mGenerationAllocReady = true; + } +} + +// --------------------------------------------------------------------------- +// Capacity management +// --------------------------------------------------------------------------- + +void KvCache::_increaseCapacity(BlockOrdinal newNumBlocks, int newHistoryLength) +{ + BlockOrdinal curNumBlocks = mBlocks.size(); + LifeCycleId numLc = mManager->storage().numLifeCycles(); + auto const& lcs = mManager->lifeCycles(); + auto ssmLcId = lcs.ssmLifeCycleId(); + + // Compute stale ranges using new history length so stale SWA blocks get no pages. + TypedVec> staleRanges(numLc); + TypedVec numSlotsPerLc(numLc, 0); + for (LifeCycleId lc{0}; lc < numLc; ++lc) + { + // SSM slots are handled separately below — don't allocate block-level slots for SSM. + if (ssmLcId.has_value() && lc == *ssmLcId) + continue; + staleRanges[lc] = _getStaleRange(newHistoryLength, lcs.getLifeCycle(lc)); + auto [staleBeg, staleEnd] = staleRanges[lc]; + int numNewBlocks; + if (curNumBlocks < staleBeg) + { + TLLM_CHECK_DEBUG(newNumBlocks >= staleEnd); + numNewBlocks = (staleBeg - curNumBlocks) + (newNumBlocks - staleEnd); + } + else + { + numNewBlocks = newNumBlocks - std::max(staleEnd, curNumBlocks); + } + numSlotsPerLc[lc] = static_cast(numNewBlocks) * mBeamWidth.value(); + } + + // SSM slots are now allocated lazily in resume() via deferred copy, not here. + + MigrationRecorder const migrationRecorder + = [this](std::vector> const& pages, std::vector const& slots, CacheLevel srcLevel, + CacheLevel dstLevel) { _recordMigratedSlots(pages, slots, srcLevel, dstLevel); }; + DropRecorder const dropRecorder = [this](std::vector> const& pages, CacheLevel cacheLevel) + { _recordDroppedPages(pages, cacheLevel); }; + auto allSlots = mManager->storage().newGpuSlots(numSlotsPerLc, migrationRecorder, dropRecorder); + + // Assert that internal index buffer sizes match expected old_num_blocks (mirrors Python line ~463). + TLLM_CHECK_DEBUG(std::all_of(mBasePageIndices.begin(), mBasePageIndices.end(), + [curNumBlocks](auto const& beamIndices) + { + return std::all_of(beamIndices.begin(), beamIndices.end(), + [curNumBlocks](auto const& buf) + { + auto const* vec = std::get_if>(&buf); + return !vec || vec->size() == toSizeT(curNumBlocks); + }); + })); + + // Create SeqBlocks for the new ordinals. + TypedVec slotCounters(numLc, 0); + _resizePageIndexBuffers(newNumBlocks); + for (BlockOrdinal ord = curNumBlocks; ord < newNumBlocks; ++ord) + { + SeqBlock sb; + sb.pages.resize(mBeamWidth); + for (auto& row : sb.pages) + row.resize(numLc); // default-constructs to monostate + for (LifeCycleId lc{0}; lc < numLc; ++lc) + { + // SSM pages live in mSsmBlocks, not in _blocks. + if (ssmLcId.has_value() && lc == *ssmLcId) + continue; + + auto [staleBeg, staleEnd] = staleRanges[lc]; + if (staleBeg <= ord && ord < staleEnd) + continue; // stale block for this lc — no page allocated + + size_t si = slotCounters[lc]++; + auto& slot = allSlots[lc][si]; + auto page = makeShared(*this, ord, lc, kGpuLevel, kDefaultBeamIndex); + page->setSlot(slot); + sb.pages[kDefaultBeamIndex][lc] = page->lock(*this, kDefaultBeamIndex, ord, lc); + } + mBlocks.push_back(std::move(sb)); + } + // Assert all allocated slots were consumed (mirrors Python line ~488). + if (TLLM_UNLIKELY(gDebug)) + { + for (LifeCycleId lc{0}; lc < numLc; ++lc) + TLLM_CHECK(slotCounters[lc] == allSlots[lc].size()); + } +} + +void KvCache::_decreaseCapacity(BlockOrdinal newNumBlocks) +{ + while (mBlocks.size() > newNumBlocks) + { + auto& sb = mBlocks.back(); + for (auto& beamPages : sb.pages) + for (auto& bp : beamPages) + bp = std::monostate{}; + sb.treeBlock.reset(); + mBlocks.pop_back(); + } + _resizePageIndexBuffers(newNumBlocks); +} + +HalfOpenRange KvCache::_getStaleRange(int historyLength, LifeCycle const& lc) const +{ + return kv_cache_manager_v2::getStaleRange(lc, historyLength, mTokensPerBlock); +} + +std::vector KvCache::_unlockStaleBlocks(int newHistoryLength) +{ + std::vector ret; + if (newHistoryLength == mHistoryLength) + return ret; + + auto scope = recordEventScope(); + auto const& lcs = mManager->lifeCycles(); + LifeCycleId numLc = mManager->storage().numLifeCycles(); + + for (LifeCycleId lcIdx{0}; lcIdx < numLc; ++lcIdx) + { + LifeCycle const& lc = lcs.getLifeCycle(lcIdx); + // SSM pages live in mSsmBlocks, not _blocks — skip. + if (std::holds_alternative(lc)) + continue; + // Full-attention (no SWA) has empty stale range — skip. + auto const& alc = std::get(lc); + if (!alc.windowSize.has_value()) + continue; + + auto oldRange = _getStaleRange(mHistoryLength, lc); + auto newRange = _getStaleRange(newHistoryLength, lc); + + BlockOrdinal unlockStart = std::max(oldRange.end, newRange.beg); + BlockOrdinal unlockEnd = std::min(mBlocks.size(), newRange.end); + + for (BlockOrdinal ord = unlockStart; ord < unlockEnd; ++ord) + { + auto& sb = mBlocks[ord]; + bool isCommitted = sb.isCommitted(); + bool holdForCommit + = !mManager->commitMinSnapshot() && !isCommitted && (mCommitState == CommitState::ALLOWED); + + for (BeamIndex bi{0}; bi < sb.pages.size(); ++bi) + { + auto& bp = sb.pages[bi][lcIdx]; + if (blockPageIsNull(bp)) + { + // No page to unlock: scratch block, commit_min_snapshot early + // release, or a stale block created unallocated by resize(). + continue; + } + TLLM_CHECK_DEBUG(std::holds_alternative(bp)); + auto holder = blockPageGetPage(bp)->hold(); + ret.push_back({ord, bi, lcIdx, holder}); + bp = holdForCommit ? BlockPage{std::move(holder)} : BlockPage{std::monostate{}}; + } + } + } + return ret; +} + +void KvCache::_lockHeldBlocks(std::vector const& backup) +{ + std::vector targets; + targets.reserve(backup.size()); + for (auto const& b : backup) + targets.push_back({b.holder->page, b.beamIdx, b.ordinal, b.lcId}); + + auto locks = batchedLockToGpu(*this, targets); + for (size_t i = 0; i < locks.size(); ++i) + { + auto const& t = backup[i]; + mBlocks[t.ordinal].pages[t.beamIdx][t.lcId] = std::move(locks[i]); + } +} + +// --------------------------------------------------------------------------- +// _takeUncommittedPage — extract uncommitted pages from a SeqBlock. +// Mirrors Python's _take_uncommitted_page(). +// --------------------------------------------------------------------------- + +TypedVec KvCache::_takeUncommittedPage( + SeqBlock& sb, BeamIndex beamIdx, std::optional skipLc) +{ + LifeCycleId numLc = mManager->storage().numLifeCycles(); + TypedVec result(numLc, TakenPage{nullptr, false}); + for (LifeCycleId lc{0}; lc < numLc; ++lc) + { + if (skipLc.has_value() && lc == *skipLc) + continue; + auto& bp = sb.pages[beamIdx][lc]; + if (auto* lock = std::get_if(&bp)) + { + auto up = dynamicPointerCast(lock->page()); + TLLM_CHECK_DEBUG_WITH_INFO(up, "page must be UncommittedPage"); + result[lc] = {up, true}; + } + else if (auto* holder = std::get_if>(&bp)) + { + TLLM_CHECK_DEBUG(*holder); + auto up = dynamicPointerCast((*holder)->page); + TLLM_CHECK_DEBUG_WITH_INFO(up, "page must be UncommittedPage"); + result[lc] = {up, false}; + } + bp = std::monostate{}; + } + return result; +} + +// --------------------------------------------------------------------------- +// _commitBlock — shared logic for committing a single block. +// Mirrors Python's _commit_block(ordinal, is_last). +// Caller must have recordEventScope() open so finishEvent() works. +// On VIRTUAL_STOP or when isLast is true, transitions to USER_STOP and +// calls _onStopCommitting() — callers do not need post-call cleanup. +// --------------------------------------------------------------------------- + +void KvCache::_commitBlock(int ord, bool isLast, bool commitSsm, bool moveSsm) +{ + TLLM_CHECK_DEBUG(mCommitState == CommitState::ALLOWED); + TLLM_CHECK_DEBUG(ord == mNumCommittedBlocks); + + auto& sb = mBlocks.at(BlockOrdinal{ord}); + TLLM_CHECK_DEBUG_WITH_INFO(sb.pages.size() == BeamIndex{1}, "Must have 1 beam only"); + + // Build token block — always slice up to tokens_per_block; is_full tells us + // whether we got a full block's worth. Mirrors Python's: + // tokens = self._committed_tokens[start : start + tokens_per_block] + // is_full = len(tokens) == tokens_per_block + int start = ord * mTokensPerBlock; + int end = std::min(start + mTokensPerBlock, static_cast(mCommittedTokens.size())); + std::vector tokenBlock(mCommittedTokens.begin() + start, mCommittedTokens.begin() + end); + bool isFull = static_cast(tokenBlock.size()) == mTokensPerBlock; + + if (!isLast && !isFull) + throw LogicError("Cannot commit block that is not full except last block"); + + // Prev node lookup (root or previous committed block). + RootBlock& root = mManager->radixTree().addOrGetExisting(mReuseScope); + LifeCycleId numLc = mManager->storage().numLifeCycles(); + + NodeBase* prevNode = &root; + if (ord > 0) + { + TLLM_CHECK_DEBUG_WITH_INFO(mBlocks[BlockOrdinal{ord - 1}].treeBlock, "prev block must be committed"); + prevNode = mBlocks[BlockOrdinal{ord - 1}].treeBlock.get(); + } + + // Try to find or create a block in the radix tree. + // Mirrors Python's try/except UselessBlockError pattern. + // TODO: Replace with if-condition once Python is removed and C++ is the primary codebase. + bool blockIsNew = false; + SharedPtr newBlock; + try + { + newBlock = addOrGetExistingBlock(prevNode, numLc, tokenBlock, &blockIsNew); + } + catch (UselessBlockError const& e) + { + newBlock = e.block; + blockIsNew = false; + } + TLLM_CHECK_DEBUG(newBlock); + TLLM_CHECK_DEBUG(newBlock->tokensPerBlock() == mTokensPerBlock); + // In reuse case, verify token match (mirrors Python: tree_block.tokens[:num_tokens] == tokens). + TLLM_CHECK_DEBUG(blockIsNew || std::equal(tokenBlock.begin(), tokenBlock.end(), newBlock->tokens.begin())); + + auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); + bool didCommit = false; + + if (blockIsNew) + { + // New block: take uncommitted pages, convert to committed. + // Mirrors Python's _take_uncommitted_page + convert path. + auto taken = _takeUncommittedPage(sb, kDefaultBeamIndex, ssmLcId); + sb.treeBlock = newBlock; + for (LifeCycleId lc{0}; lc < numLc; ++lc) + { + auto& [up, locked] = taken[lc]; + if (!up) + continue; + auto committed = up->convertToCommitted(newBlock, finishEvent()); + if (locked) + sb.pages[kDefaultBeamIndex][lc] + = committed->lock(*this, kDefaultBeamIndex, static_cast(ord), lc); + else + sb.pages[kDefaultBeamIndex][lc] = committed->hold(); + } + TLLM_CHECK_DEBUG(_getTreeBlock(static_cast(ord)) == newBlock); + ++mNumCommittedBlocks; + if (newBlock->eventSink) + { + newBlock->eventSink->addStoredBlock(*newBlock); + } + didCommit = true; + } + else if (newBlock->isFull() && mManager->allowSeqRebasing() && isFull) + { + // Existing block: rebase — reuse existing block's committed pages. + // Mirrors Python's `elif tree_block.is_full and allow_seq_rebasing and is_full` path. + std::vector reuseTasks; + for (LifeCycleId lc{0}; lc < numLc; ++lc) + { + if (ssmLcId.has_value() && lc == *ssmLcId) + continue; + auto& bp = sb.pages[kDefaultBeamIndex][lc]; + if (blockPageIsNull(bp)) + continue; + auto* existingPage = newBlock->storage.at(lc); + bool isLocked = std::holds_alternative(bp); + if (existingPage == nullptr) + { + // Existing page gone — put our uncommitted page into the tree block. + if (auto* lock = std::get_if(&bp)) + { + auto up = dynamicPointerCast(lock->page()); + if (up) + { + bp = std::monostate{}; + auto committed = up->convertToCommitted(newBlock, finishEvent()); + if (newBlock->eventSink) + { + newBlock->eventSink->addStoredLifeCycle(*newBlock, lc); + } + bp = isLocked + ? BlockPage{committed->lock(*this, kDefaultBeamIndex, static_cast(ord), lc)} + : BlockPage{committed->hold()}; + } + } + else if (auto* holder = std::get_if>(&bp)) + { + if (*holder) + { + auto up = dynamicPointerCast((*holder)->page); + if (up) + { + bp = std::monostate{}; + auto committed = up->convertToCommitted(newBlock, finishEvent()); + if (newBlock->eventSink) + { + newBlock->eventSink->addStoredLifeCycle(*newBlock, lc); + } + bp = committed->hold(); + } + } + } + } + else + { + // Downgrade lock to holder for our page; reuse the existing page. + if (isLocked) + { + auto holder = blockPageGetPage(bp)->hold(); + bp = std::move(holder); + } + reuseTasks.push_back( + {existingPage->sharedFromThis(), kDefaultBeamIndex, static_cast(ord), lc}); + } + } + if (!reuseTasks.empty()) + { + auto locks = batchedLockToGpu(*this, reuseTasks); + for (size_t ri = 0; ri < reuseTasks.size(); ++ri) + { + LifeCycleId lc = reuseTasks[ri].lifeCycle; + sb.pages[kDefaultBeamIndex][lc] = std::move(locks[ri]); + } + } + // Don't clear SSM storage on rebase — the existing block may have a valid snapshot. + sb.treeBlock = newBlock; + TLLM_CHECK_DEBUG(_getTreeBlock(static_cast(ord)) == newBlock); + ++mNumCommittedBlocks; + didCommit = true; + } + else + { + // Can't commit and can't reuse existing block. Just stop committing. + mCommitState = CommitState::VIRTUAL_STOP; + } + + if (didCommit && commitSsm) + { + TLLM_CHECK_DEBUG(ssmLcId.has_value()); + _snapshotSsmToTreeBlock(newBlock, *ssmLcId, start + static_cast(tokenBlock.size()), moveSsm); + } + + if (sb.isCommitted()) + { + auto const& lifeCycles = mManager->lifeCycles(); + for (LifeCycleId lcIdx{0}; lcIdx < numLc; ++lcIdx) + { + if (ssmLcId.has_value() && lcIdx == *ssmLcId) + { + continue; + } + LifeCycle const& lc = lifeCycles.getLifeCycle(lcIdx); + if (!std::holds_alternative(lc)) + { + continue; + } + auto const staleRange = _getStaleRange(mHistoryLength, lc); + if (staleRange.contains(BlockOrdinal{ord})) + { + for (auto& beamBlock : sb.pages) + { + beamBlock[lcIdx] = std::monostate{}; + } + } + } + } + + // Mirrors Python's tail of _commit_block: + // if is_last or self._commit_state == self.CommitState.VIRTUAL_STOP: + // self._commit_state = self.CommitState.USER_STOP + // self._on_stop_committing() + if (isLast || mCommitState == CommitState::VIRTUAL_STOP) + { + mCommitState = CommitState::USER_STOP; + _onStopCommitting(); + } +} + +// --------------------------------------------------------------------------- +// commit +// --------------------------------------------------------------------------- + +void KvCache::commit(std::vector const& tokens, bool isEnd) +{ + TLLM_CHECK_DEBUG(mStatus == Status::ACTIVE); + if (mBeamWidth != BeamIndex{1}) + throw LogicError("Not implemented yet for beam search"); + if (tokens.empty()) + { + if (isEnd) + stopCommitting(); + return; + } + if (mCommitState == CommitState::USER_STOP) + throw LogicError("Cannot commit tokens after stop_committing()"); + + bool const commitMinSnapshot = mManager->commitMinSnapshot(); + auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); + if (commitMinSnapshot) + { + int const newNumCommitted = static_cast(mCommittedTokens.size()) + static_cast(tokens.size()); + if (mHistoryLength != static_cast(mCommittedTokens.size()) && mHistoryLength != newNumCommitted) + { + throw AssertionError("commit_min_snapshot requires commit() to start or end at history_length"); + } + } + + // Append tokens to committed list. + mCommittedTokens.insert(mCommittedTokens.end(), tokens.begin(), tokens.end()); + if (mCommitState == CommitState::VIRTUAL_STOP) + { + if (isEnd) + mCommitState = CommitState::USER_STOP; + return; + } + + // Bump history_length to cover newly committed tokens (mirrors Python — done + // BEFORE the commit loop so stale-range computation sees the new history). + int const numCommitted = static_cast(mCommittedTokens.size()); + if (mHistoryLength < numCommitted) + setHistoryLength(numCommitted); + + int const numCommittedBlocksBefore = mNumCommittedBlocks; + int const newNumFullBlocks = numCommitted / mTokensPerBlock; + bool const hasPartialSnapshot = commitMinSnapshot && (numCommitted % mTokensPerBlock != 0) && numCommitted > 0; + bool const hasNewFullBlocks = newNumFullBlocks > numCommittedBlocksBefore; + if (hasNewFullBlocks || hasPartialSnapshot) + { + // Block whose end is the last committed token — where the SSM snapshot lives. + int const ssmSnapshotOrdinal = (numCommitted - 1) / mTokensPerBlock; + // Wrapped in recordEventScope() so SharedPageLock::unlock() shares one finish + // event (mirrors Python's `with self._record_event()`). + auto scope = recordEventScope(); + for (int ordinal = numCommittedBlocksBefore; ordinal < newNumFullBlocks; ++ordinal) + { + bool const commitSsm = commitMinSnapshot && ssmLcId.has_value() && ordinal == ssmSnapshotOrdinal; + _commitBlock(ordinal, /*isLast=*/false, commitSsm, /*moveSsm=*/isEnd); + // _commitBlock transitions to USER_STOP on VIRTUAL_STOP internally. + if (mCommitState != CommitState::ALLOWED) + break; + } + if (hasPartialSnapshot && mCommitState == CommitState::ALLOWED) + { + BlockOrdinal const partialOrdinal{newNumFullBlocks}; + if (isEnd) + { + _commitBlock(newNumFullBlocks, /*isLast=*/true, /*commitSsm=*/ssmLcId.has_value(), + /*moveSsm=*/ssmLcId.has_value()); + } + else + { + _snapshotPartialBlockToTree(partialOrdinal, /*commitSsm=*/ssmLcId.has_value()); + } + } + } + + if (isEnd && mCommitState != CommitState::USER_STOP) + stopCommitting(); +} + +void KvCache::stopCommitting() +{ + TLLM_CHECK_DEBUG(mStatus != Status::CLOSED); + if (mCommitState == CommitState::USER_STOP) + return; + TLLM_CHECK_DEBUG(_checkSanity()); + + // Mirrors Python's stop_committing() which calls _commit_block(ordinal, True). + if (mCommitState == CommitState::VIRTUAL_STOP) + { + mCommitState = CommitState::USER_STOP; + return; + } + + TLLM_CHECK_DEBUG(mCommitState == CommitState::ALLOWED); + + int tokensLeft = static_cast(mCommittedTokens.size()) - mNumCommittedBlocks * mTokensPerBlock; + if (tokensLeft > 0) + { + TLLM_CHECK_DEBUG(BlockOrdinal{mNumCommittedBlocks} < mBlocks.size()); + auto scope = recordEventScope(); + // isLast=true: _commitBlock handles USER_STOP + _onStopCommitting() internally. + _commitBlock(mNumCommittedBlocks, /*isLast=*/true); + } + else + { + mCommitState = CommitState::USER_STOP; + _onStopCommitting(); + } + TLLM_CHECK_DEBUG(mCommitState == CommitState::USER_STOP); +} + +// --------------------------------------------------------------------------- +// PlannedDropHandle — mirrors Python's PlannedDropHandle. +// --------------------------------------------------------------------------- + +PlannedDropHandle::PlannedDropHandle(std::vector const& pages) +{ + // Deduplicate by identity (mirrors Python's {id(page): page} dict). + std::vector unique; + std::unordered_set seen; + unique.reserve(pages.size()); + for (auto* page : pages) + { + if (seen.insert(page).second) + unique.push_back(page); + } + + std::vector> refs; + refs.reserve(unique.size()); + for (auto* page : unique) + { + refs.emplace_back(dynamicPointerCast(page->sharedFromThis())); + page->plannedDropCount += 1; + } + mPageRefs = std::move(refs); +} + +void PlannedDropHandle::drop() +{ + if (!mPageRefs.has_value()) + throw std::invalid_argument("Planned drop handle has already been dropped"); + + std::vector> pages; + for (auto const& ref : *mPageRefs) + { + auto page = ref.lock(); + if (page) + { + if (page->plannedDropCount <= 0) + throw std::invalid_argument("Committed page has no planned drop"); + pages.push_back(std::move(page)); + } + } + + mPageRefs.reset(); + for (auto const& page : pages) + { + page->plannedDropCount -= 1; + if (page->plannedDropCount == 0 && page->status() == PageStatus::DROPPABLE && page->scheduledForEviction()) + { + page->manager->excludeFromEviction(*page); + } + } +} + +PlannedDropHandle::~PlannedDropHandle() +{ + if (mPageRefs.has_value()) + { + // Mirror Python's __del__: apply the plan if not already dropped. + // Destructors must not throw; swallow any error. + try + { + drop(); + } + catch (...) + { + } + } +} + +std::unique_ptr KvCache::planCommittedBlockDrop() +{ + if (mCommitState != CommitState::USER_STOP) + throw LogicError("plan_committed_block_drop() requires stop_committing()"); + + // A cache with no committed tokens has no preceding turn to drop. + if (numCommittedTokens() == 0) + return nullptr; + + auto const match = mManager->matchReuse(mReuseScope, mCommittedTokens); + if (match.numTokens != numCommittedTokens() || match.blocks.empty()) + return nullptr; + + BlockOrdinal const end = match.blocks.size(); + std::vector pagesToDrop; + for (auto const item : mManager->lifeCycles()) + { + LifeCycleId const lcIdx = item.id; + LifeCycle const& lc = item.lc; + BlockOrdinal windowStart; + if (auto const* attn = std::get_if(&lc)) + { + // Full-attention blocks may still be needed by later turns. + if (!attn->windowSize.has_value()) + continue; + auto const staleRange = _getStaleRange(numCommittedTokens(), lc); + windowStart = std::min(staleRange.end, end); + } + else + { + // SSM: only the last committed block carries the reusable snapshot. + windowStart = BlockOrdinal{end.value() - 1}; + } + for (BlockOrdinal ordinal = windowStart; ordinal < end; ++ordinal) + { + CommittedPage* page = match.blocks[ordinal]->getPage(lcIdx); + if (page == nullptr) + return nullptr; + pagesToDrop.push_back(page); + } + } + return std::make_unique(pagesToDrop); +} + +// --------------------------------------------------------------------------- +// _onStopCommitting: release stale held uncommitted pages for SWA layers. +// Mirrors Python's _on_stop_committing(). +// --------------------------------------------------------------------------- + +void KvCache::_onStopCommitting() +{ + auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); + auto const& lcs = mManager->lifeCycles(); + + for (auto [lcIdx, lc] : lcs) + { + if (ssmLcId.has_value() && lcIdx == *ssmLcId) + continue; // SSM pages live in _ssm_blocks, not in _blocks + + auto staleRange = _getStaleRange(mHistoryLength, lc); + BlockOrdinal start = std::max(staleRange.beg, BlockOrdinal{mNumCommittedBlocks}); + BlockOrdinal end = staleRange.end; + + TLLM_CHECK_DEBUG(end <= mBlocks.size()); + for (BlockOrdinal ord = start; ord < end; ++ord) + { + auto& sb = mBlocks[ord]; + TLLM_CHECK_DEBUG(!sb.isCommitted()); + for (auto& beamPages : sb.pages) + { + auto& bp = beamPages[lcIdx]; + if (blockPageIsNull(bp)) + { + // Nothing to release: scratch block, commit_min_snapshot early + // release, or a stale block created unallocated by resize(). + continue; + } + TLLM_CHECK_DEBUG(std::holds_alternative>(bp)); + bp = std::monostate{}; + } + } + } + TLLM_CHECK_DEBUG(_checkSanity()); +} + +// --------------------------------------------------------------------------- +// _setupForReuse: find existing blocks in radix tree matching input tokens. +// --------------------------------------------------------------------------- + +void KvCache::_setupForReuse(BlockRadixTree::ReuseMatch const& match) +{ + auto const& matched = match.blocks; + int const numTokens = match.numTokens; + + auto& lifeCycles = mManager->lifeCycles(); + auto const& allLc = lifeCycles.getAll(); + LifeCycleId numLc = lifeCycles.size(); + auto ssmLcId = lifeCycles.ssmLifeCycleId(); + BlockOrdinal const fullReusedEnd{numTokens / mTokensPerBlock}; + bool const hasPartialMatch = numTokens % mTokensPerBlock != 0; + bool const shouldRecordStats = _shouldRecordStats(); + + // --- Build blocks with stale range handling --- + + BlockOrdinal const numMatchedBlocks = matched.size(); + _resizePageIndexBuffers(numMatchedBlocks); + + for (BlockOrdinal i{0}; i < numMatchedBlocks; ++i) + { + SeqBlock sb; + sb.treeBlock = matched[i]->sharedFromThis(); + sb.pages.resize(BeamIndex{1}); + sb.pages[kDefaultBeamIndex].resize(numLc); + mBlocks.push_back(std::move(sb)); + } + + BeamIndex beamIdx = kDefaultBeamIndex; + + for (LifeCycleId lcId{0}; lcId < numLc; ++lcId) + { + // SSM is handled separately below. + if (ssmLcId.has_value() && lcId == *ssmLcId) + continue; + + auto staleRange = getStaleRange(allLc[lcId], numTokens, mTokensPerBlock); + BlockOrdinal staleStart = staleRange.beg; + BlockOrdinal staleEnd = staleRange.end; + bool const isAttention = std::holds_alternative(allLc[lcId]); + int fullReusedBlocks = 0; + int partialReusedBlocks = 0; + + // Process a non-stale ordinal: hold the page. + // For partial blocks (last block, not full), defer the copy to first resume(). + auto processOrdinal = [&](BlockOrdinal ordinal) + { + auto& blk = *matched.at(ordinal); + auto* page = blk.storage.at(lcId); + TLLM_CHECK_DEBUG_WITH_INFO(page, "Expected page in non-stale block"); + auto& bpSlot = mBlocks[ordinal].pages[beamIdx][lcId]; + bpSlot = page->hold(); + if (shouldRecordStats && isAttention) + { + if (ordinal < fullReusedEnd) + { + ++fullReusedBlocks; + } + else if (hasPartialMatch && ordinal == fullReusedEnd && _hasReuseSource(bpSlot)) + { + partialReusedBlocks = 1; + } + } + }; + + for (BlockOrdinal ord{0}; ord < staleStart; ++ord) + processOrdinal(ord); + for (BlockOrdinal ord = staleEnd; ord < numMatchedBlocks; ++ord) + processOrdinal(ord); + + if (shouldRecordStats && isAttention && mPendingStats.recordReuse(lcId, fullReusedBlocks, partialReusedBlocks)) + { + mManager->markStatsDirty(id); + } + } + + // SSM reuse: hold the snapshot from the last matched block. Copy is deferred to first resume(). + if (ssmLcId.has_value() && !matched.empty()) + { + auto& snapshotBlock = *matched.back(); + auto* snapshotPage = snapshotBlock.storage[*ssmLcId]; + TLLM_CHECK_DEBUG_WITH_INFO(snapshotPage, "Last matched block must have SSM snapshot after truncation"); + mSsmBlocks[kDefaultBeamIndex][*ssmLcId] = snapshotPage->hold(); + } + // Record one SSM snapshot lookup for this reuse-match onboarding (a miss when + // nothing matched). Mirrors Python's record_ssm_snapshot_lookup call. + if (shouldRecordStats && ssmLcId.has_value()) + { + if (mPendingStats.recordSsmSnapshotLookup(*ssmLcId, match.numLookupTokens, numTokens, mTokensPerBlock)) + { + mManager->markStatsDirty(id); + } + } + + // Append matched tokens (reconstructed from the matched blocks). + mCommittedTokens = _getMatchedTokens(match); + mNumCommittedBlocks = numTokens / mTokensPerBlock; + mHistoryLength = numTokens; + mCapacity = numTokens; +} + +// --------------------------------------------------------------------------- +// _getMatchedTokens — reconstruct the committed token sequence from a match. +// Mirrors Python's _get_matched_tokens(). +// --------------------------------------------------------------------------- + +std::vector KvCache::_getMatchedTokens(BlockRadixTree::ReuseMatch const& match) const +{ + std::vector ret; + ret.reserve(static_cast(match.numTokens)); + int remaining = match.numTokens; + for (auto const* block : match.blocks) + { + TLLM_CHECK_DEBUG(remaining > 0); + int const numBlockTokens = std::min(remaining, static_cast(block->tokens.size())); + ret.insert(ret.end(), block->tokens.begin(), block->tokens.begin() + numBlockTokens); + remaining -= numBlockTokens; + } + TLLM_CHECK_DEBUG(remaining == 0); + return ret; +} + +// --------------------------------------------------------------------------- +// _getTreeBlock — get and validate the tree block at a committed ordinal. +// Mirrors Python's _get_tree_block(). +// --------------------------------------------------------------------------- + +SharedPtr const& KvCache::_getTreeBlock(BlockOrdinal ordinal) const +{ + TLLM_CHECK_DEBUG(mBlocks[ordinal].isCommitted()); + auto const& ret = mBlocks[ordinal].treeBlock; + TLLM_CHECK_DEBUG(ret); + if (TLLM_UNLIKELY(gDebug)) + { + auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); + auto const& beamBlock = mBlocks[ordinal].pages[kDefaultBeamIndex]; + for (LifeCycleId lcId{0}; lcId < beamBlock.size(); ++lcId) + { + if (ssmLcId.has_value() && lcId == *ssmLcId) + { + TLLM_CHECK_WITH_INFO(blockPageIsNull(beamBlock[lcId]), "SSM pages live in mSsmBlocks"); + } + else if (!blockPageIsNull(beamBlock[lcId])) + { + auto page = blockPageGetPage(beamBlock[lcId]); + auto committed = dynamicPointerCast(page); + TLLM_CHECK(committed && committed->block == ret.get()); + } + } + } + return ret; +} + +// --------------------------------------------------------------------------- +// _checkSanity — comprehensive invariant check. +// Mirrors Python's _check_sanity(). +// --------------------------------------------------------------------------- + +bool KvCache::_checkSanity() const +{ + if (mStatus == Status::CLOSED) + return numBlocks() == BlockOrdinal{0}; + + TLLM_CHECK_DEBUG(numCommittedTokens() <= mHistoryLength && mHistoryLength <= mCapacity); + TLLM_CHECK_DEBUG(numBlocks() == BlockOrdinal{divUp(mCapacity, mTokensPerBlock)}); + + auto const& lcs = mManager->lifeCycles(); + LifeCycleId numLc = mManager->storage().numLifeCycles(); + auto ssmLcId = lcs.ssmLifeCycleId(); + + // Precompute stale and scratch ranges for each lifecycle. + TypedVec> staleRanges(numLc); + TypedVec> scratchRangesVec(numLc); + for (LifeCycleId lc{0}; lc < numLc; ++lc) + { + auto const& lifecycle = lcs.getLifeCycle(lc); + staleRanges[lc] = _getStaleRange(mHistoryLength, lifecycle); + scratchRangesVec[lc] = _getScratchRange(lifecycle); + } + + for (BlockOrdinal ordinal{0}; ordinal < numBlocks(); ++ordinal) + { + auto const& block = mBlocks[ordinal]; + bool isCommitted = mNeverResumed || ordinal < BlockOrdinal{mNumCommittedBlocks}; + TLLM_CHECK_DEBUG(isCommitted == block.isCommitted()); + + for (auto const& beamBlock : block.pages) + { + TLLM_CHECK_DEBUG(beamBlock.size() == numLc); + for (LifeCycleId lc{0}; lc < numLc; ++lc) + { + auto const& bp = beamBlock[lc]; + if (ssmLcId.has_value() && lc == *ssmLcId) + { + // SSM pages live in mSsmBlocks, not in mBlocks. + // When mNeverResumed and SSM snapshot is held, the block is committed + // but SSM page entry remains null (SSM is in mSsmBlocks). + TLLM_CHECK_DEBUG(blockPageIsNull(bp)); + continue; + } + + auto const& staleRange = staleRanges[lc]; + auto const& scratchRange = scratchRangesVec[lc]; + + if (scratchRange.contains(ordinal)) + { + // Scratch blocks have no per-block pages. + TLLM_CHECK_DEBUG(blockPageIsNull(bp)); + } + else if (staleRange.beg <= ordinal && ordinal < staleRange.end) + { + if (isCommitted || mCommitState != CommitState::ALLOWED) + { + TLLM_CHECK_DEBUG(blockPageIsNull(bp)); + } + else + { + // For the decoder-side disagg case, for the first step, we will skip the + // out-of-window blocks. + TLLM_CHECK_DEBUG(std::holds_alternative>(bp) + || (blockPageIsNull(bp) && (mCommittedTokens.empty() || mManager->commitMinSnapshot()))); + } + } + else + { + if (mStatus == Status::ACTIVE) + TLLM_CHECK_DEBUG(std::holds_alternative(bp)); + else + TLLM_CHECK_DEBUG(std::holds_alternative>(bp)); + } + + if (!blockPageIsNull(bp)) + { + auto page = blockPageGetPage(bp); + TLLM_CHECK_DEBUG(isCommitted == (dynamicPointerCast(page) != nullptr)); + } + } + } + } + + // Check SSM blocks (mirrors Python lines 1342-1353). + if (ssmLcId.has_value()) + { + for (BeamIndex bi{0}; bi < mBeamWidth; ++bi) + { + auto const& bp = mSsmBlocks[bi][*ssmLcId]; + if (!blockPageIsNull(bp)) + { + if (mNeverResumed) + { + // Deferred copy: SSM holds CommittedPage from matched snapshot. + TLLM_CHECK_DEBUG(std::holds_alternative>(bp)); + auto page = blockPageGetPage(bp); + TLLM_CHECK_DEBUG(dynamicPointerCast(page) != nullptr); + } + else + { + TLLM_CHECK_DEBUG(std::holds_alternative(bp)); + auto page = blockPageGetPage(bp); + TLLM_CHECK_DEBUG(dynamicPointerCast(page) != nullptr); + } + } + } + } + + return true; +} + +// --------------------------------------------------------------------------- +// Page index tables +// --------------------------------------------------------------------------- + +void KvCache::_checkPageIndexBufferCapacity(BlockOrdinal newNumBlocks) const +{ + for (auto const& beamIndices : mBasePageIndices) + { + for (auto const& buf : beamIndices) + { + if (auto const* ext = std::get_if>(&buf)) + { + if (ext->len < newNumBlocks.value()) + { + throw std::invalid_argument("User-provided base page indices is too short"); + } + } + } + } +} + +void KvCache::_resizePageIndexBuffers(BlockOrdinal newNumBlocks) +{ + for (BeamIndex bi{0}; bi < mBeamWidth; ++bi) + { + for (LifeCycleId lcId{0}; lcId < mBasePageIndices[bi].size(); ++lcId) + { + auto& buf = mBasePageIndices[bi][lcId]; + if (auto* vec = std::get_if>(&buf)) + { + // When shrinking, assert tail entries are already BAD (mirrors Python line ~432). + auto const newSize = toSizeT(newNumBlocks); + TLLM_CHECK_DEBUG(newSize >= vec->size() + || std::all_of(vec->begin() + static_cast(newSize), vec->end(), + [](int idx) { return idx == kBadPageIndex.value(); })); + // Growing fills new entries with kBadPageIndex; shrinking truncates. + vec->resize(newSize, kBadPageIndex.value()); + } + else + { + // Span: caller-provided buffer must be large enough, + // and tail beyond active blocks must already be BAD + // (lock destructors set indices via updateBasePageIndex). + auto& ext = std::get>(buf); + int const newLen = newNumBlocks.value(); + if (ext.len < newLen) + { + throw std::invalid_argument("User-provided base page indices is too short"); + } + for (int i = newLen; i < ext.len; ++i) + TLLM_CHECK_DEBUG(ext[i] == kBadPageIndex.value()); + } + } + } +} + +int KvCache::updateBasePageIndex(BeamIndex bi, BlockOrdinal ord, LifeCycleId lc, int value) +{ + if (ord == kBadBlockOrdinal) + return kBadPageIndex.value(); // SSM pages use BAD_BLOCK_ORDINAL + auto& buf = mBasePageIndices[bi][lc]; + return std::visit( + [&](auto& b) -> int + { + using T = std::decay_t; + if constexpr (std::is_same_v>) + { + TLLM_CHECK_DEBUG(b.size() > toSizeT(ord)); + int old = b[toSizeT(ord)]; + b[toSizeT(ord)] = value; + return old; + } + else + { + TLLM_CHECK_DEBUG(ord < b.len); + int old = b[ord.value()]; + b[ord.value()] = value; + return old; + } + }, + buf); +} + +Span KvCache::getBasePageIndices(LayerGroupId lgId, BeamIndex beamIdx) const +{ + auto const& buf = mBasePageIndices.at(beamIdx).at(lgId); + auto result = std::visit( + [](auto const& b) -> Span { + return {b.data(), static_cast(b.size())}; + }, + buf); + // Cross-validate cached indices against freshly computed reference (mirrors Python lines ~350-354). + if (TLLM_UNLIKELY(gDebug) && isActive()) + { + auto ref = getAggregatedPageIndices(lgId, beamIdx); + auto len = static_cast(std::min(result.len, static_cast(ref.size()))); + TLLM_CHECK(std::equal(result.data(), result.data() + len, ref.begin())); + } + return result; +} + +std::vector KvCache::getAggregatedPageIndices(LayerGroupId lgId, BeamIndex beamIdx, bool validOnly) const +{ + std::vector result; + result.reserve(mBlocks.stdSize()); + for (auto const& sb : mBlocks) + { + auto const& pg = blockPageGetPage(sb.pages[beamIdx][lgId]); + if (!pg) + { + if (!validOnly) + result.push_back(kBadPageIndex.value()); + } + else + { + result.push_back(slotIdToPageIndexValue(pg->slotId())); + } + } + return result; +} + +void KvCache::setBasePageIndexBuf(BeamIndex beamIdx, LayerGroupId lgId, int32_t* buf, int len) +{ + auto& slot = mBasePageIndices[beamIdx][lgId]; + BlockOrdinal const numBlocks = mBlocks.size(); + + if (!buf || len == 0) + { + // Revert to internal vector: copy current data out of the old buffer. + auto const& old = slot; + if (auto const* ext = std::get_if>(&old)) + { + auto const n = std::min(toSizeT(numBlocks), static_cast(ext->len)); + std::vector vec(ext->ptr, ext->ptr + n); + slot = std::move(vec); + } + // If already a vector, nothing to do. + return; + } + + if (len < numBlocks.value()) + { + throw std::invalid_argument("setBasePageIndexBuf: buffer length must be >= num_blocks"); + } + + // Copy current indices into the external buffer (mirrors Python: buf[:length] = old_indices[:length]). + int const* oldData = nullptr; + int oldLen = 0; + if (auto* vec = std::get_if>(&slot)) + { + oldData = vec->data(); + oldLen = static_cast(vec->size()); + } + else + { + auto& span = std::get>(slot); + oldData = span.data(); + oldLen = span.len; + } + auto const copyLen = std::min(static_cast(oldLen), toSizeT(numBlocks)); + std::copy(oldData, oldData + copyLen, buf); + std::fill(buf + copyLen, buf + len, kBadPageIndex.value()); + slot = Span{buf, len}; +} + +int KvCache::getSsmBlockBaseIndex(LayerGroupId lgId, BeamIndex beamIdx) const +{ + auto const& bp = mSsmBlocks.at(beamIdx).at(lgId); + if (blockPageIsNull(bp)) + return kBadPageIndex.value(); + TLLM_CHECK_DEBUG(std::holds_alternative(bp)); + auto const& pg = blockPageGetPage(bp); + TLLM_CHECK_DEBUG_WITH_INFO(pg, "SSM block must have a valid page"); + return slotIdToPageIndexValue(pg->slotId()); // asserts valid slot +} + +// --------------------------------------------------------------------------- +// SWA scratch slot methods +// --------------------------------------------------------------------------- + +bool KvCache::hasScratchSlots() const +{ + return std::any_of(mScratchSlots.begin(), mScratchSlots.end(), [](auto const& v) { return !v.empty(); }); +} + +bool KvCache::isSwaScratchReuseEnabled() const noexcept +{ + return mEnableSwaScratchReuse; +} + +bool KvCache::supportsIndexMode(PageIndexMode mode) const +{ + switch (mode) + { + case PageIndexMode::PER_LAYER: return true; + case PageIndexMode::SHARED: return !hasScratchSlots(); + } + return false; +} + +HalfOpenRange KvCache::_getScratchRange( + LifeCycle const& lc, std::optional hlOverride, std::optional capOverride) const +{ + if (!mEnableSwaScratchReuse) + return {0, 0}; + int hist = hlOverride.value_or(mHistoryLength); + int cap = capOverride.value_or(mCapacity); + return computeScratchRange(lc, hist, cap, mTokensPerBlock, _swaScratchMaxRewindLen()); +} + +bool KvCache::_wouldUseSwaScratchBlocks() const +{ + int const maxRewindLen = _swaScratchMaxRewindLen(); + for (auto const& [lcId, lc] : mManager->lifeCycles()) + { + if (computeScratchRange(lc, mHistoryLength, mCapacity, mTokensPerBlock, maxRewindLen)) + return true; + } + return false; +} + +int KvCache::_swaScratchMaxRewindLen() const +{ + auto const& cfg = mManager->config().swaScratchReuse; + TLLM_CHECK_DEBUG(cfg.has_value()); + return cfg->maxRewindLen; +} + +std::optional KvCache::getScratchDesc(LayerGroupId lgId) const +{ + auto const& lc = mManager->lifeCycles().getLifeCycle(lgId); + auto sr = _getScratchRange(lc); + if (!sr) + return std::nullopt; + std::vector slotIds; + slotIds.reserve(mScratchSlots[lgId].size()); + for (auto const& lock : mScratchSlots[lgId]) + slotIds.push_back(slotIdToPageIndexValue(lock.slot().slotId())); + return ScratchDesc{sr, std::move(slotIds)}; +} + +void KvCache::setEnableSwaScratchReuse(bool enable) +{ + if (enable == mEnableSwaScratchReuse) + return; + if (enable) + { + if (!mManager->isSwaScratchReuseEnabled()) + throw std::invalid_argument( + "Cannot enable SWA scratch reuse for a request when it is disabled in KV cache manager config"); + if (_wouldUseSwaScratchBlocks()) + throw std::invalid_argument( + "Cannot enable SWA scratch reuse while the current request state would need scratch blocks"); + mEnableSwaScratchReuse = true; + return; + } + if (_wouldUseSwaScratchBlocks()) + throw std::invalid_argument("Cannot disable SWA scratch reuse while scratch blocks are needed"); + TLLM_CHECK_DEBUG(!hasScratchSlots()); + mEnableSwaScratchReuse = false; +} + +KvCache::DeltaScratchSlots KvCache::_takeExcessScratchSlots(int capacity, int historyLength) +{ + LifeCycleId numLc = mManager->storage().numLifeCycles(); + DeltaScratchSlots result; + result.excess.resize(numLc); + result.deltaCnt.resize(numLc, 0); + result.scratchRanges.resize(numLc); + + for (auto const& [lcIdx, lc] : mManager->lifeCycles()) + { + auto scratchRange = _getScratchRange(lc, historyLength, capacity); + result.scratchRanges[lcIdx] = scratchRange; + int numScratchBlocks = scratchRange.length(); + auto const& fracMax = mManager->storage().slotUtilFracMax(lcIdx); + int neededSlots = fracMax.ceilMul(numScratchBlocks); + int existingSlots = static_cast(mScratchSlots[lcIdx].size()); + int delta = neededSlots - existingSlots; + result.deltaCnt[lcIdx] = delta; + + if (delta < 0) + { + for (int i = 0; i < -delta; ++i) + { + result.excess[lcIdx].push_back(std::move(mScratchSlots[lcIdx].back())); + mScratchSlots[lcIdx].pop_back(); + } + } + } + return result; +} + +void KvCache::_recoverExcessScratchSlots(TypedVec>& excess) +{ + for (LifeCycleId lcId{0}; lcId < excess.size(); ++lcId) + { + for (auto& lock : excess[lcId]) + { + mScratchSlots[lcId].push_back(std::move(lock)); + } + excess[lcId].clear(); + } +} + +void KvCache::_freeScratchSlots() +{ + for (auto& lcSlots : mScratchSlots) + { + for (auto& lock : lcSlots) + lock.unlock(); + lcSlots.clear(); + } +} + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h new file mode 100644 index 000000000000..bd090197433b --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h @@ -0,0 +1,629 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "kv_cache_manager_v2/blockRadixTree.h" +#include "kv_cache_manager_v2/common.h" +#include "kv_cache_manager_v2/lifeCycleRegistry.h" +#include "kv_cache_manager_v2/movingAverage.h" +#include "kv_cache_manager_v2/page.h" +#include "kv_cache_manager_v2/pendingStats.h" +#include "kv_cache_manager_v2/utils/cudaEvent.h" + +#include "tensorrt_llm/common/assert.h" +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// Forward declarations. +class KvCacheIntrospection; +class KvCacheManager; +class StorageManager; +struct ScratchDesc; + +// --------------------------------------------------------------------------- +// BlockPage — what a SeqBlock holds per (beamIndex, lifeCycleId): +// - nullptr → no page (block not allocated for this lifecycle/beam) +// - SharedPageLock → locked (ACTIVE inference) +// - shared_ptr → held (suspended, waiting for activation) +// --------------------------------------------------------------------------- +using BlockPage = std::variant // held + >; + +inline bool blockPageIsNull(BlockPage const& bp) noexcept +{ + return std::holds_alternative(bp); +} + +inline SharedPtr const& blockPageGetPage(BlockPage const& bp) noexcept +{ + static SharedPtr const sNull{}; + if (auto* lock = std::get_if(&bp)) + return lock->isValid() ? lock->page() : sNull; + if (auto* holder = std::get_if>(&bp)) + return holder->get() ? (*holder)->page : sNull; + return sNull; +} + +// --------------------------------------------------------------------------- +// SeqBlock — one block-slot in a KvCache sequence. +// pages[beamIdx][lcId] tracks who holds/locks each page. +// treeBlock: non-null only for committed blocks (strong ref in rare cases). +// --------------------------------------------------------------------------- +using LifeCycleBlockPages = TypedVec; +using BeamBlockPages = TypedVec; + +struct SeqBlock +{ + BeamBlockPages pages; + SharedPtr treeBlock; // non-null iff committed + + bool isCommitted() const noexcept + { + bool ret = treeBlock != nullptr; + if (TLLM_UNLIKELY(gDebug)) + { + // When committed: must have 1 beam, all non-null pages must be CommittedPage. + if (ret) + { + TLLM_CHECK(pages.size() == BeamIndex{1}); + for (auto const& beamBlock : pages) + for (auto const& bp : beamBlock) + if (!blockPageIsNull(bp)) + { + auto pg = blockPageGetPage(bp); + TLLM_CHECK(!pg || dynamicPointerCast(pg)); + } + } + else + { + // When not committed: all non-null pages must be UncommittedPage. + for (auto const& beamBlock : pages) + for (auto const& bp : beamBlock) + if (!blockPageIsNull(bp)) + { + auto pg = blockPageGetPage(bp); + TLLM_CHECK(!pg || dynamicPointerCast(pg)); + } + } + } + return ret; + } +}; + +// --------------------------------------------------------------------------- +// Span — non-owning view into a contiguous buffer. +// Supports operator[] for uniform access with std::vector. +// --------------------------------------------------------------------------- +template +struct Span +{ + T* ptr; + int32_t len; + + T& operator[](int idx) + { + return ptr[idx]; + } + + T operator[](int idx) const + { + return ptr[idx]; + } + + int size() const noexcept + { + return len; + } + + T* data() const noexcept + { + return ptr; + } + + T* begin() const noexcept + { + return ptr; + } + + T* end() const noexcept + { + return ptr + len; + } +}; + +// --------------------------------------------------------------------------- +// PlannedDropHandle — tracks committed pages planned for dropping without +// owning them. Mirrors Python's PlannedDropHandle in _core/_kv_cache.py. +// +// The handle stores weak references and does not keep pages alive. Dropping it +// decrements each live page's planned-drop count and removes an already- +// droppable page from eviction tracking when no plans remain. +// --------------------------------------------------------------------------- +class PlannedDropHandle +{ +public: + // Deduplicates `pages` by identity, stores weak references, and increments + // each page's plannedDropCount. + explicit PlannedDropHandle(std::vector const& pages); + + // Mirrors Python's __del__: applies the plan if not already dropped. + ~PlannedDropHandle(); + + PlannedDropHandle(PlannedDropHandle const&) = delete; + PlannedDropHandle& operator=(PlannedDropHandle const&) = delete; + + // Apply this drop plan and invalidate the handle. + // + // A live page is removed from eviction tracking only when this is its final + // plan and it is already droppable and queued for eviction. Calling this + // method twice throws (translated to Python ValueError). + void drop(); + +private: + // nullopt once dropped (mirrors Python's `_page_refs is None`). + std::optional>> mPageRefs; +}; + +// --------------------------------------------------------------------------- +// KvCache — manages the per-sequence KV cache state. +// Mirrors Python's _KVCache. +// --------------------------------------------------------------------------- +class KvCache : public std::enable_shared_from_this +{ +public: + enum class Status + { + ACTIVE, + SUSPENDED, + CLOSED + }; + enum class CommitState + { + ALLOWED, + VIRTUAL_STOP, + USER_STOP + }; + + // Priority callback: (blockOrdinal, lifeCycleId) → Priority. + using PriorityCb = std::function; + + KvCache(KvCacheManager& manager, ReuseScope reuseScope, std::optional reuseMatch, + std::optional id, PriorityCb priorityCb, std::optional expectedPromptLength = std::nullopt); + + ~KvCache(); + + KvCache(KvCache const&) = delete; + KvCache& operator=(KvCache const&) = delete; + + // ---- State machine ----------------------------------------------------- + + // Resume: check utilization and lock all pages to GPU. + // Optionally sets a new CUDA stream; if nullopt, uses the existing one. + // Returns false if utilization too high or out of memory. + bool resume(std::optional stream = std::nullopt); + + // Suspend: detach from CUDA stream, unlock pages → PageHolder. + void suspend(); + + // Close: release all blocks back to KvCacheManager. + void close(); + + // Commit or discard request-local statistics accumulated since the previous scheduler commit. + KVCacheStatsDelta commitPendingStats(); + void discardPendingStats(); + + // Best-effort prefetch active pages to the target cache level. + bool prefetch(CacheLevel target); + + // ---- Capacity / history ------------------------------------------------ + + // Resize capacity and/or history_length. + // Returns true if the resize was a no-op shortcut. + bool resize(std::optional capacity, std::optional historyLength = std::nullopt); + + // Convenience: set only capacity or history length. + void setCapacity(int capacity); + void setHistoryLength(int historyLength); + + // ---- Committing tokens ------------------------------------------------- + + // Commit tokens: finalises the oldest uncommitted block and makes it + // available for reuse by other KvCaches. + // tokens must contain exactly tokensPerBlock tokens per call (until the last). + // is_end: if true, records a final reusable snapshot and stops committing. + // This is a terminal-memory contract: callers must not perform later writes + // to this KvCache's memory. The final live pages may be moved into the radix + // tree instead of copied (SSM state and the last partial block). + void commit(std::vector const& tokens, bool isEnd = false); + + // Stop committing (called by close() automatically). + void stopCommitting(); + + // ---- Page index queries ------------------------------------------------ + + // Get base page indices (slot_id) for beamIdx × layerGroupId. + // Returns a non-owning Span into the internal page-index buffer. + Span getBasePageIndices(LayerGroupId lgId, BeamIndex beamIdx = kDefaultBeamIndex) const; + + // Get aggregated (slot-level) page indices for one layer group + beam. + // Returns one entry per block; bad blocks yield kBadPageIndex. + // If valid_only=true, bad-index blocks are skipped entirely. + std::vector getAggregatedPageIndices( + LayerGroupId lgId, BeamIndex beamIdx = kDefaultBeamIndex, bool validOnly = false) const; + + // Zero-copy page index buffer: copy current base indices into [buf, buf+len) + // and arrange that future updateBasePageIndex calls write there too. + // Pass buf=nullptr / len=0 to revert to the internal vector. + void setBasePageIndexBuf(BeamIndex beamIdx, LayerGroupId lgId, int32_t* buf, int len); + + // ---- Introspection ----------------------------------------------------- + + Status status() const noexcept + { + return mStatus; + } + + CommitState commitState() const noexcept + { + return mCommitState; + } + + bool isActive() const noexcept + { + return mStatus == Status::ACTIVE; + } + + bool isClosed() const noexcept + { + return mStatus == Status::CLOSED; + } + + BlockOrdinal numBlocks() const noexcept + { + return mBlocks.size(); + } + + TypedVec const& blocks() const noexcept + { + return mBlocks; + } + + int numCommittedBlocks() const noexcept + { + return mNumCommittedBlocks; + } + + int numCommittedTokens() const noexcept + { + return static_cast(mCommittedTokens.size()); + } + + std::vector const& committedTokens() const noexcept + { + return mCommittedTokens; + } + + ReuseScope const& reuseScope() const noexcept + { + return mReuseScope; + } + + // Plan dropping SWA blocks needed only by the next conversation turn. + // + // The plan covers committed pages in each SWA life cycle's current attention + // window. Full-attention and attention-sink blocks are excluded because + // later turns may still need them. SSM state is not yet supported. Must be + // called after stopCommitting(). Returns nullptr without creating a plan if + // any required SWA page is unavailable. Mirrors Python's + // _KVCache.plan_committed_block_drop(). + std::unique_ptr planCommittedBlockDrop(); + + int historyLength() const noexcept + { + return mHistoryLength; + } + + int capacity() const noexcept + { + return mCapacity; + } + + int tokensPerBlock() const noexcept + { + return mTokensPerBlock; + } + + BeamIndex beamWidth() const noexcept + { + return mBeamWidth; + } + + CUstream cudaStream() const; + + // Mirrors Python's cuda_stream setter: if already on a stream AND active, + // make the new stream wait for the old one before switching (cross-stream sync). + void setCudaStream(CUstream stream) + { + if (mCudaStream.has_value()) + { + if (mStatus == Status::ACTIVE) + { + CachedCudaEvent ev(reinterpret_cast(*mCudaStream)); + ev.waitInStream(reinterpret_cast(stream)); + } + } + else + { + TLLM_CHECK_DEBUG(mStatus == Status::SUSPENDED && !mFinishEvent.has_value()); + } + mCudaStream = stream; + } + + CachedCudaEvent finishEvent() const; + + // RAII scope guard for _record_event() context manager. + // Sets mFinishEvent on construction, clears it on destruction (= Python's finally). + [[nodiscard]] auto recordEventScope() + { + TLLM_CHECK_DEBUG(!mFinishEvent.has_value()); + // When mCudaStream is nullopt the cache was never resumed — no GPU work + // was performed, so no CUDA event synchronization is needed. Blocks only + // contain PageHolders (not SharedPageLocks) whose destructors do not read + // finishEvent. Mirrors Python's _record_event() early-return path. + if (mCudaStream.has_value()) + { + mFinishEvent = CachedCudaEvent(reinterpret_cast(*mCudaStream)); + } + return FuncGuard([this]() { mFinishEvent.reset(); }); + } + + // Priority for (blockOrdinal, lifeCycleId) based on the callback. + Priority getPriority(BlockOrdinal ordinal, LifeCycleId lc) const; + + // Reference to StorageManager (for page acquisition/release). + StorageManager* storageManager() const; + + KvCacheManager& manager() const noexcept + { + return *mManager; + } + + // ---- SSM support -------------------------------------------------------- + + // Return the slot ID for the SSM block at the given layer group / beam. + // Returns kBadPageIndex if no SSM blocks are allocated. + int getSsmBlockBaseIndex(LayerGroupId lgId, BeamIndex beamIdx = kDefaultBeamIndex) const; + + // ---- SWA scratch slot management ------------------------------------------ + + // Return scratch metadata for a layer group, or nullopt if no scratch blocks. + std::optional getScratchDesc(LayerGroupId lgId) const; + + // True if any lifecycle has scratch slots allocated. + bool hasScratchSlots() const; + + // Whether SWA scratch reuse is enabled for this KvCache. + bool isSwaScratchReuseEnabled() const noexcept; + + // Enable or disable SWA scratch reuse. Throws if the transition is invalid. + void setEnableSwaScratchReuse(bool enable); + + // Whether the given page index mode is supported (SHARED requires no scratch slots). + bool supportsIndexMode(PageIndexMode mode) const; + + // ---- Internal callbacks (called by SharedPageLock) ---------------------- + + int updateBasePageIndex(BeamIndex bi, BlockOrdinal ord, LifeCycleId lc, int value); + + std::optional id; // opaque identifier (mirrors Python's id field) + +private: + friend class KvCacheIntrospection; + friend std::vector batchedLockToGpu( + KvCache& kvCache, std::vector const& targets); + + // Activate: lock all pages to GPU. mCudaStream must already be set. + // Internal — called by resume(). Not public (mirrors Python where activate() doesn't exist). + void activate(); + + // Internal helpers. + void _setupForReuse(BlockRadixTree::ReuseMatch const& match); + // Reconstruct the committed token sequence from a match's blocks (mirrors + // Python's _get_matched_tokens); used when reuse-matching no longer has the + // raw input tokens in scope. + std::vector _getMatchedTokens(BlockRadixTree::ReuseMatch const& match) const; + void _clearBlocks(); + // Copy `srcPage` into a new committed page attached to `treeBlock` for lifecycle + // `lcIdx`. When `ssmNumTokensInBlock` is set, the copy is an SsmCommittedPage + // covering that many tokens; otherwise a plain attention CommittedPage. No-op if + // the block already holds a page for this lifecycle, or on OOM in all levels. + void _copyPageToTreeBlock(SharedPtr const& treeBlock, LifeCycleId lcIdx, SharedPtr const& srcPage, + std::optional ssmNumTokensInBlock = std::nullopt); + + // Snapshot live SSM state to `treeBlock` reusable at `numTokens` committed tokens. + // If `move`, the live SSM page is moved (not copied) into the tree — the caller + // must guarantee no later writes to this KvCache's memory. + void _snapshotSsmToTreeBlock( + SharedPtr const& treeBlock, LifeCycleId ssmLcId, int numTokens, bool move = false); + + // Snapshot a partial (non-full) final block at `ordinal` into the radix tree, + // copying partial attention pages and optionally the SSM snapshot. + void _snapshotPartialBlockToTree(BlockOrdinal ordinal, bool commitSsm); + // Returns [stale_begin, stale_end) block ordinal range for a SWA lifecycle. + HalfOpenRange _getStaleRange(int historyLength, LifeCycle const& lc) const; + + // Backup entry for rollback after _unlockStaleBlocks. + struct StaleBackup + { + BlockOrdinal ordinal; + BeamIndex beamIdx; + LifeCycleId lcId; + SharedPtr holder; + }; + + // Unlock stale SWA blocks. Returns backup holders for rollback. + std::vector _unlockStaleBlocks(int historyLength); + + // Re-lock previously unlocked stale blocks (rollback on OOM). + void _lockHeldBlocks(std::vector const& backup); + + // Iterator over (ordinal, beamIdx, lcIdx) tuples for active (non-stale) pages. + // Mirrors Python's _active_pages(). Used by activate() for efficient lock. + struct ActivePage + { + BlockOrdinal ordinal; + BeamIndex beamIdx; + LifeCycleId lcId; + }; + + std::vector _activePages() const; + SharedPtr _page(BlockOrdinal ordinal, BeamIndex beamIdx, LifeCycleId lcId) const; + + bool _shortcutSetCapacity(int capacity); + bool _shortcutSetHistoryLength(int historyLength); + bool _shouldRecordStats() const; + void _refreshStatsDirtyState(); + void _recordDirectIterationStats(LifeCycleId lifeCycle, KVCacheIterationStatsDelta const& iterationStats); + void _recordMigratedSlots(std::vector> const& pages, std::vector const& slots, + CacheLevel srcLevel, CacheLevel dstLevel); + void _recordDroppedPages(std::vector> const& pages, CacheLevel cacheLevel); + void _refreshGenerationAllocReady(); + void _recordResizePendingAllocations(BlockOrdinal blockBegin, BlockOrdinal blockEnd, + TypedVec> const& excludedRanges, bool countAsGeneration); + void _subtractPendingAllocationRange(BlockOrdinal blockBegin, BlockOrdinal blockEnd); + static bool _hasReuseSource(BlockPage const& page); + void _increaseCapacity(BlockOrdinal newNumBlocks, int newHistoryLength); + void _decreaseCapacity(BlockOrdinal newNumBlocks); + + void _evictOutOfWindowBlocks(int historyLength) + { + (void) historyLength; + } // handled by _unlockStaleBlocks + + // Release stale held uncommitted pages for SWA layers after committing stops. + // Mirrors Python's _on_stop_committing(). + void _onStopCommitting(); + + // Commit a single block at ordinal `ord`. + // `isLast` mirrors Python's is_last parameter: when True (or on VIRTUAL_STOP), + // transitions to USER_STOP and calls _onStopCommitting() internally. + // Caller must have recordEventScope() open so finishEvent() works. + // `commitSsm` snapshots the current SSM state for this block; `moveSsm` + // moves (vs copies) the live SSM page into the tree (caller must guarantee + // no later writes to this KvCache's memory). Mirrors Python's _commit_block. + void _commitBlock(int ord, bool isLast, bool commitSsm = false, bool moveSsm = false); + + struct TakenPage + { + SharedPtr page; + bool locked; + }; + + // Extract uncommitted pages from a SeqBlock, resetting block page entries. + // Returns one TakenPage per lifecycle. Mirrors Python's _take_uncommitted_page(). + TypedVec _takeUncommittedPage( + SeqBlock& sb, BeamIndex beamIdx, std::optional skipLc = std::nullopt); + + // Get and validate the tree block at a committed ordinal. + // Mirrors Python's _get_tree_block(). Asserts committed pages reference the correct block. + SharedPtr const& _getTreeBlock(BlockOrdinal ordinal) const; + + // Comprehensive sanity check of KvCache invariants. + // Mirrors Python's _check_sanity(). Returns true on success (asserts internally). + bool _checkSanity() const; + + // ---- SWA scratch private helpers ------------------------------------------ + + // Compute the scratch block range for a lifecycle. + HalfOpenRange _getScratchRange(LifeCycle const& lc, std::optional hlOverride = std::nullopt, + std::optional capOverride = std::nullopt) const; + + // Result of _takeExcessScratchSlots: excess locks, per-lc delta counts, and scratch ranges. + struct DeltaScratchSlots + { + TypedVec> excess; + TypedVec deltaCnt; + TypedVec> scratchRanges; + }; + + // Compute and remove excess scratch slots for a new capacity/historyLength. + DeltaScratchSlots _takeExcessScratchSlots(int capacity, int historyLength); + + // Recover previously taken excess scratch slots back into mScratchSlots. + void _recoverExcessScratchSlots(TypedVec>& excess); + + // Release all scratch slots back to storage. + void _freeScratchSlots(); + + // Whether any lifecycle would require scratch blocks at the current state. + bool _wouldUseSwaScratchBlocks() const; + int _swaScratchMaxRewindLen() const; + + // Page index table management. + // _basePageIndices[beamIdx][lcId][blockOrdinal] = slotId or BAD + void _checkPageIndexBufferCapacity(BlockOrdinal newNumBlocks) const; + void _resizePageIndexBuffers(BlockOrdinal newNumBlocks); + + std::shared_ptr mManager; + ReuseScope mReuseScope; + PriorityCb mPriorityCb; + std::optional mCudaStream; + Status mStatus; + CommitState mCommitState; + BeamIndex mBeamWidth; + int mCapacity; + int mHistoryLength; + std::optional mExpectedPromptLength; + bool mGenerationAllocReady = false; + + // Page index tables: [beamIdx][lcId] → either an internal vector or an external span. + // Mirrors Python's IndexSeq = array.array | memoryview. + using PageIndexBuf = std::variant, Span>; + using LifeCyclePageIndexBuffers = TypedVec; + using BeamPageIndexBuffers = TypedVec; + BeamPageIndexBuffers mBasePageIndices; + + TypedVec mBlocks; + + std::vector mCommittedTokens; + int mNumCommittedBlocks; + std::optional mFinishEvent; + int mTokensPerBlock; + Average mAvgHistoryLength; + Average mAvgCapacity; + + // SSM pages: [beamIdx][lcId] — always initialized (empty entries = monostate). + BeamBlockPages mSsmBlocks; + bool mNeverResumed = true; + + PendingStats mPendingStats; + + // SWA scratch slot support. + bool mEnableSwaScratchReuse = false; + TypedVec> mScratchSlots; +}; + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp new file mode 100644 index 000000000000..ae136b0306f7 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp @@ -0,0 +1,860 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "kv_cache_manager_v2/kvCacheManager.h" +#include "kv_cache_manager_v2/common.h" +#include "kv_cache_manager_v2/exceptions.h" +#include "kv_cache_manager_v2/storage/config.h" +#include "kv_cache_manager_v2/storage/core.h" +#include "kv_cache_manager_v2/utils/math.h" + +#include "tensorrt_llm/common/assert.h" +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +static double nowSeconds() +{ + return static_cast(std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()) + .count()) + / 1e6; +} + +// --------------------------------------------------------------------------- +// PageIndexConverter +// --------------------------------------------------------------------------- + +std::vector PageIndexConverter::operator()( + std::vector const& baseIndices, std::optional indexMode, ScratchDesc const* scratch) const +{ + if (!indexMode.has_value()) + { + if (scratch) + { + throw std::invalid_argument("index_mode must be provided when scratch is active"); + } + indexMode = PageIndexMode::SHARED; + } + + int appliedLayerOffset = (*indexMode == PageIndexMode::PER_LAYER) ? layerOffset : 0; + int scratchPages = scratchPagesPerBlock; + + std::vector result; + result.reserve(baseIndices.size() * static_cast(expansion)); + + for (BlockOrdinal ordinal{0}; ordinal < BlockOrdinal{static_cast(baseIndices.size())}; ++ordinal) + { + int index; + if (scratch && scratch->range.contains(ordinal)) + { + // Scratch block: slot IDs come from ScratchDesc, not base_indices. + int blockPos = ordinal - scratch->range.beg; + int totalOffset = blockPos * scratchPages; + int slotIdx = totalOffset / scale; + int slotId = scratch->slotIds[static_cast(slotIdx)]; + int offset = totalOffset % scale; + index = slotId * scale + (offset + appliedLayerOffset) % scale; + } + else if (baseIndices[toSizeT(ordinal)] == kBadPageIndex.value()) + { + index = kBadPageIndex.value(); + } + else + { + index = baseIndices[toSizeT(ordinal)] * scale + appliedLayerOffset; + } + for (int i = 0; i < expansion; ++i) + { + result.push_back(index != kBadPageIndex.value() ? index * expansion + i : kBadPageIndex.value()); + } + } + return result; +} + +std::vector PageIndexConverter::operator()(int baseIndex) const +{ + return operator()(std::vector{baseIndex}, std::nullopt, nullptr); +} + +// --------------------------------------------------------------------------- +// KvCacheManager +// --------------------------------------------------------------------------- + +KvCacheManager::KvCacheManager(KVCacheManagerConfig const& config, std::shared_ptr eventSink) + : mConfig(config) + , mLifeCycles(config) + , mEventSink(std::move(eventSink)) + , mAvgReusedLength(0.9999) + , mAvgSqrCapacity(0.9999) + , mAvgSqrHistoryLength(0.9999) +{ + mConfig.validate(); + + mRadixTree = std::make_shared(mLifeCycles, mConfig.tokensPerBlock, mEventSink); + + StorageConfig storageConfig = createStorageConfig(mConfig); + mStorage + = std::make_shared(mLifeCycles, storageConfig, mConfig.tokensPerBlock, mConfig.swaScratchReuse, + mConfig.typicalStep, mConfig.constraints, mConfig.initialPoolRatio, mEventSink, mConfig.maxUtilForResume); + + mTargetRatioListGpu = _currentGpuRatio(); + mTargetRatioListOther = _currentOtherRatios(); + _resetIterationPeakNumBlocks(); + + mLastAdjustmentTime = nowSeconds(); +} + +KvCacheManager::~KvCacheManager() +{ + shutdown(); +} + +void KvCacheManager::shutdown() +{ + clearReusableBlocks(); + TLLM_CHECK_DEBUG(mStorage); + + // Post-condition: after clearing all reusable blocks and with no active KvCaches, + // no evictable pages should remain. + for (auto const& lvl : mStorage->mLevels) + { + for (PoolGroupIndex pgIdx{0}; pgIdx < lvl.numPoolGroups(); ++pgIdx) + { + TLLM_CHECK_DEBUG(lvl.controller.numEvictablePages(pgIdx) == 0); + } + } + + mStorage->destroy(); +} + +void KvCacheManager::clearReusableBlocks() +{ + TLLM_CHECK_DEBUG(mRadixTree); + mRadixTree->clear(); +} + +std::shared_ptr KvCacheManager::createKvCache(ReuseScope reuseScope, + std::vector const& inputTokens, std::optional id, KvCache::PriorityCb priorityCb, + std::optional expectedPromptLength) +{ + if (!priorityCb) + { + priorityCb = [](BlockOrdinal, LifeCycleId) { return kPriorityDefault; }; + } + + if (!expectedPromptLength.has_value() && !inputTokens.empty()) + { + expectedPromptLength = static_cast(inputTokens.size()); + } + + // Compute the reuse match once here (shared conceptually with probeReuse) and + // hand it to the KvCache, rather than having the cache re-walk the radix tree. + // Mirrors Python KVCacheManager.allocate() passing a ReuseMatch into _KVCache. + std::optional reuseMatch; + if (!inputTokens.empty()) + { + reuseMatch = matchReuse(reuseScope, inputTokens); + } + + return std::make_shared(*this, std::move(reuseScope), std::move(reuseMatch), std::move(id), + std::move(priorityCb), expectedPromptLength); +} + +BlockRadixTree::ReuseMatch KvCacheManager::matchReuse( + ReuseScope const& reuseScope, std::vector const& inputTokens) const +{ + return mRadixTree->match(reuseScope, inputTokens, enablePartialMatch()); +} + +int KvCacheManager::probeReuse(ReuseScope reuseScope, std::vector const& inputTokens) const +{ + return matchReuse(reuseScope, inputTokens).numTokens; +} + +// ---- Memory pool queries -------------------------------------------------- + +MemAddress KvCacheManager::getMemPoolBaseAddress( + LayerId layerId, DataRole role, std::optional indexMode) const +{ + auto const& attr = mStorage->getBufferAttr(layerId, role); + + if (!indexMode.has_value()) + { + if (mConfig.enableSwaScratchReuse()) + { + throw std::invalid_argument("index_mode must be provided when SWA scratch reuse is enabled"); + } + indexMode = PageIndexMode::SHARED; + } + + PoolGroupIndex pgIdx = mStorage->getPoolGroupIndex(attr.lifeCycleId); + MemAddress addr = mStorage->getMemPoolBaseAddress(pgIdx, attr.poolIndex); + if (*indexMode == PageIndexMode::SHARED) + { + addr = MemAddress(addr + attr.offset); + } + return addr; +} + +int KvCacheManager::getPageStride(LayerId layerId, DataRole role) const +{ + auto const& attr = mStorage->getBufferAttr(layerId, role); + return exactDiv(static_cast(attr.size), attr.expansion); +} + +size_t KvCacheManager::getPageIndexUpperBound(LayerId layerId, DataRole role) const +{ + auto const& attr = mStorage->getBufferAttr(layerId, role); + LifeCycleId lc = attr.lifeCycleId; + PoolGroupIndex pg = mStorage->getPoolGroupIndex(lc); + SlotCount const numSlots = mStorage->numSlots(pg, kGpuLevel); + auto slotSizes = mStorage->slotSize(pg); + size_t slotSize = slotSizes.at(attr.poolIndex); + return (exactDiv(slotSize, attr.size) * slotCountToSizeT(numSlots) - exactDiv(attr.offset, attr.size)) + * static_cast(attr.expansion); +} + +int KvCacheManager::getPageIndexScale(LayerId layerId, DataRole role) const +{ + auto const& attr = mStorage->getBufferAttr(layerId, role); + return mStorage->mSlotToPageIndices.at(attr.lifeCycleId).at(attr.poolIndex); +} + +PageIndexConverter KvCacheManager::getPageIndexConverter(LayerId layerId, DataRole role) const +{ + auto const& attr = mStorage->getBufferAttr(layerId, role); + auto const& layerAttr = mStorage->getLayerAttr(layerId); + int scale = mStorage->mSlotToPageIndices.at(attr.lifeCycleId).at(attr.poolIndex); + int offset = exactDiv(static_cast(attr.offset), static_cast(attr.size)); + int scratchPages = layerAttr.slotUtil.at(attr.poolIndex); + return PageIndexConverter{scale, attr.expansion, offset, scratchPages}; +} + +std::optional KvCacheManager::supportsIndexMode(PageIndexMode mode) const +{ + switch (mode) + { + case PageIndexMode::PER_LAYER: return true; + case PageIndexMode::SHARED: return mConfig.enableSwaScratchReuse() ? std::optional(std::nullopt) : true; + } + return std::nullopt; +} + +// ---- getAggregatedPages --------------------------------------------------- + +std::vector KvCacheManager::getAggregatedPages(std::vector const& buffers) const +{ + using Key = std::pair; + + struct Entry + { + size_t start; + size_t end; + ExpandedBuffer eb; + }; + + std::map> groups; + + for (auto const& bufferId : buffers) + { + auto it = mStorage->mBufferAttr.find(bufferId); + if (it == mStorage->mBufferAttr.end()) + throw std::out_of_range("getAggregatedPages: unknown buffer id"); + + auto const& attr = it->second; + size_t start = attr.offset; + size_t end = attr.offset + attr.size; + Key key{attr.lifeCycleId, attr.poolIndex}; + groups[key].push_back({start, end, ExpandedBuffer{bufferId, attr.expansion}}); + } + + std::vector result; + for (auto& [key, entries] : groups) + { + auto [lifeCycleId, poolIdx] = key; + auto pgIdx = mStorage->getPoolGroupIndex(lifeCycleId); + + std::sort(entries.begin(), entries.end(), [](Entry const& a, Entry const& b) { return a.start < b.start; }); + + auto const poolBase = mStorage->getMemPoolBaseAddress(pgIdx, poolIdx); + size_t stride = mStorage->slotSize(pgIdx).at(poolIdx); + + auto flush + = [&, lifeCycleId = lifeCycleId](size_t start, size_t end, std::vector& buffersInRange) + { + result.push_back(AggregatedPageDesc{ + MemAddress(poolBase + start), end - start, stride, lifeCycleId, std::move(buffersInRange)}); + }; + + size_t currentStart = entries.front().start; + size_t currentEnd = entries.front().end; + std::vector currentBuffers{entries.front().eb}; + for (size_t i = 1; i < entries.size(); ++i) + { + if (entries[i].start == currentEnd) + { + currentEnd = entries[i].end; + currentBuffers.push_back(entries[i].eb); + continue; + } + + flush(currentStart, currentEnd, currentBuffers); + currentStart = entries[i].start; + currentEnd = entries[i].end; + currentBuffers = {entries[i].eb}; + } + flush(currentStart, currentEnd, currentBuffers); + } + + return result; +} + +// ---- Pool group layout ----------------------------------------------------- + +TypedVec KvCacheManager::poolGroupDescs() const +{ + auto const& slotDescList = mStorage->slotDescList(); + TypedVec result; + result.reserve(slotDescList.size()); + + for (PoolGroupIndex pgIdx{0}; pgIdx < slotDescList.size(); ++pgIdx) + { + auto const& slotDesc = slotDescList.at(pgIdx); + auto slotSizeList = mStorage->slotSize(pgIdx); + + TypedVec pools; + pools.reserve(slotSizeList.size()); + for (PoolIndex poolIdx{0}; poolIdx < slotSizeList.size(); ++poolIdx) + { + pools.push_back( + PoolDesc{poolIdx, mStorage->getMemPoolBaseAddress(pgIdx, poolIdx), slotSizeList.at(poolIdx)}); + } + + result.push_back(PoolGroupDesc{pgIdx, mStorage->numSlots(pgIdx, kGpuLevel), slotDesc, std::move(pools)}); + } + + return result; +} + +// ---- Query / info --------------------------------------------------------- + +int KvCacheManager::tokensPerBlock() const noexcept +{ + return mRadixTree->tokensPerBlock(); +} + +bool KvCacheManager::enablePartialMatch() const noexcept +{ + return mConfig.enablePartialReuse; +} + +int KvCacheManager::numLayers() const noexcept +{ + return static_cast(mStorage->layerToLifeCycleIds().size()); +} + +std::vector KvCacheManager::layerIds() const +{ + std::vector ids; + for (auto const& [lid, lc] : mStorage->layerToLifeCycleIds()) + ids.push_back(lid); + return ids; +} + +LayerGroupId KvCacheManager::getLayerGroupId(LayerId layerId) const +{ + return mStorage->layerToLifeCycleIds().at(layerId); +} + +TypedVec> KvCacheManager::layerGrouping() const +{ + LifeCycleId numLc = mLifeCycles.size(); + TypedVec> result(numLc); + for (auto const& [lid, lc] : mStorage->layerToLifeCycleIds()) + { + result.at(lc).push_back(lid); + } + return result; +} + +// ---- Resize --------------------------------------------------------------- + +bool KvCacheManager::resize(CacheLevel level, size_t quota, bool bestEfforts) +{ + if (bestEfforts) + throw std::runtime_error("best_efforts resize not implemented"); + try + { + _adjustLevel(level, quota); + return true; + } + catch (std::exception const& e) + { + return false; + } +} + +size_t KvCacheManager::getQuota(CacheLevel level) const +{ + return mStorage->mLevels.at(level).storage->totalQuota(); +} + +// ---- Statistics ---------------------------------------------------------- + +void KvCacheManager::commitStats( + KVCacheStatsDelta const& stats, IterationStatsByLifeCycle const& iterationStatsByLifeCycle) +{ + if (!mConfig.enableStats) + { + return; + } + + _updateIterationPeakNumBlocks(); + mCommittedStats.add(stats); + for (auto const& [lifeCycle, iterationStats] : iterationStatsByLifeCycle) + { + if (!iterationStats.empty()) + { + mIterationStatsByLifeCycle[lifeCycle].add(iterationStats); + } + } +} + +KVCacheStatsDelta KvCacheManager::getCommittedStats() const +{ + return mCommittedStats.copy(); +} + +IterationStatsByLifeCycle KvCacheManager::getAndResetIterationStats() +{ + IterationStatsByLifeCycle stats; + for (auto const& [lifeCycle, delta] : mIterationStatsByLifeCycle) + { + if (!delta.empty()) + { + stats.emplace(lifeCycle, delta.copy()); + } + } + mIterationStatsByLifeCycle.clear(); + return stats; +} + +void KvCacheManager::commitSsmSnapshotIterationStats(SsmSnapshotIterationStatsByLifeCycle const& statsByLifeCycle) +{ + if (!mConfig.enableStats) + { + return; + } + for (auto const& [lifeCycle, iterationStats] : statsByLifeCycle) + { + if (!iterationStats.empty()) + { + mSsmSnapshotIterationStatsByLifeCycle[lifeCycle].add(iterationStats); + } + } +} + +SsmSnapshotIterationStatsByLifeCycle KvCacheManager::getAndResetSsmSnapshotIterationStats() +{ + SsmSnapshotIterationStatsByLifeCycle stats; + for (auto const& [lifeCycle, delta] : mSsmSnapshotIterationStatsByLifeCycle) + { + if (!delta.empty()) + { + stats.emplace(lifeCycle, delta.copy()); + } + } + mSsmSnapshotIterationStatsByLifeCycle.clear(); + return stats; +} + +PeakBlockStatsByCacheLevel KvCacheManager::_currentBlockStatsByCacheLevel() const +{ + PeakBlockStatsByCacheLevel result(mStorage->numCacheLevels()); + for (CacheLevel cacheLevel{0}; cacheLevel < mStorage->numCacheLevels(); ++cacheLevel) + { + auto& levelStats = result[cacheLevel]; + levelStats.resize(mStorage->numPoolGroups()); + for (PoolGroupIndex poolGroup{0}; poolGroup < mStorage->numPoolGroups(); ++poolGroup) + { + auto const stats = mStorage->getStatistics(cacheLevel, poolGroup); + levelStats[poolGroup] = {stats.available(), stats.unavailable(), stats.evictable}; + } + } + return result; +} + +void KvCacheManager::_resetIterationPeakNumBlocks(std::optional cacheLevel) +{ + if (!cacheLevel.has_value()) + { + mIterationPeakNumBlocksByCacheLevel = _currentBlockStatsByCacheLevel(); + return; + } + + PeakBlockStatsByPoolGroup levelStats(mStorage->numPoolGroups()); + for (PoolGroupIndex poolGroup{0}; poolGroup < mStorage->numPoolGroups(); ++poolGroup) + { + auto const stats = mStorage->getStatistics(*cacheLevel, poolGroup); + levelStats[poolGroup] = {stats.available(), stats.unavailable(), stats.evictable}; + } + mIterationPeakNumBlocksByCacheLevel[*cacheLevel] = std::move(levelStats); +} + +void KvCacheManager::_updateIterationPeakNumBlocks() +{ + auto const current = _currentBlockStatsByCacheLevel(); + for (CacheLevel cacheLevel{0}; cacheLevel < current.size(); ++cacheLevel) + { + auto& peakLevel = mIterationPeakNumBlocksByCacheLevel[cacheLevel]; + for (PoolGroupIndex poolGroup{0}; poolGroup < current[cacheLevel].size(); ++poolGroup) + { + auto& peak = peakLevel[poolGroup]; + auto const& value = current[cacheLevel][poolGroup]; + peak.available = std::max(peak.available, value.available); + peak.unavailable = std::max(peak.unavailable, value.unavailable); + peak.evictable = std::max(peak.evictable, value.evictable); + } + } +} + +PeakBlockStatsByPoolGroup KvCacheManager::getAndResetIterationPeakBlockStats(CacheLevel cacheLevel) +{ + _updateIterationPeakNumBlocks(); + PeakBlockStatsByPoolGroup peak = mIterationPeakNumBlocksByCacheLevel.at(cacheLevel); + _resetIterationPeakNumBlocks(cacheLevel); + return peak; +} + +void KvCacheManager::markStatsDirty(std::optional kvCacheId) +{ + if (kvCacheId.has_value()) + { + mDirtyStatsKvCacheIds.insert(*kvCacheId); + } +} + +void KvCacheManager::clearStatsDirty(std::optional kvCacheId) +{ + if (kvCacheId.has_value()) + { + mDirtyStatsKvCacheIds.erase(*kvCacheId); + } +} + +std::unordered_set KvCacheManager::getDirtyStatsKvCacheIds() const +{ + return mDirtyStatsKvCacheIds; +} + +void KvCacheManager::markStatsExcluded(std::optional kvCacheId) +{ + if (kvCacheId.has_value()) + { + mStatsExcludedKvCacheIds.insert(*kvCacheId); + clearStatsDirty(kvCacheId); + } +} + +void KvCacheManager::clearStatsExcluded(std::optional kvCacheId) +{ + if (kvCacheId.has_value()) + { + mStatsExcludedKvCacheIds.erase(*kvCacheId); + } +} + +bool KvCacheManager::isStatsExcluded(std::optional kvCacheId) const +{ + return kvCacheId.has_value() && mStatsExcludedKvCacheIds.find(*kvCacheId) != mStatsExcludedKvCacheIds.end(); +} + +TypedVec KvCacheManager::cacheTierList() const +{ + TypedVec result; + result.reserve(mStorage->mLevels.size()); + for (auto const& lvl : mStorage->mLevels) + { + result.push_back(lvl.cacheTier); + } + return result; +} + +std::vector KvCacheManager::allBufferIds() const +{ + std::vector result; + result.reserve(mStorage->mBufferAttr.size()); + for (auto const& item : mStorage->mBufferAttr) + result.push_back(item.first); + return result; +} + +int KvCacheManager::clampMaxSeqLenForMem(int batchSize, int tokenNumUpperBound) const +{ + TLLM_CHECK_DEBUG(batchSize > 0); + int tokPerBlock = tokensPerBlock(); + PoolGroupIndex numPg = mStorage->numPoolGroups(); + auto const& lcs = mLifeCycles; + auto const& lcGrouping = mStorage->mLifeCycleGrouping; + + // Remaining slot counts per pool group. + TypedVec remainingSlots(numPg); + for (PoolGroupIndex pgIdx{0}; pgIdx < numPg; ++pgIdx) + { + remainingSlots[pgIdx] = mStorage->numSlots(pgIdx); + } + + // Compute required slot counts per pool group for a given seq_len. + auto getNumSlots = [&](int seqLen) -> TypedVec + { + TypedVec ret(numPg, 0); + for (LifeCycleId lifeCycleId{0}; lifeCycleId < lcs.size(); ++lifeCycleId) + { + auto staleRange = getStaleRange(lcs[lifeCycleId], seqLen, tokPerBlock); + int numStaleBlocks = staleRange.end - staleRange.beg; + int numSlots = divUp(seqLen, tokPerBlock) - numStaleBlocks; + auto pgIdx = lcGrouping[lifeCycleId]; + ret[pgIdx] += numSlots; + } + return ret; + }; + + // Reserve slots for (batch_size - 1) minimal sequences. + auto minSlots = getNumSlots(1); + for (PoolGroupIndex pgIdx{0}; pgIdx < numPg; ++pgIdx) + { + TLLM_CHECK_DEBUG(minSlots[pgIdx] >= 0); + SlotCount const reservedSlots = minSlots[pgIdx] * (batchSize - 1); + remainingSlots[pgIdx] -= reservedSlots; + if (remainingSlots[pgIdx] < 0) + { + return 0; + } + } + + auto isEnough = [&](int numBlocks) -> bool + { + auto needed = getNumSlots(numBlocks * tokPerBlock); + for (PoolGroupIndex pgIdx{0}; pgIdx < numPg; ++pgIdx) + { + TLLM_CHECK_DEBUG(needed[pgIdx] >= 0); + if (needed[pgIdx] > remainingSlots[pgIdx]) + { + return false; + } + } + return true; + }; + + if (!isEnough(1)) + { + return 0; + } + int lb = 1; + int ub = divUp(tokenNumUpperBound, tokPerBlock); + if (isEnough(ub)) + { + return tokenNumUpperBound; + } + while (lb < ub - 1) + { + int mid = (lb + ub) / 2; + if (isEnough(mid)) + { + lb = mid; + } + else + { + ub = mid; + } + } + return std::min(lb * tokPerBlock, tokenNumUpperBound); +} + +void KvCacheManager::_adjustLevel(CacheLevel level, size_t quota) +{ + auto const& ratioList = _getTargetRatioList(level); + TypedVec>> const* persistent = nullptr; + TypedVec>> persistentPages; + if (mStorage->isLastLevel(level)) + { + persistentPages = _gatherPersistentPages(); + persistent = &persistentPages; + } + mStorage->adjustCacheLevel(level, quota, ratioList, persistent); +} + +TypedVec>> KvCacheManager::_gatherPersistentPages() const +{ + CacheLevel lastLevel = mStorage->numCacheLevels() - 1; + PoolGroupIndex numPg = mStorage->numPoolGroups(); + TypedVec>> result(numPg); + + for (KvCache* kvc : mLivingKvCaches) + { + TLLM_CHECK_DEBUG(kvc->status() == KvCache::Status::SUSPENDED); + for (auto const& sb : kvc->blocks()) + { + for (auto const& beamPages : sb.pages) + { + for (LifeCycleId lc{0}; lc < beamPages.size(); ++lc) + { + // Mirrors Python: holder must be _PageHolder type (suspended state). + if (blockPageIsNull(beamPages[lc])) + { + continue; + } + TLLM_CHECK_DEBUG_WITH_INFO(std::holds_alternative>(beamPages[lc]), + "Non-null holder must be PageHolder in suspended state"); + auto const& pg = blockPageGetPage(beamPages[lc]); + if (!pg) + { + continue; + } + // Mirrors Python assertions for invariant checking. + TLLM_CHECK_DEBUG_WITH_INFO( + pg->status() == PageStatus::HELD, "Page in suspended KvCache must be HELD"); + TLLM_CHECK_DEBUG_WITH_INFO((pg->scheduledForEviction() == (pg->cacheLevel != lastLevel)), + "Eviction scheduling invariant violated"); + if (pg->scheduledForEviction()) + { + continue; + } + PoolGroupIndex pgIdx = mStorage->getPoolGroupIndex(lc); + result[pgIdx].push_back(pg); + } + } + } + } + return result; +} + +// ---- KvCache registry ----------------------------------------------------- + +void KvCacheManager::registerKvCache(KvCache* kvc) +{ + mLivingKvCaches.insert(kvc); + ++mNumCreatedKvCaches; +} + +void KvCacheManager::unregisterKvCache(KvCache* kvc) +{ + mLivingKvCaches.erase(kvc); +} + +void KvCacheManager::tryUpdateTargetRatios() +{ + if (mNumSampledKvCaches - mLastUpdateNumSampledKvCaches < 100) + return; + mLastUpdateNumSampledKvCaches = mNumSampledKvCaches; + + int tokensPerBlock = mConfig.tokensPerBlock; + int avgReusedLength = static_cast(std::round(mAvgReusedLength.value())); + int avgCapacity = static_cast(std::round(std::sqrt(mAvgSqrCapacity.value()))); + int avgHistoryLength = static_cast(std::round(std::sqrt(mAvgSqrHistoryLength.value()))); + if (avgCapacity > 0) + mTargetRatioListGpu + = mStorage->constrainRatio(mStorage->ratioFromLength(tokensPerBlock, avgHistoryLength, avgCapacity)); + if (avgReusedLength > 0) + mTargetRatioListOther = mStorage->ratioFromLength(tokensPerBlock, avgReusedLength, avgReusedLength); +} + +TypedVec KvCacheManager::_currentGpuRatio() const +{ + return mStorage->getRatioList(kGpuLevel); +} + +TypedVec KvCacheManager::_currentOtherRatios() const +{ + CacheLevel numLevels = mStorage->numCacheLevels(); + if (numLevels == CacheLevel{1}) + { + return _currentGpuRatio(); + } + PoolGroupIndex numPg = mStorage->numPoolGroups(); + TypedVec result(numPg, 0.f); + for (CacheLevel lvl{1}; lvl < numLevels; ++lvl) + { + auto ratios = mStorage->getRatioList(lvl); + for (PoolGroupIndex pgIdx{0}; pgIdx < numPg; ++pgIdx) + { + result[pgIdx] += ratios[pgIdx]; + } + } + float denom = static_cast((numLevels - 1).value()); + for (auto& r : result) + { + r /= denom; + } + return result; +} + +// ---- needAdjustment / adjust ----------------------------------------------- + +TypedVec const& KvCacheManager::_getTargetRatioList(CacheLevel level) const +{ + return (level == kGpuLevel) ? mTargetRatioListGpu : mTargetRatioListOther; +} + +bool KvCacheManager::_needAdjustment(CacheLevel level) const +{ + auto const& target = _getTargetRatioList(level); + auto current = (level == kGpuLevel) ? _currentGpuRatio() : _currentOtherRatios(); + constexpr float kThreshold = 1.25f; + for (PoolGroupIndex pgIdx{0}; pgIdx < target.size() && pgIdx < current.size(); ++pgIdx) + { + TLLM_CHECK_DEBUG_WITH_INFO(current[pgIdx] > 0.f && target[pgIdx] > 0.f, "ratios must not be zero"); + float ratio = target[pgIdx] / current[pgIdx]; + if (ratio < 1.f / kThreshold || ratio > kThreshold) + return true; + } + return false; +} + +bool KvCacheManager::needAdjustment() const +{ + if (mNumSampledKvCaches < 2000) + return false; + double now = nowSeconds(); + if (now - mLastAdjustmentTime < 120.0) + return false; + CacheLevel lastLevel = mStorage->numCacheLevels() - 1; + return _needAdjustment(kGpuLevel) || _needAdjustment(lastLevel); +} + +void KvCacheManager::adjust() +{ + for (KvCache* kvc : mLivingKvCaches) + TLLM_CHECK_DEBUG(kvc->status() == KvCache::Status::SUSPENDED); + + CacheLevel numLevels = mStorage->numCacheLevels(); + for (CacheLevel level{0}; level < numLevels; ++level) + { + if (_needAdjustment(level)) + _adjustLevel(level, getQuota(level)); + } + mLastAdjustmentTime = nowSeconds(); +} + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h new file mode 100644 index 000000000000..a9737b5b282c --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h @@ -0,0 +1,344 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "kv_cache_manager_v2/blockRadixTree.h" +#include "kv_cache_manager_v2/common.h" +#include "kv_cache_manager_v2/config.h" +#include "kv_cache_manager_v2/eventSink.h" +#include "kv_cache_manager_v2/kvCache.h" +#include "kv_cache_manager_v2/lifeCycleRegistry.h" +#include "kv_cache_manager_v2/movingAverage.h" +#include "kv_cache_manager_v2/stats.h" +#include "kv_cache_manager_v2/storageManager.h" + +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// --------------------------------------------------------------------------- +// PoolDesc / PoolGroupDesc — describe GPU memory pool layout. +// --------------------------------------------------------------------------- +struct PoolDesc +{ + PoolIndex poolIndex{0}; + MemAddress baseAddress = 0; + size_t slotBytes = 0; +}; + +struct PoolGroupDesc +{ + PoolGroupIndex poolGroupIndex{0}; + SlotCount numSlots = 0; + SlotDesc slotDesc; + TypedVec pools; +}; + +// --------------------------------------------------------------------------- +// ExpandedBuffer / AggregatedPageDesc — returned by getAggregatedPages(). +// --------------------------------------------------------------------------- +struct ExpandedBuffer +{ + BufferId id; + int expansion; // expansion factor (tokens_per_block / tokens_per_block_override) +}; + +struct AggregatedPageDesc +{ + MemAddress base; // pool base address + buffer offset + size_t size; // byte span of this aggregated buffer group + size_t stride; // slot size (bytes per slot in the pool group) + LifeCycleId layerGroupId; // pool group / life-cycle id + std::vector buffers; // constituent buffers in offset order +}; + +// --------------------------------------------------------------------------- +// ScratchDesc — scratch metadata for one layer group of one sequence. +// Scratch blocks store ephemeral KV data using shared coalesced slots. +// Mirrors _kv_cache_manager.py::ScratchDesc. +// --------------------------------------------------------------------------- +struct ScratchDesc +{ + HalfOpenRange range; // block ordinal range [beg, end) + std::vector slotIds; // scratch slot IDs, length = ceil(numScratchBlocks / scale) + + explicit operator bool() const noexcept + { + return static_cast(range); + } +}; + +// --------------------------------------------------------------------------- +// PageIndexConverter — convert base page index → kernel page indices. +// --------------------------------------------------------------------------- +struct PageIndexConverter +{ + int scale; + int expansion; + int layerOffset = 0; // sub-page offset within coalesced slot + int scratchPagesPerBlock = 1; // sub-pages per block for scratch allocation + + // Convert a sequence of base page indices to per-layer page indices. + // indexMode: SHARED (default) or PER_LAYER. When scratch is active, must be PER_LAYER. + // scratch: optional scratch descriptor from KvCache::getScratchDesc(). + // Mirrors _kv_cache_manager.py::PageIndexConverter.__call__. + std::vector operator()(std::vector const& baseIndices, + std::optional indexMode = std::nullopt, ScratchDesc const* scratch = nullptr) const; + + // Backward-compatible single-index overload. + std::vector operator()(int baseIndex) const; +}; + +// --------------------------------------------------------------------------- +// KvCacheManager — top-level KV cache manager. +// Mirrors Python's KVCacheManager. +// --------------------------------------------------------------------------- +class KvCacheManager : public std::enable_shared_from_this +{ +public: + explicit KvCacheManager(KVCacheManagerConfig const& config, std::shared_ptr eventSink = nullptr); + ~KvCacheManager(); + + KvCacheManager(KvCacheManager const&) = delete; + KvCacheManager& operator=(KvCacheManager const&) = delete; + + // ---- Lifecycle -------------------------------------------------------- + + void shutdown(); + + // Clear all reusable (committed) blocks from the radix tree. + void clearReusableBlocks(); + + // ---- KvCache creation ------------------------------------------------- + + // Create a new KvCache. Returned cache is SUSPENDED; call activate() with a stream. + // input_tokens: optional sequence to match against existing cached blocks. + // priorityCb: optional priority override per block. + std::shared_ptr createKvCache(ReuseScope reuseScope = {}, std::vector const& inputTokens = {}, + std::optional id = std::nullopt, KvCache::PriorityCb priorityCb = {}, + std::optional expectedPromptLength = std::nullopt); + + BlockRadixTree::ReuseMatch matchReuse( + ReuseScope const& reuseScope, std::vector const& inputTokens) const; + int probeReuse(ReuseScope reuseScope = {}, std::vector const& inputTokens = {}) const; + + // ---- Memory pool queries ----------------------------------------------- + + // Base address of the memory pool. When indexMode is PER_LAYER, returns pool group base + // (without per-layer offset). When SHARED, returns per-layer base (with offset baked in). + MemAddress getMemPoolBaseAddress( + LayerId layerId, DataRole role, std::optional indexMode = std::nullopt) const; + + int getPageStride(LayerId layerId, DataRole role) const; + size_t getPageIndexUpperBound(LayerId layerId, DataRole role) const; + + // Scale factor: base_page_index * scale → kernel page index. + int getPageIndexScale(LayerId layerId, DataRole role) const; + + // Composite converter (scale + expansion). + PageIndexConverter getPageIndexConverter(LayerId layerId, DataRole role) const; + + // Group a set of BufferIds into contiguous AggregatedPageDesc descriptors. + // Mirrors Python's KVCacheManager.get_aggregated_pages(). + std::vector getAggregatedPages(std::vector const& buffers) const; + + TypedVec poolGroupDescs() const; + + // ---- Query / info ------------------------------------------------------ + + int tokensPerBlock() const noexcept; + bool enablePartialMatch() const noexcept; + + bool commitMinSnapshot() const noexcept + { + return mConfig.commitMinSnapshot; + } + + bool isSwaScratchReuseEnabled() const noexcept + { + return mConfig.enableSwaScratchReuse(); + } + + // Whether managed KV caches support the given page index mode. + // Returns true/false for a definitive answer, nullopt for per-instance check. + std::optional supportsIndexMode(PageIndexMode mode) const; + + bool allowSeqRebasing() const noexcept + { + return true; + } + + int numLayers() const noexcept; + + std::vector layerIds() const; + LayerGroupId getLayerGroupId(LayerId layerId) const; + + // Layer grouping: layers with the same lifecycle share pool allocation. + // NOTE: the iteration order of the layer lists (and of the groups) is NOT + // part of the API contract and may differ across backends/runs. Do not rely + // on it to infer buffer/pool memory order — query poolGroupDescs() + // (PoolGroupDesc::pools[i].baseAddress + coalescedBuffers) for that. + TypedVec> layerGrouping() const; + + // Iterator over all buffer identifiers. Mirrors Python's all_buffer_ids property. + std::vector allBufferIds() const; + + // Sorted by CacheLevel from warm to cold. Mirrors Python's cache_tier_list property. + TypedVec cacheTierList() const; + + // Get the max possible sequence length limited by GPU memory pools. + // Mirrors Python's clamp_max_seq_len_for_mem(). + int clampMaxSeqLenForMem(int batchSize, int tokenNumUpperBound) const; + + // ---- Resize ----------------------------------------------------------- + + bool resize(CacheLevel level, size_t quota, bool bestEfforts = false); + size_t getQuota(CacheLevel level) const; + + // ---- Statistics ------------------------------------------------------- + + void commitStats(KVCacheStatsDelta const& stats, IterationStatsByLifeCycle const& iterationStatsByLifeCycle = {}); + KVCacheStatsDelta getCommittedStats() const; + IterationStatsByLifeCycle getAndResetIterationStats(); + PeakBlockStatsByPoolGroup getAndResetIterationPeakBlockStats(CacheLevel cacheLevel); + + void commitSsmSnapshotIterationStats(SsmSnapshotIterationStatsByLifeCycle const& statsByLifeCycle); + SsmSnapshotIterationStatsByLifeCycle getAndResetSsmSnapshotIterationStats(); + + void markStatsDirty(std::optional kvCacheId); + void clearStatsDirty(std::optional kvCacheId); + std::unordered_set getDirtyStatsKvCacheIds() const; + void markStatsExcluded(std::optional kvCacheId); + void clearStatsExcluded(std::optional kvCacheId); + bool isStatsExcluded(std::optional kvCacheId) const; + + // Mirrors Python's need_adjustment property and adjust() method. + // All KvCaches must be suspended before calling adjust(). + bool needAdjustment() const; + void adjust(); + + // ---- Internals used by KvCache ---------------------------------------- + + StorageManager& storage() noexcept + { + return *mStorage; + } + + KVCacheManagerConfig const& config() const noexcept + { + return mConfig; + } + + LifeCycleRegistry const& lifeCycles() const noexcept + { + return mLifeCycles; + } + + BlockRadixTree& radixTree() noexcept + { + return *mRadixTree; + } + + std::shared_ptr const& eventSink() const noexcept + { + return mEventSink; + } + + // Called by KvCache constructor/destructor. + void registerKvCache(KvCache* kvc); + void unregisterKvCache(KvCache* kvc); + + // Moving-average updates from closed KvCaches. + void updateAvgReusedLength(double v) + { + mAvgReusedLength.update(v); + } + + void updateAvgSqrCapacity(double v) + { + mAvgSqrCapacity.update(v); + } + + void updateAvgSqrHistoryLength(double v) + { + mAvgSqrHistoryLength.update(v); + } + + void incrementNumSampledKvCaches() + { + ++mNumSampledKvCaches; + } + + // Try to rebalance memory pool ratios based on usage statistics. + void tryUpdateTargetRatios(); + + // White-box introspection (incl. test-only auto-tuner state mutation) reaches + // private members directly rather than widening the public API. + friend class KvCacheIntrospection; + +private: + void _adjustLevel(CacheLevel level, size_t quota); + bool _needAdjustment(CacheLevel level) const; + TypedVec const& _getTargetRatioList(CacheLevel level) const; + TypedVec>> _gatherPersistentPages() const; + + PeakBlockStatsByCacheLevel _currentBlockStatsByCacheLevel() const; + void _resetIterationPeakNumBlocks(std::optional cacheLevel = std::nullopt); + void _updateIterationPeakNumBlocks(); + + // Current per-pool-group GPU utilization ratios. + TypedVec _currentGpuRatio() const; + TypedVec _currentOtherRatios() const; + + KVCacheManagerConfig mConfig; + LifeCycleRegistry mLifeCycles; + std::shared_ptr mEventSink; + std::shared_ptr mRadixTree; + std::shared_ptr mStorage; + + // Weak references to all living KvCaches. + std::set mLivingKvCaches; + + // Moving averages used for ratio rebalancing. + MovingAverage mAvgReusedLength; + MovingAverage mAvgSqrCapacity; + MovingAverage mAvgSqrHistoryLength; + + TypedVec mTargetRatioListGpu; + TypedVec mTargetRatioListOther; + + int mNumCreatedKvCaches{0}; + int mNumSampledKvCaches{0}; + double mLastAdjustmentTime{0.0}; + int mLastUpdateNumSampledKvCaches{0}; + + KVCacheStatsDelta mCommittedStats; + IterationStatsByLifeCycle mIterationStatsByLifeCycle; + SsmSnapshotIterationStatsByLifeCycle mSsmSnapshotIterationStatsByLifeCycle; + PeakBlockStatsByCacheLevel mIterationPeakNumBlocksByCacheLevel; + std::unordered_set mDirtyStatsKvCacheIds; + std::unordered_set mStatsExcludedKvCacheIds; +}; + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp new file mode 100644 index 000000000000..c26e77f61b2f --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp @@ -0,0 +1,99 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "kv_cache_manager_v2/lifeCycleRegistry.h" + +#include "tensorrt_llm/common/assert.h" +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// --------------------------------------------------------------------------- +// makeLifeCycle — factory dispatching on LayerConfig variant. +// --------------------------------------------------------------------------- + +LifeCycle makeLifeCycle(LayerConfig const& layer, int tokensPerBlock) +{ + return std::visit( + [&](auto const& cfg) -> LifeCycle + { + cfg.validate(); + using T = std::decay_t; + if constexpr (std::is_same_v) + return SsmLifeCycle{}; + else + return AttnLifeCycle::make(cfg.slidingWindowSize, cfg.numSinkTokens, tokensPerBlock); + }, + layer); +} + +// --------------------------------------------------------------------------- +// LifeCycleRegistry +// --------------------------------------------------------------------------- + +LifeCycleRegistry::LifeCycleRegistry(KVCacheManagerConfig const& config) +{ + for (auto const& layer : config.layers) + { + LifeCycle lc = makeLifeCycle(layer, config.tokensPerBlock); + if (mLifeCycleIdMap.find(lc) == mLifeCycleIdMap.end()) + { + check(); + LifeCycleId id = mLifeCycleList.size(); + mLifeCycleList.push_back(lc); + mLifeCycleIdMap[lc] = id; + if (std::holds_alternative(lc)) + mSsmLifeCycleId = id; + } + } + check(); +} + +LifeCycle const& LifeCycleRegistry::operator[](LifeCycleId id) const +{ + return mLifeCycleList.at(id); +} + +LifeCycle const& LifeCycleRegistry::getLifeCycle(LifeCycleId id) const +{ + return (*this)[id]; +} + +LifeCycleId LifeCycleRegistry::getId(LifeCycle const& lc) const +{ + auto it = mLifeCycleIdMap.find(lc); + if (it == mLifeCycleIdMap.end()) + { + throw std::out_of_range("LifeCycleRegistry::getId: life cycle not found"); + } + return it->second; +} + +LifeCycleId LifeCycleRegistry::size() const noexcept +{ + check(); + return mLifeCycleList.size(); +} + +inline void LifeCycleRegistry::check() const +{ + TLLM_CHECK_DEBUG_WITH_INFO( + toSizeT(mLifeCycleList.size()) == mLifeCycleIdMap.size(), "corrupted life cycle registry"); +} + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h new file mode 100644 index 000000000000..f655fa7ce092 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h @@ -0,0 +1,251 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "kv_cache_manager_v2/common.h" +#include "kv_cache_manager_v2/config.h" +#include "kv_cache_manager_v2/utils/math.h" +#include "tensorrt_llm/common/assert.h" + +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// --------------------------------------------------------------------------- +// AttnLifeCycle — lifecycle for attention layers (SWA + sink blocks). +// Mirrors _life_cycle_registry.py::AttnLifeCycle. +// --------------------------------------------------------------------------- +struct AttnLifeCycle +{ + std::optional windowSize; // nullopt = no sliding window + int numSinkBlocks = 0; // divUp(numSinkTokens, tokensPerBlock) + + HalfOpenRange getStaleRange(int historyLength, int tokensPerBlock) const + { + int numBlocks = divUp(historyLength, tokensPerBlock); + BlockOrdinal start{std::min(numBlocks, numSinkBlocks)}; + if (!windowSize.has_value()) + return {start, start}; + // `+ 1` is intentional: attention always runs for >= 1 in-flight input + // token at position `historyLength`, so the live window is + // [historyLength + 1 - windowSize, historyLength]. Do not drop it. + BlockOrdinal windowStart{(historyLength + 1 - *windowSize) / tokensPerBlock}; + return {start, std::max(start, windowStart)}; + } + + bool operator==(AttnLifeCycle const& o) const noexcept + { + return windowSize == o.windowSize && numSinkBlocks == o.numSinkBlocks; + } + + bool operator<(AttnLifeCycle const& o) const noexcept + { + if (windowSize != o.windowSize) + return windowSize < o.windowSize; + return numSinkBlocks < o.numSinkBlocks; + } + + static AttnLifeCycle make(std::optional ws, std::optional numSinkTokens, int tokensPerBlock) + { + TLLM_CHECK_DEBUG(tokensPerBlock > 0); + TLLM_CHECK_DEBUG(!ws.has_value() || *ws > 0); + TLLM_CHECK_DEBUG(!numSinkTokens.has_value() || *numSinkTokens >= 0); + TLLM_CHECK_DEBUG((!numSinkTokens.has_value() || *numSinkTokens == 0) || ws.has_value()); + int sinkBlocks = divUp(numSinkTokens.value_or(0), tokensPerBlock); + return AttnLifeCycle{ws, sinkBlocks}; + } +}; + +// --------------------------------------------------------------------------- +// SsmLifeCycle — lifecycle for SSM (State Space Model) layers. +// All blocks before the last full block are stale (recurrent state). +// Singleton: all SSM layers share the same lifecycle. +// --------------------------------------------------------------------------- +struct SsmLifeCycle +{ + HalfOpenRange getStaleRange(int historyLength, int tokensPerBlock) const + { + return {BlockOrdinal{0}, BlockOrdinal{historyLength / tokensPerBlock}}; + } + + bool operator==(SsmLifeCycle const&) const noexcept + { + return true; + } + + bool operator<(SsmLifeCycle const&) const noexcept + { + return false; + } +}; + +// --------------------------------------------------------------------------- +// LifeCycle — variant of attention or SSM lifecycle. +// --------------------------------------------------------------------------- +using LifeCycle = std::variant; + +// Free function: compute stale range via std::visit. +inline HalfOpenRange getStaleRange(LifeCycle const& lc, int historyLength, int tokensPerBlock) +{ + return std::visit([&](auto const& v) { return v.getStaleRange(historyLength, tokensPerBlock); }, lc); +} + +// Compute the range of blocks that should use scratch (shared) slots during SWA prefill. +// Scratch = stale_at_capacity ∩ input_blocks, where: +// stale_at_capacity: blocks out-of-window when all capacity tokens become history. +// input_blocks: [divUp(historyLength, tpb), divUp(capacity, tpb)) — new blocks +// for the current chunk. +// Mirrors _life_cycle_registry.py::compute_scratch_range(). +inline HalfOpenRange computeScratchRange( + LifeCycle const& lc, int historyLength, int capacity, int tokensPerBlock, int maxRewindLen) +{ + auto const* attn = std::get_if(&lc); + if (!attn || !attn->windowSize.has_value()) + { + return {BlockOrdinal{0}, BlockOrdinal{0}}; + } + int const nonRewindableCapacity = std::max(0, capacity - maxRewindLen); + auto capStale = attn->getStaleRange(nonRewindableCapacity, tokensPerBlock); + HalfOpenRange inputRange{divUp(historyLength, tokensPerBlock), divUp(capacity, tokensPerBlock)}; + return intersect(capStale, inputRange); +} + +// Factory: create a LifeCycle from a LayerConfig variant. +LifeCycle makeLifeCycle(LayerConfig const& layer, int tokensPerBlock); + +// Integer id assigned to each unique LifeCycle. +using LifeCycleId = StrongIndex; + +// Alias for public exposure (same meaning as LifeCycleId). +using LayerGroupId = LifeCycleId; + +// --------------------------------------------------------------------------- +// LifeCycleRegistry — deduplicates LifeCycle objects and assigns integer ids. +// Mirrors _life_cycle_registry.py::LifeCycleRegistry. +// --------------------------------------------------------------------------- +class LifeCycleRegistry +{ +public: + explicit LifeCycleRegistry(KVCacheManagerConfig const& config); + + // Look up a LifeCycle by id. + LifeCycle const& operator[](LifeCycleId id) const; + LifeCycle const& getLifeCycle(LifeCycleId id) const; + + // Look up the id for a LifeCycle (throws if not found). + LifeCycleId getId(LifeCycle const& lc) const; + + // Number of unique life cycles. + LifeCycleId size() const noexcept; + + // Iteration over LifeCycles in registration order. + TypedVec const& getAll() const noexcept + { + return mLifeCycleList; + } + + bool contains(LifeCycle const& lc) const noexcept + { + return mLifeCycleIdMap.count(lc) > 0; + } + + // SSM helpers. + std::optional ssmLifeCycleId() const noexcept + { + return mSsmLifeCycleId; + } + + bool hasSSM() const noexcept + { + return mSsmLifeCycleId.has_value(); + } + + // Return (id, AttnLifeCycle*) pairs for attention lifecycles only. Used by _setupForReuse. + std::vector> attentionLifeCycles() const + { + std::vector> result; + for (LifeCycleId lcId{0}; lcId < mLifeCycleList.size(); ++lcId) + { + if (auto const* attn = std::get_if(&mLifeCycleList[lcId])) + result.emplace_back(lcId, attn); + } + return result; + } + + // Iterate: (id, lifecycle) pairs — all entries. + struct Item + { + LifeCycleId id; + LifeCycle const& lc; + }; + + class ItemIterator + { + public: + ItemIterator(TypedVec const& list, LifeCycleId pos) + : mList(&list) + , mPos(pos) + { + } + + Item operator*() const + { + return {mPos, (*mList)[mPos]}; + } + + ItemIterator& operator++() + { + ++mPos; + return *this; + } + + bool operator!=(ItemIterator const& o) const + { + return mPos != o.mPos; + } + + private: + TypedVec const* mList; + LifeCycleId mPos; + }; + + ItemIterator begin() const + { + return {mLifeCycleList, LifeCycleId{0}}; + } + + ItemIterator end() const + { + return {mLifeCycleList, mLifeCycleList.size()}; + } + +private: + void check() const; + + TypedVec mLifeCycleList; + std::map mLifeCycleIdMap; + std::optional mSsmLifeCycleId; +}; + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/movingAverage.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/movingAverage.h new file mode 100644 index 000000000000..c2f53af570ff --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/movingAverage.h @@ -0,0 +1,99 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/assert.h" + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// --------------------------------------------------------------------------- +// Exponential moving average (mirrors _core/_moving_average.py::MovingAverage). +// weight-corrected bias: avg += (value - avg) / weight, weight = 1 + decay*weight +// --------------------------------------------------------------------------- +class MovingAverage +{ +public: + explicit MovingAverage(double decay = 0.9999) noexcept + : mDecay(decay) + , mAvg(0.0) + , mWeight(0.0) + , mNumUpdates(0) + { + } + + double update(double value) noexcept + { + mWeight = 1.0 + mDecay * mWeight; + mAvg += (value - mAvg) / mWeight; + ++mNumUpdates; + return mAvg; + } + + double value() const noexcept + { + return mAvg; + } + + int numUpdates() const noexcept + { + return mNumUpdates; + } + +private: + double mDecay; + double mAvg; + double mWeight; + int mNumUpdates; +}; + +// --------------------------------------------------------------------------- +// Simple arithmetic mean (mirrors _core/_moving_average.py::Average). +// --------------------------------------------------------------------------- +class Average +{ +public: + Average() noexcept + : mSum(0.0) + , mCount(0) + { + } + + void update(double value) noexcept + { + mSum += value; + ++mCount; + } + + double value() const noexcept + { + TLLM_CHECK_DEBUG(mCount > 0); + return mSum / mCount; + } + + int count() const noexcept + { + return mCount; + } + +private: + double mSum; + int mCount; +}; + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp new file mode 100644 index 000000000000..91f9921df948 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp @@ -0,0 +1,570 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "kv_cache_manager_v2/page.h" +#include "kv_cache_manager_v2/common.h" +#include "kv_cache_manager_v2/exceptions.h" +#include "kv_cache_manager_v2/kvCache.h" // for KvCache +#include "kv_cache_manager_v2/storageManager.h" // for StorageManager + +#include "tensorrt_llm/common/assert.h" +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// --------------------------------------------------------------------------- +// Page +// --------------------------------------------------------------------------- + +Page::Page(StorageManager* mgr, LifeCycleId lc, CacheLevel level, Priority prio) + : manager(mgr) + , lifeCycle(lc) + , cacheLevel(level) + , priority(prio) + , nodeRef(std::nullopt) +{ +} + +Page::~Page() +{ + TLLM_CHECK_DEBUG_WITH_INFO(status() == PageStatus::DROPPABLE && !scheduledForEviction(), + "Page destroyed while still held or scheduled for eviction"); + if (hasValidSlot()) + { + Slot s; + s.setSlotId(slotId()); + s.readyEvent = std::move(readyEvent); + resetSlot(); + manager->releaseSlot(lifeCycle, cacheLevel, std::move(s)); + } +} + +PageStatus Page::status() const noexcept +{ + auto h = holder.lock(); + if (!h) + return PageStatus::DROPPABLE; + if (h->uniqLock.expired()) + return PageStatus::HELD; + return PageStatus::LOCKED; +} + +SharedPtr Page::hold() +{ + // Return existing holder if any. + auto h = holder.lock(); + if (h) + return h; + + auto self = sharedFromThis(); + h = makeShared(self); + holder = h; + + // If we were scheduled for eviction but are no longer evictable (just got held), remove. + if (scheduledForEviction()) + { + if (!manager->isEvictable(*this)) + { + manager->excludeFromEviction(*this); + TLLM_CHECK_DEBUG(!scheduledForEviction()); + } + } + return h; +} + +SharedPageLock Page::lock(KvCache& kvCache, BeamIndex beamIndex, BlockOrdinal ordinal, LifeCycleId lc, bool skipWait) +{ + return hold()->lock(kvCache, beamIndex, ordinal, lc, skipWait); +} + +// --------------------------------------------------------------------------- +// CommittedPage +// --------------------------------------------------------------------------- + +CommittedPage::CommittedPage(StorageManager* mgr, SharedPtr blk, LifeCycleId lc, CacheLevel level, Priority prio) + : Page(mgr, lc, level, prio) + , block(blk.get()) +{ +} + +SsmCommittedPage::SsmCommittedPage( + StorageManager* mgr, SharedPtr blk, LifeCycleId lc, CacheLevel level, Priority prio, int numTokensInBlock_) + : CommittedPage(mgr, std::move(blk), lc, level, prio) + , numTokensInBlock(numTokensInBlock_) +{ + TLLM_CHECK_DEBUG(numTokensInBlock_ > 0); +} + +CommittedPage::~CommittedPage() +{ + if (block != nullptr) + { + // unlinkPage nulls storage[lc]->block (i.e. our member), so capture + // the block pointer first. Pass `this` as the expected page: if a newer + // page already replaced us in the slot (e.g. a larger SSM snapshot), the + // slot is left alone and prev is nullptr — skip stale-block cleanup then. + Block* blk = block; + auto* prev = blk->unlinkPage(lifeCycle, this); + if (prev != nullptr) + { + TLLM_CHECK_DEBUG_WITH_INFO(prev == this, "unlinkPage returned unexpected page"); + LifeCycle const& lc = manager->lifeCycles().getLifeCycle(lifeCycle); + Block::clearStaleBlocksAfterPageUnlink(*blk, lifeCycle, lc); + } + } + // Delegate slot release to Page::~Page(). +} + +// --------------------------------------------------------------------------- +// UncommittedPage +// --------------------------------------------------------------------------- + +UncommittedPage::UncommittedPage(KvCache& kvc, BlockOrdinal ord, LifeCycleId lc, CacheLevel level, BeamIndex bi) + : Page(kvc.storageManager(), lc, level, kvc.getPriority(ord, lc)) + , kvCache(&kvc) + , ordinal(ord) + , beamIndex(bi) +{ +} + +UncommittedPage::~UncommittedPage() +{ + // Mirrors Python UncommittedPage.__del__: for attention LCs, the page must be either: + // - part of an SSM lifecycle (different rules), + // - at an ordinal beyond the current block list (block already removed), + // - the slot at this position is null, CommittedPage, or this page itself (self-destruction). + // The "p == this" condition is C++-specific: std::variant destroys the old value before + // switching to monostate, so during destruction the slot still references this page. + if (TLLM_UNLIKELY(gDebug)) + { + auto ssmLcId = manager->lifeCycles().ssmLifeCycleId(); + bool isSsm = ssmLcId.has_value() && lifeCycle == *ssmLcId; + if (!isSsm) + { + [[maybe_unused]] bool blockRemoved = kvCache->blocks().size() <= ordinal; + [[maybe_unused]] bool pageOk = true; + if (!blockRemoved) + { + auto const& bp = kvCache->blocks()[ordinal].pages[beamIndex][lifeCycle]; + auto page = blockPageGetPage(bp); + pageOk + = blockPageIsNull(bp) || page.get() == this || dynamicPointerCast(page) != nullptr; + } + TLLM_CHECK_WITH_INFO( + blockRemoved || pageOk, "UncommittedPage destroyed but slot still holds a different uncommitted page"); + } + } + // Delegate slot release to Page::~Page(). +} + +SharedPtr UncommittedPage::convertToCommitted(SharedPtr blk, CachedCudaEvent readyEv) +{ + TLLM_CHECK_DEBUG(!scheduledForEviction()); + TLLM_CHECK_DEBUG_WITH_INFO( + blk->storage.at(lifeCycle) == nullptr, "Block slot for this lifecycle already has a committed page"); + TLLM_CHECK_DEBUG_WITH_INFO(status() == PageStatus::DROPPABLE, "Release holder/lock before converting"); + + // Set the ready event before transfer (matches Python: self.ready_event = ready_event). + this->readyEvent = std::move(readyEv); + + auto committed = makeShared(manager, blk, lifeCycle, cacheLevel, priority); + // Move slot id to the committed page; invalidate our slot. + committed->setSlotId(slotId()); // asserts valid + committed->readyEvent = std::move(readyEvent); + resetSlot(); + readyEvent = CachedCudaEvent::makeNull(); + + TLLM_CHECK_DEBUG(!hasValidSlot() && readyEvent.isClosed()); + TLLM_CHECK_DEBUG_WITH_INFO(committed->hasValidSlot(), "committed page must have a valid slot after transfer"); + + // Register in block storage. + blk->storage.at(lifeCycle) = committed.get(); + + return committed; +} + +SharedPtr UncommittedPage::convertToSsmCommitted( + SharedPtr blk, CachedCudaEvent readyEv, int numTokensInBlock) +{ + TLLM_CHECK_DEBUG(!scheduledForEviction()); + TLLM_CHECK_DEBUG_WITH_INFO( + blk->storage.at(lifeCycle) == nullptr, "Block slot for this lifecycle already has a committed page"); + TLLM_CHECK_DEBUG_WITH_INFO(status() == PageStatus::DROPPABLE, "Release holder/lock before converting"); + + this->readyEvent = std::move(readyEv); + + auto committed = makeShared(manager, blk, lifeCycle, cacheLevel, priority, numTokensInBlock); + committed->setSlotId(slotId()); // asserts valid + committed->readyEvent = std::move(readyEvent); + resetSlot(); + readyEvent = CachedCudaEvent::makeNull(); + + TLLM_CHECK_DEBUG(!hasValidSlot() && readyEvent.isClosed()); + TLLM_CHECK_DEBUG_WITH_INFO(committed->hasValidSlot(), "committed page must have a valid slot after transfer"); + + blk->storage.at(lifeCycle) = committed.get(); + + return committed; +} + +// --------------------------------------------------------------------------- +// PageHolder +// --------------------------------------------------------------------------- + +PageHolder::PageHolder(SharedPtr p) + : page(std::move(p)) +{ +} + +PageHolder::~PageHolder() +{ + TLLM_CHECK_DEBUG_WITH_INFO(uniqLock.expired(), "PageHolder destroyed while lock still active"); + + page->holder.reset(); // clear back-reference + auto const manager = page->manager; + + // If it's a committed page, schedule for eviction (if evictable). + if (page->isCommitted()) + { + if (!page->scheduledForEviction()) + manager->scheduleForEviction(*page); + + // If the block is orphan, exclude from eviction immediately. + auto* cp = dynamic_cast(page.get()); + if (cp) + { + if (cp->block == nullptr || cp->block->isOrphan()) + manager->excludeFromEviction(*page); + } + } + else + { + // Uncommitted page: if scheduled for eviction, remove it. + if (page->scheduledForEviction()) + manager->excludeFromEviction(*page); + } +} + +SharedPageLock PageHolder::lock( + KvCache& kvCache, BeamIndex beamIndex, BlockOrdinal ordinal, LifeCycleId lc, bool skipWait) +{ + // Create or reuse UniqPageLock. + auto ul = uniqLock.lock(); + if (!ul) + { + ul = makeShared(sharedFromThis()); + uniqLock = ul; + } + + // Remove from eviction queue if scheduled. + if (page->scheduledForEviction()) + { + page->manager->excludeFromEviction(*page); + TLLM_CHECK_DEBUG(!page->scheduledForEviction()); + } + + return ul->share(kvCache, beamIndex, ordinal, lc, skipWait); +} + +// --------------------------------------------------------------------------- +// UniqPageLock +// --------------------------------------------------------------------------- + +UniqPageLock::UniqPageLock(SharedPtr h) + : holder(std::move(h)) +{ + if (holder->page->cacheLevel != kGpuLevel) + throw LogicError("Lock can only be applied to GPU-memory pages"); +} + +UniqPageLock::~UniqPageLock() +{ + Page& p = *page(); + TLLM_CHECK_DEBUG(p.cacheLevel == kGpuLevel && !p.scheduledForEviction()); + // Set readyEvent to the merged finish events of all readers. For committed (read-only) + // pages, this means the next reader will wait for prior reads to complete, which is + // unnecessary but correct. See the CommittedPage comment in page.h for rationale. + p.readyEvent = mergeEvents(finishEvents); + + // Clear the holder's lock reference. + TLLM_CHECK_DEBUG(holder); + holder->uniqLock.reset(); + + // Optimized path (mirrors Python): set holder=nullptr, then check if still evictable. + auto holderCopy = std::move(holder); + holder = nullptr; + + // If the page is not droppable (still held by someone else) and evictable, + // schedule for eviction. + if (p.status() != PageStatus::DROPPABLE) + { + auto const manager = p.manager; + if (manager->isEvictable(p)) + manager->scheduleForEviction(p); + } +} + +void UniqPageLock::notifyFinish(CachedCudaEvent event) +{ + finishEvents.push_back(std::move(event)); + // Avoid unbounded growth for system prompt pages shared by all requests. + if (finishEvents.size() > 32) + { + CachedCudaEvent merged = mergeEvents(finishEvents); + finishEvents.clear(); + finishEvents.push_back(std::move(merged)); + } +} + +SharedPtr const& UniqPageLock::page() const +{ + TLLM_CHECK_DEBUG(holder && holder->page); + return holder->page; +} + +SharedPageLock UniqPageLock::share( + KvCache& kvCache, BeamIndex beamIndex, BlockOrdinal ordinal, LifeCycleId lc, bool skipWait) +{ + return SharedPageLock(sharedFromThis(), kvCache, beamIndex, ordinal, lc, skipWait); +} + +// --------------------------------------------------------------------------- +// SharedPageLock +// --------------------------------------------------------------------------- + +SharedPageLock::SharedPageLock(SharedPtr ul, KvCache& kvCache, BeamIndex beamIndex, BlockOrdinal ordinal, + LifeCycleId lc, bool skipWait) + : mUniqLock(std::move(ul)) + , mUser{&kvCache, beamIndex, ordinal, lc} +{ + if (!skipWait) + page()->readyEvent.waitInStream(reinterpret_cast(kvCache.cudaStream())); + + acquirePageIndex(); +} + +SharedPageLock::~SharedPageLock() +{ + if (mUniqLock) + unlock(); +} + +SharedPageLock::SharedPageLock(SharedPageLock&& other) noexcept + : mUniqLock(std::move(other.mUniqLock)) + , mUser(std::move(other.mUser)) +{ +} + +SharedPageLock& SharedPageLock::operator=(SharedPageLock&& other) noexcept +{ + if (this != &other) + { + if (mUniqLock) + unlock(); + mUniqLock = std::move(other.mUniqLock); + mUser = std::move(other.mUser); + } + return *this; +} + +SharedPtr const& SharedPageLock::page() const +{ + TLLM_CHECK_DEBUG(mUniqLock); + return mUniqLock->page(); +} + +SharedPtr SharedPageLock::unlock() +{ + TLLM_CHECK_DEBUG(mUniqLock); + + // Record finish event from the KvCache stream. + mUniqLock->notifyFinish(mUser.kvCache->finishEvent()); + + releasePageIndex(); + auto p = page(); // copy shared_ptr before reset + mUniqLock.reset(); + return p; +} + +void SharedPageLock::acquirePageIndex() +{ + auto* kvc = mUser.kvCache; + auto& pg = *page(); + int old = kvc->updateBasePageIndex( + mUser.beamIndex, mUser.ordinal, mUser.lifeCycle, slotIdToPageIndexValue(pg.slotId())); + // Mirrors Python assertion: old base index must be BAD (prevents double-locking same slot). + TLLM_CHECK_DEBUG_WITH_INFO( + old == kBadPageIndex.value(), "Double-lock: page index already acquired for this (beam, ordinal, lc)"); + (void) old; +} + +void SharedPageLock::releasePageIndex() +{ + int oldBaseIndex + = mUser.kvCache->updateBasePageIndex(mUser.beamIndex, mUser.ordinal, mUser.lifeCycle, kBadPageIndex.value()); + // Mirrors Python assertion: old base index must match this page's slot ID. + TLLM_CHECK_DEBUG(oldBaseIndex == slotIdToPageIndexValue(page()->slotId())); + (void) oldBaseIndex; +} + +// --------------------------------------------------------------------------- +// batchedLockToGpu +// --------------------------------------------------------------------------- + +std::vector batchedLockToGpu(KvCache& kvCache, std::vector const& targets) +{ + auto* storeMgr = kvCache.storageManager(); + TLLM_CHECK_DEBUG(storeMgr); + // All pages must belong to the same storage manager. + TLLM_CHECK_DEBUG(targets.empty() + || std::all_of(targets.begin(), targets.end(), [&](auto const& t) { return t.page->manager == storeMgr; })); + + // Determine how many GPU slots are needed per pool group. + TypedVec requirements(storeMgr->numPoolGroups(), 0); + std::vector wasScheduled(targets.size(), false); + + for (size_t i = 0; i < targets.size(); ++i) + { + auto const& t = targets[i]; + wasScheduled[i] = t.page->scheduledForEviction(); + if (wasScheduled[i]) + storeMgr->excludeFromEviction(*t.page); + if (t.page->cacheLevel != kGpuLevel) + { + PoolGroupIndex pgIdx = storeMgr->getPoolGroupIndex(t.lifeCycle); + requirements[pgIdx] += 1; + } + } + + try + { + MigrationRecorder const migrationRecorder + = [&kvCache](std::vector> const& pages, std::vector const& slots, CacheLevel srcLevel, + CacheLevel dstLevel) { kvCache._recordMigratedSlots(pages, slots, srcLevel, dstLevel); }; + DropRecorder const dropRecorder = [&kvCache](std::vector> const& pages, CacheLevel cacheLevel) + { kvCache._recordDroppedPages(pages, cacheLevel); }; + storeMgr->prepareFreeSlots(kGpuLevel, requirements, migrationRecorder, dropRecorder); + // Migrate non-GPU pages. + storeMgr->batchedMigrateToGpu(targets, kvCache, migrationRecorder); + } + catch (...) + { + // Restore eviction scheduling. + for (size_t i = 0; i < targets.size(); ++i) + if (wasScheduled[i]) + storeMgr->scheduleForEviction(*targets[i].page); + throw; + } + + // Wait for all ready events on KvCache's stream (deduplicated). + { + std::vector readyEvents; + readyEvents.reserve(targets.size()); + for (auto const& t : targets) + readyEvents.push_back(&t.page->readyEvent); + streamWaitEvents(reinterpret_cast(kvCache.cudaStream()), readyEvents); + } + + // Lock all pages. + std::vector locks; + locks.reserve(targets.size()); + for (auto const& t : targets) + locks.emplace_back(t.page->lock(kvCache, t.beamIndex, t.ordinal, t.lifeCycle, + /*skipWait=*/true)); + return locks; +} + +// --------------------------------------------------------------------------- +// ScratchSlotLock +// --------------------------------------------------------------------------- + +ScratchSlotLock::ScratchSlotLock(Slot slot, KvCache& owner, LifeCycleId lifeCycle, bool skipWait) + : mOwner(&owner) + , mLifeCycle(lifeCycle) +{ + if (!skipWait) + { + slot.readyEvent.waitInStream(reinterpret_cast(owner.cudaStream())); + } + mSlot.setSlot(slot); +} + +ScratchSlotLock::~ScratchSlotLock() +{ + if (mSlot.hasValidSlot()) + { + try + { + unlock(); + } + catch (...) + { + } + } +} + +ScratchSlotLock::ScratchSlotLock(ScratchSlotLock&& other) noexcept + : mSlot(std::move(other.mSlot)) + , mOwner(other.mOwner) + , mLifeCycle(other.mLifeCycle) +{ + // Invalidate moved-from: mSlot move only transfers readyEvent, not slotId (trivially-copyable). + other.mSlot.resetSlot(); + other.mOwner = nullptr; +} + +ScratchSlotLock& ScratchSlotLock::operator=(ScratchSlotLock&& other) noexcept +{ + if (this != &other) + { + if (mSlot.hasValidSlot()) + { + try + { + unlock(); + } + catch (...) + { + } + } + mSlot = std::move(other.mSlot); + other.mSlot.resetSlot(); // Invalidate moved-from slotId. + mOwner = other.mOwner; + mLifeCycle = other.mLifeCycle; + other.mOwner = nullptr; + } + return *this; +} + +Slot ScratchSlotLock::detachSlot() +{ + TLLM_CHECK_DEBUG(mSlot.hasValidSlot()); + Slot result; + result.setSlot(mSlot); + return result; +} + +void ScratchSlotLock::unlock() +{ + TLLM_CHECK_DEBUG(mSlot.hasValidSlot()); + mSlot.readyEvent = mOwner->finishEvent(); + mOwner->storageManager()->releaseSlot(mLifeCycle, kGpuLevel, std::move(mSlot)); + TLLM_CHECK_DEBUG(!mSlot.hasValidSlot()); +} + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h new file mode 100644 index 000000000000..9f52dca1f0ee --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h @@ -0,0 +1,307 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "kv_cache_manager_v2/blockRadixTree.h" +#include "kv_cache_manager_v2/common.h" +#include "kv_cache_manager_v2/evictionController.h" +#include "kv_cache_manager_v2/lifeCycleRegistry.h" +#include "kv_cache_manager_v2/storage/core.h" +#include "kv_cache_manager_v2/utils/cudaEvent.h" +#include "kv_cache_manager_v2/utils/sharedPtr.h" + +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// Forward declarations to break circular includes. +class StorageManager; +class KvCache; +class PageHolder; +class UniqPageLock; +class SharedPageLock; + +// --------------------------------------------------------------------------- +// Page — base class for all KV-cache pages. +// Inherits from Slot (holds slotId + readyEvent). +// Mirrors Python's Page(Slot) dataclass. +// --------------------------------------------------------------------------- +class Page : public Slot, public EnableSharedFromThis +{ +public: + StorageManager* manager; + LifeCycleId lifeCycle; + CacheLevel cacheLevel; + Priority priority; + WeakPtr holder; // empty → DROPPABLE + std::optional nodeRef; // present → scheduled for eviction + + Page(StorageManager* mgr, LifeCycleId lc, CacheLevel level, Priority prio); + + virtual ~Page(); + + virtual bool isCommitted() const = 0; + + PageStatus status() const noexcept; + + bool scheduledForEviction() const noexcept + { + return nodeRef.has_value(); + } + + // Prevent the page from being dropped (returns/creates a PageHolder). + SharedPtr hold(); + + // Acquire a shared lock (migrates to GPU if needed). + // skip_wait: caller guarantees the page is ready on kvCache's stream. + SharedPageLock lock( + KvCache& kvCache, BeamIndex beamIndex, BlockOrdinal ordinal, LifeCycleId lifeCycle, bool skipWait = false); +}; + +// --------------------------------------------------------------------------- +// CommittedPage — page associated with a Block in the radix tree. +// +// A committed page is immutable — all access after commit is read-only. +// +// We intentionally do not add a separate read event to track read completion. +// The inherited Slot::readyEvent serves double duty: after commit or migration +// it represents write completion; after UniqPageLock is destroyed it is set to +// the merged finish events of all prior readers. This means a new reader may +// unnecessarily wait for a prior reader (read-after-read on immutable data), +// but this is functionally correct, only occurs when the lock is fully released +// between reuses, and saves one event field per committed page — a worthwhile +// tradeoff given the potentially huge number of committed pages in the system. +// --------------------------------------------------------------------------- +class CommittedPage : public Page +{ +public: + Block* block; + + // Number of outstanding PlannedDropHandles that intend to drop this page. + // Mirrors Python's CommittedPage.planned_drop_count. + int plannedDropCount{0}; + + CommittedPage(StorageManager* mgr, SharedPtr blk, LifeCycleId lc, CacheLevel level, Priority prio); + + ~CommittedPage() override; + + bool isCommitted() const override + { + return true; + } +}; + +// --------------------------------------------------------------------------- +// SsmCommittedPage — a committed SSM snapshot page. +// +// Unlike attention CommittedPages (which always cover a full block), an SSM +// snapshot may cover only a prefix of its block. `numTokensInBlock` records how +// many tokens of the block this snapshot is reusable for. +// --------------------------------------------------------------------------- +class SsmCommittedPage : public CommittedPage +{ +public: + int numTokensInBlock; + + SsmCommittedPage(StorageManager* mgr, SharedPtr blk, LifeCycleId lc, CacheLevel level, Priority prio, + int numTokensInBlock); +}; + +// --------------------------------------------------------------------------- +// UncommittedPage — page associated with a live KvCache sequence. +// --------------------------------------------------------------------------- +class UncommittedPage : public Page +{ +public: + KvCache* kvCache; + BlockOrdinal ordinal; + BeamIndex beamIndex; + std::vector tokens; + + UncommittedPage(KvCache& kvc, BlockOrdinal ord, LifeCycleId lc, CacheLevel level, BeamIndex bi = kDefaultBeamIndex); + + ~UncommittedPage() override; + + bool isCommitted() const override + { + return false; + } + + // Convert this UncommittedPage into a CommittedPage and attach to `block`. + // The UncommittedPage becomes invalid (slot transferred to CommittedPage). + SharedPtr convertToCommitted(SharedPtr block, CachedCudaEvent readyEvent); + + // Convert this UncommittedPage into an SsmCommittedPage covering + // `numTokensInBlock` tokens and attach to `block`. Invalidates this page. + SharedPtr convertToSsmCommitted( + SharedPtr block, CachedCudaEvent readyEvent, int numTokensInBlock); +}; + +// --------------------------------------------------------------------------- +// PageHolder — prevents a page from being dropped (HELD status). +// Mirrors Python's _PageHolder. +// --------------------------------------------------------------------------- +class PageHolder : public EnableSharedFromThis +{ +public: + explicit PageHolder(SharedPtr page); + ~PageHolder(); + + PageHolder(PageHolder const&) = delete; + PageHolder& operator=(PageHolder const&) = delete; + + // Acquire a shared lock (creates or reuses the UniqPageLock). + SharedPageLock lock( + KvCache& kvCache, BeamIndex beamIndex, BlockOrdinal ordinal, LifeCycleId lifeCycle, bool skipWait = false); + + SharedPtr page; + WeakPtr uniqLock; // non-null → LOCKED +}; + +// --------------------------------------------------------------------------- +// UniqPageLock — locks a page to prevent eviction (LOCKED status). +// Owns finish events from all SharedPageLocks it issued. +// Mirrors Python's _UniqPageLock. +// --------------------------------------------------------------------------- +class UniqPageLock : public EnableSharedFromThis +{ +public: + explicit UniqPageLock(SharedPtr holder); + ~UniqPageLock(); + + UniqPageLock(UniqPageLock const&) = delete; + UniqPageLock& operator=(UniqPageLock const&) = delete; + + // Issue a SharedPageLock to a specific (kvCache, beam, ordinal, lifecycle). + SharedPageLock share( + KvCache& kvCache, BeamIndex beamIndex, BlockOrdinal ordinal, LifeCycleId lifeCycle, bool skipWait); + + SharedPtr const& page() const; + + // Append a finish event, merging when count exceeds 32 to prevent unbounded growth. + void notifyFinish(CachedCudaEvent event); + + SharedPtr holder; + std::vector finishEvents; +}; + +// --------------------------------------------------------------------------- +// LockOwner — identifies who holds a SharedPageLock. +// --------------------------------------------------------------------------- +struct LockOwner +{ + KvCache* kvCache; + BeamIndex beamIndex; + BlockOrdinal ordinal; + LifeCycleId lifeCycle; +}; + +// --------------------------------------------------------------------------- +// SharedPageLock — one user's hold on an active page lock. +// Mirrors Python's _SharedPageLock. +// --------------------------------------------------------------------------- +class SharedPageLock +{ +public: + SharedPageLock(SharedPtr uniqLock, KvCache& kvCache, BeamIndex beamIndex, BlockOrdinal ordinal, + LifeCycleId lifeCycle, bool skipWait); + + ~SharedPageLock(); + + SharedPageLock(SharedPageLock&&) noexcept; + SharedPageLock& operator=(SharedPageLock&&) noexcept; + + SharedPageLock(SharedPageLock const&) = delete; + SharedPageLock& operator=(SharedPageLock const&) = delete; + + // Explicitly release the lock (called by destructor if not already released). + SharedPtr unlock(); + + SharedPtr const& page() const; + + bool isValid() const noexcept + { + return mUniqLock != nullptr; + } + +private: + // Internal helpers that update KvCache page index tables. + void acquirePageIndex(); + void releasePageIndex(); + + SharedPtr mUniqLock; + LockOwner mUser; +}; + +// --------------------------------------------------------------------------- +// BatchedLockTarget — input for batched_lock_to_gpu. +// --------------------------------------------------------------------------- +struct BatchedLockTarget +{ + SharedPtr page; + BeamIndex beamIndex; + BlockOrdinal ordinal; + LifeCycleId lifeCycle; +}; + +// --------------------------------------------------------------------------- +// batchedLockToGpu — migrate pages to GPU then lock them. +// Returns one SharedPageLock per target. +// Mirrors Python's batched_lock_to_gpu(). +// --------------------------------------------------------------------------- +std::vector batchedLockToGpu(KvCache& kvCache, std::vector const& targets); + +// --------------------------------------------------------------------------- +// ScratchSlotLock — manages a scratch slot for SWA prefill memory reuse. +// Wraps a Slot with owner (KvCache) and lifecycle references. +// On destruction, releases the slot back to the StorageManager. +// Mirrors _page.py::ScratchSlotLock. +// --------------------------------------------------------------------------- +class ScratchSlotLock +{ +public: + ScratchSlotLock(Slot slot, KvCache& owner, LifeCycleId lifeCycle, bool skipWait = false); + ~ScratchSlotLock(); + + ScratchSlotLock(ScratchSlotLock&& other) noexcept; + ScratchSlotLock& operator=(ScratchSlotLock&& other) noexcept; + + ScratchSlotLock(ScratchSlotLock const&) = delete; + ScratchSlotLock& operator=(ScratchSlotLock const&) = delete; + + // Detach and return the slot (transfers ownership to caller). + Slot detachSlot(); + + // Release the slot back to storage manager. + void unlock(); + + Slot const& slot() const noexcept + { + return mSlot; + } + +private: + Slot mSlot; + KvCache* mOwner; + LifeCycleId mLifeCycle; +}; + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/pendingStats.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/pendingStats.h new file mode 100644 index 000000000000..3ebb785dfb5f --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/pendingStats.h @@ -0,0 +1,255 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "kv_cache_manager_v2/stats.h" + +#include "tensorrt_llm/common/assert.h" +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +struct PendingAllocationSegment +{ + LifeCycleId lifeCycle; + BlockOrdinal blockBegin; + BlockOrdinal blockEnd; + int beamWidth; + bool countAsMissed; + bool countAsGeneration; +}; + +struct PendingStatsDelta +{ + KVCacheStatsDelta globalStats; + KVCacheStatsDelta requestStats; + KVCacheIterationStatsDelta iterationStats; + std::optional lifeCycle; + + [[nodiscard]] bool empty() const noexcept + { + return globalStats.empty() && requestStats.empty() && iterationStats.empty(); + } +}; + +class PendingStats +{ +public: + [[nodiscard]] bool empty() const noexcept + { + return mRequestStats.empty() && mGlobalStats.empty() && mIterationStatsByLifeCycle.empty() + && mSsmSnapshotIterationStatsByLifeCycle.empty(); + } + + void clear() noexcept + { + mRequestStats.clear(); + mGlobalStats.clear(); + mIterationStatsByLifeCycle.clear(); + mSsmSnapshotIterationStatsByLifeCycle.clear(); + mAllocationSegments.clear(); + } + + bool recordAllocationRange(LifeCycleId lifeCycle, BlockOrdinal blockBegin, BlockOrdinal blockEnd, int beamWidth, + bool countAsMissed, bool countAsGeneration = false) + { + if (blockBegin >= blockEnd) + { + return false; + } + PendingAllocationSegment segment{lifeCycle, blockBegin, blockEnd, beamWidth, countAsMissed, countAsGeneration}; + if (!add(allocationDelta(segment, blockBegin, blockEnd))) + { + return false; + } + mAllocationSegments.push_back(segment); + return true; + } + + bool recordReuse(LifeCycleId lifeCycle, int fullReusedBlocks, int partialReusedBlocks) + { + int const reusedBlocks = fullReusedBlocks + partialReusedBlocks; + if (reusedBlocks == 0) + { + return false; + } + + PendingStatsDelta delta; + delta.globalStats.reusedBlocks = reusedBlocks; + delta.requestStats.reusedBlocks = reusedBlocks; + delta.iterationStats.iterReusedBlocks = reusedBlocks; + delta.iterationStats.iterFullReusedBlocks = fullReusedBlocks; + delta.iterationStats.iterPartialReusedBlocks = partialReusedBlocks; + delta.lifeCycle = lifeCycle; + return add(delta); + } + + // Record one SSM snapshot lookup for a lifecycle. Mirrors Python's + // _PendingStats.record_ssm_snapshot_lookup(). Returns false if lookupTokens == 0. + bool recordSsmSnapshotLookup(LifeCycleId lifeCycle, int lookupTokens, int reusedTokens, int tokensPerBlock) + { + if (lookupTokens == 0) + { + return false; + } + TLLM_CHECK_DEBUG(lookupTokens > 0); + TLLM_CHECK_DEBUG(0 <= reusedTokens && reusedTokens <= lookupTokens); + TLLM_CHECK_DEBUG(tokensPerBlock > 0); + + bool const isHit = reusedTokens > 0; + SsmSnapshotIterationStatsDelta delta; + delta.iterSnapshotLookups = 1; + delta.iterSnapshotHits = isHit ? 1 : 0; + delta.iterSnapshotMisses = isHit ? 0 : 1; + delta.iterReusedTokens = reusedTokens; + delta.iterUnreusedTokens = lookupTokens - reusedTokens; + delta.iterAlignedSnapshotHits = (isHit && reusedTokens % tokensPerBlock == 0) ? 1 : 0; + delta.iterUnalignedSnapshotHits = (isHit && reusedTokens % tokensPerBlock != 0) ? 1 : 0; + mSsmSnapshotIterationStatsByLifeCycle[lifeCycle].add(delta); + return true; + } + + bool subtractAllocationRange(BlockOrdinal blockBegin, BlockOrdinal blockEnd) + { + if (blockBegin >= blockEnd || mAllocationSegments.empty()) + { + return false; + } + + bool changed = false; + int index = static_cast(mAllocationSegments.size()) - 1; + while (index >= 0) + { + auto& segment = mAllocationSegments[static_cast(index)]; + if (segment.blockEnd <= blockBegin) + { + break; + } + BlockOrdinal const removedBegin = std::max(blockBegin, segment.blockBegin); + BlockOrdinal const removedEnd = std::min(blockEnd, segment.blockEnd); + if (removedBegin >= removedEnd) + { + --index; + continue; + } + + changed = true; + subtract(allocationDelta(segment, removedBegin, removedEnd)); + if (removedBegin <= segment.blockBegin) + { + mAllocationSegments.erase(mAllocationSegments.begin() + index); + } + else + { + TLLM_CHECK_DEBUG(removedEnd == segment.blockEnd); + segment.blockEnd = removedBegin; + } + --index; + } + return changed; + } + + KVCacheStatsDelta const& globalStats() const noexcept + { + return mGlobalStats; + } + + KVCacheStatsDelta const& requestStats() const noexcept + { + return mRequestStats; + } + + IterationStatsByLifeCycle const& iterationStatsByLifeCycle() const noexcept + { + return mIterationStatsByLifeCycle; + } + + SsmSnapshotIterationStatsByLifeCycle const& ssmSnapshotIterationStatsByLifeCycle() const noexcept + { + return mSsmSnapshotIterationStatsByLifeCycle; + } + +private: + static PendingStatsDelta allocationDelta( + PendingAllocationSegment const& segment, BlockOrdinal blockBegin, BlockOrdinal blockEnd) + { + int64_t const numBlocks = static_cast(std::max(0, blockEnd - blockBegin)) * segment.beamWidth; + PendingStatsDelta delta; + delta.globalStats.allocTotalBlocks = numBlocks; + delta.globalStats.allocNewBlocks = numBlocks; + delta.globalStats.missedBlocks = segment.countAsMissed ? numBlocks : 0; + delta.requestStats = delta.globalStats.copy(); + delta.iterationStats.iterAllocTotalBlocks = numBlocks; + delta.iterationStats.iterAllocNewBlocks = numBlocks; + delta.iterationStats.iterMissedBlocks = segment.countAsMissed ? numBlocks : 0; + delta.iterationStats.iterGenAllocBlocks = segment.countAsGeneration ? numBlocks : 0; + delta.lifeCycle = segment.lifeCycle; + return delta; + } + + bool add(PendingStatsDelta const& delta) + { + if (delta.empty()) + { + return false; + } + mGlobalStats.add(delta.globalStats); + mRequestStats.add(delta.requestStats); + if (!delta.iterationStats.empty()) + { + TLLM_CHECK_DEBUG(delta.lifeCycle.has_value()); + mIterationStatsByLifeCycle[*delta.lifeCycle].add(delta.iterationStats); + } + return true; + } + + bool subtract(PendingStatsDelta const& delta) + { + if (delta.empty()) + { + return false; + } + mGlobalStats.subtract(delta.globalStats); + mRequestStats.subtract(delta.requestStats); + if (!delta.iterationStats.empty()) + { + TLLM_CHECK_DEBUG(delta.lifeCycle.has_value()); + auto const it = mIterationStatsByLifeCycle.find(*delta.lifeCycle); + if (it != mIterationStatsByLifeCycle.end()) + { + it->second.subtract(delta.iterationStats); + if (it->second.empty()) + { + mIterationStatsByLifeCycle.erase(it); + } + } + } + return true; + } + + KVCacheStatsDelta mRequestStats; + KVCacheStatsDelta mGlobalStats; + IterationStatsByLifeCycle mIterationStatsByLifeCycle; + SsmSnapshotIterationStatsByLifeCycle mSsmSnapshotIterationStatsByLifeCycle; + std::vector mAllocationSegments; +}; + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/stats.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/stats.h new file mode 100644 index 000000000000..e81e936f6e30 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/stats.h @@ -0,0 +1,265 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "kv_cache_manager_v2/common.h" +#include "kv_cache_manager_v2/lifeCycleRegistry.h" +#include "kv_cache_manager_v2/storage/config.h" + +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +struct KVCacheStatsDelta +{ + int64_t allocTotalBlocks = 0; + int64_t allocNewBlocks = 0; + int64_t reusedBlocks = 0; + int64_t missedBlocks = 0; + + void add(KVCacheStatsDelta const& other) noexcept + { + allocTotalBlocks += other.allocTotalBlocks; + allocNewBlocks += other.allocNewBlocks; + reusedBlocks += other.reusedBlocks; + missedBlocks += other.missedBlocks; + } + + void subtract(KVCacheStatsDelta const& other) noexcept + { + allocTotalBlocks -= other.allocTotalBlocks; + allocNewBlocks -= other.allocNewBlocks; + reusedBlocks -= other.reusedBlocks; + missedBlocks -= other.missedBlocks; + } + + void clear() noexcept + { + *this = {}; + } + + [[nodiscard]] KVCacheStatsDelta copy() const noexcept + { + return *this; + } + + [[nodiscard]] bool empty() const noexcept + { + return allocTotalBlocks == 0 && allocNewBlocks == 0 && reusedBlocks == 0 && missedBlocks == 0; + } + + [[nodiscard]] bool operator==(KVCacheStatsDelta const& other) const noexcept + { + return allocTotalBlocks == other.allocTotalBlocks && allocNewBlocks == other.allocNewBlocks + && reusedBlocks == other.reusedBlocks && missedBlocks == other.missedBlocks; + } +}; + +struct KVCacheIterationStatsDelta +{ + int64_t iterAllocTotalBlocks = 0; + int64_t iterAllocNewBlocks = 0; + int64_t iterReusedBlocks = 0; + int64_t iterFullReusedBlocks = 0; + int64_t iterPartialReusedBlocks = 0; + int64_t iterMissedBlocks = 0; + int64_t iterGenAllocBlocks = 0; + int64_t iterOnboardBlocks = 0; + int64_t iterOnboardBytes = 0; + int64_t iterOffloadBlocks = 0; + int64_t iterOffloadBytes = 0; + int64_t iterIntraDeviceCopyBlocks = 0; + int64_t iterIntraDeviceCopyBytes = 0; + int64_t iterHostDroppedBlocks = 0; + int64_t iterHostDroppedBytes = 0; + + void add(KVCacheIterationStatsDelta const& other) noexcept + { + iterAllocTotalBlocks += other.iterAllocTotalBlocks; + iterAllocNewBlocks += other.iterAllocNewBlocks; + iterReusedBlocks += other.iterReusedBlocks; + iterFullReusedBlocks += other.iterFullReusedBlocks; + iterPartialReusedBlocks += other.iterPartialReusedBlocks; + iterMissedBlocks += other.iterMissedBlocks; + iterGenAllocBlocks += other.iterGenAllocBlocks; + iterOnboardBlocks += other.iterOnboardBlocks; + iterOnboardBytes += other.iterOnboardBytes; + iterOffloadBlocks += other.iterOffloadBlocks; + iterOffloadBytes += other.iterOffloadBytes; + iterIntraDeviceCopyBlocks += other.iterIntraDeviceCopyBlocks; + iterIntraDeviceCopyBytes += other.iterIntraDeviceCopyBytes; + iterHostDroppedBlocks += other.iterHostDroppedBlocks; + iterHostDroppedBytes += other.iterHostDroppedBytes; + } + + void subtract(KVCacheIterationStatsDelta const& other) noexcept + { + iterAllocTotalBlocks -= other.iterAllocTotalBlocks; + iterAllocNewBlocks -= other.iterAllocNewBlocks; + iterReusedBlocks -= other.iterReusedBlocks; + iterFullReusedBlocks -= other.iterFullReusedBlocks; + iterPartialReusedBlocks -= other.iterPartialReusedBlocks; + iterMissedBlocks -= other.iterMissedBlocks; + iterGenAllocBlocks -= other.iterGenAllocBlocks; + iterOnboardBlocks -= other.iterOnboardBlocks; + iterOnboardBytes -= other.iterOnboardBytes; + iterOffloadBlocks -= other.iterOffloadBlocks; + iterOffloadBytes -= other.iterOffloadBytes; + iterIntraDeviceCopyBlocks -= other.iterIntraDeviceCopyBlocks; + iterIntraDeviceCopyBytes -= other.iterIntraDeviceCopyBytes; + iterHostDroppedBlocks -= other.iterHostDroppedBlocks; + iterHostDroppedBytes -= other.iterHostDroppedBytes; + } + + void clear() noexcept + { + *this = {}; + } + + [[nodiscard]] KVCacheIterationStatsDelta copy() const noexcept + { + return *this; + } + + [[nodiscard]] bool empty() const noexcept + { + return iterAllocTotalBlocks == 0 && iterAllocNewBlocks == 0 && iterReusedBlocks == 0 + && iterFullReusedBlocks == 0 && iterPartialReusedBlocks == 0 && iterMissedBlocks == 0 + && iterGenAllocBlocks == 0 && iterOnboardBlocks == 0 && iterOnboardBytes == 0 && iterOffloadBlocks == 0 + && iterOffloadBytes == 0 && iterIntraDeviceCopyBlocks == 0 && iterIntraDeviceCopyBytes == 0 + && iterHostDroppedBlocks == 0 && iterHostDroppedBytes == 0; + } + + [[nodiscard]] double iterCacheHitRate() const noexcept + { + int64_t const total = iterReusedBlocks + iterMissedBlocks; + if (iterReusedBlocks == 0 || total == 0) + { + return 0.0; + } + return static_cast(iterReusedBlocks) / static_cast(total); + } + + [[nodiscard]] bool operator==(KVCacheIterationStatsDelta const& other) const noexcept + { + return iterAllocTotalBlocks == other.iterAllocTotalBlocks && iterAllocNewBlocks == other.iterAllocNewBlocks + && iterReusedBlocks == other.iterReusedBlocks && iterFullReusedBlocks == other.iterFullReusedBlocks + && iterPartialReusedBlocks == other.iterPartialReusedBlocks && iterMissedBlocks == other.iterMissedBlocks + && iterGenAllocBlocks == other.iterGenAllocBlocks && iterOnboardBlocks == other.iterOnboardBlocks + && iterOnboardBytes == other.iterOnboardBytes && iterOffloadBlocks == other.iterOffloadBlocks + && iterOffloadBytes == other.iterOffloadBytes + && iterIntraDeviceCopyBlocks == other.iterIntraDeviceCopyBlocks + && iterIntraDeviceCopyBytes == other.iterIntraDeviceCopyBytes + && iterHostDroppedBlocks == other.iterHostDroppedBlocks + && iterHostDroppedBytes == other.iterHostDroppedBytes; + } +}; + +using IterationStatsByLifeCycle = std::unordered_map; + +// --------------------------------------------------------------------------- +// SsmSnapshotIterationStatsDelta — per-lifecycle counters for SSM snapshot +// reuse in one iteration. Mirrors Python's SsmSnapshotIterationStatsDelta. +// --------------------------------------------------------------------------- +struct SsmSnapshotIterationStatsDelta +{ + int64_t iterSnapshotLookups = 0; + int64_t iterSnapshotHits = 0; + int64_t iterSnapshotMisses = 0; + int64_t iterReusedTokens = 0; + int64_t iterUnreusedTokens = 0; + int64_t iterAlignedSnapshotHits = 0; + int64_t iterUnalignedSnapshotHits = 0; + + void add(SsmSnapshotIterationStatsDelta const& other) noexcept + { + iterSnapshotLookups += other.iterSnapshotLookups; + iterSnapshotHits += other.iterSnapshotHits; + iterSnapshotMisses += other.iterSnapshotMisses; + iterReusedTokens += other.iterReusedTokens; + iterUnreusedTokens += other.iterUnreusedTokens; + iterAlignedSnapshotHits += other.iterAlignedSnapshotHits; + iterUnalignedSnapshotHits += other.iterUnalignedSnapshotHits; + } + + void subtract(SsmSnapshotIterationStatsDelta const& other) noexcept + { + iterSnapshotLookups -= other.iterSnapshotLookups; + iterSnapshotHits -= other.iterSnapshotHits; + iterSnapshotMisses -= other.iterSnapshotMisses; + iterReusedTokens -= other.iterReusedTokens; + iterUnreusedTokens -= other.iterUnreusedTokens; + iterAlignedSnapshotHits -= other.iterAlignedSnapshotHits; + iterUnalignedSnapshotHits -= other.iterUnalignedSnapshotHits; + } + + void clear() noexcept + { + *this = {}; + } + + [[nodiscard]] SsmSnapshotIterationStatsDelta copy() const noexcept + { + return *this; + } + + [[nodiscard]] bool empty() const noexcept + { + return iterSnapshotLookups == 0 && iterSnapshotHits == 0 && iterSnapshotMisses == 0 && iterReusedTokens == 0 + && iterUnreusedTokens == 0 && iterAlignedSnapshotHits == 0 && iterUnalignedSnapshotHits == 0; + } + + [[nodiscard]] double iterSnapshotHitRate() const noexcept + { + if (iterSnapshotHits == 0 || iterSnapshotLookups == 0) + { + return 0.0; + } + return static_cast(iterSnapshotHits) / static_cast(iterSnapshotLookups); + } + + [[nodiscard]] bool operator==(SsmSnapshotIterationStatsDelta const& other) const noexcept + { + return iterSnapshotLookups == other.iterSnapshotLookups && iterSnapshotHits == other.iterSnapshotHits + && iterSnapshotMisses == other.iterSnapshotMisses && iterReusedTokens == other.iterReusedTokens + && iterUnreusedTokens == other.iterUnreusedTokens + && iterAlignedSnapshotHits == other.iterAlignedSnapshotHits + && iterUnalignedSnapshotHits == other.iterUnalignedSnapshotHits; + } +}; + +using SsmSnapshotIterationStatsByLifeCycle = std::unordered_map; + +struct PoolGroupPeakBlockStats +{ + SlotCount available = 0; + SlotCount unavailable = 0; + SlotCount evictable = 0; + + [[nodiscard]] bool operator==(PoolGroupPeakBlockStats const& other) const noexcept + { + return available == other.available && unavailable == other.unavailable && evictable == other.evictable; + } +}; + +using PeakBlockStatsByPoolGroup = TypedVec; +using PeakBlockStatsByCacheLevel = TypedVec; + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp new file mode 100644 index 000000000000..7411577d03b1 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp @@ -0,0 +1,247 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "kv_cache_manager_v2/storage/config.h" +#include "kv_cache_manager_v2/common.h" +#include "kv_cache_manager_v2/utils/math.h" + +#include "tensorrt_llm/common/assert.h" +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// --------------------------------------------------------------------------- +// StorageConfig methods +// --------------------------------------------------------------------------- + +LifeCycleId StorageConfig::numLifeCycles() const +{ + LifeCycleId count{0}; + for (auto const& sd : slotDescList) + { + count += static_cast(sd.variants.size()); + } + return count; +} + +TypedVec StorageConfig::lifeCycleGrouping() const +{ + LifeCycleId n = numLifeCycles(); + TypedVec ret(n, PoolGroupIndex{-1}); + for (PoolGroupIndex poolGroupIndex{0}; poolGroupIndex < slotDescList.size(); ++poolGroupIndex) + { + for (auto const& variant : slotDescList[poolGroupIndex].variants) + { + ret[variant.lifeCycleId] = poolGroupIndex; + } + } + return ret; +} + +std::map StorageConfig::bufferAttributes() const +{ + std::map ret; + for (auto const& sd : slotDescList) + { + for (auto const& variant : sd.variants) + { + LifeCycleId lcId = variant.lifeCycleId; + for (PoolIndex poolIndex{0}; poolIndex < variant.coalescedBuffers.size(); ++poolIndex) + { + auto const& cb = variant.coalescedBuffers[poolIndex]; + size_t offset = 0; + for (auto const& bufId : cb.bufferIds) + { + int exp = 1; + auto it = expansion.find(bufId); + if (it != expansion.end()) + exp = it->second; + + ret[bufId] = BufferAttr{lcId, poolIndex, offset, cb.singleBufferSize, exp}; + offset += cb.singleBufferSize; + } + } + } + } + return ret; +} + +std::unordered_map StorageConfig::layerToLifeCycleIds() const +{ + std::unordered_map map; + for (auto const& [bufId, attr] : bufferAttributes()) + { + auto [it, inserted] = map.emplace(bufId.layerId, attr.lifeCycleId); + if (!inserted) + { + TLLM_CHECK_DEBUG(it->second == attr.lifeCycleId); + } + } + return map; +} + +std::map StorageConfig::layerAttributes() const +{ + std::map ret; + for (auto const& pg : slotDescList) + { + for (auto const& variant : pg.variants) + { + LifeCycleId lcId = variant.lifeCycleId; + for (PoolIndex poolIndex{0}; poolIndex < variant.coalescedBuffers.size(); ++poolIndex) + { + auto const& cb = variant.coalescedBuffers[poolIndex]; + // Count how many buffers per layer in this coalesced buffer. + std::unordered_map slotUtilPerLayer; + for (auto const& bufId : cb.bufferIds) + { + slotUtilPerLayer[bufId.layerId] += 1; + } + int buffersPerSlot = cb.numBuffers(); + + for (auto const& [layerId, count] : slotUtilPerLayer) + { + auto it = ret.find(layerId); + if (it == ret.end()) + { + LayerAttr attr; + attr.lifeCycleId = lcId; + attr.slotUtil.resize(variant.coalescedBuffers.size(), 0); + attr.slotUtilFracMax = Rational{0, 1}; + it = ret.emplace(layerId, std::move(attr)).first; + } + auto& attr = it->second; + TLLM_CHECK_DEBUG(attr.lifeCycleId == lcId); + attr.slotUtil[poolIndex] = count; + Rational frac{count, buffersPerSlot}; + if (frac > attr.slotUtilFracMax) + { + attr.slotUtilFracMax = frac; + } + } + } + } + } + return ret; +} + +// --------------------------------------------------------------------------- +// createStorageConfig — factory function. +// Mirrors _storage/_config.py::create_storage_config. +// --------------------------------------------------------------------------- +StorageConfig createStorageConfig(KVCacheManagerConfig const& config) +{ + LifeCycleRegistry registry{config}; + int tokensPerBlock = config.tokensPerBlock; + + // Map: lifeCycleId → (bufferSize → list of BufferId). + // Outer map key is LifeCycleId; inner map key is expanded buffer size. + // NOTE: Python uses insertion-ordered dict/defaultdict here. The std::map + // grouping below intentionally remains sorted, which can make observable + // pool-group indices/layout order differ from Python. Later lookups use + // StorageConfig mappings, so this is not expected to affect correctness. + std::map>> bufferGroups; + std::unordered_map expansionMap; + + for (auto const& layer : config.layers) + { + LifeCycle lc = makeLifeCycle(layer, tokensPerBlock); + LifeCycleId lcId = registry.getId(lc); + + std::visit( + [&](auto const& cfg) + { + for (auto const& buf : cfg.buffers) + { + int tpbo = buf.tokensPerBlockOverride.value_or(tokensPerBlock); + int exp = exactDiv(tokensPerBlock, tpbo); + size_t expandedSize = buf.size * static_cast(exp); + + BufferId bid{cfg.layerId, buf.role}; + expansionMap[bid] = exp; + bufferGroups[lcId][expandedSize].push_back(bid); + } + }, + layer); + } + + // Build one SlotDescVariant per life cycle. + // Each variant: coalesced buffers sorted by size descending. + std::vector slotGroups; + slotGroups.reserve(bufferGroups.size()); + for (auto const& [lcId, sizeMap] : bufferGroups) + { + SlotDescVariant var; + var.lifeCycleId = lcId; + for (auto const& [sz, bufIds] : sizeMap) + { + CoalescedBuffer cb; + cb.singleBufferSize = sz; + cb.bufferIds = bufIds; + var.coalescedBuffers.push_back(std::move(cb)); + } + // Sort descending by size. + std::sort(var.coalescedBuffers.begin(), var.coalescedBuffers.end(), + [](CoalescedBuffer const& a, CoalescedBuffer const& b) { return a.size() > b.size(); }); + slotGroups.push_back(std::move(var)); + } + + // Merge SlotDescVariants that share the same slotSizeList. + // Key: tuple of sizes (sorted desc). + std::map, std::vector> poolGroupsBySizes; + for (auto& sg : slotGroups) + { + auto sizes = sg.slotSizeList(); + poolGroupsBySizes[sizes.raw()].push_back(std::move(sg)); + } + + StorageConfig out; + out.cacheTiers = TypedVec{config.cacheTiers}; + out.expansion = expansionMap; + + for (auto& [sizes, variants] : poolGroupsBySizes) + { + SlotDesc sd; + sd.variants = std::move(variants); + out.slotDescList.push_back(std::move(sd)); + } + + // A21: Assert all life_cycle_ids across all SlotDescVariants are unique. + // Mirrors Python StorageConfig.__post_init__: + // all_life_cycle_ids = [lc_id for variant in self.slot_desc_list for lc_id in variant.life_cycle_ids] + // assert len(all_life_cycle_ids) == len(set(all_life_cycle_ids)) + if (TLLM_UNLIKELY(gDebug)) + { + std::unordered_set allLcIds; + for (auto const& sd : out.slotDescList) + { + for (auto const& variant : sd.variants) + { + [[maybe_unused]] bool inserted = allLcIds.insert(variant.lifeCycleId).second; + TLLM_CHECK_WITH_INFO(inserted, "Duplicate life_cycle_id across SlotDescVariants"); + } + } + } + + return out; +} + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.h new file mode 100644 index 000000000000..b09d0dbbc9a0 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.h @@ -0,0 +1,249 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "kv_cache_manager_v2/common.h" +#include "kv_cache_manager_v2/config.h" +#include "kv_cache_manager_v2/lifeCycleRegistry.h" + +#include "tensorrt_llm/common/assert.h" +#include +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// --------------------------------------------------------------------------- +// Index types (mirrors _storage/_core.py) +// --------------------------------------------------------------------------- +// Index of a pool group (= set of pools with mirrored allocation). +using PoolGroupIndex = StrongIndex; +// Index of a pool within a pool group. +using PoolIndex = StrongIndex; +// Plain count of slots when the value is not itself a slot id. +using SlotCount = std::int64_t; +// Slot index within a pool. +using SlotId = StrongIndex; + +[[nodiscard]] inline SlotCount slotCountValueFromSize(std::size_t count) noexcept +{ + TLLM_CHECK_DEBUG(count <= static_cast(std::numeric_limits::max())); + return static_cast(count); +} + +[[nodiscard]] inline std::size_t slotCountToSizeT(SlotCount count) noexcept +{ + TLLM_CHECK_DEBUG_WITH_INFO(count >= 0, "Slot count must be non-negative for size_t conversion"); + return static_cast(count); +} + +[[nodiscard]] inline int slotIdToPageIndexValue(SlotId slotId) +{ + if (slotId.value() < 0 || slotId.value() > std::numeric_limits::max()) + { + throw std::overflow_error("SlotId does not fit in int page-index storage"); + } + return static_cast(slotId.value()); +} + +// --------------------------------------------------------------------------- +// BufferId — (layer_id, role) pair identifying one buffer in a layer. +// Mirrors _storage/_config.py::BufferId. +// --------------------------------------------------------------------------- +struct BufferId +{ + LayerId layerId = 0; + DataRole role; + + bool operator==(BufferId const& o) const noexcept + { + return layerId == o.layerId && role == o.role; + } + + bool operator<(BufferId const& o) const noexcept + { + if (layerId != o.layerId) + return layerId < o.layerId; + return role < o.role; + } +}; + +struct BufferIdHash +{ + size_t operator()(BufferId const& id) const noexcept + { + size_t h = std::hash{}(id.layerId); + h ^= std::hash{}(id.role) + 0x9e3779b9 + (h << 6) + (h >> 2); + return h; + } +}; + +// --------------------------------------------------------------------------- +// CoalescedBuffer — several buffers of the same size and life cycle, +// laid out contiguously within one slot. +// Mirrors _storage/_config.py::CoalescedBuffer. +// --------------------------------------------------------------------------- +struct CoalescedBuffer +{ + size_t singleBufferSize = 0; + std::vector bufferIds; + + size_t size() const noexcept + { + return singleBufferSize * bufferIds.size(); + } + + int numBuffers() const noexcept + { + return static_cast(bufferIds.size()); + } +}; + +// --------------------------------------------------------------------------- +// SlotDescVariant — one life cycle's view of a pool group. +// Mirrors _storage/_config.py::SlotDescVariant. +// --------------------------------------------------------------------------- +struct SlotDescVariant +{ + LifeCycleId lifeCycleId{0}; + TypedVec coalescedBuffers; // sorted size desc + + // Slot size for each pool in this group. + TypedVec slotSizeList() const + { + TypedVec out; + out.reserve(coalescedBuffers.size()); + for (auto const& cb : coalescedBuffers) + out.push_back(cb.size()); + // Coalesced buffers must be sorted in descending size order. + TLLM_CHECK_DEBUG(std::is_sorted(out.begin(), out.end(), std::greater<>())); + return out; + } +}; + +// --------------------------------------------------------------------------- +// SlotDesc — a pool group descriptor. Variants share the same slotSizeList +// but may differ in which buffers they contain. +// Mirrors _storage/_config.py::SlotDesc. +// --------------------------------------------------------------------------- +struct SlotDesc +{ + std::vector variants; // different life cycles sharing this pool group + + TypedVec slotSizeList() const + { + return getUniformAttribute(variants, [](auto const& v) { return v.slotSizeList(); }); + } +}; + +// --------------------------------------------------------------------------- +// BufferAttr — metadata for one buffer within storage. +// Mirrors _storage/_config.py::BufferAttr. +// --------------------------------------------------------------------------- +struct BufferAttr +{ + LifeCycleId lifeCycleId{0}; + PoolIndex poolIndex{0}; + size_t offset = 0; // byte offset within the slot + size_t size = 0; // expanded size of the buffer (after expansion) + int expansion = 1; // expansion factor (tokens_per_block / tokens_per_block_override) +}; + +// --------------------------------------------------------------------------- +// Rational — exact fraction for slot utilization computations. +// Mirrors Python's fractions.Fraction used in _storage/_config.py::LayerAttr. +// --------------------------------------------------------------------------- +struct Rational +{ + int num = 0; + int den = 1; + + // ceil(n * num / den). Uses int64_t intermediates to avoid overflow. + [[nodiscard]] int ceilMul(int value) const noexcept + { + auto n = static_cast(value) * num; + return static_cast((n + den - 1) / den); + } + + bool operator>(Rational const& other) const noexcept + { + return static_cast(num) * other.den > static_cast(other.num) * den; + } + + bool operator==(Rational const& other) const noexcept + { + return static_cast(num) * other.den == static_cast(other.num) * den; + } +}; + +// --------------------------------------------------------------------------- +// LayerAttr — per-layer storage attributes for scratch slot management. +// Mirrors _storage/_config.py::LayerAttr. +// --------------------------------------------------------------------------- +struct LayerAttr +{ + LifeCycleId lifeCycleId{0}; + + // Number of sub-pages within a single coalesced slot that belong to this layer, + // per pool within the pool group. + TypedVec slotUtil; + + // Fraction of slot_util to total number of buffers in the slot, max over all pools. + Rational slotUtilFracMax; +}; + +// --------------------------------------------------------------------------- +// StorageConfig — complete storage layout for the KV cache. +// Mirrors _storage/_config.py::StorageConfig. +// --------------------------------------------------------------------------- +struct StorageConfig +{ + TypedVec cacheTiers; + TypedVec slotDescList; + + // Expansion factor per buffer (for heterogeneous tokens_per_block). + std::unordered_map expansion; + + // Map each LifeCycleId → its PoolGroupIndex. + TypedVec lifeCycleGrouping() const; + + // Attribute map for each buffer. + std::map bufferAttributes() const; + + // Map LayerId → LifeCycleId. + std::unordered_map layerToLifeCycleIds() const; + + // Per-layer storage attributes (slot utilization within coalesced slots). + // Mirrors _storage/_config.py::StorageConfig.layer_attributes(). + std::map layerAttributes() const; + + LifeCycleId numLifeCycles() const; +}; + +// --------------------------------------------------------------------------- +// Factory: create StorageConfig from KVCacheManagerConfig. +// Mirrors _storage/_config.py::create_storage_config. +// --------------------------------------------------------------------------- +StorageConfig createStorageConfig(KVCacheManagerConfig const& config); + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.cpp new file mode 100644 index 000000000000..892687304452 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.cpp @@ -0,0 +1,946 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "kv_cache_manager_v2/storage/core.h" + +#include "tensorrt_llm/common/assert.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// --------------------------------------------------------------------------- +// SlotAllocator +// --------------------------------------------------------------------------- + +SlotAllocator::SlotAllocator(SlotCount capacity) + : mCapacity(capacity) + , mTargetCapacity(capacity) + , mNumActiveSlots(0) + , mNumReadyRecycledSlots(0) + , mOccupiedMask(slotCountToSizeT(capacity)) +{ +} + +SlotAllocator::~SlotAllocator() +{ + // Mirrors Python SlotAllocator.__del__ (assert_critical checks). + if (TLLM_UNLIKELY(gDebug)) + { + TLLM_CHECK_WITH_INFO(slotCountToSizeT(mNumReadyRecycledSlots) == mRecycledSlots.size(), + "SlotAllocator destroyed with unfinished events — did you call synchronize()?"); + TLLM_CHECK_WITH_INFO(mTargetCapacity == mCapacity && mOverflowSlots.empty(), + "SlotAllocator destroyed while resize is in progress"); + TLLM_CHECK_WITH_INFO( + mOccupiedMask.numSetBits() == 0, "SlotAllocator destroyed with occupied slots still in use"); + TLLM_CHECK_WITH_INFO(mRecycledSlots.size() == slotCountToSizeT(mNumActiveSlots), + "SlotAllocator destroyed with some slots not recycled"); + } +} + +SlotCount SlotAllocator::numFreeSlots() const noexcept +{ + SlotCount const inactiveSlots = mTargetCapacity > mNumActiveSlots ? mTargetCapacity - mNumActiveSlots : 0; + return slotCountValueFromSize(mRecycledSlots.size()) + inactiveSlots; +} + +SlotCount SlotAllocator::numOccupiedSlots() const noexcept +{ + return slotCountValueFromSize(mOccupiedMask.numSetBits()); +} + +Slot SlotAllocator::allocate() +{ + if (numFreeSlots() == 0) + { + throw OutOfPagesError("SlotAllocator: no free slots"); + } + scrubEvents(); + + Slot slot; + if (mNumReadyRecycledSlots > 0) + { + TLLM_CHECK_DEBUG_WITH_INFO(!mRecycledSlots.empty(), "ready recycled slots > 0 but deque is empty"); + slot = std::move(mRecycledSlots.front()); + mRecycledSlots.pop_front(); + TLLM_CHECK_DEBUG_WITH_INFO(slot.hasValidSlot(), "ready recycled slot has no valid id"); + --mNumReadyRecycledSlots; + TLLM_CHECK_DEBUG_WITH_INFO(slot.readyEvent.isClosed(), "ready recycled slot has non-null event"); + } + else if (mNumActiveSlots < std::min(mCapacity, mTargetCapacity)) + { + slot.setSlotId(SlotId{mNumActiveSlots++}); + } + else + { + slot = std::move(mRecycledSlots.front()); + mRecycledSlots.pop_front(); + TLLM_CHECK_DEBUG_WITH_INFO(slot.hasValidSlot(), "non-ready recycled slot has no valid id"); + } + mOccupiedMask.set(toSizeT(slot.slotId())); + return slot; +} + +std::vector SlotAllocator::allocateMultiple(SlotCount numSlots) +{ + if (numSlots < 0) + { + throw LogicError("SlotAllocator::allocateMultiple: slot count must be non-negative"); + } + if (numFreeSlots() < numSlots) + { + throw OutOfPagesError("SlotAllocator: not enough free slots"); + } + std::vector result; + result.reserve(slotCountToSizeT(numSlots)); + for (SlotCount slotCount{0}; slotCount < numSlots; ++slotCount) + { + result.push_back(allocate()); + } + return result; +} + +void SlotAllocator::release(Slot slot) +{ + if (!slot.hasValidSlot()) + { + throw LogicError("SlotAllocator::release: slot has no valid id"); + } + SlotId const slotId = slot.slotId(); + if (slotId >= numSlots() || !mOccupiedMask.get(toSizeT(slotId))) + { + throw LogicError("SlotAllocator::release: slot is not occupied"); + } + mOccupiedMask.clear(toSizeT(slotId)); + if (slotId < mTargetCapacity) + { + mRecycledSlots.push_back(std::move(slot)); + } + else + { + mOverflowSlots.push_back(std::move(slot)); + } + scrubEvents(); + TLLM_CHECK_DEBUG(check()); +} + +void SlotAllocator::expand(SlotCount newNumSlots) +{ + TLLM_CHECK_DEBUG(check()); + TLLM_CHECK_DEBUG(mTargetCapacity == mCapacity); + TLLM_CHECK_DEBUG(newNumSlots > mCapacity); + mOccupiedMask.resize(slotCountToSizeT(newNumSlots)); + mCapacity = newNumSlots; + mTargetCapacity = newNumSlots; + TLLM_CHECK_DEBUG(check()); +} + +void SlotAllocator::prepareForShrink(SlotCount newNumSlots) +{ + TLLM_CHECK_DEBUG(check()); + TLLM_CHECK_DEBUG(mTargetCapacity == mCapacity); + TLLM_CHECK_DEBUG(newNumSlots < mCapacity); + std::deque newRecycled; + SlotCount newNumReady = 0; + SlotCount const oldNumReady = mNumReadyRecycledSlots; + SlotCount idx = 0; + for (auto& s : mRecycledSlots) + { + if (s.slotId() < newNumSlots) + { + if (idx < oldNumReady) + ++newNumReady; + newRecycled.push_back(std::move(s)); + } + else + { + mOverflowSlots.push_back(std::move(s)); + } + ++idx; + } + mRecycledSlots = std::move(newRecycled); + mNumReadyRecycledSlots = newNumReady; + mTargetCapacity = newNumSlots; + TLLM_CHECK_DEBUG(check()); +} + +bool SlotAllocator::finishShrink() +{ + TLLM_CHECK_DEBUG(check()); + // Overflow-range IDs that were ever issued are exactly + // max(0, _num_active_slots - _target_capacity); the underused case + // (_num_active_slots <= _target_capacity) collapses to zero. + SlotCount const expectedOverflow = std::max(SlotCount{0}, mNumActiveSlots - mTargetCapacity); + if (shrinkInProgress() && slotCountValueFromSize(mOverflowSlots.size()) == expectedOverflow) + { + // Validate uniqueness of slot IDs in overflow (debug only). + if (TLLM_UNLIKELY(gDebug)) + { + std::set ids; + for (auto const& s : mOverflowSlots) + { + TLLM_CHECK(s.hasValidSlot()); + ids.insert(s.slotId()); + } + TLLM_CHECK_WITH_INFO(ids.size() == mOverflowSlots.size(), "Duplicate slot IDs in overflow slots"); + } + // Synchronize overflow events (deduplicated — slots often share events). + { + std::vector overflowEvents; + overflowEvents.reserve(mOverflowSlots.size()); + for (auto& s : mOverflowSlots) + overflowEvents.push_back(&s.readyEvent); + synchronizeAll(overflowEvents); + } + for (auto& s : mOverflowSlots) + s.resetSlot(); + mOverflowSlots.clear(); + mCapacity = mTargetCapacity; + mNumActiveSlots = std::min(mNumActiveSlots, mCapacity); + scrubEvents(); + TLLM_CHECK_DEBUG(check()); + return true; + } + throw std::runtime_error("SlotAllocator::finishShrink: cannot finish shrink yet"); +} + +std::vector SlotAllocator::getSlotsBlockingShrink() const +{ + std::vector result; + for (SlotCount id = mTargetCapacity; id < mCapacity; ++id) + { + if (mOccupiedMask.get(slotCountToSizeT(id))) + result.push_back(SlotId{id}); + } + return result; +} + +void SlotAllocator::synchronize() +{ + while (slotCountToSizeT(mNumReadyRecycledSlots) != mRecycledSlots.size()) + { + scrubEvents(); + } +} + +bool SlotAllocator::check() const noexcept +{ + // Mirrors Python SlotAllocator._check(). + if (mNumActiveSlots > mCapacity) + return false; + if (mTargetCapacity > mCapacity) + return false; + if (!shrinkInProgress() && !mOverflowSlots.empty()) + return false; + for (auto const& slot : mOverflowSlots) + { + if (!slot.hasValidSlot()) + return false; + SlotCount id = slot.slotId().value(); + if (id < mTargetCapacity || id >= mCapacity) + return false; + } + SlotCount const accountedSlots = slotCountValueFromSize(mRecycledSlots.size()) + + slotCountValueFromSize(mOverflowSlots.size()) + numOccupiedSlots(); + if (accountedSlots != mNumActiveSlots) + return false; + return true; +} + +void SlotAllocator::scrubEvents() +{ + for (size_t i = slotCountToSizeT(mNumReadyRecycledSlots); i < mRecycledSlots.size(); ++i) + { + if (mRecycledSlots[i].queryReady()) + { + ++mNumReadyRecycledSlots; + } + else + { + break; + } + } +} + +// --------------------------------------------------------------------------- +// GpuSlotPool +// --------------------------------------------------------------------------- + +GpuSlotPool::GpuSlotPool(size_t slotSize, size_t vmSize, PooledPhysMemAllocator& physMemAllocator, SlotCount numSlots) + : SlotPoolBase(slotSize) + , mVirtMem(vmSize, physMemAllocator) +{ + TLLM_CHECK_DEBUG_WITH_INFO( + vmSize % physMemAllocator.physMemSize() == 0, "vm_size must be aligned to phys_mem_size"); + resize(numSlots); +} + +size_t GpuSlotPool::computeNumPhysMem(size_t slotSize, SlotCount numSlots, size_t physMemSize) noexcept +{ + return divUp(slotCountToSizeT(numSlots) * slotSize, physMemSize); +} + +SlotCount GpuSlotPool::computeNumSlots(size_t slotSize, size_t numPhysMem, size_t physMemSize) noexcept +{ + return slotCountValueFromSize(numPhysMem * physMemSize / slotSize); +} + +SlotCount GpuSlotPool::numSlots() const noexcept +{ + return computeNumSlots(mSlotSize, mVirtMem.numPhysMem(), mVirtMem.physMemSize()); +} + +void GpuSlotPool::destroy() +{ + mVirtMem.destroy(); +} + +void GpuSlotPool::resize(SlotCount newNumSlots) +{ + size_t physSize = mVirtMem.physMemSize(); + size_t newPhysMem = computeNumPhysMem(mSlotSize, newNumSlots, physSize); + mVirtMem.realloc(physSize * newPhysMem); +} + +SlotCount GpuSlotPool::extendByOnePhysMem() +{ + mVirtMem.extend(1); + return numSlots(); +} + +Address GpuSlotPool::slotAddress(SlotId slot) const +{ + TLLM_CHECK_DEBUG_WITH_INFO(slot < numSlots(), "GpuSlotPool::slotAddress: slot index out of bounds"); + return MemAddress(mVirtMem.address() + mSlotSize * toSizeT(slot)); +} + +// --------------------------------------------------------------------------- +// HostSlotPool +// --------------------------------------------------------------------------- + +HostSlotPool::HostSlotPool(size_t slotSize, SlotCount numSlots) + : SlotPoolBase(slotSize) + , mHostMem(alignedSize(numSlots)) +{ +} + +size_t HostSlotPool::alignedSize(SlotCount numSlots) const noexcept +{ + return roundUp(slotCountToSizeT(numSlots) * mSlotSize, HostMem::kAlignment); +} + +SlotCount HostSlotPool::numSlots() const noexcept +{ + return slotCountValueFromSize(mHostMem.size() / mSlotSize); +} + +void HostSlotPool::destroy() +{ + mHostMem.destroy(); +} + +void HostSlotPool::resize(SlotCount newNumSlots) +{ + mHostMem.resize(alignedSize(newNumSlots)); +} + +Address HostSlotPool::slotAddress(SlotId slot) const +{ + TLLM_CHECK_DEBUG_WITH_INFO(slot < numSlots(), "HostSlotPool::slotAddress: slot index out of bounds"); + return MemAddress(mHostMem.address() + mSlotSize * toSizeT(slot)); +} + +// --------------------------------------------------------------------------- +// DiskSlotPool +// --------------------------------------------------------------------------- + +DiskSlotPool::DiskSlotPool(std::string const& directory, size_t slotSize, SlotCount numSlots) + : SlotPoolBase(slotSize) +{ + // Try O_TMPFILE first, fall back to mkstemp. + mFd = ::open(directory.c_str(), O_TMPFILE | O_RDWR | O_EXCL, 0664); + if (mFd < 0) + { + if (errno == EOPNOTSUPP) + { + char tmpl[PATH_MAX]; + snprintf(tmpl, sizeof(tmpl), "%s/kvXXXXXX", directory.c_str()); + mFd = ::mkstemp(tmpl); + if (mFd < 0) + { + throw DiskOOMError("DiskSlotPool: mkstemp failed: " + std::string(std::strerror(errno))); + } + ::unlink(tmpl); + } + else + { + throw DiskOOMError("DiskSlotPool: open O_TMPFILE failed: " + std::string(std::strerror(errno))); + } + } + resize(numSlots); +} + +DiskSlotPool::~DiskSlotPool() +{ + destroy(); +} + +SlotCount DiskSlotPool::numSlots() const noexcept +{ + TLLM_CHECK_DEBUG(mFd != kBadFileDescriptor); + off_t sz = ::lseek(mFd, 0, SEEK_END); + return (sz < 0 || mSlotSize == 0) ? 0 : slotCountValueFromSize(static_cast(sz) / mSlotSize); +} + +void DiskSlotPool::destroy() +{ + if (mFd != kBadFileDescriptor) + { + ::close(mFd); + mFd = kBadFileDescriptor; + } +} + +void DiskSlotPool::resize(SlotCount newNumSlots) +{ + resizeFile(mFd, slotCountToSizeT(newNumSlots) * mSlotSize); +} + +Address DiskSlotPool::slotAddress(SlotId slot) const +{ + TLLM_CHECK_DEBUG_WITH_INFO(slot < numSlots(), "DiskSlotPool::slotAddress: slot index out of bounds"); + size_t const byteOffset = toSizeT(slot) * mSlotSize; + TLLM_CHECK_DEBUG_WITH_INFO(byteOffset <= static_cast(std::numeric_limits::max()), + "DiskSlotPool::slotAddress: byte offset out of range"); + return DiskAddress{mFd, static_cast(byteOffset)}; +} + +// --------------------------------------------------------------------------- +// PoolGroupBase +// --------------------------------------------------------------------------- + +PoolGroupBase::PoolGroupBase(SlotCount numSlots) + : mSlotAllocator(numSlots) +{ +} + +SlotCount PoolGroupBase::getNumSlotsFromPools() const noexcept +{ + if (mPools.empty()) + return 0; + SlotCount minSlots = mPools.front()->numSlots(); + for (PoolIndex poolIdx{1}; poolIdx < mPools.size(); ++poolIdx) + { + minSlots = std::min(minSlots, mPools[poolIdx]->numSlots()); + } + return minSlots; +} + +PoolGroupBase::~PoolGroupBase() +{ + destroy(); +} + +SlotCount PoolGroupBase::numSlots() const noexcept +{ + SlotCount n = mSlotAllocator.numSlots(); + if (TLLM_UNLIKELY(gDebug)) + { + // Mirrors Python PoolGroupBase.num_slots: assert num_slots <= self._get_num_slots_from_pools() + [[maybe_unused]] SlotCount poolSlots = getNumSlotsFromPools(); + TLLM_CHECK_WITH_INFO(n <= poolSlots, "SlotAllocator capacity exceeds pool capacity"); + } + return n; +} + +Slot PoolGroupBase::allocate() +{ + return mSlotAllocator.allocate(); +} + +std::vector PoolGroupBase::allocateMultiple(SlotCount numSlots) +{ + return mSlotAllocator.allocateMultiple(numSlots); +} + +void PoolGroupBase::release(Slot slot) +{ + mSlotAllocator.release(std::move(slot)); +} + +void PoolGroupBase::destroy() +{ + if (mDestroyed) + return; + if (mSlotAllocator.numSlots() != 0) + { + mSlotAllocator.synchronize(); + mSlotAllocator.prepareForShrink(0); + mSlotAllocator.finishShrink(); + } + for (auto& p : mPools) + p->destroy(); + mDestroyed = true; +} + +void PoolGroupBase::resizePools(std::optional newNumSlots) +{ + SlotCount n = newNumSlots.value_or(mSlotAllocator.numSlots()); + for (auto& p : mPools) + p->resize(n); + // Mirrors Python PoolGroupBase.resize_pools: assert NDEBUG or self._check(True) + // After resize, allocator capacity must not exceed pool capacity (allow mismatch). + TLLM_CHECK_DEBUG_WITH_INFO(mSlotAllocator.numSlots() <= getNumSlotsFromPools(), + "After resizePools: allocator capacity exceeds pool capacity"); +} + +TypedVec PoolGroupBase::slotAddress(SlotId slotId) const +{ + TypedVec addrs; + addrs.reserve(mPools.size()); + for (auto const& p : mPools) + addrs.push_back(p->slotAddress(slotId)); + return addrs; +} + +TypedVec PoolGroupBase::slotSize() const +{ + TypedVec sizes; + sizes.reserve(mPools.size()); + for (auto const& p : mPools) + sizes.push_back(p->slotSize()); + return sizes; +} + +// --------------------------------------------------------------------------- +// GpuPoolGroup +// --------------------------------------------------------------------------- + +GpuPoolGroup::GpuPoolGroup( + SlotCount numSlots, TypedVec const& slotSizeList, PooledPhysMemAllocator& physMemAllocator) + : PoolGroupBase(numSlots) +{ + size_t physMemSize = physMemAllocator.physMemSize(); + // Query total GPU memory to size virtual address space (mirrors Python GpuPoolGroup). + size_t totalGpuMem = 0; + { + CUdevice dev{}; + cuCheck(cuCtxGetDevice(&dev)); + cuCheck(cuDeviceTotalMem(&totalGpuMem, dev)); + } + // @TODO: We should replace maxSlotSize with sum. This should also be updated in Python. Will do it later. + TLLM_CHECK_WITH_INFO(!slotSizeList.empty(), "GpuPoolGroup: slotSizeList must not be empty"); + size_t maxSlotSize = *std::max_element(slotSizeList.begin(), slotSizeList.end()); + for (size_t sz : slotSizeList) + { + // VA proportional to GPU memory, scaled by slot size ratio (mirrors Python). + // Compute ratio as double first to avoid size_t overflow. + double sizeRatio = static_cast(sz) / static_cast(maxSlotSize); + size_t vmSize = roundDown(static_cast(static_cast(totalGpuMem) * sizeRatio), physMemSize); + vmSize = std::max(vmSize, roundUp(slotCountToSizeT(numSlots) * sz, physMemSize)); + mPools.push_back(std::make_unique(sz, vmSize, physMemAllocator, numSlots)); + } +} + +// --------------------------------------------------------------------------- +// HostPoolGroup +// --------------------------------------------------------------------------- + +HostPoolGroup::HostPoolGroup(SlotCount numSlots, TypedVec const& slotSizeList) + : PoolGroupBase(numSlots) +{ + for (size_t sz : slotSizeList) + { + mPools.push_back(std::make_unique(sz, numSlots)); + } +} + +// --------------------------------------------------------------------------- +// DiskPoolGroup +// --------------------------------------------------------------------------- + +DiskPoolGroup::DiskPoolGroup( + SlotCount numSlots, TypedVec const& slotSizeList, std::string const& directory) + : PoolGroupBase(numSlots) +{ + for (size_t sz : slotSizeList) + { + mPools.push_back(std::make_unique(directory, sz, numSlots)); + } +} + +// --------------------------------------------------------------------------- +// CacheLevelStorage +// --------------------------------------------------------------------------- + +std::vector CacheLevelStorage::allocateMultiple(PoolGroupIndex pgIdx, SlotCount numSlots) +{ + return mPoolGroups.at(pgIdx)->allocateMultiple(numSlots); +} + +void CacheLevelStorage::release(PoolGroupIndex pgIdx, Slot slot) +{ + mPoolGroups.at(pgIdx)->release(std::move(slot)); +} + +SlotCount CacheLevelStorage::numFreeSlots(PoolGroupIndex pgIdx) const +{ + return mPoolGroups.at(pgIdx)->numFreeSlots(); +} + +TypedVec CacheLevelStorage::slotAddress(PoolGroupIndex pgIdx, SlotId slotId) const +{ + return mPoolGroups.at(pgIdx)->slotAddress(slotId); +} + +// --------------------------------------------------------------------------- +// GpuCacheLevelStorage +// --------------------------------------------------------------------------- + +GpuCacheLevelStorage::GpuCacheLevelStorage( + StorageConfig const& storageCfg, TypedVec const& slotCountList, size_t physMemSize) +{ + TLLM_CHECK_DEBUG_WITH_INFO(slotCountList.size() == storageCfg.slotDescList.size(), + "GpuCacheLevelStorage: slotCountList and slotDescList must have the same length"); + mPhysMemAllocator = std::make_unique(physMemSize); + + for (PoolGroupIndex pgIdx{0}; pgIdx < storageCfg.slotDescList.size(); ++pgIdx) + { + mPoolGroups.push_back(std::make_unique( + slotCountList[pgIdx], storageCfg.slotDescList[pgIdx].slotSizeList(), *mPhysMemAllocator)); + } +} + +// --------------------------------------------------------------------------- +// HostCacheLevelStorage +// --------------------------------------------------------------------------- + +HostCacheLevelStorage::HostCacheLevelStorage( + StorageConfig const& storageCfg, TypedVec const& slotCountList) +{ + TLLM_CHECK_DEBUG_WITH_INFO(slotCountList.size() == storageCfg.slotDescList.size(), + "HostCacheLevelStorage: slotCountList and slotDescList must have the same length"); + for (PoolGroupIndex pgIdx{0}; pgIdx < storageCfg.slotDescList.size(); ++pgIdx) + { + mPoolGroups.push_back( + std::make_unique(slotCountList[pgIdx], storageCfg.slotDescList[pgIdx].slotSizeList())); + } +} + +// --------------------------------------------------------------------------- +// DiskCacheLevelStorage +// --------------------------------------------------------------------------- + +DiskCacheLevelStorage::DiskCacheLevelStorage( + StorageConfig const& storageCfg, TypedVec const& slotCountList, std::string directory) + : mDirectory(std::move(directory)) +{ + TLLM_CHECK_DEBUG_WITH_INFO(slotCountList.size() == storageCfg.slotDescList.size(), + "DiskCacheLevelStorage: slotCountList and slotDescList must have the same length"); + for (PoolGroupIndex pgIdx{0}; pgIdx < storageCfg.slotDescList.size(); ++pgIdx) + { + mPoolGroups.push_back(std::make_unique( + slotCountList[pgIdx], storageCfg.slotDescList[pgIdx].slotSizeList(), mDirectory)); + } +} + +// --------------------------------------------------------------------------- +// Factory +// --------------------------------------------------------------------------- + +std::unique_ptr createCacheLevelStorage(CacheTierConfig const& tierCfg, + StorageConfig const& storageCfg, TypedVec const& slotCountList) +{ + return std::visit( + [&](auto const& cfg) -> std::unique_ptr + { + using T = std::decay_t; + if constexpr (std::is_same_v) + { + // Compute phys mem size (granularity) from quota. + constexpr size_t kPageSize = 2ULL << 20; + // Guard std::log2(0) (UB when cast to int) for quotas below 1 GiB, + // where the integer ratio is 0 and the exponent floor is used. + size_t const ratio = cfg.quota / (kPageSize * 512); + int const exponent = ratio == 0 ? 0 : std::min(4, std::max(0, static_cast(std::log2(ratio)))); + size_t physMemSize = kPageSize << exponent; + return std::make_unique(storageCfg, slotCountList, physMemSize); + } + else if constexpr (std::is_same_v) + { + return std::make_unique(storageCfg, slotCountList); + } + else + { + return std::make_unique(storageCfg, slotCountList, cfg.path); + } + }, + tierCfg); +} + +// --------------------------------------------------------------------------- +// CacheLevelStorage helper: grainsToSlots +// Distribute grains among pools within a pool group. +// Returns {num_slots, grains_consumed}. +// Mirrors Python CacheLevelStorage._grains_to_slots. +// --------------------------------------------------------------------------- + +std::pair CacheLevelStorage::grainsToSlots( + size_t pgGrains, TypedVec const& slotSizeList, size_t granularity) +{ + TypedVec minPoolGrains(slotSizeList.size()); + for (PoolIndex poolIdx{0}; poolIdx < slotSizeList.size(); ++poolIdx) + { + minPoolGrains[poolIdx] = divUp(slotSizeList[poolIdx], granularity); + } + + size_t minTotal = 0; + for (auto g : minPoolGrains) + { + minTotal += g; + } + if (pgGrains < minTotal) + return {0, 0}; + + SlotCount numSlots{std::numeric_limits::max()}; + size_t remainingPgGrains = pgGrains; + + // Sort pools by slot size ascending. + std::vector poolOrder; + poolOrder.reserve(slotSizeList.stdSize()); + for (PoolIndex poolIdx{0}; poolIdx < slotSizeList.size(); ++poolIdx) + { + poolOrder.push_back(poolIdx); + } + std::sort(poolOrder.begin(), poolOrder.end(), + [&](PoolIndex a, PoolIndex b) { return slotSizeList[a] < slotSizeList[b]; }); + + for (size_t j = 0; j < poolOrder.size(); ++j) + { + PoolIndex const poolIdx = poolOrder[j]; + size_t slotSz = slotSizeList[poolIdx]; + size_t poolSzSum = 0; + for (size_t k = j; k < poolOrder.size(); ++k) + poolSzSum += slotSizeList[poolOrder[k]]; + double poolFrac = (poolSzSum > 0) ? static_cast(slotSz) / static_cast(poolSzSum) : 1.0; + size_t roundedGrains = static_cast(std::nearbyint(static_cast(remainingPgGrains) * poolFrac)); + size_t poolGrains = std::max(minPoolGrains[poolIdx], roundedGrains); + SlotCount const poolSlots = slotCountValueFromSize(poolGrains * granularity / slotSz); + numSlots = std::min(numSlots, poolSlots); + TLLM_CHECK_DEBUG(poolGrains <= remainingPgGrains); + remainingPgGrains -= poolGrains; + } + TLLM_CHECK_DEBUG(remainingPgGrains == 0); + TLLM_CHECK_DEBUG(numSlots > 0); + + auto slotsToGrains = [&](SlotCount slots) { return grainsForSlots(slots, slotSizeList, granularity); }; + SlotCount lo = numSlots; + SlotCount step = 1; + SlotCount hi = lo + step; + while (slotsToGrains(hi) <= pgGrains) + { + lo = hi; + step *= 2; + hi = lo + step; + } + while (lo + 1 < hi) + { + SlotCount const mid = lo + ((hi - lo) / 2); + if (slotsToGrains(mid) <= pgGrains) + { + lo = mid; + } + else + { + hi = mid; + } + } + size_t const used = slotsToGrains(lo); + TLLM_CHECK_DEBUG(used <= pgGrains); + TLLM_CHECK_DEBUG(slotsToGrains(lo + 1) > pgGrains); + return {lo, used}; +} + +// --------------------------------------------------------------------------- +// CacheLevelStorage helper: grainsForSlots +// Compute minimum grains needed for numSlots in a pool group. +// Mirrors Python CacheLevelStorage._grains_for_slots. +// --------------------------------------------------------------------------- + +size_t CacheLevelStorage::grainsForSlots( + SlotCount numSlots, TypedVec const& slotSizeList, size_t granularity) +{ + size_t total = 0; + for (auto s : slotSizeList) + total += divUp(slotCountToSizeT(numSlots) * s, granularity); + return total; +} + +// --------------------------------------------------------------------------- +// CacheLevelStorage::ratioToSlotCountList (static) +// Mirrors Python CacheLevelStorage.ratio_to_slot_count_list. +// --------------------------------------------------------------------------- + +TypedVec CacheLevelStorage::ratioToSlotCountList(size_t quota, + TypedVec> const& sizeLists, + TypedVec const& ratioList, size_t granularity, + TypedVec const& minSlots) +{ + PoolGroupIndex numPg = sizeLists.size(); + TLLM_CHECK_DEBUG(ratioList.size() == numPg); + TLLM_CHECK_DEBUG_WITH_INFO(std::all_of(ratioList.begin(), ratioList.end(), [](auto x) { return x > 0; }), + "ratioToSlotCountList: all ratios must be positive"); + TLLM_CHECK_DEBUG(quota % granularity == 0); + size_t totalGrains = quota / granularity; + if (TLLM_UNLIKELY(gDebug)) + { + [[maybe_unused]] size_t minGrains = 0; + for (auto const& sizes : sizeLists) + minGrains += toSizeT(sizes.size()); + TLLM_CHECK_WITH_INFO(totalGrains >= minGrains, + "ratioToSlotCountList: insufficient total grains for at least 1 slot per pool group"); + } + + TypedVec slotCntList(numPg, 0); + size_t remainingGrains = totalGrains; + std::vector activePgs(toSizeT(numPg)); + std::iota(activePgs.begin(), activePgs.end(), PoolGroupIndex{0}); + + // Iteratively peel off constrained PGs until all active PGs are + // unconstrained: + // 1. Distribute remaining quota among active PGs by ratio. + // 2. Any PG with slots <= min_slots is constrained — pin it to + // min_slots and subtract its grains from the budget. + // 3. Repeat with the remaining PGs and re-normalized ratios. + // Each iteration removes at least one PG, so this terminates. + while (!activePgs.empty()) + { + // Distribute remainingGrains among active PGs by ratio. + size_t nActive = activePgs.size(); + std::vector activeRatio(nActive); + for (size_t i = 0; i < nActive; ++i) + activeRatio[i] = ratioList[activePgs[i]]; + + std::vector slotsForActive(nActive, 0); + std::vector grainsForActive(nActive, 0); + size_t budget = remainingGrains; + + // Sort indices by ratio ascending. + std::vector idxLst(nActive); + std::iota(idxLst.begin(), idxLst.end(), size_t{0}); + std::sort(idxLst.begin(), idxLst.end(), [&](size_t a, size_t b) { return activeRatio[a] < activeRatio[b]; }); + + for (size_t i = 0; i < idxLst.size(); ++i) + { + size_t idx = idxLst[i]; + double ratioSum = 0.0; + for (size_t j = i; j < idxLst.size(); ++j) + ratioSum += static_cast(activeRatio[idxLst[j]]); + double pct = (ratioSum > 0.0) ? static_cast(activeRatio[idx]) / ratioSum : 1.0; + auto [slots, used] = grainsToSlots(static_cast(std::nearbyint(static_cast(budget) * pct)), + sizeLists[activePgs[idx]], granularity); + slotsForActive[idx] = slots; + grainsForActive[idx] = used; + TLLM_CHECK_DEBUG(used <= budget); + budget -= used; + } + + // Identify constrained PGs (slots <= min_slots). + std::vector constrained; + std::vector unconstrained; + for (size_t idx = 0; idx < nActive; ++idx) + { + PoolGroupIndex pgIdx = activePgs[idx]; + SlotCount const minSlotCount = minSlots[pgIdx]; + if (slotsForActive[idx] <= minSlotCount) + constrained.push_back(idx); + else + unconstrained.push_back(idx); + } + + if (constrained.empty()) + { + // All active PGs are unconstrained — accept their allocations. + for (size_t idx = 0; idx < nActive; ++idx) + slotCntList[activePgs[idx]] = slotsForActive[idx]; + break; + } + + // Pin constrained PGs to min_slots and subtract from budget. + for (size_t idx : constrained) + { + PoolGroupIndex pgIdx = activePgs[idx]; + SlotCount const minSlotCount = minSlots[pgIdx]; + size_t minGrains = grainsForSlots(minSlotCount, sizeLists[pgIdx], granularity); + auto [slots, used] = grainsToSlots(minGrains, sizeLists[pgIdx], granularity); + slotCntList[pgIdx] = slots; + TLLM_CHECK_DEBUG(used <= remainingGrains); + remainingGrains -= used; + } + + if (unconstrained.empty()) + { + // All PGs are constrained — nothing left to redistribute. + break; + } + + if (remainingGrains == 0) + throw std::runtime_error("Insufficient quota to satisfy min_slots constraints"); + + // Continue with unconstrained PGs only. + std::vector newActivePgs; + newActivePgs.reserve(unconstrained.size()); + for (size_t idx : unconstrained) + newActivePgs.push_back(activePgs[idx]); + activePgs = std::move(newActivePgs); + } + + // _g2s may under-count slots due to imperfect grain distribution + // across pools. Try bumping each PG's slot count while it still fits + // within the same grain budget. + for (PoolGroupIndex pgIdx{0}; pgIdx < sizeLists.size(); ++pgIdx) + { + size_t grainsNow = grainsForSlots(slotCntList[pgIdx], sizeLists[pgIdx], granularity); + while (grainsForSlots(slotCntList[pgIdx] + 1, sizeLists[pgIdx], granularity) <= grainsNow) + slotCntList[pgIdx] += 1; + } + + return slotCntList; +} + +// Instance convenience wrapper. +TypedVec CacheLevelStorage::computeSlotCountList( + TypedVec const& ratioList, TypedVec const& minSlots, + std::optional quota) const +{ + size_t q = quota.value_or(totalQuota()); + TLLM_CHECK_DEBUG(ratioList.size() == mPoolGroups.size()); + return ratioToSlotCountList(q, slotSizeLists(), ratioList, poolSizeGranularity(), minSlots); +} + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.h new file mode 100644 index 000000000000..eaea0c4a26b2 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.h @@ -0,0 +1,563 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "kv_cache_manager_v2/common.h" +#include "kv_cache_manager_v2/cudaVirtMem.h" +#include "kv_cache_manager_v2/exceptions.h" +#include "kv_cache_manager_v2/storage/config.h" +#include "kv_cache_manager_v2/utils/cudaEvent.h" +#include "kv_cache_manager_v2/utils/hostMem.h" +#include "kv_cache_manager_v2/utils/math.h" +#include "tensorrt_llm/common/assert.h" + +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// --------------------------------------------------------------------------- +// Slot — represents ownership of one allocated slot in a pool group. +// Mirrors _storage/_core.py::Slot. +// +// ready_event: completes when the slot's data is safe to use. +// - Newly allocated: completes when previous users are done. +// - After migration: completes when the copy is done. +// - Passed to release(): completes when current users are done. +// --------------------------------------------------------------------------- +struct Slot +{ + CachedCudaEvent readyEvent = CachedCudaEvent::makeNull(); + + // Mirrors Python @property slot_id: asserts valid, returns unwrapped value. + [[nodiscard]] SlotId slotId() const + { + return mSlotId.value(); + } + + [[nodiscard]] bool hasValidSlot() const noexcept + { + return mSlotId.has_value(); + } + + void setSlotId(SlotId id) noexcept + { + mSlotId = id; + } + + void resetSlot() noexcept + { + mSlotId.reset(); + } + + bool queryReady() + { + return readyEvent.queryComplete(); + } + + // Transfer slot ownership: moves slotId and readyEvent from src to this. + void setSlot(Slot& src) + { + if (hasValidSlot()) + { + throw LogicError("Slot::setSlot: already has a valid slot"); + } + mSlotId = src.mSlotId; + readyEvent = std::move(src.readyEvent); + src.mSlotId.reset(); + } + +private: + std::optional mSlotId; +}; + +// --------------------------------------------------------------------------- +// SlotAllocator — manages a fixed-capacity array of slot ids. +// Mirrors _storage/_core.py::SlotAllocator. +// --------------------------------------------------------------------------- +class SlotAllocator +{ +public: + explicit SlotAllocator(SlotCount capacity); + ~SlotAllocator(); + + [[nodiscard]] SlotCount numFreeSlots() const noexcept; + [[nodiscard]] SlotCount numOccupiedSlots() const noexcept; + + [[nodiscard]] SlotCount numSlots() const noexcept + { + return mCapacity; + } + + Slot allocate(); + std::vector allocateMultiple(SlotCount numSlots); + void release(Slot slot); + + void expand(SlotCount newNumSlots); + void prepareForShrink(SlotCount newNumSlots); + bool finishShrink(); + + [[nodiscard]] bool shrinkInProgress() const noexcept + { + return mTargetCapacity < mCapacity; + } + + [[nodiscard]] std::vector getSlotsBlockingShrink() const; + + // Read-only accessors for debug assertions (mirrors Python's direct attribute access). + [[nodiscard]] SlotCount numOverflowSlots() const noexcept + { + return slotCountValueFromSize(mOverflowSlots.size()); + } + + [[nodiscard]] SlotCount numActiveSlots() const noexcept + { + return mNumActiveSlots; + } + + [[nodiscard]] SlotCount targetCapacity() const noexcept + { + return mTargetCapacity; + } + + void synchronize(); + +private: + void scrubEvents(); + [[nodiscard]] bool check() const noexcept; + + SlotCount mCapacity; + SlotCount mTargetCapacity; + SlotCount mNumActiveSlots; + SlotCount mNumReadyRecycledSlots; + std::deque mRecycledSlots; + std::vector mOverflowSlots; + DynamicBitset mOccupiedMask; +}; + +// --------------------------------------------------------------------------- +// SlotPoolBase — abstract base for a single memory pool. +// Mirrors _storage/_core.py::SlotPoolBase. +// --------------------------------------------------------------------------- +class SlotPoolBase +{ +public: + explicit SlotPoolBase(size_t slotSize) + : mSlotSize(slotSize) + { + } + + virtual ~SlotPoolBase() = default; + + size_t slotSize() const noexcept + { + return mSlotSize; + } + + virtual SlotCount numSlots() const noexcept = 0; + + size_t numBytes() const noexcept + { + return mSlotSize * slotCountToSizeT(numSlots()); + } + + virtual void destroy() = 0; + virtual void resize(SlotCount newNumSlots) = 0; + virtual Address slotAddress(SlotId slot) const = 0; + +protected: + size_t mSlotSize; +}; + +// --------------------------------------------------------------------------- +// GpuSlotPool — GPU virtual memory pool. +// --------------------------------------------------------------------------- +class GpuSlotPool : public SlotPoolBase +{ +public: + GpuSlotPool(size_t slotSize, size_t vmSize, PooledPhysMemAllocator& physMemAllocator, SlotCount numSlots); + + SlotCount numSlots() const noexcept override; + void destroy() override; + void resize(SlotCount newNumSlots) override; + Address slotAddress(SlotId slot) const override; + + // Extend by exactly one physical memory chunk; returns new numSlots. + SlotCount extendByOnePhysMem(); + + static size_t computeNumPhysMem(size_t slotSize, SlotCount numSlots, size_t physMemSize) noexcept; + static SlotCount computeNumSlots(size_t slotSize, size_t numPhysMem, size_t physMemSize) noexcept; + +private: + VirtMem mVirtMem; +}; + +// --------------------------------------------------------------------------- +// HostSlotPool — pinned host memory pool. +// --------------------------------------------------------------------------- +class HostSlotPool : public SlotPoolBase +{ +public: + HostSlotPool(size_t slotSize, SlotCount numSlots); + + SlotCount numSlots() const noexcept override; + void destroy() override; + void resize(SlotCount newNumSlots) override; + Address slotAddress(SlotId slot) const override; + + size_t alignedSize(SlotCount numSlots) const noexcept; + +private: + HostMem mHostMem; +}; + +// --------------------------------------------------------------------------- +// DiskSlotPool — temp-file backed disk pool. +// --------------------------------------------------------------------------- +class DiskSlotPool : public SlotPoolBase +{ +public: + // directory: path under which to create the temp file. + DiskSlotPool(std::string const& directory, size_t slotSize, SlotCount numSlots); + ~DiskSlotPool() override; + + SlotCount numSlots() const noexcept override; + void destroy() override; + void resize(SlotCount newNumSlots) override; + Address slotAddress(SlotId slot) const override; + + int fd() const noexcept + { + return mFd; + } + +private: + int mFd = kBadFileDescriptor; +}; + +// --------------------------------------------------------------------------- +// PoolGroupBase — manages multiple pools with mirrored slot allocation. +// Mirrors _storage/_core.py::PoolGroupBase. +// --------------------------------------------------------------------------- +class PoolGroupBase +{ +public: + explicit PoolGroupBase(SlotCount numSlots); + virtual ~PoolGroupBase(); + + PoolIndex numPools() const noexcept + { + return mPools.size(); + } + + SlotCount numSlots() const noexcept; + + SlotCount numFreeSlots() const noexcept + { + return mSlotAllocator.numFreeSlots(); + } + + SlotAllocator& slotAllocator() noexcept + { + return mSlotAllocator; + } + + SlotAllocator const& slotAllocator() const noexcept + { + return mSlotAllocator; + } + + Slot allocate(); + std::vector allocateMultiple(SlotCount numSlots); + void release(Slot slot); + void destroy(); + void resizePools(std::optional newNumSlots = std::nullopt); + + // Addresses for all pools at a given slot id. + TypedVec slotAddress(SlotId slotId) const; + // Sizes per pool. + TypedVec slotSize() const; + + // Total bytes across all pools for one slot. + size_t numBytes() const noexcept + { + size_t total = 0; + for (auto const& pool : mPools) + total += pool->numBytes(); + return total; + } + + // Total bytes across all pools, rounding each pool's bytes up to granularity. + // Mirrors Python total_quota: sum(round_up(p.num_bytes, granularity) for p in pg._pools). + size_t roundedNumBytes(size_t granularity) const noexcept + { + size_t total = 0; + for (auto const& pool : mPools) + total += roundUp(pool->numBytes(), granularity); + return total; + } + +protected: + // Mirrors Python _get_num_slots_from_pools: min(p.num_slots for p in pools). + SlotCount getNumSlotsFromPools() const noexcept; + + SlotAllocator mSlotAllocator; + TypedVec> mPools; + bool mDestroyed = false; +}; + +// --------------------------------------------------------------------------- +// GpuPoolGroup / HostPoolGroup / DiskPoolGroup +// --------------------------------------------------------------------------- +class GpuPoolGroup : public PoolGroupBase +{ +public: + GpuPoolGroup( + SlotCount numSlots, TypedVec const& slotSizeList, PooledPhysMemAllocator& physMemAllocator); +}; + +class HostPoolGroup : public PoolGroupBase +{ +public: + HostPoolGroup(SlotCount numSlots, TypedVec const& slotSizeList); +}; + +class DiskPoolGroup : public PoolGroupBase +{ +public: + DiskPoolGroup(SlotCount numSlots, TypedVec const& slotSizeList, std::string const& directory); +}; + +// --------------------------------------------------------------------------- +// CacheLevelStorage — manages all pool groups for one cache tier. +// Mirrors _storage/_core.py::CacheLevelStorage. +// --------------------------------------------------------------------------- +class CacheLevelStorage +{ +public: + virtual ~CacheLevelStorage() = default; + + virtual CacheTier cacheTier() const noexcept = 0; + + PoolGroupIndex numPoolGroups() const noexcept + { + return mPoolGroups.size(); + } + + std::vector allocateMultiple(PoolGroupIndex pgIdx, SlotCount numSlots); + void release(PoolGroupIndex pgIdx, Slot slot); + SlotCount numFreeSlots(PoolGroupIndex pgIdx) const; + TypedVec slotAddress(PoolGroupIndex pgIdx, SlotId slotId) const; + + // Additional accessors used by StorageManager and KvCacheManager. + virtual void destroy() + { + for (auto& pg : mPoolGroups) + pg->destroy(); + } + + PoolIndex numPools(PoolGroupIndex pgIdx) const + { + return mPoolGroups.at(pgIdx)->numPools(); + } + + SlotCount numSlots(PoolGroupIndex pgIdx) const + { + return mPoolGroups.at(pgIdx)->numSlots(); + } + + TypedVec slotSize(PoolGroupIndex pgIdx) const + { + auto szl = mPoolGroups.at(pgIdx)->slotSize(); + TypedVec ret; + ret.reserve(szl.size()); + for (auto sz : szl) + ret.push_back(sz); + return ret; + } + + PoolGroupBase& poolGroup(PoolGroupIndex pgIdx) + { + return *mPoolGroups.at(pgIdx); + } + + MemAddress getBaseAddress(PoolGroupIndex pgIdx, PoolIndex poolIdx, SlotId slotId) const + { + return std::get(mPoolGroups.at(pgIdx)->slotAddress(slotId).at(poolIdx)); + } + + size_t totalQuota() const noexcept + { + size_t granularity = poolSizeGranularity(); + size_t total = 0; + for (auto const& pg : mPoolGroups) + total += pg->roundedNumBytes(granularity); + return total; + } + + // Returns numSlots() per pool group. + TypedVec slotCountList() const + { + TypedVec ret; + ret.reserve(mPoolGroups.size()); + for (auto const& pg : mPoolGroups) + ret.push_back(pg->numSlots()); + return ret; + } + + TypedVec> slotSizeLists() const + { + TypedVec> ret; + ret.reserve(mPoolGroups.size()); + for (auto const& pg : mPoolGroups) + { + auto szl = pg->slotSize(); + TypedVec sizes; + sizes.reserve(szl.size()); + for (auto sz : szl) + sizes.push_back(sz); + ret.push_back(std::move(sizes)); + } + return ret; + } + + // Current ratio list: proportion of bytes per pool group. + TypedVec ratioList() const + { + TypedVec ret(mPoolGroups.size(), 0.f); + float total = 0.f; + for (PoolGroupIndex pgIdx{0}; pgIdx < mPoolGroups.size(); ++pgIdx) + { + auto sz = static_cast(mPoolGroups[pgIdx]->numBytes()); + ret[pgIdx] = sz; + total += sz; + } + TLLM_CHECK_DEBUG(total > 0.f); + for (auto& r : ret) + r /= total; + return ret; + } + + virtual size_t poolSizeGranularity() const noexcept + { + return size_t{2} << 20; + } + + // Compute slot counts per pool group for a given ratio, min_slots, and optional quota. + // Instance convenience method — delegates to the static version below. + TypedVec computeSlotCountList(TypedVec const& ratioList, + TypedVec const& minSlots, std::optional quota = std::nullopt) const; + + // Static version: compute slot counts from quota, slot size lists, ratio, granularity, and min_slots. + // Mirrors Python CacheLevelStorage.ratio_to_slot_count_list. + static TypedVec ratioToSlotCountList(size_t quota, + TypedVec> const& slotSizeLists, + TypedVec const& ratioList, size_t granularity, + TypedVec const& minSlots); + + // Distribute grains among pools within a pool group. + // Returns {num_slots, grains_consumed}. + static std::pair grainsToSlots( + size_t pgGrains, TypedVec const& slotSizeList, size_t granularity); + + // Compute minimum grains needed for numSlots in a pool group. + static size_t grainsForSlots( + SlotCount numSlots, TypedVec const& slotSizeList, size_t granularity); + + virtual void postResize() {} + +protected: + TypedVec> mPoolGroups; +}; + +class GpuCacheLevelStorage : public CacheLevelStorage +{ +public: + GpuCacheLevelStorage( + StorageConfig const& storageCfg, TypedVec const& slotCountList, size_t physMemSize); + + CacheTier cacheTier() const noexcept override + { + return CacheTier::GPU_MEM; + } + + size_t poolSizeGranularity() const noexcept override + { + return mPhysMemAllocator->physMemSize(); + } + + void postResize() override + { + CacheLevelStorage::postResize(); + mPhysMemAllocator->clear(); + } + + void destroy() override + { + CacheLevelStorage::destroy(); + mPhysMemAllocator->clear(); + } + +private: + std::unique_ptr mPhysMemAllocator; +}; + +class HostCacheLevelStorage : public CacheLevelStorage +{ +public: + HostCacheLevelStorage(StorageConfig const& storageCfg, TypedVec const& slotCountList); + + CacheTier cacheTier() const noexcept override + { + return CacheTier::HOST_MEM; + } + + size_t poolSizeGranularity() const noexcept override + { + return HostMem::kAlignment; + } +}; + +class DiskCacheLevelStorage : public CacheLevelStorage +{ +public: + DiskCacheLevelStorage(StorageConfig const& storageCfg, TypedVec const& slotCountList, + std::string directory); + + CacheTier cacheTier() const noexcept override + { + return CacheTier::DISK; + } + + std::string const& directory() const noexcept + { + return mDirectory; + } + +private: + std::string mDirectory; +}; + +// Factory: create appropriate CacheLevelStorage for a given tier config. +std::unique_ptr createCacheLevelStorage(CacheTierConfig const& tierCfg, + StorageConfig const& storageCfg, TypedVec const& slotCountList); + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp new file mode 100644 index 000000000000..724a33e308b0 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp @@ -0,0 +1,1292 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "kv_cache_manager_v2/storageManager.h" +#include "kv_cache_manager_v2/common.h" +#include "kv_cache_manager_v2/copyEngine.h" +#include "kv_cache_manager_v2/exceptions.h" +#include "kv_cache_manager_v2/page.h" +#include "kv_cache_manager_v2/utils/hostMem.h" +#include "kv_cache_manager_v2/utils/math.h" +#include "tensorrt_llm/common/logger.h" + +#include "tensorrt_llm/common/assert.h" +#include +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// --------------------------------------------------------------------------- +// CacheLevelManager +// --------------------------------------------------------------------------- + +CacheLevelManager::CacheLevelManager(TypedVec const& lifeCycleGrouping, CacheLevel cl, + CacheTierConfig const& tierConfig, StorageConfig const& storageConfig, + TypedVec const& slotCountList) + : cacheLevel(cl) + , cacheTier(CacheTier(tierConfig.index())) + , controller(lifeCycleGrouping, cl) +{ + storage = createCacheLevelStorage(tierConfig, storageConfig, slotCountList); +} + +size_t CacheLevelManager::cacheTierGranularity(CacheTier tier, size_t quota) +{ + switch (tier) + { + case CacheTier::GPU_MEM: + { + constexpr size_t kPageSize = 2ULL << 20; + return kPageSize << std::min(4, std::max(0, static_cast(std::log2(quota / (kPageSize * 512))))); + } + case CacheTier::HOST_MEM: return HostMem::kAlignment; // 4 KiB + case CacheTier::DISK: return size_t{2} << 20; // DiskCacheLevelStorage::POOL_SIZE_GRANULARITY + default: throw std::invalid_argument("Invalid cache tier"); + } +} + +// --------------------------------------------------------------------------- +// StorageManager constructor helpers +// --------------------------------------------------------------------------- + +namespace +{ + +// Compute the slot-to-page-indices scale factors. +// For each (lcId, poolIdx), scale = numBuffersInCoalescedSlot. +// Python: _slot_to_page_indices[lc_id][pool_idx] = numBuffers +TypedVec> computeSlotToPageIndices(StorageConfig const& config) +{ + LifeCycleId numLc = config.numLifeCycles(); + TypedVec> result(numLc); + + auto const& slotDescList = config.slotDescList; + auto const& grouping = config.lifeCycleGrouping(); + + for (LifeCycleId lcId{0}; lcId < result.size(); ++lcId) + { + PoolGroupIndex pgIdx = grouping[lcId]; + SlotDesc const& sd = slotDescList.at(pgIdx); + // Find the variant that corresponds to this lifecycle. + for (auto const& variant : sd.variants) + { + if (variant.lifeCycleId == lcId) + { + // Each coalesced buffer contributes its numBuffers as the scale. + result[lcId].reserve(variant.coalescedBuffers.size()); + for (auto const& cb : variant.coalescedBuffers) + result[lcId].push_back(cb.numBuffers()); + break; + } + } + if (result[lcId].empty()) + result[lcId].push_back(1); // fallback + } + return result; +} + +} // namespace + +// --------------------------------------------------------------------------- +// StorageManager +// --------------------------------------------------------------------------- + +StorageManager::StorageManager(LifeCycleRegistry const& lifeCycles, StorageConfig const& config, int tokensPerBlock, + std::optional swaScratchReuse, std::optional const& typicalBatch, + std::vector const& constraints, std::optional> const& initialPoolRatio, + std::shared_ptr eventSink, float maxUtilForResume) + : mLifeCycles(lifeCycles) + , mEventSink(std::move(eventSink)) + , mStorageConfig(config) + , mSwaScratchReuse(std::move(swaScratchReuse)) +{ + mLifeCycleGrouping = config.lifeCycleGrouping(); + mLayerToLifeCycleIds = config.layerToLifeCycleIds(); + mSlotToPageIndices = computeSlotToPageIndices(config); + mBufferAttr = config.bufferAttributes(); + mSlotDescList = config.slotDescList; + + // Compute layer attributes and slot utilization fractions for scratch support. + mLayerAttributes = config.layerAttributes(); + mSlotUtilFracMax.resize(lifeCycles.size(), Rational{0, 1}); + for (auto const& [layerId, layerAttr] : mLayerAttributes) + { + LifeCycleId const lcIdx = layerAttr.lifeCycleId; + if (layerAttr.slotUtilFracMax > mSlotUtilFracMax[lcIdx]) + { + mSlotUtilFracMax[lcIdx] = layerAttr.slotUtilFracMax; + } + } + + TLLM_CHECK_DEBUG(std::all_of(mLifeCycleGrouping.begin(), mLifeCycleGrouping.end(), + [this](PoolGroupIndex pg) { return pg < numPoolGroups(); })); + TLLM_CHECK_DEBUG(numPoolGroups() + == PoolGroupIndex{ + static_cast(std::set(mLifeCycleGrouping.begin(), mLifeCycleGrouping.end()).size())}); + + // Build one CacheLevelManager per tier. + TLLM_CHECK_DEBUG(!config.cacheTiers.empty()); + TLLM_CHECK_DEBUG_WITH_INFO( + std::holds_alternative(config.cacheTiers[kGpuLevel]), "First cache tier must be GPU"); + + // Compute slot size lists for all pool groups. + TypedVec> slotSizeLists; + slotSizeLists.reserve(mSlotDescList.size()); + for (auto const& sd : mSlotDescList) + { + slotSizeLists.push_back(sd.slotSizeList()); + } + + size_t gpuQuota = cacheTierQuota(config.cacheTiers[kGpuLevel]); + size_t gpuGranularity = CacheLevelManager::cacheTierGranularity(CacheTier::GPU_MEM, gpuQuota); + + // Constraints stay feasibility floors even under an explicit initial pool + // ratio (a share below what a declared batch needs is clamped up), and the + // floors are scaled by 1/maxUtilForResume because KvCache::resume rejects any + // pool group above that utilization. Mirrors PR#16269 on the Python side. + mMinSlots = computeMinSlotsFromConstraints(constraints, tokensPerBlock, mSwaScratchReuse, maxUtilForResume); + + // Compute init_ratio from explicit config, typical_batch, constraints, or fallback. + TypedVec initRatio; + if (initialPoolRatio.has_value()) + { + if (initialPoolRatio->size() != toSizeT(numPoolGroups())) + { + throw std::invalid_argument("initial_pool_ratio length must match number of pool groups (" + + std::to_string(toSizeT(numPoolGroups())) + "), got " + std::to_string(initialPoolRatio->size())); + } + if (std::any_of(initialPoolRatio->begin(), initialPoolRatio->end(), [](float ratio) { return ratio <= 0.0F; })) + { + throw std::invalid_argument("initial_pool_ratio values must be positive"); + } + + constexpr double kExpectedRatioSum = 1.0; + constexpr double kRatioSumTolerance = 1e-6; + double const ratioSum = std::accumulate(initialPoolRatio->begin(), initialPoolRatio->end(), 0.0); + if (!std::isfinite(ratioSum) || std::abs(ratioSum - kExpectedRatioSum) > kRatioSumTolerance) + { + throw std::invalid_argument("initial_pool_ratio values must sum to 1.0"); + } + initRatio = TypedVec(*initialPoolRatio); + } + else if (typicalBatch.has_value()) + { + initRatio = ratioFromBatch(*typicalBatch, tokensPerBlock, mSwaScratchReuse, gpuGranularity); + } + else if (!constraints.empty()) + { + // Use the constraint slot counts as the ratio basis. + auto minBytes = slotsToBytes(mMinSlots, gpuGranularity); + initRatio = normalizeToRatio(minBytes); + } + else + { + // Fallback: average history length 2048. + BatchDesc fallback; + fallback.kvCaches.push_back(KVCacheDesc{2049, 2048}); + initRatio = ratioFromBatch(fallback, tokensPerBlock, mSwaScratchReuse, gpuGranularity); + } + + mLevels.reserve(config.cacheTiers.size()); + for (CacheLevel level{0}; level < config.cacheTiers.size(); ++level) + { + auto slotCountList = computeSlotCountForLevel(config.cacheTiers[level], slotSizeLists, initRatio); + mLevels.emplace_back(mLifeCycleGrouping, level, config.cacheTiers[level], config, slotCountList); + } + + TLLM_CHECK_DEBUG(mLevels.empty() + || numPoolGroups() + == getUniformAttribute(mLevels, [](auto const& lvl) { return lvl.storage->numPoolGroups(); })); +} + +StorageManager::~StorageManager() +{ + destroy(); +} + +void StorageManager::destroy() +{ + for (auto& lvl : mLevels) + { + TLLM_CHECK_DEBUG(lvl.storage); + lvl.storage->destroy(); + } + mLevels.clear(); +} + +// --------------------------------------------------------------------------- +// newSlots +// --------------------------------------------------------------------------- + +TypedVec> StorageManager::newSlots(CacheLevel level, + TypedVec const& numSlotsPerLc, MigrationRecorder const& migrationRecorder, + DropRecorder const& dropRecorder) +{ + TLLM_CHECK_DEBUG(numSlotsPerLc.size() == numLifeCycles()); + auto& storage = *mLevels.at(level).storage; + + // Aggregate by pool group. + TypedVec pgNumSlots(numPoolGroups(), 0); + for (LifeCycleId lcId{0}; lcId < numSlotsPerLc.size(); ++lcId) + { + SlotCount const numSlots = numSlotsPerLc[lcId]; + if (numSlots < 0) + { + throw LogicError("StorageManager::newSlots: slot count must be non-negative"); + } + pgNumSlots[mLifeCycleGrouping[lcId]] += numSlots; + } + + // Prepare free slots if needed. + bool needMore = false; + for (PoolGroupIndex pgIdx{0}; pgIdx < pgNumSlots.size(); ++pgIdx) + { + if (pgNumSlots[pgIdx] > storage.numFreeSlots(pgIdx)) + { + needMore = true; + break; + } + } + + if (needMore) + { + prepareFreeSlots(level, pgNumSlots, migrationRecorder, dropRecorder); + } + + // A14: post-condition — free-slot counts satisfy requirements. + for (PoolGroupIndex pgIdx{0}; pgIdx < pgNumSlots.size(); ++pgIdx) + { + TLLM_CHECK_DEBUG_WITH_INFO(pgNumSlots[pgIdx] <= storage.numFreeSlots(pgIdx), + "Free slot count does not satisfy requirement after prepareFreeSlots"); + } + + // Allocate. + TypedVec> ret(numLifeCycles()); + try + { + for (LifeCycleId lcId{0}; lcId < ret.size(); ++lcId) + { + PoolGroupIndex pg = mLifeCycleGrouping[lcId]; + ret[lcId] = storage.allocateMultiple(pg, numSlotsPerLc[lcId]); + } + } + catch (...) + { + for (LifeCycleId lcId{0}; lcId < ret.size(); ++lcId) + { + PoolGroupIndex pg = mLifeCycleGrouping[lcId]; + for (auto& s : ret[lcId]) + storage.release(pg, std::move(s)); + } + throw; + } + return ret; +} + +TypedVec> StorageManager::newGpuSlots( + TypedVec const& numSlotsPerLc, MigrationRecorder const& migrationRecorder, + DropRecorder const& dropRecorder) +{ + return newSlots(kGpuLevel, numSlotsPerLc, migrationRecorder, dropRecorder); +} + +std::vector StorageManager::newSlotsForPoolGroup(CacheLevel level, PoolGroupIndex pgIdx, SlotCount numSlots, + MigrationRecorder const& migrationRecorder, DropRecorder const& dropRecorder) +{ + if (numSlots < 0) + { + throw LogicError("StorageManager::newSlotsForPoolGroup: numSlots must be non-negative"); + } + auto& storage = *mLevels.at(level).storage; + if (numSlots > storage.numFreeSlots(pgIdx)) + { + TypedVec requirements(numPoolGroups(), 0); + requirements.at(pgIdx) = numSlots; + prepareFreeSlots(level, requirements, migrationRecorder, dropRecorder); + } + TLLM_CHECK_DEBUG(numSlots <= storage.numFreeSlots(pgIdx)); + return storage.allocateMultiple(pgIdx, numSlots); +} + +Address StorageManager::slotAddress(CacheLevel level, PoolGroupIndex pgIdx, SlotId slotId, PoolIndex poolIdx) const +{ + return mLevels.at(level).storage->slotAddress(pgIdx, slotId).at(poolIdx); +} + +CacheTier StorageManager::cacheTier(CacheLevel level) const +{ + return mLevels.at(level).cacheTier; +} + +void StorageManager::releaseSlot(LifeCycleId lc, CacheLevel level, Slot slot) +{ + PoolGroupIndex pg = mLifeCycleGrouping.at(lc); + mLevels.at(level).storage->release(pg, std::move(slot)); +} + +// --------------------------------------------------------------------------- +// isEvictable +// --------------------------------------------------------------------------- + +bool StorageManager::isEvictable(Page const& page, std::optional level) const noexcept +{ + PageStatus s = page.status(); + CacheLevel lvl = level.value_or(page.cacheLevel); + return (s == PageStatus::DROPPABLE && page.isCommitted()) || (s == PageStatus::HELD && lvl < numCacheLevels() - 1); +} + +// --------------------------------------------------------------------------- +// scheduleForEviction / excludeFromEviction +// --------------------------------------------------------------------------- + +void StorageManager::scheduleForEviction(Page& page) +{ + if (isEvictable(page)) + mLevels.at(page.cacheLevel).controller.scheduleForEviction(page); +} + +void StorageManager::excludeFromEviction(Page& page) +{ + TLLM_CHECK_DEBUG(page.nodeRef.has_value()); + mLevels.at(page.cacheLevel).controller.remove(*page.nodeRef); +} + +// --------------------------------------------------------------------------- +// prepareFreeSlots +// --------------------------------------------------------------------------- + +void StorageManager::prepareFreeSlots(CacheLevel level, TypedVec const& requirements, + MigrationRecorder const& migrationRecorder, DropRecorder const& dropRecorder) +{ + TypedVec> goals(numCacheLevels()); + for (CacheLevel lvl{0}; lvl < goals.size(); ++lvl) + { + goals[lvl].resize(numPoolGroups(), 0); + } + for (PoolGroupIndex pgIdx{0}; pgIdx < requirements.size(); ++pgIdx) + { + goals.at(level).at(pgIdx) = requirements.at(pgIdx); + } + + TypedVec>> fallenPages(numPoolGroups()); + _prepareFreeSlots(goals, level, fallenPages, migrationRecorder, dropRecorder); +} + +void StorageManager::forceEvict( + CacheLevel level, TypedVec const& minNumPages, DropRecorder const& dropRecorder) +{ + auto evicted = mLevels.at(level).controller.evict(minNumPages); + + if (isLastLevel(level)) + { + // Last level: all evicted pages must be DROPPABLE (they get dropped, not migrated). + for (auto const& pages : evicted) + { + for (auto const& page : pages) + { + TLLM_CHECK_DEBUG_WITH_INFO(page->status() == PageStatus::DROPPABLE, "Corrupted eviction controller"); + } + } + if (dropRecorder) + { + for (auto const& pages : evicted) + { + if (!pages.empty()) + { + dropRecorder(pages, level); + } + } + } + return; + } + + TypedVec> goals(numCacheLevels()); + for (CacheLevel lvl{0}; lvl < goals.size(); ++lvl) + { + goals[lvl].resize(numPoolGroups(), 0); + } + CacheLevel nextLvl = level + 1; + + TypedVec>> fallen(numPoolGroups()); + for (PoolGroupIndex pgIdx{0}; pgIdx < fallen.size(); ++pgIdx) + { + for (auto& sp : evicted.at(pgIdx)) + fallen.at(pgIdx).push_back(sp); + } + _prepareFreeSlots(goals, nextLvl, fallen, MigrationRecorder{}, dropRecorder); +} + +// --------------------------------------------------------------------------- +// _prepareFreeSlots (recursive) +// --------------------------------------------------------------------------- + +void StorageManager::_prepareFreeSlots(TypedVec>& goals, + CacheLevel lvlId, TypedVec>>& fallenPages, + MigrationRecorder const& migrationRecorder, DropRecorder const& dropRecorder) +{ + // A7: goals dimensions must match [numCacheLevels][numPoolGroups]. + if (TLLM_UNLIKELY(gDebug)) + { + TLLM_CHECK_WITH_INFO(goals.size() == numCacheLevels(), "goals.rows must equal numCacheLevels"); + TLLM_CHECK_DEBUG_WITH_INFO( + std::all_of(goals.begin(), goals.end(), [this](auto const& row) { return row.size() == numPoolGroups(); }), + "goals.cols must equal numPoolGroups"); + } + + // A8: all fallen pages must come from upper cache levels (cache_level < lvlId). + TLLM_CHECK_DEBUG_WITH_INFO(std::all_of(fallenPages.begin(), fallenPages.end(), + [lvlId](auto const& pages) { + return std::all_of(pages.begin(), pages.end(), + [lvlId](auto const& p) { return p->cacheLevel < lvlId; }); + }), + "Fallen pages must come from upper cache levels"); + + auto& lvl = mLevels.at(lvlId); + auto& storage = *lvl.storage; + auto& ctrl = lvl.controller; + bool isLast = isLastLevel(lvlId); + + TypedVec numToEvict(numPoolGroups(), 0); + TypedVec>> heldPages(numPoolGroups()); + + for (PoolGroupIndex pgIdx{0}; pgIdx < numToEvict.size(); ++pgIdx) + { + SlotCount const goal = goals.at(lvlId).at(pgIdx); + SlotCount const fallen = slotCountValueFromSize(fallenPages.at(pgIdx).size()); + SlotCount const oldFree = storage.numFreeSlots(pgIdx); + SlotCount const evictableCount = ctrl.numEvictablePages(pgIdx); + SlotCount const required = goal + fallen; + SlotCount const shortage = required > oldFree ? required - oldFree : 0; + numToEvict.at(pgIdx) = std::min(shortage, evictableCount); + + SlotCount fallenHeld = 0; + if (isLast) + { + // Separate held pages from fallen_pages (mirrors Python's remove_if). + auto& fp = fallenPages.at(pgIdx); + heldPages.at(pgIdx) = stealIf(fp, [](SharedPtr const& p) { return p->status() == PageStatus::HELD; }); + fallenHeld = slotCountValueFromSize(heldPages.at(pgIdx).size()); + + if (fallenHeld > oldFree + evictableCount) + throw OutOfPagesError( + "Too many held pages falling to last-level cache for group " + std::to_string(pgIdx.value())); + } + + if (oldFree + evictableCount < fallenHeld + goal) + throw OutOfPagesError("Impossible to meet free-slot goal " + std::to_string(goal) + " for group " + + std::to_string(pgIdx.value())); + } + + auto evicted = ctrl.evict(numToEvict); + TypedVec>> acceptedPages(numPoolGroups()); + + if (isLast) + { + for (PoolGroupIndex pgIdx{0}; pgIdx < evicted.size(); ++pgIdx) + { + auto& ev = evicted.at(pgIdx); + SlotCount const oldFree = storage.numFreeSlots(pgIdx); + SlotCount const numEvicted = slotCountValueFromSize(ev.size()); + // A9: all evicted pages at last level must be DROPPABLE. + TLLM_CHECK_DEBUG_WITH_INFO( + std::all_of(ev.begin(), ev.end(), [](auto const& p) { return p->status() == PageStatus::DROPPABLE; }), + "Evicted page at last level must be DROPPABLE"); + // Drop droppable evicted pages (GC). + if (dropRecorder && !ev.empty()) + { + dropRecorder(ev, lvlId); + } + ev.clear(); + SlotCount const newFree = storage.numFreeSlots(pgIdx); + TLLM_CHECK_DEBUG(newFree >= numEvicted + oldFree); + + // A10: held_pages count must not exceed new_free. + TLLM_CHECK_DEBUG_WITH_INFO(slotCountValueFromSize(heldPages.at(pgIdx).size()) <= newFree, + "held_pages count exceeds new free slot count"); + + // Add held pages from upper levels. + auto& hp = heldPages.at(pgIdx); + auto& fp = fallenPages.at(pgIdx); + fp.insert(fp.end(), hp.begin(), hp.end()); + hp.clear(); + + SlotCount const goal = goals.at(lvlId).at(pgIdx); + SlotCount const freeAfterGoal = newFree > goal ? newFree - goal : 0; + SlotCount const numAccepted = std::min(freeAfterGoal, slotCountValueFromSize(fp.size())); + if (numAccepted > 0) + { + acceptedPages.at(pgIdx).assign(fp.end() - static_cast(numAccepted), fp.end()); + } + fp.clear(); + } + } + else + { + // A12: no held pages at non-last level. + TLLM_CHECK_DEBUG_WITH_INFO( + std::all_of(heldPages.begin(), heldPages.end(), [](auto const& hp) { return hp.empty(); }), + "held_pages must be empty at non-last level"); + + CacheLevel nextLvl = lvlId + 1; + for (PoolGroupIndex pgIdx{0}; pgIdx < evicted.size(); ++pgIdx) + { + auto& ev = evicted.at(pgIdx); + SlotCount const oldFree = storage.numFreeSlots(pgIdx); + SlotCount const numEvicted = slotCountValueFromSize(ev.size()); + auto& fp = fallenPages.at(pgIdx); + fp.insert(fp.begin(), ev.begin(), ev.end()); // prepend evicted to fallen (preserving order) + ev.clear(); + + SlotCount const goal = goals.at(lvlId).at(pgIdx); + SlotCount const availableAfterGoal = oldFree + numEvicted > goal ? oldFree + numEvicted - goal : 0; + SlotCount const numAccepted = std::min(availableAfterGoal, slotCountValueFromSize(fp.size())); + if (numAccepted > 0) + { + acceptedPages.at(pgIdx).assign(fp.end() - static_cast(numAccepted), fp.end()); + fp.erase(fp.end() - static_cast(numAccepted), fp.end()); + } + } + _prepareFreeSlots(goals, nextLvl, fallenPages, migrationRecorder, dropRecorder); + } + + // A13: all fallen pages must have been consumed. + TLLM_CHECK_DEBUG_WITH_INFO( + std::all_of(fallenPages.begin(), fallenPages.end(), [](auto const& fp) { return fp.empty(); }), + "All fallen pages must be consumed after level loop"); + + // Migrate accepted pages into lvlId. + for (PoolGroupIndex pgIdx{0}; pgIdx < acceptedPages.size(); ++pgIdx) + { + // Group by source level (mirrors Python's partition()). + auto bySrcLevel = partition(acceptedPages.at(pgIdx), [](SharedPtr const& p) { return p->cacheLevel; }); + + for (auto& [srcLvl, pages] : bySrcLevel) + { + _batchedMigrate(pgIdx, lvlId, srcLvl, pages, /*updateSrc=*/true, migrationRecorder); + for (auto const& p : pages) + { + if (isLast && p->status() == PageStatus::HELD) + continue; + lvl.controller.scheduleForEviction(*p); + } + } + } +} + +// --------------------------------------------------------------------------- +// _batchedMigrate +// --------------------------------------------------------------------------- + +void StorageManager::_batchedMigrate(PoolGroupIndex pgIdx, CacheLevel dstLevel, CacheLevel srcLevel, + std::vector> const& srcPages, bool updateSrc, MigrationRecorder const& migrationRecorder, + bool defrag) +{ + TLLM_CHECK_DEBUG(defrag || dstLevel != srcLevel); + SlotCount const numSlots = slotCountValueFromSize(srcPages.size()); + + auto& srcPoolGroup = poolGroup(srcLevel, pgIdx); + auto& dstPoolGroup = poolGroup(dstLevel, pgIdx); + + if (dstPoolGroup.numFreeSlots() < numSlots) + throw OutOfPagesError("Not enough free slots for migration"); + + auto dstSlots = dstPoolGroup.allocateMultiple(numSlots); + // A15: allocated slot count must match the request. + TLLM_CHECK_DEBUG_WITH_INFO(slotCountValueFromSize(dstSlots.size()) == numSlots, "dst_slots size mismatch"); + try + { + CacheTier dstTier = mLevels.at(dstLevel).cacheTier; + CacheTier srcTier = mLevels.at(srcLevel).cacheTier; + + PoolIndex numPools = mNumPools(pgIdx); + + // Build copy tasks per pool. + TypedVec> tasksPerPool(numPools); + for (std::size_t i = 0; i < srcPages.size(); ++i) + { + auto const& src = srcPages.at(i); + auto const& dst = dstSlots.at(i); + // Fix #8: assert non-defrag migrations only accept pages not scheduled for eviction. + TLLM_CHECK_DEBUG(defrag || !src->scheduledForEviction()); + for (PoolIndex poolIdx{0}; poolIdx < tasksPerPool.size(); ++poolIdx) + { + Address dstAddr = dstPoolGroup.slotAddress(dst.slotId()).at(poolIdx); + Address srcAddr = srcPoolGroup.slotAddress(src->slotId()).at(poolIdx); + tasksPerPool.at(poolIdx).push_back({dstAddr, srcAddr}); + } + } + + // Collect prior events (src + dst ready events) — mirrors Python's prior_events set. + std::vector priorEvents; + priorEvents.reserve(2 * srcPages.size()); + for (std::size_t i = 0; i < srcPages.size(); ++i) + { + priorEvents.push_back(&srcPages.at(i)->readyEvent); + priorEvents.push_back(&dstSlots.at(i).readyEvent); + } + + // Create a temporary CUDA stream that waits for all prior events before copying. + TemporaryCudaStream tempStream(priorEvents); + { + auto scope = tempStream.enter(); + CUstream stream = tempStream.get(); + auto slotSizes = slotSize(pgIdx); + for (PoolIndex poolIdx{0}; poolIdx < numPools; ++poolIdx) + { + batchedCopy(dstTier, srcTier, slotSizes.at(poolIdx), tasksPerPool.at(poolIdx), stream); + } + } // ~Scope records finish event + + CachedCudaEvent finishEvent = tempStream.takeFinishEvent(); + if (migrationRecorder && !defrag) + { + migrationRecorder(srcPages, dstSlots, srcLevel, dstLevel); + } + std::set> emittedCacheLevelUpdates; + bool const emitCacheLevelUpdates + = updateSrc && !defrag && srcLevel != dstLevel && static_cast(mEventSink); + for (std::size_t i = 0; i < srcPages.size(); ++i) + { + dstSlots.at(i).readyEvent = finishEvent; + // Fix #6: set src.ready_event unconditionally — compulsory for the next owner + // getting this slot from the pool. Mirrors Python: `src.ready_event = finish_event`. + srcPages.at(i)->readyEvent = finishEvent; + if (updateSrc) + { + bool wasScheduled = srcPages.at(i)->scheduledForEviction(); + if (wasScheduled) + excludeFromEviction(*srcPages.at(i)); + // Extract source slot from the page and release it back to the pool. + Slot srcSlot; + srcSlot.setSlotId(srcPages.at(i)->slotId()); // asserts valid + srcSlot.readyEvent = finishEvent; + srcPages.at(i)->resetSlot(); + srcPoolGroup.release(std::move(srcSlot)); + // Transfer dst slot ownership to the page. + srcPages.at(i)->setSlot(dstSlots.at(i)); + srcPages.at(i)->cacheLevel = dstLevel; + if (emitCacheLevelUpdates && srcPages.at(i)->isCommitted()) + { + auto const& page = static_cast(*srcPages.at(i)); + Block const* block = page.block; + std::string const blockKey = block + ? std::string(reinterpret_cast(block->key.data()), block->key.size()) + : std::string{}; + if (block && !block->isOrphan() + && emittedCacheLevelUpdates.insert({blockKey, page.lifeCycle.value()}).second) + { + mEventSink->addCacheLevelUpdated(block->key, srcLevel, dstLevel, page.lifeCycle); + } + } + if (wasScheduled) + scheduleForEviction(*srcPages.at(i)); + } + } + } + catch (...) + { + for (auto& s : dstSlots) + dstPoolGroup.release(std::move(s)); + throw; + } +} + +// --------------------------------------------------------------------------- +// batchedMigrateToGpu +// --------------------------------------------------------------------------- + +void StorageManager::batchedMigrateToGpu( + std::vector const& targets, KvCache& /*kvCache*/, MigrationRecorder const& migrationRecorder) +{ + // Group by (srcLevel, pgIdx). + std::map, std::vector>> groups; + for (auto const& t : targets) + { + if (t.page->cacheLevel == kGpuLevel) + continue; + PoolGroupIndex pg = mLifeCycleGrouping.at(t.lifeCycle); + groups[{t.page->cacheLevel, pg}].push_back(t.page); + } + for (auto& [key, pages] : groups) + _batchedMigrate(key.second, kGpuLevel, key.first, pages, /*updateSrc=*/true, migrationRecorder); +} + +void StorageManager::prefetch( + CacheLevel dstLevel, TypedVec>>> const& pages) +{ + TypedVec numSlotsToMigrate(numPoolGroups(), 0); + std::vector> scheduled; + + struct ReschedulePagesGuard + { + StorageManager& storageManager; + std::vector>& scheduled; + + ~ReschedulePagesGuard() + { + for (auto const& page : scheduled) + { + storageManager.scheduleForEviction(*page); + } + scheduled.clear(); + } + } reschedulePagesGuard{*this, scheduled}; + + for (PoolGroupIndex pgIndex{0}; pgIndex < pages.size(); ++pgIndex) + { + auto const& poolGroupPages = pages.at(pgIndex); + for (CacheLevel level{0}; level < poolGroupPages.size(); ++level) + { + auto const& levelPages = poolGroupPages.at(level); + TLLM_CHECK_DEBUG(level >= dstLevel || levelPages.empty()); + for (auto const& page : levelPages) + { + if (page->scheduledForEviction()) + { + excludeFromEviction(*page); + scheduled.push_back(page); + } + else if (isEvictable(*page, dstLevel)) + { + scheduled.push_back(page); + } + TLLM_CHECK_DEBUG(level >= dstLevel); + if (level == dstLevel) + { + continue; + } + numSlotsToMigrate.at(pgIndex) += 1; + } + } + } + + prepareFreeSlots(dstLevel, numSlotsToMigrate); + for (PoolGroupIndex pgIndex{0}; pgIndex < pages.size(); ++pgIndex) + { + auto const& poolGroupPages = pages.at(pgIndex); + for (CacheLevel lvl = dstLevel + 1; lvl < numCacheLevels(); ++lvl) + { + _batchedMigrate(pgIndex, dstLevel, lvl, poolGroupPages.at(lvl), /*updateSrc=*/true); + } + } +} + +// --------------------------------------------------------------------------- +// Query helpers +// --------------------------------------------------------------------------- + +LifeCycle const& StorageManager::getLifeCycle(LifeCycleId lc) const +{ + return mLifeCycles[lc]; +} + +PoolGroupIndex StorageManager::getPoolGroupIndex(LifeCycleId lc) const +{ + return mLifeCycleGrouping.at(lc); +} + +PoolIndex StorageManager::mNumPools(PoolGroupIndex pgIdx) const +{ + TLLM_CHECK_DEBUG(!mLevels.empty()); + return getUniformAttribute(mLevels, [pgIdx](auto const& lvl) { return lvl.storage->numPools(pgIdx); }); +} + +PoolIndex StorageManager::numPools(PoolGroupIndex pgIdx) const +{ + return mNumPools(pgIdx); +} + +TypedVec StorageManager::slotSize(PoolGroupIndex pgIdx) const +{ + return mSlotDescList.at(pgIdx).slotSizeList(); +} + +PoolGroupBase& StorageManager::poolGroup(CacheLevel lvl, PoolGroupIndex pgIdx) +{ + return mLevels.at(lvl).storage->poolGroup(pgIdx); +} + +MemAddress StorageManager::getMemPoolBaseAddress(LayerId layerId, DataRole role) const +{ + auto it = mBufferAttr.find(BufferId{layerId, role}); + if (it == mBufferAttr.end()) + throw std::out_of_range("Unknown BufferId"); + auto const& attr = it->second; + PoolGroupIndex pgIdx = mLifeCycleGrouping.at(attr.lifeCycleId); + return mLevels[kGpuLevel].storage->getBaseAddress(pgIdx, attr.poolIndex, SlotId{0}) + attr.offset; +} + +MemAddress StorageManager::getMemPoolBaseAddress(PoolGroupIndex pgIdx, PoolIndex poolIdx) const +{ + return mLevels[kGpuLevel].storage->getBaseAddress(pgIdx, poolIdx, SlotId{0}); +} + +LayerAttr const& StorageManager::getLayerAttr(LayerId layerId) const +{ + auto it = mLayerAttributes.find(layerId); + if (it == mLayerAttributes.end()) + throw std::out_of_range("Unknown LayerId for LayerAttr"); + return it->second; +} + +SlotCount StorageManager::numSlots(PoolGroupIndex pgIdx, CacheLevel level) const +{ + return mLevels.at(level).storage->numSlots(pgIdx); +} + +StorageStatistics StorageManager::getStatistics(CacheLevel level, PoolGroupIndex pgIdx) const +{ + auto const& lvl = mLevels.at(level); + SlotCount freeSlots = lvl.storage->numFreeSlots(pgIdx); + SlotCount totalSlots = lvl.storage->numSlots(pgIdx); + SlotCount evictable = lvl.controller.numEvictablePages(pgIdx); + auto sizes = lvl.storage->slotSize(pgIdx); + return StorageStatistics{sizes, totalSlots, freeSlots, evictable}; +} + +TypedVec StorageManager::getUtilization(CacheLevel level) const +{ + TypedVec result; + result.reserve(numPoolGroups()); + for (PoolGroupIndex pgIdx{0}; pgIdx < numPoolGroups(); ++pgIdx) + { + auto const s = getStatistics(level, pgIdx); + TLLM_CHECK_DEBUG(s.total > 0); + result.push_back(static_cast(s.unavailable()) / static_cast(s.total)); + } + return result; +} + +float StorageManager::getOverallUtilization(CacheLevel level) const +{ + float num = 0.f, den = 0.f; + for (PoolGroupIndex pgIdx{0}; pgIdx < numPoolGroups(); ++pgIdx) + { + auto s = getStatistics(level, pgIdx); + float sz = 0.f; + for (auto v : s.slotSizes) + sz += static_cast(v); + num += sz * static_cast(s.unavailable()); + den += sz * static_cast(s.total); + } + TLLM_CHECK_DEBUG(den > 0.f); + return num / den; +} + +// --------------------------------------------------------------------------- +// expandPoolGroup +// --------------------------------------------------------------------------- + +void StorageManager::expandPoolGroup(CacheLevel level, PoolGroupIndex pgIdx, SlotCount newNumSlots) +{ + auto& pg = poolGroup(level, pgIdx); + TLLM_CHECK_DEBUG(newNumSlots > pg.numSlots()); + pg.resizePools(newNumSlots); + pg.slotAllocator().expand(newNumSlots); +} + +// --------------------------------------------------------------------------- +// shrinkPoolGroup — mirrors Python _storage_manager.py::shrink_pool_group +// --------------------------------------------------------------------------- + +void StorageManager::shrinkPoolGroup( + CacheLevel level, PoolGroupIndex pgIdx, SlotCount newNumSlots, std::vector> const& persistentPages) +{ + auto& pg = poolGroup(level, pgIdx); + auto& allocator = pg.slotAllocator(); + auto& ctrl = mLevels.at(level).controller; + TLLM_CHECK_DEBUG(newNumSlots < pg.numSlots()); + + // A16: persistent_pages preconditions. + TLLM_CHECK_DEBUG_WITH_INFO( + persistentPages.size() <= slotCountToSizeT(newNumSlots), "Not enough slots to hold all persistent pages"); + TLLM_CHECK_DEBUG_WITH_INFO(std::all_of(persistentPages.begin(), persistentPages.end(), + [this, level, pgIdx](auto const& p) + { return p->cacheLevel == level && mLifeCycleGrouping.at(p->lifeCycle) == pgIdx; }), + "Persistent page cache level or pool group mismatch"); + + // Fast path: when no slot id has ever been issued in the to-be-removed + // range [newNumSlots, capacity), there is nothing to migrate. + // numActiveSlots() is a monotone high-water mark of issued ids. + if (allocator.numActiveSlots() <= newNumSlots) + { + allocator.prepareForShrink(newNumSlots); + allocator.finishShrink(); + pg.resizePools(newNumSlots); + return; + } + + // Find overflow pages: scheduled pages with slot_id >= newNumSlots. + auto gen = ctrl.pageGenerator(pgIdx); + std::deque>> overflowSlots; + { + SlotCount idx = 0; + while (auto const* page = gen()) + { + if ((*page)->slotId() >= newNumSlots) + overflowSlots.emplace_back(idx, *page); + ++idx; + } + } + + // Persistent pages in overflow range. + std::vector> overflowPersistent; + for (auto const& p : persistentPages) + { + if (p->slotId() >= newNumSlots) + overflowPersistent.push_back(p); + } + SlotCount numOverflowPersistent = slotCountValueFromSize(overflowPersistent.size()); + + // A2: RUNTIME check — persistent overflow pages must fit in the new capacity. + if (numOverflowPersistent > newNumSlots) + { + throw OutOfPagesError("Not enough slots to hold all persistent pages"); + } + + // Mark the allocator for shrink. + allocator.prepareForShrink(newNumSlots); + + // Calculate minimum number of lowest-priority pages to evict. + // Need numEvictedOverflowSlots because evicted overflow pages won't become free, + // because only free non-overflow slots can be used for defragmentation. + SlotCount minNumEvicted = 0; + SlotCount numEvictedOverflowSlots = 0; + while (!overflowSlots.empty() + && slotCountValueFromSize(overflowSlots.size()) + numOverflowPersistent + > std::min(newNumSlots, overflowSlots.front().first + allocator.numFreeSlots() - numEvictedOverflowSlots)) + { + minNumEvicted = overflowSlots.front().first + 1; + overflowSlots.pop_front(); + ++numEvictedOverflowSlots; + } + + // Force-evict the required pages. + TypedVec evictReqs(numPoolGroups(), 0); + evictReqs[pgIdx] = minNumEvicted; + forceEvict(level, evictReqs); + + // Remaining overflow pages to defragment. + std::vector> overflowPages; + overflowPages.reserve(overflowSlots.size() + overflowPersistent.size()); + for (auto& [idx, p] : overflowSlots) + overflowPages.push_back(p); + for (auto& p : overflowPersistent) + overflowPages.push_back(p); + + // Ensure free slots for the overflow pages. + TypedVec reqs(numPoolGroups(), 0); + reqs[pgIdx] = slotCountValueFromSize(overflowPages.size()); + prepareFreeSlots(level, reqs); + + // A17: all overflow pages must be at the expected cache level. + TLLM_CHECK_DEBUG_WITH_INFO(std::all_of(overflowPages.begin(), overflowPages.end(), + [level](auto const& p) { return p->cacheLevel == level; }), + "Overflow page cache level mismatch"); + + // Defragment: migrate overflow pages to free slots within the same level. + _batchedMigrate(pgIdx, level, level, overflowPages, /*updateSrc=*/true, MigrationRecorder{}, /*defrag=*/true); + + // A18: post-defrag overflow assertion — overflow slot count matches expectations. + TLLM_CHECK_DEBUG_WITH_INFO(allocator.numOverflowSlots() == allocator.numActiveSlots() - allocator.targetCapacity(), + "Post-defrag overflow slot count mismatch"); + + // Finalize shrink and resize pools. + allocator.finishShrink(); + pg.resizePools(newNumSlots); +} + +// --------------------------------------------------------------------------- +// adjustCacheLevel — mirrors Python _storage_manager.py::adjust_cache_level +// --------------------------------------------------------------------------- + +void StorageManager::adjustCacheLevel(CacheLevel level, std::optional newQuota, + TypedVec const& ratioList, + TypedVec>> const* persistentPages) +{ + auto& lvlStorage = *mLevels.at(level).storage; + auto oldNumSlots = lvlStorage.slotCountList(); + size_t quota = newQuota.has_value() + ? roundUp(newQuota.value(), static_cast(lvlStorage.poolSizeGranularity())) + : lvlStorage.totalQuota(); + size_t minQuota = minQuotaForLevel(lvlStorage.slotSizeLists(), lvlStorage.poolSizeGranularity()); + if (quota < minQuota) + { + throw std::invalid_argument("Quota " + std::to_string(quota) + + " is insufficient for min_slots constraints (requires at least " + std::to_string(minQuota) + ")"); + } + auto newNumSlots = lvlStorage.computeSlotCountList(ratioList, mMinSlots, quota); + + if (!isLastLevel(level)) + TLLM_CHECK_DEBUG(persistentPages == nullptr); + + // Shrink first. + for (PoolGroupIndex pgIdx{0}; pgIdx < newNumSlots.size(); ++pgIdx) + { + if (newNumSlots[pgIdx] >= oldNumSlots[pgIdx]) + continue; + std::vector> pages; + if (persistentPages) + pages = (*persistentPages)[pgIdx]; + shrinkPoolGroup(level, pgIdx, newNumSlots[pgIdx], pages); + } + // Then expand. + for (PoolGroupIndex pgIdx{0}; pgIdx < newNumSlots.size(); ++pgIdx) + { + if (newNumSlots[pgIdx] <= oldNumSlots[pgIdx]) + continue; + expandPoolGroup(level, pgIdx, newNumSlots[pgIdx]); + } + lvlStorage.postResize(); +} + +TypedVec StorageManager::getRatioList(CacheLevel level) const +{ + return mLevels.at(level).storage->ratioList(); +} + +TypedVec StorageManager::ratioFromLength( + int tokensPerBlock, int historyLength, int capacity) const +{ + if (capacity < historyLength) + { + TLLM_LOG_WARNING("Bad sampling for capacity and history_length"); + capacity = historyLength; + } + int numBlocks = divUp(capacity, tokensPerBlock); + TypedVec numBytes(numPoolGroups(), 0); + auto ssmLcId = mLifeCycles.ssmLifeCycleId(); + auto const& lifecycles = mLifeCycles.getAll(); + for (LifeCycleId lcId{0}; lcId < lifecycles.size(); ++lcId) + { + PoolGroupIndex pgIdx = mLifeCycleGrouping[lcId]; + auto ss = slotSize(pgIdx); + size_t slotSizeSum = 0; + for (auto s : ss) + slotSizeSum += s; + int numRequiredBlocks; + if (ssmLcId.has_value() && lcId == *ssmLcId) + { + numRequiredBlocks = 1; + } + else + { + auto stale = getStaleRange(lifecycles[lcId], historyLength, tokensPerBlock); + numRequiredBlocks = std::max(numBlocks - stale.length(), 1); + } + numBytes[pgIdx] += static_cast(numRequiredBlocks) * slotSizeSum; + } + return normalizeToRatio(numBytes); +} + +// --------------------------------------------------------------------------- +// ratioFromBatch +// --------------------------------------------------------------------------- + +TypedVec StorageManager::ratioFromBatch(BatchDesc const& batch, int tokensPerBlock, + std::optional const& swaScratchReuse, size_t granularity) const +{ + auto numSlots = computeSlotsForBatch(batch, tokensPerBlock, swaScratchReuse); + auto numBytes = slotsToBytes(numSlots, granularity); + return normalizeToRatio(numBytes); +} + +// --------------------------------------------------------------------------- +// computeMinSlotsFromConstraints +// --------------------------------------------------------------------------- + +TypedVec StorageManager::computeMinSlotsFromConstraints( + std::vector const& constraints, int tokensPerBlock, + std::optional const& swaScratchReuse, float maxUtilForResume) const +{ + TLLM_CHECK_DEBUG(maxUtilForResume > 0.0f && maxUtilForResume <= 1.0f); + // All returned elements are positive. Constraint-derived floors include headroom + // for the utilization gate checked by KvCache::resume. + TypedVec maxSlots(numPoolGroups(), 0); + + auto swaFloorBlocks = [tokensPerBlock](AttnLifeCycle const& lc) -> int + { + int window = *lc.windowSize; + // Handle oscillation of slot count required by SWA while the window slides. + return lc.numSinkBlocks + (window + tokensPerBlock - 2) / tokensPerBlock + 1; + }; + + // Full-attention lifecycles share the largest SWA floor: all attention + // lifecycles see the same seq_len, so this is a valid lower bound. + int floorNumBlocks = 1; + for (auto const& [lcId, attn] : mLifeCycles.attentionLifeCycles()) + { + if (attn->windowSize.has_value()) + floorNumBlocks = std::max(floorNumBlocks, swaFloorBlocks(*attn)); + } + for (auto const& [lcIdx, lc] : mLifeCycles) + { + PoolGroupIndex pgIdx = getPoolGroupIndex(lcIdx); + auto const* attn = std::get_if(&lc); + if (attn == nullptr) + { + // SSM / non-attention: 1 slot floor per life cycle. + maxSlots[pgIdx] += 1; + } + else if (attn->windowSize.has_value()) + { + maxSlots[pgIdx] += swaFloorBlocks(*attn); + } + else + { + maxSlots[pgIdx] += floorNumBlocks; + } + } + for (auto const& batch : constraints) + { + auto slots = computeSlotsForBatch(batch, tokensPerBlock, swaScratchReuse); + for (PoolGroupIndex pgIdx{0}; pgIdx < slots.size(); ++pgIdx) + { + auto const scaledSlots = static_cast( + std::ceil(static_cast(slots[pgIdx]) / static_cast(maxUtilForResume))); + maxSlots[pgIdx] = std::max(maxSlots[pgIdx], scaledSlots); + } + } + return maxSlots; +} + +// --------------------------------------------------------------------------- +// computeSlotsForBatch +// --------------------------------------------------------------------------- + +TypedVec StorageManager::computeSlotsForBatch( + BatchDesc const& batch, int tokensPerBlock, std::optional const& swaScratchReuse) const +{ + TypedVec numSlots(numPoolGroups(), 0); + auto ssmLcId = mLifeCycles.ssmLifeCycleId(); + int sysBlocks = batch.systemPromptLength / tokensPerBlock; + + for (auto const& [lcIdx, lc] : mLifeCycles) + { + PoolGroupIndex pgIdx = mLifeCycleGrouping[lcIdx]; + if (ssmLcId.has_value() && lcIdx == *ssmLcId) + { + // SSM: always 1 dedicated block per request, never shared. + numSlots[pgIdx] += slotCountValueFromSize(batch.kvCaches.size()); + continue; + } + // Shared sys blocks (counted once): union of non-stale sys blocks across all requests. + HalfOpenRange sysRange{0, sysBlocks}; + HalfOpenRange staleIntersection = sysRange; + for (auto const& kv : batch.kvCaches) + { + auto stale = getStaleRange(lc, kv.historyLength, tokensPerBlock); + staleIntersection = intersect(staleIntersection, stale); + } + numSlots[pgIdx] += sysBlocks - staleIntersection.length(); + + // Per-request unique blocks (excluding shared sys blocks already counted above). + for (auto const& kv : batch.kvCaches) + { + int totalBlocks = divUp(kv.capacity, tokensPerBlock); + auto stale = getStaleRange(lc, kv.historyLength, tokensPerBlock); + int nonStale = totalBlocks - stale.length(); + int nonStaleSys = sysBlocks - intersect(stale, sysRange).length(); + int uniqueNonStale = std::max(0, nonStale - nonStaleSys); + if (swaScratchReuse.has_value()) + { + auto scratch = computeScratchRange( + lc, kv.historyLength, kv.capacity, tokensPerBlock, swaScratchReuse->maxRewindLen); + int numScratch = scratch.length(); + // Scratch blocks share coalesced slots: actual slots = ceil(numScratch * fracMax). + numSlots[pgIdx] += (uniqueNonStale - numScratch) + mSlotUtilFracMax[lcIdx].ceilMul(numScratch); + } + else + { + numSlots[pgIdx] += uniqueNonStale; + } + } + } + return numSlots; +} + +// --------------------------------------------------------------------------- +// slotsToBytes +// --------------------------------------------------------------------------- + +TypedVec StorageManager::slotsToBytes( + TypedVec const& numSlots, size_t granularity) const +{ + TypedVec numBytes(numPoolGroups(), 0); + for (PoolGroupIndex pgIdx{0}; pgIdx < numSlots.size(); ++pgIdx) + { + for (auto poolSize : slotSize(pgIdx)) + { + numBytes[pgIdx] += roundUp(slotCountToSizeT(numSlots[pgIdx]) * poolSize, granularity); + } + } + return numBytes; +} + +// --------------------------------------------------------------------------- +// computeSlotCountForLevel +// --------------------------------------------------------------------------- + +TypedVec StorageManager::computeSlotCountForLevel(CacheTierConfig const& tierConfig, + TypedVec> const& slotSizeLists, + TypedVec const& ratio) const +{ + CacheTier tier = cacheTierOf(tierConfig); + size_t quota = cacheTierQuota(tierConfig); + size_t granularity = CacheLevelManager::cacheTierGranularity(tier, quota); + quota = std::max(minQuotaForLevel(slotSizeLists, granularity), roundUp(quota, granularity)); + return CacheLevelStorage::ratioToSlotCountList(quota, slotSizeLists, ratio, granularity, mMinSlots); +} + +// --------------------------------------------------------------------------- +// minQuotaForLevel +// --------------------------------------------------------------------------- + +size_t StorageManager::minQuotaForLevel( + TypedVec> const& slotSizeLists, size_t granularity) const +{ + size_t total = 0; + for (PoolGroupIndex pgIdx{0}; pgIdx < slotSizeLists.size(); ++pgIdx) + { + for (auto slotSize : slotSizeLists[pgIdx]) + { + total += roundUp(slotCountToSizeT(mMinSlots[pgIdx]) * slotSize, granularity); + } + } + return total; +} + +// --------------------------------------------------------------------------- +// constrainRatio +// --------------------------------------------------------------------------- + +TypedVec StorageManager::constrainRatio(TypedVec const& ratio) const +{ + auto& gpuStorage = *mLevels[kGpuLevel].storage; + size_t granularity = gpuStorage.poolSizeGranularity(); + auto slotCountList = gpuStorage.computeSlotCountList(ratio, mMinSlots); + auto numBytes = slotsToBytes(slotCountList, granularity); + return normalizeToRatio(numBytes); +} + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h new file mode 100644 index 000000000000..27979516059e --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h @@ -0,0 +1,331 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "kv_cache_manager_v2/common.h" +#include "kv_cache_manager_v2/config.h" +#include "kv_cache_manager_v2/eventSink.h" +#include "kv_cache_manager_v2/evictionController.h" +#include "kv_cache_manager_v2/lifeCycleRegistry.h" +#include "kv_cache_manager_v2/storage/config.h" +#include "kv_cache_manager_v2/storage/core.h" +#include "tensorrt_llm/common/assert.h" + +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// Forward declarations. +class Page; +class KvCache; +struct BatchedLockTarget; + +using MigrationRecorder + = std::function> const&, std::vector const&, CacheLevel, CacheLevel)>; +using DropRecorder = std::function> const&, CacheLevel)>; + +// --------------------------------------------------------------------------- +// StorageStatistics — per-pool-group slot counts. +// --------------------------------------------------------------------------- +struct StorageStatistics +{ + TypedVec slotSizes; + SlotCount total; // total slots + SlotCount free; // free (unallocated) slots + SlotCount evictable; // scheduled for eviction + + SlotCount available() const noexcept + { + return free + evictable; + } + + SlotCount unavailable() const noexcept + { + return total - available(); + } +}; + +// --------------------------------------------------------------------------- +// CacheLevelManager — one storage level (GPU/Host/Disk) + its eviction controller. +// --------------------------------------------------------------------------- +class CacheLevelManager +{ +public: + CacheLevelManager(TypedVec const& lifeCycleGrouping, CacheLevel cacheLevel, + CacheTierConfig const& tierConfig, StorageConfig const& storageConfig, + TypedVec const& slotCountList); + + // Compute pool size granularity for a given cache tier and quota. + static size_t cacheTierGranularity(CacheTier tier, size_t quota); + + CacheLevel cacheLevel; + CacheTier cacheTier; + std::unique_ptr storage; + PerLevelEvictionController controller; + + PoolGroupIndex numPoolGroups() const noexcept + { + TLLM_CHECK_DEBUG_WITH_INFO( + storage->numPoolGroups() == controller.numPoolGroups(), "Storage and controller disagree on numPoolGroups"); + return controller.numPoolGroups(); + } +}; + +// --------------------------------------------------------------------------- +// StorageManager — manages all cache levels and the eviction pipeline. +// Mirrors Python's StorageManager. +// --------------------------------------------------------------------------- +class StorageManager : public std::enable_shared_from_this +{ +public: + StorageManager(LifeCycleRegistry const& lifeCycles, StorageConfig const& config, int tokensPerBlock, + std::optional swaScratchReuse = std::nullopt, + std::optional const& typicalBatch = std::nullopt, std::vector const& constraints = {}, + std::optional> const& initialPoolRatio = std::nullopt, + std::shared_ptr eventSink = nullptr, float maxUtilForResume = 1.0f); + ~StorageManager(); + + StorageManager(StorageManager const&) = delete; + StorageManager& operator=(StorageManager const&) = delete; + + void destroy(); + + // ---- Allocation ------------------------------------------------------- + + // Allocate slots for all life cycles at the given cache level. + // numSlotsPerLc[lcId] = how many slots to allocate for that life cycle. + // Returns a vector indexed by lcId. + TypedVec> newSlots(CacheLevel level, + TypedVec const& numSlotsPerLc, MigrationRecorder const& migrationRecorder = {}, + DropRecorder const& dropRecorder = {}); + + TypedVec> newGpuSlots(TypedVec const& numSlotsPerLc, + MigrationRecorder const& migrationRecorder = {}, DropRecorder const& dropRecorder = {}); + + // Allocate slots for a single pool group at the given cache level. + // Returns numSlots Slot objects. Throws OutOfPagesError if allocation fails. + std::vector newSlotsForPoolGroup(CacheLevel level, PoolGroupIndex pgIdx, SlotCount numSlots, + MigrationRecorder const& migrationRecorder = {}, DropRecorder const& dropRecorder = {}); + + // Release a slot back to its pool. + void releaseSlot(LifeCycleId lc, CacheLevel level, Slot slot); + + // ---- Eviction ---------------------------------------------------------- + + // Schedule a page for eviction (if evictable). + void scheduleForEviction(Page& page); + + // Remove a page from the eviction queue. + void excludeFromEviction(Page& page); + + // Check if a page is evictable (optionally at a target level). + bool isEvictable(Page const& page, std::optional level = std::nullopt) const noexcept; + + // Ensure numFreeSlots[pgIdx] free GPU slots exist (evicting pages as needed). + void prepareFreeSlots(CacheLevel level, TypedVec const& requirements, + MigrationRecorder const& migrationRecorder = {}, DropRecorder const& dropRecorder = {}); + + // Force-evict pages from a level to free space. + void forceEvict(CacheLevel level, TypedVec const& minNumPages, + DropRecorder const& dropRecorder = {}); + + // Dynamic cache level resizing. + void adjustCacheLevel(CacheLevel level, std::optional newQuota, + TypedVec const& ratioList, + TypedVec>> const* persistentPages); + void shrinkPoolGroup(CacheLevel level, PoolGroupIndex pgIdx, SlotCount newNumSlots, + std::vector> const& persistentPages); + void expandPoolGroup(CacheLevel level, PoolGroupIndex pgIdx, SlotCount newNumSlots); + + // ---- Migration --------------------------------------------------------- + + // Migrate a batch of pages to GPU (used by batchedLockToGpu). + void batchedMigrateToGpu( + std::vector const& targets, KvCache& kvCache, MigrationRecorder const& migrationRecorder); + + // Best-effort migration of grouped pages to a destination cache level. + void prefetch( + CacheLevel dstLevel, TypedVec>>> const& pages); + + // ---- Query helpers ----------------------------------------------------- + + LifeCycleRegistry const& lifeCycles() const noexcept + { + return mLifeCycles; + } + + LifeCycle const& getLifeCycle(LifeCycleId lc) const; + + LifeCycleId numLifeCycles() const noexcept + { + return mLifeCycleGrouping.size(); + } + + PoolGroupIndex numPoolGroups() const noexcept + { + return mSlotDescList.size(); + } + + TypedVec const& slotDescList() const noexcept + { + return mSlotDescList; + } + + CacheLevel numCacheLevels() const noexcept + { + return mLevels.size(); + } + + bool isLastLevel(CacheLevel lvl) const noexcept + { + return lvl == numCacheLevels() - 1; + } + + PoolGroupIndex getPoolGroupIndex(LifeCycleId lc) const; + PoolIndex numPools(PoolGroupIndex pgIdx) const; + + // Return the byte size of each pool in a pool group. + TypedVec slotSize(PoolGroupIndex pgIdx) const; + + // Current ratio list for a cache level (proportional to byte usage per pool group). + TypedVec getRatioList(CacheLevel level) const; + + // Compute init ratio from an assumed average history length and capacity. + TypedVec ratioFromLength(int tokensPerBlock, int historyLength, int capacity) const; + + // Compute ratio from a BatchDesc. + TypedVec ratioFromBatch(BatchDesc const& batch, int tokensPerBlock, + std::optional const& swaScratchReuse, size_t granularity) const; + + // Apply stored min_slots constraint to a ratio list for GPU level. + TypedVec constrainRatio(TypedVec const& ratio) const; + + // Byte address of a slot's buffer in GPU memory (per-layer, with offset). + MemAddress getMemPoolBaseAddress(LayerId layerId, DataRole role) const; + + // Pool group base address without per-layer offset. + MemAddress getMemPoolBaseAddress(PoolGroupIndex pgIdx, PoolIndex poolIdx) const; + + // Per-layer storage attributes. + LayerAttr const& getLayerAttr(LayerId layerId) const; + + // Address of a slot's buffer in a specific pool at a cache level. + Address slotAddress(CacheLevel level, PoolGroupIndex pgIdx, SlotId slotId, PoolIndex poolIdx) const; + + // Cache tier for a given level. + CacheTier cacheTier(CacheLevel level) const; + + // NOTE: Python's get_statistics(level) returns a list over all pool groups. + // C++ takes a single pgIdx for flexibility; the nanobind wrapper loops over + // all pool groups to match Python's signature. + StorageStatistics getStatistics(CacheLevel level = kGpuLevel, PoolGroupIndex pgIdx = PoolGroupIndex{0}) const; + TypedVec getUtilization(CacheLevel level = kGpuLevel) const; + float getOverallUtilization(CacheLevel level = kGpuLevel) const; + + // Pool-group slot count (number of pages). + SlotCount numSlots(PoolGroupIndex pgIdx, CacheLevel level = kGpuLevel) const; + + // Layer-to-lifecycle mapping (for KvCacheManager queries). + std::unordered_map const& layerToLifeCycleIds() const noexcept + { + return mLayerToLifeCycleIds; + } + + // Per-lifecycle max slot utilization fraction (for scratch slot computation). + Rational const& slotUtilFracMax(LifeCycleId lcId) const + { + return mSlotUtilFracMax.at(lcId); + } + + friend class KvCacheManager; + // White-box test introspection reaches computeSlotsForBatch() via friendship + // (mirrors Python's _compute_slots_for_batch). + friend class KvCacheIntrospection; + +private: + // Minimum per-pool-group slot counts to support a BatchDesc. + TypedVec computeSlotsForBatch( + BatchDesc const& batch, int tokensPerBlock, std::optional const& swaScratchReuse) const; + + // Constraint-based partitioning helpers. + TypedVec computeMinSlotsFromConstraints(std::vector const& constraints, + int tokensPerBlock, std::optional const& swaScratchReuse, + float maxUtilForResume = 1.0f) const; + TypedVec slotsToBytes( + TypedVec const& numSlots, size_t granularity) const; + TypedVec computeSlotCountForLevel(CacheTierConfig const& tierConfig, + TypedVec> const& slotSizeLists, + TypedVec const& ratio) const; + size_t minQuotaForLevel( + TypedVec> const& slotSizeLists, size_t granularity) const; + + PoolIndex mNumPools(PoolGroupIndex pgIdx) const; + + // Internal helpers. + void _prepareFreeSlots(TypedVec>& goals, CacheLevel lvlId, + TypedVec>>& fallenPages, + MigrationRecorder const& migrationRecorder = {}, DropRecorder const& dropRecorder = {}); + + void _batchedMigrate(PoolGroupIndex pgIdx, CacheLevel dstLevel, CacheLevel srcLevel, + std::vector> const& srcPages, bool updateSrc, MigrationRecorder const& migrationRecorder = {}, + bool defrag = false); + + PoolGroupBase& poolGroup(CacheLevel lvl, PoolGroupIndex pgIdx); + + LifeCycleRegistry const& mLifeCycles; + std::shared_ptr mEventSink; + TypedVec mLifeCycleGrouping; // lcId → pgIdx + std::unordered_map mLayerToLifeCycleIds; + StorageConfig mStorageConfig; + + // slot-to-page-index scale factors: [lcId][poolIdx] + TypedVec> mSlotToPageIndices; + + // Per-layer storage attributes for scratch slot management. + std::map mLayerAttributes; + + // Max slot utilization fraction per lifecycle (across all layers in that lifecycle). + TypedVec mSlotUtilFracMax; + + // Whether SWA scratch reuse is enabled. + std::optional mSwaScratchReuse; + + // Get buffer attributes for a (LayerId, DataRole) pair. Throws std::out_of_range if not found. + // Mirrors Python's get_buffer_attr(). + BufferAttr const& getBufferAttr(LayerId layerId, DataRole role) const + { + auto it = mBufferAttr.find(BufferId{layerId, role}); + if (it == mBufferAttr.end()) + throw std::out_of_range("Unknown buffer id"); + return it->second; + } + + // Buffer attributes keyed by BufferId. + std::map mBufferAttr; + + TypedVec mSlotDescList; + TypedVec mMinSlots; + TypedVec mLevels; +}; + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/cudaEvent.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/cudaEvent.cpp new file mode 100644 index 000000000000..865374f765ea --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/cudaEvent.cpp @@ -0,0 +1,187 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "kv_cache_manager_v2/utils/cudaEvent.h" +#include "kv_cache_manager_v2/exceptions.h" + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// --------------------------------------------------------------------------- +// CudaEventPool / CudaStreamPool singleton implementations +// --------------------------------------------------------------------------- + +CudaEventPool::CudaEventPool() + : SimplePool( + []() -> CUevent + { + CUevent ev; + cuCheck(cuEventCreate(&ev, CU_EVENT_DISABLE_TIMING)); + return ev; + }, + [](CUevent ev) { cuEventDestroy(ev); }, + /*initSize=*/1024) +{ +} + +CudaEventPool& CudaEventPool::instance() +{ + static CudaEventPool pool; + return pool; +} + +CudaStreamPool::CudaStreamPool() + : SimplePool( + []() -> CUstream + { + CUstream s; + cuCheck(cuStreamCreate(&s, CU_STREAM_NON_BLOCKING)); + return s; + }, + [](CUstream s) { cuStreamDestroy(s); }, + /*initSize=*/128) +{ +} + +CudaStreamPool& CudaStreamPool::instance() +{ + static CudaStreamPool pool; + return pool; +} + +// --------------------------------------------------------------------------- +// CachedCudaEvent implementation +// --------------------------------------------------------------------------- + +CachedCudaEvent CachedCudaEvent::makeNull() noexcept +{ + return CachedCudaEvent{}; +} + +CachedCudaEvent::CachedCudaEvent(CudaStream stream) + : mEvent(std::make_shared(CudaEventPool::instance().get())) +{ + cuCheck(cuEventRecord(mEvent->get(), reinterpret_cast(stream))); +} + +bool CachedCudaEvent::queryComplete() +{ + if (isClosed()) + { + return true; + } + CUresult result = cuEventQuery(mEvent->get()); + if (result == CUDA_SUCCESS) + { + close(); + return true; + } + if (result == CUDA_ERROR_NOT_READY) + { + return false; + } + throw CuError(result); +} + +void CachedCudaEvent::synchronize() +{ + if (isClosed()) + { + return; + } + cuCheck(cuEventSynchronize(mEvent->get())); + close(); +} + +void CachedCudaEvent::waitInStream(CudaStream stream) const +{ + if (isClosed()) + { + return; + } + cuCheck(cuStreamWaitEvent(reinterpret_cast(stream), mEvent->get(), 0)); +} + +void CachedCudaEvent::close() +{ + if (mEvent) + { + mEvent->reset(); + } +} + +// --------------------------------------------------------------------------- +// CachedCudaStream implementation +// --------------------------------------------------------------------------- + +CachedCudaStream::CachedCudaStream() + : mPoolItem(CudaStreamPool::instance().get()) +{ +} + +CachedCudaEvent CachedCudaStream::recordEvent() +{ + return CachedCudaEvent{reinterpret_cast(handle())}; +} + +void CachedCudaStream::synchronize() +{ + cuCheck(cuStreamSynchronize(handle())); +} + +// --------------------------------------------------------------------------- +// TemporaryCudaStream implementation +// --------------------------------------------------------------------------- + +TemporaryCudaStream::TemporaryCudaStream(std::vector const& priorEvents) + : mStream() +{ + CudaStream cs = reinterpret_cast(mStream.handle()); + streamWaitEvents(cs, priorEvents); +} + +// --------------------------------------------------------------------------- +// mergeEvents — merge multiple CUDA events into one. +// Mirrors Python's merge_events() in _utils.py. +// --------------------------------------------------------------------------- + +CachedCudaEvent mergeEvents(std::vector& events) +{ + // Filter out closed events (optimization: skip cuStreamWaitEvent calls). + std::vector live; + for (auto& ev : events) + { + if (!ev.isClosed()) + live.push_back(&ev); + } + if (live.empty()) + return CachedCudaEvent::makeNull(); + if (live.size() == 1) + return std::move(*live[0]); + // Multiple live events: merge via TemporaryCudaStream. + std::vector priors; + priors.reserve(live.size()); + for (auto* ev : live) + priors.push_back(ev); + TemporaryCudaStream tempStream(priors); + { + auto scope = tempStream.enter(); + } + return tempStream.takeFinishEvent(); +} + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/cudaEvent.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/cudaEvent.h new file mode 100644 index 000000000000..5621ca4b04b5 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/cudaEvent.h @@ -0,0 +1,443 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "kv_cache_manager_v2/common.h" +#include "kv_cache_manager_v2/exceptions.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// --------------------------------------------------------------------------- +// FuncGuard — generic RAII scope guard that calls a void() callable on destruction. +// Movable (moved-from instance is disarmed). Not copyable. +// --------------------------------------------------------------------------- +template +class FuncGuard +{ +public: + explicit FuncGuard(F&& func) + : mFunc(std::forward(func)) + , mActive(true) + { + } + + ~FuncGuard() + { + if (mActive) + { + mFunc(); + } + } + + FuncGuard(FuncGuard&& other) noexcept + : mFunc(std::move(other.mFunc)) + , mActive(other.mActive) + { + other.mActive = false; + } + + FuncGuard(FuncGuard const&) = delete; + FuncGuard& operator=(FuncGuard const&) = delete; + FuncGuard& operator=(FuncGuard&&) = delete; + +private: + F mFunc; + bool mActive; +}; + +// --------------------------------------------------------------------------- +// SimplePool — generic resource pool for opaque handle types. +// Mirrors _utils.py::SimplePool. +// +// T is the pointed-to type (e.g. CUevent_st, CUstream_st). +// CreateFn returns T*, DestroyFn takes T*. +// get() returns a PoolItem (unique_ptr) with zero extra allocation — +// the unique_ptr directly wraps the handle pointer. +// +// Derived (CRTP, default void): +// - void: instance pool — Deleter stores a SimplePool* pointer (8 bytes). +// - non-void: singleton pool — Deleter is stateless (0 bytes), calls +// Derived::instance() to find the pool. PoolItem is pointer-sized. +// --------------------------------------------------------------------------- + +// Forward declare so Deleters can reference it. +template +class SimplePool; + +// Deleter for instance pools (Derived == void): stores a pool pointer. +template +struct InstancePoolDeleter +{ + SimplePool* pool = nullptr; + + void operator()(T* ptr) const noexcept; +}; + +// Deleter for singleton pools (Derived != void): stateless, zero-size. +template +struct SingletonPoolDeleter +{ + void operator()(T* ptr) const noexcept; +}; + +template +class SimplePool +{ +public: + using CreateFn = std::function; + using DestroyFn = std::function; + + using Deleter + = std::conditional_t, InstancePoolDeleter, SingletonPoolDeleter>; + using PoolItem = std::unique_ptr; + + SimplePool(CreateFn createFn, DestroyFn destroyFn, int initSize = 0, std::optional maxSize = std::nullopt) + : mCreateFn(std::move(createFn)) + , mDestroyFn(std::move(destroyFn)) + , mMaxSize(maxSize) + , mOutstandingCount(0) + { + for (int i = 0; i < initSize; ++i) + { + mItems.push_back(mCreateFn()); + } + } + + ~SimplePool() + { + clear(); + } + + SimplePool(SimplePool const&) = delete; + SimplePool& operator=(SimplePool const&) = delete; + + // Get a resource wrapped in a PoolItem that auto-returns to pool on destruction. + [[nodiscard]] PoolItem get() + { + // Increment only after the item is successfully obtained, so a throwing + // mCreateFn() leaves mOutstandingCount unchanged (no leak in stats). + T* item = mItems.empty() ? mCreateFn() : popFront(); + ++mOutstandingCount; + if constexpr (std::is_void_v) + { + return PoolItem(item, Deleter{this}); + } + else + { + return PoolItem(item, Deleter{}); + } + } + + void clear() + { + while (!mItems.empty()) + { + mDestroyFn(popFront()); + } + } + + [[nodiscard]] int outstandingCount() const noexcept + { + return mOutstandingCount; + } + + [[nodiscard]] int cachedCount() const noexcept + { + return static_cast(mItems.size()); + } + +private: + friend struct InstancePoolDeleter; + friend struct SingletonPoolDeleter; + + T* popFront() + { + T* item = mItems.front(); + mItems.pop_front(); + return item; + } + + void put(T* item) + { + --mOutstandingCount; + if (mMaxSize.has_value() && static_cast(mItems.size()) >= *mMaxSize) + { + mDestroyFn(item); + } + else + { + mItems.push_back(item); + } + } + + CreateFn mCreateFn; + DestroyFn mDestroyFn; + std::optional mMaxSize; + std::deque mItems; + int mOutstandingCount; +}; + +// Deleter implementations (after SimplePool is fully defined). +template +void InstancePoolDeleter::operator()(T* ptr) const noexcept +{ + if (pool) + { + pool->put(ptr); + } +} + +template +void SingletonPoolDeleter::operator()(T* ptr) const noexcept +{ + Derived::instance().put(ptr); +} + +// --------------------------------------------------------------------------- +// CudaEventPool — singleton CRTP pool for CUevent handles. +// --------------------------------------------------------------------------- +class CudaEventPool : public SimplePool +{ +public: + static CudaEventPool& instance(); + +private: + CudaEventPool(); +}; + +// --------------------------------------------------------------------------- +// CudaStreamPool — singleton CRTP pool for CUstream handles. +// --------------------------------------------------------------------------- +class CudaStreamPool : public SimplePool +{ +public: + static CudaStreamPool& instance(); + +private: + CudaStreamPool(); +}; + +// --------------------------------------------------------------------------- +// CachedCudaEvent — pooled CUevent (no timing). +// Mirrors _utils.py::CachedCudaEvent. +// +// On construction: gets an event from the global pool and records it to stream. +// Copyable: copies share the same underlying CUevent via shared_ptr. +// Last copy returns the event to the pool. +// NULL sentinel: always considered complete, no event in flight. +// --------------------------------------------------------------------------- +class CachedCudaEvent +{ +public: + // NULL sentinel: always considered complete, no event in flight. + static CachedCudaEvent makeNull() noexcept; + + // Normal constructor: gets an event and records it on stream. + explicit CachedCudaEvent(CudaStream stream); + + // Copyable and movable (shared ownership of the underlying CUevent). + CachedCudaEvent(CachedCudaEvent const&) = default; + CachedCudaEvent& operator=(CachedCudaEvent const&) = default; + CachedCudaEvent(CachedCudaEvent&&) noexcept = default; + CachedCudaEvent& operator=(CachedCudaEvent&&) noexcept = default; + ~CachedCudaEvent() = default; + + // Query if the recorded work is done. + bool queryComplete(); + + // Block until complete. + void synchronize(); + + // Insert a stream dependency on this event. + void waitInStream(CudaStream stream) const; + + // True if no CUevent is held (NULL or already closed by any copy). + [[nodiscard]] bool isClosed() const noexcept + { + return !mEvent || !*mEvent; + } + + // Release the event back to pool. Visible to ALL copies sharing this event. + void close(); + + // Raw CUevent handle. Returns nullptr for NULL/closed events. + // Also serves as identity key for deduplication. + [[nodiscard]] CUevent handle() const noexcept + { + return isClosed() ? nullptr : mEvent->get(); + } + +private: + explicit CachedCudaEvent() noexcept = default; // used by makeNull() + + // Shared ownership of the PoolItem. close() resets the inner unique_ptr, + // visible to all copies. Last shared_ptr drop is a no-op (inner already empty). + std::shared_ptr mEvent; +}; + +// --------------------------------------------------------------------------- +// Stream-level helpers. +// --------------------------------------------------------------------------- + +// Wait for all events on the given stream. Deduplicates internally. +// Mirrors Python's stream_wait_events() which converts to set() before iterating. +inline void streamWaitEvents(CudaStream stream, std::vector const& events) +{ + thread_local std::vector handles; + handles.clear(); + handles.reserve(events.size()); + for (auto const* ev : events) + { + if (ev && !ev->isClosed()) + handles.push_back(ev->handle()); + } + std::sort(handles.begin(), handles.end()); + handles.erase(std::unique(handles.begin(), handles.end()), handles.end()); + for (CUevent h : handles) + cuCheck(cuStreamWaitEvent(reinterpret_cast(stream), h, 0)); +} + +// Synchronize and close all events. Deduplicates internally. +// Mirrors Python's set()-based synchronization pattern. +inline void synchronizeAll(std::vector const& events) +{ + thread_local std::vector handles; + handles.clear(); + handles.reserve(events.size()); + for (auto* ev : events) + { + if (!ev->isClosed()) + handles.push_back(ev->handle()); + } + std::sort(handles.begin(), handles.end()); + handles.erase(std::unique(handles.begin(), handles.end()), handles.end()); + for (CUevent h : handles) + cuCheck(cuEventSynchronize(h)); + for (auto* ev : events) + ev->close(); +} + +// --------------------------------------------------------------------------- +// CachedCudaStream — pooled non-blocking CUstream. +// Mirrors _utils.py::CachedCudaStream. +// --------------------------------------------------------------------------- +class CachedCudaStream +{ +public: + CachedCudaStream(); + + CachedCudaStream(CachedCudaStream&&) noexcept = default; + CachedCudaStream& operator=(CachedCudaStream&&) noexcept = default; + CachedCudaStream(CachedCudaStream const&) = delete; + CachedCudaStream& operator=(CachedCudaStream const&) = delete; + + [[nodiscard]] CUstream handle() const noexcept + { + return mPoolItem.get(); // CUstream = CUstream_st* + } + + // Wait for a single event on this stream. + void waitEvent(CachedCudaEvent const& event) const + { + event.waitInStream(reinterpret_cast(handle())); + } + + // Wait for all events on this stream. Deduplicates internally. + void waitEvents(std::vector const& events) + { + streamWaitEvents(reinterpret_cast(handle()), events); + } + + CachedCudaEvent recordEvent(); + void synchronize(); + +private: + CudaStreamPool::PoolItem mPoolItem; // returns to pool on destruction +}; + +// --------------------------------------------------------------------------- +// TemporaryCudaStream — pooled stream with finish-event tracking. +// Mirrors Python's TemporaryCudaStream context manager. +// +// Usage (matches Python's `with TemporaryCudaStream(events) as stream:`): +// +// TemporaryCudaStream tempStream(priorEvents); +// { +// auto scope = tempStream.enter(); // __enter__ +// launchKernel(tempStream.get()); +// } // ~Scope → __exit__ records finish event +// auto ev = tempStream.takeFinishEvent(); // after with block +// +// --------------------------------------------------------------------------- +class TemporaryCudaStream +{ +public: + // Acquire a stream from pool and issue cuStreamWaitEvent for each prior event. + explicit TemporaryCudaStream(std::vector const& priorEvents); + + // Begin a scoped block. Destructor records the finish event (= Python __exit__). + // Skips recording during stack unwinding to match Python's `if not exc_type:` guard. + [[nodiscard]] auto enter() + { + int const exCount = std::uncaught_exceptions(); + return FuncGuard( + [this, exCount]() + { + if (std::uncaught_exceptions() == exCount) + mFinishEvent = mStream.recordEvent(); + }); + } + + [[nodiscard]] CUstream get() const noexcept + { + return mStream.handle(); + } + + // Consume the finish event recorded by Scope destructor. + [[nodiscard]] CachedCudaEvent takeFinishEvent() + { + auto result = std::move(mFinishEvent); + mFinishEvent = CachedCudaEvent::makeNull(); + return result; + } + + TemporaryCudaStream(TemporaryCudaStream const&) = delete; + TemporaryCudaStream& operator=(TemporaryCudaStream const&) = delete; + +private: + CachedCudaStream mStream; + CachedCudaEvent mFinishEvent = CachedCudaEvent::makeNull(); +}; + +// Merge multiple CUDA events into one. +// Returns makeNull() for 0 live events, the single live event for 1, +// or a TemporaryCudaStream-merged event for many. +// Mirrors Python's merge_events() utility. +CachedCudaEvent mergeEvents(std::vector& events); + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMem.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMem.cpp new file mode 100644 index 000000000000..e51cc2ceaa4b --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMem.cpp @@ -0,0 +1,334 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "kv_cache_manager_v2/utils/hostMem.h" +#include "kv_cache_manager_v2/exceptions.h" + +#include "tensorrt_llm/common/assert.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// --------------------------------------------------------------------------- +// Low-level helpers +// --------------------------------------------------------------------------- + +MemAddress hostMmap(size_t size) +{ + void* ptr = ::mmap(nullptr, size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + if (ptr == MAP_FAILED || ptr == nullptr) + { + throw HostOOMError(std::string("mmap failed: ") + std::strerror(errno)); + } + return reinterpret_cast(ptr); +} + +void hostMunmap(MemAddress ptr, size_t size) noexcept +{ + int ret = ::munmap(reinterpret_cast(ptr), size); + if (ret != 0) + { + std::fprintf(stderr, "munmap failed with errno %d\n", errno); + } +} + +MemAddress hostMremap(MemAddress ptr, size_t oldSize, size_t newSize) +{ + void* newPtr = ::mremap(reinterpret_cast(ptr), oldSize, newSize, MREMAP_MAYMOVE); + if (newPtr == MAP_FAILED || newPtr == nullptr) + { + throw HostOOMError(std::string("mremap failed: ") + std::strerror(errno)); + } + return reinterpret_cast(newPtr); +} + +void resizeFile(int fd, size_t newSize) +{ + off_t oldSize = ::lseek(fd, 0, SEEK_END); + if (static_cast(oldSize) < newSize) + { + int ret = ::posix_fallocate(fd, oldSize, static_cast(newSize - oldSize)); + if (ret != 0) + { + throw DiskOOMError("posix_fallocate failed: " + std::string(std::strerror(ret))); + } + } + else if (static_cast(oldSize) > newSize) + { + if (::ftruncate(fd, static_cast(newSize)) != 0) + { + throw DiskOOMError("ftruncate failed: " + std::string(std::strerror(errno))); + } + } +} + +bool hostUseThp() +{ + char const* value = std::getenv("TLLM_KV_CACHE_MANAGER_V2_THP"); + return value == nullptr || std::string_view(value) == "1"; +} + +int hostPrefaultThreads() +{ + char const* value = std::getenv("TLLM_KV_CACHE_MANAGER_V2_PREFAULT_THREADS"); + if (value != nullptr) + { + return std::stoi(value); + } + unsigned int const detectedCpuCount = std::thread::hardware_concurrency(); + unsigned int const cpuCount = detectedCpuCount == 0 ? 32 : detectedCpuCount; + return static_cast(std::min(64U, cpuCount / 2)); +} + +void hostMadvisePageMode(MemAddress ptr, size_t size, bool useThp, HostMadviseFn madviseFn) noexcept +{ + HostMadviseFn const fn = madviseFn != nullptr ? madviseFn : ::madvise; + int const advice = useThp ? MADV_HUGEPAGE : MADV_NOHUGEPAGE; + if (fn(reinterpret_cast(ptr), size, advice) != 0) + { + std::fprintf(stderr, "madvise failed with errno %d\n", errno); + } +} + +void hostPrefaultChunk(MemAddress ptr, size_t size, HostMadviseFn madviseFn, HostMemsetFn memsetFn) +{ + HostMadviseFn const advise = madviseFn != nullptr ? madviseFn : ::madvise; + HostMemsetFn const touch = memsetFn != nullptr ? memsetFn : ::memset; +#ifdef MADV_POPULATE_WRITE + if (advise(reinterpret_cast(ptr), size, MADV_POPULATE_WRITE) == 0) + { + return; + } + + int const errorCode = errno; + if (errorCode == EINVAL || errorCode == ENOSYS) + { + touch(reinterpret_cast(ptr), 0, size); + return; + } + if (errorCode == ENOMEM) + { + throw HostOOMError("madvise(MADV_POPULATE_WRITE) failed: " + std::string(std::strerror(errorCode))); + } + throw std::system_error(errorCode, std::generic_category(), "madvise(MADV_POPULATE_WRITE) failed"); +#else + // MADV_POPULATE_WRITE requires glibc >= 2.34 / Linux >= 5.14 headers and is not defined in + // older build environments (e.g. Rocky8 package-sanity images). Fall back to explicitly + // touching the pages to force population, matching the EINVAL/ENOSYS runtime path above. + (void) advise; + touch(reinterpret_cast(ptr), 0, size); +#endif +} + +// --------------------------------------------------------------------------- +// HostMem implementation +// --------------------------------------------------------------------------- + +bool HostMem::shouldUseChunkedRegistration() +{ + struct utsname u + { + }; + + if (::uname(&u) != 0) + { + return false; + } + // Check for Linux kernel 6.11, 6.12, 6.13 prefix. + std::string_view rel{u.release}; + for (auto prefix : {"6.11", "6.12", "6.13"}) + { + if (rel.substr(0, 4) == prefix) + { + return true; + } + } + return false; +} + +HostMem::HostMem(size_t size) + : mUseThp(hostUseThp()) +{ + if (size == 0) + { + return; + } + mAddr = hostMmap(size); + TLLM_CHECK_DEBUG(mAddr % kAlignment == 0); + mSize = size; + try + { + madvisePageMode(); + int const prefaultThreads = hostPrefaultThreads(); + if (prefaultThreads > 0) + { + parallelPrefault(prefaultThreads); + } + registerToCuda(); + } + catch (...) + { + unregisterFromCuda(); + hostMunmap(mAddr, mSize); + mAddr = 0; + mSize = 0; + throw; + } +} + +HostMem::~HostMem() +{ + destroy(); +} + +void HostMem::resize(size_t newSize) +{ + unregisterFromCuda(); + try + { + mAddr = hostMremap(mAddr, mSize, newSize); + TLLM_CHECK_DEBUG(mAddr % kAlignment == 0); + mSize = newSize; + madvisePageMode(); + } + catch (...) + { + registerToCuda(); + throw; + } + registerToCuda(); +} + +void HostMem::destroy() +{ + if (mAddr == 0) + { + return; + } + unregisterFromCuda(); + hostMunmap(mAddr, mSize); + mAddr = 0; + mSize = 0; +} + +void HostMem::madvisePageMode() +{ + TLLM_CHECK_DEBUG(mAddr && mSize); + hostMadvisePageMode(mAddr, mSize, mUseThp); +} + +void HostMem::parallelPrefault(int numThreads) +{ + size_t const numChunks = (mSize + kPrefaultChunkSize - 1) / kPrefaultChunkSize; + int const workerCount = std::min(numThreads, static_cast(numChunks)); + std::atomic_size_t nextChunk{0}; + std::atomic_bool failed{false}; + std::exception_ptr error; + std::mutex errorMutex; + + auto worker = [&]() + { + while (!failed.load()) + { + size_t const chunkIndex = nextChunk.fetch_add(1); + if (chunkIndex >= numChunks) + { + return; + } + size_t const offset = chunkIndex * kPrefaultChunkSize; + size_t const chunkSize = std::min(kPrefaultChunkSize, mSize - offset); + try + { + hostPrefaultChunk(mAddr + offset, chunkSize); + } + catch (...) + { + failed.store(true); + std::lock_guard lock(errorMutex); + if (error == nullptr) + { + error = std::current_exception(); + } + return; + } + } + }; + + std::vector workers; + workers.reserve(static_cast(workerCount)); + for (int i = 0; i < workerCount; ++i) + { + workers.emplace_back(worker); + } + for (auto& thread : workers) + { + thread.join(); + } + if (error != nullptr) + { + std::rethrow_exception(error); + } +} + +void HostMem::registerToCuda() +{ + TLLM_CHECK_DEBUG(mNumRegisteredChunks == 0); + static bool chunked = shouldUseChunkedRegistration(); + + size_t chunkSize = (chunked && mSize > kChunkSize) ? kChunkSize : mSize; + for (size_t offset = 0; offset < mSize; offset += chunkSize) + { + size_t sz = std::min(chunkSize, mSize - offset); + CUresult res = cuMemHostRegister( + reinterpret_cast(mAddr + offset), sz, CU_MEMHOSTREGISTER_PORTABLE | CU_MEMHOSTREGISTER_DEVICEMAP); + cuCheck(res); + ++mNumRegisteredChunks; + } +} + +void HostMem::unregisterFromCuda() +{ + static bool chunked = shouldUseChunkedRegistration(); + size_t chunkSize = (chunked && mSize > kChunkSize) ? kChunkSize : mSize; + for (size_t offset = 0; offset < mSize && mNumRegisteredChunks > 0; offset += chunkSize) + { + cuMemHostUnregister(reinterpret_cast(mAddr + offset)); + --mNumRegisteredChunks; + } + TLLM_CHECK_DEBUG(mNumRegisteredChunks == 0); +} + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMem.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMem.h new file mode 100644 index 000000000000..cf498653d2ee --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMem.h @@ -0,0 +1,100 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "kv_cache_manager_v2/common.h" + +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// --------------------------------------------------------------------------- +// HostMem — mmap-backed pinned host memory, resizable via mremap. +// Mirrors _utils.py::HostMem. +// +// Memory is: +// - Anonymous private mmap +// - Advised according to TLLM_KV_CACHE_MANAGER_V2_THP +// - Optionally prefaulted in parallel before CUDA registration +// - Registered to CUDA as page-locked (CU_MEMHOSTREGISTER_DEVICEMAP) +// +// On kernels 6.11/6.12/6.13, pinning is chunked in 2GB pieces to work around +// a kernel bug that prevents pinning more than 2GB in one call. +// --------------------------------------------------------------------------- +class HostMem +{ +public: + static constexpr size_t kAlignment = 4096; // 4 KB + static constexpr size_t kChunkSize = 2ULL << 30; // 2 GB + static constexpr size_t kPrefaultChunkSize = 512ULL << 20; // 512 MB + + explicit HostMem(size_t size); + ~HostMem(); + + HostMem(HostMem const&) = delete; + HostMem& operator=(HostMem const&) = delete; + + // Resize in-place (mremap, preserves data). Unregisters and re-registers with CUDA. + void resize(size_t newSize); + + // Unregister from CUDA and unmap. Safe to call multiple times. + void destroy(); + + MemAddress address() const noexcept + { + return mAddr; + } + + size_t size() const noexcept + { + return mSize; + } + +private: + void registerToCuda(); + void unregisterFromCuda(); + void madvisePageMode(); + void parallelPrefault(int numThreads); + + MemAddress mAddr = 0; + size_t mSize = 0; + int mNumRegisteredChunks = 0; + bool mUseThp = true; + + // Detect kernel version once at startup. + static bool shouldUseChunkedRegistration(); +}; + +// --------------------------------------------------------------------------- +// Low-level wrappers used internally (also exposed for storage pool use). +// --------------------------------------------------------------------------- +MemAddress hostMmap(size_t size); // throws HostOOMError +void hostMunmap(MemAddress ptr, size_t size) noexcept; +MemAddress hostMremap(MemAddress ptr, size_t oldSize, size_t newSize); // throws HostOOMError +void resizeFile(int fd, size_t newSize); // throws DiskOOMError + +using HostMadviseFn = int (*)(void*, size_t, int); +using HostMemsetFn = void* (*) (void*, int, size_t); + +bool hostUseThp(); +int hostPrefaultThreads(); +void hostMadvisePageMode(MemAddress ptr, size_t size, bool useThp, HostMadviseFn madviseFn = nullptr) noexcept; +void hostPrefaultChunk(MemAddress ptr, size_t size, HostMadviseFn madviseFn = nullptr, HostMemsetFn memsetFn = nullptr); + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/math.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/math.h new file mode 100644 index 000000000000..bdd1fd503708 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/math.h @@ -0,0 +1,384 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "kv_cache_manager_v2/utils/typedIndex.h" + +#include "tensorrt_llm/common/assert.h" +#include +#include +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// --------------------------------------------------------------------------- +// Integer math helpers (mirrors _utils.py) +// --------------------------------------------------------------------------- + +template +[[nodiscard]] inline T divUp(T x, T y) noexcept +{ + return (x + y - 1) / y; +} + +template +[[nodiscard]] inline T roundUp(T x, T y) noexcept +{ + return divUp(x, y) * y; +} + +template +[[nodiscard]] inline T roundDown(T x, T y) noexcept +{ + return (x / y) * y; +} + +template +[[nodiscard]] inline T exactDiv(T x, T y) +{ + TLLM_CHECK_DEBUG(x % y == 0); + return x / y; +} + +template +[[nodiscard]] inline bool inRange(T x, T lower, T upper) noexcept +{ + return lower <= x && x < upper; +} + +// Returns the intersection of [a.first, a.second) and [b.first, b.second). +// If the ranges are disjoint, the result has first > second; the caller must +// check first < second before using it as a non-empty range. +template +[[nodiscard]] inline std::pair overlap(std::pair a, std::pair b) noexcept +{ + T lo = a.first > b.first ? a.first : b.first; + T hi = a.second < b.second ? a.second : b.second; + return {lo, hi}; +} + +// Extract an attribute from the first element and assert (debug) all elements agree. +// Mirrors Python's get_uniform_attribute(_utils.py:202). +template +[[nodiscard]] auto getUniformAttribute(Range const& range, Func&& func) -> decltype(func(*range.begin())) +{ + auto it = range.begin(); + TLLM_CHECK_DEBUG(it != range.end()); + auto result = func(*it); + TLLM_CHECK_DEBUG(std::all_of(range.begin(), range.end(), [&](auto const& item) { return func(item) == result; })); + return result; +} + +// Find the first index in [begin, end) where predicate is true. +// Returns distance(begin, end) if not found. +// Mirrors Python's find_index(_utils.py:366). +template +[[nodiscard]] int findIndex(Iter begin, Iter end, Pred pred) +{ + return static_cast(std::distance(begin, std::find_if(begin, end, pred))); +} + +// Steal items matching predicate from the container and return them (stable). +// Single-pass O(n), mirrors Python's remove_if(_utils.py:174). +template +std::vector stealIf(std::vector& original, Pred pred) +{ + std::vector removed; + size_t writeIdx = 0; + for (size_t i = 0; i < original.size(); ++i) + { + if (pred(original[i])) + removed.push_back(std::move(original[i])); + else + original[writeIdx++] = std::move(original[i]); + } + original.erase(original.begin() + static_cast(writeIdx), original.end()); + return removed; +} + +// Group items by a classifier function, returning a map of key → vector of items. +// Mirrors Python's partition(_utils.py:195). +template +auto partition(std::vector const& items, Classifier classifier) + -> std::map())), std::vector> +{ + using Key = decltype(classifier(std::declval())); + std::map> result; + for (auto const& item : items) + result[classifier(item)].push_back(item); + return result; +} + +// Normalize a vector of values to a ratio vector summing to 1.0. +// Mirrors Python's typed_map(values, lambda x: x / total). +template +[[nodiscard]] std::vector normalizeToRatio(std::vector const& values) +{ + auto total = std::accumulate(values.begin(), values.end(), static_cast(0)); + TLLM_CHECK_DEBUG(total > 0); + std::vector ratio(values.size()); + for (size_t i = 0; i < values.size(); ++i) + ratio[i] = static_cast(values[i]) / static_cast(total); + return ratio; +} + +template +[[nodiscard]] TypedVec normalizeToRatio(TypedVec const& values) +{ + auto total = std::accumulate(values.begin(), values.end(), static_cast(0)); + TLLM_CHECK_DEBUG(total > 0); + TypedVec ratio(values.size()); + for (Index index{0}; index < values.size(); ++index) + { + ratio[index] = static_cast(values[index]) / static_cast(total); + } + return ratio; +} + +// --------------------------------------------------------------------------- +// HalfOpenRange — a half-open range [beg, end). Empty when beg >= end. +// Mirrors _utils.py::HalfOpenRange. +// --------------------------------------------------------------------------- +template +struct HalfOpenRange +{ + using IndexType = Index; + using DifferenceType = decltype(std::declval() - std::declval()); + + Index beg{0}; + Index end{0}; + + constexpr HalfOpenRange() noexcept = default; + + template + constexpr HalfOpenRange(Beg b, End e) noexcept + : beg(Index{b}) + , end(Index{e}) + { + } + + [[nodiscard]] DifferenceType length() const noexcept + { + return beg < end ? end - beg : DifferenceType{0}; + } + + [[nodiscard]] bool empty() const noexcept + { + return beg >= end; + } + + explicit operator bool() const noexcept + { + return beg < end; + } + + bool operator==(HalfOpenRange const& o) const noexcept + { + if (beg >= end && o.beg >= o.end) + return true; + return beg == o.beg && end == o.end; + } + + bool operator!=(HalfOpenRange const& o) const noexcept + { + return !(*this == o); + } + + // Membership test: is value in [beg, end)? + // Mirrors Python's HalfOpenRange.__contains__. + [[nodiscard]] bool contains(Index value) const noexcept + { + return beg <= value && value < end; + } +}; + +// Returns the intersection of two half-open ranges. +// The result may be empty (beg >= end), which is safe to chain. +template +[[nodiscard]] inline HalfOpenRange intersect(HalfOpenRange a, HalfOpenRange b) noexcept +{ + return {std::max(a.beg, b.beg), std::min(a.end, b.end)}; +} + +// --------------------------------------------------------------------------- +// DynamicBitset — resizable bitset using 64-bit words. +// Mirrors _utils.py::DynamicBitset. +// --------------------------------------------------------------------------- +class DynamicBitset +{ +public: + explicit DynamicBitset(size_t capacity) + : mWords(divUp(capacity, size_t{64}), uint64_t{0}) + , mNumSetBits(0) + { + } + + void set(size_t index) + { + if (!get(index)) + { + mWords[index / 64] |= (uint64_t{1} << (index % 64)); + ++mNumSetBits; + } + } + + [[nodiscard]] bool get(size_t index) const noexcept + { + return (mWords[index / 64] & (uint64_t{1} << (index % 64))) != 0; + } + + void clear(size_t index) + { + if (get(index)) + { + mWords[index / 64] &= ~(uint64_t{1} << (index % 64)); + --mNumSetBits; + } + } + + [[nodiscard]] size_t numSetBits() const noexcept + { + return mNumSetBits; + } + + void resize(size_t newCapacity) + { + size_t const oldWords = mWords.size(); + size_t const newWords = divUp(newCapacity, size_t{64}); + + // When the capacity shrinks, every set bit at or above newCapacity is + // dropped. Account for those bits so numSetBits() stays accurate, and + // mask the retained partial word so anySet() cannot observe stale bits. + // This covers both fewer-words and same-word-count (newWords == oldWords + // with a smaller newCapacity) shrinks. + if (newWords <= oldWords) + { + for (size_t w = newWords; w < oldWords; ++w) + { + mNumSetBits -= static_cast(__builtin_popcountll(mWords[w])); + } + if (newWords >= 1 && newCapacity % 64 != 0) + { + uint64_t const keepMask = (uint64_t{1} << (newCapacity % 64)) - 1; + uint64_t& word = mWords[newWords - 1]; + mNumSetBits -= static_cast(__builtin_popcountll(word & ~keepMask)); + word &= keepMask; + } + } + + // Grow (zero-filled) or shrink storage to the new word count. + mWords.resize(newWords, uint64_t{0}); + } + + // Returns true if any bit in [start, end) is set. + [[nodiscard]] bool anySet(size_t start, size_t end) const noexcept + { + if (start >= end) + { + return false; + } + size_t startWord = start / 64; + size_t endWord = (end - 1) / 64; + uint64_t startMask = ~uint64_t{0} << (start % 64); + if (startWord == endWord) + { + size_t bitsInWord = end % 64; + uint64_t endMask = bitsInWord ? ((uint64_t{1} << bitsInWord) - 1) : ~uint64_t{0}; + return (mWords[startWord] & startMask & endMask) != 0; + } + if (mWords[startWord] & startMask) + { + return true; + } + for (size_t w = startWord + 1; w < endWord; ++w) + { + if (mWords[w]) + { + return true; + } + } + size_t bitsInLastWord = end % 64; + if (bitsInLastWord == 0) + { + return mWords[endWord] != 0; + } + return (mWords[endWord] & ((uint64_t{1} << bitsInLastWord) - 1)) != 0; + } + +private: + std::vector mWords; + size_t mNumSetBits; +}; + +// --------------------------------------------------------------------------- +// Array2D — row-major 2D array with typed row/column indices. +// Mirrors _utils.py::Array2D. +// --------------------------------------------------------------------------- +template +class Array2D +{ +public: + Array2D(int rows, int cols, T initVal = T{}) + : mData(static_cast(rows * cols), initVal) + , mCols(cols) + { + } + + T& operator()(int row, int col) noexcept + { + return mData[static_cast(row * mCols + col)]; + } + + T const& operator()(int row, int col) const noexcept + { + return mData[static_cast(row * mCols + col)]; + } + + [[nodiscard]] int rows() const noexcept + { + return static_cast(mData.size()) / mCols; + } + + [[nodiscard]] int cols() const noexcept + { + return mCols; + } + + // Pointer to start of row (for slicing / iteration). + T* rowData(int row) noexcept + { + return mData.data() + row * mCols; + } + + T const* rowData(int row) const noexcept + { + return mData.data() + row * mCols; + } + +private: + std::vector mData; + int mCols; +}; + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/sharedPtr.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/sharedPtr.h new file mode 100644 index 000000000000..604a8a0e9c12 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/sharedPtr.h @@ -0,0 +1,593 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Non-atomic SharedPtr / WeakPtr — drop-in replacements for std::shared_ptr +// and std::weak_ptr that use plain int refcounts instead of atomics. +// +// Motivation: std::shared_ptr uses atomic increments/decrements for thread +// safety. On ARM (Grace), atomics are significantly more expensive than on +// x86. The KV cache manager's shared_ptr usage is entirely single-threaded, +// so the atomics are pure overhead. +// +// Usage: +// SharedPtr replaces std::shared_ptr +// WeakPtr replaces std::weak_ptr +// EnableSharedFromThis replaces std::enable_shared_from_this +// makeShared(args...) replaces std::make_shared(args...) +// dynamicPointerCast(p) replaces std::dynamic_pointer_cast(p) +// toStd(SharedPtr) bridges SharedPtr → std::shared_ptr + +#pragma once + +#include "tensorrt_llm/common/assert.h" +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// Forward declarations. +template +class SharedPtr; +template +class WeakPtr; +template +class EnableSharedFromThis; + +// ========================================================================= +// Control block +// ========================================================================= + +namespace detail +{ + +struct ControlBlockBase +{ + int strongCount = 1; + int weakCount = 1; // +1 bias while strongCount > 0 + + virtual void destroyObject() noexcept = 0; + virtual void deallocate() noexcept = 0; + + void addStrongRef() noexcept + { + ++strongCount; + } + + void releaseStrongRef() noexcept + { + if (--strongCount == 0) + { + destroyObject(); + releaseWeakRef(); // release the bias + } + } + + void addWeakRef() noexcept + { + ++weakCount; + } + + void releaseWeakRef() noexcept + { + if (--weakCount == 0) + { + deallocate(); + } + } + +protected: + ~ControlBlockBase() = default; // prevent polymorphic delete via base +}; + +template +struct InplaceControlBlock final : ControlBlockBase +{ + // Aligned uninitialized storage for T. + alignas(T) unsigned char storage[sizeof(T)]; + + T* ptr() noexcept + { + return std::launder(reinterpret_cast(&storage)); + } + + void destroyObject() noexcept override + { + ptr()->~T(); + } + + void deallocate() noexcept override + { + delete this; + } +}; + +// --------------------------------------------------------------------------- +// EnableSharedFromThis detection — handles polymorphic inheritance. +// E.g. makeShared() detects EnableSharedFromThis +// inherited through Page. +// --------------------------------------------------------------------------- +template +auto detectEnableImpl(EnableSharedFromThis*) -> EnableSharedFromThis; + +auto detectEnableImpl(...) -> void; + +template +using EnableBase = decltype(detectEnableImpl(std::declval())); + +template +inline constexpr bool hasEnable = !std::is_void_v>; + +} // namespace detail + +// ========================================================================= +// SharedPtr +// ========================================================================= + +template +class SharedPtr +{ +public: + // -- Constructors ------------------------------------------------------- + + SharedPtr() noexcept + : mPtr(nullptr) + , mCb(nullptr) + { + } + + SharedPtr(std::nullptr_t) noexcept // NOLINT(google-explicit-constructor) + : mPtr(nullptr) + , mCb(nullptr) + { + } + + SharedPtr(SharedPtr const& other) noexcept + : mPtr(other.mPtr) + , mCb(other.mCb) + { + if (mCb) + mCb->addStrongRef(); + } + + SharedPtr(SharedPtr&& other) noexcept + : mPtr(other.mPtr) + , mCb(other.mCb) + { + other.mPtr = nullptr; + other.mCb = nullptr; + } + + // Converting copy (U* implicitly convertible to T*). + template , int> = 0> + SharedPtr(SharedPtr const& other) noexcept // NOLINT(google-explicit-constructor) + : mPtr(other.mPtr) + , mCb(other.mCb) + { + if (mCb) + mCb->addStrongRef(); + } + + // Converting move. + template , int> = 0> + SharedPtr(SharedPtr&& other) noexcept // NOLINT(google-explicit-constructor) + : mPtr(other.mPtr) + , mCb(other.mCb) + { + other.mPtr = nullptr; + other.mCb = nullptr; + } + + ~SharedPtr() + { + if (mCb) + mCb->releaseStrongRef(); + } + + // -- Assignment --------------------------------------------------------- + + SharedPtr& operator=(SharedPtr const& other) noexcept + { + if (this != &other) + { + SharedPtr tmp(other); + swap(tmp); + } + return *this; + } + + SharedPtr& operator=(SharedPtr&& other) noexcept + { + if (this != &other) + { + SharedPtr tmp(std::move(other)); + swap(tmp); + } + return *this; + } + + template , int> = 0> + SharedPtr& operator=(SharedPtr const& other) noexcept + { + SharedPtr tmp(other); + swap(tmp); + return *this; + } + + template , int> = 0> + SharedPtr& operator=(SharedPtr&& other) noexcept + { + SharedPtr tmp(std::move(other)); + swap(tmp); + return *this; + } + + SharedPtr& operator=(std::nullptr_t) noexcept + { + reset(); + return *this; + } + + // -- Observers ---------------------------------------------------------- + + T* get() const noexcept + { + return mPtr; + } + + T& operator*() const noexcept + { + return *mPtr; + } + + T* operator->() const noexcept + { + return mPtr; + } + + explicit operator bool() const noexcept + { + return mPtr != nullptr; + } + + int useCount() const noexcept + { + return mCb ? mCb->strongCount : 0; + } + + // -- Modifiers ---------------------------------------------------------- + + void reset() noexcept + { + SharedPtr().swap(*this); + } + + void swap(SharedPtr& other) noexcept + { + std::swap(mPtr, other.mPtr); + std::swap(mCb, other.mCb); + } + + // -- Comparisons -------------------------------------------------------- + + bool operator==(SharedPtr const& other) const noexcept + { + return mPtr == other.mPtr; + } + + bool operator!=(SharedPtr const& other) const noexcept + { + return mPtr != other.mPtr; + } + + bool operator==(std::nullptr_t) const noexcept + { + return mPtr == nullptr; + } + + bool operator!=(std::nullptr_t) const noexcept + { + return mPtr != nullptr; + } + +private: + // Aliasing constructor — shares control block from `owner`, stores `ptr`. + // Used by dynamicPointerCast and EnableSharedFromThis. + SharedPtr(detail::ControlBlockBase* cb, T* ptr) noexcept + : mPtr(ptr) + , mCb(cb) + { + if (mCb) + mCb->addStrongRef(); + } + + template + friend class SharedPtr; + template + friend class WeakPtr; + template + friend class EnableSharedFromThis; + template + friend SharedPtr makeShared(Args&&... args); + template + friend SharedPtr dynamicPointerCast(SharedPtr const&); + template + friend std::shared_ptr toStd(SharedPtr const&); + + T* mPtr; + detail::ControlBlockBase* mCb; +}; + +// Free-standing comparison with nullptr (reversed operand order). +template +bool operator==(std::nullptr_t, SharedPtr const& sp) noexcept +{ + return sp == nullptr; +} + +template +bool operator!=(std::nullptr_t, SharedPtr const& sp) noexcept +{ + return sp != nullptr; +} + +// ========================================================================= +// WeakPtr +// ========================================================================= + +template +class WeakPtr +{ +public: + WeakPtr() noexcept + : mPtr(nullptr) + , mCb(nullptr) + { + } + + WeakPtr(SharedPtr const& sp) noexcept // NOLINT(google-explicit-constructor) + : mPtr(sp.mPtr) + , mCb(sp.mCb) + { + if (mCb) + mCb->addWeakRef(); + } + + // Converting constructor from SharedPtr. + template , int> = 0> + WeakPtr(SharedPtr const& sp) noexcept // NOLINT(google-explicit-constructor) + : mPtr(sp.mPtr) + , mCb(sp.mCb) + { + if (mCb) + mCb->addWeakRef(); + } + + WeakPtr(WeakPtr const& other) noexcept + : mPtr(other.mPtr) + , mCb(other.mCb) + { + if (mCb) + mCb->addWeakRef(); + } + + WeakPtr(WeakPtr&& other) noexcept + : mPtr(other.mPtr) + , mCb(other.mCb) + { + other.mPtr = nullptr; + other.mCb = nullptr; + } + + ~WeakPtr() + { + if (mCb) + mCb->releaseWeakRef(); + } + + // -- Assignment --------------------------------------------------------- + + WeakPtr& operator=(SharedPtr const& sp) noexcept + { + WeakPtr tmp(sp); + swap(tmp); + return *this; + } + + template , int> = 0> + WeakPtr& operator=(SharedPtr const& sp) noexcept + { + WeakPtr tmp(sp); + swap(tmp); + return *this; + } + + WeakPtr& operator=(WeakPtr const& other) noexcept + { + if (this != &other) + { + WeakPtr tmp(other); + swap(tmp); + } + return *this; + } + + WeakPtr& operator=(WeakPtr&& other) noexcept + { + if (this != &other) + { + WeakPtr tmp(std::move(other)); + swap(tmp); + } + return *this; + } + + // -- Observers ---------------------------------------------------------- + + bool expired() const noexcept + { + return !mCb || mCb->strongCount == 0; + } + + SharedPtr lock() const noexcept + { + if (expired()) + return SharedPtr(); + // Object still alive — construct a SharedPtr sharing the control block. + SharedPtr result; + result.mPtr = mPtr; + result.mCb = mCb; + mCb->addStrongRef(); + return result; + } + + // -- Modifiers ---------------------------------------------------------- + + void reset() noexcept + { + WeakPtr().swap(*this); + } + + void swap(WeakPtr& other) noexcept + { + std::swap(mPtr, other.mPtr); + std::swap(mCb, other.mCb); + } + +private: + template + friend class SharedPtr; + template + friend class WeakPtr; + template + friend class EnableSharedFromThis; + + T* mPtr; + detail::ControlBlockBase* mCb; +}; + +// ========================================================================= +// EnableSharedFromThis +// ========================================================================= + +template +class EnableSharedFromThis +{ +public: + SharedPtr sharedFromThis() + { + auto sp = mWeakThis.lock(); + TLLM_CHECK_DEBUG_WITH_INFO(sp, "sharedFromThis() called on object not owned by SharedPtr"); + return sp; + } + + SharedPtr sharedFromThis() const + { + auto sp = mWeakThis.lock(); + TLLM_CHECK_DEBUG_WITH_INFO(sp, "sharedFromThis() called on object not owned by SharedPtr"); + // Convert SharedPtr to SharedPtr via the converting constructor. + return sp; + } + +protected: + EnableSharedFromThis() noexcept = default; + + // Copy/move must NOT copy mWeakThis — the new object is a distinct entity. + EnableSharedFromThis(EnableSharedFromThis const&) noexcept {} + + EnableSharedFromThis& operator=(EnableSharedFromThis const&) noexcept + { + return *this; + } + + ~EnableSharedFromThis() = default; + +private: + template + friend SharedPtr makeShared(Args&&... args); + + mutable WeakPtr mWeakThis; +}; + +// ========================================================================= +// makeShared(args...) +// ========================================================================= + +template +SharedPtr makeShared(Args&&... args) +{ + auto* cb = new detail::InplaceControlBlock(); + try + { + new (cb->storage) T(std::forward(args)...); + } + catch (...) + { + // T's constructor threw — control block was never fully initialized. + // Release directly; destroyObject() must not run. + delete cb; + throw; + } + + SharedPtr result; + result.mPtr = cb->ptr(); + result.mCb = cb; + + // Wire up EnableSharedFromThis if T (or a base) derives from it. + if constexpr (detail::hasEnable) + { + result.mPtr->detail::template EnableBase::mWeakThis = result; + } + + return result; +} + +// ========================================================================= +// dynamicPointerCast(SharedPtr) +// ========================================================================= + +template +SharedPtr dynamicPointerCast(SharedPtr const& src) +{ + To* raw = dynamic_cast(src.get()); + if (!raw) + return SharedPtr(); + // Aliasing: share control block, different pointer. + return SharedPtr(src.mCb, raw); +} + +// ========================================================================= +// toStd — one-way bridge to std::shared_ptr for nanobind boundary +// ========================================================================= + +template +std::shared_ptr toStd(SharedPtr const& sp) +{ + if (!sp) + return nullptr; + // The custom deleter captures a copy of the SharedPtr, keeping + // the non-atomic refcount alive. When the last std::shared_ptr + // copy dies, the captured SharedPtr is destroyed and releases its + // non-atomic reference. + SharedPtr copy(sp); + return std::shared_ptr(sp.get(), [captured = std::move(copy)](T*) mutable { captured.reset(); }); +} + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/typedIndex.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/typedIndex.h new file mode 100644 index 000000000000..bc50645d4941 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/typedIndex.h @@ -0,0 +1,377 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/assert.h" +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +template +class StrongIndex +{ + static_assert(std::is_integral::value, "StrongIndex requires an integral underlying type"); + +public: + using ValueType = T; + + constexpr StrongIndex() noexcept = default; + + explicit constexpr StrongIndex(T value) noexcept + : mValue(value) + { + } + + [[nodiscard]] constexpr T value() const noexcept + { + return mValue; + } + + template ::value>> + constexpr StrongIndex& operator+=(U rhs) noexcept + { + mValue = static_cast(mValue + static_cast(rhs)); + return *this; + } + + template ::value>> + constexpr StrongIndex& operator-=(U rhs) noexcept + { + mValue = static_cast(mValue - static_cast(rhs)); + return *this; + } + + constexpr StrongIndex& operator++() noexcept + { + ++mValue; + return *this; + } + + constexpr StrongIndex operator++(int) noexcept + { + StrongIndex old{*this}; + ++(*this); + return old; + } + + constexpr StrongIndex& operator--() noexcept + { + --mValue; + return *this; + } + + constexpr StrongIndex operator--(int) noexcept + { + StrongIndex old{*this}; + --(*this); + return old; + } + +private: + T mValue{DefaultValue}; +}; + +template +constexpr bool operator==(StrongIndex lhs, StrongIndex rhs) noexcept +{ + return lhs.value() == rhs.value(); +} + +template +constexpr bool operator!=(StrongIndex lhs, StrongIndex rhs) noexcept +{ + return !(lhs == rhs); +} + +template +constexpr bool operator<(StrongIndex lhs, StrongIndex rhs) noexcept +{ + return lhs.value() < rhs.value(); +} + +template +constexpr bool operator<(StrongIndex lhs, T rhs) noexcept +{ + return lhs.value() < rhs; +} + +template +constexpr bool operator>(StrongIndex lhs, StrongIndex rhs) noexcept +{ + return rhs < lhs; +} + +template +constexpr bool operator<=(StrongIndex lhs, StrongIndex rhs) noexcept +{ + return !(rhs < lhs); +} + +template +constexpr bool operator>=(StrongIndex lhs, StrongIndex rhs) noexcept +{ + return !(lhs < rhs); +} + +template +constexpr bool operator>=(StrongIndex lhs, T rhs) noexcept +{ + return lhs.value() >= rhs; +} + +template ::value>> +constexpr StrongIndex operator+(StrongIndex lhs, U rhs) noexcept +{ + lhs += rhs; + return lhs; +} + +template ::value>> +constexpr StrongIndex operator+(U lhs, StrongIndex rhs) noexcept +{ + rhs += lhs; + return rhs; +} + +template ::value>> +constexpr StrongIndex operator-(StrongIndex lhs, U rhs) noexcept +{ + lhs -= rhs; + return lhs; +} + +template +constexpr T operator-(StrongIndex lhs, StrongIndex rhs) noexcept +{ + return static_cast(lhs.value() - rhs.value()); +} + +template +[[nodiscard]] std::size_t toSizeT(StrongIndex index) noexcept +{ + if constexpr (std::is_signed::value) + { + TLLM_CHECK_DEBUG_WITH_INFO(index.value() >= 0, "StrongIndex value must be non-negative for size_t conversion"); + } + return static_cast(index.value()); +} + +template +class TypedVec +{ +public: + using IndexType = Index; + using ValueType = T; + using ContainerType = std::vector; + using iterator = typename ContainerType::iterator; + using const_iterator = typename ContainerType::const_iterator; + + TypedVec() = default; + + explicit TypedVec(Index count) + : mData(toSizeT(count)) + { + } + + TypedVec(Index count, T const& value) + : mData(toSizeT(count), value) + { + } + + TypedVec(std::initializer_list values) + : mData(values) + { + } + + explicit TypedVec(ContainerType data) + : mData(std::move(data)) + { + } + + [[nodiscard]] T& operator[](Index index) noexcept + { + return mData[toSizeT(index)]; + } + + [[nodiscard]] T const& operator[](Index index) const noexcept + { + return mData[toSizeT(index)]; + } + + [[nodiscard]] T& at(Index index) + { + return mData.at(toSizeT(index)); + } + + [[nodiscard]] T const& at(Index index) const + { + return mData.at(toSizeT(index)); + } + + [[nodiscard]] Index size() const noexcept + { + return Index{static_cast(mData.size())}; + } + + [[nodiscard]] std::size_t stdSize() const noexcept + { + return mData.size(); + } + + [[nodiscard]] bool empty() const noexcept + { + return mData.empty(); + } + + void clear() noexcept + { + mData.clear(); + } + + void reserve(Index count) + { + mData.reserve(toSizeT(count)); + } + + void resize(Index count) + { + mData.resize(toSizeT(count)); + } + + void resize(Index count, T const& value) + { + mData.resize(toSizeT(count), value); + } + + void push_back(T const& value) + { + mData.push_back(value); + } + + void push_back(T&& value) + { + mData.push_back(std::move(value)); + } + + template + T& emplace_back(Args&&... args) + { + return mData.emplace_back(std::forward(args)...); + } + + void pop_back() + { + mData.pop_back(); + } + + [[nodiscard]] T& front() + { + return mData.front(); + } + + [[nodiscard]] T const& front() const + { + return mData.front(); + } + + [[nodiscard]] T& back() + { + return mData.back(); + } + + [[nodiscard]] T const& back() const + { + return mData.back(); + } + + [[nodiscard]] iterator begin() noexcept + { + return mData.begin(); + } + + [[nodiscard]] const_iterator begin() const noexcept + { + return mData.begin(); + } + + [[nodiscard]] const_iterator cbegin() const noexcept + { + return mData.cbegin(); + } + + [[nodiscard]] iterator end() noexcept + { + return mData.end(); + } + + [[nodiscard]] const_iterator end() const noexcept + { + return mData.end(); + } + + [[nodiscard]] const_iterator cend() const noexcept + { + return mData.cend(); + } + + [[nodiscard]] ContainerType& raw() noexcept + { + return mData; + } + + [[nodiscard]] ContainerType const& raw() const noexcept + { + return mData; + } + + friend bool operator==(TypedVec const& lhs, TypedVec const& rhs) + { + return lhs.mData == rhs.mData; + } + + friend bool operator!=(TypedVec const& lhs, TypedVec const& rhs) + { + return !(lhs == rhs); + } + +private: + ContainerType mData; +}; + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 + +namespace std +{ + +template +struct hash> +{ + size_t operator()( + tensorrt_llm::batch_manager::kv_cache_manager_v2::StrongIndex index) const noexcept + { + return std::hash{}(index.value()); + } +}; + +} // namespace std diff --git a/cpp/tensorrt_llm/common/sha256/README.md b/cpp/tensorrt_llm/common/sha256/README.md new file mode 100644 index 000000000000..99c61e7a76ef --- /dev/null +++ b/cpp/tensorrt_llm/common/sha256/README.md @@ -0,0 +1,53 @@ +# SHA-256 (vendored from Bitcoin Core, modified) + +This directory contains the SHA-256 implementation from the +[Bitcoin Core](https://github.com/bitcoin/bitcoin) project, **reduced and +adapted by NVIDIA** to the minimal single-block hasher needed by TensorRT-LLM. +It provides a portable scalar SHA-256 with runtime dispatch to a +hardware-accelerated transform (x86 SHA-NI or ARMv8 crypto extensions), exposed +through the `CSHA256` class. + +TensorRT-LLM uses it in the C++ KVCacheManagerV2 (`blockRadixTree`) to hash +token sequences into block keys, byte-identically to the Python backend's +`hashlib.sha256` block keys. + +## Provenance + +- **Upstream project:** Bitcoin Core — https://github.com/bitcoin/bitcoin +- **Source path upstream:** `src/`, `src/crypto/`, and `src/compat` +- **Upstream commit:** `70d9ec7f3d452789d04dce81dc02db0b3b778bb5` (branch `master`) + +## Contents + +| File | Notes | +|------|-------| +| `sha256.h` | `CSHA256` API; `SHA256D64` declaration removed | +| `sha256.cpp` | Scalar core + `CSHA256` + slim runtime dispatch | +| `sha256_x86_shani.cpp` | x86 SHA-NI transform (guard changed) | +| `sha256_arm_shani.cpp` | ARMv8 crypto transform (guard changed) | +| `attributes.h` | `ALWAYS_INLINE` macro used by the x86 transform | +| `sha256_endian.h` | Big-endian helpers replacing upstream `common.h` | + +## NVIDIA modifications + +The files are reduced to the single-block path TensorRT-LLM needs. + +Changes vs. upstream: + +- **Removed** the public `SHA256D64` double-hash API (and its scalar + `TransformD64`) and the standalone SSE4 / SSE4.1 / AVX2 multi-block transform + files (unused — TensorRT-LLM only calls single-block `CSHA256`). The two + SHA-NI files retain their upstream 2-way `Transform_2way` helpers; these are + now unreferenced but were left in place to keep the transforms byte-close to + upstream. +- **Removed** the upstream support-header chain (`crypto/common.h`, + `compat/endian.h`, `compat/byteswap.h`, `compat/cpuid.h`), which pulled in + C++20 (``, ``). Endian helpers are now the small, + NVIDIA-authored `sha256_endian.h`; CPU detection uses a `CPUID` leaf-7 check + via `__get_cpuid_count` (x86) / `getauxval(AT_HWCAP)` (aarch64). The vendored + sources now build as plain **C++17**. +- **Flattened** the directory (no `crypto/` / `compat/` subtree) and switched + the HW-transform build guards from the upstream `ENABLE_*` macros to target + architecture macros. Build wiring (per-file `-msha` / `-march=armv8-a+crypto`) + lives in `cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/CMakeLists.txt`. + diff --git a/cpp/tensorrt_llm/common/sha256/attributes.h b/cpp/tensorrt_llm/common/sha256/attributes.h new file mode 100644 index 000000000000..275dad9f8ede --- /dev/null +++ b/cpp/tensorrt_llm/common/sha256/attributes.h @@ -0,0 +1,27 @@ +// Copyright (c) 2009-2010 Satoshi Nakamoto +// Copyright (c) 2009-present The Bitcoin Core developers +// Distributed under the MIT software license, see the accompanying +// file COPYING or http://www.opensource.org/licenses/mit-license.php. + +#ifndef BITCOIN_ATTRIBUTES_H +#define BITCOIN_ATTRIBUTES_H + +#if defined(__clang__) +# if __has_attribute(lifetimebound) +# define LIFETIMEBOUND [[clang::lifetimebound]] +# else +# define LIFETIMEBOUND +# endif +#else +# define LIFETIMEBOUND +#endif + +#if defined(__GNUC__) +# define ALWAYS_INLINE inline __attribute__((always_inline)) +#elif defined(_MSC_VER) +# define ALWAYS_INLINE __forceinline +#else +# error No known always_inline attribute for this platform. +#endif + +#endif // BITCOIN_ATTRIBUTES_H diff --git a/cpp/tensorrt_llm/common/sha256/sha256.cpp b/cpp/tensorrt_llm/common/sha256/sha256.cpp new file mode 100644 index 000000000000..c2200bc71586 --- /dev/null +++ b/cpp/tensorrt_llm/common/sha256/sha256.cpp @@ -0,0 +1,302 @@ +// SPDX-FileCopyrightText: Copyright (c) 2014-present The Bitcoin Core developers +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: MIT +// +// Permission is hereby granted, free of charge, to any person obtaining a +// copy of this software and associated documentation files (the "Software"), +// to deal in the Software without restriction, including without limitation +// the rights to use, copy, modify, merge, publish, distribute, sublicense, +// and/or sell copies of the Software, and to permit persons to whom the +// Software is furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in +// all copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +// DEALINGS IN THE SOFTWARE. +// +// Modifications for TensorRT-LLM: +// - Reduced to the single-block CSHA256 hasher. +// - Removed the SHA256D64 double-hash API and the SSE4/SSE4.1/AVX2 multi-block +// transforms; runtime dispatch selects scalar / x86 SHA-NI / ARMv8 crypto. +// - Replaced / endian helpers with sha256_endian.h. +// - CPU detection via CPUID leaf 7 (x86) / getauxval (aarch64); plain C++17 (no ). + +#include "sha256.h" +#include "sha256_endian.h" + +#include +#include +#include +#include + +#if defined(__x86_64__) || defined(__amd64__) || defined(__i386__) +#include +namespace sha256_x86_shani { +void Transform(uint32_t* s, const unsigned char* chunk, size_t blocks); +} +#endif + +#if defined(__aarch64__) +#include +#include +namespace sha256_arm_shani { +void Transform(uint32_t* s, const unsigned char* chunk, size_t blocks); +} +#endif + +// Internal implementation code. +namespace +{ +/// Internal SHA-256 implementation. +namespace sha256 +{ +uint32_t inline Ch(uint32_t x, uint32_t y, uint32_t z) { return z ^ (x & (y ^ z)); } +uint32_t inline Maj(uint32_t x, uint32_t y, uint32_t z) { return (x & y) | (z & (x | y)); } +uint32_t inline Sigma0(uint32_t x) { return (x >> 2 | x << 30) ^ (x >> 13 | x << 19) ^ (x >> 22 | x << 10); } +uint32_t inline Sigma1(uint32_t x) { return (x >> 6 | x << 26) ^ (x >> 11 | x << 21) ^ (x >> 25 | x << 7); } +uint32_t inline sigma0(uint32_t x) { return (x >> 7 | x << 25) ^ (x >> 18 | x << 14) ^ (x >> 3); } +uint32_t inline sigma1(uint32_t x) { return (x >> 17 | x << 15) ^ (x >> 19 | x << 13) ^ (x >> 10); } + +/** One round of SHA-256. */ +void inline Round(uint32_t a, uint32_t b, uint32_t c, uint32_t& d, uint32_t e, uint32_t f, uint32_t g, uint32_t& h, uint32_t k) +{ + uint32_t t1 = h + Sigma1(e) + Ch(e, f, g) + k; + uint32_t t2 = Sigma0(a) + Maj(a, b, c); + d += t1; + h = t1 + t2; +} + +/** Initialize SHA-256 state. */ +void inline Initialize(uint32_t* s) +{ + s[0] = 0x6a09e667ul; + s[1] = 0xbb67ae85ul; + s[2] = 0x3c6ef372ul; + s[3] = 0xa54ff53aul; + s[4] = 0x510e527ful; + s[5] = 0x9b05688cul; + s[6] = 0x1f83d9abul; + s[7] = 0x5be0cd19ul; +} + +/** Perform a number of SHA-256 transformations, processing 64-byte chunks. */ +void Transform(uint32_t* s, const unsigned char* chunk, size_t blocks) +{ + while (blocks--) { + uint32_t a = s[0], b = s[1], c = s[2], d = s[3], e = s[4], f = s[5], g = s[6], h = s[7]; + uint32_t w0, w1, w2, w3, w4, w5, w6, w7, w8, w9, w10, w11, w12, w13, w14, w15; + + Round(a, b, c, d, e, f, g, h, 0x428a2f98 + (w0 = ReadBE32(chunk + 0))); + Round(h, a, b, c, d, e, f, g, 0x71374491 + (w1 = ReadBE32(chunk + 4))); + Round(g, h, a, b, c, d, e, f, 0xb5c0fbcf + (w2 = ReadBE32(chunk + 8))); + Round(f, g, h, a, b, c, d, e, 0xe9b5dba5 + (w3 = ReadBE32(chunk + 12))); + Round(e, f, g, h, a, b, c, d, 0x3956c25b + (w4 = ReadBE32(chunk + 16))); + Round(d, e, f, g, h, a, b, c, 0x59f111f1 + (w5 = ReadBE32(chunk + 20))); + Round(c, d, e, f, g, h, a, b, 0x923f82a4 + (w6 = ReadBE32(chunk + 24))); + Round(b, c, d, e, f, g, h, a, 0xab1c5ed5 + (w7 = ReadBE32(chunk + 28))); + Round(a, b, c, d, e, f, g, h, 0xd807aa98 + (w8 = ReadBE32(chunk + 32))); + Round(h, a, b, c, d, e, f, g, 0x12835b01 + (w9 = ReadBE32(chunk + 36))); + Round(g, h, a, b, c, d, e, f, 0x243185be + (w10 = ReadBE32(chunk + 40))); + Round(f, g, h, a, b, c, d, e, 0x550c7dc3 + (w11 = ReadBE32(chunk + 44))); + Round(e, f, g, h, a, b, c, d, 0x72be5d74 + (w12 = ReadBE32(chunk + 48))); + Round(d, e, f, g, h, a, b, c, 0x80deb1fe + (w13 = ReadBE32(chunk + 52))); + Round(c, d, e, f, g, h, a, b, 0x9bdc06a7 + (w14 = ReadBE32(chunk + 56))); + Round(b, c, d, e, f, g, h, a, 0xc19bf174 + (w15 = ReadBE32(chunk + 60))); + + Round(a, b, c, d, e, f, g, h, 0xe49b69c1 + (w0 += sigma1(w14) + w9 + sigma0(w1))); + Round(h, a, b, c, d, e, f, g, 0xefbe4786 + (w1 += sigma1(w15) + w10 + sigma0(w2))); + Round(g, h, a, b, c, d, e, f, 0x0fc19dc6 + (w2 += sigma1(w0) + w11 + sigma0(w3))); + Round(f, g, h, a, b, c, d, e, 0x240ca1cc + (w3 += sigma1(w1) + w12 + sigma0(w4))); + Round(e, f, g, h, a, b, c, d, 0x2de92c6f + (w4 += sigma1(w2) + w13 + sigma0(w5))); + Round(d, e, f, g, h, a, b, c, 0x4a7484aa + (w5 += sigma1(w3) + w14 + sigma0(w6))); + Round(c, d, e, f, g, h, a, b, 0x5cb0a9dc + (w6 += sigma1(w4) + w15 + sigma0(w7))); + Round(b, c, d, e, f, g, h, a, 0x76f988da + (w7 += sigma1(w5) + w0 + sigma0(w8))); + Round(a, b, c, d, e, f, g, h, 0x983e5152 + (w8 += sigma1(w6) + w1 + sigma0(w9))); + Round(h, a, b, c, d, e, f, g, 0xa831c66d + (w9 += sigma1(w7) + w2 + sigma0(w10))); + Round(g, h, a, b, c, d, e, f, 0xb00327c8 + (w10 += sigma1(w8) + w3 + sigma0(w11))); + Round(f, g, h, a, b, c, d, e, 0xbf597fc7 + (w11 += sigma1(w9) + w4 + sigma0(w12))); + Round(e, f, g, h, a, b, c, d, 0xc6e00bf3 + (w12 += sigma1(w10) + w5 + sigma0(w13))); + Round(d, e, f, g, h, a, b, c, 0xd5a79147 + (w13 += sigma1(w11) + w6 + sigma0(w14))); + Round(c, d, e, f, g, h, a, b, 0x06ca6351 + (w14 += sigma1(w12) + w7 + sigma0(w15))); + Round(b, c, d, e, f, g, h, a, 0x14292967 + (w15 += sigma1(w13) + w8 + sigma0(w0))); + + Round(a, b, c, d, e, f, g, h, 0x27b70a85 + (w0 += sigma1(w14) + w9 + sigma0(w1))); + Round(h, a, b, c, d, e, f, g, 0x2e1b2138 + (w1 += sigma1(w15) + w10 + sigma0(w2))); + Round(g, h, a, b, c, d, e, f, 0x4d2c6dfc + (w2 += sigma1(w0) + w11 + sigma0(w3))); + Round(f, g, h, a, b, c, d, e, 0x53380d13 + (w3 += sigma1(w1) + w12 + sigma0(w4))); + Round(e, f, g, h, a, b, c, d, 0x650a7354 + (w4 += sigma1(w2) + w13 + sigma0(w5))); + Round(d, e, f, g, h, a, b, c, 0x766a0abb + (w5 += sigma1(w3) + w14 + sigma0(w6))); + Round(c, d, e, f, g, h, a, b, 0x81c2c92e + (w6 += sigma1(w4) + w15 + sigma0(w7))); + Round(b, c, d, e, f, g, h, a, 0x92722c85 + (w7 += sigma1(w5) + w0 + sigma0(w8))); + Round(a, b, c, d, e, f, g, h, 0xa2bfe8a1 + (w8 += sigma1(w6) + w1 + sigma0(w9))); + Round(h, a, b, c, d, e, f, g, 0xa81a664b + (w9 += sigma1(w7) + w2 + sigma0(w10))); + Round(g, h, a, b, c, d, e, f, 0xc24b8b70 + (w10 += sigma1(w8) + w3 + sigma0(w11))); + Round(f, g, h, a, b, c, d, e, 0xc76c51a3 + (w11 += sigma1(w9) + w4 + sigma0(w12))); + Round(e, f, g, h, a, b, c, d, 0xd192e819 + (w12 += sigma1(w10) + w5 + sigma0(w13))); + Round(d, e, f, g, h, a, b, c, 0xd6990624 + (w13 += sigma1(w11) + w6 + sigma0(w14))); + Round(c, d, e, f, g, h, a, b, 0xf40e3585 + (w14 += sigma1(w12) + w7 + sigma0(w15))); + Round(b, c, d, e, f, g, h, a, 0x106aa070 + (w15 += sigma1(w13) + w8 + sigma0(w0))); + + Round(a, b, c, d, e, f, g, h, 0x19a4c116 + (w0 += sigma1(w14) + w9 + sigma0(w1))); + Round(h, a, b, c, d, e, f, g, 0x1e376c08 + (w1 += sigma1(w15) + w10 + sigma0(w2))); + Round(g, h, a, b, c, d, e, f, 0x2748774c + (w2 += sigma1(w0) + w11 + sigma0(w3))); + Round(f, g, h, a, b, c, d, e, 0x34b0bcb5 + (w3 += sigma1(w1) + w12 + sigma0(w4))); + Round(e, f, g, h, a, b, c, d, 0x391c0cb3 + (w4 += sigma1(w2) + w13 + sigma0(w5))); + Round(d, e, f, g, h, a, b, c, 0x4ed8aa4a + (w5 += sigma1(w3) + w14 + sigma0(w6))); + Round(c, d, e, f, g, h, a, b, 0x5b9cca4f + (w6 += sigma1(w4) + w15 + sigma0(w7))); + Round(b, c, d, e, f, g, h, a, 0x682e6ff3 + (w7 += sigma1(w5) + w0 + sigma0(w8))); + Round(a, b, c, d, e, f, g, h, 0x748f82ee + (w8 += sigma1(w6) + w1 + sigma0(w9))); + Round(h, a, b, c, d, e, f, g, 0x78a5636f + (w9 += sigma1(w7) + w2 + sigma0(w10))); + Round(g, h, a, b, c, d, e, f, 0x84c87814 + (w10 += sigma1(w8) + w3 + sigma0(w11))); + Round(f, g, h, a, b, c, d, e, 0x8cc70208 + (w11 += sigma1(w9) + w4 + sigma0(w12))); + Round(e, f, g, h, a, b, c, d, 0x90befffa + (w12 += sigma1(w10) + w5 + sigma0(w13))); + Round(d, e, f, g, h, a, b, c, 0xa4506ceb + (w13 += sigma1(w11) + w6 + sigma0(w14))); + Round(c, d, e, f, g, h, a, b, 0xbef9a3f7 + (w14 + sigma1(w12) + w7 + sigma0(w15))); + Round(b, c, d, e, f, g, h, a, 0xc67178f2 + (w15 + sigma1(w13) + w8 + sigma0(w0))); + + s[0] += a; + s[1] += b; + s[2] += c; + s[3] += d; + s[4] += e; + s[5] += f; + s[6] += g; + s[7] += h; + chunk += 64; + } +} + +} // namespace sha256 + +typedef void (*TransformType)(uint32_t*, const unsigned char*, size_t); +TransformType Transform = sha256::Transform; + +[[maybe_unused]] bool SelfTest() { + // Input state (equal to the initial SHA256 state) + static const uint32_t init[8] = { + 0x6a09e667ul, 0xbb67ae85ul, 0x3c6ef372ul, 0xa54ff53aul, 0x510e527ful, 0x9b05688cul, 0x1f83d9abul, 0x5be0cd19ul + }; + // Some random input data to test with + static const unsigned char data[641] = "-" // Intentionally not aligned + "Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do " + "eiusmod tempor incididunt ut labore et dolore magna aliqua. Et m" + "olestie ac feugiat sed lectus vestibulum mattis ullamcorper. Mor" + "bi blandit cursus risus at ultrices mi tempus imperdiet nulla. N" + "unc congue nisi vita suscipit tellus mauris. Imperdiet proin fer" + "mentum leo vel orci. Massa tempor nec feugiat nisl pretium fusce" + " id velit. Telus in metus vulputate eu scelerisque felis. Mi tem" + "pus imperdiet nulla malesuada pellentesque. Tristique magna sit."; + // Expected output state for hashing the i*64 first input bytes above (excluding SHA256 padding). + static const uint32_t result[9][8] = { + {0x6a09e667ul, 0xbb67ae85ul, 0x3c6ef372ul, 0xa54ff53aul, 0x510e527ful, 0x9b05688cul, 0x1f83d9abul, 0x5be0cd19ul}, + {0x91f8ec6bul, 0x4da10fe3ul, 0x1c9c292cul, 0x45e18185ul, 0x435cc111ul, 0x3ca26f09ul, 0xeb954caeul, 0x402a7069ul}, + {0xcabea5acul, 0x374fb97cul, 0x182ad996ul, 0x7bd69cbful, 0x450ff900ul, 0xc1d2be8aul, 0x6a41d505ul, 0xe6212dc3ul}, + {0xbcff09d6ul, 0x3e76f36eul, 0x3ecb2501ul, 0x78866e97ul, 0xe1c1e2fdul, 0x32f4eafful, 0x8aa6c4e5ul, 0xdfc024bcul}, + {0xa08c5d94ul, 0x0a862f93ul, 0x6b7f2f40ul, 0x8f9fae76ul, 0x6d40439ful, 0x79dcee0cul, 0x3e39ff3aul, 0xdc3bdbb1ul}, + {0x216a0895ul, 0x9f1a3662ul, 0xe99946f9ul, 0x87ba4364ul, 0x0fb5db2cul, 0x12bed3d3ul, 0x6689c0c7ul, 0x292f1b04ul}, + {0xca3067f8ul, 0xbc8c2656ul, 0x37cb7e0dul, 0x9b6b8b0ful, 0x46dc380bul, 0xf1287f57ul, 0xc42e4b23ul, 0x3fefe94dul}, + {0x3e4c4039ul, 0xbb6fca8cul, 0x6f27d2f7ul, 0x301e44a4ul, 0x8352ba14ul, 0x5769ce37ul, 0x48a1155ful, 0xc0e1c4c6ul}, + {0xfe2fa9ddul, 0x69d0862bul, 0x1ae0db23ul, 0x471f9244ul, 0xf55c0145ul, 0xc30f9c3bul, 0x40a84ea0ul, 0x5b8a266cul}, + }; + // Test Transform() for 0 through 8 transformations. + for (size_t i = 0; i <= 8; ++i) { + uint32_t state[8]; + std::copy(init, init + 8, state); + Transform(state, data + 1, i); + if (!std::equal(state, state + 8, result[i])) return false; + } + + return true; +} +} // namespace + +std::string SHA256AutoDetect(sha256_implementation::UseImplementation use_implementation) +{ + std::string ret = "standard"; + Transform = sha256::Transform; + +#if defined(__x86_64__) || defined(__amd64__) || defined(__i386__) + if (use_implementation & sha256_implementation::USE_SHANI) { + // Intel SHA extensions: CPUID.(EAX=7,ECX=0):EBX[bit 29]. Portable across + // GCC and Clang (unlike __builtin_cpu_supports("sha")). + unsigned int eax = 0, ebx = 0, ecx = 0, edx = 0; + if (__get_cpuid_count(7, 0, &eax, &ebx, &ecx, &edx) && (ebx & (1u << 29))) { + Transform = sha256_x86_shani::Transform; + ret = "x86_shani"; + } + } +#endif + +#if defined(__aarch64__) + if ((use_implementation & sha256_implementation::USE_SHANI) && (getauxval(AT_HWCAP) & HWCAP_SHA2)) { + Transform = sha256_arm_shani::Transform; + ret = "arm_shani"; + } +#endif + + assert(SelfTest()); + return ret; +} + +////// SHA-256 + +CSHA256::CSHA256() +{ + sha256::Initialize(s); +} + +CSHA256& CSHA256::Write(const unsigned char* data, size_t len) +{ + const unsigned char* end = data + len; + size_t bufsize = bytes % 64; + if (bufsize && bufsize + len >= 64) { + // Fill the buffer, and process it. + memcpy(buf + bufsize, data, 64 - bufsize); + bytes += 64 - bufsize; + data += 64 - bufsize; + Transform(s, buf, 1); + bufsize = 0; + } + if (end - data >= 64) { + size_t blocks = (end - data) / 64; + Transform(s, data, blocks); + data += 64 * blocks; + bytes += 64 * blocks; + } + if (end > data) { + // Fill the buffer with what remains. + memcpy(buf + bufsize, data, end - data); + bytes += end - data; + } + return *this; +} + +void CSHA256::Finalize(unsigned char hash[OUTPUT_SIZE]) +{ + static const unsigned char pad[64] = {0x80}; + unsigned char sizedesc[8]; + WriteBE64(sizedesc, bytes << 3); + Write(pad, 1 + ((119 - (bytes % 64)) % 64)); + Write(sizedesc, 8); + WriteBE32(hash, s[0]); + WriteBE32(hash + 4, s[1]); + WriteBE32(hash + 8, s[2]); + WriteBE32(hash + 12, s[3]); + WriteBE32(hash + 16, s[4]); + WriteBE32(hash + 20, s[5]); + WriteBE32(hash + 24, s[6]); + WriteBE32(hash + 28, s[7]); +} + +CSHA256& CSHA256::Reset() +{ + bytes = 0; + sha256::Initialize(s); + return *this; +} diff --git a/cpp/tensorrt_llm/common/sha256/sha256.h b/cpp/tensorrt_llm/common/sha256/sha256.h new file mode 100644 index 000000000000..ff89b71459fc --- /dev/null +++ b/cpp/tensorrt_llm/common/sha256/sha256.h @@ -0,0 +1,67 @@ +// SPDX-FileCopyrightText: Copyright (c) 2014-present The Bitcoin Core developers +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: MIT +// +// Permission is hereby granted, free of charge, to any person obtaining a +// copy of this software and associated documentation files (the "Software"), +// to deal in the Software without restriction, including without limitation +// the rights to use, copy, modify, merge, publish, distribute, sublicense, +// and/or sell copies of the Software, and to permit persons to whom the +// Software is furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in +// all copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +// DEALINGS IN THE SOFTWARE. +// +// Modifications for TensorRT-LLM: +// removed the SHA256D64 double-hash declaration (unused). + +#ifndef BITCOIN_CRYPTO_SHA256_H +#define BITCOIN_CRYPTO_SHA256_H + +#include +#include +#include + +/** A hasher class for SHA-256. */ +class CSHA256 +{ +private: + uint32_t s[8]; + unsigned char buf[64]; + uint64_t bytes{0}; + +public: + static const size_t OUTPUT_SIZE = 32; + + CSHA256(); + CSHA256& Write(const unsigned char* data, size_t len); + void Finalize(unsigned char hash[OUTPUT_SIZE]); + CSHA256& Reset(); +}; + +namespace sha256_implementation { +enum UseImplementation : uint8_t { + STANDARD = 0, + USE_SSE4 = 1 << 0, + USE_AVX2 = 1 << 1, + USE_SHANI = 1 << 2, + USE_SSE4_AND_AVX2 = USE_SSE4 | USE_AVX2, + USE_SSE4_AND_SHANI = USE_SSE4 | USE_SHANI, + USE_ALL = USE_SSE4 | USE_AVX2 | USE_SHANI, +}; +} + +/** Autodetect the best available SHA256 implementation. + * Returns the name of the implementation. + */ +std::string SHA256AutoDetect(sha256_implementation::UseImplementation use_implementation = sha256_implementation::USE_ALL); + +#endif // BITCOIN_CRYPTO_SHA256_H diff --git a/cpp/tensorrt_llm/common/sha256/sha256_arm_shani.cpp b/cpp/tensorrt_llm/common/sha256/sha256_arm_shani.cpp new file mode 100644 index 000000000000..228ea3908e73 --- /dev/null +++ b/cpp/tensorrt_llm/common/sha256/sha256_arm_shani.cpp @@ -0,0 +1,916 @@ +// SPDX-FileCopyrightText: Copyright (c) 2022-present The Bitcoin Core developers +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: MIT +// +// Permission is hereby granted, free of charge, to any person obtaining a +// copy of this software and associated documentation files (the "Software"), +// to deal in the Software without restriction, including without limitation +// the rights to use, copy, modify, merge, publish, distribute, sublicense, +// and/or sell copies of the Software, and to permit persons to whom the +// Software is furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in +// all copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +// DEALINGS IN THE SOFTWARE. +// +// Based on https://github.com/noloader/SHA-Intrinsics/blob/master/sha256-arm.c, +// Written and placed in public domain by Jeffrey Walton. +// Based on code from ARM, and by Johannes Schneiders, Skip Hovsmith and +// Barry O'Rourke for the mbedTLS project. +// Variant specialized for 64-byte inputs added by Pieter Wuille. + +#if defined(__aarch64__) + +#include +#include +#include +#include + +namespace { +alignas(uint32x4_t) static constexpr std::array K = +{ + 0x428A2F98, 0x71374491, 0xB5C0FBCF, 0xE9B5DBA5, + 0x3956C25B, 0x59F111F1, 0x923F82A4, 0xAB1C5ED5, + 0xD807AA98, 0x12835B01, 0x243185BE, 0x550C7DC3, + 0x72BE5D74, 0x80DEB1FE, 0x9BDC06A7, 0xC19BF174, + 0xE49B69C1, 0xEFBE4786, 0x0FC19DC6, 0x240CA1CC, + 0x2DE92C6F, 0x4A7484AA, 0x5CB0A9DC, 0x76F988DA, + 0x983E5152, 0xA831C66D, 0xB00327C8, 0xBF597FC7, + 0xC6E00BF3, 0xD5A79147, 0x06CA6351, 0x14292967, + 0x27B70A85, 0x2E1B2138, 0x4D2C6DFC, 0x53380D13, + 0x650A7354, 0x766A0ABB, 0x81C2C92E, 0x92722C85, + 0xA2BFE8A1, 0xA81A664B, 0xC24B8B70, 0xC76C51A3, + 0xD192E819, 0xD6990624, 0xF40E3585, 0x106AA070, + 0x19A4C116, 0x1E376C08, 0x2748774C, 0x34B0BCB5, + 0x391C0CB3, 0x4ED8AA4A, 0x5B9CCA4F, 0x682E6FF3, + 0x748F82EE, 0x78A5636F, 0x84C87814, 0x8CC70208, + 0x90BEFFFA, 0xA4506CEB, 0xBEF9A3F7, 0xC67178F2, +}; +} + +namespace sha256_arm_shani { +void Transform(uint32_t* s, const unsigned char* chunk, size_t blocks) +{ + uint32x4_t STATE0, STATE1, ABCD_SAVE, EFGH_SAVE; + uint32x4_t MSG0, MSG1, MSG2, MSG3; + uint32x4_t TMP0, TMP2; + + // Load state + STATE0 = vld1q_u32(&s[0]); + STATE1 = vld1q_u32(&s[4]); + + while (blocks--) + { + // Save state + ABCD_SAVE = STATE0; + EFGH_SAVE = STATE1; + + // Load and convert input chunk to Big Endian + MSG0 = vreinterpretq_u32_u8(vrev32q_u8(vld1q_u8(chunk + 0))); + MSG1 = vreinterpretq_u32_u8(vrev32q_u8(vld1q_u8(chunk + 16))); + MSG2 = vreinterpretq_u32_u8(vrev32q_u8(vld1q_u8(chunk + 32))); + MSG3 = vreinterpretq_u32_u8(vrev32q_u8(vld1q_u8(chunk + 48))); + chunk += 64; + + // Original implementation preloaded message and constant addition which was 1-3% slower. + // Now included as first step in quad round code saving one Q Neon register + // "TMP0 = vaddq_u32(MSG0, vld1q_u32(&K[0]));" + + // Rounds 1-4 + TMP0 = vaddq_u32(MSG0, vld1q_u32(&K[0])); + TMP2 = STATE0; + MSG0 = vsha256su0q_u32(MSG0, MSG1); + STATE0 = vsha256hq_u32(STATE0, STATE1, TMP0); + STATE1 = vsha256h2q_u32(STATE1, TMP2, TMP0); + MSG0 = vsha256su1q_u32(MSG0, MSG2, MSG3); + + // Rounds 5-8 + TMP0 = vaddq_u32(MSG1, vld1q_u32(&K[4])); + TMP2 = STATE0; + MSG1 = vsha256su0q_u32(MSG1, MSG2); + STATE0 = vsha256hq_u32(STATE0, STATE1, TMP0); + STATE1 = vsha256h2q_u32(STATE1, TMP2, TMP0); + MSG1 = vsha256su1q_u32(MSG1, MSG3, MSG0); + + // Rounds 9-12 + TMP0 = vaddq_u32(MSG2, vld1q_u32(&K[8])); + TMP2 = STATE0; + MSG2 = vsha256su0q_u32(MSG2, MSG3); + STATE0 = vsha256hq_u32(STATE0, STATE1, TMP0); + STATE1 = vsha256h2q_u32(STATE1, TMP2, TMP0); + MSG2 = vsha256su1q_u32(MSG2, MSG0, MSG1); + + // Rounds 13-16 + TMP0 = vaddq_u32(MSG3, vld1q_u32(&K[12])); + TMP2 = STATE0; + MSG3 = vsha256su0q_u32(MSG3, MSG0); + STATE0 = vsha256hq_u32(STATE0, STATE1, TMP0); + STATE1 = vsha256h2q_u32(STATE1, TMP2, TMP0); + MSG3 = vsha256su1q_u32(MSG3, MSG1, MSG2); + + // Rounds 17-20 + TMP0 = vaddq_u32(MSG0, vld1q_u32(&K[16])); + TMP2 = STATE0; + MSG0 = vsha256su0q_u32(MSG0, MSG1); + STATE0 = vsha256hq_u32(STATE0, STATE1, TMP0); + STATE1 = vsha256h2q_u32(STATE1, TMP2, TMP0); + MSG0 = vsha256su1q_u32(MSG0, MSG2, MSG3); + + // Rounds 21-24 + TMP0 = vaddq_u32(MSG1, vld1q_u32(&K[20])); + TMP2 = STATE0; + MSG1 = vsha256su0q_u32(MSG1, MSG2); + STATE0 = vsha256hq_u32(STATE0, STATE1, TMP0); + STATE1 = vsha256h2q_u32(STATE1, TMP2, TMP0); + MSG1 = vsha256su1q_u32(MSG1, MSG3, MSG0); + + // Rounds 25-28 + TMP0 = vaddq_u32(MSG2, vld1q_u32(&K[24])); + TMP2 = STATE0; + MSG2 = vsha256su0q_u32(MSG2, MSG3); + STATE0 = vsha256hq_u32(STATE0, STATE1, TMP0); + STATE1 = vsha256h2q_u32(STATE1, TMP2, TMP0); + MSG2 = vsha256su1q_u32(MSG2, MSG0, MSG1); + + // Rounds 29-32 + TMP0 = vaddq_u32(MSG3, vld1q_u32(&K[28])); + TMP2 = STATE0; + MSG3 = vsha256su0q_u32(MSG3, MSG0); + STATE0 = vsha256hq_u32(STATE0, STATE1, TMP0); + STATE1 = vsha256h2q_u32(STATE1, TMP2, TMP0); + MSG3 = vsha256su1q_u32(MSG3, MSG1, MSG2); + + // Rounds 33-36 + TMP0 = vaddq_u32(MSG0, vld1q_u32(&K[32])); + TMP2 = STATE0; + MSG0 = vsha256su0q_u32(MSG0, MSG1); + STATE0 = vsha256hq_u32(STATE0, STATE1, TMP0); + STATE1 = vsha256h2q_u32(STATE1, TMP2, TMP0); + MSG0 = vsha256su1q_u32(MSG0, MSG2, MSG3); + + // Rounds 37-40 + TMP0 = vaddq_u32(MSG1, vld1q_u32(&K[36])); + TMP2 = STATE0; + MSG1 = vsha256su0q_u32(MSG1, MSG2); + STATE0 = vsha256hq_u32(STATE0, STATE1, TMP0); + STATE1 = vsha256h2q_u32(STATE1, TMP2, TMP0); + MSG1 = vsha256su1q_u32(MSG1, MSG3, MSG0); + + // Rounds 41-44 + TMP0 = vaddq_u32(MSG2, vld1q_u32(&K[40])); + TMP2 = STATE0; + MSG2 = vsha256su0q_u32(MSG2, MSG3); + STATE0 = vsha256hq_u32(STATE0, STATE1, TMP0); + STATE1 = vsha256h2q_u32(STATE1, TMP2, TMP0); + MSG2 = vsha256su1q_u32(MSG2, MSG0, MSG1); + + // Rounds 45-48 + TMP0 = vaddq_u32(MSG3, vld1q_u32(&K[44])); + TMP2 = STATE0; + MSG3 = vsha256su0q_u32(MSG3, MSG0); + STATE0 = vsha256hq_u32(STATE0, STATE1, TMP0); + STATE1 = vsha256h2q_u32(STATE1, TMP2, TMP0); + MSG3 = vsha256su1q_u32(MSG3, MSG1, MSG2); + + // Rounds 49-52 + TMP0 = vaddq_u32(MSG0, vld1q_u32(&K[48])); + TMP2 = STATE0; + STATE0 = vsha256hq_u32(STATE0, STATE1, TMP0); + STATE1 = vsha256h2q_u32(STATE1, TMP2, TMP0); + + // Rounds 53-56 + TMP0 = vaddq_u32(MSG1, vld1q_u32(&K[52])); + TMP2 = STATE0; + STATE0 = vsha256hq_u32(STATE0, STATE1, TMP0); + STATE1 = vsha256h2q_u32(STATE1, TMP2, TMP0); + + // Rounds 57-60 + TMP0 = vaddq_u32(MSG2, vld1q_u32(&K[56])); + TMP2 = STATE0; + STATE0 = vsha256hq_u32(STATE0, STATE1, TMP0); + STATE1 = vsha256h2q_u32(STATE1, TMP2, TMP0); + + // Rounds 61-64 + TMP0 = vaddq_u32(MSG3, vld1q_u32(&K[60])); + TMP2 = STATE0; + STATE0 = vsha256hq_u32(STATE0, STATE1, TMP0); + STATE1 = vsha256h2q_u32(STATE1, TMP2, TMP0); + + // Update state + STATE0 = vaddq_u32(STATE0, ABCD_SAVE); + STATE1 = vaddq_u32(STATE1, EFGH_SAVE); + } + + // Save final state + vst1q_u32(&s[0], STATE0); + vst1q_u32(&s[4], STATE1); +} +} + +namespace sha256d64_arm_shani { +void Transform_2way(unsigned char* output, const unsigned char* input) +{ + /* Initial state. */ + alignas(uint32x4_t) static constexpr std::array INIT = { + 0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a, + 0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19 + }; + + /* Precomputed message schedule for the 2nd transform. */ + alignas(uint32x4_t) static constexpr std::array MIDS = { + 0xc28a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5, + 0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5, + 0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3, + 0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf374, + 0x649b69c1, 0xf0fe4786, 0x0fe1edc6, 0x240cf254, + 0x4fe9346f, 0x6cc984be, 0x61b9411e, 0x16f988fa, + 0xf2c65152, 0xa88e5a6d, 0xb019fc65, 0xb9d99ec7, + 0x9a1231c3, 0xe70eeaa0, 0xfdb1232b, 0xc7353eb0, + 0x3069bad5, 0xcb976d5f, 0x5a0f118f, 0xdc1eeefd, + 0x0a35b689, 0xde0b7a04, 0x58f4ca9d, 0xe15d5b16, + 0x007f3e86, 0x37088980, 0xa507ea32, 0x6fab9537, + 0x17406110, 0x0d8cd6f1, 0xcdaa3b6d, 0xc0bbbe37, + 0x83613bda, 0xdb48a363, 0x0b02e931, 0x6fd15ca7, + 0x521afaca, 0x31338431, 0x6ed41a95, 0x6d437890, + 0xc39c91f2, 0x9eccabbd, 0xb5c9a0e6, 0x532fb63c, + 0xd2c741c6, 0x07237ea3, 0xa4954b68, 0x4c191d76 + }; + + /* A few precomputed message schedule values for the 3rd transform. */ + alignas(uint32x4_t) static constexpr std::array FINS = { + 0x5807aa98, 0x12835b01, 0x243185be, 0x550c7dc3, + 0x80000000, 0x00000000, 0x00000000, 0x00000000, + 0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf274 + }; + + /* Padding processed in the 3rd transform (byteswapped). */ + alignas(uint32x4_t) static constexpr std::array FINAL = {0x80000000, 0, 0, 0, 0, 0, 0, 0x100}; + + uint32x4_t STATE0A, STATE0B, STATE1A, STATE1B, ABCD_SAVEA, ABCD_SAVEB, EFGH_SAVEA, EFGH_SAVEB; + uint32x4_t MSG0A, MSG0B, MSG1A, MSG1B, MSG2A, MSG2B, MSG3A, MSG3B; + uint32x4_t TMP0A, TMP0B, TMP2A, TMP2B, TMP; + + // Transform 1: Load state + STATE0A = vld1q_u32(&INIT[0]); + STATE0B = STATE0A; + STATE1A = vld1q_u32(&INIT[4]); + STATE1B = STATE1A; + + // Transform 1: Load and convert input chunk to Big Endian + MSG0A = vreinterpretq_u32_u8(vrev32q_u8(vld1q_u8(input + 0))); + MSG1A = vreinterpretq_u32_u8(vrev32q_u8(vld1q_u8(input + 16))); + MSG2A = vreinterpretq_u32_u8(vrev32q_u8(vld1q_u8(input + 32))); + MSG3A = vreinterpretq_u32_u8(vrev32q_u8(vld1q_u8(input + 48))); + MSG0B = vreinterpretq_u32_u8(vrev32q_u8(vld1q_u8(input + 64))); + MSG1B = vreinterpretq_u32_u8(vrev32q_u8(vld1q_u8(input + 80))); + MSG2B = vreinterpretq_u32_u8(vrev32q_u8(vld1q_u8(input + 96))); + MSG3B = vreinterpretq_u32_u8(vrev32q_u8(vld1q_u8(input + 112))); + + // Transform 1: Rounds 1-4 + TMP = vld1q_u32(&K[0]); + TMP0A = vaddq_u32(MSG0A, TMP); + TMP0B = vaddq_u32(MSG0B, TMP); + TMP2A = STATE0A; + TMP2B = STATE0B; + MSG0A = vsha256su0q_u32(MSG0A, MSG1A); + MSG0B = vsha256su0q_u32(MSG0B, MSG1B); + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP0A); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP0B); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP0A); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP0B); + MSG0A = vsha256su1q_u32(MSG0A, MSG2A, MSG3A); + MSG0B = vsha256su1q_u32(MSG0B, MSG2B, MSG3B); + + // Transform 1: Rounds 5-8 + TMP = vld1q_u32(&K[4]); + TMP0A = vaddq_u32(MSG1A, TMP); + TMP0B = vaddq_u32(MSG1B, TMP); + TMP2A = STATE0A; + TMP2B = STATE0B; + MSG1A = vsha256su0q_u32(MSG1A, MSG2A); + MSG1B = vsha256su0q_u32(MSG1B, MSG2B); + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP0A); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP0B); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP0A); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP0B); + MSG1A = vsha256su1q_u32(MSG1A, MSG3A, MSG0A); + MSG1B = vsha256su1q_u32(MSG1B, MSG3B, MSG0B); + + // Transform 1: Rounds 9-12 + TMP = vld1q_u32(&K[8]); + TMP0A = vaddq_u32(MSG2A, TMP); + TMP0B = vaddq_u32(MSG2B, TMP); + TMP2A = STATE0A; + TMP2B = STATE0B; + MSG2A = vsha256su0q_u32(MSG2A, MSG3A); + MSG2B = vsha256su0q_u32(MSG2B, MSG3B); + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP0A); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP0B); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP0A); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP0B); + MSG2A = vsha256su1q_u32(MSG2A, MSG0A, MSG1A); + MSG2B = vsha256su1q_u32(MSG2B, MSG0B, MSG1B); + + // Transform 1: Rounds 13-16 + TMP = vld1q_u32(&K[12]); + TMP0A = vaddq_u32(MSG3A, TMP); + TMP0B = vaddq_u32(MSG3B, TMP); + TMP2A = STATE0A; + TMP2B = STATE0B; + MSG3A = vsha256su0q_u32(MSG3A, MSG0A); + MSG3B = vsha256su0q_u32(MSG3B, MSG0B); + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP0A); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP0B); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP0A); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP0B); + MSG3A = vsha256su1q_u32(MSG3A, MSG1A, MSG2A); + MSG3B = vsha256su1q_u32(MSG3B, MSG1B, MSG2B); + + // Transform 1: Rounds 17-20 + TMP = vld1q_u32(&K[16]); + TMP0A = vaddq_u32(MSG0A, TMP); + TMP0B = vaddq_u32(MSG0B, TMP); + TMP2A = STATE0A; + TMP2B = STATE0B; + MSG0A = vsha256su0q_u32(MSG0A, MSG1A); + MSG0B = vsha256su0q_u32(MSG0B, MSG1B); + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP0A); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP0B); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP0A); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP0B); + MSG0A = vsha256su1q_u32(MSG0A, MSG2A, MSG3A); + MSG0B = vsha256su1q_u32(MSG0B, MSG2B, MSG3B); + + // Transform 1: Rounds 21-24 + TMP = vld1q_u32(&K[20]); + TMP0A = vaddq_u32(MSG1A, TMP); + TMP0B = vaddq_u32(MSG1B, TMP); + TMP2A = STATE0A; + TMP2B = STATE0B; + MSG1A = vsha256su0q_u32(MSG1A, MSG2A); + MSG1B = vsha256su0q_u32(MSG1B, MSG2B); + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP0A); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP0B); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP0A); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP0B); + MSG1A = vsha256su1q_u32(MSG1A, MSG3A, MSG0A); + MSG1B = vsha256su1q_u32(MSG1B, MSG3B, MSG0B); + + // Transform 1: Rounds 25-28 + TMP = vld1q_u32(&K[24]); + TMP0A = vaddq_u32(MSG2A, TMP); + TMP0B = vaddq_u32(MSG2B, TMP); + TMP2A = STATE0A; + TMP2B = STATE0B; + MSG2A = vsha256su0q_u32(MSG2A, MSG3A); + MSG2B = vsha256su0q_u32(MSG2B, MSG3B); + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP0A); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP0B); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP0A); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP0B); + MSG2A = vsha256su1q_u32(MSG2A, MSG0A, MSG1A); + MSG2B = vsha256su1q_u32(MSG2B, MSG0B, MSG1B); + + // Transform 1: Rounds 29-32 + TMP = vld1q_u32(&K[28]); + TMP0A = vaddq_u32(MSG3A, TMP); + TMP0B = vaddq_u32(MSG3B, TMP); + TMP2A = STATE0A; + TMP2B = STATE0B; + MSG3A = vsha256su0q_u32(MSG3A, MSG0A); + MSG3B = vsha256su0q_u32(MSG3B, MSG0B); + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP0A); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP0B); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP0A); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP0B); + MSG3A = vsha256su1q_u32(MSG3A, MSG1A, MSG2A); + MSG3B = vsha256su1q_u32(MSG3B, MSG1B, MSG2B); + + // Transform 1: Rounds 33-36 + TMP = vld1q_u32(&K[32]); + TMP0A = vaddq_u32(MSG0A, TMP); + TMP0B = vaddq_u32(MSG0B, TMP); + TMP2A = STATE0A; + TMP2B = STATE0B; + MSG0A = vsha256su0q_u32(MSG0A, MSG1A); + MSG0B = vsha256su0q_u32(MSG0B, MSG1B); + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP0A); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP0B); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP0A); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP0B); + MSG0A = vsha256su1q_u32(MSG0A, MSG2A, MSG3A); + MSG0B = vsha256su1q_u32(MSG0B, MSG2B, MSG3B); + + // Transform 1: Rounds 37-40 + TMP = vld1q_u32(&K[36]); + TMP0A = vaddq_u32(MSG1A, TMP); + TMP0B = vaddq_u32(MSG1B, TMP); + TMP2A = STATE0A; + TMP2B = STATE0B; + MSG1A = vsha256su0q_u32(MSG1A, MSG2A); + MSG1B = vsha256su0q_u32(MSG1B, MSG2B); + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP0A); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP0B); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP0A); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP0B); + MSG1A = vsha256su1q_u32(MSG1A, MSG3A, MSG0A); + MSG1B = vsha256su1q_u32(MSG1B, MSG3B, MSG0B); + + // Transform 1: Rounds 41-44 + TMP = vld1q_u32(&K[40]); + TMP0A = vaddq_u32(MSG2A, TMP); + TMP0B = vaddq_u32(MSG2B, TMP); + TMP2A = STATE0A; + TMP2B = STATE0B; + MSG2A = vsha256su0q_u32(MSG2A, MSG3A); + MSG2B = vsha256su0q_u32(MSG2B, MSG3B); + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP0A); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP0B); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP0A); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP0B); + MSG2A = vsha256su1q_u32(MSG2A, MSG0A, MSG1A); + MSG2B = vsha256su1q_u32(MSG2B, MSG0B, MSG1B); + + // Transform 1: Rounds 45-48 + TMP = vld1q_u32(&K[44]); + TMP0A = vaddq_u32(MSG3A, TMP); + TMP0B = vaddq_u32(MSG3B, TMP); + TMP2A = STATE0A; + TMP2B = STATE0B; + MSG3A = vsha256su0q_u32(MSG3A, MSG0A); + MSG3B = vsha256su0q_u32(MSG3B, MSG0B); + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP0A); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP0B); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP0A); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP0B); + MSG3A = vsha256su1q_u32(MSG3A, MSG1A, MSG2A); + MSG3B = vsha256su1q_u32(MSG3B, MSG1B, MSG2B); + + // Transform 1: Rounds 49-52 + TMP = vld1q_u32(&K[48]); + TMP0A = vaddq_u32(MSG0A, TMP); + TMP0B = vaddq_u32(MSG0B, TMP); + TMP2A = STATE0A; + TMP2B = STATE0B; + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP0A); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP0B); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP0A); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP0B); + + // Transform 1: Rounds 53-56 + TMP = vld1q_u32(&K[52]); + TMP0A = vaddq_u32(MSG1A, TMP); + TMP0B = vaddq_u32(MSG1B, TMP); + TMP2A = STATE0A; + TMP2B = STATE0B; + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP0A); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP0B); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP0A); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP0B); + + // Transform 1: Rounds 57-60 + TMP = vld1q_u32(&K[56]); + TMP0A = vaddq_u32(MSG2A, TMP); + TMP0B = vaddq_u32(MSG2B, TMP); + TMP2A = STATE0A; + TMP2B = STATE0B; + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP0A); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP0B); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP0A); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP0B); + + // Transform 1: Rounds 61-64 + TMP = vld1q_u32(&K[60]); + TMP0A = vaddq_u32(MSG3A, TMP); + TMP0B = vaddq_u32(MSG3B, TMP); + TMP2A = STATE0A; + TMP2B = STATE0B; + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP0A); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP0B); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP0A); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP0B); + + // Transform 1: Update state + TMP = vld1q_u32(&INIT[0]); + STATE0A = vaddq_u32(STATE0A, TMP); + STATE0B = vaddq_u32(STATE0B, TMP); + TMP = vld1q_u32(&INIT[4]); + STATE1A = vaddq_u32(STATE1A, TMP); + STATE1B = vaddq_u32(STATE1B, TMP); + + // Transform 2: Save state + ABCD_SAVEA = STATE0A; + ABCD_SAVEB = STATE0B; + EFGH_SAVEA = STATE1A; + EFGH_SAVEB = STATE1B; + + // Transform 2: Rounds 1-4 + TMP = vld1q_u32(&MIDS[0]); + TMP2A = STATE0A; + TMP2B = STATE0B; + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP); + + // Transform 2: Rounds 5-8 + TMP = vld1q_u32(&MIDS[4]); + TMP2A = STATE0A; + TMP2B = STATE0B; + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP); + + // Transform 2: Rounds 9-12 + TMP = vld1q_u32(&MIDS[8]); + TMP2A = STATE0A; + TMP2B = STATE0B; + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP); + + // Transform 2: Rounds 13-16 + TMP = vld1q_u32(&MIDS[12]); + TMP2A = STATE0A; + TMP2B = STATE0B; + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP); + + // Transform 2: Rounds 17-20 + TMP = vld1q_u32(&MIDS[16]); + TMP2A = STATE0A; + TMP2B = STATE0B; + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP); + + // Transform 2: Rounds 21-24 + TMP = vld1q_u32(&MIDS[20]); + TMP2A = STATE0A; + TMP2B = STATE0B; + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP); + + // Transform 2: Rounds 25-28 + TMP = vld1q_u32(&MIDS[24]); + TMP2A = STATE0A; + TMP2B = STATE0B; + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP); + + // Transform 2: Rounds 29-32 + TMP = vld1q_u32(&MIDS[28]); + TMP2A = STATE0A; + TMP2B = STATE0B; + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP); + + // Transform 2: Rounds 33-36 + TMP = vld1q_u32(&MIDS[32]); + TMP2A = STATE0A; + TMP2B = STATE0B; + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP); + + // Transform 2: Rounds 37-40 + TMP = vld1q_u32(&MIDS[36]); + TMP2A = STATE0A; + TMP2B = STATE0B; + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP); + + // Transform 2: Rounds 41-44 + TMP = vld1q_u32(&MIDS[40]); + TMP2A = STATE0A; + TMP2B = STATE0B; + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP); + + // Transform 2: Rounds 45-48 + TMP = vld1q_u32(&MIDS[44]); + TMP2A = STATE0A; + TMP2B = STATE0B; + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP); + + // Transform 2: Rounds 49-52 + TMP = vld1q_u32(&MIDS[48]); + TMP2A = STATE0A; + TMP2B = STATE0B; + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP); + + // Transform 2: Rounds 53-56 + TMP = vld1q_u32(&MIDS[52]); + TMP2A = STATE0A; + TMP2B = STATE0B; + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP); + + // Transform 2: Rounds 57-60 + TMP = vld1q_u32(&MIDS[56]); + TMP2A = STATE0A; + TMP2B = STATE0B; + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP); + + // Transform 2: Rounds 61-64 + TMP = vld1q_u32(&MIDS[60]); + TMP2A = STATE0A; + TMP2B = STATE0B; + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP); + + // Transform 2: Update state + STATE0A = vaddq_u32(STATE0A, ABCD_SAVEA); + STATE0B = vaddq_u32(STATE0B, ABCD_SAVEB); + STATE1A = vaddq_u32(STATE1A, EFGH_SAVEA); + STATE1B = vaddq_u32(STATE1B, EFGH_SAVEB); + + // Transform 3: Pad previous output + MSG0A = STATE0A; + MSG0B = STATE0B; + MSG1A = STATE1A; + MSG1B = STATE1B; + MSG2A = vld1q_u32(&FINAL[0]); + MSG2B = MSG2A; + MSG3A = vld1q_u32(&FINAL[4]); + MSG3B = MSG3A; + + // Transform 3: Load state + STATE0A = vld1q_u32(&INIT[0]); + STATE0B = STATE0A; + STATE1A = vld1q_u32(&INIT[4]); + STATE1B = STATE1A; + + // Transform 3: Rounds 1-4 + TMP = vld1q_u32(&K[0]); + TMP0A = vaddq_u32(MSG0A, TMP); + TMP0B = vaddq_u32(MSG0B, TMP); + TMP2A = STATE0A; + TMP2B = STATE0B; + MSG0A = vsha256su0q_u32(MSG0A, MSG1A); + MSG0B = vsha256su0q_u32(MSG0B, MSG1B); + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP0A); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP0B); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP0A); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP0B); + MSG0A = vsha256su1q_u32(MSG0A, MSG2A, MSG3A); + MSG0B = vsha256su1q_u32(MSG0B, MSG2B, MSG3B); + + // Transform 3: Rounds 5-8 + TMP = vld1q_u32(&K[4]); + TMP0A = vaddq_u32(MSG1A, TMP); + TMP0B = vaddq_u32(MSG1B, TMP); + TMP2A = STATE0A; + TMP2B = STATE0B; + MSG1A = vsha256su0q_u32(MSG1A, MSG2A); + MSG1B = vsha256su0q_u32(MSG1B, MSG2B); + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP0A); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP0B); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP0A); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP0B); + MSG1A = vsha256su1q_u32(MSG1A, MSG3A, MSG0A); + MSG1B = vsha256su1q_u32(MSG1B, MSG3B, MSG0B); + + // Transform 3: Rounds 9-12 + TMP = vld1q_u32(&FINS[0]); + TMP2A = STATE0A; + TMP2B = STATE0B; + MSG2A = vld1q_u32(&FINS[4]); + MSG2B = MSG2A; + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP); + MSG2A = vsha256su1q_u32(MSG2A, MSG0A, MSG1A); + MSG2B = vsha256su1q_u32(MSG2B, MSG0B, MSG1B); + + // Transform 3: Rounds 13-16 + TMP = vld1q_u32(&FINS[8]); + TMP2A = STATE0A; + TMP2B = STATE0B; + MSG3A = vsha256su0q_u32(MSG3A, MSG0A); + MSG3B = vsha256su0q_u32(MSG3B, MSG0B); + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP); + MSG3A = vsha256su1q_u32(MSG3A, MSG1A, MSG2A); + MSG3B = vsha256su1q_u32(MSG3B, MSG1B, MSG2B); + + // Transform 3: Rounds 17-20 + TMP = vld1q_u32(&K[16]); + TMP0A = vaddq_u32(MSG0A, TMP); + TMP0B = vaddq_u32(MSG0B, TMP); + TMP2A = STATE0A; + TMP2B = STATE0B; + MSG0A = vsha256su0q_u32(MSG0A, MSG1A); + MSG0B = vsha256su0q_u32(MSG0B, MSG1B); + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP0A); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP0B); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP0A); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP0B); + MSG0A = vsha256su1q_u32(MSG0A, MSG2A, MSG3A); + MSG0B = vsha256su1q_u32(MSG0B, MSG2B, MSG3B); + + // Transform 3: Rounds 21-24 + TMP = vld1q_u32(&K[20]); + TMP0A = vaddq_u32(MSG1A, TMP); + TMP0B = vaddq_u32(MSG1B, TMP); + TMP2A = STATE0A; + TMP2B = STATE0B; + MSG1A = vsha256su0q_u32(MSG1A, MSG2A); + MSG1B = vsha256su0q_u32(MSG1B, MSG2B); + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP0A); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP0B); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP0A); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP0B); + MSG1A = vsha256su1q_u32(MSG1A, MSG3A, MSG0A); + MSG1B = vsha256su1q_u32(MSG1B, MSG3B, MSG0B); + + // Transform 3: Rounds 25-28 + TMP = vld1q_u32(&K[24]); + TMP0A = vaddq_u32(MSG2A, TMP); + TMP0B = vaddq_u32(MSG2B, TMP); + TMP2A = STATE0A; + TMP2B = STATE0B; + MSG2A = vsha256su0q_u32(MSG2A, MSG3A); + MSG2B = vsha256su0q_u32(MSG2B, MSG3B); + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP0A); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP0B); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP0A); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP0B); + MSG2A = vsha256su1q_u32(MSG2A, MSG0A, MSG1A); + MSG2B = vsha256su1q_u32(MSG2B, MSG0B, MSG1B); + + // Transform 3: Rounds 29-32 + TMP = vld1q_u32(&K[28]); + TMP0A = vaddq_u32(MSG3A, TMP); + TMP0B = vaddq_u32(MSG3B, TMP); + TMP2A = STATE0A; + TMP2B = STATE0B; + MSG3A = vsha256su0q_u32(MSG3A, MSG0A); + MSG3B = vsha256su0q_u32(MSG3B, MSG0B); + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP0A); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP0B); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP0A); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP0B); + MSG3A = vsha256su1q_u32(MSG3A, MSG1A, MSG2A); + MSG3B = vsha256su1q_u32(MSG3B, MSG1B, MSG2B); + + // Transform 3: Rounds 33-36 + TMP = vld1q_u32(&K[32]); + TMP0A = vaddq_u32(MSG0A, TMP); + TMP0B = vaddq_u32(MSG0B, TMP); + TMP2A = STATE0A; + TMP2B = STATE0B; + MSG0A = vsha256su0q_u32(MSG0A, MSG1A); + MSG0B = vsha256su0q_u32(MSG0B, MSG1B); + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP0A); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP0B); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP0A); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP0B); + MSG0A = vsha256su1q_u32(MSG0A, MSG2A, MSG3A); + MSG0B = vsha256su1q_u32(MSG0B, MSG2B, MSG3B); + + // Transform 3: Rounds 37-40 + TMP = vld1q_u32(&K[36]); + TMP0A = vaddq_u32(MSG1A, TMP); + TMP0B = vaddq_u32(MSG1B, TMP); + TMP2A = STATE0A; + TMP2B = STATE0B; + MSG1A = vsha256su0q_u32(MSG1A, MSG2A); + MSG1B = vsha256su0q_u32(MSG1B, MSG2B); + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP0A); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP0B); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP0A); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP0B); + MSG1A = vsha256su1q_u32(MSG1A, MSG3A, MSG0A); + MSG1B = vsha256su1q_u32(MSG1B, MSG3B, MSG0B); + + // Transform 3: Rounds 41-44 + TMP = vld1q_u32(&K[40]); + TMP0A = vaddq_u32(MSG2A, TMP); + TMP0B = vaddq_u32(MSG2B, TMP); + TMP2A = STATE0A; + TMP2B = STATE0B; + MSG2A = vsha256su0q_u32(MSG2A, MSG3A); + MSG2B = vsha256su0q_u32(MSG2B, MSG3B); + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP0A); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP0B); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP0A); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP0B); + MSG2A = vsha256su1q_u32(MSG2A, MSG0A, MSG1A); + MSG2B = vsha256su1q_u32(MSG2B, MSG0B, MSG1B); + + // Transform 3: Rounds 45-48 + TMP = vld1q_u32(&K[44]); + TMP0A = vaddq_u32(MSG3A, TMP); + TMP0B = vaddq_u32(MSG3B, TMP); + TMP2A = STATE0A; + TMP2B = STATE0B; + MSG3A = vsha256su0q_u32(MSG3A, MSG0A); + MSG3B = vsha256su0q_u32(MSG3B, MSG0B); + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP0A); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP0B); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP0A); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP0B); + MSG3A = vsha256su1q_u32(MSG3A, MSG1A, MSG2A); + MSG3B = vsha256su1q_u32(MSG3B, MSG1B, MSG2B); + + // Transform 3: Rounds 49-52 + TMP = vld1q_u32(&K[48]); + TMP0A = vaddq_u32(MSG0A, TMP); + TMP0B = vaddq_u32(MSG0B, TMP); + TMP2A = STATE0A; + TMP2B = STATE0B; + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP0A); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP0B); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP0A); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP0B); + + // Transform 3: Rounds 53-56 + TMP = vld1q_u32(&K[52]); + TMP0A = vaddq_u32(MSG1A, TMP); + TMP0B = vaddq_u32(MSG1B, TMP); + TMP2A = STATE0A; + TMP2B = STATE0B; + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP0A); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP0B); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP0A); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP0B); + + // Transform 3: Rounds 57-60 + TMP = vld1q_u32(&K[56]); + TMP0A = vaddq_u32(MSG2A, TMP); + TMP0B = vaddq_u32(MSG2B, TMP); + TMP2A = STATE0A; + TMP2B = STATE0B; + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP0A); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP0B); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP0A); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP0B); + + // Transform 3: Rounds 61-64 + TMP = vld1q_u32(&K[60]); + TMP0A = vaddq_u32(MSG3A, TMP); + TMP0B = vaddq_u32(MSG3B, TMP); + TMP2A = STATE0A; + TMP2B = STATE0B; + STATE0A = vsha256hq_u32(STATE0A, STATE1A, TMP0A); + STATE0B = vsha256hq_u32(STATE0B, STATE1B, TMP0B); + STATE1A = vsha256h2q_u32(STATE1A, TMP2A, TMP0A); + STATE1B = vsha256h2q_u32(STATE1B, TMP2B, TMP0B); + + // Transform 3: Update state + TMP = vld1q_u32(&INIT[0]); + STATE0A = vaddq_u32(STATE0A, TMP); + STATE0B = vaddq_u32(STATE0B, TMP); + TMP = vld1q_u32(&INIT[4]); + STATE1A = vaddq_u32(STATE1A, TMP); + STATE1B = vaddq_u32(STATE1B, TMP); + + // Store result + vst1q_u8(output, vrev32q_u8(vreinterpretq_u8_u32(STATE0A))); + vst1q_u8(output + 16, vrev32q_u8(vreinterpretq_u8_u32(STATE1A))); + vst1q_u8(output + 32, vrev32q_u8(vreinterpretq_u8_u32(STATE0B))); + vst1q_u8(output + 48, vrev32q_u8(vreinterpretq_u8_u32(STATE1B))); +} +} + +#endif diff --git a/cpp/tensorrt_llm/common/sha256/sha256_endian.h b/cpp/tensorrt_llm/common/sha256/sha256_endian.h new file mode 100644 index 000000000000..82af31f00060 --- /dev/null +++ b/cpp/tensorrt_llm/common/sha256/sha256_endian.h @@ -0,0 +1,61 @@ +// SPDX-FileCopyrightText: Copyright (c) 2014-present The Bitcoin Core developers +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: MIT +// +// Permission is hereby granted, free of charge, to any person obtaining a +// copy of this software and associated documentation files (the "Software"), +// to deal in the Software without restriction, including without limitation +// the rights to use, copy, modify, merge, publish, distribute, sublicense, +// and/or sell copies of the Software, and to permit persons to whom the +// Software is furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in +// all copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +// DEALINGS IN THE SOFTWARE. +// +// Minimal big-endian read/write helpers for the vendored Bitcoin Core SHA-256 +// (sha256.cpp). This NVIDIA-authored header replaces the upstream +// / / dependency chain +// (which required C++20) so the vendored sources build as plain C++17. The +// helper names and byte order match the upstream common.h functions they +// replace, so the SHA-256 output is unchanged. + +#ifndef TENSORRT_LLM_3RDPARTY_SHA256_ENDIAN_H +#define TENSORRT_LLM_3RDPARTY_SHA256_ENDIAN_H + +#include + +inline uint32_t ReadBE32(unsigned char const* ptr) +{ + return (static_cast(ptr[0]) << 24) | (static_cast(ptr[1]) << 16) + | (static_cast(ptr[2]) << 8) | static_cast(ptr[3]); +} + +inline void WriteBE32(unsigned char* ptr, uint32_t x) +{ + ptr[0] = static_cast(x >> 24); + ptr[1] = static_cast(x >> 16); + ptr[2] = static_cast(x >> 8); + ptr[3] = static_cast(x); +} + +inline void WriteBE64(unsigned char* ptr, uint64_t x) +{ + ptr[0] = static_cast(x >> 56); + ptr[1] = static_cast(x >> 48); + ptr[2] = static_cast(x >> 40); + ptr[3] = static_cast(x >> 32); + ptr[4] = static_cast(x >> 24); + ptr[5] = static_cast(x >> 16); + ptr[6] = static_cast(x >> 8); + ptr[7] = static_cast(x); +} + +#endif // TENSORRT_LLM_3RDPARTY_SHA256_ENDIAN_H diff --git a/cpp/tensorrt_llm/common/sha256/sha256_x86_shani.cpp b/cpp/tensorrt_llm/common/sha256/sha256_x86_shani.cpp new file mode 100644 index 000000000000..5f3477c692f4 --- /dev/null +++ b/cpp/tensorrt_llm/common/sha256/sha256_x86_shani.cpp @@ -0,0 +1,377 @@ +// SPDX-FileCopyrightText: Copyright (c) 2018-present The Bitcoin Core developers +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: MIT +// +// Permission is hereby granted, free of charge, to any person obtaining a +// copy of this software and associated documentation files (the "Software"), +// to deal in the Software without restriction, including without limitation +// the rights to use, copy, modify, merge, publish, distribute, sublicense, +// and/or sell copies of the Software, and to permit persons to whom the +// Software is furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in +// all copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +// DEALINGS IN THE SOFTWARE. +// +// Based on https://github.com/noloader/SHA-Intrinsics/blob/master/sha256-x86.c, +// Written and placed in public domain by Jeffrey Walton. +// Based on code from Intel, and by Sean Gulley for the miTLS project. + +#if defined(__x86_64__) || defined(__amd64__) || defined(__i386__) + +#include +#include +#include + +#include + +namespace { + +alignas(__m128i) const uint8_t MASK[16] = {0x03, 0x02, 0x01, 0x00, 0x07, 0x06, 0x05, 0x04, 0x0b, 0x0a, 0x09, 0x08, 0x0f, 0x0e, 0x0d, 0x0c}; +alignas(__m128i) const uint8_t INIT0[16] = {0x8c, 0x68, 0x05, 0x9b, 0x7f, 0x52, 0x0e, 0x51, 0x85, 0xae, 0x67, 0xbb, 0x67, 0xe6, 0x09, 0x6a}; +alignas(__m128i) const uint8_t INIT1[16] = {0x19, 0xcd, 0xe0, 0x5b, 0xab, 0xd9, 0x83, 0x1f, 0x3a, 0xf5, 0x4f, 0xa5, 0x72, 0xf3, 0x6e, 0x3c}; + +void ALWAYS_INLINE QuadRound(__m128i& state0, __m128i& state1, uint64_t k1, uint64_t k0) +{ + const __m128i msg = _mm_set_epi64x(k1, k0); + state1 = _mm_sha256rnds2_epu32(state1, state0, msg); + state0 = _mm_sha256rnds2_epu32(state0, state1, _mm_shuffle_epi32(msg, 0x0e)); +} + +void ALWAYS_INLINE QuadRound(__m128i& state0, __m128i& state1, __m128i m, uint64_t k1, uint64_t k0) +{ + const __m128i msg = _mm_add_epi32(m, _mm_set_epi64x(k1, k0)); + state1 = _mm_sha256rnds2_epu32(state1, state0, msg); + state0 = _mm_sha256rnds2_epu32(state0, state1, _mm_shuffle_epi32(msg, 0x0e)); +} + +void ALWAYS_INLINE ShiftMessageA(__m128i& m0, __m128i m1) +{ + m0 = _mm_sha256msg1_epu32(m0, m1); +} + +void ALWAYS_INLINE ShiftMessageC(__m128i& m0, __m128i m1, __m128i& m2) +{ + m2 = _mm_sha256msg2_epu32(_mm_add_epi32(m2, _mm_alignr_epi8(m1, m0, 4)), m1); +} + +void ALWAYS_INLINE ShiftMessageB(__m128i& m0, __m128i m1, __m128i& m2) +{ + ShiftMessageC(m0, m1, m2); + ShiftMessageA(m0, m1); +} + +void ALWAYS_INLINE Shuffle(__m128i& s0, __m128i& s1) +{ + const __m128i t1 = _mm_shuffle_epi32(s0, 0xB1); + const __m128i t2 = _mm_shuffle_epi32(s1, 0x1B); + s0 = _mm_alignr_epi8(t1, t2, 0x08); + s1 = _mm_blend_epi16(t2, t1, 0xF0); +} + +void ALWAYS_INLINE Unshuffle(__m128i& s0, __m128i& s1) +{ + const __m128i t1 = _mm_shuffle_epi32(s0, 0x1B); + const __m128i t2 = _mm_shuffle_epi32(s1, 0xB1); + s0 = _mm_blend_epi16(t1, t2, 0xF0); + s1 = _mm_alignr_epi8(t2, t1, 0x08); +} + +__m128i ALWAYS_INLINE Load(const unsigned char* in) +{ + return _mm_shuffle_epi8(_mm_loadu_si128((const __m128i*)in), _mm_load_si128((const __m128i*)MASK)); +} + +void ALWAYS_INLINE Save(unsigned char* out, __m128i s) +{ + _mm_storeu_si128((__m128i*)out, _mm_shuffle_epi8(s, _mm_load_si128((const __m128i*)MASK))); +} +} + +namespace sha256_x86_shani { +void Transform(uint32_t* s, const unsigned char* chunk, size_t blocks) +{ + __m128i m0, m1, m2, m3, s0, s1, so0, so1; + + /* Load state */ + s0 = _mm_loadu_si128((const __m128i*)s); + s1 = _mm_loadu_si128((const __m128i*)(s + 4)); + Shuffle(s0, s1); + + while (blocks--) { + /* Remember old state */ + so0 = s0; + so1 = s1; + + /* Load data and transform */ + m0 = Load(chunk); + QuadRound(s0, s1, m0, 0xe9b5dba5b5c0fbcfull, 0x71374491428a2f98ull); + m1 = Load(chunk + 16); + QuadRound(s0, s1, m1, 0xab1c5ed5923f82a4ull, 0x59f111f13956c25bull); + ShiftMessageA(m0, m1); + m2 = Load(chunk + 32); + QuadRound(s0, s1, m2, 0x550c7dc3243185beull, 0x12835b01d807aa98ull); + ShiftMessageA(m1, m2); + m3 = Load(chunk + 48); + QuadRound(s0, s1, m3, 0xc19bf1749bdc06a7ull, 0x80deb1fe72be5d74ull); + ShiftMessageB(m2, m3, m0); + QuadRound(s0, s1, m0, 0x240ca1cc0fc19dc6ull, 0xefbe4786E49b69c1ull); + ShiftMessageB(m3, m0, m1); + QuadRound(s0, s1, m1, 0x76f988da5cb0a9dcull, 0x4a7484aa2de92c6full); + ShiftMessageB(m0, m1, m2); + QuadRound(s0, s1, m2, 0xbf597fc7b00327c8ull, 0xa831c66d983e5152ull); + ShiftMessageB(m1, m2, m3); + QuadRound(s0, s1, m3, 0x1429296706ca6351ull, 0xd5a79147c6e00bf3ull); + ShiftMessageB(m2, m3, m0); + QuadRound(s0, s1, m0, 0x53380d134d2c6dfcull, 0x2e1b213827b70a85ull); + ShiftMessageB(m3, m0, m1); + QuadRound(s0, s1, m1, 0x92722c8581c2c92eull, 0x766a0abb650a7354ull); + ShiftMessageB(m0, m1, m2); + QuadRound(s0, s1, m2, 0xc76c51A3c24b8b70ull, 0xa81a664ba2bfe8a1ull); + ShiftMessageB(m1, m2, m3); + QuadRound(s0, s1, m3, 0x106aa070f40e3585ull, 0xd6990624d192e819ull); + ShiftMessageB(m2, m3, m0); + QuadRound(s0, s1, m0, 0x34b0bcb52748774cull, 0x1e376c0819a4c116ull); + ShiftMessageB(m3, m0, m1); + QuadRound(s0, s1, m1, 0x682e6ff35b9cca4full, 0x4ed8aa4a391c0cb3ull); + ShiftMessageC(m0, m1, m2); + QuadRound(s0, s1, m2, 0x8cc7020884c87814ull, 0x78a5636f748f82eeull); + ShiftMessageC(m1, m2, m3); + QuadRound(s0, s1, m3, 0xc67178f2bef9A3f7ull, 0xa4506ceb90befffaull); + + /* Combine with old state */ + s0 = _mm_add_epi32(s0, so0); + s1 = _mm_add_epi32(s1, so1); + + /* Advance */ + chunk += 64; + } + + Unshuffle(s0, s1); + _mm_storeu_si128((__m128i*)s, s0); + _mm_storeu_si128((__m128i*)(s + 4), s1); +} +} + +namespace sha256d64_x86_shani { + +void Transform_2way(unsigned char* out, const unsigned char* in) +{ + __m128i am0, am1, am2, am3, as0, as1, aso0, aso1; + __m128i bm0, bm1, bm2, bm3, bs0, bs1, bso0, bso1; + + /* Transform 1 */ + bs0 = as0 = _mm_load_si128((const __m128i*)INIT0); + bs1 = as1 = _mm_load_si128((const __m128i*)INIT1); + am0 = Load(in); + bm0 = Load(in + 64); + QuadRound(as0, as1, am0, 0xe9b5dba5b5c0fbcfull, 0x71374491428a2f98ull); + QuadRound(bs0, bs1, bm0, 0xe9b5dba5b5c0fbcfull, 0x71374491428a2f98ull); + am1 = Load(in + 16); + bm1 = Load(in + 80); + QuadRound(as0, as1, am1, 0xab1c5ed5923f82a4ull, 0x59f111f13956c25bull); + QuadRound(bs0, bs1, bm1, 0xab1c5ed5923f82a4ull, 0x59f111f13956c25bull); + ShiftMessageA(am0, am1); + ShiftMessageA(bm0, bm1); + am2 = Load(in + 32); + bm2 = Load(in + 96); + QuadRound(as0, as1, am2, 0x550c7dc3243185beull, 0x12835b01d807aa98ull); + QuadRound(bs0, bs1, bm2, 0x550c7dc3243185beull, 0x12835b01d807aa98ull); + ShiftMessageA(am1, am2); + ShiftMessageA(bm1, bm2); + am3 = Load(in + 48); + bm3 = Load(in + 112); + QuadRound(as0, as1, am3, 0xc19bf1749bdc06a7ull, 0x80deb1fe72be5d74ull); + QuadRound(bs0, bs1, bm3, 0xc19bf1749bdc06a7ull, 0x80deb1fe72be5d74ull); + ShiftMessageB(am2, am3, am0); + ShiftMessageB(bm2, bm3, bm0); + QuadRound(as0, as1, am0, 0x240ca1cc0fc19dc6ull, 0xefbe4786E49b69c1ull); + QuadRound(bs0, bs1, bm0, 0x240ca1cc0fc19dc6ull, 0xefbe4786E49b69c1ull); + ShiftMessageB(am3, am0, am1); + ShiftMessageB(bm3, bm0, bm1); + QuadRound(as0, as1, am1, 0x76f988da5cb0a9dcull, 0x4a7484aa2de92c6full); + QuadRound(bs0, bs1, bm1, 0x76f988da5cb0a9dcull, 0x4a7484aa2de92c6full); + ShiftMessageB(am0, am1, am2); + ShiftMessageB(bm0, bm1, bm2); + QuadRound(as0, as1, am2, 0xbf597fc7b00327c8ull, 0xa831c66d983e5152ull); + QuadRound(bs0, bs1, bm2, 0xbf597fc7b00327c8ull, 0xa831c66d983e5152ull); + ShiftMessageB(am1, am2, am3); + ShiftMessageB(bm1, bm2, bm3); + QuadRound(as0, as1, am3, 0x1429296706ca6351ull, 0xd5a79147c6e00bf3ull); + QuadRound(bs0, bs1, bm3, 0x1429296706ca6351ull, 0xd5a79147c6e00bf3ull); + ShiftMessageB(am2, am3, am0); + ShiftMessageB(bm2, bm3, bm0); + QuadRound(as0, as1, am0, 0x53380d134d2c6dfcull, 0x2e1b213827b70a85ull); + QuadRound(bs0, bs1, bm0, 0x53380d134d2c6dfcull, 0x2e1b213827b70a85ull); + ShiftMessageB(am3, am0, am1); + ShiftMessageB(bm3, bm0, bm1); + QuadRound(as0, as1, am1, 0x92722c8581c2c92eull, 0x766a0abb650a7354ull); + QuadRound(bs0, bs1, bm1, 0x92722c8581c2c92eull, 0x766a0abb650a7354ull); + ShiftMessageB(am0, am1, am2); + ShiftMessageB(bm0, bm1, bm2); + QuadRound(as0, as1, am2, 0xc76c51A3c24b8b70ull, 0xa81a664ba2bfe8a1ull); + QuadRound(bs0, bs1, bm2, 0xc76c51A3c24b8b70ull, 0xa81a664ba2bfe8a1ull); + ShiftMessageB(am1, am2, am3); + ShiftMessageB(bm1, bm2, bm3); + QuadRound(as0, as1, am3, 0x106aa070f40e3585ull, 0xd6990624d192e819ull); + QuadRound(bs0, bs1, bm3, 0x106aa070f40e3585ull, 0xd6990624d192e819ull); + ShiftMessageB(am2, am3, am0); + ShiftMessageB(bm2, bm3, bm0); + QuadRound(as0, as1, am0, 0x34b0bcb52748774cull, 0x1e376c0819a4c116ull); + QuadRound(bs0, bs1, bm0, 0x34b0bcb52748774cull, 0x1e376c0819a4c116ull); + ShiftMessageB(am3, am0, am1); + ShiftMessageB(bm3, bm0, bm1); + QuadRound(as0, as1, am1, 0x682e6ff35b9cca4full, 0x4ed8aa4a391c0cb3ull); + QuadRound(bs0, bs1, bm1, 0x682e6ff35b9cca4full, 0x4ed8aa4a391c0cb3ull); + ShiftMessageC(am0, am1, am2); + ShiftMessageC(bm0, bm1, bm2); + QuadRound(as0, as1, am2, 0x8cc7020884c87814ull, 0x78a5636f748f82eeull); + QuadRound(bs0, bs1, bm2, 0x8cc7020884c87814ull, 0x78a5636f748f82eeull); + ShiftMessageC(am1, am2, am3); + ShiftMessageC(bm1, bm2, bm3); + QuadRound(as0, as1, am3, 0xc67178f2bef9A3f7ull, 0xa4506ceb90befffaull); + QuadRound(bs0, bs1, bm3, 0xc67178f2bef9A3f7ull, 0xa4506ceb90befffaull); + as0 = _mm_add_epi32(as0, _mm_load_si128((const __m128i*)INIT0)); + bs0 = _mm_add_epi32(bs0, _mm_load_si128((const __m128i*)INIT0)); + as1 = _mm_add_epi32(as1, _mm_load_si128((const __m128i*)INIT1)); + bs1 = _mm_add_epi32(bs1, _mm_load_si128((const __m128i*)INIT1)); + + /* Transform 2 */ + aso0 = as0; + bso0 = bs0; + aso1 = as1; + bso1 = bs1; + QuadRound(as0, as1, 0xe9b5dba5b5c0fbcfull, 0x71374491c28a2f98ull); + QuadRound(bs0, bs1, 0xe9b5dba5b5c0fbcfull, 0x71374491c28a2f98ull); + QuadRound(as0, as1, 0xab1c5ed5923f82a4ull, 0x59f111f13956c25bull); + QuadRound(bs0, bs1, 0xab1c5ed5923f82a4ull, 0x59f111f13956c25bull); + QuadRound(as0, as1, 0x550c7dc3243185beull, 0x12835b01d807aa98ull); + QuadRound(bs0, bs1, 0x550c7dc3243185beull, 0x12835b01d807aa98ull); + QuadRound(as0, as1, 0xc19bf3749bdc06a7ull, 0x80deb1fe72be5d74ull); + QuadRound(bs0, bs1, 0xc19bf3749bdc06a7ull, 0x80deb1fe72be5d74ull); + QuadRound(as0, as1, 0x240cf2540fe1edc6ull, 0xf0fe4786649b69c1ull); + QuadRound(bs0, bs1, 0x240cf2540fe1edc6ull, 0xf0fe4786649b69c1ull); + QuadRound(as0, as1, 0x16f988fa61b9411eull, 0x6cc984be4fe9346full); + QuadRound(bs0, bs1, 0x16f988fa61b9411eull, 0x6cc984be4fe9346full); + QuadRound(as0, as1, 0xb9d99ec7b019fc65ull, 0xa88e5a6df2c65152ull); + QuadRound(bs0, bs1, 0xb9d99ec7b019fc65ull, 0xa88e5a6df2c65152ull); + QuadRound(as0, as1, 0xc7353eb0fdb1232bull, 0xe70eeaa09a1231c3ull); + QuadRound(bs0, bs1, 0xc7353eb0fdb1232bull, 0xe70eeaa09a1231c3ull); + QuadRound(as0, as1, 0xdc1eeefd5a0f118full, 0xcb976d5f3069bad5ull); + QuadRound(bs0, bs1, 0xdc1eeefd5a0f118full, 0xcb976d5f3069bad5ull); + QuadRound(as0, as1, 0xe15d5b1658f4ca9dull, 0xde0b7a040a35b689ull); + QuadRound(bs0, bs1, 0xe15d5b1658f4ca9dull, 0xde0b7a040a35b689ull); + QuadRound(as0, as1, 0x6fab9537a507ea32ull, 0x37088980007f3e86ull); + QuadRound(bs0, bs1, 0x6fab9537a507ea32ull, 0x37088980007f3e86ull); + QuadRound(as0, as1, 0xc0bbbe37cdaa3b6dull, 0x0d8cd6f117406110ull); + QuadRound(bs0, bs1, 0xc0bbbe37cdaa3b6dull, 0x0d8cd6f117406110ull); + QuadRound(as0, as1, 0x6fd15ca70b02e931ull, 0xdb48a36383613bdaull); + QuadRound(bs0, bs1, 0x6fd15ca70b02e931ull, 0xdb48a36383613bdaull); + QuadRound(as0, as1, 0x6d4378906ed41a95ull, 0x31338431521afacaull); + QuadRound(bs0, bs1, 0x6d4378906ed41a95ull, 0x31338431521afacaull); + QuadRound(as0, as1, 0x532fb63cb5c9a0e6ull, 0x9eccabbdc39c91f2ull); + QuadRound(bs0, bs1, 0x532fb63cb5c9a0e6ull, 0x9eccabbdc39c91f2ull); + QuadRound(as0, as1, 0x4c191d76a4954b68ull, 0x07237ea3d2c741c6ull); + QuadRound(bs0, bs1, 0x4c191d76a4954b68ull, 0x07237ea3d2c741c6ull); + as0 = _mm_add_epi32(as0, aso0); + bs0 = _mm_add_epi32(bs0, bso0); + as1 = _mm_add_epi32(as1, aso1); + bs1 = _mm_add_epi32(bs1, bso1); + + /* Extract hash */ + Unshuffle(as0, as1); + Unshuffle(bs0, bs1); + am0 = as0; + bm0 = bs0; + am1 = as1; + bm1 = bs1; + + /* Transform 3 */ + bs0 = as0 = _mm_load_si128((const __m128i*)INIT0); + bs1 = as1 = _mm_load_si128((const __m128i*)INIT1); + QuadRound(as0, as1, am0, 0xe9b5dba5B5c0fbcfull, 0x71374491428a2f98ull); + QuadRound(bs0, bs1, bm0, 0xe9b5dba5B5c0fbcfull, 0x71374491428a2f98ull); + QuadRound(as0, as1, am1, 0xab1c5ed5923f82a4ull, 0x59f111f13956c25bull); + QuadRound(bs0, bs1, bm1, 0xab1c5ed5923f82a4ull, 0x59f111f13956c25bull); + ShiftMessageA(am0, am1); + ShiftMessageA(bm0, bm1); + bm2 = am2 = _mm_set_epi64x(0x0ull, 0x80000000ull); + QuadRound(as0, as1, 0x550c7dc3243185beull, 0x12835b015807aa98ull); + QuadRound(bs0, bs1, 0x550c7dc3243185beull, 0x12835b015807aa98ull); + ShiftMessageA(am1, am2); + ShiftMessageA(bm1, bm2); + bm3 = am3 = _mm_set_epi64x(0x10000000000ull, 0x0ull); + QuadRound(as0, as1, 0xc19bf2749bdc06a7ull, 0x80deb1fe72be5d74ull); + QuadRound(bs0, bs1, 0xc19bf2749bdc06a7ull, 0x80deb1fe72be5d74ull); + ShiftMessageB(am2, am3, am0); + ShiftMessageB(bm2, bm3, bm0); + QuadRound(as0, as1, am0, 0x240ca1cc0fc19dc6ull, 0xefbe4786e49b69c1ull); + QuadRound(bs0, bs1, bm0, 0x240ca1cc0fc19dc6ull, 0xefbe4786e49b69c1ull); + ShiftMessageB(am3, am0, am1); + ShiftMessageB(bm3, bm0, bm1); + QuadRound(as0, as1, am1, 0x76f988da5cb0a9dcull, 0x4a7484aa2de92c6full); + QuadRound(bs0, bs1, bm1, 0x76f988da5cb0a9dcull, 0x4a7484aa2de92c6full); + ShiftMessageB(am0, am1, am2); + ShiftMessageB(bm0, bm1, bm2); + QuadRound(as0, as1, am2, 0xbf597fc7b00327c8ull, 0xa831c66d983e5152ull); + QuadRound(bs0, bs1, bm2, 0xbf597fc7b00327c8ull, 0xa831c66d983e5152ull); + ShiftMessageB(am1, am2, am3); + ShiftMessageB(bm1, bm2, bm3); + QuadRound(as0, as1, am3, 0x1429296706ca6351ull, 0xd5a79147c6e00bf3ull); + QuadRound(bs0, bs1, bm3, 0x1429296706ca6351ull, 0xd5a79147c6e00bf3ull); + ShiftMessageB(am2, am3, am0); + ShiftMessageB(bm2, bm3, bm0); + QuadRound(as0, as1, am0, 0x53380d134d2c6dfcull, 0x2e1b213827b70a85ull); + QuadRound(bs0, bs1, bm0, 0x53380d134d2c6dfcull, 0x2e1b213827b70a85ull); + ShiftMessageB(am3, am0, am1); + ShiftMessageB(bm3, bm0, bm1); + QuadRound(as0, as1, am1, 0x92722c8581c2c92eull, 0x766a0abb650a7354ull); + QuadRound(bs0, bs1, bm1, 0x92722c8581c2c92eull, 0x766a0abb650a7354ull); + ShiftMessageB(am0, am1, am2); + ShiftMessageB(bm0, bm1, bm2); + QuadRound(as0, as1, am2, 0xc76c51a3c24b8b70ull, 0xa81a664ba2bfe8A1ull); + QuadRound(bs0, bs1, bm2, 0xc76c51a3c24b8b70ull, 0xa81a664ba2bfe8A1ull); + ShiftMessageB(am1, am2, am3); + ShiftMessageB(bm1, bm2, bm3); + QuadRound(as0, as1, am3, 0x106aa070f40e3585ull, 0xd6990624d192e819ull); + QuadRound(bs0, bs1, bm3, 0x106aa070f40e3585ull, 0xd6990624d192e819ull); + ShiftMessageB(am2, am3, am0); + ShiftMessageB(bm2, bm3, bm0); + QuadRound(as0, as1, am0, 0x34b0bcb52748774cull, 0x1e376c0819a4c116ull); + QuadRound(bs0, bs1, bm0, 0x34b0bcb52748774cull, 0x1e376c0819a4c116ull); + ShiftMessageB(am3, am0, am1); + ShiftMessageB(bm3, bm0, bm1); + QuadRound(as0, as1, am1, 0x682e6ff35b9cca4full, 0x4ed8aa4a391c0cb3ull); + QuadRound(bs0, bs1, bm1, 0x682e6ff35b9cca4full, 0x4ed8aa4a391c0cb3ull); + ShiftMessageC(am0, am1, am2); + ShiftMessageC(bm0, bm1, bm2); + QuadRound(as0, as1, am2, 0x8cc7020884c87814ull, 0x78a5636f748f82eeull); + QuadRound(bs0, bs1, bm2, 0x8cc7020884c87814ull, 0x78a5636f748f82eeull); + ShiftMessageC(am1, am2, am3); + ShiftMessageC(bm1, bm2, bm3); + QuadRound(as0, as1, am3, 0xc67178f2bef9a3f7ull, 0xa4506ceb90befffaull); + QuadRound(bs0, bs1, bm3, 0xc67178f2bef9a3f7ull, 0xa4506ceb90befffaull); + as0 = _mm_add_epi32(as0, _mm_load_si128((const __m128i*)INIT0)); + bs0 = _mm_add_epi32(bs0, _mm_load_si128((const __m128i*)INIT0)); + as1 = _mm_add_epi32(as1, _mm_load_si128((const __m128i*)INIT1)); + bs1 = _mm_add_epi32(bs1, _mm_load_si128((const __m128i*)INIT1)); + + /* Extract hash into out */ + Unshuffle(as0, as1); + Unshuffle(bs0, bs1); + Save(out, as0); + Save(out + 16, as1); + Save(out + 32, bs0); + Save(out + 48, bs1); +} + +} + +#endif diff --git a/cpp/tensorrt_llm/nanobind/CMakeLists.txt b/cpp/tensorrt_llm/nanobind/CMakeLists.txt index 4d6fbf9c2607..5dc1de84a308 100755 --- a/cpp/tensorrt_llm/nanobind/CMakeLists.txt +++ b/cpp/tensorrt_llm/nanobind/CMakeLists.txt @@ -11,6 +11,7 @@ set(SRCS batch_manager/kvCacheConnector.cpp batch_manager/kvCacheManager.cpp batch_manager/kvCacheManagerV2Utils.cpp + batch_manager/kvCacheManagerV2.cpp batch_manager/llmRequest.cpp common/tllmExceptions.cpp executor/bindings.cpp @@ -29,6 +30,15 @@ set(SRCS include_directories(${PROJECT_SOURCE_DIR}/include) +# KV Cache Manager V2 include directory (global, needed during compilation of +# SRCS) NOTE: PROJECT_SOURCE_DIR = cpp/ (build_wheel.py passes cpp/ as -S), so +# batch_manager is one level down and the vendored SHA-256 is at +# tensorrt_llm/common/sha256. +set(KV_CACHE_MANAGER_V2_INCLUDE + ${PROJECT_SOURCE_DIR}/tensorrt_llm/batch_manager) +include_directories(${KV_CACHE_MANAGER_V2_INCLUDE} + ${PROJECT_SOURCE_DIR}/tensorrt_llm/common/sha256) + # NOSTRIP: Disable stripping for BOLT compatibility (requires --emit-relocs) if(ENABLE_BOLT_COMPATIBLE) nanobind_add_module(${TRTLLM_NB_MODULE} NOSTRIP ${SRCS}) @@ -36,6 +46,12 @@ else() nanobind_add_module(${TRTLLM_NB_MODULE} ${SRCS}) endif() +# SHA-256 symbols are provided by tensorrt_llm_batch_manager_static (via the +# shared lib). Only add the include directory so headers resolve during +# compilation of kvCacheManagerV2.cpp. +target_include_directories( + ${TRTLLM_NB_MODULE} PRIVATE ${PROJECT_SOURCE_DIR}/tensorrt_llm/common/sha256) + set_property(TARGET ${TRTLLM_NB_MODULE} PROPERTY POSITION_INDEPENDENT_CODE ON) target_link_directories(${TRTLLM_NB_MODULE} PUBLIC diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp new file mode 100644 index 000000000000..7cda4942c104 --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -0,0 +1,1903 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.h" + +#include "kv_cache_manager_v2/blockRadixTree.h" +#include "kv_cache_manager_v2/common.h" +#include "kv_cache_manager_v2/config.h" +#include "kv_cache_manager_v2/eventManager.h" +#include "kv_cache_manager_v2/exceptions.h" +#include "kv_cache_manager_v2/introspection.h" +#include "kv_cache_manager_v2/kvCache.h" +#include "kv_cache_manager_v2/kvCacheManager.h" +#include "kv_cache_manager_v2/lifeCycleRegistry.h" +#include "kv_cache_manager_v2/page.h" +#include "kv_cache_manager_v2/stats.h" +#include "kv_cache_manager_v2/storage/config.h" +#include "kv_cache_manager_v2/storage/core.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace nb = nanobind; +namespace kv = tensorrt_llm::batch_manager::kv_cache_manager_v2; + +namespace tensorrt_llm::nanobind::batch_manager +{ + +// Helper: convert a Python iterable of int|bytes to vector. +// nanobind's variant caster can't auto-convert bytes → DigestToken. +static std::vector castTokenIterable(nb::handle tokens) +{ + std::vector vec; + for (auto item : nb::cast(tokens)) + { + if (nb::isinstance(item)) + { + auto b = nb::cast(item); + if (nb::len(b) != kv::kDIGEST_LEN) + { + throw std::invalid_argument("Token bytes must have length kDIGEST_LEN"); + } + kv::Digest d; + std::memcpy(d.data(), b.c_str(), kv::kDIGEST_LEN); + vec.emplace_back(kv::DigestToken(d)); + } + else + { + vec.emplace_back(nb::cast(item)); + } + } + return vec; +} + +static kv::TypedVec typedPoolSizeList(std::vector const& slotSizeList) +{ + return kv::TypedVec{slotSizeList}; +} + +static kv::TypedVec> typedSlotSizeLists( + std::vector> const& slotSizeLists) +{ + kv::TypedVec> result; + for (auto const& slotSizeList : slotSizeLists) + { + result.push_back(typedPoolSizeList(slotSizeList)); + } + return result; +} + +static kv::TypedVec typedSlotCounts(std::vector const& slotCounts) +{ + kv::TypedVec result; + for (kv::SlotCount slotCount : slotCounts) + { + if (slotCount < 0) + { + throw std::invalid_argument("slot counts must be non-negative"); + } + result.push_back(slotCount); + } + return result; +} + +static nb::object castLifeCycle(kv::LifeCycle const& lifeCycle) +{ + return std::visit([](auto const& concreteLifeCycle) { return nb::cast(concreteLifeCycle); }, lifeCycle); +} + +static kv::KvCache::PriorityCb castPriorityCallback(kv::KvCacheManager const& manager, nb::object callback) +{ + if (callback.is_none()) + { + return {}; + } + if (!PyCallable_Check(callback.ptr())) + { + throw std::invalid_argument("custom_priority_callback must be callable"); + } + + kv::LifeCycleRegistry const* lifeCycles = &manager.lifeCycles(); + return [callback = std::move(callback), lifeCycles](kv::BlockOrdinal ordinal, kv::LifeCycleId lifeCycleId) + { + nb::gil_scoped_acquire acquire; + return nb::cast(callback(ordinal.value(), castLifeCycle(lifeCycles->getLifeCycle(lifeCycleId)))); + }; +} + +static nb::tuple bufferIdTuple(kv::BufferId const& self) +{ + return nb::make_tuple(self.layerId, self.role); +} + +static nb::object optionalIntToObject(std::optional value) +{ + if (!value.has_value()) + { + return nb::none(); + } + return nb::cast(*value); +} + +static nb::tuple reuseScopeTuple(kv::ReuseScope const& self) +{ + return nb::make_tuple(optionalIntToObject(self.loraId), optionalIntToObject(self.salt)); +} + +static nb::list tokenList(std::vector const& tokens) +{ + nb::list result; + for (auto const& tok : tokens) + { + if (auto* id = std::get_if(&tok)) + { + result.append(*id); + } + else + { + auto const& d = std::get(tok); + result.append(nb::bytes(reinterpret_cast(d.data()), d.size())); + } + } + return result; +} + +static nb::list committedTokensList(kv::KvCache const& self) +{ + return tokenList(self.committedTokens()); +} + +static nb::bytes digestBytes(kv::Digest const& digest) +{ + return nb::bytes(reinterpret_cast(digest.data()), digest.size()); +} + +static kv::Digest castDigest(nb::handle value) +{ + if (!nb::isinstance(value)) + { + throw std::invalid_argument("block hash must be bytes"); + } + auto bytes = nb::cast(value); + if (nb::len(bytes) != kv::kDIGEST_LEN) + { + throw std::invalid_argument("block hash bytes must have length kDIGEST_LEN"); + } + kv::Digest digest; + std::memcpy(digest.data(), bytes.c_str(), digest.size()); + return digest; +} + +static bool hasDigestSize(nb::handle value) +{ + return nb::len(value) == kv::kDIGEST_LEN; +} + +static kv::EventBlockHash castEventBlockHash(nb::handle value) +{ + if (PyLong_Check(value.ptr())) + { + return nb::cast(value); + } + return nb::cast(value); +} + +static std::optional castOptionalEventBlockHash(nb::handle value) +{ + return value.is_none() ? std::nullopt : std::optional{castEventBlockHash(value)}; +} + +static std::vector castStoredBlocks(nb::handle values) +{ + std::vector result; + for (nb::handle value : nb::cast(values)) + { + result.push_back(nb::cast(value)); + } + return result; +} + +static std::vector castMmKeys(nb::handle values) +{ + std::vector result; + for (nb::handle value : nb::cast(values)) + { + nb::tuple tuple = nb::cast(value); + if (tuple.size() != 2 && tuple.size() != 3) + { + throw std::invalid_argument("mm_key must have two or three entries"); + } + std::optional uuid; + if (tuple.size() == 3 && !tuple[2].is_none()) + { + uuid = nb::cast(tuple[2]); + } + if (!nb::isinstance(tuple[0])) + { + throw std::invalid_argument("mm_key hash must be bytes"); + } + auto hash = nb::cast(tuple[0]); + result.push_back(kv::MmKey{std::string(hash.c_str(), static_cast(nb::len(hash))), + nb::cast(tuple[1]), std::move(uuid), tuple.size() == 3}); + } + return result; +} + +static nb::list castMmKeys(kv::KVCacheStoredBlockData const& data) +{ + nb::list result; + for (auto const& mmKey : data.mmKeys) + { + auto hash = nb::bytes(mmKey.hash.data(), mmKey.hash.size()); + if (mmKey.hasUuidField) + { + result.append(nb::make_tuple(std::move(hash), mmKey.startOffset, mmKey.uuid)); + } + else + { + result.append(nb::make_tuple(std::move(hash), mmKey.startOffset)); + } + } + return result; +} + +static nb::object castEventData(kv::KVCacheEventData const& data) +{ + return std::visit([](auto const& concreteData) { return nb::cast(concreteData); }, data); +} + +static kv::KVCacheEventData castEventData(nb::handle data) +{ + if (nb::isinstance(data)) + { + return nb::cast(data); + } + if (nb::isinstance(data)) + { + return nb::cast(data); + } + if (nb::isinstance(data)) + { + return nb::cast(data); + } + if (nb::isinstance(data)) + { + return nb::cast(data); + } + throw std::invalid_argument("Unsupported KV cache event data type"); +} + +static std::string pythonHandleRepr(nb::handle value) +{ + nb::object result = nb::steal(PyObject_Repr(value.ptr())); + if (!result) + { + throw nb::python_error(); + } + return nb::cast(result); +} + +template +static std::string pythonRepr(T const& value) +{ + return pythonHandleRepr(nb::cast(value)); +} + +class PythonAttentionDpGather +{ +public: + explicit PythonAttentionDpGather(nb::handle callback) + : mCallback(callback.ptr()) + { + Py_INCREF(mCallback); + } + + ~PythonAttentionDpGather() + { + if (Py_IsInitialized()) + { + nb::gil_scoped_acquire acquire; + Py_DECREF(mCallback); + } + } + + PythonAttentionDpGather(PythonAttentionDpGather const&) = delete; + PythonAttentionDpGather& operator=(PythonAttentionDpGather const&) = delete; + + std::vector> operator()(std::vector const& events) const + { + nb::gil_scoped_acquire acquire; + return nb::cast>>(nb::borrow(mCallback)(events)); + } + +private: + PyObject* mCallback; +}; + +static kv::EventManager::AttentionDpGatherFn castAttentionDpGather(nb::handle callback) +{ + if (callback.is_none()) + { + return {}; + } + if (!PyCallable_Check(callback.ptr())) + { + throw std::invalid_argument("attention_dp_gather must be callable"); + } + auto gather = std::make_shared(callback); + return [gather = std::move(gather)](std::vector const& events) { return (*gather)(events); }; +} + +static nb::object castStatsDelta(kv::KVCacheStatsDelta const& stats) +{ + return nb::cast(stats); +} + +static nb::object castIterationStatsDelta(kv::KVCacheIterationStatsDelta const& stats) +{ + return nb::cast(stats); +} + +static nb::dict castIterationStatsByLifeCycle(kv::IterationStatsByLifeCycle const& statsByLifeCycle) +{ + nb::dict result; + for (auto const& [lifeCycle, stats] : statsByLifeCycle) + { + result[nb::int_(lifeCycle.value())] = castIterationStatsDelta(stats); + } + return result; +} + +static nb::dict castSsmSnapshotIterationStatsByLifeCycle( + kv::SsmSnapshotIterationStatsByLifeCycle const& statsByLifeCycle) +{ + nb::dict result; + for (auto const& [lifeCycle, stats] : statsByLifeCycle) + { + result[nb::int_(lifeCycle.value())] = nb::cast(stats); + } + return result; +} + +static nb::list castPeakBlockStats(kv::PeakBlockStatsByPoolGroup const& statsByPoolGroup) +{ + nb::list result; + for (auto const& stats : statsByPoolGroup) + { + result.append(nb::cast(stats)); + } + return result; +} + +static std::string statsDeltaRepr(kv::KVCacheStatsDelta const& stats) +{ + std::ostringstream stream; + stream << "KVCacheStatsDelta(alloc_total_blocks=" << stats.allocTotalBlocks + << ", alloc_new_blocks=" << stats.allocNewBlocks << ", reused_blocks=" << stats.reusedBlocks + << ", missed_blocks=" << stats.missedBlocks << ')'; + return stream.str(); +} + +static std::string iterationStatsDeltaRepr(kv::KVCacheIterationStatsDelta const& stats) +{ + std::ostringstream stream; + stream << "KVCacheIterationStatsDelta(iter_alloc_total_blocks=" << stats.iterAllocTotalBlocks + << ", iter_alloc_new_blocks=" << stats.iterAllocNewBlocks + << ", iter_reused_blocks=" << stats.iterReusedBlocks + << ", iter_full_reused_blocks=" << stats.iterFullReusedBlocks + << ", iter_partial_reused_blocks=" << stats.iterPartialReusedBlocks + << ", iter_missed_blocks=" << stats.iterMissedBlocks + << ", iter_gen_alloc_blocks=" << stats.iterGenAllocBlocks + << ", iter_onboard_blocks=" << stats.iterOnboardBlocks << ", iter_onboard_bytes=" << stats.iterOnboardBytes + << ", iter_offload_blocks=" << stats.iterOffloadBlocks << ", iter_offload_bytes=" << stats.iterOffloadBytes + << ", iter_intra_device_copy_blocks=" << stats.iterIntraDeviceCopyBlocks + << ", iter_intra_device_copy_bytes=" << stats.iterIntraDeviceCopyBytes + << ", iter_host_dropped_blocks=" << stats.iterHostDroppedBlocks + << ", iter_host_dropped_bytes=" << stats.iterHostDroppedBytes << ')'; + return stream.str(); +} + +static std::string ssmSnapshotStatsDeltaRepr(kv::SsmSnapshotIterationStatsDelta const& stats) +{ + std::ostringstream stream; + stream << "SsmSnapshotIterationStatsDelta(iter_snapshot_lookups=" << stats.iterSnapshotLookups + << ", iter_snapshot_hits=" << stats.iterSnapshotHits << ", iter_snapshot_misses=" << stats.iterSnapshotMisses + << ", iter_reused_tokens=" << stats.iterReusedTokens << ", iter_unreused_tokens=" << stats.iterUnreusedTokens + << ", iter_aligned_snapshot_hits=" << stats.iterAlignedSnapshotHits + << ", iter_unaligned_snapshot_hits=" << stats.iterUnalignedSnapshotHits << ')'; + return stream.str(); +} + +static std::string peakBlockStatsRepr(kv::PoolGroupPeakBlockStats const& stats) +{ + std::ostringstream stream; + stream << "PoolGroupPeakBlockStats(available=" << stats.available << ", unavailable=" << stats.unavailable + << ", evictable=" << stats.evictable << ')'; + return stream.str(); +} + +static nb::object castRequestIds(std::unordered_set const& requestIds) +{ + nb::object result = nb::module_::import_("builtins").attr("set")(); + for (kv::RequestIdType const requestId : requestIds) + { + result.attr("add")(requestId); + } + return result; +} + +static std::optional castOptionalIntAttr(nb::handle obj, char const* attrName) +{ + nb::object attr = nb::steal(PyObject_GetAttrString(obj.ptr(), attrName)); + if (!attr) + { + throw nb::python_error(); + } + if (attr.is_none()) + { + return std::nullopt; + } + return nb::cast(attr); +} + +static kv::ReuseScope castReuseScope(nb::object reuseScope) +{ + if (reuseScope.is_none()) + { + return {}; + } + if (nb::isinstance(reuseScope)) + { + return nb::cast(reuseScope); + } + // Backward-compatible bridge for old callers that still pass lora_task_id as the first argument. + if (PyLong_Check(reuseScope.ptr())) + { + return {nb::cast(reuseScope), std::nullopt}; + } + if (PyObject_HasAttrString(reuseScope.ptr(), "lora_id") && PyObject_HasAttrString(reuseScope.ptr(), "salt")) + { + return {castOptionalIntAttr(reuseScope, "lora_id"), castOptionalIntAttr(reuseScope, "salt")}; + } + throw std::invalid_argument( + "reuse_scope must be None, ReuseScope, an int lora_task_id, or an object with lora_id and salt"); +} + +void KvCacheManagerV2Bindings::initBindings(nb::module_& m) +{ + // Export the C++ debug mode as an immutable Python bool snapshot. + m.attr("NDEBUG") = nb::bool_(!kv::gDebug); + + // ---- Exceptions -------------------------------------------------------- + static nb::object sOutOfMemoryError = nb::exception(m, "OutOfMemoryError"); + static nb::object sHostOOMError = nb::exception(m, "HostOOMError"); + static nb::object sDiskOOMError = nb::exception(m, "DiskOOMError"); + static nb::object sCuOOMError = nb::exception(m, "CuOOMError"); + static nb::object sLogicError = nb::exception(m, "LogicError"); + static nb::object sResourceBusyError = nb::exception(m, "ResourceBusyError"); + static nb::object sOutOfPagesError = nb::exception(m, "OutOfPagesError"); + + // Map kv::AssertionError to Python's builtin AssertionError so shared tests see + // the same exception type as the pure-Python backend (which uses `assert`). + // Registered last so it is tried before the generic std::exception fallback. + nb::register_exception_translator( + [](std::exception_ptr const& p, void*) + { + try + { + if (p) + std::rethrow_exception(p); + } + catch (kv::AssertionError const& e) + { + PyErr_SetString(PyExc_AssertionError, e.what()); + } + }); + + // ---- Enums ------------------------------------------------------------- + nb::enum_(m, "PageStatus") + .value("LOCKED", kv::PageStatus::LOCKED) + .value("HELD", kv::PageStatus::HELD) + .value("DROPPABLE", kv::PageStatus::DROPPABLE) + .export_values(); + + nb::enum_(m, "CacheTier") + .value("GPU_MEM", kv::CacheTier::GPU_MEM) + .value("HOST_MEM", kv::CacheTier::HOST_MEM) + .value("DISK", kv::CacheTier::DISK) + .export_values(); + + // ---- KvCache::Status enum (also accessible as _KVCache.Status) --------- + auto kvCacheStatus = nb::enum_(m, "KvCacheStatus") + .value("ACTIVE", kv::KvCache::Status::ACTIVE) + .value("SUSPENDED", kv::KvCache::Status::SUSPENDED) + .value("CLOSED", kv::KvCache::Status::CLOSED); + + // ---- KV cache events ---------------------------------------------------- + nb::class_(m, "UniqueToken") + .def(nb::init(), nb::arg("token_id"), nb::arg("token_extra_id") = 0) + .def_ro("token_id", &kv::UniqueToken::tokenId) + .def_ro("token_extra_id", &kv::UniqueToken::tokenExtraId) + .def("__eq__", &kv::UniqueToken::operator==, nb::arg("other")) + .def("__repr__", + [](kv::UniqueToken const& self) + { + return "UniqueToken(token_id=" + pythonRepr(self.tokenId) + + ", token_extra_id=" + pythonRepr(self.tokenExtraId) + ")"; + }) + .def("__reduce__", + [](kv::UniqueToken const& self) + { return nb::make_tuple(nb::type(), nb::make_tuple(self.tokenId, self.tokenExtraId)); }); + + nb::class_(m, "KVCacheCreatedData") + .def(nb::init>(), nb::arg("num_blocks_per_cache_level")) + .def_ro("num_blocks_per_cache_level", &kv::KVCacheCreatedData::numBlocksPerCacheLevel) + .def("__eq__", &kv::KVCacheCreatedData::operator==, nb::arg("other")) + .def("__repr__", + [](kv::KVCacheCreatedData const& self) { + return "KVCacheCreatedData(num_blocks_per_cache_level=" + pythonRepr(self.numBlocksPerCacheLevel) + ")"; + }) + .def("__reduce__", + [](kv::KVCacheCreatedData const& self) { + return nb::make_tuple(nb::type(), nb::make_tuple(self.numBlocksPerCacheLevel)); + }); + + nb::class_(m, "KVCacheStoredBlockData") + .def( + "__init__", + [](kv::KVCacheStoredBlockData* self, kv::EventBlockHash blockHash, std::vector tokens, + int cacheLevel, int priority, nb::handle mmKeys, std::optional cacheSalt) + { + new (self) kv::KVCacheStoredBlockData{std::move(blockHash), std::move(tokens), cacheLevel, priority, + castMmKeys(mmKeys), std::move(cacheSalt)}; + }, + nb::arg("block_hash"), nb::arg("tokens"), nb::arg("cache_level"), nb::arg("priority"), + nb::arg("mm_keys") = nb::make_tuple(), nb::arg("cache_salt") = std::nullopt) + .def_ro("block_hash", &kv::KVCacheStoredBlockData::blockHash) + .def_ro("tokens", &kv::KVCacheStoredBlockData::tokens) + .def_ro("cache_level", &kv::KVCacheStoredBlockData::cacheLevel) + .def_ro("priority", &kv::KVCacheStoredBlockData::priority) + .def_prop_ro("mm_keys", [](kv::KVCacheStoredBlockData const& self) { return castMmKeys(self); }) + .def_ro("cache_salt", &kv::KVCacheStoredBlockData::cacheSalt) + .def("__eq__", &kv::KVCacheStoredBlockData::operator==, nb::arg("other")) + .def("__repr__", + [](kv::KVCacheStoredBlockData const& self) + { + return "KVCacheStoredBlockData(block_hash=" + pythonRepr(self.blockHash) + + ", tokens=" + pythonRepr(self.tokens) + ", cache_level=" + pythonRepr(self.cacheLevel) + + ", priority=" + pythonRepr(self.priority) + ", mm_keys=" + pythonHandleRepr(castMmKeys(self)) + + ", cache_salt=" + pythonRepr(self.cacheSalt) + ")"; + }) + .def("__reduce__", + [](kv::KVCacheStoredBlockData const& self) + { + return nb::make_tuple(nb::type(), + nb::make_tuple( + self.blockHash, self.tokens, self.cacheLevel, self.priority, castMmKeys(self), self.cacheSalt)); + }); + + nb::class_(m, "KVCacheStoredData") + .def( + "__init__", + [](kv::KVCacheStoredData* self, nb::object parentHash, nb::object blocks) { + new (self) kv::KVCacheStoredData{castOptionalEventBlockHash(parentHash), castStoredBlocks(blocks)}; + }, + nb::arg("parent_hash").none(), nb::arg("blocks")) + .def_ro("parent_hash", &kv::KVCacheStoredData::parentHash) + .def_ro("blocks", &kv::KVCacheStoredData::blocks) + .def("__eq__", &kv::KVCacheStoredData::operator==, nb::arg("other")) + .def("__repr__", + [](kv::KVCacheStoredData const& self) + { + return "KVCacheStoredData(parent_hash=" + pythonRepr(self.parentHash) + + ", blocks=" + pythonRepr(self.blocks) + ")"; + }) + .def("__reduce__", + [](kv::KVCacheStoredData const& self) { + return nb::make_tuple(nb::type(), nb::make_tuple(self.parentHash, self.blocks)); + }); + + nb::class_(m, "KVCacheRemovedData") + .def(nb::init>(), nb::arg("block_hashes")) + .def_ro("block_hashes", &kv::KVCacheRemovedData::blockHashes) + .def("__eq__", &kv::KVCacheRemovedData::operator==, nb::arg("other")) + .def("__repr__", + [](kv::KVCacheRemovedData const& self) + { return "KVCacheRemovedData(block_hashes=" + pythonRepr(self.blockHashes) + ")"; }) + .def("__reduce__", + [](kv::KVCacheRemovedData const& self) + { return nb::make_tuple(nb::type(), nb::make_tuple(self.blockHashes)); }); + + nb::class_(m, "KVCacheEventDiff") + .def(nb::init(), nb::arg("old_value"), nb::arg("new_value")) + .def_ro("old_value", &kv::KVCacheEventDiff::oldValue) + .def_ro("new_value", &kv::KVCacheEventDiff::newValue) + .def("__eq__", &kv::KVCacheEventDiff::operator==, nb::arg("other")) + .def("__repr__", + [](kv::KVCacheEventDiff const& self) + { + return "KVCacheEventDiff(old_value=" + pythonRepr(self.oldValue) + + ", new_value=" + pythonRepr(self.newValue) + ")"; + }) + .def("__reduce__", + [](kv::KVCacheEventDiff const& self) + { return nb::make_tuple(nb::type(), nb::make_tuple(self.oldValue, self.newValue)); }); + + nb::class_(m, "KVCacheUpdatedData") + .def(nb::init, std::optional>(), + nb::arg("block_hash"), nb::arg("cache_level").none(), nb::arg("priority").none()) + .def_ro("block_hash", &kv::KVCacheUpdatedData::blockHash) + .def_ro("cache_level", &kv::KVCacheUpdatedData::cacheLevel) + .def_ro("priority", &kv::KVCacheUpdatedData::priority) + .def("__eq__", &kv::KVCacheUpdatedData::operator==, nb::arg("other")) + .def("__repr__", + [](kv::KVCacheUpdatedData const& self) + { + return "KVCacheUpdatedData(block_hash=" + pythonRepr(self.blockHash) + + ", cache_level=" + pythonRepr(self.cacheLevel) + ", priority=" + pythonRepr(self.priority) + ")"; + }) + .def("__reduce__", + [](kv::KVCacheUpdatedData const& self) + { + return nb::make_tuple( + nb::type(), nb::make_tuple(self.blockHash, self.cacheLevel, self.priority)); + }); + + nb::class_(m, "KVCacheEvent") + .def( + "__init__", + [](kv::KVCacheEvent* self, int64_t eventId, nb::handle data, int windowSize, + std::optional hashAlgo, std::optional attentionDpRank, + kv::EventLayerGroupId layerGroupId) + { + new (self) kv::KVCacheEvent{ + eventId, castEventData(data), windowSize, std::move(hashAlgo), attentionDpRank, layerGroupId}; + }, + nb::arg("event_id"), nb::arg("data"), nb::arg("window_size"), nb::arg("hash_algo") = std::nullopt, + nb::arg("attention_dp_rank") = std::nullopt, nb::arg("layer_group_id") = std::nullopt) + .def_ro("event_id", &kv::KVCacheEvent::eventId) + .def_prop_ro("data", [](kv::KVCacheEvent const& self) { return castEventData(self.data); }) + .def_ro("window_size", &kv::KVCacheEvent::windowSize) + .def_ro("hash_algo", &kv::KVCacheEvent::hashAlgo) + .def_ro("attention_dp_rank", &kv::KVCacheEvent::attentionDpRank) + .def_ro("layer_group_id", &kv::KVCacheEvent::layerGroupId) + .def("__eq__", &kv::KVCacheEvent::operator==, nb::arg("other")) + .def("__repr__", + [](kv::KVCacheEvent const& self) + { + return "KVCacheEvent(event_id=" + pythonRepr(self.eventId) + ", data=" + + pythonHandleRepr(castEventData(self.data)) + ", window_size=" + pythonRepr(self.windowSize) + + ", hash_algo=" + pythonRepr(self.hashAlgo) + ", attention_dp_rank=" + + pythonRepr(self.attentionDpRank) + ", layer_group_id=" + pythonRepr(self.layerGroupId) + ")"; + }) + .def("__reduce__", + [](kv::KVCacheEvent const& self) + { + return nb::make_tuple(nb::type(), + nb::make_tuple(self.eventId, castEventData(self.data), self.windowSize, self.hashAlgo, + self.attentionDpRank, self.layerGroupId)); + }); + + nb::class_(m, "KVCacheEventManager") + .def( + "__init__", + [](kv::EventManager* self, int maxKvEventEntries, int windowSize, std::optional attentionDpRank, + nb::handle attentionDpGather, std::string hashAlgo, nb::handle windowSizeByLayerGroup) + { + std::map windowSizes; + if (!windowSizeByLayerGroup.is_none()) + { + windowSizes = nb::cast>(windowSizeByLayerGroup); + } + new (self) kv::EventManager(maxKvEventEntries, windowSize, attentionDpRank, + castAttentionDpGather(attentionDpGather), std::move(hashAlgo), std::move(windowSizes)); + }, + nb::arg("max_kv_event_entries"), nb::kw_only(), nb::arg("window_size") = 0, + nb::arg("attention_dp_rank") = std::nullopt, nb::arg("attention_dp_gather") = nb::none(), + nb::arg("hash_algo") = "v2_sha256", nb::arg("window_size_by_layer_group").none() = nb::none()) + .def("add_created_event", &kv::EventManager::addCreatedEvent, nb::arg("num_blocks_per_cache_level"), + nb::arg("layer_group_ids") = std::nullopt, nb::call_guard()) + .def("set_layer_group_window_sizes", &kv::EventManager::setLayerGroupWindowSizes, nb::arg("window_sizes"), + nb::call_guard()) + .def( + "add_stored_event", + [](kv::EventManager& self, nb::object parentHash, nb::object blocks, kv::EventLayerGroupId layerGroupId) + { + self.addStoredEvent( + kv::KVCacheStoredData{castOptionalEventBlockHash(parentHash), castStoredBlocks(blocks)}, + layerGroupId); + }, + nb::arg("parent_hash").none(), nb::arg("blocks"), nb::arg("layer_group_id") = std::nullopt) + .def( + "add_removed_event", + [](kv::EventManager& self, nb::handle blockHashes) + { + if (nb::isinstance(blockHashes)) + { + if (hasDigestSize(blockHashes)) + { + self.addRemovedBlock(castDigest(blockHashes)); + } + return; + } + if (PyLong_Check(blockHashes.ptr()) || nb::isinstance(blockHashes)) + { + self.addRemovedEvent({castEventBlockHash(blockHashes)}); + return; + } + for (nb::handle blockHash : nb::cast(blockHashes)) + { + if (nb::isinstance(blockHash)) + { + if (hasDigestSize(blockHash)) + { + self.addRemovedBlock(castDigest(blockHash)); + } + } + else + { + self.addRemovedEvent({castEventBlockHash(blockHash)}); + } + } + }, + nb::arg("block_hashes")) + .def( + "add_removed_life_cycle_event", + [](kv::EventManager& self, nb::handle blockHash, int lifeCycleId) + { + if (!nb::isinstance(blockHash)) + { + throw std::invalid_argument("block hash must be bytes"); + } + if (hasDigestSize(blockHash)) + { + self.addRemovedLifeCycle(castDigest(blockHash), kv::LifeCycleId{lifeCycleId}); + } + }, + nb::arg("block_hash"), nb::arg("life_cycle_id")) + .def( + "add_updated_event", + [](kv::EventManager& self, nb::handle blockHash, std::optional cacheLevel, + std::optional priority, kv::EventLayerGroupId layerGroupId) + { + if (!cacheLevel.has_value() && !priority.has_value()) + { + return; + } + if (nb::isinstance(blockHash)) + { + if (!hasDigestSize(blockHash)) + { + return; + } + auto digest = castDigest(blockHash); + nb::gil_scoped_release release; + self.addUpdatedEvent(digest, cacheLevel, priority, layerGroupId); + return; + } + auto eventBlockHash = castEventBlockHash(blockHash); + nb::gil_scoped_release release; + self.addUpdatedEvent(std::move(eventBlockHash), cacheLevel, priority, layerGroupId); + }, + nb::arg("block_hash"), nb::kw_only(), nb::arg("cache_level") = std::nullopt, + nb::arg("priority") = std::nullopt, nb::arg("layer_group_id") = std::nullopt) + .def( + "flush_iteration_events", &kv::EventManager::flushIterationEvents, nb::call_guard()) + .def("get_latest_events", &kv::EventManager::getLatestEvents, nb::arg("timeout_ms") = std::nullopt, + nb::call_guard()) + .def_prop_ro("hash_algo", &kv::EventManager::hashAlgorithm) + .def_prop_ro("_hash_algo", &kv::EventManager::hashAlgorithm) + .def_static("_hash_block_key", &kv::EventManager::hashV1BlockKey, nb::arg("tokens"), nb::arg("parent_hash") = 0, + nb::arg("lora_task_id") = std::nullopt, nb::arg("cache_salt_id") = std::nullopt); + + // ---- Statistics -------------------------------------------------------- + nb::class_(m, "KVCacheStatsDelta") + .def( + "__init__", + [](kv::KVCacheStatsDelta* self, int64_t allocTotalBlocks, int64_t allocNewBlocks, int64_t reusedBlocks, + int64_t missedBlocks) + { + new (self) kv::KVCacheStatsDelta{}; + self->allocTotalBlocks = allocTotalBlocks; + self->allocNewBlocks = allocNewBlocks; + self->reusedBlocks = reusedBlocks; + self->missedBlocks = missedBlocks; + }, + nb::arg("alloc_total_blocks") = 0, nb::arg("alloc_new_blocks") = 0, nb::arg("reused_blocks") = 0, + nb::arg("missed_blocks") = 0) + .def_rw("alloc_total_blocks", &kv::KVCacheStatsDelta::allocTotalBlocks) + .def_rw("alloc_new_blocks", &kv::KVCacheStatsDelta::allocNewBlocks) + .def_rw("reused_blocks", &kv::KVCacheStatsDelta::reusedBlocks) + .def_rw("missed_blocks", &kv::KVCacheStatsDelta::missedBlocks) + .def("add", &kv::KVCacheStatsDelta::add, nb::arg("other")) + .def("subtract", &kv::KVCacheStatsDelta::subtract, nb::arg("other")) + .def("clear", &kv::KVCacheStatsDelta::clear) + .def("copy", &kv::KVCacheStatsDelta::copy) + .def_prop_ro("empty", &kv::KVCacheStatsDelta::empty) + .def("__eq__", &kv::KVCacheStatsDelta::operator==, nb::arg("other")) + .def("__repr__", &statsDeltaRepr); + + nb::class_(m, "KVCacheIterationStatsDelta") + .def( + "__init__", + [](kv::KVCacheIterationStatsDelta* self, int64_t iterAllocTotalBlocks, int64_t iterAllocNewBlocks, + int64_t iterReusedBlocks, int64_t iterFullReusedBlocks, int64_t iterPartialReusedBlocks, + int64_t iterMissedBlocks, int64_t iterGenAllocBlocks, int64_t iterOnboardBlocks, + int64_t iterOnboardBytes, int64_t iterOffloadBlocks, int64_t iterOffloadBytes, + int64_t iterIntraDeviceCopyBlocks, int64_t iterIntraDeviceCopyBytes, int64_t iterHostDroppedBlocks, + int64_t iterHostDroppedBytes) + { + new (self) kv::KVCacheIterationStatsDelta{}; + self->iterAllocTotalBlocks = iterAllocTotalBlocks; + self->iterAllocNewBlocks = iterAllocNewBlocks; + self->iterReusedBlocks = iterReusedBlocks; + self->iterFullReusedBlocks = iterFullReusedBlocks; + self->iterPartialReusedBlocks = iterPartialReusedBlocks; + self->iterMissedBlocks = iterMissedBlocks; + self->iterGenAllocBlocks = iterGenAllocBlocks; + self->iterOnboardBlocks = iterOnboardBlocks; + self->iterOnboardBytes = iterOnboardBytes; + self->iterOffloadBlocks = iterOffloadBlocks; + self->iterOffloadBytes = iterOffloadBytes; + self->iterIntraDeviceCopyBlocks = iterIntraDeviceCopyBlocks; + self->iterIntraDeviceCopyBytes = iterIntraDeviceCopyBytes; + self->iterHostDroppedBlocks = iterHostDroppedBlocks; + self->iterHostDroppedBytes = iterHostDroppedBytes; + }, + nb::arg("iter_alloc_total_blocks") = 0, nb::arg("iter_alloc_new_blocks") = 0, + nb::arg("iter_reused_blocks") = 0, nb::arg("iter_full_reused_blocks") = 0, + nb::arg("iter_partial_reused_blocks") = 0, nb::arg("iter_missed_blocks") = 0, + nb::arg("iter_gen_alloc_blocks") = 0, nb::arg("iter_onboard_blocks") = 0, nb::arg("iter_onboard_bytes") = 0, + nb::arg("iter_offload_blocks") = 0, nb::arg("iter_offload_bytes") = 0, + nb::arg("iter_intra_device_copy_blocks") = 0, nb::arg("iter_intra_device_copy_bytes") = 0, + nb::arg("iter_host_dropped_blocks") = 0, nb::arg("iter_host_dropped_bytes") = 0) + .def_rw("iter_alloc_total_blocks", &kv::KVCacheIterationStatsDelta::iterAllocTotalBlocks) + .def_rw("iter_alloc_new_blocks", &kv::KVCacheIterationStatsDelta::iterAllocNewBlocks) + .def_rw("iter_reused_blocks", &kv::KVCacheIterationStatsDelta::iterReusedBlocks) + .def_rw("iter_full_reused_blocks", &kv::KVCacheIterationStatsDelta::iterFullReusedBlocks) + .def_rw("iter_partial_reused_blocks", &kv::KVCacheIterationStatsDelta::iterPartialReusedBlocks) + .def_rw("iter_missed_blocks", &kv::KVCacheIterationStatsDelta::iterMissedBlocks) + .def_rw("iter_gen_alloc_blocks", &kv::KVCacheIterationStatsDelta::iterGenAllocBlocks) + .def_rw("iter_onboard_blocks", &kv::KVCacheIterationStatsDelta::iterOnboardBlocks) + .def_rw("iter_onboard_bytes", &kv::KVCacheIterationStatsDelta::iterOnboardBytes) + .def_rw("iter_offload_blocks", &kv::KVCacheIterationStatsDelta::iterOffloadBlocks) + .def_rw("iter_offload_bytes", &kv::KVCacheIterationStatsDelta::iterOffloadBytes) + .def_rw("iter_intra_device_copy_blocks", &kv::KVCacheIterationStatsDelta::iterIntraDeviceCopyBlocks) + .def_rw("iter_intra_device_copy_bytes", &kv::KVCacheIterationStatsDelta::iterIntraDeviceCopyBytes) + .def_rw("iter_host_dropped_blocks", &kv::KVCacheIterationStatsDelta::iterHostDroppedBlocks) + .def_rw("iter_host_dropped_bytes", &kv::KVCacheIterationStatsDelta::iterHostDroppedBytes) + .def("add", &kv::KVCacheIterationStatsDelta::add, nb::arg("other")) + .def("subtract", &kv::KVCacheIterationStatsDelta::subtract, nb::arg("other")) + .def("clear", &kv::KVCacheIterationStatsDelta::clear) + .def("copy", &kv::KVCacheIterationStatsDelta::copy) + .def_prop_ro("empty", &kv::KVCacheIterationStatsDelta::empty) + .def_prop_ro("iter_cache_hit_rate", &kv::KVCacheIterationStatsDelta::iterCacheHitRate) + .def("__eq__", &kv::KVCacheIterationStatsDelta::operator==, nb::arg("other")) + .def("__repr__", &iterationStatsDeltaRepr); + + m.attr("KVCacheIterationStatsDelta").attr("_field_names") = nb::make_tuple("iter_alloc_total_blocks", + "iter_alloc_new_blocks", "iter_reused_blocks", "iter_full_reused_blocks", "iter_partial_reused_blocks", + "iter_missed_blocks", "iter_gen_alloc_blocks", "iter_onboard_blocks", "iter_onboard_bytes", + "iter_offload_blocks", "iter_offload_bytes", "iter_intra_device_copy_blocks", "iter_intra_device_copy_bytes", + "iter_host_dropped_blocks", "iter_host_dropped_bytes"); + + nb::class_(m, "SsmSnapshotIterationStatsDelta") + .def( + "__init__", + [](kv::SsmSnapshotIterationStatsDelta* self, int64_t iterSnapshotLookups, int64_t iterSnapshotHits, + int64_t iterSnapshotMisses, int64_t iterReusedTokens, int64_t iterUnreusedTokens, + int64_t iterAlignedSnapshotHits, int64_t iterUnalignedSnapshotHits) + { + new (self) kv::SsmSnapshotIterationStatsDelta{}; + self->iterSnapshotLookups = iterSnapshotLookups; + self->iterSnapshotHits = iterSnapshotHits; + self->iterSnapshotMisses = iterSnapshotMisses; + self->iterReusedTokens = iterReusedTokens; + self->iterUnreusedTokens = iterUnreusedTokens; + self->iterAlignedSnapshotHits = iterAlignedSnapshotHits; + self->iterUnalignedSnapshotHits = iterUnalignedSnapshotHits; + }, + nb::arg("iter_snapshot_lookups") = 0, nb::arg("iter_snapshot_hits") = 0, + nb::arg("iter_snapshot_misses") = 0, nb::arg("iter_reused_tokens") = 0, nb::arg("iter_unreused_tokens") = 0, + nb::arg("iter_aligned_snapshot_hits") = 0, nb::arg("iter_unaligned_snapshot_hits") = 0) + .def_rw("iter_snapshot_lookups", &kv::SsmSnapshotIterationStatsDelta::iterSnapshotLookups) + .def_rw("iter_snapshot_hits", &kv::SsmSnapshotIterationStatsDelta::iterSnapshotHits) + .def_rw("iter_snapshot_misses", &kv::SsmSnapshotIterationStatsDelta::iterSnapshotMisses) + .def_rw("iter_reused_tokens", &kv::SsmSnapshotIterationStatsDelta::iterReusedTokens) + .def_rw("iter_unreused_tokens", &kv::SsmSnapshotIterationStatsDelta::iterUnreusedTokens) + .def_rw("iter_aligned_snapshot_hits", &kv::SsmSnapshotIterationStatsDelta::iterAlignedSnapshotHits) + .def_rw("iter_unaligned_snapshot_hits", &kv::SsmSnapshotIterationStatsDelta::iterUnalignedSnapshotHits) + .def("add", &kv::SsmSnapshotIterationStatsDelta::add, nb::arg("other")) + .def("subtract", &kv::SsmSnapshotIterationStatsDelta::subtract, nb::arg("other")) + .def("clear", &kv::SsmSnapshotIterationStatsDelta::clear) + .def("copy", &kv::SsmSnapshotIterationStatsDelta::copy) + .def_prop_ro("empty", &kv::SsmSnapshotIterationStatsDelta::empty) + .def_prop_ro("iter_snapshot_hit_rate", &kv::SsmSnapshotIterationStatsDelta::iterSnapshotHitRate) + .def("__eq__", &kv::SsmSnapshotIterationStatsDelta::operator==, nb::arg("other")) + .def("__repr__", &ssmSnapshotStatsDeltaRepr); + + nb::class_(m, "PoolGroupPeakBlockStats") + .def( + "__init__", + [](kv::PoolGroupPeakBlockStats* self, kv::SlotCount available, kv::SlotCount unavailable, + kv::SlotCount evictable) { + new (self) kv::PoolGroupPeakBlockStats{available, unavailable, evictable}; + }, + nb::arg("available"), nb::arg("unavailable"), nb::arg("evictable")) + .def_ro("available", &kv::PoolGroupPeakBlockStats::available) + .def_ro("unavailable", &kv::PoolGroupPeakBlockStats::unavailable) + .def_ro("evictable", &kv::PoolGroupPeakBlockStats::evictable) + .def("__eq__", &kv::PoolGroupPeakBlockStats::operator==, nb::arg("other")) + .def("__repr__", &peakBlockStatsRepr); + + // ---- Life cycle helpers ------------------------------------------------ + using BlockRange = kv::HalfOpenRange; + nb::class_(m, "HalfOpenRange") + .def(nb::init(), nb::arg("beg"), nb::arg("end")) + .def_prop_ro("beg", [](BlockRange const& self) { return self.beg.value(); }) + .def_prop_ro("end", [](BlockRange const& self) { return self.end.value(); }) + .def("__bool__", [](BlockRange const& self) { return static_cast(self); }) + .def("__len__", &BlockRange::length) + .def("__eq__", [](BlockRange const& self, BlockRange const& other) { return self == other; }) + .def( + "__contains__", [](BlockRange const& self, int item) { return self.contains(kv::BlockOrdinal{item}); }, + nb::arg("item")); + + nb::enum_(m, "PageIndexMode") + .value("SHARED", kv::PageIndexMode::SHARED) + .value("PER_LAYER", kv::PageIndexMode::PER_LAYER); + + nb::class_(m, "ScratchDesc") + .def_ro("range", &kv::ScratchDesc::range) + .def_ro("slot_ids", &kv::ScratchDesc::slotIds) + .def("__bool__", [](kv::ScratchDesc const& self) { return static_cast(self); }); + + nb::class_(m, "AttnLifeCycle") + .def(nb::init, int>(), nb::arg("window_size"), nb::arg("num_sink_blocks")) + .def_prop_ro("window_size", [](kv::AttnLifeCycle const& self) { return self.windowSize; }) + .def_ro("num_sink_blocks", &kv::AttnLifeCycle::numSinkBlocks) + .def("get_stale_range", &kv::AttnLifeCycle::getStaleRange, nb::arg("history_length"), + nb::arg("tokens_per_block")) + .def("__eq__", &kv::AttnLifeCycle::operator==); + + nb::class_(m, "SsmLifeCycle") + .def(nb::init<>()) + .def( + "get_stale_range", &kv::SsmLifeCycle::getStaleRange, nb::arg("history_length"), nb::arg("tokens_per_block")) + .def("__eq__", &kv::SsmLifeCycle::operator==); + + // ---- CUDA event -------------------------------------------------------- + auto cachedCudaEvent + = nb::class_(m, "CachedCudaEvent") + .def(nb::init(), nb::arg("stream")) + .def("query_complete", &kv::CachedCudaEvent::queryComplete, nb::call_guard()) + .def("synchronize", &kv::CachedCudaEvent::synchronize, nb::call_guard()) + .def( + "wait_in_stream", + [](kv::CachedCudaEvent const& self, kv::CudaStream stream) { self.waitInStream(stream); }, + nb::arg("stream"), nb::call_guard()) + .def("close", &kv::CachedCudaEvent::close) + .def("is_closed", &kv::CachedCudaEvent::isClosed); + cachedCudaEvent.attr("NULL") = kv::CachedCudaEvent::makeNull(); + + // ---- ReuseScope -------------------------------------------------------- + nb::class_(m, "ReuseScope") + .def(nb::init, std::optional>(), + nb::arg("lora_id").none() = std::nullopt, nb::arg("salt").none() = std::nullopt) + .def_ro("lora_id", &kv::ReuseScope::loraId) + .def_ro("salt", &kv::ReuseScope::salt) + .def("__len__", [](kv::ReuseScope const&) { return 2; }) + .def( + "__getitem__", + [](kv::ReuseScope const& self, int index) -> nb::object + { + if (index < 0) + { + index += 2; + } + if (index == 0) + { + return optionalIntToObject(self.loraId); + } + if (index == 1) + { + return optionalIntToObject(self.salt); + } + throw nb::index_error("ReuseScope index out of range"); + }, + nb::arg("index")) + .def("__iter__", + [](kv::ReuseScope const& self) + { return nb::steal(PyObject_GetIter(reuseScopeTuple(self).ptr())); }) + .def("__hash__", [](kv::ReuseScope const& self) { return PyObject_Hash(reuseScopeTuple(self).ptr()); }) + .def( + "__eq__", + [](kv::ReuseScope const& self, nb::object other) -> nb::object + { + if (nb::isinstance(other)) + { + return nb::bool_(self == nb::cast(other)); + } + if (PyTuple_Check(other.ptr()) && PyTuple_GET_SIZE(other.ptr()) == 2) + { + return nb::bool_(PyObject_RichCompareBool(reuseScopeTuple(self).ptr(), other.ptr(), Py_EQ) == 1); + } + return nb::not_implemented(); + }, + nb::arg("other")) + .def("__repr__", + [](kv::ReuseScope const& self) + { + std::string repr = "ReuseScope(lora_id="; + repr += self.loraId.has_value() ? std::to_string(*self.loraId) : "None"; + repr += ", salt="; + repr += self.salt.has_value() ? std::to_string(*self.salt) : "None"; + repr += ")"; + return repr; + }); + + // ---- BufferId ---------------------------------------------------------- + nb::class_(m, "BufferId") + .def(nb::init(), nb::arg("layer_id"), nb::arg("role")) + .def_ro("layer_id", &kv::BufferId::layerId) + .def_ro("role", &kv::BufferId::role) + .def("__len__", [](kv::BufferId const&) { return 2; }) + .def( + "__getitem__", + [](kv::BufferId const& self, int index) -> nb::object + { + if (index < 0) + { + index += 2; + } + if (index == 0) + { + return nb::cast(self.layerId); + } + if (index == 1) + { + return nb::cast(self.role); + } + throw nb::index_error("BufferId index out of range"); + }, + nb::arg("index")) + .def("__iter__", + [](kv::BufferId const& self) + { return nb::steal(PyObject_GetIter(bufferIdTuple(self).ptr())); }) + .def("__hash__", [](kv::BufferId const& self) { return PyObject_Hash(bufferIdTuple(self).ptr()); }) + .def( + "__eq__", + [](kv::BufferId const& self, nb::object other) -> nb::object + { + if (nb::isinstance(other)) + { + return nb::bool_(self == nb::cast(other)); + } + if (PyTuple_Check(other.ptr()) && PyTuple_GET_SIZE(other.ptr()) == 2) + { + return nb::bool_(PyObject_RichCompareBool(bufferIdTuple(self).ptr(), other.ptr(), Py_EQ) == 1); + } + return nb::not_implemented(); + }, + nb::arg("other")); + + // ---- Storage layout structs ------------------------------------------- + nb::class_(m, "CoalescedBuffer") + .def_ro("single_buffer_size", &kv::CoalescedBuffer::singleBufferSize) + .def_ro("buffer_ids", &kv::CoalescedBuffer::bufferIds) + .def_prop_ro("size", &kv::CoalescedBuffer::size) + .def_prop_ro("num_buffers", &kv::CoalescedBuffer::numBuffers); + + nb::class_(m, "SlotDescVariant") + .def_prop_ro("layer_group_id", [](kv::SlotDescVariant const& self) { return self.lifeCycleId.value(); }) + .def_prop_ro("coalesced_buffers", [](kv::SlotDescVariant const& self) { return self.coalescedBuffers.raw(); }) + .def_prop_ro("slot_size_list", + [](kv::SlotDescVariant const& self) + { + auto sizes = self.slotSizeList(); + return std::vector(sizes.begin(), sizes.end()); + }); + + nb::class_(m, "SlotDesc") + .def_ro("variants", &kv::SlotDesc::variants) + .def_prop_ro("slot_size_list", + [](kv::SlotDesc const& self) + { + auto sizes = self.slotSizeList(); + return std::vector(sizes.begin(), sizes.end()); + }); + + nb::class_(m, "PoolDesc") + .def_prop_ro("pool_index", [](kv::PoolDesc const& self) { return self.poolIndex.value(); }) + .def_ro("base_address", &kv::PoolDesc::baseAddress) + .def_ro("slot_bytes", &kv::PoolDesc::slotBytes); + + nb::class_(m, "PoolGroupDesc") + .def_prop_ro("pool_group_index", [](kv::PoolGroupDesc const& self) { return self.poolGroupIndex.value(); }) + .def_ro("num_slots", &kv::PoolGroupDesc::numSlots) + .def_ro("slot_desc", &kv::PoolGroupDesc::slotDesc) + .def_prop_ro("pools", [](kv::PoolGroupDesc const& self) { return self.pools.raw(); }); + + nb::class_(m, "ExpandedBuffer") + .def_ro("id", &kv::ExpandedBuffer::id) + .def_ro("expansion", &kv::ExpandedBuffer::expansion) + .def("__eq__", + [](kv::ExpandedBuffer const& a, kv::ExpandedBuffer const& b) + { return a.id == b.id && a.expansion == b.expansion; }); + + nb::class_(m, "AggregatedPageDesc") + .def_ro("base", &kv::AggregatedPageDesc::base) + .def_ro("size", &kv::AggregatedPageDesc::size) + .def_ro("stride", &kv::AggregatedPageDesc::stride) + .def_prop_ro("layer_group_id", [](kv::AggregatedPageDesc const& self) { return self.layerGroupId.value(); }) + .def_ro("buffers", &kv::AggregatedPageDesc::buffers); + + // ---- Config structs ---------------------------------------------------- + // Helper: add __copy__ and __deepcopy__ for aggregate config types. + // All config structs are simple aggregates — default copy construction works. +#define DEF_COPY(cls) \ + .def("__copy__", [](cls const& self) { return cls(self); }) \ + .def( \ + "__deepcopy__", [](cls const& self, nb::dict) { return cls(self); }, nb::arg("memo")) + + nb::class_(m, "GpuCacheTierConfig") + .def(nb::init(), nb::arg("quota")) + .def_rw("quota", &kv::GpuCacheTierConfig::quota) + .def_prop_ro("tier", &kv::GpuCacheTierConfig::tier) + .def("assert_valid", &kv::GpuCacheTierConfig::assertValid) DEF_COPY(kv::GpuCacheTierConfig); + + nb::class_(m, "HostCacheTierConfig") + .def(nb::init(), nb::arg("quota")) + .def_rw("quota", &kv::HostCacheTierConfig::quota) + .def_prop_ro("tier", &kv::HostCacheTierConfig::tier) + .def("assert_valid", &kv::HostCacheTierConfig::assertValid) DEF_COPY(kv::HostCacheTierConfig); + + nb::class_(m, "DiskCacheTierConfig") + .def(nb::init(), nb::arg("quota"), nb::arg("path")) + .def_rw("quota", &kv::DiskCacheTierConfig::quota) + .def_rw("path", &kv::DiskCacheTierConfig::path) + .def_prop_ro("tier", &kv::DiskCacheTierConfig::tier) + .def("assert_valid", &kv::DiskCacheTierConfig::assertValid) DEF_COPY(kv::DiskCacheTierConfig); + + nb::class_(m, "BufferConfig") + .def(nb::init>(), nb::arg("role"), nb::arg("size"), + nb::arg("tokens_per_block_override") = std::nullopt) + .def_rw("role", &kv::BufferConfig::role) + .def_rw("size", &kv::BufferConfig::size) + .def_rw("tokens_per_block_override", &kv::BufferConfig::tokensPerBlockOverride) DEF_COPY(kv::BufferConfig); + + nb::class_(m, "AttentionLayerConfig") + .def(nb::init, std::optional, std::optional>(), + nb::arg("layer_id"), nb::arg("buffers"), nb::arg("sliding_window_size") = std::nullopt, + nb::arg("num_sink_tokens") = std::nullopt) + .def_rw("layer_id", &kv::AttentionLayerConfig::layerId) + .def_rw("buffers", &kv::AttentionLayerConfig::buffers) + .def_rw("sliding_window_size", &kv::AttentionLayerConfig::slidingWindowSize) + .def_rw("num_sink_tokens", &kv::AttentionLayerConfig::numSinkTokens) + .def_prop_ro("window_size", &kv::AttentionLayerConfig::windowSize) DEF_COPY(kv::AttentionLayerConfig); + + nb::enum_(m, "LayerType") + .value("ATTENTION", kv::LayerType::ATTENTION) + .value("SSM", kv::LayerType::SSM); + + nb::class_(m, "SsmLayerConfig") + .def(nb::init>(), nb::arg("layer_id"), nb::arg("buffers")) + .def_rw("layer_id", &kv::SsmLayerConfig::layerId) + .def_rw("buffers", &kv::SsmLayerConfig::buffers) DEF_COPY(kv::SsmLayerConfig); + + nb::class_(m, "KVCacheDesc") + .def(nb::init(), nb::arg("capacity"), nb::arg("history_length")) + .def_rw("capacity", &kv::KVCacheDesc::capacity) + .def_rw("history_length", &kv::KVCacheDesc::historyLength) + .def("__eq__", + [](kv::KVCacheDesc const& self, nb::handle other) + { + if (!nb::isinstance(other)) + { + return false; + } + return self == nb::cast(other); + }) + .def("__repr__", + [](kv::KVCacheDesc const& self) + { + return "KVCacheDesc(capacity=" + std::to_string(self.capacity) + + ", history_length=" + std::to_string(self.historyLength) + ")"; + }) DEF_COPY(kv::KVCacheDesc); + + nb::class_(m, "BatchDesc") + .def( + "__init__", + [](kv::BatchDesc* bd, std::vector kvCaches, int systemPromptLength) + { + new (bd) kv::BatchDesc(); + bd->kvCaches = std::move(kvCaches); + bd->systemPromptLength = systemPromptLength; + }, + nb::arg("kv_caches"), nb::arg("system_prompt_length") = 0) + .def_rw("kv_caches", &kv::BatchDesc::kvCaches) + .def_rw("system_prompt_length", &kv::BatchDesc::systemPromptLength) + .def("__eq__", + [](kv::BatchDesc const& self, nb::handle other) + { + if (!nb::isinstance(other)) + { + return false; + } + return self == nb::cast(other); + }) + .def("__repr__", + [](kv::BatchDesc const& self) + { + std::string repr = "BatchDesc(kv_caches=["; + for (size_t i = 0; i < self.kvCaches.size(); ++i) + { + if (i != 0) + { + repr += ", "; + } + repr += "KVCacheDesc(capacity=" + std::to_string(self.kvCaches[i].capacity) + + ", history_length=" + std::to_string(self.kvCaches[i].historyLength) + ")"; + } + repr += "], system_prompt_length=" + std::to_string(self.systemPromptLength) + ")"; + return repr; + }) DEF_COPY(kv::BatchDesc); + + nb::class_(m, "SwaScratchReuseConfig") + .def( + "__init__", + [](kv::SwaScratchReuseConfig* cfg, int maxRewindLen) + { + new (cfg) kv::SwaScratchReuseConfig(); + cfg->maxRewindLen = maxRewindLen; + cfg->validate(); + }, + nb::arg("max_rewind_len") = 0) + .def_rw("max_rewind_len", &kv::SwaScratchReuseConfig::maxRewindLen) DEF_COPY(kv::SwaScratchReuseConfig); + + nb::class_(m, "KVCacheManagerConfig") + .def( + "__init__", + [](kv::KVCacheManagerConfig* cfg, int tokensPerBlock, std::vector cacheTiers, + nb::list layers, float maxUtilForResume, bool enablePartialReuse, + std::optional typicalStep, std::vector constraints, + std::optional> initialPoolRatio, + std::optional swaScratchReuse, bool commitMinSnapshot, bool enableStats) + { + new (cfg) kv::KVCacheManagerConfig(); + cfg->tokensPerBlock = tokensPerBlock; + cfg->cacheTiers = std::move(cacheTiers); + // Convert Python list of AttentionLayerConfig|SsmLayerConfig to vector. + for (auto item : layers) + { + if (nb::isinstance(item)) + cfg->layers.push_back(nb::cast(item)); + else + cfg->layers.push_back(nb::cast(item)); + } + cfg->maxUtilForResume = maxUtilForResume; + cfg->enablePartialReuse = enablePartialReuse; + cfg->typicalStep = std::move(typicalStep); + cfg->constraints = std::move(constraints); + cfg->initialPoolRatio = std::move(initialPoolRatio); + cfg->swaScratchReuse = std::move(swaScratchReuse); + cfg->commitMinSnapshot = commitMinSnapshot; + cfg->enableStats = enableStats; + // Mirror Python's __post_init__: validate at construction. Config-integrity + // failures raise AssertionError (translated below). + cfg->validate(); + }, + nb::arg("tokens_per_block"), nb::arg("cache_tiers"), nb::arg("layers"), + nb::arg("max_util_for_resume") = 0.97f, nb::arg("enable_partial_reuse") = true, + nb::arg("typical_step") = std::nullopt, nb::arg("constraints") = std::vector{}, + nb::arg("initial_pool_ratio").none() = std::nullopt, nb::arg("swa_scratch_reuse").none() = std::nullopt, + nb::arg("commit_min_snapshot") = false, nb::arg("enable_stats") = true) + .def_rw("tokens_per_block", &kv::KVCacheManagerConfig::tokensPerBlock) + .def_rw("cache_tiers", &kv::KVCacheManagerConfig::cacheTiers) + .def_rw("layers", &kv::KVCacheManagerConfig::layers) + .def_rw("max_util_for_resume", &kv::KVCacheManagerConfig::maxUtilForResume) + .def_rw("enable_partial_reuse", &kv::KVCacheManagerConfig::enablePartialReuse) + .def_rw("typical_step", &kv::KVCacheManagerConfig::typicalStep) + .def_rw("constraints", &kv::KVCacheManagerConfig::constraints) + .def_rw("initial_pool_ratio", &kv::KVCacheManagerConfig::initialPoolRatio) + .def_rw("swa_scratch_reuse", &kv::KVCacheManagerConfig::swaScratchReuse) + .def_rw("commit_min_snapshot", &kv::KVCacheManagerConfig::commitMinSnapshot) + .def_rw("enable_stats", &kv::KVCacheManagerConfig::enableStats) + .def_prop_ro("enable_swa_scratch_reuse", &kv::KVCacheManagerConfig::enableSwaScratchReuse) + .def("validate", &kv::KVCacheManagerConfig::validate) DEF_COPY(kv::KVCacheManagerConfig); + +#undef DEF_COPY + + // ---- PlannedDropHandle ------------------------------------------------- + nb::class_(m, "PlannedDropHandle") + .def("drop", &kv::PlannedDropHandle::drop, nb::call_guard()); + + // ---- KvCache ----------------------------------------------------------- + nb::class_(m, "_KVCache") + .def( + "resume", + [](kv::KvCache& self, nb::object stream) + { + std::optional optStream; + if (!stream.is_none()) + optStream = reinterpret_cast(nb::cast(stream)); + nb::gil_scoped_release rel; + return self.resume(optStream); + }, + nb::arg("cuda_stream") = nb::none()) + .def("suspend", &kv::KvCache::suspend, nb::call_guard()) + .def( + "prefetch", [](kv::KvCache& self, int target) { return self.prefetch(kv::CacheLevel{target}); }, + nb::arg("target"), nb::call_guard()) + .def("close", &kv::KvCache::close, nb::call_guard()) + .def("commit_pending_stats", [](kv::KvCache& self) { return castStatsDelta(self.commitPendingStats()); }) + .def("discard_pending_stats", &kv::KvCache::discardPendingStats) + .def( + "resize", + [](kv::KvCache& self, std::optional capacity, std::optional historyLength) -> bool + { return self.resize(capacity, historyLength); }, + nb::arg("capacity") = std::nullopt, nb::arg("history_length") = std::nullopt, + nb::call_guard()) + .def( + "commit", + [](kv::KvCache& self, nb::object acceptedInputTokens, nb::object beamSearchIndices, bool isEnd) + { + auto vec = castTokenIterable(acceptedInputTokens); + if (!beamSearchIndices.is_none()) + { + PyErr_SetString(PyExc_AssertionError, "beam_search_indices must be None"); + throw nb::python_error(); + } + // Note: an empty token list with is_end=True must still stop committing, + // so we do not early-return on empty; commit() handles it. + nb::gil_scoped_release release; + self.commit(vec, isEnd); + }, + nb::arg("accepted_input_tokens"), nb::arg("beam_search_indices").none() = nb::none(), + nb::arg("is_end") = false) + .def("stop_committing", &kv::KvCache::stopCommitting, nb::call_guard()) + .def( + "get_base_page_indices", + [](kv::KvCache const& self, int layerGroupId, int beamIdx) + { + kv::Span span; + { + nb::gil_scoped_release release; + span = self.getBasePageIndices(kv::LayerGroupId{layerGroupId}, kv::BeamIndex{beamIdx}); + } + // Zero-copy: return a read-only numpy ndarray referencing the internal buffer. + // nb::handle() = no owner; the returned array does not own the data. + // Contract: callers must keep the KvCache alive and must not mutate/resize it + // while using this view. The view is intended for read-only use, matching the + // practical use of Python's array.array/memoryview index buffers. + // TODO(yaoy): switch to nb::ndarray when we have nanobind >= 2.9.0, + // or nb::memoryview for nanobind >= 2.12.0. + return nb::ndarray>( + span.ptr, {static_cast(span.len)}, nb::handle()); + }, + nb::arg("layer_group_id"), nb::arg("beam_idx") = 0) + .def( + "get_ssm_block_base_index", + [](kv::KvCache const& self, int layerGroupId, int beamId) + { return self.getSsmBlockBaseIndex(kv::LayerGroupId{layerGroupId}, kv::BeamIndex{beamId}); }, + nb::arg("layer_group_id"), nb::arg("beam_id") = 0, nb::call_guard()) + .def( + "get_scratch_desc", + [](kv::KvCache const& self, int layerGroupId) + { return self.getScratchDesc(kv::LayerGroupId{layerGroupId}); }, + nb::arg("layer_group_id"), nb::call_guard()) + .def_prop_ro("has_scratch_slots", &kv::KvCache::hasScratchSlots) + .def_prop_rw("enable_swa_scratch_reuse", &kv::KvCache::isSwaScratchReuseEnabled, + [](kv::KvCache& self, bool enable) { self.setEnableSwaScratchReuse(enable); }) + .def("supports_index_mode", &kv::KvCache::supportsIndexMode, nb::arg("mode")) + .def_prop_ro("status", [](kv::KvCache const& kvc) { return kvc.status(); }) + .def_prop_ro("is_active", &kv::KvCache::isActive) + .def_prop_ro("finish_event", + [](kv::KvCache const& self) + { + try + { + return self.finishEvent(); + } + catch (std::bad_optional_access const&) + { + // Python unwrap_optional(None) raises ValueError with this message. + PyErr_SetString(PyExc_ValueError, "Expected non-None value"); + throw nb::python_error(); + } + }) + .def_prop_ro("num_blocks", [](kv::KvCache const& self) { return self.numBlocks().value(); }) + .def_prop_ro("num_committed_blocks", &kv::KvCache::numCommittedBlocks) + .def_prop_ro("num_committed_tokens", &kv::KvCache::numCommittedTokens) + .def_prop_rw("history_length", &kv::KvCache::historyLength, + [](kv::KvCache& self, int hist) { self.setHistoryLength(hist); }) + .def_prop_rw("capacity", &kv::KvCache::capacity, [](kv::KvCache& self, int cap) { self.setCapacity(cap); }) + .def_prop_ro("tokens_per_block", &kv::KvCache::tokensPerBlock) + .def_prop_ro("beam_width", [](kv::KvCache const& self) { return self.beamWidth().value(); }) + .def_prop_rw( + "cuda_stream", + [](kv::KvCache const& self) -> intptr_t { return reinterpret_cast(self.cudaStream()); }, + [](kv::KvCache& self, intptr_t stream) { self.setCudaStream(reinterpret_cast(stream)); }) + .def_rw("id", &kv::KvCache::id) + .def_prop_ro( + "manager", [](kv::KvCache& self) -> kv::KvCacheManager& { return self.manager(); }, + nb::rv_policy::reference_internal) + .def_prop_ro("committed_tokens", &committedTokensList) + .def_prop_ro("reuse_scope", &kv::KvCache::reuseScope) + .def( + "plan_committed_block_drop", [](kv::KvCache& self) { return self.planCommittedBlockDrop(); }, + nb::call_guard()) + .def( + "get_aggregated_page_indices", + [](kv::KvCache const& self, int layerGroupId, int beamIdx, bool validOnly) { + return self.getAggregatedPageIndices(kv::LayerGroupId{layerGroupId}, kv::BeamIndex{beamIdx}, validOnly); + }, + nb::arg("layer_group_id"), nb::arg("beam_idx") = 0, nb::arg("valid_only") = false) + .def( + "set_base_page_index_buf", + [](kv::KvCache& self, int beamIdx, int layerGroupId, nb::object bufObj) + { + kv::BeamIndex const typedBeamIdx{beamIdx}; + kv::LayerGroupId const typedLayerGroupId{layerGroupId}; + if (bufObj.is_none()) + { + self.setBasePageIndexBuf(typedBeamIdx, typedLayerGroupId, nullptr, 0); + return; + } + // Accept any object exporting a 1-D writable int32 ('i') buffer. + // PyBUF_ND is required so that memoryview exports shape + format together; + // without it, Python refuses to present a non-byte format as a flat buffer. + Py_buffer view; + if (PyObject_GetBuffer(bufObj.ptr(), &view, PyBUF_WRITABLE | PyBUF_FORMAT | PyBUF_ND) != 0) + throw nb::python_error(); + struct Cleanup + { + Py_buffer* v; + ~Cleanup() + { + PyBuffer_Release(v); + } + } cleanup{&view}; + if (std::string(view.format) != "i" || view.ndim != 1) + throw std::invalid_argument("set_base_page_index_buf: buffer must be 1-D int32 ('i')"); + self.setBasePageIndexBuf(typedBeamIdx, typedLayerGroupId, static_cast(view.buf), + static_cast(view.len / sizeof(int32_t))); + }, + nb::arg("beam_idx"), nb::arg("layer_group_id"), nb::arg("buf").none()); + + // Make Status accessible as _KVCache.Status. + m.attr("_KVCache").attr("Status") = kvCacheStatus; + + // ---- Introspection ------------------------------------------------------- + auto mIntrospection = m.def_submodule("_introspection", "KV cache manager v2 introspection helpers"); + nb::class_(mIntrospection, "StorageStatistics") + .def_prop_ro("slot_sizes", [](kv::StorageStatistics const& self) { return self.slotSizes.raw(); }) + .def_ro("total", &kv::StorageStatistics::total) + .def_ro("free", &kv::StorageStatistics::free) + .def_ro("evictable", &kv::StorageStatistics::evictable) + .def_prop_ro("available", &kv::StorageStatistics::available) + .def_prop_ro("unavailable", &kv::StorageStatistics::unavailable); + mIntrospection.def( + "active_page_stats", + [](kv::KvCache const& kvCache) + { + auto [counts, unscheduledEvictable] = kv::KvCacheIntrospection::activePageStats(kvCache); + return std::make_tuple(std::move(counts.raw()), std::move(unscheduledEvictable.raw())); + }, + nb::arg("kv_cache"), nb::call_guard()); + mIntrospection.def("all_tree_pages_droppable", &kv::KvCacheIntrospection::allTreePagesDroppable, nb::arg("manager"), + nb::call_guard()); + mIntrospection.def( + "is_commit_allowed", + [](kv::KvCache const& kvCache) { return kvCache.commitState() == kv::KvCache::CommitState::ALLOWED; }, + nb::arg("kv_cache"), nb::call_guard()); + mIntrospection.def( + "current_gpu_ratio", + [](kv::KvCacheManager& manager) + { + auto ratio = manager.storage().getRatioList(kv::kGpuLevel); + return std::move(ratio.raw()); + }, + nb::arg("manager"), nb::call_guard()); + // White-box test hooks: mutate auto-tuner state so accuracy tests can force + // a pool rebalance. Mirror the Python manager's internal attributes. + mIntrospection.def( + "set_num_sampled_kv_caches", + [](kv::KvCacheManager& manager, int value) { kv::KvCacheIntrospection::setNumSampledKvCaches(manager, value); }, + nb::arg("manager"), nb::arg("value"), nb::call_guard()); + mIntrospection.def( + "set_last_adjustment_time", + [](kv::KvCacheManager& manager, double value) + { kv::KvCacheIntrospection::setLastAdjustmentTime(manager, value); }, + nb::arg("manager"), nb::arg("value"), nb::call_guard()); + mIntrospection.def( + "set_target_ratio_list_gpu", + [](kv::KvCacheManager& manager, std::vector ratios) + { + kv::KvCacheIntrospection::setTargetRatioListGpu( + manager, kv::TypedVec{std::move(ratios)}); + }, + nb::arg("manager"), nb::arg("ratios"), nb::call_guard()); + mIntrospection.def( + "storage_statistics", + [](kv::KvCacheManager& manager, int cacheLevel) + { + auto stats = kv::KvCacheIntrospection::storageStatistics(manager, kv::CacheLevel{cacheLevel}); + return std::move(stats.raw()); + }, + nb::arg("manager"), nb::arg("cache_level") = kv::kGpuLevel.value(), nb::call_guard()); + mIntrospection.def( + "life_cycle_pool_group_indices", + [](kv::KvCacheManager& manager) + { + std::vector result; + result.reserve(manager.lifeCycles().size().value()); + for (kv::LifeCycleId lifeCycle{0}; lifeCycle < manager.lifeCycles().size(); ++lifeCycle) + { + result.push_back(manager.storage().getPoolGroupIndex(lifeCycle).value()); + } + return result; + }, + nb::arg("manager"), nb::call_guard()); + // White-box reuse-tree introspection: mirror the Python manager's _life_cycles + // and _radix_tree attributes so shared tests can inspect reuse state. + mIntrospection.def( + "attention_life_cycle_ids", + [](kv::KvCacheManager& manager) + { + std::vector result; + for (auto const& [lcId, attn] : manager.lifeCycles().attentionLifeCycles()) + result.push_back(lcId.value()); + return result; + }, + nb::arg("manager"), nb::call_guard()); + mIntrospection.def( + "swa_life_cycle_ids", + [](kv::KvCacheManager& manager) + { + std::vector result; + for (auto const& [lcId, attn] : manager.lifeCycles().attentionLifeCycles()) + if (attn->windowSize.has_value()) + result.push_back(lcId.value()); + return result; + }, + nb::arg("manager"), nb::call_guard()); + mIntrospection.def( + "ssm_life_cycle_id", + [](kv::KvCacheManager& manager) -> std::optional + { + auto id = manager.lifeCycles().ssmLifeCycleId(); + if (id.has_value()) + return id->value(); + return std::nullopt; + }, + nb::arg("manager"), nb::call_guard()); + // Returns (num_tokens, pages) where pages[i] is None for a block with no page in + // this lifecycle, else (slot_id, num_tokens_in_block) with num_tokens_in_block = -1 + // for a non-SSM (attention) page. + mIntrospection.def( + "reuse_match_pages", + [](kv::KvCacheManager& manager, nb::object reuseScope, nb::object tokens, int lcId, bool enablePartial) + { + auto rs = castReuseScope(reuseScope); + auto vec = castTokenIterable(tokens); + int numTokens = 0; + std::vector>> pages; + { + nb::gil_scoped_release release; + auto matchResult = manager.radixTree().match(rs, vec, enablePartial); + numTokens = matchResult.numTokens; + kv::LifeCycleId lc{lcId}; + pages.reserve(matchResult.blocks.stdSize()); + for (auto* block : matchResult.blocks) + { + auto* page = block->storage.at(lc); + if (page == nullptr) + { + pages.emplace_back(std::nullopt); + continue; + } + int const slotId = page->slotId().value(); + int numTokensInBlock = -1; + if (auto* ssm = dynamic_cast(page)) + numTokensInBlock = ssm->numTokensInBlock; + pages.emplace_back(std::make_pair(slotId, numTokensInBlock)); + } + } + return std::make_tuple(numTokens, std::move(pages)); + }, + nb::arg("manager"), nb::arg("reuse_scope"), nb::arg("tokens"), nb::arg("lc_id"), + nb::arg("enable_partial") = false); + // Returns (num_tokens, counts) where counts[i] is None for a block with no page in + // this lifecycle, else the matched page's planned_drop_count. + mIntrospection.def( + "reuse_match_planned_drop_counts", + [](kv::KvCacheManager& manager, nb::object reuseScope, nb::object tokens, int lcId, bool enablePartial) + { + auto rs = castReuseScope(reuseScope); + auto vec = castTokenIterable(tokens); + int numTokens = 0; + std::vector> counts; + { + nb::gil_scoped_release release; + auto matchResult = manager.radixTree().match(rs, vec, enablePartial); + numTokens = matchResult.numTokens; + kv::LifeCycleId lc{lcId}; + counts.reserve(matchResult.blocks.stdSize()); + for (auto* block : matchResult.blocks) + { + auto* page = block->storage.at(lc); + if (page == nullptr) + counts.emplace_back(std::nullopt); + else + counts.emplace_back(page->plannedDropCount); + } + } + return std::make_tuple(numTokens, std::move(counts)); + }, + nb::arg("manager"), nb::arg("reuse_scope"), nb::arg("tokens"), nb::arg("lc_id"), + nb::arg("enable_partial") = false); + mIntrospection.def( + "pool_group_index", + [](kv::KvCacheManager& manager, int lcId) + { return manager.storage().getPoolGroupIndex(kv::LifeCycleId{lcId}).value(); }, + nb::arg("manager"), nb::arg("lc_id")); + mIntrospection.def( + "compute_slots_for_batch", + [](kv::KvCacheManager& manager, kv::BatchDesc const& batch, int tokensPerBlock, + std::optional const& swaScratchReuse) { + return kv::KvCacheIntrospection::computeSlotsForBatch(manager, batch, tokensPerBlock, swaScratchReuse) + .raw(); + }, + nb::arg("manager"), nb::arg("batch"), nb::arg("tokens_per_block"), nb::arg("swa_scratch_reuse") = std::nullopt); + mIntrospection.def( + "storage_utilization", + [](kv::KvCacheManager& manager, int cacheLevel) + { + auto utilization = manager.storage().getUtilization(kv::CacheLevel{cacheLevel}); + return std::move(utilization.raw()); + }, + nb::arg("manager"), nb::arg("cache_level") = kv::kGpuLevel.value(), nb::call_guard()); + mIntrospection.def( + "grains_for_slots", + [](kv::SlotCount numSlots, std::vector const& slotSizeList, size_t granularity) + { + if (numSlots < 0) + { + throw std::invalid_argument("num_slots must be non-negative"); + } + return kv::CacheLevelStorage::grainsForSlots(numSlots, typedPoolSizeList(slotSizeList), granularity); + }, + nb::arg("num_slots"), nb::arg("slot_size_list"), nb::arg("granularity"), + nb::call_guard()); + mIntrospection.def( + "grains_to_slots", + [](size_t pgGrains, std::vector const& slotSizeList, size_t granularity) + { + auto [slots, used] + = kv::CacheLevelStorage::grainsToSlots(pgGrains, typedPoolSizeList(slotSizeList), granularity); + return std::make_tuple(slots, used); + }, + nb::arg("pg_grains"), nb::arg("slot_size_list"), nb::arg("granularity"), + nb::call_guard()); + mIntrospection.def( + "ratio_to_slot_count_list", + [](size_t quota, std::vector> const& slotSizeLists, std::vector const& ratioList, + size_t granularity, std::vector const& minSlots) + { + auto slotCountList = kv::CacheLevelStorage::ratioToSlotCountList(quota, typedSlotSizeLists(slotSizeLists), + kv::TypedVec{ratioList}, granularity, typedSlotCounts(minSlots)); + std::vector result; + result.reserve(slotCountList.stdSize()); + for (kv::SlotCount slotCount : slotCountList) + { + result.push_back(slotCount); + } + return result; + }, + nb::arg("quota"), nb::arg("slot_size_lists"), nb::arg("ratio_list"), nb::arg("granularity"), + nb::arg("min_slots"), nb::call_guard()); + + // ---- KvCacheManager ---------------------------------------------------- + nb::class_(m, "KVCacheManager") + .def( + "__init__", + [](kv::KvCacheManager* self, kv::KVCacheManagerConfig const& config, nb::object eventManager) + { + std::shared_ptr eventSink; + if (!eventManager.is_none()) + { + eventSink = nb::cast>(eventManager); + } + nb::gil_scoped_release release; + new (self) kv::KvCacheManager(config, std::move(eventSink)); + }, + nb::arg("config"), nb::arg("event_manager").none() = nb::none()) + .def("shutdown", &kv::KvCacheManager::shutdown, nb::call_guard()) + .def( + "clear_reusable_blocks", &kv::KvCacheManager::clearReusableBlocks, nb::call_guard()) + .def( + "create_kv_cache", + [](std::shared_ptr self, nb::object reuseScopeObj, nb::object inputTokens, + std::optional id, nb::object customPriorityCallback, + std::optional expectedPromptLength) + { + kv::ReuseScope reuseScope = castReuseScope(std::move(reuseScopeObj)); + std::vector tokens; + bool const hasInputTokens = !inputTokens.is_none(); + if (!inputTokens.is_none()) + { + tokens = castTokenIterable(inputTokens); + } + if (!expectedPromptLength.has_value() && hasInputTokens) + { + expectedPromptLength = static_cast(tokens.size()); + } + kv::KvCache::PriorityCb priorityCb = castPriorityCallback(*self, std::move(customPriorityCallback)); + nb::gil_scoped_release release; + return self->createKvCache( + std::move(reuseScope), tokens, id, std::move(priorityCb), expectedPromptLength); + }, + nb::arg("reuse_scope") = nb::none(), nb::arg("input_tokens") = nb::none(), nb::arg("id") = std::nullopt, + nb::arg("custom_priority_callback") = nb::none(), nb::arg("expected_prompt_length") = std::nullopt) + .def( + "probe_reuse", + [](std::shared_ptr self, nb::object reuseScopeObj, nb::object inputTokens) + { + kv::ReuseScope reuseScope = castReuseScope(std::move(reuseScopeObj)); + std::vector tokens; + if (!inputTokens.is_none()) + { + tokens = castTokenIterable(inputTokens); + } + nb::gil_scoped_release release; + return self->probeReuse(std::move(reuseScope), tokens); + }, + nb::arg("reuse_scope") = nb::none(), nb::arg("input_tokens") = nb::none()) + .def("get_mem_pool_base_address", &kv::KvCacheManager::getMemPoolBaseAddress, nb::arg("layer_id"), + nb::arg("data_role"), nb::arg("index_mode") = std::nullopt, nb::call_guard()) + .def("get_page_stride", &kv::KvCacheManager::getPageStride, nb::arg("layer_id"), nb::arg("data_role")) + .def("get_page_index_scale", &kv::KvCacheManager::getPageIndexScale, nb::arg("layer_id"), nb::arg("data_role")) + .def("get_page_index_upper_bound", &kv::KvCacheManager::getPageIndexUpperBound, nb::arg("layer_id"), + nb::arg("data_role")) + .def( + "resize", + [](kv::KvCacheManager& self, int cacheLevel, size_t quota, bool bestEfforts) + { return self.resize(kv::CacheLevel{cacheLevel}, quota, bestEfforts); }, + nb::arg("cache_level"), nb::arg("quota"), nb::arg("best_efforts") = false, + nb::call_guard()) + .def( + "get_quota", + [](kv::KvCacheManager const& self, int cacheLevel) { return self.getQuota(kv::CacheLevel{cacheLevel}); }, + nb::arg("cache_level")) + .def("get_committed_stats", + [](kv::KvCacheManager const& self) { return castStatsDelta(self.getCommittedStats()); }) + .def("get_and_reset_iteration_stats", + [](kv::KvCacheManager& self) { return castIterationStatsByLifeCycle(self.getAndResetIterationStats()); }) + .def("get_and_reset_ssm_snapshot_iteration_stats", + [](kv::KvCacheManager& self) + { return castSsmSnapshotIterationStatsByLifeCycle(self.getAndResetSsmSnapshotIterationStats()); }) + .def( + "get_and_reset_iteration_peak_block_stats", + [](kv::KvCacheManager& self, int cacheLevel) + { return castPeakBlockStats(self.getAndResetIterationPeakBlockStats(kv::CacheLevel{cacheLevel})); }, + nb::arg("cache_level")) + .def("mark_stats_dirty", &kv::KvCacheManager::markStatsDirty, nb::arg("kv_cache_id").none()) + .def("clear_stats_dirty", &kv::KvCacheManager::clearStatsDirty, nb::arg("kv_cache_id").none()) + .def("get_dirty_stats_kv_cache_ids", + [](kv::KvCacheManager const& self) { return castRequestIds(self.getDirtyStatsKvCacheIds()); }) + .def("mark_stats_excluded", &kv::KvCacheManager::markStatsExcluded, nb::arg("kv_cache_id").none()) + .def("clear_stats_excluded", &kv::KvCacheManager::clearStatsExcluded, nb::arg("kv_cache_id").none()) + .def("is_stats_excluded", &kv::KvCacheManager::isStatsExcluded, nb::arg("kv_cache_id").none()) + .def_prop_ro("tokens_per_block", &kv::KvCacheManager::tokensPerBlock) + .def_prop_ro("event_manager", + [](kv::KvCacheManager const& self) + { return std::dynamic_pointer_cast(self.eventSink()); }) + .def_prop_ro("init_config", [](kv::KvCacheManager const& self) { return self.config(); }) + .def_prop_ro("cache_tier_list", [](kv::KvCacheManager const& self) { return self.cacheTierList().raw(); }) + .def_prop_ro("all_buffer_ids", &kv::KvCacheManager::allBufferIds) + .def_prop_ro("pool_group_descs", [](kv::KvCacheManager const& self) { return self.poolGroupDescs().raw(); }) + .def("clamp_max_seq_len_for_mem", &kv::KvCacheManager::clampMaxSeqLenForMem, nb::arg("batch_size"), + nb::arg("token_num_upper_bound"), nb::call_guard()) + .def_prop_ro("allow_seq_rebasing", &kv::KvCacheManager::allowSeqRebasing) + .def_prop_ro("enable_partial_match", &kv::KvCacheManager::enablePartialMatch) + .def_prop_ro("enable_swa_scratch_reuse", &kv::KvCacheManager::isSwaScratchReuseEnabled) + .def( + "supports_index_mode", + [](kv::KvCacheManager const& self, kv::PageIndexMode mode) -> nb::object + { + auto result = self.supportsIndexMode(mode); + if (!result.has_value()) + return nb::none(); + return nb::cast(*result); + }, + nb::arg("mode")) + .def_prop_ro("commit_min_snapshot", &kv::KvCacheManager::commitMinSnapshot) + .def_prop_ro("num_layers", &kv::KvCacheManager::numLayers) + .def_prop_ro("layer_ids", &kv::KvCacheManager::layerIds) + .def_prop_ro( + "layer_grouping", [](kv::KvCacheManager const& self) { return self.layerGrouping().raw(); }, + "Layers grouped by shared lifecycle/pool allocation. The iteration order of the " + "layer lists (and of the groups) is NOT an API contract and may differ across " + "backends/runs; do not rely on it for buffer/pool memory order -- use " + "pool_group_descs (PoolGroupDesc.pools[i].base_address + coalesced_buffers) instead.") + .def( + "get_layer_group_id", + [](kv::KvCacheManager const& self, kv::LayerId layerId) { return self.getLayerGroupId(layerId).value(); }, + nb::arg("layer_id")) + .def("get_page_index_converter", &kv::KvCacheManager::getPageIndexConverter, nb::arg("layer_id"), + nb::arg("data_role")) + .def( + "get_aggregated_pages", + [](kv::KvCacheManager const& self, nb::object buffers) + { + std::vector ids; + for (auto item : nb::cast(buffers)) + ids.push_back(nb::cast(item)); + nb::gil_scoped_release release; + return self.getAggregatedPages(ids); + }, + nb::arg("buffers")) + .def("adjust", &kv::KvCacheManager::adjust, nb::call_guard()) + .def_prop_ro("need_adjustment", &kv::KvCacheManager::needAdjustment); + + // ---- PageIndexConverter ------------------------------------------------ + nb::class_(m, "PageIndexConverter") + .def_ro("scale", &kv::PageIndexConverter::scale) + .def_ro("expansion", &kv::PageIndexConverter::expansion) + .def_ro("layer_offset", &kv::PageIndexConverter::layerOffset) + .def_ro("scratch_pages_per_block", &kv::PageIndexConverter::scratchPagesPerBlock) + .def( + "__call__", + [](kv::PageIndexConverter const& self, std::vector const& baseIndices, + std::optional indexMode, nb::object scratchObj) -> std::vector + { + kv::ScratchDesc const* scratch = nullptr; + std::optional scratchHolder; + if (!scratchObj.is_none()) + { + scratchHolder = nb::cast(scratchObj); + scratch = &*scratchHolder; + } + return self(baseIndices, indexMode, scratch); + }, + nb::arg("base_indices"), nb::arg("index_mode") = nb::none(), nb::arg("scratch") = nb::none()); +} + +} // namespace tensorrt_llm::nanobind::batch_manager diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.h b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.h new file mode 100644 index 000000000000..c2bb4a51c1b6 --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.h @@ -0,0 +1,30 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +namespace tensorrt_llm::nanobind::batch_manager +{ + +struct KvCacheManagerV2Bindings +{ + static void initBindings(::nanobind::module_& m); +}; + +} // namespace tensorrt_llm::nanobind::batch_manager diff --git a/cpp/tensorrt_llm/nanobind/bindings.cpp b/cpp/tensorrt_llm/nanobind/bindings.cpp index 2c724adc279c..a2054dbd7217 100644 --- a/cpp/tensorrt_llm/nanobind/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/bindings.cpp @@ -39,6 +39,7 @@ #include "tensorrt_llm/nanobind/batch_manager/cacheTransceiver.h" #include "tensorrt_llm/nanobind/batch_manager/kvCacheConnector.h" #include "tensorrt_llm/nanobind/batch_manager/kvCacheManager.h" +#include "tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.h" #include "tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2Utils.h" #include "tensorrt_llm/nanobind/batch_manager/llmRequest.h" #include "tensorrt_llm/nanobind/common/tllmExceptions.h" @@ -135,6 +136,9 @@ NB_MODULE(TRTLLM_NB_MODULE, m) auto mInternalBatchManager = mInternal.def_submodule("batch_manager", "Batch manager internal bindings"); auto mInternalBatchManagerKvCacheV2Utils = mInternalBatchManager.def_submodule("kv_cache_manager_v2_utils", "KV Cache Manager V2 Utils bindings"); + auto mInternalBatchManagerKvCacheV2 + = mInternalBatchManager.def_submodule("kv_cache_manager_v2", "KV Cache Manager V2 bindings"); + tensorrt_llm::nanobind::batch_manager::KvCacheManagerV2Bindings::initBindings(mInternalBatchManagerKvCacheV2); auto mInternalThop = mInternal.def_submodule("thop", "Torch op internal bindings"); auto mExceptions = m.def_submodule("exceptions", "Exceptions internal bindings"); diff --git a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt index 75329d192a8c..02f214986712 100644 --- a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt +++ b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt @@ -24,6 +24,19 @@ add_gtest(contextTransferCoordinatorTest contextTransferCoordinatorTest.cpp) add_gtest(evictionPolicyTest evictionPolicyTest.cpp) add_gtest(kvCacheManagerTest kvCacheManagerTest.cpp) add_gtest(kvCacheManagerFabricMemoryTest kvCacheManagerFabricMemoryTest.cpp) +add_gtest(kvCacheManagerV2TypedIndexTest kvCacheManagerV2TypedIndexTest.cpp) +target_include_directories( + kvCacheManagerV2TypedIndexTest + PRIVATE ${PROJECT_SOURCE_DIR}/tensorrt_llm/batch_manager) +add_gtest(kvCacheManagerV2HostMemTest kvCacheManagerV2HostMemTest.cpp) +target_include_directories( + kvCacheManagerV2HostMemTest + PRIVATE ${PROJECT_SOURCE_DIR}/tensorrt_llm/batch_manager) +add_gtest(kvCacheManagerV2StatsTest kvCacheManagerV2StatsTest.cpp) +target_include_directories( + kvCacheManagerV2StatsTest + PRIVATE ${PROJECT_SOURCE_DIR}/tensorrt_llm/batch_manager + ${PROJECT_SOURCE_DIR}/tensorrt_llm/common/sha256) add_gtest(kvCacheUtilsTest kvCacheUtilsTest.cpp) add_gtest(llmRequestTest llmRequestTest.cpp) add_gtest(microBatchSchedulerTest microBatchSchedulerTest.cpp) diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2HostMemTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2HostMemTest.cpp new file mode 100644 index 000000000000..82ba4c049f86 --- /dev/null +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2HostMemTest.cpp @@ -0,0 +1,206 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/exceptions.h" +#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMem.h" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace +{ + +using namespace tensorrt_llm::batch_manager::kv_cache_manager_v2; + +class ScopedEnv +{ +public: + ScopedEnv(char const* name, std::optional value) + : mName(name) + { + if (char const* oldValue = std::getenv(name); oldValue != nullptr) + { + mOldValue = oldValue; + } + if (value.has_value()) + { + if (::setenv(name, value->c_str(), /*overwrite=*/1) != 0) + { + throw std::system_error(errno, std::generic_category(), "setenv failed"); + } + } + else if (::unsetenv(name) != 0) + { + throw std::system_error(errno, std::generic_category(), "unsetenv failed"); + } + } + + ~ScopedEnv() + { + if (mOldValue.has_value()) + { + ::setenv(mName.c_str(), mOldValue->c_str(), /*overwrite=*/1); + } + else + { + ::unsetenv(mName.c_str()); + } + } + +private: + std::string mName; + std::optional mOldValue; +}; + +int gMadviseErrno = 0; +int gCapturedAdvice = 0; +int gMemsetCalls = 0; + +int captureMadvise(void*, size_t, int advice) +{ + gCapturedAdvice = advice; + return 0; +} + +int failMadvise(void*, size_t, int) +{ + errno = gMadviseErrno; + return -1; +} + +void* countMemset(void* ptr, int value, size_t size) +{ + ++gMemsetCalls; + return std::memset(ptr, value, size); +} + +TEST(KvCacheManagerV2HostMemTest, SelectsConfiguredPageMode) +{ + hostMadvisePageMode(MemAddress{1}, HostMem::kAlignment, true, captureMadvise); + EXPECT_EQ(gCapturedAdvice, MADV_HUGEPAGE); + hostMadvisePageMode(MemAddress{1}, HostMem::kAlignment, false, captureMadvise); + EXPECT_EQ(gCapturedAdvice, MADV_NOHUGEPAGE); + + ScopedEnv defaultThp("TLLM_KV_CACHE_MANAGER_V2_THP", std::nullopt); + EXPECT_TRUE(hostUseThp()); + { + ScopedEnv disableThp("TLLM_KV_CACHE_MANAGER_V2_THP", "0"); + EXPECT_FALSE(hostUseThp()); + } +} + +TEST(KvCacheManagerV2HostMemTest, ReadsPrefaultThreadConfiguration) +{ + ScopedEnv disablePrefault("TLLM_KV_CACHE_MANAGER_V2_PREFAULT_THREADS", "0"); + EXPECT_EQ(hostPrefaultThreads(), 0); + { + ScopedEnv threeThreads("TLLM_KV_CACHE_MANAGER_V2_PREFAULT_THREADS", "3"); + EXPECT_EQ(hostPrefaultThreads(), 3); + } +} + +class PrefaultFallbackTest : public testing::TestWithParam +{ +}; + +TEST_P(PrefaultFallbackTest, TouchesMemory) +{ + std::vector data(HostMem::kAlignment, 0xFF); + gMadviseErrno = GetParam(); + gMemsetCalls = 0; + hostPrefaultChunk(reinterpret_cast(data.data()), data.size(), failMadvise, countMemset); + EXPECT_EQ(gMemsetCalls, 1); + EXPECT_TRUE(std::all_of(data.begin(), data.end(), [](unsigned char value) { return value == 0; })); +} + +INSTANTIATE_TEST_SUITE_P(UnsupportedPopulateWrite, PrefaultFallbackTest, testing::Values(EINVAL, ENOSYS)); + +TEST(KvCacheManagerV2HostMemTest, ConvertsPrefaultEnomem) +{ + std::vector data(HostMem::kAlignment); + gMadviseErrno = ENOMEM; + EXPECT_THROW(hostPrefaultChunk(reinterpret_cast(data.data()), data.size(), failMadvise, countMemset), + HostOOMError); +} + +TEST(KvCacheManagerV2HostMemTest, PropagatesOtherPrefaultErrors) +{ + std::vector data(HostMem::kAlignment); + gMadviseErrno = EIO; + try + { + hostPrefaultChunk(reinterpret_cast(data.data()), data.size(), failMadvise, countMemset); + FAIL() << "Expected std::system_error"; + } + catch (std::system_error const& error) + { + EXPECT_EQ(error.code().value(), EIO); + } +} + +class HostMemPageModeTest : public testing::TestWithParam +{ +}; + +TEST_P(HostMemPageModeTest, RegistersAndResizesWithPrefaultDisabled) +{ + ScopedEnv thp("TLLM_KV_CACHE_MANAGER_V2_THP", std::string(GetParam())); + ScopedEnv prefault("TLLM_KV_CACHE_MANAGER_V2_PREFAULT_THREADS", "0"); + ASSERT_EQ(cudaSetDevice(0), cudaSuccess); + HostMem memory(HostMem::kAlignment); + EXPECT_NE(memory.address(), 0); + EXPECT_EQ(memory.size(), HostMem::kAlignment); + memory.resize(2 * HostMem::kAlignment); + EXPECT_EQ(memory.size(), 2 * HostMem::kAlignment); +} + +INSTANTIATE_TEST_SUITE_P(ThpModes, HostMemPageModeTest, testing::Values("0", "1")); + +TEST(KvCacheManagerV2HostMemTest, PrefaultedAllocationSupportsGpuRoundTrip) +{ + ScopedEnv thp("TLLM_KV_CACHE_MANAGER_V2_THP", "1"); + ScopedEnv prefault("TLLM_KV_CACHE_MANAGER_V2_PREFAULT_THREADS", "2"); + ASSERT_EQ(cudaSetDevice(0), cudaSuccess); + + constexpr size_t kSize = 4 << 20; + HostMem memory(kSize); + std::memset(reinterpret_cast(memory.address()), 0x5A, kSize); + + void* devicePtr = nullptr; + ASSERT_EQ(cudaMalloc(&devicePtr, kSize), cudaSuccess); + ASSERT_EQ( + cudaMemcpy(devicePtr, reinterpret_cast(memory.address()), kSize, cudaMemcpyHostToDevice), cudaSuccess); + std::memset(reinterpret_cast(memory.address()), 0, kSize); + ASSERT_EQ( + cudaMemcpy(reinterpret_cast(memory.address()), devicePtr, kSize, cudaMemcpyDeviceToHost), cudaSuccess); + + auto const* bytes = reinterpret_cast(memory.address()); + EXPECT_TRUE(std::all_of(bytes, bytes + kSize, [](unsigned char value) { return value == 0x5A; })); + EXPECT_EQ(cudaFree(devicePtr), cudaSuccess); +} + +} // namespace diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp new file mode 100644 index 000000000000..cd91444e6028 --- /dev/null +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp @@ -0,0 +1,371 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h" +#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h" +#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h" +#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h" +#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/pendingStats.h" +#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/stats.h" +#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h" + +#include +#include + +#include +#include + +namespace +{ + +using namespace tensorrt_llm::batch_manager::kv_cache_manager_v2; + +KVCacheManagerConfig makeConfig(bool enableStats = true) +{ + KVCacheManagerConfig config; + config.tokensPerBlock = 4; + config.cacheTiers.emplace_back(GpuCacheTierConfig{4 << 20}); + AttentionLayerConfig layer; + layer.layerId = 0; + layer.buffers.push_back(BufferConfig{"key", 4096, std::nullopt}); + config.layers.emplace_back(std::move(layer)); + config.enableStats = enableStats; + return config; +} + +KVCacheManagerConfig makeTieredConfig() +{ + KVCacheManagerConfig config; + config.tokensPerBlock = 4; + config.cacheTiers.emplace_back(GpuCacheTierConfig{4 << 20}); + config.cacheTiers.emplace_back(HostCacheTierConfig{4 << 20}); + AttentionLayerConfig layer; + layer.layerId = 0; + layer.buffers.push_back(BufferConfig{"key", 2 << 20, std::nullopt}); + config.layers.emplace_back(std::move(layer)); + return config; +} + +TEST(KvCacheManagerV2StatsTest, StatsDeltaArithmetic) +{ + KVCacheStatsDelta stats{4, 3, 2, 1}; + KVCacheStatsDelta const delta{1, 2, 3, 4}; + stats.add(delta); + EXPECT_EQ(stats.allocTotalBlocks, 5); + EXPECT_EQ(stats.allocNewBlocks, 5); + EXPECT_EQ(stats.reusedBlocks, 5); + EXPECT_EQ(stats.missedBlocks, 5); + + KVCacheStatsDelta const copy = stats.copy(); + stats.subtract(delta); + EXPECT_EQ(stats.allocTotalBlocks, 4); + EXPECT_EQ(copy.allocTotalBlocks, 5); + stats.clear(); + EXPECT_TRUE(stats.empty()); +} + +TEST(KvCacheManagerV2StatsTest, IterationStatsDeltaArithmeticAndHitRate) +{ + KVCacheIterationStatsDelta stats; + stats.iterReusedBlocks = 3; + stats.iterFullReusedBlocks = 2; + stats.iterPartialReusedBlocks = 1; + stats.iterMissedBlocks = 1; + stats.iterOnboardBytes = 1024; + EXPECT_DOUBLE_EQ(stats.iterCacheHitRate(), 0.75); + + KVCacheIterationStatsDelta delta = stats.copy(); + stats.add(delta); + EXPECT_EQ(stats.iterReusedBlocks, 6); + EXPECT_EQ(stats.iterOnboardBytes, 2048); + stats.subtract(delta); + EXPECT_EQ(stats.iterReusedBlocks, 3); + stats.clear(); + EXPECT_TRUE(stats.empty()); + EXPECT_DOUBLE_EQ(stats.iterCacheHitRate(), 0.0); +} + +TEST(KvCacheManagerV2StatsTest, PendingAllocationRangesAreReversibleAndScoped) +{ + PendingStats pending; + EXPECT_TRUE(pending.recordAllocationRange( + LifeCycleId{0}, BlockOrdinal{0}, BlockOrdinal{3}, /*beamWidth=*/2, /*countAsMissed=*/true)); + EXPECT_TRUE(pending.recordAllocationRange(LifeCycleId{1}, BlockOrdinal{3}, BlockOrdinal{5}, + /*beamWidth=*/1, /*countAsMissed=*/false, /*countAsGeneration=*/true)); + + EXPECT_EQ(pending.globalStats().allocTotalBlocks, 8); + EXPECT_EQ(pending.globalStats().allocNewBlocks, 8); + EXPECT_EQ(pending.globalStats().missedBlocks, 6); + EXPECT_EQ(pending.requestStats().allocTotalBlocks, 8); + ASSERT_EQ(pending.iterationStatsByLifeCycle().size(), 2); + EXPECT_EQ(pending.iterationStatsByLifeCycle().at(LifeCycleId{0}).iterMissedBlocks, 6); + EXPECT_EQ(pending.iterationStatsByLifeCycle().at(LifeCycleId{1}).iterGenAllocBlocks, 2); + + EXPECT_TRUE(pending.subtractAllocationRange(BlockOrdinal{2}, BlockOrdinal{5})); + EXPECT_EQ(pending.globalStats().allocTotalBlocks, 4); + EXPECT_EQ(pending.globalStats().missedBlocks, 4); + ASSERT_EQ(pending.iterationStatsByLifeCycle().size(), 1); + EXPECT_EQ(pending.iterationStatsByLifeCycle().at(LifeCycleId{0}).iterAllocTotalBlocks, 4); + + EXPECT_TRUE(pending.subtractAllocationRange(BlockOrdinal{0}, BlockOrdinal{2})); + EXPECT_TRUE(pending.empty()); +} + +TEST(KvCacheManagerV2StatsTest, PendingReuseSurvivesAllocationRollbackUntilClear) +{ + PendingStats pending; + EXPECT_TRUE(pending.recordAllocationRange( + LifeCycleId{0}, BlockOrdinal{0}, BlockOrdinal{1}, /*beamWidth=*/1, /*countAsMissed=*/true)); + EXPECT_TRUE(pending.recordReuse(LifeCycleId{0}, /*fullReusedBlocks=*/2, /*partialReusedBlocks=*/1)); + + EXPECT_TRUE(pending.subtractAllocationRange(BlockOrdinal{0}, BlockOrdinal{1})); + EXPECT_EQ(pending.globalStats().allocTotalBlocks, 0); + EXPECT_EQ(pending.globalStats().reusedBlocks, 3); + auto const& iteration = pending.iterationStatsByLifeCycle().at(LifeCycleId{0}); + EXPECT_EQ(iteration.iterReusedBlocks, 3); + EXPECT_EQ(iteration.iterFullReusedBlocks, 2); + EXPECT_EQ(iteration.iterPartialReusedBlocks, 1); + + pending.clear(); + EXPECT_TRUE(pending.empty()); +} + +TEST(KvCacheManagerV2StatsTest, ManagerCommitResetAndRequestIdTracking) +{ + ASSERT_EQ(cudaSetDevice(0), cudaSuccess); + auto manager = std::make_shared(makeConfig()); + + KVCacheStatsDelta globalStats{4, 3, 2, 1}; + KVCacheIterationStatsDelta iterationStats; + iterationStats.iterAllocTotalBlocks = 4; + iterationStats.iterReusedBlocks = 2; + manager->commitStats(globalStats, {{LifeCycleId{0}, iterationStats}}); + + EXPECT_EQ(manager->getCommittedStats().allocTotalBlocks, 4); + auto firstIteration = manager->getAndResetIterationStats(); + ASSERT_EQ(firstIteration.size(), 1); + EXPECT_EQ(firstIteration.at(LifeCycleId{0}).iterReusedBlocks, 2); + EXPECT_TRUE(manager->getAndResetIterationStats().empty()); + + manager->markStatsDirty(11); + manager->markStatsDirty(std::nullopt); + EXPECT_EQ(manager->getDirtyStatsKvCacheIds().count(11), 1); + manager->markStatsExcluded(11); + EXPECT_TRUE(manager->isStatsExcluded(11)); + EXPECT_TRUE(manager->getDirtyStatsKvCacheIds().empty()); + manager->clearStatsExcluded(11); + EXPECT_FALSE(manager->isStatsExcluded(11)); + + auto cache = manager->createKvCache({}, {}, 17, {}, 8); + manager->markStatsDirty(17); + EXPECT_TRUE(cache->commitPendingStats().empty()); + EXPECT_TRUE(manager->getDirtyStatsKvCacheIds().empty()); + cache->close(); + + RequestIdType const cudaGraphDummyRequestId = std::numeric_limits::max(); + auto dummyCache = manager->createKvCache({}, {}, cudaGraphDummyRequestId); + ASSERT_TRUE(dummyCache->id.has_value()); + EXPECT_EQ(*dummyCache->id, cudaGraphDummyRequestId); + manager->markStatsDirty(cudaGraphDummyRequestId); + EXPECT_EQ(manager->getDirtyStatsKvCacheIds(), std::unordered_set{cudaGraphDummyRequestId}); + manager->markStatsExcluded(cudaGraphDummyRequestId); + EXPECT_TRUE(manager->isStatsExcluded(cudaGraphDummyRequestId)); + EXPECT_TRUE(manager->getDirtyStatsKvCacheIds().empty()); + dummyCache->close(); +} + +TEST(KvCacheManagerV2StatsTest, DisabledStatsSuppressManagerCommit) +{ + ASSERT_EQ(cudaSetDevice(0), cudaSuccess); + auto manager = std::make_shared(makeConfig(false)); + manager->commitStats(KVCacheStatsDelta{4, 3, 2, 1}); + EXPECT_TRUE(manager->getCommittedStats().empty()); + EXPECT_TRUE(manager->getAndResetIterationStats().empty()); +} + +TEST(KvCacheManagerV2StatsTest, PeakBlockStatsResetStartsNextIntervalFromCurrentSnapshot) +{ + ASSERT_EQ(cudaSetDevice(0), cudaSuccess); + auto manager = std::make_shared(makeTieredConfig()); + auto& storage = manager->storage(); + LifeCycleId const lifeCycle{0}; + + TypedVec twoSlots(LifeCycleId{1}, 2); + auto gpuSlots = storage.newGpuSlots(twoSlots); + manager->commitStats({}); + + RootBlock& root = manager->radixTree().addOrGetExisting({}); + std::vector> pages; + NodeBase* previous = &root; + int token = 0; + for (auto& slot : gpuSlots[lifeCycle]) + { + std::vector tokens; + for (int i = 0; i < manager->tokensPerBlock(); ++i) + { + tokens.emplace_back(TokenId{token++}); + } + auto block = addOrGetExistingBlock(previous, LifeCycleId{1}, std::move(tokens)); + auto page = makeShared(&storage, block, lifeCycle, kGpuLevel, kPriorityDefault); + page->setSlot(slot); + block->storage[lifeCycle] = page.get(); + storage.scheduleForEviction(*page); + pages.push_back(page); + previous = block.get(); + } + manager->commitStats({}); + + TypedVec oneSlot(LifeCycleId{1}, 1); + auto hostSlots = storage.newSlots(CacheLevel{1}, oneSlot); + manager->commitStats({}); + storage.releaseSlot(lifeCycle, CacheLevel{1}, std::move(hostSlots[lifeCycle].front())); + manager->clearReusableBlocks(); + pages.clear(); + + auto primaryPeak = manager->getAndResetIterationPeakBlockStats(kGpuLevel); + auto secondaryPeak = manager->getAndResetIterationPeakBlockStats(CacheLevel{1}); + ASSERT_EQ(primaryPeak.size(), PoolGroupIndex{1}); + ASSERT_EQ(secondaryPeak.size(), PoolGroupIndex{1}); + EXPECT_EQ(primaryPeak[PoolGroupIndex{0}].available, 2); + EXPECT_EQ(primaryPeak[PoolGroupIndex{0}].unavailable, 2); + EXPECT_EQ(primaryPeak[PoolGroupIndex{0}].evictable, 2); + EXPECT_EQ(secondaryPeak[PoolGroupIndex{0}].available, 2); + EXPECT_EQ(secondaryPeak[PoolGroupIndex{0}].unavailable, 1); + EXPECT_EQ(secondaryPeak[PoolGroupIndex{0}].evictable, 0); + + primaryPeak = manager->getAndResetIterationPeakBlockStats(kGpuLevel); + secondaryPeak = manager->getAndResetIterationPeakBlockStats(CacheLevel{1}); + EXPECT_EQ(primaryPeak[PoolGroupIndex{0}].available, 2); + EXPECT_EQ(primaryPeak[PoolGroupIndex{0}].unavailable, 0); + EXPECT_EQ(primaryPeak[PoolGroupIndex{0}].evictable, 0); + EXPECT_EQ(secondaryPeak[PoolGroupIndex{0}].available, 2); + EXPECT_EQ(secondaryPeak[PoolGroupIndex{0}].unavailable, 0); + EXPECT_EQ(secondaryPeak[PoolGroupIndex{0}].evictable, 0); + + auto nextIntervalSlots = storage.newSlots(kGpuLevel, oneSlot); + manager->commitStats({}); + storage.releaseSlot(lifeCycle, kGpuLevel, std::move(nextIntervalSlots[lifeCycle].front())); + primaryPeak = manager->getAndResetIterationPeakBlockStats(kGpuLevel); + EXPECT_EQ(primaryPeak[PoolGroupIndex{0}].available, 2); + EXPECT_EQ(primaryPeak[PoolGroupIndex{0}].unavailable, 1); + EXPECT_EQ(primaryPeak[PoolGroupIndex{0}].evictable, 0); +} + +TEST(KvCacheManagerV2StatsTest, MigrationAndLastTierDropRecordersReceiveExactPages) +{ + ASSERT_EQ(cudaSetDevice(0), cudaSuccess); + auto manager = std::make_shared(makeTieredConfig()); + auto& storage = manager->storage(); + LifeCycleId const lifeCycle{0}; + ASSERT_EQ(storage.getStatistics(kGpuLevel).total, 2); + ASSERT_EQ(storage.getStatistics(CacheLevel{1}).total, 2); + + int offloaded = 0; + int onboarded = 0; + int dropped = 0; + MigrationRecorder const migrationRecorder + = [&](std::vector> const& pages, std::vector const& slots, CacheLevel srcLevel, + CacheLevel dstLevel) + { + EXPECT_EQ(pages.size(), slots.size()); + if (srcLevel == kGpuLevel && dstLevel == CacheLevel{1}) + { + offloaded += static_cast(pages.size()); + } + else if (srcLevel == CacheLevel{1} && dstLevel == kGpuLevel) + { + onboarded += static_cast(pages.size()); + } + }; + DropRecorder const dropRecorder = [&](std::vector> const& pages, CacheLevel level) + { + EXPECT_EQ(level, CacheLevel{1}); + dropped += static_cast(pages.size()); + }; + + RootBlock& root = manager->radixTree().addOrGetExisting({}); + int tokenBase = 0; + auto makeCommittedPages = [&](std::vector slots) + { + std::vector> pages; + NodeBase* previous = &root; + for (auto& slot : slots) + { + std::vector tokens; + for (int i = 0; i < manager->tokensPerBlock(); ++i) + { + tokens.emplace_back(TokenId{tokenBase++}); + } + auto block = addOrGetExistingBlock(previous, LifeCycleId{1}, std::move(tokens)); + auto page = makeShared(&storage, block, lifeCycle, kGpuLevel, kPriorityDefault); + page->setSlot(slot); + block->storage[lifeCycle] = page.get(); + storage.scheduleForEviction(*page); + pages.push_back(page); + previous = block.get(); + } + return pages; + }; + + TypedVec twoSlots(LifeCycleId{1}, 2); + auto initialSlots = storage.newGpuSlots(twoSlots); + auto firstPages = makeCommittedPages(std::move(initialSlots[lifeCycle])); + + auto temporarySlots = storage.newGpuSlots(twoSlots, migrationRecorder, dropRecorder); + EXPECT_EQ(offloaded, 2); + EXPECT_EQ(onboarded, 0); + EXPECT_EQ(dropped, 0); + for (auto& slot : temporarySlots[lifeCycle]) + { + storage.releaseSlot(lifeCycle, kGpuLevel, std::move(slot)); + } + + auto cache = manager->createKvCache(); + std::vector targets; + for (BlockOrdinal ordinal{0}; ordinal < BlockOrdinal{2}; ++ordinal) + { + auto const& page = firstPages[toSizeT(ordinal)]; + ASSERT_TRUE(page->scheduledForEviction()); + storage.excludeFromEviction(*page); + targets.push_back({page, kDefaultBeamIndex, ordinal, lifeCycle}); + } + storage.batchedMigrateToGpu(targets, *cache, migrationRecorder); + EXPECT_EQ(onboarded, 2); + for (auto const& page : firstPages) + { + storage.scheduleForEviction(*page); + } + + temporarySlots = storage.newGpuSlots(twoSlots, migrationRecorder, dropRecorder); + EXPECT_EQ(offloaded, 4); + auto secondPages = makeCommittedPages(std::move(temporarySlots[lifeCycle])); + (void) secondPages; + firstPages.clear(); + targets.clear(); + + auto finalSlots = storage.newGpuSlots(twoSlots, migrationRecorder, dropRecorder); + EXPECT_EQ(offloaded, 6); + EXPECT_EQ(onboarded, 2); + EXPECT_EQ(dropped, 2); + for (auto& slot : finalSlots[lifeCycle]) + { + storage.releaseSlot(lifeCycle, kGpuLevel, std::move(slot)); + } + cache->close(); +} + +} // namespace diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2TypedIndexTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2TypedIndexTest.cpp new file mode 100644 index 000000000000..ac94dc063e72 --- /dev/null +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2TypedIndexTest.cpp @@ -0,0 +1,333 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h" +#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h" +#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.h" +#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.h" +#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h" +#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/math.h" +#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/typedIndex.h" + +#include + +#include +#include +#include +#include + +namespace +{ + +using namespace tensorrt_llm::batch_manager::kv_cache_manager_v2; + +template +struct HasIndexOperator : std::false_type +{ +}; + +template +struct HasIndexOperator()[std::declval()])>> + : std::true_type +{ +}; + +template +struct HasContains : std::false_type +{ +}; + +template +struct HasContains().contains(std::declval()))>> + : std::true_type +{ +}; + +template +struct HasEqual : std::false_type +{ +}; + +template +struct HasEqual() == std::declval())>> : std::true_type +{ +}; + +template +struct HasNotEqual : std::false_type +{ +}; + +template +struct HasNotEqual() != std::declval())>> : std::true_type +{ +}; + +template +struct HasLess : std::false_type +{ +}; + +template +struct HasLess() < std::declval())>> : std::true_type +{ +}; + +template +struct HasGreater : std::false_type +{ +}; + +template +struct HasGreater() > std::declval())>> : std::true_type +{ +}; + +template +struct HasLessEqual : std::false_type +{ +}; + +template +struct HasLessEqual() <= std::declval())>> : std::true_type +{ +}; + +template +struct HasGreaterEqual : std::false_type +{ +}; + +template +struct HasGreaterEqual() >= std::declval())>> : std::true_type +{ +}; + +TEST(KvCacheManagerV2TypedIndexTest, StrongIndexSupportsIntegerArithmetic) +{ + LifeCycleId lc{3}; + + EXPECT_EQ((lc + 2).value(), 5); + EXPECT_EQ((2 + lc).value(), 5); + EXPECT_EQ((lc - 2).value(), 1); + EXPECT_EQ((lc + std::size_t{2}).value(), 5); + EXPECT_EQ((std::int64_t{2} + lc).value(), 5); + EXPECT_EQ((lc - std::int64_t{2}).value(), 1); + + LifeCycleId begin{4}; + LifeCycleId end{10}; + static_assert(std::is_same::value, "index difference must be an integer"); + EXPECT_EQ(end - begin, 6); + + lc += 4; + EXPECT_EQ(lc.value(), 7); + + lc -= 2; + EXPECT_EQ(lc.value(), 5); + + EXPECT_EQ((++lc).value(), 6); + EXPECT_EQ((lc++).value(), 6); + EXPECT_EQ(lc.value(), 7); + EXPECT_EQ((--lc).value(), 6); + EXPECT_EQ((lc--).value(), 6); + EXPECT_EQ(lc.value(), 5); +} + +TEST(KvCacheManagerV2TypedIndexTest, StrongIndexSupportsValueTypeUpperBoundComparison) +{ + static_assert(HasEqual::value, "matching strong index equality must work"); + static_assert(HasNotEqual::value, "matching strong index inequality must work"); + static_assert(HasLess::value, "matching strong index ordering must work"); + static_assert(HasGreater::value, "matching strong index ordering must work"); + static_assert(HasLessEqual::value, "matching strong index ordering must work"); + static_assert(HasGreaterEqual::value, "matching strong index ordering must work"); + + static_assert(HasLess::value, + "strong index must support upper-bound comparison against its value type"); + static_assert(HasLess::value, "slot id must support upper-bound comparison against slot counts"); + static_assert(HasGreaterEqual::value, + "strong index must support upper-bound rejection against its value type"); + static_assert( + HasGreaterEqual::value, "slot id must support upper-bound rejection against slot counts"); + + static_assert( + !HasEqual::value, "strong index must not compare equal to raw value type"); + static_assert( + !HasEqual::value, "raw value type must not compare equal to strong index"); + static_assert(!HasNotEqual::value, + "strong index must not compare unequal to raw value type"); + static_assert(!HasNotEqual::value, + "raw value type must not compare unequal to strong index"); + static_assert( + !HasLess::value, "raw value type must not order against strong index"); + static_assert(!HasGreater::value, + "strong index must not use raw value type for greater-than comparison"); + static_assert( + !HasGreater::value, "raw value type must not order against strong index"); + static_assert(!HasLessEqual::value, + "strong index must not use raw value type for less-equal comparison"); + static_assert(!HasLessEqual::value, + "raw value type must not order against strong index"); + static_assert(!HasGreaterEqual::value, + "raw value type must not order against strong index"); +} + +TEST(KvCacheManagerV2TypedIndexTest, StrongIndexDefaultsMatchSentinels) +{ + static_assert(CacheLevel{}.value() == kGpuLevel.value(), "CacheLevel default should name the GPU level"); + static_assert(BeamIndex{}.value() == kDefaultBeamIndex.value(), "BeamIndex default should name the default beam"); + static_assert(BlockOrdinal{}.value() == kBadBlockOrdinal.value(), "BlockOrdinal default should be invalid"); + static_assert(PageIndex{}.value() == kBadPageIndex.value(), "PageIndex default should be invalid"); +} + +TEST(KvCacheManagerV2TypedIndexTest, SlotCapacityAccessorsUsePlainSlotCount) +{ + static_assert(std::is_same().numSlots()), SlotCount>::value, + "SlotAllocator::numSlots must return a plain slot count"); + static_assert(std::is_same().numSlots()), SlotCount>::value, + "PoolGroupBase::numSlots must return a plain slot count"); + static_assert( + std::is_same().numSlots(PoolGroupIndex{0})), SlotCount>::value, + "CacheLevelStorage::numSlots must return a plain slot count"); + static_assert( + std::is_same().numSlots(PoolGroupIndex{0})), SlotCount>::value, + "StorageManager::numSlots must return a plain slot count"); +} + +TEST(KvCacheManagerV2TypedIndexTest, SlotIdUsesInt64) +{ + static_assert(std::is_same::value, "SlotId must use int64_t"); + static_assert(std::is_same::value, "SlotCount must use int64_t"); +} + +TEST(KvCacheManagerV2TypedIndexTest, SlotCountsUsePlainSlotCount) +{ + static_assert(std::is_same().numFreeSlots()), SlotCount>::value, + "SlotAllocator::numFreeSlots must return a plain slot count"); + static_assert(std::is_same().numOccupiedSlots()), SlotCount>::value, + "SlotAllocator::numOccupiedSlots must return a plain slot count"); + static_assert(std::is_same().numOverflowSlots()), SlotCount>::value, + "SlotAllocator::numOverflowSlots must return a plain slot count"); + static_assert(std::is_same().numFreeSlots()), SlotCount>::value, + "PoolGroupBase::numFreeSlots must return a plain slot count"); + static_assert(std::is_same().numFreeSlots(PoolGroupIndex{0})), + SlotCount>::value, + "CacheLevelStorage::numFreeSlots must return a plain slot count"); + static_assert(std::is_same().size()), SlotCount>::value, + "PrioritizedEvictionPolicy::size must return a plain count"); + static_assert(std::is_same().total), SlotCount>::value, + "StorageStatistics::total must be a plain count"); + static_assert(std::is_same().available()), SlotCount>::value, + "StorageStatistics::available must return a plain count"); + static_assert(std::is_same().unavailable()), SlotCount>::value, + "StorageStatistics::unavailable must return a plain count"); +} + +TEST(KvCacheManagerV2TypedIndexTest, AllocationRequestsUsePlainSlotCount) +{ + using SlotAllocatorAllocateMultiple = std::vector (SlotAllocator::*)(SlotCount); + static_assert(std::is_same::value, + "SlotAllocator::allocateMultiple must take a plain slot count"); + static_assert(!std::is_invocable::value, + "SlotAllocator::allocateMultiple must not take a slot id"); + + using StorageManagerNewGpuSlots = TypedVec> (StorageManager::*)( + TypedVec const&, MigrationRecorder const&, DropRecorder const&); + static_assert(std::is_same::value, + "StorageManager::newGpuSlots must take plain slot counts"); + + using StorageManagerNewSlotsForPoolGroup = std::vector (StorageManager::*)( + CacheLevel, PoolGroupIndex, SlotCount, MigrationRecorder const&, DropRecorder const&); + static_assert( + std::is_same::value, + "StorageManager::newSlotsForPoolGroup must take a plain slot count"); + static_assert( + std::is_same().numEvictablePages(PoolGroupIndex{0})), + SlotCount>::value, + "PerLevelEvictionController::numEvictablePages must return a plain count"); +} + +TEST(KvCacheManagerV2TypedIndexTest, SlotAllocatorRejectsOutOfRangeSlotId) +{ + SlotAllocator allocator(1); + Slot invalidSlot; + invalidSlot.setSlotId(SlotId{1}); + + EXPECT_THROW(allocator.release(std::move(invalidSlot)), LogicError); +} + +TEST(KvCacheManagerV2TypedIndexTest, HalfOpenRangeCarriesIndexType) +{ + using BlockRange = HalfOpenRange; + + static_assert(std::is_same::value, + "range endpoints must carry the requested index type"); + static_assert(std::is_same().length()), int>::value, + "range length must remain an integer count"); + static_assert(HasContains::value, "matching index type must work"); + static_assert(!HasContains::value, "different strong index must not work"); + static_assert(!HasContains::value, "raw int must not be accepted by typed contains"); + + BlockRange empty; + EXPECT_EQ(empty.beg, BlockOrdinal{0}); + EXPECT_EQ(empty.end, BlockOrdinal{0}); + EXPECT_EQ(empty.length(), 0); + EXPECT_FALSE(empty); + + BlockRange range{2, 5}; + EXPECT_EQ(range.beg, BlockOrdinal{2}); + EXPECT_EQ(range.end, BlockOrdinal{5}); + EXPECT_EQ(range.length(), 3); + EXPECT_TRUE(range.contains(BlockOrdinal{2})); + EXPECT_FALSE(range.contains(BlockOrdinal{5})); + + BlockRange overlap = intersect(range, BlockRange{4, 8}); + EXPECT_EQ(overlap.beg, BlockOrdinal{4}); + EXPECT_EQ(overlap.end, BlockOrdinal{5}); +} + +TEST(KvCacheManagerV2TypedIndexTest, TypedVecRequiresMatchingIndexType) +{ + using LifeCycleVector = TypedVec; + + static_assert(HasIndexOperator::value, "matching index type must work"); + static_assert(!HasIndexOperator::value, "different strong index must not work"); + static_assert(!HasIndexOperator::value, "raw int must not index typed vector"); + + LifeCycleVector values(LifeCycleId{2}, 0); + values[LifeCycleId{0}] = 11; + values[LifeCycleId{1}] = 17; + + EXPECT_EQ(values.size(), LifeCycleId{2}); + EXPECT_EQ(values[LifeCycleId{0}], 11); + EXPECT_EQ(values.at(LifeCycleId{1}), 17); +} + +TEST(KvCacheManagerV2TypedIndexTest, TypedVecSupportsVectorInterop) +{ + TypedVec values; + values.reserve(PoolGroupIndex{2}); + values.push_back(3); + values.push_back(5); + + EXPECT_EQ(values.stdSize(), 2); + EXPECT_EQ(values.size(), PoolGroupIndex{2}); + EXPECT_EQ(values.raw().at(0), 3); + EXPECT_EQ(values[PoolGroupIndex{1}], 5); +} + +} // namespace diff --git a/docs/source/features/kvcache.md b/docs/source/features/kvcache.md index 82f7fbe58143..b8b395abdc4d 100644 --- a/docs/source/features/kvcache.md +++ b/docs/source/features/kvcache.md @@ -123,6 +123,8 @@ KV cache salting provides a security mechanism to control which requests can reu To use cache salting, specify the `cache_salt` parameter as a string when creating requests. Only requests with matching cache salt values can share cached KV blocks. The salt value can be any non-empty string, such as a user ID, tenant ID, or hash string. +This isolation is enforced entirely by the block-key hash: the salt is mixed into the hashed input and prefix matching is decided by digest equality alone (blocks are not re-compared token-by-token). The block-key hash is therefore required to be a cryptographic hash with strong collision resistance and a 256-bit digest (SHA-256 provides ~128-bit collision resistance, which is ample here); substituting a non-cryptographic hash would allow crafted collisions to bypass salt isolation and must not be done. + ### Multimodal UUID Support for Cache Identification When working with multimodal models (e.g., vision-language models), the KV cache system needs to identify which cached blocks correspond to which multimodal inputs (images, videos, etc.). By default, the system uses content-based hashing to generate unique identifiers for each multimodal input. However, this approach has limitations for cache management across sessions, as the same content must be re-processed to generate the same hash. diff --git a/jenkins/license_cpp.json b/jenkins/license_cpp.json index bba247b97f50..408a1217408c 100644 --- a/jenkins/license_cpp.json +++ b/jenkins/license_cpp.json @@ -44,6 +44,12 @@ "tensorrt_llm/kernels/causalConv1d/causalConv1d.h": "dual license", "tensorrt_llm/thop/causalConv1dOp.cpp": "dual license", "tensorrt_llm/common/vec_dtypes.cuh": "dual license", + "tensorrt_llm/common/sha256/attributes.h": "external (Bitcoin Core MIT)", + "tensorrt_llm/common/sha256/sha256.cpp": "external (Bitcoin Core MIT)", + "tensorrt_llm/common/sha256/sha256.h": "external (Bitcoin Core MIT)", + "tensorrt_llm/common/sha256/sha256_arm_shani.cpp": "external (Bitcoin Core MIT)", + "tensorrt_llm/common/sha256/sha256_endian.h": "external (Bitcoin Core MIT)", + "tensorrt_llm/common/sha256/sha256_x86_shani.cpp": "external (Bitcoin Core MIT)", "_": "don't remove, for trailing comma" } } diff --git a/scripts/build_wheel.py b/scripts/build_wheel.py index 462b74d8e511..a38a1ddb686e 100755 --- a/scripts/build_wheel.py +++ b/scripts/build_wheel.py @@ -404,6 +404,8 @@ def generate_python_stubs_linux(venv_python: Path, deep_ep: bool, binding_lib_name: str): build_run(f"\"{venv_python}\" -m pip install nanobind") build_run(f"\"{venv_python}\" -m pip install pybind11-stubgen") + nanobind_stubgen_patterns = get_project_dir( + ) / "scripts" / "nanobind_stubgen.patterns" env_stub_gen = os.environ.copy() cuda_home_dir = env_stub_gen.get("CUDA_HOME") or env_stub_gen.get( @@ -421,8 +423,10 @@ def generate_python_stubs_linux(venv_python: Path, deep_ep: bool, link_dir = None try: - build_run(f"\"{venv_python}\" -m nanobind.stubgen -m bindings -r -O .", - env=env_stub_gen) + build_run( + f"\"{venv_python}\" -m nanobind.stubgen -m bindings -r -O . " + f"-p \"{nanobind_stubgen_patterns}\" -q", + env=env_stub_gen) # Pre-import torch so deep_gemm_cpp_tllm's FP4 scalar-type registration # succeeds; CLI args after `-c ...` land in sys.argv[1:] for argparse. build_run( @@ -767,14 +771,26 @@ def main(*, if cache_dir.exists(): clear_folder(cache_dir) - install_file = copy + def safe_copy(src, dst): + """Copy a file, replacing a destination symlink with a real file.""" + src_path = Path(src) + dst_path = Path(dst) + if dst_path.is_dir(): + dst_path = dst_path / src_path.name + if dst_path.is_symlink(): + dst_path.unlink() + return copy(src_path, dst_path) + + install_file = safe_copy # Wrapper for copytree that checks if source and destination are the same def safe_copytree(src, dst, dirs_exist_ok=True): """Copy tree, but skip if source and destination resolve to the same directory.""" src_path = Path(src).resolve() - dst_path = Path(dst).resolve() - if src_path == dst_path: + dst_path = Path(dst) + if dst_path.is_symlink(): + dst_path.unlink() + elif src_path == dst_path.resolve(): # Source and destination are the same, skip copying return if dst_path.exists() and dirs_exist_ok: @@ -789,7 +805,7 @@ def symlink_remove_dst(src, dst): dst = os.path.abspath(dst) if os.path.isdir(dst): dst = os.path.join(dst, os.path.basename(src)) - if os.path.exists(dst): + if os.path.lexists(dst): os.remove(dst) os.symlink(src, dst) @@ -798,7 +814,7 @@ def symlink_remove_dst(src, dst): def symlink_remove_dst_tree(src, dst, dirs_exist_ok=True): src = os.path.abspath(src) dst = os.path.abspath(dst) - if dirs_exist_ok and os.path.exists(dst): + if dirs_exist_ok and os.path.lexists(dst): os.remove(dst) os.symlink(src, dst) diff --git a/scripts/nanobind_stubgen.patterns b/scripts/nanobind_stubgen.patterns new file mode 100644 index 000000000000..278dc4b92207 --- /dev/null +++ b/scripts/nanobind_stubgen.patterns @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +^bindings\.internal\.batch_manager\.kv_cache_manager_v2\.CachedCudaEvent\.NULL$: + \from typing import ClassVar + NULL: ClassVar[CachedCudaEvent] diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 518c7b711162..4b7bde587dbf 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -60,13 +60,13 @@ KVCacheEventManager, KVCacheIterationStatsDelta, LayerId, - LifeCycleId, PageIndexMode, PlannedDropHandle, PoolGroupPeakBlockStats, ReuseScope, SwaScratchReuseConfig, TokenIdExt, + _introspection, _KVCache, exact_div, gen_multimodal_cache_key_tokens, @@ -1217,24 +1217,39 @@ def _build_pool_mapping_tensors(self): for pool_id in range(self.num_pools): layer_id = self.impl.layer_grouping[pool_id][0] role_a, _ = self._get_pool_roles(pool_id) - kv_cache_pool_pointers_list.append( - [ - self.impl.get_mem_pool_base_address(layer_id, role_a, PageIndexMode.SHARED), - 0, - ] + key_base_addr = self.impl.get_mem_pool_base_address( + layer_id, role_a, PageIndexMode.SHARED ) + kv_cache_pool_pointers_list.append([key_base_addr, 0]) if self.dtype == DataType.NVFP4: + # The KEY/scale pointers are a (rep-layer, 0-offset) origin + # against which each layer's kv_cache_pool_mapping offset is + # resolved. The block-scale origin must reproduce the SAME + # per-layer offset() as KEY, so mirror the KEY base for the + # same representative layer and shift it back by that layer's + # offset. For the base manager offset(rep) == 0, so this is + # just the rep layer's scale base; for address-ranked + # subclasses (MiniMax-M3) offset(rep) may be non-zero, and the + # shift lands the origin on the pool's slot-0 scale address. + # This keeps block_scale_offset == offset without depending on + # the non-contractual layer_grouping order. block_scale_role = self._get_block_scale_role(role_a) - block_scale_pool_pointers_list.append( - [ + if block_scale_role is not None: + rep_offset = self._kv_pool_mapping_offset(layer_id, pool_id, key_base_addr) + scale_stride = ( + self.get_layer_bytes_per_token(layer_id, block_scale_role) + * self.kv_factor + * self.tokens_per_block + ) + scale_base_addr = ( self.impl.get_mem_pool_base_address( layer_id, block_scale_role, PageIndexMode.SHARED ) - if block_scale_role is not None - else 0, - 0, - ] - ) + - rep_offset * scale_stride + ) + else: + scale_base_addr = 0 + block_scale_pool_pointers_list.append([scale_base_addr, 0]) for layer_id in typed_range(LayerId(self.num_local_layers)): layer_group_id = self.impl.get_layer_group_id(layer_id) @@ -1472,14 +1487,19 @@ def get_event_window_size(layer_id: int) -> int: } def _format_kv_cache_pool_lifecycle_entry(self, layer_id: LayerId, role: DataRole) -> str: - attr = self.impl._storage.get_buffer_attr(layer_id, role) - pool_group_id = self.impl._storage.get_pool_group_index(attr.life_cycle_id) - lifecycle = self.impl._life_cycles.get_life_cycle(attr.life_cycle_id) - return ( - f"role={str(role)}, pool_group_id={int(pool_group_id)}, " - f"lifecycle_id={int(attr.life_cycle_id)}, " - f"lifecycle={lifecycle}" - ) + for pool_group in self.impl.pool_group_descs: + for variant in pool_group.slot_desc.variants: + for coalesced in variant.coalesced_buffers: + for buffer_id in coalesced.buffer_ids: + if int(buffer_id.layer_id) == int(layer_id) and str(buffer_id.role) == str( + role + ): + return ( + f"role={role!s}, " + f"pool_group_id={int(pool_group.pool_group_index)}, " + f"layer_group_id={int(variant.layer_group_id)}" + ) + return f"role={role!s}, pool_group_id=?, layer_group_id=?" def _log_kv_cache_pool_lifecycle_mapping(self) -> None: entries = OrderedDict() @@ -2000,15 +2020,23 @@ def get_index_k_buffer( layer_offset = self.layer_offsets[layer_idx] try: addr = self.impl.get_mem_pool_base_address(layer_offset, Role.INDEX_KEY) - page_stride = self.impl.get_page_stride(layer_offset, Role.INDEX_KEY) - page_upper = self.impl.get_page_index_upper_bound(layer_offset, Role.INDEX_KEY) - converter = self.impl.get_page_index_converter(layer_offset, Role.INDEX_KEY) - except KeyError: + except (KeyError, IndexError): # INDEX_KEY not registered for this layer (default V2 manager # registers only K/V/scale; sparse subclasses register # INDEX_KEY only on sparse layers via # ``_extra_buffers_per_layer``). + # + # The python backend raises ``KeyError`` from the missing dict + # lookup; the C++ backend's ``getBufferAttr`` throws + # ``std::out_of_range`` for an unknown buffer id, which nanobind + # maps to ``IndexError``. Catch both so the "role not registered + # -> None" contract holds on either backend. return None + # INDEX_KEY is registered; the remaining lookups must succeed. Keep them + # outside the try so genuine failures surface instead of returning None. + page_stride = self.impl.get_page_stride(layer_offset, Role.INDEX_KEY) + page_upper = self.impl.get_page_index_upper_bound(layer_offset, Role.INDEX_KEY) + converter = self.impl.get_page_index_converter(layer_offset, Role.INDEX_KEY) if isinstance(dtype, DataType): torch_dtype = binding_to_torch_dtype(dtype) @@ -2558,13 +2586,16 @@ def _stats_life_cycle_window_size(self, life_cycle) -> Optional[int]: return self._stats_window_size(life_cycle.window_size) def _get_storage_statistics(self, cache_level: CacheLevel): - return self.impl._storage.get_statistics(cache_level) + return _introspection.storage_statistics(self.impl, cache_level) def _stats_life_cycle_metadata(self) -> dict[int, tuple[int, Optional[int], str]]: - pool_groups_by_life_cycle = [ - self.impl._storage.get_pool_group_index(LifeCycleId(life_cycle_id)) - for life_cycle_id in range(len(self.impl.layer_grouping)) - ] + # life cycle (== layer group) -> pool group is static structure exposed by + # the public pool_group_descs API; no introspection needed. + pool_groups_by_life_cycle = { + int(variant.layer_group_id): int(pool_group.pool_group_index) + for pool_group in self.impl.pool_group_descs + for variant in pool_group.slot_desc.variants + } metadata: dict[int, tuple[int, Optional[int], str]] = {} for life_cycle_id, layer_ids in enumerate(self.impl.layer_grouping): diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index e7cab986c2ef..a670074dd23b 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -2909,6 +2909,25 @@ def _build_cache_config( typical_step = BatchDesc(request_descs * self._max_resident_sequences() + dummy_requests) + # The recurrent (SSM) state pool must hold one slot per resident + # sequence plus every reserved dummy slot. Unlike attention pages, a + # Mamba state is fixed-size per sequence, so this floor is independent + # of sequence length. The base config only emits constraints when + # ``avg_seq_len`` is set, and speculative decoding inflates the reserved + # dummy slots (CUDA-graph padding), so without an explicit floor the SSM + # pool can be undersized (see the live/dummy-slot check in _setup_states + # / __init__). Add a min-slots constraint of zero-capacity requests: + # these cost no attention pages but reserve one SSM slot each. + if any(isinstance(layer, SsmLayerConfig) for layer in layers): + ssm_floor_slots = (self._max_resident_sequences() + + self._num_reserved_dummy_slots) + constraints = [ + *constraints, + BatchDesc([ + KVCacheDesc(capacity=0, history_length=0) + for _ in range(ssm_floor_slots) + ]), + ] return replace( config, layers=layers, diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 34ebd340a5ee..2efc017e0e4c 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -42,7 +42,7 @@ from tensorrt_llm.llmapi.llm_args import PeftCacheConfig, WaitingQueuePolicy from tensorrt_llm.logger import logger from tensorrt_llm.mapping import CpType -from tensorrt_llm.runtime.kv_cache_manager_v2._exceptions import OutOfPagesError +from tensorrt_llm.runtime.kv_cache_manager_v2 import OutOfPagesError from tensorrt_llm.tools.layer_wise_benchmarks import get_calibrator from tensorrt_llm.tools.profiler.host_profile_tools.host_profiler import ( get_global_profiler, host_profiler_context) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py index e96c62b5501c..2d5b90468bdf 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py @@ -13,77 +13,280 @@ # See the License for the specific language governing permissions and # limitations under the License. -from . import rawref # noqa: F401 -from ._block_radix_tree import ReuseScope, gen_multimodal_cache_key_tokens # noqa: F401 -from ._common import ( # noqa: F401 - BAD_PAGE_INDEX, - CACHE_LEVEL1, - GPU_LEVEL, - NDEBUG, - CacheLevel, - CacheTier, - CudaStream, - LayerId, - MemAddress, - PageIndexMode, - PageStatus, - Priority, - SlidingWindowSize, - TokenId, - TokenIdExt, -) -from ._config import ( # noqa: F401 - AttentionLayerConfig, - BatchDesc, - BufferConfig, - CacheTierConfig, - DataRole, - DiskCacheTierConfig, - GpuCacheTierConfig, - HostCacheTierConfig, - KVCacheDesc, - KVCacheManagerConfig, - SsmLayerConfig, - SwaScratchReuseConfig, -) -from ._core import ( # noqa: F401 - DEFAULT_BEAM_INDEX, - AggregatedPageDesc, - BeamIndex, - ExpandedBuffer, - KVCacheManager, - PageIndexConverter, - PlannedDropHandle, - PoolDesc, - PoolGroupDesc, - PoolGroupPeakBlockStats, - ScratchDesc, - _KVCache, -) -from ._core._kv_cache import _Status as KvCacheStatus # noqa: F401 -from ._event_manager import ( # noqa: F401 - KVCacheCreatedData, - KVCacheEvent, - KVCacheEventDiff, - KVCacheEventManager, - KVCacheRemovedData, - KVCacheStoredBlockData, - KVCacheStoredData, - KVCacheUpdatedData, - UniqueToken, -) -from ._exceptions import CuError, OutOfMemoryError, OutOfPagesError # noqa: F401 -from ._life_cycle_registry import AttnLifeCycle, LayerGroupId, LifeCycleId # noqa: F401 -from ._stats import ( # noqa: F401 - _KV_CACHE_ITERATION_STATS_DELTA_FIELDS, - KVCacheIterationStatsDelta, - KVCacheStatsDelta, - SsmSnapshotIterationStatsDelta, -) -from ._storage import BufferId # noqa: F401 -from ._storage._config import CoalescedBuffer, SlotDesc, SlotDescVariant # noqa: F401 -from ._storage._core import PoolGroupIndex, PoolIndex # noqa: F401 -from ._utils import HalfOpenRange, exact_div, typed_range # noqa: F401 +import os +import sys +from importlib.util import find_spec +from pathlib import Path +from types import ModuleType +from typing import NamedTuple, Optional, Union + +_BACKEND = os.environ.get("TLLM_KV_CACHE_MANAGER_V2_BACKEND", "cpp").lower() + +if _BACKEND == "python": + from . import rawref # noqa: F401 + from ._block_radix_tree import ReuseScope # noqa: F401 + from ._cache_key import ( # noqa: F401 + gen_multimodal_cache_key_tokens, + sequence_to_blockchain_keys, + ) + from ._common import ( # noqa: F401 + BAD_PAGE_INDEX, + CACHE_LEVEL1, + GPU_LEVEL, + NDEBUG, + CacheLevel, + CacheTier, + CudaStream, + LayerId, + MemAddress, + PageIndexMode, + PageStatus, + Priority, + SlidingWindowSize, + TokenId, + TokenIdExt, + ) + from ._config import ( # noqa: F401 + AttentionLayerConfig, + BatchDesc, + BufferConfig, + CacheTierConfig, + DataRole, + DiskCacheTierConfig, + GpuCacheTierConfig, + HostCacheTierConfig, + KVCacheDesc, + KVCacheManagerConfig, + SsmLayerConfig, + SwaScratchReuseConfig, + ) + from ._core import ( # noqa: F401 + DEFAULT_BEAM_INDEX, + AggregatedPageDesc, + BeamIndex, + ExpandedBuffer, + KVCacheManager, + PageIndexConverter, + PlannedDropHandle, + PoolDesc, + PoolGroupDesc, + PoolGroupPeakBlockStats, + ScratchDesc, + _KVCache, + ) + from ._core._kv_cache import _Status as KvCacheStatus # noqa: F401 + from ._event_manager import ( # noqa: F401 + KVCacheCreatedData, + KVCacheEvent, + KVCacheEventDiff, + KVCacheEventManager, + KVCacheRemovedData, + KVCacheStoredBlockData, + KVCacheStoredData, + KVCacheUpdatedData, + UniqueToken, + ) + from ._exceptions import CuError, OutOfMemoryError, OutOfPagesError # noqa: F401 + from ._life_cycle_registry import AttnLifeCycle, LayerGroupId, LifeCycleId # noqa: F401 + from ._stats import ( # noqa: F401 + _KV_CACHE_ITERATION_STATS_DELTA_FIELDS, + KVCacheIterationStatsDelta, + KVCacheStatsDelta, + SsmSnapshotIterationStatsDelta, + ) + from ._storage import BufferId # noqa: F401 + from ._storage._config import CoalescedBuffer, SlotDesc, SlotDescVariant # noqa: F401 + from ._storage._core import PoolGroupIndex, PoolIndex # noqa: F401 + from ._utils import HalfOpenRange, exact_div, typed_range # noqa: F401 + + _cpp_introspection = None +else: + + class ReuseScope(NamedTuple): + lora_id: int | None = None + salt: int | None = None + + def to_bytes(self) -> bytes: + ret = sum((value is not None) << i for i, value in enumerate(self)).to_bytes( + 1, "little", signed=False + ) + for value in self: + if value is not None: + ret += value.to_bytes(8, "little", signed=False) + return ret + + def _load_cpp_module(): + if "tensorrt_llm" in sys.modules: + from tensorrt_llm.bindings.internal.batch_manager import kv_cache_manager_v2 + + return kv_cache_manager_v2 + + spec = find_spec("kv_cache_manager_v2") + assert spec is not None and spec.origin is not None + trtllm_root = str(Path(spec.origin).parent.parent.parent) + sys.path.insert(0, trtllm_root) + try: + from bindings.internal.batch_manager import kv_cache_manager_v2 + + return kv_cache_manager_v2 + finally: + sys.path.remove(trtllm_root) + + _cpp = _load_cpp_module() + + AggregatedPageDesc = _cpp.AggregatedPageDesc + AttentionLayerConfig = _cpp.AttentionLayerConfig + BatchDesc = _cpp.BatchDesc + # BatchDesc is also consumed via dataclasses.replace(): MambaCacheManager's + # _build_cache_config appends dummy KVCacheDesc slots to each constraint with + # replace(batch, kv_caches=[...]). Like KVCacheManagerConfig below, the C++ + # binding replaces the Python @dataclass, so advertise the dataclass field + # set (replace() is keyed on __dataclass_fields__: reads fields via getattr, + # rebuilds via BatchDesc(**fields)). The binding already has a keyword + # __init__ and readable kv_caches / system_prompt_length fields. + import dataclasses as _dataclasses_bd + + @_dataclasses_bd.dataclass + class _BatchDescFieldSpec: + kv_caches: object = None + system_prompt_length: int = 0 + + BatchDesc.__dataclass_fields__ = _BatchDescFieldSpec.__dataclass_fields__ + del _BatchDescFieldSpec, _dataclasses_bd + BufferConfig = _cpp.BufferConfig + BufferId = _cpp.BufferId + CoalescedBuffer = _cpp.CoalescedBuffer + CacheTier = _cpp.CacheTier + DiskCacheTierConfig = _cpp.DiskCacheTierConfig + GpuCacheTierConfig = _cpp.GpuCacheTierConfig + ExpandedBuffer = _cpp.ExpandedBuffer + HostCacheTierConfig = _cpp.HostCacheTierConfig + KVCacheDesc = _cpp.KVCacheDesc + KVCacheCreatedData = _cpp.KVCacheCreatedData + KVCacheEvent = _cpp.KVCacheEvent + KVCacheEventDiff = _cpp.KVCacheEventDiff + KVCacheEventManager = _cpp.KVCacheEventManager + KVCacheIterationStatsDelta = _cpp.KVCacheIterationStatsDelta + KVCacheManager = _cpp.KVCacheManager + KVCacheManagerConfig = _cpp.KVCacheManagerConfig + # The C++ KVCacheManagerConfig binding replaces the Python @dataclass, but + # callers (the DeepSeek-V4 cache manager's _build_cache_config and our own + # host-tier fallback) use dataclasses.replace() on it. dataclasses.replace() + # is a free function keyed on __dataclass_fields__: it reads each field via + # getattr and rebuilds via cls(**fields). The binding already has a full + # keyword __init__ and readable fields, so we only need to advertise the + # dataclass field set. Field defaults/types are irrelevant here — replace() + # only uses the field names + init flag. The read-only + # enable_swa_scratch_reuse property is intentionally excluded (not a ctor + # field), matching the Python dataclass. + import dataclasses as _dataclasses + + @_dataclasses.dataclass + class _KVCacheManagerConfigFieldSpec: + tokens_per_block: int = 0 + cache_tiers: object = None + layers: object = None + max_util_for_resume: float = 0.97 + enable_partial_reuse: bool = True + constraints: object = None + typical_step: object = None + initial_pool_ratio: object = None + swa_scratch_reuse: object = None + commit_min_snapshot: bool = False + enable_stats: bool = True + + KVCacheManagerConfig.__dataclass_fields__ = _KVCacheManagerConfigFieldSpec.__dataclass_fields__ + del _KVCacheManagerConfigFieldSpec, _dataclasses + KVCacheRemovedData = _cpp.KVCacheRemovedData + KVCacheStatsDelta = _cpp.KVCacheStatsDelta + KVCacheStoredBlockData = _cpp.KVCacheStoredBlockData + KVCacheStoredData = _cpp.KVCacheStoredData + KVCacheUpdatedData = _cpp.KVCacheUpdatedData + KvCacheStatus = _cpp.KvCacheStatus + OutOfPagesError = _cpp.OutOfPagesError + PageStatus = _cpp.PageStatus + PoolDesc = _cpp.PoolDesc + PoolGroupDesc = _cpp.PoolGroupDesc + PoolGroupPeakBlockStats = _cpp.PoolGroupPeakBlockStats + SlotDesc = _cpp.SlotDesc + SlotDescVariant = _cpp.SlotDescVariant + SsmLayerConfig = _cpp.SsmLayerConfig + _KVCache = _cpp._KVCache + _cpp_introspection = getattr(_cpp, "_introspection", None) + _KV_CACHE_ITERATION_STATS_DELTA_FIELDS = tuple(KVCacheIterationStatsDelta._field_names) + PlannedDropHandle = _cpp.PlannedDropHandle + + # Symbols added on main that are not yet ported to the C++ backend. + # TODO(kvCacheManagerV2-cpp): port these and replace the fallbacks. + AttnLifeCycle = getattr(_cpp, "AttnLifeCycle", None) + CuError = getattr(_cpp, "CuError", RuntimeError) + OutOfMemoryError = getattr(_cpp, "OutOfMemoryError", MemoryError) + PageIndexConverter = getattr(_cpp, "PageIndexConverter", None) + ReuseScope = getattr(_cpp, "ReuseScope", ReuseScope) + ScratchDesc = getattr(_cpp, "ScratchDesc", None) + SsmSnapshotIterationStatsDelta = _cpp.SsmSnapshotIterationStatsDelta + SwaScratchReuseConfig = getattr(_cpp, "SwaScratchReuseConfig", None) + UniqueToken = _cpp.UniqueToken + + BeamIndex = int + CacheLevel = int + CacheTierConfig = Union[GpuCacheTierConfig, HostCacheTierConfig, DiskCacheTierConfig] + CudaStream = int + DataRole = str + HalfOpenRange = getattr(_cpp, "HalfOpenRange", tuple) + LayerGroupId = int + LayerId = int + LifeCycleId = int + MemAddress = int + PoolGroupIndex = int + PoolIndex = int + Priority = int + SlidingWindowSize = Optional[int] + TokenId = int + TokenIdExt = Union[int, bytes] + + BAD_PAGE_INDEX = -1 + DEFAULT_BEAM_INDEX = 0 + GPU_LEVEL = 0 + CACHE_LEVEL1 = 1 + NDEBUG = os.environ.get("TLLM_DEBUG_MODE", "")[0:1] != "1" + + class _RawRef: + def __init__(self, obj=None): + self._obj = obj + + def __call__(self): + return self._obj + + def invalidate(self) -> None: + self._obj = None + + @classmethod + def __class_getitem__(cls, _item): + return cls + + rawref = ModuleType(f"{__name__}.rawref") + rawref.ReferenceType = _RawRef + rawref.ref = _RawRef + rawref.NULL = _RawRef() + sys.modules.setdefault(f"{__name__}.rawref", rawref) + + class PageIndexMode(int): + SHARED = 0 + PER_LAYER = 1 + + from ._cache_key import ( # noqa: F401 + gen_multimodal_cache_key_tokens, + sequence_to_blockchain_keys, + ) + + def exact_div(x: int, y: int) -> int: + assert x % y == 0 + return x // y + + def typed_range(*args: int) -> range: + return range(*args) + __all__ = [ "AggregatedPageDesc", @@ -120,13 +323,13 @@ "KVCacheUpdatedData", "KvCacheStatus", "LayerGroupId", - "PlannedDropHandle", "LayerId", "LifeCycleId", "MemAddress", "NDEBUG", "OutOfPagesError", "PageIndexConverter", + "PlannedDropHandle", "PoolGroupPeakBlockStats", "PageIndexMode", "PageStatus", @@ -139,11 +342,11 @@ "ScratchDesc", "KVCacheIterationStatsDelta", "KVCacheStatsDelta", + "SsmSnapshotIterationStatsDelta", "SlidingWindowSize", "SlotDesc", "SlotDescVariant", "SsmLayerConfig", - "SsmSnapshotIterationStatsDelta", "SwaScratchReuseConfig", "TokenId", "TokenIdExt", @@ -154,6 +357,7 @@ "_KVCache", "exact_div", "gen_multimodal_cache_key_tokens", + "sequence_to_blockchain_keys", "rawref", "typed_range", ] diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index 1a03c885b73a..506290e52586 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -287,13 +287,18 @@ class KVCacheEventManager: def flush_iteration_events(self) -> None: ... def get_latest_events(self, timeout_ms: float | None = None) -> list[KVCacheEvent]: ... -# From _block_radix_tree.py +# From _cache_key.py def gen_multimodal_cache_key_tokens( id_offset: int, multi_modal_data_digest: bytes, num_tokens: int, token_offset: int = 0, ) -> list[TokenIdExt]: ... +def sequence_to_blockchain_keys( + tokens_per_block: int, + reuse_scope: ReuseScope, + tokens: Sequence[TokenIdExt], +) -> Iterator[tuple[list[TokenIdExt], bytes]]: ... # From _core/_kv_cache.py class _Status(enum.Enum): diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py index 110912a07ddb..c47c591f0d06 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py @@ -13,17 +13,21 @@ # See the License for the specific language governing permissions and # limitations under the License. -import hashlib -from array import array from typing import TYPE_CHECKING, Iterable, Iterator, NamedTuple, Sequence, TypeVar, cast from . import rawref -from ._common import NDEBUG, BlockOrdinal, PageStatus, TokenId, TokenIdExt +from ._cache_key import ( # noqa: F401 + BlockKey, + Hasher, + TokenBlock, + gen_multimodal_cache_key_tokens, + reuse_scope_to_bytes, + sequence_to_blockchain_keys, +) +from ._common import NDEBUG, BlockOrdinal, PageStatus, TokenIdExt from ._life_cycle_registry import AttnLifeCycle, LifeCycle, LifeCycleId, LifeCycleRegistry from ._utils import ( TypedIndexList, - chunked, - div_up, expect_type, filled_list, find_index, @@ -36,8 +40,6 @@ from ._event_manager import KVCacheEventManager from ._page import CommittedPage -BlockKey = bytes - class ReuseScope(NamedTuple): """Per-request namespace for prefix reuse.""" @@ -45,21 +47,8 @@ class ReuseScope(NamedTuple): lora_id: int | None = None salt: int | None = None - def _mask(self) -> bytes: - return sum((value is not None) << i for i, value in enumerate(self)).to_bytes( - div_up(len(self), 8), "little", signed=False - ) - def to_bytes(self) -> bytes: - ret = self._mask() - for value in self: - if type(value) is int: - ret += value.to_bytes(8, "little", signed=False) - else: - assert value is None, ( - "Did you forget to update to_bytes() when adding new non-int fields to ReuseScope?" - ) - return ret + return reuse_scope_to_bytes(self) class ReuseMatch(NamedTuple): @@ -70,74 +59,6 @@ class ReuseMatch(NamedTuple): num_lookup_tokens: int -# id_offset is usually vocab_size -def gen_multimodal_cache_key_tokens( - id_offset: int, multi_modal_data_digest: bytes, num_tokens: int, token_offset: int = 0 -) -> list[TokenIdExt]: - """Create synthetic tokens used only when building multimodal KV-cache keys. - - Item-local token 0 carries the content digest; later offsets use deterministic IDs above the vocab. - """ - assert num_tokens > 0 - assert token_offset >= 0 - return [ - multi_modal_data_digest if token_offset + i == 0 else TokenId(id_offset + token_offset + i) - for i in range(num_tokens) - ] - - -class Hasher: - __slots__ = "_hasher" - _hasher: "hashlib._Hash" - - def __init__(self, data: int | bytes | None | Sequence[int | bytes] = None) -> None: - self._hasher = hashlib.sha256() - if data is not None: - self.update(data) - - # This function is perf-critical. Expect compromised code quality. - def update(self, data: int | bytes | Sequence[int | bytes]) -> "Hasher": - if type(data) is int: - assert NDEBUG or (data >= 0 and data < (1 << 64)) - self._hasher.update(data.to_bytes(8, "little")) - elif type(data) is bytes: - self._hasher.update(data) - else: - # Hash the whole token block in one C call instead of one per token. - # array("Q", data).tobytes() packs each int as 8 native-endian bytes; - # all NVIDIA GPU host platforms (x86_64, aarch64/Grace) are little-endian - # so this is byte-identical to the per-token to_bytes(8, "little") loop. - # Falls back to that loop for multimodal blocks (which contain bytes items). - try: - self._hasher.update(array("Q", data).tobytes()) # type: ignore - except (TypeError, OverflowError): - for item in data: # type: ignore - assert ( - NDEBUG - or (type(item) is int and (0 <= item < (1 << 64))) - or type(item) is bytes - ) - self._hasher.update(item.to_bytes(8, "little") if (type(item) is int) else item) # type: ignore - return self - - @property - def digest(self) -> bytes: - return self._hasher.digest() - - -TokenBlock = list[TokenIdExt] - - -def sequence_to_blockchain_keys( - tokens_per_block: int, reuse_scope: ReuseScope, tokens: Sequence[TokenIdExt] -) -> Iterator[tuple[TokenBlock, BlockKey]]: - digest = Hasher(reuse_scope.to_bytes()).digest - yield [], digest - for token_block in chunked(tokens, tokens_per_block): - digest = Hasher(digest).update(token_block).digest - yield token_block, digest - - Child = TypeVar("Child", bound="Block | RootBlock") Children = dict[BlockKey, Child] @@ -417,6 +338,9 @@ def num_life_cycles(self) -> LifeCycleId: def prev(self) -> "Block | RootBlock": return unwrap_rawref(self._prev) + def get_page(self, lc_idx: LifeCycleId) -> "CommittedPage | None": + return map_optional(self.storage[lc_idx], lambda f: f()) + def unlink_page( self, lc_idx: LifeCycleId, expected_page: "CommittedPage | None" = None ) -> bool: diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_cache_key.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_cache_key.py new file mode 100644 index 000000000000..c418c4a9708a --- /dev/null +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_cache_key.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Backend-neutral helpers for deriving KV-cache reuse keys. + +These are the pure-Python reference implementation shared by both the Python +and C++ backends. They depend only on the light-weight, backend-neutral +``_common`` module (no CUDA / bindings), so they can be imported and re-exported +as public API regardless of the active backend. +""" + +import hashlib +import itertools +from array import array +from typing import Iterable, Iterator, Sequence + +from ._common import NDEBUG, TokenId, TokenIdExt + +BlockKey = bytes +TokenBlock = list[TokenIdExt] + + +# id_offset is usually vocab_size +def gen_multimodal_cache_key_tokens( + id_offset: int, multi_modal_data_digest: bytes, num_tokens: int, token_offset: int = 0 +) -> list[TokenIdExt]: + """Create synthetic tokens used only when building multimodal KV-cache keys. + + Item-local token 0 carries the content digest; later offsets use deterministic IDs above the vocab. + """ + assert num_tokens > 0 + assert token_offset >= 0 + return [ + multi_modal_data_digest if token_offset + i == 0 else TokenId(id_offset + token_offset + i) + for i in range(num_tokens) + ] + + +class Hasher: + # SECURITY INVARIANT: the block-key hash MUST stay cryptographically + # collision-resistant and >= 256-bit. The radix tree is a globally shared, + # cross-request/cross-tenant cache index; prefix matches are decided purely by + # digest equality with NO re-check of the underlying tokens; and the hashed + # input (tokens, the user-supplied cache_salt, multimodal content bytes) is + # attacker-influenceable. A collision therefore silently reuses another + # request's KV blocks (cross-request corruption / data leak), and cache_salt + # tenant isolation relies entirely on this hash's collision resistance. Do NOT + # swap in a non-cryptographic hash (xxHash, HighwayHash, ...) or truncate below + # 256 bits without first adding a token-content equality check on match. The + # C++ backend (blockRadixTree) mirrors this with SHA-256 (CSHA256). + __slots__ = "_hasher" + _hasher: "hashlib._Hash" + + def __init__(self, data: int | bytes | None | Sequence[int | bytes] = None) -> None: + self._hasher = hashlib.sha256() + if data is not None: + self.update(data) + + # This function is perf-critical. Expect compromised code quality. + def update(self, data: int | bytes | Sequence[int | bytes]) -> "Hasher": + if type(data) is int: + assert NDEBUG or (data >= 0 and data < (1 << 64)) + self._hasher.update(data.to_bytes(8, "little")) + elif type(data) is bytes: + self._hasher.update(data) + else: + # Hash the whole token block in one C call instead of one per token. + # array("Q", data).tobytes() packs each int as 8 native-endian bytes; + # all NVIDIA GPU host platforms (x86_64, aarch64/Grace) are little-endian + # so this is byte-identical to the per-token to_bytes(8, "little") loop. + # Falls back to that loop for multimodal blocks (which contain bytes items). + try: + self._hasher.update(array("Q", data).tobytes()) # type: ignore + except (TypeError, OverflowError): + for item in data: # type: ignore + assert ( + NDEBUG + or (type(item) is int and (0 <= item < (1 << 64))) + or type(item) is bytes + ) + self._hasher.update(item.to_bytes(8, "little") if (type(item) is int) else item) # type: ignore + return self + + @property + def digest(self) -> bytes: + return self._hasher.digest() + + +def reuse_scope_to_bytes(reuse_scope: Iterable[int | None]) -> bytes: + """Serialize a reuse scope to its reuse-namespace bytes. + + Backend-neutral: reads the scope's fields by iteration, so it works for both + the pure-Python ``ReuseScope`` NamedTuple and the C++ binding without relying + on a ``to_bytes()`` method. The layout mirrors the C++ ``emitReuseScopeBytes``: + a mask byte (one bit per field, set when the field is present) followed by one + little-endian ``uint64`` per present field (``signed=False``). + """ + values = list(reuse_scope) + mask = sum((value is not None) << i for i, value in enumerate(values)) + ret = mask.to_bytes((len(values) + 7) // 8, "little", signed=False) + for value in values: + if value is not None: + ret += int(value).to_bytes(8, "little", signed=False) + return ret + + +def sequence_to_blockchain_keys( + tokens_per_block: int, reuse_scope: Iterable[int | None], tokens: Sequence[TokenIdExt] +) -> Iterator[tuple[TokenBlock, BlockKey]]: + """Yield ``(token_block, key)`` pairs seeding a blockchain of KV-cache keys. + + The first pair is the root (``[]``, reuse-scope digest); each subsequent pair + hashes one ``tokens_per_block`` chunk on top of the previous digest. + """ + digest = Hasher(reuse_scope_to_bytes(reuse_scope)).digest + yield [], digest + iterator = iter(tokens) + while True: + token_block = list(itertools.islice(iterator, tokens_per_block)) + if not token_block: + break + digest = Hasher(digest).update(token_block).digest + yield token_block, digest diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py index 303d79e9ee13..da0500dc4c06 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py @@ -1020,7 +1020,12 @@ def commit( ), move_ssm=is_end, ) - if has_partial_snapshot: + # _commit_block transitions out of ALLOWED (to USER_STOP) when a + # block cannot be committed (VIRTUAL_STOP). Stop here so we don't + # re-enter _commit_block on an already-stopped cache. + if self._commit_state != self.CommitState.ALLOWED: + break + if has_partial_snapshot and self._commit_state == self.CommitState.ALLOWED: partial_ordinal = BlockOrdinal(new_num_full_blocks) if is_end: self._commit_block( @@ -1061,49 +1066,36 @@ def plan_committed_block_drop(self) -> PlannedDropHandle | None: if self._commit_state != self.CommitState.USER_STOP: raise LogicError("plan_committed_block_drop() requires stop_committing()") - ssm_lc_id = self.manager._life_cycles.ssm_life_cycle_id - matched_blocks: list[Block] = [] - if self.num_committed_tokens > 0: - # Locate pages through the radix tree rather than - # SeqBlock.tree_block: the latter is not guaranteed to identify a - # partial snapshot after reuse. Requiring an exact match keeps the - # preceding conversation plan intact if this turn no longer has a - # complete reusable endpoint. All PP ranks use the same lookup so - # attention-only ranks include the final partial SWA block too. - match = self.manager._match_reuse(self.reuse_scope, self._committed_tokens) - if match.num_tokens != self.num_committed_tokens or not match.blocks: - return None - matched_blocks = match.blocks - elif ssm_lc_id is not None: + if self.num_committed_tokens == 0: return None - def get_committed_page(block: Block, life_cycle_id: LifeCycleId) -> CommittedPage | None: - page_ref = block.storage[life_cycle_id] - return page_ref() if page_ref is not None else None + # Locate pages through the radix tree rather than + # SeqBlock.tree_block: the latter is not guaranteed to identify a + # partial snapshot after reuse. Requiring an exact match keeps the + # preceding conversation plan intact if this turn no longer has a + # complete reusable endpoint. All PP ranks use the same lookup so + # attention-only ranks include the final partial SWA block too. + match = self.manager._match_reuse(self.reuse_scope, self._committed_tokens) + if match.num_tokens != self.num_committed_tokens or not match.blocks: + return None - end = BlockOrdinal(len(matched_blocks)) + end = BlockOrdinal(len(match.blocks)) pages_to_drop: list[CommittedPage] = [] for lc_idx, lc in self.manager._life_cycles.items(): - if isinstance(lc, SsmLifeCycle): - continue - if lc.window_size is None: - continue - stale_range = _KVCache._get_stale_range( - self.tokens_per_block, self.num_committed_tokens, lc - ) - window_start = min(stale_range.end, end) + if isinstance(lc, AttnLifeCycle): + if lc.window_size is None: + continue + stale_range = _KVCache._get_stale_range( + self.tokens_per_block, self.num_committed_tokens, lc + ) + window_start = min(stale_range.end, end) + else: + window_start = BlockOrdinal(end - 1) for ordinal in typed_range(window_start, end): - tree_block = matched_blocks[ordinal] - page = get_committed_page(tree_block, lc_idx) + page = match.blocks[ordinal].get_page(lc_idx) if page is None: return None pages_to_drop.append(page) - - if ssm_lc_id is not None: - page = get_committed_page(matched_blocks[-1], ssm_lc_id) - if not isinstance(page, SsmCommittedPage): - return None - pages_to_drop.append(page) return PlannedDropHandle(pages_to_drop) # Users promise to not commit any more tokens. For cases where we shouldn't reuse generated tokens diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py index 2492278bc9df..74c73700d3e6 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py @@ -671,6 +671,11 @@ def layer_grouping(self) -> HomoTuple[HomoTuple[LayerId]]: """ Layers are divided into multiple groups. Buffers in the same layer group for the same token block are always allocated/deallocated together. + + NOTE: the iteration order of the layer lists (and of the groups) is NOT part of + the API contract and may differ across backends/runs. Do not rely on it for + buffer/pool memory order -- query ``pool_group_descs`` (PoolGroupDesc.pools[i] + .base_address + coalesced_buffers) for that. """ layer_to_life_cycle_ids = self._storage._layer_to_life_cycle_ids num_life_cycles = self._life_cycles.size diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py index 8b063677111d..4a5bdc44394f 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py @@ -197,3 +197,127 @@ def ratio_to_slot_count_list( quota, slot_size_lists, ratio_list, granularity, min_slots ) ) + + +def attention_life_cycle_ids(manager: Any) -> list[int]: + """Return the lifecycle ids of all attention lifecycles, in order.""" + cpp_introspection = _cpp_introspection_module() + if cpp_introspection is not None: + return list(cpp_introspection.attention_life_cycle_ids(manager)) + return [lc_id for lc_id, _ in manager._life_cycles.attention_life_cycles()] + + +def swa_life_cycle_ids(manager: Any) -> list[int]: + """Return the lifecycle ids of attention lifecycles that use a sliding window.""" + cpp_introspection = _cpp_introspection_module() + if cpp_introspection is not None: + return list(cpp_introspection.swa_life_cycle_ids(manager)) + return [ + lc_id + for lc_id, lc in manager._life_cycles.attention_life_cycles() + if lc.window_size is not None + ] + + +def ssm_life_cycle_id(manager: Any) -> int | None: + """Return the SSM lifecycle id, or None if there is no SSM lifecycle.""" + cpp_introspection = _cpp_introspection_module() + if cpp_introspection is not None: + return cpp_introspection.ssm_life_cycle_id(manager) + return manager._life_cycles.ssm_life_cycle_id + + +def reuse_match_pages( + manager: Any, + reuse_scope: Any, + tokens: Any, + lc_id: int, + enable_partial: bool = False, +) -> tuple[int, list[tuple[int, int | None] | None]]: + """Match ``tokens`` against the radix tree and report reusable pages per block. + + Returns ``(num_tokens, pages)`` where ``pages[i]`` is ``None`` when block ``i`` + holds no page for lifecycle ``lc_id``, otherwise ``(slot_id, num_tokens_in_block)`` + with ``num_tokens_in_block`` set only for SSM pages (``None`` for attention pages). + """ + cpp_introspection = _cpp_introspection_module() + if cpp_introspection is not None: + num_tokens, raw_pages = cpp_introspection.reuse_match_pages( + manager, reuse_scope, list(tokens), lc_id, enable_partial + ) + pages: list[tuple[int, int | None] | None] = [] + for entry in raw_pages: + if entry is None: + pages.append(None) + else: + slot_id, num_tokens_in_block = entry + pages.append((slot_id, None if num_tokens_in_block < 0 else num_tokens_in_block)) + return num_tokens, pages + + from ._utils import unwrap_rawref + + match = manager._radix_tree.match(reuse_scope, list(tokens), enable_partial) + py_pages: list[tuple[int, int | None] | None] = [] + for block in match.blocks: + ref = block.storage[lc_id] + if ref is None: + py_pages.append(None) + else: + page = unwrap_rawref(ref) + py_pages.append((page.slot_id, getattr(page, "num_tokens_in_block", None))) + return match.num_tokens, py_pages + + +def reuse_match_planned_drop_counts( + manager: Any, + reuse_scope: Any, + tokens: Any, + lc_id: int, + enable_partial: bool = False, +) -> tuple[int, list[int | None]]: + """Match ``tokens`` and report each matched block's ``planned_drop_count`` for lifecycle ``lc_id``. + + Returns ``(num_tokens, counts)`` where ``counts[i]`` is ``None`` when block ``i`` holds no page + for lifecycle ``lc_id``, otherwise the matched page's ``planned_drop_count``. + """ + cpp_introspection = _cpp_introspection_module() + if cpp_introspection is not None: + return cpp_introspection.reuse_match_planned_drop_counts( + manager, reuse_scope, list(tokens), lc_id, enable_partial + ) + + from ._utils import unwrap_rawref + + match = manager._radix_tree.match(reuse_scope, list(tokens), enable_partial) + counts: list[int | None] = [] + for block in match.blocks: + ref = block.storage[lc_id] + counts.append(None if ref is None else unwrap_rawref(ref).planned_drop_count) + return match.num_tokens, counts + + +def pool_group_index(manager: Any, lc_id: int) -> int: + """Return the storage pool-group index for lifecycle ``lc_id``.""" + cpp_introspection = _cpp_introspection_module() + if cpp_introspection is not None: + return cpp_introspection.pool_group_index(manager, lc_id) + return manager._storage.get_pool_group_index(lc_id) + + +def compute_slots_for_batch( + manager: Any, + batch: Any, + tokens_per_block: int, + swa_scratch_reuse: Any = None, +) -> list[int]: + """Return the minimum per-pool-group slot counts to support ``batch``.""" + cpp_introspection = _cpp_introspection_module() + if cpp_introspection is not None: + return list( + cpp_introspection.compute_slots_for_batch( + manager, batch, tokens_per_block, swa_scratch_reuse + ) + ) + return list( + manager._storage._compute_slots_for_batch(batch, tokens_per_block, swa_scratch_reuse) + ) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py index 69836ab0efcf..6d93c6ec3220 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py @@ -244,12 +244,12 @@ def __init__( gpu_quota = config.cache_tiers[GPU_LEVEL].quota gpu_granularity = CacheLevelManager.cache_tier_granularity(CacheTier.GPU_MEM, gpu_quota) - constraints_for_min_slots = [] if initial_pool_ratio is not None else constraints or [] + # Constraints stay feasibility floors even under an explicit initial pool + # ratio (a share below what a declared batch needs is clamped up), and the + # floors are scaled by 1/max_util_for_resume because _KVCache.resume rejects + # any pool group above that utilization. Mirrors PR #16269. self._min_slots = self._compute_min_slots_from_constraints( - constraints_for_min_slots, - tokens_per_block, - swa_scratch_reuse, - max_util_for_resume, + constraints or [], tokens_per_block, swa_scratch_reuse, max_util_for_resume ) # Compute init_ratio from explicit config, typical_batch, constraints, or fallback. @@ -918,6 +918,8 @@ def _compute_min_slots_from_constraints( All returned elements are positive. Constraint-derived floors include headroom for the utilization gate checked by ``_KVCache.resume``. """ + if not 0 < max_util_for_resume <= 1: + raise ValueError(f"max_util_for_resume must be in (0, 1], got {max_util_for_resume}") max_slots = filled_list(0, self.num_pool_groups) def swa_floor_blocks(lc: AttnLifeCycle) -> int: diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_utils.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_utils.py index da750eeca0e4..2f970ca347b7 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_utils.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_utils.py @@ -673,13 +673,28 @@ def num_set_bits(self) -> int: return self._num_set_bits def resize(self, new_capacity: int) -> None: - extra_elems = div_up(new_capacity, 64) - len(self._bits) - if extra_elems > 0: - self._bits.extend(array.array(self.TYPE_CODE, [0] * extra_elems)) - elif extra_elems < 0: - self._bits = self._bits[:extra_elems] - if new_capacity % 64 != 0: - self._bits[-1] &= self.ALL_SET_MASK >> (64 - (new_capacity % 64)) + old_elems = len(self._bits) + new_elems = div_up(new_capacity, 64) + + # When the capacity shrinks, every set bit at or above new_capacity is + # dropped. Account for those bits so num_set_bits stays accurate, and + # mask the retained partial word so any_set() cannot observe stale bits. + # This covers both fewer-words and same-word-count (new_elems == + # old_elems with a smaller new_capacity) shrinks. + if new_elems <= old_elems: + for w in range(new_elems, old_elems): + self._num_set_bits -= bin(self._bits[w]).count("1") + if new_elems >= 1 and new_capacity % 64 != 0: + keep_mask = self.ALL_SET_MASK >> (64 - (new_capacity % 64)) + word = self._bits[new_elems - 1] + dropped = word & ~keep_mask & self.ALL_SET_MASK + self._num_set_bits -= bin(dropped).count("1") + self._bits[new_elems - 1] = word & keep_mask + + if new_elems > old_elems: + self._bits.extend(array.array(self.TYPE_CODE, [0] * (new_elems - old_elems))) + elif new_elems < old_elems: + self._bits = self._bits[:new_elems] # check if any bit in the range [start, end) is set def any_set(self, start: int, end: int) -> bool: diff --git a/tests/integration/test_lists/test-db/l0_cpu_arm.yml b/tests/integration/test_lists/test-db/l0_cpu_arm.yml index 07b1a2767da0..9b35d56db7bd 100644 --- a/tests/integration/test_lists/test-db/l0_cpu_arm.yml +++ b/tests/integration/test_lists/test-db/l0_cpu_arm.yml @@ -14,4 +14,5 @@ l0_cpu_arm: orchestrator: mpi tests: - unittest/executor/test_rpc.py + - unittest/executor/test_event_loop_error_broadcast.py - unittest/others/test_http_utils_fail_fast.py diff --git a/tests/integration/test_lists/test-db/l0_cpu_x86.yml b/tests/integration/test_lists/test-db/l0_cpu_x86.yml index 6d196908b41f..072f17dcc00a 100644 --- a/tests/integration/test_lists/test-db/l0_cpu_x86.yml +++ b/tests/integration/test_lists/test-db/l0_cpu_x86.yml @@ -16,3 +16,4 @@ l0_cpu_x86: - unittest/executor/test_rpc.py - unittest/others/test_http_utils_fail_fast.py - unittest/executor/test_multi_frontend_routing.py + - unittest/executor/test_event_loop_error_broadcast.py diff --git a/tests/unittest/_torch/executor/test_kv_pool_rebalance.py b/tests/unittest/_torch/executor/test_kv_pool_rebalance.py index 59f4f3eb39ea..ed0cc5636e7b 100644 --- a/tests/unittest/_torch/executor/test_kv_pool_rebalance.py +++ b/tests/unittest/_torch/executor/test_kv_pool_rebalance.py @@ -32,7 +32,7 @@ import pytest from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor -from tensorrt_llm.runtime.kv_cache_manager_v2._exceptions import OutOfPagesError +from tensorrt_llm.runtime.kv_cache_manager_v2 import OutOfPagesError # --------------------------------------------------------------------------- # # Helpers diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index e7d61248e7b4..e80554f354b3 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -74,6 +74,7 @@ KVCacheManagerConfig, LayerId, SsmLayerConfig, + _introspection, ) from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManager as RuntimeKVCacheManager @@ -1266,7 +1267,8 @@ def test_v2_hybrid_typical_batch_splits_capacity_across_ssm_states_and_dummies() config = mgr._build_cache_config(base_config) assert isinstance(config.layers[0], SsmLayerConfig) - assert config.layers[1] is base_layers[1] + assert isinstance(config.layers[1], AttentionLayerConfig) + assert int(config.layers[1].layer_id) == int(base_layers[1].layer_id) assert config.typical_step == BatchDesc( [KVCacheDesc(capacity=32, history_length=31)] * 6 + [KVCacheDesc(capacity=0, history_length=0)] @@ -1566,12 +1568,17 @@ def allocated_memory(pool_ratio): config = mgr._build_cache_config(base_config) runtime_manager = RuntimeKVCacheManager(config) try: - statistics = runtime_manager._storage.get_statistics() + statistics = _introspection.storage_statistics(runtime_manager) + + def _slot_sizes(stat): + # cpp binding exposes `slot_sizes`; the Python backend `slot_size`. + return stat.slot_sizes if hasattr(stat, "slot_sizes") else stat.slot_size + allocated_bytes = [ - int(stats.total) * sum(int(size) for size in stats.slot_size) + int(stats.total) * sum(int(size) for size in _slot_sizes(stats)) for stats in statistics ] - return allocated_bytes, list(runtime_manager._current_gpu_ratio) + return allocated_bytes, _introspection.current_gpu_ratio(runtime_manager) finally: runtime_manager.shutdown() @@ -2135,7 +2142,7 @@ def test_v2_hybrid_disagg_page_table_preserves_lifecycle_indices(): try: page_table = build_page_table_from_manager(mgr) - assert len(page_table.layer_groups) == mgr.impl._storage.num_life_cycles + assert len(page_table.layer_groups) == len(mgr.impl.layer_grouping) assert isinstance(page_table.layer_groups[0], MambaLayerGroup) assert isinstance(page_table.layer_groups[1], AttentionLayerGroup) diff --git a/tests/unittest/disaggregated/test_router.py b/tests/unittest/disaggregated/test_router.py index f7bf10adf3a7..dd0b6d5facb3 100644 --- a/tests/unittest/disaggregated/test_router.py +++ b/tests/unittest/disaggregated/test_router.py @@ -14,7 +14,7 @@ from tensorrt_llm.runtime.kv_cache_hash import (get_cache_salt_id, hash_v1_block_key, truncate_sha256_hash_to_int64) -from tensorrt_llm.runtime.kv_cache_manager_v2._block_radix_tree import ( +from tensorrt_llm.runtime.kv_cache_manager_v2 import ( ReuseScope, sequence_to_blockchain_keys) # yapf: disable from tensorrt_llm.serve.openai_protocol import (ChatCompletionRequest, diff --git a/tests/unittest/executor/test_base_worker.py b/tests/unittest/executor/test_base_worker.py index a60922edb72f..7bee69ef87f9 100644 --- a/tests/unittest/executor/test_base_worker.py +++ b/tests/unittest/executor/test_base_worker.py @@ -35,6 +35,9 @@ def raise_load_error(lora_request): raise RuntimeError("bad adapter") worker = object.__new__(BaseWorker) + # GC-time __del__ -> shutdown() reads this; __init__ is bypassed here, so + # seed it to keep teardown a clean no-op. + worker.doing_shutdown = False worker._lora_manager = LoraManager() worker._load_lora_adapter = raise_load_error request = type( diff --git a/tests/unittest/executor/test_event_loop_error_broadcast.py b/tests/unittest/executor/test_event_loop_error_broadcast.py index 04851b8e651b..df8db031cd9e 100644 --- a/tests/unittest/executor/test_event_loop_error_broadcast.py +++ b/tests/unittest/executor/test_event_loop_error_broadcast.py @@ -13,9 +13,16 @@ import datetime import queue as _stdlib_queue +import pytest + from tensorrt_llm.executor.base_worker import AwaitResponseHelper from tensorrt_llm.executor.utils import ErrorResponse +# CI's CPU stages select tests with ``-m "cpu_only and not disabled"``; without +# this marker the whole file is deselected and pytest exits 5 (no tests ran), +# which the runner reports as a failure. These are pure stub-based unit tests. +pytestmark = pytest.mark.cpu_only + class _EngineStub: """Stub for self.worker.engine: returns whatever the test plugged in.""" @@ -56,6 +63,9 @@ def __init__(self, engine, num_pending: int = 1): self.popped = [] self.result_queue = None self.postproc_queues = None + # responses_handler() reads this unguarded (base_worker.py); BaseWorker + # sets it in __init__, which this stub bypasses, so seed it to None. + self.frontend_result_queues = None # Echoed straight back so __call__'s filter is a no-op. def _engine_response_callback(r): diff --git a/tests/unittest/executor/test_proxy_postproc_terminate.py b/tests/unittest/executor/test_proxy_postproc_terminate.py index 00b6883c8734..0043360a7a32 100644 --- a/tests/unittest/executor/test_proxy_postproc_terminate.py +++ b/tests/unittest/executor/test_proxy_postproc_terminate.py @@ -38,7 +38,12 @@ def put(self, res): def _make_proxy(): # Avoid GenerationExecutorProxy.__init__ (it spawns workers); we only # exercise the pure dispatch logic with hand-set attributes. - return object.__new__(GenerationExecutorProxy) + proxy = object.__new__(GenerationExecutorProxy) + # GC-time __del__ -> shutdown() reads these; __init__ is bypassed here, so + # seed them to keep teardown a clean no-op. + proxy.workers_started = False + proxy._multi_frontend_ipc_dir = None + return proxy def test_late_response_after_terminate_is_dropped_without_keyerror(): diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py index db5921867760..c3d88dc0433b 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py @@ -15,6 +15,7 @@ import gc import os +import pickle import threading import time from importlib.util import find_spec @@ -28,10 +29,28 @@ KV_CACHE_HASH_ALGO_V2_SHA256_64, truncate_sha256_hash_to_int64, ) +from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheCreatedData as NativeKVCacheCreatedData +from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheEvent as NativeKVCacheEvent +from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheEventDiff as NativeKVCacheEventDiff +from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + KVCacheEventManager as NativeKVCacheEventManager, +) +from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheRemovedData as NativeKVCacheRemovedData +from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + KVCacheStoredBlockData as NativeKVCacheStoredBlockData, +) +from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheStoredData as NativeKVCacheStoredData +from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheUpdatedData as NativeKVCacheUpdatedData +from tensorrt_llm.runtime.kv_cache_manager_v2 import UniqueToken as NativeUniqueToken from tensorrt_llm.runtime.kv_cache_manager_v2._event_manager import ( + KVCacheCreatedData, + KVCacheEvent, KVCacheEventDiff, KVCacheEventManager, + KVCacheRemovedData, KVCacheStoredBlockData, + KVCacheStoredData, + KVCacheUpdatedData, UniqueToken, ) @@ -64,6 +83,7 @@ _DEFAULT_CACHE_LEVEL = CacheLevel(0) +_USING_CPP_BACKEND = os.environ.get("TLLM_KV_CACHE_MANAGER_V2_BACKEND", "cpp").lower() != "python" class _FakePage: @@ -138,6 +158,13 @@ def _flush_serialized_events(event_manager): return KVCacheEventSerializer.serialize(event_manager.get_latest_events(0)) +def _flush_contract_events(event_manager): + events = _flush_serialized_events(event_manager) + for event in events: + event["hash_algo"] = "" + return events + + def _stored_events(events): return [event for event in events if event["data"]["type"] == "stored"] @@ -169,6 +196,251 @@ def _commit_and_close(manager, stream, tokens, *, input_tokens=None, reuse_scope gc.collect() +@pytest.mark.skipif(not _USING_CPP_BACKEND, reason="requires the native C++ event manager") +def test_native_event_manager_queue_and_stored_coalescing(): + event_manager = NativeKVCacheEventManager( + max_kv_event_entries=8, + window_size=128, + hash_algo="v2_sha256", + ) + event_manager.add_stored_event( + None, + [ + NativeKVCacheStoredBlockData( + "block0", + [NativeUniqueToken(1), NativeUniqueToken(2)], + cache_level=0, + priority=35, + ) + ], + ) + event_manager.add_stored_event( + "block0", + [ + NativeKVCacheStoredBlockData( + "block1", + [NativeUniqueToken(3), NativeUniqueToken(4)], + cache_level=0, + priority=35, + ) + ], + ) + + event_manager.flush_iteration_events() + event_objects = pickle.loads(pickle.dumps(event_manager.get_latest_events(0))) + events = KVCacheEventSerializer.serialize(event_objects) + + assert len(events) == 1 + assert events[0]["hash_algo"] == "v2_sha256" + assert [block["block_hash"] for block in events[0]["data"]["blocks"]] == [ + "block0", + "block1", + ] + + +def test_native_event_manager_v1_hash_matches_legacy_cpp_hasher(): + _tb = pytest.importorskip("tensorrt_llm.bindings") + block_key = _tb.internal.batch_manager.BlockKey + block_key_hasher = _tb.internal.batch_manager.BlockKeyHasher + parent_hash = block_key_hasher.hash(block_key([1, 2, 3, 4])) + + assert NativeKVCacheEventManager._hash_block_key([1, 2, 3, 4], 0, None, None) == parent_hash + assert NativeKVCacheEventManager._hash_block_key( + [5, 6], parent_hash, None, None + ) == block_key_hasher.hash(block_key([5, 6]), parent_hash) + + +@pytest.mark.skipif(not _USING_CPP_BACKEND, reason="requires the native C++ event manager") +def test_native_event_manager_attention_dp_gather_callback(): + gathered_events = [] + + def gather(local_events): + gathered_events.extend(local_events) + return [pickle.loads(pickle.dumps(local_events))] + + event_manager = NativeKVCacheEventManager( + max_kv_event_entries=2, + window_size=128, + attention_dp_rank=0, + attention_dp_gather=gather, + ) + event_manager.add_created_event([4]) + + events = _flush_serialized_events(event_manager) + + assert len(gathered_events) == 1 + assert events[0]["attention_dp_rank"] == 0 + assert events[0]["data"]["num_blocks_per_cache_level"] == [4] + + +@pytest.mark.skipif(not _USING_CPP_BACKEND, reason="compares native and Python event managers") +def test_native_event_data_value_semantics_match_python_reference(): + native_token = NativeUniqueToken("token", 7) + native_block = NativeKVCacheStoredBlockData( + "block", + [native_token], + cache_level=1, + priority=35, + mm_keys=[(b"short-mm-key", 3), (b"another-key", 5, "uuid")], + cache_salt="salt", + ) + native_diff = NativeKVCacheEventDiff(0, 1) + native_objects = [ + native_token, + NativeKVCacheCreatedData([2, 3]), + native_block, + NativeKVCacheStoredData(None, [native_block]), + NativeKVCacheRemovedData(["block"]), + native_diff, + NativeKVCacheUpdatedData("block", native_diff, None), + ] + native_event = NativeKVCacheEvent( + 4, + native_objects[3], + 128, + "same-hash-label", + 1, + 2, + ) + native_objects.append(native_event) + + python_token = UniqueToken("token", 7) + python_block = KVCacheStoredBlockData( + "block", + [python_token], + cache_level=1, + priority=35, + mm_keys=[(b"short-mm-key", 3), (b"another-key", 5, "uuid")], + cache_salt="salt", + ) + python_diff = KVCacheEventDiff(0, 1) + python_objects = [ + python_token, + KVCacheCreatedData([2, 3]), + python_block, + KVCacheStoredData(None, [python_block]), + KVCacheRemovedData(["block"]), + python_diff, + KVCacheUpdatedData("block", python_diff, None), + ] + python_objects.append( + KVCacheEvent( + 4, + python_objects[3], + 128, + "same-hash-label", + 1, + 2, + ) + ) + + for value, python_value in zip(native_objects, python_objects): + assert pickle.loads(pickle.dumps(value)) == value + assert repr(value) == repr(python_value) + + python_event = KVCacheEventManager(max_kv_event_entries=1) + native_manager = NativeKVCacheEventManager(max_kv_event_entries=1) + python_event.add_stored_event(None, [python_block], layer_group_id=2) + native_manager.add_stored_event(None, [native_block], layer_group_id=2) + assert _flush_contract_events(native_manager) == _flush_contract_events(python_event) + + +@pytest.mark.skipif(not _USING_CPP_BACKEND, reason="compares native and Python event managers") +def test_native_event_manager_public_methods_match_python_reference(): + python_manager = KVCacheEventManager( + max_kv_event_entries=16, + window_size=128, + window_size_by_layer_group=None, + ) + native_manager = NativeKVCacheEventManager( + max_kv_event_entries=16, + window_size=128, + window_size_by_layer_group=None, + ) + python_manager.set_layer_group_window_sizes({0: 64, 1: 96}) + native_manager.set_layer_group_window_sizes({0: 64, 1: 96}) + + for manager, block_type, token_type, diff_type in ( + (python_manager, KVCacheStoredBlockData, UniqueToken, KVCacheEventDiff), + ( + native_manager, + NativeKVCacheStoredBlockData, + NativeUniqueToken, + NativeKVCacheEventDiff, + ), + ): + manager.add_created_event([2, 3], layer_group_ids=[0, 1]) + manager.add_stored_event( + None, + [block_type("block-0", [token_type(1)], cache_level=0, priority=35)], + layer_group_id=0, + ) + manager.add_stored_event( + "block-0", + [block_type("block-1", [token_type(2)], cache_level=0, priority=35)], + layer_group_id=0, + ) + manager.add_removed_event(block_hash for block_hash in ("removed-0", "removed-1")) + manager.add_updated_event( + "updated", + cache_level=diff_type(0, 1), + priority=diff_type(35, 50), + layer_group_id=1, + ) + + assert _flush_contract_events(native_manager) == _flush_contract_events(python_manager) + assert native_manager.get_latest_events(0) == python_manager.get_latest_events(0) == [] + assert native_manager.get_latest_events(-1) == python_manager.get_latest_events(-1) == [] + + +@pytest.mark.skipif(not _USING_CPP_BACKEND, reason="compares native and Python event managers") +def test_native_event_manager_unknown_raw_keys_match_python_noop_behavior(): + python_manager = KVCacheEventManager(max_kv_event_entries=4) + native_manager = NativeKVCacheEventManager(max_kv_event_entries=4) + + for manager, diff_type in ( + (python_manager, KVCacheEventDiff), + (native_manager, NativeKVCacheEventDiff), + ): + manager.add_removed_event([b"short", b"still-not-a-radix-key"]) + manager.add_removed_life_cycle_event(b"short", 0) + manager.add_updated_event(b"short") + manager.add_updated_event(b"short", cache_level=diff_type(0, 1)) + + assert _flush_contract_events(native_manager) == _flush_contract_events(python_manager) == [] + + +@pytest.mark.skipif(not _USING_CPP_BACKEND, reason="compares native and Python event managers") +def test_native_event_manager_disabled_queue_and_blocking_read_match_python_reference(): + for manager_type in (KVCacheEventManager, NativeKVCacheEventManager): + disabled = manager_type(max_kv_event_entries=0) + disabled.add_created_event([1]) + disabled.flush_iteration_events() + assert disabled.get_latest_events(0) == [] + + manager = manager_type(max_kv_event_entries=1) + result = [] + reader = threading.Thread(target=lambda: result.extend(manager.get_latest_events())) + reader.start() + time.sleep(0.05) + assert reader.is_alive() + manager.add_created_event([1]) + manager.flush_iteration_events() + reader.join(timeout=1) + assert not reader.is_alive() + assert len(result) == 1 + + +@pytest.mark.skipif(not _USING_CPP_BACKEND, reason="compares native and Python event managers") +def test_native_event_manager_constructor_errors_match_python_reference(): + for manager_type in (KVCacheEventManager, NativeKVCacheEventManager): + with pytest.raises(ValueError, match="Unsupported V2 KV cache event hash algorithm"): + manager_type(max_kv_event_entries=1, hash_algo="unsupported") + + with pytest.raises(TypeError): + NativeKVCacheUpdatedData("block") + + def test_v2_kv_cache_event_manager_serialization(): event_manager = KVCacheEventManager(max_kv_event_entries=4, window_size=128) event_manager.add_created_event([2, 3]) @@ -587,7 +859,7 @@ def test_v1_and_v2_managers_emit_same_v1_hash_stored_events(): manager_v1.flush_iteration_events() v1_events = KVCacheEventSerializer.serialize(manager_v1.get_latest_events(10)) - event_manager_v2 = KVCacheEventManager( + event_manager_v2 = NativeKVCacheEventManager( max_kv_event_entries=event_buffer_max_size, window_size=max_seq_len, hash_algo=KV_CACHE_HASH_ALGO_V1, @@ -907,7 +1179,7 @@ def test_v2_stored_events_match_block_hash_chain(): gc.collect() gc.disable() - event_manager = KVCacheEventManager(max_kv_event_entries=16, window_size=128) + event_manager = NativeKVCacheEventManager(max_kv_event_entries=16, window_size=128) manager = None try: tokens_per_block = 4 @@ -922,12 +1194,16 @@ def test_v2_stored_events_match_block_hash_chain(): stored_events = _stored_events(events) assert len(stored_events) == 1 + # Both the Python and C++ backends hash blocks with SHA-256 over the same + # little-endian token encoding, so the block-key chain matches the + # pure-Python Block.make_key implementation on either backend. root_key = RootBlock.make_key(ReuseScope()) block0_key = Block.make_key(root_key, tokens[:tokens_per_block]) block1_key = Block.make_key(block0_key, tokens[tokens_per_block:]) expected_hashes = [block0_key.hex(), block1_key.hex()] assert _stored_block_hashes(stored_events) == expected_hashes + assert stored_events[0]["hash_algo"] == "v2_sha256" assert stored_events[0]["data"]["parent_hash"] is None assert [ block["block_hash"] for block in stored_events[0]["data"]["blocks"] @@ -938,6 +1214,15 @@ def test_v2_stored_events_match_block_hash_chain(): assert [ token["token_id"] for token in stored_events[0]["data"]["blocks"][1]["tokens"] ] == list(range(tokens_per_block, 2 * tokens_per_block)) + + if _USING_CPP_BACKEND: + event_manager.add_updated_event( + bytes.fromhex(expected_hashes[0]), + cache_level=NativeKVCacheEventDiff(old_value=0, new_value=1), + layer_group_id=0, + ) + updated_events = _flush_serialized_events(event_manager) + assert updated_events[0]["data"]["block_hash"] == expected_hashes[0] finally: gc.enable() if manager is not None: @@ -950,7 +1235,7 @@ def test_v2_v1_hash_events_include_cache_salt_from_kv_cache(): gc.collect() gc.disable() - event_manager = KVCacheEventManager( + event_manager = NativeKVCacheEventManager( max_kv_event_entries=16, window_size=128, hash_algo=KV_CACHE_HASH_ALGO_V1, @@ -986,7 +1271,7 @@ def test_v2_reused_prefix_does_not_emit_duplicate_stored_events(): gc.collect() gc.disable() - event_manager = KVCacheEventManager(max_kv_event_entries=16, window_size=128) + event_manager = NativeKVCacheEventManager(max_kv_event_entries=16, window_size=128) manager = None try: tokens_per_block = 4 @@ -1010,12 +1295,14 @@ def test_v2_reused_prefix_does_not_emit_duplicate_stored_events(): reuse_events = _flush_serialized_events(event_manager) reused_hashes = _stored_block_hashes(reuse_events) + # SHA-256 block-key chain is identical across both backends. root_key = RootBlock.make_key(ReuseScope()) block0_key = Block.make_key(root_key, prefix_tokens[:tokens_per_block]) block1_key = Block.make_key(block0_key, prefix_tokens[tokens_per_block:]) block2_key = Block.make_key(block1_key, new_tokens) + prefix_keys = [block0_key, block1_key] - assert prefix_hashes == [block0_key.hex(), block1_key.hex()] + assert prefix_hashes == [block_key.hex() for block_key in prefix_keys] assert reused_hashes == [block2_key.hex()] assert not (set(prefix_hashes) & set(reused_hashes)) finally: @@ -1030,7 +1317,7 @@ def test_v2_removed_events_match_stored_hashes(): gc.collect() gc.disable() - event_manager = KVCacheEventManager(max_kv_event_entries=16, window_size=128) + event_manager = NativeKVCacheEventManager(max_kv_event_entries=16, window_size=128) manager = None try: tokens_per_block = 4 @@ -1075,7 +1362,7 @@ def test_v2_removed_event_emitted_when_last_level_page_is_dropped(): tokens_per_block = 8 window_size = 8 num_blocks = 4 - event_manager = KVCacheEventManager(max_kv_event_entries=16, window_size=window_size) + event_manager = NativeKVCacheEventManager(max_kv_event_entries=16, window_size=window_size) manager = None try: manager = _create_test_manager( @@ -1133,17 +1420,23 @@ def test_v2_kv_cache_event_manager_emits_updated_on_level_migration(): # a small window is used: it keeps the min_slots floor low enough that the # GPU level can be shrunk below the committed block count, forcing surplus # reusable pages to migrate to host. + # + # More committed blocks and a deeper shrink than the removed-event test are + # used here on purpose: the GPU eviction/migration walks the two attention + # pool groups in insertion order, so the surplus must exceed the first pool + # group's block count for the second pool group to also migrate — only then + # do both layer groups emit an "updated" (GPU -> host) event. tokens_per_block = 8 window_size = 8 - num_blocks = 4 - event_manager = KVCacheEventManager(max_kv_event_entries=16, window_size=window_size) + num_blocks = 8 + event_manager = NativeKVCacheEventManager(max_kv_event_entries=64, window_size=window_size) manager = None try: manager = _create_test_manager( event_manager, tokens_per_block=tokens_per_block, - gpu_quota=16 << 20, - host_quota=8 << 20, + gpu_quota=32 << 20, + host_quota=32 << 20, window_size=window_size, kv_buf_size=1 << 20, ) @@ -1163,7 +1456,7 @@ def test_v2_kv_cache_event_manager_emits_updated_on_level_migration(): del kv_cache gc.collect() - assert manager.resize(CacheLevel(0), 8 << 20) + assert manager.resize(CacheLevel(0), 12 << 20) event_manager.flush_iteration_events() events = KVCacheEventSerializer.serialize(event_manager.get_latest_events()) diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py index 469c66a61f3f..e5b17c6bece3 100755 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py @@ -80,7 +80,6 @@ round_up, temporary_sys_path, typed_range, - unwrap_rawref, ) else: from tensorrt_llm.runtime.kv_cache_manager_v2 import ( @@ -138,7 +137,6 @@ round_up, temporary_sys_path, typed_range, - unwrap_rawref, ) from copy import deepcopy @@ -151,6 +149,19 @@ def get_cached_cuda_event_type(): + backend = os.environ.get("TLLM_KV_CACHE_MANAGER_V2_BACKEND", "cpp").lower() + if backend == "cpp": + try: + from bindings.internal.batch_manager.kv_cache_manager_v2 import CachedCudaEvent + + return CachedCudaEvent + except ImportError: + from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2 import ( + CachedCudaEvent, + ) + + return CachedCudaEvent + if find_spec("kv_cache_manager_v2") is not None: from kv_cache_manager_v2._utils import CachedCudaEvent @@ -673,21 +684,18 @@ def test_commit_min_snapshot_reuses_swa_post_commit_prefix(self) -> None: kv1.close() stream_holder.take_finish_event().synchronize() - match = self.manager._radix_tree.match(ReuseScope(), prompt) - self.assertEqual(match.num_tokens, len(prompt)) - self.assertEqual(len(match.blocks), 4) - - swa_lc_id = next( - lc_id - for lc_id, lc in self.manager._life_cycles.attention_life_cycles() - if lc.window_size is not None + swa_lc_id = _introspection.swa_life_cycle_ids(self.manager)[0] + num_tokens, pages = _introspection.reuse_match_pages( + self.manager, ReuseScope(), prompt, swa_lc_id ) + self.assertEqual(num_tokens, len(prompt)) + self.assertEqual(len(pages), 4) # The committed snapshot is reusable at the post-commit token count, but # old SWA blocks outside that window should not keep reusable pages. - self.assertIsNone(match.blocks[0].storage[swa_lc_id]) - self.assertIsNone(match.blocks[1].storage[swa_lc_id]) - self.assertIsNotNone(match.blocks[2].storage[swa_lc_id]) - self.assertIsNotNone(match.blocks[3].storage[swa_lc_id]) + self.assertIsNone(pages[0]) + self.assertIsNone(pages[1]) + self.assertIsNotNone(pages[2]) + self.assertIsNotNone(pages[3]) self.assertEqual( self.manager.probe_reuse(input_tokens=prompt[: tokens_per_block * 3]), 0, @@ -2034,8 +2042,8 @@ def test_ssm_snapshot_iteration_stats( reused.commit_pending_stats() self.assertEqual(self.manager.get_dirty_stats_kv_cache_ids(), set()) - assert self.manager._life_cycles.ssm_life_cycle_id is not None - ssm_life_cycle_id = self.manager._life_cycles.ssm_life_cycle_id + ssm_life_cycle_id = _introspection.ssm_life_cycle_id(self.manager) + assert ssm_life_cycle_id is not None snapshot_stats = self.manager.get_and_reset_ssm_snapshot_iteration_stats() self.assertEqual(set(snapshot_stats), {ssm_life_cycle_id}) stats = snapshot_stats[ssm_life_cycle_id] @@ -2279,24 +2287,20 @@ def test_ssm_planned_drop_includes_partial_swa_window(self) -> None: drop_handle = kv_cache.plan_committed_block_drop() self.assertIsNotNone(drop_handle) - match = self.manager._radix_tree.match( - ReuseScope(), prompt, self.manager.enable_partial_match - ) - self.assertEqual(match.num_tokens, len(prompt)) - attn_lc_id = next(iter(self.manager._life_cycles.attention_life_cycles()))[0] - assert self.manager._life_cycles.ssm_life_cycle_id is not None - ssm_lc_id = self.manager._life_cycles.ssm_life_cycle_id - planned_pages = [] - for block in match.blocks: - page_ref = block.storage[attn_lc_id] - assert page_ref is not None - planned_pages.append(unwrap_rawref(page_ref)) - ssm_page_ref = match.blocks[-1].storage[ssm_lc_id] - assert ssm_page_ref is not None - planned_pages.append(unwrap_rawref(ssm_page_ref)) - self.assertTrue(all(page.planned_drop_count == 1 for page in planned_pages)) - planned_pages.clear() - del page_ref, ssm_page_ref + attn_lc_id = _introspection.attention_life_cycle_ids(self.manager)[0] + ssm_lc_id = _introspection.ssm_life_cycle_id(self.manager) + assert ssm_lc_id is not None + num_tokens, attn_counts = _introspection.reuse_match_planned_drop_counts( + self.manager, ReuseScope(), prompt, attn_lc_id, self.manager.enable_partial_match + ) + self.assertEqual(num_tokens, len(prompt)) + # Every partial SWA-window attention page is planned for drop exactly once. + self.assertTrue(attn_counts and all(count == 1 for count in attn_counts)) + _, ssm_counts = _introspection.reuse_match_planned_drop_counts( + self.manager, ReuseScope(), prompt, ssm_lc_id, self.manager.enable_partial_match + ) + # The SSM snapshot on the last committed block is planned for drop. + self.assertEqual(ssm_counts[-1], 1) kv_cache.close() assert drop_handle is not None @@ -2434,17 +2438,21 @@ def test_ssm_partial_snapshot_respects_partial_reuse_setting(self) -> None: exact.resume(stream) exact.close() - match = self.manager._radix_tree.match( - ReuseScope(), prompt[:48], self.manager.enable_partial_match + ssm_lc_id = _introspection.ssm_life_cycle_id(self.manager) + assert ssm_lc_id is not None + num_tokens, pages = _introspection.reuse_match_pages( + self.manager, + ReuseScope(), + prompt[:48], + ssm_lc_id, + self.manager.enable_partial_match, ) - self.assertEqual(match.num_tokens, 48) - assert self.manager._life_cycles.ssm_life_cycle_id is not None - page_ref = match.blocks[-1].storage[self.manager._life_cycles.ssm_life_cycle_id] - assert page_ref is not None - page = unwrap_rawref(page_ref) - self.assertEqual(page.num_tokens_in_block, 16) - - del exact, kv1, longer, match, page, page_ref + self.assertEqual(num_tokens, 48) + last_page = pages[-1] + assert last_page is not None + self.assertEqual(last_page[1], 16) + + del exact, kv1, longer gc.collect() stream_holder.synchronize() self.manager.shutdown() @@ -2464,30 +2472,30 @@ def test_commit_is_end_moves_partial_attention_and_ssm_pages(self) -> None: kv_cache.capacity = len(prompt) kv_cache.history_length = len(prompt) - attn_lc_id = next(iter(self.manager._life_cycles.attention_life_cycles()))[0] - assert self.manager._life_cycles.ssm_life_cycle_id is not None - ssm_lc_id = self.manager._life_cycles.ssm_life_cycle_id + attn_lc_id = _introspection.attention_life_cycle_ids(self.manager)[0] + ssm_lc_id = _introspection.ssm_life_cycle_id(self.manager) + assert ssm_lc_id is not None attn_tail_slot = kv_cache.get_base_page_indices(LayerGroupId(attn_lc_id))[1] ssm_slot = kv_cache.get_ssm_block_base_index(LayerGroupId(ssm_lc_id)) kv_cache.commit(prompt, is_end=True) kv_cache.close() - match = self.manager._radix_tree.match( - ReuseScope(), prompt, self.manager.enable_partial_match + _, attn_pages = _introspection.reuse_match_pages( + self.manager, ReuseScope(), prompt, attn_lc_id, self.manager.enable_partial_match ) - self.assertEqual(match.num_tokens, len(prompt)) - tree_block = match.blocks[-1] + num_tokens, ssm_pages = _introspection.reuse_match_pages( + self.manager, ReuseScope(), prompt, ssm_lc_id, self.manager.enable_partial_match + ) + self.assertEqual(num_tokens, len(prompt)) - attn_ref = tree_block.storage[attn_lc_id] - ssm_ref = tree_block.storage[ssm_lc_id] - assert attn_ref is not None - assert ssm_ref is not None - attn_page = unwrap_rawref(attn_ref) - ssm_page = unwrap_rawref(ssm_ref) - self.assertEqual(attn_page.slot_id, attn_tail_slot) - self.assertEqual(ssm_page.slot_id, ssm_slot) - self.assertEqual(ssm_page.num_tokens_in_block, 16) + attn_page = attn_pages[-1] + ssm_page = ssm_pages[-1] + assert attn_page is not None + assert ssm_page is not None + self.assertEqual(attn_page[0], attn_tail_slot) + self.assertEqual(ssm_page[0], ssm_slot) + self.assertEqual(ssm_page[1], 16) def test_commit_min_snapshot_requires_history_alignment(self) -> None: """commit_min_snapshot requires commit() to start or end at history length.""" @@ -2767,9 +2775,9 @@ def test_typical_step_long_sequences(self): def test_zero_capacity_request_reserves_only_an_ssm_slot(self): """Every request reserves one SSM slot, including a zero-token dummy.""" manager = KVCacheManager(self._make_hybrid_config()) - ssm_lc = manager._life_cycles.ssm_life_cycle_id + ssm_lc = _introspection.ssm_life_cycle_id(manager) assert ssm_lc is not None - ssm_pg = manager._storage.get_pool_group_index(ssm_lc) + ssm_pg = _introspection.pool_group_index(manager, ssm_lc) attn_pg = 1 - ssm_pg batch = BatchDesc( @@ -2778,7 +2786,7 @@ def test_zero_capacity_request_reserves_only_an_ssm_slot(self): KVCacheDesc(capacity=0, history_length=0), ] ) - slots = manager._storage._compute_slots_for_batch(batch, self.TOKENS_PER_BLOCK, None) + slots = _introspection.compute_slots_for_batch(manager, batch, self.TOKENS_PER_BLOCK, None) self.assertEqual(slots[ssm_pg], 2) self.assertEqual(slots[attn_pg], 2) manager.shutdown() @@ -2825,8 +2833,16 @@ def test_constraint_reserves_resume_headroom(self): kv_cache.close() manager.shutdown() - def test_initial_pool_ratio_overrides_typical_step_and_constraints(self): - """Explicit initial_pool_ratio takes precedence over inferred sizing inputs.""" + def test_constraint_floor_overrides_infeasible_initial_pool_ratio(self): + """A constraint's feasibility floor overrides an infeasible initial_pool_ratio. + + initial_pool_ratio is the target split and still overrides typical_step, but + constraints stay feasibility floors (mirrors PR #16269): if a declared batch + needs more slots than its target share can hold, that pool group's share is + clamped up so the batch can be resumed, rather than starving it during warmup. + Here pool group 1's 0.2 target cannot satisfy the 256-request constraint, so + its share is clamped above 0.2 and pool group 0 gives up the remainder. + """ typical = BatchDesc(kv_caches=[KVCacheDesc(capacity=4096, history_length=4000)] * 32) constraint = BatchDesc(kv_caches=[KVCacheDesc(capacity=256, history_length=128)] * 256) cfg = self._make_config( @@ -2837,7 +2853,8 @@ def test_initial_pool_ratio_overrides_typical_step_and_constraints(self): manager = KVCacheManager(cfg) ratio = _introspection.current_gpu_ratio(manager) - self.assertGreater(ratio[0], ratio[1]) + self.assertGreater(ratio[1], 0.2) + self.assertLess(ratio[0], 0.8) self.assertAlmostEqual(sum(ratio), 1.0, places=6) manager.shutdown() @@ -3705,11 +3722,7 @@ def test_reuse_across_prefill_turns_keeps_only_window_minus_one(self) -> None: tokens_per_block=tokens_per_block, gpu_quota=64 << 20, ) - swa_lc_id = next( - lc_id - for lc_id, lc in self.manager._life_cycles.attention_life_cycles() - if lc.window_size is not None - ) + swa_lc_id = _introspection.swa_life_cycle_ids(self.manager)[0] prompt1 = [TokenId(i) for i in range(2 * window_size - 1)] # 127 tokens prompt2 = [TokenId(i) for i in range(200)] # first 127 tokens identical to prompt1 @@ -3732,11 +3745,12 @@ def test_reuse_across_prefill_turns_keeps_only_window_minus_one(self) -> None: # match() is documented volatile, so read what we need and drop the reference before # mutating the tree again (holding live Block refs across a later clear would leave the # eviction accounting inconsistent). - match1 = self.manager._radix_tree.match(ReuseScope(), prompt1) - self.assertEqual(match1.num_tokens, len(prompt1)) - self.assertEqual(len(match1.blocks), 4) - has_page = [b.storage[swa_lc_id] is not None for b in match1.blocks] - del match1 + num_tokens1, pages1 = _introspection.reuse_match_pages( + self.manager, ReuseScope(), prompt1, swa_lc_id + ) + self.assertEqual(num_tokens1, len(prompt1)) + self.assertEqual(len(pages1), 4) + has_page = [p is not None for p in pages1] self.assertEqual( has_page, [False, False, True, True], # positions 0..63 out of window; 64..126 in window diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_salting.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_salting.py index c6d1d4e76a43..4c2dc9615686 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_salting.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_salting.py @@ -20,19 +20,15 @@ from typing import TYPE_CHECKING, cast if not TYPE_CHECKING and find_spec("kv_cache_manager_v2") is not None: - from kv_cache_manager_v2 import ReuseScope, TokenId - from kv_cache_manager_v2._block_radix_tree import ( - Block, - BlockRadixTree, - sequence_to_blockchain_keys, - ) + from kv_cache_manager_v2 import TokenId, sequence_to_blockchain_keys + from kv_cache_manager_v2._block_radix_tree import Block, BlockRadixTree, ReuseScope from kv_cache_manager_v2._life_cycle_registry import LifeCycleRegistry else: - from tensorrt_llm.runtime.kv_cache_manager_v2 import ReuseScope, TokenId + from tensorrt_llm.runtime.kv_cache_manager_v2 import TokenId, sequence_to_blockchain_keys from tensorrt_llm.runtime.kv_cache_manager_v2._block_radix_tree import ( Block, BlockRadixTree, - sequence_to_blockchain_keys, + ReuseScope, ) from tensorrt_llm.runtime.kv_cache_manager_v2._life_cycle_registry import LifeCycleRegistry @@ -49,7 +45,11 @@ def attention_life_cycles(self) -> Iterator[tuple[object, object]]: class TestReuseScope(unittest.TestCase): - def test_to_bytes_distinguishes_scope_fields(self) -> None: + def test_reuse_scope_seeds_distinct_keys(self) -> None: + # Distinct reuse scopes -- including the None-vs-0 cases for each field -- + # must seed distinct radix-tree keys. The first blockchain key is the + # root (the reuse-scope digest), so hashing the same tokens under each + # scope isolates the scope's contribution to the key. scopes = [ ReuseScope(), ReuseScope(lora_id=0), @@ -57,11 +57,15 @@ def test_to_bytes_distinguishes_scope_fields(self) -> None: ReuseScope(lora_id=0, salt=0), ReuseScope(lora_id=7, salt=11), ] + tokens = [TokenId(1), TokenId(2)] + + def root_key(scope: "ReuseScope") -> bytes: + return next(iter(sequence_to_blockchain_keys(2, scope, tokens)))[1] - serialized = [scope.to_bytes() for scope in scopes] - self.assertEqual(len(set(serialized)), len(scopes)) - self.assertEqual(serialized[0], b"\x00") - self.assertEqual(serialized, [scope.to_bytes() for scope in scopes]) + roots = [root_key(scope) for scope in scopes] + self.assertEqual(len(set(roots)), len(scopes)) + # Deterministic across repeated derivation. + self.assertEqual(roots, [root_key(scope) for scope in scopes]) def test_blockchain_keys_are_seeded_by_reuse_scope(self) -> None: tokens = [TokenId(1), TokenId(2), TokenId(3), TokenId(4)] diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_api.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_api.py index ea25d28e541c..a74723c95e6e 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_api.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_api.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os import pytest import torch @@ -26,12 +27,18 @@ KVCacheManager, KVCacheManagerConfig, KVCacheStatsDelta, + PoolGroupPeakBlockStats, SsmSnapshotIterationStatsDelta, ) pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.fixture(scope="module", autouse=True) +def initialize_cuda_context() -> None: + torch.empty(1, device="cuda") + + def _make_config(*, enable_stats: bool = True) -> KVCacheManagerConfig: return KVCacheManagerConfig( tokens_per_block=4, @@ -79,6 +86,45 @@ def test_stats_delta_arithmetic() -> None: assert snapshot.iter_snapshot_hit_rate == 0.0 +def test_cpp_stats_types_are_native() -> None: + if os.environ.get("TLLM_KV_CACHE_MANAGER_V2_BACKEND", "cpp").lower() != "cpp": + pytest.skip("C++ backend only") + + from tensorrt_llm.bindings.internal.batch_manager import kv_cache_manager_v2 as cpp + + assert KVCacheStatsDelta is cpp.KVCacheStatsDelta + assert KVCacheIterationStatsDelta is cpp.KVCacheIterationStatsDelta + assert PoolGroupPeakBlockStats is cpp.PoolGroupPeakBlockStats + + stats = KVCacheStatsDelta(alloc_total_blocks=1, reused_blocks=2) + assert repr(stats) == ( + "KVCacheStatsDelta(alloc_total_blocks=1, alloc_new_blocks=0, " + "reused_blocks=2, missed_blocks=0)" + ) + peak = PoolGroupPeakBlockStats(available=3, unavailable=4, evictable=5) + assert peak == PoolGroupPeakBlockStats(3, 4, 5) + with pytest.raises(AttributeError): + peak.available = 6 + + +def test_manager_accepts_uint64_max_request_id() -> None: + manager = KVCacheManager(_make_config()) + cache = None + cuda_graph_dummy_request_id = (1 << 64) - 1 + try: + cache = manager.create_kv_cache(id=cuda_graph_dummy_request_id) + assert cache.id == cuda_graph_dummy_request_id + manager.mark_stats_dirty(cuda_graph_dummy_request_id) + assert manager.get_dirty_stats_kv_cache_ids() == {cuda_graph_dummy_request_id} + manager.mark_stats_excluded(cuda_graph_dummy_request_id) + assert manager.is_stats_excluded(cuda_graph_dummy_request_id) + assert manager.get_dirty_stats_kv_cache_ids() == set() + finally: + if cache is not None: + cache.close() + manager.shutdown() + + @pytest.mark.parametrize("enable_stats", [False, True]) def test_manager_stats_config_and_api(enable_stats: bool) -> None: manager = KVCacheManager(_make_config(enable_stats=enable_stats))