训练营简介 2025年昇腾CANN训练营第二季，基于CANN开源开放全场景，推出0基础入门系列、码力全开特辑、开发者案例等专题课程，助力不同阶段开发者快速提升算子开发技能。获得Ascend C算子中级认证，即可领取精美证书，完成社区任务更有机会赢取华为手机，平板、开发板等大奖。

报名链接：https://www.hiascend.com/developer/activities/cann20252#cann-camp-2502-intro

前言

在完成了 RMSNorm 和 Int8 量化后，我们的 LLaMA 算子库还缺少一个关键组件——位置编码。

RoPE (Rotary Positional Embedding) 是目前大模型的标配。它的核心思想极其优美：通过旋转角度来表示相对位置。如果两个 Token 的距离是 $k$，那么它们的向量在空间中就相差一个旋转角 $k\theta$。

但对于 Ascend C 开发者来说，RoPE 的实现有两个工程痛点：

1. 逻辑复杂：涉及向量元素的交叉计算（$x_1 \cos - x_2 \sin$）。

2. 数据排布：LLaMA 使用的 "Half-Rotate" 模式（前半段和后半段配对）需要对数据进行大跨度的搬运和重排。

本期文章，我们将深入 RoPE 的数学本质，并利用 Ascend C 的地址偏移技巧，高效实现这一逻辑。

一、 核心原理：让向量在空间中起舞

RoPE 将向量看作复平面上的点。对于向量中的每一对数值 $(x_1, x_2)$，乘以一个旋转矩阵：

$$\begin{pmatrix} x'_1 \\ x'_2 \end{pmatrix} = \begin{pmatrix} \cos \theta & -\sin \theta \\ \sin \theta & \cos \theta \end{pmatrix} \begin{pmatrix} x_1 \\ x_2 \end{pmatrix}$$

展开计算：

$$x'_1 = x_1 \cos \theta - x_2 \sin \theta$$$$x'_2 = x_1 \sin \theta + x_2 \cos \theta$$

在 LLaMA 的实现中，为了利用 SIMD 效率，通常不把相邻元素配对，而是将长度为 $D$ 的向量切分为两半：$[x_0, ..., x_{d/2-1}]$ 和 $[x_{d/2}, ..., x_{d-1}]$。 前一半 $x_i$ 与后一半 $x_{i+d/2}$ 配对。

变换后的“旋转项” $\text{Rotate}(X)$ 变为：

$$[-x_{d/2}, ..., -x_{d-1}, x_0, ..., x_{d/2-1}]$$

也就是：交换前后两半，并将前半部分（现在的）取负。

二、 算法映射：如何在 UB 中“乾坤大挪移”？

我们的目标公式是：

$$X_{out} = X \odot \text{Cos} + \text{Rotate}(X) \odot \text{Sin}$$

其中 $\odot$ 是逐元素乘法。难点在于构造 $\text{Rotate}(X)$。 假设向量长度 len，我们要把 X[0...half] 变成 -X[half...end]，把 X[half...end] 变成 X[0...half]。

Ascend C 高效实现技巧： 利用 DataCopy 或者 Vector 指令的源操作数地址偏移。如果数据都在 UB 里，我们可以直接操作指针！

三、 代码实战：Ascend C 实现 RoPE

3.1 Kernel 类定义

输入除了 $X$，还需要 Host 侧预计算好的 $\cos$ 和 $\sin$ 表。

class KernelRoPE {
public:
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR out, 
                                uint32_t totalLen, uint32_t tileLen) {
        // Init 逻辑...
        // 假设 tileLen 包含了完整的 Hidden Dim (例如 128 或 4096)
        // RoPE 通常是对 Head Dim 进行操作
    }
    // ...
};

3.2 Compute 核心逻辑

