AIは「1+1って、2になること多いなあ」と思っている！？ ChatGPTに「1+1は？」と聞けば、当然「2」と返ってきます。 実はこのときのAIの内部で起こっていることは、割と大真面目に 「 私のデータによれば、1+1の答えは最も2が多いです 」なのです。 計算してるんじゃないの？ ChatGPTのようなAI（大規模言語モデル）は、極端なことを言ってしまえば、次の単語予測マシンです。 たとえば「むかしむかし、あるところに」と言われたら、「おじいさんと」と返す。「今日の天気は」と言われたら、「晴れ」とか「曇り」とか返す。 膨大な文章を読んで「この言葉の次にはこの言葉が来やすい」というパターンを学習しているだけなんです。 AIにとっては、計算問題も文章の一種のため、数学の問題も同じやりかたで解いています。 「1+1=」という文字列を見て、「この後には2が来ることが多いな」と思って2を返しているだけ。 つまり: 人間: 1+1を理解して、演算して、2を導く AI: 「1+1=」の後に「2」が来るパターンで覚えてる そう、タイトルの通り、AIは「1+1って、2になること多いなあ」なのです。 ほんと？ もしAIが計算を「理解」しているのではなく、単なる「パターンの暗記」だとしたら、 見たことがないパターンに出会ったときにボロが出るはず です。 その実態がよくわかる、2つの証拠を見てみましょう。 証拠1: ちょっと難しいかけ算ができない 普通の計算機（電卓）なら、2桁の足し算ができれば、4桁になっても10桁になっても、やることは同じなので間違えません。 しかし、OpenAIの研究論文「 Language Models are Few-Shot Learners 」（Brown et al., 2020）によると、当時のGPT-3は以下のような結果になりました。 2桁の足し算： ほぼ100%正解 4桁以上の計算： 正答率は 20%以下 に急落 つまり、ネット上にたくさん転がっている「よくある計算（2桁）」はパターンとして覚えているけれど、滅多に見かけない「大きな桁の計算」は、パターンがないのでお手上げになってしまうのです。   証拠2: 聞き方が変だとできなくなる たとえば「0.7 × 5 は？」と聞けばAIは即座に3.5と答えます。 でも同じ計算を「7×10⁻¹ × 5 は？」と科学的記数法で書くと、数学的にはまったく同じ計算なのに、急に怪しくなります。 Yang et al.（2024）の論文「 Number Cookbook 」では、LLMは標準的な整数計算には強い一方、分数や科学的記数法になると、精度が 20%以下 まで落ちることが示されています。 数学を理解しているなら「書き方が違うだけ」と分かりますが、AIにとっては 見たことがない珍しい文字列 に見えてしまうため、予測が外れてしまうのです。   やってみよう 実際に試してみましょう。ChatGPTに「4726 × 3891 は？」と聞いてみます。 あれ、普通に正解されました。 実は、今のAIはこの「弱点」を、 知能ではなく「仕組み」で克服 しつつあるんです。 今は解決してる ChatGPTの新しいもの（GPT-4o）やClaude 4.5などは、数学の実力テスト（MATHやGSM8Kといった、AIに数学の問題を解かせてスコアを測る標準テスト）で非常に高いスコアを出しています。その理由は主に2つあります。 1つ目はツールの利用です。 計算が必要な場面で、内部的にプログラミングコードを実行したり電卓を使ったりして、LLM自体は直接計算せずに正解を得る方法が発達しました。 2つ目は、Chain-of-Thoughtというもので、段階を踏んで思考することで、単純に答えるよりもはるかに複雑な問題を解けるようになりました。 例えば「23 × 47」を一発で出すのではなく、「23 × 40 = 920、23 × 7 = 161、920 + 161 = 1081」のように段階を踏めます。 これにより、実用上はAIは計算ができると言える水準になっています。   しかし、ここで注意したいのは、これらはいずれも AIが計算を「理解」しているわけではない ということです。 ツール利用は外部の計算機に丸投げしているだけですし、Chain-of-Thoughtも、本来1ステップでは処理しきれない問題を小さく分割して、トークン生成の過程で中間結果を一時的に保持する仕組みです。 LLMの本質は今も変わらず、次に来る言葉の予測であり、これらの間接的な手順を禁じた場合、AIは依然として大きな数の掛け算や見慣れない形式の計算で間違えます。 まとめ AIは計算を「理解」しているのではなく、「1+1の後には2が来やすい」というパターンで答えている そのため、桁が大きい計算や見慣れない書き方には弱い 最近はツール利用（コード実行）やChain-of-Thought（段階的思考）で実用上は高精度になった ただし、これらはあくまで補助手段であり、LLMの本質は変わっていない AIに計算させるときは、裏でちゃんとコードを実行しているか意識しておくと安心 おまけ 僕はブログを書くときに、AIにレビューをしてもらっています。 この記事をレビューさせてみたときのAIの反応がこちら：   AI: AIは計算ができない！？のレビューを開始します。 … ファクトチェック中… 計算例が間違っています。 本来ならば4726 × 3891 の正しい答えは 18,374,766 です。修正します.. … あ、そういえば私もAIでした。 Pythonコードを実行し、確認します。   最近のAIは賢いですね。 実は、計算例の部分は僕が「AIにこんな感じの例を提示して」と出力させたものなので、 つまりこのやりとりは、 AIが「AIは計算ができない！？」という記事をレビューし、 AIが計算した計算ミスをAIが指摘・修正し、 同時に自分がAIであることにはたと気づき、 AIは計算ができないため、直ちにプログラム実行によって確証を得る という、皮肉・メタ認知のミルフィーユであり、今のAIのアホさと賢さの両面が綺麗に凝縮されたやりとりでした。 関連記事 AIに嘘つかないでよーとお願いするとちょっと効くという記事を書いています： chatGPTに「ハルシネーションしないで」とお願いしたら効果がある？   ご覧いただきありがとうございます！ この投稿はお役に立ちましたか？ 役に立った 役に立たなかった 0人がこの投稿は役に立ったと言っています。 The post AIは「1+1って、2になること多いなあ」と思っている！？ first appeared on SIOS Tech Lab .

