1
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

PYNQ-Z2でLLMを動かした

1
Posted at

前提

使用するボードはPYNQ-Z2

以下の実行結果がこれ

cat /etc/sysconfig/board.info 2>/dev/null || cat /etc/os-release
python3 -c "import pynq; print(pynq.__version__)"
python3 --version
NAME="PynqLinux"
VERSION="3.0 (Belfast)"
ID=pynqlinux
ID_LIKE=ubuntu
PRETTY_NAME="PynqLinux, based on Ubuntu 22.04"
VERSION_ID="3.0"
HOME_URL="https://www.pynq.io/"
SUPPORT_URL="https://discuss.pynq.io/"
BUG_REPORT_URL="https://www.pynq.io"
PRIVACY_POLICY_URL="https://www.pynq.io"
VERSION_CODENAME=Belfast
UBUNTU_CODENAME=jammy
3.0.1
Python 3.10.4

使用するVivadoとVitis HLSは2023.2.2

使用するLLMモデルはstories15M.bin
PYNQ-Z2はFPGAボードとしては小型らしく,このくらい小型のモデルじゃないと載らないらしい.

hwhファイルがPYNQ-Z2で使用できるかどうか確認する

PYNQ-Z2のOSが3系だとVivadoの2023に対応していません.
hwhファイルが使用できるかどうか,簡単なプロジェクトで確認しました.

# PYNQ-Z2向け .hwh パース互換性確認用の最小構成プロジェクト
# Zynq PS7 + AXI GPIO(LED 4bit)のみのブロックデザインを作成し,
# ビットストリームと.hwhを生成する

# 変数定義
set proj_name "smoketest"
set proj_dir  "./smoketest_proj"
set part_name "xc7z020clg400-1"
set board_vlnv "tul.com.tw:pynq-z2:part0:1.0"

# プロジェクト作成
create_project $proj_name $proj_dir -part $part_name -force
set_property board_part $board_vlnv [current_project]

# ブロックデザイン作成
create_bd_design "smoketest_bd"

# 一つ目のIP.PS
# PSに合わせた規定設定を自動で適用
create_bd_cell -type ip -vlnv xilinx.com:ip:processing_system7:5.5 processing_system7_0
apply_bd_automation -rule xilinx.com:bd_rule:processing_system7 \
  -config {make_external "FIXED_IO, DDR" apply_board_preset "1" Master "Disable" Slave "Disable"} \
  [get_bd_cells processing_system7_0]

# 二つ目のIP.AXI GPIO
# ボードのLEDに繋いでいるだけ
create_bd_cell -type ip -vlnv xilinx.com:ip:axi_gpio:2.0 axi_gpio_led
set_property -dict [list \
  CONFIG.C_GPIO_WIDTH {4} \
  CONFIG.C_ALL_OUTPUTS {1} \
] [get_bd_cells axi_gpio_led]

apply_bd_automation -rule xilinx.com:bd_rule:axi4 \
  -config { Master "/processing_system7_0/M_AXI_GP0" Clk "Auto" } \
  [get_bd_intf_pins axi_gpio_led/S_AXI]

apply_bd_automation -rule xilinx.com:bd_rule:board \
  -config { Board_Interface "leds_4bits" } \
  [get_bd_intf_pins axi_gpio_led/GPIO]

# ブロック配置を整形し,配線・アドレスマップに矛盾がないか検証
regenerate_bd_layout
validate_bd_design

# ブロックデザインのラップファイルを作成し,トップモジュールとして合成対象に登録
make_wrapper -files [get_files $proj_dir/$proj_name.srcs/sources_1/bd/smoketest_bd/smoketest_bd.bd] -top
add_files -norecurse $proj_dir/$proj_name.gen/sources_1/bd/smoketest_bd/hdl/smoketest_bd_wrapper.v
update_compile_order -fileset sources_1
set_property top smoketest_bd_wrapper [current_fileset]

# 論理合成
launch_runs synth_1 -jobs 4
wait_on_run synth_1

# bitファイル生成
launch_runs impl_1 -to_step write_bitstream -jobs 4
wait_on_run impl_1

# .bit と .hwh を出力ディレクトリにまとめて配置
file mkdir "./output"
file copy -force "$proj_dir/$proj_name.runs/impl_1/smoketest_bd_wrapper.bit" "./output/smoketest.bit"
file copy -force "$proj_dir/$proj_name.gen/sources_1/bd/smoketest_bd/hw_handoff/smoketest_bd.hwh" "./output/smoketest.hwh"

puts "SMOKETEST_BUILD_DONE"

vivado -mode batch -source build_smoketest.tcl -log build.log -journal build.jou

で実行したらbitファイルとhwhファイルをPYNQ-Z2に配置し,以下のソースを実行します

from pynq import Overlay

ol = Overlay("smoketest.bit")
print(ol.ip_dict.keys())
print(ol.axi_gpio_led.register_map)

結果がこんな感じになったらOK

dict_keys(['axi_gpio_led', 'processing_system7_0'])
RegisterMap {
  GPIO_DATA = Register(Channel_1_GPIO_DATA=0),
  GPIO_TRI = Register(Channel_1_GPIO_TRI=15),
  GPIO2_DATA = Register(Channel_2_GPIO_DATA=0),
  GPIO2_TRI = Register(Channel_2_GPIO_TRI=4294967295),
  GIER = Register(Global_Interrupt_Enable=1),
  IP_ISR = Register(Channel_1_Interrupt_Status=0, Channel_2_Interrupt_Status=0),
  IP_IER = Register(Channel_1_Interrupt_Enable=0, Channel_2_Interrupt_Enable=0)
}

必要なファイルをダウンロード

モデルファイル

curl -L -o stories15M.bin https://huggingface.co/karpathy/tinyllamas/resolve/main/stories15M.bin

このモデルのカタログスペック? は以下.

  • セルフアテンション+FFNの層が6階層ある.なお,今回はこれを一層だけPLに実装する
  • トークナイザの語彙数は32000.ただし英語オンリー
  • 入力ベクトルの次元数は288
  • FFNの隠れ層の次元数は768
  • セルフアテンションのhead数(並列実行数)は6
  • 一度に扱える最大トークン数(ざっくり言うと扱える文章の長さ)が256
  • 正規化にはRMSNormを使用.アテンションの直前とFFNの直前と出力層の直前にかかる
  • FFNはSwiGLUを使用

あと,推論に使用するllama2.cをクローンします.

git clone https://github.com/karpathy/llama2.c.git llama2c_upstream

ビルドして動作確認

llama2.cをビルドして動作確認します.

@rem 速度優先,Windowsようにコンパイル
gcc -Ofast -D_WIN32 -o run.exe -I. run.c win.c
@rem プロンプトを投げて実際に話しかけてみる
./run.exe ../weights/stories15M.bin -n 60 -i "hello!"

ちなみに返答がこれでした.

hello! Open, opened! I found a chicken in the garden. He was so happy! He was always happy. He ran around the garden looking for food. Then he found a peckie. He was so excited! He loved the peckie.
But then it was time
achieved tok/s: 198.653199

何と言っているのかよく分からないのでGoogle翻訳すると以下のようなことを言っているようです.

こんにちは!パカッ、開いた!庭でニワトリを見つけました。彼はとても嬉しそうでした!彼はいつもご機嫌だったのです。彼は食べ物を探して庭を走り回りました。すると、彼は「ペッキー」を見つけました。彼は大興奮!そのペッキーが大好きだったのです。
でも、やがてその時がやってきました。

終わり方怖い.これは60トークンで切ってるからですが.
ローカルLLMに比べるとやっぱり速いですが,出力される文章はゴミみたいなクオリティです.

ap_fixed.hをHLS合成なしでビルドできるか確認

これが出来たら高位合成しなくても固定小数点演算をシミュレーションできます.
シミュレーションというのはPC上で行う事なので,実機でさえ動けばいいやって場合はこのステップは省略可能.

// Vitis HLS付属のap_fixed.hがHLS合成なし(ホストg++単体)でビルド・実行できるかの確認用
#include <cstdio>
#include "ap_fixed.h"

int main() {
    // 計画で採用予定の型: 16bit全体,整数部6bit,小数部10bit
    ap_fixed<16, 6> a = 3.14159;
    ap_fixed<16, 6> b = -1.5;
    ap_fixed<16, 6> c = a * b; // 固定小数点の乗算

    printf("a = %f\n", a.to_float());
    printf("b = %f\n", b.to_float());
    printf("a * b = %f\n", c.to_float());
    return 0;
}
g++ -std=c++11 -I"(Vitis HLS内にあるincludeディレクトリのパスだよ)" -o ap_fixed_smoketest.exe ap_fixed_smoketest.cpp
./ap_fixed_smoketest.exe

やっていることはたんに計算しているだけなので,こういう実行結果になればOK

a = 3.140625
b = -1.500000
a * b = -4.710938

run.cを分割

先ほどLLMに話しかけるのに使用したファイルを機能ごとに分割して改造しやすくします.

加えて,確率的サンプリングではなく貪欲法を使用するよう変更しました.
これで,同じ入力ならfp32で計算しても上記固定小数点で計算しても同じ出力になりやすくなります(一位と二位がよほど僅差でなければ絶対に同じ).確率的サンプリングしてしまうと,似ている数字が出力されやすくなり乱数のシード値が同じでも計算結果がずれやすくなります.

※以下のソースは最終版の,固定小数点を使用しているものです.

model.h
// Transformerのforward計算(RMSNorm/Attention/SwiGLU FFN)本体
#ifndef MODEL_H
#define MODEL_H

#include "weights_loader.h"

// 1トークン分の計算に使うバッファ群(毎トークンで使い回す作業領域)
typedef struct {
    float* x;   // 現在の本流ベクトル (dim,)
    float* xb;  // 残差ブロック内の作業バッファ (dim,)
    float* xb2; // 同上,予備 (dim,)
    float* hb;  // FFN中間層の作業バッファ (hidden_dim,)
    float* hb2; // 同上,ゲート側 (hidden_dim,)
    float* q;   // Queryベクトル (dim,)
    float* att; // Attentionスコア/重み (n_heads, seq_len)
    float* logits; // 出力ロジット (vocab_size,)
    float* key_cache;   // KVキャッシュ(Key)   (layer, seq_len, kv_dim)
    float* value_cache; // KVキャッシュ(Value) (layer, seq_len, kv_dim)
} RunState;

// モデル全体(構造情報+重み+作業バッファ+ファイルハンドル)
typedef struct {
    Config config;
    TransformerWeights weights;
    RunState state;
    int fd;             // チェックポイントファイルのディスクリプタ
    float* data;         // メモリマップされた生データ
    ssize_t file_size;   // チェックポイントファイルのバイト数
} Transformer;

// checkpoint_pathからモデルを読み込み,計算用バッファも確保する
void build_transformer(Transformer* t, const char* checkpoint_path);

// build_transformer()で確保したリソースを解放する
void free_transformer(Transformer* t);

// 位置posにあるトークンtokenを1つ処理し,次トークンの予測ロジットを返す(vocab_size,)
// KVキャッシュへの書き込みを伴うため,pos=0から順番に呼び出す必要がある
float* forward(Transformer* transformer, int token, int pos);

#endif
mode.c
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include "model.h"

// 二乗平均平方根で正規化し,学習済みスケールを掛ける(Attention/FFN直前で使用)
static void rmsnorm(float* o, float* x, float* weight, int size) {
    float ss = 0.0f;
    for (int j = 0; j < size; j++) {
        ss += x[j] * x[j];
    }
    ss /= size;
    ss += 1e-5f; // ゼロ割り防止のイプシロン
    ss = 1.0f / sqrtf(ss);
    for (int j = 0; j < size; j++) {
        o[j] = weight[j] * (ss * x[j]);
    }
}

// 数値安定のため最大値を引いてからexpし,合計1になるよう正規化する
static void softmax(float* x, int size) {
    float max_val = x[0];
    for (int i = 1; i < size; i++) {
        if (x[i] > max_val) { max_val = x[i]; }
    }
    float sum = 0.0f;
    for (int i = 0; i < size; i++) {
        x[i] = expf(x[i] - max_val);
        sum += x[i];
    }
    for (int i = 0; i < size; i++) {
        x[i] /= sum;
    }
}

// 行列wと入力xの積(w: d行n列,x: n次元 -> 出力: d次元).Q/K/V/O/FFN射影で共通利用
static void matmul(float* xout, float* x, float* w, int n, int d) {
    for (int i = 0; i < d; i++) {
        float val = 0.0f;
        for (int j = 0; j < n; j++) {
            val += w[i * n + j] * x[j];
        }
        xout[i] = val;
    }
}

