Optimize custom machine learning operations with Metal tensors

Highlight

Metal Tensor API 和 MPP TensorOps 库在 iOS/macOS 27 中新增了对量化数据类型（FP8、4-bit/2-bit 整数）的原生支持，并允许 Cooperative Tensor 直接作为矩阵乘输入，开发者可以在 GPU 上实现 FlashAttention 这类复杂算子，无需在 Threadgroup Memory 中来回搬运数据，同时自动利用 M5/A19 芯片内置的 Neural Accelerator 硬件加速。

核心内容

大模型推理的瓶颈已经从算力转向了内存带宽。一个 70B 参数的模型，即便用 16-bit 半精度存储，也需要 140GB 显存——这远远超出了任何移动设备的容量。量化（Quantization）是唯一的出路：把 16-bit 权重压缩到 8-bit、4-bit 甚至 2-bit，既能减小模型体积，又能减少推理时的内存流量。

但量化在 GPU 上的实现一直是个苦差事。数据和 Scale Factor（比例因子）通常放在两个独立的 Buffer 里，Shader 里要手动算索引去读 Scale，再反量化回 FP16 才能参与计算。中间结果如果写回 Threadgroup Memory 再读出来，带宽和延迟根本受不了。FlashAttention 这类算子为了省掉内存往返，不得不在寄存器里倒腾数据，SIMD group 之间的同步写起来像走钢丝。

Apple 在这次 WWDC 上给出了系统性的解决方案。首先，MTLTensor 现在可以把量化数据和 Scale Factor 打包成同一个对象，通过 Auxiliary Plane（辅助平面）管理。Shader 端用 tensor_blockwise 声明类型后，TensorOps 的 matmul2d 会自动处理反量化，并调用 M5 芯片 Neural Accelerator 的硬件加速。其次，Cooperative Tensor（数据分布在 SIMD group 寄存器中的张量）现在可以直接作为下一次矩阵乘的输入，彻底消灭了 Threadgroup Memory 的往返搬运。

这套 API 的目标很明确：让写自定义 ML 算子和调用高层框架一样简单，同时性能不打折扣。无论是给 Core AI 写自定义算子，还是给 MLX、llama.cpp 贡献内核，都可以直接受益。

量化数据类型的原生支持

（03:16）TensorOps 在 iOS/macOS 26 的更新中加入了 4-bit 和 8-bit 整数类型的支持，iOS/macOS 27 进一步扩展到了 4-bit/8-bit 浮点类型和 2-bit 整数类型。开发者只需要在创建 MTLTensor 时指定量化 dataType，TensorOps 就会自动利用可用的硬件加速。

多平面张量：数据和 Scale Factor 一体化

（04:21）在 iOS/macOS 27 中，单个 MTLTensor 对象可以通过 Auxiliary Plane 同时容纳量化数据和 Scale Factor。Scale Plane 支持业界常用的 FP8 E8M0 block-wise 格式，每个 Scale 元素作用于数据平面中的一个块。Host 端配置好 MTLTensorAuxiliaryPlaneDescriptor 后，数据、Scale 和元数据会打包进同一个张量对象，Shader 端直接按块读取即可。

Cooperative Tensor 直连矩阵乘

（12:20）以前实现 FlashAttention 时，QxK 的中间结果必须先从 Cooperative Tensor 写回 Threadgroup Memory，再读出来传给下一次矩阵乘。iOS/macOS 27 新增了 get_left_input_cooperative_tensor 方法，允许 Cooperative Tensor 直接作为 matmul2d 的输入。数据从”寄存器 -> 共享内存 -> 寄存器”变成了”寄存器 -> 寄存器”，对访存密集型算子是实打实的倍数级提升。

详细内容

创建带 Scale Factor 的量化张量

（03:53）Host 端创建量化张量的流程和普通张量几乎一样，区别在于指定量化 dataType，并附加 Auxiliary Plane 描述 Scale Factor：

#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];

关键点：

• MTLTensorDataTypeMetalFloat8E4M3 指定数据平面使用 FP8 E4M3 格式存储
• dimensions 用 MTLTensorExtents 包装，支持任意秩（rank）
• 错误通过 NSError 返回，创建失败时 tensor 为 nil

如果张量需要 Scale Factor，继续配置 Auxiliary Plane：

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];

tensorDesc.auxiliaryPlanes = auxiliaryPlanes;

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

关键点：

• blockFactors 定义了 block-wise 缩放的块大小，这里每 32 列共享一个 Scale
• MTLTensorPlaneTypeScales 表明这个辅助平面用于存储 Scale Factor
• 数据、Scale 和元数据最终打包进同一个 MTLTensor 对象
• 量化数据类型有额外的内存对齐要求，需参考 Metal 文档确认

Shader 中的 MXFP8 张量类型声明

（06:07）在 Metal Shading Language（MSL）中，通过类型别名声明带 Scale Plane 的张量类型：

#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)]])
{
    // ...
}

