Elamur Research
☰
home / papers / 2025 / KernelBench: Can LLMs Write Efficient GPU Kernels?

KernelBench: Can LLMs Write Efficient GPU Kernels?

Written for the Research library (claude-opus-5-5) from the paper text. Original — arxiv.org/pdf/2502.10517

Summary

Hand-writing fast GPU kernels takes deep expertise and a lot of time. New ML architectures and new accelerators often wait years for well-tuned implementations; FlashAttention is a well-known example. The authors ask whether language models can take over this work.

They present KernelBench, an open-source evaluation framework with 250 PyTorch workloads. For each task, the model receives a reference Model class and must produce a ModelNew replacement. The model chooses which operations to rewrite and how, using any language or library. Outputs are checked automatically for correctness and wall-clock speed. Results are scored with fast_p, the fraction of outputs that are correct and also beat the baseline by more than a chosen speedup factor p.

Tested out of the box, even the strongest reasoning models match PyTorch Eager on fewer than one task in five. The main failures are compilation and correctness errors. Iterative refinement using execution and profiler feedback raises scores considerably. The benchmark gets harder as p increases, and gains on it carry over directly to faster real-world kernels.

Key points

  • The 250 tasks are split into three levels:
    • Level 1: 100 single operators, such as convolutions, matmuls, norms and losses.
    • Level 2: 100 fusable sequences of 3–6 operators.
    • Level 3: 50 complete architectures, such as AlexNet and MiniGPT, drawn from pytorch, huggingface/transformers and pytorch-image-models.
  • OpenAI-o1 and DeepSeek-R1 perform best without help, but still reach PyTorch Eager speed on under 20% of tasks.
  • Feeding execution and profiling results back to the model raised fast_1 by level: 12%→43%, 36%→72% and 12%→18%.
  • The authors attribute weak correctness partly to scarce CUDA training data, which makes up only 0.073% of The Stack v1.2. R1 tried tensor-core wmma instructions in about half of the Level 1 matmul tasks, but most of these attempts failed to compile.
  • Notable generated kernels include:
    • a dense-times-diagonal matrix product that scales rows instead of loading zeros (13× speedup);
    • a fused GELU kernel (2.9×);
    • a cosine similarity kernel using shared memory (2.8×).

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

Fields: AI & machine learning · Computing & GPUs

Cited by

168 papers of our library cite it (of the 140673 whose bibliographies we read).

Cite

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2502.10517 · arXiv 2502.10517

@misc{ouyang2025kernelbench,
  title = {{KernelBench: Can LLMs Write Efficient GPU Kernels?}},
  author = {Ouyang, Anne and Guo, Simon and Arora, Simran and Zhang, Alex L. and Hu, William and Re, Christopher and Mirhoseini, Azalia},
  year = {2025},
  doi = {10.48550/arxiv.2502.10517},
  eprint = {2502.10517},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2502.10517}
}

References