static void malloc_run_state(RunState* s, Config* p) {
    int kv_dim = (p->dim * p->n_kv_heads) / p->n_heads;
    s->x = calloc(p->dim, sizeof(float));
    s->xb = calloc(p->dim, sizeof(float));
    s->xb2 = calloc(p->dim, sizeof(float));
    s->hb = calloc(p->hidden_dim, sizeof(float));
    s->hb2 = calloc(p->hidden_dim, sizeof(float));
    s->q = calloc(p->dim, sizeof(float));
    s->key_cache = calloc((size_t)p->n_layers * p->seq_len * kv_dim, sizeof(float));
    s->value_cache = calloc((size_t)p->n_layers * p->seq_len * kv_dim, sizeof(float));
    s->att = calloc((size_t)p->n_heads * p->seq_len, sizeof(float));
    s->logits = calloc(p->vocab_size, sizeof(float));
    if (!s->x || !s->xb || !s->xb2 || !s->hb || !s->hb2 || !s->q
        || !s->key_cache || !s->value_cache || !s->att || !s->logits) {
        fprintf(stderr, "malloc failed!\n");
        exit(EXIT_FAILURE);
    }
}

static void free_run_state(RunState* s) {
    free(s->x);
    free(s->xb);
    free(s->xb2);
    free(s->hb);
    free(s->hb2);
    free(s->q);
    free(s->att);
    free(s->logits);
    free(s->key_cache);
    free(s->value_cache);
}

void build_transformer(Transformer* t, const char* checkpoint_path) {
    read_checkpoint(checkpoint_path, &t->config, &t->weights, &t->fd, &t->data, &t->file_size);
    malloc_run_state(&t->state, &t->config);
}

void free_transformer(Transformer* t) {
    free_checkpoint(t->fd, t->data, t->file_size);
    free_run_state(&t->state);
}

float* forward(Transformer* transformer, int token, int pos) {
    Config* p = &transformer->config;
    TransformerWeights* w = &transformer->weights;
    RunState* s = &transformer->state;
    float* x = s->x;
    int dim = p->dim;
    int kv_dim = (p->dim * p->n_kv_heads) / p->n_heads;
    int kv_mul = p->n_heads / p->n_kv_heads; // Query側ヘッド数がKV側の何倍か(multiquery対応)
    int hidden_dim = p->hidden_dim;
    int head_size = dim / p->n_heads;

    // トークン埋め込みをxにコピー
    float* content_row = w->token_embedding_table + (size_t)token * dim;
    memcpy(x, content_row, dim * sizeof(*x));

    // 全レイヤーを順番に処理
    for (int l = 0; l < p->n_layers; l++) {

        // Attention前のRMSNorm
        rmsnorm(s->xb, x, w->rms_att_weight + (size_t)l * dim, dim);

        // このレイヤー・この位置に対応するKVキャッシュ領域を指す
        size_t loff = (size_t)l * p->seq_len * kv_dim;
        float* k = s->key_cache + loff + (size_t)pos * kv_dim;
        float* v = s->value_cache + loff + (size_t)pos * kv_dim;

        // Q/K/Vそれぞれへの射影
        matmul(s->q, s->xb, w->wq + (size_t)l * dim * dim, dim, dim);
        matmul(k, s->xb, w->wk + (size_t)l * dim * kv_dim, dim, kv_dim);
        matmul(v, s->xb, w->wv + (size_t)l * dim * kv_dim, dim, kv_dim);

        // RoPE: 位置posに応じてQ/Kを2要素ずつペアにして回転させる(相対位置エンコーディング)
        for (int i = 0; i < dim; i += 2) {
            int head_dim = i % head_size;
            float freq = 1.0f / powf(10000.0f, head_dim / (float)head_size);
            float val = pos * freq;
            float fcr = cosf(val);
            float fci = sinf(val);
            int rotn = i < kv_dim ? 2 : 1; // 2=Q&K,1=Qのみ(kv_dim < dimのとき)
            for (int v_idx = 0; v_idx < rotn; v_idx++) {
                float* vec = v_idx == 0 ? s->q : k;
                float v0 = vec[i];
                float v1 = vec[i + 1];
                vec[i] = v0 * fcr - v1 * fci;
                vec[i + 1] = v0 * fci + v1 * fcr;
            }
        }

        // 各ヘッドごとにAttentionを計算
        for (int h = 0; h < p->n_heads; h++) {
            float* q = s->q + h * head_size;
            float* att = s->att + h * p->seq_len;

            // 現在位置までの各時刻とのスコア(内積)を計算
            for (int t = 0; t <= pos; t++) {
                float* kt = s->key_cache + loff + (size_t)t * kv_dim + (h / kv_mul) * head_size;
                float score = 0.0f;
                for (int i = 0; i < head_size; i++) {
                    score += q[i] * kt[i];
                }
                score /= sqrtf(head_size);
                att[t] = score;
            }

            // スコアをsoftmaxしてAttention重みに変換
            softmax(att, pos + 1);

            // Attention重みでValueを加重和し,xbに書き戻す
            float* xb = s->xb + h * head_size;
            memset(xb, 0, head_size * sizeof(float));
            for (int t = 0; t <= pos; t++) {
                float* vt = s->value_cache + loff + (size_t)t * kv_dim + (h / kv_mul) * head_size;
                float a = att[t];
                for (int i = 0; i < head_size; i++) {
                    xb[i] += a * vt[i];
                }
            }
        }

        // Attention出力を288次元に戻す射影
        matmul(s->xb2, s->xb, w->wo + (size_t)l * dim * dim, dim, dim);

        // 残差加算(Attentionサブレイヤー)
        for (int i = 0; i < dim; i++) {
            x[i] += s->xb2[i];
        }

        // FFN前のRMSNorm
        rmsnorm(s->xb, x, w->rms_ffn_weight + (size_t)l * dim, dim);

        // SwiGLU: ゲート側(w1)と素通し側(w3)をそれぞれ射影
        matmul(s->hb, s->xb, w->w1 + (size_t)l * dim * hidden_dim, dim, hidden_dim);
        matmul(s->hb2, s->xb, w->w3 + (size_t)l * dim * hidden_dim, dim, hidden_dim);

        // SiLU(w1側)を計算してw3側と要素ごとに掛け算(ゲート)
        for (int i = 0; i < hidden_dim; i++) {
            float val = s->hb[i];
            val *= (1.0f / (1.0f + expf(-val))); // SiLU(x) = x * sigmoid(x)
            val *= s->hb2[i];
            s->hb[i] = val;
        }

        // FFN出力を288次元に戻す射影
        matmul(s->xb, s->hb, w->w2 + (size_t)l * dim * hidden_dim, hidden_dim, dim);

        // 残差加算(FFNサブレイヤー)
        for (int i = 0; i < dim; i++) {
            x[i] += s->xb[i];
        }
    }

    // 最終RMSNorm
    rmsnorm(x, x, w->rms_final_weight, dim);

    // 分類器で32000語彙分のロジットに変換
    matmul(s->logits, x, w->wcls, p->dim, p->vocab_size);
    return s->logits;
}
weights_loader.h
// stories15M.bin(llama2.c形式のチェックポイント)の読み込み
#ifndef WEIGHTS_LOADER_H
#define WEIGHTS_LOADER_H

#include <sys/types.h>

// モデルの構造情報(チェックポイント先頭のヘッダから読み込む)
typedef struct {
    int dim;         // 本流ベクトルの次元数(residual streamの太さ)
    int hidden_dim;   // FFN内部の中間層の次元数
    int n_layers;     // Transformerブロックの数
    int n_heads;       // Attentionのヘッド数
    int n_kv_heads;    // Key/Valueのヘッド数(Query側と同数なら通常のMulti-Head Attention)
    int vocab_size;    // 語彙数(分類先の数)
    int seq_len;       // 最大コンテキスト長
} Config;

// 各層の重み(すべてConfigのdim/hidden_dim/n_layers等から決まるサイズの配列)
typedef struct {
    float* token_embedding_table; // 埋め込み行列 (vocab_size, dim)
    float* rms_att_weight;        // Attention前RMSNormのスケール (layer, dim)
    float* rms_ffn_weight;        // FFN前RMSNormのスケール (layer, dim)
    float* wq; // Query射影行列 (layer, dim, n_heads*head_size)
    float* wk; // Key射影行列 (layer, dim, n_kv_heads*head_size)
    float* wv; // Value射影行列 (layer, dim, n_kv_heads*head_size)
    float* wo; // Attention出力射影行列 (layer, n_heads*head_size, dim)
    float* w1; // SwiGLU FFNのゲート側射影 (layer, hidden_dim, dim)
    float* w2; // SwiGLU FFNの出力射影 (layer, dim, hidden_dim)
    float* w3; // SwiGLU FFNの素通し側射影 (layer, hidden_dim, dim)
    float* rms_final_weight; // 最終RMSNormのスケール (dim,)
    float* wcls;             // 分類器(vocab_size, dim).埋め込み行列と共有される場合あり
} TransformerWeights;

// チェックポイントファイルをメモリマップし,configとweightsに値を設定する
// fd/data/file_sizeはfree_checkpoint()での後片付けに必要なので呼び出し側で保持する
void read_checkpoint(const char* checkpoint_path, Config* config, TransformerWeights* weights,
                      int* fd, float** data, ssize_t* file_size);

// read_checkpoint()で確保したメモリマッピングとファイルディスクリプタを解放する
void free_checkpoint(int fd, float* data, ssize_t file_size);

#endif
weigths_loader.c
#include <stdio.h>
#include <stdlib.h>
#include <fcntl.h>
#if defined _WIN32
    #include "win.h"
#else
    #include <unistd.h>
    #include <sys/mman.h>
#endif
#include "weights_loader.h"

// ptrを先頭からConfigのサイズに応じて割り振り,各重み配列の開始位置を決める
static void memory_map_weights(TransformerWeights* w, Config* p, float* ptr, int shared_weights) {
    int head_size = p->dim / p->n_heads;
    // 13B以上の大きいモデルでも桁あふれしないよう64bitで計算
    unsigned long long n_layers = p->n_layers;

    w->token_embedding_table = ptr;
    ptr += p->vocab_size * p->dim;
    w->rms_att_weight = ptr;
    ptr += n_layers * p->dim;
    w->wq = ptr;
    ptr += n_layers * p->dim * (p->n_heads * head_size);
    w->wk = ptr;
    ptr += n_layers * p->dim * (p->n_kv_heads * head_size);
    w->wv = ptr;
    ptr += n_layers * p->dim * (p->n_kv_heads * head_size);
    w->wo = ptr;
    ptr += n_layers * (p->n_heads * head_size) * p->dim;
    w->rms_ffn_weight = ptr;
    ptr += n_layers * p->dim;
    w->w1 = ptr;
    ptr += n_layers * p->dim * p->hidden_dim;
    w->w2 = ptr;
    ptr += n_layers * p->hidden_dim * p->dim;
    w->w3 = ptr;
    ptr += n_layers * p->dim * p->hidden_dim;
    w->rms_final_weight = ptr;
    ptr += p->dim;
    ptr += p->seq_len * head_size / 2; // 旧バージョンのRoPE用テーブル(freq_cis_real)分を読み飛ばす
    ptr += p->seq_len * head_size / 2; // 同上(freq_cis_imag)
    w->wcls = shared_weights ? w->token_embedding_table : ptr;
}

void read_checkpoint(const char* checkpoint_path, Config* config, TransformerWeights* weights,
                      int* fd, float** data, ssize_t* file_size) {
    FILE* file = fopen(checkpoint_path, "rb");
    if (!file) { fprintf(stderr, "Couldn't open file %s\n", checkpoint_path); exit(EXIT_FAILURE); }

    // ヘッダ(Config構造体分)を読み込む
    if (fread(config, sizeof(Config), 1, file) != 1) { exit(EXIT_FAILURE); }
    // vocab_sizeが負の場合は分類器の重みが埋め込み行列と共有されていない合図
    int shared_weights = config->vocab_size > 0 ? 1 : 0;
    config->vocab_size = abs(config->vocab_size);

    // ファイルサイズを取得
    fseek(file, 0, SEEK_END);
    *file_size = ftell(file);
    fclose(file);

    // 残りの重みデータをメモリマップで読み込む(コピーせずファイルを直接メモリとして扱う)
    *fd = open(checkpoint_path, O_RDONLY);
    if (*fd == -1) { fprintf(stderr, "open failed!\n"); exit(EXIT_FAILURE); }
    *data = mmap(NULL, *file_size, PROT_READ, MAP_PRIVATE, *fd, 0);
    if (*data == MAP_FAILED) { fprintf(stderr, "mmap failed!\n"); exit(EXIT_FAILURE); }

    float* weights_ptr = *data + sizeof(Config) / sizeof(float);
    memory_map_weights(weights, config, weights_ptr, shared_weights);
}

void free_checkpoint(int fd, float* data, ssize_t file_size) {
    if (data != MAP_FAILED) { munmap(data, file_size); }
    if (fd != -1) { close(fd); }
}
tokenizer.h
// 文字列とトークンID列を相互変換するBPE(Byte Pair Encoding)トークナイザ
#ifndef TOKENIZER_H
#define TOKENIZER_H

// 語彙を文字列でソートした状態で保持し,2分探索できるようにするための対応表
typedef struct {
    char* str;
    int id;
} TokenIndex;

