Skip to content

[TRTLLM-11513][fix] eliminate dangling reference and TOCTOU races in KVCacheManager #12251

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
thorjohnsen wants to merge 17 commits into from
Closed

[TRTLLM-11513][fix] eliminate dangling reference and TOCTOU races in KVCacheManager #12251

Show file tree
Hide file tree
Changes from 15 commits
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
5a29d81
fix(kv_cache): eliminate dangling reference race in KVCacheManager::g…
thorjohnsen Mar 16, 2026
c3a9589
fix(kv_cache): update kvCacheUtils.h getSequence dot-chain to arrow
thorjohnsen Mar 16, 2026
177bfc2
fix(kv-cache): add per-sequence mutex to eliminate removeToken/remove…
thorjohnsen Mar 16, 2026
13aa0e1
style: add braces to all single-line if bodies in modified code
thorjohnsen Mar 16, 2026
f604c45
Resolve merge conflicts
thorjohnsen Mar 16, 2026
11acc93
precommit run
thorjohnsen Mar 16, 2026
e4773f4
fix: expose GenerationRequest mutex/removed state via public API
thorjohnsen Mar 17, 2026
11dfe67
fix: update nanobind trampoline getSequence return type to shared_ptr
thorjohnsen Mar 17, 2026
4b1f83d
Fix syntax error
thorjohnsen Mar 17, 2026
350b6e1
fix(test): pass GenerationRequest by const ref in assertBlocks lambda
thorjohnsen Mar 17, 2026
606a745
Merge branch 'main' into thorjohnsen/fix_kvcache_race_conditions
thorjohnsen Mar 17, 2026
d39d595
precommit run
thorjohnsen Mar 17, 2026
782ddc3
fix(kv-cache): hold per-sequence lock for full rewindKVCache loop
thorjohnsen Mar 17, 2026
5b1b93f
Merge branch 'main' into thorjohnsen/fix_kvcache_race_conditions
thorjohnsen Mar 17, 2026
5b3ce02
precommit run
thorjohnsen Mar 17, 2026
6529a63
Merge remote-tracking branch 'upstream/main' into thorjohnsen/fix_kvc…
thorjohnsen Mar 23, 2026
5753193
refactor: remove dead mFreeBlockIterator member from KVCacheBlock
thorjohnsen Mar 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 31 additions & 5 deletions cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@
#include <limits>
#include <list>
#include <memory>
#include <mutex>
#include <optional>
#include <set>
#include <unordered_map>
Expand Down Expand Up @@ -520,6 +521,30 @@ class GenerationRequest
std::set<KVCacheBlock::IdType> mUsedBlocks;
// Current prepopulated prompt length
SizeType32 mCurrentPrepopulatedPromptLen;
// Serialises per-sequence mutations (addToken, removeToken, storeNewBlock, etc.)
// against removeSequence to prevent use of freed blocks.
mutable std::mutex mMutex;
// Set by removeSequence under mMutex before releasing blocks; checked by all
// mutating callers so they can exit cleanly rather than touch freed state.
bool mRemoved{false};

public:
// Acquires mMutex; the returned unique_lock is moveable so callers hold RAII ownership.
[[nodiscard]] std::unique_lock<std::mutex> getLock() const
{
return std::unique_lock<std::mutex>(mMutex);
}

bool isRemoved() const noexcept
{
return mRemoved;
}

// Called by KVCacheManager::removeSequence while holding the lock.
void markRemoved() noexcept
{
mRemoved = true;
}
};

// attach metadata to a pool pointer
Expand Down Expand Up @@ -1608,8 +1633,7 @@ class BaseKVCacheManager
// void removeToken(SizeType32 seqSlotIdx);
virtual void rewindKVCache(LlmRequest::RequestIdType requestId, SizeType32 rewindLengths) = 0;

[[nodiscard]] virtual GenerationRequest const& getSequence(LlmRequest::RequestIdType requestId) const = 0;
[[nodiscard]] virtual GenerationRequest& getSequence(LlmRequest::RequestIdType requestId) = 0;
[[nodiscard]] virtual std::shared_ptr<GenerationRequest> getSequence(LlmRequest::RequestIdType requestId) const = 0;

