Preserve draft NSA state during CP disaggregated transfer
CP shared KV already registered the draft model main KV buffer with the prefill/decode Mooncake managers, but NSA draft state buffers were not part of the state registration set. HiCache/cache-hit traffic could then transfer pages from the draft pool without transferring the matching draft index/scale state, which is a plausible cause of the EAGLE/MTP accept-length collapse after cache hits. This appends compatible draft NSA state buffers to the existing state transfer registration on both prefill and decode, and extends transfer-side diagnostics so source/destination state-buffer counts are visible. The mismatch guard degrades to the common prefix of registered state buffers instead of crashing if a rolling deployment exposes asymmetric registration. Constraint: Scope is intentionally limited to target_state_type=nsa and draft_state_type=nsa. Rejected: Treat draft main KV transfer as sufficient | NSA attention also needs draft index/scale state for transferred pages. Rejected: Add Mamba/SWA draft-state semantics now | those state layouts need separate correctness analysis. Confidence: medium Scope-risk: moderate Directive: Do not remove the draft_state_buffer_start/count fields without checking Mooncake source/destination registration symmetry. Tested: PYTHONDONTWRITEBYTECODE=1 python3 -m py_compile python/sglang/srt/disaggregation/prefill.py python/sglang/srt/disaggregation/decode.py python/sglang/srt/disaggregation/mooncake/conn.py Tested: git diff --check Tested: Remote prefill log showed registered_state_bufs=79 and maybe_send_extra_state src_state_bufs=79 dst_state_bufs=79 with no state-buffer mismatch. Not-tested: Full accept-length recovery; latest remote run hit an unrelated prefill KV allocator idle-check leak after transfer registration succeeded.
This commit is contained in:
laoyao0822
379
docs/superpowers/plans/2026-05-24-cp-draft-nsa-state-transfer.md
Normal file
379
docs/superpowers/plans/2026-05-24-cp-draft-nsa-state-transfer.md
Normal file
@@ -0,0 +1,379 @@
|
||||
# CP Draft NSA State Transfer Implementation Plan
|
||||
|
||||
> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
|
||||
|
||||
**Goal:** Restore MTP/EAGLE accept length under CP shared KV + HiCache cache-hit by transferring draft model NSA index/scale state along with target model state.
|
||||
|
||||
**Architecture:** Target and draft main KV already share token/page indices and are appended into `kv_args.kv_data_*`. The missing piece is draft NSA state (`index_k_with_scale_buffer`) registration and PD transfer: append compatible draft state buffers into `kv_args.state_*` on both prefill and decode, then let the existing generic state transfer path copy the same page indices into all registered state buffers.
|
||||
|
||||
**Tech Stack:** Python, SGLang disaggregation prefill/decode, Mooncake KV transfer, CP shared KV, NSA KV pool, MTP/EAGLE draft model.
|
||||
|
||||
---
|
||||
|
||||
## Root-cause evidence
|
||||
|
||||
Remote log `/mnt/beegfs/qjs/sglang_cp_hicache_20260524_125458.log` shows:
|
||||
|
||||
- `prefill_state_manager`: target NSA state has `target_state_bufs=78`; draft NSA state exists with `draft_state_bufs=1`; registered state remains `registered_state_bufs=78`.
|
||||
- `prefill_draft_state_not_registered`: explicitly reports `draft_state_buffers_are_not_in_pd_state_transfer`.
|
||||
- `register_buffers`: main KV registers `total_kv_bufs=79` with `draft_count=1`, but state registers only `state_bufs=78`.
|
||||
- `maybe_send_extra_state`: transfers only `src_state_bufs=78`, `dst_state_bufs=78`, while `draft_state_bufs=1` exists.
|
||||
- `prefill_send_cachehit_draft_prefix`: cache-hit prefix is sent from draft pool for decode, so stale/missing draft NSA state can directly lower accept length.
|
||||
|
||||
Current code confirms this:
|
||||
|
||||
- `python/sglang/srt/disaggregation/prefill.py:227-236` appends draft main KV to `kv_args.kv_data_*`.
|
||||
- `python/sglang/srt/disaggregation/prefill.py:272-297` registers only target `token_to_kv_pool.get_state_buf_infos()` into `kv_args.state_*`.
|
||||
- `python/sglang/srt/disaggregation/decode.py:392-401` appends draft main KV to `kv_args.kv_data_*`.
|
||||
- `python/sglang/srt/disaggregation/decode.py:430-455` registers only target state into `kv_args.state_*`.
|
||||
- `python/sglang/srt/mem_cache/memory_pool.py:1955-1965` shows `NSATokenToKVPool.get_state_buf_infos()` returns per-layer `index_k_with_scale_buffer` pointers/lens/item lens.
|
||||
- `python/sglang/srt/disaggregation/mooncake/conn.py:729-736` uses `kv_args.state_data_ptrs`, destination state ptrs, and state page indices generically; once draft state is appended on both sides, the existing path can transfer it.
|
||||
|
||||
## Scope
|
||||
|
||||
In scope:
|
||||
|
||||
1. NSA draft state only (`state_type == "nsa"` and `draft_state_type == "nsa"`).
|
||||
2. Prefill and decode `KVArgs` state registration.
|
||||
3. Mooncake verification logs and mismatch warnings.
|
||||
4. Remote validation on `g0034` container `sglang-glm5-dev-2`.
|
||||
|
||||
Out of scope for this patch:
|
||||
|
||||
1. Mamba/SWA draft state semantics.
|
||||
2. New HiCache host storage format for draft state.
|
||||
3. CP TBO behavior changes.
|
||||
4. Performance tuning after correctness is restored.
|
||||
|
||||
## Files
|
||||
|
||||
- Modify: `python/sglang/srt/disaggregation/prefill.py`
|
||||
- Add reusable state-info helper near `_state_buf_debug_summary`.
|
||||
- Append draft NSA state into `kv_args.state_*` in `PrefillBootstrapQueue._init_kv_manager`.
|
||||
- Change diagnostic log from `prefill_draft_state_not_registered` to `prefill_draft_state_registered` or `prefill_draft_state_skipped`.
|
||||
- Modify: `python/sglang/srt/disaggregation/decode.py`
|
||||
- Mirror the prefill state append in `DecodePreallocQueue._init_kv_manager`.
|
||||
- Change diagnostic log from `decode_draft_state_not_registered` to `decode_draft_state_registered` or `decode_draft_state_skipped`.
|
||||
- Modify: `python/sglang/srt/disaggregation/mooncake/conn.py`
|
||||
- Keep existing debug logs.
|
||||
- Add a guarded warning in `maybe_send_extra` when source and destination state buffer counts differ.
|
||||
- Test/verify: remote compile and E2E logs on `g0034`.
|
||||
|
||||
---
|
||||
|
||||
### Task 1: Add deterministic state buffer helper
|
||||
|
||||
**Files:**
|
||||
- Modify: `python/sglang/srt/disaggregation/prefill.py`
|
||||
- Modify: `python/sglang/srt/disaggregation/decode.py`
|
||||
|
||||
- [ ] **Step 1: Add helper returning actual state buffers**
|
||||
|
||||
Add this next to `_state_buf_debug_summary` in both files:
|
||||
|
||||
```python
|
||||
def _state_buf_infos(pool):
|
||||
state_type = _pool_state_type(pool)
|
||||
if state_type == "none":
|
||||
return state_type, [], [], []
|
||||
state_data_ptrs, state_data_lens, state_item_lens = pool.get_state_buf_infos()
|
||||
return state_type, state_data_ptrs, state_data_lens, state_item_lens
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Update `_state_buf_debug_summary` to use helper**
|
||||
|
||||
Change the successful path to:
|
||||
|
||||
```python
|
||||
try:
|
||||
state_type, state_data_ptrs, state_data_lens, state_item_lens = _state_buf_infos(pool)
|
||||
except Exception as exc: # pragma: no cover - diagnostics must not fail setup.
|
||||
return state_type, -1, f"error={type(exc).__name__}: {exc}", "error"
|
||||
return (
|
||||
state_type,
|
||||
len(state_data_ptrs),
|
||||
_seq_summary(state_data_lens),
|
||||
_seq_summary(state_item_lens),
|
||||
)
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Compile check**
|
||||
|
||||
Run:
|
||||
|
||||
```bash
|
||||
python3 -m py_compile python/sglang/srt/disaggregation/prefill.py python/sglang/srt/disaggregation/decode.py
|
||||
```
|
||||
|
||||
Expected: command exits `0`.
|
||||
|
||||
### Task 2: Append draft NSA state in prefill KVArgs
|
||||
|
||||
**Files:**
|
||||
- Modify: `python/sglang/srt/disaggregation/prefill.py`
|
||||
|
||||
- [ ] **Step 1: Insert draft state append after target state initialization**
|
||||
|
||||
After current `kv_args.state_type` setup in `PrefillBootstrapQueue._init_kv_manager`, add:
|
||||
|
||||
```python
|
||||
draft_state_type = "none"
|
||||
draft_state_data_ptrs = []
|
||||
draft_state_data_lens = []
|
||||
draft_state_item_lens = []
|
||||
if self.draft_token_to_kv_pool is not None:
|
||||
(
|
||||
draft_state_type,
|
||||
draft_state_data_ptrs,
|
||||
draft_state_data_lens,
|
||||
draft_state_item_lens,
|
||||
) = _state_buf_infos(self.draft_token_to_kv_pool)
|
||||
|
||||
kv_args.draft_state_type = draft_state_type
|
||||
kv_args.draft_state_buffer_start = len(kv_args.state_data_ptrs)
|
||||
kv_args.draft_state_buffer_count = 0
|
||||
|
||||
if draft_state_data_ptrs:
|
||||
if kv_args.state_type == "nsa" and draft_state_type == "nsa":
|
||||
kv_args.state_data_ptrs += draft_state_data_ptrs
|
||||
kv_args.state_data_lens += draft_state_data_lens
|
||||
kv_args.state_item_lens += draft_state_item_lens
|
||||
kv_args.draft_state_buffer_count = len(draft_state_data_ptrs)
|
||||
else:
|
||||
_cp_draft_shared_kv_debug(
|
||||
"prefill_draft_state_skipped cp_rank=%s target_state_type=%s "
|
||||
"draft_state_type=%s draft_state_bufs=%s reason=unsupported_state_type",
|
||||
self.tp_rank,
|
||||
kv_args.state_type,
|
||||
draft_state_type,
|
||||
len(draft_state_data_ptrs),
|
||||
)
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Replace stale debug block**
|
||||
|
||||
Remove the `prefill_draft_state_not_registered` log. The debug block should report:
|
||||
|
||||
```python
|
||||
_cp_draft_shared_kv_debug(
|
||||
"prefill_state_manager cp_rank=%s target_state_type=%s "
|
||||
"draft_state_type=%s draft_state_bufs=%s draft_state_start=%s "
|
||||
"registered_state_bufs=%s registered_state_lens=%s registered_state_item_lens=%s",
|
||||
self.tp_rank,
|
||||
kv_args.state_type,
|
||||
kv_args.draft_state_type,
|
||||
kv_args.draft_state_buffer_count,
|
||||
kv_args.draft_state_buffer_start,
|
||||
len(kv_args.state_data_ptrs),
|
||||
_seq_summary(kv_args.state_data_lens),
|
||||
_seq_summary(kv_args.state_item_lens),
|
||||
)
|
||||
```
|
||||
|
||||
Expected runtime after fix: `registered_state_bufs=79`, `draft_state_bufs=1`, `draft_state_start=78`.
|
||||
|
||||
- [ ] **Step 3: Compile check**
|
||||
|
||||
Run:
|
||||
|
||||
```bash
|
||||
python3 -m py_compile python/sglang/srt/disaggregation/prefill.py
|
||||
```
|
||||
|
||||
Expected: command exits `0`.
|
||||
|
||||
### Task 3: Append draft NSA state in decode KVArgs
|
||||
|
||||
**Files:**
|
||||
- Modify: `python/sglang/srt/disaggregation/decode.py`
|
||||
|
||||
- [ ] **Step 1: Mirror Task 2 in `DecodePreallocQueue._init_kv_manager`**
|
||||
|
||||
After target state initialization, add the same draft-state append logic, replacing `self.tp_rank` logging context and `decode_draft_state_*` labels:
|
||||
|
||||
```python
|
||||
draft_state_type = "none"
|
||||
draft_state_data_ptrs = []
|
||||
draft_state_data_lens = []
|
||||
draft_state_item_lens = []
|
||||
if self.draft_token_to_kv_pool is not None:
|
||||
(
|
||||
draft_state_type,
|
||||
draft_state_data_ptrs,
|
||||
draft_state_data_lens,
|
||||
draft_state_item_lens,
|
||||
) = _state_buf_infos(self.draft_token_to_kv_pool)
|
||||
|
||||
kv_args.draft_state_type = draft_state_type
|
||||
kv_args.draft_state_buffer_start = len(kv_args.state_data_ptrs)
|
||||
kv_args.draft_state_buffer_count = 0
|
||||
|
||||
if draft_state_data_ptrs:
|
||||
if kv_args.state_type == "nsa" and draft_state_type == "nsa":
|
||||
kv_args.state_data_ptrs += draft_state_data_ptrs
|
||||
kv_args.state_data_lens += draft_state_data_lens
|
||||
kv_args.state_item_lens += draft_state_item_lens
|
||||
kv_args.draft_state_buffer_count = len(draft_state_data_ptrs)
|
||||
else:
|
||||
_cp_draft_shared_kv_debug(
|
||||
"decode_draft_state_skipped cp_rank=%s target_state_type=%s "
|
||||
"draft_state_type=%s draft_state_bufs=%s reason=unsupported_state_type",
|
||||
self.tp_rank,
|
||||
kv_args.state_type,
|
||||
draft_state_type,
|
||||
len(draft_state_data_ptrs),
|
||||
)
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Ensure decode registration sends appended state pointers**
|
||||
|
||||
No Mooncake receiver packing change should be needed because `MooncakeKVReceiver._register_kv_args` already packs `self.kv_mgr.kv_args.state_data_ptrs` and `state_item_lens`. After Task 3, decode registration should log `state_bufs=79`.
|
||||
|
||||
- [ ] **Step 3: Compile check**
|
||||
|
||||
Run:
|
||||
|
||||
```bash
|
||||
python3 -m py_compile python/sglang/srt/disaggregation/decode.py
|
||||
```
|
||||
|
||||
Expected: command exits `0`.
|
||||
|
||||
### Task 4: Add state-count mismatch guard in Mooncake
|
||||
|
||||
**Files:**
|
||||
- Modify: `python/sglang/srt/disaggregation/mooncake/conn.py`
|
||||
|
||||
- [ ] **Step 1: Add warning before NSA/SWA generic transfer**
|
||||
|
||||
Inside `maybe_send_extra`, before `_send_kvcache_generic`, add:
|
||||
|
||||
```python
|
||||
if len(self.kv_args.state_data_ptrs) != len(dst_state_data_ptrs):
|
||||
logger.warning(
|
||||
"State buffer count mismatch during PD transfer: src=%s dst=%s "
|
||||
"state_type=%s draft_state_type=%s draft_state_bufs=%s room=%s session=%s",
|
||||
len(self.kv_args.state_data_ptrs),
|
||||
len(dst_state_data_ptrs),
|
||||
state_type,
|
||||
getattr(self.kv_args, "draft_state_type", None),
|
||||
getattr(self.kv_args, "draft_state_buffer_count", None),
|
||||
req.room,
|
||||
req.mooncake_session_id,
|
||||
)
|
||||
```
|
||||
|
||||
Do not raise yet; raising would convert accept-rate degradation into request failure during validation if only one side is stale.
|
||||
|
||||
- [ ] **Step 2: Compile check**
|
||||
|
||||
Run:
|
||||
|
||||
```bash
|
||||
python3 -m py_compile python/sglang/srt/disaggregation/mooncake/conn.py
|
||||
```
|
||||
|
||||
Expected: command exits `0`.
|
||||
|
||||
### Task 5: Remote sync and verification
|
||||
|
||||
**Files:**
|
||||
- Sync modified local files to remote container path `/sgl-workspace/sglang-tai` on `g0034`.
|
||||
|
||||
- [ ] **Step 1: Sync exact changed files first**
|
||||
|
||||
Run from local repo:
|
||||
|
||||
```bash
|
||||
scp python/sglang/srt/disaggregation/prefill.py g0034:/mnt/beegfs/cjy/sglang-dev/python/sglang/srt/disaggregation/prefill.py
|
||||
scp python/sglang/srt/disaggregation/decode.py g0034:/mnt/beegfs/cjy/sglang-dev/python/sglang/srt/disaggregation/decode.py
|
||||
scp python/sglang/srt/disaggregation/mooncake/conn.py g0034:/mnt/beegfs/cjy/sglang-dev/python/sglang/srt/disaggregation/mooncake/conn.py
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Remote compile inside container**
|
||||
|
||||
Run:
|
||||
|
||||
```bash
|
||||
ssh g0034 "docker exec sglang-glm5-dev-2 bash -lc 'cd /sgl-workspace/sglang-tai && python3 -m py_compile python/sglang/srt/disaggregation/prefill.py python/sglang/srt/disaggregation/decode.py python/sglang/srt/disaggregation/mooncake/conn.py'"
|
||||
```
|
||||
|
||||
Expected: command exits `0`.
|
||||
|
||||
- [ ] **Step 3: Restart workload with diagnostics**
|
||||
|
||||
Use the current launch path with:
|
||||
|
||||
```bash
|
||||
export SGLANG_CP_DRAFT_SHARED_KV_DEBUG=1
|
||||
```
|
||||
|
||||
Expected at startup:
|
||||
|
||||
```text
|
||||
register_buffers ... total_kv_bufs=79 ... state_bufs=79 ... draft_state_bufs=1
|
||||
```
|
||||
|
||||
Expected during transfer:
|
||||
|
||||
```text
|
||||
maybe_send_extra_state ... src_state_bufs=79 dst_state_bufs=79 ... target_registration_state_bufs=79
|
||||
```
|
||||
|
||||
Unexpected and blocking:
|
||||
|
||||
```text
|
||||
prefill_draft_state_not_registered
|
||||
State buffer count mismatch during PD transfer: src=79 dst=78
|
||||
State buffer count mismatch during PD transfer: src=78 dst=79
|
||||
```
|
||||
|
||||
- [ ] **Step 4: Validate accept length**
|
||||
|
||||
Run the same HiCache+MTP cache-hit workload. Monitor decode logs:
|
||||
|
||||
```bash
|
||||
ssh g0034 "grep -h 'accept len' /mnt/beegfs/qjs/sglang_cp_hicache_*.log | tail -n 80"
|
||||
```
|
||||
|
||||
Pass condition: after cache-hit requests, `accept len` does not collapse persistently to about `1.00` / `accept rate=0.25`; it returns to the pre-regression range for this workload.
|
||||
|
||||
### Task 6: Clean diagnostics and commit
|
||||
|
||||
**Files:**
|
||||
- Modify: same files as above.
|
||||
|
||||
- [ ] **Step 1: Keep only useful gated logs**
|
||||
|
||||
Keep logs under `SGLANG_CP_DRAFT_SHARED_KV_DEBUG`. Remove false-positive `*_not_registered` logs once the registration path is fixed.
|
||||
|
||||
- [ ] **Step 2: Git status check**
|
||||
|
||||
Run:
|
||||
|
||||
```bash
|
||||
git status --short
|
||||
```
|
||||
|
||||
Expected: only intended source files and this plan file are modified/added.
|
||||
|
||||
- [ ] **Step 3: Commit with Lore protocol after remote validation**
|
||||
|
||||
Use a Lore-style commit message:
|
||||
|
||||
```text
|
||||
Transfer draft NSA state with CP shared KV disaggregation
|
||||
|
||||
MTP accept length collapsed after HiCache cache-hit because prefill transferred
|
||||
draft main KV pages but not the draft NSA index/scale state buffer. Target and
|
||||
draft main KV share token/page indices, so the state transfer path can register
|
||||
and copy the draft NSA state buffer alongside target per-layer state buffers.
|
||||
|
||||
Constraint: First fix is scoped to NSA draft state used by GLM5 MTP/EAGLE.
|
||||
Rejected: Add draft state to HiCache host format first | PD transfer currently fails before decode can receive correct draft state.
|
||||
Confidence: medium
|
||||
Scope-risk: moderate
|
||||
Directive: Do not extend this to SWA/Mamba draft state without proving their state indices match target page indices.
|
||||
Tested: py_compile and remote HiCache+MTP E2E accept-length validation on g0034
|
||||
Not-tested: NIXL/MORI backends and non-NSA draft state
|
||||
```
|
||||
@@ -113,6 +113,26 @@ def _pool_summary(pool) -> str:
|
||||
return " ".join(parts)
|
||||
|
||||
|
||||
def _pool_state_type(pool) -> str:
|
||||
if pool is None or not hasattr(pool, "get_state_buf_infos"):
|
||||
return "none"
|
||||
if isinstance(pool, SWAKVPool):
|
||||
return "swa"
|
||||
if isinstance(pool, HybridLinearKVPool):
|
||||
return "mamba"
|
||||
if isinstance(pool, NSATokenToKVPool):
|
||||
return "nsa"
|
||||
return "unknown"
|
||||
|
||||
|
||||
def _state_buf_infos(pool):
|
||||
state_type = _pool_state_type(pool)
|
||||
if state_type == "none":
|
||||
return state_type, [], [], []
|
||||
state_data_ptrs, state_data_lens, state_item_lens = pool.get_state_buf_infos()
|
||||
return state_type, state_data_ptrs, state_data_lens, state_item_lens
|
||||
|
||||
|
||||
def _kv_locs_to_page_indices_cpu(
|
||||
kv_locs: torch.Tensor,
|
||||
page_size: int,
|
||||
@@ -426,6 +446,68 @@ class DecodePreallocQueue:
|
||||
kv_args.state_item_lens = []
|
||||
kv_args.state_type = "none"
|
||||
|
||||
draft_state_type = "none"
|
||||
draft_state_data_ptrs = []
|
||||
draft_state_data_lens = []
|
||||
draft_state_item_lens = []
|
||||
if self.draft_token_to_kv_pool is not None:
|
||||
(
|
||||
draft_state_type,
|
||||
draft_state_data_ptrs,
|
||||
draft_state_data_lens,
|
||||
draft_state_item_lens,
|
||||
) = _state_buf_infos(self.draft_token_to_kv_pool)
|
||||
|
||||
kv_args.draft_state_type = draft_state_type
|
||||
kv_args.draft_state_buffer_start = len(kv_args.state_data_ptrs)
|
||||
kv_args.draft_state_buffer_count = 0
|
||||
if draft_state_data_ptrs:
|
||||
if kv_args.state_type == "nsa" and draft_state_type == "nsa":
|
||||
kv_args.state_data_ptrs += draft_state_data_ptrs
|
||||
kv_args.state_data_lens += draft_state_data_lens
|
||||
kv_args.state_item_lens += draft_state_item_lens
|
||||
kv_args.draft_state_buffer_count = len(draft_state_data_ptrs)
|
||||
else:
|
||||
_cp_draft_shared_kv_debug(
|
||||
"decode_draft_state_skipped cp_rank=%s target_state_type=%s "
|
||||
"draft_state_type=%s draft_state_bufs=%s "
|
||||
"reason=unsupported_state_type",
|
||||
self.tp_rank,
|
||||
kv_args.state_type,
|
||||
draft_state_type,
|
||||
len(draft_state_data_ptrs),
|
||||
)
|
||||
|
||||
if envs.SGLANG_CP_DRAFT_SHARED_KV_DEBUG.get():
|
||||
_cp_draft_shared_kv_debug(
|
||||
"decode_state_manager cp_rank=%s target_state_type=%s "
|
||||
"draft_state_type=%s draft_state_bufs=%s draft_state_lens=%s "
|
||||
"draft_state_item_lens=%s draft_state_start=%s registered_state_bufs=%s "
|
||||
"registered_state_lens=%s registered_state_item_lens=%s",
|
||||
self.tp_rank,
|
||||
kv_args.state_type,
|
||||
kv_args.draft_state_type,
|
||||
kv_args.draft_state_buffer_count,
|
||||
_seq_summary(draft_state_data_lens),
|
||||
_seq_summary(draft_state_item_lens),
|
||||
kv_args.draft_state_buffer_start,
|
||||
len(kv_args.state_data_ptrs),
|
||||
_seq_summary(kv_args.state_data_lens),
|
||||
_seq_summary(kv_args.state_item_lens),
|
||||
)
|
||||
if kv_args.draft_state_buffer_count > 0:
|
||||
_cp_draft_shared_kv_debug(
|
||||
"decode_draft_state_registered cp_rank=%s "
|
||||
"draft_state_type=%s draft_state_bufs=%s draft_state_start=%s "
|
||||
"registered_state_type=%s registered_state_bufs=%s",
|
||||
self.tp_rank,
|
||||
kv_args.draft_state_type,
|
||||
kv_args.draft_state_buffer_count,
|
||||
kv_args.draft_state_buffer_start,
|
||||
kv_args.state_type,
|
||||
len(kv_args.state_data_ptrs),
|
||||
)
|
||||
|
||||
kv_args.ib_device = self.scheduler.server_args.disaggregation_ib_device
|
||||
kv_args.gpu_id = self.scheduler.gpu_id
|
||||
kv_manager_class = get_kv_class(self.transfer_backend, KVClassType.MANAGER)
|
||||
|
||||
@@ -259,7 +259,9 @@ class MooncakeKVManager(CommonKVManager):
|
||||
if self.kv_args.kv_data_ptrs and self.kv_args.kv_data_lens:
|
||||
_cp_draft_shared_kv_debug(
|
||||
"register_buffers mode=%s cp_rank=%s total_kv_bufs=%s "
|
||||
"draft_start=%s draft_count=%s kv_lens=%s kv_item_lens=%s",
|
||||
"draft_start=%s draft_count=%s kv_lens=%s kv_item_lens=%s "
|
||||
"state_type=%s state_bufs=%s state_lens=%s state_item_lens=%s "
|
||||
"draft_state_type=%s draft_state_bufs=%s",
|
||||
self.disaggregation_mode,
|
||||
self.attn_cp_rank,
|
||||
len(self.kv_args.kv_data_ptrs),
|
||||
@@ -267,6 +269,12 @@ class MooncakeKVManager(CommonKVManager):
|
||||
getattr(self.kv_args, "draft_kv_buffer_count", None),
|
||||
_np_summary(self.kv_args.kv_data_lens),
|
||||
_np_summary(self.kv_args.kv_item_lens),
|
||||
getattr(self.kv_args, "state_type", None),
|
||||
len(getattr(self.kv_args, "state_data_ptrs", []) or []),
|
||||
_np_summary(getattr(self.kv_args, "state_data_lens", [])),
|
||||
_np_summary(getattr(self.kv_args, "state_item_lens", [])),
|
||||
getattr(self.kv_args, "draft_state_type", None),
|
||||
getattr(self.kv_args, "draft_state_buffer_count", None),
|
||||
)
|
||||
self.engine.batch_register(
|
||||
self.kv_args.kv_data_ptrs, self.kv_args.kv_data_lens
|
||||
@@ -645,6 +653,33 @@ class MooncakeKVManager(CommonKVManager):
|
||||
):
|
||||
"""Send state or extra pool data with type-specific handling."""
|
||||
state_type = getattr(self.kv_args, "state_type", "none")
|
||||
_cp_draft_shared_kv_debug(
|
||||
"maybe_send_extra_state cp_rank=%s room=%s session=%s state_type=%s "
|
||||
"src_state_bufs=%s dst_state_bufs=%s src_state_lens=%s src_state_item_lens=%s "
|
||||
"prefill_state_indices=%s dst_state_indices=%s draft_state_type=%s "
|
||||
"draft_state_bufs=%s target_registration_state_bufs=%s",
|
||||
self.attn_cp_rank,
|
||||
req.room,
|
||||
req.mooncake_session_id,
|
||||
state_type,
|
||||
len(getattr(self.kv_args, "state_data_ptrs", []) or []),
|
||||
len(dst_state_data_ptrs or []),
|
||||
_np_summary(getattr(self.kv_args, "state_data_lens", [])),
|
||||
_np_summary(getattr(self.kv_args, "state_item_lens", [])),
|
||||
_np_summary(prefill_state_indices),
|
||||
_np_summary(
|
||||
dst_state_indices
|
||||
if dst_state_indices is not None
|
||||
else getattr(req, "dst_state_indices", [])
|
||||
),
|
||||
getattr(self.kv_args, "draft_state_type", None),
|
||||
getattr(self.kv_args, "draft_state_buffer_count", None),
|
||||
(
|
||||
len(target_rank_registration_info.dst_state_data_ptrs)
|
||||
if target_rank_registration_info is not None
|
||||
else None
|
||||
),
|
||||
)
|
||||
|
||||
if state_type == "mamba":
|
||||
# Check if we need slice transfer for different TP sizes
|
||||
@@ -689,13 +724,34 @@ class MooncakeKVManager(CommonKVManager):
|
||||
prefill_state_indices = prefill_state_indices[
|
||||
: len(effective_dst_state_indices)
|
||||
]
|
||||
src_state_data_ptrs = self.kv_args.state_data_ptrs
|
||||
dst_state_ptrs = dst_state_data_ptrs
|
||||
state_item_lens = self.kv_args.state_item_lens
|
||||
if len(src_state_data_ptrs) != len(dst_state_ptrs):
|
||||
transfer_buf_count = min(len(src_state_data_ptrs), len(dst_state_ptrs))
|
||||
logger.warning(
|
||||
"State buffer count mismatch during PD transfer: src=%s dst=%s "
|
||||
"state_type=%s draft_state_type=%s draft_state_bufs=%s room=%s "
|
||||
"session=%s; transferring first %s buffers only",
|
||||
len(src_state_data_ptrs),
|
||||
len(dst_state_ptrs),
|
||||
state_type,
|
||||
getattr(self.kv_args, "draft_state_type", None),
|
||||
getattr(self.kv_args, "draft_state_buffer_count", None),
|
||||
req.room,
|
||||
req.mooncake_session_id,
|
||||
transfer_buf_count,
|
||||
)
|
||||
src_state_data_ptrs = src_state_data_ptrs[:transfer_buf_count]
|
||||
dst_state_ptrs = dst_state_ptrs[:transfer_buf_count]
|
||||
state_item_lens = state_item_lens[:transfer_buf_count]
|
||||
# Reuse _send_kvcache_generic interface to send extra pool data
|
||||
prefill_state_indices = np.array(prefill_state_indices, dtype=np.int32)
|
||||
return self._send_kvcache_generic(
|
||||
mooncake_session_id=req.mooncake_session_id,
|
||||
src_data_ptrs=self.kv_args.state_data_ptrs,
|
||||
dst_data_ptrs=dst_state_data_ptrs,
|
||||
item_lens=self.kv_args.state_item_lens,
|
||||
src_data_ptrs=src_state_data_ptrs,
|
||||
dst_data_ptrs=dst_state_ptrs,
|
||||
item_lens=state_item_lens,
|
||||
prefill_data_indices=prefill_state_indices,
|
||||
dst_data_indices=effective_dst_state_indices,
|
||||
executor=executor,
|
||||
@@ -1454,6 +1510,24 @@ class MooncakeKVReceiver(CommonKVReceiver):
|
||||
dst_tp_rank = str(tp_rank).encode("ascii")
|
||||
dst_attn_tp_size = str(self.kv_mgr.attn_tp_size).encode("ascii")
|
||||
dst_kv_item_len = str(kv_item_len).encode("ascii")
|
||||
_cp_draft_shared_kv_debug(
|
||||
"decode_register_kv_args cp_rank=%s room=%s session=%s "
|
||||
"kv_bufs=%s aux_bufs=%s state_type=%s state_bufs=%s "
|
||||
"state_item_lens=%s draft_start=%s draft_count=%s "
|
||||
"draft_state_type=%s draft_state_bufs=%s",
|
||||
self.kv_mgr.attn_cp_rank,
|
||||
self.bootstrap_room,
|
||||
self.session_id,
|
||||
len(self.kv_mgr.kv_args.kv_data_ptrs),
|
||||
len(self.kv_mgr.kv_args.aux_data_ptrs),
|
||||
getattr(self.kv_mgr.kv_args, "state_type", None),
|
||||
len(self.kv_mgr.kv_args.state_data_ptrs),
|
||||
_np_summary(self.kv_mgr.kv_args.state_item_lens),
|
||||
getattr(self.kv_mgr.kv_args, "draft_kv_buffer_start", None),
|
||||
getattr(self.kv_mgr.kv_args, "draft_kv_buffer_count", None),
|
||||
getattr(self.kv_mgr.kv_args, "draft_state_type", None),
|
||||
getattr(self.kv_mgr.kv_args, "draft_state_buffer_count", None),
|
||||
)
|
||||
|
||||
sock, lock = self._connect_to_bootstrap_server(bootstrap_info)
|
||||
with lock:
|
||||
|
||||
@@ -98,6 +98,26 @@ def _pool_summary(pool) -> str:
|
||||
return " ".join(parts)
|
||||
|
||||
|
||||
def _pool_state_type(pool) -> str:
|
||||
if pool is None or not hasattr(pool, "get_state_buf_infos"):
|
||||
return "none"
|
||||
if isinstance(pool, SWAKVPool):
|
||||
return "swa"
|
||||
if isinstance(pool, HybridLinearKVPool):
|
||||
return "mamba"
|
||||
if isinstance(pool, NSATokenToKVPool):
|
||||
return "nsa"
|
||||
return "unknown"
|
||||
|
||||
|
||||
def _state_buf_infos(pool):
|
||||
state_type = _pool_state_type(pool)
|
||||
if state_type == "none":
|
||||
return state_type, [], [], []
|
||||
state_data_ptrs, state_data_lens, state_item_lens = pool.get_state_buf_infos()
|
||||
return state_type, state_data_ptrs, state_data_lens, state_item_lens
|
||||
|
||||
|
||||
def _kv_locs_to_page_indices_cpu(
|
||||
kv_locs: torch.Tensor,
|
||||
page_size: int,
|
||||
@@ -267,6 +287,68 @@ class PrefillBootstrapQueue:
|
||||
kv_args.state_item_lens = []
|
||||
kv_args.state_type = "none"
|
||||
|
||||
draft_state_type = "none"
|
||||
draft_state_data_ptrs = []
|
||||
draft_state_data_lens = []
|
||||
draft_state_item_lens = []
|
||||
if self.draft_token_to_kv_pool is not None:
|
||||
(
|
||||
draft_state_type,
|
||||
draft_state_data_ptrs,
|
||||
draft_state_data_lens,
|
||||
draft_state_item_lens,
|
||||
) = _state_buf_infos(self.draft_token_to_kv_pool)
|
||||
|
||||
kv_args.draft_state_type = draft_state_type
|
||||
kv_args.draft_state_buffer_start = len(kv_args.state_data_ptrs)
|
||||
kv_args.draft_state_buffer_count = 0
|
||||
if draft_state_data_ptrs:
|
||||
if kv_args.state_type == "nsa" and draft_state_type == "nsa":
|
||||
kv_args.state_data_ptrs += draft_state_data_ptrs
|
||||
kv_args.state_data_lens += draft_state_data_lens
|
||||
kv_args.state_item_lens += draft_state_item_lens
|
||||
kv_args.draft_state_buffer_count = len(draft_state_data_ptrs)
|
||||
else:
|
||||
_cp_draft_shared_kv_debug(
|
||||
"prefill_draft_state_skipped cp_rank=%s target_state_type=%s "
|
||||
"draft_state_type=%s draft_state_bufs=%s "
|
||||
"reason=unsupported_state_type",
|
||||
self.tp_rank,
|
||||
kv_args.state_type,
|
||||
draft_state_type,
|
||||
len(draft_state_data_ptrs),
|
||||
)
|
||||
|
||||
if envs.SGLANG_CP_DRAFT_SHARED_KV_DEBUG.get():
|
||||
_cp_draft_shared_kv_debug(
|
||||
"prefill_state_manager cp_rank=%s target_state_type=%s "
|
||||
"draft_state_type=%s draft_state_bufs=%s draft_state_lens=%s "
|
||||
"draft_state_item_lens=%s draft_state_start=%s registered_state_bufs=%s "
|
||||
"registered_state_lens=%s registered_state_item_lens=%s",
|
||||
self.tp_rank,
|
||||
kv_args.state_type,
|
||||
kv_args.draft_state_type,
|
||||
kv_args.draft_state_buffer_count,
|
||||
_seq_summary(draft_state_data_lens),
|
||||
_seq_summary(draft_state_item_lens),
|
||||
kv_args.draft_state_buffer_start,
|
||||
len(kv_args.state_data_ptrs),
|
||||
_seq_summary(kv_args.state_data_lens),
|
||||
_seq_summary(kv_args.state_item_lens),
|
||||
)
|
||||
if kv_args.draft_state_buffer_count > 0:
|
||||
_cp_draft_shared_kv_debug(
|
||||
"prefill_draft_state_registered cp_rank=%s "
|
||||
"draft_state_type=%s draft_state_bufs=%s draft_state_start=%s "
|
||||
"registered_state_type=%s registered_state_bufs=%s",
|
||||
self.tp_rank,
|
||||
kv_args.draft_state_type,
|
||||
kv_args.draft_state_buffer_count,
|
||||
kv_args.draft_state_buffer_start,
|
||||
kv_args.state_type,
|
||||
len(kv_args.state_data_ptrs),
|
||||
)
|
||||
|
||||
kv_manager_class = get_kv_class(self.transfer_backend, KVClassType.MANAGER)
|
||||
kv_manager = kv_manager_class(
|
||||
kv_args,
|
||||
|
||||
Reference in New Issue
Block a user