Las Cinco Capas

Cuando escribes torch.matmul(A, B) , una pila sorprendentemente profunda de software transforma esa llamada Python de una sola línea en código máquina ejecutándose en una GPU. Hay al menos cinco capas distintas entre tu script y los transistores que realizan la aritmética real, y cada una existe por una buena razón. Este artículo traza el camino completo — de Python al silicio — para que puedas ver exactamente dónde vive tu cómputo en cada etapa. 

┌───────────────────────────────────────────────────────────┐
│  Capa 1: Frontend Python                                   │
│  torch.matmul(A, B)                                        │
├───────────────────────────────────────────────────────────┤
│  Capa 2: Enlace pybind11                                   │
│  Objetos Python → Objetos C++ (punteros, formas)           │
├───────────────────────────────────────────────────────────┤
│  Capa 3: Backend C++ (ATen)                                │
│  at::matmul → dispatcher → enrutar por dispositivo/dtype   │
├───────────────────────────────────────────────────────────┤
│  Capa 4: Bibliotecas de Kernels                            │
│  cuBLAS (álgebra lineal) / cuDNN (ops DNN)                 │
│  Binarios SASS precompilados con pip install torch         │
├───────────────────────────────────────────────────────────┤
│  Capa 5: Hardware                                          │
│  SASS → microcódigo → señales eléctricas → transistores    │
└───────────────────────────────────────────────────────────┘

Aquí está el resumen de una oración para cada capa. Capa 1 (Frontend Python) es la API familiar con la que interactúas — valida las entradas y registra operaciones de autograd pero no realiza cálculos. Capa 2 (pybind11) es el puente de lenguaje que traduce objetos Python en objetos C++, pasando punteros de memoria y metadatos del tensor a través de la frontera. Capa 3 (ATen + Dispatcher) es la capa de enrutamiento C++ que examina el dispositivo, dtype y disposición de cada tensor y selecciona la implementación correcta del kernel. Capa 4 (Bibliotecas de Kernels) es donde ocurre el cómputo real — bibliotecas optimizadas por el fabricante como cuBLAS y cuDNN que se envían como ensamblador GPU precompilado. Capa 5 (Hardware) es el sustrato físico: instrucciones de ensamblador GPU decodificadas en señales eléctricas que conmutan transistores en las unidades aritméticas del chip.

El resto de este artículo recorre cada capa en detalle, construyendo un modelo mental de lo que sucede desde el momento en que presionas Enter hasta el momento en que un resultado aparece en tu tensor.

Capa 1: Frontend Python

torch.matmul(A, B) es una función Python definida en el paquete Python de PyTorch. Vive en un archivo que puedes inspeccionar tú mismo ( torch/functional.py o similar, dependiendo de la versión), y no realiza casi ningún cómputo. Su trabajo principal es validar las entradas : ¿son las formas compatibles para multiplicación de matrices? ¿Están ambos tensores en el mismo dispositivo? ¿Es el dtype soportado? Si algo está mal, obtienes un mensaje de error claro a nivel Python en lugar de un segfault críptico de C++.

Una vez que la validación pasa, la función Python llama a C++ a través de pybind11 . pybind11 es una biblioteca ligera de solo encabezados que genera enlaces Python para código C++. Se encarga del trabajo engorroso de traducir objetos Python ( torch.Tensor ) en sus contrapartes C++ ( at::Tensor ), pasando a través de la frontera del lenguaje todo lo que el lado C++ necesita: punteros de memoria en bruto al almacenamiento del tensor, metadatos de forma y stride, información de dtype y la etiqueta de dispositivo (CPU, CUDA, MPS, etc.). Esta capa de enlace es la razón por la que PyTorch se siente como Python — escribes código Pythonico y obtienes errores Pythonicos — pero se ejecuta a velocidad de C++.

💡 La capa Python es también donde viven los hooks de autograd. Cuando un tensor tiene requires_grad=True, el envoltorio del lado Python registra la operación en el grafo computacional antes de pasar el cálculo real a C++. Esto significa que la sobrecarga de construcción del grafo está en Python, pero el cómputo del paso forward está en C++. Para la mayoría de las cargas de trabajo, el costo de construcción del grafo es insignificante comparado con el tiempo real de ejecución del kernel.