typedef struct {
    char** vocab;             // トークンID -> 文字列
    float* vocab_scores;       // マージ優先度(値が大きいほど優先してマージされる)
    TokenIndex* sorted_vocab;  // str_lookup用にソート済みの語彙(遅延初期化)
    int vocab_size;
    unsigned int max_token_length;
    unsigned char byte_pieces[512]; // 1バイトずつの生バイト文字列(256種×2バイト)
} Tokenizer;

// tokenizer.binを読み込み,vocab_size個の語彙を構築する
void build_tokenizer(Tokenizer* t, const char* tokenizer_path, int vocab_size);

// build_tokenizer()で確保したメモリを解放する
void free_tokenizer(Tokenizer* t);

// 文字列textをトークンID列tokensに変換する(呼び出し側でtokensの領域を確保しておくこと)
// bos!=0なら先頭にBOS(=1),eos!=0なら末尾にEOS(=2)を付与する
void encode(Tokenizer* t, const char* text, int bos, int eos, int* tokens, int* n_tokens);

// トークンID(token)を文字列に変換する.prev_tokenは直前のトークン(先頭空白除去の判定に使用)
char* decode(Tokenizer* t, int prev_token, int token);

// 表示不能な制御文字を除いてpieceを標準出力に印字する
void safe_printf(char* piece);

#endif
tokenizer.c
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <ctype.h>
#include "tokenizer.h"

// TokenIndexを文字列の辞書順で比較(qsort/bsearch用)
static int compare_tokens(const void* a, const void* b) {
    return strcmp(((TokenIndex*)a)->str, ((TokenIndex*)b)->str);
}

// ソート済み語彙からstrを2分探索し,見つかればトークンID,なければ-1を返す
static int str_lookup(char* str, TokenIndex* sorted_vocab, int vocab_size) {
    TokenIndex tok = { .str = str };
    TokenIndex* res = bsearch(&tok, sorted_vocab, vocab_size, sizeof(TokenIndex), compare_tokens);
    return res != NULL ? res->id : -1;
}

void build_tokenizer(Tokenizer* t, const char* tokenizer_path, int vocab_size) {
    t->vocab_size = vocab_size;
    t->vocab = (char**)malloc(vocab_size * sizeof(char*));
    t->vocab_scores = (float*)malloc(vocab_size * sizeof(float));
    t->sorted_vocab = NULL; // 実際にencode()が呼ばれるまで構築を遅延

    // 1バイト単位の生バイト文字列テーブルを用意(0x00-0xFFそれぞれ1文字+終端)
    for (int i = 0; i < 256; i++) {
        t->byte_pieces[i * 2] = (unsigned char)i;
        t->byte_pieces[i * 2 + 1] = '\0';
    }

    FILE* file = fopen(tokenizer_path, "rb");
    if (!file) { fprintf(stderr, "couldn't load %s\n", tokenizer_path); exit(EXIT_FAILURE); }
    if (fread(&t->max_token_length, sizeof(int), 1, file) != 1) { fprintf(stderr, "failed read\n"); exit(EXIT_FAILURE); }
    int len;
    for (int i = 0; i < vocab_size; i++) {
        if (fread(t->vocab_scores + i, sizeof(float), 1, file) != 1) { fprintf(stderr, "failed read\n"); exit(EXIT_FAILURE); }
        if (fread(&len, sizeof(int), 1, file) != 1) { fprintf(stderr, "failed read\n"); exit(EXIT_FAILURE); }
        t->vocab[i] = (char*)malloc(len + 1);
        if (fread(t->vocab[i], len, 1, file) != 1) { fprintf(stderr, "failed read\n"); exit(EXIT_FAILURE); }
        t->vocab[i][len] = '\0';
    }
    fclose(file);
}

void free_tokenizer(Tokenizer* t) {
    for (int i = 0; i < t->vocab_size; i++) { free(t->vocab[i]); }
    free(t->vocab);
    free(t->vocab_scores);
    free(t->sorted_vocab);
}

char* decode(Tokenizer* t, int prev_token, int token) {
    char* piece = t->vocab[token];
    // BOS(=1)直後のトークンは先頭の空白を取り除く(SentencePieceの仕様)
    if (prev_token == 1 && piece[0] == ' ') { piece++; }
    // "<0x01>"のような生バイト表現のトークンは対応するバイト文字列に変換する
    unsigned char byte_val;
    if (sscanf(piece, "<0x%02hhX>", &byte_val) == 1) {
        piece = (char*)t->byte_pieces + byte_val * 2;
    }
    return piece;
}

void safe_printf(char* piece) {
    if (piece == NULL) { return; }
    if (piece[0] == '\0') { return; }
    if (piece[1] == '\0') {
        unsigned char byte_val = piece[0];
        // 制御コードなど表示不能なバイトは出力しない
        if (!(isprint(byte_val) || isspace(byte_val))) {
            return;
        }
    }
    printf("%s", piece);
}

void encode(Tokenizer* t, const char* text, int bos, int eos, int* tokens, int* n_tokens) {
    if (text == NULL) { fprintf(stderr, "cannot encode NULL text\n"); exit(EXIT_FAILURE); }

    if (t->sorted_vocab == NULL) {
        // str_lookup用に語彙を文字列順にソートしておく(初回のみ)
        t->sorted_vocab = malloc(t->vocab_size * sizeof(TokenIndex));
        for (int i = 0; i < t->vocab_size; i++) {
            t->sorted_vocab[i].str = t->vocab[i];
            t->sorted_vocab[i].id = i;
        }
        qsort(t->sorted_vocab, t->vocab_size, sizeof(TokenIndex), compare_tokens);
    }

    // 2トークン分の結合候補を一時的に保持するバッファ
    char* str_buffer = malloc((t->max_token_length * 2 + 1 + 2) * sizeof(char));
    size_t str_len = 0;

    *n_tokens = 0;

    // BOS(=1)トークンを先頭に付与
    if (bos) tokens[(*n_tokens)++] = 1;

    // 空文字列でなければ,SentencePiece仕様に合わせてダミーの先頭空白トークンを追加
    if (text[0] != '\0') {
        int dummy_prefix = str_lookup(" ", t->sorted_vocab, t->vocab_size);
        tokens[(*n_tokens)++] = dummy_prefix;
    }

    // UTF-8のバイト列を1コードポイントずつ処理する
    for (const char* c = text; *c != '\0'; c++) {
        // 0x80(10xxxxxx)は継続バイト.継続バイトでなければ新しいコードポイントの開始
        if ((*c & 0xC0) != 0x80) {
            str_len = 0;
        }

        str_buffer[str_len++] = *c;
        str_buffer[str_len] = '\0';

        // 次のバイトも継続バイトならまだコードポイントの途中なので読み進める
        if ((*(c + 1) & 0xC0) == 0x80 && str_len < 4) {
            continue;
        }

        // 1コードポイント分読み終えたので語彙にあるか検索する
        int id = str_lookup(str_buffer, t->sorted_vocab, t->vocab_size);
        if (id != -1) {
            tokens[(*n_tokens)++] = id;
        } else {
            // 語彙になければ1バイトずつ個別トークンにフォールバックする(先頭3つは特殊トークン)
            for (size_t i = 0; i < str_len; i++) {
                tokens[(*n_tokens)++] = (unsigned char)str_buffer[i] + 3;
            }
        }
        str_len = 0;
    }

    // スコアが最も高い隣接トークンのペアを繰り返しマージしていく(BPEの本体)
    while (1) {
        float best_score = -1e10f;
        int best_id = -1;
        int best_idx = -1;

        for (int i = 0; i < (*n_tokens - 1); i++) {
            sprintf(str_buffer, "%s%s", t->vocab[tokens[i]], t->vocab[tokens[i + 1]]);
            int id = str_lookup(str_buffer, t->sorted_vocab, t->vocab_size);
            if (id != -1 && t->vocab_scores[id] > best_score) {
                best_score = t->vocab_scores[id];
                best_id = id;
                best_idx = i;
            }
        }

        if (best_idx == -1) { break; } // マージできるペアがなければ終了

        // best_idxとbest_idx+1を1つのトークンbest_idにまとめ,後続を1つずつ詰める
        tokens[best_idx] = best_id;
        for (int i = best_idx + 1; i < (*n_tokens - 1); i++) {
            tokens[i] = tokens[i + 1];
        }
        (*n_tokens)--;
    }

    // EOS(=2)トークンを末尾に付与
    if (eos) tokens[(*n_tokens)++] = 2;

    free(str_buffer);
}
main.c
// fp32リファレンス実装のCLI.常に貪欲法(greedy)でトークンを選ぶため,
// 同じチェックポイント・同じプロンプトなら常に同じトークン列が得られる(再現性のある回帰テスト用)
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "model.h"
#include "tokenizer.h"

// ロジットの中から最も値が大きい語彙(=貪欲法での次トークン)のインデックスを返す
static int sample_argmax(float* logits, int n) {
    int max_i = 0;
    float max_p = logits[0];
    for (int i = 1; i < n; i++) {
        if (logits[i] > max_p) {
            max_i = i;
            max_p = logits[i];
        }
    }
    return max_i;
}

static void error_usage(void) {
    fprintf(stderr, "Usage:   reference_c <checkpoint> [options]\n");
    fprintf(stderr, "Example: reference_c ../weights/stories15M.bin -n 60 -i \"Once upon a time\" -o out.txt\n");
    fprintf(stderr, "Options:\n");
    fprintf(stderr, "  -n <int>    生成する最大トークン数,デフォルト256\n");
    fprintf(stderr, "  -i <string> 入力プロンプト\n");
    fprintf(stderr, "  -z <string> トークナイザファイルのパス,デフォルトtokenizer.bin\n");
    fprintf(stderr, "  -o <string> 生成したトークンID列を書き出すファイル(回帰テスト用,省略可)\n");
    exit(EXIT_FAILURE);
}

int main(int argc, char* argv[]) {
    char* checkpoint_path = NULL;
    char* tokenizer_path = "tokenizer.bin";
    int steps = 256;
    char* prompt = "";
    char* out_path = NULL;

    if (argc >= 2) { checkpoint_path = argv[1]; } else { error_usage(); }
    for (int i = 2; i < argc; i += 2) {
        if (i + 1 >= argc) { error_usage(); }
        if (argv[i][0] != '-' || strlen(argv[i]) != 2) { error_usage(); }
        if (argv[i][1] == 'n') { steps = atoi(argv[i + 1]); }
        else if (argv[i][1] == 'i') { prompt = argv[i + 1]; }
        else if (argv[i][1] == 'z') { tokenizer_path = argv[i + 1]; }
        else if (argv[i][1] == 'o') { out_path = argv[i + 1]; }
        else { error_usage(); }
    }

    // モデルとトークナイザを読み込む
    Transformer transformer;
    build_transformer(&transformer, checkpoint_path);
    if (steps <= 0 || steps > transformer.config.seq_len) { steps = transformer.config.seq_len; }

    Tokenizer tokenizer;
    build_tokenizer(&tokenizer, tokenizer_path, transformer.config.vocab_size);

    // プロンプトをトークンID列にエンコード(先頭にBOSを付与)
    int num_prompt_tokens = 0;
    int* prompt_tokens = (int*)malloc((strlen(prompt) + 3) * sizeof(int));
    encode(&tokenizer, prompt, 1, 0, prompt_tokens, &num_prompt_tokens);
    if (num_prompt_tokens < 1) {
        fprintf(stderr, "prompt tokenization failed\n");
        exit(EXIT_FAILURE);
    }

    // 生成したトークンID列を控えておく(回帰テスト用出力のため)
    int* generated_tokens = (int*)malloc(steps * sizeof(int));
    int n_generated = 0;

    FILE* out_file = NULL;
    if (out_path != NULL) {
        out_file = fopen(out_path, "w");
        if (!out_file) { fprintf(stderr, "couldn't open output file %s\n", out_path); exit(EXIT_FAILURE); }
    }

    int token = prompt_tokens[0]; // プロンプトの先頭トークンから開始
    int pos = 0;
    while (pos < steps) {
        float* logits = forward(&transformer, token, pos);

        int next;
        if (pos < num_prompt_tokens - 1) {
            // プロンプトの残りがあればそのまま強制的に次トークンとする
            next = prompt_tokens[pos + 1];
        } else {
            // プロンプトを消化しきったら貪欲法で次トークンを決定
            next = sample_argmax(logits, transformer.config.vocab_size);
        }
        pos++;

        if (next == 1) { break; } // BOS(=1)が出たらシーケンス終端とみなす

        char* piece = decode(&tokenizer, token, next);
        safe_printf(piece);
        fflush(stdout);

        generated_tokens[n_generated++] = next;
        token = next;
    }
    printf("\n");

    // 回帰テスト用にトークンID列をファイルへ書き出す(1行1トークンID)
    if (out_file != NULL) {
        for (int i = 0; i < n_generated; i++) {
            fprintf(out_file, "%d\n", generated_tokens[i]);
        }
        fclose(out_file);
    }

    free(prompt_tokens);
    free(generated_tokens);
    free_tokenizer(&tokenizer);
    free_transformer(&transformer);
    return 0;
}

この時の回答を「正解」としてファイル保存しておきます.

