Efficient Matrix Implementation for Rotary Position Embedding

Chen Minqi; Zhongqi Yue; Shihao Zhang; Yun Xu; Peng Wu; kaixiang Xu; Zeyi Huang; Hanwang Zhang

cs.LGApr 10, 2026

Efficient Matrix Implementation for Rotary Position Embedding

Authors: Chen Minqi, Zhongqi Yue, Shihao Zhang, Yun Xu, Peng Wu, kaixiang Xu, Zeyi Huang, Hanwang Zhang

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

Figure 1: (a) Steps to compute

\text{RoPE}(\mathbf{x},p)

given an input feature

\mathbf{x}\in\mathcal{R}^{d}

and its absolute position

p

for a 1-D sequence; (b) Steps for an

n

-D sequence, which involves an extra pair of split and merge . Note that

\sum_{i=1}^{n}d_{i}=d^{\prime}

Panel of Table 1 — Table 1: A performance profiling analysis based on Flux [ 12 ] inference on Ascend 910B3. RoPE-related computations contributes $44\%$ of the total vector operator cost, making it a primary performance limiter.

Figure 2: Illustration of our efficient matrix RoPE implementation (RoME). RoME replaces the complex and memory-intensive operations commonly used in various standard RoPE implementations, like rearrange , chunk , cat , and flatten , with a single structured rotation matrix

{\bf M}

. This unified matrix formulation substantially accelerates both small and fused operators, while also simplifying the overall implementation.

Figure 3: An example of transforming input

\mathbf{x}

\mathbf{M}_{\text{3D}}

. Sequence length and hidden dimension is 4 and 20, respectively.

Figure 4: Illustration of our Cube Vector co-parallel computation.

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
interleave	3D (dim=44+44+40)	5.7ms	2.0ms		3.6x	1.3x
half	2D (dim=64+64)	5.0ms	2.2ms		2.1x	1.4x
half	3D (dim=44+44+40)	5.8ms	9.9ms		3.6x	6.2x

Table 2: Time taken for a single operator (applying position embedding) and speedup times for RoME compared with different RoPE implementations. RoPE refers to it’s origin application and 2D/3D sequels [ 20 , 9 , 17 ] . RoPE+fuse refers to model-specific optimization implemented in [ 1 , 2 ] , where 2D and 3D operators are fused for efficient computation.

Task	Model	Parameters	QKShape[BNSD]	Baseline(ms)	RoME(ms)	Speedup
LLM	Llama3.1	8B	$[1,24,8192,128]$	$50.8$	$48.1$	$4.7\%$
VLM	internvl3	8B	$[1,28,1846,128]$	$174$	$166$	$4.5\%$
Edit	FLUX.1-Kontext	12B	$[1,24,8704,128]$	$1059$	$978$	$8.3\%$
T2I	Qwen-Image	20B	$[1,24,6032,128]$	$1457$	$1401$	$4\%$
T2V	HunyuanVideo	13B	$[1,24,42840,128]/540P$	$14410$	$14400$	$0.1\%$
	HunyuanVideo	13B	$[1,24,8192,128]/720P$	$40120$	$40100$	$0.1\%$

Table 3: Acceleration in inference by applying RoME vs. RoPE.

Inputs	RoPE (s)	RoME (s)	Speedup
$272P$	$1.27$	$1.22$	$4\%$
$480P$	$3.37$	$3.30$	$2.3\%$

Table 4: Time taken for each training step using RoPE vs RoME. Model is wan2.2/1.3B .

	Chunk+Cat	Matrix ( $\mathbf{M}$ )	Direct Rotation ( $\mathbf{R}$ )
Algorithmic complexity	$\mathcal{O}(d)$	$\mathcal{O}(d)$	$\mathcal{O}(d^{2})$
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

Table 5: Comparison of RoPE implementations. Our matrix-based approach achieves the best trade-off between efficiency and simplicity, reducing vector-operator cost that accounts for over 50% of the forward-time MFU bottleneck.

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

Table 6: Ablation study for each modules in RoME. Using none of them corresponds to standard RoPE.

Figure A5 : An example of RoPE-half.

Figure A6 : An example of RoPE-interleave.

Figure A7 : An example of RoPE-interleave-half.

Figure A8 : An example of RoPE-quarter.

Figure A9 : An example of RoPE-3D algorithm.

Figure A10 : An example of RoPE-3D matrix.

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

Higher-Dimensional Rotary Position Embedding
Aug 30, 2026Yixing Li, Ruobing Xie, Yudong Zhang +3Rotary Positional Embeddings Transformer Attention
The Chinese University of Hong Kong · Tencent Hunyuan · Tsinghua University +1
nD-RoPE: A Generalized RoPE for n-Dimensional Position Embedding
Jun 10, 2026Boyang Li, Yulin Wu, Sizhe Xu +5Rotary Positional Embeddings Positional Encoding
New York University, New York, NY, USA · Peking University, Beijing, China · Imperial College London, London, UK +1
Fractional Rotation, Full Potential? Investigating Performance and Convergence of Partial RoPE
Mar 12, 2026Mohammad Aflah Khan, Krishna P. Gummadi, Manish Gupta +1Rotary Positional Embeddings Positional Encoding
Max Planck Institute for Software Systems · Microsoft, Hyderabad

Aug 30, 2026J/K move · Enter open · S save

Higher-Dimensional Rotary Position Embedding

Yixing Li, Ruobing Xie, Yudong Zhang, Yushi Bai+2

The Chinese University of Hong Kong · Tencent Hunyuan · Tsinghua University+1

Rotary Positional Embeddings Transformer Attention Positional Encoding

Open paper AskPDF arXiv

Transformers rely on position embedding mechanisms in long context modeling in most cases. Rotary Position Embedding (RoPE) embeds positional information with independent 2D rotations, forming relative position terms in self-attention. However, its pairwise, block-based, and decoupled structure limits deep mixing and robustness across channels. We propose HD-RoPE, which extends RoPE from independent 2D rotations to higher-dimensional rotations and introduces a Paley-I orthogonal basis to obtain balanced, isotropic, and dense phase mixing within each rotation subspace. This significantly enhances channel coupling and rotational degrees of freedom while maintaining orthogonal stability and the relative position closure property. Furthermore, HD-RoPE is easily optimized for engineering efficiency without introducing additional trainable parameters. We have conducted extensive evaluation results demonstrating that HD-RoPE achieves significant performance improvements over standard RoPE across various popular benchmarks and in both long and short contexts.

Abstract

Rotary Positional Embeddings

Explore similar work

Higher-Dimensional Rotary Position Embedding

nD-RoPE: A Generalized RoPE for n-Dimensional Position Embedding

Fractional Rotation, Full Potential? Investigating Performance and Convergence of Partial RoPE