[None][feat] Add BaseResourceManager-based KV-cache compression manager framework - #15106
Merged
Merged
Conversation
…nager-based) Adds a standalone L2 KV-cache compression framework, decoupled from any concrete algorithm: - BaseKVCacheCompressionManager(BaseResourceManager) with 8 semantic lifecycle hooks (on_request_init; on_context_attention[_end]; on_context_end; on_generation_attention[_end]; on_generation_step_end; on_request_finish), auto-driven by PyExecutor's prepare/update/ free_resources -- no PyExecutor changes needed. - SparseAttentionManager / KVCacheStorageManager axis subclasses. - A single compression manager is registered in the resource-manager registry; the per-layer attention hooks fire from TrtllmAttention.forward via metadata.compression_manager. - Factory: create_compression_manager / create_sparse_attention_manager. - Unit tests (test_compression_manager.py). Framework only: no concrete behavior-layer algorithm ships here, and multi-method stacking (composing several axes) is intentionally future work. Legacy rocket/dsa/skip_softmax paths are untouched. Signed-off-by: Hu Tianrui <tianruih@nvidia.com>
Hudayday
force-pushed
the
kvcache-compression-framework
branch
from
June 8, 2026 13:50
0c9ebc8 to
3d7a89e
Compare
Contributor
📝 WalkthroughWalkthroughThis PR introduces a behavior-layer KV-cache compression manager framework for sparse attention. It defines lifecycle hooks for request initialization, per-layer attention (context/generation), phase boundaries, and completion, integrates these hooks into the attention backend and PyExecutor resource management, and provides factory dispatch routing based on a new ChangesKV-Cache Compression Manager Framework
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Suggested labels
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: 6
🧹 Nitpick comments (1)
tests/unittest/_torch/attention/sparse/test_compression_manager.py (1)
89-226: Coverage assessment: unit scope here is sufficient; integration can be follow-up.For
tests/unittest/_torch/attention/sparse/test_compression_manager.py, coverage is sufficient for the unit contracts (ABC defaults, RM hook translation, factoryNonebehavior, and canonical exports). If you want to extend coverage later, do it as follow-up integration tests aroundtensorrt_llm/_torch/attention_backend/trtllm.pyandtensorrt_llm/_torch/pyexecutor/model_engine.pywiring paths.🤖 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/test_compression_manager.py` around lines 89 - 226, The review says unit coverage in tests/unittest/_torch/attention/sparse/test_compression_manager.py is sufficient and no code changes are required; leave the test file as-is (classes TestBaseABC, TestSubclasses, TestResourceManagerAPI, TestFactories, TestCanonicalImports and factory functions like create_compression_manager/create_sparse_attention_manager need no modifications) and defer any additional coverage to follow-up integration tests around tensorrt_llm/_torch/attention_backend/trtllm.py and tensorrt_llm/_torch/pyexecutor/model_engine.py.Source: Coding guidelines
🤖 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/__init__.py`:
- Around line 1-6: Update the module docstring to accurately describe the
current public API: state that concrete manager subclasses are re-exported at
package level (list the re-exported subclass names) and that factory dispatchers
create_sparse_attention_manager and create_compression_manager still exist and
use the Pydantic discriminator SparseAttentionConfig to select implementations;
keep guidance that users should prefer the factory dispatchers for production
code while acknowledging the convenience of direct imports of the re-exported
subclasses.
In `@tensorrt_llm/_torch/attention_backend/sparse/utils.py`:
- Around line 86-89: The behavior-layer configs are being allowed to run on
backends that don't implement the compression hooks (so TrtllmAttention.forward
invokes on_*_attention* but vanilla/flashinfer never see them), causing silent
degradation; update the per-method dispatch in utils.py to detect when a
behavior-layer config is being used with a non-TRTLLM backend (use the existing
backend identity check — e.g., inspect the backend returned by
get_attention_backend or type/name of the backend) and immediately raise a clear
RuntimeError indicating that behavior-layer attention requires the TRTLLM
backend; apply the same fail-fast check to the other affected block(s) around
the 100-118 region so behavior-layer configs cannot be routed to
vanilla/flashinfer until those backends implement the on_*_attention* hooks
(reference TrtllmAttention.forward and the on_*_attention* hook names).
In `@tensorrt_llm/_torch/attention_backend/trtllm.py`:
- Around line 1738-1757: Compression hooks are being called with k=None
(fused-QKV path) and both on_context_attention and on_generation_attention are
invoked every forward; change the logic so you only call compression_manager
hooks when k is non-None and when the hook matches the current phase (call
on_context_attention only during context phase and on_generation_attention only
during generation phase using the phase indicator on metadata, e.g.,
metadata.is_generation or metadata.phase), and for fused-QKV where k is None
skip compression_manager and fall back to the existing sparse_kv_predict /
sparse_attn_predict paths (functions: metadata.compression_manager,
on_context_attention, on_generation_attention, sparse_kv_predict,
sparse_attn_predict, get_local_layer_idx) to avoid passing None into the hook
API and to prevent double/incorrect state updates.
In `@tensorrt_llm/_torch/attention_backend/utils.py`:
- Around line 22-31: The code currently sets sparse_attn_config = None when
sparse_attn_config.is_behavior_layer_method, silently disabling behavior-layer
configs; instead, detect the backend and raise an explicit error when a
behavior-layer sparse config is used on a non-TRTLLM backend. Update the branch
that checks sparse_attn_config.is_behavior_layer_method to: if the active
backend (check the variable representing the backend or a helper like
is_trtllm_backend) is TRTLLM allow the config to proceed (or keep existing
behavior), otherwise raise a ValueError (or other appropriate exception)
indicating that behavior-layer sparse_attn_config is invalid on non-TRTLLM
backends; reference sparse_attn_config and its is_behavior_layer_method
attribute to locate the logic to change.
In `@tensorrt_llm/_torch/pyexecutor/_util.py`:
- Around line 1670-1691: The new KV cache compression manager must not run after
the KV cache manager; update the registration so
ResourceManagerType.KV_CACHE_MANAGER remains the last entry. After creating
compression_manager (in the block that uses create_compression_manager,
ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER, kv_cache_manager,
model_engine.compression_manager), insert the compression manager into
resource_manager.resource_managers and then, if
ResourceManagerType.KV_CACHE_MANAGER is present, pop and reinsert that existing
KV_CACHE_MANAGER entry so it becomes the final item (preserving the documented
prepare/update/free_resources ordering).
In `@tests/unittest/_torch/attention/sparse/test_compression_manager.py`:
- Line 1: Add the required NVIDIA copyright/license header at the very top of
the file so it appears before the existing module docstring in
test_compression_manager.py; ensure the header matches the repo's standard
template and includes the current modification year, then keep the existing
docstring and tests unchanged below it.
---
Nitpick comments:
In `@tests/unittest/_torch/attention/sparse/test_compression_manager.py`:
- Around line 89-226: The review says unit coverage in
tests/unittest/_torch/attention/sparse/test_compression_manager.py is sufficient
and no code changes are required; leave the test file as-is (classes
TestBaseABC, TestSubclasses, TestResourceManagerAPI, TestFactories,
TestCanonicalImports and factory functions like
create_compression_manager/create_sparse_attention_manager need no
modifications) and defer any additional coverage to follow-up integration tests
around tensorrt_llm/_torch/attention_backend/trtllm.py and
tensorrt_llm/_torch/pyexecutor/model_engine.py.
🪄 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: c2478dfa-b3e1-4eda-b235-45fae8142df5
📒 Files selected for processing (11)
tensorrt_llm/_torch/attention_backend/interface.pytensorrt_llm/_torch/attention_backend/sparse/__init__.pytensorrt_llm/_torch/attention_backend/sparse/kv_cache_compression_manager.pytensorrt_llm/_torch/attention_backend/sparse/utils.pytensorrt_llm/_torch/attention_backend/trtllm.pytensorrt_llm/_torch/attention_backend/utils.pytensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/_torch/pyexecutor/resource_manager.pytensorrt_llm/llmapi/llm_args.pytests/unittest/_torch/attention/sparse/test_compression_manager.py
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #52796 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #52796 [ run ] completed with state |
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #52798 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #52798 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #52903 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #52903 [ run ] completed with state |
Hudayday
force-pushed
the
kvcache-compression-framework
branch
from
June 10, 2026 01:44
d0dcc81 to
3d7a89e
Compare
heyuhhh
reviewed
Jun 10, 2026
heyuhhh
reviewed
Jun 10, 2026
heyuhhh
reviewed
Jun 10, 2026
heyuhhh
reviewed
Jun 10, 2026
heyuhhh
reviewed
Jun 10, 2026
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #54561 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #54561 [ run ] completed with state
|
…ramework Signed-off-by: Hu Tianrui <tianruih@nvidia.com> # Conflicts: # tensorrt_llm/_torch/pyexecutor/_util.py # tensorrt_llm/_torch/pyexecutor/resource_manager.py
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #54708 [ run ] triggered by Bot. Commit: |
DomBrown
approved these changes
Jun 17, 2026
DomBrown
left a comment
DomBrown
left a comment
Collaborator
There was a problem hiding this comment.
LGTM for the API stability part.
Collaborator
|
PR_Github #54708 [ run ] completed with state
|
The executor factory create_py_executor_instance reads llm_args.kv_cache_compression_config to decide whether to build a compression manager. Real LlmArgs always defines this optional field, but minimal unit-test mocks (test_dual_pool_kv_cache) do not, so the direct attribute access raised AttributeError before the scheduler was built. Look it up with getattr defaulting to None so an absent field is treated the same as None. Signed-off-by: Hu Tianrui <tianruih@nvidia.com>
…ramework Signed-off-by: Hu Tianrui <tianruih@nvidia.com>
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #54851 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #54851 [ run ] completed with state
|
…ramework Signed-off-by: Hu Tianrui <tianruih@nvidia.com>
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #54926 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #54926 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #54950 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #54950 [ run ] completed with state |
lowsfer
reviewed
Jun 23, 2026
BrianLi23
pushed a commit
to BrianLi23/TensorRT-LLM
that referenced
this pull request
Jul 9, 2026
…er framework (NVIDIA#15106) Signed-off-by: Hu Tianrui <tianruih@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
10 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
Adds a standalone KV-cache compression framework: a
BaseResourceManagerbase class that lets a KV-reduction algorithm (e.g. periodic token eviction /
KV compaction) run alongside the normal KV cache manager, driven entirely by
PyExecutor's existing resource-manager loop. No PyExecutor changes, no
attention-backend changes; fully dormant unless a compression config is set.
Framework only — no concrete algorithm ships in this PR.
Design
KV-cache compression changes which KV is stored, not the attention
computation — so it is a resource manager, kept entirely separate from
SparseAttentionConfigand the attention backend.BaseKVCacheCompressionManager(BaseResourceManager)(inresource_manager.py)— four KV-cache lifecycle hooks, all default no-op:
on_request_init— per-request setup (e.g. allocate scoring buffers)on_context_step_end— end of prefill (one-shot prefill-end eviction)on_generation_step_end— each decode step (periodic / budget eviction)on_request_finish— release per-request stateBecause it is a
BaseResourceManager, PyExecutor's main loop already drivesit:
prepare_resources/update_resources/free_resourcestranslate intothe four hooks, gated on the same signals the peer resource managers use
(
is_first_context_chunk,context_requests_last_chunk) — no manager-siderequest bookkeeping. It holds a
KVCacheManagerV2as a tool and never inheritsfrom it (the cache manager owns the physical KV).
__init__refuses KV-cacheblock reuse (a method that rewrites stored K/V can't share prefix blocks — the
same guard
RocketKVCacheManagermakes);get_max/needed_resource_to_completionreturn 0 (it owns no physical resource, so it never gates the scheduler).
create_kv_cache_compression_manager(config, kv_cache_manager)— factorydispatched from LLM init; framework-only, so it warns + returns None for any
algorithm (concrete methods add a dispatch branch in a follow-up PR).
KvCacheCompressionConfig— a top-levelLlmArgsfield, separate fromSparseAttentionConfig.Wiring
create_py_executor(_util.py) builds the manager fromkv_cache_compression_configwhen set and registers it in the resource-managerregistry (
ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER) the same way theKV cache manager is registered; it runs after the cache manager so it reconciles
once the cache is resized.
Scope
method drives standard attention over the (physically compacted) KV; it is not
sparse attention.
existing inference.
Changes per file
_torch/pyexecutor/resource_manager.pyBaseKVCacheCompressionManager(BaseResourceManager)(4 lifecycle hooks + RM-API→hook translation + zero resource counts + block-reuse guard) +create_kv_cache_compression_managerfactory +ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER._torch/pyexecutor/_util.pykv_cache_compression_config, before warmup.llmapi/llm_args.pyKvCacheCompressionConfig+LlmArgs.kv_cache_compression_config.tests/unittest/_torch/pyexecutor/test_kv_cache_compression_manager.pyTest Coverage
tests/unittest/_torch/pyexecutor/test_kv_cache_compression_manager.py:BaseResourceManagerinheritance; four hooks default no-op + accept extrakwargs; zero resource counts.
prepare_resourcesfireson_request_initon thefirst prefill chunk only;
update_resourcesfireson_context_step_endforcontext_requests_last_chunk+ exactly oneon_generation_step_endperiteration;
free_resourcesfireson_request_finish.resource_manager, not inthe sparse module.
PR Checklist