Optimize custom machine learning operations with Metal tensors
AI & Machine Learning 高级 16m

使用 Metal 张量优化自定机器学习运算

Optimize custom machine learning operations with Metal tensors

2026年6月10日

在 Apple 官方观看视频

一句话判断

这场 Session 最值得关注的一件事是：Metal Tensor 原生支持了多平面量化张量（Multi-plane tensors）和 Cooperative Tensors 直连 MatMul，这意味着在 GPU 上写 FlashAttention 这种复杂算子时，终于不用在 Threadgroup Memory（线程组内存）里来回倒腾数据了，M5 芯片的神经网络加速器（Neural Accelerator）能被彻底喂饱。

这场 Session 讲了什么

大模型推理的瓶颈早就从算力变成了内存带宽。为了把几十 B 的模型塞进 iPhone 和 Mac，量化（Quantization）是必经之路。以前在 Metal 里手写量化算子，你得自己处理解包、缩放，还要在 Threadgroup Memory 里做各种数据搬运，写得头疼还跑不满 GPU。

现在 Apple 直接把量化数据类型（比如 FP8、4-bit 整数）和 Scale Factor（比例因子）打包进了一个 MTLTensor 对象里。更狠的是，在写 FlashAttention 时，Cooperative Tensors（协作张量）现在可以直接喂给 matmul2d 算子，干掉了中间写回 Shared Memory 的步骤。

这对搞端侧大模型（比如 llama.cpp 贡献者、Core AI 开发者）的人来说是巨大的利好。你不用再跟底层内存对齐死磕，直接调 TensorOps 就能榨干 M5 和 A19 芯片里新加的硬件加速单元。

值得深挖的点

Multi-plane Tensors 让量化不再“缺胳膊少腿”

大模型权重压缩是刚需，MXFP8 和 block-wise scaling 已经是业界标准。以前在 Metal 里搞量化，数据和 Scale Factor 往往是分开的两个 Buffer，Shader 里得手动算索引去读 Scale，再反量化回 FP16。这不仅代码写得像天书，还白白浪费了内存带宽。

Apple 这次把数据和 Scale Factor 塞进了同一个 MTLTensor，通过 Auxiliary Plane（辅助平面）来管理。你在 Host 端配好 MTLTensorAuxiliaryPlaneDescriptor，Shader 端直接用 tensor_blockwise 声明类型。最爽的是，当你把这个多平面张量传给 TensorOps 的 matmul2d 时，底层会自动调用 M5 芯片的 Neural Accelerator 进行硬件级反量化。

这里的 Trade-off 很明显：代码极简，性能拉满，但硬件对量化类型的 Alignment（内存对齐）要求极严。比如 4-bit 数据类型，如果你的 Tensor 维度不是 32 的倍数，或者 Buffer 的起始地址没对齐，大概率会直接 Fallback 到慢速路径，甚至引发 Shader 编译错误。用这个 API 前，务必把 Metal 文档里的对齐表格背下来。

Cooperative Tensors 直连 MatMul 消灭内存搬运

Attention 机制需要连续做 QxK 矩阵乘、SoftMax、再乘 V。中间结果如果写回 Threadgroup Memory 再读出来，带宽和延迟根本受不了。以前为了优化，大家绞尽脑汁在寄存器里倒腾数据，但 SIMD group 之间的数据交换依然很痛苦。

这次 TensorOps 允许将 Cooperative Tensors（数据分布在 SIMD group 寄存器里的张量）直接作为下一个 matmul2d 的输入。配合 execution_simdgroups 作用域，每个 SIMD group 独立计算自己拥有的行，算完 SoftMax 后，数据直接在寄存器里流转给下一次矩阵乘。

这彻底消灭了 Threadgroup memory round-trip（线程组内存往返）。从“寄存器 -> 共享内存 -> 寄存器”变成了纯粹的“寄存器 -> 寄存器”。对于 FlashAttention 这种访存密集型算子，这不仅是几行代码的简化，而是实打实的倍数级性能提升。坑点在于，你需要精确控制每个 SIMD group 映射到中间矩阵的哪一行，map_iterator 这个新 API 就是为了解决这种多维索引映射而生的，用错了迭代器会导致数据错位。

代码片段

1. 创建带 Scale Factor 的多平面量化张量

场景：在 Host 端分配一个 FP8 量化张量，并绑定 E8M0 格式的 Scale Factor 平面。

