BonsaiのWebGPUカーネルを読み解く

1-bit / 三値 (ternary) 量子化した小さいLLMを、ブラウザの中だけでどこまで速く動かせるか気になって、しばらく手を動かしていました。
きっかけは Bonsai という小さいLLMと、それを開発している PrismML が公開している WebGPU デモです。
このデモが思ったより速くて、中で何が起きているのか気になって、実際に使われているWGSLシェーダのソースを読みに行きました。
この記事はその読み解きの記録です（読んだ後に自分でも手書きカーネルを実装して本家と対決させた話は別記事にまとめました）。

Bonsaiという小さいLLM

Bonsai は重みを {+1, -1} や {+1, 0, -1} のような極端に小さいビット幅へ圧縮した1-bit / 三値量子化のLLMです。
8Bパラメータのモデルが1GB前後に収まり、ブラウザのタブにも載るサイズになります。

このフォーマットには、サイズが小さくなる以上のうまみがあります。
三値の重みを2本のビットプレーン（Pビット: 値が+1ならbitが立つ、Nビット: 値が-1ならbitが立つ）で表現すると、内積が次の恒等式で計算できます。

dot = 2 * (popcount(P & X) - popcount(N & X)) - popcount(P) + popcount(N)

（X は活性化を二値化したビットプレーン）
ビットごとに条件分岐して足し引きする代わりに、AND と popcount というビット単位のバルク演算に置き換えられる、つまり乗算を一切使わずに内積が計算できるのがこの恒等式の核です。
CPUのSIMDにとっては、乗算器を使わずにすむぶんだけ有利になる、特別な構造を持ったフォーマットということになります。

PrismMLがWebGPUで動かせるデモを紹介していた

PrismML は Bonsai を実際にブラウザで動かせるデモを公開しています。

• webml-community/bonsai-ternary-webgpu（実際に動かせるデモ）
• Try It Without Installing / Bonsai Demos コレクション

ヘッドレス Chrome で実際にロードしてプロンプトを送ってみると、UI上に表示される速度はこうでした。

プロンプト: "What is the capital of France?"
応答: "The capital of France is Paris. ..."
表示された速度: 43.0 TOKENS / SEC

ブラウザの中だけで43 tok/sは、正直かなり速い数字です。
このモデルを手元でネイティブビルドしたPrismML自身の llama.cpp フォーク（Apple M3, NEON, シングルスレッド）が16.31 tok/sだったので、ブラウザ版がネイティブのシングルスレッド版を上回っていることになります。
何をしたらブラウザでここまで出るのか、気になったので中を覗いてみることにしました。

WebGPUを利用しているコードの確認

このデモのJSバンドルには、実際に使われているWGSLシェーダのソースがテンプレート文字列としてそのまま埋め込まれていました。
数えると52個。RMSNorm、QK-Norm+RoPE、KVキャッシュ書き込み、Flash Attention風のsplit-K decode attention、argmax付きlm_head、SwiGLU……推論に必要な処理がひととおり自前のWGSLで実装されています。
外部ライブラリではなく、PrismML独自のcompute-opランタイムのようで、デバイスの対応状況（subgroup対応の有無など）を見て実行時にWGSLをレンダリングする作りになっていました。

その中でいちばん気になったのが、いちばん呼ばれる回数が多いはずの三値行列積のカーネルです。
覗いてみると、1トークンずつ処理するdecodeと、複数トークンをまとめて処理するprefillとで、まったく違う計算方法を使っていました。

decode側（llama-decode-q1-proj.wgsl）は、1本のu32にパックされた32個の符号ビットをvec4<f32>単位で±1.0に展開し、GPUネイティブのdot()（ハードウェアのFMA）にそのまま渡していました。
自分がCPU側で書いていた「popcountで乗算を避ける」発想とは逆で、乗算を避けるどころか、わざわざ実数へ展開してGPUの得意な乗算器に投げているわけです。
理にかなった話で、CPUは乗算器のスループットがボトルネックになりやすいのに対し、GPUには専用の行列演算ハードウェアがあって乗算の生スループットはそもそも桁違いに大きい。
ボトルネックはむしろメモリ帯域とハードウェアへのデータ供給効率のほうにあるので、乗算を避けること自体には価値がありません。

一方prefill側（binary-lut4-matmul.wgsl）はさらに踏み込んでいて、乗算どころか加算まで減らしにいく設計でした。
活性化4個ぶんが取りうる16通りの部分和（0, x, y, x+y, z, ...）をワークグループ共有メモリ上のLUTへあらかじめ計算しておき、重みの4bitニブル1個をLUT参照1回に置き換えます。
これは同じM個の出力行に対してLUTを使い回せる、つまり複数トークンをまとめて処理する（=計算律速になる）prefillだからこそ割に合う設計で、1トークンずつのdecodeでLUTを作っても使い回す相手がおらず割に合いません。

さらにもう一段、subgroupMatrixMultiplyAccumulate（GPUのテンソルコア相当の協調行列演算命令）を使う第3の経路もありました。
つまり1つのカーネルではなく、形状とハードウェアの対応状況に応じて3系統のカーネルを実行時に自動選択しているということになります。

行列積カーネル以外にも面白い実装があったので、2つだけ実際のコードを引用して紹介します。

KVキャッシュ書き込み: RoPEと同じdispatchの中で終わらせる

decode-partial-rope-cache-kv.wgslは、新しいトークン1個ぶんのK/VをRoPEで回転させながら、そのままキャッシュ配列へ書き込みます。
キャッシュ上の位置計算は別ファイル（kv-cache-index.wgsl）に切り出されていて、こうなっていました。

