May 23, 2025
在上一篇 比 tiktoken 还快？BPE 分词优化小结 中提到，我们假设 Pre-Tokenizer 后的子串都是比较短的，所以类似于 tiktoken 中的实现，我们只用了 vector 遍历的方式来实现 bpe 分词（详见源码~/src/tokenizer/BPE.cpp#L157），即使改成链表，时间复杂度仍然是 O(nm) 的，通常 m(merge次数) 与 n(字符串长度) 接近，所以也可以看作是 O(n^2) ，这在 n 较小的情况下问题不大，但如果 n 很大呢？我们用 tiktoken 实测一下：
tiktoken_enc = tiktoken.Encoding(...)
start = time.perf_counter_ns()
tiktoken_enc.encode("a" * 102400)
end = time.perf_counter_ns()
cost = (end - start) / 1e6
print(f"{cost:.3f} ms")5182 ms也就是说一个 100k 长度的 token，encode 耗时 5s，这显然不合理，而如果用 huggingface tokenizers，相同的机器上只需要 56.094 ms，相差 100 倍，翻看源码发现，huggingface 就是用了上一篇文章提到的堆结构（Rust BinaryHeap）来实现，将时间复杂度降低到 O(mlogn)，经过一番调试，我在 tinygpt::tokenizer 中用 C++ 也实现了一个版本（源码 ~/src/tokenizer/BPE.cpp#L223），先看一下性能对比数据：
| token 长度 | tiktoken | huggingface | tinygpt |
|---|---|---|---|
| 32 B | 0.242 ms | 0.089 ms | 0.027 ms |
| 128 B | 0.065 ms | 0.081 ms | 0.056 ms |
| 1 KB | 0.526 ms | 0.444 ms | 0.264 ms |
| 10 KB | 44.048 ms | 4.791 ms | 2.901 ms |
| 100 KB | 5182.736 ms | 56.094 ms | 35.890 ms |
| 400 KB | 94965.138 ms | 250.234 ms | 162.074 ms |
可以看到，在 token 较长情况下，tiktoken 性能下降非常明显。接下来看下我们的 C++ 优化版本的具体实现 (BPE::bpeV2)：
整体流程上和 BPE::bpeV1 是一致的，区别主要是两点：
- 用双链表来记录子串的 rank 信息
- 用 std::priority_queue 来实现 O(logN) 的最小 rank 查找
首先定义链表节点：
struct Node {
uint32_t pos;
uint32_t rank;
Node* prev;
Node* next;
bool active;
};这里多了个 active 标记，是因为在发生 merge 时，需要删除被 merge 的节点，而在堆中是不方便操作的，因此通过修改 active 标记来实现。
然后是双链表的初始化：
std::vector<Node> nodes(words.size() + 1);
for (uint32_t i = 0; i < nodes.size(); i++) {
nodes[i].pos = (i < words.size()) ? (words[i].data() - text.data()) : text.size();
nodes[i].prev = (i > 0) ? &nodes[i - 1] : nullptr;
nodes[i].next = (i < nodes.size() - 1) ? &nodes[i + 1] : nullptr;
nodes[i].rank = std::numeric_limits<uint32_t>::max();
nodes[i].active = true;
}接下来是初始化基于 std::priority_queue 的最小堆
using QueueElem = std::pair<uint32_t, Node*>; // <rank, node>
auto cmp = [](const QueueElem& a, const QueueElem& b) {
// If ranks are equal, prioritize the smaller index
if (a.first == b.first) {
return a.second->pos > b.second->pos;
}
return a.first > b.first;
};
std::priority_queue<QueueElem, std::vector<QueueElem>, decltype(cmp)> pq(cmp);
for (auto& node : nodes) {
if (node.next) {
node.rank = getRank(&node);
if (node.rank != std::numeric_limits<uint32_t>::max()) {
pq.emplace(node.rank, &node);
}
}
}这里需要注意的是 compare 的实现，如果两个节点 rank 相等，则比较下标，因为在 bpeV1 版本中我们的线性扫描是下标从小到大 for 循环的，这样才能保证分词结果的一致性。
然后是在 merge 迭代中，每次只需要从堆中 pop 出最小 rank 节点即可：
while (!pq.empty()) {
auto [minRank, node] = pq.top();
pq.pop();
if (!node->active || node->rank != minRank) {
continue;
}
Node* nextNode = node->next;
if (!nextNode || !nextNode->active) {
continue;
}
...
}后面具体的 merge 只需要操作双链表的指针，设置 active 标记：
// mark not active
nextNode->active = false;
activeCount--;
// merge with next
node->next = nextNode->next;
if (node->next) {
node->next->prev = node;
}然后更新 merge 节点前后的 rank，将新产生的节点添加到堆中
// update rank
const uint32_t oldRank = node->rank;
node->rank = getRank(node);
if (node->rank != oldRank && node->rank != std::numeric_limits<uint32_t>::max()) {
pq.emplace(node->rank, node);
}
// update prev node
if (node->prev && node->prev->active) {
const uint32_t prevOldRank = node->prev->rank;
node->prev->rank = getRank(node->prev);
if (node->prev->rank != prevOldRank && node->prev->rank != std::numeric_limits<uint32_t>::max()) {
pq.emplace(node->prev->rank, node->prev);
}
}在 merge 迭代完成后，我们遍历双链表，仍然 active 的节点就是最终 bpe 分词的结果：
std::vector<std::string_view> ret;
ret.reserve(activeCount - 1);
for (Node* curr = &nodes[0]; curr; curr = curr->next) {
if (curr->active && curr->next) {
ret.emplace_back(text.data() + curr->pos, curr->next->pos - curr->pos);
}
}实际测试发现，短 token 情况下还是 bpeV1 更快，原因可能是数据结构更精简，cpu cache 命中率更高，因此只在 token 较长情况下才切换到 bpeV2：
constexpr size_t FAST_BPE_THRESHOLD = 32;
auto bpePieces = token.size() > FAST_BPE_THRESHOLD ? bpeV2(token) : bpeV1(token);