关键点：

• tensor_blockwise 声明 block-wise Scale Plane，模板参数依次是：平面类型、存储位置、数据格式、块大小
• mxfp8_tensor 的主类型参数：存储位置、维度扩展（dextents<int, 2> 表示 2D 动态维度）、句柄类型、Scale Plane
• tensor_handle 表示从 Host 端 MTLTensor 绑定的张量
• Buffer binding 和普通 MSL kernel 完全一致

如果不需要创建完整的 Host 端 MTLTensor，也可以在 Shader 栈上直接构造内联张量：

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

mxfp8_tensor_inline matrixA(dataBufferA,
                             dextents<int, 2>(K, M),
                             array<int, 2>({ 1, K }),
                             scales_plane(scalesBufferA));

关键点：

• tensor_inline 替换 tensor_handle，表示栈上分配
• 构造函数参数：数据 Buffer 指针、维度、步长（stride）数组、Scale Plane
• 适合临时张量，减少 Host 端对象创建开销

切片与量化矩阵乘

（07:19）按 Threadgroup ID 切片输入输出张量，然后调用 TensorOps 执行矩阵乘：

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);

constexpr auto descriptor = matmul2d_descriptor(TILEM,
                                                TILEN,
                                                dynamic_length_v<int>,
                                                false,
                                                false);

matmul2d<descriptor, execution_simdgroups<4>> op;
op.run(tA, tB, tC);

关键点：

• slice 按 Threadgroup ID 提取每个线程组负责的 tile
• 数据平面和 Scale Plane 会同时按 block size 切片
• matmul2d_descriptor 的模板参数：M、N、K（dynamic_length_v<int> 表示运行时确定）、左矩阵是否转置、右矩阵是否转置
• execution_simdgroups<4> 指定线程组内用 4 个 SIMD group 并行执行
• TensorOps 自动处理量化张量的反量化，无需手动干预

用 Cooperative Tensor 实现 FlashAttention

（10:27）FlashAttention 的核心是把 QxK、SoftMax、乘 V 三个步骤融合进一个 kernel，避免中间结果写回内存。首先设置 SIMD group 级别的矩阵乘：

constexpr auto mul_qk_op_desc = matmul2d_descriptor(/* ... */);
matmul2d<mul_qk_op_desc, execution_simdgroups> mul_qk_op;

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);

auto ctQK = mul_qk_op.get_destination_cooperative_tensor<decltype(tQSlice),
                                                         decltype(tKSlice),
                                                         float>();

mul_qk_op.run(tQSlice, tKSlice, ctQK);

关键点：

• execution_simdgroups 让每个 SIMD group 独立计算，各自拥有完整的中间矩阵行
• sgid（SIMD group ID）用于切片，确保 SoftMax 计算无需跨 SIMD group 交换数据
• get_destination_cooperative_tensor 创建 Cooperative Tensor，数据分布在参与计算的线程私有寄存器中
• Cooperative Tensor 避免了中间结果写回 Threadgroup Memory

接下来计算行维度的最大值归约，为 SoftMax 做准备：

（11:18）

auto ctTileRowMax = mul_qk_op.get_row_reduction_destination_cooperative_tensor<
                        decltype(tQSlice),
                        decltype(tKSlice),
                        float>();

reduce_rows(ctQK, ctTileRowMax, reduction_operation::max, -INFINITY);

关键点：

• get_row_reduction_destination_cooperative_tensor 创建用于存储归约结果的 Cooperative Tensor
• reduce_rows 计算每行的最大值，线程间自动交换数据
• 初始值设为 -INFINITY，确保最大值计算正确

然后用 map_iterator 计算 SoftMax：

（11:56）

#pragma clang loop unroll(full)
for (auto it = ctQK.begin(); it != ctQK.end(); it++) {
    auto row_it = ctRowMax.map_iterator(it);
    *it = exp(*it - *row_it);
}

关键点：

• map_iterator 把 2D Cooperative Tensor 的元素映射到对应的行最大值位置
• 两个 Cooperative Tensor 形状不同，map_iterator 自动处理索引映射
• 先减最大值再求指数，防止数值溢出

最后把 Cooperative Tensor 直接作为下一次矩阵乘的输入：

（12:33）

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)) {
    auto ctQKIn = mul_sv_op.get_left_input_cooperative_tensor<float, half, float>(ctQK);
    mul_sv_op.run(ctQKIn, tVSlice, ctO);
} else {
    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);
}

关键点：

• is_compatible_as_left_input 检查 Cooperative Tensor 的布局是否兼容作为左输入
• 模板参数 <float, half, float> 分别表示左输入、右输入、输出的数据类型
• 兼容时直接复用，数据全程在寄存器中流转
• 不兼容时回退到 Threadgroup Memory 存储再加载
• simdgroup_barrier 确保所有线程完成写入后再读取

集成到 Core AI 模型

