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

// cfc_step_per_branch.cu — fused per-branch CfC step (forward + backward).
//
// Per spec §5.4 point 1 (docs/superpowers/specs/2026-05-21-per-horizon-cfc-inference-design.md):
//   Single fused kernel covering all N_HORIZONS branches × n_batch in one launch.
//   Grid: (B, N_HORIZONS, 1), block: (MAX_BUCKET_DIM=96, 1, 1) uniform predicate.
//
// Uniform predicate (`threadIdx.x >= bucket_dim_k[branch]` early-return) handles
// uneven bucket sizes [43, 43, 42] without warp divergence — the
// comparison is against a per-block constant, so all threads in a warp take
// the same branch.
//
// Per pearl_cooperative_staging_eliminates_redundant_reads: the per-batch
// `x` and `h_old` rows are SHARED across all output channels in a block
// (every thread reads HIDDEN_DIM elements of both in the K-loop). Stage
// them into shared memory once at block entry. W_in / W_rec rows are
// per-channel (different per thread c), so they remain direct DRAM reads.
//
// Per feedback_no_atomicadd.md: no atomicAdd anywhere. Each thread writes
// to its own (batch, c) slot in forward, and to its own per-batch grad
// slice in backward. Cross-batch reduction is the caller's responsibility
// via the existing `reduce_axis0` infrastructure.
//
// Per pearl_no_host_branches_in_captured_graph: no host branching; all
// control flow uses device-resident metadata read from device tensors.
//
// Per feedback_nvidia_grade_perf_for_kernels.md: warp-uniform branches,
// no atomicAdd, coalesced loads via cooperative staging.
//
// ALPHA fix (2026-05-21): the kernel now indexes into ORIGINAL channel
// layout via `channels_in_bucket[branch][tid]` — abandons the prior
// bucket-grouped τ layout that mis-aligned with W_in/W_rec/b/tau_all
// (which all stayed in original-channel layout). See
// `bucket_transition_kernels.cu::channels_in_bucket_kernel` for the
// lookup table.

#define HIDDEN_DIM 128
#define N_HORIZONS 3
#define MAX_BUCKET_DIM 96

// ─────────────────────────────────────────────────────────────────────
// cfc_step_per_branch_fwd: forward pass.
//
// Launch:
//   grid  = (B, N_HORIZONS, 1)
//   block = (MAX_BUCKET_DIM = 96, 1, 1)
//   shared_mem_bytes = 2 * HIDDEN_DIM * sizeof(float) = 1024 bytes
//
// Per (batch, branch, thread-in-bucket) thread computes:
//   c = channels_in_bucket[branch * MAX_BUCKET_DIM + tid]  (original channel index)
//   pre = b[c] + Σ_k W_in[c, k] * x[batch, k] + Σ_k W_rec[c, k] * h_old[batch, k]
//   decay = exp(-dt / max(tau_all[c], 1e-6))
//   h_new[batch, c] = h_old[batch, c] * decay + (1 - decay) * tanh(pre)
//
// x_local[HIDDEN_DIM] and h_old_local[HIDDEN_DIM] are cooperative-staged
// in shared memory; without them every thread re-reads the full row
// (HIDDEN_DIM × bucket_dim threads → up to 96 × redundant reads per row).
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void cfc_step_per_branch_fwd(
    const float* __restrict__ w_in,                          // [HIDDEN_DIM × HIDDEN_DIM] original channel layout
    const float* __restrict__ w_rec,                         // [HIDDEN_DIM × HIDDEN_DIM] original channel layout
    const float* __restrict__ b,                             // [HIDDEN_DIM] original channel layout
    const float* __restrict__ tau_all,                       // [HIDDEN_DIM] original channel layout
    const float* __restrict__ x,                             // [B × HIDDEN_DIM] original channel layout
    const float* __restrict__ h_old,                         // [B × HIDDEN_DIM] original channel layout
    float dt,
    int B,
    const unsigned int* __restrict__ channels_in_bucket,     // [N_HORIZONS × MAX_BUCKET_DIM]
    const unsigned int* __restrict__ bucket_dim_k,           // [N_HORIZONS]
    float* __restrict__ h_new                                // [B × HIDDEN_DIM] original channel layout
) {
    extern __shared__ float smem[];
    float* x_local     = smem;                  // [HIDDEN_DIM]
    float* h_old_local = smem + HIDDEN_DIM;     // [HIDDEN_DIM]

    int batch  = blockIdx.x;
    int branch = blockIdx.y;
    int tid    = threadIdx.x;

    if (batch >= B) return;

    // Cooperative staging of per-batch x[batch, *] and h_old[batch, *] rows.
    // All threads (including those that will early-return on the per-branch
    // predicate) participate in the staging — the row is shared across the
    // block's output channels, so we need every thread to help load.
    //
    // Each thread loads HIDDEN_DIM / blockDim.x = 128 / 96 ≈ 2 elements
    // (rounded up via the stride loop).
    for (int i = tid; i < HIDDEN_DIM; i += blockDim.x) {
        x_local[i]     = x[batch * HIDDEN_DIM + i];
        h_old_local[i] = h_old[batch * HIDDEN_DIM + i];
    }
    __syncthreads();

    // Uniform per-branch predicate: idle threads in shorter buckets early-return.
    // The comparison is against a per-block constant (bucket_dim_k[branch]),
    // so all threads in a warp take the same path → no warp divergence.
    unsigned int bucket_dim = bucket_dim_k[branch];
    if ((unsigned int)tid >= bucket_dim) return;

    // Look up the ORIGINAL channel index for this (branch, tid) position.
    unsigned int c_u = channels_in_bucket[branch * MAX_BUCKET_DIM + tid];
    // Defensive: sentinel (0xFFFFFFFF) or out-of-range indices must not
    // write h_new. By construction of channels_in_bucket_kernel, all
    // tid < bucket_dim_k[branch] entries hold valid channel indices.
    if (c_u >= (unsigned int)HIDDEN_DIM) return;
    int c = (int)c_u;

    // pre = b[c] + Σ_k W_in[c, k] * x_local[k] + Σ_k W_rec[c, k] * h_old_local[k]
    float pre = b[c];
    for (int k = 0; k < HIDDEN_DIM; ++k) {
        pre += w_in[c * HIDDEN_DIM + k] * x_local[k];
    }
    for (int k = 0; k < HIDDEN_DIM; ++k) {
        pre += w_rec[c * HIDDEN_DIM + k] * h_old_local[k];
    }

    float decay = expf(-dt / fmaxf(tau_all[c], 1e-6f));
    h_new[batch * HIDDEN_DIM + c] =
        h_old_local[c] * decay + (1.0f - decay) * tanhf(pre);
}

// ─────────────────────────────────────────────────────────────────────
// cfc_step_per_branch_bwd: backward pass.
//
// Launch:
//   grid  = (B, N_HORIZONS, 1)
//   block = (MAX_BUCKET_DIM = 96, 1, 1)
//   shared_mem_bytes = 2 * HIDDEN_DIM * sizeof(float) = 1024 bytes
//
// Per-batch grad slices are written by this kernel; cross-batch
// reduction is the caller's responsibility (existing `reduce_axis0`
// pattern; see crates/ml-alpha/cuda/cfc_step.cu Phase B comment block).
//
// grad shapes:
//   grad_w_in   : [B × HIDDEN_DIM × HIDDEN_DIM]  per-batch scratch (overwrite)
//   grad_w_rec  : [B × HIDDEN_DIM × HIDDEN_DIM]  per-batch scratch (overwrite)
//   grad_b      : [B × HIDDEN_DIM]               per-batch scratch (overwrite)
//   grad_tau_all: [B × HIDDEN_DIM]               per-batch scratch (overwrite)
//   grad_h_old  : [B × HIDDEN_DIM]               per-batch (overwrite)
//
// Note: this kernel uses OVERWRITE semantics (the caller zeros scratch
// per training step). The single-step CfC backward in the existing
// `cfc_step_backward_batched` uses += for K-loop accumulation; for
// per-branch usage in Phase 2 the K-loop's per-position accumulation is
// the caller's concern (matching existing scratch + reduce pattern).
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void cfc_step_per_branch_bwd(
    const float* __restrict__ w_in,                           // [HIDDEN_DIM × HIDDEN_DIM] original channel layout
    const float* __restrict__ w_rec,                          // [HIDDEN_DIM × HIDDEN_DIM] original channel layout
    const float* __restrict__ b,                              // [HIDDEN_DIM] original channel layout
    const float* __restrict__ tau_all,                        // [HIDDEN_DIM] original channel layout
    const float* __restrict__ x,                              // [B × HIDDEN_DIM] original channel layout
    const float* __restrict__ h_old,                          // [B × HIDDEN_DIM] original channel layout
    const float* __restrict__ grad_h_new,                     // [B × HIDDEN_DIM] original channel layout
    float dt,
    int B,
    const unsigned int* __restrict__ channels_in_bucket,      // [N_HORIZONS × MAX_BUCKET_DIM]
    const unsigned int* __restrict__ bucket_dim_k,            // [N_HORIZONS]
    float* __restrict__ grad_w_in,    // [B × HIDDEN_DIM × HIDDEN_DIM]
    float* __restrict__ grad_w_rec,   // [B × HIDDEN_DIM × HIDDEN_DIM]
    float* __restrict__ grad_b,       // [B × HIDDEN_DIM]
    float* __restrict__ grad_tau_all, // [B × HIDDEN_DIM]
    float* __restrict__ grad_h_old    // [B × HIDDEN_DIM]
) {
    extern __shared__ float smem[];
    float* x_local     = smem;                  // [HIDDEN_DIM]
    float* h_old_local = smem + HIDDEN_DIM;     // [HIDDEN_DIM]

    int batch  = blockIdx.x;
    int branch = blockIdx.y;
    int tid    = threadIdx.x;

    if (batch >= B) return;

    // Cooperative staging — same pattern as forward.
    for (int i = tid; i < HIDDEN_DIM; i += blockDim.x) {
        x_local[i]     = x[batch * HIDDEN_DIM + i];
        h_old_local[i] = h_old[batch * HIDDEN_DIM + i];
    }
    __syncthreads();

    unsigned int bucket_dim = bucket_dim_k[branch];
    if ((unsigned int)tid >= bucket_dim) return;

    // Look up the ORIGINAL channel index for this (branch, tid) position.
    unsigned int c_u = channels_in_bucket[branch * MAX_BUCKET_DIM + tid];
    if (c_u >= (unsigned int)HIDDEN_DIM) return;
    int c = (int)c_u;

    // Recompute forward pre + tanh to derive d_pre.
    float pre = b[c];
    for (int k = 0; k < HIDDEN_DIM; ++k) {
        pre += w_in[c * HIDDEN_DIM + k] * x_local[k];
    }
    for (int k = 0; k < HIDDEN_DIM; ++k) {
        pre += w_rec[c * HIDDEN_DIM + k] * h_old_local[k];
    }
    const float tau_eps = 1e-6f;
    const float tau_raw = tau_all[c];
    const float tau_eff = fmaxf(tau_raw, tau_eps);
    const float decay   = expf(-dt / tau_eff);
    const float s       = tanhf(pre);
    const float dh      = grad_h_new[batch * HIDDEN_DIM + c];
    const float d_pre   = dh * (1.0f - decay) * (1.0f - s * s);
    const float d_decay = dh * (h_old_local[c] - s);

    // Per-batch grad scratch writes — each thread (batch, branch, c) is the
    // sole writer of grad_*[batch, c, ...]. No race; no atomicAdd.
    grad_b[batch * HIDDEN_DIM + c] = d_pre;

    // grad_tau receives 0 when tau is at the clamp floor (strict d(max)/d(tau) = 0).
    const float gate = (tau_raw > tau_eps) ? 1.0f : 0.0f;
    grad_tau_all[batch * HIDDEN_DIM + c] =
        gate * d_decay * decay * dt / (tau_eff * tau_eff);

    // grad_w_in[batch, c, k] = d_pre * x_local[k]
    // grad_w_rec[batch, c, k] = d_pre * h_old_local[k]
    //
    // Layout note: thread serves the (batch, c) role, sweeping k in the inner
    // loop. Writes along k are coalesced within a thread but strided across
    // threads (each thread writes its own [c, *] row in a per-(batch) slab).
    // Per pearl_coalesce_via_thread_role_swap: this is single-thread-writes-
    // single-row, the warp-coalescing concern is on cross-thread writes to
    // the SAME row, which we don't do here. No swap needed.
    for (int k = 0; k < HIDDEN_DIM; ++k) {
        const long long off =
            (long long)batch * HIDDEN_DIM * HIDDEN_DIM
            + (long long)c * HIDDEN_DIM + k;
        grad_w_in[off]  = d_pre * x_local[k];
        grad_w_rec[off] = d_pre * h_old_local[k];
    }

    // grad_h_old[batch, c] receives the direct decay contribution. The cross-
    // channel term Σ_j d_pre[j] * W_rec[j, c] requires a reduction across
    // threads in the block — the caller may run a second-pass reduction
    // kernel when grad_h_old needs to feed backward through a prior CfC
    // step (currently CfC is K=1 in per-branch Phase 2, so this direct
    // contribution is the only one).
    grad_h_old[batch * HIDDEN_DIM + c] = dh * decay;
}