PaperlayerView exact revision on arXiv

Source-linked AI summary

D2K-Bench: Can LLM Agents Turn Expert Designs into Efficient GPU Kernels?

Daifeng Li, Huiqiang Jiang, Chengruidong Zhang, Wei Wu, Xudong Guo, Jianhong Tu, Jianwei Zhang, Binhang Yuan, Dayiheng Liu

arXiv:2610.03226v1cs.LGcs.AIcs.DC

TL;DR

Runtime alone does not show how the gap between LLM-generated and expert GPU kernels relates to design discovery and implementation. D2K-Bench pairs performance comparisons with assessments of proposed designs and generated code, with and without expert guidance. Across five models, guidance improved correctness and performance, while generated code still left assessed design properties unimplemented.

  • Problem

    Runtime alone does not reveal how the gap between LLM-generated and expert GPU kernels relates to design discovery and implementation.

    Source: §1 · p. 1

  • Method

    D2K-Bench pairs performance comparisons with and without expert design guidance with assessments of independently proposed designs and generated code.

    Source: §1 · p. 2

  • Results

    Across five models, guidance raises correctness from 93.1% to 98.5% over 130 model–task pairs, while Performance Score increases from 1.46 to 1.95 over all 26 tasks.

    Source: §6 · p. 10

  • Takeaways & Limitations

    Guidance improves performance, but agents still leave assessed design properties unimplemented.

    Source: §6 · p. 10

    The SOL score is not reported because algorithmic or dataflow changes can invalidate the bound derived from the PyTorch reference.

    Source: §3.2 · p. 7

Abstract

from arXiv · show

GPU kernels generated by large language model (LLM) agents can remain less efficient than expert implementations, but runtime alone does not reveal how the gap relates to design discovery and implementation. We introduce D2K-Bench, a diagnostic benchmark of 26 tasks and 85 workloads that measures how effectively agents translate expert design guidance into efficient GPU kernels. The guidance covers L1: high-level algorithmic insights, L2: dataflow design, and L3: low-level optimization tricks, including dependencies among these levels. Pairwise runs with and without guidance share task descriptions, workloads, tools, hardware, and a 350-turn budget. Complementary assessments examine independently proposed designs and the design properties implemented in generated code. Across five models on NVIDIA B200 GPUs, guidance raises correctness over 130 model-task pairs from 93.1% to 98.5% and increases the Performance Score over all 26 tasks from 1.46 to 1.95. For the three frontier models with correct submissions on all 26 tasks in both runs (GPT-6-Astra, Claude-Opus-4.8, and GPT-5.6-Sol), geometric mean speedup increases from $1.69\times$ to $2.49\times$. Across all five models, the mean combined implementation score increases from 57 to 70 out of 100. These results show the value of expert design guidance while identifying design properties that remain unimplemented.

1 Introduction

D2K-Bench examines how well LLM agents translate expert GPU-kernel design guidance into efficient implementations. Its diagnostic approach connects performance with design discovery and code implementation, and guidance improves results while leaving some design properties unimplemented.

1 Introduction

  • Runtime alone cannot show which design choices generated code implements or how well, so the paper assesses proposed designs and executable source alongside performance.Effective algorithms and dataflows must also be implemented through correct code, GPU execution mechanisms, and tuning.Source: §1 · p. 2
  • D2K-Bench pairs task-specific expert guidance with comparisons of runs with and without guidance, making guidance benefits and unimplemented design properties observable.Guidance covers L1 algorithmic insights, L2 dataflow design, and L3 low-level optimization tricks, including dependencies among levels.Source: §1 · p. 2
  • Paired runs share task descriptions, workloads, Triton, tools, and hardware; an LLM judge assesses independent designs and source evidence, crediting correct alternatives.Runtime is hidden during source assessment, and comparisons cover all 26 tasks in both runs.Source: §1 · p. 2
  • Across five models, guidance raises correctness from 93.1% to 98.5%, Performance Score from 1.46 to 1.95, and mean combined implementation score from 57 to 70 out of 100.The Performance Score is computed over all 26 tasks for each model.Source: §1 · p. 2

