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

Reduce CP shared-KV request-boundary stalls

Browse Source 
CP shared KV now avoids the PyTorch sort/search remap for the single-request current-only path by deriving compact rows from page-level inverse mapping. The same change keeps sort NVTX attribution gated and splits high-frequency MoE sort markers behind a separate env var so profiling does not perturb normal runs.

Decode-side disaggregation prealloc also avoids rebuilding large token index tensors and records finer allocation timing, while compute-owner allocation/free tests cover the shared-KV page-lane behavior.

Constraint: The runtime tree used for validation is the remote /sgl-workspace/sglang-tai mount, which is not itself a Git repository, so these tracked files were synchronized into the local repo before commit.

Rejected: Keep torch.sort/searchsorted for current remap | it emits ATen/CCCL radixSortKVInPlace kernels in the attention hot path.

Rejected: Enable MoE sort NVTX under the generic sort env | the MoE preprocess sort is too frequent and can make profiling look like a hang.

Confidence: medium

Scope-risk: moderate

Directive: Do not reintroduce token-level torch.sort/searchsorted in CP shared-KV current remap without profiling the attention hot path under Nsight.

Tested: Remote container py_compile for modified runtime files; git diff --cached --check.

Not-tested: Full multi-node GLM5 PD throughput/profile rerun after the page-inverse current remap.
This commit is contained in:
laoyao0822
2026-05-05 05:18:35 +08:00
parent a638d71d53
commit 49eaf9ffde
11 changed files with 551 additions and 92 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/decode.py
+57 -24
View File
@@ -43,7 +43,6 @@ from sglang.srt.disaggregation.utils import (
TransferBackend,
get_kv_class,
is_mla_backend,
kv_to_page_indices,
poll_and_all_reduce,
prepare_abort,
)
@@ -78,6 +77,29 @@ if TYPE_CHECKING:
CLIP_MAX_NEW_TOKEN = envs.SGLANG_CLIP_MAX_NEW_TOKENS_ESTIMATION.get()
def _kv_locs_to_page_indices_cpu(
kv_locs: torch.Tensor,
page_size: int,
*,
num_tokens: Optional[int] = None,
):
"""Return compact int32 page indices for a token-location tensor.
PD bootstrap only needs page-level destination indices. Avoid materializing
a full token-level CPU copy before converting to pages.
"""
if num_tokens is not None:
kv_locs = kv_locs[:num_tokens]
if page_size == 1:
page_locs = kv_locs
else:
page_locs = kv_locs[::page_size] // page_size
return page_locs.to(dtype=torch.int32).cpu().numpy()
def _is_fake_transfer(req: Req, server_args: ServerArgs) -> bool:
return req.bootstrap_host == FAKE_BOOTSTRAP_HOST or (
req.bootstrap_host is None
@@ -409,7 +431,11 @@ class DecodePreallocQueue:
)
self.queue.append(
DecodeRequest(req=req, kv_receiver=kv_receiver, waiting_for_input=False)
DecodeRequest(
req=req,
kv_receiver=kv_receiver,
waiting_for_input=False,
)
)
def _check_if_req_exceed_kv_capacity(self, req: Req) -> bool:
@@ -675,16 +701,17 @@ class DecodePreallocQueue:
continue
allocatable_tokens -= required_tokens_for_request
self._pre_alloc(decode_req.req)
kv_indices = (
self.req_to_token_pool.req_to_token[decode_req.req.req_pool_idx][
: len(decode_req.req.origin_input_ids)
]
.cpu()
.numpy()
kv_loc = self._pre_alloc(
decode_req.req,
write_req_to_token=False,
)
page_size = self.token_to_kv_pool_allocator.page_size
page_indices = _kv_locs_to_page_indices_cpu(
kv_loc,
page_size,
num_tokens=origin_input_len,
)
# Prepare extra pool indices for hybrid models
if isinstance(self.token_to_kv_pool, HybridLinearKVPool):
@@ -703,9 +730,7 @@ class DecodePreallocQueue:
window_start = max(0, seq_len - window_size)
window_start = (window_start // page_size) * page_size
window_kv_indices_full = self.req_to_token_pool.req_to_token[
decode_req.req.req_pool_idx, window_start:seq_len
]
window_kv_indices_full = kv_loc[window_start:seq_len]
# Translate to SWA pool indices
window_kv_indices_swa = (
@@ -713,23 +738,23 @@ class DecodePreallocQueue:
window_kv_indices_full
)
)
state_indices = window_kv_indices_swa.cpu().numpy()
state_indices = kv_to_page_indices(state_indices, page_size)
state_indices = _kv_locs_to_page_indices_cpu(
window_kv_indices_swa,
page_size,
)
elif isinstance(self.token_to_kv_pool, NSATokenToKVPool):
seq_len = len(decode_req.req.origin_input_ids)
kv_indices_full = self.req_to_token_pool.req_to_token[
decode_req.req.req_pool_idx, :seq_len
]
state_indices = kv_indices_full.cpu().numpy()
state_indices = kv_to_page_indices(state_indices, page_size)
state_indices = page_indices
else:
state_indices = None
self.req_to_token_pool.write(
(decode_req.req.req_pool_idx, slice(0, len(kv_loc))), kv_loc
)
decode_req.metadata_buffer_index = (
self.req_to_metadata_buffer_idx_allocator.alloc()
)
assert decode_req.metadata_buffer_index is not None
page_indices = kv_to_page_indices(kv_indices, page_size)
decode_req.kv_receiver.init(
page_indices, decode_req.metadata_buffer_index, state_indices
)
@@ -840,7 +865,12 @@ class DecodePreallocQueue:
last_loc,
)
def _pre_alloc(self, req: Req) -> torch.Tensor:
def _pre_alloc(
self,
req: Req,
*,
write_req_to_token: bool = True,
) -> torch.Tensor:
"""Pre-allocate the memory for req_to_token and token_kv_pool"""
req_pool_indices = self.req_to_token_pool.alloc([req])
@@ -875,7 +905,10 @@ class DecodePreallocQueue:
"KV cache is full! There is a bug in memory estimation."
)
self.req_to_token_pool.write((req.req_pool_idx, slice(0, len(kv_loc))), kv_loc)
if write_req_to_token:
self.req_to_token_pool.write(
(req.req_pool_idx, slice(0, len(kv_loc))), kv_loc
)
# populate metadata
req.fill_ids = req.origin_input_ids + req.output_ids
python/sglang/srt/disaggregation/prefill.py
+25 -13
View File
@@ -37,7 +37,6 @@ from sglang.srt.disaggregation.utils import (
TransferBackend,
get_kv_class,
is_mla_backend,
kv_to_page_indices,
kv_to_page_num,
poll_and_all_reduce_attn_cp_tp_group,
prepare_abort,
@@ -62,6 +61,20 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
def _kv_locs_to_page_indices_cpu(
kv_locs: torch.Tensor,
page_size: int,
):
"""Return int32 page indices without materializing token-level CPU indices."""
if page_size == 1:
page_locs = kv_locs
else:
page_locs = kv_locs[::page_size] // page_size
return page_locs.to(dtype=torch.int32).cpu().numpy()
def release_req_to_metadata_buffer(
req: Req, allocator: ReqToMetadataIdxAllocator
) -> None:
@@ -723,10 +736,9 @@ class SchedulerDisaggregationPrefillMixin:
# if not the last chunk and the last page is partial, delay the last partial page to the next send
end_idx = end_idx - end_idx % page_size
kv_indices = (
self.req_to_token_pool.req_to_token[req.req_pool_idx, start_idx:end_idx]
.cpu()
.numpy()
page_indices = _kv_locs_to_page_indices_cpu(
self.req_to_token_pool.req_to_token[req.req_pool_idx, start_idx:end_idx],
page_size,
)
req.start_send_idx = end_idx
state_indices = None
@@ -762,19 +774,19 @@ class SchedulerDisaggregationPrefillMixin:
window_kv_indices_full
)
)
state_indices = window_kv_indices_swa.cpu().numpy()
state_indices = kv_to_page_indices(state_indices, page_size)
state_indices = _kv_locs_to_page_indices_cpu(
window_kv_indices_swa,
page_size,
)
elif isinstance(
self.token_to_kv_pool_allocator.get_kvcache(), NSATokenToKVPool
):
seq_len = len(req.fill_ids)
kv_indices_full = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :seq_len
]
state_indices = kv_indices_full.cpu().numpy()
state_indices = kv_to_page_indices(state_indices, page_size)
state_indices = _kv_locs_to_page_indices_cpu(
self.req_to_token_pool.req_to_token[req.req_pool_idx, :seq_len],
page_size,
)
page_indices = kv_to_page_indices(kv_indices, page_size)
if len(page_indices) == 0:
logger.info(
f"Skip sending kv chunk for request {req.rid=} {req.bootstrap_room=} because page_indices is empty"