Taalas-Style On-Chip Weights on a $250 FPGA: a Language Model at 60k tok/s

316万パラメータのINT4トランスフォーマーが、250ドルのAMD Kria KV260 FPGAのオンチップメモリだけで動きます。トークンのループにDRAMは一度も登場しません。トークンを1つずつ生成する処理は演算ではなくメモリに縛られるため、重みがチップ外のDDRにあると専用回路をいくら足しても速度はそこで頭打ちになります。そこで重みも活性値もKVキャッシュもFPGA自身のURAMとブロックRAMに収まるまでモデルを小さくしました。

実測値は、公開中のチャット版が毎秒約21,300トークン、同じボードのArmコアが11トークン、RTX 3050 Ti搭載ノートPCが719トークン、何も記憶しない16ストリーム版は59,965トークンです。ブラウザのウィジェットはWebSocketで実機につながっており、「むかしむかし」と投げれば本物のチップが続きを書きます。学習コーパスのTinyStoriesを内容語だけに削ってあるため、応答はぶつ切りで、質問に答えるより物語の続きを書いてしまいます。