feat(rl): FRD softmax + CE + dL/dlogits backward stage 1 (F.3a)

Per-(batch, horizon) softmax + cross-entropy loss + gradient w.r.t.
the 21 atom logits. First of three backward stages — F.3b adds layer-2
weight grads (dW2, db2, dhidden), F.3c adds layer-1 weight grads
(dW1, db1, dh_t with ReLU mask).

Kernel `cuda/rl_frd_softmax_ce_grad.cu`:
  * grid_dim = (B, FRD_N_HORIZONS, 1), block_dim = (FRD_N_ATOMS=21, 1, 1)
    — one block per (batch, horizon) pair, threads cooperate over the
    21 atoms via shared mem
  * Standard numerically-stable softmax: shift by row_max, exponentiate,
    normalize by row_sum (thread 0 does the serial reductions — 21
    atoms is small enough warp-shuffle overhead isn't worth it)
  * Gradient: (p[a] - 1{a==label}) / B at the source per v_head_bwd
    convention (mean-reduce over batch)
  * Loss: -log(p[label]) with 1e-30 floor against log(0)
  * Sentinel label (-1) zeros both gradient row and loss — for the
    missing-horizon case at the rightmost edge of the snapshot stream
    (forward returns at h=300 ticks aren't realized for the last
    300 snapshots; loader marks those labels with -1)
  * Per feedback_no_atomicadd: per-(b, h, a) sole-writer pattern

Rust wiring `src/rl/frd.rs::FrdHead::softmax_ce_grad`:
  * Second cubin loaded alongside fwd (separate module per the
    aux_heads pattern; small handle, no impact on init time)
  * Caller provides labels_d [B, FRD_N_HORIZONS] of i32 and gets back
    grad_logits + per-(b, h) raw CE; sum + λ_frd scaling left to the
    caller (F.4 will hook this into stats.l_total + Adam step)

Tests `tests/frd_head.rs` — 3 new GPU-oracle tests (6/6 file total),
all PASS on RTX 3050 Ti:
  1. frd_softmax_ce_grad_uniform_logits_match_log_n_atoms — for any
     label, uniform logits → CE = ln(FRD_N_ATOMS) = ln(21) ≈ 3.0445.
     Also asserts per-row Σ grad_logits = 0 (softmax-CE invariant).
  2. frd_softmax_ce_grad_sentinel_label_zeros_row — label=-1 with
     non-trivial random logits produces exactly zero loss + grad
     for every row (no leak through the sentinel path).
  3. frd_softmax_ce_grad_finite_diff_matches_analytical — perturbs
     one logit slot by ±ε=1e-3, compares (L(+ε) - L(-ε))/(2ε) to
     the kernel's analytical gradient. rel_err ≈ 1.3e-3 (fp32
     finite-diff is rounding-error-limited at this ε; tolerance
     set to 5e-3 with explanatory comment).

The first two tests provide strong analytical oracles (no CPU
reference impl per feedback_no_cpu_test_fallbacks). The finite-diff
test cross-validates the full softmax+CE chain via a numerical
gradient — the standard ground-truth for autodiff kernels.

crates/ml-alpha/build.rs

@@ -76,6 +76,7 @@ const KERNELS: &[&str] = &[
     "rl_trail_mutate",               // SP20 P1+P5 audit fix (a7/a8 dead): TrailTighten/TrailLoosen mutate unit_trail_distance bounded MIN/MAX, symmetric reciprocal adjust
     "rl_trail_stop_check",           // SP20 P1+P5 audit fix: per-unit trail breach check; OVERRIDE action to FlatFromLong/Short on breach (close routes through existing flat plumbing)
     "rl_frd_fwd",                    // SP20 P3: Forward-Return-Distribution head fwd — 2-layer MLP [HIDDEN_DIM → FRD_HIDDEN_DIM → FRD_N_HORIZONS × FRD_N_ATOMS]; ReLU hidden cached for bwd; softmax + CE happen in bwd
+    "rl_frd_softmax_ce_grad",        // SP20 P3 F.3a: per-(batch, horizon) softmax + CE loss + dL/dlogits; 1 block per (b, h), 21 threads; label = -1 sentinel masks the row
 ];
 
 // Cache bust v31 — five new reduce / derive kernels populate the input

