[Paper] AscendCraft:通过 DSL 引导的转译编译实现自动 Ascend NPU Kernel 生成

发布: (2026年1月30日 GMT+8 17:34)
7 分钟阅读
原文: arXiv

Source: arXiv - 2601.22760v1

概述

深度学习的性能依赖于在专用硬件上运行的高度调优的内核。虽然大型语言模型(LLM)在自动生成 GPU 内核方面展现出潜力,但在华为的 Ascend NPU 上实现同样的功能仍然是一个难题,因为 Ascend 的编程模型(AscendC)不透明且文档匮乏。AscendCraft 通过引入一种轻量级领域特定语言(DSL),捕获 Ascend 内核的核心语义,然后利用 LLM 将 DSL 代码“转译”为功能完整的 AscendC 内核,从而填补了这一空白。

关键贡献

  • DSL 抽象用于 AscendC – 一种简洁、易读的语言,隐藏低层样板代码,同时暴露 Ascend NPU 独有的执行语义。
  • 两阶段生成流水线 – (1) LLM 从高级算子描述生成 DSL 代码;(2) 基于约束的 LLM 降级过程将 DSL 转换为优化后的 AscendC。
  • 高成功率 – 98.1 % 的生成内核能够编译,且 90.4 % 在多样化基准套件 (MultiKernelBench) 上通过功能正确性测试。
  • 竞争性性能 – 46.2 % 的生成内核在相同硬件上运行时间达到或超过 PyTorch 的即时执行。
  • 展示可扩展性 – 系统成功为全新 “mHC” 架构生成内核,性能优于现有的 PyTorch 实现。

方法论

  1. 设计 DSL – 作者们将 AscendC 的核心概念(张量平铺、内存层次结构、向量指令)提炼为一小套高级原语。DSL 故意保持“轻量”,以便 LLM 能够从少量专家示例中学习。
  2. 提示工程与示例选择 – 对于每个算子类别(例如卷积、矩阵乘法、逐元素操作),精选的 DSL 代码片段被作为 few‑shot 示例输入 LLM。这引导模型生成符合语法且遵循 Ascend 执行模型的 DSL 代码。
  3. 结构化转译 – 第二次 LLM 传递接收 DSL 输出以及一组形式约束(例如寄存器限制、向量宽度)。它逐步将 DSL 降低为 AscendC,插入必要的 pragma、内存拷贝和循环嵌套,并在每一步检查是否满足约束。
  4. 自动化验证 – 生成的 AscendC kernel 使用 Ascend 工具链编译,在真实 Ascend 设备上运行,并与参考输出对比,以评估功能正确性。性能则在相同输入下与 PyTorch eager 执行进行对比测量。

Source:

结果与发现

指标数值
编译成功率98.1 %
功能正确性(通过参考测试)90.4 %
内核 ≥ PyTorch eager 性能46.2 %
覆盖的算子(7 类)Convolution, GEMM, Pooling, Element‑wise, Softmax, Normalization, Reduction
新架构(mHC)生成的内核2 个正确的内核,均显著快于 PyTorch eager

这意味着:

  • DSL 有效弥合了高层算子意图与 Ascend 低层编程模型之间的“语义鸿沟”。
  • 在结构化 DSL 与约束检查的引导下,LLM 不仅能够生成语法正确的代码,还能为此前难以实现自动化生成的硬件平台产出高性能内核。
  • 该方法能够在算子族之间扩展,并且能够在无需手写内核的情况下适配新硬件扩展。

实际意义

  • 加速内核开发周期 – 团队可以在数小时内而非数周内为 Ascend NPU 原型化自定义算子,从而让资深编译器工程师专注于更高层次的优化。
  • 降低入门门槛 – 缺乏深厚 AscendC 专业知识的初创公司和研究实验室仍然可以通过 AscendCraft 的 DSL + LLM 工作流利用 Ascend 硬件的性能优势。
  • 集成到 CI 流水线 – 高编译成功率使得在持续集成过程中自动生成并测试内核成为可能,确保新模型变体始终拥有已优化的实现。
  • 可移植到其他 NPU – DSL‑first 的理念是硬件无关的;只需对 DSL 的原语进行适度调整,同一流水线即可面向其他专有加速器(如 Cambricon、Graphcore)。
  • 性价比高的性能调优 – 由于生成的内核在许多情况下已超越 eager PyTorch,开发者可以在无需手动微调的情况下实现近乎最优的推理速度,从而降低云计算成本。

限制与未来工作

  • 覆盖空白 – 虽然测试了七类算子,但更为奇特的内核(例如自定义注意力机制)仍未验证。
  • 性能上限 – 虽然 46 % 的内核匹配或超过 PyTorch eager,但仍未达到利用每个微架构细节的手工优化内核的水平。
  • 对提示质量的依赖 – 系统的成功取决于精心构造的 few‑shot 示例;扩展到全新领域可能需要额外的提示工程工作。
  • 硬件特定约束 – 当前的约束集合针对 Ascend 定制;要扩展到其他 NPU 需要新的约束模型,甚至可能需要更丰富的 DSL 构造。
  • 未来方向 包括将 DSL 扩展到覆盖控制流构造,集成基于强化学习的内核调优,并开源该流水线以促进社区驱动的扩展。

作者

  • Zhongzhen Wen
  • Shudi Shao
  • Zhong Li
  • Yu Ge
  • Tongtong Xu
  • Yuanyi Lin
  • Tian Zhang

论文信息

  • arXiv ID: 2601.22760v1
  • 分类: cs.DC, cs.LG, cs.PF, cs.SE
  • 出版日期: 2026年1月30日
  • PDF: 下载 PDF
Back to Blog

相关文章

阅读更多 »