bpe_emit_piece merge loop is O(n²): large prompts take minutes to tokenize
- Lingua principale
- C
- Stelle
- 22.3k
- Fork
- 2.1k
- Merge medio
- 1g 3h
- PR unite (30g)
- 4
Descrizione
# 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:
```c
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×.
Guida per i contributori
Apri la guida per i contributori
Valutazione
Questa issue non è ancora stata valutata.