LARK: Lean-Assisted Reasoning for LLM-Generated GPU Kernels

Zixi Zhang; Jiayi Nie; Haoran Wu; Yujie Yan; Yann Herklotz; Rika Antonova; Robert Mullins; Yiren Zhao

cs.AINov 1, 2026

LARK: Lean-Assisted Reasoning for LLM-Generated GPU Kernels

Authors: Zixi Zhang, Jiayi Nie, Haoran Wu, Yujie Yan, Yann Herklotz, Rika Antonova, Robert Mullins, Yiren Zhao,

Organizations: Department of Electrical and Electronic Engineering, Imperial College London, London, United Kingdom · Department of Computer Science and Technology, University of Cambridge, Cambridge, United Kingdom · Quettaflop AI · School of Informatics, University of Edinburgh, Edinburgh, United Kingdom ·

Abstract

Efficient GPU kernels are typically low-level and hardware-specific, and LLM agents can now generate them faster than experts can verify. Existing approaches rely on testing, which provides evidence but not a guarantee of correctness. We introduce LARK, an agentic framework that automatically and formally verifies GPU kernels against their PyTorch reference implementations. LARK proves equivalence in functional and memory behavior under exact arithmetic and establishes a new trust boundary: human effort is only needed to verify operator-level Lean semantics ahead of kernel generation time, instead of checking every single line of generated code. LARK formalizes the GPU machine model and the semantics of each operation in Lean 4 once per target language, and the Lean 4 semantics of any incoming kernel is derived from these pre-built definitions. LLM agents then automate the proof in two steps. First, an agent decomposes the whole kernel semantics into smaller segments and Lean 4 verifies this segmentation, which addresses the intractability of proving equivalence for large kernels directly. Second, agents write Lean proofs that each kernel segment matches its high-level PyTorch counterpart. Together, the two steps give whole-program equivalence. We demonstrate that LARK can be used 1) to ensure correctness in existing code-generation agents and 2) to automatically verify human-written kernel libraries. On the KernelCraft GPU tasks, combined with a simple ReAct agent, LARK enables GLM 5.3 to fully verify 10 of 12 Level-1 kernels and Opus 5.5 all 5 Level-2 kernels. Moreover, over 10 production kernels randomly sampled from popular open-source LLM serving frameworks, LARK verifies 9 and identifies a new data race bug in the tenth. We believe LARK represents a significant step toward verifiable code generation in the agentic era.

Figures & tables

Figure 1: Where trust sits in LLM kernel generation. Top: existing pipelines check a generated kernel by sampled tests or review, so the whole kernel is taken on faith. Bottom: LARK fixes what is assumed (dashed) before the kernel is seen, the machine model and operation semantics formalized once per language and the task’s interface; every property of the kernel is proved in Lean or the kernel is rejected, and the assumptions are reported with every result.

Figure 2: Running example: general matrix multiplication (GEMM). (a) The PyTorch reference. (b) The equivalence theorem: the task’s fixed premise ( StdPre : inputs placed, buffers apart, bases aligned, launch geometry) implies fault-free completion ( some s’ ) and

AB

in the region of

C

, over the reals. (c) The PTX kernel implementing (a); its per-lane addressing, shared-memory staging, barrier, and warp-collective ldmatrix / mma (lines 4–11) each need formal meaning before (b) can be stated. Line 13 (red), in place of line 15 (green), omits the barrier before the next slab is staged. Testing can only sample (b) and accepts the bug; LARK proves (b) with line 15 and locates the race with line 13.

Figure 3: The LARK pipeline. Offline, the machine model and per-operation semantics are formalized once per target language and measured against the hardware, and the Torch module once per module. Per kernel, the translator composes the flat whole-kernel semantics in Lean, then three Lean-checked stages follow: (1) the race-freedom check; (2) an agent segments the flat model and proves the segment-level model equivalent to it; (3) agents prove the segment-level model equivalent to the Torch semantics. Lean’s error reports drive the agent loops. Green: Lean-checked. Yellow: trusted, listed in the report ( Section 4.3 ). Red: agent-generated, trusted for nothing.

