From cb47a687170d29d250cc71df883593bcf3867755 Mon Sep 17 00:00:00 2001 From: cctry Date: Tue, 14 Jul 2026 10:08:30 -0700 Subject: [PATCH] [PD] Stride KV token->page indices on device before D2H copy (#31173) Co-authored-by: cctry --- python/sglang/srt/disaggregation/decode.py | 31 +++++++-------------- python/sglang/srt/disaggregation/prefill.py | 14 ++++------ python/sglang/srt/mem_cache/common.py | 8 ++---- 3 files changed, 17 insertions(+), 36 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index bfdf49096..d0995088a 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -1076,22 +1076,12 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): "hisparse_page_size", page_size, ) - # Must cast to int32 for ZMQ serialization -- from_zmq reads np.int32. - kv_indices = ( - dst_kv_indices[: origin_input_len - prefix_len] - .cpu() - .numpy() - .astype(np.int32) - ) + kv_indices = dst_kv_indices[: origin_input_len - prefix_len] else: # Only send delta indices (beyond prefix) to prefill. - kv_indices = ( - self.req_to_token_pool.req_to_token[decode_req.req.req_pool_idx][ - total_prefix_len:origin_input_len - ] - .cpu() - .numpy() - ) + kv_indices = self.req_to_token_pool.req_to_token[ + decode_req.req.req_pool_idx + ][total_prefix_len:origin_input_len] seq_len = origin_input_len @@ -1116,9 +1106,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): window_kv_indices_full ) ) - return kv_to_page_indices( - window_kv_indices_swa.cpu().numpy(), page_size - ) + return kv_to_page_indices(window_kv_indices_swa, page_size) def _dsa_payload(): kv_indices_full = self.req_to_token_pool.req_to_token[ @@ -1126,9 +1114,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): ] # Indexer lives on device pool; always use device page_size device_page_size = self.token_to_kv_pool.page_size - return kv_to_page_indices( - kv_indices_full.cpu().numpy(), device_page_size - ) + return kv_to_page_indices(kv_indices_full, device_page_size) def _swa_ring_payload(): # Mirror of prefill _swa_ring_payload using this side's req_pool_idx. @@ -1181,7 +1167,10 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): 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, kv_transfer_page_size) + # int32 for ZMQ serialization -- from_zmq reads np.int32. + page_indices = kv_to_page_indices(kv_indices, kv_transfer_page_size).astype( + np.int32 + ) decode_req.kv_receiver.send_metadata( page_indices, decode_req.metadata_buffer_index, diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 3360c858c..d8036d36f 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -1059,11 +1059,6 @@ class SchedulerDisaggregationPrefillMixin: ) return - kv_indices = ( - self.req_to_token_pool.req_to_token[req.req_pool_idx, start_idx:end_idx] - .cpu() - .numpy() - ) state_indices: Optional[List] = None if last_chunk: self.disagg_metadata_buffers.set_buf(req) @@ -1096,15 +1091,13 @@ class SchedulerDisaggregationPrefillMixin: window_kv_indices_full ) ) - return kv_to_page_indices( - window_kv_indices_swa.cpu().numpy(), page_size - ) + return kv_to_page_indices(window_kv_indices_swa, page_size) def _dsa_payload(): kv_indices_full = self.req_to_token_pool.req_to_token[ req.req_pool_idx, :seq_len ] - return kv_to_page_indices(kv_indices_full.cpu().numpy(), page_size) + return kv_to_page_indices(kv_indices_full, page_size) def _swa_ring_payload(): # Unified_kv SWA ring rows (req_pool_idx*ring_stride + pos%ring_stride) @@ -1157,6 +1150,9 @@ class SchedulerDisaggregationPrefillMixin: else: state_indices.append(None) + kv_indices = self.req_to_token_pool.req_to_token[ + req.req_pool_idx, start_idx:end_idx + ] page_indices = kv_to_page_indices(kv_indices, page_size) if not req.disagg_kv_sender.should_send_kv_chunk(len(page_indices), last_chunk): return diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index e2cedf6ae..902519614 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -48,12 +48,8 @@ MAMBA_STATE_PER_REQ_NO_CACHE = 1 logger = logging.getLogger(__name__) -def kv_to_page_indices(kv_indices: np.ndarray, page_size: int): - # The page is guaranteed to be full except the last page. - if page_size == 1: - return kv_indices - - return kv_indices[::page_size] // page_size +def kv_to_page_indices(kv_indices: torch.Tensor, page_size: int) -> np.ndarray: + return (kv_indices[::page_size] // page_size).cpu().numpy() def kv_to_page_num(num_kv_indices: int, page_size: int):