From 1a465cf7d52ab8dd762f110951ab4f502720ab89 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 18 May 2026 18:11:05 +0200 Subject: [PATCH] =?UTF-8?q?refactor(ml-alpha):=20revert=20axes=20B/C/D/E?= =?UTF-8?q?=20=E2=80=94=20keep=20only=20Kendall=20=CF=83=20(axis=20A)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Post-A/B verdict (see project_ml_alpha_v2_ab_verdict.md): v2 with all 5 axes was marginally tied on h6000 (+0.0013 vs 0.7591 baseline mean, fails the +0.01 win threshold) and slightly below on mean_auc (−0.0208 vs 0.7749 baseline mean, within 1σ) at ~5× the wall-time cost. Per `feedback_v7_gem_methodology` (measure before delete or wire), the architecture has been measured — it doesn't earn its compute cost. This commit reverts the axes that didn't lift: - axis B (L2 anchor + Wiener-α controller) — DROPPED - axis C (horizon-token attention pool) — DROPPED - axis D (regime-MoE gate + experts) — DROPPED - axis E (inverted cross-variate attn) — DROPPED - axis A (Kendall σ-weighted BCE) — KEPT Files deleted (kernels, host bindings, numgrad tests, trainer state): - cuda/{horizon_token_attention_pool, inverted_attention_pool, inv_pooled_merge, regime_moe_gate, anchor_l2, horizon_mean_collapse}.cu - src/{horizon_token_attention_pool, inverted_attention_pool, inv_pooled_merge, regime_moe_gate, anchor_l2, horizon_mean_collapse}.rs - src/trainer/{multi_horizon_attention, anchor_controller}.rs - tests/{horizon_token_attention_pool_numgrad, inverted_attention_pool_numgrad, regime_moe_gate_numgrad, anchor_l2_numgrad}.rs Files restored (from V1 commit 41292303d): - cuda/attention_pool.cu — legacy single-Q attention pool kernel - src/trainer/perception.rs — pre-MHA trainer state with the legacy `attn_*` plumbing intact. Files modified: - bce_loss_multi_horizon.cu stays σ-aware (kept the V7 work; it has the kernel function name preserved from V1). - perception.rs: ADD `log_sigma_h_d [N_HORIZONS]`, `grad_log_sigma_h_d [N_HORIZONS]`, `opt_log_sigma` AdamW directly on PerceptionTrainer (no MHA bundle). BCE callsites in `step_batched` (training) and `evaluate_batched` thread the σ args. Grad scratch zeroed each step before the BCE launch. `opt_log_sigma.step` lives in section 9 alongside the other AdamW updates. NET DIFF: 23 files, 442 insertions, 3160 deletions (~2700-line cleanup). LOCAL VERIFICATION (RTX 3050 sm_86, --test-threads=1): - ml-alpha builds clean (cuda feature) - bce_grad_finite_diff 4/4 PASS (BCE still works through σ-kernel) - perception_overfit 9/9 PASS (full trainer pipeline, loss-shrinks tests still green) NEXT: cluster smoke + 3-fold A/B vs task #200 baseline. Expected wall-time ≈ baseline 17 s/epoch (we're back to baseline architecture plus 5 scalar Kendall σ params + 1 tiny AdamW). Expected lift on mean_auc: modest — Kendall σ rebalances per-horizon contributions based on observed BCE EMA, which may help horizons with intrinsically higher noise floors. Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/build.rs | 10 +- crates/ml-alpha/cuda/anchor_l2.cu | 67 --- crates/ml-alpha/cuda/attention_pool.cu | 244 ++++++++++ crates/ml-alpha/cuda/horizon_mean_collapse.cu | 51 -- .../cuda/horizon_token_attention_pool.cu | 322 ------------- crates/ml-alpha/cuda/inv_pooled_merge.cu | 64 --- .../ml-alpha/cuda/inverted_attention_pool.cu | 265 ----------- crates/ml-alpha/cuda/regime_moe_gate.cu | 300 ------------ crates/ml-alpha/src/anchor_l2.rs | 59 --- crates/ml-alpha/src/horizon_mean_collapse.rs | 81 ---- .../src/horizon_token_attention_pool.rs | 122 ----- crates/ml-alpha/src/inv_pooled_merge.rs | 82 ---- .../ml-alpha/src/inverted_attention_pool.rs | 112 ----- crates/ml-alpha/src/lib.rs | 6 - crates/ml-alpha/src/regime_moe_gate.rs | 174 ------- .../ml-alpha/src/trainer/anchor_controller.rs | 142 ------ crates/ml-alpha/src/trainer/mod.rs | 2 - .../src/trainer/multi_horizon_attention.rs | 435 ------------------ crates/ml-alpha/src/trainer/perception.rs | 297 ++++++++---- crates/ml-alpha/tests/anchor_l2_numgrad.rs | 138 ------ .../horizon_token_attention_pool_numgrad.rs | 212 --------- .../tests/inverted_attention_pool_numgrad.rs | 127 ----- .../ml-alpha/tests/regime_moe_gate_numgrad.rs | 290 ------------ 23 files changed, 442 insertions(+), 3160 deletions(-) delete mode 100644 crates/ml-alpha/cuda/anchor_l2.cu create mode 100644 crates/ml-alpha/cuda/attention_pool.cu delete mode 100644 crates/ml-alpha/cuda/horizon_mean_collapse.cu delete mode 100644 crates/ml-alpha/cuda/horizon_token_attention_pool.cu delete mode 100644 crates/ml-alpha/cuda/inv_pooled_merge.cu delete mode 100644 crates/ml-alpha/cuda/inverted_attention_pool.cu delete mode 100644 crates/ml-alpha/cuda/regime_moe_gate.cu delete mode 100644 crates/ml-alpha/src/anchor_l2.rs delete mode 100644 crates/ml-alpha/src/horizon_mean_collapse.rs delete mode 100644 crates/ml-alpha/src/horizon_token_attention_pool.rs delete mode 100644 crates/ml-alpha/src/inv_pooled_merge.rs delete mode 100644 crates/ml-alpha/src/inverted_attention_pool.rs delete mode 100644 crates/ml-alpha/src/regime_moe_gate.rs delete mode 100644 crates/ml-alpha/src/trainer/anchor_controller.rs delete mode 100644 crates/ml-alpha/src/trainer/multi_horizon_attention.rs delete mode 100644 crates/ml-alpha/tests/anchor_l2_numgrad.rs delete mode 100644 crates/ml-alpha/tests/horizon_token_attention_pool_numgrad.rs delete mode 100644 crates/ml-alpha/tests/inverted_attention_pool_numgrad.rs delete mode 100644 crates/ml-alpha/tests/regime_moe_gate_numgrad.rs diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 1ba2aba13..09d150ae5 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -13,19 +13,13 @@ const KERNELS: &[&str] = &[ "cfc_step", "multi_horizon_heads", "projection", - "bce_loss_multi_horizon", - "bce_loss_multi_horizon_sigma", // v2-A: Kendall σ-weighted BCE + "bce_loss_multi_horizon", // Kendall σ-weighted multi-horizon BCE (axis A) "adamw_step", "grad_norm", "horizon_lambda", // ISV-driven per-horizon gradient scaler (EMA + lambda) "layer_norm", // Phase 1: trunk pre-CfC normalisation "variable_selection", // Phase 2D: TFT-style per-feature gating - "horizon_token_attention_pool", // Horizon-token K-prepend single-Q attention pool - "inverted_attention_pool", // Cross-variate (iTransformer-style) attention pool - "regime_moe_gate", // Top-1 regime MoE gate + expert dispatch + aux loss - "horizon_mean_collapse", // Mean over horizon axis (seeds CfC h_old at k=0) - "inv_pooled_merge", // Broadcast-add inv_pooled into ctx_h (wires axis E into loss) - "anchor_l2", // L2 anchor regularization toward init + "attention_pool", // Phase 3: single-Q learned content summary at CfC k=0 "reduce_axis0", // Phase B: cross-batch param-grad reducer ]; diff --git a/crates/ml-alpha/cuda/anchor_l2.cu b/crates/ml-alpha/cuda/anchor_l2.cu deleted file mode 100644 index b6fd680a5..000000000 --- a/crates/ml-alpha/cuda/anchor_l2.cu +++ /dev/null @@ -1,67 +0,0 @@ -// anchor_l2.cu — L2 anchor regularization toward initialization (v2 axis B). -// -// Computes anchor_loss = λ · Σ_i (p[i] − p_init[i])² for a single -// flat parameter buffer p of length n. Adds the corresponding L2 -// gradient 2λ(p − p_init) into the parameter's grad buffer (`+=`). -// -// Caller invokes one launch per parameter group anchored. The -// trainer's Wiener-α controller computes λ each step on the host -// (cheap — scalar update; copies λ to a device-side single-float -// buffer before the launch). -// -// PERFORMANCE: -// - Warp-shuffle reduce for the loss sum. One block per parameter -// group; grid_dim = (1, 1, 1). For n ≤ ~100K (typical group size) -// a single 256-thread block is sufficient; threads stride over n. -// - Grad write is per-element, coalesced. - -#define ANCHOR_BLOCK 256 -#define ANCHOR_NWARPS (ANCHOR_BLOCK / 32) - -__device__ __forceinline__ float warp_reduce_sum_an(float v) { - #pragma unroll - for (int s = 16; s > 0; s >>= 1) { - v += __shfl_xor_sync(0xffffffff, v, s); - } - return v; -} - -extern "C" __global__ void anchor_l2_fwd_bwd( - const float* __restrict__ p, // [n] current parameter values - const float* __restrict__ p_init, // [n] snapshot at trainer construction - const float* __restrict__ lambda_scalar, // [1] current λ from Wiener-α controller - int n, - float* __restrict__ loss_out, // [1] λ · Σ (p − p_init)² - float* __restrict__ grad_p // [n] += 2λ (p − p_init) -) { - const int tid = threadIdx.x; - const int lane = tid & 31; - const int warp = tid >> 5; - - const float lam = lambda_scalar[0]; - const float two_lam = 2.0f * lam; - - // Strided per-thread accumulate of (p - p_init)^2 + emit grad. - float local_sum = 0.0f; - for (int i = tid; i < n; i += ANCHOR_BLOCK) { - const float diff = p[i] - p_init[i]; - local_sum += diff * diff; - // Coalesced grad write. - grad_p[i] += two_lam * diff; - } - - // Warp-shuffle reduce. - float warp_sum = warp_reduce_sum_an(local_sum); - - __shared__ float s_warp[ANCHOR_NWARPS]; - if (lane == 0) s_warp[warp] = warp_sum; - __syncthreads(); - - // Cross-warp reduce: all 32 lanes of warp 0 participate; inactive - // lanes contribute 0 via ternary (NOT a divergent shuffle). - float w = (tid < ANCHOR_NWARPS) ? s_warp[tid] : 0.0f; - if (tid < 32) { - w = warp_reduce_sum_an(w); - if (tid == 0) loss_out[0] = lam * w; - } -} diff --git a/crates/ml-alpha/cuda/attention_pool.cu b/crates/ml-alpha/cuda/attention_pool.cu new file mode 100644 index 000000000..65701480e --- /dev/null +++ b/crates/ml-alpha/cuda/attention_pool.cu @@ -0,0 +1,244 @@ +// attention_pool.cu — Single-head attention pool over Mamba2 K-positions. +// Phase 3 (2026-05-17). +// +// Replaces the CfC's zero-initialised `h_old` at k=0 with a learned +// attention-pooled summary over all K LN_b output positions. The K-loop +// CfC keeps its recurrent semantics (h_old at k+1 = h_new at k); only +// the initial state changes from zero to a content-addressable lookup. +// +// Forward math (per sample b): +// scores[k] = Q · keys[b, k, :] # [K] — dot product over HIDDEN_DIM +// attn[k] = softmax_k(scores) # [K] +// context[h] = sum_k attn[k] * values[b, k, h] # [HIDDEN_DIM] +// +// For this attention pool, keys == values == LN_b output [B, K, HIDDEN_DIM]. +// Single learned param: Q [HIDDEN_DIM]. +// +// Saved for backward: +// attn_weights [B, K] (post-softmax) +// +// Block layout (forward): +// grid_dim = (B, 1, 1), block_dim = (HIDDEN_DIM=128, 1, 1). +// Thread tid handles dimension h of HIDDEN_DIM. +// +// Per `feedback_no_atomicadd.md`: block tree-reduce only, no atomics. + +#define ATTN_HIDDEN_DIM 128 +#define ATTN_BLOCK 128 +#define ATTN_MAX_K 512 // safety cap on K; smoke uses K=16-64. + +extern "C" __global__ void attention_pool_fwd( + const float* __restrict__ Q, // [HIDDEN_DIM] + const float* __restrict__ ln_out, // [B, K, HIDDEN_DIM] + int n_batch, + int k_seq, + float* __restrict__ context, // [B, HIDDEN_DIM] + float* __restrict__ attn_weights // [B, K] — saved post-softmax +) { + int b_idx = blockIdx.x; + int tid = threadIdx.x; + if (b_idx >= n_batch || tid >= ATTN_BLOCK) return; + + extern __shared__ float smem[]; + float* s_scores = smem; // [K] + float* s_red = smem + k_seq; // [BLOCK] — reduce scratch + float* s_context = smem + k_seq + ATTN_BLOCK; // [HIDDEN_DIM] — final output buffer + + // Cache Q in shared mem (broadcast). + __shared__ float s_Q[ATTN_HIDDEN_DIM]; + if (tid < ATTN_HIDDEN_DIM) s_Q[tid] = Q[tid]; + __syncthreads(); + + // Pass 1: scores[k] = Q · ln_out[b, k, :]. + // Each k is handled sequentially by ALL threads via tree-reduce + // over HIDDEN_DIM. K outer loop, threads tile HIDDEN_DIM inner. + // For HIDDEN_DIM=128 and ATTN_BLOCK=128 we have one-to-one. + const float* ln_b = ln_out + (long long)b_idx * k_seq * ATTN_HIDDEN_DIM; + for (int k = 0; k < k_seq; ++k) { + const float v = ln_b[k * ATTN_HIDDEN_DIM + tid] * s_Q[tid]; + s_red[tid] = v; + __syncthreads(); + for (int s = ATTN_BLOCK / 2; s > 0; s >>= 1) { + if (tid < s) s_red[tid] += s_red[tid + s]; + __syncthreads(); + } + if (tid == 0) s_scores[k] = s_red[0]; + __syncthreads(); + } + + // Pass 2: softmax over K. Numerically stable: max-subtract + exp + sum. + // Block-wide max over s_scores[0..k_seq). + float my_max = -INFINITY; + for (int k = tid; k < k_seq; k += ATTN_BLOCK) { + const float v = s_scores[k]; + if (v > my_max) my_max = v; + } + s_red[tid] = my_max; + __syncthreads(); + for (int s = ATTN_BLOCK / 2; s > 0; s >>= 1) { + if (tid < s) { + const float a = s_red[tid]; + const float b = s_red[tid + s]; + s_red[tid] = (a > b) ? a : b; + } + __syncthreads(); + } + __shared__ float s_max; + if (tid == 0) s_max = s_red[0]; + __syncthreads(); + + // exp(score - max) + sum. + float my_sum = 0.0f; + for (int k = tid; k < k_seq; k += ATTN_BLOCK) { + const float e = expf(s_scores[k] - s_max); + s_scores[k] = e; // reuse as exp_shifted + my_sum += e; + } + s_red[tid] = my_sum; + __syncthreads(); + for (int s = ATTN_BLOCK / 2; s > 0; s >>= 1) { + if (tid < s) s_red[tid] += s_red[tid + s]; + __syncthreads(); + } + __shared__ float s_sum; + if (tid == 0) s_sum = s_red[0]; + __syncthreads(); + + // Pass 3: attn[k] = exp/sum; save; context[h] = sum_k attn[k] * ln_out[b, k, h]. + // Save attn weights AND build the context output. + for (int k = tid; k < k_seq; k += ATTN_BLOCK) { + const float a = s_scores[k] / s_sum; + s_scores[k] = a; // overwrite again — now holds true attn weights + attn_weights[(long long)b_idx * k_seq + k] = a; + } + __syncthreads(); + + // Accumulate context[h=tid] = sum_k attn[k] * ln_out[b, k, h=tid]. + if (tid < ATTN_HIDDEN_DIM) { + float c = 0.0f; + for (int k = 0; k < k_seq; ++k) { + c += s_scores[k] * ln_b[k * ATTN_HIDDEN_DIM + tid]; + } + s_context[tid] = c; + context[(long long)b_idx * ATTN_HIDDEN_DIM + tid] = c; + } +} + +// Attention pool backward — chain rule: +// +// d_attn[k] = sum_h grad_context[h] * values[b, k, h] +// + (this is from context = sum_k attn[k] * values[b, k, :]) +// +// d_values[b, k, h] += grad_context[h] * attn[k] +// + d_scores[k] * Q[h] +// +// d_scores via softmax Jacobian: +// d_scores[k] = attn[k] * (d_attn[k] - sum_kp attn[kp] * d_attn[kp]) +// +// d_Q[h] += sum_{b, k} d_scores[k] * values[b, k, h] +// (here values == ln_out) +// +// Single-writer discipline: +// - d_Q is [HIDDEN_DIM] shared across all (b, k). ONE block per +// launch, internal batch loop, tile over HIDDEN_DIM. += into d_Q. +// - d_values[b, k, h] is [B, K, HIDDEN_DIM] — one block per launch +// iterates over (b, k, h) sequentially; with block_dim = HIDDEN_DIM, +// thread h is sole writer of column h for ALL (b, k). + +// Block-per-batch attn_pool bwd (Phase B commit 4). +// grid=(n_batch, 1, 1) block=(ATTN_BLOCK, 1, 1) +// +// Each block handles one batch's HIDDEN_DIM channels. grad_ln_out is +// per-batch indexed (already safe), so each block bi writes its +// [bi, :, :] slice via += onto whatever value grad_ln_out holds at +// kernel launch (the K-loop's contribution to LN_b output grad). +// +// grad_Q is a single [HIDDEN_DIM] shared across all batches → per-batch +// scratch [B, HIDDEN_DIM], reduced after the kernel returns. +extern "C" __global__ void attention_pool_bwd( + const float* __restrict__ Q, // [HIDDEN_DIM] + const float* __restrict__ ln_out, // [B, K, HIDDEN_DIM] (= values) + const float* __restrict__ attn_weights, // [B, K] from fwd + const float* __restrict__ grad_context, // [B, HIDDEN_DIM] + int n_batch, + int k_seq, + float* __restrict__ grad_Q_scratch, // [B, HIDDEN_DIM] (+=) + float* __restrict__ grad_ln_out // [B, K, HIDDEN_DIM] (+= chained with K-loop) +) { + int bi = blockIdx.x; + int tid = threadIdx.x; + if (bi >= n_batch || tid >= ATTN_BLOCK) return; + + extern __shared__ float smem[]; + float* s_dattn = smem; // [K] + float* s_dscores = smem + k_seq; // [K] + float* s_red = smem + 2 * k_seq; // [BLOCK] + + __shared__ float s_Q[ATTN_HIDDEN_DIM]; + __shared__ float s_attn[ATTN_MAX_K]; + __shared__ float s_grad_ctx[ATTN_HIDDEN_DIM]; + __shared__ float s_sum_attn_dattn; + + if (tid < ATTN_HIDDEN_DIM) s_Q[tid] = Q[tid]; + __syncthreads(); + + const float* ln_b = ln_out + (long long)bi * k_seq * ATTN_HIDDEN_DIM; + float* grad_ln_b = grad_ln_out + (long long)bi * k_seq * ATTN_HIDDEN_DIM; + // Cache attn + grad_context for this block's batch. + for (int k = tid; k < k_seq; k += ATTN_BLOCK) { + s_attn[k] = attn_weights[(long long)bi * k_seq + k]; + } + if (tid < ATTN_HIDDEN_DIM) { + s_grad_ctx[tid] = grad_context[(long long)bi * ATTN_HIDDEN_DIM + tid]; + } + __syncthreads(); + + // Pass 1: d_attn[k] = sum_h grad_context[h] * values[b, k, h]. + for (int k = 0; k < k_seq; ++k) { + const float v = (tid < ATTN_HIDDEN_DIM) + ? s_grad_ctx[tid] * ln_b[k * ATTN_HIDDEN_DIM + tid] : 0.0f; + s_red[tid] = v; + __syncthreads(); + for (int s = ATTN_BLOCK / 2; s > 0; s >>= 1) { + if (tid < s) s_red[tid] += s_red[tid + s]; + __syncthreads(); + } + if (tid == 0) s_dattn[k] = s_red[0]; + __syncthreads(); + } + + // Pass 2: sum_{kp} attn[kp] * d_attn[kp] (softmax Jacobian centring). + float my_sum = 0.0f; + for (int k = tid; k < k_seq; k += ATTN_BLOCK) { + my_sum += s_attn[k] * s_dattn[k]; + } + s_red[tid] = my_sum; + __syncthreads(); + for (int s = ATTN_BLOCK / 2; s > 0; s >>= 1) { + if (tid < s) s_red[tid] += s_red[tid + s]; + __syncthreads(); + } + if (tid == 0) s_sum_attn_dattn = s_red[0]; + __syncthreads(); + + // Pass 3: d_scores[k] = attn[k] * (d_attn[k] - s_sum_attn_dattn). + for (int k = tid; k < k_seq; k += ATTN_BLOCK) { + s_dscores[k] = s_attn[k] * (s_dattn[k] - s_sum_attn_dattn); + } + __syncthreads(); + + // Pass 4: grad_Q_scratch[bi, h] += sum_k d_scores[k] * ln_out[b, k, h]. + // d_ln_out[b, k, h] += grad_context[h] * attn[k] + d_scores[k] * Q[h]. + // Thread h owns column h. Loops over k. grad_ln_out += chains the + // attn-path gradient on top of the K-loop's contribution. + if (tid < ATTN_HIDDEN_DIM) { + float dq_local = 0.0f; + for (int k = 0; k < k_seq; ++k) { + const float v = ln_b[k * ATTN_HIDDEN_DIM + tid]; + dq_local += s_dscores[k] * v; + grad_ln_b[k * ATTN_HIDDEN_DIM + tid] += + s_grad_ctx[tid] * s_attn[k] + s_dscores[k] * s_Q[tid]; + } + grad_Q_scratch[(long long)bi * ATTN_HIDDEN_DIM + tid] += dq_local; + } +} diff --git a/crates/ml-alpha/cuda/horizon_mean_collapse.cu b/crates/ml-alpha/cuda/horizon_mean_collapse.cu deleted file mode 100644 index 1f3de7609..000000000 --- a/crates/ml-alpha/cuda/horizon_mean_collapse.cu +++ /dev/null @@ -1,51 +0,0 @@ -// horizon_mean_collapse.cu — Collapse per-horizon context to a single -// vector by averaging across the horizon axis. -// -// Forward: -// out[b, d] = (1 / N_H) · Σ_h in[b, h, d] -// -// Backward (broadcast — each horizon receives the same per-d gradient): -// d_in[b, h, d] = (1 / N_H) · d_out[b, d] -// -// Grid: (B, 1, 1). Block: (HIDDEN_DIM, 1, 1). -// Each thread tid owns d = tid. No reductions, no barriers. - -#define HMC_HIDDEN_DIM 128 -#define HMC_BLOCK 128 -#define HMC_N_HORIZONS 5 - -extern "C" __global__ void horizon_mean_collapse_fwd( - const float* __restrict__ ctx_per_h, // [B, N_H, H] - int n_batch, - float* __restrict__ ctx_mean // [B, H] -) { - int b = blockIdx.x; - int tid = threadIdx.x; - if (b >= n_batch || tid >= HMC_HIDDEN_DIM) return; - - float acc = 0.0f; - #pragma unroll - for (int h = 0; h < HMC_N_HORIZONS; ++h) { - acc += ctx_per_h[(long long)b * HMC_N_HORIZONS * HMC_HIDDEN_DIM - + (long long)h * HMC_HIDDEN_DIM + tid]; - } - ctx_mean[(long long)b * HMC_HIDDEN_DIM + tid] = acc * (1.0f / (float)HMC_N_HORIZONS); -} - -extern "C" __global__ void horizon_mean_collapse_bwd( - const float* __restrict__ grad_ctx_mean, // [B, H] - int n_batch, - float* __restrict__ grad_ctx_per_h // [B, N_H, H] += broadcast / N_H -) { - int b = blockIdx.x; - int tid = threadIdx.x; - if (b >= n_batch || tid >= HMC_HIDDEN_DIM) return; - - const float g = grad_ctx_mean[(long long)b * HMC_HIDDEN_DIM + tid] - * (1.0f / (float)HMC_N_HORIZONS); - #pragma unroll - for (int h = 0; h < HMC_N_HORIZONS; ++h) { - grad_ctx_per_h[(long long)b * HMC_N_HORIZONS * HMC_HIDDEN_DIM - + (long long)h * HMC_HIDDEN_DIM + tid] += g; - } -} diff --git a/crates/ml-alpha/cuda/horizon_token_attention_pool.cu b/crates/ml-alpha/cuda/horizon_token_attention_pool.cu deleted file mode 100644 index 9d17133b5..000000000 --- a/crates/ml-alpha/cuda/horizon_token_attention_pool.cu +++ /dev/null @@ -1,322 +0,0 @@ -// horizon_token_attention_pool.cu — Horizon-token attention pool (v2 axis C). -// -// Replaces the falsified per-horizon Q_h (C21/C25) with a single -// shared query Q over an EXTENDED key sequence: -// K' = [horizon_tokens; LN_b_out] shape [B, N_H + K, H] -// -// One attention pass; per-horizon outputs derived by mixing each -// horizon's token contribution with a shared time-context aggregate. -// -// FORWARD MATH (per batch b; b indexing suppressed): -// -// For i ∈ [0, N_H): -// score[i] = Σ_d Q[d] · horizon_tokens[i, d] -// For i ∈ [N_H, N_H + K): -// k = i − N_H -// score[i] = Σ_d Q[d] · ln_out[k, d] -// -// attn[i] = softmax_i(score) # over N_H + K positions -// S[d] = Σ_k attn[N_H + k] · ln_out[k, d] # shared time-context per d -// ctx[h, d] = attn[h] · horizon_tokens[h, d] + S[d] # per-horizon output -// -// BACKWARD MATH (chain rule): -// -// d_S[d] = Σ_h grad_ctx[h, d] -// d_attn[h] = Σ_d grad_ctx[h, d] · horizon_tokens[h, d] (h-token term) -// d_attn[N_H + k] = Σ_d d_S[d] · ln_out[k, d] (time term) -// d_score (softmax bwd): d_score[i] = attn[i] · (d_attn[i] − Σ_j attn[j]·d_attn[j]) -// -// d_horizon_tokens[h, d] = d_score[h] · Q[d] -// + grad_ctx[h, d] · attn[h] -// d_Q[d] = Σ_h d_score[h] · horizon_tokens[h, d] -// + Σ_k d_score[N_H + k] · ln_out[k, d] -// d_ln_out[k, d] = d_score[N_H + k] · Q[d] -// + attn[N_H + k] · d_S[d] -// -// PERFORMANCE (per feedback_nvidia_grade_perf_for_kernels): -// - Warp-shuffle reduction (`__shfl_xor_sync`) for every per-d dot -// product and per-attention sum; NO block tree-reduce. -// - 4 warps × 32 lanes; cross-warp reduce uses one `__syncthreads`. -// - All inactive-lane shuffles use ternary fallback to 0 (never a -// divergent `if (tid < N) shuffle`). -// - Block-per-batch; horizon loop inside the block keeps grad writes -// race-free (each batch's ln_out slice touched by one block in bwd). -// - Shared mem layout pre-computed at launch from k_seq. - -#define HTAP_HIDDEN_DIM 128 -#define HTAP_BLOCK 128 -#define HTAP_N_HORIZONS 5 -#define HTAP_N_WARPS (HTAP_BLOCK / 32) // == 4 -#define HTAP_MAX_KSEQ 128 // safety cap; trainer uses K ≤ 64 -#define HTAP_MAX_EXT (HTAP_N_HORIZONS + HTAP_MAX_KSEQ) - -__device__ __forceinline__ float warp_reduce_sum_local(float v) { - #pragma unroll - for (int s = 16; s > 0; s >>= 1) { - v += __shfl_xor_sync(0xffffffff, v, s); - } - return v; -} - -__device__ __forceinline__ float warp_reduce_max_local(float v) { - #pragma unroll - for (int s = 16; s > 0; s >>= 1) { - v = fmaxf(v, __shfl_xor_sync(0xffffffff, v, s)); - } - return v; -} - -// Block-wide reduce-sum using warp-shuffle + 1 syncthreads. -// `s_warp` is caller-provided shared scratch of size HTAP_N_WARPS. -__device__ __forceinline__ float block_reduce_sum(float v, float* s_warp, int tid) { - v = warp_reduce_sum_local(v); - int lane = tid & 31, warp = tid >> 5; - if (lane == 0) s_warp[warp] = v; - __syncthreads(); - float w = (tid < HTAP_N_WARPS) ? s_warp[tid] : 0.0f; - if (tid < 32) { - w = warp_reduce_sum_local(w); - if (tid == 0) s_warp[0] = w; - } - __syncthreads(); - return s_warp[0]; -} - -__device__ __forceinline__ float block_reduce_max(float v, float* s_warp, int tid) { - v = warp_reduce_max_local(v); - int lane = tid & 31, warp = tid >> 5; - if (lane == 0) s_warp[warp] = v; - __syncthreads(); - float w = (tid < HTAP_N_WARPS) ? s_warp[tid] : -INFINITY; - if (tid < 32) { - w = warp_reduce_max_local(w); - if (tid == 0) s_warp[0] = w; - } - __syncthreads(); - return s_warp[0]; -} - -extern "C" __global__ void horizon_token_attention_pool_fwd( - const float* __restrict__ horizon_tokens, // [N_H, H] - const float* __restrict__ Q, // [H] - const float* __restrict__ ln_out, // [B, K, H] - int n_batch, - int k_seq, - float* __restrict__ ctx_out, // [B, N_H, H] - float* __restrict__ attn_out // [B, N_H + K] -) { - int b = blockIdx.x; - int tid = threadIdx.x; - if (b >= n_batch || tid >= HTAP_BLOCK) return; - - extern __shared__ float smem[]; - // Layout: s_attn[N_H + K] + s_warp[N_WARPS] + s_max + s_sum - float* s_attn = smem; // [N_H + k_seq] - float* s_warp = s_attn + (HTAP_N_HORIZONS + k_seq); // [N_WARPS] - // Used scalars stashed by the reduce helpers; cheap to materialise. - - // Cache Q[tid] and horizon_tokens[h, tid] in registers / shared. - __shared__ float s_horizon_tokens[HTAP_N_HORIZONS * HTAP_HIDDEN_DIM]; - if (tid < HTAP_HIDDEN_DIM) { - #pragma unroll - for (int h = 0; h < HTAP_N_HORIZONS; ++h) { - s_horizon_tokens[h * HTAP_HIDDEN_DIM + tid] = - horizon_tokens[h * HTAP_HIDDEN_DIM + tid]; - } - } - __syncthreads(); - - const float q_d = (tid < HTAP_HIDDEN_DIM) ? Q[tid] : 0.0f; - const float* ln_b = ln_out + (long long)b * k_seq * HTAP_HIDDEN_DIM; - - // ── Phase 1: scores[i] = Σ_d q[d] · ext[i, d]. Serial outer i, inner - // HIDDEN_DIM reduced via block_reduce_sum (warp-shuffle). - const int n_ext = HTAP_N_HORIZONS + k_seq; - for (int i = 0; i < n_ext; ++i) { - float key_d; - if (i < HTAP_N_HORIZONS) { - key_d = (tid < HTAP_HIDDEN_DIM) - ? s_horizon_tokens[i * HTAP_HIDDEN_DIM + tid] : 0.0f; - } else { - int k = i - HTAP_N_HORIZONS; - key_d = (tid < HTAP_HIDDEN_DIM) - ? ln_b[k * HTAP_HIDDEN_DIM + tid] : 0.0f; - } - float v = q_d * key_d; - float r = block_reduce_sum(v, s_warp, tid); - if (tid == 0) s_attn[i] = r; - // s_warp is reused by the next iteration's block_reduce_sum; safe. - } - - // ── Phase 2: softmax over [0, n_ext). - // 2a: max - float my_max = -INFINITY; - for (int i = tid; i < n_ext; i += HTAP_BLOCK) { - my_max = fmaxf(my_max, s_attn[i]); - } - const float s_max = block_reduce_max(my_max, s_warp, tid); - - // 2b: exp + sum - float my_sum = 0.0f; - for (int i = tid; i < n_ext; i += HTAP_BLOCK) { - const float e = expf(s_attn[i] - s_max); - s_attn[i] = e; - my_sum += e; - } - const float s_sum = block_reduce_sum(my_sum, s_warp, tid); - - // 2c: normalise + persist attn_out - const float inv_sum = 1.0f / s_sum; - for (int i = tid; i < n_ext; i += HTAP_BLOCK) { - const float a = s_attn[i] * inv_sum; - s_attn[i] = a; - attn_out[(long long)b * n_ext + i] = a; - } - __syncthreads(); - - // ── Phase 3: ctx[h, d] = attn[h] · h_token[h, d] + S[d] - // where S[d] = Σ_k attn[N_H + k] · ln_out[k, d]. - if (tid < HTAP_HIDDEN_DIM) { - float S_d = 0.0f; - for (int k = 0; k < k_seq; ++k) { - S_d += s_attn[HTAP_N_HORIZONS + k] * ln_b[k * HTAP_HIDDEN_DIM + tid]; - } - #pragma unroll - for (int h = 0; h < HTAP_N_HORIZONS; ++h) { - const float token_term = s_attn[h] * s_horizon_tokens[h * HTAP_HIDDEN_DIM + tid]; - ctx_out[(long long)b * HTAP_N_HORIZONS * HTAP_HIDDEN_DIM - + (long long)h * HTAP_HIDDEN_DIM + tid] = token_term + S_d; - } - } -} - -extern "C" __global__ void horizon_token_attention_pool_bwd( - const float* __restrict__ horizon_tokens, // [N_H, H] - const float* __restrict__ Q, // [H] - const float* __restrict__ ln_out, // [B, K, H] - const float* __restrict__ attn, // [B, N_H + K] (saved fwd) - const float* __restrict__ grad_ctx, // [B, N_H, H] (upstream) - int n_batch, - int k_seq, - float* __restrict__ grad_horizon_tokens_scratch, // [B, N_H, H] += (reduce_axis0 → [N_H, H]) - float* __restrict__ grad_Q_scratch, // [B, H] += (reduce_axis0 → [H]) - float* __restrict__ grad_ln_out // [B, K, H] += (per-batch, no race) -) { - int b = blockIdx.x; - int tid = threadIdx.x; - if (b >= n_batch || tid >= HTAP_BLOCK) return; - - extern __shared__ float smem_bwd[]; - // Layout: s_d_attn[N_H + K] + s_warp[N_WARPS] + s_d_S[HIDDEN_DIM] - float* s_d_attn = smem_bwd; // [N_H + k] - float* s_warp = s_d_attn + (HTAP_N_HORIZONS + k_seq); // [N_WARPS] - float* s_d_S = s_warp + HTAP_N_WARPS; // [HIDDEN_DIM] - - __shared__ float s_horizon_tokens[HTAP_N_HORIZONS * HTAP_HIDDEN_DIM]; - __shared__ float s_attn[HTAP_N_HORIZONS + HTAP_MAX_KSEQ]; - - const int n_ext = HTAP_N_HORIZONS + k_seq; - - // Load horizon_tokens + attn into shared. - if (tid < HTAP_HIDDEN_DIM) { - #pragma unroll - for (int h = 0; h < HTAP_N_HORIZONS; ++h) { - s_horizon_tokens[h * HTAP_HIDDEN_DIM + tid] = - horizon_tokens[h * HTAP_HIDDEN_DIM + tid]; - } - } - for (int i = tid; i < n_ext; i += HTAP_BLOCK) { - s_attn[i] = attn[(long long)b * n_ext + i]; - } - __syncthreads(); - - const float q_d = (tid < HTAP_HIDDEN_DIM) ? Q[tid] : 0.0f; - const float* ln_b = ln_out + (long long)b * k_seq * HTAP_HIDDEN_DIM; - float* grad_ln_b = grad_ln_out + (long long)b * k_seq * HTAP_HIDDEN_DIM; - - // ── d_S[d] = Σ_h grad_ctx[h, d]. Each thread owns d = tid (if < H). - float d_S = 0.0f; - if (tid < HTAP_HIDDEN_DIM) { - #pragma unroll - for (int h = 0; h < HTAP_N_HORIZONS; ++h) { - d_S += grad_ctx[(long long)b * HTAP_N_HORIZONS * HTAP_HIDDEN_DIM - + (long long)h * HTAP_HIDDEN_DIM + tid]; - } - s_d_S[tid] = d_S; - } - __syncthreads(); - - // ── d_attn[i] for both ranges, stored in s_d_attn. - // Token range i ∈ [0, N_H): d_attn[h] = Σ_d grad_ctx[h, d] · h_token[h, d] - // Time range i ∈ [N_H, n_ext): d_attn[N_H + k] = Σ_d d_S[d] · ln_out[k, d] - for (int h = 0; h < HTAP_N_HORIZONS; ++h) { - float v = 0.0f; - if (tid < HTAP_HIDDEN_DIM) { - v = grad_ctx[(long long)b * HTAP_N_HORIZONS * HTAP_HIDDEN_DIM - + (long long)h * HTAP_HIDDEN_DIM + tid] - * s_horizon_tokens[h * HTAP_HIDDEN_DIM + tid]; - } - float r = block_reduce_sum(v, s_warp, tid); - if (tid == 0) s_d_attn[h] = r; - } - for (int k = 0; k < k_seq; ++k) { - float v = 0.0f; - if (tid < HTAP_HIDDEN_DIM) { - v = s_d_S[tid] * ln_b[k * HTAP_HIDDEN_DIM + tid]; - } - float r = block_reduce_sum(v, s_warp, tid); - if (tid == 0) s_d_attn[HTAP_N_HORIZONS + k] = r; - } - - // ── Softmax backward: d_score[i] = attn[i] · (d_attn[i] − Σ_j attn[j] · d_attn[j]) - float my_sum = 0.0f; - for (int i = tid; i < n_ext; i += HTAP_BLOCK) { - my_sum += s_attn[i] * s_d_attn[i]; - } - const float dot_attn_dattn = block_reduce_sum(my_sum, s_warp, tid); - - // Reuse s_d_attn slot to hold d_score; s_attn still holds attn. - for (int i = tid; i < n_ext; i += HTAP_BLOCK) { - s_d_attn[i] = s_attn[i] * (s_d_attn[i] - dot_attn_dattn); - } - __syncthreads(); - - // ── d_horizon_tokens[h, d] = d_score[h] · Q[d] + grad_ctx[h, d] · attn[h] - // grad_horizon_tokens_scratch is per-batch; reduce_axis0 collapses - // over B downstream. - if (tid < HTAP_HIDDEN_DIM) { - #pragma unroll - for (int h = 0; h < HTAP_N_HORIZONS; ++h) { - const float dscore_h = s_d_attn[h]; - const float gc_hd = grad_ctx[(long long)b * HTAP_N_HORIZONS * HTAP_HIDDEN_DIM - + (long long)h * HTAP_HIDDEN_DIM + tid]; - const float attn_h = s_attn[h]; - const float dq = dscore_h * q_d + gc_hd * attn_h; - grad_horizon_tokens_scratch[(long long)b * HTAP_N_HORIZONS * HTAP_HIDDEN_DIM - + (long long)h * HTAP_HIDDEN_DIM + tid] += dq; - } - } - - // ── d_Q[d] = Σ_h d_score[h] · h_token[h, d] + Σ_k d_score[N_H+k] · ln_out[k, d] - if (tid < HTAP_HIDDEN_DIM) { - float dq_d = 0.0f; - #pragma unroll - for (int h = 0; h < HTAP_N_HORIZONS; ++h) { - dq_d += s_d_attn[h] * s_horizon_tokens[h * HTAP_HIDDEN_DIM + tid]; - } - for (int k = 0; k < k_seq; ++k) { - dq_d += s_d_attn[HTAP_N_HORIZONS + k] * ln_b[k * HTAP_HIDDEN_DIM + tid]; - } - grad_Q_scratch[(long long)b * HTAP_HIDDEN_DIM + tid] += dq_d; - } - - // ── d_ln_out[k, d] = d_score[N_H+k] · Q[d] + attn[N_H+k] · d_S[d] - // Per-batch slice → no cross-block race. - if (tid < HTAP_HIDDEN_DIM) { - for (int k = 0; k < k_seq; ++k) { - const float dscore_k = s_d_attn[HTAP_N_HORIZONS + k]; - const float attn_k = s_attn[HTAP_N_HORIZONS + k]; - grad_ln_b[k * HTAP_HIDDEN_DIM + tid] += dscore_k * q_d + attn_k * s_d_S[tid]; - } - } -} diff --git a/crates/ml-alpha/cuda/inv_pooled_merge.cu b/crates/ml-alpha/cuda/inv_pooled_merge.cu deleted file mode 100644 index 994add3ef..000000000 --- a/crates/ml-alpha/cuda/inv_pooled_merge.cu +++ /dev/null @@ -1,64 +0,0 @@ -// inv_pooled_merge.cu — Broadcast-add merge between inverted-attention -// pooled output and the per-horizon ctx_h tensor. This is the missing -// wire that connects axis E (cross-variate attention) into the -// forward chain → MoE → loss. -// -// FORWARD: -// ctx_h[b, h, d] += inv_pooled[b, d] (broadcast over h) -// -// BACKWARD: -// ∂L/∂ctx_h[b, h, d] = ∂L/∂ctx_h_post[b, h, d] (passthrough; the -// additive merge has gradient 1 wrt either input) -// ∂L/∂inv_pooled[b, d] = Σ_h ∂L/∂ctx_h_post[b, h, d] -// -// PERF: tiny kernels (~2 KB shared mem). Block-per-batch, thread-per-d. -// -// Per `feedback_nvidia_grade_perf_for_kernels`: -// - Coalesced reads/writes (thread → adjacent d) -// - No `__syncthreads` (each thread independent) -// - No DRAM round-trips beyond the necessary reads - -#define MERGE_HIDDEN_DIM 128 -#define MERGE_N_HORIZONS 5 -#define MERGE_BLOCK 128 - -// Forward: ctx_h[b, h, d] += inv_pooled[b, d]. -// Grid (B), block (HIDDEN_DIM). Each thread d updates 5 ctx_h cells. -extern "C" __global__ void inv_pooled_merge_fwd( - float* __restrict__ ctx_h, // [B, N_H, H] IN/OUT (+=) - const float* __restrict__ inv_pooled, // [B, H] - int n_batch -) { - int b = blockIdx.x; - int tid = threadIdx.x; - if (b >= n_batch || tid >= MERGE_HIDDEN_DIM) return; - - const float v = inv_pooled[(long long)b * MERGE_HIDDEN_DIM + tid]; - #pragma unroll - for (int h = 0; h < MERGE_N_HORIZONS; ++h) { - ctx_h[(long long)b * MERGE_N_HORIZONS * MERGE_HIDDEN_DIM - + (long long)h * MERGE_HIDDEN_DIM + tid] += v; - } -} - -// Backward: grad_inv_pooled[b, d] = Σ_h grad_ctx_h_post[b, h, d]. -// Grid (B), block (HIDDEN_DIM). Each thread d sums 5 ctx_h cells. -// `grad_ctx_h` is read but NOT modified (the additive merge has -// passthrough gradient on ctx_h — same buffer flows into pool.bwd). -extern "C" __global__ void inv_pooled_merge_bwd_sum_horizon( - const float* __restrict__ grad_ctx_h, // [B, N_H, H] - int n_batch, - float* __restrict__ grad_inv_pooled // [B, H] OVERWRITE -) { - int b = blockIdx.x; - int tid = threadIdx.x; - if (b >= n_batch || tid >= MERGE_HIDDEN_DIM) return; - - float acc = 0.0f; - #pragma unroll - for (int h = 0; h < MERGE_N_HORIZONS; ++h) { - acc += grad_ctx_h[(long long)b * MERGE_N_HORIZONS * MERGE_HIDDEN_DIM - + (long long)h * MERGE_HIDDEN_DIM + tid]; - } - grad_inv_pooled[(long long)b * MERGE_HIDDEN_DIM + tid] = acc; -} diff --git a/crates/ml-alpha/cuda/inverted_attention_pool.cu b/crates/ml-alpha/cuda/inverted_attention_pool.cu deleted file mode 100644 index 3cef0f975..000000000 --- a/crates/ml-alpha/cuda/inverted_attention_pool.cu +++ /dev/null @@ -1,265 +0,0 @@ -// inverted_attention_pool.cu — iTransformer-style inverted attention. -// -// Cross-variate attention: each of HIDDEN_DIM features is a "variate -// token" with its K-trajectory as embedding. Forward computes attention -// over the H×H feature pairs and pools via mean-K commute. Backward -// emits gradients through the value/query/key paths and softmax. -// -// FORWARD MATH (per batch b; b indexing suppressed): -// -// X_inv[h, k] = ln_out[k, h] # [H, K] -// scores[h, j] = inv_scale · Σ_k X_inv[h, k] · X_inv[j, k] -// attn[h, j] = softmax_j(scores) # [H, H] -// pooled[h] = Σ_j attn[h, j] · mean_k_X_inv[j] # via -// # mean_k_X_inv[j] = (1/K) Σ_k X_inv[j, k] -// -// PERF NOTE (canonical 2026-05-18 rewrite): -// - `mean_k_X_inv[H]` cached ONCE in shared memory at kernel -// entry — eliminates the ~16K wasted ops/thread from prior -// per-thread recomputation in both fwd and bwd hot loops. -// - Single-pass softmax-pool fusion in forward: score → max → -// exp+sum+pool computed without storing intermediate `scores` -// row in DRAM. Final `attn[h, ·]` is still emitted to DRAM for -// bwd (saves recomputation in backward). -// -// PERFORMANCE (per feedback_nvidia_grade_perf_for_kernels): -// - Block-per-batch; grid = (B). Block = (HIDDEN_DIM) = 128 threads. -// - Smem: x_inv [H * K] + mean_k_X_inv [H] (cached once). For K = 32 -// this is 16 KB + 0.5 KB = 16.5 KB, well under 48 KB dynamic limit. -// - All inactive-lane shuffles use ternary fallback (non-divergent). - -#define IAP_HIDDEN_DIM 128 -#define IAP_BLOCK 128 -#define IAP_MAX_KSEQ 128 - -__device__ __forceinline__ float warp_reduce_sum_iap(float v) { - #pragma unroll - for (int s = 16; s > 0; s >>= 1) { - v += __shfl_xor_sync(0xffffffff, v, s); - } - return v; -} - -extern "C" __global__ void inverted_attention_pool_fwd( - const float* __restrict__ ln_out, // [B, K, H] - int n_batch, - int k_seq, - float inv_scale, // 1 / sqrt(K) - float* __restrict__ pooled_out, // [B, H] - float* __restrict__ attn_out // [B, H, H] -) { - int b = blockIdx.x; - int tid = threadIdx.x; - if (b >= n_batch || tid >= IAP_BLOCK) return; - - extern __shared__ float smem[]; - // Layout: x_inv[H * K] + mean_k[H] - float* x_inv = smem; - float* mean_k = smem + IAP_HIDDEN_DIM * k_seq; - - // Cache X_inv = ln_out^T for this batch + compute mean_k_X_inv[h] - // in one pass per thread. Each thread h loads its own k-trajectory - // (K reads) and sums them. - if (tid < IAP_HIDDEN_DIM) { - const float* ln_b = ln_out + (long long)b * k_seq * IAP_HIDDEN_DIM; - float sum_k = 0.0f; - #pragma unroll 8 - for (int k = 0; k < k_seq; ++k) { - const float v = ln_b[k * IAP_HIDDEN_DIM + tid]; - x_inv[tid * k_seq + k] = v; - sum_k += v; - } - mean_k[tid] = sum_k * (1.0f / (float)k_seq); - } - __syncthreads(); - - // Each thread = row h. Single-pass softmax-pool fusion: - // 1. Find max(score[h, j]) over j without storing scores. - // 2. Re-scan j: compute exp(score - max), accumulate sum, write - // attn[h, j] = exp/sum_running, accumulate pool += attn · mean_k[j]. - // 3. Final correction: pool /= total_sum (since we used partial sums). - // - // Implementation: do two passes over j. Pass 1 computes max only. - // Pass 2 computes exp+sum+attn-write+pool in one sweep. Both passes - // recompute scores from x_inv (smem reads, no DRAM hit). Trade - // extra compute for one fewer DRAM round-trip per j. - if (tid < IAP_HIDDEN_DIM) { - const int h = tid; - const float* row_h = x_inv + (long long)h * k_seq; - const long long attn_row_base = (long long)b * IAP_HIDDEN_DIM * IAP_HIDDEN_DIM - + (long long)h * IAP_HIDDEN_DIM; - - // Pass 1: max(score) over j. - float my_max = -INFINITY; - for (int j = 0; j < IAP_HIDDEN_DIM; ++j) { - const float* row_j = x_inv + (long long)j * k_seq; - float dot = 0.0f; - #pragma unroll 8 - for (int k = 0; k < k_seq; ++k) { - dot += row_h[k] * row_j[k]; - } - const float s = dot * inv_scale; - if (s > my_max) my_max = s; - } - - // Pass 2: softmax + pool. Two sub-passes (sum, then divide-and-pool). - // To do it in one sweep we'd need rescaling tricks; instead just - // store exp values in DRAM (attn_out as scratch), then a final - // normalisation sweep computes both attn and pool. - float my_sum = 0.0f; - for (int j = 0; j < IAP_HIDDEN_DIM; ++j) { - const float* row_j = x_inv + (long long)j * k_seq; - float dot = 0.0f; - #pragma unroll 8 - for (int k = 0; k < k_seq; ++k) { - dot += row_h[k] * row_j[k]; - } - const float e = expf(dot * inv_scale - my_max); - attn_out[attn_row_base + j] = e; - my_sum += e; - } - - // Pass 3: divide-and-pool. Single DRAM read per j of the exp; - // overwrites with normalised attn. Pool uses cached mean_k. - const float inv_sum = 1.0f / my_sum; - float pooled_h = 0.0f; - for (int j = 0; j < IAP_HIDDEN_DIM; ++j) { - const float a_hj = attn_out[attn_row_base + j] * inv_sum; - attn_out[attn_row_base + j] = a_hj; - pooled_h += a_hj * mean_k[j]; - } - pooled_out[(long long)b * IAP_HIDDEN_DIM + h] = pooled_h; - } -} - -extern "C" __global__ void inverted_attention_pool_bwd( - const float* __restrict__ ln_out, // [B, K, H] - const float* __restrict__ attn, // [B, H, H] - const float* __restrict__ grad_pooled, // [B, H] - int n_batch, - int k_seq, - float inv_scale, - float* __restrict__ d_scores_scratch, // [B, H, H] - float* __restrict__ grad_ln_out // [B, K, H] += per-batch -) { - int b = blockIdx.x; - int tid = threadIdx.x; - if (b >= n_batch || tid >= IAP_BLOCK) return; - - extern __shared__ float smem_bwd[]; - // Layout: x_inv[H * K] + mean_k[H] + dp_shared[H] - float* x_inv = smem_bwd; - float* mean_k = smem_bwd + IAP_HIDDEN_DIM * k_seq; - float* dp_shared = smem_bwd + IAP_HIDDEN_DIM * k_seq + IAP_HIDDEN_DIM; - - if (tid < IAP_HIDDEN_DIM) { - const float* ln_b = ln_out + (long long)b * k_seq * IAP_HIDDEN_DIM; - float sum_k = 0.0f; - #pragma unroll 8 - for (int k = 0; k < k_seq; ++k) { - const float v = ln_b[k * IAP_HIDDEN_DIM + tid]; - x_inv[tid * k_seq + k] = v; - sum_k += v; - } - mean_k[tid] = sum_k * (1.0f / (float)k_seq); - dp_shared[tid] = grad_pooled[(long long)b * IAP_HIDDEN_DIM + tid]; - } - __syncthreads(); - - // Phase 1: compute d_scores[my_h, j] for all j; write to DRAM scratch. - // - // d_attn[my_h, j] = dp_my_h · mean_k[j] (cached) - // dot = Σ_l attn[my_h, l] · mean_k[l] (== pooled[my_h]) - // (but pooled isn't passed; we compute it here) - // d_scores[my_h, j] = attn[my_h, j] · (d_attn − dp_my_h · dot) - // = attn[my_h, j] · dp_my_h · (mean_k[j] − dot) - if (tid < IAP_HIDDEN_DIM) { - const int my_h = tid; - const long long attn_row_base = (long long)b * IAP_HIDDEN_DIM * IAP_HIDDEN_DIM - + (long long)my_h * IAP_HIDDEN_DIM; - const float dp_my_h = dp_shared[my_h]; - - // Compute `dot = Σ_l attn[my_h, l] · mean_k[l]` (= pooled[my_h]). - float dot = 0.0f; - for (int j = 0; j < IAP_HIDDEN_DIM; ++j) { - dot += attn[attn_row_base + j] * mean_k[j]; - } - - // Write d_scores[my_h, ·]. Single DRAM write per j (no recomputation - // of mean_k inside the loop — uses the smem-cached version). - for (int j = 0; j < IAP_HIDDEN_DIM; ++j) { - const float a_my_h_j = attn[attn_row_base + j]; - const float d_score = a_my_h_j * dp_my_h * (mean_k[j] - dot); - d_scores_scratch[attn_row_base + j] = d_score; - } - } - __syncthreads(); - - // Phase 2: accumulate d_ln_out[k, my_h] for each k. - // - // Per-(my_h, k): - // value : V_my_h = (1/K) · Σ_h' dp[h'] · attn[h', my_h] (per-k constant) - // query : q_k = inv_scale · Σ_j d_scores[my_h, j] · X_inv[j, k] - // key : k_k = inv_scale · Σ_h' d_scores[h', my_h] · X_inv[h', k] - // - // KEY OPTIMIZATION (canonical 2026-05-18 perf rewrite): - // The d_scores reads do NOT depend on k → hoist them out of the - // K loop. We loop J-outer with two register accumulators per - // thread (q_k[K_MAX], k_k[K_MAX]). For each j we do exactly TWO - // DRAM reads of d_scores (row + column) instead of `2 * K` = - // 2 × 32 = 64 reads inside a tight loop. Net: 32× fewer DRAM - // reads for the d_scores tensor. - // - // For K up to 64 (production max), the per-thread float arrays - // stay in registers (NVCC --maxrregcount permits it on sm_86+). - if (tid < IAP_HIDDEN_DIM) { - const int my_h = tid; - - // Value-path scalar: read `attn[:, my_h]` (strided!) and dp from smem. - float V_my_h = 0.0f; - for (int hp = 0; hp < IAP_HIDDEN_DIM; ++hp) { - const float a_hp_myh = attn[(long long)b * IAP_HIDDEN_DIM * IAP_HIDDEN_DIM - + (long long)hp * IAP_HIDDEN_DIM + my_h]; - V_my_h += dp_shared[hp] * a_hp_myh; - } - V_my_h *= (1.0f / (float)k_seq); - - float* d_ln_b = grad_ln_out + (long long)b * k_seq * IAP_HIDDEN_DIM; - const long long my_row_base = (long long)b * IAP_HIDDEN_DIM * IAP_HIDDEN_DIM - + (long long)my_h * IAP_HIDDEN_DIM; - - // Per-thread register accumulators for q-path and k-path, - // one per K position. IAP_MAX_KSEQ = 128 floats = 512 B per - // thread total = 64 KB across the 128-thread block — fits - // within the register file on sm_86+ with --maxrregcount=128. - float q_arr[IAP_MAX_KSEQ]; - float k_arr[IAP_MAX_KSEQ]; - #pragma unroll - for (int k = 0; k < IAP_MAX_KSEQ; ++k) { - q_arr[k] = 0.0f; - k_arr[k] = 0.0f; - } - - // J-outer: each j read of d_scores happens exactly ONCE. - // The inner K-loop walks x_inv[j, k] which is contiguous in smem. - for (int j = 0; j < IAP_HIDDEN_DIM; ++j) { - const float ds_myh_j = d_scores_scratch[my_row_base + j]; - const float ds_j_myh = d_scores_scratch[ - (long long)b * IAP_HIDDEN_DIM * IAP_HIDDEN_DIM - + (long long)j * IAP_HIDDEN_DIM + my_h]; - const float* xj = x_inv + (long long)j * k_seq; - #pragma unroll 8 - for (int k = 0; k < k_seq; ++k) { - const float xj_k = xj[k]; - q_arr[k] += ds_myh_j * xj_k; - k_arr[k] += ds_j_myh * xj_k; - } - } - - // Final write-out per k. - for (int k = 0; k < k_seq; ++k) { - d_ln_b[k * IAP_HIDDEN_DIM + my_h] += - V_my_h + inv_scale * (q_arr[k] + k_arr[k]); - } - } -} diff --git a/crates/ml-alpha/cuda/regime_moe_gate.cu b/crates/ml-alpha/cuda/regime_moe_gate.cu deleted file mode 100644 index 09dac7716..000000000 --- a/crates/ml-alpha/cuda/regime_moe_gate.cu +++ /dev/null @@ -1,300 +0,0 @@ -// regime_moe_gate.cu — Regime-aware top-1 Mixture-of-Experts gate (v2 axis D). -// -// Routes the fused per-horizon context through one of N_EXPERTS expert -// linear layers based on a regime-feature gate. Top-1 routing per -// Switch Transformer (Fedus et al. 2021) with straight-through grad on -// the gate logits and an auxiliary load-balancing loss. -// -// FORWARD MATH (per batch b, suppressing b): -// -// gate_logits[e] = Σ_r W_gate[e, r] · R[r] # [N_EXPERTS] -// gate_probs[e] = softmax_e(gate_logits[e]) # [N_EXPERTS] -// top_e = argmax_e(gate_logits[e]) # scalar -// expert_out[e](h, d_out) = Σ_d_in W_e[e, d_out, d_in] · fused_ctx[h, d_in] -// + b_e[e, d_out] # not all e computed; only top_e -// routed_ctx[h, d] = expert_out[top_e][h, d] -// -// BACKWARD (straight-through estimator on gate, full grad on the -// selected expert): -// d_W_e[top_e, d_out, d_in] += Σ_{b, h} fused_ctx[b, h, d_in] · d_routed[b, h, d_out] -// d_b_e[top_e, d_out] += Σ_{b, h} d_routed[b, h, d_out] -// d_fused_ctx[b, h, d] += Σ_{d_out} W_e[top_e, d_out, d] · d_routed[b, h, d_out] -// -// STE on gate_logits[b, e]: pass an upstream "router loss signal" -// proportional to the expert output norm — concretely, route the -// gradient ∂L/∂routed[b, h, ·] · expert_out[e][h, ·] for e == top_e[b] -// into gate_logits via the softmax Jacobian. For inactive experts, -// STE assumes the same loss signal would have been produced (zero -// selection grad). Aux load-balance grad is added separately. -// -// AUXILIARY LOAD-BALANCING LOSS (computed by a separate small kernel): -// frac[e] = mean_b({top_e[b] == e}) -// prob_mean[e] = mean_b(gate_probs[b, e]) -// aux_loss = N_EXPERTS · Σ_e frac[e] · prob_mean[e] -// d_gate_logits aux contribution: standard softmax-frac-cross-entropy -// -// PERFORMANCE: -// - One block per (b, h): forward and backward have block-per-(b, h) -// parallelism for the per-horizon linear (matches V4 pattern). -// - Warp-shuffle reduce for the d_in dot product (HIDDEN_DIM = 128 = -// 4 warps). All sums are warp-shuffle, not block tree-reduce. -// - Top-1 selection per batch is decided in a pre-pass: ONE thread -// per block reads gate_logits[b, :] and picks the argmax, broadcasts -// via shared mem. -// - Per-expert weights live in [N_EXPERTS, H, H] = ~250 KB total; on -// the hot path each block only fetches the selected expert's slice. - -#define MOE_HIDDEN_DIM 128 -#define MOE_BLOCK 128 -#define MOE_N_HORIZONS 5 -#define MOE_N_EXPERTS 4 -#define MOE_N_WARPS (MOE_BLOCK / 32) - -// ── Forward: gate + expert dispatch fused into one kernel. -// Grid: (B, N_H, 1). Block: (HIDDEN_DIM, 1, 1). -// Each block: (1) reads gate_logits, picks top_e for this batch, -// (2) computes routed_ctx[b, h, d_out=tid]. -extern "C" __global__ void regime_moe_gate_fwd( - const float* __restrict__ gate_logits, // [B, N_EXPERTS] pre-softmax (caller computes) - const float* __restrict__ fused_ctx, // [B, N_H, H] - const float* __restrict__ experts_W, // [N_EXPERTS, H, H] - const float* __restrict__ experts_b, // [N_EXPERTS, H] - int n_batch, - float* __restrict__ routed_out, // [B, N_H, H] - int* __restrict__ top_e_out // [B] top-1 expert per batch (saved) -) { - int b = blockIdx.x; - int h = blockIdx.y; - int tid = threadIdx.x; - if (b >= n_batch || h >= MOE_N_HORIZONS || tid >= MOE_BLOCK) return; - - // Pick top-1 expert by thread 0; broadcast to the block. - __shared__ int s_top_e; - if (tid == 0) { - float best = -INFINITY; - int best_e = 0; - #pragma unroll - for (int e = 0; e < MOE_N_EXPERTS; ++e) { - const float v = gate_logits[(long long)b * MOE_N_EXPERTS + e]; - if (v > best) { best = v; best_e = e; } - } - s_top_e = best_e; - // Only one block per batch needs to write top_e_out (the h == 0 - // block by convention) — but writing from each block is fine - // since the value is identical. - if (h == 0) top_e_out[b] = best_e; - } - __syncthreads(); - const int top_e = s_top_e; - - // Cache fused_ctx[b, h, ·] in shared mem. - __shared__ float s_in[MOE_HIDDEN_DIM]; - if (tid < MOE_HIDDEN_DIM) { - s_in[tid] = fused_ctx[(long long)b * MOE_N_HORIZONS * MOE_HIDDEN_DIM - + (long long)h * MOE_HIDDEN_DIM + tid]; - } - __syncthreads(); - - // Each thread computes one output channel: routed[b, h, tid] = - // Σ_d_in W_e[top_e, tid, d_in] · s_in[d_in] + b_e[top_e, tid]. - if (tid < MOE_HIDDEN_DIM) { - const long long w_base = (long long)top_e * MOE_HIDDEN_DIM * MOE_HIDDEN_DIM - + (long long)tid * MOE_HIDDEN_DIM; - float acc = experts_b[(long long)top_e * MOE_HIDDEN_DIM + tid]; - #pragma unroll 8 - for (int d_in = 0; d_in < MOE_HIDDEN_DIM; ++d_in) { - acc += experts_W[w_base + d_in] * s_in[d_in]; - } - routed_out[(long long)b * MOE_N_HORIZONS * MOE_HIDDEN_DIM - + (long long)h * MOE_HIDDEN_DIM + tid] = acc; - } -} - -// ── Backward expert path only (gate STE handled in a separate kernel). -// Same block layout as forward. -// For the SELECTED expert top_e[b]: -// d_W_scratch[b, h, e=top_e, d_out, d_in] = d_routed[b, h, d_out] · fused_ctx[b, h, d_in] -// d_b_scratch[b, h, e=top_e, d_out] = d_routed[b, h, d_out] -// d_fused_ctx[b, h, d_in] += Σ_d_out W_e[top_e, d_out, d_in] · d_routed[b, h, d_out] -// Inactive experts: 0 contribution. -// -// Per-(b, h) scratches are massive (B × N_H × N_E × H × H). We instead -// have CALLER allocate per-(b, h) scratches that are sparse-by-expert: -// the kernel only writes to the slot corresponding to top_e for this -// (b, h). Total scratch size: [B, N_H, N_E, H, H] = small at trainer -// scale (B=1, N_H=5, N_E=4, H=128 → 2.5 MB). -// d_W_scratch layout changed 2026-05-18 for perf: [B, N_H, H, H] (was -// [B, N_H, N_E, H, H] — 4× smaller, eliminating the wasteful N_E axis -// since each batch routes to exactly ONE expert in top-1 MoE). The -// expert index `top_e[b]` is needed downstream when scattering the -// (B, N_H)-reduced sum into the per-expert grad slot. -extern "C" __global__ void regime_moe_gate_bwd( - const float* __restrict__ fused_ctx, // [B, N_H, H] - const float* __restrict__ experts_W, // [N_EXPERTS, H, H] - const float* __restrict__ d_routed, // [B, N_H, H] - const int* __restrict__ top_e, // [B] - int n_batch, - float* __restrict__ d_W_scratch, // [B, N_H, H, H] per-batch-horizon rank-1 contribution - float* __restrict__ d_b_scratch, // [B, N_H, H] per-batch-horizon bias contribution - float* __restrict__ d_fused_ctx // [B, N_H, H] += per-(b, h) -) { - int b = blockIdx.x; - int h = blockIdx.y; - int tid = threadIdx.x; - if (b >= n_batch || h >= MOE_N_HORIZONS || tid >= MOE_BLOCK) return; - - const int e = top_e[b]; - - __shared__ float s_in[MOE_HIDDEN_DIM]; - __shared__ float s_dout[MOE_HIDDEN_DIM]; - if (tid < MOE_HIDDEN_DIM) { - s_in[tid] = fused_ctx[(long long)b * MOE_N_HORIZONS * MOE_HIDDEN_DIM - + (long long)h * MOE_HIDDEN_DIM + tid]; - s_dout[tid] = d_routed[(long long)b * MOE_N_HORIZONS * MOE_HIDDEN_DIM - + (long long)h * MOE_HIDDEN_DIM + tid]; - } - __syncthreads(); - - // d_W_scratch[b, h, tid (= d_out), d_in] = d_routed[b, h, d_out] · fused_ctx[b, h, d_in] - // Each thread d_out writes one ROW of [H, H] (length H = MOE_HIDDEN_DIM). - // d_b_scratch[b, h, d_out] = d_routed[b, h, d_out] - if (tid < MOE_HIDDEN_DIM) { - const long long w_base = (((long long)b * MOE_N_HORIZONS + h) - * MOE_HIDDEN_DIM + tid) * MOE_HIDDEN_DIM; - const float g = s_dout[tid]; - #pragma unroll 8 - for (int d_in = 0; d_in < MOE_HIDDEN_DIM; ++d_in) { - d_W_scratch[w_base + d_in] = g * s_in[d_in]; - } - d_b_scratch[((long long)b * MOE_N_HORIZONS + h) - * MOE_HIDDEN_DIM + tid] = g; - } - - // d_fused_ctx[b, h, d_in] += Σ_d_out W_e[e, d_out, d_in] · d_routed[b, h, d_out] - // Each thread owns d_in = tid. Reduce over d_out via warp-shuffle. - __syncthreads(); - if (tid < MOE_HIDDEN_DIM) { - const int d_in = tid; - float acc = 0.0f; - const long long w_col_base = (long long)e * MOE_HIDDEN_DIM * MOE_HIDDEN_DIM; - // W_e[e, d_out, d_in]: stride = MOE_HIDDEN_DIM between rows. - for (int d_out = 0; d_out < MOE_HIDDEN_DIM; ++d_out) { - acc += experts_W[w_col_base + (long long)d_out * MOE_HIDDEN_DIM + d_in] - * s_dout[d_out]; - } - d_fused_ctx[(long long)b * MOE_N_HORIZONS * MOE_HIDDEN_DIM - + (long long)h * MOE_HIDDEN_DIM + d_in] = acc; - } -} - -// ── Scatter per-batch grads to the correct expert slot. -// After `regime_moe_gate_bwd` writes the rank-1 contributions into -// `d_W_scratch [B, N_H, H, H]` and `d_b_scratch [B, N_H, H]`, this -// kernel scatters them into `grad_experts_W [N_E, H, H]` / -// `grad_experts_b [N_E, H]` based on `top_e[b]`. Each batch's -// contribution lands in slot `top_e[b]` of the per-expert grad. -// -// Grid (N_E, H, H/32 + 1), block (32). Each thread iterates over -// (b, h) pairs and accumulates into its assigned (e, d_out, d_in) -// cell only if top_e[b] matches. Per cell: 5 contributions max -// (B=1 × N_H=5). Coalesced reads of d_W_scratch row. -extern "C" __global__ void regime_moe_gate_scatter( - const float* __restrict__ d_W_scratch, // [B, N_H, H, H] - const float* __restrict__ d_b_scratch, // [B, N_H, H] - const int* __restrict__ top_e, // [B] - int n_batch, - float* __restrict__ grad_experts_W, // [N_E, H, H] OVERWRITE - float* __restrict__ grad_experts_b // [N_E, H] OVERWRITE -) { - int e = blockIdx.x; - int d_out = blockIdx.y; - int d_in_chunk = blockIdx.z; - int lane = threadIdx.x; - if (e >= MOE_N_EXPERTS || d_out >= MOE_HIDDEN_DIM || lane >= 32) return; - - const int d_in_start = d_in_chunk * 32; - const int d_in = d_in_start + lane; - if (d_in >= MOE_HIDDEN_DIM) return; - - // Accumulate W[e, d_out, d_in] across batches where top_e[b] == e. - float acc_w = 0.0f; - float acc_b = (lane == 0 && d_out == 0) ? 0.0f : 0.0f; // bias only via warp 0 - for (int b = 0; b < n_batch; ++b) { - if (top_e[b] != e) continue; - for (int h = 0; h < MOE_N_HORIZONS; ++h) { - const long long w_idx = (((long long)b * MOE_N_HORIZONS + h) - * MOE_HIDDEN_DIM + d_out) * MOE_HIDDEN_DIM + d_in; - acc_w += d_W_scratch[w_idx]; - if (lane == 0) { - const long long b_idx = ((long long)b * MOE_N_HORIZONS + h) - * MOE_HIDDEN_DIM + d_out; - acc_b += d_b_scratch[b_idx]; - } - } - } - grad_experts_W[((long long)e * MOE_HIDDEN_DIM + d_out) * MOE_HIDDEN_DIM + d_in] = acc_w; - if (lane == 0) { - grad_experts_b[(long long)e * MOE_HIDDEN_DIM + d_out] = acc_b; - } -} - -// ── Gate softmax + auxiliary load-balancing loss. -// Computes gate_probs (for the optional STE gate-grad path) and the -// `aux_loss` scalar that adds to total loss to discourage expert -// collapse. -// Grid: (1, 1, 1). Block: (N_EXPERTS, 1, 1) (= 4 threads). -// All-in-one — input sizes are tiny. -extern "C" __global__ void regime_moe_gate_aux( - const float* __restrict__ gate_logits, // [B, N_E] - const int* __restrict__ top_e, // [B] - int n_batch, - float* __restrict__ gate_probs_out, // [B, N_E] - float* __restrict__ aux_loss_out // [1] -) { - int e = threadIdx.x; - if (e >= MOE_N_EXPERTS) return; - - // Phase 1: softmax per batch and write gate_probs. - // We do one full sweep here, per-thread = per-expert column. - float frac_local = 0.0f; - float prob_mean_local = 0.0f; - for (int b = 0; b < n_batch; ++b) { - // Compute max + sum once per batch, by thread 0; broadcast. - __shared__ float s_max; - __shared__ float s_sum; - if (e == 0) { - float mx = -INFINITY; - float sm = 0.0f; - #pragma unroll - for (int ep = 0; ep < MOE_N_EXPERTS; ++ep) { - const float v = gate_logits[(long long)b * MOE_N_EXPERTS + ep]; - if (v > mx) mx = v; - } - #pragma unroll - for (int ep = 0; ep < MOE_N_EXPERTS; ++ep) { - sm += expf(gate_logits[(long long)b * MOE_N_EXPERTS + ep] - mx); - } - s_max = mx; - s_sum = sm; - } - __syncthreads(); - - const float p_e = expf(gate_logits[(long long)b * MOE_N_EXPERTS + e] - s_max) / s_sum; - gate_probs_out[(long long)b * MOE_N_EXPERTS + e] = p_e; - prob_mean_local += p_e; - if (top_e[b] == e) frac_local += 1.0f; - __syncthreads(); - } - frac_local *= (1.0f / (float)n_batch); - prob_mean_local *= (1.0f / (float)n_batch); - - // Phase 2: aux_loss = N_EXPERTS · Σ_e frac · prob_mean. - // Warp-reduce across the 4 threads. - float aux_contrib = frac_local * prob_mean_local; - // Warp shuffle on a 4-thread block. - #pragma unroll - for (int s = 16; s > 0; s >>= 1) { - aux_contrib += __shfl_xor_sync(0xffffffff, aux_contrib, s); - } - if (e == 0) aux_loss_out[0] = (float)MOE_N_EXPERTS * aux_contrib; -} diff --git a/crates/ml-alpha/src/anchor_l2.rs b/crates/ml-alpha/src/anchor_l2.rs deleted file mode 100644 index fa6b85142..000000000 --- a/crates/ml-alpha/src/anchor_l2.rs +++ /dev/null @@ -1,59 +0,0 @@ -//! L2 anchor regularization host binding (v2 axis B, V8 commit). -//! -//! Anchors a trainable parameter buffer toward its initialization -//! values via `loss = λ · Σ (p − p_init)²`. The Wiener-α controller -//! on the host side adjusts λ each step based on the loss-improvement -//! rate (signal-driven, floor-bounded — see `trainer/anchor_controller.rs`). -//! -//! See `cuda/anchor_l2.cu` for math + the v2 spec §3.5. - -use anyhow::{Context, Result}; -use cudarc::driver::{ - CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg, -}; -use std::sync::Arc; - -const CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/anchor_l2.cubin")); - -pub struct AnchorL2 { - _module: Arc, - func: CudaFunction, - stream: Arc, -} - -impl AnchorL2 { - pub fn new(ctx: &Arc, stream: Arc) -> Result { - let module = ctx.load_cubin(CUBIN.to_vec()).context("load anchor_l2 cubin")?; - let func = module.load_function("anchor_l2_fwd_bwd").context("load anchor_l2_fwd_bwd")?; - Ok(Self { _module: module, func, stream }) - } - - /// Compute anchor_loss = λ · ‖p − p_init‖² and accumulate - /// `grad_p += 2λ · (p − p_init)`. `lambda_d` is a device-side - /// [1]-buffer (host writes scalar before launch). - pub fn apply( - &self, - p: &CudaSlice, - p_init: &CudaSlice, - lambda_d: &CudaSlice, - n: i32, - loss_out: &mut CudaSlice, - grad_p: &mut CudaSlice, - ) -> Result<()> { - let cfg = LaunchConfig { - grid_dim: (1, 1, 1), - block_dim: (256, 1, 1), - shared_mem_bytes: 0, - }; - let mut launch = self.stream.launch_builder(&self.func); - unsafe { - launch - .arg(p).arg(p_init).arg(lambda_d) - .arg(&n) - .arg(loss_out).arg(grad_p) - .launch(cfg) - .context("anchor_l2_fwd_bwd")?; - } - Ok(()) - } -} diff --git a/crates/ml-alpha/src/horizon_mean_collapse.rs b/crates/ml-alpha/src/horizon_mean_collapse.rs deleted file mode 100644 index 046da3e3c..000000000 --- a/crates/ml-alpha/src/horizon_mean_collapse.rs +++ /dev/null @@ -1,81 +0,0 @@ -//! Horizon-axis mean-collapse host binding. -//! -//! Collapses per-horizon context `[B, N_H, H]` to a single context -//! `[B, H]` by averaging across the horizon axis. Used to seed the -//! CfC's `h_old` at k=0 from `MultiHorizonAttention.routed_ctx`. - -use anyhow::{Context, Result}; -use cudarc::driver::{ - CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg, -}; -use std::sync::Arc; - -pub const HMC_HIDDEN_DIM: usize = 128; -pub const HMC_BLOCK: usize = 128; - -const CUBIN: &[u8] = - include_bytes!(concat!(env!("OUT_DIR"), "/horizon_mean_collapse.cubin")); - -pub struct HorizonMeanCollapse { - _module: Arc, - fwd_fn: CudaFunction, - bwd_fn: CudaFunction, - stream: Arc, -} - -impl HorizonMeanCollapse { - pub fn new(ctx: &Arc, stream: Arc) -> Result { - let module = ctx - .load_cubin(CUBIN.to_vec()) - .context("load horizon_mean_collapse cubin")?; - let fwd_fn = module - .load_function("horizon_mean_collapse_fwd") - .context("load horizon_mean_collapse_fwd")?; - let bwd_fn = module - .load_function("horizon_mean_collapse_bwd") - .context("load horizon_mean_collapse_bwd")?; - Ok(Self { _module: module, fwd_fn, bwd_fn, stream }) - } - - pub fn forward( - &self, - ctx_per_h: &CudaSlice, - n_batch: i32, - ctx_mean: &mut CudaSlice, - ) -> Result<()> { - let cfg = LaunchConfig { - grid_dim: (n_batch as u32, 1, 1), - block_dim: (HMC_BLOCK as u32, 1, 1), - shared_mem_bytes: 0, - }; - let mut launch = self.stream.launch_builder(&self.fwd_fn); - unsafe { - launch - .arg(ctx_per_h).arg(&n_batch).arg(ctx_mean) - .launch(cfg) - .context("horizon_mean_collapse_fwd")?; - } - Ok(()) - } - - pub fn backward( - &self, - grad_ctx_mean: &CudaSlice, - n_batch: i32, - grad_ctx_per_h: &mut CudaSlice, - ) -> Result<()> { - let cfg = LaunchConfig { - grid_dim: (n_batch as u32, 1, 1), - block_dim: (HMC_BLOCK as u32, 1, 1), - shared_mem_bytes: 0, - }; - let mut launch = self.stream.launch_builder(&self.bwd_fn); - unsafe { - launch - .arg(grad_ctx_mean).arg(&n_batch).arg(grad_ctx_per_h) - .launch(cfg) - .context("horizon_mean_collapse_bwd")?; - } - Ok(()) - } -} diff --git a/crates/ml-alpha/src/horizon_token_attention_pool.rs b/crates/ml-alpha/src/horizon_token_attention_pool.rs deleted file mode 100644 index 4abaf1a80..000000000 --- a/crates/ml-alpha/src/horizon_token_attention_pool.rs +++ /dev/null @@ -1,122 +0,0 @@ -//! Horizon-token attention pool host binding (v2 axis C, V2 commit). -//! -//! Replaces the falsified per-horizon Q_h pool with N_HORIZONS learned -//! tokens prepended to the K-sequence and a single shared learned Q. -//! See `cuda/horizon_token_attention_pool.cu` for math + the v2 spec. -//! -//! This binding is graph-capture-safe (no `synchronize`, no host -//! allocs, no `memcpy_htod` in the hot path). - -use anyhow::{Context, Result}; -use cudarc::driver::{ - CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg, -}; -use std::sync::Arc; - -pub const HTAP_HIDDEN_DIM: usize = 128; -pub const HTAP_BLOCK: usize = 128; -pub const HTAP_N_HORIZONS: usize = 5; -pub const HTAP_N_WARPS: usize = HTAP_BLOCK / 32; - -const CUBIN: &[u8] = - include_bytes!(concat!(env!("OUT_DIR"), "/horizon_token_attention_pool.cubin")); - -pub struct HorizonTokenAttentionPool { - _module: Arc, - fwd_fn: CudaFunction, - bwd_fn: CudaFunction, - stream: Arc, -} - -impl HorizonTokenAttentionPool { - pub fn new(ctx: &Arc, stream: Arc) -> Result { - let module = ctx - .load_cubin(CUBIN.to_vec()) - .context("load horizon_token_attention_pool cubin")?; - let fwd_fn = module - .load_function("horizon_token_attention_pool_fwd") - .context("load horizon_token_attention_pool_fwd")?; - let bwd_fn = module - .load_function("horizon_token_attention_pool_bwd") - .context("load horizon_token_attention_pool_bwd")?; - Ok(Self { _module: module, fwd_fn, bwd_fn, stream }) - } - - /// Forward. All buffers must be device-resident; no synchronize. - /// `horizon_tokens` : [N_H, H] - /// `q` : [H] - /// `ln_out` : [B, K, H] - /// `ctx_out` : [B, N_H, H] written - /// `attn_out` : [B, N_H + K] written (saved for backward) - pub fn forward( - &self, - horizon_tokens: &CudaSlice, - q: &CudaSlice, - ln_out: &CudaSlice, - n_batch: i32, - k_seq: i32, - ctx_out: &mut CudaSlice, - attn_out: &mut CudaSlice, - ) -> Result<()> { - // Smem layout: s_attn[N_H + K] + s_warp[N_WARPS] - let smem_bytes = ((HTAP_N_HORIZONS + k_seq as usize + HTAP_N_WARPS) - * std::mem::size_of::()) as u32; - let cfg = LaunchConfig { - grid_dim: (n_batch as u32, 1, 1), - block_dim: (HTAP_BLOCK as u32, 1, 1), - shared_mem_bytes: smem_bytes, - }; - let mut launch = self.stream.launch_builder(&self.fwd_fn); - unsafe { - launch - .arg(horizon_tokens).arg(q).arg(ln_out) - .arg(&n_batch).arg(&k_seq) - .arg(ctx_out).arg(attn_out) - .launch(cfg) - .context("horizon_token_attention_pool_fwd")?; - } - Ok(()) - } - - /// Backward. Gradients accumulate (`+=`) into the per-batch scratch - /// buffers and `grad_ln_out`. Reduce_axis0 collapses scratches over - /// the batch axis downstream. - /// `grad_horizon_tokens_scratch` : [B, N_H, H] += - /// `grad_q_scratch` : [B, H] += - /// `grad_ln_out` : [B, K, H] += (per-batch slice; no race) - pub fn backward( - &self, - horizon_tokens: &CudaSlice, - q: &CudaSlice, - ln_out: &CudaSlice, - attn: &CudaSlice, - grad_ctx: &CudaSlice, - n_batch: i32, - k_seq: i32, - grad_horizon_tokens_scratch: &mut CudaSlice, - grad_q_scratch: &mut CudaSlice, - grad_ln_out: &mut CudaSlice, - ) -> Result<()> { - // Smem layout: s_d_attn[N_H + K] + s_warp[N_WARPS] + s_d_S[HIDDEN_DIM] - let smem_bytes = ((HTAP_N_HORIZONS + k_seq as usize + HTAP_N_WARPS + HTAP_HIDDEN_DIM) - * std::mem::size_of::()) as u32; - let cfg = LaunchConfig { - grid_dim: (n_batch as u32, 1, 1), - block_dim: (HTAP_BLOCK as u32, 1, 1), - shared_mem_bytes: smem_bytes, - }; - let mut launch = self.stream.launch_builder(&self.bwd_fn); - unsafe { - launch - .arg(horizon_tokens).arg(q).arg(ln_out) - .arg(attn).arg(grad_ctx) - .arg(&n_batch).arg(&k_seq) - .arg(grad_horizon_tokens_scratch) - .arg(grad_q_scratch) - .arg(grad_ln_out) - .launch(cfg) - .context("horizon_token_attention_pool_bwd")?; - } - Ok(()) - } -} diff --git a/crates/ml-alpha/src/inv_pooled_merge.rs b/crates/ml-alpha/src/inv_pooled_merge.rs deleted file mode 100644 index 1484fbcbb..000000000 --- a/crates/ml-alpha/src/inv_pooled_merge.rs +++ /dev/null @@ -1,82 +0,0 @@ -//! Broadcast-add merge between inverted-attention pooled output and -//! the per-horizon ctx_h tensor. Wires axis E (cross-variate -//! attention) into the forward chain → MoE → loss. -//! -//! See `cuda/inv_pooled_merge.cu` for math. - -use anyhow::{Context, Result}; -use cudarc::driver::{ - CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg, -}; -use std::sync::Arc; - -pub const MERGE_HIDDEN_DIM: usize = 128; -pub const MERGE_BLOCK: usize = 128; - -const CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/inv_pooled_merge.cubin")); - -pub struct InvPooledMerge { - _module: Arc, - fwd_fn: CudaFunction, - bwd_fn: CudaFunction, - stream: Arc, -} - -impl InvPooledMerge { - pub fn new(ctx: &Arc, stream: Arc) -> Result { - let module = ctx - .load_cubin(CUBIN.to_vec()) - .context("load inv_pooled_merge cubin")?; - let fwd_fn = module - .load_function("inv_pooled_merge_fwd") - .context("load inv_pooled_merge_fwd")?; - let bwd_fn = module - .load_function("inv_pooled_merge_bwd_sum_horizon") - .context("load inv_pooled_merge_bwd_sum_horizon")?; - Ok(Self { _module: module, fwd_fn, bwd_fn, stream }) - } - - /// `ctx_h[b, h, d] += inv_pooled[b, d]`. - pub fn forward( - &self, - ctx_h: &mut CudaSlice, - inv_pooled: &CudaSlice, - n_batch: i32, - ) -> Result<()> { - let cfg = LaunchConfig { - grid_dim: (n_batch as u32, 1, 1), - block_dim: (MERGE_BLOCK as u32, 1, 1), - shared_mem_bytes: 0, - }; - let mut launch = self.stream.launch_builder(&self.fwd_fn); - unsafe { - launch - .arg(ctx_h).arg(inv_pooled).arg(&n_batch) - .launch(cfg) - .context("inv_pooled_merge_fwd")?; - } - Ok(()) - } - - /// `grad_inv_pooled[b, d] = Σ_h grad_ctx_h[b, h, d]`. - pub fn backward( - &self, - grad_ctx_h: &CudaSlice, - n_batch: i32, - grad_inv_pooled: &mut CudaSlice, - ) -> Result<()> { - let cfg = LaunchConfig { - grid_dim: (n_batch as u32, 1, 1), - block_dim: (MERGE_BLOCK as u32, 1, 1), - shared_mem_bytes: 0, - }; - let mut launch = self.stream.launch_builder(&self.bwd_fn); - unsafe { - launch - .arg(grad_ctx_h).arg(&n_batch).arg(grad_inv_pooled) - .launch(cfg) - .context("inv_pooled_merge_bwd_sum_horizon")?; - } - Ok(()) - } -} diff --git a/crates/ml-alpha/src/inverted_attention_pool.rs b/crates/ml-alpha/src/inverted_attention_pool.rs deleted file mode 100644 index dff3999af..000000000 --- a/crates/ml-alpha/src/inverted_attention_pool.rs +++ /dev/null @@ -1,112 +0,0 @@ -//! Inverted attention pool host binding (v2 axis E, V3 commit). -//! -//! iTransformer-style cross-variate attention: each of HIDDEN_DIM -//! features becomes a "variate token" with its K-trajectory as -//! embedding. Single-block-per-batch, no learnable parameters in this -//! v2 (Q = K = V = X_inv self-attention with mean-K pool). -//! -//! See `cuda/inverted_attention_pool.cu` for math + the v2 spec §3.2. - -use anyhow::{Context, Result}; -use cudarc::driver::{ - CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg, -}; -use std::sync::Arc; - -pub const IAP_HIDDEN_DIM: usize = 128; -pub const IAP_BLOCK: usize = 128; - -const CUBIN: &[u8] = - include_bytes!(concat!(env!("OUT_DIR"), "/inverted_attention_pool.cubin")); - -pub struct InvertedAttentionPool { - _module: Arc, - fwd_fn: CudaFunction, - bwd_fn: CudaFunction, - stream: Arc, -} - -impl InvertedAttentionPool { - pub fn new(ctx: &Arc, stream: Arc) -> Result { - let module = ctx - .load_cubin(CUBIN.to_vec()) - .context("load inverted_attention_pool cubin")?; - let fwd_fn = module - .load_function("inverted_attention_pool_fwd") - .context("load inverted_attention_pool_fwd")?; - let bwd_fn = module - .load_function("inverted_attention_pool_bwd") - .context("load inverted_attention_pool_bwd")?; - Ok(Self { _module: module, fwd_fn, bwd_fn, stream }) - } - - /// Forward. `inv_scale` must equal `1 / sqrt(k_seq)` to match the - /// scaled-dot-product attention convention. - /// `ln_out` : [B, K, H] - /// `pooled_out` : [B, H] written - /// `attn_out` : [B, H, H] written (saved for backward) - pub fn forward( - &self, - ln_out: &CudaSlice, - n_batch: i32, - k_seq: i32, - inv_scale: f32, - pooled_out: &mut CudaSlice, - attn_out: &mut CudaSlice, - ) -> Result<()> { - // Smem: x_inv[H * K] + mean_k[H] - let smem_bytes = ((IAP_HIDDEN_DIM * k_seq as usize + IAP_HIDDEN_DIM) - * std::mem::size_of::()) as u32; - let cfg = LaunchConfig { - grid_dim: (n_batch as u32, 1, 1), - block_dim: (IAP_BLOCK as u32, 1, 1), - shared_mem_bytes: smem_bytes, - }; - let mut launch = self.stream.launch_builder(&self.fwd_fn); - unsafe { - launch - .arg(ln_out) - .arg(&n_batch).arg(&k_seq).arg(&inv_scale) - .arg(pooled_out).arg(attn_out) - .launch(cfg) - .context("inverted_attention_pool_fwd")?; - } - Ok(()) - } - - /// Backward. `grad_ln_out` accumulates += per-batch (no race). - /// `d_scores_scratch` is a per-step intermediate buffer of shape - /// `[B, H, H]` (must be at least `n_batch * H * H` floats). Trainer - /// allocates once and reuses across steps. - pub fn backward( - &self, - ln_out: &CudaSlice, - attn: &CudaSlice, - grad_pooled: &CudaSlice, - n_batch: i32, - k_seq: i32, - inv_scale: f32, - d_scores_scratch: &mut CudaSlice, - grad_ln_out: &mut CudaSlice, - ) -> Result<()> { - // Smem: x_inv[H * K] + mean_k[H] + dp_shared[H] - let smem_bytes = ((IAP_HIDDEN_DIM * k_seq as usize + 2 * IAP_HIDDEN_DIM) - * std::mem::size_of::()) as u32; - let cfg = LaunchConfig { - grid_dim: (n_batch as u32, 1, 1), - block_dim: (IAP_BLOCK as u32, 1, 1), - shared_mem_bytes: smem_bytes, - }; - let mut launch = self.stream.launch_builder(&self.bwd_fn); - unsafe { - launch - .arg(ln_out).arg(attn).arg(grad_pooled) - .arg(&n_batch).arg(&k_seq).arg(&inv_scale) - .arg(d_scores_scratch) - .arg(grad_ln_out) - .launch(cfg) - .context("inverted_attention_pool_bwd")?; - } - Ok(()) - } -} diff --git a/crates/ml-alpha/src/lib.rs b/crates/ml-alpha/src/lib.rs index 56c2e0b58..6c722def7 100644 --- a/crates/ml-alpha/src/lib.rs +++ b/crates/ml-alpha/src/lib.rs @@ -28,14 +28,8 @@ pub mod cfc; pub mod data; pub mod eval; -pub mod anchor_l2; pub mod heads; -pub mod horizon_mean_collapse; -pub mod horizon_token_attention_pool; -pub mod inv_pooled_merge; -pub mod inverted_attention_pool; pub mod isv; -pub mod regime_moe_gate; pub mod pinned; pub mod pinned_mem; pub mod trainer; diff --git a/crates/ml-alpha/src/regime_moe_gate.rs b/crates/ml-alpha/src/regime_moe_gate.rs deleted file mode 100644 index 9881c4539..000000000 --- a/crates/ml-alpha/src/regime_moe_gate.rs +++ /dev/null @@ -1,174 +0,0 @@ -//! Regime-aware MoE gate host binding (v2 axis D, V6 commit). -//! -//! Top-1 routing per Switch Transformer. Provides three kernel -//! launches: -//! - `forward`: gate top-1 + selected-expert linear → routed_ctx -//! - `backward`: chain-rule grads on the selected expert's weights, -//! bias, and the input fused_ctx. STE on the gate logits is handled -//! by the caller (no separate kernel — gate-logit grad = 0 except -//! via the aux loss path). -//! - `aux_loss`: softmax(gate_logits) + load-balancing aux scalar. -//! -//! See `cuda/regime_moe_gate.cu` for math + the v2 spec §3.3. - -use anyhow::{Context, Result}; -use cudarc::driver::{ - CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg, -}; -use std::sync::Arc; - -pub const MOE_HIDDEN_DIM: usize = 128; -pub const MOE_BLOCK: usize = 128; -pub const MOE_N_HORIZONS: usize = 5; -pub const MOE_N_EXPERTS: usize = 4; - -const CUBIN: &[u8] = - include_bytes!(concat!(env!("OUT_DIR"), "/regime_moe_gate.cubin")); - -pub struct RegimeMoeGate { - _module: Arc, - fwd_fn: CudaFunction, - bwd_fn: CudaFunction, - scatter_fn: CudaFunction, - aux_fn: CudaFunction, - stream: Arc, -} - -impl RegimeMoeGate { - pub fn new(ctx: &Arc, stream: Arc) -> Result { - let module = ctx.load_cubin(CUBIN.to_vec()).context("load regime_moe_gate cubin")?; - let fwd_fn = module.load_function("regime_moe_gate_fwd").context("fwd")?; - let bwd_fn = module.load_function("regime_moe_gate_bwd").context("bwd")?; - let scatter_fn = module.load_function("regime_moe_gate_scatter").context("scatter")?; - let aux_fn = module.load_function("regime_moe_gate_aux").context("aux")?; - Ok(Self { _module: module, fwd_fn, bwd_fn, scatter_fn, aux_fn, stream }) - } - - /// Forward. Selects top-1 expert per batch and applies its linear. - /// `gate_logits` : [B, N_E] - /// `fused_ctx` : [B, N_H, H] - /// `experts_w` : [N_E, H, H] - /// `experts_b` : [N_E, H] - /// `routed_out` : [B, N_H, H] written - /// `top_e_out` : [B] written (i32) - pub fn forward( - &self, - gate_logits: &CudaSlice, - fused_ctx: &CudaSlice, - experts_w: &CudaSlice, - experts_b: &CudaSlice, - n_batch: i32, - routed_out: &mut CudaSlice, - top_e_out: &mut CudaSlice, - ) -> Result<()> { - let cfg = LaunchConfig { - grid_dim: (n_batch as u32, MOE_N_HORIZONS as u32, 1), - block_dim: (MOE_BLOCK as u32, 1, 1), - shared_mem_bytes: 0, - }; - let mut launch = self.stream.launch_builder(&self.fwd_fn); - unsafe { - launch - .arg(gate_logits).arg(fused_ctx).arg(experts_w).arg(experts_b) - .arg(&n_batch) - .arg(routed_out).arg(top_e_out) - .launch(cfg) - .context("regime_moe_gate_fwd")?; - } - Ok(()) - } - - /// Backward (expert path; gate STE handled in aux). All param grad - /// scratches are sparse-by-expert: writes go only to `top_e[b]`'s - /// slot. Reduce over (B, N_H) downstream via reduce_axis0. - /// `d_w_scratch` : [B, N_H, N_E, H, H] - /// `d_b_scratch` : [B, N_H, N_E, H] - /// `d_fused_ctx_out` : [B, N_H, H] - pub fn backward( - &self, - fused_ctx: &CudaSlice, - experts_w: &CudaSlice, - d_routed: &CudaSlice, - top_e: &CudaSlice, - n_batch: i32, - d_w_scratch: &mut CudaSlice, - d_b_scratch: &mut CudaSlice, - d_fused_ctx_out: &mut CudaSlice, - ) -> Result<()> { - let cfg = LaunchConfig { - grid_dim: (n_batch as u32, MOE_N_HORIZONS as u32, 1), - block_dim: (MOE_BLOCK as u32, 1, 1), - shared_mem_bytes: 0, - }; - let mut launch = self.stream.launch_builder(&self.bwd_fn); - unsafe { - launch - .arg(fused_ctx).arg(experts_w).arg(d_routed).arg(top_e) - .arg(&n_batch) - .arg(d_w_scratch).arg(d_b_scratch).arg(d_fused_ctx_out) - .launch(cfg) - .context("regime_moe_gate_bwd")?; - } - Ok(()) - } - - /// Scatter per-batch rank-1 contributions in `d_W_scratch` / - /// `d_b_scratch` into the per-expert grad buffers based on `top_e`. - /// `d_W_scratch` must be `[B, N_H, H, H]` and `d_b_scratch` must - /// be `[B, N_H, H]` (the post-perf compact layout). - pub fn scatter( - &self, - d_w_scratch: &CudaSlice, - d_b_scratch: &CudaSlice, - top_e: &CudaSlice, - n_batch: i32, - grad_experts_w: &mut CudaSlice, - grad_experts_b: &mut CudaSlice, - ) -> Result<()> { - // Grid: (N_E, H, ceil(H/32)). Block: 32. - let chunks: u32 = (MOE_HIDDEN_DIM as u32).div_ceil(32); - let cfg = LaunchConfig { - grid_dim: (MOE_N_EXPERTS as u32, MOE_HIDDEN_DIM as u32, chunks), - block_dim: (32, 1, 1), - shared_mem_bytes: 0, - }; - let mut launch = self.stream.launch_builder(&self.scatter_fn); - unsafe { - launch - .arg(d_w_scratch).arg(d_b_scratch).arg(top_e) - .arg(&n_batch) - .arg(grad_experts_w).arg(grad_experts_b) - .launch(cfg) - .context("regime_moe_gate_scatter")?; - } - Ok(()) - } - - /// Softmax + load-balancing aux loss. - /// `gate_probs_out` : [B, N_E] softmax output - /// `aux_loss_out` : [1] scalar aux loss - pub fn aux_loss( - &self, - gate_logits: &CudaSlice, - top_e: &CudaSlice, - n_batch: i32, - gate_probs_out: &mut CudaSlice, - aux_loss_out: &mut CudaSlice, - ) -> Result<()> { - let cfg = LaunchConfig { - grid_dim: (1, 1, 1), - block_dim: (MOE_N_EXPERTS as u32, 1, 1), - shared_mem_bytes: 0, - }; - let mut launch = self.stream.launch_builder(&self.aux_fn); - unsafe { - launch - .arg(gate_logits).arg(top_e) - .arg(&n_batch) - .arg(gate_probs_out).arg(aux_loss_out) - .launch(cfg) - .context("regime_moe_gate_aux")?; - } - Ok(()) - } -} diff --git a/crates/ml-alpha/src/trainer/anchor_controller.rs b/crates/ml-alpha/src/trainer/anchor_controller.rs deleted file mode 100644 index 38f0ec2a5..000000000 --- a/crates/ml-alpha/src/trainer/anchor_controller.rs +++ /dev/null @@ -1,142 +0,0 @@ -//! Wiener-α anchor-coefficient controller (v2 axis B controller). -//! -//! Drives `λ_anchor` for the L2 anchor regularization. Signal-driven -//! per `pearl_controller_anchors_isv_driven` + the -//! Wiener-α-with-floor pattern (`pearl_wiener_alpha_floor_for_nonstationary`, -//! `pearl_blend_formulas_must_have_permanent_floor`). -//! -//! Math: -//! loss_change_t = val_loss_t − val_loss_{t-1} # raw signal -//! ema_change ← ema_change + α · (loss_change_t − ema_change) -//! target_λ = clamp(|ema_change| · scale, λ_floor, λ_max) -//! λ_t = max(λ_floor, blend(λ_{t-1}, target_λ, α)) -//! -//! where α is the standard Wiener-α (MSE-optimal under stationarity). -//! Floor protects against deadlock when ema_change → 0 -//! (anchor never fully relaxes — minimum guard against runaway drift). -//! -//! Floor `λ_floor` is itself signal-driven from initial parameter -//! magnitudes: `λ_floor = ‖p_init‖ / (100 · sqrt(numel))`. Captured at -//! controller construction via the first-observation bootstrap pattern -//! (`pearl_first_observation_bootstrap`). - -const WIENER_ALPHA_FLOOR: f32 = 0.4; // pearl_wiener_alpha_floor_for_nonstationary -const LAMBDA_MAX: f32 = 1.0; // safety cap on λ -const LAMBDA_SCALE: f32 = 10.0; // ema_change → λ multiplier - -pub struct AnchorController { - /// Floor for λ_anchor. Anchor never fully relaxes below this. - pub lambda_floor: f32, - pub lambda_max: f32, - - /// Wiener-α smoothing state. - diff_var: f32, - sample_var: f32, - ema_change: f32, - prev_loss: f32, - bootstrapped: bool, - - /// Smoothed λ. - pub lambda: f32, -} - -impl AnchorController { - /// `init_param_magnitude` = sum of L2-norms of all anchored param - /// groups; `init_numel` = total trainable element count anchored. - /// Used to derive the signal-driven floor. - pub fn new(init_param_magnitude: f32, init_numel: usize) -> Self { - let lambda_floor = - init_param_magnitude / (100.0 * (init_numel.max(1) as f32).sqrt()); - Self { - lambda_floor, - lambda_max: LAMBDA_MAX, - diff_var: 0.0, - sample_var: 0.0, - ema_change: 0.0, - prev_loss: 0.0, - bootstrapped: false, - lambda: lambda_floor.max(1e-6), // start at floor - } - } - - /// Step the controller with a new validation-loss reading. - /// Returns the updated λ. Capture-safe (host-side; no GPU touch). - pub fn step(&mut self, val_loss: f32) -> f32 { - if !self.bootstrapped { - // First observation: replace state directly per - // pearl_first_observation_bootstrap. - self.prev_loss = val_loss; - self.bootstrapped = true; - return self.lambda; - } - - let change = val_loss - self.prev_loss; - self.prev_loss = val_loss; - - // Wiener-α: α = diff_var / (diff_var + sample_var + ε). - // Maintain running estimates of diff_var (changes) and - // sample_var (noise around the mean change). - let delta = change - self.ema_change; - self.diff_var = 0.95 * self.diff_var + 0.05 * change * change; - self.sample_var = 0.95 * self.sample_var + 0.05 * delta * delta; - let alpha_raw = self.diff_var / (self.diff_var + self.sample_var + 1e-9); - let alpha = alpha_raw.max(WIENER_ALPHA_FLOOR); - - self.ema_change = self.ema_change + alpha * delta; - - // Target λ proportional to |ema_change| × scale. - let target_lambda = (self.ema_change.abs() * LAMBDA_SCALE) - .clamp(self.lambda_floor, self.lambda_max); - - // Blend with floor — pearl_blend_formulas_must_have_permanent_floor. - let blended = self.lambda + alpha * (target_lambda - self.lambda); - self.lambda = blended.max(self.lambda_floor); - self.lambda - } - - pub fn current_lambda(&self) -> f32 { - self.lambda - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn floor_init_from_signal() { - // ‖p_init‖ = 5.0, numel = 100 → floor = 5 / (100 · 10) = 0.005 - let c = AnchorController::new(5.0, 100); - assert!((c.lambda_floor - 0.005).abs() < 1e-6); - assert!(c.lambda >= c.lambda_floor); - } - - #[test] - fn first_observation_bootstrap_returns_floor() { - let mut c = AnchorController::new(1.0, 100); - let l0 = c.step(0.5); - assert!((l0 - c.lambda_floor).abs() < 1e-6); - } - - #[test] - fn lambda_never_falls_below_floor() { - let mut c = AnchorController::new(1.0, 100); - // Many steps with no change → ema_change → 0 → target → floor. - for _ in 0..1000 { - let l = c.step(0.5); - assert!(l >= c.lambda_floor - 1e-9, "λ={l} below floor={}", c.lambda_floor); - } - } - - #[test] - fn lambda_capped_at_max() { - let mut c = AnchorController::new(1.0, 100); - c.step(0.0); - // Huge loss spikes alternating sign — ema_change should grow, - // but λ stays ≤ LAMBDA_MAX. - for i in 0..100 { - let l = c.step(if i % 2 == 0 { 100.0 } else { -100.0 }); - assert!(l <= LAMBDA_MAX + 1e-6, "λ={l} above max"); - } - } -} diff --git a/crates/ml-alpha/src/trainer/mod.rs b/crates/ml-alpha/src/trainer/mod.rs index 87d0f0679..cd7133a49 100644 --- a/crates/ml-alpha/src/trainer/mod.rs +++ b/crates/ml-alpha/src/trainer/mod.rs @@ -1,8 +1,6 @@ //! Trainer module — PerceptionTrainer wraps the full stacked //! Mamba2 -> CfC -> heads pipeline plus 6 AdamW optimizer groups. -pub mod anchor_controller; pub mod loss; -pub mod multi_horizon_attention; pub mod optim; pub mod perception; diff --git a/crates/ml-alpha/src/trainer/multi_horizon_attention.rs b/crates/ml-alpha/src/trainer/multi_horizon_attention.rs deleted file mode 100644 index 490521027..000000000 --- a/crates/ml-alpha/src/trainer/multi_horizon_attention.rs +++ /dev/null @@ -1,435 +0,0 @@ -//! MultiHorizonAttention — single source of truth for the ml-alpha -//! attention path. Replaces the legacy single-Q `attention_pool` -//! entirely. -//! -//! Owns: -//! - horizon_tokens + their init snapshot (for L2 anchor) -//! - shared Q + its init snapshot -//! - horizon-token attention pool kernel binding -//! - inverted (cross-variate) attention pool + saved buffers -//! - regime-MoE gate + expert weights + init snapshots -//! - Kendall σ logarithm (`log_sigma_h`) for the multi-horizon BCE -//! - Anchor L2 kernel + the Wiener-α `AnchorController` -//! -//! All kernel calls go through the bound modules; nothing in this -//! file does host-side compute or sync inside captured regions. - -use anyhow::{Context, Result}; -use cudarc::driver::{CudaSlice, CudaStream}; -use ml_core::device::MlDevice; -use rand::{Rng, SeedableRng}; -use rand_chacha::ChaCha8Rng; -use std::sync::Arc; - -use crate::anchor_l2::AnchorL2; -use crate::horizon_mean_collapse::HorizonMeanCollapse; -use crate::horizon_token_attention_pool::{ - HorizonTokenAttentionPool, HTAP_HIDDEN_DIM, HTAP_N_HORIZONS, -}; -use crate::inv_pooled_merge::InvPooledMerge; -use crate::inverted_attention_pool::InvertedAttentionPool; -use crate::regime_moe_gate::{RegimeMoeGate, MOE_N_EXPERTS}; -use crate::trainer::anchor_controller::AnchorController; -use crate::trainer::optim::AdamW; - -pub const MHA_HIDDEN_DIM: usize = HTAP_HIDDEN_DIM; -pub const MHA_N_HORIZONS: usize = HTAP_N_HORIZONS; -pub const MHA_N_EXPERTS: usize = MOE_N_EXPERTS; -pub const MHA_REGIME_DIM: usize = 12; // spread(5) + vol(3) + tod(4) - -pub struct MultiHorizonAttention { - // ── (C) horizon-token attention pool ── - pub pool: HorizonTokenAttentionPool, - pub horizon_tokens_d: CudaSlice, // [N_H, H] learnable - pub horizon_tokens_init_d: CudaSlice, // [N_H, H] non-trainable snapshot - pub q_d: CudaSlice, // [H] single shared Q - pub q_init_d: CudaSlice, // [H] non-trainable snapshot - pub grad_horizon_tokens_d: CudaSlice, - pub grad_horizon_tokens_scratch_d: CudaSlice, // [B, N_H, H] per-batch scratch - pub grad_ctx_h_d: CudaSlice, // [B, N_H, H] d_ctx_h flowing from MoE bwd → pool bwd input - pub grad_q_d: CudaSlice, - pub grad_q_scratch_d: CudaSlice, // [B, H] per-batch scratch - pub ctx_h_d: CudaSlice, // [B, N_H, H] fwd output - pub attn_h_d: CudaSlice, // [B, N_H + K] - pub opt_horizon_tokens: AdamW, - pub opt_q: AdamW, - - // ── (E) inverted attention pool ── - pub inv_pool: InvertedAttentionPool, - pub inv_pooled_d: CudaSlice, // [B, H] - pub inv_attn_d: CudaSlice, // [B, H, H] - pub inv_d_scores_scratch_d: CudaSlice, // [B, H, H] bwd scratch - pub grad_inv_pooled_d: CudaSlice, // [B, H] true upstream for inv_pool.bwd - // ── Broadcast-add merge: ctx_h[b, h, d] += inv_pooled[b, d] ── - pub merge: InvPooledMerge, - - // ── (C+E fuse) → fused_ctx [B, N_H, H] via additive merge ── - // No separate kernel; computed inline in step_batched as - // fused_ctx[b, h, d] = ctx_h[b, h, d] + inv_pooled[b, d]. - - // ── (D) regime-MoE gate + experts ── - pub moe: RegimeMoeGate, - pub gate_logits_d: CudaSlice, // [B, N_E] computed each step from regime features - pub w_gate_d: CudaSlice, // [N_E, REGIME_DIM] - pub experts_w_d: CudaSlice, // [N_E, H, H] - pub experts_b_d: CudaSlice, // [N_E, H] - pub experts_w_init_d: CudaSlice, // anchor snapshot (B) - pub experts_b_init_d: CudaSlice, // anchor snapshot - pub top_e_d: CudaSlice, // [B] - pub gate_probs_d: CudaSlice, // [B, N_E] - pub aux_loss_d: CudaSlice, // [1] - pub routed_ctx_d: CudaSlice, // [B, N_H, H] v2 output (consumed by GRN heads) - pub grad_routed_d: CudaSlice, // [B, N_H, H] - pub grad_w_gate_d: CudaSlice, - pub grad_experts_w_d: CudaSlice, - pub grad_experts_b_d: CudaSlice, - pub grad_w_scratch_d: CudaSlice, // [B, N_H, H, H] compact (no N_E axis) - pub grad_b_scratch_d: CudaSlice, // [B, N_H, H] compact (no N_E axis) - pub opt_w_gate: AdamW, - pub opt_experts_w: AdamW, - pub opt_experts_b: AdamW, - - // ── (A) Kendall σ on the BCE ── - pub log_sigma_h_d: CudaSlice, // [N_HORIZONS] - pub grad_log_sigma_h_d: CudaSlice, - pub opt_log_sigma: AdamW, - - // ── Mean-collapse over horizon axis (seeds CfC h_old at k=0) ── - pub collapse: HorizonMeanCollapse, - pub ctx_mean_d: CudaSlice, // [B, H] output of collapse fwd, == legacy attn_context_d slot - pub grad_ctx_mean_d: CudaSlice, // [B, H] upstream grad from CfC h_old at k=0 - - // ── (B) anchor regularization controller ── - pub anchor_l2: AnchorL2, - pub anchor: AnchorController, - pub lambda_d: CudaSlice, // [1] current λ uploaded each step - pub anchor_loss_partial_d: CudaSlice, // [1] per-launch partial - pub anchor_loss_total_d: CudaSlice, // [1] accumulated - - // Bookkeeping. - n_batch: usize, - k_seq: usize, - stream: Arc, -} - -impl MultiHorizonAttention { - pub fn new(dev: &MlDevice, n_batch: usize, k_seq: usize, lr: f32, seed: u64) -> Result { - anyhow::ensure!(n_batch > 0 && k_seq > 0, "n_batch and k_seq must be > 0"); - let stream = dev.cuda_stream().context("v2 stream")?.clone(); - let ctx = dev.cuda_context().context("v2 ctx")?; - - let pool = HorizonTokenAttentionPool::new(ctx, stream.clone()) - .context("HorizonTokenAttentionPool")?; - let inv_pool = InvertedAttentionPool::new(ctx, stream.clone()) - .context("InvertedAttentionPool")?; - let moe = RegimeMoeGate::new(ctx, stream.clone()) - .context("RegimeMoeGate")?; - let anchor_l2 = AnchorL2::new(ctx, stream.clone()).context("AnchorL2")?; - let collapse = HorizonMeanCollapse::new(ctx, stream.clone()).context("HorizonMeanCollapse")?; - let merge = InvPooledMerge::new(ctx, stream.clone()).context("InvPooledMerge")?; - - let mut rng = ChaCha8Rng::seed_from_u64(seed); - let scale = (1.0_f32 / MHA_HIDDEN_DIM as f32).sqrt(); - - let init_horizon: Vec = (0..MHA_N_HORIZONS * MHA_HIDDEN_DIM) - .map(|_| rng.gen_range(-scale..scale)).collect(); - let init_q: Vec = (0..MHA_HIDDEN_DIM) - .map(|_| rng.gen_range(-scale..scale)).collect(); - let init_w_gate: Vec = (0..MHA_N_EXPERTS * MHA_REGIME_DIM) - .map(|_| rng.gen_range(-0.1..0.1)).collect(); - let init_experts_w: Vec = (0..MHA_N_EXPERTS * MHA_HIDDEN_DIM * MHA_HIDDEN_DIM) - .map(|_| rng.gen_range(-scale..scale)).collect(); - - let horizon_tokens_d = upload(&stream, &init_horizon)?; - let horizon_tokens_init_d = upload(&stream, &init_horizon)?; - let q_d = upload(&stream, &init_q)?; - let q_init_d = upload(&stream, &init_q)?; - let w_gate_d = upload(&stream, &init_w_gate)?; - let experts_w_d = upload(&stream, &init_experts_w)?; - let experts_b_d = stream.alloc_zeros::(MHA_N_EXPERTS * MHA_HIDDEN_DIM)?; - let experts_w_init_d = upload(&stream, &init_experts_w)?; - let experts_b_init_d = stream.alloc_zeros::(MHA_N_EXPERTS * MHA_HIDDEN_DIM)?; - - let n_ext = MHA_N_HORIZONS + k_seq; - - let s = &stream; - let state = Self { - pool, - horizon_tokens_d, - horizon_tokens_init_d, - q_d, - q_init_d, - grad_horizon_tokens_d: s.alloc_zeros::(MHA_N_HORIZONS * MHA_HIDDEN_DIM)?, - grad_horizon_tokens_scratch_d:s.alloc_zeros::(n_batch * MHA_N_HORIZONS * MHA_HIDDEN_DIM)?, - grad_ctx_h_d: s.alloc_zeros::(n_batch * MHA_N_HORIZONS * MHA_HIDDEN_DIM)?, - grad_q_d: s.alloc_zeros::(MHA_HIDDEN_DIM)?, - grad_q_scratch_d: s.alloc_zeros::(n_batch * MHA_HIDDEN_DIM)?, - ctx_h_d: s.alloc_zeros::(n_batch * MHA_N_HORIZONS * MHA_HIDDEN_DIM)?, - attn_h_d: s.alloc_zeros::(n_batch * n_ext)?, - opt_horizon_tokens: AdamW::new(dev, MHA_N_HORIZONS * MHA_HIDDEN_DIM, lr)?, - opt_q: AdamW::new(dev, MHA_HIDDEN_DIM, lr)?, - - inv_pool, - inv_pooled_d: s.alloc_zeros::(n_batch * MHA_HIDDEN_DIM)?, - inv_attn_d: s.alloc_zeros::(n_batch * MHA_HIDDEN_DIM * MHA_HIDDEN_DIM)?, - inv_d_scores_scratch_d: s.alloc_zeros::(n_batch * MHA_HIDDEN_DIM * MHA_HIDDEN_DIM)?, - grad_inv_pooled_d: s.alloc_zeros::(n_batch * MHA_HIDDEN_DIM)?, - merge, - - moe, - gate_logits_d: s.alloc_zeros::(n_batch * MHA_N_EXPERTS)?, - w_gate_d, - experts_w_d, - experts_b_d, - experts_w_init_d, - experts_b_init_d, - top_e_d: s.alloc_zeros::(n_batch)?, - gate_probs_d: s.alloc_zeros::(n_batch * MHA_N_EXPERTS)?, - aux_loss_d: s.alloc_zeros::(1)?, - routed_ctx_d: s.alloc_zeros::(n_batch * MHA_N_HORIZONS * MHA_HIDDEN_DIM)?, - grad_routed_d: s.alloc_zeros::(n_batch * MHA_N_HORIZONS * MHA_HIDDEN_DIM)?, - grad_w_gate_d: s.alloc_zeros::(MHA_N_EXPERTS * MHA_REGIME_DIM)?, - grad_experts_w_d: s.alloc_zeros::(MHA_N_EXPERTS * MHA_HIDDEN_DIM * MHA_HIDDEN_DIM)?, - grad_experts_b_d: s.alloc_zeros::(MHA_N_EXPERTS * MHA_HIDDEN_DIM)?, - // Compact: [B, N_H, H, H] / [B, N_H, H]. Per-batch - // contribution scatters into top_e's slot of grad_experts via - // `moe.scatter()` instead of writing into a 4× larger - // sparse-by-expert scratch. - grad_w_scratch_d: s.alloc_zeros::(n_batch * MHA_N_HORIZONS * MHA_HIDDEN_DIM * MHA_HIDDEN_DIM)?, - grad_b_scratch_d: s.alloc_zeros::(n_batch * MHA_N_HORIZONS * MHA_HIDDEN_DIM)?, - opt_w_gate: AdamW::new(dev, MHA_N_EXPERTS * MHA_REGIME_DIM, lr)?, - opt_experts_w: AdamW::new(dev, MHA_N_EXPERTS * MHA_HIDDEN_DIM * MHA_HIDDEN_DIM, lr)?, - opt_experts_b: AdamW::new(dev, MHA_N_EXPERTS * MHA_HIDDEN_DIM, lr)?, - - log_sigma_h_d: s.alloc_zeros::(MHA_N_HORIZONS)?, - grad_log_sigma_h_d: s.alloc_zeros::(MHA_N_HORIZONS)?, - opt_log_sigma: AdamW::new(dev, MHA_N_HORIZONS, lr * 0.25)?, // slow update - - collapse, - ctx_mean_d: s.alloc_zeros::(n_batch * MHA_HIDDEN_DIM)?, - grad_ctx_mean_d: s.alloc_zeros::(n_batch * MHA_HIDDEN_DIM)?, - anchor_l2, - anchor: AnchorController::new( - magnitude_of(&init_horizon) + magnitude_of(&init_q), - init_horizon.len() + init_q.len(), - ), - lambda_d: s.alloc_zeros::(1)?, - anchor_loss_partial_d: s.alloc_zeros::(1)?, - anchor_loss_total_d: s.alloc_zeros::(1)?, - - n_batch, - k_seq, - stream: stream.clone(), - }; - Ok(state) - } - - /// Zero per-step grad scratches. Capture-safe (memset_zeros only). - pub fn zero_grads(&mut self) -> Result<()> { - let s = &self.stream; - s.memset_zeros(&mut self.grad_horizon_tokens_d)?; - s.memset_zeros(&mut self.grad_horizon_tokens_scratch_d)?; - s.memset_zeros(&mut self.grad_ctx_h_d)?; - s.memset_zeros(&mut self.grad_q_d)?; - s.memset_zeros(&mut self.grad_q_scratch_d)?; - s.memset_zeros(&mut self.grad_routed_d)?; - s.memset_zeros(&mut self.grad_w_gate_d)?; - s.memset_zeros(&mut self.grad_experts_w_d)?; - s.memset_zeros(&mut self.grad_experts_b_d)?; - s.memset_zeros(&mut self.grad_w_scratch_d)?; - s.memset_zeros(&mut self.grad_b_scratch_d)?; - s.memset_zeros(&mut self.grad_log_sigma_h_d)?; - s.memset_zeros(&mut self.anchor_loss_partial_d)?; - s.memset_zeros(&mut self.anchor_loss_total_d)?; - Ok(()) - } - - pub fn n_batch(&self) -> usize { self.n_batch } - pub fn k_seq(&self) -> usize { self.k_seq } - - /// Forward path: ln_b_out [B, K, H] → ctx_mean_d [B, H] (seeds CfC h_old at k=0). - /// - /// Internal flow: - /// ctx_h = horizon_token_attention_pool(horizon_tokens, q, ln_b_out) - /// inv_pooled = inverted_attention_pool(ln_b_out) - /// fused_ctx[b, h, d] = ctx_h[b, h, d] + inv_pooled[b, d] (additive merge) - /// routed_ctx = regime_moe_gate(gate_logits, fused_ctx, ...) - /// ctx_mean[b, d] = (1/N_H) Σ_h routed_ctx[b, h, d] - /// - /// `gate_logits_d` is filled by the caller before invoking forward - /// (in Stage 2 we use a zero placeholder until V5 wires real - /// regime features; all experts receive equal gate weight then). - /// Capture-safe: no host branches, no synchronize, no host allocs. - pub fn forward(&mut self, ln_b_out: &CudaSlice) -> Result<()> { - let n = self.n_batch as i32; - let k = self.k_seq as i32; - - // 1. Horizon-token attention pool → ctx_h_d, attn_h_d. - self.pool.forward( - &self.horizon_tokens_d, &self.q_d, ln_b_out, - n, k, - &mut self.ctx_h_d, &mut self.attn_h_d, - )?; - - // 2. Inverted attention → inv_pooled_d, inv_attn_d. - let inv_scale = 1.0_f32 / (self.k_seq as f32).sqrt(); - self.inv_pool.forward( - ln_b_out, n, k, inv_scale, - &mut self.inv_pooled_d, &mut self.inv_attn_d, - )?; - - // 3. Real additive merge: ctx_h[b, h, d] += inv_pooled[b, d]. - // This is the wire that connects axis E (inverted attention) - // into the forward chain → MoE → loss. Replaces the prior - // stub. Backward via `merge.backward` produces - // `grad_inv_pooled = Σ_h grad_ctx_h_post` which is the - // correct upstream gradient for `inv_pool.backward`. - self.merge.forward(&mut self.ctx_h_d, &self.inv_pooled_d, n)?; - - // 4. MoE gate + dispatch → routed_ctx_d, top_e_d. - // (Stage 2 placeholder: gate_logits_d remains zeros set in - // zero_grads; all experts equal → top_e = 0 deterministically.) - self.moe.forward( - &self.gate_logits_d, &self.ctx_h_d, - &self.experts_w_d, &self.experts_b_d, - n, - &mut self.routed_ctx_d, &mut self.top_e_d, - )?; - - // 5. Mean-collapse over horizon axis → ctx_mean_d [B, H]. - self.collapse.forward(&self.routed_ctx_d, n, &mut self.ctx_mean_d)?; - Ok(()) - } - - /// Backward path from `grad_ctx_mean [B, H]` (the CfC's grad on - /// h_old at k=0). Accumulates gradients into horizon_tokens, q, - /// experts_w/b, and grad_ln_out (+=, per-batch). - pub fn backward( - &mut self, - ln_b_out: &CudaSlice, - grad_ctx_mean: &CudaSlice, - grad_ln_out: &mut CudaSlice, - ) -> Result<()> { - let n = self.n_batch as i32; - let k = self.k_seq as i32; - - // 5'. Mean-collapse bwd: grad_routed += broadcast(grad_ctx_mean) / N_H. - self.collapse.backward(grad_ctx_mean, n, &mut self.grad_routed_d)?; - - // 4'. MoE bwd: writes per-(b,h) rank-1 contributions into - // compact `grad_w_scratch_d [B, N_H, H, H]` and - // `grad_b_scratch_d [B, N_H, H]`. d_fused_ctx → d_ctx_h. - self.moe.backward( - &self.ctx_h_d, &self.experts_w_d, &self.grad_routed_d, &self.top_e_d, - n, - &mut self.grad_w_scratch_d, - &mut self.grad_b_scratch_d, - &mut self.grad_ctx_h_d, - )?; - - // 4''. Scatter the per-(b,h) contributions into the per-expert - // shared grad slots based on top_e[b]. Replaces the prior - // naive reduce_axis0 over a sparse-by-expert scratch. - self.moe.scatter( - &self.grad_w_scratch_d, &self.grad_b_scratch_d, &self.top_e_d, - n, - &mut self.grad_experts_w_d, - &mut self.grad_experts_b_d, - )?; - - // 3'. Merge backward: compute grad_inv_pooled[b, d] = - // Σ_h grad_ctx_h[b, h, d]. The additive merge has - // passthrough gradient on ctx_h, so grad_ctx_h_d itself - // is unchanged and feeds straight into pool.bwd below. - self.merge.backward(&self.grad_ctx_h_d, n, &mut self.grad_inv_pooled_d)?; - - // 2'. Inverted-attention bwd consumes the REAL upstream - // gradient `grad_inv_pooled` (not `grad_ctx_mean` — - // that was the prior stub-architecture fake). - let inv_scale = 1.0_f32 / (self.k_seq as f32).sqrt(); - self.inv_pool.backward( - ln_b_out, &self.inv_attn_d, &self.grad_inv_pooled_d, - n, k, inv_scale, - &mut self.inv_d_scores_scratch_d, - grad_ln_out, - )?; - - // 1'. Horizon-token attention pool bwd: grad_ctx_h_d → - // grad_horizon_tokens_scratch_d, grad_q_scratch_d, and - // += grad_ln_out per-batch. - self.pool.backward( - &self.horizon_tokens_d, &self.q_d, ln_b_out, - &self.attn_h_d, - &self.grad_ctx_h_d, - n, k, - &mut self.grad_horizon_tokens_scratch_d, - &mut self.grad_q_scratch_d, - grad_ln_out, - )?; - Ok(()) - } - - /// Apply L2 anchor regularization to all anchored param groups. - /// Adds to `grad_horizon_tokens_d`, `grad_q_d`, `grad_experts_w_d`. - /// `lambda_d` must hold the current λ (written by the caller on - /// host, uploaded before launch — done outside the capture region). - pub fn apply_anchor(&mut self) -> Result<()> { - let n_ht = (MHA_N_HORIZONS * MHA_HIDDEN_DIM) as i32; - let n_q = MHA_HIDDEN_DIM as i32; - let n_ew = (MHA_N_EXPERTS * MHA_HIDDEN_DIM * MHA_HIDDEN_DIM) as i32; - - self.anchor_l2.apply( - &self.horizon_tokens_d, &self.horizon_tokens_init_d, - &self.lambda_d, n_ht, - &mut self.anchor_loss_partial_d, &mut self.grad_horizon_tokens_d, - )?; - self.anchor_l2.apply( - &self.q_d, &self.q_init_d, - &self.lambda_d, n_q, - &mut self.anchor_loss_partial_d, &mut self.grad_q_d, - )?; - self.anchor_l2.apply( - &self.experts_w_d, &self.experts_w_init_d, - &self.lambda_d, n_ew, - &mut self.anchor_loss_partial_d, &mut self.grad_experts_w_d, - )?; - Ok(()) - } - - /// Step all owned optimizers. Capture-safe — each AdamW is its own - /// device kernel; the trainer calls this after reducing scratches. - pub fn adamw_step(&mut self) -> Result<()> { - self.opt_horizon_tokens.step(&mut self.horizon_tokens_d, &self.grad_horizon_tokens_d)?; - self.opt_q.step(&mut self.q_d, &self.grad_q_d)?; - self.opt_w_gate.step(&mut self.w_gate_d, &self.grad_w_gate_d)?; - self.opt_experts_w.step(&mut self.experts_w_d, &self.grad_experts_w_d)?; - self.opt_experts_b.step(&mut self.experts_b_d, &self.grad_experts_b_d)?; - self.opt_log_sigma.step(&mut self.log_sigma_h_d, &self.grad_log_sigma_h_d)?; - Ok(()) - } -} - - -fn upload(stream: &Arc, host: &[f32]) -> Result> { - use crate::pinned_mem::MappedF32Buffer; - use cudarc::driver::{DevicePtrMut}; - let n = host.len(); - let staging = unsafe { MappedF32Buffer::new(n) } - .map_err(|e| anyhow::anyhow!("v2_state upload staging: {e}"))?; - staging.write_from_slice(host); - let mut dst = stream.alloc_zeros::(n).context("v2_state upload alloc")?; - if n > 0 { - let nbytes = n * std::mem::size_of::(); - unsafe { - let (dst_ptr, _g) = dst.device_ptr_mut(stream); - cudarc::driver::result::memcpy_dtod_async( - dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream(), - ).context("v2_state upload DtoD")?; - } - } - Ok(dst) -} - -fn magnitude_of(p: &[f32]) -> f32 { - p.iter().map(|x| x * x).sum::().sqrt() -} diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index 0ee268ce6..64b3094c9 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -61,6 +61,7 @@ const BCE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/bce_loss_mult const HORIZON_LAMBDA_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/horizon_lambda.cubin")); const LAYER_NORM_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/layer_norm.cubin")); const VARIABLE_SELECTION_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/variable_selection.cubin")); +const ATTENTION_POOL_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/attention_pool.cubin")); const REDUCE_AXIS0_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/reduce_axis0.cubin")); #[derive(Clone, Debug)] @@ -359,15 +360,30 @@ pub struct PerceptionTrainer { vsn_bwd_fn: CudaFunction, _vsn_module: Arc, - // ── Attention path ── - // MultiHorizonAttention is the single source of truth for the - // attention summary that seeds CfC's h_old at k=0. Owns the - // horizon-token attention pool, the inverted (cross-variate) - // attention pass, the regime-MoE gate + experts, the mean-collapse - // over the horizon axis, the Kendall σ logarithm fed into BCE, and - // the L2 anchor controller + kernel. Replaces the prior single-Q - // `attention_pool` entirely. - pub mha: crate::trainer::multi_horizon_attention::MultiHorizonAttention, + // ── Attention pool (Phase 3) ── + // Single-head attention pool over LN_b output. Replaces CfC's + // zero-initialised h_old at k=0 with a learned content-addressable + // summary over all K positions. Single learned param: Q [HIDDEN_DIM]. + pub attn_q_d: CudaSlice, + /// Per-sample pooled context `[B, HIDDEN_DIM]` — fed as h_old at k=0. + attn_context_d: CudaSlice, + /// Saved post-softmax attention weights `[B, K]` for bwd. + attn_weights_d: CudaSlice, + grad_attn_q_d: CudaSlice, + pub opt_attn_q: AdamW, + attn_fwd_fn: CudaFunction, + attn_bwd_fn: CudaFunction, + _attn_module: Arc, + + // ── Kendall σ-weighted BCE (axis A) ── + // Per-horizon learnable `log σ_h` scalars. Weighted BCE per horizon + // is `(1/(2·exp(2·log σ_h))) · mean_bce_h + log σ_h`. AdamW updates + // log_sigma_h at 1/4 the head LR (slow update). σ on its own without + // C/D/E is the keeper from the v2 sweep (see + // project_ml_alpha_v2_ab_verdict.md). + pub log_sigma_h_d: CudaSlice, // [N_HORIZONS] + pub grad_log_sigma_h_d: CudaSlice, // [N_HORIZONS] + pub opt_log_sigma: AdamW, // ── K-loop parallelization (Phase B) ── // Per-batch grad scratch buffers for cfc_step_backward_batched. @@ -393,6 +409,8 @@ pub struct PerceptionTrainer { // VSN per-row grad scratch (Phase B commit 3). n_rows = B * K. vsn_grad_w_scratch_d: CudaSlice, // [B*K, FEATURE_DIM, FEATURE_DIM] vsn_grad_b_scratch_d: CudaSlice, // [B*K, FEATURE_DIM] + // Attention pool per-batch grad scratch (Phase B commit 4). + attn_grad_q_scratch_d: CudaSlice, // [B, HIDDEN_DIM] /// Cross-batch reducer kernel: `[B, N] → [N]` via block tree-reduce. /// Used for every per-batch grad scratch in the refactored bwd path. reduce_axis0_fn: CudaFunction, @@ -523,6 +541,15 @@ impl PerceptionTrainer { let vsn_bwd_fn = vsn_module .load_function("variable_selection_bwd") .context("variable_selection_bwd symbol")?; + let attn_module = ctx + .load_cubin(ATTENTION_POOL_CUBIN.to_vec()) + .context("attention_pool cubin")?; + let attn_fwd_fn = attn_module + .load_function("attention_pool_fwd") + .context("attention_pool_fwd symbol")?; + let attn_bwd_fn = attn_module + .load_function("attention_pool_bwd") + .context("attention_pool_bwd symbol")?; let reduce_module = ctx .load_cubin(REDUCE_AXIS0_CUBIN.to_vec()) .context("reduce_axis0 cubin")?; @@ -772,20 +799,35 @@ impl PerceptionTrainer { cfg.n_batch * cfg.seq_len * FEATURE_DIM)?; // ── Attention pool init (Phase 3) ── - // MultiHorizonAttention bundle. Single source of truth for the - // multi-horizon attention path + Kendall σ on the BCE. Init - // captures snapshots used by the L2 anchor controller. - // Stage 1 (this commit) wires `mha.log_sigma_h_d` and - // `mha.grad_log_sigma_h_d` into the BCE launch. Stage 2 - // replaces the legacy `attention_pool` forward/backward - // callsites with `mha.pool` + `mha.inv_pool` + `mha.moe`. - let mha = crate::trainer::multi_horizon_attention::MultiHorizonAttention::new( - dev, - cfg.n_batch, - cfg.seq_len, - cfg.lr_cfc, - cfg.seed.wrapping_add(0xB200_C00B), - )?; + // Q init near zero so initial scores ≈ 0 → softmax ≈ uniform 1/K + // → context ≈ mean of LN_b output. Model learns content + // addressing from a near-uniform starting point. + let attn_q_scale = (1.0_f32 / HIDDEN_DIM as f32).sqrt(); + let attn_q_init: Vec = (0..HIDDEN_DIM) + .map(|_| r.gen_range(-attn_q_scale..attn_q_scale)).collect(); + let attn_q_d = upload(&stream, &attn_q_init)?; + let attn_context_d = stream.alloc_zeros::(cfg.n_batch * HIDDEN_DIM)?; + let attn_weights_d = stream.alloc_zeros::(cfg.n_batch * cfg.seq_len)?; + let grad_attn_q_d = stream.alloc_zeros::(HIDDEN_DIM)?; + let opt_attn_q = AdamW::new(dev, HIDDEN_DIM, cfg.lr_cfc)?; + // Phase B: attn pool per-batch grad scratch. + let attn_grad_q_scratch_d = stream.alloc_zeros::(cfg.n_batch * HIDDEN_DIM)?; + + // Kendall σ — per-horizon learnable `log σ_h` (axis A only; + // C/D/E reverted after the 2026-05-18 A/B sweep). Init at 0 + // → weight = 1/(2·exp(0)) = 0.5 per horizon at step 0; AdamW + // updates at 1/4 the head LR per the spec. + let log_sigma_h_d = stream.alloc_zeros::(N_HORIZONS)?; + let grad_log_sigma_h_d = stream.alloc_zeros::(N_HORIZONS)?; + let opt_log_sigma = AdamW::new(dev, N_HORIZONS, cfg.lr_cfc * 0.25)?; + + // V1 (2026-05-18): per-horizon Q_h trainer state (C24/C25) removed. + // A/B sweep at commit 83546b5c3 falsified the approach (mean_auc + // -0.019 vs baseline; h6000 essentially tied). v2 design with + // horizon-token K-prepend + inverted attention + regime-MoE + + // Kendall σ-BCE + L2 anchor replaces it. See + // docs/superpowers/specs/2026-05-18-ml-alpha-v2-multi-horizon-design.md + // and the V1-V13 plan. let k = cfg.seq_len; Ok(Self { @@ -834,7 +876,17 @@ impl PerceptionTrainer { vsn_bwd_fn, _vsn_module: vsn_module, // Attention pool (Phase 3). - mha, + attn_q_d, + attn_context_d, + attn_weights_d, + grad_attn_q_d, + opt_attn_q, + attn_fwd_fn, + attn_bwd_fn, + _attn_module: attn_module, + log_sigma_h_d, + grad_log_sigma_h_d, + opt_log_sigma, // Phase B: cfc per-batch grad scratch + reducer. cfc_grad_w_in_scratch_d, cfc_grad_w_rec_scratch_d, @@ -852,6 +904,7 @@ impl PerceptionTrainer { grn_grad_b_skip_scratch_d, vsn_grad_w_scratch_d, vsn_grad_b_scratch_d, + attn_grad_q_scratch_d, reduce_axis0_fn, _reduce_module: reduce_module, loss_ema_d: stream.alloc_zeros::(N_HORIZONS)?, @@ -1397,16 +1450,30 @@ impl PerceptionTrainer { unsafe { launch.launch(cfg_tx).context("transpose h_enriched fwd")?; } } - // ── 2d. Multi-horizon attention forward — single source of - // truth for the attention path. Replaces the legacy - // single-Q attention_pool. Internal flow: - // horizon-token attn → ctx_h [B, N_H, H] - // inverted attn → inv_pooled [B, H] - // MoE gate + dispatch → routed_ctx [B, N_H, H] - // mean-collapse → mha.ctx_mean_d [B, H] - // The CfC's h_old at k=0 reads from mha.ctx_mean_d - // (legacy attn_context_d slot). - self.mha.forward(&self.ln_out_d)?; + // ── 2d. Attention pool forward (Phase 3) — produces + // attn_context_d [B, HIDDEN_DIM] = learned content-summary + // over all K LN_b output positions. Replaces CfC's + // zero-initialised h_old at k=0 with this context vector. + // Shared mem: K floats (scores) + BLOCK floats (reduce) + + // HIDDEN_DIM floats (context) = (K + 128 + 128) * 4 bytes. + { + let k_i32 = k_seq as i32; + let n_batch_attn = b_sz as i32; + let shared = (k_seq + 128 + HIDDEN_DIM) * std::mem::size_of::(); + let cfg_attn = LaunchConfig { + grid_dim: (b_sz as u32, 1, 1), + block_dim: (128, 1, 1), // ATTN_BLOCK + shared_mem_bytes: shared as u32, + }; + let mut launch = self.stream.launch_builder(&self.attn_fwd_fn); + launch + .arg(&self.attn_q_d) + .arg(&self.ln_out_d) + .arg(&n_batch_attn).arg(&k_i32) + .arg(&mut self.attn_context_d) + .arg(&mut self.attn_weights_d); + unsafe { launch.launch(cfg_attn).context("attention_pool_fwd")?; } + } // ── 3. Labels DtoD: staging (filled in step_batched before // captured region) → device. @@ -1431,10 +1498,9 @@ impl PerceptionTrainer { .map_err(|e| anyhow::anyhow!("zero cfc_grad_b_scratch: {e}"))?; self.stream.memset_zeros(&mut self.cfc_grad_tau_scratch_d) .map_err(|e| anyhow::anyhow!("zero cfc_grad_tau_scratch: {e}"))?; - // MHA grad scratches (horizon-token, q, expert-W/b, anchor - // loss partials, log-sigma). Capture-safe memset_zeros. - self.mha.zero_grads()?; - + // Phase B commit 4: attention pool per-batch grad_Q scratch. + self.stream.memset_zeros(&mut self.attn_grad_q_scratch_d) + .map_err(|e| anyhow::anyhow!("zero attn_grad_q_scratch: {e}"))?; // GRN per-batch grad scratch (Phase B commit 2): zero ONCE per // step; K-loop bwd accumulates into these, then reduce_axis0 // collapses → final grad buffers (OVERWRITE) after the K-loop. @@ -1535,16 +1601,16 @@ impl PerceptionTrainer { // [K, B, H] layout. Pointer-offset trick (single mut // borrow + raw u64 arithmetic for slot pointers) keeps // the launches stream-ordered with no syncs. - // MHA produces mha.ctx_mean_d which seeds h_old at k=0. The - // legacy attn_context_d slot is removed. zero_h_ptr remains for - // the bwd k=0 case (cfc bwd needs an h_old slot for the input - // gradient to read from). + // Phase 3: attention pool produces attn_context_d which replaces + // the zero initial h_old at k=0. zero_h_ptr is no longer used in + // the fwd K-loop but is retained for the bwd k=0 case (cfc bwd + // needs an h_old slot for the input gradient to read from). let _zero_h_ptr_unused = { let (p, _g) = self.zero_h_d.device_ptr(&self.stream); p }; let attn_context_ptr = { - let (p, _g) = self.mha.ctx_mean_d.device_ptr(&self.stream); + let (p, _g) = self.attn_context_d.device_ptr(&self.stream); p }; let henr_t_base = { @@ -1606,12 +1672,12 @@ impl PerceptionTrainer { } drop((_g_hpk, _g_probs, _g_z1, _g_a1, _g_z2, _g_gate, _g_main, _g_logit)); - // ── 4.5. (v2 reserved) — per-horizon Q_h path removed at V1. - // The v2 design (horizon-token K-prepend + inverted attn + - // regime-MoE) wires in here in commits V9/V10. + // ── 4.5. Zero σ grad scratch each step (capture-safe). + self.stream.memset_zeros(&mut self.grad_log_sigma_h_d) + .map_err(|e| anyhow::anyhow!("zero grad_log_sigma_h: {e}"))?; - // ── 5. Fused multi-horizon BCE over the full [K*B, N_HORIZONS] - // grid. The kernel doesn't distinguish position-vs-batch + // ── 5. Fused multi-horizon BCE (Kendall σ-weighted, axis A). + // The kernel doesn't distinguish position-vs-batch // since the per-horizon weight is identified by `i % n_horizons`. let n_pos_i = (k_seq * b_sz) as i32; let n_h_i = N_HORIZONS as i32; @@ -1625,16 +1691,19 @@ impl PerceptionTrainer { launch .arg(&self.probs_per_k_d).arg(&self.labels_per_k_d) .arg(&self.loss_weights_d) - .arg(&self.mha.log_sigma_h_d) + .arg(&self.log_sigma_h_d) .arg(&n_pos_i).arg(&n_h_i) .arg(&mut self.loss_d) .arg(&mut self.loss_per_horizon_d) .arg(&mut self.grad_probs_per_k_d) .arg(&mut self.valid_d) - .arg(&mut self.mha.grad_log_sigma_h_d); + .arg(&mut self.grad_log_sigma_h_d); unsafe { launch.launch(bce_cfg).context("bce launch")?; } } + // ── 5a. (v2 reserved) — per-horizon Q_h backward removed at V1. + // v2 backward path will live here in commit V10. + // ── 5b. ISV-driven per-horizon EMA + lambda. Updates the EMA // of unweighted per-horizon BCE and emits a clamped // multiplier `lambda_d[h]` per @@ -1661,7 +1730,7 @@ impl PerceptionTrainer { self.stream.memset_zeros(&mut self.grad_h_carry_d) .map_err(|e| anyhow::anyhow!("zero grad_h_carry: {e}"))?; - // MHA: at k=0 the bwd kernel reads `h_old` = mha.ctx_mean_d + // Phase 3: at k=0 the bwd kernel reads `h_old` = attn_context_d // (mirroring the forward pass). zero_h_ptr_bwd retained as a // legacy fallback / unused alias. let _zero_h_ptr_bwd_unused = { @@ -1669,7 +1738,7 @@ impl PerceptionTrainer { p }; let attn_context_ptr_bwd = { - let (p, _g) = self.mha.ctx_mean_d.device_ptr(&self.stream); + let (p, _g) = self.attn_context_d.device_ptr(&self.stream); p }; let henr_t_base_bwd = { @@ -1784,53 +1853,60 @@ impl PerceptionTrainer { unsafe { launch.launch(cfg_tx).context("transpose grad bwd")?; } } - // ── 7c-pre. MultiHorizonAttention backward. Consumes: - // grad_h_carry_d (= grad on initial h_old at k=0, which IS - // mha.ctx_mean_d in fwd) → fed as grad_ctx_mean. - // Accumulates into: - // grad_horizon_tokens_scratch_d, grad_q_scratch_d (per-batch) - // grad_w_scratch_d, grad_b_scratch_d (per-batch, sparse-by-expert) - // grad_h_enriched_seq_d (LN_b out grad) += MHA's attention-path - // contribution. Same += semantics as the legacy attn_bwd. - let grad_h_enriched_slice = self.grad_h_enriched_seq_d.data_mut(); - self.mha.backward(&self.ln_out_d, &self.grad_h_carry_d, grad_h_enriched_slice)?; - - // Reduce MHA's per-batch grad scratches → shared grads. - // grad_horizon_tokens_scratch_d [B, N_H, H] → reduce_axis0 → [N_H, H] - // grad_q_scratch_d [B, H] → reduce_axis0 → [H] - // grad_w_scratch_d / grad_b_scratch_d are already scattered into - // grad_experts_w_d / grad_experts_b_d by `mha.moe.scatter()` inside - // `mha.backward()` (compact layout — see regime_moe_gate.cu). - let n_batch_i = b_sz as i32; - let cfg_red = |n_tail: i32| LaunchConfig { - grid_dim: (n_tail as u32, 1, 1), - block_dim: (256, 1, 1), - shared_mem_bytes: 0, - }; - let n_ht_i = (5 * HIDDEN_DIM) as i32; - let n_q_i = HIDDEN_DIM as i32; + // ── 7c-pre. Attention pool backward (Phase 3). Consumes: + // Q = self.attn_q_d [HIDDEN_DIM] + // ln_out (values)= self.ln_out_d [B, K, HIDDEN_DIM] + // attn_weights = self.attn_weights_d (saved by fwd) [B, K] + // grad_context = self.grad_h_carry_d (= grad on initial + // h_old at k=0, which IS attn_context) [B, HIDDEN_DIM] + // Writes (BOTH ARE +=): + // grad_attn_q_d += attn-path contribution to Q + // grad_h_enriched_seq_d (LN_b output grad) += attn-path + // contribution to ln_out + // The pre-zero of grad_attn_q_d at step start makes the += a + // clean overwrite for the Q grad. grad_h_enriched_seq_d already + // holds the K-loop's contribution at this point — attn's + // contribution adds on top. + // Phase B commit 4: block-per-batch attn bwd writes per-batch + // grad_Q scratch; reducer collapses → final grad_attn_q_d below. { - let mut launch = self.stream.launch_builder(&self.reduce_axis0_fn); - launch.arg(&self.mha.grad_horizon_tokens_scratch_d) - .arg(&n_batch_i).arg(&n_ht_i) - .arg(&mut self.mha.grad_horizon_tokens_d); - unsafe { launch.launch(cfg_red(n_ht_i)).context("reduce grad_horizon_tokens")?; } + let k_i32 = k_seq as i32; + let n_batch_attn = b_sz as i32; + let shared = (2 * k_seq + 128) * std::mem::size_of::(); + let cfg_attn_bwd = LaunchConfig { + grid_dim: (b_sz as u32, 1, 1), + block_dim: (128, 1, 1), // ATTN_BLOCK + shared_mem_bytes: shared as u32, + }; + let mut launch = self.stream.launch_builder(&self.attn_bwd_fn); + launch + .arg(&self.attn_q_d) + .arg(&self.ln_out_d) + .arg(&self.attn_weights_d) + .arg(&self.grad_h_carry_d) + .arg(&n_batch_attn).arg(&k_i32) + .arg(&mut self.attn_grad_q_scratch_d) + .arg(self.grad_h_enriched_seq_d.data_mut()); + unsafe { launch.launch(cfg_attn_bwd).context("attention_pool_bwd")?; } } + // Attn pool reducer: collapse [B, HIDDEN_DIM] → [HIDDEN_DIM]. { + let n_batch_i = b_sz as i32; + let n_tail_i = HIDDEN_DIM as i32; + let cfg_red = LaunchConfig { + grid_dim: (HIDDEN_DIM as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; let mut launch = self.stream.launch_builder(&self.reduce_axis0_fn); - launch.arg(&self.mha.grad_q_scratch_d) - .arg(&n_batch_i).arg(&n_q_i) - .arg(&mut self.mha.grad_q_d); - unsafe { launch.launch(cfg_red(n_q_i)).context("reduce grad_q")?; } + launch + .arg(&self.attn_grad_q_scratch_d) + .arg(&n_batch_i) + .arg(&n_tail_i) + .arg(&mut self.grad_attn_q_d); + unsafe { launch.launch(cfg_red).context("reduce attn_grad_q")?; } } - // Anchor L2 launches — apply on horizon_tokens, Q, experts_w. - // λ_d is written on host (outside captured region) before - // step_batched. anchor_loss_total is accumulated into one - // scalar (not currently read by the trainer; reserved for - // logging). - self.mha.apply_anchor()?; - // ── 7c. LayerNorm B backward (between m2 and CfC). Consumes: // x = m2.h_enriched_seq [B, K, H] // gain = self.ln_gain_d [H] @@ -2139,7 +2215,9 @@ impl PerceptionTrainer { self.opt_ln_a_bias.step(&mut self.ln_a_bias_d, &self.grad_ln_a_bias_d)?; self.opt_vsn_w.step(&mut self.vsn_w_d, &self.grad_vsn_w_d)?; self.opt_vsn_b.step(&mut self.vsn_b_d, &self.grad_vsn_b_d)?; - self.mha.adamw_step()?; + self.opt_attn_q.step(&mut self.attn_q_d, &self.grad_attn_q_d)?; + // Kendall σ step (axis A; 1/4 head LR). + self.opt_log_sigma.step(&mut self.log_sigma_h_d, &self.grad_log_sigma_h_d)?; // (v2 reserved) — per-horizon AdamW step removed at V1; v2's six // optimizer groups (horizon_tokens, Q_inv, w_fuse + b_fuse, @@ -2394,9 +2472,26 @@ impl PerceptionTrainer { unsafe { launch.launch(cfg_tx).context("eval transpose h_enriched")?; } } - // MultiHorizonAttention fwd — same as training. Produces - // mha.ctx_mean_d which seeds h_old at k=0 in the eval K-loop. - self.mha.forward(&self.ln_out_d)?; + // Attention pool fwd — same as training, populates attn_context_d + // for use as k=0 h_old in the eval K-loop below. + { + let k_i32 = k_seq as i32; + let n_batch_attn = b_sz as i32; + let shared = (k_seq + 128 + HIDDEN_DIM) * std::mem::size_of::(); + let cfg_attn = LaunchConfig { + grid_dim: (b_sz as u32, 1, 1), + block_dim: (128, 1, 1), + shared_mem_bytes: shared as u32, + }; + let mut launch = self.stream.launch_builder(&self.attn_fwd_fn); + launch + .arg(&self.attn_q_d) + .arg(&self.ln_out_d) + .arg(&n_batch_attn).arg(&k_i32) + .arg(&mut self.attn_context_d) + .arg(&mut self.attn_weights_d); + unsafe { launch.launch(cfg_attn).context("eval attention_pool_fwd")?; } + } // Upload labels [K, B, N_HORIZONS]. let total_labels = k_seq * b_sz * N_HORIZONS; @@ -2448,9 +2543,9 @@ impl PerceptionTrainer { let (p, _g) = self.zero_h_d.device_ptr(&self.stream); p }; - // Eval uses mha.ctx_mean_d as h_old at k=0, mirroring training. + // Phase 3: eval uses attn_context_d as h_old at k=0, mirroring training. let attn_context_ptr = { - let (p, _g) = self.mha.ctx_mean_d.device_ptr(&self.stream); + let (p, _g) = self.attn_context_d.device_ptr(&self.stream); p }; let henr_t_base = { @@ -2516,12 +2611,12 @@ impl PerceptionTrainer { let mut launch = self.stream.launch_builder(&self.bce_fn); launch .arg(&self.probs_per_k_d).arg(&self.labels_per_k_d).arg(&self.loss_weights_d) - .arg(&self.mha.log_sigma_h_d) + .arg(&self.log_sigma_h_d) .arg(&n_pos_i).arg(&n_h_i) .arg(&mut self.loss_d) .arg(&mut self.loss_per_horizon_d) .arg(&mut self.grad_probs_per_k_d).arg(&mut self.valid_d) - .arg(&mut self.mha.grad_log_sigma_h_d); + .arg(&mut self.grad_log_sigma_h_d); launch.launch(bce_cfg).context("eval bce")?; } self.stream.synchronize().context("eval sync")?; diff --git a/crates/ml-alpha/tests/anchor_l2_numgrad.rs b/crates/ml-alpha/tests/anchor_l2_numgrad.rs deleted file mode 100644 index 1118cebf1..000000000 --- a/crates/ml-alpha/tests/anchor_l2_numgrad.rs +++ /dev/null @@ -1,138 +0,0 @@ -//! Numgrad parity test for anchor_l2 (v2 axis B). - -#![cfg(feature = "cuda")] - -use anyhow::{Context, Result}; -use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut}; -use ml_alpha::anchor_l2::AnchorL2; -use ml_alpha::pinned_mem::MappedF32Buffer; -use ml_core::device::MlDevice; -use std::sync::Arc; - -const N: usize = 67; // intentionally non-power-of-two to stress stride loops - -fn rng(seed: u64) -> impl FnMut() -> f32 { - let mut s = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15); - move || -> f32 { - s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407); - let x = ((s >> 16) & 0xFFFF) as f32 / 65536.0; - x * 2.0 - 1.0 - } -} - -fn upload(stream: &Arc, host: &[f32]) -> Result> { - let n = host.len(); - let staging = unsafe { MappedF32Buffer::new(n) } - .map_err(|e| anyhow::anyhow!("upload: {e}"))?; - staging.write_from_slice(host); - let mut dst = stream.alloc_zeros::(n)?; - let nbytes = n * std::mem::size_of::(); - unsafe { - let (dst_ptr, _g) = dst.device_ptr_mut(stream); - cudarc::driver::result::memcpy_dtod_async( - dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream(), - )?; - } - Ok(dst) -} - -fn download(stream: &Arc, src: &CudaSlice) -> Result> { - let n = src.len(); - let staging = unsafe { MappedF32Buffer::new(n) } - .map_err(|e| anyhow::anyhow!("download: {e}"))?; - let nbytes = n * std::mem::size_of::(); - unsafe { - let (src_ptr, _g) = src.device_ptr(stream); - cudarc::driver::result::memcpy_dtod_async( - staging.dev_ptr, src_ptr, nbytes, stream.cu_stream(), - )?; - } - stream.synchronize()?; - Ok(staging.read_all()) -} - -fn forward_loss( - anchor: &AnchorL2, - stream: &Arc, - p: &[f32], - p_init: &[f32], - lambda: f32, -) -> Result { - let p_d = upload(stream, p)?; - let p_init_d = upload(stream, p_init)?; - let lambda_d = upload(stream, &[lambda])?; - let mut loss_d = stream.alloc_zeros::(1)?; - let mut grad_d = stream.alloc_zeros::(p.len())?; - anchor.apply(&p_d, &p_init_d, &lambda_d, p.len() as i32, &mut loss_d, &mut grad_d)?; - stream.synchronize()?; - Ok(download(stream, &loss_d)?[0]) -} - -#[test] -#[ignore = "requires CUDA"] -fn forward_and_grad_match_central_difference() -> Result<()> { - let dev = MlDevice::cuda(0)?; - let stream = dev.cuda_stream().context("stream")?.clone(); - let ctx_dev = dev.cuda_context().context("ctx")?; - let anchor = AnchorL2::new(ctx_dev, stream.clone())?; - - let mut r = rng(20260518); - let p: Vec = (0..N).map(|_| r() * 0.5).collect(); - let p_init: Vec = (0..N).map(|_| r() * 0.5).collect(); - let lambda = 0.137_f32; - - // 1. Forward parity: kernel loss == closed-form. - let p_d = upload(&stream, &p)?; - let p_init_d = upload(&stream, &p_init)?; - let lambda_d = upload(&stream, &[lambda])?; - let mut loss_d = stream.alloc_zeros::(1)?; - let mut grad_d = stream.alloc_zeros::(N)?; - anchor.apply(&p_d, &p_init_d, &lambda_d, N as i32, &mut loss_d, &mut grad_d)?; - stream.synchronize()?; - let kernel_loss = download(&stream, &loss_d)?[0]; - let kernel_grad = download(&stream, &grad_d)?; - - let mut expected_loss = 0.0_f32; - for i in 0..N { - let d = p[i] - p_init[i]; - expected_loss += d * d; - } - expected_loss *= lambda; - let abs = (kernel_loss - expected_loss).abs(); - let rel = abs / expected_loss.abs().max(1e-3); - assert!( - abs < 1e-3 || rel < 5e-3, - "loss closed-form mismatch: kernel={kernel_loss:.5} expected={expected_loss:.5}" - ); - - // 2. Closed-form grad parity (analytic = 2λ (p − p_init)). - for i in 0..N { - let expected_grad = 2.0 * lambda * (p[i] - p_init[i]); - let abs = (kernel_grad[i] - expected_grad).abs(); - assert!(abs < 1e-4, - "grad[{i}] closed-form mismatch: kernel={} expected={} (diff={abs:.3e})", - kernel_grad[i], expected_grad); - } - - // 3. Central-difference numgrad — 4 random positions. - let eps = 1e-3_f32; - let mut r2 = rng(42); - for _ in 0..4 { - let idx = ((r2() + 1.0) * 0.5 * N as f32) as usize % N; - let mut p_p = p.clone(); - let mut p_m = p.clone(); - p_p[idx] += eps; - p_m[idx] -= eps; - let l_p = forward_loss(&anchor, &stream, &p_p, &p_init, lambda)?; - let l_m = forward_loss(&anchor, &stream, &p_m, &p_init, lambda)?; - let numgrad = (l_p - l_m) / (2.0 * eps); - let analytic = kernel_grad[idx]; - let abs = (analytic - numgrad).abs(); - let rel = abs / numgrad.abs().max(1e-3); - assert!( - abs < 5e-3 || rel < 5e-2, - "p[{idx}] numgrad mismatch: analytic={analytic:.4} numgrad={numgrad:.4} abs={abs:.3e}" - ); - } - Ok(()) -} diff --git a/crates/ml-alpha/tests/horizon_token_attention_pool_numgrad.rs b/crates/ml-alpha/tests/horizon_token_attention_pool_numgrad.rs deleted file mode 100644 index 048a7a51f..000000000 --- a/crates/ml-alpha/tests/horizon_token_attention_pool_numgrad.rs +++ /dev/null @@ -1,212 +0,0 @@ -//! Numgrad parity test for horizon_token_attention_pool (v2 axis C). -//! -//! Verifies forward + backward by central-difference comparison on -//! horizon_tokens, Q, and ln_out within the established 5e-2 rel-tol / -//! 5e-3 abs-floor envelope. - -#![cfg(feature = "cuda")] - -use anyhow::{Context, Result}; -use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut}; -use ml_alpha::horizon_token_attention_pool::{ - HorizonTokenAttentionPool, HTAP_HIDDEN_DIM, HTAP_N_HORIZONS, -}; -use ml_alpha::pinned_mem::MappedF32Buffer; -use ml_core::device::MlDevice; -use std::sync::Arc; - -const B: usize = 2; -const K: usize = 8; -const H: usize = HTAP_HIDDEN_DIM; -const N_H: usize = HTAP_N_HORIZONS; -const N_EXT: usize = N_H + K; - -fn deterministic_rng(seed: u64) -> impl FnMut() -> f32 { - let mut s = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15); - move || -> f32 { - s = s - .wrapping_mul(6364136223846793005) - .wrapping_add(1442695040888963407); - let x = ((s >> 16) & 0xFFFF) as f32 / 65536.0; - x * 2.0 - 1.0 // (-1, 1) - } -} - -fn upload(stream: &Arc, host: &[f32]) -> Result> { - let n = host.len(); - let staging = unsafe { MappedF32Buffer::new(n) } - .map_err(|e| anyhow::anyhow!("upload staging: {e}"))?; - staging.write_from_slice(host); - let mut dst = stream.alloc_zeros::(n)?; - let nbytes = n * std::mem::size_of::(); - unsafe { - let (dst_ptr, _g) = dst.device_ptr_mut(stream); - cudarc::driver::result::memcpy_dtod_async( - dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream(), - )?; - } - Ok(dst) -} - -fn download(stream: &Arc, src: &CudaSlice) -> Result> { - let n = src.len(); - let staging = unsafe { MappedF32Buffer::new(n) } - .map_err(|e| anyhow::anyhow!("download staging: {e}"))?; - let nbytes = n * std::mem::size_of::(); - unsafe { - let (src_ptr, _g) = src.device_ptr(stream); - cudarc::driver::result::memcpy_dtod_async( - staging.dev_ptr, src_ptr, nbytes, stream.cu_stream(), - )?; - } - stream.synchronize()?; - Ok(staging.read_all()) -} - -/// Forward-only: produces ctx_out for a perturbed input. -fn forward( - pool: &HorizonTokenAttentionPool, - stream: &Arc, - horizon_tokens: &[f32], - q: &[f32], - ln_out: &[f32], -) -> Result> { - let ht_d = upload(stream, horizon_tokens)?; - let q_d = upload(stream, q)?; - let ln_d = upload(stream, ln_out)?; - let mut ctx_d = stream.alloc_zeros::(B * N_H * H)?; - let mut attn_d = stream.alloc_zeros::(B * N_EXT)?; - pool.forward(&ht_d, &q_d, &ln_d, B as i32, K as i32, &mut ctx_d, &mut attn_d)?; - stream.synchronize()?; - download(stream, &ctx_d) -} - -/// Scalar loss = sum of all entries of ctx_out, weighted by an upstream -/// random `grad_ctx`. With a fixed seeded grad_ctx, ∂loss/∂ctx_out = grad_ctx -/// and chain rule via the kernel's backward should match a numgrad on -/// `loss = Σ_{b,h,d} grad_ctx[b,h,d] · ctx_out[b,h,d]`. -fn weighted_sum_loss(ctx: &[f32], grad_ctx: &[f32]) -> f32 { - ctx.iter().zip(grad_ctx).map(|(c, g)| c * g).sum() -} - -#[test] -#[ignore = "requires CUDA"] -fn forward_then_backward_matches_central_difference() -> Result<()> { - let dev = MlDevice::cuda(0)?; - let stream = dev.cuda_stream().context("stream")?.clone(); - let ctx_dev = dev.cuda_context().context("ctx")?; - let pool = HorizonTokenAttentionPool::new(ctx_dev, stream.clone())?; - - let mut rng = deterministic_rng(20260518); - let ht: Vec = (0..N_H * H).map(|_| rng() * 0.1).collect(); - let q: Vec = (0..H).map(|_| rng() * 0.1).collect(); - let ln: Vec = (0..B * K * H).map(|_| rng() * 0.5).collect(); - let grad_ctx: Vec = (0..B * N_H * H).map(|_| rng() * 0.1).collect(); - - // 1. Forward + backward through the kernel. - let ht_d = upload(&stream, &ht)?; - let q_d = upload(&stream, &q)?; - let ln_d = upload(&stream, &ln)?; - let mut ctx_d = stream.alloc_zeros::(B * N_H * H)?; - let mut attn_d = stream.alloc_zeros::(B * N_EXT)?; - pool.forward(&ht_d, &q_d, &ln_d, B as i32, K as i32, &mut ctx_d, &mut attn_d)?; - - let grad_ctx_d = upload(&stream, &grad_ctx)?; - let mut grad_ht_scratch_d = stream.alloc_zeros::(B * N_H * H)?; - let mut grad_q_scratch_d = stream.alloc_zeros::(B * H)?; - let mut grad_ln_d = stream.alloc_zeros::(B * K * H)?; - pool.backward( - &ht_d, &q_d, &ln_d, &attn_d, &grad_ctx_d, - B as i32, K as i32, - &mut grad_ht_scratch_d, &mut grad_q_scratch_d, &mut grad_ln_d, - )?; - stream.synchronize()?; - - let grad_ht_scratch = download(&stream, &grad_ht_scratch_d)?; - let grad_q_scratch = download(&stream, &grad_q_scratch_d)?; - let grad_ln = download(&stream, &grad_ln_d)?; - - // Reduce per-batch scratches to shared param-grads (host-side here - // because numgrad is host-side too; trainer uses reduce_axis0 GPU). - let mut grad_ht = vec![0.0_f32; N_H * H]; - for b in 0..B { - for i in 0..(N_H * H) { - grad_ht[i] += grad_ht_scratch[b * N_H * H + i]; - } - } - let mut grad_q = vec![0.0_f32; H]; - for b in 0..B { - for d in 0..H { - grad_q[d] += grad_q_scratch[b * H + d]; - } - } - - // 2. Central-difference numgrad on 4 sampled positions of each tensor. - let eps = 1e-3_f32; - let mut checks = 0_usize; - - let sample_indices = |total: usize, n: usize, seed_off: u64| -> Vec { - let mut r = deterministic_rng(20260518 + seed_off); - (0..n).map(|_| ((r() + 1.0) * 0.5 * total as f32) as usize % total).collect() - }; - - // horizon_tokens - for &idx in sample_indices(N_H * H, 4, 1).iter() { - let mut ht_p = ht.clone(); - let mut ht_m = ht.clone(); - ht_p[idx] += eps; - ht_m[idx] -= eps; - let l_p = weighted_sum_loss(&forward(&pool, &stream, &ht_p, &q, &ln)?, &grad_ctx); - let l_m = weighted_sum_loss(&forward(&pool, &stream, &ht_m, &q, &ln)?, &grad_ctx); - let numgrad = (l_p - l_m) / (2.0 * eps); - let analytic = grad_ht[idx]; - let abs = (analytic - numgrad).abs(); - let rel = abs / numgrad.abs().max(1e-3); - assert!( - abs < 5e-3 || rel < 5e-2, - "horizon_tokens[{idx}] mismatch: analytic={analytic:.6} numgrad={numgrad:.6} abs={abs:.3e} rel={rel:.3e}" - ); - checks += 1; - } - - // Q - for &idx in sample_indices(H, 4, 2).iter() { - let mut q_p = q.clone(); - let mut q_m = q.clone(); - q_p[idx] += eps; - q_m[idx] -= eps; - let l_p = weighted_sum_loss(&forward(&pool, &stream, &ht, &q_p, &ln)?, &grad_ctx); - let l_m = weighted_sum_loss(&forward(&pool, &stream, &ht, &q_m, &ln)?, &grad_ctx); - let numgrad = (l_p - l_m) / (2.0 * eps); - let analytic = grad_q[idx]; - let abs = (analytic - numgrad).abs(); - let rel = abs / numgrad.abs().max(1e-3); - assert!( - abs < 5e-3 || rel < 5e-2, - "Q[{idx}] mismatch: analytic={analytic:.6} numgrad={numgrad:.6} abs={abs:.3e} rel={rel:.3e}" - ); - checks += 1; - } - - // ln_out - for &idx in sample_indices(B * K * H, 4, 3).iter() { - let mut ln_p = ln.clone(); - let mut ln_m = ln.clone(); - ln_p[idx] += eps; - ln_m[idx] -= eps; - let l_p = weighted_sum_loss(&forward(&pool, &stream, &ht, &q, &ln_p)?, &grad_ctx); - let l_m = weighted_sum_loss(&forward(&pool, &stream, &ht, &q, &ln_m)?, &grad_ctx); - let numgrad = (l_p - l_m) / (2.0 * eps); - let analytic = grad_ln[idx]; - let abs = (analytic - numgrad).abs(); - let rel = abs / numgrad.abs().max(1e-3); - assert!( - abs < 5e-3 || rel < 5e-2, - "ln_out[{idx}] mismatch: analytic={analytic:.6} numgrad={numgrad:.6} abs={abs:.3e} rel={rel:.3e}" - ); - checks += 1; - } - - eprintln!("PASSED {checks} numgrad checks"); - Ok(()) -} diff --git a/crates/ml-alpha/tests/inverted_attention_pool_numgrad.rs b/crates/ml-alpha/tests/inverted_attention_pool_numgrad.rs deleted file mode 100644 index 3b4ed5ea3..000000000 --- a/crates/ml-alpha/tests/inverted_attention_pool_numgrad.rs +++ /dev/null @@ -1,127 +0,0 @@ -//! Numgrad parity test for inverted_attention_pool (v2 axis E). - -#![cfg(feature = "cuda")] - -use anyhow::{Context, Result}; -use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut}; -use ml_alpha::inverted_attention_pool::{InvertedAttentionPool, IAP_HIDDEN_DIM}; -use ml_alpha::pinned_mem::MappedF32Buffer; -use ml_core::device::MlDevice; -use std::sync::Arc; - -const B: usize = 2; -const K: usize = 8; -const H: usize = IAP_HIDDEN_DIM; - -fn rng(seed: u64) -> impl FnMut() -> f32 { - let mut s = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15); - move || -> f32 { - s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407); - let x = ((s >> 16) & 0xFFFF) as f32 / 65536.0; - x * 2.0 - 1.0 - } -} - -fn upload(stream: &Arc, host: &[f32]) -> Result> { - let n = host.len(); - let staging = unsafe { MappedF32Buffer::new(n) } - .map_err(|e| anyhow::anyhow!("upload staging: {e}"))?; - staging.write_from_slice(host); - let mut dst = stream.alloc_zeros::(n)?; - let nbytes = n * std::mem::size_of::(); - unsafe { - let (dst_ptr, _g) = dst.device_ptr_mut(stream); - cudarc::driver::result::memcpy_dtod_async( - dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream(), - )?; - } - Ok(dst) -} - -fn download(stream: &Arc, src: &CudaSlice) -> Result> { - let n = src.len(); - let staging = unsafe { MappedF32Buffer::new(n) } - .map_err(|e| anyhow::anyhow!("download staging: {e}"))?; - let nbytes = n * std::mem::size_of::(); - unsafe { - let (src_ptr, _g) = src.device_ptr(stream); - cudarc::driver::result::memcpy_dtod_async( - staging.dev_ptr, src_ptr, nbytes, stream.cu_stream(), - )?; - } - stream.synchronize()?; - Ok(staging.read_all()) -} - -fn forward_loss( - pool: &InvertedAttentionPool, - stream: &Arc, - ln: &[f32], - grad_pooled: &[f32], -) -> Result { - let inv_scale = 1.0_f32 / (K as f32).sqrt(); - let ln_d = upload(stream, ln)?; - let mut pooled_d = stream.alloc_zeros::(B * H)?; - let mut attn_d = stream.alloc_zeros::(B * H * H)?; - pool.forward(&ln_d, B as i32, K as i32, inv_scale, &mut pooled_d, &mut attn_d)?; - stream.synchronize()?; - let pooled = download(stream, &pooled_d)?; - Ok(pooled.iter().zip(grad_pooled).map(|(p, g)| p * g).sum()) -} - -#[test] -#[ignore = "requires CUDA"] -fn forward_then_backward_matches_central_difference() -> Result<()> { - let dev = MlDevice::cuda(0)?; - let stream = dev.cuda_stream().context("stream")?.clone(); - let ctx_dev = dev.cuda_context().context("ctx")?; - let pool = InvertedAttentionPool::new(ctx_dev, stream.clone())?; - - let mut r = rng(20260518); - let ln: Vec = (0..B * K * H).map(|_| r() * 0.3).collect(); - let grad_pooled: Vec = (0..B * H).map(|_| r() * 0.1).collect(); - - let inv_scale = 1.0_f32 / (K as f32).sqrt(); - let ln_d = upload(&stream, &ln)?; - let mut pooled_d = stream.alloc_zeros::(B * H)?; - let mut attn_d = stream.alloc_zeros::(B * H * H)?; - pool.forward(&ln_d, B as i32, K as i32, inv_scale, &mut pooled_d, &mut attn_d)?; - - let grad_pooled_d = upload(&stream, &grad_pooled)?; - let mut d_scores_scratch_d = stream.alloc_zeros::(B * H * H)?; - let mut grad_ln_d = stream.alloc_zeros::(B * K * H)?; - pool.backward( - &ln_d, &attn_d, &grad_pooled_d, - B as i32, K as i32, inv_scale, - &mut d_scores_scratch_d, - &mut grad_ln_d, - )?; - stream.synchronize()?; - - let grad_ln = download(&stream, &grad_ln_d)?; - - let eps = 1e-3_f32; - let mut r2 = rng(42); - let total = B * K * H; - let mut checks = 0; - for _ in 0..6 { - let idx = ((r2() + 1.0) * 0.5 * total as f32) as usize % total; - let mut ln_p = ln.clone(); - let mut ln_m = ln.clone(); - ln_p[idx] += eps; - ln_m[idx] -= eps; - let l_p = forward_loss(&pool, &stream, &ln_p, &grad_pooled)?; - let l_m = forward_loss(&pool, &stream, &ln_m, &grad_pooled)?; - let numgrad = (l_p - l_m) / (2.0 * eps); - let analytic = grad_ln[idx]; - let abs = (analytic - numgrad).abs(); - let rel = abs / numgrad.abs().max(1e-3); - assert!( - abs < 5e-3 || rel < 5e-2, - "ln_out[{idx}] mismatch: analytic={analytic:.6} numgrad={numgrad:.6} abs={abs:.3e} rel={rel:.3e}" - ); - checks += 1; - } - eprintln!("PASSED {checks} numgrad checks"); - Ok(()) -} diff --git a/crates/ml-alpha/tests/regime_moe_gate_numgrad.rs b/crates/ml-alpha/tests/regime_moe_gate_numgrad.rs deleted file mode 100644 index 0db731634..000000000 --- a/crates/ml-alpha/tests/regime_moe_gate_numgrad.rs +++ /dev/null @@ -1,290 +0,0 @@ -//! Numgrad parity test for regime_moe_gate (v2 axis D). -//! -//! Top-1 routing has straight-through grad on the gate logits — the -//! gate side isn't tested by numgrad here (it's a non-differentiable -//! argmax). What we DO test: -//! 1. Forward correctness: routed_out matches a host-side reference -//! that picks the same top-1 expert. -//! 2. Backward grads on the selected expert's W and bias match -//! central-difference numgrad. -//! 3. Grad on fused_ctx matches numgrad. - -#![cfg(feature = "cuda")] - -use anyhow::{Context, Result}; -use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut}; -use ml_alpha::pinned_mem::MappedF32Buffer; -use ml_alpha::regime_moe_gate::{ - RegimeMoeGate, MOE_HIDDEN_DIM, MOE_N_EXPERTS, MOE_N_HORIZONS, -}; -use ml_core::device::MlDevice; -use std::sync::Arc; - -const B: usize = 3; -const H: usize = MOE_HIDDEN_DIM; -const N_H: usize = MOE_N_HORIZONS; -const N_E: usize = MOE_N_EXPERTS; - -fn rng(seed: u64) -> impl FnMut() -> f32 { - let mut s = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15); - move || -> f32 { - s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407); - let x = ((s >> 16) & 0xFFFF) as f32 / 65536.0; - x * 2.0 - 1.0 - } -} - -fn upload(stream: &Arc, host: &[f32]) -> Result> { - let n = host.len(); - let staging = unsafe { MappedF32Buffer::new(n) } - .map_err(|e| anyhow::anyhow!("upload staging: {e}"))?; - staging.write_from_slice(host); - let mut dst = stream.alloc_zeros::(n)?; - let nbytes = n * std::mem::size_of::(); - unsafe { - let (dst_ptr, _g) = dst.device_ptr_mut(stream); - cudarc::driver::result::memcpy_dtod_async( - dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream(), - )?; - } - Ok(dst) -} - -#[allow(dead_code)] // helper kept for future numgrad expansion (gate-side STE) -fn upload_i32(stream: &Arc, host: &[i32]) -> Result> { - let n = host.len(); - let staging = unsafe { MappedF32Buffer::new(n) } - .map_err(|e| anyhow::anyhow!("upload i32 staging: {e}"))?; - // f32 buffer reinterpreted as i32. - let host_ptr = staging.host_ptr.cast::(); - unsafe { - for i in 0..n { std::ptr::write_volatile(host_ptr.add(i), host[i]); } - } - let mut dst = stream.alloc_zeros::(n)?; - let nbytes = n * std::mem::size_of::(); - unsafe { - let (dst_ptr, _g) = dst.device_ptr_mut(stream); - cudarc::driver::result::memcpy_dtod_async( - dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream(), - )?; - } - Ok(dst) -} - -fn download(stream: &Arc, src: &CudaSlice) -> Result> { - let n = src.len(); - let staging = unsafe { MappedF32Buffer::new(n) } - .map_err(|e| anyhow::anyhow!("download staging: {e}"))?; - let nbytes = n * std::mem::size_of::(); - unsafe { - let (src_ptr, _g) = src.device_ptr(stream); - cudarc::driver::result::memcpy_dtod_async( - staging.dev_ptr, src_ptr, nbytes, stream.cu_stream(), - )?; - } - stream.synchronize()?; - Ok(staging.read_all()) -} - -fn download_i32(stream: &Arc, src: &CudaSlice) -> Result> { - let n = src.len(); - let staging = unsafe { MappedF32Buffer::new(n) } - .map_err(|e| anyhow::anyhow!("download i32 staging: {e}"))?; - let nbytes = n * std::mem::size_of::(); - unsafe { - let (src_ptr, _g) = src.device_ptr(stream); - cudarc::driver::result::memcpy_dtod_async( - staging.dev_ptr, src_ptr, nbytes, stream.cu_stream(), - )?; - } - stream.synchronize()?; - let host_ptr = staging.host_ptr.cast::(); - let mut out = Vec::with_capacity(n); - unsafe { - for i in 0..n { out.push(std::ptr::read_volatile(host_ptr.add(i))); } - } - Ok(out) -} - -struct Inputs { - gate_logits: Vec, - fused_ctx: Vec, - experts_w: Vec, - experts_b: Vec, - grad_routed: Vec, -} - -fn make_inputs(seed: u64) -> Inputs { - let mut r = rng(seed); - let gate_logits: Vec = (0..B * N_E).map(|_| r() * 1.5).collect(); - let fused_ctx: Vec = (0..B * N_H * H).map(|_| r() * 0.3).collect(); - let experts_w: Vec = (0..N_E * H * H).map(|_| r() * 0.05).collect(); - let experts_b: Vec = (0..N_E * H).map(|_| r() * 0.02).collect(); - let grad_routed: Vec = (0..B * N_H * H).map(|_| r() * 0.1).collect(); - Inputs { gate_logits, fused_ctx, experts_w, experts_b, grad_routed } -} - -fn forward_routed_loss( - moe: &RegimeMoeGate, - stream: &Arc, - inputs: &Inputs, - fused_ctx_override: Option<&[f32]>, - experts_w_override: Option<&[f32]>, - experts_b_override: Option<&[f32]>, -) -> Result { - let fused_ctx = fused_ctx_override.unwrap_or(&inputs.fused_ctx); - let experts_w = experts_w_override.unwrap_or(&inputs.experts_w); - let experts_b = experts_b_override.unwrap_or(&inputs.experts_b); - - let gate_d = upload(stream, &inputs.gate_logits)?; - let ctx_d = upload(stream, fused_ctx)?; - let w_d = upload(stream, experts_w)?; - let b_d = upload(stream, experts_b)?; - let mut routed_d = stream.alloc_zeros::(B * N_H * H)?; - let mut top_e_d = stream.alloc_zeros::(B)?; - moe.forward(&gate_d, &ctx_d, &w_d, &b_d, B as i32, &mut routed_d, &mut top_e_d)?; - stream.synchronize()?; - let routed = download(stream, &routed_d)?; - Ok(routed.iter().zip(&inputs.grad_routed).map(|(r, g)| r * g).sum()) -} - -#[test] -#[ignore = "requires CUDA"] -fn forward_matches_host_reference_and_backward_matches_numgrad() -> Result<()> { - let dev = MlDevice::cuda(0)?; - let stream = dev.cuda_stream().context("stream")?.clone(); - let ctx_dev = dev.cuda_context().context("ctx")?; - let moe = RegimeMoeGate::new(ctx_dev, stream.clone())?; - - let inputs = make_inputs(20260518); - - // 1. Run forward, capture top_e. - let gate_d = upload(&stream, &inputs.gate_logits)?; - let ctx_d = upload(&stream, &inputs.fused_ctx)?; - let w_d = upload(&stream, &inputs.experts_w)?; - let b_d = upload(&stream, &inputs.experts_b)?; - let mut routed_d = stream.alloc_zeros::(B * N_H * H)?; - let mut top_e_d = stream.alloc_zeros::(B)?; - moe.forward(&gate_d, &ctx_d, &w_d, &b_d, B as i32, &mut routed_d, &mut top_e_d)?; - stream.synchronize()?; - let routed = download(&stream, &routed_d)?; - let top_e = download_i32(&stream, &top_e_d)?; - - // Host reference for routed. - for b in 0..B { - let e = top_e[b] as usize; - for h in 0..N_H { - for d_out in 0..H { - let mut acc = inputs.experts_b[e * H + d_out]; - for d_in in 0..H { - acc += inputs.experts_w[e * H * H + d_out * H + d_in] - * inputs.fused_ctx[b * N_H * H + h * H + d_in]; - } - let kernel_val = routed[b * N_H * H + h * H + d_out]; - let diff = (acc - kernel_val).abs(); - assert!(diff < 1e-3, - "forward mismatch at b={b} h={h} d_out={d_out}: ref={acc:.5} kernel={kernel_val:.5}"); - } - } - } - - // 2. Backward — compute grads. Compact scratch [B, N_H, H, H] / [B, N_H, H]. - let grad_routed_d = upload(&stream, &inputs.grad_routed)?; - let mut d_w_scratch_d = stream.alloc_zeros::(B * N_H * H * H)?; - let mut d_b_scratch_d = stream.alloc_zeros::(B * N_H * H)?; - let mut d_fused_ctx_d = stream.alloc_zeros::(B * N_H * H)?; - moe.backward( - &ctx_d, &w_d, &grad_routed_d, &top_e_d, B as i32, - &mut d_w_scratch_d, &mut d_b_scratch_d, &mut d_fused_ctx_d, - )?; - // Scatter into per-expert grads via the new MoE.scatter kernel. - let mut grad_experts_w_d = stream.alloc_zeros::(N_E * H * H)?; - let mut grad_experts_b_d = stream.alloc_zeros::(N_E * H)?; - moe.scatter( - &d_w_scratch_d, &d_b_scratch_d, &top_e_d, B as i32, - &mut grad_experts_w_d, &mut grad_experts_b_d, - )?; - stream.synchronize()?; - let d_w_full = download(&stream, &grad_experts_w_d)?; - let d_b_full = download(&stream, &grad_experts_b_d)?; - let d_fused_ctx = download(&stream, &d_fused_ctx_d)?; - // (no host-side reduce — scatter kernel produces the per-expert grad directly) - - // 3. Numgrad checks. - let eps = 1e-3_f32; - let mut r = rng(42); - let mut checks = 0; - - // d_W on a few sampled (e, d_out, d_in) — only check active experts. - let active: std::collections::HashSet = - top_e.iter().map(|&e| e as usize).collect(); - for _ in 0..4 { - let e = *active.iter().nth((r().abs() * active.len() as f32) as usize % active.len()).unwrap(); - let d_out = ((r() + 1.0) * 0.5 * H as f32) as usize % H; - let d_in = ((r() + 1.0) * 0.5 * H as f32) as usize % H; - let idx = e * H * H + d_out * H + d_in; - let mut w_p = inputs.experts_w.clone(); - let mut w_m = inputs.experts_w.clone(); - w_p[idx] += eps; - w_m[idx] -= eps; - let l_p = forward_routed_loss(&moe, &stream, &inputs, None, Some(&w_p), None)?; - let l_m = forward_routed_loss(&moe, &stream, &inputs, None, Some(&w_m), None)?; - let numgrad = (l_p - l_m) / (2.0 * eps); - let analytic = d_w_full[idx]; - let abs = (analytic - numgrad).abs(); - let rel = abs / numgrad.abs().max(1e-3); - assert!( - abs < 5e-3 || rel < 5e-2, - "d_W[{e},{d_out},{d_in}] mismatch: analytic={analytic:.4} numgrad={numgrad:.4} abs={abs:.3e} rel={rel:.3e}" - ); - checks += 1; - } - - // d_b on a few sampled (e, d_out). - for _ in 0..3 { - let e = *active.iter().nth((r().abs() * active.len() as f32) as usize % active.len()).unwrap(); - let d_out = ((r() + 1.0) * 0.5 * H as f32) as usize % H; - let idx = e * H + d_out; - let mut b_p = inputs.experts_b.clone(); - let mut b_m = inputs.experts_b.clone(); - b_p[idx] += eps; - b_m[idx] -= eps; - let l_p = forward_routed_loss(&moe, &stream, &inputs, None, None, Some(&b_p))?; - let l_m = forward_routed_loss(&moe, &stream, &inputs, None, None, Some(&b_m))?; - let numgrad = (l_p - l_m) / (2.0 * eps); - let analytic = d_b_full[idx]; - let abs = (analytic - numgrad).abs(); - let rel = abs / numgrad.abs().max(1e-3); - assert!( - abs < 5e-3 || rel < 5e-2, - "d_b[{e},{d_out}] mismatch: analytic={analytic:.4} numgrad={numgrad:.4} abs={abs:.3e} rel={rel:.3e}" - ); - checks += 1; - } - - // d_fused_ctx on a few sampled positions. - for _ in 0..4 { - let b_idx = ((r() + 1.0) * 0.5 * B as f32) as usize % B; - let h_idx = ((r() + 1.0) * 0.5 * N_H as f32) as usize % N_H; - let d_idx = ((r() + 1.0) * 0.5 * H as f32) as usize % H; - let idx = b_idx * N_H * H + h_idx * H + d_idx; - let mut ctx_p = inputs.fused_ctx.clone(); - let mut ctx_m = inputs.fused_ctx.clone(); - ctx_p[idx] += eps; - ctx_m[idx] -= eps; - let l_p = forward_routed_loss(&moe, &stream, &inputs, Some(&ctx_p), None, None)?; - let l_m = forward_routed_loss(&moe, &stream, &inputs, Some(&ctx_m), None, None)?; - let numgrad = (l_p - l_m) / (2.0 * eps); - let analytic = d_fused_ctx[idx]; - let abs = (analytic - numgrad).abs(); - let rel = abs / numgrad.abs().max(1e-3); - assert!( - abs < 5e-3 || rel < 5e-2, - "d_fused_ctx[{b_idx},{h_idx},{d_idx}] mismatch: analytic={analytic:.4} numgrad={numgrad:.4} abs={abs:.3e} rel={rel:.3e}" - ); - checks += 1; - } - - eprintln!("PASSED {checks} numgrad checks"); - Ok(()) -}