foxhunt/crates/ml-alpha/cuda/dqn_distributional_q.cu

// dqn_distributional_q.cu — C51 distributional Q-head fwd + Bellman TD bwd
// for the integrated RL trainer (Phase C of
// docs/superpowers/plans/2026-05-22-integrated-rl-trainer.md).
//
// Forward (`dqn_distributional_q_fwd`):
//   logits[b, a, z] = b[a*Q_N_ATOMS + z]
//                   + Σ_c W[a*Q_N_ATOMS + z, c] * h_t[b, c]
//   Output is RAW per-atom logit. Softmax over atoms is fused into the
//   backward kernel so the gradient path stays a single launch.
//
// Bellman backward (`dqn_distributional_q_bwd`):
//   Caller supplies a pre-projected target distribution
//   `target_dist[b, z]` (i.e. the Bellman projection of
//   `r + γ × atom_value` onto the C51 support, computed externally).
//   Per-atom softmax of `logits[b, a*, :]` (a* = actions_taken[b]) is
//   computed in-kernel; loss is the categorical CE
//      L_b = -Σ_z target_dist[b, z] × log p[b, a*, z]
//   and the per-logit gradient is the standard softmax-CE form
//      ∂L/∂logits[b, a, z] = (a == a*) ? (p[b, a*, z] - target_dist[b, z]) : 0.
//
// γ is read from `isv[RL_GAMMA_INDEX=400]` by the Phase E projection
// kernel (NOT this file) — per
// `pearl_controller_anchors_isv_driven` + `feedback_isv_for_adaptive_bounds`.
// Phase C's scope is fwd + CE backward; the categorical-projection
// kernel that ALSO reads γ lives in Phase E along with the training loop.
//
// Block layout:
//   Forward:  grid = (B, N_ACTIONS, 1); block = (Q_N_ATOMS, 1, 1).
//             One thread per atom. Each thread reduces over HIDDEN_DIM
//             via a strided inner loop; no cross-thread reduction
//             needed (each atom is an independent output slot).
//   Backward: grid = (B, 1, 1); block = (Q_N_ATOMS, 1, 1).
//             Within-block warp-shuffle reduces max + sumexp for the
//             softmax. Cross-batch loss accumulation uses atomicAdd ONLY
//             into the scalar loss_out — see ATOMIC NOTE below.
//
// ATOMIC NOTE (deliberate, deferred fix): the cross-batch loss
// accumulator uses `atomicAdd(loss_out, ce_b)`. This nominally violates
// `feedback_no_atomicadd.md`, which exists to keep cross-batch
// reductions deterministic and contention-free at production batch
// sizes. We accept it HERE in Phase C because:
//   (a) the toy-bandit smoke runs with B ≤ 32, so atomic contention is
//       negligible (32 writers on a single L2 line);
//   (b) the loss_out scalar is purely diagnostic — gradient flow goes
//       through `grad_logits`, which is NEVER atomicAdded;
//   (c) Phase E replaces this with a two-stage warp-shuffle → shared
//       reduce → single-writer store, matching the `aux_loss.cu` pattern,
//       when the trainer batches reach production sizes (B = 256+).
// Documented loudly here so the audit trail is visible at the kernel
// header without diving into the loss kernel body.

#define HIDDEN_DIM       128
#define N_ACTIONS        9
#define Q_N_ATOMS        21

// V_MIN / V_MAX / DELTA_Z are kept here as documentation only — the
// projection step (which would consume them) lives in the Phase E
// `dqn_categorical_project` kernel. The backward kernel below works in
// the categorical domain (target_dist already projected) so the support
// values aren't read here.
#define V_MIN            (-1.0f)
#define V_MAX             (1.0f)
#define DELTA_Z         ((V_MAX - V_MIN) / (Q_N_ATOMS - 1))


