Speed up Python bigram key conversion with C-level zip
The EAGLE bigram fallback built 64K-token keys with an index-based Python comprehension (6.3 ms per call at 64K tokens). zip + islice runs the pairing loop in C and avoids copying the shifted operand: 3.6-4.2 ms per call (~1.6x). Output is byte-identical (same tuples); the tai-kernel fast path is unaffected. Packing bigrams into int64s via numpy was measured and rejected: the list<->array boxing makes it slower (4.1 ms) than zip until token ids are numpy end-to-end, and it would change HiCache storage hash inputs. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -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:
|
||||
|
||||
@@ -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))
|
||||
|
||||
Reference in New Issue
Block a user