Cross-layer KV sharing

def _source_map(source_ids: list[int], n_layer: int) -> dict[int, int]:
    """Map each layer to the nearest source layer that is not later than that layer."""
    ids = sorted(source_ids)
    return {layer: max(s for s in ids if s <= layer) for layer in range(ids[0], n_layer)}


def _shared_layer_cfg(text: dict[str, Any]) -> AttrDict:
    n_layer = text["num_hidden_layers"]
    candidate_source = text["candidate_source_layer_id"]
    shared = AttrDict(
        layer_index_module_map=_source_map(text["index_source_layer_ids"], n_layer),
        layer_kv_map=_source_map(text["kv_source_layer_ids"], n_layer),
        layer_index_k_map=_source_map(text.get("index_k_source_layer_ids", text["kv_source_layer_ids"]), n_layer),
    )
    if candidate_source is not None:
        shared.layer_candidate_source_map = {layer: candidate_source for layer in range(candidate_source, n_layer)}
        shared.candidate_source_cfg = {
            candidate_source: AttrDict(
                candidate_topk_blocks=text["candidate_topk_blocks"],
                candidate_block_size=text["candidate_block_size"],
                topk_block_score_reduce="max",
            )
        }
    return shared

        raw = shared_cfg.get(name)
        # The default value of layer_index_k_map is layer_index_module_map.
        if raw is None and name == "layer_index_k_map":
            raw = shared_cfg.get("layer_index_module_map")
        assert raw is not None, f"shared layer config missing {name}"
        assert all(0 <= layer_id < num_layers and 0 <= src < num_layers for layer_id, src in raw.items()), (
            f"{name} contains invalid layer ids or sources, num_layers={num_layers}"
        )
        layer_map = [raw.get(layer_id, layer_id) for layer_id in range(num_layers)]
        assert all(layer_map[layer_id] == layer_id for layer_id in range(n_main_layer, num_layers)), (
            f"{name} sharing on MTP layers is not supported"
        )
        return layer_map

        # ---- Index Attention ----
        # NOTE index K 是否复用只由 index-K map 决定，和整个 index module / main kv map 解耦：
        #   index_mapped_layer_id != layer_id 时复用别层 indexer 的 ica_cache，本层不重复存 index K；
        #   否则本层拥有 index kv；只有独立 compressor indexer 需要 index state。
        if GPUCacheInterface._is_sparse(attn_module):
            if index_mapped_layer_id != layer_id:
                # share-index-K：复用别层 indexer 的 K cache，不推进 kv offset / 不 alloc index state
                ica_cache = cache_list[index_mapped_layer_id].ica_cache
            else:
                token_per_block = block_size // ratio if ratio > 0 else block_size
                delta = token_per_block * self._qk_topk_bytes
                if gpu_blocks is not None:
                    ica_cache = gpu_blocks[:, self._kv_offset : self._kv_offset + delta].view(-1, token_per_block, 1, self._qk_topk_bytes)
                padded_delta = _align_to_multiple(delta, 1, 32)
                self._kv_offset += padded_delta
                kv_length += padded_delta

                if _needs_compress_state(ratio) and _needs_index_state(attn_module):
                    num_state_token = self._state_block.all_compress_block_size[ratio]
                    ica_state = self._alloc_state(num_state_token, self._qk_topk_head_dim)

        # ---- Main Attention ----
        # NOTE main（compress）kv 是否复用只由 kv map 决定，和 index map 解耦。
        if ratio > 0:
            if kv_mapped_layer_id != layer_id:
                ca_cache = cache_list[kv_mapped_layer_id].ca_cache
            else:
                # 所有压缩率（c1/c2）统一处理
                token_per_block = block_size // ratio
                delta = token_per_block * self._bytes_per_token
                if gpu_blocks is not None:
                    ca_cache = gpu_blocks[:, self._kv_offset : self._kv_offset + delta].as_strided(
                        (gpu_blocks.shape[0], token_per_block, 1, self._bytes_per_token),
                        (gpu_blocks.stride(0), self._bytes_per_token, self._bytes_per_token, 1),
                    )
                padded_delta = _align_to_multiple(delta, 1, 32)
                self._kv_offset += padded_delta
                kv_length += padded_delta

                if _needs_compress_state(ratio):
                    num_state_token = self._state_block.all_compress_block_size[ratio]
                    ca_state = self._alloc_state(num_state_token, self._compress_dim)

        # The sharing of kv and index (topk) can be decoupled: the two sources are not necessarily the same layer.
        self.use_shared_kv = shared_topology is not None
        if shared_topology is None:
            self.is_kv_source = True
            self.is_index_source = True
            self.is_index_k_source = True
        else:
            self.is_kv_source = shared_topology.is_kv_source_layer(layer_idx)
            self.is_index_source = shared_topology.is_index_source_layer(layer_idx)
            self.is_index_k_source = shared_topology.is_index_k_source_layer(layer_idx)

        elif self.is_sparse:
            kv_ca_index = cross_layer_runtime.topk_indices
            assert kv_ca_index is not None, f"layer {self.layer_idx}: index-source topk indices not stored"

        if self.has_local_index:
            cross_layer_runtime.topk_indices = kv_ca_index