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

// rl_iqn_forward.cu — IQN forward pass: split pipeline for cuBLAS SGEMM.
//
// The monolithic kernel is replaced by three custom kernels interleaved
// with two cuBLAS SGEMM calls driven from Rust. The pipeline:
//
//   1. rl_iqn_tau_cos_features (custom):
//      - Thread 0 of each (batch, tau) block samples tau ~ U(0,1) via
//        inline xorshift32 and writes to tau[batch, tau_idx].
//      - All threads compute cos_features[j] = cos((j+1) * pi * tau)
//        for j = 0..EMBED_DIM-1.
//      Output: tau[B, N_TAU], cos_features[B*N_TAU, EMBED_DIM].
//
//   2. cuBLAS SGEMM (from Rust):
//      embed_out[B*N_TAU, HIDDEN_DIM] = cos_features[B*N_TAU, EMBED_DIM]
//                                       @ W_embed[EMBED_DIM, HIDDEN_DIM]
//
//   3. rl_iqn_relu_hadamard (custom):
//      embed_out += b_embed (bias add)
//      phi = ReLU(embed_out)
//      combined = h_t ⊙ phi (hadamard product with broadcast over tau)
//      Output: combined[B*N_TAU, HIDDEN_DIM].
//
//   4. cuBLAS SGEMM (from Rust):
//      q_raw[B*N_TAU, N_ACTIONS] = combined[B*N_TAU, HIDDEN_DIM]
//                                  @ W_out[HIDDEN_DIM, N_ACTIONS]
//
//   5. rl_iqn_bias_add_q (custom):
//      q_values = q_raw + b_out (broadcast bias)
//      Output: q_values[B, N_TAU, N_ACTIONS].
//
//   6. rl_iqn_expected_q (unchanged):
//      E[Q(s,a)] = mean_tau Q(s, tau, a).
//
// Per `feedback_no_atomicadd.md`: no atomicAdd.
// Per `feedback_cpu_is_read_only.md`: all compute on GPU.
// Per `feedback_no_nvrtc.md`: pre-compiled cubin via build.rs.

#include <stdint.h>

#define HIDDEN_DIM   128
#define N_ACTIONS    11
#define N_TAU_MAX    64
#define EMBED_DIM    64
#define PI_F         3.14159265f

// Inline xorshift32 PRNG — identical to rl_sample_tau.cu.
__device__ static uint32_t xorshift32_iqn(uint32_t* state) {
    uint32_t x = *state;
    x ^= x << 13;
    x ^= x >> 17;
    x ^= x << 5;
    *state = x;
    return x;
}