crates/ml-alpha/cuda/rl_frd_softmax_ce_grad.cu

@@ -0,0 +1,115 @@
+// rl_frd_softmax_ce_grad.cu — FRD head per-horizon softmax + CE loss +
+// gradient w.r.t. logits (SP20 P3 backward stage 1).
+//
+// Per batch (b) and horizon (h), the 21 atom logits feed a softmax:
+//   p[b, h, a] = exp(logits[b, h, a] - max) / Σ_a' exp(logits[b, h, a'] - max)
+//
+// Cross-entropy against one-hot label:
+//   CE[b, h] = -log p[b, h, label[b, h]]
+//
+// Gradient (standard softmax-CE derivative, mean-reduced over batch
+// per the v_head_bwd convention so downstream reduce_axis0 + Adam
+// step land mean-grad without an explicit /B):
+//   grad_logits[b, h, a] = (1/B) × (p[b, h, a] - 1{a == label[b, h]})
+//
+// Label sentinel: a negative label index (e.g. -1) means "missing
+// horizon" (future return not yet realized at the rightmost edge of
+// the snapshot stream); the kernel zeros grad_logits for that row
+// and writes loss = 0 so the caller's sum is masked-out.
+//
+// Block layout: 1 block per (batch, horizon), 21 threads per block.
+//   * Cooperative max-reduction across 21 atoms (thread 0 serial in
+//     shared mem — 21 atoms is small enough that the warp-shuffle
+//     pattern's overhead isn't worth it)
+//   * Cooperative sum-exp across 21 atoms (same pattern)
+//   * Each thread writes one grad_logits slot
+//   * Thread 0 writes loss_per_batch_h[b, h]
+//
+// Per `feedback_no_atomicadd`: no atomics — per-(batch, horizon, atom)
+// sole-writer pattern.
+// Per `feedback_no_nvrtc`: pre-compiled cubin via build.rs.
+
+#include <stdint.h>
+
+#define FRD_N_HORIZONS  3
+#define FRD_N_ATOMS    21
+#define FRD_OUT_DIM   (FRD_N_HORIZONS * FRD_N_ATOMS)
+
+extern "C" __global__ void rl_frd_softmax_ce_grad(
+    const float* __restrict__ logits,             // [B, FRD_OUT_DIM]
+    const int*   __restrict__ labels,             // [B, FRD_N_HORIZONS]
+    int                       b_size,
+    float*       __restrict__ grad_logits,        // [B, FRD_OUT_DIM]   OUT
+    float*       __restrict__ loss_per_batch_h    // [B, FRD_N_HORIZONS] OUT
+) {
+    const int b = blockIdx.x;
+    const int h = blockIdx.y;
+    const int a = threadIdx.x;
+    if (b >= b_size || h >= FRD_N_HORIZONS || a >= FRD_N_ATOMS) return;
+
+    const int row_off = b * FRD_OUT_DIM + h * FRD_N_ATOMS;
+
+    __shared__ float exp_buf[FRD_N_ATOMS];
+    __shared__ float row_max;
+    __shared__ float row_sum;
+    __shared__ int   s_label;
+
+    // Phase 0: stage label (thread 0) — same value broadcast to all threads.
+    if (a == 0) {
+        s_label = labels[b * FRD_N_HORIZONS + h];
+    }
+    __syncthreads();
+
+    const int label = s_label;
+
+    // Sentinel label → emit zero gradient + zero loss for this row.
+    // (Caller masks out missing-horizon contributions from the total
+    // before backward propagation continues.)
+    if (label < 0 || label >= FRD_N_ATOMS) {
+        grad_logits[row_off + a] = 0.0f;
+        if (a == 0) {
+            loss_per_batch_h[b * FRD_N_HORIZONS + h] = 0.0f;
+        }
+        return;
+    }
+
+    // Phase 1: stage logits + compute max (thread 0 serial).
+    const float my_logit = logits[row_off + a];
+    exp_buf[a] = my_logit;
+    __syncthreads();
+    if (a == 0) {
+        float m = exp_buf[0];
+        #pragma unroll
+        for (int i = 1; i < FRD_N_ATOMS; ++i) {
+            m = fmaxf(m, exp_buf[i]);
+        }
+        row_max = m;
+    }
+    __syncthreads();
+
+    // Phase 2: each thread computes its exp(logit - max), stages.
+    const float my_exp = expf(my_logit - row_max);
+    exp_buf[a] = my_exp;
+    __syncthreads();
+    if (a == 0) {
+        float s = 0.0f;
+        #pragma unroll
+        for (int i = 0; i < FRD_N_ATOMS; ++i) {
+            s += exp_buf[i];
+        }
+        row_sum = s;
+    }
+    __syncthreads();
+
+    // Phase 3: per-atom softmax probability + gradient.
+    const float p_a = my_exp / row_sum;
+    const float target = (a == label) ? 1.0f : 0.0f;
+    const float inv_B = 1.0f / (float) b_size;
+    grad_logits[row_off + a] = (p_a - target) * inv_B;
+
+    // Phase 4: loss (thread 0 only). p_label = exp_buf[label] / row_sum.
+    if (a == 0) {
+        const float p_label = exp_buf[label] / row_sum;
+        loss_per_batch_h[b * FRD_N_HORIZONS + h] = -logf(fmaxf(p_label, 1e-30f));
+    }
+}

