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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user