Kind (examples)	Guard	Update	$\mathrm{cut}$
Scalar ( add , mul )	tags, $\vec{x}\in\mathrm{dom}\,\widehat{\mathrm{op}}$	$\mathit{reg}(t,d)\mapsto(\tau,\widehat{\mathrm{op}}(\vec{x}))$	$\bot$
Memory ( ld , st )	tag, aligned, $\alpha\geq 0$	$\mathit{reg}(t,d)\leftrightarrow\mathit{mem}(\alpha)$	$\bot$
Control ( bra , ret )	$\mathit{reg}(t,p)=(\mathrm{pred},b)$	$\mathit{pc}_{t}\mapsto q$ if $b$	$\bot$ , halt
Barrier ( bar.sync )	—	parks	warp, block, grid
Collective ( mma )	tags on all lanes	all lanes’ $\mathit{reg}(t^{\prime},d)$	coll

Table 1: Operation schemas by kind: guard (failed means a fault), update, and cut; all but Control also advance

\mathit{pc}_{t}

\widehat{\mathrm{op}}

is the operation’s own function (

+

, an address

\alpha

, a branch target

q

, a tile product);

\mathrm{sem}

is defined per operation, the kinds only group them.

Target Lang.	Metric	GLM 5.3	GLM 5.3 Flash	DeepSeek V4.1 Flash	GPT-6 Luna
PTX	pass rate	83%	58%	25%	25%
	input tokens (M)	314.53	714.60	148.47	445.93
	output tokens (M)	8.05	10.54	1.93	3.26
CUDA	pass rate	75%	50%	25%	25%
	input tokens (M)	280.49	626.53	80.45	476.20
	output tokens (M)	6.18	9.09	1.12	3.42

Table 2: KernelCraft Level 1 results per target language and model: pass rate over the twelve tasks, and mean input and output tokens (millions) over the proven tasks, since a failed run spends to its budget. Per-task results in Appendix A.4.1 .

Model	Metric	FFN	SwiGLU	RoPE	FlashAttn	GQA	Total
Opus 5.5	pass	✓	✓	✓	✓	✓	100%
	input token (M)	78.20	66.60	15.50	172.50	203.90	107.34
	output token (M)	1.43	1.01	0.37	2.77	3.09	1.73
GLM 5.3	pass	✗	✗	✓	✗	✗	20%
	input token (M)	(459.60)	(512.30)	71.90	(214.10)	(611.60)	373.90
	output token (M)	(9.13)	(15.96)	2.70	(4.49)	(14.16)	9.29

Table 3: KernelCraft Level 2 results on PTX: a check mark is a certificate; tokens are input and output tokens (millions), in parentheses for runs that failed or stopped at their budget and left out of the Total average.

			Rejected / tested
ID	Rule	Change	CUDA	PTX	HIP
R1	Arithmetic	One operator or immediate ( $+$ for $-$ , a constant)	3/3	3/3	3/3
R2	Load/store	One address, buffer, or stored value	3/3	3/3	3/3
R3	Barrier	One barrier removed or moved one statement	2/2	2/2	2/2
R4	Register	One variable or source register replaced at a use	3/3	3/3	3/3
R5	Loop range	One loop bound, increment, or stride	3/3	3/3	3/3

Table 4: Mutation testing: per rule and target language, the mutants LARK rejects out of those tested. Every mutant is rejected; R3 has two mutants per language because ReLU has no barrier.

Metric	vllm-silu-and-mul	vllm-rms-norm	sglang-rmsnorm	vllm-rotary-embedding	llmc-softmax-forward
model used	GLM 5.3	GLM 5.3	GLM 5.3	GLM 5.3	GLM 5.3
verified	✓	✓	✓	✓	✗
input tokens (M)	192.90	294.00	529.00	152.80	(0)
output tokens (M)	4.56	8.12	12.23	4.44	(0)
Metric	sglang-fused-add-rmsnorm	flashinfer-rmsnorm	vllm-fused-qknorm-rope	fam-flash-attn	vllm-awq-gemm
model used	Opus 5.5	Opus 5.5	Opus 5.5	Opus 5.5	Opus 5.5

Table 5: Verification of the ten production kernels (adapted as Appendix A.14 records), each in its library’s language (CUDA, except vllm-awq-gemm on PTX) by the model in its column. A check mark is a certificate; tokens are input and output tokens, in millions.

