技術サマリー：ブロック浮動小数点スケールを検索せよ！

問題定義

量子化は生成モデルの推論を加速するための標準的な手法であるが、超低ビット幅において精度を維持することは依然として課題である。ブロック浮動小数点（BFP）形式は、ダイナミックレンジ、メモリ効率、スループットの間のバランスの取れたトレードオフを提供するが、標準的な BFP アルゴリズムは、ブロック内の最大絶対値によって決定される固定スケールに依存している。著者らは、この「最大値スケーリング」ヒューリスティックがしばしば最適ではなく、不要な量子化誤差を引き起こすことを観察している。この問題は、マイクロスケーリング BFP 形式（NVFP4 など）に対するハードウェアサポートを導入する NVIDIA の Blackwell アーキテクチャの登場に伴い、特に重要である。従来の BFP と異なり、NVFP4 は 8 ビットブロックスケール（E4M3 形式）内に仮数ビットを含んでおり、標準的な最大値スケーリングが活用できていない、より細かい粒度のスケール選択を可能にしている。

手法

ScaleSearch

中核的な貢献は、マイクロスケーリング形式の仮数ビットを活用して最適なブロックスケールを微細に検索するアルゴリズムであるScaleSearchである。

• メカニズム: ブロックの最大値に基づいてスケールを厳格に選択する（$s = \text{round}(\|x\|_\infty / 6.0)$）のではなく、ScaleSearch はこのデフォルト値の周囲のスケールの近傍を探査する。スケールを、デフォルトスケールのビット表現に適用される整数オフセット $f$ として扱い、$s(f) = \text{reinterpret}(\text{reinterpret}(s, \text{int8}) + f, \text{fp8}_{UE4M3})$ とする。
• 検索戦略: アルゴリズムは、与えられたブロック分布に対する平均二乗誤差（MSE）を最小化するスケールを見つけるために、オフセットの小さな範囲（例：$f \in [-2, 6]$）を網羅的に検索する。
• 根拠: 合成ガウスデータと実際のニューラルネットワーク活性化（例：Llama 3.1 の Key 状態）の両方に対する実証的分析により、最適なオフセットの二峰性分布が明らかになった。デフォルトスケール（オフセット 0）は最大絶対値を FP4 値 6 にマッピングするが、ブロックの相当部分は最大絶対値を FP4 値 4 にマッピングすることでよりよく表現され、これは約 1.5 倍大きなスケールファクタを必要とする。ScaleSearch はこれらの最適なマッピングを特定する。

ScaleSearchAttention

これらの利点をアテンション機構と KV キャッシュに拡張するために、著者らはScaleSearchAttentionを提案する。これは、デ量子化のオーバーヘッドを回避するように設計された、因果言語モデリング向けのエンドツーエンド FP4 アテンションアルゴリズムである。

• ネイティブ NVFP4 実行: QKT および PV 行列乗算を、FP32 累積を用いて NVFP4 精度で Tensor コア上で直接実行する。
• 非干渉処理（IP）と絶対値低減: Q および K テンソルにおける外れ値を軽減し、量子化誤差を低減するために、Hadamard 変換（QuIP# に準拠）と、アテンションスコアを保持しつつ投影されたテンソルの平均二乗絶対値を低減する追加の線形変換ペア（$Q' = QR^{-T}, K' = KR$）を採用する。
• アテンション・シンク対応混合精度キャッシング: アテンションスコアが初期トークンと最近のトークン（アテンション・シンク）に集中することを認識し、KV キャッシュの最初のブロックと最後のブロックを全精度で保存し、中間ブロックは圧縮された NVFP4 で保存する。これにより、シーケンス長に応じてメモリ使用量を拡大することなく、重要なトークンに対して高精度を確保する。

主要な貢献

1. ScaleSearch アルゴリズム: ブロック浮動小数点量子化のためのスケール検索手法。隣接する表現可能なスケールを探査することでブロックごとの量子化誤差を最小化する。特に、NVFP4 E4M3 スケール形式の仮数分解能を活用する。
2. ScaleSearchAttention: ほぼゼロの精度劣化でエンドツーエンド推論を可能にする共最適化された FP4 アテンション機構。ScaleSearch、非干渉処理、混合精度 KV キャッシングを統合し、Tensor コア上で NVFP4 で直接 QKT および PV 操作を実行する。
3. 実証的検証: 大規模言語モデルおよび拡散モデル全体にわたって、先行する FP4 ベースライン（SageAttention3 や Naive-FP4 など）に対して、パープレキシティとベンチマーク精度で一貫した改善を示す広範な実験。

実験結果

学習後量子化（PTQ）

• 言語モデル: Qwen3-8B モデルにおいて、ScaleSearch は標準的な NVFP4 量子化と比較して、MATH-500 ベンチマークで最大15 ポイントの性能向上を達成した。また、GPQA および MMLU ベンチマークにおけるギャップも大幅に縮小した。
• 拡散モデル: SageAttention3 を使用して Mochi および CogVideoX モデルに適用した際、ScaleSearch は VQA-a、VQA-t、FScore 指標を改善し、全精度のアラインメントスコア（CLIPSIM、CLIP-T）と一致した。

アテンションとパープレキシティ

• パープレキシティ（Wikitext-2）: ScaleSearchAttention は、Llama 3.1 70B において、従来の FP4 量子化と比較してパープレキシティを最大22%削減した。Llama 3.1 70B において、PPL は2.6348を達成し、SageAttention3（3.3899）および Naive-FP4（3.4000）に対して大幅な改善を示した。
• GPQA Diamond: Llama 3.1 8B Instruct モデルにおいて、ScaleSearchAttention は**32.32%**の精度を達成し、SageAttention3（26.26%）を上回り、全精度ベースライン（31.81%）と同等の性能を示した。

効率性とオーバーヘッド

• 量子化オーバーヘッド: 制限された検索範囲において、ScaleSearch は FP32 から NVFP4 への量子化フェーズ中に最小限のオーバーヘッド1.74 倍（0.0449 ms 対 ベースライン 0.0258 ms）を導入する。
• スループット: アテンションカーネルにおいて、ScaleSearch は 32K シーケンス長での非因果アテンションにおいてベースラインスループット（SageAttention3）の**98.3%**を達成し、実行時への影響が negligible であることを示している。
• レイテンシ: テキストからビデオ生成までのエンドツーエンドレイテンシはわずかに増加したのみ（例：Mochi で 353.40 秒から 364.68 秒へ）であり、この手法が現実世界のワークロードに対して実用的であることを確認した。

重要性

本論文は、マイクロスケール形式でこれまで未活用であった仮数ビットを利用することで、標準的な BFP スケーリング戦略における根本的な非最適性に対処していると主張している。ヒューリスティックな最大値に依存するのではなく、最適なスケールを体系的に検索することで、この手法は量子化誤差を大幅に低減する。ScaleSearchAttention の導入は、これらの改善がメモリおよび計算集約的なアテンション機構に拡張可能であることを示しており、ほぼゼロの精度損失で高スループットかつ低精度な推論を可能にする。著者らは、この研究を、効率的な生成 AI のために emerging な 4 ビットハードウェア機能（NVIDIA Blackwell など）を完全に活用するための一歩として位置づけている。