Reuse IPC descriptors across CP shared-KV layers

CP shared-KV slot remaps already have forward-batch lifetime, but the IPC materialize path rebuilt owner/source/dense descriptor tensors on every layer. Cache prefix/current IPC descriptors on the token and paged slot-remap objects, keyed by layout, spans, device, descriptor kind, and prefix capacity, so all model layers can reuse the same request/batch-plan descriptors.

Constraint: Small-extend cache-hit workloads showed descriptor setup could exceed the all-reduce baseline before any IPC kernel work ran.

Rejected: Global descriptor cache | slot-remap lifetime is safer and avoids stale entries across request/batch-plan changes.

Rejected: Cache without physical page capacity | prefix descriptors encode capacity-invalid pages and must miss when capacity changes.

Confidence: high

Scope-risk: moderate

Directive: Do not reuse descriptors across different slot_logical_pages identity, CP layout, spans, device, or prefix capacity; stale descriptors can alias dense slots across requests.

Tested: Local py_compile; local git diff --check; remote g0034 cjy-glm5-new targeted descriptor tests 2 passed; remote full test_cp_shared_kv_runtime.py 146 passed, 21 warnings, 2 subtests passed.

Not-tested: Full ETE throughput/accuracy after descriptor cache; CUDA service benchmark still needed to quantify speedup.
(cherry picked from commit addd1ca1571e41458315d15304a0e841682fe8fa)

docs/advanced_features/nsa_prefill_cp_shared_kv_ipc_collective_replacement_plan_zh.md

@@ -352,3 +352,91 @@ staging_nbytes >= num_slots * page_nbytes
 - SGLang runtime 单测覆盖 token/index current helper：publish 使用 dense destination pages，peer materialize 使用 compact source pages。
 - `benchmark_cp_shared_kv_ipc_gather.py` 已同步改成 compact current staging 合同；quick smoke：
   `bs=2 cached=4096 extend=1024 fp8/uint8` 下 dense all_reduce p50 0.459ms，IPC compose p50 0.310ms。
