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

2026-06-01

AI Summary

This work addresses the challenges of high manual development costs and poor cross-platform portability in generating high-performance kernels for heterogeneous AI accelerators. The authors propose an iterative optimization framework based on two collaborative large language model (LLM) agents, which alternately perform functional correctness verification and performance tuning to enable reference-free, cross-backend kernel generation. By integrating compiler feedback, performance profiling, operator fusion, and mixed-precision execution, the method achieves a 2.12% improvement in inference throughput on the NVIDIA B200 and delivers an average 5.13× speedup over Triton kernels across 37 GEMM-plus-epilogue tasks on the Intel Arc B580.

Abstract

Production inference increasingly targets a heterogeneous mix of accelerators. Agentic pipelines interleave reasoning, tool calls, and multi-agent coordination, each with distinct compute and memory profiles. For optimal efficiency, each stage should run on the accelerator best suited to it. This creates a systems challenge: each pipeline now requires high-performance kernels across a growing set of hardware backends and programming models. Writing these kernels by hand is time-consuming, demands deep low-level expertise, and does not scale as kernel complexity grows. Recently, Large Language Models (LLMs) have been leveraged for automatic kernel generation, but challenges in low-level code generation and cross-backend generalization persist. We present KForge, a cross-platform framework built around an iterative refinement loop driven by two collaborating LLM-based agents: a generation agent that produces and progressively refines kernels using compilation and correctness feedback, and a performance-analysis agent that interprets profiling data, from programmatic APIs to GUI-based tools, and emits recommendations that steer the next round of synthesis. The loop alternates between functional passes, which drive a candidate to correctness, and optimization passes, which close the performance gap to hand-tuned baselines. We evaluate KForge on two backends with very different baseline reference availability. On NVIDIA B200, KForge achieves a 2.12$\%$ improvement in end-to-end throughput compared to TensorRT-LLM on the gpt-oss-20b inference speed benchmark. On Intel Arc B580, KForge generates Triton kernels achieving a 5.13$\times$ geometric mean speedup over the faster of PyTorch eager and torch.compile on 37 GEMM + tail-ops workloads from KernelBench Level 2, primarily via operator fusion and mixed-precision execution.

Problem

• AI accelerators
• kernel generation
• heterogeneous systems
• cross-platform
• performance optimization

Innovation

• LLM-driven kernel generation
• cross-platform optimization
• iterative refinement loop
• multi-agent collaboration
• operator fusion