[Feature]: UCM Store integration with HiCache for Hybrid models

关联实现PR：https://github.com/ModelEngine-Group/unified-cache-management/pull/1392

🚀 The feature, motivation and pitch

SGLang HiCache v2 适配 Hybrid 模型方案

1. 核心结论

UCM 保留 v1 处理“有实体数据的主 KV”，使用 v2 动态处理 Hybrid sidecar pool；每个物理 pool 对应一个 PosixStore，普通 pool 通过原始指针零拷贝传输，只有 MAMBA 将异构状态拼成一个定长向量后传输。

实体主 KV             → v1 store
Hybrid sidecar pool   → v2 动态 store
DeepSeek V4 逻辑 KV   → 不创建 store

需要特别区分：batch_exists_v2() 是统一的 Hybrid 命中查询入口，但这不代表主 KV 一定存在一个 v2 store。普通主 KV 仍复用 v1 store；DeepSeek V4 的逻辑 KV 不查询任何 store。

2. 关键设计决策

2.1 一个物理 pool 对应一个 store

SWA                    → 1 store
MAMBA                  → 1 store
INDEXER                → 1 store
DEEPSEEK_V4_C4         → 1 store
DEEPSEEK_V4_C4_STATE   → 1 store
其他实体 PoolName      → 各 1 store

“一个 pool 一个 store”不等于“一个 pool 只能有一个 pointer”。普通 MHA 的 K、V 是两个独立 pointer，但大小相同，可以共同组成一个 store block。

不同 PoolName 不共享 store，即使 tensor 大小相同，以此隔离数据语义、namespace 和生命周期。

2.2 普通 pool 零拷贝，MAMBA 拼接

• 普通 pool：通过 get_page_buffer_meta() 获取原始 pointer，直接交给 UCM。
• MAMBA：通过 get_data_page(flat=True) 拼接 temporal/SSM/conv 状态。

本方案暂不支持非对称 MHA。非 MAMBA pool 内的非空 component 必须等长，否则注册失败。

2.3 主 KV 不重复建 v2 store

普通 Hybrid 模型保留现有主 KV v1 store。register_mem_host_pool_v2() 收到 PoolName.KV 时跳过，避免同时出现 v1 KV store 和 v2 KV store。

3. 接口职责

v1：实体主 KV

batch_get_v1(keys, host_indices, extra_info)
batch_set_v1(keys, host_indices, extra_info)
batch_exists(keys, extra_info)

普通 MHA 的等长 K、V 通过同一个 block 传输；MLA 主 KV 通常每页只有一个连续区域。

v2：动态 Hybrid pool

register_mem_host_pool_v2(host_pool, pool_name)
batch_exists_v2(keys, pool_transfers, extra_info)
batch_get_v2(transfers, extra_info)
batch_set_v2(transfers, extra_info)

PoolTransfer 携带 PoolName、逻辑 keys、Host Pool indices、索引来源和命中策略。

4. 总体路径

SGLang HybridCacheController
        │
        ├─ 主 KV v1 ──────────────→ 主 PosixStore
        │
        └─ sidecar v2
             └─ UnifiedCacheStore
                  └─ SglangUcmConnector
                       ├─ 普通 pool：原始 pointers，零拷贝
                       └─ MAMBA：flat staging page
                                  │
                                  └─ 每个 pool 一个 PosixStore

• unifiedcache_store.py：对接 SGLang HiCacheStorage 接口。
• ucm_connector.py：创建 store、解析内存元数据、生成 key 并执行 I/O。

5. Pool 注册

5.1 实体主 KV

当 mem_pool_host.kv_buffer is not None 时，初始化阶段根据主 KV page 大小创建 v1 PosixStore。后续 sidecar 再通过 v2 单独注册，原有普通模型行为不变。

5.2 DeepSeek V4 LogicalHostPool

DeepSeek V4 的主 KV entry 是 LogicalHostPool：

kv_buffer = None
size_per_token = 0
只管理逻辑 page 索引
不保存实际 tensor

UCM 对它执行以下处理：

• 不调用 get_size_per_token()。
• 不创建主 KV store。
• 保存 UCM 基础配置，供实体 v2 pool 创建 store。
• logical anchor 的 v1 get/set 作为成功的 no-op。
• 实际数据由 SWA、C4、INDEXER、STATE、C128 等 v2 pool 存取。

no-op 只在以下条件同时成立时生效，不影响普通主 KV：

store is None and mem_pool_host.kv_buffer is None

5.3 sidecar pool

register_pool_v2() 的步骤：

1. 保存 pool_name -> host_pool。
2. 用一个逻辑 page 调用 get_page_buffer_meta() 探测布局。
3. 标准化该 page 的 pointer 和 size 列表。
4. MAMBA 进入拼接路径，其他 pool 进入零拷贝路径。
5. 创建该 pool 专属的 PosixStore。
6. 保存 component size，实际传输时再次校验。

6. 普通 pool：一个 store、多 pointer、零拷贝

假设每页返回：

pointers = [ptr0, ptr1, ..., ptrN]
sizes    = [S,    S,    ..., S]

注册时要求所有非空 size 相同。设 component 数量为 N，PosixStore 配置为：

