NVIDIA / NVIDIA/TensorRT-LLM

[Bug] Vanilla attention uses incorrect KV indices for layer-specific cache layouts

Open
#16,801 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

KV-Cache Management
Dominant language
Python
Stars
14.7k
Forks
2.8k
Avg merge
2d 23h
Merged PRs (30d)
489

Description

Description

Vanilla attention metadata is prepared once and shared by all attention
layers. It currently resolves KV block indices during metadata preparation,
where no executing layer_idx is available.

That argumentless lookup is invalid for layer-specific layouts:

  • variable sliding-window attention;
  • hybrid linear/attention caches;
  • managers with multiple physical pools;
  • KVCacheManagerV2 layouts where layers in one pool use different page-index
    scales.

The result can be an exception from an ambiguous lookup or block indices from
the wrong window, pool, or index scale. Both ordinary and MLA generation use
the prepared indices.

Expected behavior

Layer-invariant managers should retain the existing prepared-once lookup.
Layer-specific managers should defer the lookup until the attention layer
executes and call:

get_batch_cache_indices(request_ids, layer_idx=self.layer_idx)

Reproduction

Prepare shared Vanilla metadata with a VSWA, multi-pool, or heterogeneous-scale
V2 KV manager, then execute two attention layers mapped to different cache
layouts. The current argumentless lookup cannot select the correct
layer-specific index space.

Proposed fix

Detect and cache whether the manager requires layer-specific indices. Preserve
the prepared-once fast path for uniform single-pool managers; otherwise resolve
indices in ordinary and MLA generation using the executing layer.

Focused tests cover VSWA, linear, multi-pool, heterogeneous and uniform V2
scale cases, ordinary generation, and MLA generation. The final suite passes
on SM121 and the focused behavior is validated on SM120.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with Vanilla attention metadata preparation and the ordinary and MLA generation paths, then trace get_batch_cache_indices and how layer_idx is available during execution. Use the focused tests covering VSWA, linear, multi-pool, heterogeneous and uniform V2 scale cases; done means both generation paths select correct layer-specific indices while uniform managers retain the prepared-once path.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
ai, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.