Sony PSPでLLMを動かす

Sony PSP-2000は2007年製の333MHz MIPSハンドヘルドで、RAMは64MB。今週、私のPSPは1500万パラメータのTransformerを動かし、毎秒1〜2トークンの速度で英語テキストをLCDに出力している。

モデルはKarpathyのstories15M（TinyStoriesチェックポイント）で、int8量子化により約17MB。ランタイムはDockerでpspdev/pspdevを使ってクロスコンパイルされた純粋なCコード約1100行。PythonもなければLibtorchもなく、デバイス上に便利なランタイムも存在しない——PSPはメモリースティックから単一のEBOOT.PBPをロードするシングルプロセスボックスで、sceIo*、フレームバッファ、VFPUが与えられるだけだ。それ以外はすべて自分で構築する。

この記事は予算の話だ。各バイトの行き先、カーネルの実装、そして残されている課題について。

ハードウェア

「mmapなし」の行が最も痛い制約だ。Linuxボックスなら重みファイルをmmapしてページキャッシュに任せられる。PSPではsceIoLseek + sceIoReadと単一のmallocアリーナしかない。フォワードパス#1の前に17MB全体をRAMに読み込むか、USB 1.1フラッシュドライブ程度の速度でメモリースティックからストリーミングしてスループットの崩壊を見守るかだ。

PSP-1000の32MBでは、重みとKVキャッシュと作業バッファを合わせたヒープ領域が足りない。2000と3000は64MBを搭載している。64MBが必要だ。

モデル

stories15MはKarpathyのTinyStoriesチェックポイントの中で最小——Transformerレイヤー6、隠れ次元288、アテンションヘッド6、語彙32000。合計約1500万パラメータ。fp32で約57MB。int8 q80——対称的なグループ量子化、グループサイズ64、グループごとにfp32スケール1つ——で約17MB。

アーキテクチャ:  Llamaスタイルデコーダー、RoPE、SwiGLU FFN
レイヤー:       6
隠れ次元:       288
ヘッド:         6 (head_dim 48)
語彙:           32000
コンテキスト:   256 トークン
量子化:         int8 q80 (group=64、対称)
ディスクサイズ: 17 MB

モデルの準備は専用のDockerイメージで行う：python:3.11-slim + CPUのみのtorch + karpathy/llama2.cのピン留めコミット。stories15M.ptをダウンロードし、export.py --version 2を実行してq80のmodel.binを生成し、BPEのtokenizer.binをビルドする。そして——重要な点——-ffp-contract=off -fno-fast-mathでKarpathyのrunq.cリファレンスもビルドし、固定プロンプトで実行してtests/expected.txtを生成する。このファイルがPSPとのdiff対象となるバイト単位の正確なx86リファレンスだ。詳しくは後ほど。

メモリ予算

モジュールロード時に一度だけ宣言される24MBのヒープ：

PSP_HEAP_SIZE_KB(24576);

内訳：

特筆すべきトリック：トークン埋め込みテーブルはアリーナ内に量子化されたまま保持される。素朴な実装ではロード時に一度全体を逆量子化するが、それには約36MBかかり即座にOOMになる。代わりに、各フォワードでは現在のトークンに対応する1行だけを小さなfp32バッファに逆量子化する。コストはフォワードごとに1回余分な逆量子化が発生することだが、得られるものは持っていない36MBだ。

カーネル

transformer.cには定番の処理が揃っている：rmsnorm、softmax、quantize/dequantize、matmul、RoPE、アテンション、SwiGLU、サンプラー。各々がx86上のrunq.cと演算順序を一致させるために-ffp-contract=offを強制したテキストブック実装だ。これはテスト検証（後述）において重要な意味を持つ。

現在のmatmulはスカラーfp32——三重のネストループで、一度に1つのfp32積和演算。実機で約1〜2 tok/sを達成する。64トークンの補完に約1分かかる遅さだ。

matmulは交換可能な関数ポインタとしても設計されている。v1計画では4×4ベクター演算を使うVFPUカーネルを実装し、同じハードウェアで約5〜15 tok/sを目指す。（VFPUはPSPハードウェアの中で最も価値が持続するもの——128のレジスタを8つの4×4行列としてアドレス指定でき、単一命令で4×4行列乗算を実行できるベクターコプロセッサだ。）ドロップインするには1ファイルの変更で済む。

UI

PSPにはsceUtilityOsk*で起動するシステムのオンスクリーンキーボードがある。テキストはUTF-16LEで返され、UTF-8に変換する（BMPのみ——PSPのOSKはサロゲートペアに対応していない）。それをBPEトークナイザーに渡す。

チャットUIはpspDebugScreen——PSPのフレームバッファ上の組み込みデバッグフォント。480×272ディスプレイ上でモノスペース8×8ピクセル、60列×34行。2色のレイアウト：上部にプロンプト、下部に生成されたトークンが1文字ずつストリーミング表示される。バッファが画面下端に達するとレンダリングが折り返す。見た目は地味だが読みやすく、画面上のすべての文字がモデルが実際に出力したものだ。

デモ

プロンプト：

Once upon a time, there was a little girl named Layla

生成結果（T=0、64トークン）：

She was three years old and loved to explore. One day, she decided to go on an adventure. She put on her shoes and grabbed her bag. Layla walked outside and saw a big, tall tree.

この出力は、同じモデル、同じプロンプト、同じ温度、一致したFPフラグでx86_64上のrunq.cが生成するものとバイト単位で完全に一致している。この等価性がテスト検証の根拠だ——diff -q state.txt tests/expected.txtがクリーンを返すか、推論エンジンの各レイヤーが誤っているかのどちらかだ。同じリポジトリで以前のPongビルドが採用していたOCRベースのテストよりはるかに堅牢だ。PSPの画面は今や装飾的なものに過ぎず、真実はエミュレートされたメモリースティック上のテキストファイルにある。

実機へのサイドロード

上記のすべてはテストループでPPSSPP上で動作する。実機に載せるには：

PSP/
└── GAME/
    └── PspLlm/
        ├── EBOOT.PBP
        ├── model.bin
        └── tokenizer.bin

3つのファイルをメモリースティックに置き、XMBでゲームに移動してXボタンを押す。PSP-2000またはPSP-3000のみ対応——PSP-1000の32MBではヒープが足りない——PSPには署名されていないEBOOT.PBPを実行するカスタムファームウェアが必要だ（PSPなら6.61 PRO-C2または6.61 Infinity、PS VitaならAdrenaline）。コールドスタートからOSKまで約3秒。初回トークンのレイテンシはほぼプロンプトの長さに依存し、それ以降のトークンごとの時間はmatmulによる。

今後の課題

制約は計算能力ではない——VFPUアップグレードで1桁分の改善が見込める。制約はRAMだ。64MBが予算であり、OS、ヒープ、KVキャッシュ、作業バッファに支払い終えると、17MBのモデルには正確にぴったりの空間しか残らず、1バイトの余裕もない。SonyはMP3を再生しWipeoutを動かすためにこのハードウェアを出荷した。今週使った他のすべてのLLMはPSPより高価なGPUで動いていた。これはPSPで動く。