Kernels
relu / OPTIMIZATION.md
superexpai's picture
Align remaining docs with v3 corrections — errata for position-biased L2 ratios, H100 table fix (MK, powered by Claude)
78bf5c2 verified
|
Raw
History Blame Contribute Delete
4.31 kB

ReLU Optimization for RTX 4090 (sm_89)

Erratum (2026-07-31, see L2_BAND_CLIFF.md): this doc records the original v1 tuning. Two of its conclusions were later corrected: the vs-torch/vs-Hub margins measured with shared-buffer interleaves were position-biased (v3 numbers: ~parity vs torch for fp32 on the 4090), and the "streaming everywhere ≥ 32 MB" config had a real 0.34–0.65x cliff vs torch for working sets in (32 MiB, L2] — the sweep below never sampled that band (its sizes jump 33.6 MB → 134 MB). Fixed in v3 with default loads + streaming stores in-band. The sweep's DRAM-wall findings and config rankings within the same measurement layout remain valid.

Float32 ReLU (out[i] = max(in[i], 0)) tuned for the NVIDIA RTX 4090 (Ada, sm_89), built via kernel-builder (torch212-cu126) and benchmarked with the official kernels ReluBenchmark plus controlled timing. Method: empirical design-space sweep + independent design fan-out + roofline analysis.

TL;DR

Deployed kernel (relu_cuda/relu.cu): float4 vectorized + streaming load/store (__ldcs/__stcs) + adaptive block size (128 L2-resident / 512 DRAM), unroll 1, flat grid-stride, scalar tail.

Workload Hub (orig) Optimized torch.relu Optimized result
base 1024² (L2) 1836 GB/s 2335 GB/s 2285 1.27× vs Hub, beats torch
large 4096² (DRAM) 940 GB/s 945 GB/s 936 at the roofline (1.01×)

Bit-identical to the upstream kernel on all inputs (incl. NaN→0). Correctness verified incl. tail sizes 4097/333/1.

Method (ultracode)

  1. Empirical sweep (bench/sweep.cu, nvcc -arch=sm_89 -O3): 31 configs × 8 sizes (64K → 134M elems), each correctness-checked, CUDA-event timed. Configs = {vec2,vec4,scalar} × unroll{1,2,4} × store{default,__stcs} × load{default,__ldcs} × block{128,256,512}.
  2. Design fan-out (5 independent senior-engineer agents + roofline analyst). The streaming agent verified in SASS that __ldcs/__stcs lower to LDG.E.EF.128 / STG.E.EF.128 (Evict-First, not a hard L2 bypass on sm_89) — explaining why the large-case gain is small.
  3. Roofline: theoretical 1008 GB/s; practical achievable copy ~960; ReLU target ~955. So 942→ceiling headroom is only ~+1.5–2% (in the noise).

What the sweep proved

  • DRAM-bound large (working set > ~50 MB): hard wall ~944 GB/s (93.7% of peak, ~98% of the ~960 practical achievable). Every config — scalar to float4, streaming or not, any block/unroll — lands 940–945. Nothing beats the wall.
  • unroll ≥ 4 is consistently SLOWER on large (929–936) — register pressure. This is why the earlier "v2" (ILP-4 + occupancy-capped grid) regressed.
  • Under-subscribing the grid (v2's other bug) costs ~2% — keep a full grid.
  • L2-resident (≤ ~33 MB): streaming load+store wins. On base 1024², float4 + __ldcs + __stcs + block128 = 2357 GB/s in the microbench (2335 in the kernels harness) vs 2218 for plain float4 — +6%, and ahead of torch. Streaming avoids L2 write-allocate pollution; smaller blocks balance SMs.
  • Block size only matters in the L2 regime: 128 best for ≤~2–4 M elems, 512 for the 16–33 MB band; irrelevant on DRAM. Hence the adaptive threshold.

Why not "more"?

The large case is memory-bound at the GDDR6X wall — a read+write elementwise op cannot exceed copy bandwidth. ILP, wider vectors, bigger grids, streaming: all measured flat-or-worse on large. The only real headroom is L2-resident sizes, which the streaming + block-128 path captures. Knowing when to stop is the optimization.

Semantics

out = x > 0 ? x : 0 (matches upstream kernels-community/relu exactly, including NaN→0). PyTorch's F.relu propagates NaN→NaN; the upstream kernel and this one return 0 for NaN. This is a pre-existing upstream semantic, intentionally preserved — the optimization is performance-only, not a numerics change.

Reproduce

# microbench sweep
cd bench && nvcc -arch=sm_89 -O3 -o sweep sweep.cu && ./sweep
# build + official benchmark
cd relu-build && kernel-builder build --variant torch212-cxx11-cu126-x86_64-linux
kernels benchmark kernels-community/relu --version 1     # official Hub baseline