Appendix figures & tables16 assets

Supplementary material from the paper’s appendix.

Appendix

KernelCraft Level	Tasks
Level 1: Primitive Operations	silu , relu , gelu , softmax , layernorm , rmsnorm , gemv , gemm , batchmatmul , linear , conv2d , depthwise_conv
Level 2: Composite Operations	rope , ffn , swiglu , flash_attn , gqa

Appendix

Table 6: The evaluated KernelCraft tasks by level.

		Rounds / ceiling ($)
Group	Workloads	DeepSeek V4.1 Flash	GPT-6 Luna	GLM 5.3 Flash	GLM 5.3	Opus 5.5
L1	relu , silu , gelu , softmax , layernorm , rmsnorm , gemv , gemm , linear , batchmatmul , conv2d , depthwise_conv	16 / 50	16 / 50	16 / 75	8 / 100	—
L2	rope	—	—	—	16 / 100	10 / 100
	ffn , swiglu , flash_attn , gqa	—	—	—	16 / 100	10 / 200
Library	vllm_silu_and_mul , sglang_rmsnorm , vllm_rms_norm , vllm_rotary_embedding	—	—	—	8 / 100	—
	llmc_softmax_forward , sglang_fused_add_rmsnorm , flashinfer_rmsnorm , vllm_fused_qknorm_rope	—	—	—	—	10 / 100

Appendix

Table 7: Per-verification budgets by model and task group: the cap on dialogue rounds per proof step ( --max-auto-tries ) and the spending ceiling ( --max-usd ). The round caps were read off preliminary runs in which the failure analyst always answered CONTINUE (GEMM on PTX with DeepSeek V4.1 Flash and GLM 5.3, GQA on PTX with Opus 5.5): the cap is the round of the last proof closure at the 75th to 90th percentile, plus one. The ceiling is the dollar amount at which a run was stopped; it is given for reference only.

Highest cut	Core set	Library kernels
None	relu , silu , gelu , conv2d , rope	vllm_silu_and_mul , vllm_rotary_embedding
Warp	—	vllm_fused_qknorm_rope
Block	softmax , layernorm , rmsnorm , gemv , gemm , linear , batchmatmul , depthwise_conv , flash_attn , gqa	sglang_rmsnorm , vllm_rms_norm , flashinfer_rmsnorm , sglang_fused_add_rmsnorm , llmc_softmax_forward , fam_flash_attn , vllm_awq_gemm
Grid	ffn , swiglu	—

Appendix

Table 8: The evaluation kernels by the highest cut level of their phase structure ( Section 4.1 ): the 17 kernels of the core set (KernelCraft Levels 1 and 2) and the 10 production kernels. A kernel with no cut runs as one phase; a warp-collective operation alone does not add a level.

Task (parameters)	Configuration
Level 1
SiLU, ReLU, GELU $(n)$	$(16384)$
Softmax, LayerNorm, RMSNorm $(n)$	$(16384)$
GEMV $(M,N)$	$(512,1024)$
GEMM $(M,K,N)$	$(64,256,512)$
BatchMatmul $(b,M,K,N)$	$(32,64,256,512)$

Appendix

Table 9: Configuration of each evaluated task, in the notation of KernelCraft ( Nie et al., 2026 ) : every task runs at the largest of the benchmark’s five configurations, and all tensors are FP16.

n

= number of elements (input tokens),

M,K,N

= matrix dimensions,

b

= batch size,

h_{\mathrm{in}},h_{\mathrm{out}}

= input and output features,

h

= hidden size,

h_{i}

= intermediate size,

s

= sequence length,

n_{h}

= attention heads,

n_{kv}

= key-value heads,

d_{\mathrm{head}}

= head dimension,

H

= spatial size (height

=

width),

c_{\mathrm{in}},c_{\mathrm{out}}

= input and output channels,

c

= channels,

k

= kernel size.

