Files
foxhunt/crates/ml-alpha/src/lib.rs
jgrusewski d94696620d feat(ml-alpha): inverted_attention_pool kernel + numgrad (v2 E) [V3]
iTransformer-style cross-variate attention pass: each of HIDDEN_DIM
features becomes a "variate token" with its K-trajectory as embedding.

FORWARD:
  X_inv[h, k]    = ln_out[k, h]                       # transpose
  scores[h, j]   = inv_scale · Σ_k X_inv[h, k] · X_inv[j, k]
  attn[h, j]     = softmax_j(scores)
  pooled[h]      = Σ_j attn[h, j] · mean_k_X_inv[j]   # mean-K commutes out

BACKWARD: three independent chains into ln_out:
  - value path:  (1/K) · attn[h', my_h] · d_pooled[h']  (per-k constant)
  - query path:  inv_scale · Σ_j d_scores[my_h, j] · X_inv[j, k]
  - key path:    inv_scale · Σ_h' d_scores[h', my_h] · X_inv[h', k]
  Softmax bwd: d_scores[h, j] = attn · (d_attn - Σ_l attn · d_attn)

IMPLEMENTATION NOTES:
  - First attempt cached attn [H, H] = 64 KB in shared mem → tripped
    the 48 KB dynamic-shared limit on sm_86 (CUDA_ERROR_INVALID_VALUE).
  - Fixed by moving d_scores to a DRAM scratch buffer; shared mem
    holds only X_inv [H, K] (≤ 16 KB at K = 32). One block-wide barrier
    between the d_scores write and the value/query/key accumulation.
  - All per-batch slice writes; no atomicAdd, no cross-block race.
  - Pooled computation uses the mean-K commute (Σ_k attn · X_inv =
    attn · mean_k_X_inv), saving an entire H×K accumulation pass.

LOCAL VERIFICATION (RTX 3050 sm_86):
  forward_then_backward_matches_central_difference PASSES 6 numgrad
  checks on ln_out positions within 5e-2 rel / 5e-3 abs.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-05-18 14:16:26 +02:00

53 lines
1.6 KiB
Rust

//! ml-alpha — CfC perception + multi-horizon alpha heads.
//!
//! Phase A (this crate): snapshot-level CfC trunk + 5 horizon heads,
//! trained supervised on MBP-10 with multi-horizon BCE. The CfC trunk
//! must beat the Mamba2 baseline (`mamba2_block`) at every horizon
//! before Phase B (PPO) is allowed to start.
//!
//! All hot-path arithmetic runs on GPU. The crate owns no CPU compute
//! on model tensors. See `docs/superpowers/specs/2026-05-16-ml-alpha-cfc-ppo-design.md`.
#![warn(clippy::all)]
#![allow(
clippy::module_name_repetitions,
clippy::cast_precision_loss,
clippy::cast_possible_truncation,
clippy::missing_errors_doc
)]
// Module declarations are added as each task lands its source.
// Tasks 3-17 progressively turn these on.
// Task 3: pub mod isv;
// Task 4: pub mod pinned;
// Task 5+: pub mod cfc;
// Task 7+: pub mod heads;
// Task 10+: pub mod trainer;
// Task 12: pub mod data;
// Task 16+: pub mod gate;
pub mod cfc;
pub mod data;
pub mod eval;
pub mod heads;
pub mod horizon_token_attention_pool;
pub mod inverted_attention_pool;
pub mod isv;
pub mod pinned;
pub mod pinned_mem;
pub mod trainer;
// Preserved from prior crate state — used by Phase A.
pub mod fxcache_reader;
pub mod purged_split;
pub mod multi_horizon_labels;
// Gate reference only — Mamba2 baseline against which CfC must compete.
pub mod mamba2_block;
pub use isv::IsvBus;
pub use multi_horizon_labels::{generate_labels, LongHorizonLabels};
// Re-exports added as the relevant tasks land:
// pub use cfc::CfcTrunk;
// pub use heads::{MultiHorizonHeads, Projection};
// pub use pinned::{MappedPinnedSnapshotSlot, MappedPinnedFillSlot};