[#16801][fix] Use layer-aware KV indices in Vanilla attention

[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

• Vanilla attention now defers KV index lookup for layer-specific cache layouts.
• The change covers VSWA, linear-attention, multi-pool, and heterogeneous V2 scaling.
• Layer-invariant layouts retain the prepared-once fast path.
• Ordinary and MLA generation use the executing layer index.
• No public API, kernel, configuration, or test-list changes were identified.
• No correctness, performance, or consistency issues were found.

QA Engineer Review

Added unit tests for:

• Layer-specific pool index deferral.
• Layer-specific page-index scale deferral.
• Uniform-scale fast-path behavior.
• Single-pool prepared-once behavior.
• Layer-aware cache and buffer selection.
• MLA generation with layer-specific cache indices.

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

Verdict: needs follow-up — CBTS or integration coverage is not documented for the new unit tests.

[coderabbitai]

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: 42275caa-7a32-4284-8a84-34f045dc9d0a

📥 Commits

Reviewing files that changed from the base of the PR and between f12c5e5 and aaaf44a.

📒 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

---

Walkthrough

Vanilla attention now detects layer-specific KV cache layouts, defers cache-index lookup until execution, and passes the executing layer index for ordinary and MLA generation. Tests cover deferred lookup, eager fast paths, layer selection, and MLA behavior.

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 identifies layer-specific pools and scales, defers cache-index preparation when needed, and preserves eager lookup for uniform layouts. Tests cover these paths.
Resolve indices during attention execution tensorrt_llm/_torch/attention_backend/vanilla.py, tests/unittest/_torch/attention/test_vanilla_attention.py	Ordinary and MLA generation paths retrieve cache indices through the executing attention layer. Tests validate layer-specific cache access and MLA generation.

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

Possibly related PRs

• NVIDIA/TensorRT-LLM#16558: Implements related layer-specific indexer and KV-cache pool mappings.

Suggested reviewers: pengbowang-nv, yunruis

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 5.56% which is insufficient. The required threshold is 80.00%.	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 identifies the fix for layer-aware KV indices in Vanilla attention.
Description check	✅ Passed	The description explains the issue, implementation, validation, scope, and checklist with relevant test coverage.
Linked Issues check	✅ Passed	The changes address all linked issue requirements, including deferred layer-specific lookup, fast-path preservation, and regression tests.
Out of Scope Changes check	✅ Passed	The changes are limited to Vanilla attention logic and focused regression tests related to the linked issue.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{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?