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

Expand prefill CP KV capacity by sharding persistent NSA KV

Browse Source 
Prefill CP previously replicated NSA/MLA persistent KV on every CP rank, so CP8 consumed eight copies of KV memory while exposing only one rank of logical cache capacity. This change splits logical KV locs from per-rank physical storage, shards MLA latent KV and NSA index K/scale by deterministic page ownership, and keeps existing NSA attention kernels working through a full-view runtime materialization layer.

Mooncake PD transfer now sends each prefill CP rank's owned physical pages with explicit logical page positions so non-CP decode can reconstruct full-layout KV. The implementation is guarded by an explicit server flag and startup checks, and the design documentation records the implemented scope, debug environment, and Phase 3 boundary.

Constraint: Phase 2 must preserve existing NSA attention/index kernels via runtime full-view materialization
Constraint: Decode side remains non-CP and receives full KV through Mooncake
Rejected: Shard-aware NSA attention in this change | belongs to Phase 3 because it requires distributed topk/softmax/output contracts
Rejected: Request-contiguous CP ownership | unstable under chunked prefill and tied to attention split mode
Confidence: medium
Scope-risk: broad
Directive: Do not enable round-robin CP shared KV without wiring runtime materialization/PD transfer contracts for that split mode
Directive: Keep SGLANG_DEBUG_CP_SHARED_KV disabled for perf measurements; it intentionally enables CUDA-syncing diagnostics
Tested: Remote py_compile for shared-KV touched Python files in g0034 container
Tested: Remote pytest selected cp_shared/shared_kv/nsa suite: 37 passed, 34 deselected
Not-tested: Full GLM5 multi-node throughput/regression run after final doc update
Not-tested: Phase 3 shard-aware runtime, round-robin CP mode, and non-Mooncake PD backends
This commit is contained in:
laoyao0822
2026-04-26 04:11:17 +08:00
parent 76c830dda8
commit f8fca72635
22 changed files with 2489 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/mooncake/conn.py
+140 -9
View File
@@ -38,6 +38,30 @@ from sglang.srt.server_args import ServerArgs
from sglang.srt.utils.network import NetworkAddress
logger = logging.getLogger(__name__)
_CP_SHARED_DEBUG_COUNTS: dict[str, int] = {}
def _cp_shared_debug_log(key: str, message: str, *args, limit: int = 64) -> None:
if not envs.SGLANG_DEBUG_CP_SHARED_KV.get():
return
count = _CP_SHARED_DEBUG_COUNTS.get(key, 0)
if count >= limit:
return
_CP_SHARED_DEBUG_COUNTS[key] = count + 1
logger.info("[CP_SHARED_KV_DEBUG] " + message, *args)
def _np_summary(arr) -> str:
if arr is None:
return "None"
arr = np.asarray(arr)
if arr.size == 0:
return f"shape={arr.shape} dtype={arr.dtype} size=0"
head = arr.reshape(-1)[: min(8, arr.size)].tolist()
return (
f"shape={arr.shape} dtype={arr.dtype} size={arr.size} "
f"min={int(arr.min())} max={int(arr.max())} head={head}"
)
class KVTransferError(Exception):
@@ -55,7 +79,9 @@ class KVTransferError(Exception):
class TransferKVChunk:
room: int
prefill_kv_indices: npt.NDArray[np.int32]
index_slice: slice
index_slice: Optional[slice]
logical_page_positions: Optional[npt.NDArray[np.int32]]
state_logical_page_positions: Optional[npt.NDArray[np.int32]]
is_last_chunk: bool
prefill_aux_index: Optional[int]
state_indices: Optional[List[int]]
@@ -593,6 +619,7 @@ class MooncakeKVManager(CommonKVManager):
dst_state_data_ptrs: list[int],
executor: concurrent.futures.ThreadPoolExecutor,
target_rank_registration_info: Optional[KVArgsRegisterInfo] = None,
dst_state_indices: Optional[npt.NDArray[np.int32]] = None,
):
"""Send state or extra pool data with type-specific handling."""
state_type = getattr(self.kv_args, "state_type", "none")
@@ -628,23 +655,27 @@ class MooncakeKVManager(CommonKVManager):
raise RuntimeError(
f"PD Disaggregation does NOT support PD different TP sizes for non-MLA {state_type.upper()} hybrid models yet."
)
if len(prefill_state_indices) < len(req.dst_state_indices):
effective_dst_state_indices = (
np.asarray(dst_state_indices, dtype=np.int32)
if dst_state_indices is not None
else np.asarray(req.dst_state_indices, dtype=np.int32)
)
if len(prefill_state_indices) < len(effective_dst_state_indices):
logger.warning(
f"len(prefill_state_indices) = {len(prefill_state_indices)}, len(dst_state_indices) = {len(req.dst_state_indices)}"
f"len(prefill_state_indices) = {len(prefill_state_indices)}, len(dst_state_indices) = {len(effective_dst_state_indices)}"
)
prefill_state_indices = prefill_state_indices[
: len(req.dst_state_indices)
: len(effective_dst_state_indices)
]
# Reuse _send_kvcache_generic interface to send extra pool data
prefill_state_indices = np.array(prefill_state_indices, dtype=np.int32)
dst_state_indices = np.array(req.dst_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,
prefill_data_indices=prefill_state_indices,
dst_data_indices=dst_state_indices,
dst_data_indices=effective_dst_state_indices,
executor=executor,
)
else:
@@ -807,7 +838,27 @@ class MooncakeKVManager(CommonKVManager):
)
break
chunked_dst_kv_indice = req.dst_kv_indices[kv_chunk.index_slice]
if kv_chunk.logical_page_positions is not None:
chunked_dst_kv_indice = req.dst_kv_indices[
kv_chunk.logical_page_positions
]
else:
assert kv_chunk.index_slice is not None
chunked_dst_kv_indice = req.dst_kv_indices[
kv_chunk.index_slice
]
if envs.SGLANG_DEBUG_CP_SHARED_KV.get():
_cp_shared_debug_log(
"transfer_worker_kv",
"transfer worker cp_rank=%s room=%s prefill_pages=%s "
"logical_positions=%s dst_pages=%s is_last=%s",
self.attn_cp_rank,
kv_chunk.room,
_np_summary(kv_chunk.prefill_kv_indices),
_np_summary(kv_chunk.logical_page_positions),
_np_summary(chunked_dst_kv_indice),
kv_chunk.is_last_chunk,
)
# NOTE: This is temporarily a workaround to deal with the case where the prefill_kv_indices
# is mismatched with the dst_kv_indices when page size > 1, this should never happen.
@@ -871,12 +922,36 @@ class MooncakeKVManager(CommonKVManager):
if kv_chunk.is_last_chunk:
if kv_chunk.state_indices is not None:
dst_state_indices = None
if kv_chunk.state_logical_page_positions is not None:
from sglang.srt.disaggregation.utils import (
select_pages_by_request_positions,
)
dst_state_indices = select_pages_by_request_positions(
req.dst_state_indices,
kv_chunk.state_logical_page_positions,
)
if envs.SGLANG_DEBUG_CP_SHARED_KV.get():
_cp_shared_debug_log(
"transfer_worker_state",
"transfer worker state cp_rank=%s room=%s prefill_state_pages=%s "
"state_positions=%s dst_state_pages=%s",
self.attn_cp_rank,
kv_chunk.room,
_np_summary(kv_chunk.state_indices),
_np_summary(
kv_chunk.state_logical_page_positions
),
_np_summary(dst_state_indices),
)
self.maybe_send_extra(
req,
kv_chunk.state_indices,
target_rank_registration_info.dst_state_data_ptrs,
executor,
target_rank_registration_info,
dst_state_indices,
)
# Only the last chunk we need to send the aux data
@@ -1049,10 +1124,12 @@ class MooncakeKVManager(CommonKVManager):
self,
bootstrap_room: int,
kv_indices: npt.NDArray[np.int32],
index_slice: slice,
index_slice: Optional[slice],
is_last_chunk: bool,
aux_index: Optional[int] = None,
state_indices: Optional[List[int]] = None,
logical_page_positions: Optional[npt.NDArray[np.int32]] = None,
state_logical_page_positions: Optional[npt.NDArray[np.int32]] = None,
):
assert self.disaggregation_mode == DisaggregationMode.PREFILL
assert not is_last_chunk or (is_last_chunk and aux_index is not None)
@@ -1084,6 +1161,8 @@ class MooncakeKVManager(CommonKVManager):
room=bootstrap_room,
prefill_kv_indices=kv_indices,
index_slice=index_slice,
logical_page_positions=logical_page_positions,
state_logical_page_positions=state_logical_page_positions,
is_last_chunk=is_last_chunk,
prefill_aux_index=aux_index,
state_indices=state_indices,
@@ -1147,9 +1226,57 @@ class MooncakeKVSender(CommonKVSender):
index_slice = slice(self.curr_idx, self.curr_idx + len(kv_indices))
self.curr_idx += len(kv_indices)
is_last_chunk = self.curr_idx == self.num_kv_indices
logical_page_positions = None
state_logical_page_positions = None
orig_kv_indices = None
orig_state_indices = None
if self.kv_mgr.server_args.enable_nsa_prefill_cp_shared_kv:
from sglang.srt.disaggregation.utils import filter_kv_pages_for_cp_shared_kv
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
chunk_page_start = index_slice.start
orig_kv_indices = np.asarray(kv_indices, dtype=np.int32).copy()
if state_indices is not None:
orig_state_indices = np.asarray(state_indices, dtype=np.int32).copy()
layout = CpSharedKVLayout(
page_size=self.kv_mgr.kv_args.page_size,
cp_size=self.kv_mgr.attn_cp_size,
cp_rank=self.kv_mgr.attn_cp_rank,
)
kv_indices, logical_page_positions = filter_kv_pages_for_cp_shared_kv(
layout=layout,
logical_pages=kv_indices,
chunk_page_start=index_slice.start,
)
if state_indices is not None:
state_indices, state_logical_page_positions = (
filter_kv_pages_for_cp_shared_kv(
layout=layout,
logical_pages=state_indices,
chunk_page_start=0,
)
)
index_slice = None
if envs.SGLANG_DEBUG_CP_SHARED_KV.get():
_cp_shared_debug_log(
"sender_filter",
"sender filter cp_rank=%s room=%s page_start=%s orig_kv_pages=%s "
"filtered_kv_pages=%s kv_positions=%s orig_state_pages=%s "
"filtered_state_pages=%s state_positions=%s is_last=%s",
self.kv_mgr.attn_cp_rank,
self.bootstrap_room,
chunk_page_start,
_np_summary(orig_kv_indices),
_np_summary(kv_indices),
_np_summary(logical_page_positions),
_np_summary(orig_state_indices),
_np_summary(state_indices),
_np_summary(state_logical_page_positions),
is_last_chunk,
)
# Special handling for cp
if self.kv_mgr.enable_all_cp_ranks_for_transfer:
elif self.kv_mgr.enable_all_cp_ranks_for_transfer:
kv_indices, index_slice = filter_kv_indices_for_cp_rank(
self.kv_mgr,
kv_indices,
@@ -1168,6 +1295,8 @@ class MooncakeKVSender(CommonKVSender):
kv_indices,
index_slice,
False,
logical_page_positions=logical_page_positions,
state_logical_page_positions=state_logical_page_positions,
)
else:
self.kv_mgr.add_transfer_request(
@@ -1177,6 +1306,8 @@ class MooncakeKVSender(CommonKVSender):
True,
aux_index=self.aux_index,
state_indices=state_indices,
logical_page_positions=logical_page_positions,
state_logical_page_positions=state_logical_page_positions,
)
def poll(self) -> KVPoll:
python/sglang/srt/disaggregation/utils.py
+38
View File
@@ -505,6 +505,44 @@ def filter_kv_indices_for_cp_rank(
return new_kv_indices, new_index_slice
def filter_kv_pages_for_cp_shared_kv(
layout,
logical_pages: np.ndarray,
chunk_page_start: int,
) -> Tuple[np.ndarray, np.ndarray]:
logical_pages = np.asarray(logical_pages, dtype=np.int32)
if logical_pages.size > 0 and int(logical_pages.min()) < 0:
bad_pages = logical_pages[logical_pages < 0]
raise ValueError(
"CP shared KV transfer got negative logical_pages. "
f"bad_min={int(bad_pages.min())} bad_max={int(bad_pages.max())} "
f"chunk_page_start={chunk_page_start} num_pages={len(logical_pages)}"
)
request_positions = np.arange(
chunk_page_start,
chunk_page_start + len(logical_pages),
dtype=np.int32,
)
return layout.filter_owned_pages_np(logical_pages, request_positions)
def select_pages_by_request_positions(
pages: np.ndarray,
request_positions: np.ndarray,
) -> np.ndarray:
"""Select per-request page entries by absolute request page positions."""
pages = np.asarray(pages, dtype=np.int32)
request_positions = np.asarray(request_positions, dtype=np.int32)
if request_positions.size == 0:
return pages[:0]
if int(request_positions.max()) >= len(pages):
raise ValueError(
"request_positions contains an entry outside pages: "
f"max_position={int(request_positions.max())} pages={len(pages)}"
)
return pages[request_positions]
#########################
# Misc
#########################