Elixir 製 LLM 推論エンジン Llamex を Livebook で動かす

はじめに

Llamex は、 @piacerex さんが作っている
Elixir で書かれた LLM 推論エンジン です

LLM の API クライアントではありません
llama.cpp のような推論パイプラインそのものを Elixir で最小構成に書いたものです

構成要素	内容
モデル形式	JSON / GGUF（llama / mistral / gemma3）
トークナイザー	whitespace / BPE
位置情報	RoPE（回転位置エンコーディング）
正規化	RMS 正規化
フィードフォワード	SwiGLU
キャッシュ	KV キャッシュ（スライディングウィンドウ対応）
サンプラー	greedy / temperature / top-k / top-p / min-p
バックエンド	純 Elixir（List）/ Nx / EXLA / FPGA

この記事では、Llamex を Livebook 上で動かしながら、その内部を表とグラフで可視化していきます

前半は手作りの極小モデルを Llamex のデータ構造で組み立てて動かします
そうすることで、エンジンの内部を数値で追えます

最後に実際の学習済みモデル（GGUF）をダウンロードして動かします

対象読者は次のような方です

• Elixir はある程度書けるが、LLM の推論がどう実装されているのかは知らない
• llama.cpp の中で何が起きているのかを、読めるコードで確かめたい

ノートブックは GitHub で公開しています

事前準備

Livebook のインストール方法は、公式サイトを参照してください
デスクトップアプリ・Docker・Fly.io など複数の方法があります

以下のバッジから、ノートブックを直接 Livebook で開けます

実行前の注意

Llamex は Elixir 1.19 以上を要求します（mix.exs の elixir: "~> 1.19"）
お使いの Livebook が古い Elixir で動いている場合、Mix.install が失敗します

System.version() で確認できます
1.19 未満の場合は Livebook 本体を更新してください

また Llamex は Hex には公開されていないため、GitHub から直接取得します
初回は clone とコンパイルで少し時間がかかります

Mix.install([
  {:llamex, github: "piacerex/llamex"},
  {:nx, "~> 0.12.1"},
  {:kino, "~> 0.15"},
  {:kino_vega_lite, "~> 0.1"},
  # 最後の節で GGUF モデルをダウンロードするために使う
  {:req, "~> 0.5"}
])

Llamex の推論パイプライン

LLM の文章生成は、次の1トークンを予測して末尾に足す の繰り返しです
Llamex もその流れをそのまま実装しています

ロジット とは、確率へ変換する前の「次トークン候補ごとの生の点数」です
この点数の並びから1つを選ぶのがサンプラーの役目になります

Llamex を構成する主なモジュールは次の通りです

モジュール	役割
Llamex	公開 API（new_model / generate / stream / eval など）
Llamex.Config	語彙数・埋め込み次元・ヘッド数などモデルの設計値
Llamex.Model	埋め込み表・層・出力層・トークナイザーをまとめた入れ物
Llamex.Context	推論中の状態（読んだトークン列と KV キャッシュ）
Llamex.Engine	1トークン分の順伝播
Llamex.Sampler	ロジットから次トークンを選ぶ戦略
Llamex.KVCache	過去の Key / Value の保存
Llamex.Backend.*	数値計算の実体（List / Nx / NxEXLA / FPGA）

1. 語彙とトークナイザーを作る

モデルは文字列をそのまま扱えません
まず文章をトークンという単位へ切り分け、それぞれに整数の ID を振ります

Llamex のトークナイザーは 語彙マップ と 未知語トークン から作れます

vocab_tokens = [
  "<unk>",
  "<eos>",
  "ねこ",
  "いぬ",
  "は",
  "ひるね",
  "さかな",
  "さんぽ",
  "が",
  "すき"
]

# トークン文字列 -> ID の対応表
vocab = vocab_tokens |> Enum.with_index() |> Map.new()
vocab_size = length(vocab_tokens)

# 空白区切りのトークナイザー
tokenizer = Llamex.Tokenizer.whitespace(vocab, "<unk>")

vocab
|> Enum.map(fn {token, id} -> %{トークン番号: id, トークン: token} end)
|> Enum.sort_by(& &1.トークン番号)
|> Kino.DataTable.new(keys: [:トークン番号, :トークン])

<unk> は語彙にない語を受け止める未知語トークン、<eos> は文の終わりを表すトークンです

まずは単体で encode / decode の往復を見ます

sample_text = "ねこ は ひるね"

encoded = Llamex.Tokenizer.encode(tokenizer, sample_text)
decoded = Llamex.Tokenizer.decode(tokenizer, encoded)

実行結果

段階	内容
入力	ねこ は ひるね
encode	[2, 4, 5]
decode	ねこ は ひるね

語彙にない語は未知語トークンへ落ちます

["ねこ", "うさぎ", "は"]
|> Enum.map(fn word ->
  ids = Llamex.Tokenizer.encode(tokenizer, word)

  %{
    入力: word,
    トークン番号: inspect(ids),
    復元: Llamex.Tokenizer.decode(tokenizer, ids)
  }
end)

実行結果

入力	トークン番号	復元
ねこ	[2]	ねこ
うさぎ	[0]	<unk>
は	[4]	は