2 Benchmark Design and Construction

D2K-Bench combines expert-selected GPU-kernel tasks and workloads with shared task specifications and three-level expert design guidance. Its Muon example illustrates how algorithmic transformations, dataflow choices, and execution mechanisms are connected.

2.1 Benchmark Task Selection

  • The expert-selected suite contains 26 tasks and 85 workloads spanning kernel development scenarios in LLM training and inference, including attention, mixture of experts, and quantization.Tasks are selected for expert implementations that reduce computation or memory traffic through algorithmic transformations or dataflow changes, with low-level optimizations providing further gains.Source: §2.1 · p. 3

2.2 Task Specification

  • Shared task files specify outputs, tolerances, state updates, workload shapes, data types, and layouts, and require custom kernels rather than existing operator implementations.Each task also has an input generator, PyTorch reference, implementation requirements, and a solution template with an unimplemented run(...) entry point.Source: §2.2 · p. 3

2.3 Expert Design Guidance

  • Guidance has three levels: algorithmic transformations and their rationale, dataflow choices for variable placement and reuse, and low-level execution mechanisms with required ordering and synchronization.Dependencies connect the algorithmic transformation to its data representation and movement, then to GPU execution mechanisms.Source: §2.3 · pp. 3–4
  • For Muon, the guidance composes updates on smaller square states R and Q, keeping rectangular X fixed until a segment boundary instead of updating X after every step.At the fixed restart before step 2, it materializes and rounds X to FP16, rebuilds R to limit numerical drift, and starts a fresh composition.Source: §2.3 · p. 4
  • Muon’s low-level schedule runs square-state products within each segment and applies rectangular products at restart and final boundaries; symmetric kernels reuse one triangle’s accumulators.Diagonal tiles are written only once.Source: §2.3 · pp. 4–5
  • An LLM judge assesses independently written designs and generated source against task-specific criteria describing expert choices, the costs they remove, and requirements for combining them.Implementation assessment includes the guidance, expert implementation, and candidate source code.Source: §2.3 · p. 5

3 Evaluation Protocol

D2K-Bench pairs guided and unguided agent runs, measures correctness, latency, and performance, and uses an LLM judge to assess proposed designs and generated implementations. Its protocol also notes that the SOL score can become uninterpretable when generated implementations change the computation or memory traffic assumed by the bound.

3.1 Agentic Evaluation Pipeline

  • Paired guided and unguided runs use the same pipeline, tools, hardware, implementation language, and optimization budget; agents iteratively implement, test, profile, and revise kernels.Final selection is the eligible candidate with the highest geometric mean speedup across task workloads.Source: §3.1 · p. 5

3.2 Objective Evaluation Metrics

  • Correctness requires passing numerical, state-update, repeated-execution, implementation, and applicable distributed checks on every workload; latency timing excludes setup and compilation.Short-running tasks use CUDA Graphs; other tasks use CUDA events, with repeated per-call measurements averaged in both modes.Source: §3.2 · p. 6
  • The Performance Score compares final submissions with Fable-5’s first correct unguided submission, averages within and then across tasks, and uses a fixed 0.5 penalty for failed tasks.The baseline has Sperf = 1.0 by definition; overall Sperf is the geometric mean of the five model scores.Source: §3.2 · pp. 6–7
  • Algorithmic and dataflow changes can invalidate the PyTorch-reference SOL bound; transformed Muon implementations can fall below it, yielding scores above one or negative scores, so SOL is omitted.Source: §3.2 · p. 7

3.3 LLM Judge Scores for Designs and Implementations

  • The LLM judge scores generated code at L1–L3 from 0 to 100 using code evidence, hides runtime, credits alternative designs, and assigns zero to incorrect submissions.The three level scores are averaged; functional correctness is assessed separately, and the scores describe design quality and implementation completeness rather than directly estimating speedup.Source: §3.3 · p. 7
  • Independent design proposals are elicited without guidance or code generation and scored against task-specific L1–L3 criteria; model judge scores are averaged over all 26 tasks.Source: §3.3 · p. 7