gcc -Ofast -D_WIN32 -I. -o reference_c.exe main.c model.c tokenizer.c weights_loader.c win.c -lm
./reference_c.exe ../weights/stories15M.bin -n 60 -i "Once upon a time" -z ../llama2c_upstream/tokenizer.bin -o test_vectors/stories15M_greedy_once_upon_a_time.txt

固定小数の場合の精度調査

先ほど実行したのはfp32なので,固定小数でも一度計算してみて,どのくらいの精度が出るのかを確認します.

gcc -Ofast -D_WIN32 -I. -c tokenizer.c -o tokenizer.o
gcc -Ofast -D_WIN32 -I. -c weights_loader.c -o weights_loader.o
gcc -Ofast -D_WIN32 -I. -c win.c -o win.o
g++ -O3 -std=c++11 -D_WIN32 -I. -I"(Vitis HLSにあるincludeディレクトリのパスだよ)" -c model_fixed.cpp -o model_fixed.o
g++ -O3 -std=c++11 -D_WIN32 -I. -I"(Vitis HLSにあるincludeディレクトリのパスだよ)" -c main_fixed.cpp -o main_fixed.o
g++ -O3 -o reference_c_fixed.exe main_fixed.o model_fixed.o tokenizer.o weights_loader.o win.o -lm
./reference_c_fixed.exe ../weights/stories15M.bin -n 60 -i "Once upon a time" -z ../llama2c_upstream/tokenizer.bin -o test_vectors/stories15M_fixed16_6_once_upon_a_time.txt
diff test_vectors/stories15M_greedy_once_upon_a_time.txt test_vectors/stories15M_fixed16_6_once_upon_a_time.txt

浮動小数点でやった時と,出力の最初の方は同じ結果になりました.途中からずれますが.
これは小数の精度が違うことが原因です.FPGA上で実行する以上これ以上は精度を上げられないので,ここで終了とします.
固定小数点のビット数を減らすと,途中からずれ始めるトークン数が早めになります.

ちなみに,最初は飽和モードオフで実行していて全然違う結果になっていました.
飽和モードにして,オーバーフローしそうな場合は最小または最大に張り付かせることで解決しました.

Vitis HLSでの実装に入る

高位合成に使用するファイルを作成します.

llama_layer.h
// PLに実装するTransformer 1レイヤー分(Attention+FFN)のHLSカーネル
#pragma once

#include "ap_fixed.h"

// 重み・活性化の基本型.reference_c/fixed_types.hのM2検証結果と同じ型を採用
// (16bitはトークン7個目で発散,20bitで34個目まで後退,24bitに広げても変化なしだったため20bitを採用)
typedef ap_fixed<20, 6, AP_TRN, AP_SAT> fx16_t;

// 内積の累積用の広い型
typedef ap_fixed<32, 10, AP_TRN, AP_SAT> acc_t;

// stories15Mのモデル構造(コンパイル時定数.配列サイズをHLSで静的に決めるために固定する)
static const int DIM = 288;         // 本流ベクトルの次元数
static const int HIDDEN_DIM = 768;   // FFN中間層の次元数
static const int N_HEADS = 6;        // Attentionヘッド数
static const int N_KV_HEADS = 6;     // KVヘッド数(stories15MはN_HEADSと同数)
static const int HEAD_SIZE = DIM / N_HEADS;              // 1ヘッドあたりの次元数(48)
static const int KV_DIM = DIM * N_KV_HEADS / N_HEADS;     // KVキャッシュの1トークンあたりの次元数
static const int SEQ_LEN = 256;      // 最大コンテキスト長

// 1レイヤー分の重みブロックの先頭からのオフセット(fx16_t単位).
// PS側でこの並び順のままDDR上に重みを配置する(rms_att, wq, wk, wv, wo, rms_ffn, w1, w2, w3)
static const int OFFSET_RMS_ATT = 0;
static const int OFFSET_WQ = OFFSET_RMS_ATT + DIM;
static const int OFFSET_WK = OFFSET_WQ + DIM * DIM;
static const int OFFSET_WV = OFFSET_WK + DIM * KV_DIM;
static const int OFFSET_WO = OFFSET_WV + DIM * KV_DIM;
static const int OFFSET_RMS_FFN = OFFSET_WO + DIM * DIM;
static const int OFFSET_W1 = OFFSET_RMS_FFN + DIM;
static const int OFFSET_W2 = OFFSET_W1 + DIM * HIDDEN_DIM;
static const int OFFSET_W3 = OFFSET_W2 + HIDDEN_DIM * DIM;
static const int LAYER_WEIGHTS_SIZE = OFFSET_W3 + DIM * HIDDEN_DIM; // 1レイヤー分の重み総数

// PLに実装するTransformer 1レイヤー分のトップ関数
// x:        入出力の本流ベクトル(dim,).呼び出し前は入力,呼び出し後は残差加算済みの出力に上書きされる
// weights:  このレイヤー分の重みブロック(LAYER_WEIGHTS_SIZE,).上記オフセット順に並んでいること
// kv_cache: このレイヤーのKVキャッシュ(2 * SEQ_LEN * KV_DIM,).前半がKey,後半がValue
// pos:      現在のトークン位置(0起点).RoPEの回転角とKVキャッシュの書き込み位置に使う
void llama_layer(fx16_t x[DIM], fx16_t weights[LAYER_WEIGHTS_SIZE], fx16_t kv_cache[2 * SEQ_LEN * KV_DIM], int pos);

ちなみに,ループを並列実行するようにするとLUTを大量に消費したのでパイプ欄処理を全部オフにしています.

llama_layer.cpp
#include <cmath>
#include "llama_layer.h"

// 二乗平均平方根で正規化し,学習済みスケールを掛ける
// 内積(二乗和)はacc_tで累積し,最後にfx16_tへ丸める
static void rmsnorm(fx16_t out[DIM], fx16_t x[DIM], fx16_t weight[DIM]) {
    acc_t ss = 0;
    for (int j = 0; j < DIM; j++) {
#pragma HLS PIPELINE off
        ss += (acc_t)x[j] * (acc_t)x[j];
    }
    // sqrt/割り算はdoubleで計算する(HLS向け近似(CORDIC/LUT)への置き換えはM5の合成結果次第で検討)
    double ss_d = (double)ss / DIM + 1e-5;
    float inv_rms = (float)(1.0 / sqrt(ss_d));
    for (int j = 0; j < DIM; j++) {
#pragma HLS PIPELINE off
        out[j] = (fx16_t)((float)weight[j] * (inv_rms * (float)x[j]));
    }
}

// softmax.expの計算はfloatに変換して行う(HLS向け近似はM5の合成結果次第で検討)
static void softmax(fx16_t x[SEQ_LEN], int size) {
    float max_val = (float)x[0];
    for (int i = 1; i < size; i++) {
#pragma HLS PIPELINE off
        if ((float)x[i] > max_val) { max_val = (float)x[i]; }
    }
    float tmp[SEQ_LEN];
    float sum = 0.0f;
    for (int i = 0; i < size; i++) {
#pragma HLS PIPELINE off
        tmp[i] = expf((float)x[i] - max_val);
        sum += tmp[i];
    }
    for (int i = 0; i < size; i++) {
#pragma HLS PIPELINE off
        x[i] = (fx16_t)(tmp[i] / sum);
    }
}

// 行列積.w(d行n列)とx(n次元)の積をxout(d次元)へ書き込む.acc_tで積算する
static void matmul(fx16_t* xout, fx16_t* x, fx16_t* w, int n, int d) {
    for (int i = 0; i < d; i++) {
#pragma HLS PIPELINE off
        acc_t val = 0;
        for (int j = 0; j < n; j++) {
#pragma HLS PIPELINE off
            val += (acc_t)w[i * n + j] * (acc_t)x[j];
        }
        xout[i] = (fx16_t)val;
    }
}

void llama_layer(fx16_t x[DIM], fx16_t weights[LAYER_WEIGHTS_SIZE], fx16_t kv_cache[2 * SEQ_LEN * KV_DIM], int pos) {
// m_axi: 配列の実データ本体をDDRとバースト転送でやり取りする.
// bundleごとに別の物理AXI4 Masterポートが1本ずつ生成される(activations/weights/kvcacheの3本).
// depthは「このポインタが実際にアクセスする要素数」をHLSに申告する値.
// ポインタ自体には配列長の情報が無いため,csim/cosimでのメモリ確保やアドレス範囲の把握に必要
#pragma HLS INTERFACE m_axi port=x offset=slave bundle=activations depth=DIM
#pragma HLS INTERFACE m_axi port=weights offset=slave bundle=weights depth=LAYER_WEIGHTS_SIZE
#pragma HLS INTERFACE m_axi port=kv_cache offset=slave bundle=kvcache depth=2*SEQ_LEN*KV_DIM

// s_axilite: PS(ARM)からIPへ少量の値をレジスタ経由で渡す.5行とも1本のAXI4-Liteポート(bundle=control)
// にまとまるが,中身は下記の通りそれぞれ別レジスタとして並ぶ
//   x/weights/kv_cache用レジスタ: 配列の中身ではなく,DDR上の先頭アドレス(32bit)を格納する
//   pos用レジスタ: ただの整数値そのものを格納する
//   return用レジスタ: 計算結果ではなく,start/done/idleといった実行制御ビットを格納する(戻り値がvoidでも必要)
#pragma HLS INTERFACE s_axilite port=x bundle=control
#pragma HLS INTERFACE s_axilite port=weights bundle=control
#pragma HLS INTERFACE s_axilite port=kv_cache bundle=control
#pragma HLS INTERFACE s_axilite port=pos bundle=control
#pragma HLS INTERFACE s_axilite port=return bundle=control

    // 計算用の作業バッファ(BRAMに割り付けられる想定)
    fx16_t xb[DIM];
    fx16_t xb2[DIM];
    fx16_t hb[HIDDEN_DIM];
    fx16_t hb2[HIDDEN_DIM];
    fx16_t q[DIM];
    fx16_t att[SEQ_LEN];

    // kv_cacheの前半をKey,後半をValueとして扱う(このレイヤー専用の領域)
    fx16_t* key_cache = kv_cache;
    fx16_t* value_cache = kv_cache + SEQ_LEN * KV_DIM;

    // Attention前のRMSNorm
    rmsnorm(xb, x, weights + OFFSET_RMS_ATT);

    // このトークン位置に対応するKVキャッシュ領域を指す
    fx16_t* k = key_cache + pos * KV_DIM;
    fx16_t* v = value_cache + pos * KV_DIM;

    // Q/K/Vへの射影
    matmul(q, xb, weights + OFFSET_WQ, DIM, DIM);
    matmul(k, xb, weights + OFFSET_WK, DIM, KV_DIM);
    matmul(v, xb, weights + OFFSET_WV, DIM, KV_DIM);

    // RoPE: 位置posに応じてQ/Kを2要素ずつペアにして回転させる(相対位置エンコーディング)
    for (int i = 0; i < DIM; i += 2) {
#pragma HLS PIPELINE off
        int head_dim = i % HEAD_SIZE;
        float freq = 1.0f / powf(10000.0f, head_dim / (float)HEAD_SIZE);
        float val = pos * freq;
        fx16_t fcr = (fx16_t)cosf(val);
        fx16_t fci = (fx16_t)sinf(val);
        int rotn = i < KV_DIM ? 2 : 1; // 2=Q&K,1=Qのみ
        for (int v_idx = 0; v_idx < rotn; v_idx++) {
#pragma HLS PIPELINE off
            fx16_t* vec = v_idx == 0 ? q : k;
            fx16_t v0 = vec[i];
            fx16_t v1 = vec[i + 1];
            vec[i] = (fx16_t)((acc_t)v0 * (acc_t)fcr - (acc_t)v1 * (acc_t)fci);
            vec[i + 1] = (fx16_t)((acc_t)v0 * (acc_t)fci + (acc_t)v1 * (acc_t)fcr);
        }
    }

    // 各ヘッドごとにAttentionを計算(N_HEADS==N_KV_HEADSなのでkv_mulは常に1)
    for (int h = 0; h < N_HEADS; h++) {
#pragma HLS PIPELINE off
        fx16_t* qh = q + h * HEAD_SIZE;
        fx16_t* atth = att; // headごとに使い回すスコアバッファ(SEQ_LEN分)

        // 現在位置までの各時刻とのスコア(内積)を計算
        for (int t = 0; t <= pos; t++) {
#pragma HLS PIPELINE off
            fx16_t* kt = key_cache + t * KV_DIM + h * HEAD_SIZE;
            acc_t score = 0;
            for (int i = 0; i < HEAD_SIZE; i++) {
#pragma HLS PIPELINE off
                score += (acc_t)qh[i] * (acc_t)kt[i];
            }
            atth[t] = (fx16_t)((float)score / sqrtf((float)HEAD_SIZE));
        }

        // スコアをsoftmaxしてAttention重みに変換
        softmax(atth, pos + 1);

        // Attention重みでValueを加重和し,xbに書き戻す
        fx16_t* xbh = xb + h * HEAD_SIZE;
        for (int i = 0; i < HEAD_SIZE; i++) {
#pragma HLS PIPELINE off
            xbh[i] = 0;
        }
        for (int t = 0; t <= pos; t++) {
#pragma HLS PIPELINE off
            fx16_t* vt = value_cache + t * KV_DIM + h * HEAD_SIZE;
            fx16_t a = atth[t];
            for (int i = 0; i < HEAD_SIZE; i++) {
#pragma HLS PIPELINE off
                xbh[i] = (fx16_t)((acc_t)xbh[i] + (acc_t)a * (acc_t)vt[i]);
            }
        }
    }

    // Attention出力をDIM次元に戻す射影
    matmul(xb2, xb, weights + OFFSET_WO, DIM, DIM);

    // 残差加算(Attentionサブレイヤー)
    for (int i = 0; i < DIM; i++) {
#pragma HLS PIPELINE off
        x[i] = (fx16_t)((acc_t)x[i] + (acc_t)xb2[i]);
    }

    // FFN前のRMSNorm
    rmsnorm(xb, x, weights + OFFSET_RMS_FFN);

    // SwiGLU: ゲート側(w1)と素通し側(w3)をそれぞれ射影
    matmul(hb, xb, weights + OFFSET_W1, DIM, HIDDEN_DIM);
    matmul(hb2, xb, weights + OFFSET_W3, DIM, HIDDEN_DIM);

    // SiLU(w1側)を計算してw3側と要素ごとに掛け算(ゲート)
    for (int i = 0; i < HIDDEN_DIM; i++) {
#pragma HLS PIPELINE off
        float val = (float)hb[i];
        val *= (1.0f / (1.0f + expf(-val)));
        hb[i] = (fx16_t)(val * (float)hb2[i]);
    }

    // FFN出力をDIM次元に戻す射影
    matmul(xb, hb, weights + OFFSET_W2, HIDDEN_DIM, DIM);

    // 残差加算(FFNサブレイヤー)
    for (int i = 0; i < DIM; i++) {
#pragma HLS PIPELINE off
        x[i] = (fx16_t)((acc_t)x[i] + (acc_t)xb[i]);
    }
}
llama_layer_tb.cpp
// llama_layer()のcsim用テストベンチ.
// reference_c/のfp32チェックポイント読み込み・トークナイザをそのまま流用し,
// 埋め込み参照・レイヤーループ・最終分類器はPS側の役割としてこのテストベンチ内で行い,
// 1レイヤー分(Attention+FFN)の計算だけをllama_layer()に委譲する
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include "../src/llama_layer.h"

