Metalテンソルによるカスタム機械学習オペレーションの最適化

Metalテンソルによるカスタム機械学習オペレーションの最適化

Metal Tensor APIとMetal Performance Primitives（MPP）Tensor Opsライブラリを使用すると、パワフルな機械学習パフォーマンスを実現できます。Apple M5およびA19 GPUのNeural Acceleratorを活用する、ポータブルなオペレーションの構築方法を習得しましょう。Core AIアプリ向けのカスタム機械学習カーネルの構築方法と、量子化データ形式やGPUメモリ最適化を効果的に扱う方法も確認します。

関連する章
- 0:00 - Introduction
- 0:21 - Apple's ML software stack
- 2:25 - Managing quantized data
- 4:23 - Multi-plane tensors
- 5:17 - Quantized matrix multiplication
- 9:31 - Building advanced ops
- 13:35 - Integrating custom ops into Core AI
- 15:25 - Next steps
リソース
関連ビデオ

Tech Talks
- M5とA19のGPUで機械学習のワークロードを加速
WWDC25
- Metal 4による機械学習とグラフィックスの統合
こんにちは、私はShiyaoです。 GPUソフトウェアエンジニアです。
本日は、Metal テンソルの探求をご案内し、TensorOpsを使った最適化されたカスタムMLカーネルの書き方をご紹介します。
Apple プラットフォームはソフトウェアスタックのあらゆる層で MLモデルを実行するためのファーストクラスのサポートを提供しています。 Core AIやMLXといった高レベルフレームワークは最小限のコードでモデルを簡単にデプロイでき、 Metal Performance Shadersのような低レベルAPIは高性能なMetalカーネルへのアクセスを提供します。これらの層はすべて、 Metal Performance Primitives と TensorOps ライブラリが提供する低レベルアクセラレーション上に構築されています。 Metalレベルで作業したい理由はいくつかあります。 ML研究は急速に進歩しており、 Core AIなどの高レベルフレームワークにプラグインできるカスタム操作を実装したい場合があります。また、MLXやllama.cppなどの MLフレームワークに貢献している場合や、 Metal ベースのアプリケーションに取り組んでいる場合にも Metalカーネルを書く必要があるかもしれません。始める最も簡単な方法は TensorOps ライブラリを使うことです。 TensorOps は Metal Shading Language の API で、 GPU上のテンソル演算を高速化します。行列積や畳み込みなどが含まれます。利用可能なハードウェアアクセラレーションを自動的に活用し、すべての Apple Silicon GPU世代にわたってハードウェア世代間の違いを気にする必要がありません。特に、M5チップファミリーのニューラルアクセラレーターをフル活用します。
ニューラルアクセラレーターは M5の新しいハードウェアブロックで、各シェーダーコアに直接配置されています。他のGPUパイプラインと並んで配置され、密度の高い計算バウンドワーク（LLMのプリフィル段階など）を高速化するために設計されています。
関連セッションをご覧ください。 TensorOpsを使い始める基礎を学べます。このセッションでは、その基礎の上に構築し、量子化データを扱うためのベストプラクティスから始めます。次に、FlashAttentionのような高度なカスタム操作の構築方法をご紹介します。
最初のトピック、量子化データの扱い方から始めましょう。
ご存知のとおり、最先端の機械学習モデルは大型化しています。推論段階は通常メモリ帯域幅がボトルネックとなるため、重みを圧縮することが必要になります。モデルをメモリに収め、メモリ帯域幅を節約するためです。
重みを圧縮する標準的なアプローチは量子化です。アイデアはシンプルです — より高精度な重みを取り、低精度のデータ型に縮小します。例えば、16ビットの半精度の重みをわずか4ビットに圧縮できます。これらの量子化された重みはスケールファクターとペアになっており、量子化された値を計算時に元の範囲にスケールバックできます。
16ビットと32ビットの浮動小数点型に加えて、 TensorOps は量子化データ型をネイティブにサポートするようになりました。 macOS と iOS 26のアップデートで 4ビットと8ビットの整数型のサポートを追加し、 macOS と iOS 27ではさらに多くのデータ型にサポートを拡張しています。これには4ビットと8ビットの浮動小数点型と 2ビット整数型が含まれます。アプリの量子化テンソルを作成して TensorOps に渡すだけで、利用可能なハードウェアアクセラレーションを自動的に活用します。
量子化データ型を持つテンソルの作成は、通常のテンソルの作成と非常に似ています。他のテンソルと同様にディスクリプタのプロパティを設定し、量子化された dataType を指定するだけです。次に、Metal デバイス上で newTensorWithDescriptor を呼び出してテンソルを作成します。
これが量子化された要素データを保存する方法です。次に、スケールファクターについて話しましょう。 macOS と iOS 27では、
単一の MTLTensor オブジェクトでスケールを表現できるようになりました。テンソルの量子化データと並んで追加のスケールプレーンとして扱われます。このプレーンは一般的なFP8 E8M0 ブロックワイズスケールファクター形式をサポートします。スケールプレーンの各要素はデータプレーンの要素ブロックに適用されます。スケールプレーンの宣言はテンソルの宣言と似ています。まず、スケールプレーンのディスクリプタオブジェクトを作成します。次に、dataTypeと blockFactorsを設定します。最後に、このプレーンがスケール用であることを指定する補助プレーンマップを作成します。
次に、補助プレーンマップを元の tensorDescriptor に付加するだけです。量子化データ、スケール、メタデータはすべて単一のテンソルオブジェクトにパックされます。
では、これを実践しましょう。基本的な行列積カーネルを拡張して量子化をサポートします。行列積は機械学習ワークロードの中核となる演算です。例えば、LLMは推論中に何百万もの行列積を実行します。
M5機械学習トークで、 TensorOpsを使った高性能な行列積カーネルの書き方の基礎を説明しました。基本的なアプローチは、入力行列をより小さなタイルにスライスし、 TensorOps を使ってタイルごとの行列積を実行することです。これにより並列処理が最大化され、データがキャッシュに保持されます。
量子化を使ってメモリトラフィックをさらに削減し、より大きなモデルをメモリに収めることができます。カーネル内では、テンソルをバインドする前に型エイリアスを先に定義すると便利です。ここでは、 fp8_e8m0_ データ型を持つスケールファクタープレーンを宣言し、ブロックサイズを32×1とします。つまり、データプレーンの 32要素ごとに scales_plane の単一要素が共有されます。次に、FP8データ型と scales_plane を指定した完全なテンソル型を宣言します。これらのテンソルをバッファバインディングポイントにバインドするだけです。カーネルはその後、ホスト側で割り当てたテンソルにアクセスできます。あるいは、ホスト上で完全な MTLTensor を作成したくない場合は、シェーダーのスタック上で一時テンソルを作成することもできます。構文はほぼ同じで、 tensor_handle タグを tensor_inline に交換するだけです。次に、バッファポインタとその他のメタデータをテンソルコンストラクタに渡してスタック上にテンソルを作成します。
先ほど述べたように、より良い並列処理のために多くのスレッドグループに問題を分散させます。まず、各スレッドグループのタイルをスライスし、次に TensorOps で乗算を実行します。
これを行うには、スレッドグループIDを使用して入力テンソルと出力テンソルに対して単純に slice を呼び出します。データプレーンとスケールプレーンはどちらもブロックサイズに従って同時にスライスされます。量子化テンソルを使った行列積のセットアップは通常のテンソルと同じです。まず、matmul2d_descriptor をセットアップし、タイルサイズとその他のパラメータを指定します。次に matmul2d op を作成し、スレッドグループ内の simdgroup 数を指定します。次に量子化テンソルを渡すだけで、 TensorOps が逆量子化を処理してくれます。
ほとんどの場合、量子化データを直接 TensorOps に渡すべきで、利用可能なハードウェアアクセラレーションを自動的に活用できます。ただし、カスタムフォーマットを逆量子化する必要がある場合も、 TensorOps がカバーしています。最もシンプルなアプローチは、各スレッドがデバイスメモリから量子化データのチャンクを読み込み、スレッドグループメモリで f16 値に逆量子化することです。その後、インラインスレッドグループテンソルとして TensorOps に渡せます。ただし、このアプローチでは追加のロードとストアが必要で、スレッドグループメモリを経由します。理想的には、すべてのデータをスレッドレジスタに保持したいところです。データをコオペラティブテンソルに逆量子化することで実現でき、 matmul2d op への入力として渡せるようになりました。コオペラティブテンソルはストレージをスレッドプライベートメモリに matmul 操作に参加するスレッド間で分散します。そのため、量子化テンソルを直接使えない場合でも、スレッドグループメモリへの往復を省くことができます。まとめると — Metalテンソルは幅広い量子化データ型をネイティブサポートし、新しいMXスケーリングフォーマットと iOS および macOS 27で登場する E8M0スケールファクターも含まれます。これらの新しいデータ型は大きなデータ型と比べて追加のアライメント要件があることに注意してください。詳細は Metal ドキュメントを必ずご確認ください。
さらにレベルを上げましょう — TensorOps を使って完全なより複雑なカスタム操作を構築します。アテンションはすべてのトランスフォーマーネットワークの中核で、 LLMも含まれます。アテンションを計算するには、まずQとKと呼ばれる 2つの行列を掛け合わせます。次に、中間行列の行に対するリダクションを使ってSoftMaxを計算します。
最後に、Vと呼ばれる 3番目の行列を掛けます。人気のある FlashAttention アルゴリズムはこれらの操作をすべて融合させ、単一のカーネルにします。
TensorOps でこれを実装するには、まずカスタム simdgroup マッピングを設定する必要があります。各 simdgroup が中間行列の完全な行を所有するようにします。これにより、simdgroup間でデータを交換せずに SoftMax を計算できます。 execution_simdgroup 操作スコープを使ってこれを実現できます。つまり、各 simdgroup が並列で独立した行列積を実行します。
simdgroup ID を使って入力タイルをスライスできます。中間行列を保存するためにコオペラティブテンソルを使い、メモリへの書き込みなしに次のステップへの入力として使用できます。結果に対して SoftMax を計算します。
そのために、いくつかのリダクションを計算する必要があります。コオペラティブテンソルに対して。 TensorOps にはこれを助ける reduce_rows 関数が含まれています。スレッドがお互いにデータを交換して各行の最大値を計算します。結果は別のコオペラティブテンソルに返されます。
設定してみましょう。まず、リダクション出力を保存するコオペラティブテンソルを作成します。次に、ソースとデスティネーションを reduce_rows 関数に渡します。ここでは、負の無限大の初期値を使った max reduction_operation を使用します。
これら2つのコオペラティブテンソルは異なる形状を持っているため、マッピングを助けるために、 TensorOps には map_iterator 関数も含まれています。 2Dテンソルの要素を指すイテレータが与えられると、リダクションデスティネーションの対応する要素を指すイテレータを返します。
まず、イテレータを使って 2Dコオペラティブテンソルのループを設定します。次に map_iterator を呼び出して、各要素を対応する行の最大値にマッピングします。最後に、これらのイテレータを逆参照して SoftMax を計算し、結果をコオペラティブテンソルに書き戻します。
これで、このコオペラティブテンソルを V と掛け合わせる準備ができました。 macOS 26では、まずスレッドグループメモリに保存する必要がありましたが、コオペラティブテンソルを matmul 操作への直接入力として使用できるようになりました。
これを行うには、 get_left_input_cooperative_tensor メソッドを呼び出し、ソースのコオペラティブテンソルを引数として渡します。結果を2番目の matmul 操作への入力として渡せます。注意すべき点があります：すべてのコオペラティブテンソルが入力として再利用できるわけではありません。データ型やその他の要因によってレイアウトが異なる場合があります。そのため、実行前に is_compatible_as_left または right_input メソッドを呼び出して互換性を確認してください。 true が返れば、そのまま進められます。そうでない場合は、データをスレッドグループメモリを通じて保存・再読み込みして正しいレイアウトに変換する必要があります。いずれの場合も、 op.run の呼び出しは同じです。これらが TensorOps を使って高度な操作を構築するために必要な主要な TensorOps 機能です。 FlashAttention のような操作に使えます。この操作の構築方法を説明したので、 Core AI を使った実際のモデルでどのように動作するか見てみましょう。
Core AI は Python 開発者向けのツールを提供し、 PyTorch モデルを Core AI モデルに変換できます。カスタム Metalカーネルのサポートも含まれます。 "Deep Dive into Core AI Model authoring and Optimization" セッションをご覧ください。 Metal カーネルを Core AI モデルに統合する方法の詳細が説明されています。
そのセッションで説明された手順に従って、カスタムの FlashAttention カーネルを Sam3 画像セグメンテーションモデルに統合しました。カスタムアテンションカーネルの本体を Python で文字列として定義し、ここに示す TorchMetalKernel オブジェクトを登録します。
次に、デフォルトの Hugging Face アテンション実装をここに示す、カーネルを呼び出す実装に置き換えます。
最後に、Hugging Face からモデルを読み込み、最適化された Core AI アセットとして PyTorch からエクスポートします。エクスポートには少し時間がかかります。
これで推論の準備ができました。 Sam3 はプロンプト可能なコンセプトセグメンテーションを実行し、モデルに画像とテキストを提供すると、画像内のオブジェクトの位置を示すセグメンテーションマスクで応答します。ここでは、この画像で車を含むすべてのピクセルにラベルを付けるようモデルに指示しています。
では、セグメンテーションを実行します。
最終結果を見ると、モデルが画像を正しくセグメント化していることがわかります。車が青でハイライトされており、アテンションカーネルが期待通りにモデルに完全に統合されています。
本日は、Apple Silicon で最適化されたカスタム ML カーネルを構築するためのすべてのツールを説明しました。量子化データ型から、高度な TensorOps 機能（コオペラティブテンソルやリダクションなど）、 Core AI との統合まで。さらに学ぶには、完全な API リファレンスについては Metal Performance Primitives のドキュメントを参照し、パフォーマンス最適化のガイドラインはプログラミングガイドをご覧ください。 TensorOps のサンプルコードもダウンロードできます。ここでは説明しきれなかった詳細を確認できます。また、関連セッションもぜひご確認ください。 Core AI と Metal について詳しく学べます。ありがとうございました！

