{"$schema": "https://raw.githubusercontent.com/guitaripod/wwdc-sessions/master/schema/transcript.schema.json", "id": "wwdc2026-330", "language": "eng", "source": "https://developer.apple.com/videos/play/wwdc2026/330", "wordCount": 2116, "segments": [{"start": 8.0, "text": "Hello, my name is Shiyao. I'm a GPU Software Engineer."}, {"start": 11.0, "text": "Today, I am pleased to guide you through an exploration of Metal tensors,"}, {"start": 16.0, "text": "and show you how to write optimized custom ML kernels with TensorOps."}, {"start": 21.0, "text": "Apple platforms provide first-class support"}, {"start": 24.0, "text": "for running ML models at every layer of the software stack."}, {"start": 28.0, "text": "High-level frameworks like Core AI and MLX"}, {"start": 31.0, "text": "make it easy to deploy your models with minimal code,"}, {"start": 35.0, "text": "while lower-level APIs like Metal Performance Shaders"}, {"start": 38.0, "text": "provide access to high-performance Metal kernels."}, {"start": 42.0, "text": "These layers all build on the low-level acceleration"}, {"start": 45.0, "text": "provided by Metal Performance Primitives and the TensorOps library."}, {"start": 49.0, "text": "There are a few reasons why you might want to work at the Metal level."}, {"start": 53.0, "text": "ML research moves quickly,"}, {"start": 56.0, "text": "so you might want to implement custom operations"}, {"start": 58.0, "text": "which can plug into a higher level frameworks such as Core AI."}, {"start": 62.0, "text": "You may also need to write Metal kernels if you're contributing to an ML framework"}, {"start": 67.0, "text": "such as MLX or llama.cpp"}, {"start": 70.0, "text": "or if you're working on a Metal-based application."}, {"start": 74.0, "text": "The easiest way to get started is using the TensorOps library."}, {"start": 79.0, "text": "TensorOps is a Metal Shading Language API"}, {"start": 82.0, "text": "which accelerates tensor operations on the GPU,"}, {"start": 85.0, "text": "including matrix multiplication and convolution."}, {"start": 89.0, "text": "It automatically uses any available hardware acceleration"}, {"start": 93.0, "text": "across all Apple Silicon GPU generations,"}, {"start": 96.0, "text": "so you don't need to worry about the differences"}, {"start": 98.0, "text": "between hardware generations."}, {"start": 100.0, "text": "In particular, it takes full advantage of the neural accelerator"}, {"start": 104.0, "text": "in the M5 chip family."}, {"start": 107.0, "text": "The neural accelerator is a new hardware block in M5,"}, {"start": 111.0, "text": "located directly in each shader core."}, {"start": 114.0, "text": "It sits alongside the other GPU pipelines and is designed to accelerate dense"}, {"start": 118.0, "text": "compute-bound work such as the prefill stage of an LLM."}, {"start": 125.0, "text": "You can check out the related sessions"}, {"start": 127.0, "text": "to learn the basics of getting started with TensorOps."}, {"start": 131.0, "text": "In this session, I'll build on those basics,"}, {"start": 134.0, "text": "starting with best practices for working with quantized data."}, {"start": 138.0, "text": "Then, I'll show you how to build advanced custom operations"}, {"start": 142.0, "text": "such as FlashAttention."}, {"start": 145.0, "text": "Let's dive into the first topic — working with quantized data."}, {"start": 150.0, "text": "As we know, state-of-the-art machine learning models are getting larger."}, {"start": 154.0, "text": "The inference stage is typically memory bandwidth bound,"}, {"start": 158.0, "text": "so compressing the weights becomes necessary"}, {"start": 162.0, "text": "both to better fit models into memory and to save memory bandwidth."}, {"start": 167.0, "text": "The standard approach for compressing weights is quantization."}, {"start": 171.0, "text": "The idea is simple — take higher-precision weights"}, {"start": 175.0, "text": "and reduce them into lower-precision data types."}, {"start": 178.0, "text": "For example, 16-bit half-precision weights could be compressed down to just 4-bits."}, {"start": 185.0, "text": "These quantized weights are paired with scale factors,"}, {"start": 188.0, "text": "which let us scale the quantized values"}, {"start": 191.0, "text": "back into the original range when it's time to compute."}, {"start": 196.0, "text": "In addition to 16- and 32-bit floating point types,"}, {"start": 200.0, "text": "TensorOps now natively supports quantized data types."}, {"start": 204.0, "text": "We added a support for 4- and 8-bit integer types"}, {"start": 207.0, "text": "in an update to macOS and iOS 26,"}, {"start": 211.0, "text": "and we're extending support to even more data types in macOS and iOS 27."}, {"start": 217.0, "text": "This includes 4- and 8-bit floating point types"}, {"start": 220.0, "text": "and 2-bit integer types."}, {"start": 222.0, "text": "You can simply create and pass your app's quantized tensors to TensorOps"}, {"start": 227.0, "text": "and it will automatically take advantage of any available hardware acceleration."}, {"start": 233.0, "text": "Creating a tensor with a quantized data type"}, {"start": 236.0, "text": "is very similar to creating a regular tensor."}, {"start": 239.0, "text": "You fill in your descriptor's properties like any other tensor,"}, {"start": 243.0, "text": "but simply specify a quantized dataType."}, {"start": 246.0, "text": "Then create the tensor by calling newTensorWithDescriptor"}, {"start": 250.0, "text": "on your Metal device."}, {"start": 254.0, "text": "So that's how you can store your quantized element data."}, {"start": 258.0, "text": "Next, let's talk about the scale factors."}, {"start": 261.0, "text": "In macOS and iOS 27,"}, {"start": 263.0, "text": "a single MTLTensor object can now represent your scales"}, {"start": 268.0, "text": "alongside your tensor's quantized data as an additional scale plane."}, {"start": 273.0, "text": "This plane supports the popular FP8 E8M0 block-wise scale factor format."}, {"start": 278.0, "text": "Each element of the scale plane applies to a block of elements in the data plane."}, {"start": 284.0, "text": "Declaring the scale plane is similar to declaring a tensor."}, {"start": 289.0, "text": "First, create a descriptor object for the scale plane."}, {"start": 293.0, "text": "Then fill in the dataType and blockFactors."}, {"start": 296.0, "text": "Finally, create an auxiliary plane map to specify that this plane is for scales."}, {"start": 304.0, "text": "Then simply attach the auxiliary planes map to your original tensorDescriptor."}, {"start": 309.0, "text": "The quantized data, scales,"}, {"start": 312.0, "text": "and metadata will all be packed into a single tensor object."}, {"start": 317.0, "text": "Now let's put this into practice"}, {"start": 319.0, "text": "by extending a basic matrix multiplication kernel"}, {"start": 322.0, "text": "to support quantization."}, {"start": 325.0, "text": "Matrix multiplication is the core operation"}, {"start": 327.0, "text": "in machine learning workloads."}, {"start": 329.0, "text": "For instance, LLMs perform millions of matrix multiplications during inference."}, {"start": 336.0, "text": "We covered the basics"}, {"start": 338.0, "text": "of how to write a high performance matrix multiplication kernel"}, {"start": 341.0, "text": "with TensorOps in the M5 machine learning talk."}, {"start": 344.0, "text": "The basic approach"}, {"start": 346.0, "text": "is to slice the input matrices into smaller tiles,"}, {"start": 349.0, "text": "and then perform tile-wise matrix multiplications using TensorOps."}, {"start": 354.0, "text": "This maximizes parallelism and keeps data in the cache."}, {"start": 360.0, "text": "We can use quantization to further reduce memory traffic"}, {"start": 364.0, "text": "and fit larger models in memory."}, {"start": 367.0, "text": "In the kernel, it helps to define type aliases up front"}, {"start": 370.0, "text": "before binding the tensors."}, {"start": 372.0, "text": "Here we declare a scales factor plane with fp8_e8m0_ data type,"}, {"start": 377.0, "text": "and a block size of 32 by 1."}, {"start": 380.0, "text": "That means every 32 elements in the data plane"}, {"start": 383.0, "text": "share a single element in the scales_plane."}, {"start": 386.0, "text": "Then we declare a full tensor type, specifying an FP8 data type"}, {"start": 392.0, "text": "along with the scales_plane."}, {"start": 394.0, "text": "You can simply bind these tensors to buffer binding points."}, {"start": 398.0, "text": "The kernel will then have access to the tensors"}, {"start": 400.0, "text": "you've allocated on the host side."}, {"start": 403.0, "text": "Alternatively, if you don't want to create a full MTLTensor on the host,"}, {"start": 407.0, "text": "you can create a temporary tensor right on the shader's stack."}, {"start": 411.0, "text": "The syntax is almost identical,"}, {"start": 414.0, "text": "just swap the tag tensor_handle with tensor_inline."}, {"start": 418.0, "text": "Then pass your buffer pointers and other metadata"}, {"start": 421.0, "text": "to the tensor constructor to create a tensor on the stack."}, {"start": 426.0, "text": "As I mentioned earlier, we'll divide the problem"}, {"start": 428.0, "text": "over many threadgroups for better parallelism."}, {"start": 432.0, "text": "First, we'll slice out the tile for each threadgroup"}, {"start": 435.0, "text": "and then perform the multiplication with TensorOps."}, {"start": 439.0, "text": "To do this, simply call slice on your input and output tensors"}, {"start": 444.0, "text": "using the threadgroup ID."}, {"start": 446.0, "text": "The data and scales plane will both be sliced simultaneously"}, {"start": 450.0, "text": "according to the block size."}, {"start": 452.0, "text": "Setting up the matrix multiplication with quantized tensors"}, {"start": 456.0, "text": "is identical to normal tensors."}, {"start": 458.0, "text": "First, set up the matmul2d_descriptor,"}, {"start": 461.0, "text": "specifying the tile sizes and other parameters."}, {"start": 465.0, "text": "Then create a matmul2d op,"}, {"start": 467.0, "text": "specifying the number of simdgroups in the threadgroup."}, {"start": 471.0, "text": "Then simply pass in your quantized tensors"}, {"start": 474.0, "text": "and TensorOps will handle dequantization for you."}, {"start": 479.0, "text": "In most cases, you should feed your quantized data straight into TensorOps"}, {"start": 484.0, "text": "so that it can automatically utilize any available hardware acceleration."}, {"start": 488.0, "text": "However, if you need to dequantize a custom format,"}, {"start": 492.0, "text": "TensorOps still have you covered."}, {"start": 494.0, "text": "The simplest approach"}, {"start": 496.0, "text": "is to have each thread load a chunk of quantized data from device memory"}, {"start": 501.0, "text": "and dequantize it to f16 values in threadgroup memory."}, {"start": 506.0, "text": "You can then pass it as an inline threadgroup tensor to TensorOps."}, {"start": 510.0, "text": "However, this approach requires extra loads and stores"}, {"start": 514.0, "text": "through threadgroup memory."}, {"start": 516.0, "text": "Ideally, we would keep all this data in thread registers instead."}, {"start": 521.0, "text": "You can do this by dequantizing the data into a cooperative tensor,"}, {"start": 526.0, "text": "which can now be passed as an input to the matmul2d op."}, {"start": 530.0, "text": "Cooperative tensors distribute their storage"}, {"start": 533.0, "text": "across the thread private memory"}, {"start": 535.0, "text": "of the threads participating in the matmul operation."}, {"start": 538.0, "text": "So if you can't use quantized tensors directly,"}, {"start": 542.0, "text": "you can still skip the round trip through threadgroup memory."}, {"start": 546.0, "text": "To recap — Metal tensors natively support a wide range of quantized data types,"}, {"start": 552.0, "text": "including the new MX scaling formats"}, {"start": 554.0, "text": "and E8M0 scale factors coming in iOS and macOS 27."}, {"start": 560.0, "text": "Note that these new data types have additional alignment requirements"}, {"start": 564.0, "text": "compared to the larger data types,"}, {"start": 566.0, "text": "so be sure to check the Metal documentation for details."}, {"start": 571.0, "text": "Now let's take it up a notch —"}, {"start": 573.0, "text": "building a full, more complex custom operation with TensorOps."}, {"start": 579.0, "text": "Attention is at the core of every transformer network,"}, {"start": 582.0, "text": "including LLMs."}, {"start": 584.0, "text": "To compute attention,"}, {"start": 586.0, "text": "you first multiply two matrices together called Q and K."}, {"start": 591.0, "text": "Next, you compute SoftMax using reductions on the rows of the intermediate matrix."}, {"start": 597.0, "text": "Finally, you multiply by a third matrix called V."}, {"start": 601.0, "text": "The popular FlashAttention algorithm fuses all of these operations together"}, {"start": 606.0, "text": "into a single kernel."}, {"start": 609.0, "text": "To implement this with TensorOps,"}, {"start": 612.0, "text": "you'll first need to set up a custom simd group mapping"}, {"start": 615.0, "text": "so that each simd group owns complete rows of the intermediate matrix."}, {"start": 621.0, "text": "This allows you to compute the SoftMax"}, {"start": 623.0, "text": "without exchanging data between simd groups."}, {"start": 627.0, "text": "You can do this using the execution_simdgroup operation scope."}, {"start": 632.0, "text": "This means that each simd group"}, {"start": 634.0, "text": "will perform an independent matrix multiplication in parallel."}, {"start": 639.0, "text": "You can use the simd group ID to slice your input tiles."}, {"start": 643.0, "text": "We'll use a cooperative tensor to store the intermediate matrix"}, {"start": 646.0, "text": "so that we can use it as an input"}, {"start": 649.0, "text": "to the next step without writing it to the memory."}, {"start": 652.0, "text": "We'll compute SoftMax on the result."}, {"start": 656.0, "text": "To do this, we'll need to compute a couple of reductions"}, {"start": 659.0, "text": "on the cooperative tensor."}, {"start": 660.0, "text": "TensorOps includes a reduce_rows function to help with this."}, {"start": 665.0, "text": "Threads will exchange data amongst themselves"}, {"start": 668.0, "text": "to calculate the max for each row."}, {"start": 671.0, "text": "The result is returned in another cooperative tensor."}, {"start": 676.0, "text": "Let's set it up."}, {"start": 678.0, "text": "First, create a cooperative tensor to store the reduction output."}, {"start": 683.0, "text": "Then pass the source and destination to the reduce_rows function."}, {"start": 688.0, "text": "Here we'll use the max reduction_operation with an initial value of negative INFINITY."}, {"start": 696.0, "text": "These two cooperative tensors have different shapes,"}, {"start": 699.0, "text": "so to help map between them,"}, {"start": 701.0, "text": "TensorOps also includes a map_iterator function."}, {"start": 705.0, "text": "Given an iterator pointing to an element in the 2D tensor,"}, {"start": 709.0, "text": "it returns an iterator pointing to the corresponding element"}, {"start": 713.0, "text": "in the reduction destination."}, {"start": 715.0, "text": "First, set up a loop over the 2D cooperative tensor using iterators."}, {"start": 720.0, "text": "Then call map_iterator to map each element to its corresponding row max."}, {"start": 727.0, "text": "Finally, dereference these iterators to compute SoftMax"}, {"start": 731.0, "text": "and store the result back into the cooperative tensor."}, {"start": 735.0, "text": "Now we're ready to multiply this cooperative tensor by V."}, {"start": 740.0, "text": "In macOS 26, you would have had to first store it to threadgroup memory."}, {"start": 745.0, "text": "But it's now possible to use cooperative tensors"}, {"start": 748.0, "text": "directly as inputs to matmul operations."}, {"start": 753.0, "text": "To do this, call get_left_input_cooperative_tensor method,"}, {"start": 757.0, "text": "passing the source cooperative tensor as an argument."}, {"start": 761.0, "text": "You can then pass the result as an input to the second matmul operation."}, {"start": 767.0, "text": "One thing to watch out for:"}, {"start": 769.0, "text": "not every cooperative tensor can be reused as an input."}, {"start": 773.0, "text": "The layouts may different depending on the data types and other factors."}, {"start": 778.0, "text": "So before you do this,"}, {"start": 779.0, "text": "call the is_compatible_as_left or right _input method"}, {"start": 783.0, "text": "to check for compatibility."}, {"start": 786.0, "text": "If it returns true, you're good to go."}, {"start": 788.0, "text": "If not, you'll need to store and reload the data"}, {"start": 792.0, "text": "through threadgroup memory to convert it to the correct layout."}, {"start": 796.0, "text": "Either way, the call to op.run is the same."}, {"start": 800.0, "text": "Those are the key TensorOps features you'll need to build an advanced operation"}, {"start": 805.0, "text": "like FlashAttention using TensorOps."}, {"start": 808.0, "text": "Now that we've walked through how to build this operation,"}, {"start": 811.0, "text": "let's see how it runs in a real model using Core AI."}, {"start": 815.0, "text": "Core AI provides tools for Python developers"}, {"start": 818.0, "text": "to convert Pytorch models to Core AI models,"}, {"start": 821.0, "text": "including support for custom Metal kernels."}, {"start": 825.0, "text": "Check out the \"Deep Dive into Core AI Model authoring and Optimization\" session"}, {"start": 829.0, "text": "for the details of how to integrate a Metal kernel into a Core AI model."}, {"start": 836.0, "text": "I've followed the steps outlined in that session"}, {"start": 839.0, "text": "to integrate our custom FlashAttention kernel"}, {"start": 841.0, "text": "into a Sam3 image segmentation model."}, {"start": 845.0, "text": "We define the body of our custom attention kernel"}, {"start": 848.0, "text": "as a string in Python and register the TorchMetalKernel object, shown here."}, {"start": 854.0, "text": "Then, we replace the default huggingface attention implementation"}, {"start": 858.0, "text": "with one that calls our kernel, shown here."}, {"start": 865.0, "text": "Finally, we load the model from huggingface"}, {"start": 868.0, "text": "and export it from PyTorch as an optimized Core AI asset."}, {"start": 873.0, "text": "The export will take a moment to finish."}, {"start": 880.0, "text": "Now we're ready to do inference."}, {"start": 883.0, "text": "Sam3 performs promptable concept segmentation,"}, {"start": 887.0, "text": "so we provide the model with an image and text,"}, {"start": 890.0, "text": "and then it responds with a segmentation mask"}, {"start": 893.0, "text": "indicating where objects are located in the image."}, {"start": 897.0, "text": "Here, I'm prompting the model to label all pixels"}, {"start": 900.0, "text": "containing a car in this image."}, {"start": 904.0, "text": "Ok, now, I'll run the segmentation."}, {"start": 911.0, "text": "Looking at the final result,"}, {"start": 913.0, "text": "we can see the model correctly segmented the image."}, {"start": 916.0, "text": "The car is highlighted in blue,"}, {"start": 919.0, "text": "so our attention kernel is fully integrated"}, {"start": 921.0, "text": "into the model as expected."}, {"start": 925.0, "text": "Today, I've covered all the tools you can use"}, {"start": 928.0, "text": "to build optimized custom ML kernels on Apple Silicon."}, {"start": 932.0, "text": "From quantized data types, to advanced TensorOps features"}, {"start": 936.0, "text": "like cooperative tensors and reductions, to integrating with Core AI."}, {"start": 942.0, "text": "To go further,"}, {"start": 943.0, "text": "explore the Metal Performance Primitives documentation for the full API reference,"}, {"start": 948.0, "text": "and the programming guide for more performance optimization guidelines."}, {"start": 953.0, "text": "You can also download the TensorOps sample code"}, {"start": 956.0, "text": "to see the details that I couldn't cover here."}, {"start": 958.0, "text": "And be sure to check out the related sessions"}, {"start": 961.0, "text": "to learn more about Core AI and Metal."}, {"start": 964.0, "text": "Thank you!"}]}