[perf] reduce radix cache match overhead by changing the match algorithm (#27364)

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``."""

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."""