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:
|
||||
|
||||
Reference in New Issue
Block a user