sglang/docs/superpowers/plans/2026-05-24-cp-draft-nsa-state-transfer.md

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:

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:

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:

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:

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:

_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:

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:

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:

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:

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:

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:

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:

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:

export SGLANG_CP_DRAFT_SHARED_KV_DEBUG=1

Expected at startup:

register_buffers ... total_kv_bufs=79 ... state_bufs=79 ... draft_state_bufs=1

Expected during transfer:

maybe_send_extra_state ... src_state_bufs=79 dst_state_bufs=79 ... target_registration_state_bufs=79

Unexpected and blocking:

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:

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:

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:

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