AstraAI：LLM、リトリーバル、およびASTガイド付き支援によるHPCコードベースの支援

Abstract

高性能計算（HPC）ソフトウェア開発のためのコマンドラインインタフェース（CLI）コーディングフレームワークであるAstraAIを提案します。AstraAIはLinuxターミナル上で直接動作し、大規模言語モデル（LLM）を、Retrieval-Augmented Generation（RAG）および抽象構文木（AST）に基づく構造解析と統合することで、複雑な科学コードベースに対する文脈に応じたコード生成を可能にします。中核となる考え方は、推論のためにLLMへ渡される高忠実度プロンプトを構築することです。このプロンプトは、RAGによって基盤となるフレームワークのコードベースから取得した関連コードスニペットと、AST解析から抽出した構造的文脈を用いてユーザ要求を補強し、関連する関数、データ構造、およびコード全体の組織に関する正確な情報をモデルに提供します。このフレームワークは、周辺コードとの構造的整合性を保ちながら、ソースコードに対してスコープ付きの修正を行うよう設計されています。AstraAIは、Hugging Faceからローカルホストできるモデルと、American Science Cloud経由でアクセス可能なAPIベースのフロンティアモデルの両方をサポートしており、HPC環境全体で柔軟な導入を可能にします。システムは、既存プロジェクトの構造およびプログラミングパターンに整合したコードを生成します。AMReX内で、エクサスケール応用向けのDOE支援によるHPCソフトウェア基盤であるAMReXにおいて、代表的なHPCコード生成タスクでAstraAIを実演します。