extern "C" {
#include "weights_loader.h"
#include "tokenizer.h"
}

// fp32のTransformerWeightsから,layer層目の重みをllama_layer.hのオフセット順に量子化して並べる
static void pack_layer_weights(fx16_t* dst, TransformerWeights* w, Config* p, int layer) {
    int dim = p->dim;
    int hidden_dim = p->hidden_dim;
    int kv_dim = (p->dim * p->n_kv_heads) / p->n_heads;

    float* src;
    src = w->rms_att_weight + (size_t)layer * dim;
    for (int i = 0; i < dim; i++) dst[OFFSET_RMS_ATT + i] = (fx16_t)src[i];

    src = w->wq + (size_t)layer * dim * dim;
    for (int i = 0; i < dim * dim; i++) dst[OFFSET_WQ + i] = (fx16_t)src[i];

    src = w->wk + (size_t)layer * dim * kv_dim;
    for (int i = 0; i < dim * kv_dim; i++) dst[OFFSET_WK + i] = (fx16_t)src[i];

    src = w->wv + (size_t)layer * dim * kv_dim;
    for (int i = 0; i < dim * kv_dim; i++) dst[OFFSET_WV + i] = (fx16_t)src[i];

    src = w->wo + (size_t)layer * dim * dim;
    for (int i = 0; i < dim * dim; i++) dst[OFFSET_WO + i] = (fx16_t)src[i];

    src = w->rms_ffn_weight + (size_t)layer * dim;
    for (int i = 0; i < dim; i++) dst[OFFSET_RMS_FFN + i] = (fx16_t)src[i];

    src = w->w1 + (size_t)layer * dim * hidden_dim;
    for (int i = 0; i < dim * hidden_dim; i++) dst[OFFSET_W1 + i] = (fx16_t)src[i];

    src = w->w2 + (size_t)layer * hidden_dim * dim;
    for (int i = 0; i < hidden_dim * dim; i++) dst[OFFSET_W2 + i] = (fx16_t)src[i];

    src = w->w3 + (size_t)layer * dim * hidden_dim;
    for (int i = 0; i < dim * hidden_dim; i++) dst[OFFSET_W3 + i] = (fx16_t)src[i];
}

// 最終RMSNorm(PS側で行う想定).llama_layer.cpp内のrmsnormと同じ計算をここでも行う
static void final_rmsnorm(fx16_t out[DIM], fx16_t x[DIM], float* weight) {
    acc_t ss = 0;
    for (int j = 0; j < DIM; j++) { ss += (acc_t)x[j] * (acc_t)x[j]; }
    double ss_d = (double)ss / DIM + 1e-5;
    float inv_rms = (float)(1.0 / sqrt(ss_d));
    for (int j = 0; j < DIM; j++) { out[j] = (fx16_t)(weight[j] * (inv_rms * (float)x[j])); }
}

// 分類器(PS側).fp32のwclsをその場でfx16_tへキャストしつつacc_tで積算する
static int classify_argmax(fx16_t x[DIM], float* wcls, int vocab_size) {
    int best_i = 0;
    acc_t best_v;
    for (int i = 0; i < vocab_size; i++) {
        acc_t val = 0;
        for (int j = 0; j < DIM; j++) {
            val += (acc_t)((fx16_t)wcls[i * DIM + j]) * (acc_t)x[j];
        }
        if (i == 0 || val > best_v) { best_v = val; best_i = i; }
    }
    return best_i;
}

int main(int argc, char* argv[]) {
    const char* checkpoint_path = argc > 1 ? argv[1] : "../../weights/stories15M.bin";
    const char* tokenizer_path = argc > 2 ? argv[2] : "../../llama2c_upstream/tokenizer.bin";
    const char* prompt = argc > 3 ? argv[3] : "Once upon a time";
    int steps = argc > 4 ? atoi(argv[4]) : 60;
    const char* out_path = argc > 5 ? argv[5] : NULL;
    // cosimはRTLシミュレーションで非常に遅いため,1トークンあたりに実際に呼び出すレイヤー数を絞れるようにする
    // (省略時はconfig.n_layers全部呼ぶ.途中で打ち切るので生成される文章自体は正しくなくなるが,
    //  llama_layer()自体の検証が目的なので問題ない)
    int max_layers_arg = argc > 6 ? atoi(argv[6]) : -1;

    Config config;
    TransformerWeights weights;
    int fd; float* data; ssize_t file_size;
    read_checkpoint(checkpoint_path, &config, &weights, &fd, &data, &file_size);

    if (config.dim != DIM || config.hidden_dim != HIDDEN_DIM || config.n_heads != N_HEADS
        || config.n_kv_heads != N_KV_HEADS || config.seq_len != SEQ_LEN) {
        fprintf(stderr, "checkpointの構造がllama_layer.hのコンパイル時定数と一致しません\n");
        exit(EXIT_FAILURE);
    }

    Tokenizer tokenizer;
    build_tokenizer(&tokenizer, tokenizer_path, config.vocab_size);

    // 6層分の重みを事前に量子化して並べておく(位置によらず不変なので1回だけ行う)
    fx16_t* layer_weights[6];
    for (int l = 0; l < config.n_layers; l++) {
        layer_weights[l] = new fx16_t[LAYER_WEIGHTS_SIZE];
        pack_layer_weights(layer_weights[l], &weights, &config, l);
    }

    // 各レイヤー専用のKVキャッシュ領域を確保
    fx16_t* kv_cache[6];
    for (int l = 0; l < config.n_layers; l++) {
        kv_cache[l] = new fx16_t[2 * SEQ_LEN * KV_DIM];
    }

    int num_prompt_tokens = 0;
    int* prompt_tokens = (int*)malloc((strlen(prompt) + 3) * sizeof(int));
    encode(&tokenizer, prompt, 1, 0, prompt_tokens, &num_prompt_tokens);

    int* generated_tokens = (int*)malloc(steps * sizeof(int));
    int n_generated = 0;

    int token = prompt_tokens[0];
    int pos = 0;
    fx16_t x[DIM];
    while (pos < steps) {
        // 埋め込み参照(PS側).fp32埋め込みテーブルをfx16_tへキャストしてxにコピー
        float* content_row = weights.token_embedding_table + (size_t)token * config.dim;
        for (int i = 0; i < DIM; i++) { x[i] = (fx16_t)content_row[i]; }

        // 6層分,llama_layer()を順番に呼び出す(PS側がループ制御を担う想定を模擬)
        int layers_to_run = (max_layers_arg >= 0) ? max_layers_arg : config.n_layers;
        for (int l = 0; l < layers_to_run; l++) {
            llama_layer(x, layer_weights[l], kv_cache[l], pos);
        }

        // 最終RMSNorm + 分類器(PS側)
        fx16_t x_norm[DIM];
        final_rmsnorm(x_norm, x, weights.rms_final_weight);
        int next_argmax = classify_argmax(x_norm, weights.wcls, config.vocab_size);

        int next;
        if (pos < num_prompt_tokens - 1) {
            next = prompt_tokens[pos + 1];
        } else {
            next = next_argmax;
        }
        pos++;

        if (next == 1) { break; }

        char* piece = decode(&tokenizer, token, next);
        safe_printf(piece);
        fflush(stdout);

        generated_tokens[n_generated++] = next;
        token = next;
    }
    printf("\n");

    if (out_path != NULL) {
        FILE* out_file = fopen(out_path, "w");
        for (int i = 0; i < n_generated; i++) { fprintf(out_file, "%d\n", generated_tokens[i]); }
        fclose(out_file);
    }

    for (int l = 0; l < config.n_layers; l++) { delete[] layer_weights[l]; delete[] kv_cache[l]; }
    free(prompt_tokens);
    free(generated_tokens);
    free_tokenizer(&tokenizer);
    free_checkpoint(fd, data, file_size);
    return 0;
}

テスト実行します.

g++ -O2 -std=c++11 -D_WIN32 -I. -I../reference_c -I"(Vitis HLSのincludeディレクトリのパスだよ)" -c src/llama_layer.cpp -o src/llama_layer.o
g++ -O2 -std=c++11 -D_WIN32 -I. -Isrc -I../reference_c -I"(Vitis HLSのincludeディレクトリのパスだよ)" -c test/llama_layer_tb.cpp -o test/llama_layer_tb.o
g++ -O2 -o test/llama_layer_tb.exe test/llama_layer_tb.o src/llama_layer.o ../reference_c/tokenizer.o ../reference_c/weights_loader.o ../reference_c/win.o -lm
./test/llama_layer_tb.exe ../weights/stories15M.bin ../llama2c_upstream/tokenizer.bin "Once upon a time" 60 test/llama_layer_tb_output.txt
diff ../reference_c/test_vectors/stories15M_fixed20_6_once_upon_a_time.txt test/llama_layer_tb_output.txt

固定小数点で実行したのと同じ結果になりました.

テスト実行のソースを見てもらえれば分かるんですが,PL上に実装するのはセルフアテンションとその後のFFNの一層だけです.
PS側からこれを六回繰り返し呼んで推論する予定です.
これはテストファイルなので,PS側の処理を全部cppファイル内に書いていますが.
なんでこんな実装にしているかというと,ニューラルネットワーク全部をPLに載せようとするとPYNQ-Z2だと入りきらないからです.もっと計算資源がないと全部は入り切りません.もっと言うと,PYNQ-Z2の限られた計算資源では並列実行して効率化することもできず,全部順番に実行します.そのため高速化も望めず,正直FPGA上に実装するメリットはほとんどありません.
なのになぜやっているかというとロマンです.

では,さっきの実行はPC内蔵のコンパイラを使用したので,次はVitis HLSのプロジェクトを作って,Vitis HLSのコンパイラを使用してコンパイルして実行してみます.

run_hls.tcl
# llama_layer HLSプロジェクトのビルドスクリプト(csim/synth/cosim/export)
set proj_name "llama_layer_proj"
set top_func "llama_layer"
set part_name "xc7z020clg400-1"

# このスクリプトが置かれているディレクトリ(hls/scripts)を基準に各パスを解決する
set script_dir [file dirname [info script]]
set hls_dir [file normalize "$script_dir/.."]
set ref_dir [file normalize "$hls_dir/../reference_c"]
set weights_dir [file normalize "$hls_dir/../weights"]
set tokenizer_bin [file normalize "$hls_dir/../llama2c_upstream/tokenizer.bin"]
set checkpoint_bin [file normalize "$weights_dir/stories15M.bin"]

open_project -reset $proj_name

# 合成対象(DUT): 1レイヤー分のHLSカーネル
add_files "$hls_dir/src/llama_layer.cpp" -cflags "-I$hls_dir/src -I$ref_dir"

