[None][fix] Fix window vector layer indexing - #17466
Conversation
WalkthroughThe change resolves attention windows in pipeline-local layer order. It updates KV-cache sizing, configuration, relocation, attention integration, VSWA detection, distributed capacity reduction, and regression coverage. ChangesPipeline-local KV cache windows
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant KVCacheManagerV2
participant ResourceManager
participant AttentionBackend
participant RelocationKernel
KVCacheManagerV2->>ResourceManager: Resolve configured windows for local layers
ResourceManager-->>KVCacheManagerV2: Return local window vector and cache sizing
KVCacheManagerV2->>AttentionBackend: Provide local-layer cache metadata
AttentionBackend->>RelocationKernel: Use local window and pool metadata
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/attention_backend/trtllm.py (1)
1802-1807: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd regression coverage for both layer-index contracts.
Use distinct window values and a nonzero pipeline-parallel stage. Assert that V2 uses
self.layer_idxand V1 usesget_local_layer_idx(metadata). This protects the corrected behavior from a future simplification.🤖 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/trtllm.py` around lines 1802 - 1807, Add regression tests covering the window selection logic around the V2/V1 branch in the attention backend, using distinct window values and a nonzero pipeline-parallel stage. Assert that KVCacheManagerV2 selects via self.layer_idx, while the V1 path selects via get_local_layer_idx(metadata), preserving both layer-index contracts.
🤖 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.
Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/trtllm.py`:
- Around line 1802-1807: Add regression tests covering the window selection
logic around the V2/V1 branch in the attention backend, using distinct window
values and a nonzero pipeline-parallel stage. Assert that KVCacheManagerV2
selects via self.layer_idx, while the V1 path selects via
get_local_layer_idx(metadata), preserving both layer-index contracts.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 84a563a0-4010-4672-b735-4807b796ca5c
📒 Files selected for processing (1)
tensorrt_llm/_torch/attention_backend/trtllm.py
e590ab0 to
b018ef5
Compare
| window = window_vec[self.local_layer_idx % len(window_vec)] | ||
| # V2 retains the global window pattern, while V1 may shard a | ||
| # full per-layer vector into cache-local order. | ||
| window_layer_idx = self.layer_idx if isinstance( |
There was a problem hiding this comment.
If window_vec is using repeating pattern, looks like the vector will not use cache-local order for V1, should self.layer_idx be used in this case?
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
b018ef5 to
7dcd4fb
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py (1)
693-733: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winReject speculative decoding with
enable_swa_scratch_reuse. The target manager still invokes_update_kv_cache_draft_token_location(), while scratch reuse maps each local layer to a distinct pool. Thelocal_pool_idsassertion therefore fails when accepted draft tokens exist. Add configuration-time validation or support relocation across per-layer pools.🤖 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/pyexecutor/kv_cache_manager_v2.py` around lines 693 - 733, Reject the incompatible configuration during cache-manager setup when speculative decoding is enabled together with enable_swa_scratch_reuse, before _update_kv_cache_draft_token_location can run. Add validation at the relevant initialization/configuration symbol and provide a clear assertion or error; do not rely on the local_pool_ids single-pool assertion in the relocation path.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/attention_backend/trtllm.py (1)
1379-1393: 🗄️ Data Integrity & Integration | 🔵 Trivial | 💤 Low valueUse the cache-local layer index.
V1 and V2 resolve
max_attention_window_vecinto local-layer order.self.local_layer_idxis correct for both managers. Update the PR description to say “cache-local layer index,” not “model-global layer index.”🤖 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/trtllm.py` around lines 1379 - 1393, The implementation already uses self.local_layer_idx correctly in _set_cache_attention_window; update the PR description wording to refer to the “cache-local layer index” instead of the “model-global layer index,” without changing the code.Source: 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/pyexecutor/kv_cache_manager_v2.py`:
- Around line 306-323: Update _resolve_v2_max_attention_window_vec so configured
windows less than or equal to zero are normalized to None, matching
normalize_window_size in _get_static_cache_size_layer_components; retain the
existing max_seq_len-to-None normalization and projection behavior for positive
windows.
---
Outside diff comments:
In `@tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py`:
- Around line 693-733: Reject the incompatible configuration during
cache-manager setup when speculative decoding is enabled together with
enable_swa_scratch_reuse, before _update_kv_cache_draft_token_location can run.
Add validation at the relevant initialization/configuration symbol and provide a
clear assertion or error; do not rely on the local_pool_ids single-pool
assertion in the relocation path.
---
Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/trtllm.py`:
- Around line 1379-1393: The implementation already uses self.local_layer_idx
correctly in _set_cache_attention_window; update the PR description wording to
refer to the “cache-local layer index” instead of the “model-global layer
index,” without changing the code.
🪄 Autofix
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: 1eb33bf8-32f0-4c94-a48b-b34b3663e2fc
📒 Files selected for processing (11)
tensorrt_llm/_torch/attention_backend/trtllm.pytensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/_torch/pyexecutor/resource_manager.pytensorrt_llm/_torch/speculative/eagle3_dynamic_tree.pytests/unittest/_torch/attention/test_trtllm_attention_window.pytests/unittest/_torch/executor/test_kv_cache_estimation.pytests/unittest/_torch/executor/test_kv_cache_manager_v2.pytests/unittest/_torch/executor/test_py_executor.pytests/unittest/_torch/executor/test_resource_manager.pytests/unittest/_torch/speculative/test_eagle3.py
| def _resolve_v2_max_attention_window_vec( | ||
| max_attention_window_vec: Optional[Sequence[int]], | ||
| max_seq_len: int, | ||
| pp_layers: Sequence[int], | ||
| num_layers: int, | ||
| layer_mask: Optional[Sequence[bool]] = None, | ||
| ) -> List[Optional[int]]: | ||
| """Resolve a V2 window pattern into exact cache-local layer order.""" | ||
| configured_windows = _clamp_max_attention_window_vec(max_attention_window_vec, max_seq_len) | ||
| normalized_windows = [ | ||
| None if window == max_seq_len else window for window in configured_windows | ||
| ] | ||
| return _project_max_attention_window_vec( | ||
| normalized_windows, | ||
| pp_layers, | ||
| num_layers, | ||
| layer_mask, | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Normalize non-positive windows to None, as the static path does.
_resolve_v2_max_attention_window_vec maps only window == max_seq_len to None. A configured window of 0 or a negative value survives clamping and reaches AttentionLayerConfig.sliding_window_size at line 1898.
_get_static_cache_size_layer_components treats the same input differently: its normalize_window_size maps window_size <= 0 to None (line 375-376). The static estimator and the runtime layer configuration therefore disagree for the same kv_cache_config.max_attention_window.
KvCacheConfig.max_attention_window declares only min_length=1, so a non-positive entry is accepted by validation.
🔧 Proposed fix to align normalization
configured_windows = _clamp_max_attention_window_vec(max_attention_window_vec, max_seq_len)
normalized_windows = [
- None if window == max_seq_len else window for window in configured_windows
+ None if window is None or window <= 0 or window == max_seq_len else window
+ for window in configured_windows
]📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def _resolve_v2_max_attention_window_vec( | |
| max_attention_window_vec: Optional[Sequence[int]], | |
| max_seq_len: int, | |
| pp_layers: Sequence[int], | |
| num_layers: int, | |
| layer_mask: Optional[Sequence[bool]] = None, | |
| ) -> List[Optional[int]]: | |
| """Resolve a V2 window pattern into exact cache-local layer order.""" | |
| configured_windows = _clamp_max_attention_window_vec(max_attention_window_vec, max_seq_len) | |
| normalized_windows = [ | |
| None if window == max_seq_len else window for window in configured_windows | |
| ] | |
| return _project_max_attention_window_vec( | |
| normalized_windows, | |
| pp_layers, | |
| num_layers, | |
| layer_mask, | |
| ) | |
| def _resolve_v2_max_attention_window_vec( | |
| max_attention_window_vec: Optional[Sequence[int]], | |
| max_seq_len: int, | |
| pp_layers: Sequence[int], | |
| num_layers: int, | |
| layer_mask: Optional[Sequence[bool]] = None, | |
| ) -> List[Optional[int]]: | |
| """Resolve a V2 window pattern into exact cache-local layer order.""" | |
| configured_windows = _clamp_max_attention_window_vec(max_attention_window_vec, max_seq_len) | |
| normalized_windows = [ | |
| None if window is None or window <= 0 or window == max_seq_len else window | |
| for window in configured_windows | |
| ] | |
| return _project_max_attention_window_vec( | |
| normalized_windows, | |
| pp_layers, | |
| num_layers, | |
| layer_mask, | |
| ) |
🤖 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/pyexecutor/kv_cache_manager_v2.py` around lines 306 -
323, Update _resolve_v2_max_attention_window_vec so configured windows less than
or equal to zero are normalized to None, matching normalize_window_size in
_get_static_cache_size_layer_components; retain the existing max_seq_len-to-None
normalization and projection behavior for positive windows.
2 participants
Description
KvCacheConfig.max_attention_windowis a model-global repeating pattern, while KV-cache pools and attention metadata are pipeline-stage local. The previous mixed contract left V1 partially sharded, kept V2 global, and required consumers to guess which layer index to use.Resolve the configured pattern once into an exact per-local-layer
max_attention_window_vecfor both managers. Cache, sizing, and FMHA consumers now consistently use the local layer index, including nonzero and uneven pipeline-parallel stages. The projection preserves the compact enabled-layer convention used by masked hybrid caches and leaves the public configuration unchanged.The change also:
There are no public API, configuration-schema, dependency, or test-list changes.
Test Coverage
pre-commit run --files <all changed files>python3 -m py_compile <all changed Python files>pytestcollection was attempted, but this worktree does not contain the compiledtensorrt_llm.bindingsmodule; collection stops intests/unittest/conftest.pywithModuleNotFoundError.PR Checklist
Dev Engineer Review
calculate_max_num_blockswith an optional synchronization parameter.QA Engineer Review
test_trtllm_attention_window.py.test_kv_cache_estimation.py.test_kv_cache_manager_v2.py.test_py_executor.py.test_resource_manager.py.test_eagle3.py.test-db/,qa/, ortests/integration/test_lists/coverage changes were reported.