[mihai-chiorean]

Summary

Fixes #16801.

Resolve Vanilla-attention KV block indices with the executing layer whenever
the cache layout is layer-specific.

Vanilla metadata is shared across layers. Resolving indices during metadata
preparation without layer_idx can select the wrong sliding window, physical
pool, or V2 page-index scale. The same problem affects ordinary and MLA
generation.

Implementation

• Defer block-index lookup for VSWA, hybrid linear-attention, multi-pool, and
  heterogeneous-scale V2 managers.
• Resolve deferred indices with layer_idx=self.layer_idx in ordinary and MLA
  generation.
• Detect differing V2 layer scales once and cache the capability result.
• Preserve the existing prepared-once lookup for uniform single-pool managers.

Validation

Validated on DGX Spark / GB10 (SM121) with Gemma 3 1B, the PyTorch VANILLA
attention backend, KVCacheManager V1, a 661-token prompt, and 32 decoded tokens.

VSWA production path

• max_attention_window=[512, 512, 512, 512, 512, 1024]
• is_vswa=True
• 1,040 layer-specific cache-index lookups
• zero prepared-once lookups
• prefill and decode completed with coherent output
• PASS

Uniform single-pool control

• max_attention_window=[1024]
• is_vswa=False
• 40 prepared-once cache-index lookups
• zero layer-specific lookups
• prefill and decode completed with coherent output
• PASS

This confirms that layer-dependent layouts resolve indices with the executing
layer while uniform layouts retain the existing prepared-once fast path.

Additional validation:

• tests/unittest/_torch/attention/test_vanilla_attention.py: 8 tests and
  3 subtests passed on SM121.
• SM120 RTX PRO 6000 focused layer-aware and uniform-fast-path regression:
  passed.
• Touched-file pre-commit: passed locally and in GitHub Actions.
• Three independent review passes completed. The initial V2 single-pool scale
  gap and potential uniform-layout overhead were addressed before submission.

Coverage includes VSWA, linear attention, multi-pool, heterogeneous and uniform
V2 scales, ordinary generation, and MLA generation.

Scope

No public API or kernel changes. Layer-invariant layouts retain their existing
prepared-once behavior.

PR Checklist

• The PR follows the TensorRT-LLM coding guidelines.
• New behavior has focused regression coverage.
• No public API or dependency changes.
• DCO sign-off is present on every commit.

Dev Engineer Review

• Defers KV index lookup for layer-specific cache layouts.
• Preserves the prepared-once fast path for uniform layouts.
• Uses layer-aware lookup for ordinary and MLA generation.
• Covers VSWA, linear attention, multi-pool, and heterogeneous V2 scaling.
• No public API, kernel, configuration, or test-list changes were found.
• No correctness or consistency issues were found.

QA Engineer Review

Added or updated unit tests for:

• Layer-specific pool index deferral.
• Layer-specific page-index scale deferral.
• Uniform-scale fast-path behavior.
• Single-pool prepared-once behavior.
• Non-attention layer scales.
• Cache-index recomputation after manager replacement.

The tests are outside tests/integration/test_lists/. No corresponding test-db/ or qa/ entries were found.

Verdict: needs follow-up — CBTS or integration coverage is not documented.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

• Last reviewed commit: 93e6400
• Latest commit, not reviewed: 7eb6406

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

No actionable comments were generated in the recent review.

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: e329e2ca-15ad-4a76-8884-1bffe82f25a2

📥 Commits

Reviewing files that changed from the base of the PR and between 1acd695 and 93e6400.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/attention_backend/vanilla.py
• tests/unittest/_torch/attention/test_vanilla_attention.py

🚧 Files skipped from review as they are similar to previous changes (2)

• tensorrt_llm/_torch/attention_backend/vanilla.py
• tests/unittest/_torch/attention/test_vanilla_attention.py

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

Walkthrough

