diff --git a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_l1_prefetch_zero_sm_plan_zh.md b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_l1_prefetch_zero_sm_plan_zh.md new file mode 100644 index 000000000..67cdf844a --- /dev/null +++ b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_l1_prefetch_zero_sm_plan_zh.md @@ -0,0 +1,736 @@ +# NSA Prefill CP Shared-KV bs>1 L1 Prefetch 与 0SM 传输计划 + +> 日期:2026-06-10 +> 分支:`cjy-cp-refactor` +> 当前基线:`c8c736e75` +> 范围:target model 的 bs>1 CP shared-KV L1 prefetch;MLA KV 与 NSA index 两条路径;后续 0SM/低 SM TAI kernel 设计。 + +## 0. 目标和非目标 + +### 0.1 目标 + +1. 在 `--enable-cp-shared-kv-prefill-bs-gt1` 下恢复 target model 的 L1 shared-KV prefetch。 +2. 保持两级流水线合同: + +```text +L2 -> L1 HiCache load: 提前 2 个 layer +L1 shared-KV prefetch: 提前 1 个 layer +当前 layer consume/wait: 只等自己需要的 layer-ready event +``` + +3. bs>1 下 prefetch 必须使用 request-aware slot spans,不能继续使用单个 scalar `prefix_pages` 把整个 flattened page table 切成 prefix/suffix。 +4. 第一版先支持 target model: + - MLA KV prefix L1 prefetch; + - NSA index prefix L1 prefetch; + - current suffix 仍沿用当前同步 compose/reduce 路径。 +5. 后续为 prefix prefetch 提供 0SM/低 SM transport primitive,避免 prefetch materialize 与当前 layer attention/MLA/MOE 抢 SM。 + +### 0.2 非目标 + +1. 第一版不打开 EAGLE/draft prefetch。draft 只有少量 executable layer,且 draft KV 生命周期应跟随 target mirror,不能简单套 target 的 next-layer prefetch。 +2. 第一版不改变 decode 路径。 +3. 第一版不改变 HiCache storage layout,也不强制切换到 layer-page-first。 +4. 第一版不把 current suffix 的跨 rank compose 改成 0SM。current suffix 是当前 layer 刚产生的数据,peer-read ordering 比历史 prefix 更严格,先保留已有同步语义。 +5. 不新增 all-reduce/all-gather 来同步 “prefetch plan 是否一致”。plan 必须由本地 deterministic metadata 推导。 + +--- + +## 1. 两级流水线合同 + +我们期望的 target layer 流水线如下: + +```text +warmup / batch prepare: + 为 layer 0、layer 1 至少启动 L2->L1 load + +处理 layer L-2: + 启动 layer L 的 L2->L1 load + +处理 layer L-1: + 在 prefetch stream 上等待 layer L 的 L2->L1 load event + 启动 layer L 的 L1 shared-KV prefetch + +处理 layer L: + 当前 layer index/MLA consume 时等待 layer L 的 L1 prefetch event + current suffix 仍在当前 layer 生成后同步 fill/reduce + attention 使用 prefix-prefetched + current-filled dense view +``` + +对应窗口: + +| 时间点 | 允许做的工作 | 不允许做的工作 | +| --- | --- | --- | +| `layer L-2` | 启动 `L2->L1(layer L)` | 阻塞 forward stream 等 H2D | +| `layer L-1` attention 前 | 启动 `L1 prefetch(layer L)` | 与当前 layer 必要 KV reduce 争抢 collective 顺序 | +| `layer L` consume | 等待 `L1 prefetch(layer L)` event | 重新 materialize 已经 prefetched 的 prefix | + +关键点: + +1. L2->L1 load 比 L1 prefetch 早一个 layer,避免 L1 prefetch 读到尚未从 host load back 的 device pages。 +2. L1 prefetch 的启动点应保持在当前 layer attention 前,即当前实现中的 `_maybe_start_cp_shared_kv_attention_prefetch()` 语义。 +3. L1 prefetch 不应该提前到 MQA/index materialize 前,否则会与当前 layer 的必要 materialize/reduce 重叠,造成之前观察到的性能回退。 +4. L2->L1 load event 与 L1 prefetch event 是两个不同层次的 readiness: + - L2->L1 event:device pool 中的真实 physical pages 已经可读; + - L1 prefetch event:dense prefix buffer 已经可被 consumer 使用。 + +--- + +## 2. 当前代码事实 + +### C1. L1 prefetch hook 已经在 attention 前 + +相关文件: + +- `python/sglang/srt/layers/attention/nsa_backend.py` + - `_maybe_start_cp_shared_kv_attention_prefetch(forward_batch, layer_id)` + +当前 hook 做的事情: + +```text +next_layer_id = layer_id + 1 +index_prefetcher.start_next_layer_prefix(next_layer_id) +mla_prefetcher.start_next_layer_prefix(next_layer_id) +``` + +结论:L1 prefetch 的 “提前 1 layer” 启动位置基本符合目标;bs>1 工作重点不是移动 hook,而是让 prefetcher 的 plan 和 materialize 支持 batch。 + +### C2. MLA/index prefetcher 仍强制 batch_size == 1 + +相关文件: + +- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py` + - `CpSharedKVMlaPrefetcher.maybe_create()` + - `CpSharedKVIndexPrefetcher.maybe_create()` + +当前行为: + +```text +if int(forward_batch.batch_size) != 1: + create_skip / fallback reason=batch_size + return None +``` + +结论:不能直接删除 guard。pre-batch 代码内部仍假设 prefix 是 flattened slot 的 `[0:prefix_pages)`。 + +### C3. 当前 prefetcher 使用 scalar prefix_pages + +当前 prefetcher 字段: + +```text +prefix_pages: int +slot_logical_pages: Tensor +page_inverse: Tensor +total_slots = slot_logical_pages.numel() +``` + +当前 consume 逻辑: + +```text +prefix slots: [0, prefix_pages) +suffix slots: [prefix_pages, total_slots) +``` + +bs=1 时成立;bs>1 时不成立。 + +bs>1 的 row-major page table 形态是: + +```text +req0: [req0 prefix pages][req0 current pages][req0 padding/empty slots] +req1: [req1 prefix pages][req1 current pages][req1 padding/empty slots] +... +``` + +如果仍用 `[0, prefix_pages)`: + +1. 只能覆盖 req0 prefix; +2. 会把 req0 current / padding gap / req1 prefix 混在同一个 suffix bounding span 里; +3. 可能重复 reduce 无效 slot; +4. 更严重时会把不同 request 的 logical page 映射压到同一 dense view 中,复现之前 bs>1 cache-hit 精度下降类问题。 + +### C4. sync compose 已经有 bs>1 span helper,可复用语义 + +相关文件: + +- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py` + - `build_batch_prefix_slot_spans(...)` + - `build_batch_current_slot_spans(...)` + +已有语义: + +```text +build_batch_prefix_slot_spans: + 返回每个 request prefix 在 flattened real_page_table 里的精确 slot spans + +build_batch_current_slot_spans: + 返回每个 request current suffix 在 flattened real_page_table 里的精确 slot spans +``` + +结论:bs>1 prefetch 第一版应复用这两个 helper。不要重新发明 bounding span,也不要引入 batch-global scalar prefix。 + +### C5. 当前 TAI IPC materialize 是 SM kernel,不是 0SM + +相关文件: + +- SGLang: + - `python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py` + - `_try_tai_ipc_materialize_token_kv_page_slots_into(...)` + - `_try_tai_ipc_materialize_paged_buffer_page_slots_into(...)` +- tai-kernel: + - `python/tai_kernel/nsa_prefill/ipc.py` + - `materialize_cuda_ipc_peer_pages_slot_dense(...)` + - `python/tai_kernel/nsa_prefill/csrc/kvcacheio_lf_pf.cu` + - `ipc_materialize_peer_pages_slot_dense_kernel` + - `ipc_materialize_peer_pages_slot_dense_multi_warp_kernel` + +当前优势: + +1. 避免原来的 local materialize + all-reduce 冗余通信; +2. 可以直接 peer-read owner rank 的 L1 page; +3. 已支持 fp8 packed page 与 bf16 page 形态。 + +当前不足: + +1. kernel 使用 SM copy,不是 Copy Engine; +2. `start_slot != 0` 会 fallback; +3. 输入是一个连续 slot range,不是 spans list; +4. 用于 sync compose,不是 prefetcher 的 prefix pipeline。 + +结论:0SM 计划不能把现有 IPC kernel 直接称为 0SM。它是可用的 SM baseline,0SM 需要新增 CE/batched-copy primitive。 + +### C6. index active-layer skip 已经改变 index prefetch 的 layer 集合 + +index skip 后,index K/scale 只在 active index layers 存储。L1 index prefetch 必须满足: + +```text +如果 next_layer_id 不是 index active layer: + 不启动 index prefetch + 不访问 compacted index cache slot + +如果 next_layer_id 是 index active layer: + 使用 state_layer_id / active index layer id 映射后的 index buffer +``` + +MLA KV prefetch 仍然每个 target layer 都可能需要。 + +### C7. draft/EAGLE prefetch 当前应保持关闭 + +之前已经确认:draft KV 生命周期应是 target KV 的镜像,draft 不应该维护独立的 page ownership / HiCache state。draft prefetch 若直接套 target 的 `next_layer_id=layer+1`,容易出现: + +1. draft 只有一层或层数很少,`next_layer` 语义不稳定; +2. target/draft layer id 空间不同; +3. draft consume 时机在 target forward 之后,和 target prefetch window 不同。 + +结论:bs>1 L1 prefetch 第一版只打开 target。draft 后续作为单独阶段处理。 + +--- + +## 3. bs>1 L1 prefetch 语义设计 + +### 3.1 新的 prefetch plan 形态 + +建议在 `cp_shared_kv_prefetch.py` 内部引入一个小的 batch-aware plan dataclass,或者直接把同等字段放入两个 prefetcher: + +```python +@dataclass(frozen=True) +class CpSharedKVBatchPrefetchPlan: + batch_size: int + page_size: int + prefix_slot_spans: tuple[tuple[int, int], ...] + current_slot_spans: tuple[tuple[int, int], ...] + total_slots: int + dense_num_pages: int + prefix_slot_count: int + current_slot_count: int +``` + +字段合同: + +- `prefix_slot_spans`:只包含历史 cache-hit prefix pages。 +- `current_slot_spans`:只包含本次 extend 产生的 current pages。 +- `total_slots`:`slot_logical_pages.numel()`。 +- `dense_num_pages`:dense buffer 的 page 数,沿用现有 remap helper 结果。 +- `prefix_slot_count/current_slot_count`:用于日志、阈值、benchmark,不用于切分 tensor。 + +构造方式: + +```text +logical_pages = metadata.real_page_table +prefix_slot_spans = build_batch_prefix_slot_spans( + logical_pages=logical_pages, + prefix_lens_cpu=forward_batch.extend_prefix_lens_cpu, + page_size=page_size, +) +current_slot_spans = build_batch_current_slot_spans( + logical_pages=logical_pages, + prefix_lens_cpu=forward_batch.extend_prefix_lens_cpu, + extend_lens_cpu=forward_batch.extend_seq_lens_cpu, + page_size=page_size, +) +``` + +### 3.2 prefetch create gate + +替换当前 `batch_size != 1` guard 后,create gate 应变成: + +```text +if batch_size < 1: + skip/fail + +if no prefix_slot_spans: + skip, reason=no_prefix_pages + +if sum(prefix span pages) < min_prefix_pages: + skip, reason=prefix_below_min + +if sum(extend_seq_lens_cpu) < min_async_extend_tokens: + skip, reason=extend_below_min +``` + +注意:threshold 应基于 batch aggregate work,而不是单个 request。bs>1 的目的就是把多个 200-2000 token extend 合起来打满 compute;如果逐 request threshold,会错误关闭有收益的 batch prefetch。 + +### 3.3 prefix materialize + +bs>1 第一版可先使用现有 SM IPC / local materialize 语义,但必须接受 spans: + +```text +for (start_slot, end_slot) in prefix_slot_spans: + materialize slot range into dense buffer 对应位置 +``` + +不能把 spans 合并成一个包含 current/gap 的 bounding span,除非合并后仍能证明: + +1. 中间 gap 都是有效 prefix; +2. 不包含 current suffix; +3. 不包含 empty/padded page slots; +4. dense remap 对 gap 的写入不会影响 attention 可见范围。 + +当前 row-major bs>1 下通常不能证明这些条件,因此默认使用 exact spans。 + +### 3.4 consume/current fill + +consume 时: + +```text +wait prefix_prefetch_event(layer L) +for (start_slot, end_slot) in current_slot_spans: + materialize current pages / fill current rows + reduce current rows if仍使用 collective path +remap logical locs to dense locs +return dense_kv_cache, dense_locs +``` + +当前 suffix 不应再用: + +```text +start_slot = prefix_pages +end_slot = total_slots +``` + +这会在 bs>1 下把多个 request 的 prefix/current/gap 混在一起。 + +### 3.5 page_inverse 要保持 per-request 语义 + +之前 cache-hit 掉点的强 root cause 是 batch-global logical page inverse 把不同 request 的相同 logical page value 合并到同一个 dense page。prefetch 也必须遵守同一规则: + +```text +page_inverse key = (req_id, logical_page) +不能只用 logical_page +``` + +现有 remap helper 已经有 per-row sorted/dense metadata,prefetch 不应绕开它创建 batch-global inverse。 + +--- + +## 4. 0SM/低 SM transport 设计 + +### 4.1 为什么现有 IPC kernel 不够 + +现有 `materialize_cuda_ipc_peer_pages_slot_dense` 是 GPU SM copy kernel。它适合低延迟 SM copy baseline,但作为 prefetch 长 prefix 时有两个问题: + +1. 会和当前 layer attention / MQA / MOE 抢 SM; +2. 如果 prefetch 过长,可能拖慢当前 layer,抵消 overlap 收益。 + +目标 0SM primitive 应把 prefix page materialize 变成 Copy Engine 或 CPU gather + CE DMA 路径。 + +### 4.2 0SM 第一版 API 建议 + +在 tai-kernel 中新增独立 API,保留现有 SM IPC kernel: + +```python +materialize_cuda_ipc_peer_pages_slot_dense_ce( + peer_ptrs, + dst, + owner_ranks, + src_page_indices, + dst_page_indices, + page_nbytes: int, + stream, +) -> None +``` + +语义: + +```text +for i in [0, page_count): + src = peer_ptrs[owner_ranks[i]] + src_page_indices[i] * page_nbytes + dst = dst_base + dst_page_indices[i] * page_nbytes + copy page_nbytes bytes +``` + +实现路线: + +1. C++/CUDA extension 中构造 CPU descriptor arrays。 +2. 对连续 `(src, dst)` runs 做 merge,减少 CE command 数。 +3. 优先使用 `cudaMemcpyBatchAsync`(CUDA >= 12.8)或多次 `cudaMemcpyPeerAsync` baseline。 +4. 如果运行环境不支持 batch API,显式 warning/fail-fast 到 SM IPC baseline,不 silent fallback 到 all-reduce。 + +### 4.3 descriptor 必须避免 hidden D2H + +0SM CE path 不能每层把 CUDA tensor `.cpu()` 成 descriptor。descriptor 来源应是: + +```text +ForwardBatch CPU metadata + -> request prefix/current slot spans + -> owner rank + physical page index descriptors + -> cached CPU pinned descriptor arrays +``` + +要求: + +1. `owner_ranks/src_page_indices/dst_page_indices` 尽量在 batch prepare/create 阶段构造; +2. layer 间只换 source base pointer,不重建 page mapping; +3. descriptor cache key 应包含: + - `slot_logical_pages` identity/version; + - layout cp_rank/cp_size; + - page size; + - prefix/current spans; + - dtype/page_nbytes。 + +### 4.4 readiness / ordering 风险 + +all-reduce 路径天然有跨 rank 同步语义;0SM peer-read 路径没有 collective barrier。必须明确 ordering: + +```text +L2->L1 load finished on owner rank + -> owner rank device page visible to peer CE read + -> non-owner rank prefetch CE copy starts +``` + +可选实现: + +1. 第一版保守:L1 prefetch 只读本 rank owner pages,跨 rank prefix 仍走现有 all-reduce/SM IPC path。 +2. 更完整:为 L2->L1 load 暴露 per-layer CUDA IPC event 或 ready flag,prefetch stream 在读 peer page 前等待对应 peer ready。 +3. 依赖 lockstep layer schedule:所有 CP ranks 在同一个 layer 同步推进,L2->L1 提前 2 层足够覆盖。但这需要 ETE/Nsight 验证,不能作为唯一正确性保证。 + +推荐实现顺序:先接入 SM IPC spans baseline 验证 bs>1 correctness,再做 0SM CE path;0SM 首版只在能证明 ready 的 prefix pages 上启用。 + +--- + +## 5. 实施 phases + +### P0. 文档与现状锁定 + +**文件:** + +- 新增:`docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_l1_prefetch_zero_sm_plan_zh.md` + +**验收:** + +- 文档明确 L2->L1 `+2 layer`、L1 prefetch `+1 layer`。 +- 文档明确当前 prefetcher 的 bs=1 假设和风险。 +- 文档明确 0SM 与现有 SM IPC kernel 的区别。 + +### P1. 补 bs>1 prefetch plan 单测 + +**文件:** + +- 修改:`test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py` +- 修改:`test/registered/unit/mem_cache/test_cp_shared_kv_prefetch.py`(如果不存在则新增同级测试文件) + +**测试用例:** + +1. `batch_size=2`: + +```text +page_size=64 +req0 prefix=128, extend=65 +req1 prefix=64, extend=129 +``` + +期望: + +```text +prefix_slot_spans = [(0, 2), (pages_per_req, pages_per_req + 1)] +current_slot_spans = [(2, 4), (pages_per_req + 1, pages_per_req + 4)] +``` + +2. prefix 为 0 的 request 不产生 prefix span,但 current span 仍存在。 +3. 不允许 `prefix_pages=sum(prefix_pages)` 后直接切 `[0:prefix_pages)`。 +4. page inverse 按 `(req_id, logical_page)` 区分两个 request 的相同 logical page value。 + +**远端命令:** + +```bash +cd /sgl-workspace/sglang-tai +PYTHONPATH=python python -m pytest -q \ + test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py \ + test/registered/unit/mem_cache/test_cp_shared_kv_prefetch.py +``` + +### P2. 改造 MLA prefetcher 支持 bs>1 spans + +**文件:** + +- 修改:`python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py` +- 修改:`python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py`(仅在缺少 helper 时修改) + +**实现点:** + +1. `CpSharedKVMlaPrefetcher.maybe_create()` 不再因为 `batch_size != 1` 直接 skip。 +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。 + +**第一版允许:** + +- prefix 使用现有 local materialize + all-reduce 或 SM IPC baseline。 + +**第一版不允许:** + +- silent fallback 到 `[prefix_pages:total_slots)`。 +- batch-global page inverse。 + +### P3. 改造 index prefetcher 支持 bs>1 spans + +**文件:** + +- 修改:`python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py` +- 修改:`python/sglang/srt/layers/attention/nsa/nsa_indexer.py`(只在 consume 接口需要更明确的 active-layer 映射时修改) + +**实现点:** + +1. `CpSharedKVIndexPrefetcher.maybe_create()` 使用同一套 batch prefetch plan。 +2. index prefetch 只对 active index layer 启动。 +3. index `dense_page_buffer` 的 remap 和 current spans 必须 request-aware。 +4. index skip layer 不创建 handle,不产生 warning。 + +**风险控制:** + +- 如果 active index layer 映射缺失,fail-fast: + +```text +[CP_SHARED_KV_FAIL_FAST][index_prefetch] reason=missing_active_state_layer +``` + +### P4. 接入 L2->L1 +2 layer readiness + +**文件:** + +- 修改:HiCache load/prefetch 管理相关文件,优先从已有文档定位: + - `docs/advanced_features/nsa_prefill_cp_hicache_load_prefetch_overlap_notes.md` + - `docs/advanced_features/nsa_prefill_cp_hicache_async_state_sync.md` +- 预计涉及: + - `python/sglang/srt/mem_cache/hiradix_cache.py` + - `python/sglang/srt/mem_cache/hicache_storage.py` + - `python/sglang/srt/layers/attention/nsa_backend.py` + +**实现点:** + +1. 暴露 `ensure_hicache_load_started(layer_id)` 或等价机制。 +2. 在处理 layer `L-2` 时启动 layer `L` 的 L2->L1 load。 +3. L1 prefetch stream 在处理 layer `L-1` 时等待 layer `L` 的 L2->L1 load event。 +4. 如果当前实现已经一次性提交 all-layer load,应记录为 `load_window=all_layers`,并先只接 wait event,不额外提交重复 load。 + +**不能做:** + +- forward stream 上同步等待 L2->L1。 +- 用新的 all-reduce 确认所有 rank load 完成。 + +### P5. TAI SM IPC spans baseline + +**文件:** + +- SGLang: + - `python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py` +- tai-kernel: + - `python/tai_kernel/nsa_prefill/ipc.py` + - `python/tai_kernel/nsa_prefill/csrc/kvcacheio_ops.cc` + - `python/tai_kernel/nsa_prefill/csrc/kvcacheio_lf_pf.cu` + - `tests/nsa_prefill/test_cuda_ipc_gather.py` + - `benchmark/nsa_prefill/benchmark_cp_shared_kv_ipc_gather.py` + +**实现点:** + +1. 支持 spans list,而不是只支持 `start_slot == 0` 的连续 range。 +2. 支持 `dst_page_indices`,使 prefix pages 可以写到 dense buffer 的正确 page 槽位。 +3. 支持 bf16/fp8 同一接口,通过 `page_nbytes` 控制。 +4. benchmark 覆盖: + - bs=1/2/5/10; + - prefix pages 对应 4k/8k/16k/32k/64k/120k/192k tokens; + - contiguous / fragmented / owner-lane / zigzag page pattern。 + +### P6. 0SM CE prototype + +**文件:** + +- tai-kernel 新增 CE/batched-copy op;具体文件沿用 P5。 + +**实现点:** + +1. C++ extension 内部构造并提交 CE copy。 +2. 支持 run merge,避免每 page 一个 API 调用。 +3. 提供 `--mode sm_ipc|ce_batch|cuda_memcpy_batch` benchmark 对比。 +4. 如果 CUDA runtime 不支持 `cudaMemcpyBatchAsync`,测试中显式标记 skip;生产中 warning 后回 SM IPC,不回 all-reduce。 + +**验收指标:** + +1. Nsight 中 CE path 不出现长时间 copy kernel SM occupancy。 +2. 对大 prefix,当前 layer attention/MOE SM occupancy 不被 prefetch kernel 明显挤占。 +3. 小 prefix 下 CE descriptor overhead 不应超过 SM IPC baseline;否则按阈值选择 SM IPC。 + +### P7. ETE 验证与默认开关 + +**远端验证:** + +1. GSM8K 两轮: + +```bash +cd /sgl-workspace/sglang +python benchmark/gsm8k/bench_sglang.py \ + --host g0034 \ + --port 17100 \ + --backend srt \ + --data-path /mnt/beegfs/cjy/data/gsm8k_test.jsonl \ + --num-questions 1319 \ + --parallel 64 \ + --temperature 0.0 \ + --top-p 1.0 \ + --max-new-tokens 512 +``` + +验收: + +```text +第一轮 accuracy 接近 0.955 +第二轮 cache-hit accuracy 不下降到 0.68/0.69 +Invalid 接近 0 +``` + +2. replay 压测: + +```bash +ssh g0033 +source /mnt/beegfs/khr/bench/.venv/bin/activate +python3 /mnt/beegfs/khr/bench/replay_openai_request_log.py \ + --request-log /mnt/beegfs/khr/bench/workloads/request-chain-mixed-1500-15min-2rps/requests.jsonl \ + --base-url http://g0034:17100 \ + --endpoint /v1/chat/completions \ + --chain-prefill-scheduler \ + --prefill-target-concurrency 30 \ + --max-concurrency 60 \ + --metric-percentiles 50,90,95,99 \ + --timeout-sec 240 \ + --save-result /mnt/beegfs/results/replay-request-chain-mixed-1500-k20-prefill.json +``` + +验收: + +```text +无 detokenizer hang +无 cache-hit 后乱码/accept len 大幅掉到 1 的回归 +无 CP shared-KV fallback 热路径刷屏 +prefill 吞吐不低于当前 bs>1 sync baseline +``` + +--- + +## 6. 日志与 fail-fast 要求 + +### 6.1 必须 warning/fail-fast 的情况 + +1. bs>1 prefetcher 发现缺少 `prefix_slot_spans/current_slot_spans`。 +2. `prefix_slot_spans` 与 `real_page_table` row 数不匹配。 +3. `page_inverse` 不是 request-aware。 +4. index active layer 映射缺失。 +5. L1 prefetch 试图读取尚未完成 L2->L1 load 的 layer。 +6. 0SM CE path 不支持当前 dtype/page layout,但配置要求强制 0SM。 + +### 6.2 不应 warning 刷屏的情况 + +1. prefix 为 0。 +2. prefix/extend 低于阈值,按策略跳过 prefetch。 +3. index skip layer 不需要 prefetch。 +4. draft/EAGLE prefetch 第一版未启用。 + +这些只应进入 debug/timing 日志,不能在生产路径刷 warning。 + +--- + +## 7. 剩余风险 + +### R1. peer-read ordering 不是 collective ordering + +如果 0SM/IPC prefix prefetch 直接读 peer L1 pages,需要证明 owner rank 的 L2->L1 load 和当前 layer writes 已经完成。否则会读到旧 KV 或未初始化 KV。 + +第一版规避方式: + +```text +只对历史 prefix pages 启用 prefetch; +current suffix 仍由当前 layer 同步 compose/reduce; +L2->L1 load 至少提前两个 layer; +必要时 prefetch stream 等 local load event。 +``` + +### R2. CE descriptor CPU overhead 可能抵消收益 + +如果每 layer 都构造 Python list / CPU descriptor / H2D metadata,0SM copy 本身省下的 SM 可能被 CPU overhead 抵消。 + +要求: + +```text +descriptor 在 ForwardBatch/create 阶段缓存; +layer 只更新 base pointer / stream event; +benchmark 单独统计 descriptor build 与 submit latency。 +``` + +### R3. fp8 page_nbytes 必须来自真实 tensor stride + +fp8 KV cache 之前出现过 paged/ragged shape 不一致问题。prefetch/0SM 不能假设 `hidden_dim * dtype_size`,必须使用真实 page tensor stride 计算: + +```text +page_nbytes = stride(page_dim) * element_size +``` + +### R4. bs>1 tiny extend 不能因为 prefetch skip 丢失 current/full reuse + +prefetch 只是 overlap 优化。即使 prefix/extend 太短跳过 prefetch,也必须保持: + +```text +current reuse / partial-current reuse / full current sync compose 正确可用 +``` + +不能让 “不创建 prefetcher” 变成 “不走 reuse”。 + +### R5. index skip 下的 index prefetch 不应访问被压缩掉的 layer + +P3-P6 已经压缩 L1/L2 index cache。prefetch 必须复用 active-layer 映射,否则会访问错误 index state layer 或越界。 + +--- + +## 8. 推荐落地顺序 + +最小可验证路径: + +```text +P0 文档 +P1 单测锁住 spans/page_inverse 语义 +P2 MLA bs>1 prefetch spans baseline +P3 index bs>1 prefetch spans baseline +P4 接 L2->L1 +2 readiness +P5 TAI SM IPC spans baseline benchmark +P6 0SM CE prototype +P7 GSM8K/replay/Nsight 验证 +``` + +不要先做 P6 再修 P2/P3。0SM kernel 只解决 transport 代价,不解决 bs>1 prefix/current slot 语义;如果语义仍是 scalar prefix,kernel 越快只会越快地产生错误。