[ischencheng]

Fixes #882.

The KV-shared layers of Gemma-4 E2B / E4B still stored their own copy of the source layer's k, v, which is 57% of E2B's KV per token and 43% of E4B's. Now they read the source layer's slab and store nothing, so the same KV memory holds about 7/3x the tokens on E2B and 7/4x on E4B. Outputs are unchanged: on main the copies hold the same rows as the source's slab.

• ModelConfig.kv_shared_layers lists (layer, source) pairs: each shared layer and the last earlier layer of its attention type (14 / 13 on E2B, 23 / 22 on E4B; llama.cpp's kv_source_impl gives the same). The HF and GGUF parsers fill it from Gemma4EArgs.kv_shared_layers, which raises NotImplementedError for a group whose shared layers have no earlier layer of their type; no release has that shape.
• Gemma4Attention takes its shared and source flags from that list and keys shared_kv by the source layer, so the module and the pool read one list.
• kv_cache_group_specs() moves those layers out of each group's layer_ids into shared_layers, so the pool allocates slabs and kv_cost prices bytes only for layers that own KV. The budget code itself does not change.
• HybridSWAKVCache maps a shared layer onto its source's slab, and store_kv returns early for it: the source stored the same rows earlier in the forward. The attention module still passes k, v to the backend, since the extend kernel takes the current chunk from k_extend / v_extend.
• With --kv-cache-dtype fp8_per_token_head / int8_per_token_head (feat(kvcache): add fp8 and int8 per-token-head KV cache formats #894) the scales are looked up through the same layers_mapping, so a shared layer reads its source's scales, and the early return skips the quantized store as well. Nothing else is needed for those formats.
• The SWA radix cache tracks tokens and slots, not layers, and the window slots come from the one full_to_swa_index_mapping, so a shared SWA layer reads the same slots as its source and the cache bookkeeping does not change.
• The non-SWA path of create_kvcache_pool raises on shared layers. Only E2B / E4B have them, and they always have a window group.
• The Gemma-4 assistant draft (feat(gemma4): speculative decoding with the assistant draft #895) reads the target's slabs through ReadOnlyTargetKVView, which routes a shared target layer to its source through the same list. The draft's own layers have no source in its config, so they stay marked shared by its first_kv_shared_layer = 0.

End to end

A10 (sm_86, 24 GB), driver 610.57.04, CUDA 13.0, in the test container. Both arms run in the same container: the main arm copies the changed files back from main. One server start per arm and case.

KV capacity from /v1/cache/status (num_pages, page size 1; window pool in brackets), same cache budget in both arms, on 7796a0d with main 66190a9 (this head adds only a test fix on top). The last three rows run the assistant draft from #895 (--text-model-only --spec-algo mtp --spec-draft-model mtp-gemma-4-E{2,4}B-it-BF16.gguf from unsloth's MTP/):

checkpoint	setup	main	this PR
google/gemma-4-E2B-it	whole GPU	553,187 (110,637)	1,299,063 (259,812), 2.35x
google/gemma-4-E4B-it	whole GPU	123,106 (24,621)	221,259 (44,252), 1.80x
google/gemma-4-E2B-it	--kv-cache-dtype int8_per_token_head	1,101,438 (220,287)	2,576,278 (515,255), 2.34x
unsloth/gemma-4-E2B-it-GGUF Q4_K_M	assistant draft	859,623 (171,924)	2,016,682 (403,336), 2.35x
unsloth/gemma-4-E4B-it-GGUF Q4_K_M	assistant draft	353,416 (70,683)	626,180 (125,236), 1.77x
google/gemma-4-E4B-it	assistant draft	117,355 (23,471)	213,130 (42,626), 1.82x

On the previous head (3496477, on 27740fe), with 7.5 GiB free (a side process holding the rest of the A10, about what an 8 GB card leaves) and --embed-weights host: E2B 86,205 -> 209,954 (2.44x), E4B GGUF Q4_K_M 77,436 -> 141,342 (1.83x), E2B with fp8_per_token_head 177,136 -> 421,830 (2.38x); and on the whole GPU, E2B with --cache-type naive 696,444 -> 1,740,158 (2.50x), E4B with int8_per_token_head 251,286 -> 445,506 (1.77x). Per token, unit_bytes goes from 14336 / 28672 (full / window) to 6144 / 12288 on E2B, and from 28672 / 71680 to 16384 / 40960 on E4B; with the 8-bit formats from 7224 / 14560 to 3096 / 6240 on E2B and from 14448 / 36400 to 8256 / 20800 on E4B. The ratios come out a bit above 7/3 and 7/4 at small budgets because the window pool's fixed floor shrinks too.

Greedy outputs, compared as text: 11 sequential requests (4 raw completions of 128 tokens with ignore_eos, for the decode rate; a 4391-token prompt with one odd record in the middle and a question about it; a second question on the same prompt; the first one again, served from the prefix cache under radix; 4 chat requests), then 4 concurrent requests (both long ones and two chat ones). Every case matches main 11/11 and 4/4, the assistant-draft cases included, and so does a second E4B GGUF main start. The long prompt is well past the 512-token window. On the previous head the five cases above, and a runtime /v1/cache/rebuild to 60,000 pages on E2B followed by the same 15 requests, also matched main.

Decode at bs = 1, predicted_per_second, median of the 4 raw completions, main -> this PR: E2B 80.48 -> 81.11, E4B 43.39 -> 43.62, E2B int8 80.07 -> 80.79; with the assistant draft, E2B GGUF 158.0 -> 163.3, E4B GGUF 133.9 / 133.3 (two starts) -> 134.0, E4B 116.4 -> 117.3. Likely the 20 / 18 skipped stores per step, I did not profile it. On the previous head, prefill on the 4391-token prompt was the same within noise (E2B 374.5 / 379.1 -> 373.4 ms, E4B 657.4 -> 665.4 ms).

ft serve --model google/gemma-4-E2B-it --port 1919 --enable-timings-report --max-running-requests 4 [--embed-weights host] [--cache-type naive] [--kv-cache-dtype int8_per_token_head|fp8_per_token_head]
ft serve --model gemma-4-E4B-it-Q4_K_M.gguf --port 1919 --enable-timings-report --max-running-requests 4 --embed-weights host
ft serve --model gemma-4-E4B-it-Q4_K_M.gguf --port 1919 --enable-timings-report --max-running-requests 4 --text-model-only --spec-algo mtp --spec-draft-model mtp-gemma-4-E4B-it-BF16.gguf

Test

test_shared_layers_read_their_source_slab_and_own_none in tests/models/gemma4/test_e_series.py builds the pool from the E2B and E4B layer layouts, with the default format and with int8_per_token_head, and checks the slab counts and that each shared layer reads its source's buffers (and its scales for int8). It is red on main (7 / 28 and 7 / 35 slabs). It goes red when the copies come back, a shared layer gets the wrong source, or its scales stop following the source, and none of those crashes. The GGUF parser fills the field through the same Gemma4EArgs.kv_shared_layers; that line is covered by the GGUF runs above, not by this test.

test_attention_reads_the_pairs_the_pool_maps checks that every E2B / E4B attention module's shared and source flags are the list's, test_draft_layers_project_q_only_and_hand_over_nothing that the assistant draft's layers still build q only, and test_kv_shared_layers_with_no_earlier_source_are_refused the refusal. The two loader tests that build a 3-layer config now give the shared layer an earlier layer of its type, the shape every release has.

tests/models tests/kvcache tests/scheduler tests/engine tests/attention tests/spec tests/checkpoint on the A10, this head with main 66190a9: main 1003 passed, 234 skipped, 5 failed; this PR 1011 passed, 234 skipped, 5 failed. The same 5 fail in both (qwen4_exp/test_qsa_backend.py::test_chunked_prefill_matches_one_shot x3, dflash/test_draft.py::test_block_forward_matches_reference fp32 x2).

Not tested: TP, ROCm, a real 8 GB card (the A10 with memory held is a stand-in), RTX 40/50.

[jason-fxz]

@ischencheng Looks good to me. Every KV write goes through store_kv and every read through layers_mapping, including the #894 scales, so skipping the store on a shared layer loses nothing. This also matches vLLM, which gives shared layers no KV either (E2B full layers: 6 vs 14 KiB per token).

Two small asks.

First, Gemma4Attention re-derives kv_shared / kv_source with is_kv_shared / is_kv_source while the pool reads config.kv_shared_layers. Please take both from config.kv_shared_layers so the module and the pool read one list (vLLM likewise declares the target once, through kv_sharing_target_layer_name, and allocation follows it); is_kv_source can then go.

Second, the fallback in Gemma4EArgs.kv_shared_layers says layers with no earlier source "keep their own slabs", but Gemma4Attention still marks them kv_shared, skips their k / v projection and reads shared_kv[group], which nothing fills, so such a config would fail with a KeyError on the first forward. No release has that shape, so please raise there like the other unsupported shapes in gemma4_e_args (SGLang raises a ValueError for the same case):

source = max((i for i in layer_ids if i < self.first_kv_shared_layer), default=None)
if source is None:
    raise NotImplementedError("gemma4 KV-shared layers with no earlier layer of their attention type are not supported")

[ischencheng]

Thx, both are in 7796a0d: Gemma4Attention takes its shared and source flags from config.kv_shared_layers (shared_kv is keyed by the source layer, is_kv_source is gone), and Gemma4EArgs.kv_shared_layers raises the NotImplementedError you suggested. I merged main for #895, and its read-only draft view now routes through the same list too; the draft's own layers have no source in their config, so they stay shared through first_kv_shared_layer = 0. With the assistant draft, greedy outputs on E2B / E4B still match main, numbers in the body.

[jason-fxz]

Thanks, both are in. On an RTX PRO 6000, E2B now gets 2.33x the KV pages of main and the greedy outputs match main (8/8).

One last ask: please drop test_kv_shared_layers_with_no_earlier_source_are_refused and test_draft_layers_project_q_only_and_hand_over_nothing. Each asserts a branch this PR just wrote, and what it would catch already fails loudly: without the raise, the unsourced layer's qkv_proj is missing at load (KeyError), and without the first_kv_shared_layer == 0 case the assistant GGUF is refused at startup (unexpected MTP tensor 'layers.0.self_attn.q_proj.weight'). Please keep test_attention_reads_the_pairs_the_pool_maps: a module that picks a different source than the pool changes 5 of 8 E2B outputs with no error, and only that test catches it.