うさぎ は語彙にないので <unk> になりました

2. 最小モデル: 埋め込みだけのモデル

Llamex のモデルは Llamex.new_model/1 で組み立てます
最低限必要なのは config（語彙数と埋め込み次元）と token_embeddings（トークンIDごとのベクトル）だけです

ここでは中身を表として読めるように、埋め込みを one-hot（自分の位置だけが 1 で、残りが 0 のベクトル）にします
後で「出力層＝遷移表」がそのまま読めるようになります

# トークンIDと同じ位置だけが1になる埋め込み
one_hot_embeddings =
  0..(vocab_size - 1)
  |> Map.new(fn id ->
    {id, for(j <- 0..(vocab_size - 1), do: if(j == id, do: 1.0, else: 0.0))}
  end)

embeddings_only_model =
  Llamex.new_model(%{
    config: %{vocab_size: vocab_size, embedding_size: vocab_size},
    token_embeddings: one_hot_embeddings,
    tokenizer: tokenizer
  })

layers は空、output は nil です
この状態で Llamex は何を「次トークンのスコア」として使うのでしょうか

lib/llamex/engine.ex を読むと、出力層がない場合は 現在のベクトルと各候補トークンの埋め込みとの内積をロジットにしています
内積は2つのベクトルの各要素を掛けて足した値で、向きが近いほど大きくなるため、「相性の良さ」の指標として使えます

Llamex.new_context/2 で推論状態を作り、Llamex.eval/2 に1トークン渡すとロジットが返ります

# 1トークンを入力したときのロジット（＝各候補との内積）を取り出す
logits_for = fn model, token_id ->
  context = Llamex.new_context(model, Llamex.Backend.List)
  {_context, logits} = Llamex.eval(context, token_id)
  Llamex.Backend.List.to_list(logits)
end

logits_for.(embeddings_only_model, vocab["ねこ"])

実行結果: [0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]

ねこ の位置（ID 2）だけが 1.0 です
全トークンについて同じことをして、ヒートマップにしてみます

対角線しか光りません
one-hot の埋め込みは自分自身としか重ならないため、**どの2語を比べても「まったく似ていない」**ことになるからです

この状態で生成するとどうなるでしょうか

Llamex.generate(embeddings_only_model, "ねこ", %{
  backend: Llamex.Backend.List,
  max_new_tokens: 5,
  sampler: :greedy
})

実行結果

項目	値
プロンプト	ねこ
生成結果	ねこ ねこ ねこ ねこ ねこ
終了理由	:length

自分自身との内積がいちばん大きいので、同じ語を繰り返すだけになります

3. 出力層を足す: 遷移表を持ったモデル

実際の言語モデルは、最後に「隠れ状態 → 語彙のスコア」へ変換する出力層を持ちます
Llamex では output: %{weight: 行列} で渡します

いま埋め込みは one-hot なので、出力層の重み W は次のように読めます

ロジット[次トークン j] = W[j] ・ one_hot(現在トークン i) = W[j][i]

つまり W[j][i] は「現在が i のとき、次に j を出す点数」です
直前の1語だけから次を決めるこの形は bigram モデルと呼ばれ、言語モデルのいちばん単純な形にあたります

# 「現在トークン -> 次トークン」に点数を付ける
transitions = %{
  {"ねこ", "は"} => 4.0,
  {"いぬ", "は"} => 4.0,
  {"は", "ひるね"} => 3.0,
  {"は", "さかな"} => 2.4,
  {"は", "さんぽ"} => 2.0,
  {"ひるね", "が"} => 4.0,
  {"さかな", "が"} => 4.0,
  {"さんぽ", "が"} => 4.0,
  {"が", "すき"} => 4.0,
  {"すき", "<eos>"} => 4.0
}

# 出力層の重み: 行 = 次トークン、列 = 現在トークン
output_weight =
  for next_token <- vocab_tokens do
    for current_token <- vocab_tokens do
      Map.get(transitions, {current_token, next_token}, 0.0)
    end
  end

transition_rows =
  for {current, current_id} <- Enum.with_index(vocab_tokens),
      {next, next_id} <- Enum.with_index(vocab_tokens) do
    %{
      現在のトークン: "#{current_id}:#{current}",
      次トークンの候補: "#{next_id}:#{next}",
      点数: Map.get(transitions, {current, next}, 0.0)
    }
  end

LlamexTour.Visuals.heatmap(
  transition_rows,
  "出力層の重み（＝遷移表）",
  :次トークンの候補,
  :現在のトークン,
  :点数
)

は の行だけ、ひるね / さかな / さんぽ の3か所に点数が分かれていることに注目してください
ここが後でサンプラーの実験に効いてきます

bigram_model =
  Llamex.new_model(%{
    config: %{vocab_size: vocab_size, embedding_size: vocab_size},
    token_embeddings: one_hot_embeddings,
    tokenizer: tokenizer,
    output: %{weight: output_weight}
  })

