Kimi K3のトークン化を18倍高速化、百万トークンのエージェントワークロード向け
Basetenは新しいトークナイザー（Basetenkenizer）を導入し、Kimi K3のトークン化を百万トークンのエージェントワークロードで最大18倍高速化します。Rustベースの最適化（特殊化された事前トークン化、BPEマージ、マルチコアチャンク処理、ゼロコピーNumPy転送など）を組み合わせ、tiktokenと完全に同じトークンIDを維持します。この改善により、長い入力シーケンスの最初のトークンまでの時間が大幅に短縮され、特にプレフィックスキャッシュヒット率が高い場合に効果的です。
長年にわたり、推論エンジニアはトークン化にかかる時間を無視してきました。なぜなら、短い入力シーケンスでは数ミリ秒で済み、プリフィルやデコードに必要な時間のごく一部に過ぎなかったからです。しかし、この状況は変わりました。Kimi K3のようなオープンフロンティアモデルは、最大100万トークンの入力シーケンスをサポートしており、これはエージェントワークロードでますます一般的になっています。エージェントループは、ツールの結果、観測結果、検索ドキュメント、中間状態、推論トレース、更新された指示を次のリクエストに繰り返し追加します。入力シーケンスの長さとプレフィックスキャッシュヒット率が急速に増加するにつれて、トークン化時間は無視できなくなります。
この問題に対処するため、BasetenはBaseten Tokenizer（Basetenkenizer）を構築し、まずKimi K3向けに最適化しました。Baseten推論スタックでは、カスタムのPython tiktoken実装からRustベースのBasetenkenizerに移行しました。トークン化時間が最も重要な長い入力シーケンスでは、完全なサービングパスがtiktokenよりも最大18倍高速化され、正確なトークンIDの一貫性が維持されました。
Kimi K3のトークン化は簡単ではありません。byte-pair encoding（BPE）語彙と、構造トークン用の正規表現による事前トークン化を使用します。構造化部分では、</s>のような文字列は制御トークンとしてエンコードされるべきですが、ユーザーやツールのテキストでは、同じ文字列を通常のテキストとして扱わなければなりません。古いトークナイザーは、チャットをセグメントにレンダリングすることでこれを処理していました。各セグメントにはallow_specialフラグがあります。Basetenkenizerは、順序付けられた(text, allow_special)セグメントを1回のネイティブ呼び出しで受け入れ、トークン配列を返すため、Python/ネイティブ境界のオーバーヘッドを回避します。
パフォーマンスの向上は、複数の最適化によるものです。特殊化された事前トークン化スキャナーは、Kimiの分割正規表現パターンを事前に認識し、汎用正規表現エンジンの代わりに手書き関数にディスパッチします。スタック常駐BPEマージ層は、32バイト以下の短い事前トークンに対してスタック割り当てのリンクリストを使用し、ヒープやプライオリティキューのオーバーヘッドを回避します。マルチコアセマンティクスは、同一の事前トークンを同じCPUコアにルーティングし、キャッシュの再利用を実現します。さらに、ネイティブ型セグメントと安全なチャンキング、ゼロコピーNumPy所有権転送、スマートポインタPyO3バインディングなども採用されています。
ベンチマーク結果は顕著な進歩を示しています。54の型付きセグメントからなるKimi K3チャットプロンプトにおいて、Baseten Tokenizerとtiktoken、fastokens、gigatokenを比較しました。長い入力シーケンスでは、Baseten Tokenizerがエンドツーエンドのエンコードでリードし、100万トークンシーケンスでtiktoken比18倍、gigatoken比1.41倍の高速化を達成しました。ただし、gigatokenはオフラインのファイル処理で優れた性能を発揮するため、Baseten Tokenizerはオンラインサービング向けに最適化されている点に注意が必要です。
結論として、Baseten TokenizerはKimi K3のオンラインサービングにおいて、特にプレフィックスキャッシュヒット率が高い長コンテキストのエージェントワークロードで大幅に高速なパスを提供し、トークン化時間の節約が複数ターンにわたって累積され、ユーザーエクスペリエンスを向上させます。