// Creating a tensor with a quantized data type from device

#define RANK 2

MTLTensorDescriptor *tensorDesc = [MTLTensorDescriptor new];

tensorDesc.dataType = MTLTensorDataTypeMetalFloat8E4M3;
tensorDesc.usage = MTLTensorUsageCompute;

NSInteger dimensions[RANK] = {NumCols, NumRows};
tensorDesc.dimensions = [[MTLTensorExtents alloc] initWithRank:RANK values:dimensions];

NSError *err = nil;
id <MTLTensor> tensor = [device newTensorWithDescriptor:tensorDesc error:&err];

4:48 - Declare a multi-plane tensor with scale factors

// Creating a tensor with a scales auxiliary plane from device

#define RANK 2

MTLTensorAuxiliaryPlaneDescriptor *planeDesc = [MTLTensorAuxiliaryPlaneDescriptor new];
planeDesc.dataType = MTLTensorDataTypeMetalFloat8UE8M0;

NSInteger blockFactors[RANK] = {32, 1};
planeDesc.blockFactors = [[MTLTensorExtents alloc] initWithRank:RANK values:blockFactors];

MTLTensorAuxiliaryPlaneDescriptorMap *auxiliaryPlanes =
    [MTLTensorAuxiliaryPlaneDescriptorMap new];
