Software Advancement
Berkeley Lab researchers have developed a highly specialized AI assistant named AstraAI, specifically designed to help scientists write code for the world’s most powerful supercomputers. By studying the complex structure and rules of an existing scientific software project before suggesting any new code, AstraAI ensures that the code it generates integrates seamlessly with the project. To demonstrate its capabilities, the team successfully deployed AstraAI on AMReX, a Department of Energy-supported software infrastructure for exascale supercomputing applications. This breakthrough makes it easier, faster, and more reliable to develop the software needed to tackle major scientific challenges, from modeling turbulence to next-generation materials science.
Background
Developing scientific software for high-performance computing (HPC) is incredibly difficult and often takes decades, as it involves specialized programming conventions and complex structures where a single misplaced line can yield scientifically incorrect results. Commercial AI tools are trained on broad, open source data, which means they frequently fail in these niche setups. These models might generate perfectly good C++ or Python, but if that code does not tap into the custom internal tools required to juggle data across thousands of processors, it simply will not work. When that happens, the new code clashes with the overall system architecture, leading to structural errors, immediate crashes, or compilation failures.
Breakdown
AstraAI solves this problem by operating as a “context-aware” coder directly within the scientist’s terminal. Instead of sending a blind request to a generic AI, it builds a highly detailed prompt using two key techniques. First, through Retrieval-Augmented Generation, AstraAI instantly scans the project to find and extract relevant code examples. Second, using Abstract Syntax Tree analysis, it maps the structural blueprint of the existing code to identify how functions connect and what custom data structures are used.
By bundling the scientist’s request with these specific examples and structural rules, AstraAI gives the AI exactly what it needs to succeed. The framework acts as a flexible bridge, seamlessly connecting complex scientific codes with both massive commercial LLMs and open-source, locally hosted models. Crucially, by empowering those smaller local models to match verified benchmarks, AstraAI allows scientists to run the AI entirely on their own secure machine; this means they do not have to send sensitive, unreleased code over the internet, ensuring complete data privacy for specialized or proprietary projects.
Co-authors:
Mahesh Natarajan, Xiaoye Li, Weiqun Zhang
Publication:
AstraAI: LLMs, Retrieval, and AST-Guided Assistance for HPC Codebases
AstraAI is open sourced under the BSD 3-Clause license.
Codebase:
https://github.com/AIForHPC/AstraAI
Funding:
Department of Energy, Office of Science, Office of Advanced Scientific Computing Research’s ModCon. This effort is an early example of how DOE’s Genesis Mission is transforming science.
User Facilities:
This research used computing resources of the National Energy Research Scientific Computing Center (NERSC).
About Computing Sciences at Berkeley Lab
High performance computing plays a critical role in scientific discovery. Researchers increasingly rely on advances in computer science, mathematics, computational science, data science, and large-scale computing and networking to increase our understanding of ourselves, our planet, and our universe. Berkeley Lab's Computing Sciences Area researches, develops, and deploys new foundations, tools, and technologies to meet these needs and to advance research across a broad range of scientific disciplines.
