feat(dqn-v2): Plan 4 Task 2c.1 — GRN kernel module (forward + backward, additive)
Plan 4 Task 2c.1. Additive module — kernels compile and load via the existing kernel-loading infrastructure but have NO production callers in this commit. Task 2c.3+4 wires them into the trunk encoder. Eight kernels in grn_kernel.cu (Linear_a/Linear_b/Linear_residual are cuBLAS GEMMs, not new kernels): - grn_elu_inplace: element-wise ELU (canonical α=1.0) - grn_glu_forward: GLU split + sigmoid (saves sigmoid for backward) - grn_residual_layernorm_forward: residual add + LN, saves mean/rstd/normed - grn_layernorm_backward_dx: FULL JACOBIAN (not the simplified attn_layer_norm_bwd_dx which would silently propagate approximation into trunk gradients) - grn_layernorm_backward_dgamma_dbeta_p1: per-block partial reduction (no atomicAdd per feedback_no_atomicadd.md) - grn_layernorm_backward_dgamma_dbeta_p2: final reduce across blocks - grn_glu_backward - grn_elu_backward LN backward formula (full Jacobian per pearl_cold_path): d_out_g = d_out * gamma s1 = sum_d(d_out_g) s2 = sum_d(d_out_g * normed) d_x = (1/H) * rstd * (H * d_out_g - s1 - normed * s2) Layout convention: row-major [B, H] (sample-major, matches spec §4.E.2 and dt_layernorm_kernel) — intentionally different from attention_kernel.cu's [D, B] col-major; the trunk encoder downstream of GRN uses row-major buffers, so per-row LN avoids a transpose. build.rs registers grn_kernel.cu (kernel count: 57 → 58). Verified nvcc compiles cleanly via `cargo build -p ml --lib`. Module is dead code until 2c.3+4. Wire-up audit updated with the additive entry per Invariant 7. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -127,6 +127,10 @@ fn main() {
|
||||
// per_branch_target_drift(). 2-block kernel computing
|
||||
// RMS(target - online) for mag + dir branches → ISV[92,93].
|
||||
"target_drift_kernel.cu",
|
||||
// Plan 4 Task 2c.1: Gated Residual Network kernels (forward + backward).
|
||||
// Module is additive — no production callers in this commit; Task 2c.3+4
|
||||
// wires it into the trunk encoder.
|
||||
"grn_kernel.cu",
|
||||
];
|
||||
|
||||
// ALL kernels get common header (BF16 types + wrappers)
|
||||
|
||||
517
crates/ml/src/cuda_pipeline/grn_kernel.cu
Normal file
517
crates/ml/src/cuda_pipeline/grn_kernel.cu
Normal file
@@ -0,0 +1,517 @@
|
||||
/*
|
||||
* grn_kernel.cu — Gated Residual Network forward + backward kernels.
|
||||
*
|
||||
* Plan 4 Task 2c.1. Additive module — these kernels compile and load
|
||||
* but have NO production callers in this commit. Task 2c.3+4 wires
|
||||
* them into the trunk encoder/backward path.
|
||||
*
|
||||
* GRN block formula (per spec §4.E.2):
|
||||
* GRN(x) = LayerNorm(P(x) + GLU(Linear_b(ELU(Linear_a(x)))))
|
||||
*
|
||||
* where:
|
||||
* Linear_a: [hidden, in_dim]
|
||||
* ELU: x if x >= 0 else exp(x) - 1 (canonical α = 1.0)
|
||||
* Linear_b: [2*hidden, hidden] (output 2× for GLU split)
|
||||
* GLU(z): split z into [a, b] each [hidden]; output = a * sigmoid(b)
|
||||
* P(x): Linear_residual(x) if dim_in != dim_out, else identity
|
||||
* LayerNorm: per-sample mean+var; gamma, beta both [hidden]
|
||||
*
|
||||
* Composition: cuBLAS GEMMs (Linear_a, Linear_b, Linear_residual) +
|
||||
* the 6 custom kernels in this file (grn_elu_inplace, grn_glu_forward,
|
||||
* grn_residual_layernorm_forward, grn_layernorm_backward,
|
||||
* grn_glu_backward, grn_elu_backward).
|
||||
*
|
||||
* IMPORTANT: the LayerNorm backward in this file is the FULL JACOBIAN,
|
||||
* not the simplified element-wise d_x = d_out * gamma * rstd that
|
||||
* attention_backward_kernel.cu uses. The simplified version works in
|
||||
* attention's residual+LN composition because the residual gradient
|
||||
* absorbs the missing terms — for GRN's trunk-encoder backward the
|
||||
* full Jacobian is required:
|
||||
*
|
||||
* d_out_g[b, d] = d_out[b, d] * gamma[d]
|
||||
* s1[b] = sum_d(d_out_g[b, :])
|
||||
* s2[b] = sum_d(d_out_g[b, :] * normed[b, :])
|
||||
* d_x[b, d] = (1/H) * rstd[b] *
|
||||
* (H * d_out_g[b, d] - s1[b] - normed[b, d] * s2[b])
|
||||
*
|
||||
* Layout convention for GRN: tensors are ROW-MAJOR [B, H] (sample-major,
|
||||
* matches spec §4.E.2 and dt_layernorm_kernel in dt_kernels.cu). This is
|
||||
* intentionally different from attention_kernel.cu's [D, B] col-major —
|
||||
* the trunk encoder downstream of GRN uses row-major buffers, so doing
|
||||
* the LN per row in row-major avoids a transpose.
|
||||
*
|
||||
* Pearls applied:
|
||||
* - pearl_cold_path_no_exception_to_gpu_drives.md — every reduction
|
||||
* (s1, s2, d_gamma, d_beta) is GPU-side; no host roundtrip even
|
||||
* in cold-path use.
|
||||
* - feedback_no_atomicadd.md — d_gamma, d_beta accumulation across
|
||||
* batch uses per-block partial reduction + final reduction pass,
|
||||
* not atomicAdd.
|
||||
*/
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <math_constants.h>
|
||||
|
||||
/* =====================================================================
|
||||
* grn_elu_inplace — Element-wise ELU activation (canonical α = 1.0).
|
||||
*
|
||||
* f(x) = x if x >= 0
|
||||
* = exp(x) - 1 if x < 0
|
||||
*
|
||||
* In-place mutation of `x`. The original value is NOT saved here — the
|
||||
* backward path needs the post-activation value (via x_saved) and uses
|
||||
* the identity exp(x_pre) = x_post + 1 for x_pre < 0 to avoid extra
|
||||
* forward stash. (Equivalently the caller passes the post-activation
|
||||
* tensor as x_saved to grn_elu_backward.)
|
||||
*
|
||||
* Launch: grid=ceil(n/256), block=256.
|
||||
* Shapes: x is any flat tensor of length n.
|
||||
* ===================================================================== */
|
||||
extern "C" __global__ void grn_elu_inplace(
|
||||
float* __restrict__ x,
|
||||
int n
|
||||
) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx >= n) return;
|
||||
float v = x[idx];
|
||||
/* Canonical ELU with α = 1.0; expm1f is numerically preferable to
|
||||
exp(v) - 1 for small negative v. */
|
||||
x[idx] = (v >= 0.0f) ? v : expm1f(v);
|
||||
}
|
||||
|
||||
/* =====================================================================
|
||||
* grn_glu_forward — Gated Linear Unit forward.
|
||||
*
|
||||
* in_2h is row-major [B, 2*H]. Per sample b, the row is split into
|
||||
* chunk_a = in_2h[b, 0..H] and chunk_b = in_2h[b, H..2*H]. Output:
|
||||
*
|
||||
* out_h[b, d] = chunk_a[b, d] * sigmoid(chunk_b[b, d])
|
||||
*
|
||||
* sigmoid(chunk_b) is saved per element for backward — GLU's d/dx of
|
||||
* the sigmoid factor needs both chunk_a and σ(b)·(1-σ(b)) which is
|
||||
* cheaper to recompute from σ(b) than to reload b and re-apply σ.
|
||||
*
|
||||
* Why GLU: per spec §4.E.2 the GRN's gating mechanism. The sigmoid gate
|
||||
* lets the model attenuate its own non-linear branch independently of
|
||||
* the residual, which empirically stabilises depth.
|
||||
*
|
||||
* Launch: grid=(ceil(B*H/256)), block=256.
|
||||
* Flat 1D dispatch over (B*H) since the split is per element.
|
||||
* ===================================================================== */
|
||||
extern "C" __global__ void grn_glu_forward(
|
||||
const float* __restrict__ in_2h, /* [B, 2*H] row-major */
|
||||
int B, int H,
|
||||
float* __restrict__ out_h, /* [B, H] row-major */
|
||||
float* __restrict__ save_sigmoid_b /* [B, H] for backward */
|
||||
) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int total = B * H;
|
||||
if (idx >= total) return;
|
||||
int b = idx / H;
|
||||
int d = idx - b * H;
|
||||
/* Row stride is 2*H, chunk_a starts at column 0, chunk_b at column H. */
|
||||
const float* row = in_2h + (size_t)b * (2 * H);
|
||||
float a_val = row[d];
|
||||
float b_val = row[H + d];
|
||||
/* Numerically stable sigmoid: 1 / (1 + exp(-x)). For very large |b_val|
|
||||
this saturates safely to 0/1 in fp32. */
|
||||
float sig_b = 1.0f / (1.0f + expf(-b_val));
|
||||
save_sigmoid_b[idx] = sig_b;
|
||||
out_h[idx] = a_val * sig_b;
|
||||
}
|
||||
|
||||
/* =====================================================================
|
||||
* grn_residual_layernorm_forward — Fused residual + LayerNorm.
|
||||
*
|
||||
* Per spec §4.E.2 GRN: y = LayerNorm(gated + residual). Implemented
|
||||
* fused so we never materialize the (gated + residual) intermediate
|
||||
* to global memory before LN.
|
||||
*
|
||||
* For each sample b:
|
||||
* z[b, d] = gated_h[b, d] + residual_h[b, d]
|
||||
* mean[b] = (1/H) * sum_d(z[b, :])
|
||||
* var[b] = (1/H) * sum_d((z[b, d] - mean[b])^2)
|
||||
* rstd[b] = 1 / sqrt(var[b] + eps)
|
||||
* normed[b, d] = (z[b, d] - mean[b]) * rstd[b]
|
||||
* y[b, d] = gamma[d] * normed[b, d] + beta[d]
|
||||
*
|
||||
* save_mean / save_rstd / save_normed are stashed for backward. Saving
|
||||
* `normed` (already affine-naked) means the backward kernel doesn't
|
||||
* need to reload (gated, residual) and recompute the centring step —
|
||||
* memory is cheap, recompute is not.
|
||||
*
|
||||
* Reduction strategy (matches dt_layernorm_kernel):
|
||||
* - 1 block per sample (grid.x = B)
|
||||
* - block size = min(H, 256)
|
||||
* - Two-pass: pass 1 = mean (warp shuffle + smem cross-warp); pass 2
|
||||
* = variance using the published mean.
|
||||
* - eps = 1e-5f (matches existing LN kernels in attention/dt).
|
||||
*
|
||||
* Launch: grid=(B), block=(min(H, 256)).
|
||||
* Shapes: gated_h, residual_h, y, save_normed are [B, H] row-major.
|
||||
* gamma, beta are [H]. save_mean, save_rstd are [B].
|
||||
* ===================================================================== */
|
||||
extern "C" __global__ void grn_residual_layernorm_forward(
|
||||
const float* __restrict__ gated_h, /* [B, H] */
|
||||
const float* __restrict__ residual_h, /* [B, H] */
|
||||
const float* __restrict__ gamma, /* [H] */
|
||||
const float* __restrict__ beta, /* [H] */
|
||||
int B, int H,
|
||||
float* __restrict__ y, /* [B, H] */
|
||||
float* __restrict__ save_mean, /* [B] */
|
||||
float* __restrict__ save_rstd, /* [B] */
|
||||
float* __restrict__ save_normed /* [B, H] */
|
||||
) {
|
||||
int b = blockIdx.x;
|
||||
if (b >= B) return;
|
||||
int tid = threadIdx.x;
|
||||
|
||||
const float* g_row = gated_h + (size_t)b * H;
|
||||
const float* r_row = residual_h + (size_t)b * H;
|
||||
float* y_row = y + (size_t)b * H;
|
||||
float* n_row = save_normed + (size_t)b * H;
|
||||
|
||||
/* ----- Pass 1: mean of (gated + residual) per row. ----- */
|
||||
float local_sum = 0.0f;
|
||||
for (int d = tid; d < H; d += blockDim.x) {
|
||||
local_sum += g_row[d] + r_row[d];
|
||||
}
|
||||
/* Warp-level reduction. */
|
||||
for (int off = 16; off > 0; off >>= 1)
|
||||
local_sum += __shfl_xor_sync(0xFFFFFFFF, local_sum, off);
|
||||
/* Cross-warp reduction via shared memory (up to 8 warps for blockDim 256). */
|
||||
__shared__ float warp_sums[8];
|
||||
int warp_id = tid >> 5;
|
||||
int lane = tid & 31;
|
||||
if (lane == 0 && warp_id < 8) warp_sums[warp_id] = local_sum;
|
||||
__syncthreads();
|
||||
float row_sum;
|
||||
if (warp_id == 0) {
|
||||
int n_warps = (blockDim.x + 31) >> 5;
|
||||
float v = (lane < n_warps) ? warp_sums[lane] : 0.0f;
|
||||
for (int off = 16; off > 0; off >>= 1)
|
||||
v += __shfl_xor_sync(0xFFFFFFFF, v, off);
|
||||
row_sum = v;
|
||||
}
|
||||
__shared__ float sh_mean;
|
||||
if (tid == 0) sh_mean = row_sum / (float)H;
|
||||
__syncthreads();
|
||||
float mean = sh_mean;
|
||||
|
||||
/* ----- Pass 2: variance using published mean. ----- */
|
||||
float local_var = 0.0f;
|
||||
for (int d = tid; d < H; d += blockDim.x) {
|
||||
float diff = (g_row[d] + r_row[d]) - mean;
|
||||
local_var += diff * diff;
|
||||
}
|
||||
for (int off = 16; off > 0; off >>= 1)
|
||||
local_var += __shfl_xor_sync(0xFFFFFFFF, local_var, off);
|
||||
if (lane == 0 && warp_id < 8) warp_sums[warp_id] = local_var;
|
||||
__syncthreads();
|
||||
float row_var;
|
||||
if (warp_id == 0) {
|
||||
int n_warps = (blockDim.x + 31) >> 5;
|
||||
float v = (lane < n_warps) ? warp_sums[lane] : 0.0f;
|
||||
for (int off = 16; off > 0; off >>= 1)
|
||||
v += __shfl_xor_sync(0xFFFFFFFF, v, off);
|
||||
row_var = v;
|
||||
}
|
||||
__shared__ float sh_rstd;
|
||||
if (tid == 0) sh_rstd = rsqrtf(row_var / (float)H + 1e-5f);
|
||||
__syncthreads();
|
||||
float rstd = sh_rstd;
|
||||
|
||||
if (tid == 0) {
|
||||
save_mean[b] = mean;
|
||||
save_rstd[b] = rstd;
|
||||
}
|
||||
|
||||
/* ----- Pass 3: normalise + affine, write y and save_normed. ----- */
|
||||
for (int d = tid; d < H; d += blockDim.x) {
|
||||
float z = g_row[d] + r_row[d];
|
||||
float normed = (z - mean) * rstd;
|
||||
n_row[d] = normed;
|
||||
y_row[d] = gamma[d] * normed + beta[d];
|
||||
}
|
||||
}
|
||||
|
||||
/* =====================================================================
|
||||
* grn_layernorm_backward_dx — Full-Jacobian LN backward (d_x only).
|
||||
*
|
||||
* Computes the FULL LayerNorm Jacobian, NOT the simplified
|
||||
* d_x = d_out * gamma * rstd
|
||||
* pattern that attn_layer_norm_bwd_dx uses. The simplified form is
|
||||
* correct only when there is a residual gradient flowing in parallel
|
||||
* that absorbs the missing s1/s2 terms — GRN's trunk-encoder backward
|
||||
* has no such absorbing channel, so we need the exact Jacobian:
|
||||
*
|
||||
* d_out_g[b, d] = d_y[b, d] * gamma[d]
|
||||
* s1[b] = sum_d(d_out_g[b, :])
|
||||
* s2[b] = sum_d(d_out_g[b, :] * save_normed[b, :])
|
||||
* d_x[b, d] = (1/H) * rstd[b] *
|
||||
* (H * d_out_g[b, d] - s1[b] - save_normed[b, d] * s2[b])
|
||||
*
|
||||
* d_x flows backward through both the residual add (so the caller
|
||||
* routes the same d_x to BOTH the gated branch and the residual
|
||||
* branch — that is just an outer-edge wiring detail handled in 2c.3+4,
|
||||
* not here).
|
||||
*
|
||||
* Reduction strategy: 1 block per sample, two-pass.
|
||||
* - Pass 1: per-thread accumulate (d_out_g_d) and (d_out_g_d * normed_d),
|
||||
* then warp + cross-warp reduce → publish s1, s2 in shared memory.
|
||||
* - Pass 2: per-thread compute d_x[b, d] using s1, s2.
|
||||
*
|
||||
* Launch: grid=(B), block=(min(H, 256)).
|
||||
* Shapes: d_y, save_normed, d_x are [B, H] row-major.
|
||||
* gamma is [H]. save_rstd is [B].
|
||||
* ===================================================================== */
|
||||
extern "C" __global__ void grn_layernorm_backward_dx(
|
||||
const float* __restrict__ d_y, /* [B, H] grad in */
|
||||
const float* __restrict__ save_normed, /* [B, H] from forward */
|
||||
const float* __restrict__ save_rstd, /* [B] from forward */
|
||||
const float* __restrict__ gamma, /* [H] */
|
||||
int B, int H,
|
||||
float* __restrict__ d_x /* [B, H] grad out */
|
||||
) {
|
||||
int b = blockIdx.x;
|
||||
if (b >= B) return;
|
||||
int tid = threadIdx.x;
|
||||
|
||||
const float* dy_row = d_y + (size_t)b * H;
|
||||
const float* nm_row = save_normed + (size_t)b * H;
|
||||
float* dx_row = d_x + (size_t)b * H;
|
||||
float rstd = save_rstd[b];
|
||||
|
||||
/* ----- Pass 1: compute s1 and s2 via two parallel reductions. ----- */
|
||||
float local_s1 = 0.0f;
|
||||
float local_s2 = 0.0f;
|
||||
for (int d = tid; d < H; d += blockDim.x) {
|
||||
float dyg = dy_row[d] * gamma[d]; /* d_out_g */
|
||||
local_s1 += dyg;
|
||||
local_s2 += dyg * nm_row[d];
|
||||
}
|
||||
/* Warp-level reductions for s1 and s2 in lockstep. */
|
||||
for (int off = 16; off > 0; off >>= 1) {
|
||||
local_s1 += __shfl_xor_sync(0xFFFFFFFF, local_s1, off);
|
||||
local_s2 += __shfl_xor_sync(0xFFFFFFFF, local_s2, off);
|
||||
}
|
||||
/* Cross-warp reduction via shared memory: 2 × 8 entries (one each for s1, s2). */
|
||||
__shared__ float warp_s1[8];
|
||||
__shared__ float warp_s2[8];
|
||||
int warp_id = tid >> 5;
|
||||
int lane = tid & 31;
|
||||
if (lane == 0 && warp_id < 8) {
|
||||
warp_s1[warp_id] = local_s1;
|
||||
warp_s2[warp_id] = local_s2;
|
||||
}
|
||||
__syncthreads();
|
||||
float row_s1, row_s2;
|
||||
if (warp_id == 0) {
|
||||
int n_warps = (blockDim.x + 31) >> 5;
|
||||
float v1 = (lane < n_warps) ? warp_s1[lane] : 0.0f;
|
||||
float v2 = (lane < n_warps) ? warp_s2[lane] : 0.0f;
|
||||
for (int off = 16; off > 0; off >>= 1) {
|
||||
v1 += __shfl_xor_sync(0xFFFFFFFF, v1, off);
|
||||
v2 += __shfl_xor_sync(0xFFFFFFFF, v2, off);
|
||||
}
|
||||
row_s1 = v1;
|
||||
row_s2 = v2;
|
||||
}
|
||||
__shared__ float sh_s1;
|
||||
__shared__ float sh_s2;
|
||||
if (tid == 0) { sh_s1 = row_s1; sh_s2 = row_s2; }
|
||||
__syncthreads();
|
||||
float s1 = sh_s1;
|
||||
float s2 = sh_s2;
|
||||
|
||||
/* ----- Pass 2: full-Jacobian d_x. -----
|
||||
* d_x[b, d] = (1/H) * rstd * (H * d_out_g[b, d] - s1 - normed[b, d] * s2)
|
||||
*/
|
||||
float inv_H = 1.0f / (float)H;
|
||||
for (int d = tid; d < H; d += blockDim.x) {
|
||||
float dyg = dy_row[d] * gamma[d];
|
||||
float normed = nm_row[d];
|
||||
dx_row[d] = inv_H * rstd * ((float)H * dyg - s1 - normed * s2);
|
||||
}
|
||||
}
|
||||
|
||||
/* =====================================================================
|
||||
* grn_layernorm_backward_dgamma_dbeta_p1 — Per-block partial reduction.
|
||||
*
|
||||
* d_gamma[d] = sum_b(d_y[b, d] * save_normed[b, d]) (over batch)
|
||||
* d_beta[d] = sum_b(d_y[b, d]) (over batch)
|
||||
*
|
||||
* No atomicAdd per feedback_no_atomicadd.md. Uses two-phase reduction:
|
||||
* Phase 1 (this kernel): grid=(num_blocks_b, H), block=256.
|
||||
* Each block reduces a stripe of B for a single d, writes
|
||||
* gamma_partials[block_b * H + d] and beta_partials[...].
|
||||
* Phase 2 (next kernel): final reduce over num_blocks_b.
|
||||
*
|
||||
* Layout note: the partials are stored stripe-major
|
||||
* gamma_partials[block_b, d] = block_b * H + d
|
||||
* so the phase-2 final reduce iterates with unit stride across blocks
|
||||
* for the same d.
|
||||
*
|
||||
* Shared memory: 2 × blockDim.x floats (one buffer for gamma partial,
|
||||
* one for beta partial; they reduce in lockstep).
|
||||
*
|
||||
* Launch: grid=(num_blocks_b, H), block=256, smem=2*256*sizeof(float).
|
||||
* ===================================================================== */
|
||||
extern "C" __global__ void grn_layernorm_backward_dgamma_dbeta_p1(
|
||||
const float* __restrict__ d_y, /* [B, H] */
|
||||
const float* __restrict__ save_normed, /* [B, H] */
|
||||
int B, int H,
|
||||
float* __restrict__ gamma_partials, /* [num_blocks_b, H] */
|
||||
float* __restrict__ beta_partials /* [num_blocks_b, H] */
|
||||
) {
|
||||
extern __shared__ float sdata[];
|
||||
float* sg = sdata;
|
||||
float* sb = sdata + blockDim.x;
|
||||
|
||||
int d = blockIdx.y; /* hidden dim this block handles */
|
||||
int tid = threadIdx.x;
|
||||
int gid = blockIdx.x * blockDim.x + tid; /* batch index */
|
||||
|
||||
float gv = 0.0f;
|
||||
float bv = 0.0f;
|
||||
if (gid < B) {
|
||||
size_t off = (size_t)gid * H + d; /* row-major [B, H] */
|
||||
float dy_val = d_y[off];
|
||||
float nm_val = save_normed[off];
|
||||
gv = dy_val * nm_val;
|
||||
bv = dy_val;
|
||||
}
|
||||
sg[tid] = gv;
|
||||
sb[tid] = bv;
|
||||
__syncthreads();
|
||||
|
||||
/* Smem tree reduction for both partials in lockstep. */
|
||||
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
|
||||
if (tid < s) {
|
||||
sg[tid] += sg[tid + s];
|
||||
sb[tid] += sb[tid + s];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
if (tid == 0) {
|
||||
gamma_partials[blockIdx.x * H + d] = sg[0];
|
||||
beta_partials[blockIdx.x * H + d] = sb[0];
|
||||
}
|
||||
}
|
||||
|
||||
/* =====================================================================
|
||||
* grn_layernorm_backward_dgamma_dbeta_p2 — Final reduce over partials.
|
||||
*
|
||||
* For each d, sums gamma_partials[0..num_blocks_b, d] → d_gamma[d] and
|
||||
* beta_partials[0..num_blocks_b, d] → d_beta[d]. This kernel produces
|
||||
* the final batch-summed d_gamma and d_beta.
|
||||
*
|
||||
* Launch: grid=(ceil(H/256)), block=256.
|
||||
* One thread per hidden dim; sequential sum over blocks.
|
||||
* ===================================================================== */
|
||||
extern "C" __global__ void grn_layernorm_backward_dgamma_dbeta_p2(
|
||||
const float* __restrict__ gamma_partials, /* [num_blocks_b, H] */
|
||||
const float* __restrict__ beta_partials, /* [num_blocks_b, H] */
|
||||
int H, int num_blocks_b,
|
||||
float* __restrict__ d_gamma, /* [H] */
|
||||
float* __restrict__ d_beta /* [H] */
|
||||
) {
|
||||
int d = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (d >= H) return;
|
||||
float dg = 0.0f;
|
||||
float db = 0.0f;
|
||||
for (int i = 0; i < num_blocks_b; i++) {
|
||||
dg += gamma_partials[(size_t)i * H + d];
|
||||
db += beta_partials[(size_t)i * H + d];
|
||||
}
|
||||
d_gamma[d] = dg;
|
||||
d_beta[d] = db;
|
||||
}
|
||||
|
||||
/* =====================================================================
|
||||
* grn_glu_backward — Gated Linear Unit backward.
|
||||
*
|
||||
* Forward was: out_h[b, d] = chunk_a[b, d] * sigmoid(chunk_b[b, d])
|
||||
*
|
||||
* Let σ = sigmoid(chunk_b). The forward saved σ in save_sigmoid_b.
|
||||
* Standard chain rule gives:
|
||||
*
|
||||
* d_chunk_a[b, d] = d_out_h[b, d] * σ
|
||||
* d_chunk_b[b, d] = d_out_h[b, d] * chunk_a * σ * (1 - σ)
|
||||
*
|
||||
* The forward stored chunk_a as the first H columns of in_2h_saved
|
||||
* (the Linear_b output before GLU split). chunk_b lives in the second H
|
||||
* columns but is NOT needed for backward — σ already encapsulates it.
|
||||
*
|
||||
* Output layout: d_in_2h is [B, 2*H] row-major; the first H columns
|
||||
* receive d_chunk_a, the second H columns receive d_chunk_b.
|
||||
*
|
||||
* Launch: grid=(ceil(B*H/256)), block=256.
|
||||
* Flat 1D dispatch over (B*H); each thread writes both
|
||||
* d_chunk_a[b, d] and d_chunk_b[b, d] for its (b, d).
|
||||
* ===================================================================== */
|
||||
extern "C" __global__ void grn_glu_backward(
|
||||
const float* __restrict__ d_out_h, /* [B, H] */
|
||||
const float* __restrict__ in_2h_saved, /* [B, 2*H] from forward (Linear_b out) */
|
||||
const float* __restrict__ save_sigmoid_b, /* [B, H] */
|
||||
int B, int H,
|
||||
float* __restrict__ d_in_2h /* [B, 2*H] grad out */
|
||||
) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int total = B * H;
|
||||
if (idx >= total) return;
|
||||
int b = idx / H;
|
||||
int d = idx - b * H;
|
||||
|
||||
/* Read forward-saved σ and the original chunk_a (column d in the first
|
||||
half of row b of in_2h_saved). */
|
||||
float sig_b = save_sigmoid_b[idx];
|
||||
const float* row_2h = in_2h_saved + (size_t)b * (2 * H);
|
||||
float chunk_a = row_2h[d];
|
||||
float dout = d_out_h[idx];
|
||||
|
||||
float d_chunk_a = dout * sig_b;
|
||||
/* d/dx of σ(x) = σ(x) * (1 - σ(x)); the inner derivative passes
|
||||
through chunk_a unmodified (chunk_a does not depend on chunk_b). */
|
||||
float d_chunk_b = dout * chunk_a * sig_b * (1.0f - sig_b);
|
||||
|
||||
/* Scatter to the same row-major [B, 2*H] layout: columns [0..H) for
|
||||
d_chunk_a, [H..2H) for d_chunk_b. */
|
||||
float* drow_2h = d_in_2h + (size_t)b * (2 * H);
|
||||
drow_2h[d] = d_chunk_a;
|
||||
drow_2h[H + d] = d_chunk_b;
|
||||
}
|
||||
|
||||
/* =====================================================================
|
||||
* grn_elu_backward — Element-wise ELU backward.
|
||||
*
|
||||
* Forward was: f(x) = x if x >= 0, else exp(x) - 1.
|
||||
* Derivative: f'(x) = 1 if x >= 0
|
||||
* = exp(x) if x < 0
|
||||
*
|
||||
* `x_saved` is the forward POST-activation value (i.e. f(x_pre)). We
|
||||
* recover f'(x_pre) without needing the pre-activation:
|
||||
* - If x_saved >= 0, then x_pre >= 0 too → f'(x_pre) = 1.
|
||||
* - If x_saved < 0, then exp(x_pre) = x_saved + 1 (since
|
||||
* x_saved = exp(x_pre) - 1) → f'(x_pre) = x_saved + 1.
|
||||
*
|
||||
* This avoids a separate pre-activation stash — the post-activation is
|
||||
* already on the forward path so we reuse it.
|
||||
*
|
||||
* Launch: grid=ceil(n/256), block=256.
|
||||
* ===================================================================== */
|
||||
extern "C" __global__ void grn_elu_backward(
|
||||
const float* __restrict__ d_out,
|
||||
const float* __restrict__ x_saved, /* post-activation from forward */
|
||||
int n,
|
||||
float* __restrict__ d_x
|
||||
) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx >= n) return;
|
||||
float v = x_saved[idx];
|
||||
/* For x_pre < 0: f(x_pre) = exp(x_pre) - 1 < 0, and f'(x_pre) = exp(x_pre) = v + 1.
|
||||
For x_pre >= 0: f(x_pre) = x_pre >= 0, and f'(x_pre) = 1. */
|
||||
float deriv = (v >= 0.0f) ? 1.0f : (v + 1.0f);
|
||||
d_x[idx] = d_out[idx] * deriv;
|
||||
}
|
||||
@@ -277,6 +277,8 @@ Plan 4 Task 5 Mode A E.5 (2026-04-24): `attention_focus_ema_kernel.cu` + `Attent
|
||||
|
||||
Plan 4 Task 4 E.4 (2026-04-24): `state_encoder.rs` + `value_decoder.rs` Rust API split. PURE Rust API refactor — no kernel changes, no parameter changes, no checkpoint break, no new ISV slot. New types: `StateEncoder` (wraps trunk encoder `h_s1` + `h_s2`), `ValueDecoder` (per-branch FC + adv-logits, one instance per `Branch::{Dir,Mag,Ord,Urg}`), `EncoderOutput`, `BranchDecoderOutput` (carries `q_per_action_dev_ptr` + Plan 2 D.3 `v_short_dev_ptr` / `v_long_dev_ptr` pass-through fields). `BatchedForward` gains `encoder_forward_only(...)` and `decoder_forward_only(branch_idx, ...)` helpers, both lossless extractions of the existing dispatch sequence in `forward_online_raw`. The monolithic `forward_online_raw` stays callable and now routes through `encoder_forward_only` internally for the trunk portion (multi-stream branch fork/join unchanged). The new types are ADDITIVE attachment points for Plan 4 Task 1 (Full VSN, pre-encoder), Task 3 (multi-Q IQN, decoder), and Task 6 (auxiliary heads, decoder peer). 2 new Wired rows.
|
||||
|
||||
Plan 4 Task 2c.1 (2026-04-24): `grn_kernel.cu` added — Gated Residual Network forward + backward kernels (8 entry points: `grn_elu_inplace`, `grn_glu_forward`, `grn_residual_layernorm_forward`, `grn_layernorm_backward_dx`, `grn_layernorm_backward_dgamma_dbeta_p1`, `grn_layernorm_backward_dgamma_dbeta_p2`, `grn_glu_backward`, `grn_elu_backward`). Linear_a / Linear_b / Linear_residual remain cuBLAS GEMMs (no new kernels). Module is **additive — ZERO production callers in this commit**; classified Orphan-by-design. Task 2c.3+4 wires it into the trunk encoder forward + backward path. LayerNorm backward implements the FULL Jacobian (`d_x = (1/H) * rstd * (H * d_out_g - s1 - normed * s2)`), NOT the simplified element-wise form `d_x = d_out * gamma * rstd` from `attn_layer_norm_bwd_dx` — the simplified form is correct only when a parallel residual gradient absorbs the missing s1/s2 terms, which the GRN trunk-encoder backward does not have. `d_gamma` / `d_beta` use a two-phase per-block partial reduction + final-reduce pass (`feedback_no_atomicadd.md` compliant). Layout convention: row-major `[B, H]` (matches `dt_layernorm_kernel`, intentionally different from `attention_kernel.cu`'s `[D, B]` col-major — trunk buffers downstream of GRN are row-major, so per-row LN avoids a transpose). All reductions GPU-side per `pearl_cold_path_no_exception_to_gpu_drives.md`. Kernel registered in `build.rs` (kernel count 57 → 58, all compile clean under nvcc sm_80). No checkpoint break (no new params, no new ISV slot, no consumer wired). 1 new Orphan-by-design row (will reclassify Wired at 2c.3+4).
|
||||
|
||||
Plan 4 Task 2b (2026-04-24): `layout_fingerprint_seed()` extended to cover **param-tensor structural layout** in addition to the existing ISV-slot section. 86 new entries `PARAM_<NAME>=<idx>` plus `PARAM_TOTAL_TENSORS=86` appended to the seed string, mirroring `compute_param_sizes()` order one-for-one. Names match the in-code comments at each tensor index (`PARAM_W_S1=0`, `PARAM_B_S1=1`, … `PARAM_W_PLAN_OUT=84`, `PARAM_B_PLAN_OUT=85`). Pragmatic Option A scope: tensor names + positions only, NOT runtime sizes — `shared_h1`/`shared_h2`/`adv_h`/`value_h`/`num_atoms`/`bottleneck_dim`/`market_dim` are runtime config and can't be embedded in a `const fn`. Size mismatches between checkpoint and current binary are caught separately by safetensors deserialization. After this commit, ANY structural reshuffle (insert/delete/reorder/rename a tensor) changes `LAYOUT_FINGERPRINT_CURRENT` and triggers fail-fast at checkpoint load. Prerequisite for Plan 4 Task 2c (GRN ADOPT) — without 2b, GRN's tensor insertion would pass silently through checkpoint load and corrupt downstream state. ISV-slot portion of the seed unchanged. New fingerprint value: `0xa504d3c2f275b8af`. Behavior change zero; this commit IS a checkpoint break by intent. No new module / kernel / ISV slot — pure contract-coverage extension.
|
||||
|
||||
| Classification | Count |
|
||||
|
||||
Reference in New Issue
Block a user