Autograd node for Rotary Position Embedding (RoPE, Su et al. 2021).
Rotates pairs of features in the last dimension of input by
position-dependent angles, encoding absolute position information
in a way that lets dot-product attention recover relative position.
During the forward, the backend returns a triple
{rotated_input, cos, sin}; the cos/sin tables are saved into
saved_cos_ and saved_sin_ so the backward can
apply the conjugate rotation without recomputing the trigonometry.
interleaved_ selects the pair layout:
true— adjacent pairs(d=0, d=1), (d=2, d=3), …false— split-half(d=0, d=D/2), (d=1, d=D/2+1), …
The chosen layout must match how the model was trained or attention scores will be silently corrupted.
Attributes
schema_v1OpSchemaRegistered schema (name
"rotary_pos_embedding", version 1, AmpPolicy::ForceFP32). Forced to F32 because the rotation math is precision-sensitive at long sequence lengths.interleaved_boolPair-layout flag (see above).
orig_shape_ShapeSaved input shape; used by backward to size the gradient.
saved_cos_StorageCosine table produced by the forward kernel.
saved_sin_StorageSine table produced by the forward kernel.
Static methods
1static
forward
→TensorImplPtrint forward(const int & input, const int & position_ids_or_null, bool interleaved)Apply RoPE to input.
Parameters
inputTensorImplPtrInput of shape with
D even. L is the sequence axis and D is rotated in pairs.position_ids_or_nullTensorImplPtrOptional integer position-index tensor of shape
(L,). When null, the kernel uses [0, 1, …, L-1] implicitly.interleavedboolPair-layout flag (see class docstring).
Returns
TensorImplPtrRotated output with the same shape as input.
Raises
ShapeMismatchIf
input is fewer than 2-D or its last dim is odd.DeviceMismatchIf
position_ids_or_null lives on a different device than input.