Vanilla attention detects layer-specific KV cache layouts, invalidates stale cache-index state after manager changes, and defers lookup until execution. Ordinary and MLA generation use the executing layer index. Tests cover uniform and layer-specific layouts.

Changes

Vanilla cache index resolution

Layer / File(s)	Summary
Detect and defer layer-specific indices tensorrt_llm/_torch/attention_backend/vanilla.py, tests/unittest/_torch/attention/test_vanilla_attention.py	Metadata tracks cache-manager changes, ignores page scales for non-attention layers, defers lookup for layer-specific layouts, and preserves eager lookup for uniform layouts. Tests cover these cases and single-pool caching.
Resolve indices during attention execution tensorrt_llm/_torch/attention_backend/vanilla.py, tests/unittest/_torch/attention/test_vanilla_attention.py	Ordinary and MLA generation use a layer-aware cache-index helper. Tests verify requests for each attention layer and MLA output shape.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Suggested reviewers: pengbowang-nv, yunruis

Merge Risk: Minimal · up to 93e64

The PR updates layer-aware KV index resolution while preserving the existing fast path for uniform layouts, and no actionable merge-blocking risk remains.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 4.55% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 22 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely describes the primary change: using layer-aware KV indices in Vanilla attention.
Description check	✅ Passed	The description explains the issue, implementation, validation, scope, and checklist. It provides focused test coverage and relevant runtime validation, although it uses equivalent headings instead of…
Linked Issues check	✅ Passed	The changes address issue #16801 by deferring KV index lookup for layer-specific layouts, resolving indices with the executing layer for ordinary and MLA generation, preserving the uniform-manager fas…
Out of Scope Changes check	✅ Passed	The code and test changes are directly related to correcting layer-aware KV index resolution in Vanilla attention. No unrelated or out-of-scope changes are evident.

Full details: Description check

Explanation

The description explains the issue, implementation, validation, scope, and checklist. It provides focused test coverage and relevant runtime validation, although it uses equivalent headings instead of the template headings.

Full details: Linked Issues check

Explanation

The changes address issue #16801 by deferring KV index lookup for layer-specific layouts, resolving indices with the executing layer for ordinary and MLA generation, preserving the uniform-manager fast path, and adding focused regression coverage.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

[yunruis]

I think Vanilla Attention @yihwang-nv is more familiar, could you please help to review?

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

[mihai-chiorean]

Addressed the latest review note in 93e6400 by reusing the KV cache manager already validated in prepare(). Current main is merged, pre-commit is green, and there are no unresolved conversations. @yihwang-nv, could you take another look when convenient?

[mihai-chiorean]

The latest review feedback is addressed in 93e6400604. The PR is mergeable, all lightweight checks pass, there are no unresolved conversations, and the focused layer-aware KV-index regressions passed on both SM121 and SM120. @yihwang-nv, could you take another look and approve when convenient? @pengbowang-nv @yunruis, once approved, could one of you trigger full CI?

[yuxianq]

Keeping @yihwang-nv as the trt-llm-torch-attention-devs reviewer based on their existing review activity on this PR.

This confirms existing ownership; it does not request another review or add review points.

Why this reviewer is retained

Evidence: inline review comment at 2026-08-19T06:50:00Z.

The bot keeps the first eligible attention reviewer who started reviewing. Current target-branch CODEOWNERS includes attention-owned changes in this PR.

[mihai-chiorean]

@yihwang-nv friendly follow-up. I rebased onto current main (7eb64063fe) and the PR merges cleanly. vanilla.py hasn't changed upstream since the last rebase. The one new commit updates the MLA test to patch _torch.attention.backends.utils after the module move. All earlier review threads have replies.

On CPU the focused tests in test_vanilla_attention.py pass (11 passed; the 2 CUDA-only cases weren't run), and they fail against main's vanilla.py. If this looks good to you, could you approve? @yuxianq, could you then run /bot run? I'm an external contributor and can't trigger CI myself. Thanks!