[Bug] Vanilla attention uses incorrect KV indices for layer-specific cache layouts
Nobody has claimed this yet.
- 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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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