mla-analysis

MLA のコスト分析とレジームガイド

レジームの選択

レジーム	s の範囲	最適なカーネル	理由
デコード	s=1	FlashMLA	FlashAttention と比較してレイテンシを 16 倍削減（圧縮 KV）
投機的デコード	s=2-32	MLAvar6+ または FlashMLA	MLAvar6+ は FlashMLA と FlashAttention を上回る性能を発揮できるはずです
プリフィル	s>128	FlashAttention	潜在空間計算による 4 倍の FLOP ペナルティを回避します

クロスオーバーポイント：DeepSeek-V3 のパラメーターでは、FlashAttention はおよそ s=16-32 で FlashMLA より高速になります。

コストモデル（DeepSeek-V3 相当：h=128、d=128、k=512、p=64）

FlashAttention

• FLOP 数：2bhst(2d + p) = 2bhst * 320
• バイト数：w * bh(s+t)(2d + p) = w * bh(s+t) * 320
• s=1 の場合：AI ≈ 1 FLOP/byte（非常にメモリ依存）
• s=1024 の場合：AI ≈ 819 FLOP/byte（非常に計算依存）

FlashMLA（吸収による潜在空間アテンション）

• FLOP 数：2bhst(2k + p) = 2bhst * 1088
• バイト数：w * (bhs(2k+p) + bt(k+p))
• s=1 の場合：AI ≈ 228 FLOP/byte（デコード時でも計算依存）
• 問題：k=512 では FLOP 数が s に比例して増加し、d=128 の場合より 1 トークンあたりの FLOP 数が 3.4 倍になります

MLAvar6+（潜在表現と展開表現の分割）

• FLOP 数：2bhstp + 4bhsnd + 4bhsok
• バイト数：w * (bhsp + bhsd + bhsk + bok + btp + bhnd + bhnd + bhsd + bhsk)
• 要点：n を選ぶことで、オペレーショナルインテンシティをルーフラインのリッジポイント付近に調整できます

吸収のトリック（FlashMLA）

Score:  Q @ K^T = Q @ (Z @ Wk)^T = (Q @ Wk^T) @ Z^T = Qz @ Z^T
Value:  softmax(A) @ V = softmax(A) @ (Z @ Wv) = (softmax(A) @ Z) @ Wv

潜在空間（k 次元）だけで処理し、出力の展開 O_latent @ W_kvb2^T は別のカーネルで実行します。

MLAvar6+ の設計

KV キャッシュを次のように分割します：

• o 個の古いトークン：潜在表現 Z（b, o, k）— 計算負荷が高く、帯域幅の使用量は少ない
• n 個の新しいトークン：展開済み K,V（b, n, h, d）— 帯域幅の使用量が多く、計算負荷は低い
• t = o + n

n=0 なら FlashMLA に近く、o=0 なら FlashAttention に近くなります。

過去のベースライン（RTX 5090、b=32、t=4096、bfloat16）

注：これらの結果は RTX 5090 での開発時のものです。最適化目標として使用する前に、現在のデバイスで再プロファイルしてください。リッジポイント、帯域幅、計算性能の上限はデバイスごとに異なります。src/mla_var3/conf/devices.json を参照してください。

構成	FlashMLA	FlashAttention	MLAvar6+ V3（最良）
s=1	419 μs	6,781 μs	829 μs（V2）
s=16	5,161 μs	6,727 μs	4,444 μs

MLAvar6+ V3 は投機的デコード（s=16）で現在最良の性能を示しており、FlashMLA を 14% 上回ります。

テンソルの形状（DeepSeek-V3 のデフォルト）

記号	説明	デフォルト
b	バッチサイズ	64
h	ヘッド数	128
s	クエリシーケンス長	可変
t	KV コンテキスト長	4096
d	ヘッド次元	128
p	位置埋め込み次元	64
k	潜在（圧縮）次元	512

テンソルのレイアウト：

• Q、Qc：[B, H, S, D] または [B, H, S, K]
• K、V：[B, T, H, D]（展開済み）
• Z：[B, T, K]（潜在 KV キャッシュ）
• QPE、KPE：[B, H, S, P] または [B, T, P]

詳細な分析

コストモデルの完全な導出とルーフライン分析：docs/cost-analysis.md