# 貪欲法（毎回1位を選ぶ）で生成する
generate_greedy = fn model, prompt ->
  Llamex.generate(model, prompt, %{
    backend: Llamex.Backend.List,
    max_new_tokens: 8,
    sampler: :greedy,
    stop_token: vocab["<eos>"]
  })
end

実行結果

プロンプト	生成結果	終了理由
ねこ	は ひるね が すき <eos>	:stop
いぬ	は ひるね が すき <eos>	:stop

文が最後まで生成され、<eos> に達したので finish_reason が :stop になりました
stop_token を渡さない場合は max_new_tokens まで生成し続け、:length で終わります

なお ねこ から始めても いぬ から始めても、続きは同じになります
このモデルは直前の1語しか見ないので、は の前に何があったかを予測へ活かせないためです

生成結果には、プロンプト側と生成側のトークンが分解されて入っています

detail = generate_greedy.(bigram_model, "ねこ")

実行結果

項目	値
prompt_tokens	[2]
prompt_pieces	["ねこ"]
generated_tokens	[4, 5, 8, 9, 1]
generated_pieces	["は", "ひるね", "が", "すき", "<eos>"]
finish_reason	:stop

4. サンプラーを可視化する

ここが Llamex を触っていて一番おもしろいところです
Llamex.Sampler.candidates/4 を使うと、ロジットが確率分布へ変わる過程をそのまま取り出せます

分岐のある は のロジットを材料にします

branch_logits = logits_for.(bigram_model, vocab["は"])

実行結果: [0.0, 0.0, 0.0, 0.0, 0.0, 3.0, 2.4, 2.0, 0.0, 0.0]

ひるね が 3.0、さかな が 2.4、さんぽ が 2.0 です

temperature: 分布の尖り方を変える

temperature は、ロジットを確率へ変換する前に割る値です
小さくすると差が強調され、大きくすると差がならされます

temperature_rows =
  [0.5, 1.0, 2.0]
  |> Enum.flat_map(fn temperature ->
    branch_logits
    |> Llamex.Sampler.candidates(Llamex.Backend.List, %{temperature: temperature}, vocab_size)
    |> Enum.map(fn candidate ->
      %{
        トークン: Enum.at(vocab_tokens, candidate.token),
        確率: candidate.probability,
        temperature: "temperature=#{temperature}"
      }
    end)
  end)

LlamexTour.Visuals.bar_chart(
  temperature_rows,
  "temperature による確率分布の変化",
  :トークン,
  :確率,
  color_field: :temperature,
  width: 640
)

実行結果（上位3候補の確率）

temperature	ひるね	さかな	さんぽ
0.5	0.688	0.207	0.093
1.0	0.442	0.242	0.162
2.0	0.256	0.190	0.155

temperature が小さいほど1位へ確率が集中し、大きいほど候補全体へ広がります
生成文の「堅さ」と「多様さ」を調整するつまみが、そのまま数値として見えています

top-k / top-p: 候補そのものを絞る

temperature は分布の形を変えるだけで候補は減りません
top_k（上位 k 個だけ残す）と top_p（確率の合計が p に達するまで残す）は候補そのものを切り落とします

count_rows =
  [
    {"制限なし", %{temperature: 1.0}},
    {"top_k=2", %{temperature: 1.0, top_k: 2}},
    {"top_p=0.6", %{temperature: 1.0, top_p: 0.6}},
    {"min_p=0.5", %{temperature: 1.0, min_p: 0.5}}
  ]
  |> Enum.map(fn {label, opts} ->
    candidates =
      Llamex.Sampler.candidates(branch_logits, Llamex.Backend.List, opts, vocab_size)

    %{
      設定: label,
      残った候補数: length(candidates),
      候補: candidates |> Enum.map(&Enum.at(vocab_tokens, &1.token)) |> Enum.join(" / ")
    }
  end)

実行結果

設定	残った候補数	候補
制限なし	10	ひるね / さかな / さんぽ / ほか7件
top_k=2	2	ひるね / さかな
top_p=0.6	2	ひるね / さかな
min_p=0.5	2	ひるね / さかな

min_p は「1位の確率の何割以上か」で足切りする方式です

切り落とした後は残った候補だけで確率を計算し直すため、合計はきちんと 1 になります

実際に抽選して生成する

サンプラーに seed を渡すと、乱数を固定して再現できます
ねこ は まで与えて、その先の分岐を何度も引いてみます

1..6
|> Enum.map(fn seed ->
  result =
    Llamex.generate(bigram_model, "ねこ は", %{
      backend: Llamex.Backend.List,
      max_new_tokens: 8,
      sampler: %{temperature: 1.0, top_k: 3, seed: seed},
      stop_token: vocab["<eos>"]
    })

  %{seed: seed, 生成結果: result.text}
end)

実行結果

seed	生成結果
greedy（参考）	ひるね が すき <eos>
1	さかな が すき <eos>
2	さんぽ が すき <eos>
3	さかな が すき <eos>
4	さんぽ が すき <eos>
5	ひるね が すき <eos>
6	ひるね が すき <eos>

greedy は必ず点数1位の ひるね を選びますが、サンプリングでは さかな や さんぽ も選ばれます
「同じ入力でも実行のたびに違う文章が出る」という LLM の挙動は、この抽選から生まれています