Capa 2: El Dispatcher C++

Dentro del backend C++, la llamada a la función aterriza en ATen (abreviatura de A Tensor library ) — la biblioteca central de operaciones con tensores de PyTorch. ATen proporciona la firma canónica en C++ para cada operación (matmul, add, conv2d, y cientos más), pero no implementa matmul para cada combinación de dispositivo/dtype directamente. En cambio, delega a un dispatcher .

El dispatcher es esencialmente una tabla de enrutamiento. Cuando se llama a at::matmul , el dispatcher examina los metadatos del tensor — dispositivo, dtype, disposición, y si el rastreo de autograd está activo — y selecciona la implementación correcta del kernel. El mismo punto de entrada C++ at::matmul puede enrutar a:

  • Una implementación CPU (usando Intel MKL u OpenBLAS) si los tensores residen en memoria principal
  • Una llamada a cuBLAS si los tensores están en una GPU CUDA
  • Un kernel completamente diferente si se involucra precisión mixta, disposición dispersa o un dtype cuantizado

Este diseño es lo que hace a PyTorch extensible. Nuevos backends — un backend para TPU, un backend para Apple MPS, un acelerador personalizado de una startup de hardware — pueden añadirse registrando nuevas claves de dispatch . El código Python del frontend no cambia en absoluto. El usuario sigue escribiendo torch.matmul(A, B) ; el dispatcher se encarga de enrutarlo al lugar correcto según dónde se encuentren A y B . Esta separación de interfaz e implementación es una de las decisiones arquitectónicas más importantes del código base de PyTorch.

El mecanismo de dispatch también maneja transformaciones componibles . Autograd, el rastreo de torch.compile, vmap (batching vectorizado) y funcionalización están todos implementados como claves de dispatch que pueden apilarse. Cuando llamas a torch.matmul en un tensor que requiere gradientes y está siendo compilado, el dispatcher encadena a través de la clave Autograd (que registra la operación para retropropagación), luego la clave de compilación (que captura la operación en un grafo), y finalmente la clave específica del dispositivo (que ejecuta el cálculo real). Cada clave hace su trabajo y reenvía a la siguiente.

Capa 3: CUDA y las Bibliotecas de Kernels

Para tensores en GPU, el dispatcher enruta a kernels CUDA. PyTorch no implementa la mayoría de los kernels CUDA desde cero — llama a bibliotecas optimizadas por NVIDIA, que han sido afinadas durante muchos años para cada microarquitectura de GPU. Las dos bibliotecas más importantes son:

cuBLAS (CUDA Basic Linear Algebra Subprograms) proporciona implementaciones altamente optimizadas de operaciones fundamentales de álgebra lineal: multiplicación de matrices, productos punto, multiplicaciones matriz-vector, resoluciones triangulares y más. Estos son los caballos de batalla detrás de cada capa lineal en una red neuronal. Cuando calculas output = weight @ input + bias , la multiplicación de matrices aterriza finalmente en una llamada cuBLAS sgemm (multiplicación general de matrices de precisión simple) o hgemm (media precisión). Los ingenieros de NVIDIA han pasado años optimizando estas rutinas, y típicamente logran más del 80% de los FLOPS teóricos pico de la GPU — mucho mejor de lo que lograría un kernel CUDA ingenuo.

cuDNN (CUDA Deep Neural Network library) proporciona implementaciones optimizadas de primitivas de aprendizaje profundo de nivel superior: convolución, pooling, normalización por lotes, softmax y (en versiones más recientes) atención multi-cabeza. Estos son los caballos de batalla detrás de las capas CNN y los bloques de atención de transformers. cuDNN no solo implementa estas operaciones de forma ingenua — selecciona entre múltiples algoritmos (por ejemplo, convolución Winograd, convolución basada en FFT, convolución directa) y los evalúa en tiempo de ejecución para encontrar el más rápido para tus formas de tensor y GPU específicas.

Tanto cuBLAS como cuDNN se envían como binarios SASS precompilados dentro del paquete Python torch . Cuando ejecutas pip install torch , estás descargando no solo código Python sino cientos de megabytes de ensamblador GPU precompilado — código máquina optimizado listo para ejecutarse en tu GPU sin ningún paso de compilación de tu parte. 

