Reuse CP shared KV remaps across layer materialization

CP shared KV materialization repeatedly rebuilt the same logical-page slot remaps and page inverse metadata for each layer. Cache the token and paged remap metadata on the forward batch so MLA KV, index K/scale, and prefetch paths can reuse the layer-independent mapping while still materializing layer-specific data through the existing tai/torch runtime paths.

Constraint: Only mapping metadata is batch-scoped; dense KV/index contents remain layer-specific and are not reused.
Rejected: Cache fully materialized dense KV/index buffers | would add large per-layer memory residency and invalidation complexity.
Confidence: medium
Scope-risk: moderate
Directive: Do not assume this removes materialize or CP all-reduce cost; profile tai fallback logs and Nsight kernels before attributing E2E gains or losses.
Tested: git diff --check
Tested: remote g0034 container PYTHONPATH=python python3 -m pytest test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -q (52 passed, 5 warnings)
Not-tested: Full GLM-5 disaggregated E2E performance run
This commit is contained in:
laoyao0822
2026-05-12 20:02:51 +08:00
parent 099bcfb41e
commit c5c30a3f50
7 changed files with 476 additions and 25 deletions
@@ -18,6 +18,14 @@ Layer L attention 返回时不强制等待 prefetch
Layer L+1 consume prefetched buffer 时 wait event,再补齐 suffix/current pages
```
后续在 HiCache 分支补充了 batch-scoped slot remap 复用:`metadata.real_page_table`
和 index page table 在同一个 forward 内跨 layer 不变,因此只在第一次 materialize /
prefetch 时构造 `slot_logical_pages``page_inverse`、paged `dense_pages`,后续 layer
复用该计划;每层仍重新 materialize 本层 KV/index 数据并执行 CP all-reduce。这个优化
只复用“逻辑页到 dense slot 的映射计划”,不复用 dense KV 内容,避免跨层 KV 数据串用。
如果同一个 `ForwardBatch` 已经带有 remap cache 但无法复用(key mismatch 或 cache
状态不完整),运行时会打印 rate-limited warning;第一次冷启动构造 cache 不打日志。
保留的环境变量:
```text