__aicore__ inline void Compute(int32_t i) {
    LocalTensor<half> xLoc = inQueueX.DeQue<half>();
    LocalTensor<half> cLoc = inQueueCos.DeQue<half>();
    LocalTensor<half> sLoc = inQueueSin.DeQue<half>();
    LocalTensor<half> outLoc = outQueueOut.AllocTensor<half>();

    // 申请临时空间
    LocalTensor<half> term1 = tmpQueue.AllocTensor<half>(); // X * Cos
    LocalTensor<half> term2 = tmpQueue.AllocTensor<half>(); // Rot(X) * Sin
    
    // 还需要一个 buffer 存 Rotate 后的 X
    // 注意：Rotate 需要交换前后半段，这在 UB 内部可以通过 copy 实现
    LocalTensor<half> xRot = tmpQueue.AllocTensor<half>();

    // ---------------------------------------------------
    // Step 1: 计算第一项 X * Cos
    // ---------------------------------------------------
    Mul(term1, xLoc, cLoc, tileLength);

    // ---------------------------------------------------
    // Step 2: 构造 Rotate(X) = [-x_tail, x_head]
    // ---------------------------------------------------
    uint32_t halfLen = tileLength / 2;

    // 技巧：利用 UB 地址偏移进行 DataCopy
    // 2.1 搬运后半段到前半段：xRot[0] = xLoc[halfLen]
    // 注意：Ascend C 的 DataCopy 支持 UB->UB (部分芯片) 或通过 Vector Move
    // 这里为了通用性，假设使用 Muls/Add 等指令的 stride 机制或者直接 Copy
    
    // 方案 A: 既然 xRot 是临时空间，我们可以分段拷贝
    // 将 xLoc 的后半段拷贝到 xRot 的前半段
    DataCopy(xRot[0], xLoc[halfLen], halfLen);
    
    // 将 xRot 的前半段取反 (* -1)
    Muls(xRot, xRot, (half)-1.0, halfLen);

    // 2.2 搬运前半段到后半段：xRot[halfLen] = xLoc[0]
    DataCopy(xRot[halfLen], xLoc[0], halfLen);

    // 此时 xRot 里已经是 [-x_tail, x_head] 了

    // ---------------------------------------------------
    // Step 3: 计算第二项 Rotate(X) * Sin
    // ---------------------------------------------------
    Mul(term2, xRot, sLoc, tileLength);

    // ---------------------------------------------------
    // Step 4: 结果相加
    // ---------------------------------------------------
    Add(outLoc, term1, term2, tileLength);

    // ... 释放内存 ...
    outQueueOut.EnQue(outLoc);
    inQueueX.FreeTensor(xLoc);
    // ...
}

3.3 进阶优化：Sin/Cos 的广播机制

在实际模型中，输入 Shape 往往是 [Batch, SeqLen, Head, Dim]。 而位置编码对于同一个 SeqLen 中的不同 Head 是共享的，或者是部分共享的。

如果 Sin/Cos 表的 Shape 是 [SeqLen, Dim]，而我们一次处理的数据包含了多个 Head。 我们需要利用 Ascend C 的 Repeat 和 Stride 功能，让 Sin/Cos 在计算时自动广播，而不是显式地用 DataCopy 去复制数据填满 UB。

// 伪代码：广播乘法
// src0 (X) 是连续的
// src1 (Cos) 需要重复使用
// Mul(dst, src0, src1, repeat, src0Stride=8, src1Stride=0, ...)
// 将 src1Stride 设为 0，即可实现同一个 Cos 向量乘遍所有 Head

四、 总结

RoPE 是大模型算子中“计算”与“搬运”结合得最紧密的一个。

1. 数学转换：理解复数旋转在实数域的投影。

2. 数据重排：利用 UB 内部的 DataCopy 或地址偏移，实现向量的高效切分与拼接。

3. 广播优化：识别数据中的共享维度，利用 Stride 机制节省内存和带宽。

掌握了 RoPE，配合之前的 RMSNorm 和 MatMul，你已经可以拼凑出一个简易版的 LlamaAttention 模块了！