From 4b0f6290822da0f503790d94eac7953a0d25fdc1 Mon Sep 17 00:00:00 2001 From: Qiaolin Yu Date: Sat, 6 Jun 2026 15:40:28 -0700 Subject: [PATCH] [perf] reduce radix cache match overhead by changing the match algorithm (#27364) --- python/sglang/srt/mem_cache/radix_cache.py | 47 ++++++++-------- .../unit/mem_cache/test_radix_cache_unit.py | 54 +++++++++++++++++++ 2 files changed, 79 insertions(+), 22 deletions(-) diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 65c5b70e4..e7d5bade0 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -135,37 +135,40 @@ class RadixKey: f"{self.extra_key=} != {other.extra_key=}" ) - # TODO(Jialin): replace zip with numpy to skip per-element PyLong boxing def match(self, other: "RadixKey", page_size: int = 1) -> int: """Logical-unit prefix length shared with ``other``. Result is rounded down to ``page_size``.""" self._check_compatible(other) t0, t1 = self.token_ids, other.token_ids + assert type(t0) is type(t1), (type(t0), type(t1)) + n = min(len(t0), len(t1)) + + # Exponential search for the first diverging token: gallop in doubling + # windows (one C-level slice compare each), then binary-search the window + # holding the divergence -- no per-token Python loop on long shared prefixes. + matched_tokens = n + lo = 0 + step = 1 + while lo < n: + hi = lo + step if lo + step < n else n + if t0[lo:hi] != t1[lo:hi]: + while hi - lo > 1: + mid = (lo + hi) // 2 + if t0[lo:mid] == t1[lo:mid]: + lo = mid + else: + hi = mid + matched_tokens = lo + break + lo = hi + step *= 2 if self.is_bigram: - # Walk raw tokens; L matching tokens imply L-1 matching bigrams. - i = 0 - for a, b in zip(t0, t1): - if a != b: - break - i += 1 - matched = max(0, min(i - 1, len(self), len(other))) + matched = max(0, min(matched_tokens - 1, len(self), len(other))) return (matched // page_size) * page_size if page_size > 1 else matched if page_size == 1: - i = 0 - for a, b in zip(t0, t1): - if a != b: - break - i += 1 - return i - - min_len = min(len(self), len(other)) - i = 0 - while i < min_len: - if t0[i : i + page_size] != t1[i : i + page_size]: - break - i += page_size - return i + return matched_tokens + return (matched_tokens // page_size) * page_size def child_key(self, page_size: int = 1): """Hashable dict-key for the first ``page_size`` logical units, namespaced by ``extra_key``.""" diff --git a/test/registered/unit/mem_cache/test_radix_cache_unit.py b/test/registered/unit/mem_cache/test_radix_cache_unit.py index 263670fb7..c09997436 100644 --- a/test/registered/unit/mem_cache/test_radix_cache_unit.py +++ b/test/registered/unit/mem_cache/test_radix_cache_unit.py @@ -142,6 +142,60 @@ class TestRadixKey(unittest.TestCase): repr_str = repr(key) self.assertIn("...", repr_str) # Should be truncated + def _assert_match(self, a, b, page_size, expected, is_bigram=False): + key_a = RadixKey(array("q", a), is_bigram=is_bigram) + key_b = RadixKey(array("q", b), is_bigram=is_bigram) + self.assertEqual(key_a.match(key_b, page_size=page_size), expected) + + def test_match_page_size_1(self): + """match() with page_size=1: full, partial, none, prefix, and empty keys.""" + self._assert_match([1, 2, 3, 4], [1, 2, 3, 4], 1, 4) # identical + self._assert_match([1, 2, 3, 4], [1, 2, 9, 9], 1, 2) # diverge at index 2 + self._assert_match([9, 2, 3], [1, 2, 3], 1, 0) # diverge at index 0 + self._assert_match([1, 2, 3, 4], [1, 2, 3], 1, 3) # other is a prefix + self._assert_match([], [1, 2], 1, 0) # empty self + self._assert_match([1, 2], [], 1, 0) # empty other + self._assert_match([], [], 1, 0) # both empty + + def test_match_page_size_gt_1_rounds_down(self): + """match() with page_size>1 rounds the shared length down to a page.""" + self._assert_match([1, 2, 3, 4, 5, 6, 7, 8], [1, 2, 3, 4, 5, 6, 9, 8], 4, 4) + self._assert_match( + [1, 2, 3, 4], [1, 9, 3, 4], 4, 0 + ) # diverge inside first page + self._assert_match([1, 2, 3, 4, 5, 6, 7, 8], [1, 2, 3, 4, 9, 6, 7, 8], 4, 4) + self._assert_match([1, 2, 3, 4, 5, 6, 7, 8], [1, 2, 3, 4, 5, 6, 7, 8], 4, 8) + self._assert_match([1, 2, 3], [1, 2, 3], 4, 0) # shorter than one page + + def test_match_long_keys_exponential_search(self): + """Deep divergences exercise the doubling gallop windows + binary search. + + ``base`` has distinct values, so flipping one position diverges the prefix + exactly there; the shared length is that index rounded down to the page. + """ + base = list(range(2000)) + for div in (1, 2, 63, 64, 65, 127, 128, 511, 512, 513, 1234, 1999): + b = base[:] + b[div] = -1 + for page_size in (1, 4, 64): + with self.subTest(div=div, page_size=page_size): + self._assert_match( + base, b, page_size, (div // page_size) * page_size + ) + # Full match of a long key: the gallop must reach the end. + self._assert_match(base, base[:], 64, (2000 // 64) * 64) + + def test_match_bigram(self): + """is_bigram: L matching raw tokens imply L-1 matching bigrams.""" + self._assert_match([1, 2, 3, 4, 5], [1, 2, 3, 9, 5], 1, 2, is_bigram=True) + self._assert_match([1, 2, 3, 4, 5], [1, 2, 3, 4, 5], 1, 4, is_bigram=True) + self._assert_match([1, 2], [1, 2], 1, 1, is_bigram=True) + # Raw diverge at token 70 -> 69 matching bigrams -> rounded down to 64. + long_a = list(range(130)) + long_b = list(range(130)) + long_b[70] = -1 + self._assert_match(long_a, long_b, 64, 64, is_bigram=True) + class TestTreeNode(unittest.TestCase): """Test cases for TreeNode class."""