Speed up assistant span prefix comparisons

[bradhilton]

Long conversation histories repeatedly scan prompt prefixes in Python while locating assistant-generated spans. Replace those two element-by-element loops with a private helper that compares growing blocks using native slice equality, then searches within the first unequal block. On the complete canonical validation corpus, tensorization averages 250.9 → 140.0 seconds (44.2% shorter) with byte-identical outputs.

The helper handles strings and integer-token lists. An empty/first-mismatch fast path avoids copying a long history when the answer is immediately known. Temporary slices grow with the compared prefix; there is no persistent cache. Rendering, span fallbacks, selected tokens, flags, log-probabilities and error paths retain their existing behavior. The runtime diff is 24 added/14 removed lines, with no dependency or public API change.

Repeated profiling at base 077abf6fa7a3a51f9ff255c37f8886c6c8256082 and candidate 45eaed3259e179ca9068dd0851723a643536a85f, in A/B/B/A order:

Run	Source	Tensorization
A1	Base	251.214 s
B1	Candidate	140.745 s
B2	Candidate	139.225 s
A2	Base	250.516 s

This saves 110.880 seconds on average, a 1.792× speedup. Every run retained all 128 sources and 256 real/canonical-generated views, including the 68 recorded model-error returns. All 768 ordered token/flag/log-probability fields match the fresh base byte for byte, including shapes, dtypes and NaN bit patterns. Historical frozen tensor parity also passed separately. Token counts and template/encoding/span call counts are unchanged.

CPU profiling localizes the savings to span finding. These are summed worker-thread CPU seconds, averaged across both runs; each row excludes its instrumented child stages:

Stage	Calls/run	Base CPU	Candidate CPU
Assistant span finding	3,524	152.808 s	42.887 s
Chat template rendering	99,660	58.296 s	56.990 s
Encoding	103,440	84.198 s	82.271 s

Total worker conversion CPU falls from 325.382 to 211.631 seconds. The campaign used fresh sequential processes, the same four concurrent work slots and timing instrumentation, frozen captured inputs and an offline tokenizer. Imports, capture loading and parity-artifact work are excluded from the tensorization timer. All four runs finished within the fixed 20-minute campaign, and all 16 owned process identities were independently verified absent.

Validation on the final revision:

• 515 local trajectory tests pass: 458 existing plus 57 new deterministic tests. Coverage includes differential string/list oracles, Unicode, mismatch/truncation boundaries, both span call sites, suffix/removal fallbacks, overlap/unmappable errors, input preservation and full text/token-ID tokenization flags.
• Independent differential review passed 11,907 final-helper cases. Removing the helper and restoring the two original loops reproduces the original module AST exactly.
• Scoped ty, Ruff, formatting, whitespace and offline lock checks pass. CI is green: 773 Megatron lightweight tests and 1,263 unit tests pass, with 49 unit tests skipped.
• An earlier local suite attempt reached its 240-second external cap; that incomplete outcome is retained. The final 458-test rerun passed in 167.72 seconds. Its observed subprocess/filesystem waits do not prove the earlier timeout's cause.

McCarthy's independent source review is clear: 515 tests pass at this head, 557 pass with the held tokenizer composition, and 37,468 helper differential executions pass. The reviewer also rehashed the complete profiling evidence and tensor artifacts.

The committed standalone microbenchmark loads the exact production helper AST without importing the ML runtime, validates results before timing and records source hashes, iterations and medians. It needs no private corpus:

uv run --no-project --python 3.12 python scripts/benchmark_assistant_prefix.py \
  --sizes 0 16 1024 65536 --repeats 7

These are two runs per variant on one shared host and one captured corpus, not a general latency guarantee or a new end-to-end inference/probe qualification. Microbenchmarks show that very short nonzero prefixes can incur approximately 0.2–0.3 µs additional overhead; long shared prefixes benefit substantially. There are no timing assertions in tests.

Evidence on the shared machine: /home/brad/.local/share/thanos/art-prefix-optimization-20260910/, including profiling/REPORT.md, comparison.json, per-view/stage timings, source/input manifests, raw tensor comparisons and cleanup receipts. Canonical checkpoint: wandb-artifact:///wandb/049-retail49-program/r49-program-cmapping-he8cb0aa602dcc0b9:v1, step 31, sources 0–127; corpus manifest SHA256 e21f304bab8b76bc14442b2839148417f0327f1cdfdc143192e43cb376c25381. Final profiling manifest SHA256 a6dfbcda3c71eab9987825f34a3db5b25f7df1107017cb1a4477949f6bba6429 verifies all 822 evidence files.

Maintained 049 adoption is separate: its reviewed raw G/V helper correctly refuses the changed tokenizer source hash and moved exception line. It needs a separately reviewed helper revision before updating the ART pin; this PR does not bypass that fence.

[bradhilton]

