言語モデルを高位合成でFPGAに実装してみた

Turing株式会社のリサーチチームでインターンしているM1の内山です。

Turing株式会社では大規模基盤モデルによる完全自動運転を目指しており、その実現に欠かせない技術として大規模言語モデルの研究開発を行っています。

Generative AI

LLMの広範な知識と思考能力に加え、視覚情報やセンサーデータなどの多様な入力を受け入れることで、車の周囲の状況を正確に認識します。さらに、世界モデルを適用することで、高度な空間認知と身体性を獲得し、実世界に対応した生成AIを実現します。

https://tur.ing/ より引用

しかしながら、従来の大規模モデルはデータセンターという大量のGPU・潤沢な電源・安定した地盤を備えた豊かな環境で処理されるものであり、対して自動車というものは余りにも狭く、電源が乏しく、振動が大きいという劣悪極まりない環境です。

そんな大規模基盤モデルにとって地獄とも言える環境で、実用的なレベルの推論速度を出すためには、GPUのような汎用性を備えた生半可な計算機ではなく、完全自動運転用の大規模基盤モデルの実行に特化した専用計算機、アクセラレーターが間違いなく必要になります。

そこでTuringのリサーチチームでは、大規模基盤モデルの研究開発と共に、大規模基盤モデル向けアクセラレータの研究開発を日夜行っています。

本記事ではその小手調べとして中規模FPGAのKV260の上で、HLSを活用して言語モデルを動作させましたので、その作業日誌としてここに残します。

CPUによるLLM推論

LLMの推論にはNVIDIAの高性能なGPUが必要なのが社会常識になりつつある昨今、手元の計算機でLLMを動かすことはある種のロマンと化しています。しかし実際には、モデルが十分に小規模であれば、ノートパソコンのCPUであっても高速にLLMを動かすことが可能です。

手元の計算機でLLMを動かせるソフトウェアとして llama.cpp が代表的です。これはリポジトリをGitHubからクローンしてMakeするだけで、手元のパソコンでLLMが動き出すという素晴らしい実装になっています。

https://github.com/ggerganov/llama.cpp

ノートパソコンで動作するSwallow 7B on llama.cpp

非力な計算機でも llama.cpp のお陰で気軽にLLMを動かせるようになった一方で、その実装はLLMの発展と共に複雑化しており、LLM実装の入門にはもはや不向きとなっています。

LLM実装を追う上でおすすめなのが llama2.c です。これはAndrej Karpathy氏によって作られたLLMの実装であり、700行程度のC言語で書かれているため、実装を追うのも非常に容易です。また使用しているモデルもTinyStoriesで訓練した15Mモデル(60MB程度)であり、どんなに非力な計算機でも軽快に動作するのが特徴です。

https://github.com/karpathy/llama2.c

llama2.c

FPGAによるLLM推論：Swan

我々はこのこのllama2.cを参考に、Transformerとその本質的な演算だけをFPGA用に実装したswanを作成しました。

https://github.com/turingmotors/swan

Swanのターゲットとして、エントリ向けのFPGAボード「KV260」で試していきたいと思います。

Kria KV260

以降ではこのswanを利用しますので、事前にリポジトリをcloneしてCPUでの動作確認を完了させておいてください。

$ git clone git@github.com:turingmotors/swan.git
$ cd swan
$ mkdir build && cd build
$ cmake ..
$ make
$ cd ..
$ ./build/swan

実装方針

Kria KV260はZynq UltraScale+ MPSoCが搭載された、45,000円と比較的安価でかつ容易に入手可能なFPGAボードです。ZynqはFPGAにArm Cortex-A53が結合した構造をしているため、Armで動作しているLinuxから、FPGAに生成されたハードウェアを制御することが可能です。

FPGAで推論を行うにあたり、LLMが全てFPGAのBRAMに収まってくれるのであれば、かなり自由に推論器のアーキテクチャを考えることが出来るのですが、残念ながらKV260ではBRAMとUltraRAMを合わせても2.7MB程度であり、60MBのモデルは到底入りません。

