CuBridge: An LLM-Based Framework for Understanding and Reconstructing High-Performance Attention Kernels

📄 arXiv: 2605.05023v1 📥 PDF

作者: Xing Ma, Yangjie Zhou, Wu Sun, Zihan Liu, Jingwen Leng, Yun Lin, Shixuan Sun, Minyi Guo, Jin Song Dong

分类: cs.LG

发布日期: 2026-05-06

备注: Accepted to ACL 2026

💡 一句话要点

CuBridge：基于LLM的高性能Attention Kernel理解与重构框架

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: Attention机制 CUDA内核 大型语言模型 代码生成 深度学习 性能优化 中间表示

📋 核心要点

  1. 现有Attention机制的CUDA实现，在灵活性和性能之间存在trade-off，专家编写的内核难以适应新的Attention变体。
  2. CuBridge通过LLM将专家编写的CUDA内核转换为中间表示，再根据用户需求生成新的优化CUDA代码，实现高性能和灵活性。
  3. 实验结果表明，CuBridge在多种Attention变体和GPU平台上，性能优于通用框架、编译器方法和之前的LLM方法。

📝 摘要（中文）

高效的Attention机制CUDA实现对于现代深度学习系统至关重要，但支持多样化和不断发展的Attention变体仍然具有挑战性。现有的框架和编译器牺牲了性能来换取灵活性，而专家编写的内核虽然效率很高，但难以适应。最近的研究探索了使用大型语言模型（LLM）生成GPU内核，但先前的研究表明，对于像Attention这样复杂的算子，其正确性不稳定且性能差距显著。我们提出了CuBridge，一个基于LLM的框架，通过结构化的提升-转移-降低工作流程来适配专家编写的Attention内核。CuBridge从专家编写的CUDA Attention内核开始，将其提升为可执行的中间表示，该表示使执行编排显式化，同时抽象出低级CUDA语法。给定用户提供的PyTorch规范，CuBridge生成并验证目标IR程序，然后通过参考引导的降低来重建优化的CUDA代码。在不同的Attention变体和GPU平台上，CuBridge始终如一地生成正确的内核，并且显著优于通用框架、基于编译器的方法和先前的基于LLM的方法。

🔬 方法详解

问题定义：论文旨在解决高性能Attention机制CUDA内核的快速适配问题。现有方法要么是通用框架和编译器，性能较差；要么是专家编写的内核，难以修改和扩展以支持新的Attention变体。这限制了Attention机制在各种深度学习模型中的应用。

核心思路：CuBridge的核心思路是利用LLM的理解和生成能力，将专家编写的CUDA内核转换为一种中间表示（IR），然后根据用户提供的PyTorch规范，生成并验证新的IR程序，最后将新的IR程序转换回优化的CUDA代码。这种方法结合了专家知识和LLM的灵活性。

技术框架：CuBridge框架包含三个主要阶段：提升（Lift）、转移（Transfer）和降低（Lower）。提升阶段将专家编写的CUDA内核转换为中间表示（IR），该IR抽象了低级CUDA语法，并显式地表示了执行编排。转移阶段利用LLM，根据用户提供的PyTorch规范，生成并验证目标IR程序。降低阶段将目标IR程序重建为优化的CUDA代码，该过程由参考引导，以确保性能。

关键创新：CuBridge的关键创新在于其结构化的lift-transfer-lower工作流程，以及利用LLM进行IR程序生成和验证。通过将专家知识编码到初始内核中，并利用LLM的生成能力，CuBridge能够快速生成高性能的Attention内核，而无需从头开始编写。

关键设计：CuBridge的中间表示（IR）的设计是关键。它需要足够抽象，以便LLM能够理解和操作，同时又需要足够具体，以便能够重建高性能的CUDA代码。此外，参考引导的降低过程也至关重要，它利用专家编写的内核作为参考，以确保生成的CUDA代码的性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

CuBridge在多种Attention变体和GPU平台上进行了评估，实验结果表明，CuBridge生成的内核始终是正确的，并且显著优于通用框架、基于编译器的方法和先前的基于LLM的方法。具体性能提升数据未知，但摘要强调了“substantially outperforms”，表明性能提升显著。

🎯 应用场景

CuBridge可应用于各种需要高性能Attention机制的深度学习任务，例如自然语言处理、计算机视觉和语音识别。它可以加速模型的训练和推理，并支持新的Attention变体，从而推动相关领域的研究和发展。该框架的实际价值在于降低了开发和维护高性能Attention内核的难度，使得研究人员可以更专注于模型设计和算法创新。

📄 摘要（原文）

Efficient CUDA implementations of attention mechanisms are critical to modern deep learning systems, yet supporting diverse and evolving attention variants remains challenging. Existing frameworks and compilers trade performance for flexibility, while expert-written kernels achieve high efficiency but are difficult to adapt. Recent work explores large language models (LLMs) for GPU kernel generation, but prior studies report unstable correctness and significant performance gaps for complex operators such as attention. We present CuBridge, an LLM-based framework that adapts expert-written attention kernels through a structured lift-transfer-lower workflow. CuBridge starts from expert-written CUDA attention kernels and lifts them into an executable intermediate representation that makes execution orchestration explicit while abstracting low-level CUDA syntax. Given a user-provided PyTorch specification, CuBridge generates and verifies a target IR program, then reconstructs optimized CUDA code via reference-guided lowering. Across diverse attention variants and GPU platforms, CuBridge consistently produces correct kernels and substantially outperforms general frameworks, compiler-based approaches, and prior LLM-based methods.