Skip to content

Harness Engineering for LLM-Driven GPU Kernel Generation

🕒 Published (v1): 2026-07-20 14:14 UTC · Source: Arxiv · link

Why this paper was selected

Harness-centered LLM agent system for GPU kernels; constraint, validate, profile, select loop

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

This paper presents a harness-centered system for LLM-driven GPU kernel optimization, separating an evaluation harness (compilation, correctness, profiling, artifact archival) from a profile-backed optimization controller (bottleneck state, candidate generation decisions). Applied to the MLSys 2026 FlashInfer contest on NVIDIA B200 GPUs, Agent-Assisted kernels achieved mean-latency speedups of 1.62×–29.68× over FlashInfer baselines across five operator definitions, consistently outperforming fully autonomous (Full-Agent) approaches.

Problem

LLM coding agents can broaden kernel implementation search, but naïve generate-and-test loops fail in GPU kernel optimization due to: stale baselines, incomplete workload coverage, packaging drift, noisy promotion decisions, and loss of profiler/provenance information. Existing agent systems lack the structured harness infrastructure to make LLM search auditable, reproducible, and reliable for correctness-and-latency-gated contest kernels.

Method

The system implements a closed-loop Agent-Assisted Harness (Algorithm 1) with four layers:

  1. Evaluation harness: packages candidates, compiles in the B200 target environment (CUDA 13.2, PyTorch 2.12, Triton 3.6), runs correctness checks, measures official-aligned latency via cupti-python, archives profiler artifacts (NCU, Torch Profiler).

  2. Optimization controller: converts profiler evidence into a bottleneck state signature, selects one bounded optimization direction per round, detects plateaus/regressions, and maintains a trajectory memory of accepted and rejected probes.

  3. Skill-grounded prompting: operator constraints, profiling procedures, workload shapes, and promotion rules are encoded in reusable SKILL.md playbooks rather than ad-hoc prompts. Two skill layers—a generic CUDA optimizer and a FlashInfer B200 contest specialization—govern prompt construction and round contracts.

  4. Workload-grounded promotion gate: candidates pass a representative gate (paired against same-round baseline on sampled workloads), then a full sweep over all workload UUIDs. Promotion requires full-distribution improvement without correctness regressions (\(S_d = c_d \cdot \frac{1}{|W_d|}\sum_{w \in W_d} \frac{b_{d,w}}{\ell_{d,w}}\)).

Agents (Codex/GPT-5.3 and Claude Opus 4.6) generate CUDA C++, Triton, and CuTe/CUTLASS candidates; humans design the harness, curate references, and steer plateau recovery by switching model families or supplying reference kernels.

Key Contributions

  • Formalization of harness engineering as the critical reliability mechanism for LLM-driven kernel search, distinct from model capability.
  • A profile-backed optimization controller that converts NCU/Torch Profiler evidence into a reusable bottleneck state, enabling reproducible multi-round search.
  • Workload-grounded shape dispatch: optimization regimes derived from all contest workload UUIDs and measured latency distributions, not single-input proxies.
  • Noise-resistant promotion: all-workload latency evidence with correctness pass-rate gating; rejected probes archived as negative trajectory memory.
  • Empirical comparison of Agent-Assisted vs. Full-Agent (LoongFlow PES) under matched evaluation protocol on five contest operators.
  • Per-operator case studies showing that structural route changes (split KV flash-decode for DSA attention, Blackwell chunked CuTe for GDN prefill) drive the largest speedups.

Results

All speedups are mean-latency vs. supplied FlashInfer baseline on official-aligned B200 workloads:

  • DSA Sparse Attention: Agent-Assisted 29.68× (0.011 ms vs. 0.332 ms FlashInfer); Full-Agent 14.54×.
  • DSA Top-k Indexer: Agent-Assisted 18.05× (0.007 ms vs. 0.124 ms); Full-Agent 3.81×.
  • GDN Prefill: Agent-Assisted 13.70× (0.052 ms vs. 0.712 ms); Full-Agent 1.03×.
  • MoE FP8: Agent-Assisted 1.62× (0.286 ms vs. 0.464 ms); Full-Agent 0.27× (below FlashInfer baseline).
  • GDN Decode: Agent-Assisted 1.12× (0.006 ms vs. 0.007 ms); Full-Agent 0.83× (below FlashInfer baseline).
  • Full-Agent artifacts are 1.35×–13.25× slower than Agent-Assisted under the same normalization; two Full-Agent results fall below the FlashInfer baseline.

Limitations

  • Trial-sensitive correctness: the inner loop uses \(n=3\) trials per workload for speed; rare boundary cases (exposed in DSA top-k post-submission) can pass low-trial gates but fail high-trial replay.
  • Human labor required: harness design, reference curation, plateau steering, and final promotion decisions are not automated; the system is Agent-Assisted, not autonomous.
  • No transferable model or algorithm: results are contest-specific; the claim is about harness design patterns, not a deployable optimizer.
  • Hardware scope: all evaluations on B200 only; generalization to other GPU architectures is unvalidated.
  • Full-Agent baselines are restricted: only selected LoongFlow PES artifacts are evaluated; no systematic search-budget comparison.
  • Paper text is truncated before the full language-selection and ablation discussion.

Relevance to Agentic AI / LLM Agents

This paper directly operationalizes the emerging principle that agent reliability depends on harness design, not just model capability—a claim relevant to any agentic coding system. The harness/controller separation (measurement layer vs. decision layer) is a concrete instantiation of the human-as-environment-designer paradigm that parallels AlphaEvolve and OpenAI's harness engineering framing. The negative result—that Full-Agent autonomous search (LoongFlow PES) falls below the FlashInfer baseline on two of five operators—provides a controlled datapoint for the limits of prompt-only or evolutionary LLM search without structured feedback. For researchers tracking agentic AI, the artifact-memory and trajectory-supervision mechanisms here are directly applicable to any long-horizon code-generation agent operating under evaluation gates.