Optimize custom machine learning operations with Metal tensors

Watch on Apple Developer

TL;DR

Where Metal TensorOps fits

Apple's ML stack ranges from high-level frameworks such as Core AI and MLX to lower-level Metal Performance Shaders, Metal Performance Primitives, and the TensorOps library. Working at the Metal level is useful when implementing fast custom operations for higher-level frameworks, contributing to ML runtimes such as MLX or llama.cpp, or building Metal-based applications.

TensorOps is a Metal Shading Language API for accelerated tensor operations, including matrix multiplication and convolution. It is designed to use available hardware acceleration across Apple Silicon GPU generations and specifically takes advantage of the neural accelerator in the M5 chip family.

• Use Core AI or MLX when high-level deployment is enough.
• Use Metal TensorOps when a model needs custom kernels, lower-level control, or framework/runtime integration.
• The M5 neural accelerator is located in each shader core and targets dense compute-bound work such as LLM prefill.

Quantized tensors and scale planes

Quantization reduces model memory footprint and bandwidth by storing weights in lower-precision formats paired with scale factors. TensorOps can consume quantized tensors directly and use available hardware acceleration instead of requiring manual dequantization for standard formats.

Metal tensor support includes 4- and 8-bit integer types in macOS and iOS 26, with additional support in macOS and iOS 27 for 4- and 8-bit floating point types and 2-bit integer types. macOS and iOS 27 also allow one MTLTensor to represent quantized element data plus an auxiliary scale plane using the FP8 E8M0 block-wise scale factor format.

• Create a quantized tensor by setting MTLTensorDescriptor.dataType to a quantized MTLTensorDataType.
• Attach scale factors with MTLTensorAuxiliaryPlaneDescriptorMap and MTLTensorPlaneTypeScales.
• Each scale-plane element applies to a block of data-plane elements according to blockFactors.
• New smaller data types have additional alignment requirements; consult the Metal documentation for exact constraints.

Create a quantized MTLTensor

A quantized tensor is allocated like a regular MTLTensor, but with a quantized Metal tensor data type.

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

Declare an auxiliary scales plane

A multi-plane tensor can pack quantized element data, block-wise scales, and metadata into one tensor object.

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

tensorDesc.auxiliaryPlanes = auxiliaryPlanes;

Quantized matrix multiplication

The session extends a tiled TensorOps matrix multiplication kernel to quantized inputs. The host can allocate full MTLTensor objects and bind them as tensor handles, or a shader can construct inline tensor views from raw buffers when a full host-side tensor object is not needed.

When slicing multi-plane tensors, the data plane and scale plane are sliced together according to the declared block size. Once the tiles are prepared, the matmul2d setup is the same as for non-quantized tensors, and TensorOps handles dequantization for supported formats.

• Define MSL tensor aliases for the quantized element format and scale plane.
• Use tensor_handle when binding host-created tensors to buffer slots.
• Use tensor_inline to construct temporary tensor views on the shader stack from buffers and metadata.
• Prefer feeding quantized tensors directly into TensorOps; use custom dequantization only when the format is not supported directly.

MSL aliases for an MXFP8 tensor handle

The quantized element type and auxiliary scales plane are encoded in the MSL tensor type.

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

Slice tiles and run quantized matmul

TensorOps accepts quantized tensor tiles directly and performs the needed dequantization for supported formats.

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

Custom dequantization and cooperative tensors

For unsupported or custom quantization formats, one option is to have each thread load quantized data, dequantize to f16, store to threadgroup memory, and pass an inline threadgroup tensor to TensorOps. That is simple but adds extra loads and stores.

A faster approach is to dequantize into a cooperative tensor. Cooperative tensors distribute storage across the private memory of participating threads and can now be passed as inputs to matmul2d, avoiding a round trip through threadgroup memory when layouts are compatible.

• Use cooperative tensors when intermediate data should remain in registers/private storage.
• Check compatibility before reusing a cooperative tensor as a matmul input.
• Fallback to store/reload through threadgroup memory if the cooperative tensor layout is not compatible.

Reuse a cooperative tensor as matmul input

The new direct cooperative-tensor input path avoids threadgroup memory when the layout can be reused safely.

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

Building FlashAttention with TensorOps

FlashAttention fuses the Q×K multiplication, row-wise SoftMax, and multiplication by V into one kernel. The TensorOps implementation uses SIMD-group-scoped execution so each SIMD group owns complete rows of the intermediate matrix, allowing SoftMax to be computed without exchanging data between SIMD groups.

The intermediate Q×K tile is held in a cooperative tensor. TensorOps row reductions compute per-row values such as the maximum, and map_iterator maps each element of the 2D cooperative tensor to its corresponding row-reduction element for element-wise SoftMax work.

• Use execution_simdgroup for independent SIMD-group matrix multiplications.
• Use get_destination_cooperative_tensor for the Q×K intermediate.
• Use reduce_rows for row-wise reductions needed by SoftMax.
• Use map_iterator to relate 2D elements to row-reduction results.

Compute Q×K into a cooperative tensor

The Q×K tile is produced into cooperative storage rather than written to memory.

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

Row reduction and element-wise SoftMax step

reduce_rows and map_iterator support the row-wise SoftMax portion of the fused attention kernel.

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

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

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

Core AI integration path

Custom Metal TensorOps kernels can be integrated into Core AI applications. The session demonstrates the workflow conceptually with a custom FlashAttention kernel registered from Python, replacement of a default Hugging Face attention implementation, and export from PyTorch to an optimized Core AI asset.

The example uses a SAM3 image segmentation model and validates that the custom attention kernel is part of the Core AI inference path. For details of the conversion and authoring workflow, the session points to the Core AI model authoring and optimization material.

• Core AI Python tools can convert PyTorch models to Core AI models.
• Custom Metal kernels can be registered and called from model code before export.
• Use this path when a high-level model needs a custom fused Metal operation rather than a default framework implementation.

Resources

• Running inline ML operations in a shader with Metal 4
• Machine learning passes
• Download the Metal Performance Primitives (MPP) Programming Guide
• Metal Performance Shaders

Related Sessions

• Accelerate your machine learning workloads with the M5 and A19 GPUs
• Combine Metal 4 machine learning and graphics

Unofficial, not associated with Apple. By