Stabilize CP HiCache residency under L1/L2 pressure
CP shared KV now keeps explicit L1 and host free-room targets so pressure is handled by planned eviction instead of repeated capacity-edge retries. The host allocator gains contiguous-preferred page reservation, L1 owner-lane allocation prefers contiguous physical pages, and CP HiCache metadata preserves pending backup safety for page-granular radix updates. Mooncake transfer stats and allocator microbenchmarks are included to make the remaining transfer bottlenecks measurable rather than inferred. Constraint: CP shared KV uses decode CP size 1 with all prefill CP ranks participating in transfer, so L1/L2 cache residency must remain page-granular and avoid extra collectives.\nConstraint: Production HiCache can be hundreds of GB, so allocator metadata overhead must be visible before enabling aggressive contiguous allocation broadly.\nRejected: Evict only the exact deficit | this keeps the cache at the cliff and causes repeated evict/allocate pressure.\nRejected: Rely on allocator scans alone for contiguity | remote microbenchmarks show fragmented 220GB-equivalent host metadata can make contiguous-preferred scans multi-ms.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not increase L1/L2 free-room defaults or add new CP collectives without ETE evidence and transfer/allocator measurements.\nTested: python -m py_compile on touched runtime/test/benchmark files.\nTested: PYTHONPATH=. python -m pytest -q test/registered/unit/benchmark/test_cp_hicache_allocator_bench.py => 4 passed, 1 warning.\nTested: Remote g0034 log /mnt/beegfs/cjy/log/sglang_cp_hicache_20260601_233723.log shows active prefill process with L1/L2 free-room args, 702 HTTP 200 chat completions, 6272 prefill batches, and no fatal scheduler traceback in latest scan.\nTested: User-reported L1/L2 cache ETE validation passed on remote run.\nNot-tested: Full local pytest suite; local environment is missing several runtime dependencies.\nNot-tested: CUDA allocator microbenchmark during active production prefill process.\nNot-tested: Mooncake straggler fix; stats show transfer tail latency remains a separate bottleneck.
This commit is contained in:
laoyao0822
@@ -40,3 +40,36 @@ def group_concurrent_contiguous(
|
||||
dst_groups = [g.tolist() for g in dst_groups]
|
||||
|
||||
return src_groups, dst_groups
|
||||
|
||||
|
||||
def contiguous_group_stats(
|
||||
src_indices: npt.NDArray[np.int32],
|
||||
dst_indices: npt.NDArray[np.int32],
|
||||
src_groups: List[npt.NDArray[np.int32]],
|
||||
dst_groups: List[npt.NDArray[np.int32]],
|
||||
) -> dict[str, object]:
|
||||
"""Summarize contiguous-group coalescing without exposing full index arrays."""
|
||||
|
||||
del dst_groups # Same grouping shape as src_groups; keep arg for call-site clarity.
|
||||
|
||||
src_indices = np.asarray(src_indices).reshape(-1)
|
||||
dst_indices = np.asarray(dst_indices).reshape(-1)
|
||||
group_lens = [len(group) for group in src_groups]
|
||||
page_count = int(src_indices.size)
|
||||
group_count = int(len(group_lens))
|
||||
|
||||
def _diff_head(indices: npt.NDArray[np.int32], limit: int = 16) -> list[int]:
|
||||
if indices.size <= 1:
|
||||
return []
|
||||
sample = indices[: min(indices.size, limit + 1)]
|
||||
return np.diff(sample).astype(np.int64).tolist()
|
||||
|
||||
return {
|
||||
"pages": page_count,
|
||||
"groups": group_count,
|
||||
"min_group_pages": int(min(group_lens)) if group_lens else 0,
|
||||
"max_group_pages": int(max(group_lens)) if group_lens else 0,
|
||||
"avg_group_pages": float(page_count / group_count) if group_count else 0.0,
|
||||
"src_diff_head": _diff_head(src_indices),
|
||||
"dst_diff_head": _diff_head(dst_indices),
|
||||
}
|
||||
|
||||
@@ -23,6 +23,7 @@ from sglang.srt.disaggregation.common.conn import (
|
||||
)
|
||||
from sglang.srt.disaggregation.common.utils import (
|
||||
FastQueue,
|
||||
contiguous_group_stats,
|
||||
group_concurrent_contiguous,
|
||||
)
|
||||
from sglang.srt.disaggregation.mooncake.utils import (
|
||||
@@ -62,6 +63,24 @@ def _cp_draft_shared_kv_debug(message: str, *args, limit: int = 64) -> None:
|
||||
logger.info("[CP_DRAFT_SHARED_KV] " + message, *args)
|
||||
|
||||
|
||||
def _mooncake_transfer_stats_enabled() -> bool:
|
||||
return envs.SGLANG_DISAGGREGATION_TRANSFER_STATS.get()
|
||||
|
||||
|
||||
def _mooncake_transfer_stats_log(message: str, *args) -> None:
|
||||
if not _mooncake_transfer_stats_enabled():
|
||||
return
|
||||
limit = envs.SGLANG_DISAGGREGATION_TRANSFER_STATS_LIMIT.get()
|
||||
if limit is not None and limit <= 0:
|
||||
return
|
||||
key = "mooncake_transfer_stats"
|
||||
count = _CP_SHARED_DEBUG_COUNTS.get(key, 0)
|
||||
if limit is not None and count >= limit:
|
||||
return
|
||||
_CP_SHARED_DEBUG_COUNTS[key] = count + 1
|
||||
logger.info("[Mooncake-transfer-stats] " + message, *args)
|
||||
|
||||
|
||||
def _np_summary(arr) -> str:
|
||||
if arr is None:
|
||||
return "None"
|
||||
@@ -310,6 +329,7 @@ class MooncakeKVManager(CommonKVManager):
|
||||
prefill_data_indices: npt.NDArray[np.int32],
|
||||
dst_data_indices: npt.NDArray[np.int32],
|
||||
executor: concurrent.futures.ThreadPoolExecutor,
|
||||
debug_room: Optional[int] = None,
|
||||
) -> int:
|
||||
"""
|
||||
Generic KV cache transfer supporting both MHA and MLA architectures.
|
||||
@@ -319,6 +339,15 @@ class MooncakeKVManager(CommonKVManager):
|
||||
prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous(
|
||||
prefill_data_indices, dst_data_indices
|
||||
)
|
||||
transfer_stats_enabled = _mooncake_transfer_stats_enabled()
|
||||
grouping_stats = None
|
||||
if transfer_stats_enabled:
|
||||
grouping_stats = contiguous_group_stats(
|
||||
prefill_data_indices,
|
||||
dst_data_indices,
|
||||
prefill_kv_blocks,
|
||||
dst_kv_blocks,
|
||||
)
|
||||
|
||||
layers_params = None
|
||||
|
||||
@@ -387,6 +416,7 @@ class MooncakeKVManager(CommonKVManager):
|
||||
transfer_blocks.extend(set_transfer_blocks(src_ptr, dst_ptr, item_len))
|
||||
return self._transfer_data(mooncake_session_id, transfer_blocks)
|
||||
|
||||
start_time = time.perf_counter() if transfer_stats_enabled else 0.0
|
||||
if self.enable_custom_mem_pool:
|
||||
futures = [
|
||||
executor.submit(
|
||||
@@ -402,12 +432,84 @@ class MooncakeKVManager(CommonKVManager):
|
||||
if status != 0:
|
||||
for f in futures:
|
||||
f.cancel()
|
||||
if transfer_stats_enabled:
|
||||
self._log_kvcache_transfer_stats(
|
||||
mooncake_session_id=mooncake_session_id,
|
||||
debug_room=debug_room,
|
||||
grouping_stats=grouping_stats,
|
||||
layers_params=layers_params,
|
||||
elapsed_ms=(time.perf_counter() - start_time) * 1000,
|
||||
status=status,
|
||||
custom_mem_pool=True,
|
||||
)
|
||||
return status
|
||||
return 0
|
||||
status = 0
|
||||
else:
|
||||
# Combining all layers' params in one batch transfer is more efficient
|
||||
# compared to using multiple threads
|
||||
return process_layers(layers_params)
|
||||
status = process_layers(layers_params)
|
||||
|
||||
if transfer_stats_enabled:
|
||||
self._log_kvcache_transfer_stats(
|
||||
mooncake_session_id=mooncake_session_id,
|
||||
debug_room=debug_room,
|
||||
grouping_stats=grouping_stats,
|
||||
layers_params=layers_params,
|
||||
elapsed_ms=(time.perf_counter() - start_time) * 1000,
|
||||
status=status,
|
||||
custom_mem_pool=self.enable_custom_mem_pool,
|
||||
)
|
||||
return status
|
||||
|
||||
def _log_kvcache_transfer_stats(
|
||||
self,
|
||||
mooncake_session_id: str,
|
||||
debug_room: Optional[int],
|
||||
grouping_stats: Optional[dict[str, object]],
|
||||
layers_params: List[Tuple[int, int, int]],
|
||||
elapsed_ms: float,
|
||||
status: int,
|
||||
custom_mem_pool: bool,
|
||||
) -> None:
|
||||
if grouping_stats is None:
|
||||
return
|
||||
page_count = int(grouping_stats["pages"])
|
||||
group_count = int(grouping_stats["groups"])
|
||||
layer_count = len(layers_params)
|
||||
item_bytes_per_page = sum(int(item_len) for _, _, item_len in layers_params)
|
||||
total_bytes = page_count * item_bytes_per_page
|
||||
transfer_blocks = group_count * layer_count
|
||||
avg_block_bytes = (
|
||||
float(total_bytes / transfer_blocks) if transfer_blocks else 0.0
|
||||
)
|
||||
bandwidth_gbps = (
|
||||
float(total_bytes / elapsed_ms / 1e6) if elapsed_ms > 0 else 0.0
|
||||
)
|
||||
|
||||
_mooncake_transfer_stats_log(
|
||||
"cp_rank=%s room=%s session=%s status=%s custom_mem_pool=%s "
|
||||
"pages=%s groups=%s avg_group_pages=%.2f max_group_pages=%s "
|
||||
"layers=%s transfer_blocks=%s total_bytes=%.3fGiB "
|
||||
"avg_block_bytes=%.1fKiB elapsed_ms=%.3f bandwidth=%.2fGB/s "
|
||||
"src_diff_head=%s dst_diff_head=%s",
|
||||
self.attn_cp_rank,
|
||||
debug_room,
|
||||
mooncake_session_id,
|
||||
status,
|
||||
custom_mem_pool,
|
||||
page_count,
|
||||
group_count,
|
||||
float(grouping_stats["avg_group_pages"]),
|
||||
grouping_stats["max_group_pages"],
|
||||
layer_count,
|
||||
transfer_blocks,
|
||||
total_bytes / (1024**3),
|
||||
avg_block_bytes / 1024,
|
||||
elapsed_ms,
|
||||
bandwidth_gbps,
|
||||
grouping_stats["src_diff_head"],
|
||||
grouping_stats["dst_diff_head"],
|
||||
)
|
||||
|
||||
def send_kvcache(
|
||||
self,
|
||||
@@ -416,6 +518,7 @@ class MooncakeKVManager(CommonKVManager):
|
||||
dst_kv_ptrs: list[int],
|
||||
dst_kv_indices: npt.NDArray[np.int32],
|
||||
executor: concurrent.futures.ThreadPoolExecutor,
|
||||
debug_room: Optional[int] = None,
|
||||
):
|
||||
return self._send_kvcache_generic(
|
||||
mooncake_session_id=mooncake_session_id,
|
||||
@@ -425,6 +528,7 @@ class MooncakeKVManager(CommonKVManager):
|
||||
prefill_data_indices=prefill_kv_indices,
|
||||
dst_data_indices=dst_kv_indices,
|
||||
executor=executor,
|
||||
debug_room=debug_room,
|
||||
)
|
||||
|
||||
def send_kvcache_slice(
|
||||
@@ -973,6 +1077,7 @@ class MooncakeKVManager(CommonKVManager):
|
||||
target_rank_registration_info.dst_kv_ptrs,
|
||||
chunked_dst_kv_indice,
|
||||
executor,
|
||||
debug_room=kv_chunk.room,
|
||||
)
|
||||
else:
|
||||
ret = self.send_kvcache_slice(
|
||||
|
||||
Reference in New Issue
Block a user