Transformer 算子加速：CANN ops-transformer 仓库，深度优化大型模型计算

池央

767人浏览 · 2026-02-10 16:41:09

池央 · 2026-02-10 16:41:09 发布

CANN 组织链接： https://atomgit.com/cann
ops-transformer 仓库链接： https://atomgit.com/cann/ops-transformer

近年来，以 Transformer 架构为基石的深度学习模型，如 BERT、GPT 系列、Llama 等，在自然语言处理、计算机视觉乃至多模态领域取得了突破性进展。这些模型通常拥有数十亿甚至上千亿的参数，其巨大的计算量和内存需求对硬件加速器提出了极高的要求。为了将这些大型模型高效地部署到 AI 处理器上，仅仅依靠通用的计算优化已远远不够，我们需要针对 Transformer 模型的特性，设计和实现高度优化的专用算子。

ops-transformer 仓库正是 CANN 软件栈为解决这一挑战而构建的。它专注于提供一系列针对 Transformer 架构中核心计算单元（如自注意力、前馈网络、层归一化等）高度优化的算子实现。这些算子充分利用了 AI 处理器的独特硬件能力，通过精密的算子融合、内存访问优化、并行调度等技术，旨在大幅提升大型 Transformer 模型的训练和推理性能。ops-transformer 不仅提供了优化的算子，更是将这些算子深度融入 CANN 编译与运行时系统，为 AI 处理器赋能大型模型的高效计算奠定了坚实的基础。本文将深入探讨 ops-transformer 的技术内涵、其在 AI 处理器生态中的价值以及如何加速前沿 AI 应用的开发。

1. 引言：大型语言模型与 Transformer 的计算挑战

Transformer 架构的崛起彻底改变了深度学习领域，但也带来了前所未有的计算和优化需求。

1.1 Transformer 架构的崛起与计算需求

Transformer 模型凭借其自注意力机制和并行计算能力，在处理序列数据方面展现出强大优势：

自注意力机制：Transformer 的核心在于自注意力，它允许模型在处理序列的每个元素时，动态地加权关注序列中的其他元素。这使得模型能够捕捉长距离依赖关系，但代价是计算复杂度通常与序列长度的平方成正比。
多头注意力与前馈网络：为了增强模型的表达能力，Transformer 引入了多头注意力（Multi-Head Attention）和多层前馈网络（Feed-Forward Networks），这些模块都包含大量的矩阵乘法、向量加法和非线性激活函数。
规模爆炸：现代 Transformer 模型，特别是大型语言模型（LLM），参数规模从数十亿到千亿不等，这导致了天文数字般的浮点运算（FLOPs）和巨大的内存带宽需求。

1.2 专用 AI 处理器加速的必要性

通用 CPU 无法满足 Transformer 模型的巨大计算需求，专用 AI 处理器成为必然选择：

高并行计算能力：AI 处理器通常内置了大规模的矩阵乘法单元（如张量计算单元）和向量计算单元，能够高效执行 Transformer 中大量的矩阵运算。
高带宽内存：为满足大型模型参数和中间激活值存储需求，AI 处理器通常配备了 HBM (高带宽内存)，以提供足够的内存带宽来避免数据饥饿。
定制化优化：AI 处理器允许更深层次的定制化和软件栈优化，可以针对 Transformer 模型的特定计算模式进行深度优化，而非仅依赖通用指令集。

1.3 ops-transformer 仓库的核心使命

ops-transformer 仓库正是为满足 Transformer 模型在 AI 处理器上的极致性能需求而生：

聚焦核心算子：它专注于 Transformer 架构中最计算密集、内存敏感的算子，提供这些算子的最优实现。
硬件亲和性设计：所有的算子都围绕 AI 处理器的架构特点进行设计，旨在最大限度地利用其计算和内存资源。
性能瓶颈突破：通过创新的融合策略、内存管理和并行调度，ops-transformer 致力于突破 Transformer 模型在 AI 处理器上的性能瓶颈。

2. ops-transformer 核心：Transformer 专用算子库

ops-transformer 仓库的核心在于其提供的 Transformer 架构中关键模块的优化实现。

2.1 关键 Transformer 模块的拆解与实现

Transformer 模型的复杂性在于其模块化结构，ops-transformer 对这些模块进行了深度优化：

多头注意力（Multi-Head Attention, MHA）：这是 Transformer 的核心，包含 Query、Key、Value 的投影（Linear Projection）、缩放点积注意力（Scaled Dot-Product Attention）和最终输出投影。ops-transformer 对这些子操作进行融合优化，减少中间数据的内存读写，例如将 QKV 投影与后续的 Attention Score 计算融合。
前馈网络（Feed-Forward Network, FFN）：通常由两个线性层和一个激活函数（如 GELU）组成。ops-transformer 会将这两个线性层、中间的激活函数进行融合，形成一个更高效的计算单元。
层归一化（Layer Normalization）：在 MHA 和 FFN 之后通常会进行层归一化。ops-transformer 常常将 LayerNorm 与其前后的线性层进行融合，减少内存访问次数，例如 “LayerNorm + Linear” 融合算子。

2.2 算子功能的全面覆盖

ops-transformer 旨在提供 Transformer 训练和推理所需的全面算子支持：

编码器与解码器核心：仓库中包含了编码器（Encoder）和解码器（Decoder）中所有必需的算子，包括自注意力、交叉注意力、前馈网络、残差连接和位置编码等。
不同场景下的变体：考虑到推理和训练场景的差异，以及不同 Transformer 模型可能存在的变体（如针对长序列的稀疏注意力），ops-transformer 可能提供不同的算子实现或优化路径。
梯度计算支持：对于训练场景，ops-transformer 也包含了相应算子的反向传播（梯度计算）实现，并同样进行了深度优化，以确保训练效率。

2.3 标准化与兼容性

为了更好地集成到 CANN 生态中，ops-transformer 遵循严格的标准化和兼容性原则：

与 metadef 协同：ops-transformer 中的每一个算子都严格按照 metadef 仓库的规范，定义其输入输出、数据类型、形状推导规则和属性，确保其能被 CANN 编译器正确理解和优化。
与主流框架兼容：虽然是底层优化算子，但它们的设计考虑了与 PyTorch、TensorFlow、MindSpore 等主流深度学习框架的兼容性，确保模型能够平滑迁移和部署。
统一的接口规范：通过遵循 CANN 的算子开发规范，ops-transformer 提供的算子能够与其他 CANN 算子无缝协作，构成高效的计算图。

3. 深度优化策略：榨取 AI 处理器算力

ops-transformer 采用了多种先进的优化技术，旨在最大限度地榨取 AI 处理器的计算潜能。

3.1 细粒度硬件映射与并行化

将 Transformer 算子高效地映射到 AI 处理器是性能优化的基石：

张量核心单元利用：Transformer 中的大量矩阵乘法（如 QKV 投影、Attention Score 计算）被映射到 AI 处理器的高性能张量核心单元上，以实现极致的矩阵计算吞吐。
向量计算单元协调：对于像 Layer Normalization、激活函数等元素级或向量级操作，则利用 AI 处理器中的向量计算单元进行高效处理。
指令级并行与流水线：在底层 Kernel 实现中，通过精心的指令调度、内存预取和流水线操作，最大限度地提高硬件单元的利用率，减少空闲时间。

3.2 内存访问优化与数据流管理

Transformer 模型对内存带宽极为敏感，ops-transformer 对此进行了深入优化：

片上缓存优化：通过对数据进行分块（Tiling）和重新组织，确保中间数据能够尽可能地驻留在 AI 处理器的高速片上缓存中，减少对带宽有限的全局显存的访问。
零拷贝与内存复用：在算子融合时，尽可能避免中间结果从寄存器或片上缓存写回到全局显存，实现数据在计算单元之间的零拷贝传递。
数据布局转换：根据 AI 处理器对特定数据布局（如 NHWC、NC1HWC0 等）的偏好，ops-transformer 在必要时自动进行数据布局转换，或在算子设计时直接采用最优的数据布局。

3.3 算子融合与计算图优化

ops-transformer 最显著的优化之一是其广泛应用的算子融合技术：

垂直算子融合：将多个连续的、数据依赖强的算子融合成一个大的 Kernel。例如，一个 Transformer 层内部的“Query投影 + Key投影 + Value投影 + Attention计算”可以被融合，显著减少中间数据的读写和 Kernel 启动开销。
水平算子融合：将模型中多个独立的、具有相同计算模式的算子融合在一起，增加并行度。例如，在多头注意力中，多个头的计算可以在一定程度上进行水平融合。
计算图重构：ops-transformer 配合 CANN 的计算图引擎 (ge)，对包含 Transformer 算子的整个计算图进行重构，识别并消除冗余计算、优化数据流。

4. 融合 CANN 生态：从定义到执行的无缝衔接

ops-transformer 并非孤立存在，它深度融入 CANN 软件栈，实现从算子定义到模型执行的端到端优化。

4.1 与 metadef 的元数据协同

算子的精确元数据定义是 ops-transformer 与 CANN 生态协作的基础：

算子原型规范：ops-transformer 中的所有算子都严格遵循 metadef 定义的算子原型规范，包括输入输出张量的名称、数据类型、形状、格式等。
形状、类型、格式推导：每个 Transformer 算子都提供了精确的形状、类型和格式推导函数，这些函数在模型编译时被 CANN 编译器调用，以正确推导计算图中的张量属性。
属性与约束：metadef 还定义了 Transformer 算子支持的属性（如 num_heads、head_dim）及其取值范围和默认值，确保算子的正确性和稳定性。

4.2 与 ge 的计算图引擎集成

ge 作为 CANN 的计算图引擎，是 ops-transformer 发挥威力的舞台：

识别与匹配：ge 能够识别计算图中由 ops-transformer 提供的优化算子。当模型中存在特定的算子序列（如 LayerNorm + MatMul + GELU），ge 会将其替换为 ops-transformer 中高度融合的算子。
全局优化决策：ge 在进行全局计算图优化时，会考虑 ops-transformer 算子的特性，例如它们对内存布局的偏好、融合的合法性等，从而做出最优的图优化决策。
Kernel 调度与资源分配：编译后的 Transformer 算子作为底层的 Kernel，由 ge 运行时调度到 AI 处理器上执行，ge 负责管理其资源分配和并行执行。

4.3 通过 pyasc 暴露 Python 接口

为了方便 Python 开发者使用这些优化算子，ops-transformer 提供了必要的 Python 接口：

Pythonic 封装：ops-transformer 的算子通过 CANN 的 Python 接口层（如 pyasc）进行封装，使得 Python 开发者能够像调用普通深度学习框架算子一样，直接调用这些优化后的 Transformer 算子。
框架桥接：深度学习框架（如 MindSpore）能够通过其后端或自定义算子机制，桥接到 ops-transformer 提供的底层优化实现，从而在 Python 层面实现透明的性能加速。
开发与调试便利：Python 接口使得 Transformer 模型的原型开发、训练和调试更加便捷，开发者可以专注于模型逻辑，而底层性能优化则由 ops-transformer 和 CANN 自动完成。

5. 性能优势与应用价值：加速 LLM 训练与推理

ops-transformer 带来的显著性能提升，对于推动大型语言模型的发展和应用具有重大价值。

5.1 大幅提升模型吞吐与降低延迟

通过精细的算子优化和融合，ops-transformer 能够实现：

训练加速：在训练大型 Transformer 模型时，ops-transformer 显著减少了每个迭代的计算时间，从而缩短了模型的收敛周期，使得研究人员可以更快地进行模型实验和迭代。
推理加速：对于在线推理服务，优化后的 Transformer 算子能够大幅降低模型的推理延迟，提高每秒处理的请求数量（吞吐量），满足高并发、低延迟的应用需求。
资源利用率提升：由于算子更好地利用了 AI 处理器的硬件资源，不仅提升了绝对性能，也使得单位硬件资源的效率更高，降低了计算成本。

5.2 赋能大型模型的高效部署

ops-transformer 的优化使得超大规模 Transformer 模型在 AI 处理器上的部署成为可能：

内存优化：通过有效的内存复用和数据布局优化，ops-transformer 减少了 Transformer 模型在 AI 处理器上的显存占用，使得更大的模型能够在有限的硬件资源上运行。
能耗优化：更高效的计算意味着在完成相同计算任务时，消耗的电能更少，这对于大型模型在数据中心和边缘侧的部署都具有重要意义。
离线模型生成：结合 CANN 的 ge 编译器，ops-transformer 优化后的模型可以被编译成高效的离线模型文件，方便在多种 AI 处理器设备上进行部署。

5.3 缩短开发周期与降低成本

通过提供高性能的 Transformer 专用算子，ops-transformer 间接促进了 AI 领域的发展：

降低开发门槛：开发者无需深入 AI 处理器底层硬件细节，即可通过框架或上层 API 使用高性能的 Transformer 算子，极大地降低了开发和优化大型模型的门槛。
加速创新：更高的训练和推理效率意味着更快的实验周期和更低的成本，从而鼓励研究人员和开发者尝试更复杂、更大规模的 Transformer 模型，加速 AI 技术的创新。
促进生态繁荣：作为 CANN 软件栈的重要组成部分，ops-transformer 为 AI 处理器生态系统提供了强大的算力支撑，吸引更多开发者和应用加入。

6. 开发者体验与定制化：灵活扩展与持续创新

ops-transformer 的设计不仅关注性能，也注重为开发者提供良好的体验和灵活的扩展能力。

6.1 提供可调优的算子参数

为了适应不同的模型和应用场景，ops-transformer 的算子通常提供了一系列可调优的参数：

精度选择：支持 FP32、FP16 甚至 INT8 等不同精度模式，开发者可以根据性能和精度需求进行选择。
策略配置：对于某些复杂的融合算子，可能会提供不同的融合策略或数据分块参数，允许开发者进行精细调优。
API 接口统一：这些可调优参数通过统一的 API 接口暴露，方便开发者在不修改底层实现的情况下，调整算子行为。

6.2 支持自定义算子接入

对于无法满足通用需求的特殊 Transformer 算子或变体，ops-transformer 支持开发者自定义：

开放式框架：ops-transformer 遵循 CANN 的算子开发规范，开发者可以基于此框架，编写针对 AI 处理器的自定义 Transformer 算子 Kernel。
元数据注册：自定义算子同样需要提供详细的元数据（通过 metadef 规范），包括形状推导、类型推导等，确保其能被 CANN 编译器正确识别和处理。
无缝集成：一旦自定义算子注册成功并提供了优化实现，它就可以像 ops-transformer 内置算子一样，被计算图引擎 ge 纳入到全局优化和调度中。

6.3 未来发展与生态共建

ops-transformer 将持续演进，与 AI 社区共同推动 Transformer 领域的发展：

支持新型 Transformer 架构：随着 Transformer 架构的不断演变（如稀疏注意力、门控注意力、线性注意力），ops-transformer 将及时更新，提供对这些新型架构中关键算子的优化支持。
自动化优化工具链：结合 CANN 的编译优化能力，未来 ops-transformer 可能集成更智能的自动化工具，能够根据模型和硬件特性，自动生成或选择最优的 Transformer 算子实现。
社区协作与开源：ops-transformer 鼓励社区开发者贡献新的优化思路、实现和测试用例，共同丰富和完善这一重要的算子库。

附录：ops-transformer 概念性 C++ 算子实现示例

以下是一个概念性的 C++ 代码片段，旨在说明 ops-transformer 仓库中一个高度融合的 Transformer 算子（例如：一个融合了 Layer Normalization 和 QKV 投影的算子）可能的定义和 Kernel 启动模式。此示例着重于展示 API 的调用模式和核心概念，它并非直接可编译运行的代码，因为它省略了所有必要的头文件定义、完整的错误处理、以及具体的内部实现细节。其目的仅仅是展示如何通过 ops-transformer 以 C++ 方式组织和调用针对 AI 处理器优化的复杂 Transformer 算子。

#include <iostream>
#include <vector>
#include <string>
#include <map>
#include <functional>
#include <chrono> // For std::this_thread::sleep_for
#include <thread> // For std::this_thread::sleep_for

// 概念性：张量描述，通常与 metadef 的定义对应
struct TensorDesc {
    std::string name;
    std::vector<long> dims; // 张量维度
    std::string dtype;      // 数据类型，如 "float32", "float16"
    std::string format;     // 数据格式，如 "NCHW", "NHWC", "FRACTAL_NZ"
    void* device_ptr = nullptr; // AI 处理器上的数据指针
    size_t GetSizeBytes() const {
        // 简单模拟，实际会根据 dtype 准确计算
        size_t size = 1;
        for (long dim : dims) {
            size *= dim;
        }
        return size * (dtype == "float16" ? 2 : 4); // 假设 float16 是2字节，float32是4字节
    }
};

// 概念性：AI 处理器设备和内存管理接口 (通常由 pyasc/CANN Runtime 提供)
namespace AiProcessorRuntime {
    void* MallocDeviceMemory(size_t size_bytes) {
        std::cout << "[AiProcessorRuntime] 在设备上分配 " << size_bytes << " 字节内存..." << std::endl;
        void* ptr = reinterpret_cast<void*>(0x30000000ULL + (rand() % 0x100000ULL)); // 概念性地址
        std::this_thread::sleep_for(std::chrono::milliseconds(1));
        return ptr;
    }
    void FreeDeviceMemory(void* dev_ptr) {
        if (dev_ptr) {
            std::cout << "[AiProcessorRuntime] 释放设备内存 " << dev_ptr << "..." << std::endl;
            std::this_thread::sleep_for(std::chrono::milliseconds(1));
        }
    }
    void MemcpyHostToDevice(void* dev_ptr, const void* host_ptr, size_t size_bytes) {
        std::cout << "[AiProcessorRuntime] 拷贝 " << size_bytes << " 字节从 Host " << host_ptr << " 到 Device " << dev_ptr << "..." << std::endl;
        std::this_thread::sleep_for(std::chrono::milliseconds(5));
    }
    void MemcpyDeviceToHost(void* host_ptr, const void* dev_ptr, size_t size_bytes) {
        std::cout << "[AiProcessorRuntime] 拷贝 " << size_bytes << " 字节从 Device " << dev_ptr << " 到 Host " << host_ptr << "..." << std::endl;
        std::this_thread::sleep_for(std::chrono::milliseconds(5));
    }
    // 概念性：AI 处理器上的 Stream 对象，用于异步执行
    class Stream {
    public:
        Stream() { std::cout << "[AiProcessorRuntime.Stream] 创建异步流." << std::endl; }
        void Synchronize() {
            std::cout << "[AiProcessorRuntime.Stream] 同步异步流，等待所有任务完成." << std::endl;
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
        }
        ~Stream() { std::cout << "[AiProcessorRuntime.Stream] 销毁异步流." << std::endl; }
    };
} // namespace AiProcessorRuntime

// 概念性：ops-transformer 库中优化算子的 C++ 接口
namespace ops_transformer {

    // 概念性：算子参数结构体，实际会非常复杂，这里只做示意
    struct FusedLayerNormQKVParam {
        int batch_size;
        int seq_len;
        int hidden_size;
        int num_heads;
        int head_dim;
        float epsilon; // for LayerNorm
        // ... 其他权重/偏置等参数的 device_ptr
    };

    // 概念性：启动融合 LayerNorm + QKV 投影 Kernel 的函数
    // 实际会涉及到大量模板参数和底层指针
    void LaunchFusedLayerNormQKVKernel(
        const FusedLayerNormQKVParam& param,
        const TensorDesc& input_tensor,         // Batch, SeqLen, HiddenSize
        const TensorDesc& gamma_tensor,         // HiddenSize
        const TensorDesc& beta_tensor,          // HiddenSize
        const TensorDesc& q_weight_tensor,      // HiddenSize, HiddenSize
        const TensorDesc& k_weight_tensor,      // HiddenSize, HiddenSize
        const TensorDesc& v_weight_tensor,      // HiddenSize, HiddenSize
        const TensorDesc& q_bias_tensor,        // HiddenSize
        const TensorDesc& k_bias_tensor,        // HiddenSize
        const TensorDesc& v_bias_tensor,        // HiddenSize
        TensorDesc& output_q_tensor,            // Batch, SeqLen, NumHeads, HeadDim
        TensorDesc& output_k_tensor,            // Batch, SeqLen, NumHeads, HeadDim
        TensorDesc& output_v_tensor,            // Batch, SeqLen, NumHeads, HeadDim
        AiProcessorRuntime::Stream& stream      // AI 处理器上的 Stream
    ) {
        std::cout << "[ops_transformer] 启动融合 LayerNorm + QKV Projection Kernel..." << std::endl;
        std::cout << "  - Batch: " << param.batch_size << ", SeqLen: " << param.seq_len << ", HiddenSize: " << param.hidden_size << std::endl;
        std::cout << "  - Input Device Ptr: " << input_tensor.device_ptr << std::endl;
        std::cout << "  - Output Q Device Ptr: " << output_q_tensor.device_ptr << std::endl;
        // 实际 Kernel 会在这里被异步调度到 AI 处理器上
        // 这里只是模拟 Kernel 执行时间
        std::this_thread::sleep_for(std::chrono::milliseconds(50));
        std::cout << "[ops_transformer] 融合 Kernel 执行完成 (在流中)." << std::endl;
    }

