diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index e20f83b4a..6b8ac9cfa 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -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 diff --git a/test/registered/unit/disaggregation/test_cp_rank_kv_filter.py b/test/registered/unit/disaggregation/test_cp_rank_kv_filter.py new file mode 100644 index 000000000..ec831ebfb --- /dev/null +++ b/test/registered/unit/disaggregation/test_cp_rank_kv_filter.py @@ -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()