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

87
test/registered/unit/disaggregation/test_cp_rank_kv_filter.py Normal file
View File

@@ -0,0 +1,87 @@
"""Differential test: filter_kv_indices_for_cp_rank range mask vs old np.isin.
The old implementation computed rank_page_indices = kv_indices[range_mask]
via page_indices_to_cp_rank_page_indices and then np.isin(kv_indices,
rank_page_indices). For any kv_indices, a value is in the filtered subset iff
it passes the same range test, so np.isin was provably identical to the range
mask. The new implementation applies the range mask directly. This test pins
the equivalence against a verbatim copy of the old logic.
"""
import random
import unittest
from types import SimpleNamespace
import numpy as np
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=2, suite="stage-a-test-cpu")
from sglang.srt.disaggregation.utils import (
filter_kv_indices_for_cp_rank,
page_indices_to_cp_rank_page_indices,
)
def _old_filter(kv_mgr, kv_indices, index_slice):
"""Verbatim copy of the pre-refactor implementation."""
total_pages = len(kv_indices)
rank_page_indices = page_indices_to_cp_rank_page_indices(
page_indices=kv_indices,
total_pages=total_pages,
cp_rank=kv_mgr.attn_cp_rank,
cp_size=kv_mgr.attn_cp_size,
)
if rank_page_indices.size == 0:
return kv_indices[:0], slice(index_slice.start, index_slice.start)
mask = np.isin(kv_indices, rank_page_indices)
if not mask.any():
return kv_indices[:0], slice(index_slice.start, index_slice.start)
first_pos = int(mask.argmax())
last_pos = len(mask) - int(mask[::-1].argmax())
return kv_indices[first_pos:last_pos], slice(
index_slice.start + first_pos, index_slice.start + last_pos
)
class TestCpRankKvFilterEquivalence(unittest.TestCase):
def _check(self, kv_indices, cp_rank, cp_size, start=7):
mgr = SimpleNamespace(attn_cp_rank=cp_rank, attn_cp_size=cp_size)
sl = slice(start, start + len(kv_indices))
old_kv, old_slice = _old_filter(mgr, kv_indices, sl)
new_kv, new_slice = filter_kv_indices_for_cp_rank(mgr, kv_indices, sl)
np.testing.assert_array_equal(old_kv, new_kv)
self.assertEqual(old_slice, new_slice)
def test_contiguous_pages_all_ranks(self):
for total in (1, 2, 7, 8, 64, 513):
kv = np.arange(1000, 1000 + total, dtype=np.int64)
for cp_size in (1, 2, 3, 8):
for cp_rank in range(cp_size):
self._check(kv, cp_rank, cp_size)
def test_random_unique_pages_all_ranks(self):
rng = random.Random(11)
for _ in range(30):
total = rng.randrange(1, 300)
vals = rng.sample(range(100000), total)
kv = np.asarray(vals, dtype=np.int64)
cp_size = rng.choice([2, 4, 8])
for cp_rank in range(cp_size):
self._check(kv, cp_rank, cp_size)
def test_empty(self):
self._check(np.asarray([], dtype=np.int64), cp_rank=1, cp_size=8)
def test_cp_size_one_passthrough(self):
kv = np.arange(5, dtype=np.int64)
mgr = SimpleNamespace(attn_cp_rank=0, attn_cp_size=1)
sl = slice(3, 8)
new_kv, new_slice = filter_kv_indices_for_cp_rank(mgr, kv, sl)
np.testing.assert_array_equal(new_kv, kv)
self.assertEqual(new_slice, sl)
if __name__ == "__main__":
unittest.main()