[auxiliaryPlanes setDescriptor:planeDesc forPlane:MTLTensorPlaneTypeScales];

MTLTensorDescriptor *tensorDesc = [MTLTensorDescriptor new];
tensorDesc.dataType = MTLTensorDataTypeMetalFloat8E4M3;
tensorDesc.usage = MTLTensorUsageCompute;

NSInteger dimensions[RANK] = {NumCols, NumRows};
tensorDesc.dimensions = [[MTLTensorExtents alloc] initWithRank:RANK values:dimensions];
tensorDesc.auxiliaryPlanes = auxiliaryPlanes;

NSError *err = nil;
id <MTLTensor> tensor = [device newTensorWithDescriptor:tensorDesc error:&err];

6:07 - MSL type aliases for an MXFP8 tensor handle

// Type aliases for a MXFP8 multi-plane tensor handle

#include <metal_tensor>

using namespace metal;

using scales_plane = tensor_blockwise<tensor_plane_scales,
                                      device metal_fp8_ue8m0_format,
                                      32, 1>;

using mxfp8_tensor = tensor<device metal_fp8_e4m3_format,
                            dextents<int, 2>,
                            tensor_handle,
                            scales_plane>;

kernel void matmul(mxfp8_tensor matrixA [[buffer(0)]],
                   mxfp8_tensor matrixB [[buffer(1)]],
                   tensor<device half, dextents<int, 2>> matrixC [[buffer(2)]])
{
    // ...
}

