10月21日、AWSが「Beyond the KV cache: Training and serving i14's Laplace 3B linear attention model on AWS Trainium2 with native PyTorch」と題した記事を公開した。KVキャッシュを持たないLinear Attentionモデル「Laplace 3B」をAWS Trainium2上でトレーニング・サービングした実装レポートで、1,280並列シーケンス・シーケンスあたり29.5 MBの固定メモリという結果が示されている。
KVキャッシュの呪縛
LLMの推論コストを語るとき、KV（Key-Value）キャッシュは避けて通れない問題だ。Transformerのsoftmax attentionはトークンを全文脈と照合するため、会話や文書が長くなるにつれてKVキャッシュがHBM（高帯域メモリ）を圧迫し、並列処理できるシーケンス数が減り、コストが上がる。
これに対し、AI研究企業i14が設計したLaplace 3Bは、Gated Linear Attention（GLA）を採用することで、推論状態をシーケンスあたり29.5 MBの固定サイズに抑えた。文脈長が増えても状態は膨らまない。
Laplace 3Bの設計
Laplace 3Bは24層のデコーダーで、全層のsoftmax attentionをSimple GLAに置き換えている。GLAの核心は「過去のKVペアを固定サイズの行列状態に圧縮する」ことだ。
ここでいうSimple GLAは、元のGLAから入力依存のゲーティング機構を省略したシンプル版で、ゲート計算のオーバーヘッドがない分、Trainium2のような専用ハードウェアへの移植が容易になる。元のGLAはトークンごとにデータ依存の減衰率を計算するのに対し、Simple GLAではper-headの固定減衰率 λh のみを用いる。表現力とのトレードオフはあるが、実装のシンプルさと推論コストの予測可能性を優先した設計だ。
各ヘッド h は dk × dv = 96 × 192 の行列状態 S(h) を保持し、トークンごとに次の式で更新される：
$$S_t^{(h)} = \lambda_h S_{t-1}^{(h)} + k_t v_t^\top$$
λh は固定の減衰率（per-head decay）で、各ヘッドの「記憶の長さ」を決める。16ヘッドに対して幾何学的スケジュールで設定され、ヘッド0は半減期1トークン（λ=0.5）、ヘッド15は半減期2,048トークン（λ≈0.9997）と幅広い時間スケールをカバーする。
| パラメータ | 値 |
|---|---|
| 総パラメータ数 | 3.028B |
| 学習データ | 100B tokens（FineWeb-Edu） |
| シーケンス長 | 2,048 |
| オプティマイザ | Muon（2D重み行列）+ AdamW（埋め込み・正規化層） |
| FFN内部次元 | 8,192（GLU / SiLU） |
Softmax attentionとの比較をまとめると以下の通りだ：
| 比較軸 | Transformerのsoftmax attention | Laplace 3B（GLA） |
|---|---|---|
| 訓練時の計算量 | O(T²) | O(T) |
| 推論時のメモリ | 文脈長に比例して増加 | 固定（29.5 MB/シーケンス） |
| 長文コンテキスト | シーケンスごとにHBMを消費 | 文脈長が増えても状態量は変わらない |
ただし、GLAは履歴を圧縮するため、完全なキーとバリューを保持するわけではない。検索精度はワークロードごとに評価が必要だ。
Trainium2への移植：3つのエンジニアリング課題
もともとLaplace 3BはGPU向けのTritonカーネル（fused_chunk_simple_gla）で実装されていた。これをTrainium2で動かすため、AWS Generative AI Innovation Center（GenAIIC）チームはネイティブPyTorchバックエンドであるTorchNeuron Nativeを採用し、3つの課題を解決した。
課題1：GLAカーネルのネイティブPyTorch化
Trainium2の各NeuronCoreは28 MiBのSBUF（State Buffer）を128パーティションに分割して管理する。このSBUFはキャッシュではなく、コンパイラがHBMからのデータ転送を事前にスケジューリングするソフトウェア管理メモリだ。
GLAをチャンク分割方式で実装する際、チャンクサイズ（1チャンクあたりのトークン数）を128トークン・チャンク数16（C=128はチャンクサイズを指す）とすることで、128トークンのチャンクがSBUFの128パーティションに1対1でマッピングされる。この整合により、スループットが14%向上した。
課題2：FSDP2下でのMuon分散直交化
Laplace 3BはMuonオプティマイザを採用している。Muonは各勾配をNewton–Schulz反復（5回）で直交化してからパラメータを更新する手法で、2D重み行列2.718Bパラメータに適用される。
問題はFSDP2（Fully Sharded Data Parallel）との相性だ。64ランクで重みを分割すると、各ランクは行列の1/64しか持たないため、ランクローカルなGram行列 XX⊤ はブロック対角成分しか計算できない。
素朴な解法（全勾配をgather→1ランクでNewton–Schulz→broadcast）では、ステップごとに10.87 GBの通信量が発生し、他の63ランクが待機状態になる。
本番解法として実装したのがall_to_allを使った転置アプローチだ。1回目のall_to_allで各ランクが完全な行列を数枚受け取り、64ランクが並列にNewton–Schulzを実行、2回目のall_to_allで結果を返す。この方式でランクあたりの通信量は0.45 GBと24分の1に削減され、スループットは73,378→119,934トークン/秒（1.63倍）に向上した。損失曲線はビット単位で一致している。
課題3：逐次後退パスのアンロール
GLA再帰の後退パスは逐次計算からなる。GPUではautogradがランタイムにグラフを構築するが、Trainium2のコンパイラはこのステップを事前に静的展開（unroll）したグラフを要求する。チャンクサイズの選択はここでも効いてくる：
- チャンクサイズ64（チャンク数32）：グラフ深度がスケジューラの上限を超過
- チャンクサイズ256（チャンク数8）：中間データがHBMへのスピルを引き起こす
- チャンクサイズ128（チャンク数16）：スケジューラ制限内に収まり、中間データがオンチップ収容、SBUFパーティションとも整合
トレーニング・推論の結果
trn2.48xlarge（Trainium2チップ16個、論理NeuronCore 64基）1インスタンスで実行した結果：
- トレーニングスループット：コンパイルウォームアップ後、約120,000トークン/秒で安定推移
- ステップ84,210（約440億トークン）時点での損失：11.4 → 2.3729
- 勾配ノルムは初期スパイク後に約0.05へ収束。NaN/Inf値の発生なし
サービング面では、1,280並列シーケンスで11,544出力トークン/秒を記録した。固定サイズのリカレント状態により、シーケンス数を増やしてもメモリ使用量が比例して増えない特性が、高い並列度を実現している。
Linear Attentionはここ数年で研究が活発化しており、RWKVやMambaなど複数のアーキテクチャが提案されている。Laplace 3Bはその中でも専用ハードウェアへの最適化を含む実装詳細が公開されている点で、実務的な参考価値が高い。
詳細はBeyond the KV cache: Training and serving i14's Laplace 3B linear attention model on AWS Trainium2 with native PyTorchを参照していただきたい。