5. Llama 系の部品を覗く

Transformer の解説でよく見る GPT 系に対し、Llamex が実装しているのは Llama 系です
3つの部品が違います

部品	GPT 系（GPT-2 など）	Llama 系（Llamex）
正規化	レイヤー正規化	RMS 正規化
位置情報	sin / cos を埋め込みへ足す	RoPE で回転させる
フィードフォワード	ReLU	SwiGLU（SiLU でゲート）

バックエンド経由で個別に呼べるので、1つずつ動かします

RMS 正規化: 平均を引かない正規化

sample_vector = [1.0, 0.5, -0.5, 2.0]
norm_weight = [1.0, 1.0, 1.0, 1.0]

rms_normalized = Llamex.Backend.List.rms_norm(sample_vector, norm_weight, 1.0e-6)

実行結果

段階	値	平均	二乗平均平方根
正規化前	[1.0, 0.5, -0.5, 2.0]	0.75	1.1726
RMS正規化後	[0.8528, 0.4264, -0.4264, 1.7056]	0.6396	1.0

レイヤー正規化は「平均を引いて分散で割る」処理ですが、RMS 正規化は 平均を引きません
表を見ると、正規化後も平均は 0 になっていない一方、二乗平均平方根はきっちり 1.0 になります

引き算を1つ省くぶん計算が軽く、実用上の精度も落ちにくいことから、Llama 系ではこちらが使われます

RoPE: 位置を「回転」として入れる

アテンションの計算は、そのままでは「どのトークンが何番目か」を持ちません
GPT 系では sin / cos の波を埋め込みに足して位置を伝えます
Llama 系の RoPE は、ベクトルを位置に応じて 回転 させます

# Llamexの実装では、前半と後半が対になって回転する
# 4次元なら (0,2) と (1,3) がペア
rope_base = [1.0, 0.0, 0.0, 0.0]

rope_rows =
  0..11
  |> Enum.map(fn position ->
    rotated = Llamex.Backend.List.rope(rope_base, position, 10_000.0, nil)

    %{
      位置: position,
      x: Float.round(Enum.at(rotated, 0), 4),
      y: Float.round(Enum.at(rotated, 2), 4)
    }
  end)

# 原点から各位置のベクトルへ線を引き、12本を重ねて表示する
LlamexTour.Visuals.vector_fan(
  rope_rows,
  "RoPE: 位置ごとのベクトル（原点からの矢印を重ねたもの）",
  :x,
  :y,
  :位置
)

同じベクトルが、位置が進むにつれて円周上を回っていきます
長さは変わらず、向きだけが変わります

RoPE の重要な性質は、Query と Key の内積が「絶対位置」ではなく「相対距離」だけで決まることです
アテンションは Query と Key の内積で「どこを見るか」を決めるので、この性質はそのまま「位置関係の捉え方」に効きます

実験の組み立てはこうです
Key を置く位置（基準位置）を 0 から 4 まで変えながら、Query をそこから 相対距離 だけ右に置いて内積を測ります
もし内積が絶対位置に依存するなら、基準位置ごとに違うグラフになるはずです

rope = fn vector, position ->
  Llamex.Backend.List.rope(vector, position, 10_000.0, nil)
end

query_vector = [0.9, 0.2, -0.4, 0.7]
key_vector = [0.3, -0.6, 0.8, 0.1]

base_positions = Enum.to_list(0..4)
distances = Enum.to_list(0..4)

# Key を base_position に、Query をそこから distance だけ右に置いて内積を測る
dot_at = fn base_position, distance ->
  q = rope.(query_vector, base_position + distance)
  k = rope.(key_vector, base_position)

  Float.round(Llamex.Backend.List.dot(q, k), 6)
end

# 基準位置ごとに1枚ずつグラフを作り、タブで切り替えられるようにする
rope_tabs =
  Enum.map(base_positions, fn base_position ->
    rows =
      Enum.map(distances, fn distance ->
        %{相対距離: distance, 内積: dot_at.(base_position, distance)}
      end)

    chart =
      LlamexTour.Visuals.line_chart(
        rows,
        "Key を位置 #{base_position} に置いたときの内積",
        :相対距離,
        :内積,
        # 縦軸を固定しないと、タブごとに軸が伸縮して比較できなくなる
        y_domain: [-1.0, 1.0],
        width: 560
      )

    {"基準位置=#{base_position}", chart}
  end)

Kino.Layout.tabs(rope_tabs)

タブを切り替えてみてください
折れ線がまったく動きません
Key を系列のどこに置いても、相対距離が同じなら内積は同じ値になるということです

グラフだけでは「本当に同じ値か」が分かりにくいので、数値でも確かめます

# 相対距離ごとに、5つの基準位置で得られた内積を集めて重複を除く
relative_check_rows =
  Enum.map(distances, fn distance ->
    values = Enum.map(base_positions, &dot_at.(&1, distance))
    unique_values = Enum.uniq(values)

    %{
      相対距離: distance,
      異なる値の数: length(unique_values),
      内積: unique_values |> Enum.map(&to_string/1) |> Enum.join(" / ")
    }
  end)