# Lo que "pip install torch" pone en tu sistema:
#
# torch/
# ├── __init__.py              ← Frontend Python
# ├── _C.so                    ← Backend C++ (compilado vía pybind11)
# ├── lib/
# │   ├── libtorch.so          ← ATen, autograd, dispatcher
# │   ├── libtorch_cuda.so     ← Implementaciones de dispatch CUDA
# │   ├── libcublas.so         ← cuBLAS (SASS precompilado)
# │   ├── libcudnn.so          ← cuDNN (SASS precompilado)
# │   └── ...
# └── nn/, optim/, ...         ← Módulos Python de alto nivel

Esto explica por qué pip install torch descarga más de 2 GB: estás obteniendo una pila completa de computación numérica, desde envoltorios Python hasta código máquina GPU, en un solo paquete.

Capa 4: Del Código Fuente al Silicio

Los binarios de cuBLAS y cuDNN no aparecieron de la nada — alguien en NVIDIA escribió código fuente CUDA en C++ y lo compiló a través de un pipeline de múltiples etapas. Entender este pipeline es valioso porque las mismas etapas se aplican cuando cualquier persona escribe un kernel CUDA personalizado (incluyendo los kernels que torch.compile genera vía Triton, que cubriremos en un artículo posterior). 

Código fuente              Compilador NVCC              Hardware
┌──────────┐    ┌──────────────────────────┐    ┌──────────────┐
│ C++ CUDA │ →  │ .cu → PTX → SASS        │ →  │ SASS → μcód. │
│ (.cu)    │    │                          │    │ → transistores│
└──────────┘    │ PTX: IR portable         │    └──────────────┘
                │   (como LLVM IR p/ GPUs) │
                │                          │
                │ SASS: ASM específico GPU │
                │   (diferente por chip:   │
                │    Ampere, Hopper, etc.)  │
                └──────────────────────────┘

PTX (Parallel Thread Execution) es la arquitectura de conjunto de instrucciones virtual de NVIDIA. Es una representación intermedia portable — el mismo código PTX puede, en principio, ejecutarse en cualquier GPU NVIDIA, independientemente de la generación. Piensa en ello como algo análogo a LLVM IR o al bytecode de Java: una capa de abstracción estable que desacopla el lenguaje fuente del hardware destino. Las instrucciones PTX describen operaciones como "multiplicar dos flotantes" o "cargar desde memoria global" sin especificar exactamente qué unidades de hardware las ejecutarán.

SASS (Streaming ASSembly) es el código máquina GPU real, específico para una microarquitectura GPU particular (identificada por un número de versión SM). PTX se compila a SASS mediante ptxas , el ensamblador de NVIDIA (parte del CUDA Toolkit). Diferentes arquitecturas GPU producen diferente SASS a partir del mismo PTX — una GPU Ampere (SM 8.0) y una GPU Hopper (SM 9.0) generarán secuencias de instrucciones diferentes, cada una afinada para los anchos de pipeline, tamaños de archivo de registros y jerarquía de memoria de ese chip.

Microcódigo es la etapa final. En el momento de carga, la unidad de control de la GPU decodifica las instrucciones SASS en señales eléctricas que conmutan transistores específicos en la ALU (Unidad Aritmético Lógica) y los archivos de registros del Streaming Multiprocessor. Aquí es donde el cómputo se convierte en física: una instrucción de multiplicación-acumulación se transforma en cambios de voltaje que se propagan a través de puertas de transistores, produciendo la suma de productos que compone una multiplicación de matrices.

¿Por qué dos etapas de compilación (PTX y luego SASS) en lugar de compilar directamente a SASS? La respuesta es un compromiso clásico entre portabilidad y rendimiento . PTX da portabilidad — escribe tu kernel una vez, y puede ejecutarse en cualquier GPU NVIDIA actual o futura. SASS da rendimiento — la programación de instrucciones, asignación de registros y patrones de acceso a memoria están optimizados para el chip específico. NVIDIA envía ambos PTX y SASS en sus bibliotecas. Si tu arquitectura GPU no fue anticipada en tiempo de compilación (por ejemplo, estás ejecutando una biblioteca compilada para Ampere en una GPU Hopper más nueva), el driver compila JIT el PTX embebido a SASS en tiempo de ejecución. Obtienes un kernel funcional inmediatamente, aunque el primer lanzamiento puede ser ligeramente más lento mientras se ejecuta el JIT.