となるとKV260に載っているDRAM(DDR4 4GB)にモデルを格納し、必要に応じてウェイトを取り出す実装になります。これではメモリ律速になりますが、そこはエントリ向けのハードウェアとして割り切り、今回はLLMを動作させる事だけに集中しましょう。

Kria KV260でも実装可能なLLM計算機のアーキテクチャを考えてみると、KV260のFPGAにTransformerのデコーダ部で使う各種関数をオフロードし、CPUがそれらのハードウェアを利用するアーキテクチャが最も容易に実装できそうです。

LLM計算機のアーキテクチャ

今回はこのアーキテクチャで進めていきます。実装には開発の効率化のため、プログラミング言語からハードウェアを生成する、高位合成(High Level Synthesis, HLS)を使用します。

KV260の開発環境を整える

対象環境：Ubuntu 20.04

Kria KV260とHLSで開発を進める前に、叩き台となるVitisワークスペースを作成します。

後述する注意事項を頭に入れ、以下のXilinxが用意しているチュートリアルをvaddを実行する所まで進めます。

注意事項１：使用するツールはVitis Classicです。Vitisではありません。--classic オプションを付けて起動できます。

注意事項２：KV260で起動するLinuxのSDカードイメージは以下からダウンロードが可能です。事前にbalenaEtcher等でSDカードに書き込みます。

注意事項３：Xilinxが提供しているチュートリアルでは、以下の画像の赤矢印で指している場所に見られるように、コードブロックの改行が消えている箇所が複数存在します。チュートリアルに従ってコマンドを打つ場合は、左上にある「ページソースの表示」のリンクを押し、そこにあるコマンドを参考にしてください。

壊れているコードブロック

注意事項４：手順２のソフトウェアコンポーネントの作成の中の、”5. rootfs に XRT を追加の操作”において、xrtの他にzocl, opencl-clhpp, opencl-headersも追加してください。

Xilinxのチュートリアルを完走したら、LLM計算機の実装に移ります。

swanアプリケーションプロジェクトを作る

最初に、CPUで動かすホストプログラムとFPGAに生成するカーネルを統括するアプリケーションプロジェクトを作成します。

適当な場所に kv260_swan ディレクトリを作成し、そこでvitis classicを起動します。

$ source /tools/Xilinx/Vitis/2023.2/settings64.sh
$ mkdir kv260_swan
$ cd kv260_swan
$ vitis --classic --workspace .

Vitis classicが起動したら、チュートリアルで作成したカスタムプラットフォームを追加します。起動画面の PLATFORM > Add Custom Platform から、チュートリアルで作成したkv260_customディレクトリを選択して追加します。

カスタムプラットフォームの追加が完了したら、PROJECT > Create Application Projectを選択してアプリケーションプロジェクトの作成に入ります。

アプリケーションの作成ウィンドウが出てきたらNextを押し、Select a platform from repository で先程追加したkv260_customを選択します。

kv260_customを選択

続いて、Application Project Name を swan にします。この時、System project nameは自動的に swan_system になります。

次の画面では、Sysroot Path, Root FS, Kernel Imageにそれぞれ cortexa72-cortexa53-xilinx-linux, rootfs.ext4, Imageを設定します。これらのファイルはチュートリアルで作成したものを使用します。

次にテンプレートの選択画面に入りますが、ここではチュートリアルとは異なり Emtpy Application を指定します。Finishを押せばアプリケーションプロジェクトの作成が完了します。

swanのデコーダ部で使われている関数

swanのデコーダであるdecode.cppを覗いてみると、llama2のデコーダで使われる関数は主に６つです。

matmul - 行列ベクトル積
RMS norm - 正規化
Softmax - ソフトマックス
Add - ベクトル加算
Mul - ベクトル乗算
RoPE - 位置エンコーディング

これら６つの関数をHLSでFPGAにオフロードすることにします。

カーネルプログラムとホストプログラム

KV260に対してHLSで開発を進める上で、OpenCL を避けて通ることは出来ません。

vaddのサンプルコードを見ると、krnl_vadd.cppはHLSらしいコードで書かれているのに対して、vadd.cppについては見慣れないAPIを大量に使用しています。

