Sakana AI, a company headquartered in Tokyo specializing in artificial intelligence (AI), has unveiled a novel agentic framework designed to enhance the development and deployment speeds of large language models (LLMs). The announcement was made on Thursday, revealing the AI CUDA Engineer, which optimizes the codebase to accelerate both pre-training and inference of AI models. The firm emphasized that the entire operation is automated and driven by AI agents. Last year, Sakana AI had already introduced The AI Scientist, an AI capable of conducting scientific research.
Sakana AI Unveils AI CUDA Engineer
In a statement, the firm outlined its journey towards enhancing the deployment and inference speeds of LLMs after successfully developing AI systems that create new models and fully automate the AI research process.
This research ultimately led to the creation of the AI CUDA Engineer, a fully automated framework that focuses on CUDA (Compute Unified Device Architecture) kernel discovery and optimization.
CUDA kernels function as specialized operations that execute on Nvidia GPUs, promoting the parallel execution of code across numerous threads. This approach is superior to traditional methods, allowing for faster computational task execution, particularly when handling large datasets. Consequently, optimizing AI model deployment and inference through this methodology is considered highly effective.
Sakana AI stated that the AI CUDA Engineer can automatically transform PyTorch modules into optimized CUDA kernels, achieving significant improvements in deployment speed. It can produce kernels reported to be 10 to 100 times faster than their PyTorch counterparts.
The framework operates through a four-step process. Initially, the agent converts the PyTorch code into functional kernels. Following this, optimization techniques are applied to produce only the highest-performing kernels. Next, kernel crossover is introduced, blending multiple optimized kernels to create new variants. The final step involves archiving these high-performance CUDA kernels to enhance future performance improvements. The company has also disseminated a study that elaborates on this process.
In conjunction with the research paper, Sakana AI has released the AI CUDA Engineer Archive, a dataset containing over 30,000 kernels generated by the AI, available under the CC-By-4.0 license and accessible via Hugging Face.
Furthermore, the Japanese company launched a new website that enables users to interactively explore 17,000 verified kernels and their profiles. This platform allows visitors to navigate through these kernels across 230 tasks, as well as compare individual CUDA kernel experiments.