Flag	DeepSeek V4.1 Flash	GPT-6 Luna	GLM 5.3 Flash	GLM 5.3	Opus 5.5	What it controls
--model	deepseek-v4.1-flash-1017	gpt-6-luna	glm-5.3-flash	glm-5.3	opus-5.5	Model of the verification agents
--isa	PTX, CUDA, HIP	PTX, CUDA, HIP	PTX, CUDA, HIP	L1: all three; L2: PTX; library: CUDA	L2: PTX; library: CUDA, PTX	Target language of the run
--kernel	L1	L1	L1	L1, L2, production libraries	L2, production libraries	Task verified ( Table 6 )
--backward-effort	xhigh	xhigh	default	default	xhigh	Reasoning effort of every verification agent
--proof-effort	xhigh	xhigh	default ( max )	default ( max )	xhigh	Effort of the two proof agents
--other-effort	default ( xhigh )	default ( xhigh )	default ( high )	default ( high )	default ( xhigh )	Effort of the statement, review, and diagnosis agents

Appendix

Table 10: Configuration of the verification runs, per model. Every run passes --verify-only , so the kernel is the registry’s pinned program for its task and language and the generation flags are unused. “default” means the flag is not passed. Round caps and spending ceilings are per verification of one kernel; the ceilings are in dollars because they are what the runs were stopped at, and their derivation from the preliminary runs is in Table 7 . The last column says what each flag controls; the settings shared by every run are listed below the per-model rows.

Model	Target Lang.	Metric	ReLU	SiLU	GELU	Softmax	LayerNorm	RMSNorm	GEMV	GEMM	Linear	BatchMatmul	Conv2D	DepthwiseConv	Total
GLM 5.3	CUDA	pass/fail/budget	pass	pass	pass	pass	pass	pass	budget	pass	pass	budget	pass	budget	75%
		input tokens (M)	44.67	61.61	87.60	332.71	291.45	481.23	(700.51)	389.51	430.82	(549.06)	404.79	(562.16)	280.49
		output tokens (M)	1.78	1.78	2.38	7.97	6.32	5.93	(15.60)	8.95	11.25	(13.53)	9.23	(13.47)	6.18
	PTX	pass/fail/budget	pass	pass	pass	pass	pass	pass	pass	pass	pass	budget	pass	budget	83%
		input tokens (M)	119.20	77.30	56.50	391.20	483.20	288.60	481.00	277.70	508.70	(752.70)	461.90	(616.80)	314.53
		output tokens (M)	2.66	2.77	1.91	10.63	11.50	6.68	13.79	8.09	10.54	(15.33)	11.94	(15.93)	8.05

Appendix

Table 11: Verification of the twelve KernelCraft GPU Level 1 tasks, per model and target language. A run ends “pass” with a certificate, “budget” at its spending ceiling, or “fail” at its round cap or by an agent failure. Tokens are input and output tokens of the verification, in millions; the “Total” column gives the pass rate and the average tokens over the proven kernels. A count in parentheses belongs to a run that did not prove its kernel, reflects the budget it was given rather than the cost of a proof, and is left out of the average.

Set	Model	Segmentation proof $\mathrm{SEG}(K)$	Equivalence proof $\mathrm{EQ}(K,S)$	Diagnosis	Race-freedom fallback	Other steps
Level 1 (PTX, CUDA)	DeepSeek V4.1 Flash	271.21 (68.7%)	164.53 (36.2%; 15 of 24)	14.85 (3.6%)	—	$\leq 0.2\%$
	GPT-6 Luna	646.92 (93.6%)	267.78 (38.5%; 9 of 24)	35.95 (4.4%)	—	$\leq 0.2\%$
	GLM 5.3 Flash	719.16 (55.1%)	454.39 (48.3%; 21 of 24)	10.34 (0.7%)	—	$\leq 0.1\%$
	GLM 5.3	116.39 (42.8%)	186.60 (54.8%)	4.12 (1.3%)	—	$\leq 0.3\%$
Level 2 (PTX)	GLM 5.3	103.96 (29.6%)	388.21 (72.5%; 4 of 5)	3.29 (1.0%)	—	$\leq 0.2\%$
Level 2 (PTX)	Opus 5.5	16.26 (27.2%)	60.76 (69.4%)	2.71 (1.4%; 2 of 5)	—	$\leq 1.0\%$

Appendix

Table 12: Tokens per stage of a verification (millions), per model and evaluation set: the median over the runs that reached the stage and, in parentheses, the median share of the run’s tokens; “

k

N