4 Experiments

Across five models, expert guidance improves correctness and performance, while runtime-blind design and implementation scores broadly align with measured performance. Cumulative guidance raises performance at every level, with different levels contributing differently across models.

4 Experiments

  • The experiments compare five models across 26 tasks using a 350-turn budget per task, with and without comprehensive expert guidance.For cumulative ablations, three models with correct implementations on every task receive fixed inputs of no guidance, L1, L1+L2, or L1+L2+L3.Source: §4 · p. 8, §4.3 · p. 8

4.2 RQ2: Alignment of LLM Judge Scores with Measured Performance

  • Runtime-blind implementation scores align with broad performance differences and guidance gains, while independent design scores reflect broad performance differences.These alignments are approximate and demonstrated across model aggregates.Source: §4.2 · p. 8, Table 2 · p. 9

4.1 RQ1: Correctness and Performance Gains with Expert Guidance

  • Correctness rises from 93.1% to 98.5% and overall Sperf from 1.46 to 1.95 with comprehensive guidance; three all-task models’ geometric mean speedup rises from 1.69× to 2.49×.The overall Sperf increase is 33.9%; Sperf weights models equally across all 26 tasks.Source: §4.1 · p. 8, Table 1 · p. 9

4.3 RQ3: Effects of Design Guidance Levels

  • Across three models, performance rises at every cumulative guidance step; averaged across them, L1 yields the largest incremental gain, with L2 and L3 adding improvements.GPT-6-Astra gains most from L1, while GPT-5.6-Sol and Claude-Opus-4.8 gain more from L3 than L2.Source: §4.3 · pp. 8–9, Figure 3 · p. 9

5 Related Work

Related work spans kernel-generation benchmarks with varied tasks and evaluation methods, and agentic systems that optimize kernels through iterative feedback, reusable knowledge, and hardware guidance.

GPU kernel generation benchmarks

  • Kernel-generation benchmarks span GPU operators, heterogeneous hardware, database queries, production interfaces, serving workloads, and inference traces.Examples include TritonBench, MultiKernelBench [11], DataKernelBench [12], FastKernels [3], FlashInfer-Bench, and Atrex-Bench [20].Source: §5 · pp. 9–10
  • Benchmark evaluation methods range from PyTorch speedup and analytical hardware bounds to verification, hidden testing, and hardware-efficiency or generalization analyses.KernelBench established the task of producing a correct kernel faster than its PyTorch reference; SOL-ExecBench [5] uses analytically derived hardware bounds.Source: §5 · pp. 9–10

Kernel optimization by LLM agents

  • Kernel-optimization agents use iterative workflows, reusable knowledge, learned or synthetic examples, staged search, and expert-provided hardware guidance.Hari et al. [2] use Speed-of-Light guidance, while Shui et al. [6] report that expert directions can improve final implementations in a B200 contest harness.Source: §5 · p. 10

6 Conclusion

D2K-BENCH evaluates how LLM agents translate expert design guidance into efficient GPU kernels. Across five models, guidance improves correctness, performance, and implementation scores, while assessed design properties remain unimplemented.

6 Conclusion

  • D2K-BENCH pairs 26 tasks and 85 workloads with guidance on algorithmic insights, dataflow design, and low-level optimization tricks.Its evaluation combines pairwise performance comparisons with assessments of independently proposed designs and generated code.Source: §6 · p. 10
  • Correctness rises from 93.1% to 98.5% across 130 model–task pairs, while Performance Score increases from 1.46 to 1.95 over all 26 tasks.Across five models, mean combined implementation score rises from 57 to 70.Source: §6 · p. 10
  • Guidance improves performance, but agents still leave assessed design properties unimplemented.Source: §6 · p. 10

A.1 Benchmark Composition

