Reduce CP shared-KV batch overhead without reverting bs1 planning

CP shared-KV bs>1 exposed two separate overhead sources: HiCache load-back could thrash near capacity, and partial-current sync compose could all-reduce row-major page-table gaps between request prefixes. Keep the intended batch-plan path for bs=1, but make host/L1 free-room handling less reactive and teach the MLA/index sync compose path to use exact per-request prefix slot spans instead of one bounding span.\n\nThe exact-span path preserves the single-span IPC fast path for bs=1/single-span cases, while avoiding over-communication for heterogeneous cache-hit batches. The HiCache metadata tests cover host/L1 free-room propagation and load-back batching behavior.\n\nConstraint: bs=1 using CPSharedKVBatchPlan is the expected steady-state path and must not be treated as a regression.\nConstraint: Remote production-like validation runs inside g0034 container /sgl-workspace/sglang-tai.\nRejected: Disable batch-plan for bs=1 | user confirmed this is intended behavior and it would hide the actual bs>1 overhead.\nRejected: Keep one bounding prefix span for bs>1 | row-major page tables can include current/gap slots and inflate per-layer all-reduce work.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not replace exact prefix spans with a single row-major bounding span unless ETE data proves collective launch count dominates gap over-communication.\nTested: g0034 docker py_compile for changed Python/test files.\nTested: g0034 docker PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py test/registered/unit/mem_cache/test_cp_hicache_metadata.py — 241 passed, 5 warnings, 2 subtests passed.\nNot-tested: Full ETE bs>1 throughput after the exact-prefix-span change.\nNot-tested: CUDA kernel benchmark and live traffic replay.
2026-06-06 02:03:16 +08:00
parent 1b99de7459
commit 81bfd4bfd3
10 changed files with 544 additions and 92 deletions
@@ -230,3 +230,49 @@ Implications:
  separate ratios because L1 hit value and host capacity pressure are different.
 - Free-room alone is not sufficient for contiguous allocation. `HostKVCache` and
  CP owner-lane device allocation also need contiguous-preferred selection.
+
+## 2026-06-05 correction: load-back must not synchronously evict for L1 free room
+
+Remote evidence from `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260605_152940.log`
+showed prefill slowdown and death caused by a CP0 HiCache load-back planning
+straggler:
+
+```text
+node_id=677 load_back CP starts on all CP ranks at 15:48:26
+CP1-7 report SUCCESS by 15:48:27
+CP0 reports SUCCESS at 15:48:49
+health check fires at 15:48:47 due no detokenizer heartbeat for 20s
+```
+
+Important detail: `HiCache-load load_back CP SUCCESS` is emitted before
+`cache_controller.start_loading()`, so this 23s gap is not the actual L2->L1 H2D
+copy.  It is in the scheduler-side admission/planning path:
+
+```text
+HiRadixCache.load_back
+  -> _build_cp_load_back_plan
+  -> optional _evict_cp_load_back_owner_lanes
+  -> HiCacheController.load_cp
+      -> alloc_pages_with_owners
+      -> build host/device index descriptors
+```
+
+The failing node had exact capacity available but non-zero L1 free-room deficit:
+
+```text
+required_by_owner=[41,45,45,45,46,46,46,46]
+available_by_owner=[221,194,197,198,237,285,290,308]
+free-room deficit on lanes 1-3
+```
+
+Therefore the reactive free-room rule must be different for synchronous
+load-back admission:
+
+- exact owner-lane deficit is blocking and may trigger synchronous eviction;
+- free-room deficit is advisory/observability for load-back and must not force
+  synchronous eviction when exact capacity already fits.
+
+This preserves correctness while avoiding heavy eviction planning on the cache-hit
+scheduler hot path.  Free-room maintenance for L1 should happen on real extend
+allocation pressure or a later non-blocking/proactive path, not inside a
+cache-hit load-back that already has enough exact capacity.
@@ -857,6 +857,58 @@ runtime / kernel 都消费 CPSharedKVBatchPlan descriptors

 ### W6A 完成状态