” where not every run reached the stage. “Other steps” is the largest median share among segmentation, statement, and review steps. The static compile and the Lean checks inside the proof loop spend time, not tokens (text). Rows cover the same runs as Table 13 ; the rejected llm.c run (GLM 5.3) stopped in the race-freedom fallback, so it reached no proof stage.

Figure 4: Tokens of each proven verification against the compiled program’s length, over the PTX and CUDA runs, one colour per model (the number of runs in parentheses). The dashed lines are least-squares fits in log coordinates for the models with at least ten proven runs, a guide to the eye, not a fitted law; DeepSeek V4.1 Flash and GPT-6 Luna prove only the elementwise kernels and get none.

Panel of Figure 5 — Figure 5: Level 1 pass rate and tokens against the number of active parameters of the three small models with published sizes, over PTX and CUDA.

Model	Calls analyst / probe / 2nd op.	Share of cell input aggregate / median	Verdict stops	False positives refuted / issued	2nd op. refuted
DeepSeek V4.1 Flash	610 / 755 / 3	3.7% / 3.5%	1	1 / 1	3 / 3
GPT-6 Luna	599 / 815 / 5	4.7% / 4.3%	6	1 / 1	2 / 5
GLM 5.3 Flash	336 / 303 / 8	0.7% / 0.6%	1	1 / 1	6 / 8
GLM 5.3	154 / 176 / 4	1.1% / 1.1%	2	0 / 0	0 / 4
Opus 5.5	3 / 10 / 0	0.8% / 0.5%	0	0 / 0	0 / 0

Appendix

Table 13: The failure analysis inside the verification runs, per model: the Level 1 runs on PTX and CUDA, the Level 2 runs, and the production kernels. “Analysis” is every analyst, falsity-probe, and second-opinion call. Share: the analysis’s input tokens over the run’s, summed over a model’s runs and as the per-run median. Verdict stops: runs an analyst verdict ended before their spending ceiling. False positives: kernel-bug or race verdicts on a correct kernel, those refuted by the second analyst over those issued; every one was refuted, so none was acted on (two were issued before the second analyst existed and are counted as reviewed and refuted). Refuted: reviews in which the second analyst disputed the verdict, out of those it gave; a disputed verdict is not acted on.

Panel of Figure 6 — Figure 6: Verdicts of the failure analysis, as shares of the verdicts given.

	PTX (H100)				CUDA (H100)				HIP (MI300X)
Kernel	Ours ( $\mu$ s)	vs eager	vs compile	vs vendor	Ours ( $\mu$ s)	vs eager	vs compile	vs vendor	Ours ( $\mu$ s)	vs eager	vs compile	vs vendor
relu	2.05	3.53	0.96	—	2.01	3.60	0.98	—	1.73	4.35	1.07	—
silu	2.08	3.55	0.97	—	2.06	3.58	0.98	—	1.75	3.81	1.10	—
gelu	1.99	3.52	1.02	2.38	1.98	3.54	1.03	2.39	1.77	4.17	3.02	1.72
softmax	4.12	2.93	3.06	3.00	4.37	2.76	2.88	2.83	5.25	2.23	2.19	1.32
layernorm	5.03	6.63	1.78	—	4.78	6.98	1.87	—	5.05	3.79	1.69	0.85

Appendix

Table 14: Latency of the verified evaluation kernels at their pinned problem sizes (median of 7 rounds): PTX and CUDA on an H100 80GB SXM, HIP on an MI300X. The Level 2 kernels were verified on PTX only. “Ours” is the kernel’s time in microseconds; a ratio is the baseline’s time over ours on the same device, so a value above 1 means the verified kernel is faster, and the targets are not compared with each other. Baselines: torch eager (the task’s Torch module), torch.compile with max-autotune, and the fastest vendor library that offers the operation (cuBLAS, cuDNN, vLLM, FlashInfer, and the torch attention backends on the H100; hipBLASLt, rocBLAS, MIOpen, and AITER on the MI300X; a dash where none does).

