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

// heads_block_diagonal_fwd.cu — heads forward with compact ragged w_skip.
//
// Per spec §5.4 point 2 (docs/superpowers/specs/2026-05-21-per-horizon-cfc-inference-design.md):
//   heads_w_skip storage is compacted from [N_HORIZONS × HIDDEN_DIM] = 384
//   floats with off-bucket zeros to a ragged buffer of size HIDDEN_DIM = 128
//   floats. Each horizon head reads ONLY its bucket's contiguous slice via
//   heads_w_skip_offset[N_HORIZONS+1].
//
// Layout (for HIDDEN_DIM=128, BUCKET_DIM_K=[43,43,42]):
//   w_skip_compact: [43 (head 0) | 43 (head 1) | 42 (head 2)]
//   heads_w_skip_offset: [0, 43, 86, 128]
//
// Skip projection for horizon h, batch b:
//   skip_logit[b, h] = b_skip[h]
//                    + Σ_{c=0..bucket_dim_k[h]} h_state[b, bucket_channel_offset[h] + c]
//                                              * w_skip_compact[heads_w_skip_offset[h] + c]
//
// Launch: grid = (B, N_HORIZONS, 1), block = (MAX_BUCKET_DIM = 96, 1, 1).
// Uniform predicate (threadIdx.x < bucket_dim) handles uneven bucket sizes
// without warp divergence — the comparison is against a per-block constant,
// so all threads in a warp take the same branch.
//
// Per feedback_no_atomicadd.md: block-tree reduction with warp-shuffle-style
// shared-memory pairwise pattern; never atomicAdd.
// Per pearl_no_host_branches_in_captured_graph.md: no host branching; all
// control flow uses device-resident metadata.
// Per feedback_nvidia_grade_perf_for_kernels.md: warp-uniform branches,
// block-tree reduction, no atomicAdd.

#define HIDDEN_DIM 128
#define N_HORIZONS 3
#define MAX_BUCKET_DIM 96
// Power-of-two padding for the block-tree reduction. Lanes
// [MAX_BUCKET_DIM .. REDUCE_PAD) are zero-initialized in shared memory so
// the halving-stride reduction is correct for the non-power-of-two
// MAX_BUCKET_DIM=96 case. Next power of two above 96 is 128.
#define REDUCE_PAD 128

// ─────────────────────────────────────────────────────────────────────
// heads_block_diagonal_fwd: compact ragged w_skip projection.
//
// Launch:
//   grid  = (B, N_HORIZONS, 1)
//   block = (MAX_BUCKET_DIM = 96, 1, 1)
//   shared_mem_bytes = 0  (uses statically-sized __shared__ buffer)
//
// Output:
//   skip_logit[batch, horizon] = b_skip[horizon] + Σ_c h_state[batch, bucket_start+c] *
//                                                      w_skip_compact[w_start+c]
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void heads_block_diagonal_fwd(
    // Compact ragged w_skip:
    const float* __restrict__ w_skip_compact,                  // [HIDDEN_DIM = sum(bucket_dim_k)]
    const unsigned int* __restrict__ heads_w_skip_offset,      // [N_HORIZONS + 1]
    const unsigned int* __restrict__ bucket_channel_offset,    // [N_HORIZONS + 1]
    const unsigned int* __restrict__ bucket_dim_k,             // [N_HORIZONS]
    const float* __restrict__ b_skip,                          // [N_HORIZONS]
    // Inputs
    const float* __restrict__ h_state,                         // [B × HIDDEN_DIM]
    int B,
    // Outputs
    float* __restrict__ skip_logit                             // [B × N_HORIZONS]
) {
    int batch   = blockIdx.x;
    int horizon = blockIdx.y;
    int tid     = threadIdx.x;

    if (batch >= B) return;

    unsigned int bucket_dim   = bucket_dim_k[horizon];
    unsigned int bucket_start = bucket_channel_offset[horizon];
    unsigned int w_start      = heads_w_skip_offset[horizon];

    // Shared accumulator sized to REDUCE_PAD (128, next power of two ≥
    // MAX_BUCKET_DIM=96). Lanes [bucket_dim .. REDUCE_PAD) hold 0.0 so the
    // standard power-of-two halving-stride reduction is correct.
    //
    // Only blockDim.x = MAX_BUCKET_DIM = 96 threads exist; lanes
    // 96..127 are written by tid<32 (those threads write BOTH their own slot
    // and the padding slot at tid+MAX_BUCKET_DIM).
    __shared__ float sdata[REDUCE_PAD];
    sdata[tid] = 0.0f;
    if (tid + MAX_BUCKET_DIM < REDUCE_PAD) {
        sdata[tid + MAX_BUCKET_DIM] = 0.0f;
    }
    __syncthreads();

    if ((unsigned int)tid < bucket_dim) {
        unsigned int c     = bucket_start + (unsigned int)tid;
        unsigned int w_idx = w_start + (unsigned int)tid;
        // Defensive: c < HIDDEN_DIM by construction (bucket_channel_offset[N_HORIZONS] == HIDDEN_DIM).
        sdata[tid] = h_state[batch * HIDDEN_DIM + c] * w_skip_compact[w_idx];
    }
    __syncthreads();

    // Block-tree reduction over REDUCE_PAD=128 padded lanes (no atomicAdd per
    // feedback_no_atomicadd.md). Power-of-two halving stride is clean since
    // padding lanes are zero-initialized.
    for (int s = REDUCE_PAD / 2; s > 0; s >>= 1) {
        if (tid < s) {
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }

    if (tid == 0) {
        skip_logit[batch * N_HORIZONS + horizon] = sdata[0] + b_skip[horizon];
    }
}