diff --git a/python/sglang/srt/mem_cache/utils.py b/python/sglang/srt/mem_cache/utils.py index 7b9e8c87f..168090b56 100644 --- a/python/sglang/srt/mem_cache/utils.py +++ b/python/sglang/srt/mem_cache/utils.py @@ -14,6 +14,7 @@ """Common utilities.""" import logging +from itertools import islice from typing import Any, List, Optional, Tuple import torch @@ -274,7 +275,10 @@ def _python_convert_to_bigram_key(tokens: List[int]) -> List[Tuple[int, int]]: return tokens if len(tokens) < 2: return [] - return [(tokens[i], tokens[i + 1]) for i in range(len(tokens) - 1)] + # zip runs the pairing loop in C; islice avoids copying the list for the + # shifted operand. This is called per match_prefix/insert with EAGLE, on + # the full key length, so constant factors matter. + return list(zip(tokens, islice(tokens, 1, None))) try: diff --git a/test/registered/unit/mem_cache/test_radix_key_views.py b/test/registered/unit/mem_cache/test_radix_key_views.py index 3cd1d68c2..4738cea24 100644 --- a/test/registered/unit/mem_cache/test_radix_key_views.py +++ b/test/registered/unit/mem_cache/test_radix_key_views.py @@ -44,6 +44,20 @@ from sglang.srt.mem_cache.radix_cache import ( ) +class TestBigramConversionEquivalence(unittest.TestCase): + def test_zip_impl_matches_index_comprehension(self): + from sglang.srt.mem_cache.utils import _python_convert_to_bigram_key + + rng = random.Random(7) + for n in (0, 1, 2, 3, 17, 256): + tokens = [rng.randrange(1000) for _ in range(n)] + want = [(tokens[i], tokens[i + 1]) for i in range(max(len(tokens) - 1, 0))] + self.assertEqual(_python_convert_to_bigram_key(tokens), want) + # Already-converted (tuple) input is returned unchanged. + pairs = [(1, 2), (2, 3)] + self.assertEqual(_python_convert_to_bigram_key(pairs), pairs) + + class TestRadixKeyViewSemantics(unittest.TestCase): def test_slice_is_view_sharing_backing(self): backing = list(range(100))