bpe_emit_piece merge loop is O(n²): large prompts take minutes to tokenize

[graysonren-dot]

Bug report: bpe_emit_piece merge loop is O(n²) — large prompts take minutes to tokenize

Repo: antirez/ds4
Commit tested: 84cc882 (current main)
Hardware: Mac Studio M3 Ultra / 512GB, Metal backend

Symptom

Any request with a large prompt (e.g. 24k tokens / 39k chars) to ds4-server
spends 175–250 seconds of near-100% CPU before the prefill even starts
(the live kv cache miss / prompt start log lines appear minutes after the
request arrives). Small prompts are unaffected, so this looks like a "stall",
but it is tokenization.

Root cause

In bpe_emit_piece (ds4.c:37174) the byte-level BPE merge loop rescans all
adjacent symbol pairs on every merge iteration:

for (;;) {
    int best_i = -1;
    int best_rank = INT32_MAX;
    for (int i = 0; i + 1 < n_sym; i++) {   /* full rescan each merge */
        int rank = bpe_rank(vocab, &sym[i], &sym[i + 1]);
        if (rank >= 0 && rank < best_rank) { best_rank = rank; best_i = i; }
    }
    if (best_i < 0) break;
    /* merge sym[best_i] and sym[best_i+1], shift array, n_sym-- */
}

Merging n symbols this way is O(n²). Real CJK prompts are single long pieces
after the pre-tokenizer, so n is large and the quadratic term dominates.

Measurements

ds4 --dump-tokens --prompt-file on the same model (Q4KExperts 0731):

Input	Tokens	Tokenize time (excl. model load)
39 000 chars	~24 000	~90–180 s
4 000 chars	~2 400	~1.1 s

10× the text ⇒ 80×+ the time ⇒ clearly super-linear (O(n²)).

End-to-end effect on ds4-server (24k-token chat requests): request arrives →
~3–4 min of CPU → then prompt start → normal prefill (~50 s).

Suggested fix

Only the two adjacent pairs around a merge point can change rank after a merge,
so recompute just those and keep a min-heap (or equivalent) of pair ranks
instead of rescanning all pairs each iteration: O(n log n) per piece. Must be
validated against the official test vectors (tests/test-vectors/) since any
change in merge order changes tokenization.

Workaround (no code change)

Keep prompts small: tokenization cost scales quadratically, so halving prompt
size cuts tokenization ~4×.