46 references, 7 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Apple. Apple ml compute framework (mlx), 2020. URL https://developer.apple.com/metal/. ↗
  2. Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, and Christopher Ré. Simple linear attention language models balance the recall-throughput tradeoff. International Conference on Machine Learning, 2024. find
  3. Bradley Brown, Jordan Juravsky, Ryan Ehrlich, Ronald Clark, Quoc V. Le, Christopher Ré, and Azalia Mirhoseini. Large language monkeys: Scaling inference compute with repeated sampling, 2024. URL https://arxiv.org/abs/2407.21787. ↗ our card
  4. Cerebras. Cerebras wafer-scale engine wse architecture. Online. https://cerebras.ai/product-chip/. ↗
  5. Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhai… find our card
  6. Tri Dao. FlashAttention-2: Faster attention with better parallelism and work partitioning. International Conference on Learning Representations, 2024. find
  7. Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. International Conference on Machine Learning (ICML), 2024. find
  8. Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. In Advances in Neural Information Processing Systems, 2022. find
  9. DeepSeek-AI. DeepSeek-V3. Deepseek-v3 technical report, 2025. URL https://github.com/deepseek-ai/ ↗
  10. Graphcore. Graphcore IPU architecture. Online. https://www.graphcore.ai/products/ipu. ↗
  11. Groq. Groq architecture. Online. https://groq.com/. ↗
  12. Dejan Grubisic, Chris Cummins, Volker Seeker, and Hugh Leather. Priority sampling of large language models for compilers, 2024. URL https://arxiv.org/abs/2402.18734. — [arXiv:2402.18734](https://arxiv.org/abs/2402.18734) ↗
  13. Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus), 2023. URL https://arxiv.org/ abs/1606.08415. ↗
  14. Norman P. Jouppi, George Kurian, Sheng Li, Peter Ma, Rahul Nagarajan, Lifeng Nai, Nishant Patil, Suvinay Subramanian, Andy Swing, Brian Towles, Cliff Young, Xiang Zhou, Zongwei Zhou, and David Patterson. Tpu v4: An optically reconfigurable supercomputer for machine learning with hardware support for… find
  15. Peter Kim. Flashattention minimal. flash-attention-minimal. Online, 2024. https://github.com/tspeterkim/ ↗
  16. Denis Kocetkov, Raymond Li, Loubna Ben Allal, Jia Li, Chenghao Mou, Carlos Muñoz Ferrandis, Yacine Jernite, Margaret Mitchell, Sean Hughes, Thomas Wolf, Dzmitry Bahdanau, Leandro von Werra, and Harm de Vries. The stack: 3 tb of permissively licensed source code, 2022. URL https: //arxiv.org/abs/221… find
  17. Yuhang Lai, Chengxi Li, Yiming Wang, Tianyi Zhang, Ruiqi Zhong, Luke Zettlemoyer, Scott Wen tau Yih, Daniel Fried, Sida Wang, and Tao Yu. Ds-1000: A natural and reliable benchmark for data science code generation, 2022. URL https://arxiv.org/abs/2211.11501. ↗
  18. Raymond Li, Loubna Ben Allal, Yangtian Zi, Niklas Muennighoff, Denis Kocetkov, Chenghao Mou, Marc Marone, Christopher Akiki, Jia Li, Jenny Chim, Qian Liu, Evgenii Zheltonozhskii, Terry Yue Zhuo, Thomas Wang, Olivier Dehaene, Mishig Davaadorj, Joel Lamy-Poirier, João Monteiro, Oleh Shliazhko, Nicola… find
  19. Yujia Li, David Choi, Junyoung Chung, Nate Kushman, Julian Schrittwieser, Rémi Leblond, Tom Eccles, James Keeling, Felix Gimeno, Agustin Dal Lago, Thomas Hubert, Peter Choy, Cyprien de Masson d’Autume, Igor Babuschkin, Xinyun Chen, Po-Sen Huang, Johannes Welbl, Sven Gowal, Alexey Cherepanov, James… find our card
  20. Christian J. Mills. Cuda mode notes - lecture 004. Online, 2024. https://christianjmills.com/ posts/cuda-mode-notes/lecture-004/. ↗
  21. Daniel Nichols, Pranav Polasam, Harshitha Menon, Aniruddha Marathe, Todd Gamblin, and Abhinav Bhatele. Performance-aligned llms for generating fast code, 2024. URL https://arxiv.org/abs/2404. 18864. ↗
  22. NVIDIA. cudnn: Gpu-accelerated library for deep neural networks, 2014. URL https://developer. nvidia.com/cudnn. find
  23. NVIDIA. Cuda templates for linear algebra subroutines, 2017. URL https://github.com/NVIDIA/ cutlass. ↗
  24. NVIDIA. Nvidia Tesla V100 GPU architecture, 2017. find
  25. NVIDIA. Nvidia A100 tensor core GPU architecture, 2020. find
  26. NVIDIA. Nvidia H100 tensor core GPU architecture, 2022. find
  27. NVIDIA. cuBLAS, 2023. URL https://docs.nvidia.com/cuda/cublas/. ↗
  28. Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang, Zach DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Sou… find our card
  29. Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, Xuzheng He, Haowen Hou, Przemyslaw Kazienko, Jan Kocon, and Jiaming et al. Kong. Rwkv: Reinventing rnns for the transformer era. Findings of the Association… find
  30. Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao. Flashattention-3: Fast and accurate attention with asynchrony and low-precision, 2024. URL https://arxiv.org/ abs/2407.08608. ↗
  31. Quan Shi, Michael Tang, Karthik Narasimhan, and Shunyu Yao. Can language models solve olympiad programming?, 2024. URL https://arxiv.org/abs/2404.10952. ↗
  32. Benjamin Spector, Simran Arora, Aaryan Singhal, Daniel Fu, and Christopher Ré. Thunderkittens: Simple, fast, and adorable ai kernels. International Conference on Learning Representations (ICLR), 2024. find
  33. Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. ACM Computing Surveys, 55(6):1–28, 2022. find our card
  34. Team PyTorch, Horace He, Driss Guessous, Yanbo Liang, and Joy Dong. FlexAttention: The flexibility of PyTorch with the performance of FlashAttention, 2024. URL https://pytorch.org/blog/ flexattention/. ↗
  35. Ali TehraniJamsaz, Arijit Bhattacharjee, Le Chen, Nesreen K. Ahmed, Amir Yazdanbakhsh, and Ali Jannesari. Coderosetta: Pushing the boundaries of unsupervised code translation for parallel programming. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https://… find
  36. Philippe Tillet, H. T. Kung, and David Cox. Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, 2019. find our card
  37. Alan M. Turing. On computable numbers, with an application to the Entscheidungsproblem. Proceedings of the London Mathematical Society, 2(42):230–265, 1936. URL http://www.cs.helsinki.fi/u/ gionis/cc05/OnComputableNumbers.pdf. ↗
  38. Pedro Valero-Lara, Alexis Huante, Mustafa Al Lail, William F. Godoy, Keita Teranishi, Prasanna Balaprakash, and Jeffrey S. Vetter. Comparing llama-2 and gpt-3 llms for hpc kernels generation, 2023. URL https://arxiv.org/abs/2309.07103. ↗
  39. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. 31st Conference on Neural Information Processing Systems (NIPS 2017), 2017. find
  40. Siddhant Waghjale, Vishruth Veerendranath, Zhiruo Wang, and Daniel Fried. ECCO: Can we improve model-generated code efficiency without sacrificing functional correctness? In Yaser Al-Onaizan, Mohit Bansal, and Yun-Nung Chen (eds.), Proceedings of the 2024 Conference on Empirical Methods in Natural L… find
  41. Yuanbo Wen, Qi Guo, Qiang Fu, Xiaqing Li, Jianxing Xu, Yanlin Tang, Yongwei Zhao, Xing Hu, Zidong Du, Ling Li, Chao Wang, Xuehai Zhou, and Yunji Chen. BabelTower: Learning to auto-parallelized program translation. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Siva… find
  42. Hjalmar Wijk, Tao Lin, Joel Becker, Sami Jawhar, Neev Parikh, Thomas Broadley, Lawrence Chan, Michael Chen, Josh Clymer, Jai Dhyani, Elena Ericheva, Katharyn Garcia, Brian Goodrich, Nikola Jurkovic, Megan Kinniment, Aron Lajko, Seraphina Nix, Lucas Sato, William Saunders, Maksym Taran, Ben West, and… find
  43. John Yang, Carlos E. Jimenez, Alexander Wettig, Kilian Lieret, Shunyu Yao, Karthik Narasimhan, and Ofir Press. Swe-agent: Agent-computer interfaces enable automated software engineering. arXiv:2405.15793, 2024. find our card
  44. John Yang, Carlos E. Jimenez, Alex L. Zhang, Kilian Lieret, Joyce Yang, Xindi Wu, Ori Press, Niklas Muennighoff, Gabriel Synnaeve, Karthik R. Narasimhan, Diyi Yang, Sida I. Wang, and Ofir Press. Swe-bench multimodal: Do ai systems generalize to visual software domains?, 2024. URL https://arxiv.org/a… — [arXiv:2410.03859](https://arxiv.org/abs/2410.03859) ↗
  45. Songlin Yang and Yu Zhang. Fla: A triton-based library for hardware-efficient implementations of linear attention mechanism, January 2024. URL https://github.com/sustcsonglin/ flash-linear-attention. ↗
  46. Pengcheng Yin, Wen-Ding Li, Kefan Xiao, Abhishek Rao, Yeming Wen, Kensen Shi, Joshua Howland, Paige Bailey, Michele Catasta, Henryk Michalewski, Alex Polozov, and Charles Sutton. Natural language to code generation in interactive data science notebooks, 2022. URL https://arxiv.org/abs/2212. 09248. ↗