[None][feat] Muse Spark Open modeling_v2 TEP target: a per-deployment lever resolver and the Vera Rubin high-concurrency, mid-concurrency and low-latency fast paths

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp

@@ -27,6 +27,7 @@
 #include "tensorrt_llm/common/assert.h"
 #include <algorithm>
 #include <cassert>
+#include <cstdlib>
 #include <cstring>
 #include <stdexcept>
 #include <utility>
@@ -696,9 +697,127 @@ BlockRadixTree::BlockRadixTree(
     : mLifeCycles(lifeCycles)
     , mTokensPerBlock(tokensPerBlock)
     , mEventSink(std::move(eventSink))
+    , mKeyChainMemo(tokensPerBlock)
 {
 }
 
+// ---------------------------------------------------------------------------
+// BlockKeyChainMemo
+// ---------------------------------------------------------------------------
+
+bool keyChainMemoEnabled() noexcept
+{
+    static bool const enabled = []()
+    {
+        char const* value = std::getenv("TLLM_KV_CACHE_MANAGER_V2_KEY_CHAIN_MEMO");
+        return value != nullptr && value[0] == '1' && value[1] == '\0';
+    }();
+    return enabled;
+}
+
+BlockKeyChainMemo::BlockKeyChainMemo(int tokensPerBlock, size_t capacity)
+    : mTokensPerBlock(static_cast<size_t>(tokensPerBlock))
+    , mCapacity(capacity)
+{
+    TLLM_CHECK(tokensPerBlock > 0 && capacity > 0);
+}
+
+namespace
+{
+//! Number of leading full blocks on which `a` and `b` (both at least `numBlocks` blocks long) agree.
+size_t commonFullBlocks(uint32_t const* a, uint32_t const* b, size_t numBlocks, size_t tokensPerBlock)
+{
+    // Compare in large chunks first (memcmp is vectorised), then find the first differing block.
+    constexpr size_t kBlocksPerChunk = 64;
+    size_t block = 0;
+    while (block < numBlocks)
+    {
+        size_t const n = std::min(kBlocksPerChunk, numBlocks - block);
+        size_t const offset = block * tokensPerBlock;
+        if (std::memcmp(a + offset, b + offset, n * tokensPerBlock * sizeof(uint32_t)) == 0)
+        {
+            block += n;
+            continue;
+        }
+        for (size_t i = 0; i < n; ++i)
+        {
+            size_t const o = (block + i) * tokensPerBlock;
+            if (std::memcmp(a + o, b + o, tokensPerBlock * sizeof(uint32_t)) != 0)
+            {
+                return block + i;
+            }
+        }
+        TLLM_CHECK_DEBUG_WITH_INFO(false, "chunk differs but no block does");
+        return block;
+    }
+    return numBlocks;
+}
+} // namespace
+
+BlockKeyChainMemo::Hit BlockKeyChainMemo::find(
+    BlockKey const& rootKey, TokenIdExt const* tokens, size_t numFullBlocks) const
+{
+    static_assert(sizeof(TokenIdExt) == sizeof(uint32_t));
+    std::vector<std::shared_ptr<Entry const>> entries;
+    {
+        std::lock_guard<std::mutex> const lock(mMutex);
+        entries = mEntries;
+    }
+    auto const* bits = reinterpret_cast<uint32_t const*>(tokens);
+    Hit best;
+    for (auto it = entries.rbegin(); it != entries.rend(); ++it)
+    {
+        Entry const& entry = **it;
+        if (entry.rootKey != rootKey)
+        {
+            continue;
+        }
+        size_t const limit = std::min(numFullBlocks, entry.keys.size());
+        if (limit <= best.numBlocks)
+        {
+            continue;
+        }
+        size_t const n = commonFullBlocks(bits, entry.tokenBits.data(), limit, mTokensPerBlock);
+        if (n > best.numBlocks)
+        {
+            best = Hit{*it, n};
+        }
+    }
+    return best;
+}
+
+void BlockKeyChainMemo::remember(
+    BlockKey const& rootKey, TokenIdExt const* tokens, Hit const& hit, std::vector<BlockKey> newKeys)
+{
+    if (newKeys.empty())
+    {
+        return;
+    }
+    auto entry = std::make_shared<Entry>();
+    entry->rootKey = rootKey;
+    size_t const numBlocks = hit.numBlocks + newKeys.size();
+    entry->keys.reserve(numBlocks);
+    if (hit.numBlocks > 0)
+    {
+        entry->keys.assign(hit.entry->keys.begin(), hit.entry->keys.begin() + static_cast<ptrdiff_t>(hit.numBlocks));
+    }
+    entry->keys.insert(entry->keys.end(), newKeys.begin(), newKeys.end());
+    auto const* bits = reinterpret_cast<uint32_t const*>(tokens);
+    entry->tokenBits.assign(bits, bits + numBlocks * mTokensPerBlock);
+
+    std::lock_guard<std::mutex> const lock(mMutex);
+    // The new chain extends the hit entry when it covered all of that entry: replace it.
+    if (hit.entry && hit.numBlocks == hit.entry->keys.size())
+    {
+        std::erase(mEntries, hit.entry);
+    }
+    if (mEntries.size() >= mCapacity)
+    {
+        mEntries.erase(mEntries.begin());
+    }
+    mEntries.push_back(std::move(entry));
+}
+
 BlockRadixTree::~BlockRadixTree()
 {
     // Detach blocks leaf-first in O(1) extra space. Dropping mRoots directly would instead
@@ -790,6 +909,12 @@ std::vector<BlockRadixTree::MatchResult> BlockRadixTree::matchTokenPath(
 {
     // Read-only: probeReuse() runs this under a shared lock, so probes may overlap. A root still
     // awaiting erasure is childless, so it yields no match.
+    if (keyChainMemoEnabled()
+        && (knownNoDigest
+            || std::none_of(tokens.begin(), tokens.end(), [](TokenIdExt const& t) { return t.isDigest(); })))
+    {
+        return matchTokenPathMemo(reuseScope, tokens, enablePartialMatch);
+    }
     std::vector<MatchResult> results;
 
     // Lazily compute one key per iteration — no wasted hashing on early miss.
@@ -837,6 +962,68 @@ std::vector<BlockRadixTree::MatchResult> BlockRadixTree::matchTokenPath(
     return results;
 }
 
+std::vector<BlockRadixTree::MatchResult> BlockRadixTree::matchTokenPathMemo(
+    ReuseScope const& reuseScope, TokenSpan tokens, bool enablePartialMatch) const
+{
+    // The same walk as matchTokenPath; only where each block key comes from differs. Keys of the
+    // memoised prefix are read from the memo, the rest are hashed exactly as the generator does
+    // (Hasher over the previous key and the block's tokens) and recorded afterwards.
+    std::vector<MatchResult> results;
+    BlockKey const rootKey = RootBlock::makeKey(reuseScope);
+    auto rootIt = mRoots.find(rootKey);
+    if (rootIt == mRoots.end())
+        return results;
+
+    size_t const tokensPerBlock = static_cast<size_t>(mTokensPerBlock);
+    size_t const numTokens = static_cast<size_t>(tokens.size());
+    auto const hit = mKeyChainMemo.find(rootKey, tokens.begin(), numTokens / tokensPerBlock);
+    std::vector<BlockKey> newKeys;
+
+    std::unordered_map<BlockKey, SharedPtr<Block>> const* currentNext = &rootIt->second->next;
+    HalfOpenRange<size_t> missedRange;
+    bool missed = false;
+    BlockKey digest = rootKey;
+    for (size_t ordinal = 0, beg = 0; beg < numTokens; ++ordinal, beg += tokensPerBlock)
+    {
+        size_t const end = std::min(beg + tokensPerBlock, numTokens);
+        if (ordinal < hit.numBlocks)
+        {
+            digest = hit.entry->keys[ordinal];
+        }
+        else
+        {
+            Hasher h;
+            h.update(digest);
+            h.update(tokens.begin() + beg, end - beg, /*knownNoDigest=*/true);
+            digest = h.digest();
+            if (end - beg == tokensPerBlock)
+            {
+                newKeys.push_back(digest);
+            }
+        }
+        auto blockIt = currentNext->find(digest);
+        if (blockIt == currentNext->end())
+        {
+            missedRange = {beg, end};
+            missed = true;
+            break;
+        }
+        Block* block = blockIt->second.get();
+        results.push_back({block, static_cast<int>(end - beg)});
+        currentNext = &block->next;
+    }
+    mKeyChainMemo.remember(rootKey, tokens.begin(), hit, std::move(newKeys));
+
+    if (missed && enablePartialMatch)
+    {
+        auto [best, bestMatch]
+            = findBestPartialMatchInNextNodes(*currentNext, tokens.begin() + missedRange.beg, missedRange.length());
+        if (best)
+            results.push_back({best, bestMatch});
+    }
+    return results;
+}
+
 std::vector<BlockRadixTree::MatchResult> BlockRadixTree::pruneMatch(
     std::vector<MatchResult> matched, std::optional<LifeCycleId> ssmLcId) const
 {

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h

@@ -31,6 +31,7 @@
 #include <cstdint>
 #include <iterator>
 #include <memory>
+#include <mutex>
 #include <optional>
 #include <unordered_map>
 #include <vector>
@@ -347,6 +348,54 @@ struct Block : NodeBase, EnableSharedFromThis<Block>
     std::shared_ptr<Digest const> mLastTokenDigest;
 };
 
+// ---------------------------------------------------------------------------
+// BlockKeyChainMemo — exact memo of recent block-key chains (opt-in:
+// TLLM_KV_CACHE_MANAGER_V2_KEY_CHAIN_MEMO=1, read once per process).
+//
+// A block key depends only on the reuse scope and the tokens up to the end of that block, so the
+// chain of a prompt is the same in every radix walk of that prompt. One admission walks it
+// several times (shortest-prefill-first probe, first-new-block probe, cache creation), each time
+// paying one SHA-256 per block. The memo keeps the full-block keys of the last few walked
+// sequences together with a copy of their tokens; a later walk reuses the keys of its longest
+// token-for-token equal prefix and hashes only the rest. Only digest-free (plain id) token
+// sequences are memoised, so equal token bits mean equal content. Thread-safe: probes walk the
+// tree concurrently under the manager's shared lock.
+// ---------------------------------------------------------------------------
+[[nodiscard]] bool keyChainMemoEnabled() noexcept;
+
+class BlockKeyChainMemo
+{
+public:
+    struct Entry
+    {
+        BlockKey rootKey;
+        std::vector<uint32_t> tokenBits; // exactly keys.size() full blocks
+        std::vector<BlockKey> keys;
+    };
+
+    struct Hit
+    {
+        std::shared_ptr<Entry const> entry;
+        size_t numBlocks = 0; // leading full blocks whose keys entry->keys supplies
+    };
+
+    explicit BlockKeyChainMemo(int tokensPerBlock, size_t capacity = 32);
+
+    //! The entry sharing the longest prefix of full blocks with `tokens` (digest-free) under
+    //! `rootKey`. numBlocks is 0 when nothing matches.
+    [[nodiscard]] Hit find(BlockKey const& rootKey, TokenIdExt const* tokens, size_t numFullBlocks) const;
+
+    //! Record the chain of `tokens`: the hit's leading keys followed by `newKeys`, the keys of the
+    //! next full blocks. A no-op when there is nothing new.
+    void remember(BlockKey const& rootKey, TokenIdExt const* tokens, Hit const& hit, std::vector<BlockKey> newKeys);
+
+private:
+    size_t mTokensPerBlock;
+    size_t mCapacity;
+    mutable std::mutex mMutex;
+    std::vector<std::shared_ptr<Entry const>> mEntries; // least recently recorded first
+};
+
 // ---------------------------------------------------------------------------
 // BlockRadixTree — the global cache index.
 // next: reuseScope digest → RootBlock.
@@ -433,6 +482,9 @@ class BlockRadixTree
     // knownNoDigest: from external text_only knowledge, never a scan (see Hasher::update).
     std::vector<MatchResult> matchTokenPath(
         ReuseScope const& reuseScope, TokenSpan tokens, bool knownNoDigest, bool enablePartialMatch) const;
+    // matchTokenPath for a digest-free sequence, with block keys taken from / recorded in mKeyChainMemo.
+    std::vector<MatchResult> matchTokenPathMemo(
+        ReuseScope const& reuseScope, TokenSpan tokens, bool enablePartialMatch) const;
     // Shorten `matched` to the prefix that is actually reusable. Passing
     // std::nullopt for `ssmLcId` skips the recurrent-snapshot constraint and
     // yields the attention-only prefix (used for numReusableTokensBeforeHybridPruning).
@@ -448,6 +500,8 @@ class BlockRadixTree
 
     std::unordered_map<BlockKey, SharedPtr<RootBlock>> mRoots;
     mutable std::vector<BlockKey> mPendingRootErases;
+    // Pure function of (scope, tokens), so it never goes stale as the tree changes.
+    mutable BlockKeyChainMemo mKeyChainMemo;
 };
 
 // ---------------------------------------------------------------------------

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.cpp

@@ -22,6 +22,7 @@
 #include <cerrno>
 #include <climits>
 #include <cmath>
+#include <cstdlib>
 #include <cstring>
 #include <fcntl.h>
 #include <limits>
@@ -86,14 +87,36 @@ SlotCount SlotAllocator::numOccupiedSlots() const noexcept
     return slotCountValueFromSize(mOccupiedMask.numSetBits());
 }
 
+namespace
+{
+//! TLLM_KV_CACHE_MANAGER_V2_LAZY_SLOT_SCRUB=1 (read once): scrub recycled-slot events once per
+//! allocateMultiple() instead of once per slot, and not on release(). Scrubbing only decides which
+//! free slot is handed out first (ready ones are preferred); every user of a slot still waits for
+//! its readyEvent, so this changes neither correctness nor contents, only how many cuEventQuery
+//! calls a migration of N pages issues (about 2 instead of 2N).
+bool lazySlotScrubEnabled() noexcept
+{
+    static bool const enabled = []()
+    {
+        char const* value = std::getenv("TLLM_KV_CACHE_MANAGER_V2_LAZY_SLOT_SCRUB");
+        return value != nullptr && value[0] == '1' && value[1] == '\0';
+    }();
+    return enabled;
+}
+} // namespace
+
 Slot SlotAllocator::allocate()
 {
     if (numFreeSlots() == 0)
     {
         throw OutOfPagesError("SlotAllocator: no free slots");
     }
     scrubEvents();
+    return takeFreeSlot();
+}
 
+Slot SlotAllocator::takeFreeSlot()
+{
     Slot slot;
     if (mNumReadyRecycledSlots > 0)
     {
@@ -130,6 +153,16 @@ std::vector<Slot> SlotAllocator::allocateMultiple(SlotCount numSlots)
     }
     std::vector<Slot> result;
     result.reserve(slotCountToSizeT(numSlots));
+    if (lazySlotScrubEnabled())
+    {
+        // numFreeSlots() >= numSlots was checked above, and each take consumes exactly one.
+        scrubEvents();
+        for (SlotCount slotCount{0}; slotCount < numSlots; ++slotCount)
+        {
+            result.push_back(takeFreeSlot());
+        }
+        return result;
+    }
     for (SlotCount slotCount{0}; slotCount < numSlots; ++slotCount)
     {
         result.push_back(allocate());
@@ -160,7 +193,10 @@ void SlotAllocator::release(Slot slot)
     {
         mOverflowSlots.push_back(std::move(slot));
     }
-    scrubEvents();
+    if (!lazySlotScrubEnabled())
+    {
+        scrubEvents();
+    }
     TLLM_CHECK_DEBUG(check());
 }
 

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.h

@@ -189,6 +189,9 @@ class SlotAllocator
 
 private:
     void scrubEvents();
+    // allocate() after its scrub: the next ready recycled slot, else a never-used slot, else the
+    // oldest recycled slot (its readyEvent still pending). Precondition: numFreeSlots() > 0.
+    Slot takeFreeSlot();
     [[nodiscard]] bool check() const noexcept;
 
     SlotCount mCapacity;