前の資料の最後に「言えないこと」を 3 つ書きました。 通しで 1 トークンも生成していない、流しているのはランダムな数値、 指数表の温度は仮の値。この資料はその 3 つを片付けた記録です。
いまは文章を渡すと、ボードが 1 トークンずつ計算して続きを返します。 行列積と attention は自作の回路、残りは CPU。 Mac で同じ計算をさせると、logits が最後のビットまで一致します。
stage 6 · the circuit was right, the model was not
最初にやったのは、これまで使ってきた SmolLM2-135M の本物の重みを三値 (−1 / 0 / +1)に丸めて回路に流すことでした。回路の答えは 239,232 語すべて numpy と一致しました。回路は正しく動いています。
ところがモデルとしては完全に壊れていました。
| SmolLM2-135M | perplexity | The capital of France is → |
|---|---|---|
| そのまま(fp32) | 18.31 | まともな英語 |
| 三値に丸めた | 9,176,689 | makers emergmakermakermaker… |
| 参考:語彙 49,152 から無作為に選ぶ | 49,152 | — |
でたらめに選ぶより 187 倍悪い。丸め方を行ごとに変えてもほとんど変わりません。 重み 1 個あたりの誤差が 5〜6 割あり、それが 30 層ぶん積み重なるからです。
BitNet b1.58 は三値のまま学習したモデルの名前です。 普通に学習した重みを後から三値に丸めたものとは別物でした。 回路の正しさとは関係のない話ですが、載せるモデルを選び直す必要があるということです。
この段階ではもう一つ見つかりました。回路は道幅の 100.0% で動いているのに、 1 トークンの時間の 83% が DMA の呼び出しの固定費でした。 当時の形式は「1 転送に x(入力ベクトル)が 1 つ」なので、行列が変わるたびに転送を切るしかなく、 1 トークンで 181 回も呼んでいたからです。これが 3 節の宿題になります。
stage 6.5 · the only one that fits in 512 MB
三値で学習済みで、このボードの 512 MB に載るものを探すと、候補は
1bitLLM/bitnet_b1_58-large(0.7B)だけでした。
numpy で回路と同じ量子化手順を書いて確かめると、
perplexity 17.90。三値のまま、SmolLM2 の fp32(18.31)より良い値です。
ただし形が違うので、回路もデータの詰め方も作り直しになります。
| SmolLM2-135M | bitnet_b1_58-large | 回路への影響 | |
|---|---|---|---|
| hidden / FFN / 層 | 576 / 1536 / 30 | 1536 / 4096 / 24 | 重みが 6 倍(三値で 136 MB) |
| head_dim | 64 | 96 | attention の回路を作り直し |
| q : kv ヘッド | 9 : 3 | 16 : 16 | GQA が無い。1 バイトの再利用が 3 回→1 回 |
| lm_head | — | fp16(三値ではない) | そのままでは 98 MB を毎トークン読む |
lm_head(最後に語彙 32,002 個の点数を出す行列)はメモリの都合で三値に落としました。 INT8 なら品質はそのまま(17.90 → 17.89)ですが、DMA で読める領域(CMA)に収まりません。 三値にすると perplexity は 17.90 → 35.59 と倍になりますが、文章は壊れません。 これはこの構成で払っている品質の代償で、最後の節でもう一度触れます。
stages 7 & 8 · a header, then a variable row length
181 回の呼び出しをどう減らすか。回路は速いので、回路への渡し方を変えます。 ストリームの頭に小さな見出しを付けて「この後 R 行続く」と書けば、 x が変わっても転送を切らずに次の組へ進めます。次に、見出しに1 行の長さも書けるようにしました。
行長を可変にした効果は 3 つ重なりました。詰め物が減って読む量が 17% 減り、 組の切れ目が 4 分の 1 になり、CPU で部分和を足す後始末が消えました。 そしてもう一つ、重みが 141.3 MiB に縮んだおかげで、後で attention と同じチップに同居できるようになります(5 節)。
stages 7c & 9 · head_dim 96, no GQA, and a multiplier
bitnet-large の attention を CPU で測ったら 96.22 ms(文脈 512)でした。
それまで使っていた換算値は 180 ms。
古いベンチを NEON 無効(-O2)で焼いていたためで、前の資料も訂正しました。
これで行列積(96.51 ms)と attention(96.22 ms)がほぼ半々だと分かりました。
段階5b の回路は head_dim 64・GQA 3:1 専用なので使えません。
作り直す前に numpy でもう一つ確かめたことがあります。
段階5b はスコアを s >> 8 という固定のシフトで 8 bit にしていましたが、
これは q と K の量子化の尺度を無視しています。本物のモデルに当てると
指数の温度が層・ヘッドごとにずれます。
| attention の作り(lm_head は三値) | perplexity |
|---|---|
| float の本来の softmax | 35.56 |
| 段階5b のまま(固定シフト) | 58.25 |
| 見出しに倍率 m を載せる(掛け算器 1 個) | 35.90 |
m = TAU · sq · sk / √96 を 224 倍した整数を見出しに書き、
スコアに掛けてから指数表を引きます。実際の値の範囲は 10,818〜79,505 で、
16 bit には入らず、DSP の 18 bit 入力にちょうど収まりました。
これで前の資料の「指数表の温度は仮の値」が片付きました。
| 段階9 の attention 回路(HP 2 本) | 値 |
|---|---|
| 照合 | 37,248 語すべて一致 |
| 文脈 512・1 トークン(24 層 × 16 ヘッド) | 20.22 ms |
| 読み出し | 1.872 GB/s(DDR の壁の 93.2%) |
| 1 バイトあたり | 1.00 MAC(GQA が無いので) |
| 掛け算器 | 17 個(段階5b は 48 個) |
| CPU(NEON 有効)との比 | 4.8 倍 |
142.86 MHz では 2 回ともタイミング違反になりました。違反したのは自作の回路ではなく AXI DMA の中の転送長カウンタです。125 MHz に下げました。 2 本ぶんの道幅は 2.000 GB/s で、DDR の壁 2.009 とほぼ同じです。 帯域で頭打ちの回路なので、クロックを下げても損はほとんどありません。
stage 10 · four DMAs, no switch
行列積の回路と attention の回路を、同じビットストリームに載せました。 1 トークンの中で両者は必ず交互に動きます。 だから切り替え回路は作らず、DMA を 4 台置いて、使うほうを叩くだけにしました。 同時には動かないので、DDR の取り合いも起きません。
同居の代償は 1.4% でした(別々のビットストリームなら 107.8 ms、同居で 109.3 ms の見込み)。 その差のほぼ全部は、DMA の転送長を 8 MiB 以下に切ったぶん往復が増えたことによるものです。 資源は LUT 34%。チップにはまだ 3 分の 2 が空いています。
stage 11 · scatter-gather instead of copying
通しで動かすには、1 トークンごとに K と V を 1 行ずつ積み増す必要があります(KV キャッシュ)。 ここで attention の回路の入力形式が問題になりました。
回路が受け取るのは [見出し][q][K を T 行][V を T 行] を 1 ヘッドずつつなげたストリームです。
T が 1 増えると、V の開始位置と、後ろのヘッド全部の位置がずれます。
単純な DMA は「連続した 1 領域」しか読めないので、毎トークン KV を丸ごと詰め直す(T=512 で 36 MB)か、
ヘッドごとに転送を切る(1 トークン 384 回)しかありません。
段階9 までは、K と V を INT8 にする尺度を「全位置の最大値」から決めていました。 これはまだ来ていない未来の K を見ています。しかも T が伸びるたびに最大値が変わるので、 過去の INT8 を全部作り直すことになり、積み増しができません。
そこで層・ヘッドごとに尺度を固定しました。評価文とは別の文章 144 トークンを流して max|K| と max|V| を測り、それを使い続けます。perplexity は 35.90 → 35.61 と、固定したほうがむしろ良くなりました。
bit-identical logits on two machines
通しで動かすと、確かめるべきものが増えます。 回路の答えが正しいかだけでなく、CPU 側の浮動小数の計算まで含めて、1 トークン全体が正しいかです。 そこで同じ計算を Mac でも走らせ、両方の logits(語彙 32,002 個ぶんの点数)を比べることにしました。
Mac(Apple M4)もボード(Cortex-A9)も ARM なので、arm_neon.h の同じ命令で書けます。
それでも、何もしなければ結果は揃いません。揃えるために決めたことはこれです。
-ffp-contract=off)。Mac の clang は既定で FMA にまとめてしまうexpf や rintf は libm ごとに最下位ビットが違いうるvrecpsq_f32 を使わない。同じ名前でも AArch64 では融合演算になり、ARMv7 と結果が変わるthe hardware was fine all along
最初に通しで動いたとき、1 トークンは 225.6 ms かかりました。 段階10 の見込み(109 ms)の倍です。原因を一つずつ潰した記録がこの表です。
| 何を変えたか | 1 トークン | 効いた理由 |
|---|---|---|
| (参考)CPU の小物を numpy で | 小物だけで 305 ms | ボードの numpy は NEON を使わず、1 要素 20〜100 ns |
| (参考)素直な C に | 小物だけで 55 ms | rintf が関数呼び出し、浮動小数の比較で流れが止まる |
| (参考)C + NEON に | 小物だけで 25 ms | 丸めは小細工、exp は多項式、割り算は逆数の推定 |
| 通しの初版 | 225.6 ms | — |
| PYNQ の配列を素の numpy 配列に | 164〜172 ms | PYNQ の配列は切り出すたびに約 100 µs かかっていた |
| DMA の操作を C に | 125.9 ms | Python からのレジスタ操作は 1 回 6〜14 µs。1 転送で数百 µs になっていた |
DMA の 1 転送が理論値より 150 µs ほど長くかかる、という測定が出ました。 状態レジスタを読み続けると転送そのものが遅れるのだろうと考え、 「終わりそうな時刻までは時計だけ見て待つ」版を作りましたが、何も変わりませんでした。
DMA の操作を C に移すと、転送は理論どおりの時間で終わりました。 ハードは最初から正常で、遅れていたのは Python の側でした。 ただ、「時計で待ってから 1 回だけ読む」実験で出た 250 µs の余分は、いまも説明がついていません。
段階6 でも「回路は道幅の 100% で動いているのに、呼び出しの固定費が 83%」という話がありました。 ここでも同じことが起きています。回路を速くするのと同じくらい、回路への渡し方に手間がかかる、 というのがこの段階でいちばん時間を使ったところです。
only attention grows with T
文脈 512(ボードの KV の置き場の上限)まで生成を続け、 部品ごとの時間を 1 トークンずつ記録しました。
| 文脈長 T | 行列積 | attention | CPU | 合計 | tok/s |
|---|---|---|---|---|---|
| 64 | 87.3 | 5.3 | 26.6 | 119.2 | 8.39 |
| 128 | 87.3 | 7.7 | 26.6 | 121.6 | 8.22 |
| 256 | 87.2 | 12.7 | 26.7 | 126.7 | 7.89 |
| 384 | 87.3 | 17.4 | 26.7 | 131.5 | 7.60 |
| 500〜512 | 87.4 | 22.0 | 26.7 | 136.0 | 7.35 |
単位は ms。各行は前後 8 トークンの中央値。同じ条件でもう一度走らせた別の回では、文脈 512 で 138.0 ms(7.25 tok/s)でした。
where the budget went
段階10 では、部品を別々に測って足し、文脈 512 で 109.2 ms と見込んでいました。 実物は 136.0 ms。差の 26.8 ms がどこから来たかを並べます。
what is measured and what is not
言えること。すべて実機での実測です。
言えないこと。