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:
2026-06-10 04:30:59 +00:00
parent 8979c81a22
commit 0d065a8ab0
2 changed files with 19 additions and 1 deletions

View File

@@ -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:

View File

@@ -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))