6:51 - Declare an inline MXFP8 tensor on the stack

// Type aliases for a MXFP8 multi-plane tensor inline

#include <metal_tensor>

using namespace metal;

using scales_plane = tensor_blockwise<tensor_plane_scales,
                                      device metal_fp8_ue8m0_format,
                                      32, 1>;

using mxfp8_tensor_inline = tensor<device metal_fp8_e4m3_format,
                                   dextents<int, 2>,
                                   tensor_inline,
                                   scales_plane>;

// Construct tensor on the stack from buffer pointers
mxfp8_tensor_inline matrixA(dataBufferA,
                             dextents<int, 2>(K, M),
                             array<int, 2>({ 1, K }),
                             scales_plane(scalesBufferA));

7:19 - Slice tensors and run a quantized matmul

// Slice the tensors to extract the relevant tile
auto tA = matrixA.slice(0, tgid.y * TILEM);
auto tB = matrixB.slice(tgid.x * TILEN, 0);
auto tC = matrixC.slice(tgid.x * TILEN, tgid.y * TILEM);

// Set up the matmul descriptor
constexpr auto descriptor = matmul2d_descriptor(TILEM,                  // M
                                                TILEN,                  // N
                                                dynamic_length_v<int>,  // K
                                                false,   // Left matrix transposed
                                                false);  // Right matrix transposed

