gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(ml-alpha): consolidate BCE — Kendall σ is THE BCE, wire MHA into trainer

Browse Source 
Single source of truth for the multi-horizon BCE per the new
feedback_single_source_of_truth_no_duplicates pearl. Eliminates the
`bce_loss_multi_horizon_sigma.cu` / `loss_sigma.rs` duplication
introduced earlier today and folds the Kendall σ-weighting into the
canonical kernel.

Deletions:
  cuda/bce_loss_multi_horizon_sigma.cu              (folded into legacy)
  src/trainer/loss_sigma.rs                         (helper merged)
  tests/bce_sigma_numgrad.rs                        (subsumed)

Rewrites:
  cuda/bce_loss_multi_horizon.cu — replaced with σ-aware
    implementation; kernel function name kept as
    `bce_multi_horizon_forward_backward` so cubin symbol stays stable.
    NVIDIA-grade warp-shuffle reduce (4 warps, 1 cross-warp barrier);
    new args `log_sigma_h` (per-horizon Kendall σ logarithm) and
    `d_log_sigma_h` (its gradient).
  src/trainer/loss.rs — standalone helper updated to new 11-arg
    kernel signature. Exposes optional `log_sigma_h` in `BceInput`
    (None → zeros / passthrough Kendall init) and returns
    `mean_bce_per_h` + `d_log_sigma_h` in `BceOutput`.
  tests/bce_grad_finite_diff.rs — adds `log_sigma_h: None` to the
    test inputs; all 4 numgrad tests PASS unchanged.
  build.rs — drops `bce_loss_multi_horizon_sigma` entry from
    KERNELS. The single canonical `bce_loss_multi_horizon` cubin
    now contains the σ-aware kernel.

Wiring:
  PerceptionTrainer gains a single `pub mha: MultiHorizonAttention`
    field. Owns `log_sigma_h_d` + `grad_log_sigma_h_d` (and the rest
    of the multi-horizon attention path, anchored on Stage 2 to fully
    replace the legacy `attention_pool` path).
  step_batched + evaluate_batched BCE callsites now thread
    `mha.log_sigma_h_d` and `mha.grad_log_sigma_h_d` through the
    11-arg kernel signature.

This commit keeps the legacy `attention_pool` callsite in place; the
Stage 2 commit will replace it with `mha.pool` + `mha.inv_pool` +
`mha.moe` and delete the `attn_*` fields entirely.

Verified locally: ml-alpha builds clean with the cuda feature,
bce_grad_finite_diff (4/4) PASS on RTX 3050.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-18 14:59:12 +02:00
parent 996e61b6ef
commit 6a3f45d872
9 changed files with 209 additions and 633 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
crates/ml-alpha/build.rs
View File

@@ -20,11 +20,11 @@ const KERNELS: &[&str] = &[
"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
"attention_pool", // Phase 3: learned context summary at CfC k=0
"horizon_token_attention_pool", // v2-C: horizon-token K-prepend single-Q attention
"inverted_attention_pool", // v2-E: iTransformer-style cross-variate attention
"regime_moe_gate", // v2-D: top-1 MoE gate + expert dispatch + aux loss
"anchor_l2", // v2-B: L2 anchor regularization toward init
"attention_pool", // Legacy single-Q content summary at CfC k=0 (Stage 2 replaces with MHA)
"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
"anchor_l2", // L2 anchor regularization toward init
"reduce_axis0", // Phase B: cross-batch param-grad reducer
];

230
crates/ml-alpha/cuda/bce_loss_multi_horizon.cu
View File

