Bias + GELU and neighbor aggregation: the mage-001 kernels, 1.4-1.7x behind Triton
These two operations still run the kernels written for the first comparison. Everything else has been revisited; these have not, and they are now the largest systematic deficit against a tuned baseline.
Numbers (kernel time, Nsight Systems, 100 iterations, same-session controls)
| Operation |
cuda-oxide (mage-006) |
Triton |
PyTorch |
cuTile (mage-004) |
| Bias + GELU 4096x768 |
11.02 |
7.76 |
16.12 |
7.97 |
| Neighbor aggregation 4096x64x65536 |
10.28 |
5.97 |
120.93 |
62.12 |
Event spans (mean of 300 samples, mage-006): GELU 13.96 against Triton's 28.35; neighbor 15.88 against Triton's 27.62. So the span column favours our kernels while the kernel column does not — the same pattern as LayerNorm before its rewrite, and the same caveat applies: the two columns come from different runs.
What both kernels do now
bias_gelu: one element per thread, thread::index_1d(), grid sized to the element count divided by 256. Each thread loads x[i] and bias[i % width], computes the tanh approximation, stores. All scalar, one pass, no reuse.neighbor: one thread per output element with a serial loop over the row's edges, reading x, weights and the CSR arrays directly from global memory per element.
Neither has a tiled or vectorised path, and neither uses more than a thread's own registers.
What to try, in order
- Vectorised element access in both kernels.
vector::as_vectors::<F32x4> gives 128-bit loads and stores; the GELU kernel is a pure elementwise map, so four elements per thread with one load and one store is the obvious shape. The bias index pattern (i % width) stays correct if the quad index is computed per lane group.
- Block per row for neighbor aggregation. The row's contribution is a reduction over
rowptr-delimited edges; a block or warp per row with a shuffle reduction removes the per-element CSR reads that the current kernel repeats, and matches how PyTorch's index_add_ reference and the Triton kernel are shaped.
- Match the control's occupancy. The LayerNorm rewrite closed 18.64 -> 8.97 by copying Triton's shape (one block per row, 31 registers, one pass, minimal shared). Measure the register count and block shape of a candidate against Triton's, as that analysis did with
ASTSource + cuobjdump on an offline compile.
- Only after those: cache hints on read-only loads, and a fused variant that writes in place.
Protocol
Same as the other gap issues: one measurement at a time on the shared GPU, a control kernel captured in the same session, kernel time and event span reported separately and named, and worse variants recorded here rather than dropped. See issue #53 for the shared-resource rules and docs/research/parallel-tracks-handoff.md for the current state of every track.
Bias + GELU and neighbor aggregation: the mage-001 kernels, 1.4-1.7x behind Triton
These two operations still run the kernels written for the first comparison. Everything else has been revisited; these have not, and they are now the largest systematic deficit against a tuned baseline.
Numbers (kernel time, Nsight Systems, 100 iterations, same-session controls)
Event spans (mean of 300 samples, mage-006): GELU 13.96 against Triton's 28.35; neighbor 15.88 against Triton's 27.62. So the span column favours our kernels while the kernel column does not — the same pattern as LayerNorm before its rewrite, and the same caveat applies: the two columns come from different runs.
What both kernels do now
bias_gelu: one element per thread,thread::index_1d(), grid sized to the element count divided by 256. Each thread loadsx[i]andbias[i % width], computes the tanh approximation, stores. All scalar, one pass, no reuse.neighbor: one thread per output element with a serial loop over the row's edges, readingx,weightsand the CSR arrays directly from global memory per element.Neither has a tiled or vectorised path, and neither uses more than a thread's own registers.
What to try, in order
vector::as_vectors::<F32x4>gives 128-bit loads and stores; the GELU kernel is a pure elementwise map, so four elements per thread with one load and one store is the obvious shape. The bias index pattern (i % width) stays correct if the quad index is computed per lane group.rowptr-delimited edges; a block or warp per row with a shuffle reduction removes the per-element CSR reads that the current kernel repeats, and matches how PyTorch'sindex_add_reference and the Triton kernel are shaped.ASTSource+cuobjdumpon an offline compile.Protocol
Same as the other gap issues: one measurement at a time on the shared GPU, a control kernel captured in the same session, kernel time and event span reported separately and named, and worse variants recorded here rather than dropped. See issue #53 for the shared-resource rules and
docs/research/parallel-tracks-handoff.mdfor the current state of every track.