matmul2d<descriptor, execution_simdgroups<4>> op;

// Run the op — TensorOps handles dequantization automatically
op.run(tA, tB, tC);

10:27 - Set up simdgroup-scoped QxK multiplication

// Setup QxK matrix multiplication op
constexpr auto mul_qk_op_desc = matmul2d_descriptor(/* ... */);
matmul2d<mul_qk_op_desc, execution_simdgroups> mul_qk_op;

// Slice Q, K, V
auto tQSlice = tQ.slice<D, ROWS_PER_SIMD>(0, sgid * ROWS_PER_SIMD);
auto tKSlice = tK.slice<D, BK>(0, k);
auto tVSlice = tV.slice<D, BK>(0, k);

// Create cooperative tensor to store tile of QxK
auto ctQK = mul_qk_op.get_destination_cooperative_tensor<decltype(tQSlice),
                                                         decltype(tKSlice),
                                                         float>();

// Multiply QxK
mul_qk_op.run(tQSlice, tKSlice, ctQK);

11:18 - Compute row-wise reduction for SoftMax

// Create a cooperative tensor to store row reduction output
auto ctTileRowMax = mul_qk_op.get_row_reduction_destination_cooperative_tensor<
                        decltype(tQSlice),
                        decltype(tKSlice),
                        float>();

