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>
53 lines
1.6 KiB
Rust
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};
|