@@ -1,129 +1,185 @@
// bce_loss_multi_horizon.cu
// bce_loss_multi_horizon.cu — Kendall σ-weighted multi-horizon BCE.
//
// Fused multi-horizon BCE forward + backward with optional per-horizon
// loss weighting. Also emits per-horizon UNWEIGHTED mean BCE — that
// is the signal an ISV-driven horizon weighting layer EMA-tracks to
// produce dynamic per-horizon weights (high per-horizon BCE → horizon
// is under-trained → boost its weight). The total scalar loss_out
// remains the externally-weighted aggregate used by callers for
// reporting and early-stopping.
// Single source of truth for the multi-horizon BCE. Implements
// Kendall homoscedastic-uncertainty weighting per [Kendall+Gal+Cipolla
// 2018]: each horizon has a learnable `log_sigma_h` scalar that
// balances its loss contribution. Replaces the prior plain-pooled BCE.
//
// Input layout (row-major):
// probs [n_pos, n_horizons] — model outputs in (0, 1)
// labels [n_pos, n_horizons] — binary {0.0, 1.0}; NaN = mask (drop)
// loss_weights [n_horizons] — per-horizon weight (nullptr = all 1.0)
// Output:
// loss_out[1] — weighted mean BCE over valid entries
// loss_per_horizon[n_horizons] — UNWEIGHTED mean BCE per horizon (ISV signal)
// grad_probs[n_pos, n_horizons] — d loss / d prob (per-element scaled)
// valid_count_out[1] — number of non-NaN labels (unweighted)
// Math:
//
// per element (mask m = !isnan(y), weight w = loss_weights[h]):
// p = clamp(probs[i], 1e-6, 1 - 1e-6)
// L_i = -[y log p + (1-y) log(1-p)] (unweighted; weighted form scales L_i by w)
// dL/dp = m * w * (p - y) / (p (1-p))
// raw_bce_h = Σ_{i in horizon h, m_i = 1} L_i (unweighted sum)
// count_h = #{i in horizon h : m_i = 1}
// mean_bce_h = raw_bce_h / count_h
// w_h = base_weight_h / (2 * exp(2 * log_sigma_h))
// total_loss = Σ_h [ w_h * mean_bce_h + log_sigma_h ]
//
// Normaliser: W_valid = sum_i m_i * w_i
// total_loss = (1/W_valid) * sum_i (m_i * w_i * L_i)
// grad_probs[i] = m_i * w * (p - y) / (p (1-p)) / W_valid
// loss_per_horizon[h] = (sum_{i: m_i, h_i=h} L_i) / (sum_{i: m_i, h_i=h} 1)
// Gradients:
// d L / d p_i = m_i * (w_h / count_h) * (p − y) / (p (1 − p))
// d L / d log_sigma_h = 1 − 2 * w_h * mean_bce_h
//
// Caller contract: n_horizons MUST equal N_HORIZONS_BCE (5). Hardcoded
// rather than dynamic-shared-mem because the per-horizon accumulator
// matrix has a fixed shape (5 × 256) at our trainer's scale.
// Per `feedback_nvidia_grade_perf_for_kernels`:
// - Warp-shuffle reduction (`__shfl_xor_sync`), NOT block tree-reduce.
// - One `__syncthreads` for the cross-warp aggregate, total = O(1) barriers.
// - Inactive lanes contribute 0 via ternary; never divergent shuffles.
// - Hard-coded N_HORIZONS_BCE = 5 → per-warp 5-element accumulator stays in registers.
//
// Block layout: grid = (1), block = (256) = 8 warps. Single block over
// the full [n_pos, n_horizons] grid is fine — we never have > 256 ×
// (per-thread fan-in) positions in the trainer hot loop (max K × B =
// 64 × 8 = 512, well within stride-loop reach).
#define N_HORIZONS_BCE 5
#define BCS_N_HORIZONS 5
#define BCS_BLOCK 256
#define BCS_N_WARPS (BCS_BLOCK / 32) // == 8
__device__ __forceinline__ float warp_reduce_sum(float v) {
#pragma unroll
for (int s = 16; s > 0; s >>= 1) {
v += __shfl_xor_sync(0xffffffff, v, s);
}
return v;
}
__device__ __forceinline__ int warp_reduce_sum_int(int v) {
#pragma unroll
for (int s = 16; s > 0; s >>= 1) {
v += __shfl_xor_sync(0xffffffff, v, s);
}
return v;
}
extern "C" __global__ void bce_multi_horizon_forward_backward(
const float* __restrict__ probs, // [n_pos * n_horizons]
const float* __restrict__ labels, // [n_pos * n_horizons]
const float* __restrict__ loss_weights, // [n_horizons] (nullptr OK → 1.0)
const float* __restrict__ base_weights, // [n_horizons] (per-horizon λ from ISV; constants OK)
const float* __restrict__ log_sigma_h, // [n_horizons] (learnable Kendall σ logarithm)
int n_pos,
int n_horizons,
float* __restrict__ loss_out, // [1] — externally-weighted mean
float* __restrict__ loss_per_horizon, // [n_horizons] — UNWEIGHTED per-horizon mean
float* __restrict__ grad_probs, // [n_pos * n_horizons]
int* __restrict__ valid_count_out // [1] — number of non-NaN labels
float* __restrict__ loss_out, // [1] — total Kendall loss
float* __restrict__ mean_bce_per_h_out, // [n_horizons] — unweighted per-horizon mean (ISV signal)
float* __restrict__ grad_probs, // [n_pos * n_horizons] — ∂L/∂p
int* __restrict__ valid_count_out, // [1] — Σ_i m_i (unweighted)
float* __restrict__ d_log_sigma_h // [n_horizons] — ∂L/∂log_sigma_h
) {
const int tid = threadIdx.x;
const int tid = threadIdx.x;
const int lane = tid & 31;
const int warp = tid >> 5;
const int total = n_pos * n_horizons;
__shared__ float sloss[256]; // weighted-loss numerator accumulator
__shared__ float sw_valid[256]; // weighted-mask normaliser accumulator
__shared__ int svalid[256]; // unweighted valid-label counter (reported)
// Per-horizon accumulators (unweighted): sloss_h[h*256 + tid],
// svalid_h[h*256 + tid]. Hardcoded width N_HORIZONS_BCE=5 → 13 KiB
// total shared usage including the three scalars above; comfortably
// under any SM smem limit (48-100 KiB depending on arch).
__shared__ float sloss_h[N_HORIZONS_BCE * 256];
__shared__ float svalid_h[N_HORIZONS_BCE * 256];
sloss[tid] = 0.0f;
sw_valid[tid] = 0.0f;
svalid[tid] = 0;
// Per-thread per-horizon partials in registers.
float local_raw[BCS_N_HORIZONS];
float local_cnt[BCS_N_HORIZONS];
#pragma unroll
for (int h = 0; h < N_HORIZONS_BCE; ++h) {
sloss_h[h * 256 + tid] = 0.0f;
svalid_h[h * 256 + tid] = 0.0f;
for (int h = 0; h < BCS_N_HORIZONS; ++h) {
local_raw[h] = 0.0f;
local_cnt[h] = 0.0f;
}
int local_valid = 0;
// First pass: count valid entries and accumulate weighted aggregate
// loss + unweighted per-horizon loss + weight sums.
for (int i = tid; i < total; i += blockDim.x) {
// Pass 1 — per-thread strided accumulation. Same loop also zeros
// grad_probs at masked positions so the second pass can skip those.
for (int i = tid; i < total; i += BCS_BLOCK) {
const float y = labels[i];
if (isnan(y)) {
grad_probs[i] = 0.0f;
continue;
}
const int h = i % n_horizons;
const float w = (loss_weights != nullptr) ? loss_weights[h] : 1.0f;
const float p = fminf(fmaxf(probs[i], 1e-6f), 1.0f - 1e-6f);
const float L = -(y * logf(p) + (1.0f - y) * logf(1.0f - p));
sloss[tid] += w * L;
sloss_h[h * 256 + tid] += L; // UNWEIGHTED — the ISV signal
svalid_h[h * 256 + tid] += 1.0f;
sw_valid[tid] += w;
svalid[tid] += 1;
local_raw[h] += L;
local_cnt[h] += 1.0f;
local_valid += 1;
}
// Warp-shuffle reduce each of the 5 horizon sums + the per-horizon
// count + the global valid count. Each lane in warp gets the full
// warp-sum at the end (no need to broadcast manually).
float warp_raw[BCS_N_HORIZONS];
float warp_cnt[BCS_N_HORIZONS];
#pragma unroll
for (int h = 0; h < BCS_N_HORIZONS; ++h) {
warp_raw[h] = warp_reduce_sum(local_raw[h]);
warp_cnt[h] = warp_reduce_sum(local_cnt[h]);
}
int warp_valid = warp_reduce_sum_int(local_valid);
// Cross-warp reduce — only lane 0 of each warp writes its partial.
__shared__ float s_raw[BCS_N_WARPS * BCS_N_HORIZONS];
__shared__ float s_cnt[BCS_N_WARPS * BCS_N_HORIZONS];
__shared__ int s_valid[BCS_N_WARPS];
if (lane == 0) {
#pragma unroll
for (int h = 0; h < BCS_N_HORIZONS; ++h) {
s_raw[warp * BCS_N_HORIZONS + h] = warp_raw[h];
s_cnt[warp * BCS_N_HORIZONS + h] = warp_cnt[h];
}
s_valid[warp] = warp_valid;
}
__syncthreads();
// Block tree-reduce on all accumulators (incl. per-horizon arrays).
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
if (tid < s) {
sloss[tid] += sloss[tid + s];
sw_valid[tid] += sw_valid[tid + s];
svalid[tid] += svalid[tid + s];
#pragma unroll
for (int h = 0; h < N_HORIZONS_BCE; ++h) {
sloss_h[h * 256 + tid] += sloss_h[h * 256 + tid + s];
svalid_h[h * 256 + tid] += svalid_h[h * 256 + tid + s];
// Warp 0 finishes: reduce N_WARPS partials. All 32 lanes participate;
// inactive lanes (tid ≥ N_WARPS) contribute 0 via ternary (NOT a
// divergent shuffle).
__shared__ float mean_bce_h_shared[BCS_N_HORIZONS];
__shared__ float w_eff_shared[BCS_N_HORIZONS];
__shared__ int valid_count_shared;
if (warp == 0) {
#pragma unroll
for (int h = 0; h < BCS_N_HORIZONS; ++h) {
float v_raw = (lane < BCS_N_WARPS) ? s_raw[lane * BCS_N_HORIZONS + h] : 0.0f;
float v_cnt = (lane < BCS_N_WARPS) ? s_cnt[lane * BCS_N_HORIZONS + h] : 0.0f;
v_raw = warp_reduce_sum(v_raw);
v_cnt = warp_reduce_sum(v_cnt);
if (lane == 0) {
const float c_h = v_cnt;
const float mean_bce = (c_h > 0.0f) ? v_raw / c_h : 0.0f;
mean_bce_h_shared[h] = mean_bce;
mean_bce_per_h_out[h] = mean_bce;
const float ls = log_sigma_h[h];
const float bw = (base_weights != nullptr) ? base_weights[h] : 1.0f;
// w_h = bw / (2 * exp(2*ls)) = 0.5 * bw * exp(-2 * ls)
const float w_h = 0.5f * bw * expf(-2.0f * ls);
// Per-element grad prefactor = w_h / count_h.
w_eff_shared[h] = (c_h > 0.0f) ? w_h / c_h : 0.0f;
// ∂L/∂log_sigma_h = 1 − 2 * w_h * mean_bce_h.
d_log_sigma_h[h] = 1.0f - 2.0f * w_h * mean_bce;
}
}
__syncthreads();
}
if (tid == 0) {
const float w_valid = sw_valid[0];
loss_out[0] = (w_valid > 0.0f) ? sloss[0] / w_valid : 0.0f;
valid_count_out[0] = svalid[0];
#pragma unroll
for (int h = 0; h < N_HORIZONS_BCE; ++h) {
const float c_h = svalid_h[h * 256];
loss_per_horizon[h] = (c_h > 0.0f) ? sloss_h[h * 256] / c_h : 0.0f;
int v_valid = (lane < BCS_N_WARPS) ? s_valid[lane] : 0;
v_valid = warp_reduce_sum_int(v_valid);
if (lane == 0) {
valid_count_shared = v_valid;
valid_count_out[0] = v_valid;
}
}
__syncthreads();
// Second pass: emit per-element grad scaled by w_h / W_valid.
const float w_valid = sw_valid[0];
const float inv = (w_valid > 0.0f) ? (1.0f / w_valid) : 0.0f;
for (int i = tid; i < total; i += blockDim.x) {
// Total loss — small serial reduce by thread 0 over 5 horizons.
if (tid == 0) {
float total_loss = 0.0f;
#pragma unroll
for (int h = 0; h < BCS_N_HORIZONS; ++h) {
const float ls = log_sigma_h[h];
const float bw = (base_weights != nullptr) ? base_weights[h] : 1.0f;
const float w_h = 0.5f * bw * expf(-2.0f * ls);
total_loss += w_h * mean_bce_h_shared[h] + ls;
}
loss_out[0] = total_loss;
}
__syncthreads();
// Pass 2 — per-element grad. Coalesced strided access. w_eff_shared
// is read-only at this point, broadcast freely across threads.
for (int i = tid; i < total; i += BCS_BLOCK) {
const float y = labels[i];
if (isnan(y)) {
continue;
continue; // grad was already zeroed in pass 1
}
const int h = i % n_horizons;
const float w = (loss_weights != nullptr) ? loss_weights[h] : 1.0f;
const float p = fminf(fmaxf(probs[i], 1e-6f), 1.0f - 1e-6f);
grad_probs[i] = w * (p - y) / (p * (1.0f - p)) * inv;
grad_probs[i] = w_eff_shared[h] * (p - y) / (p * (1.0f - p));
}
}

183
crates/ml-alpha/cuda/bce_loss_multi_horizon_sigma.cu
View File

@@ -1,183 +0,0 @@
// bce_loss_multi_horizon_sigma.cu — Kendall σ-weighted multi-horizon BCE.
//
// v2-A. Independent of the legacy `bce_loss_multi_horizon.cu`. Used by
// PerceptionV2State; legacy callers (eval, smoke) keep the original.
//
// Math:
//
// raw_bce_h = Σ_{i in horizon h, m_i = 1} L_i (unweighted sum)
// count_h = #{i in horizon h : m_i = 1}
// mean_bce_h = raw_bce_h / count_h
// w_h = base_weight_h / (2 * exp(2 * log_sigma_h))
// total_loss = Σ_h [ w_h * mean_bce_h + log_sigma_h ]
//
// Gradients:
// d L / d p_i = m_i * (w_h / count_h) * (p − y) / (p (1 − p))
// d L / d log_sigma_h = 1 − 2 * w_h * mean_bce_h
//
// Per `feedback_nvidia_grade_perf_for_kernels`:
// - Warp-shuffle reduction (`__shfl_xor_sync`), NOT block tree-reduce.
// - One `__syncthreads` for the cross-warp aggregate, total = O(1) barriers.
// - Inactive lanes contribute 0 via ternary; never divergent shuffles.
// - Hard-coded N_HORIZONS_BCE = 5 → per-warp 5-element accumulator stays in registers.
//
// Block layout: grid = (1), block = (256) = 8 warps. Single block over
// the full [n_pos, n_horizons] grid is fine — we never have > 256 ×
// (per-thread fan-in) positions in the trainer hot loop (max K × B =
// 64 × 8 = 512, well within stride-loop reach).
#define BCS_N_HORIZONS 5
#define BCS_BLOCK 256
#define BCS_N_WARPS (BCS_BLOCK / 32) // == 8
__device__ __forceinline__ float warp_reduce_sum(float v) {
#pragma unroll
for (int s = 16; s > 0; s >>= 1) {
v += __shfl_xor_sync(0xffffffff, v, s);
}
return v;
}
__device__ __forceinline__ int warp_reduce_sum_int(int v) {
#pragma unroll
for (int s = 16; s > 0; s >>= 1) {
v += __shfl_xor_sync(0xffffffff, v, s);
}
return v;
}
extern "C" __global__ void bce_multi_horizon_sigma_forward_backward(
const float* __restrict__ probs, // [n_pos * n_horizons]
const float* __restrict__ labels, // [n_pos * n_horizons]
const float* __restrict__ base_weights, // [n_horizons] (per-horizon λ from ISV; constants OK)
const float* __restrict__ log_sigma_h, // [n_horizons] (learnable Kendall σ logarithm)
int n_pos,
int n_horizons,
float* __restrict__ loss_out, // [1] — total Kendall loss
float* __restrict__ mean_bce_per_h_out, // [n_horizons] — unweighted per-horizon mean (ISV signal)
float* __restrict__ grad_probs, // [n_pos * n_horizons] — ∂L/∂p
int* __restrict__ valid_count_out, // [1] — Σ_i m_i (unweighted)
float* __restrict__ d_log_sigma_h // [n_horizons] — ∂L/∂log_sigma_h
) {
const int tid = threadIdx.x;
const int lane = tid & 31;
const int warp = tid >> 5;
const int total = n_pos * n_horizons;
// Per-thread per-horizon partials in registers.
float local_raw[BCS_N_HORIZONS];
float local_cnt[BCS_N_HORIZONS];
#pragma unroll
for (int h = 0; h < BCS_N_HORIZONS; ++h) {
local_raw[h] = 0.0f;
local_cnt[h] = 0.0f;
}
int local_valid = 0;
// Pass 1 — per-thread strided accumulation. Same loop also zeros
// grad_probs at masked positions so the second pass can skip those.
for (int i = tid; i < total; i += BCS_BLOCK) {
const float y = labels[i];
if (isnan(y)) {
grad_probs[i] = 0.0f;
continue;
}
const int h = i % n_horizons;
const float p = fminf(fmaxf(probs[i], 1e-6f), 1.0f - 1e-6f);
const float L = -(y * logf(p) + (1.0f - y) * logf(1.0f - p));
local_raw[h] += L;
local_cnt[h] += 1.0f;
local_valid += 1;
}
// Warp-shuffle reduce each of the 5 horizon sums + the per-horizon
// count + the global valid count. Each lane in warp gets the full
// warp-sum at the end (no need to broadcast manually).
float warp_raw[BCS_N_HORIZONS];
float warp_cnt[BCS_N_HORIZONS];
#pragma unroll
for (int h = 0; h < BCS_N_HORIZONS; ++h) {
warp_raw[h] = warp_reduce_sum(local_raw[h]);
warp_cnt[h] = warp_reduce_sum(local_cnt[h]);
}
int warp_valid = warp_reduce_sum_int(local_valid);
// Cross-warp reduce — only lane 0 of each warp writes its partial.
__shared__ float s_raw[BCS_N_WARPS * BCS_N_HORIZONS];
__shared__ float s_cnt[BCS_N_WARPS * BCS_N_HORIZONS];
__shared__ int s_valid[BCS_N_WARPS];
if (lane == 0) {
#pragma unroll
for (int h = 0; h < BCS_N_HORIZONS; ++h) {
s_raw[warp * BCS_N_HORIZONS + h] = warp_raw[h];
s_cnt[warp * BCS_N_HORIZONS + h] = warp_cnt[h];
}
s_valid[warp] = warp_valid;
}
__syncthreads();
// Warp 0 finishes: reduce N_WARPS partials. All 32 lanes participate;
// inactive lanes (tid ≥ N_WARPS) contribute 0 via ternary (NOT a
// divergent shuffle).
__shared__ float mean_bce_h_shared[BCS_N_HORIZONS];
__shared__ float w_eff_shared[BCS_N_HORIZONS];
__shared__ int valid_count_shared;
if (warp == 0) {
#pragma unroll
for (int h = 0; h < BCS_N_HORIZONS; ++h) {
float v_raw = (lane < BCS_N_WARPS) ? s_raw[lane * BCS_N_HORIZONS + h] : 0.0f;
float v_cnt = (lane < BCS_N_WARPS) ? s_cnt[lane * BCS_N_HORIZONS + h] : 0.0f;
v_raw = warp_reduce_sum(v_raw);
v_cnt = warp_reduce_sum(v_cnt);
if (lane == 0) {
const float c_h = v_cnt;
const float mean_bce = (c_h > 0.0f) ? v_raw / c_h : 0.0f;
mean_bce_h_shared[h] = mean_bce;
mean_bce_per_h_out[h] = mean_bce;
const float ls = log_sigma_h[h];
const float bw = (base_weights != nullptr) ? base_weights[h] : 1.0f;
// w_h = bw / (2 * exp(2*ls)) = 0.5 * bw * exp(-2 * ls)
const float w_h = 0.5f * bw * expf(-2.0f * ls);
// Per-element grad prefactor = w_h / count_h.
w_eff_shared[h] = (c_h > 0.0f) ? w_h / c_h : 0.0f;
// ∂L/∂log_sigma_h = 1 − 2 * w_h * mean_bce_h.
d_log_sigma_h[h] = 1.0f - 2.0f * w_h * mean_bce;
}
}
int v_valid = (lane < BCS_N_WARPS) ? s_valid[lane] : 0;
v_valid = warp_reduce_sum_int(v_valid);
if (lane == 0) {
valid_count_shared = v_valid;
valid_count_out[0] = v_valid;
}
}
__syncthreads();
// Total loss — small serial reduce by thread 0 over 5 horizons.
if (tid == 0) {
float total_loss = 0.0f;
#pragma unroll
for (int h = 0; h < BCS_N_HORIZONS; ++h) {
const float ls = log_sigma_h[h];
const float bw = (base_weights != nullptr) ? base_weights[h] : 1.0f;
const float w_h = 0.5f * bw * expf(-2.0f * ls);
total_loss += w_h * mean_bce_h_shared[h] + ls;
}
loss_out[0] = total_loss;
}
__syncthreads();
// Pass 2 — per-element grad. Coalesced strided access. w_eff_shared
// is read-only at this point, broadcast freely across threads.
for (int i = tid; i < total; i += BCS_BLOCK) {
const float y = labels[i];
if (isnan(y)) {
continue; // grad was already zeroed in pass 1
}
const int h = i % n_horizons;
const float p = fminf(fmaxf(probs[i], 1e-6f), 1.0f - 1e-6f);
grad_probs[i] = w_eff_shared[h] * (p - y) / (p * (1.0f - p));
}
}

41
crates/ml-alpha/src/trainer/loss.rs
View File

@@ -1,4 +1,8 @@
//! Fused multi-horizon BCE loss + gradient w.r.t. probs.
//! Multi-horizon BCE — standalone helper that exercises
//! `bce_multi_horizon_forward_backward`. Single source of truth: the
//! same kernel the trainer hot path calls.
//!
//! See `cuda/bce_loss_multi_horizon.cu` for Kendall σ math.
use std::sync::Arc;
@@ -15,12 +19,17 @@ pub struct BceInput {
pub labels: Vec<f32>, // NaN entries are masked out
pub n_horizons: usize,
pub n_pos: usize,
/// Optional per-horizon Kendall σ logarithm. Pass `None` → σ=0 →
/// each horizon contributes `0.5 · mean_bce_h + 0` to the loss.
pub log_sigma_h: Option<Vec<f32>>,
}
pub struct BceOutput {
pub loss: f32,
pub n_valid: i32,
pub mean_bce_per_h: Vec<f32>,
pub grad_probs: Vec<f32>,
pub d_log_sigma_h: Vec<f32>,
}
pub fn bce_multi_horizon_loss_and_grad_gpu(dev: &MlDevice, input: &BceInput) -> Result<BceOutput> {
@@ -31,18 +40,27 @@ pub fn bce_multi_horizon_loss_and_grad_gpu(dev: &MlDevice, input: &BceInput) ->
let stream: &Arc<CudaStream> = dev.cuda_stream().context("bce stream")?;
let ctx = dev.cuda_context().context("bce ctx")?;
let module = ctx.load_cubin(CUBIN.to_vec()).context("load bce cubin")?;
let func = module.load_function("bce_multi_horizon_forward_backward").context("bce fn")?;
let func = module
.load_function("bce_multi_horizon_forward_backward")
.context("load bce_multi_horizon_forward_backward")?;
let probs_d = upload(stream, &input.probs)?;
let labels_d = upload(stream, &input.labels)?;
// Standalone helper invokes the kernel with uniform (all-ones)
// per-horizon weights — matches the production trainer's
// `horizon_weights: [1.0; 5]` default semantics.
// Default base weights = 1 per horizon (Kendall σ does the rebalancing).
let weights_host: Vec<f32> = vec![1.0; input.n_horizons];
let weights_d = upload(stream, &weights_host)?;
// σ default: zeros (passthrough Kendall init).
let sigma_host: Vec<f32> = input
.log_sigma_h
.clone()
.unwrap_or_else(|| vec![0.0; input.n_horizons]);
let sigma_d = upload(stream, &sigma_host)?;
let mut loss_d = stream.alloc_zeros::<f32>(1).context("loss alloc")?;
let mut mean_bce_h_d = stream.alloc_zeros::<f32>(input.n_horizons)?;
let mut grad_d = stream.alloc_zeros::<f32>(total).context("grad alloc")?;
let mut valid_d = stream.alloc_zeros::<i32>(1).context("valid alloc")?;
let mut d_sigma_d = stream.alloc_zeros::<f32>(input.n_horizons)?;
let n_pos_i = input.n_pos as i32;
let n_horizons_i = input.n_horizons as i32;
@@ -53,20 +71,25 @@ pub fn bce_multi_horizon_loss_and_grad_gpu(dev: &MlDevice, input: &BceInput) ->
};
let mut launch = stream.launch_builder(&func);
launch
.arg(&probs_d).arg(&labels_d).arg(&weights_d)
.arg(&probs_d).arg(&labels_d).arg(&weights_d).arg(&sigma_d)
.arg(&n_pos_i).arg(&n_horizons_i)
.arg(&mut loss_d).arg(&mut grad_d).arg(&mut valid_d);
.arg(&mut loss_d).arg(&mut mean_bce_h_d).arg(&mut grad_d).arg(&mut valid_d)
.arg(&mut d_sigma_d);
unsafe { launch.launch(cfg).context("bce launch")?; }
stream.synchronize().context("bce sync")?;
let loss = download_f32(stream, &loss_d)?[0];
let mean_bce_per_h = download_f32(stream, &mean_bce_h_d)?;
let grad = download_f32(stream, &grad_d)?;
let valid = download_i32(stream, &valid_d)?[0];
let d_log_sigma_h = download_f32(stream, &d_sigma_d)?;
Ok(BceOutput {
loss,
n_valid: valid,
mean_bce_per_h,
grad_probs: grad,
d_log_sigma_h,
})
}
@@ -102,9 +125,6 @@ fn download_f32(stream: &Arc<CudaStream>, src: &CudaSlice<f32>) -> Result<Vec<f3
}
fn download_i32(stream: &Arc<CudaStream>, src: &CudaSlice<i32>) -> Result<Vec<i32>> {
// i32 path through a transient host vec; mapped-pinned i32 buffer
// lives in the parent ml crate which we can't import (cycle). For
// a single i32 readback this is acceptable slow-path cost.
let n = src.len();
let staging = unsafe { MappedF32Buffer::new(n) }
.map_err(|e| anyhow::anyhow!("bce download i32: {e}"))?;
@@ -115,7 +135,6 @@ fn download_i32(stream: &Arc<CudaStream>, src: &CudaSlice<i32>) -> Result<Vec<i3
.context("bce i32 download DtoD")?;
}
stream.synchronize().context("bce i32 download sync")?;
// Reinterpret f32 buffer as i32.
let host_ptr = staging.host_ptr.cast::<i32>();
let mut out = Vec::with_capacity(n);
unsafe {

166
crates/ml-alpha/src/trainer/loss_sigma.rs
View File

@@ -1,166 +0,0 @@
//! Kendall σ-weighted multi-horizon BCE (v2 axis A).
//!
//! Standalone helper that exercises `bce_loss_multi_horizon_sigma.cu`.
//! Used by the numgrad parity test in
//! `tests/bce_sigma_numgrad.rs`. The hot-path trainer integration
//! comes in commit V10.
//!
//! See `docs/superpowers/specs/2026-05-18-ml-alpha-v2-multi-horizon-design.md`
//! §3.4 for the math and motivation.
use std::sync::Arc;
use anyhow::{Context, Result};
use cudarc::driver::{
CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtr, DevicePtrMut,
LaunchConfig, PushKernelArg,
};
use ml_core::device::MlDevice;
use crate::pinned_mem::MappedF32Buffer;
const SIGMA_CUBIN: &[u8] =
include_bytes!(concat!(env!("OUT_DIR"), "/bce_loss_multi_horizon_sigma.cubin"));
pub struct BceSigmaInput {
pub probs: Vec<f32>,
pub labels: Vec<f32>, // NaN entries are masked out
pub base_weights: Vec<f32>, // [n_horizons]
pub log_sigma_h: Vec<f32>, // [n_horizons]
pub n_horizons: usize,
pub n_pos: usize,
}
pub struct BceSigmaOutput {
pub total_loss: f32,
pub n_valid: i32,
pub mean_bce_per_h: Vec<f32>,
pub grad_probs: Vec<f32>,
pub d_log_sigma_h: Vec<f32>,
}
/// Loads the σ-weighted BCE kernel from the embedded cubin.
pub struct BceSigmaKernel {
_module: Arc<CudaModule>,
pub func: CudaFunction,
}
impl BceSigmaKernel {
pub fn new(ctx: &Arc<CudaContext>) -> Result<Self> {
let module = ctx
.load_cubin(SIGMA_CUBIN.to_vec())
.context("load bce_sigma cubin")?;
let func = module
.load_function("bce_multi_horizon_sigma_forward_backward")
.context("load bce_sigma fn")?;
Ok(Self { _module: module, func })
}
}
/// Standalone helper: upload inputs, launch kernel, download outputs.
/// **Not** used in the trainer hot path (that's a V10 wiring task) — this
/// is the numgrad-test entry point.
pub fn bce_sigma_loss_and_grad_gpu(
dev: &MlDevice,
input: &BceSigmaInput,
) -> Result<BceSigmaOutput> {
let total = input.n_pos * input.n_horizons;
anyhow::ensure!(input.probs.len() == total, "probs len mismatch");
anyhow::ensure!(input.labels.len() == total, "labels len mismatch");
anyhow::ensure!(input.base_weights.len() == input.n_horizons, "base_weights len");
anyhow::ensure!(input.log_sigma_h.len() == input.n_horizons, "log_sigma_h len");
let stream: &Arc<CudaStream> = dev.cuda_stream().context("bce_sigma stream")?;
let ctx = dev.cuda_context().context("bce_sigma ctx")?;
let kernel = BceSigmaKernel::new(ctx)?;
let probs_d = upload(stream, &input.probs)?;
let labels_d = upload(stream, &input.labels)?;
let base_weights_d = upload(stream, &input.base_weights)?;
let log_sigma_d = upload(stream, &input.log_sigma_h)?;
let mut loss_d = stream.alloc_zeros::<f32>(1).context("loss alloc")?;
let mut mean_bce_h_d = stream.alloc_zeros::<f32>(input.n_horizons)?;
let mut grad_d = stream.alloc_zeros::<f32>(total).context("grad alloc")?;
let mut valid_d = stream.alloc_zeros::<i32>(1).context("valid alloc")?;
let mut d_log_sigma_d = stream.alloc_zeros::<f32>(input.n_horizons)?;
let n_pos_i = input.n_pos as i32;
let n_horizons_i = input.n_horizons as i32;
let cfg = LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
};
let mut launch = stream.launch_builder(&kernel.func);
launch
.arg(&probs_d).arg(&labels_d).arg(&base_weights_d).arg(&log_sigma_d)
.arg(&n_pos_i).arg(&n_horizons_i)
.arg(&mut loss_d).arg(&mut mean_bce_h_d).arg(&mut grad_d).arg(&mut valid_d)
.arg(&mut d_log_sigma_d);
unsafe { launch.launch(cfg).context("bce_sigma launch")?; }
stream.synchronize().context("bce_sigma sync")?;
Ok(BceSigmaOutput {
total_loss: download_f32(stream, &loss_d)?[0],
n_valid: download_i32(stream, &valid_d)?[0],
mean_bce_per_h: download_f32(stream, &mean_bce_h_d)?,
grad_probs: download_f32(stream, &grad_d)?,
d_log_sigma_h: download_f32(stream, &d_log_sigma_d)?,
})
}
fn upload(stream: &Arc<CudaStream>, host: &[f32]) -> Result<CudaSlice<f32>> {
let n = host.len();
let staging = unsafe { MappedF32Buffer::new(n) }
.map_err(|e| anyhow::anyhow!("bce_sigma upload staging: {e}"))?;
staging.write_from_slice(host);
let mut dst = stream.alloc_zeros::<f32>(n).context("bce_sigma upload alloc")?;
if n > 0 {
let nbytes = n * std::mem::size_of::<f32>();
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("bce_sigma upload DtoD")?;
}
}
Ok(dst)
}
fn download_f32(stream: &Arc<CudaStream>, src: &CudaSlice<f32>) -> Result<Vec<f32>> {
let n = src.len();
let staging = unsafe { MappedF32Buffer::new(n) }
.map_err(|e| anyhow::anyhow!("bce_sigma download f32: {e}"))?;
let nbytes = n * std::mem::size_of::<f32>();
unsafe {
let (src_ptr, _g) = src.device_ptr(stream);
cudarc::driver::result::memcpy_dtod_async(
staging.dev_ptr, src_ptr, nbytes, stream.cu_stream(),
).context("bce_sigma download DtoD")?;
}
stream.synchronize().context("bce_sigma download sync")?;
Ok(staging.read_all())
}
fn download_i32(stream: &Arc<CudaStream>, src: &CudaSlice<i32>) -> Result<Vec<i32>> {
let n = src.len();
let staging = unsafe { MappedF32Buffer::new(n) }
.map_err(|e| anyhow::anyhow!("bce_sigma download i32: {e}"))?;
let nbytes = n * std::mem::size_of::<i32>();
unsafe {
let (src_ptr, _g) = src.device_ptr(stream);
cudarc::driver::result::memcpy_dtod_async(
staging.dev_ptr, src_ptr, nbytes, stream.cu_stream(),
).context("bce_sigma i32 download DtoD")?;
}
stream.synchronize().context("bce_sigma i32 download sync")?;
let host_ptr = staging.host_ptr.cast::<i32>();
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)
}

3
crates/ml-alpha/src/trainer/mod.rs
View File

@@ -3,7 +3,6 @@
pub mod anchor_controller;
pub mod loss;
pub mod loss_sigma;
pub mod multi_horizon_attention;
pub mod optim;
pub mod perception;
pub mod multi_horizon_attention;

41
crates/ml-alpha/src/trainer/perception.rs
View File

@@ -375,6 +375,14 @@ pub struct PerceptionTrainer {
attn_bwd_fn: CudaFunction,
_attn_module: Arc<CudaModule>,
/// MultiHorizonAttention bundle owns the v2 attention path
/// (horizon-token + inverted + MoE), the Kendall σ logarithm for
/// BCE, the anchor controller, and all corresponding optimizers.
/// At present `log_sigma_h_d` and `grad_log_sigma_h_d` are
/// consumed by the BCE callsite; full forward/backward replacement
/// of the legacy `attn_*` path lands in the Stage 2 commit.
pub mha: crate::trainer::multi_horizon_attention::MultiHorizonAttention,
// ── K-loop parallelization (Phase B) ──
// Per-batch grad scratch buffers for cfc_step_backward_batched.
// Zeroed once per training step; the K-loop's 64 bwd calls accumulate
@@ -803,13 +811,20 @@ impl PerceptionTrainer {
// Phase B: attn pool per-batch grad scratch.
let attn_grad_q_scratch_d = stream.alloc_zeros::<f32>(cfg.n_batch * HIDDEN_DIM)?;
// 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.
// 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),
)?;
let k = cfg.seq_len;
Ok(Self {
@@ -866,6 +881,7 @@ impl PerceptionTrainer {
attn_fwd_fn,
attn_bwd_fn,
_attn_module: attn_module,
mha,
// Phase B: cfc per-batch grad scratch + reducer.
cfc_grad_w_in_scratch_d,
cfc_grad_w_rec_scratch_d,
@@ -1670,17 +1686,16 @@ 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(&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.valid_d)
.arg(&mut self.mha.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
@@ -2586,10 +2601,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(&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.grad_probs_per_k_d).arg(&mut self.valid_d)
.arg(&mut self.mha.grad_log_sigma_h_d);
launch.launch(bce_cfg).context("eval bce")?;
}
self.stream.synchronize().context("eval sync")?;

2
crates/ml-alpha/tests/bce_grad_finite_diff.rs
View File

@@ -14,7 +14,7 @@ fn test_device() -> MlDevice {
}
fn make_input(probs: Vec<f32>, labels: Vec<f32>) -> BceInput {
BceInput { probs, labels, n_horizons: 5, n_pos: 4 }
BceInput { probs, labels, n_horizons: 5, n_pos: 4, log_sigma_h: None }
}
#[test]

166
crates/ml-alpha/tests/bce_sigma_numgrad.rs
View File

@@ -1,166 +0,0 @@
//! Numgrad parity test for the Kendall σ-weighted multi-horizon BCE
//! kernel (`bce_multi_horizon_sigma_forward_backward`, v2 axis A).
//!
//! Verifies:
//! 1. Analytical d_log_sigma_h matches central-difference numgrad
//! within 5e-2 rel-tol / 5e-3 abs-floor.
//! 2. Analytical grad_probs matches central-difference numgrad on
//! a random subset of positions.
//! 3. Total loss equals the closed-form Σ_h (w_h · mean_bce_h + log σ_h)
//! reconstructed from `mean_bce_per_h`.
#![cfg(feature = "cuda")]
use anyhow::Result;
use ml_alpha::trainer::loss_sigma::{bce_sigma_loss_and_grad_gpu, BceSigmaInput};
use ml_core::device::MlDevice;
const N_HORIZONS: usize = 5;
const N_POS: usize = 24;
fn make_inputs(seed: u64) -> BceSigmaInput {
// Deterministic linear-congruential mock RNG; avoids a rand dep.
let mut s = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15);
let mut nxt = || -> f32 {
s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
let x = ((s >> 16) & 0xFFFF) as f32 / 65536.0;
x
};
let total = N_POS * N_HORIZONS;
let mut probs = Vec::with_capacity(total);
let mut labels = Vec::with_capacity(total);
for _ in 0..total {
let p = 0.05_f32 + 0.9 * nxt(); // (0.05, 0.95)
let y = if nxt() < 0.5 { 0.0_f32 } else { 1.0_f32 };
probs.push(p);
labels.push(y);
}
// Non-uniform base weights, non-uniform log_sigma.
let base_weights: Vec<f32> = (0..N_HORIZONS)
.map(|i| 0.5 + i as f32 * 0.25)
.collect();
let log_sigma_h: Vec<f32> = (0..N_HORIZONS)
.map(|i| -0.4 + i as f32 * 0.15)
.collect();
BceSigmaInput {
probs,
labels,
base_weights,
log_sigma_h,
n_horizons: N_HORIZONS,
n_pos: N_POS,
}
}
fn forward_only_loss(dev: &MlDevice, input: &BceSigmaInput) -> Result<f32> {
Ok(bce_sigma_loss_and_grad_gpu(dev, input)?.total_loss)
}
#[test]
#[ignore = "requires CUDA"]
fn d_log_sigma_matches_central_difference() -> Result<()> {
let dev = MlDevice::cuda(0)?;
let input = make_inputs(20260518);
let out = bce_sigma_loss_and_grad_gpu(&dev, &input)?;
assert!(out.total_loss.is_finite(), "loss not finite: {}", out.total_loss);
let eps = 1e-3_f32;
for h in 0..N_HORIZONS {
let mut input_plus = input.clone_for_numgrad();
let mut input_minus = input.clone_for_numgrad();
input_plus.log_sigma_h[h] += eps;
input_minus.log_sigma_h[h] -= eps;
let l_plus = forward_only_loss(&dev, &input_plus)?;
let l_minus = forward_only_loss(&dev, &input_minus)?;
let numgrad = (l_plus - l_minus) / (2.0 * eps);
let analytic = out.d_log_sigma_h[h];
let abs_err = (analytic - numgrad).abs();
let rel_err = abs_err / numgrad.abs().max(1e-3);
assert!(
abs_err < 5e-3 || rel_err < 5e-2,
"d_log_sigma[{h}] mismatch: analytic={analytic:.6} numgrad={numgrad:.6} abs={abs_err:.3e} rel={rel_err:.3e}"
);
}
Ok(())
}
#[test]
#[ignore = "requires CUDA"]
fn grad_probs_matches_central_difference() -> Result<()> {
let dev = MlDevice::cuda(0)?;
let input = make_inputs(42);
let out = bce_sigma_loss_and_grad_gpu(&dev, &input)?;
assert!(out.total_loss.is_finite(), "loss not finite: {}", out.total_loss);
// Sample 8 random positions per the LCG.
let eps = 1e-3_f32;
for k in 0..8 {
let i = (k * 7919 + 3) % (N_POS * N_HORIZONS);
let mut input_plus = input.clone_for_numgrad();
let mut input_minus = input.clone_for_numgrad();
input_plus.probs[i] += eps;
input_minus.probs[i] -= eps;
let l_plus = forward_only_loss(&dev, &input_plus)?;
let l_minus = forward_only_loss(&dev, &input_minus)?;
let numgrad = (l_plus - l_minus) / (2.0 * eps);
let analytic = out.grad_probs[i];
let abs_err = (analytic - numgrad).abs();
let rel_err = abs_err / numgrad.abs().max(1e-3);
assert!(
abs_err < 5e-3 || rel_err < 5e-2,
"grad_probs[{i}] mismatch: analytic={analytic:.6} numgrad={numgrad:.6} abs={abs_err:.3e} rel={rel_err:.3e}"
);
}
Ok(())
}
#[test]
#[ignore = "requires CUDA"]
fn total_loss_equals_closed_form() -> Result<()> {
let dev = MlDevice::cuda(0)?;
let input = make_inputs(123);
let out = bce_sigma_loss_and_grad_gpu(&dev, &input)?;
// Closed-form Σ_h (w_h · mean_bce_h + log σ_h)
let mut expected = 0.0_f64;
for h in 0..N_HORIZONS {
let ls = input.log_sigma_h[h] as f64;
let bw = input.base_weights[h] as f64;
let w_h = 0.5 * bw * (-2.0 * ls).exp();
let mean = out.mean_bce_per_h[h] as f64;
expected += w_h * mean + ls;
}
let abs_err = (out.total_loss as f64 - expected).abs();
let rel_err = abs_err / expected.abs().max(1e-3);
assert!(
abs_err < 1e-3 || rel_err < 5e-3,
"total_loss mismatch: kernel={} closed_form={expected} abs={abs_err:.3e} rel={rel_err:.3e}",
out.total_loss
);
Ok(())
}
// Small helper trait — clones the input by hand because BceSigmaInput
// doesn't implement Clone (deliberately to keep the test crate's
// dependency surface clean).
trait CloneForNumgrad {
fn clone_for_numgrad(&self) -> Self;
}
impl CloneForNumgrad for BceSigmaInput {
fn clone_for_numgrad(&self) -> Self {
Self {
probs: self.probs.clone(),
labels: self.labels.clone(),
base_weights: self.base_weights.clone(),
log_sigma_h: self.log_sigma_h.clone(),
n_horizons: self.n_horizons,
n_pos: self.n_pos,
}
}
}