# テストベンチ: reference_c/の重み読み込み・トークナイザをそのまま流用する
add_files -tb "$hls_dir/test/llama_layer_tb.cpp" -cflags "-I$hls_dir/src -I$ref_dir -D_WIN32"
add_files -tb "$ref_dir/weights_loader.c" -cflags "-I$ref_dir -D_WIN32"
add_files -tb "$ref_dir/tokenizer.c" -cflags "-I$ref_dir -D_WIN32"
add_files -tb "$ref_dir/win.c" -cflags "-I$ref_dir -D_WIN32"

set_top $top_func

open_solution -reset "solution1"
set_part $part_name
create_clock -period 10 -name default_clk ;# 100MHz(PYNQ-Z2の標準的なPL動作クロックを想定)

# C++コードとして普通に実行するだけ.ハードウェアの記述はない
csim_design -argv [list $checkpoint_bin $tokenizer_bin "Once upon a time" 60]

puts "CSIM_DONE"
exit
vitis_hls -f scripts/run_hls.tcl

問題なく文章生成されました.固定小数点で実行した時と同じ結果になりました.
ついに論理合成を行います.これでやっと,FPGA上で動かすためのファイルが作成されます.

run_hls.tcl
# llama_layer HLSプロジェクトのビルドスクリプト(csim/synth/cosim/export)
set proj_name "llama_layer_proj"
set top_func "llama_layer"
set part_name "xc7z020clg400-1"

# このスクリプトが置かれているディレクトリ(hls/scripts)を基準に各パスを解決する
set script_dir [file dirname [info script]]
set hls_dir [file normalize "$script_dir/.."]
set ref_dir [file normalize "$hls_dir/../reference_c"]
set weights_dir [file normalize "$hls_dir/../weights"]
set tokenizer_bin [file normalize "$hls_dir/../llama2c_upstream/tokenizer.bin"]
set checkpoint_bin [file normalize "$weights_dir/stories15M.bin"]

open_project -reset $proj_name

# 合成対象(DUT): 1レイヤー分のHLSカーネル
add_files "$hls_dir/src/llama_layer.cpp" -cflags "-I$hls_dir/src -I$ref_dir"

# テストベンチ: reference_c/の重み読み込み・トークナイザをそのまま流用する
add_files -tb "$hls_dir/test/llama_layer_tb.cpp" -cflags "-I$hls_dir/src -I$ref_dir -D_WIN32"
add_files -tb "$ref_dir/weights_loader.c" -cflags "-I$ref_dir -D_WIN32"
add_files -tb "$ref_dir/tokenizer.c" -cflags "-I$ref_dir -D_WIN32"
add_files -tb "$ref_dir/win.c" -cflags "-I$ref_dir -D_WIN32"

set_top $top_func

open_solution -reset "solution1"
set_part $part_name
create_clock -period 10 -name default_clk ;# 100MHz(PYNQ-Z2の標準的なPL動作クロックを想定)

# csimは前回のvitis_hls実行で確認済み(CSim done with 0 errors,約26分かかる)なので今回はスキップする.
# 再度確認したい場合はこの行のコメントを外す
# csim_design -argv [list $checkpoint_bin $tokenizer_bin "Once upon a time" 60]

puts "CSIM_SKIPPED"

# 論理合成(C synthesis).C++コードから実際のRTL回路(Verilog)を生成し,
# LUT/FF/BRAM/DSPのリソース使用量とタイミングのレポートを得る
csynth_design

puts "CSYNTH_DONE"
exit
vitis_hls -f scripts/run_hls.tcl

ここで,全部並列実行するようになっている場合はLUTを113%も使用することになっていました.
パイプラインをオフにしたことで79%まで少なくなりました.

これで高位合成が完了.Cで書いたソースがハードウェア記述言語に翻訳されました.

シミュレーションしてみる

実機で動かす前にシミュレーションしてみます.

run_hls.tcl
# llama_layer HLSプロジェクトのビルドスクリプト(csim/synth/cosim/export)
set proj_name "llama_layer_proj"
set top_func "llama_layer"
set part_name "xc7z020clg400-1"

# このスクリプトが置かれているディレクトリ(hls/scripts)を基準に各パスを解決する
set script_dir [file dirname [info script]]
set hls_dir [file normalize "$script_dir/.."]
set ref_dir [file normalize "$hls_dir/../reference_c"]
set weights_dir [file normalize "$hls_dir/../weights"]
set tokenizer_bin [file normalize "$hls_dir/../llama2c_upstream/tokenizer.bin"]
set checkpoint_bin [file normalize "$weights_dir/stories15M.bin"]

open_project -reset $proj_name

# 合成対象(DUT): 1レイヤー分のHLSカーネル
add_files "$hls_dir/src/llama_layer.cpp" -cflags "-I$hls_dir/src -I$ref_dir"

# テストベンチ: reference_c/の重み読み込み・トークナイザをそのまま流用する
add_files -tb "$hls_dir/test/llama_layer_tb.cpp" -cflags "-I$hls_dir/src -I$ref_dir -D_WIN32"
add_files -tb "$ref_dir/weights_loader.c" -cflags "-I$ref_dir -D_WIN32"
add_files -tb "$ref_dir/tokenizer.c" -cflags "-I$ref_dir -D_WIN32"
add_files -tb "$ref_dir/win.c" -cflags "-I$ref_dir -D_WIN32"

set_top $top_func

open_solution -reset "solution1"
set_part $part_name
create_clock -period 10 -name default_clk ;# 100MHz(PYNQ-Z2の標準的なPL動作クロックを想定)

# csimは前回のvitis_hls実行で確認済み(CSim done with 0 errors,約26分かかる)なので今回はスキップする.
# 再度確認したい場合はこの行のコメントを外す
# csim_design -argv [list $checkpoint_bin $tokenizer_bin "Once upon a time" 60]

puts "CSIM_SKIPPED"

# 論理合成(C synthesis).C++コードから実際のRTL回路(Verilog)を生成し,
# LUT/FF/BRAM/DSPのリソース使用量とタイミングのレポートを得る
csynth_design

puts "CSYNTH_DONE"

# C/RTLコシミュレーション.生成されたRTL回路を実際にシミュレータで動かし,csimと同じ結果になるか確認する.
# RTLシミュレーションは非常に遅い(浮動小数点IPコアをDSP48E1レベルで詳細にシミュレートするため)ため,
# llama_layer()を1回だけ呼び出す最小構成(1トークン・1レイヤー)に絞って実行する
cosim_design -argv [list $checkpoint_bin $tokenizer_bin "Once upon a time" 1 "cosim_out.txt" 1]

puts "COSIM_DONE"
exit
vitis_hls -f scripts/run_hls.tcl

ちなみに「高位合成はもうしているからスキップ」とは出来ないそうです.
open_solution -resetで前回の合成結果を消してしまっているため.

最初は2トークンだけ出力しようかと思ったのですが,そんなことしていると数時間かかりそうだったので1レイヤーだけの実行にしました.
これでも17分くらいかかったんですけどね.

IPパッケージにする

とんでもなく遅かったですが,ともかく動くことは分かったのでIPパッケージにします.

run_hls.tcl
# llama_layer HLSプロジェクトのビルドスクリプト(csim/synth/cosim/export)
set proj_name "llama_layer_proj"
set top_func "llama_layer"
set part_name "xc7z020clg400-1"

# このスクリプトが置かれているディレクトリ(hls/scripts)を基準に各パスを解決する
set script_dir [file dirname [info script]]
set hls_dir [file normalize "$script_dir/.."]
set ref_dir [file normalize "$hls_dir/../reference_c"]
set weights_dir [file normalize "$hls_dir/../weights"]
set tokenizer_bin [file normalize "$hls_dir/../llama2c_upstream/tokenizer.bin"]
set checkpoint_bin [file normalize "$weights_dir/stories15M.bin"]

open_project -reset $proj_name

# 合成対象(DUT): 1レイヤー分のHLSカーネル
add_files "$hls_dir/src/llama_layer.cpp" -cflags "-I$hls_dir/src -I$ref_dir"

# テストベンチ: reference_c/の重み読み込み・トークナイザをそのまま流用する
add_files -tb "$hls_dir/test/llama_layer_tb.cpp" -cflags "-I$hls_dir/src -I$ref_dir -D_WIN32"
add_files -tb "$ref_dir/weights_loader.c" -cflags "-I$ref_dir -D_WIN32"
add_files -tb "$ref_dir/tokenizer.c" -cflags "-I$ref_dir -D_WIN32"
add_files -tb "$ref_dir/win.c" -cflags "-I$ref_dir -D_WIN32"

set_top $top_func

open_solution -reset "solution1"
set_part $part_name
create_clock -period 10 -name default_clk ;# 100MHz(PYNQ-Z2の標準的なPL動作クロックを想定)

# csimは前回のvitis_hls実行で確認済み(CSim done with 0 errors,約26分かかる)なので今回はスキップする.
# 再度確認したい場合はこの行のコメントを外す
# csim_design -argv [list $checkpoint_bin $tokenizer_bin "Once upon a time" 60]

puts "CSIM_SKIPPED"

# 論理合成(C synthesis).C++コードから実際のRTL回路(Verilog)を生成し,
# LUT/FF/BRAM/DSPのリソース使用量とタイミングのレポートを得る.
# open_solution -resetで毎回ソリューションを作り直すため,cosimに使うRTLを用意するには必須(csynth自体は約21秒と軽い)
csynth_design

puts "CSYNTH_DONE"

# cosimは前回確認済み(C/RTL co-simulation finished: PASS,約17分かかる)なので今回はスキップする
# cosim_design -argv [list $checkpoint_bin $tokenizer_bin "Once upon a time" 1 "cosim_out.txt" 1]
puts "COSIM_SKIPPED"

# HLS IPとしてパッケージ化(Vivado IP Catalogに取り込める形式でエクスポート)
export_design -rtl verilog -format ip_catalog

puts "EXPORT_DONE"
exit
vitis_hls -f scripts/run_hls.tcl

ブロックデザインを作成

さっき作ったIPパッケージを利用して,ブロックデザインを作成します.

build_bd.tcl
# llama_layer IPをPS(Zynq)と接続するVivadoブロックデザイン作成スクリプト
# M7: IPパッケージ化とVivadoブロックデザイン作成(ビットストリーム生成はまだ行わない)

set proj_name "llama_accel"
set proj_dir  "./llama_accel_proj"
set part_name "xc7z020clg400-1"
set board_vlnv "tul.com.tw:pynq-z2:part0:1.0"

set script_dir [file dirname [info script]]
set vivado_dir [file normalize "$script_dir/.."]
set ip_repo_path [file normalize "$vivado_dir/../hls/llama_layer_proj/solution1/impl/ip"]

create_project $proj_name $proj_dir -part $part_name -force
set_property board_part $board_vlnv [current_project]

# llama_layer IPをカタログに追加
set_property ip_repo_paths $ip_repo_path [current_project]
update_ip_catalog -rebuild

create_bd_design "llama_accel_bd"

# PS(Zynq)を配置しボードプリセットを適用
create_bd_cell -type ip -vlnv xilinx.com:ip:processing_system7:5.5 processing_system7_0
apply_bd_automation -rule xilinx.com:bd_rule:processing_system7 \
  -config {make_external "FIXED_IO, DDR" apply_board_preset "1" Master "Disable" Slave "Disable"} \
  [get_bd_cells processing_system7_0]

# m_axi(activations/weights/kvcache)を3本とも受けられるよう,HPポートを3つ有効化する
set_property -dict [list \
  CONFIG.PCW_USE_S_AXI_HP0 {1} \
  CONFIG.PCW_USE_S_AXI_HP1 {1} \
  CONFIG.PCW_USE_S_AXI_HP2 {1} \
] [get_bd_cells processing_system7_0]

# 注: CONFIG.PCW_M_AXI_GP0_FREQMHZ/PCW_S_AXI_HP{0,1,2}_FREQMHZは値が10のまま.
# set_propertyで100を指定するとread-only(CRITICAL WARNING: BD 41-737)で拒否され,
# reset_propertyもこのプロパティでは非対応(ERROR: Common 17-229)で使えなかった.
# 一方,実際にPLへ供給されるクロックであるCONFIG.PCW_CLK0_FREQは100000000(100MHz),
# CONFIG.PCW_ACT_FPGA0_PERIPHERAL_FREQMHZも100.000000で,実配線のクロックは100MHzになっている.
# これらFREQMHZプロパティは表示・見積もり用の値で,実際の合成・タイミング制約は
# 配線されたクロックネットから導出されると推測されるが,未確認

# llama_layer IPを配置
create_bd_cell -type ip -vlnv xilinx.com:hls:llama_layer:1.0 llama_layer_0

# m_axi(データ本体,3本)をそれぞれ別のHPポートに接続(独立したバースト転送のため)
apply_bd_automation -rule xilinx.com:bd_rule:axi4 \
  -config { Master "/llama_layer_0/m_axi_activations" Clk "Auto" } \
  [get_bd_intf_pins processing_system7_0/S_AXI_HP0]

