Implement Rotary Position Embedding (RoPE) in Triton. Given an input x of shape (N,D) with D even, and pre-computed cos and sin tables of shape (N,D/2), write the rotated output. For each row i and pair index j:
out[i,2j]=x[i,2j]⋅cos[i,j]−x[i,2j+1]⋅sin[i,j] out[i,2j+1]=x[i,2j]⋅sin[i,j]+x[i,2j+1]⋅cos[i,j]RoPE encodes absolute positions by rotating pairs of activation channels in 2D by a position-dependent angle. The pattern is a per-row map with no reduction: each row is independent, and within a row each pair is independent. One program per token row keeps the kernel simple and the launch overhead negligible.
Input: x = [[1.0, 2.0, 3.0, 4.0]], cos = [[1.0, 0.0]], sin = [[0.0, 1.0]]
Output: [[1.0, 2.0, -4.0, 3.0]]
Input: x = [[1.0, 0.0], [0.0, 1.0]], cos = [[1.0], [0.0]], sin = [[0.0], [1.0]]
Output: [[1.0, 0.0], [-1.0, 0.0]]
One program per row. Address x at offsets 2j and 2j+1 across the pair axis; the same column mask works for both halves.
Load cos and sin from the row of the precomputed tables and reuse them for both output halves. No transcendentals are computed inside the kernel.
Form out_even=xeven⋅c−xodd⋅s and out_odd=xeven⋅s+xodd⋅c in registers, then store each at its own offset.
torch.float32 on cuda.atol=1e-3.Accepts: array
Accepts: array
Accepts: array
Submit a correct solution to see your kernel performance.
Implement Rotary Position Embedding (RoPE) in Triton. Given an input x of shape (N,D) with D even, and pre-computed cos and sin tables of shape (N,D/2), write the rotated output. For each row i and pair index j:
out[i,2j]=x[i,2j]⋅cos[i,j]−x[i,2j+1]⋅sin[i,j] out[i,2j+1]=x[i,2j]⋅sin[i,j]+x[i,2j+1]⋅cos[i,j]RoPE encodes absolute positions by rotating pairs of activation channels in 2D by a position-dependent angle. The pattern is a per-row map with no reduction: each row is independent, and within a row each pair is independent. One program per token row keeps the kernel simple and the launch overhead negligible.
Input: x = [[1.0, 2.0, 3.0, 4.0]], cos = [[1.0, 0.0]], sin = [[0.0, 1.0]]
Output: [[1.0, 2.0, -4.0, 3.0]]
Input: x = [[1.0, 0.0], [0.0, 1.0]], cos = [[1.0], [0.0]], sin = [[0.0], [1.0]]
Output: [[1.0, 0.0], [-1.0, 0.0]]
One program per row. Address x at offsets 2j and 2j+1 across the pair axis; the same column mask works for both halves.
Load cos and sin from the row of the precomputed tables and reuse them for both output halves. No transcendentals are computed inside the kernel.
Form out_even=xeven⋅c−xodd⋅s and out_odd=xeven⋅s+xodd⋅c in registers, then store each at its own offset.
torch.float32 on cuda.atol=1e-3.Accepts: array
Accepts: array
Accepts: array
Submit a correct solution to see your kernel performance.