[#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

• Defers KV index lookup for layer-specific cache layouts.
• Covers VSWA, linear attention, multi-pool, and heterogeneous V2 scaling.
• Preserves the prepared-once fast path for uniform layouts.
• Uses layer-aware lookup for ordinary and MLA generation.
• 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]

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: b2df9d4c-0a6a-4b7b-b85b-9a39685f697b

📥 Commits

Reviewing files that changed from the base of the PR and between 94b5f9b and 66d6949.

📒 Files selected for processing (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 now detects layer-specific KV cache layouts, invalidates stale cache-index state after manager changes, and defers lookup until execution. Ordinary and MLA generation pass 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

Merge Risk: ⚪ Minimal · up to 66d69

The change defers KV index resolution only for layer-dependent cache layouts while preserving the existing fast path for uniform layouts; no actionable merge-blocking risk remains after normal checks and review.

Suggested reviewers: bowenfu, 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 and concisely identifies the layer-aware KV-index fix in Vanilla attention.
Description check	✅ Passed	The description covers the bug, implementation, validation, scope, and key checklist items; only some optional template items are omitted.
Linked Issues check	✅ Passed	The changes address all linked issue objectives, including deferred layer-aware lookup, uniform fast-path preservation, and focused regression coverage.
Out of Scope Changes check	✅ Passed	The code and test changes are directly related to correcting layer-specific KV-index handling in Vanilla attention.

✨ 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?

[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.