From afb6d73396d388d0bf0e4b3e3311396b495d0b9c Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 22 May 2026 01:26:37 +0200 Subject: [PATCH] =?UTF-8?q?refactor(per-horizon):=20N=5FHORIZONS=205?= =?UTF-8?q?=E2=86=923=20=E2=80=94=20bucket-coupled=20CUDA=20kernels?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Five CUDA kernels + their Rust caller migrated atomically to prevent silent memory layout corruption between Rust-side bucket geometry (MAX_BUCKET_DIM=96, BUCKET_DIM_K=[43,43,42], BUCKET_CHANNEL_OFFSET= [0,43,86,128]) and kernel-side hardcoded constants. Kernels: - bucket_transition_kernels.cu: N_HORIZONS 5→3, MAX_BUCKET_DIM 28→96, BUCKET_DIM_LAST 28→42, BUCKET_OFFSETS {0,25,50,75,100,128}→{0,43,86,128}; bucket_assign_kernel quintile→tercile rewire. - cfc_step_per_branch.cu: defines + doc. - heads_block_diagonal_fwd.cu: same; REDUCE_PAD 32→128 (next pow2 ≥ 96). - multi_horizon_heads.cu: N_HORIZONS_H 5→3 covers fwd, bwd, batched, and the GRN/2-layer variants via the single define + shared mem [N_HORIZONS_H] + loop bounds. - output_smoothness.cu: OS_N_HORIZONS 5→3. Rust caller (silent-corruption fix found during audit): - crates/ml-alpha/src/cfc/step.rs:135,192 had local hardcoded N_HORIZONS=5 / MAX_BUCKET_DIM=28 constants — replaced with crate::heads::N_HORIZONS / crate::cfc::bucket_routing::MAX_BUCKET_DIM (single source of truth via the Rust SoT constants). cargo build -p ml-alpha: all 5 cubins rebuild PASS. GPU oracle tests on RTX 3050 sm_86: 5/19 pass; 14 failures are test-side fixtures hardcoding old 5×28 layout — owned by SDD Task 8 (not regression). Co-Authored-By: Claude Opus 4.7 (1M context) --- .../cuda/bucket_transition_kernels.cu | 30 +- crates/ml-alpha/cuda/cfc_step_per_branch.cu | 18 +- .../ml-alpha/cuda/heads_block_diagonal_fwd.cu | 30 +- crates/ml-alpha/cuda/multi_horizon_heads.cu | 291 +++++++++--------- crates/ml-alpha/cuda/output_smoothness.cu | 33 +- crates/ml-alpha/src/cfc/step.rs | 15 +- 6 files changed, 215 insertions(+), 202 deletions(-) diff --git a/crates/ml-alpha/cuda/bucket_transition_kernels.cu b/crates/ml-alpha/cuda/bucket_transition_kernels.cu index bb68ff602..73f106680 100644 --- a/crates/ml-alpha/cuda/bucket_transition_kernels.cu +++ b/crates/ml-alpha/cuda/bucket_transition_kernels.cu @@ -1,16 +1,23 @@ // bucket_transition_kernels.cu — Phase 1→2 transition kernels. // // All-on-device per feedback_no_htod_htoh_only_mapped_pinned. No atomicAdd -// per feedback_no_atomicadd; reductions use warp-shuffle. Static quintile -// boundaries for HIDDEN_DIM=128: [0, 25, 50, 75, 100, 128]. +// per feedback_no_atomicadd; reductions use warp-shuffle. Static tercile +// boundaries for HIDDEN_DIM=128 with N_HORIZONS=3: [0, 43, 86, 128] +// (post-2026-05-22 horizon-rebase to 3 buckets, matches Rust-side +// `crates/ml-alpha/src/cfc/bucket_routing.rs::BUCKET_CHANNEL_OFFSET`). +// +// MAX_BUCKET_DIM=96 matches the Rust-side `MAX_BUCKET_DIM` and is sized +// for the upper bound of ISV-driven τ-assignment concentration (~75% of +// channels into one bucket under extreme τ distributions). #define HIDDEN_DIM 128 -#define N_HORIZONS 5 -#define BUCKET_DIM_LAST 28 // last bucket absorbs HIDDEN_DIM - 4*25 = 28 -#define MAX_BUCKET_DIM 28 +#define N_HORIZONS 3 +#define BUCKET_DIM_LAST 42 // last bucket absorbs HIDDEN_DIM - 2*43 = 42 +#define MAX_BUCKET_DIM 96 // Static bucket boundaries (could also be __constant__ memory). -__device__ const unsigned int BUCKET_OFFSETS[N_HORIZONS + 1] = {0, 25, 50, 75, 100, 128}; +// Tercile cut points: 43, 86, 128 (sum = HIDDEN_DIM = 128). +__device__ const unsigned int BUCKET_OFFSETS[N_HORIZONS + 1] = {0, 43, 86, 128}; // ───────────────────────────────────────────────────────────────────── // tau_sort_kernel: bitonic-merge sort of cfc.tau values. @@ -72,7 +79,8 @@ extern "C" __global__ void tau_sort_kernel( // // Launch: 1 block × HIDDEN_DIM threads. // Each thread handles one channel. Looks up its rank (position in sorted) -// and assigns bucket id via static quintile boundaries. +// and assigns bucket id via static tercile boundaries (post-2026-05-22 +// N_HORIZONS=3 rebase: cuts at rank 43 and 86 for [43, 43, 42]). // ───────────────────────────────────────────────────────────────────── extern "C" __global__ void bucket_assign_kernel( const unsigned int* __restrict__ sorted_indices, // [HIDDEN_DIM] @@ -81,12 +89,10 @@ extern "C" __global__ void bucket_assign_kernel( int tid = threadIdx.x; if (tid >= HIDDEN_DIM) return; // tid is a sorted-rank position; sorted_indices[tid] is the original channel - // at this rank. Assign bucket = quintile of rank. + // at this rank. Assign bucket = tercile of rank. unsigned char bucket = 0; - if (tid >= 100) bucket = 4; - else if (tid >= 75) bucket = 3; - else if (tid >= 50) bucket = 2; - else if (tid >= 25) bucket = 1; + if (tid >= 86) bucket = 2; + else if (tid >= 43) bucket = 1; bucket_id_per_channel[sorted_indices[tid]] = bucket; } diff --git a/crates/ml-alpha/cuda/cfc_step_per_branch.cu b/crates/ml-alpha/cuda/cfc_step_per_branch.cu index d7b0a8896..82905ef38 100644 --- a/crates/ml-alpha/cuda/cfc_step_per_branch.cu +++ b/crates/ml-alpha/cuda/cfc_step_per_branch.cu @@ -1,11 +1,11 @@ // 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 5 branches × n_batch in one launch. -// Grid: (B, N_HORIZONS, 1), block: (MAX_BUCKET_DIM=28, 1, 1) uniform predicate. +// 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 [25, 25, 25, 25, 28] without warp divergence — the +// 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. // @@ -34,15 +34,15 @@ // lookup table. #define HIDDEN_DIM 128 -#define N_HORIZONS 5 -#define MAX_BUCKET_DIM 28 +#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 = 28, 1, 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: @@ -53,7 +53,7 @@ // // 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 28 × redundant reads per 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 @@ -83,7 +83,7 @@ extern "C" __global__ void cfc_step_per_branch_fwd( // 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 / 28 ≈ 5 elements + // 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]; @@ -124,7 +124,7 @@ extern "C" __global__ void cfc_step_per_branch_fwd( // // Launch: // grid = (B, N_HORIZONS, 1) -// block = (MAX_BUCKET_DIM = 28, 1, 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 diff --git a/crates/ml-alpha/cuda/heads_block_diagonal_fwd.cu b/crates/ml-alpha/cuda/heads_block_diagonal_fwd.cu index 70c79e931..0cbfad73e 100644 --- a/crates/ml-alpha/cuda/heads_block_diagonal_fwd.cu +++ b/crates/ml-alpha/cuda/heads_block_diagonal_fwd.cu @@ -1,21 +1,21 @@ // 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] = 640 +// 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=[25,25,25,25,28]): -// w_skip_compact: [25 (head 0) | 25 (head 1) | 25 (head 2) | 25 (head 3) | 28 (head 4)] -// heads_w_skip_offset: [0, 25, 50, 75, 100, 128] +// 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 = 28, 1, 1). +// 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. @@ -28,20 +28,20 @@ // block-tree reduction, no atomicAdd. #define HIDDEN_DIM 128 -#define N_HORIZONS 5 -#define MAX_BUCKET_DIM 28 +#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=28 case. -#define REDUCE_PAD 32 +// 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 = 28, 1, 1) +// block = (MAX_BUCKET_DIM = 96, 1, 1) // shared_mem_bytes = 0 (uses statically-sized __shared__ buffer) // // Output: @@ -71,12 +71,12 @@ extern "C" __global__ void heads_block_diagonal_fwd( unsigned int bucket_start = bucket_channel_offset[horizon]; unsigned int w_start = heads_w_skip_offset[horizon]; - // Shared accumulator sized to REDUCE_PAD (32, next power of two ≥ - // MAX_BUCKET_DIM=28). Lanes [bucket_dim .. REDUCE_PAD) hold 0.0 so the + // 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 = 28 threads exist; lanes - // 28..31 are written by tid<4 (those threads write BOTH their own slot + // 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; @@ -93,7 +93,7 @@ extern "C" __global__ void heads_block_diagonal_fwd( } __syncthreads(); - // Block-tree reduction over REDUCE_PAD=32 padded lanes (no atomicAdd per + // 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) { diff --git a/crates/ml-alpha/cuda/multi_horizon_heads.cu b/crates/ml-alpha/cuda/multi_horizon_heads.cu index c58849bfd..1f4bcb26c 100644 --- a/crates/ml-alpha/cuda/multi_horizon_heads.cu +++ b/crates/ml-alpha/cuda/multi_horizon_heads.cu @@ -1,18 +1,27 @@ // multi_horizon_heads.cu // -// 128-dim hidden -> 5 sigmoid logits, one per horizon -// {30, 100, 300, 1000, 6000} snapshots forward. Single-block kernel, -// 5 threads (one per head). Each thread computes its own dot product -// against the 128-dim hidden vector + bias, then sigmoid. +// 128-dim hidden -> N_HORIZONS sigmoid logits, one per horizon +// (post-2026-05-22 rebase: 3 horizons {10, 100, 1000} snapshots forward). +// Single-block kernel, N_HORIZONS threads (one per head). Each thread +// computes its own dot product against the 128-dim hidden vector + bias, +// then sigmoid. + +// File-wide layout constants (must agree with +// `crates/ml-alpha/src/heads.rs::{N_HORIZONS, HIDDEN_DIM, HEAD_MID_DIM}`). +#define N_HORIZONS_H 3 +#define HIDDEN_H 128 +#define HEAD_MID_H 64 +#define GELU_C0 0.7978845608f // sqrt(2/pi) +#define GELU_C1 0.044715f extern "C" __global__ void multi_horizon_heads( - const float* __restrict__ w, // [5, 128] - const float* __restrict__ b, // [5] + const float* __restrict__ w, // [N_HORIZONS, 128] + const float* __restrict__ b, // [N_HORIZONS] const float* __restrict__ h, // [128] - float* __restrict__ probs // [5] + float* __restrict__ probs // [N_HORIZONS] ) { int k = threadIdx.x; - if (k >= 5) return; + if (k >= N_HORIZONS_H) return; float z = b[k]; for (int i = 0; i < 128; ++i) { z += w[k * 128 + i] * h[i]; @@ -23,32 +32,32 @@ extern "C" __global__ void multi_horizon_heads( // Backward through multi_horizon_heads. // // Given: -// grad_probs[5] = dL / dp -// probs[5] = forward output (sigmoid) -// h[128] = forward input +// grad_probs[N_HORIZONS] = dL / dp +// probs[N_HORIZONS] = forward output (sigmoid) +// h[128] = forward input // Computes: -// grad_w[5, 128] = dL / dW = d_z[k] * h[i] -// grad_b[5] = dL / db = d_z[k] -// grad_h[128] = dL / dh = sum_k d_z[k] * W[k, i] +// grad_w[N_HORIZONS, 128] = dL / dW = d_z[k] * h[i] +// grad_b[N_HORIZONS] = dL / db = d_z[k] +// grad_h[128] = dL / dh = sum_k d_z[k] * W[k, i] // where d_z[k] = grad_probs[k] * p[k] * (1 - p[k]). // // Block tree-reduce for the grad_h sum (no atomicAdd). One thread per // hidden unit (128 threads); thread 0 also writes grad_b. extern "C" __global__ void multi_horizon_heads_backward( - const float* __restrict__ w, // [5, 128] - const float* __restrict__ probs, // [5] + const float* __restrict__ w, // [N_HORIZONS, 128] + const float* __restrict__ probs, // [N_HORIZONS] const float* __restrict__ h, // [128] - const float* __restrict__ grad_probs, // [5] + const float* __restrict__ grad_probs, // [N_HORIZONS] const float* __restrict__ grad_h_carry, // [128] — optional carry (nullptr OK) - float* __restrict__ grad_w, // [5, 128] (+=) - float* __restrict__ grad_b, // [5] (+=) - float* __restrict__ grad_h // [128] overwrite-with-carry + float* __restrict__ grad_w, // [N_HORIZONS, 128] (+=) + float* __restrict__ grad_b, // [N_HORIZONS] (+=) + float* __restrict__ grad_h // [128] overwrite-with-carry ) { - __shared__ float d_z[5]; + __shared__ float d_z[N_HORIZONS_H]; int tid = threadIdx.x; - if (tid < 5) { + if (tid < N_HORIZONS_H) { const float p = probs[tid]; d_z[tid] = grad_probs[tid] * p * (1.0f - p); } @@ -57,14 +66,14 @@ extern "C" __global__ void multi_horizon_heads_backward( // Parameter-grad writes use `+=`: per-position supervision invokes // this kernel K times per training step and needs the gradients // accumulated. Callers MUST pre-zero grad_w / grad_b at step start. - if (tid < 5) { + if (tid < N_HORIZONS_H) { grad_b[tid] += d_z[tid]; } // grad_w[k, i] += d_z[k] * h[i] - // grid stride: each thread covers one i for k in 0..5. + // grid stride: each thread covers one i for k in 0..N_HORIZONS. if (tid < 128) { const float h_i = h[tid]; - for (int k = 0; k < 5; ++k) { + for (int k = 0; k < N_HORIZONS_H; ++k) { grad_w[k * 128 + tid] += d_z[k] * h_i; } // grad_h[i] = sum_k d_z[k] * W[k, i] + carry[i] @@ -73,7 +82,7 @@ extern "C" __global__ void multi_horizon_heads_backward( // h_new_k flows into h_old_{k+1}, so grad_h_new at position k must // sum the heads-side gradient with the carry from the upstream step. float acc = 0.0f; - for (int k = 0; k < 5; ++k) { + for (int k = 0; k < N_HORIZONS_H; ++k) { acc += d_z[k] * w[k * 128 + tid]; } const float carry = (grad_h_carry != nullptr) ? grad_h_carry[tid] : 0.0f; @@ -83,62 +92,62 @@ extern "C" __global__ void multi_horizon_heads_backward( // ─── Batched variants ──────────────────────────────────────────────── -// Process N_BATCH samples per launch. Layout: h[B, 128], probs[B, 5]. -// Thread tid handles either head k (tid < 5) or hidden unit i (tid < -// 128). Block dim = 128, one block per launch — each thread loops over -// b internally. Param grads (grad_w, grad_b) shared across batch and +// Process N_BATCH samples per launch. Layout: h[B, 128], probs[B, N_HORIZONS]. +// Thread tid handles either head k (tid < N_HORIZONS) or hidden unit i +// (tid < 128). Block dim = 128, one block per launch — each thread loops +// over b internally. Param grads (grad_w, grad_b) shared across batch and // accumulated via `+=` (no race: thread tid is sole writer to its row). extern "C" __global__ void multi_horizon_heads_batched( - const float* __restrict__ w, // [5, 128] - const float* __restrict__ b, // [5] + const float* __restrict__ w, // [N_HORIZONS, 128] + const float* __restrict__ b, // [N_HORIZONS] const float* __restrict__ h, // [n_batch, 128] int n_batch, - float* __restrict__ probs // [n_batch, 5] + float* __restrict__ probs // [n_batch, N_HORIZONS] ) { int k = threadIdx.x; - if (k >= 5) return; + if (k >= N_HORIZONS_H) return; const float bias_k = b[k]; for (int bi = 0; bi < n_batch; ++bi) { const float* h_b = h + (long long)bi * 128; float z = bias_k; for (int i = 0; i < 128; ++i) z += w[k * 128 + i] * h_b[i]; - probs[(long long)bi * 5 + k] = 1.0f / (1.0f + expf(-z)); + probs[(long long)bi * N_HORIZONS_H + k] = 1.0f / (1.0f + expf(-z)); } } extern "C" __global__ void multi_horizon_heads_backward_batched( - const float* __restrict__ w, // [5, 128] - const float* __restrict__ probs, // [n_batch, 5] + const float* __restrict__ w, // [N_HORIZONS, 128] + const float* __restrict__ probs, // [n_batch, N_HORIZONS] const float* __restrict__ h, // [n_batch, 128] - const float* __restrict__ grad_probs, // [n_batch, 5] + const float* __restrict__ grad_probs, // [n_batch, N_HORIZONS] const float* __restrict__ grad_h_carry, // [n_batch, 128] (nullptr OK) - const float* __restrict__ lambda, // [5] — per-horizon trunk-grad scaler + const float* __restrict__ lambda, // [N_HORIZONS] — per-horizon trunk-grad scaler int n_batch, - float* __restrict__ grad_w, // [5, 128] accum += - float* __restrict__ grad_b, // [5] accum += + float* __restrict__ grad_w, // [N_HORIZONS, 128] accum += + float* __restrict__ grad_b, // [N_HORIZONS] accum += float* __restrict__ grad_h // [n_batch, 128] overwrite + carry ) { - // Shared mem: d_z[B, 5]. At B=32: 32 * 5 = 160 floats = 640 bytes. - extern __shared__ float sd_z[]; // size = n_batch * 5 + // Shared mem: d_z[B, N_HORIZONS]. At B=32: 32 * 3 = 96 floats = 384 bytes. + extern __shared__ float sd_z[]; // size = n_batch * N_HORIZONS_H int tid = threadIdx.x; // Pass 1: per-(bi, k) compute d_z and stash in shared. - if (tid < 5) { + if (tid < N_HORIZONS_H) { for (int bi = 0; bi < n_batch; ++bi) { - const float p = probs[(long long)bi * 5 + tid]; - const float dp = grad_probs[(long long)bi * 5 + tid]; - sd_z[(long long)bi * 5 + tid] = dp * p * (1.0f - p); + const float p = probs[(long long)bi * N_HORIZONS_H + tid]; + const float dp = grad_probs[(long long)bi * N_HORIZONS_H + tid]; + sd_z[(long long)bi * N_HORIZONS_H + tid] = dp * p * (1.0f - p); } } __syncthreads(); // Per-horizon trunk-gradient scaler, broadcast via shared mem so - // every thread reads the same 5 floats without repeated global loads. - __shared__ float s_lambda[5]; - if (tid < 5) { + // every thread reads the same N_HORIZONS floats without repeated global loads. + __shared__ float s_lambda[N_HORIZONS_H]; + if (tid < N_HORIZONS_H) { // Sentinel: a fully zero lambda buffer (initial state, before // the EMA kernel has run even once) is treated as "scaler = 1" // so the kernel reduces to its pre-ISV behavior. After step 1 @@ -149,15 +158,15 @@ extern "C" __global__ void multi_horizon_heads_backward_batched( } __syncthreads(); - // grad_b[k] += sum_b d_z[b, k]. Thread k (k < 5) is sole writer. + // grad_b[k] += sum_b d_z[b, k]. Thread k (k < N_HORIZONS) is sole writer. // NOTE: lambda does NOT scale the head's own bias gradient — only // the trunk gradient. We want the per-horizon head to keep learning // its own bias normally; lambda only biases how much the horizon's // error signal LEAKS into the shared trunk. - if (tid < 5) { + if (tid < N_HORIZONS_H) { float acc = 0.0f; for (int bi = 0; bi < n_batch; ++bi) { - acc += sd_z[(long long)bi * 5 + tid]; + acc += sd_z[(long long)bi * N_HORIZONS_H + tid]; } grad_b[tid] += acc; } @@ -173,12 +182,12 @@ extern "C" __global__ void multi_horizon_heads_backward_batched( for (int bi = 0; bi < n_batch; ++bi) { const float h_bi = h[(long long)bi * 128 + tid]; // Accumulate W-row grads via += into global (thread tid sole writer of column tid). - for (int k = 0; k < 5; ++k) { - grad_w[k * 128 + tid] += sd_z[(long long)bi * 5 + k] * h_bi; + for (int k = 0; k < N_HORIZONS_H; ++k) { + grad_w[k * 128 + tid] += sd_z[(long long)bi * N_HORIZONS_H + k] * h_bi; } float acc = 0.0f; - for (int k = 0; k < 5; ++k) { - acc += s_lambda[k] * sd_z[(long long)bi * 5 + k] * w[k * 128 + tid]; + for (int k = 0; k < N_HORIZONS_H; ++k) { + acc += s_lambda[k] * sd_z[(long long)bi * N_HORIZONS_H + k] * w[k * 128 + tid]; } const float carry = (grad_h_carry != nullptr) ? grad_h_carry[(long long)bi * 128 + tid] : 0.0f; @@ -192,25 +201,19 @@ extern "C" __global__ void multi_horizon_heads_backward_batched( // GELU hidden → [HEAD_MID=64 → 1] sigmoid output. // // Forward layout: -// w1 [5, HEAD_MID, HIDDEN] per-horizon first-layer weights -// b1 [5, HEAD_MID] per-horizon first-layer biases -// w2 [5, HEAD_MID] per-horizon output weights (scalar per horizon) -// b2 [5] per-horizon output bias -// h [n_batch, HIDDEN] trunk output (LayerNorm-normalised) -// probs [n_batch, 5] sigmoid output -// z1_out [n_batch, 5, HEAD_MID] saved pre-GELU for backward -// a1_out [n_batch, 5, HEAD_MID] saved post-GELU for backward +// w1 [N_HORIZONS, HEAD_MID, HIDDEN] per-horizon first-layer weights +// b1 [N_HORIZONS, HEAD_MID] per-horizon first-layer biases +// w2 [N_HORIZONS, HEAD_MID] per-horizon output weights (scalar per horizon) +// b2 [N_HORIZONS] per-horizon output bias +// h [n_batch, HIDDEN] trunk output (LayerNorm-normalised) +// probs [n_batch, N_HORIZONS] sigmoid output +// z1_out [n_batch, N_HORIZONS, HEAD_MID] saved pre-GELU for backward +// a1_out [n_batch, N_HORIZONS, HEAD_MID] saved post-GELU for backward // // GELU(x) ≈ 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))) // // Launch: grid=(n_batch, 1, 1), block=(HEAD_MID, 1, 1). -// Each block handles ONE sample and loops over the 5 horizons. - -#define N_HORIZONS_H 5 -#define HIDDEN_H 128 -#define HEAD_MID_H 64 -#define GELU_C0 0.7978845608f // sqrt(2/pi) -#define GELU_C1 0.044715f +// Each block handles ONE sample and loops over the N_HORIZONS horizons. __device__ __forceinline__ float gelu_act(float x) { const float u = GELU_C0 * (x + GELU_C1 * x * x * x); @@ -218,15 +221,15 @@ __device__ __forceinline__ float gelu_act(float x) { } extern "C" __global__ void multi_horizon_heads_2layer_fwd_batched( - const float* __restrict__ w1, // [5, HEAD_MID, HIDDEN] - const float* __restrict__ b1, // [5, HEAD_MID] - const float* __restrict__ w2, // [5, HEAD_MID] - const float* __restrict__ b2, // [5] + const float* __restrict__ w1, // [N_HORIZONS, HEAD_MID, HIDDEN] + const float* __restrict__ b1, // [N_HORIZONS, HEAD_MID] + const float* __restrict__ w2, // [N_HORIZONS, HEAD_MID] + const float* __restrict__ b2, // [N_HORIZONS] const float* __restrict__ h, // [n_batch, HIDDEN] int n_batch, - float* __restrict__ probs, // [n_batch, 5] - float* __restrict__ z1_out, // [n_batch, 5, HEAD_MID] - float* __restrict__ a1_out // [n_batch, 5, HEAD_MID] + float* __restrict__ probs, // [n_batch, N_HORIZONS] + float* __restrict__ z1_out, // [n_batch, N_HORIZONS, HEAD_MID] + float* __restrict__ a1_out // [n_batch, N_HORIZONS, HEAD_MID] ) { int b_idx = blockIdx.x; int m = threadIdx.x; @@ -234,11 +237,11 @@ extern "C" __global__ void multi_horizon_heads_2layer_fwd_batched( const float* h_row = h + (long long)b_idx * HIDDEN_H; - // Shared per-horizon post-GELU activations [5 * 64], indexed as [k*64 + m]. + // Shared per-horizon post-GELU activations [N_HORIZONS_H * 64], indexed as [k*64 + m]. __shared__ float s_a1[N_HORIZONS_H * HEAD_MID_H]; // Pass 1: per-horizon hidden unit. Thread `m` computes the m-th - // hidden unit for each of the 5 horizons sequentially. + // hidden unit for each of the N_HORIZONS horizons sequentially. #pragma unroll for (int k = 0; k < N_HORIZONS_H; ++k) { float z = b1[k * HEAD_MID_H + m]; @@ -252,8 +255,8 @@ extern "C" __global__ void multi_horizon_heads_2layer_fwd_batched( } __syncthreads(); - // Pass 2: per-horizon output logit. Threads m=0..4 each handle one - // horizon by computing the dot product over a1[k, :] and W2[k, :]. + // Pass 2: per-horizon output logit. Threads m=0..(N_HORIZONS-1) each + // handle one horizon by computing the dot product over a1[k, :] and W2[k, :]. if (m < N_HORIZONS_H) { const int k = m; float z2 = b2[k]; @@ -266,7 +269,7 @@ extern "C" __global__ void multi_horizon_heads_2layer_fwd_batched( } // 2-layer heads backward — chain rule through sigmoid → linear → -// GELU → linear → trunk. ISV `lambda[5]` scales the trunk-grad +// GELU → linear → trunk. ISV `lambda[N_HORIZONS]` scales the trunk-grad // contribution per horizon (per `pearl_adam_normalizes_loss_weights.md` // the effective lever is the trunk gradient, not loss weight). // @@ -285,7 +288,7 @@ extern "C" __global__ void multi_horizon_heads_2layer_fwd_batched( // the per-block sample loop and the per-k loop are serial within one // thread; no cross-block races since each block handles ONE sample. // Same for grad_w2[k, m] / grad_b2[k] / grad_b1[k, m]: each is written -// by exactly one thread per block (m for grad_w2 / grad_b1 dims, m<5 +// by exactly one thread per block (m for grad_w2 / grad_b1 dims, m= HEAD_MID_H) return; - // Broadcast lambda[5] via shared mem (5 floats). + // Broadcast lambda[N_HORIZONS] via shared mem (N_HORIZONS floats). __shared__ float s_lambda[N_HORIZONS_H]; if (tid < N_HORIZONS_H) { const float l = lambda[tid]; @@ -348,7 +351,7 @@ extern "C" __global__ void multi_horizon_heads_2layer_bwd_batched( __syncthreads(); for (int bi = 0; bi < n_batch; ++bi) { - // Pass 1: compute d_z2[k] for k=0..4 (only first 5 threads). + // Pass 1: compute d_z2[k] for k=0..(N_HORIZONS-1) (only first N_HORIZONS threads). if (tid < N_HORIZONS_H) { const int k = tid; const float p = probs[(long long)bi * N_HORIZONS_H + k]; @@ -376,7 +379,7 @@ extern "C" __global__ void multi_horizon_heads_2layer_bwd_batched( // grad_b1[k, m] += d_z1[k, m] grad_b1[k * HEAD_MID_H + m] += d_z1_km; } - // grad_b2[k] (only first 5 threads). + // grad_b2[k] (only first N_HORIZONS threads). if (tid < N_HORIZONS_H) { grad_b2[tid] += s_d_z2[tid]; } @@ -447,28 +450,28 @@ extern "C" __global__ void multi_horizon_heads_2layer_bwd_batched( // // Thread layout (block_dim = HEAD_MID = 64): // m = threadIdx.x -- owns "output_row" dim of (k, m_out=m) for w1/w2/eta_1/eta_2 -// for skip / gate / main scalars: thread tid<5 handles horizon k=tid +// for skip / gate / main scalars: thread tid> 5; const int total = K * B * OS_N_HORIZONS; - const int stride_k = B * OS_N_HORIZONS; // stride between k and k+1 at fixed (b,h) + const int stride_k = B * OS_N_HORIZONS; // stride between k and k+1 at fixed (b,h); B*N_HORIZONS const int n_pairs = (K - 1) * B; // N_pairs == 0 when K==1; in that case we still need to zero outputs. const float inv_n_pairs = (n_pairs > 0) ? (1.0f / (float)n_pairs) : 0.0f; @@ -124,7 +127,7 @@ extern "C" __global__ void output_smoothness_loss_and_grad( } __syncthreads(); - // Thread 0 sums the 5 λ-weighted terms into loss_out. + // Thread 0 sums the N_HORIZONS λ-weighted terms into loss_out. if (tid == 0) { float total_loss = 0.0f; #pragma unroll diff --git a/crates/ml-alpha/src/cfc/step.rs b/crates/ml-alpha/src/cfc/step.rs index b4be8ee66..53aeb8f38 100644 --- a/crates/ml-alpha/src/cfc/step.rs +++ b/crates/ml-alpha/src/cfc/step.rs @@ -92,9 +92,10 @@ pub fn cfc_step_backward_gpu( /// Fused per-branch CfC forward step (Phase 2 dispatch). /// /// Per spec §5.4 point 1 (docs/superpowers/specs/2026-05-21-per-horizon-cfc-inference-design.md): -/// single launch covers all 5 branches × n_batch with grid=(B, N_HORIZONS, 1) -/// and block=(MAX_BUCKET_DIM=28, 1, 1). Uniform-predicate handles uneven -/// bucket sizes [25, 25, 25, 25, 28] without warp divergence. +/// single launch covers all N_HORIZONS branches × n_batch with grid=(B, N_HORIZONS, 1) +/// and block=(MAX_BUCKET_DIM=96, 1, 1). Uniform-predicate handles uneven +/// bucket sizes [43, 43, 42] without warp divergence (post-2026-05-22 +/// horizon-rebase to N_HORIZONS=3). /// /// ALPHA fix (2026-05-21): the kernel now reads channels via /// `channels_in_bucket_d` (a [N_HORIZONS × MAX_BUCKET_DIM] lookup table @@ -132,8 +133,8 @@ pub fn cfc_step_per_branch_fwd_gpu( bucket_dim_k_d: &CudaSlice, h_new_d: &mut CudaSlice, ) -> Result<()> { - const N_HORIZONS: u32 = 5; - const MAX_BUCKET_DIM: u32 = 28; + const N_HORIZONS: u32 = crate::heads::N_HORIZONS as u32; + const MAX_BUCKET_DIM: u32 = crate::cfc::bucket_routing::MAX_BUCKET_DIM as u32; const HIDDEN_DIM: u32 = 128; let cfg = LaunchConfig { grid_dim: (b_sz as u32, N_HORIZONS, 1), @@ -189,8 +190,8 @@ pub fn cfc_step_per_branch_bwd_gpu( grad_tau_all_d: &mut CudaSlice, grad_h_old_d: &mut CudaSlice, ) -> Result<()> { - const N_HORIZONS: u32 = 5; - const MAX_BUCKET_DIM: u32 = 28; + const N_HORIZONS: u32 = crate::heads::N_HORIZONS as u32; + const MAX_BUCKET_DIM: u32 = crate::cfc::bucket_routing::MAX_BUCKET_DIM as u32; const HIDDEN_DIM: u32 = 128; let cfg = LaunchConfig { grid_dim: (b_sz as u32, N_HORIZONS, 1),