OCL_CHECK(err, cl::Buffer buffer_a(context, CL_MEM_READ_ONLY, size_in_bytes, NULL, &err));
OCL_CHECK(err, cl::Buffer buffer_b(context, CL_MEM_READ_ONLY, size_in_bytes, NULL, &err));
OCL_CHECK(err, cl::Buffer buffer_result(context, CL_MEM_WRITE_ONLY, size_in_bytes, NULL, &err));

これらは OpenCL と呼ばれるAPIです。

OpenCLとは、様々な形態の計算機システムで並列計算を行うためのAPIであり、マルチコアCPU、GPGPU、FPGAなど様々な計算資源に対応しています。KV260において、FPGAにHLSで実装したモジュールをCPUから制御したい場合、このOpenCL APIを使います。

https://www.khronos.org/opencl/

OpenCLにおいて、計算機システムは制御用のホストと計算用のデバイスに別れており、ホスト側で動く制御用のプログラムの事をホストプログラムと呼び、デバイスで動くプログラムの事をカーネルプログラムと呼びます。

CPUとGPGPUが載った計算機システムを例に取ると、CPUはGPGPUを制御する立場にあるので、CPUで動くプログラムはホストプログラムであり、GPGPUで走るプログラムはカーネルプログラムに対応します。

FPGAにおいては、KV260に載っている Zynq™ UltraScale+™ MPSoC は Arm の Cortex-A53 とFPGAが結合した構造をしており、HLSで実装されたFPGA上のハードウェアがデバイス、Armのコアがホストに対応しています。

ホストとカーネル

以上より、HLSをKV260で使うためには、ホストプログラムとカーネルプログラムの２種類のプログラムを書く必要がありますが、多くのコードをサンプルコードから流用できるため、OpenCL APIを深く理解しなくてもFPGAにハードウェアとして生成されたカーネルを、ホストプログラムから利用できます。

各種カーネルプログラムをHLSで実装する

ここでは各種カーネルプログラムを実装していきます。なお、ここで示されている関数の実装は全てswanリポジトリに置いてありますので、適宜参照してください。

https://github.com/turingmotors/swan

matmul - 行列ベクトル積

HLSを用いて実装した行列ベクトル積が以下の通りです。

kernel_matmul()は引数として、ベクトルと行列の入力データが格納されているDRAMのアドレスを指すポインタと、演算結果を格納するDRAMのアドレスを指すポインタ、そして行と列のサイズを受け取っています。

compute_matmul()が実際に演算を行っている部分で、ベクトルの入力をローカルのレジスタに保存した後、ストリーミングされてくる行列データに対して積和演算を行っています。それ以外はDRAMからデータの受け渡しを行う関数です。

#include <stdint.h>
#include <hls_stream.h>

#define MAX_DATA_SIZE 1024

static void load_vec(float* i_vec, hls::stream<float>& inStream, int vec_size) {
mem_rd:
    for (int i = 0; i < vec_size; i++) {
    	inStream << i_vec[i];
    }
}

static void load_mat(float *i_mat, hls::stream<float>& inStream, int vec_size, int col_size) {
mem_rd:
    for (int i = 0; i < col_size; i++) {
      for (int j = 0; j < vec_size; j++) {
        inStream << i_mat[vec_size * i + j];
      }
    }
}

static void compute_matmul(hls::stream<float>& in1_stream,
                        hls::stream<float>& in2_stream,
                        hls::stream<float>& out_stream,
                        int vec_size,
                        int col_size) {

	float vec_local[MAX_DATA_SIZE];
	float sum_local = 0;
	for(int i = 0; i < vec_size; i++) {
		vec_local[i] = in1_stream.read();
	}
execute:
    for (int i = 0; i < col_size; i++) {
      for(int j = 0; j < vec_size; j++) {
        #pragma HLS UNROLL
        sum_local += vec_local[j] * in2_stream.read();
      }
      out_stream << sum_local;
      sum_local = 0;
    }
}

static void store_result(float* out, hls::stream<float>& out_stream, int col_size) {
mem_wr:
    for (int i = 0; i < col_size; i++) {
        out[i] = out_stream.read();
    }
}

