用大模型自动优化注意力核函数，兼顾正确性与高性能。
CuBridge: An LLM-Based Framework for Understanding and Reconstructing High-Performance Attention Kernels
- 通过提升-迁移-降低流程，将专家代码转化为可复用中间表示
- 在多类注意力变体上生成正确代码，性能超越主流框架和前代方法
- 适合需要高效自定义注意力的深度学习研究者和工程师
高效实现注意力机制的CUDA代码对现代深度学习系统至关重要，但支持多样且不断演进的注意力变体仍具挑战。现有框架和编译器在灵活性与性能间权衡，而手工编写的核函数虽高效却难以适配。近期工作尝试用大语言模型（LLMs）生成GPU核函数，但先前研究在复杂算子如注意力上报告了稳定性差和性能差距大的问题。我们提出CuBridge，一个基于大语言模型的框架，通过结构化的提升-迁移-降低流程，适配专家编写的注意力核函数。CuBridge从专家级CUDA注意力代码出发，将其提升至可执行的中间表示（IR），显式表达执行调度并抽象底层CUDA语法。给定用户提供的PyTorch规范，CuBridge生成并验证目标IR程序，再通过参考引导的降低过程重建优化后的CUDA代码。在多种注意力变体和GPU平台上，CuBridge持续生成正确代码，并显著优于通用框架、编译器方法及先前的基于大模型的方法。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。