@@ -2368,12 +2368,39 @@ extern "C" __global__ void compute_expected_q(
branch_logit_offset += N * n_d * num_atoms;
}
/* Accumulate atom stats across all samples (monitoring only) */
/* Accumulate atom stats across all samples (monitoring only).
* Uses warp+block reduction to minimize atomicAdd non-determinism.
* One atomicAdd per BLOCK (not per thread). */
if (atom_stats != NULL) {
/* Average per-action: divide by total_actions so the final mean is per-action entropy */
float inv_actions = 1.0f / (float)total_actions;
atomicAdd(&atom_stats[0], sample_entropy * inv_actions) ;
atomicAdd(&atom_stats[1], (float)sample_utilized * inv_actions) ;
float my_entropy = sample_entropy * inv_actions;
float my_util = (float)sample_utilized * inv_actions;
/* Warp-level reduction */
for (int offset = 16; offset > 0; offset >>= 1) {
my_entropy += __shfl_xor_sync(0xFFFFFFFF, my_entropy, offset);
my_util += __shfl_xor_sync(0xFFFFFFFF, my_util, offset);
}
/* Block-level reduction via shared memory */
__shared__ float s_entropy[8];
__shared__ float s_util[8];
int warp_id = threadIdx.x / 32;
int lane = threadIdx.x % 32;
if (lane == 0) { s_entropy[warp_id] = my_entropy; s_util[warp_id] = my_util; }
__syncthreads();
if (warp_id == 0) {
float ve = (lane < blockDim.x / 32) ? s_entropy[lane] : 0.0f;
float vu = (lane < blockDim.x / 32) ? s_util[lane] : 0.0f;
for (int off = 16; off > 0; off >>= 1) {
ve += __shfl_xor_sync(0xFFFFFFFF, ve, off);
vu += __shfl_xor_sync(0xFFFFFFFF, vu, off);
}
if (lane == 0) {
atomicAdd(&atom_stats[0], ve);
atomicAdd(&atom_stats[1], vu);
}
}
}
}
@@ -3040,15 +3067,16 @@ extern "C" __global__ void selectivity_forward(
/**
* Mamba-2 selectivity gate — backward pass (BCE gradient).
* DETERMINISTIC: one thread per weight element, loops over batch.
*
* target[b] = clamp(per_sample_loss[b] / mean_loss, 0, 1)
* d_z[b] = sel[b] - target[b]
* dW[k] + = d_z[b] * h_s2[b, k] (atomicAdd )
* db + = d_z[b] (atomicAdd )
* dW[k] = sum_b( d_z[b] * h_s2[b, k])
* db = sum_b( d_z[b])
*
* d_sel_params layout: [dW[SH2], db] (SH2+1 floats, must be zeroed before call )
* d_sel_params layout: [dW[SH2], db] (SH2+1 floats)
*
* Grid: ceil(B /256), Block: 256
* Grid: ceil((SH2+1) /256), Block: 256
*/
extern "C" __global__ void selectivity_backward(
const float* __restrict__ h_s2, /* [B, SH2] */
@@ -3058,18 +3086,29 @@ extern "C" __global__ void selectivity_backward(
float mean_loss,
int B, int SH2)
{
int b = blockIdx.x * blockDim.x + threadIdx.x;
if (b >= B) return ;
int widx = blockIdx.x * blockDim.x + threadIdx.x;
int total_params = SH2 + 1 ;
if (widx >= total_params) return;
float safe_mean = fmaxf(mean_loss, 1e-8f);
float target = fminf(fmaxf(per_sample_loss[b] / safe_mean, 0.0f), 1.0f) ;
float d_z = sel_out[b] - target;
float sum = 0.0f ;
const float* h = h_s2 + b * SH2;
for (int k = 0; k < SH2; k++) {
atomicAdd(&d_sel_params[k], d_z * h[k]);
if (widx < SH2) {
/* dW[widx] = sum_b(d_z[b] * h_s2[b, widx]) */
for (int b = 0; b < B; b++) {
float target = fminf(fmaxf(per_sample_loss[b] / safe_mean, 0.0f), 1.0f);
float d_z = sel_out[b] - target;
sum += d_z * h_s2[(long long)b * SH2 + widx];
}
} else {
/* db = sum_b(d_z[b]) */
for (int b = 0; b < B; b++) {
float target = fminf(fmaxf(per_sample_loss[b] / safe_mean, 0.0f), 1.0f);
float d_z = sel_out[b] - target;
sum += d_z;
}
}
atomicAdd(& d_sel_params[SH2], d_z) ; /* db */
d_sel_params[widx] = sum ; /* single deterministic write */
}
/**
@@ -3494,14 +3533,17 @@ extern "C" __global__ void kan_gate_backward(
d_gate_pre[i] = d_raw * dgate_dx;
/* Accumulate d_coeff via atomicAdd (batch dimension reduction) */
/* Accumulate d_coeff and d_residual_w via atomicAdd (batch reduction).
* DETERMINISM NOTE: Each thread has a unique (neuron, span) combination
* so warp-level pre-reduction is not applicable (different target indices).
* Full elimination requires a separate per-weight reduction kernel.
* Impact: KAN spline is a small auxiliary gating component (~adv_h*8
* coefficients + adv_h residual weights), not on the primary gradient path. */
float* dc = d_spline_coeff + neuron * NUM_BASES;
if (k0 >= 0 && k0 < NUM_BASES) atomicAdd(&dc[k0], d_raw * b0);
if (span >= 0 && span < NUM_BASES) atomicAdd(&dc[span], d_raw * b1);
if (k2 >= 0 && k2 < NUM_BASES) atomicAdd(&dc[k2], d_raw * b2);
if (k3_idx >= 0 && k3_idx < NUM_BASES) atomicAdd(&dc[k3_idx], d_raw * b3);
/* Accumulate d_residual_w via atomicAdd */
atomicAdd(&d_residual_w[neuron], d_raw * x);
}
@@ -3820,32 +3862,16 @@ extern "C" __global__ void liquid_tau_rk4_step(
/**
* Backward through 2-step FC-SiLU-FC diffusion denoiser.
* DETERMINISTIC: one thread per weight element, loops over batch.
*
* Loss: L = mean_b((Q_refined[b] - Q_target[b])^2) summed over 12 actions.
* d_Q_refined = 2*(Q_refined - Q_target) / B per element.
*
* Then backprop through 2 FC-SiLU-FC steps (reverse order) to accumulate
* d_params via atomicAdd across the batch dimension .
* d_params by summing over all samples with fixed iteration order .
*
* For each step (reverse):
* d_W2 += outer(d_out, h_silu), d_b2 += d_out
* d_h = W2^T @ d_out
* d_pre_silu = d_h * silu'(pre_act)
* d_W1 += outer(d_pre_silu, input), d_b1 += d_pre_silu
*
* silu'(x) = sigmoid(x) + x * sigmoid(x) * (1 - sigmoid(x))
* = sigmoid(x) * (1 + x * (1 - sigmoid(x)))
*
* One thread per sample. Grid: ceil(B/256), Block: 256.
* Uses atomicAdd for weight gradient accumulation across batch.
*
* @param Q_refined [B, 12] denoiser output (final q_coord_buf)
* @param Q_target [B, 12] target network Q-values
* @param denoise_params [1800] current weights (2 steps ×
* @param d_denoise_params [1800] gradient accumulator (caller must zero)
* @param Q_input [B, 12] denoiser input (pre-denoise Q-values)
* @param var_q [B, 12] Var[Q] for noise conditioning
* @param B batch size
* Total params: 1800 (2 steps ×
* Grid: ceil(1800/256), Block: 256.
*/
extern "C" __global__ void q_denoise_backward(
const float* __restrict__ Q_refined,
@@ -3856,52 +3882,45 @@ extern "C" __global__ void q_denoise_backward(
const float* __restrict__ var_q,
int B)
{
int b = blockIdx.x * blockDim.x + threadIdx.x;
if (b >= B) return;
int widx = blockIdx.x * blockDim.x + threadIdx.x;
const int D = 12;
const int H = 24;
const int STEP_PARAMS = 900; /* W1[24,24]+b1[24]+W2[12,24]+b2[12] */
const int num_steps = 2;
const int TOTAL_PARAMS = num_steps * STEP_PARAMS;
if (widx >= TOTAL_PARAMS) return;
float inv_B = 1.0f / (float)B;
/* Compute loss gradient: d_out[i] = 2*(Q_refined[b,i] - Q_target[b,i]) / B */
float d_out[12] ;
for (int i = 0; i < D; i++) {
d_out[i] = 2.0f * (Q_refined[b * D + i] - Q_target[b * D + i]) * inv_B;
}
/* Determine which step and which sub-parameter this thread handles */
int step = widx / STEP_PARAMS ;
int local = widx % STEP_PARAMS;
/* Backprop through steps in reverse order (step 1 then step 0).
* Each step: input_step[24] → h=SiLU(W1@input+b1)[24] → residual=W2@h+b2[12]
* Q_next = Q_prev + 0.1 * residual
* So d_residual = 0.1 * d_out, and d_Q_prev += d_out (identity skip connection). */
/* Sub-parameter layout within one step (900 total):
* W1: [H, H] = 576 elements (offsets 0..575)
* b1: [H] = 24 elements (offsets 576..599)
* W2: [D, H] = 288 elements (offsets 600..887)
* b2: [D] = 12 elements (offsets 888..899) */
int is_W1 = (local < 576);
int is_b1 = (local >= 576 && local < 600);
int is_W2 = (local >= 600 && local < 888);
/* int is_b2 = (local >= 888); */
for (int step = num_steps - 1; step >= 0; step--) {
int off = step * STEP_PARAMS;
const float* W1 = denoise_params + off;
const float* b1_ptr = denoise_params + off + 576;
const float* W2 = denoise_params + off + 600;
const float* b2_ptr = denoise_params + off + 888;
float acc = 0.0f;
float* dW1 = d_denoise_params + off;
float* db1 = d_denoise_params + off + 576;
float* dW2 = d_denoise_params + off + 600 ;
float* db2 = d_denoise_params + off + 888;
/* d_residual = 0.1 * d_out */
float d_res[12];
for (int b = 0; b < B; b++) {
/* Compute loss gradient for this sample */
float d_out_b[12] ;
for (int i = 0; i < D; i++) {
d_res [i] = 0.1 f * d_out[i] ;
d_out_b [i] = 2.0 f * (Q_refined[b * D + i] - Q_target[b * D + i]) * inv_B ;
}
/* Reconstruct forward: build input for this step.
* Step 0: input = [Q_input; noise_level_step0]
* Step 1: input = [Q_after_step0; noise_level_step1]
* We need to reconstruct the intermediate Q-values.
* For simplicity, re-run the forward to get intermediates. */
/* Process steps in reverse. The gradient for the current step is needed. */
/ * For each step s from (num_steps-1) down to 0, d_out propagates through
* the skip connection unchanged. We only need the gradient at `step`. */
/* Reconstruct Q at start of this step by running forward from Q_input
* through steps 0..step-1. */
/* Reconstruct Q at start of ` step` by running forward from Q_input
* through steps 0..( step-1) . */
float q_curr[12];
for (int i = 0; i < D; i++) q_curr[i] = Q_input[b * D + i];
@@ -3920,19 +3939,19 @@ extern "C" __global__ void q_denoise_backward(
}
float h_s[24];
for (int i = 0; i < H; i++) {
float acc = sb1[i];
for (int j = 0; j < H; j++) acc += sW1[i * H + j] * inp_s[j];
float sig = 1.0f / (1.0f + expf(-acc ));
h_s[i] = acc * sig;
float v = sb1[i];
for (int j = 0; j < H; j++) v += sW1[i * H + j] * inp_s[j];
float sig = 1.0f / (1.0f + expf(-v ));
h_s[i] = v * sig;
}
for (int i = 0; i < D; i++) {
float acc = sb2[i];
for (int j = 0; j < H; j++) acc += sW2[i * H + j] * h_s[j];
q_curr[i] += 0.1f * acc ;
float v = sb2[i];
for (int j = 0; j < H; j++) v += sW2[i * H + j] * h_s[j];
q_curr[i] += 0.1f * v ;
}
}
/* Now q_curr = Q at start of this step. Build the input. */
/* Build input for this step */
float schedule = 1.0f - (float)(step + 1) / (float)num_steps;
float input[24];
for (int i = 0; i < D; i++) {
@@ -3941,54 +3960,67 @@ extern "C" __global__ void q_denoise_backward(
}
/* Forward through this step to get pre-activations */
int off = step * STEP_PARAMS;
const float* W1 = denoise_params + off;
const float* b1_ptr = denoise_params + off + 576;
const float* W2 = denoise_params + off + 600;
float pre_act[24], h_silu[24];
for (int i = 0; i < H; i++) {
float acc = b1_ptr[i];
for (int j = 0; j < H; j++) acc += W1[i * H + j] * input[j];
pre_act[i] = acc ;
float sig = 1.0f / (1.0f + expf(-acc ));
h_silu[i] = acc * sig;
float v = b1_ptr[i];
for (int j = 0; j < H; j++) v += W1[i * H + j] * input[j];
pre_act[i] = v ;
float sig = 1.0f / (1.0f + expf(-v ));
h_silu[i] = v * sig;
}
/* Backward: d_W2 += outer(d_res, h_silu), d_b2 += d_res */
/* d_residual = 0.1 * d_out */
float d_res[12];
for (int i = 0; i < D; i++) {
atomicAdd(&db2[i], d_res [i]) ;
for (int j = 0; j < H; j++) {
atomicAdd(&dW2[i * H + j], d_res[i] * h_silu[j]);
d_res[i] = 0.1f * d_out_b [i];
}
/* Compute the gradient for this thread's specific parameter */
if (is_W2) {
/* dW2[i, j] += d_res[i] * h_silu[j] */
int idx = local - 600;
int row = idx / H;
int col = idx % H;
acc += d_res[row] * h_silu[col];
} else if (local >= 888) {
/* db2[i] += d_res[i] */
int idx = local - 888;
acc += d_res[idx];
} else {
/* Need d_pre for W1/b1 gradients */
float d_h[24];
for (int i = 0; i < H; i++) {
float v = 0.0f;
for (int j = 0; j < D; j++) v += W2[j * H + i] * d_res[j];
d_h[i] = v;
}
float d_pre[24];
for (int i = 0; i < H; i++) {
float sig = 1.0f / (1.0f + expf(-pre_act[i]));
float silu_grad = sig * (1.0f + pre_act[i] * (1.0f - sig));
d_pre[i] = d_h[i] * silu_grad;
}
if (is_W1) {
/* dW1[i, j] += d_pre[i] * input[j] */
int idx = local;
int row = idx / H;
int col = idx % H;
acc += d_pre[row] * input[col];
} else if (is_b1) {
/* db1[i] += d_pre[i] */
int idx = local - 576;
acc += d_pre[idx];
}
}
/* d_h = W2^T @ d_res */
float d_h[24];
for (int i = 0; i < H; i++) {
float acc = 0.0f;
for (int j = 0; j < D; j++) {
acc += W2[j * H + i] * d_res[j];
}
d_h[i] = acc;
}
/* d_pre_silu = d_h * silu'(pre_act)
* silu'(x) = sigmoid(x) * (1 + x * (1 - sigmoid(x))) */
float d_pre[24];
for (int i = 0; i < H; i++) {
float sig = 1.0f / (1.0f + expf(-pre_act[i]));
float silu_grad = sig * (1.0f + pre_act[i] * (1.0f - sig));
d_pre[i] = d_h[i] * silu_grad;
}
/* d_W1 += outer(d_pre, input), d_b1 += d_pre */
for (int i = 0; i < H; i++) {
atomicAdd(&db1[i], d_pre[i]);
for (int j = 0; j < H; j++) {
atomicAdd(&dW1[i * H + j], d_pre[i] * input[j]);
}
}
/* Propagate gradient to previous step's output:
* d_out_prev = d_out (skip connection carries gradient through) */
/* d_out already holds the gradient for the previous step. */
}
d_denoise_params[widx] = acc; /* single deterministic write */
}
/* ================================================================== */
@@ -4328,21 +4360,22 @@ extern "C" __global__ void atom_position_gradient(
/**
* Mamba2 BPTT: backpropagation through K=8 scan steps.
* DETERMINISTIC: one thread per weight element, loops over batch.
*
* Given d_h_enriched [B, SH2] (gradient from downstream):
* The temporal_context = (W_C @ h_current) * x_K, so:
* d_x_K = d_context * (W_C @ h_current) element-wise
* d_W_C + = sum_b d_context * x_K^T (outer product )
* d_W_C[j,s] = sum_b( d_out[b,j] * x_K[b,s] )
*
* Reverse scan t = K-1..0:
* d_A_t = d_x_{t+1} * x_t * sigmoid'(a_raw_t)
* d_x_t = d_x_{t+1} * A_t (gate propagation)
* d_B_t = d_x_{t+1}
* d_W_A += sum_b d_A_t @ h_t^T
* d_W_B += sum_b d_B_t @ h_t^T
* d_x_t = d_x_{t+1} * A_t
* d_W_A[j,s] += sum_b( d_gate[b,t,s] * h_t[b,j] )
* d_W_B[j,s] += sum_b( d_x[b,t,s] * h_t[b,j] )
*
* Grid: ceil(B /256), Block: 256. One thread per sample .
* Accumulates weight gradients via atomicAdd (batch reduction) .
* Grid: ceil(total_weight_params /256), Block: 256. One thread per weight .
* total_weight_params = 3 * sh2 * state_d .
* Weight layout: [d_w_a: sh2*state_d | d_w_b: sh2*state_d | d_w_c: sh2*state_d]
*/
extern "C" __global__ void mamba2_scan_backward(
const float* __restrict__ h_history, /* [B, K, SH2] saved activations */
@@ -4360,91 +4393,87 @@ extern "C" __global__ void mamba2_scan_backward(
int state_d,
const float* __restrict__ isv_signals /* [12] pinned. NULL = no regime modulation. */
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= B) return ;
int widx = blockIdx.x * blockDim.x + threadIdx.x;
int wc_size = sh2 * state_d ;
int total_params = 3 * wc_size;
if (widx >= total_params) return;
const float* h_current = h_s2 + (long long)i * sh2;
const float* d_out = d_h_enriched + (long long)i * sh2;
/* Determine which weight matrix and which (j, s) indices */
int mat; /* 0=W_A, 1=W_B, 2=W_C */
int local;
if (widx < wc_size) {
mat = 0; local = widx;
} else if (widx < 2 * wc_size) {
mat = 1; local = widx - wc_size;
} else {
mat = 2; local = widx - 2 * wc_size;
}
int j = local / state_d;
int s = local % state_d;
/* ── Forward pass replay to get intermediate states ── */
float x_states[9][16]; /* x_states[t] for t=0..K, state_d<=16 */
float a_raw[8][16]; /* raw gate values before sigmoid */
float acc = 0.0f;
/* x_states[0] = zeros */
for (int s = 0; s < state_d; s++) x_states[0][s] = 0.0f ;
for (int b = 0; b < B; b++) {
const float* d_out = d_h_enriched + (long long)b * sh2 ;
for (int t = 0; t < K; t++) {
const float* h_t = h_history + ((long long)i * K + t) * sh2;
for (int s = 0; s < state_d; s++) {
/* ── Forward pass replay to get intermediate states ── */
float x_states_local[9]; /* x_states[t] for t=0..K, only state_dim s */
float a_raw_local[8]; /* raw gate values before sigmoid, only s */
x_states_local[0] = 0.0f;
for (int t = 0; t < K; t++) {
const float* h_t = h_history + ((long long)b * K + t) * sh2;
float a_val = 0.0f, b_val = 0.0f;
for (int j = 0; j < sh2; j++) {
a_val += w_a[(long long)j * state_d + s] * h_t[j];
b_val += w_b[(long long)j * state_d + s] * h_t[j];
for (int jj = 0; jj < sh2; j j++) {
a_val += w_a[(long long)jj * state_d + s] * h_t[jj ];
b_val += w_b[(long long)jj * state_d + s] * h_t[jj ];
}
a_raw[t][s ] = a_val;
a_raw_local[t ] = a_val;
float gate = 1.0f / (1.0f + expf(-a_val));
/* Regime-conditioned decay — must match forward kernel exactly */
if (isv_signals != NULL) {
gate *= isv_signals[11]; /* regime_stability */
if (isv_signals != NULL) gate *= isv_signals[11];
x_states_local[t + 1] = gate * x_states_local[t] + b_val;
}
if (mat == 2) {
/* d_W_C[j, s] += d_out[j] * x_states[K][s] */
acc += d_out[j] * x_states_local[K];
} else {
/* Need reverse scan for d_W_A and d_W_B */
/* Compute d_x for state dimension s via reverse scan */
float d_x_s = 0.0f;
for (int jj = 0; jj < sh2; jj++) {
d_x_s += d_out[jj] * w_c[(long long)jj * state_d + s];
}
for (int t = K - 1; t >= 0; t--) {
const float* h_t = h_history + ((long long)b * K + t) * sh2;
float a_val = a_raw_local[t];
float gate = 1.0f / (1.0f + expf(-a_val));
float sigmoid_deriv = gate * (1.0f - gate);
float stability = 1.0f;
if (isv_signals != NULL) stability = isv_signals[11];
if (mat == 0) {
/* d_W_A[j, s] += d_gate * h_t[j] */
float d_gate = d_x_s * x_states_local[t] * sigmoid_deriv * stability;
acc += d_gate * h_t[j];
} else {
/* d_W_B[j, s] += d_x_s * h_t[j] */
acc += d_x_s * h_t[j];
}
d_x_s = d_x_s * gate * stability;
}
x_states[t+1][s] = gate * x_states[t][s] + b_val;
}
}
/* ── Backward through output projection ── */
/* temporal_context[j] = sum_s(w_c[j,s] * x_K[s])
* d_x_K[s] += sum_j(d_out[j] * w_c[j,s])
* d_w_c[j,s] += d_out[j] * x_K[s] */
float d_x[16];
for (int s = 0; s < state_d; s++) {
float dx = 0.0f;
for (int j = 0; j < sh2; j++) {
dx += d_out[j] * w_c[(long long)j * state_d + s];
}
d_x[s] = dx;
}
/* d_W_C gradient */
for (int j = 0; j < sh2; j++) {
for (int s = 0; s < state_d; s++) {
atomicAdd(&d_w_c[(long long)j * state_d + s], d_out[j] * x_states[K][s]);
}
}
/* ── Reverse scan ── */
for (int t = K - 1; t >= 0; t--) {
const float* h_t = h_history + ((long long)i * K + t) * sh2;
for (int s = 0; s < state_d; s++) {
float a_val = a_raw[t][s];
float gate = 1.0f / (1.0f + expf(-a_val));
float sigmoid_deriv = gate * (1.0f - gate);
/* Regime-conditioned decay — must match forward kernel exactly.
* effective_gate = sigmoid(a_val) * stability.
* d(effective_gate)/d(a_val) = sigmoid'(a_val) * stability. */
float stability = 1.0f;
if (isv_signals != NULL) {
stability = isv_signals[11]; /* regime_stability */
}
/* d_A_t = d_x_{t+1} * x_t * sigmoid'(a_raw_t) * stability */
float d_gate = d_x[s] * x_states[t][s] * sigmoid_deriv * stability;
/* d_W_A: d_gate is scalar per (sample, state_dim), h_t is [SH2] */
for (int j = 0; j < sh2; j++) {
atomicAdd(&d_w_a[(long long)j * state_d + s], d_gate * h_t[j]);
}
/* d_W_B: d_x[s] directly (input projection gradient) */
for (int j = 0; j < sh2; j++) {
atomicAdd(&d_w_b[(long long)j * state_d + s], d_x[s] * h_t[j]);
}
/* Propagate d_x backward through effective gate:
* d_x_t = d_x_{t+1} * sigmoid(a_val) * stability */
d_x[s] = d_x[s] * gate * stability;
}
}
/* Single deterministic write per weight element */
if (mat == 0) d_w_a[(long long)j * state_d + s] = acc;
else if (mat == 1) d_w_b[(long long)j * state_d + s] = acc;
else d_w_c[(long long)j * state_d + s] = acc;
}
/* ================================================================== */
@@ -4656,7 +4685,7 @@ extern "C" __global__ void branch_independence_penalty(
* exceeds sim_threshold, penalizes mean |Q_i - Q_{i+1}| weighted
* by how far similarity exceeds the threshold.
*
* Uses atomicAdd for single-scalar accumulation .
* Uses warp+block reduction → one atomicAdd per BLOCK (deterministic) .
*
* Grid: ceil((N-1) / 256), Block: 256.
*/
@@ -4671,27 +4700,48 @@ extern "C" __global__ void temporal_consistency_penalty(
float sim_threshold
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= N - 1) return ;
float my_penalty = 0.0f ;
float dot = 0.0f, norm_a = 0.0f, norm_b = 0.0f;
for (int j = 0; j < sh2; j++) {
floa t a = h_s2[(long long)i * sh2 + j];
float b = h_s2[(long long)(i + 1) * sh2 + j];
dot += a * b ;
norm_a += a * a ;
norm_b += b * b ;
}
float sim = dot / (sqrtf(norm_a) * sqrtf(norm_b) + 1e-8f);
if (sim > sim_threshold) {
float q_diff = 0.0f;
for (int a = 0; a < total_actions; a++) {
q_diff += fabsf(q_values[(long long)i * total_actions + a]
- q_values[(long long)(i + 1) * total_actions + a]);
if (i < N - 1) {
float dot = 0.0f, norm_a = 0.0f, norm_b = 0.0f;
for (in t j = 0; j < sh2; j++) {
float a = h_s2[(long long)i * sh2 + j];
float b = h_s2[(long long)(i + 1) * sh2 + j] ;
dot += a * b ;
norm_a += a * a ;
norm_b += b * b;
}
q_diff /= (float)total_actions ;
float weight = (sim - sim_threshold) / (1.0f - sim_threshold);
atomicAdd(penalty_out, lambda_tc * q_diff * weight / (float)(N - 1));
float sim = dot / (sqrtf(norm_a) * sqrtf(norm_b) + 1e-8f) ;
if (sim > sim_threshold) {
float q_diff = 0.0f;
for (int a = 0; a < total_actions; a++) {
q_diff += fabsf(q_values[(long long)i * total_actions + a]
- q_values[(long long)(i + 1) * total_actions + a]);
}
q_diff /= (float)total_actions;
float weight = (sim - sim_threshold) / (1.0f - sim_threshold);
my_penalty = lambda_tc * q_diff * weight / (float)(N - 1);
}
}
/* Warp-level reduction */
for (int offset = 16; offset > 0; offset >>= 1)
my_penalty += __shfl_xor_sync(0xFFFFFFFF, my_penalty, offset);
/* Block-level reduction via shared memory */
__shared__ float warp_sums[8];
int warp_id = threadIdx.x / 32;
int lane = threadIdx.x % 32;
if (lane == 0) warp_sums[warp_id] = my_penalty;
__syncthreads();
if (warp_id == 0) {
float val = (lane < blockDim.x / 32) ? warp_sums[lane] : 0.0f;
for (int off = 16; off > 0; off >>= 1)
val += __shfl_xor_sync(0xFFFFFFFF, val, off);
if (lane == 0)
atomicAdd(penalty_out, val);
}
}
@@ -4708,28 +4758,50 @@ extern "C" __global__ void temporal_consistency_penalty(
*
* Total loss = lambda_pred * sum(per_sample_mse) / (N-1).
*
* Uses warp+block reduction → one atomicAdd per BLOCK (deterministic).
*
* Grid: ceil((N-1) / 256), Block: 256.
*/
extern "C" __global__ void predictive_coding_loss(
const float* __restrict__ h_s2, /* [N, SH2] enriched trunk (after Mamba2) */
float* __restrict__ loss_out, /* [1] scalar MSE loss (atomicAdd) */
float* __restrict__ loss_out, /* [1] scalar MSE loss */
int N,
int sh2,
float lambda_pred /* loss weight (0.1) */
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= N - 1) return ;
float my_loss = 0.0f ;
float mse = 0.0f;
for (int j = 0; j < sh2; j++) {
float current = h_s2[(long long)i * sh2 + j];
float next = h_s2[(long long)(i + 1) * sh2 + j];
float diff = current - next ;
mse += diff * diff ;
if (i < N - 1) {
float mse = 0.0f;
for (int j = 0; j < sh2; j++) {
float current = h_s2[(long long)i * sh2 + j];
float next = h_s2[(long long)(i + 1) * sh2 + j] ;
float diff = current - next ;
mse += diff * diff;
}
mse /= (float)sh2;
my_loss = lambda_pred * mse / (float)(N - 1);
}
mse /= (float)sh2;
atomicAdd(loss_out, lambda_pred * mse / (float)(N - 1));
/* Warp-level reduction */
for (int offset = 16; offset > 0; offset >>= 1)
my_loss += __shfl_xor_sync(0xFFFFFFFF, my_loss, offset);
/* Block-level reduction via shared memory */
__shared__ float warp_sums[8];
int warp_id = threadIdx.x / 32;
int lane = threadIdx.x % 32;
if (lane == 0) warp_sums[warp_id] = my_loss;
__syncthreads();
if (warp_id == 0) {
float val = (lane < blockDim.x / 32) ? warp_sums[lane] : 0.0f;
for (int off = 16; off > 0; off >>= 1)
val += __shfl_xor_sync(0xFFFFFFFF, val, off);
if (lane == 0)
atomicAdd(loss_out, val);
}
}
/* ================================================================== */
@@ -4847,6 +4919,14 @@ extern "C" __global__ void apply_risk_budget(
commit_lambda_buf[i] = 0.01f * (1.0f - R);
}
/**
* Risk budget backward — DETERMINISTIC: one thread per weight element.
*
* Weight layout (flat): [d_w_risk_fc: AH*(SH2+13) | d_b_risk_fc: AH | d_w_risk_out: AH | d_b_risk_out: 1]
* Total: AH*(SH2+13) + AH + AH + 1 = AH*(SH2+15) + 1
*
* Grid: ceil(total_risk_params / 256), Block: 256.
*/
extern "C" __global__ void risk_budget_backward(
const float* __restrict__ h_s2,
const float* __restrict__ isv_signals_dev_ptr, /* [12] raw ISV */
@@ -4862,38 +4942,74 @@ extern "C" __global__ void risk_budget_backward(
float* __restrict__ d_b_risk_out,
int B, int SH2, int AH, int b1_size
) {
int i = blockIdx.x * blockDim.x + threadIdx.x ;
if (i >= B) return ;
int input_dim = SH2 + 13 ;
int wfc_size = AH * input_dim ;
int bfc_size = AH;
int wout_size = AH;
int bout_size = 1;
int total_params = wfc_size + bfc_size + wout_size + bout_size;
int input_dim = SH2 + 13; /* SH2 + 12 ISV + 1 predicted_error */
float R = risk_budget[i] ;
int widx = blockIdx.x * blockDim.x + threadIdx.x;
if (widx >= total_params) return ;
float d_R = 0.0f;
const float* dq = d_q_mag + (long long)i * b1_size;
const float* qp = q_mag_pre + (long long)i * b1_size;
if (b1_size > 1) d_R += dq[1] * qp[1] * 0.5f / fmaxf(sqrtf(fmaxf(R, 1e-6f)), 1e-6f);
if (b1_size > 2) d_R += dq[2] * qp[2];
float acc = 0.0f;
float d_raw = d_R * R * (1.0f - R);
for (int i = 0; i < B; i++) {
float R = risk_budget[i];
const float* h_risk = risk_hidden + (long long)i * AH ;
atomicAdd(d_b_risk_out, d_raw) ;
for (int j = 0; j < AH; j++) {
atomicAdd(&d_w_risk_out[j], d_raw * h_risk[j] );
float d_R = 0.0f ;
const float* dq = d_q_mag + (long long)i * b1_size ;
const float* qp = q_mag_pre + (long long)i * b1_size;
if (b1_size > 1) d_R += dq[1] * qp[1] * 0.5f / fmaxf(sqrtf(fmaxf(R, 1e-6f)), 1e-6f );
if (b1_size > 2) d_R += dq[2] * qp[2];
float d_raw = d_R * R * (1.0f - R);
if (widx < wfc_size) {
/* d_w_risk_fc[j * input_dim + k] */
int j = widx / input_dim;
int k = widx % input_dim;
float d_hj = d_raw * w_risk_out[j];
const float* h_risk = risk_hidden + (long long)i * AH;
if (h_risk[j] <= 0.0f) d_hj = 0.0f;
float input_val;
if (k < SH2) {
input_val = h_s2[(long long)i * SH2 + k];
} else if (k < SH2 + 12) {
input_val = isv_signals_dev_ptr[k - SH2];
} else {
input_val = predicted_error[i];
}
acc += d_hj * input_val;
} else if (widx < wfc_size + bfc_size) {
/* d_b_risk_fc[j] */
int j = widx - wfc_size;
float d_hj = d_raw * w_risk_out[j];
const float* h_risk = risk_hidden + (long long)i * AH;
if (h_risk[j] <= 0.0f) d_hj = 0.0f;
acc += d_hj;
} else if (widx < wfc_size + bfc_size + wout_size) {
/* d_w_risk_out[j] */
int j = widx - wfc_size - bfc_size;
const float* h_risk = risk_hidden + (long long)i * AH;
acc += d_raw * h_risk[j];
} else {
/* d_b_risk_out */
acc += d_raw;
}
}
const float* h = h_s2 + (long long)i * SH2;
for (int j = 0; j < AH; j++) {
float d_hj = d_raw * w_risk_out[j] ;
if (h_risk[j] <= 0.0f) d_hj = 0.0f;
atomicAdd(&d_b_risk_fc[j], d_hj);
for (int k = 0; k < SH2; k++)
atomicAdd(&d_w_risk_fc[(long long)j * input_dim + k], d_hj * h[k]);
for (int k = 0; k < 12; k++)
atomicAdd(&d_w_risk_fc[(long long)j * input_dim + SH2 + k], d_hj * isv_signals_dev_ptr[k]);
atomicAdd(&d_w_risk_fc[(long long)j * input_dim + SH2 + 12], d_hj * predicted_error[i]);
}
/* Single deterministic write per weight element */
if (widx < wfc_size)
d_w_risk_fc[widx] = acc ;
else if (widx < wfc_size + bfc_size)
d_b_risk_fc[widx - wfc_size] = acc;
else if (widx < wfc_size + bfc_size + wout_size)
d_w_risk_out[widx - wfc_size - bfc_size] = acc;
else
d_b_risk_out[0] = acc;
}
/* ================================================================== */
@@ -5147,6 +5263,13 @@ extern "C" __global__ void recursive_confidence_forward(
/* ================================================================== */
/* Kernel: recursive_confidence_backward — MSE grad into trunk */
/* ================================================================== */
/**
* Phase 1: Per-sample trunk gradient + warp-reduced weight gradient.
* d_h_s2[i, k] += d_sigmoid * w_conf[k] (plain write — single writer per (i,k)).
* d_w_conf and d_b_conf use warp+block reduction → one atomicAdd per BLOCK.
*
* Grid: ceil(B/256), Block: 256.
*/
extern "C" __global__ void recursive_confidence_backward(
const float* __restrict__ h_s2, /* [B, SH2] */
const float* __restrict__ predicted_error, /* [B] */
@@ -5159,18 +5282,56 @@ extern "C" __global__ void recursive_confidence_backward(
float loss_weight /* 0.01 */
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= B) return;
float pred = predicted_error[i] ;
float target = lagged_td_error[0];
float d_loss = loss_weight * 2.0f * (pred - target) / (float)B ;
float d_sigmoid = d_loss * pred * (1.0f - pred) ;
float d_sigmoid = 0.0f ;
if (i < B) {
float pred = predicted_error[i] ;
float target = lagged_td_error[0] ;
float d_loss = loss_weight * 2.0f * (pred - target) / (float)B;
d_sigmoid = d_loss * pred * (1.0f - pred);
const float* h = h_s2 + (long long)i * SH2;
atomicAdd(d_b_conf, d_sigmoid);
/* Per-sample trunk gradient — no cross-sample contention, plain write */
for (int k = 0; k < SH2; k++) {
d_h_s2[(long long)i * SH2 + k] += d_sigmoid * w_conf[k];
}
}
/* Weight gradient reduction: d_b_conf = sum_b(d_sigmoid_b) */
float my_db = d_sigmoid;
for (int offset = 16; offset > 0; offset >>= 1)
my_db += __shfl_xor_sync(0xFFFFFFFF, my_db, offset);
__shared__ float ws_db[8];
int warp_id = threadIdx.x / 32;
int lane = threadIdx.x % 32;
if (lane == 0) ws_db[warp_id] = my_db;
__syncthreads();
if (warp_id == 0) {
float val = (lane < blockDim.x / 32) ? ws_db[lane] : 0.0f;
for (int off = 16; off > 0; off >>= 1)
val += __shfl_xor_sync(0xFFFFFFFF, val, off);
if (lane == 0) atomicAdd(d_b_conf, val);
}
/* Weight gradient reduction: d_w_conf[k] = sum_b(d_sigmoid_b * h_s2[b, k]) */
for (int k = 0; k < SH2; k++) {
atomicAdd(&d_w_conf[k], d_sigmoid * h[k]) ;
atomicAdd(&d_h_s2[(long long)i * SH2 + k], d_sigmoid * w_conf[k]);
float my_dw = (i < B) ? d_sigmoid * h_s2[(long long)i * SH2 + k] : 0.0f ;
for (int offset = 16; offset > 0; offset >>= 1)
my_dw += __shfl_xor_sync(0xFFFFFFFF, my_dw, offset);
__shared__ float ws_dw[8];
if (lane == 0) ws_dw[warp_id] = my_dw;
__syncthreads();
if (warp_id == 0) {
float val = (lane < blockDim.x / 32) ? ws_dw[lane] : 0.0f;
for (int off = 16; off > 0; off >>= 1)
val += __shfl_xor_sync(0xFFFFFFFF, val, off);
if (lane == 0) atomicAdd(&d_w_conf[k], val);
}
__syncthreads();
}
}
jgrusewski