perf: curiosity training — cuBLAS GEMMs (1085ms → <5ms)

Replace the serial curiosity_fwd_bwd_per_block kernel (CUR_TOTAL_PARAMS=11306
loop iterations × block-level shared-memory reduction per-iteration = 1085ms)
with a cuBLAS GEMM pipeline matching the existing curiosity inference path:

Forward: curiosity_prepare_input → GEMM1(W1) → bias_leaky_relu → GEMM2(W2) → mse_fwd_grad
Backward: gemm_dw(dW2) → bias_grad_reduce(db2) → gemm_dx(d_hidden) → leaky_relu_bwd → gemm_dw(dW1) → bias_grad_reduce(db1)

New CUDA kernels added to curiosity_training_kernel.cu:
  - curiosity_mse_fwd_grad: +b2 in-place, d_pred = 2/CUR_OUTPUT*(pred-target)
  - curiosity_leaky_relu_bwd: gates d_hidden by sign of post-activation hidden
  - curiosity_bias_grad_reduce: sum dy[N, D] over batch → grad_b[D]

GpuCuriosityTrainer rewritten with CuriosityGemm (dedicated cuBLAS+cublasLt
handle) + intermediate buffers (input_buf, hidden_buf, pred_buf, d_hidden_buf).
Reuses forward kernels from curiosity_inference_kernel.cu. Keeps curiosity_adam_step.
Drops partial_grads buffer (max_blocks*11306 floats saved).

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu

@@ -1,16 +1,31 @@
 /**
- * GPU-resident curiosity forward model training kernel.
+ * GPU-resident curiosity forward model training kernels.
  *
- * Trains the curiosity forward model (2-layer MLP) entirely on GPU with zero
- * CPU involvement. Experience data (states, actions, next_states) is already
- * on GPU from the experience collector. Weights are updated in-place via Adam.
+ * cuBLAS GEMM pipeline (replaces the 1085ms curiosity_fwd_bwd_per_block kernel):
+ *
+ *   Forward:
+ *     1. curiosity_shift_states     — build next_states from states[+1]
+ *     2. curiosity_prepare_input    — build input[N, CUR_INPUT] (from inference cubin)
+ *     3. cuBLAS GEMM1               — hidden[N, CUR_HIDDEN] = input @ W1^T
+ *     4. curiosity_bias_leaky_relu  — in-place +b1, LeakyReLU (from inference cubin)
+ *     5. cuBLAS GEMM2               — pred[N, CUR_OUTPUT] = hidden @ W2^T
+ *     6. curiosity_mse_fwd_grad     — +b2, compute d_pred = 2/N*(pred+b2-target)
+ *
+ *   Backward:
+ *     7. cuBLAS dW2                 — grad_w2 += d_pred^T @ hidden
+ *     8. curiosity_bias_grad_reduce — grad_b2 = sum_batch(d_pred)
+ *     9. cuBLAS d_hidden            — d_hidden = d_pred @ W2
+ *    10. curiosity_leaky_relu_bwd   — mask d_hidden by post-activation sign
+ *    11. cuBLAS dW1                 — grad_w1 += d_hidden^T @ input
+ *    12. curiosity_bias_grad_reduce — grad_b1 = sum_batch(d_hidden)
+ *    13. curiosity_adam_step        — Adam update per parameter group
  *
  * Architecture: [CUR_INPUT=MARKET_DIM+3] -> [CUR_HIDDEN=128] LeakyReLU(0.01) -> [CUR_OUTPUT=MARKET_DIM]
- * Total params: 128*(MARKET_DIM+3) + 128 + MARKET_DIM*128 + MARKET_DIM
+ * Total params: 128*(MARKET_DIM+3) + 128 + MARKET_DIM*128 + MARKET_DIM = 11306
  *
  * Requires common_device_functions.cuh to be prepended for CUR_INPUT/CUR_HIDDEN/CUR_OUTPUT.
  *
- * Native float everywhere. Fully deterministic (no atomicAdd).
+ * Native float everywhere.
  */
 
 /* Total trainable parameters */
@@ -327,3 +342,97 @@ extern "C" __global__ void curiosity_adam_step(
     float v_hat = v[i] / (1.0f - powf(beta2, (float)step));
     params[i] = params[i] - lr_bf * m_hat / (sqrtf(v_hat) + eps_bf);
 }