fn cache_index(pos: u32, head: u32, dim: u32) -> u32 {
  return (((LAYER * CACHE_LEN + pos) * NUM_KV_HEADS + head) * HEAD_DIM + dim);
}

これを使う本体側（K/Vの書き込み部分だけ抜粋、条件分岐は一部省略）はこうです。

fn rope_index(pos: u32, d: u32) -> u32 {
  return pos * ROTARY_HALF + d;
}

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let index = gid.x;
  let pos = params.past_len;
  // ...（Qの回転はここより前、index < Q_SIZE の分岐で行われる）

  let head = index / HEAD_DIM;
  let d = index % HEAD_DIM;
  if (d < ROTARY_HALF) {
    let c = rope_cos[rope_index(pos, d)];
    let s = rope_sin[rope_index(pos, d)];
    let first = k[base + d];
    let second = k[base + d + ROTARY_HALF];
    cache_keys[cache_index(pos, head, d)] = first * c - second * s;
    cache_keys[cache_index(pos, head, d + ROTARY_HALF)] = second * c + first * s;
  }
  cache_values[cache_index(pos, head, d)] = v[index];
}

RoPEの回転とキャッシュへの書き込みが、同じ64スレッドのdispatch 1回の中で完結しています。
「回転してから書く」という2つの処理を別々のシェーダに分けない、地味だけど確実にディスパッチ回数を減らす工夫です。

Flash Attention風のsplit-K decode attention

decodeは新しいクエリが1トークンぶんしかないので、通常のFlash Attentionのようにクエリ方向で並列化する余地がありません。
代わりにPrismMLは、KVキャッシュに積まれた過去のトークン列（KV方向）をNUM_SPLITS個に分割し、split同士を並列に走らせていました（attn-flash-decode-splitk.wgsl）。

// KV列をNUM_SPLITS個のワークグループに分割し、このワークグループは
// [splitStart, splitEnd) の区間だけを担当する
let keysPerSplit = (kvSeq + NUM_SPLITS - 1u) / NUM_SPLITS;
let splitStart = split * keysPerSplit;
var splitEnd = min(splitStart + keysPerSplit, kvSeq);

var runningMax = -FLT_MAX;
var runningDenom = 0.0;

var kjBase = splitStart;
loop {
  if (kjBase >= splitEnd) { break; }

  // このタイル(最大WG個のキー)についてQ・Kの内積スコアを計算
  // ...(score, m, dPart を求める。実装は省略)
  let tile = combine_partials(m, dPart, tid, sgSize);

  // オンラインsoftmax: 新しい最大値に合わせて分母と出力を再スケールする
  let newMax = max(runningMax, tile.x);
  let correction = exp_shift(runningMax, newMax);
  runningDenom = runningDenom * correction + tile.y * exp_shift(tile.x, newMax);
  runningMax = newMax;

  // running_out も同じ correction で再スケールしてからVの加重和を足し込む
  running_out[d4] = running_out[d4] * correction + vSum;

  kjBase = kjBase + WG;
}

// このsplitが担当した区間の部分結果(部分出力・部分max・部分denom)を書き出す
partial_out[partialBase + d4] = running_out[d4];
partial_max[mdBase] = runningMax;
partial_denom[mdBase] = runningDenom;

各splitは自分が担当する区間だけを見てオンラインsoftmax（ランニング最大値とランニング分母）を独立に計算し、部分結果を書き出して終わります。
最後にattn-flash-decode-splitk-merge.wgslが、各splitの部分結果をlog-sum-exp形式で正しく合成します。

var globalMax = -FLT_MAX;
for (var s: u32 = 0u; s < NUM_SPLITS; s = s + 1u) {
  globalMax = max(globalMax, partial_max[mdBase + s]);
}
var globalDenom = 0.0;
var outV = vec4<f32>(0.0);
for (var s: u32 = 0u; s < NUM_SPLITS; s = s + 1u) {
  let w = exp_shift(partial_max[mdBase + s], globalMax);
  globalDenom = globalDenom + partial_denom[mdBase + s] * w;
  outV = outV + partial_out[pBase + d4] * w;
}
output[qBaseV4 + d4] = vec4<f32>(outV / globalDenom);

「1つの最大値・1つの分母で正規化する」というsoftmaxの定義上、区間ごとに別々に計算した部分結果を後からナイーブに足し合わせることはできません（区間ごとに基準にした最大値がずれているため）。
exp_shift(partial_max[s], globalMax)で全splitの基準をglobalMaxに揃え直してから合成する、というのがこのmergeカーネルの役目です。
KVキャッシュが長くなるほどNUM_SPLITSを増やして並列度を上げられる、decodeならではの設計だと思います。

正直、ここまで作り込まれているとは思っていなかったので、単純に感心しました。
そして同時に、「じゃあCPU側の"乗算を避ける"発想は、GPUの土俵では完全に無意味なのか、それともCPU/Wasm側では依然として意味があるのか」という素朴な疑問が湧いてきました。

decode側がGPUネイティブのFMAに展開して投げ、prefill側がLUTに置き換える——CPU側で自分が実装しようとしている「popcountで乗算を避ける」発想とは、そもそも刺さる場所が違うのではないか。
実際にWasm SIMDでカーネルを書いて確かめた話は、別記事にまとめました。

まだ読めていないところ

subgroupMatrixMultiplyAccumulateを使う第3の経路

デバイスに合わせてWGSLを組み立てるランタイム