[rishabhsinha17]

Summary

Fixes #417. Models with cross-layer KV sharing (gemma E2B) crashed at startup with KeyError: '...self_attn.attn' in kvcached's layer-to-tensor binding. Skip the sharing layers wherever a layer is resolved against kv_cache_tensors, then alias each one to its target layer's cache, mirroring vanilla vLLM's initialize_kv_cache_tensors.

Root cause

vLLM's maybe_add_kv_sharing_layers_to_kv_cache_groups appends sharing-layer names to kv_cache_groups[*].layer_names and records them in the runner's runner_only_attn_layers WITHOUT adding them to any kv_cache_tensors[*].shared_by (verified identical on 0.22.1/0.23.0/0.24.0, v1/worker/utils.py). Vanilla vLLM copes by skipping runner_only_attn_layers in _allocate_kv_cache_tensors / _reshape_kv_cache_tensors and aliasing kv_caches[layer] = kv_caches[target] after allocation (gpu_model_runner.py:7064/:7115/:7266 at v0.24.0). kvcached's _allocate_kv_cache_from_kvcached did neither and direct-indexed every grp.layer_names entry, which is the KeyError site.

The appended names also inflated _get_group_size, so the worker would have allocated more pools than the scheduler-side num_layers (the scheduler's config never contains the appended names, the append is runner-side only).

Changes

kvcached/integration/vllm/patches.py:

• New helpers _get_runner_only_attn_layers, _tensor_backed_layer_names, _alias_shared_kv_layers; _get_group_size takes an optional runner-only set (default keeps the old behavior; the scheduler-side call is unchanged).
• _allocate_kv_cache_from_kvcached: resolve only tensor-backed layer names in the validation loop, representative-layer selection, pool count, and the per-group hetero views.
• _reshape_kv_cache_tensors_from_kvcached: bind pools to tensor-backed names only, then alias sharing layers to their target's cache (fail-loud if the target has no cache). On vLLM releases whose vanilla initialize_kv_cache_tensors also aliases after the patched reshape returns, the second pass re-assigns the same objects.
• The V8 path (kv_cache_config.tensors schema) predates cross-layer sharing and is untouched.

tests/test_shared_kv_layer_binding.py (added to tests/manifests/cpu.txt): 14 CPU tests over the helpers, including a scenario shaped like the #417 report that reproduces the KeyError without the filter and a full binding+aliasing pass with it.

Validation

• python3.11 -m pytest tests/test_shared_kv_layer_binding.py: 14 passed.
• Full CPU manifest: 190 passed (176 on main + 14 new); the 10 local errors are the pre-existing macOS /dev/shm + posix_ipc environment issues, unchanged from main.
• tools/check_test_classification.py: valid (cpu=23).
• ruff 0.11.7, isort 6.0.1 clean; mypy --python-version 3.10 introduces no new errors vs main (error set identical).
• git merge-tree --write-tree: merges clean against open PRs fix: make KV tensor VMM batches transactional #418/fix: coordinate distributed physical KV mapping outcomes #455/fix: split page create from map to avoid a remap race on Volta #465 (2026-08-26).
• GPU, 1x L40S, vLLM 0.24.0, transformers 5.14.1, google/gemma-4-E2B-it, autopatch engaged (all 9 patches logged; note the issue's earlier PASS reports were measured with autopatch silently inactive, see Editable install does not install kvcached_autopatch.pth; KVCACHED_AUTOPATCH silently no-ops #470):
  • main + a diagnostic log at the KeyError site: sharing_active=True num_shared=20, group_layer_names=35 tensor_backed(shared_by)=15, 20 layer names in groups but absent from every shared_by, then startup dies with KeyError: 'language_model.model.layers.15.self_attn.attn' in EngineCore.
  • this branch, same env and run: starts and completes the generation cell cleanly, zero KeyErrors.
  • Byte-level output parity vs vanilla was not re-measured in this session (main cannot start at all, so there is no kvcached-on baseline to compare before this fix); happy to add a parity run if wanted.

[RixinLiu]

@shipiyouniao can you help me to review this PR?

[shipiyouniao]

@RixinLiu, I went through f3b8a00. I haven't found a new functional bug in this change. Filtering the runner-only layers before counting/binding pools, then aliasing them to their targets, matches vLLM's handling of cross-layer KV sharing.

The 14 new tests pass. I also exercised the actual patched allocation/reshape entry points with a small review harness: 12 cases passed on CPU and the same 12 passed with T4 CUDA tensors. These cover homogeneous/heterogeneous groups, both layouts, exact target aliasing and readback, a missing target followed by recovery, and delegation when kvcached is disabled. The harness substitutes the allocator/backend; it is not a native VMM or Gemma serving test.

Two follow-ups would make this much easier to sign off:

• Could you include tests that call the patched allocation/reshape methods? The current tests rebuild the binding loop around the helpers, so they would still pass if a helper stopped being used at a call site. In particular, assert that the pool count matches the scheduler's pre-append group size and that a sharing layer gets its target's exact tensor/view in the heterogeneous path.
• Could you add the Gemma output comparison you offered? The reference can be vanilla vLLM with kvcached disabled; main's startup failure doesn't prevent that comparison. A small fixed greedy prompt set, including repeated-prefix requests and more than one decode step, would check that fixing the startup exception hasn't left incorrect KV bindings. Please record the versions/settings and token outputs, and investigate any mismatch rather than requiring only successful generation.

These are validation gaps, not a reproduced correctness defect. I'm leaving this as a comment pending that final model-level check; no code redesign requested.

[RixinLiu]

LGTM. I run the two layouts on my A100, the output matches exactly with vanilla vLLM.