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: