(W.I.P) 富岳(a64fx)で gemm, attention 極めたいメモ
富岳で俺得 gemm kernel 書いたけど(ほとんど Claude code まかせだけど), 性能でえへん... 40 ~ 45 % くらいが限界...
なんでぇ
エコモードで FLOPS 性能半減してた
デフォ(エコモード)だと FPU は 1pipe のみ実行なので最高 50% peak flops になります!
ブーストモード: クロック 10% マシ.
デフォルトではブーストエコなので, クロックマシの FLOPS 性能半減.
peak flops 50% 切るアプリではこれでいいかもですが,
gemm とか attention みたいに peak flops rate 70~90% になるようなコードだと, エコモード off にしたいところ.
ブーストしつつエコモード off は基本出来ないようなので,
演算メインのアプリは非ブーストで非エコの設定で動かすことになるでしょう => これが HPC であれば本来の「通常」モードといえるか.
公開されているガイド
全体的に
SVE は int 操作でも FP unit 使うので, 変にデータ変換(fp <-> int や, and/select)かけるなどするとすぐ対ピーク性能が落ちます.
L1/L2 のメモリサイズは小さいけど, HBM2 の帯域生かして(latency は 260 くらいあるけど)あまりデータ変換しないほうがたぶん性能出るでしょう.
予想としては頑張って最適化してピーク 80% 出たとして FP32 で 5 TF くらい, FP16 で 10 TF くらい?
Instruction set
Int gemm?
Int madd は vector(SVE)命令の場合 Integer unit ではなく FP unit(FLA, FLB)使います.
INT -> FP 変換でもサイクルかかるので気を付けて.
a64fx は INT8 までハードウェアサポート.
(22 TOPS ほど)
BF16?
BF16 はハードウェアサポートされていません.
(SVE2 から導入っぽい?)
FP32 に変換して処理が楽でしょうか.
fdiv, fsqrt(softmax, LayerNorm)
150 cycles くらいかかります...
おとなしく近似 div, 近似 sqrt(どちらも latency 9) + newton-raphton rounding 使いましょう...
(2 NR でほぼ 32bit fp 精度)
exp
fexpa : exp 計算の一部を高速化(小数部のテーブル引き). latency 4.
fscale : 2^n で乗算を行う. latency 9
y = x × log₂(e) を計算
y を整数部 n と小数部 f に分離 (y = n + f, 0 ≤ f < 1)
2^f を FEXPA でテーブル参照
2^n × 2^f を FSCALE で計算
svfloat32_t sve_exp_f32(svbool_t pg, svfloat32_t x) {
const float log2e = 1.442695041f;
const float ln2_hi = 0.693359375f;
const float ln2_lo = -2.12194440e-4f;
// Step 1: y = x * log2(e)
svfloat32_t y = svmul_f32_x(pg, x, svdup_f32(log2e));
// Step 2: n = floor(y), f = y - n
svfloat32_t n_f = svrintm_f32_x(pg, y); // floor
svint32_t n = svcvt_s32_f32_x(pg, n_f);
svfloat32_t f = svsub_f32_x(pg, y, n_f);
// Step 3: 2^f via FEXPA + polynomial refinement
// FEXPA用に f を整数表現に変換
// FEXPAは入力の下位6bit(FP32)をインデックスとして使う
svuint32_t f_bits = svreinterpret_u32_f32(f);
svfloat32_t exp_f = svexpa_f32(svlsl_n_u32_x(pg, f_bits, 17));
// 多項式で精度を補正
// exp(f*ln2) ≈ 1 + f*ln2 + (f*ln2)^2/2 + ...
svfloat32_t r = svmul_f32_x(pg, f, svdup_f32(0.6931472f)); // f * ln(2)
svfloat32_t r2 = svmul_f32_x(pg, r, r);
svfloat32_t poly = svmla_f32_x(pg, svdup_f32(1.0f), r, svdup_f32(1.0f));
poly = svmla_f32_x(pg, poly, r2, svdup_f32(0.5f));
// Step 4: 2^n * 2^f via FSCALE
svfloat32_t result = svscale_f32_x(pg, poly, n);
return result;
}
log
FLOGB は SVE2 から. a64fx では使えない
sin/cos
sin/cos に使える Trigonometric 関数 FTMAD/FTSMUL(latency 9), FTSSEL(latency 4) あり.
Sector cache(L1, L2)
L1 : 4 セクタ(仕様上. 実際は使い勝手を考えて 2 セクタの設定になっている), L2 : 2 セクタまで作成可能
セクタ0は insn cache などにも使われる
L2 はアシスタントコアで 2 way, 1MB 消費される. したがってアプリでは 14 way, 7MB まで利用可能.
HW の機能というわけではなく, コンパイラでのソフトウェア対応の模様.
(numactl とかで制御できるわけではない)
Online softmax
Softmax の e^x / sum(e^x_i) はブロック単位などで逐次計算可能.
実際は数値安定性のためにブロック単位で min/max 取って補正
(Flash attention 1).
あとは除算を block 最後に行う(Flash Attention 2)
Gradient checkpointing
ややこしい名称(Gradient の特定位置にチェックポイントを置くみたいに受け取れる)だが, 一部のレイヤーの activation(中間結果)だけをメモリに保存する.
メモリに保存するより必要に応じて Layer 再計算のほうがコストがかかならいのに適している.
Attention 出力 -> 保持したほうが有利
LayerNorm, MLP など -> 計算コスト低い, メモリサイズ大きいなどで破棄(再計算)が有利
fp16 madd + fp32 accum
BF16 FDOT がないので fp16 で頑張ってみる.
fp64 gemm kernel(ピーク 93% くらい) をベースに, Claude code Opus 4.5 クンに頑張ってもらった.
ただ fp16 で accum すると精度が気になるので, K 個単位で fp32 にキャストし accum.
ロードとストアはとりあえず考えないものとした(同一アドレスなどにした)
K 少ない(fp32 にする頻度多い)とストールや変換コスト大で性能出ない
K 多いとパイプラインは埋まりやすいが, 蓄積誤差がきになる.
汎用だと K=24 ~ K=32 くらいが精度と速度のトレードオフ
fp16 ピーク性能に対して 64 ~ 70%.
fp32 性能換算だと 1.3 ~ 1.4 倍くらい.
int8/int16 gemm
INT8/INT16 の SDOT/UDOT は accum が INT32/INT64 になる.
レイテンシは FP と同じく 9
TF32(BF19)
中身は 19bit の BF.
BF16 より精度ほしいが FP32 ほどでもない場合.
TF32 -> FP32 は楽(NV 互換の丸めに合わせる場合)
FP32 -> TF32 はちょっと処理必要
通信や HBM 上では TF32 で持たせて約 40% のデータ削減見込める.
パックするとビットが中途半端なので,
TF32 仕様に合わせなくてもよければ, CPU(or FP32 強いアクセラレータ)なら FP18 にしたり FP20 にしたりでもよかろう.
64bits に 3 つパックするとして [21, 22, 21] とか, 128bits に 18bit で 7 個パック(2bit あまり)とか.
Load latency
L1D
L1D -> reg
scalar だと 5 cycles
arm64v8 fp&SIMD(128bit)では 8
SVE では VL によらず 11 cycles!
L2(shared)
L2 -> L1 が 37 ~ 47 cycles くらい. average 42
HBM2
HBM2 <-> L2 が 260 cycles くらい.
Sector cache
ポインタにタグ付けで制御
way 数指定はコンパイラレベルでの処理なのかしら
(tag bitfield はなし?)
// Sector cache hints on pointers
mov x22, 0x2
lsl x22, x22, 56
orr x6, x6, x22 // C pointer with sector hint
mov x22, 0x41
lsl x22, x22, 56
orr x3, x3, x22 // B pointer with sector hint
┌─────────────────────────────────────────────────────────────┐
│ A行列 (x3): 0x41 << 56 │
│ │
│ 0x41 = 0100 0001 (binary) │
│ │
│ bits 63:60 = 0100 = 4 → pf_func = 4 (強プリフェッチ?) │
│ bits 59:58 = 00 → 予約 │
│ bits 57:56 = 01 = 1 → sector = 1 │
│ │
├─────────────────────────────────────────────────────────────┤
│ C行列 (x6): 0x2 << 56 │
│ │
│ 0x2 = 0000 0010 (binary) │
│ │
│ bits 63:60 = 0000 = 0 → pf_func = 0 (通常) │
│ bits 59:58 = 00 → 予約 │
│ bits 57:56 = 10 = 2 → sector = 2 │
│ │
├─────────────────────────────────────────────────────────────┤
│ B行列 (x4): 設定なし │
│ │
│ → sector = 0 (デフォルト) │
│ → pf_func = 0 (通常) │
│ │
└─────────────────────────────────────────────────────────────┘
┌─────────────────────────────────────────────────────────────┐
│ A64FX pf_func フィールド (bits 63:60) │
│ │
│ 値 │ 動作 │
│─────┼───────────────────────────────────────────────────────│
│ 0 │ 通常動作 │
│ 1 │ L1 ハードウェアプリフェッチ抑制 │
│ 2 │ L2 ハードウェアプリフェッチ抑制 │
│ 3 │ L1 + L2 プリフェッチ抑制 │
│ 4 │ 強いプリフェッチ (aggressive) │
│ 5 │ ... │
│ │
│ A行列に pf_func=4: │
│ → ハードウェアプリフェッチを積極的に行う │
│ → A は同じデータを複数回再利用するため有効 │
└─────────────────────────────────────────────────────────────┘
┌──────────┬────────┬──────────┬─────────────────────────────┐
│ 行列 │ セクタ │ pf_func │ 備考 │
├──────────┼────────┼──────────┼─────────────────────────────┤
│ B (pb) │ 0 │ 0 │ パック済み、ストリーム読み │
│ A (a) │ 1 │ 4 │ 強プリフェッチ │
│ C (pc) │ 2 │ 0 │ Read-Modify-Write │
│ (未使用) │ 3 │ - │ │
└──────────┴────────┴──────────┴─────────────────────────────┘
L2 (8MB, 16-way) のデフォルト分割:
Sector 0: 4-way (2MB) → B行列
Sector 1: 4-way (2MB) → A行列
Sector 2: 4-way (2MB) → C行列
Sector 3: 4-way (2MB) → 未使用
// A行列 (x3 → x21 に格納)
str x3, [sp,136+16] // スタックに保存
...
ldr x21, [sp,136+16] // x21 として使用
// A要素のブロードキャストロード
ld1rd z22.d, p1/z, [x0,0] // x0 = A のオフセット付きポインタ
ld1rd z23.d, p1/z, [x11,0] // x11 = A + lda*1
ld1rd z24.d, p1/z, [x10,0] // x10 = A + lda*2
ld1rd z25.d, p1/z, [x9,0] // x9 = A + lda*3
ld1rd z26.d, p1/z, [x7,0] // x7 = A + lda*4
// C行列 (x6)
ld1d z2.d, p1/z, [x6,0,mul vl] // C ロード
st1d z2.d, p1, [x6,0,mul vl] // C ストア
prfm pstl1keep, [x6,xzr,lsl #3] // C プリフェッチ
// B行列 (x4 → x1 経由)
ld1d z0.d, p1/z, [x1,0,mul vl] // B ロード
prfm pldl1keep, [x1,256*9] // B プリフェッチ
キャッシュ動作の図解
┌─────────────────────────────────────────────────────────────┐
│ L2 Cache (8MB) │
│ │
│ ┌─────────────┐ ┌─────────────┐ ┌─────────────┐ ┌─────────┐│
│ │ Sector 0 │ │ Sector 1 │ │ Sector 2 │ │Sector 3 ││
│ │ (B行列) │ │ (A行列) │ │ (C行列) │ │ (空き) ││
│ │ 4-way │ │ 4-way │ │ 4-way │ │ 4-way ││
│ │ 2MB │ │ 2MB │ │ 2MB │ │ 2MB ││
│ │ │ │ pf_func=4 │ │ │ │ ││
│ │ ストリーム │ │ 強PF+再利用│ │ R-M-W │ │ ││
│ └─────────────┘ └─────────────┘ └─────────────┘ └─────────┘│
│ │
│ 効果: A,B,C が互いのキャッシュラインを追い出さない │
└─────────────────────────────────────────────────────────────┘
┌─────────────────────────────────────────────────────────────┐
│ DGEMM のアクセスパターン: │
│ │
│ B行列 (Sector 0): │
│ - パック済みで連続アクセス │
│ - K方向にストリーム読み │
│ - 通常プリフェッチで十分 │
│ │
│ A行列 (Sector 1, pf_func=4): │
│ - 同じ列を複数の C タイルで再利用 │
│ - 強プリフェッチで L2 に確実にプリロード │
│ - セクタ分離で B に追い出されない │
│ │
│ C行列 (Sector 2): │
│ - Read-Modify-Write パターン │
│ - 計算後に書き戻し │
│ - A, B と競合しないようにセクタ分離 │
└─────────────────────────────────────────────────────────────┘
dgemm.s 解析
5x2 x 2 unroll のマイクロカーネル
A からは ld1rd で scalar -> vector
B は ld1d で vector load(5 回再利用)
C accum 20 本(10 x 2)
A 10 本(10 x 2)
B load 2 本(ping-pong)
B load 命令の発行(レジスタ依存)が fmla の 5~6 cycle 前.
ld1d, ld1rd のレイテンシは 11(ld1rd でも long mode), fmla は 9 で間に合わないとおもったけど, out-of-order なりレジスタリネーミングなりでうまくパイプラインうまるようになっているのかしら.
Out of order
コミットスタック: 128 entries
FPU 物理レジスタ: 128 entries(論理 32 本, リネーミング・レジスタ 96 本)
RSE(Reservation Station for EXE): 20 x 2 split
(FLA/FLB ごとに 20 本)
Fetch port: virtual 160, real 40
Store port: virtual 160, real 40
というわけで実レジスタは 4 倍ほどある.
FPU RSE 20 本なので, 20 FPU 命令(5x4) をダブルバッファリングする程度の余裕はありそう?
Fetch port/store port はよくわからない.
少なくとも 20 lds/sts をダブルバッファリングする余裕はあるのカナ.
gemm など演算ネックな(演算命令に対して ld/st 命令が少ない)な kernelであれば特に問題にならないと fetch port/store port 数は思われる.
参考
BF16 だけど B200(?) で 1.8 PFLOPS, 実 LLM の matrix 構成で 1.5 PF くらい.
(peak BF16 2.2 TFLOPS(dense))
B200 は FP32 は 74 TF くらいで a64fx の 12 倍くらい.
BF16 とかの低精度では a64fx は B200 に比べて 100 倍くらい効率悪いですが, FP32 だとまあスペック相当の比率なのでワンチャン a64fx にもすこし利点ありそ?(あとは電力効率とか, 大規模での安定運用とか)
2025 末の今は FP8 attention(Flash Attention 3) や INT gemm で低精度 LLM training が流行りですが, Ozaki-scheme で fp64 性能だせるようになったり, Rubin で fp64 強化(?)の噂もあるようなので, 将来は LLM training でも FP32, FP64 回帰になる... かもしれません.
Discussion