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

[PD] Stride KV token->page indices on device before D2H copy (#31173)

Browse Source 
Co-authored-by: cctry <cctry@fb.com>
This commit is contained in:
cctry
2026-07-14 10:08:30 -07:00
committed by GitHub
co-authored by cctry
parent a5c3e0283f
commit cb47a68717
3 changed files with 17 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/decode.py
+10 -21
View File
@@ -1076,22 +1076,12 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
"hisparse_page_size", "hisparse_page_size",
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]
kv_indices = (
dst_kv_indices[: origin_input_len - prefix_len]
.cpu()
.numpy()
.astype(np.int32)
)
else: else:
# Only send delta indices (beyond prefix) to prefill. # Only send delta indices (beyond prefix) to prefill.
kv_indices = ( kv_indices = self.req_to_token_pool.req_to_token[
self.req_to_token_pool.req_to_token[decode_req.req.req_pool_idx][ decode_req.req.req_pool_idx
total_prefix_len:origin_input_len ][total_prefix_len:origin_input_len]
]
.cpu()
.numpy()
)
seq_len = origin_input_len seq_len = origin_input_len
@@ -1116,9 +1106,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
window_kv_indices_full window_kv_indices_full
) )
) )
return kv_to_page_indices( return kv_to_page_indices(window_kv_indices_swa, page_size)
window_kv_indices_swa.cpu().numpy(), page_size
)
def _dsa_payload(): def _dsa_payload():
kv_indices_full = self.req_to_token_pool.req_to_token[ 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 # Indexer lives on device pool; always use device page_size
device_page_size = self.token_to_kv_pool.page_size device_page_size = self.token_to_kv_pool.page_size
return kv_to_page_indices( return kv_to_page_indices(kv_indices_full, device_page_size)
kv_indices_full.cpu().numpy(), device_page_size
)
def _swa_ring_payload(): def _swa_ring_payload():
# Mirror of prefill _swa_ring_payload using this side's req_pool_idx. # 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() self.req_to_metadata_buffer_idx_allocator.alloc()
) )
assert decode_req.metadata_buffer_index is not None 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( decode_req.kv_receiver.send_metadata(
page_indices, page_indices,
decode_req.metadata_buffer_index, decode_req.metadata_buffer_index,
python/sglang/srt/disaggregation/prefill.py
+5 -9
View File
@@ -1059,11 +1059,6 @@ class SchedulerDisaggregationPrefillMixin:
) )
return 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 state_indices: Optional[List] = None
if last_chunk: if last_chunk:
self.disagg_metadata_buffers.set_buf(req) self.disagg_metadata_buffers.set_buf(req)
@@ -1096,15 +1091,13 @@ class SchedulerDisaggregationPrefillMixin:
window_kv_indices_full window_kv_indices_full
) )
) )
return kv_to_page_indices( return kv_to_page_indices(window_kv_indices_swa, page_size)
window_kv_indices_swa.cpu().numpy(), page_size
)
def _dsa_payload(): def _dsa_payload():
kv_indices_full = self.req_to_token_pool.req_to_token[ kv_indices_full = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :seq_len 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(): def _swa_ring_payload():
# Unified_kv SWA ring rows (req_pool_idx*ring_stride + pos%ring_stride) # Unified_kv SWA ring rows (req_pool_idx*ring_stride + pos%ring_stride)
@@ -1157,6 +1150,9 @@ class SchedulerDisaggregationPrefillMixin:
else: else:
state_indices.append(None) 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) 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): if not req.disagg_kv_sender.should_send_kv_chunk(len(page_indices), last_chunk):
return return
python/sglang/srt/mem_cache/common.py
+2 -6
View File
@@ -48,12 +48,8 @@ MAMBA_STATE_PER_REQ_NO_CACHE = 1
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
def kv_to_page_indices(kv_indices: np.ndarray, page_size: int): def kv_to_page_indices(kv_indices: torch.Tensor, page_size: int) -> np.ndarray:
# The page is guaranteed to be full except the last page. return (kv_indices[::page_size] // page_size).cpu().numpy()
if page_size == 1:
return kv_indices
return kv_indices[::page_size] // page_size
def kv_to_page_num(num_kv_indices: int, page_size: int): def kv_to_page_num(num_kv_indices: int, page_size: int):