Skip to content

[None][fix] Fix VSWA gate in KVCacheManager window-size resolution#16233

Open
eopXD wants to merge 1 commit into
NVIDIA:mainfrom
eopXD:user/yuehtingc/kv-vswa-window-resolution
Open

[None][fix] Fix VSWA gate in KVCacheManager window-size resolution#16233
eopXD wants to merge 1 commit into
NVIDIA:mainfrom
eopXD:user/yuehtingc/kv-vswa-window-resolution

Conversation

@eopXD

@eopXD eopXD commented Jul 10, 2026

Copy link
Copy Markdown
Collaborator

@coderabbitai summary

Description

KVCacheManager.get_cache_indices, get_priority_by_block_id, and
get_batch_cache_indices gated the "window_size required" check on
len(max_attention_window_vec) > 1. That spuriously rejects a uniform
per-layer window vector such as [4096, 4096, ...], which has len > 1 but a
single effective window (is_vswa is False). commit_and_get_block_hashes
already used the correct is_vswa signal (distinct positive windows); the
others did not — so callers that legitimately omit window_size (e.g. the KV
cache connector) broke on any model with a uniform per-layer window vector.

This extracts a single _resolve_window_size helper keyed on is_vswa and
routes all four per-window accessors through it, removing the copy-pasted guard
and its divergence. Behavior for genuine VSWA (raise when window_size is
omitted) and for an explicitly-provided window_size is unchanged.

Test Coverage

tests/unittest/_torch/executor/test_resource_manager.py::TestResolveWindowSize

  • test_uniform_multi_layer_vector_does_not_raise — the regression: a uniform
    [4096, 4096, 4096] vector resolves to 4096 instead of raising.
  • test_genuine_vswa_requires_window_size — genuine VSWA still raises when
    window_size is omitted, and an explicit window_size passes through.

PR Checklist

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

get_cache_indices, get_priority_by_block_id, and get_batch_cache_indices
gated the "window_size required" check on len(max_attention_window_vec) > 1,
which spuriously rejects a uniform per-layer window vector (e.g.
[4096, 4096, ...]) that has a single effective window.
commit_and_get_block_hashes already used the correct is_vswa signal (distinct
positive windows); the others did not, so callers that omit window_size (e.g.
the KV cache connector) broke on such models.

Extract a shared _resolve_window_size helper keyed on is_vswa and route all
four per-window accessors through it.

Signed-off-by: Yueh-Ting Chen <yueh.ting.chen@gmail.com>
@eopXD eopXD requested a review from a team as a code owner July 10, 2026 08:58
@eopXD

eopXD commented Jul 10, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58649 [ run ] triggered by Bot. Commit: 403dfdb Link to invocation

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.

2 participants