Implementation:NVIDIA TransformerEngine Fused RoPE C API

Field	Value
Sources	TransformerEngine
Domains	Deep_Learning, Optimization
Last Updated	2026-02-07 14:00 GMT

Overview

Declares the C API for applying Rotary Position Embedding (RoPE) to input tensors, supporting both individual tensor and combined QKV tensor variants with forward and backward passes.

Description

fused_rope.h provides four extern "C" functions:

• nvte_fused_rope_forward / nvte_fused_rope_backward: Apply RoPE to a single tensor with configurable strides, supporting SBHD, BSHD, and THD QKV formats, interleaved RoPE mode, context parallelism (cp_size/cp_rank), and custom start positions.
• nvte_fused_qkv_rope_forward / nvte_fused_qkv_rope_backward: Handle combined QKV inputs with separate frequency tensors for Q and K, applying RoPE to Q and K while passing V through unchanged.

Key features:

• Multiple QKV format support (SBHD, BSHD, THD)
• Interleaved rotary position embedding mode
• Context parallelism with cp_size/cp_rank parameters
• Custom start_positions for offset-based embedding

Usage

Use for applying rotary position embeddings in the attention computation pipeline. The fused kernel avoids extra memory reads/writes compared to separate operations.

Code Reference

Source Location

Repository

NVIDIA/TransformerEngine

File

transformer_engine/common/include/transformer_engine/fused_rope.h

Lines

1--146

Signature

void nvte_fused_rope_forward(
    const NVTETensor input, const NVTETensor cu_seqlens,
    const NVTETensor freqs, const NVTETensor start_positions,
    NVTETensor output, const NVTE_QKV_Format qkv_format,
    const bool interleaved, const int cp_size, const int cp_rank,
    const int s, const int b, const int h, const int d, const int d2,
    const int stride_s_or_t, const int stride_b,
    const int stride_h, const int stride_d, cudaStream_t stream);

void nvte_fused_rope_backward(...);
void nvte_fused_qkv_rope_forward(...);
void nvte_fused_qkv_rope_backward(...);

Import

#include "transformer_engine/fused_rope.h"

I/O Contract

Inputs

Name	Type	Required	Description
input	NVTETensor	Yes	Input tensor for RoPE
freqs	NVTETensor	Yes	Frequency tensor for rotary embedding
cu_seqlens	NVTETensor	Yes (THD)	Cumulative sequence lengths (for THD format)
qkv_format	NVTE_QKV_Format	Yes	QKV memory format
interleaved	bool	Yes	Whether to use interleaved RoPE

Outputs

Name	Type	Description
output	NVTETensor	Tensor with RoPE applied

Usage Examples

#include "transformer_engine/fused_rope.h"

// Apply RoPE to query tensor
nvte_fused_rope_forward(q_input, cu_seqlens, freqs, start_positions,
                        q_output, NVTE_SBHD, /*interleaved=*/false,
                        cp_size, cp_rank, s, b, h, d, d2,
                        stride_s, stride_b, stride_h, stride_d, stream);

Related Pages

• Environment:NVIDIA_TransformerEngine_CUDA_Toolkit_Requirements