Technical Summary: Kernel Forge

Problem Statement

Machine learning models increasingly rely on a small set of compute kernels (e.g., matrix multiplication, convolution, normalization) for the majority of their runtime. While optimizing these kernels is a direct path to reducing latency and cost, traditional methods require expert engineers to hand-write low-level GPU code, a process that is slow, expensive, and expertise-intensive.

Recent agentic systems based on Large Language Models (LLMs) have shown promise in generating and optimizing kernels with less human effort. However, existing tools suffer from four critical limitations when applied to real-world scenarios:

1. Integration Gap: They often emit standalone CUDA code, leaving developers to manually reintegrate optimized kernels into unmodified PyTorch models.
2. Workload Narrowness: Many tools focus exclusively on LLM workloads or isolated benchmark kernels, failing to address the diverse shapes, layouts, and arguments found in vision, diffusion, and audio models.
3. Search Limitations: Current approaches often rely on linear refinement or beam-style search, which may prune promising optimization paths that initially appear suboptimal or fail to explore the solution space effectively.
4. Observability: Most systems are research prototypes or scripts lacking integrated graphical interfaces for monitoring progress, inspecting candidates, or debugging failures.

Furthermore, existing benchmarks frequently evaluate correctness and latency using randomly generated tensors and isolated kernel invocations. This fails to capture real-world model-specific conditions, such as tensor shapes, activation distributions, and neighboring operator behaviors, leading to a gap between isolated kernel performance and actual model execution.

Methodology: Kernel Forge

Kernel Forge is an open-source, end-to-end agentic harness designed to optimize real PyTorch model executions without requiring users to modify model code. The system operates through a multi-stage pipeline:

1. Model Ingestion and Operator Capture

Kernel Forge accepts any standard PyTorch model and user-provided example inputs. It executes the model to capture the specific operators running during inference. Instead of treating operators as generic families, it groups them into runtime-compatible variants based on concrete invocation patterns (tensor shapes, dtypes, layouts, arguments, and device placement). These variants are encapsulated as "operator cards," which serve as the optimization targets and preserve the PyTorch eager execution as the reference for correctness.

2. Agentic Kernel Generation

The system employs an LLM-based agent to generate specialized CUDA kernels.

• Input: The agent receives an operator card and hardware/runtime configuration.
• Process: For the first candidate, the LLM generates CUDA source code and a launch wrapper. In subsequent iterations, the agent receives a parent kernel and its lineage from the search controller to propose revisions aimed at reducing latency while preserving behavior.
• Constraint: The agent is prompted to produce custom CUDA code rather than delegating to existing libraries (e.g., cuDNN, cuBLAS), though the system acknowledges that generated code may sometimes effectively wrap these backends.

3. Validation and Repair

Generated candidates undergo a rigorous validation loop:

• Compilation: Source is compiled with nvcc. Compilation errors trigger a bounded repair loop.
• Execution: Successful candidates are launched with the captured inputs. Runtime failures (crashes, illegal memory access) are fed back to the repair loop.
• Numerical Correctness: Candidates are compared against the PyTorch eager outputs captured during ingestion, using operator-specific numerical tolerances. Only candidates passing this workload-specific filter enter the search tree.

4. Monte Carlo Tree Search (MCTS) Controller

Unlike linear refinement, Kernel Forge uses MCTS to explore multiple optimization paths.

• Tree Structure: The controller maintains a persistent revision tree where nodes represent candidate kernels, storing source code, lineage, visit counts, and measured latency.
• Selection Strategy: The controller uses a minimization version of Upper Confidence bounds applied to Trees (UCT) combined with progressive widening. It balances exploration of less-visited branches with the exploitation of promising paths, allowing the system to revisit candidates that may not show immediate gains but could lead to better solutions later.
• Feedback: Latency measurements are propagated up the tree to guide the selection of the next parent kernel for refinement.

5. Guarded Dispatch and Packaging

The final output is a guarded package (.cast) that integrates optimized kernels back into the model.

• Guarded Logic: A generated kernel is dispatched only if it is valid, compatible with the captured call pattern, numerically correct, and faster than the measured PyTorch eager latency for that specific variant.
• Fallback: If a candidate fails any check or is slower, the system automatically falls back to the original PyTorch eager path. This ensures safety and prevents performance regressions.
• Auditability: The system distinguishes between custom generated CUDA, backend wrappers, and fallback paths, providing a clear audit trail of what was actually optimized.

User Interface

Kernel Forge provides both a Graphical User Interface (GUI) and Command-Line Interface (CLI). These tools allow users to monitor optimization progress, inspect candidate kernels, view the MCTS revision tree, debug failures, and compare results across different model providers.

Key Contributions

1. End-to-End Agentic System: The introduction of Kernel Forge, a tool that optimizes real PyTorch model executions by capturing operators, generating specialized CUDA kernels, exploring multiple paths via MCTS, and automatically integrating results back into the model.
2. Observability and Usability: The provision of a GUI and CLI for monitoring, inspecting, and debugging the optimization process, bridging the gap between research prototypes and practical tools.
3. Real-World Evaluation: An evaluation across four diverse PyTorch models (ResNet-50, Stable Diffusion 3.5 Medium, Gemma 4 E2B, and Qwen 3.5 35B-A3B) representing vision, diffusion, and LLM workloads.
4. Open-Source Release: The release of the tool to support reproducible research and practical adoption of agentic kernel optimization.

Results

Evaluated on an NVIDIA DGX Spark with GB10 GPU, Kernel Forge was tested with 50 optimization iterations per kernel. The results highlight both the potential and the challenges of LLM-driven optimization in real workloads:

• Performance Gains: Kernel Forge successfully optimized 14 kernels to outperform PyTorch eager mode. Notable speedups include:
  • 1.52× on adaptive_avgpool2d in ResNet-50.
  • 1.70× on group_norm in Stable Diffusion 3.5 Medium.
  • 2.83× on softmax in Gemma 4 E2B.
  • 1.54× on softmax in Qwen 3.5 35B-A3B.
• Operator Specificity: Significant speedups were predominantly found in open-source/native PyTorch paths (e.g., tensor_iadd, relu, softmax). In contrast, operators backed by mature proprietary or vendor implementations (e.g., conv2d, linear, scaled_dot_product_attention) often remained at or below PyTorch eager performance.
• Runtime Impact: The operators that achieved the highest local speedups often accounted for a small fraction of the total captured runtime. Conversely, the dominant runtime contributors were frequently the highly optimized vendor-backed paths that were difficult to improve upon.
• Search Cost: The evaluation noted that search costs (LLM API usage) increased significantly at higher iteration counts (opt50), but this did not always correlate with proportional gains in aggregate operator-region performance, particularly for dominant operators.

Significance and Claims

The paper claims that Kernel Forge demonstrates that LLM-driven kernel optimization can be effectively applied beyond isolated benchmarks to operators captured from full PyTorch models.

The authors emphasize the importance of deployment-aware evaluation. They argue that while LLMs can generate fast standalone kernels, the true value lies in how these kernels behave within real model executions. The "guarded dispatch" mechanism is presented as a critical innovation, allowing the system to apply generated kernels only where evidence supports a benefit while preserving the reliability of existing framework implementations.

The paper concludes that agentic kernel optimization should be evaluated not just by its ability to generate fast code, but by its ability to navigate the trade-offs between optimization depth, search cost, and the specific characteristics of real-world workloads, particularly the prevalence of mature backend implementations for dominant operators.