extern "C" {
void kernel_matmul(float *i_vec, float *i_mat, float *o_vec, int vec_size, int col_size) {
#pragma HLS INTERFACE m_axi port = i_vec bundle = gmem0
#pragma HLS INTERFACE m_axi port = i_mat bundle = gmem1
#pragma HLS INTERFACE m_axi port = o_vec bundle = gmem0

    static hls::stream<float> vec_stream("vec_stream");
    static hls::stream<float> mat_stream("mat_stream");
    static hls::stream<float> out_stream("out_stream");

#pragma HLS dataflow
    load_vec(i_vec, vec_stream, vec_size);
    load_mat(i_mat, mat_stream, vec_size, col_size);
    compute_matmul(vec_stream, mat_stream, out_stream, vec_size, col_size);
    store_result(o_vec, out_stream, col_size);
}
}

RMS norm - 正規化

続いて正規化関数の実装であるkernel_rmsnorm() ですが、これは２つのベクトル入力を受け取り、１つのベクトルを出力する関数であり、引数としてそれぞれのポインタを受け取っています。

compute_rmsnorm() が計算を行う関数であり、両方ともローカルのジレスタに格納した上で計算を行っています。

#include <stdint.h>
#include <math.h>
#include <hls_stream.h>

static void load_vec(float* i_vec, hls::stream<float>& inStream, int vec_size) {
mem_rd:
    for (int i = 0; i < vec_size; i++) {
    	inStream << i_vec[i];
    }
}

static void compute_rmsnorm(hls::stream<float>& in1_stream,
                        hls::stream<float>& in2_stream,
                        hls::stream<float>& out_stream) {

	float vec_local_1[288];
	float vec_local_2[288];
	float sum_local = 0;
	for(int i = 0; i < 288; i++) {
		vec_local_1[i] = in1_stream.read();
	}
	for(int i = 0; i < 288; i++) {
		vec_local_2[i] = in2_stream.read();
	}

  for(int i = 0; i < 288; i++) {
    sum_local += vec_local_1[i] * vec_local_1[i];
  }

  constexpr float eps = 1e-5;
  const float norm = 1 / std::sqrt(sum_local / 288 + eps);

  for (size_t i = 0; i < 288; i++) {
    out_stream << vec_local_1[i] * norm * vec_local_2[i];
  }
}

static void store_result(float* out, hls::stream<float>& out_stream, int vec_size) {
mem_wr:
    for (int i = 0; i < vec_size; i++) {
        out[i] = out_stream.read();
    }
}

extern "C" {
void kernel_rmsnorm(float *i_vec_1, float *i_vec_2, float *o_vec, int vec_size) {
#pragma HLS INTERFACE m_axi port = i_vec_1 bundle = gmem0
#pragma HLS INTERFACE m_axi port = i_vec_2 bundle = gmem1
#pragma HLS INTERFACE m_axi port = o_vec bundle = gmem0

    static hls::stream<float> vec_stream_1("vec_stream_1");
    static hls::stream<float> vec_stream_2("mat_stream_2");
    static hls::stream<float> out_stream("out_stream");

#pragma HLS dataflow
    load_vec(i_vec_1, vec_stream_1, vec_size);
    load_vec(i_vec_2, vec_stream_2, vec_size);
    compute_rmsnorm(vec_stream_1, vec_stream_2, out_stream);
    store_result(o_vec, out_stream, vec_size);
}
}

Softmax, Add, Mul, RoPEも同様に実装します。

ビットストリームをビルドする

ここではKV260のPLにハードウェアを生成するためのビットストリームのビルドを行います。

最初に、左上の Explorer から swan_kernels/src を右クリックし、Import Sources を選択します。すると選択ウィンドウが開くので、Browse… で swanリポジトリの srcディレクトリを選択し、kernel_*.cppと名前の付くファイルを追加します。

ファイルを追加できたら、右上のExplorerから swan_kernels/swan_kernels.prj を選択し、丸いアイコンの Add Hardware Functionsを押してkernel_*と名の付く関数を選択します。

カーネル関数を指定出来たら、右上の Explorer から swan_system_hw_link/swan_system_hw_link.prj を選択した上で、右上の Active build configuration を Hardware に指定した後、金槌のアイコンをクリックしてビットストリームのビルドを開始します。これには数分掛かります。

