Verigrad: Verification-Driven Multi-Agent GPU Kernel Generation for High-Order MLIP Derivatives
yao liu ⋅ Yuanchang Zhou ⋅ Hongtao Xu ⋅ Mingzhen Li
Abstract
High-order derivative kernels in machine learning interatomic potentials (MLIPs) are a dominant cost in scientific ML, yet current autonomous GPU-kernel agents do not reliably target this workload. They generate kernels against local tensor-level references, whereas MLIP derivative workloads span forward, backward, and higher-order automatic differentiation (AD) phases with cross-phase saved tensors and reconstruction rules. We introduce \textsc{Verigrad}, a verification-driven multi-agent harness that makes high-order derivative generation a first-class kernel-generation workload. \textsc{Verigrad} first exposes derivative semantics through \texttt{DerivativeTask} and a derivative-aware IR, and then decomposes the workload into verifiable kernel contracts. An execution-based derivative verifier validates the generated artifacts; the same verifier then gates hardware-guided refinement, so optimized kernels replace earlier candidates only when the derivative checks continue to pass. We integrate the generated kernels into MatRIS training and inference, where \textsc{Verigrad} attains 1.33--1.55$\times$ end-to-end speedup over PyTorch eager across energy-only, energy--force, and full-derivative inference and training workloads.
Chat is not available.
Successful Page Load