[[nodiscard]] virtual bool isCrossKv() const = 0;

Expand Down Expand Up @@ -1982,8 +2006,7 @@ class KVCacheManager : public BaseKVCacheManager
void removeToken(LlmRequest::RequestIdType requestId);
void rewindKVCache(LlmRequest::RequestIdType requestId, SizeType32 rewindLengths) override;

[[nodiscard]] GenerationRequest const& getSequence(LlmRequest::RequestIdType requestId) const override;
[[nodiscard]] GenerationRequest& getSequence(LlmRequest::RequestIdType requestId) override;
[[nodiscard]] std::shared_ptr<GenerationRequest> getSequence(LlmRequest::RequestIdType requestId) const override;

[[nodiscard]] bool isCrossKv() const override
{
Expand Down Expand Up @@ -2104,6 +2127,9 @@ class KVCacheManager : public BaseKVCacheManager
SizeType32 sinkTokenLength, SizeType32 blockCapacity, SizeType32 beamWidth, SizeType32 tokensPerBlock);

private:
// Removes one token from sequence; caller must hold seqPtr->getLock().
void removeTokenLocked(GenerationRequest& sequence);

// Maximum number of sequences
SizeType32 mMaxNumSequences;
// Maximum beam width
Expand All @@ -2120,7 +2146,7 @@ class KVCacheManager : public BaseKVCacheManager
// Block manager
BlockManager mBlockManager;
// Map of all sequences
std::unordered_map<LlmRequest::RequestIdType, GenerationRequest> mSequences;
std::unordered_map<LlmRequest::RequestIdType, std::shared_ptr<GenerationRequest>> mSequences;
// Whether to cache KV pages for reuse
bool mEnableBlockReuse;
// Mutex to protect access to mSequences
Expand Down
2 changes: 1 addition & 1 deletion cpp/include/tensorrt_llm/batch_manager/kvCacheUtils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -200,7 +200,7 @@ class BlockRange
auto windowSize = cacheManager.getBlockManager().getPoolWindowSize(poolIdx);
mPoolsPerWindow[windowSize].push_back(cacheManager.getBlockManager().getPrimaryPool(poolIdx));
mBlockIdsPerWindow[windowSize]
= cacheManager.getSequence(mRequestId).getCacheBlockIds(windowSize).at(kFIRST_AND_ONLY_BEAM);
= cacheManager.getSequence(mRequestId)->getCacheBlockIds(windowSize).at(kFIRST_AND_ONLY_BEAM);
}
if (cacheManager.isEnableIndexerKCache())
{
Expand Down
166 changes: 112 additions & 54 deletions cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2270,7 +2270,9 @@ SizeType32 KVCacheManager::getNeededBlocksOneStep(
return 0;
}

auto const numCurrTokens = getSequence(req.mRequestId).getNumTokens();
auto const seqPtr_ = getSequence(req.mRequestId);

@eopXD eopXD Mar 19, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: Why the extra _ in the variable name?

@thorjohnsen thorjohnsen Mar 19, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It was generated like that by Claude.

TLLM_CHECK_WITH_INFO(seqPtr_, "No sequence found for request %lu", req.mRequestId);
auto const numCurrTokens = seqPtr_->getNumTokens();
auto const generatedTokens = numCurrTokens - req.getPromptLen();
auto const maxTokensToAddToKVCache = req.mMaxNewTokens - generatedTokens;
auto const tokensPerStep = req.getNumDraftTokens() + 1;
Expand Down Expand Up @@ -2322,7 +2324,7 @@ SizeType32 KVCacheManager::getRemainingBlocksToCompletion(LlmRequest const& req,
if (seqIt != mSequences.end())
{
auto const& seq = seqIt->second;
numAllocBlocksPerBeam = seq.getCacheBlockIds(windowSize).at(0).size();
numAllocBlocksPerBeam = seq->getCacheBlockIds(windowSize).at(0).size();
}
}

Expand Down Expand Up @@ -2426,7 +2428,17 @@ void BlockManager::updateCacheBlockOffsetsAtIdx(GenerationRequest& sequence, Siz
void KVCacheManager::addToken(RequestIdType requestId)
{
// TODO: add streamLLM support
auto& sequence = getSequence(requestId);
auto seqPtr = getSequence(requestId);
if (!seqPtr)
{
return;
}
auto seqLock = seqPtr->getLock();
if (seqPtr->isRemoved())
{
return;
}
auto& sequence = *seqPtr;

@eopXD eopXD Mar 19, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we shouldn't return silently here for the conditions above. If the conditions happen, it should be a logic bug in the higher level and we should not fail silently.

@thorjohnsen thorjohnsen Mar 19, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The issue here is that we have a separate thread that calls removeSequence, which can potentially remove the sequence before we can lock it. The silent return is necessary because the rest of the code will crash if the sequence has been removed. The higher level code should also perform the same check and take the sequence out of the active batch if it has been removed. Perhaps I should add a warning message saying the sequence has been removed?

@thorjohnsen thorjohnsen Mar 19, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am really not a fan of the design that has a separate thread decide if sequence should be removed or not. It near guarantees that you will have one thread preparing the sequence for another step while the other thread is removing it. The silent returns are there to prevent the server from crashing when this happens.

@eopXD eopXD Mar 19, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we should have a guard at the higher level caller then. Failing silently here is likely not a good convention.

sequence.addNewTokens(1);
mBlockManager.adjustBlocksIfNeeded(sequence);
}
Expand Down Expand Up @@ -2488,14 +2500,18 @@ void KVCacheManager::addSequence(
? llmRequest->getKvCacheRetentionConfig().value_or(executor::KvCacheRetentionConfig())
: executor::KvCacheRetentionConfig();

auto const [seqIt, emplaceDone] = [&]
auto seq = std::make_shared<GenerationRequest>(
requestId, inputLength, beamWidth, mBlockManager.getWindowSizesMetadata(), std::move(kvCacheRetentionConfig));
{
auto lck = std::scoped_lock(mSequencesMtx);
return mSequences.try_emplace(requestId, requestId, inputLength, beamWidth,
mBlockManager.getWindowSizesMetadata(), kvCacheRetentionConfig);
}();
TLLM_CHECK(emplaceDone);
auto& sequence = seqIt->second;
auto const [_, emplaceDone] = mSequences.emplace(requestId, seq);
if (!emplaceDone)
{

@eopXD eopXD Mar 19, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The assertion TLLM_CHECK(emplaceDone) is now a silent return. Any rationale for it?

@thorjohnsen thorjohnsen Mar 19, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Calling addSequence more than once for the same sequence should not be an error. There are some cases in the python executor were addSequence would be called more than once, if not for some book-keeping that keeps track of whether addSequence has already been called. KVCacheManager already knows if addSequence has been called for a given sequence, doing a silent return when emplaceDone is false makes it safe to call addSequence more than once for the same sequence. We can remove this book-keeping logic in executor with this change.

return;
}
}

auto& sequence = *seq;

// Get statistics for block allocations/reuse pre request.
SizeType32 const numAllocTotalBlocksPreRequest = mBlockManager.getNumAllocTotalBlocks();
Expand Down Expand Up @@ -2574,18 +2590,12 @@ void KVCacheManager::addSequence(
void KVCacheManager::storeContextBlocks(LlmRequest const& llmRequest)
{
auto const requestId = llmRequest.mRequestId;
bool found = false;
{
// protect the mSequences
std::scoped_lock lock(mSequencesMtx);
found = mSequences.find(requestId) != mSequences.end();
}
if (found)
auto seqPtr = getSequence(requestId);
if (seqPtr)
{
auto& sequence = getSequence(requestId);
if (mEnableBlockReuse && !llmRequest.isDummyRequest())
{
mBlockManager.storeContextBlocks(sequence, llmRequest);
mBlockManager.storeContextBlocks(*seqPtr, llmRequest);
}
}
else
Expand All @@ -2600,7 +2610,14 @@ void KVCacheManager::storeNewBlock(LlmRequest const& llmRequest)
// - Beam search is NOT enabled
// - Block reuse is enabled.
auto const requestId = llmRequest.mRequestId;
auto& sequence = getSequence(requestId);
auto seqPtr = getSequence(requestId);
TLLM_CHECK_WITH_INFO(seqPtr, "No sequence found for request %lu", requestId);
auto seqLock = seqPtr->getLock();
if (seqPtr->isRemoved())

@eopXD eopXD Mar 19, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should this also be a TLLM_CHECK too?

@thorjohnsen thorjohnsen Mar 19, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No. If the sequence was removed before we could lock it, we can't store the new block for reuse, which is a minor inconvenience, not a fatal error. The silent return is necessary to prevent a crash.

{
return;
}
auto& sequence = *seqPtr;
if (sequence.getBeamWidth() > 1 || !mEnableBlockReuse)
{
return;
Expand All @@ -2612,21 +2629,30 @@ std::optional<KVCacheBlock::IdType> KVCacheManager::removeSequence(
RequestIdType requestId, OptionalRef<LlmRequest const> llmRequest, bool pinBlocks)
{
TLLM_LOG_TRACE("[%s]::%s start", isCrossKv() ? "CROSS" : "SELF", __PRETTY_FUNCTION__);
auto sequenceNode = [this, requestId]
auto seqPtr = [this, requestId]() -> std::shared_ptr<GenerationRequest>
{
std::scoped_lock lock(mSequencesMtx);
return mSequences.extract(requestId);
auto it = mSequences.find(requestId);
if (it == mSequences.end())
{
return nullptr;
}
auto ptr = it->second;
mSequences.erase(it);
return ptr;
}();
Comment thread
coderabbitai[bot] marked this conversation as resolved.
std::optional<KVCacheBlock::IdType> lastStoredId = std::nullopt;
if (!sequenceNode.empty())
if (seqPtr)
{
auto seqLock = seqPtr->getLock();
seqPtr->markRemoved();
if (mEnableBlockReuse)
{
lastStoredId = mBlockManager.releaseBlocks(sequenceNode.mapped(), llmRequest, pinBlocks);
lastStoredId = mBlockManager.releaseBlocks(*seqPtr, llmRequest, pinBlocks);
}
else
{
lastStoredId = mBlockManager.releaseBlocks(sequenceNode.mapped(), std::nullopt, pinBlocks);
lastStoredId = mBlockManager.releaseBlocks(*seqPtr, std::nullopt, pinBlocks);
}
}
if (mBlockManager.isSequenceHeld(requestId))
Expand All @@ -2643,7 +2669,14 @@ std::vector<KVCacheBlock::IdType> KVCacheManager::storeBlocksForReuse(
RequestIdType requestId, OptionalRef<LlmRequest const> llmRequest, bool pinBlocks)
{
TLLM_LOG_TRACE("[%s]::%s start", isCrossKv() ? "CROSS" : "SELF", __PRETTY_FUNCTION__);
auto& sequence = getSequence(requestId);
auto seqPtr = getSequence(requestId);
TLLM_CHECK_WITH_INFO(seqPtr, "No sequence found for request %lu", requestId);
auto seqLock = seqPtr->getLock();
if (seqPtr->isRemoved())
{
return {};
}
auto& sequence = *seqPtr;
auto pinnedBlockIds = mBlockManager.storeBlocksForReuse(sequence, llmRequest, pinBlocks);
TLLM_LOG_TRACE("[%s]::%s stop", isCrossKv() ? "CROSS" : "SELF", __PRETTY_FUNCTION__);
return pinnedBlockIds;
Expand All @@ -2657,8 +2690,14 @@ void KVCacheManager::schedulingRemoveSequence(RequestIdType requestId)

void KVCacheManager::pinBlocks(RequestIdType requestId)
{
auto& sequence = getSequence(requestId);
mBlockManager.pinBlocks(sequence);
auto seqPtr = getSequence(requestId);
TLLM_CHECK_WITH_INFO(seqPtr, "No sequence found for request %lu", requestId);
auto seqLock = seqPtr->getLock();
if (seqPtr->isRemoved())
{
return;
}
mBlockManager.pinBlocks(*seqPtr);
}

void KVCacheManager::unpinBlocksById(std::vector<KVCacheBlock::IdType> const& blockIds)
Expand Down Expand Up @@ -2688,7 +2727,9 @@ tle::RetentionPriority KVCacheManager::getPriorityByBlockId(KVCacheBlock::IdType

SizeType32 KVCacheManager::copyBlockOffsets(ITensor& output, SizeType32 outputSlotOffset, RequestIdType requestId) const
{
auto const& sequence = getSequence(requestId);
auto const seqPtr = getSequence(requestId);
TLLM_CHECK_WITH_INFO(seqPtr, "No sequence found for request %lu", requestId);
auto const& sequence = *seqPtr;
auto const beamWidth = sequence.getBeamWidth();

auto* dstPtr = bufferCast<tk::KVCacheIndex>(output);
Expand Down Expand Up @@ -2998,10 +3039,8 @@ BlocksPerWindow BaseKVCacheManager::calculateMaxNumBlocks(executor::KvCacheConfi
return windowSizeToBlocks;
}

void KVCacheManager::removeToken(RequestIdType requestId)
void KVCacheManager::removeTokenLocked(GenerationRequest& sequence)
{
// TODO: add streamLLM support
auto& sequence = getSequence(requestId);
if (sequence.getNumTokens() == 0)

@eopXD eopXD Mar 19, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Add TLLM_CHECK to ensure lock acquired. Or encapsulate the lock acquisition and token removal loop in the same function to fool-proof this function to be mistakenly called.

@thorjohnsen thorjohnsen Mar 19, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this would defeat the purpose of this method. The method declaration states that it should not be called unless caller already owns the lock. It is a private method that is only used in some loops where the caller already owns the lock, calling std::mutex::lock again would cause a deadlock. There is no way to check lock status, calling try_lock when thread already owns the lock would cause undefined behavior.

@eopXD eopXD Mar 19, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think you mistaken a bit.
I was asking to encapsulate the locking and the loop of removeTokenLocked call into a single function, not adding another lock into the removeTokenLocked function.

{
return;
Expand All @@ -3018,36 +3057,48 @@ void KVCacheManager::removeToken(RequestIdType requestId)
}
}

void KVCacheManager::rewindKVCache(RequestIdType requestId, SizeType32 rewindLengths)
void KVCacheManager::removeToken(RequestIdType requestId)
{
// Check if the sequence still exists before rewinding
// In overlap mode with MTP, the request may have been terminated and removed
// from mSequences before rewindKVCache is called
// TODO: add streamLLM support
auto seqPtr = getSequence(requestId);
if (!seqPtr)
{
std::scoped_lock lck(mSequencesMtx);
if (mSequences.find(requestId) == mSequences.end())
{
TLLM_LOG_DEBUG("Request %lu has already been removed from KV cache manager, skipping rewind", requestId);
return;
}
return;
}

for (SizeType32 si = 0; si < rewindLengths; ++si)
auto seqLock = seqPtr->getLock();
if (seqPtr->isRemoved())
{
removeToken(requestId);
return;
}
removeTokenLocked(*seqPtr);
}

GenerationRequest const& KVCacheManager::getSequence(RequestIdType requestId) const
void KVCacheManager::rewindKVCache(RequestIdType requestId, SizeType32 rewindLengths)
{
auto lck = std::scoped_lock(mSequencesMtx);
return mSequences.at(requestId);
// Hold the per-sequence lock for the entire loop so removeSequence cannot
// slip in between iterations and cause a partial rewind.
auto seqPtr = getSequence(requestId);
if (!seqPtr)
{
TLLM_LOG_DEBUG("Request %lu has already been removed from KV cache manager, skipping rewind", requestId);
return;
}
auto seqLock = seqPtr->getLock();
if (seqPtr->isRemoved())
{
return;
}
for (SizeType32 si = 0; si < rewindLengths; ++si)
{
removeTokenLocked(*seqPtr);
}
}

GenerationRequest& KVCacheManager::getSequence(RequestIdType requestId)
std::shared_ptr<GenerationRequest> KVCacheManager::getSequence(RequestIdType requestId) const
{
auto lck = std::scoped_lock(mSequencesMtx);
return mSequences.at(requestId);
auto it = mSequences.find(requestId);
return it != mSequences.end() ? it->second : nullptr;
}

SizeType32 BaseKVCacheManager::getSinkBubbleLength(SizeType32 sinkTokenLen, SizeType32 tokensPerBlock)
Expand All @@ -3060,7 +3111,9 @@ SizeType32 BaseKVCacheManager::getSinkBubbleLength(SizeType32 sinkTokenLen, Size
std::vector<std::vector<SizeType32>> const& KVCacheManager::getCacheBlockIds(
RequestIdType requestId, SizeType32 windowSize) const
{
return getSequence(requestId).getCacheBlockIds(windowSize);
auto seqPtr = getSequence(requestId);
TLLM_CHECK_WITH_INFO(seqPtr, "No sequence found for request %lu", requestId);
return seqPtr->getCacheBlockIds(windowSize);
}
Comment on lines +3130 to 3133

@coderabbitai coderabbitai Bot Mar 16, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

getCacheBlockIds returns a dangling-prone reference under concurrent removal.

Line 3072 returns const& to data owned by GenerationRequest, but the function only keeps ownership via a local shared_ptr. After return, concurrent removeSequence can drop the final owner and invalidate the returned reference. This should return by value (or expose ownership with the returned view).

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp` around lines 3070 - 3073,
The function currently returns a const& from
seqPtr->getCacheBlockIds(windowSize) while holding only a local shared_ptr (from
getSequence), which allows removeSequence to drop the last owner and leave the
reference dangling; change the implementation and signature of getCacheBlockIds
to return by value (make a copy of the container while holding seqPtr) or
otherwise return an owning view (e.g., shared_ptr to the container) so the
returned object remains valid after the local shared_ptr is released; locate the
call site using getSequence, GenerationRequest, and removeSequence and ensure
callers accept the by-value return (or the new owning view) to avoid
use-after-free.

@eopXD eopXD Mar 19, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The concern looks legit.

@coderabbitai coderabbitai Bot Mar 19, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧩 Analysis chain

@eopXD, agreed — the concern is valid. Here's the fix: return by value so the copy is made while seqPtr is still alive, eliminating the dangling-reference window.

  • Change getCacheBlockIds return type from const& to by-value in both the header declaration and the .cpp definition
  • Verify callers to ensure they still compile (binding auto const& to a temporary is safe via lifetime extension; binding a non-const reference would be a compile error and must be fixed)

🏁 Script executed:

# Find the getCacheBlockIds declaration(s) in the header
rg -n "getCacheBlockIds" --type=cpp --type=h -C2

Length of output: 55400

🏁 Script executed:

# Find all callers of getCacheBlockIds on a KVCacheManager/BaseKVCacheManager to check binding patterns
rg -n "getCacheBlockIds" --type=cpp -C3

Length of output: 75259

@eopXD, here is the complete fix across all four affected files. The root issue is that seqPtr (the local shared_ptr) is destroyed when getCacheBlockIds returns, so the const& it hands back is immediately dangling. Returning by value makes a copy while seqPtr is still alive.

1. cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h — two signature changes

In BaseKVCacheManager (around Line 1671):

-    [[nodiscard]] virtual std::vector<std::vector<SizeType32>> const& getCacheBlockIds(
+    [[nodiscard]] virtual std::vector<std::vector<SizeType32>> getCacheBlockIds(
         LlmRequest::RequestIdType requestId, SizeType32 windowSize) const
         = 0;

In KVCacheManager (around Line 2075):

-    std::vector<std::vector<SizeType32>> const& getCacheBlockIds(
+    std::vector<std::vector<SizeType32>> getCacheBlockIds(
         LlmRequest::RequestIdType requestId, SizeType32 windowSize) const override;

2. cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp — definition (Line 3111)

-std::vector<std::vector<SizeType32>> const& KVCacheManager::getCacheBlockIds(
+std::vector<std::vector<SizeType32>> KVCacheManager::getCacheBlockIds(
     RequestIdType requestId, SizeType32 windowSize) const
 {
     auto seqPtr = getSequence(requestId);
     TLLM_CHECK_WITH_INFO(seqPtr, "No sequence found for request %lu", requestId);
-    return seqPtr->getCacheBlockIds(windowSize);
+    return seqPtr->getCacheBlockIds(windowSize);   // copy made while seqPtr is alive
 }

3. cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp — trampoline (Line 196)

-    std::vector<std::vector<SizeType32>> const& getCacheBlockIds(
+    std::vector<std::vector<SizeType32>> getCacheBlockIds(
         tb::LlmRequest::RequestIdType requestId, SizeType32 windowSize) const override
     {
         NB_OVERRIDE_PURE(getCacheBlockIds, requestId, windowSize);
     }

4. cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp — caller fix (Line 4228)

This one needs special attention. After the fix, getCacheBlockIds returns a temporary vector<vector<SizeType32>>. Indexing it with [0] yields a reference into that temporary, and C++ does not extend the temporary's lifetime when binding through a subscript — so auto const& allBlockIds would still dangle:

-    auto const& allBlockIds = kvCacheManager.getCacheBlockIds(requestId, maxAttentionWindow)[0];
+    auto allBlockIds = kvCacheManager.getCacheBlockIds(requestId, maxAttentionWindow)[0];

Failed to handle agent chat message. Please try again.

@thorjohnsen thorjohnsen Mar 19, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It is a legit concern.

@coderabbitai coderabbitai Bot Mar 19, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Seems like the humans are having a chat. I'll hop back into my burrow for now. If you need me again, just tag @coderabbitai in a new comment, and I'll come hopping out!


std::vector<std::vector<std::vector<SizeType32>>> KVCacheManager::getBatchCacheBlockIds(
Expand All @@ -3070,22 +3123,27 @@ std::vector<std::vector<std::vector<SizeType32>>> KVCacheManager::getBatchCacheB
result.reserve(requestIds.size());
for (auto const& requestId : requestIds)
{
auto const& sequence = getSequence(requestId);
result.emplace_back(sequence.getCacheBlockIds(windowSize));
auto const seqPtr = getSequence(requestId);
TLLM_CHECK_WITH_INFO(seqPtr, "No sequence found for request %lu", requestId);
result.emplace_back(seqPtr->getCacheBlockIds(windowSize));
}
return result;
}

std::optional<KVCacheBlock::IdType> KVCacheManager::getLastBlockId(LlmRequest::RequestIdType requestId) const
{
auto const& seq = getSequence(requestId);
auto const seq = getSequence(requestId);
if (!seq)
{
return std::nullopt;
}
// Use the first window size
auto firstWindowSize = mBlockManager.getFirstWindowSize();
if (firstWindowSize == 0)
{
return std::nullopt;
}
auto const& perBeam = seq.getCacheBlockIds(firstWindowSize);
auto const& perBeam = seq->getCacheBlockIds(firstWindowSize);
if (perBeam.empty() || perBeam[0].empty())
{
return std::nullopt;
Expand Down
2 changes: 1 addition & 1 deletion cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,7 +130,7 @@ class PyKvCacheManager : public tbk::BaseKVCacheManager
NB_OVERRIDE_PURE(storeBlocksForReuse, requestId, llmRequest, pinBlocks);
}

tbk::GenerationRequest const& getSequence(tb::LlmRequest::RequestIdType requestId) const override
std::shared_ptr<tbk::GenerationRequest> getSequence(tb::LlmRequest::RequestIdType requestId) const override
{
NB_OVERRIDE_PURE(getSequence, requestId);
}
Expand Down
Loading
Loading