+
+### 2026-06-12 小 extend / bs1-2 / 200k context 性能边界
+
+用户指出线上模型 context 上限约 200k token。因此早先用 `cached=307200` 得到的 scaling 结论只能说明 kernel 趋势，不能作为线上策略依据。
+
+远端用 `benchmark_cp_shared_kv_ipc_gather.py --cache-hit-only` 按 `cached + extend <= 200k` 重新测了 bs=1/2、小 extend。关键结果：
+
+```text
+bs=1 cached=65536  extend=256..4096: all_reduce p50 0.286..0.297ms, IPC p50 0.359..0.361ms  => IPC 更慢
+bs=1 cached=102400 extend=256..4096: all_reduce p50 0.414..0.437ms, IPC p50 0.411..0.417ms  => 基本持平/略快
+bs=1 cached=160000 extend=256..4096: all_reduce p50 0.622..0.634ms, IPC p50 0.521..0.524ms  => IPC 明显更快
+bs=1 cached=190000 extend=256..4096: all_reduce p50 0.731..0.746ms, IPC p50 0.566..0.572ms  => IPC 明显更快
+bs=2 cached=65536  extend=256..4096: all_reduce p50 0.518..0.544ms, IPC p50 0.455..0.477ms  => IPC 更快
+bs=2 cached=102400 extend=256..4096: all_reduce p50 0.785..0.812ms, IPC p50 0.590..0.610ms  => IPC 明显更快
+```
+
+结论：
+
+1. “小 extend + bs1-2 效果差”在 **bs=1 且 cached 约 64k** 时成立；此时 dense all_reduce 本身只有约 0.29ms，IPC 的固定开销无法摊薄。
+2. **bs=1 cached 约 100k** 是近似 break-even 区间；真实 ETE 若包含 Python descriptor 构造、stream wait、index+MLA 双路径等额外开销，IPC 可能从 kernel 持平变成端到端劣化。
+3. **cached >= 160k 或 bs=2** 时，纯 materialize kernel 仍显示 IPC 优于 all_reduce。
+4. 下一步不能只看 kernel microbenchmark，需要补 runtime-style benchmark，把 descriptor 构造、Python/Torch tensor 准备、index/MLA 两条路径、stream 同步一起计入。
+5. 可能的工程策略不是简单回退，而是：
+   - 做 fused current fill+publish，去掉 current path 的额外 HBM copy / kernel 固定开销；
+   - 对 bs=1 + cached 较小的区间建立显式 cost model，必要时保留 all_reduce fast path，但必须是可解释的策略选择，不是 silent fallback。
+
+### 2026-06-12 runtime-style benchmark 补充
+
+TAI benchmark 已新增 runtime-overhead 模式：
+
+```bash
+PYTHONPATH=python torchrun --standalone --nproc_per_node=8 \
+  benchmark/nsa_prefill/benchmark_cp_shared_kv_ipc_gather.py \
+  --runtime-overhead-only \
+  --cache-hit-cached-tokens 65536 102400 160000 190000 \
+  --cache-hit-extend-tokens 256 1024 4096 \
+  --cache-hit-batch-requests 1 2 \
+  --cache-hit-max-context-tokens 200000 \
+  --dtype uint8 --kv-dim 656 --warmup 2 --repeat 5 --no-check
+```
+
+其中 `--cache-hit-max-context-tokens` 按 per-request context 过滤无效点，避免把超过模型 200k 上限的 benchmark 混进结论。
+
+新增输出路径：
+
+- `cache_hit_dense_all_reduce_full`：dense all_reduce baseline。
+- `cache_hit_runtime_descriptor_setup_only`：每次 forward 重建 owner/src/dst descriptor tensor 的 Python/Torch/H2D 成本。
+- `cache_hit_runtime_ipc_prefix_current_compose`：descriptor 重建 + prefix IPC materialize + current compact publish/wait-gather。
+
+关键结果：
+
+```text
+bs=1 cached=65k:  all_reduce ~0.34ms, descriptor setup ~0.46ms, runtime IPC ~0.84ms
+bs=1 cached=102k: all_reduce ~0.45ms, descriptor setup ~0.61ms, runtime IPC ~1.05ms
+bs=1 cached=160k: all_reduce ~0.65ms, descriptor setup ~0.87ms, runtime IPC ~1.40ms
+bs=1 cached=190k: all_reduce ~0.75ms, descriptor setup ~1.00ms, runtime IPC ~1.58ms
+bs=2 cached=65k:  all_reduce ~0.55ms, descriptor setup ~0.74ms, runtime IPC ~1.22ms
+bs=2 cached=102k: all_reduce ~0.81ms, descriptor setup ~1.05ms, runtime IPC ~1.66ms
+```
+
+这解释了线上“小 extend / bs1-2 / cache hit”效果差的主要来源：纯 TAI IPC kernel 在较大 prefix 下可以比 all_reduce 快，但当前 runtime 若每次 forward 都重建 descriptor tensor，固定开销已经超过 all_reduce 本身。该 benchmark 是一个上界模型：它刻意把 descriptor 构造放进 timed region，用于暴露未缓存 descriptor 的最坏 hot-path 成本。
+
+后续优化优先级：
+
+1. descriptor 需要 request/batch-plan 级缓存或复用，不能每层/每次 materialize 从 Python list 重建 GPU tensor。
+2. current path 仍需要 fused current fill+publish，减少额外 HBM copy 和 kernel launch。
+3. 在 descriptor 缓存完成前，不应默认假设 IPC 替换 all_reduce 会改善 bs1/2 小 extend ETE；需要显式 cost model 或 gate。
+
+### 2026-06-12 runtime descriptor 跨 layer 复用实现
+
+已把 prefix/current IPC slot descriptor 缓存在 `SharedTokenKVSlotRemap` 与 `SharedPagedBufferSlotRemap` 上，生命周期与 forward batch 的 slot remap 一致。这样同一个 request/batch-plan 在多层 forward 中不会每层重复构造 `slot_indices / owner_ranks / src_page_indices / dense_page_indices` GPU tensor。
+
+缓存 key 包含：descriptor 类型（prefix/current）、cache 类型（token/paged）、`slot_logical_pages` storage identity、CP layout、merged slot spans、device，以及 prefix 的 physical page capacity。capacity 被纳入 key 是为了避免 host/L1 capacity 变化时错误复用已经标 invalid 的 source page descriptor。
+
+运行时路径：
+
+- MLA KV prefix：`materialize_prefix_and_reuse_current_kv_page_slots` -> `_get_or_build_prefix_ipc_slot_descriptors`。
+- MLA KV current：同一函数的 current staging IPC 路径 -> `_get_or_build_current_ipc_slot_descriptors`。
+- index prefix/current：`materialize_prefix_and_reuse_current_index_page_slots` 走同一套 descriptor cache，但用 `cache_kind="paged"` 与 token 描述符隔离。
+
+当前实现只复用 Python/Torch descriptor tensor，不改变 TAI kernel 合同，也不新增 collective。后续仍需要 fused current fill+publish 来降低 current path 的额外 HBM copy / kernel launch。
+
+验证：
+
+```text
+远端 cjy-glm5-new：PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
+146 passed, 21 warnings, 2 subtests passed
+```