# 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
```
