RustによるHOF認知Flash Attention実装の手引き

Asher Bond

合成可能な HOF サービスとしての注意機構のメカナイズ ― Rust による実装

著者: Asher Bond (asher.bond@distillative.ai)

要旨

HTTP エンドポイントの背後に公開された注意計算は、単なるカーネルではなく、一個のシステムオブジェクトである。本稿は、スケール化ドット積注意 [1] を、Rust における合成可能な高階パイプライン ― 射影 → スコアリング → 集約 → 配信 ― としてメカナイズする。各段は、関数型プログラミングの伝統 [2, 3, 4] に連なる、小さく再利用可能な関数である。カーネルは厳密なスケール化ドット積注意 [1] を計算する。設計の要点は、その周囲に引かれた段の境界にある ― この境界は十分に明瞭であり、FlashAttention [5] のタイル化・IO 認識カーネルは、それを包むサービスに一切手を触れることなく、スコアリング段へそのまま差し替えられる。実装言語は Rust であり、アーキテクチャは Input–Process–Output の語彙で記述する。

1. はじめに

注意計算を HTTP エンドポイントの背後で利用可能にすることは、ありふれた、しかし基幹を担うエンジニアリングであり、Rust はそのための堅実な選択である。すなわち、型付けされたメモリ安全なカーネルを、ガベージコレクタの停止なしに配信する。本設計は、サービスを小さく再利用可能な段へと分解する ― 入力を query/key/value へ射影し、スコアリングし、正規化し、集約し、配信する ― その各々が、次段へと合成される関数である。カーネルは厳密なスケール化ドット積注意 [1] を実装する。FlashAttention [5] は、n × n のスコア行列を高帯域メモリから排除する、タイル化・IO 認識の厳密注意アルゴリズムである。本稿の合成は段の境界を精密に引いており、そのカーネルは周囲のサービスを乱すことなくスコアリング段へと収まる。この差し替え境界こそが、本稿の貢献である。

2. 関連研究

注意機構。 カーネルはスケール化ドット積注意 [1] を計算する ― Q/K/V へ射影し、Q·Kᵀ をスコアリングし、スケールし、softmax を取り、V に対して集約する。

FlashAttention。 IO 認識の厳密注意アルゴリズム [5] は、計算をタイル化し、n × n のスコア行列を高帯域メモリ上に実体化することを回避する。これは、スコアリング段が受け入れるべく作られたカーネルにほかならない。サービスは段のインタフェースを露出し、タイル化カーネルは、密なカーネルに代わってそれを満たす。

合成。 サービスを小さく再利用可能な段(射影 → スコアリング → 集約 → 配信)として構成することは、モジュール性についての関数合成的な論拠 [2, 3, 4] である。この合成こそが、カーネルを、融合したモノリスではなく差し替え可能な部品たらしめる。

3. 設計(IPO)

Input(入力)。 /process へ POST された浮動小数点の JSON 配列を、1 × d の行へと整形する。
Process(処理)。 入力を query/key/value 行列へ射影し、スケール化ドット積スコアを計算し、行方向に正規化し、V に対して集約し、結果を返す。
Output(出力)。 JSON 配列(注意出力)、および、インメモリのマップから既存の格納済み表現を返す /retrieve エンドポイント。

4. カーネル

厳密なスケール化ドット積注意 [1] は、exp(scores) を行方向の総和で正規化して行確率的な重み行列 A を得たのち、A · V を返す。

pub fn compute_attention(&self, x: &Array2<f32>) -> Array2<f32> {
 let q = x.dot(&self.query);
 let k = x.dot(&self.key);
 let v = x.dot(&self.value);

 let scores = q.dot(&k.t()) / (x.shape()[1] as f32).sqrt();
 let exp = scores.mapv(f32::exp);
 let row_sums = exp.sum_axis(Axis(1));
 let weights = &exp / &row_sums.insert_axis(Axis(1)); // row-stochastic A
 weights.dot(&v)
}

この密なカーネルは n × n のスコア行列を全面的に実体化する ― これは厳密注意 [1] の標準的なメモリプロファイルである。FlashAttention [5] は、その行列を高帯域メモリに一度も保持することなく、同一の厳密な結果を計算するタイル化カーネルである。両者がここで交換可能であるのは、§3 の段の境界ゆえである。サービスは特定のカーネルではなく注意出力を消費する。したがってメモリプロファイルは、スコアリング段に置かれたカーネルの性質であり、その基準として [5] が参照される。

5. 配信する

プロジェクト構成は、カーネルを呼び出す HOF ラッパ、表現を JSON へ格納・読み出しする DistributedRepresenter、そして /process と /retrieve を露出する hyper ベースのサーバからなる。その各々が一つの段である。すなわち attention(§4 のカーネル)、hof_cognition(合成を担うラッパ)、distributed_representation(格納庫)、server(トランスポート)である。小さなモジュールへの合成は、カーネルを差し替え可能な単位に保つ、関数型的な構成の選択 [2, 3] である。[ILLUSTRATIVE] な結線 ― この例は 127.0.0.1:3000 上で動作し、curl -X POST .../process -d '[0.5,0.1,0.4,0.8]' がカーネルを起動し、curl .../retrieve が格納済み表現を返す。

根拠と適用範囲

本稿の貢献は、合成と差し替え境界であり、スループットの数値ではない。§4 のカーネルは厳密なスケール化ドット積注意 [1] であり、スコア行列を全面的に実体化する。FlashAttention [5] は、その代わりにスコアリング段が受け入れるべく設計されたタイル化カーネルであり、メモリプロファイルを測る基準でもある。配信された例は、例示的な結線 ― 単一行の入力、インメモリの表現格納庫 ― であり、段のインタフェースを端から端まで示すものである。アーキテクチャ上の主張は、境界の上に立つ。すなわち、固定されたカーネルではなく注意出力を消費するサービスは、カーネルを純粋な差し替え対象たらしめる。

参考文献

Ashish Vaswani et al. (2017). Attention Is All You Need. Advances in Neural Information Processing Systems (NeurIPS). arXiv:1706.03762.
John Backus (1978). Can Programming Be Liberated from the von Neumann Style? A Functional Style and Its Algebra of Programs. Communications of the ACM. [1977 ACM Turing Award Lecture]
John Hughes (1989). Why Functional Programming Matters. The Computer Journal.
Christopher Strachey (2000). Fundamental Concepts in Programming Languages. Higher-Order and Symbolic Computation. [Reprint of 1967 lecture notes]
Tri Dao et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. Advances in Neural Information Processing Systems (NeurIPS). arXiv:2205.14135.