（13:35）Core AI 提供了 Python 工具链，可以把 PyTorch 模型转换为 Core AI 格式，并在转换过程中注入自定义 Metal kernel。演讲者以 SAM3 图像分割模型为例，演示了完整的集成流程：

1. 在 Python 中把自定义 Attention kernel 的 MSL 代码定义为字符串
2. 注册 TorchMetalKernel 对象
3. 替换 Hugging Face 默认的 Attention 实现为调用自定义 kernel 的版本
4. 从 Hugging Face 加载模型，导出为优化的 Core AI asset

导出后的模型可以直接在 Apple 设备上运行，推理时自动调用自定义的 FlashAttention kernel。

核心启发

给现有推理框架写自定义量化算子

做什么：为 llama.cpp 或 MLX 写一个支持 FP8 权重的自定义 Linear 层内核。

为什么值得做：现有框架的 Metal backend 可能还没适配 iOS/macOS 27 的新量化类型。用 MTLTensor 的多平面张量 + TensorOps，可以用不到 100 行代码实现一个硬件加速的量化矩阵乘，性能远超手动 bit-shift 解包的实现。

怎么开始：从 MTLTensorDescriptor 创建带 MTLTensorPlaneTypeScales 的量化张量开始，Shader 端用 tensor_blockwise 声明类型，然后直接丢给 matmul2d。参考 TensorOps 的示例代码项目。

在 iPhone 上跑压缩版视觉大模型

做什么：把一个 7B 参数的视觉语言模型（如 LLaVA）量化到 4-bit，部署到 iPhone 上实现实时图像理解。

为什么值得做：4-bit 量化能把模型体积压缩到原来的 1/4，配合 M5/A19 的 Neural Accelerator，Prefill 阶段的密集计算可以被硬件直接加速。以前移动端跑不了的大模型，现在有了落地的可能。

怎么开始：用 Core AI 的 Python 转换工具把 PyTorch 模型导出，在转换时指定量化配置。如果模型里有自定义 Attention 实现，参考本 Session 的 FlashAttention 集成方式注入 Metal kernel。

给 Core AI 模型写自定义后处理算子

做什么：为分割、检测类模型写一个自定义的 NMS（非极大值抑制）或 Mask 后处理内核。

为什么值得做：后处理往往是推理 pipeline 的瓶颈，尤其是涉及排序、过滤这类不适合 GPU 并行化的操作。用 Cooperative Tensor 和 reduce_rows 可以把部分逻辑搬到 GPU 上，减少 CPU-GPU 往返。

怎么开始：先分析后处理中的可并行部分（如置信度排序可以用并行归约），然后用 Cooperative Tensor 在 Shader 中实现，最后通过 TorchMetalKernel 注册到 Core AI 的转换流程中。

探索 2-bit 量化的极限压缩

做什么：尝试把模型权重压缩到 2-bit，测试在 Apple Silicon 上的精度和速度 trade-off。

为什么值得做：iOS/macOS 27 新增了对 2-bit 整数类型的支持。对于某些对精度不敏感的层（如 Embedding 层、部分 FFN 层），2-bit 量化可能在不明显损失精度的情况下进一步压缩模型，让更大的模型能塞进移动设备。

怎么开始：从模型的权重分布分析入手，找出数值范围小、对精度不敏感的层。用 MTLTensorDataType 中对应的 2-bit 类型创建张量，配合合适的 block-wise Scale Factor，在小型模型上验证精度后再扩展。入口：MTLTensorDataType 2-bit 类型 + MTLTensorAuxiliaryPlaneDescriptor

给推理框架写 FlashAttention 内核

做什么：用 Cooperative Tensor 和 TensorOps 的 matmul2d 实现一个完整的 FlashAttention kernel，替代框架默认的 Attention 实现。

为什么值得做：FlashAttention 通过融合 QxK、SoftMax、乘 V 三个步骤避免了中间结果写回内存，对访存密集型的大模型推理有倍数级提升。iOS/macOS 27 的 Cooperative Tensor 直连矩阵乘让实现更简单，数据全程在寄存器中流转。

怎么开始：参考 Session 中的代码结构，先用 matmul2d 计算 QxK 到 Cooperative Tensor，用 reduce_rows 做行维度最大值归约，再用 map_iterator 计算 SoftMax，最后把 Cooperative Tensor 直接作为下一次矩阵乘的输入。入口：matmul2d + get_destination_cooperative_tensor + get_left_input_cooperative_tensor

关联 Session

• Session 324: Accelerate machine learning with Core AI — Core AI 框架的整体介绍，包括模型转换和部署流程
• Session 328: Build local AI features with MLX Swift — 用 MLX Swift 在 Apple 设备上构建本地 AI 功能，可与自定义 Metal kernel 结合
• Session 388: Optimize Metal performance for Apple GPUs — Metal 性能优化的通用原则，包括内存管理和着色器调优
• Session 359: Enhance your game with Metal neural rendering — Metal 神经网络渲染，涉及类似的 GPU 计算优化思路