+---
+
+## 25. 2026-06-06 bs>1 打开后性能回退排查记录
+
+用户确认：`batch_size == 1` 也生成并消费 `CPSharedKVBatchPlan` 是预期收敛方向，
+不是本轮性能问题本身。后续不要再把“bs=1 走 batch-plan”当成需要回退的 bug。
+
+最新远端日志（`/mnt/beegfs/cjy/log/sglang_cp_hicache_20260605_164245.log`）
+显示：
+
+- 没有 `[HiCache-load] slow load_back planning` 或容量等待类日志；
+- fallback 主要是 `[CP_SHARED_KV_FALLBACK][tai_ipc_materialize] reason=paged_start_slot_nonzero`，
+  每 rank 限频后共 64 条，说明 index/page-buffer materialize 的非 0 起始 slot 会退到
+  local materialize + collective；
+- CP0 batch 分布仍以 bs=1 为主，少量 bs=2/3，因此不能简单用“真实 batch 数变大”
+  解释全部慢点；
+- debug 日志本身会污染性能测试，但不是一个足够可靠的根因解释。
+
+当前高优先级性能嫌疑：
+
+1. **bs>1 partial-current sync compose 的 row-major prefix slot span 可能过度覆盖。**
+   `build_batch_prefix_slot_span()` 为了用一个 contiguous span 覆盖所有 request prefix，
+   在 row-major page table 下会包含 request 之间的 current/空洞 slot。cache-hit 高、
+   prefix 分布不均时，prefix materialize/all-reduce 的页数可能远大于真实 prefix 页数。
+   这会落在每层 critical path 上，尤其在 bs>1 async prefetch 还没恢复时更明显。
+
+2. **index/top-k batch path 仍有 per-layer Python/Torch descriptor overhead。**
+   `_get_topk_in_seq_cp_pair_batch()` 每层构造 Python list、`torch.tensor(..., device)`、
+   `torch.cat`、`torch.full`，再 scatter compact topk。TAI batch kernel 已减少部分
+   K/S copy，但 descriptor 构造还不是 per-forward 预计算。
+
+3. **tiny extend compute padding 会把 `< cp_size` pages 的 request 补到 `cp_size` pages。**
+   例如 page=64、cp=8 时，200 token 会以 512 token compute 形态参与 CP split。
+   这对 200-512 token 的线上短 extend 可能抵消 batch 的 compute 填充收益。
+
+4. **bs>1 L1 prefix prefetch 仍未实现。**
+   `CpSharedKVMlaPrefetcher` / `CpSharedKVIndexPrefetcher` 仍要求
+   `forward_batch.batch_size == 1`。如果开启真实 bs>1 后 cache-hit 高，prefix 准备
+   会更多落到同步 compose 路径，而不是目标的 async L2->L1/L1 materialize 流水线。
+
+下一步验证顺序：
+
+1. 加或复用低频 perf counter，记录每个 forward 的 prefix true pages、prefix span pages、
+   current span pages、batch size、extend/prefix lens，不在每层刷屏；
+2. microbench `_get_topk_in_seq_cp_pair_batch()` descriptor/cat/scatter 与
+   `split_tensor_by_cp_batch_plan()` 在 200/512/1k/2k extend、bs=1/2/4/5 下的 CPU submit
+   和 GPU elapsed；
+3. 若证实 prefix span 过度覆盖，优先改为 compact prefix descriptor 或支持 nonzero-slot
+   TAI IPC materialize，避免 row-major gap 被 all-reduce；
+4. 若证实 descriptor overhead 主导，先把 index batch descriptor 缓存在 `ForwardBatch`
+   上，避免 78 层重复构造。
+
 已补 characterization tests，锁住现有 L2->L1 load-back batch 行为：

 - `test_cp_start_loading_batches_multiple_load_cp_requests_with_draft`
@@ -1616,3 +1668,28 @@ PYTHONPATH=python python -m pytest -q \
  `batch_gt1_index_q_length_mismatch`；
 - 日志中的 `[CP_SHARED_KV_FALLBACK][tai_index_mqa_prepare] current_index_k must be uint8`
  是另一个性能 fast-path dtype 问题，本次未修；它当前是 warning fallback，不是本次进程退出原因。
+
+## 26. 2026-06-06 bs>1 partial-current prefix span 修正
+
+用户确认：`bs=1` 走 batch-plan 路径是预期行为，不应作为回退或禁用目标。
+
+本轮性能排查确认一个明确问题：bs>1 的 partial-current 同步 compose 之前使用
+`build_batch_prefix_slot_span()` 把 row-major page table 中多个 request 的 prefix 区间压成一个
+bounding span。这个 span 会覆盖 request 行之间的 current/gap slot，导致 index/MLA prefix
+materialize 后对并非 prefix 的页一起做 all-reduce。在线上 cache-hit-heavy、prefix 长度不一致的
+batch 中，这会把本应只覆盖 prefix pages 的通信量放大。
+
+修正：
+- 保留旧的 `build_batch_prefix_slot_span()` 作为单 bounding span helper。
+- 新增 `build_batch_prefix_slot_spans()`，返回每个 request prefix 的精确 slot span，并只 merge 相邻区间。
+- `materialize_prefix_and_reuse_current_kv_page_slots()` /
+  `materialize_prefix_and_reuse_current_index_page_slots()` 支持 `prefix_slot_spans`。
+- MLA 与 index 的 bs>1 partial-current sync compose 使用精确 prefix spans；bs=1 或单 span 仍保留 IPC fast path。
+
+权衡：
+- 可能把一次 prefix all-reduce 变成至多 batch_size 次 range all-reduce；但这不超过逐 request 执行的 collective 次数，且避免 row-major gap 过通信。
+- 这只修正同步 compose 的过通信；bs>1 仍有 per-layer Python descriptor 开销、bs>1 L1 async prefetch 未支持等潜在性能项，后续需要用 ETE/benchmark 继续确认。
+
+验证：
+- 远端 `g0034` container `/sgl-workspace/sglang-tai` py_compile 通过。
+- 远端 `test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py`：113 passed, 5 warnings, 2 subtests passed。