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

Drop redundant np.isin from CP-rank KV page filter

Browse Source 
filter_kv_indices_for_cp_rank built rank_page_indices =
kv_indices[range_mask] and then ran np.isin(kv_indices,
rank_page_indices) — provably identical to range_mask itself (a value
is in the filtered subset iff it passes the same range test), at the
cost of an extra sort per chunk send under
SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER=1. Apply the range mask
directly: 30.5 -> 8.1 us per 1024-page chunk. Differential test pins
equivalence against a verbatim copy of the old logic.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
孙远航
2026-06-10 05:15:16 +00:00
parent b742371fe3
commit 5bec69f40a
2 changed files with 117 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

49
python/sglang/srt/disaggregation/utils.py
View File

@@ -679,30 +679,41 @@ def filter_kv_indices_for_cp_rank(
cp_rank = kv_mgr.attn_cp_rank
cp_size = kv_mgr.attn_cp_size
rank_page_indices = page_indices_to_cp_rank_page_indices(
page_indices=kv_indices,
total_pages=total_pages,
cp_rank=cp_rank,
cp_size=cp_size,
)
if cp_size <= 1:
return kv_indices, index_slice
if total_pages == 0:
return kv_indices[:0], slice(index_slice.start, index_slice.start)
if rank_page_indices.size == 0:
# This rank owns the global page-id range [start_page, end_page); see
# page_indices_to_cp_rank_page_indices. Test the range directly: the old
# np.isin(kv_indices, kv_indices[range_mask]) was identical to range_mask
# (a value is in the filtered subset iff it passes the same range test)
# but cost an extra O(N log N) sort per chunk send.
first_page = int(kv_indices.min())
base = total_pages // cp_size
rem = total_pages % cp_size
if rem == 0:
local_start = cp_rank * base
local_end = local_start + base
else:
local_start = cp_rank * base + min(cp_rank, rem)
local_end = local_start + base + (1 if cp_rank < rem else 0)
start_page = first_page + local_start
end_page = first_page + local_end
mask = (kv_indices >= start_page) & (kv_indices < end_page)
if not mask.any():
new_kv_indices = kv_indices[:0]
new_index_slice = slice(index_slice.start, index_slice.start)
else:
mask = np.isin(kv_indices, rank_page_indices)
if not mask.any():
new_kv_indices = kv_indices[:0]
new_index_slice = slice(index_slice.start, index_slice.start)
else:
first_pos = int(mask.argmax())
last_pos = len(mask) - int(mask[::-1].argmax())
first_pos = int(mask.argmax())
last_pos = len(mask) - int(mask[::-1].argmax())
new_kv_indices = kv_indices[first_pos:last_pos]
new_index_slice = slice(
index_slice.start + first_pos,
index_slice.start + last_pos,
)
new_kv_indices = kv_indices[first_pos:last_pos]
new_index_slice = slice(
index_slice.start + first_pos,
index_slice.start + last_pos,
)
return new_kv_indices, new_index_slice