ビルドが完了すると、swan_system_hw_link/Hardware 以下に binary_container_1.xclbinが生成されています。これがビットストリームです。

Vitis analyzerにて生成されたビットストリームを確認すると、各種関数がFPGA上に配置配線されているのが確認できます。

$ vitis_analyer --classic

ホストプログラムを書く

続いてCPUで動作する制御用のプログラム、ホストプログラムを作成します。

ホストプログラムは大別して以下の５ステップに別れています。

初期設定
メモリ確保
引数設定
実行及び結果読み出し

知らないAPIを使うのに、最低でも５ステップ掛かるのは面倒に思えるかもしれませんが、多くの部分をサンプルコードから流用できるため、OpenCL APIを手書きする量は大して多くありません。

1. 初期設定

初期設定ではOpenCLの各種オブジェクトを宣言します。以下のコードではデバイスやコンテキスト、キューやプログラムオブジェクト等を宣言していますが、ここは殆どのOpenCLホストプログラムで共通している部分であり、自分で書く必要があるのはcl::Kernelで宣言されているカーネルオブジェクトの部分のみです。

ここでは先程作成したカーネル関数に対応するカーネルオブジェクトを宣言しています。

  std::string xclbinFilename = "./binary_container_1.bin";

  size_t size_in_bytes = MAX_DATA_SIZE * sizeof(float);

  std::vector<cl::Device> devices;
  cl_int err;
  cl::Context context;
  cl::CommandQueue q;
  cl::Kernel kernel_matmul;
  cl::Kernel kernel_mul;
  cl::Kernel kernel_rmsnorm;
  cl::Kernel kernel_softmax;
  cl::Kernel kernel_add;
  cl::Kernel kernel_rope;
  cl::Program program;
  std::vector<cl::Platform> platforms;

続いて、先程宣言したオブジェクトに情報を入力していきます。

ビットストリームファイルを読み込んでプログラムオブジェクトを作成し、そのプログラムオブジェクトを用いて各種カーネルオブジェクトがFPGA上のカーネルにアクセス出来るようにしています。

 // Load xclbin
  std::cout << "Loading: '" << xclbinFilename << "'\n";
  std::ifstream bin_file(xclbinFilename, std::ifstream::binary);
  bin_file.seekg(0, bin_file.end);
  unsigned nb = bin_file.tellg();
  bin_file.seekg(0, bin_file.beg);
  char* buf = new char[nb];
  bin_file.read(buf, nb);
  cl::Program::Binaries bins;
  bins.push_back({buf, nb});
  
  bool valid_device = false;
  for (unsigned int i = 0; i < devices.size(); i++) {
      auto device = devices[i];
      // Creating Context and Command Queue for selected Device
      OCL_CHECK(err, context = cl::Context(device, nullptr, nullptr, nullptr, &err));
      OCL_CHECK(err, q = cl::CommandQueue(context, device, CL_QUEUE_PROFILING_ENABLE, &err));
      std::cout << "Trying to program device[" << i << "]: " << device.getInfo<CL_DEVICE_NAME>() << std::endl;
      cl::Program program(context, {device}, bins, nullptr, &err);
      if (err != CL_SUCCESS) {
          std::cout << "Failed to program device[" << i << "] with xclbin file!\n";
      } else {
          std::cout << "Device[" << i << "]: program successful!\n";
          OCL_CHECK(err, kernel_matmul = cl::Kernel(program, "kernel_matmul", &err));
          std::cout << "load : kernel_matmul" << std::endl;
          OCL_CHECK(err, kernel_mul = cl::Kernel(program, "kernel_mul", &err));
          std::cout << "load : kernel_mul" << std::endl;
          OCL_CHECK(err, kernel_rmsnorm = cl::Kernel(program, "kernel_rmsnorm", &err));
          std::cout << "load : kernel_rmsnorm" << std::endl;
          OCL_CHECK(err, kernel_softmax = cl::Kernel(program, "kernel_softmax", &err));
          std::cout << "load : kernel_softmax" << std::endl;
          OCL_CHECK(err, kernel_add = cl::Kernel(program, "kernel_add", &err));
          std::cout << "load : kernel_add" << std::endl;
          OCL_CHECK(err, kernel_rope = cl::Kernel(program, "kernel_rope", &err));
          std::cout << "load : kernel_rope" << std::endl;
          valid_device = true;
          break; // we break because we found a valid device
      }
  }

