chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Compact skipped NSA index-cache state safely

Browse Source 
Index skip reduces the number of target layers that own NSA index state,
but PD transfer and HiCache still assumed dense full-layer state buffers.
This change carries explicit state layer IDs through prefill/decode
registration, compacts device and host index buffers to active layers,
and maps logical layer IDs to compact slots on transfer paths.

The PD side fails fast when prefill/decode disagree on NSA state layer
identity instead of silently truncating or copying mismatched buffers.
Host direct tests now use the same CPU-index descriptor contract required
by the TAI cudaMemcpyBatchAsync path, and host registered memory is
unregistered on tensor finalization to avoid stale cudaHostRegister state
across CUDA tests.

Constraint: CP shared-KV with index_topk skip must keep target/draft state identity explicit before compacting buffers
Constraint: Direct HiCache TAI transfer rejects CUDA indices to avoid hidden D2H copies on the control path
Rejected: Keep full-layer L1/L2 index buffers | wastes the memory/bandwidth that index skip is meant to save
Rejected: Infer state buffer order by count only | can silently corrupt cache when active layer sets differ
Confidence: high
Scope-risk: moderate
Directive: Do not compact or reorder NSA state buffers without carrying logical layer IDs through PD registration and validating both sides
Tested: Remote container py_compile for touched runtime files
Tested: Remote container pytest: test_nsa_pool_host_unit.py, test_model_runner_kv_cache_mixin.py, test_cp_shared_kv_transfer_mapping.py, test_pd_state_layer_ids.py, test_cp_per_layer_transfer.py, test_cp_shared_kv_runtime.py -> 200 passed, 2 subtests passed
Not-tested: Full ETE GSM8K/replay after compacted P3-P6 changes
Co-authored-by: OmX <omx@oh-my-codex.dev>
This commit is contained in:
laoyao0822
2026-06-10 05:37:06 +08:00
co-authored by OmX
parent d21952b903
commit 1ebde44e59
14 changed files with 611 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/base/conn.py
+1
View File
@@ -31,6 +31,7 @@ class KVArgs:
state_data_ptrs: List[int]
state_data_lens: List[int]
state_item_lens: List[int]
state_layer_ids: List[int]
state_type: str # "none", "mamba", "swa"
# for mamba state different tp slice transfer
state_dim_per_tensor: List[int] # dimension to slice for each state tensor
python/sglang/srt/disaggregation/decode.py
+9
View File
@@ -136,6 +136,13 @@ def _state_buf_infos(pool):
return state_type, state_data_ptrs, state_data_lens, state_item_lens
def _state_layer_ids(pool):
get_state_layer_ids = getattr(pool, "get_state_layer_ids", None)
if get_state_layer_ids is None:
return []
return list(get_state_layer_ids())
def _kv_locs_to_page_indices_cpu(
kv_locs: torch.Tensor,
page_size: int,
@@ -429,6 +436,7 @@ class DecodePreallocQueue:
kv_args.state_data_ptrs = state_data_ptrs
kv_args.state_data_lens = state_data_lens
kv_args.state_item_lens = state_item_lens
kv_args.state_layer_ids = _state_layer_ids(self.token_to_kv_pool)
if isinstance(self.token_to_kv_pool, SWAKVPool):
kv_args.state_type = "swa"
@@ -447,6 +455,7 @@ class DecodePreallocQueue:
kv_args.state_data_ptrs = []
kv_args.state_data_lens = []
kv_args.state_item_lens = []
kv_args.state_layer_ids = []
kv_args.state_type = "none"
draft_state_type = "none"
python/sglang/srt/disaggregation/mooncake/conn.py
+35
View File
@@ -175,6 +175,7 @@ class KVArgsRegisterInfo:
# for mamba state different tp slice transfer
dst_state_item_lens: list[int]
dst_state_dim_per_tensor: list[int]
dst_state_layer_ids: list[int]
@classmethod
def from_zmq(cls, msg: List[bytes]):
@@ -199,6 +200,11 @@ class KVArgsRegisterInfo:
if len(msg) > 11 and len(msg[11]) > 0
else []
),
dst_state_layer_ids=(
list(struct.unpack(f"{len(msg[12])//4}i", msg[12]))
if len(msg) > 12 and len(msg[12]) > 0
else []
),
)
@@ -957,6 +963,22 @@ class MooncakeKVManager(CommonKVManager):
raise RuntimeError(
f"PD Disaggregation does NOT support PD different TP sizes for non-MLA {state_type.upper()} hybrid models yet."
)
src_state_layer_ids = list(
getattr(self.kv_args, "state_layer_ids", []) or []
)
dst_state_layer_ids = (
list(getattr(target_rank_registration_info, "dst_state_layer_ids", []) or [])
if target_rank_registration_info is not None
else []
)
if src_state_layer_ids or dst_state_layer_ids:
if src_state_layer_ids != dst_state_layer_ids:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][state_layer_ids] "
f"prefill={src_state_layer_ids} decode={dst_state_layer_ids} "
f"state_type={state_type} room={req.room} "
f"session={req.mooncake_session_id}"
)
effective_dst_state_indices = (
np.asarray(dst_state_indices, dtype=np.int32)
if dst_state_indices is not None
@@ -984,6 +1006,14 @@ class MooncakeKVManager(CommonKVManager):
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):
if src_state_layer_ids or dst_state_layer_ids:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][state_buffer_count] "
f"src={len(src_state_data_ptrs)} dst={len(dst_state_ptrs)} "
f"src_layers={src_state_layer_ids} "
f"dst_layers={dst_state_layer_ids} state_type={state_type} "
f"room={req.room} session={req.mooncake_session_id}"
)
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 "
@@ -1839,6 +1869,10 @@ class MooncakeKVReceiver(CommonKVReceiver):
packed_state_dim_per_tensor = b"".join(
struct.pack("I", dim) for dim in state_dim_per_tensor
)
packed_state_layer_ids = b"".join(
struct.pack("i", int(layer_id))
for layer_id in getattr(self.kv_mgr.kv_args, "state_layer_ids", [])
)
# Note(shangming): No need to add pp rank here since decode pp size should be equal to prefill pp size or 1
tp_rank = self.kv_mgr.kv_args.engine_rank
kv_item_len = self.kv_mgr.kv_args.kv_item_lens[0]
@@ -1880,6 +1914,7 @@ class MooncakeKVReceiver(CommonKVReceiver):
dst_kv_item_len,
packed_state_item_lens,
packed_state_dim_per_tensor,
packed_state_layer_ids,
]
)
python/sglang/srt/disaggregation/nixl/conn.py
+33
View File
@@ -84,6 +84,7 @@ class KVArgsRegisterInfo:
decode_tp_size: int
decode_tp_rank: int
dst_kv_item_len: int
dst_state_layer_ids: list[int]
@classmethod
def from_zmq(cls, msg: List[bytes]):
@@ -106,6 +107,11 @@ class KVArgsRegisterInfo:
decode_tp_size=int(msg[9].decode("ascii")),
decode_tp_rank=int(msg[10].decode("ascii")),
dst_kv_item_len=int(msg[11].decode("ascii")),
dst_state_layer_ids=(
list(struct.unpack(f"{len(msg[12]) // 4}i", msg[12]))
if len(msg) > 12 and msg[12] != b""
else []
),
)
@@ -680,6 +686,7 @@ class NixlKVManager(CommonKVManager):
dst_gpu_id: int,
notif: str,
decode_tp_size: int,
dst_state_layer_ids: Optional[List[int]] = None,
):
"""Send state or extra pool data with type-specific handling."""
state_type = getattr(self.kv_args, "state_type", "none")
@@ -702,6 +709,26 @@ class NixlKVManager(CommonKVManager):
raise RuntimeError(
f"PD Disaggregation does NOT support PD different TP sizes for non-MLA {state_type.upper()} hybrid models yet."
)
src_state_layer_ids = list(
getattr(self.kv_args, "state_layer_ids", []) or []
)
dst_state_layer_ids = list(dst_state_layer_ids or [])
if src_state_layer_ids or dst_state_layer_ids:
if src_state_layer_ids != dst_state_layer_ids:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][state_layer_ids] "
f"prefill={src_state_layer_ids} decode={dst_state_layer_ids} "
f"state_type={state_type} peer={peer_name} notif={notif}"
)
if len(self.kv_args.state_data_ptrs) != len(dst_state_data_ptrs):
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][state_buffer_count] "
f"src={len(self.kv_args.state_data_ptrs)} "
f"dst={len(dst_state_data_ptrs)} "
f"src_layers={src_state_layer_ids} "
f"dst_layers={dst_state_layer_ids} "
f"state_type={state_type} peer={peer_name} notif={notif}"
)
if len(prefill_state_indices) != len(dst_state_indices):
raise RuntimeError(
f"State index length mismatch: prefill={len(prefill_state_indices)}, "
@@ -791,6 +818,7 @@ class NixlKVManager(CommonKVManager):
dst_info.gpu_id,
f"{req.room}_state_{self.kv_args.pp_rank}",
decode_tp_size,
dst_info.dst_state_layer_ids,
)
if state_xfer_handle is not None:
handles.append(state_xfer_handle)
@@ -1069,6 +1097,10 @@ class NixlKVReceiver(CommonKVReceiver):
packed_state_data_ptrs = b"".join(
struct.pack("Q", ptr) for ptr in self.kv_mgr.kv_args.state_data_ptrs
)
packed_state_layer_ids = b"".join(
struct.pack("i", int(layer_id))
for layer_id in getattr(self.kv_mgr.kv_args, "state_layer_ids", [])
)
with lock:
sock.send_multipart(
@@ -1086,6 +1118,7 @@ class NixlKVReceiver(CommonKVReceiver):
str(self.kv_mgr.attn_tp_size).encode("ascii"),
str(self.kv_mgr.kv_args.engine_rank).encode("ascii"),
str(self.kv_mgr.kv_args.kv_item_lens[0]).encode("ascii"),
packed_state_layer_ids,
]
)
python/sglang/srt/disaggregation/prefill.py
+9
View File
@@ -191,6 +191,13 @@ def _state_buf_infos(pool):
return state_type, state_data_ptrs, state_data_lens, state_item_lens
def _state_layer_ids(pool):
get_state_layer_ids = getattr(pool, "get_state_layer_ids", None)
if get_state_layer_ids is None:
return []
return list(get_state_layer_ids())
def _kv_locs_to_page_indices_cpu(
kv_locs: torch.Tensor,
page_size: int,
@@ -376,6 +383,7 @@ class PrefillBootstrapQueue:
kv_args.state_data_ptrs = state_data_ptrs
kv_args.state_data_lens = state_data_lens
kv_args.state_item_lens = state_item_lens
kv_args.state_layer_ids = _state_layer_ids(self.token_to_kv_pool)
if isinstance(self.token_to_kv_pool, SWAKVPool):
kv_args.state_type = "swa"
@@ -394,6 +402,7 @@ class PrefillBootstrapQueue:
kv_args.state_data_ptrs = []
kv_args.state_data_lens = []
kv_args.state_item_lens = []
kv_args.state_layer_ids = []
kv_args.state_type = "none"
draft_state_type = "none"
python/sglang/srt/disaggregation/utils.py
+2
View File
@@ -259,6 +259,8 @@ def append_cp_draft_state_buffers(
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
if hasattr(kv_args, "state_layer_ids"):
kv_args.state_layer_ids += [-(i + 1) for i in range(draft_state_bufs)]
kv_args.draft_state_buffer_count = draft_state_bufs
return True