実行結果

相対距離	異なる値の数	内積
0	1	-0.1
1	1	0.634223
2	1	0.743427
3	1	0.131261
4	1	-0.635397

異なる値の数 がすべて 1 になりました
5つの基準位置から計算した内積が、丸め誤差もなく完全に一致しているということです

この性質のおかげで、学習時より長い系列でも位置の扱いが破綻しにくくなります
位置を足し込む方式では、学習時に見たことのない大きな位置の値がそのまま入力へ混ざりますが、RoPE では回転角が変わるだけで、内積は常に相対距離の関数のままだからです

SwiGLU: ゲート付きフィードフォワード

GPT 系のフィードフォワード層は ReLU を使いますが、Llama 系は SiLU（Swish）でゲートをかけます
Llamex.Backend.List.silu_multiply/2 は SiLU(gate) * up を計算します

# up をすべて1にすると、SiLU そのものの形が見える
xs = for i <- -60..60, do: i / 10.0
ones = for _ <- xs, do: 1.0

silu_values = Llamex.Backend.List.silu_multiply(xs, ones)

ReLU は 0 で折れ曲がり、負の入力を完全に捨てます
SiLU はなめらかに曲がり、負の側もわずかに値を残します
この「なめらかさ」が学習の安定に効くとされています

# SwiGLU は gate 側で up 側の通り具合を調節する
gate_values = [-2.0, -1.0, 0.0, 1.0, 2.0]
up_values = [1.0, 2.0, 3.0, 4.0, 5.0]

gated = Llamex.Backend.List.silu_multiply(gate_values, up_values)

実行結果

gate	up	出力
-2.0	1.0	-0.2384
-1.0	2.0	-0.5379
0.0	3.0	0.0
1.0	4.0	2.9242
2.0	5.0	8.808

gate が負のときは出力がほぼ止まり、正のときは up の値がそのまま通ります
文字通り「ゲート（門）」として働いていることが分かります

6. Transformer ブロックと KV キャッシュ

ここまでのモデルには層がありませんでした
layers を渡すと、Llamex は各ブロックで「RMS 正規化 → 因果マスク付き自己アテンション → 残差 → RMS 正規化 → SwiGLU → 残差」を実行します

ここで作るブロックは、アテンションとフィードフォワードの寄与が 0 になるよう wo と w_down をゼロ行列にしています
手作りの重みでは意味のある文脈処理は起きないため、遷移表の結果を保ったまま、KV キャッシュの動きだけを観察する狙いです

build_block = fn sliding_window ->
  block = %{
    head_count: 2,
    kv_head_count: 2,
    attention_norm: ones_vector.(vocab_size),
    feed_forward_norm: ones_vector.(vocab_size),
    wq: identity_matrix.(vocab_size),
    wk: identity_matrix.(vocab_size),
    wv: identity_matrix.(vocab_size),
    # 出力側をゼロにして、アテンションの寄与を打ち消す
    wo: zero_matrix.(vocab_size, vocab_size),
    w_gate: identity_matrix.(vocab_size),
    w_up: identity_matrix.(vocab_size),
    w_down: zero_matrix.(vocab_size, vocab_size)
  }

  if sliding_window do
    Map.put(block, :sliding_window, sliding_window)
  else
    block
  end
end

KV キャッシュはどう伸びるか

自己回帰生成では、過去のトークンの Key / Value を毎回計算し直すのは無駄です
そこで KV キャッシュに貯めます

Llamex.KVCache.entry_count/1 で貯まった数を数えられます

# 1トークンずつ入力し、そのたびにKVキャッシュの大きさを測る
measure_kv_cache = fn model, steps ->
  context = Llamex.new_context(model, Llamex.Backend.List)

  {_context, rows} =
    Enum.reduce(1..steps, {context, []}, fn step, {context, acc} ->
      token_id = rem(step, vocab_size)
      {context, _logits} = Llamex.eval(context, token_id)

      row = %{
        入力したトークン数: step,
        キャッシュのエントリ数: Llamex.KVCache.entry_count(context.kv_cache)
      }

      {context, acc ++ [row]}
    end)

  rows
end

実行結果

入力したトークン数	制限なし	sliding_window=3
1	2	2
2	4	4
3	6	6
4	8	6
5	10	6
6	12	6
7	14	6
8	16	6

制限なし では、トークンを1つ読むたびにキャッシュが増え続けます（2ブロックあるので1トークンにつき2エントリ）
長い文章を生成するほどメモリを食う、というのがそのまま形に出ています

sliding_window=3 を指定したブロックでは、直近3件だけを残して古いものを捨てるため、途中から頭打ちになります
これは Mistral などが採用している仕組みで、Llamex では層ごとの設定として持っています

7. バックエンドを切り替える

Llamex は数値計算の実体を差し替えられます

バックエンド	実体	想定用途
Llamex.Backend.List	純 Elixir のリスト演算	AtomVM など制約の多い環境
Llamex.Backend.Nx	Nx テンソル	標準的な BEAM 環境
Llamex.Backend.NxEXLA	XLA コンパイル	大きなモデルでの高速化
Llamex.Backend.FPGA	FPGA 境界（未対応時は List へ）	ハードウェア実験