2. メモリ確保

次にホストとカーネル間のデータ転送に用いるメモリ空間を確保します。

以下ではcl::Bufferでメモリのアクセス管理、メモリサイズを指定してメモリオブジェクトを作成しており、buffer_aからbuffer_dはデータ入力に使うWrite Onlyな領域を確保し、buffer_result, buffer_result2ではデータ出力に使うRead Onlyな領域を確保しています。

  OCL_CHECK(err, cl::Buffer buffer_a(context, CL_MEM_READ_ONLY, size_in_bytes, NULL, &err));
  OCL_CHECK(err, cl::Buffer buffer_b(context, CL_MEM_READ_ONLY, size_in_bytes * size_in_bytes, NULL, &err));
  OCL_CHECK(err, cl::Buffer buffer_c(context, CL_MEM_READ_ONLY, size_in_bytes, NULL, &err));
  OCL_CHECK(err, cl::Buffer buffer_d(context, CL_MEM_READ_ONLY, size_in_bytes, NULL, &err));
  OCL_CHECK(err, cl::Buffer buffer_result(context, CL_MEM_WRITE_ONLY, size_in_bytes, NULL, &err));
  OCL_CHECK(err, cl::Buffer buffer_result2(context, CL_MEM_WRITE_ONLY, size_in_bytes, NULL, &err));

メモリオブジェクトを生成したら、確保したメモリ空間にホストプログラムがアクセスするためのポインタを、enqueueMapBuffer() を用いて取得しています。

 // We then need to map our OpenCL buffer5 to get the pointers
  float* ptr_a;
  float* ptr_b;
  float* ptr_c;
  float* ptr_d;
  float* ptr_result;
  float* ptr_result2;
  OCL_CHECK(err,
            ptr_a = (float*)q.enqueueMapBuffer(buffer_a, CL_TRUE, CL_MAP_WRITE, 0, size_in_bytes, NULL, NULL, &err));
  OCL_CHECK(err,
            ptr_b = (float*)q.enqueueMapBuffer(buffer_b, CL_TRUE, CL_MAP_WRITE, 0, size_in_bytes * size_in_bytes, NULL, NULL, &err));
  OCL_CHECK(err,
            ptr_c = (float*)q.enqueueMapBuffer(buffer_c, CL_TRUE, CL_MAP_WRITE, 0, size_in_bytes, NULL, NULL, &err));
  OCL_CHECK(err,
            ptr_d = (float*)q.enqueueMapBuffer(buffer_d, CL_TRUE, CL_MAP_WRITE, 0, size_in_bytes, NULL, NULL, &err));
  OCL_CHECK(err, ptr_result = (float*)q.enqueueMapBuffer(buffer_result, CL_TRUE, CL_MAP_READ, 0, size_in_bytes, NULL,
                                                         NULL, &err));
  OCL_CHECK(err, ptr_result2 = (float*)q.enqueueMapBuffer(buffer_result2, CL_TRUE, CL_MAP_READ, 0, size_in_bytes, NULL,
                                                         NULL, &err));

3. 引数設定