// 创建主数据平面的描述符，指定 FP8 数据类型
MTLTensorDescriptor *tensorDesc = [MTLTensorDescriptor new];
tensorDesc.dataType = MTLTensorDataTypeMetalFloat8E4M3;
tensorDesc.usage = MTLTensorUsageCompute;
tensorDesc.dimensions = [[MTLTensorExtents alloc] initWithRank:2 values:(NSInteger[]){NumCols, NumRows}];

// 创建 Scale Factor 辅助平面的描述符
MTLTensorAuxiliaryPlaneDescriptor *planeDesc = [MTLTensorAuxiliaryPlaneDescriptor new];
planeDesc.dataType = MTLTensorDataTypeMetalFloat8UE8M0;
// 设置 block-wise 缩放的块大小，每 32 个元素共享一个 Scale
planeDesc.blockFactors = [[MTLTensorExtents alloc] initWithRank:2 values:(NSInteger[]){32, 1}];

// 将辅助平面映射到 Scales 类型，并附加到主描述符
MTLTensorAuxiliaryPlaneDescriptorMap *auxiliaryPlanes = [MTLTensorAuxiliaryPlaneDescriptorMap new];
[auxiliaryPlanes setDescriptor:planeDesc forPlane:MTLTensorPlaneTypeScales];
tensorDesc.auxiliaryPlanes = auxiliaryPlanes;

// 坑：blockFactors 的乘积必须能整除主平面的对应维度，否则创建 Tensor 时会直接报错。

2. 在 Shader 中计算 Row-wise Reduction (SoftMax)

场景：在 FlashAttention 中，对 QxK 的结果进行行维度的 Max 归约，为 SoftMax 做准备。

#include <metal_tensor>
using namespace metal;

// 假设 ctQK 是存储 QxK 结果的 2D Cooperative Tensor
// 创建一个 1D Cooperative Tensor 来存储每一行的最大值
auto ctTileRowMax = mul_qk_op.get_row_reduction_destination_cooperative_tensor<
                        decltype(tQSlice), decltype(tKSlice), float>();

// 调用 TensorOps 内置的 reduce_rows，初始值设为负无穷
reduce_rows(ctQK, ctTileRowMax, reduction_operation::max, -INFINITY);

// 使用 map_iterator 将 2D 元素映射到对应的 1D 行最大值，计算 exp
#pragma clang loop unroll(full)
for (auto it = ctQK.begin(); it != ctQK.end(); it++) {
    // 自动找到当前元素所在行的 Max 值
    auto row_it = ctTileRowMax.map_iterator(it);
    // 减去最大值并求指数，防止数值溢出
    *it = exp(*it - *row_it);
}
// 坑：reduce_rows 会隐式触发 SIMD group 内的数据交换，确保调用前所有线程都已同步写入 ctQK。

最佳实践

如果你正在维护基于 MLX 或 llama.cpp 的端侧推理项目，不要自己急着改底层，等社区大佬把这些新 API 适配好。如果是自己从头写 Metal 推理引擎，立刻把以前那些手写 bit-shift 解包和 Threadgroup 搬运的“祖传代码”删掉，全面拥抱 Multi-plane tensor 和 TensorOps。

新项目强烈建议直接基于 Core AI 构建。Apple 这次提供了完整的 Python 工具链，你可以直接把 PyTorch 模型转成 Core AI 格式，并在转换过程中无缝插入自定义的 Metal TensorOps 算子。这比自己在 Swift 层封装一堆 C++ 桥接代码要优雅得多。

实战中最容易踩的坑是内存对齐。量化张量（尤其是 2-bit 和 4-bit）对 Buffer 的起始地址和行列维度有严格的硬件对齐要求。在 Host 端分配 MTLBuffer 时，务必使用 MTLResourceStorageModePrivate 并确保尺寸向上取整到硬件要求的 Block Size，否则你会看到性能断崖式下跌。

还有什么值得关注

  • M5/A19 新增 Neural Accelerator：这个硬件模块直接集成在 Shader Core 旁边，专门用来加速 LLM Prefill 阶段的密集计算，不再和图形渲染抢 ALU 资源。
  • Core AI 的 Python 转换工具：官方下场做 PyTorch 到 Core AI 的模型转换了，支持自定义算子注入，省去了写一堆 ONNX 中间件的麻烦。
  • Inline Tensors 支持栈上分配：如果不想在 Host 端创建完整的 MTLTensor，现在可以直接在 Shader 栈上用 tensor_inline 从 Buffer 指针构造张量，少了一次对象创建的开销。
Metal TensorOps Quantization FlashAttention Core AI