Uniendo Todo

Tracemos una sola llamada de extremo a extremo, desde el momento en que presionas Enter hasta el momento en que un resultado aparece en tu tensor. Este es el viaje completo de result = torch.matmul(A, B) donde A y B son tensores float32 de 1024x1024 en una GPU CUDA: 

Escribes:       result = torch.matmul(A, B)

1. Python:      torch.matmul valida formas, llama a C++ vía pybind11
2. C++:         at::matmul despacha según device=CUDA, dtype=float32
3. Dispatch:    Enruta a cuBLAS sgemm (multiplicación general de matrices simple)
4. cuBLAS:      Ejecuta SASS precompilado en los Streaming Multiprocessors
5. Hardware:    Instrucciones SASS conmutan transistores ALU, multiplicando y acumulando
6. Retorno:     El almacenamiento del tensor resultado se llena, metadatos se fijan, se retorna a Python

Tiempo total:        ~0.1 ms para un matmul 1024×1024 en una GPU moderna
Sobrecarga Python:   ~5 μs (la llamada pybind11 y dispatch)
Cómputo real:        ~95 μs (ejecución GPU)

Los números anteriores son aproximados, pero la proporción es lo que importa: la sobrecarga del lado Python (validación de argumentos, cruce de pybind11, búsqueda de dispatch) es típicamente alrededor de 5 microsegundos, mientras que el cómputo GPU para una operación de tamaño razonable es de decenas a cientos de microsegundos o más. La capa Python toma aproximadamente el 5% del tiempo total para un matmul de 1024x1024, y una fracción aún menor para operaciones más grandes.

Esta proporción explica una observación práctica importante: el modo eager funciona suficientemente bien para la mayoría de las cargas de entrenamiento . Cuando cada operación toma cientos de microsegundos en la GPU, la sobrecarga de 5 microsegundos de Python es ruido. Solo empiezas a sentir la sobrecarga Python cuando las operaciones individuales son muy pequeñas (unos pocos microsegundos de trabajo GPU cada una) y muy numerosas — miles de operaciones diminutas por paso forward. En ese régimen, la sobrecarga Python por operación empieza a dominar, y puedes acumular milisegundos de tiempo desperdiciado por iteración.

💡 Este es exactamente el escenario donde torch.compile (artículo 4) ayuda. Al trazar tu código Python en un grafo y fusionar muchas operaciones pequeñas en menos kernels grandes, torch.compile elimina la sobrecarga Python por operación por completo. En lugar de cruzar la frontera Python-a-C++ mil veces, la cruzas una vez, entregas el grafo completo, y dejas que el compilador genere un único kernel optimizado.

Hay una sutileza más que vale la pena notar. Las operaciones GPU son asíncronas por defecto. Cuando Python llama a cuBLAS, no espera a que la GPU termine — encola la operación en la cola de comandos de la GPU (llamada CUDA stream) y retorna inmediatamente. Python puede entonces continuar ejecutando la siguiente línea de código (típicamente encolando la siguiente operación) mientras la GPU sigue trabajando en la anterior. La GPU y CPU corren en paralelo, como un pipeline. Solo pagas un costo de sincronización si lees explícitamente el resultado de vuelta al CPU (por ejemplo, result.item() o print(result) ), lo que fuerza al CPU a esperar hasta que la GPU termine. Este modelo de ejecución asíncrona reduce aún más el impacto práctico de la sobrecarga Python, porque el CPU generalmente está alimentando la cola de la GPU más rápido de lo que la GPU puede vaciarla.

Quiz

Pon a prueba tu comprensión de la pila de ejecución de PyTorch, desde el frontend Python hasta el hardware GPU.

¿Cuál es el rol de pybind11 en la pila de PyTorch?

¿Cómo decide el dispatcher de ATen qué kernel ejecutar para una operación dada?

¿Cuál es la diferencia entre PTX y SASS en el pipeline de compilación CUDA?

¿Por qué pip install torch incluye binarios SASS precompilados (cuBLAS, cuDNN)?