Implementation:NVIDIA TransformerEngine PyTorch Ext RoPE
| Field | Value |
|---|---|
| Sources | TransformerEngine |
| Domains | Deep_Learning, PyTorch, Attention |
| Last Updated | 2026-02-07 14:00 GMT |
Overview
Implements fused Rotary Position Embedding (RoPE) forward and backward passes for SBHD, BSHD, and THD memory layouts with context parallelism support.
Description
Provides fused_rope_forward, fused_rope_backward, fused_qkv_rope_forward, and fused_qkv_rope_backward. Handles multiple QKV memory formats (SBHD, BSHD, THD) by extracting strides and dimensions accordingly and passing them to nvte_fused_rope_forward/nvte_fused_rope_backward CUDA kernels. The QKV variant applies separate frequency tensors for Q and K heads, splitting the packed QKV tensor. Supports context parallelism (cp_size/cp_rank), cumulative sequence lengths (THD format), start positions for inference, and interleaved frequency layouts.
Usage
Used by the RotaryPositionEmbedding module to apply position encoding in a single fused kernel, avoiding separate Python-level operations.
Code Reference
Source Location
- Repository
NVIDIA/TransformerEngine- File
transformer_engine/pytorch/csrc/extensions/apply_rope.cpp- Lines
- 1--301
Signature
namespace transformer_engine::pytorch {
at::Tensor fused_rope_forward(
at::Tensor input, at::Tensor freqs,
bool transpose_output_memory, int cp_size, int cp_rank,
at::Tensor cu_seqlens, at::Tensor start_positions);
at::Tensor fused_rope_backward(
at::Tensor output_grads, at::Tensor freqs,
bool transpose_output_memory, int cp_size, int cp_rank,
at::Tensor cu_seqlens, at::Tensor start_positions);
std::vector<at::Tensor> fused_qkv_rope_forward(
at::Tensor qkv, at::Tensor freqs_q, at::Tensor freqs_k,
int num_q_heads, bool transpose_output_memory,
int cp_size, int cp_rank, at::Tensor cu_seqlens);
std::vector<at::Tensor> fused_qkv_rope_backward(
at::Tensor qkv_grads, at::Tensor freqs_q, at::Tensor freqs_k,
int num_q_heads, bool transpose_output_memory,
int cp_size, int cp_rank, at::Tensor cu_seqlens);
}
Import
#include "../extensions.h"
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| input | at::Tensor |
Yes | Input tensor in SBHD, BSHD, or THD format |
| freqs | at::Tensor |
Yes | Rotary frequency tensor |
| cp_size | int |
No | Context parallelism world size (default 1) |
| cp_rank | int |
No | Context parallelism rank (default 0) |
| cu_seqlens | at::Tensor |
No | Cumulative sequence lengths for THD format |
| start_positions | at::Tensor |
No | Start positions for inference RoPE |
Outputs
| Name | Type | Description |
|---|---|---|
| output | at::Tensor |
Tensor with rotary position embedding applied |
Usage Examples
import transformer_engine_torch as tex
# Apply fused RoPE to input tensor
output = tex.fused_rope_forward(input_tensor, freqs, False, 1, 0, cu_seqlens, start_pos)