tensor_size = S
shard_size  = S × N
block_size  = S × N

一个逻辑 key 对应一个 block，整个 page 的 pointer 列表一次性交给同一个 store：

pointers = [
    [page0_ptr0, page0_ptr1],
    [page1_ptr0, page1_ptr1],
]

UCM 直接访问 L2 Host Pool 的原始 buffer，不创建中间 tensor。若 size 不一致则拒绝注册，防止错误布局损坏数据。

7. MAMBA：一个 store、一个拼接向量

MAMBA 内可能包含不同大小的 temporal/SSM/conv 状态，不能作为等长 tensor 直接交给同一 PosixStore。

写入：

get_data_page(offset, flat=True)
  → [temporal bytes | conv0 bytes | conv1 bytes | ...]
  → 固定长度 staging page
  → PosixStore block

加载：

get_dummy_flat_data_page()
  → PosixStore load
  → set_from_flat_data_page(offset, page)
  → 恢复内部状态

该路径增加一次拼接/拆分拷贝，换取一个 pool 只使用一个 store，并保证 page 级完整性。

8. 命中查询：batch_exists_v2()

8.1 普通 Hybrid 模型

batch_exists_v2()
  ├─ 查询主 KV 的 v1 store，得到候选 KV 前缀
  ├─ 查询每个 sidecar v2 store
  └─ 按 PoolHitPolicy 对所有必需 pool 求交集

所以当前 batch_exists_v2() 会间接使用 v1 查询，但不会为主 KV 创建第二个 v2 store。

8.2 DeepSeek V4

batch_exists_v2()
  ├─ logical KV：不查询 store，只提供全部候选 key 范围
  ├─ 查询 SWA/C4/INDEXER/STATE/C128 等实体 v2 stores
  └─ 对实体 pool 命中结果求交集

“logical KV 提供全部候选范围”不代表已经命中。最终可恢复范围完全由必需的实体 v2 pool 决定。

8.3 HitPolicy

• ALL_PAGES：要求从头开始连续存在。
• TRAILING_PAGES：按依赖 pool 的尾部 page 范围判断，例如从 SWA 索引派生的 state pool。

同一 pool 的全部 component 位于同一个 block，一次 lookup 即可表示该 pool page 是否完整存在。

9. 写入与加载顺序

HybridCacheController 的备份顺序是：

batch_set_v2(sidecar pools)
  → batch_set_v1(primary KV)

• 普通模型：v2 sidecar 和 v1 主 KV 都写实际数据。
• DeepSeek V4：v2 sidecar 写实际数据，logical KV 的 v1 调用成功 no-op。

加载时，普通 pool 直接读入 L2 原始 pointer；MAMBA 先读入 staging page 再拆分；DeepSeek V4 logical KV 的 v1 调用不搬运数据。

10. Key 与 namespace

v2 物理 key 的逻辑组成：

{logical_key}{model_config_suffix}
__v2_tp_{tp_rank}_{tp_size}
__pool_{pool_name}
__component_0

随后转换为 MD5 二进制 key。普通 MHA 即使有 K、V 两个 pointer，也只使用一个 block key。

Store namespace：

sglang_v2/{pool_name}_{component_size}x{component_count}

例如 deepseek_v4_c4_786240x1。x1/x2 表示一个 block 内的 component 数量，不表示 TP、page 或副本数。

11. DeepSeek V4 Flash 实例

当前模型每个 TP rank 创建 6 个实体 store：

Pool namespace	Page 大小	Store/TP
swa_6589440x1	44 × 149760	1
deepseek_v4_c4_786240x1	21 × 37440	1
deepseek_v4_c4_indexer_177408x1	21 × 8448	1
deepseek_v4_c4_state_2752512x1	21 × 131072	1
deepseek_v4_c4_indexer_state_688128x1	21 × 32768	1
deepseek_v4_c128_34560x1	20 × 1728	1

12. 保护与限制

• 仅支持 page_first Host Pool 布局。
• 暂不支持 Asymmetric MHA。
• MAMBA 不是零拷贝路径。
• 不同 PoolName 不合并 store。
• 注册后 component 数量和大小必须保持不变。
• 不支持store的GC配置

13. 串讲总结

这次适配没有把所有缓存强制迁移到 v2，而是在保留主 KV v1 兼容性的基础上，用 v2 解决 Hybrid sidecar 类型和数量动态变化的问题。存储粒度统一为“一个物理 pool 一个 PosixStore”：普通 pool 将一个 page 的等长原始 pointers 直接组成一个 block，实现零拷贝；MAMBA 因内部状态大小不一致，先拼接成固定长度 page。DeepSeek V4 的主 KV 只是逻辑索引，不创建 store，也不参与数据搬运，最终命中和恢复由各实体 v2 pool 共同决定。

Alternatives

No response

Additional context

No response

Start with unifiedcache_store.py and ucm_connector.py, then trace SglangUcmConnector and register_pool_v2(). The linked implementation PR is the primary reference. Done means Hybrid sidecar pools use per-pool v2 stores while primary KV retains v1 behavior, with the documented MAMBA and DeepSeek V4 exceptions supported.