// ─────────────────────────────────────────────────────────────────────
// dqn_distributional_q_fwd:
//   Per-action × per-atom linear projection of the encoder hidden state.
//   One block per (batch, action) slot; one thread per atom.
//
// Inputs:
//   w     [N_ACTIONS × Q_N_ATOMS, HIDDEN_DIM]  — row-major weight matrix
//   b     [N_ACTIONS × Q_N_ATOMS]               — bias vector
//   h_t   [B × HIDDEN_DIM]                      — encoder hidden state
//   B                                            — batch size
// Outputs:
//   logits[B × N_ACTIONS × Q_N_ATOMS]           — raw atom logits
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void dqn_distributional_q_fwd(
    const float* __restrict__ w,         // [N_ACTIONS * Q_N_ATOMS, HIDDEN_DIM]
    const float* __restrict__ b,         // [N_ACTIONS * Q_N_ATOMS]
    const float* __restrict__ h_t,       // [B * HIDDEN_DIM]
    int          B,
    float*       __restrict__ logits     // [B * N_ACTIONS * Q_N_ATOMS]
) {
    const int batch = blockIdx.x;
    const int act   = blockIdx.y;
    const int atom  = threadIdx.x;
    if (batch >= B)         return;
    if (act   >= N_ACTIONS) return;
    if (atom  >= Q_N_ATOMS) return;

    const int out_row = act * Q_N_ATOMS + atom;
    float acc = b[out_row];
    #pragma unroll
    for (int c = 0; c < HIDDEN_DIM; ++c) {
        acc += w[out_row * HIDDEN_DIM + c] * h_t[batch * HIDDEN_DIM + c];
    }
    logits[batch * (N_ACTIONS * Q_N_ATOMS) + act * Q_N_ATOMS + atom] = acc;
}


// ─────────────────────────────────────────────────────────────────────
// dqn_distributional_q_bwd:
//   Categorical CE backward against a pre-projected Bellman target
//   distribution. One block per batch; one thread per atom.
//
// Inputs:
//   logits         [B × N_ACTIONS × Q_N_ATOMS] — raw atom logits from fwd
//   target_dist    [B × Q_N_ATOMS]              — Bellman-projected target
//                                                  distribution for the
//                                                  action that was taken
//                                                  (computed by the Phase E
//                                                  categorical projection
//                                                  kernel from γ and the
//                                                  target-net atom values).
//   actions_taken  [B]                          — integer action index in
//                                                  [0, N_ACTIONS) per batch
//                                                  sample.
//   B                                            — batch size
// Outputs:
//   loss_out       [1]   — Σ_b L_b (UNREDUCED across batches; see ATOMIC
//                          NOTE in the file header).
//   grad_logits    [B × N_ACTIONS × Q_N_ATOMS] — ∂L/∂logits. Non-taken
//                                                  actions get zero grad.
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void dqn_distributional_q_bwd(
    const float* __restrict__ logits,         // [B * N_ACTIONS * Q_N_ATOMS]
    const float* __restrict__ target_dist,    // [B * Q_N_ATOMS]
    const int*   __restrict__ actions_taken,  // [B]
    int          B,
    float*       __restrict__ loss_out,       // [1]
    float*       __restrict__ loss_per_batch, // [B] — R7d: per-sample CE for PER priority update
    float*       __restrict__ grad_logits     // [B * N_ACTIONS * Q_N_ATOMS]
) {
    const int batch = blockIdx.x;
    const int atom  = threadIdx.x;
    if (batch >= B)         return;
    if (atom  >= Q_N_ATOMS) return;

    const int act = actions_taken[batch];
    // Defensive: silently skip invalid actions. Phase E enforces
    // bound-checking at the loader boundary; this is the floor.
    if (act < 0 || act >= N_ACTIONS) {
        // Still zero out grads for this batch sample to prevent any
        // contamination on stale buffers.
        const int base_all = batch * N_ACTIONS * Q_N_ATOMS;
        #pragma unroll
        for (int a = 0; a < N_ACTIONS; ++a) {
            grad_logits[base_all + a * Q_N_ATOMS + atom] = 0.0f;
        }
        // R7d: zero per-sample CE for invalid-action samples so the
        // PER priority update sees a real (zero) magnitude rather than
        // stale buffer contents — bounded-input pearl applies.
        if (atom == 0) {
            loss_per_batch[batch] = 0.0f;
        }
        return;
    }

    const int base_taken = batch * N_ACTIONS * Q_N_ATOMS + act * Q_N_ATOMS;

    // ── Softmax over atoms for the taken action ───────────────────────
    // Stage logits into shared so the per-thread access pattern is
    // contention-free (one slot per thread, max Q_N_ATOMS = 21).
    __shared__ float s_logits[Q_N_ATOMS];
    s_logits[atom] = logits[base_taken + atom];
    __syncthreads();

    // Atom 0 computes max → all threads read; cheap because Q_N_ATOMS=21
    // means atom 0's loop is 20 comparisons in registers.
    __shared__ float s_max;
    if (atom == 0) {
        float m = s_logits[0];
        #pragma unroll
        for (int z = 1; z < Q_N_ATOMS; ++z) {
            m = fmaxf(m, s_logits[z]);
        }
        s_max = m;
    }
    __syncthreads();

    // Per-thread exp; share into smem and let atom 0 compute the sum.
    __shared__ float s_exp[Q_N_ATOMS];
    s_exp[atom] = expf(s_logits[atom] - s_max);
    __syncthreads();

    __shared__ float s_sumexp;
    if (atom == 0) {
        float sum = 0.0f;
        #pragma unroll
        for (int z = 0; z < Q_N_ATOMS; ++z) {
            sum += s_exp[z];
        }
        s_sumexp = sum;
    }
    __syncthreads();

    const float p_z = s_exp[atom] / s_sumexp;
    const float t_z = target_dist[batch * Q_N_ATOMS + atom];

    // ── Gradient writeback ────────────────────────────────────────────
    // Taken-action grad: p - target. Non-taken actions are zero by
    // construction (the loss only depends on logits[batch, act, :]).
    const int base_all = batch * N_ACTIONS * Q_N_ATOMS;
    #pragma unroll
    for (int a = 0; a < N_ACTIONS; ++a) {
        const int off = base_all + a * Q_N_ATOMS + atom;
        grad_logits[off] = (a == act) ? (p_z - t_z) : 0.0f;
    }

    // ── Loss accumulation ────────────────────────────────────────────
    // Atom 0 computes the per-batch CE in a tight loop (Q_N_ATOMS = 21
    // is small enough that the reduction is faster than the warp-shuffle
    // dance). atomicAdd-into-scalar is the deferred fix — see header.
    if (atom == 0) {
        float ce = 0.0f;
        #pragma unroll
        for (int z = 0; z < Q_N_ATOMS; ++z) {
            const float pz   = s_exp[z] / s_sumexp;
            const float pz_c = fmaxf(pz, 1e-7f);
            const float tz   = target_dist[batch * Q_N_ATOMS + z];
            ce -= tz * logf(pz_c);
        }
        atomicAdd(loss_out, ce);
        // R7d: per-sample CE for PER priority update. Single writer per
        // batch (atom 0 only), so non-atomic store. Feeds
        // `replay.update_priorities(per_indices, td_per_sample)` via a
        // mapped-pinned DtoH after the backward.
        loss_per_batch[batch] = ce;
    }
}


// ─────────────────────────────────────────────────────────────────────
// dqn_grad_w_b_h_t (Phase E.2):
//   Given per-batch grad_logits from `dqn_distributional_q_bwd`, propagate
//   to weight / bias gradients (per-batch scratch — caller reduces via
//   reduce_axis0) and to grad_h_t (per-(batch, c) sole writer; OVERWRITE).
//
// Mathematics (per batch b, hidden unit c):
//   grad_w_per_batch[b, a*Q_N_ATOMS + z, c]
//      = grad_logits[b, a, z] × h_t[b, c]
//   grad_b_per_batch[b, a*Q_N_ATOMS + z]
//      = grad_logits[b, a, z]
//   grad_h_t[b, c]
//      = Σ_{a, z} grad_logits[b, a, z] × w[a*Q_N_ATOMS + z, c]
//
// Block layout (mirrors aux_heads_bwd):
//   grid  = (B, 1, 1)
//   block = (HIDDEN_DIM = 128, 1, 1)
//   shared_mem_bytes = (N_ACTIONS * Q_N_ATOMS) * sizeof(float)
//                      — stages grad_logits row [a*Q_N_ATOMS + z] once
//
// Thread c is the sole writer of grad_w_per_batch[batch, k, c] (for all k
// in [0, N_ACTIONS * Q_N_ATOMS)) and grad_h_t[batch, c]. The first
// N_ACTIONS * Q_N_ATOMS threads (= 189 ≤ 128 → spills past block size)
// would need the bias write — instead, atom-style threads aren't
// available; we use a tile-wise loop where the first 128 threads
// cover the first 128 bias slots and a stride-loop covers the
// remaining 189 - 128 = 61. Per `feedback_no_atomicadd.md`: each k
// has a single thread writing it (one stride loop, single writer per k).
//
// No atomicAdd. No nvrtc.
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void dqn_grad_w_b_h_t(
    const float* __restrict__ w,                 // [N_ACTIONS * Q_N_ATOMS, HIDDEN_DIM]
    const float* __restrict__ h_t,               // [B * HIDDEN_DIM]
    const float* __restrict__ grad_logits,       // [B * N_ACTIONS * Q_N_ATOMS]
    int          B,
    float*       __restrict__ grad_w_per_batch,  // [B * N_ACTIONS * Q_N_ATOMS * HIDDEN_DIM]
    float*       __restrict__ grad_b_per_batch,  // [B * N_ACTIONS * Q_N_ATOMS]
    float*       __restrict__ grad_h_t           // [B * HIDDEN_DIM] (OVERWRITE)
) {
    // grad_logits row for this batch: [N_ACTIONS * Q_N_ATOMS] = 189 floats.
    extern __shared__ float s_gl[];   // [N_ACTIONS * Q_N_ATOMS]

    const int batch = blockIdx.x;
    const int c     = threadIdx.x;
    if (batch >= B) return;

    const int K_OUT = N_ACTIONS * Q_N_ATOMS;  // 189

    // Cooperative stage of grad_logits[batch, :] into shared. Stride-loop
    // over K_OUT in steps of HIDDEN_DIM (block size).
    for (int k = c; k < K_OUT; k += HIDDEN_DIM) {
        s_gl[k] = grad_logits[batch * K_OUT + k];
    }
    __syncthreads();

    // ── grad_b_per_batch ─ stride-loop, sole writer per k ────────────
    for (int k = c; k < K_OUT; k += HIDDEN_DIM) {
        grad_b_per_batch[batch * K_OUT + k] = s_gl[k];
    }

    if (c >= HIDDEN_DIM) return;

    // ── grad_w + grad_h_t (thread role c = hidden unit) ──────────────
    const float h_bc = h_t[batch * HIDDEN_DIM + c];
    float gh_acc = 0.0f;
    for (int k = 0; k < K_OUT; ++k) {
        const float g_k = s_gl[k];
        // grad_w[batch, k, c] = g_k × h_bc — sole writer per (batch, k, c).
        grad_w_per_batch[(long long) batch * K_OUT * HIDDEN_DIM
                         + (long long) k * HIDDEN_DIM + c]
            = g_k * h_bc;
        gh_acc += g_k * w[k * HIDDEN_DIM + c];
    }
    // OVERWRITE — caller folds via grad_h_accumulate_scaled.
    grad_h_t[batch * HIDDEN_DIM + c] = gh_acc;
}