b23f8f2efa
Local L40S profile (perception_overfit smoke) GPU kernel time:
1589ms → 739ms (53.5% reduction, 2.15× speedup).
Wall-clock smoke: 9.6s → 4.94s (1.94× faster).
Per-kernel deltas (nsys --cuda-graph-trace=node):
reduce_axis0: 362ms → 10ms (36× faster)
Block layout: per-column (1 block / output) → 32-wide column tile
(block_dim = 32 × 8). Cross-thread reads were strided by n_tail
(~40K floats = 160KB stride) — one cache line per thread, 8× HBM
bandwidth wasted. New tile gives coalesced 128B transactions per
warp. Block tree-reduce kept (no atomicAdd, per feedback_no_atomicadd).
+1 shared-mem pad to eliminate 32-way bank conflict on the ty reduce.
multi_horizon_heads_grn_bwd_batched: 540ms → 113ms (4.8× faster)
1. Stage h_row[HIDDEN] and a1[5,HEAD_MID] in shared at block entry.
Eliminates ~28K redundant DRAM reads/block across Pass 3 + Pass 5.
2. Pass 5 reorder: k outer / i inner with d_z1[k,m] pinned in
register; writes to grad_w1_scratch are sequential per-thread.
3. Block size 64 → 128 threads. Pass 5/6 now partition over i
(output column): cross-thread writes become COALESCED 128B/warp
(was stride-128 = 512B). Passes 2/3/4 gate on (tid < HEAD_MID).
4. Pass 3 thread role: m_out → m_in/n. Same coalescing fix on
grad_w2 writes AND w2 reads in the d_eta_2 sum.
cfc_step_backward_batched: 351ms → 271ms (1.3× faster)
1. Stage x_b[n_in] and h_old_b[n_hid] in shared (was 128× redundant
DRAM reads per block; now 1× cooperative load).
2. Pass 1 thread role: i (output row) → k (output col). For each
i loop iteration, the warp writes grad_w_in[..., tid] /
grad_w_rec[..., tid] — COALESCED 128B/warp (was stride-128
non-coalesced).
multi_horizon_heads_grn_fwd_batched: 211ms → 204ms
Stage h_row[HIDDEN] in shared — Pass 1 and Pass 3 both consume.
cfc_step_batched (fwd): 95ms → 94ms
Stage x_b and h_old_b in shared.
Shared-mem budgets fit comfortably under the 48KB SM cap (~6KB / ~2KB
respectively). All 9 perception_overfit tests pass — gradient
correctness validated end-to-end (constant-signal overfit, K-loop
capture/replay, stride-4 path, evaluate-only paths).
Discipline:
- Block tree-reduce only, never atomicAdd
- No nvrtc; pre-compiled cubins via build.rs
- Mapped-pinned-only is unaffected (CPU↔GPU contract untouched)
- Single source of truth: replaced kernels in place, no v2 suffixes
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
65 lines
2.5 KiB
Plaintext
65 lines
2.5 KiB
Plaintext
// reduce_axis0.cu — sum [B, N] → [N] along the leading axis.
|
||
//
|
||
// Layout: each block reduces TILE_J=32 contiguous columns of `out`,
|
||
// using a 32×8 thread grid. The (tx) dimension covers the column tile
|
||
// (output index `j`), the (ty) dimension covers a B-stride loop.
|
||
//
|
||
// DRAM access pattern:
|
||
// thread (tx, ty) reads per_batch[(b_base + ty*step + s) * n_tail + (j_base + tx)]
|
||
// for s = 0..(n_batch/TILE_B), where (tx varies, ty fixed) walks a warp.
|
||
// → 32 consecutive `j` per warp → COALESCED 128-byte transactions.
|
||
//
|
||
// vs the prior block-per-column layout where adjacent threads strided
|
||
// the input by `n_tail` (often >40K floats = 160KB stride), forcing
|
||
// one cache line per thread and wasting 8× HBM bandwidth.
|
||
//
|
||
// Block reduce: per-thread partial sums land in s[ty][tx]; the first
|
||
// warp (ty=0) sums the 8 partials per column and writes `out[j]`.
|
||
// No atomicAdd, no warp-shuffle across warps; block-tree-reduce only.
|
||
//
|
||
// Launch contract:
|
||
// grid_dim = (ceil(n_tail / 32), 1, 1)
|
||
// block_dim = (32, 8, 1)
|
||
// shared = 0 (compile-time static)
|
||
|
||
#define REDAX0_TILE_J 32
|
||
#define REDAX0_TILE_B 8
|
||
#define REDAX0_BLOCK (REDAX0_TILE_J * REDAX0_TILE_B) // 256 threads/block
|
||
|
||
extern "C" __global__ __launch_bounds__(REDAX0_BLOCK, 4)
|
||
void reduce_axis0(
|
||
const float* __restrict__ per_batch, // [B, N]
|
||
int n_batch,
|
||
int n_tail,
|
||
float* __restrict__ out // [N] — OVERWRITE
|
||
) {
|
||
const int tx = threadIdx.x; // column-tile lane
|
||
const int ty = threadIdx.y; // batch-tile lane
|
||
const int j = (int)blockIdx.x * REDAX0_TILE_J + tx; // output column
|
||
const bool active = (j < n_tail);
|
||
|
||
float my_sum = 0.0f;
|
||
if (active) {
|
||
// Stride-loop along B in chunks of TILE_B. Within a warp (fixed ty,
|
||
// varying tx) we step `j` by +1 → coalesced 32-float load per warp.
|
||
for (int bi = ty; bi < n_batch; bi += REDAX0_TILE_B) {
|
||
my_sum += per_batch[(long long)bi * n_tail + j];
|
||
}
|
||
}
|
||
|
||
// +1 pad eliminates 32-way shared-mem bank conflict on the ty reduce.
|
||
__shared__ float s[REDAX0_TILE_B][REDAX0_TILE_J + 1];
|
||
s[ty][tx] = my_sum;
|
||
__syncthreads();
|
||
|
||
// First warp sums the 8 per-batch-tile partials for its column.
|
||
if (ty == 0 && active) {
|
||
float total = 0.0f;
|
||
#pragma unroll
|
||
for (int t = 0; t < REDAX0_TILE_B; ++t) {
|
||
total += s[t][tx];
|
||
}
|
||
out[j] = total;
|
||
}
|
||
}
|