メモリ確保が完了したら、各カーネルへ引数を設定します。カーネルへの引数の設定にはカーネルオブジェクトのsetArgメソッドを用います。setArgメソッドでは第一引数は設定先の引数のインデックスを取り、第二引数には引数そのものを取ります。

  // set the kernel Arguments
  OCL_CHECK(err, err = kernel_rmsnorm.setArg(0, buffer_a));
  OCL_CHECK(err, err = kernel_rmsnorm.setArg(1, buffer_b));
  OCL_CHECK(err, err = kernel_rmsnorm.setArg(2, buffer_result));

  OCL_CHECK(err, err = kernel_matmul.setArg(0, buffer_a));
  OCL_CHECK(err, err = kernel_matmul.setArg(1, buffer_b));
  OCL_CHECK(err, err = kernel_matmul.setArg(2, buffer_result));
  OCL_CHECK(err, err = kernel_matmul.setArg(3, swan::kDim));
  OCL_CHECK(err, err = kernel_matmul.setArg(4, swan::kDim));

  OCL_CHECK(err, err = kernel_mul.setArg(0, buffer_a));
  OCL_CHECK(err, err = kernel_mul.setArg(1, buffer_b));
  OCL_CHECK(err, err = kernel_mul.setArg(2, buffer_result));

  OCL_CHECK(err, err = kernel_add.setArg(0, buffer_a));
  OCL_CHECK(err, err = kernel_add.setArg(1, buffer_b));
  OCL_CHECK(err, err = kernel_add.setArg(2, buffer_result));

  OCL_CHECK(err, err = kernel_softmax.setArg(0, buffer_a));
  OCL_CHECK(err, err = kernel_softmax.setArg(1, buffer_result));

  OCL_CHECK(err, err = kernel_rope.setArg(0, buffer_a));
  OCL_CHECK(err, err = kernel_rope.setArg(1, buffer_b));
  OCL_CHECK(err, err = kernel_rope.setArg(2, buffer_c));
  OCL_CHECK(err, err = kernel_rope.setArg(3, buffer_d));
  OCL_CHECK(err, err = kernel_rope.setArg(4, buffer_result));
  OCL_CHECK(err, err = kernel_rope.setArg(5, buffer_result2));
  OCL_CHECK(err, err = kernel_rope.setArg(6, 1));

上記のコードではメモリ節約のため、作成した６つのカーネル関数で同じメモリ空間を使い回しています。

カーネルでメモリ空間を共有している

これでカーネルを実行するための前準備は終了です。

4. 実行及び結果読み出し

以上の手順で作成したOpenCLのオブジェクトを、main.cppではDecode()に渡しています。

    // 6-1. Load the context input and decode the next token.
    swan::CopyTensor1d(ctx_input, tok_emb_table[token]);
    swan::Decode(token, pos, ctx_input, ctx_k_cache, ctx_v_cache,
                 ctx_final_norm, weights, q, kernel_matmul, kernel_mul, kernel_rmsnorm, kernel_softmax, kernel_add, kernel_rope, 
                 ptr_a, ptr_b, ptr_c, ptr_d, ptr_result, ptr_result2,
                 buffer_a, buffer_b, buffer_c, buffer_d, buffer_result, buffer_result2);
    
    // 6-2. Calculate the logits and softmax.
    swan::MutmulVocab(ctx_logits, ctx_final_norm, tok_emb_table);

そしてdecode.cppでは、渡されたOpenCLのオブジェクトを、各種演算を行う関数へ更に渡しています。

    // -- Attention --

    // 1. RMS Normalize
    RMSNormFPGA(ctx.attn_norm[i_layer], attn_input, w.rms_att_w[i_layer], q, kernel_rmsnorm, ptr_a, ptr_b, ptr_result, buffer_a, buffer_b, buffer_result);

    // 2. Weight Multiple
    MatmulFPGA(ctx.attn_wqx[i_layer], ctx.attn_norm[i_layer], w.attn_wq[i_layer], q, kernel_matmul, ptr_a, ptr_b, ptr_result, buffer_a, buffer_b, buffer_result);
    MatmulFPGA(ctx.attn_wkx[i_layer], ctx.attn_norm[i_layer], w.attn_wk[i_layer], q, kernel_matmul, ptr_a, ptr_b, ptr_result, buffer_a, buffer_b, buffer_result);
    MatmulFPGA(ctx.attn_wvx[i_layer], ctx.attn_norm[i_layer], w.attn_wv[i_layer], q, kernel_matmul, ptr_a, ptr_b, ptr_result, buffer_a, buffer_b, buffer_result);

OpenCLのオブジェクトを渡された各種関数では、OpenCL APIを用いて対応するカーネルの実行を行っています。