同じモデル・同じプロンプトで速度を比べます

prepared_model = Llamex.prepare_model(layered_model, Llamex.Backend.Nx)

backend_rows = [
  %{構成: "Backend.List", 平均ミリ秒: benchmark.(layered_model, Llamex.Backend.List, 5)},
  %{構成: "Backend.Nx", 平均ミリ秒: benchmark.(layered_model, Llamex.Backend.Nx, 5)},
  %{構成: "Backend.Nx + prepare_model", 平均ミリ秒: benchmark.(prepared_model, nil, 5)}
]

実行結果（筆者環境での一例。実行ごとにばらつきます）

構成	平均ミリ秒
Backend.List	1.95
Backend.Nx	5.53
Backend.Nx + prepare_model	5.62

結果を見て意外に思うかもしれませんが、この極小モデルでは Backend.List がいちばん速くなります
語彙10・埋め込み10次元では、Nx のテンソル1回ごとのオーバーヘッドのほうが計算そのものより大きいためです

prepare_model/2 は、モデルの重みをバックエンドの形式へあらかじめ変換しておく関数です
同じモデルへ何度も推論するときに毎回の変換を省けます
ただしこの規模では変換コスト自体が小さいため、Backend.Nx 単体との差は実行ごとのばらつきに埋もれます

つまり、バックエンドの選択と prepare_model が効いてくるのは、行列が十分に大きくなってからです

8. 1トークンずつ進める

Llamex.generate/3 は最後までまとめて生成しますが、Llamex.step/3 を使うと1トークンずつ手で進められます

# 「ねこ」から始めて、1トークンずつ進める
initial_context = Llamex.new_context(bigram_model, Llamex.Backend.List)

{_context, _token, step_rows} =
  Enum.reduce(1..5, {initial_context, vocab["ねこ"], []}, fn step, {context, current_token, acc} ->
    result = Llamex.step(context, current_token, %{sampler: :greedy})

    row = %{
      ステップ: step,
      入力トークン: Enum.at(vocab_tokens, current_token),
      選ばれたトークン番号: result.token,
      選ばれたトークン: result.text
    }

    {result.context, result.token, acc ++ [row]}
  end)

実行結果

ステップ	入力トークン	選ばれたトークン番号	選ばれたトークン
1	ねこ	4	は
2	は	5	ひるね
3	ひるね	8	が
4	が	9	すき
5	すき	1	<eos>

前のステップで選ばれたトークンが、次のステップの入力になっています
これが自己回帰生成の実体です

Llamex.stream/3 を使うと、同じ処理を Stream として受け取れます
1トークン出るたびに画面へ流し込めます

stream_frame = Kino.Frame.new()
Kino.render(stream_frame)

prepared_model
|> Llamex.stream("ねこ", %{max_new_tokens: 8, sampler: :greedy, stop_token: vocab["<eos>"]})
|> Enum.reduce([], fn chunk, pieces ->
  # チャンクは1トークンずつ届くので、表示用に空白で連結する
  updated = if chunk.text == "", do: pieces, else: pieces ++ [chunk.text]

  Kino.Frame.render(
    stream_frame,
    Kino.Text.new("""
    生成中: #{Enum.join(updated, " ")}
    受け取ったトークン数: #{length(updated)}
    終了理由: #{inspect(chunk.finish_reason)}
    """)
  )

  Process.sleep(200)
  updated
end)

ストリームの最後のチャンクは token: nil、text: "" で、finish_reason だけを持ちます
生成が「なぜ終わったか」を受け取るためのものです

9. 実際のモデル（GGUF）を動かす

ここまでは手作りの極小モデルでした
Llamex は GGUF ファイル（llama.cpp などで使われる量子化済みモデル形式）を読めます

対応項目	内容
アーキテクチャ	llama / mistral / gemma3
量子化形式	F32 / F16 / BF16 / Q2_K〜Q8_K
トークナイザー	whitespace / BPE

使うのは stories260K.gguf です
TinyStories（子ども向けの短い物語）データセットで訓練された、パラメータ数 26 万・ファイルサイズ約 1.1MB の Llama モデルです

実用品ではありませんが、アーキテクチャは本物の Llama で、英語の物語らしい続きを生成します

# モデルの保存先。すでにダウンロード済みなら再取得しない
models_dir = Path.join(System.tmp_dir!(), "llamex_models")
File.mkdir_p!(models_dir)

model_path = Path.join(models_dir, "stories260K.gguf")

model_url =
  "https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories260K.gguf"

unless File.exists?(model_path) do
  Req.get!(model_url, into: File.stream!(model_path), redirect: true)
end

実行結果: ファイルサイズ 1185376 バイト（約 1.1MB）、ダウンロードは数秒で完了します

モデルを読む前に中身を調べる

Llamex は、読み込む前に GGUF の中身を診断できます
「そもそも Llamex で動かせるモデルか」を確かめるための機能です

summary = Llamex.GGUF.Diagnostic.inspect_summary_file(model_path)

