第26章 MAX を導入する意味

26.1 この章で学ぶこと

この章で学ぶ内容は以下のとおりです。

表26-1: この章で学ぶこと
学ぶこと	ポイント
Mojo 版の限界	スカラーテープ方式では推論 1 ステップで数万ノードが積み上がる
MAX の役割	計算グラフをコンパイル・融合・最適化してハードウェアに送る実行エンジン
三つの中心概念	`Graph`（設計図）・`Module`（層の束）・`Weight`（学習される値）
MAX と Mojo の役割分担	Mojo = 演算カーネル記述言語、MAX = グラフ最適化と実行管理

26.2 microgpt.py と Mojo 版の限界

第25章 microgpt.py を Mojo で書き直すで実装した Mojo 版は、すべての演算をスカラー単位でテープに追記していきます。

t.add(a, b)   → テープにノードを 1 個追記
t.mul(c, d)   → テープにノードを 1 個追記
...           → 推論 1 ステップで数千〜数万ノードが積み上がる

これは自動微分（学習）のためには正しい設計ですが、推論だけを速く動かしたい場合には非効率です。

実装	特徴	問題
Python 版	短く読みやすい	Python インタプリタのオーバーヘッド
Mojo 版（Tape）	型安全・自動微分可能	スカラーループ、JIT/SIMD 最適化なし
MAX	グラフを先に構築してから実行	→ 次節で説明

26.3 MAX とは何か

MAX（Modular Accelerated Execution）は、モデルの計算をグラフとして表し、実行しやすい形へ整理する基盤です。

microgpt.py のような実行方式と MAX の実行方式を比べると次のようになります。

MAX では先に「何をどんな順で計算するか」というグラフを組み立て、その後でコンパイル・最適化してから実行します。このため：

演算の融合（複数の小さな演算を1回のカーネルにまとめる）
SIMD/GPU への自動マッピング
メモリレイアウトの最適化

が可能になります。

26.4 三つの中心概念

26.4.1 Graph：計算の設計図

Graph は「何を入力にし、どんな演算を経て、何を出力するか」を実行前に定義する設計図です。

Graph
├── 入力ノード（TensorType で型・形状を宣言）
├── 演算ノード（ops.matmul, ops.relu, ...）
└── 出力ノード

グラフを定義しただけでは何も実行されません。 InferenceSession に渡してコンパイルして初めて動きます。

26.4.2 Module：層や重みを持つまとまり

Module は、PyTorch の nn.Module に近い概念です。層、重み、計算の流れをひとまとまりで表します。

microgpt の例で言えば、次のような単位を Module として整理できます。

GPTModule
├── wte  (Weight: トークン埋め込み)
├── wpe  (Weight: 位置埋め込み)
├── AttentionModule (サブモジュール)
├── MLPModule       (サブモジュール)
└── lm_head (Weight: 出力線形層)

26.4.3 Weight：学習される値を明示する

Weight は、モデルの学習可能なパラメータを表します。入力テンソルと区別して明示することで、「どれが固定値でどれが更新される値か」がグラフ上で明確になります。

26.5 MAX と Mojo の役割分担

MAX と Mojo は競合ではなく、役割を分けて組み合わせる設計になっています。

役割	担当
モデル全体の計算グラフを組む	MAX（Python API）
標準にない演算を高速実装する	Mojo（custom ops）
学習（自動微分・Adam）	Mojo の Tape（または別途 MAX Engine）
推論の高速実行	MAX（InferenceSession）

典型的な進め方は「まず標準の MAX ops でグラフを組み、ボトルネックになった部分だけを Mojo の custom ops に置き換える」です。

26.6 まとめ

Python ループ / Mojo Tape のスカラー実行は学習には適しているが、推論の高速化には限界がある
MAX は計算グラフを先に構築し、コンパイル・最適化してから実行する
Graph（設計図）、Module（まとまり）、Weight（学習パラメータ）が三つの中心概念
MAX と Mojo は競合せず役割分担する

次章では MAX の Python API を使った具体的なコード例を示します。

26.7 次に読む章

第27章 MAX でコードを書くへ進みます。