行列ベクトル積の関数であるMatmulFPGA() を例にカーネルの実行方法を説明すると、確保したメモリ空間を示すポインタを用いて、DRAMに入力データを格納し、enqueueMigrateMemObjects()でDRAMからデバイスへデータ転送を行います。その後enqueueTask()に実行したいカーネルオブジェクトを指定することでカーネルを実行できます。

実行後、enqueueMigrateMemObjets()を用いてデバイスからDRAMへデータ転送を行い、finish()で全ての処理が完了するまで待機します。

その後、ホストプログラムは確保したメモリ空間を指すポインタからデータを読み出せます。

void MatmulFPGA(Tensor1d& out, const Tensor1d& in, const Tensor2dAttn& w,
                 cl::CommandQueue q, cl::Kernel kernel_matmul, float *ptr_a, float *ptr_b, float *ptr_result,
                 cl::Buffer buffer_a, cl::Buffer buffer_b, cl::Buffer buffer_result) {
  for(int i = 0; i < kDim; i++) {
    ptr_a[i] = in[i];
  }
  for(int i = 0; i < kDim; i++) {
    for(int j = 0; j < kDim; j++) {
      ptr_b[i * kDim + j] = w[i][j];
    }
  }
  q.enqueueMigrateMemObjects({buffer_a, buffer_b}, 0);
  kernel_matmul.setArg(3, kDim);
  kernel_matmul.setArg(4, kDim);
  q.enqueueTask(kernel_matmul);
  q.enqueueMigrateMemObjects({buffer_result}, CL_MIGRATE_MEM_OBJECT_HOST);
  q.finish();
  for(int i = 0; i < kDim; i++) {
    out[i] = ptr_result[i];
  }
}

ホストプログラムをビルドする

ここではホストプログラムのビルドを行います。

まずは swan/src を右クリックし、swan リポジトリのsrcディレクトリにあるホストプログラムを追加します。必要なファイルを追加した状態が以下になります。

ファイルの追加が完了したら、swan/swan.prjを選択した後に金槌のアイコンをクリックして、ホストプログラムのビルドを開始します。

ビルドが完了すると、swan/Hardware以下に実行ファイルが生成されます。

KV260に各種ファイルを転送する

最初に、以下でダウンロードしたSDカードイメージを書き込んだSDカードをKV260に挿入し、KV260でPetalinuxが動作している事を確認してください。

その後、ifconfig等でKV260のIPアドレスを確認した後、scpコマンドを用いて以下のファイルをKV260に転送します。

kv260_swan/swan/Hardware/swan
kv260_swan/swan_system_hw_link/Hardware/binary_container_1.xclbin をbinary_container_1.bin にリネームしたもの
pl.dtbo (チュートリアルにて生成)
shell.json (チュートリアルにて生成)

KV260にswanハードウェアを生成する

KV260に接続し、/lib/firmware/xilinx/ 以下に swanディレクトリを作成し、そこにビットストリーム、pl.dtbo、shell.jsonをコピーします。

$ sudo mkdir /lib/firmware/xilinx/swan
$ sudo cp binary_container_1.bin pl.dtbo shell.json /lib/firmware/xilinx/swan

その後、xmutilコマンドを用いてPLにハードウェアを生成します。

$ sudo xmutil listapps
$ sudo xmutil unloadapp
$ sudo xmutil loadapp swan

Loaded to Slot 0と出れば成功です。

swan on FPGA

これで、ホームディレクトリにある swanを実行すると、KV260上でLLMが動き出します。

$ ./swan

KV260でLLMが動作する様子

平均して3.6 token/sで動作しています。

まとめと今後

LLMはメモリ律速なアプリケーション
Kria KV260とHLSでLLMのデコーダ内部の関数をハードウェアオフロードして動作させた
大胆な計算機アーキテクチャの改善でLLM推論の高速化を目指す

VPK180が届きました。本番はこれからです！

採用情報

Turing 株式会社では大規模GPUクラスタの構築や、LLM 推論アクセラレータの開発など、完全自動運転の実現に必要な様々な技術開発を行っています。興味がある方は、Turing の公式 Web サイト、採用情報などをご覧ください。話を聞きたいという方はCTOの青木さんの X(旧Twitter) DMや採用ページの応募フォームからでもお気軽にご連絡ください。

Discussion