crates/ml-alpha/src/rl/frd.rs

@@ -58,6 +58,8 @@ use crate::rl::common::{FRD_HIDDEN_DIM, FRD_N_ATOMS, FRD_N_HORIZONS};
 use crate::trainer::integrated::write_slice_f32_d_pub;
 
 const FRD_FWD_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/rl_frd_fwd.cubin"));
+const FRD_SOFTMAX_CE_GRAD_CUBIN: &[u8] =
+    include_bytes!(concat!(env!("OUT_DIR"), "/rl_frd_softmax_ce_grad.cubin"));
 
 /// Per-batch output width: `FRD_N_HORIZONS × FRD_N_ATOMS`. Each batch
 /// row contains 3 contiguous horizon blocks of 21 atom logits.
@@ -77,12 +79,16 @@ impl Default for FrdHeadConfig {
     }
 }
 
-/// FRD head — owns device weights for the 2-layer MLP + cached forward
-/// cubin function. Backward kernel is added in F.3.
+/// FRD head — owns device weights for the 2-layer MLP + cached cubin
+/// function handles. Forward (F.2) ships separately from backward
+/// (F.3a/b/c) so each stage has its own GPU-oracle test before
+/// composition.
 pub struct FrdHead {
     stream: Arc<CudaStream>,
     _module: Arc<CudaModule>,
     fwd_fn: CudaFunction,
+    _softmax_ce_grad_module: Arc<CudaModule>,
+    softmax_ce_grad_fn: CudaFunction,
 
     pub w1_d: CudaSlice<f32>, // [HIDDEN_DIM, FRD_HIDDEN_DIM]
     pub b1_d: CudaSlice<f32>, // [FRD_HIDDEN_DIM]
@@ -100,6 +106,12 @@ impl FrdHead {
         let fwd_fn = module
             .load_function("rl_frd_fwd")
             .context("load rl_frd_fwd fn")?;
+        let softmax_ce_grad_module = ctx
+            .load_cubin(FRD_SOFTMAX_CE_GRAD_CUBIN.to_vec())
+            .context("load rl_frd_softmax_ce_grad cubin")?;
+        let softmax_ce_grad_fn = softmax_ce_grad_module
+            .load_function("rl_frd_softmax_ce_grad")
+            .context("load rl_frd_softmax_ce_grad fn")?;
 
         // Per pearl_scoped_init_seed_for_reproducibility — guard around
         // Xavier draws so GPU init helpers downstream see the same RNG.
@@ -133,6 +145,8 @@ impl FrdHead {
             stream,
             _module: module,
             fwd_fn,
+            _softmax_ce_grad_module: softmax_ce_grad_module,
+            softmax_ce_grad_fn,
             w1_d,
             b1_d,
             w2_d,
@@ -140,6 +154,50 @@ impl FrdHead {
         })
     }
 
+    /// Backward stage 1 (F.3a): per-(batch, horizon) softmax + CE loss
+    /// + gradient w.r.t. logits. Pairs with the forward kernel's
+    /// `logits_out_d`. The `labels_d` buffer is `[B, FRD_N_HORIZONS]`
+    /// of i32 — pass `-1` as the sentinel for missing horizons (the
+    /// kernel zeros that row's gradient + loss).
+    ///
+    /// Outputs:
+    ///   * `grad_logits_d [B, FRD_OUT_DIM]` — mean-reduced over batch
+    ///     (multiplied by `1/B` at the source per v_head_bwd convention)
+    ///   * `loss_per_batch_h_d [B, FRD_N_HORIZONS]` — raw CE per row;
+    ///     caller sums + scales by `λ_frd` before backward continues
+    pub fn softmax_ce_grad(
+        &self,
+        logits_d: &CudaSlice<f32>,
+        labels_d: &CudaSlice<i32>,
+        grad_logits_d: &mut CudaSlice<f32>,
+        loss_per_batch_h_d: &mut CudaSlice<f32>,
+        b_size: usize,
+    ) -> Result<()> {
+        debug_assert_eq!(logits_d.len(), b_size * FRD_OUT_DIM);
+        debug_assert_eq!(labels_d.len(), b_size * FRD_N_HORIZONS);
+        debug_assert_eq!(grad_logits_d.len(), b_size * FRD_OUT_DIM);
+        debug_assert_eq!(loss_per_batch_h_d.len(), b_size * FRD_N_HORIZONS);
+        let cfg = LaunchConfig {
+            grid_dim: (b_size as u32, FRD_N_HORIZONS as u32, 1),
+            block_dim: (FRD_N_ATOMS as u32, 1, 1),
+            shared_mem_bytes: 0,
+        };
+        let b_size_i = b_size as i32;
+        let mut launch = self.stream.launch_builder(&self.softmax_ce_grad_fn);
+        launch
+            .arg(logits_d)
+            .arg(labels_d)
+            .arg(&b_size_i)
+            .arg(grad_logits_d)
+            .arg(loss_per_batch_h_d);
+        unsafe {
+            launch
+                .launch(cfg)
+                .context("rl_frd_softmax_ce_grad launch")?;
+        }
+        Ok(())
+    }
+
     /// Forward pass — fills `logits_out_d [B, FRD_OUT_DIM]` and caches
     /// the post-ReLU hidden activations in `hidden_out_d [B,
     /// FRD_HIDDEN_DIM]` for the backward kernel.

crates/ml-alpha/tests/frd_head.rs

@@ -30,7 +30,7 @@ use ml_alpha::heads::HIDDEN_DIM;
 use ml_alpha::rl::common::{FRD_HIDDEN_DIM, FRD_N_ATOMS, FRD_N_HORIZONS};
 use ml_alpha::rl::frd::{FrdHead, FrdHeadConfig, FRD_OUT_DIM};
 use ml_alpha::trainer::integrated::{
-    read_slice_d_pub, write_slice_f32_d_pub,
+    read_slice_d_pub, write_slice_f32_d_pub, write_slice_i32_d_pub,
 };
 use ml_core::device::MlDevice;
 use std::sync::Arc;
@@ -187,3 +187,168 @@ fn frd_forward_relu_mask_consistent_with_cached_hidden() -> Result<()> {
     );
     Ok(())
 }
