English · Español

01 — El dispatcher y aten¶

El dispatcher de PyTorch es una tabla. Toma (nombre_op, conjunto_de_claves) y devuelve el kernel a ejecutar. El conjunto de claves codifica device (CPU/CUDA), dtype, layout, requires_grad, autocast, y media docena más. aten es la biblioteca de kernels de backend que se referencian. Esta página formaliza la tabla, las claves, y traza un linear(x, W, b) a través de ella.

Esta página explica mecánicamente el dispatcher de PyTorch. Tras leerla podrás leer la salida de TORCH_LOGS=dispatcher y explicar cada línea. El laboratorio instrumentará una llamada real y registrará las decisiones.

El modelo mental: un `dict` gigante¶

Conceptualmente:

dispatcher: dict[(op_name, KeySet), Callable] = {
    ("aten::matmul", {"CPU", "Float", "Strided"}):       cpu_matmul_fp32,
    ("aten::matmul", {"CUDA", "Float", "Strided"}):      cuda_matmul_fp32,
    ("aten::matmul", {"CUDA", "Half", "Strided"}):       cuda_matmul_fp16,
    ("aten::matmul", {"CUDA", "Float", "AutogradCUDA"}): autograd_wrapper,
    ...
}

En la práctica no es literalmente un dict de Python (es una estructura de datos en C++ por rendimiento), pero el lookup lógico es idéntico. Dada una llamada a op:

y = torch.matmul(x, W.T)   # both fp32 CUDA tensors with requires_grad=True

El dispatcher calcula el conjunto de claves a partir de las entradas:

{CUDA, Float, Strided, AutogradCUDA}

Después hace lookup de ("aten::matmul", key_set) y dispatcha.

El conjunto de claves¶

Una clave de dispatch codifica uno de varios ejes:

Eje	Ejemplos
Backend	`CPU`, `CUDA`, `MPS` (Apple), `XLA`, `Meta` (sólo formas)
Layout	`Strided`, `Sparse`, `SparseCSR`, `MkldnnCPU`
"Clave" de dtype	`Float`, `Half`, `BFloat16`, `Long`, `Bool`
Bandera de autograd	`AutogradCPU`, `AutogradCUDA`, `AutogradFunctionality`
Autocast	`AutocastCPU`, `AutocastCUDA`
Fake / Meta	`FuncTorchBatched`, `BackendSelect`
...	varias más (ver `torch._C._dispatch_keys`)

El conjunto de claves completo para un tensor es la unión de todas las claves que aplican: un tensor CUDA fp32 con requires_grad=True tiene {Backend.CUDA, Layout.Strided, AutogradCUDA, ...}.

Cuando dos o más tensores entran en una op, se toma la unión a través de los tensores. El dispatcher elige la clave de mayor prioridad con un kernel registrado.

Orden de prioridad¶

Las claves tienen un orden de prioridad. Las claves de alta prioridad se comprueban primero. A grandes rasgos:

Claves de autograd (para que el grafo de backward se construya antes del cómputo real).
Claves de autocast (para que la promoción de dtype ocurra antes del backend).
Claves de Functorch / vmap (para que vmap pueda hacer batching).
Claves Sparse / Mkldnn (para que backends sparse/MKL sobrescriban).
Claves de backend (CPU, CUDA, MPS).

El recorrido descendente es: la clave de mayor prioridad con un kernel registrado gana. Ese kernel se ejecuta; si hace redispatch (un patrón común), la siguiente clave de mayor prioridad toma el relevo.

Por eso un torch.matmul con requires_grad=True dispatcha primero al kernel de autograd (que registra el grafo de backward), que re-dispatcha sin la clave de autograd al kernel CUDA (que realmente calcula).

Trazando una llamada `linear(x, W, b)`¶

torch.nn.functional.linear(x, W, b) es internamente:

# Approximately:
return torch.addmm(b, x, W.T)
# Where addmm(b, x, W^T) computes b + x @ W^T

Con la forma del ejemplo corriente de la Fase 25: x: (2, 64) fp32 CUDA requires_grad, W: (600, 64) fp32 CUDA requires_grad, b: (600,) fp32 CUDA requires_grad.

La traza del dispatch (TORCH_LOGS=dispatcher):

[step 1] aten::linear   key=AutogradCUDA → autograd wrapper
[step 2]   aten::linear key=CUDA         → calls aten::addmm internally
[step 3]   aten::addmm  key=AutogradCUDA → autograd wrapper (records AddmmBackward0)
[step 4]     aten::addmm key=CUDA        → calls cuBLAS gemm
[step 5]     cuBLAS launches kernel sm80_xmma_gemm_f32f32_*
[step 6]   y returned (shape (2, 600), grad_fn=<AddmmBackward0>)

Seis pasos. Cada uno es una sola decisión del dispatcher. El laboratorio 00 instrumenta exactamente esta secuencia con TORCH_LOGS=dispatcher y tú anotas cada línea.

