CPUが得意なことをCPUにまかせて少ないVRAMでも大きめのLLMを速く動かす

Redditに「VRAM足りないとき一部のレイヤーをCPUに任せるんではなく、レイヤー全部をGPUに載せてレイヤー内部のFFNだけCPUに持っていったら速くなった、なんでこれが標準じゃないんだ」というのがあったので、おうちのRTX 4060 Ti 16GBで試してみたら微妙に速くなりました。
https://www.reddit.com/r/LocalLLaMA/comments/1ki7tg7/dont_offload_gguf_layers_offload_tensors_200_gen/

Qwen3 30B A3Bで試してみる

こういった指定がOllamaやLM Studioではできないので、今回はKoboldCPPというので試してます。
https://github.com/LostRuins/koboldcpp

KoboldCPPでは実用が厳しいので、llama.cppで試すほうがよさそう。

とりあえず、LM StudioでQwen3 30B A3Bのq3_k_xlを動かしたときは15.58tok/sec

48中38レイヤーをGPUに割り当てています。

ということで、koboldcppの実行。ダウンロードした実行ファイルに--overridetensorsと--modelと--gpulayersを指定して起動します。

koboldcpp.exe --overridetensors "blk\.([0-9]*[05])\.ffn_.*_exps\.=CPU" --model "D:\dev\gguf\unsloth\Qwen3-30B-A3B-GGUF\Qwen3-30B-A3B-UD-Q3_K_XL.gguf" --gpulayers 48

--overridetensors "blk\.([0-9]*[05])\.ffn_.*_exps\.=CPU"という指定が肝ですね。

0と5で終わるffn内の層がCPUに乗ります。

今回はRedditに書いてあった指定を使っているのだけど、層の名前を確認したいときは正規表現で.*を指定すれば全部CPUに乗るので確認できそう。

http://localhost:5001にアクセスして「bertとgptの違いは」と聞いてみます。

17.55tok/sec！12%速くなりましたね。

メモリ消費はこのくらい。

落としたときに2.2GB使っていたので、11.4GBほど消費してます。これはLM Studioで36レイヤー読み込んだときと同じ。

Llama4 ScoutのQ2_KをLM Studioで16レイヤーをGPUにオフロードした場合とKoboldCPPで--overridetensors "blk\.([0-9]*[0124578])\.ffn_.*_exps\.=CPU"としてFFNだけ2/3ほどCPUに残した場合では、4.1tok/secだったのが4.9tok/secと20%速くなりました。

ただ、思ったより効果がでてないのは、うちのCPUがちょっと弱いからではないかと。強いCPUならもっと効果が出ると思います。
Qwen3 32Bで試したら性能向上できなかったけど、CPUが強ければそれなりに効果が出そう。

何をやっているのか

では何をやっているのか見るためにLLMの基本構造を確認してみましょう。

いまのLLMはトランスフォーマという構造をベースにして、だいたいこんな感じになってる。位置エンコーディング(Posional Encoding)からFeed Forwardまでで一層で、それがQwen 30B A3Bなら48層、Qwen 32Bなら64層という風になってる。

で、LM Studioをはじめ、LLMの実行系の設定では、層単位でGPUにどれだけ乗せるか、逆にCPUにどれだけ残すかというのを設定するようになってる。
でも、層全体で決めるんじゃなくて、層のなかの役割によってCPUでも効率化できるか、GPUじゃないとだめかって決まるんで、CPUでも効率化できるところはCPUに残して、GPUのメリットあるところはなるべくGPUに乗せたほうがいいんでは、って話ですね。

なぜそれがいいのか

じゃあなぜそれがいいのか、って見るのには、実際のコード見るのがいいと思います。

ということで、llama2.cをJavaで書き直したやつをベースに。
https://gist.github.com/kishida/05656bfcbe840f269784f7dbbee5928e

LLMの処理を見るのはforwardメソッド。
https://gist.github.com/kishida/05656bfcbe840f269784f7dbbee5928e#file-llama-java-L300

まず後段になるFeedForwardを見てみます。今回CPUに乗せようというのはこの部分です。

rmsnorm(s.xb, x, w.rms_ffn_weight[l], dim);

