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

Prevent stale CP shared-KV contracts from corrupting prefill

Browse Source 
CP shared-KV now uses CP-local current rows consistently across MLA/index current reuse, passes fp8 current-index K through the tai-kernel uint8 ABI, and clears the transient EAGLE CP-local hidden marker after draft capture. The disaggregation bootstrap also fingerprints the runtime source contract so prefill/decode mismatches fail fast instead of silently exchanging incompatible KV metadata.

Constraint: CP shared-KV batch paths flatten current K/V rows in CP-rank-local valid order, not global request order.

Constraint: tai-kernel current-index prepare validates current_index_k as uint8 bytes for fp8 payloads.

Rejected: Keep using global extend offsets for bs>1 current-index reuse | corrupts request-local bases once current_index_kv is CP-local.

Rejected: Infer CP-local EAGLE hidden semantics from tensor shape | static padding and bs>1 can make shape-based inference unsafe.

Confidence: medium

Scope-risk: moderate

Directive: Do not reintroduce forward_batch.out_cache_loc slicing in CP shared-KV current reuse without verifying CP-local owner-lane layout.

Tested: Remote container py_compile for touched runtime/test files.

Tested: Remote PYTHONPATH=python pytest -q test/registered/unit/layers/test_nsa_cp_utils.py test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py test/registered/unit/disaggregation/test_common_conn_runtime_fingerprint.py (198 passed, 2 subtests passed).

Not-tested: Full remote ETE traffic after this commit; accept length and garbage-output recovery still require a fresh prefill/decode run.

