Blackwell + CUTLASS FP8 で踏む「2種類の misalignment」

はじめに

Blackwell 世代 GPU（SM120）で CUTLASS の FP8 GEMM を試したところ、
次のようなエラーに遭遇しました。

GEMM sync failed: misaligned address

一見すると単純なアライメント問題に見えますが、調査を進めると
まったく異なるレイヤで発生する 2 種類の misalignment が存在することが分かりました。

本記事では、実際に踏んだ事例をもとに、

何が misaligned なのか
なぜ混同されやすいのか
それぞれどう解決できるのか

を整理します。

検証環境

GPU: RTX 5090（Blackwell, SM120）
CUTLASS: 4.3.3
Build: -arch=sm_120a
Precision: FP8 (E4M3)
Debug tool: compute-sanitizer

結論の要約

Blackwell + CUTLASS FP8 では、2 種類の misalignment が存在します。

種類	Issue	レイヤ	状態
shared memory misalignment	#2902	ldmatrix / LDSM	解決方法あり
TMA descriptor misalignment	#2905	TMA descriptor	解決方法あり

エラーメッセージは似ていますが、原因・影響範囲・対処法は完全に別物です。

① shared memory 側の misalignment（Issue #2902）

症状

ビルドは成功
実行時に以下のエラーでクラッシュ

misaligned shared or local address

compute-sanitizer を使うと、ldmatrix（LDSM）命令で失敗していることが分かります。

原因

CUTLASS 内部で使用されている partition_S() が原因でした。

元の shared memory tensor は 16B 以上の alignment
partition_S() により 8B alignment に低下
しかし ldmatrix.sync.aligned.* は 16B strict

つまり、

shared memory の alignment contract が途中で破壊される

という問題です。

特定方法

compute-sanitizer --tool memcheck ./binary

これにより、

misaligned address
該当する ldmatrix 命令

を正確に特定できます。

解決方法（Issue #2902）

実行時に alignment を判定し、fast path / slow path を分岐します。

16B aligned
→ ldmatrix fast path
misaligned
→ scalar shared load fallback

この方法により、

misaligned address エラーは解消
correctness を維持
fast path の性能はそのまま

shared memory 側の misalignment は 設計上の前提違反だが、実装で回避可能です。

② TMA descriptor 側の misalignment（Issue #2905）

症状

shared memory 側の問題を解消した後、
FP8 + TMA kernel では 別の misalignment エラーが発生しました。

これは ldmatrix とは無関係で、
TMA（Tensor Memory Accelerator）descriptor の構築段階で起きていました。

原因

TMA は descriptor に 厳密な alignment contract を要求
CUTLASS が生成する descriptor が、
特定条件下で alignment を満たしていない
結果として、TMA load/store 実行時に失敗

つまり、

global → shared 転送の設定情報自体が misaligned

という、shared memory より前段の問題です。

解決方法（Issue #2905）

Issue #2905 では、以下の修正コードを提示しました。

TMA descriptor の alignment を明示的に保証
必要な padding / alignment を descriptor 構築時に追加
kernel 側の前提と descriptor の alignment contract を一致させる

この修正により、

TMA descriptor 起因の misalignment は解消
FP8 + TMA kernel が正しく実行可能になることを確認

👉 こちらは 回避ではなく、根本的な解決です。

なぜ混同されやすいのか

エラーメッセージが似ている
FP8 / Blackwell / TMA 周辺で同時に発生する
どちらも CUTLASS 内部実装が原因

しかし実際には、

Issue	レイヤ
#2902	shared memory load（ldmatrix）
#2905	TMA descriptor（DMA 設定）

と、問題の階層が異なります。

misalignment を一括りにすると、正しい対処に辿り着けません。

補足：12% 相対誤差はバグではない

FP8 E4M3 の仕様は以下の通りです。

exponent: 4 bits
mantissa: 3 bits
最大相対誤差 ≈ 12.5%

実測された 約12.6% の相対誤差は、
FP8 の理論限界そのものでした。

CPU 側を FP32 reference にして比較すると、この差は必ず出ます。

👉 これは精度仕様であり、実装バグではありません。

まとめ

Blackwell + CUTLASS FP8 では 2 種類の misalignment が存在する
shared memory 側（#2902）は runtime 分岐で解決可能
TMA descriptor 側（#2905）は descriptor 構築修正で解決可能
FP8 の約12%誤差は 仕様通り

重要なのは、

「どこが misaligned なのか」をレイヤごとに切り分けること

です。

関連 Issue

shared memory misalignment
https://github.com/NVIDIA/cutlass/issues/2902
TMA descriptor misalignment
https://github.com/NVIDIA/cutlass/issues/2905

おわりに

Blackwell + FP8 + CUTLASS は不安定に見えがちですが、
実際には 問題は具体的で、対処可能です。

本記事が、同じエラーに遭遇した人の
compute-sanitizer 地獄を少しでも短縮できれば幸いです。

【追記】

入力レンジと CPU reference の精度を FP8 に合わせたところ、
相対誤差は 12.6% → 0.14% に改善した。

これは FP8 E4M3 の数値特性によるものであり、
FP8 GEMM の実装自体は正しく動作していることが確認できた。

サイズ	修正前	修正後
128x128	12.6%	0.14%
256x256	12.8%	0.14%
512x512	12.8%	0.14%