Apple Developer onepage/WWDC26
#WWDC26#Metal#TensorOps#CoreAI2026-06-08 · 16:13

用 Metal tensors 把自訂 ML operation 做到 Apple Silicon 底層

這場 session 是給 ML framework、Core AI model 與 Metal app 團隊的底層優化路線圖：用 TensorOps 處理 quantized data、multi-plane tensors、FlashAttention 類 fused operation，最後把 custom Metal kernel 接回 Core AI 模型轉換流程。

// 一句話結論

當 Core AI 或 MLX 的高階 API 不夠時，Apple 希望你用 TensorOps 寫可攜的 Metal kernel：直接吃 quantized tensors、少搬資料，並把自訂 operation 封進 Core AI asset。

Apple Developer · WWDC26 · Video ID: toO9j5NKg0cApple Developer 官方頁
Content map

這支影片在拆四個底層 ML kernel 問題

定位

什麼時候值得下到 Metal

高階 Core AI / MLX 適合部署模型；當你要實作研究新 op、貢獻 MLX / llama.cpp，或在 Metal app 裡跑自訂模型，就需要 TensorOps。

Quantization

不要先解壓再算

TensorOps 支援 4/8-bit integer、4/8-bit float、2-bit integer 與 MX scaling formats，目標是直接用 quantized tensor 進 matmul，降低 memory bandwidth。

Memory

Multi-plane tensor 把 data 和 scale 綁在一起

iOS / macOS 27 起，單一 MTLTensor 可同時代表 quantized data plane 與 FP8 E8M0 scale plane，避免散落的 buffer metadata。

Fusion

FlashAttention 是 TensorOps 高階能力展示

用 cooperative tensors、row reductions、map_iterator 與 direct matmul input，把 QK、SoftMax、V 的工作融合，減少 threadgroup memory round trip。

Implementation flow

從 quantized tensor 到 Core AI custom op 的工作流

1

定義 quantized MTLTensorDescriptor

在 descriptor 指定 MetalFloat8、integer 量化型別、usage 與 dimensions；新 data types 有額外 alignment 要求，必須查 Metal 文件。

2

為 scales 建 auxiliary plane

使用 MTLTensorAuxiliaryPlaneDescriptor 描述 FP8 E8M0 scale plane 與 block factors，讓 data、scale、metadata 一起進 tensor object。

3

在 shader 中切 tile 並交給 TensorOps matmul

用 threadgroup ID slice input/output tensors；大多情況直接把 quantized tensor 餵給 TensorOps，讓它自動處理 dequantization 與 hardware acceleration。

4

用 cooperative tensors fusion 複雜 op

FlashAttention 類 op 需要讓 SIMD group 擁有完整 row，做 row max / SoftMax，再把 cooperative tensor 直接當第二次 matmul input。

5

把 Metal kernel 接回 Core AI conversion

在 Python 轉換流程註冊 TorchMetalKernel，替換 PyTorch attention implementation，再 export 成含 custom kernel 的 Core AI asset。

重點 1：TensorOps 是 Apple Silicon GPU 的 ML kernel 抽象層

影片把 Apple ML stack 分成高階 Core AI / MLX、低階 Metal Performance Shaders / Metal Performance Primitives，以及 TensorOps。TensorOps 是 Metal Shading Language API，讓你寫 tensor operation 時自動使用不同 Apple Silicon GPU generation 的硬體加速，特別是 M5 shader core 裡的 Neural Accelerator。

這不是要一般 app 團隊全部改寫 kernel；它適合需要自訂 operation、追新研究 operator、或把 PyTorch / Core AI 模型裡的瓶頸 fused 成 Metal kernel 的團隊。

Quantized data 應直接進 TensorOps若格式受支援，直接把 quantized tensor 傳入 matmul，讓 TensorOps 自動 dequantize 並使用硬體加速。
Scale plane 是 tensor 的一部分FP8 E8M0 block-wise scale factor 可作為 auxiliary plane，與 data plane 一起被 slice 與傳遞。
自訂格式才手動 dequantize若必須處理自訂格式，可 dequantize 到 cooperative tensor，避免多一次 threadgroup memory load/store。
Alignment 是 release 風險2-bit、4-bit、FP8 與 MX scaling formats 有額外 alignment 要求；錯誤常不會在高階 API 層被自動修正。

重點 2：FlashAttention 示範的是「少搬資料」的設計哲學

Attention 原本需要 Q x K、row-wise SoftMax、再乘 V。FlashAttention 的核心是把這些 operation fuse 成單一 kernel，避免把中間矩陣反覆寫回 memory。影片示範使用 execution_simdgroup scope 讓每個 SIMD group 擁有完整 rows，接著用 cooperative tensor 儲存 QK 中間結果並做 row reduction。

01
先確保 simd group mapping 對齊演算法需求
SoftMax 需要 row reductions，所以每個 SIMD group 應持有完整 row，否則會被迫跨 group 溝通。
02
用 cooperative tensors 承載中間結果
它把 storage 分散在參與 thread 的 private memory 中，可直接餵給 TensorOps，而不是先落到 threadgroup memory。
03
重用前先檢查 layout compatibility
不是每個 cooperative tensor 都能直接當 matmul input；需用 is_compatible_as_left_input / right input 判斷，失敗時才 fallback store/reload。

重點 3：Core AI integration 讓自訂 kernel 進入模型資產

最後 demo 將 custom FlashAttention kernel 整合進 SAM3 image segmentation 模型。流程是在 Python 端把 Metal kernel body 註冊成 TorchMetalKernel，替換 Hugging Face attention implementation，載入模型並從 PyTorch export 成最佳化 Core AI asset。推論時，模型可用文字提示標出圖片中的車輛區域，證明 custom attention kernel 已進入模型流程。

風險與遷移注意事項

高階 framework 仍是預設入口若 Core AI / MLX 已能達成目標，不要過早下到 Metal；底層 kernel 會增加 QA、profiling 與硬體相容性成本。
量化格式要和模型轉換管線一致host side tensor descriptor、shader type alias、scale block size 與 PyTorch export 都要對齊。
Threadgroup fallback 要測cooperative tensor layout 不相容時會走 store/reload；這條 fallback path 也要測效能與正確性。
Profiling 要看 memory bandwidth，不只看 FLOPS影片的主軸是減少資料搬移與中間結果寫回；單看 kernel 計算量可能抓不到真正瓶頸。

對 ML / Metal 團隊的實作啟發

01
先用 reference PyTorch 保護正確性
custom Metal kernel 應該和 PyTorch reference op 一起放進轉換與 regression tests。
02
把 quantized tensor descriptor 做成可審查 config
data type、auxiliary plane、block factors、alignment 不要散在 shader 與 host code 多處。
03
用 Instruments / Metal tooling 驗證資料搬移
FlashAttention 類融合的價值在減少中間結果寫回；profiling 應對照 threadgroup memory、cache 與 device memory 行為。

來源

主要來源為 Apple Developer 官方影片頁、transcript、chapters 與 code snippets；YouTube 影片頁用於 iframe 與發布 metadata。官方頁：Optimize custom machine learning operations with Metal tensors - WWDC26。YouTube：WWDC26: Optimize custom machine learning operations with Metal tensors | Apple。

// TL;DR