RoPE

本页暂无中文版。以下为英文原文。

5 ops, 23 workloads.

One table per op, one row per workload. Ratio is the fastest other implementation's device time divided by ours, so green is faster than it, plain is level with it, red is slower. Times are in ms. How these numbers are taken.

RoPEFwd

x: [S, D]

• neox-1d-2kx: [2048, 64]dtype=f16, f32
• neox-1d-4kx: [4096, 128]dtype=bf16
• interleaved-1d-2kx: [2048, 64]rope_layout=interleaveddtype=f16
• interleaved-1d-5kx: [5120, 128]rope_layout=interleaveddtype=f32

x: [B, S, H, D]

input_layout=2d

• neox-2d-2kx: [2, 2048, 32, 128]dtype=f16
• interleaved-2d-2kx: [2, 2048, 32, 128]rope_layout=interleaveddtype=bf16

Workload	dtype	Ratio	Device time	Alternatives		Throughput	SOL	Bound
		alt / ours	ms	name	ms	TFLOP/s	of ceiling	by
neox-1d-2k	f16	1.09×	0.00186	torch-compile flashinfer torch-ref	0.00202 0.00279 0.0256	0.212	6%	lat
	f32	1.36×	0.00205	torch-compile torch-ref	0.00278 0.00966	0.192	11%	lat
neox-1d-4k	bf16	1.16×	0.00269	torch-compile flashinfer torch-ref	0.00311 0.00422 0.0321	0.585	17%	lat
interleaved-1d-2k	f16	1.26×	0.00181	torch-compile flashinfer torch-ref	0.00227 0.00272 0.0276	0.217	6%	lat
interleaved-1d-5k	f32	1.02×	0.00368	torch-compile torch-ref	0.00374 0.0209	0.534	32%	lat
neox-2d-2k	f16	1.05×	0.0177	torch-compile flashinfer torch-ref	0.0186 0.0252 0.4525	2.84	84%	mem
interleaved-2d-2k	bf16	1.16×	0.0185	flashinfer torch-compile torch-ref	0.0214 0.0442 0.3824	2.72	80%	mem

RoPENeoxPositionIdsFwd

x: [N, H, D]position_ids: [N], i32

• prefill-2kx: [2048, 32, 128]position_ids: [2048], i32seq_lens=[2048]max_position=4096dtype=f16/i32
• prefill-4kx: [4096, 32, 128]position_ids: [4096], i32seq_lens=[4096]max_position=8192dtype=bf16/i32
• qwen3-8b-decode-b32x: [32, 32, 128]position_ids: [32], i32seq_lens=repeat(1, 32)max_position=32768dtype=bf16/i32

x: [N, H, D]position_ids: [N], i64

• prefill-varlenx: [5120, 32, 128]position_ids: [5120], i64seq_lens=[1000,3000,1120]max_position=4096dtype=f32/i64

Workload	dtype	Ratio	Device time	Alternatives		Throughput	SOL	Bound
		alt / ours	ms	name	ms	TFLOP/s	of ceiling	by
prefill-2k	f16/i32	1.10×	0.0108	torch-compile vllm torch-ref	0.0120 0.0248 0.2304	2.32	55%	mem
prefill-4k	bf16/i32	1.17×	0.0192	torch-compile vllm torch-ref	0.0225 0.0456 0.4371	2.63	68%	mem
qwen3-8b-decode-b32	bf16/i32	1.07×	0.00182	torch-compile vllm torch-ref	0.00195 0.00426 0.0303	0.216	3%	lat
prefill-varlen	f32/i64	1.06×	0.0445	torch-compile vllm torch-ref	0.0471 0.0567 0.3516	1.41	79%	mem

RoPELlama31Fwd

x: [S, D]

• 1d-8kx: [8192, 128]dtype=bf16
• 1d-5kx: [5120, 128]dtype=f32

x: [B, S, H, D]

• 2d-8kx: [1, 8192, 32, 128]input_layout=2ddtype=f16

Workload	dtype	Ratio	Device time	Alternatives		Throughput	SOL	Bound
		alt / ours	ms	name	ms	TFLOP/s	of ceiling	by
1d-8k	bf16	1.28×	0.00352	torch-compile flashinfer torch-ref	0.00451 0.00618 0.0403	0.894	26%	lat
1d-5k	f32	1.14×	0.0041	torch-compile torch-ref	0.00467 0.0155	0.48	28%	lat
2d-8k	f16	1.07×	0.0330	torch-compile flashinfer torch-ref	0.0355 0.0482 0.8508	3.05	90%	mem

YaRNFwd

x: [S, D]

• 1d-8kx: [8192, 128]dtype=bf16
• 1d-5kx: [5120, 128]dtype=f32

x: [B, S, H, D]

• 2d-8kx: [1, 8192, 32, 128]input_layout=2ddtype=f16

Workload	dtype	Ratio	Device time	Alternatives		Throughput	SOL	Bound
		alt / ours	ms	name	ms	TFLOP/s	of ceiling	by
1d-8k	bf16	1.29×	0.00352	torch-compile flashinfer torch-ref	0.00454 0.00618 0.0402	0.894	26%	lat
1d-5k	f32	1.15×	0.00413	torch-compile torch-ref	0.00474 0.0155	0.476	28%	lat
2d-8k	f16	1.07×	0.0331	torch-compile flashinfer torch-ref	0.0355 0.0481 0.8508	3.05	90%	mem

LongRoPEFwd

x: [S, D]

• 1d-8kx: [8192, 128]dtype=bf16
• 1d-5kx: [5120, 128]dtype=f32
• 1d-8k-rescaledx: [8192, 128]dtype=bf16/f32, bf16/f16, bf16/bf16

x: [B, S, H, D]

• 2d-8kx: [1, 8192, 32, 128]input_layout=2ddtype=f16

Workload	dtype	Ratio	Device time	Alternatives		Throughput	SOL	Bound
		alt / ours	ms	name	ms	TFLOP/s	of ceiling	by
1d-8k	bf16	1.29×	0.00352	torch-compile flashinfer torch-ref	0.00454 0.00618 0.0401	0.894	26%	lat
1d-5k	f32	1.16×	0.0041	torch-compile torch-ref	0.00474 0.0154	0.48	28%	lat
1d-8k-rescaled	bf16/f32	1.28×	0.00355	torch-compile flashinfer torch-ref	0.00454 0.00618 0.0402	0.886	26%	lat
	bf16/f16	1.30×	0.00349	torch-compile flashinfer torch-ref	0.00454 0.00618 0.0401	0.902	27%	lat
	bf16/bf16	1.30×	0.00349	torch-compile flashinfer torch-ref	0.00454 0.00618 0.0402	0.902	27%	lat
2d-8k	f16	1.07×	0.0331	torch-compile flashinfer torch-ref	0.0354 0.0480 0.8508	3.05	90%	mem