apply_bd_automation -rule xilinx.com:bd_rule:axi4 \
  -config { Master "/llama_layer_0/m_axi_weights" Clk "Auto" } \
  [get_bd_intf_pins processing_system7_0/S_AXI_HP1]

apply_bd_automation -rule xilinx.com:bd_rule:axi4 \
  -config { Master "/llama_layer_0/m_axi_kvcache" Clk "Auto" } \
  [get_bd_intf_pins processing_system7_0/S_AXI_HP2]

# s_axi_control(制御レジスタ)をPSのGP0マスターに接続
apply_bd_automation -rule xilinx.com:bd_rule:axi4 \
  -config { Master "/processing_system7_0/M_AXI_GP0" Clk "Auto" } \
  [get_bd_intf_pins llama_layer_0/s_axi_control]

# ブロック配置を整形(見た目のみ,動作には影響しない)
regenerate_bd_layout
# 未接続ポート・アドレスマップの衝突・クロック配線などの整合性を検証
validate_bd_design

# ここまでの変更はメモリ上にしかないため,明示的に.bdファイルへ保存する
save_bd_design

puts "BD_VALIDATE_DONE"
exit
vivado -mode batch -source scripts/build_bd.tcl -log build_bd.log -journal build_bd.jou

bitファイル生成

ついに回路が完成したので合成します.

build_bd.tcl
# llama_layer IPをPS(Zynq)と接続するVivadoブロックデザイン作成スクリプト
# M7: IPパッケージ化とVivadoブロックデザイン作成
# M8: 論理合成・配置配線・ビットストリーム生成

set proj_name "llama_accel"
set proj_dir  "./llama_accel_proj"
set part_name "xc7z020clg400-1"
set board_vlnv "tul.com.tw:pynq-z2:part0:1.0"

set script_dir [file dirname [info script]]
set vivado_dir [file normalize "$script_dir/.."]
set ip_repo_path [file normalize "$vivado_dir/../hls/llama_layer_proj/solution1/impl/ip"]

create_project $proj_name $proj_dir -part $part_name -force
set_property board_part $board_vlnv [current_project]

# llama_layer IPをカタログに追加
set_property ip_repo_paths $ip_repo_path [current_project]
update_ip_catalog -rebuild

create_bd_design "llama_accel_bd"

# PS(Zynq)を配置しボードプリセットを適用
create_bd_cell -type ip -vlnv xilinx.com:ip:processing_system7:5.5 processing_system7_0
apply_bd_automation -rule xilinx.com:bd_rule:processing_system7 \
  -config {make_external "FIXED_IO, DDR" apply_board_preset "1" Master "Disable" Slave "Disable"} \
  [get_bd_cells processing_system7_0]

# m_axi(activations/weights/kvcache)を3本とも受けられるよう,HPポートを3つ有効化する
set_property -dict [list \
  CONFIG.PCW_USE_S_AXI_HP0 {1} \
  CONFIG.PCW_USE_S_AXI_HP1 {1} \
  CONFIG.PCW_USE_S_AXI_HP2 {1} \
] [get_bd_cells processing_system7_0]

# 注: CONFIG.PCW_M_AXI_GP0_FREQMHZ/PCW_S_AXI_HP{0,1,2}_FREQMHZは値が10のまま.
# set_propertyで100を指定するとread-only(CRITICAL WARNING: BD 41-737)で拒否され,
# reset_propertyもこのプロパティでは非対応(ERROR: Common 17-229)で使えなかった.
# 一方,実際にPLへ供給されるクロックであるCONFIG.PCW_CLK0_FREQは100000000(100MHz),
# CONFIG.PCW_ACT_FPGA0_PERIPHERAL_FREQMHZも100.000000で,実配線のクロックは100MHzになっている.
# これらFREQMHZプロパティは表示・見積もり用の値で,実際の合成・タイミング制約は
# 配線されたクロックネットから導出されると推測されるが,未確認

# llama_layer IPを配置
create_bd_cell -type ip -vlnv xilinx.com:hls:llama_layer:1.0 llama_layer_0

# m_axi(データ本体,3本)をそれぞれ別のHPポートに接続(独立したバースト転送のため)
apply_bd_automation -rule xilinx.com:bd_rule:axi4 \
  -config { Master "/llama_layer_0/m_axi_activations" Clk "Auto" } \
  [get_bd_intf_pins processing_system7_0/S_AXI_HP0]

apply_bd_automation -rule xilinx.com:bd_rule:axi4 \
  -config { Master "/llama_layer_0/m_axi_weights" Clk "Auto" } \
  [get_bd_intf_pins processing_system7_0/S_AXI_HP1]

apply_bd_automation -rule xilinx.com:bd_rule:axi4 \
  -config { Master "/llama_layer_0/m_axi_kvcache" Clk "Auto" } \
  [get_bd_intf_pins processing_system7_0/S_AXI_HP2]

# s_axi_control(制御レジスタ)をPSのGP0マスターに接続
apply_bd_automation -rule xilinx.com:bd_rule:axi4 \
  -config { Master "/processing_system7_0/M_AXI_GP0" Clk "Auto" } \
  [get_bd_intf_pins llama_layer_0/s_axi_control]

# ブロック配置を整形(見た目のみ,動作には影響しない)
regenerate_bd_layout
# 未接続ポート・アドレスマップの衝突・クロック配線などの整合性を検証
validate_bd_design

# ここまでの変更はメモリ上にしかないため,明示的に.bdファイルへ保存する
save_bd_design

puts "BD_VALIDATE_DONE"

# ブロックデザインを論理合成できるHDLのトップファイルとしてラップし,合成対象に登録
make_wrapper -files [get_files $proj_dir/$proj_name.srcs/sources_1/bd/llama_accel_bd/llama_accel_bd.bd] -top
add_files -norecurse $proj_dir/$proj_name.gen/sources_1/bd/llama_accel_bd/hdl/llama_accel_bd_wrapper.v
update_compile_order -fileset sources_1
set_property top llama_accel_bd_wrapper [current_fileset]

# 論理合成
launch_runs synth_1 -jobs 4
wait_on_run synth_1

# 配置配線してビットストリームを生成
launch_runs impl_1 -to_step write_bitstream -jobs 4
wait_on_run impl_1

# タイミングクロージャ(negative slackがないか)を確認
open_run impl_1
set slack [get_property SLACK [get_timing_paths -delay_type min_max]]
puts "TIMING_SLACK: $slack"

# .bitと.hwhを出力ディレクトリにまとめて配置
file mkdir "./output"
file copy -force "$proj_dir/$proj_name.runs/impl_1/llama_accel_bd_wrapper.bit" "./output/llama_accel.bit"
file copy -force "$proj_dir/$proj_name.gen/sources_1/bd/llama_accel_bd/hw_handoff/llama_accel_bd.hwh" "./output/llama_accel.hwh"

puts "BITSTREAM_DONE"
exit
vivado -mode batch -source scripts/build_bd.tcl -log build_bd.log -journal build_bd.jou

正確な時間は測っていませんが,かなり時間がかかります.

実機で確認

さて,ついに実機で動作させます.
先ほどの合成で作成されたbitファイルとhwhファイルをPYNQ-Z2にアップロードします.

そうしたら以下を実行します.

from pynq import Overlay

ol = Overlay("llama_accel.bit")
print(ol.ip_dict.keys())
print(ol.llama_layer_0.register_map)

実行結果がこうなればOKです

dict_keys(['llama_layer_0', 'processing_system7_0'])
RegisterMap {
  CTRL = Register(AP_START=0, AP_DONE=0, AP_IDLE=1, AP_READY=0, RESERVED_1=0, AUTO_RESTART=0, RESERVED_2=0, INTERRUPT=0, RESERVED_3=0),
  GIER = Register(Enable=0, RESERVED=0),
  IP_IER = Register(CHAN0_INT_EN=0, CHAN1_INT_EN=0, RESERVED_0=0),
  IP_ISR = Register(CHAN0_INT_ST=0, CHAN1_INT_ST=0, RESERVED_0=0),
  x_1 = Register(x=write-only),
  x_2 = Register(x=write-only),
  weights_offset_1 = Register(weights_offset=write-only),
  weights_offset_2 = Register(weights_offset=write-only),
  kv_cache_1 = Register(kv_cache=write-only),
  kv_cache_2 = Register(kv_cache=write-only),
  pos_r = Register(pos_r=write-only)
}

これで,bitファイルがPYNQ-Z2上でロードできることが分かりました.

モデルの重みを量子化する

以下のソースを実行します.
浮動小数点で保存されているモデルの重みを固定小数点にして,FPGA上で実行しやすいようにします.

dump_weights.cpp
// stories15M.binのfp32重みを,llama_layer.hのオフセット順にap_fixed<20,6>へ量子化し,
// 6レイヤー分を連結した1つのバイナリファイルへ書き出す(M9: PYNQ-Z2のDDRにそのまま転送する用)
#include <cstdio>
#include <cstdlib>
#include "../src/llama_layer.h"

extern "C" {
#include "weights_loader.h"
}

// llama_layer_tb.cppのpack_layer_weights()と同じロジック(重複コードだが,
// テストベンチとこのツールで依存関係を分離するためあえてここにも実装する)
static void pack_layer_weights(fx16_t* dst, TransformerWeights* w, Config* p, int layer) {
    int dim = p->dim;
    int hidden_dim = p->hidden_dim;
    int kv_dim = (p->dim * p->n_kv_heads) / p->n_heads;

    float* src;
    src = w->rms_att_weight + (size_t)layer * dim;
    for (int i = 0; i < dim; i++) dst[OFFSET_RMS_ATT + i] = (fx16_t)src[i];

    src = w->wq + (size_t)layer * dim * dim;
    for (int i = 0; i < dim * dim; i++) dst[OFFSET_WQ + i] = (fx16_t)src[i];

    src = w->wk + (size_t)layer * dim * kv_dim;
    for (int i = 0; i < dim * kv_dim; i++) dst[OFFSET_WK + i] = (fx16_t)src[i];

    src = w->wv + (size_t)layer * dim * kv_dim;
    for (int i = 0; i < dim * kv_dim; i++) dst[OFFSET_WV + i] = (fx16_t)src[i];

    src = w->wo + (size_t)layer * dim * dim;
    for (int i = 0; i < dim * dim; i++) dst[OFFSET_WO + i] = (fx16_t)src[i];

    src = w->rms_ffn_weight + (size_t)layer * dim;
    for (int i = 0; i < dim; i++) dst[OFFSET_RMS_FFN + i] = (fx16_t)src[i];

    src = w->w1 + (size_t)layer * dim * hidden_dim;
    for (int i = 0; i < dim * hidden_dim; i++) dst[OFFSET_W1 + i] = (fx16_t)src[i];

    src = w->w2 + (size_t)layer * hidden_dim * dim;
    for (int i = 0; i < hidden_dim * dim; i++) dst[OFFSET_W2 + i] = (fx16_t)src[i];

    src = w->w3 + (size_t)layer * dim * hidden_dim;
    for (int i = 0; i < dim * hidden_dim; i++) dst[OFFSET_W3 + i] = (fx16_t)src[i];
}

int main(int argc, char* argv[]) {
    if (argc < 3) {
        fprintf(stderr, "Usage: dump_weights <checkpoint.bin> <output.bin>\n");
        return EXIT_FAILURE;
    }
    const char* checkpoint_path = argv[1];
    const char* output_path = argv[2];

    Config config;
    TransformerWeights weights;
    int fd; float* data; ssize_t file_size;
    read_checkpoint(checkpoint_path, &config, &weights, &fd, &data, &file_size);

    if (config.dim != DIM || config.hidden_dim != HIDDEN_DIM || config.n_heads != N_HEADS
        || config.n_kv_heads != N_KV_HEADS || config.seq_len != SEQ_LEN) {
        fprintf(stderr, "checkpointの構造がllama_layer.hのコンパイル時定数と一致しません\n");
        return EXIT_FAILURE;
    }

    FILE* out = fopen(output_path, "wb");
    if (!out) { fprintf(stderr, "couldn't open output file %s\n", output_path); return EXIT_FAILURE; }

    fx16_t* layer_buf = new fx16_t[LAYER_WEIGHTS_SIZE];
    for (int l = 0; l < config.n_layers; l++) {
        pack_layer_weights(layer_buf, &weights, &config, l);
        fwrite(layer_buf, sizeof(fx16_t), LAYER_WEIGHTS_SIZE, out);
        printf("layer %d packed (%d elements, %zu bytes)\n", l, LAYER_WEIGHTS_SIZE, LAYER_WEIGHTS_SIZE * sizeof(fx16_t));
    }
    delete[] layer_buf;
    fclose(out);

    printf("done: %s (%d layers x %d elements x %zu bytes)\n", output_path, config.n_layers, LAYER_WEIGHTS_SIZE, sizeof(fx16_t));

    free_checkpoint(fd, data, file_size);
    return 0;
}
g++ -O2 -std=c++11 -D_WIN32 -I. -I../src -I../../reference_c -I"(Vitis HLSのincludeディレクトリのパスだよ)" -c dump_weights.cpp -o dump_weights.o
g++ -O2 -o dump_weights.exe dump_weights.o ../../reference_c/weights_loader.o ../../reference_c/win.o -lm
./dump_weights.exe ../../weights/stories15M.bin ../../weights/stories15M_fixed.bin