McCarthy: CLEAR at 45eaed3259e179ca9068dd0851723a643536a85f against 077abf6fa7a3a51f9ff255c37f8886c6c8256082, CPU/source scope. No blocking finding.

The helper preserves the common-prefix result for the built-in strings and integer lists used here: growing blocks advance a proven prefix; a differing block bounds the answer and the inner search strictly narrows that interval. Empty, immediate-mismatch and proper-prefix cases are correct; inputs remain unchanged. Native slices still allocate O(prefix) temporary storage, and short prefixes need not become faster.

Independent checks: 515 trajectory tests pass at this exact head; 557 pass after applying this patch to held ART 749938576aa7baf78ac295bc15bdd5e26a9a45ce (#880/#882/#885/#886). 37,468 differential executions cover exhaustive short Unicode/integer-list pairs, seeded large integers, truncation and power-of-two boundaries, plus non-mutation and geometric slice-work controls. Reversing just the helper and two substitutions recovers the entire original module AST on both base and composition. Ruff/format clean.

049 adoption boundary: the composed tokenizer hashes to 1917cb970fc88f9a9227bc6dc119b8429e1b8350e31e5eacd203dcfad2d76c88; its native mask raise moves from line 491 to 501. I executed the reviewed raw G/V helper 89193ba6 against it: the existing source fence correctly refuses before AST parsing. A separately sealed/reviewed helper revision must establish the new source and exception-location contracts before adoption. This source clearance does not override that fence or authorize an ART pin update.

Profiling audit: both 367-file source archives match Git. Rehashed the final 822-file manifest and all four actual tensor artifacts: 768 fields / 256 ordered views / 128 sources, including all 68 recorded model-error returns, match byte-for-byte. A/B/B/A CPU tensorization is 251.214 / 140.745 / 139.225 / 250.516 s; all four supervisor receipts finish before the unchanged 8:19:29 AM Mountain cap and record cleanup. I audited artifacts rather than independently replaying the full corpus. The 44.20% reduction is this shared-host CPU-stage comparison, not an end-to-end or GPU guarantee. The prior 240-second incomplete test run and unproved cause remain preserved. Hosted quality checks were still running at my final read.

Evidence: /home/brad/.local/share/mccarthy/art887-review-20260910/REVIEW.md, SHA256 3abee134c80fc41e599e00e690d39df436b623c5e8b50d34619c5319e7a76091. No merge, deployment or live resource action performed.

[bradhilton]

Minsky review — CLEAR at 45eaed3259e179ca9068dd0851723a643536a85f (base 077abf6f), CPU/source scope. Second dedicated pass-off; McCarthy cleared this same SHA.

Finding: the runtime delta is one private helper, _common_prefix_length, plus its two substitutions for the element-wise while loops in _assistant_char_spans (prompt-vs-completed shared, prompt-vs-rendered start). The helper compares doubling slices natively and, on the first unequal block, binary-searches for the largest middle with equal left[matched:middle] == right[matched:middle]. Invariant left[:matched] == right[:matched] holds throughout; high = end − 1 is valid because the block [matched:end) is known to differ; the upper-midpoint search terminates with matched == high = the common prefix length; end = min(end, limit) keeps slices in range; empty inputs or a first-element mismatch return 0 and full-prefix cases return min(len), exactly as the loops did. Slicing is code-point based for str and element based for list[int], matching the original indexing. Nothing else in the module changes (reversing the two substitutions recovers the original, as McCarthy also verified).

Verification at this exact head (ART venv):

• tests/unit/trajectories: 515 passed. Ruff check and format clean on the three changed files.
• Negative control: the new test_assistant_prefix.py against 077abf6f → 54 failed / 3 passed (the helper does not exist there; the three passing cases are structural controls), so the 57 new tests target the change.
• Own differential against a linear oracle: exhaustive all pairs over a 3-symbol alphabet up to length 6 for both lists and strings, 60,000 randomized pairs with shared-prefix lengths placed at and around every power of two up to 1025 plus random lengths to 5,000 using the alphabet a b \0 é 雪 🙂 𐀀 and small ints, and identical/prefix cases with 2^65-magnitude ints up to 4,097 elements: 2,509,316 checks, all equal.
• Profiling receipts (Thanos, independently re-hashed by McCarthy): A/B/B/A 251.214 / 140.745 / 139.225 / 250.516 s on the captured 128-source corpus with all 768 tensor fields byte-identical; I read and did not rerun them. This is captured-corpus CPU tensorization evidence, not an end-to-end or fresh-inference claim.

Composition boundary (unchanged from McCarthy's note): this changes _tokenize.py bytes and moves the mask raise site, so the maintained 049 raw G/V helper's byte/AST fence correctly refuses the composed source; adoption there needs a separately sealed and reviewed helper revision. No ART pin change is part of this PR.

Non-draft; CI green per the thread. Merge/adoption remain Brad's decision.