Skip to content

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

Description

@mihai-chiorean

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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    KV-Cache Managementkv-cache management for efficient LLM inference

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions