当 token 计数“看不见”时

本次请求的 token 计数会在请求末尾时返回,但正如漏算的 Token:AI 网关限额机制的攻防博弈提到的,有些时候请求并不能正常结束,导致中间件无法获知 token 的计数。我们之前遇到过一个场景,高达 10% 左右的请求是提前中断的。即使不考虑这种异常请求,有些时候我们也需要提早知道 token 计费请求,比如基于 token 数的限流。如果需要等到上个请求结束时才更新计数器,有可能已经是明日黄花。

这个时候,就需要有不依赖于推理引擎的 token 计数方式。调用服务提供商或推理引擎的 count token 接口是一个办法,但考虑到请求放大等原因,这个方法并不实用。我们需要在本地完成计算。于是,网关里通常需要集成 tokenizer。

tiktoken 还是 hf tokenizer,这是个问题

主流的 tokenizer 库有两个流派:来自于 OpenAI 的 tiktoken,或来自于 HuggingFace 的 tokenizer(以下简称 hf tokenizer)。许多 tokenizer 库,要么是它们的移植版本(如 tiktoken-go),要么是对它们的 binding。这两个项目都是用 Rust 实现的(tiktoken 没有对外提供对应的 crate,只有 Python package,但它是 Rust 核 Python 皮的项目)。

有趣的是,这两个项目各自在 README 的头顶显眼位置贴上和对家比较的 benchmark,简直是相爱相杀。

谁的陈述更贴近真相?我自己也写了个 benchmark(代码见文后)。和 hf tokenizer 不同的是,我测试的模型是 gpt-4o,因为这个模型用户更为广泛,而且作为 tiktoken 官方支持的模型,相对更能体现 tiktoken 的真实实力。注意由于 tiktoken 没有提供直接的 crate,而我又想做尽可能贴近核心的 benchmark,所以 benchmark 是在 tiktoken 项目里实现的,需要在 tiktoken 代码库目录下运行。

这里直接贴性能结论:在全部四个用例里(短、中、长、多语言),tiktoken 的吞吐量都是 hf tokenizer 的 1.5 倍以上。

tiktoken 为什么更快

现代 tokenizer 大致分成几类:

  • WordPiece:合并规则不只看频率,还看“合并后能提升训练数据概率多少”。
  • Unigram:反过来,从大词表开始,逐步删掉对模型贡献最小的 token。
  • BPE:从字符(或字节)开始,反复合并最常见的相邻对。它是 Byte-Pair Encoding 的缩写。通常使用的是它的一个变种:Byte-Level BPE。

Tiktoken 的实现即是 Byte-Level BPE。编码操作如下:

Full text → [Regex split] → Chunks → [BPE merge per chunk] → Tokens

在 Regex split 时,针对 gpt-4o 它会采用下列的正则:

