gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: DT kernels — eliminate atomicAdd, deterministic gradient reduction

Browse Source 
Three atomicAdd anti-patterns removed from Decision Transformer kernels:

1. dt_linear_backward_kernel (critical): Split into dt_linear_backward_kernel
   (input gradient only, per-sample deterministic) + dt_linear_grad_kernel
   (one thread per weight, loop over samples — zero atomics).

2. dt_causal_attention_kernel: Replace cross-head atomicAdd output projection
   with per-head output buffer [B, H, T, E] + separate dt_sum_heads_kernel
   that sums heads in fixed order.

3. dt_cross_entropy_kernel: Remove atomicAdd for total_loss, add
   dt_reduce_loss_kernel (single-thread sequential sum for full determinism).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-18 10:05:05 +02:00
parent 020e9ba460
commit 3065ae1f3b
2 changed files with 238 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files

137
crates/ml/src/cuda_pipeline/decision_transformer.rs
View File

@@ -246,12 +246,15 @@ struct DtKernels {
embed: CudaFunction,
qkv_projection: CudaFunction,
causal_attention: CudaFunction,
sum_heads: CudaFunction,
layernorm: CudaFunction,
ffn: CudaFunction,
action_head: CudaFunction,
cross_entropy: CudaFunction,
reduce_loss: CudaFunction,
ce_backward: CudaFunction,
linear_backward: CudaFunction,
linear_grad: CudaFunction,
zero: CudaFunction,
residual_add: CudaFunction,
return_to_go: CudaFunction,
@@ -279,12 +282,15 @@ fn compile_dt_kernels(
embed: load("dt_embed_kernel")?,
qkv_projection: load("dt_qkv_projection_kernel")?,
causal_attention: load("dt_causal_attention_kernel")?,
sum_heads: load("dt_sum_heads_kernel")?,
layernorm: load("dt_layernorm_kernel")?,
ffn: load("dt_ffn_kernel")?,
action_head: load("dt_action_head_kernel")?,
cross_entropy: load("dt_cross_entropy_kernel")?,
reduce_loss: load("dt_reduce_loss_kernel")?,
ce_backward: load("dt_ce_backward_kernel")?,
linear_backward: load("dt_linear_backward_kernel")?,
linear_grad: load("dt_linear_grad_kernel")?,
zero: load("dt_zero_kernel")?,
residual_add: load("dt_residual_add_kernel")?,
return_to_go: load("dt_return_to_go_kernel")?,
@@ -302,6 +308,8 @@ struct DtScratch {
q_buf: CudaSlice<f32>,
k_buf: CudaSlice<f32>,
v_buf: CudaSlice<f32>,
/// Per-head attention output: [B, num_heads, T, E]
per_head_out: CudaSlice<f32>,
/// Attention output (pre-LN): [B, T, E]
attn_out: CudaSlice<f32>,
/// Post-LN1 output: [B, T, E]
@@ -343,6 +351,7 @@ fn alloc_dt_scratch(
let bt = config.batch_size * config.context_len;
let bt_input = config.batch_size * config.context_len * config.input_dim();
let total_params = config.total_params();
let per_head = config.batch_size * config.num_heads * config.context_len * config.embed_dim;
let alloc_f = |n: usize, name: &str| -> Result<CudaSlice<f32>, MLError> {
stream.alloc_zeros::<f32>(n)
@@ -354,6 +363,7 @@ fn alloc_dt_scratch(
q_buf: alloc_f(bte, "q_buf")?,
k_buf: alloc_f(bte, "k_buf")?,
v_buf: alloc_f(bte, "v_buf")?,
per_head_out: alloc_f(per_head, "per_head_out")?,
attn_out: alloc_f(bte, "attn_out")?,
ln1_out: alloc_f(bte, "ln1_out")?,
ffn_out: alloc_f(bte, "ffn_out")?,
@@ -578,6 +588,7 @@ impl DecisionTransformer {
let q_ptr = raw_ptr_f32_mut(&mut scratch.q_buf, stream);
let k_ptr = raw_ptr_f32_mut(&mut scratch.k_buf, stream);
let v_ptr = raw_ptr_f32_mut(&mut scratch.v_buf, stream);
let per_head_out_ptr = raw_ptr_f32_mut(&mut scratch.per_head_out, stream);
let attn_ptr = raw_ptr_f32_mut(&mut scratch.attn_out, stream);
let ln1_ptr = raw_ptr_f32_mut(&mut scratch.ln1_out, stream);
let ffn_ptr = raw_ptr_f32_mut(&mut scratch.ffn_out, stream);
@@ -633,10 +644,9 @@ impl DecisionTransformer {
}
}
// Causal attention with output projection
// Causal attention: per-head projected outputs
{
let w_o = params_ptr + byte_off(lb + lo.w_o);
let b_o = params_ptr + byte_off(lb + lo.b_o);
let dh = e / self.config.num_heads;
let shmem = (t * dh * 2 * 4) as u32; // K + V in shmem
let lc = LaunchConfig {
@@ -646,14 +656,31 @@ impl DecisionTransformer {
};
unsafe {
stream.launch_builder(&kernels.causal_attention)
.arg(&q_ptr).arg(&k_ptr).arg(&v_ptr).arg(&w_o).arg(&b_o)
.arg(&attn_ptr)
.arg(&q_ptr).arg(&k_ptr).arg(&v_ptr).arg(&w_o)
.arg(&per_head_out_ptr)
.arg(&b_i32).arg(&t_i32).arg(&e_i32).arg(&num_heads_i32)
.launch(lc)
.map_err(|e| MLError::ModelError(format!("DT attn L{layer_idx}: {e}")))?;
}
}
// Sum per-head outputs into attn_out (deterministic — fixed head order)
{
let b_o = params_ptr + byte_off(lb + lo.b_o);
let lc = LaunchConfig {
grid_dim: (bt_grid, 1, 1),
block_dim: (embed_block, 1, 1),
shared_mem_bytes: 0,
};
unsafe {
stream.launch_builder(&kernels.sum_heads)
.arg(&per_head_out_ptr).arg(&b_o).arg(&attn_ptr)
.arg(&b_i32).arg(&t_i32).arg(&e_i32).arg(&num_heads_i32)
.launch(lc)
.map_err(|e| MLError::ModelError(format!("DT sum_heads L{layer_idx}: {e}")))?;
}
}
// Residual add: attn_out += layer_input
{
let n_elems = (b * t * e) as i32;
@@ -730,28 +757,27 @@ impl DecisionTransformer {
}
}
// 5. Cross-entropy loss
// 5. Cross-entropy loss (per-sample) + deterministic reduction
{
// Zero total_loss
let one = 1_i32;
let lc_z = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1, 1, 1), shared_mem_bytes: 0 };
unsafe {
stream.launch_builder(&kernels.zero)
.arg(&total_loss_ptr).arg(&one)
.launch(lc_z)
.map_err(|e| MLError::ModelError(format!("DT zero loss: {e}")))?;
}
let action_block = a.min(256).max(1) as u32;
let lc = LaunchConfig { grid_dim: (bt_grid, 1, 1), block_dim: (action_block, 1, 1), shared_mem_bytes: 0 };
unsafe {
stream.launch_builder(&kernels.cross_entropy)
.arg(&logits_ptr).arg(&target_ptr)
.arg(&psl_ptr).arg(&total_loss_ptr)
.arg(&psl_ptr)
.arg(&bt).arg(&a_i32)
.launch(lc)
.map_err(|e| MLError::ModelError(format!("DT CE loss: {e}")))?;
}
// Deterministic sequential reduction of per-sample losses
let lc_r = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1, 1, 1), shared_mem_bytes: 0 };
unsafe {
stream.launch_builder(&kernels.reduce_loss)
.arg(&psl_ptr).arg(&total_loss_ptr).arg(&bt)
.launch(lc_r)
.map_err(|e| MLError::ModelError(format!("DT reduce loss: {e}")))?;
}
}
// ── BACKWARD PASS (Phase 1 — simplified) ──────────────────────
@@ -783,28 +809,50 @@ impl DecisionTransformer {
}
}
// Action head backward: d_logits -> d_layer
// Action head backward: d_logits -> d_layer (input gradient only)
{
let w_h = params_ptr + byte_off(self.layout.w_head_off);
let dw_h = d_params_ptr + byte_off(self.layout.w_head_off);
let db_h = d_params_ptr + byte_off(self.layout.b_head_off);
let io_max = e.max(a).min(256) as u32;
let lc = LaunchConfig { grid_dim: (bt_grid, 1, 1), block_dim: (io_max, 1, 1), shared_mem_bytes: 0 };
// Input to action head was ln2_out (or embedded if 0 layers)
let head_input = if num_layers > 0 { ln2_ptr } else { emb_ptr };
// Step 1: Input gradient (per-sample, deterministic by construction)
let lc_dx = LaunchConfig {
grid_dim: (bt_grid, 1, 1),
block_dim: (e.min(256) as u32, 1, 1),
shared_mem_bytes: 0,
};
unsafe {
stream.launch_builder(&kernels.linear_backward)
.arg(&d_logits_ptr) // d_output [N, A]
.arg(&head_input) // input [N, E]
.arg(&head_input) // input [N, E] (unused by kernel)
.arg(&w_h) // W [E, A]
.arg(&dw_h) // dW
.arg(&db_h) // db
.arg(&dw_h) // dW (unused by kernel)
.arg(&db_h) // db (unused by kernel)
.arg(&d_layer_ptr) // d_input [N, E]
.arg(&bt).arg(&e_i32).arg(&a_i32)
.launch(lc)
.map_err(|e| MLError::ModelError(format!("DT head bw: {e}")))?;
.launch(lc_dx)
.map_err(|e| MLError::ModelError(format!("DT head bw dx: {e}")))?;
}
// Step 2: Weight + bias gradients (one thread per weight, deterministic)
let total_head_weights = (e * a) as i32;
let lc_dw = LaunchConfig {
grid_dim: ((total_head_weights as u32 + 255) / 256, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
};
unsafe {
stream.launch_builder(&kernels.linear_grad)
.arg(&d_logits_ptr) // d_output [N, A]
.arg(&head_input) // input [N, E]
.arg(&dw_h) // dW [E, A]
.arg(&db_h) // db [A]
.arg(&bt).arg(&e_i32).arg(&a_i32)
.launch(lc_dw)
.map_err(|e| MLError::ModelError(format!("DT head bw dw: {e}")))?;
}
}
@@ -813,24 +861,47 @@ impl DecisionTransformer {
// The gradient passes through unchanged via residual connections.
// No parameter gradients for transformer layers in Phase 1.
// Embed backward: d_layer -> dW_embed, db_embed
// Embed backward: d_layer -> dW_embed, db_embed (deterministic)
{
let w_e = params_ptr + byte_off(self.layout.w_embed_off);
let dw_e = d_params_ptr + byte_off(self.layout.w_embed_off);
let db_e = d_params_ptr + byte_off(self.layout.b_embed_off);
let io_max = e.max(input_dim).min(256) as u32;
let lc = LaunchConfig { grid_dim: (bt_grid, 1, 1), block_dim: (io_max, 1, 1), shared_mem_bytes: 0 };
// Step 1: Input gradient (per-sample)
let lc_dx = LaunchConfig {
grid_dim: (bt_grid, 1, 1),
block_dim: (input_dim.min(256) as u32, 1, 1),
shared_mem_bytes: 0,
};
unsafe {
stream.launch_builder(&kernels.linear_backward)
.arg(&d_layer_ptr) // d_output [N, E]
.arg(&traj_ptr) // input [N, input_dim]
.arg(&traj_ptr) // input [N, input_dim] (unused by kernel)
.arg(&w_e) // W [input_dim, E]
.arg(&dw_e) // dW
.arg(&db_e) // db
.arg(&d_embed_ptr) // d_input [N, input_dim] (unused but kernel writes it)
.arg(&dw_e) // dW (unused by kernel)
.arg(&db_e) // db (unused by kernel)
.arg(&d_embed_ptr) // d_input [N, input_dim]
.arg(&bt).arg(&input_dim_i32).arg(&e_i32)
.launch(lc)
.map_err(|e| MLError::ModelError(format!("DT embed bw: {e}")))?;
.launch(lc_dx)
.map_err(|e| MLError::ModelError(format!("DT embed bw dx: {e}")))?;
}
// Step 2: Weight + bias gradients (one thread per weight, deterministic)
let total_embed_weights = (input_dim * e) as i32;
let lc_dw = LaunchConfig {
grid_dim: ((total_embed_weights as u32 + 255) / 256, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
};
unsafe {
stream.launch_builder(&kernels.linear_grad)
.arg(&d_layer_ptr) // d_output [N, E]
.arg(&traj_ptr) // input [N, input_dim]
.arg(&dw_e) // dW [input_dim, E]
.arg(&db_e) // db [E]
.arg(&bt).arg(&input_dim_i32).arg(&e_i32)
.launch(lc_dw)
.map_err(|e| MLError::ModelError(format!("DT embed bw dw: {e}")))?;
}
}

200
crates/ml/src/cuda_pipeline/dt_kernels.cu
View File

@@ -130,8 +130,7 @@ extern "C" __global__ void dt_causal_attention_kernel(
const float* __restrict__ K, /* [B, T, E] */
const float* __restrict__ V, /* [B, T, E] */
const float* __restrict__ W_O, /* [E, E] output projection */
const float* __restrict__ b_O, /* [E] output bias */
float* __restrict__ output, /* [B, T, E] */
float* __restrict__ per_head_out, /* [B, num_heads, T, E] per-head projected output */
int B, int T, int E, int num_heads
) {
int b = blockIdx.x;
@@ -198,43 +197,54 @@ extern "C" __global__ void dt_causal_attention_kernel(
attn_out[dd] = attn_out[dd] * inv_sum;
}
/* Store per-head attention output into temp shared buffer for output projection.
* We use shmem after the sync point (K/V no longer needed for this thread). */
__syncthreads();
/* Write per-head attention output to a per-head buffer.
* A separate dt_sum_heads_kernel will combine all heads deterministically. */
/* Reuse shmem for the concatenated head output: [T, E] would be too much.
* Instead, each thread writes its head's output, then we project per-element. */
/* Write per-head result into a shared buffer [T, Dh] */
float* sh_head_out = shmem; /* reuse: [T, Dh] */
for (int dd = 0; dd < Dh; dd++) {
sh_head_out[t_idx * Dh + dd] = attn_out[dd];
}
__syncthreads();
/* per_head_out layout: [B, num_heads, T, E] — each head writes its
* projected output to its own slice, avoiding cross-head atomicAdd. */
float* head_slice = per_head_out + ((b * num_heads + h) * T + t_idx) * E;
/* Output projection: for this head's contribution to output[b][t_idx][d]
* Each head contributes: sum_dh(head_out[t_idx][dh] * W_O[(h*Dh+dh)*E + d])
* Since other heads aren't available here, we use atomicAdd on the output. */
float* out = output + (b * T + t_idx) * E;
/* Zero output on first head */
if (h == 0) {
for (int dd = 0; dd < E; dd++) {
out[dd] = b_O[dd]; /* Initialize with bias */
}
}
__syncthreads(); /* Ensure bias initialization is visible */
/* Each head's projection contribution via atomicAdd.
* DETERMINISM NOTE: num_heads blocks (typically 4) accumulate into
* the same output element. With only 4 concurrent writers per element,
* non-determinism is minimal. Full elimination would require all heads
* in one block (changing grid layout). DT is a separate model. */
/* Output projection: head_out[dh] -> projected[d] via W_O
* out[d] = sum_dh(attn_out[dh] * W_O[(h*Dh+dh)*E + d]) */
for (int dd = 0; dd < E; dd++) {
float proj = 0.0f;
for (int dh = 0; dh < Dh; dh++) {
proj = proj + sh_head_out[t_idx * Dh + dh] * W_O[(head_off + dh) * E + dd];
proj = proj + attn_out[dh] * W_O[(head_off + dh) * E + dd];
}
atomicAdd(&out[dd], proj);
head_slice[dd] = proj;
}
}
/* ══════════════════════════════════════════════════════════════════════
* KERNEL 3b: SUM PER-HEAD ATTENTION OUTPUTS
*
* Sums the per-head projected outputs into the final attention output
* and adds the output bias. Fully deterministic — fixed summation order.
*
* output[b][t][d] = b_O[d] + sum_h(per_head_out[b][h][t][d])
*
* Grid: (B*T, 1, 1)
* Block: (min(E, 256), 1, 1)
* ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_sum_heads_kernel(
const float* __restrict__ per_head_out, /* [B, num_heads, T, E] */
const float* __restrict__ b_O, /* [E] output bias */
float* __restrict__ output, /* [B, T, E] */
int B, int T, int E, int num_heads
) {
int bt = blockIdx.x;
if (bt >= B * T) return;
int d = threadIdx.x;
int b = bt / T;
int t = bt % T;
for (int dd = d; dd < E; dd += blockDim.x) {
float sum = b_O[dd];
for (int h = 0; h < num_heads; h++) {
sum += per_head_out[((b * num_heads + h) * T + t) * E + dd];
}
output[bt * E + dd] = sum;
}
}
@@ -440,7 +450,6 @@ extern "C" __global__ void dt_cross_entropy_kernel(
const float* __restrict__ logits, /* [B*T, A] */
const int* __restrict__ targets, /* [B*T] */
float* __restrict__ per_sample_loss, /* [B*T] */
float* __restrict__ total_loss, /* [1] atomicAdd */
int N, /* B*T */
int A /* num_actions */
) {
@@ -471,11 +480,30 @@ extern "C" __global__ void dt_cross_entropy_kernel(
loss = fminf(fmaxf(loss, 0.0f), 100.0f);
per_sample_loss[n] = loss;
/* total_loss accumulated via warp+block reduction — one atomicAdd per BLOCK.
* Since grid=(B*T) with 1 thread per block, each block has 1 thread, so
* atomicAdd is from exactly N sources with fixed order. For large N, use
* a separate reduction kernel for full determinism. */
atomicAdd(total_loss, loss / (float)N);
/* Total loss is computed by a separate dt_reduce_loss_kernel for
* full determinism — no atomicAdd across samples. */
}
/* ══════════════════════════════════════════════════════════════════════
* KERNEL 7b: REDUCE PER-SAMPLE LOSS TO TOTAL LOSS
*
* Deterministic sequential sum of per-sample losses to compute mean.
* Single-thread kernel — N is small (B*T, typically 256*20=5120) so
* the serial sum is fast and perfectly reproducible.
*
* Grid: (1, 1, 1)
* Block: (1, 1, 1)
* ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_reduce_loss_kernel(
const float* __restrict__ per_sample_loss, /* [N] */
float* __restrict__ total_loss, /* [1] */
int N
) {
float sum = 0.0f;
for (int i = 0; i < N; i++) {
sum += per_sample_loss[i];
}
total_loss[0] = sum / (float)N;
}
/* ══════════════════════════════════════════════════════════════════════
@@ -527,24 +555,23 @@ extern "C" __global__ void dt_ce_backward_kernel(
}
/* ══════════════════════════════════════════════════════════════════════
* KERNEL 9: LINEAR BACKWARD
* KERNEL 9: LINEAR BACKWARD — INPUT GRADIENT
*
* Given d_output[B*T, O] and input[B*T, I] and W[I, O]:
* dW[i][o] += sum_n(input[n][i] * d_output[n][o]) (weight gradient)
* db[o] += sum_n(d_output[n][o]) (bias gradient)
* d_input[n][i] = sum_o(d_output[n][o] * W[i][o]) (input gradient)
* Computes the input gradient for a linear layer:
* d_input[n][i] = sum_o(d_output[n][o] * W[i][o])
*
* Grid: (B*T, 1, 1)
* Block: (min(max(I,O), 256), 1, 1)
* This is per-sample with no cross-sample accumulation — fully
* deterministic by construction.
*
* Note: weight gradients use atomicAdd across samples.
* Grid: (N, 1, 1)
* Block: (min(I, 256), 1, 1)
* ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_linear_backward_kernel(
const float* __restrict__ d_output, /* [N, O] */
const float* __restrict__ input, /* [N, I] */
const float* __restrict__ input, /* [N, I] (unused — kept for API compat) */
const float* __restrict__ W, /* [I, O] */
float* __restrict__ dW, /* [I, O] atomicAdd */
float* __restrict__ db, /* [O] atomicAdd */
float* __restrict__ dW, /* [I, O] (unused — computed by dt_linear_grad_kernel) */
float* __restrict__ db, /* [O] (unused — computed by dt_linear_grad_kernel) */
float* __restrict__ d_input, /* [N, I] */
int N,
int I,
@@ -555,10 +582,9 @@ extern "C" __global__ void dt_linear_backward_kernel(
int tid = threadIdx.x;
const float* dout = d_output + n * O;
const float* x = input + n * I;
float* dx = d_input + n * I;
/* Compute d_input[n][i] = sum_o(d_output[n][o] * W[i][o]) */
/* d_input[n][i] = sum_o(d_output[n][o] * W[i][o]) */
for (int i = tid; i < I; i += blockDim.x) {
float val = 0.0f;
for (int o = 0; o < O; o++) {
@@ -566,23 +592,65 @@ extern "C" __global__ void dt_linear_backward_kernel(
}
dx[i] = val;
}
}
/* Accumulate weight + bias gradients via warp-reduced atomicAdd.
* Each sample's contribution is pre-reduced within the warp before
* writing, yielding one atomicAdd per warp per weight instead of
* one per thread.
* DETERMINISM NOTE: Decision Transformer is a separate model from
* the main DQN; residual cross-block atomicAdd is acceptable. */
for (int i = tid; i < I; i += blockDim.x) {
float xi = x[i];
for (int o = 0; o < O; o++) {
atomicAdd(&dW[i * O + o], xi * dout[o]);
/* ══════════════════════════════════════════════════════════════════════
* KERNEL 9b: LINEAR BACKWARD — WEIGHT + BIAS GRADIENTS (DETERMINISTIC)
*
* One thread per weight parameter. Each thread loops over all N samples
* and accumulates the gradient contribution in a register — fully
* deterministic with zero atomics.
*
* dW[i][o] = sum_n(input[n][i] * d_output[n][o])
* db[o] = sum_n(d_output[n][o])
*
* Weight grid: (ceil(I*O / 256), 1, 1)
* Weight block: (256, 1, 1)
*
* Bias is computed by threads in the first ceil(O/256) blocks where
* the thread's linear index maps to a valid bias element (idx < O).
* Since I*O >> O, we piggyback the bias computation onto the weight
* kernel: thread idx maps to a (i, o) pair, and if i==0 we also
* compute db[o]. This avoids a separate kernel launch.
*
* For DT dimensions (E=128, A=9 -> I*O = 1152 or E=128, input_dim=74
* -> I*O = 9472), this fits in a handful of blocks. The loop over N
* (N = B*T, typically 5120) is the inner loop — cache-friendly on
* d_output (sequential reads of O floats per sample).
* ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_linear_grad_kernel(
const float* __restrict__ d_output, /* [N, O] */
const float* __restrict__ input, /* [N, I] */
float* __restrict__ dW, /* [I, O] */
float* __restrict__ db, /* [O] */
int N,
int I,
int O
) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total_weights = I * O;
/* ── Weight gradient: one thread per (i, o) pair ── */
if (idx < total_weights) {
int i = idx / O;
int o = idx % O;
float sum_w = 0.0f;
for (int n = 0; n < N; n++) {
sum_w += input[n * I + i] * d_output[n * O + o];
}
}
dW[idx] = sum_w;
/* Bias gradient */
for (int o = tid; o < O; o += blockDim.x) {
atomicAdd(&db[o], dout[o]);
/* ── Bias gradient: piggyback on i==0 threads ──
* When i==0, this thread also computes db[o] = sum_n(d_output[n][o]).
* Since O is small, all bias elements are covered by the first O threads. */
if (i == 0) {
float sum_b = 0.0f;
for (int n = 0; n < N; n++) {
sum_b += d_output[n * O + o];
}
db[o] = sum_b;
}
}
}