Optimizando Multiplicación de Matrices para LLMs en Swift en Apple Silicon

La multiplicación de matrices es la operación computacional fundamental en el entrenamiento de redes neuronales y Large Language Models (LLMs), representando la mayor parte de la carga de trabajo. El problema fundamental de la computación que se aborda es cómo ejecutar eficientemente esta operación de álgebra lineal a gran escala en arquitecturas de hardware modernas, específicamente Apple Silicon, para reducir el tiempo de entrenamiento y la latencia de inferencia. La relevancia actual radica en la explosión de los LLMs y la necesidad de optimizar su ejecución en plataformas de consumo y edge, donde las bibliotecas de alto rendimiento no siempre están disponibles o son deseables por razones de control.

Históricamente, la optimización de la multiplicación de matrices ha sido un campo de estudio intenso, desde los algoritmos clásicos como Strassen hasta las implementaciones altamente optimizadas de BLAS (Basic Linear Algebra Subprograms) y LAPACK. La clave siempre ha sido explotar la jerarquía de memoria, el paralelismo a nivel de instrucción (SIMD) y el paralelismo a nivel de datos (GPU, coprocesadores especializados) para maximizar el throughput y minimizar la latencia. Este artículo busca replicar y entender estas técnicas de optimización desde cero en Swift, comparando su rendimiento con una implementación de referencia en C y explorando las capacidades de las unidades de procesamiento de Apple Silicon.

Arquitectura del Sistema

El sistema se centra en la implementación de la función matmul_forward (multiplicación de matrices) que es el corazón de la pasada hacia adelante (forward pass) de un LLM. La arquitectura evoluciona a través de varias etapas de optimización. Inicialmente, una implementación básica de Swift replica el código C de llm.c de Karpathy, utilizando Array<Float> para el almacenamiento de datos. La primera optimización clave es el uso de MutableSpan (Swift 6.2) para evitar el overhead de Copy-On-Write (COW) de Array y las comprobaciones de unicidad, mejorando el rendimiento de las operaciones de escritura.

Posteriormente, se introduce Relaxed.multiplyAdd de Swift-Numerics para habilitar las instrucciones Fused-Multiply-Add (FMA) a nivel de hardware (como fmla.4s en ARM), que combinan una multiplicación y una suma en una sola operación, reduciendo la latencia y aumentando el throughput. Para igualar la optimización del compilador C, se utiliza InlineArray (Swift 6.2) para la asignación de arrays en el stack, permitiendo el desenrollado de bucles (loop unrolling) y el tiling de datos, similar a cómo C utiliza buffers locales para mejorar la localidad de caché. El paralelismo se introduce con DispatchQueue.concurrentPerform, dividiendo la carga de trabajo en chunks para ser procesados por múltiples núcleos de CPU, aunque con la necesidad de withUnsafeMutableBufferPointer y SendableUnsafeMutableBuffer para manejar la seguridad de acceso a memoria concurrente.

Finalmente, se explora el uso de unidades de hardware especializadas: el coprocesador AMX (Apple Matrix Coprocessor) a través de instrucciones inversamente diseñadas como AMX_MATFP, AMX_LDX, AMX_LDY, AMX_STZ, que realizan operaciones de matriz 16x16 de forma nativa. La última etapa utiliza Metal, el framework de computación GPU de Apple, donde el kernel de multiplicación de matrices se escribe en Metal Shading Language (MSL) y se ejecuta en la GPU. Aquí, el paralelismo se gestiona a través de thread_position_in_grid y thread_position_in_threadgroup, y la optimización se logra mediante el tiling de matrices en threadgroup memory para explotar la localidad de caché de la GPU.

var out = out.mutableSpan

val = Relaxed.multiplyAdd(inp[bt * C + i], weight[o * C + i], val)

out.withUnsafeMutableBufferPointer { outBuffer in
  let outStorage = SendableUnsafeMutableBuffer(baseAddress: outBuffer.baseAddress!)
  DispatchQueue.concurrentPerform(iterations: chunkCount) { chunk in
    // ... lógica de procesamiento de chunk ...
  }
}

kernel void matmul_forward_kernel(
  device float* out [[buffer(0)]],
  const device float* inp [[buffer(1)]],
  const device float* weight [[buffer(2)]],
  constant uint& BT [[buffer(3)]],
  constant uint& C [[buffer(4)]],
  constant uint& OC [[buffer(5)]],
  uint2 gid [[thread_position_in_grid]],
  uint2 lid [[thread_position_in_threadgroup]]
) {
  threadgroup float inpTile[MATMUL_TILE][MATMUL_TILE];
  threadgroup float weightTile[MATMUL_TILE][MATMUL_TILE];
  // ... lógica de tiling y multiplicación ...
}

Fundamentos Teóricos

El problema de la multiplicación eficiente de matrices es un pilar de la computación numérica y el álgebra lineal, con raíces en trabajos fundamentales como el algoritmo de Strassen (1969) que demostró que la multiplicación de matrices no requiere O(N^3) operaciones. Sin embargo, en la práctica, las implementaciones optimizadas de BLAS (Basic Linear Algebra Subprograms) como OpenBLAS o Intel MKL, que se basan en algoritmos de tiling y blocking para explotar la jerarquía de caché y las instrucciones SIMD, son las más utilizadas. Estos principios se reflejan directamente en las optimizaciones presentadas, desde el desenrollado de bucles para SIMD hasta el tiling de matrices para la memoria de grupo de hilos en la GPU.

La necesidad de fused-multiply-add (FMA) se relaciona con el estándar IEEE 754 para aritmética de punto flotante, donde FMA puede mejorar la precisión y el rendimiento al realizar una multiplicación y una suma como una sola operación atómica. La gestión de la concurrencia y la seguridad de memoria en Swift, especialmente con MutableSpan y withUnsafeMutableBufferPointer, aborda los desafíos de la programación concurrente y la consistencia de datos, un campo ampliamente estudiado en sistemas distribuidos y programación paralela, donde conceptos como la coherencia de caché y los modelos de memoria son críticos para el rendimiento y la corrección.