Análisis de rendimiento en PyTorch (Parte 2): De nn.Linear a un MLP fusionado

Este artículo, la segunda parte de una serie sobre el perfilado en PyTorch, profundiza en la optimización de bloques de red neuronal explorando la transición de la operación básica nn.Linear a un Perceptrón Multicapa (MLP) fusionado. Se explica cómo la adición de sesgo se integra en los kernels de multiplicación de matrices utilizando epílogos y cómo torch.compile puede fusionar operaciones puntuales, como GeLU y multiplicación, en un solo kernel de Triton. El análisis se extiende a la comparación con kernels escritos a mano, destacando las ventajas y desventajas de la especialización de kernels y la importancia de la fusión para reducir el tráfico de memoria de alto ancho de banda (HBM). Los autores también resaltan la utilidad de la biblioteca `kernels` para la distribución eficiente de kernels optimizados.