// Compute max over each row of QxK tile
reduce_rows(ctQK, ctTileRowMax, reduction_operation::max, -INFINITY);

11:56 - Compute element-wise SoftMax with map_iterator

// Iterate over elements of QxK tile
#pragma clang loop unroll(full)
for (auto it = ctQK.begin(); it != ctQK.end(); it++) {
    // Fetch row max corresponding to this element
    auto row_it = ctRowMax.map_iterator(it);

    // Subtract row max from each element and compute exponent
    *it = exp(*it - *row_it);
}

12:33 - Reuse cooperative tensor as matmul input

constexpr auto mul_sv_op_desc = matmul2d_descriptor(/* ... */);
matmul2d<mul_sv_op_desc, metal::execution_simdgroup> mul_sv_op;

if (mul_sv_op.is_compatible_as_left_input<float, half, float>(ctQK)) {
    // Directly reuse cooperative tensor as input
    auto ctQKIn = mul_sv_op.get_left_input_cooperative_tensor<float, half, float>(ctQK);
    mul_sv_op.run(ctQKIn, tVSlice, ctO);
} else {
    // Store and reload through threadgroup memory if layout is not compatible
    ctQK.store(tgTensor);
    simdgroup_barrier(mem_flags::mem_threadgroup);

    auto ctQKIn = mul_sv_op.get_left_input_cooperative_tensor<float, half, float>();
    ctQKIn.load(tgTensor);
    mul_sv_op.run(ctQKIn, tVSlice, ctO);
}

- 0:00 - Introduction
- Overview of how Metal tensors and TensorOps enable you to write optimized custom ML kernels on Apple Silicon.
- 0:21 - Apple's ML software stack
- A tour of Apple's ML software stack, from high-level frameworks like Core AI and MLX down to Metal Performance Shaders, Metal Performance Primitives, and the TensorOps library — and why you might want to work at the Metal level.
- 2:25 - Managing quantized data
- How quantization reduces memory bandwidth requirements for large models, and the new quantized data types natively supported by TensorOps, including MX scaling formats.
- 4:23 - Multi-plane tensors
- How a single MTLTensor object can now represent both quantized element data and scale factors as separate planes, and how to configure multi-plane tensor descriptors in your Metal shaders.
- 5:17 - Quantized matrix multiplication
- How to extend a tiled matrix multiplication kernel to support quantized inputs, including binding scales planes, using inline tensors, slicing with threadgroup IDs, and handling custom dequantization formats.
- 9:31 - Building advanced ops
- How to implement Flash Attention with TensorOps, covering custom SIMD group mappings, cooperative tensors, row reductions, SoftMax, and the new API for passing cooperative tensors directly as matrix multiplication inputs — eliminating the threadgroup memory round-trip.
- 13:35 - Integrating custom ops into Core AI
- How to integrate a custom Metal TensorOps kernel into a Core AI application, using Core AI's Python tools to convert PyTorch models and plug in custom Metal operations.
- 15:25 - Next steps
- A summary of the TensorOps features covered — quantized types, multi-plane tensors, Flash Attention, and Core AI integration — with pointers to sample code and related sessions on Core AI and Metal.

「今すぐ始める」を詳しく見る

最新情報

プラットフォームを詳しく見る

特集

テクノロジーを詳しく見る

特集

コミュニティを詳しく見る

特集

ドキュメントを詳しく見る

リリースノート

ダウンロードを詳しく見る

特集

サポートを詳しく見る

特集

クイックリンク

関連する章

リソース

関連ビデオ

Tech Talks

WWDC25