perf(gemma4): read the source layer's kv slab on kv-shared layers

python/freetoken/kvcache/__init__.py

@@ -177,6 +177,8 @@ def create_kvcache_pool(
     # slabs.
     layer_ids: tuple[int, ...] | None = None
     kv_specs = [s for s in model_config.kv_cache_group_specs() if s.num_layers > 0]
+    if any(s.shared_layers for s in kv_specs):
+        raise NotImplementedError("KV-shared layers are only mapped by the hybrid full/SWA pool")
     if model_config.has_linear_attention:
         assert len(kv_specs) == 1, (
             f"hybrid-linear models support one paged-KV group, got "

python/freetoken/kvcache/hybrid_swa_pool.py

@@ -87,6 +87,7 @@ def __init__(
             "swa": self.swa_kv_pool,
         }
         self.layers_mapping = self._build_layers_mapping(num_layers, specs)
+        self._shared_layers = frozenset(layer_id for spec in specs.values() for layer_id, _ in spec.shared_layers)
         if self._swa_paged:
             self._init_swa_paged_state()
 
@@ -116,6 +117,11 @@ def _build_layers_mapping(
                 if mapping[layer_id] is not None:
                     raise ValueError(f"KV layer id {layer_id} appears in more than one group")
                 mapping[layer_id] = _LayerRef(group=group_name, index=local_index)
+        for group_name in ("full", "swa"):
+            for layer_id, source in specs[group_name].shared_layers:
+                if source not in specs[group_name].layer_ids or mapping[layer_id] is not None:
+                    raise ValueError(f"KV-shared layer {layer_id} cannot read layer {source}'s {group_name} slab")
+                mapping[layer_id] = mapping[source]
 
         missing = [layer_id for layer_id, ref in enumerate(mapping) if ref is None]
         if missing:
@@ -220,6 +226,8 @@ def store_kv(
         out_loc: torch.Tensor,
         layer_id: int,
     ) -> None:
+        if layer_id in self._shared_layers:
+            return  # its source layer stored these rows earlier in the same forward
         ref = self.layers_mapping[layer_id]
         storage = self._storages[ref.group]
         indices = out_loc

python/freetoken/models/config.py

@@ -130,6 +130,9 @@ class KVCacheGroupSpec:
     # Attention-type taxonomy value for this group; drives the backend capability
     # matrix and (with the pool factory) selects the KV pool family.
     attn_type: AttnType = AttnType.FULL
+    # (layer, source) pairs: layers of this group that read the source's slab and own none.
+    # They are not in ``layer_ids``, so neither the pool nor the cost model gives them a slab.
+    shared_layers: Tuple[Tuple[int, int], ...] = ()
 
     @property
     def num_layers(self) -> int:
@@ -382,6 +385,8 @@ class ModelConfig:
     # Generic execution-path capability flags (set by a model's parse_config) so the engine and
     # factories stay model-agnostic instead of branching on dsv4_args:
     single_stream_only: bool = False  # model runs one sequence at a time -> force bs=1
+    # (layer, source) pairs: the layer attends over the KV the source layer stored and stores none
+    kv_shared_layers: Tuple[Tuple[int, int], ...] = ()
     # Extra per-request tensors riding the LinearStatePool slots (see SlotStateSpec);
     # () for models without any. Requires a linear-attention group to ride on.
     slot_states: Tuple[SlotStateSpec, ...] = ()
@@ -498,13 +503,21 @@ def kv_cache_group_specs(self) -> Tuple[KVCacheGroupSpec, ...]:
                 ),
             )
 
+        sources = dict(self.kv_shared_layers)
+
+        def split(group) -> dict:
+            return dict(
+                layer_ids=tuple(i for i in group.layer_ids if i not in sources),
+                shared_layers=tuple((i, sources[i]) for i in group.layer_ids if i in sources),
+            )
+
         specs: list[KVCacheGroupSpec] = []
         for group in self.attention_groups:
             if isinstance(group, FullAttentionGroupConfig):
                 specs.append(
                     KVCacheGroupSpec(
                         name=group.name,
-                        layer_ids=group.layer_ids,
+                        **split(group),
                         num_kv_heads=group.num_kv_heads,
                         head_dim=group.head_dim,
                         sliding_window=None,
@@ -519,7 +532,7 @@ def kv_cache_group_specs(self) -> Tuple[KVCacheGroupSpec, ...]:
                 specs.append(
                     KVCacheGroupSpec(
                         name=group.name,
-                        layer_ids=group.layer_ids,
+                        **split(group),
                         num_kv_heads=group.num_kv_heads,
                         head_dim=group.head_dim,
                         sliding_window=group.sliding_window,

python/freetoken/models/gemma4/attention.py

@@ -19,8 +19,8 @@ class Gemma4Attention(BaseOP):
 
     On E2B / E4B the KV-shared tail projects q only and attends over the k, v of the last
     earlier layer of its attention type. That layer hands its rows over in ``shared_kv``, a
-    dict living for one forward, so a captured graph records the same tensors. Each shared
-    layer still stores them in its own KV slab; reading the source's slab would save that memory.
+    dict living for one forward, so a captured graph records the same tensors. The KV pool
+    maps a shared layer onto its source's slab and skips its store (``kv_shared_layers``).
     """
 
     def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""):
@@ -34,10 +34,13 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""):
         self.num_kv_heads = group.num_kv_heads
         self.num_qo_heads = config.num_qo_heads
         self.k_eq_v = isinstance(group, FullAttentionGroupConfig) and group.k_eq_v
+        # the pool maps the same pairs, so a shared layer reads the rows of the slab it maps onto
+        sources = dict(config.kv_shared_layers)
+        self.kv_source_layer = sources.get(layer_id)
+        self.kv_source = layer_id in sources.values()
         e_args = config.gemma4_e_args
-        self.kv_shared = e_args is not None and e_args.is_kv_shared(layer_id)
-        self.kv_source = e_args is not None and e_args.is_kv_source(layer_id, group.layer_ids)
-        self._group_name = group.name
+        # the MTP draft (first_kv_shared_layer 0) has no source of its own: it reads the target's slabs
+        self.kv_shared = self.kv_source_layer is not None or (e_args is not None and e_args.first_kv_shared_layer == 0)
 
         self.q_dim = self.num_qo_heads * self.head_dim
         self.kv_dim = self.num_kv_heads * self.head_dim
@@ -95,7 +98,7 @@ def _apply_rope(
         return q_view.view_as(q), k_view.view_as(k)
 
     @nvtx_annotate("MHA")
-    def forward(self, x: torch.Tensor, shared_kv: dict[str, tuple[torch.Tensor, torch.Tensor]] | None) -> torch.Tensor:
+    def forward(self, x: torch.Tensor, shared_kv: dict[int, tuple[torch.Tensor, torch.Tensor]] | None) -> torch.Tensor:
         ctx = get_global_ctx()
         positions = ctx.batch.positions
         T = x.shape[0]
@@ -106,7 +109,7 @@ def forward(self, x: torch.Tensor, shared_kv: dict[str, tuple[torch.Tensor, torc
             # the rope kernels rotate a key along with the query: give them one head of zeros
             q, _ = self._apply_rope(positions, q, q.new_zeros(T, 1, self.head_dim))
             # absent means the read-only draft: it attends the target's paged slabs and contributes no rows
-            k, v = shared_kv[self._group_name] if shared_kv is not None else (None, None)
+            k, v = shared_kv[self.kv_source_layer] if shared_kv is not None else (None, None)
         else:
             qkv = self.qkv_proj.forward(x)
             q_lin, k_lin, v_lin = qkv.split(
@@ -127,7 +130,7 @@ def forward(self, x: torch.Tensor, shared_kv: dict[str, tuple[torch.Tensor, torc
             k = k.reshape(T, self.num_kv_heads * self.head_dim)
             v = v.reshape(T, self.num_kv_heads * self.head_dim)
             if self.kv_source:
-                shared_kv[self._group_name] = (k, v)
+                shared_kv[self.layer_id] = (k, v)
         o = ctx.attn_backend.forward(
             q.contiguous(),
             None if k is None else k.contiguous(),

python/freetoken/models/gemma4/config.py

@@ -58,8 +58,20 @@ class Gemma4EArgs:
     def is_kv_shared(self, layer_id: int) -> bool:
         return layer_id >= self.first_kv_shared_layer
 
-    def is_kv_source(self, layer_id: int, group_layer_ids: tuple[int, ...]) -> bool:
-        return layer_id == max((i for i in group_layer_ids if i < self.first_kv_shared_layer), default=None)
+    def kv_shared_layers(self, *groups_layer_ids: tuple[int, ...]) -> tuple[tuple[int, int], ...]:
+        """``ModelConfig.kv_shared_layers``: each shared layer paired with the last earlier layer of its group."""
+        pairs = []
+        for layer_ids in groups_layer_ids:
+            shared = [i for i in layer_ids if self.is_kv_shared(i)]
+            if not shared:
+                continue
+            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"
+                )
+            pairs += [(i, source) for i in shared]
+        return tuple(sorted(pairs))
 
 
 def gemma4_e_args(
@@ -201,6 +213,15 @@ def parse_config(hf_config: Any) -> ModelConfig:
     # checkpoint's own ignore-list.
     expert_quant = detect_expert_quant(top_cfg)
     dense_quant = "nvfp4" if (expert_quant == "nvfp4" and not num_experts > 0) else "none"
+    # Gemma4UnifiedTextConfig omits the E-series fields
+    e_args = gemma4_e_args(
+        cfg.num_hidden_layers,
+        getattr(cfg, "hidden_size_per_layer_input", None) or 0,
+        getattr(cfg, "vocab_size_per_layer_input", None) or 0,
+        getattr(cfg, "num_kv_shared_layers", None) or 0,
+        num_experts,
+        cfg.intermediate_size * (2 if getattr(cfg, "use_double_wide_mlp", False) else 1),
+    )
 
     return ModelConfig(
         num_layers=cfg.num_hidden_layers,
@@ -229,15 +250,8 @@ def parse_config(hf_config: Any) -> ModelConfig:
         embedding_scale=float(cfg.hidden_size) ** 0.5,
         vision_config=_parse_vision_config(top_cfg, cfg.hidden_size),
         image_token_id=getattr(top_cfg, "image_token_id", None),
-        # Gemma4UnifiedTextConfig omits the E-series fields
-        gemma4_e_args=gemma4_e_args(
-            cfg.num_hidden_layers,
-            getattr(cfg, "hidden_size_per_layer_input", None) or 0,
-            getattr(cfg, "vocab_size_per_layer_input", None) or 0,
-            getattr(cfg, "num_kv_shared_layers", None) or 0,
-            num_experts,
-            cfg.intermediate_size * (2 if getattr(cfg, "use_double_wide_mlp", False) else 1),
-        ),
+        gemma4_e_args=e_args,
+        kv_shared_layers=e_args.kv_shared_layers(full_layer_ids, sliding_layer_ids) if e_args is not None else (),
         attention_groups=(
             FullAttentionGroupConfig(
                 name="full",

python/freetoken/models/gemma4/gguf.py

@@ -312,6 +312,7 @@ def g(key: str):
         vision_config=vision_config,
         image_token_id=getattr(shim, "image_token_id", None),
         gemma4_e_args=e_args,
+        kv_shared_layers=e_args.kv_shared_layers(full_layer_ids, swa_layer_ids) if e_args is not None else (),
         attention_groups=(
             FullAttentionGroupConfig(
                 name="full",

python/freetoken/models/gemma4/model.py

@@ -65,7 +65,7 @@ def forward(
         self,
         x: torch.Tensor,
         per_layer_input: torch.Tensor | None,
-        shared_kv: dict[str, tuple[torch.Tensor, torch.Tensor]],
+        shared_kv: dict[int, tuple[torch.Tensor, torch.Tensor]],
     ) -> torch.Tensor:
         # --- attention sandwich ---
         residual = x
@@ -131,7 +131,7 @@ def _per_layer_inputs(self, input_ids: torch.Tensor, x: torch.Tensor) -> torch.T
     def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
         x = embed_input_ids(self.embed_tokens, input_ids, get_global_ctx().batch)
         per_layer = self._per_layer_inputs(input_ids, x) if self._has_per_layer_inputs else None
-        shared_kv: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
+        shared_kv: dict[int, tuple[torch.Tensor, torch.Tensor]] = {}
         for i, layer in enumerate(self.layers.op_list):
             x = layer.forward(x, None if per_layer is None else per_layer[:, i], shared_kv)
         return self.norm.forward(x)

python/freetoken/spec/mtp/shared_kv.py

@@ -48,20 +48,12 @@ def __init__(self, target: BaseKVCachePool, draft_config: ModelConfig, target_co
         # llama.cpp's share callback: draft swa -> the target's last SWA layer, draft full -> its last full layer
         last_swa = max(swa_ids) if swa_ids else None
         last_full = max(full_ids) if full_ids else None
-        e_args = getattr(target_config, "gemma4_e_args", None)
+        # a KV-shared target layer owns no slab: read its source's, as the target's pool maps it
+        sources = dict(target_config.kv_shared_layers)
         self._route: list[int] = []
         for layer_id in range(draft_config.num_layers):
             layer = last_swa if draft_config.is_swa_layer(layer_id) else last_full
-            # the hop is a no-op while shared target layers copy their rows into their own slab
-            if e_args is not None and e_args.is_kv_shared(layer):
-                sources = [i for i in target_config.attention_group_for_layer(layer).layer_ids if i < e_args.first_kv_shared_layer]
-                if not sources:
-                    raise RuntimeError(
-                        f"target layer {layer} is KV-shared but its attention group has no source below the "
-                        f"sharing boundary {e_args.first_kv_shared_layer} to route draft layer {layer_id}"
-                    )
-                layer = max(sources)
-            self._route.append(layer)
+            self._route.append(sources.get(layer, layer))
 
     def k_cache(self, index: int) -> torch.Tensor:
         return self._target.k_cache(self._route[index])

tests/models/gemma4/test_e_series.py

@@ -100,10 +100,56 @@ def test_kv_source_is_the_last_unshared_layer_of_each_attention_type():
     """E2B shares from layer 15: full attention hands over from 14, sliding from 13."""
     config = parse_config(_hf_config(E2B))
     full, swa = config.attention_groups
-    args = config.gemma4_e_args
+    sources = dict(config.kv_shared_layers)
 
-    assert [i for i in full.layer_ids if args.is_kv_source(i, full.layer_ids)] == [14]
-    assert [i for i in swa.layer_ids if args.is_kv_source(i, swa.layer_ids)] == [13]
+    assert sorted(sources) == list(range(15, config.num_layers))
+    assert {sources[i] for i in full.layer_ids if i in sources} == {14}
+    assert {sources[i] for i in swa.layer_ids if i in sources} == {13}
+
+
+def test_kv_shared_layers_with_no_earlier_source_are_refused():
+    from freetoken.models.gemma4.config import Gemma4EArgs
+
+    args = Gemma4EArgs(per_layer_hidden_size=256, per_layer_vocab_size=262144, first_kv_shared_layer=4, kv_shared_intermediate_size=1)
+    with pytest.raises(NotImplementedError, match="no earlier layer of their attention type"):
+        args.kv_shared_layers((0, 1, 2, 3), (4, 5))
+
+
+@pytest.mark.parametrize("kv_format", ["none", "int8_per_token_head"])
+@pytest.mark.parametrize("release, full_slabs, swa_slabs", [(E2B, 3, 12), (E4B, 4, 20)])
+def test_shared_layers_read_their_source_slab_and_own_none(release, full_slabs, swa_slabs, kv_format):
+    from freetoken.kvcache import create_kvcache_pool
+    from freetoken.kvcache.format import KV_FORMATS
+
+    config = parse_config(_hf_config(release))
+    pool = create_kvcache_pool(
+        model_config=config, num_pages=2, page_size=16, num_swa_tokens=8,
+        dtype=torch.bfloat16, device=torch.device("cpu"), kv_format=KV_FORMATS[kv_format],
+    )
+
+    assert pool.full_kv_pool.buffer.shape[1] == full_slabs
+    assert pool.swa_kv_pool.buffer.shape[1] == swa_slabs
+    for i, source in config.kv_shared_layers:
+        assert pool.k_cache(i).data_ptr() == pool.k_cache(source).data_ptr()
+        assert pool.v_cache(i).data_ptr() == pool.v_cache(source).data_ptr()
+        if kv_format != "none":
+            # a shared layer reading zero scales would only show up as worse 8-bit outputs
+            assert pool.k_scale(i).data_ptr() == pool.k_scale(source).data_ptr()
+            assert pool.v_scale(i).data_ptr() == pool.v_scale(source).data_ptr()
+
+
+@pytest.mark.parametrize("release", [E2B, E4B])
+def test_attention_reads_the_pairs_the_pool_maps(release):
+    """A shared layer attends over the rows of the source the pool maps its slab onto."""
+    config = parse_config(_hf_config(release))
+    sources = dict(config.kv_shared_layers)
+    model = _meta_model(config)
+
+    for i, layer in enumerate(model.model.layers.op_list):
+        attn = layer.self_attn
+        assert (attn.kv_shared, attn.kv_source_layer) == (i in sources, sources.get(i))
+        assert attn.kv_source == (i in sources.values())
+        assert hasattr(attn, "q_proj") == attn.kv_shared
 
 
 def test_model_shapes_match_the_e2b_checkpoint():
@@ -139,7 +185,8 @@ def test_loader_drops_shared_k_v_but_not_the_vision_towers(tmp_path, monkeypatch
     """The checkpoint ships k, v and k_norm for shared layers too; the tower's layers reuse those names and indices."""
     from freetoken.models.gemma4 import weight
 
-    release = dict(E2B, layers=3, shared=1, period=3)
+    # layers 0 and 2 slide, 1 is full; the shared layer 2 reads layer 0
+    release = dict(E2B, layers=3, shared=1, period=2)
     prefix = "model.language_model.layers"
     layers = {
         f"{prefix}.{i}.self_attn.{leaf}.weight": torch.zeros(2, 2)
@@ -230,7 +277,7 @@ def test_a_fused_qkv_file_gives_a_kv_shared_layer_its_q_rows(monkeypatch):
     from freetoken.models.gemma4 import gguf as gemma4_gguf
     import freetoken.utils
 
-    release = dict(E2B, layers=3, shared=1, period=3, hidden=8, ffn=16, double_wide=False)
+    release = dict(E2B, layers=3, shared=1, period=2, hidden=8, ffn=16, double_wide=False)
     config = parse_config(_hf_config(release, vocab_size=32, vocab_size_per_layer_input=32, hidden_size_per_layer_input=4))
     H, P, L, V = 8, 4, 3, 32
     index = {"token_embd.weight": _f32((H, V)), "output_norm.weight": _f32((H,)),
@@ -281,6 +328,7 @@ def _draft_view_config(target):
         target,
         num_layers=4,
         gemma4_e_args=None,
+        kv_shared_layers=(),
         attention_groups=(
             dataclasses.replace(swa, layer_ids=(0, 1, 2)),
             dataclasses.replace(full, layer_ids=(3,)),
@@ -310,3 +358,21 @@ def test_read_only_draft_view_routes_to_the_shared_target_slabs(release, overrid
     target = parse_config(_hf_config(release, **overrides))
     view = ReadOnlyTargetKVView(_SlabProbe(), _draft_view_config(target), target)
     assert [view.k_cache(i) for i in range(4)] == expected
+
+
+def test_draft_layers_project_q_only_and_hand_over_nothing():
+    """The sidecar's own e_args shares every layer from the first; it has no source of its own."""
+    import dataclasses
+
+    from freetoken.models.gemma4.attention import Gemma4Attention
+    from freetoken.models.gemma4.config import Gemma4EArgs
+    from freetoken.utils.torch_utils import torch_dtype
+
+    target = parse_config(_hf_config(E4B))
+    draft = dataclasses.replace(
+        _draft_view_config(target), gemma4_e_args=Gemma4EArgs(0, 0, 0, target.intermediate_size)
+    )
+    with torch.device("meta"), torch_dtype(torch.bfloat16):
+        layers = [Gemma4Attention(draft, i, prefix=f"layers.{i}.self_attn") for i in range(4)]
+
+    assert [(a.kv_shared, a.kv_source, hasattr(a, "qkv_proj")) for a in layers] == [(True, False, False)] * 4