// Now for FFN in PyTorch we have: self.w2(F.silu(self.w1(x)) * self.w3(x))
// first calculate self.w1(x) and self.w3(x)
matmul(s.hb, s.xb, w.w1[l], dim, hidden_dim);
matmul(s.hb2, s.xb, w.w3[l], dim, hidden_dim);

// SwiGLU non-linearity
for (int i = 0; i < hidden_dim; i++) {
  // 省略
}

// final matmul to get the output of the ffn
matmul(s.xb, s.hb, w.w2[l], hidden_dim, dim);

SwiGLUのところは省略してますが1重ループです。rmsnormも1重ループになってます。1重ループは基本的に時間がかからないので、高速化の必要性も薄いです。あとはmatmulです。FFNの処理時間はmatmul部分にかかります。

そのmatmulはこんな感じ。

static void matmul(float[] xout, float[] x, FloatBuffer ws, int n, int d) {
    MemorySegment w = MemorySegment.ofBuffer(ws);
    IntStream.range(0, d).parallel().forEach(i -> {
        FloatVector val = FloatVector.zero(SPECIES);
        for (int j = 0; j < n; j+=SIMD_SIZE) {
            FloatVector a = FloatVector.fromMemorySegment(
               SPECIES, w, (i * n + j + SIMD_SIZE) * FLOAT_SIZE, ByteOrder.LITTLE_ENDIAN);
            FloatVector b = FloatVector.fromArray(SPECIES, x, j + 0*SIMD_SIZE);
            val = a.fma(b, val);
        }
        xout[i] = val.reduceLanes(VectorOperators.ADD);
    });
}

細かいところは置いておいて、IntStreamでparallelとしてマルチスレッド化してるところと、その中にループがあってFloatVectorを使ってAVXなどSIMD命令を使うようにしてることだけ見てください。

つまり、スレッドを動かすコア数がそれなりにあってAVXのように1命令で複数のデータを処理できれば、CPUでも速く処理ができます。

一方でマルチヘッドアテンションはこんな感じですね。

// multihead attention. iterate over all heads
final var fl = l;
IntStream.range(0, p.n_heads).parallel().forEach(h -> {
    int qpos = h * head_size;
    int kvpos = h / kv_mul * head_size;
    float[] att = s.att[h];
    for (int t = 0; t <= pos; t++) {
        float score = 0;
        FloatVector val = FloatVector.zero(SPECIES);
        for (int i = 0; i < head_size; i+=SIMD_SIZE) {
            FloatVector a = FloatVector.fromArray(SPECIES, s.q, qpos + i);
            FloatVector b = FloatVector.fromArray(SPECIES, s.key_cache[fl][t], kvpos + i);
            val = a.fma(b, val);
        }
        score = val.reduceLanes(VectorOperators.ADD);
        score /= head_aqrt;
        // save the score to the attention buffer
        att[t] = score;
    }
    ・・・

IntStreamのparallelでマルチスレッド化して、内部にFloatVectorを使ったループがあるのはmatmulと似てるのだけど、FloatVectorを使ったループがループで囲まれて、全体で3重ループになってます。

そして、真ん中のループは特にハードウェアでの高速化がされてないです。CPUだとこれを高速化する仕組みがない。

Intel AMXとかあるけど4世代Xeonにようやく搭載されたところで、普及してない。使えるとLLMが速くなるようです。
インテルの AI 対応 AMX CPU アクセラレータのテスト結果について | Google Cloud 公式ブログ

一方でGPUだと3重ループを速くすることができます。
GPU処理の共通フレームワークであるOpenCLの説明に次のように書いてます。

解きたい問題には全て、直線状やキューブ状や平面状のようにある程度の次元性が存在している。 OpenCLでは最大3次元までを指定してカーネルを展開する。

ここで、サッと3重ループをGPUで効率よく処理したソースが出せるといいんだけど、ディープラーニングをGPU使って速くしようとした処理では、ちゃんと3重ループの処理が書けてなくて高速化できてなかった。
https://github.com/kishida/neuralnet/blob/use_jocl/src/main/resources/kernels/convolution_forward.cl#L15

次のようにiのループとjのループもGPUの並列化に任せるようにすると速くなるはず。

int fxy = get_global_id(0);
int i = get_global_id(1);
int j = get_global_id(2);

[増補改訂]GPUを支える技術 ――超並列ハードウェアの快進撃[技術基礎] (WEB+DB PRESS plus)

作者:Hisa Ando
技術評論社

Amazon