pat_str = "|".join(
          [
              r"""[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*[\p{Ll}
  \p{Lm}\p{Lo}\p{M}]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?""",
              r"""[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+[\p{Ll}
  \p{Lm}\p{Lo}\p{M}]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?""",
              r"""\p{N}{1,3}""",
              r""" ?[^\s\p{L}\p{N}]+[\r\n/]*""",
              r"""\s*[\r\n]+""",
              r"""\s+(?!\S)""",
              r"""\s+""",
          ]

解释一下:

Pattern 1:小写结尾的单词

[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*[\p{Ll}\p{Lm}\p{Lo}\p{M}]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?

[^\r\n\p{L}\p{N}]?:可选,单个“非换行/字母/数字”的字符(空格、标点)

[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*:0 次或多次:大写/标题大小写/修饰/其他字母 + 组合符号

[\p{Ll}\p{Lm}\p{Lo}\p{M}]+:1 次或多次:小写/修饰/其他字母 + 组合符号

(?i:'s|'t|...):可选的英语缩写(大小写不敏感)

匹配:以小写结尾的单词
"hello" → "hello"
"Hello" → "Hello"
"HELLo" → "HELLo"
" hello" → " hello" (带前导空格)
".Hello's" → ".Hello's" (带前导标点 + 英语缩写)
"中文" → "中文" (Lo 字符)

Pattern 2:大写开头的单词

[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+[\p{Ll}\p{Lm}\p{Lo}\p{M}]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?

[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+:1 次或多次:大写字母(必需)
[\p{Ll}\p{Lm}\p{Lo}\p{M}]*:0 次或多次:小写字母

匹配:以大写开头的单词,末尾可继续大写
"HELLO" → "HELLO" (全大写)
"USA" → "USA"
"NASA's" → "NASA's"
" HTTP" → " HTTP"

Pattern 1 vs 2 的区分:
"Hello" → Pattern 1(upper* lower+)✓
"HELLO" → Pattern 2(upper+ lower*)✓
"hello" → Pattern 1(upper*=∅, lower+)✓

Pattern 3:数字(1-3 位)

\p{N}{1,3}

匹配:每次 1 到 3 位数字
"123456" → ["123", "456"]
"42" → ["42"]
"1000" → ["100", "0"]

避免长数字被当作单个 token,有利于算术相关的处理。

Pattern 4:标点/符号

?[^\s\p{L}\p{N}]+[\r\n/]*

?:可选空格

[^\s\p{L}\p{N}]+:1 次或多次:非空白/字母/数字(标点、符号)

[\r\n/]*:0 次或多次:换行或斜杠

匹配:连续的标点序列
"..." → "..."
" !!!" → " !!!"
"===" → "==="
"---\n" → "---\n"

Pattern 5:行尾换行

\s*[\r\n]+

匹配:带可选前导空白的换行
"\n" → "\n"
" \n\n" → " \n\n"
"\t\r\n" → "\t\r\n"

Pattern 6:末尾空白

\s+(?!\S)

\s+: 1 次或多次空白
(?!\S): 负向前瞻:后面不是非空白字符

匹配:字符串末尾(或换行前)的空白
"hello " → the " " at the end

Pattern 7:其他空白

\s+

匹配:其余所有空白(单词之间的空格)
"a b" → the " " between a and b

最后需要强调的是,顺序很重要。

正则的 alternation(|)从左到右尝试,先匹配到的胜出:

1. 单词(小写结尾)  ─┐
2. 单词(大写占优)  ─┼── 先匹配字母
3. 数字(1-3 位)   ─── 再匹配数字
4. 标点符号         ─── 再匹配符号
5. 换行            ─┐
6. 末尾空白         ─┼── 最后处理空白
7. 其他空白         ─┘

完整示例

输入:"Hello WORLD! 12345\n "

Token 列表:
"Hello" → Pattern 1
" WORLD" → Pattern 2
"!" → Pattern 4(这里没有前导空格,因为空格已经被消耗掉了)
" " → Pattern 7(数字前的空格)
"123" → Pattern 3
"45" → Pattern 3
"\n" → Pattern 5
" " → Pattern 6(末尾空白)

在切割成多个 chunks 后,tiktoken 会在每个 chunk 内执行 _byte_pair_merge_byte_pair_merge 包含以下步骤:

  • O(mn),n 为 parts 数量,m 为 merges 次数
  • 每次迭代线性扫描,找最小 rank 的 pair
  • 使用 Vec::remove(),会触发元素移动

Parts 的数量(n):
初始值:输入片段的长度。
例如,输入 "hello" → 初始 parts 为 5 个:['h', 'e', 'l', 'l', 'o']
合并之后:每次合并让 parts -1,合并 m 次后剩下 n - m 个 parts(也就是最终 tokens)。

合并次数(m):
上界:m ≤ n - 1(n 个 parts 最多合并成 1 个 token)

实际值:取决于有多少相邻 pair 在词表里有合并规则。由下列因素决定:

  1. 训练词表 - BPE 训练阶段学到哪些 byte-pair merges
  2. 具体输入 - 哪些 pair 实际上相邻出现
  3. 合并优先级 - rank 越小越先合并,影响后续相邻关系

示例:
输入:"hello"(n=5)
Vocab 有:"ll"→rank 1, "he"→rank 2, "lo"→rank 3

步骤 0:[h, e, l, l, o](5 parts)
步骤 1:[h, e, ll, o](合并 "ll",4 parts)
步骤 2:[he, ll, o](合并 "he",3 parts)
步骤 3:[he, llo](如果 vocab 里有 "llo" 就合并,否则停止)

m = 2 或 3,取决于 vocab

有趣的是,由于每个 chunk 都很小,导致性能的瓶颈实际上不在 BPE 算法上,而在正则表达式切割的部分。以下是 tiktoken encode 时的火焰图:

tiktoken火焰图

在 tiktoken 代码里也提到:

// Most of the time is spent in regex. The easiest way to speed this up is by using less fancy
// regex features. For instance, using a regex parse-able by `regex` crate is 3x faster than
// the usual regex we use

它之所以没法用更快的 regex crate,是因为 split 用到的正则表达式里需要 \s+(?!\S) 这样的 lookahead 操作。在注释里,tiktoken 的作者还抱怨了下因为 tokenizer 的这种设计导致性能无法进一步优化。所以模型推理是一个系统工程,在训练时处理数据的 token pattern 会给推理时的 tokenizer 设立一个上界。

hf tokenizer 真的慢吗?

hf tokenizer 作为一个更加通用的 tokenizer 库,即使实现了 Byte-Level BPE,也不能像 tiktoken 那样直接操作在 byte 序列的层次。虽然 hf tokenizer 的架构和 tiktoken 类似:

Text → Normalizer → PreTokenizer → BPE → PostProcessor → Tokens

其中 PreTokenizer 相当于 tiktoken 的 Regex split。但是为了更好地适应不同的模型,hf tokenizer 必须实现其他的流程。

以 hf tokenizer 的火焰图为例,我们可以看到两张火焰图都花了很多时间在 <fancy_regex::Matches as core::iter::traits::iterator::Iterator>::next 上。

hf tokenizer 的火焰图

但 tiktoken 的这个函数直接就是 tiktoken::CoreBPE::encode_ordinary 的下级函数,而 hf tokenizer 还需要执行同样占用大量时间的 normalization。这里面额外的内存分配,也许解释了那多出来的 50% 时间花在了哪里。

作为代价,tiktoken 只适合用来解析 GPT 系列模型,如果有模型需要不一样的 PreTokenizer 或其他操作,就没办法直接用 tiktoken 这种简单的流程。

所以 hf tokenizer 尽管在支持 GPT tokenizer 上不如 tiktoken 快,但是有些时候你不得不用它。

举一个实际的例子。Qwen3-235B-A22B-Instruct-2507 的 tokenizer 和 GPT 的基本上差不多。我们曾经测试过,即使在差异较大的中文场景,Qwen 和 GPT 在输出长度的偏差上仍然只有 10% 左右。所以如果只关心 token 长度,tiktoken 可以当作 Qwen tokenizer 使用。

但是 DeepSeek 和 GPT 的差异较大。在中文场景的实验中,DeepSeek 和 GPT 的结果相差一倍。我估计是因为 DeepSeek 的词表和分割时用的正则表达式较 GPT 的差别很多。

这也是为什么不能用 tiktoken 来模拟 DeepSeek 的 tokenizer。

眼尖的读者在看完 benchmark 代码后可能会发现,benchmark 时我是直接指定词表和分割的正则表达式。虽然 Python 层次上的 API 没有暴露指定正则和词表的接口,但是底层的 Rust 核心是有的。那是不是可以给 DeepSeek 的词表做一些预处理,让它和 GPT 差不多,这样就能基于 tiktoken 的底层接口来解析呢?也许可以,也许不行。注意 DeepSeek 除了 pretokenizer 不一样外,还定义了额外的 PostProcessor:https://huggingface.co/deepseek-ai/DeepSeek-V3.2/resolve/main...。我不确定这种额外的 PostProcessor 是否已经在 tiktoken 里面对齐。

注意 Qwen3.5 的 pretokenize 较 Qwen3-235B-A22B-Instruct-2507 有所变更,优化了中文编码,预计在中文场景下和 tiktoken 的结果会有更大的不同。

结语

当 token 计数不能依赖推理引擎回传时,网关侧的本地 tokenizer 是绕不开的基础设施。tiktoken 和 hf tokenizer 的差别不只在速度,还在“能否复现模型的预处理语义”:前者以极简流程换高吞吐,后者以更强的可配置性换更广泛的模型适配。实践中更关键的是知道自己在模拟哪一种模型行为,以及吞吐量和准确度之间的取舍。

benchmark 代码

use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
use sha2::Digest;
use std::collections::HashMap;
use std::io::Read;
use std::time::Duration;
use tiktoken::CoreBPE;
use tokenizers::Tokenizer;

const O200K_BASE_URL: &str =
    "https://openaipublic.blob.core.windows.net/encodings/o200k_base.tiktoken";
const O200K_BASE_HASH: &str = "446a9538cb6c348e3516120d7c08b09f57c36495e2acfffe59a5bf8b0cfb1a2d";

// Pattern from openai_public.py o200k_base()
const O200K_PAT_STR: &str = r#"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*[\p{Ll}\p{Lm}\p{Lo}\p{M}]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?|[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+[\p{Ll}\p{Lm}\p{Lo}\p{M}]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n/]*|\s*[\r\n]+|\s+(?!\S)|\s+"#;

const ENDOFTEXT: &str = "<|endoftext|>";
const ENDOFPROMPT: &str = "<|endofprompt|>";

/// Number of times to encode each text in a single benchmark iteration
/// This amortizes the measurement overhead and focuses on encoding performance
const ITERATIONS_PER_SAMPLE: usize = 100;

/// Downloads and caches the tiktoken file
fn download_tiktoken_bpe(
    url: &str,
    expected_hash: &str,
) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
    let cache_dir = std::env::var("TIKTOKEN_CACHE_DIR").unwrap_or_else(|_| {
        let tmp = std::env::temp_dir();
        tmp.join("data-gym-cache").to_string_lossy().to_string()
    });

    std::fs::create_dir_all(&cache_dir)?;

    let cache_key = format!("{:x}", md5::compute(url.as_bytes()));
    let cache_path = std::path::Path::new(&cache_dir).join(&cache_key);

    // Check cache
    if cache_path.exists() {
        let data = std::fs::read(&cache_path)?;
        let hash = format!("{:x}", sha2::Sha256::digest(&data));
        if hash == expected_hash {
            return Ok(data);
        }
        // Hash mismatch, remove cached file
        let _ = std::fs::remove_file(&cache_path);
    }

    // Download
    let resp = ureq::get(url).call()?;
    let mut data = Vec::new();
    resp.into_reader().read_to_end(&mut data)?;

    // Verify hash
    let hash = format!("{:x}", sha2::Sha256::digest(&data));
    if hash != expected_hash {
        return Err(format!("Hash mismatch: expected {}, got {}", expected_hash, hash).into());
    }

    // Cache
    std::fs::write(&cache_path, &data)?;

    Ok(data)
}

/// Parses tiktoken BPE format (base64-encoded token + rank)
fn parse_tiktoken_bpe(data: &[u8]) -> Result<HashMap<Vec<u8>, u32>, Box<dyn std::error::Error>> {
    use data_encoding::BASE64;

    let mut ranks = HashMap::new();
    for line in data.split(|&b| b == b'\n') {
        if line.is_empty() {
            continue;
        }
        let parts: Vec<&[u8]> = line.splitn(2, |&b| b == b' ').collect();
        if parts.len() != 2 {
            continue;
        }
        let token = BASE64.decode(parts[0])?;
        let rank: u32 = std::str::from_utf8(parts[1])?.parse()?;
        ranks.insert(token, rank);
    }
    Ok(ranks)
}

/// Loads o200k_base encoding
fn load_o200k_base() -> Result<CoreBPE, Box<dyn std::error::Error>> {
    let data = download_tiktoken_bpe(O200K_BASE_URL, O200K_BASE_HASH)?;
    let mergeable_ranks = parse_tiktoken_bpe(&data)?;

    let mut special_tokens = HashMap::new();
    special_tokens.insert(ENDOFTEXT.to_string(), 199999);
    special_tokens.insert(ENDOFPROMPT.to_string(), 200018);

    CoreBPE::new::<_, _, Vec<(String, (u32, u32))>>(mergeable_ranks, special_tokens, O200K_PAT_STR)
        .map_err(|e| format!("Failed to create CoreBPE: {}", e).into())
}

/// Loads HuggingFace tokenizer
fn load_hf_tokenizer() -> Result<Tokenizer, Box<dyn std::error::Error>> {
    // Try to load from HuggingFace Hub using the http feature
    // This requires network access on first run
    match Tokenizer::from_pretrained("Xenova/gpt-4o", None) {
        Ok(tokenizer) => Ok(tokenizer),
        Err(e) => {
            // Fallback: try to load from local cache
            let cache_path = dirs::home_dir()
                .ok_or("Could not determine home directory")?
                .join(".cache/huggingface/hub/models--Xenova--gpt-4o/snapshots")
                .join("main/tokenizer.json");

            if cache_path.exists() {
                Tokenizer::from_file(cache_path)
                    .map_err(|e| format!("Failed to load from cache: {}", e).into())
            } else {
                Err(format!("Failed to load HF tokenizer: {}. Try running: python -c 'from tokenizers import Tokenizer; Tokenizer.from_pretrained(\"Xenova/gpt-4o\")'", e).into())
            }
        }
    }
}

/// Sample texts for benchmarking
fn get_sample_texts() -> Vec<(&'static str, String)> {
    vec![
        (
            "short",
            "Hello, world! This is a test.".to_string(),
        ),
        (
            "medium",
            "The quick brown fox jumps over the lazy dog. ".repeat(50),
        ),
        (
            "long",
            "Artificial intelligence (AI) is intelligence demonstrated by machines, in contrast to the natural intelligence displayed by humans and animals. Leading AI textbooks define the field as the study of \"intelligent agents\": any device that perceives its environment and takes actions that maximize its chance of successfully achieving its goals. ".repeat(100),
        ),
        (
            "multilingual",
            "Hello 世界 مرحبا Привет こんにちは 안녕하세요 ".repeat(50),
        ),
    ]
}

fn bench_tiktoken(c: &mut Criterion) {
    let enc = load_o200k_base().expect("Failed to load o200k_base");
    let texts = get_sample_texts();

    let mut group = c.benchmark_group("tiktoken_encode");
    for (name, text) in &texts {
        let text_bytes = text.len();
        // Throughput is total bytes processed (iterations * text size)
        group.throughput(Throughput::Bytes(
            (text_bytes * ITERATIONS_PER_SAMPLE) as u64,
        ));
        group.bench_with_input(BenchmarkId::from_parameter(name), text, |b, text| {
            b.iter(|| {
                // Encode multiple times per iteration to amortize measurement overhead
                for _ in 0..ITERATIONS_PER_SAMPLE {
                    let tokens = enc.encode_ordinary(black_box(text));
                    black_box(tokens);
                }
            })
        });
    }
    group.finish();
}

fn bench_hf_tokenizer(c: &mut Criterion) {
    let tokenizer = load_hf_tokenizer().expect("Failed to load HF tokenizer");
    let texts = get_sample_texts();

    let mut group = c.benchmark_group("hf_tokenizer_encode");
    for (name, text) in &texts {
        let text_bytes = text.len();
        // Throughput is total bytes processed (iterations * text size)
        group.throughput(Throughput::Bytes(
            (text_bytes * ITERATIONS_PER_SAMPLE) as u64,
        ));
        group.bench_with_input(BenchmarkId::from_parameter(name), text, |b, text| {
            b.iter(|| {
                // Encode multiple times per iteration to amortize measurement overhead
                for _ in 0..ITERATIONS_PER_SAMPLE {
                    let encoding = tokenizer.encode(black_box(text.clone()), true).unwrap();
                    let token_ids: Vec<u32> = encoding.get_ids().to_vec();
                    black_box(token_ids);
                }
            })
        });
    }
    group.finish();
}

fn bench_comparison(c: &mut Criterion) {
    let tiktoken_enc = load_o200k_base().expect("Failed to load o200k_base");
    let hf_tokenizer = load_hf_tokenizer().expect("Failed to load HF tokenizer");

    let texts = get_sample_texts();

    let mut group = c.benchmark_group("comparison");
    group.measurement_time(Duration::from_secs(10));

    for (name, text) in &texts {
        let text_bytes = text.len();
        // Throughput is total bytes processed (iterations * text size)
        group.throughput(Throughput::Bytes(
            (text_bytes * ITERATIONS_PER_SAMPLE) as u64,
        ));

        group.bench_with_input(BenchmarkId::new("tiktoken", name), text, |b, text| {
            b.iter(|| {
                // Encode multiple times per iteration to amortize measurement overhead
                for _ in 0..ITERATIONS_PER_SAMPLE {
                    let tokens = tiktoken_enc.encode_ordinary(black_box(text));
                    black_box(tokens);
                }
            })
        });

        group.bench_with_input(BenchmarkId::new("hf_tokenizer", name), text, |b, text| {
            b.iter(|| {
                // Encode multiple times per iteration to amortize measurement overhead
                for _ in 0..ITERATIONS_PER_SAMPLE {
                    let encoding = hf_tokenizer.encode(black_box(text.clone()), true).unwrap();
                    let token_ids: Vec<u32> = encoding.get_ids().to_vec();
                    black_box(token_ids);
                }
            })
        });
    }
    group.finish();
}

criterion_group!(
    benches,
    bench_tiktoken,
    bench_hf_tokenizer,
    bench_comparison,
);
criterion_main!(benches);

标签: none

添加新评论