旋转位置嵌入的高效矩阵实现

旋转位置嵌入 (RoPE) 已成为跨语言、视觉和 3D 领域的现代 Transformer 架构的核心组件。然而，现有的实现依赖于向量级的分割和合并操作，这会引入不可忽略的计算开销，而这在注意力优化中经常被忽视。该问题在多维设置（例如 2D 和 3D RoPE）中进一步放大，其中额外的向量运算和不均匀的特征分区会降低硬件利用率。为了克服这些限制，我们提出了 RoME（旋转矩阵位置嵌入），这是一种数学上等效但计算高效的 RoPE 重新表述，用统一矩阵变换代替向量运算。 RoME 消除了特定于维度的操作，简化了实施，并支持现代 NPU 上跨立方体和向量单元的融合并行执行。实验表明，RoME 在操作员级别和全模型级别上都提供了显着的加速。该实现可在 https://gitcode.com/cann/ops-Transformer/blob/master/experimental/poseembedding/rope_matrix/README.md 获取。