Table 3 names operator categories and representative design content in D2K-BENCH.

A.1 Benchmark Composition

  • Table 3 presents operator categories alongside representative design content.Source: Table 3 · p. 13
  • GEMM denotes general matrix multiplication.Source: Table 3 · p. 13

A.2 Experimental Environment

Experiments ran on an eight-GPU NVIDIA B200 system, with specified software and agent tooling.

A.2 Experimental Environment

  • The evaluation system used eight NVIDIA B200 GPUs connected via NVLink and NVSwitch on x86_64 Linux.Source: §A.2 · p. 13
  • Experiments used Python 3.12 in Docker containers, CUDA Toolkit 13.0, and a prerelease PyTorch 2.9 build (cu130).Source: §A.2 · p. 13
  • Claude Code CLI 2.1.197 served four models, Codex CLI 0.147.0 served GPT models and the Qwen3.8-Max-0902 judge.Claude Code was used for Claude-Opus-4.8, GLM-5.3, Kimi-K3, and Fable-5; Codex was used for GPT-6-Astra, GPT-5.6-Sol, and the judge.Source: §A.2 · p. 13

A.3 Speedup for Fully Correct Models

This section reports speedup for fully correct models and a secondary thresholded-coverage analysis across all model–task pairs.

A.3 Speedup for Fully Correct Models

  • Table 4 reports speedup over all 26 tasks against Fable-5’s first-correct baseline for models correct on every workload in both configurations.The speedup is aggregated geometrically within each task and then across all 26 tasks.Source: Table 4 · p. 13, §A.3 · p. 13
  • Thresholded coverage counts tasks whose task-level geometric mean speedup over Fable-5 exceeds threshold p, among tasks with correct final submissions.Source: §A.3 · p. 13
  • The adapted fast_p is secondary to Sperf and is reported at p = 1, 1.5, and 2, with its denominator remaining all 26 tasks regardless of correctness.Source: §A.3 · pp. 13–14
  • Guidance raises the share of model–task pairs that exceed the corresponding speedup threshold at all three thresholds; at p = 1.5, model ordering matches Sperf.Source: §A.3 · p. 14, Table 5 · p. 14

A.4 Token Efficiency

The token-efficiency analysis compares guided and unguided optimization progress under the same turn budget and reports associated token and estimated-cost increases.

A.4 Token Efficiency

  • 2.31 versus 1.38: the guided trajectory reaches a higher final S′perf score than the unguided trajectory under the same 350-turn budget.Figure 4 uses a 0.1 fallback for S′perf.Source: §A.4 · p. 14, Figure 4 · p. 14
  • The smaller 0.1 fallback applies only until a task’s first eligible submission, after which its highest eligible score so far is used.Scores below 0.1 are used as measured.Source: §A.4 · p. 14
  • Guidance increases cache-read tokens by 57.6% and cache-write tokens by 29.0% across 26 tasks.The additional token usage is concentrated in cached inputs.Source: §A.4 · p. 14, Table 6 · p. 15
  • Estimated API cost rises from US$631.56 to US$820.10 (+29.9%), with cached input processing accounting for 97.3% of the increase.Output tokens rise 2.4%; token totals combine guidance text and interaction history.Source: §A.4 · p. 14

B Generation Inputs and Prompts

Appendix B documents the shared task inputs and prompts for guided and no-guidance generation, including Muon kernel design guidance and independent, task-specific GPU design proposals. It also describes the prompts for evaluating proposals against expert criteria and scoring generated implementations.