# 診断結果は項目が多いので、要点だけ取り出す
diagnostic_rows =
  [
    :architecture,
    :loadable?,
    :compatibility_issues,
    :architecture_runtime_status,
    :tokenizer_model,
    :tokenizer_kind,
    :tokenizer_token_count,
    :rope_variant,
    :attention_variant,
    :chat_usable
  ]
  |> Enum.map(fn key ->
    %{項目: Atom.to_string(key), 値: inspect(Map.get(summary, key))}
  end)

実行結果

項目	値
architecture	"llama"
loadable?	true
compatibility_issues	[]
architecture_runtime_status	"supported"
tokenizer_model	"llama"
tokenizer_kind	"whitespace"
tokenizer_token_count	512
rope_variant	%{type: "default"}
attention_variant	%{type: "full"}
chat_usable	false

loadable? が true、compatibility_issues が空なら、そのまま読み込めます
対応していないアーキテクチャや量子化形式のモデルを渡すと、ここに理由が並びます

chat_usable が false なのは、このモデルがチャットテンプレートを持たない素の言語モデルだからです
true のモデルなら Llamex.generate_chat/3 でチャット形式のやり取りができます

設計値（config）だけを取り出すこともできます

config_map = Llamex.GGUF.ModelLoader.model_config_summary_file(model_path)

実行結果

設定	値
attention_head_count	8
attention_head_count_kv	4
block_count	5
context_size	2048
embedding_size	64
feed_forward_size	172
rope_dimension_count	8
rope_theta	10000.0
vocab_size	512

attention_head_count が 8 なのに attention_head_count_kv は 4 です
これは GQA（Grouped Query Attention） といって、Key と Value のヘッド数を Query より減らして KV キャッシュを節約する仕組みです

第6節で kv_head_count を指定できたのは、この構成に対応するためでした

読み込んで生成する

{load_microseconds, gguf_model} =
  :timer.tc(fn -> Llamex.GGUF.ModelLoader.load(model_path) end)

実行結果: 読み込み時間 約 0.25 秒

本物のサブワード分割を見る

第1節では空白区切りのトークナイザーを自作しましたが、このモデルは SentencePiece 形式の語彙を持っています

gguf_pieces = fn text ->
  ids = Llamex.encode(gguf_model, text)
  pieces = Enum.map(ids, &Map.fetch!(gguf_model.tokenizer.id_to_token, &1))

  %{
    入力: text,
    分割数: length(ids),
    トークン: Enum.join(pieces, " | "),
    復元: Llamex.decode(gguf_model, ids)
  }
end

["Once upon a time", "dragon", "happily", "unbelievable"]
|> Enum.map(gguf_pieces)

実行結果

入力	分割数	トークン	復元
Once upon a time	4	▁Once ▁upon ▁a ▁time	Once upon a time
dragon	5	▁d r a g on	dragon
happily	2	▁happ ily	happily
unbelievable	11	▁u n b e l i e v a b le	unbelievable

▁（アンダースコアに似た記号）は単語の先頭に空白があったことを表します
Once のようによく出る語は1トークンですが、dragon や unbelievable は語彙にないため、モデルが知っている小さな部品へ分解されます

語彙がわずか 512 個しかないモデルなので、分割はかなり細かくなります

本物のモデルの「次トークン確率」を見る

第4節ではサンプラーを手作りの遷移表で試しましたが、同じことを本物のモデルのロジットに対して行えます

Llamex.prefill/3 はプロンプトを読み込んだ状態の Context を返します
そこから Llamex.eval/2 を呼べば、プロンプト直後のロジットが手に入ります

gguf_prompt = "Once upon a time"

prefilled = Llamex.prefill(gguf_model, gguf_prompt, %{backend: Llamex.Backend.List})
{_context, gguf_logits} = Llamex.eval(prefilled.context, prefilled.current_token)

gguf_candidates =
  Llamex.Sampler.candidates(gguf_logits, Llamex.Backend.List, %{temperature: 1.0}, 10)

gguf_candidate_rows =
  Enum.map(gguf_candidates, fn candidate ->
    %{
      次トークン: Map.fetch!(gguf_model.tokenizer.id_to_token, candidate.token),
      確率: candidate.probability
    }
  end)

実行結果

次トークン	確率
,	0.7437
▁she	0.1009
▁they	0.0462
▁he	0.0223
.	0.0169
▁the	0.0116
▁her	0.0055
▁time	0.0046
▁it	0.0045
▁there	0.0044

「昔むかしあるところに」の続きとして、カンマや she / they / he が上位に来ます
26 万パラメータのモデルでも、物語の書き出しらしい続きを確率として持っていることが分かります

temperature を実モデルのロジットへ適用すると、第4節で見たのと同じ変化が起きます

文章を生成する

gguf_backend_rows =
  [Llamex.Backend.List, Llamex.Backend.Nx]
  |> Enum.map(fn backend ->
    {microseconds, result} =
      :timer.tc(fn ->
        Llamex.generate(gguf_model, gguf_prompt, %{
          backend: backend,
          max_new_tokens: 10,
          sampler: :greedy
        })
      end)

    %{
      バックエンド: inspect(backend),
      秒: Float.round(microseconds / 1_000_000, 2),
      生成結果: result.text
    }
  end)

