diff --git a/docs/superpowers/plans/2026-05-24-cp-draft-nsa-state-transfer.md b/docs/superpowers/plans/2026-05-24-cp-draft-nsa-state-transfer.md new file mode 100644 index 000000000..3c84517e6 --- /dev/null +++ b/docs/superpowers/plans/2026-05-24-cp-draft-nsa-state-transfer.md @@ -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 +``` diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index a3349a9a3..09bcf81ac 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -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) diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 3d3a90021..0875704c0 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -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: diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 809f596f1..d3dc71721 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -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,