Co-authored-by: OmX <omx@oh-my-codex.dev>
This commit is contained in:
laoyao0822
2026-06-04 20:22:29 +08:00
co-authored by OmX
parent 3d6007246b
commit f50e2b1e00
9 changed files with 1463 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/common/conn.py
+117
View File
@@ -2,11 +2,14 @@ from __future__ import annotations
import asyncio
import dataclasses
import hashlib
import importlib.util
import logging
import threading
import time
from collections import defaultdict
from functools import cache
from pathlib import Path
from typing import Dict, List, Optional, Set, Tuple, Union
import numpy as np
@@ -44,6 +47,75 @@ from sglang.srt.utils.network import (
logger = logging.getLogger(__name__)
_RUNTIME_FINGERPRINT_MODULES = (
"sglang.srt.layers.attention.nsa_backend",
"sglang.srt.layers.attention.nsa.nsa_indexer",
"sglang.srt.layers.attention.nsa.utils",
"sglang.srt.mem_cache.memory_pool",
"sglang.srt.disaggregation.mooncake.conn",
)
@cache
def _get_runtime_python_root() -> Optional[Path]:
spec = importlib.util.find_spec("sglang")
origin = getattr(spec, "origin", None) if spec is not None else None
if origin is None:
return None
path = Path(origin).resolve()
# .../<root>/python/sglang/__init__.py
for parent in path.parents:
if parent.name == "python":
return parent
return path.parent.parent
@cache
def get_runtime_source_root() -> Optional[str]:
python_root = _get_runtime_python_root()
if python_root is None:
return None
return str(python_root.parent)
@cache
def get_runtime_source_fingerprint() -> Optional[str]:
"""Return a stable fingerprint for code that defines the P/D KV contract.
Prefill and decode may run from different mounted source trees in long-lived
containers. CP shared-KV + FP8 + page-tail handling is not wire-compatible
across our recent development snapshots, so decode must fail fast when the
source contract differs instead of silently consuming malformed KV.
Do not import the target modules while fingerprinting: some package
``__init__`` files import the connection modules themselves. Build paths
from the already-imported ``sglang`` package root instead.
"""
python_root = _get_runtime_python_root()
if python_root is None:
return None
digest = hashlib.sha256()
for module_name in _RUNTIME_FINGERPRINT_MODULES:
origin = python_root.joinpath(*module_name.split(".")).with_suffix(".py")
digest.update(module_name.encode("utf-8"))
digest.update(b"\0")
try:
digest.update(origin.read_bytes())
except OSError as exc:
logger.warning(
"Failed to read source for disaggregation runtime fingerprint: "
"module=%s origin=%s error=%s",
module_name,
origin,
exc,
)
return None
digest.update(b"\0")
return digest.hexdigest()
@dataclasses.dataclass
class PrefillServerInfo:
# Topology fields (fetched from bootstrap server)
@@ -54,6 +126,8 @@ class PrefillServerInfo:
page_size: Optional[int]
kv_cache_dtype: Optional[str]
follow_bootstrap_room: bool
runtime_source_fingerprint: Optional[str] = None
runtime_source_root: Optional[str] = None
# Pre-computed rank mapping (set by try_ensure_parallel_info on decode side)
target_tp_rank: Optional[int] = None
@@ -73,6 +147,16 @@ class PrefillServerInfo:
str(self.kv_cache_dtype) if self.kv_cache_dtype is not None else None
)
self.follow_bootstrap_room = bool(self.follow_bootstrap_room)
self.runtime_source_fingerprint = (
str(self.runtime_source_fingerprint)
if self.runtime_source_fingerprint is not None
else None
)
self.runtime_source_root = (
str(self.runtime_source_root)
if self.runtime_source_root is not None
else None
)
@dataclasses.dataclass
@@ -236,6 +320,22 @@ class CommonKVManager(BaseKVManager):
f"Both servers must use the same --kv-cache-dtype value."
)
local_source_fingerprint = get_runtime_source_fingerprint()
if (
info.runtime_source_fingerprint is not None
and local_source_fingerprint is not None
and info.runtime_source_fingerprint != local_source_fingerprint
):
raise RuntimeError(
"Runtime source fingerprint mismatch between prefill and decode. "
"This can corrupt disaggregated KV transfer contracts. "
f"prefill_fingerprint={info.runtime_source_fingerprint}, "
f"decode_fingerprint={local_source_fingerprint}, "
f"prefill_source_root={info.runtime_source_root}, "
f"decode_source_root={get_runtime_source_root()}. "
"Restart both prefill and decode from the same source tree."
)
self._resolve_rank_mapping(info)
self.prefill_info_table[bootstrap_addr] = info
logger.debug(f"Prefill parallel info for [{bootstrap_addr}]: {info}")
@@ -348,6 +448,8 @@ class CommonKVManager(BaseKVManager):
"page_size": self.kv_args.page_size,
"kv_cache_dtype": self.server_args.kv_cache_dtype,
"load_balance_method": self.server_args.load_balance_method,
"runtime_source_fingerprint": get_runtime_source_fingerprint(),
"runtime_source_root": get_runtime_source_root(),
}
try:
@@ -668,6 +770,8 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer):
self.page_size = None
self.kv_cache_dtype: Optional[str] = None
self.follow_bootstrap_room: Optional[bool] = None
self.runtime_source_fingerprint: Optional[str] = get_runtime_source_fingerprint()
self.runtime_source_root: Optional[str] = get_runtime_source_root()
self.prefill_port_table: Dict[
int, Dict[int, Dict[int, Dict[int, PrefillRankInfo]]]
] = {}
@@ -734,6 +838,8 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer):
rank_port = int(data["rank_port"])
page_size = int(data["page_size"])
kv_cache_dtype = data["kv_cache_dtype"]
runtime_source_fingerprint = data.get("runtime_source_fingerprint")
runtime_source_root = data.get("runtime_source_root")
if self.attn_tp_size is None:
self.attn_tp_size = attn_tp_size
@@ -753,6 +859,15 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer):
if self.kv_cache_dtype is None and kv_cache_dtype is not None:
self.kv_cache_dtype = kv_cache_dtype
if (
self.runtime_source_fingerprint is None
and runtime_source_fingerprint is not None
):
self.runtime_source_fingerprint = str(runtime_source_fingerprint)
if self.runtime_source_root is None and runtime_source_root is not None:
self.runtime_source_root = str(runtime_source_root)
if self.follow_bootstrap_room is None:
load_balance_method = data.get(
"load_balance_method", "follow_bootstrap_room"
@@ -822,6 +937,8 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer):
if self.follow_bootstrap_room is not None
else True
),
runtime_source_fingerprint=self.runtime_source_fingerprint,
runtime_source_root=self.runtime_source_root,
)
return web.json_response(dataclasses.asdict(info), status=200)