Journal / 10 September 2026 / mage-002
Rewriting the layer norm and matmul kernels
Two Rust kernels restructured so values stay in registers and partial answers meet inside a warp, with the measured GPU time before and after.
The first comparison measured five operations through PyTorch, Triton, and Rust and left two of the Rust kernels behind the alternatives. The source showed where the time went. Layer normalization passed every partial sum through shared memory and a sequence of barriers, and matrix multiplication computed one output per thread from small shared tiles. Both arrangements spend time on work the arithmetic does not ask for.
Two rewrites change where the values are held and which threads have to meet. This entry presents the kernels, the measured GPU time before and after, and the three untouched operations that stayed in the harness as controls.
A reduction has to meet somewhere
Layer normalization reduces each row twice, once for the mean and once for the squared deviations around it:
\[\mu_i=\frac{1}{d}\sum_{j=1}^{d}x_{ij}, \qquad \sigma_i^2=\frac{1}{d}\sum_{j=1}^{d}\left(x_{ij}-\mu_i\right)^2 .\]Each reduction ends in a single value per row, so the threads holding partial answers must combine them. Dividing a row across 256 threads makes that combination the expensive part: the first kernel wrote one partial sum per thread into shared memory and halved the block through five barriers, with every thread waiting at each step.
The rewrite gives one warp to a row. Thirty-two lanes read the row in 128-bit quads — four consecutive floats per load — and add each quad into a single register; the warp then reduces its lanes with shuffle_down offsets of 16, 8, 4, 2, and 1. The partial answers meet inside registers instead of shared memory, and the same structure computes the centered variance. The resource request in the profile changes with it: 0 bytes of shared memory and 40 registers per thread, against 1024 bytes and 27 registers before, and 512 blocks of eight warps replace 4096 blocks of one row each.
Removing the barrier also removes the reason the block had to be sized to a row. The scalar layer_norm stays in the file for widths that are not a multiple of four.
A tile decides how often shared memory is read
A matrix product reuses both inputs along its contraction:
\[C_{ij}=\sum_k A_{ik}B_{kj}.\]One output needs a row of $A$ and a column of $B$. Neighboring outputs reuse most of that data, which is what a tile makes explicit: a block owning a 64 × 64 region of $C$ needs 64 rows of $A$ and 64 columns of $B$, reused across the whole region.
The first tiled kernel gave each thread one output from 16 × 16 tiles, so a thread read a full row and column segment of shared memory for every output: sixteen shared-memory reads per multiply-add, from 4096 blocks.
The rewrite keeps a 4 × 4 tile of outputs per thread: sixteen accumulators fed by four values of $A$ and four of $B$ per step, which is sixteen multiply-adds from eight shared-memory reads. The contraction moves in steps of 32, and before each step the block loads its two tiles through 128-bit quads — a 64 × 32 region of $A$ and a 32 × 64 region of $B$ — into 16384 bytes of shared memory. Blocks fall from 4096 to 256 and registers per thread rise from 37 to 55.
A deeper step in the contraction means fewer barriers: the tile is loaded once per 32 columns instead of once per 16. The trade is explicit, since more values per thread means fewer threads resident, which is why the kernel keeps the 4 × 4 shape instead of a wider one.
What the measured time shows
Both rewritten kernels move the Rust numbers closer to the library baselines in the same capture session. Matrix multiplication falls from 344.0 µs to 141.7 µs of GPU kernel time, against 54.6 µs for PyTorch’s library-backed call and 84.7 µs for Triton. Layer normalization falls from 18.6 µs to 11.1 µs, against 11.4 µs for PyTorch and 8.1 µs for Triton. The spans around the call move in the same direction, from 326.7 µs to 144.4 µs and from 19.5 µs to 15.4 µs.
The three operations that keep their earlier kernels — bias + GELU, triangle contraction, and neighbor aggregation — reproduce their earlier kernel times within 2%, which is what this harness can resolve between runs. Their event spans move further, by up to 66% for Triton’s neighbor aggregation, so the span is the noisier of the two measurements at these durations.
A reduction that needed a block and now needs a warp is a structural change, and it is the part that can be explained without hardware counters. The kernel time follows it. The counters that would separate instruction count from memory traffic from occupancy remain unavailable on this host, so that decomposition is not measured here.
What these numbers do not establish
The rewrite closes a gap without closing it. The new matrix multiply is 2.60 × PyTorch’s library call and 1.67 × Triton’s kernel at this shape, and layer normalization is 1.38 × Triton’s kernel while roughly level with PyTorch’s.
These are five fixed FP32 shapes on one WSL workstation with unlocked clocks and no exclusive-use guarantee. Kernel time comes from a single Nsight Systems capture per implementation and operation, so it carries no round-to-round interval of its own; the LayerNorm event span for the rewritten kernel spread from 13.5 µs to 19.0 µs across three rounds, wider than the difference under discussion. Tile shapes and the contraction step were chosen by reasoning about reuse and barriers, not by a tuning sweep, and compiled PyTorch, graph replay, batched launches, lower precision, tile tails, backward passes, and end-to-end service behavior are untested.
Two variants were measured and rejected along the way, both in single runs that are not retained evidence. A layer norm variant that held a row of 128-bit quads in a [F32x4; 8] array measured 22.12 µs of kernel time against 11.02 µs for the warp kernel in the same build, so keeping the row per thread cost roughly twice the time. An 8 × 4 register tile on 128-row blocks measured 147.74 µs around the call against 145.61 µs for the retained 4 × 4 tile in the same session, which is inside the variation this harness cannot resolve.
The full method, the retained samples, and the reproduction commands are in the measurement record. The first comparison and its graphs are in field note 001.
Measured values
Both Rust kernels recompiled against the same five FP32 shapes, one RTX 4090 under WSL2, with PyTorch and Triton measured beside them in the same session. Kernel time is the sum of captured kernel durations over 100 iterations; spans are means of 300 warmed CUDA-event samples in three rotating rounds.
GPU kernel time for the two rewritten kernels, with the unchanged operations shown as controls:
Each panel has its own scale. Lower is better. PyTorch, Triton, and the rewritten Rust kernel come from one capture session; the mage-001 Rust bar is the earlier session.
Values (µs)
| Operation | Rust mage-001 | Rust mage-002 | Change | Status |
|---|---|---|---|---|
| Matrix multiplication | 344.0 | 141.7 | −58.8% | rewritten |
| Bias + GELU | 11.2 | 11.0 | −1.7% | control |
| LayerNorm | 18.6 | 11.1 | −40.5% | rewritten |
| Triangle contraction | 81.2 | 80.4 | −0.9% | control |
| Neighbor aggregation | 10.3 | 10.3 | 0.0% | control |
Kernel time and time around the call for the same five operations in the new capture:
Top row: time inside the kernels. Bottom row: time around the call. Each column has its own scale, so implementations compare within a column.
Rust values, both views (µs)
| Operation | Kernel mage-001 | Kernel mage-002 | Span mage-001 | Span mage-002 | Status |
|---|---|---|---|---|---|
| Matrix multiplication | 344.0 | 141.7 | 326.7 | 144.4 | rewritten |
| Bias + GELU | 11.2 | 11.0 | 13.0 | 13.0 | control |
| LayerNorm | 18.6 | 11.1 | 19.5 | 15.4 | rewritten |
| Triangle contraction | 81.2 | 80.4 | 78.5 | 83.4 | control |
| Neighbor aggregation | 10.3 | 10.3 | 13.2 | 12.9 | control |
GPU kernel time, summed per operation. Separate Nsight Systems capture, 100 iterations per measurement; gaps between launches are excluded. The WSL timestamp fallback has reduced precision.
Values (µs)
| Operation | PyTorch | Triton | Rust |
|---|---|---|---|
| Matrix multiplication | 54.6 | 84.7 | 141.7 |
| Bias + GELU | 15.8 | 7.7 | 11.0 |
| LayerNorm | 11.4 | 8.1 | 11.1 |
| Triangle contraction | 28.4 | 93.8 | 80.4 |
| Neighbor aggregation | 75.8 | 8.0 | 10.3 |
Mean of 300 warmed CUDA-event spans, collected in three rounds with rotating implementation order. Whiskers show the range of the three round means, not a confidence interval. A span can include gaps while the host submits work.
Values (µs)
| Operation | PyTorch | Triton | Rust |
|---|---|---|---|
| Matrix multiplication | 57.3 | 91.5 | 144.4 |
| Bias + GELU | 24.4 | 26.1 | 13.0 |
| LayerNorm | 22.9 | 20.9 | 15.4 |
| Triangle contraction | 57.7 | 100.9 | 83.4 |
| Neighbor aggregation | 101.1 | 36.2 | 12.9 |
Each row has its own scale, so implementations are comparable within a row and not across operations. The two views come from separate runs with different launch rhythms, so subtracting one from the other does not isolate host overhead. Compilation, transfers, and service startup are excluded.