CuBridge: An LLM-Based Framework for Understanding and Reconstructing High-Performance Attention Kernels
Paper · 2026 · Xing Ma, Yangjie Zhou, Wu Sun et al.
Summary
Fast CUDA kernels for new attention variants, such as custom masks, altered score functions or different normalization, are hard to build. PyTorch-style frameworks and template compilers like FlexAttention are flexible but slow. Expert libraries such as FlashAttention are fast but difficult to extend. LLMs that write kernels from scratch, or edit expert code directly, often produce incorrect or slow results.
CuBridge instead treats expert-written kernels as performance references. An LLM first lifts the source CUDA into CuIR, an executable intermediate representation. CuIR exposes execution orchestration (work assignment, synchronization, data movement) while hiding low-level syntax. Given a PyTorch specification of the target variant, the system then transforms the IR program and checks it by executing it. Finally, it lowers the program back to CUDA through reference-guided reconstruction and refinement.
On A100 and H100 GPUs, across many variants and real-model configurations, CuBridge produced correct kernels. Relative to general frameworks, template-based compilers and earlier LLM methods, its average speedups were 16.03×, 1.39× and 3.33× respectively. The work shows that a structured intermediate layer makes LLM-driven kernel adaptation both reliable and efficient.
Key points
- Motivation: when GPT-5 edits FlashAttention-3 directly for a PrefixLM mask, it can apply the mask before the asynchronous GEMM has finished.
- Ablation: tested on H100 with 96 cases (8 variants × 12 sequence lengths).
- Pass@5 is 0.38 for vanilla GPT-5, 0.58 for GPT-5 + ReAct and 1.00 for CuBridge.
- Normalized speedups are 1.00×, 1.23× and 4.19× respectively.
- Comparison with Qimeng-Attention:
- On simple relative-position attention on A100, CuBridge is about 1.24× faster.
- On complex "comb" masking it reaches 5.37× on A100 and 11.47× on H100.
- LLM backends: Claude, DeepSeek-V3 and Qwen-3-235B stay within 5% of GPT-5's throughput (e.g., ~577 TFLOPS at 4k tokens, PrefixLM, Llama2-7B config). Qwen-3-32B fails to produce valid kernels.
- Pipeline components: Lifter, IR Transfer, a CuIR Executor-based Verifier, IR Differencer, Rewriter and Refiner, with profiling feedback.
Institutions
Fields, topics and concepts
📚 Shelves 5The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 worksHardware and Architecture · 81 works
🧭 Topics 1Research subjects and languages the work is about.GPU & parallel computing · 19 works
Cite
Cite as Xing Ma, Yangjie Zhou, Wu Sun, Zihan Liu, Jingwen Leng, Yun Lin, Shixuan Sun, Minyi Guo and Jin Song Dong. 2026. CuBridge: An LLM-Based Framework for Understanding and Reconstructing High-Performance Attention Kernels. Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers).
@inproceedings{ma2026cubridge,
title = {{CuBridge: An LLM-Based Framework for Understanding and Reconstructing High-Performance Attention Kernels}},
author = {Xing Ma, Yangjie Zhou, Wu Sun, Zihan Liu, Jingwen Leng, Yun Lin, Shixuan Sun, Minyi Guo and Jin Song Dong},
booktitle = {Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)},
year = {2026},
doi = {10.18653/v1/2026.acl-long.500},
url = {https://doi.org/10.18653/v1/2026.acl-long.500}
}References
- Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, Manjunath Kudlur, Josh Levenberg, Rajat Monga, Sherry Moore, Derek Gordon Murray, Benoit Steiner, Paul A. Tucker, Vijay Vasudevan, Pete Warden, and 3 othe… — TensorFlow: A system for large-scale machine learning
- Anthropic. 2025. Claude 3.7 sonnet and claude code. find
- Carlo Baronio, Pietro Marsella, Ben Pan, Simon Guo, and Silas Alberti. 2025. Kevin: Multi-turn RL for generating CUDA kernels. CoRR, abs/2507.11948. find
- Michael Bauer, Sean Treichler, and Alex Aiken. 2014. Singe: leveraging warp specialization for high performance on gpus. In ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming, PPoPP ’14, Orlando, FL, USA, February 15-19, 2014, pages 119–130. ACM. find
Updated
This page in Markdown → https://md.research.elamur.ai/free/en/papers/ma-2026-cubridge-llm-based-framework-understanding/