B.1 Muon Orthogonalization Kernel

  • Both configurations receive the first six shared task files; only the guided configuration receives expert design guidance, and the expert source code is withheld.The guidance records L1 algorithmic insights, L2 dataflow, L3 optimizations, and dependencies across those levels.Source: B Generation Inputs and Prompts · p. 15, §B.1 · p. 18
  • The observable operator independently normalizes each matrix in FP32, stores the iterate in FP16, and applies five distinct coefficient triples in order; it is an orthogonalization surrogate, not an exact polar factor.The steps argument must equal five, and the declared coefficient sequence cannot be shortened or repeated.Source: §B.1 · pp. 15–16
  • The registered workloads are three BF16 matrix batches: [32, 2048, 5464], [32, 4096, 11008], and [8, 7168, 18432].Each registered shape may use a specialized kernel, but matrix values remain runtime data.Source: §B.1 · p. 16
  • The timed computation must use custom Triton kernels, not an existing target-operator implementation, and every output must derive from the current invocation’s inputs.PyTorch is limited to interface glue, allocation, and shape or stride handling; caching or reusing dynamic-input-derived values is prohibited.Source: §B.1 · p. 18
  • The dataflow keeps X unchanged while square R and Q states evolve; symmetric square-state products and GEMM-add kernels handle interior phases, with rectangular GEMMs at boundaries.Independent matrices and output tiles can run in parallel, while polynomial step order and segment boundaries remain sequential.Source: §B.1 · pp. 18–19
  • The main algorithmic insight is to compose polynomial actions on the smaller Gram side and apply the accumulated action to rectangular X only at restart and final commit boundaries.This replaces repeated large rectangular products with smaller square products and fewer rectangular commits.Source: §B.1 · pp. 18–19
  • The fixed restart follows coefficients 0 and 1 and precedes coefficient 2: commit Q to X in FP16, rebuild R from committed X, and discard the old Q and R.The guidance treats restart as necessary to control numerical drift while retaining most deferred-composition savings.Source: §B.1 · pp. 18–19
  • At restart, the design rebuilds Gram state from committed X and clears old Q/R; at the final boundary, it returns the updated tensor.The restart occurs before coefficient 2, while the final rectangular application follows coefficient 4.Source: §B.1 · p. 20
  • For registered wide-matrix workloads (M<N), the design forms the smaller left Gram R=XX^T and composes polynomial steps in square R/Q state without interior X updates.Q is applied to X after coefficients 0 and 1 at the fixed restart before coefficient 2, and again after coefficient 4 at the final boundary.Source: §B.1 · pp. 20–21
  • Computing one triangular set of square output tiles and publishing off-diagonal values to both halves avoids nearly half the duplicated off-diagonal product work.Diagonal tiles publish once; a one-triangle-only store would leave later dense consumers without a complete matrix.Source: §B.1 · p. 21

B.2 Independent Design Proposal Prompt

  • Independent design runs receive no expert guidance or solution template, cannot write or execute kernel code, and submit only model_design.json through submit_design.The brief is staged as CLAUDE.md for all models; the target is NVIDIA B200 (Blackwell), sm_100, and agents read shared task inputs plus the output schema.Source: §B.2 · p. 21
  • The prompt requires task-specific causal chains linking shape-specific bottlenecks to L1-L3 changes, cost reductions, conditions, and evidence that distinguishes complete from partial adoption.It rejects generic bottleneck claims and mechanisms stated without their conditions.Source: §B.2 · pp. 21–22
  • The proposal schema asks L1 entries to explain algorithmic changes, correctness, performance, and dataflow consequences, while L2 specifies value placement, layout, ownership, lifetimes, reuse, and adoption checks.Parallelism must be described at registered shapes, including available parallelism; stages and edges record residency, handoffs, and readiness.Source: §B.2 · p. 22
  • L3 proposals must connect execution mechanisms to L1 or L2, identify challenges and straightforward-implementation pitfalls, and state what must be measured rather than assumed.The schema organizes annotations into three levels and cross-level rationale, while leaving implementation and tuning to agents.Source: §B.2 · p. 22
  • The independent proposal schema separates dataflow stages and edges, recording outputs, layout, ownership and residency, consumer lifetime, performance conditions, storage lifetime, and readiness.It also includes restrictions and adoption criteria.Source: §B.2 · pp. 24–25
  • The proposal prompt treats L3 tricks as execution mechanisms that preserve L1 and L2 benefits, while leaving exact schedules and tuning parameters open.Examples include pipelining, specialization and roles, synchronization, load balancing, and layout-transition instructions when source-backed.Source: §B.2 · p. 25

