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.
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_idxis available.That argumentless lookup is invalid for layer-specific layouts:
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:
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.