PS側の処理を作成する

以前説明した通りPS側でかなりの処理を行いますので,それを実装します.

llama_infer.py
# PYNQ-Z2のJupyter上で実行するPS側オーケストレーション(M10).
# トークナイズ・埋め込み参照・分類器はPython(NumPy)で実装し,
# Transformer 1レイヤー分(Attention+FFN)の計算だけをFPGA上のllama_layer_0 IPに投げる.
import struct
import time
import numpy as np
from pynq import Overlay, allocate

# stories15Mのモデル構造(reference_c/hlsのコンパイル時定数と同じ値)
DIM = 288
HIDDEN_DIM = 768
N_LAYERS = 6
N_HEADS = 6
KV_DIM = DIM  # n_kv_heads == n_heads のため
SEQ_LEN = 256
VOCAB_SIZE = 32000
LAYER_WEIGHTS_SIZE = 995904  # hls/src/llama_layer.h と同じ値


# --- トークナイザ(tokenizer.cのPython移植.BPEマージのロジックはCと同一) ---
class Tokenizer:
    def __init__(self, path, vocab_size):
        self.vocab_size = vocab_size
        self.vocab = []
        self.scores = []
        with open(path, "rb") as f:
            self.max_token_length = struct.unpack("<i", f.read(4))[0]
            for _ in range(vocab_size):
                score = struct.unpack("<f", f.read(4))[0]
                length = struct.unpack("<i", f.read(4))[0]
                piece = f.read(length)
                self.scores.append(score)
                self.vocab.append(piece)
        self.str_to_id = {v: i for i, v in enumerate(self.vocab)}

    def encode(self, text, bos=True, eos=False):
        tokens = []
        if bos:
            tokens.append(1)
        if len(text) > 0:
            tokens.append(self.str_to_id[b" "])

        data = text.encode("utf-8")
        i = 0
        while i < len(data):
            # UTF-8の1コードポイント分の長さを判定
            b0 = data[i]
            if b0 & 0x80 == 0:
                clen = 1
            elif b0 & 0xE0 == 0xC0:
                clen = 2
            elif b0 & 0xF0 == 0xE0:
                clen = 3
            else:
                clen = 4
            piece = data[i:i + clen]
            if piece in self.str_to_id:
                tokens.append(self.str_to_id[piece])
            else:
                # 語彙になければ1バイトずつ個別トークンにフォールバック(先頭3つは特殊トークン)
                for b in piece:
                    tokens.append(b + 3)
            i += clen

        # スコアが最も高い隣接トークンのペアを繰り返しマージ(BPE本体)
        while True:
            best_score = -1e10
            best_id = -1
            best_idx = -1
            for i in range(len(tokens) - 1):
                merged = self.vocab[tokens[i]] + self.vocab[tokens[i + 1]]
                mid = self.str_to_id.get(merged)
                if mid is not None and self.scores[mid] > best_score:
                    best_score = self.scores[mid]
                    best_id = mid
                    best_idx = i
            if best_idx == -1:
                break
            tokens[best_idx] = best_id
            del tokens[best_idx + 1]

        if eos:
            tokens.append(2)
        return tokens

    def decode(self, prev_token, token):
        piece = self.vocab[token]
        if prev_token == 1 and piece[:1] == b" ":
            piece = piece[1:]
        # "<0xXX>"形式の生バイト表現を実際のバイトに変換
        if piece.startswith(b"<0x") and piece.endswith(b">"):
            try:
                byte_val = int(piece[3:-1], 16)
                piece = bytes([byte_val])
            except ValueError:
                pass
        return piece


# --- 重み読み込み(埋め込み行列・最終RMSNormのみPython側で直接読む,分類器は埋め込みと共有) ---
def load_embedding_and_final_norm(checkpoint_path):
    with open(checkpoint_path, "rb") as f:
        f.seek(28)  # Configヘッダ(7 x int32)を読み飛ばす
        embedding = np.fromfile(f, dtype="<f4", count=VOCAB_SIZE * DIM).reshape(VOCAB_SIZE, DIM)
        # rms_final_weightは「埋め込み + 6層分のLAYER_WEIGHTS_SIZE」だけ先にある
        skip_floats = N_LAYERS * LAYER_WEIGHTS_SIZE
        f.seek(skip_floats * 4, 1)
        rms_final_weight = np.fromfile(f, dtype="<f4", count=DIM)
    return embedding, rms_final_weight


# --- fx16_t(ap_fixed<20,6,AP_TRN,AP_SAT>)との相互変換 ---
# 実測により,32bitコンテナの中身は「value * 2^14」を負方向へ切り捨て(floor)し,
# ±32の範囲で飽和させた符号付き整数そのものだと確認済み(hls/src/llama_layer.hのfx16_tと同一)
FX_FRAC_BITS = 14
FX_SCALE = 1 << FX_FRAC_BITS       # 16384
FX_RAW_MIN = -(1 << 19)            # -524288 ( = -32.0)
FX_RAW_MAX = (1 << 19) - 1         #  524287 ( =  31.999939...)


def to_fx16(values):
    raw = np.floor(np.asarray(values, dtype=np.float64) * FX_SCALE)
    raw = np.clip(raw, FX_RAW_MIN, FX_RAW_MAX)
    return raw.astype(np.int32)


def from_fx16(raw):
    # HLSはap_fixed<20,...>の書き込み時,32bitコンテナのうち意味のある下位20bit分の
    # バイトにしかAXIの書き込みストローブ(WSTRB)を立てないため,上位12bit(4バイト目)には
    # 「書き込む前にそこにあった古い値の名残」が残ったままになる.
    # そのため読み出す側で下位20bitだけを取り出し,自分で符号拡張し直す必要がある
    raw20 = np.asarray(raw, dtype=np.int64) & 0xFFFFF
    raw20 = np.where(raw20 & 0x80000, raw20 - 0x100000, raw20)
    return raw20.astype(np.float64) / FX_SCALE


def rmsnorm(x, weight):
    ss = np.sum(x.astype(np.float64) ** 2) / DIM + 1e-5
    inv_rms = 1.0 / np.sqrt(ss)
    return (weight * (inv_rms * x)).astype(np.float32)


def classify_argmax(x_norm, embedding):
    logits = embedding.astype(np.float32) @ x_norm.astype(np.float32)
    return int(np.argmax(logits))


# --- FPGA IP呼び出し ---
def call_llama_layer(ip, x_buf, weights_phys, kv_cache_phys, pos):
    # ベースアドレスと位置をレジスタへ書き込む(64bitアドレスを32bit x 2に分割)
    ip.register_map.x_1 = x_buf.physical_address & 0xFFFFFFFF
    ip.register_map.x_2 = (x_buf.physical_address >> 32) & 0xFFFFFFFF
    ip.register_map.weights_offset_1 = weights_phys & 0xFFFFFFFF
    ip.register_map.weights_offset_2 = (weights_phys >> 32) & 0xFFFFFFFF
    ip.register_map.kv_cache_1 = kv_cache_phys & 0xFFFFFFFF
    ip.register_map.kv_cache_2 = (kv_cache_phys >> 32) & 0xFFFFFFFF
    ip.register_map.pos_r = pos

    # 計算開始を指示し,完了(AP_DONE)を待つ
    ip.register_map.CTRL.AP_START = 1
    while ip.register_map.CTRL.AP_DONE == 0:
        pass
    # AP_STARTを立てたままだと,次にAP_IDLEに戻った瞬間にIPが勝手に再スタートしてしまうため,
    # ここで明示的にクリアする
    ip.register_map.CTRL.AP_START = 0


def load_layer_weight_buffers(weights_bin_path):
    """stories15M_fixed.bin(dump_weights.exeで生成済み)を6層分,pynq.allocateしたバッファへ読み込む"""
    buffers = []
    with open(weights_bin_path, "rb") as f:
        for l in range(N_LAYERS):
            raw = np.fromfile(f, dtype="<i4", count=LAYER_WEIGHTS_SIZE)
            buf = allocate(shape=(LAYER_WEIGHTS_SIZE,), dtype=np.int32)
            buf[:] = raw
            buf.flush()
            buffers.append(buf)
    return buffers


def alloc_kv_cache_buffers():
    """各レイヤー専用のKVキャッシュ領域をゼロ初期化して確保する"""
    buffers = []
    for l in range(N_LAYERS):
        buf = allocate(shape=(2 * SEQ_LEN * KV_DIM,), dtype=np.int32)
        buf[:] = 0
        buf.flush()
        buffers.append(buf)
    return buffers


def generate(overlay, tokenizer, embedding, rms_final_weight, weight_bufs, kv_cache_bufs,
             prompt="Once upon a time", steps=60):
    ip = overlay.llama_layer_0
    x_buf = allocate(shape=(DIM,), dtype=np.int32)  # fx16_tは4byteコンテナなのでint32として確保

    prompt_tokens = tokenizer.encode(prompt, bos=True, eos=False)
    num_prompt_tokens = len(prompt_tokens)

    token = prompt_tokens[0]
    pos = 0
    output_text = b""
    start_time = None  # 1トークン目は初期化コストが乗るので,計測は2トークン目以降で開始する
    while pos < steps:
        # 埋め込み参照(PS側).fp32の埋め込み行をfx16_tへ量子化してxバッファへ書き込む
        x_fp32 = embedding[token]
        x_buf[:] = to_fx16(x_fp32)
        x_buf.flush()

        # 6層分,llama_layer IPを順番に呼び出す
        for l in range(N_LAYERS):
            call_llama_layer(ip, x_buf, weight_bufs[l].physical_address, kv_cache_bufs[l].physical_address, pos)
            x_buf.invalidate()  # IPがDDRへ書き戻した最新の値をキャッシュから読み直す

        # 最終RMSNorm + 分類器(PS側,fp32で計算.fx16_tへの量子化はここでは行わない簡略版)
        x_after = from_fx16(np.array(x_buf)).astype(np.float32)
        x_norm = rmsnorm(x_after, rms_final_weight)
        next_argmax = classify_argmax(x_norm, embedding)

        if pos < num_prompt_tokens - 1:
            next_id = prompt_tokens[pos + 1]
        else:
            next_id = next_argmax
        pos += 1

        if next_id == 1:  # BOSが出たら終端とみなす
            break

        piece = tokenizer.decode(token, next_id)
        output_text += piece
        print(piece.decode("utf-8", errors="replace"), end="", flush=True)
        token = next_id

        if start_time is None:
            start_time = time.time()  # 1トークン目の生成が終わった直後からタイマー開始

    print()
    if pos > 1:
        elapsed = time.time() - start_time
        print(f"achieved tok/s: {(pos - 1) / elapsed:.6f}")
    x_buf.close()
    return output_text


if __name__ == "__main__":
    overlay = Overlay("llama_accel.bit")
    tokenizer = Tokenizer("tokenizer.bin", VOCAB_SIZE)
    embedding, rms_final_weight = load_embedding_and_final_norm("stories15M.bin")
    weight_bufs = load_layer_weight_buffers("stories15M_fixed.bin")
    kv_cache_bufs = alloc_kv_cache_buffers()

    generate(overlay, tokenizer, embedding, rms_final_weight, weight_bufs, kv_cache_bufs,
             prompt="Once upon a time", steps=60)

FPGA上で推論する

以下のファイルをbitファイルと同じディレクトリにアップロードします.

  • stories15M.bin (量子化前の,浮動小数点での重み.トークンIDからのベクトルへの変換と,出力層に使用する)
  • stories15M_fixed.bin (6層ぶんの量子化済み重み.セルフアテンションとFFNに使用する)
  • tokenizer.bin (文字列とトークンIDを相互変換するための情報.git cloneしたときについでに取得されています)
  • llama_infer.py (実行ファイル)

ついに推論実行です.

# jupyterノートブックのセルで実行する
%run llama_infer.py

結果がこちら.

Once upon a time, there was a little girl named Lily. She loved to play outside in the sunshine. One day, she saw a big, red ball in the sky. It was the sun! She thought it was so pretty.
Lily wanted to play with the ball
achieved tok/s: 1.355500

日本語訳がこちら.

昔々、リリーという名の小さな女の子がいました。彼女は太陽の光を浴びて外で遊ぶのが大好きでした。ある日、彼女は空に大きな赤いボールを見つけました。それは太陽でした!彼女はそれをとてもきれいだと思いました。
リリーはそのボールで遊びたいと思いました。

PC上で実行した時と若干文章が違いますが,ついにFPGA上でLLMを動作させることが出来ました!
速度は遅いですが,文章は一応意味が通っています.
※文章が最後ぶつ切りされているのは60トークンで切っているからです

ついにこれで完成になります.
実際に動いている動画がこちら.

llama2.cの改造について

MITライセンスなので一応貼っておきますね.

MIT License

Copyright (c) 2023 Andrej

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
1
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
1
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?