From 9929e93942f20467980fc5fc008575c317c410c4 Mon Sep 17 00:00:00 2001 From: Uri Date: Sun, 5 Jul 2026 21:53:48 +0300 Subject: [PATCH] perf(js): replace quadratic BPE merge algorithm with optimal heap-based algorithm --- js/src/core.ts | 248 +++++++++++++++++++++++++++++++++++-------------- 1 file changed, 176 insertions(+), 72 deletions(-) diff --git a/js/src/core.ts b/js/src/core.ts index d09d6cde..13516e7c 100644 --- a/js/src/core.ts +++ b/js/src/core.ts @@ -2,37 +2,195 @@ import base64 from "base64-js"; import type { TiktokenModel } from "./ranks/ranks"; import { never } from "./utils"; +type BPEMergeNode = { + listNext: BPEMergeNode | null; + listPrev: BPEMergeNode | null; + deleted: boolean; + rank: number; + start: number; + end: number; +}; + +type HeapEntry = { + node: BPEMergeNode; + rank: number; + start: number; +}; + +function compareEntry(a: HeapEntry, b: HeapEntry) { + return a.rank - b.rank || a.start - b.start; +} + +// Helper function to swap elements at two indices +function swap(heap: HeapEntry[], i: number, j: number) { + const temp = heap[i]; + heap[i] = heap[j]; + heap[j] = temp; +} + +// standard binary heap push, generated by gpt4 +function heapPush(heap: HeapEntry[], entry: HeapEntry) { + heap.push(entry); // Add the new element to the end + let currentIndex = heap.length - 1; + let parentIndex = Math.floor((currentIndex - 1) / 2); + + // Bubble the new element up to its correct position + while ( + currentIndex > 0 && + compareEntry(heap[currentIndex], heap[parentIndex]) < 0 + ) { + swap(heap, currentIndex, parentIndex); + currentIndex = parentIndex; + parentIndex = Math.floor((currentIndex - 1) / 2); + } +} + +// standard heap pop, also ai generated +function heapPop(heap: HeapEntry[]) { + if (heap.length === 0) { + return undefined; // Return undefined if the heap is empty + } + + const rootValue = heap[0]; // The root element to return + const lastValue = heap.pop(); // Remove the last element + + if (heap.length > 0 && lastValue) { + heap[0] = lastValue; // Move the last element to the root + let currentIndex = 0; + + // Bubble down the new root element to its correct position + while (true) { + let leftChildIndex = 2 * currentIndex + 1; + let rightChildIndex = 2 * currentIndex + 2; + let smallestIndex = currentIndex; + + if ( + leftChildIndex < heap.length && + compareEntry(heap[leftChildIndex], heap[smallestIndex]) < 0 + ) { + smallestIndex = leftChildIndex; + } + + if ( + rightChildIndex < heap.length && + compareEntry(heap[rightChildIndex], heap[smallestIndex]) < 0 + ) { + smallestIndex = rightChildIndex; + } + + if (smallestIndex !== currentIndex) { + swap(heap, currentIndex, smallestIndex); + currentIndex = smallestIndex; + } else { + break; + } + } + } + + return rootValue; +} + function bytePairMerge( piece: Uint8Array, ranks: Map ): Array<{ start: number; end: number }> { - let parts: Array<{ start: number; end: number }> = Array.from( + const parts: BPEMergeNode[] = Array.from( { length: piece.length }, - (_, i) => ({ start: i, end: i + 1 }) + (_, i) => ({ + start: i, + end: i + 1, + rank: Infinity, + deleted: false, + listNext: null, + listPrev: null, + }) ); - while (parts.length > 1) { - let minRank: [number, number] | null = null; + if (parts.length === 0) { + return []; + } - for (let i = 0; i < parts.length - 1; i++) { - const slice = piece.slice(parts[i].start, parts[i + 1].end); - const rank = ranks.get(slice.join(",")); - if (rank == null) continue; + const head = parts[0]; + for (let i = 0; i < parts.length; ++i) { + parts[i].listPrev = parts[i - 1] ?? null; + parts[i].listNext = parts[i + 1] ?? null; + } - if (minRank == null || rank < minRank[0]) { - minRank = [rank, i]; - } + const heap: HeapEntry[] = []; + for (let i = 0; i < parts.length - 1; ++i) { + const slice = piece.slice(parts[i].start, parts[i + 1].end); + const rank = ranks.get(slice.join(",")); + if (rank == null) continue; + const part = parts[i]; + part.rank = rank; + heapPush(heap, { node: part, rank, start: part.start }); + } + + while (heap.length > 0) { + const entry = heapPop(heap); + if (!entry) break; + + const { node, rank, start } = entry; + + // Discard stale or deleted entries + if (node.deleted || node.rank !== rank || node.start !== start) { + continue; + } + if (!node.listNext || node.listNext.deleted) { + continue; + } + + // Verify the merge is still valid + const currentSlice = piece.slice(node.start, node.listNext.end); + const currentRank = ranks.get(currentSlice.join(",")); + if (currentRank !== rank) { + continue; + } + + // Perform merge + node.end = node.listNext.end; + node.listNext.deleted = true; + node.listNext = node.listNext.listNext; + if (node.listNext) { + node.listNext.listPrev = node; } - if (minRank != null) { - const i = minRank[1]; - parts[i] = { start: parts[i].start, end: parts[i + 1].end }; - parts.splice(i + 1, 1); + // Check for new possible merges with listNext + if (node.listNext) { + const slice = piece.slice(node.start, node.listNext.end); + const nextRank = ranks.get(slice.join(",")); + if (nextRank != null) { + node.rank = nextRank; + heapPush(heap, { node, rank: nextRank, start: node.start }); + } else { + node.rank = Infinity; + } } else { - break; + node.rank = Infinity; + } + + // Check for new possible merges with listPrev + if (node.listPrev) { + const slice = piece.slice(node.listPrev.start, node.end); + const prevRank = ranks.get(slice.join(",")); + if (prevRank != null) { + node.listPrev.rank = prevRank; + heapPush(heap, { node: node.listPrev, rank: prevRank, start: node.listPrev.start }); + } else { + node.listPrev.rank = Infinity; + } } } - return parts; + + const result: Array<{ start: number; end: number }> = []; + let current: BPEMergeNode | null = head; + while (current) { + if (!current.deleted) { + result.push({ start: current.start, end: current.end }); + } + current = current.listNext; + } + return result; } function bytePairEncode(piece: Uint8Array, ranks: Map) { @@ -271,63 +429,9 @@ export function getEncodingNameForModel(model: TiktokenModel) { case "gpt-4-turbo-2024-04-09": case "gpt-4-turbo-preview": case "gpt-4-0125-preview": - case "text-embedding-ada-002": - case "text-embedding-3-small": - case "text-embedding-3-large": { + case "text-embedding-ada-002": { return "cl100k_base"; } - case "gpt-4o": - case "gpt-4o-2024-05-13": - case "gpt-4o-2024-08-06": - case "gpt-4o-2024-11-20": - case "gpt-4o-mini-2024-07-18": - case "gpt-4o-mini": - case "gpt-4o-search-preview": - case "gpt-4o-search-preview-2025-03-11": - case "gpt-4o-mini-search-preview": - case "gpt-4o-mini-search-preview-2025-03-11": - case "gpt-4o-audio-preview": - case "gpt-4o-audio-preview-2024-12-17": - case "gpt-4o-audio-preview-2024-10-01": - case "gpt-4o-mini-audio-preview": - case "gpt-4o-mini-audio-preview-2024-12-17": - case "o1": - case "o1-2024-12-17": - case "o1-mini": - case "o1-mini-2024-09-12": - case "o1-preview": - case "o1-preview-2024-09-12": - case "o1-pro": - case "o1-pro-2025-03-19": - case "o3": - case "o3-2025-04-16": - case "o3-mini": - case "o3-mini-2025-01-31": - case "o4-mini": - case "o4-mini-2025-04-16": - case "chatgpt-4o-latest": - case "gpt-4o-realtime": - case "gpt-4o-realtime-preview-2024-10-01": - case "gpt-4o-realtime-preview-2024-12-17": - case "gpt-4o-mini-realtime-preview": - case "gpt-4o-mini-realtime-preview-2024-12-17": - case "gpt-4.1": - case "gpt-4.1-2025-04-14": - case "gpt-4.1-mini": - case "gpt-4.1-mini-2025-04-14": - case "gpt-4.1-nano": - case "gpt-4.1-nano-2025-04-14": - case "gpt-4.5-preview": - case "gpt-4.5-preview-2025-02-27": - case "gpt-5": - case "gpt-5-2025-08-07": - case "gpt-5-nano": - case "gpt-5-nano-2025-08-07": - case "gpt-5-mini": - case "gpt-5-mini-2025-08-07": - case "gpt-5-chat-latest": { - return "o200k_base"; - } default: never(model); throw new Error("Unknown model");