実行結果

バックエンド	秒	生成結果
Llamex.Backend.List	0.15	, there was a little gothip
Llamex.Backend.Nx	3.34	, there was a little gothip

どちらのバックエンドでも greedy なので生成結果は同じです
速度は第7節と同じ傾向で、この規模では純 Elixir の Backend.List が優位です

サンプリングにすると、同じプロンプトから毎回違う物語が始まります

1..3
|> Enum.map(fn seed ->
  result =
    Llamex.generate(gguf_model, gguf_prompt, %{
      backend: Llamex.Backend.List,
      max_new_tokens: 20,
      sampler: %{temperature: 0.9, top_k: 40, top_p: 0.9, seed: seed}
    })

  %{seed: seed, 生成結果: result.text}
end)

実行結果

seed	生成結果
1	, the dogars. The blt the birdd int the b
2	they were a time, it wilbold bry plannee
3	, Maleom ater cared!" said Lily. She cally

文法が崩れているのは、モデルが 26 万パラメータしかないためです
それでも「物語の続きを1トークンずつ選んでいる」という動きは実物と変わりません

ストリーミング表示

実モデルでもストリーミングできます
ただし1トークンずつ復号すると ▁（単語先頭の空白）の情報が落ちるため、そこまでの全トークンをまとめて復号します

gguf_frame = Kino.Frame.new()
Kino.render(gguf_frame)

gguf_model
|> Llamex.stream(gguf_prompt, %{
  backend: Llamex.Backend.List,
  max_new_tokens: 20,
  sampler: %{temperature: 0.9, top_k: 40, seed: 7}
})
|> Enum.each(fn chunk ->
  # チャンク単体ではなく、生成済みトークン列全体を復号する
  text = Llamex.decode(gguf_model, chunk.generated_tokens)

  Kino.Frame.render(
    gguf_frame,
    Kino.Text.new("""
    #{gguf_prompt}#{text}

    生成トークン数: #{length(chunk.generated_tokens)}
    終了理由: #{inspect(chunk.finish_reason)}
    """)
  )

  Process.sleep(100)
end)

実行結果（1トークンずつ伸びていきます）

もっと大きなモデルを動かすには

同じリポジトリには stories15M-q4_0.gguf（約 19MB、1500 万パラメータ、Q4_0 量子化）もあります
URL のファイル名を差し替えれば試せますが、時間の見積もりが大きく変わります

項目	stories260K	stories15M
ファイルサイズ	約 1.1MB	約 19MB
ダウンロード	数秒	数秒
診断（inspect_summary_file）	一瞬	一瞬。loadable? は true
読み込み（load）	約 0.25 秒	非常に長い

inspect_summary_file/1 の時点では 15M 版も loadable?: true と出ます
しかし load/1 は量子化されたテンソルを Elixir の数値リストへ展開するため、パラメータ数に比例して時間とメモリを消費します

筆者の環境では 5 分待っても読み込みが完了しませんでした

Llamex.GGUF.ModelLoader.load/2 には tensor_format: :compact オプションがあり、量子化されたまま保持することもできます
ただし対応する演算が限られるため、まずは小さいモデルで挙動を確かめるのが確実です

コマンドラインからは mix llamex.generate などのタスクも用意されています

まとめ

• Llamex は「LLM API のクライアント」ではなく、Elixir で書かれた推論エンジンそのもの
• 出力層がないモデルでは、ロジットは埋め込み同士の内積になる
• 出力層の重みは「現在トークン → 次トークン」の遷移表として読める
• Llamex.Sampler.candidates/4 を使うと、temperature / top_k / top_p / min_p が確率分布に与える影響を数値で見られる
• Llama 系の3つの部品は、GPT 系とここが違う
  • RMS 正規化: 平均を引かない。二乗平均平方根だけを 1 にそろえる
  • RoPE: 位置を足すのではなく回転で入れる。内積が相対距離だけで決まる
  • SwiGLU: ReLU の代わりに、なめらかな SiLU でゲートをかける
• KV キャッシュはトークンを読むたびに伸び、sliding_window を設定すると頭打ちになる
• バックエンドは差し替え可能。ただし小さいモデルでは純 Elixir の Backend.List が最速
• 学習済みの GGUF モデル（約 1.1MB）を読み込むと、手作りモデルで見てきた仕組みがそのまま実物で動く

Llamex はモジュールが小さく分かれているので、コードを読み進めやすい構成になっています

ファイル	見どころ
lib/llamex/engine.ex	1トークン分の順伝播。出力層がないときの内積フォールバックもここ
lib/llamex/sampler.ex	temperature / top-k / top-p / min-p / 繰り返しペナルティの実装
lib/llamex/layers/rope.ex	RoPE の回転が数十行で書かれている
lib/llamex/kv_cache.ex	キャッシュの追加とスライディングウィンドウでの切り詰め
lib/llamex/backend/list.ex	純 Elixir だけで書かれたテンソル演算の全体像

推論エンジンの中身を知りたいけれど C++ は読みたくない、という場合の入口としてちょうど良い規模だと感じました