Figure 7: Latency of the verified evaluation kernels relative to the baselines, from the same measurement as Table 14 : PTX and CUDA on the H100, HIP on the MI300X. The vendor libraries are cuBLAS and cuBLASLt (GEMM, GEMV, batched GEMM, linear), cuDNN (convolution), FlashInfer (RMSNorm, softmax, RoPE, prefill attention), vLLM (SiLU-and-mul, GELU, FlashAttention-2), and the torch attention backends on the H100, and hipBLASLt, rocBLAS, MIOpen, and AITER on the MI300X; “fastest” is the best of these per kernel. A dagger marks a kernel whose time is within the launch overhead.

Kernel	Lang.	Mutants a	Static b	Race-free c	Kernel bug d	Premise e	Escaped f
ReLU	CUDA	4	1	0	2	1	0
ReLU	PTX	4	0	1	2	1	0
ReLU	HIP	4	0	1	3	0	0
Softmax	CUDA	5	1	0	4	0	0
Softmax	PTX	5	1	0	4	0	0
Softmax	HIP	5	2	2	0	1	0

Appendix

Table 15: Mutation study: how each non-equivalent mutant was rejected. One mutant per rule, each counted once by its final run; ReLU has no barrier, so R3 does not apply to it.

Component	What it contains	Status	Per-language part
Machine model $M$
State and memory hierarchy	One state for the launch: registers keyed by thread and name, global, shared, and local memory with type tags, coordinates and geometry	Shared	Coordinate reads and address spellings
Value representation	Typed values: exact reals with $-\infty$ , exact integers with a width tag, predicates, packed halves; partial operations (faults)	Shared
Cut hierarchy and folds	$\mathrm{cut}$ kinds, the folds of Equation 2 , convergence, fuels	Shared	Which operations define which cut; the warp width (32 or 64)
Race freedom	Definition over footprints, span lemmas, pair-lemma library, the reduction metatheory relied on	Shared
Premise schemas	Placement, separation, parameters, alignment, geometry, fuel, entry (Appendix A.8 )	Shared

Appendix

Table 16: Components of LARK by whether they are shared across target languages or written per language, as instantiated for PTX, CUDA, and HIP. Shared components are one artifact used verbatim by every language; shared core components have one algorithm and one emitted Lean text with a thin per-language front end; per language components are written for each target and, where marked, measured against its hardware.

Stage / artifact	Produced by	Basis for acceptance	Status	If it fails
Offline, once per target language or per module
Machine model $M$	Humans, agents	Review; pinned by hash; its idealizations declared as scope	Trusted	(pre-checked)
Per-operation semantics	Generated from pinned tables	Measured against the hardware; pinned by hash	Trusted	Operation not in table: refuse
Torch module formalization	Agents, humans	Tested against the Torch module on random inputs	Trusted (tested)	Mismatch: fix before any run
Premise and extractor definitions	Pinned in $M$	One definition each over a signature; pre-checked	Trusted	(pre-checked)
Reduction metatheory	Literature	Race-free phases $\Rightarrow$ one final state for every interleaving	Trusted	Reported

Appendix

Table 17: Basis for acceptance at each stage of LARK. Verified : a Lean proof for the kernel at hand. Trusted : pre-checked or declared, kernel-independent, listed in the trusted-assumption report ( Section 4.3 ). Judged : data a run supplies, quoted in the report. Untrusted : agent-generated, covered by a later Verified row. Figure 1 uses the same statuses.

GPU Kernel Correctness Verification

See all 18 papers

Jul 20, 2026

Harness Engineering for LLM-Driven GPU Kernel Generation

Large language models (LLMs) can assist GPU kernel generation, but their practical effectiveness depends on whether generated code can be reliably constrained, validated, profiled, and selected. This paper presents a harness-centered system for LLM-driven GPU kernel optimization in the MLSys 2026 FlashInfer AI Kernel Generation Contest on NVIDIA Blackwell B200 GPUs. The system separates an evaluation harness from a profile-backed optimization controller: the harness enforces compilation, correctness, official-aligned timing, and artifact archival, while the controller turns profiler and workload evidence into bounded candidate-generation decisions. Human-authored skills capture operator constraints, references, profiling procedures, and promotion rules, while Codex and Claude Code agents generate candidate kernels inside those constraints. Across five operator definitions, the retained official-aligned artifacts achieved mean-latency speedups over supplied FlashInfer baselines of 1.62x, 18.05x, 29.68x, 1.12x, and 13.70x. The Agent-Assisted kernels outperform the Full-Agent artifacts across the evaluated definitions, indicating that expert-provided optimization directions, high-quality references, and workload context remain critical for reliable AI-driven kernel optimization.

