Skip to content

Conversation

@Tabrizian
Copy link
Member

@Tabrizian Tabrizian commented Nov 24, 2025

Minor changes to support KVCache reuse.

Summary by CodeRabbit

  • Refactor

    • Optimized K-V cache shape management for improved consistency and efficiency.
    • Enhanced prefill chunking logic to support additional caching scenarios.
  • New Features

    • Enabled block reuse functionality for K-cache configurations in sparse attention modes.

✏️ Tip: You can customize this high-level summary in your review settings.

Description

Test Coverage

PR Checklist

Please review the following before submitting your PR:

  • PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

  • PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

  • Test cases are provided for new code paths (see test instructions)

  • Any new dependencies have been scanned for license and vulnerabilities

  • CODEOWNERS updated if ownership changes

  • Documentation updated as needed

  • Update tava architecture diagram if there is a significant design change in PR.

  • The reviewers assigned automatically/manually are appropriate for the PR.

  • Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...

Provide a user friendly way for developers to interact with a Jenkins server.

Run /bot [-h|--help] to print this help message.

See details below for each supported subcommand.

run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]

Launch build/test pipelines. All previously running jobs will be killed.

--reuse-test (optional)pipeline-id (OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.

--disable-reuse-test (OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.

--disable-fail-fast (OPTIONAL) : Disable fail fast on build/tests/infra failures.

--skip-test (OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.

--stage-list "A10-PyTorch-1, xxx" (OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.

--gpu-type "A30, H100_PCIe" (OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.

--test-backend "pytorch, cpp" (OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.

--only-multi-gpu-test (OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.

--disable-multi-gpu-test (OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.

--add-multi-gpu-test (OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.

--post-merge (OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.

--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" (OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".

--detailed-log (OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.

--debug (OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in the stage-list parameter to access the appropriate container environment. Note: Does NOT update GitHub check status.

For guidance on mapping tests to stage names, see docs/source/reference/ci-overview.md
and the scripts/test_to_stage_mapping.py helper.

kill

kill

Kill all running builds associated with pull request.

skip

skip --comment COMMENT

Skip testing for latest commit on pull request. --comment "Reason for skipping build/test" is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

reuse-pipeline

reuse-pipeline

Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

@Tabrizian Tabrizian force-pushed the user/imant/indexerCacheReuse branch 2 times, most recently from 233e38a to 28ccd29 Compare November 24, 2025 18:02
@Tabrizian Tabrizian force-pushed the user/imant/indexerCacheReuse branch from 28ccd29 to f4fc826 Compare November 26, 2025 19:56
@Tabrizian Tabrizian force-pushed the user/imant/indexerCacheReuse branch from f4fc826 to b41a20b Compare November 26, 2025 19:59
@Tabrizian Tabrizian marked this pull request as ready for review November 26, 2025 20:00
@Tabrizian Tabrizian requested review from a team as code owners November 26, 2025 20:00
@Tabrizian
Copy link
Member Author

/bot run --disable-fail-fast

@Tabrizian Tabrizian changed the title [None][feat] Add support for KVCache reuse for DSAv32 [None][feat] Add support for KVCache reuse for DSv32 Nov 26, 2025
@coderabbitai
Copy link
Contributor

coderabbitai bot commented Nov 26, 2025

📝 Walkthrough

Walkthrough

KV cache shape construction in C++ was simplified by removing conditional branching on containsIndexerKCache to use a uniform shape. Concurrently, sparse attention handling in Python was broadened to support enable_context_mla_with_cached_kv, and NotImplementedError guards for block reuse were removed. Test configurations were updated to use tokens_per_block instead of enable_block_reuse.

Changes

Cohort / File(s) Change Summary
KV Cache Manager Shape Simplification
cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
Removed conditional cache shape construction for primary pools; now unconditionally applies {mNumPrimaryBlocks, pool.numLayers, mKVFactor, blockSize} shape regardless of containsIndexerKCache state.
Sparse DSA Attention Logic
tensorrt_llm/_torch/attention_backend/sparse/dsa.py
Expanded prefill chunk creation trigger to include enable_context_mla_with_cached_kv flag; updated KV gathering dependency to use same flag; removed NotImplementedError guards in two locations to permit block reuse configurations.
Test Configuration Updates
tests/integration/defs/accuracy/test_llm_api_pytorch.py
Replaced enable_block_reuse=False parameter with tokens_per_block=64 across multiple KvCacheConfig constructor calls in test blocks.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

  • Rationale for guard removal: Verify why NotImplementedError guards for block reuse are now safe to remove and confirm no unintended side effects.
  • Cache shape uniformity: Ensure the unconditional shape change in kvCacheManager.cpp doesn't break downstream logic that previously relied on conditional shapes for indexer K-cache.
  • Flag coordination: Confirm that enable_context_mla_with_cached_kv correctly replaces the previous split logic between prefill chunking and KV gathering decisions.

Pre-merge checks and finishing touches

❌ Failed checks (2 warnings)
Check name Status Explanation Resolution
Description check ⚠️ Warning The description is largely incomplete. While a brief summary is provided, the Description, Test Coverage, and most checklist items are unfilled or only minimally addressed. Fill in the Description section with details about what KVCache reuse for DSAv32 entails and why it's needed. Add specific test coverage details in the Test Coverage section.
Docstring Coverage ⚠️ Warning Docstring coverage is 16.67% which is insufficient. The required threshold is 80.00%. You can run @coderabbitai generate docstrings to improve docstring coverage.
✅ Passed checks (1 passed)
Check name Status Explanation
Title check ✅ Passed The title clearly summarizes the main change: adding KVCache reuse support for DSAv32, which aligns with the file modifications removing NotImplementedError guards and adjusting cache-related logic.
✨ Finishing touches
  • 📝 Generate docstrings
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Post copyable unit tests in a comment

Tip

📝 Customizable high-level summaries are now available in beta!

You can now customize how CodeRabbit generates the high-level summary in your pull requests — including its content, structure, tone, and formatting.

  • Provide your own instructions using the high_level_summary_instructions setting.
  • Format the summary however you like (bullet lists, tables, multi-section layouts, contributor stats, etc.).
  • Use high_level_summary_in_walkthrough to move the summary from the description to the walkthrough section.

Example instruction:

"Divide the high-level summary into five sections:

  1. 📝 Description — Summarize the main change in 50–60 words, explaining what was done.
  2. 📓 References — List relevant issues, discussions, documentation, or related PRs.
  3. 📦 Dependencies & Requirements — Mention any new/updated dependencies, environment variable changes, or configuration updates.
  4. 📊 Contributor Summary — Include a Markdown table showing contributions:
    | Contributor | Lines Added | Lines Removed | Files Changed |
  5. ✔️ Additional Notes — Add any extra reviewer context.
    Keep each section concise (under 200 words) and use bullet or numbered lists for clarity."

Note: This feature is currently in beta for Pro-tier users, and pricing will be announced later.


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.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

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

Actionable comments posted: 0

🧹 Nitpick comments (1)
tests/integration/defs/accuracy/test_llm_api_pytorch.py (1)

2420-2427: Explicit tokens_per_block=64 for DeepSeek‑V3.2 tests looks good; consider de‑duplicating the constant

Overriding KvCacheConfig’s default tokens_per_block (32) to 64 in the DeepSeek‑V3.2 tests is consistent with the PR goal of exercising the DSAv32 KV cache layout and reuse, and the patterns at Lines 2420/2426/2489/2548 are logically sound.

To avoid a scattered magic number and make future block‑size changes easier, consider factoring 64 into a single module‑ or class‑level constant (e.g., DEEPSEEK_V32_TOKENS_PER_BLOCK = 64) and using that in all three KvCacheConfig initializations for this model family.

Also applies to: 2489-2491, 2548-2550

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 356f67c and b41a20b.

📒 Files selected for processing (3)
  • cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp (1 hunks)
  • tensorrt_llm/_torch/attention_backend/sparse/dsa.py (2 hunks)
  • tests/integration/defs/accuracy/test_llm_api_pytorch.py (3 hunks)
🧰 Additional context used
📓 Path-based instructions (3)
**/*.{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)
Prefer const or constexpr variables over #define whenever possible, as the latter are not visible to the compiler
A variable that is not modified after its initialization should be declared as const
Except 0 (only used in comparison for checking signness/existence/emptiness) and nullptr, 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 empty for or while loop in a new line
The statement forming the body of a switch, while, do .. while or for statement shall be a compound statement (use brace-delimited statements)
If and else should 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., thisIsASubDir and thisIsAFilename.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
  • tensorrt_llm/_torch/attention_backend/sparse/dsa.py
  • tests/integration/defs/accuracy/test_llm_api_pytorch.py
**/*.py

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

**/*.py: The code developed for TensorRT-LLM should conform to Python 3.8+
Indent Python code with 4 spaces; do not use tabs
Always maintain the namespace when importing in Python, even if only one class or function from a module is used (e.g., use from package.subpackage import foo and then foo.SomeClass() instead of from package.subpackage.foo import SomeClass)
Python filenames should use snake_case (e.g., some_file.py)
Python class names should use PascalCase (e.g., class SomeClass)
Python function and method names should use snake_case (e.g., def my_awesome_function():)
Python local variable names should use snake_case, with prefix k for variable names that start with a number (e.g., k_99th_percentile = ...)
Python global variables should use upper snake_case with prefix G (e.g., G_MY_GLOBAL = ...)
Python constants should use upper snake_case (e.g., MY_CONSTANT = ...)
Avoid shadowing variables declared in an outer scope in Python
Initialize all externally visible members of a Python class in the constructor
For Python interfaces that may be used outside a file, prefer docstrings over comments
Python comments should be reserved for code within a function, or interfaces that are local to a file
Use Google style docstrings for Python classes and functions, which can be parsed by Sphinx
Python attributes and variables can be documented inline with type and description (e.g., self.x = 5 followed by """<type>: Description of 'x'""" )
Avoid using reflection in Python when functionality can be easily achieved without reflection
When using try-except blocks in Python, limit the except clause to the smallest set of specific errors possible instead of catching all exceptions
When using try-except blocks in Python to handle multiple possible variable types (duck-typing), keep the body of the try as small as possible and use the else block to implement the logic

Files:

  • tensorrt_llm/_torch/attention_backend/sparse/dsa.py
  • tests/integration/defs/accuracy/test_llm_api_pytorch.py
🧠 Learnings (10)
📓 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: 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.
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.
📚 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-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
  • tensorrt_llm/_torch/attention_backend/sparse/dsa.py
  • tests/integration/defs/accuracy/test_llm_api_pytorch.py
📚 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
  • tensorrt_llm/_torch/attention_backend/sparse/dsa.py
  • tests/integration/defs/accuracy/test_llm_api_pytorch.py
📚 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-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-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: There is a planned refactoring to move cache block bookkeeping utilities from BlockManager/WindowBlockManager into the GenerationRequest class itself to improve code organization and make responsibilities clearer.

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
📚 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:

  • tensorrt_llm/_torch/attention_backend/sparse/dsa.py
📚 Learning: 2025-08-26T09:37:10.463Z
Learnt from: jiaganc
Repo: NVIDIA/TensorRT-LLM PR: 7031
File: tensorrt_llm/bench/dataclasses/configuration.py:90-104
Timestamp: 2025-08-26T09:37:10.463Z
Learning: In TensorRT-LLM's bench configuration, the `get_pytorch_perf_config()` method returns `self.pytorch_config` which is a Dict[str, Any] that can contain default values including `cuda_graph_config`, making the fallback `llm_args["cuda_graph_config"]` safe to use.

Applied to files:

  • tests/integration/defs/accuracy/test_llm_api_pytorch.py
🧬 Code graph analysis (2)
tensorrt_llm/_torch/attention_backend/sparse/dsa.py (1)
tensorrt_llm/_torch/attention_backend/interface.py (2)
  • num_contexts (204-205)
  • num_contexts (208-211)
tests/integration/defs/accuracy/test_llm_api_pytorch.py (2)
tensorrt_llm/llmapi/llm_args.py (1)
  • KvCacheConfig (1430-1574)
tensorrt_llm/_torch/attention_backend/trtllm.py (1)
  • tokens_per_block (657-661)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
  • GitHub Check: Pre-commit Check
🔇 Additional comments (3)
tensorrt_llm/_torch/attention_backend/sparse/dsa.py (2)

933-942: LGTM - Condition correctly broadened for cached KV support.

The condition now creates indexer_prefill_chunks when either:

  1. Multiple chunk groups exist (original behavior), or
  2. enable_context_mla_with_cached_kv is enabled (new behavior)

This ensures that when KV cache reuse is active, the chunk metadata with proper cached token offsets (via prepare_one_prefill_chunk) is always used, rather than falling back to the single-pass path that doesn't account for cached tokens correctly.


1021-1023: LGTM - Simplified condition aligns with feature flag.

The condition for _need_full_kv_gathering is now directly tied to enable_context_mla_with_cached_kv, which is the appropriate flag for determining when cached KV tokens need to be included in the slot mappings. The updated comment accurately reflects that this covers both chunked prefill and KVCache reuse scenarios.

cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp (1)

878-879: I'm unable to verify the review comment as the repository cloning is failing consistently. Without access to the codebase, I cannot search for:

  1. How the indexer K-cache pool is actually used by consumers
  2. Whether Python's get_indexer_k_cache_buffers function handles the shape
  3. What C++ accessors consume the pool and if they correctly interpret the 4D shape
  4. Whether the shape change is actually breaking or neutral for these consumers

Since the review comment raises a legitimate concern about verifying consumer behavior after a structural change (3D→4D shape), but I cannot access the codebase to confirm whether this verification is necessary or if consumers handle it correctly, this requires manual verification.

Verify that indexer K-cache pool consumers handle the new 4D shape {mNumPrimaryBlocks, pool.numLayers, mKVFactor, blockSize} where mKVFactor=1 for K-only pools.

The shape construction is now uniform across all pools, including indexer K-cache pools which previously used a 3D shape. Ensure that all consumers of the indexer K-cache pool correctly handle the 4D shape after this change.

@tensorrt-cicd
Copy link
Collaborator

PR_Github #25884 [ run ] triggered by Bot. Commit: b41a20b

Copy link
Collaborator

@chang-l chang-l left a comment

Choose a reason for hiding this comment

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

Do you mind also updating the doc here:

- Support for KV Cache Reuse and Chunked Prefill in DeepSeek-V3.2-Exp is currently under development. When running `quickstart_advanced.py`, please include `--disable_kv_cache_reuse` to disable KV Cache Reuse. When using `trtllm-eval`/`trtllm-serve`/`trtllm-bench`, please include the following configuration in the extra llm_api options:

Signed-off-by: Iman Tabrizian <[email protected]>
@Tabrizian Tabrizian requested review from a team as code owners November 26, 2025 22:47
Copy link
Collaborator

@lfr-0531 lfr-0531 left a comment

Choose a reason for hiding this comment

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

LGTM!

Copy link
Collaborator

Choose a reason for hiding this comment

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

Could you also update the test_disaggregated_serving.py::TestDeepSeekV32Exp to enable the KV cache reuse?

@tensorrt-cicd
Copy link
Collaborator

PR_Github #25884 [ run ] completed with state SUCCESS. Commit: b41a20b
/LLM/main/L0_MergeRequest_PR pipeline #19629 completed with status: 'FAILURE'

@chang-l
Copy link
Collaborator

chang-l commented Nov 27, 2025

/bot run

@tensorrt-cicd
Copy link
Collaborator

PR_Github #26070 [ run ] triggered by Bot. Commit: cb48a3a

@tensorrt-cicd
Copy link
Collaborator

PR_Github #26070 [ run ] completed with state SUCCESS. Commit: cb48a3a
/LLM/main/L0_MergeRequest_PR pipeline #19794 completed with status: 'FAILURE'

Signed-off-by: Iman Tabrizian <[email protected]>
@lfr-0531
Copy link
Collaborator

/bot run

@tensorrt-cicd
Copy link
Collaborator

PR_Github #26134 [ run ] triggered by Bot. Commit: cb48a3a

@tensorrt-cicd
Copy link
Collaborator

PR_Github #26134 [ run ] completed with state SUCCESS. Commit: cb48a3a
/LLM/main/L0_MergeRequest_PR pipeline #19850 completed with status: 'FAILURE'

@lfr-0531
Copy link
Collaborator

/bot run

@tensorrt-cicd
Copy link
Collaborator

PR_Github #26222 [ run ] triggered by Bot. Commit: df52073

@tensorrt-cicd
Copy link
Collaborator

PR_Github #26222 [ run ] completed with state SUCCESS. Commit: df52073
/LLM/main/L0_MergeRequest_PR pipeline #19921 completed with status: 'FAILURE'

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants