ops-transformer：AI 处理器上的高性能 Transformer 算子库

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

自注意力（Self-Attention）机制的引入，使得 Transformer 模型在自然语言处理（NLP）领域掀起了一场革命，并迅速扩展到计算机视觉、语音识别乃至多模态学习等诸多前沿领域。然而，Transformer 模型，尤其是其核心的注意力机制，带来了巨大的计算开销和内存需求。为了在 AI 处理器上高效运行这些模型，一套高度优化的底层算子库变得至关重要。ops-transformer 正是 CANN 软件栈中专注于为 Transformer 模型提供高性能算子的库。

ops-transformer 旨在将 Transformer 架构中的关键模块（如多头注意力、前馈网络、层归一化等）转化为在 AI 处理器上能够极致高效执行的底层算子。它深入挖掘 AI 处理器独特的并行计算能力，通过硬件亲和性设计、内存访问优化和混合精度支持，确保 Transformer 模型能够以最低的延迟和最高的吞吐量进行训练和推理。通过使用 ops-transformer，开发者可以显著加速其 Transformer 模型的运行效率，降低部署成本，并赋能更广泛的 AI 应用场景。本文将深入探讨 ops-transformer 的技术内涵、其在 AI Transformer 生态中的价值以及如何赋能高性能 AI Transformer 应用的开发。

1. 引言：Transformer 模型与 ops-transformer 的重要性

Transformer 模型已成为现代人工智能的基石，但其计算特性也对其高效执行提出了严峻挑战。ops-transformer 应运而生，旨在克服这些挑战。

1.1 Transformer 模型的崛起与计算挑战

Transformer 模型的成功源于其强大的序列建模能力和对长距离依赖的捕捉机制。然而，这种能力也伴随着显著的计算和内存成本：

自注意力机制：其核心是 QKV（Query, Key, Value）矩阵乘法和 Softmax 运算，计算复杂度与序列长度的平方成正比，这在处理长序列时会迅速成为性能瓶颈。
模型规模庞大：从 BERT、GPT 系列到最新的 LLaMA、Gemini 等，Transformer 模型的参数量已达数十亿乃至万亿级别，带来了巨大的存储和计算需求。
高并发推理需求：在实际部署中，尤其是在云服务中，需要同时处理大量的用户请求，这要求 Transformer 模型的推理具备极高的吞吐量和低延迟。

1.2 ops-transformer 的核心价值

ops-transformer 的出现，正是为了在 AI 处理器上解决这些计算挑战，其核心价值在于：

极致性能优化：将 Transformer 模型的各个子模块抽象成高度优化的算子，充分利用 AI 处理器独特的张量计算和向量计算能力，将计算效率最大化。
降低开发门槛：提供一套统一、易用的 API 接口，让开发者能够便捷地调用高性能的 Transformer 算子，而无需深入了解 AI 处理器底层的硬件架构和优化细节。
加速创新进程：通过显著缩短模型的训练和推理时间，ops-transformer 使得研究人员和开发者能够更快地迭代模型、验证想法，加速 Transformer 技术在各领域的落地。

1.3 AI 处理器上 Transformer 应用的关键使能者

ops-transformer 不仅仅是一个算子库，它是 AI 处理器在 Transformer 领域发挥最大效能的关键使能者：

充分释放 AI 算力：通过将 Transformer 的核心计算卸载到 AI 处理器执行，并进行深度优化，CPU 可以专注于其他任务，从而释放整个系统的算力潜力。
构建高效模型流水线：ops-transformer 使得在 AI 处理器上构建从数据输入到结果输出的完整、高性能 Transformer 模型流水线成为可能。
赋能广泛应用场景：无论是大规模语言模型（LLM）的训练和部署、高性能机器翻译、智能问答系统，还是多模态内容理解，ops-transformer 都为这些应用提供了坚实高效的底层支持。

2. ops-transformer 核心算子集：Transformer 架构的构建块

ops-transformer 提供了 Transformer 架构中所有关键组件的优化算子，确保模型各部分的计算都能够高效运行。

2.1 自注意力机制的关键算子

自注意力是 Transformer 的核心，ops-transformer 对其关键运算进行了深度优化：

QKV 投影与矩阵乘法：高效执行输入张量与 Query、Key、Value 权重矩阵的乘法操作，通常利用 AI 处理器上的矩阵计算单元（Matrix Unit）实现极高的吞吐量。
Scaled Dot-Product Attention ：将 QKV 矩阵乘法、缩放（除以 d k \sqrt{d_k} dk ）、掩码（Masking）以及 Softmax 运算融合，减少中间数据读写，从而降低延迟。
Dropout 与注意力融合：将注意力权重上的 Dropout 操作与前述计算进行融合，避免单独的 Kernel 调用和数据传输开销。

2.2 多头注意力与前馈网络加速

多头注意力（Multi-Head Attention）和前馈网络（Feed-Forward Network）是 Transformer Block 的另外两个重要组成部分，ops-transformer 也对它们进行了精细优化：

多头并行处理：将多个注意力头（Head）的计算以并行方式调度到 AI 处理器上，最大限度地利用多核并行能力，同时优化各个头的计算融合。
Concat 与线性投影：将多头注意力的输出结果进行拼接（Concatenation），并高效地通过最终的线性投影层，将结果映射回目标维度。
前馈网络层：对两个线性层（Linear Layer）和激活函数（通常是 GELU 或 ReLU）组成的串行计算进行优化，包括算子融合和内存访问模式优化。

2.3 位置编码与残差连接优化

除了核心的计算模块，ops-transformer 也关注模型中其他看似简单但对性能同样重要的操作：

位置编码的生成与应用：高效生成正弦/余弦位置编码，并将其融合到输入 Embedding 中，避免额外的内存拷贝和计算开销。对于相对位置编码或可学习位置编码，也提供相应的优化。
层归一化 (Layer Normalization)：针对 LayerNorm 这种常见的逐元素（Element-wise）操作，进行高度并行化和融合优化，减少其对带宽的依赖。
残差连接 (Residual Connection)：将残差连接中的加法操作与其他逐元素操作（如激活函数）进行融合，减少 Kernel 启动次数，提高数据局部性。

3. 深度优化策略：极致性能的秘密

ops-transformer 的卓越性能并非偶然，而是基于对 AI 处理器架构的深刻理解和一系列先进优化策略的结晶。

3.1 硬件亲和性设计与定制化指令

ops-transformer 的算子实现深度契合 AI 处理器独特的计算单元和指令集：

矩阵计算单元 (Matrix Unit) 的充分利用：Transformer 模型中大量的矩阵乘法操作（如 QKV 投影、自注意力得分计算、前馈网络中的线性变换）都被精心映射到 AI 处理器的矩阵计算单元上，实现超高的浮点运算吞吐量。
向量计算单元 (Vector Unit) 的高效调度：对于逐元素操作（如 Softmax、LayerNorm、激活函数）和数据搬移，ops-transformer 充分利用向量计算单元的并行能力，一次性处理多个数据。
定制化指令集与 Kernel 优化：针对 AI 处理器特有的指令集，ops-transformer 团队开发了高度优化的 Kernel，甚至可能利用底层微码优化，最大程度地压榨硬件性能。

3.2 内存访问与数据流优化

高效的内存访问是高性能计算的关键。ops-transformer 采用了多种技术来优化数据流和内存利用率：

算子融合 (Operator Fusion)：将多个连续的、内存密集型的小算子融合成一个大算子（Kernel），例如将 QKV 投影、Scaled Dot-Product Attention、Dropout 和 Add/LayerNorm 等操作融合成一个大的"多头注意力"算子。这可以显著减少中间数据的生成、显存读写次数和 Kernel 启动开销。
Tiling 与分块计算：针对大尺寸矩阵乘法和注意力计算，采用 Tiling 技术将数据分割成小块，循环地在 AI 处理器的高速片上缓存中进行计算，提高数据局部性，减少对全局显存的访问带宽需求。
内存重排与布局优化：根据 AI 处理器内存访问模式的特点，对张量的数据布局进行优化（例如 NHWC 到 NCHW 或其他定制布局），以实现更高效的内存带宽利用。

3.3 混合精度计算加速

混合精度计算已成为深度学习领域的标准实践，ops-transformer 为此提供了全面的支持：

FP16/BF16 数据类型支持：算子能够高效处理 FP16（半精度浮点数）和 BF16（Brain Float 16）数据类型，利用 AI 处理器对这些数据类型的原生加速支持，提升计算速度并减少显存占用。
精度损失控制：通过动态范围缩放（Dynamic Range Scaling）等技术，确保在降低精度的同时，最小化模型精度损失，保持模型训练和推理的稳定性。
自动精度转换与管理：提供机制自动管理不同数据类型之间的转换，并与上层深度学习框架协作，使得开发者可以便捷地开启混合精度训练和推理。

4. 易用性与集成：赋能高效模型开发

ops-transformer 旨在为开发者提供一套功能强大且易于使用的接口，并与现有 AI 软件栈无缝集成，简化 Transformer 模型的开发和部署。

4.1 统一的 API 接口与编程模型

ops-transformer 提供了清晰、一致且易于使用的 API 接口，这对于开发者来说至关重要：

模块化设计：将 Transformer 模型的各个逻辑模块（如 MultiHeadAttention、FeedForward、LayerNorm）封装成独立的算子，每个算子都拥有简洁明确的输入和输出。
参数配置灵活性：算子通常提供丰富的参数选项，允许开发者根据具体需求精细调整行为，例如指定注意力头的数量、Dropout 率、激活函数类型等。
C++ 接口：核心功能通过高性能的 C++ 接口暴露，方便底层库集成和追求极致性能的开发者直接调用。

4.2 与 CANN 软件栈的无缝对接

作为 CANN 软件栈的一部分，ops-transformer 与其他组件实现了深度集成，形成一个完整的 AI 计算生态：

与 CANN Runtime 协同：ops-transformer 算子可以直接被 CANN Runtime 调用和调度执行。这意味着在 AI 处理器上运行的深度学习模型，可以无缝地利用 ops-transformer 提供的优化能力。
设备内存共享：ops-transformer 算子直接操作 AI 处理器设备上的内存。它接收来自模型的输入张量，并在设备内存上生成输出张量，避免了主机-设备之间不必要的数据拷贝开销。
流与事件同步：ops-transformer 的操作可以被提交到 CANN Runtime 的计算流中，并利用事件机制与其他计算任务进行异步同步，实现计算与内存操作的并行。

4.3 主流深度学习框架的适配

为了方便开发者将其整合到现有项目中，ops-transformer 提供了与主流深度学习框架集成的能力：

适配框架生态：ops-transformer 算子可以通过扩展或插件的形式，被 PyTorch、TensorFlow、MindSpore 等框架调用。开发者可以通过框架提供的自定义算子机制或特定后端接口来使用这些优化算子。
统一数据表示：ops-transformer 致力于与框架内部的张量数据表示兼容，减少数据转换的开销和复杂性，例如自动处理张量维度布局的适配。
简化模型部署：在部署 Transformer 模型时，ops-transformer 可以作为模型核心组件的加速后端，从而实现更高效、更完整的端到端 AI 解决方案。

5. 性能分析与调试支持：洞察与优化

ops-transformer 不仅提供高性能算子，还配套提供工具和接口，帮助开发者深入理解 Transformer 模型在 AI 处理器上的行为，并进行性能分析和问题诊断。

5.1 细粒度性能监测与事件追踪

为了帮助开发者精确地定位性能瓶颈，ops-transformer 提供了详细的监测能力：

算子执行时间分解：能够记录每个 Transformer 算子（如 MultiHeadAttention、LayerNorm）的启动时间、完成时间以及内部各个阶段的耗时。这有助于识别哪些算子占据了大部分执行时间。
内存访问模式分析：工具可以统计每个算子的内存读写带宽、缓存命中率以及访存模式，揭示是否存在内存访问瓶颈或不合理的数据布局。
任务流事件追踪：结合 CANN Runtime 的事件追踪功能，开发者可以生成详细的时间线视图，展示不同算子在 AI 处理器上的执行顺序、重叠情况以及等待时间，从而识别并行度不足或同步开销。

5.2 瓶颈定位与优化建议

基于收集到的性能数据，ops-transformer 的工具可以帮助开发者进行瓶颈分析：

自动瓶颈识别：工具能够自动分析性能数据，识别出Transformer 模型中的潜在瓶颈，例如某个算子计算量过大、某个数据传输路径过长、或者某个 Kernel 的启动开销过高。
优化建议生成：根据瓶颈类型，工具可以给出针对性的优化建议，例如建议调整 Batch Size、序列长度、注意力头数、数据布局，或尝试使用不同的混合精度策略。
图优化分析：结合 CANN 编译器和运行时，可以分析 Transformer 模型的计算图，识别算子融合的机会，或是否存在不必要的内存拷贝。

5.3 错误报告与诊断辅助

当 Transformer 模型在 AI 处理器上运行时出现问题时，ops-transformer 提供详细的错误信息：

算子级错误码：定义了一套全面、规范的错误码体系，涵盖了算子输入参数错误、内存分配失败、Kernel 执行异常、硬件错误等各种问题。
详细错误信息与上下文：除了错误码，ops-transformer 还会提供详细的错误描述、发生错误时的设备 ID、涉及的算子名称、输入张量形状等上下文信息，极大加速了问题的诊断过程。
日志记录与调试接口：提供丰富的日志记录功能，可以将底层的运行状态、警告和错误信息记录下来，并可能提供一些调试接口，允许开发者查看算子执行前的输入数据和执行后的输出数据。

6. 面对挑战与未来展望：持续演进的 Transformer 引擎

Transformer 模型仍在快速发展，ops-transformer 也将持续演进，以适应未来的计算需求和技术挑战。

6.1 超长序列与多模态支持

Transformer 模型的应用场景不断扩展，对更复杂数据的处理能力提出了更高要求：

超长序列优化：针对自注意力机制在处理超长序列时的计算瓶颈，ops-transformer 将持续探索并支持各种高效注意力机制，如稀疏注意力、线性注意力、局部注意力等，以降低计算复杂度。
多模态数据融合：随着 Transformer 模型在多模态（如文本+图像、语音+文本）领域的应用，ops-transformer 将优化跨模态数据对齐、融合和处理的算子，支持更复杂的特征交互。
高效缓存机制：在推理阶段，特别是对于生成式模型，Attention 的 Key 和 Value 矩阵需要被缓存。ops-transformer 将提供高效的 KV Cache 管理算子，以支持更大 Batch Size 和更长生成序列。

6.2 稀疏化与量化优化

为了进一步提升 Transformer 模型的效率，稀疏化和量化是重要的方向：

模型稀疏化支持：ops-transformer 将支持结构化或非结构化稀疏模型，通过稀疏矩阵乘法等优化算子，跳过不必要的计算，加速稀疏模型的执行。
模型量化加速：提供更深度的量化支持，从 INT8 到 INT4 甚至二值化网络 (BNN)，确保量化后的模型在 AI 处理器上能够以极高的效率运行，同时尽量保持精度。
硬件协同优化：与未来 AI 处理器硬件的稀疏计算单元和高精度量化加速单元紧密结合，实现稀疏化和量化模型的原生加速。

6.3 软硬件协同的深度融合

ops-transformer 的未来发展将更加注重与 AI 处理器硬件的深度协同设计，实现软硬件一体化优化：

编译器与运行时联动：ops-transformer 将与 CANN 编译器和运行时系统紧密合作，共同实现端到端的计算图优化、数据流分析和调度，进一步减少通信和内存开销。
特定领域架构 (DSA) 适配：随着 AI 处理器不断演进，出现更多针对 Transformer 模型或其特定子结构优化的专用硬件加速模块，ops-transformer 将及时适配这些新架构，发挥其最大潜力。
动态图与即时编译：探索支持更灵活的动态计算图和即时编译技术，以适应 Transformer 模型在研究和部署中日益增长的动态特性。

附录：ops-transformer 概念性 C++ API 交互示例

以下是一个概念性的 C++ 代码片段，旨在说明一个深度学习框架或上层应用如何可能与 ops-transformer 库提供的 API 进行交互，以完成一个简化的多头注意力（Multi-Head Attention）操作。此示例着重于展示 API 的调用模式和核心概念，它并非直接可编译运行的代码，因为它省略了所有必要的头文件定义、完整的错误处理、设备初始化和上下文配置。其目的仅仅是展示如何通过 ops-transformer 抽象地调用高性能 Transformer 算子。

cpp 复制代码

#include <iostream>
#include <vector>
#include <string>
#include <memory> // For std::unique_ptr
#include <thread> // For std::this_thread::sleep_for
#include <chrono> // For std::chrono::milliseconds
#include <numeric> // For std::iota

// 概念性：AscDevKit 库，用于底层设备和内存管理
namespace AscDevKit {
    // 假设 AscDevKit 提供了 Context、Stream 等类
    // 以及 MallocDeviceMemory、FreeDeviceMemory、MemcpyHostToDeviceAsync 等函数
    // 这里只声明需要的概念性类和函数，以简化示例
    class Context { public: int device_id; std::string handle; Context(int dev_id) : device_id(dev_id), handle("Context_Dev" + std::to_string(dev_id)) {} };
    class Stream { public: std::string handle; Stream(const std::string& ctx_h) : handle("Stream_" + ctx_h) {} };
    enum Status { OK = 0, ERROR_GENERAL = 1 };

    Status InitDevice(int device_id) {
        std::cout << "[AscDevKit] 初始化设备 " << device_id << "..." << std::endl;
        std::this_thread::sleep_for(std::chrono::milliseconds(50));
        return OK;
    }

    Status CreateContext(int device_id, std::unique_ptr<Context>& context_out) {
        std::cout << "[AscDevKit] 为设备 " << device_id << " 创建上下文..." << std::endl;
        std::this_thread::sleep_for(std::chrono::milliseconds(20));
        context_out = std::make_unique<Context>(device_id);
        return OK;
    }

    Status CreateStream(const Context& context, std::unique_ptr<Stream>& stream_out) {
        std::cout << "[AscDevKit] 在上下文 " << context.handle << " 中创建流..." << std::endl;
        std::this_thread::sleep_for(std::chrono::milliseconds(15));
        stream_out = std::make_unique<Stream>(context.handle);
        return OK;
    }

    Status MallocDeviceMemory(void** dev_ptr, size_t size_bytes) {
        std::cout << "[AscDevKit] 在设备上分配 " << size_bytes << " 字节内存..." << std::endl;
        *dev_ptr = reinterpret_cast<void*>(0x30000000ULL + (rand() % 0x100000ULL)); // 概念性地址
        std::this_thread::sleep_for(std::chrono::milliseconds(5));
        return OK;
    }

    Status FreeDeviceMemory(void* dev_ptr) {
        if (!dev_ptr) return OK;
        std::cout << "[AscDevKit] 释放设备内存 " << dev_ptr << "..." << std::endl;
        std::this_thread::sleep_for(std::chrono::milliseconds(2));
        return OK;
    }

    Status MemcpyHostToDeviceAsync(void* dev_ptr, const void* host_ptr, size_t size_bytes, const Stream& stream) {
        std::cout << "[AscDevKit] 异步拷贝 " << size_bytes << " 字节从 Host " << host_ptr << " 到 Device " << dev_ptr << "..." << std::endl;
        std::this_thread::sleep_for(std::chrono::milliseconds(15));
        return OK;
    }

    Status MemcpyDeviceToHostAsync(void* host_ptr, const void* dev_ptr, size_t size_bytes, const Stream& stream) {
        std::cout << "[AscDevKit] 异步拷贝 " << size_bytes << " 字节从 Device " << dev_ptr << " 到 Host " << host_ptr << "..." << std::endl;
        std::this_thread::sleep_for(std::chrono::milliseconds(15));
        return OK;
    }

    Status StreamSynchronize(const Stream& stream) {
        std::cout << "[AscDevKit] 同步流 " << stream.handle << "..." << std::endl;
        std::this_thread::sleep_for(std::chrono::milliseconds(10));
        return OK;
    }
} // namespace AscDevKit

// 概念性：ops-transformer 库的 API 接口头文件
namespace OpsTransformer {

    // 概念性：代表一个 AI 处理器上的张量
    struct Tensor {
        void* device_ptr; // AI 处理器内存地址
        std::vector<long> dims; // 张量维度
        // ... 其他属性如数据类型、步长等
        size_t GetSizeBytes() const {
            size_t size = 1;
            for (long dim : dims) {
                size *= dim;
            }
            // 假设是 float 类型
            return size * sizeof(float);
        }
    };

    // 概念性：函数返回状态
    enum Status {
        OK = 0,
        ERROR_INVALID_INPUT,
        ERROR_UNSUPPORTED_CONFIGURATION,
        // ... 其他错误码
    };

    // --- Transformer 算子 API 示例 ---

    // 多头注意力 (Multi-Head Attention) 算子
    // query, key, value: 输入 Query, Key, Value 张量 (在设备上)
    // output: 输出张量 (在设备上)
    // head_num: 注意力头的数量
    // head_dim: 每个头的维度
    // attn_mask: 注意力掩码 (可选，在设备上)
    // dropout_prob: Dropout 概率
    // stream: 用于异步提交任务的 AI 处理器流
    Status MultiHeadAttention(const Tensor& query,
                              const Tensor& key,
                              const Tensor& value,
                              Tensor& output,
                              int head_num,
                              int head_dim,
                              const Tensor* attn_mask, // 可选
                              float dropout_prob,
                              const AscDevKit::Stream& stream) {
        std::cout << "[OpsTransformer] 在流 " << stream.handle << " 上执行 MultiHeadAttention 算子:" << std::endl;
        std::cout << "  - Query Dims: "; for(long d : query.dims) std::cout << d << " "; std::cout << std::endl;
        std::cout << "  - Key Dims: "; for(long d : key.dims) std::cout << d << " "; std::cout << std::endl;
        std::cout << "  - Value Dims: "; for(long d : value.dims) std::cout << d << " "; std::cout << std::endl;
        std::cout << "  - Head Num: " << head_num << ", Head Dim: " << head_dim << std::endl;
        std::cout << "  - Dropout Prob: " << dropout_prob << std::endl;

        // 实际会在这里调用 AI 处理器底层的 Kernel 来执行 MHA
        // 这个过程是异步的，并将任务添加到 stream 中
        std::this_thread::sleep_for(std::chrono::milliseconds(100)); // 模拟计算时间
        std::cout << "[OpsTransformer] MultiHeadAttention Kernel 已提交到流。" << std::endl;
        return OK;
    }

    // 概念性：层归一化 (Layer Normalization) 算子
    // Status LayerNorm(const Tensor& input, const Tensor& gamma, const Tensor& beta, Tensor& output, float epsilon, const AscDevKit::Stream& stream);

} // namespace OpsTransformer

int main() {
    std::cout << "--- ops-transformer 概念性 C++ API 交互演示 ---" << std::endl;

    int target_device_id = 0; // 目标 AI 处理器设备 ID
    AscDevKit::Status dev_status;
    OpsTransformer::Status transformer_status;

    // 使用智能指针管理 AscDevKit 资源
    std::unique_ptr<AscDevKit::Context> context_ptr;
    std::unique_ptr<AscDevKit::Stream> stream_ptr;

    // 定义 Transformer MHA 的输入尺寸
    const long batch_size = 2;
    const long sequence_length = 512;
    const long hidden_size = 768; // Embed_dim
    const int head_num = 12;
    const int head_dim = hidden_size / head_num; // 768 / 12 = 64

    // 准备主机端数据 (模拟 Query, Key, Value)
    // 假设 Q, K, V 都来自同一个输入，这里简化为尺寸相同
    std::vector<long> qkv_dims = {batch_size, sequence_length, hidden_size};
    size_t qkv_size_bytes = batch_size * sequence_length * hidden_size * sizeof(float);

    std::vector<float> host_q_data(qkv_size_bytes / sizeof(float));
    std::iota(host_q_data.begin(), host_q_data.end(), 0.0f); // 填充一些概念性数据

    // AI 处理器设备内存指针
    void* q_dev_ptr = nullptr;
    void* k_dev_ptr = nullptr;
    void* v_dev_ptr = nullptr;
    void* output_dev_ptr = nullptr;

    try {
        // 1. 初始化 AI 处理器设备（通过 AscDevKit）
        dev_status = AscDevKit::InitDevice(target_device_id);
        if (dev_status != AscDevKit::OK) throw std::runtime_error("AscDevKit 设备初始化失败");

        // 2. 创建设备上下文和流（通过 AscDevKit）
        dev_status = AscDevKit::CreateContext(target_device_id, context_ptr);
        if (dev_status != AscDevKit::OK || !context_ptr) throw std::runtime_error("AscDevKit 上下文创建失败");

        dev_status = AscDevKit::CreateStream(*context_ptr, stream_ptr);
        if (dev_status != AscDevKit::OK || !stream_ptr) throw std::runtime_error("AscDevKit 流创建失败");

        // 3. 在 AI 处理器上分配输入 (Q, K, V) 和输出内存
        dev_status = AscDevKit::MallocDeviceMemory(&q_dev_ptr, qkv_size_bytes);
        if (dev_status != AscDevKit::OK) throw std::runtime_error("Query 设备内存分配失败");
        dev_status = AscDevKit::MallocDeviceMemory(&k_dev_ptr, qkv_size_bytes);
        if (dev_status != AscDevKit::OK) throw std::runtime_error("Key 设备内存分配失败");
        dev_status = AscDevKit::MallocDeviceMemory(&v_dev_ptr, qkv_size_bytes);
        if (dev_status != AscDevKit::OK) throw std::runtime_error("Value 设备内存分配失败");
        dev_status = AscDevKit::MallocDeviceMemory(&output_dev_ptr, qkv_size_bytes); // 输出与输入 Query 尺寸相同
        if (dev_status != AscDevKit::OK) throw std::runtime_error("Output 设备内存分配失败");

        // 4. 将主机数据异步拷贝到 AI 处理器设备
        // 实际中 Q, K, V 可能有独立的权重投影，这里简化为直接拷贝
        dev_status = AscDevKit::MemcpyHostToDeviceAsync(q_dev_ptr, host_q_data.data(), qkv_size_bytes, *stream_ptr);
        if (dev_status != AscDevKit::OK) throw std::runtime_error("Host到Device Query拷贝失败");
        dev_status = AscDevKit::MemcpyHostToDeviceAsync(k_dev_ptr, host_q_data.data(), qkv_size_bytes, *stream_ptr);
        if (dev_status != AscDevKit::OK) throw std::runtime_error("Host到Device Key拷贝失败");
        dev_status = AscDevKit::MemcpyHostToDeviceAsync(v_dev_ptr, host_q_data.data(), qkv_size_bytes, *stream_ptr);
        if (dev_status != AscDevKit::OK) throw std::runtime_error("Host到Device Value拷贝失败");


        // 5. 准备 ops-transformer 的 Tensor 结构
        OpsTransformer::Tensor q_tensor = {q_dev_ptr, qkv_dims};
        OpsTransformer::Tensor k_tensor = {k_dev_ptr, qkv_dims};
        OpsTransformer::Tensor v_tensor = {v_dev_ptr, qkv_dims};
        OpsTransformer::Tensor output_tensor = {output_dev_ptr, qkv_dims}; // 输出张量尺寸与 Query 相同

        // 6. 调用 ops-transformer 的 MultiHeadAttention 算子
        transformer_status = OpsTransformer::MultiHeadAttention(
            q_tensor, k_tensor, v_tensor, output_tensor,
            head_num, head_dim,
            nullptr, // 概念性：不使用注意力掩码
            0.1f,    // 概念性：Dropout 概率
            *stream_ptr);
        if (transformer_status != OpsTransformer::OK) throw std::runtime_error("ops-transformer MultiHeadAttention 算子执行失败");

        // 7. 将 AI 处理器上的结果异步拷贝回主机内存
        std::vector<float> host_output_data(qkv_size_bytes / sizeof(float), 0.0f);
        dev_status = AscDevKit::MemcpyDeviceToHostAsync(
            host_output_data.data(), output_dev_ptr, qkv_size_bytes, *stream_ptr);
        if (dev_status != AscDevKit::OK) throw std::runtime_error("Device到Host拷贝失败");

        // 8. 同步流，等待所有操作完成
        dev_status = AscDevKit::StreamSynchronize(*stream_ptr);
        if (dev_status != AscDevKit::OK) throw std::runtime_error("AscDevKit 流同步失败");

        std::cout << "\nTransformer MultiHeadAttention 操作概念性完成。主机输出数据示例 (前10个元素):" << std::endl;
        for (int i = 0; i < 10 && i < host_output_data.size(); ++i) {
            std::cout << host_output_data[i] << " ";
        }
        std::cout << std::endl;

    } catch (const std::runtime_error& e) {
        std::cerr << "程序遇到错误: " << e.what() << std::endl;
    }

    // 9. 清理 AI 处理器设备内存（通过 AscDevKit）
    std::cout << "\n--- 清理 ops-transformer 及 AscDevKit 资源 ---" << std::endl;
    AscDevKit::FreeDeviceMemory(q_dev_ptr);
    AscDevKit::FreeDeviceMemory(k_dev_ptr);
    AscDevKit::FreeDeviceMemory(v_dev_ptr);
    AscDevKit::FreeDeviceMemory(output_dev_ptr);

    // 智能指针会在作用域结束时自动销毁对象
    AscDevKit::DestroyStream(stream_ptr);
    AscDevKit::DestroyContext(context_ptr);
    // AscDevKit::UninitDevice(target_device_id); // 概念性反初始化设备函数

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

    return 0;
}