[None][feat] DSv4: sparse cache manager adapter - #15394
Merged
longlee0622 merged 8 commits intoJun 25, 2026
Merged
Conversation
Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
jiaganc
force-pushed
the
user/fanrongl/dsv4-sparse-cache-manager
branch
from
June 24, 2026 15:06
8f6a51a to
3532b22
Compare
jiaganc
marked this pull request as ready for review
June 24, 2026 15:09
jiaganc
added
the
api-compatible
jiaganc
approved these changes
Jun 24, 2026
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
Contributor
📝 WalkthroughWalkthroughIntroduces ChangesDeepSeek-V4 sparse attention cache manager
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 7
🧹 Nitpick comments (4)
tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py (2)
44-50: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winFinish the required function annotations.
has_fp8_kv_cache, nested_add_layer(), andget_cache_size_per_token()are missing annotations. As per coding guidelines, “Always annotate functions. Make the return typeNoneif the function does not return anything.”Proposed fix
def _estimate_bytes_per_token( head_dim: int, index_head_dim: int, compress_ratios: List[int], - has_fp8_kv_cache, + has_fp8_kv_cache: bool, @@ def _add_layer( layer_idx: int, attn_type: DeepseekV4AttentionType, sliding_window_size: int | None - ): + ) -> None: @@ - def get_cache_size_per_token(model_config: ModelConfig, mapping: Mapping, **kwargs): + def get_cache_size_per_token(model_config: ModelConfig, mapping: Mapping, **kwargs) -> int:Also applies to: 369-371, 812-812
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py` around lines 44 - 50, The missing type annotations need to be completed in the cache manager helpers: add an explicit annotation for has_fp8_kv_cache in _estimate_bytes_per_token, annotate the nested _add_layer() helper, and add a return type to get_cache_size_per_token(). Update the relevant function signatures in cache_manager.py so every function is fully annotated, using None for helpers that do not return a value.Source: Coding guidelines
91-97: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMake the class-level attention set immutable.
Ruff flags this mutable class attribute. Since it is used as a constant, make it immutable and update the references.
Proposed fix
- fixed_size_attention = { + FIXED_SIZE_ATTENTION = frozenset({ DeepseekV4AttentionType.SWA, DeepseekV4AttentionType.COMPRESSOR_STATE, DeepseekV4AttentionType.COMPRESSOR_SCORE, DeepseekV4AttentionType.INDEXER_COMPRESSOR_STATE, DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE, - } + }) @@ - elif attn_type in self.fixed_size_attention: + elif attn_type in self.FIXED_SIZE_ATTENTION: @@ - if attn_type in self.fixed_size_attention: + if attn_type in self.FIXED_SIZE_ATTENTION:Also applies to: 283-291, 674-677
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py` around lines 91 - 97, The class-level fixed_size_attention set in cache_manager should be made immutable since it is used as a constant. Update the DeepseekV4CacheManager definition to use an immutable collection type and adjust all references in the methods that rely on it (including the other occurrences noted in the file) so membership checks and any related logic continue to work without mutating shared state.Source: Linters/SAST tools
tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py (2)
913-917: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winRemove the manual
__main__runner.Running this file directly bypasses pytest markers like
skip_pre_blackwelland the memory skip, so it can execute the large GPU case on unsupported machines. Use pytest invocation instead.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py` around lines 913 - 917, Remove the manual __main__ runner from TestDeepseekV4CacheManager so the test file is not executed directly outside pytest. Keep the test entry points under pytest control in test_deepseek_v4_cache_manager and rely on pytest markers such as skip_pre_blackwell and the memory skip to gate the large GPU case.
2-2: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAlign the new test annotations with the Python 3.10+ style.
Use built-in generics (
dict,list,tuple,int | None) instead oftyping.Dict/List/Optional/Tuple, add-> Noneto test functions, and make_get_window_size()returnint | Nonesince it returnsNonefor compressed cache types. As per coding guidelines, Python code should use 3.10+ typing style and always annotate function return types. Based on learnings, Python 3.10+ features are expected throughout tests.Also applies to: 27-27, 111-146, 689-720, 833-873
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py` at line 2, Update the DeepSeek V4 cache manager test annotations to Python 3.10+ style: replace typing imports and uses in the relevant test module with built-in generics and union syntax, add explicit -> None return annotations to test functions, and change _get_window_size() to return int | None since it can return None for compressed cache types. Use the existing symbols such as _get_window_size and the affected test methods to locate and update all matching annotations consistently across the file.Sources: Coding guidelines, Learnings
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/__init__.py`:
- Line 4: The __all__ export list in __init__.py needs to be alphabetically
sorted to satisfy RUF022. Update the existing DeepseekV4CacheManager and
DeepseekV4AttentionType entries in __all__ so the export order is lexicographic,
keeping the same symbols but rearranging them in sorted order.
In `@tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py`:
- Around line 149-168: Normalize the public compress_ratios values before
computing block sizes in DeepSeekV4SparseAttentionCacheManager initialization so
that a documented SWA-only ratio of 0 is converted to the internal ratio of 1.
Update the logic around sparse_attn_config.compress_ratios, the MTP extension
block, and the compressed_block_sizes calculation to use normalized ratios
everywhere, and apply the same normalization in the static estimator path
referenced by DeepSeekV4SparseAttentionConfig so tokens_per_block is never
divided by zero.
- Around line 650-678: The context estimate in _get_context_bytes() is
over-counting fixed-window pools because it always uses the full prompt length,
unlike _get_generation_bytes() which caps fixed_size_attention tokens with
_get_window_size(). Update _get_context_bytes() in cache_manager.py to apply the
same window-size cap for fixed_size_attention entries, using the existing
fixed_size_attention set and _get_window_size() helper so SWA/compressor-state
pools are only charged for their actual windowed context.
In `@tensorrt_llm/llmapi/llm_args.py`:
- Around line 803-812: Validate the DSv4 config contract in the llm_args
Pydantic model by tightening the DeepSeekV4-specific fields in the class itself:
make sure index_head_dim is explicitly required for this subclass, and add a
`@field_validator` on compress_ratios to reject or normalize zero values before
they reach cache setup. Update the Field descriptions for compress_ratios,
window_size, and index_topk so they match the actual DSv4 cache-manager
behavior, and add a focused test covering compress_ratios with 0 values to lock
the intended validation behavior.
In
`@tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py`:
- Line 1: The new test file is missing the required NVIDIA copyright header. Add
the standard NVIDIA copyright/licensing header at the top of the file before the
first import in test_deepseek_v4_cache_manager.py so it matches the repository’s
source file guidelines.
- Around line 59-60: The class-level skip_less_device_memory(80000) on
DeepSeekV4CacheManagerTest is too broad and is skipping smaller cache coverage
tests unnecessarily. Move the 80GB memory guard off the class decorator and
apply it only to the heavy path in test_write_read_cache or the large
prompt-length parametrization, while keeping the lighter layout, pool-mapping,
and invalid-value tests runnable. Use the existing DeepSeekV4CacheManagerTest
and test_write_read_cache symbols to relocate the skip without changing the
smaller test cases.
- Around line 725-827: The tests that create and use DeepseekV4CacheManager need
guaranteed cleanup even when assertions fail. In the affected test flow around
cache_manager creation, wrap the manager lifecycle in a try/finally so
cache_manager.shutdown() always runs after free_resources, and apply the same
pattern to the other manager-owning tests mentioned so parametrized runs do not
leak GPU/cache resources. Use the existing DeepseekV4CacheManager,
free_resources, and shutdown symbols to place the cleanup in the right scope.
---
Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py`:
- Around line 44-50: The missing type annotations need to be completed in the
cache manager helpers: add an explicit annotation for has_fp8_kv_cache in
_estimate_bytes_per_token, annotate the nested _add_layer() helper, and add a
return type to get_cache_size_per_token(). Update the relevant function
signatures in cache_manager.py so every function is fully annotated, using None
for helpers that do not return a value.
- Around line 91-97: The class-level fixed_size_attention set in cache_manager
should be made immutable since it is used as a constant. Update the
DeepseekV4CacheManager definition to use an immutable collection type and adjust
all references in the methods that rely on it (including the other occurrences
noted in the file) so membership checks and any related logic continue to work
without mutating shared state.
In
`@tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py`:
- Around line 913-917: Remove the manual __main__ runner from
TestDeepseekV4CacheManager so the test file is not executed directly outside
pytest. Keep the test entry points under pytest control in
test_deepseek_v4_cache_manager and rely on pytest markers such as
skip_pre_blackwell and the memory skip to gate the large GPU case.
- Line 2: Update the DeepSeek V4 cache manager test annotations to Python 3.10+
style: replace typing imports and uses in the relevant test module with built-in
generics and union syntax, add explicit -> None return annotations to test
functions, and change _get_window_size() to return int | None since it can
return None for compressed cache types. Use the existing symbols such as
_get_window_size and the affected test methods to locate and update all matching
annotations consistently across the file.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 30fd78a9-48f2-45eb-b69c-dc0773b48e1d
📒 Files selected for processing (8)
tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/__init__.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.pytensorrt_llm/llmapi/__init__.pytensorrt_llm/llmapi/llm_args.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.pytests/unittest/api_stability/references/llm.yaml
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
Collaborator
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #55531 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #55531 [ run ] completed with state
|
Collaborator
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #55676 [ run ] triggered by Bot. Commit: |
longlee0622
enabled auto-merge (squash)
June 25, 2026 04:17
Collaborator
|
PR_Github #55676 [ run ] completed with state
|
Collaborator
|
/bot run |
Collaborator
|
PR_Github #55707 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #55707 [ run ] completed with state |
1 task
BrianLi23
pushed a commit
to BrianLi23/TensorRT-LLM
that referenced
this pull request
Jul 9, 2026
Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com> Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com> Co-authored-by: Fanrong Li <lfr-0531@users.noreply.github.com> Co-authored-by: Jiagan Cheng <jiaganc@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
6 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Part of the DSv4 split from umbrella PR #14751.
This PR adds the DSv4 sparse cache manager adapter layer, including:
DeepSeekV4SparseAttentionConfigwiring and API stability reference update.This PR is opened as draft because it depends on earlier split PRs:
Rebase note: rebased onto latest
maincommit09449d488171436c60e79fd71b29d35b64ac2516; no git conflicts were encountered. The latest main touched KV cache transfer C++ files outside this PR scope, so I rebuilt and reinstalled before verification.Test Coverage
CCACHE_DIR=/home/scratch.fanrongl_coreai/ccache_trtllm python3 ./scripts/build_wheel.py --trt_root /usr/local/tensorrt --benchmarks --use_ccache --cuda_architectures "90-real;100-real" --configure_cmake./.venv-3.12/bin/pip install --force-reinstall --no-deps build/tensorrt_llm-1.3.0rc19-cp312-cp312-linux_x86_64.whltensorrt_llmfrom this PR worktree and did not import the full DSv4 backend module.nvidia-smi; ran tests on idle GPU 1.CUDA_VISIBLE_DEVICES=1 timeout 20m ./.venv-3.12/bin/python -m pytest tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.pyCUDA_VISIBLE_DEVICES=1 timeout 20m ./.venv-3.12/bin/python -m pytest tests/unittest/llmapi/test_llm_args.py tests/unittest/api_stabilitygit diff --check HEAD~1..HEAD./.venv-3.12/bin/pre-commit run --files ...on all PR6 changed files.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)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.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
To see a list of available CI bot commands, please comment
/bot help.Summary by CodeRabbit