Python se ha consolidado como el líder en aprendizaje automático gracias a su ergonomía. Sin embargo, la escritura de código GPU realmente rápido ha requerido históricamente recurrir a C++ para crear núcleos personalizados y mantener las vinculaciones con Python. Esto representa una barrera significativa para muchos desarrolladores e investigadores de Python.
Frameworks como PyTorch abordan este desafío implementando núcleos en CUDA C++, ya sea de forma manual o mediante bibliotecas como las NVIDIA CUDA Core Compute Libraries. La creación de núcleos manuales es laboriosa y requiere un profundo conocimiento de la arquitectura. Utilizar CUB, una biblioteca de C++ dentro de CCCL, puede ser preferible, ya que sus primitivas están altamente optimizadas por arquitectura y se prueban rigurosamente. Sin embargo, exponer CUB a Python tradicionalmente implica construir y mantener vinculaciones, lo que limita la flexibilidad en Python.
La biblioteca cuda.compute de NVIDIA supera estas limitaciones al ofrecer una API de alto nivel y Pythonic para primitivas de CUB a nivel de dispositivo.
Resultados en la competencia GPU MODE
Utilizando cuda.compute, el equipo de CCCL de NVIDIA logró liderar la tabla de clasificación de GPU MODE, una competencia de núcleos organizada por una comunidad en línea con más de 20,000 miembros, enfocada en aprender y mejorar la programación en GPU. GPU MODE organiza competiciones para identificar las mejores implementaciones en diversas tareas, desde adiciones de vectores simples hasta multiplicaciones de matrices más complejas.
El equipo de CCCL de NVIDIA se centra en proporcionar implementaciones de primitivas paralelas “speed-of-light” (SOL) en diversas arquitecturas de GPU mediante abstracciones de alto nivel. Logró la mayor cantidad de primeros lugares en las arquitecturas de GPU probadas: NVIDIA B200, NVIDIA H100, NVIDIA A100, y NVIDIA L4.
Código optimizado en Python
CUB ofrece núcleos CUDA optimizados para operaciones paralelas comunes, incluidas las que aparecen en la competencia GPU MODE. Estos núcleos están ajustados arquitectónicamente y se consideran implementaciones cercanas al límite de velocidad.
La biblioteca cuda.compute admite tipos y operadores personalizados definidos directamente en Python. En su funcionamiento interno, compila en tiempo de ejecución (JIT) núcleos especializados y aplica optimización en tiempo de enlace para ofrecer un rendimiento cercano al de CUDA C++. Esto permite que los desarrolladores permanezcan en Python, al tiempo que obtienen la flexibilidad de las plantillas y el rendimiento de los núcleos CUDA ajustados.
Con cuda.compute, se obtiene:
- Flujos de trabajo CUDA rápidos y combinables en Python: Desarrolle aplicaciones CUDA eficientes y modulares directamente en Python.
- Tipos de datos y operadores personalizados: Utilice tipos de datos y operadores personalizados sin necesidad de vínculos con C++.
- Rendimiento optimizado: Logre un rendimiento consciente de la arquitectura a través de primitivas CUB probadas.
- Iteración rápida: Acelere el desarrollo con compilación JIT sin comprometer el rendimiento.
