Keep CP shared-KV prefetch warnings actionable
Expected no-prefetch paths were polluting production logs: no cache prefix, tiny/first-layer windows, and FP8 RAGGED top-k were being reported as fallback warnings. The prefetch contract now treats zero-prefix and first-layer misses as normal skips, while preserving warnings for non-zero misaligned prefixes and real consume misses after the first layer. The same change keeps RAGGED cache-hit prefetch eligible and records the CE/IPM prefetch contract in the plan doc. Constraint: FP8 sparse prefill uses RAGGED top-k, but CP shared-KV prefix materialization is still page-slot based Constraint: Layer 0 has no previous attention-window hook that can have prefetched the layer Rejected: Warn whenever a prefetcher is absent | no-cache and too-short requests are expected synchronous paths and make logs unusable Confidence: high Scope-risk: moderate Directive: Keep CP_SHARED_KV_FALLBACK warnings for unexpected contract failures only; use debug logs for expected skip paths Tested: Local py_compile for cp_shared_kv_prefetch.py, nsa_indexer.py, nsa_backend.py Tested: Remote cjy-glm5-new targeted regression: 3 passed, 21 warnings Tested: Remote cjy-glm5-new full test_cp_shared_kv_runtime.py: 156 passed, 21 warnings, 2 subtests passed Not-tested: New ETE run after prefill restart to confirm log volume reduction in production traffic (cherry picked from commit e08e321e5929fdbb30102ec0b19c6ff0ecac7e7e)
This commit is contained in:
laoyao0822
@@ -430,15 +430,37 @@ L2->L1 load finished on owner rank
|
||||
|
||||
### P0. 文档与现状锁定
|
||||
|
||||
**文件：**
|
||||
**状态（2026-06-12）：已完成。**
|
||||
|
||||
- 新增：`docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_l1_prefetch_zero_sm_plan_zh.md`
|
||||
完成内容：
|
||||
|
||||
**验收：**
|
||||
1. L1 prefetch create 路径不再以 `get_attention_cp_group().pynccl_comm` 作为启用条件。
|
||||
prefetch 的新合同是 IPC/CE materialize，不再依赖 CP collective communicator。
|
||||
2. `cp_shared_kv_prefetch.py` 与 `nsa_backend.py` 中 prefetch hot path 已清理旧
|
||||
`launch_pending_reduce` / `prefix_reduce` / `reduce_enqueue` / `collective_disallowed`
|
||||
命名，统一改成 `finalize_pending_materialize` / `prefix_materialize` /
|
||||
`materialize_enqueue` / `materialize_event_missing`。
|
||||
3. prefetch 文件中不再 import/use `_all_reduce_materialized_buffer_*` 或
|
||||
`get_attention_cp_group`。如果 IPC/CE materialize 不可用，CP>1 下应显式
|
||||
fail-fast，而不是 silent fallback 到 all-reduce。
|
||||
4. 测试已覆盖：`maybe_create()` 不依赖 pynccl，attention-window 不再触发
|
||||
prefix reduce 旧路径，tiny extend skip 测试也不再 mock pynccl。
|
||||
|
||||
- 文档明确 L2->L1 `+2 layer`、L1 prefetch `+1 layer`。
|
||||
- 文档明确当前 prefetcher 的 bs=1 假设和风险。
|
||||
- 文档明确 0SM 与现有 SM IPC kernel 的区别。
|
||||
验收证据：
|
||||
|
||||
```text
|
||||
远端 cjy-glm5-new:
|
||||
PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
|
||||
150 passed, 21 warnings, 2 subtests passed
|
||||
```
|
||||
|
||||
静态检查：
|
||||
|
||||
```text
|
||||
rg "pynccl|launch_pending_reduce|prefix_reduce|collective_disallowed|reduce_enqueue|get_attention_cp_group|_all_reduce" \
|
||||
cp_shared_kv_prefetch.py nsa_backend.py
|
||||
# no matches
|
||||
```
|
||||
|
||||
### P1. 补 bs>1 prefetch plan 单测
|
||||
|
||||
@@ -508,17 +530,17 @@ PYTHONPATH=python python -m pytest -q \
|
||||
- `current_slot_spans`
|
||||
- `prefix_page_count`
|
||||
- `current_page_count`
|
||||
3. `start_next_layer_prefix()` 只 materialize/reduce `prefix_slot_spans`，
|
||||
3. `start_next_layer_prefix()` 只 materialize `prefix_slot_spans`，
|
||||
不再把 batch flattened page table 当成 `[0:prefix_pages)`。
|
||||
4. `consume_prefix_with_current()` 只 reduce `current_slot_spans`，避免把
|
||||
row gap / 其他 request prefix 一起 reduce。
|
||||
4. `consume_prefix_with_current()` 只处理 `current_slot_spans`，避免把
|
||||
row gap / 其他 request prefix 一起纳入 suffix compose。
|
||||
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。
|
||||
- MLA prefix spans 已不再走 async all-reduce fallback；当前要求 IPC/CE materialize
|
||||
生成 completion event，CP>1 下不可用则 fail-fast。
|
||||
- index prefetcher 仍未改造，继续由 P3 处理。
|
||||
|
||||
**文件：**
|
||||
@@ -531,13 +553,13 @@ PYTHONPATH=python python -m pytest -q \
|
||||
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。
|
||||
4. `consume_prefix_with_current()` 只处理 current spans。
|
||||
5. `consume()` 如果仍存在 legacy suffix path，必须使用 `current_slot_spans`
|
||||
或 fail-fast，不能回到错误 bounding suffix。
|
||||
|
||||
**第一版允许：**
|
||||
|
||||
- prefix 使用现有 local materialize + all-reduce 或 SM IPC baseline。
|
||||
- prefix 使用现有 SM IPC 或 CE materialize baseline。
|
||||
|
||||
**第一版不允许：**
|
||||
|
||||
@@ -567,8 +589,8 @@ RED 证据：旧代码在 `batch_size=2` 时以
|
||||
- `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。
|
||||
3. `start_next_layer_prefix()` 只 materialize index prefix spans。
|
||||
4. `consume_prefix_with_current()` 只处理 index current spans。
|
||||
5. `consume()` 的 legacy suffix 路径也改为 `current_slot_spans`，
|
||||
不再使用 batch bounding suffix。
|
||||
|
||||
@@ -581,8 +603,8 @@ test_cp_shared_kv_runtime.py
|
||||
|
||||
当前限制：
|
||||
|
||||
- index prefetch 仍使用现有 materialize + async all-reduce baseline；
|
||||
spans-list TAI IPC / 0SM CE 留给后续 P5/P6。
|
||||
- index prefetch 已不再走 async all-reduce fallback；当前要求 IPC/CE materialize
|
||||
生成 completion event，CP>1 下不可用则 fail-fast。
|
||||
- active index layer / index skip 的 runtime hook 当前沿用已有
|
||||
`nsa_backend.py` 调用点；本阶段没有修改 skip 参数语义。
|
||||
|
||||
@@ -637,11 +659,11 @@ test_cp_shared_kv_runtime.py
|
||||
- `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 上异步提交。
|
||||
- 修正策略：MLA/index 的 `start_next_layer_prefix()` 先取得 `current_stream`，把它传给 prefetch-safe getter，使 L2->L1 ready event 挂到实际读取 KV/index page 的 stream；随后仍保持 stream ordering，由 IPC/CE materialize event 表示 prefix ready。
|
||||
- 没有新增 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`
|
||||
- `test_mla_prefetch_waits_l2_l1_on_materialize_stream`
|
||||
- `test_index_prefetch_waits_l2_l1_on_materialize_stream`
|
||||
|
||||
### P5. TAI SM IPC spans baseline
|
||||
|
||||
@@ -995,3 +1017,79 @@ SYMM CE prefix + SM current: gpu_p50=0.521ms cpu_p50=0.546ms
|
||||
2. 当前 PyTorch symmetric-memory prototype 没有优于 CUDA IPC CE；推测原因包括 symmetric memory backend/registration overhead、统一 window size 带来的空间浪费、以及仍使用 cudaMemcpyBatchAsync per-span submit。
|
||||
3. 短期最有价值路径是继续用 IPC CE 做 prefix prefetch baseline，同时保留 symmetric window helper 作为后续 NCCL CE/NVSHMEM transport 的地址合同实验入口。
|
||||
4. 生产接入前必须继续做 span coalescing、跨 layer descriptor 复用，以及更大 bs/不同 cached/extend sweep；不能只凭 symmetric window 概念直接替换现有路径。
|
||||
|
||||
---
|
||||
|
||||
## 10. FP8/RAGGED 与 index-skip prefetch 修正（2026-06-12）
|
||||
|
||||
### 10.1 远端日志现象
|
||||
|
||||
最新 prefill 日志中 `index_prefetch` fallback 被刷屏：
|
||||
|
||||
```text
|
||||
[CP_SHARED_KV_FALLBACK][index_prefetch]
|
||||
current_missing_prefetcher: 53136
|
||||
not_paged_topk: 1296
|
||||
```
|
||||
|
||||
排查结论：
|
||||
|
||||
1. `current_missing_prefetcher` 大多来自无 cache / tiny extend / prefetcher 未创建的正常同步路径，不是运行时错误；consume 侧不应每层每 rank warning。
|
||||
2. `not_paged_topk` 是错误判定：FP8 sparse prefill 走 `TopkTransformMethod.RAGGED` 是预期行为，但 index/MLA prefix materialize 仍是 page-slot 语义，不能因为 top-k transform 不是 PAGED 就关闭 prefetch。
|
||||
|
||||
### 10.2 当前修正
|
||||
|
||||
1. `CpSharedKVIndexPrefetcher.maybe_create()` 对 RAGGED top-k 不再 fallback，而是继续创建 index prefetcher。
|
||||
2. `CpSharedKVMlaPrefetcher.maybe_create()` 对 RAGGED top-k 也继续创建 MLA prefetcher。
|
||||
3. `nsa_indexer._maybe_materialize_shared_index_buffer()` 不再对缺失 prefetcher 的正常同步路径打 warning；仅在 prefetcher 已存在但 consume miss 时保留 fallback warning。
|
||||
4. `nsa_backend.py` 的 FP8/RAGGED cache-hit compose 现在先尝试 `mla_prefetcher.consume_prefix_with_current()`，只有 prefetcher 存在但 miss 时才 warning 并回到同步 `materialize_prefix_and_reuse_current_kv_page_slots()`。
|
||||
5. attention-window 启动 index prefetch 时检查 `token_to_kv_pool.index_layer_to_slot`，index-skip/shared 层没有 index-cache slot 时只跳过 index prefetch，不影响 MLA prefetch。
|
||||
|
||||
### 10.3 验证
|
||||
|
||||
远端容器 `cjy-glm5-new` 已验证：
|
||||
|
||||
```text
|
||||
python -m py_compile \
|
||||
python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py \
|
||||
python/sglang/srt/layers/attention/nsa_backend.py
|
||||
|
||||
PYTHONPATH=python python -m pytest -q \
|
||||
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py::TestCpSharedKVRuntimeHelpers::test_fp8_ragged_mla_uses_page_slot_current_compose_for_cache_hit \
|
||||
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py::TestCpSharedKVRuntimeHelpers::test_mla_prefetch_create_allows_ragged_topk_transform \
|
||||
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py::TestCpSharedKVRuntimeHelpers::test_index_prefetch_create_allows_ragged_topk_transform \
|
||||
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py::TestCpSharedKVTaiMaterializeIntegration::test_attention_window_prefetch_skips_inactive_index_cache_layer
|
||||
结果：4 passed
|
||||
```
|
||||
|
||||
### 10.4 剩余风险
|
||||
|
||||
1. 旧进程日志不会反映本次降噪和 RAGGED prefetch 修正；需要重启 prefill 后再观察 warning 频率。
|
||||
2. RAGGED 当前层 materialize 的历史代码位置仍在 flashmla_sparse 分支内部；后续如果严格追求“当前层 materialize 完成后再启动下一层 prefetch”，需要把 RAGGED current compose 前移到 attention-window hook 之前。
|
||||
3. 本次只修正 create/consume 合同和单测；还需要 ETE 确认 FP8 cache-hit 下 index/MLA prefetch 命中率和精度不回退。
|
||||
|
||||
### 10.5 预期 skip 不再 warning（2026-06-12）
|
||||
|
||||
远端新进程仍出现大量 fallback warning 后，继续按日志聚合定位：
|
||||
|
||||
```text
|
||||
[CP_SHARED_KV_FALLBACK][index_prefetch] reason=current_consume_miss 主要集中在 layer=0
|
||||
[CP_SHARED_KV_FALLBACK][mla_prefetch] reason=ragged_current_consume_miss 主要集中在 layer=0
|
||||
[CP_SHARED_KV_FALLBACK][index_prefetch|mla_prefetch] reason=prefix_not_page_aligned prefix_lens=[0]
|
||||
```
|
||||
|
||||
结论：这些都不是异常 fallback：
|
||||
|
||||
1. `layer=0` 没有上一层 attention-window 能提前发起本层 prefetch，consume miss 是预期同步路径。
|
||||
2. `prefix_lens=[0]` 表示 request 本身没有命中 cache，prefix prefetch 不应创建。
|
||||
3. 真实需要 warning 的场景应收窄为：非零 prefix 不 page-aligned、prefetcher 已创建且 `layer > start_layer` 仍 consume miss、metadata/slot 合同异常等。
|
||||
|
||||
修正合同：
|
||||
|
||||
1. MLA/index prefetch create 遇到全零 prefix 时只走 debug skip，不再打 `[CP_SHARED_KV_FALLBACK]` warning。
|
||||
2. index partial-current prefetcher 在 `layer == token_to_kv_pool.start_layer` miss 时不 warning。
|
||||
3. FP8/RAGGED MLA prefetcher 在 `layer == start_layer` miss 时不 warning。
|
||||
|
||||
这类 skip 仍保留 debug 级别路径，方便需要时通过 debug 环境变量观察，但默认生产日志不应被无 cache/tiny/首层预期路径污染。
|
||||
|
||||
验证：远端 `cjy-glm5-new` 已通过 targeted regression（3 passed）以及完整 `test_cp_shared_kv_runtime.py`（156 passed, 21 warnings, 2 subtests passed）。
|
||||
|
||||
Reference in New Issue
Block a user