KForge: Program Synthesis for Diverse AI Hardware Accelerators
This addresses the challenge of hardware-specific optimization for ML practitioners, though it appears incremental as an agent-based extension of existing program synthesis approaches.
The authors tackled the problem of optimizing GPU kernels for diverse AI hardware accelerators by developing KForge, a platform-agnostic framework using two collaborative LLM-based agents for program synthesis and performance analysis. They demonstrated effective synthesis across NVIDIA CUDA and Apple Metal platforms with a single-shot example requirement.
GPU kernels are critical for ML performance but difficult to optimize across diverse accelerators. We present KForge, a platform-agnostic framework built on two collaborative LLM-based agents: a generation agent that produces and iteratively refines programs through compilation and correctness feedback, and a performance analysis agent that interprets profiling data to guide optimization. This agent-based architecture requires only a single-shot example to target new platforms. We make three key contributions: (1) introducing an iterative refinement system where the generation agent and performance analysis agent collaborate through functional and optimization passes, interpreting diverse profiling data (from programmatic APIs to GUI-based tools) to generate actionable recommendations that guide program synthesis for arbitrary accelerators; (2) demonstrating that the generation agent effectively leverages cross-platform knowledge transfer, where a reference implementation from one architecture substantially improves generation quality for different hardware targets; and (3) validating the platform-agnostic nature of our approach by demonstrating effective program synthesis across fundamentally different parallel computing platforms: NVIDIA CUDA and Apple Metal.