Speed up assistant span prefix comparisons

scripts/benchmark_assistant_prefix.py

@@ -0,0 +1,136 @@
+"""CPU microbenchmark; no private corpus or timing assertions.
+
+Run: uv run --no-project --python 3.12 python scripts/benchmark_assistant_prefix.py
+The exact production helper AST is loaded without importing the ML runtime.
+"""
+
+import argparse
+import ast
+import hashlib
+import json
+from pathlib import Path
+import platform
+import statistics
+import time
+
+
+def linear(left, right):
+    matched = 0
+    while (
+        matched < len(left) and matched < len(right) and left[matched] == right[matched]
+    ):
+        matched += 1
+    return matched
+
+
+def inputs(kind, size, position):
+    if kind == "tokens":
+        left = [257 + (i * 17) % 100003 for i in range(size)]
+        right = [257 + (i * 17) % 100003 for i in range(size)]
+    else:
+        alphabet = "abcdefghijk" if kind == "ascii" else "α中文🌀́"
+        left = (alphabet * (size // len(alphabet) + 1))[:size]
+        right = (left + "!")[:-1]
+    if position == "proper_prefix":
+        right = right[:-1]
+    elif position != "equal" and size:
+        index = {
+            "early": 0,
+            "near_start": min(3, size - 1),
+            "middle": size // 2,
+            "late": size - 1,
+        }[position]
+        if isinstance(right, list):
+            right[index] = -1009
+        else:
+            right = right[:index] + "!" + right[index + 1 :]
+    return left, right
+
+
+def measure(function, left, right, repeats):
+    iterations = 1
+    while True:
+        start = time.perf_counter()
+        for _ in range(iterations):
+            function(left, right)
+        if time.perf_counter() - start >= 0.003 or iterations >= 16384:
+            break
+        iterations *= 2
+    samples = []
+    for _ in range(repeats):
+        start = time.perf_counter()
+        for _ in range(iterations):
+            function(left, right)
+        samples.append((time.perf_counter() - start) / iterations)
+    return {"median_seconds": statistics.median(samples), "iterations": iterations}
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--sizes", nargs="+", type=int, default=[0, 16, 1024, 65536])
+    parser.add_argument("--repeats", type=int, default=5)
+    args = parser.parse_args()
+    if args.repeats < 1 or any(size < 0 for size in args.sizes):
+        parser.error("sizes must be nonnegative; repeats must be positive")
+    source = Path(__file__).resolve().parents[1] / "src/art/trajectories/_tokenize.py"
+    raw = source.read_bytes()
+    node = next(
+        node
+        for node in ast.parse(raw).body
+        if isinstance(node, ast.FunctionDef) and node.name == "_common_prefix_length"
+    )
+    namespace = {}
+    exec(
+        compile(ast.Module(body=[node], type_ignores=[]), str(source), "exec"),
+        namespace,
+    )
+    candidate = namespace["_common_prefix_length"]
+    rows = []
+    for kind in ("ascii", "unicode", "tokens"):
+        for size in args.sizes:
+            positions = (
+                ("equal",)
+                if size == 0
+                else ("early", "near_start", "middle", "late", "equal", "proper_prefix")
+            )
+            for position in positions:
+                left, right = inputs(kind, size, position)
+                expected = linear(left, right)
+                assert candidate(left, right) == expected
+                order = [("linear", linear), ("candidate", candidate)]
+                if len(rows) % 2:
+                    order.reverse()
+                measurements = {
+                    name: measure(function, left, right, args.repeats)
+                    for name, function in order
+                }
+                baseline, improved = measurements["linear"], measurements["candidate"]
+                rows.append(
+                    dict(
+                        kind=kind,
+                        size=size,
+                        position=position,
+                        prefix=expected,
+                        linear=baseline,
+                        candidate=improved,
+                        speedup=baseline["median_seconds"] / improved["median_seconds"],
+                    )
+                )
+    assert source.read_bytes() == raw, "production source changed during measurement"
+    print(
+        json.dumps(
+            dict(
+                python=platform.python_version(),
+                implementation=platform.python_implementation(),
+                source_sha256=hashlib.sha256(raw).hexdigest(),
+                helper_ast_sha256=hashlib.sha256(ast.dump(node).encode()).hexdigest(),
+                repeats=args.repeats,
+                results=rows,
+            ),
+            indent=2,
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()

src/art/trajectories/_tokenize.py

@@ -170,6 +170,28 @@ def _unique_prompt_suffix_end(prompt: str, rendered: str, *, after: int) -> int
     return after + start + best
 
 
+def _common_prefix_length(left: str | list[int], right: str | list[int]) -> int:
+    """Compare growing blocks natively, then locate the first unequal element."""
+    if not left or not right or left[0] != right[0]:
+        return 0
+    limit = min(len(left), len(right))
+    matched, end = 1, 2
+    while matched < limit:
+        end = min(end, limit)
+        if left[matched:end] != right[matched:end]:
+            high = end - 1
+            while matched < high:
+                middle = (matched + high + 1) // 2
+                if left[matched:middle] == right[matched:middle]:
+                    matched = middle
+                else:
+                    high = middle - 1
+            return matched
+        matched = end
+        end *= 2
+    return matched
+
+
 def _assistant_char_spans(
     messages: list[dict[str, Any]],
     rendered: str,
@@ -190,13 +212,7 @@ def _assistant_char_spans(
             generated = completed[len(prompt) :]
         else:
             prior = render(messages[:message_index], add_generation_prompt=False)
-            shared = 0
-            while (
-                shared < len(prompt)
-                and shared < len(completed)
-                and prompt[shared] == completed[shared]
-            ):
-                shared += 1
+            shared = _common_prefix_length(prompt, completed)
             anchored = not (
                 shared < len(prior)
                 or prompt[: len(prior)] != prior
@@ -249,13 +265,7 @@ def _assistant_char_spans(
                 continue
         if not generated:
             continue
-        start = 0
-        while (
-            start < len(prompt)
-            and start < len(rendered)
-            and prompt[start] == rendered[start]
-        ):
-            start += 1
+        start = _common_prefix_length(prompt, rendered)
         retained_start = next(
             (
                 offset