前の資料の最後で、行列積だけ速くしても意味がないと分かりました。 そこで attention を回路にしたら 15.6 倍になり、 softmax は帯域を 1 バイトも食いませんでした。 4 本の回路はどれも 2.0 GB/s の壁に張り付いています。
そして最後に、それまでの測り方が間違っていたことが分かりました。 結論の向きは変わりませんでしたが、数字は全部引き直しています(6 節)。
公開後に 比較相手の CPU 実装を -O2(NEON 無効)で焼いていたことが分かりました。
-mfpu=neon を足すと同じコードが 1.78 倍速くなります。
そのため本文の「CPU の何倍」はすべて 1.8 倍ほど過大でした。
数字はこのページ全体で訂正済みです(attention は 28 倍 → 15.6 倍)。
6 節で「測定式を疑え」と書いた直後に、今度は比較相手の作り込みで同じ偏りをやりました。 自分が作ったものは丁寧に測り、比べる相手は雑に測る——これが 2 回続いたということです。
the 78% that was still on the CPU
前の資料の 最後に足した実測が、計画の前提を壊しました。行列積を回路にして 14.2 ms まで縮めても、attention を CPU に残すと文脈 512 で 46 ms。 系全体の 78% を CPU が占めてしまい、せっかくの回路がほとんど効きません。
だから段階5でやることは一つでした。attention も回路にする。
two phases, one memory in between
行列積の回路は単純でした。流れてくる重みに次々と掛けて足すだけで、 後ろを振り返る必要がありません。attention はそうはいきません。
まず K を全部流してスコアを作り、そのスコアを重みとしてV をもう一度流す。 K を読み終わるまで V の計算が始められないので、途中にスコアの置き場が要ります。 ここが構造上いちばん大きな違いです。
段階5a では、この回路のたった1サイクルのずれを実機で追いかけて丸一日使いました。 しかも3 回とも違う原因を自信を持って断言し、2 回は間違っていました。 真因は「データが 2 段なのに有効フラグを 3 段で受けていた」——数え間違いです。
段階5b では順番を変えました。Vivado のシミュレータで先に通す。 文脈長を 1・2・3・5・8・12・16・64 と振り、上流のデータを 1 割の確率で欠かし、 下流を 4 回に 1 回詰まらせて、1,560 語を numpy と突き合わせます。 一発で通り、実機でも一発で一致しました。 論理の誤りはシミュレータで、タイミングと帯域だけ実機で——この切り分けです。
why attention is denser than matmul
ここが意外なところでした。attention は行列積より「濃い」計算です。
行列積では、読んだ重み 1 バイトは 1 回しか使いません。使ったら捨てます。 だから 1 バイトにつき掛け算 1 回。壁の高さがそのまま計算の上限になります。
attention は違います。このモデル(SmolLM2-135M)は GQA という構造で、 K と V の 1 ヘッドを、q の 3 ヘッドが共有しています。 つまり読んだ K の 1 バイトは、3 つの別々の計算に同時に配れる。
the exponent never leaves the chip
softmax は CPU では重い処理です。実測で attention 46 ms のうち 6.9 ms。 指数関数を文脈の長さぶん計算するからです。
回路では、これが表引き 1 回になります。スコアを 8 bit に丸めてあるので、
取りうる値は 256 通りしかありません。256 通りぶんの答えを先に焼いておけばいい。
exp という演算そのものが回路から消えます。
そして肝心なのは、この表がチップの中にあることです。 DDR との間をやりとりするデータは 1 バイトも増えません。
softmax には割り算が要りますが、回路ではやっていません。 分子 192 語と分母 3 語を出して、割るのは下流に任せます。 1 グループに 3 回しかない割り算のために回路を持つのは割に合わないからです。
ただ、出力を 192 語から 195 語に増やすのが意外と厄介でした。 出力段には 192 個の値から 1 つ選ぶ巨大な選択回路があり、ここは 1 クロックの余裕が 0.436 ns しかありません。2 択を 1 つ足すだけで落ちかねない。 そこで分母を、選択回路のアドレス計算式にちょうど出てくる空き番地 (192・200・208)に置きました。式を変えずに 3 語増やせて、回路は 1 段も深くなっていません。
when the number refuses to make sense
ここまでは順調でした。HP ポートを 2 本にして帯域を倍にする—— 段階3・4 でやった通りのことをやるだけのはずでした。
ところが 1.46 GB/s しか出ません。 1 本が 1.00 なので 1.46 倍。行列積のときは 1.72 倍伸びたのに、です。 照合は全部一致しているので、回路は正しく動いています。速度だけが合わない。
すぐに理由は思いつきました。attention の回路は、出力の 195 語を吐く間だけ上流を止めます。 行列積の回路は一度も止まりません。これが効いているに違いない——と。
でも数が合いませんでした。止まるのは 1 グループ 8,217 拍のうち 216 拍で、 たかだか 2.6% です。26% の説明にはなりません。
そこで回路をいじる前に、文脈長を 128・512・2048 と振って測り、 「1 グループあたり止まる時間 S」と「1 回の測定あたりの固定費 P」の 2 つを未知数にして連立方程式を解きました。答えはこうです。
measure the slope, not the ratio
P の正体は、測定プログラム自身です。転送を仕掛けて完了を待つ処理に 1 チャネル対あたり 0.7〜0.8 ms かかっていました。 8 ms の測定なら、それだけで 1 割が消えます。
そして、これは段階1 から全部の測定に乗っていました。
全部測り直しました。結論の向きは一つも変わりませんでしたが、数字はこうなります。
| 測定 | 旧(1 点) | 新(傾き) | 意味 |
|---|---|---|---|
| 段階2 ・ 計算なし・HP 1 本 | 1.13 | 1.143 | 道幅ちょうど 100.0% |
| 段階2 ・ 計算なし・HP 2 本 | 1.97 | 2.009 | DDR 理論 2.1 の 95.7% |
| 段階3 ・ INT8 MAC 16 個 | 1.91 | 1.982 | 壁の 98.7% |
| 段階4 ・ 三値 | 1.83 | 1.894 | 9.47 G重み/s |
| 段階5b ・ attention + softmax | 1.46 | 1.99 | 壁の 99.3% |
いちばん驚いたのは 1 行目です。HP 1 本の回路は、道幅ちょうど 100.0% で動いていました。 「98.7% だから 1.3% はどこかで損している」と 2 週間思っていたのですが、 損していたのは測定プログラムの中だけでした。
数字が合わないとき、まず疑うべきは測定式です。 今回たまたま助かったのは、回路を作り直す前に文脈長を振って連立方程式を解いたからです。 あのまま「影の累算器を足せば停止が消えるはずだ」と作り始めていたら、 2 週間かけて 2.6% しか改善しないものを作って、それでも数字が合わずに途方に暮れていました。
four different circuits, one number
測り直したうえで、これまで作った回路を全部並べます。
| 段階5b の回路 | 値 |
|---|---|
| 読み出し(HP 2 本) | 1.99 GB/s |
| DDR の壁に対して | 99.3% |
| 掛け算 | 5.95 G MAC/s |
| 1 バイトあたり | 2.99 MAC |
| LUT | 12,820 / 53,200 (24%) |
| DSP | 108 / 220 (49%) |
| タイミングの余裕 | +0.201 ns |
| 文脈 512 の attention(1 トークン) | 2.97 ms ── CPU の 15.6 倍 |
the optimisation that is not worth it
段階5b には、やる予定だった 3 つ目の項目がありました。 出力の 195 語を吐く間の停止(2.6%)を消すという改良です。 累算器をもう 1 組持って、出力しながら次のグループを受ければ消せます。
やめました。理由は測定が教えてくれました。 HP 1 本の回路は 97.4% で動いていて、これは停止から計算した理論値そのものです。 ところが HP 2 本ではすでに壁の 99.3%。停止を消しても、DDR がそれ以上出してくれません。
「律速は面積ではなく帯域」——この計画の出発点にあった読みが、 ここでは作らない理由になりました。 192 個の 32 bit 累算器をもう 1 組(レジスタ 6,144 個 × 2 組)積んでも、 壁が許さないのだから意味がありません。
チップにはまだ LUT が 76%、DSP が 51% 余っています。 余っているのに使い道がない、というのがこのボードの現状です。
what is measured and what is not
言えること。すべて実機での実測です。
言えないこと。ここは正直に書いておきます。