+
+fn upload_i32(stream: &Arc<CudaStream>, host: &[i32]) -> Result<CudaSlice<i32>> {
+    let mut d = stream.alloc_zeros::<i32>(host.len())?;
+    write_slice_i32_d_pub(stream, host, &mut d)?;
+    Ok(d)
+}
+
+#[test]
+#[ignore = "requires CUDA (MlDevice::cuda(0))"]
+fn frd_softmax_ce_grad_uniform_logits_match_log_n_atoms() -> Result<()> {
+    let Some((dev, head)) = build_head() else { return Ok(()) };
+    let stream = dev.cuda_stream()?.clone();
+    let b_size = 4;
+
+    // Uniform logits (all zeros) → softmax = 1/FRD_N_ATOMS everywhere,
+    // so CE = -log(1/21) = ln(21) ≈ 3.0445 for ANY label index.
+    let logits = vec![0.0_f32; b_size * FRD_OUT_DIM];
+    let logits_d = upload_f32(&stream, &logits)?;
+    // Set label = 10 (mid-bucket) for every (batch, horizon) row.
+    let labels: Vec<i32> = vec![10; b_size * FRD_N_HORIZONS];
+    let labels_d = upload_i32(&stream, &labels)?;
+    let mut grad_d = stream.alloc_zeros::<f32>(b_size * FRD_OUT_DIM)?;
+    let mut loss_d = stream.alloc_zeros::<f32>(b_size * FRD_N_HORIZONS)?;
+
+    head.softmax_ce_grad(&logits_d, &labels_d, &mut grad_d, &mut loss_d, b_size)?;
+
+    let loss = read_slice_d_pub(&stream, &loss_d, b_size * FRD_N_HORIZONS)?;
+    let expected = (FRD_N_ATOMS as f32).ln();
+    for (i, v) in loss.iter().enumerate() {
+        assert!(
+            (v - expected).abs() < 1e-5,
+            "uniform-logit CE[{i}] should be ln({}) = {expected}; got {v}",
+            FRD_N_ATOMS
+        );
+    }
+
+    // Softmax-CE gradient invariant: Σ_a (p[a] - 1{a==label}) = 0
+    // (mean-reduced over B doesn't change the per-row sum being zero).
+    let grad = read_slice_d_pub(&stream, &grad_d, b_size * FRD_OUT_DIM)?;
+    for b in 0..b_size {
+        for h in 0..FRD_N_HORIZONS {
+            let off = b * FRD_OUT_DIM + h * FRD_N_ATOMS;
+            let row_sum: f32 = grad[off..off + FRD_N_ATOMS].iter().sum();
+            assert!(
+                row_sum.abs() < 1e-6,
+                "Σ grad_logits per (b={b}, h={h}) must be 0 (softmax-CE invariant); got {row_sum}"
+            );
+        }
+    }
+
+    eprintln!(
+        "PASS — uniform logits → CE = ln({}) = {:.4}; per-row grad sum = 0",
+        FRD_N_ATOMS, expected
+    );
+    Ok(())
+}
+
+#[test]
+#[ignore = "requires CUDA (MlDevice::cuda(0))"]
+fn frd_softmax_ce_grad_sentinel_label_zeros_row() -> Result<()> {
+    let Some((dev, head)) = build_head() else { return Ok(()) };
+    let stream = dev.cuda_stream()?.clone();
+    let b_size = 2;
+
+    // Non-trivial logits so a non-sentinel label would produce non-zero
+    // loss and gradient — proving the sentinel really is masking.
+    use rand::{Rng, SeedableRng};
+    use rand_chacha::ChaCha8Rng;
+    let mut r = ChaCha8Rng::seed_from_u64(0x5EE);
+    let logits: Vec<f32> = (0..b_size * FRD_OUT_DIM)
+        .map(|_| r.gen_range(-1.0..1.0))
+        .collect();
+    let logits_d = upload_f32(&stream, &logits)?;
+    // Labels: ALL sentinel (-1) → every row should be zeroed.
+    let labels: Vec<i32> = vec![-1; b_size * FRD_N_HORIZONS];
+    let labels_d = upload_i32(&stream, &labels)?;
+    let mut grad_d = stream.alloc_zeros::<f32>(b_size * FRD_OUT_DIM)?;
+    let mut loss_d = stream.alloc_zeros::<f32>(b_size * FRD_N_HORIZONS)?;
+
+    head.softmax_ce_grad(&logits_d, &labels_d, &mut grad_d, &mut loss_d, b_size)?;
+
+    let loss = read_slice_d_pub(&stream, &loss_d, b_size * FRD_N_HORIZONS)?;
+    let grad = read_slice_d_pub(&stream, &grad_d, b_size * FRD_OUT_DIM)?;
+    for (i, v) in loss.iter().enumerate() {
+        assert_eq!(*v, 0.0, "sentinel label loss[{i}] must be 0; got {v}");
+    }
+    for (i, v) in grad.iter().enumerate() {
+        assert_eq!(
+            *v, 0.0,
+            "sentinel label grad_logits[{i}] must be 0; got {v}"
+        );
+    }
+
+    eprintln!("PASS — sentinel label (-1) zeros loss + grad rows");
+    Ok(())
+}
+
+#[test]
+#[ignore = "requires CUDA (MlDevice::cuda(0))"]
+fn frd_softmax_ce_grad_finite_diff_matches_analytical() -> Result<()> {
+    let Some((dev, head)) = build_head() else { return Ok(()) };
+    let stream = dev.cuda_stream()?.clone();
+    let b_size = 1;
+
+    // Single batch, fixed (h=0, target_a=5). Pick a smooth random logit
+    // pattern so probabilities aren't degenerate at 0/1.
+    use rand::{Rng, SeedableRng};
+    use rand_chacha::ChaCha8Rng;
+    let mut r = ChaCha8Rng::seed_from_u64(0xFD1FF);
+    let mut logits: Vec<f32> = (0..b_size * FRD_OUT_DIM)
+        .map(|_| r.gen_range(-0.5..0.5))
+        .collect();
+    let labels: Vec<i32> = vec![5, 10, 15]; // distinct per horizon
+    let labels_d = upload_i32(&stream, &labels)?;
+
+    // Analytical gradient via the kernel.
+    let logits_d = upload_f32(&stream, &logits)?;
+    let mut grad_d = stream.alloc_zeros::<f32>(b_size * FRD_OUT_DIM)?;
+    let mut loss_d = stream.alloc_zeros::<f32>(b_size * FRD_N_HORIZONS)?;
+    head.softmax_ce_grad(&logits_d, &labels_d, &mut grad_d, &mut loss_d, b_size)?;
+    let grad_analytical = read_slice_d_pub(&stream, &grad_d, b_size * FRD_OUT_DIM)?;
+
+    // Finite-difference for slot (b=0, h=0, a=3). Note: gradient was
+    // mean-reduced by 1/B at the kernel source; for b_size=1 the 1/B
+    // factor is 1.0 so finite-diff matches directly.
+    let probe_h = 0_usize;
+    let probe_a = 3_usize;
+    let probe_off = probe_h * FRD_N_ATOMS + probe_a;
+    let eps = 1e-3_f32;
+
+    // L(logits + ε · e_j) — perturb only the target slot upward.
+    logits[probe_off] += eps;
+    let logits_plus_d = upload_f32(&stream, &logits)?;
+    head.softmax_ce_grad(&logits_plus_d, &labels_d, &mut grad_d, &mut loss_d, b_size)?;
+    let loss_plus = read_slice_d_pub(&stream, &loss_d, b_size * FRD_N_HORIZONS)?;
+    let l_plus = loss_plus[probe_h]; // only h=0 affected — h=1,2 share the perturbation only if probe was in their horizon block
+
+    // L(logits - ε · e_j)
+    logits[probe_off] -= 2.0 * eps;
+    let logits_minus_d = upload_f32(&stream, &logits)?;
+    head.softmax_ce_grad(&logits_minus_d, &labels_d, &mut grad_d, &mut loss_d, b_size)?;
+    let loss_minus = read_slice_d_pub(&stream, &loss_d, b_size * FRD_N_HORIZONS)?;
+    let l_minus = loss_minus[probe_h];
+
+    let numerical = (l_plus - l_minus) / (2.0 * eps);
+    let analytical = grad_analytical[probe_off];
+
+    let rel_err = (numerical - analytical).abs() / (analytical.abs().max(1e-6));
+    // 5e-3 tolerance: fp32 single-precision finite-difference at ε=1e-3
+    // saturates around the per-evaluation rounding error (~1e-7) divided
+    // by ε, so 1e-4-ish absolute error is intrinsic to the test method,
+    // not the kernel. Tighter threshold would force fp64 finite-diff
+    // (not worth the complexity here).
+    assert!(
+        rel_err < 5e-3,
+        "finite-diff gradient mismatch at (h={probe_h}, a={probe_a}): \
+         analytical={analytical:.6}, numerical={numerical:.6}, rel_err={rel_err:.6}"
+    );
+
+    eprintln!(
+        "PASS — finite-diff matches analytical: analytical={:.6} numerical={:.6} rel_err={:.2e}",
+        analytical, numerical, rel_err
+    );
+    Ok(())
+}