aten, brevemente¶

aten es "A TENsor library" — la biblioteca de kernels en C++ a la que el dispatcher hace routing. Los operadores se nombran con prefijo aten::: aten::matmul, aten::softmax, aten::linear. Cada uno tiene potencialmente docenas de kernels registrados — uno por combinación de (backend, dtype, layout) más wrappers de autograd.

No escribes kernels de aten en la Fase 25 (eso requiere compilar PyTorch desde fuentes). Pero lees qué op de aten fue dispatcheada, y mediante qué kernel.

El equivalente del lado de Python — torch._C._dispatch_print_registrations_for_dispatch_key("CUDA") — imprime cada op que tiene un kernel registrado para CUDA. Aproximadamente 2.500 operadores. La complejidad aparente del framework es sobre todo el tamaño de la tabla.

¿Cómo entra una op nueva en el dispatcher?¶

Dos caminos:

In-tree (fuentes de PyTorch): añadir un kernel en C++ + un bloque TORCH_LIBRARY_IMPL. Requiere compilar PyTorch desde fuentes. No es tu camino.
Out-of-tree (lado de Python): torch.library.custom_op(...). Añade entradas a la tabla del dispatcher desde Python, en tiempo de import. Esta es la ruta de promoción del softmax de Triton de la Fase 24 en el laboratorio 02 de la Fase 25.

@torch.library.custom_op("mylib::softmax_triton", mutates_args=())
def softmax_triton(x: torch.Tensor) -> torch.Tensor:
    # call into Triton kernel
    return ...

# Register meta (shape inference)
@softmax_triton.register_fake
def _softmax_triton_fake(x):
    return torch.empty_like(x)

# Register backward
def _softmax_triton_backward(ctx, grad_out):
    ...
softmax_triton.register_autograd(_softmax_triton_backward, setup_context=...)

Tras el registro, torch.ops.mylib.softmax_triton(x) dispatcha igual que cualquier otra op de aten. Al dispatcher no le importa que sea "personalizada".

Autocast: el dispatch implícito por dtype¶

torch.autocast(device_type='cuda', dtype=torch.float16) activa una clave de dispatch de Autocast. Dentro de la región de autocast, las ops dispatchan primero al kernel de autocast, que puede castear las entradas a fp16 antes de re-dispatchar al kernel CUDA.

Qué ops castean y a qué dtype: una tabla hardcoded en aten/src/ATen/autocast_mode.cpp. A grandes rasgos: los GEMMs van a fp16, las reducciones se quedan en fp32, softmax se queda en fp32 (estabilidad numérica). No ves esto en el código de usuario; el dispatcher lo hace.

Esta es la primera capa que envuelve tu llamada cruda. Después autograd. Después backend. El dispatcher recorre cada capa por orden de prioridad.

Coste de rendimiento del dispatch¶

Cada decisión de dispatch cuesta ~1–2 μs de overhead en C++. Para un modelo con 100 ops por forward, eso son 100–200 μs de sólo dispatch. Para un modelo que corre a 10 ms/forward, eso es 1–2% de overhead. Para un modelo que corre a 100 μs/forward (inferencia de grammar MiniGPT pequeño), eso es 100–200% de overhead — el dispatch domina.

Por eso existe torch.compile: captura el grafo una vez, después ejecuta el kernel optimizado sin dispatch por op. CUDA Graphs (Fase 33) hacen lo mismo a nivel de lanzamiento de kernel. El laboratorio 03 de la Fase 25 mide el overhead del dispatch concretamente.

Lo que deberías ahora ser capaz de hacer¶

Enunciar el conjunto de claves de un tensor (backend, dtype, layout, autograd) inspeccionándolo.
Leer una línea de TORCH_LOGS=dispatcher y explicar qué se está dispatchando.
Encontrar a qué op de aten mapea una llamada del frontend de Python (linear → addmm, etc.).
Predecir la secuencia de dispatch para un forward a través de nn.Linear(64, 600).
Explicar por qué el dispatch de autograd ocurre antes que el dispatch de backend.

Lo que esta página NO cubre¶

__torch_dispatch__. Una sobrescritura de dispatch en modo de usuario (por tensor). Potente pero de nicho; mencionada en theory/02 sólo para el contraste con autograd.
__torch_function__. Aún más de nicho; sobrescritura a nivel de subclase. Saltarlo.
Implementación de dispatch del lado de fuentes de PyTorch (Dispatcher.cpp). Fuera de alcance.
Ruta de dispatch MPS (Apple Silicon). Mencionada; la GPU en la nube de Borja es NVIDIA, así que la ruta CUDA es el foco.

Siguiente: theory/02-autograd-engine.md — la segunda mitad: captura de grafo en forward, recorrido en backward, registro de autograd para op personalizada.