はじめに こんにちは！ セーフィーの開発本部エンジニアリングオフィスで26卒内定者インターンをしている吉田・水野・小石川です。 今回はインターン活動の一環で2025/11/15(土)に開催された「BTCONJP2025」というカンファレンスに私たちインターン生がゼロからWebアプリを作成し展示させていただいたので、その技術的挑戦と当日のトラブル対応について振り返りたいと思います。 当日の写真 導入・背景 自己紹介 吉田🧖)開発本部で内定者インターンをしている吉田和司です！結構多趣味で最近はドローンに興味を持って資格を取ったりしました。他にもサウナが大好きで、この前は名古屋の

初めまして。 AWS上で大規模な会員向けアプリケーションの構築・運用に携わっています。 本記事では、 OpenSearchとSageMakerを組み合わせたセマンティック検索基盤の構築事例 を紹介します。 既存のOpenSearchにカスタムAIモデルを連携する構成について、検討の過程と構成例を整理しています。 また、当初のデプロイ時に直面した課題や、最終的にSageMakerを採用するに至った技術選定の背景についても触れています。 2. 現状の検索構成 先日、アプリケーション開発チームより、次のような相談を受けました。 「クーポン検索機能を高度化するカスタムAIモデル（DistilUSE）を作ったので、こちらをAWS上にデプロイして既存のOpenSearchと連携してほしい」 現状の検索機能は、 OpenSearchを中心とした比較的シンプルな構成 になっていました。 アプリケーションユーザーからのクーポン検索リクエストは、まず ECS上で稼働しているアプリケーション（コンテンツ取得API） に送信されます。このAPIが検索条件を受け取り、 OpenSearch Serviceの検索APIを呼び出すことで検索処理を実行 します。 OpenSearch側ではクーポン情報をインデックスとして保持しており、検索処理は主に キーワード検索（完全一致・部分一致） によって行われていました。検索結果はAPIを経由してアプリケーションに返却され、最終的にユーザー画面へ表示されます。 また、OpenSearch ServiceおよびECSは同一VPC内に配置されており、通信はALBを介して行われる構成です。 この構成では、検索処理の責務はすべてOpenSearchが担っており、 検索クエリの解釈や意味的な理解といった処理は行っていない状態 でした。 図1：ECS上のAPIからOpenSearchを直接呼び出し、検索結果を返す構成 3.構成検討の過程 アプリケーションチームから依頼を受け、まずは 既存の構成を大きく変えずにカスタムAIモデルを組み込めないか を検討しました。 その選択肢の一つとして、OpenSearchが提供する ML Commons 機能の利用を検討しました。 ML Commonsでは、S3に配置したモデルをOpenSearch側でロードし、検索処理の中で推論を実行する仕組みが提供されています。 この仕組みを利用すれば、推論用のサーバーを新たに用意する必要がなく、アーキテクチャを比較的シンプルに保てると考えました。 そこで、アプリケーション開発チームから受領した カスタムAIモデル（DistilUSE / tar.gz） をS3に配置し、OpenSearch Serviceからモデルのロードを試みました。 しかし、結果としてこの構成は実現できませんでした。 今回受領したカスタムモデルは、 OpenSearch Service（マネージド環境）におけるML Commonsの制約により、そのまま有効化することができない仕様だったからです。 このため、 OpenSearch内部でモデルを実行する構成は断念 し、別のアーキテクチャを検討することにしました。 図2：S3に配置したモデルをOpenSearch内部で実行する構成（検討時） ML Commons を利用できなかった理由 ML Commons の利用を見送った理由は、主に 実行環境とマネージドサービスとしての制約 にあります。 受領したカスタムAIモデルは Python（PyTorch）ベース で構築されていましたが、OpenSearch Service の ML Commons は Javaベースの実行環境 を前提としています。そのため、モデルをそのまま実行することが難しく、形式変換などの追加対応が必要となりました。 これらの対応は工数や運用面の負荷が大きく、今回は現実的ではないと判断しました。 また、OpenSearch Service はマネージドサービスであるため、外部から持ち込んだカスタムモデルを柔軟に実行すること自体に制約があります。 以上の検討を通じて、 検索エンジンに計算処理を持たせる構成は適切ではない と判断し、計算処理は専用の実行基盤に切り出す方針としました。 次章では、この方針を踏まえて採用した SageMakerを用いた構成 について説明します。 4.最終アーキテクチャ 検討の結果、 計算処理は計算専用のリソースに切り出すべき と判断し、カスタムAIモデルの実行基盤として Amazon SageMaker を採用しました。 最終的なアーキテクチャを図3に示します。 本構成では、検索機能そのものは引き続きOpenSearchが担い、検索クエリやドキュメントのベクトル化といった 計算負荷の高い処理をSageMakerの推論エンドポイントに委譲 しています。 これにより、OpenSearchに過度な処理を持たせることなく、役割を明確に分離した構成とすることができました。 図3：SageMaker推論エンドポイントと連携した最終構成 5. SageMaker採用の背景 ここは今回の構築において、最も検討に時間を要したポイントです。 「Amazon Bedrock（サーバーレス）を利用すれば、より簡単に実装できたのではないか」 という選択肢もありましたが、最終的には SageMakerを用いた自前構築が最適 と判断しました。 理由は大きく2点あります。 理由1：カスタムモデルの利用要件 Amazon Bedrockは、あらかじめ用意された基盤モデルをAPI経由で利用できるサービスです。 一方で、利用できるモデルはBedrock側で提供されているものに限定されます。今回のケースでは、アプリケーション開発チームより 特定のカスタムモデル（DistilUSE）を利用したい という明確な要件がありました。 この要件を満たすには、モデルや実行環境を自由に構成できるSageMakerを採用する必要がありました。 理由2：スケールとコストの観点 今回対象としたシステムは、 全国規模のECサイト であり、検索リクエストが高頻度かつ継続的に発生することが想定されます。 BedrockのようなAPIベースのサービスは従量課金であるため、 リクエスト数の増加に伴うコストの増大 レート制限（スロットリング）による影響 といった点が懸念されました。 一方、SageMakerの推論エンドポイントは、 インスタンスをプロビジョニングすることで、一定のコストで安定した処理能力を確保 できます。大規模なトラフィックを前提とした場合、コストの見通しやすさと安定性の面でSageMakerの方が適していると判断しました。 6. 構築時のポイント 構築にあたっては、以下の点を重視しました。 役割の分離 検索処理はOpenSearch、ベクトル化などの計算処理はSageMakerとし、各コンポーネントの責務を明確に分離しました。 閉域網でのセキュリティ確保 SageMakerの推論エンドポイントはVPC内に配置し、OpenSearchとはVPCエンドポイント経由で通信させています。 データがインターネットを経由しない構成としました。 アプリケーションへの透過性 OpenSearchのIngestionパイプラインを活用することで、アプリケーション（ECS）側は裏でSageMakerが動作していることを意識せず、従来通り検索リクエストを送信するだけで利用可能としています。 なお、AIモデルのデプロイには AWS公式のDeep Learning Container を使用し、アプリケーションチームから受領した model.tar.gz （モデル本体・設定・辞書）をそのまま推論環境として利用しています。 7. まとめ 本構成により、コストや運用リスクを抑えつつ、柔軟な検索機能を実現できるアーキテクチャを構築することができました。 また、補足として、S3に配置したシノニム辞書をOpenSearchと連携させることで、検索精度の底上げも行っています。 今回の取り組みを通じて、AI機能の実装においては 用途や規模に応じたサービスの使い分けが重要 であると改めて感じました。 Amazon Bedrock 手軽に始めたい場合や、小〜中規模でモデルに強い制約がないケース Amazon SageMaker 特定のモデルを利用したい場合や、大規模アクセスを前提とするケース 単に動作する仕組みを作るだけでなく、将来的なスケールや運用まで見据えた設計を行うことの重要性を学ぶ良い機会となりました。

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