La inferencia de modelos de lenguaje grandes (LLMs) en el edge o en entornos con restricciones de costo y energía presenta un desafío fundamental en la arquitectura de computadoras. Los procesadores de propósito general (CPUs) carecen de la eficiencia energética y el throughput necesario para las operaciones intensivas en multiplicación de matrices, mientras que las GPUs de alta gama son costosas y consumen mucha energía. La tesis de openTPU es que un diseño de hardware especializado, optimizado para las operaciones dominantes en LLMs (principalmente multiplicación de matrices y operaciones vectoriales), puede ofrecer una solución eficiente y de bajo costo. Al ser open-source y diseñado con la ayuda de agentes de IA, también explora la automatización del diseño de hardware, un problema que se remonta a los inicios del diseño de circuitos integrados.
Este proyecto aborda la necesidad de aceleradores de IA accesibles y transparentes, permitiendo a los ingenieros Staff+ y Arquitectos comprender la pila completa, desde el algoritmo de ML hasta los "wires" del hardware. La elección de FPGAs permite una rápida iteración y validación del diseño, así como la implementación en hardware real sin los costos y tiempos de fabricación de ASICs, sirviendo como un banco de pruebas para futuras arquitecturas de inferencia de IA.
Arquitectura del Sistema
La arquitectura de openTPU es deliberadamente simple, centrada en la eficiencia y la transparencia. El diseño reside en un monorepo que abarca SystemVerilog para el hardware, un Instruction Set Architecture (ISA) de 8x32 bits, un simulador bit-exacto en Python, un lenguaje kernel (ol) con su compilador, y software host para la interfaz PCIe. El hardware principal consiste en un secuenciador que emite una instrucción por ciclo a unidades especializadas: una unidad DMA para movimiento de datos, una unidad de matriz sistólica de cuatro columnas para multiplicaciones int8, una unidad vectorial para aritmética fp32, y un cuantificador para convertir resultados a int8. No hay caché ni scheduling oculto; cada movimiento de datos es una instrucción explícita, lo que facilita el profiling.
La comunicación con la DRAM se gestiona mediante controladores LiteDRAM, calibrados por una pequeña CPU embebida en el core de memoria. Esto es crítico para alcanzar la eficiencia de ancho de banda necesaria (82-85% del pico de DDR3-1066). El software host (otpu-chat, otpu-smi, otpu-lens) interactúa con la tarjeta FPGA a través de PCIe. El compilador de ol traduce kernels de alto nivel (como mlp o attention) a la ISA de openTPU, optimizando layouts y direccionamiento de bucles afines. Para modelos que exceden los 4 GiB de DRAM de la tarjeta (como los Mixture-of-Experts), se implementa un mecanismo de offload donde los expertos se transmiten desde el almacenamiento del host a slots en la DRAM de la tarjeta, con el host copiando expertos faltantes a la tasa del enlace PCIe. La cuantización de pesos a 4-bit (FP4 con escalas de bloque de dos niveles) y activaciones a int8 es una decisión clave para reducir el ancho de banda de memoria y aumentar la velocidad de decodificación, manteniendo la cabeza del LM en int8 para precisión.
Flujo de Inferencia de Token (Decode)
- 1 Host Envía prompt inicial y programa de decodificación a la FPGA.
- 2 FPGA - DMA Carga embeddings y filas RoPE desde DRAM.
- 3 FPGA - Matrix Unit Realiza multiplicaciones de matrices (matmul) con pesos int8.
- 4 FPGA - Vector Unit Ejecuta operaciones fp32 (activaciones, normalización).
- 5 FPGA - Quantizer Convierte resultados a int8 para la siguiente capa.
- 6 FPGA - DMA Transmite logits de vuelta al host.
- 7 Host Realiza argmax para seleccionar el siguiente token y lo envía de vuelta a la ...
Flujo de Offload de Expertos (MoE)
- 1 FPGA Procesa token y determina qué expertos son necesarios.
- 2 FPGA - DRAM Busca expertos en slots de DRAM locales.
- 3 FPGA - PCIe Request Si el experto no está en DRAM, solicita al host.
- 4 Host - Storage Lee experto del pool de archivos en disco.
- 5 Host - PCIe Transmite experto a la FPGA.
- 6 FPGA - DMA Copia experto recibido a un slot de DRAM.
- 7 FPGA - Matrix Unit Ejecuta computación del experto.
| Capa | Tecnología | Justificación |
|---|---|---|
| compute | Xilinx Kintex-7 FPGA | Plataforma de hardware reconfigurable para implementar el acelerador de IA. Permite la validación rápida del diseño y la ejecución en hardware real. vs ASICs (Application-Specific Integrated Circuits), GPUs de propósito general, CPUs xc7k480t, dos canales DDR3-1066, 4 GiB de DRAM. |
| compute | SystemVerilog | Lenguaje de descripción de hardware (HDL) para el diseño RTL del acelerador openTPU, incluyendo la unidad de matriz sistólica, unidad vectorial y DMA. vs VHDL, Chisel |
| data-processing | LiteDRAM controllers | Gestión de la interfaz con la memoria DDR3, crucial para alcanzar alta eficiencia de ancho de banda (82-85% del pico) en operaciones de inferencia. vs Xilinx MIG (Memory Interface Generator) Calibrados por una CPU interna en el core de memoria. |
| orchestration | PCIe | Interfaz de comunicación de alta velocidad entre el host (CPU) y la tarjeta FPGA, utilizada para cargar modelos, enviar prompts y recibir logits. vs USB, Ethernet |
| data-processing | openTPU kernel language (ol) + compiler | Lenguaje de alto nivel para escribir kernels de ML (ej. MLP, atención) que se compilan a la ISA de openTPU, optimizando el uso del hardware. vs Programación directa en ensamblador ISA, Compiladores de ML existentes (ej. TVM, XLA) Soporta layouts, direccionamiento de bucles afines y fusión. |
| observability | otpu-lens | Herramienta de profiling que registra y visualiza la ejecución del acelerador (RTL, simulador o tarjeta), mostrando el uso de unidades y los ciclos de espera. vs Waveform viewers (ej. GTKWave) Visualización de roofline, timeline y tablas por instrucción. |
Trade-offs
Ganancias
- ▲ Velocidad de decodificación (throughput)
- ▲ Eficiencia energética (implícita por hardware especializado)
- ▲ Transparencia y capacidad de profiling
- ▲ Costo de hardware (FPGA vs GPU de alta gama)
Costes
- △ Precisión del modelo (por cuantización 4-bit)
- ▲ Flexibilidad (hardware especializado vs GPU de propósito general)
- ▲ Complejidad de desarrollo (diseño de hardware)
from opentpu import language as ol
@ol.jit
def mlp(h, gamma, w_gate, w_up, w_down, out, eps):
x = ol.load(h)
xs = ol.quantize(rmsnorm(x, ol.load(gamma), eps))
g = ol.dot(xs, w_gate)
u = ol.dot(xs, w_up)
a = ol.all_gather(silu(g) * u)
y = ol.all_gather(ol.dot(a, w_down))
if ol.program_id() == 0:
ol.store(out, x + y)Fundamentos Teóricos
El diseño de aceleradores especializados para cargas de trabajo específicas, como las redes neuronales, tiene sus raíces en el concepto de arquitecturas de dominio específico (DSA). Este enfoque fue popularizado por trabajos como el de Google con sus Tensor Processing Units (TPUs), que demostraron la eficiencia de hardware customizado para operaciones de machine learning. La idea de una unidad de matriz sistólica, central en openTPU, se remonta a los trabajos pioneros de H.T. Kung y Charles E. Leiserson en los años 70, que exploraron arquitecturas paralelas para el procesamiento de señales y álgebra lineal. Su paper "Systolic Arrays (for VLSI)" (1979) sentó las bases para el diseño de arrays de procesadores interconectados que operan en un flujo de datos rítmico, minimizando el movimiento de datos y maximizando el reuso.
La cuantización de modelos, especialmente a formatos de baja precisión como int8 o 4-bit, es un área activa de investigación en la academia y la industria, buscando reducir la huella de memoria y el costo computacional sin una pérdida significativa de precisión. Trabajos como "Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference" de Jacob et al. (2018) de Google, o "Deep Learning with Low-Precision Number Formats" de Courbariaux et al. (2016), han explorado los compromisos entre precisión y eficiencia. openTPU aplica estos principios, utilizando cuantización de pesos a 4-bit y activaciones a int8, y valida su impacto en la precisión mediante métricas como la divergencia KL y el acuerdo top-1, conectando directamente con la investigación en la robustez de la cuantización.