arXiv:2505.06302cs.LGcs.AI2025-05被引 7

用一句话提示,自动生成适配多硬件的高性能张量运算

QiMeng-TensorOp: Automatically Generating High-Performance Tensor Operators with Hardware Primitives

  • 通过指令级优化让大模型理解硬件特性,自动生成高效张量操作
  • 在RISC-V和NVIDIA GPU上分别达到OpenBLAS和cuBLAS的2.51倍和1.24倍性能
  • 开发效率提升200倍,适合需要快速部署的AI系统研发人员

计算密集型张量操作占大型语言模型(LLMs)和深度神经网络计算总量的90%以上。针对RISC-V、ARM和GPU等多样化硬件架构,自动高效地利用硬件原语生成高性能张量操作至关重要,而人工优化需数月且难以移植。大语言模型虽擅长生成高层代码,但难以充分理解硬件特性并产出高性能张量操作。本文提出仅需一句用户提示即可启动的张量操作自动生成框架QiMeng-TensorOp,使大模型能自动利用硬件特征生成基于硬件原语的张量操作,并调优参数以实现跨平台最优性能。在多种硬件平台、主流大模型及典型张量操作上的实验表明,该框架有效释放了各类硬件的计算潜力,自动生成的张量操作性能显著优于基线。相比基础大模型,最高提升达1291倍;与人类专家相比,在RISC-V CPU上超越OpenBLAS 251%,在NVIDIA GPU上达到cuBLAS的124%。此外,开发成本较人工降低200倍。

原文摘要 · Abstract (English)

Computation-intensive tensor operators constitute over 90\% of the computations in Large Language Models (LLMs) and Deep Neural Networks.Automatically and efficiently generating high-performance tensor operators with hardware primitives is crucial for diverse and ever-evolving hardware architectures like RISC-V, ARM, and GPUs, as manually optimized implementation takes at least months and lacks portability.LLMs excel at generating high-level language codes, but they struggle to fully comprehend hardware characteristics and produce high-performance tensor operators. We introduce a tensor-operator auto-generation framework with a one-line user prompt (QiMeng-TensorOp), which enables LLMs to automatically exploit hardware characteristics to generate tensor operators with hardware primitives, and tune parameters for optimal performance across diverse hardware. Experimental results on various hardware platforms, SOTA LLMs, and typical tensor operators demonstrate that QiMeng-TensorOp effectively unleashes the computing capability of various hardware platforms, and automatically generates tensor operators of superior performance. Compared with vanilla LLMs, QiMeng-TensorOp achieves up to $1291 \times$ performance improvement. Even compared with human experts, QiMeng-TensorOp could reach $251 \%$ of OpenBLAS on RISC-V CPUs, and $124 \%$ of cuBLAS on NVIDIA GPUs. Additionally, QiMeng-TensorOp also significantly reduces development costs by $200 \times$ compared with human experts.

张量运算自动代码生成硬件优化大模型应用

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。