Keep CP shared-KV fast paths off dense current collectives
CP shared-KV cache-hit batches should compose long prefix pages and short current pages through page-slot IPC instead of falling back to dense all_reduce. Wire the runtime and prefetch consume paths to the TAI current-staging helpers, fail fast when the configured CUDA fast path cannot run, and document the bs>1 cache-hit benchmark evidence. Constraint: bs>1 prefill must preserve the page-slot contract across fp8/bf16 and zero-lane current tails. Rejected: Silent all_reduce fallback | hides correctness and performance regressions in production. Confidence: medium Scope-risk: moderate Directive: Any future fallback in CP shared-KV CUDA fast paths must be explicit warning/fail-fast and covered by runtime tests. Tested: Local py_compile cp_shared_kv_runtime.py and cp_shared_kv_prefetch.py; remote PYTHONPATH=python pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py (144 passed, 21 warnings, 2 subtests passed); remote TAI IPC benchmark fp8 bs>1 cache-hit matrix recorded in docs. Not-tested: Full ETE mixed replay after replacing all current collectives with IPC. (cherry picked from commit 8aa3b4ce59e0ebef5da5b0d07499a5f1d9785997)
This commit is contained in:
laoyao0822
+131
-39
@@ -442,6 +442,23 @@ L2->L1 load finished on owner rank
|
||||
|
||||
### P1. 补 bs>1 prefetch plan 单测
|
||||
|
||||
**状态（2026-06-12）：已完成 MLA baseline。**
|
||||
|
||||
已新增：
|
||||
|
||||
- `test_mla_prefetch_create_batch_uses_exact_prefix_and_current_spans`
|
||||
- `test_mla_prefetch_batch_consume_reduces_exact_current_spans`
|
||||
|
||||
RED 证据：旧代码在 `batch_size=2` 时 `maybe_create()` 直接返回 `None`，且
|
||||
`CpSharedKVMlaPrefetcher.__init__()` 不接受 `prefix_slot_spans/current_slot_spans`。
|
||||
|
||||
GREEN 证据：远端 `cjy-glm5-new` 容器内
|
||||
`test_cp_shared_kv_runtime.py` 全文件通过：
|
||||
|
||||
```text
|
||||
132 passed, 21 warnings, 2 subtests passed
|
||||
```
|
||||
|
||||
**文件：**
|
||||
|
||||
- 修改：`test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py`
|
||||
@@ -479,6 +496,31 @@ PYTHONPATH=python python -m pytest -q \
|
||||
|
||||
### P2. 改造 MLA prefetcher 支持 bs>1 spans
|
||||
|
||||
**状态（2026-06-12）：已完成第一版 spans baseline。**
|
||||
|
||||
当前实现：
|
||||
|
||||
1. `CpSharedKVMlaPrefetcher.maybe_create()` 不再以 `batch_size != 1` 为
|
||||
skip 条件。
|
||||
2. create 阶段基于 `metadata.real_page_table`、`extend_prefix_lens_cpu`、
|
||||
`extend_seq_lens_cpu` 构造：
|
||||
- `prefix_slot_spans`
|
||||
- `current_slot_spans`
|
||||
- `prefix_page_count`
|
||||
- `current_page_count`
|
||||
3. `start_next_layer_prefix()` 只 materialize/reduce `prefix_slot_spans`，
|
||||
不再把 batch flattened page table 当成 `[0:prefix_pages)`。
|
||||
4. `consume_prefix_with_current()` 只 reduce `current_slot_spans`，避免把
|
||||
row gap / 其他 request prefix 一起 reduce。
|
||||
5. `consume()` 的 legacy full-materialize suffix 路径也改为使用
|
||||
`current_slot_spans`，避免 bs>1 bounding suffix。
|
||||
|
||||
当前限制：
|
||||
|
||||
- MLA prefix spans 仍走现有 materialize + async all-reduce baseline；还没有
|
||||
接入 spans-list TAI IPC 或 0SM CE。
|
||||
- index prefetcher 仍未改造，继续由 P3 处理。
|
||||
|
||||
**文件：**
|
||||
|
||||
- 修改：`python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py`
|
||||
@@ -490,7 +532,8 @@ PYTHONPATH=python python -m pytest -q \
|
||||
2. create 阶段构造 `prefix_slot_spans/current_slot_spans`。
|
||||
3. `start_next_layer_prefix()` 只 materialize prefix spans。
|
||||
4. `consume_prefix_with_current()` 只 fill/reduce current spans。
|
||||
5. `consume()` 如果仍存在 legacy suffix path，bs>1 下 fail-fast，避免错误 bounding suffix。
|
||||
5. `consume()` 如果仍存在 legacy suffix path，必须使用 `current_slot_spans`
|
||||
或 fail-fast，不能回到错误 bounding suffix。
|
||||
|
||||
**第一版允许：**
|
||||
|
||||
@@ -503,6 +546,46 @@ PYTHONPATH=python python -m pytest -q \
|
||||
|
||||
### P3. 改造 index prefetcher 支持 bs>1 spans
|
||||
|
||||
**状态（2026-06-12）：已完成第一版 spans baseline。**
|
||||
|
||||
已新增：
|
||||
|
||||
- `test_index_prefetch_create_batch_uses_exact_prefix_and_current_spans`
|
||||
- `test_index_prefetch_batch_consume_reduces_exact_current_spans`
|
||||
|
||||
RED 证据：旧代码在 `batch_size=2` 时以
|
||||
`[CP_SHARED_KV_FALLBACK][index_prefetch] reason=batch_size` 返回 `None`，
|
||||
且 `CpSharedKVIndexPrefetcher.__init__()` 不接受
|
||||
`prefix_slot_spans/current_slot_spans`。
|
||||
|
||||
当前实现：
|
||||
|
||||
1. `CpSharedKVIndexPrefetcher.maybe_create()` 不再以 `batch_size != 1`
|
||||
为 skip 条件。
|
||||
2. create 阶段复用 MLA 同一套 deterministic spans：
|
||||
- `prefix_slot_spans`
|
||||
- `current_slot_spans`
|
||||
- `prefix_page_count`
|
||||
- `current_page_count`
|
||||
3. `start_next_layer_prefix()` 只 materialize/reduce index prefix spans。
|
||||
4. `consume_prefix_with_current()` 只 fill/reduce index current spans。
|
||||
5. `consume()` 的 legacy suffix 路径也改为 `current_slot_spans`，
|
||||
不再使用 batch bounding suffix。
|
||||
|
||||
远端验证：
|
||||
|
||||
```text
|
||||
test_cp_shared_kv_runtime.py
|
||||
134 passed, 21 warnings, 2 subtests passed
|
||||
```
|
||||
|
||||
当前限制：
|
||||
|
||||
- index prefetch 仍使用现有 materialize + async all-reduce baseline；
|
||||
spans-list TAI IPC / 0SM CE 留给后续 P5/P6。
|
||||
- active index layer / index skip 的 runtime hook 当前沿用已有
|
||||
`nsa_backend.py` 调用点；本阶段没有修改 skip 参数语义。
|
||||
|
||||
**文件：**
|
||||
|
||||
- 修改：`python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py`
|
||||
@@ -547,6 +630,19 @@ PYTHONPATH=python python -m pytest -q \
|
||||
- forward stream 上同步等待 L2->L1。
|
||||
- 用新的 all-reduce 确认所有 rank load 完成。
|
||||
|
||||
**2026-06-12 P4 实施记录：**
|
||||
|
||||
- 当前 HiCache load 已通过 `LayerDoneCounter` 暴露 per-layer readiness：
|
||||
- `LayerDoneCounter.wait_until_on_stream(layer_id - start_layer, stream)`
|
||||
- `MLATokenToKVPool.get_key_buffer_for_prefetch(layer_id, stream)`
|
||||
- `MLATokenToKVPool.get_index_k_with_scale_buffer_for_prefetch(layer_id, stream)`
|
||||
- 发现并修正一个 ordering 疏漏：P1-P3 后 prefix materialize 仍在 current stream 上执行，但 `start_next_layer_prefix()` 把 L2->L1 readiness wait 绑定到了 prefetch stream。这样只能保护后续 reduce，不能保护实际读取 L1 raw pages 的 materialize。
|
||||
- 修正策略：MLA/index 的 `start_next_layer_prefix()` 先取得 `current_stream`，把它传给 prefetch-safe getter，使 L2->L1 ready event 挂到实际读取 KV/index page 的 stream；随后仍保持 `prefetch_stream.wait_stream(current_stream)`，reduce 在 prefetch stream 上异步提交。
|
||||
- 没有新增 collective；没有把 forward stream 改成 CPU 同步等待。
|
||||
- 单测锁住：
|
||||
- `test_mla_prefetch_waits_l2_l1_on_materialize_stream_and_reduces_on_prefetch_stream`
|
||||
- `test_index_prefetch_waits_l2_l1_on_materialize_stream_and_reduces_on_prefetch_stream`
|
||||
|
||||
### P5. TAI SM IPC spans baseline
|
||||
|
||||
**文件：**
|
||||
@@ -735,52 +831,48 @@ P7 GSM8K/replay/Nsight 验证
|
||||
|
||||
不要先做 P6 再修 P2/P3。0SM kernel 只解决 transport 代价，不解决 bs>1 prefix/current slot 语义；如果语义仍是 scalar prefix，kernel 越快只会越快地产生错误。
|
||||
|
||||
## 2026-06-12 追加：symm compose 与旧 IPC kernel 的对比口径
|
||||
---
|
||||
|
||||
`symm-syh` 分支的 compose/symm kernel 可能比旧 CUDA IPC current-staging kernel 单次更快，但不能只比较 kernel elapsed time。需要把同步点作为一等指标，否则可能出现“单 kernel 快，但每层/每 buffer 多一次同步，ETE 更慢”的误判。
|
||||
## 9. P5 当前实现状态：bs>1 prefix/suffix IPC spans
|
||||
|
||||
对比 benchmark / Nsight trace 必须至少拆分以下事件：
|
||||
已补齐一个 TAI SM IPC baseline kernel，用于先替代 bs>1 prefix/suffix 上的
|
||||
“local materialize + all_reduce”：
|
||||
|
||||
1. **IPC capability agreement**
|
||||
- `_agreed_tai_ipc_peer_ptrs()` 里的 group agreement / all-reduce。
|
||||
- 应确认是每个 pool tensor 一次，还是每个 forward/layer 重复触发。
|
||||
```text
|
||||
materialize_cuda_ipc_peer_pages_slot_indices(
|
||||
peer_ptrs,
|
||||
dst,
|
||||
owner_ranks,
|
||||
src_page_indices,
|
||||
dst_page_indices,
|
||||
page_nbytes,
|
||||
)
|
||||
```
|
||||
|
||||
2. **symm barrier**
|
||||
- `cp_symm_barrier()` 调用次数、耗时、等待方差。
|
||||
- 需要按 token KV / index buffer 分开统计。
|
||||
- 当前 tai-kernel 实现是 1 个 block / 1 个 warp 的 CUDA spin barrier，
|
||||
不是 copy-engine/0SM 路径；它占用很少 SM，但会在当前 stream 上形成
|
||||
明确同步点。判断 symm 是否优于旧 IPC 时，必须把这个 barrier 的次数和
|
||||
rank 间等待方差算进去。
|
||||
语义：
|
||||
|
||||
3. **symm mega gather**
|
||||
- `materialize_cuda_ipc_peer_pages_slot_dense()` 在 symm combined ptr table 上的耗时。
|
||||
- 统计 prefix pages、current pages、request 数、dense pages。
|
||||
```text
|
||||
dst[dst_page_indices[i]] = peer_ptrs[owner_ranks[i]][src_page_indices[i]]
|
||||
```
|
||||
|
||||
4. **compact current reduce**
|
||||
- symm 未开启或不可用时的 `_reduce_current_pages_compact()`。
|
||||
- 这是 compact current collective，不是 dense full fallback；但仍会同步/占用通信资源，需要单独计数。
|
||||
与旧 `slot_dense` kernel 的区别：
|
||||
|
||||
5. **dense full fallback reduce**
|
||||
- `_all_reduce_materialized_buffer(... v2_full ...)`。
|
||||
- 生产 CUDA + TAI materialize 开启时不允许静默发生；应 fail-fast 暴露 `CP_SHARED_KV_FAIL_FAST][compose_v2]`。
|
||||
1. 旧 kernel 只能写 `slot i -> dense page i+1`，等价于只支持从 slot 0 开始的一段连续 prefix。
|
||||
2. 新 kernel 显式传入 `dst_page_indices`，支持 bs>1 的多个 prefix spans 和 suffix spans。
|
||||
3. invalid owner/src 会 zero-fill 对应 dst page；未被 descriptor 指向的 dst page 保持原值。
|
||||
4. 该 kernel 仍是 SM copy baseline，不是最终 0SM copy-engine queue。
|
||||
|
||||
6. **CPU descriptor / plan 成本**
|
||||
- `get_or_build_compose_plan()` cache hit/miss。
|
||||
- per-forward 是否复用 descriptor；不要把一次性 build 成本误算到每层 steady state。
|
||||
SGLang 接入位置：
|
||||
|
||||
建议 benchmark 矩阵：
|
||||
1. `cp_shared_kv_runtime.py`
|
||||
- 新增 token/index span IPC helper。
|
||||
- partial-current prefix 支持多 span IPC，prefix 不再因为 bs>1 退回 all_reduce。
|
||||
2. `cp_shared_kv_prefetch.py`
|
||||
- MLA/index prefix prefetch 先尝试 span IPC；成功时在 current stream record event，不再 enqueue prefix all_reduce。
|
||||
- full-cache-hit suffix consume 先尝试 span IPC；成功时跳过 suffix all_reduce。
|
||||
|
||||
- dtype：bf16 / fp8_e4m3；
|
||||
- batch size：1, 2, 5, 10；
|
||||
- extend：1k, 2k, 10k, 40k, 65k；
|
||||
- cached/prefix：100k, 200k, 300k；
|
||||
- case：cache-hit partial-current、current-only、multi-request shared prefix；
|
||||
- mode：legacy dense fallback（只作为基线，不允许生产静默）、old CUDA IPC current staging、symm staging、symm prefetch。
|
||||
尚未覆盖：
|
||||
|
||||
结论标准：保留默认路径必须同时满足：
|
||||
|
||||
- 无 dense full fallback；
|
||||
- 同步点数量不高于旧 IPC 路径，或同步耗时被更少 kernel/更高带宽抵消；
|
||||
- ETE replay 在 cache-hit-heavy 短 extend 场景提升，而不是只在 micro benchmark 提升；
|
||||
- GSM8K cache-hit 二轮精度不掉点。
|
||||
1. partial-current 的 current rows 仍是本 rank 当前 forward 产生的临时 buffer，source layout 不是长期 L1 page buffer；不能直接复用 peer page IPC。
|
||||
2. 该 current rows all_reduce 需要单独设计 owner-aware current-source IPC/fused compose kernel。
|
||||
3. 0SM CE path 仍未实现，本阶段只是先消除 bs>1 prefix/suffix 上不必要的 collective。
|
||||
|
||||
+271
@@ -0,0 +1,271 @@
|
||||
# NSA Prefill CP shared-KV：用自研 IPC collective 替换 materialize all_reduce
|
||||
|
||||
## 目标
|
||||
|
||||
彻底移除 CP shared-KV materialize 热路径上的 NCCL/Gloo `all_reduce`：
|
||||
|
||||
1. prefix / suffix / full-cache-hit：继续使用 L1 page buffer 上的 IPC page gather，失败必须显式 warning / fail-fast，不能静默回退。
|
||||
2. partial-current / current reuse：新增 current staging + ready flag IPC collective，替换当前 dense page fill 后的 slot-range `all_reduce`。
|
||||
3. bs>1：descriptor 必须一次覆盖 batch 内多个 request 的 slot spans，不允许 per-request 循环发射 kernel。
|
||||
4. fp8 / bf16：token KV 与 index page buffer 两条路径都要支持。
|
||||
5. benchmark：必须覆盖 all_reduce baseline、现有 IPC prefix/suffix、current staging IPC，并输出 CPU submit、GPU elapsed、有效带宽、kernel launch 数。
|
||||
|
||||
## 当前 all_reduce 分类
|
||||
|
||||
### 已可用 IPC 替换的路径
|
||||
|
||||
- `materialize_shared_token_kv_buffer` / `materialize_shared_paged_buffer`：full materialize fallback all_reduce。
|
||||
- `materialize_prefix_and_reuse_current_*` 的 prefix spans：已有 `_try_tai_ipc_materialize_*_page_slot_spans_into`。
|
||||
- `CpSharedKV*Mla/IndexPrefetcher.consume()` 的 suffix spans：已有 IPC span gather。
|
||||
- `start_next_layer_prefix()` 的 prefix prefetch：已有 IPC span gather。
|
||||
|
||||
这些路径的源数据是长期存在的 L1 `kv_cache` / `page_buffer`，IPC handle 可以按 storage 缓存，只需在分配/扩容后重新 open。
|
||||
|
||||
### 仍依赖 all_reduce 的路径
|
||||
|
||||
- `materialize_prefix_and_reuse_current_kv_page_slots()`：`fill_current_kv_page_slots_and_remap_locs()` 后，对 current slot spans 做 `_all_reduce_materialized_buffer_range()`。
|
||||
- `materialize_prefix_and_reuse_current_index_page_slots()`：`fill_current_index_page_slots()` 后，对 current page spans 做 `_all_reduce_materialized_buffer_range()`。
|
||||
- `CpSharedKVMlaPrefetcher.consume_prefix_with_current()` / `CpSharedKVIndexPrefetcher.consume_prefix_with_current()`：prefetched prefix + current fill 后仍 reduce current slot spans。
|
||||
|
||||
这些路径的源是每层 forward 产生的临时 `current_kv_cache` / `current_index_k` / `current_index_scale`。不能直接对临时 tensor 做 per-layer IPC handle all_gather，否则只是把 all_reduce 换成另一个高频 collective。
|
||||
|
||||
## 设计选择
|
||||
|
||||
### 方案 A：临时 tensor IPC handle all_gather（拒绝）
|
||||
|
||||
每层对 `current_*` tensor open IPC handle，然后 peer-read current rows。
|
||||
|
||||
拒绝原因：
|
||||
- data_ptr/shape 每层/每 batch 变化，handle cache 命中率低。
|
||||
- 仍需要高频 `all_gather` 交换 handle/offset。
|
||||
- CPU submit 和同步开销不可控，违背“彻底干掉 collective”的目标。
|
||||
|
||||
### 方案 B：persistent current staging + ready flag（采用）
|
||||
|
||||
每个 CP rank 维护长期 CUDA staging buffer 和 ready counter buffer：
|
||||
|
||||
1. current fill kernel 同时把本 rank owner-lane current pages 写入本 rank staging buffer，布局与 dense slot page 对齐。
|
||||
2. publish 完成后在同 stream 写 ready seq：`__threadfence_system()` 后 store seq。
|
||||
3. 所有 rank 用 IPC peer ptrs 读取 owner rank 的 staging pages，gather 到本地 dense buffer；gather kernel 在读取每个 owner 前等待 `peer_ready[owner] >= seq`。
|
||||
4. descriptor 以 dense slot page 为单位，跨 batch request 合并为一个 owner/page/slot list，一次 kernel launch 完成多个 request。
|
||||
|
||||
采用原因：
|
||||
- IPC handle 只在 staging buffer 分配/扩容时交换，热路径无 NCCL/Gloo collective。
|
||||
- current 数据仍按 page 最小单位发布，符合当前 page-aligned cache 合同。
|
||||
- bs>1 可以复用 slot span descriptor，一次 launch 覆盖多个 request。
|
||||
- 可与现有 TAI current fill kernel 融合，避免重复 remap/row-mask 逻辑。
|
||||
|
||||
## staging buffer 合同
|
||||
|
||||
### token KV staging
|
||||
|
||||
- 形状语义：flat bytes，容量至少覆盖 `dense_num_pages * page_size * kv_row_bytes`。
|
||||
- 写入地址：`dense_slot_page * page_nbytes + row_offset * row_nbytes`。
|
||||
- 每次 publish 只保证 valid current rows 正确；为了避免 stale tail，publish kernel 需要对 touched current slot pages 的 tail slack 清零，或 SGLang 必须保证返回 locs 不引用 tail slack。第一版建议在 publish kernel 内按 touched page 清零，优先正确性。
|
||||
|
||||
### index staging
|
||||
|
||||
- 形状语义：flat bytes，容量至少覆盖 `dense_num_pages * index_page_bytes`。
|
||||
- 写入地址：`dense_slot_page * page_bytes`，内部包含 K rows + scale rows。
|
||||
- current K/scale 的 valid rows 写入对应 row offset；tail slack 清零。
|
||||
|
||||
### ready flag
|
||||
|
||||
- 每 rank 一个 `uint64/int64` counter buffer，通过 IPC peer ptrs 打开。
|
||||
- 每次 current publish 使用递增 seq。
|
||||
- gather kernel 对需要读取的 owner 执行 device-side wait，直到 `peer_ready[owner] >= seq`。
|
||||
- wait kernel 需要 watchdog/iteration bound，debug build 可 fail-fast；生产第一版可以保留有限 spin + error flag，避免死锁静默挂住。
|
||||
|
||||
## descriptor 合同
|
||||
|
||||
输入：`current_slot_spans` / `slot_logical_pages` / layout / physical capacity。
|
||||
|
||||
输出：
|
||||
- `owner_ranks[num_pages]`
|
||||
- `src_page_indices[num_pages]`：对 staging 来说等于 dense slot page id；对 persistent L1 prefix/suffix 来说是 physical page index。
|
||||
- `dst_slot_indices[num_pages]`：dense buffer 1-based slot index，保持现有 kernel 合同。
|
||||
|
||||
bs>1 要求：
|
||||
- span list 可以覆盖多个 request。
|
||||
- descriptor 构造只按 merged slot spans 生成一次，不允许 request loop + 多次 kernel。
|
||||
- 如果某些 request 没 current page，descriptor 为空时直接成功。
|
||||
|
||||
## TAI kernel 阶段
|
||||
|
||||
### P1：RED tests / benchmark skeleton
|
||||
|
||||
- SGLang unit：partial-current compose 在 cp_size>1 时应调用 IPC current gather helper，不应调用 `_all_reduce_materialized_buffer_range`。
|
||||
- tai-kernel CUDA test：声明期望 API `publish_current_*_to_staging` 和 `materialize_cuda_ipc_peer_pages_slot_indices_wait_ready`，先验证缺失失败。
|
||||
- benchmark skeleton：同一输入比较 `local fill + all_reduce` 与 `publish + IPC gather`。
|
||||
|
||||
### P2：current publish kernel
|
||||
|
||||
- token：扩展/新增 TAI op，复用 `cp_fill_current_kv_page_slots` remap 逻辑，同时写 staging 和 ready seq。
|
||||
- index：扩展/新增 TAI op，复用 `cp_fill_current_index_page_slots`，同时写 staging 和 ready seq。
|
||||
|
||||
### P3：ready-wait IPC gather kernel
|
||||
|
||||
- 基于现有 `materialize_cuda_ipc_peer_pages_slot_indices` 增加 ready peer ptrs + seq 参数。
|
||||
- 支持 page_nbytes 变长，owner/dst descriptor 长度变长。
|
||||
- 支持 self-rank peer ptr 快路径，便于单机 CUDA unit test。
|
||||
|
||||
### P4：SGLang runtime 接入
|
||||
|
||||
- 新增 runtime helper：`_try_tai_ipc_materialize_current_token_kv_page_slot_spans_into`。
|
||||
- 新增 runtime helper：`_try_tai_ipc_materialize_current_paged_buffer_page_slot_spans_into`。
|
||||
- current compose：先 fill+publish，再 IPC gather current spans；失败时 warning/fail-fast，不再静默 all_reduce。
|
||||
- prefetch consume_prefix_with_current：同样走 current staging IPC。
|
||||
|
||||
### P5：严格化 fallback
|
||||
|
||||
- prefix/suffix/full materialize：IPC 失败在生产 CP shared-KV fast path 下 fail-fast；仅在显式 debug env 下允许 fallback，用 warning 标记。
|
||||
- all_reduce helper 保留给非 CP shared-KV 或测试 reference，不在 fast path 默认触发。
|
||||
|
||||
### P6：验证
|
||||
|
||||
- CPU unit：descriptor / fallback contract / no-allreduce path。
|
||||
- CUDA unit：self-rank staging roundtrip、multi-process cp=2/8 roundtrip、fp8/bf16、bs=1/5/10。
|
||||
- benchmark：4k/16k/40k/80k/160k prefix+current，bs=1/5/10，报告 CPU/GPU 时间和有效带宽。
|
||||
- ETE：GSM8K 两轮 cache-hit 精度不掉点；mixed replay 不出现 detokenizer hang / all_reduce collective mismatch。
|
||||
|
||||
## 风险与约束
|
||||
|
||||
- device-side ready wait 如果某 rank 没有 publish 会死等；必须确保所有 rank 都按同一 seq 进入 publish/gather，即使本 rank current rows 为空也要 publish ready。
|
||||
- staging buffer 扩容会触发一次 IPC handle exchange；必须高水位缓存，不能每层分配。
|
||||
- current staging tail slack 不能污染 attention/index；第一版应清零 touched current pages，后续再优化成 valid-locs 完全约束。
|
||||
- 一次 kernel 同时 publish+peer-gather 在跨进程场景没有全局同步，容易死锁；第一版采用 publish kernel + gather kernel 两阶段。
|
||||
|
||||
## 当前结论
|
||||
|
||||
先实现 staging+ready 的 SM IPC collective，彻底移除 current compose all_reduce。0SM/copy-engine 版本后续单独做；当前优先解决 correctness、CPU collective overhead、bs>1 一次 launch。
|
||||
|
||||
## 2026-06-11 实现与 benchmark 更新
|
||||
|
||||
### 已完成
|
||||
|
||||
- TAI 已新增 `publish_cuda_ipc_slot_pages_and_mark_ready`：把本 rank 已经填好的 dense slot pages 复制到 persistent staging，并写 ready seq。
|
||||
- TAI 已新增 `materialize_cuda_ipc_peer_pages_slot_indices_wait_ready`：按 owner/src/dst descriptor 等待 peer ready 后从 IPC peer staging 复制到本地 dense buffer。
|
||||
- SGLang runtime 的以下 current compose 已接入 current-staging IPC helper：
|
||||
- token KV：`materialize_prefix_and_reuse_current_kv_page_slots()`。
|
||||
- index page buffer：`materialize_prefix_and_reuse_current_index_page_slots()`。
|
||||
- MLA/index prefetch `consume_prefix_with_current()`。
|
||||
- bs>1 descriptor 合同已覆盖：一次 launch 可以用 flattened slot-page descriptor 覆盖多个 request spans；benchmark 用 `--current-batch-requests 5` 验证该形态。
|
||||
|
||||
### 与原计划的差异
|
||||
|
||||
P2 暂时没有把 publish 融入 `cp_fill_current_*` row-fill kernel，而是采用：
|
||||
|
||||
1. current fill 先写 dense buffer；
|
||||
2. `publish_cuda_ipc_slot_pages_and_mark_ready` 再按 page 复制本 rank owner pages 到 staging；
|
||||
3. wait-ready IPC gather 从 peer staging 拉取所有 current pages。
|
||||
|
||||
这样 correctness 风险低，接入面小，但 current path 多一次 page copy + 一个额外 kernel。benchmark 也证明小 current span 下该版本不一定优于 NCCL/Gloo all_reduce；后续要进一步追性能，需要做“current fill 同时写 staging + mark ready”的融合版本。
|
||||
|
||||
### 远端 benchmark 证据
|
||||
|
||||
环境：`g0034` / `cjy-glm5-new` / `torchrun --nproc_per_node=8` / `benchmark/nsa_prefill/benchmark_cp_shared_kv_ipc_gather.py`。
|
||||
|
||||
命令示例：
|
||||
|
||||
```bash
|
||||
PYTHONPATH=python torchrun --standalone --nproc_per_node=8 \
|
||||
benchmark/nsa_prefill/benchmark_cp_shared_kv_ipc_gather.py \
|
||||
--tokens 16384 32768 65536 98304 122880 \
|
||||
--dtype uint8 --kv-dim 656 \
|
||||
--warmup 3 --repeat 8 \
|
||||
--include-current-staging --current-batch-requests 5 --no-check
|
||||
```
|
||||
|
||||
FP8/uint8 MLA page（`kv_dim=656`）current-staging IPC vs current all_reduce：
|
||||
|
||||
| tokens | pages | all_reduce p50 | current IPC p50 | 结论 |
|
||||
| --- | ---: | ---: | ---: | --- |
|
||||
| 16k | 256 | 0.184 ms | 0.199 ms | 当前两阶段 IPC 小幅变慢 |
|
||||
| 32k | 512 | 0.227 ms | 0.217 ms | 基本持平/略快 |
|
||||
| 65k | 1024 | 0.337 ms | 0.253 ms | IPC 快约 25% |
|
||||
| 98k | 1536 | 0.450 ms | 0.307 ms | IPC 快约 32% |
|
||||
| 122k | 1920 | 0.527 ms | 0.362 ms | IPC 快约 31% |
|
||||
|
||||
BF16 MLA page（`kv_dim=576`）current-staging IPC vs current all_reduce：
|
||||
|
||||
| tokens | pages | all_reduce p50 | current IPC p50 | 结论 |
|
||||
| --- | ---: | ---: | ---: | --- |
|
||||
| 16k | 256 | 0.175 ms | 0.175 ms | 持平 |
|
||||
| 65k | 1024 | 0.397 ms | 0.414 ms | 两阶段 IPC 小幅变慢 |
|
||||
| 122k | 1920 | 0.671 ms | 0.594 ms | IPC 快约 11% |
|
||||
|
||||
同一 benchmark 中 prefix/L1 persistent IPC 的 `cuda_ipc_peer_pages_materialize_slot_dense` 仍明显优于 dense all_reduce，例如 FP8：65k tokens 从 0.334 ms 降到 0.164 ms，122k tokens 从 0.517 ms 降到 0.274 ms。这说明 persistent L1 prefix/suffix 场景适合直接 IPC；current 临时数据场景的瓶颈在 publish 阶段，下一步应优先融合 current fill + staging publish。
|
||||
|
||||
### 下一步必须处理的点
|
||||
|
||||
1. fast path fallback 已收窄：当 tensor 在 CUDA 上且 `SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE=1` 时，prefix/current IPC 失败会 fail-fast，不再静默走 all_reduce；CPU reference/TAI 显式关闭路径仍保留测试 fallback。
|
||||
2. wait-ready kernel 现在有 `max_spins`，但没有 device error flag；超过 spin 后可能复制 stale data。生产化前应增加 error flag 或明确 fail-fast 检测，避免 silent corruption。
|
||||
3. current publish 未融合，短 current span 可能回退性能。应新增 TAI fused current fill+publish op：fill dense 的同时写 staging，随后只用一个 tiny mark-ready kernel，再 wait-gather。
|
||||
|
||||
## 2026-06-11 bs>1 cache-hit compose benchmark 补充
|
||||
|
||||
### 为什么补这个 benchmark
|
||||
|
||||
之前的 IPC benchmark 只按“总 pages”测 current staging 或 prefix gather，不能覆盖真实线上 cache-hit 形态：
|
||||
|
||||
- 每个 request 有很长 cached prefix（100k-300k tokens）。
|
||||
- 每个 request 的 extend/current 较短（10k-65k tokens）。
|
||||
- prefill bs>1 时一次 batch 内有 2-10 个 request。
|
||||
- 部分短 current tail 只落在少数 CP owner lanes，上游 zero-lane rank 仍必须 publish ready，不能退出 collective 合同。
|
||||
|
||||
因此新增 `--include-cache-hit-compose` / `--cache-hit-only`，显式构造 request-major dense slot layout：
|
||||
|
||||
```text
|
||||
page0(dummy) |
|
||||
req0 cached pages | req0 current pages |
|
||||
req1 cached pages | req1 current pages | ...
|
||||
```
|
||||
|
||||
对比两条路径：
|
||||
|
||||
1. `cache_hit_dense_all_reduce_full`：每 rank 构造本 owner pages 的 dense buffer，然后对整个 dense slot buffer 做 all_reduce。
|
||||
2. `cache_hit_ipc_prefix_current_compose`：cached prefix 从 persistent compact owner staging 走 IPC slot-index materialize；current/extend 从 dense current staging publish + ready-wait IPC gather。
|
||||
|
||||
注意：current owner 分布按完整 request page positions 计算，所以短 extend 可能只触达少数 owner lanes。这是预期现象，不再要求每个 rank 都拥有 current page；zero-lane rank 仍会 publish ready，避免 peer wait 死锁。
|
||||
|
||||
### 远端命令
|
||||
|
||||
```bash
|
||||
PYTHONPATH=python torchrun --standalone --nproc_per_node=8 \
|
||||
benchmark/nsa_prefill/benchmark_cp_shared_kv_ipc_gather.py \
|
||||
--cache-hit-only \
|
||||
--cache-hit-cached-tokens 102400 204800 307200 \
|
||||
--cache-hit-extend-tokens 10240 32768 65536 \
|
||||
--cache-hit-batch-requests 2 5 10 \
|
||||
--dtype uint8 --kv-dim 656 \
|
||||
--warmup 2 --repeat 5 --no-check
|
||||
```
|
||||
|
||||
### FP8/uint8 结果摘要
|
||||
|
||||
环境：`g0034` / `cjy-glm5-new` / 8 ranks / `kv_dim=656` / page size 64。
|
||||
|
||||
| bs | cached/req | extend/req | dense all_reduce p50 | IPC compose p50 | 收益 |
|
||||
| ---: | ---: | ---: | ---: | ---: | ---: |
|
||||
| 2 | 100k | 10k | 0.858 ms | 0.686 ms | 1.25x |
|
||||
| 2 | 200k | 65k | 1.960 ms | 1.503 ms | 1.30x |
|
||||
| 2 | 300k | 65k | 2.669 ms | 1.999 ms | 1.34x |
|
||||
| 5 | 100k | 10k | 2.043 ms | 1.525 ms | 1.34x |
|
||||
| 5 | 200k | 65k | 4.792 ms | 3.517 ms | 1.36x |
|
||||
| 5 | 300k | 65k | 6.605 ms | 4.798 ms | 1.38x |
|
||||
| 10 | 100k | 10k | 4.036 ms | 2.951 ms | 1.37x |
|
||||
| 10 | 200k | 65k | 9.441 ms | 6.999 ms | 1.35x |
|
||||
| 10 | 300k | 65k | 13.023 ms | 9.552 ms | 1.36x |
|
||||
|
||||
完整矩阵结论：在 fp8 cache-hit bs>1 形态下，IPC compose 对 dense all_reduce 稳定约 1.25x-1.6x；cached 越长、bs 越大收益越稳定。这个 benchmark 比单独 current staging 更贴近实际，因为真实 cache-hit 主要成本来自长 cached prefix 的 materialize，而这部分 persistent IPC 收益明显。
|
||||
|
||||
### BF16 spot-check
|
||||
|
||||
命令只测代表性 case：`bs=5 cached=200k extend=10k/65k dtype=bf16 kv_dim=656`。
|
||||
|
||||
| bs | cached/req | extend/req | dense all_reduce p50 | IPC compose p50 | 结论 |
|
||||
| ---: | ---: | ---: | ---: | ---: | --- |
|
||||
| 5 | 200k | 10k | 5.929 ms | 5.595 ms | 略快 |
|
||||
| 5 | 200k | 65k | 7.387 ms | 7.366 ms | 基本持平 |
|
||||
|
||||
BF16 下收益不明显，原因是两阶段 current publish 的额外 copy 被放大；当前线上 GLM5 使用 fp8 KV cache，因此优先级仍然是把 fp8 bs>1 fast path 接稳。后续若要兼顾 BF16，需要 fused current fill+publish，减少 current staging 的额外 HBM copy。
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user