LLM kernels 2: RoPE, GELU and elementwise fusion
Two files, both implementing rotary position embedding, both matching the paper. Give them the same tensor and the same tables and almost every output element differs. One pairs dimension 0 with dimension 1; the other pairs dimension 0 with dimension 64. Neither is a typo. The paper rotates adjacent pairs, a Llama-style implementation rotates halves, and a kernel written against one and dropped into a model trained with the other produces plausible garbage that no elementwise test catches.
The choice also decides what the hardware does: one pairing gives every lane two loads two floats apart, the other gives it two contiguous runs. This page writes both, prices them against a copy, then fuses bias, GELU and the residual add into one pass and asks whether the win tracks the bytes or the kernel count.
What RoPE does to a head, and to an address
Rotary position embedding turns a position into a rotation. Inside one head
of width D, the D numbers are read as D/2 two-dimensional vectors, and
the vector at pair index j is rotated by angle pos * theta_j, with
theta_j = 10000^(-2j/D). That is the block-diagonal matrix of the RoFormer
paper's equation 15 (https://arxiv.org/abs/2104.09864 , checked 2026-09-01),
and each 2 by 2 block is the rotation everyone writes the same way:
out0 = x0 * cos - x1 * sin
out1 = x0 * sin + x1 * cos
Two multiplies and an add each, so four FMA-shaped operations for four floats of traffic. That ratio is the whole performance story: RoPE moves two floats in and two out per pair and does almost nothing to them, so it is memory bound by a wide margin and the only thing worth optimising is how the addresses fall.
The angles come from a cos and a sin table of seq rows by D/2
columns, and both are indexed by position and pair, never by batch or by
head. Every sequence in the batch and every head in the layer reads the same
table entry. At a sequence length of 2048 and a head width of 128 that is
512 KiB per table, so after the first few blocks those reads are
L2 hits, not
global memory traffic, which is why the byte
model in the program leaves them out and says so.
The pairing is where the addresses part company. Give one thread one pair
and let lane L of a warp own pair p:
// adjacent pairs: two loads, two floats apart
const float x0 = in[2 * p];
const float x1 = in[2 * p + 1];
// half split: two loads, each one a coalesced run of 32
const float x0 = in[headBase + j];
const float x1 = in[headBase + j + half];
Under adjacent pairing the warp's first load reaches 32 addresses spread across 256 bytes. It asks for eight 32-byte sectors and uses half of each, which is day 11's stride-2 row. The second load asks for those same eight sectors and uses the other half. So the DRAM bytes are identical to a coalesced read and the cost is the extra request, not extra traffic. Under half-split pairing each load is already a contiguous run of 32 floats, fully coalesced, and the question does not arise.
Diagram: one warp, three ways to fetch a pair. Three bands. Each band has 32 lane boxes on the left and a strip of 32-byte sectors on the right, with one arrow per load request. Band 1, adjacent pairs, two scalar loads: the first load's 32 arrows land on eight sectors using the even floats, the second load's 32 arrows land on the same eight sectors using the odd floats. Caption "2 requests, 8 sectors, every byte used." Band 2, adjacent pairs, one
float2load: 32 arrows, eight sectors. Caption "1 request, the same 8 sectors." Band 3, half split, two scalar loads: 32 arrows onto four sectors low in the head and 32 onto four sectors half a head away. Caption "2 requests, 8 sectors, two runs." Alt text: "Three ways one warp fetches a rotary pair. All three touch eight thirty-two-byte sectors and waste nothing. Only the number of requests changes, from two down to one."
The two GELUs are two functions
GELU is x times the standard normal CDF, and it ships two ways. PyTorch's
approximate='none' computes it with erf; approximate='tanh' computes
0.5x(1 + tanh(sqrt(2/pi)(x + 0.044715x^3))), and the default is 'none'
(https://docs.pytorch.org/docs/stable/generated/torch.nn.GELU.html , checked
2026-09-01). GPT-2 shipped the tanh form and plenty of models still do, so
"use GELU" does not name a function.
Readers arrive expecting the tanh version to be the same function with a
cheaper implementation, the way __expf is a cheaper expf. It is not.
Day 47 priced that kind of substitution,
where a fast intrinsic differs from its accurate twin by a handful of ulps,
and erff and tanhf are each documented at a maximum error of 2 ULP (CUDA
Programming Guide, mathematical functions appendix,
https://docs.nvidia.com/cuda/cuda-programming-guide/05-appendices/mathematical-functions.html
, checked 2026-09-01). The gap between the two GELU forms is a curve-fitting
error instead: it survives in double precision and does not shrink when you
compute it more carefully. The program measures it over a sweep of x and
reports where it is worst.
// The exact GELU: x times the standard normal CDF, which is what PyTorch
// computes for approximate='none'.
__device__ __forceinline__ float geluErf(float x) {
return 0.5f * x * (1.0f + erff(x * kInvSqrt2));
}
// The tanh GELU, the approximation from the original paper and what GPT-2
// shipped. Same shape, different function: the cubic inside the tanh is not
// a rounding of the erf form.
__device__ __forceinline__ float geluTanh(float x) {
const float inner = kGeluA * (x + kGeluB * x * x * x);
return 0.5f * x * (1.0f + tanhf(inner));
}
Which one your kernel should compute is decided by the checkpoint you are serving, not by taste, and picking the other one is a correctness bug that looks like a precision quibble.
Three kernels that must agree, and one that must not
Full program in
code/day96-elementwise/elementwise.cu.
Three properties keep the measurement honest.
The byte counts print before the times, from the same constants the launches use, so the model the run is judged against cannot drift away from the program.
Two RoPE kernels must agree bit for bit, and the third must not.
ropeInterleavedVec2 runs the same arithmetic in the same order as
ropeInterleaved with one vectorized load
instead of two scalar ones, so it is compared for exact equality; a
tolerance there would hide the class of bug a vectorised rewrite
introduces. ropeHalfSplit computes a different tensor, and the program
fails if it does not.
Everything is checked against a double-precision reference on a ragged
shape before anything is timed, 3 by 17 by 5 by 64 for RoPE and 101 rows
of 97 for the chain, and timed with
events against a copyFloor kernel measured in the
same process.
__global__ void ropeInterleaved(const float* __restrict__ in,
const float* __restrict__ cosTab,
const float* __restrict__ sinTab,
float* __restrict__ out, size_t pairs,
int halfDim, int heads, int seq) {
const size_t p = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
if (p < pairs) {
const size_t half = static_cast<size_t>(halfDim);
const size_t j = p % half;
const size_t pos = (p / (half * static_cast<size_t>(heads))) %
static_cast<size_t>(seq);
const float c = cosTab[pos * half + j];
const float s = sinTab[pos * half + j];
const float x0 = in[2 * p];
const float x1 = in[2 * p + 1];
out[2 * p] = x0 * c - x1 * s;
out[2 * p + 1] = x0 * s + x1 * c;
}
}
Then the chain, which is the tail of a transformer's feed-forward block:
add the bias, apply GELU, add the residual. As three kernels it reads x
and writes a temporary, reads that and writes another, then reads that plus
the residual and writes the output. Two, two and three floats, so 7 per
element. As one kernel it reads x, reads the residual, writes the output:
3 per element. The launch ratio is 3 and the byte ratio is 7 to 3, and
those being different numbers is what lets the run say which one the card
was charging for.
__global__ void biasGeluResidualFused(const float* __restrict__ x,
const float* __restrict__ bias,
const float* __restrict__ r,
float* __restrict__ out, size_t cols) {
const size_t col =
blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
if (col < cols) {
const size_t i = blockIdx.y * cols + col;
out[i] = geluErf(x[i] + bias[col]) + r[i];
}
}
The row is blockIdx.y, so the column index falls out of threadIdx.x and
the bias index is the column. Write this kernel over a flat index and you
need i % cols in the inner loop to find the bias entry, which buys an
integer division on every element to save nothing.
Note. The bias vector is 1024 floats, 4 KiB against 64 MiB per buffer, and every row reads it, so it is in neither byte count. That is an assumption about the cache, not an arithmetic fact, and Nsight Compute counters are what settle it. The required Nsight Systems trace checks the launch timeline but does not collect those counters; day 48's
dram__bytesrecipe does the job unchanged.
Results
Re-verified on hardware. The CUDA program ran on the same Tesla T4 with driver 580.173.02 and CUDA 13.0 (V13.0.88), passed every gate and exited 0. The measured shape barely moved: staged over fused changed from 2.544 to 2.530. Both run transcripts and both Nsight Systems reports are listed in front matter. The CUDA 13 capture text contains the kernel summary; the separate published CSV remains the CUDA 12.6 export.
| kernel | ms | GB/s | x copy |
|---|---|---|---|
ropeInterleaved |
0.6081 | 220.7 | 0.89 |
ropeInterleavedVec2 |
0.5630 | 238.4 | 0.96 |
ropeHalfSplit |
0.5681 | 236.3 | 0.95 |
copyFloor |
0.5392 | 248.9 | 1.00 |
| version | kernels | floats/elem | ms | GB/s |
|---|---|---|---|---|
addBiasRow |
1 | 2 | 0.5796 | 231.6 |
applyGeluErf |
1 | 2 | 0.6236 | 215.2 |
addResidual |
1 | 3 | 0.7696 | 261.6 |
| staged chain, one region | 3 | 7 | 1.9798 | 237.3 |
| fused chain | 1 | 3 | 0.7825 | 257.3 |
copyFloor |
1 | 2 | 0.5601 | 239.6 |
The four predictions resolved as follows.
- Refuted. Staged over fused was 2.530, above the 2.333 byte ratio and between the byte and launch models. The program's weaker prediction, merely that fusion wins, held.
- Held on the byte-scored rows. Every printed effective-bandwidth row landed within 15 percent of its local 239.6 GB/s copy floor, from 215.2 to 261.6 GB/s.
- Held under CUDA 13.0. All three RoPE times remained close. Vectorized interleaved and half-split both beat scalar interleaved; the half-split direction reverses the CUDA 12.6 run's small loss.
- Held. The GELU gap was 4.735273e-04 at x = 2.701655, inside both predicted bands.
The staged and fused correctness outputs differed at 181 of 9,797 elements,
by at most 2.384e-07, and both passed their double-reference gates. The CUDA
13 Nsight Systems report captures the staged kernels and the fused kernel in
the same traced run. Its captured summary contains all three staged kernel names
and biasGeluResidualFused, so the timeline can be inspected without
inferring it from CUDA-event totals alone.
Run it yourself
A free Colab T4 is enough, and so is any card at sm_75 or above. The program run needs no root, profiler, library or second GPU; reproducing the published timeline additionally needs unprivileged Nsight Systems:
nvcc -std=c++17 -O3 -arch=sm_75 -o elementwise elementwise.cu
No Compiler Explorer embed. The program allocates five 64 MiB buffers and times ten kernels thirteen times each, which spends the 20 second execution cap before the timing starts, and shrinking the buffers would shrink the traffic the page exists to price. If you have no GPU at all, read /setup/learn-cuda-without-a-gpu.
Exercise
Write biasGeluFused, one kernel that adds the bias and applies the exact
GELU, replacing addBiasRow and applyGeluErf. Add it to the chain in
front of the existing residual add, add its row to the part 3 table with its
own byte count, and report the fraction of the copy floor it reached.
Time: 30 to 45 minutes. Submit: the kernel, your table row, and one sentence saying what fraction of 244.7 GB/s it reached and why that number cannot exceed 1.
Check: the program already compares every chain version against the
double-precision reference at 101 rows of 97, a shape no launch
configuration covers, and prints the first bad index with both values, so
adding your version to that comparison is the whole correctness test.
Grading is the bandwidth your kernel achieves on its own byte count against
the 244.7 GB/s copy ceiling day 49 measured on this card: 90 percent or more
is gold, 75 to 90 is silver, below 75 means look at your bias index first.
Two mistakes account for most failures, i % cols from a flat index, and
counting the bias as traffic, which inflates GB/s without moving the
kernel.
Hint 1
You are not deleting a launch. You are deleting a store and the load that follows it. Write down what your kernel reads and writes per element before you write the kernel, and compare that to the 4 floats the two staged kernels move.
Hint 2
The fused chain in the repo already solves the bias-index problem. Look at
what its grid shape is doing with blockIdx.y, and at what that saves in
the inner expression.
Solution
Two floats per element, x in and the GELU result out, so 8 bytes, against
16 for the two kernels it replaces. The whole chain goes from 7 floats per
element to 5, and the residual stage is untouched at 3.
The fraction cannot exceed 1 because your kernel reads one buffer and
writes one buffer, exactly what copyFloor does, plus an erff. A result
above 1 means the byte count in the numerator is wrong, and the usual reason
is counting the bias vector once per element instead of once per row.
A fused elementwise kernel is a copy with arithmetic stapled on, and its ceiling is the copy. Count the floats it moves, divide by the copy floor, and you have predicted its time before writing it.
Pitfalls
Your RoPE kernel matches the reference and the model outputs nonsense. The reference and the model disagree about pairing. Print the first eight outputs of one head from both, and if element 0 matches while element 1 does not, you are rotating adjacent pairs where the weights expect halves, or the reverse. Fix the kernel or permute the projection weights once at load time; do not fix it by adjusting the tables.
Your fused kernel is slower and you blame the fusion. Check the index
arithmetic first. A flat index plus i % cols for the bias is an integer
division on every element, and on a kernel this close to the copy floor
there is nothing else for it to hide behind. Day 48's register argument
comes second, along with its point that the deleted
launch was only ever worth microseconds.
Your two GELUs disagree in the fourth decimal and you widen the
tolerance. A tolerance that admits the tanh form admits nearly any bug in
the exact one. Pick the form the checkpoint was trained with, keep the
tolerance at the table's f32 row, and report the difference between the
forms as its own number.
You added the bias to the byte count and your kernel beat the copy floor. A per-column vector read by every row is one cache-resident 4 KiB buffer, not 4 bytes per element. Counting it as traffic inflates GB/s by a quarter and makes an impossible number look like a good one. Day 30 shipped that mistake in the other direction and documented it.
You compared the staged and fused outputs bit for bit and gated on it. They may differ. An intermediate that goes to memory is rounded to float on the way; one that stays in a register can be carried into an FMA the staged version could not reach. Report the count and the widest gap, as the program does, and leave day 68 to make bitwise reproducibility its own subject.
Go deeper
- RoFormer, the paper RoPE comes from, with the rotary matrix in equation 15 and the elementwise form in section 3.4.2: https://arxiv.org/abs/2104.09864 (checked 2026-09-01)
torch.nn.GELU, for the two forms and which one is the default: https://docs.pytorch.org/docs/stable/generated/torch.nn.GELU.html (checked 2026-09-01)- CUDA Programming Guide, mathematical functions appendix, for the ULP
bounds on
erffandtanhf: https://docs.nvidia.com/cuda/cuda-programming-guide/05-appendices/mathematical-functions.html (checked 2026-09-01) - CUDA C++ Best Practices Guide, "Coalesced Access to Global Memory", which is the rule the pairing question reduces to: https://docs.nvidia.com/cuda/cuda-c-best-practices-guide/index.html (checked 2026-09-01)
- Programming Massively Parallel Processors, 4th edition, chapter 6, on performance considerations, for the traffic argument behind fusion: https://shop.elsevier.com/books/programming-massively-parallel-processors/hwu/978-0-323-91231-0
Next
Day 97 feeds the rotated tensor to attention, where the softmax cannot be fused the way an elementwise chain can, because every output depends on every input in its row. Its stretch combines that attention kernel with this elementwise chain and day 95's normalisation into a tiny transformer forward pass. Day 99 then assembles a smaller forward and backward pass so the capstone's gradients are not first contact. The memory bandwidth argument does not change.