Efficient Matrix Implementation for Rotary Position Embedding
Organizations: Huawei Technologies · Nanyang Technological University
Abstract
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.
Figures & tables
| Mode | Dimension | RoPE | RoPE+fuse | RoME | vs. RoPE | vs. RoPE+fuse |
|---|---|---|---|---|---|---|
| interleave | 2D (dim=64+64) | 5.2ms | 2.0ms | 1.6ms | 3.3x | 1.3x |
| 3D (dim=44+44+40) | 5.7ms | 3.6x | ||||
| half | 2D (dim=64+64) | 5.0ms | 2.2ms | 2.1x | 1.4x | |
| 3D (dim=44+44+40) | 5.8ms | 9.9ms | 3.6x | 6.2x |
| Task | Model | Parameters | QKShape[BNSD] | Baseline(ms) | RoME(ms) | Speedup |
| LLM | Llama3.1 | 8B | ||||
| VLM | internvl3 | 8B | ||||
| Edit | FLUX.1-Kontext | 12B | ||||
| T2I | Qwen-Image | 20B | ||||
| T2V | HunyuanVideo | 13B | ||||
| Inputs | RoPE (s) | RoME (s) | Speedup |
|---|---|---|---|
| Chunk+Cat | Matrix ( ) | Direct Rotation ( ) | |
|---|---|---|---|
| Algorithmic complexity | |||
| Memory access pattern | Non-contiguous | Contiguous | Dense matmul |
| Kernel launches | 3–4 ops | 1 fused op | 1 heavy op |
| Parallel efficiency | Low | High | Moderate |
| Implementation simplicity | Moderate | High | Low |
| Fusibility into attention | No | Yes | Difficult |
| Method | RoPE | RoME (Ours) | ||||||
| Matrix transfer | - | ✓ | - | - | ✓ | ✓ | - | ✓ |
| Input fusion | - | - | ✓ | - | ✓ | - | ✓ | ✓ |
| Operator fusion | - | - | - | ✓ | - | ✓ | ✓ | ✓ |
| Infer Cost(ms) | 6.0 | 5.7 | 4.7 | 7.7 | 2.8 | 3.7 | 2.1 | 1.6 |
| vs. RoPE | - | 1.1x | 1.3x | 0.8x | 2.1x | 1.6x | 2.8x | 3.7x |
Rotary Positional Embeddings
Rotary positional embeddings (RoPE) encode token positions by rotating attention features, supporting relative-position reasoning in language and multimodal models.
See all 98 papers