// ─────────────────────────────────────────────────────────────────────
// rl_iqn_tau_cos_features:
//   Stage 1: inline tau sampling + cosine basis computation.
//
//   Grid  = (B, N_TAU, 1)
//   Block = (EMBED_DIM, 1, 1)  — one thread per embedding dimension.
//
// DETERMINISM (Phase 2.6, 2026-06-02): `prng_state` is now READ-ONLY
// in this kernel. The previous version had a read-write race: all
// N_TAU blocks per batch read `prng_state[batch]` at the head of the
// kernel, and the (tau_idx == 0) block wrote back the advanced state
// at the end — with NO inter-block barrier. Blocks with tau_idx > 0
// running on a different SM could read the post-write value (if the
// tau_idx == 0 block finished first) OR the pre-write value
// (otherwise), depending on SM scheduling. Different runs picked
// different orderings, producing run-dependent τ values for
// tau_idx > 0 and hence run-dependent `iqn_q_values` even at step 0
// (no upstream RL-loop divergence required). Diagnosed via Phase 2.6
// `dump_backward_state_for_debug` Group H verdict (iqn_q_values
// DIVERGE at step 0 idx 55, max|Δ|=4.5e-5) while Groups G/I/J stayed
// EQUAL through steps 0/1 — see
// `docs/superpowers/notes/2026-06-02-determinism-phase2.6-backward
// -kernel-fix.md`.
//
// Fix: kernel now READS prng_state[batch] only (no writes); the
// state advancement is performed by `rl_iqn_advance_prng_state`
// in a separate single-thread-per-batch launch that runs AFTER
// `rl_iqn_tau_cos_features` finishes via stream ordering — no race
// possible. Per `feedback_no_atomicadd.md` + canonical Phase-2
// PER-rebuild rule: "fixed accumulation order requires a real
// barrier between stages, and kernel-launch ordering on a single
// stream is the cleanest grid-wide barrier we have".
//
// Inputs:
//   prng_state    [B]                       — per-batch xorshift32 seed (READ-ONLY)
//   B                                       — batch size
//   N_TAU                                   — number of quantile samples
// Outputs:
//   tau           [B, N_TAU]                — sampled U(0,1) quantile fractions
//   cos_features  [B*N_TAU, EMBED_DIM]     — cos((j+1) * pi * tau)
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void rl_iqn_tau_cos_features(
    const uint32_t* __restrict__ prng_state,  // [B] READ-ONLY (advanced by sibling kernel)
    int             B,
    int             N_TAU,
    float*       __restrict__ tau,            // [B, N_TAU]
    float*       __restrict__ cos_features    // [B*N_TAU, EMBED_DIM]
) {
    const int batch = blockIdx.x;
    const int tau_idx = blockIdx.y;
    const int j = threadIdx.x;  // embedding dim index
    if (batch >= B) return;
    if (tau_idx >= N_TAU) return;
    if (j >= EMBED_DIM) return;

    // Thread 0 samples tau for this (batch, tau_idx) pair.
    __shared__ float s_tau_val;
    if (j == 0) {
        uint32_t seed = prng_state[batch];
        if (seed == 0u) seed = (uint32_t)(batch + 1) * 2654435761u + 0xBEEFu;

        uint32_t local_state = seed ^ (uint32_t)(tau_idx * 2654435761u);

        #pragma unroll
        for (int w = 0; w < 4; ++w) xorshift32_iqn(&local_state);

        const uint32_t r = xorshift32_iqn(&local_state);
        const float u = (float)(r >> 8) * (1.0f / 16777216.0f);

        tau[batch * N_TAU + tau_idx] = u;
        s_tau_val = u;
        // prng_state advancement moved to `rl_iqn_advance_prng_state`
        // (see kernel below); no writes to prng_state in this kernel.
    }
    __syncthreads();

    const float tau_val = s_tau_val;

    // cos_features[row, j] = cos((j+1) * pi * tau_val)
    const int row = batch * N_TAU + tau_idx;
    cos_features[row * EMBED_DIM + j] = cosf((float)(j + 1) * PI_F * tau_val);
}

// ─────────────────────────────────────────────────────────────────────
// rl_iqn_advance_prng_state:
//   Companion to `rl_iqn_tau_cos_features` — advances the per-batch
//   xorshift32 PRNG state by exactly 8 steps. Launched AFTER
//   `rl_iqn_tau_cos_features` so the read-only sampling sees a stable
//   value for prng_state[batch] across all (tau_idx) blocks, then
//   this kernel applies a single deterministic update per batch.
//
// Determinism rationale: by separating "read state for tau sampling"
// from "advance state for next call", we eliminate the read-write
// race the previous monolithic kernel had. Both kernels run on the
// same stream, so the launch ordering is a grid-wide barrier
// (no __threadfence required).
//
// Grid  = (ceil(B/256), 1, 1)
// Block = (256, 1, 1)        — one thread per batch element.
//
// Inputs:
//   B                       — batch size
// Outputs:
//   prng_state [B]          — advanced 8 xorshift32 steps in place.
//                              `0` seed bootstrap matches the
//                              `rl_iqn_tau_cos_features` convention
//                              `(batch + 1) * 2654435761u + 0xBEEFu`
//                              so the first advance from the
//                              cold-start sentinel is identical to
//                              what the prior monolithic kernel
//                              produced for tau_idx==0.
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void rl_iqn_advance_prng_state(
    uint32_t*    __restrict__ prng_state,    // [B]
    int          B
) {
    const int batch = blockIdx.x * blockDim.x + threadIdx.x;
    if (batch >= B) return;

    uint32_t seed = prng_state[batch];
    if (seed == 0u) seed = (uint32_t)(batch + 1) * 2654435761u + 0xBEEFu;

    uint32_t adv = seed;
    #pragma unroll
    for (int w = 0; w < 8; ++w) xorshift32_iqn(&adv);
    prng_state[batch] = adv;
}

// ─────────────────────────────────────────────────────────────────────
// rl_iqn_relu_hadamard:
//   Stage 3: bias-add → ReLU → element-wise product with h_t.
//
//   Grid  = (B*N_TAU, ceil(HIDDEN_DIM / 256), 1)
//   Block = (min(HIDDEN_DIM, 256), 1, 1)
//
// Inputs:
//   embed_out     [B*N_TAU, HIDDEN_DIM]  — output of cuBLAS SGEMM (no bias)
//   b_embed       [HIDDEN_DIM]           — embedding bias
//   h_t           [B, HIDDEN_DIM]        — encoder hidden state
//   M                                    — total rows = B * N_TAU
//   B                                    — batch size (for h_t indexing)
//   N_TAU                                — quantile count
// Outputs:
//   combined      [B*N_TAU, HIDDEN_DIM]  — h_t ⊙ ReLU(embed_out + b_embed)
//   embed_pre_relu [B*N_TAU, HIDDEN_DIM] — embed_out + b_embed (before ReLU,
//                                          saved for backward)
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void rl_iqn_relu_hadamard(
    const float* __restrict__ embed_out,       // [M, HIDDEN_DIM]
    const float* __restrict__ b_embed,         // [HIDDEN_DIM]
    const float* __restrict__ h_t,             // [B, HIDDEN_DIM]
    int          M,                            // B * N_TAU
    int          B,
    int          N_TAU,
    float*       __restrict__ combined,        // [M, HIDDEN_DIM]
    float*       __restrict__ embed_pre_relu   // [M, HIDDEN_DIM] (saved for bwd)
) {
    const int row = blockIdx.x;
    const int c = blockIdx.y * blockDim.x + threadIdx.x;
    if (row >= M) return;
    if (c >= HIDDEN_DIM) return;

    const int batch = row / N_TAU;

    float val = embed_out[row * HIDDEN_DIM + c] + b_embed[c];
    embed_pre_relu[row * HIDDEN_DIM + c] = val;

    float phi_c = fmaxf(val, 0.0f);
    combined[row * HIDDEN_DIM + c] = h_t[batch * HIDDEN_DIM + c] * phi_c;
}

// ─────────────────────────────────────────────────────────────────────
// rl_iqn_bias_add_q:
//   Stage 5: add b_out bias to the cuBLAS SGEMM output.
//
//   Grid  = (B*N_TAU, 1, 1)
//   Block = (N_ACTIONS, 1, 1)
//
// In-place: q_values[row, a] += b_out[a]
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void rl_iqn_bias_add_q(
    float*       __restrict__ q_values,   // [M, N_ACTIONS] (mutated in-place)
    const float* __restrict__ b_out,      // [N_ACTIONS]
    int          M                        // B * N_TAU
) {
    const int row = blockIdx.x;
    const int a = threadIdx.x;
    if (row >= M) return;
    if (a >= N_ACTIONS) return;

    q_values[row * N_ACTIONS + a] += b_out[a];
}

// ─────────────────────────────────────────────────────────────────────
// rl_iqn_expected_q:
//   Compute E[Q(s, a)] = (1/N_TAU) × Σ_{tau} Q(s, tau, a)
//   for ensemble action selection.
//
//   Unchanged from the original monolithic kernel.
//
// Inputs:
//   q_values    [B, N_TAU, N_ACTIONS]     — full quantile Q tensor
//   B                                     — batch size
//   N_TAU                                 — number of quantile samples
// Outputs:
//   expected_q  [B, N_ACTIONS]            — mean Q per action
//
// Block layout: grid = (B, 1, 1); block = (N_ACTIONS, 1, 1).
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void rl_iqn_expected_q(
    const float* __restrict__ q_values,    // [B, N_TAU, N_ACTIONS]
    int          B,
    int          N_TAU,
    float*       __restrict__ expected_q    // [B, N_ACTIONS]
) {
    const int batch = blockIdx.x;
    const int act = threadIdx.x;
    if (batch >= B) return;
    if (act >= N_ACTIONS) return;

    float sum = 0.0f;
    const int base = batch * N_TAU * N_ACTIONS;
    for (int t = 0; t < N_TAU; ++t) {
        sum += q_values[base + t * N_ACTIONS + act];
    }
    expected_q[batch * N_ACTIONS + act] = sum / (float)N_TAU;
}