Metal 텐서로 맞춤형 머신 러닝 연산 최적화하기

Metal 텐서로 맞춤형 머신 러닝 연산 최적화하기

Metal Tensor API와 MPP(Metal Performance Primitive) 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 프레임워크에 기여하고 있거나 MLX나 llama.cpp에 기여하는 경우, 또는 Metal 기반 애플리케이션을 작업하는 경우에도 Metal 커널을 작성해야 할 수 있습니다. 시작하는 가장 쉬운 방법은 TensorOps 라이브러리를 사용하는 것입니다. TensorOps는 Metal Shading Language API로, GPU에서 텐서 연산을 가속화하며, 행렬 곱셈 및 합성곱을 포함합니다. 모든 Apple Silicon GPU 세대에 걸쳐 사용 가능한 하드웨어 가속을 자동으로 활용하므로, 모든 Apple Silicon GPU 세대에 걸쳐 하드웨어 세대 간의 차이에 대해 걱정할 필요가 없습니다. 하드웨어 세대 간 차이를 걱정하지 않아도 됩니다. 특히, M5 칩 패밀리의 뉴럴 가속기를 최대한 활용합니다. M5 칩 패밀리의 뉴럴 가속기를.
뉴럴 가속기는 M5의 새로운 하드웨어 블록으로, 각 셰이더 코어 내부에 직접 위치합니다. 다른 GPU 파이프라인 옆에 위치하며, 밀집된 LLM의 프리필 단계와 같은 컴퓨팅 집약적인 작업을 가속화하도록 설계되었습니다.
관련 세션을 확인하시면 TensorOps 시작의 기초를 배울 수 있습니다. 이 세션에서는 그 기초를 바탕으로, 양자화된 데이터 작업을 위한 모범 사례부터 시작하겠습니다. 그런 다음 FlashAttention과 같은 고급 커스텀 연산을 구축하는 방법을 보여드리겠습니다. FlashAttention과 같은 고급 커스텀 연산을. 첫 번째 주제인 양자화된 데이터 작업으로 바로 들어가 보겠습니다.
알다시피, 최신 머신러닝 모델은 점점 더 커지고 있습니다. 추론 단계는 일반적으로 메모리 대역폭에 의해 제한되므로, 가중치를 압축하는 것이 필요합니다. 모델을 메모리에 더 잘 맞추고 메모리 대역폭을 절약하기 위해서입니다.
가중치를 압축하는 표준 방법은 양자화입니다. 아이디어는 간단합니다 — 더 높은 정밀도의 가중치를 취하여 더 낮은 정밀도의 데이터 타입으로 줄입니다. 예를 들어, 16비트 반정밀도 가중치를 4비트로 압축할 수 있습니다. 이러한 양자화된 가중치는 스케일 팩터와 함께 사용되며, 양자화된 값을 계산할 때 원래 범위로 다시 스케일링할 수 있게 해줍니다.
16비트 및 32비트 부동소수점 타입 외에도, TensorOps는 이제 양자화된 데이터 타입을 기본적으로 지원합니다. macOS 및 iOS 26 업데이트에서 4비트 및 8비트 정수 타입 지원을 추가했으며, macOS 및 iOS 26 업데이트에서 macOS 및 iOS 27에서는 더 많은 데이터 타입으로 지원을 확장합니다. 여기에는 4비트 및 8비트 부동소수점 타입과 2비트 정수 타입이 포함됩니다. 앱의 양자화된 텐서를 TensorOps에 간단히 생성하고 전달하면 사용 가능한 모든 하드웨어 가속을 자동으로 활용합니다.
양자화된 데이터 타입으로 텐서를 생성하는 것은 일반 텐서를 생성하는 것과 매우 유사합니다. 다른 텐서와 마찬가지로 디스크립터의 속성을 채우되, 단순히 양자화된 dataType을 지정합니다. 그런 다음 Metal 디바이스에서 newTensorWithDescriptor를 호출하여 텐서를 생성합니다.
이것이 양자화된 요소 데이터를 저장하는 방법입니다. 다음으로, 스케일 팩터에 대해 이야기해 보겠습니다. macOS 및 iOS 27에서는, 단일 MTLTensor 객체가 이제 스케일을 표현할 수 있으며 텐서의 양자화된 데이터와 함께 추가 스케일 플레인으로 나타낼 수 있습니다. 이 플레인은 인기 있는 FP8 E8M0 블록 단위 스케일 팩터 형식을 지원합니다. 스케일 플레인의 각 요소는 데이터 플레인의 요소 블록에 적용됩니다. 스케일 플레인 선언은 텐서 선언과 유사합니다. 먼저 스케일 플레인에 대한 디스크립터 객체를 생성합니다. 그런 다음 dataType과 blockFactors를 채웁니다. 마지막으로, 이 플레인이 스케일용임을 지정하는 보조 플레인 맵을 생성합니다.
그런 다음 원래 tensorDescriptor에 보조 플레인 맵을 첨부합니다. 양자화된 데이터, 스케일, 및 메타데이터가 모두 단일 텐서 객체에 패킹됩니다.
이제 이를 실제로 적용해 보겠습니다. 기본 행렬 곱셈 커널을 확장하여 양자화를 지원하도록 합니다. 행렬 곱셈은 머신러닝 작업의 핵심 연산입니다. 예를 들어, LLM은 추론 중에 수백만 번의 행렬 곱셈을 수행합니다.
기초를 다루었으며, TensorOps로 고성능 행렬 곱셈 커널을 작성하는 방법을 M5 머신러닝 강연에서 다뤘습니다. 기본 접근 방식은 입력 행렬을 더 작은 타일로 분할하고, TensorOps를 사용하여 타일 단위 행렬 곱셈을 수행하는 것입니다. 이는 병렬성을 극대화하고 데이터를 캐시에 유지합니다.
양자화를 사용하여 메모리 트래픽을 더욱 줄이고 더 큰 모델을 메모리에 맞출 수 있습니다. 커널에서는 텐서를 바인딩하기 전에 미리 타입 별칭을 정의하는 것이 도움이 됩니다. 텐서를 바인딩하기 전에. 여기서는 fp8_e8m0_ 데이터 타입으로 스케일 팩터 플레인을 선언하고, 32 x 1의 블록 크기를 사용합니다. 즉, 데이터 플레인의 32개 요소마다 scales_plane의 단일 요소를 공유합니다. 그런 다음 FP8 데이터 타입을 지정하는 전체 텐서 타입을 선언하고 scales_plane과 함께 사용합니다. 이러한 텐서를 버퍼 바인딩 포인트에 간단히 바인딩할 수 있습니다. 그러면 커널이 텐서에 접근할 수 있게 되며, 호스트 측에서 할당한 텐서에 접근합니다. 또는 호스트에서 전체 MTLTensor를 생성하지 않으려면, 셰이더의 스택에서 바로 임시 텐서를 생성할 수 있습니다. 구문은 거의 동일하며, tensor_handle 태그를 tensor_inline으로 교체하면 됩니다. 그런 다음 버퍼 포인터와 기타 메타데이터를 텐서 생성자에 전달하여 스택에 텐서를 생성합니다.
앞서 언급했듯이, 더 나은 병렬성을 위해 많은 스레드그룹에 걸쳐 문제를 분할합니다. 먼저 각 스레드그룹을 위한 타일을 슬라이스하고, 그런 다음 TensorOps로 곱셈을 수행합니다.
이렇게 하려면 입력 및 출력 텐서에서 slice를 호출하면 됩니다. 스레드그룹 ID를 사용하여. 데이터와 스케일 플레인 모두 동시에 슬라이스되며 블록 크기에 따라 처리됩니다. 양자화된 텐서로 행렬 곱셈을 설정하는 것은 일반 텐서와 동일합니다. 먼저 matmul2d_descriptor를 설정하고, 타일 크기와 기타 매개변수를 지정합니다. 그런 다음 matmul2d op를 생성하고, 스레드그룹의 simdgroup 수를 지정합니다. 그런 다음 양자화된 텐서를 전달하면 TensorOps가 역양자화를 자동으로 처리합니다.
대부분의 경우, 양자화된 데이터를 TensorOps에 직접 전달하여 사용 가능한 모든 하드웨어 가속을 자동으로 활용하도록 해야 합니다. 그러나 커스텀 형식을 역양자화해야 하는 경우에도, TensorOps가 여전히 지원합니다. 가장 간단한 접근 방식은 각 스레드가 디바이스 메모리에서 양자화된 데이터 청크를 로드하고 스레드그룹 메모리에서 f16 값으로 역양자화하는 것입니다. 그런 다음 인라인 스레드그룹 텐서로 TensorOps에 전달할 수 있습니다. 그러나 이 접근 방식은 추가적인 로드 및 저장이 필요하며 스레드그룹 메모리를 통해 이루어집니다. 이상적으로는 모든 데이터를 스레드 레지스터에 유지하는 것이 좋습니다. 데이터를 cooperative 텐서로 역양자화하면 이를 달성할 수 있으며, 이제 matmul2d op의 입력으로 전달할 수 있습니다. cooperative 텐서는 저장소를 스레드 프라이빗 메모리에 분산시키며, matmul 연산에 참여하는 스레드들에 걸쳐 분산됩니다. 따라서 양자화된 텐서를 직접 사용할 수 없는 경우에도, 스레드그룹 메모리를 통한 왕복을 건너뛸 수 있습니다. 요약하면 — Metal 텐서는 다양한 양자화된 데이터 타입을 기본적으로 지원하며, 새로운 MX 스케일링 형식을 포함하여 iOS 및 macOS 27에서 제공되는 E8M0 스케일 팩터도 지원합니다. 이러한 새로운 데이터 타입에는 추가적인 정렬 요구 사항이 있음을 참고하세요. 더 큰 데이터 타입에 비해 자세한 내용은 Metal 문서를 확인하세요.
이제 한 단계 높여보겠습니다 — TensorOps로 더 복잡한 전체 커스텀 연산을 구축해 보겠습니다. 어텐션은 모든 트랜스포머 네트워크의 핵심이며, LLM도 포함됩니다. 어텐션을 계산하려면, 먼저 Q와 K라는 두 행렬을 함께 곱합니다. 다음으로 중간 행렬의 행에 대한 리덕션을 사용하여 SoftMax를 계산합니다.
마지막으로 V라는 세 번째 행렬과 곱합니다. 인기 있는 FlashAttention 알고리즘은 이러한 모든 연산을 단일 커널로 융합합니다.
TensorOps로 이를 구현하려면, 먼저 커스텀 simd group 매핑을 설정하여 각 simd group이 중간 행렬의 완전한 행을 소유하도록 해야 합니다. 이를 통해 SoftMax를 계산할 수 있으며 simd group 간에 데이터를 교환하지 않아도 됩니다. execution_simdgroup operation scope를 사용하여 이를 수행할 수 있습니다. 즉, 각 simd group이 병렬로 독립적인 행렬 곱셈을 수행합니다.
simd group ID를 사용하여 입력 타일을 슬라이스할 수 있습니다. 중간 행렬을 저장하기 위해 cooperative 텐서를 사용하여 다음 단계의 입력으로 사용할 수 있으며 메모리에 쓰지 않아도 됩니다. 결과에 대해 SoftMax를 계산합니다.
이를 위해 몇 가지 리덕션을 계산해야 합니다. cooperative 텐서에 대해. TensorOps에는 이를 돕기 위한 reduce_rows 함수가 포함되어 있습니다. 스레드들이 서로 데이터를 교환하여 각 행의 최댓값을 계산합니다. 결과는 또 다른 cooperative 텐서에 반환됩니다.
설정해 보겠습니다. 먼저 리덕션 출력을 저장할 cooperative 텐서를 생성합니다. 그런 다음 소스와 목적지를 reduce_rows 함수에 전달합니다. 여기서는 음의 INFINITY 초기값으로 max reduction_operation을 사용합니다.
이 두 cooperative 텐서는 형태가 다르므로, 둘 간의 매핑을 돕기 위해 TensorOps에는 map_iterator 함수도 포함되어 있습니다. 2D 텐서의 요소를 가리키는 이터레이터가 주어지면, 리덕션 대상의 해당 요소를 가리키는 이터레이터를 반환합니다. 리덕션 대상의 해당 요소를.
먼저 이터레이터를 사용하여 2D cooperative 텐서에 대한 루프를 설정합니다. 그런 다음 map_iterator를 호출하여 각 요소를 해당 행의 최댓값에 매핑합니다. 마지막으로 이 이터레이터를 역참조하여 SoftMax를 계산하고 결과를 cooperative 텐서에 다시 저장합니다.
이제 이 cooperative 텐서를 V와 곱할 준비가 되었습니다. macOS 26에서는 먼저 스레드그룹 메모리에 저장해야 했을 것입니다. 그러나 이제 cooperative 텐서를 matmul 연산의 입력으로 직접 사용할 수 있습니다.
이렇게 하려면 get_left_input_cooperative_tensor 메서드를 호출하고, 소스 cooperative 텐서를 인수로 전달합니다. 그런 다음 결과를 두 번째 matmul 연산의 입력으로 전달할 수 있습니다. 주의해야 할 한 가지 사항: 모든 cooperative 텐서가 입력으로 재사용될 수 있는 것은 아닙니다. 레이아웃은 데이터 타입과 기타 요인에 따라 달라질 수 있습니다. 따라서 이 작업을 수행하기 전에, is_compatible_as_left 또는 right_input 메서드를 호출하여 호환성을 확인합니다. true를 반환하면 진행하면 됩니다. 그렇지 않으면 데이터를 저장하고 다시 로드해야 하며, 스레드그룹 메모리를 통해 올바른 레이아웃으로 변환합니다. 어떤 경우든 op.run 호출은 동일합니다. 이것이 TensorOps를 사용하여 고급 연산을 구축하는 데 필요한 핵심 TensorOps 기능입니다. FlashAttention과 같은 고급 연산을 TensorOps로. 이 연산을 구축하는 방법을 살펴봤으므로, Core AI를 사용하는 실제 모델에서 어떻게 실행되는지 살펴보겠습니다. Core AI는 Python 개발자를 위한 도구를 제공하여 PyTorch 모델을 Core AI 모델로 변환하며, 커스텀 Metal 커널 지원을 포함합니다. Metal 커널을 Core AI 모델에 통합하는 방법의 세부 사항은 "Deep Dive into Core AI Model authoring and Optimization" 세션을 확인하세요.
해당 세션에 설명된 단계를 따라 커스텀 FlashAttention 커널을 Sam3 이미지 세분화 모델에 통합했습니다. 커스텀 어텐션 커널의 본문을 Python에서 문자열로 정의하고 여기에 표시된 TorchMetalKernel 객체를 등록합니다.
그런 다음 기본 huggingface 어텐션 구현을 여기에 표시된 커널을 호출하는 구현으로 교체합니다.
마지막으로 huggingface에서 모델을 로드하고 PyTorch에서 최적화된 Core AI 에셋으로 내보냅니다. 내보내기가 완료되는 데 잠시 시간이 걸립니다.
이제 추론을 실행할 준비가 되었습니다. Sam3는 프롬프트 가능한 개념 세분화를 수행하며, 모델에 이미지와 텍스트를 제공하면, 세분화 마스크로 응답하여 이미지에서 객체가 위치한 곳을 나타냅니다. 여기서 저는 모델에게 이 이미지에서 자동차가 포함된 모든 픽셀에 레이블을 지정하도록 프롬프트합니다.
자, 이제 세분화를 실행하겠습니다.
최종 결과를 보면, 모델이 이미지를 올바르게 세분화했음을 알 수 있습니다. 자동차가 파란색으로 강조 표시되어 있으며, 어텐션 커널이 완전히 통합되어 예상대로 모델에 적용되었습니다.
오늘 Apple Silicon에서 최적화된 커스텀 ML 커널을 구축하는 데 사용할 수 있는 모든 도구를 다뤘습니다. 양자화된 데이터 타입부터 cooperative 텐서 및 리덕션과 같은 고급 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