KernelForge

Autonomous GPU kernel engineer powered by Nemotron on Nebius. Synthesizes, compiles, and verifies custom Triton kernels with closed-loop telemetry and zero silent errors.

Comment

Inspiration

Writing performant GPU kernels in Triton or CUDA is notoriously difficult. Developers spend hours fine-tuning block sizes, managing register spilling, handling thread synchronization, and debugging cryptic numerical drift or CUDA illegal memory accesses on real hardware.

With the emergence of powerful open reasoning models like nvidia/nemotron-3-super-120b-a12b, we asked a direct engineering question: Can an LLM write, profile, and iteratively repair production-grade Triton kernels directly on physical accelerator silicon without human intervention?

We didn't want a toy syntax-checker. We set out to build an autonomous hardware-in-the-loop agent that runs candidate code on real GPUs, captures sub-millisecond execution telemetry, performs strict floating-point numerical verification, and feeds error autopsies directly back into the model context to drive convergence toward a verified champion kernel.

What It Does

KernelForge automates the GPU optimization cycle end-to-end:

  1. Hardware-in-the-Loop Synthesis: Dispatches kernel synthesis prompts to nvidia/nemotron-3-super-120b-a12b via the Nebius Token Factory API (https://api.tokenfactory.nebius.com/v1).
  2. Subprocess Isolation Harness: Launches generated kernels in isolated subprocesses (runner.py) on physical NVIDIA Tesla T4 silicon to safely isolate memory leaks, illegal memory faults, and kernel panics from the controller agent.
  3. Strict Numerical Verification: Validates candidate tensors against an eager reference under rigorous FP16 tolerances: $$|y_{\text{cand}} - y_{\text{ref}}| \le 5\times 10^{-3} + 10^{-2} \times |y_{\text{ref}}|$$
  4. Autonomous Diagnostic Autopsy: If a kernel produces numerical drift or fails at runtime, KernelForge parses the error trace, tensor dimensions, and maximum absolute error, injecting structured diagnostic feedback into Nemotron’s next generation turn.
  5. Multi-Round Interleaved Benchmarking: Runs 7 interleaved rounds of triton.testing.do_bench following an active 1-second GPU warmup to eliminate thermal and frequency fluctuations, benchmarking against both baseline PyTorch Eager and TorchInductor (torch.compile).

How We Built It

The entire pipeline was built and validated on an NVIDIA Tesla T4 GPU in a Kaggle accelerator environment:

  • LLM Reasoning Engine: nvidia/nemotron-3-super-120b-a12b hosted on Nebius Token Factory.
  • Compute & Tooling: NVIDIA Tesla T4 (16 GB GDDR6), running PyTorch 2.11.0+cu128, Triton 3.6.0, and CUDA 12.8.
  • Agent Orchestrator (agent.py): Handles context windows, prompt conditioning, AST-based security verification (rejecting arbitrary system-level calls), and early-stopping rules (halting after 2 non-improving iterations to prevent token thrashing).
  • Execution & Profiling (runner.py): Drives GPU execution, captures nanosecond-precision timings, and checks output tensors.
  • Static Replay Dashboard (dashboard/index.html): A lightweight, dependency-free dashboard that visualizes benchmark distributions, speedup deltas, and iteration-by-iteration diagnostic autopsies.

Real-World Engineering Insights & Autopsies

1. The Naive From-Scratch GEMM Trap

In early experiments, we tasked the model with synthesizing a complete matrix multiplication kernel in Triton. While mathematically correct, it took 46.06 ms (~$23\times$ slower than cuBLAS). On Turing SM75 architectures, naive global memory loads cannot match hardware-tuned micro-assembly.

2. Pragmatic Hybrid Design (cuBLAS GEMM + Triton Epilogue)

We restructured the task into an architectural hybrid:

  • Compute-Bound GEMM: Kept in cuBLAS (torch.mm), stacking projection weights (torch.cat([w1, w2], dim=0)) into a single GEMM call.
  • Bandwidth-Bound Epilogue: Targeted the SwiGLU activation and bias additions ($out = (x W_1^T + b_1) \odot \text{silu}(x W_2^T + b_2)$) inside an @triton.jit kernel, fusing memory accesses into FP32 registers and writing back once to global memory in FP16.

3. Captured Autonomous Self-Repair

During optimization loops, we observed genuine self-healing:

  • Iteration 1: Candidate failed numerical validation with a massive max error of 14.453 due to bias broadcasting and precision errors.
  • Iteration 2: Injected feedback caused Nemotron to rewrite accumulator casting (tl.float32) and coordinate indexing, yielding a numerically correct champion with a max absolute error of 0.0078 and a reproducible 1.341× speedup (1.340 ms vs 1.797 ms eager).

4. Transparent Compiler Lowering Analysis vs torch.compile

In multi-round benchmarking, the KernelForge champion ran at 1.340 ms, matching TorchInductor (1.335 ms). TorchInductor achieved parity by avoiding intermediate allocations through internal stride aliasing. In contrast, our Python wrapper incurred slight overhead copying tensor slices prior to kernel launch. Documenting this real compiler boundary provided deep, actionable insights into end-to-end memory lowering.

Feedback on Nebius Token Factory & NVIDIA Tooling

  • Nebius Token Factory: API response latency and throughput on nvidia/nemotron-3-super-120b-a12b were exceptionally fast, enabling rapid, interactive agent iteration cycles without timeouts.
  • NVIDIA Nemotron-3 Super 120B: Exhibited superior instruction-following for Triton dialect specifics (e.g., tl.constexpr, memory tiling, pointer masking) compared to general-purpose coding LLMs, quickly parsing tracebacks to fix numerical bugs.

Built With

Share this project:

Updates

Submission history