Yue Shui, Chenyu Ma, Hangfei Xu +2

Baidu, Inc.

GPU Kernel Optimization Large Language Model-Guided Optimization GPU Kernel Generation

Ask PDF

Jul 16, 2026

Are LLM-Generated GPU Kernels Production-Ready? A Trace-Driven Benchmark and Optimization Agent

Existing GPU kernel generation benchmarks draw problems from synthetic or curated sources that diverge from deployed workloads. We present Atrex-Bench, a benchmark whose 30 operators and 440 shapes are sampled directly from full-cluster production inference traces of compute-limited, memory-rich GPUs. Each problem carries an importance weight derived from its share of observed GPU time, weighted by application card-hours and computed separately for the serving phases in which it runs, together with a per-problem roofline ceiling, so the aggregate score emphasizes the kernels that consume the most serving time. Evaluating six frontier coding agents on Atrex-Bench shows that even the best vanilla model reaches only

{\sim}10\%

of the hardware roofline on production operators; and correctness alone overstates capability, since much of the apparent pass rate comes from PyTorch fallbacks rather than kernels the model wrote. To close this gap, we co-release Atrex-Kernel-Agent (AKA), a profile-driven kernel-optimization agent that combines iterative measure-revise search, optimization dropout for escaping stalled search contexts, and a layered GPU-optimization knowledge base (298 reference-kernel files and 244 optimization-knowledge documents, plus external upstream reference projects for API/ISA lookup). In a controlled case study, the agent converts zero-FlyDSL fallbacks into real kernels that match or exceed hand-tuned production baselines.

Lingyun Yang, Yuxiao Wang, Shenghao Liang +8

ATREX Team†, Alibaba Group

GPU Kernel Optimization Large Language Model-Guided Optimization GPU Kernel Generation

Ask PDF

May 22, 2026

FastKernels: Benchmarking GPU Kernel Generation in Production

LLM-based agents for GPU kernel generation are advancing rapidly, yet their progress is fundamentally constrained by the benchmarks they optimize against. Existing benchmarks are poorly aligned with production inference frameworks: they evaluate kernels on a single GPU with synthetic inputs, ignore the surrounding compilation stack, and reward replicating known optimizations rather than discovering new ones. The resulting reward signals are misleading: agents learn to generate kernels that score well in sandboxes but introduce interface incompatibilities, compilation-stack conflicts, and silent correctness degradation when integrated into real systems. We introduce FastKernels, a kernel benchmark built around a minimal set of 46 representative architectures spanning 8 categories, whose kernels collectively subsume those of 96.2% (409/425) of HuggingFace Transformers architectures. FastKernels doubles as a minimalistic, production-grade inference framework that runs at parity with hardened systems such as vLLM and SGLang on mainstream LLM serving and substantially exceeds upstream references on under-served architectures; each task's interface mirrors the corresponding module in the state-of-the-art library for its architecture family, enabling direct deployment of optimized kernels into production codebases. Evaluating state-of-the-art kernel agents on FastKernels, we find that even the strongest agent achieves only 0.94

\times

aggregate speedup over production baselines, with weaker agents at

0.78\times

and

0.53\times

-- confirming that benchmark-production misalignment is a critical bottleneck for the field. We release FastKernels as a stepping stone toward kernel agents whose benchmark gains translate directly into production throughput improvements. Code is available at https://github.com/Snowflake-AI-Research/fastkernels

Gabriele Oliaro, Yichao Fu, May Jiang +5

1Snowflake AI Research · 2CMU · 3UCSD +1

GPU Kernel Optimization GPU Acceleration LLM Inference High-Performance Computing Large Language Model-Guided Optimization GPU Kernel Generation+1

Ask PDF

Abstract

GPU Kernel Correctness Verification

Explore similar work

Harness Engineering for LLM-Driven GPU Kernel Generation

Are LLM-Generated GPU Kernels Production-Ready? A Trace-Driven Benchmark and Optimization Agent

FastKernels: Benchmarking GPU Kernel Generation in Production