Technical Summary: KForge – LLM-Driven Cross-Platform Kernel Generation for AI Accelerators

Problem Statement

The proliferation of agentic AI applications has shifted inference workloads from monolithic execution to heterogeneous pipelines. These pipelines interleave reasoning, tool calls, and multi-agent coordination, each stage possessing distinct compute and memory profiles. Consequently, optimal deployment requires mapping specific pipeline stages to the most suitable accelerators. This heterogeneity creates a significant systems challenge: developers must produce high-performance kernels across a growing diversity of hardware backends (NVIDIA, AMD, Intel, Apple) and programming models (CUDA, Triton, SYCL, Metal, etc.).

Manually writing and porting these kernels is time-consuming, requires deep low-level expertise, and fails to scale as kernel complexity increases. While Large Language Models (LLMs) have shown promise in code generation, applying them to low-level kernel optimization remains difficult due to the brittleness of performance code (where minor changes in memory layout or precision can break correctness) and the scarcity of training data for non-CUDA platforms. Furthermore, existing benchmarks often evaluate kernels in isolation, failing to capture the impact of kernel optimization on full-model end-to-end execution within production runtimes.

Methodology: The KForge Framework

KForge is a cross-platform, multi-stage autonomous program synthesis framework designed to mirror the iterative workflow of a kernel engineer. It employs a collaborative loop between two specialized LLM-based agents:

1. Generation Agent: Responsible for producing and progressively refining kernel candidates. It operates within a prompt-driven loop that incorporates:

   • Iterative Refinement: Using compilation errors, runtime errors, and numerical mismatches from previous iterations to correct the code.
   • Cross-Platform Translation: Leveraging existing functional implementations (e.g., CUDA) as references to generate code for underrepresented platforms (e.g., Metal or SYCL).
   • Candidate Guardrails: Enforcing constraints to prevent unwanted behaviors like data type downcasting or precision changes.

2. Performance-Analysis Agent: A specialized agent distinct from the generator to handle the sparsity of optimization signals within large profiling datasets. It interprets diverse profiling inputs (programmatic metrics from tools like Nsight Systems or visual outputs from Xcode Instruments) to generate concrete code-level recommendations. These recommendations target hardware utilization metrics such as memory bandwidth, warp occupancy, and arithmetic intensity.

The Synthesis Loop:
The system alternates between two phases:

• Functional Passes: The generation agent iterates until the code compiles, executes without errors, and produces numerically correct outputs.
• Optimization Passes: Once correctness is established, the performance-analysis agent provides feedback to steer the next round of synthesis toward closing the performance gap with hand-tuned baselines.

The framework supports a uniform interface across four vendors (NVIDIA, AMD, Intel, Apple) and six programming models, utilizing a model registry to remain agnostic to the underlying LLM.

Key Contributions

• KForge Framework: Introduction of a multi-stage autonomous synthesis system where generation and performance-analysis agents collaborate using compilation, correctness, and profiling feedback.
• Cross-Platform Architecture: A system design supporting four accelerator vendors and six programming models through a unified interface, enabling the transfer of algorithmic ideas across platforms.
• Dual-Mode Evaluation: Case studies demonstrating the framework's efficacy in two distinct scenarios: competing against mature, vendor-tuned runtimes and serving as a bring-up tool for emerging hardware lacking reference implementations.

Experimental Results

The authors evaluated KForge on two backends with different baseline availabilities:

1. NVIDIA B200 (Competing with Vendor-Tuned Runtimes)

• Context: Optimizing decode-path kernels for the gpt-oss-20b model against TensorRT-LLM v1.3.0rc9.
• Micro-benchmarks: KForge-generated kernels consistently outperformed the baseline across various batch sizes. Notably, the "MoE finalize" kernel showed a speedup increasing from 1.06× at small batch sizes to 1.43× at batch size 128.
• End-to-End Performance: Integrating these optimized kernels into the full inference pipeline yielded a 2.12% improvement in end-to-end throughput (2656.61 vs. 2601.55 tokens/sec) and a 2.07% reduction in wall-clock time. The authors note that while these gains appear modest, they are significant given the years of engineering effort already invested in the vendor baseline.

2. Intel Arc B580 (Emerging Hardware Bring-Up)

• Context: Generating Triton kernels for the Intel Arc B580 (Battlemage) where comparable hand-tuned references do not exist.
• Benchmark: Evaluated on 37 GEMM + tail-ops workloads from KernelBench Level 2.
• Results: KForge achieved a 5.13× geometric mean speedup over the faster of PyTorch eager mode and torch.compile.
• Optimization Techniques: Analysis of the generated kernels revealed the successful application of:
  • Layer-wise Fusion: Collapsing post-GEMM operations into single kernels to eliminate full-tensor round-trips.
  • Mixed-Precision Execution: Utilizing FP16 for matmuls on XMX units with FP32 accumulation.
  • Single-Pass Reductions: Replacing multi-pass GroupNorm and LogSumExp with streaming online algorithms.

Significance and Claims

The paper positions KForge as a practical tool for both production-grade optimization and rapid hardware bring-up. The authors claim that:

• Agentic Synthesis is Viable: The separation of generation and performance analysis allows LLMs to effectively navigate the complex trade-offs of low-level kernel optimization.
• Scalability of Gains: Even marginal percentage improvements (e.g., ~2%) against highly optimized vendor runtimes are meaningful at deployment scale, compounding to significant efficiency gains.
• Bridging the Gap: The framework successfully addresses the "cold start" problem for new hardware (like Intel Arc) by generating performant kernels without requiring pre-existing hand-tuned references, while still competing with established runtimes on mature hardware.

The work concludes that agentic kernel synthesis is transitioning from a research concept to a functional component of the AI systems stack, capable of handling the heterogeneity of modern inference workloads.