    // 概念性：启动融合 FFN (Linear + GELU + Linear) Kernel 的函数
    void LaunchFusedFFNKernel(
        const FusedLayerNormQKVParam& param, // 共享一些参数
        const TensorDesc& input_tensor,
        // ... 其他 FFN 权重/偏置等参数的 device_ptr
        TensorDesc& output_tensor,
        AiProcessorRuntime::Stream& stream
    ) {
        std::cout << "[ops_transformer] 启动融合 FFN (Linear+GELU+Linear) Kernel..." << std::endl;
        std::cout << "  - Input Device Ptr: " << input_tensor.device_ptr << std::endl;
        std::this_thread::sleep_for(std::chrono::milliseconds(40));
        std::cout << "[ops_transformer] 融合 FFN Kernel 执行完成 (在流中)." << std::endl;
    }

} // namespace ops_transformer


int main() {
    std::cout << "--- ops-transformer 概念性 C++ 算子实现示例 ---" << std::endl;

    // 1. 定义模型和算子参数
    ops_transformer::FusedLayerNormQKVParam params = {
        .batch_size = 4,
        .seq_len = 128,
        .hidden_size = 768,
        .num_heads = 12,
        .head_dim = 64,
        .epsilon = 1e-5f
    };

    // 2. 准备输入张量描述 (假定数据类型为 float16)
    TensorDesc input_desc = {"input_data", {params.batch_size, params.seq_len, params.hidden_size}, "float16", "NHWC"};

    // 3. 准备权重/偏置张量描述 (仅 QKV 投影相关，LayerNorm 的权重偏置也类似)
    TensorDesc gamma_desc = {"gamma", {params.hidden_size}, "float16", "ND"};
    TensorDesc beta_desc = {"beta", {params.hidden_size}, "float16", "ND"};
    TensorDesc q_weight_desc = {"q_weight", {params.hidden_size, params.hidden_size}, "float16", "FRACTAL_NZ"};
    TensorDesc k_weight_desc = {"k_weight", {params.hidden_size, params.hidden_size}, "float16", "FRACTAL_NZ"};
    TensorDesc v_weight_desc = {"v_weight", {params.hidden_size, params.hidden_size}, "float16", "FRACTAL_NZ"};
    TensorDesc q_bias_desc = {"q_bias", {params.hidden_size}, "float16", "ND"};
    TensorDesc k_bias_desc = {"k_bias", {params.hidden_size}, "float16", "ND"};
    TensorDesc v_bias_desc = {"v_bias", {params.hidden_size}, "float16", "ND"};

    // 4. 准备输出张量描述 (Q, K, V)
    TensorDesc output_q_desc = {"output_q", {params.batch_size, params.seq_len, params.num_heads, params.head_dim}, "float16", "NHWC"};
    TensorDesc output_k_desc = {"output_k", {params.batch_size, params.seq_len, params.num_heads, params.head_dim}, "float16", "NHWC"};
    TensorDesc output_v_desc = {"output_v", {params.batch_size, params.seq_len, params.num_heads, params.head_dim}, "float16", "NHWC"};

    // 5. 在 AI 处理器设备上分配内存
    input_desc.device_ptr = AiProcessorRuntime::MallocDeviceMemory(input_desc.GetSizeBytes());
    gamma_desc.device_ptr = AiProcessorRuntime::MallocDeviceMemory(gamma_desc.GetSizeBytes());
    beta_desc.device_ptr = AiProcessorRuntime::MallocDeviceMemory(beta_desc.GetSizeBytes());
    q_weight_desc.device_ptr = AiProcessorRuntime::MallocDeviceMemory(q_weight_desc.GetSizeBytes());
    k_weight_desc.device_ptr = AiProcessorRuntime::MallocDeviceMemory(k_weight_desc.GetSizeBytes());
    v_weight_desc.device_ptr = AiProcessorRuntime::MallocDeviceMemory(v_weight_desc.GetSizeBytes());
    q_bias_desc.device_ptr = AiProcessorRuntime::MallocDeviceMemory(q_bias_desc.GetSizeBytes());
    k_bias_desc.device_ptr = AiProcessorRuntime::MallocDeviceMemory(k_bias_desc.GetSizeBytes());
    v_bias_desc.device_ptr = AiProcessorRuntime::MallocDeviceMemory(v_bias_desc.GetSizeBytes());
    output_q_desc.device_ptr = AiProcessorRuntime::MallocDeviceMemory(output_q_desc.GetSizeBytes());
    output_k_desc.device_ptr = AiProcessorRuntime::MallocDeviceMemory(output_k_desc.GetSizeBytes());
    output_v_desc.device_ptr = AiProcessorRuntime::MallocDeviceMemory(output_v_desc.GetSizeBytes());

    // 6. 准备主机端数据并拷贝到设备（模拟）
    std::vector<char> host_input_data(input_desc.GetSizeBytes(), 0); // 模拟数据
    AiProcessorRuntime::MemcpyHostToDevice(input_desc.device_ptr, host_input_data.data(), host_input_data.size());
    // ... 其他权重/偏置数据拷贝

    // 7. 创建 AI 处理器上的异步流
    AiProcessorRuntime::Stream stream;

    // 8. 启动融合算子 Kernel
    ops_transformer::LaunchFusedLayerNormQKVKernel(
        params,
        input_desc,
        gamma_desc, beta_desc,
        q_weight_desc, k_weight_desc, v_weight_desc,
        q_bias_desc, k_bias_desc, v_bias_desc,
        output_q_desc, output_k_desc, output_v_desc,
        stream
    );

    // 9. 启动另一个融合算子 (例如 FFN)
    TensorDesc ffn_input_desc = output_q_desc; // 假设FFN输入是Attention的输出之一
    TensorDesc ffn_output_desc = {"ffn_output", ffn_input_desc.dims, "float16", "NHWC"};
    ffn_output_desc.device_ptr = AiProcessorRuntime::MallocDeviceMemory(ffn_output_desc.GetSizeBytes());

    ops_transformer::LaunchFusedFFNKernel(
        params, // 共享部分参数
        ffn_input_desc,
        ffn_output_desc,
        stream
    );

    // 10. 同步流，等待所有算子执行完成
    stream.Synchronize();

    // 11. 将结果从设备拷贝回主机（模拟）
    std::vector<char> host_output_q_data(output_q_desc.GetSizeBytes());
    AiProcessorRuntime::MemcpyDeviceToHost(host_output_q_data.data(), output_q_desc.device_ptr, host_output_q_data.size());
    // ... 其他输出数据拷贝

    std::cout << "\n--- 所有操作概念性完成。 ---" << std::endl;

    // 12. 释放设备内存
    AiProcessorRuntime::FreeDeviceMemory(input_desc.device_ptr);
    AiProcessorRuntime::FreeDeviceMemory(gamma_desc.device_ptr);
    AiProcessorRuntime::FreeDeviceMemory(beta_desc.device_ptr);
    AiProcessorRuntime::FreeDeviceMemory(q_weight_desc.device_ptr);
    AiProcessorRuntime::FreeDeviceMemory(k_weight_desc.device_ptr);
    AiProcessorRuntime::FreeDeviceMemory(v_weight_desc.device_ptr);
    AiProcessorRuntime::FreeDeviceMemory(q_bias_desc.device_ptr);
    AiProcessorRuntime::FreeDeviceMemory(k_bias_desc.device_ptr);
    AiProcessorRuntime::FreeDeviceMemory(v_bias_desc.device_ptr);
    AiProcessorRuntime::FreeDeviceMemory(output_q_desc.device_ptr);
    AiProcessorRuntime::FreeDeviceMemory(output_k_desc.device_ptr);
    AiProcessorRuntime::FreeDeviceMemory(output_v_desc.device_ptr);
    AiProcessorRuntime::FreeDeviceMemory(ffn_output_desc.device_ptr);


    std::cout << "\n--- ops-transformer 概念性 API 演示结束 ---" << std::endl;

    return 0;
}