Align CP shared-KV prefetch with the attention overlap window
Launching CP shared-KV prefetch from MLA prepare or the indexer can make next-layer prefix work overlap current-layer MQA/materialization instead of the attention window. Centralize the launch in the NSA backend after current-layer materialization and before attention, and leave the early indexer hook inert so the call site cannot regress silently. The accompanying notes capture the draft-as-forward-layer follow-up plan and the latest OOM diagnosis: the observed 178945-token failure matches the CP in-seq MQA logits allocation, so the follow-up fix is q-dimension chunking in _get_topk_ragged_with_cp(), not a max-prefetch-size gate. Constraint: CP shared-KV prefetch must avoid overlapping current-layer MQA/materialization and must not add silent fallback behavior. Rejected: Limit maximum prefetch size | hides the CP logits peak and can reduce cache/prefetch effectiveness. Rejected: Keep prepare/indexer launch sites | they place next-layer collectives in the wrong overlap window. Confidence: medium Scope-risk: moderate Directive: Do not reintroduce early CP prefetch launch without checking Nsight overlap and CP MQA memory peaks. Tested: Local git diff --check and py_compile for touched Python files. Tested: Remote container py_compile plus targeted pytest: 3 passed, 5 warnings. Not-tested: Full ETE under production traffic after this commit. Not-tested: CP in-seq MQA logits chunking; documented as follow-up.
This commit is contained in:
@@ -472,3 +472,112 @@ CP draft shared KV fallback: unsupported draft architecture
|
||||
4. draft persistent KV 使用 CP shared-KV physical shard,而不是每 rank full KV。
|
||||
5. prefill->decode draft KV transfer 正常。
|
||||
6. fallback 路径安全、可观测,非 GLM-5 draft 不误走新路径。
|
||||
|
||||
## 2026-06-01:Draft layer as forward-layer in prefetch/backup pipeline plan
|
||||
|
||||
### 背景
|
||||
|
||||
当前 EAGLE / NextN draft 已经作为 target HiCache 的 shadow payload 接入了一部分生命周期:
|
||||
|
||||
- `CacheController.attach_draft_pool()` 会把 draft device/host pool 绑定到 target controller,并注册同一个 `LayerDoneCounter`。
|
||||
- draft decoder layer end 会通过 `notify_layer_end_for_backup()` 进入 `CacheController.on_layer_end(layer_id, source="draft")`,再提交 draft D2H backup。
|
||||
- `CacheController.start_loading()` 已经会把 target load queue 与 draft load queue 一起处理。
|
||||
|
||||
但 CP shared-KV 的 index/MLA prefetch 当前仍把 draft 排除在外:
|
||||
|
||||
- `nsa_backend.init_forward_metadata()` 对 `cp_shared_kv_is_draft_input(forward_batch)` 显式禁用 prefetcher 创建。
|
||||
- `cp_shared_kv_should_prefetch_next_layer()` 对 draft input 返回 `False`。
|
||||
- draft/NextN 只有一个 executable layer,`layer_id=0`,不能直接复用 target 的 `next_layer` prefetch 语义。
|
||||
|
||||
### 设计目标
|
||||
|
||||
把 draft 层视为 target forward 后的一个逻辑 forward layer:
|
||||
|
||||
```text
|
||||
target layer 0
|
||||
...
|
||||
target layer N-1
|
||||
draft layer 0
|
||||
```
|
||||
|
||||
但第一阶段不重构全局 event/counter timeline;先让 draft layer 拥有独立的 same-layer prefetch 合同:
|
||||
|
||||
```text
|
||||
draft metadata ready
|
||||
-> start draft current-layer index/MLA prefix prefetch
|
||||
-> draft layer 0 consumes prefix + current KV
|
||||
-> draft layer 0 end submits async backup
|
||||
```
|
||||
|
||||
### 分阶段计划
|
||||
|
||||
#### P0:锁住当前 draft 合同
|
||||
|
||||
- 测试 draft input 不应走 target-style next-layer prefetch。
|
||||
- 测试 draft layer id 只允许 `0`。
|
||||
- 测试 draft backup notifier 仍通过 `source="draft"` 触发。
|
||||
- 保留 fail-fast,不允许静默 fallback 到 target prefetch 路径。
|
||||
|
||||
#### P1:新增 draft same-layer prefetch API
|
||||
|
||||
新增独立语义,避免误用 target `start_next_layer_prefix()`:
|
||||
|
||||
```python
|
||||
start_current_layer_prefix(layer_id=0, token_to_kv_pool=draft_pool)
|
||||
```
|
||||
|
||||
约束:
|
||||
|
||||
- 只允许 draft / NextN path 使用。
|
||||
- 只允许 `layer_id == 0`。
|
||||
- metadata 必须已经构造完成。
|
||||
- 不满足条件时明确 warning/fail-fast,不做静默慢 fallback。
|
||||
|
||||
#### P2:把 draft prefetch 启动点放到 draft metadata ready 后、attention 前
|
||||
|
||||
目标顺序:
|
||||
|
||||
```text
|
||||
draft init_forward_metadata
|
||||
pre-attention hook:
|
||||
start index prefix prefetch
|
||||
start MLA prefix prefetch
|
||||
draft indexer / attention consume:
|
||||
wait pending event if needed
|
||||
```
|
||||
|
||||
这和 target 当前的 attention-front prefetch 位置保持一致,但语义是 current-layer,不是 next-layer。
|
||||
|
||||
#### P3:先不重构 backup 主流程
|
||||
|
||||
当前 draft backup 已经在 draft layer end 触发。第一阶段只验证:
|
||||
|
||||
- 是否确实 per-layer async 提交;
|
||||
- 是否还有 catch-up fallback;
|
||||
- final ack 是否仍等待 target + draft payload 完成。
|
||||
|
||||
若发现 draft backup 延迟到 request 结束,再单独修复。
|
||||
|
||||
#### P4:后续再评估 target-last-layer -> draft-layer bridge
|
||||
|
||||
更激进的优化是让 target 最后一层提前发起 draft layer 0 prefetch。但这需要更早构造 draft metadata,而 draft forward batch 当前依赖 target 输出 hidden / verified token,风险较高。
|
||||
|
||||
暂不作为第一阶段实现项。
|
||||
|
||||
#### P5:后续再统一虚拟 layer timeline
|
||||
|
||||
理想模型是:
|
||||
|
||||
```text
|
||||
event 0..N-1: target layers
|
||||
event N: draft layer 0
|
||||
```
|
||||
|
||||
但这会影响 `LayerDoneCounter`、draft H2D wait、backup ack、target/draft strong-sync 等多个路径。只有当 P1-P3 证明收益不足或现有 event 绑定导致 correctness/perf 问题时再做。
|
||||
|
||||
### 风险记录
|
||||
|
||||
- 直接打开 draft prefetcher 可能复现之前 EAGLE hang;draft 没有 next-layer,必须单独建 current-layer API。
|
||||
- draft metadata 生成时间晚于 target metadata,target last-layer 提前 prefetch 需要额外桥接。
|
||||
- index/MLA prefix materialize 会占 SM,短 draft 层可能出现启动开销大于 overlap 收益,需要 benchmark/trace 验证。
|
||||
- 任何 fallback 必须 warning/fail-fast,不能静默降级。
|
||||
|
||||
@@ -1028,3 +1028,22 @@ Fallback:
|
||||
debug enabled 时自动回到现有同步 materialize
|
||||
pynccl unavailable 时自动回到现有同步 materialize
|
||||
```
|
||||
|
||||
## 2026-06-01:attention-front prefetch OOM observation
|
||||
|
||||
Remote log:
|
||||
|
||||
- `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260531_161740.log`
|
||||
- Failure time: `2026-05-31 16:21:21`
|
||||
- Stack: target prefill `eagle_worker.forward_target_extend()` -> `DeepseekV2AttentionMLA.forward_prepare()` -> `Indexer._get_topk_in_seq_cp_pair()` -> `_get_topk_ragged_with_cp()` -> `deep_gemm.fp8_mqa_logits()`.
|
||||
- Error: CUDA OOM while allocating `7.58 GiB`; GPU0 had `6.67 GiB` free, process used `132.84 GiB`.
|
||||
- Immediately preceding successful batch included very large requests, e.g. `#new-token: 140672`, `#cached-token: 13824`, `#inflight-req: 2`; next failed request prepared backup for `logical_len=178945`.
|
||||
|
||||
Confirmed root cause direction:
|
||||
|
||||
- The OOM is not a host-cache capacity error; it is a CUDA activation/temp-buffer peak inside current-layer NSA index MQA/topk.
|
||||
- CP in-seq `_get_topk_ragged_with_cp()` calls `deep_gemm.fp8_mqa_logits()` without the q-dimension chunking guard that the non-CP ragged path already has.
|
||||
- For `logical_len=178945`, CP0's zigzag tail segment is about `11184` query rows against about `178945` KV rows, so the fp32 logits allocation is about `7.45 GiB`, matching the `7.58 GiB` OOM.
|
||||
- Do not solve this by limiting max prefetch size. The follow-up fix should add chunked MQA logits handling to the CP in-seq path.
|
||||
|
||||
Do not re-debug from scratch before checking this section.
|
||||
|
||||
Reference in New Issue
Block a user