diff --git a/crates/ml/src/cuda_pipeline/decision_transformer.rs b/crates/ml/src/cuda_pipeline/decision_transformer.rs index b6897c562..b57f8df7e 100644 --- a/crates/ml/src/cuda_pipeline/decision_transformer.rs +++ b/crates/ml/src/cuda_pipeline/decision_transformer.rs @@ -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, k_buf: CudaSlice, v_buf: CudaSlice, + /// Per-head attention output: [B, num_heads, T, E] + per_head_out: CudaSlice, /// Attention output (pre-LN): [B, T, E] attn_out: CudaSlice, /// 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, MLError> { stream.alloc_zeros::(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}")))?; } } diff --git a/crates/ml/src/cuda_pipeline/dt_kernels.cu b/crates/ml/src/cuda_pipeline/dt_kernels.cu index c11997561..0aa6cb71a 100644 --- a/crates/ml/src/cuda_pipeline/dt_kernels.cu +++ b/crates/ml/src/cuda_pipeline/dt_kernels.cu @@ -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; + } } }