用矩阵运算替代向量操作，让旋转位置编码更快更高效。
Efficient Matrix Implementation for Rotary Position Embedding
- 将旋转位置编码改用统一矩阵变换，避免拆分合并开销。
- 在现代NPU上实现跨立方体与向量单元的并行融合执行。
- 支持多维场景（如2D/3D），提升硬件利用率和推理速度。
旋转位置编码（RoPE）已成为语言、视觉及3D领域中现代Transformer架构的核心组件。然而，现有实现依赖于向量级的拆分与合并操作，引入显著计算开销，常被忽略于注意力优化中。该问题在多维场景（如2D和3D RoPE）中尤为突出，因额外向量操作与不均特征分割导致硬件利用率下降。为此，我们提出RoME（Rotary Matrix position Embedding），一种数学等价但计算更高效的RoPE重表述方法，以统一矩阵变换取代向量操作。RoME消除维度特异性运算，简化实现，并支持在现代NPUs上通过立方体与向量单元的融合并行执行。实验表明，RoME在算子级与全模型层面均实现显著加速。代码已开源：https://gitcode.com/cann/ops-transformer/blob/master/experimental/posembedding/rope_matrix/README.md。
原文摘要 · Abstract (English)
Rotary Position Embedding (RoPE) has become a core component of modern Transformer architectures across language, vision, and 3D domains. However, existing implementations rely on vector-level split and merge operations that introduce non-negligible computational overhead, often overlooked in attention optimization. The problem is further amplified in multi-dimensional settings (e.g., 2D and 3D RoPE), where additional vector operations and uneven feature partitions degrade hardware utilization. To overcome these limitations, we propose RoME (Rotary Matrix position Embedding), a mathematically equivalent yet computationally efficient reformulation of RoPE that replaces vector operations with unified matrix transformations. RoME eliminates dimension-specific operations, simplifies implementation, and enables fused parallel execution across Cube and Vector units on modern NPUs. Experiments show that RoME delivers substantial acceleration at both the operator and full-model levels. The implementation is available at https://gitcode.com/cann/ops-transformer/blob/master/experimental/posembedding/rope_matrix/README.md.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。