Introduction

Welcome to Inside Akunu: Building a High-Performance Inference Engine for Apple Silicon. This is a book about going fast on hardware that most people barely understand. It is also a book about taking apart an inference engine – piece by piece, kernel by kernel, buffer by buffer – until you understand every decision that went into making it one of the fastest ways to run large language models on a Mac.

If you have ever wondered why your MacBook Pro can run a 70-billion-parameter model at conversational speed, or how a Metal compute shader can squeeze every last byte of bandwidth out of Apple’s unified memory, or what exactly happens between the moment you press Enter on a prompt and the moment the first token appears on screen, this book is for you.

What Is Akunu?

Akunu is a high-performance inference engine purpose-built for Apple Silicon. It runs large language models (LLMs), speech-to-text models (Whisper), and other transformer-based architectures entirely on the GPU cores of Apple’s M-series chips. It is written in C++17 (with a C API for FFI compatibility) and Metal Shading Language (MSL), with no dependencies on heavyweight frameworks like PyTorch, CoreML, or even Apple’s own MPS (Metal Performance Shaders) graph library.

The numbers speak for themselves: in decode throughput – the metric that determines how fast tokens appear on screen during generation – akunu achieves 1.83x faster decode than llama.cpp,¹ the most popular open-source inference engine for consumer hardware. This is not a marginal improvement squeezed out of micro-optimizations. It is the result of fundamentally rethinking how to map the inference workload onto Apple Silicon’s unique architecture.

Here is a snapshot of decode performance on an M4 Pro (16 GPU cores, 273 GB/s):

How does akunu achieve this? Not through any single trick, but through a relentless focus on the hardware. Every kernel is hand-tuned for Apple GPU microarchitecture. Every buffer allocation is designed around unified memory. Every dispatch decision accounts for the chip generation, the System Level Cache size, and the specific GPU core count of the machine it is running on. The entire inference pipeline – from loading model weights off disk to sampling the next token – is a single, carefully orchestrated sequence of Metal compute commands with minimal CPU overhead.

Who Is This Book For?

This book is written for people with a computer science background who want to understand how high-performance GPU programming works on Apple Silicon, and specifically how that knowledge is applied to build a state-of-the-art inference engine.

You should be comfortable with:

• C/C++ programming. Akunu’s core is C++17, with a C API for external use. You should know your way around pointers, structs, and manual memory management.

• Basic linear algebra. Matrix multiplication, dot products, transpose operations. You do not need a PhD in numerical methods, but you should know what a matrix-vector product is and why it matters for neural networks.

• General computer architecture. Caches, memory hierarchies, pipelining, parallelism. A systems programming or computer architecture course is sufficient background.

• The basics of how neural networks work. You should know what a transformer is at a high level – attention, feed-forward layers, embeddings, softmax. We will not be training models in this book; we will be running them as fast as physically possible.

You do not need prior experience with:

• Apple’s Metal API or Metal Shading Language
• GPU programming (CUDA, OpenCL, or otherwise)
• Apple Silicon internals
• The specific architectures of LLMs like Llama, Mistral, or Qwen

We will teach you all of that from the ground up.

What You Will Learn

This book is organized into nine parts, each building on the last. By the time you finish, you will have a deep understanding of every layer of the stack, from the silicon to the HTTP API.

Part I: Apple Silicon Fundamentals

Before you can write fast code for Apple Silicon, you need to understand what Apple Silicon actually is. Part I takes you on a tour of the hardware, starting with why Apple moved from Intel to ARM, then drilling into the SoC architecture, the GPU core design, the unified memory system, and the generational improvements from M1 through M4.

By the end of Part I, you will understand:

• Why Apple’s unified memory architecture is a game-changer for ML inference
• How Apple’s GPU cores differ from NVIDIA’s streaming multiprocessors
• What a SIMD group is and why it is the fundamental unit of GPU execution
• How the System Level Cache (SLC) acts as a shared last-level cache for all on-chip agents
• Why memory bandwidth – not compute – is the bottleneck for LLM decode

Part II: Metal Programming

Metal is Apple’s low-level GPU programming framework, and it is the only way to access the full power of the Apple GPU. Part II teaches you Metal from scratch: how to create compute pipelines, write shaders in the Metal Shading Language, dispatch threadgroups, manage buffers, and use SIMD group matrix operations. This is not a toy tutorial – by the end, you will understand the same programming model that akunu’s kernels use.

By the end of Part II, you will understand:

• How to set up a Metal compute pipeline and dispatch work to the GPU
• The Metal Shading Language (MSL) – a C++14-based language for writing GPU kernels
• How threadgroups, SIMD groups, and threads map to the hardware
• The Metal memory model: device memory, threadgroup memory, and how they interact
• SIMD group matrix operations (simdgroup_matrix) – Apple’s answer to NVIDIA’s Tensor Cores
• Performance optimization patterns: coalesced access, occupancy tuning, avoiding bank conflicts

Part III: Machine Learning on Metal

Part III bridges the gap between GPU programming and machine learning. We cover how tensors are represented in GPU memory, the different strategies for matrix multiplication, the attention mechanism, FlashAttention and how it maps to Metal, and quantization – the art of making models small enough to fit in memory without destroying quality.

Part IV: Akunu Architecture

Now we get to akunu itself. Part IV covers the design philosophy, the build system, the C API, the device abstraction layer, architecture descriptors (which let akunu support new model architectures without code changes), and chip configuration.

Part V: The Inference Pipeline

Part V walks through the complete inference pipeline: model loading, the dispatch table (akunu’s precompiled GPU command sequence), the prefill phase, the decode loop, and the decoding strategies: greedy/chain decode, sampled, speculative, and grammar-constrained.

Part VI: Metal Kernels Deep Dive

This is the heart of the book. Part VI takes you through every Metal kernel: GEMV, GEMM, FlashAttention, normalization, RoPE, embedding and activation, and sampling. For each, we explain the algorithm, the memory access pattern, the SIMD group coordination, and include interactive animations showing GPU execution.

Part VII: Weight Management

Models come in many file formats. Part VII covers the weight provider abstraction, the GGUF file format and akunu’s parser, SafeTensors and MLX format support, and the zoo of quantization formats (Q4_0, Q4_1, Q8_0, K-quants, MLX 3/4/6/8-bit, and more).

Part VIII: Supporting Systems

Inference engines need more than just GPU kernels. Part VIII covers the KV cache, scratch buffer architecture, the tokenizer (BPE and SentencePiece), the grammar engine (for structured output), Whisper (speech-to-text), and the HTTP server.

Part IX: Contributing to Akunu

The final part is a contributor guide: dev setup, testing, adding kernels, adding architectures, profiling, and architectural decision records.

Why This Book Exists

There are many resources for learning CUDA programming on NVIDIA GPUs. There are tutorials for PyTorch, guides for TensorRT, deep dives into NVIDIA’s Tensor Core architecture. The NVIDIA ecosystem is mature, well-documented, and widely understood.

The Apple Silicon ecosystem has… almost none of that.

If you want to understand how Apple’s GPU works at the microarchitectural level, you will find sparse documentation, a handful of WWDC sessions, and a lot of educated guesswork from the reverse-engineering community. If you want to write high-performance compute shaders for Metal, the official guides are thin on practical advice. If you want to understand how to map an LLM inference workload onto Apple Silicon efficiently, you are largely on your own.

This book exists to fill that gap. We have spent hundreds of hours profiling, benchmarking, reverse-engineering, and optimizing akunu’s Metal kernels on Apple Silicon. We have learned things about Apple’s GPU that are not documented anywhere. We want to share that knowledge so that the next person who wants to build something fast on Apple hardware does not have to start from zero.

How to Read This Book

This book is designed to be read sequentially. Each chapter builds on concepts introduced in previous chapters, and later parts assume familiarity with the hardware and programming model covered in earlier parts.

That said, here are some suggested paths depending on your background:

“I just want to understand akunu’s codebase so I can contribute.” Read Part I (skim if you already know Apple Silicon), skim Part II (read Chapter 8 on MSL carefully), then jump to Part IV and read sequentially through Part IX.

“I’m a CUDA programmer and I want to learn Metal.” Start with Part I to understand how Apple’s hardware differs from NVIDIA’s. Then read Part II carefully – it is the Metal equivalent of a CUDA programming guide. Chapter 3 (GPU architecture) and Chapter 9 (threadgroups and dispatch) are especially relevant for mapping your CUDA mental model to Metal.

“I want to understand the ML/inference side.” Skim Part I and II for context, then read Part III carefully. Then jump to Part V (the inference pipeline) and Part VI (the kernels).

“I want to understand everything.” Read the whole book, front to back. That is what it is for.

Conventions Used in This Book

Throughout this book, we use the following conventions:

• Code listings are shown in monospaced font. Metal Shading Language code is annotated with comments explaining non-obvious constructs.

• ASCII diagrams are used extensively. We chose ASCII art over images because it renders correctly in every format (web, PDF, terminal), is easy to modify, and can be included in code review comments and commit messages.

• Performance numbers are given for specific hardware configurations. Unless otherwise noted, benchmarks were run on an M4 Pro (16 GPU cores, 273 GB/s) with macOS 15 unless otherwise noted. Your numbers will differ on different hardware.

• “Apple GPU” refers to the GPU cores on Apple Silicon (M1, M2, M3, M4 families), not the older Intel integrated graphics on pre-2020 Macs.

• Register types in Metal are explained when first used. Metal uses uint, float, half (16-bit float), and others. We will explain the implications of each.

• Chip-specific details are called out in notes. When behavior differs between M1 and M4, we will tell you.

A Note on Apple’s Documentation (or Lack Thereof)

Apple is famously secretive about its hardware. The company does not publish detailed microarchitectural specifications for its GPUs, does not release ISA references for its shader cores, and does not provide the kind of performance tuning guides that NVIDIA publishes for CUDA.

This means that some of what we describe in this book – particularly regarding the internal structure of GPU cores, the exact sizes of register files, and the behavior of the instruction pipeline – is based on a combination of:

1. Apple’s public documentation² (Metal Best Practices Guide, WWDC sessions, Metal Feature Set tables)
2. Reverse engineering (running carefully constructed microbenchmarks to measure latencies, throughputs, and cache sizes)
3. Community knowledge³ (the excellent work of Dougall Johnson, Alyssa Rosenzweig, and others who have reverse-engineered Apple GPU internals)
4. Empirical observation (running akunu’s kernels with different configurations and measuring what works best)

Where we are confident in our claims, we state them as facts. Where we are making educated inferences from indirect evidence, we say so. We encourage you to verify our claims through your own experiments – that is part of the fun.

Let’s Begin

Apple Silicon is, in our opinion, the most interesting computing platform to emerge in the last decade. It combines extraordinary hardware – a unified memory architecture, a surprisingly powerful GPU, and a custom interconnect fabric – with a programming model (Metal) that is powerful but underexplored. The gap between what this hardware can do and what most software actually does with it is enormous.

Akunu exists to close that gap for inference workloads. This book exists to show you how.

Turn the page. Let’s talk about the silicon.

“Any sufficiently advanced technology is indistinguishable from magic.” — Arthur C. Clarke

Our goal is to make the magic distinguishable.

1. Benchmarks run on M4 Pro (16 GPU cores, 273 GB/s) with llama.cpp b8610, 3 reps per config. Decode measured at tg128 (128-token generation). The 1.83x is the average across 19 GGUF models. Speedup is highest on small models (3.66x on Qwen3-0.6B Q3_K_S) where per-token dispatch overhead dominates, and converges to parity on larger models (0.98x on Qwen3-8B Q8_0). Prefill averages 0.91x. See BENCHMARKS.md for the full data and Chapter 55 for methodology. ↩

2. Apple. “Metal Best Practices Guide.” developer.apple.com. The primary official reference for Metal compute optimization. See https://developer.apple.com/library/archive/documentation/3DDrawing/Conceptual/MTLBestPracticesGuide/index.html. ↩

3. Grinberg, D. et al. “Reverse-engineering Apple GPU cores.” Asahi Linux project, 2022. The most detailed public analysis of Apple GPU internals. See https://dougallj.github.io/applegpu/. ↩

The Apple Silicon Revolution

On June 22, 2020, at the Worldwide Developers Conference, Apple CEO Tim Cook announced that the Mac was transitioning from Intel x86 processors to Apple’s own custom ARM-based chips.¹ This was not a surprise to industry watchers – rumors had been circulating for years – but the ambition and speed of the transition stunned everyone.² Apple was not just switching CPU vendors. They were bringing the entire system-on-chip (SoC) approach that had made the iPhone and iPad dominant in mobile computing to the desktop and laptop. Within two years, every Mac in the lineup would run on Apple Silicon.

This chapter tells the story of how we got here, why it matters, and what it means for the kind of workload we care about most in this book: running large language models as fast as possible on local hardware.

Why Apple Left Intel

To understand why Apple Silicon exists, you need to understand why Apple was unhappy with Intel. The relationship between the two companies began in 2005, when Steve Jobs announced the Mac’s transition from PowerPC to Intel x86. At the time, Intel’s chips offered a compelling combination of performance and power efficiency, and the x86 ecosystem was the dominant force in personal computing.

For about a decade, this partnership worked well. Intel’s tick-tock cadence – alternating between process shrinks and microarchitecture improvements – delivered steady performance gains. But starting around 2015, things began to go wrong.

Intel’s Process Stagnation

Intel’s manufacturing process, once the envy of the semiconductor industry, hit a wall. The company’s 10nm node (which they originally planned to ship in 2016) was delayed repeatedly. The 14nm process that was supposed to be a one-generation stopgap ended up being stretched across four generations of products:

  Intel's 14nm Purgatory
  =======================

  2014: Broadwell    (14nm)   -- on schedule
  2015: Skylake      (14nm)   -- ok, next year we move to 10nm
  2016: Kaby Lake    (14nm+)  -- 10nm delayed, here's a refinement
  2017: Coffee Lake  (14nm++) -- still no 10nm, more cores this time
  2018: Coffee Lake  (14nm++) -- yeah, 10nm is still not ready
  2019: Ice Lake     (10nm)   -- finally! but only in low-power laptops
  2020: Rocket Lake  (14nm)   -- back to 14nm for desktops. really.

Each year that Intel stayed on 14nm, the performance improvements got smaller. The company resorted to increasing power consumption and adding cores to show benchmark improvements, but the fundamental per-core performance was stagnating.

Meanwhile, Apple’s in-house chip team – born from the 2008 acquisition of PA Semi, a boutique chip design firm – was making ARM-based processors that improved by 20-40% in performance every single year. By 2020, Apple’s A14 chip (in the iPhone 12) was competitive with Intel’s laptop processors in single-threaded performance, despite running at a fraction of the power.

The Power Problem

For a laptop maker like Apple, power efficiency is not a nice-to-have; it is the single most important metric. Every watt of power consumed becomes a watt of heat that must be dissipated. More heat means bigger fans, thicker chassis, and shorter battery life. Intel’s chips were designed primarily for desktops and servers, then adapted for laptops. Apple wanted chips designed for laptops first.

  Power Efficiency Comparison (circa 2020)
  =========================================

  Intel Core i7 (10th gen laptop):
  +----------------------------+
  | TDP: 28W (actual: 35-50W)  |
  | Performance: ████████      |
  | Per-watt:    ███           |
  +----------------------------+

  Apple A14 (iPad/iPhone):
  +----------------------------+
  | TDP: ~5-6W                 |
  | Performance: ██████        |
  | Per-watt:    █████████████ |
  +----------------------------+

  Apple wanted to bring that per-watt advantage to the Mac.

The gap in power efficiency was not just about process technology (though TSMC’s nodes were ahead of Intel’s). It was about fundamental architectural choices. ARM’s instruction set, which we will discuss shortly, is inherently more power-efficient than x86. And Apple’s custom microarchitecture was designed from the ground up for efficiency in a way that Intel’s x86 legacy made difficult.

The Integration Advantage

Perhaps the most important reason Apple moved to its own silicon was integration. With Intel chips, a Mac was a collection of discrete components: an Intel CPU, a separate AMD or NVIDIA GPU (or Intel integrated graphics), a separate ISP (image signal processor), a separate Thunderbolt controller, a separate security chip (T1, then T2). Each of these communicated over various buses and protocols, with all the latency and power overhead that implies.

  Intel-era Mac Architecture (simplified)
  =========================================

  +--------+     PCIe     +------------+
  | Intel  |<------------>| AMD/NVIDIA |
  |  CPU   |              | Discrete   |
  |        |              |   GPU      |
  +---+----+              +------------+
      |                        |
      | DDR4                   | GDDR6
      v                        v
  +--------+              +-----------+
  | System |              |   VRAM    |
  | Memory |              | (separate)|
  | (RAM)  |              |           |
  +--------+              +-----------+

  Data must cross PCIe bus to go between CPU and GPU.
  GPU has its own memory (VRAM), separate from system RAM.
  Bandwidth between CPU and GPU is limited (~16 GB/s PCIe 3.0).

Apple’s SoC approach puts everything on a single chip, sharing a single pool of memory. We will explore this unified memory architecture in detail in Chapter 4, but the key insight is this: when the CPU and GPU share the same memory, you eliminate the need to copy data between them. For machine learning workloads, where models can be tens of gigabytes, this is transformative.

The ARM Instruction Set Architecture

Before we dive into Apple Silicon specifically, let’s talk about ARM – the instruction set architecture (ISA) that forms the foundation of every Apple Silicon chip.

ARM stands for “Advanced RISC Machines” (originally “Acorn RISC Machine,” after the British company that designed the first ARM processor in 1985). ARM does not manufacture chips; it designs instruction set architectures and licenses them to chip makers. Apple licenses the ARMv8-A architecture (ARMv8.5-A for M1, ARMv8.6-A for M2/M3, ARMv9.2-A for M4) and then designs its own custom microarchitecture that implements that ISA.

This distinction is important: the ISA defines what instructions the processor understands. The microarchitecture defines how those instructions are executed. Two chips can implement the same ISA with radically different performance characteristics. Apple’s custom ARM cores (codenamed Firestorm, Icestorm, Avalanche, Blizzard, Everest, Sawtooth, etc.) are not the same as Qualcomm’s Kryo cores or ARM’s own Cortex cores, even though they all execute ARM instructions.

RISC vs. CISC

ARM is a RISC (Reduced Instruction Set Computer) architecture.³ Intel’s x86 is a CISC (Complex Instruction Set Computer) architecture. This is one of the most important distinctions in processor design, and it has deep implications for power efficiency and performance.

CISC (x86) philosophy: Provide a large number of complex instructions. A single instruction might load data from memory, perform an arithmetic operation, and store the result back to memory. The idea is to minimize the number of instructions the compiler needs to generate, reducing code size.

RISC (ARM) philosophy: Provide a smaller number of simple instructions. Each instruction does one thing: load data, perform arithmetic, or store data, but not a combination. The idea is that simpler instructions can be executed faster, and the hardware can be simpler (and therefore more power-efficient).

  CISC vs. RISC: Adding Two Numbers from Memory
  ===============================================

  x86 (CISC):
  +-----------------------------------------+
  | ADD [mem_a], [mem_b]                    |  <-- One instruction
  |                                         |      but internally decoded
  | Internally becomes:                     |      into multiple micro-ops
  |   load temp1, [mem_a]                   |
  |   load temp2, [mem_b]                   |
  |   add  temp1, temp1, temp2              |
  |   store [mem_a], temp1                  |
  +-----------------------------------------+

  ARM (RISC):
  +-----------------------------------------+
  | LDR  X0, [mem_a]                        |  <-- Load first operand
  | LDR  X1, [mem_b]                        |  <-- Load second operand
  | ADD  X0, X0, X1                         |  <-- Add them
  | STR  X0, [mem_a]                        |  <-- Store result
  +-----------------------------------------+

  ARM: 4 instructions, each simple and predictable
  x86: 1 instruction, but complex internal decoding required

In practice, modern x86 processors bridge the gap by decoding CISC instructions into internal RISC-like micro-operations (micro-ops). But this decoding step consumes die area and power. ARM processors skip this step entirely because their instructions are already simple. This is one of the main reasons ARM chips are more power-efficient.

Fixed-Width Instructions

One of ARM’s most important characteristics is that all instructions are the same width: 32 bits (4 bytes). (ARM also supports a 16-bit “Thumb” mode for code density, but Apple Silicon runs in AArch64 mode where all instructions are 32 bits.)

This might seem like a minor detail, but it has profound implications for the processor’s front end – the part of the chip that fetches and decodes instructions.

  x86 Variable-Length Instructions
  =================================

  Memory layout of x86 code:
  +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
  |  1 byte |  3 bytes   |   7 bytes         | 2 bytes|
  |  instr  |  instr     |   instr           | instr  |
  +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+

  The processor cannot tell where one instruction ends and
  the next begins without decoding each instruction sequentially.
  This is hard to parallelize.


  ARM Fixed-Width Instructions
  ==============================

  Memory layout of ARM code:
  +----------+----------+----------+----------+----------+
  | 4 bytes  | 4 bytes  | 4 bytes  | 4 bytes  | 4 bytes  |
  |  instr   |  instr   |  instr   |  instr   |  instr   |
  +----------+----------+----------+----------+----------+

  Every instruction starts at a 4-byte boundary.
  The processor can easily fetch and decode multiple
  instructions in parallel. No ambiguity about boundaries.

For Intel, the variable-length instruction format means the front-end decode logic is one of the most complex (and power-hungry) parts of the chip. The decoder must examine each byte to determine the instruction length before it can find the start of the next instruction. Intel has thrown enormous engineering resources at this problem (instruction caches, micro-op caches, complex decode pipelines), but it remains an inherent disadvantage.

The Register File

ARM’s AArch64 architecture provides 31 general-purpose 64-bit registers (X0-X30), compared to x86-64’s 16 general-purpose registers (RAX, RBX, RCX, RDX, RSI, RDI, RBP, RSP, R8-R15). Having more registers means less “register spilling” – the costly process of saving register values to the stack when the processor runs out of registers.

  Register Files Compared
  ========================

  ARM AArch64:                    x86-64:
  +------+------+------+         +------+------+
  |  X0  |  X1  |  X2  |         | RAX  | RBX  |
  +------+------+------+         +------+------+
  |  X3  |  X4  |  X5  |         | RCX  | RDX  |
  +------+------+------+         +------+------+
  |  X6  |  X7  |  X8  |         | RSI  | RDI  |
  +------+------+------+         +------+------+
  |  X9  |  X10 |  X11 |         | RBP  | RSP  |
  +------+------+------+         +------+------+
  |  X12 |  X13 |  X14 |         | R8   | R9   |
  +------+------+------+         +------+------+
  |  X15 |  X16 |  X17 |         | R10  | R11  |
  +------+------+------+         +------+------+
  |  X18 |  X19 |  X20 |         | R12  | R13  |
  +------+------+------+         +------+------+
  |  X21 |  X22 |  X23 |         | R14  | R15  |
  +------+------+------+         +------+------+
  |  X24 |  X25 |  X26 |           16 registers
  +------+------+------+
  |  X27 |  X28 |  X29 |
  +------+------+------+
  |  X30 (LR)   |  SP  |
  +------+------+------+
    31 registers + SP

  Plus 32 SIMD/FP registers (V0-V31) on both architectures,
  but ARM's are 128-bit (NEON) with optional SVE extensions.

ARM also specifies a dedicated link register (X30/LR) for function return addresses, which simplifies function call conventions and branch prediction.

NEON and Advanced SIMD

ARM includes NEON, a SIMD (Single Instruction, Multiple Data) extension that operates on 128-bit vector registers. NEON can process multiple data elements simultaneously:

  NEON SIMD Operation: Adding 4 floats at once
  ==============================================

  V0: | float_a0 | float_a1 | float_a2 | float_a3 |  (128 bits)
      +----------+----------+----------+----------+
         +            +            +            +
  V1: | float_b0 | float_b1 | float_b2 | float_b3 |  (128 bits)
      +----------+----------+----------+----------+
         =            =            =            =
  V2: | float_c0 | float_c1 | float_c2 | float_c3 |  (128 bits)
      +----------+----------+----------+----------+

  FADD V2.4S, V0.4S, V1.4S   -- One instruction, four additions

Apple’s implementation of NEON is extremely wide and high-throughput. The performance cores (P-cores) on M-series chips can execute up to four 128-bit NEON operations per cycle, giving them enormous SIMD throughput for CPU-side workloads.

For the purposes of this book, NEON is relevant primarily for CPU-side preprocessing (tokenizer operations, weight format conversion) rather than the main inference workload, which runs on the GPU. But it is good to know it is there.

AMX: Apple’s Secret Matrix Coprocessor

Here is something you will not find in the official ARM specification: Apple’s M-series chips include a custom matrix coprocessor called AMX (Apple Matrix Extensions). AMX is not part of the ARM ISA – it is a proprietary Apple extension that provides hardware-accelerated matrix multiplication on the CPU side.

AMX operates on its own set of registers and can perform operations like multiplying two 16x16 matrices of FP16 values in a single instruction. Apple uses AMX internally in the Accelerate framework and CoreML, but it is not officially documented or exposed as a public API.

  AMX: Apple's Matrix Coprocessor (undocumented)
  ================================================

  +---------------------+
  | Apple CPU Core      |
  |                     |
  |  +------+  +------+ |
  |  | ALU  |  | NEON | |
  |  +------+  +------+ |
  |                     |
  |  +---------------+  |
  |  |      AMX      |  |  <-- Custom matrix coprocessor
  |  | 16x16 FP16    |  |      Not part of ARM ISA
  |  | matrix mul    |  |      Undocumented by Apple
  |  | in one cycle  |  |      Used by Accelerate/CoreML
  |  +---------------+  |
  |                     |
  +---------------------+

  AMX registers are separate from the ARM register file.
  AMX instructions are encoded as system register writes.

We mention AMX here for completeness, but akunu does not use it. Akunu runs the entire inference workload on the GPU, where the much larger number of ALUs and the higher memory bandwidth make GPU compute far more efficient for the matrix operations that dominate LLM inference. AMX is more useful for smaller matrix operations that need low latency (like real-time audio processing).

The A-Series Lineage

Apple Silicon did not spring into existence fully formed. It is the culmination of over a decade of chip design work that began with the original iPhone.

The Mobile Era (2010-2019)

Apple started designing its own chips with the A4, which powered the original iPad and iPhone 4 in 2010. Each generation brought significant improvements:

  The A-Series Evolution
  =======================

  A4  (2010) -- First Apple-designed SoC. ARM Cortex-A8 based.
       |        Single core, 45nm Samsung process.
       v
  A5  (2011) -- Dual-core ARM Cortex-A9. First Apple GPU (PowerVR).
       |
       v
  A6  (2012) -- First CUSTOM Apple CPU core (Swift). No more ARM reference
       |        designs. This is where Apple diverged from everyone else.
       v
  A7  (2013) -- First 64-bit mobile processor. EVER. Caught the industry
       |        off guard. Desktop-class architecture in a phone.
       v
  A8  (2014) -- 20nm TSMC. Performance and efficiency improvements.
       |
       v
  A9  (2015) -- Custom "Twister" core. Huge IPC gains.
       |
       v
  A10 (2016) -- First big.LITTLE configuration (2 perf + 2 efficiency cores)
       |        "Fusion" architecture.
       v
  A11 (2017) -- First Apple-designed GPU. Dropped PowerVR.
       |        First Neural Engine (2-core, 600B ops/sec).
       v
  A12 (2018) -- 7nm TSMC. 8-core Neural Engine (5 trillion ops/sec).
       |        "Bionic" branding begins.
       v
  A13 (2019) -- "Bionic." Fastest mobile chip by a huge margin.
       |        8-core Neural Engine (6 trillion ops/sec).
       v
  A14 (2020) -- 5nm TSMC. First 5nm chip in any consumer device.
       |        Performance competitive with Intel laptop chips.
       v
  M1  (2020) -- A14 core design, scaled up for Mac.
                The revolution begins.

Several key milestones in this lineage are directly relevant to Apple Silicon and akunu:

A6 (2012): Custom CPU cores. Apple stopped using ARM’s reference CPU designs and started designing its own microarchitectures. This gave Apple the freedom to make architectural decisions optimized for their specific use cases, and it is why Apple’s ARM cores consistently outperform everyone else’s ARM cores.

A7 (2013): 64-bit ARM. Apple was the first company to ship a 64-bit ARM processor in a consumer device. Qualcomm and Samsung scrambled to catch up. The 64-bit address space would later be essential for the large memory configurations in M-series chips.

A11 (2017): Custom GPU. Apple designed its own GPU, replacing the PowerVR cores it had licensed for years. This custom GPU design is the direct ancestor of the GPU cores in every M-series chip, and understanding its architecture is critical for understanding akunu’s Metal kernels.

A11 (2017): First Neural Engine. Apple’s first dedicated ML accelerator appeared in the A11. While akunu does not use the Neural Engine (it targets the GPU for maximum flexibility and performance), the Neural Engine’s existence shows how seriously Apple takes on-device ML.

From Phone to Mac: The M1 Announcement

On November 10, 2020, Apple announced the M1 chip and the first three Macs to use it: the MacBook Air, MacBook Pro 13-inch, and Mac mini. The M1 was not just an A14 chip with a different name – it was the A14’s core design scaled up for the thermal and power envelope of a laptop:

  A14 (iPhone) vs. M1 (Mac)
  ===========================

  A14:                              M1:
  +---------------------+          +-------------------------------+
  | 2 Perf + 4 Eff cores|          | 4 Perf + 4 Eff cores          |
  | 4 GPU cores         |          | 7 or 8 GPU cores              |
  | 16-core Neural Eng  |          | 16-core Neural Engine         |
  | 4GB/6GB RAM         |          | 8GB/16GB Unified Memory       |
  | ~4W power budget    |          | ~15-20W power budget          |
  | LPDDR4X             |          | LPDDR4X, 128-bit bus          |
  | 11.8B transistors   |          | 16B transistors               |
  | 88mm^2 die size     |          | 120mm^2 die size              |
  +---------------------+          +-------------------------------+

  Same CPU core design (Firestorm + Icestorm)
  Same GPU core design
  But more of everything, with more memory and bandwidth

The M1 was a revelation. It offered:

• Better CPU performance than Intel’s Core i9 in most workloads
• Better GPU performance than Intel’s Iris integrated graphics (and competitive with low-end discrete GPUs)
• 20-hour battery life in the MacBook Air
• Fanless operation in the MacBook Air (the chip ran cool enough to not need a fan)
• Unified Memory Architecture with zero-copy sharing between CPU and GPU

The tech industry was stunned. Not because Apple had made a fast chip – the A-series trajectory made that predictable – but because the first-generation Mac chip was this good. It was not a compromise “good enough for the Mac” chip; it was genuinely better than Intel’s best laptop offerings in almost every measurable way.

What Makes Apple Silicon Different

Now that we understand the history, let’s talk about what makes Apple Silicon architecturally unique. There are several key differentiators that matter for our purposes:

1. System-on-Chip (SoC) Integration

Unlike traditional PCs where the CPU, GPU, memory controller, I/O controllers, and other components are separate chips on a motherboard, Apple Silicon puts everything on a single die (or, for Ultra variants, two dies connected by a high-speed interconnect).

  Traditional PC Architecture          Apple Silicon SoC
  ==========================          ==================

  +------+   +------+   +------+      +---------------------------+
  | CPU  |   | GPU  |   | WiFi |      |  CPU   GPU    Neural Eng  |
  +--+---+   +--+---+   +--+---+      |                           |
     |          |          |          |  Media  ISP   Secure Enc  |
  ===+====PCIe==+====USB===+====      |                           |
     |                                |  Memory    Thunderbolt    |
  +--+---------------------------+    |  Controller  Controller   |
  |        Motherboard           |    |                           |
  |  +--------+  +-----------+   |    |  I/O     SLC    Fabric    |
  |  | Memory |  | Storage   |   |    +---------------------------+
  |  |  DIMMs |  | Controller|   |          |
  |  +--------+  +-----------+   |    +---------------------------+
  +------------------------------+    |    Unified Memory         |
                                      |    (LPDDR on-package)     |
  Multiple chips, multiple buses,     +---------------------------+
  multiple memory pools
                                      One chip, one memory pool,
                                      one interconnect fabric

This integration has several benefits:

• Lower latency: Components communicate over an on-die fabric instead of PCIe or USB. On-die communication can be an order of magnitude lower latency than crossing a PCIe bus.

• Lower power: Driving signals across a PCIe bus or memory DIMMs requires much more power than on-die communication. The shorter the wire, the less power it takes.

• Higher bandwidth: Apple can build wide, custom interconnects between components because they are all on the same die. The M4 Max, for example, has a 512-bit memory bus – wider than what you would typically see on a discrete GPU.

2. Unified Memory Architecture (UMA)

This is the single most important differentiator for ML workloads, and we will devote all of Chapter 4 to it. The short version: on Apple Silicon, the CPU and GPU share the same physical memory. There is no “system RAM” and “VRAM” – there is just memory, and both the CPU and GPU can access it.

For LLM inference, this means:

• No copying model weights between CPU and GPU. When akunu loads a model from disk, the CPU reads the file into a Metal buffer, and the GPU can immediately access those weights without any data transfer.

• No VRAM limitation. On a discrete GPU, your model must fit in VRAM (12GB, 24GB, etc.). On Apple Silicon, your model can use the entire unified memory pool – up to 192GB on the M4 Max.

• CPU post-processing is free. After the GPU generates logits, the CPU can read the output buffer directly without waiting for a DMA transfer.

  Why UMA Matters for LLM Inference
  ===================================

  Discrete GPU (NVIDIA):
  1. Load model from disk -> CPU memory (RAM)      [slow: disk I/O]
  2. Copy weights from RAM -> GPU VRAM (PCIe)      [slow: PCIe ~32 GB/s]
  3. GPU computes inference using VRAM              [fast: HBM ~900 GB/s]
  4. Copy results from VRAM -> RAM (PCIe)           [slow: PCIe ~32 GB/s]
  5. CPU reads results from RAM                     [fast: DDR5 ~50 GB/s]

  Apple Silicon (UMA):
  1. Load model from disk -> unified memory         [slow: disk I/O]
  2. GPU computes inference using same memory       [fast: ~273-546 GB/s]
  3. CPU reads results from same memory             [fast: same bus]

  Steps 2 and 4 from the discrete GPU path simply do not exist on Apple Silicon.

3. Efficiency Cores and Performance Cores (big.LITTLE)

Apple Silicon uses a heterogeneous CPU design with two types of cores:

• Performance cores (P-cores): Wide, deep-pipeline, out-of-order cores designed for maximum single-threaded performance. These are among the fastest CPU cores in the world. They consume more power.

• Efficiency cores (E-cores): Narrower, simpler cores designed for maximum performance per watt. They handle background tasks and light workloads, allowing the P-cores to stay idle (and powered down) when they are not needed.

  CPU Cluster Architecture
  =========================

  +-------------------------------------------+
  |  Performance Cluster                      |
  |  +---------+ +---------+ +---------+ +--+ |
  |  | P-core  | | P-core  | | P-core  | |..| |
  |  | Wide    | | Wide    | | Wide    | |  | |
  |  | OoO     | | OoO     | | OoO     | |  | |
  |  | 192KB   | | 192KB   | | 192KB   | |  | |
  |  | L1      | | L1      | | L1      | |  | |
  |  +---------+ +---------+ +---------+ +--+ |
  |         Shared L2 Cache (12-16MB)         |
  +-------------------------------------------+

  +-------------------------------------------+
  |  Efficiency Cluster                       |
  |  +-------+ +-------+ +-------+ +-------+  |
  |  |E-core | |E-core | |E-core | |E-core |  |
  |  |Narrow | |Narrow | |Narrow | |Narrow |  |
  |  |In-ord | |In-ord | |In-ord | |In-ord |  |
  |  | 128KB | | 128KB | | 128KB | | 128KB |  |
  |  | L1    | | L1    | | L1    | | L1    |  |
  |  +-------+ +-------+ +-------+ +-------+  |
  |         Shared L2 Cache (4MB)             |
  +-------------------------------------------+

For akunu, the P-cores matter during model loading and weight rearrangement (CPU-side work), while the GPU handles the actual inference computation. The E-cores might handle I/O and tokenization during inference.

4. The Custom GPU

Apple’s GPU is a tile-based deferred renderer (TBDR) designed primarily for mobile graphics, but it turns out to be surprisingly capable for compute workloads as well. We will cover the GPU architecture in detail in Chapter 3, but the key points are:

• Apple’s GPU uses 32-thread SIMD groups (similar to NVIDIA’s 32-thread warps)
• It has on-chip tile memory (similar to NVIDIA’s shared memory)
• It supports SIMD group matrix operations for hardware-accelerated matrix multiplication
• It shares the same unified memory as the CPU, with no separate VRAM

The GPU is where akunu spends the vast majority of its time, and understanding its architecture is the foundation for understanding akunu’s performance.

5. The Neural Engine

Apple Silicon includes a dedicated Neural Engine – a specialized accelerator for neural network inference. The Neural Engine is optimized for a specific set of operations (convolutions, matrix multiplications, etc.) and can be more power-efficient than the GPU for those operations.

However, akunu does not use the Neural Engine. Here is why:

• Limited programmability: The Neural Engine is accessed through CoreML, which abstracts away the hardware details. You cannot write custom kernels for it.
• Limited precision support: The Neural Engine is optimized for FP16 and INT8 operations. Many of akunu’s kernels need mixed-precision arithmetic.
• Limited flexibility: The Neural Engine is designed for standard neural network layers. Custom operations (like akunu’s fused kernels) cannot run on it.
• GPU is fast enough: Apple’s GPU, with SIMD group matrix operations and high memory bandwidth, is more than fast enough for inference when properly optimized.

  Why Akunu Uses the GPU, Not the Neural Engine
  ===============================================

  Neural Engine:
  +-------------------+
  | + Power efficient  |
  | + Good for std ops |
  | - No custom kernels|
  | - CoreML only      |
  | - Fixed data types |
  | - Limited control  |
  +-------------------+

  GPU:
  +-------------------+
  | + Full control     |
  | + Custom kernels   |
  | + Any data type    |
  | + High bandwidth   |
  | + SIMD group ops   |
  | + Threadgroup mem  |
  | - More power       |
  +-------------------+

  For maximum performance with custom operations, the GPU wins.

How Apple Silicon Changes the Game for Local Inference

With the background in place, let’s talk about why Apple Silicon is uniquely well-suited for running LLMs locally – and why akunu exists.

The Memory Wall

The fundamental bottleneck in LLM inference is not compute; it is memory bandwidth. During the decode phase (generating one token at a time), the model must read its entire weight matrix for every single token. For a 7-billion-parameter model in 4-bit quantization, that is roughly 3.5GB of data that must be read from memory for every single token generated.

The speed at which you can read data from memory – the memory bandwidth – determines the upper bound on your decode speed. This is the “memory wall.”

  The Memory Wall: Why Bandwidth Matters
  ========================================

  Model: Llama 3.1 8B (Q4_0 quantization)
  Weight size: ~4.3 GB
  Operation per token: read all weights + compute

  If memory bandwidth is B (GB/s) and model size is S (GB):
  Maximum theoretical tokens/sec = B / S

  +-------------------+----------+-----------------------+
  | Hardware          | BW (GB/s)| Max decode (tok/s)    |
  +-------------------+----------+-----------------------+
  | DDR4 laptop       |    38    |   38/4.3 =  ~9 tok/s  |
  | M1 (LPDDR4X)      |    68    |   68/4.3 = ~16 tok/s  |
  | RTX 3090 (HBM)    |   936    |  936/4.3 = ~218 tok/s |
  | M4 Pro (LPDDR5X)  |   273    |  273/4.3 = ~63 tok/s  |
  | M4 Max (LPDDR5X)  |   546    |  546/4.3 = ~127 tok/s |
  +-------------------+----------+-----------------------+

  Note: These are THEORETICAL MAXIMUMS. Actual performance depends on
  how efficiently the software uses the available bandwidth.
  Akunu gets remarkably close to these theoretical limits.

Apple Silicon’s advantage here is threefold:

1. High bandwidth LPDDR5X memory. The M4 Max provides 546 GB/s of memory bandwidth, which is in the range of older high-end discrete GPUs.

2. No PCIe bottleneck. On an NVIDIA system, even if the GPU has 900+ GB/s of HBM bandwidth, the model weights must first cross the 32-64 GB/s PCIe bus to get from system RAM to VRAM. If the model does not fit in VRAM, you hit the PCIe wall. On Apple Silicon, there is no such bottleneck – the full memory bandwidth is available to the GPU.

3. Large memory capacity. The M4 Max supports up to 128GB of unified memory. You can run a 70B model in 4-bit quantization (~35GB) without any model splitting or offloading.

The Total-Cost-of-Ownership Argument

There is also a practical economic argument for Apple Silicon inference. An NVIDIA A100 or H100 GPU costs thousands of dollars (or hundreds per month in cloud rental), requires a dedicated server, consumes hundreds of watts, and needs active cooling. A MacBook Pro with an M4 Max chip can run significant models while sitting on your lap, running on battery, making no noise.

For personal use, development, testing, and small-scale deployment, Apple Silicon offers a compelling total cost of ownership that discrete GPU setups cannot match.

What Akunu Exploits

Akunu is designed from the ground up to exploit Apple Silicon’s unique characteristics:

  Akunu's Hardware-Aware Design
  ==============================

  Apple Silicon Feature         Akunu Exploitation
  ========================      ================================
  Unified Memory (UMA)     -->  Zero-copy weight loading.
                                CPU rearranges weights directly
                                in GPU-accessible buffers.

  High memory bandwidth    -->  Bandwidth-optimized kernels.
                                GEMV kernels saturate the
                                memory bus.

  SIMD group matrix ops    -->  Hardware-accelerated matrix
                                multiplication in GEMM and
                                attention kernels.

  System Level Cache (SLC) -->  Kernel tiling tuned to SLC
                                size per chip generation.

  GPU core count varies    -->  Dispatch parameters tuned per
                                chip (M1 vs M4 Pro vs M4 Max).

  Threadgroup memory       -->  Cooperative kernels that share
                                data between SIMD groups via
                                fast on-chip memory.

A Roadmap of What’s Coming

Let’s close this chapter with a preview of the journey ahead. Here is the conceptual stack we will build up over the course of this book:

  The Akunu Stack (bottom to top)
  ================================

  Layer 7: Application
  +----------------------------------------------------------+
  |  HTTP Server  |  CLI  |  Swift/Python Bindings           |
  +----------------------------------------------------------+
  Layer 6: Decoding Strategies
  +----------------------------------------------------------+
  |  Greedy | Sampled | Speculative | Grammar-Constrained    |
  +----------------------------------------------------------+
  Layer 5: Inference Pipeline
  +----------------------------------------------------------+
  |  Model Loading | Prefill | Decode Loop | KV Cache        |
  +----------------------------------------------------------+
  Layer 4: Metal Kernels
  +----------------------------------------------------------+
  |  GEMV | GEMM | FlashAttn | RMSNorm | RoPE | Sampling     |
  +----------------------------------------------------------+
  Layer 3: Metal Compute Framework
  +----------------------------------------------------------+
  |  Pipelines | Buffers | Command Encoders | Threadgroups   |
  +----------------------------------------------------------+
  Layer 2: Apple GPU Architecture
  +----------------------------------------------------------+
  |  GPU Cores | SIMD Groups | Tile Memory | Register File   |
  +----------------------------------------------------------+
  Layer 1: Apple Silicon SoC
  +----------------------------------------------------------+
  |  CPU | GPU | Neural Engine | Memory | SLC | Fabric       |
  +----------------------------------------------------------+
  Layer 0: Silicon
  +----------------------------------------------------------+
  |  TSMC 3nm/5nm Process | Transistors | Interconnect       |
  +----------------------------------------------------------+

  We start at Layer 0/1 (this chapter and the next four)
  and work our way up to Layer 7.

In the next chapter, we will zoom into Layer 1 and examine the System-on-Chip architecture in detail. We will look at every component on the Apple Silicon die and understand how they work together.

Summary

Let’s recap what we covered in this chapter:

• Apple transitioned from Intel to ARM because Intel’s process technology stagnated, ARM offers better power efficiency, and Apple’s SoC approach enables superior integration.

• The ARM ISA is a RISC architecture with fixed-width 32-bit instructions, 31 general- purpose registers, NEON SIMD extensions, and (on Apple chips) the undocumented AMX matrix coprocessor.

• Apple’s chip design lineage stretches from the A4 (2010) through the A14 (2020) to the M1 and beyond, with key milestones including custom CPU cores (A6), 64-bit ARM (A7), custom GPU (A11), and the Neural Engine (A11).

• Apple Silicon’s key differentiators include SoC integration, unified memory architecture, heterogeneous CPU cores, a custom TBDR GPU, and a dedicated Neural Engine.

• For LLM inference, the critical factors are memory bandwidth (the memory wall), the absence of a PCIe bottleneck (UMA), and the large unified memory pool.

• Akunu exploits UMA for zero-copy weight loading, high bandwidth for saturating GEMV kernels, SIMD group matrix ops for hardware-accelerated matmul, and per-chip tuning for optimal dispatch parameters.

Next up: let’s crack open the chip and see what is inside.

1. Apple. “Apple announces M1.” apple.com, November 2020. The original announcement detailing unified memory architecture and performance-per-watt claims. See https://www.apple.com/newsroom/2020/11/apple-unleashes-m1/. ↩

2. Turley, J. “Apple Ignites the ARM Mac.” Microprocessor Report, 2020. Analysis of Apple’s transition from Intel to ARM and the architectural advantages of the M1. See https://www.linleygroup.com/. ↩

3. Patterson, D. and Hennessy, J. “Computer Architecture: A Quantitative Approach.” 6th Edition. The foundational text on RISC vs CISC tradeoffs, pipeline design, and the memory wall. See https://www.elsevier.com/books/computer-architecture/hennessy/978-0-12-811905-1. ↩

System-on-Chip Architecture

If you’ve spent your career thinking about computers as a collection of separate chips on a motherboard — a CPU here, a GPU there, RAM sticks in their slots — then Apple Silicon is going to fundamentally rewire how you think about hardware. Everything lives on one die (or two, in the Ultra variants). And that changes everything about how we write high-performance inference code.

What Is a System-on-Chip?

A System-on-Chip (SoC) integrates what traditionally were separate components — processor, graphics, memory controller, I/O — onto a single piece of silicon. This isn’t a new idea; your phone has been running on SoCs for over a decade. What Apple did with the M-series was bring this approach to laptop and desktop-class performance.

In a traditional PC, the CPU has its own RAM (DDR5, ~90 GB/s) and the GPU has its own VRAM (GDDR6X, ~1 TB/s). Data must cross the PCIe bus (~32 GB/s) to move between them. Two separate memory pools, two separate bandwidth domains.

On Apple Silicon, everything — CPU, GPU, Neural Engine — shares one pool of LPDDR5 memory (120-819 GB/s depending on chip). No copy, no PCIe bottleneck. The GPU reads the same bytes the CPU wrote.

┌─────────────────────────────────────────────────────────────┐
│                    TRADITIONAL PC ARCHITECTURE              │
│                                                             │
│  ┌──────────┐    PCIe x16    ┌────────────────┐             │
│  │   CPU    │◄──────────────►│  Discrete GPU  │             │
│  │ (Intel/  │    ~32 GB/s    │  (NVIDIA/AMD)  │             │
│  │  AMD)    │                │                │             │
│  └────┬─────┘                └───────┬────────┘             │
│       │ DDR5                         │ GDDR6X               │
│       │ ~90 GB/s                     │ ~1 TB/s              │
│  ┌────┴─────┐                ┌───────┴───────┐              │
│  │ System   │                │   Video RAM   │              │
│  │ RAM      │                │   (VRAM)      │              │
│  │ 32-128GB │                │   8-24 GB     │              │
│  └──────────┘                └───────────────┘              │
│                                                             │
│  Two separate memory pools. Data must be COPIED between them│
└─────────────────────────────────────────────────────────────┘

┌─────────────────────────────────────────────────────────────┐
│                 APPLE SILICON ARCHITECTURE                  │
│                                                             │
│  ┌──────────────────────────────────────────┐               │
│  │            SINGLE SoC DIE                │               │
│  │                                          │               │
│  │  ┌─────┐ ┌─────┐ ┌──────┐ ┌──────────┐   │               │
│  │  │ CPU │ │ GPU │ │Neural│ │  Media   │   │               │
│  │  │     │ │     │ │Engine│ │  Engine  │   │               │
│  │  └──┬──┘ └──┬──┘ └──┬───┘ └────┬─────┘   │               │
│  │     └───────┴───────┴──────────┘         │               │
│  │              │ Fabric                    │               │
│  │     ┌────────┴────────┐                  │               │
│  │     │ System Level    │                  │               │
│  │     │ Cache (SLC)     │                  │               │
│  │     └────────┬────────┘                  │               │
│  └──────────────┼───────────────────────────┘               │
│          ┌──────┴──────┐                                    │
│          │  Unified    │                                    │
│          │  Memory     │                                    │
│          │  16-192 GB  │                                    │
│          └─────────────┘                                    │
│                                                             │
│  ONE memory pool. CPU and GPU access it with ZERO copies.   │
└─────────────────────────────────────────────────────────────┘

The animation below shows this difference. Watch how data flows in each architecture — pay attention to the PCIe bottleneck in the traditional model vs the direct access in UMA.

Anatomy of the Die

Let’s map out every major component on an M4 Pro die:¹

┌──────────────────────────────────────────────────────────────────┐
│                        M4 Pro SoC Die                            │
│                                                                  │
│  ┌─────────────────────────┐  ┌──────────────────────────────┐   │
│  │   CPU Cluster           │  │      GPU (20 cores)          │   │
│  │                         │  │                              │   │
│  │  ┌────┐┌────┐┌────┐     │  │  ┌───┐┌───┐┌───┐┌───┐┌───┐   │   │
│  │  │ P0 ││ P1 ││ P2 │     │  │  │ 0 ││ 1 ││ 2 ││ 3 ││ 4 │   │   │
│  │  └────┘└────┘└────┘     │  │  └───┘└───┘└───┘└───┘└───┘   │   │
│  │  ┌────┐┌────┐┌────┐     │  │  ┌───┐┌───┐┌───┐┌───┐┌───┐   │   │
│  │  │ P3 ││ P4 ││ P5 │     │  │  │ 5 ││ 6 ││ 7 ││ 8 ││ 9 │   │   │
│  │  └────┘└────┘└────┘     │  │  └───┘└───┘└───┘└───┘└───┘   │   │
│  │  ┌────┐┌────┐┌────┐     │  │  ┌───┐┌───┐┌───┐┌───┐┌───┐   │   │
│  │  │ P6 ││ P7 ││ P8 │     │  │  │10 ││11 ││12 ││13 ││14 │   │   │
│  │  └────┘└────┘└────┘     │  │  └───┘└───┘└───┘└───┘└───┘   │   │
│  │  ┌────┐                 │  │  ┌───┐┌───┐┌───┐┌───┐┌───┐   │   │
│  │  │ P9 │ P=Performance   │  │  │15 ││16 ││17 ││18 ││19 │   │   │
│  │  └────┘                 │  │  └───┘└───┘└───┘└───┘└───┘   │   │
│  │  ┌────┐┌────┐┌────┐     │  │         20 GPU Cores         │   │
│  │  │ E0 ││ E1 ││ E2 │     │  │                              │   │
│  │  └────┘└────┘└────┘     │  │                              │   │
│  │  ┌────┐ E=Efficiency    │  │                              │   │
│  │  │ E3 │                 │  │                              │   │
│  │  └────┘                 │  │                              │   │
│  └─────────────────────────┘  └──────────────────────────────┘   │
│                                                                  │
│  ┌──────────────┐  ┌──────────┐  ┌────────────┐  ┌──────────┐    │
│  │ Neural Engine│  │  Media   │  │  Display   │  │  Secure  │    │
│  │  16 Cores    │  │  Engine  │  │  Engine    │  │ Enclave  │    │
│  │  38 TOPS     │  │ ProRes   │  │            │  │          │    │
│  └──────────────┘  └──────────┘  └────────────┘  └──────────┘    │
│                                                                  │
│  ┌──────────────────────────────────────────────────────────┐    │
│  │                    Fabric / Interconnect                 │    │
│  └──────────────────────────────────────────────────────────┘    │
│  ┌──────────────────────────────────────────────────────────┐    │
│  │              System Level Cache (SLC) — 36 MB            │    │
│  └──────────────────────────────────────────────────────────┘    │
│  ┌──────────┐  ┌───────────┐  ┌──────────┐  ┌──────────────┐     │
│  │ Memory   │  │Thunderbolt│  │  PCIe    │  │    USB       │     │
│  │Controller│  │Controller │  │Controller│  │  Controller  │     │
│  └──────────┘  └───────────┘  └──────────┘  └──────────────┘     │
└──────────────────────────────────────────────────────────────────┘

CPU Clusters: Performance and Efficiency

Apple Silicon uses ARM’s big.LITTLE concept with two core types:

Performance Cores (P-cores):

• Wide, out-of-order execution with deep pipelines
• High clock speeds (up to ~4.5 GHz on M4)
• Large L1 caches (192 KB instruction, 128 KB data per core)
• Large shared L2 cache (16-32 MB per cluster)
• Used for compute-intensive tasks

Efficiency Cores (E-cores):

• Narrower, simpler pipeline
• Much lower power consumption
• Lower clock speeds (~2.8 GHz)
• Handle background tasks, I/O-bound work

For akunu, the CPU mainly handles model loading, tokenization, dispatch table building, grammar constraint checking, and control flow between GPU dispatches. The GPU does the neural network computation.

The GPU

This is where akunu spends most of its time. Each GPU core contains ALUs, a texture sampling unit, and tile memory. Core counts vary: M4 has 10, M4 Pro has 20, M4 Max has 40, M4 Ultra has 80. We’ll cover GPU architecture in extreme detail in Chapter 3.

The Neural Engine

Apple’s dedicated matrix multiplication accelerator: 16 cores on M4 Pro, 38 TOPS at INT8. However, akunu does NOT use the Neural Engine. It’s only accessible through CoreML, which doesn’t give the fine-grained control needed for optimized inference. Metal compute shaders give us full control over memory layout, kernel design, and dispatch — which is why akunu achieves 1.83x faster decode than llama.cpp.

System Level Cache (SLC)

The SLC is a large, shared last-level cache between all die components and DRAM:

         CPU L2     GPU L2     NE Cache
            \        |        /
             ▼       ▼       ▼
         ┌─────────────────────────┐
         │   System Level Cache    │
         │   M4:      16 MB        │
         │   M4 Pro:  36 MB        │
         │   M4 Max:  48 MB        │
         │   M4 Ultra: 96 MB       │
         └────────────┬────────────┘
                      ▼
              ┌───────────────┐
              │  LPDDR5 DRAM  │
              └───────────────┘

The SLC is critical for inference: weight tiles that fit in SLC get accessed much faster, and intermediate activations often fit entirely.² Akunu’s ChipConfig::slc_size tunes behavior based on available SLC.

Memory Controller

For decode, token generation time is approximately model_size_bytes / memory_bandwidth. For a 7B Q4_0 model (~3.8 GB): M4 Pro gives ~71 tok/s theoretical, M4 Max ~143 tok/s. Akunu gets close to these limits.

The Chip Hierarchy: Binning and Variants

Apple creates chip families from the same base design:

          BASE DIE ──────► M4 (binned, fewer cores)
              │
              ├──────────► M4 Pro (mid-range)
              │
              └──────────► M4 Max (full config)
                               │
                               │ UltraFusion (2 dies)
                               ▼
                           M4 Ultra (2x Max)

The Ultra variant connects two Max dies via UltraFusion (~2.5 TB/s die-to-die bandwidth). Software sees it as one unified chip — no special programming needed.

What This Means for Akunu

1. Zero-copy weight loading: GGUF files are memory-mapped; the GPU reads directly from mmap’d regions
2. CPU-side operations on GPU buffers: Whisper’s cross-attention K/V rearrangement happens on CPU, directly on GPU-accessible memory
3. Minimal CPU-GPU sync: Precompiled dispatch table minimizes CPU’s role in the hot path
4. SLC-aware tuning: ChipConfig adjusts tile and batch sizes based on SLC
5. Bandwidth as bottleneck: During decode, akunu is memory-bandwidth bound — quantization (Q4_0) has dramatic impact

In the next chapter, we’ll zoom into the GPU architecture specifically.

1. Apple. “Apple M4 Pro chip.” apple.com, 2024. Die specifications including GPU core count, memory bandwidth, and SLC size. See https://www.apple.com/newsroom/2024/10/apple-introduces-m4-pro-and-m4-max/. ↩

2. Frumusanu, A. “Apple’s M4 Family: A Deep Dive.” AnandTech / Chips and Cheese, 2024. Independent analysis of Apple Silicon die layout, cache hierarchy, and fabric bandwidth. See https://chipsandcheese.com/. ↩

The GPU: Cores, SIMD Groups, and Threadgroups

This is the most important chapter for understanding akunu’s Metal kernels. Every kernel in akunu — every GEMV, every FlashAttention variant, every normalization — is designed around the execution model we’re about to explore. If you’ve worked with CUDA, many concepts will feel familiar, but the terminology and some architectural details differ. If you haven’t, don’t worry — we’ll build up from first principles.

The Big Picture: Why GPUs?

A CPU is designed to execute a single thread of complex instructions as fast as possible. It has deep pipelines, branch prediction, out-of-order execution, and large caches. A single CPU core might run at 4+ GHz but can only do one or two multiply-accumulate operations per clock.

A GPU takes the opposite approach. It trades single-thread performance for massive parallelism. Each individual “thread” is simpler and slower, but there are thousands of them running simultaneously.

CPU (M4 Pro Performance Core):
┌──────────────────────────────────┐
│  Complex OoO pipeline            │
│  Branch predictor                │
│  192 KB L1I + 128 KB L1D         │
│  ~4.5 GHz                        │
│  1-2 FMA per clock               │
│  ≈ 9 GFLOPS FP32 per core        │
└──────────────────────────────────┘
  × 10 P-cores = ~90 GFLOPS

GPU (M4 Pro, 20 cores):
┌──────────┐┌──────────┐┌──────────┐  ...  ┌──────────┐
│  Core 0  ││  Core 1  ││  Core 2  │       │  Core 19 │
│  128 ALUs││  128 ALUs││  128 ALUs│       │  128 ALUs│
│ ~1.5 GHz ││ ~1.5 GHz ││ ~1.5 GHz │       │ ~1.5 GHz │
└──────────┘└──────────┘└──────────┘       └──────────┘
  20 cores × 128 ALUs × 2 FMA × 1.5 GHz ≈ 7,680 GFLOPS FP32
  FP16: ~15,360 GFLOPS (2x throughput)

For matrix multiplication — the core operation in neural networks — you need to do the same thing (multiply and add) billions of times with different data. GPUs are purpose-built for exactly this.

GPU Core Architecture

Each Apple GPU core is a self-contained processing unit. Let’s look inside one:

┌─────────────────────────────────────────────────────────┐
│                    GPU CORE                             │
│                                                         │
│  ┌────────────────────────────────────────────────────┐ │
│  │              Execution Units                       │ │
│  │                                                    │ │
│  │  ┌────────┐ ┌────────┐ ┌────────┐ ┌────────┐       │ │
│  │  │ ALU×32 │ │ ALU×32 │ │ ALU×32 │ │ ALU×32 │       │ │
│  │  │ (SG 0) │ │ (SG 1) │ │ (SG 2) │ │ (SG 3) │       │ │
│  │  └────────┘ └────────┘ └────────┘ └────────┘       │ │
│  │    SIMD       SIMD       SIMD       SIMD           │ │
│  │   Group 0    Group 1    Group 2    Group 3         │ │
│  │                                                    │ │
│  │  Up to 32 SIMD groups can be resident              │ │
│  │  (1024 threads max per core)                       │ │
│  └────────────────────────────────────────────────────┘ │
│                                                         │
│  ┌─────────────────┐  ┌──────────────────────────────┐  │
│  │ Register File    │  │  Threadgroup Memory (32 KB) │  │
│  │ (per thread)     │  │  (shared within threadgroup)│  │
│  └─────────────────┘  └──────────────────────────────┘  │
│                                                         │
│  ┌──────────────────────────────────────────────────┐   │
│  │              L1 Cache / Tile Memory              │   │
│  └──────────────────────────────────────────────────┘   │
│                          │                              │
│                     L2 Cache (shared across cores)      │
└─────────────────────────────────────────────────────────┘

Key facts:

• Each core has 128 ALUs organized into groups of 32¹
• Each group of 32 ALUs forms a SIMD group (Apple’s term for what NVIDIA calls a “warp”)
• A core can have up to 1024 threads resident (32 SIMD groups × 32 threads)
• 32 KB of threadgroup memory (what NVIDIA calls “shared memory”)

The SIMD Group: 32 Threads in Lockstep

The SIMD group is the fundamental unit of execution on the GPU. It’s a group of 32 threads that execute the exact same instruction at the exact same time, but on different data. This is called SIMT (Single Instruction, Multiple Threads).

SIMD Group (32 threads):
┌──────────────────────────────────────────────────────┐
│  Thread 0   Thread 1   Thread 2   ...   Thread 31    │
│                                                      │
│  All execute: result = a[tid] * b[tid] + c[tid]      │
│                                                      │
│  t0: a[0]*b[0]+c[0]                                  │
│  t1: a[1]*b[1]+c[1]                                  │
│  t2: a[2]*b[2]+c[2]                                  │
│  ...                                                 │
│  t31: a[31]*b[31]+c[31]                              │
│                                                      │
│  ALL happen in ONE clock cycle (on 32 ALUs)          │
└──────────────────────────────────────────────────────┘

This has profound implications:

• No divergence penalty if all threads take the same branch. If thread 0 takes the if path and thread 1 takes the else path, the SIMD group must execute BOTH paths (with inactive threads masked out). This wastes cycles.
• Memory coalescing: When all 32 threads access consecutive memory addresses, the hardware can combine them into a single wide memory transaction. Random access patterns are much slower.
• SIMD intrinsics: Threads within a SIMD group can communicate directly via simd_sum, simd_max, simd_shuffle, etc. These are essentially free — no memory access needed.

SIMD Group Communication

This is one of the most powerful features for reduction operations (like computing a dot product). Threads in a SIMD group can share data without going through memory:

simd_sum example (sum 32 values across a SIMD group):

Thread:   0    1    2    3   ...  31
Value:   1.0  2.0  3.0  4.0 ... 32.0

simd_sum → Every thread gets: 528.0

No shared memory needed! Hardware does this in ~5 cycles (log2(32) shuffle steps).

simd_shuffle_down (shift values):

Thread:   0    1    2    3    4    5   ...
Before:  a0   a1   a2   a3   a4   a5  ...
After (offset=2):  a2   a3   a4   a5   a6   a7  ...

simd_shuffle_xor (butterfly pattern):

Thread:  0    1    2    3    4    5    6    7
XOR(1): a1   a0   a3   a2   a5   a4   a7   a6
XOR(2): a2   a3   a0   a1   a6   a7   a4   a5
XOR(4): a4   a5   a6   a7   a0   a1   a2   a3

Akunu uses these extensively. For example, in GEMV kernels, each thread in a SIMD group computes a partial dot product, then simd_sum combines them. In FlashAttention decode, simd_sum computes the Q·K dot product across threads that each hold different head dimensions.

The Threadgroup

A threadgroup is a collection of threads that:

1. Execute on the same GPU core
2. Can share threadgroup memory (fast on-chip SRAM)
3. Can synchronize with threadgroup_barrier()

Threadgroup (128 threads = 4 SIMD groups):
┌─────────────────────────────────────────────────────────┐
│                                                         │
│  ┌──────────────┐  ┌──────────────┐                     │
│  │  SIMD Group 0│  │  SIMD Group 1│                     │
│  │  threads 0-31│  │ threads 32-63│                     │
│  └──────────────┘  └──────────────┘                     │
│  ┌──────────────┐  ┌──────────────┐                     │
│  │  SIMD Group 2│  │  SIMD Group 3│                     │
│  │ threads 64-95│  │threads 96-127│                     │
│  └──────────────┘  └──────────────┘                     │
│                                                         │
│  ┌──────────────────────────────────────────────────┐   │
│  │         Threadgroup Memory (up to 32 KB)         │   │
│  │    Accessible by ALL threads in this threadgroup │   │
│  │    ~1 cycle latency (vs ~100+ cycles for DRAM)   │   │
│  └──────────────────────────────────────────────────┘   │
│                                                         │
│  threadgroup_barrier(mem_flags::mem_threadgroup)        │
│  ↑ Ensures all threads have finished their writes       │
│    before any thread reads                              │
└─────────────────────────────────────────────────────────┘

Common threadgroup sizes in akunu:²

• 32 threads (1 SG): FlashAttention decode fast — no barriers needed!
• 128 threads (4 SGs): GEMV, GEMM, RMSNorm, RoPE
• 256 threads (8 SGs): Large GEMV variants
• 1024 threads (32 SGs): FlashAttention decode parallel, Gumbel-max sampling

Why Threadgroup Memory Matters

Threadgroup memory is like a programmer-managed L1 cache. It’s fast (~1 cycle) but small (32 KB). The classic use is tiled matrix multiply:

Without threadgroup memory:
  Thread 0 reads A[0,0] from DRAM (100+ cycles)
  Thread 1 reads A[0,0] from DRAM (100+ cycles) — SAME DATA!
  Thread 2 reads A[0,0] from DRAM — again!
  ... massive bandwidth waste

With threadgroup memory:
  1. Thread 0 loads A[0,0..31] into shared[0..31]
  2. threadgroup_barrier() — wait for all loads
  3. ALL threads read from shared[] (1 cycle each)
  4. One DRAM fetch serves 32 threads

Akunu’s GEMM kernel (simd_gemm_f16.metal, covered in Chapter 35) allocates 6144 bytes of threadgroup memory: sa[4096] for weight tiles and sb[2048] for activation tiles. This lets 128 threads share weight data loaded once from DRAM.

The Execution Hierarchy

Click on any level to zoom in and see what’s inside. This is the most important mental model for GPU programming — every kernel you write maps onto this hierarchy.

Thread identification variables (used in every kernel):

kernel void my_kernel(
    uint tid       [[thread_position_in_grid]],      // Global thread ID
    uint tgid      [[threadgroup_position_in_grid]],  // Which threadgroup
    uint tiisg     [[thread_index_in_simdgroup]],     // Lane within SG (0-31)
    uint sgitg     [[simdgroup_index_in_threadgroup]], // Which SG in TG (0-N)
    uint tiitg     [[thread_position_in_threadgroup]]  // Thread within TG
) { ... }

CUDA vs Metal Terminology

If you’re coming from CUDA, here’s the translation table:

Key differences:

1. No warp divergence penalty tracking in Metal³ — the hardware handles it, but you should still avoid it
2. SIMD group matrix operations use 8×8 tiles (vs NVIDIA’s 16×16 Tensor Core tiles)
3. Threadgroup memory is explicitly sized per dispatch, not declared with __shared__
4. No dynamic parallelism — kernels can’t launch other kernels

Occupancy and Register Pressure

Each GPU core has a fixed amount of resources. The more resources a threadgroup uses, the fewer threadgroups can be resident simultaneously:

GPU Core Resources:
┌────────────────────────────────┐
│  Register File:  ~32K regs     │
│  Threadgroup Mem: 32 KB        │
│  Max Threads:     1024         │
│  Max SGs:         32           │
└────────────────────────────────┘

Example: Kernel uses 32 registers per thread, 4096 bytes TG mem

  Threads per TG: 128
  Registers per TG: 128 × 32 = 4096
  TG memory per TG: 4096 bytes

  Max concurrent TGs: min(
    32K / 4096 = 8 (register limited),
    32K / 4096 = 8 (TG mem limited),
    1024 / 128 = 8 (thread limited)
  ) = 8 threadgroups, 1024 threads

  Occupancy: 1024 / 1024 = 100%

High occupancy means more threads to hide memory latency. When one SIMD group is waiting for a memory fetch (100+ cycles), the core can switch to another SIMD group that has work ready. This is how GPUs tolerate high memory latency — they never run out of work.

Dispatch Models

Metal offers two dispatch models:

dispatchThreadgroups(gridSize, threadgroupSize):

• You specify how many threadgroups to launch and their size
• Total threads = gridSize × threadgroupSize
• May launch more threads than needed (you handle bounds checking)

dispatchThreads(totalThreads, threadgroupSize):

• You specify the total number of threads needed
• Metal handles partial threadgroups at the edges
• Cleaner for simple 1:1 thread-to-data mappings

Akunu’s DispatchCmd has a use_dispatch_threads flag that selects between these. Most kernels use dispatchThreadgroups for precise control.

The Tile-Based Deferred Renderer (TBDR)

Apple’s GPU was originally designed for mobile graphics, which uses a Tile-Based Deferred Rendering architecture. In graphics, the screen is divided into tiles and each tile is rendered entirely in fast tile memory before writing to DRAM.

For compute shaders (which akunu uses exclusively), TBDR doesn’t apply directly. Compute shaders bypass the tiling hardware and operate like a traditional GPU compute model. However, the tile memory architecture means:

• The GPU has fast on-chip storage (threadgroup memory)
• Memory access patterns that fit in cache lines are rewarded
• The GPU is efficient at processing data in blocks/tiles

This is why akunu’s kernels are organized around tiles: 32-element K-blocks in GEMV, 32×64 output tiles in GEMM, 32-position KV tiles in FlashAttention.

Summary

The Apple GPU execution model is:

1. Grid dispatches many threadgroups
2. Each threadgroup runs on one GPU core
3. Each threadgroup contains multiple SIMD groups of 32 threads
4. Threads within a SIMD group execute in lockstep and communicate via SIMD intrinsics
5. Threads within a threadgroup share threadgroup memory and synchronize via barriers
6. High occupancy hides memory latency

Akunu’s kernels are designed around this hierarchy: SIMD-level reductions for dot products, threadgroup-level cooperation for tiled matrix multiply, and grid-level parallelism across output rows and attention heads.

Next, we’ll look at the unified memory architecture that makes Apple Silicon uniquely suited for inference workloads.

1. Grinberg, D. “Reverse-engineering Apple GPU cores.” Asahi Linux project, 2022. The most detailed public analysis of Apple GPU core internals, SIMD group behavior, and threadgroup memory layout. See https://dougallj.github.io/applegpu/. ↩

2. Apple. “Metal Best Practices Guide.” developer.apple.com. Official guidance on threadgroup sizing, occupancy, and memory access patterns for Metal compute shaders. See https://developer.apple.com/library/archive/documentation/3DDrawing/Conceptual/MTLBestPracticesGuide/index.html. ↩

3. Apple. “Metal Feature Set Tables.” developer.apple.com. Defines GPU family capabilities including simdgroup_matrix support (Family 7+), max threadgroup size, and threadgroup memory limits. See https://developer.apple.com/metal/Metal-Feature-Set-Tables.pdf. ↩

Unified Memory Architecture

If there’s one hardware feature that makes Apple Silicon uniquely suited for running large language models locally, it’s the Unified Memory Architecture (UMA). It’s the reason a MacBook with 64 GB of RAM can run a 70B model that would require a $2,000+ NVIDIA GPU on a PC. Let’s understand exactly how it works and why it matters for inference.

The Traditional Memory Problem

On a discrete GPU system, there are two completely separate memory pools:

┌─────────────────────────────────────────────────────────┐
│             TRADITIONAL DISCRETE GPU SYSTEM             │
│                                                         │
│   CPU Side                        GPU Side              │
│   ┌──────────┐                    ┌──────────────┐      │
│   │  System   │    PCIe 4.0 x16   │    VRAM      │      │
│   │  RAM      │◄──────────────────►│    (GDDR6X)  │      │
│   │  64 GB    │    ~32 GB/s        │    24 GB     │      │
│   │  DDR5     │    bidirectional   │              │      │
│   │           │                    │  Bandwidth:  │      │
│   │  BW: ~90  │                    │  ~1 TB/s     │      │
│   │   GB/s    │                    │              │      │
│   └──────────┘                    └──────────────┘      │
│                                                          │
│   To run inference:                                      │
│   1. Load model into system RAM (disk → RAM)             │
│   2. Copy weights to VRAM (RAM → PCIe → VRAM)          │
│   3. GPU computes on VRAM                                │
│   4. Copy results back (VRAM → PCIe → RAM)              │
│                                                          │
│   PROBLEM: Model must fit in VRAM (24 GB max typical)    │
│   PROBLEM: PCIe transfer is slow (~32 GB/s)              │
│   PROBLEM: Double the memory usage (model in both pools) │
└─────────────────────────────────────────────────────────┘

A 70B model in Q4_0 quantization is about 40 GB. It literally doesn’t fit in a 24 GB RTX 4090’s VRAM. You need model parallelism across multiple GPUs, or offloading to CPU (which is painfully slow over PCIe).

Apple’s Unified Memory

On Apple Silicon, there’s one memory pool, shared by everything:

┌─────────────────────────────────────────────────────────┐
│             APPLE SILICON UNIFIED MEMORY                  │
│                                                          │
│  ┌────────────────────────────────────────────────────┐  │
│  │                   SoC Die                           │  │
│  │                                                     │  │
│  │   CPU ──┐                                          │  │
│  │         ├──► Fabric ──► SLC ──► Memory Controller  │  │
│  │   GPU ──┤                                          │  │
│  │         ├──►                                       │  │
│  │   NE  ──┘                                          │  │
│  └────────────────────────────────────────────────────┘  │
│                          │                               │
│                   ┌──────┴──────┐                        │
│                   │   LPDDR5    │                        │
│                   │   Unified   │                        │
│                   │   Memory    │                        │
│                   │   Pool      │                        │
│                   │             │                        │
│                   │  16-192 GB  │                        │
│                   └─────────────┘                        │
│                                                          │
│   To run inference:                                      │
│   1. Load model into memory (disk → unified memory)      │
│   2. GPU computes on it directly. Done.                  │
│                                                          │
│   NO COPY. NO VRAM LIMIT. NO PCIe BOTTLENECK.           │
└─────────────────────────────────────────────────────────┘

The implications:

1. No VRAM limit: A Mac Studio with 192 GB can run models that would need 8× NVIDIA A100s
2. No copy overhead: Weights loaded once are accessible by both CPU and GPU
3. Zero-copy operations: CPU can rearrange GPU data directly (akunu uses this for Whisper)

The Memory Hierarchy

Even though memory is “unified,” there’s still a hierarchy of caches for performance. The animation below shows a memory request falling through the layers. Click a layer to see what happens when data is found there (cache hit) vs when it must go deeper (cache miss).

For inference, the critical numbers are:

• SLC size: determines how much of the model weights are “hot” in cache
• DRAM bandwidth: determines maximum decode speed for memory-bound operations

Why Memory Bandwidth Is Everything for Decode

During token generation (decode), the model processes one token at a time. For each token, it needs to:

1. Read the embedding for the input token: tiny
2. For each layer (e.g., 32 layers in a 7B model):
   • Read Q/K/V projection weights: (dim + 2 × kv_dim) × dim × bytes_per_weight (KV is smaller with GQA)
   • Read KV cache: 2 × seq_len × kv_dim × 2 bytes (FP16)
   • Read output projection weights: dim × dim × bytes_per_weight
   • Read FFN gate+up weights: 2 × dim × ffn_dim × bytes_per_weight
   • Read FFN down weights: ffn_dim × dim × bytes_per_weight

For a 7B Q4_0 model, that’s roughly 3.8 GB of weights read per token. The computation (multiply-accumulate) is fast. The bottleneck is reading the weights from memory.

Token generation speed = memory bandwidth / model size (bytes).¹

LLaMA 7B Q4_0 (~3.8 GB) theoretical decode speed:

The RTX 4090 has higher bandwidth, but only 24 GB of VRAM. For models that fit, it’s faster. For large models (70B+), it can’t run them at all without multi-GPU setups.

The Roofline Model

The roofline model helps visualize whether an operation is compute-bound or memory-bound:

Performance
(FLOPS)
    │
    │            ┌─────────────────────── Compute Ceiling
    │           /│                        (peak TFLOPS)
    │          / │
    │         /  │
    │        /   │
    │       /    │
    │      /     │
    │     /      │
    │    /       │
    │   / Memory │ Compute
    │  /  Bound  │  Bound
    │ /          │
    │/           │
    └────────────┼──────────────────────
                 Ridge Point
           Operational Intensity
           (FLOPS / byte)

GEMV (decode, M=1, see Chapter 34): Each weight is read once and used for 1 multiply-add. Operational intensity = 2 FLOPS / (0.5 bytes for Q4_0) = 4 FLOPS/byte. This is well below the ridge point → memory-bound.

GEMM (prefill, M=seq_len): Each weight is read once but used for M multiply-adds. Operational intensity = 2M FLOPS / (0.5 bytes). For M=512: 2048 FLOPS/byte → compute-bound.

This is why prefill is fast (compute-bound, GPU cores are busy) and decode is slower (memory-bound, waiting for data).

How Akunu Exploits UMA

1. Memory-Mapped Weight Loading

GGUF files are opened with mmap(). The OS maps the file directly into the process’s virtual address space. When a page is first accessed, the kernel loads it from disk. The GPU can then access these pages directly — no explicit copy needed.

Traditional:  Disk → fread() → RAM buffer → cudaMemcpy() → VRAM
Akunu UMA:    Disk → mmap() → Unified Memory ← GPU reads directly

2. CPU-Side Buffer Operations

For Whisper’s cross-attention, akunu precomputes K/V projections on the GPU, then rearranges the memory layout on the CPU:

// CPU directly reads/writes GPU buffer contents (UMA makes this efficient)
for (int h = 0; h < n_heads; h++)
    for (int p = 0; p < enc_seq; p++)
        memcpy(&dst[(h*enc_seq + p)*head_dim],
               &src[p*dec_dim + h*head_dim],
               head_dim * sizeof(__fp16));

On a discrete GPU, this would require: GPU→VRAM→PCIe→RAM→CPU rearrange→RAM→PCIe→VRAM→GPU. On Apple Silicon: CPU reads and writes directly to the same memory the GPU uses. Hundreds of microseconds vs potentially milliseconds.

3. Minimal Synchronization

Because CPU and GPU share memory, akunu can submit GPU work and immediately prepare the next operation. There’s no need to wait for a PCIe transfer to complete — just a lightweight event/fence for GPU command completion.

4. KV Cache Persistence

The KV cache lives in unified memory and persists across requests. In akunu’s server mode, if a new request shares a prefix with the previous one, the cached KV entries are already in memory — no need to recompute or retransfer.

Memory Bandwidth in Practice

Real bandwidth utilization is always less than theoretical peak. Several factors reduce effective bandwidth:

Theoretical Peak:     273 GB/s (M4 Pro)
                        │
Bank conflicts:         │ -5%
Cache misses:           │ -10%
Non-coalesced access:   │ -15%
TLB misses:             │ -5%
                        ▼
Practical Peak:       ~175-195 GB/s (65-70% utilization)

Akunu achieves high bandwidth utilization through:

• Vectorized loads: half4 reads (8 bytes at once) instead of individual half reads
• Coalesced access: Adjacent threads read adjacent memory addresses
• Block-strided K-loops: Process K-dimension in blocks of 512 to maximize cache line reuse
• Weight layout optimization: Quantized blocks are laid out for sequential access

Comparison: UMA vs Discrete vs Cloud

Apple Silicon’s sweet spot is models that don’t quite fit in consumer GPU VRAM (13B-70B). The ability to run these models locally, at reasonable speed, with no cloud dependency, is akunu’s core value proposition.

Summary

Unified Memory Architecture means:

• One memory pool shared by CPU and GPU
• No data copying between host and device memory
• Total system RAM available for model weights
• Memory bandwidth is the primary bottleneck for decode
• Quantization directly translates to speed (less data to read)

The key equation for decode performance: tokens/sec ≈ memory_bandwidth / model_size_bytes

This is why akunu focuses heavily on quantization support (Q4_0 is 4× smaller than FP16, so 4× faster decode) and why its dispatch table eliminates every unnecessary overhead — when you’re memory-bound, every wasted cycle is a wasted byte of bandwidth.

1. Pope, R., et al. “Efficiently Scaling Transformer Inference.” MLSys 2023. This paper provides a thorough analysis of the memory-bound vs compute-bound regimes of transformer inference and derives the bandwidth-limited decode throughput formula. See https://arxiv.org/abs/2211.05102. ↩

The Apple GPU Family: M1 through M4

Now that we understand the SoC architecture, unified memory, and GPU execution model, let’s survey the actual hardware across Apple Silicon generations. Each generation brought meaningful improvements for ML inference, and akunu’s ChipConfig system tunes its behavior for each.

Generation Overview1

M1 Family (2020-2021)

The M1 was the first Apple Silicon chip for Mac. It proved the concept works.

M1 Family Specifications:
┌──────────┬──────────┬──────────┬──────────┬──────────┐
│          │   M1     │ M1 Pro   │ M1 Max   │ M1 Ultra │
├──────────┼──────────┼──────────┼──────────┼──────────┤
│ GPU Cores│   7-8    │  14-16   │  24-32   │  48-64   │
│ CPU Cores│   8      │  8-10    │  10      │  20      │
│ Max RAM  │  16 GB   │  32 GB   │  64 GB   │  128 GB  │
│ Mem BW   │ 68 GB/s  │ 200 GB/s │ 400 GB/s │ 800 GB/s │
│ SLC      │  16 MB   │  24 MB   │  48 MB   │  96 MB   │
│ GPU Fam  │ Apple 7  │ Apple 7  │ Apple 7  │ Apple 7  │
│ FP32 TF  │  2.6     │  5.2     │  10.4    │  20.8    │
│ FP16 TF  │  5.2     │  10.4    │  20.8    │  41.6    │
│ NE TOPS  │  11      │  11      │  11      │  22      │
└──────────┴──────────┴──────────┴──────────┴──────────┘

For inference:

• M1 base: Usable for small models (1-3B Q4). 68 GB/s bandwidth limits decode speed
• M1 Max 64GB: First chip that could comfortably run 13B Q4 models
• M1 Ultra 128GB: Could run 70B Q4 models, albeit slowly

Key capability: Apple GPU Family 7 introduced simdgroup_matrix operations (8×8 matrix tiles), which akunu uses for GEMM kernels during prefill.

M2 Family (2022-2023)

An evolutionary improvement: more cores, more bandwidth, same architecture.

M2 Family Specifications:
┌──────────┬──────────┬──────────┬──────────┬──────────┐
│          │   M2     │ M2 Pro   │ M2 Max   │ M2 Ultra │
├──────────┼──────────┼──────────┼──────────┼──────────┤
│ GPU Cores│   8-10   │  16-19   │  30-38   │  60-76   │
│ CPU Cores│   8      │  10-12   │  12      │  24      │
│ Max RAM  │  24 GB   │  32 GB   │  96 GB   │  192 GB  │
│ Mem BW   │ 100 GB/s │ 200 GB/s │ 400 GB/s │ 800 GB/s │
│ SLC      │  16 MB   │  24 MB   │  48 MB   │  96 MB   │
│ GPU Fam  │ Apple 8  │ Apple 8  │ Apple 8  │ Apple 8  │
│ FP32 TF  │  3.6     │  6.8     │  13.6    │  27.2    │
│ FP16 TF  │  7.2     │  13.6    │  27.2    │  54.4    │
│ NE TOPS  │  15.8    │  15.8    │  15.8    │  31.6    │
└──────────┴──────────┴──────────┴──────────┴──────────┘

For inference:

• M2 Max 96GB: Sweet spot for 70B Q4 models
• M2 Ultra 192GB: Could run quantized 100B+ models
• ~30% faster than M1 at equivalent tier

M3 Family (2023-2024)

The jump to 3nm brought significant GPU architectural changes.

M3 Family Specifications:
┌──────────┬──────────┬──────────┬──────────┬──────────┐
│          │   M3     │ M3 Pro   │ M3 Max   │ M3 Ultra │
├──────────┼──────────┼──────────┼──────────┼──────────┤
│ GPU Cores│   8-10   │  14-18   │  30-40   │  60-80   │
│ CPU Cores│   8      │  11-12   │  14-16   │  28-32   │
│ Max RAM  │  24 GB   │  36 GB   │  128 GB  │  192 GB  │
│ Mem BW   │ 100 GB/s │ 150 GB/s │ 400 GB/s │ 800 GB/s │
│ SLC      │  16 MB   │  24 MB   │  48 MB   │  96 MB   │
│ GPU Fam  │ Apple 9  │ Apple 9  │ Apple 9  │ Apple 9  │
│ FP32 TF  │  4.1     │  7.4     │  16.4    │  32.8    │
│ FP16 TF  │  8.2     │  14.8    │  32.8    │  65.6    │
│ NE TOPS  │  18      │  18      │  18      │  36      │
└──────────┴──────────┴──────────┴──────────┴──────────┘

M3’s GPU innovations:

• Dynamic Caching: GPU dynamically allocates register and threadgroup memory per kernel, instead of statically at dispatch time. This improves occupancy for kernels that don’t use much threadgroup memory.
• Ray Tracing Hardware: Not relevant for inference, but shows architectural maturation
• Mesh Shading: Graphics feature, not relevant here
• Apple GPU Family 9: same SIMD group matrix ops as Family 7, but with improved scheduling

M4 Family (2024-2025)

The latest generation, optimized for AI workloads.

M4 Family Specifications:
┌──────────┬──────────┬──────────┬──────────┬──────────┐
│          │   M4     │ M4 Pro   │ M4 Max   │ M4 Ultra │
├──────────┼──────────┼──────────┼──────────┼──────────┤
│ GPU Cores│   10     │  20      │  40      │  80      │
│ CPU Cores│ 4P+6E    │ 10P+4E   │ 12P+4E   │ 24P+8E  │
│ Max RAM  │  32 GB   │  48 GB   │  128 GB  │  192 GB  │
│ Mem BW   │ 120 GB/s │ 273 GB/s │ 546 GB/s │ 819 GB/s │
│ SLC      │  16 MB   │  36 MB   │  48 MB   │  96 MB   │
│ GPU Fam  │ Apple 9  │ Apple 9  │ Apple 9  │ Apple 9  │
│ FP32 TF  │  4.6     │  9.2     │  18.4    │  36.8    │
│ FP16 TF  │  9.2     │  18.4    │  36.8    │  73.6    │
│ NE TOPS  │  38      │  38      │  38      │  76      │
└──────────┴──────────┴──────────┴──────────┴──────────┘

M4 highlights for inference:

• Significantly higher memory bandwidth at each tier (M4 Pro: 273 vs M3 Pro: 150 GB/s)
• Larger SLC (M4 Pro: 36 MB vs M3 Pro: 24 MB)
• More GPU cores per tier
• Enhanced Neural Engine (38 TOPS)

Inference Performance by Chip

Here are theoretical decode speeds for common model sizes:

Tokens/sec (theoretical max, Q4_0 quantization):

Model Size:     3B (~1.7GB)  7B (~3.8GB)  13B (~7.3GB)  70B (~40GB)
─────────────────────────────────────────────────────────────────────
M4              70           31           16            —*
M4 Pro          160          71           37            6.8
M4 Max          321          143          74            13.6
M4 Ultra        481          215          112           20.4
─────────────────────────────────────────────────────────────────────
*Doesn't fit in 32 GB RAM

Actual performance will be 60-80% of theoretical due to cache misses, overhead, and non-ideal memory access patterns.² Akunu typically achieves 70-85% of theoretical bandwidth utilization.

Akunu’s ChipConfig

Akunu detects the hardware at startup and selects tuning parameters via the ChipConfig struct. Here’s how different chips get different configurations:

ChipConfig Parameters:
┌────────────────────┬─────────────────────────────────────┐
│ Parameter          │ Purpose                              │
├────────────────────┼─────────────────────────────────────┤
│ slc_size           │ SLC size in bytes. Affects tile     │
│                    │ sizes and prefetch strategies.       │
├────────────────────┼─────────────────────────────────────┤
│ gemv_k_threshold   │ K dimension where GEMV switches     │
│                    │ from small (4 SGs) to large (8 SGs) │
│                    │ variant. Higher bandwidth chips can  │
│                    │ benefit from more parallelism.       │
├────────────────────┼─────────────────────────────────────┤
│ prefill_chunk      │ Max tokens per prefill batch.       │
│                    │ Limited by threadgroup memory for    │
│                    │ attention. Larger chips handle more. │
├────────────────────┼─────────────────────────────────────┤
│ chain_decode_count │ Tokens per chained decode GPU       │
│                    │ submission. More cores → more tokens │
│                    │ can be chained without overhead.     │
└────────────────────┴─────────────────────────────────────┘

Example configurations:

M4 (10 GPU cores, 120 GB/s, 16 MB SLC):
  gemv_k_threshold = 2048
  prefill_chunk = 512
  chain_decode_count = 4

M4 Pro (20 GPU cores, 273 GB/s, 36 MB SLC):
  gemv_k_threshold = 2048
  prefill_chunk = 2048
  chain_decode_count = 6

M4 Max (40 GPU cores, 546 GB/s, 48 MB SLC):
  gemv_k_threshold = 4096
  prefill_chunk = 4096
  chain_decode_count = 8

The key insight: the same kernels run on all chips, but the dispatch table uses different configurations. The GEMV kernel for Q4_0 on an M4 uses 4 SIMD groups (128 threads), while on an M4 Max it might use 8 SIMD groups (256 threads) for K dimensions above the threshold.

Choosing the Right Hardware for Your Workload

┌─────────────────────────────────────────────────────────────┐
│                  MODEL SIZE vs CHIP GUIDE                     │
│                                                              │
│  Model Size    Recommended Minimum   Sweet Spot              │
│  ──────────    ────────────────────  ──────────              │
│  1-3B          M4 (16 GB)           M4 Pro (24 GB)          │
│  7B            M4 Pro (24 GB)       M4 Pro (36 GB)          │
│  13B           M4 Pro (36 GB)       M4 Max (64 GB)          │
│  34B           M4 Max (64 GB)       M4 Max (128 GB)         │
│  70B           M4 Max (128 GB)      M4 Ultra (192 GB)       │
│  100B+         M4 Ultra (192 GB)    —                       │
│                                                              │
│  Rule of thumb: you need ~1.2x the model size in Q4 as RAM  │
│  (weights + KV cache + scratch buffers + OS overhead)        │
└─────────────────────────────────────────────────────────────┘

Summary

Across four generations:

• M1 introduced UMA and SIMD group matrix ops — the foundation
• M2 increased bandwidth and core counts — evolutionary improvement
• M3 added dynamic caching and 3nm efficiency — architectural refinement
• M4 pushed bandwidth dramatically higher — the best for inference

Akunu’s ChipConfig abstracts these differences into tuning parameters, so the same codebase runs optimally on everything from an M1 MacBook Air to an M4 Ultra Mac Studio. The key variables are always the same: memory bandwidth, GPU core count, and SLC size.

In the next part, we’ll learn how to actually program these GPUs using the Metal framework.

1. Apple. “Apple M1 chip”, “Apple M2 chip”, “Apple M3 chip”, “Apple M4 chip.” apple.com, 2020-2024. Official specifications for each chip generation. See https://www.apple.com/newsroom/2024/05/apple-introduces-m4-chip/. ↩

2. Frumusanu, A. and Smith, R. “Apple M-series deep dives.” AnandTech, 2020-2023. Independent benchmarks and die analysis across generations. See https://chipsandcheese.com/. ↩

Introduction to Metal

Welcome to Part II of this book. In Part I, we explored the hardware that makes Apple Silicon so compelling for machine learning inference: the system-on-chip design, the GPU architecture with its SIMD groups and threadgroups, and the unified memory architecture that eliminates the CPU-to-GPU copy bottleneck. Now it is time to learn how to actually program that hardware.

The answer is Metal – Apple’s low-level GPU programming framework. Over the next seven chapters, we will go from zero to writing high-performance compute shaders that form the backbone of ML inference engines like akunu. By the end of Part II, you will understand every layer of the Metal compute stack, from the API calls in Swift/Objective-C all the way down to the individual threads executing on GPU cores.

Let us begin.

What Is Metal?

Metal is Apple’s unified graphics and compute API.¹ Introduced at WWDC 2014, it replaced the aging OpenGL ES on iOS and eventually OpenGL and OpenCL on macOS. Think of Metal as Apple’s answer to Vulkan or DirectX 12 – a modern, low-overhead GPU programming interface that gives you explicit control over the hardware.²

But Metal is more than just a graphics API. It has three major facets:

+-----------------------------------------------------------+
|                      Metal Framework                       |
+-----------------------------------------------------------+
|                                                           |
|  +----------------+  +----------------+  +-------------+  |
|  |  Metal API     |  |  Metal Shading |  |  Metal      |  |
|  |  (Swift/ObjC)  |  |  Language      |  |  Performance|  |
|  |                |  |  (C++14-based) |  |  Shaders    |  |
|  |  - Device      |  |                |  |  (MPS)      |  |
|  |  - Buffers     |  |  - Compute     |  |             |  |
|  |  - Queues      |  |    kernels     |  |  - MatMul   |  |
|  |  - Encoders    |  |  - Vertex/     |  |  - Conv     |  |
|  |  - Pipelines   |  |    Fragment    |  |  - Image    |  |
|  |  - Textures    |  |    shaders     |  |    ops      |  |
|  +----------------+  +----------------+  +-------------+  |
|                                                           |
+-----------------------------------------------------------+

1. The Metal API (Swift/Objective-C): The host-side interface. You use this to create GPU devices, allocate memory, build command buffers, and submit work to the GPU. This runs on the CPU.

2. The Metal Shading Language (MSL): The language you write GPU programs in. It is based on C++14 with Apple-specific extensions for GPU concepts like threadgroups, SIMD operations, and address spaces. Your compute kernels are written in MSL.

3. Metal Performance Shaders (MPS): A library of pre-built, highly optimized GPU kernels for common operations – matrix multiplication, convolution, image processing, neural network layers, and more. Think of MPS as Apple’s cuDNN equivalent.

For ML inference, we primarily care about the compute side of Metal. We will barely touch graphics (vertex/fragment shaders, render passes, etc.). Our world is compute kernels, buffers, and dispatches.

Metal vs. The Competition

If you are coming from CUDA, OpenCL, or Vulkan, you will find Metal familiar in some ways and different in others. Let us compare:

Metal vs. CUDA

+-------------------+------------------------------------------+
|     Aspect        |          CUDA           |     Metal      |
+-------------------+-------------------------+----------------+
| Vendor            | NVIDIA only             | Apple only     |
| Language          | CUDA C/C++ (extended)   | MSL (C++14     |
|                   |                         |   extended)    |
| Host API          | CUDA Runtime/Driver     | Metal API      |
|                   | (C/C++)                 | (Swift/ObjC)   |
| Execution model   | Grid → Block → Thread   | Grid →         |
|                   |                         |  Threadgroup → |
|                   |                         |  Thread        |
| SIMD width        | 32 (warp)               | 32 (SIMD group)|
| Shared memory     | __shared__              | threadgroup    |
|                   |                         |   address space|
| Sync primitive    | __syncthreads()         | threadgroup_   |
|                   |                         |   barrier()    |
| Matrix accel.     | Tensor Cores (wmma)     | simdgroup_     |
|                   |                         |   matrix       |
| Memory model      | Discrete + UVA          | Unified (UMA)  |
| Ecosystem         | Massive (cuDNN, cuBLAS, | Smaller (MPS,  |
|                   |  TensorRT, Triton)      |  MPSGraph,     |
|                   |                         |  Core ML)      |
| Maturity for ML   | 15+ years               | ~5 years       |
+-------------------+-------------------------+----------------+

The biggest conceptual difference: CUDA’s programming model treats the GPU as a separate device with its own memory.³ You explicitly copy data between CPU and GPU. On Apple Silicon with Metal, the CPU and GPU share the same physical memory (UMA). There is no copy – you allocate a buffer once, and both CPU and GPU can access it.

The biggest practical difference: CUDA has a massive ecosystem for ML. Libraries like cuDNN, cuBLAS, TensorRT, and Triton make it possible to write high-performance ML code without touching raw CUDA kernels. Metal’s ecosystem is smaller. MPS provides some building blocks, but for state-of-the-art inference, you often need to write custom kernels – which is exactly what akunu does.

Metal vs. OpenCL

+-------------------+------------------------------------------+
|     Aspect        |         OpenCL          |     Metal      |
+-------------------+-------------------------+----------------+
| Portability       | Cross-platform          | Apple only     |
| API style         | C-based, verbose        | ObjC/Swift,    |
|                   |                         |   modern       |
| Shader language   | OpenCL C (C99-based)    | MSL (C++14)    |
| Runtime compile   | Yes (common)            | Yes + offline  |
| Performance       | Driver-dependent        | Tuned for      |
|                   |                         |   Apple HW     |
| Status on Apple   | Deprecated since        | Active, primary|
|                   |   macOS 10.14           |   GPU API      |
+-------------------+-------------------------+----------------+

OpenCL was once available on macOS, but Apple deprecated it in 2018. Metal is the only supported path for GPU compute on Apple platforms. If you have existing OpenCL kernels, they need to be ported to MSL.

The good news: the conceptual mapping is straightforward. OpenCL work-groups are Metal threadgroups. OpenCL __local memory is Metal threadgroup memory. OpenCL __global is Metal device. The languages are different (C99 vs C++14), but the GPU programming model is fundamentally the same.

Metal vs. Vulkan

+-------------------+------------------------------------------+
|     Aspect        |         Vulkan          |     Metal      |
+-------------------+-------------------------+----------------+
| Portability       | Cross-platform          | Apple only     |
| Verbosity         | Extremely verbose       | Moderate       |
| Shader language   | SPIR-V (usually from    | MSL (write     |
|                   |   GLSL/HLSL)            |   directly)    |
| Compute support   | Full                    | Full           |
| Validation layers | External (very helpful) | Metal API      |
|                   |                         |   Validation   |
| Driver overhead   | Very low                | Very low       |
| On Apple          | Via MoltenVK (wrapper   | Native         |
|                   |   over Metal)           |                |
+-------------------+-------------------------+----------------+

Vulkan and Metal share the same philosophy: explicit, low-overhead GPU control. Both require you to manage command buffers, synchronization, and pipeline states yourself. Vulkan is more verbose (creating a compute pipeline in Vulkan can take hundreds of lines), while Metal strikes a balance between control and usability.

Fun fact: MoltenVK, the Vulkan-on-Apple implementation, is actually a translation layer that converts Vulkan calls to Metal calls underneath. So Metal is the true native API.

Summary: Why Metal for ML on Apple?

  Why Metal?
  ==========

  1. It is the ONLY way to access Apple GPU compute
     (OpenCL is deprecated, no CUDA, Vulkan is via MoltenVK)

  2. Unified Memory Architecture means ZERO-COPY buffer sharing
     between CPU and GPU -- huge for inference

  3. Metal Shading Language is pleasant to write
     (C++14 with nice extensions, not as painful as GLSL)

  4. simdgroup_matrix operations give you hardware-accelerated
     matrix multiply (like Tensor Cores)

  5. Apple tunes Metal drivers specifically for their hardware
     (you get the best possible performance)

The Metal Ecosystem

Let us zoom in on the three pillars of the Metal ecosystem and understand how they fit together for ML workloads.

The Metal API (Host Side)

The Metal API is an Objective-C/Swift framework. You use it on the CPU to orchestrate GPU work. The key objects are:

  +----------------------------------------------------+
  |                  Your Application                   |
  |                  (Swift / ObjC / C++)               |
  +----------------------------------------------------+
              |
              v
  +----------------------------------------------------+
  |                  MTLDevice                          |
  |  Represents the GPU. Entry point for everything.   |
  |  - makeCommandQueue()                              |
  |  - makeBuffer(length:options:)                     |
  |  - makeComputePipelineState(function:)             |
  |  - makeLibrary(source:options:)                    |
  +----------------------------------------------------+
              |
              +------------------+------------------+
              |                  |                  |
              v                  v                  v
  +-----------------+  +-----------------+  +-----------------+
  | MTLCommandQueue |  | MTLBuffer       |  | MTLLibrary      |
  | Ordered queue   |  | GPU memory      |  | Collection of   |
  | of cmd buffers  |  | allocation      |  | compiled shaders|
  +-----------------+  +-----------------+  +-----------------+
              |                                     |
              v                                     v
  +-----------------+                     +-----------------+
  | MTLCommandBuffer|                     | MTLFunction     |
  | A batch of GPU  |                     | A single shader |
  | commands        |                     | entry point     |
  +-----------------+                     +-----------------+
              |                                     |
              v                                     v
  +-----------------------+          +---------------------------+
  |MTLComputeCommandEncoder|          |MTLComputePipelineState   |
  | Records compute cmds  |          | Compiled, ready-to-run   |
  | (set buffers, dispatch)|          | version of a kernel      |
  +-----------------------+          +---------------------------+

We will explore each of these objects in detail in Chapter 7. For now, just know the flow: you create a device, create a command queue, create command buffers, encode commands into them, and commit them to the GPU.

The Metal Shading Language (MSL)

MSL is the language you write GPU programs in. It looks like C++14 with some extra keywords and types:

// A simple MSL compute kernel
kernel void add_arrays(
    device const float* a [[buffer(0)]],
    device const float* b [[buffer(1)]],
    device float* result   [[buffer(2)]],
    uint id [[thread_position_in_grid]]
) {
    result[id] = a[id] + b[id];
}

Key MSL features for compute:

• Address space qualifiers: device, constant, threadgroup, thread
• Attribute syntax: [[buffer(0)]], [[thread_position_in_grid]], [[kernel]]
• Built-in vector types: half, half4, float4, uint2, etc.
• SIMD group intrinsics: simd_sum(), simd_shuffle(), etc.
• Threadgroup memory: shared memory within a threadgroup
• SIMD group matrix ops: simdgroup_half8x8, simdgroup_multiply_accumulate()

MSL is covered in depth in Chapter 8.

Metal Performance Shaders (MPS)

MPS is Apple’s library of optimized GPU kernels. For ML, the most relevant parts are:

• MPSMatrixMultiplication: Optimized GEMM
• MPSImageConvolution: 2D convolution
• MPSNNGraph: Neural network inference graph (older API)
• MPSGraph: A more modern compute graph framework

MPS is useful as a starting point, but for state-of-the-art inference performance, custom kernels often outperform MPS. This is because:

1. MPS kernels are general-purpose. A custom kernel can be specialized for exact matrix dimensions, quantization formats, and fusion patterns.
2. MPS cannot fuse arbitrary operations. A custom kernel can fuse (say) dequantization + matrix multiply + bias add + activation into a single pass, saving memory bandwidth.
3. MPS does not support all quantization formats used in modern LLM inference.

This is exactly why akunu writes its own Metal kernels rather than relying on MPS.

When to Use Metal Compute vs. Metal Graphics

Metal supports two kinds of GPU work:

1. Graphics (render pipelines): Vertex shaders, fragment shaders, rasterization, render passes. This is for drawing things on screen.
2. Compute (compute pipelines): General-purpose computation on the GPU. No rendering, no pixels – just data in, data out.

For ML inference, we use compute exclusively. Here is why:

  Graphics Pipeline:                 Compute Pipeline:
  ==================                ==================

  Vertices → Vertex Shader          Data → Compute Kernel → Data
           → Rasterizer
           → Fragment Shader
           → Framebuffer

  - Fixed-function stages           - Fully programmable
  - Designed for rendering          - Designed for GPGPU
  - Data flows through              - You control data flow
    a rigid pipeline                  completely
  - Output is pixels                - Output is whatever you want

Compute pipelines give us:

• Arbitrary data access patterns: Read from and write to any buffer location
• Threadgroup shared memory: Fast scratchpad for inter-thread communication
• Flexible dispatch: 1D, 2D, or 3D grids of arbitrary size
• No rendering overhead: No rasterizer, no framebuffer, no blend state

There are rare cases where graphics shaders are (ab)used for compute – for example, some older GPU compute techniques use fragment shaders to process textures. But on modern Apple GPUs, compute shaders are the right tool for ML.

The Metal Programming Model

Now let us build up the mental model for how Metal compute works. This is the single most important section in this chapter, so take your time with it.

The Big Picture

Here is the full pipeline from your application to GPU execution:

  YOUR APPLICATION (CPU)
  ======================

  1. Get a reference to the GPU
     +------------------+
     | MTLDevice        |  <-- Represents the GPU hardware
     +------------------+
              |
  2. Create a command queue (once, reuse it)
              |
              v
     +------------------+
     | MTLCommandQueue  |  <-- FIFO queue of command buffers
     +------------------+
              |
  3. Create a command buffer (one per "batch" of work)
              |
              v
     +------------------+
     | MTLCommandBuffer |  <-- Container for GPU commands
     +------------------+
              |
  4. Create a compute command encoder
              |
              v
     +--------------------------+
     | MTLComputeCommandEncoder |  <-- Records compute commands
     +--------------------------+
              |
  5. Set the pipeline state (which kernel to run)
  6. Set buffers (input/output data)
  7. Dispatch threadgroups (how many threads to launch)
  8. End encoding
              |
              v
  9. Commit the command buffer to the GPU
              |
              v

  GPU EXECUTION
  =============

  The GPU picks up the command buffer from the queue,
  executes the recorded commands:
    - Binds the kernel
    - Binds the buffers
    - Launches threadgroups across GPU cores
    - Each thread runs the kernel function
    - Results are written to output buffers

              |
              v

  RESULTS AVAILABLE
  =================
  (In the output buffer, which on UMA is already
   accessible to the CPU -- no copy needed!)

Let us walk through each component.

MTLDevice – The GPU

Everything starts with a MTLDevice. This object represents the GPU hardware. You get it like this:

// Swift
guard let device = MTLCreateSystemDefaultDevice() else {
    fatalError("Metal is not supported on this device")
}

Or in Objective-C:

id<MTLDevice> device = MTLCreateSystemDefaultDevice();

On a Mac with Apple Silicon, this gives you the built-in GPU. On a Mac with multiple GPUs (e.g., an older Mac Pro with an eGPU), you can enumerate all devices with MTLCopyAllDevices().

The device is your factory for creating everything else: buffers, command queues, pipeline states, libraries.

MTLCommandQueue – The Submission Highway

A command queue is an ordered sequence of command buffers. You typically create one queue at startup and reuse it for the lifetime of your application:

guard let commandQueue = device.makeCommandQueue() else {
    fatalError("Could not create command queue")
}

Think of it as a highway on-ramp. Command buffers you commit to the queue will be executed in order (mostly – Metal can reorder independent work for efficiency, but the observable results respect submission order for resources).

MTLCommandBuffer – A Batch of Work

A command buffer is a container for GPU commands. You create one whenever you have work to submit:

guard let commandBuffer = commandQueue.makeCommandBuffer() else {
    fatalError("Could not create command buffer")
}

A command buffer can contain multiple encoder passes. For compute work, each pass uses a MTLComputeCommandEncoder. The command buffer is not executed until you call commit().

  Command Buffer Lifecycle:
  =========================

  Created ──> Encoding ──> Committed ──> Scheduled ──> Completed
    |            |             |             |              |
    |   (you record    (you call     (Metal       (GPU has
    |    commands)     .commit())    schedules     finished)
    |                               execution)

MTLComputeCommandEncoder – Recording Commands

The encoder is how you record commands into the command buffer. For compute work:

guard let encoder = commandBuffer.makeComputeCommandEncoder() else {
    fatalError("Could not create compute encoder")
}

// Record commands:
encoder.setComputePipelineState(pipelineState)
encoder.setBuffer(inputBuffer, offset: 0, index: 0)
encoder.setBuffer(outputBuffer, offset: 0, index: 1)
encoder.dispatchThreadgroups(gridSize, threadsPerThreadgroup: groupSize)
encoder.endEncoding()

Important: the encoder does not execute anything. It records commands into the command buffer. Execution happens later when you commit.

MTLComputePipelineState – The Compiled Kernel

Before you can dispatch a kernel, you need to compile it into a pipeline state object (PSO). This involves:

1. Loading your MSL source code (or a pre-compiled .metallib binary)
2. Getting a MTLFunction from the library
3. Creating a MTLComputePipelineState from the function

// Load the default library (compiled from .metal files in your project)
let library = device.makeDefaultLibrary()!

// Get the kernel function by name
let function = library.makeFunction(name: "add_arrays")!

// Create the pipeline state (this compiles the function for the GPU)
let pipelineState = try device.makeComputePipelineState(function: function)

Creating a PSO can be expensive (it involves final compilation and optimization), so you typically do it once at startup and cache the result. We will explore this in detail in Chapter 7.

Dispatch – Launching Threads

The final piece is telling the GPU how many threads to launch. Metal gives you two options:

// Option 1: Dispatch by threadgroup count
// You specify: (number of threadgroups) x (threads per threadgroup)
let gridSize = MTLSize(width: 64, height: 1, depth: 1)
let groupSize = MTLSize(width: 256, height: 1, depth: 1)
encoder.dispatchThreadgroups(gridSize, threadsPerThreadgroup: groupSize)
// Total threads = 64 * 256 = 16,384

// Option 2: Dispatch by total thread count (Metal adjusts automatically)
let totalThreads = MTLSize(width: 16384, height: 1, depth: 1)
let groupSize = MTLSize(width: 256, height: 1, depth: 1)
encoder.dispatchThreads(totalThreads, threadsPerThreadgroup: groupSize)

Option 1 (dispatchThreadgroups) requires you to calculate the grid dimensions yourself. Option 2 (dispatchThreads) lets you specify the total number of threads, and Metal handles the math. We will discuss the tradeoffs in Chapter 9.

Putting It All Together: Hello World Compute Shader

Let us write a complete example that adds two arrays on the GPU. This is the “Hello World” of GPU computing.

Step 1: The Metal Shader (MSL)

Create a file called compute.metal:

// compute.metal
// A simple kernel that adds two arrays element-wise.

#include <metal_stdlib>
using namespace metal;

kernel void add_arrays(
    device const float* inA  [[buffer(0)]],
    device const float* inB  [[buffer(1)]],
    device float*       out  [[buffer(2)]],
    uint id [[thread_position_in_grid]]
) {
    out[id] = inA[id] + inB[id];
}

Let us break down every piece:

• #include <metal_stdlib> – Includes standard Metal functions and types.
• using namespace metal; – Avoids having to prefix everything with metal::.
• kernel void add_arrays(...) – The kernel keyword marks this as a compute kernel entry point. It must return void.
• device const float* inA [[buffer(0)]] – A pointer to read-only data in the device address space (GPU-accessible memory). The [[buffer(0)]] attribute tells Metal this is bound at buffer index 0.
• device float* out [[buffer(2)]] – A writable output buffer at index 2.
• uint id [[thread_position_in_grid]] – A built-in variable that gives each thread its unique index in the dispatch grid. Thread 0 gets id=0, thread 1 gets id=1, etc.

The kernel body is trivial: each thread reads one element from inA and inB, adds them, and writes the result to out.

  Visualization of execution:
  ===========================

  Thread 0:  out[0] = inA[0] + inB[0]
  Thread 1:  out[1] = inA[1] + inB[1]
  Thread 2:  out[2] = inA[2] + inB[2]
  ...
  Thread N:  out[N] = inA[N] + inB[N]

  Each thread handles exactly one element.
  All threads execute in parallel across GPU cores.

Step 2: The Host Code (Swift)

Here is the complete Swift code to set up Metal, compile the kernel, prepare data, dispatch the kernel, and read the results:

import Metal
import Foundation

// ============================================================
// STEP 1: Get the GPU device
// ============================================================
guard let device = MTLCreateSystemDefaultDevice() else {
    fatalError("Metal is not supported on this device")
}
print("Using GPU: \(device.name)")

// ============================================================
// STEP 2: Create a command queue
// ============================================================
guard let commandQueue = device.makeCommandQueue() else {
    fatalError("Could not create command queue")
}

// ============================================================
// STEP 3: Load and compile the shader
// ============================================================

// Option A: Load from a .metal file in the project bundle
// let library = device.makeDefaultLibrary()!

// Option B: Compile from source string at runtime
let shaderSource = """
#include <metal_stdlib>
using namespace metal;

kernel void add_arrays(
    device const float* inA  [[buffer(0)]],
    device const float* inB  [[buffer(1)]],
    device float*       out  [[buffer(2)]],
    uint id [[thread_position_in_grid]]
) {
    out[id] = inA[id] + inB[id];
}
"""

let library = try! device.makeLibrary(source: shaderSource, options: nil)
let function = library.makeFunction(name: "add_arrays")!
let pipelineState = try! device.makeComputePipelineState(function: function)

// ============================================================
// STEP 4: Prepare the data
// ============================================================
let arrayLength = 1_000_000
let bufferSize = arrayLength * MemoryLayout<Float>.size

// Create Metal buffers with shared storage mode (CPU + GPU access)
let bufferA = device.makeBuffer(length: bufferSize, options: .storageModeShared)!
let bufferB = device.makeBuffer(length: bufferSize, options: .storageModeShared)!
let bufferOut = device.makeBuffer(length: bufferSize, options: .storageModeShared)!

// Fill input buffers with data
let pointerA = bufferA.contents().bindMemory(to: Float.self, capacity: arrayLength)
let pointerB = bufferB.contents().bindMemory(to: Float.self, capacity: arrayLength)

for i in 0..<arrayLength {
    pointerA[i] = Float(i)
    pointerB[i] = Float(i) * 2.0
}

// ============================================================
// STEP 5: Create a command buffer and encoder
// ============================================================
let commandBuffer = commandQueue.makeCommandBuffer()!
let encoder = commandBuffer.makeComputeCommandEncoder()!

// ============================================================
// STEP 6: Encode the compute command
// ============================================================
encoder.setComputePipelineState(pipelineState)
encoder.setBuffer(bufferA, offset: 0, index: 0)
encoder.setBuffer(bufferB, offset: 0, index: 1)
encoder.setBuffer(bufferOut, offset: 0, index: 2)

// Calculate dispatch sizes
let threadGroupSize = MTLSize(width: 256, height: 1, depth: 1)
let threadGroups = MTLSize(
    width: (arrayLength + 255) / 256,  // Ceiling division
    height: 1,
    depth: 1
)
encoder.dispatchThreadgroups(threadGroups, threadsPerThreadgroup: threadGroupSize)
encoder.endEncoding()

// ============================================================
// STEP 7: Commit and wait
// ============================================================
commandBuffer.commit()
commandBuffer.waitUntilCompleted()

// ============================================================
// STEP 8: Read the results
// ============================================================
let pointerOut = bufferOut.contents().bindMemory(to: Float.self, capacity: arrayLength)

// Verify a few results
for i in 0..<5 {
    print("out[\(i)] = \(pointerOut[i])  (expected: \(Float(i) + Float(i) * 2.0))")
}
// Output:
// out[0] = 0.0  (expected: 0.0)
// out[1] = 3.0  (expected: 3.0)
// out[2] = 6.0  (expected: 6.0)
// out[3] = 9.0  (expected: 9.0)
// out[4] = 12.0  (expected: 12.0)

The Full Flow Visualized

Here is what happens when you run this code:

  CPU Side                              GPU Side
  ========                              ========

  1. MTLCreateSystemDefaultDevice()
     +---> Gets handle to GPU

  2. makeCommandQueue()
     +---> Creates submission queue

  3. makeLibrary(source:)               Compiler: MSL → AIR → GPU ISA
     makeFunction(name:)
     makeComputePipelineState()

  4. makeBuffer() x 3                   Allocates in unified memory
     Fill bufferA, bufferB              (same physical RAM)

  5. makeCommandBuffer()
     makeComputeCommandEncoder()

  6. setComputePipelineState()           |
     setBuffer() x 3                    | All recorded, not
     dispatchThreadgroups()             | yet executed
     endEncoding()                      |

  7. commandBuffer.commit()
     +----------------------------------->  GPU picks up work
                                            |
     commandBuffer.waitUntilCompleted()     v
     (CPU blocks here)                      GPU launches 3,907
                                            threadgroups of 256
                                            threads each
                                            |
                                            v
                                            Each thread runs
                                            add_arrays kernel
                                            |
                                            v
                                            Results written to
                                            bufferOut
                                            |
     <-----(completion signal)------------- Done!

  8. Read pointerOut[i]                 (Same physical memory,
                                         no copy needed!)

Notice step 8: we read the results directly from the buffer’s contents() pointer. There is no “copy back from GPU” step. This is the UMA advantage – the buffer lives in unified memory, accessible to both CPU and GPU.

The Objective-C Version

Since akunu is written in C/C++ and uses the Metal Objective-C API (via Objective-C++), here is the same example in Objective-C:

#import <Metal/Metal.h>
#import <Foundation/Foundation.h>

int main(int argc, const char * argv[]) {
    @autoreleasepool {
        // 1. Get the device
        id<MTLDevice> device = MTLCreateSystemDefaultDevice();
        NSLog(@"Using GPU: %@", device.name);

        // 2. Create command queue
        id<MTLCommandQueue> commandQueue = [device newCommandQueue];

        // 3. Compile the shader
        NSString *shaderSource = @
            "#include <metal_stdlib>\n"
            "using namespace metal;\n"
            "kernel void add_arrays(\n"
            "    device const float* inA  [[buffer(0)]],\n"
            "    device const float* inB  [[buffer(1)]],\n"
            "    device float*       out  [[buffer(2)]],\n"
            "    uint id [[thread_position_in_grid]]\n"
            ") {\n"
            "    out[id] = inA[id] + inB[id];\n"
            "}\n";

        NSError *error = nil;
        id<MTLLibrary> library = [device newLibraryWithSource:shaderSource
                                                      options:nil
                                                        error:&error];
        id<MTLFunction> function = [library newFunctionWithName:@"add_arrays"];
        id<MTLComputePipelineState> pso =
            [device newComputePipelineStateWithFunction:function error:&error];

        // 4. Create buffers
        NSUInteger arrayLength = 1000000;
        NSUInteger bufferSize = arrayLength * sizeof(float);

        id<MTLBuffer> bufA = [device newBufferWithLength:bufferSize
                                options:MTLResourceStorageModeShared];
        id<MTLBuffer> bufB = [device newBufferWithLength:bufferSize
                                options:MTLResourceStorageModeShared];
        id<MTLBuffer> bufOut = [device newBufferWithLength:bufferSize
                                  options:MTLResourceStorageModeShared];

        // Fill input data
        float *ptrA = (float *)bufA.contents;
        float *ptrB = (float *)bufB.contents;
        for (NSUInteger i = 0; i < arrayLength; i++) {
            ptrA[i] = (float)i;
            ptrB[i] = (float)i * 2.0f;
        }

        // 5-7. Encode and dispatch
        id<MTLCommandBuffer> cmdBuf = [commandQueue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmdBuf computeCommandEncoder];

        [enc setComputePipelineState:pso];
        [enc setBuffer:bufA offset:0 atIndex:0];
        [enc setBuffer:bufB offset:0 atIndex:1];
        [enc setBuffer:bufOut offset:0 atIndex:2];

        MTLSize groupSize = MTLSizeMake(256, 1, 1);
        MTLSize gridSize = MTLSizeMake((arrayLength + 255) / 256, 1, 1);
        [enc dispatchThreadgroups:gridSize threadsPerThreadgroup:groupSize];
        [enc endEncoding];

        [cmdBuf commit];
        [cmdBuf waitUntilCompleted];

        // 8. Read results
        float *ptrOut = (float *)bufOut.contents;
        for (int i = 0; i < 5; i++) {
            NSLog(@"out[%d] = %.1f (expected: %.1f)", i, ptrOut[i],
                  (float)i + (float)i * 2.0f);
        }
    }
    return 0;
}

The Objective-C API maps almost 1:1 to Swift. The main difference is syntax ([device newCommandQueue] vs device.makeCommandQueue()). akunu uses Objective-C++ so it can mix C++ code with Metal API calls seamlessly.

How akunu Uses Metal

Now that you understand the basics, let us peek at how akunu structures its Metal usage. We will go much deeper in Part IV, but a high-level overview helps connect the dots:

  akunu Architecture (simplified):
  =================================

  +----------------------------------------------------------+
  |                     akunu C API                          |
  |  ak_context_create()  ak_generate()  ak_decode_step()   |
  +----------------------------------------------------------+
              |
              v
  +----------------------------------------------------------+
  |                   MetalDevice                            |
  |                                                          |
  |  - device_: id<MTLDevice>                                |
  |  - queue_: id<MTLCommandQueue>                           |
  |  - pso_cache_: HashMap<string, MTLComputePipelineState>  |
  |                                                          |
  |  Methods:                                                |
  |  - allocate(size) → Buffer                               |
  |  - begin_encoding() → starts recording                   |
  |  - set_pipeline(Pipeline) → sets the kernel              |
  |  - set_buffer(Buffer, offset, index) → binds data        |
  |  - dispatch(grid, threadgroup) → launches threads        |
  |  - end_encoding() → finishes recording                   |
  |  - commit() → submits to GPU                             |
  +----------------------------------------------------------+
              |
              v
  +----------------------------------------------------------+
  |                   Metal Shaders                          |
  |  (.metal files compiled to .metallib)                    |
  |                                                          |
  |  - gemv_f16.metal     (matrix-vector multiply)           |
  |  - gemm_f16.metal     (matrix-matrix multiply)           |
  |  - attention.metal    (flash attention)                   |
  |  - rms_norm.metal     (RMS normalization)                |
  |  - rope.metal         (rotary position embedding)        |
  |  - ...                                                   |
  +----------------------------------------------------------+

akunu wraps the Metal API in a MetalDevice class that provides a cleaner, C++-friendly interface. Instead of creating MTLCommandBuffer and MTLComputeCommandEncoder objects directly, you call methods like begin_encoding(), set_pipeline(), set_buffer(), and dispatch().

The pipeline state objects (PSOs) are cached in a hash map (pso_cache_) so they are only compiled once. The cache key includes the kernel name and any specialization constants (e.g., "gemv_k128" for a GEMV kernel specialized for K=128).

Key Takeaways

Before moving to Chapter 7 where we dive deep into compute pipelines, let us summarize what we have learned:

1. Metal is Apple’s low-level GPU API – the only way to do GPU compute on Apple platforms.

2. Metal has three parts: the host API (Swift/ObjC), the Metal Shading Language (MSL), and Metal Performance Shaders (MPS).

3. For ML inference, we use compute pipelines, not graphics pipelines.

4. The programming model flows:

   Device → Queue → Command Buffer → Encoder → Dispatch → Commit
   

5. On Apple Silicon, UMA means zero-copy: buffers allocated with .storageModeShared are accessible to both CPU and GPU without any data transfer.

6. MSL is C++14 with GPU extensions: address spaces, vector types, SIMD intrinsics, and thread indexing built-ins.

7. akunu wraps Metal in a MetalDevice class that caches pipeline states and provides a streamlined C++ interface.

In the next chapter, we will go deep on compute pipelines: how kernels are compiled, how pipeline states are created and cached, and how the command encoding flow works in practice.

Exercises

1. Run the Hello World example: If you have a Mac with Apple Silicon, create an Xcode project (macOS Command Line Tool), add a .metal file with the add_arrays kernel, and run the Swift host code. Verify the output.

2. Modify the kernel: Change the kernel to compute out[i] = inA[i] * inB[i] + 1.0 (fused multiply-add). How does the host code change? (Hint: it does not, only the shader changes.)

3. Explore MTLDevice properties: Print device.maxThreadgroupMemoryLength, device.maxThreadsPerThreadgroup, and device.name. What values do you get on your hardware?

4. Think about error handling: Our example uses fatalError and force-unwraps everywhere. In production code (like akunu), what error handling strategy would you use?

5. Compare with CUDA: If you have experience with CUDA, write down the mapping between CUDA concepts and Metal concepts. For example: cudaMalloc maps to device.makeBuffer, <<<grid, block>>> maps to dispatchThreadgroups, etc.

1. Apple. “Metal Programming Guide.” developer.apple.com. The official reference for Metal’s programming model, including device creation, command queues, and compute pipelines. See https://developer.apple.com/documentation/metal/performing-calculations-on-a-gpu. ↩

2. Apple. “Metal Shading Language Specification.” developer.apple.com. The formal MSL specification covering types, address spaces, and built-in functions. See https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf. ↩

3. Apple. “Metal Best Practices Guide.” developer.apple.com. Performance guidance for buffer management, pipeline caching, and dispatch strategies. See https://developer.apple.com/library/archive/documentation/3DDrawing/Conceptual/MTLBestPracticesGuide/index.html. ↩

Compute Pipelines and Shaders

If you have ever worked with graphics APIs, the word “pipeline” probably conjures images of vertex shaders, rasterizers, fragment shaders, and blend states all wired together into a monolithic object. Compute pipelines are refreshingly simpler. A compute pipeline is, at its core, just one thing: a compiled GPU function plus the metadata the driver needs to launch it. That is it. No rasterizer, no depth test, no blend modes. One function, ready to run.

In this chapter we are going to trace the full lifecycle of a compute pipeline on Metal – from writing a kernel function in the Metal Shading Language, through the compilation stages that turn it into executable GPU code, all the way to encoding a dispatch command that actually runs the thing. Along the way we will see how akunu wraps this machinery to keep things fast and ergonomic from Rust.

What Is a Compute Pipeline State Object?

A MTLComputePipelineState (PSO) is an opaque, immutable object that represents a fully compiled and optimized GPU kernel. Once you have one, you can use it over and over again to dispatch work – the driver does not need to recompile anything. Creating a PSO is expensive (potentially tens of milliseconds), but using one is cheap (microseconds).

Think of it like compiling a C++ program. The compilation step is slow, but once you have the binary, running it is fast. You would never recompile the same source file every time you want to run the program. Same idea here.

+------------------+       +------------------+       +-------------------+
|  MSL Source Code | ----> |  MTLLibrary      | ----> | MTLFunction       |
|  (kernel void    |       |  (collection of  |       | (one specific     |
|   my_kernel...)  |       |   compiled       |       |  kernel function) |
+------------------+       |   functions)     |       +-------------------+
                           +------------------+              |
                                                             v
                                                    +-------------------+
                                                    | MTLComputePipeline|
                                                    | State (PSO)       |
                                                    | (ready to dispatch|
                                                    |  on GPU)          |
                                                    +-------------------+

Let us walk through each step.

The Compilation Flow: MSL to PSO

Metal has a multi-stage compilation pipeline. Understanding these stages matters because each one is a knob you can turn for performance.

Stage 1: MSL Source to AIR (Metal Intermediate Representation)

The Metal Shading Language is a dialect of C++14 with GPU-specific extensions (we will cover MSL in depth in the next chapter). When you compile MSL, the first output is AIR – Apple’s Intermediate Representation. AIR is analogous to LLVM IR (and in fact the Metal compiler is built on LLVM). It is a platform-independent representation of your shader.

+-----------+     metal compiler      +----------+
| .metal    |  ------------------->   | .air     |
| (MSL src) |    (clang frontend)     | (LLVM IR |
+-----------+                         |  variant)|
                                      +----------+

Stage 2: AIR to metallib (Metal Library)

Multiple AIR files are linked together into a .metallib file. This is a binary archive that contains all your compiled functions. You can think of it as the GPU equivalent of a .a static library or a .dylib.

+----------+
| .air     |--+
+----------+  |    metallib tool      +------------+
              +------------------->   | .metallib  |
+----------+  |                       | (binary    |
| .air     |--+                       |  archive)  |
+----------+                          +------------+

Stage 3: metallib to PSO (at runtime)

At runtime, you load the .metallib into a MTLLibrary, look up a function by name, and then create a pipeline state. This final step performs GPU-specific optimizations – register allocation, instruction scheduling, occupancy tuning – that depend on the specific GPU model the code will run on.

+------------+    MTLDevice           +-------------+     MTLDevice
| .metallib  | ------------------>   | MTLFunction  | ----------------->
| (binary)   |  newLibrary(data:)    | (looked up   |  newComputePipeline
+------------+                       |  by name)    |  State(function:)
                                     +-------------+
                                                          +----------+
                                                          | PSO      |
                                                          | (GPU-    |
                                                          |  ready)  |
                                                          +----------+

Putting It All Together

Here is the full pipeline:

MSL Source (.metal)
       |
       | metal compiler (offline or runtime)
       v
AIR (.air)
       |
       | metallib linker
       v
Metal Library (.metallib)
       |
       | Load at runtime: device.newLibrary(data:)
       v
MTLLibrary (in-memory)
       |
       | library.newFunction(name: "my_kernel")
       v
MTLFunction
       |
       | device.newComputePipelineState(function:)
       v
MTLComputePipelineState (PSO)  <-- this is what you dispatch with

Runtime vs Offline Compilation

You have two choices for when compilation happens.

Runtime Compilation (from source)

You embed MSL source code as a string and compile it at runtime:

// Runtime compilation from source string
let source = """
kernel void add_arrays(device float* a [[buffer(0)]],
                       device float* b [[buffer(1)]],
                       device float* c [[buffer(2)]],
                       uint id [[thread_position_in_grid]]) {
    c[id] = a[id] + b[id];
}
"""
let library = try device.makeLibrary(source: source, options: nil)
let function = library.makeFunction(name: "add_arrays")!
let pso = try device.makeComputePipelineState(function: function)

This is convenient for development and for cases where you need to generate shader code dynamically. The downside is that compilation happens at application startup, which can take tens to hundreds of milliseconds per kernel.

Offline Compilation (precompiled metallib)

You compile your shaders ahead of time using Apple’s command-line tools:

# Step 1: Compile MSL to AIR
xcrun -sdk macosx metal -c my_kernels.metal -o my_kernels.air

# Step 2: Link AIR into a metallib
xcrun -sdk macosx metallib my_kernels.air -o my_kernels.metallib

# Or do it in one step:
xcrun -sdk macosx metal my_kernels.metal -o my_kernels.metallib

Then at runtime, you just load the precompiled binary:

let libraryURL = Bundle.main.url(forResource: "my_kernels",
                                  withExtension: "metallib")!
let library = try device.makeLibrary(URL: libraryURL)

This is what production applications should do. The runtime cost drops from “compile everything” to “just load a binary and create PSOs.” Akunu takes the offline compilation approach – its Metal kernels are precompiled into metallib files that ship with the library.

When Would You Use Runtime Compilation?

Runtime compilation is not just for lazy developers. There are legitimate reasons to compile at runtime:

1. Generated kernels: When the kernel code depends on runtime parameters like tensor shapes, data types, or hardware capabilities.
2. JIT specialization: Generating code paths optimized for specific input sizes that are only known at inference time.
3. Plugin systems: Loading user-provided shader code.

That said, for a production ML inference engine, offline compilation is almost always the right default.

Function Specialization Constants

Here is where things get interesting for ML workloads. Metal supports function constants – compile-time constants that let you create specialized variants of the same kernel without duplicating source code.

Consider a GEMV (matrix-vector multiply) kernel. The optimal implementation depends on the matrix dimensions:

// Define function constants
constant uint K [[function_constant(0)]];
constant bool USE_BIAS [[function_constant(1)]];

kernel void gemv(device const half* W [[buffer(0)]],
                 device const half* x [[buffer(1)]],
                 device half* y [[buffer(2)]],
                 device const half* bias [[buffer(3)]],
                 uint row [[thread_position_in_grid]]) {
    
    half sum = 0;
    // The compiler knows K at compile time -- it can fully
    // unroll this loop, choose optimal vector widths, etc.
    for (uint i = 0; i < K; i += 4) {
        half4 w = *((device const half4*)(W + row * K + i));
        half4 v = *((device const half4*)(x + i));
        sum += dot(w, v);
    }
    
    // Dead code elimination removes this branch entirely
    // when USE_BIAS is false
    if (USE_BIAS) {
        sum += bias[row];
    }
    
    y[row] = sum;
}

At PSO creation time, you provide the constant values:

let constants = MTLFunctionConstantValues()

var k: UInt32 = 128
constants.setConstantValue(&k, type: .uint, index: 0)

var useBias: Bool = true
constants.setConstantValue(&useBias, type: .bool, index: 1)

let function = try library.makeFunction(name: "gemv",
                                         constantValues: constants)
let pso = try device.makeComputePipelineState(function: function)

The compiler treats these constants exactly like #define values – it can unroll loops, eliminate dead branches, choose optimal vector sizes, and perform constant propagation. The result is a specialized binary for this specific combination of constant values.

Why does this matter for ML? Because ML models have fixed dimensions. A LLaMA-7B model always has a hidden dimension of 4096. An attention head is always 128-dimensional. Once you know the model architecture, you know exactly what constants to specialize on, and the compiler can produce code that is perfectly tuned for those dimensions.

Pipeline Caching Strategies

Creating a PSO involves GPU-specific compilation. On an M1 Max, creating a single pipeline takes somewhere between 5-50ms depending on kernel complexity. If you have 50 kernels, that is potentially seconds of startup time. And if you use function constants, each unique combination of constants is a separate PSO.

The solution is caching. Metal provides MTLBinaryArchive for persisting compiled pipelines to disk, but the more common approach in practice is to cache PSOs in memory and be smart about when you create them.

How Akunu Caches PSOs

Akunu maintains a PSO cache in its MetalDevice struct. The cache is a hash map keyed by a string that encodes the kernel name plus all specialization constant values:

PSO Cache (HashMap<String, MTLComputePipelineState>)
+-------------------------------------------+
| Key                    | Value (PSO)       |
|------------------------+-------------------|
| "gemv_k128"            | PSO_0x7f8a...     |
| "gemv_k4096"           | PSO_0x7f8b...     |
| "gemm_tm32_tn64_tk16"  | PSO_0x7f8c...     |
| "softmax_n128"         | PSO_0x7f8d...     |
| "rmsnorm_n4096"        | PSO_0x7f8e...     |
| "rope_dim128_base10000"| PSO_0x7f8f...     |
+-------------------------------------------+

The first time you need a PSO with specific constants, it gets compiled and stored. Every subsequent use is a hash table lookup – essentially free. In Rust, the cache lives inside the MetalDevice:

#![allow(unused)]
fn main() {
pub struct MetalDevice {
    device: metal::Device,
    command_queue: metal::CommandQueue,
    library: metal::Library,
    pso_cache: HashMap<String, metal::ComputePipelineState>,
    // ...
}

impl MetalDevice {
    fn get_or_create_pso(
        &mut self,
        kernel_name: &str,
        constants: &[(u32, FunctionConstantValue)],
    ) -> &metal::ComputePipelineState {
        // Build cache key: "kernel_name_const0val_const1val_..."
        let key = build_cache_key(kernel_name, constants);
        
        self.pso_cache.entry(key).or_insert_with(|| {
            let func_constants = metal::FunctionConstantValues::new();
            for (index, value) in constants {
                match value {
                    FunctionConstantValue::UInt(v) => {
                        func_constants.set_constant_value_at_index(
                            v as *const _ as *const _,
                            metal::MTLDataType::UInt,
                            *index as u64,
                        );
                    }
                    // ... other types
                }
            }
            
            let function = self.library
                .get_function(kernel_name, Some(func_constants))
                .expect("kernel not found");
            
            self.device
                .new_compute_pipeline_state_with_function(&function)
                .expect("failed to create PSO")
        })
    }
}
}

The key insight is that for ML inference, the set of PSOs you need is determined entirely by the model architecture. A LLaMA model uses the same kernels with the same dimensions for every token it generates. So in practice, all PSOs get created during model loading, and inference itself never hits a cache miss.

Model Loading Phase:             Inference Phase:
+------------------------+      +------------------------+
| For each layer:        |      | For each token:        |
|   - Create GEMM PSOs   |      |   - Look up PSO (hit)  |
|   - Create norm PSOs   |      |   - Encode dispatch     |
|   - Create attn PSOs   |      |   - Execute on GPU      |
|   (one-time cost)      |      |   (no compilation!)     |
+------------------------+      +------------------------+
     ~100-500ms total                ~0ms PSO overhead

Command Encoding Flow

Now that we have a PSO, let us actually use it. In Metal, you do not call GPU functions directly. Instead, you encode commands into a command buffer, and then submit the entire buffer to the GPU for execution.

The Full Encoding Stack

MTLCommandQueue
       |
       | makeCommandBuffer()
       v
MTLCommandBuffer
       |
       | makeComputeCommandEncoder()
       v
MTLComputeCommandEncoder
       |
       | setComputePipelineState()    -- which kernel to run
       | setBuffer(buffer, offset, index) -- bind data
       | setBytes(ptr, length, index)     -- inline small data
       | dispatchThreadgroups(...)        -- launch!
       |
       | endEncoding()
       v
MTLCommandBuffer
       |
       | commit()      -- send to GPU
       | waitUntilCompleted()  -- block until done
       v
[GPU executes the commands]

Let us walk through this with a concrete example. Say we want to run our add_arrays kernel on two arrays of 1024 floats.

Step 1: Create a Command Buffer

let commandBuffer = commandQueue.makeCommandBuffer()!

A command buffer is a container for GPU commands. You can think of it as a recording of work to be done. Nothing executes until you call commit().

Step 2: Create a Compute Command Encoder

let encoder = commandBuffer.makeComputeCommandEncoder()!

The encoder is how you record compute commands. Metal separates encoding (recording) from execution (running) so that you can build up a batch of work and submit it all at once.

Step 3: Set the Pipeline State

encoder.setComputePipelineState(addArraysPSO)

This tells the GPU which kernel function to run for subsequent dispatch calls.

Step 4: Bind Data

encoder.setBuffer(bufferA, offset: 0, index: 0)  // buffer(0) in MSL
encoder.setBuffer(bufferB, offset: 0, index: 1)  // buffer(1) in MSL
encoder.setBuffer(bufferC, offset: 0, index: 2)  // buffer(2) in MSL

Each setBuffer call binds a MTLBuffer to a buffer index. These indices correspond to the [[buffer(N)]] attributes in your MSL kernel.

Step 5: Dispatch

let gridSize = MTLSize(width: 1024, height: 1, depth: 1)
let threadgroupSize = MTLSize(width: 256, height: 1, depth: 1)
encoder.dispatchThreadgroups(
    MTLSize(width: 1024 / 256, height: 1, depth: 1),  // 4 threadgroups
    threadsPerThreadgroup: threadgroupSize
)

We will cover dispatch geometry in detail in Chapter 9. For now, just know that we are launching 1024 threads organized into groups of 256.

Step 6: End Encoding and Commit

encoder.endEncoding()
commandBuffer.commit()
commandBuffer.waitUntilCompleted()

endEncoding() finalizes the encoder – you cannot add more commands after this. commit() submits the command buffer to the GPU. waitUntilCompleted() blocks the CPU until the GPU finishes (in production you would typically use completion handlers or addCompletedHandler instead of blocking).

Encoding Multiple Dispatches

For ML inference, a single command buffer typically contains many dispatch calls – one per layer operation:

let encoder = commandBuffer.makeComputeCommandEncoder()!

// Layer 1: RMSNorm
encoder.setComputePipelineState(rmsNormPSO)
encoder.setBuffer(hiddenState, offset: 0, index: 0)
encoder.setBuffer(normWeights, offset: 0, index: 1)
encoder.setBuffer(normalizedOutput, offset: 0, index: 2)
encoder.dispatchThreadgroups(normGrid, threadsPerThreadgroup: normThreadgroup)

// Layer 1: Q,K,V projection (GEMM)
encoder.setComputePipelineState(gemmPSO)
encoder.setBuffer(normalizedOutput, offset: 0, index: 0)
encoder.setBuffer(qkvWeights, offset: 0, index: 1)
encoder.setBuffer(qkvOutput, offset: 0, index: 2)
encoder.dispatchThreadgroups(gemmGrid, threadsPerThreadgroup: gemmThreadgroup)

// Layer 1: Attention
encoder.setComputePipelineState(attentionPSO)
// ... more setBuffer calls ...
encoder.dispatchThreadgroups(attnGrid, threadsPerThreadgroup: attnThreadgroup)

// ... and so on for all layers ...

encoder.endEncoding()
commandBuffer.commit()

The GPU executes these dispatches in order (within the same encoder), which gives you implicit synchronization between dependent operations. The output of RMSNorm flows into GEMM which flows into attention – no explicit barriers needed within a single encoder.

The setBytes Optimization

For small amounts of data (under 4KB), Metal provides setBytes() as an alternative to setBuffer(). Instead of allocating a GPU buffer and copying data into it, setBytes() inlines the data directly into the command buffer:

struct GEMMParams {
    var M: UInt32
    var N: UInt32
    var K: UInt32
    var alpha: Float
}

var params = GEMMParams(M: 4096, N: 4096, K: 4096, alpha: 1.0)
encoder.setBytes(&params, length: MemoryLayout<GEMMParams>.size, index: 3)

This avoids the overhead of buffer allocation for small, frequently-changing parameter structs. Akunu uses this pattern extensively – kernel parameters like matrix dimensions, strides, and scalar values are all passed via setBytes().

setBuffer path (for large data):
  CPU Memory --> MTLBuffer (GPU-accessible) --> Kernel reads buffer(N)
  [allocate]    [copy or share]                 [dereference pointer]

setBytes path (for small params):
  CPU Memory --> Embedded in command buffer --> Kernel reads buffer(N)
  [no alloc]    [inline copy, < 4KB]            [same interface!]

From the kernel’s perspective, there is no difference – both appear as [[buffer(N)]] arguments. The optimization is entirely on the CPU side.

How Akunu Wraps the Metal Pipeline

Akunu’s Rust code wraps all of this Metal machinery behind a clean interface. The MetalDevice struct owns the Metal device, command queue, library, and PSO cache. Individual operations (GEMM, softmax, RMSNorm, etc.) are methods on MetalDevice that handle all the encoding details internally.

Here is a simplified view of how a GEMM dispatch looks in akunu:

#![allow(unused)]
fn main() {
impl MetalDevice {
    pub fn gemm(
        &mut self,
        a: &MetalBuffer,    // M x K matrix
        b: &MetalBuffer,    // K x N matrix
        c: &MetalBuffer,    // M x N output matrix
        m: u32,
        n: u32,
        k: u32,
    ) -> Result<()> {
        // 1. Get or create the PSO (cached)
        let cache_key = format!("gemm_tm32_tn64_tk{}", k);
        let pso = self.get_or_create_pso("gemm", &[
            (0, FunctionConstantValue::UInt(k)),
        ]);
        
        // 2. Create a command buffer and encoder
        let command_buffer = self.command_queue.new_command_buffer();
        let encoder = command_buffer.new_compute_command_encoder();
        
        // 3. Set pipeline and bind buffers
        encoder.set_compute_pipeline_state(&pso);
        encoder.set_buffer(0, Some(&a.buffer), 0);
        encoder.set_buffer(1, Some(&b.buffer), 0);
        encoder.set_buffer(2, Some(&c.buffer), 0);
        
        // 4. Set parameters via setBytes
        let params = GemmParams { m, n, k };
        encoder.set_bytes(
            3,
            std::mem::size_of::<GemmParams>() as u64,
            &params as *const _ as *const _,
        );
        
        // 5. Calculate dispatch geometry
        let threadgroup_size = MTLSize::new(128, 1, 1);  // 4 SIMD groups
        let grid_size = MTLSize::new(
            (n + 63) / 64,   // tiles along N
            (m + 31) / 32,   // tiles along M
            1,
        );
        
        // 6. Dispatch and finalize
        encoder.dispatch_threadgroups(grid_size, threadgroup_size);
        encoder.end_encoding();
        command_buffer.commit();
        
        Ok(())
    }
}
}

Notice the pattern: get PSO, create encoder, bind resources, dispatch, finalize. Every operation follows this same template. The differences are in which kernel is used, what buffers are bound, and how the dispatch geometry is calculated.

Batching Multiple Operations

For inference, akunu batches all the operations for a single forward pass into one command buffer:

#![allow(unused)]
fn main() {
impl MetalDevice {
    pub fn forward_pass(
        &mut self,
        model: &Model,
        input: &MetalBuffer,
        output: &MetalBuffer,
    ) -> Result<()> {
        let command_buffer = self.command_queue.new_command_buffer();
        let encoder = command_buffer.new_compute_command_encoder();
        
        for layer in &model.layers {
            // RMSNorm
            self.encode_rmsnorm(&encoder, &layer.norm_weights, ...);
            
            // QKV projection
            self.encode_gemm(&encoder, &layer.qkv_weights, ...);
            
            // Attention
            self.encode_attention(&encoder, ...);
            
            // Feed-forward network
            self.encode_gemm(&encoder, &layer.ff_weights, ...);
        }
        
        encoder.end_encoding();
        command_buffer.commit();
        command_buffer.wait_until_completed();
        
        Ok(())
    }
}
}

This is efficient because:

1. One command buffer: Only one submission to the GPU driver, reducing CPU-side overhead.
2. Sequential execution: Within a single encoder, dispatches execute in order, providing implicit synchronization.
3. Cached PSOs: All pipeline states are already compiled from model loading.
4. No CPU-GPU round trips: The CPU records all the work and submits it in one shot. The GPU executes everything without waiting for the CPU.

Advanced: Pipeline Reflection and Validation

When you create a PSO, you can ask Metal for reflection information – metadata about the kernel’s resource usage:

var reflection: MTLAutoreleasedComputePipelineReflection?
let pso = try device.makeComputePipelineState(
    function: function,
    options: .argumentInfo,
    reflection: &reflection
)

// Now you can inspect resource usage
if let reflection = reflection {
    for arg in reflection.arguments {
        print("Argument: \(arg.name), index: \(arg.index)")
        print("  Type: \(arg.type), access: \(arg.access)")
        print("  Size: \(arg.bufferDataSize)")
    }
}

This is useful for debugging and for building generic dispatch systems that can validate buffer bindings at runtime. If you pass a buffer that is too small or bind the wrong type, the validation layer will catch it.

In debug builds, Metal’s validation layer (enabled via the MTL_DEBUG_LAYER environment variable) performs extensive checking:

# Enable Metal validation
export MTL_DEBUG_LAYER=1

# Enable shader validation (slower but catches more bugs)
export MTL_SHADER_VALIDATION=1

These are invaluable during development. They catch things like:

• Buffer overruns (reading/writing past the end of a buffer)
• Missing buffer bindings
• Threadgroup memory size mismatches
• Resource hazards (reading a buffer that is still being written)

Async Pipeline Creation

PSO creation can be done asynchronously, which is important when you need to create many pipelines at startup:

let group = DispatchGroup()

for kernelConfig in allKernelConfigs {
    group.enter()
    device.makeComputePipelineState(function: function) { pso, error in
        if let pso = pso {
            cache[kernelConfig.key] = pso
        }
        group.leave()
    }
}

group.wait()  // All PSOs ready

This lets the Metal runtime parallelize compilation across CPU cores. For a model with dozens of kernel variants, this can significantly reduce startup time.

Summary

Let us recap the key points:

1. A PSO is a compiled GPU kernel – expensive to create, cheap to use.
2. Compilation is multi-stage: MSL source -> AIR -> metallib -> PSO.
3. Offline compilation (precompiled metallib) eliminates runtime compilation cost.
4. Function constants create specialized kernel variants without code duplication – critical for ML where dimensions are fixed per model.
5. PSO caching (like akunu’s pso_cache_) ensures each kernel variant is compiled only once.
6. Command encoding follows a simple pattern: begin, set pipeline, bind resources, dispatch, end.
7. setBytes is an optimization for small parameter structs (< 4KB).
8. Batch multiple operations into a single command buffer for efficiency.

In the next chapter, we will dive into the Metal Shading Language itself – the language you use to write those kernel functions.

The Metal Shading Language

The Metal Shading Language – MSL for short – is what you write GPU kernels in. If you already know C++, you are going to feel right at home, because MSL is essentially C++14 with some things removed (no exceptions, no virtual functions, no RTTI, no recursive function calls) and a bunch of GPU-specific features added. It is not a weird domain-specific language. It is not a visual graph editor. It is just C++ with extra address spaces, vector types, and thread-coordination primitives bolted on.

In this chapter, we are going to cover the MSL features that matter most for ML kernel development. We will write several complete kernels along the way, starting simple and building up to real patterns used in inference engines.

The Basics: Kernel Functions

A compute kernel in MSL is declared with the kernel keyword (or equivalently, [[kernel]]):

#include <metal_stdlib>
using namespace metal;

kernel void add_arrays(
    device const float* a [[buffer(0)]],
    device const float* b [[buffer(1)]],
    device float* c       [[buffer(2)]],
    uint id               [[thread_position_in_grid]]
) {
    c[id] = a[id] + b[id];
}

Let us break down every piece of this.

• kernel void: This is a GPU entry point. Kernels always return void – they communicate results by writing to buffers.
• device const float* a: A pointer into device memory (more on address spaces below). The [[buffer(0)]] attribute tells Metal which buffer binding index this parameter corresponds to.
• uint id [[thread_position_in_grid]]: A built-in variable that gives this thread its unique ID in the dispatch grid.

That is the most minimal kernel possible: read two values, add them, write the result. Every thread processes one element. A thousand threads process a thousand elements. The GPU runs them all in parallel.

Address Spaces

In regular C++, all pointers live in one big flat address space. On the GPU, memory is divided into distinct address spaces with different performance characteristics and access rules. Every pointer in MSL must be qualified with its address space.

+------------------------------------------------------------------+
|                        GPU Memory Map                            |
+------------------------------------------------------------------+
|                                                                  |
|  device          Large, slow-ish (DRAM / unified memory)         |
|  [read/write]    All threads can access. This is where your      |
|                  weight matrices and activation tensors live.     |
|                  Hundreds of GB/s bandwidth on Apple Silicon.     |
|                                                                  |
+------------------------------------------------------------------+
|                                                                  |
|  constant        Same physical memory as device, but goes        |
|  [read-only]     through the constant cache. Best for small,     |
|                  uniform data read by all threads (like params).  |
|                  Limited to 64KB per argument.                    |
|                                                                  |
+------------------------------------------------------------------+
|                                                                  |
|  threadgroup     Fast on-chip SRAM (like CUDA shared memory).    |
|  [read/write]    Only threads in the same threadgroup can access. |
|                  32KB per threadgroup on Apple Silicon.           |
|                  Much faster than device memory.                  |
|                                                                  |
+------------------------------------------------------------------+
|                                                                  |
|  thread          Registers. Private to each thread.              |
|  [read/write]    Fastest possible access.                        |
|                  Limited by register file size.                   |
|                                                                  |
+------------------------------------------------------------------+

Here is how they look in practice:

kernel void address_space_demo(
    device float* data         [[buffer(0)]],   // device address space
    constant Params& params    [[buffer(1)]],   // constant address space
    threadgroup float* shared  [[threadgroup(0)]], // threadgroup (rarely used this way)
    uint tid [[thread_position_in_grid]]
) {
    // 'thread' address space -- local variable, lives in registers
    float local_val = data[tid];
    
    // Read from constant (cached, good for broadcast reads)
    float scale = params.scale;
    
    // Write to device (main memory)
    data[tid] = local_val * scale;
}

When to Use Each Address Space

The most common mistake newcomers make is putting everything in device. If you have a small struct of parameters (dimensions, strides, scaling factors), put it in constant – the constant cache will broadcast it to all threads efficiently.

Vector and Matrix Types

MSL provides built-in vector types that map directly to the GPU’s SIMD hardware. These are not library abstractions – they are native types that the hardware processes in a single instruction.

Scalar Types

half   h = 1.0h;    // 16-bit float (FP16)
float  f = 1.0f;    // 32-bit float (FP32)
int    i = 42;      // 32-bit signed integer
uint   u = 42u;     // 32-bit unsigned integer
short  s = 42;      // 16-bit signed integer
ushort us = 42u;    // 16-bit unsigned integer
bool   b = true;    // boolean

Vector Types

Vectors come in sizes 2, 3, and 4:

half2  h2 = half2(1.0h, 2.0h);          // 2 x FP16
half4  h4 = half4(1.0h, 2.0h, 3.0h, 4.0h);  // 4 x FP16
float2 f2 = float2(1.0f, 2.0f);         // 2 x FP32
float4 f4 = float4(1.0f, 2.0f, 3.0f, 4.0f); // 4 x FP32
uint2  u2 = uint2(10, 20);              // 2 x uint32

You can swizzle components (rearrange them) using .xyzw or .rgba:

float4 v = float4(1.0, 2.0, 3.0, 4.0);
float2 xy = v.xy;      // (1.0, 2.0)
float  z  = v.z;       // 3.0
float4 rev = v.wzyx;   // (4.0, 3.0, 2.0, 1.0)

Why Vectors Matter for ML

Loading a half4 from memory is a single 64-bit load instead of four separate 16-bit loads. This is huge for memory-bound kernels (which most ML kernels are):

Single loads (bad):                 Vectorized load (good):
  load h[0]  -- 16 bits              load h4  -- 64 bits, one instruction
  load h[1]  -- 16 bits              
  load h[2]  -- 16 bits              4x fewer load instructions
  load h[3]  -- 16 bits              Better memory bus utilization
  4 instructions total                1 instruction total

We will come back to vectorized loads in the performance chapter. For now, just remember: prefer half4 and float4 over scalar loads whenever possible.

Thread Identification

Every thread in a compute dispatch has several built-in identifiers. Understanding which one to use is crucial.

kernel void thread_id_demo(
    uint tid    [[thread_position_in_grid]],
    uint tpg    [[threads_per_grid]],
    uint gid    [[threadgroup_position_in_grid]],
    uint tgpg   [[threadgroups_per_grid]],
    uint lid    [[thread_position_in_threadgroup]],
    uint tptg   [[threads_per_threadgroup]],
    uint simd_lane [[thread_index_in_simdgroup]],
    uint simd_id   [[simdgroup_index_in_threadgroup]]
) {
    // tid  = global thread ID (0, 1, 2, ..., total_threads - 1)
    // gid  = which threadgroup this thread belongs to
    // lid  = local thread ID within the threadgroup (0, 1, ..., tptg - 1)
    // simd_lane = lane within the SIMD group (0 - 31)
    // simd_id   = which SIMD group within the threadgroup
}

Here is how they relate:

Grid (all threads)
+------------------------------------------------------------------+
| Threadgroup 0 (gid=0)     | Threadgroup 1 (gid=1)     | ...     |
| +------------------------+ | +------------------------+ |         |
| | SIMD Group 0 (simd=0)  | | | SIMD Group 0           | |         |
| | lanes 0..31             | | | lanes 0..31             | |         |
| +------------------------+ | +------------------------+ |         |
| | SIMD Group 1 (simd=1)  | | | SIMD Group 1           | |         |
| | lanes 0..31             | | | lanes 0..31             | |         |
| +------------------------+ | +------------------------+ |         |
| | SIMD Group 2 (simd=2)  | | | SIMD Group 2           | |         |
| | lanes 0..31             | | | lanes 0..31             | |         |
| +------------------------+ | +------------------------+ |         |
| | SIMD Group 3 (simd=3)  | | | SIMD Group 3           | |         |
| | lanes 0..31             | | | lanes 0..31             | |         |
| +------------------------+ | +------------------------+ |         |
+------------------------------------------------------------------+

For threadgroup size = 128:
  lid = simd_id * 32 + simd_lane
  tid = gid * 128 + lid

Which ID Should You Use?

• thread_position_in_grid: When each thread processes one element of an array. By far the most common.
• threadgroup_position_in_grid + thread_position_in_threadgroup: When threads cooperate within a threadgroup (reductions, tiled algorithms).
• thread_index_in_simdgroup + simdgroup_index_in_threadgroup: When you are doing SIMD-level tricks (shuffles, reductions, matrix ops).

The half Type: FP16 for ML

The half type is a 16-bit floating-point number. It has:

• 1 sign bit
• 5 exponent bits (range: ~6 x 10^-8 to 65504)
• 10 mantissa bits (~3.3 decimal digits of precision)

For ML inference, half is the workhorse type. Here is why:

1. 2x throughput: Apple GPUs can process twice as many half operations per clock compared to float.
2. 2x memory bandwidth: A half is 2 bytes vs 4 bytes for float. Since most ML kernels are memory-bound, this effectively doubles your throughput.
3. Sufficient precision: Neural network weights and activations rarely need more than 3 digits of precision. FP16 is more than enough for inference.

// FP16 dot product -- 2x the throughput of FP32
kernel void dot_product_fp16(
    device const half4* a [[buffer(0)]],
    device const half4* b [[buffer(1)]],
    device half* result   [[buffer(2)]],
    uint id [[thread_position_in_grid]]
) {
    // Each thread computes dot product of one half4 pair
    // That is 4 multiply-adds in a single vector instruction
    half4 va = a[id];
    half4 vb = b[id];
    result[id] = dot(va, vb);  // va.x*vb.x + va.y*vb.y + ...
}

Mixed Precision

Sometimes you want FP16 for storage and memory transfer but FP32 for accumulation to avoid precision loss:

kernel void gemv_mixed_precision(
    device const half4* W [[buffer(0)]],   // Weights in FP16
    device const half4* x [[buffer(1)]],   // Input in FP16
    device float* y       [[buffer(2)]],   // Output in FP32
    constant uint& K      [[buffer(3)]],
    uint row [[thread_position_in_grid]]
) {
    float sum = 0.0f;  // Accumulate in FP32 for precision
    
    uint k4 = K / 4;   // Process 4 elements at a time via half4
    for (uint i = 0; i < k4; i++) {
        half4 w = W[row * k4 + i];
        half4 v = x[i];
        // Multiply in FP16, accumulate in FP32
        sum += float(w.x) * float(v.x);
        sum += float(w.y) * float(v.y);
        sum += float(w.z) * float(v.z);
        sum += float(w.w) * float(v.w);
    }
    
    y[row] = sum;
}

This is a common pattern in ML inference: load data as FP16 (saving bandwidth), accumulate in FP32 (preserving precision), and write the result back as whatever the downstream consumer needs.

SIMD Group Intrinsics

Now we get to the really fun stuff. SIMD group intrinsics let threads within a 32-thread SIMD group communicate directly through the register file – no shared memory, no barriers, no synchronization overhead.

On Apple GPUs, a SIMD group (also called a warp on NVIDIA or a wavefront on AMD) is always 32 threads. These 32 threads execute in lockstep – the same instruction, at the same time, on adjacent data. SIMD intrinsics let these threads share values with each other.

simd_sum: Reduction Within a SIMD Group

// Sum 32 values (one from each lane) in a single operation
float lane_value = data[tid];
float total = simd_sum(lane_value);
// Now ALL 32 lanes have the same total

Without simd_sum, you would need a tree reduction with shared memory and barriers. With it, the hardware does it for you in a few cycles.

Before simd_sum:          After simd_sum:
Lane 0:  3.0              Lane 0:  sum of all 32
Lane 1:  1.0              Lane 1:  sum of all 32
Lane 2:  4.0              Lane 2:  sum of all 32
...                       ...
Lane 31: 2.0              Lane 31: sum of all 32

simd_max / simd_min: Finding Extremes

float lane_value = data[tid];
float max_val = simd_max(lane_value);  // Max across all 32 lanes
float min_val = simd_min(lane_value);  // Min across all 32 lanes

This is essential for softmax, which needs the maximum value before computing exponentials.

simd_shuffle: Direct Lane-to-Lane Communication

// Read the value from a specific lane
float val = simd_shuffle(my_value, target_lane);

simd_shuffle lets any lane read any other lane’s value. It is like having a crossbar switch between all 32 registers.

simd_shuffle_down: Shift Values Down

float val = simd_shuffle_down(my_value, delta);
// Lane i gets the value from lane (i + delta)

Before simd_shuffle_down(val, 1):    After:
Lane 0: A                           Lane 0: B  (got lane 1's value)
Lane 1: B                           Lane 1: C  (got lane 2's value)
Lane 2: C                           Lane 2: D  (got lane 3's value)
Lane 3: D                           Lane 3: E  (got lane 4's value)
...                                  ...

This is the building block for prefix sums and sequential reductions.

simd_shuffle_xor: Butterfly Pattern

float val = simd_shuffle_xor(my_value, mask);
// Lane i gets the value from lane (i XOR mask)

simd_shuffle_xor(val, 1):    // Swap adjacent pairs
Lane 0 <-> Lane 1
Lane 2 <-> Lane 3
Lane 4 <-> Lane 5
...

simd_shuffle_xor(val, 2):    // Swap pairs of pairs
Lane 0 <-> Lane 2
Lane 1 <-> Lane 3
Lane 4 <-> Lane 6
...

This butterfly pattern is used in parallel reductions and FFTs.

Putting It Together: SIMD-Level Softmax

Here is a practical example – computing softmax over 32 elements using only SIMD intrinsics (no shared memory):

kernel void softmax_simd(
    device const float* input  [[buffer(0)]],
    device float* output       [[buffer(1)]],
    uint lane [[thread_index_in_simdgroup]],
    uint group [[simdgroup_index_in_threadgroup]],
    uint gid [[threadgroup_position_in_grid]]
) {
    // Each SIMD group processes one row of 32 elements
    uint row = gid * (threads_per_threadgroup / 32) + group;
    
    float val = input[row * 32 + lane];
    
    // Step 1: Find max across all 32 lanes
    float max_val = simd_max(val);
    
    // Step 2: Subtract max and exponentiate (for numerical stability)
    float exp_val = exp(val - max_val);
    
    // Step 3: Sum all exponentials
    float sum_exp = simd_sum(exp_val);
    
    // Step 4: Normalize
    output[row * 32 + lane] = exp_val / sum_exp;
}

That is a complete softmax in about 10 lines, with no shared memory and no barriers. Each SIMD group handles one 32-element row entirely within its register file. You cannot write this more concisely or more efficiently.

Threadgroup Memory and Barriers

When you need cooperation between threads in different SIMD groups (but within the same threadgroup), you use threadgroup memory and barriers.

Threadgroup memory is a fast on-chip SRAM shared by all threads in a threadgroup. Think of it as a programmer-managed cache.

kernel void reduction_with_shared(
    device const float* input [[buffer(0)]],
    device float* output      [[buffer(1)]],
    uint tid  [[thread_position_in_grid]],
    uint lid  [[thread_position_in_threadgroup]],
    uint gid  [[threadgroup_position_in_grid]],
    uint tptg [[threads_per_threadgroup]]
) {
    // Declare threadgroup memory
    threadgroup float shared[256];  // One slot per thread
    
    // Each thread loads one element into shared memory
    shared[lid] = input[tid];
    
    // BARRIER: Wait for all threads to finish writing
    threadgroup_barrier(mem_flags::mem_threadgroup);
    
    // Tree reduction in shared memory
    for (uint stride = tptg / 2; stride > 0; stride /= 2) {
        if (lid < stride) {
            shared[lid] += shared[lid + stride];
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }
    
    // Thread 0 writes the result
    if (lid == 0) {
        output[gid] = shared[0];
    }
}

The threadgroup_barrier(mem_flags::mem_threadgroup) call is critical. It does two things:

1. Execution barrier: All threads in the threadgroup must reach this point before any can proceed.
2. Memory fence: All writes to threadgroup memory by threads before the barrier are visible to all threads after the barrier.

Without the barrier, you would read stale data – thread 5 might read shared[10] before thread 10 has written to it. Race conditions on the GPU are just as dangerous as on the CPU, and they are much harder to debug.

When to Use Threadgroup Memory vs SIMD Intrinsics

+---------------------------------------+---------------------------------------+
| SIMD Intrinsics                       | Threadgroup Memory                    |
+---------------------------------------+---------------------------------------+
| 32 threads only (one SIMD group)      | Up to 1024 threads (one threadgroup)  |
| No barriers needed (lockstep)         | Barriers required                     |
| Fastest (register-to-register)        | Fast (on-chip SRAM)                   |
| Simple patterns (reduce, broadcast)   | Arbitrary access patterns             |
| Use FIRST if possible                 | Use when SIMD is not enough           |
+---------------------------------------+---------------------------------------+

The rule of thumb: if you can do it with SIMD intrinsics alone, do that. Fall back to threadgroup memory only when you need cross-SIMD-group communication.

Example: RMSNorm Kernel

Let us write a real kernel used in transformer inference. RMSNorm (Root Mean Square Normalization) is used in LLaMA and many modern models:

RMSNorm(x) = x * weight / sqrt(mean(x^2) + epsilon)

Here is the full kernel:

#include <metal_stdlib>
using namespace metal;

constant uint N [[function_constant(0)]];  // hidden dimension

kernel void rmsnorm(
    device const half* input   [[buffer(0)]],
    device const half* weight  [[buffer(1)]],
    device half* output        [[buffer(2)]],
    constant float& epsilon    [[buffer(3)]],
    uint tid  [[thread_position_in_grid]],
    uint lid  [[thread_position_in_threadgroup]],
    uint simd_lane [[thread_index_in_simdgroup]],
    uint simd_id   [[simdgroup_index_in_threadgroup]]
) {
    // Each threadgroup processes one row of N elements
    // Step 1: Compute sum of squares (each thread handles N/tptg elements)
    uint threads = 256;  // threadgroup size
    float sum_sq = 0.0f;
    
    for (uint i = lid; i < N; i += threads) {
        float val = float(input[i]);
        sum_sq += val * val;
    }
    
    // Step 2: Reduce within SIMD group
    sum_sq = simd_sum(sum_sq);
    
    // Step 3: Reduce across SIMD groups using threadgroup memory
    threadgroup float simd_sums[8];  // max 8 SIMD groups (256/32)
    if (simd_lane == 0) {
        simd_sums[simd_id] = sum_sq;
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);
    
    // First SIMD group reduces the partial sums
    float total = 0.0f;
    if (simd_id == 0 && simd_lane < (threads / 32)) {
        total = simd_sum(simd_sums[simd_lane]);
    }
    
    // Broadcast the RMS scale factor
    threadgroup float rms_scale;
    if (lid == 0) {
        rms_scale = rsqrt(total / float(N) + epsilon);
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);
    
    // Step 4: Apply normalization
    float scale = rms_scale;
    for (uint i = lid; i < N; i += threads) {
        float val = float(input[i]);
        output[i] = half(val * scale * float(weight[i]));
    }
}

This kernel demonstrates several important MSL patterns:

1. Function constants (N) for compile-time specialization.
2. Mixed precision: Load as half, compute in float, store as half.
3. Two-stage reduction: First within SIMD groups (simd_sum), then across SIMD groups via threadgroup memory.
4. Stride loop pattern: Each thread processes multiple elements (for (uint i = lid; i < N; i += threads)).

Example: Vectorized Element-Wise Operations

For simple element-wise operations (add, multiply, activation functions), vectorized loads are the key optimization:

kernel void silu_activation(
    device const half4* input  [[buffer(0)]],
    device half4* output       [[buffer(1)]],
    uint id [[thread_position_in_grid]]
) {
    half4 x = input[id];
    // SiLU(x) = x * sigmoid(x) = x / (1 + exp(-x))
    half4 sigmoid_x = half4(1.0h) / (half4(1.0h) + exp(-x));
    output[id] = x * sigmoid_x;
}

By processing half4 (4 elements per thread), each thread does useful work on 8 bytes of data per load. With a threadgroup size of 256, one threadgroup processes 1024 elements. This keeps the memory system saturated.

Example: Fused RoPE (Rotary Position Embeddings)

RoPE is used in virtually every modern LLM. It applies a rotation to pairs of dimensions based on their position in the sequence:

constant uint HEAD_DIM [[function_constant(0)]];
constant float ROPE_BASE [[function_constant(1)]];

kernel void rope(
    device half* q          [[buffer(0)]],  // query tensor
    device half* k          [[buffer(1)]],  // key tensor
    constant uint& seq_pos  [[buffer(2)]],  // position in sequence
    uint tid [[thread_position_in_grid]]
) {
    // Each thread processes one pair of dimensions
    uint pair_idx = tid;  // Which (cos, sin) pair
    uint d = pair_idx % (HEAD_DIM / 2);
    uint head = pair_idx / (HEAD_DIM / 2);
    
    // Compute rotation angle: theta = pos * base^(-2d/dim)
    float freq = 1.0f / pow(ROPE_BASE, float(2 * d) / float(HEAD_DIM));
    float angle = float(seq_pos) * freq;
    float cos_angle = cos(angle);
    float sin_angle = sin(angle);
    
    // Apply rotation to Q
    uint base_idx = head * HEAD_DIM;
    float q0 = float(q[base_idx + d]);
    float q1 = float(q[base_idx + d + HEAD_DIM / 2]);
    q[base_idx + d]              = half(q0 * cos_angle - q1 * sin_angle);
    q[base_idx + d + HEAD_DIM/2] = half(q0 * sin_angle + q1 * cos_angle);
    
    // Apply same rotation to K
    float k0 = float(k[base_idx + d]);
    float k1 = float(k[base_idx + d + HEAD_DIM / 2]);
    k[base_idx + d]              = half(k0 * cos_angle - k1 * sin_angle);
    k[base_idx + d + HEAD_DIM/2] = half(k0 * sin_angle + k1 * cos_angle);
}

Notice the function constants: HEAD_DIM and ROPE_BASE. These are fixed per model (e.g., 128 and 10000.0 for LLaMA), so the compiler can optimize the frequency computation and potentially pre-compute parts of it.

Built-in Math Functions

MSL provides a comprehensive set of math functions. Here are the ones you will use most often in ML kernels:

// Exponential and logarithm
float e = exp(x);        // e^x
float l = log(x);        // natural log
half  e_fast = exp(h);   // FP16 exp -- 2x throughput

// Trigonometric
float s = sin(x);
float c = cos(x);

// Power and roots
float p = pow(base, exponent);
float s = sqrt(x);
float r = rsqrt(x);      // 1/sqrt(x) -- faster than 1.0/sqrt(x)

// Min, max, clamp
float m = min(a, b);
float M = max(a, b);
float c = clamp(x, lo, hi);  // max(lo, min(x, hi))

// Absolute value
float a = abs(x);

// Fused multiply-add (one rounding instead of two)
float f = fma(a, b, c);  // a*b + c

// Dot product of vectors
float d = dot(v1, v2);   // v1.x*v2.x + v1.y*v2.y + ...

// Type conversion (explicit)
half  h = half(f);        // float -> half (may lose precision)
float f = float(h);       // half -> float (lossless)

The rsqrt function deserves special mention. You will see it everywhere in ML kernels (normalization layers, attention scaling). It computes 1/sqrt(x) in a single instruction, which is faster than computing sqrt(x) and then dividing.

Structs and Parameter Passing

For kernels with many parameters, pack them into a struct:

struct GEMMParams {
    uint M;        // rows of A, rows of C
    uint N;        // cols of B, cols of C
    uint K;        // cols of A, rows of B
    float alpha;   // scaling factor
};

kernel void gemm(
    device const half* A  [[buffer(0)]],
    device const half* B  [[buffer(1)]],
    device half* C        [[buffer(2)]],
    constant GEMMParams& params [[buffer(3)]],
    uint2 gid [[threadgroup_position_in_grid]],
    uint lid  [[thread_position_in_threadgroup]]
) {
    uint M = params.M;
    uint N = params.N;
    uint K = params.K;
    // ...
}

On the CPU side, this struct gets passed via setBytes() (for structs under 4KB) or via a MTLBuffer. The struct layout must match between MSL and the host language. This means you need to be careful about alignment – MSL follows standard C++ alignment rules, and your Rust/Swift struct must match.

#![allow(unused)]
fn main() {
// Rust side -- must match the MSL struct layout
#[repr(C)]
struct GEMMParams {
    m: u32,
    n: u32,
    k: u32,
    alpha: f32,
}
}

The #[repr(C)] attribute ensures Rust uses C-compatible layout, which matches MSL’s struct layout.

Conditional Compilation with Function Constants

We touched on function constants in the previous chapter. Here is a deeper look at how they enable conditional compilation in MSL:

constant bool HAS_BIAS    [[function_constant(0)]];
constant bool HAS_RESIDUAL[[function_constant(1)]];
constant uint BLOCK_SIZE  [[function_constant(2)]];

kernel void linear_layer(
    device const half* input   [[buffer(0)]],
    device const half* weights [[buffer(1)]],
    device half* output        [[buffer(2)]],
    device const half* bias    [[buffer(3)]],
    device const half* residual[[buffer(4)]],
    constant uint& N           [[buffer(5)]],
    uint tid [[thread_position_in_grid]]
) {
    half result = /* ... compute matmul ... */ ;
    
    // These branches are resolved at compile time!
    // No runtime branch penalty.
    if (HAS_BIAS) {
        result += bias[tid];
    }
    
    if (HAS_RESIDUAL) {
        result += residual[tid];
    }
    
    output[tid] = result;
}

When HAS_BIAS is false, the compiler completely removes the bias addition and the bias buffer binding. The generated code is identical to a kernel that never had bias support in the first place. This lets you write one kernel source that generates many specialized variants.

Atomic Operations

Sometimes multiple threads need to update the same memory location. MSL provides atomic operations for this:

kernel void histogram(
    device const uint* data       [[buffer(0)]],
    device atomic_uint* bins      [[buffer(1)]],
    uint tid [[thread_position_in_grid]]
) {
    uint value = data[tid];
    uint bin = value % 256;
    atomic_fetch_add_explicit(&bins[bin], 1, memory_order_relaxed);
}

Atomics are slow – they serialize access. Avoid them in hot paths. If you find yourself using atomics in a performance-critical kernel, there is almost certainly a better algorithm (like per-threadgroup histograms followed by a merge).

Putting It All Together: A Complete Softmax Kernel

Let us write a production-quality softmax kernel that handles arbitrary row sizes using all the MSL features we have covered:

#include <metal_stdlib>
using namespace metal;

constant uint COLS [[function_constant(0)]];

kernel void softmax(
    device const half* input  [[buffer(0)]],
    device half* output       [[buffer(1)]],
    uint gid  [[threadgroup_position_in_grid]],
    uint lid  [[thread_position_in_threadgroup]],
    uint simd_lane [[thread_index_in_simdgroup]],
    uint simd_id   [[simdgroup_index_in_threadgroup]]
) {
    // Each threadgroup processes one row
    uint row = gid;
    uint threads = 256;
    
    // Threadgroup memory for cross-SIMD-group reduction
    threadgroup float simd_max_vals[8];
    threadgroup float simd_sum_vals[8];
    
    // ---- Pass 1: Find max ----
    float local_max = -INFINITY;
    for (uint i = lid; i < COLS; i += threads) {
        local_max = max(local_max, float(input[row * COLS + i]));
    }
    
    // Reduce within SIMD group
    local_max = simd_max(local_max);
    
    // Reduce across SIMD groups
    if (simd_lane == 0) {
        simd_max_vals[simd_id] = local_max;
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);
    
    float row_max;
    if (simd_id == 0) {
        float v = (simd_lane < threads / 32) ? 
                  simd_max_vals[simd_lane] : -INFINITY;
        row_max = simd_max(v);
    }
    
    // Broadcast max to all threads
    threadgroup float shared_max;
    if (lid == 0) shared_max = row_max;
    threadgroup_barrier(mem_flags::mem_threadgroup);
    row_max = shared_max;
    
    // ---- Pass 2: Compute exp and sum ----
    float local_sum = 0.0f;
    for (uint i = lid; i < COLS; i += threads) {
        float val = exp(float(input[row * COLS + i]) - row_max);
        local_sum += val;
    }
    
    // Reduce sum (same two-stage pattern)
    local_sum = simd_sum(local_sum);
    if (simd_lane == 0) {
        simd_sum_vals[simd_id] = local_sum;
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);
    
    float row_sum;
    if (simd_id == 0) {
        float v = (simd_lane < threads / 32) ? 
                  simd_sum_vals[simd_lane] : 0.0f;
        row_sum = simd_sum(v);
    }
    
    threadgroup float shared_sum;
    if (lid == 0) shared_sum = row_sum;
    threadgroup_barrier(mem_flags::mem_threadgroup);
    row_sum = shared_sum;
    
    // ---- Pass 3: Normalize and write output ----
    float inv_sum = 1.0f / row_sum;
    for (uint i = lid; i < COLS; i += threads) {
        float val = exp(float(input[row * COLS + i]) - row_max);
        output[row * COLS + i] = half(val * inv_sum);
    }
}

This kernel has three passes over the data (find max, compute exponentials and sum, normalize). Each pass uses the two-stage reduction pattern (SIMD-level reduction followed by threadgroup-level reduction). The function constant COLS lets the compiler optimize the loop bounds.

Summary

MSL is a practical language for writing GPU kernels. Here are the key takeaways:

1. MSL is C++14 with GPU extensions. No surprises if you know C++.
2. Four address spaces: device (main memory), constant (cached read-only), threadgroup (fast shared SRAM), thread (registers).
3. Vector types (half4, float4) are essential for memory throughput.
4. The half type gives you 2x throughput and 2x bandwidth – use it everywhere in ML kernels.
5. SIMD intrinsics (simd_sum, simd_max, simd_shuffle, etc.) enable fast intra-SIMD-group communication without shared memory or barriers.
6. Threadgroup memory + barriers enable cross-SIMD-group communication within a threadgroup.
7. Function constants create specialized kernel variants at compile time.
8. The two-stage reduction (SIMD reduce, then threadgroup reduce) is the most common pattern in ML kernels.

In the next chapter, we will zoom out from individual threads to look at how threadgroups and SIMD groups are organized and dispatched.

Threadgroups, SIMD Groups, and Dispatch

So far we have talked about individual threads and what they can do. But a single GPU thread is weak – it is only useful because there are thousands of them running simultaneously. The way you organize those thousands of threads matters enormously for performance. Get the geometry wrong and you leave half the GPU idle. Get it right and your kernel saturates every execution unit.

This chapter is about the organizational hierarchy of GPU threads in Metal, and how to dispatch them effectively.

The Thread Hierarchy

Metal organizes threads into a three-level hierarchy:

Grid (all threads in the dispatch)
+------------------------------------------------------------------+
|                                                                  |
|  Threadgroup 0         Threadgroup 1         Threadgroup 2       |
|  +-----------------+   +-----------------+   +-----------------+ |
|  | SIMD Group 0    |   | SIMD Group 0    |   | SIMD Group 0    | |
|  | [32 threads]    |   | [32 threads]    |   | [32 threads]    | |
|  +-----------------+   +-----------------+   +-----------------+ |
|  | SIMD Group 1    |   | SIMD Group 1    |   | SIMD Group 1    | |
|  | [32 threads]    |   | [32 threads]    |   | [32 threads]    | |
|  +-----------------+   +-----------------+   +-----------------+ |
|  | SIMD Group 2    |   | SIMD Group 2    |   | SIMD Group 2    | |
|  | [32 threads]    |   | [32 threads]    |   | [32 threads]    | |
|  +-----------------+   +-----------------+   +-----------------+ |
|  | SIMD Group 3    |   | SIMD Group 3    |   | SIMD Group 3    | |
|  | [32 threads]    |   | [32 threads]    |   | [32 threads]    | |
|  +-----------------+   +-----------------+   +-----------------+ |
|                                                                  |
+------------------------------------------------------------------+

Let us define each level.

SIMD Group (32 threads)

A SIMD group is 32 threads that execute in lockstep – the same instruction, at the same clock cycle. This is the fundamental unit of execution on Apple GPUs. You do not choose the SIMD group size; it is always 32.

Threads in a SIMD group can communicate through SIMD intrinsics (simd_sum, simd_shuffle, etc.) with zero overhead – no barriers, no shared memory, just direct register reads. This is the fastest form of inter-thread communication available.

Threadgroup (up to 1024 threads)

A threadgroup is a collection of SIMD groups that are co-scheduled on the same GPU compute unit. All threads in a threadgroup:

• Share access to threadgroup memory (fast on-chip SRAM)
• Can synchronize with threadgroup_barrier()
• Are guaranteed to execute concurrently (they are all resident on the same compute unit)

A threadgroup can contain up to 1024 threads, which means up to 32 SIMD groups. In practice, common sizes are 128 (4 SIMD groups), 256 (8 SIMD groups), or sometimes 512 or 1024.

Grid (the entire dispatch)

The grid is the collection of all threadgroups in a dispatch. Threadgroups in the grid are independent – there is no way for threads in different threadgroups to communicate during a single dispatch (except through device memory, but there are no cross-threadgroup barriers).

Two Ways to Dispatch: Threadgroups vs Threads

Metal gives you two dispatch methods, and the difference is subtle but important.

dispatchThreadgroups

You specify how many threadgroups to launch and how many threads per threadgroup:

// Launch a grid of threadgroups
let threadgroupsPerGrid = MTLSize(width: 16, height: 8, depth: 1)
let threadsPerThreadgroup = MTLSize(width: 256, height: 1, depth: 1)

encoder.dispatchThreadgroups(threadgroupsPerGrid,
                             threadsPerThreadgroup: threadsPerThreadgroup)
// Total threads = 16 * 8 * 256 = 32,768

With this method, you are responsible for computing the grid dimensions. If your problem size is not a perfect multiple of the threadgroup size, you need to handle the boundary condition in the kernel.

dispatchThreads

You specify the total number of threads you want, and Metal figures out the threadgroup count:

// Launch exactly this many threads
let threadsPerGrid = MTLSize(width: 4000, height: 3000, depth: 1)
let threadsPerThreadgroup = MTLSize(width: 16, height: 16, depth: 1)

encoder.dispatchThreads(threadsPerGrid,
                        threadsPerThreadgroup: threadsPerThreadgroup)
// Metal launches enough threadgroups to cover 4000 x 3000,
// and threads outside that range are automatically disabled

With dispatchThreads, Metal handles the boundary for you. If you request 4000 x 3000 threads with 16x16 threadgroups, Metal launches 250 x 188 threadgroups (25016 = 4000, 18816 = 3008), but the extra 8 rows of threads simply do not execute.

Which Should You Use?

For ML kernels, dispatchThreadgroups is almost always the right choice. Here is why:

1. Explicit control: You know exactly how many threadgroups you are launching, which makes reasoning about resource usage straightforward.
2. Kernel design: ML kernels are typically designed around threadgroup-level cooperation (tiled GEMM, reductions). The kernel code already assumes a specific threadgroup structure.
3. No wasted threads: With dispatchThreadgroups, you design your kernel so that every thread does useful work. There are no “out-of-bounds” threads.

dispatchThreads is nice for simple element-wise kernels where you just want one thread per element and do not care about threadgroup structure. But for anything involving shared memory, reductions, or tiling, use dispatchThreadgroups.

Choosing Threadgroup Sizes

The threadgroup size is one of the most important performance decisions you make. Here are the rules:

Rule 1: Always Use Multiples of 32

Since SIMD groups are 32 threads, your threadgroup size should always be a multiple of 32. If you use, say, 100 threads per threadgroup, Metal will still allocate 4 SIMD groups (128 threads worth of resources), but 28 threads in the last SIMD group will be idle. That is 22% waste.

Threadgroup size = 100 (bad):
+----------------+ +----------------+ +----------------+ +----------------+
| SIMD Group 0   | | SIMD Group 1   | | SIMD Group 2   | | SIMD Group 3   |
| 32 active      | | 32 active      | | 32 active      | | 4 active       |
| 0 idle         | | 0 idle         | | 0 idle         | | 28 IDLE        |
+----------------+ +----------------+ +----------------+ +----------------+
                                                            ^^^^^^^^^^^^
                                                            22% waste!

Threadgroup size = 128 (good):
+----------------+ +----------------+ +----------------+ +----------------+
| SIMD Group 0   | | SIMD Group 1   | | SIMD Group 2   | | SIMD Group 3   |
| 32 active      | | 32 active      | | 32 active      | | 32 active      |
| 0 idle         | | 0 idle         | | 0 idle         | | 0 idle         |
+----------------+ +----------------+ +----------------+ +----------------+
                                0% waste

Rule 2: Common Sizes and When to Use Them

Rule 3: Balance Occupancy and Resources

Here is the tension: larger threadgroups use more resources (registers, threadgroup memory), which means fewer threadgroups can be resident on a compute unit at the same time. Smaller threadgroups use fewer resources but may not have enough threads for efficient cooperation.

Compute Unit (simplified)
+----------------------------------------------------------+
| Register File: 64KB      Threadgroup Memory: 32KB        |
|                                                          |
| Option A: 4 threadgroups of 128 threads (512 threads)    |
|   Each TG gets: 16KB registers, 8KB shared memory        |
|   Good latency hiding (many threads to switch between)   |
|                                                          |
| Option B: 2 threadgroups of 256 threads (512 threads)    |
|   Each TG gets: 32KB registers, 16KB shared memory       |
|   Same total threads, but more resources per TG          |
|                                                          |
| Option C: 1 threadgroup of 1024 threads (1024 threads)   |
|   Gets all: 64KB registers, 32KB shared memory           |
|   Lots of threads, but if this TG stalls, nothing else   |
|   can run on this compute unit                           |
+----------------------------------------------------------+

The sweet spot for ML kernels is usually 128-256 threads. This gives you enough SIMD groups for cooperation (4-8) while leaving room for multiple threadgroups per compute unit (good for latency hiding).

Rule 4: Consult maxTotalThreadsPerThreadgroup

Every PSO has a maxTotalThreadsPerThreadgroup property that tells you the maximum threadgroup size the kernel supports, based on its resource usage:

let pso = try device.makeComputePipelineState(function: function)
print(pso.maxTotalThreadsPerThreadgroup)  // e.g., 1024, or maybe 512

If your kernel uses a lot of registers or threadgroup memory, this number may be less than 1024. Always check.

SIMD Group Cooperation Patterns

Now let us look at how SIMD groups work together to solve problems. These patterns show up over and over in ML kernels.

Pattern 1: SIMD Reduction

The simplest and most common pattern. Each thread has a value, and you want the sum (or max, or min) across all threads in the SIMD group.

float my_value = data[tid];
float total = simd_sum(my_value);
// All 32 lanes now hold the same total

Lane:  0    1    2    3    ...  31
Val:   3.0  1.0  4.0  1.5  ...  2.7
                    |
              simd_sum
                    |
             total = 87.3 (in all lanes)

Pattern 2: SIMD Broadcast

One lane has a value that all other lanes need.

float value;
if (simd_lane == 0) {
    value = compute_something_expensive();
}
// Broadcast lane 0's value to all lanes
value = simd_broadcast_first(value);
// Or from a specific lane:
value = simd_shuffle(value, source_lane);

Pattern 3: SIMD Prefix Sum (Inclusive Scan)

Each lane gets the sum of all values from lane 0 through its own lane. Useful for compaction, histograms, and stream processing.

float val = data[tid];
float prefix = simd_prefix_inclusive_sum(val);

Lane:    0    1    2    3    4    5    ...
Input:   3    1    4    1    5    9    ...
Output:  3    4    8    9   14   23    ...
         ^    ^    ^
         |    |    +-- 3+1+4
         |    +-- 3+1
         +-- 3

Pattern 4: SIMD Shuffle for Data Reuse

When adjacent threads need overlapping data (like in convolution), shuffles avoid redundant memory loads:

// Each lane loads one element
float my_val = data[base + simd_lane];

// Now lane i can access neighboring values without memory loads:
float left  = simd_shuffle_up(my_val, 1);   // lane i gets lane (i-1)'s value
float right = simd_shuffle_down(my_val, 1); // lane i gets lane (i+1)'s value

// Stencil computation (e.g., 1D convolution with kernel [0.25, 0.5, 0.25])
float result = 0.25f * left + 0.5f * my_val + 0.25f * right;

Before shuffle:
Lane:   0    1    2    3    4    ...
Val:    A    B    C    D    E    ...

After simd_shuffle_down(val, 1):
Lane:   0    1    2    3    4    ...
Val:    B    C    D    E    F    ...
        ^    ^    ^
        Each lane got the next lane's value

Pattern 5: Two-Stage Reduction (Cross-SIMD-Group)

When you need a reduction across the entire threadgroup (not just one SIMD group), you do it in two stages:

Stage 1: SIMD-level reduction (fast, no barriers)
+------------------+  +------------------+  +------------------+  +------------------+
| SIMD Group 0     |  | SIMD Group 1     |  | SIMD Group 2     |  | SIMD Group 3     |
| 32 values        |  | 32 values        |  | 32 values        |  | 32 values        |
| --> simd_sum     |  | --> simd_sum     |  | --> simd_sum     |  | --> simd_sum     |
| Result: S0       |  | Result: S1       |  | Result: S2       |  | Result: S3       |
+------------------+  +------------------+  +------------------+  +------------------+

Stage 2: Write partial sums to threadgroup memory, then reduce
+-------------------------------------------+
| threadgroup float partials[4];            |
| partials = [S0, S1, S2, S3]              |
|                                           |
| threadgroup_barrier(...)                  |
|                                           |
| SIMD Group 0 reads all 4 values          |
| --> simd_sum(partials[simd_lane])         |
| Result: S0 + S1 + S2 + S3 = TOTAL        |
+-------------------------------------------+

Here is the code:

// Assume threadgroup size = 128 (4 SIMD groups)
threadgroup float partials[4];

// Stage 1: Each SIMD group reduces its 32 values
float local_sum = simd_sum(my_value);

// One thread per SIMD group writes to shared memory
if (simd_lane == 0) {
    partials[simd_id] = local_sum;
}
threadgroup_barrier(mem_flags::mem_threadgroup);

// Stage 2: First SIMD group reduces the partial sums
float total;
if (simd_id == 0 && simd_lane < 4) {
    total = simd_sum(partials[simd_lane]);
}

// Broadcast to all threads if needed
threadgroup float shared_total;
if (lid == 0) {
    shared_total = total;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
total = shared_total;

This two-stage pattern is the workhorse of ML kernels. You will find it in normalization (RMSNorm, LayerNorm), softmax, and anywhere else you need a global reduction within a threadgroup.

Thread Mapping: From Grid to Data

One of the most important design decisions is how you map threads to data. There are several common patterns.

Pattern A: One Thread Per Element

The simplest mapping. Each thread processes exactly one element:

kernel void scale(
    device half* data [[buffer(0)]],
    constant float& factor [[buffer(1)]],
    uint tid [[thread_position_in_grid]]
) {
    data[tid] = half(float(data[tid]) * factor);
}

Thread:  0    1    2    3    4    5    6    7  ...
Data:   [0]  [1]  [2]  [3]  [4]  [5]  [6]  [7] ...
         |    |    |    |    |    |    |    |
         v    v    v    v    v    v    v    v

Dispatch: one thread per element, trivial.

Pattern B: One Thread Per Vector (Vectorized)

Each thread processes a vector of elements:

kernel void scale_vec4(
    device half4* data [[buffer(0)]],
    constant float& factor [[buffer(1)]],
    uint tid [[thread_position_in_grid]]
) {
    half4 v = data[tid];
    data[tid] = half4(float4(v) * factor);
}

Thread:  0         1         2         3      ...
Data:   [0,1,2,3] [4,5,6,7] [8,9,10,11] [12,13,14,15] ...
         |         |         |            |
         v         v         v            v

Dispatch: one thread per 4 elements. Total threads = N / 4.

Pattern C: Stride Loop (One Threadgroup Per Row)

Each threadgroup processes an entire row, with threads striding through the elements:

kernel void row_sum(
    device const float* matrix [[buffer(0)]],
    device float* sums [[buffer(1)]],
    constant uint& cols [[buffer(2)]],
    uint gid [[threadgroup_position_in_grid]],
    uint lid [[thread_position_in_threadgroup]]
) {
    uint row = gid;
    float sum = 0.0f;
    
    // Each thread strides through the row
    for (uint col = lid; col < cols; col += 256) {
        sum += matrix[row * cols + col];
    }
    
    // Reduce within threadgroup...
    sum = simd_sum(sum);
    // ... (two-stage reduction as before)
    
    if (lid == 0) {
        sums[row] = sum;
    }
}

Row 0 (threadgroup 0):
Thread 0 reads:  [0], [256], [512], [768], ...
Thread 1 reads:  [1], [257], [513], [769], ...
Thread 2 reads:  [2], [258], [514], [770], ...
...
Thread 255 reads: [255], [511], [767], [1023], ...

Dispatch: one threadgroup per row.

Pattern D: 2D Tiling (GEMM)

For matrix multiplication, threads are organized in a 2D grid of tiles:

kernel void gemm_tiled(
    device const half* A [[buffer(0)]],
    device const half* B [[buffer(1)]],
    device half* C       [[buffer(2)]],
    constant uint& M     [[buffer(3)]],
    constant uint& N     [[buffer(4)]],
    constant uint& K     [[buffer(5)]],
    uint2 gid [[threadgroup_position_in_grid]],
    uint  lid [[thread_position_in_threadgroup]],
    uint  simd_id [[simdgroup_index_in_threadgroup]]
) {
    // Each threadgroup computes a TM x TN tile of C
    // gid.y = tile row, gid.x = tile column
    uint tile_row = gid.y;  // Which TM-row block
    uint tile_col = gid.x;  // Which TN-column block
    
    // Each SIMD group within the threadgroup computes a sub-tile
    // ...
}

Matrix C (M x N):
+--------+--------+--------+--------+
| TG(0,0)| TG(0,1)| TG(0,2)| TG(0,3)|  <-- threadgroup row 0
| 32x64  | 32x64  | 32x64  | 32x64  |
+--------+--------+--------+--------+
| TG(1,0)| TG(1,1)| TG(1,2)| TG(1,3)|  <-- threadgroup row 1
| 32x64  | 32x64  | 32x64  | 32x64  |
+--------+--------+--------+--------+
| TG(2,0)| TG(2,1)| TG(2,2)| TG(2,3)|
| 32x64  | 32x64  | 32x64  | 32x64  |
+--------+--------+--------+--------+

Each threadgroup computes a TM x TN tile of C (e.g., 32 x 64).
Grid dimensions: (N/TN, M/TM, 1) threadgroups.

This is the tiling pattern used in akunu’s GEMM kernels. We will cover it in detail in the SIMD matrix operations chapter.

Dispatch Geometry Calculation

Let us put this together with concrete dispatch calculations for common ML operations.

Element-wise Operations (ReLU, SiLU, Add)

Problem: Apply activation to N elements
Approach: One thread per element (or per half4)

Vectorized (half4):
  threads_needed = N / 4
  threadgroup_size = 256
  threadgroups = ceil(threads_needed / 256)

Example: N = 4096
  threads_needed = 1024
  threadgroup_size = 256
  threadgroups = 4
  Grid: MTLSize(4, 1, 1)
  TG:   MTLSize(256, 1, 1)

Row-wise Operations (Softmax, RMSNorm)

Problem: Process M rows of N elements each
Approach: One threadgroup per row

  threadgroup_size = 256
  threadgroups = M

Example: M = 32 (batch), N = 4096 (hidden dim)
  threadgroups = 32
  Grid: MTLSize(32, 1, 1)
  TG:   MTLSize(256, 1, 1)
  
  Each of the 256 threads in a threadgroup strides through
  4096 / 256 = 16 elements.

Matrix Multiplication (GEMM)

Problem: C[M,N] = A[M,K] * B[K,N]
Approach: 2D grid, each threadgroup computes a TM x TN tile

  TM = 32, TN = 64 (tile sizes)
  threadgroup_size = 128 (4 SIMD groups)
  threadgroups_x = ceil(N / TN)
  threadgroups_y = ceil(M / TM)

Example: M = 4096, N = 4096, K = 4096
  threadgroups_x = 4096 / 64 = 64
  threadgroups_y = 4096 / 32 = 128
  Grid: MTLSize(64, 128, 1)
  TG:   MTLSize(128, 1, 1)
  Total threadgroups: 8192
  Total threads: 8192 * 128 = 1,048,576

Batched Operations

For batched operations, you can use the third grid dimension:

Problem: Batch of B matrices, each M x N, apply softmax per row
Approach: 3D grid -- batch x rows x 1

  Grid: MTLSize(M, B, 1)
  TG:   MTLSize(256, 1, 1)

Kernel sees:
  batch_idx = threadgroup_position_in_grid.y
  row_idx   = threadgroup_position_in_grid.x

Real-World Example: Dispatch for Multi-Head Attention

Let us trace through the dispatch geometry for multi-head attention, a critical component of transformer inference.

Suppose we have:

• Batch size B = 1 (single sequence, typical for inference)
• Number of heads H = 32
• Sequence length S = 2048
• Head dimension D = 128

The attention computation involves:

1. Q * K^T -> scores [H, 1, S] (for single-token generation)
2. softmax(scores / sqrt(D)) -> weights [H, 1, S]
3. weights * V -> output [H, 1, D]

Step 1: Score computation (GEMV -- each head is a vector-matrix multiply)
  For each head: q[1, D] * K[S, D]^T -> scores[1, S]
  
  Grid: MTLSize(ceil(S/256), H, 1)   -- one TG row per head
  TG:   MTLSize(256, 1, 1)
  
  Each threadgroup computes 256 elements of the score vector
  for one attention head.

Step 2: Softmax over scores
  For each head: softmax(scores[1, S])
  
  Grid: MTLSize(H, 1, 1)   -- one TG per head
  TG:   MTLSize(256, 1, 1)
  
  Each threadgroup processes one row (one head's scores).

Step 3: Weighted sum (another GEMV)
  For each head: weights[1, S] * V[S, D] -> output[1, D]
  
  Grid: MTLSize(ceil(D/32), H, 1)
  TG:   MTLSize(256, 1, 1)

In practice, akunu fuses some of these steps together to reduce memory traffic, but the dispatch geometry follows this general pattern.

Common Pitfalls

Pitfall 1: Threadgroup Size Not a Multiple of 32

// BAD: wastes 22% of GPU resources
encoder.dispatchThreadgroups(grid, threadsPerThreadgroup: MTLSize(100, 1, 1))

// GOOD: full utilization
encoder.dispatchThreadgroups(grid, threadsPerThreadgroup: MTLSize(128, 1, 1))

Pitfall 2: Too Few Threadgroups

If you dispatch only 2 threadgroups on a GPU with 10 compute units, 8 compute units sit idle. You want at least as many threadgroups as compute units, and ideally several times more for latency hiding.

M2 Ultra: 76 compute units
Minimum threadgroups for full utilization: 76
Better: 76 * 4 = 304+ threadgroups (multiple waves)

Pitfall 3: Forgetting Bounds Checks

When using dispatchThreadgroups, your total thread count may exceed your data size. Always check bounds:

kernel void safe_scale(
    device half* data [[buffer(0)]],
    constant uint& count [[buffer(1)]],
    uint tid [[thread_position_in_grid]]
) {
    if (tid >= count) return;  // Bounds check!
    data[tid] = data[tid] * 2.0h;
}

Pitfall 4: Divergent Execution Within a SIMD Group

All 32 threads in a SIMD group execute together. If some threads take one branch and others take a different branch, the SIMD group executes both branches (masking inactive threads). This is called divergence and it wastes cycles:

// BAD: divergent branches within a SIMD group
if (tid % 2 == 0) {
    // Even threads do expensive path A
    result = expensive_computation_A(data[tid]);
} else {
    // Odd threads do expensive path B
    result = expensive_computation_B(data[tid]);
}
// Both paths execute for every SIMD group -- 2x the cost!

// BETTER: organize work so entire SIMD groups take the same branch
uint group = tid / 32;
if (group % 2 == 0) {
    // Entire SIMD groups do path A
    result = expensive_computation_A(data[tid]);
} else {
    // Entire SIMD groups do path B
    result = expensive_computation_B(data[tid]);
}

Pitfall 5: Not Accounting for Non-Uniform Threadgroups

When using dispatchThreads, the last threadgroup in each dimension may be smaller than requested. If your kernel assumes a fixed threadgroup size (e.g., uses threadgroup float shared[256]), it will still work, but some of those shared memory slots will contain garbage. Be careful when reducing – only reduce over the actual active thread count.

Visualizing Execution on Apple Silicon

Let us trace how a dispatch actually executes on an Apple Silicon GPU.

An M1 Pro has 16 compute units. Each compute unit can execute multiple threadgroups concurrently (depending on resource usage). Here is a simplified timeline:

Dispatch: 64 threadgroups, 128 threads each (4 SIMD groups per TG)

Compute Unit 0:  [TG 0][TG 16][TG 32][TG 48]  (4 TGs across time)
Compute Unit 1:  [TG 1][TG 17][TG 33][TG 49]
Compute Unit 2:  [TG 2][TG 18][TG 34][TG 50]
...
Compute Unit 15: [TG 15][TG 31][TG 47][TG 63]

Wave 0: TG 0-15  (one TG per compute unit)
Wave 1: TG 16-31 (launched as wave 0 TGs finish)
Wave 2: TG 32-47
Wave 3: TG 48-63

Total: 4 waves to process all 64 threadgroups

In reality, the scheduler is more dynamic than this – it can have multiple threadgroups resident per compute unit if there are enough resources. But the wave model gives you the right intuition. More threadgroups = more waves = more opportunity for the GPU to hide latency.

Choosing Grid Dimensions for ML Workloads

Here is a decision flowchart for common ML operations:

What kind of operation?
    |
    +-- Element-wise (ReLU, add, scale)?
    |     Threads per element: 1 (or 1 per half4 for vectorization)
    |     Grid: (ceil(N/4/TG_SIZE), 1, 1)
    |     TG: (256, 1, 1)
    |
    +-- Row-wise reduction (softmax, norm)?
    |     One threadgroup per row
    |     Grid: (num_rows, 1, 1)
    |     TG: (256, 1, 1)
    |
    +-- Matrix multiply (GEMM)?
    |     2D grid of tiles
    |     Grid: (ceil(N/TN), ceil(M/TM), 1)
    |     TG: (128, 1, 1)   -- 4 SIMD groups
    |
    +-- Batched operation?
    |     Use 3rd dimension for batch
    |     Grid: (spatial_x, spatial_y, batch)
    |     TG: depends on inner operation
    |
    +-- Vector-matrix multiply (GEMV)?
          Grid: (ceil(M/TG_SIZE), 1, 1)   -- one TG per chunk of rows
          TG: (256, 1, 1)

Summary

1. Thread hierarchy: Grid -> Threadgroups -> SIMD Groups -> Threads. Each level has different communication capabilities.
2. SIMD groups (32 threads) communicate via intrinsics – fast, no barriers.
3. Threadgroups (up to 1024 threads) communicate via threadgroup memory + barriers.
4. Grid threadgroups cannot communicate during a dispatch.
5. Always use multiples of 32 for threadgroup sizes.
6. Common sizes: 128 for GEMM, 256 for reductions, 32 for trivial ops.
7. dispatchThreadgroups for kernels with threadgroup cooperation (most ML kernels).
8. dispatchThreads for simple per-element kernels.
9. Two-stage reduction (SIMD reduce + threadgroup reduce) is the fundamental pattern.
10. Dispatch enough threadgroups to keep all compute units busy (at least as many as there are compute units).

The next chapter covers the memory model – how data flows between CPU and GPU, and why memory bandwidth is the critical bottleneck for ML inference.

Memory Model and Buffer Management

If you have ever profiled a GPU workload and been disappointed by the utilization numbers, the root cause was almost certainly memory. Not compute. Memory. On Apple Silicon, the story is both simpler and subtler than on discrete GPUs: the CPU and GPU share a single physical memory pool (Unified Memory Architecture, or UMA), which eliminates an entire class of problems – but introduces a few new ones. This chapter walks through Metal’s memory model from first principles, shows how akunu maps it to inference, and explains why bandwidth – not FLOPS – is the number you should be worrying about.

The Unified Memory Architecture

Traditional GPU programming on NVIDIA or AMD involves two physically separate memory pools. The CPU has its DDR/LPDDR; the GPU has its GDDR/HBM. Every time you want the GPU to see CPU data, you perform an explicit copy across the PCIe bus (16 GB/s for PCIe 4.0 x16). Every time you want CPU results, you copy back. These copies are slow, asynchronous, and a constant source of bugs.¹

Apple Silicon threw this model out. Starting with M1, the CPU, GPU, and Neural Engine all share a single pool of LPDDR memory with a single set of page tables.² There is no PCIe bus. There is no copy. When the CPU writes to address 0x1234, the GPU can read from that same address – because it is the same physical page.

┌──────────────────────────────────────────────────┐
│               Unified Memory (LPDDR5/5X)         │
│  ┌─────────┐  ┌─────────┐  ┌─────────┐          │
│  │  CPU     │  │  GPU    │  │  Neural │          │
│  │  Cores   │  │  Cores  │  │  Engine │          │
│  └────┬─────┘  └────┬────┘  └────┬────┘          │
│       │             │            │                │
│       └─────────────┼────────────┘                │
│                     │                             │
│           ┌─────────┴─────────┐                   │
│           │  System Level     │                   │
│           │  Cache (SLC)      │                   │
│           └─────────┬─────────┘                   │
│                     │                             │
│           ┌─────────┴─────────┐                   │
│           │  Fabric / Memory  │                   │
│           │  Controller       │                   │
│           └───────────────────┘                   │
└──────────────────────────────────────────────────┘

This has profound implications for inference engines:

1. Zero-copy buffer sharing. A MTLBuffer allocated with MTLResourceStorageModeShared is readable and writable by both CPU and GPU without any explicit transfer. The CPU gets a raw pointer (buf.contents); the GPU gets the same backing pages.

2. No staging buffers needed. On CUDA, you allocate a “host-pinned” staging buffer, memcpy into it, then launch a cuMemcpyHtoD. On Metal with shared mode, you just write directly and dispatch.

3. Coherency is automatic (mostly). After GPU work completes (i.e., waitUntilCompleted returns), the CPU sees all GPU writes immediately. The hardware’s cache coherency protocol handles this.³

Metal Storage Modes

Metal offers three storage modes for buffers. Understanding them is essential because choosing wrong means either unnecessary copies or corrupted data.

On Apple Silicon, Shared mode is the overwhelmingly correct choice. It provides direct CPU and GPU access with no copies. Private mode can be slightly faster for buffers the CPU never touches (the GPU may have more freedom in cache management), but in practice the difference is negligible for inference workloads – and you lose the ability to read results without a blit encoder.⁴

Managed mode exists for macOS systems with discrete AMD GPUs and is irrelevant to Apple Silicon. If you see Managed in Metal sample code, it is targeting a different hardware class.

Why Akunu Uses Shared Mode Everywhere

Look at how MetalDevice::allocate works in akunu:

Buffer MetalDevice::allocate(size_t bytes) {
    id<MTLBuffer> buf = [STATE.device newBufferWithLength:MAX(bytes, 16)
                                                 options:MTLResourceStorageModeShared];
    void *h = (void *)CFBridgingRetain(buf);
    allocated_bytes_ += MAX(bytes, 16);
    return {h, bytes, [buf contents]};
}

Every allocation uses MTLResourceStorageModeShared. The returned Buffer struct stores both the opaque Metal handle (h) and the CPU-accessible pointer ([buf contents]). This means:

• Weight loading is zero-copy. When akunu parses a GGUF file, it can mmap the file and allocate Metal buffers that point at the same pages. The GPU reads weights directly from the memory-mapped file. No memcpy, no staging buffer, no DMA transfer.

• KV cache is CPU-readable. After decode completes, the CPU can inspect KV cache contents directly through buf.contents – useful for debugging and speculative decoding verification.

• Scratch buffers just work. The ScratchBuffers struct allocates everything up front with device.allocate(). During the hot path, nothing is allocated or freed.

The Buffer Struct: Akunu’s Abstraction

Akunu wraps Metal buffers in a simple POD struct defined in device.h:

struct Buffer {
    void *handle;    // Backend-specific pointer (MTLBuffer*, CUdeviceptr, etc.)
    size_t size;     // Size in bytes
    void *contents;  // CPU-accessible pointer (for UMA or mapped memory)
};

This is deliberately minimal. No reference counting, no smart pointers, no virtual methods. Just three fields that fit in 24 bytes. The handle is an opaque pointer – on Metal it is a CFBridgingRetain’d id<MTLBuffer>, on a future CUDA backend it would be a CUdeviceptr. The contents pointer is the CPU-visible address; on discrete GPUs it would be nullptr.

Why no reference counting? Because buffer lifetimes in inference are trivially static. You allocate all buffers at model load time, use them for the entire session, and free them at shutdown. There is no dynamic buffer creation in the hot path. This means the overhead of shared_ptr or arc reference counting is pure waste.⁵

The allocate-with-data Overload

There is a second allocate overload that accepts initial data:

Buffer MetalDevice::allocate(const void *data, size_t bytes) {
    id<MTLBuffer> buf = [STATE.device newBufferWithBytes:data
                                                  length:bytes
                                                 options:MTLResourceStorageModeShared];
    void *h = (void *)CFBridgingRetain(buf);
    allocated_bytes_ += bytes;
    return {h, bytes, [buf contents]};
}

This is used for two things: uploading initial weight data and creating pre-filled parameter buffers. Metal’s newBufferWithBytes: internally memcpys the data into the newly allocated buffer. On UMA this is a simple memcpy; there is no bus transfer.

The setBytes() Optimization: Inline Small Data

Not everything needs a buffer. Metal provides setBytes:length:atIndex: for small, frequently-changing data. Instead of allocating a buffer, writing to it, and binding it, you just hand the encoder a pointer and a length. Metal copies the bytes directly into the command buffer’s inline data area.⁶

The rules are:

Akunu uses this aggressively. Every kernel’s parameter struct (dimensions, epsilon values, strides) is under 64 bytes. The DispatchCmd struct stores these inline:

// Inline params (up to 64 bytes -- covers all kernel param structs)
uint8_t param_bytes[64];
int param_size;
int param_index;  // buffer index for setBytes/setBuffer

During dispatch table replay, static params use pre-allocated GPU buffers via setBuffer() (zero per-token work), while position-patched params use setBytes() (the patched bytes get copied inline into the command buffer). This split is the key insight from akunu’s Metal backend – setBytes() is perfect for the few parameters that change per token, while setBuffer() avoids redundant work for the many parameters that do not.

// From MetalDevice::encode_dispatch_table:
// Static params: use pre-allocated GPU buffer (no per-token work)
if (cmd.param_buf.handle)
    [enc setBuffer:(__bridge id<MTLBuffer>)cmd.param_buf.handle
            offset:0
           atIndex:(NSUInteger)cmd.param_index];
else
    [enc setBytes:cmd.param_bytes
           length:(NSUInteger)cmd.param_size
          atIndex:(NSUInteger)cmd.param_index];

Memory Alignment

Metal requires 16-byte alignment for buffer offsets when binding with setBuffer:offset:atIndex:.⁷ This is a hardware constraint – the GPU’s memory controller fetches aligned 16-byte chunks, and unaligned offsets would require splitting a fetch across two cache lines.

Akunu ensures alignment in several ways:

1. Buffer sizes are rounded up. The minimum allocation is MAX(bytes, 16), ensuring every buffer is at least one alignment unit.

2. QKV sub-offsets are naturally aligned. The ScratchBuffers struct computes QKV offsets as q_dim * 2 and (q_dim + kv_dim) * 2. Since q_dim and kv_dim are always multiples of head_dim (typically 64 or 128), and each element is 2 bytes (FP16), these offsets are always multiples of 128 or 256 bytes – far exceeding the 16-byte requirement.

3. GGUF block alignment. Quantized formats like Q4_0 use 32-element blocks (18 bytes each). Weight tensors are stored as contiguous arrays of these blocks, and GGUF pads to alignment boundaries.

Resource Hazards and Synchronization

On discrete GPUs with separate memory, you have explicit copy commands that create natural synchronization points. On UMA with shared mode, the GPU and CPU can both touch the same memory at any time. This creates resource hazards: what happens if the CPU writes to a buffer while the GPU is reading it?

Metal handles this through command buffer completion:

1. CPU-to-GPU: The CPU writes data before calling begin_encoding(). The encoder captures buffer references at encode time. As long as writes complete before encoding begins, the GPU will see them.

2. GPU-to-CPU: After end_encoding_sync() returns (or waitUntilCompleted signals), all GPU writes are visible to the CPU. The hardware flushes caches as part of command buffer completion.

3. GPU-to-GPU (within a command buffer): Metal guarantees that compute dispatches within a single command encoder execute in order. Dispatch N sees all writes from dispatch N-1. No barriers needed.⁸

4. GPU-to-GPU (across command buffers): You need explicit synchronization. Akunu’s end_encoding_event() / begin_encoding_after_event() uses MTLSharedEvent for GPU-to-GPU signaling across command buffers:

// Signal after this command buffer completes
[STATE.cmdBuffer encodeSignalEvent:STATE.pipelineEvent value:signal_val];
[STATE.cmdBuffer commit];

// Next command buffer waits for the signal
[STATE.cmdBuffer encodeWaitForEvent:STATE.pipelineEvent value:STATE.eventValue];

This is critical for pipelined chain decode, where one command buffer is executing on the GPU while the CPU encodes the next one.

The Bandwidth Bottleneck

Here is the uncomfortable truth about LLM inference on Apple Silicon: you will almost never be compute-bound during token generation. You will be memory-bandwidth-bound.

Why? Consider a single decode step for a 7B parameter model with Q4_0 quantization. Each parameter is 4.5 bits on average (4 bits of data plus amortized scale/min). The total weight data is roughly:

$$7 \times 10^9 \times 4.5 / 8 \approx 3.9 \text{ GB}$$

Every decode step reads every weight exactly once (each layer’s Q, K, V, O, gate, up, down projections). That is 3.9 GB of memory reads per token. On an M4 Pro with ~273 GB/s memory bandwidth, the theoretical floor is:

$$3.9 \text{ GB} / 273 \text{ GB/s} \approx 14.3 \text{ ms/token} \approx 70 \text{ tok/s}$$

The actual compute work (multiply-accumulate operations) takes a fraction of this time. The GPU cores are waiting for data to arrive from DRAM, not crunching numbers.⁹

This table reveals a fundamental reality: your token generation speed is almost entirely determined by how fast your chip can feed data to the GPU cores. More compute cores help for prefill (which is compute-bound), but for autoregressive decode, bandwidth is king.¹⁰

Implications for Engine Design

This bandwidth bottleneck drives several akunu design decisions:

1. Minimize weight reads. Read each weight exactly once per token. Fuse operations where possible (SiLU + down projection, QKV projection) to reduce the number of kernel launches and avoid re-reading intermediate buffers.

2. Use the System Level Cache (SLC). Apple Silicon chips have a large last-level cache shared between CPU and GPU. On M4 Pro, it is estimated at 32 MB. Weights that fit in SLC are served at cache bandwidth (~2 TB/s on M4), not DRAM bandwidth. This is why akunu fuses QKV and gate+up weights on chips with large SLC – the fused weight matrix is more likely to be in cache for the second read.¹¹

3. Quantize aggressively. Q4_0 reads half the bytes of FP16. Q2_K reads a quarter. Lower precision means less bandwidth, which directly translates to higher tok/s.

4. Avoid unnecessary reads. The dispatch table pre-computes everything that can be pre-computed. Buffer bindings, pipeline states, parameter structs – all resolved at build time, not at dispatch time.

Pre-Allocated Scratch Buffers

Akunu allocates all intermediate buffers once at model load time via the ScratchBuffers struct:

struct ScratchBuffers {
    Buffer h0;         // [dim] FP16
    Buffer h1;         // [dim] FP16
    Buffer residual;   // [dim] FP16
    Buffer qkv;        // [q_dim + 2*kv_dim] FP16
    Buffer attn_out;   // [max(q_dim, dim)] FP16
    Buffer ffn_gate;   // [2 * ffn_dim] FP16 (2x for fused gate+up)
    Buffer ffn_up;     // [ffn_dim] FP16
    Buffer ffn_act;    // [ffn_dim] FP16
    Buffer logits;     // [vocab_size] FP16
    Buffer token_ids;  // [max_chain] U32
    // ... plus batch buffers for prefill
};

A few things to notice:

• ffn_gate is 2x ffn_dim. This accommodates fused gate+up projections, where the output of a single GEMV contains both the gate and up vectors contiguously.

• qkv is contiguous. Q, K, and V are stored in a single buffer at computed offsets (qkv_q_offset, qkv_k_offset, qkv_v_offset). When QKV fusion is enabled, a single GEMV writes all three projections into this buffer. When not fused, three separate GEMVs write to their respective sub-regions.

• No dynamic allocation. The ScratchBuffers::create factory is called once. After that, every forward pass reuses the same buffers. This means zero memory allocation overhead in the hot path, zero fragmentation, and deterministic memory usage.

• Prefill buffers are separate. Batch buffers (batch_h0, batch_q, etc.) are sized for the maximum prefill chunk. They are larger than decode buffers by a factor of prefill_chunk (typically 4096).

KV Cache Memory Layout

The KV cache is another critical memory structure, defined in kv_cache.h:

struct KVCache {
    int n_layers;
    int n_kv_heads;
    int head_dim;
    int max_length;
    int current_length;
    std::vector<Buffer> k_buffers;  // one per layer
    std::vector<Buffer> v_buffers;  // one per layer
    int kv_stride;  // max_length * head_dim
};

Each layer gets two buffers: one for K, one for V. The layout is [n_kv_heads, max_length, head_dim] in FP16. The kv_stride (= max_length * head_dim) is the number of elements between consecutive KV heads.

For a 32-layer model with 8 KV heads, 128 head dim, and 4096 max context:

$$\text{KV bytes per layer} = 8 \times 4096 \times 128 \times 2 = 8 \text{ MB}$$ $$\text{Total KV cache} = 32 \times 2 \times 8 = 512 \text{ MB}$$

This is significant – half a gigabyte just for KV cache on a model that “only” has 7B parameters. For 70B models with 80 layers, the KV cache can easily exceed 4 GB. This is why max_context is capped at 4096 by default in akunu_load_model().

Buffer Memory Layout (Shared Mode — UMA)

  ┌──────────────────────────────────────────────┐
  │  Model Weights (3.9 GB)                      │  GPU reads, CPU writes at load
  ├──────────────────────────────────────────────┤
  │  KV Cache (512 MB)                           │  GPU reads+writes per token
  ├──────────────────────────────────────────────┤
  │  Scratch Buffers (~10 MB)                    │  GPU reads+writes, reused every step
  ├──────────────────────────────────────────────┤
  │  Prefill Batch Buffers (~200 MB)             │  Used only during prefill
  └──────────────────────────────────────────────┘

Write Buffer and Read Buffer: Portability

The Device base class provides default implementations of write_buffer and read_buffer that use plain memcpy:

virtual void write_buffer(Buffer dst, const void *src, size_t bytes, size_t offset = 0) {
    if (dst.contents)
        memcpy((char *)dst.contents + offset, src, bytes);
}

virtual void read_buffer(void *dst, Buffer src, size_t bytes, size_t offset = 0) {
    if (src.contents)
        memcpy(dst, (const char *)src.contents + offset, bytes);
}

On UMA, this is literally a memcpy – there is no DMA, no bus, no asynchronous transfer. The comment says “Override for discrete GPU backends (CUDA cuMemcpyHtoD, etc.).” This is the portability escape hatch: a future CUDA backend would override these methods to use proper device-to-host copies.

Memory Tracking

MetalDevice tracks total bytes allocated:

size_t allocated_bytes_ = 0;
// In allocate():
allocated_bytes_ += MAX(bytes, 16);
// In free_buffer():
if (buf.size <= allocated_bytes_) allocated_bytes_ -= buf.size;

This is exposed through akunu_model_memory() in the C API, letting callers see how much GPU memory a loaded model uses. It is a simple counter, not a full memory allocator – because akunu does not need a full memory allocator. Buffers are allocated at init, freed at shutdown, and nothing happens in between.

Common Pitfalls

Before we leave the memory model, let me highlight a few traps that catch even experienced Metal developers:

Pitfall 1: Reading GPU Results Before Completion

// WRONG: GPU may still be writing
device.end_encoding_async();
float *result = (float *)logits_buf.contents;  // Race condition!

Always call wait() or end_encoding_sync() before reading GPU-written buffers from the CPU. The UMA does not mean coherent-at-all-times; it means coherent-after-completion.

Pitfall 2: Buffer Lifetime with Unretained References

Akunu uses commandBufferWithUnretainedReferences for performance. This means Metal will NOT retain buffer references – if you free a buffer before the GPU finishes, you get a GPU fault. This is safe in akunu because all buffers outlive the GPU work (they are freed only at model destruction), but adding dynamic buffer management would require switching to commandBuffer (with retain) or careful lifetime tracking.¹²

Pitfall 3: Assuming Private Mode is Faster

On Apple Silicon, Private mode offers minimal benefit over Shared mode for inference workloads. The GPU’s cache hierarchy works the same way for both. Private prevents CPU access, which can be slightly more efficient for GPU-only temporaries, but the difference is typically <1% for bandwidth-bound kernels. Akunu chose simplicity over micro-optimization here.

Pitfall 4: Ignoring the 16-Byte Alignment Rule

If you bind a buffer with an offset that is not a multiple of 16, Metal will either silently produce garbage or crash with a validation error (if Metal validation is enabled). The fix is always the same: pad your data structures to 16-byte boundaries. Notice how akunu’s param structs include _p0, _p1 padding fields:

struct { uint32_t dim; float eps; uint32_t _p0, _p1; } norm_params;

Those _p0, _p1 fields pad the struct to 16 bytes, ensuring alignment when passed through setBytes().

Summary

Apple Silicon’s UMA simplifies GPU programming enormously: no copies, no staging buffers, no DMA. But it does not eliminate the fundamental bottleneck of memory bandwidth. For LLM inference, where every token requires reading the entire weight matrix, bandwidth determines throughput.

Akunu’s memory strategy is:

• Shared mode everywhere for zero-copy CPU-GPU access
• Static allocation of all buffers at model load time
• Pre-allocated param buffers with the setBytes/setBuffer split for per-token patching
• 16-byte aligned everything
• Bandwidth-aware design driving quantization, weight fusion, and SLC exploitation

In the next chapter, we will look at the compute side: how Apple’s SIMD group matrix operations turn those bandwidth-fed bytes into actual matrix multiplications.

1. NVIDIA CUDA Programming Guide, Section 3.2.2, “Device Memory.” The PCIe bus bottleneck is well-documented; PCIe 4.0 x16 provides ~32 GB/s bidirectional, but only ~16 GB/s in one direction. See https://docs.nvidia.com/cuda/cuda-c-programming-guide/. ↩

2. Apple, “Apple M1 Chip,” 2020. The Unified Memory Architecture was first introduced with M1, providing up to 68.25 GB/s bandwidth with a single memory pool. See https://www.apple.com/newsroom/2020/11/apple-unleashes-m1/. ↩

3. Apple, “Metal Best Practices Guide: Resource Storage Modes,” 2024. On Apple Silicon, shared mode buffers are coherent after command buffer completion. See https://developer.apple.com/documentation/metal/choosing-a-resource-storage-mode-for-apple-gpus. ↩

4. Apple, “Metal Feature Set Tables.” On Apple Silicon (Apple GPU family 7+), MTLResourceStorageModeShared is the recommended mode for most buffers. Private mode is useful for render targets and textures that the CPU never accesses. See https://developer.apple.com/metal/Metal-Feature-Set-Tables.pdf. ↩

5. This is a deliberate design decision. llama.cpp uses a similar approach with ggml_backend_buffer, where buffer lifetimes are tied to the model context. Reference counting would add atomic operations to every buffer bind – thousands per forward pass. ↩

6. Apple, “setBytes:length:atIndex: Documentation.” The data is copied inline into the command buffer. The maximum size is 4 KB. For larger data, use a buffer. See https://developer.apple.com/documentation/metal/mtlcomputecommandencoder. ↩

7. Apple, “Metal Best Practices Guide: Buffer Alignment.” Buffer offsets must be a multiple of 16 bytes for setBuffer:offset:atIndex:. See https://developer.apple.com/documentation/metal/mtlcomputecommandencoder. ↩

8. Apple, “Metal Programming Guide: Command Organization.” Within a single compute command encoder, dispatches execute in order with implicit memory barriers. See https://developer.apple.com/documentation/metal/performing_calculations_on_a_gpu. ↩

9. This analysis follows the roofline model methodology. See Williams, S., Waterman, A., & Patterson, D. (2009). “Roofline: An Insightful Visual Performance Model for Multicore Architectures.” Communications of the ACM, 52(4), 65-76. See https://doi.org/10.1145/1498765.1498785. ↩

10. Memory bandwidth numbers sourced from Apple’s product specifications and independent testing by Anandtech and Chips and Cheese. Actual achieved bandwidth is typically 70-85% of theoretical peak due to memory controller overhead. See https://chipsandcheese.com/. ↩

11. The System Level Cache (SLC) is described in various Apple Silicon die analyses. See “Apple M1 Die Shot Analysis” by Chips and Cheese, 2021. SLC sizes are estimated from die area analysis and performance profiling; Apple does not officially disclose exact SLC sizes. See https://www.apple.com/newsroom/2020/11/apple-unleashes-m1/. ↩

12. Apple, “commandBufferWithUnretainedReferences Documentation.” Using unretained references avoids the overhead of Metal retaining every buffer referenced by a command buffer. This is safe when buffer lifetimes are manually managed. See https://developer.apple.com/documentation/metal/mtlcommandqueue/1508684-makecommandbufferwithunretainedr. ↩

SIMD Group Matrix Operations

Matrix multiplication is the beating heart of every transformer. During a single forward pass of a 7B-parameter model, you perform roughly 30 matrix multiplications per layer across 32 layers – nearly a thousand matmuls per token. If your matmul is slow, everything is slow. On Apple Silicon, the key to fast matmuls is the simdgroup_matrix API: a set of hardware-accelerated operations that let 32 GPU threads cooperatively multiply 8x8 matrix tiles at full throughput.¹

This chapter explains how SIMD group matrix operations work, how they compose into larger GEMM/GEMV kernels, and how akunu uses them for both prefill and decode.

What Is a SIMD Group?

Before diving into matrix operations, let’s establish the execution model. An Apple GPU organizes threads into SIMD groups (also called warps on NVIDIA, or wavefronts on AMD). A SIMD group on Apple Silicon is always 32 threads that execute in lockstep – every thread runs the same instruction at the same cycle.²

┌─────────────────────────────────────────────┐
│              Threadgroup (128 threads)       │
│  ┌─────────┐ ┌─────────┐ ┌─────────┐ ┌───────────┐  │
│  │ SIMD 0  │ │ SIMD 1  │ │ SIMD 2  │ │ SIMD 3    │  │
│  │ 32 thds │ │ 32 thds │ │ 32 thds │ │ 32 thds   │  │
│  └─────────┘ └─────────┘ └─────────┘ └───────────┘  │
└─────────────────────────────────────────────┘

Within a SIMD group, threads can communicate through SIMD shuffle instructions – reading values from any other thread in the group without going through memory. This is extremely fast (single-cycle) and is the foundation of SIMD group matrix operations.

The simdgroup_matrix API

Starting with Apple GPU family 7 (M1 and later), Metal Shading Language provides the simdgroup_matrix type and associated operations.³ The fundamental unit is an 8x8 matrix distributed across the 32 threads of a SIMD group:

#include <metal_simdgroup_matrix>

simdgroup_half8x8   mat_a;   // 8x8 matrix of float16
simdgroup_float8x8  mat_b;   // 8x8 matrix of float32

Each thread in the SIMD group holds a portion of the 8x8 matrix. You do not control which thread holds which element – the hardware distributes the 64 elements across 32 threads (2 elements per thread) in a hardware-defined layout. This layout is opaque; you interact with the matrix through three operations:

load

simdgroup_half8x8 ma;
simdgroup_load(ma, src_ptr, stride);

Load an 8x8 tile from device or threadgroup memory. The stride is the number of elements between consecutive rows. All 32 threads participate in the load cooperatively – each thread reads its assigned 2 elements.

store

simdgroup_store(mc, dst_ptr, stride);

Store an 8x8 tile back to memory. Same cooperative pattern as load.

multiply_accumulate

simdgroup_multiply_accumulate(mc, ma, mb, mc);
// mc += ma * mb   (8x8 += 8x8 * 8x8)

This is the core operation. It multiplies two 8x8 matrices and accumulates the result into a third. The hardware performs 512 multiply-accumulate operations (8 * 8 * 8) in a single instruction across the SIMD group.⁴ On Apple GPU family 7+, this executes in a few cycles – dramatically faster than doing the same work with scalar or vector instructions.

How 32 Threads Hold One 8x8 Matrix

Let’s be precise about the data distribution. An 8x8 matrix has 64 elements. A SIMD group has 32 threads. So each thread holds exactly 2 elements. The exact mapping is hardware-defined and opaque, but conceptually:

32 Threads Holding an 8x8 Matrix (each thread holds 2 elements):

     col0  col1  col2  col3  col4  col5  col6  col7
    ┌─────┬─────┬─────┬─────┬─────┬─────┬─────┬─────┐
r0  │ T0  │ T0  │ T1  │ T1  │ T2  │ T2  │ T3  │ T3  │
    ├─────┼─────┼─────┼─────┼─────┼─────┼─────┼─────┤
r1  │ T4  │ T4  │ T5  │ T5  │ T6  │ T6  │ T7  │ T7  │
    ├─────┼─────┼─────┼─────┼─────┼─────┼─────┼─────┤
r2  │ T8  │ T8  │ T9  │ T9  │ T10 │ T10 │ T11 │ T11 │
    ├─────┼─────┼─────┼─────┼─────┼─────┼─────┼─────┤
r3  │ T12 │ T12 │ T13 │ T13 │ T14 │ T14 │ T15 │ T15 │
    ├─────┼─────┼─────┼─────┼─────┼─────┼─────┼─────┤
r4  │ T16 │ T16 │ T17 │ T17 │ T18 │ T18 │ T19 │ T19 │
    ├─────┼─────┼─────┼─────┼─────┼─────┼─────┼─────┤
r5  │ T20 │ T20 │ T21 │ T21 │ T22 │ T22 │ T23 │ T23 │
    ├─────┼─────┼─────┼─────┼─────┼─────┼─────┼─────┤
r6  │ T24 │ T24 │ T25 │ T25 │ T26 │ T26 │ T27 │ T27 │
    ├─────┼─────┼─────┼─────┼─────┼─────┼─────┼─────┤
r7  │ T28 │ T28 │ T29 │ T29 │ T30 │ T30 │ T31 │ T31 │
    └─────┴─────┴─────┴─────┴─────┴─────┴─────┴─────┘

The key insight: you never access individual elements of a simdgroup_matrix. You load, you multiply-accumulate, you store. The hardware handles the internal layout. Trying to extract element [3][5] would require SIMD shuffles and defeat the purpose.

Building Larger Matmuls from 8x8 Tiles

An 8x8 tile is too small for real work. A typical LLM linear layer might be [4096, 4096] – that is 512 x 512 tiles. The strategy is to:

1. Assign tiles to threadgroups. Each threadgroup computes a rectangular block of the output matrix.
2. Within a threadgroup, each SIMD group accumulates a column of output tiles.
3. Loop over the K dimension in steps of 8 (or larger), accumulating partial sums.

Akunu’s GEMM kernels follow the llama.cpp tile geometry, which uses TM=32 and TN=64:

Output tile per threadgroup: 32 rows (M) x 64 columns (N)
Threadgroup size: 128 threads = 4 SIMD groups
Each SIMD group handles: 32 rows x 16 columns (using 8x8 sub-tiles)

K-loop: stride 32 elements (NK=32)
  - Load A tile [32 x 32] into threadgroup memory
  - Load B tile [64 x 32] into threadgroup memory
  - Each SIMD group does: 4 rows x 2 columns x 4 K-steps of 8x8 MACs

Here is how this maps to akunu’s simd_gemm_f16.metal:

// Dispatch: grid=(ceil(N/64), ceil(M/32), 1), threads=(32,4,1)
// TG memory: 4096 + 2048 = 6144 bytes

constexpr short NR0 = 64;   // N tile (weight rows)
constexpr short NR1 = 32;   // M tile (activation rows)
constexpr short NK  = 32;   // K tile (reduction dimension)

simdgroup_half8x8 mc[8];    // 8 accumulator tiles per SIMD group
for (short i = 0; i < 8; i++) {
    mc[i] = make_filled_simdgroup_matrix<half, 8>(0.h);
}

for (uint loop_k = 0; loop_k < K; loop_k += NK) {
    // Load A and B tiles into threadgroup memory
    // ...
    threadgroup_barrier(mem_flags::mem_threadgroup);

    // Multiply: 8x8 tiles along K
    for (short k_step = 0; k_step < NK; k_step += 8) {
        simdgroup_half8x8 ma, mb;
        simdgroup_load(ma, sa + ..., ...);
        for (short tile = 0; tile < 8; tile++) {
            simdgroup_load(mb, sb + ..., ...);
            simdgroup_multiply_accumulate(mc[tile], ma, mb, mc[tile]);
        }
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);
}

// Store accumulated results
for (short tile = 0; tile < 8; tile++) {
    simdgroup_store(mc[tile], C + ..., ldc);
}

The accumulator array mc[8] holds 8 output sub-tiles (8x8 each = 8 * 64 = 512 elements), which covers a 32x16 or 8x64 region depending on the layout. The K-loop processes 32 elements at a time, with each step performing 8x8 MACs.

The Tiling Hierarchy

Let’s trace how a large matmul decomposes:

Full matrix:  C[M,N] = A[M,K] * B^T[N,K]
Example:      C[4096, 4096] = A[4096, 4096] * B^T[4096, 4096]

Grid level:   ceil(4096/64) x ceil(4096/32) = 64 x 128 threadgroups
TG level:     Each TG computes C[32, 64] using 128 threads (4 SIMDs)
SIMD level:   Each SIMD accumulates C[8, 64] or similar sub-region
Tile level:   Each 8x8 MAC produces C[8,8] += A[8,8] * B[8,8]

Akunu’s GEMM Dispatch Geometry

Looking at simd_gemm_f16.metal, the dispatch parameters are:

grid    = (ceil(N/64), ceil(M/32), 1)
threads = (32, 4, 1)   // 128 threads = 4 SIMD groups
TG memory = 6144 bytes  // 4096 for A tile + 2048 for B tile

The threadgroup of 128 threads breaks down as:

• 4 SIMD groups (sgitg = 0..3)
• Each SIMD group processes a different row band of the M-tile
• All 4 SIMD groups share the same B-tile data (loaded cooperatively into threadgroup memory)

For the K-loop:

• NL0 = NK/16 = 2: each thread loads 2 chunks of B data
• NL1 = NK/8 = 4: each thread loads 4 chunks of A data
• Data flows: device memory -> threadgroup memory -> SIMD registers -> accumulate -> device memory

Small-M Variant: simd_gemm_small_f16

During decode (M=1 to ~8), the full GEMM is wasteful – most of the 32-row M-tile is empty. Akunu provides simd_gemm_small_f16 which is optimized for M <= 8:

// From dtype_descriptor.h:
return (M >= 2 && M <= 8) ? d.gemm_small_kernel : d.gemm_kernel;

The small variant uses a reduced M-tile size and fewer SIMD groups, trading parallelism for reduced overhead. For M=1 (single-token decode), akunu uses GEMV kernels instead of GEMM entirely – the dispatch geometry is fundamentally different.

GEMV: The Decode Workhorse

During autoregressive decode, M=1. This is a matrix-vector product, not a matrix-matrix product. The access pattern changes dramatically:

• GEMM (prefill): Each weight element is reused across M activation rows. Arithmetic intensity is O(M). With M=4096, you get excellent compute utilization.
• GEMV (decode): Each weight element is used exactly once. Arithmetic intensity is O(1). You are completely bandwidth-bound.⁵

Akunu’s GEMV kernels do not use simdgroup_matrix at all – they use vectorized loads and SIMD reduction instead. The GEMV kernel for Q4_0 quantized weights, for example, dequantizes blocks of 32 weights, multiplies them by the input vector, and reduces across the SIMD group using simd_sum.

The DTypeDescriptor table maps each quantization format to its GEMV kernel and dispatch geometry:

The “wide” variants use larger threadgroups (256 threads = 8 SIMD groups) to increase occupancy on Pro/Max chips with many GPU cores. The ChipConfig controls when to switch:

c.wide_gemv_threshold = 32768;  // use wide GEMV when N exceeds this

The Accumulation Flow: From Tiles to Output

                    8x8 Tile Accumulation Loop

   K=0..7       K=8..15      K=16..23              Result
  ┌──────┐     ┌──────┐     ┌──────┐              ┌──────┐
  │A[8x8]│     │A[8x8]│     │A[8x8]│              │      │
  └──┬───┘     └──┬───┘     └──┬───┘              │C[8x8]│
     ×            ×            ×          ...      │      │
  ┌──┴───┐     ┌──┴───┐     ┌──┴───┐              │Accum │
  │B[8x8]│     │B[8x8]│     │B[8x8]│              └──────┘
  └──┬───┘     └──┬───┘     └──┬───┘
     +            +            +  ──────────────►

The accumulation loop for a single output tile works as follows:

1. Initialize mc = 0 (zero-filled 8x8 accumulator)
2. For each K-step (stride 8 or 32):
   • Load an 8x8 slice of A into ma
   • Load an 8x8 slice of B into mb
   • mc += ma * mb (hardware multiply-accumulate)
3. After K-loop completes, store mc to output

For K=4096, this loop runs 4096/8 = 512 times (or 4096/32 = 128 times if the K-stride is 32 with 4 sub-steps). Each iteration performs 512 FMA operations. Total: 512 * 512 = 262,144 FMA operations per 8x8 output tile, which is exactly 8 * 8 * 4096 – the correct number.

Quantized GEMM: Dequantize in Registers

For quantized weights (Q4_0, Q4_K, Q8_0, etc.), the GEMM kernels follow the same tiling strategy but add a dequantization step. The weight data is loaded in its quantized format and dequantized into threadgroup memory before the SIMD matrix operations:

For each K-step:
  1. Load quantized weight block from device memory
  2. Dequantize to FP16 in threadgroup memory
  3. Load A tile into threadgroup memory (already FP16)
  4. barrier()
  5. SIMD matrix multiply-accumulate from threadgroup memory
  6. barrier()

Akunu has separate GEMM kernels for each quantization format – simd_gemm_q4_0, simd_gemm_q4_k, simd_gemm_q8_0, etc. – because the dequantization logic is tightly coupled with the load pattern. A Q4_0 block is 18 bytes (16 four-bit values + 1 FP16 scale), while a Q4_K super-block is 144 bytes with nested sub-scales. The memory access pattern and threadgroup memory layout differ for each.

MLX Quantized Formats

Akunu also supports MLX SafeTensors quantized models. These use a different quantization format (affine quantization with configurable group size and bit width) and require function constants to specialize the kernel at pipeline creation time:

// From emit_gemv() in table_builder.cpp:
if (dt.is_mlx) {
    uint32_t fc_indices[] = {0, 1};
    uint32_t fc_values[] = {(uint32_t)quant_group_size, (uint32_t)K};
    pso = device.get_pipeline(kernel_name, cache_key, fc_indices, fc_values, 2);
}

Function constants are Metal’s mechanism for compile-time specialization. Instead of branching on group_size inside the kernel, you bake the value into the pipeline state object. The compiler can then optimize the kernel for that specific group size – unrolling loops, eliminating dead branches, and computing strides at compile time.⁶

Comparison to NVIDIA Tensor Cores

It is instructive to compare Apple’s SIMD group matrix operations to NVIDIA’s Tensor Cores:

The key architectural difference: NVIDIA Tensor Cores are dedicated hardware units separate from the CUDA cores. Apple’s SIMD matrix operations are executed by the same ALUs that do regular floating-point math – they are an instruction set extension, not a separate unit.⁷ This means:

1. No mode switching. You can freely interleave matrix operations with scalar/vector code. On NVIDIA, switching between Tensor Core and CUDA core work can cause pipeline bubbles.

2. Lower peak throughput per core. Apple’s matrix multiply throughput is lower than NVIDIA’s dedicated Tensor Cores. But Apple compensates with higher memory bandwidth per FLOP (critical for inference) and the zero-copy UMA advantage.

3. Simpler programming model. The simdgroup_matrix API is genuinely easier to use than NVIDIA’s WMMA or inline PTX MMA instructions. Load, store, multiply-accumulate – that’s it.

BF16 Support on M4

Starting with Apple GPU family 9 (M4), Metal supports native bfloat (BF16) as a first-class type.⁸ Akunu has dedicated BF16 kernels:

// From dtype_descriptor.h:
{31, "gemv_bf16", nullptr, nullptr,
 "simd_gemm_bf16", "simd_gemm_small_bf16",
 "embedding_lookup_bf16", ...}

BF16 has the same exponent range as FP32 (8 bits) but with only 7 bits of mantissa (vs FP16’s 10 bits). This makes it better for accumulation-heavy workloads where dynamic range matters more than precision. The simdgroup_matrix API supports simdgroup_bfloat8x8 on M4, enabling hardware-accelerated BF16 matrix multiply.

Fused SiLU + GEMV

One of akunu’s more aggressive optimizations is the fused SiLU+down GEMV kernel. In a standard SwiGLU FFN block:

gate = GEMV(gate_weight, x)
up   = GEMV(up_weight, x)
act  = SiLU(gate) * up
down = GEMV(down_weight, act)

That’s 3 GEMV dispatches + 1 activation dispatch = 4 kernel launches. The fused kernel combines the last two steps:

down[i] = sum_j( SiLU(gate[j]) * up[j] * down_weight[i,j] )

This reads the gate and up vectors, applies SiLU element-wise, and immediately multiplies by the down weight – all in one kernel, one pass over the down weight matrix. The DTypeDescriptor tracks which formats have fused kernels:

const char *fused_silu_kernel;        // "gemv_q4_0_silu", etc.
const char *fused_silu_large_kernel;  // Wide variant for Pro+

Not all formats have fused kernels (Q4_K does not, for example), so the table builder falls back to separate activation + GEMV when unavailable.

Practical Performance Numbers

To put all of this in perspective, here is what these kernel choices mean for actual throughput. On an M4 Pro (20 GPU cores, 273 GB/s bandwidth):

The GEMV kernels dominate decode time because they read the most data. Everything else – norms, activations, attention – is comparatively cheap. This is why akunu spends so much effort on GEMV kernel variants, wide vs. standard selection, and fused operations.

Summary

Apple Silicon’s simdgroup_matrix API provides hardware-accelerated 8x8 matrix multiplication that serves as the building block for larger GEMM operations. Akunu uses the llama.cpp tile geometry (TM=32, TN=64) for prefill GEMM and specialized GEMV kernels for decode. The choice between GEMM and GEMV, and between standard and wide variants, is driven by the DTypeDescriptor table and ChipConfig thresholds – all resolved at dispatch table build time, not at runtime.

The next chapter explores the broader set of performance optimization patterns that make these kernels – and the overall inference pipeline – fast.

1. Apple, “Metal Shading Language Specification,” Version 3.2, Section 6.9, “SIMD-group Matrix Functions.” Available at https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf. ↩

2. Apple, “Metal Best Practices Guide: SIMD-groups.” A SIMD group on Apple GPU is always 32 threads. This is analogous to NVIDIA’s warp size. See https://developer.apple.com/documentation/metal/compute_passes/creating_threads_and_threadgroups. ↩

3. The simdgroup_matrix API was introduced in Metal 2.4 (iOS 15, macOS 12) and requires Apple GPU family 7 or later (M1+). Earlier Apple GPUs (A-series before A14) have smaller SIMD groups (8 or 16 threads) and do not support matrix operations. See https://developer.apple.com/metal/Metal-Feature-Set-Tables.pdf. ↩

4. The actual hardware implementation likely uses a systolic array or similar structure within each GPU execution unit. Apple does not disclose microarchitectural details, but the instruction behavior (32 threads cooperatively computing an 8x8x8 MAC) is consistent with a small matrix engine per execution unit. ↩

5. This is the fundamental insight behind the roofline model applied to LLM inference. For M=1 GEMV, the arithmetic intensity is approximately 1 FLOPs/byte (2 FLOPs per weight element, ~2 bytes per weight for FP16). The roofline crossover point on Apple Silicon is typically at M=16-64, depending on the chip. See Williams et al., “Roofline: An Insightful Visual Performance Model for Multicore Architectures,” CACM 2009. See https://doi.org/10.1145/1498765.1498785. ↩

6. Apple, “Using Function Specialization to Build Pipeline Variants.” Function constants allow creating specialized versions of a shader function, enabling the compiler to optimize based on known constant values. See https://developer.apple.com/documentation/metal/using-function-specialization-to-build-pipeline-variants. ↩

7. Dougall Johnson, “Apple GPU Architecture” (reverse-engineering documentation), 2022-2024. Johnson’s analysis shows that Apple GPU ALUs execute SIMD matrix instructions directly, without dedicated tensor hardware. See https://dougallj.github.io/applegpu/. ↩

8. Apple, “What’s New in Metal,” WWDC24. M4 introduces native bfloat16 support in Metal, including simdgroup_bfloat8x8 matrix operations. See https://developer.apple.com/videos/play/wwdc2024/10220/. ↩

Performance Optimization Patterns

We have covered the memory model and the matrix operations. Now we need to talk about how to make them fast. This chapter is a catalog of GPU performance patterns – the techniques that separate a naive kernel from one that saturates the hardware. Each pattern is explained in general terms, then grounded in how akunu applies it. If you have ever stared at a GPU profiler and wondered why your kernel is at 30% utilization, this chapter is for you.

The Roofline Model: Know Your Bottleneck

Before optimizing anything, you need to know what you are optimizing for. The roofline model gives you a simple framework: every kernel is either compute-bound or memory-bandwidth-bound.¹

Performance (GFLOPS)
    ^
    |                  ┌─────────────── Peak Compute
    |                 /
    |                /
    |               / <-- Roofline
    |              /
    |             /
    |            /
    |           /
    |──────────/───────────────────── Peak Bandwidth
    |         /
    |        /
    |       /
    └──────┴──────────────────────────>
           Arithmetic Intensity (FLOPS/byte)

Arithmetic intensity = (total FLOPs) / (total bytes transferred). It tells you how much compute work you do for each byte you read from memory.

• Low arithmetic intensity (e.g., GEMV with M=1): ~1 FLOP/byte. You read a weight, multiply once, move on. You are bandwidth-bound. Adding more compute units does not help.
• High arithmetic intensity (e.g., GEMM with M=512): ~512 FLOPS/byte. You read a weight and reuse it 512 times. You are compute-bound. More bandwidth does not help.

For Apple Silicon, the crossover point depends on the chip:

This means: if your kernel’s arithmetic intensity is below ~32, it is bandwidth-bound on most Apple Silicon. Every single GEMV in the decode path (arithmetic intensity ~1-2) is bandwidth-bound. Only prefill GEMM and attention (with long sequences) are compute-bound.

What This Means for Akunu

The roofline model dictates akunu’s optimization priorities:

1. For decode (bandwidth-bound): Minimize bytes read. Use quantization. Fuse operations to avoid re-reading intermediates. Optimize for memory access patterns.
2. For prefill (compute-bound): Maximize compute utilization. Use SIMD matrix operations. Maximize occupancy. Optimize threadgroup sizes.
3. For attention (mixed): Short sequences are bandwidth-bound (small KV cache); long sequences become compute-bound.

Pattern 1: Vectorized Loads

GPU memory controllers deliver data in large chunks. On Apple Silicon, the memory bus is 128 or 256 bits wide. Reading a single float16 (2 bytes) wastes most of that bus width. Instead, you want to read 4 or 8 elements at once using vector types:

// Slow: scalar loads (2 bytes each)
half val0 = input[tid * 4 + 0];
half val1 = input[tid * 4 + 1];
half val2 = input[tid * 4 + 2];
half val3 = input[tid * 4 + 3];

// Fast: vectorized load (8 bytes at once)
half4 vals = *(device const half4 *)(input + tid * 4);

The vectorized version generates a single memory transaction instead of four. On Apple Silicon, half4 loads are the sweet spot – 8 bytes per load, which matches the register file width.²

In akunu’s GEMV kernels, you will see patterns like:

// Load 16 half values at once (32 bytes)
device const half *x = B_half + row * K + il * 16;
// Thread reads 16 consecutive halves via multiple half4 loads

The simd_gemm_f16.metal kernel has each thread load 16 consecutive half-precision values per K-step, using the thread’s position within the SIMD group to cover different portions of the tile.

Pattern 2: Coalesced Memory Access

Coalescing means that adjacent threads access adjacent memory locations, so the hardware can merge their loads into a single wide memory transaction.

Thread 0 reads address 0x1000
Thread 1 reads address 0x1002
Thread 2 reads address 0x1004
...
Thread 31 reads address 0x103E

=> Hardware merges into ONE 64-byte transaction

Uncoalesced access – where threads read scattered addresses – is catastrophic. Instead of one transaction, you get 32 individual transactions, each wasting bus bandwidth.

Thread 0 reads address 0x1000   // row 0, col 0
Thread 1 reads address 0x4000   // row 1, col 0
Thread 2 reads address 0x7000   // row 2, col 0

=> 32 separate transactions (32x slower)

How Akunu Ensures Coalescing

In akunu’s GEMV kernels, the weight matrix B is stored in row-major order with rows corresponding to output dimensions. Each SIMD group processes a contiguous block of rows. Within the K-reduction loop, threads within a SIMD group read contiguous elements along the K dimension:

Thread 0: B[row][k_base + 0..3]
Thread 1: B[row][k_base + 4..7]
Thread 2: B[row][k_base + 8..11]
...

This is perfectly coalesced along K. For the output (N dimension), different SIMD groups write to different output rows, which are widely separated – but writes are much less frequent than reads, so this is acceptable.

For quantized formats, coalescing is trickier. Q4_0 blocks are 18 bytes: 16 four-bit values + 2 bytes of scale. The block layout is designed so that adjacent threads can read adjacent blocks, maintaining coalescing despite the non-power-of-2 block size.

Pattern 3: Threadgroup Memory and Bank Conflicts

Threadgroup memory (also called shared memory on NVIDIA) is a fast scratchpad local to a threadgroup. On Apple Silicon, it is organized into banks – typically 32 banks of 4 bytes each.³

Bank conflicts occur when multiple threads in a SIMD group access the same bank simultaneously:

// Bank conflict: threads 0 and 16 both access bank 0
shmem[tid * 32]  // stride 32 = exact bank width => every access hits bank 0!

// No conflict: threads access consecutive addresses
shmem[tid]       // each thread hits a different bank

In akunu’s GEMM kernels, the A and B tiles are loaded into threadgroup memory. The tile layout is carefully chosen to avoid bank conflicts during the subsequent simdgroup_load operations. The sa (A tile) and sb (B tile) pointers are offset:

threadgroup half *sa = shmem;
threadgroup half *sb = shmem + 4096 / sizeof(half);

The A tile gets 4096 bytes (2048 halves) and the B tile gets 2048 bytes (1024 halves), for a total of 6144 bytes per threadgroup. These sizes are chosen to minimize bank conflicts when loading into SIMD matrix registers.

Pattern 4: Loop Unrolling

GPU compilers can unroll loops, but sometimes you need to help. Unrolling reduces loop overhead (branch instructions, index increments) and exposes instruction-level parallelism:

// Before: tight loop, branch every iteration
for (int i = 0; i < 8; i++) {
    simdgroup_load(mb, sb + ..., stride);
    simdgroup_multiply_accumulate(mc[i], ma, mb, mc[i]);
}

// After (compiler typically does this): 8 loads + 8 MACs, no branches
// The 'constexpr' loop bound helps the compiler unroll

In akunu’s GEMM kernel, the inner K-loop has a stride of NK=32, with 4 sub-steps of 8 elements each. The sub-steps process 8 output tiles per SIMD group. The compiler unrolls both the sub-step loop and the tile loop because the bounds are compile-time constants.

Pattern 5: Function Constants (Metal Specialization)

Metal’s function constants are a form of compile-time specialization that lets you create optimized kernel variants without code duplication:

constant uint FC_GEMM_K [[function_constant(10)]];
constant bool FC_GEMM_K_SPECIALIZED = is_function_constant_defined(FC_GEMM_K);

// In the kernel:
const uint K_dim = FC_GEMM_K_SPECIALIZED ? FC_GEMM_K : K;

When FC_GEMM_K is defined, the compiler knows the K dimension at compile time and can:

• Unroll the K-loop completely for small K
• Eliminate bounds checking
• Pre-compute strides and offsets
• Optimize register allocation

Akunu uses function constants extensively for MLX quantized kernels, where group_size and K are baked into the pipeline state object:

uint32_t fc_indices[] = {0, 1};
uint32_t fc_values[] = {(uint32_t)quant_group_size, (uint32_t)K};
pso = device.get_pipeline(kernel_name, cache_key, fc_indices, fc_values, 2);

Each unique (group_size, K) combination gets a specialized pipeline. These are cached in MetalDevice::pso_cache_ so the specialization cost is paid once at model load time.⁴

Pattern 6: Kernel Fusion

Kernel fusion combines multiple operations into a single GPU dispatch. Each dispatch has overhead: pipeline binding, buffer binding, dispatch command encoding, and GPU scheduling. More importantly, each dispatch boundary forces intermediate results to be written to and re-read from global memory.

Fusions in Akunu

Akunu applies several fusions:

1. Residual + RMSNorm Fusion

Instead of separate residual_add and rmsnorm dispatches:

// Unfused: 3 dispatches, 3 reads + 3 writes
residual = input + skip_connection     // read input, skip; write residual
norm_input = rmsnorm(residual, weight) // read residual, weight; write norm_input

Akunu uses residual_rmsnorm_f16:

// Fused: 1 dispatch, reads input + skip + weight, writes norm_output + updated residual
residual_rmsnorm_f16(input, skip, weight, norm_output, residual, params)

This saves two kernel launches and two round-trips to global memory.

2. SiLU(gate) * up + Down GEMV Fusion

As discussed in the previous chapter, the fused SiLU+down kernel combines activation and projection:

// Unfused: 2 dispatches
act = SiLU(gate) * up              // read gate, up; write act
down = GEMV(down_weight, act)      // read act, down_weight; write down

// Fused: 1 dispatch
down = fused_silu_gemv(gate, up, down_weight)  // reads gate, up, weight; writes down

This eliminates the intermediate act buffer write and re-read.

3. QK-Norm + RoPE + KV Cache Write Fusion

For architectures with per-head Q/K norms (Qwen3, Gemma), akunu fuses head normalization, rotary position encoding, and KV cache writes into a single kernel:

// From table_builder.cpp:
Pipeline fused_pso = device.get_pipeline("head_norm_rope_neox_kv_write_f16");

This replaces 3-4 separate dispatches with one, which is especially impactful because these operations are tiny (operating on a single head at a time) and the dispatch overhead would dominate.

4. QKV Projection Fusion

When Q, K, and V weight matrices share the same dtype and the SLC is large enough to benefit:

bool fuse_qkv = (chip.should_fuse_weights || is_mlx) && q_dtype == k_dtype && k_dtype == v_dtype;
if (fuse_qkv) {
    Buffer fused_w = weights.fuse_weights(q_name, k_name, v_name);
    gemv(scratch.residual, fused_w, scratch.qkv, 0, q_dtype, qkv_total, dim);
}

Three GEMV dispatches become one, reading the fused weight matrix once instead of three times.

5. Gate + Up Projection Fusion

Same principle applied to the FFN gate and up projections:

bool fuse_gate_up = (chip.should_fuse_weights || gate_is_mlx) && (gate_dtype == up_dtype);
if (fuse_gate_up) {
    Buffer fused_gate_up_w = weights.fuse_weights(gate_name, up_name);
    gemv(scratch.attn_out, fused_gate_up_w, scratch.ffn_gate, 0, gate_dtype, 2 * ffn_dim, dim);
}

When NOT to Fuse

Fusion is not always beneficial. The SLC-gated fusion decisions in akunu illustrate this: on chips with small SLC (M1 base, 8 MB estimated), the fused QKV weight matrix may be too large to fit in cache, causing more cache thrashing than the unfused version. The ChipConfig::should_fuse_weights flag controls this:

c.should_fuse_weights = (c.slc_bytes >= 16 * 1024 * 1024);  // Pro+ and M4 Base

Pattern 7: Occupancy and Threadgroup Sizing

Occupancy is the ratio of active threads to the maximum the GPU can support simultaneously. Higher occupancy generally means better latency hiding – when one SIMD group stalls on a memory access, another can execute.

On Apple Silicon, each GPU core can run multiple threadgroups concurrently (the exact limit depends on register pressure and threadgroup memory usage). The threadgroup size directly affects occupancy:

Akunu’s ChipConfig controls threadgroup sizing for normalization kernels:

c.norm_tg_size = 1024;  // max threads for RMSNorm
c.max_threads_per_tg = 1024;  // Metal's maximum

For RMSNorm on a 4096-dimensional model, the threadgroup size is min(4096, 1024) = 1024. All 1024 threads participate in the reduction (computing the root-mean-square), with the final result broadcast to all threads for the normalization step.

For GEMV, the threadgroup size trades off between parallelism and overhead:

• 128 threads (4 SIMD groups): Standard GEMV. Good for small chips with limited cores.
• 256 threads (8 SIMD groups): Wide GEMV. Better occupancy on Pro+ chips with many cores. Controlled by chip.gemv_wide_standard.

Pattern 8: Avoiding Redundant Work with Pre-Computation

The most efficient computation is the one you do not do. Akunu pre-computes everything possible at model load time:

1. Pipeline state objects are created and cached in pso_cache_. No pipeline compilation during inference.

2. Buffer bindings are resolved once in the dispatch table. The DispatchCmd stores actual Buffer handles and byte offsets, not symbolic names.

3. Kernel parameters (dimensions, epsilon values, strides) are stored as raw bytes in DispatchCmd::param_bytes[64]. For static params, a GPU buffer is pre-allocated:

// From table_builder.cpp, end of build_dispatch_table():
for (auto& cmd : cmds) {
    if (cmd.param_size > 0) {
        cmd.param_buf = device.allocate(cmd.param_bytes, cmd.param_size);
    }
}

This means the dispatch table replay loop does almost zero work per command beyond Metal API calls. No string lookups, no hash table queries, no conditional logic.

4. RoPE frequencies can be pre-computed into a GPU buffer (arch.rope_freqs) rather than computed per-token.

5. Dispatch geometry (grid size, threadgroup size) is computed once and stored in DispatchCmd::grid and DispatchCmd::threadgroup.

Pattern 9: Minimizing Command Buffer Overhead

Each Metal command buffer has submission overhead: the CPU must package the commands, the GPU command processor must parse them, and there is a synchronization cost at completion. Akunu minimizes this in two ways:

1. Chain Decode: Many Tokens Per Command Buffer

Instead of one command buffer per token, akunu’s chain decode batches multiple tokens into a single submission:

// ChipConfig determines batch size:
c.chain_decode_chunk = 128;  // M4: 128 tokens per submission
c.chain_decode_chunk = 64;   // M1 base: 64 tokens per submission

One command buffer encodes the full dispatch table N times (once per token), with per-token patching of position and offsets. This amortizes the command buffer overhead across many tokens.

2. Unretained References

Akunu uses commandBufferWithUnretainedReferences, which tells Metal not to retain buffer references. This avoids atomic reference counting on every setBuffer call – a significant savings when a single command buffer contains thousands of buffer bindings.⁵

3. Event-Based Pipelining

Akunu supports overlapping GPU execution with CPU encoding using MTLSharedEvent:

// GPU signals event after completing
[cmdBuffer encodeSignalEvent:pipelineEvent value:signalVal];
[cmdBuffer commit];

// Next command buffer waits for the event (GPU-GPU sync, no CPU involvement)
[nextCmdBuffer encodeWaitForEvent:pipelineEvent value:eventValue];

This allows the CPU to encode the next batch of tokens while the GPU is still processing the current one. The event-based synchronization is GPU-to-GPU, avoiding a CPU round-trip.

Pattern 10: The setBuffer vs setBytes Split

This pattern is specific to akunu’s dispatch table replay and deserves its own section. During chain decode, the same commands are repeated for each token. Most parameters are identical across tokens – only the position-dependent fields change.

Akunu splits parameters into two categories:

This is visible in the encode_chain fast path. For each command, the encoder patches position-dependent parameters in-place and calls set_bytes to inline the (small, <64 byte) parameter data into the command buffer:

for (int i = 0; i < n; i++) {
    for (auto& cmd : table.commands) {
        device.set_pipeline(cmd.pso);
        // Bind buffers (static — same every token)
        for (int b = 0; b < cmd.buffer_count; b++)
            device.set_buffer(cmd.buffers[b], cmd.offsets[b], b);
        // Patch and set parameters (only position/kv_len change per token)
        if (cmd.patch_type != PATCH_NONE)
            apply_patch(cmd, pos + i);
        device.set_bytes(cmd.param_bytes, cmd.param_size, cmd.param_index);
        device.dispatch(cmd.grid, cmd.threadgroup);
    }
}

For a typical 32-layer model, the dispatch table has ~260 commands. Only the RoPE, attention, and argmax commands need per-token patching (~100 commands). The rest pass through with unchanged param_bytes — but all use set_bytes (not setBuffer) since the parameter data is always small enough (<64 bytes) to inline.

Profiling Tools

Knowing these patterns is only useful if you can measure their impact. Apple provides several profiling tools:

Xcode GPU Debugger

Capture a Metal frame and inspect:

• Per-dispatch GPU time
• Memory bandwidth utilization
• Occupancy
• Wait time (stalls)

Metal System Trace (Instruments)

Part of Instruments.app. Shows:

• Command buffer submission and completion timeline
• GPU utilization over time
• CPU-GPU synchronization points
• Memory allocation events

akunu_profile

Akunu includes a profiling tool that uses per-layer command buffers to get GPU timing for each operation:

// From akunu.h:
int akunu_profile_decode_step(akunu_model_t model, uint32_t token_id,
                              int position, float *timing_out, int max_entries);
const char *akunu_profile_label(akunu_model_t model, int index);

This runs each layer in its own command buffer (much slower than normal inference) but gives you per-operation GPU timing. The output looks like:

embedding              0.012 ms
layer.0.attention      0.045 ms
layer.0.rope_kv_write  0.008 ms
layer.0.attention      0.082 ms
layer.0.o_proj         0.041 ms
layer.0.fused_ffn_norm 0.006 ms
layer.0.gate_up_proj   0.078 ms
layer.0.ffn            0.043 ms
...

Counter Sampling

Metal supports GPU hardware counter sampling for detailed performance analysis. You can measure:

• ALU utilization
• Memory read/write bytes
• Cache hit rates
• Occupancy percentages

These are available through MTLCounterSampleBuffer and are essential for diagnosing whether a kernel is compute-bound or bandwidth-bound.⁶

Putting It All Together: A Single Layer’s Performance Profile

Let’s trace through one transformer layer during decode (M=1) on M4 Pro and identify the bottleneck for each operation:

For 32 layers: ~13.8 ms/token = ~72 tok/s. This is close to the theoretical bandwidth limit of ~70 tok/s we computed in the memory chapter, confirming that decode is bandwidth-saturated.

Common Mistakes

Mistake 1: Optimizing Compute for a Bandwidth-Bound Kernel

If your GEMV kernel is at 95% bandwidth utilization and 10% compute utilization, making the math faster will not help. You need to reduce the number of bytes read (quantize to lower bits, fuse operations to eliminate intermediate buffers).

Mistake 2: Tiny Threadgroups

Using a threadgroup of 32 threads for a GEMV kernel means only 1 SIMD group per threadgroup. The GPU core has no threads to switch to when this SIMD group stalls on memory. Use at least 128 threads (4 SIMD groups) for any memory-heavy kernel.

Mistake 3: Forgetting Threadgroup Barriers

In GEMM kernels that use threadgroup memory, forgetting threadgroup_barrier(mem_flags::mem_threadgroup) between writing to and reading from shared memory causes data races. The barrier ensures all threads in the threadgroup have completed their writes before any thread reads.

Mistake 4: Over-Fusing

Fusing too many operations into one kernel can increase register pressure, reducing occupancy and hurting performance. If a fused kernel needs more registers than the hardware can provide, the GPU will “spill” registers to memory, destroying performance. The separate activation + GEMV fallback in akunu exists for exactly this reason.

Summary

Performance optimization on Apple Silicon GPU follows a clear decision tree:

1. Is the kernel bandwidth-bound or compute-bound? Use the roofline model.
2. If bandwidth-bound: Reduce bytes (quantize, fuse), improve access patterns (coalesce, vectorize).
3. If compute-bound: Maximize utilization (occupancy, SIMD matrix ops, loop unrolling).
4. Always: Pre-compute everything possible, minimize dispatch overhead, use function constants for specialization.

Akunu applies all of these patterns systematically, with the ChipConfig and DTypeDescriptor tables encoding the chip-specific and dtype-specific tuning decisions. The dispatch table pre-resolves everything at model load time so the hot path is a tight loop of Metal API calls.

1. Williams, S., Waterman, A., & Patterson, D. (2009). “Roofline: An Insightful Visual Performance Model for Multicore Architectures.” Communications of the ACM, 52(4), 65-76. The roofline model remains the most effective tool for classifying GPU kernel performance. See https://doi.org/10.1145/1498765.1498785. ↩

2. Apple, “Metal Best Practices Guide: Optimize Memory Accesses,” 2024. Vectorized loads are recommended for maximizing memory throughput. See https://developer.apple.com/documentation/xcode/analyzing-the-performance-of-your-metal-app. ↩

3. The exact bank configuration on Apple Silicon is not officially documented. The 32-bank / 4-byte-per-bank configuration is inferred from performance profiling and is consistent with other GPU architectures. See Dougall Johnson’s Apple GPU documentation: https://dougallj.github.io/applegpu/. ↩

4. Apple, “Using Function Specialization to Build Pipeline Variants.” Function constants are the recommended way to create specialized shader variants without preprocessor macros. See https://developer.apple.com/documentation/metal/using-function-specialization-to-build-pipeline-variants. ↩

5. Apple, “commandBufferWithUnretainedReferences Documentation.” Unretained references eliminate atomic retain/release overhead but require manual lifetime management. See https://developer.apple.com/documentation/metal/mtlcommandqueue/1508684-makecommandbufferwithunretainedr. ↩

6. Apple, “Optimizing Performance with the GPU Counters Instrument,” 2024. GPU hardware counters provide per-kernel metrics including ALU utilization, memory bandwidth, and cache hit rates. See https://developer.apple.com/documentation/xcode/analyzing-the-performance-of-your-metal-app. ↩

Tensors and Linear Algebra on GPU

If you have spent any time reading ML papers or browsing inference codebases, you have seen the word “tensor” more times than you can count. It gets thrown around so casually that you might suspect it means something terribly complicated. It does not. At its heart, a tensor in the machine learning context is just a multi-dimensional array of numbers. That is it. A 1D tensor is a vector. A 2D tensor is a matrix. A 3D tensor is… well, a 3D array. And so on.

But here is the thing: while the concept is simple, the implementation on a GPU is where all the interesting engineering lives. How do you lay out a 4-dimensional array in a flat slab of GPU memory? How do you make sure that when 3,000 threads all try to read different elements at once, the memory system does not choke? How do you express broadcasting rules in terms of pointer arithmetic? These are the questions that determine whether your inference engine runs at 10 tokens per second or 100.

In this chapter, we will build your intuition for tensors from the ground up, then show exactly how they map to Metal GPU buffers. By the end, you will understand the memory layout decisions that dominate every kernel we write in later chapters.

What Is a Tensor, Really?

Let us start with the basics and build up.

Scalars, Vectors, Matrices, and Beyond

Rank 0 (Scalar):     42.0

Rank 1 (Vector):     [1.0, 2.0, 3.0, 4.0]

Rank 2 (Matrix):     [[1.0, 2.0, 3.0],
                       [4.0, 5.0, 6.0]]

Rank 3 (3D Tensor):  [[[1.0, 2.0],
                        [3.0, 4.0],
                        [5.0, 6.0]],

                       [[7.0, 8.0],
                        [9.0, 10.0],
                        [11.0, 12.0]]]

The rank (or order) of a tensor is simply how many indices you need to pick out a single element:

• Rank 0: zero indices. It is just a number.
• Rank 1: one index. v[i] gives you element i of a vector.
• Rank 2: two indices. M[i][j] gives you the element at row i, column j.
• Rank 3: three indices. T[i][j][k] picks one element from a 3D block.

The shape describes how many elements exist along each dimension. For the rank-3 tensor above, the shape is (2, 3, 2) – two “slices,” each containing 3 rows of 2 elements.

Tensors in Transformers

In a typical transformer model, you encounter tensors of various ranks constantly:

Token embeddings:       shape (seq_len, d_model)         Rank 2
Weight matrix:          shape (d_model, d_out)            Rank 2
Attention scores:       shape (n_heads, seq_len, seq_len) Rank 3
Batched activations:    shape (batch, seq_len, d_model)   Rank 3
Multi-head QKV:         shape (batch, n_heads, seq, d_k)  Rank 4

For inference on a single sequence (batch size 1), we can often drop the batch dimension. But even without batching, we regularly work with rank-2 and rank-3 tensors.

Here is a concrete example. Suppose we have a small model with:

• d_model = 512 (embedding dimension)
• n_heads = 8 (number of attention heads)
• d_k = 64 (per-head dimension, since 512 / 8 = 64)
• seq_len = 128 (sequence length)

The query tensor Q after projection would have shape (128, 512), and after reshaping for multi-head attention, it becomes (8, 128, 64).

Memory Layouts: How Multi-Dimensional Data Lives in Flat Memory

Here is the fundamental problem: GPU memory (and CPU memory, for that matter) is a flat, one-dimensional address space. You have addresses 0, 1, 2, 3, and so on. But your tensor is multi-dimensional. So you need a rule for mapping multi-dimensional indices to flat memory addresses.

Row-Major Order (C-Style)

The most common layout, and the one used by virtually all ML frameworks and inference engines, is row-major order. The idea is simple: you lay out the last dimension contiguously, then the second-to-last, and so on.

For a 2D matrix with shape (3, 4):

Logical view:           Memory layout (row-major):

     col 0  col 1  col 2  col 3
row 0 [ a      b      c      d  ]     addr 0:  a
row 1 [ e      f      g      h  ]     addr 1:  b
row 2 [ i      j      k      l  ]     addr 2:  c
                                       addr 3:  d
                                       addr 4:  e
                                       addr 5:  f
                                       addr 6:  g
                                       addr 7:  h
                                       addr 8:  i
                                       addr 9:  j
                                       addr 10: k
                                       addr 11: l

The elements of row 0 are contiguous in memory (a, b, c, d), followed by row 1 (e, f, g, h), followed by row 2 (i, j, k, l). The last index changes fastest as you walk through memory.

To find element M[i][j] in a (R, C) matrix:

address(i, j) = base + (i * C + j) * element_size

Column-Major Order (Fortran-Style)

The alternative is column-major order, where the first dimension is contiguous. BLAS libraries and MATLAB use this. The same (3, 4) matrix in column-major:

Memory layout (column-major):

addr 0:  a   (row 0, col 0)
addr 1:  e   (row 1, col 0)
addr 2:  i   (row 2, col 0)
addr 3:  b   (row 0, col 1)
addr 4:  f   (row 1, col 1)
addr 5:  j   (row 2, col 1)
addr 6:  c   (row 0, col 2)
addr 7:  g   (row 1, col 2)
addr 8:  k   (row 2, col 2)
addr 9:  d   (row 0, col 3)
addr 10: h   (row 1, col 3)
addr 11: l   (row 2, col 3)

Here the first index changes fastest. To find element M[i][j]:

address(i, j) = base + (j * R + i) * element_size

Why Does This Matter for GPU Performance?

Consider what happens when 32 threads in a SIMD group all need to read one element each. If thread t reads element M[row][t] and the matrix is row-major, all 32 threads read from consecutive memory addresses. The GPU’s memory system can satisfy this with a single coalesced memory transaction.

Row-major, threads reading M[row][0..31]:

Memory: [ M[r][0] | M[r][1] | M[r][2] | ... | M[r][31] | M[r][32] | ... ]
           ^          ^          ^                ^
         thread 0   thread 1   thread 2       thread 31

All addresses are contiguous --> ONE memory transaction (coalesced)

But if the matrix were column-major and threads tried to read M[row][0..31], each thread would be reading from addresses that are R elements apart:

Column-major, threads reading M[row][0..31]:

Memory: [ M[0][0] | M[1][0] | M[2][0] | ... | M[0][1] | M[1][1] | ... ]

Thread 0 reads M[row][0] at offset row
Thread 1 reads M[row][1] at offset R + row
Thread 2 reads M[row][2] at offset 2R + row
...
Addresses are R elements apart --> MANY memory transactions (strided)

This can be 10-30x slower depending on the stride. Coalesced access is one of the single most important performance considerations on any GPU.

Strides: The General Case

The concept of strides generalizes memory layout. A stride tells you how many elements to skip in memory when you increment one index by one.

For a row-major matrix of shape (R, C):

• Stride along dimension 0 (rows): C (skip one whole row of C elements)
• Stride along dimension 1 (columns): 1 (elements are adjacent)

For a row-major 3D tensor of shape (D0, D1, D2):

• Stride along dimension 0: D1 * D2
• Stride along dimension 1: D2
• Stride along dimension 2: 1

The general address formula for any tensor:

address(i0, i1, ..., in) = base + sum(ik * stride_k) for k in 0..n

Here is a concrete example. For a tensor of shape (2, 3, 4) in row-major:

Strides: (12, 4, 1)

Because:
  stride[2] = 1
  stride[1] = shape[2] = 4
  stride[0] = shape[1] * shape[2] = 3 * 4 = 12

Element T[1][2][3]:
  address = 0 + 1*12 + 2*4 + 3*1 = 12 + 8 + 3 = 23

Let us verify by counting:
  Slice 0: [[ 0  1  2  3]    elements 0-11
            [ 4  5  6  7]
            [ 8  9 10 11]]

  Slice 1: [[12 13 14 15]    elements 12-23
            [16 17 18 19]
            [20 21 22 23]]   <-- T[1][2][3] = element 23. Correct!

Strides also let you express interesting things like transposition without copying data. If you have a matrix with strides (C, 1), its transpose has strides (1, C) – same data, just different stride interpretation.

How Tensors Map to GPU Buffers

On Metal, a tensor lives in an MTLBuffer – a contiguous block of GPU-accessible memory. There is no built-in “tensor” type in Metal. The buffer is just bytes. Your shader code interprets those bytes as a tensor by computing offsets from indices using the stride formula.

The Buffer Layout

MTLBuffer (contiguous memory):
+----+----+----+----+----+----+----+----+----+----+----+----+
| e0 | e1 | e2 | e3 | e4 | e5 | e6 | e7 | e8 | e9 |e10 |e11 |
+----+----+----+----+----+----+----+----+----+----+----+----+
  ^                                                         ^
  buffer.contents()                                   buffer end

Each element is sizeof(element_type) bytes:
  float  = 4 bytes
  half   = 2 bytes
  int8_t = 1 byte

For a half-precision (2, 3, 2) tensor:
  Total elements = 2 * 3 * 2 = 12
  Total bytes = 12 * 2 = 24 bytes

Computing Offsets in a Metal Shader

In a Metal compute shader, you typically receive the buffer pointer and the tensor’s metadata (shape, strides) as arguments. Here is what the pattern looks like:

kernel void elementwise_relu(
    device const half* input  [[buffer(0)]],
    device half*       output [[buffer(1)]],
    constant uint3&    shape  [[buffer(2)]],   // (D0, D1, D2)
    constant uint3&    strides [[buffer(3)]],  // (S0, S1, S2)
    uint3 gid [[thread_position_in_grid]]
) {
    // Bounds check
    if (gid.x >= shape.x || gid.y >= shape.y || gid.z >= shape.z) return;

    // Compute flat index from multi-dimensional position
    uint idx = gid.x * strides.x + gid.y * strides.y + gid.z * strides.z;

    // Read, compute, write
    half val = input[idx];
    output[idx] = val > 0 ? val : 0;
}

The key insight: the GPU hardware knows nothing about tensors. It only knows about flat memory addresses. All the multi-dimensional indexing is just arithmetic that we do in the shader.

Alignment Considerations

Metal has specific alignment requirements for buffer access:

Type        Size    Alignment
float       4B      4B
half        2B      2B
float4      16B     16B
half4       8B      8B

Rule of thumb: access addresses that are multiples of the type size.

When packing tensor data into buffers, you generally want each row to start at an aligned address. For a matrix of half values with 127 columns, each row takes 254 bytes. The next row starts at byte 254, which is 2-byte aligned (fine for half). But if you wanted to read rows as half4 vectors (common optimization), you would want 8-byte alignment, which means padding rows to 128 elements (256 bytes).

Multiple Tensors in One Buffer (Offset Packing)

In practice, inference engines often pack multiple tensors into a single large buffer with offsets:

Single MTLBuffer:
+------------------+------------------+------------------+
|  Weight matrix   |   Bias vector    |  LayerNorm scale |
|  (4096 x 4096)   |   (4096)         |  (4096)          |
|  offset: 0       |   offset: 33MB   |  offset: 33MB+8K |
+------------------+------------------+------------------+

This reduces the number of buffer bindings needed (Metal has a limit on how many buffers you can bind to a single kernel invocation) and can improve memory allocation efficiency.

Common Tensor Operations

Now let us look at the operations that transformers actually perform on tensors and how they parallelize on a GPU.

Element-wise Operations

These apply a function independently to each element. The output has the same shape as the input. Examples:

• ReLU: output[i] = max(0, input[i])
• SiLU (Swish): output[i] = input[i] * sigmoid(input[i])
• GELU: output[i] = input[i] * 0.5 * (1 + erf(input[i] / sqrt(2)))
• Addition: output[i] = a[i] + b[i]
• Scalar multiply: output[i] = alpha * input[i]

Parallelization is trivial – assign one thread per element:

Input tensor:  [a0, a1, a2, a3, a4, a5, a6, a7]

Thread 0 --> processes a0
Thread 1 --> processes a1
Thread 2 --> processes a2
...
Thread 7 --> processes a7

Each thread:
  1. Read input[thread_id]
  2. Apply function
  3. Write output[thread_id]

No synchronization needed! Each thread is independent.

For a tensor with N elements, you dispatch ceil(N / threads_per_threadgroup) threadgroups, each with threads_per_threadgroup threads (commonly 256 or 1024).

Dispatch for N = 10000, threads_per_group = 256:

Threadgroups needed = ceil(10000 / 256) = 40 threadgroups

Threadgroup  0: threads 0-255     --> elements 0-255
Threadgroup  1: threads 256-511   --> elements 256-511
...
Threadgroup 38: threads 9728-9983 --> elements 9728-9983
Threadgroup 39: threads 9984-10239 --> elements 9984-9999 (rest out of bounds)

Reduction Operations

Reductions collapse one or more dimensions by aggregating elements. Examples:

• Sum along a dimension: output[i] = sum(input[i][j] for all j)
• Max along a dimension: output[i] = max(input[i][j] for all j)
• Mean: sum divided by count

These are trickier to parallelize because threads need to cooperate to produce the result.

Strategy 1: One threadgroup per reduction

Matrix shape (4, 8), reduce along columns (dim 1):

Row 0: [a0  a1  a2  a3  a4  a5  a6  a7]  --> sum = a0+a1+...+a7
Row 1: [b0  b1  b2  b3  b4  b5  b6  b7]  --> sum = b0+b1+...+b7
Row 2: [c0  c1  c2  c3  c4  c5  c6  c7]  --> sum = c0+c1+...+c7
Row 3: [d0  d1  d2  d3  d4  d5  d6  d7]  --> sum = d0+d1+...+d7

Assign one threadgroup per row.
Within each threadgroup, threads cooperate to sum the row:

Threadgroup 0 (8 threads for row 0):
  Step 1: Each thread loads one element
    t0=a0, t1=a1, t2=a2, t3=a3, t4=a4, t5=a5, t6=a6, t7=a7

  Step 2: Parallel reduction tree
    t0 = a0+a4   t1 = a1+a5   t2 = a2+a6   t3 = a3+a7
    t0 = (a0+a4)+(a2+a6)      t1 = (a1+a5)+(a3+a7)
    t0 = total sum

  Step 3: Thread 0 writes result

Strategy 2: SIMD group reduction

On Metal, SIMD groups of 32 threads have special fast operations for reduction. The simd_sum() function sums a value across all threads in the SIMD group in a single instruction cycle:

SIMD group reduction (32 threads):

Before simd_sum:
  t0=3.0  t1=1.0  t2=4.0  t3=1.0  ...  t31=2.0

After val = simd_sum(val):
  t0=sum  t1=sum  t2=sum  t3=sum  ...  t31=sum
  (all threads have the same total sum)

For reductions over dimensions larger than 32, you use a combination: each SIMD group reduces its portion, then results are combined using threadgroup memory.

Broadcasting

Broadcasting lets you operate on tensors of different shapes by “stretching” the smaller tensor to match the larger one, conceptually. No data is actually copied.

Example: Add a bias vector to every row of a matrix

Matrix A shape: (1024, 512)
Bias b shape:   (512,)

Result C[i][j] = A[i][j] + b[j]    for all i, j

The bias is "broadcast" along dimension 0:

    A:                      b:                  C:
    [[a00 a01 ... a0,511]   [b0 b1 ... b511]    [[a00+b0  a01+b1  ...]
     [a10 a11 ... a1,511] +                   =   [a10+b0  a11+b1  ...]
     ...                                           ...
     [a1023,0 ...       ]]                         [a1023,0+b0 ...   ]]

In the shader, broadcasting is implemented by simply not advancing the index along the broadcast dimension:

// Adding bias (shape: [N]) to matrix (shape: [M, N])
kernel void add_bias(
    device const half* matrix [[buffer(0)]],
    device const half* bias   [[buffer(1)]],
    device half*       output [[buffer(2)]],
    uint2 gid [[thread_position_in_grid]]   // (row, col)
) {
    uint row = gid.x;
    uint col = gid.y;
    uint idx = row * N + col;

    output[idx] = matrix[idx] + bias[col];  // bias indexed only by col
}

The bias[col] access ignores the row entirely. That is broadcasting – the same bias element b[j] is used for every row.

Parallelizing Tensor Operations

The general strategies for mapping tensor operations to GPU threads fall into a few patterns. Let us formalize them.

Pattern 1: One Thread Per Element

Used for element-wise operations.

Tensor shape: (M, N)
Grid: (M, N)
Each thread (i, j) processes element [i][j]

+-------+-------+-------+-------+
| t(0,0)| t(0,1)| t(0,2)| t(0,3)|
+-------+-------+-------+-------+
| t(1,0)| t(1,1)| t(1,2)| t(1,3)|
+-------+-------+-------+-------+
| t(2,0)| t(2,1)| t(2,2)| t(2,3)|
+-------+-------+-------+-------+

Dispatch: grid_size = (M, N, 1)
          threadgroup_size = (16, 16, 1)  // 256 threads per group

Pattern 2: One Threadgroup Per Row (or Column)

Used for reductions, softmax, layer normalization – anything that operates across one dimension.

Tensor shape: (M, N)
Grid: M threadgroups, each with T threads

Threadgroup 0 --> Row 0:  [e0  e1  e2 ... eN-1]
Threadgroup 1 --> Row 1:  [e0  e1  e2 ... eN-1]
...
Threadgroup M-1 --> Row M-1

Within each threadgroup, T threads cooperate:
  Thread t handles elements: t, t+T, t+2T, ...

For N=1024, T=256:
  Thread 0: elements 0, 256, 512, 768
  Thread 1: elements 1, 257, 513, 769
  ...

Pattern 3: Tiled Processing

Used for matrix multiplication and attention. Threads in a threadgroup cooperatively load tiles into fast threadgroup memory, then compute.

Matrix multiply C = A * B
A: (M, K), B: (K, N), C: (M, N)

Tile size: (Tm, Tn) per threadgroup, iterate over K in tiles of Tk

Threadgroup (gi, gj) computes C[gi*Tm..(gi+1)*Tm][gj*Tn..(gj+1)*Tn]

       B (K x N)
       +---+---+---+---+
       |   | Bj|   |   |      Bj = tile column j
       +---+---+---+---+
       |   | Bj|   |   |
  A    +---+---+---+---+
(M x K)
+--+--+   +---+
|  |Ai| x | Bj| = Cij (Tm x Tn tile of output)
+--+--+   +---+
|  |  |
+--+--+

Ai = tile row i of A
For each k-tile: load Ai[:,k:k+Tk] and Bj[k:k+Tk,:] into shared memory
                 compute partial products
                 accumulate into Cij registers

We will explore tiled matrix multiplication in enormous detail in the next chapter.

Pattern 4: SIMD Group per Work Unit

Metal’s SIMD groups (wavefronts) of 32 threads have hardware support for fast intra-group communication. Many kernels assign one SIMD group to one logical unit of work:

GEMV: One SIMD group per output element (or a few elements)

Output vector y = W * x    (W is M x K, x is K x 1, y is M x 1)

SIMD group 0 --> y[0] = dot(W[0,:], x)
SIMD group 1 --> y[1] = dot(W[1,:], x)
...

Within each SIMD group:
  32 threads divide the K-length dot product
  Thread t handles indices: t, t+32, t+64, ...
  Each thread accumulates a partial sum
  simd_sum() gives the total

The Transformer’s Core Operations

A transformer model is built from a surprisingly small set of tensor operations, repeated many times. Let us catalog them and understand their computational characteristics.

Linear Projection (GEMM / GEMV)

The workhorse of the transformer. Every layer has multiple linear projections:

In one transformer block:
  Q = X * Wq + bq       (query projection)
  K = X * Wk + bk       (key projection)
  V = X * Wv + bv       (value projection)
  O = Attn * Wo + bo    (output projection)
  H = X * W1 + b1       (FFN first layer)
  Y = H * W2 + b2       (FFN second layer)

That's 6 matrix multiplications per layer!
A 32-layer model does 192 matmuls per forward pass.

The computation is Y = X * W where:

• During prefill: X is (seq_len, d_model), W is (d_model, d_out) – a full GEMM
• During decode: X is (1, d_model), W is (d_model, d_out) – a GEMV

Prefill (GEMM):                    Decode (GEMV):

 X (128 x 4096)                     x (1 x 4096)
 +----------+                        +----------+
 |          |                        |          |
 |          |    W (4096 x 4096)     +----------+    W (4096 x 4096)
 |          |    +----------+                        +----------+
 |          |  x |          |                      x |          |
 |          |    |          |                        |          |
 |          |    |          |                        |          |
 +----------+    +----------+                        +----------+
       =                                   =
 Y (128 x 4096)                     y (1 x 4096)
 +----------+                        +----------+
 |          |                        |          |
 |          |                        +----------+
 |          |
 +----------+

The GEMV case (decode) is memory-bandwidth bound because you read the entire weight matrix just to produce one output row. The GEMM case (prefill) is compute bound because you amortize the weight read across many input rows.

Attention (Batched Dot Products + Softmax)

The attention mechanism computes, for each query position, a weighted average of all value positions:

scores = Q * K^T / sqrt(d_k)     Matrix multiply: (seq, d_k) x (d_k, seq) = (seq, seq)
weights = softmax(scores)         Element-wise + reduction
output = weights * V              Matrix multiply: (seq, seq) x (seq, d_k) = (seq, d_k)

The middle step (softmax) is a reduction along the last dimension of each row, which requires special handling. We will devote Chapters 15 and 16 to attention.

Layer Normalization (RMSNorm)

Modern transformers use RMSNorm, which normalizes each row by its root-mean-square:

RMSNorm(x) = x / sqrt(mean(x^2) + eps) * gamma

For a row x of length d_model:
  1. Compute sum of squares:  ss = sum(x[i]^2 for i in 0..d_model)
  2. Compute RMS:             rms = sqrt(ss / d_model + eps)
  3. Normalize and scale:     output[i] = x[i] / rms * gamma[i]

Step 1 is a reduction (sum).
Step 3 is element-wise.

This follows the “one threadgroup per row” pattern. The reduction in step 1 uses SIMD group sums.

Activation Functions

Applied element-wise between the FFN layers:

SiLU(x) = x * sigmoid(x) = x * (1 / (1 + exp(-x)))

GELU(x) = x * 0.5 * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))

These are purely element-wise, one thread per element. Trivially parallel.

Residual Connections

Just element-wise addition:

output = sublayer_output + input

Every transformer block has two residual connections:
  x = attention(norm(x)) + x      <-- residual
  x = ffn(norm(x)) + x            <-- residual

Putting It All Together: One Transformer Block

Input: x (seq_len, d_model)

Step 1: RMSNorm          --> x_norm (seq_len, d_model)         [reduction per row]
Step 2: Q = x_norm * Wq  --> Q (seq_len, d_model)             [GEMM or GEMV]
Step 3: K = x_norm * Wk  --> K (seq_len, d_kv)                [GEMM or GEMV]
Step 4: V = x_norm * Wv  --> V (seq_len, d_kv)                [GEMM or GEMV]
Step 5: Reshape Q/K/V for multi-head
Step 6: Attention         --> attn (seq_len, d_model)          [batched matmul + softmax]
Step 7: O = attn * Wo     --> O (seq_len, d_model)            [GEMM or GEMV]
Step 8: residual          --> x = O + x                        [element-wise add]
Step 9: RMSNorm           --> x_norm                           [reduction per row]
Step 10: FFN up           --> h (seq_len, d_ff)                [GEMM or GEMV]
Step 11: FFN gate         --> g (seq_len, d_ff)                [GEMM or GEMV]
Step 12: SiLU + multiply  --> h = SiLU(g) * h                 [element-wise]
Step 13: FFN down         --> f (seq_len, d_model)             [GEMM or GEMV]
Step 14: residual         --> x = f + x                        [element-wise add]

Output: x (seq_len, d_model)

Count the operations: 6 matrix multiplications, 2 normalizations, 2 residual adds, 1 activation, and the attention computation. The matrix multiplications dominate compute time (90%+ of total FLOPS).

Floating-Point Precision: FP16, FP32, and BF161

The choice of numeric precision has enormous impact on both performance and quality.

FP32 (Single Precision)

FP32: 1 sign bit + 8 exponent bits + 23 mantissa bits = 32 bits

+---+----------+-----------------------+
| S | EEEEEEEE | MMMMMMMMMMMMMMMMMMMMMMM|
+---+----------+-----------------------+
 1      8              23

Range: ~1.18e-38 to ~3.4e+38
Precision: ~7 decimal digits
Size: 4 bytes per value

FP32 is the “safe” choice. It has plenty of range and precision for any ML computation. But it uses twice the memory and bandwidth of FP16, making it 2x slower for memory-bound operations.

FP16 (Half Precision)

FP16: 1 sign bit + 5 exponent bits + 10 mantissa bits = 16 bits

+---+-------+------------+
| S | EEEEE | MMMMMMMMMM |
+---+-------+------------+
 1     5          10

Range: ~6.1e-5 to ~65504
Precision: ~3.3 decimal digits
Size: 2 bytes per value

FP16 halves memory usage and doubles effective bandwidth. Apple GPUs have native FP16 support and can process FP16 values at 2x the rate of FP32 in many operations.

The downside is the limited range. Values smaller than 6.1e-5 become zero (underflow), and values larger than 65504 become infinity (overflow). During inference, weights are usually small enough that this is not a problem. But intermediate computations (especially in attention) can overflow if not handled carefully.

BF16 (Brain Float 16)

BF16: 1 sign bit + 8 exponent bits + 7 mantissa bits = 16 bits

+---+----------+---------+
| S | EEEEEEEE | MMMMMMM |
+---+----------+---------+
 1      8           7

Range: ~1.18e-38 to ~3.4e+38  (same as FP32!)
Precision: ~2.4 decimal digits
Size: 2 bytes per value

BF16 has the same exponent range as FP32 (so no overflow/underflow issues) but less precision than FP16. It was designed specifically for ML workloads where range matters more than precision.

Apple Silicon support: M1/M2 GPUs do not have native BF16 support in Metal. M3+ has some BF16 capabilities. In practice, most Apple GPU inference uses FP16 for computation with FP32 for accumulation.

The Mixed-Precision Strategy

The standard approach in inference:

Weights:        stored in FP16 (or quantized, e.g., 4-bit)
Activations:    computed in FP16
Accumulations:  done in FP32 (dot products, reductions)
Final output:   cast back to FP16

Why FP32 for accumulations?

Consider summing 4096 FP16 values, each around 0.01:
  True sum: ~40.96
  FP16 sum: after adding a few hundred values, each new addition
            changes the sum by less than FP16 precision can represent.
            You lose significant accuracy.

  FP32 sum: no problem. 7 decimal digits of precision is plenty.

In Metal shader code, this looks like:

// Dot product with mixed precision
float sum = 0.0f;    // FP32 accumulator
for (uint i = 0; i < K; i++) {
    sum += float(A[i]) * float(B[i]);   // Cast FP16 inputs to FP32
}
output[idx] = half(sum);   // Cast result back to FP16

Precision Comparison Table

+--------+------+--------+-----------+-------------+------------------+
| Format | Bits | Range  | Precision | Memory/elem | Apple GPU Native |
+--------+------+--------+-----------+-------------+------------------+
| FP32   |  32  | 1e38   | ~7 digits | 4 bytes     | Yes              |
| FP16   |  16  | 65504  | ~3 digits | 2 bytes     | Yes (2x rate)    |
| BF16   |  16  | 1e38   | ~2 digits | 2 bytes     | M3+ partial      |
| INT8   |   8  | -128   | exact     | 1 byte      | Yes              |
|        |      | to 127 | integers  |             |                  |
| INT4   |   4  | -8     | exact     | 0.5 bytes   | No (emulated)    |
|        |      | to 7   | integers  |             |                  |
+--------+------+--------+-----------+-------------+------------------+

Fused Operations: Saving Bandwidth by Combining Kernels

Here is one of the most important optimization ideas in GPU programming: kernel fusion.²

The Problem: Bandwidth Waste

Consider computing SiLU activation followed by element-wise multiply (as in the LLaMA FFN):

Step 1: gate = SiLU(x * Wg)
Step 2: up = x * Wu
Step 3: hidden = gate * up

If each step is a separate kernel:

Kernel 1 (SiLU):
  Read gate_raw from global memory     (N elements read)
  Compute SiLU
  Write gate to global memory          (N elements written)

Kernel 2 (element-wise multiply):
  Read gate from global memory          (N elements read - AGAIN!)
  Read up from global memory            (N elements read)
  Compute gate * up
  Write hidden to global memory         (N elements written)

Total memory traffic: 2N reads + N writes + 2N reads + N writes = 4N reads + 2N writes

The Solution: Fused Kernel

Combine both operations into a single kernel:

Fused kernel (SiLU + multiply):
  Read gate_raw from global memory     (N elements read)
  Read up from global memory           (N elements read)
  Compute SiLU(gate_raw) * up in registers
  Write hidden to global memory        (N elements written)

Total memory traffic: 2N reads + N writes

Savings: 2N reads + N writes eliminated!

The fused kernel avoids the round trip through global memory between operations. The intermediate value (the SiLU output) lives only in registers, which are essentially free to access.

Common Fused Operations in Transformers

+------------------------------------+---------------------------------+
| Fused Operation                    | What It Combines                |
+------------------------------------+---------------------------------+
| Fused SiLU gate                    | SiLU(x) * y                    |
| Fused RMSNorm                      | norm + scale (gamma multiply)   |
| Fused attention                    | QK^T + scale + mask + softmax   |
|   (FlashAttention)                 |   + V multiply (Ch 16)          |
| Fused dequant + matmul             | Dequantize weights + GEMV/GEMM  |
| Fused residual + norm              | x + sublayer + RMSNorm          |
| Fused RoPE + reshape               | Rotary embedding + head reshape |
+------------------------------------+---------------------------------+

When Fusion Matters and When It Does Not

Fusion helps MOST when:
  - Operations are memory-bandwidth bound (small compute per element)
  - Intermediate tensors are large
  - Operations are sequential (output of one feeds input of next)

Fusion helps LEAST when:
  - Operations are compute bound (GEMM with large matrices)
  - Intermediate tensors are small (fit in cache anyway)
  - Operations have different parallelism patterns
    (e.g., can't easily fuse a reduction with an element-wise op)

The Bandwidth Equation

To understand why fusion matters, consider the numbers for an Apple M2 GPU:

M2 GPU specs:
  Memory bandwidth: ~100 GB/s
  Compute (FP16):   ~3.6 TFLOPS

For a SiLU activation on 4096 FP16 elements:
  Data to read:  4096 * 2 bytes = 8 KB
  Data to write: 4096 * 2 bytes = 8 KB
  Compute: 4096 * ~5 FLOPs (exp, add, div, mul) = ~20 KFLOPS

  Time limited by bandwidth: 16 KB / 100 GB/s = 0.16 microseconds
  Time limited by compute:   20 KFLOPS / 3.6 TFLOPS = 0.006 microseconds

  This operation is 27x more bandwidth bound than compute bound!

Every byte you can avoid reading from or writing to global memory is a direct performance win for bandwidth-bound operations. And in transformers, the majority of non-GEMM operations are bandwidth bound.

A Complete Example: Processing One Token

Let us trace the memory traffic for processing a single token through one LLaMA-7B layer, showing where fusion helps.

Model parameters (LLaMA 7B):
  d_model = 4096
  n_heads = 32
  d_k = 128
  d_ff = 11008 (intermediate FFN dimension)
  Weights in FP16

Step 1: RMSNorm
  Read: x (4096 * 2B = 8KB) + gamma (8KB) = 16KB
  Write: x_norm (8KB)
  Total: 24KB

Step 2: Q projection (GEMV: 1 x 4096 times 4096 x 4096)
  Read: x_norm (8KB) + Wq (4096*4096*2B = 32MB)
  Write: Q (8KB)
  Total: ~32MB   <-- Weight read dominates!

Step 3: K projection
  Read: x_norm (8KB) + Wk (4096*128*32*2B = ... depends on GQA)
  ... similar to Q

Step 4-6: V projection, attention, output projection
  ... similar pattern

Step 7-8: FFN up + gate + SiLU + down
  Read: Wu (4096*11008*2B = 86MB) + Wg (86MB) + Wd (86MB)
  Total: ~258MB of weight reads for FFN alone

Grand total weight reads per layer: ~400MB
For 32 layers: ~12.8GB of weight reads PER TOKEN

At 100 GB/s bandwidth: 12.8GB / 100 GB/s = 128ms per token = ~8 tokens/sec

This is why:

1. Quantization is essential – 4-bit weights cut reads by 4x.
2. Fusion matters for the non-GEMM ops – saving even a few MB per layer adds up.
3. Decode is bandwidth-bound – the GPU spends most of its time waiting for memory.

Summary

In this chapter, we have established the foundations:

• Tensors are multi-dimensional arrays. In ML, we work with rank-2 to rank-4 tensors constantly.
• Memory layout (row-major with strides) determines how multi-dimensional indices map to flat GPU buffer addresses.
• Coalesced access – adjacent threads reading adjacent memory – is critical for GPU performance.
• Common operations include element-wise (trivially parallel), reductions (require thread cooperation), and matrix multiplications (require tiling).
• Mixed precision (FP16 values, FP32 accumulation) gives us the best of both worlds.
• Kernel fusion eliminates wasteful memory round-trips between operations.

The next chapter zooms in on the single most important operation: matrix multiplication. It consumes 90%+ of inference compute, and getting it right is the difference between a usable inference engine and a toy.

1. Apple. “Metal Shading Language Specification, v3.1.” developer.apple.com. Covers the half type, vector types, and precision guarantees for FP16 arithmetic on Apple GPUs. See https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf. ↩

2. Micikevicius, P., et al. “Mixed Precision Training.” ICLR 2018. Establishes the practice of using FP16 for computation with FP32 accumulation to maintain numerical stability. See https://arxiv.org/abs/1710.03740. ↩

Matrix Multiplication Strategies

If there is one operation you need to understand deeply when running neural networks on GPUs, it is matrix multiplication. Not because it is conceptually hard – you learned it in linear algebra class – but because virtually everything in a transformer reduces to it. Every linear layer, every projection, every feed-forward network… all matmuls. If your matmul is slow, your model is slow. Period.

In this chapter we are going to build up from the textbook definition of matrix multiplication, examine why it matters so much for inference, and then dive into the different strategies you will encounter when targeting Metal GPUs. We will cover the two main regimes – GEMV and GEMM – plus all the messy in-between cases, and finally discuss how quantization changes the picture.

Why Matmul Is the Center of the Universe

Let us start with a concrete example. Consider a single linear layer in a transformer:

y = x * W + b

Here x is your input (say, a vector of dimension 4096), W is the weight matrix (say, 4096 x 11008), and b is a bias vector. That multiplication x * W is a matrix multiplication. Now count how many linear layers exist in a single transformer block:

• Q projection: matmul
• K projection: matmul
• V projection: matmul
• Output projection: matmul
• Gate projection (FFN): matmul
• Up projection (FFN): matmul
• Down projection (FFN): matmul

That is seven matmuls per layer. A 32-layer model has 224 matmuls per forward pass. If each one takes 0.5ms, you are already at 112ms just for the matmuls – and that is before attention, normalization, or anything else.

  One Transformer Block
  =====================

  Input
    |
    v
  [LayerNorm]
    |
    +---> Q Projection  (matmul #1)
    +---> K Projection  (matmul #2)
    +---> V Projection  (matmul #3)
    |
    v
  [Attention Mechanism]
    |
    v
  [Output Projection]    (matmul #4)
    |
    v
  [Residual Add]
    |
    v
  [LayerNorm]
    |
    +---> Gate Projection (matmul #5)
    +---> Up Projection   (matmul #6)
    |
    v
  [SiLU + Elementwise Mul]
    |
    v
  [Down Projection]       (matmul #7)
    |
    v
  [Residual Add]
    |
    v
  Output

So when someone says “optimizing inference,” they mostly mean “optimizing matmul.”

The Two Regimes: GEMV vs GEMM

Here is the critical insight that changes everything about how you write GPU kernels for LLM inference: the shape of the matmul depends on what phase of inference you are in.

Token Generation (Decode): GEMV

During autoregressive decoding, you generate one token at a time. Your input x is a single vector of dimension d_model. The weight matrix W has shape [d_model, d_out]. So the multiplication is:

x:  [1, 4096]
W:  [4096, 11008]
y:  [1, 11008]

This is a matrix-vector multiplication, or GEMV (General Matrix-Vector multiply). In BLAS terminology, M=1 (only one row in the output). The defining characteristic of GEMV is that every element of the weight matrix is read exactly once. There is no data reuse. This makes GEMV fundamentally memory bandwidth bound – your performance is limited by how fast you can stream the weight matrix from memory.

  GEMV: One vector times a matrix
  ================================

  x = [x0, x1, x2, x3]    (1 x K)

       W = [ w00  w01  w02 ]
           [ w10  w11  w12 ]    (K x N)
           [ w20  w21  w22 ]
           [ w30  w31  w32 ]

  y = [y0, y1, y2]         (1 x N)

  y0 = x0*w00 + x1*w10 + x2*w20 + x3*w30
  y1 = x0*w01 + x1*w11 + x2*w21 + x3*w31
  y2 = x0*w02 + x1*w12 + x2*w22 + x3*w32

  Each weight is touched exactly once.
  Bottleneck = memory bandwidth.

Prompt Processing (Prefill): GEMM

During prefill, you process the entire input prompt at once. If the prompt has 512 tokens, your input is a matrix of shape [512, 4096]. Now the multiplication becomes:

X:  [512, 4096]
W:  [4096, 11008]
Y:  [512, 11008]

This is a full matrix-matrix multiplication, or GEMM (General Matrix-Matrix multiply). M=512, and now every element of W gets reused across 512 different input rows. This changes the arithmetic intensity dramatically – GEMM can be compute bound if you organize the data access well.

  GEMM: A matrix times a matrix
  ===============================

  X = [ x00 x01 x02 x03 ]
      [ x10 x11 x12 x13 ]    (M x K)
      [ x20 x21 x22 x23 ]

       W = [ w00  w01 ]
           [ w10  w11 ]       (K x N)
           [ w20  w21 ]
           [ w30  w31 ]

  Y = [ y00 y01 ]
      [ y10 y11 ]             (M x N)
      [ y20 y21 ]

  Each weight wij is used M times (once per row of X).
  With good tiling, this becomes compute-bound.

The Arithmetic Intensity Spectrum

Let us quantify this. Arithmetic intensity is the ratio of compute operations to memory bytes transferred:

  GEMV (M=1):
    Operations = 2 * K * N  (multiply + add for each element)
    Bytes read = K * N * sizeof(weight) + K * sizeof(input)
    Intensity  ~ 2 / sizeof(weight)
    For FP16:    2 / 2 = 1 FLOP/byte    --> bandwidth bound

  GEMM (M=512):
    Operations = 2 * M * K * N
    Bytes read = K * N * sizeof(weight) + M * K * sizeof(input)
    Intensity  ~ 2 * M / sizeof(weight)  (when K*N >> M*K)
    For FP16:    2 * 512 / 2 = 512 FLOP/byte  --> compute bound

Metal GPUs typically have 200-400 GB/s of memory bandwidth and 10-20 TFLOPS of FP16 compute. The crossover point – where you shift from bandwidth bound to compute bound – is roughly around M=16-64, depending on the specific hardware and data types.

  Arithmetic Intensity vs M
  =========================

  Intensity
  (FLOP/byte)
      |
  512 |                                          *  GEMM (M=512)
      |                                       *
      |                                    *
      |                                 *
      |                              *
   64 |                           *  <-- Compute bound above this line
      |                        *
      | - - - - - - - - - - -*- - - - - - HW Compute/BW ratio
   16 |                  *
      |               *
    4 |            *
    2 |         *
    1 |  *   *    <-- GEMV (M=1)
      +--+---+---+---+---+---+---+---+--->  M
         1   2   4   8  16  32  64  512

GEMV: Thread and SIMD Parallelism

Let us start with GEMV since it is the bread and butter of token generation. Remember, M=1, so we are computing y = x * W where x is a vector and W is a matrix. The output y has N elements, and each output element is a dot product of length K.

The Naive Approach

The simplest approach: assign one thread per output element. Each thread computes a full dot product of length K:

// Naive GEMV: one thread per output element
kernel void gemv_naive(
    device const float* x      [[buffer(0)]],   // [K]
    device const float* W      [[buffer(1)]],   // [K x N], row-major
    device float*       y      [[buffer(2)]],   // [N]
    uint                tid    [[thread_position_in_grid]])
{
    float sum = 0.0f;
    for (uint k = 0; k < K; k++) {
        sum += x[k] * W[k * N + tid];  // Column tid of W
    }
    y[tid] = sum;
}

This works but has terrible performance. Each thread reads all K elements of x (lots of redundant reads) and walks down a column of W with stride N (terrible memory access pattern, no coalescing).

Coalesced Access with Transposed Weights

First improvement: store W in column-major order (or equivalently, store W^T in row-major order). Now adjacent threads read adjacent memory locations:

// GEMV with transposed weights: coalesced reads
kernel void gemv_transposed(
    device const float* x      [[buffer(0)]],   // [K]
    device const float* Wt     [[buffer(1)]],   // [N x K], W transposed
    device float*       y      [[buffer(2)]],   // [N]
    uint                tid    [[thread_position_in_grid]])
{
    // Thread tid computes output element tid
    // Reads row tid of Wt, which is column tid of W
    float sum = 0.0f;
    for (uint k = 0; k < K; k++) {
        sum += x[k] * Wt[tid * K + k];
    }
    y[tid] = sum;
}

Better memory access pattern, but we still have N threads each independently reading the entire x vector. Let us fix that.

SIMD-Level Parallelism: Splitting the K Dimension

Here is the key idea for efficient GEMV on Metal: instead of one thread per output, use a group of threads (a SIMD group) to cooperatively compute one output element. Each thread in the SIMD group handles a slice of the K dimension, then we use simd_sum to reduce:

  SIMD GEMV: 32 threads cooperate on one dot product
  ====================================================

  x = [x0, x1, x2, ..., x31, x32, ..., x63, ...]

  Thread 0 handles: x[0]*w[0] + x[32]*w[32] + x[64]*w[64] + ...
  Thread 1 handles: x[1]*w[1] + x[33]*w[33] + x[65]*w[65] + ...
  ...
  Thread 31 handles: x[31]*w[31] + x[63]*w[63] + x[95]*w[95] + ...

  Then: simd_sum() adds all 32 partial sums in hardware
  Result: one output element computed cooperatively

// SIMD GEMV: one SIMD group per output element
kernel void gemv_simd(
    device const float* x      [[buffer(0)]],
    device const float* Wt     [[buffer(1)]],   // [N x K]
    device float*       y      [[buffer(2)]],
    constant uint&      K      [[buffer(3)]],
    uint                simd_lane  [[thread_index_in_simdgroup]],
    uint                simd_gid   [[simdgroup_index_in_threadgroup]],
    uint                tg_id      [[threadgroup_position_in_grid]])
{
    // Each SIMD group computes one output element
    uint n = tg_id * SIMD_GROUPS_PER_TG + simd_gid;
    if (n >= N) return;

    float partial = 0.0f;
    // Each lane processes every 32nd element
    for (uint k = simd_lane; k < K; k += 32) {
        partial += x[k] * Wt[n * K + k];
    }

    // Hardware reduction across the SIMD group
    float total = simd_sum(partial);

    // Only lane 0 writes the result
    if (simd_lane == 0) {
        y[n] = total;
    }
}

The simd_sum intrinsic is a hardware-level reduction. On Metal, a SIMD group (also called a “wave” on other platforms) is 32 threads that execute in lockstep. The simd_sum operation sums a value across all 32 lanes in just a few clock cycles – no shared memory, no barriers, no atomics. This is enormously powerful.

  simd_sum reduction (hardware butterfly)
  ========================================

  Lane:  0    1    2    3   ...  30   31
  Val:  2.1  0.5  1.3  0.8 ...  0.2  1.1

  Step 1: swap with neighbor, add
         2.6  2.6  2.1  2.1 ...  1.3  1.3

  Step 2: swap with stride-2 neighbor, add
         4.7  4.7  4.7  4.7 ...  ...  ...

  Step 3: stride 4 ...
  Step 4: stride 8 ...
  Step 5: stride 16 ...

  After 5 steps: all lanes hold the total sum.
  Cost: ~5 cycles. No memory traffic.

Multiple Output Elements per SIMD Group

We can do even better. Instead of one output element per SIMD group, compute several. This amortizes the cost of reading x:

// SIMD GEMV: one SIMD group computes 4 output elements
kernel void gemv_simd_multi(
    device const float* x      [[buffer(0)]],
    device const float* Wt     [[buffer(1)]],
    device float*       y      [[buffer(2)]],
    constant uint&      K      [[buffer(3)]],
    uint                simd_lane  [[thread_index_in_simdgroup]],
    uint                simd_gid   [[simdgroup_index_in_threadgroup]],
    uint                tg_id      [[threadgroup_position_in_grid]])
{
    uint n_base = (tg_id * SIMD_GROUPS_PER_TG + simd_gid) * 4;

    float4 partial = float4(0.0f);

    for (uint k = simd_lane; k < K; k += 32) {
        float xk = x[k];  // Load x[k] once, use for all 4 outputs
        partial[0] += xk * Wt[(n_base + 0) * K + k];
        partial[1] += xk * Wt[(n_base + 1) * K + k];
        partial[2] += xk * Wt[(n_base + 2) * K + k];
        partial[3] += xk * Wt[(n_base + 3) * K + k];
    }

    // Reduce each output element
    float y0 = simd_sum(partial[0]);
    float y1 = simd_sum(partial[1]);
    float y2 = simd_sum(partial[2]);
    float y3 = simd_sum(partial[3]);

    if (simd_lane == 0) {
        y[n_base + 0] = y0;
        y[n_base + 1] = y1;
        y[n_base + 2] = y2;
        y[n_base + 3] = y3;
    }
}

Now each SIMD group loads x[k] once and uses it four times. The weight reads are still one-shot (each weight used exactly once), but we have reduced the x reads by 4x.

GEMM: Tiling for Compute Efficiency

When M is large (prefill phase), we enter the world of GEMM. The key technique here is tiling: dividing the output matrix into tiles and assigning each tile to a threadgroup. Within the threadgroup, we further divide work among SIMD groups.

The Tiling Concept

  GEMM Tiling Overview
  =====================

  Output Y [M x N] is divided into tiles of size [Bm x Bn]

       N
  <----------->
  +----+----+----+----+  ^
  | TG | TG | TG | TG |  |
  | 00 | 01 | 02 | 03 |  |
  +----+----+----+----+  | M
  | TG | TG | TG | TG |  |
  | 10 | 11 | 12 | 13 |  |
  +----+----+----+----+  v

  Each TG computes a [Bm x Bn] tile of the output.

  To compute its tile, TG needs:
  - A strip of X: [Bm x K]
  - A strip of W: [K x Bn]

  But K can be huge (4096+), so we tile along K too:

       K
  <----------->
  +====+====+====+  <- X rows for this TG's tile
  | Bk | Bk | Bk |
  +====+====+====+

  Process K in chunks of Bk, accumulating partial results.

Cooperative Loading into Threadgroup Memory

The key insight is that threads within a threadgroup can cooperatively load a tile of data into fast threadgroup memory (shared memory), then all threads read from it. This converts slow device memory reads into fast threadgroup memory reads:

// Simplified tiled GEMM
kernel void gemm_tiled(
    device const half* X        [[buffer(0)]],  // [M x K]
    device const half* W        [[buffer(1)]],  // [K x N]
    device half*       Y        [[buffer(2)]],  // [M x N]
    constant uint&     M        [[buffer(3)]],
    constant uint&     N        [[buffer(4)]],
    constant uint&     K        [[buffer(5)]],
    uint2              tg_pos   [[threadgroup_position_in_grid]],
    uint               tid      [[thread_index_in_threadgroup]])
{
    // Tile dimensions
    const uint Bm = 32;  // Tile rows
    const uint Bn = 32;  // Tile cols
    const uint Bk = 16;  // K-tile size

    // Threadgroup memory for tiles
    threadgroup half X_tile[Bm * Bk];
    threadgroup half W_tile[Bk * Bn];

    // This threadgroup computes output tile at (tg_pos.y * Bm, tg_pos.x * Bn)
    uint m_base = tg_pos.y * Bm;
    uint n_base = tg_pos.x * Bn;

    // Accumulator for this thread's output elements
    float acc[4] = {0, 0, 0, 0};  // Each thread computes 4 elements

    // Walk along K in steps of Bk
    for (uint k_base = 0; k_base < K; k_base += Bk) {

        // === Cooperative load: all threads load tiles into shared memory ===
        // (simplified -- real code distributes load across all threads)
        if (tid < Bm * Bk) {
            uint row = tid / Bk;
            uint col = tid % Bk;
            X_tile[tid] = X[(m_base + row) * K + (k_base + col)];
        }
        if (tid < Bk * Bn) {
            uint row = tid / Bn;
            uint col = tid % Bn;
            W_tile[tid] = W[(k_base + row) * N + (n_base + col)];
        }

        threadgroup_barrier(mem_flags::mem_threadgroup);

        // === Compute: multiply tiles ===
        // Each thread computes its assigned output elements
        // using data from threadgroup memory (fast!)
        uint my_m = tid / 8;   // Which row within tile
        uint my_n = (tid % 8) * 4;  // Which 4 columns within tile

        for (uint kk = 0; kk < Bk; kk++) {
            half x_val = X_tile[my_m * Bk + kk];
            for (uint j = 0; j < 4; j++) {
                acc[j] += float(x_val) * float(W_tile[kk * Bn + my_n + j]);
            }
        }

        threadgroup_barrier(mem_flags::mem_threadgroup);
    }

    // Write results
    uint my_m = tid / 8;
    uint my_n = (tid % 8) * 4;
    for (uint j = 0; j < 4; j++) {
        Y[(m_base + my_m) * N + (n_base + my_n + j)] = half(acc[j]);
    }
}

Let us trace through what happens for one K-tile:

  One iteration of the K-loop (k_base = 0, Bk = 16)
  ===================================================

  Step 1: Cooperative Load
  ~~~~~~~~~~~~~~~~~~~~~~~~
  All 256 threads in the threadgroup work together to load:

  X_tile [32 x 16]:                    W_tile [16 x 32]:
  Rows m_base..m_base+31               Rows 0..15
  Cols 0..15 of X                      Cols n_base..n_base+31 of W

  Thread 0 loads X[m_base][0..15] row  Thread 128 loads W[0][n_base..+31]
  Thread 1 loads X[m_base+1][0..15]    Thread 129 loads W[1][n_base..+31]
  ...                                   ...

  Step 2: Barrier
  ~~~~~~~~~~~~~~~
  Wait for all loads to complete.

  Step 3: Compute
  ~~~~~~~~~~~~~~~
  Each thread reads from the fast threadgroup memory tiles
  and accumulates partial results.

  Thread 0 computes: Y[m_base+0][n_base+0..3] += X_tile[0][k] * W_tile[k][0..3]
  Thread 1 computes: Y[m_base+0][n_base+4..7] += X_tile[0][k] * W_tile[k][4..7]
  ...
  Thread 8 computes: Y[m_base+1][n_base+0..3] += X_tile[1][k] * W_tile[k][0..3]
  ...

  Step 4: Barrier
  ~~~~~~~~~~~~~~~
  Wait before overwriting tiles with next K-chunk.

SIMD Group Matrix Multiply-Accumulate (MMA)

Metal 3.0+ on Apple Silicon supports SIMD group matrix operations that are much faster than manually computing the multiply-accumulate. These map to the hardware’s matrix multiplication units (the AMX/matrix coprocessor):

#include <metal_simdgroup_matrix>

// Using SIMD group MMA for tiled GEMM
kernel void gemm_simdgroup_mma(
    device const half* X        [[buffer(0)]],
    device const half* W        [[buffer(1)]],
    device half*       Y        [[buffer(2)]],
    constant uint&     M        [[buffer(3)]],
    constant uint&     N        [[buffer(4)]],
    constant uint&     K        [[buffer(5)]],
    uint2              tg_pos   [[threadgroup_position_in_grid]],
    uint               sgid     [[simdgroup_index_in_threadgroup]],
    uint               lane     [[thread_index_in_simdgroup]])
{
    // Each SIMD group computes an 8x8 tile of the output
    // using 8x8 matrix operations

    // Declare SIMD group matrices
    simdgroup_matrix<half, 8, 8> x_mat;
    simdgroup_matrix<half, 8, 8> w_mat;
    simdgroup_matrix<half, 8, 8> acc_mat;

    // Initialize accumulator to zero
    simdgroup_fill(acc_mat, half(0));

    uint m_base = tg_pos.y * 32 + (sgid / 4) * 8;
    uint n_base = tg_pos.x * 32 + (sgid % 4) * 8;

    for (uint k = 0; k < K; k += 8) {
        // Load 8x8 tiles from X and W
        simdgroup_load(x_mat, X + m_base * K + k, K);
        simdgroup_load(w_mat, W + k * N + n_base, N);

        // 8x8 matrix multiply-accumulate in hardware
        simdgroup_multiply_accumulate(acc_mat, x_mat, w_mat, acc_mat);
    }

    // Store the 8x8 result tile
    simdgroup_store(acc_mat, Y + m_base * N + n_base, N);
}

The simdgroup_multiply_accumulate function computes an 8x8 matrix multiply in hardware. A single SIMD group of 32 threads cooperatively holds the 8x8 matrices (each thread holds 2 elements) and the multiplication happens in the dedicated matrix hardware. This is dramatically faster than doing the multiply-add manually.

  SIMD Group MMA: 32 threads own an 8x8 matrix
  ===============================================

  The 8x8 = 64 elements are distributed across 32 threads:
  Thread 0:  elements [0,0] and [0,1]
  Thread 1:  elements [0,2] and [0,3]
  ...
  Thread 15: elements [1,6] and [1,7]
  Thread 16: elements [2,0] and [2,1]
  ...
  Thread 31: elements [3,6] and [3,7]

  simdgroup_multiply_accumulate(C, A, B, C):
    C += A * B    (all 8x8, in ~1-2 cycles)

  This maps to Apple's matrix coprocessor hardware.

Wide GEMV: The Vocabulary Projection Problem

There is one particular GEMV that deserves special attention: the final vocabulary projection. In a typical LLM, this multiplies the hidden state (dimension 4096) by the vocabulary embedding matrix to produce logits over the entire vocabulary:

hidden: [1, 4096]
vocab_weights: [4096, 32000]    (or 128256 for Llama 3!)
logits: [1, 32000]

This is still a GEMV (M=1), but N is enormous – 32K to 128K output elements. The challenge is that you need enough parallelism to saturate the GPU, and you need to read a very large weight matrix.

The strategy is to parallelize aggressively across the N dimension:

  Wide GEMV for Vocab Projection
  ===============================

  N = 32000 output logits
  K = 4096

  Approach: assign SIMD groups to output elements

  Threadgroup 0:    SIMD groups compute outputs [0..127]
  Threadgroup 1:    SIMD groups compute outputs [128..255]
  ...
  Threadgroup 249:  SIMD groups compute outputs [31872..31999]

  With 250 threadgroups x 4 SIMD groups/TG x 32 threads/SIMD = 32000 threads
  Each SIMD group: 32 threads reduce a dot product of length 4096
  That is 4096/32 = 128 multiply-adds per thread.

  Total weight data: 4096 * 32000 * 2 bytes = 256 MB (FP16)
  At 200 GB/s: ~1.3ms to stream all weights
  This sets the floor for GEMV performance.

The key insight: for wide GEMV, you want each SIMD group to handle one (or a few) output columns, with the 32 lanes splitting the K-dimension reduction. The x vector is small enough to fit in threadgroup memory or even registers, so you load it once and reuse it for all output columns in the threadgroup.

// Wide GEMV: optimized for large N (vocab projection)
kernel void gemv_wide(
    device const half*  x       [[buffer(0)]],   // [K]
    device const half*  Wt      [[buffer(1)]],   // [N x K]
    device half*        y       [[buffer(2)]],   // [N]
    constant uint&      K       [[buffer(3)]],
    constant uint&      N       [[buffer(4)]],
    uint                sgid    [[simdgroup_index_in_threadgroup]],
    uint                lane    [[thread_index_in_simdgroup]],
    uint                tg_id   [[threadgroup_position_in_grid]])
{
    // Load x into threadgroup memory (cooperative load)
    threadgroup half x_shared[4096];  // Assumes K <= 4096
    uint tid = sgid * 32 + lane;
    for (uint i = tid; i < K; i += 256) {  // 256 threads per TG
        x_shared[i] = x[i];
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);

    // Each SIMD group handles one output element
    uint n = tg_id * 8 + sgid;  // 8 SIMD groups per TG
    if (n >= N) return;

    float sum = 0.0f;
    device const half* w_row = Wt + n * K;

    for (uint k = lane; k < K; k += 32) {
        sum += float(x_shared[k]) * float(w_row[k]);
    }

    sum = simd_sum(sum);
    if (lane == 0) {
        y[n] = half(sum);
    }
}

Batched GEMV: The Awkward Middle Ground

In practice, inference is not always purely M=1 or M=512. There are several scenarios where M is small but greater than 1:

1. Batch decoding: serving multiple users simultaneously (M = batch_size, often 2-16)
2. Speculative decoding: verifying multiple candidate tokens at once (M = num_candidates, often 4-8)
3. Small prompts: short prefill with just a few tokens

For these cases, M is too small for efficient GEMM (not enough reuse to be compute bound) but too large for pure GEMV (we want some reuse of the weight data).

The solution is batched GEMV: treat it as M independent GEMVs but share weight loads across them:

  Batched GEMV (M=4)
  ===================

  X = [ x0 ]     (4 rows, each a vector of length K)
      [ x1 ]
      [ x2 ]
      [ x3 ]

  Strategy: Each SIMD group handles multiple output columns
  and ALL M rows simultaneously.

  SIMD Group 0 processing output column n=0:
  +-----+-----+-----+-----+-----+
  | Lane|  0  |  1  |  2  | ... |
  +-----+-----+-----+-----+-----+
  | Row0| x0[0]*w[0] | x0[1]*w[1] | x0[2]*w[2] | ... |
  | Row1| x1[0]*w[0] | x1[1]*w[1] | x1[2]*w[2] | ... |
  | Row2| x2[0]*w[0] | x2[1]*w[1] | x2[2]*w[2] | ... |
  | Row3| x3[0]*w[0] | x3[1]*w[1] | x3[2]*w[2] | ... |
  +-----+-----+-----+-----+-----+

  Each weight w[k] is loaded once and used for M=4 rows.
  Weight reuse factor: 4x compared to pure GEMV.

// Batched GEMV: M rows share weight loads
kernel void gemv_batched(
    device const half*  X       [[buffer(0)]],   // [M x K]
    device const half*  Wt      [[buffer(1)]],   // [N x K]
    device half*        Y       [[buffer(2)]],   // [M x N]
    constant uint&      M       [[buffer(3)]],
    constant uint&      K       [[buffer(4)]],
    constant uint&      N       [[buffer(5)]],
    uint                sgid    [[simdgroup_index_in_threadgroup]],
    uint                lane    [[thread_index_in_simdgroup]],
    uint                tg_id   [[threadgroup_position_in_grid]])
{
    uint n = tg_id * 8 + sgid;  // Output column
    if (n >= N) return;

    // Accumulators for each row
    float sums[8] = {0};  // Support up to M=8

    device const half* w_row = Wt + n * K;

    for (uint k = lane; k < K; k += 32) {
        half w_val = w_row[k];  // Load weight ONCE

        // Apply to all M rows
        for (uint m = 0; m < M; m++) {
            sums[m] += float(X[m * K + k]) * float(w_val);
        }
    }

    // Reduce each row's sum
    for (uint m = 0; m < M; m++) {
        float total = simd_sum(sums[m]);
        if (lane == 0) {
            Y[m * N + n] = half(total);
        }
    }
}

The performance gain over running M separate GEMVs is significant because:

1. Weight data is loaded once instead of M times
2. Kernel launch overhead is reduced
3. The GPU can be better saturated with the additional work

  Performance: M separate GEMVs vs Batched GEMV
  ===============================================

  M=4, K=4096, N=4096, FP16 weights

  M separate GEMVs:
    Weight reads: 4 * 4096 * 4096 * 2 bytes = 128 MB
    Time at 200 GB/s: ~0.64 ms

  Batched GEMV:
    Weight reads: 1 * 4096 * 4096 * 2 bytes = 32 MB
    Time at 200 GB/s: ~0.16 ms

  4x improvement from weight reuse!

How Quantization Changes the Game

Everything we have discussed so far assumes FP16 weights. But in practice, most inference deployments use quantized weights – 4-bit or even 2-bit. This changes the matmul kernels significantly.

The Core Challenge: Dequantize on the Fly

Quantized weights are stored in a compressed format. You cannot directly multiply them with the input – you first need to dequantize them back to a floating-point representation. The question is: where and when do you dequantize?

The answer is on-the-fly dequantization: each thread dequantizes just the weight values it needs, right before multiplying them with the input. The dequantized values live only in registers, never written back to memory.

  Dequantize-on-the-fly GEMV
  ===========================

  Memory:  [...Q4 packed weights...]  (4 bits per weight)
                    |
                    v  (load)
  Registers: [packed_byte]
                    |
                    v  (unpack + scale)
  Registers: [fp16_weight_a, fp16_weight_b]
                    |
                    v  (multiply with input)
  Registers: [partial_sum]
                    |
                    v  (simd_sum)
  Memory:   [output]

  Key: dequantized weights never touch memory.
  We trade compute (dequantize) for bandwidth (smaller reads).

Q4_0 GEMV Example

Let us work through a concrete example with Q4_0 quantization (the simplest format). In Q4_0, every block of 32 weights shares a single FP16 scale factor. Each weight is stored as a 4-bit integer (0-15), representing the range [-8, 7] after subtracting 8:

// Q4_0 block structure
struct block_q4_0 {
    half    scale;         // 2 bytes: shared scale for 32 weights
    uchar   packed[16];   // 16 bytes: 32 x 4-bit values packed into pairs
};
// Total: 18 bytes for 32 weights (4.5 bits/weight effective)

// Q4_0 GEMV: dequantize on the fly
kernel void gemv_q4_0(
    device const half*         x       [[buffer(0)]],  // [K]
    device const block_q4_0*   W       [[buffer(1)]],  // [N x K/32] blocks
    device half*               y       [[buffer(2)]],  // [N]
    constant uint&             K       [[buffer(3)]],
    uint                       sgid    [[simdgroup_index_in_threadgroup]],
    uint                       lane    [[thread_index_in_simdgroup]],
    uint                       tg_id   [[threadgroup_position_in_grid]])
{
    uint n = tg_id * 8 + sgid;
    uint num_blocks = K / 32;

    float sum = 0.0f;

    // Each lane processes every 32nd block
    for (uint b = lane; b < num_blocks; b += 32) {
        // Load the quantized block
        device const block_q4_0& blk = W[n * num_blocks + b];
        half scale = blk.scale;

        // Dequantize and multiply all 32 weights in this block
        for (uint j = 0; j < 16; j++) {
            uchar packed = blk.packed[j];

            // Unpack two 4-bit values
            int8_t w0 = (packed & 0x0F) - 8;  // Low nibble
            int8_t w1 = (packed >> 4)   - 8;   // High nibble

            // Dequantize: weight = scale * quantized_value
            float dq0 = float(scale) * float(w0);
            float dq1 = float(scale) * float(w1);

            // Multiply with input
            uint k_idx = b * 32 + j * 2;
            sum += dq0 * float(x[k_idx]);
            sum += dq1 * float(x[k_idx + 1]);
        }
    }

    sum = simd_sum(sum);
    if (lane == 0) {
        y[n] = half(sum);
    }
}

The Bandwidth Win

The reason quantization is so important for GEMV performance is straightforward arithmetic. Since GEMV is bandwidth bound, reducing the data size directly speeds things up:

  Bandwidth Savings with Quantization
  ====================================

  Weight matrix: 4096 x 4096 = 16.7M weights

  FP16:  16.7M * 2 bytes   = 33.6 MB    (baseline)
  Q8_0:  16.7M * 1.06 bytes = 17.7 MB   (1.9x faster)
  Q4_0:  16.7M * 0.56 bytes =  9.4 MB   (3.6x faster)
  Q4_K:  16.7M * 0.57 bytes =  9.5 MB   (3.5x faster)
  Q2_K:  16.7M * 0.34 bytes =  5.7 MB   (5.9x faster)

  At 200 GB/s memory bandwidth:
  FP16:  0.168 ms per layer
  Q4_0:  0.047 ms per layer  <-- 3.6x speedup!

  Decode speed for 7B model:
  FP16:  14.0 GB / 200 GB/s = 70 ms/token  -> 14 tok/s
  Q4_0:   3.9 GB / 200 GB/s = 19 ms/token  -> 52 tok/s

The compute cost of dequantization is negligible compared to the bandwidth savings. A few shifts and multiplies per weight value are cheap when the alternative is waiting for twice as many bytes to arrive from memory.

Quantized GEMM: A Different Story

For GEMM (prefill), the situation is more nuanced. GEMM is compute bound, so reducing weight size does not automatically speed things up – you are not bandwidth limited in the first place. However, quantized GEMM is still useful because:

1. It reduces memory footprint, letting larger models fit in memory
2. For moderate M values (16-64), you might still be bandwidth bound
3. Apple’s AMX hardware has limited support for mixed-precision MMA

The typical approach is to dequantize tiles of weights into threadgroup memory in FP16, then use the standard SIMD group MMA operations on the dequantized tiles:

  Quantized GEMM: Dequantize into Threadgroup Memory
  ====================================================

  Device Memory:          Threadgroup Memory:       Registers:
  +-----------+           +----------+              +--------+
  | Q4 Weights|  --load-->| FP16 Tile|  --MMA-->   | Accum  |
  | (compact) |           | (32x32)  |              | (8x8)  |
  +-----------+           +----------+              +--------+

  1. Cooperatively load Q4 blocks from device memory
  2. Dequantize to FP16 in threadgroup memory
  3. Use simdgroup_multiply_accumulate on FP16 tiles
  4. Repeat for all K-tiles

Worked Example: Full Thread Assignment

Let us trace through a complete example to see how threads are assigned for a real GEMV. Consider:

  Problem: y = x * W
  x:  [1, 4096]   (one token's hidden state)
  W:  [4096, 4096] (Q projection weight)
  y:  [1, 4096]   (query vector)

  Configuration:
  - Threadgroup size: 256 threads = 8 SIMD groups x 32 lanes
  - Each SIMD group: 1 output element
  - Grid: 4096 / 8 = 512 threadgroups

  Memory layout: W stored transposed as Wt [4096 x 4096]

Let us follow Threadgroup 0, SIMD group 3, Lane 17:

  Thread Identity
  ===============
  Threadgroup:  0
  SIMD group:   3
  Lane:         17

  Output assignment
  =================
  output index n = threadgroup * 8 + simd_group = 0 * 8 + 3 = 3
  This thread helps compute y[3].

  Work assignment
  ===============
  Lane 17 processes every 32nd element along K:
    k = 17, 49, 81, 113, ..., 4081

  Total iterations: 4096 / 32 = 128

  For each iteration (e.g., k=17):
    partial_sum += x[17] * Wt[3 * 4096 + 17]

  After all 128 iterations:
    partial_sum = x[17]*Wt[3,17] + x[49]*Wt[3,49] + ... + x[4081]*Wt[3,4081]

  Reduction
  =========
  simd_sum(partial_sum) adds all 32 lanes' partial sums:
    y[3] = sum over all k: x[k] * Wt[3, k]
         = dot(x, column 3 of W)

  Only lane 0 of SIMD group 3 writes y[3] to memory.

  Summary for entire kernel
  =========================
  512 threadgroups * 8 SIMD groups = 4096 dot products
  Each SIMD group: 32 threads * 128 iterations = 4096 multiply-adds
  Total: 4096 * 4096 = 16.7M multiply-adds (matches 2*K*N/2 operations)

Choosing the Right Strategy

Here is a decision tree for choosing the right matmul strategy:

  Matmul Strategy Decision Tree
  ==============================

                     What is M?
                        |
          +-------------+-------------+
          |             |             |
        M = 1         2-16         > 16
          |             |             |
        GEMV      Batched GEMV      GEMM
          |             |             |
     +----+----+   +----+----+   +----+----+
     |         |   |         |   |         |
   Small N  Large N           |   Use SIMD
   (4096)  (32K+)          Share   Group
     |         |          weight    MMA
   1 SIMD   Lots of       loads  with tiling
   group    threadgroups          |
   per out                   +----+----+
     |                       |         |
   Standard               FP16     Quantized
   SIMD GEMV              tiles    (dequant
                          + MMA    into TG
                                   memory)

And here are rough performance expectations on Apple M-series GPUs:

  Performance Expectations (M2 Pro, ~200 GB/s)
  =============================================

  Operation                    Size           FP16       Q4_0
  ---------                    ----           ----       ----
  GEMV  (decode, 1 layer)     [1,4096]*[4096,4096]     0.17ms    0.05ms
  GEMV  (vocab projection)    [1,4096]*[4096,32000]    1.28ms    0.36ms
  Batched GEMV (M=4)          [4,4096]*[4096,4096]     0.17ms    0.06ms
  GEMM  (prefill, M=512)      [512,4096]*[4096,4096]   ~2ms      ~2ms

  Note: Batched GEMV with M=4 is barely slower than M=1 GEMV because
  the same weight data is read -- only the compute increases.
  GEMM times are similar for FP16 and Q4 because GEMM is compute-bound.

Summary

Matrix multiplication is the heart of neural network inference. On Metal GPUs, you need different strategies depending on the shape:

1. GEMV (M=1, decode): Bandwidth bound. Split the K dimension across SIMD lanes, reduce with simd_sum. Optimize for coalesced memory access and maximize bandwidth utilization.

2. Batched GEMV (M=2-16): Still bandwidth bound but with some weight reuse. Load weights once, apply to all M rows. Significant speedup over M separate GEMVs.

3. GEMM (M>16, prefill): Can be compute bound with proper tiling. Use cooperative loading into threadgroup memory and SIMD group MMA for maximum throughput.

4. Quantized variants: For GEMV, quantization directly translates to speedup via reduced bandwidth. Dequantize on-the-fly in registers. For GEMM, dequantize into threadgroup memory tiles, then use standard MMA.

The next chapter will look at where these matmuls are used in the context of attention – where Q, K, and V come from matmuls, and then the attention computation itself introduces a different kind of matrix multiplication.

The Attention Mechanism

If matrix multiplication is the workhorse of transformer inference, attention is the brain. It is the mechanism that allows each token to look at every other token in the sequence and decide what information is relevant. Without attention, a transformer is just a stack of feed-forward networks – powerful but unable to model relationships between positions in a sequence.

In this chapter, we are going to take attention apart piece by piece. We will start with the mathematical formulation, walk through a complete numerical example, explore multi-head and grouped-query variants, understand the KV cache, and analyze the computational complexity. By the end, you will understand every detail of how attention works, which will set us up perfectly for the FlashAttention optimization in the next chapter.

The Core Idea: Queries, Keys, and Values

Attention is fundamentally a lookup mechanism. Think of it like a fuzzy dictionary lookup: you have a query (“what am I looking for?”), a set of keys (“what is available?”), and corresponding values (“what information do those keys hold?”). The twist is that instead of an exact match, you compute a similarity score between the query and every key, then return a weighted combination of the values.

  Attention as Fuzzy Dictionary Lookup
  =====================================

  Query: "What comes after 'the'?"

  Keys:            Values:           Similarity:
  ----             ------            -----------
  "the"    --->    [0.2, 0.8, ...]     0.85  (high -- relevant!)
  "cat"    --->    [0.5, 0.1, ...]     0.72  (medium)
  "sat"    --->    [0.3, 0.6, ...]     0.15  (low)
  "on"     --->    [0.1, 0.4, ...]     0.62  (medium)

  Output = 0.85 * val("the") + 0.72 * val("cat") + 0.15 * val("sat") + ...
           (after normalizing similarities with softmax)

But where do queries, keys, and values come from? They are all produced by linear projections of the input embeddings.

Self-Attention: Q, K, V Projections

Given an input sequence of token embeddings X with shape [seq_len, d_model], we produce Q, K, and V through three separate linear projections:

Q = X * W_Q    where W_Q has shape [d_model, d_k]
K = X * W_K    where W_K has shape [d_model, d_k]
V = X * W_V    where W_V has shape [d_model, d_v]

In the original transformer paper and most modern LLMs, d_k = d_v = d_model. Each projection is a matrix multiplication – exactly the kind we optimized in the previous chapter.

  Self-Attention Projections
  ===========================

  Input X: [seq_len x d_model]

  X ----+-----> [W_Q] -----> Q  [seq_len x d_k]
        |
        +-----> [W_K] -----> K  [seq_len x d_k]
        |
        +-----> [W_V] -----> V  [seq_len x d_v]

  "Self" attention means Q, K, V all come from the SAME input X.
  This is in contrast to cross-attention where Q comes from one
  source and K, V come from another (used in encoder-decoder models).

The Attention Equation

Once we have Q, K, and V, the attention computation is:

Attention(Q, K, V) = softmax(Q * K^T / sqrt(d_k)) * V

Let us break this down step by step.

Step 1: Compute Attention Scores (Q * K^T)

We multiply Q by the transpose of K. If Q has shape [seq_len, d_k] and K has shape [seq_len, d_k], then Q * K^T has shape [seq_len, seq_len]. Each element (i, j) is the dot product of query vector i with key vector j – a measure of how much token i should attend to token j.

  Q * K^T  -->  Attention Score Matrix
  ======================================

  Q = [ q0 ]     K^T = [ k0^T  k1^T  k2^T  k3^T ]
      [ q1 ]
      [ q2 ]     (each qi, kj is a vector of dimension d_k)
      [ q3 ]

  Q * K^T = [ q0.k0  q0.k1  q0.k2  q0.k3 ]
            [ q1.k0  q1.k1  q1.k2  q1.k3 ]
            [ q2.k0  q2.k1  q2.k2  q2.k3 ]
            [ q3.k0  q3.k1  q3.k2  q3.k3 ]

  Element (i,j) = dot product of query i with key j
                = how much should token i attend to token j?

Step 2: Scale by sqrt(d_k)

We divide every score by sqrt(d_k). Why? Without scaling, the dot products grow in magnitude proportional to the dimension. If d_k = 64, the dot products might be around 64 in magnitude. When you pass large numbers into softmax, the output becomes extremely peaked – almost all the weight goes to one key, and the gradients vanish. Dividing by sqrt(64) = 8 keeps the values in a reasonable range.

  Why Scale?
  ==========

  Without scaling (d_k = 64):
    Typical dot product magnitude: ~64
    softmax([64, 2, 1, 3]) = [1.000, 0.000, 0.000, 0.000]
    --> Almost a hard lookup. Gradients vanish.

  With scaling:
    Scaled dot product: ~64/8 = ~8
    softmax([8.0, 0.25, 0.125, 0.375]) = [0.997, 0.001, 0.001, 0.001]
    --> Still peaked, but gradients can flow.

  The scaling keeps variance of dot products at ~1.0
  regardless of dimension, assuming inputs have unit variance.

Step 3: Apply Softmax

The softmax function converts the scaled scores into a probability distribution over keys for each query position. Each row of the score matrix is independently softmaxed:

softmax(z_i) = exp(z_i) / sum_j(exp(z_j))

After softmax, each row sums to 1.0, and all values are between 0 and 1. These are the attention weights.

  Softmax: Scores --> Weights
  ============================

  Before softmax (one row, after scaling):
  [ 2.1,  0.5,  -0.3,  1.2 ]

  exp each:
  [ 8.17, 1.65,  0.74,  3.32 ]

  sum = 13.88

  Divide by sum:
  [ 0.589, 0.119, 0.053, 0.239 ]
    ^                      ^
    Token 0 gets most     Token 3 gets
    attention              second most

  Sum = 1.0 (it is a probability distribution)

Step 4: Weighted Sum of Values

Finally, we multiply the attention weights by V. Each output vector is a weighted combination of value vectors, where the weights come from softmax:

  Weighted Value Sum
  ===================

  Attention weights (one row): [0.589, 0.119, 0.053, 0.239]

  V = [ v0 ]   =  [ 0.1, 0.9, 0.3 ]
      [ v1 ]      [ 0.5, 0.2, 0.8 ]
      [ v2 ]      [ 0.7, 0.4, 0.1 ]
      [ v3 ]      [ 0.3, 0.6, 0.5 ]

  output = 0.589 * v0 + 0.119 * v1 + 0.053 * v2 + 0.239 * v3
         = 0.589 * [0.1, 0.9, 0.3]
         + 0.119 * [0.5, 0.2, 0.8]
         + 0.053 * [0.7, 0.4, 0.1]
         + 0.239 * [0.3, 0.6, 0.5]
         = [0.227, 0.718, 0.388]

  The output is dominated by v0 (weight 0.589) and v3 (weight 0.239).

The Full Pipeline

  Complete Attention Computation
  ===============================

  Input: X [seq_len x d_model]
    |
    +--[W_Q]--> Q [seq_len x d_k]
    +--[W_K]--> K [seq_len x d_k]
    +--[W_V]--> V [seq_len x d_v]
    |
    v
  Scores = Q * K^T              [seq_len x seq_len]
    |
    v
  Scaled = Scores / sqrt(d_k)   [seq_len x seq_len]
    |
    v
  Masked = apply causal mask     [seq_len x seq_len]
    |
    v
  Weights = softmax(Masked)      [seq_len x seq_len]
    |
    v
  Output = Weights * V           [seq_len x d_v]
    |
    v
  Projected = Output * W_O       [seq_len x d_model]

Multi-Head Attention: Divide and Conquer

In practice, we do not compute a single attention function. Instead, we split the representation into multiple heads, each operating on a smaller dimension. This allows the model to attend to information from different representation subspaces simultaneously.

If d_model = 512 and we use h = 8 heads, then each head operates on d_k = d_model / h = 64 dimensions.

  Multi-Head Attention
  =====================

  d_model = 512, h = 8 heads, d_k = 64 per head

  Input X: [seq_len x 512]

  Full Q projection: X * W_Q  -->  Q_full [seq_len x 512]

  Split into 8 heads:
  Q_full = [ Q_head0 | Q_head1 | Q_head2 | ... | Q_head7 ]
             [s x 64]   [s x 64]  [s x 64]       [s x 64]

  Similarly for K and V.

  Each head computes attention independently:
  head_i = Attention(Q_head_i, K_head_i, V_head_i)

  Concatenate results:
  MultiHead = [ head_0 | head_1 | ... | head_7 ]  [seq_len x 512]

  Final projection:
  Output = MultiHead * W_O   [seq_len x 512]

Why multiple heads? Different heads can learn different types of relationships:

  What Different Heads Might Learn
  =================================

  Head 0: Syntactic dependencies  ("the" attends to its noun)
  Head 1: Positional patterns     (attend to previous token)
  Head 2: Semantic similarity     (similar words attend to each other)
  Head 3: Coreference            ("it" attends to its referent)
  Head 4: Long-range dependencies (closing bracket attends to opening)
  ...

  Each head has its own Q, K, V projections, so each head
  learns to look for different patterns in the data.

In terms of implementation, multi-head attention is usually done by performing the full projection to get [seq_len, d_model] shaped Q, K, V, then reshaping to [seq_len, n_heads, d_k] (or equivalently [n_heads, seq_len, d_k] after transposing). The attention is then computed in parallel across all heads – on a GPU, this is a batched operation.

  Implementation: Reshape and Batch
  ==================================

  After Q projection: Q [seq_len x d_model]

  Reshape to: Q [seq_len x n_heads x d_k]
  Transpose:  Q [n_heads x seq_len x d_k]

  Now attention is a batched operation:
  For each head h in [0, n_heads):
      scores_h = Q[h] * K[h]^T        [seq_len x seq_len]
      weights_h = softmax(scores_h)     [seq_len x seq_len]
      output_h = weights_h * V[h]       [seq_len x d_v]

  This is embarrassingly parallel -- all heads computed simultaneously.

Grouped Query Attention (GQA)

Standard multi-head attention requires storing separate K and V for every head. For a model with 32 heads generating 4096-length sequences with d_k = 128, that is:

KV cache per layer = 2 * 32 * 4096 * 128 * 2 bytes = 64 MB
For 32 layers: 2 GB just for the KV cache!

Grouped Query Attention (GQA) reduces this by sharing K and V across groups of query heads. Instead of 32 KV heads, you might use only 8. Each KV head serves a group of 4 query heads.

  Multi-Head vs Grouped Query vs Multi-Query
  ============================================

  Multi-Head Attention (MHA): n_kv_heads = n_q_heads = 32
  -------------------------------------------------------
  Q heads:  Q0  Q1  Q2  Q3  Q4  Q5  ... Q31
  K heads:  K0  K1  K2  K3  K4  K5  ... K31
  V heads:  V0  V1  V2  V3  V4  V5  ... V31
  (Each Q head has its own K,V pair)


  Grouped Query Attention (GQA): n_kv_heads = 8, group_size = 4
  ---------------------------------------------------------------
  Q heads:  Q0  Q1  Q2  Q3 | Q4  Q5  Q6  Q7 | Q8 ... | Q28 Q29 Q30 Q31
  K heads:      K0         |     K1          |  K2    |       K7
  V heads:      V0         |     V1          |  V2    |       V7
  (4 Q heads share each K,V pair)


  Multi-Query Attention (MQA): n_kv_heads = 1
  ---------------------------------------------
  Q heads:  Q0  Q1  Q2  Q3  Q4  Q5  ... Q31
  K heads:              K0
  V heads:              V0
  (All Q heads share a single K,V pair)

The memory savings are substantial:

  KV Cache Size Comparison (per layer)
  ======================================

  seq_len = 4096, d_k = 128, FP16

  MHA (32 KV heads): 2 * 32 * 4096 * 128 * 2 =  64 MB
  GQA (8 KV heads):  2 *  8 * 4096 * 128 * 2 =  16 MB  (4x reduction)
  MQA (1 KV head):   2 *  1 * 4096 * 128 * 2 =   2 MB  (32x reduction)

  For 32 layers:
  MHA:  2048 MB
  GQA:   512 MB   <-- Llama 2 70B, Llama 3
  MQA:    64 MB

In the GPU kernel, GQA changes how you index into the KV cache. For query head h, the corresponding KV head is h / group_size:

// GQA: mapping query heads to KV heads
uint q_head = head_index;                    // 0..31
uint kv_head = q_head / group_size;          // 0..7 (if group_size=4)

// Load Q from q_head's slice
device const half* q = Q + q_head * d_k;

// Load K, V from kv_head's slice (shared by 4 Q heads)
device const half* k = K_cache + kv_head * seq_len * d_k;
device const half* v = V_cache + kv_head * seq_len * d_v;

The KV Cache: Why We Store K and V

During autoregressive decoding, we generate one token at a time. At step t, we need to compute attention between the new token’s query and all previous tokens’ keys and values. Without caching, we would need to reproject all previous tokens through W_K and W_V every single step – a massive waste of computation.

The KV cache stores the K and V vectors for all previous tokens. At each decoding step, we only compute K and V for the new token and append them to the cache.

  KV Cache During Decoding
  =========================

  Step 1: Process token "The"
  +---------+
  | K: k_0  |  V: v_0
  +---------+
  Q = q_0, attend to [k_0]

  Step 2: Process token "cat"
  +---------+---------+
  | K: k_0  | K: k_1  |  V: v_0, v_1
  +---------+---------+
  Q = q_1, attend to [k_0, k_1]

  Step 3: Process token "sat"
  +---------+---------+---------+
  | K: k_0  | K: k_1  | K: k_2  |  V: v_0, v_1, v_2
  +---------+---------+---------+
  Q = q_2, attend to [k_0, k_1, k_2]

  Step 4: Process token "on"
  +---------+---------+---------+---------+
  | K: k_0  | K: k_1  | K: k_2  | K: k_3  |  V: v_0, v_1, v_2, v_3
  +---------+---------+---------+---------+
  Q = q_3, attend to [k_0, k_1, k_2, k_3]

  At each step, we compute K, V for only the NEW token
  and append to the cache. The cache grows by one entry per step.

Without KV cache:

  Step t: compute K, V for ALL t tokens, then attend
  Cost: O(t * d_model * d_k) for projections alone
  Total over T steps: O(T^2 * d_model * d_k)  -- quadratic!

With KV cache:

  Step t: compute K, V for 1 new token, append to cache
  Cost: O(d_model * d_k) for projections
  Total over T steps: O(T * d_model * d_k) for projections -- linear!
  (Attention itself is still O(T * d_k) per step, O(T^2 * d_k) total)

The KV cache is typically stored as a contiguous buffer per layer, per head, pre-allocated to the maximum sequence length:

  KV Cache Memory Layout
  =======================

  For one layer, one KV head:

  K cache: [max_seq_len x d_k]
  +------+------+------+------+------+------+------+------+
  | k_0  | k_1  | k_2  | k_3  | .... | .... | .... | .... |
  +------+------+------+------+------+------+------+------+
    ^                     ^      ^
    filled positions      |      empty (pre-allocated)
                          current position

  V cache: [max_seq_len x d_v]
  +------+------+------+------+------+------+------+------+
  | v_0  | v_1  | v_2  | v_3  | .... | .... | .... | .... |
  +------+------+------+------+------+------+------+------+

  At each decode step:
  1. Compute k_new, v_new for the new token
  2. Write k_new to K_cache[current_pos]
  3. Write v_new to V_cache[current_pos]
  4. Compute attention: q_new vs K_cache[0..current_pos]
  5. Increment current_pos

Causal Masking

In autoregressive language models, a token should only attend to tokens that came before it (or at the same position). This is called causal masking or the “no peeking into the future” constraint.

During prefill (when processing the entire prompt at once), we enforce this by adding a mask to the attention scores before softmax:

  Causal Mask (4 tokens)
  =======================

  Before masking:                After masking:
  [ 2.1  0.5  1.3  0.8 ]        [ 2.1  -inf -inf -inf ]
  [ 0.3  1.7  0.9  0.4 ]        [ 0.3   1.7 -inf -inf ]
  [ 1.1  0.2  2.3  1.5 ]        [ 1.1   0.2  2.3 -inf ]
  [ 0.7  1.4  0.6  1.9 ]        [ 0.7   1.4  0.6  1.9 ]

  Mask matrix:
  [  0   -inf -inf -inf ]
  [  0    0   -inf -inf ]
  [  0    0    0   -inf ]
  [  0    0    0    0   ]

  Token i can only see tokens 0..i (the lower triangle).
  -inf becomes 0 after softmax, so masked tokens contribute nothing.

After softmax, the -inf entries become 0, effectively removing those tokens from the weighted sum. The result is that each token’s output depends only on itself and all preceding tokens.

  After softmax with causal mask:
  ================================

  [ 1.000  0.000  0.000  0.000 ]   Token 0: only sees itself
  [ 0.197  0.803  0.000  0.000 ]   Token 1: sees tokens 0, 1
  [ 0.192  0.078  0.639  0.000 ]   Token 2: sees tokens 0, 1, 2  (skipped)
  [ 0.092  0.186  0.083  0.306 ]   Token 3: sees all tokens      (skipped)
                                                                   ^^^^^^
                                                          sum may not be 1.0
                                                          because I showed
                                                          approximate values

During decoding, causal masking is implicit – the query is for position t and the KV cache only contains positions 0..t-1, so there is nothing to mask.

Computational Complexity

Let us carefully count the operations in attention for a sequence of length n with head dimension d:

  Operation Counts for Attention
  ===============================

  1. Q * K^T:  matrix multiply [n x d] * [d x n] = [n x n]
     Operations: 2 * n * n * d  (multiply + add)

  2. Scale:    divide each of n*n elements by sqrt(d)
     Operations: n * n

  3. Softmax:  for each of n rows, compute exp and normalize over n elements
     Operations: ~3 * n * n  (exp, sum, divide)

  4. Weights * V:  matrix multiply [n x n] * [n x d] = [n x d]
     Operations: 2 * n * n * d

  Total: ~4 * n^2 * d + 4 * n^2
       = O(n^2 * d)

  For n = 4096, d = 128:
    4 * 4096^2 * 128 = ~8.6 billion operations per head
    Times 32 heads = ~274 billion operations

  Memory for attention matrix: n * n * sizeof(float) per head
    = 4096 * 4096 * 4 = 64 MB per head
    = 2 GB for 32 heads  <-- This is the problem FlashAttention solves!

The quadratic scaling in sequence length is the fundamental limitation:

  How Attention Scales with Sequence Length
  ==========================================

  n        Ops (per head)    Attn Matrix Size
  ---      ---------------   ----------------
  512         ~34M              1 MB
  1024       ~134M              4 MB
  2048       ~537M             16 MB
  4096      ~2.15B             64 MB
  8192      ~8.59B            256 MB
  16384    ~34.4B               1 GB
  32768   ~137.4B               4 GB
  131072  ~2.20T               64 GB  (Llama 3.1 context length!)

  The attention matrix alone exceeds GPU memory for long sequences.
  This is why FlashAttention (next chapter) is essential.

Worked Example: 4-Token, 2-Head Attention

Let us trace through a complete attention computation with concrete numbers. We will use a tiny example: 4 tokens, 2 heads, d_model = 8, d_k = d_v = 4 per head.

Setup

  Model parameters:
    d_model = 8
    n_heads = 2
    d_k = d_v = d_model / n_heads = 4
    seq_len = 4

  Input tokens: "The cat sat on"

  Input embeddings X [4 x 8]:
  Token 0 ("The"): [ 1.0,  0.5,  0.3,  0.8, | 0.2,  0.7,  0.4,  0.6 ]
  Token 1 ("cat"): [ 0.3,  1.2,  0.7,  0.1, | 0.9,  0.4,  0.5,  0.3 ]
  Token 2 ("sat"): [ 0.6,  0.2,  1.1,  0.4, | 0.3,  0.8,  1.0,  0.2 ]
  Token 3 ("on"):  [ 0.4,  0.7,  0.5,  1.0, | 0.6,  0.3,  0.7,  0.9 ]

Step 1: Compute Q, K, V (via projection)

For simplicity, let us assume the projections have already been computed and show the result after splitting into heads:

  After projection and head split:

  Head 0 (d_k = 4):
  Q0 = [ 1.2, 0.3, 0.5, 0.8 ]    K0 = [ 0.9, 0.4, 0.7, 0.2 ]    V0 = [ 0.3, 0.8, 0.5, 0.1 ]
       [ 0.4, 1.1, 0.2, 0.6 ]         [ 0.5, 1.0, 0.3, 0.8 ]         [ 0.7, 0.2, 0.9, 0.4 ]
       [ 0.7, 0.5, 0.9, 0.3 ]         [ 0.8, 0.6, 1.1, 0.5 ]         [ 0.4, 0.6, 0.3, 0.8 ]
       [ 0.3, 0.8, 0.4, 1.0 ]         [ 0.2, 0.7, 0.5, 1.0 ]         [ 0.9, 0.5, 0.7, 0.3 ]

  Head 1 (d_k = 4):
  Q1 = [ 0.6, 0.9, 0.2, 0.4 ]    K1 = [ 0.3, 0.7, 0.5, 0.1 ]    V1 = [ 0.5, 0.4, 0.2, 0.7 ]
       [ 0.8, 0.3, 0.7, 0.5 ]         [ 0.6, 0.2, 0.8, 0.4 ]         [ 0.2, 0.9, 0.6, 0.3 ]
       [ 0.1, 0.6, 0.4, 0.8 ]         [ 0.4, 0.5, 0.3, 0.9 ]         [ 0.8, 0.3, 0.5, 0.6 ]
       [ 0.5, 0.4, 0.9, 0.7 ]         [ 0.7, 0.3, 0.6, 0.5 ]         [ 0.3, 0.7, 0.4, 0.8 ]

Step 2: Q * K^T for Head 0

  Scores0 = Q0 * K0^T    [4 x 4]

  Score(0,0) = Q0[0] . K0[0] = 1.2*0.9 + 0.3*0.4 + 0.5*0.7 + 0.8*0.2
             = 1.08 + 0.12 + 0.35 + 0.16 = 1.71

  Score(0,1) = Q0[0] . K0[1] = 1.2*0.5 + 0.3*1.0 + 0.5*0.3 + 0.8*0.8
             = 0.60 + 0.30 + 0.15 + 0.64 = 1.69

  Score(0,2) = Q0[0] . K0[2] = 1.2*0.8 + 0.3*0.6 + 0.5*1.1 + 0.8*0.5
             = 0.96 + 0.18 + 0.55 + 0.40 = 2.09

  Score(0,3) = Q0[0] . K0[3] = 1.2*0.2 + 0.3*0.7 + 0.5*0.5 + 0.8*1.0
             = 0.24 + 0.21 + 0.25 + 0.80 = 1.50

  (Computing all 16 entries similarly...)

  Scores0 = [ 1.71  1.69  2.09  1.50 ]
            [ 1.16  1.68  1.49  1.37 ]
            [ 1.38  1.22  1.80  1.08 ]
            [ 1.04  1.49  1.33  1.42 ]

Step 3: Scale

  sqrt(d_k) = sqrt(4) = 2.0

  Scaled0 = Scores0 / 2.0

  Scaled0 = [ 0.855  0.845  1.045  0.750 ]
            [ 0.580  0.840  0.745  0.685 ]
            [ 0.690  0.610  0.900  0.540 ]
            [ 0.520  0.745  0.665  0.710 ]

Step 4: Apply Causal Mask

  Masked0 = [ 0.855  -inf   -inf   -inf  ]
            [ 0.580  0.840  -inf   -inf  ]
            [ 0.690  0.610  0.900  -inf  ]
            [ 0.520  0.745  0.665  0.710 ]

Step 5: Softmax (row by row)

  Row 0: softmax([0.855, -inf, -inf, -inf])
         = [1.000, 0.000, 0.000, 0.000]

  Row 1: softmax([0.580, 0.840])  (ignoring -inf entries)
         exp: [1.786, 2.317]  sum = 4.103
         = [0.435, 0.565, 0.000, 0.000]

  Row 2: softmax([0.690, 0.610, 0.900])
         exp: [1.994, 1.840, 2.460]  sum = 6.294
         = [0.317, 0.292, 0.391, 0.000]

  Row 3: softmax([0.520, 0.745, 0.665, 0.710])
         exp: [1.682, 2.107, 1.945, 2.034]  sum = 7.768
         = [0.217, 0.271, 0.250, 0.262]

  Weights0 = [ 1.000  0.000  0.000  0.000 ]
             [ 0.435  0.565  0.000  0.000 ]
             [ 0.317  0.292  0.391  0.000 ]
             [ 0.217  0.271  0.250  0.262 ]

Step 6: Weighted Sum of Values

  Output0 = Weights0 * V0    [4 x 4]

  Output0[0] = 1.000 * V0[0]
             = [0.300, 0.800, 0.500, 0.100]

  Output0[1] = 0.435 * V0[0] + 0.565 * V0[1]
             = 0.435 * [0.3, 0.8, 0.5, 0.1] + 0.565 * [0.7, 0.2, 0.9, 0.4]
             = [0.131, 0.348, 0.218, 0.044] + [0.396, 0.113, 0.509, 0.226]
             = [0.526, 0.461, 0.726, 0.270]

  Output0[2] = 0.317 * V0[0] + 0.292 * V0[1] + 0.391 * V0[2]
             = [0.095, 0.254, 0.159, 0.032]
             + [0.204, 0.058, 0.263, 0.117]
             + [0.156, 0.235, 0.117, 0.313]
             = [0.456, 0.547, 0.539, 0.461]

  Output0[3] = 0.217 * V0[0] + 0.271 * V0[1] + 0.250 * V0[2] + 0.262 * V0[3]
             = [0.065, 0.174, 0.109, 0.022]
             + [0.190, 0.054, 0.244, 0.108]
             + [0.100, 0.150, 0.075, 0.200]
             + [0.236, 0.131, 0.183, 0.079]
             = [0.591, 0.509, 0.611, 0.409]

Step 7: Concatenate Heads and Project

Head 1 would be computed identically (with its own Q1, K1, V1). Then the two heads’ outputs are concatenated along the last dimension and projected:

  Concatenation:
  Output0[i]: [d_v = 4 values]
  Output1[i]: [d_v = 4 values]

  Concat[i] = [Output0[i] | Output1[i]]   (d_model = 8 values)

  Final = Concat * W_O    [4 x 8] * [8 x 8] = [4 x 8]

  This produces the final attention output for this layer.

Decode Attention vs Prefill Attention

The attention computation looks quite different depending on the inference phase:

  Prefill Attention (processing entire prompt)
  =============================================

  Q: [seq_len x d_k]      (all query positions)
  K: [seq_len x d_k]      (all key positions)
  V: [seq_len x d_v]      (all value positions)

  QK^T: [seq_len x seq_len]  -- full attention matrix
  Softmax: [seq_len x seq_len]
  Output: [seq_len x d_v]

  This is a GEMM problem. Matrix-matrix multiplications.
  Causal mask applied explicitly.


  Decode Attention (generating one token at a time)
  ==================================================

  q: [1 x d_k]            (single new query)
  K: [seq_len x d_k]      (all keys from KV cache)
  V: [seq_len x d_v]      (all values from KV cache)

  q * K^T: [1 x seq_len]   -- one row of attention scores
  Softmax: [1 x seq_len]
  Output: [1 x d_v]

  This is a GEMV problem. Vector-matrix multiplications.
  No causal mask needed (KV cache only has past tokens).

This distinction is critical for optimization. Prefill attention can use tiled GEMM techniques with high data reuse. Decode attention is bandwidth-bound, reading through the entire KV cache for each head.

  Decode Attention: Memory Access Pattern
  =========================================

  For one head, one decode step at position t:

  Read q:       1 * d_k * 2 bytes        =    256 bytes (d_k=128)
  Read K cache: t * d_k * 2 bytes        = varies
  Read V cache: t * d_v * 2 bytes        = varies
  Write output: 1 * d_v * 2 bytes        =    256 bytes

  At position t = 4096:
  Read K: 4096 * 128 * 2 = 1 MB per head
  Read V: 4096 * 128 * 2 = 1 MB per head
  Total per head: ~2 MB
  Total for 32 heads: ~64 MB

  At 200 GB/s: ~0.32 ms

  This grows linearly with sequence length.
  At t = 32768: ~0.5 GB read, ~2.5 ms

Putting It All Together

Here is a Metal kernel sketch showing the structure of decode attention for one head:

// Decode attention: single query vs KV cache
kernel void decode_attention(
    device const half*  q          [[buffer(0)]],  // [d_k]
    device const half*  K_cache    [[buffer(1)]],  // [seq_len x d_k]
    device const half*  V_cache    [[buffer(2)]],  // [seq_len x d_v]
    device half*        output     [[buffer(3)]],  // [d_v]
    constant uint&      seq_len    [[buffer(4)]],
    constant uint&      d_k        [[buffer(5)]],
    uint                tid        [[thread_position_in_grid]])
{
    // Step 1: Compute attention scores (q . K[i] for each i)
    float score = 0.0f;
    for (uint d = 0; d < d_k; d++) {
        score += float(q[d]) * float(K_cache[tid * d_k + d]);
    }
    score /= sqrt(float(d_k));

    // Step 2: Softmax (requires two passes or online algorithm)
    // ... (this is where FlashAttention comes in!)

    // Step 3: Weighted sum of values
    // output[d] = sum over i: weight[i] * V_cache[i * d_v + d]
    // ...
}

This sketch is deliberately incomplete – computing softmax across all positions requires either materializing the full score vector (memory expensive for long sequences) or using the online softmax algorithm that FlashAttention employs. That is exactly what we will cover in the next chapter.

Summary

The attention mechanism is the defining feature of transformers. Here are the key takeaways:

1. Q, K, V projections are matrix multiplications that transform input embeddings into queries, keys, and values.

2. The attention equation softmax(QK^T / sqrt(d_k)) * V computes a weighted combination of values based on query-key similarity.

3. Multi-head attention splits the representation into independent heads that can learn different attention patterns.

4. Grouped Query Attention shares KV heads across multiple query heads, reducing KV cache memory by the group size factor.

5. The KV cache stores previously computed keys and values during decoding, avoiding redundant recomputation. It grows linearly with sequence length.

6. Causal masking ensures tokens only attend to past positions, enforced via -inf masking before softmax during prefill.

7. Complexity is O(n^2 * d) – quadratic in sequence length. The attention matrix itself requires O(n^2) memory, which becomes prohibitive for long sequences.

8. Prefill vs decode have very different computational profiles: prefill is a GEMM problem (compute bound), decode is a GEMV problem (bandwidth bound).

In the next chapter, we will see how FlashAttention solves the O(n^2) memory problem through a clever tiling scheme and online softmax computation.

FlashAttention: Theory and Practice

In the previous chapter, we saw that standard attention requires materializing an [n x n] attention matrix. For a sequence length of 4096, that is 64 MB per head. For 32 heads, 2 GB. For 32,768 tokens (Llama 3’s context), 32 GB just for attention weights. This is clearly untenable.

FlashAttention is the algorithmic breakthrough that makes long-context attention practical. The key idea is beautifully simple: never materialize the full attention matrix. Instead, compute attention in tiles, maintaining a running softmax that produces exactly the same result as the standard algorithm – bit for bit.

This chapter will build up the theory from first principles, work through a complete numerical example, and then show how FlashAttention adapts for both prefill and decode phases, including advanced variants for GQA and speculative decoding.

The Problem: O(n^2) Memory

Let us be very concrete about what “materializing the attention matrix” means in the standard algorithm:

  Standard Attention Memory Usage
  ================================

  Step 1: S = Q * K^T         Allocate [n x n] matrix S
  Step 2: S = S / sqrt(d)     In-place on S
  Step 3: P = softmax(S)      Allocate [n x n] matrix P (or overwrite S)
  Step 4: O = P * V           Allocate [n x d] matrix O

  Peak memory: at least one [n x n] matrix in memory simultaneously

  n = 4096, FP32:
    [n x n] = 4096 * 4096 * 4 bytes = 64 MB (per head!)
    32 heads = 2 GB

  n = 32768, FP32:
    [n x n] = 32768 * 32768 * 4 bytes = 4 GB (per head!)
    32 heads = 128 GB  <-- exceeds any GPU's memory

Even if memory were unlimited, there is a performance problem. The attention matrix is written to device memory in step 1/3 and read back in step 3/4. Those round-trips to device memory are slow:

  Memory Traffic in Standard Attention
  =====================================

  Write S to memory:    n^2 * 4 bytes      (after Q * K^T)
  Read S from memory:   n^2 * 4 bytes      (for softmax)
  Write P to memory:    n^2 * 4 bytes      (softmax output)
  Read P from memory:   n^2 * 4 bytes      (for P * V)

  Total memory traffic: 4 * n^2 * 4 bytes = 16 * n^2 bytes

  At n = 4096: 16 * 16M = 256 MB of memory traffic per head
  At 200 GB/s: 1.3 ms just for the attention matrix I/O (per head!)

FlashAttention eliminates this memory traffic entirely by keeping everything in registers and threadgroup memory (SRAM).

The Key Insight: Online Softmax

The reason we need the full [n x n] matrix is softmax. Softmax requires knowing the maximum value across an entire row (for numerical stability) and the sum of exponentials (for normalization). Both of these seem to require seeing all n values before producing any output.

The breakthrough is online softmax: a streaming algorithm that processes values one chunk at a time, maintaining running statistics that can be corrected as new data arrives.

Standard Softmax (Two-Pass)

The textbook softmax for a vector z of length n:

  Pass 1: Find maximum
    m = max(z[0], z[1], ..., z[n-1])

  Pass 2: Compute exponentials and sum
    For i in 0..n-1:
      e[i] = exp(z[i] - m)
    s = sum(e[0], e[1], ..., e[n-1])

  Output: softmax[i] = e[i] / s

  Requires storing all n values of z to make two passes.

Online Softmax (One-Pass, Streaming)

Online softmax processes elements one at a time (or one tile at a time), maintaining a running maximum m and a running sum s. When a new maximum is found, the previously accumulated sum is corrected:

  Initialize:
    m = -infinity    (running maximum)
    s = 0            (running sum of exp)

  For each new value z[i]:
    m_new = max(m, z[i])
    s = s * exp(m - m_new) + exp(z[i] - m_new)
    m = m_new

  After all values:
    softmax[i] = exp(z[i] - m) / s

  The key line: s = s * exp(m - m_new) + exp(z[i] - m_new)

  When m_new > m (new max found):
    exp(m - m_new) < 1, so the old sum is scaled DOWN
    This corrects for the fact that previous exponentials
    were computed relative to the old (smaller) maximum

  When m_new = m (no new max):
    exp(m - m_new) = exp(0) = 1, so old sum unchanged
    Just add the new exponential

Let us trace through an example:

  Online Softmax Example: z = [2.0, 5.0, 1.0, 4.0]
  ===================================================

  Step 0: m = -inf, s = 0

  Step 1: z[0] = 2.0
    m_new = max(-inf, 2.0) = 2.0
    s = 0 * exp(-inf - 2.0) + exp(2.0 - 2.0)
      = 0 + exp(0)
      = 1.0
    m = 2.0

  Step 2: z[1] = 5.0
    m_new = max(2.0, 5.0) = 5.0       <-- New max found!
    s = 1.0 * exp(2.0 - 5.0) + exp(5.0 - 5.0)
      = 1.0 * exp(-3.0) + exp(0)
      = 1.0 * 0.0498 + 1.0
      = 1.0498
    m = 5.0

    Verification: at this point, s should equal exp(2-5) + exp(5-5) = 0.0498 + 1 = 1.0498  ✓
    The correction factor exp(2.0 - 5.0) = 0.0498 adjusts the old sum.

  Step 3: z[2] = 1.0
    m_new = max(5.0, 1.0) = 5.0       <-- No new max
    s = 1.0498 * exp(5.0 - 5.0) + exp(1.0 - 5.0)
      = 1.0498 * 1.0 + exp(-4.0)
      = 1.0498 + 0.0183
      = 1.0681
    m = 5.0

  Step 4: z[3] = 4.0
    m_new = max(5.0, 4.0) = 5.0       <-- No new max
    s = 1.0681 * exp(5.0 - 5.0) + exp(4.0 - 5.0)
      = 1.0681 * 1.0 + exp(-1.0)
      = 1.0681 + 0.3679
      = 1.4360
    m = 5.0

  Final softmax:
    softmax[0] = exp(2.0 - 5.0) / 1.4360 = 0.0498 / 1.4360 = 0.0347
    softmax[1] = exp(5.0 - 5.0) / 1.4360 = 1.0000 / 1.4360 = 0.6964
    softmax[2] = exp(1.0 - 5.0) / 1.4360 = 0.0183 / 1.4360 = 0.0128
    softmax[3] = exp(4.0 - 5.0) / 1.4360 = 0.3679 / 1.4360 = 0.2562

  Sum = 0.0347 + 0.6964 + 0.0128 + 0.2562 = 1.0001 ≈ 1.0  ✓

Online Softmax with Running Output Correction

But FlashAttention needs more than just softmax values – it needs softmax(S) * V, where V is the value matrix. We need to maintain a running weighted sum of V that gets corrected each time the maximum changes:

  Online Attention (combining softmax + V multiply):
  ====================================================

  Initialize:
    m = -inf          (running max)
    s = 0             (running sum of exp)
    o = 0             (running output, vector of size d_v)

  For each position j (or tile of positions):
    score = q . K[j]                       (dot product)
    m_new = max(m, score)
    correction = exp(m - m_new)

    // Correct the old output and sum
    o = o * correction
    s = s * correction

    // Add new contribution
    weight = exp(score - m_new)
    o = o + weight * V[j]
    s = s + weight

    m = m_new

  After all positions:
    output = o / s    (normalize by total sum)

This is the core of FlashAttention. Let us trace through it:

  Online Attention Example
  =========================

  q = [1, 0]  (query, dimension 2 for simplicity)

  K = [ 0.5,  0.3 ]    V = [ 1.0,  0.0 ]
      [ 0.8, -0.2 ]        [ 0.0,  1.0 ]
      [ 0.1,  0.7 ]        [ 0.5,  0.5 ]

  Scores = q . K[j]:
    s0 = 1*0.5 + 0*0.3 = 0.5
    s1 = 1*0.8 + 0*(-0.2) = 0.8
    s2 = 1*0.1 + 0*0.7 = 0.1

  (Skip scaling by sqrt(d) for clarity)

  Step 0: m = -inf, s = 0, o = [0, 0]

  Step 1: Process position 0  (score = 0.5)
    m_new = max(-inf, 0.5) = 0.5
    correction = exp(-inf - 0.5) = 0
    o = [0,0] * 0 = [0, 0]
    s = 0 * 0 = 0
    weight = exp(0.5 - 0.5) = 1.0
    o = [0,0] + 1.0 * [1.0, 0.0] = [1.0, 0.0]
    s = 0 + 1.0 = 1.0
    m = 0.5

  Step 2: Process position 1  (score = 0.8)
    m_new = max(0.5, 0.8) = 0.8           <-- New max!
    correction = exp(0.5 - 0.8) = exp(-0.3) = 0.7408
    o = [1.0, 0.0] * 0.7408 = [0.7408, 0.0]
    s = 1.0 * 0.7408 = 0.7408
    weight = exp(0.8 - 0.8) = 1.0
    o = [0.7408, 0.0] + 1.0 * [0.0, 1.0] = [0.7408, 1.0]
    s = 0.7408 + 1.0 = 1.7408
    m = 0.8

  Step 3: Process position 2  (score = 0.1)
    m_new = max(0.8, 0.1) = 0.8           <-- No new max
    correction = exp(0.8 - 0.8) = 1.0
    o = [0.7408, 1.0] * 1.0 = [0.7408, 1.0]
    s = 1.7408 * 1.0 = 1.7408
    weight = exp(0.1 - 0.8) = exp(-0.7) = 0.4966
    o = [0.7408, 1.0] + 0.4966 * [0.5, 0.5] = [0.7408+0.2483, 1.0+0.2483]
      = [0.9891, 1.2483]
    s = 1.7408 + 0.4966 = 2.2374
    m = 0.8

  Normalize: output = o / s = [0.9891/2.2374, 1.2483/2.2374]
                             = [0.4421, 0.5579]

  Verification (standard method):
    scores = [0.5, 0.8, 0.1]
    softmax = exp([0.5, 0.8, 0.1] - 0.8) / sum
            = [exp(-0.3), exp(0), exp(-0.7)] / sum
            = [0.7408, 1.0, 0.4966] / 2.2374
            = [0.3311, 0.4470, 0.2219]
    output = 0.3311*[1,0] + 0.4470*[0,1] + 0.2219*[0.5,0.5]
           = [0.3311, 0] + [0, 0.4470] + [0.1110, 0.1110]
           = [0.4421, 0.5580]  ✓ (matches!)

The Tiling Strategy

Online softmax lets us process keys/values one at a time. But processing one position at a time would be very slow – too little parallelism. FlashAttention processes keys and values in tiles of size Bk (typically 32 or 64).

  FlashAttention Tiling
  ======================

  Instead of materializing the full [n x n] attention matrix:

  Standard:
  +--------------------------------------------------+
  |                                                    |
  |            Full [n x n] matrix in memory           |
  |                                                    |
  +--------------------------------------------------+
  Memory: O(n^2)


  FlashAttention:
  Process K/V in tiles, never store more than [Bq x Bk]:

  +------+
  | Tile |  Bq x Bk     Only this small tile is in memory
  +------+  at a time!

  Iterate over all K/V tiles:

  K/V:  [====Tile 0====][====Tile 1====][====Tile 2====]...
             Bk cols         Bk cols         Bk cols

  For each tile:
  1. Load Bk keys and values into threadgroup memory
  2. Compute Bq x Bk scores (in registers)
  3. Update running max, sum, and output (online softmax)
  4. Move to next tile

  Memory: O(Bq * Bk) for the score tile
        + O(Bq * d) for the running output
        = O(n) total (since Bq, Bk, d are all constants)

The Full Algorithm

Here is the complete FlashAttention algorithm for prefill (multiple query positions):

  FlashAttention Algorithm (Prefill)
  ====================================

  Inputs: Q [n x d], K [n x d], V [n x d]
  Output: O [n x d]
  Tile sizes: Bq (query tile), Bk (key/value tile)

  For each query tile q_tile in [0, n/Bq):
    Load Q_tile = Q[q_tile*Bq : (q_tile+1)*Bq]    // [Bq x d]

    Initialize per-row:
      m[Bq] = -infinity      // running max for each query in tile
      s[Bq] = 0              // running sum for each query in tile
      O_tile[Bq x d] = 0    // running output for each query in tile

    For each KV tile kv_tile in [0, n/Bk):

      // --- Skip if causally masked ---
      // If ALL queries in q_tile come BEFORE all keys in kv_tile,
      // then all scores are masked to -inf. Skip entirely.
      if (q_tile + 1) * Bq <= kv_tile * Bk:
        continue

      Load K_tile = K[kv_tile*Bk : (kv_tile+1)*Bk]    // [Bk x d]
      Load V_tile = V[kv_tile*Bk : (kv_tile+1)*Bk]    // [Bk x d]

      // --- Compute score tile ---
      S_tile = Q_tile * K_tile^T / sqrt(d)              // [Bq x Bk]

      // --- Apply causal mask within tile ---
      For each (i, j) in S_tile:
        if (q_tile*Bq + i) < (kv_tile*Bk + j):
          S_tile[i][j] = -infinity

      // --- Online softmax update ---
      For each row i in [0, Bq):
        m_new = max(m[i], max(S_tile[i]))
        correction = exp(m[i] - m_new)

        // Correct old accumulators
        O_tile[i] *= correction
        s[i] *= correction

        // Add new contributions
        For each j in [0, Bk):
          weight = exp(S_tile[i][j] - m_new)
          O_tile[i] += weight * V_tile[j]
          s[i] += weight

        m[i] = m_new

    // --- Normalize ---
    For each row i in [0, Bq):
      O_tile[i] /= s[i]

    Write O_tile to O[q_tile*Bq : (q_tile+1)*Bq]

Let us visualize how the tiles march through the attention matrix:

  Tiling Visualization (n=16, Bq=4, Bk=4)
  ==========================================

  Full attention matrix [16 x 16] with causal mask:

       KV tiles:  0    1    2    3
              +----+----+----+----+
  Q tile 0:  | XX |    |    |    |    XX = computed tile
              +----+----+----+----+    .. = skipped (causal)
  Q tile 1:  | XX | XX |    |    |
              +----+----+----+----+
  Q tile 2:  | XX | XX | XX |    |
              +----+----+----+----+
  Q tile 3:  | XX | XX | XX | XX |
              +----+----+----+----+

  Processing order for Q tile 2:
  1. Load Q[8:12]                    (4 queries)
  2. Load K[0:4], V[0:4]            -> compute 4x4 scores, update
  3. Load K[4:8], V[4:8]            -> compute 4x4 scores, update
  4. Load K[8:12], V[8:12]          -> compute 4x4 scores (with mask), update
  5. Skip K[12:16] (causally masked) -> all queries < all keys
  6. Normalize and write O[8:12]

  At no point do we have more than a 4x4 score tile in memory.
  Memory: O(16 * d) for O instead of O(16 * 16) for the full matrix.

Memory Analysis: O(n) Instead of O(n^2)

Let us carefully count the memory used by FlashAttention:

  FlashAttention Memory Usage
  ============================

  Threadgroup memory (SRAM):
    Q_tile:      Bq * d * sizeof(half)     = 32 * 128 * 2  =   8 KB
    K_tile:      Bk * d * sizeof(half)     = 32 * 128 * 2  =   8 KB
    V_tile:      Bk * d * sizeof(half)     = 32 * 128 * 2  =   8 KB
    S_tile:      Bq * Bk * sizeof(float)   = 32 * 32 * 4   =   4 KB
    Total SRAM:                                             ~  28 KB

  Per-query state (registers):
    m[Bq]:       Bq * sizeof(float)        = 32 * 4        = 128 bytes
    s[Bq]:       Bq * sizeof(float)        = 32 * 4        = 128 bytes
    O_tile:      Bq * d * sizeof(float)    = 32 * 128 * 4  =  16 KB
    Total registers:                                        ~  16 KB

  Device memory:
    Input Q, K, V:  3 * n * d * sizeof(half)               = O(n * d)
    Output O:       n * d * sizeof(half)                    = O(n * d)
    NO attention matrix stored!                             = O(n)

  Comparison for n = 32768, d = 128:
    Standard:      n^2 * 4 bytes = 4 GB per head
    FlashAttention: n * d * 2 bytes = 8 MB per head   (500x reduction!)

A Complete Numerical Example with Tiling

Let us trace FlashAttention on a small example: n=6, d=2, Bq=2, Bk=3.

  Setup
  ======

  Q = [ 1.0,  0.5 ]     K = [ 0.3,  0.7 ]     V = [ 1.0,  0.0 ]
      [ 0.8, -0.1 ]         [ 0.6,  0.2 ]         [ 0.0,  1.0 ]
      [ 0.2,  0.9 ]         [-0.1,  0.8 ]         [ 0.5,  0.5 ]
      [-0.3,  0.4 ]         [ 0.4, -0.3 ]         [ 0.8,  0.2 ]
      [ 0.7,  0.6 ]         [ 0.9,  0.1 ]         [ 0.3,  0.7 ]
      [ 0.1, -0.5 ]         [ 0.2,  0.5 ]         [ 0.6,  0.4 ]

  sqrt(d) = sqrt(2) = 1.414

  Tile sizes: Bq = 2 (queries), Bk = 3 (keys)

Q tile 0 (queries 0-1), KV tile 0 (keys 0-2):

  Q_tile = [ 1.0,  0.5 ]     K_tile = [ 0.3,  0.7 ]     V_tile = [ 1.0,  0.0 ]
           [ 0.8, -0.1 ]              [ 0.6,  0.2 ]              [ 0.0,  1.0 ]
                                       [-0.1,  0.8 ]              [ 0.5,  0.5 ]

  S_tile = Q_tile * K_tile^T / 1.414:

  Q[0].K[0] = 1.0*0.3 + 0.5*0.7 = 0.65    /1.414 = 0.460
  Q[0].K[1] = 1.0*0.6 + 0.5*0.2 = 0.70    /1.414 = 0.495
  Q[0].K[2] = 1.0*(-0.1)+0.5*0.8 = 0.30   /1.414 = 0.212
  Q[1].K[0] = 0.8*0.3+(-0.1)*0.7 = 0.17   /1.414 = 0.120
  Q[1].K[1] = 0.8*0.6+(-0.1)*0.2 = 0.46   /1.414 = 0.325
  Q[1].K[2] = 0.8*(-0.1)+(-0.1)*0.8=-0.16 /1.414 = -0.113

  S_tile = [ 0.460   0.495   0.212 ]
           [ 0.120   0.325  -0.113 ]

  Causal mask (query 0 sees key 0 only, query 1 sees keys 0-1):
  S_tile = [ 0.460  -inf    -inf   ]
           [ 0.120   0.325  -inf   ]

  Online softmax update (initial m=-inf, s=0, O=0):

  Row 0:
    m_new = max(-inf, 0.460) = 0.460
    correction = exp(-inf - 0.460) = 0
    O[0] = [0,0]*0 = [0, 0]
    s[0] = 0*0 = 0
    weight for key 0: exp(0.460 - 0.460) = 1.0
    O[0] += 1.0 * V[0] = [1.0, 0.0]
    s[0] += 1.0 = 1.0
    (keys 1,2 masked: weight = exp(-inf) = 0)
    m[0] = 0.460

  Row 1:
    m_new = max(-inf, max(0.120, 0.325)) = 0.325
    correction = 0
    weight for key 0: exp(0.120 - 0.325) = exp(-0.205) = 0.8146
    weight for key 1: exp(0.325 - 0.325) = 1.0
    O[1] = 0.8146*[1,0] + 1.0*[0,1] = [0.8146, 1.0]
    s[1] = 0.8146 + 1.0 = 1.8146
    m[1] = 0.325

Q tile 0 (queries 0-1), KV tile 1 (keys 3-5):

  K_tile = [ 0.4, -0.3 ]     V_tile = [ 0.8,  0.2 ]
           [ 0.9,  0.1 ]              [ 0.3,  0.7 ]
           [ 0.2,  0.5 ]              [ 0.6,  0.4 ]

  But wait -- query 0 is at position 0, and keys 3-5 are all in the future.
  Query 1 is at position 1, and keys 3-5 are also in the future.
  ALL entries would be masked to -inf.

  We can skip this entire KV tile! (Causal skip optimization)

After processing all KV tiles for Q tile 0, normalize:

  O[0] = [1.0, 0.0] / 1.0 = [1.0, 0.0]
  O[1] = [0.8146, 1.0] / 1.8146 = [0.449, 0.551]

This is exactly what standard attention would produce (query 0 only sees key 0, so it copies V[0]; query 1 sees keys 0-1 with softmax weights).

The remaining Q tiles (1 and 2) would be processed similarly, each iterating over the relevant KV tiles.

Decode Variant: FlashDecoding

During decoding, we have a single query (M=1) attending to the entire KV cache. The FlashAttention algorithm simplifies because there is only one query:

  FlashDecoding (single query)
  =============================

  q: [1 x d]                     (single query vector)
  K_cache: [seq_len x d]         (all past keys)
  V_cache: [seq_len x d]         (all past values)

  Tiling: split KV cache into tiles of Bk positions

  +-------+-------+-------+-------+-------+
  | Tile0 | Tile1 | Tile2 | Tile3 | Tile4 |   K/V cache
  +-------+-------+-------+-------+-------+
   0..31   32..63  64..95  96..127  ...

  Each tile is processed by one threadgroup (or SIMD group).
  All tiles can be processed IN PARALLEL since each
  produces an independent (m, s, o) tuple.

  Then a final reduction step merges the per-tile results.

The parallel structure is different from prefill FlashAttention, where the KV tiles must be processed sequentially (because the running softmax state is updated incrementally). In decode, each tile independently produces a partial result:

  Parallel FlashDecoding
  =======================

  Phase 1: Each tile independently computes (m_t, s_t, o_t)

  Tile 0: q vs K[0:32]   -> (m_0, s_0, o_0)
  Tile 1: q vs K[32:64]  -> (m_1, s_1, o_1)
  Tile 2: q vs K[64:96]  -> (m_2, s_2, o_2)
  ...

  Phase 2: Merge all tiles' results

  merge(tile_i, tile_j):
    m_new = max(m_i, m_j)
    corr_i = exp(m_i - m_new)
    corr_j = exp(m_j - m_new)
    s_new = s_i * corr_i + s_j * corr_j
    o_new = (o_i * s_i * corr_i + o_j * s_j * corr_j) / s_new
    (or equivalently, keep unnormalized and normalize at the end)

  This merge is associative! Can be done in a tree reduction.

// FlashDecoding: Phase 1 -- per-tile computation
kernel void flash_decode_phase1(
    device const half*  q          [[buffer(0)]],    // [d]
    device const half*  K_cache    [[buffer(1)]],    // [seq_len x d]
    device const half*  V_cache    [[buffer(2)]],    // [seq_len x d]
    device float*       tile_max   [[buffer(3)]],    // [num_tiles]
    device float*       tile_sum   [[buffer(4)]],    // [num_tiles]
    device float*       tile_out   [[buffer(5)]],    // [num_tiles x d]
    constant uint&      seq_len    [[buffer(6)]],
    constant uint&      d          [[buffer(7)]],
    uint                lane       [[thread_index_in_simdgroup]],
    uint                sgid       [[simdgroup_index_in_threadgroup]],
    uint                tg_id      [[threadgroup_position_in_grid]])
{
    const uint Bk = 32;
    uint k_start = tg_id * Bk;
    uint k_end = min(k_start + Bk, seq_len);

    // Load q into registers (small enough)
    // (In practice, load into threadgroup memory cooperatively)

    float m = -INFINITY;
    float s = 0.0f;
    float o[128] = {0};  // Assuming d <= 128

    for (uint k = k_start; k < k_end; k++) {
        // Compute score = q . K[k] / sqrt(d)
        float score = 0.0f;
        for (uint i = lane; i < d; i += 32) {
            score += float(q[i]) * float(K_cache[k * d + i]);
        }
        score = simd_sum(score);
        score /= sqrt(float(d));

        // Online softmax update
        float m_new = max(m, score);
        float correction = exp(m - m_new);
        float weight = exp(score - m_new);

        s = s * correction + weight;
        for (uint i = lane; i < d; i += 32) {
            o[i] = o[i] * correction + weight * float(V_cache[k * d + i]);
        }
        m = m_new;
    }

    // Write per-tile results
    if (lane == 0) {
        tile_max[tg_id] = m;
        tile_sum[tg_id] = s;
    }
    for (uint i = lane; i < d; i += 32) {
        tile_out[tg_id * d + i] = o[i];
    }
}

// FlashDecoding: Phase 2 -- merge tile results
kernel void flash_decode_phase2(
    device const float* tile_max   [[buffer(0)]],    // [num_tiles]
    device const float* tile_sum   [[buffer(1)]],    // [num_tiles]
    device const float* tile_out   [[buffer(2)]],    // [num_tiles x d]
    device half*        output     [[buffer(3)]],    // [d]
    constant uint&      num_tiles  [[buffer(4)]],
    constant uint&      d          [[buffer(5)]],
    uint                lane       [[thread_index_in_simdgroup]])
{
    float m = -INFINITY;
    float s = 0.0f;
    float o[128] = {0};

    for (uint t = 0; t < num_tiles; t++) {
        float m_t = tile_max[t];
        float s_t = tile_sum[t];

        float m_new = max(m, m_t);
        float corr_old = exp(m - m_new);
        float corr_new = exp(m_t - m_new);

        // Merge sums
        s = s * corr_old + s_t * corr_new;

        // Merge outputs
        for (uint i = lane; i < d; i += 32) {
            o[i] = o[i] * corr_old + tile_out[t * d + i] * corr_new;
        }

        m = m_new;
    }

    // Normalize and write
    for (uint i = lane; i < d; i += 32) {
        output[i] = half(o[i] / s);
    }
}

Tree Masking for Speculative Decoding

Speculative decoding is an optimization where a smaller “draft” model proposes several tokens ahead, and the larger model verifies them all at once. The verification step requires attention with a tree mask instead of the standard causal mask.

In speculative decoding, the draft model might produce a tree of candidates:

  Speculative Decoding Tree
  ==========================

  Position:  0  1  2  3  4  5  6  7  8
  Token:     A  B  C  D  E  F  G  H  I
  Parent:    -  0  1  1  2  2  3  3  4

  Tree structure:
       A
       |
       B
      / \
     C   D
    /|   |\
   E  F  G  H
   |
   I

  Causal mask for this tree (1 = can attend, 0 = masked):

       A  B  C  D  E  F  G  H  I
  A  [ 1  0  0  0  0  0  0  0  0 ]
  B  [ 1  1  0  0  0  0  0  0  0 ]
  C  [ 1  1  1  0  0  0  0  0  0 ]
  D  [ 1  1  0  1  0  0  0  0  0 ]
  E  [ 1  1  1  0  1  0  0  0  0 ]
  F  [ 1  1  1  0  0  1  0  0  0 ]
  G  [ 1  1  0  1  0  0  1  0  0 ]
  H  [ 1  1  0  1  0  0  0  1  0 ]
  I  [ 1  1  1  0  1  0  0  0  1 ]

  Note: D can see A and B (its ancestors) but NOT C (sibling).
  E can see A, B, C (ancestors along its path) but NOT D, F, G, H.

FlashAttention handles tree masks by passing the mask as an additional input. During the tiled computation, the mask is consulted for each score tile:

  Tree Mask in FlashAttention
  ============================

  For each score tile S_tile[Bq x Bk]:
    For each (i, j):
      query_pos = q_tile * Bq + i
      key_pos   = kv_tile * Bk + j
      if mask[query_pos][key_pos] == 0:
        S_tile[i][j] = -infinity

  The mask is typically stored as a bitmask for efficiency:
    1 bit per (query, key) pair
    For n=256: 256 * 256 / 8 = 8 KB

  With standard causal mask, this is implicit (just compare positions).
  With tree mask, we need the explicit bitmask.

GQA Handling in FlashAttention

When using Grouped Query Attention, multiple query heads share the same KV head. The FlashAttention kernel needs to route each query head to the correct KV head:

  GQA in FlashAttention
  ======================

  n_q_heads = 32, n_kv_heads = 8, group_size = 4

  Query heads  0, 1, 2, 3   --> KV head 0
  Query heads  4, 5, 6, 7   --> KV head 1
  Query heads  8, 9, 10, 11 --> KV head 2
  ...
  Query heads 28, 29, 30, 31 --> KV head 7

  Implementation options:

  Option A: One kernel launch per KV head, process 4 Q heads
  +-----------+
  | Q heads   |  Process heads 0,1,2,3 together
  | 0,1,2,3   |  They all read the SAME K,V tiles
  |           |  Load K,V once, apply to 4 query sets
  +-----------+

  Option B: Separate launches, but the KV tile loading is shared
  Each Q head kernel knows its kv_head_id = q_head / group_size
  and indexes into the correct KV cache slice.

The efficient approach processes all query heads in a group simultaneously, loading each KV tile once:

// GQA FlashAttention: process one group of query heads
kernel void flash_attn_gqa(
    device const half*  Q          [[buffer(0)]],    // [n_q_heads x seq x d]
    device const half*  K_cache    [[buffer(1)]],    // [n_kv_heads x seq x d]
    device const half*  V_cache    [[buffer(2)]],    // [n_kv_heads x seq x d]
    device half*        O          [[buffer(3)]],    // [n_q_heads x seq x d]
    constant uint&      kv_head    [[buffer(4)]],    // Which KV head
    constant uint&      group_size [[buffer(5)]],    // Typically 4
    // ... other params ...
    uint                sgid       [[simdgroup_index_in_threadgroup]],
    uint                tg_id      [[threadgroup_position_in_grid]])
{
    // Determine which query head this SIMD group handles
    // Within one threadgroup, different SIMD groups handle different Q heads
    uint q_head_in_group = sgid % group_size;
    uint q_head = kv_head * group_size + q_head_in_group;

    // K, V come from kv_head (shared by all Q heads in group)
    device const half* K = K_cache + kv_head * seq_len * d;
    device const half* V = V_cache + kv_head * seq_len * d;

    // Q comes from this specific query head
    device const half* my_Q = Q + q_head * seq_len * d;

    // Now run standard FlashAttention on (my_Q, K, V)
    // K and V tiles are loaded once and shared across SIMD groups
    // within the threadgroup (via threadgroup memory)
    // ...
}

Numerical Stability Deep Dive

The exp and division operations in softmax are numerically treacherous. Let us examine the potential pitfalls and how online softmax handles them:

  Numerical Issues in Softmax
  ============================

  Problem 1: exp overflow
    exp(100) = 2.69e43   (fine in FP32)
    exp(89)  = 4.49e38   (max FP32 is 3.40e38!)
    exp(90)  = +inf      (OVERFLOW!)

  Solution: subtract the max before exp
    If max(z) = 90 and z[i] = 85:
      exp(85) = overflow
      exp(85 - 90) = exp(-5) = 0.0067   (safe!)

  Problem 2: exp underflow (less critical)
    exp(-100) = 3.72e-44  (denormal in FP32)
    exp(-104) = 0.0       (underflow to zero)

  This is usually fine -- those entries get negligible weight anyway.

  Online softmax and stability:
    We always compute exp(z[i] - m) where m = running max.
    Since m >= z[i] for all previously seen values:
      z[i] - m <= 0    (always non-positive)
      exp(z[i] - m) <= 1.0    (never overflows!)

  The correction factor exp(m_old - m_new):
    m_old <= m_new   (max only increases)
    m_old - m_new <= 0
    exp(m_old - m_new) <= 1.0   (never overflows!)

  Summary: online softmax is ALWAYS numerically safe.
  Standard softmax without the max-subtraction trick can overflow.

Practical Performance Considerations

Tile Size Selection

The choice of Bq and Bk affects performance significantly:

  Tile Size Tradeoffs
  ====================

  Larger tiles (Bq=64, Bk=64):
    + More compute per memory load (better arithmetic intensity)
    + Better SIMD group MMA utilization
    - More threadgroup memory needed
    - Fewer threadgroups can run concurrently (occupancy)
    - More wasted compute at sequence boundaries

  Smaller tiles (Bq=16, Bk=32):
    + Less threadgroup memory
    + Better occupancy
    + Less waste at boundaries
    - Lower arithmetic intensity
    - More kernel overhead per tile

  Typical sweet spots on Apple Silicon:
    Prefill: Bq=32, Bk=32 or Bq=64, Bk=64
    Decode:  Bk=32 or Bk=64 (Bq=1)

Memory Bandwidth vs Compute

  FlashAttention Performance Model
  =================================

  Prefill (n=4096, d=128, Bq=32, Bk=32):

  Per Q-tile iteration (processing one KV tile):
    Load K_tile: 32 * 128 * 2 = 8 KB
    Load V_tile: 32 * 128 * 2 = 8 KB
    Compute S:   2 * 32 * 32 * 128 = 262K FLOPs
    Compute O:   2 * 32 * 32 * 128 = 262K FLOPs
    Total compute: ~524K FLOPs
    Total memory: ~16 KB

    Arithmetic intensity: 524K / 16K = 32.75 FLOPs/byte
    This is in the compute-bound regime! (Good!)

  Decode (n=4096, d=128, Bk=32):
    Per KV tile:
      Load K_tile: 32 * 128 * 2 = 8 KB
      Load V_tile: 32 * 128 * 2 = 8 KB
      Compute scores: 2 * 1 * 32 * 128 = 8K FLOPs
      Compute output: 2 * 1 * 32 * 128 = 8K FLOPs
      Total compute: ~16K FLOPs
      Total memory: ~16 KB

    Arithmetic intensity: 16K / 16K = 1 FLOP/byte
    Bandwidth-bound (expected for decode)

    Total KV cache read: 2 * 4096 * 128 * 2 = 2 MB per head
    At 200 GB/s: ~0.01 ms per head, ~0.32 ms for 32 heads

Summary

FlashAttention transforms attention from an O(n^2) memory algorithm to an O(n) memory algorithm without changing the mathematical result. The key ideas are:

1. Online softmax enables streaming computation by maintaining a running maximum and sum, with correction factors applied when the maximum changes. This is numerically stable by construction.

2. Tiling divides the KV sequence into blocks of size Bk. For each query tile, we iterate over KV tiles, computing a small score matrix in registers/threadgroup memory and updating the running output.

3. Memory savings are dramatic: from 4 GB per head (n=32768) to about 8 MB per head. This is what makes 128K+ context lengths feasible.

4. Decode variant (FlashDecoding) processes KV tiles in parallel since there is only one query, then merges results using the same online softmax correction. This maximizes GPU utilization for the bandwidth-bound decode phase.

5. Tree masking extends the algorithm for speculative decoding by replacing the simple causal mask comparison with a lookup into an explicit bitmask.

6. GQA integration shares KV tile loads across multiple query heads in the same group, reducing redundant memory traffic.

The combination of FlashAttention with quantized KV caches and GQA is what makes modern long-context LLM inference practical on consumer GPUs. Without these techniques, even a 4096-token context would strain memory; with them, 128K-token contexts are routine.

Quantization: Making Models Fit

Here is a problem you will hit immediately when trying to run LLMs on consumer hardware: a 7-billion parameter model in FP16 takes 14 GB of memory. An M2 MacBook Air has 8 GB of unified memory. A 70B model needs 140 GB in FP16 – more than any single GPU on the market. Even if you have enough memory, the model needs to stream through the memory bus during inference, and memory bandwidth is the bottleneck for token generation.

Quantization is the solution. By representing weights with fewer bits – 8, 4, 3, or even 2 bits instead of 16 – we shrink the model dramatically. A 7B model at 4-bit quantization fits in about 4 GB. The 70B model fits in 35-40 GB. And because there is less data to read from memory, inference gets faster too.

But quantization is not free. Fewer bits means less precision, which means some degradation in output quality. The art of quantization is finding the sweet spot: aggressive enough to fit in memory and hit target speeds, but gentle enough to preserve model quality.

In this chapter, we will cover the major quantization schemes you will encounter in practice, understand the math behind them, work through concrete examples of quantizing and dequantizing, and analyze the tradeoffs between size, speed, and quality.

Why Quantize? The Numbers

Let us start with the raw arithmetic that makes quantization essential:

  Model Sizes at Different Precisions
  =====================================

  Parameters    FP32     FP16     Q8_0     Q4_0     Q4_K_M    Q2_K
  ----------    ----     ----     ----     ----     ------    ----
  1.5B          6 GB     3 GB    1.6 GB   0.9 GB   1.0 GB   0.6 GB
  7B           28 GB    14 GB    7.5 GB   4.0 GB   4.4 GB   2.7 GB
  13B          52 GB    26 GB   13.8 GB   7.4 GB   8.0 GB   4.9 GB
  34B         136 GB    68 GB   36.1 GB  19.2 GB  21.1 GB  12.7 GB
  70B         280 GB   140 GB   74.4 GB  39.6 GB  43.4 GB  26.2 GB

  Apple Silicon Unified Memory:
  M1/M2/M3 (base):      8 GB   -> Q4_0 7B fits, FP16 does not
  M1/M2/M3 Pro:        18 GB   -> Q4_0 13B fits
  M1/M2/M3 Max:        64 GB   -> Q4_0 70B fits
  M2/M3 Ultra:        192 GB   -> FP16 70B fits

But memory capacity is only half the story. The other half is bandwidth:

  The Bandwidth Equation
  =======================

  During autoregressive decoding (generating one token at a time),
  EVERY weight in the model is read from memory exactly once per token.

  decode_speed (tokens/sec) = memory_bandwidth / model_size_in_bytes

  Apple M2 Pro: ~200 GB/s bandwidth

  FP16 7B:   200 / 14.0 = ~14 tokens/sec
  Q8_0 7B:   200 / 7.5  = ~27 tokens/sec
  Q4_0 7B:   200 / 4.0  = ~50 tokens/sec
  Q4_K_M 7B: 200 / 4.4  = ~45 tokens/sec
  Q2_K 7B:   200 / 2.7  = ~74 tokens/sec

  This is a theoretical upper bound -- real speeds are 60-80% of this
  due to attention computation, KV cache reads, and other overhead.
  But the proportionality holds: half the bytes = double the speed.

  Decode Speed vs Quantization
  =============================

  Speed                Model: 7B on M2 Pro (200 GB/s)
  (tok/s)
  80 |                                              *  Q2_K
     |
  70 |
     |
  60 |
     |                                 *  Q4_0
  50 |                              *  Q4_K_M
     |
  40 |
     |
  30 |                 *  Q8_0
     |
  20 |
     |     *  FP16
  10 |
     |
   0 +-----+------+------+------+------+------+----->
         16     8      6      4      3      2    bits/weight

The Fundamentals of Quantization

At its core, quantization maps a continuous range of floating-point values to a discrete set of integer values. The simplest form is uniform affine quantization:

  Uniform Quantization
  =====================

  Given: a set of FP16 weights in range [w_min, w_max]
  Goal:  represent each weight using b bits (0 to 2^b - 1)

  Quantize:
    scale = (w_max - w_min) / (2^b - 1)
    zero_point = round(-w_min / scale)
    q[i] = round(w[i] / scale) + zero_point

  Dequantize:
    w[i] = scale * (q[i] - zero_point)

  Example: b=4 bits (range 0..15)
    Weights: [-0.8, 0.3, -0.1, 0.5, -0.6, 0.7]
    w_min = -0.8, w_max = 0.7
    scale = (-0.8 - 0.7) / -15 = 0.1
    zero_point = round(0.8 / 0.1) = 8

    Quantize:
    q[-0.8] = round(-0.8/0.1) + 8 = -8 + 8 = 0
    q[0.3]  = round(0.3/0.1) + 8  =  3 + 8 = 11
    q[-0.1] = round(-0.1/0.1) + 8 = -1 + 8 = 7
    q[0.5]  = round(0.5/0.1) + 8  =  5 + 8 = 13
    q[-0.6] = round(-0.6/0.1) + 8 = -6 + 8 = 2
    q[0.7]  = round(0.7/0.1) + 8  =  7 + 8 = 15

    Stored: [0, 11, 7, 13, 2, 15]  (each is 4 bits)

    Dequantize:
    w[0]  = 0.1 * (0 - 8)  = -0.8   (exact)
    w[11] = 0.1 * (11 - 8) =  0.3   (exact)
    w[7]  = 0.1 * (7 - 8)  = -0.1   (exact)
    w[13] = 0.1 * (13 - 8) =  0.5   (exact)
    w[2]  = 0.1 * (2 - 8)  = -0.6   (exact)
    w[15] = 0.1 * (15 - 8) =  0.7   (exact)

  In this lucky example, all values hit exactly.
  In reality, most values get rounded, introducing error.

But this global quantization (one scale and zero-point for the entire weight matrix) is too crude for neural networks. The weight distribution varies significantly across different parts of the matrix. Modern quantization uses block quantization: the weights are divided into small blocks, each with its own scale (and possibly zero-point).

Block Quantization: GGUF Format (Q4_0)

The GGUF format (used by llama.cpp and many Metal inference engines) uses block quantization. The simplest variant, Q4_0, works as follows:

  Q4_0 Block Structure
  =====================

  Block size: 32 weights
  Storage per block:
    - 1 x FP16 scale factor (2 bytes)
    - 32 x 4-bit signed integers packed into 16 bytes
  Total: 18 bytes for 32 weights = 4.5 bits/weight

  Quantization (symmetric, no zero point):
    For each block of 32 weights:
      abs_max = max(|w[0]|, |w[1]|, ..., |w[31]|)
      scale = abs_max / 8     (maps to range [-8, 7])

      For each weight w[i]:
        q[i] = clamp(round(w[i] / scale), -8, 7)
        Store as (q[i] + 8), giving range [0, 15] (4 bits)

  Dequantization:
    w[i] = scale * (stored[i] - 8)

  Memory layout (18 bytes per block):
  +--------+--------+--------+---+--------+
  | scale  | byte 0 | byte 1 |...| byte15 |
  | (FP16) | w0|w1  | w2|w3  |   |w30|w31 |
  +--------+--------+--------+---+--------+
     2B        1B       1B          1B
             (2 nibbles packed per byte)

Worked Example: Quantizing a Block with Q4_0

  Example: Quantize 32 weights with Q4_0
  ========================================

  Weights (first 8 of 32 shown):
  [ 0.23, -0.41, 0.67, -0.12, 0.55, -0.89, 0.34, 0.08, ... ]

  Step 1: Find absolute maximum
    abs_max = max(|all 32 weights|) = 0.89  (from -0.89)

  Step 2: Compute scale
    scale = 0.89 / 8 = 0.11125

  Step 3: Quantize each weight
    q[0] = round(0.23 / 0.11125)  = round(2.067)  =  2  -> stored: 2+8 = 10
    q[1] = round(-0.41 / 0.11125) = round(-3.685)  = -4  -> stored: -4+8 = 4
    q[2] = round(0.67 / 0.11125)  = round(6.022)  =  6  -> stored: 6+8 = 14
    q[3] = round(-0.12 / 0.11125) = round(-1.079)  = -1  -> stored: -1+8 = 7
    q[4] = round(0.55 / 0.11125)  = round(4.944)  =  5  -> stored: 5+8 = 13
    q[5] = round(-0.89 / 0.11125) = round(-8.000)  = -8  -> stored: -8+8 = 0
    q[6] = round(0.34 / 0.11125)  = round(3.056)  =  3  -> stored: 3+8 = 11
    q[7] = round(0.08 / 0.11125)  = round(0.719)  =  1  -> stored: 1+8 = 9

  Step 4: Pack into bytes (two 4-bit values per byte)
    byte[0] = (stored[1] << 4) | stored[0] = (4 << 4) | 10 = 0x4A
    byte[1] = (stored[3] << 4) | stored[2] = (7 << 4) | 14 = 0x7E
    byte[2] = (stored[5] << 4) | stored[4] = (0 << 4) | 13 = 0x0D
    byte[3] = (stored[7] << 4) | stored[6] = (9 << 4) | 11 = 0x9B

  Step 5: Dequantize (to verify)
    dq[0] = 0.11125 * (10 - 8) = 0.11125 * 2  =  0.2225  (was 0.23, error: 0.008)
    dq[1] = 0.11125 * (4 - 8)  = 0.11125 * -4 = -0.4450  (was -0.41, error: 0.035)
    dq[2] = 0.11125 * (14 - 8) = 0.11125 * 6  =  0.6675  (was 0.67, error: 0.003)
    dq[3] = 0.11125 * (7 - 8)  = 0.11125 * -1 = -0.1113  (was -0.12, error: 0.009)
    dq[4] = 0.11125 * (13 - 8) = 0.11125 * 5  =  0.5563  (was 0.55, error: 0.006)
    dq[5] = 0.11125 * (0 - 8)  = 0.11125 * -8 = -0.8900  (was -0.89, error: 0.000)
    dq[6] = 0.11125 * (11 - 8) = 0.11125 * 3  =  0.3338  (was 0.34, error: 0.006)
    dq[7] = 0.11125 * (9 - 8)  = 0.11125 * 1  =  0.1113  (was 0.08, error: 0.031)

  Average absolute error: ~0.012
  Relative error: ~3-4% on average

Bit Packing Details

Two 4-bit values are packed into each byte. This is a critical operation for both quantization (packing) and dequantization (unpacking):

  Bit Packing: Two Nibbles Per Byte
  ===================================

  Stored values: a=10 (0b1010), b=4 (0b0100)

  Pack: byte = (b << 4) | a
        byte = (0b0100 << 4) | 0b1010
        byte = 0b01001010
        byte = 0x4A

  Unpack:
    a = byte & 0x0F       = 0x4A & 0x0F = 0x0A = 10
    b = (byte >> 4) & 0x0F = (0x4A >> 4) & 0x0F = 0x04 = 4

  Visual:
  +---+---+---+---+---+---+---+---+
  | b3| b2| b1| b0| a3| a2| a1| a0|  <-- one byte
  +---+---+---+---+---+---+---+---+
  |    high nibble |   low nibble  |
  |    value b     |   value a     |

In Metal shader code:

// Unpacking 4-bit values from bytes
uchar packed_byte = block.packed[j];

// Extract two 4-bit values
int8_t val_low  = (packed_byte & 0x0F) - 8;    // Low nibble, subtract offset
int8_t val_high = (packed_byte >> 4) - 8;       // High nibble, subtract offset

// Dequantize
float w0 = float(block.scale) * float(val_low);
float w1 = float(block.scale) * float(val_high);

K-Quant Family: Super-Blocks of 256

The basic Q4_0 format uses a single scale per 32 weights. The K-quant family (Q2_K, Q3_K, Q4_K, Q5_K, Q6_K) introduced by llama.cpp uses a two-level hierarchy: super-blocks of 256 weights, each containing 8 sub-blocks of 32 weights.

  K-Quant Super-Block Structure (Q4_K)
  ======================================

  Super-block: 256 weights

  +---------------------------------------------------+
  | FP16 scale_of_scales  |  FP16 scale_of_mins       |   4 bytes
  +---------------------------------------------------+
  | 8 x 6-bit sub-block scales (packed)               |  6 bytes
  +---------------------------------------------------+
  | 8 x 6-bit sub-block mins (packed)                 |  6 bytes
  +---------------------------------------------------+
  | 256 x 4-bit quantized values (packed)             | 128 bytes
  +---------------------------------------------------+
  Total: 144 bytes for 256 weights = 4.5 bits/weight

  Each sub-block of 32 weights has its own:
    - 6-bit scale (quantized, relative to super-block scale_of_scales)
    - 6-bit minimum (quantized, relative to super-block scale_of_mins)

  Dequantization for weight i in sub-block b:
    sub_scale = scale_of_scales * sub_scales_q6[b]
    sub_min   = scale_of_mins * sub_mins_q6[b]
    w[i] = sub_scale * q4[i] - sub_min

Why the two-level structure? It improves quantization accuracy by allowing each sub-block of 32 weights to have a different range, while keeping the overhead (scale and min metadata) small by quantizing the per-sub-block parameters themselves.

  Q4_0 vs Q4_K: Quantization Granularity
  ========================================

  Q4_0: One FP16 scale per 32 weights
  +---------+---------+---------+---------+
  | scale_0 | scale_1 | scale_2 | scale_3 |  ... (256 weights = 8 blocks)
  |  32 wts |  32 wts |  32 wts |  32 wts |
  +---------+---------+---------+---------+
  Each scale is independent FP16 (full precision).
  Overhead: 2 bytes / 32 weights = 0.5 bits/weight
  Total: 4.0 + 0.5 = 4.5 bits/weight

  Q4_K: Two-level hierarchy for 256 weights
  +-------------------------------------------+
  | super-block: scale_of_scales, scale_of_mins|  4 bytes for 256 weights
  |   sub-block 0: 6-bit scale, 6-bit min     |
  |   sub-block 1: 6-bit scale, 6-bit min     |
  |   ...                                      |
  |   sub-block 7: 6-bit scale, 6-bit min     |
  +-------------------------------------------+
  Sub-block parameters: 12 bytes / 256 = 0.375 bits/weight
  Super-block parameters: 4 bytes / 256 = 0.125 bits/weight
  Total: 4.0 + 0.375 + 0.125 = 4.5 bits/weight

  Same bits/weight, but Q4_K has finer-grained adaptation
  and asymmetric ranges (scale + min instead of just scale).
  Result: measurably lower perplexity.

The K-Quant Zoo

Here is the full family of K-quant formats:

  K-Quant Formats Overview
  =========================

  Format   Bits/wt  Quant    Sub-block params     Quality
  ------   -------  -----    ----------------     -------
  Q2_K     2.56     2-bit    4-bit scale+min      Poor (emergency use)
  Q3_K_S   3.44     3-bit    6-bit scale          Acceptable
  Q3_K_M   3.91     3-bit    6-bit scale+min      Fair
  Q3_K_L   4.28     3-bit    6-bit scale+min      Fair+
  Q4_K_S   4.50     4-bit    6-bit scale+min      Good
  Q4_K_M   4.84     4-bit    6-bit scale+min      Good+
  Q5_K_S   5.50     5-bit    6-bit scale+min      Very Good
  Q5_K_M   5.69     5-bit    6-bit scale+min      Very Good+
  Q6_K     6.56     6-bit    8-bit scale          Excellent

  The S/M/L suffixes mean Small/Medium/Large and refer to which
  layers get higher precision. "M" uses higher precision for the
  attention layers and output layers (most sensitive to quantization).

The _M variants are particularly clever: they use mixed precision, with more bits for the layers that matter most:

  Q4_K_M: Mixed Precision by Layer
  ==================================

  Layer Type              Quantization    Why
  ----------              ------------    ---
  Attention Q,K,V,O       Q4_K           Important for quality
  FFN gate, up            Q4_K           Standard
  FFN down                Q6_K           Output projection, sensitive
  Output (vocab) layer    Q6_K           Final prediction, very sensitive
  Embedding               Q4_K           Large but less sensitive

  The "M" variant is ~8% larger than "S" but measurably better.

Per-Channel and Per-Group Quantization (MLX Format)

The MLX framework (Apple’s machine learning framework) uses a different quantization approach: per-group affine quantization with configurable group sizes.

  MLX Quantization
  =================

  Group size: typically 64 or 128
  Per group: one FP16 scale and one FP16 bias (zero-point)

  Quantize (per group of G weights):
    w_min = min(weights in group)
    w_max = max(weights in group)
    scale = (w_max - w_min) / (2^bits - 1)
    bias = w_min

    q[i] = round((w[i] - bias) / scale)
    q[i] = clamp(q[i], 0, 2^bits - 1)

  Dequantize:
    w[i] = scale * q[i] + bias

  Example with group_size=64, bits=4:
  +------------------------------------+
  | scale (FP16) | bias (FP16) | 64x4b |
  |    2 bytes   |   2 bytes   | 32 B  |
  +------------------------------------+
  Total: 36 bytes for 64 weights = 4.5 bits/weight

  With group_size=128:
  +-------------------------------------+
  | scale (FP16) | bias (FP16) | 128x4b |
  |    2 bytes   |   2 bytes   |  64 B  |
  +-------------------------------------+
  Total: 68 bytes for 128 weights = 4.25 bits/weight

The key difference from GGUF quantization: MLX uses affine quantization (scale + bias) rather than symmetric quantization (scale only). This better handles weight distributions that are not centered around zero:

  Symmetric vs Affine Quantization
  ==================================

  Symmetric (Q4_0):
    Maps [-abs_max, +abs_max] to [-8, +7]
    scale = abs_max / 8
    w = scale * (q - 8)

    Problem: if weights are [0.1, 0.3, 0.5, 0.7, 0.9]
    abs_max = 0.9, scale = 0.1125
    Half the quantization levels (-8 to -1) are wasted!

                            Wasted range
                     <---------->
    -0.9  -0.7  -0.5  -0.3  -0.1  0.1  0.3  0.5  0.7  0.9
     |     |     |     |     |     |    |    |    |    |
     q=0   q=2   q=4   q=6   q=8  q=9  q=11 q=13 q=14 q=16
                 NO WEIGHTS HERE           ALL WEIGHTS HERE


  Affine (MLX):
    Maps [w_min, w_max] to [0, 15]
    scale = (0.9 - 0.1) / 15 = 0.0533
    bias = 0.1
    w = scale * q + bias

    0.1    0.15   0.21   0.26   ...   0.84   0.9
     |      |      |      |           |      |
     q=0    q=1    q=2    q=3  ...   q=14   q=15
              FULL RANGE UTILIZED!

  Affine quantization uses all levels for the actual data range.
  Better for asymmetric distributions (common in practice).

MLX Quantization Code Example

// MLX-style dequantization in Metal
kernel void dequantize_mlx_q4(
    device const uint32_t* packed_weights [[buffer(0)]],  // Packed 4-bit
    device const half*     scales         [[buffer(1)]],  // Per-group scales
    device const half*     biases         [[buffer(2)]],  // Per-group biases
    device half*           output         [[buffer(3)]],  // Dequantized
    constant uint&         group_size     [[buffer(4)]],
    uint                   tid            [[thread_position_in_grid]])
{
    uint group_id = tid / group_size;
    uint in_group  = tid % group_size;

    half scale = scales[group_id];
    half bias  = biases[group_id];

    // Each uint32 holds 8 x 4-bit values
    uint word_idx = tid / 8;
    uint bit_offset = (tid % 8) * 4;
    uint32_t word = packed_weights[word_idx];
    uint8_t q = (word >> bit_offset) & 0xF;

    output[tid] = scale * half(q) + bias;
}

Impact on Quality: Perplexity Analysis

Perplexity is the standard metric for measuring quantization quality. Lower is better – it measures how “surprised” the model is by test data. A perplexity increase of more than 1-2% is generally noticeable in output quality.

  Perplexity Comparison (Llama 2 7B, WikiText-2)
  ================================================

  Format      Bits/wt    Perplexity    Delta     Quality Assessment
  ------      -------    ----------    -----     ------------------
  FP16        16.00      5.796         ---       Reference (baseline)
  Q8_0         8.50      5.799        +0.003     Indistinguishable
  Q6_K         6.56      5.804        +0.008     Excellent
  Q5_K_M       5.69      5.812        +0.016     Very good
  Q5_K_S       5.50      5.819        +0.023     Very good
  Q4_K_M       4.84      5.882        +0.086     Good
  Q4_K_S       4.50      5.912        +0.116     Good
  Q4_0         4.50      5.946        +0.150     Acceptable
  Q3_K_M       3.91      6.145        +0.349     Fair (noticeable)
  Q3_K_S       3.44      6.351        +0.555     Degraded
  Q2_K         2.56      6.981        +1.185     Poor

  Rule of thumb:
  - Q4_K_M and above: quality loss rarely noticeable in practice
  - Q3_K: quality loss sometimes noticeable, especially on hard tasks
  - Q2_K: clear quality degradation, only for extreme memory constraints

  Perplexity vs Bits per Weight
  ==============================

  Perplexity
      |
  7.0 |*                                              Q2_K
      |
  6.5 |    *                                          Q3_K_S
      |       *                                       Q3_K_M
  6.0 |            *  *                               Q4_0, Q4_K_S
      |               *  *                            Q4_K_M, Q5_K_S
  5.8 |                     * *                       Q5_K_M, Q6_K
      |                          *                    Q8_0
  5.6 |                                         *    FP16
      |
  5.4 +---+---+---+---+---+---+---+---+---+---+--->
         2   3   4   5   6   7   8  ...  16     bits/weight

  The curve has a "knee" around 4-5 bits/weight.
  Below 4 bits, quality drops rapidly.
  Above 5 bits, returns diminish.
  4-5 bits is the sweet spot for most applications.

The Bandwidth Equation: Predicting Decode Speed

We mentioned this earlier, but let us formalize it. During autoregressive decoding:

  Decode Speed Model
  ====================

  Given:
    P = number of parameters
    b = bits per weight (effective, including quantization overhead)
    BW = memory bandwidth (bytes/sec)
    overhead = non-matmul time (attention, normalization, etc.)

  model_bytes = P * b / 8

  time_per_token = model_bytes / BW + overhead

  tokens_per_second = 1 / time_per_token

  Example: Llama 2 7B on M2 Pro
  P = 6.74 billion (actual parameter count)
  BW = 200 GB/s

  Q4_K_M (b = 4.84):
    model_bytes = 6.74e9 * 4.84 / 8 = 4.077 GB
    matmul_time = 4.077 / 200 = 0.0204 sec = 20.4 ms
    overhead ≈ 5 ms (attention + KV cache + norms)
    time_per_token ≈ 25.4 ms
    speed ≈ 39 tokens/sec

  Observed: ~35-40 tokens/sec (matches!)

This equation is remarkably accurate because GEMV (the matmul during decode) is almost perfectly bandwidth-bound. The only thing that breaks the model is when the KV cache becomes very large (long sequences) and the attention computation starts contributing significantly.

  When Does Attention Start to Matter?
  ======================================

  KV cache read per step (GQA, 8 KV heads, d=128):
    kv_bytes = 2 * 8 * seq_len * 128 * 2 = 4096 * seq_len bytes

  Total per token = model_bytes + kv_bytes

  seq_len      kv_bytes    model(Q4_K_M)    kv/model    Impact
  --------     --------    -------------    --------    ------
  512          2 MB        4.08 GB          0.05%       Negligible
  2048         8 MB        4.08 GB          0.2%        Negligible
  8192         32 MB       4.08 GB          0.8%        Minimal
  32768        128 MB      4.08 GB          3.1%        Small
  131072       512 MB      4.08 GB          12.5%       Noticeable (~11% slower)

  KV cache impact is negligible for typical conversations (< 8K tokens)
  but matters for very long contexts (> 32K tokens).

Comparison Table

Here is a comprehensive comparison of all the major quantization formats:

  +----------+--------+--------+----------+---------+--------+-------+
  | Format   | Bits/  | Block  | Params   | Asym-   | Mixed  | Used  |
  |          | weight | size   | per block| metric? | prec?  | by    |
  +----------+--------+--------+----------+---------+--------+-------+
  | FP16     | 16.00  | N/A    | N/A      | N/A     | No     | Base  |
  | Q8_0     |  8.50  |   32   | 1 scale  | No      | No     | GGUF  |
  | Q6_K     |  6.56  |  256   | 8-bit sc | No      | No     | GGUF  |
  | Q5_K_M   |  5.69  |  256   | 6-bit s+m| Yes     | Yes    | GGUF  |
  | Q5_K_S   |  5.50  |  256   | 6-bit s+m| Yes     | No     | GGUF  |
  | Q4_K_M   |  4.84  |  256   | 6-bit s+m| Yes     | Yes    | GGUF  |
  | Q4_K_S   |  4.50  |  256   | 6-bit s+m| Yes     | No     | GGUF  |
  | Q4_0     |  4.50  |   32   | 1 scale  | No      | No     | GGUF  |
  | Q3_K_M   |  3.91  |  256   | 6-bit s+m| Yes     | Yes    | GGUF  |
  | Q3_K_S   |  3.44  |  256   | 6-bit sc | No      | No     | GGUF  |
  | Q2_K     |  2.56  |  256   | 4-bit s+m| Yes     | No     | GGUF  |
  | MLX-4    |  4.25   |  64-128| scale+bias| Yes   | No     | MLX   |
  | MLX-8    |  8.25   |  64-128| scale+bias| Yes   | No     | MLX   |
  | MLX-2    |  2.25   |  64-128| scale+bias| Yes   | No     | MLX   |
  +----------+--------+--------+----------+---------+--------+-------+

  Legend: sc = scale only, s+m = scale + minimum, Asym = asymmetric range
  Mixed prec = different quantization levels for different layers

Worked Example: Full Dequantization Pipeline

Let us walk through dequantizing a complete Q4_K block on the GPU, showing exactly what happens in a Metal shader:

  Q4_K Dequantization Pipeline
  ==============================

  Input: One super-block of 256 quantized weights

  Step 1: Read super-block header
    d = FP16 scale_of_scales       (e.g., 0.0156)
    dmin = FP16 scale_of_mins      (e.g., 0.0078)

  Step 2: Unpack sub-block parameters (6-bit values, packed)
    The 8 sub-block scales are packed as 6-bit values in 6 bytes:
    +--------+--------+--------+--------+--------+--------+
    | byte 0 | byte 1 | byte 2 | byte 3 | byte 4 | byte 5 |
    +--------+--------+--------+--------+--------+--------+

    sub_scale[0] = byte[0] & 0x3F                    = 23
    sub_scale[1] = ((byte[0]>>6) | (byte[1]<<2)) & 0x3F = 18
    sub_scale[2] = ((byte[1]>>4) | (byte[2]<<4)) & 0x3F = 31
    ...

    Similarly for sub_mins.

  Step 3: Compute actual sub-block scale and min
    For sub-block b:
      actual_scale[b] = d * sub_scale[b]
      actual_min[b]   = dmin * sub_min[b]

    actual_scale[0] = 0.0156 * 23 = 0.3588
    actual_min[0]   = 0.0078 * 15 = 0.1170  (example)

  Step 4: Dequantize weights in sub-block
    For weight i in sub-block b:
      q = unpack_4bit(packed_data, b*32 + i)   (0..15)
      w = actual_scale[b] * q - actual_min[b]

    q[0] = 7:   w = 0.3588 * 7 - 0.1170 = 2.3946
    q[1] = 12:  w = 0.3588 * 12 - 0.1170 = 4.1886
    q[2] = 3:   w = 0.3588 * 3 - 0.1170 = 0.9594
    q[3] = 0:   w = 0.3588 * 0 - 0.1170 = -0.1170  (min value)
    q[4] = 15:  w = 0.3588 * 15 - 0.1170 = 5.2650  (max value)
    ...

And here is how this looks in Metal shader code:

// Q4_K super-block structure
struct block_q4_K {
    half    d;              // scale of scales
    half    dmin;           // scale of mins
    uchar   scales[12];    // 8 x 6-bit scales + 8 x 6-bit mins, packed
    uchar   qs[128];       // 256 x 4-bit quantized values
};

// Dequantize one Q4_K block
inline void dequantize_q4_K(
    device const block_q4_K& block,
    uint sub_block,     // 0..7
    uint index_in_sub,  // 0..31
    thread float& result)
{
    // Step 1: Read super-block scales
    float d = float(block.d);
    float dmin = float(block.dmin);

    // Step 2: Unpack 6-bit sub-block scale and min
    // (Simplified -- actual packing is more complex)
    uint8_t raw_scale, raw_min;
    if (sub_block < 4) {
        raw_scale = block.scales[sub_block] & 0x3F;
        raw_min   = block.scales[sub_block + 4] & 0x3F;
    } else {
        // Higher sub-blocks use bits from multiple bytes
        raw_scale = ((block.scales[sub_block + 4] & 0xF) |
                    ((block.scales[sub_block - 4] >> 6) << 4));
        raw_min   = ((block.scales[sub_block + 4] >> 4) |
                    ((block.scales[sub_block]     >> 6) << 4));
    }

    float scale = d * float(raw_scale);
    float min   = dmin * float(raw_min);

    // Step 3: Unpack 4-bit weight value
    uint byte_idx = (sub_block * 32 + index_in_sub) / 2;
    uchar packed = block.qs[byte_idx];
    uint8_t q;
    if (index_in_sub % 2 == 0) {
        q = packed & 0x0F;        // Low nibble
    } else {
        q = (packed >> 4) & 0x0F; // High nibble
    }

    // Step 4: Dequantize
    result = scale * float(q) - min;
}

Quantizing the KV Cache

So far we have discussed quantizing the model weights. But there is another large memory consumer during inference: the KV cache. For long context lengths, the KV cache can rival or exceed the model size.

  KV Cache Quantization
  ======================

  Standard KV cache (FP16):
    Per layer: 2 * n_kv_heads * seq_len * d_k * 2 bytes
    32 layers, 8 KV heads, d_k=128, seq_len=32768:
    = 32 * 2 * 8 * 32768 * 128 * 2 = 4 GB

  Q8 KV cache:
    Same but 1 byte per element + small scale overhead
    ≈ 2 GB  (50% reduction)

  Q4 KV cache:
    ≈ 1 GB  (75% reduction)

  KV cache quantization is trickier than weight quantization because:
  1. Values are computed dynamically (not known ahead of time)
  2. Must quantize on-the-fly as new K/V vectors are computed
  3. Range can change as new tokens arrive
  4. Quality impact is harder to predict

  Approach: per-head, per-position quantization
    For each new (key, value) vector being added to the cache:
    1. Compute the vector in FP16
    2. Quantize to Q8 or Q4 with per-vector scale
    3. Store quantized vector in cache
    4. Dequantize on-the-fly during attention computation

Practical Considerations

Choosing the Right Quantization

Here is a practical decision flowchart:

  Quantization Selection Flowchart
  ==================================

  How much memory do you have?
  |
  +-- Very limited (8 GB): Use Q4_K_M for 7B, Q2_K for 13B
  |
  +-- Moderate (16-18 GB): Use Q4_K_M for 13B, Q6_K for 7B
  |
  +-- Generous (32-64 GB): Use Q6_K or Q8_0 for 13B-34B
  |
  +-- Abundant (96+ GB): Consider FP16 for best quality

  What is your speed target?
  |
  +-- Maximum speed: Use lowest quantization that maintains acceptable quality
  |                  (usually Q4_K_M or Q4_K_S)
  |
  +-- Quality priority: Use highest quantization that meets speed requirements
  |                     (Q6_K or Q8_0)
  |
  +-- Balanced: Q4_K_M is almost always the right answer

  Task sensitivity?
  |
  +-- Creative writing, coding: Q4_K_M usually fine
  |
  +-- Math, reasoning: Consider Q5_K_M or Q6_K
  |
  +-- Simple Q&A, summarization: Q3_K_M can work

Quantization and GEMV Performance

Remember from Chapter 14: GEMV is bandwidth-bound. Quantization directly reduces the bytes read, so the speedup is nearly linear with compression ratio:

  GEMV Performance with Quantization
  ====================================

  Single matmul: [1, 4096] * [4096, 11008]
  M2 Pro, 200 GB/s

  Format    Weight Size    Read Time    Speedup
  ------    -----------    ---------    -------
  FP16      90.1 MB        0.451 ms    1.0x
  Q8_0      47.8 MB        0.239 ms    1.9x
  Q6_K      37.0 MB        0.185 ms    2.4x
  Q4_K_M    24.5 MB        0.123 ms    3.7x
  Q4_0      23.6 MB        0.118 ms    3.8x
  Q2_K      14.4 MB        0.072 ms    6.3x

  Note: Dequantization adds ~5-10% compute overhead.
  Net speedup is slightly less than the raw bandwidth reduction,
  but still very close.

GEMM (Prefill) with Quantization

For prefill (GEMM), quantization does not directly speed things up because GEMM is compute-bound. However, it does reduce memory footprint, which matters for:

1. Fitting larger models in memory
2. Keeping more of the model in cache
3. Reducing memory pressure from concurrent operations

  Prefill Performance: Less Clear-Cut
  =====================================

  Matmul: [512, 4096] * [4096, 11008]
  M2 Pro

  Format    Approach                     Time
  ------    --------                     ----
  FP16      Native FP16 MMA              ~2.0 ms
  Q4_K_M    Dequant to FP16 + MMA        ~2.5 ms  (slower!)
  Q4_K_M    Optimized mixed-precision    ~2.2 ms  (close to FP16)

  During prefill, quantization can actually be SLOWER because:
  1. Extra dequantization compute
  2. MMA hardware is optimized for FP16/BF16, not mixed precision
  3. GEMM is compute-bound, not memory-bound

  But for most LLM use cases, decode time dominates total latency,
  so optimizing decode (where quantization helps enormously) is
  more important than optimizing prefill.

Summary

Quantization is what makes LLM inference practical on consumer hardware. The key points:

1. Why quantize: A 7B FP16 model needs 14 GB; Q4 needs ~4 GB. Decode speed is proportional to 1/model_size because GEMV is bandwidth-bound.

2. Block quantization (Q4_0): Groups of 32 weights share one FP16 scale. 4 bits per weight + scale overhead = 4.5 bits/weight effective. Simple and fast.

3. K-quant family: Super-blocks of 256 with 8 sub-blocks of 32. Two-level scale hierarchy. Mixed precision variants (_M) use more bits for sensitive layers. Q4_K_M is the most popular choice.

4. Per-group affine (MLX): Groups of 64-128 with per-group scale and bias. Asymmetric quantization handles non-centered distributions better.

5. Bit packing: Two 4-bit values per byte. Unpacking is a shift and mask operation, done on-the-fly during dequantization in the GPU shader.

6. Quality impact: The perplexity curve has a knee at ~4-5 bits/weight. Q4_K_M is the sweet spot for most applications. Below 3 bits, quality degrades noticeably.

7. The bandwidth equation: decode_speed ≈ bandwidth / model_bytes. This simple formula predicts real-world performance with surprising accuracy.

8. Dequantize on the fly: During GEMV, weights are dequantized in registers as they are loaded. The compute cost of dequantization is negligible compared to the bandwidth savings.

With quantization, a $1200 MacBook Air with 16 GB of RAM can run a 7B parameter model at 40+ tokens per second – fast enough for interactive use. That is the power of trading a tiny bit of precision for an enormous reduction in memory and bandwidth requirements.

Akunu Overview and Design Philosophy

Welcome to the deep-dive section of this book. Up until now, we have been building intuition for how LLM inference works on Apple Silicon: the Metal compute pipeline, the memory hierarchy, quantized matrix math, and the attention mechanism. Now it is time to see how all of those pieces come together in a real, production-quality inference engine.

Akunu is a high-performance LLM inference engine written specifically for Apple Silicon. The name comes from the Sinhala word meaning “embers” – a fitting metaphor for a project that tries to extract every last bit of heat from the GPU silicon.

In this chapter we will survey the project at a high level: what it does, how it is organized, what design principles drive it, and what the end-to-end inference flow looks like. Subsequent chapters will zoom in on each subsystem.

What Akunu Is (and What It Is Not)

Akunu is a local inference engine. You give it a model file (GGUF or MLX SafeTensors), it loads the weights onto the Apple GPU, and it runs the full transformer forward pass – prefill and decode – entirely on-device. There is no cloud, no server round-trip, no Python runtime.

Here is what it supports today:

What it is not: it is not a training framework. It is not a general-purpose tensor library. It does not try to be cross-platform (though the architecture makes a future CUDA backend straightforward, as we will see). Every design decision optimizes for one thing: token throughput on Apple Silicon.

Performance: The Numbers

Let us start with the punchline, because performance is the reason this engine exists. All benchmarks were run on an Apple M4 Pro (16 GPU cores, 273 GB/s memory bandwidth):

Decode throughput (tg128, tokens/sec):

  vs llama.cpp:
    Average speedup:  1.83x
    Best case:        3.66x  (Qwen3-0.6B-Q3_K_S: 448 vs 123 t/s)
    Wins:             20/21 configurations

  vs MLX:
    Average speedup:  1.17x
    Best case:        1.25x  (Qwen3-0.6B-bf16: 207 vs 165 t/s)
    Wins:             11/11 configurations

These are not cherry-picked numbers. Across 19 GGUF model configurations and 11 MLX configurations, akunu wins decode throughput in 31 out of 32 tests. The speedup is most dramatic on small models (0.6B-1B parameters) with aggressive quantization (Q2_K through Q5_K), where akunu achieves 2-3.5x the throughput of llama.cpp.

Why? Because small quantized models are compute-bound during decode – the matrix multiplications finish so fast that overhead dominates. Akunu’s precompiled dispatch table and zero-allocation hot path eliminate that overhead. We will see exactly how in the sections below.

The Five Design Principles

Every non-trivial design decision in akunu traces back to one of five principles. Understanding these up front will make the rest of the codebase click.

Principle 1: Data-Driven Design (ArchDescriptor)

The naive way to support multiple architectures looks like this:

// DON'T DO THIS
if (arch == "llama") {
    activation = silu_gate;
    rope = rope_interleaved;
} else if (arch == "qwen3") {
    activation = silu_gate;
    rope = rope_neox;
    has_qk_norm = true;
} else if (arch == "gemma") {
    activation = gelu_gate;
    rope = rope_neox;
    has_qk_norm = true;
    embedding_scale = sqrt(dim);
    // ... 20 more fields
}

This approach does not scale. Every new architecture touches dozens of files. Every if/else branch is a potential bug.

Akunu takes a different approach: it captures all architecture-specific differences in a single POD struct called ArchDescriptor. The struct has about 20 fields covering activation kernels, RoPE style, embedding scaling, normalization, encoder config, and more. The entire table builder and prefill engine read from this struct and never branch on the architecture name.

Adding a new architecture means writing one factory function that fills in the struct. That is it. No code changes in the dispatch table builder, the prefill engine, or the decode loop.

+------------------+     +------------------+     +------------------+
| arch_llama()     |     | arch_qwen3()     |     | arch_gemma(dim)  |
| activation:      |     | activation:      |     | activation:      |
|   silu_gate_f16  |     |   silu_gate_f16  |     |   gelu_gate_f16  |
| rope:            |     | rope:            |     | rope:            |
|   interleaved    |     |   neox           |     |   neox           |
| qk_norm: false   |     | qk_norm: true    |     | qk_norm: true    |
| embed_scale: 0   |     | tie_embed: true  |     | embed_scale:     |
+--------+---------+     +--------+---------+     |   sqrt(dim)      |
         |                        |                +--------+---------+
         |                        |                         |
         +------------------------+-------------------------+
                                  |
                                  v
                    +----------------------------+
                    | build_dispatch_table(...)   |
                    | (reads ArchDescriptor,      |
                    |  never branches on arch)    |
                    +----------------------------+

We will cover ArchDescriptor in depth in Chapter 22.

Principle 2: Precompiled Dispatch (DispatchTable)

This is the big one. In most inference engines, every forward pass involves:

1. Looking up which kernel to run
2. Resolving buffer pointers
3. Computing dispatch geometry (grid size, threadgroup size)
4. Encoding the compute command

Akunu does all of this once, at model load time, and stores the result in a flat array of DispatchCmd structs. The decode hot path simply iterates this array, patches a couple of per-token fields (position, token offset), and submits the whole thing to the GPU.

Model load time (once):                    Decode time (every token):

  Parse weights                              for each token:
  |                                            for each cmd in dispatch_table:
  v                                              patch position
  Resolve kernel names                           patch token offset
  |                                              submit to encoder
  v
  Look up Pipeline State Objects
  |
  v
  Compute grid dimensions
  |
  v
  Bind buffers + params
  |
  v
  Store in DispatchCmd[]

The DispatchCmd struct itself is a fixed-size POD type with no heap allocations:

DispatchCmd (fixed size, no heap):
+-----------------------------------+
| Pipeline pso                      |  8 bytes
| Buffer buffers[8]                 |  8 x 24 bytes
| uint32_t offsets[8]               |  32 bytes
| int buffer_count                  |  4 bytes
| uint8_t param_bytes[64]           |  64 bytes  (inline kernel params)
| int param_size, param_index       |  8 bytes
| Buffer param_buf                  |  24 bytes
| uint8_t param2_bytes[16]          |  16 bytes  (secondary params)
| Dim3 grid, threadgroup            |  24 bytes
| bool use_dispatch_threads         |  1 byte
| PatchType patch_type              |  1 byte
| int patch_offset_1, patch_offset_2|  8 bytes
+-----------------------------------+

The entire forward pass for a single token is typically 50-100 commands (embedding + N layers * ~5 commands each + output norm + logit projection + argmax). These are stored contiguously in memory, which is great for the CPU cache.

This is why akunu’s decode is fast: the CPU-side work per token is essentially a memcpy of a few patched bytes plus a loop of setBuffer/setBytes/dispatch calls – all inlined, no virtual dispatch, no hash lookups, no string comparisons.

Principle 3: Zero-Allocation Hot Path

Once the model is loaded, the decode loop allocates zero bytes of memory. All buffers are pre-allocated at model init time in a ScratchBuffers struct:

ScratchBuffers (all pre-allocated at model load):

  Decode (single token):
    h0         [dim]          -- embedding output / residual ping
    h1         [dim]          -- residual pong
    residual   [dim]          -- norm output
    qkv        [q_dim+2*kv]  -- contiguous Q|K|V
    attn_out   [max(q_dim,dim)]
    ffn_gate   [ffn_dim]
    ffn_up     [ffn_dim]
    ffn_act    [ffn_dim]
    logits     [vocab_size]
    token_ids  [max_chain]

  Prefill (batch):
    batch_h0   [chunk * dim]
    batch_q    [chunk * q_dim]
    batch_k    [chunk * kv_dim]
    ... (same pattern)

The KV cache is also pre-allocated to the maximum context length. The decode loop never calls malloc, never calls device.allocate, never resizes a vector. This matters more than you might think – on Apple Silicon, malloc can take microseconds, and when you are generating 400+ tokens per second, every microsecond counts.

Even the thread-local error buffer is a static char[512]:

thread_local char error_buf[512] = {};

No std::string, no exceptions, no heap in the hot path.

Principle 4: Virtual Device Interface

The Device base class provides a pure-virtual interface with about 30 methods covering buffer allocation, kernel loading, command encoding, and synchronization. Today there is exactly one implementation: MetalDevice, which wraps the Metal API. But the abstraction exists for a reason.

         +------------------+
         |   Device (base)  |
         |   pure virtual   |
         +--------+---------+
                  |
     +------------+------------+
     |                         |
+----+-------+          +-----+------+
| MetalDevice|          | CudaDevice |
| (ObjC++)   |          | (future)   |
+------------+          +------------+

All backend-specific code lives behind this interface. The core engine – the dispatch table builder, the prefill encoder, the decode loop – is pure C++ with no #import, no @autoreleasepool, no id<MTLBuffer>. If someone wanted to port akunu to CUDA, they would implement CudaDevice and everything else would just work.

This is not a hypothetical – the clean separation was a deliberate design choice. We will explore the Device interface in detail in Chapter 21.

Principle 5: Lazy Weight Loading

GGUF files can be large – a Q4_0 quantized 8B model is about 4.5 GB. Loading all weights into GPU memory at once would waste time on unused tensors and spike memory usage during initialization.

Akunu’s WeightStore uses lazy loading: when you call get_tensor("layers.5.attention.q.weight"), it checks its internal cache. If the tensor has not been loaded yet, it reads the raw bytes from the GGUF file (using memory-mapped I/O) and uploads them to a GPU buffer. Subsequent calls return the cached buffer instantly.

The WeightProvider class unifies this behind a common interface for both GGUF and MLX SafeTensors formats:

                    +-------------------+
                    | WeightProvider    |
                    | (unified facade)  |
                    +--------+----------+
                             |
                +------------+------------+
                |                         |
         +------+------+          +------+------+
         | WeightStore |          | MLXWeightStore|
         | (GGUF)      |          | (SafeTensors) |
         +------+------+          +------+--------+
                |                         |
        +-------+-------+         +------+--------+
        | GGUF mmap'd   |         | SafeTensors   |
        | file on disk   |         | + config.json |
        +----------------+         +---------------+

Weight fusion also happens here: for performance, akunu can concatenate the Q, K, and V projection matrices (or gate + up FFN matrices) into a single contiguous GPU buffer, allowing one large GEMV instead of two or three small ones. This is controlled by ChipConfig.should_fuse_weights, which is true on chips with large enough SLC (System Level Cache).

Project Structure

The project is about 265 source files (C++, Objective-C++, Metal) plus around 135 Metal shader files. Here is how they are organized:

akunu/
|
+-- include/akunu/          Public C API headers
|   +-- akunu.h             Opaque handle API (load, generate, encode, etc.)
|   +-- types.h             POD structs (AkunuModelConfig, AkunuGenerationStats, etc.)
|
+-- src/
|   +-- core/               Architecture-agnostic engine core
|   |   +-- device.h        Virtual GPU device interface
|   |   +-- dispatch_table.h  DispatchCmd + DispatchTable + encode_chain()
|   |   +-- table_builder.h/cpp  Builds dispatch table from weights + config
|   |   +-- arch_descriptor.h    ArchDescriptor + factory functions
|   |   +-- dtype_descriptor.h   DTypeDescriptor + kernel lookup tables
|   |   +-- chip_config.h        ChipConfig (hardware tuning)
|   |   +-- prefill.h/cpp        Batched prefill (GEMM-based)
|   |
|   +-- inference/           High-level inference orchestration
|   |   +-- model_state.h   ModelState struct (the opaque handle's guts)
|   |   +-- model_loader.cpp  Model loading + initialization
|   |   +-- decode_loop.cpp   Top-level generate loop (prefill + decode)
|   |   +-- decode_greedy.cpp   Chain decode (greedy, zero-alloc)
|   |   +-- decode_sampled.cpp  Sampled decode (top-k/p/min-p)
|   |   +-- decode_speculative.cpp  N-gram speculative decode
|   |   +-- decode_grammar.cpp  Grammar-constrained decode
|   |   +-- sampling.cpp       CPU-side sampling (softmax, top-k, top-p)
|   |   +-- embedding.cpp      BERT-style embedding extraction
|   |
|   +-- cache/               Memory management
|   |   +-- kv_cache.h       Per-layer K/V buffer arrays
|   |   +-- scratch.h        Pre-allocated scratch buffers
|   |   +-- whisper_buffers.h  Whisper-specific encoder buffers
|   |
|   +-- weight/              Weight file I/O
|   |   +-- weight_provider.h   Unified GGUF/MLX interface
|   |   +-- weight_store.h/cpp  GGUF weight loading + fusion
|   |   +-- mlx_weight_store.h/cpp  MLX SafeTensors loading
|   |   +-- gguf_parser.h/cpp  Low-level GGUF format parsing
|   |   +-- safetensors_parser.h  SafeTensors header parsing
|   |
|   +-- tokenizer/           BPE tokenizer
|   +-- grammar/             Grammar-constrained decoding (GBNF + XGrammar)
|   +-- whisper/             Whisper encoder + decoder
|   +-- audio/               Mel spectrogram computation
|   +-- server/              OpenAI-compatible HTTP server
|   +-- speculative/         N-gram draft predictor
|   +-- akunu_api.cpp        C API implementation (thin wrappers)
|
+-- backend/
|   +-- metal/
|       +-- metal_device.h/mm   MetalDevice implementation (ObjC++)
|       +-- metal_device_impl.h AkunuMetalState (ObjC wrapper)
|       +-- metal_types.h       Metal-specific type aliases
|       +-- kernels/            ~135 Metal shader files
|           +-- metal/kernel/
|               +-- attention/   Flash attention (prefill + decode variants)
|               +-- matmul/      GEMV + GEMM for all quant types
|               +-- norm/        RMSNorm, LayerNorm, head norms
|               +-- rope/        RoPE (interleaved + NeoX + fused variants)
|               +-- activation/  SiLU, GELU, gated variants
|               +-- embedding/   Token embedding lookup (all dtypes)
|               +-- sampling/    GPU-side argmax, top-k, temperature, penalties
|               +-- convert/     Dtype conversion (F32<->F16, dequant)
|               +-- conv/        Conv1D for Whisper frontend
|               +-- fused/       Fused kernels (GEMV+norm, whisper GEMV)
|
+-- tools/                   CLI executables
|   +-- akunu_chat.cpp       Interactive chat
|   +-- akunu_bench.cpp      llama-bench compatible benchmark
|   +-- akunu_complete.cpp   Text completion
|   +-- akunu_inspect.cpp    Model weight inspector
|   +-- akunu_profile.cpp    Per-layer GPU profiler
|   +-- akunu_serve.cpp      OpenAI-compatible HTTP server
|   +-- akunu_transcribe.cpp Whisper transcription
|   +-- akunu_benchmark.cpp  Extended benchmarking
|
+-- tests/                   Unit + integration tests
+-- 3rdparty/                XGrammar submodule
+-- bindings/                Language bindings (Swift)
+-- CMakeLists.txt           Build system
+-- Makefile                 Top-level build driver

If you count the lines of actual akunu code (excluding 3rdparty), the core engine is roughly 9,750 lines of C++ and Objective-C++, plus about 135 Metal shader files. That is remarkably compact for what it does – a full inference engine supporting 6 architectures, 2 weight formats, 16+ quantization types, grammar-constrained decoding, speculative decoding, Whisper transcription, and an HTTP server.

High-Level Inference Flow

Let us trace what happens when you call akunu_generate() with a prompt. This is the 30,000-foot view; later chapters will zoom in on each step.

akunu_generate(model, prompt_tokens, n_prompt, max_tokens, sampling, callback, ...)
|
v
run_decode_loop(state, ...)
|
+-- 1. PREFILL (batched, GEMM-based)
|   |
|   |   for chunk in prompt_tokens (up to max_prefill_chunk at a time):
|   |     encode_prefill(device, weights, config, arch, kv_cache, scratch,
|   |                    chunk_tokens, chunk_size, position)
|   |     |
|   |     |   For each layer:
|   |     |     GEMM: batch_residual @ Q_weight  -> batch_q
|   |     |     GEMM: batch_residual @ K_weight  -> batch_k
|   |     |     GEMM: batch_residual @ V_weight  -> batch_v
|   |     |     RoPE + write to KV cache
|   |     |     Flash attention (prefill variant)
|   |     |     GEMM: attn_out @ O_weight        -> batch_h1
|   |     |     Residual add + RMSNorm
|   |     |     GEMM: batch_residual @ gate_weight -> batch_gate
|   |     |     GEMM: batch_residual @ up_weight   -> batch_up
|   |     |     Activation (SiLU*gate or GELU*gate)
|   |     |     GEMM: batch_act @ down_weight      -> batch_h1
|   |     |     Residual add + RMSNorm (next layer)
|   |     |
|   |     v
|   |     Output norm + logit projection + argmax -> first token
|   |
|   v
|   Return first_token + timing stats
|
+-- 2. DECODE (chain decode, GEMV-based)
    |
    |   Choose decode path:
    |     - temperature == 0 && speculation_enabled -> decode_speculative()
    |     - temperature == 0                        -> decode_greedy()
    |     - temperature > 0                         -> decode_sampled()
    |     - grammar != null                         -> decode_grammar()
    |
    |   decode_greedy (hot path):
    |     while generated < max_tokens:
    |       write next_token to token_ids buffer
    |       device.begin_encoding()
    |       device.encode_dispatch_table(&dispatch_table, position, chunk_size)
    |       device.end_encoding_sync()   (or async with double buffering)
    |       kv_cache.advance(chunk_size)
    |       read token_ids buffer -> output tokens
    |       for each token: callback(token_id, text, user_data)
    |
    v
    Return AkunuGenerationStats { prefill_time, decode_time, tokens/sec, ... }

A few things to notice:

Prefill uses GEMM, decode uses GEMV. During prefill, we process many tokens at once, so the Q/K/V projections are matrix-matrix multiplications (M > 1). During decode, we process one token at a time, so they are matrix-vector multiplications (M = 1). Akunu has separate optimized kernels for each.

Chain decode generates multiple tokens per GPU submission. Instead of submitting one command buffer per token, akunu submits a “chain” of N tokens in a single begin_encoding() / end_encoding() pair. The dispatch table is replayed N times with patched position values. This amortizes the Metal command buffer overhead across many tokens. The chain size is tuned per chip (64-128 tokens, see ChipConfig).

The callback is synchronous. After each GPU chunk completes, tokens are read back from the token_ids buffer and delivered to the user’s callback one at a time. The callback can return false to stop generation early.

No Python in the loop. The entire flow – from tokenization through GPU dispatch through token decoding – is C++. The C API boundary is a thin wrapper in akunu_api.cpp that casts the opaque void* handle to ModelState* and forwards the call.

The ModelState: What Lives Behind the Opaque Handle

When you call akunu_load_model(), it returns an akunu_model_t, which is a void* pointing to a ModelState struct. This is the central state object that ties everything together:

struct ModelState {
    std::unique_ptr<Device> device;     // GPU device (MetalDevice)
    WeightProvider *weights;            // Weight file access
    Tokenizer tokenizer;                // BPE tokenizer
    AkunuModelConfig config;            // Parsed model config
    ArchDescriptor arch;                // Architecture descriptor
    ChipConfig chip;                    // Hardware tuning params
    KVCache kv_cache;                   // Per-layer K/V buffers
    ScratchBuffers scratch;             // Pre-allocated intermediates
    DispatchTable dispatch_table;       // Precompiled decode commands
    NGramPredictor predictor;           // Speculative n-gram predictor
    bool speculation_enabled;           // Whether spec decode is on

    // Whisper-specific fields
    bool is_whisper;
    std::unique_ptr<WhisperBuffers> whisper_buf;
    std::unique_ptr<MelSpectrogram> mel_spec;
    std::unique_ptr<WhisperModel> whisper_model;
    DispatchTable whisper_decode_table;
    // ... beam search buffers
};

This is it. One struct, one allocation. The entire engine state fits in a single cache-friendly object. Compare this to inference frameworks that scatter state across dozens of Python objects, each with its own reference counting and garbage collection pressure.

Model Loading: What Happens at Init Time

The akunu_load_model() function is where all the expensive work happens. Here is the sequence:

akunu_load_model(path, metallib_path, max_context)
|
+-- 1. Create MetalDevice (MTLCreateSystemDefaultDevice)
+-- 2. Load metallib (compiled shader library)
+-- 3. Open weight file (GGUF or MLX SafeTensors)
+-- 4. Parse model config (dims, layers, heads, vocab, etc.)
+-- 5. Select ArchDescriptor (arch_from_config)
+-- 6. Detect ChipConfig (GPU cores, family, SLC estimate)
+-- 7. Handle format-specific quirks:
|      - MLX LLaMA: switch to NeoX RoPE
|      - Tie embeddings if output.weight missing
|      - Set quant_bits / quant_group_size from MLX metadata
+-- 8. Precompute RoPE frequencies (LLaMA 3 wavelen scaling)
+-- 9. Load tokenizer (from GGUF metadata or HF tokenizer.json)
+-- 10. Set context length (capped at model max or user-specified)
+-- 11. Allocate KV cache (n_layers * 2 buffers * max_seq_len)
+-- 12. Allocate ScratchBuffers (all intermediates)
+-- 13. Build DispatchTable (resolves all PSOs, binds buffers)
+-- 14. Warmup pass (compiles remaining Metal pipelines)
+-- 15. Return ModelState* as opaque handle

Steps 1-12 are straightforward initialization. Step 13 is where the magic happens: build_dispatch_table() walks through the entire forward pass – embedding, norms, projections, RoPE, attention, FFN – and for each operation, it resolves the kernel name from DTypeDescriptor, looks up or compiles the Metal pipeline state object, computes the dispatch geometry, binds the weight and scratch buffers, and stores everything in a DispatchCmd. By the time this function returns, the engine knows exactly what to do for each token – no runtime decisions left.

Supported Architectures at a Glance

Let us briefly survey how each supported architecture maps to akunu’s abstractions:

+----------+-------------+----------+----------+--------+--------+--------+
| Arch     | Activation  | RoPE     | QK Norm  | PostNm | Enc/Dec| Embed  |
|          |             |          |          |        |        | Scale  |
+----------+-------------+----------+----------+--------+--------+--------+
| LLaMA    | silu_gate   | interleav| no       | no     | no     | 0      |
| Qwen3    | silu_gate   | neox     | yes      | no     | no     | 0      |
| Gemma    | gelu_gate   | neox     | yes      | yes    | no     | sqrt(d)|
| Gemma3   | gelu_gate   | neox     | yes      | yes    | no     | sqrt(d)|
| Whisper  | gelu (plain)| none     | no       | no     | yes    | 0      |
| BERT     | silu_gate   | neox     | no       | no     | no     | 0      |
+----------+-------------+----------+----------+--------+--------+--------+

All of these differences are captured in the ArchDescriptor struct – no special code paths. Gemma 3’s sliding window attention with alternating global/local layers? That is just cfg.sliding_window_pattern > 0 in the RoPE theta selection. Whisper’s Conv1D frontend, cross-attention, and sinusoidal positional embeddings? Those are flag fields in ArchDescriptor plus a separate WhisperModel loading path.

The key insight is that most transformer architectures are minor variations on the same theme. They all have embedding, norm, QKV projection, attention, output projection, FFN, and output logits. The differences are in which activation function, which RoPE variant, whether there is an extra normalization step, and so on. Akunu exploits this regularity by parameterizing the differences rather than branching on them.

What Makes This Different from Other Engines

If you have used llama.cpp, MLX, or vLLM, you might wonder what akunu does differently. Here is a quick comparison:

vs. llama.cpp: llama.cpp builds a computation graph (ggml) at runtime and evaluates it node-by-node. Each node involves a virtual dispatch to find the right kernel, plus buffer management. Akunu eliminates this overhead by precompiling the entire forward pass into a flat command array. llama.cpp is more general (it runs on CPU, CUDA, Metal, Vulkan, etc.), but akunu squeezes more performance out of Metal specifically.

vs. MLX: MLX is a general-purpose array framework (like PyTorch) that happens to run on Metal. It has a JIT compiler, automatic differentiation, and a Python frontend. This generality comes at a cost: each operation goes through MLX’s dispatch layer, which involves hash lookups and potentially JIT compilation. Akunu bypasses all of this – it talks directly to Metal with precompiled pipelines.

vs. vLLM: vLLM targets datacenter GPU inference with features like PagedAttention, continuous batching, and multi-GPU tensor parallelism. Akunu targets single-device Apple Silicon with features like chain decode, SLC-aware weight fusion, and Metal-specific kernel optimization. Different tools for different jobs.

The common thread is specialization. Akunu does fewer things, but does them very well on one specific hardware platform.

A Note on Code Style

Before we dive deeper in the following chapters, a word about the codebase style. Akunu is written in C++17 with a strong preference for:

• POD structs over class hierarchies (DispatchCmd, KVCache, ScratchBuffers)
• Fixed-size inline storage over heap allocation (param_bytes[64], buffers[8])
• Factory functions over constructors (KVCache::create, ScratchBuffers::create)
• Explicit state over hidden globals (ModelState holds everything)
• One virtual class (Device) instead of a deep hierarchy
• Thread-local for truly per-thread state (error buffer, RNG)

The Metal backend uses Objective-C++ (.mm files) because it has to – Metal is an Objective-C API. But this is strictly quarantined behind the Device interface. The rest of the engine is pure C++ that compiles with any standard compiler.

Error handling is C-style: functions return null/false on failure and set a thread-local error string via set_error(). No exceptions in the hot path. No RAII wrappers around GPU resources (the ModelState destructor handles cleanup).

This style is not “modern C++” in the Herb Sutter sense, but it is effective for systems programming where you care about memory layout, cache behavior, and predictable performance.

Summary

Akunu is a tightly-focused inference engine that trades generality for performance on Apple Silicon. Its key design decisions are:

1. ArchDescriptor – all architecture differences as data, not branches
2. DispatchTable – precompiled GPU command sequences, replayed per token
3. Zero-allocation decode – all buffers pre-allocated, nothing on the hot path
4. Virtual Device – clean Metal abstraction, ready for future backends
5. Lazy weight loading – tensors loaded on demand, with fusion support

The result is an engine that achieves 1.8x average speedup over llama.cpp on decode and 1.17x over MLX, in about 10,000 lines of C++ plus 135 Metal shaders.

In the next chapter, we will see how to build and run akunu from source.

Building and Running Akunu

This chapter covers the practical mechanics of building akunu from source and using its CLI tools. If you have been reading the previous chapters to understand how akunu works, this is where you roll up your sleeves and actually compile it. We will walk through the CMake configuration, build options, Metal shader compilation, and each of the CLI tools.

Prerequisites

Akunu requires:

• macOS 13 (Ventura) or later – for Metal 3 support and Apple GPU family 7+
• Xcode 15+ (or at least the Command Line Tools) – for the clang++ compiler with Objective-C++ support and the metal shader compiler
• CMake 3.20+ – the build system
• Apple Silicon Mac – M1 or later. Akunu’s Metal backend requires UMA and simdgroup_matrix support (GPU family 7+)
• A model file – either GGUF format (from llama.cpp ecosystem) or MLX SafeTensors directory

Optional dependencies:

• XGrammar (v0.1.33) – for grammar-constrained JSON output. Included as a git submodule at 3rdparty/xgrammar

Project Structure

akunu/
├── CMakeLists.txt           # Top-level build configuration
├── include/
│   └── akunu/
│       ├── akunu.h          # C API header
│       └── types.h          # Shared type definitions
├── src/
│   ├── akunu_api.cpp        # C API implementation
│   ├── core/                # Backend-agnostic core
│   │   ├── device.h         # Device abstraction
│   │   ├── dispatch_table.h # Precompiled command sequence
│   │   ├── table_builder.cpp
│   │   ├── arch_descriptor.h
│   │   ├── chip_config.h
│   │   ├── dtype_descriptor.h
│   │   └── ...
│   ├── weight/              # GGUF/MLX weight loading
│   ├── tokenizer/           # BPE tokenizer
│   ├── grammar/             # GBNF/JSON schema grammar
│   ├── inference/           # Decode loops, sampling
│   ├── cache/               # KV cache, scratch buffers
│   ├── server/              # HTTP server
│   ├── speculative/         # Speculative decoding
│   └── whisper/             # Whisper speech-to-text
├── backend/
│   └── metal/
│       ├── metal_device.h
│       ├── metal_device.mm  # ObjC++ Metal implementation
│       ├── metal_device_impl.h
│       └── kernels/         # .metal shader source files
├── tools/                   # CLI executables
│   ├── akunu_chat.cpp
│   ├── akunu_bench.cpp
│   ├── akunu_complete.cpp
│   ├── akunu_inspect.cpp
│   ├── akunu_profile.cpp
│   ├── akunu_benchmark.cpp
│   ├── akunu_serve.cpp
│   └── akunu_transcribe.cpp
├── tests/                   # Test executables
│   ├── kernels/             # Per-kernel correctness tests
│   └── ...
└── 3rdparty/
    └── xgrammar/            # Git submodule

CMake Configuration

The build is configured through CMakeLists.txt. Let’s walk through the key sections.

Language Standards

cmake_minimum_required(VERSION 3.20)
project(akunu VERSION 0.1 LANGUAGES CXX OBJCXX)

set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
set(CMAKE_OBJCXX_STANDARD 17)
set(CMAKE_OBJCXX_FLAGS "${CMAKE_OBJCXX_FLAGS} -fobjc-arc")

Akunu uses C++17 and Objective-C++ 17. The -fobjc-arc flag enables Automatic Reference Counting for Objective-C objects – this is how MetalDevice manages Metal API objects (MTLDevice, MTLCommandBuffer, etc.) without manual retain/release.¹

Backend Selection

option(AKUNU_BACKEND_METAL "Build Metal backend (Apple Silicon)" ON)
option(AKUNU_BACKEND_CUDA  "Build CUDA backend (NVIDIA)" OFF)

Metal is enabled by default. CUDA exists as a placeholder for future work. The backend selection determines which source files and frameworks are linked.

Core Sources

The core engine is pure C++ (no platform dependencies):

set(CORE_SOURCES
    src/weight/gguf_parser.cpp        # GGUF file format parser
    src/weight/weight_store.cpp       # Weight management + fusion
    src/weight/mlx_weight_store.cpp   # MLX SafeTensors parser
    src/core/table_builder.cpp        # Dispatch table construction
    src/core/device_defaults.cpp      # Default Device method implementations
    src/core/prefill.cpp              # Prefill (batched) forward pass
    src/tokenizer/tokenizer.cpp       # BPE tokenizer
    src/grammar/grammar.cpp           # GBNF grammar parser
    src/grammar/json_schema_to_grammar.cpp
    src/grammar/xgrammar_wrapper.cpp  # XGrammar integration
    src/inference/model_state.cpp     # Model lifecycle
    src/inference/sampling.cpp        # Temperature, top-k, top-p, min-p
    src/inference/model_loader.cpp    # Model loading orchestration
    src/inference/decode_greedy.cpp   # Greedy decode loop
    src/inference/decode_sampled.cpp  # Sampled decode loop
    src/inference/decode_speculative.cpp
    src/inference/decode_grammar.cpp  # Grammar-constrained decode
    src/inference/decode_loop.cpp     # Common decode infrastructure
    src/inference/embedding.cpp       # Text embedding (BERT)
    src/whisper/whisper_inference.cpp # Whisper transcription
    src/akunu_api.cpp                 # C API implementation
)

Metal Backend

When AKUNU_BACKEND_METAL is ON:

if(AKUNU_BACKEND_METAL)
    list(APPEND BACKEND_SOURCES backend/metal/metal_device.mm)
endif()

The Metal backend is a single Objective-C++ file (metal_device.mm) that implements the Device virtual interface.

Framework Linking

The Metal backend links five Apple frameworks:

target_link_libraries(akunu_engine PUBLIC
    "-framework Metal"                      # GPU compute
    "-framework MetalPerformanceShaders"    # (available for future use)
    "-framework Foundation"                 # NSObject, NSString, NSURL
    "-framework Accelerate"                 # vDSP (audio processing for Whisper)
    "-framework IOKit"                      # GPU core count detection
)

XGrammar Integration

The XGrammar submodule provides grammar-constrained decoding:

set(XGRAMMAR_DIR "${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/xgrammar")
if(EXISTS "${XGRAMMAR_DIR}/include/xgrammar/xgrammar.h")
    add_subdirectory(${XGRAMMAR_DIR} ${CMAKE_BINARY_DIR}/xgrammar EXCLUDE_FROM_ALL)
    set(AKUNU_HAS_XGRAMMAR ON)
endif()

If the submodule is not initialized, XGrammar is simply disabled and grammar-constrained generation will not be available. To enable it:

git submodule update --init --recursive

Shared Library for Bindings

option(AKUNU_BUILD_SHARED "Build shared library for language bindings" OFF)

When enabled, this builds libakunu.dylib in addition to the static libakunu_engine.a. The shared library exposes the C API (akunu.h) and can be loaded by Python, Swift, or any language with C FFI support.

Building from Source

Basic Build

cd ~/Projects/akunu
mkdir build && cd build
cmake .. -DCMAKE_BUILD_TYPE=Release
make -j$(sysctl -n hw.ncpu)

This produces:

• libakunu_engine.a – the static library
• akunu_chat, akunu_bench, akunu_complete, etc. – CLI tools
• akunu_test_*, akunu_kernel_* – test executables

Build with XGrammar

git submodule update --init --recursive
mkdir build && cd build
cmake .. -DCMAKE_BUILD_TYPE=Release
make -j$(sysctl -n hw.ncpu)

The build system auto-detects XGrammar and sets AKUNU_HAS_XGRAMMAR=1.

Build Shared Library

cmake .. -DCMAKE_BUILD_TYPE=Release -DAKUNU_BUILD_SHARED=ON
make -j$(sysctl -n hw.ncpu)

Debug Build

cmake .. -DCMAKE_BUILD_TYPE=Debug
make -j$(sysctl -n hw.ncpu)

Debug builds enable assertions and disable optimizations. For Metal shader debugging, also enable the Metal validation layer:

export METAL_DEVICE_WRAPPER_TYPE=1
export MTL_DEBUG_LAYER=1
./akunu_chat model.gguf akunu.metallib

Metal Shader Compilation

The Metal shader sources live in backend/metal/kernels/. They must be compiled into a .metallib file before akunu can load them. The compilation pipeline is:

.metal source files
    │
    ▼ (metal compiler)
.air intermediate files
    │
    ▼ (metallib archiver)
akunu.metallib

The compilation command (not automated by CMake – you need to do this manually or via a script):

# Compile all .metal files to .air
xcrun -sdk macosx metal -c -target air64-apple-macos13.0 \
    -I backend/metal/kernels/metal/include \
    backend/metal/kernels/metal/kernel/**/*.metal \
    -o kernels.air

# Archive into metallib
xcrun -sdk macosx metallib kernels.air -o akunu.metallib

The -I flag adds the include directory for shared headers (ShaderTypes.h, KernelCommon.h) that are used across kernel files.

The resulting akunu.metallib file contains all GPU kernels in compiled form. At runtime, MetalDevice::load_library() loads this file and individual kernels are extracted by name via get_pipeline().²

Kernel Organization

The Metal kernels are organized by category:

Total: roughly 100+ kernel functions compiled into a single metallib.

CLI Tools

akunu_chat

Interactive chat with a loaded model. Handles conversation formatting using the model’s native chat template.

./akunu_chat path/to/model.gguf path/to/akunu.metallib

Features:

• Auto-detects chat template (ChatML, Llama 3, Gemma, etc.)
• Multi-turn conversation with KV cache reuse
• Streaming token output
• System prompt support

akunu_bench

Performance benchmarking tool. Measures prefill and decode throughput.

./akunu_bench path/to/model.gguf path/to/akunu.metallib

Reports:

• Prefill tokens/second (for various prompt lengths)
• Decode tokens/second (steady-state generation)
• Memory usage
• Model configuration summary

akunu_complete

Text completion (non-chat). Takes a prompt and generates a continuation.

./akunu_complete path/to/model.gguf path/to/akunu.metallib

Useful for testing raw model behavior without chat formatting.

akunu_inspect

Model inspection tool. Dumps model metadata and tensor information.

./akunu_inspect path/to/model.gguf

Shows:

• Architecture, vocabulary size, embedding dimension
• Number of layers, heads, KV heads
• RoPE configuration
• Tensor names, shapes, and dtypes

akunu_profile

Per-layer GPU timing profiler. Runs each operation in its own command buffer for accurate timing.

./akunu_profile path/to/model.gguf path/to/akunu.metallib

Output shows per-operation GPU time in milliseconds, useful for identifying bottlenecks.

akunu_serve

HTTP API server with OpenAI-compatible endpoints.

./akunu_serve path/to/model.gguf path/to/akunu.metallib --port 8080

Provides:

• /v1/chat/completions – streaming and non-streaming chat
• /v1/completions – text completion
• Multi-model support (load multiple models)
• Concurrent request handling with per-model mutex

akunu_transcribe

Speech-to-text using Whisper models.

./akunu_transcribe path/to/whisper-model.gguf path/to/akunu.metallib input.wav

Supports:

• WAV input (resampled to 16kHz internally)
• Language detection or forced language
• Timestamp generation
• Streaming segment callback

akunu_benchmark

Extended benchmarking tool with more detailed metrics.

./akunu_benchmark path/to/model.gguf path/to/akunu.metallib

Library Targets

The CMake build produces two main library targets:

Both expose the same C API defined in include/akunu/akunu.h. The static library is used by all CLI tools and tests. The shared library is intended for language bindings.

Test Executables

The build produces numerous test executables:

Integration Tests

Kernel Tests

Individual kernel correctness tests (each compiled as ObjC++):

These tests compare GPU kernel output against CPU reference implementations to verify correctness within FP16 tolerance.

Troubleshooting

“Failed to load metallib”

The metallib path is incorrect or the file was compiled for a different target. Ensure:

1. The metallib file exists at the specified path
2. It was compiled with -target air64-apple-macos13.0 or later
3. The Metal device supports the required GPU family

“Failed to get pipeline: kernel_name”

A kernel function is missing from the metallib. This usually means:

1. The kernel source file was not included in the metallib compilation
2. There is a naming mismatch between the kernel function name in Metal and the string in C++

“allocate: failed to allocate N bytes”

The model is too large for available memory. Options:

1. Use a smaller quantization (Q4_0 instead of FP16)
2. Reduce max_context to shrink KV cache
3. Close other applications to free memory
4. Use a chip with more unified memory

Build Errors with XGrammar

If XGrammar fails to build, you can disable it:

cmake .. -DCMAKE_BUILD_TYPE=Release
# XGrammar auto-disables if submodule is not initialized

Summary

Building akunu is a standard CMake workflow. The main moving parts are:

1. CMake configuration with AKUNU_BACKEND_METAL=ON (default)
2. Metal shader compilation into akunu.metallib (manual step)
3. Framework linking for Metal, Foundation, Accelerate, IOKit
4. CLI tools for chat, benchmark, profiling, serving, and transcription

The next chapter covers the C API that all these tools are built on.

1. Apple, “Transitioning to ARC Release Notes.” ARC (Automatic Reference Counting) eliminates manual retain/release calls for Objective-C objects. The compiler inserts retain/release operations automatically. See https://developer.apple.com/library/archive/releasenotes/ObjectiveC/RN-TransitioningToARC/. ↩

2. Apple, “Building a Library with Metal’s Command-Line Tools.” The metal and metallib command-line tools compile .metal sources into .metallib archives. See https://developer.apple.com/documentation/metal/shader_libraries/building_a_shader_library_by_precompiling_source_files. ↩

The C API

Every akunu CLI tool – akunu_chat, akunu_bench, akunu_serve, akunu_transcribe – is built on top of a single C API defined in two header files: include/akunu/akunu.h and include/akunu/types.h. This chapter explains why akunu uses a C API rather than a C++ one, walks through every major function, and shows how to build a complete application from scratch.

Why C?

This is a C++ project. The core engine is written in C++17 with Objective-C++ for the Metal backend. So why is the public API in plain C?

1. FFI compatibility. Every programming language can call C functions. Python has ctypes and cffi. Swift has direct C interop. Rust has extern "C". Java has JNI. Go has cgo. A C API is the universal adapter.¹

2. ABI stability. C++ name mangling, vtable layouts, and standard library implementations differ between compilers and versions. A C API with POD structs and opaque pointers has a stable ABI – you can swap out the shared library without recompiling the caller.

3. No header dependencies. The C API headers include only <stdint.h>, <stdbool.h>, and <stddef.h>. No C++ standard library, no Metal headers, no Objective-C. Any C or C++ compiler can parse them.

4. Opaque handles prevent misuse. The caller cannot poke at internal state because the model is just a void*. This forces all interaction through the API functions, making it possible to change internal representations without breaking callers.

Header Organization

The API is split across two files:

types.h – Shared Data Structures

// types.h contains POD structs used by both the C API and internal C++ code

typedef struct {
    uint32_t dim;              // embedding dimension
    uint32_t n_layers;         // transformer layers
    uint32_t n_heads;          // query heads
    uint32_t n_kv_heads;       // key/value heads (GQA)
    uint32_t head_dim;         // dim per head
    uint32_t q_dim;            // total Q projection output
    uint32_t kv_dim;           // total KV projection output
    uint32_t ffn_dim;          // feed-forward intermediate dimension
    uint32_t vocab_size;       // vocabulary size
    uint32_t max_seq_len;      // maximum context length
    float norm_eps;            // RMSNorm/LayerNorm epsilon
    float rope_theta;          // RoPE base frequency
    uint32_t sliding_window_pattern;
    float rope_local_theta;
    char architecture[32];     // "llama", "qwen3", "gemma", "whisper"

    // Encoder parameters (0 = decoder-only)
    uint32_t enc_n_layers;
    uint32_t enc_n_heads;
    uint32_t enc_dim;
    uint32_t enc_ffn_dim;
    uint32_t n_mels;           // mel spectrogram bins (Whisper)
    uint32_t enc_max_seq_len;
} AkunuModelConfig;

Notice that AkunuModelConfig uses fixed-size arrays (char architecture[32]) instead of std::string, and all fields are primitive types. This is a POD struct – it can be safely passed across C/C++ boundaries and even memory-mapped.

Other types in types.h:

akunu.h – The Function API

The main header declares all API functions inside an extern "C" block:

#ifdef __cplusplus
extern "C" {
#endif

typedef void *akunu_model_t;  // Opaque model handle

// ... function declarations ...

#ifdef __cplusplus
}
#endif

The extern "C" block ensures C linkage (no name mangling) when compiled as C++. The #ifdef __cplusplus guards make the header valid for both C and C++ compilers.

The Opaque Handle Pattern

The entire model state – weights, KV cache, scratch buffers, dispatch table, tokenizer, device – is wrapped behind a single opaque handle:

typedef void *akunu_model_t;

This is a pointer to an internal C++ ModelState object that the caller never sees. Every API function takes this handle as its first argument. The pattern is:

// Create
akunu_model_t model = akunu_load_model("model.gguf", "akunu.metallib", 0);

// Use
akunu_generate(model, tokens, n_tokens, 256, sampling, callback, NULL);

// Destroy
akunu_free_model(model);

No global state, no singletons. You can load multiple models simultaneously by creating multiple handles. Each handle owns its own GPU resources.²

Model Lifecycle

Loading

akunu_model_t akunu_load_model(const char *model_path,
                                const char *metallib_path,
                                int max_context);

This function does a lot of work:

1. Creates a MetalDevice (or default device)
2. Loads the metallib (device.load_library())
3. Parses the model file (GGUF or MLX SafeTensors)
4. Allocates weight buffers and uploads weights to GPU
5. Creates the ArchDescriptor from model metadata
6. Queries ChipConfig from the device
7. Allocates KV cache and scratch buffers
8. Builds the dispatch table (build_dispatch_table())
9. Initializes the tokenizer
10. Returns the opaque handle (or NULL on failure)

Parameters:

• model_path: Path to a .gguf file or MLX SafeTensors directory
• metallib_path: Path to compiled akunu.metallib. Pass NULL for auto-detection (searches common paths)
• max_context: Maximum context window. 0 = use model default (capped at 4096)

Freeing

void akunu_free_model(akunu_model_t model);

Releases all GPU buffers, KV cache, scratch buffers, cached pipeline state objects, and the Metal device. After this call, the handle is invalid.

Error Handling

const char *akunu_get_error(void);

Returns the last error message. This uses thread-local storage, so it is safe to call from multiple threads. If akunu_load_model returns NULL, call this to find out why:

akunu_model_t model = akunu_load_model("bad_path.gguf", "akunu.metallib", 0);
if (!model) {
    printf("Error: %s\n", akunu_get_error());
    // "Error: Failed to open file: bad_path.gguf"
}

Model Information

AkunuModelConfig akunu_get_config(akunu_model_t model);
size_t akunu_model_memory(akunu_model_t model);

akunu_get_config returns a copy of the model configuration struct. Since AkunuModelConfig is a POD struct, this is a simple memcpy – no dynamic allocation.

akunu_model_memory returns the total GPU memory used by the model in bytes. This includes weights, KV cache, scratch buffers, and pre-allocated parameter buffers.

Tokenization

int akunu_encode(akunu_model_t model, const char *text,
                 uint32_t *out_tokens, int max_tokens);

const char *akunu_decode_token(akunu_model_t model, uint32_t token_id);

int akunu_token_count(akunu_model_t model, const char *text);

The tokenizer is a BPE implementation built into akunu (no external dependency). Token IDs are uint32_t values.

akunu_encode writes token IDs into a caller-provided buffer. Returns the number of tokens written. If the output buffer is too small, the text is silently truncated.

akunu_decode_token returns a pointer to the token’s text representation. The pointer is valid until the model is freed – it points into the tokenizer’s vocabulary table.

akunu_token_count is a convenience function that counts tokens without allocating an output buffer.

Generation: The Callback Pattern

Generation uses a callback function for streaming output:

typedef bool (*akunu_token_callback)(uint32_t token_id,
                                     const char *text,
                                     void *user_data);

The callback is invoked for each generated token. Returning false stops generation immediately. The user_data pointer is passed through from the akunu_generate call, allowing the callback to access caller state without globals.

AkunuGenerationStats akunu_generate(
    akunu_model_t model,
    const uint32_t *prompt_tokens, int n_prompt,
    int max_tokens,
    AkunuSamplingConfig sampling,
    akunu_token_callback callback,
    void *user_data);

This is the main generation entry point. It:

1. Resets the KV cache
2. Runs prefill on the prompt tokens
3. Enters the decode loop, calling the callback for each token
4. Returns statistics (prefill time, decode time, tokens/second)

Sampling Configuration

typedef struct {
    float temperature;  // 0 = greedy (argmax)
    int top_k;          // 0 = disabled
    float top_p;        // 1.0 = disabled
    float min_p;        // 0.0 = disabled
    float repeat_penalty;  // 1.0 = disabled
} AkunuSamplingConfig;

Temperature 0 triggers the greedy decode path (argmax on GPU, no CPU sampling). Non-zero temperature runs the sampled decode path with optional top-k, top-p, and min-p filtering.

Generation Statistics

typedef struct {
    int prompt_tokens;
    int generated_tokens;
    float prefill_time_ms;
    float decode_time_ms;
    float prefill_tokens_per_sec;
    float decode_tokens_per_sec;
} AkunuGenerationStats;

This struct is returned by value from akunu_generate. It contains everything you need to report performance.

Continued Generation

For multi-turn chat, you do not want to re-process the entire conversation history each turn. akunu_generate_continue extends the existing KV cache:

AkunuGenerationStats akunu_generate_continue(
    akunu_model_t model,
    const uint32_t *new_tokens, int n_new,
    int max_tokens,
    AkunuSamplingConfig sampling,
    akunu_token_callback callback,
    void *user_data);

This prefills only the new_tokens (the latest user message) and generates from the combined context. The KV cache from previous turns is preserved.

Grammar-Constrained Generation

For structured output (JSON, specific formats), akunu supports grammar-constrained decoding:

akunu_grammar_t akunu_grammar_create(akunu_model_t model, const char *gbnf);
akunu_grammar_t akunu_grammar_create_from_schema(akunu_model_t model,
                                                  const char *json_schema);
akunu_grammar_t akunu_grammar_create_json(akunu_model_t model);
void akunu_grammar_free(akunu_grammar_t grammar);

AkunuGenerationStats akunu_generate_grammar(
    akunu_model_t model,
    const uint32_t *prompt_tokens, int n_prompt,
    int max_tokens,
    AkunuSamplingConfig sampling,
    akunu_grammar_t grammar,
    akunu_token_callback callback,
    void *user_data);

The grammar handle is opaque, like the model handle. Three factory functions create grammars from GBNF strings, JSON Schema strings, or a generic JSON grammar. The grammar masks invalid tokens at each step, guaranteeing the output conforms to the grammar.³

Low-Level API

For benchmarking and custom decode loops, akunu exposes lower-level functions:

// Run prefill, return first generated token
uint32_t akunu_prefill(akunu_model_t model,
                       const uint32_t *tokens, int n_tokens);

// Run one decode step, return next token
uint32_t akunu_decode_step(akunu_model_t model,
                           uint32_t token_id, int position);

// Chain decode: multiple tokens in one GPU submission
int akunu_chain_decode(akunu_model_t model,
                       uint32_t first_token, int start_position,
                       int count, uint32_t *out_tokens);

// Get current KV cache position
int akunu_get_position(akunu_model_t model);

// Reset KV cache
void akunu_reset(akunu_model_t model);

The akunu_chain_decode function is the key primitive for fast greedy generation. It encodes the dispatch table N times into a single command buffer, patching position fields for each token. This is how akunu achieves high throughput for greedy (temperature=0) decoding.

Speculative Decoding

void akunu_set_speculation(akunu_model_t model, bool enabled);

When enabled, the decode loop uses n-gram prediction to speculatively generate multiple tokens, then verifies them against the model. Correctly predicted tokens skip full forward passes. This only works with greedy mode (temperature=0) because the verification requires deterministic token selection.

Embeddings

For BERT-style encoder models:

int akunu_embed(akunu_model_t model,
                const uint32_t *tokens, int n_tokens,
                float *out_embedding, int max_dims);

int akunu_embed_text(akunu_model_t model, const char *text,
                     float *out_embedding, int max_dims);

int akunu_embedding_dim(akunu_model_t model);

akunu_embed runs a forward pass through the encoder, mean-pools the final hidden layer, and writes the resulting embedding vector to out_embedding. Returns the embedding dimension on success, 0 on failure.

akunu_embed_text is a convenience wrapper that tokenizes the text internally.

Whisper Transcription

const char *akunu_transcribe(akunu_model_t model,
                             const char *wav_path,
                             const char *language,
                             AkunuTranscribeStats *stats_out);

const char *akunu_transcribe_pcm(akunu_model_t model,
                                 const float *samples, int n_samples,
                                 const char *language,
                                 AkunuTranscribeStats *stats_out);

bool akunu_is_whisper(akunu_model_t model);
void akunu_set_timestamps(akunu_model_t model, bool enabled);

The transcription API supports both file-based and PCM buffer input. The returned string is valid until the next call or model free – it points to an internal buffer.

Streaming callbacks are also available:

typedef bool (*akunu_segment_callback)(int start_ms, int end_ms,
                                       const char *text, void *user_data);

const char *akunu_transcribe_stream(akunu_model_t model,
                                    const char *wav_path,
                                    const char *language,
                                    AkunuTranscribeStats *stats_out,
                                    akunu_segment_callback callback,
                                    void *user_data);

Chat Templates

const char *akunu_format_chat(akunu_model_t model,
                              const char *system_prompt,
                              const char *user_message);

const char *akunu_chat_template(akunu_model_t model);

akunu_format_chat applies the model’s native chat template to format a system prompt and user message into the expected input format (e.g., ChatML, Llama 3 format, Gemma format). The returned string is valid until the next call.

akunu_chat_template returns the template name as a string (“chatml”, “llama3”, “gemma”, or “unknown”).

Profiling

int akunu_profile_decode_step(akunu_model_t model,
                              uint32_t token_id, int position,
                              float *timing_out, int max_entries);

const char *akunu_profile_label(akunu_model_t model, int index);

The profiling API runs each operation in its own command buffer to get per-operation GPU timing. timing_out receives an array of float values (milliseconds). akunu_profile_label returns the human-readable label for each entry (e.g., “layer.0.attention”, “layer.0.o_proj”).

GPU Sampling Operations

void akunu_gpu_temperature_scale(akunu_model_t model, float temperature);
void akunu_gpu_repetition_penalty(akunu_model_t model,
                                   const uint32_t *token_ids,
                                   int n_tokens, float penalty);

These functions run sampling operations directly on the GPU, avoiding CPU readback of the logits buffer. Temperature scaling is a simple element-wise multiply; repetition penalty adjusts logits for previously seen tokens.

Model Inspection

int akunu_tensor_count(akunu_model_t model);
const char *akunu_tensor_name(akunu_model_t model, int index);
uint32_t akunu_tensor_dtype(akunu_model_t model, int index);
const char *akunu_tensor_raw_dtype(akunu_model_t model, int index);

These functions allow iterating over all tensors in the model. akunu_inspect uses them to dump the full tensor list. akunu_tensor_raw_dtype returns the original dtype string (e.g., “BF16” for SafeTensors) while akunu_tensor_dtype returns the internal GGUF dtype code.

Thread Safety

The akunu API has the following thread safety guarantees:

1. Different model handles are fully independent. You can call functions on model_A from thread 1 and model_B from thread 2 concurrently with no synchronization needed.

2. A single model handle is NOT thread-safe. You must serialize all calls to the same model. The akunu_serve server handles this with a per-model mutex.

3. akunu_get_error() is thread-safe. It uses thread-local storage.

4. Model loading (akunu_load_model) is thread-safe. Each call creates its own device and resources.

Complete Example

Here is a complete program that loads a model, generates text, and reports statistics:

#include "akunu/akunu.h"
#include <stdio.h>
#include <string.h>

static bool on_token(uint32_t token_id, const char *text, void *user_data) {
    printf("%s", text);
    fflush(stdout);
    (void)token_id;
    (void)user_data;
    return true;  // continue generating
}

int main(int argc, char **argv) {
    if (argc < 3) {
        fprintf(stderr, "Usage: %s <model.gguf> <akunu.metallib>\n", argv[0]);
        return 1;
    }

    // Load model
    akunu_model_t model = akunu_load_model(argv[1], argv[2], 4096);
    if (!model) {
        fprintf(stderr, "Failed to load model: %s\n", akunu_get_error());
        return 1;
    }

    // Print model info
    AkunuModelConfig cfg = akunu_get_config(model);
    printf("Model: %s, %u layers, %u dim, %.1f MB GPU memory\n",
           cfg.architecture, cfg.n_layers, cfg.dim,
           akunu_model_memory(model) / 1048576.0);

    // Tokenize prompt
    const char *prompt = "Explain the roofline model in one paragraph:";
    uint32_t tokens[4096];
    int n_tokens = akunu_encode(model, prompt, tokens, 4096);
    printf("Prompt: %d tokens\n\n", n_tokens);

    // Generate
    AkunuSamplingConfig sampling = {
        .temperature = 0.0f,  // greedy
        .top_k = 0,
        .top_p = 1.0f,
        .min_p = 0.0f,
        .repeat_penalty = 1.0f
    };

    AkunuGenerationStats stats = akunu_generate(
        model, tokens, n_tokens,
        256,       // max_tokens
        sampling,
        on_token,
        NULL       // user_data
    );

    // Report
    printf("\n\n--- Stats ---\n");
    printf("Prefill: %d tokens in %.1f ms (%.0f tok/s)\n",
           stats.prompt_tokens, stats.prefill_time_ms,
           stats.prefill_tokens_per_sec);
    printf("Decode:  %d tokens in %.1f ms (%.0f tok/s)\n",
           stats.generated_tokens, stats.decode_time_ms,
           stats.decode_tokens_per_sec);

    akunu_free_model(model);
    return 0;
}

Compile and run:

clang -std=c11 -I include example.c -L build -lakunu_engine \
    -framework Metal -framework Foundation -framework Accelerate \
    -framework IOKit -lstdc++ -o example

./example path/to/model.gguf path/to/akunu.metallib

API Function Reference

Summary

The C API is akunu’s external interface. It uses the opaque handle pattern, POD structs, and C linkage to provide maximum compatibility across languages and compilers. The callback-based generation pattern supports streaming output without allocating result buffers. Thread safety is per-model-handle, requiring callers to serialize access to a single model.

1. The C FFI is effectively the lingua franca of systems programming. See “Foreign Function Interface” on Wikipedia for a survey of language support. Every major language runtime supports calling C functions with zero or minimal overhead. See https://en.wikipedia.org/wiki/Foreign_function_interface. ↩

2. This “handle + function” pattern is sometimes called the “C object pattern” or “ADT (Abstract Data Type) in C.” It provides encapsulation without language-level support for classes. The Linux kernel uses this pattern extensively for device drivers. ↩

3. Grammar-constrained decoding uses the XGrammar library (v0.1.33) internally. XGrammar compiles the grammar into a token mask that can be applied at each decoding step. See the XGrammar project: https://github.com/mlc-ai/xgrammar. ↩

The Device Abstraction Layer

If you look at akunu’s core code – the table builder, the prefill engine, the decode loops – you will notice something conspicuously absent: there are no Metal API calls. No id<MTLBuffer>, no [encoder setComputePipelineState:], no MTLSizeMake. The core is pure C++17, completely agnostic to the GPU backend. All hardware interaction flows through a single abstract class: Device.

This chapter examines the Device abstraction, its concrete MetalDevice implementation, and the design decisions behind it.

The Problem

An inference engine needs to:

1. Allocate GPU memory
2. Load compiled shader libraries
3. Create compute pipelines from shader functions
4. Encode sequences of GPU commands (set buffers, set parameters, dispatch)
5. Submit and synchronize

These operations look completely different across GPU APIs. Metal uses Objective-C message passing ([encoder setBuffer:...]). CUDA uses C function calls (cuLaunchKernel()). Vulkan uses verbose descriptor sets and command buffers. If the core engine hardcodes any of these, it cannot be ported.

The solution is a pure virtual interface that captures the common operations and lets each backend implement them in its native API.

The Device Interface

The Device class in src/core/device.h defines the contract. Let’s walk through it section by section.

Handles: Buffer, Pipeline, Dim3

Before the class itself, three simple structs define the currency of GPU programming:

struct Buffer {
    void *handle;    // Backend-specific (MTLBuffer*, CUdeviceptr, etc.)
    size_t size;     // Size in bytes
    void *contents;  // CPU-accessible pointer (UMA) or nullptr (discrete)
};

struct Pipeline {
    void *handle;    // Backend-specific (MTLComputePipelineState*, CUfunction, etc.)
};

struct Dim3 {
    uint32_t x, y, z;
    Dim3(uint32_t x = 1, uint32_t y = 1, uint32_t z = 1) : x(x), y(y), z(z) {}
};

These are POD types. No virtual methods, no reference counting, no destructors. Buffer is 24 bytes. Pipeline is 8 bytes. Dim3 is 12 bytes. They are cheap to copy, pass by value, and store in arrays – which matters because the dispatch table stores hundreds of them.

The void* handle pattern is the C equivalent of generics. On Metal, Buffer::handle is a CFBridgingRetain’d id<MTLBuffer>. On CUDA, it would be a CUdeviceptr. The core code never dereferences these pointers; it just passes them back to the device.

Device Info

virtual const char *name() const = 0;
virtual const char *backend_name() const = 0;  // "metal", "cuda", "vulkan"
virtual int gpu_core_count() const = 0;
virtual int gpu_family() const { return 0; }    // Apple GPU family (7=M1, 8=M2/M3, 9=M4)
virtual size_t total_memory() const = 0;

These methods expose hardware capabilities that affect algorithmic decisions. gpu_core_count() is used by ChipConfig::from_gpu() to determine SLC size estimates, GEMV variant selection, and chain decode chunk size. gpu_family() distinguishes M1/M2/M3/M4 for generation-specific tuning.

Note that gpu_family() has a default implementation returning 0. This is because non-Apple backends do not have “GPU families” – the method is Apple-specific but exposed at the abstract level because ChipConfig needs it.

Library and Pipeline Management

virtual bool load_library(const std::string& path) = 0;
virtual Pipeline get_pipeline(const std::string& name) = 0;
virtual Pipeline get_pipeline(const std::string& name,
                              const std::string& cache_key,
                              const uint32_t *constant_indices,
                              const uint32_t *constant_values,
                              int n_constants,
                              const uint32_t *constant_types = nullptr) = 0;

load_library loads a compiled shader library (metallib on Metal, PTX/CUBIN on CUDA). get_pipeline retrieves a named compute function, optionally specialized with function constants.

The two-overload design is important. The simple overload (get_pipeline(name)) handles the common case of a kernel with no specialization. The extended overload with constant_indices/constant_values handles Metal function constants and could map to CUDA template instantiation or Vulkan specialization constants.

The cache_key parameter in the extended overload is separate from name because the same kernel function can produce multiple specialized pipelines. For example, gemv_mlx_q4 with group_size=64, K=4096 gets a cache key of "gemv_mlx_q4_gs64_k4096", while the same kernel with K=2048 gets "gemv_mlx_q4_gs64_k2048". Both pipelines come from the same source function but are compiled with different constants.

Buffer Management

virtual Buffer allocate(size_t bytes) = 0;
virtual Buffer allocate(const void *data, size_t bytes) = 0;
virtual void free_buffer(Buffer buf) = 0;

virtual void write_buffer(Buffer dst, const void *src, size_t bytes, size_t offset = 0) {
    if (dst.contents) memcpy((char *)dst.contents + offset, src, bytes);
}

virtual void read_buffer(void *dst, Buffer src, size_t bytes, size_t offset = 0) {
    if (src.contents) memcpy(dst, (const char *)src.contents + offset, bytes);
}

virtual void *buffer_contents(Buffer buf) { return buf.contents; }

The two allocate overloads handle empty allocation (for scratch buffers) and initialized allocation (for uploading weight data). free_buffer releases the GPU memory.

write_buffer and read_buffer have default implementations that use memcpy – correct for UMA where buf.contents is a CPU-accessible pointer. A CUDA backend would override these to use cuMemcpyHtoD and cuMemcpyDtoH.

Command Encoding

This is the heart of the interface – the methods that encode GPU commands:

virtual void begin_encoding() = 0;
virtual void set_pipeline(Pipeline pso) = 0;
virtual void set_buffer(Buffer buf, int offset, int index) = 0;
virtual void set_bytes(const void *data, int size, int index) = 0;
virtual void set_threadgroup_memory(int bytes, int index) = 0;
virtual void dispatch(Dim3 grid, Dim3 threadgroup) = 0;
virtual void dispatch_threads(Dim3 total, Dim3 threadgroup) = 0;
virtual double end_encoding_sync() = 0;
virtual void end_encoding_async() = 0;
virtual void wait() = 0;

The encoding sequence mirrors Metal’s command encoder pattern:

begin_encoding()
  set_pipeline(pso)
  set_buffer(buf, offset, index)
  set_bytes(params, size, index)
  dispatch(grid, threadgroup)
  // ... more dispatches ...
end_encoding_sync()  // or end_encoding_async() + wait()

On Metal, these map directly:

On CUDA, the mapping would be different but conceptually similar: begin_encoding would create a CUDA stream, set_pipeline would set the current function, and dispatch would call cuLaunchKernel.

Advanced Synchronization

The interface provides several synchronization patterns beyond simple sync/async:

virtual void end_encoding_handler() { end_encoding_sync(); }
virtual void wait_handler() { wait(); }
virtual void end_encoding_event() { end_encoding_async(); }
virtual void begin_encoding_after_event() { begin_encoding(); }
virtual void wait_event() { wait(); }
virtual void wait_async() { wait(); }

These have default implementations that fall back to simple sync/async, but MetalDevice overrides them with more sophisticated patterns:

The event-based pattern is used for chain decode pipelining: while the GPU executes command buffer N, the CPU encodes command buffer N+1. The shared event ensures N+1 does not start executing until N completes, without requiring a CPU round-trip.

Dispatch Table Fast Path

virtual void encode_dispatch_table(const void *table_ptr,
                                   int start_position, int count) {
    // Default: use generic encode_chain (virtual calls per command)
}

virtual void encode_command_range(const void *table_ptr,
                                  int position,
                                  int cmd_start, int cmd_count) {}

These methods allow the backend to bypass the virtual call overhead of the generic encode_chain function. MetalDevice overrides encode_dispatch_table with a tight loop that calls Metal API functions directly on the ObjC encoder object, avoiding the per-command virtual dispatch through set_pipeline, set_buffer, etc.

The void* parameter (pointing to a DispatchTable) is a deliberate abstraction leak – the backend needs to know the table structure to iterate it efficiently. The alternative (encoding through the virtual interface) works but is slower.

Hardware Capabilities

virtual ChipConfig chip_config() const;
virtual const DTypeDescriptor& dtype_lookup(uint32_t dtype) const;
const char *embedding_kernel_for(uint32_t dtype) const;
const char *gemm_kernel_for(uint32_t dtype, int M) const;

chip_config() returns hardware tuning parameters derived from gpu_core_count() and gpu_family(). The default implementation (in device_defaults.cpp) calls ChipConfig::from_gpu():

ChipConfig Device::chip_config() const {
    return ChipConfig::from_gpu(gpu_core_count(), gpu_family());
}

dtype_lookup returns the kernel name and dispatch geometry for a given quantization format. Both have default implementations using the global tables in dtype_descriptor.h and chip_config.h.

Factory

static std::unique_ptr<Device> Device::create_default();

This static factory method creates the default device for the current platform. On macOS, it creates a MetalDevice. On a hypothetical CUDA platform, it would create a CudaDevice.

The MetalDevice Implementation

MetalDevice in backend/metal/metal_device.h implements all pure virtual methods. It is an Objective-C++ class (compiled from .mm files) that wraps Metal API objects.

Internal State

class MetalDevice : public Device {
private:
    void *device_;      // id<MTLDevice> (via CFBridgingRetain)
    void *queue_;       // id<MTLCommandQueue>
    void *library_;     // id<MTLLibrary>
    void *cmd_buffer_;  // id<MTLCommandBuffer> (current)
    void *encoder_;     // id<MTLComputeCommandEncoder> (current)
    size_t allocated_bytes_ = 0;
    std::unordered_map<std::string, void *> pso_cache_;
    std::string device_name_;
};

Notice the void* pointers – even within the MetalDevice, the ObjC objects are stored as raw pointers. This is because the header file (metal_device.h) is included by C++ translation units that cannot parse Objective-C types. The actual ObjC types are accessed through the AkunuMetalState wrapper in metal_device_impl.h:

@interface AkunuMetalState : NSObject
@property(nonatomic, strong) id<MTLDevice> device;
@property(nonatomic, strong) id<MTLCommandQueue> queue;
@property(nonatomic, strong) id<MTLLibrary> library;
@property(nonatomic, strong) id<MTLCommandBuffer> cmdBuffer;
@property(nonatomic, strong) id<MTLComputeCommandEncoder> encoder;
@property(nonatomic, strong) id<MTLCommandBuffer> prevCmdBuffer;
@property(nonatomic, strong) id<MTLSharedEvent> pipelineEvent;
@property(nonatomic, assign) uint64_t eventValue;
@end

The device_ pointer is actually a CFBridgingRetain’d reference to an AkunuMetalState instance. In the .mm implementation, a macro bridges back to the typed object:

#define STATE ((__bridge AkunuMetalState *)device_)

This pattern – ObjC state wrapped in a C++ class with void* storage – is the standard way to mix ObjC and C++ in headers that must be parseable by both compilers.¹

Pipeline Caching

std::unordered_map<std::string, void *> pso_cache_;

Every get_pipeline call first checks the cache. Pipeline creation (compiling a Metal function into a compute pipeline state) is expensive – it involves shader compilation, register allocation, and GPU resource validation. Caching ensures this happens once per kernel variant.

For function-constant-specialized pipelines, the cache key includes the specialization parameters:

Pipeline MetalDevice::get_pipeline(const std::string &name,
                                   const std::string &cache_key, ...) {
    auto it = pso_cache_.find(cache_key);
    if (it != pso_cache_.end())
        return {it->second};
    // ... create specialized pipeline ...
    pso_cache_[cache_key] = ptr;
}

A model with MLX Q4 weights might create 20+ specialized pipelines (different K dimensions for each layer’s GEMV). These are all cached after the first build_dispatch_table() call.

GPU Core Count Detection

MetalDevice queries the actual GPU core count from IOKit, not from Metal:

static int iokit_gpu_core_count() {
    io_iterator_t iter = 0;
    IOServiceGetMatchingServices(kIOMainPortDefault,
                                 IOServiceMatching("AGXAccelerator"), &iter);
    io_service_t service = IOIteratorNext(iter);
    CFNumberRef ref = IORegistryEntryCreateCFProperty(
        service, CFSTR("gpu-core-count"), kCFAllocatorDefault, 0);
    CFNumberGetValue(ref, kCFNumberIntType, &cores);
    // ...
}

Metal’s API does not expose the GPU core count directly. The AGXAccelerator IOKit service (Apple’s GPU driver) stores it as a registry property. This is more reliable than string-parsing the device name, though a fallback heuristic exists:

if ([name containsString:@"Ultra"]) cached = 60;
else if ([name containsString:@"Max"]) cached = 30;
else if ([name containsString:@"Pro"]) cached = 16;
else cached = 10;

GPU Family Detection

int MetalDevice::gpu_family() const {
    id<MTLDevice> dev = STATE.device;
    if ([dev supportsFamily:MTLGPUFamilyApple9]) return 9; // M4
    if ([dev supportsFamily:MTLGPUFamilyApple8]) return 8; // M2/M3
    if ([dev supportsFamily:MTLGPUFamilyApple7]) return 7; // M1
    return 6;
}

This is used by ChipConfig::from_gpu() to apply generation-specific tuning. M4 (family 9) has improved cache hierarchy and native BF16 support; M1 (family 7) has smaller SLC and no BF16.

Buffer Allocation

Buffer MetalDevice::allocate(size_t bytes) {
    id<MTLBuffer> buf = [STATE.device newBufferWithLength:MAX(bytes, 16)
                                                 options:MTLResourceStorageModeShared];
    void *h = (void *)CFBridgingRetain(buf);
    allocated_bytes_ += MAX(bytes, 16);
    return {h, bytes, [buf contents]};
}

Key details:

• Minimum 16 bytes (alignment requirement)
• MTLResourceStorageModeShared (UMA zero-copy)
• CFBridgingRetain prevents ARC from releasing the buffer
• [buf contents] provides the CPU-accessible pointer
• allocated_bytes_ tracks total GPU memory usage

The Encode Dispatch Table Fast Path

The most performance-critical method is encode_dispatch_table, which replays the dispatch table for chain decode. Rather than calling the virtual set_pipeline/set_buffer/dispatch methods (which go through vtable dispatch), it operates directly on the Metal encoder:

void MetalDevice::encode_dispatch_table(const void *table_ptr,
                                        int start_position, int count) {
    const DispatchCmd *cmds = table.commands.data();
    id<MTLComputeCommandEncoder> enc = STATE.encoder;

    for (int tok = 0; tok < count; tok++) {
        for (int c = 0; c < n_cmds; c++) {
            const DispatchCmd &cmd = cmds[c];
            [enc setComputePipelineState:cmd.pso.handle];
            // ... buffer binding, param patching, dispatch ...
        }
    }
}

This avoids ~6 virtual calls per command (set_pipeline, set_buffer x N, set_bytes, dispatch), which at 200 commands per token and 128 tokens per batch would be ~150,000 virtual calls per submission. The direct Metal calls are significantly faster.²

Thread Safety Note

From the header:

/// THREAD SAFETY: This class is NOT thread-safe. All encoding operations
/// must be called from a single thread. The server serializes access
/// via the per-model mutex in ModelEntry.

Metal command encoders are inherently single-threaded. The MetalDevice does not add locking because the performance cost would be unacceptable in the hot path. Callers must ensure serialized access.

Why This Abstraction Exists

Three reasons:

1. Portability. While akunu currently only has a Metal backend, the abstraction makes it possible to add CUDA, Vulkan, or WebGPU backends without touching the core engine. The dispatch table, table builder, and decode loops work identically regardless of backend.

2. Testability. A mock Device implementation can be used for unit testing without a real GPU. The table builder can be tested by checking the commands it generates, without actually dispatching them.

3. Code organization. The separation forces a clean boundary between “what to compute” (core) and “how to compute it” (backend). Objective-C++ is confined to a single .mm file. The rest of the project compiles as standard C++17.

The abstraction is intentionally thin. It does not try to abstract away GPU programming – it just abstracts away the specific API calls. You still think in terms of buffers, pipelines, threadgroups, and dispatches. This is a pragmatic choice: the alternative (a high-level “tensor operation” abstraction) would hide the performance-critical details that make the difference between 30 tok/s and 70 tok/s.

Summary

The Device abstraction layer is a pure virtual C++ class with ~20 virtual methods covering device info, buffer management, command encoding, and synchronization. MetalDevice implements it using Objective-C++ with direct Metal API calls, an ObjC state wrapper for ARC compatibility, and a pipeline cache for kernel specialization. The fast path (encode_dispatch_table) bypasses the virtual interface entirely for maximum chain decode throughput.

1. This pattern is described in Apple’s “Mixing Objective-C and C++” technical note. The key challenge is that Objective-C types (id<MTLDevice>, etc.) cannot appear in C++ headers that might be included by pure C++ translation units. The void* + bridge cast pattern is the standard workaround. See https://developer.apple.com/library/archive/documentation/Cocoa/Conceptual/ProgrammingWithObjectiveC/Introduction/Introduction.html. ↩

2. Virtual function call overhead on modern CPUs is typically 2-5ns due to the indirect branch prediction miss on the first call and potential instruction cache miss for the vtable. At 150,000 calls per submission, this adds up to 0.3-0.75ms – significant compared to the ~14ms total decode time for a 7B model. ↩

Architecture Descriptors: Data-Driven Design

Here is a problem that every multi-architecture inference engine faces: Llama uses SiLU activation and standard RoPE. Qwen3 uses SiLU but NeoX-style RoPE and per-head QK norms. Gemma uses GELU activation, NeoX RoPE, QK norms, post-attention norms, post-FFN norms, and scales the embedding output by sqrt(dim). Whisper uses LayerNorm instead of RMSNorm, has bias terms on every linear layer, uses sinusoidal positional encoding instead of RoPE, and has an encoder-decoder architecture.

How do you handle all of this without drowning in if (arch == "llama") ... else if (arch == "qwen3") ... else if (arch == "gemma") ... branches scattered across your codebase?

Akunu’s answer is the ArchDescriptor: a plain data struct that captures all architecture-specific behavior as fields. The table builder and prefill engine read from this struct – they never branch on architecture name. Adding a new architecture means writing one factory function and one case in the lookup function. Zero changes to the core engine.

The Problem in Detail

Consider the table builder. It needs to emit dispatch commands for one transformer layer. Here is a partial list of things that vary by architecture:

Without a descriptor, you would need if-else branches for each of these in every place the table builder makes an architecture-dependent decision. That is 11+ branch points per layer, across hundreds of lines of code. It is unreadable, error-prone, and a maintenance nightmare.

The ArchDescriptor Struct

The descriptor is defined in src/core/arch_descriptor.h:

struct ArchDescriptor {
    // Activation
    const char *activation_kernel;  // "silu_gate_f16", "gelu_gate_f16", or "gelu_f16"

    // Embedding
    float embedding_scale;  // 0 = no scaling, >0 = scale by this value

    // Per-head norms
    bool has_qk_norm;

    // Post-norms (Gemma-style)
    bool has_post_attn_norm;
    bool has_post_ffn_norm;
    const char *post_attn_norm_key;  // weight suffix: "post_attention_norm"
    const char *post_ffn_norm_key;   // "post_ffw_norm"

    // MLX quantization
    int quant_bits;        // 0 = not MLX, 3/4/6/8 = MLX quant bits
    int quant_group_size;  // typically 64

    // RoPE
    const char *rope_kernel;      // fused RoPE+KV write kernel name
    const char *rope_standalone;  // standalone RoPE kernel name
    Buffer rope_freqs;            // precomputed frequency divisors

    // Output
    bool tie_embeddings;  // use token_embedding.weight for logit projection

    // Encoder-decoder
    bool is_encoder_decoder;
    bool has_cross_attention;
    bool has_conv_frontend;
    bool has_bias;
    const char *norm_type;          // "rmsnorm" or "layernorm"
    const char *encoder_activation;
    bool is_embedding_model;        // BERT-style encoder-only
};

Every field is a simple type: bool, int, float, const char*, or Buffer (a POD struct). No virtual methods. No inheritance. No dynamism. The descriptor is created once and read many times.

Field Categories

The fields fall into several categories:

Kernel selection fields (activation_kernel, rope_kernel, rope_standalone, encoder_activation): These are kernel function names that get passed directly to device.get_pipeline(). The table builder does not care what activation function is used – it just dispatches whatever kernel the descriptor says.

Boolean feature flags (has_qk_norm, has_post_attn_norm, has_post_ffn_norm, tie_embeddings, is_encoder_decoder, has_cross_attention, has_conv_frontend, has_bias, is_embedding_model): These control whether certain dispatch commands are emitted. A false flag means the corresponding block is simply skipped.

Numeric parameters (embedding_scale, quant_bits, quant_group_size): These are values that get baked into kernel parameters.

Weight key strings (post_attn_norm_key, post_ffn_norm_key): These are the suffixes used to look up weight tensors. Gemma calls its post-attention norm weights layers.N.post_attention_norm.weight, while other architectures do not have them at all.

Factory Functions

Each supported architecture has a factory function that returns a fully initialized descriptor:

arch_llama() – The Default

inline ArchDescriptor arch_llama() {
    ArchDescriptor d = {};
    d.activation_kernel = "silu_gate_f16";
    d.embedding_scale = 0.0f;
    d.has_qk_norm = false;
    d.has_post_attn_norm = false;
    d.has_post_ffn_norm = false;
    d.rope_kernel = "rope_qkv_write_f16";  // standard (interleaved)
    d.rope_standalone = "rope_f16";
    d.tie_embeddings = false;
    d.is_encoder_decoder = false;
    d.has_bias = false;
    d.norm_type = "rmsnorm";
    return d;
}

Llama is the simplest architecture and serves as the default. SiLU-gated activation, standard RoPE, no special norms, no encoder. Most LLaMA-family models (Llama 2, Llama 3, Mistral, etc.) use this descriptor.¹

arch_qwen3() – Incremental Derivation

inline ArchDescriptor arch_qwen3() {
    ArchDescriptor d = arch_llama();  // start from llama defaults
    d.has_qk_norm = true;
    d.rope_kernel = "rope_neox_qkv_write_f16";  // NeoX (split-half)
    d.rope_standalone = "rope_neox_f16";
    d.tie_embeddings = true;
    return d;
}

Qwen3 starts from Llama and overrides three things: adds QK head norms, switches to NeoX-style RoPE, and ties the embedding weights. This derivation pattern makes the differences explicit and minimizes redundancy.

arch_gemma() – The Complex One

inline ArchDescriptor arch_gemma(int dim) {
    ArchDescriptor d = arch_llama();  // start from llama defaults
    d.activation_kernel = "gelu_gate_f16";
    d.embedding_scale = sqrtf((float)dim);
    d.has_qk_norm = true;
    d.has_post_attn_norm = true;
    d.has_post_ffn_norm = true;
    d.post_attn_norm_key = "post_attention_norm";
    d.post_ffn_norm_key = "post_ffw_norm";
    d.rope_kernel = "rope_neox_qkv_write_f16";
    d.rope_standalone = "rope_neox_f16";
    d.tie_embeddings = true;
    return d;
}

Gemma is the most feature-rich decoder-only architecture. Notice that embedding_scale is computed from the model dimension – this is why the factory takes dim as a parameter. The value sqrt(dim) is specific to Gemma’s architectural design.²

arch_whisper() – Encoder-Decoder

inline ArchDescriptor arch_whisper() {
    ArchDescriptor d = {};
    d.activation_kernel = "gelu_f16";  // plain GELU (no gate)
    d.rope_kernel = nullptr;  // no RoPE
    d.rope_standalone = nullptr;
    d.tie_embeddings = true;
    d.is_encoder_decoder = true;
    d.has_cross_attention = true;
    d.has_conv_frontend = true;
    d.has_bias = true;
    d.norm_type = "layernorm";
    d.encoder_activation = "gelu_f16";
    return d;
}

Whisper does NOT derive from Llama – it starts from a zeroed struct because almost everything is different. No RoPE, no gated activation, LayerNorm instead of RMSNorm, bias on every linear layer, and an encoder-decoder architecture with cross-attention and a convolutional audio frontend.³

arch_bert() – Embedding Model

inline ArchDescriptor arch_bert() {
    ArchDescriptor d = {};
    d.activation_kernel = "silu_gate_f16";  // SwiGLU (nomic-bert, modernBERT)
    d.rope_kernel = "rope_neox_qkv_write_f16";
    d.rope_standalone = "rope_neox_f16";
    d.norm_type = "rmsnorm";
    d.is_embedding_model = true;
    return d;
}

BERT-style models (specifically modern variants like nomic-bert) use a LLaMA-like architecture with bidirectional attention. The is_embedding_model flag tells the inference engine to skip autoregressive decoding and instead return the hidden states for mean-pooling.

The Lookup Function

inline ArchDescriptor arch_from_config(const char *arch_name, int dim) {
    if (strstr(arch_name, "whisper")) return arch_whisper();
    if (strstr(arch_name, "bert"))    return arch_bert();
    if (strstr(arch_name, "qwen"))    return arch_qwen3();
    if (strstr(arch_name, "gemma"))   return arch_gemma(dim);
    return arch_llama();  // Default: LLaMA-like
}

This function reads the architecture field from the GGUF metadata (a string like “llama”, “qwen3”, “gemma2”, “whisper”) and returns the appropriate descriptor. The strstr matching is intentionally loose – “gemma2” and “gemma3” both match “gemma”, “qwen2” and “qwen3” both match “qwen”. This works because the architectural differences between Gemma 2 and Gemma 3, or between Qwen 2 and Qwen 3, are captured in the model’s weight metadata (different head counts, etc.), not in the descriptor.

The default fallback is arch_llama(), which handles the vast majority of GGUF models in the wild (Llama, Mistral, CodeLlama, etc.).

How the Table Builder Uses Descriptors

Let’s trace through the table builder to see how the descriptor eliminates branching.

Activation Kernel

// In the SiLU/GELU dispatch:
Pipeline act_pso = device.get_pipeline(arch.activation_kernel);

The table builder does not know or care whether the activation is SiLU or GELU. It dispatches whatever arch.activation_kernel says. This single line replaces:

// Without descriptors (do NOT do this):
Pipeline act_pso;
if (cfg.architecture == "llama" || cfg.architecture == "qwen3")
    act_pso = device.get_pipeline("silu_gate_f16");
else if (cfg.architecture == "gemma")
    act_pso = device.get_pipeline("gelu_gate_f16");
else if (cfg.architecture == "whisper")
    act_pso = device.get_pipeline("gelu_f16");

RoPE Selection

// In the RoPE dispatch:
if (arch.rope_kernel) {
    Pipeline rope_pso = device.get_pipeline(arch.rope_kernel);
    // ... emit RoPE command ...
}
// If rope_kernel is nullptr (Whisper), the block is skipped entirely

Post-Norms

if (arch.has_post_attn_norm) {
    snprintf(name, sizeof(name), "layers.%d.%s.weight",
             layer, arch.post_attn_norm_key);
    Buffer post_norm_w = weights.get_tensor(name);
    emit_standalone_norm(table, device, chip, scratch.residual,
                         post_norm_w, scratch.post_norm, dim, cfg.norm_eps);
}

The boolean flag controls whether the block exists; the key string controls which weight tensor is loaded. For Llama and Qwen3, has_post_attn_norm is false and this code is never reached.

Embedding Scaling

if (arch.embedding_scale > 0.0f) {
    float scale = arch.embedding_scale;
    CmdBuilder(table, device.get_pipeline("temperature_scale_f16"), ...)
        .buf(scratch.h0, 0)
        .params(scale, 1)
        .emit();
}

Only Gemma has embedding_scale > 0, so only Gemma gets this dispatch command.

Tied Embeddings

const char *logit_name = arch.tie_embeddings
    ? "token_embedding.weight"
    : "output.weight";
Buffer logit_w = weights.get_tensor(logit_name);

A single ternary replaces a multi-way branch.

Fused QK-Norm + RoPE

The table builder checks whether the fused kernel is applicable:

bool use_fused_norm_rope = arch.has_qk_norm &&
    strcmp(arch.rope_kernel, "rope_neox_qkv_write_f16") == 0;

This is true for Qwen3 and Gemma (NeoX RoPE + QK norms), false for everything else. When true, a single fused kernel replaces 3 separate dispatches.

The DTypeDescriptor: Kernel Registry

While ArchDescriptor captures architecture-level variation, DTypeDescriptor captures dtype-level variation. It is defined in src/core/dtype_descriptor.h:

struct DTypeDescriptor {
    uint32_t dtype;
    const char *gemv_kernel;
    const char *gemv_large_kernel;
    const char *gemv_wide_kernel;
    const char *gemm_kernel;
    const char *gemm_small_kernel;
    const char *embedding_kernel;
    const char *fused_silu_kernel;
    const char *fused_silu_large_kernel;
    int gemv_rows_per_tg;
    int gemv_tg_size;
    int large_rows_per_tg;
    int large_tg_size;
    int wide_rows_per_tg;
    int wide_tg_size;
    int fused_silu_rows;
    int fused_silu_tg;
    int fused_silu_large_rows;
    int fused_silu_large_tg;
    bool is_mlx;
};

This is a flat POD struct with no pointers to chase. Each supported GGUF dtype code gets an entry in a static lookup table:

static const DTypeDescriptor kDTypes[] = {
    // dtype  gemv           gemv_large     gemv_wide         gemm              gemm_small           embed                  fused_silu          fused_silu_large  rows tg  lg_r lg_tg w_r w_tg fs_r fs_tg fsl_r fsl_tg mlx
    {0,  "gemv_f16",         nullptr,       "gemv_wide_f16",  "simd_gemm_f16",  "simd_gemm_small_f16", nullptr,             nullptr,            nullptr,           16, 128, 0,   0,   64, 256, 0,   0,    0,    0,    false},
    {2,  "gemv_q4_0",        nullptr,       "gemv_wide_q4_0", "simd_gemm_q4_0", "simd_gemm_small_q4_0","embedding_lookup_q4_0","gemv_q4_0_silu", nullptr,          16, 128, 0,   0,   64, 256, 16,  128,  0,    0,    false},
    {8,  "gemv_q8_0",        nullptr,       "gemv_wide_q8_0", "simd_gemm_q8_0", "simd_gemm_small_q8_0","embedding_lookup_q8_0", nullptr,         nullptr,          32, 256, 0,   0,   64, 256, 0,   0,    0,    0,    false},
    {12, "gemv_q4_k",        nullptr,       "gemv_wide_q4_k", "simd_gemm_q4_k", "simd_gemm_small_q4_k","embedding_lookup_q4_k", nullptr,         nullptr,          16, 256, 0,   0,   32, 256, 0,   0,    0,    0,    false},
    {100,"gemv_mlx_q4",      nullptr,       "gemv_wide_mlx_q4","simd_gemm_mlx_q4","simd_gemm_small_mlx_q4","embedding_lookup_mlx_q4","gemv_mlx_q4_silu",nullptr,   16, 128, 0,   0,   32, 256, 16,  128,  0,    0,    true},
    // ... more entries ...
};

The table currently has 17 entries covering F16, BF16, Q4_0, Q4_1, Q5_0, Q8_0, Q2_K through Q6_K, and MLX Q3/Q4/Q6/Q8 formats.

How the Table Builder Uses DTypeDescriptor

The emit_gemv helper in table_builder.cpp reads from the descriptor to select the right kernel and dispatch geometry:

static void emit_gemv(DispatchTable& table, Device& device,
                      const ChipConfig& chip,
                      Buffer input, Buffer weight, Buffer output,
                      int output_offset, uint32_t dtype, int N, int K, ...) {
    const auto& dt = device.dtype_lookup(dtype);

    // Select kernel variant
    bool use_wide = (N > chip.wide_gemv_threshold && dt.gemv_wide_kernel != nullptr);
    bool use_large = (!use_wide && dt.gemv_large_kernel != nullptr
                      && K >= chip.q4_small_k_threshold);

    const char *kernel_name;
    int rows, tg;
    if (use_wide) {
        kernel_name = dt.gemv_wide_kernel;
        rows = dt.wide_rows_per_tg;
        tg = dt.wide_tg_size;
    } else if (use_large) {
        kernel_name = dt.gemv_large_kernel;
        rows = dt.large_rows_per_tg;
        tg = dt.large_tg_size;
    } else {
        kernel_name = dt.gemv_kernel;
        rows = dt.gemv_rows_per_tg;
        tg = dt.gemv_tg_size;
    }

    Pipeline pso = device.get_pipeline(kernel_name);
    int n_groups = (N + rows - 1) / rows;
    // ... create and emit DispatchCmd ...
}

The table builder never mentions “Q4_0” or “Q8_0” by name. It looks up the descriptor for the weight’s dtype code and uses whatever kernel and geometry the table specifies. Supporting a new quantization format means adding one row to kDTypes[] and writing the corresponding Metal kernel.

The is_mlx Flag

MLX quantized formats require special handling: function constants for group_size and K, a different parameter struct layout (MLXParams vs GEMMParams), and different weight byte calculations. The is_mlx boolean flag in DTypeDescriptor controls this:

if (dt.is_mlx) {
    uint32_t fc_indices[] = {0, 1};
    uint32_t fc_values[] = {(uint32_t)quant_group_size, (uint32_t)K};
    pso = device.get_pipeline(kernel_name, cache_key, fc_indices, fc_values, 2);
} else {
    pso = device.get_pipeline(kernel_name);
}

This is the one architectural branch in emit_gemv that is not purely data-driven. It could be eliminated by adding a “pipeline creation strategy” function pointer to the descriptor, but the additional complexity is not justified for two variants.

Design Principles

The ArchDescriptor and DTypeDescriptor embody several design principles:

1. Data Over Code

Instead of encoding architecture-specific behavior in if/else chains, encode it in data structures. The table builder reads the data; it does not interpret architecture names.

2. Flat Over Deep

Both descriptors are flat POD structs. No inheritance hierarchies, no virtual methods, no builder patterns. This makes them trivially copyable, inspectable in a debugger, and cacheline-friendly.

3. Explicit Over Implicit

Every architecture-specific behavior is a named field. When you read arch.has_qk_norm, you know exactly what it controls. There are no hidden side effects, no overridden methods with subtle behavior differences.

4. Default-Safe

The zero-initialization of ArchDescriptor (ArchDescriptor d = {}) produces a safe state: no activation (nullptr), no RoPE (nullptr), no norms (false), not encoder-decoder (false). Every factory function explicitly sets the fields it needs. This prevents “forgot to set X” bugs.

5. Derivation Without Inheritance

arch_qwen3() derives from arch_llama() by copying the struct and overriding fields. This is simpler and more explicit than C++ inheritance. You can see all the differences at a glance.

Adding a New Architecture

To add support for a hypothetical “Falcon” architecture:

1. Write a factory function:

inline ArchDescriptor arch_falcon() {
    ArchDescriptor d = arch_llama();  // similar to Llama
    d.has_bias = true;               // Falcon has bias
    d.rope_kernel = "rope_neox_qkv_write_f16";  // NeoX RoPE
    d.rope_standalone = "rope_neox_f16";
    return d;
}

2. Add a case to the lookup:

inline ArchDescriptor arch_from_config(const char *arch_name, int dim) {
    if (strstr(arch_name, "falcon")) return arch_falcon();
    // ... existing cases ...
}

3. That’s it. No changes to table_builder.cpp, prefill.cpp, or any other core file (assuming the existing kernel set handles the architecture’s needs).

Summary

The ArchDescriptor pattern replaces architecture-specific branching with data-driven dispatch. A flat POD struct captures all variation points (activation, RoPE, norms, embeddings, encoder/decoder). Factory functions create descriptors for each architecture, with incremental derivation from the Llama default. The DTypeDescriptor does the same for quantization formats, mapping dtype codes to kernel names and dispatch geometry. Together, these two descriptors make the table builder and prefill engine completely agnostic to both architecture and quantization format.

1. Touvron, H., et al. (2023). “LLaMA: Open and Efficient Foundation Language Models.” arXiv:2302.13971. LLaMA’s architecture (RMSNorm, SwiGLU, RoPE, no bias) has become the de facto standard for open-weight LLMs. See https://arxiv.org/abs/2302.13971. ↩

2. Google DeepMind. (2024). “Gemma: Open Models Based on Gemini Research and Technology.” arXiv:2403.08295. Gemma scales embedding outputs by sqrt(dim), which is uncommon in other LLM architectures but follows the original Transformer convention. See https://arxiv.org/abs/2403.08295. ↩

3. Radford, A., et al. (2023). “Robust Speech Recognition via Large-Scale Weak Supervision.” Proceedings of ICML 2023. Whisper’s architecture uses sinusoidal positional encoding, LayerNorm, and cross-attention, following the original Transformer encoder-decoder design. See https://arxiv.org/abs/2212.04356. ↩

Chip Configuration and Hardware Tuning

Akunu runs on everything from an M1 MacBook Air with 8 GPU cores and 68 GB/s bandwidth to an M4 Ultra Mac Studio with 80 GPU cores and over a terabyte per second of bandwidth. The same inference engine, the same kernels, the same dispatch table structure – but with very different optimal tuning parameters. A threadgroup size that saturates an M1 might leave an M4 Max starved. A chain decode chunk size that is efficient on M4 might cause command buffer overhead to dominate on M1.

The ChipConfig struct solves this: it derives all hardware-specific tuning parameters from two inputs – the GPU core count and the GPU family number. Every magic number in the engine lives here, not scattered across the codebase.

The ChipConfig Struct

Defined in src/core/chip_config.h:

struct ChipConfig {
    // === Hardware info ===
    int gpu_cores;           // GPU core count (from IOKit)
    int gpu_family;          // Apple GPU family (7=M1, 8=M2/M3, 9=M4)
    int max_threads_per_tg;  // Metal max (typically 1024)

    // === SLC (System Level Cache) ===
    int slc_bytes;             // estimated SLC size
    bool should_fuse_weights;  // fuse gate+up, QKV when SLC is large enough

    // === GEMV tuning ===
    int q4_small_k_threshold;  // Q4_0: use small-K variant below this
    int wide_gemv_threshold;   // switch to wide GEMV when N exceeds this
    bool gemv_wide_standard;   // use 256-thread standard GEMV on Pro+

    // === Chain decode ===
    int chain_decode_chunk;  // tokens per GPU submission

    // === Prefill ===
    int max_prefill_chunk;  // max tokens per prefill GEMM batch

    // === Norm dispatch ===
    int norm_tg_size;  // threadgroup size for RMSNorm
};

Every field has a clear purpose and measurable impact on performance. Let’s go through each one.

Chip Detection

The factory function ChipConfig::from_gpu(int cores, int family) builds a config from two hardware-detected values:

• cores: The GPU core count, queried from IOKit’s AGXAccelerator service via IORegistryEntryCreateCFProperty(CFSTR("gpu-core-count")). This is the most reliable indicator of chip tier.
• family: The Apple GPU family number, queried via Metal’s [device supportsFamily:]. Family 7 = M1, Family 8 = M2/M3, Family 9 = M4.

These two values together uniquely identify the chip tier and generation:

SLC Size Estimation

The System Level Cache (SLC) is the large last-level cache shared between CPU, GPU, and other IP blocks. Apple does not officially disclose SLC sizes, but they can be estimated from die analysis and performance profiling.¹

if (cores >= 60)
    c.slc_bytes = 96 * 1024 * 1024;  // Ultra
else if (cores >= 30)
    c.slc_bytes = 48 * 1024 * 1024;  // Max
else if (cores >= 16)
    c.slc_bytes = (family >= 9) ? 32 * 1024 * 1024 : 24 * 1024 * 1024;  // Pro
else
    c.slc_bytes = (family >= 9) ? 16 * 1024 * 1024 : 8 * 1024 * 1024;  // Base

Key observations:

• M4 (family 9) gets larger SLC estimates at every tier. Die analysis suggests M4 improved the cache hierarchy significantly.²
• Ultra chips (2x Max die-to-die) get double the Max SLC.

Impact on Weight Fusion

The SLC size directly controls whether QKV and gate+up weight fusion is beneficial:

c.should_fuse_weights = (c.slc_bytes >= 16 * 1024 * 1024);  // Pro+ and M4 Base

Weight fusion concatenates two or three weight matrices into one, replacing multiple GEMV dispatches with a single larger one. This reduces kernel launch overhead and can improve cache utilization – but only if the fused weight matrix fits (or mostly fits) in the SLC. On an M1 base with 8 MB SLC, a fused QKV weight matrix for a 4096-dim model might be:

$$Q + K + V = (4096 \times 4096 + 2 \times 4096 \times 1024) \times 0.5 \text{ bytes/elem} = 12.3 \text{ MB (Q4_0)}$$

This exceeds the 8 MB SLC, causing cache thrashing. On an M4 Pro with 32 MB SLC, the fused matrix fits comfortably, and the fusion saves two kernel launches.

GEMV Kernel Selection

Three parameters control GEMV variant selection:

q4_small_k_threshold

if (family >= 9 && cores >= 16)
    c.q4_small_k_threshold = 512;   // M4 Pro+
else if (cores >= 16)
    c.q4_small_k_threshold = 1024;  // M1-M3 Pro+
else
    c.q4_small_k_threshold = 2048;  // Base chips

This determines when to use a “large-K” GEMV variant that uses more SIMD groups for the K-dimension reduction. Larger chips can afford more parallelism at smaller K values because they have more GPU cores to keep busy.

wide_gemv_threshold

c.wide_gemv_threshold = 32768;

When the output dimension N exceeds this threshold (e.g., vocabulary projection with 128K tokens), switch to a “wide” GEMV kernel with more output rows per threadgroup. This is constant across all chips because the decision depends on the N dimension, not the hardware.

gemv_wide_standard

c.gemv_wide_standard = (cores >= 16);

On Pro+ chips (16+ cores), use 256-thread threadgroups for the standard GEMV instead of 128-thread. More threads per threadgroup means more SIMD groups, which means better occupancy on chips with many cores. On base chips with 8-10 cores, the extra threads would just increase register pressure without improving occupancy.

Chain Decode Chunk Size

if (cores >= 60)
    c.chain_decode_chunk = 128;      // Ultra
else if (cores >= 30)
    c.chain_decode_chunk = 128;      // Max
else if (family >= 9)
    c.chain_decode_chunk = 128;      // M4 Pro/Base
else if (cores >= 16)
    c.chain_decode_chunk = 96;       // M3 Pro
else
    c.chain_decode_chunk = 64;       // M1/M2 Base

This is the number of tokens batched into a single GPU command buffer during chain decode. The trade-offs:

• Larger chunks: Amortize command buffer overhead across more tokens. Better GPU utilization. But each submission takes longer, increasing latency to first token in a batch.
• Smaller chunks: Lower per-submission latency. But higher overhead ratio (command buffer setup / useful work).

M4’s improved GPU command processor can handle larger chunks efficiently. Older chips with smaller caches and slower command processing benefit from smaller chunks.

For a 32-layer model with ~200 commands per token:

• Chunk 128: 25,600 dispatches per command buffer
• Chunk 64: 12,800 dispatches per command buffer

The total generation throughput is virtually identical (both are bandwidth-limited), but the chunk size affects how quickly the first token in a chunk appears.

Prefill Chunk Size

c.max_prefill_chunk = 4096;

Maximum tokens per prefill GEMM batch. This is constant because prefill is compute-bound (large M), and the GEMM kernel performance is relatively insensitive to batch size above a few hundred tokens. The limit prevents excessive memory usage for the batch scratch buffers.

Norm Threadgroup Size

c.norm_tg_size = 1024;

RMSNorm and LayerNorm use a single threadgroup to process the entire dimension vector. The threadgroup size is min(dim, norm_tg_size). For a 4096-dim model, this means 1024 threads (the Metal maximum), with each thread processing 4 elements.

This is constant because normalization kernels are lightweight and the threadgroup size has minimal impact on performance.

How ChipConfig Flows Through the System

The configuration flows from hardware detection to kernel dispatch:

┌──────────────┐      ┌────────────────┐      ┌──────────────────┐
│  MetalDevice  │─────>│  ChipConfig    │─────>│  table_builder   │
│  gpu_cores()  │      │  from_gpu()    │      │  emit_gemv()     │
│  gpu_family() │      │                │      │  build_dispatch  │
└──────────────┘      └────────────────┘      └──────────────────┘
                                                      │
                                                      ▼
                                              ┌──────────────────┐
                                              │  DispatchTable   │
                                              │  (commands[]     │
                                              │   with resolved  │
                                              │   kernel+geom)   │
                                              └──────────────────┘

1. MetalDevice::chip_config() calls ChipConfig::from_gpu(gpu_core_count(), gpu_family())
2. The config is passed to build_dispatch_table() and prefill()
3. The table builder uses chip.should_fuse_weights to decide QKV/gate+up fusion
4. emit_gemv() uses chip.wide_gemv_threshold, chip.q4_small_k_threshold, and chip.gemv_wide_standard to select kernel variants
5. The chain decoder uses chip.chain_decode_chunk for batch sizing
6. The resulting dispatch table has hardware-optimal kernel choices baked in

All of this happens once at model load time. At inference time, the dispatch table is simply replayed – no hardware queries, no branching, no configuration lookups.

Real-World Impact

To make the impact concrete, here are approximate decode speeds for a 7B Q4_0 model on different chips, showing how ChipConfig-driven tuning affects performance:

The “Fuse?” column shows whether weight fusion is enabled (SLC >= 16 MB). The “Wide GEMV?” column shows whether 256-thread GEMV is used (cores >= 16). These configuration-driven choices account for a significant portion of the performance difference between “works” and “fast.”

Design Philosophy

ChipConfig embodies a few key principles:

1. All magic numbers in one place. Instead of hardcoded constants in table_builder.cpp, prefill.cpp, and decode_loop.cpp, every tuning parameter is in ChipConfig. This makes it easy to review, adjust, and reason about the tuning.

2. Derived, not configured. Users do not set these values. They are automatically derived from hardware detection. This eliminates “you need to tune parameter X for your specific hardware” friction.

3. Conservative defaults. The base chip configuration (64 tokens/chunk, no fusion, 128-thread GEMV) is conservative. It works correctly on every chip. The Pro/Max/Ultra configurations are more aggressive but only activate when the hardware supports them.

4. Chip tier matters more than generation. The core count (which determines chip tier: base/Pro/Max/Ultra) has more impact on tuning than the GPU family (which determines generation: M1/M2/M3/M4). An M1 Max (32 cores) gets more aggressive tuning than an M4 base (10 cores), even though the M4 is a newer generation.

Deep Dive: The SLC and Weight Fusion Decision

The should_fuse_weights decision deserves a more thorough analysis because it is one of the most impactful configuration choices. Let’s trace through the numbers for a concrete model.

Consider Llama 3 8B with Q4_0 quantization:

• dim = 4096, n_heads = 32, n_kv_heads = 8, head_dim = 128
• q_dim = 4096, kv_dim = 1024, ffn_dim = 14336

QKV fusion candidate:

On an M1 base (8 MB SLC), the fused QKV matrix (14.2 MB) far exceeds the cache. Reading it causes cache thrashing – the end of the matrix evicts the beginning before the next token needs it. Three separate GEMVs at 9.4 + 2.4 + 2.4 = 14.2 MB read the same total bytes but each individual matrix is smaller. The K and V matrices (2.4 MB each) may actually fit in the SLC between dispatches, providing a small cache benefit.

On an M4 Pro (32 MB SLC), the fused QKV matrix (14.2 MB) fits with room to spare. The single GEMV dispatch saves two kernel launches (~0.01ms each) and potentially benefits from better memory controller scheduling (one large sequential read vs. three separate reads with dispatch overhead between them).

Gate+Up fusion candidate:

The fused gate+up matrix (66 MB) exceeds the SLC on every Apple Silicon chip. So why fuse at all? Because the fusion still saves one kernel launch and one dispatch table entry per layer, and the memory controller can prefetch more efficiently when reading one contiguous allocation. On Pro+ chips where the SLC is large enough to hold a meaningful portion of the matrix, the cache hit rate on the “tail” of the read (after the GPU has processed the head) is better than with two separate reads.

On base chips, the gate+up fusion is skipped entirely (controlled by should_fuse_weights = false) because the kernel launch saving does not compensate for the worse cache behavior.

Deep Dive: Chain Decode Chunk Size

The chunk size determines how many tokens are packed into a single Metal command buffer. Let’s analyze the trade-offs quantitatively.

Command buffer overhead: Creating a command buffer, encoding commands, and submitting takes a fixed amount of CPU time. On Apple Silicon, this is approximately:

• Command buffer creation: ~2 us
• Encoder creation: ~1 us
• Submission (commit): ~5 us
• GPU scheduling delay: ~10-50 us (varies by chip and load)

Total fixed overhead per submission: ~20-60 us.

Per-token encoding time: Each token requires ~260 Metal API calls (set pipeline, set buffer, set bytes, dispatch). Each API call takes approximately 50-100ns. Total per-token encoding: ~13-26 us.

For different chunk sizes on M4 Pro:

The “Tokens Until Display” column shows the latency cost: with a chunk size of 128, the user must wait for all 128 tokens to be generated before any are displayed. For interactive chat at 70 tok/s, a chunk of 128 takes ~1.8 seconds – potentially noticeable.

Akunu balances this by using smaller chunks for the first few tokens (to minimize time-to-first-token) and larger chunks for sustained generation. The chain_decode_chunk in ChipConfig sets the maximum chunk size for sustained generation.

On older chips (M1 base), the per-token encoding time is higher (slower CPU) and the GPU execution time is longer (less bandwidth), so the fixed overhead is relatively smaller. A chunk of 64 achieves a good balance. On M4, the faster CPU makes encoding cheaper and the faster GPU makes execution shorter, so a chunk of 128 is optimal.

Deep Dive: GEMV Variant Selection Logic

The three GEMV-related parameters work together to form a decision tree. Let me trace through how emit_gemv in table_builder.cpp uses them.

For each GEMV dispatch, the kernel variant is selected as follows:

Input: N (output rows), K (reduction dimension), dtype

1. Look up DTypeDescriptor for dtype
2. If N > chip.wide_gemv_threshold AND wide kernel exists:
     → Use wide GEMV (more output rows per TG)
3. Else if large-K kernel exists AND K >= chip.q4_small_k_threshold:
     → Use large-K GEMV (more threads for K-reduction)
4. Else:
     → Use standard GEMV

For a concrete example with Llama 3 8B (Q4_0):

The wide GEMV is only used for the vocabulary projection, where N (vocab size) is very large. All other projections use the standard kernel. The q4_small_k_threshold would matter for models with smaller hidden dimensions or for the KV projections, where K might be below the threshold on base chips.

The Backwards Compatibility Shim

The struct includes a backwards-compatibility factory:

static ChipConfig from_gpu_cores(int cores) { return from_gpu(cores, 0); }

This allows older code that only knows the core count (not the GPU family) to still create a configuration. With family = 0, the M4-specific optimizations (larger SLC, lower K-threshold) are skipped, and the config falls back to conservative M1-era tuning.

Extending ChipConfig

If you were adding CUDA support, ChipConfig would need additional fields:

// Hypothetical CUDA additions:
int sm_count;               // number of streaming multiprocessors
int shared_mem_per_sm;      // shared memory (threadgroup memory equivalent)
int warp_size;              // always 32, but good to be explicit
bool has_tensor_cores;
int tensor_core_gen;        // 0=none, 3=Volta, 4=Ampere, 5=Hopper
int l2_cache_bytes;         // L2 cache size (analogous to SLC)

The factory function would detect these from cuDeviceGetAttribute and derive tuning parameters. The should_fuse_weights decision would key off l2_cache_bytes instead of slc_bytes. The chain_decode_chunk would adapt to the different command submission model (CUDA streams vs Metal command buffers). The table builder would use the same ChipConfig interface, just with different values.

For Vulkan, additional fields might include:

int subgroup_size;          // analogous to SIMD width (varies: 16, 32, 64)
bool has_subgroup_matrix;   // cooperative matrix support
int max_workgroup_size;     // varies by device

The beauty of the ChipConfig approach is that the table builder does not care which backend produced the values. It just reads should_fuse_weights, wide_gemv_threshold, etc. The backend-specific detection logic is encapsulated in the factory function.³

Summary

ChipConfig is a simple struct that maps hardware capabilities (GPU core count, GPU family) to optimal tuning parameters (SLC size, fusion decisions, GEMV variant selection, chain decode chunk size). It is built once at device creation and consumed at dispatch table build time. All decisions are baked into the dispatch table – nothing is queried or branched on at inference time.

1. Chips and Cheese. (2021-2024). “Apple M1/M2/M3/M4 Die Analysis.” SLC sizes are estimated from die area analysis, performance counter measurements, and benchmark profiling. Apple does not officially disclose SLC sizes. See https://chipsandcheese.com/. ↩

2. Apple, “Apple M4 chip” (2024). The M4 chip announcement highlights “next-generation GPU” with improved ray tracing and machine learning performance. Independent analysis suggests significant cache hierarchy improvements. See https://www.apple.com/newsroom/2024/05/apple-introduces-m4-chip/. ↩

3. This separation of concerns – hardware detection in the backend, tuning consumption in the core – is the same pattern used by the Device abstraction layer. ChipConfig is effectively a “capabilities descriptor” that the core engine reads, analogous to how ArchDescriptor is a “model capabilities descriptor.” ↩

Model Loading: From File to GPU

If you have ever wondered what actually happens between the moment you point an inference engine at a model file and the moment the first token appears on your screen, this chapter is for you. In akunu, that transition is orchestrated by a single function – akunu_load_model – defined in src/inference/model_loader.cpp. The function is roughly 250 lines of C++, and it touches almost every subsystem in the engine: format detection, weight I/O, configuration parsing, architecture selection, tokenizer construction, GPU memory allocation, weight fusion, KV cache creation, scratch buffer setup, RoPE precomputation, dispatch table building, shader compilation, and a warmup forward pass. That is a lot of ground to cover, so let us take it one stage at a time.

The Bird’s-Eye View

Here is the full pipeline from “user calls the C API” to “model is ready to generate text.”

  akunu_load_model(path, metallib_path, max_context)
  |
  |  1. Create GPU device
  |     Device::create_default()
  |
  |  2. Load metallib (shader library)
  |     load_metallib(device, metallib_path)
  |
  |  3. Detect model format (GGUF vs MLX/SafeTensors vs Whisper GGML)
  |     is_whisper_ggml(path) / WeightProvider::detect_format(path)
  |
  |  4. Open weights
  |     WeightProvider::open(path)
  |
  |  5. Extract model config
  |     weights.get_config() --> AkunuModelConfig
  |
  |  6. Infer architecture descriptor
  |     arch_from_config(config.architecture, config.dim) --> ArchDescriptor
  |
  |  7. Detect chip capabilities
  |     device.chip_config() --> ChipConfig
  |
  |  8. Resolve weight quirks (MLX RoPE style, tied embeddings, quant params)
  |
  |  9. Precompute RoPE frequencies
  |     init_rope_freqs(state, path)
  |
  | 10. Validate config (dim, layers, vocab, heads, head_dim all nonzero)
  |
  | 11. Load tokenizer
  |     load_tokenizer(state, path)
  |
  | 12. Allocate KV cache
  |     KVCache::create(device, n_layers, n_kv_heads, head_dim, ctx)
  |
  | 13. Allocate scratch buffers
  |     ScratchBuffers::create(device, config, ctx, max_prefill_chunk)
  |
  | 14. Build greedy dispatch table
  |     build_dispatch_table(device, weights, config, arch, chip, kv, scratch)
  |
  | 15. Build sampled dispatch table (greedy + Gumbel top-k + argmax)
  |
  | 16. Pre-allocate GPU param buffers for dispatch commands
  |
  | 17. Warmup forward pass (compile all PSOs)
  |     encode_dispatch_table(&table, 0, 1); end_encoding_sync();
  |
  | 18. Return opaque model handle
  v

That is 18 discrete stages. Some of them are trivial (one function call). Others – like the dispatch table build – are complex enough to deserve an entire chapter of their own (see Chapter 25). But they all execute sequentially, one time, at model load. Nothing in this list happens again during inference. That is the whole point: pay the cost once up front so the hot path is zero-allocation and zero-branching.

Stage 1: Creating the Metal Device

The very first thing akunu_load_model does is create a Metal device:

state->device = Device::create_default();
printf("GPU: %s\n", state->device->name());

Device is akunu’s hardware abstraction layer. On Apple Silicon, create_default() calls MTLCreateSystemDefaultDevice(), wraps the resulting id<MTLDevice> in an internal implementation class (MetalDeviceImpl), and queries the hardware for its GPU core count and Apple GPU family number. Those two integers feed into ChipConfig::from_gpu() later, which tunes every kernel dispatch in the engine to the specific chip you are running on.

Why does the device come first? Because almost everything else – the weight provider, the KV cache, the scratch buffers – needs a Device& reference to allocate GPU memory. Device construction is cheap (microseconds), and it anchors the entire object graph.

Stage 2: Loading the Metallib

if (!load_metallib(*state->device, metallib_path)) {
    set_error("Failed to load metallib. Build with: make shaders");
    delete state;
    return nullptr;
}

A .metallib is Apple’s precompiled shader archive – the GPU equivalent of a .a static library. Akunu’s Metal kernels (GEMV, GEMM, RoPE, attention, normalization, activation, argmax, etc.) are compiled offline into a single akunu.metallib file. load_metallib tries a user-provided path first, then falls back to a handful of well-known build output locations:

engine/build/akunu.metallib
.build/metallib/akunu.metallib
build/akunu.metallib
akunu.metallib

If none of these exist, the model load fails immediately. You cannot do inference without GPU kernels. The function is deliberately simple – no complex search logic, no environment variables, just a flat priority list.

Why load shaders before weights? Because there is no point spending time and memory on a multi-gigabyte weight file if we cannot even dispatch a kernel. Fail fast.

Stage 3: Format Detection

Akunu supports three model formats:

1. GGUF – the standard quantized format from llama.cpp. Most common.
2. MLX SafeTensors – Apple’s MLX framework format. Directory of .safetensors files plus a config.json.
3. Whisper GGML – legacy binary format for OpenAI Whisper models. Uses a "lmgg" or "ggjt" magic number.

The detection logic is refreshingly unsophisticated:

bool is_whisper_ggml(const char *path) {
    FILE *f = fopen(path, "rb");
    char magic[4];
    fread(magic, 1, 4, f);
    fclose(f);
    return memcmp(magic, "lmgg", 4) == 0
        || memcmp(magic, "ggjt", 4) == 0;
}

If the path is a directory or ends in .safetensors, it is MLX. Otherwise, GGUF. No content sniffing beyond the four-byte magic for Whisper. This is a case where “dumb and fast” beats “clever and fragile.”

Stage 4: Opening Weights

For standard LLM models (not Whisper), the next step is opening the weight provider:

state->weights = new WeightProvider(*state->device);
if (!state->weights->open(model_path)) {
    set_error("Failed to open model: %s", model_path);
    delete state;
    return nullptr;
}

WeightProvider is a unified facade that wraps either a WeightStore (GGUF parser) or an MLXWeightStore (SafeTensors parser). The open() call parses file headers and metadata but does NOT eagerly load all tensors into GPU memory. Weight tensors are memory-mapped and lazily materialized on first access through get_tensor(). For a 7B Q4_0 model, the GGUF file is about 4 GB – you do not want to copy all of that before you even know the model’s architecture.

The WeightProvider exposes a uniform interface regardless of format:

+-----------------+-----------+-----------+
| Method          | GGUF      | MLX       |
+-----------------+-----------+-----------+
| get_config()    | metadata  | JSON      |
| get_tensor(name)| mmap+GPU  | mmap+GPU  |
| get_dtype(name) | per-tensor| per-tensor|
| has_tensor(name)| lookup    | lookup    |
| fuse_weights()  | concat    | concat    |
+-----------------+-----------+-----------+

This abstraction is crucial because the rest of the engine – table builder, prefill, everything – never knows or cares what format the weights came from.

Stage 5: Extracting the Model Config

state->config = state->weights->get_config();

For GGUF, this reads well-known metadata keys (llama.embedding_length, llama.block_count, llama.attention.head_count, etc.) from the GGUF header. For MLX, it parses config.json in the model directory.

The result is a plain-old-data struct, AkunuModelConfig:

AkunuModelConfig
+----------------------------------+
| dim            (e.g., 4096)      |  embedding dimension
| n_layers       (e.g., 32)       |  transformer layers
| n_heads        (e.g., 32)       |  query heads
| n_kv_heads     (e.g., 8)        |  key/value heads (GQA)
| head_dim       (e.g., 128)      |  dim per head
| q_dim          (e.g., 4096)     |  n_heads * head_dim
| kv_dim         (e.g., 1024)     |  n_kv_heads * head_dim
| ffn_dim        (e.g., 14336)    |  feed-forward intermediate
| vocab_size     (e.g., 32000)    |  vocabulary size
| max_seq_len    (e.g., 8192)     |  maximum context length
| norm_eps       (e.g., 1e-5)     |  RMSNorm epsilon
| rope_theta     (e.g., 500000.0) |  RoPE base frequency
| architecture   "llama"          |  arch name string
+----------------------------------+

Every downstream allocation and dispatch geometry is driven entirely by these numbers. There are no “if model is 7B then…” branches anywhere in akunu. The config is the single source of truth.

Stage 6: Architecture Inference

state->arch = arch_from_config(state->config.architecture, state->config.dim);

This is one of akunu’s cleanest design patterns. The ArchDescriptor struct captures every architecture-specific behavior as data, not code:

ArchDescriptor
+------------------------------+------------------+
| Field                        | Example (LLaMA)  |
+------------------------------+------------------+
| activation_kernel            | "silu_gate_f16"  |
| embedding_scale              | 0.0 (no scale)   |
| has_qk_norm                  | false            |
| has_post_attn_norm           | false            |
| has_post_ffn_norm            | false            |
| rope_kernel                  | "rope_qkv_write" |
| tie_embeddings               | false            |
| quant_bits                   | 0 (GGUF native)  |
+------------------------------+------------------+

Different architectures get different descriptors:

arch_from_config("llama", ...)  --> arch_llama()
arch_from_config("qwen3", ...)  --> arch_qwen3()    // QK-norm, tied embeds
arch_from_config("gemma", ...)  --> arch_gemma(dim)  // GELU, post-norms, scale
arch_from_config("whisper",...) --> arch_whisper()   // LayerNorm, cross-attn
arch_from_config("bert", ...)   --> arch_bert()      // encoder-only, SwiGLU

The beauty of this approach is that adding a new architecture requires exactly one factory function and one case in arch_from_config. The table builder, the prefill engine, and the decode loop never branch on the architecture name. They read fields from the descriptor. This is the “data-driven polymorphism” pattern, and it produces code that is both simpler and faster than the traditional virtual-method approach.

Stage 7: Chip Configuration

state->chip = state->device->chip_config();

ChipConfig captures hardware-specific tuning parameters:

ChipConfig
+-----------------------------+-------+----------+----------+
| Field                       | M1    | M3 Pro   | M4 Ultra |
+-----------------------------+-------+----------+----------+
| gpu_cores                   | 8     | 18       | 64       |
| slc_bytes                   | 8 MB  | 24 MB    | 96 MB    |
| should_fuse_weights         | false | true     | true     |
| chain_decode_chunk          | 64    | 96       | 128      |
| max_prefill_chunk           | 4096  | 4096     | 4096     |
| q4_small_k_threshold        | 2048  | 1024     | 512      |
| wide_gemv_threshold         | 32768 | 32768    | 32768    |
+-----------------------------+-------+----------+----------+

Notice should_fuse_weights: on chips with a large System Level Cache (16+ MB, meaning Pro and above, plus M4 base), akunu will fuse Q/K/V weight matrices into a single contiguous buffer and dispatch one GEMV instead of three. The fused weights fit in SLC, so the second and third “GEMVs” hit cache instead of DRAM. On base M1/M2/M3, the SLC is too small for this to help, so fusion is disabled.

Also notice chain_decode_chunk: this controls how many tokens are chained into a single GPU command buffer submission. Larger chunks amortize the overhead of Metal command encoding. M4 can handle 128 tokens per chunk; M1 is limited to 64 due to its smaller command processor and narrower memory bus.

Stage 8: Weight Quirks

After extracting the architecture descriptor, model_loader patches it with format-specific adjustments:

// MLX Llama uses NeoX (split-half) RoPE, not interleaved
if (weights->format() == MLX_SAFETENSORS && strstr(config.architecture, "llama"))
    arch.rope_kernel = "rope_neox_qkv_write_f16";

// Tie embeddings if the architecture says so, OR if output.weight is missing
if (!arch.tie_embeddings)
    arch.tie_embeddings = !weights->has_tensor("output.weight");

// Copy MLX quantization info
arch.quant_bits = weights->quant_bits();
arch.quant_group_size = weights->quant_group_size();

This is where the messy real-world details of model formats get resolved. MLX uses a different RoPE convention than GGUF for the same architecture. Some GGUF exports include output.weight and some do not. Rather than scattering format checks throughout the engine, they are all concentrated here in the loader.

Stage 9: RoPE Frequency Precomputation

init_rope_freqs(state, model_path);

Most models use standard RoPE with a base frequency theta. But LLaMA 3 introduced a complex wavelen-based frequency scaling scheme, and some models use simple linear scaling. akunu handles both by precomputing the frequency divisors for each dimension of the rotary embedding at load time:

For each dimension i in [0, head_dim/2):
    base_freq = theta^(2i / head_dim)

    LLaMA 3 wavelen scaling:
        wavelen = 2 * pi * base_freq
        if wavelen > low_wavelen:
            freq[i] = base_freq * factor         (long wavelengths scaled)
        else if wavelen > high_wavelen:
            smooth = (orig_max_pos/wavelen - low_freq_factor)
                   / (high_freq_factor - low_freq_factor)
            freq[i] = base_freq / ((1-smooth)/factor + smooth)
        else:
            freq[i] = base_freq                  (short wavelengths unchanged)

    Linear scaling:
        freq[i] = base_freq * factor

The resulting frequency vector is uploaded to a GPU buffer (arch.rope_freqs). If no scaling is needed, the buffer stays null and the RoPE kernel computes frequencies on the fly from theta.

This precomputation avoids two problems. First, the wavelen formulas involve floating-point operations (pow, division, conditional branches) that would add latency if repeated on every token. Second, it keeps the GPU kernel simpler – it either reads precomputed frequencies from a buffer or computes the standard geometric series, never the complex LLaMA 3 formula.

Stage 10: Config Validation

if (config.dim == 0 || config.n_layers == 0 || config.vocab_size == 0 ||
    config.n_heads == 0 || config.head_dim == 0) {
    set_error("Invalid model config: ...");
    delete state;
    return nullptr;
}

A simple sanity check. If any critical dimension is zero, something went wrong during metadata parsing (corrupt file, unsupported format version, missing keys). Fail immediately rather than producing cryptic GPU errors downstream.

Stage 11: Tokenizer Loading

The load_tokenizer function is more involved than you might expect. It needs to handle two completely different tokenizer sources:

GGUF path:
    vocab  <-- weights.get_string_array("tokenizer.ggml.tokens")
    scores <-- weights.get_float_array("tokenizer.ggml.scores")
    merges <-- weights.get_string_array("tokenizer.ggml.merges")
    type   <-- weights.get_metadata_string("tokenizer.ggml.model")
    bos_id <-- weights.get_metadata_int("tokenizer.ggml.bos_token_id")
    eos_id <-- weights.get_metadata_int("tokenizer.ggml.eos_token_id")

HuggingFace path (MLX models):
    load_hf_tokenizer(model_dir, hf_data)
    --> parses tokenizer.json + tokenizer_config.json

After loading the raw vocabulary, akunu also scans for implicit stop tokens that are not the “official” EOS token but should still terminate generation:

const char *stop_tokens[] = {
    "<|im_end|>",     // ChatML
    "<|endoftext|>",  // GPT-2 style
    "<|eot_id|>",     // LLaMA 3
    "<end_of_turn>",  // Gemma
    "</s>",           // legacy
    nullptr
};

For each of these, if the string exists in the vocabulary with a different ID than the primary EOS, it is registered as an additional EOS token. This means the decode loop does not need to know about chat templates – it just checks tokenizer.is_eos(tok) and the tokenizer handles the multi-EOS logic.

Stage 12: KV Cache Allocation

int ctx = max_context > 0
    ? max_context
    : std::min((int)config.max_seq_len, chip.max_prefill_chunk);

state->kv_cache = KVCache::create(
    *state->device, config.n_layers, config.n_kv_heads, config.head_dim, ctx);

The KV cache is the largest single allocation in the system. For a 32-layer model with 8 KV heads, 128-dim heads, and a 4096-token context, the total is:

Per-layer buffer size:
    n_kv_heads * max_seq_len * head_dim * sizeof(FP16)
    = 8 * 4096 * 128 * 2
    = 8 MB

Total:
    n_layers * 2 (K + V) * 8 MB
    = 32 * 2 * 8 MB
    = 512 MB

The KVCache struct is a flat POD container – no virtual calls, no reference counting, no linked lists:

KVCache
+------------------+
| n_layers    = 32 |
| n_kv_heads  = 8  |
| head_dim    = 128|
| max_length  = 4096|
| current_length = 0|
| kv_stride   = 524288  (max_length * head_dim)
|                  |
| k_buffers[32]    |  <-- one GPU buffer per layer
| v_buffers[32]    |  <-- one GPU buffer per layer
+------------------+

All buffers are zero-filled at creation. The current_length field tracks how many positions have been written (by prefill or decode). The kv_stride is precomputed to avoid a multiplication in the attention kernel’s inner loop.

Stage 13: Scratch Buffer Allocation

state->scratch = ScratchBuffers::create(
    *state->device, state->config, ctx, chip.max_prefill_chunk);

Scratch buffers are the working memory for a single forward pass. They are allocated once and reused every time. No dynamic allocation ever happens in the hot path.

ScratchBuffers (decode -- single token)
+----------------------------------------+
| h0         [dim]         FP16   residual ping  |
| h1         [dim]         FP16   residual pong  |
| residual   [dim]         FP16   norm output    |
| qkv        [q+2*kv_dim] FP16   Q|K|V concat   |
| attn_out   [q_dim]       FP16   attention out  |
| post_norm  [dim]         FP16   Gemma temp     |
| ffn_gate   [2*ffn_dim]  FP16   gate|up fused   |
| ffn_up     [ffn_dim]    FP16   up projection   |
| ffn_act    [ffn_dim]    FP16   activation out  |
| logits     [vocab]       FP16   final logits   |
| token_ids  [max_chain]  U32    token buffer    |
+----------------------------------------+

ScratchBuffers (prefill -- batch)
+----------------------------------------+
| batch_h0        [chunk * dim]    FP16  |
| batch_h1        [chunk * dim]    FP16  |
| batch_residual  [chunk * dim]    FP16  |
| batch_q         [chunk * q_dim]  FP16  |
| batch_k         [chunk * kv_dim] FP16  |
| batch_v         [chunk * kv_dim] FP16  |
| batch_attn_out  [chunk * q_dim]  FP16  |
| batch_gate      [chunk * ffn]    FP16  |
| batch_up        [chunk * ffn]    FP16  |
| batch_act       [chunk * ffn]    FP16  |
| batch_post_norm [chunk * dim]    FP16  |
+----------------------------------------+

The dual sets of buffers – one for single-token decode, one for batched prefill – are a deliberate design choice. Decode uses GEMV (matrix-vector); prefill uses GEMM (matrix-matrix). They have completely different memory access patterns, and sharing buffers between them would either waste memory or require dynamic resizing.

Stage 14: Building the Greedy Dispatch Table

state->dispatch_table = build_dispatch_table(
    *state->device, *state->weights, state->config,
    state->arch, state->chip, state->kv_cache, state->scratch);

This is the most complex and most important step of model loading. The dispatch table is a flat array of DispatchCmd structs – one for every GPU kernel dispatch needed to process a single token through the entire transformer. Chapter 25 covers this in detail, but the high-level structure is:

Command sequence for one token:

  [0]   Embedding lookup
  [1]   Embedding scale (Gemma only)
  [2]   Initial RMSNorm

  For each layer:
    [3+N*k]   Fused QKV GEMV (or separate Q, K, V)
    [...]     Fused QK-norm + RoPE + KV write (or separate)
    [...]     Flash attention decode
    [...]     O projection GEMV
    [...]     Post-attn norm (Gemma only)
    [...]     Fused residual + FFN norm
    [...]     Fused gate+up GEMV (or separate)
    [...]     Fused SiLU+down GEMV (or separate activation + down)
    [...]     Post-FFN norm (Gemma only)
    [...]     Fused next attn norm

  [N-2] Output norm
  [N-1] Logit projection GEMV
  [N]   Argmax

For a 32-layer LLaMA model with all fusions enabled, this is typically around 160-200 commands.

Stage 15: Building the Sampled Dispatch Table

After the greedy table is built, model_loader constructs a second table for sampled (temperature > 0) decoding:

auto& greedy_cmds = state->dispatch_table.commands;
auto& sampled_cmds = state->sampled_dispatch_table.commands;

// Copy all commands except the final argmax
for (size_t i = 0; i + 1 < greedy_cmds.size(); i++)
    sampled_cmds.push_back(greedy_cmds[i]);

// Insert a Gumbel top-k noise kernel
DispatchCmd gumbel_cmd = DispatchCmd::make(gumbel_pso, Dim3(1), Dim3(1024));
gumbel_cmd.add_buffer(scratch.logits, 0, 0);
// ... set up GumbelTopKParams ...
gumbel_cmd.patch_type = DispatchCmd::PATCH_POSITION;

sampled_cmds.push_back(gumbel_cmd);
sampled_cmds.push_back(greedy_cmds.back());  // argmax still works!

The key insight: the Gumbel-max trick turns sampling into argmax. By adding temperature * Gumbel_noise to each logit before taking the argmax, you get a sample from the categorical distribution softmax(logits / temperature). This means the sampled table is identical to the greedy table except for one extra kernel inserted before the argmax. No CPU round-trip. No softmax. No probability computation. Just noise + argmax, entirely on the GPU.

Stage 16: Pre-Allocating Parameter Buffers

for (auto& cmd : cmds) {
    if (cmd.param_size > 0) {
        cmd.param_buf = device.allocate(cmd.param_bytes, cmd.param_size);
    }
}

Each DispatchCmd can carry up to 64 bytes of inline parameter data. Rather than using Metal’s setBytes (which copies the data on every dispatch), akunu pre-allocates a small GPU buffer for each command’s parameters and uses setBuffer instead. On Unified Memory Architecture (UMA) systems like Apple Silicon, these buffers are CPU-accessible, so updating a parameter (like the position for RoPE) is a simple pointer write – no copy, no upload.

Stage 17: Warmup Forward Pass

printf("Warming up...\n");
uint32_t warmup_token = state->tokenizer.bos_id();
state->device->write_buffer(state->scratch.token_ids, &warmup_token, 4);
state->device->begin_encoding();
state->device->encode_dispatch_table(&state->dispatch_table, 0, 1);
state->device->end_encoding_sync();
state->kv_cache.reset();

Metal compiles Pipeline State Objects (PSOs) lazily – the first time a kernel is dispatched, the GPU driver compiles it from the metallib. This compilation can take tens of milliseconds per kernel. If we did not warm up, the first real inference would stutter terribly.

The warmup runs a complete forward pass for one token (the BOS token), forcing every PSO in the dispatch table to compile. Then it resets the KV cache to zero length, erasing any state written during warmup. The result: when the user sends their first prompt, every kernel is already compiled and ready to go.

Stage 18: Returning the Model Handle

printf("Ready.\n\n");
return state;

The ModelState pointer is returned as an opaque akunu_model_t handle. From here on, the caller uses this handle with akunu_generate, akunu_prefill, akunu_embed, etc. All the complexity of model loading is hidden behind a single pointer.

The Complete Flow Diagram

Let us put it all together in a single visual:

   akunu_load_model("model.gguf", "akunu.metallib", 4096)
   |
   +-- Device::create_default()
   |     |
   |     +-- MTLCreateSystemDefaultDevice()
   |     +-- query GPU cores, family
   |
   +-- load_metallib(device, path)
   |     |
   |     +-- try user path, then 4 fallback paths
   |     +-- device.load_library(path) --> MTLLibrary
   |
   +-- is_whisper_ggml(path)?
   |     |
   |     +-- NO (standard LLM path)
   |
   +-- WeightProvider::open(path)
   |     |
   |     +-- detect_format(path)  --> GGUF or MLX
   |     +-- GGUF: WeightStore::open() --> parse header, mmap tensors
   |     +-- MLX:  MLXWeightStore::open() --> parse JSON, mmap safetensors
   |
   +-- weights.get_config()
   |     |
   |     +-- AkunuModelConfig { dim=4096, layers=32, heads=32, ... }
   |
   +-- arch_from_config("llama", 4096)
   |     |
   |     +-- ArchDescriptor { silu_gate, no-qk-norm, rope_qkv_write, ... }
   |
   +-- device.chip_config()
   |     |
   |     +-- ChipConfig { cores=10, slc=8MB, chunk=64, ... }
   |
   +-- Resolve MLX RoPE, tied embeds, quant params
   |
   +-- init_rope_freqs(state, path)
   |     |
   |     +-- LLaMA3 scaling? compute wavelen-based freqs
   |     +-- Linear scaling? compute simple scaled freqs
   |     +-- Neither? leave rope_freqs null (use theta at runtime)
   |
   +-- Validate config (all dims nonzero)
   |
   +-- load_tokenizer(state, path)
   |     |
   |     +-- GGUF: read vocab/scores/merges from metadata
   |     +-- MLX: load_hf_tokenizer(dir)
   |     +-- Register extra stop tokens
   |
   +-- KVCache::create(device, 32, 8, 128, 4096)
   |     |
   |     +-- 32 layers x 2 (K+V) x 8MB = 512 MB GPU memory
   |
   +-- ScratchBuffers::create(device, config, 4096, 4096)
   |     |
   |     +-- decode: h0, h1, residual, qkv, attn_out, ffn, logits
   |     +-- prefill: batch versions of all the above
   |
   +-- build_dispatch_table(device, weights, config, arch, chip, kv, scratch)
   |     |
   |     +-- ~180 DispatchCmd structs
   |     +-- all PSOs resolved, all buffers bound, all params set
   |
   +-- Build sampled_dispatch_table (greedy + Gumbel + argmax)
   |
   +-- Pre-allocate param_buf for every command
   |
   +-- Warmup: run 1 token forward pass to compile all PSOs
   |     +-- begin_encoding()
   |     +-- encode_dispatch_table(&table, 0, 1)
   |     +-- end_encoding_sync()  <-- blocks until GPU done
   |     +-- kv_cache.reset()     <-- erase warmup state
   |
   +-- return state  (as opaque akunu_model_t)

Memory Budget

Here is a concrete example for Llama 3.1 8B (Q4_0, 4096 context) on M3 Pro:

Component               Size
-------------------------------------
Model weights (mmap)    ~4.3 GB
KV cache (32 layers)    ~512 MB
Scratch (decode)        ~2.5 MB
Scratch (prefill)       ~550 MB
Metallib (shaders)      ~2 MB
Tokenizer               ~3 MB
Dispatch table           ~30 KB
-------------------------------------
Total GPU-resident      ~5.4 GB

The weights dominate, as expected. The KV cache is the second largest allocation. Everything else is a rounding error.

Error Handling Philosophy

You may have noticed that akunu_load_model uses a simple pattern:

if (something_failed) {
    set_error("...");
    delete state;
    return nullptr;
}

There are no exceptions. There are no error codes. The function either returns a valid model handle or returns nullptr and sets a thread-local error string that the caller can retrieve with akunu_last_error(). This is the standard C API pattern – it works across language boundaries (Swift, Python, Rust FFI) and avoids the overhead and complexity of C++ exception handling in a performance-critical codebase.

The Whisper Path

For completeness, akunu_load_model also handles Whisper models. The flow is similar but with encoder-decoder-specific additions:

Whisper-specific stages:
  - load_whisper_model() instead of WeightProvider
  - arch_whisper() descriptor (LayerNorm, bias, cross-attention)
  - MelSpectrogram processor (with model's precomputed mel filters)
  - WhisperBuffers (encoder/decoder scratch)
  - Copy learned positional embeddings (encoder + decoder)
  - build_whisper_decode_table() (includes GPU suppress params)
  - Beam search buffers (5 beams x KV caches + intermediate buffers)

We will not dive deep into the Whisper path in this book, but it is worth knowing that the same akunu_load_model entry point handles both LLMs and audio models. The architecture descriptor pattern makes this possible without an explosion of conditional logic.

Summary

Model loading in akunu is a carefully ordered sequence of 18 stages that transforms a file path into a ready-to-run GPU inference pipeline. Every decision – architecture-specific behavior, chip-specific tuning, format-specific quirks – is resolved during loading and encoded into data structures (the ArchDescriptor, ChipConfig, DispatchTable) that the hot path reads but never modifies.

The result: zero allocation, zero branching, and zero format-awareness during inference. The next four chapters explore the data structures this loading process creates and how they are used to actually generate text at hundreds of tokens per second.

The Dispatch Table: Precompiled GPU Command Sequences

This is the chapter about akunu’s central innovation. If you had to distill the entire engine design into a single idea, it would be this: build the GPU command sequence once at model load time, then replay it for every token.

Most inference engines construct GPU commands on-the-fly during inference. For each token, they iterate through the model layers, look up weight tensors by name, create kernel parameter structs, resolve pipeline state objects, and emit dispatch commands. This per-token overhead is small in isolation – maybe a few hundred microseconds – but it adds up. At 70 tokens per second, each token has a 14ms budget. Spending 0.5ms on command construction is 3.5% overhead, and on smaller models it can be much worse.

Akunu eliminates this overhead entirely. The DispatchTable is a flat array of DispatchCmd structs – plain old data, no pointers to chase, no virtual calls, no hash table lookups. Built once, replayed thousands of times. The only per-token work is patching a few position-dependent fields.

The Core Data Structures

DispatchCmd: A Single GPU Command

Every GPU operation in a forward pass – every GEMV, every norm, every attention, every RoPE – is represented as a single DispatchCmd:

struct DispatchCmd {
    Pipeline pso;                   // Compiled compute pipeline

    // Buffer bindings
    static constexpr int MAX_BUFFERS = 8;
    Buffer buffers[MAX_BUFFERS];
    uint32_t offsets[MAX_BUFFERS];
    int buffer_count;

    // Inline params (up to 64 bytes)
    uint8_t param_bytes[64];
    int param_size;
    int param_index;

    // Pre-allocated GPU buffer for static params
    Buffer param_buf;

    // Secondary params
    uint8_t param2_bytes[16];
    int param2_size;
    int param2_index;

    // Threadgroup memory
    int tg_mem_bytes;
    int tg_mem_index;

    // Dispatch geometry
    Dim3 grid;
    Dim3 threadgroup;
    bool use_dispatch_threads;

    // Per-token patching
    enum PatchType : uint8_t {
        PATCH_NONE = 0,
        PATCH_TOKEN_OFFSET,
        PATCH_POSITION,
        PATCH_KV_SEQ_LEN,
        PATCH_POS_AND_KV,
        PATCH_ARGMAX_OUTPUT,
    };
    PatchType patch_type;
    int patch_offset_1;
    int patch_offset_2;
};

Let’s examine each section.

Pipeline (pso): The pre-compiled compute pipeline state object. Resolved once during table building via device.get_pipeline(). No string lookup at dispatch time.

Buffer bindings (buffers[], offsets[], buffer_count): Fixed-size array of up to 8 buffer bindings. Each entry has a Buffer handle and a byte offset. These are the weight buffers, scratch buffers, and KV cache buffers that the kernel reads from and writes to. All resolved at build time.

Inline parameters (param_bytes[64], param_size, param_index): The kernel’s parameter struct (dimensions, epsilon, strides, etc.) stored inline as raw bytes. 64 bytes is enough for every kernel parameter struct in akunu. The param_index is the argument buffer index for setBytes or setBuffer.

Pre-allocated param buffer (param_buf): For commands with static (non-position-dependent) parameters, a GPU buffer is pre-allocated containing the parameter data. At replay time, setBuffer is used instead of setBytes, avoiding the per-dispatch copy entirely.

Secondary params (param2_bytes[16]): Some kernels need two separate setBytes calls at different argument indices. The secondary param slot handles this.

Threadgroup memory (tg_mem_bytes, tg_mem_index): Some kernels require threadgroup memory allocation (e.g., GEMM tile staging). This is pre-computed.

Dispatch geometry (grid, threadgroup, use_dispatch_threads): Pre-computed grid and threadgroup dimensions. use_dispatch_threads selects between dispatchThreadgroups and dispatchThreads (the latter auto-computes grid size from total threads).

Per-token patching (patch_type, patch_offset_1, patch_offset_2): This is the key mechanism that makes replay possible despite per-token variation. More on this below.

Size and Alignment

A single DispatchCmd is approximately 280 bytes. For a 32-layer Llama model, the table contains roughly 200-250 commands, totaling ~56-70 KB. This fits comfortably in L2 cache, meaning the replay loop operates almost entirely from cache.¹

DispatchTable: The Complete Sequence

struct DispatchTable {
    std::vector<DispatchCmd> commands;      // Hot path: dense command array
    std::vector<DispatchLabel> labels;      // Cold path: profiling labels
    int tokens_per_tg;
};

The hot/cold split is deliberate. Profiling labels (48-byte strings like “layer.0.attention”) are stored in a parallel vector, separate from the command array. During inference, the labels are never accessed – the inner loop iterates only the dense commands vector. During profiling, labels are accessed by index to annotate timing data.

Per-Token Patching

Here is the fundamental challenge: most of the forward pass is identical for every token – same weights, same kernels, same dispatch geometry. But a few things change:

1. Position: RoPE needs the current sequence position. Attention needs the KV sequence length.
2. Token offset: The embedding lookup reads from token_ids[token_index]. The argmax writes to token_ids[token_index + 1].

Akunu solves this with a small enum of patch types: