Need support for Gemma model which has cross-layer kv sharing

[RixinLiu]

kvcached crashes with a bare KeyError on models with cross-layer KV sharing

google/gemma-4-E2B-it fails at KV-cache init under kvcached, on both vLLM 0.22.1 and 0.24.0. The same model runs fine with kvcached disabled.

KeyError: 'language_model.model.layers.15.self_attn.attn'
  File "kvcached/integration/vllm/patches.py", line 1285, in _allocate_kv_cache_from_kvcached
    tensor_cfg = layer_to_tensor_cfg[layer_name]

[rishabhsinha17]

Repro, hit while filling the vLLM side of #425.

Env:
Commit: main 8aa9a04
vLLM: 0.23.0 (only version with gemma4_unified before the 0.24 incompatibility in #399) + transformers 5.14.1
GPU: A100-SXM4-80GB
Model: google/gemma-4-E2B-it, --max-model-len 1024, enforce_eager

Startup gets through profiling ("Available KV cache memory: 56.24 GiB") then dies in kvcached's layer-to-tensor binding:

KeyError: 'language_model.model.layers.15.self_attn.attn'

Identical on both layouts. The vLLM patches build a per-layer tensor map and assume every attention layer owns a KV tensor; cross-layer-sharing layers do not have one. The SGLang path resolves sharing inside SGLang's own pool before kvcached sees anything, consistent with E2B passing in @jeff3071's SGLang table.

I can take the fix. Repro environment is still up.

[rishabhsinha17]

Update (revised 2026-08-15; my original mechanism and version-scoping here were wrong, corrected in place): the retest data stands. Retested google/gemma-4-E2B-it on kvcached main 094481f, vLLM 0.24.0, L40S: both layouts pass, outputs byte-identical to the no-kvcached baseline (#425 has the full matrix).

The mechanism I originally gave does not hold. add_kv_sharing_layers_to_kv_cache_groups is identical and identically called from initialize_kv_cache on vLLM 0.22.1, 0.23.0, and 0.24.0 (v1/worker/utils.py), so it cannot explain a version difference, and the report above already shows the KeyError on 0.22.1 and 0.24.0. The registration also does not resolve anything for kvcached: it appends sharing-layer names to kv_cache_groups[*].layer_names and runner_only_attn_layers without adding them to any kv_cache_tensors[*].shared_by. Vanilla vLLM copes by skipping runner_only_attn_layers in its coverage check and aliasing kv_caches[layer] = kv_caches[target] after allocation; kvcached's _allocate_kv_cache_from_kvcached does neither and direct-indexes every grp.layer_names entry, which is the KeyError site. That code is unchanged between my failing repro commit (8aa9a04) and the passing retest commit (094481f).

Which leaves my pass unexplained: the lookup is reachable in the stated config on 0.24, yet the run passed. The candidate I can test next GPU session is whether cross-layer sharing was active at all in my env (transformers 5.14.1); if shared_kv_cache_layers came up empty, the pass is vacuous. Until that is settled, treat the E2B PASS rows in #425 as unconfirmed for the sharing path.

So: not 0.23-scoped, and this should stay open. The fix belongs in kvcached, and its shape follows the vanilla path: skip runner_only_attn_layers when building layer_to_tensor_cfg, and alias shared layers to their target layer's tensor after allocation, mirroring initialize_kv_cache_tensors. Still mine as claimed.