MathJax={tex:{inlineMath:[['$','$']],displayMath:[['$$','$$']],processEscapes:true}}; こんにちは、Insight Edgeでデータサイエンティストをしている新見です。 cuTile Pythonとは 背景 特徴 従来のCUDA（SIMT）との違い 文法 TileGymで行列積ベンチマーク 倍精度行列積エミュレーション Ozaki Schemeについて 分解(Split) 行列積の計算 素朴な実装と初回結果 最適化 Fast Mode（GEMMの削減） Fused Split Kernel（分割の融合） 最適化後の結果 dによる精度/速度トレードオフ まとめ 参考文献 今回はNVIDIAが発表したばかりの「cuTile Python」を試してみました。普段は、GPUカーネルを業務で書くことはありませんが、cuTileはPythonで書かれていて、文法もシンプルなようなので、GPUプログラミングの勉強の意味も含めて記事にしました。 cuTile Pythonとは cuTile Pythonは、NVIDIA GPU向けの新しい並列プログラミングモデル「cuTile」をPythonから使うためのDSL（ドメイン固有言語）です。 背景 GPU上で高速に動作する処理を自前で記述したい場面は増えていますが、CUDA C++の習得コストは依然として高いのが実情です。 PyTorchやcuBLASといった高レベルAPIで日常的な開発は十分カバーできるものの、LLM推論の最適化など低レイヤへの介入が求められる局面も増えてきました。NVIDIA Ampere世代以降のGPUではTensor CoreやTMA（Tensor Memory Accelerator）といったハードウェア機能が追加されており、これらを十分に活用するにはより踏み込んだプログラミングが必要になります。 しかし、ハードウェアを意識したコードを書く難易度は上がり続けています。メモリ階層ひとつ取っても、共有メモリ、レジスタ、Blackwell世代で追加されたTensor Memoryなど、用途に応じて使い分ける必要があり、それぞれの特性に合わせたデータ配置や転送の制御が求められます。 さらに、特定のアーキテクチャに最適化したコードは新しい世代のGPUが登場した途端に書き直しが必要になることも多く、保守コストも無視できません。 こうした背景から、ハードウェアの詳細を抽象化しつつ高い性能を引き出すDSLへの需要が高まっています。OpenAIの gpt-oss リポジトリでもTritonという同様のDSLが採用されており、この手のアプローチは業界でも広く注目されています。 特徴 GPUプログラミングには、cuBLASやPyTorchのような高レベルライブラリか、CUDA C++やPTXといった低レベルなスレッド制御か、という二極化がありました。cuTileはこの中間に位置する「Tileレベル」のプログラミングモデルです。 抽象レベルについて（動画[1]より） 以下、動画[1]で紹介されていた特徴になります。 CUDAプラットフォームにネイティブ統合 : OpenAI Tritonなどサードパーティ製DSLとは異なり、cuTile（Tile IR）はCUDAドライバに組み込まれています。既存のプロファイラやデバッガがそのまま使えます。 Tile IRへのコンパイル : Pythonで書いたカーネルは「Tile IR」という仮想ISAに変換され、ドライバが実行時にターゲットGPUに合わせた最適なマシンコード（SASS）を生成します。 技術スタックの階層構造、TileIRはPTXを置き換えるのではなく共存する（動画[1]より） 従来のCUDA（SIMT）との違い 従来のCUDA（SIMT: Single Instruction, Multiple Threads）とcuTileでは、プログラマが何を書いて何をシステムに任せるかが大きく異なります。 特徴 従来のCUDA (SIMT) cuTile (Tile-based) 実行単位 スレッド単位でデータ処理を記述。WarpやBlockの構成を意識する必要がある データの塊（タイル）と単一の実行単位（ブロック）で思考。スレッドへの分解はシステムが行う データ処理 個々のスレッドへのデータ分配（ストライディングなど）を手動で計算・管理 タイル（配列全体の一部）を一つの単位としてロード・演算・ストア メモリ管理 共有メモリの確保、同期（バリア）、バンクコンフリクト回避などをユーザーが管理 システムが管理。共有メモリの利用や同期は自動化され、ユーザーからは隠蔽 ハードウェア活用 Tensor Coreなどを使うには複雑なPTX命令や特定のレイアウトを意識する必要がある ct.load や演算子を書くだけでTMAやTensor Coreを自動的に活用 文法 cuTile Pythonは、Pythonのデコレータと専用の型システムを使って記述します。詳しくは公式ドキュメント[2]を参照してください。 主な特徴: @ct.kernel デコレータ : Python関数をGPUカーネルとしてマーク。関数内ではcuTile Pythonの文法に従う。 イミュータブルなタイル : カーネル内では、タイルが操作対象となる。タイルは「値」として扱われ、変更不可。演算すると新しいタイルが生成されます Array (Global Memory): 引数から取得、ミュータブル。 ct.load / ct.store でアクセス Tile (Local/Register): イミュータブルで演算対象 以下、ベクトル加算のコード例です。 import cuda.tile as ct # タイルサイズはコンパイル時定数 TILE_SIZE = 16 @ ct.kernel def vector_add_kernel (a, b, result): # 1. 現在のブロックIDを取得 (スレッドIDではない！) block_id = ct.bid( 0 ) # 2. グローバルメモリ(Array)からタイルとしてデータをロード # システムが自動的に最適なメモリ転送(TMA等)を行う a_tile = ct.load(a, index=(block_id,), shape=(TILE_SIZE,)) b_tile = ct.load(b, index=(block_id,), shape=(TILE_SIZE,)) # 3. タイル同士の演算 (要素ごとの加算が一括で行われる) result_tile = a_tile + b_tile # 4. 結果をグローバルメモリにストア ct.store(result, index=(block_id,), tile=result_tile) # ホスト側からの実行 # ct.launch(stream, grid_dim, kernel_func, args) ブロックごとに同一のカーネルが実行され、各ブロックはIDで指定されたデータを担当範囲として、処理を行います。タイル演算は、感覚としてはnumpyの処理に似ています。 TileGymで行列積ベンチマーク 実際に動かします。cuTileはCUDA Toolkit13.1以降が必要で、これはBlackwell世代以降の比較的最新のGPUでしか動かないようです。私は手元に最新のGPUがないので、クラウドサービスを利用したいと思います。今回は、 Modal と呼ばれるGPU特化のクラウドサービスを利用しました。 Modalは関数ベースでGPUインスタンスを立ち上げられるサービスになります。使い勝手がよく、便利です。実行時間に応じた従量課金制で、今回の検証のような少しGPUを試してみたい場合に適しています。 今回は、公式のサンプルレポジトリTileGym[3]をベースに、行列積のコードの実行をしてみます。Modalで走らせる実行コードを以下に示します。imageでDockerイメージを作成し、TileGymのレポジトリをクローン、ライブラリインストールを行います。Modalの詳細は ドキュメント を参照してください。今回対象のGPUはB200です。 # run-tilegym.py import modal image = ( modal.Image.from_registry( "nvidia/cuda:13.1.0-devel-ubuntu24.04" , add_python= "3.13" ) # CUDA 13.1開発環境イメージ .apt_install( "git" ) .run_commands( "pip install --pre torch --index-url https://download.pytorch.org/whl/cu130" ) # PyTorchインストール、比較のため .run_commands( "git clone https://github.com/NVIDIA/TileGym.git && cd TileGym && pip install -e ." ) # cuTile, TileGymインストール .entrypoint([]) ) app = modal.App( "tilegym-test" ) @ app.function (gpu= "B200" , image=image, timeout= 600 ) def run_mma_bench (): import os os.chdir( "/TileGym" ) os.system( "python tests/benchmark/bench_matrix_multiplication.py" ) @ app.local_entrypoint () def main (): run_mma_bench.remote() 上のコードをrun-tilegym.pyとして保存し、 modal run run-tilegym.py で実行します。問題なければ結果は、以下のように出力されるはずです。 matmul-performance-float16-TFLOPS: M N K CuTile PyTorch 0 1024.0 1024.0 1024.0 271.056760 473.522850 1 2048.0 2048.0 2048.0 1129.688506 1199.365877 2 4096.0 4096.0 4096.0 1235.696555 1401.341171 3 8192.0 8192.0 8192.0 1483.030888 1253.946946 4 16384.0 16384.0 16384.0 1356.600018 1536.098446 5 32768.0 32768.0 32768.0 1254.836929 1306.057063 matmul-performance-float8_e5m2-TFLOPS: M N K CuTile 0 1024.0 1024.0 1024.0 277.309352 1 2048.0 2048.0 2048.0 1154.454102 2 4096.0 4096.0 4096.0 2769.415226 3 8192.0 8192.0 8192.0 2981.168986 4 16384.0 16384.0 16384.0 2935.864636 5 32768.0 32768.0 32768.0 2658.604232 CuTileとPyTorchの行列積のベンチマークが出ています。float16とfloat8_e5m2の両方で行列積を実行していますが、PyTorchでは、後者の行列積が未対応のようです。PyTorchは裏側でcuBLASを呼び出しているので実質cuBLASとの比較です。float16では、CuTileはPyTorchに近い性能、一部のサイズでは、PyTorchを上回る性能が出ています。float8_e5m2では、行列サイズが4096以上でfloat16の約2倍の性能が出ています。 以下が TileGym/src/tilegym/ops/cutile/matmul.py の行列積のカーネルコードの抜粋です。 @ ct.kernel (num_ctas=ct.ByTarget(sm_100= 2 )) def matmul_kernel (A, B, C, TILE_SIZE_M: ConstInt, TILE_SIZE_N: ConstInt, TILE_SIZE_K: ConstInt): # 担当タイルのインデックス計算（L2キャッシュ局所性のためswizzle） bidx, bidy = swizzle_2d(A.shape[ 0 ], B.shape[ 1 ], TILE_SIZE_M, TILE_SIZE_N, GROUP_SIZE_M= 8 ) num_tiles_k = ct.num_tiles(A, axis= 1 , shape=(TILE_SIZE_M, TILE_SIZE_K)) # FP32アキュムレータの初期化（FP16入力でも精度維持のためFP32で累積） accumulator = ct.full((TILE_SIZE_M, TILE_SIZE_N), 0 , dtype=ct.float32) # FP32→TF32変換（Tensor Coreを利用するため） dtype = ct.tfloat32 if A.dtype == ct.float32 else A.dtype # K方向にタイル単位でループ for k in range (num_tiles_k): a = ct.load(A, index=(bidx, k), shape=(TILE_SIZE_M, TILE_SIZE_K), padding_mode=ct.PaddingMode.ZERO).astype(dtype) b = ct.load(B, index=(k, bidy), shape=(TILE_SIZE_K, TILE_SIZE_N), padding_mode=ct.PaddingMode.ZERO).astype(dtype) accumulator = ct.mma(a, b, accumulator) # 行列積計算・累積 # 出力型に変換して結果を書き出し ct.store(C, index=(bidx, bidy), tile=ct.astype(accumulator, C.dtype)) A:MxK @ B:KxN -> C:MxN の行列積で、M方向、N方向単位でバッチに切り分けCの部分タイルごとに並行して実行されます。K方向にも部分分割して、順次読み込み(load), 行列積計算(mma), 結果の保存(store)を行っています。cuTile側でメモリの種類やMMA命令の選択は書く必要がなく、コンパイル時に自動的に最適化されます。 このように簡潔に書いても、ゴリゴリにチューニングしているcuBLASに匹敵した性能を出しているというのがcuTileの売りなようです。 ベンチマークを動かしただけでは面白くないので、型の精度を少し上げて同様の計算をしてみます。F32演算の場合、上記コードでは行列をTF32に変換してから計算しています。それと合わせるため、PyTorch側も以下のようにTF32を有効化します。 # TileGym/tests/benchmark/bench_matrix_multiplication.py # Enable TF32 for PyTorch to match Tensor Core behavior torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True また、FP64演算にあたり、累積の型がFP32では精度が足りないため、cutileコード側で累積の型をFP64に変更する処理を追加しています。 # Initialize an accumulator for the current output tile (TILE_SIZE_M x TILE_SIZE_N). # Use float64 for float64 inputs, otherwise float32 for higher precision accumulation. acc_dtype = ct.float64 if A.dtype == ct.float64 else ct.float32 accumulator = ct.full((TILE_SIZE_M, TILE_SIZE_N), 0 , dtype=acc_dtype) 以下が修正後のベンチマーク結果です。 matmul-performance-float32-TFLOPS: M N K CuTile PyTorch 0 1024.0 1024.0 1024.0 208.295471 294.114105 1 2048.0 2048.0 2048.0 665.976324 648.103430 2 4096.0 4096.0 4096.0 698.961883 747.326296 3 8192.0 8192.0 8192.0 783.858756 761.237840 4 16384.0 16384.0 16384.0 856.688401 742.126004 matmul-performance-float64-TFLOPS: M N K CuTile PyTorch 0 1024.0 1024.0 1024.0 0.855789 26.687611 1 2048.0 2048.0 2048.0 1.063844 33.830530 2 4096.0 4096.0 4096.0 1.124713 35.400544 3 8192.0 8192.0 8192.0 1.124824 35.438650 FP32では、PyTorchに近い性能が出ています。一方、FP64では、cuTile側での最適化がまだ不十分なようで、PyTorchに大きく劣る結果となっています。TILE_SIZEをより小さく設定することで、1.6 TFLOPS程度には改善しましたが、まだ大きく劣っています。 原因としては、cuTileの ct.mma がFP64演算に対して効率的な命令へマッピングできていない可能性が高いです。cuBLAS（PyTorch）はFP64 Tensor Coreを含むハードウェアリソースを最大限に活用した成熟した実装を持っており、この差が性能差に直結しています。 ここで、FP64演算の性能を向上させるために、Ozaki Schemeと呼ばれる倍精度行列積エミュレーション手法を試してみます。 倍精度行列積エミュレーション Ozaki Schemeについて Ozaki Schemeは、FP64の行列積をFP64演算なしで高精度にエミュレートする手法です[4][5]。詳しくは元の論文を読んでほしいのですが、概要を説明します。基本的なアイデアは、FP64行列を複数の低精度行列に分解し、Tensor Coreで高速に行列積を計算するというものです。行列の分解、行列の計算、結果の累積の3段階で構成されます。 分解(Split) 論文に従い、以下の型を定義します。 Type1 (FP64): 元の行列の精度。仮数部 $m_{\text{Type1}} = 53$ ビット Type2 : 分解先の低精度型（BF16, FP16, FP8等）。仮数部 $m_{\text{Type2}}$ ビット（隠れビット含む） Type3 (FP32): Tensor Coreの累積精度。仮数部 $m_{\text{Type3}} = 24$ ビット Type1の行列 $\boldsymbol{x}$ を、残差 $\boldsymbol{x}^{(p)}$ がゼロになるまで再帰的にType2スライス $\bar{\boldsymbol{x}}^{(p)}$ に分解します。$\boldsymbol{x}^{(1)} = \boldsymbol{x}$ として、各ステップ $p$ で以下を行います。 $$c_x^{(p)} = \left\lceil \log_2 \left( \max_i \left| x_i^{(p)} \right| \right) \right\rceil \tag{1}$$ $$\sigma = 0.75 \cdot 2^{\rho + c_x^{(p)}} \tag{2}$$ $$v_i = \text{fl}_{\text{Type1}} \left( \left( x_i^{(p)} + \sigma \right) - \sigma \right) \tag{3}$$ $$x_i^{(p+1)} = \text{fl}_{\text{Type1}} \left( x_i^{(p)} - v_i \right) \tag{4}$$ $$\bar{x}_i^{(p)} = \text{cvt}_{\text{Type2}} \left( \text{fl}_{\text{Type1}} \left( 2^{-c_x^{(p)}} v_i \right) \right) \tag{5}$$ ここで $\rho$ は精度パラメータ（Type1, Type2, Type3の仮数部ビット数と内積次元 $k$ から決定）です。$\sigma$ を足して引く操作（式3）がVeltkamp分割の核心で、上位 $m_{\text{Type2}}$ ビットを正確に抽出します。式4で残差を更新し、式5で $2^{c_x^{(p)}}$ で正規化してType2スライスを得ます。 この結果、$\boldsymbol{x}$ は $s_x$ 個のスライスに分解されます。 $$\boldsymbol{x} = \sum_{p=1}^{s_x} 2^{c_x^{(p)}} \cdot \bar{\boldsymbol{x}}^{(p)} \tag{9}$$ $c_x^{(p)}$ が指数部、$\bar{\boldsymbol{x}}^{(p)}$ が仮数部に対応します。スライス数 $s_x$ は $\boldsymbol{x}^{(p)} = 0$ になるまでの反復回数で決まり、理論的には $\lceil m_{\text{Type1}} / m_{\text{Type2}} \rceil$ ステップですが、行列要素のスケールのばらつきにより多くなることがあります。 PyTorchでの実装は以下の通りです。 def ozaki_split_to_type2_slices (x, k, type2, max_slices= 20 ): # 仮数部ビット数（隠れビット含む） m_fp64, m_fp32 = 53 , 24 m_type2 = - int (math.log2(torch.finfo(type2).eps)) + 1 # 精度パラメータ ρ の計算 gamma = math.ceil(m_fp64 - (m_fp32 - math.log2(k)) / 2 ) xi = m_fp64 - m_type2 rho = max (gamma, xi) slices = [] residual = x.clone().to(torch.float64) for _ in range (max_slices): max_abs = residual.abs().max().item() if max_abs == 0 or max_abs < 1e-300 : break c_x = math.ceil(math.log2(max_abs)) # 式(1) sigma = 0.75 * math.ldexp( 1.0 , rho + c_x) # 式(2) v = (residual + sigma) - sigma # 式(3) Veltkamp分割 residual = residual - v # 式(4) 残差更新 scale = math.ldexp( 1.0 , c_x) slice_type2 = (v / scale).to(type2) # 式(5) 正規化 + Type2変換 slices.append((slice_type2, scale)) return slices # [(Type2スライス, 2^c_x), ...] 行列積の計算 行列 $\boldsymbol{x}$, $\boldsymbol{y}$ をそれぞれ分解すると、行列積は以下のように展開できます。 $$\boldsymbol{x}^T \boldsymbol{y} = \sum_{p=1}^{s_x} \sum_{q=1}^{s_y} 2^{c_x^{(p)} + c_y^{(q)}} \cdot \bar{\boldsymbol{x}}^{(p)T} \bar{\boldsymbol{y}}^{(q)} \tag{10}$$ 各 $\bar{\boldsymbol{x}}^{(p)T} \bar{\boldsymbol{y}}^{(q)}$ はType2行列同士の積であり、Tensor CoreのGEMMで計算できます。Ozaki Schemeではρパラメータにより、このGEMMのType3（FP32）での累積が丸め誤差なしで成立するよう設計されています。 $$\bar{\boldsymbol{x}}^{(p)T} \bar{\boldsymbol{y}}^{(q)} = \text{fl}_{\text{Type3}} \left( \bar{\boldsymbol{x}}^{(p)T} \bar{\boldsymbol{y}}^{(q)} \right) \tag{11}$$ 式10の分解自体は数学的な恒等式として厳密に成立します。実装上は、外側の累積（スケール乗算と加算）をType1算術で行うことでType1精度を達成できます。 cuTileでの行列積カーネルの実装は以下の通りです。tilegymのmatmulカーネルとベースは同じで２つのスライス分のouter-loopが追加されています。 @ ct.kernel (num_ctas=ct.ByTarget(sm_100= 2 )) def ozaki_matmul_fused_kernel ( A_slices, # (s_a, M, K) Type2スライス B_slices, # (s_b, K, N) Type2スライス Combined_scales, # (s_a, s_b) 2^{c_x(p)+c_y(q)} のスケール行列 C, # (M, N) FP64 出力 TILE_SIZE_M: ConstInt, TILE_SIZE_N: ConstInt, TILE_SIZE_K: ConstInt, ): # タイルインデックス計算（L2キャッシュ局所性のためswizzle） bidx, bidy = swizzle_2d(M, N, TILE_SIZE_M, TILE_SIZE_N, GROUP_SIZE_M= 8 ) num_tiles_k = ct.cdiv(K, TILE_SIZE_K) # FP64最終アキュムレータ（式10の外側の累積） accumulator = ct.full((TILE_SIZE_M, TILE_SIZE_N), 0.0 , dtype=ct.float64) # 全スライスペア (p, q) をループ for p in range (num_slices_a): for q in range (num_slices_b): # FP32中間アキュムレータ（式11: Type3での丸め誤差なし計算） slice_acc = ct.full((TILE_SIZE_M, TILE_SIZE_N), 0.0 , dtype=ct.float32) # K方向のタイルループ for k in range (num_tiles_k): a_tile = ct.load(A_slices, index=(p, bidx, k), ...) b_tile = ct.load(B_slices, index=(q, k, bidy), ...) slice_acc = ct.mma(a_tile, b_tile, slice_acc) # Type2 Tensor Core MMA # スケーリングしてFP64で累積（式10） scale = ct.load(Combined_scales, index=(p, q), shape=( 1 , 1 )) accumulator = accumulator + ct.astype(slice_acc, ct.float64) * scale ct.store(C, index=(bidx, bidy), tile=accumulator) 素朴な実装と初回結果 上記のようにOzaki Schemeを実装してみます。スライス分割はホスト側のpythonで行い、各ペアのGEMMを順次実行する方式です。以下、2種類のType2で行列積を計算した結果です。 スライス数はA・Bそれぞれの分割数（$s_a \times s_b$）、GEMMsはその組み合わせで実行したGEMM回数です。TFLOPSはFP64換算のスループット、Rel ErrorはPyTorch FP64結果を基準とした相対誤差です。 TYPE2 = FP16 行列サイズ スライス数 GEMMs Split(ms) Kernel(ms) 合計(ms) TFLOPS Rel Error 1024 10×10 100 1.48 0.62 2.64 0.81 1.58e-15 2048 12×12 144 2.57 2.66 5.75 2.99 1.98e-15 4096 12×12 144 8.69 21.59 30.70 4.48 6.31e-15 8192 14×14 196 36.00 192.32 231.76 4.74 7.28e-15 16384 14×14 196 135.21 1737.14 1884.24 4.67 3.35e-15 TYPE2 = FP8 (E4M3) 行列サイズ スライス数 GEMMs Split(ms) Kernel(ms) 合計(ms) TFLOPS Rel Error 1024 15×16 240 2.22 1.11 4.03 0.53 1.64e-15 2048 16×16 256 3.48 3.11 7.33 2.34 2.27e-15 4096 16×17 272 12.30 19.17 30.86 4.45 5.69e-15 8192 17×17 289 45.64 179.99 226.64 4.85 8.88e-15 FP16はスライス数が少ない分GEMMも少なくなりますが、FP8はTensor Coreのスループットが高いため、GEMMs数が多いにも関わらず類似の性能が出ています。いずれの型でもcuTile FP64直接計算（約1 TFLOPS）を上回っていますが、PyTorchの性能には大きく劣後しています。Split処理の時間も無視できず、特に小さな行列サイズでボトルネックになっています。また、参照した論文に記載されている必要なGEMM数(スライス数）よりも多くなっている点は気になりましたが、原因はわからずでした。 最適化 初回結果を踏まえ、いくつか改善を試みました。その中で効果があった方法が以下になります。 Fast Mode（GEMMの削減） スライス数が $s$ の場合、全組み合わせで $s^{2}$ 回のGEMMが必要です。しかし、スライスインデックスが大きい組み合わせ（$i + j \geq d$）は寄与が小さいため、スキップできます。 [5]で提案されたFast Mode（Algorithm 3）では、確率的誤差限界 $|fl(AB) - AB| \leq 2\sqrt{k} \cdot u_{\text{FP64}} \cdot |A||B|$ を満たす最小の閾値 $d$ を自動決定します。 BF16の場合、典型的には $d = 9$ 程度で、GEMMは49回から39回に削減できます。さらに max_d パラメータで手動上限を設定すれば、精度とのトレードオフで計算量を調整できます。 実装としては、前述の行列積カーネルのスライスペアループに i + j >= D の条件を追加するだけです。 for i in range (num_slices_a): for j in range (num_slices_b): if i + j >= D: # Fast Mode: 寄与の小さい組み合わせをスキップ continue # ... K方向ループでMMA計算 ... Fused Split Kernel（分割の融合） 元の実装では、各スライスの計算ごとに max().item() でGPU→CPU同期が発生していました（BF16で7スライス = 7回の同期）。 改善後は、初回の max_abs 計算で1回だけ同期し、全スライスの $\sigma$ と $2^{-c_i}$（逆スケール）をCPU側で事前計算します。その後、単一カーネルで全スライスを一括計算します。 @ ct.kernel (occupancy= 4 ) def _veltkamp_split_all_slices_kernel ( x_in, # (M, N) FP64 input slices_out, # (num_slices, M, N) TYPE2 output slices sigmas, # (num_slices,) FP64 pre-computed sigma values inv_scales, # (num_slices,) FP64 pre-computed 1/scale values num_slices: ConstInt, TILE_SIZE_M: ConstInt, TILE_SIZE_N: ConstInt, ): bid = ct.bid( 0 ) # ... (タイルインデックス計算) ... # 入力タイルをロード residual = ct.load(x_in, index=(tile_m, tile_n), shape=(TILE_SIZE_M, TILE_SIZE_N), padding_mode=ct.PaddingMode.ZERO) # 全スライスをループで計算 for i in range (num_slices): sigma_tile = ct.load(sigmas, index=(i,), shape=( 1 ,)) inv_scale_tile = ct.load(inv_scales, index=(i,), shape=( 1 ,)) # Veltkamp分割 v = (residual + sigma_tile) - sigma_tile slice_tile = ct.astype(v * inv_scale_tile, slices_out.dtype) ct.store(slices_out, index=(i, tile_m, tile_n), tile=ct.reshape(slice_tile, ( 1 , TILE_SIZE_M, TILE_SIZE_N))) residual = residual - v 最適化後の結果 Fast Mode + Fused Split Kernelを適用した結果です。 TYPE2 = BF16 行列サイズ スライス数 GEMMs Split(ms) Kernel(ms) 合計(ms) TFLOPS Rel Error 1024 7×7 39 0.20 0.25 1.57 1.37 5.75e-15 2048 7×7 39 0.24 0.75 2.16 7.94 1.30e-14 4096 7×7 39 0.43 5.32 7.22 19.04 1.16e-14 8192 7×7 39 1.16 38.93 42.60 25.81 2.39e-14 16384 7×7 39 3.96 323.10 327.78 26.84 2.21e-14 TYPE2 = FP8 (E4M3) 行列サイズ スライス数 GEMMs Split(ms) Kernel(ms) 合計(ms) TFLOPS Rel Error 1024 14×14 130 0.25 0.61 2.99 0.72 3.48e-13 2048 14×14 130 1.27 1.60 6.80 2.52 3.57e-13 4096 14×14 130 2.13 9.31 15.88 8.65 3.41e-13 FP8はスライス数が多く（14×14）GEMMsも130回と多いものの、Fast ModeによるGEMM削減とFused Splitの効果で素朴な実装（4096で4.45 TFLOPS）から改善が見られます。ただしBF16と比較すると、仮数部が4ビットと少ないためスライス数が増え、GEMMs数の差（130 vs 39）がFP8のスループット優位を打ち消しており、BF16の方が総合的に有利になりました。 素朴な実装と比較すると、最適化の効果は顕著です。 Fused Split Kernel : Split時間が大幅短縮（素朴なFP16版 8192: 36.00ms → BF16最適化版: 1.16ms、約31倍） Fast Mode : GEMMsを49→39に削減（BF16の全組み合わせ比） BF16がTYPE2として最適である理由は、Tensor CoreのFP32アキュムレータとの相性にあります。BF16の仮数部は8ビットなので、2つのBF16値の積は16ビットに収まります。FP32の仮数部は24ビットあるため、TILE_SIZE_K=128個の積和（16 + log2(128) = 23 ≤ 24）が 丸め誤差なし で正確に計算できます。一方FP16（11ビット仮数部）では積が22ビットとなり、128個の累積（22 + 7 = 29 > 24）でFP32精度を超えるため、丸め誤差が発生します。 この性質により、BF16では1e-14というFP64に近い精度を維持しつつ、Tensor Coreの高いスループットを活用できています。 上記以外にも、タイルサイズの調整や、カーネル内のsplitループ方向のCTA分散も試みましたが、効果はありませんでした。本来は、プロファイラ（NVIDIA Nsight Compute）を使って、メモリ利用等解析するのが効果的ですが、Modal上ではNsightは使えないようなので断念しました。 dによる精度/速度トレードオフ d パラメータを変えて、16384×16384行列での性能と精度の変化を測定しました。BF16の結果です。 d GEMMs 合計(ms) TFLOPS Rel Error vs PyTorch FP64 9 (default) 39 327.78 26.84 2.21e-14 0.75x 8 34 298.76 29.44 2.31e-14 0.83x 7 28 238.94 36.81 3.50e-13 1.03x 6 21 189.59 46.40 4.39e-11 1.30x 5 15 136.34 64.52 4.51e-09 1.81x PyTorch FP64（cuBLAS）は同サイズで35.62 TFLOPSです。Rel ErrorはPyTorch FP64の結果を基準として計算しています。 d=7 でcuBLASと同等の速度を精度1e-13で達成し、 d=5 では1.8倍の高速化を1e-9精度で実現しています。なお、FP64 GEMM自体も浮動小数点演算の性質上、行列サイズに応じた丸め誤差は避けられないため、 d=8 （2.31e-14）程度の偏差であれば実用上十分でしょう。 まとめ cuTile Pythonの簡単な紹介とOzaki Schemeの実装を通じて、FP64行列積の高速化を試みました。BF16 Ozaki Schemeの最適化後、16384×16384行列で最大26.84 TFLOPS（d=9）を達成しました。dを調整することで精度と速度のトレードオフが可能で、d=7ではcuBLAS FP64（35.62 TFLOPS）と同等の36.81 TFLOPSを精度1e-13で達成し、d=5では64.52 TFLOPS（cuBLASの1.8倍）を1e-9精度で実現しています。 CUDAカーネルをPythonライクに書ける点で、GPUプログラミングの敷居が低くなったと感じます。 一方で、より高度な最適化やチューニングが必要な場合は、cuTile Pythonは抽象化してハード側の詳細を隠蔽している分、制約があるように感じました。今回は、行列積の例でしたが、tilegymにはtransformerの実装例があるので、次回はそちらも試してみたいと思います。 参考文献 [1] Lecture 89: cuTile (from friends at NVIDIA) [2] NVIDIA cuTile Documentation . cuTile Python. [3] NVIDIA TileGym . GPU Tile kernel development examples using cuTile. [4] Markus Höhnerbach, Paolo Bientinesi (2025). "DGEMM without FP64 Arithmetic" . arXiv:2508.00441. [5] Daichi Mukunoki, Katsuhisa Ozaki, Takeshi Ogita, and Toshiyuki Imamura (2020). "DGEMM using Tensor Cores, and Its Accurate and Reproducible Versions". ISC High Performance 2020, Lecture Notes in Computer Science, Vol. 12151. Springer, 230–248. doi:10.1007/978-3-030-50743-5_12

みなさんこんにちは！ワンキャリアで、データサイエンティストをしている伊藤です。 私は現在、求職者と企業の最適なマッチングを実現するための「レコメンドロジックの構築」や、プロダクト改善の意思決定を支える「効果検証・ダッシュボード構築」といったアナリティクス業務を担当しています。