+
+/* ------------------------------------------------------------------ */
+/*  Kernel 5: MSE forward + d_pred computation (cuBLAS training path) */
+/* ------------------------------------------------------------------ */
+
+/**
+ * One thread per element in pred[N, CUR_OUTPUT].
+ *
+ * 1. Adds b2 in-place to pred (broadcast along batch).
+ * 2. Computes d_pred[n, o] = (2.0 / CUR_OUTPUT) * (pred[n,o] - target[n,o]).
+ *    The 1/N mean is applied later by curiosity_adam_step (divides by batch_size).
+ *
+ * cuBLAS GEMM 2 writes pred[N, CUR_OUTPUT] without bias. We add bias here
+ * and compute the MSE gradient for the backward pass.
+ *
+ * Launch: grid=(ceil(N*CUR_OUTPUT/256)), block=(256)
+ */
+extern "C" __global__ void curiosity_mse_fwd_grad(
+    float*       __restrict__ pred,       /* [N, CUR_OUTPUT] — in-place +b2, then d_pred out */
+    const float* __restrict__ b2,         /* [CUR_OUTPUT]                                    */
+    const float* __restrict__ target,     /* [N, state_dim] (uses first CUR_OUTPUT elements) */
+    int N,
+    int state_dim
+) {
+    int total = N * CUR_OUTPUT;
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i >= total) return;
+
+    int n = i / CUR_OUTPUT;   /* sample index */
+    int o = i % CUR_OUTPUT;   /* output index */
+
+    float p = pred[i] + b2[o];
+    float t = target[n * state_dim + o];
+    float d = (2.0f / (float)CUR_OUTPUT) * (p - t);
+
+    /* Write back d_pred; pred buffer is reused as gradient buffer */
+    pred[i] = d;
+}
+
+/* ------------------------------------------------------------------ */
+/*  Kernel 6: LeakyReLU backward (in-place on d_hidden)               */
+/* ------------------------------------------------------------------ */
+
+/**
+ * One thread per element in d_hidden[N, CUR_HIDDEN].
+ *
+ * Gates d_hidden by the sign of the post-activation hidden values:
+ *   d_hidden[i] *= (hidden_post[i] > 0) ? 1.0 : 0.01
+ *
+ * This is valid because LeakyReLU preserves the sign of the input:
+ * f(x) > 0 iff x > 0, so the mask is equivalent to (pre_act > 0).
+ *
+ * Launch: grid=(ceil(N*CUR_HIDDEN/256)), block=(256)
+ */
+extern "C" __global__ void curiosity_leaky_relu_bwd(
+    float*       __restrict__ d_hidden,    /* [N, CUR_HIDDEN] — modified in-place */
+    const float* __restrict__ hidden_post, /* [N, CUR_HIDDEN] — post-activation   */
+    int total_elements                     /* N * CUR_HIDDEN                       */
+) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i >= total_elements) return;
+    float mask = (hidden_post[i] > 0.0f) ? 1.0f : 0.01f;
+    d_hidden[i] *= mask;
+}
+
+/* ------------------------------------------------------------------ */
+/*  Kernel 7: Bias gradient reduction (sum over batch)                */
+/* ------------------------------------------------------------------ */
+
+/**
+ * One thread per output dimension d.
+ *
+ * Computes grad_b[d] = sum_{n=0}^{N-1} dy[n, d]
+ *
+ * Used for both grad_b1 (d=CUR_HIDDEN, dy=d_hidden) and
+ * grad_b2 (d=CUR_OUTPUT, dy=d_pred).
+ *
+ * Launch: grid=(ceil(out_dim/256)), block=(256)
+ */
+extern "C" __global__ void curiosity_bias_grad_reduce(
+    const float* __restrict__ dy,      /* [N, out_dim] */
+    float*       __restrict__ grad_b,  /* [out_dim]    */
+    int N,
+    int out_dim
+) {
+    int d = blockIdx.x * blockDim.x + threadIdx.x;
+    if (d >= out_dim) return;
+
+    float sum = 0.0f;
+    for (int n = 0; n < N; n++) {
+        sum += dy[n * out_dim + d];
+    }
+    grad_b[d] = sum;
+}