C LLM Judge Prompts

  • The two assessments use separate user-message templates, and both assign continuous L1–L3 scores whose equal-weight arithmetic mean yields JSdesign or JSimpl.The shared system message precedes each user template, and the release scorer fills braced fields before transmission.Source: C LLM Judge Prompts · p. 26

C.2 Independent Design Assessment Prompt

  • The independent design judge scores model_design.json against task-specific expert criteria without runtime measurements or expert implementation code.For the Muon example, judge_notes is empty because its rubric has no task-specific notes.Source: §C.2 · p. 28
  • The design judge scores L1, L2, and L3 separately on a continuous 0–100 scale, using five descriptive anchors as reference points rather than categorical score conversions.It evaluates required properties together within each level and credits legal alternatives that satisfy the design intent under an explicit correctness condition.Source: §C.2 · p. 28
  • Every nonzero level score requires 1–4 machine-verified exact quotes from the candidate design; the judge explains realized and missing properties and scores contradictions at the affected level.The scorer computes JS_design as the equal-weight mean of L1, L2, and L3, without a separate error deduction or cross-level cap.Source: §C.2 · p. 28

C.1 Implementation Assessment Prompt

The prompt evaluates how fully executable code realizes expert design across three levels, using a reference implementation as a concrete comparison rather than a required template. Scores require evidence from source lines and reflect the independently assessed design angles.

C.1 Implementation Assessment Prompt

  • Each level is scored continuously from 0 to 100 for how fully its required angles are realized, with all three levels weighted equally in JS_impl.The levels cover high-level algorithmic insights, dataflow design, and low-level optimization tricks; the prompt specifies no hard cap or pass/fail threshold.Source: §C.1 · p. 27
  • The assessment checks algorithmic transformations and removed work, data representation and lifetimes, and correct low-level execution mechanisms that preserve dataflow benefits.For example, the criteria include avoiding forbidden intermediates and verifying synchronization and producer-consumer scheduling.Source: §C.1 · p. 27
  • A mathematically equivalent or superior alternative can earn equal credit when executable code makes its correctness, bottleneck removal, representation, ownership, and final consumer visible.The reference is one concrete realization, not a claim of global optimality.Source: §C.1 · p. 27
  • Every level score must cite 1–4 exact executable source passages with inclusive line ranges; paraphrases, comments, or invalid ranges fail verification.The rationale must name which required angles are realized and missing.Source: §C.1 · p. 27

C.3 Muon Design and Implementation Assessment Case

The Muon case contrasts an unguided recurrence with a guided square-composition design, then assesses the proposed designs and submitted code. The guided implementation receives a higher mean implementation score, while still omitting the expert’s fixed restart.

C.3 Muon Design and Implementation Assessment Case

  • The independent proposal keeps the five-step recurrence and triangular symmetric products but omits Q composition and the fixed restart, leaving repeated rectangular updates.Its cost argument also counts a rectangular Horner multiplication absent from the supplied PyTorch reference.Source: §C.3 · p. 29
  • The unguided code retains stepwise rectangular updates, while the guided code composes polynomial actions in square Q and applies Q to the rectangular input once.The square kernels in the unguided loop exploit symmetry but leave cross-step rectangular dataflow unchanged.Source: §C.3 · p. 29
  • JSimpl is 52 for the unguided implementation and 73 for the guided implementation; guided code passed correctness checks on all three registered workloads.The guided implementation’s higher score is consistent with lower measured latency across the three workloads.Source: §C.3 · p. 30
  • The guided implementation earns L1=62, L2=70, and L3=86, but still omits the fixed restart before coefficient 2.Its scores credit square composition, evolving square states, and triangular computation with mirrored stores; the restart handoff remains absent.Source: §C.3 · p. 30

Related papers

Loading 2610.03226v1…