[None][fix] Fix MTP rewind corner case #9693
Conversation
|
/bot run |
📝 WalkthroughWalkthroughAdded a pre-emptive existence check in Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~8 minutes
Pre-merge checks and finishing touches❌ Failed checks (2 warnings)
✅ Passed checks (1 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp (1)
2845-2856: Guard against stale request IDs inrewindKVCacheis good; double‑check concurrency assumptionsThe early existence check under
mSequencesMtxcleanly handles the MTP overlap case whererewindKVCacheis invoked after the request has already been removed, avoiding an out‑of‑range onmSequences.at()viagetSequence()/removeToken.One thing to sanity‑check: if
removeSequence(requestId)can run concurrently withrewindKVCache(requestId, …), there’s still a narrow window between this pre‑check and the subsequentremoveToken()calls where the sequence could be extracted, andgetSequence()would stillat()a missing key. If the scheduler guarantees these two operations are not concurrent for the samerequestId, this implementation is sufficient; otherwise, a future refactor (e.g., splittingremoveTokeninto a locked helper and holdingmSequencesMtxacross the whole rewind) might be worth considering.
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp(1 hunks)
🧰 Additional context used
📓 Path-based instructions (2)
**/*.{cpp,h,cu}
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
**/*.{cpp,h,cu}: Closing braces of namespaces should have a comment saying the namespace it closes (e.g.,} // namespace foo)
Preferconstorconstexprvariables over#definewhenever possible, as the latter are not visible to the compiler
A variable that is not modified after its initialization should be declared asconst
Except0(only used in comparison for checking signness/existence/emptiness) andnullptr,true,false, all other literals should only be used for variable initialization and should be replaced with named constants
Use Allman indentation style for braces in C++
Put the semicolon for an emptyfororwhileloop in a new line
The statement forming the body of aswitch,while,do .. whileorforstatement shall be a compound statement (use brace-delimited statements)
Ifandelseshould always be followed by brace-delimited statements, even if empty or a single statement
C++ filenames should use camel case with first letter lowercase (e.g.,thisIsASubDirandthisIsAFilename.cpp)
All filenames involved in compilation of a compilation target must have case-insensitive unique filenames
All types (including class names) should use camel case with uppercase first letter (e.g.,FooBarClass)
Local variables, methods and namespaces should use camel case with first letter lowercase (e.g.,localFooBar)
Non-magic-number global variables that are non-static and not defined in anonymous namespace should use camel case prefixed by a lower case 'g' (e.g.,gDontUseGlobalFoos)
Non-magic-number global variables that are static or defined in an anonymous namespace should use camel case prefixed by a lower case 's' (e.g.,sMutableStaticGlobal)
Locally visible static variables should use camel case with lowercase prefix 's' as the first letter of the name (e.g.,static std::once_flag sFlag;)
Public, private and protected class member variables should use camel case prefixed with 'm' (e.g.,mNbFooValues), though the 'm' pre...
Files:
cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
**/*.{cpp,h,cu,py}
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
All TensorRT-LLM Open Source Software code files should contain an NVIDIA copyright header that includes the current year at the top
Files:
cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
🧠 Learnings (8)
📓 Common learnings
Learnt from: thorjohnsen
Repo: NVIDIA/TensorRT-LLM PR: 6910
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:0-0
Timestamp: 2025-08-14T21:04:50.248Z
Learning: In KV cache onboarding logic during prefill in cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp, when calculating which blocks fall within the attention window, use getTokensPerBlock() to advance token indices rather than block->getUniqueTokens().size(), because the calculation needs to consider the post-prefill state where blocks will be filled to capacity, not their current token count.
Learnt from: eopXD
Repo: NVIDIA/TensorRT-LLM PR: 6768
File: cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h:0-0
Timestamp: 2025-08-20T06:48:45.368Z
Learning: In cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp, updateSequenceCacheBlockOffsets is only called when adding a sequence, not during detach operations. During detach, the cache block bookkeeping is handled by GenerationRequest::removeFrontBlock.
Learnt from: eopXD
Repo: NVIDIA/TensorRT-LLM PR: 6768
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:2010-2045
Timestamp: 2025-08-21T09:41:49.347Z
Learning: In cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp, updateSequenceCacheBlockOffsets is specifically for updating bookkeeping when blocks are added during the context phase, not for refreshing offsets after detach operations. During detach operations, GenerationRequest::removeFrontBlock handles the necessary cache block bookkeeping internally.
Learnt from: eopXD
Repo: NVIDIA/TensorRT-LLM PR: 6768
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:577-579
Timestamp: 2025-08-20T06:56:02.889Z
Learning: In cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp, maxSequenceLength is now enforced as a non-optional argument in the BlockManager constructor, so concerns about std::nullopt defaulting to 0 are not applicable. When windowSize > maxSequenceLength, a warning should be added instead of handling optional parameter cases.
Learnt from: eopXD
Repo: NVIDIA/TensorRT-LLM PR: 6767
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:0-0
Timestamp: 2025-08-15T06:46:54.897Z
Learning: In cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp addToken function, newly allocated blocks are unshared by design. The beam search path in addToken (when sequence.getNumTokens() > windowSize) is currently broken/non-functional with SWA, so the block allocation doesn't follow a shared-then-unshared pattern.
📚 Learning: 2025-08-21T09:41:49.347Z
Learnt from: eopXD
Repo: NVIDIA/TensorRT-LLM PR: 6768
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:2010-2045
Timestamp: 2025-08-21T09:41:49.347Z
Learning: In cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp, updateSequenceCacheBlockOffsets is specifically for updating bookkeeping when blocks are added during the context phase, not for refreshing offsets after detach operations. During detach operations, GenerationRequest::removeFrontBlock handles the necessary cache block bookkeeping internally.
Applied to files:
cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
📚 Learning: 2025-08-20T06:48:45.368Z
Learnt from: eopXD
Repo: NVIDIA/TensorRT-LLM PR: 6768
File: cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h:0-0
Timestamp: 2025-08-20T06:48:45.368Z
Learning: In cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp, updateSequenceCacheBlockOffsets is only called when adding a sequence, not during detach operations. During detach, the cache block bookkeeping is handled by GenerationRequest::removeFrontBlock.
Applied to files:
cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
📚 Learning: 2025-08-15T06:46:54.897Z
Learnt from: eopXD
Repo: NVIDIA/TensorRT-LLM PR: 6767
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:0-0
Timestamp: 2025-08-15T06:46:54.897Z
Learning: In cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp addToken function, newly allocated blocks are unshared by design. The beam search path in addToken (when sequence.getNumTokens() > windowSize) is currently broken/non-functional with SWA, so the block allocation doesn't follow a shared-then-unshared pattern.
Applied to files:
cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
📚 Learning: 2025-08-20T06:56:02.889Z
Learnt from: eopXD
Repo: NVIDIA/TensorRT-LLM PR: 6768
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:577-579
Timestamp: 2025-08-20T06:56:02.889Z
Learning: In cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp, maxSequenceLength is now enforced as a non-optional argument in the BlockManager constructor, so concerns about std::nullopt defaulting to 0 are not applicable. When windowSize > maxSequenceLength, a warning should be added instead of handling optional parameter cases.
Applied to files:
cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
📚 Learning: 2025-08-14T21:04:50.248Z
Learnt from: thorjohnsen
Repo: NVIDIA/TensorRT-LLM PR: 6910
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:0-0
Timestamp: 2025-08-14T21:04:50.248Z
Learning: In KV cache onboarding logic during prefill in cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp, when calculating which blocks fall within the attention window, use getTokensPerBlock() to advance token indices rather than block->getUniqueTokens().size(), because the calculation needs to consider the post-prefill state where blocks will be filled to capacity, not their current token count.
Applied to files:
cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
📚 Learning: 2025-08-15T06:46:53.813Z
Learnt from: eopXD
Repo: NVIDIA/TensorRT-LLM PR: 6767
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:0-0
Timestamp: 2025-08-15T06:46:53.813Z
Learning: In the TensorRT-LLM KV cache manager, SWA (Sliding Window Attention) combined with beam search is currently in a broken/non-functional state and is planned for future rework. During preparatory refactoring phases, code related to SWA+beam search may intentionally remain in a non-working state until the broader rework is completed.
Applied to files:
cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
📚 Learning: 2025-08-06T08:18:28.669Z
Learnt from: zhengd-nv
Repo: NVIDIA/TensorRT-LLM PR: 6633
File: cpp/tensorrt_llm/batch_manager/dataTransceiverImpl.cpp:145-155
Timestamp: 2025-08-06T08:18:28.669Z
Learning: In cpp/tensorrt_llm/batch_manager/dataTransceiverImpl.cpp, the existing `mMtxForMap` mutex in DataSenderImpl is sufficient to synchronize measurement file operations in the `release` method, as all file operations occur within the same critical section that protects the `mRequestToSession` map access.
Applied to files:
cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
🧬 Code graph analysis (1)
cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp (1)
cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp (16)
requestId(108-111)requestId(108-108)requestId(113-117)requestId(113-114)requestId(119-124)requestId(119-121)requestId(126-130)requestId(126-127)requestId(132-135)requestId(132-132)requestId(137-140)requestId(137-137)requestId(169-172)requestId(169-169)requestId(190-194)requestId(190-191)
svc-trtllm-gh-bot
added
the
Community want to contribute
|
@lfr-0531 Could you please help the customer review this fix? Thanks you! |
Shixiaowei02
force-pushed
the
bugfix/mtp_with_async_scheduler
branch
from
88b27ce to
be1f481
Compare
|
/bot run |
Shixiaowei02
changed the title
|
PR_Github #29632 [ run ] triggered by Bot. Commit: |
|
PR_Github #29632 [ run ] completed with state |
eopXD
left a comment
eopXD
left a comment
There was a problem hiding this comment.
Slightly confused about the corner case here.
Rewind is only called if the generation request is not in a complete state. This should be a hard assertion rather than returning with only a warning.
# rewind kv cache
for request in scheduled_batch.generation_requests:
if request.state != LlmRequestState.GENERATION_COMPLETE:
if request.py_rewind_len > 0:
self.rewind_kv_cache(request, request.py_rewind_len)
May I ask what specific case have you encountered that requires the code change of this MR?
Hi, this corner case also happens when using TensorRT-LLM with FlexKV. In this case, when finishing kv cache offloading, the correspoding resources will be released. However, when combing with overlap scheduling + MTP, in the overlap step, if the MTP verification fails, a KV cache rewind may be triggered. At this point, the KV cache has already been released due to the earlier offloading, so the rewind operation attempts to access non-existent KV cache entries, which leads to an IndexError. |
Shixiaowei02
force-pushed
the
bugfix/mtp_with_async_scheduler
branch
from
be1f481 to
526b3ca
Compare
|
/bot run |
|
PR_Github #30972 [ run ] triggered by Bot. Commit: |
|
After discussion, we believe that the situation where KV offloading to storage overlaps with MTP is inherently conflicting. Could @jthomson04 help take a look or add some tests? Thanks! cc @pcastonguay |
|
PR_Github #30972 [ run ] completed with state
|
Shixiaowei02
force-pushed
the
bugfix/mtp_with_async_scheduler
branch
from
526b3ca to
cca24d7
Compare
|
/bot run |
1 similar comment
|
/bot run |
|
PR_Github #32838 [ run ] triggered by Bot. Commit: |
|
PR_Github #32839 [ run ] triggered by Bot. Commit: |
|
PR_Github #32838 [ run ] completed with state |
|
PR_Github #32839 [ run ] completed with state
|
|
@axxx03 Can you please run pre-commit and push the changes? |
|
/bot run --disable-fail-fast |
|
PR_Github #33211 [ run ] triggered by Bot. Commit: |
|
@axxx03 it's failing pre-commit check |
|
PR_Github #33211 [ run ] completed with state
|
|
@dagil-nvidia @jthomson04 fyi I opened #10941 to fix pre-commit and trigger CI so it can run overnight. |
|
Closing. #10941 has been merged. |
|
Sorry for the late reply. |
8 participants
Description
When using overlap scheduler with MTP, the request may have been terminated and removed from mSequences before rewindKVCache is called, we should avoid this.