feat(ml-alpha): multi_horizon_heads kernel (128->5 sigmoid)
Per-horizon P(up) at h ∈ {30, 100, 300, 1000, 6000} snapshots forward.
Single-block 5-thread kernel; each thread is its own 128-dim dot
product + sigmoid. No atomicAdd.
Tests (5/5 pass on sm_86) assert invariants only:
- sigmoid output ∈ [0, 1] for all heads
- zero weights + zero bias → 0.5 exactly
- bias = +20 → saturates near 1
- bias = -20 → saturates near 0
- per-head independence (mixed-bias configuration)
Addendum updated to explicitly state no-CPU-mirror discipline per
feedback_no_cpu_test_fallbacks.md.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -1,2 +1,21 @@
|
||||
// placeholder — real implementation in the task adding this kernel
|
||||
extern "C" __global__ void multi_horizon_heads_stub() {}
|
||||
// multi_horizon_heads.cu
|
||||
//
|
||||
// 128-dim hidden -> 5 sigmoid logits, one per horizon
|
||||
// {30, 100, 300, 1000, 6000} snapshots forward. Single-block kernel,
|
||||
// 5 threads (one per head). Each thread computes its own dot product
|
||||
// against the 128-dim hidden vector + bias, then sigmoid.
|
||||
|
||||
extern "C" __global__ void multi_horizon_heads(
|
||||
const float* __restrict__ w, // [5, 128]
|
||||
const float* __restrict__ b, // [5]
|
||||
const float* __restrict__ h, // [128]
|
||||
float* __restrict__ probs // [5]
|
||||
) {
|
||||
int k = threadIdx.x;
|
||||
if (k >= 5) return;
|
||||
float z = b[k];
|
||||
for (int i = 0; i < 128; ++i) {
|
||||
z += w[k * 128 + i] * h[i];
|
||||
}
|
||||
probs[k] = 1.0f / (1.0f + expf(-z));
|
||||
}
|
||||
|
||||
104
crates/ml-alpha/src/heads.rs
Normal file
104
crates/ml-alpha/src/heads.rs
Normal file
@@ -0,0 +1,104 @@
|
||||
//! Multi-horizon sigmoid heads (128 → 5) and projection (128 → 8 + layer-norm).
|
||||
//!
|
||||
//! Per spec Section 2 + 2026-05-16 stacked amendment, the heads sit on
|
||||
//! the CfC hidden output and emit 5 sigmoid probabilities of "up" at
|
||||
//! horizons {30, 100, 300, 1000, 6000} snapshots forward.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg};
|
||||
use ml_core::device::MlDevice;
|
||||
|
||||
use crate::pinned_mem::MappedF32Buffer;
|
||||
|
||||
const HEADS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/multi_horizon_heads.cubin"));
|
||||
|
||||
pub const N_HORIZONS: usize = 5;
|
||||
pub const HORIZONS: [usize; N_HORIZONS] = [30, 100, 300, 1000, 6000];
|
||||
pub const HIDDEN_DIM: usize = 128;
|
||||
pub const PROJ_DIM: usize = 8;
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct HeadsWeights {
|
||||
pub w: Vec<f32>, // [5, 128] row-major
|
||||
pub b: Vec<f32>, // [5]
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct ProjectionWeights {
|
||||
pub w: Vec<f32>, // [8, 128]
|
||||
pub b: Vec<f32>, // [8]
|
||||
pub ln_gain: Vec<f32>, // [8]
|
||||
pub ln_bias: Vec<f32>, // [8]
|
||||
}
|
||||
|
||||
/// Placeholder marker — concrete trunk integration in Task 9.
|
||||
pub struct MultiHorizonHeads;
|
||||
|
||||
pub struct Projection;
|
||||
|
||||
pub fn multi_horizon_heads_gpu(
|
||||
dev: &MlDevice,
|
||||
w: &HeadsWeights,
|
||||
h: &[f32],
|
||||
) -> Result<[f32; N_HORIZONS]> {
|
||||
assert_eq!(h.len(), HIDDEN_DIM);
|
||||
assert_eq!(w.w.len(), N_HORIZONS * HIDDEN_DIM);
|
||||
assert_eq!(w.b.len(), N_HORIZONS);
|
||||
|
||||
let stream: &Arc<CudaStream> = dev.cuda_stream().context("heads stream")?;
|
||||
let ctx = dev.cuda_context().context("heads ctx")?;
|
||||
let module = ctx.load_cubin(HEADS_CUBIN.to_vec()).context("load heads cubin")?;
|
||||
let func = module.load_function("multi_horizon_heads").context("load function")?;
|
||||
|
||||
let w_d = upload(stream, &w.w)?;
|
||||
let b_d = upload(stream, &w.b)?;
|
||||
let h_d = upload(stream, h)?;
|
||||
let mut probs_d = stream.alloc_zeros::<f32>(N_HORIZONS).context("probs alloc")?;
|
||||
|
||||
let cfg = LaunchConfig {
|
||||
grid_dim: (1, 1, 1),
|
||||
block_dim: (N_HORIZONS as u32, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
let mut launch = stream.launch_builder(&func);
|
||||
launch.arg(&w_d).arg(&b_d).arg(&h_d).arg(&mut probs_d);
|
||||
unsafe { launch.launch(cfg).context("heads launch")?; }
|
||||
|
||||
let v = download(stream, &probs_d)?;
|
||||
let mut out = [0f32; N_HORIZONS];
|
||||
out.copy_from_slice(&v);
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
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!("heads upload staging: {e}"))?;
|
||||
staging.write_from_slice(host);
|
||||
let mut dst = stream.alloc_zeros::<f32>(n).context("heads 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("heads upload DtoD")?;
|
||||
}
|
||||
}
|
||||
Ok(dst)
|
||||
}
|
||||
|
||||
fn download(stream: &Arc<CudaStream>, src: &CudaSlice<f32>) -> Result<Vec<f32>> {
|
||||
let n = src.len();
|
||||
let staging = unsafe { MappedF32Buffer::new(n) }
|
||||
.map_err(|e| anyhow::anyhow!("heads download staging: {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("heads download DtoD")?;
|
||||
}
|
||||
stream.synchronize().context("heads download sync")?;
|
||||
Ok(staging.read_all())
|
||||
}
|
||||
@@ -26,6 +26,7 @@
|
||||
// Task 12: pub mod data;
|
||||
// Task 16+: pub mod gate;
|
||||
pub mod cfc;
|
||||
pub mod heads;
|
||||
pub mod isv;
|
||||
pub mod pinned;
|
||||
pub mod pinned_mem;
|
||||
|
||||
87
crates/ml-alpha/tests/heads_bit_equiv.rs
Normal file
87
crates/ml-alpha/tests/heads_bit_equiv.rs
Normal file
@@ -0,0 +1,87 @@
|
||||
//! multi_horizon_heads GPU invariants.
|
||||
//!
|
||||
//! Per `feedback_no_cpu_test_fallbacks.md`: validation via analytically-
|
||||
//! known synthetic inputs + property assertions.
|
||||
|
||||
use approx::assert_relative_eq;
|
||||
use ml_alpha::heads::{multi_horizon_heads_gpu, HeadsWeights, HIDDEN_DIM, N_HORIZONS};
|
||||
use ml_core::device::MlDevice;
|
||||
|
||||
fn test_device() -> MlDevice {
|
||||
MlDevice::cuda(0).expect("CUDA 0 required for ml-alpha tests")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn probs_are_in_unit_interval() {
|
||||
let dev = test_device();
|
||||
let w = HeadsWeights {
|
||||
w: vec![0.1; N_HORIZONS * HIDDEN_DIM],
|
||||
b: vec![0.0; N_HORIZONS],
|
||||
};
|
||||
let h: Vec<f32> = (0..HIDDEN_DIM).map(|i| 0.01 * i as f32).collect();
|
||||
let probs = multi_horizon_heads_gpu(&dev, &w, &h).expect("gpu");
|
||||
for k in 0..N_HORIZONS {
|
||||
assert!((0.0..=1.0).contains(&probs[k]), "head {k} prob {} not in [0,1]", probs[k]);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn zero_weights_and_bias_give_half() {
|
||||
// logit = 0 -> sigmoid(0) = 0.5
|
||||
let dev = test_device();
|
||||
let w = HeadsWeights {
|
||||
w: vec![0.0; N_HORIZONS * HIDDEN_DIM],
|
||||
b: vec![0.0; N_HORIZONS],
|
||||
};
|
||||
let h: Vec<f32> = (0..HIDDEN_DIM).map(|i| i as f32).collect();
|
||||
let probs = multi_horizon_heads_gpu(&dev, &w, &h).expect("gpu");
|
||||
for k in 0..N_HORIZONS {
|
||||
assert_relative_eq!(probs[k], 0.5_f32, epsilon = 1e-6);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn large_positive_bias_saturates_high() {
|
||||
let dev = test_device();
|
||||
let w = HeadsWeights {
|
||||
w: vec![0.0; N_HORIZONS * HIDDEN_DIM],
|
||||
b: vec![20.0; N_HORIZONS],
|
||||
};
|
||||
let h = vec![0.0; HIDDEN_DIM];
|
||||
let probs = multi_horizon_heads_gpu(&dev, &w, &h).expect("gpu");
|
||||
for k in 0..N_HORIZONS {
|
||||
assert!(probs[k] > 0.999, "head {k} should saturate near 1 with bias=20, got {}", probs[k]);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn large_negative_bias_saturates_low() {
|
||||
let dev = test_device();
|
||||
let w = HeadsWeights {
|
||||
w: vec![0.0; N_HORIZONS * HIDDEN_DIM],
|
||||
b: vec![-20.0; N_HORIZONS],
|
||||
};
|
||||
let h = vec![0.0; HIDDEN_DIM];
|
||||
let probs = multi_horizon_heads_gpu(&dev, &w, &h).expect("gpu");
|
||||
for k in 0..N_HORIZONS {
|
||||
assert!(probs[k] < 0.001, "head {k} should saturate near 0 with bias=-20, got {}", probs[k]);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn per_head_independence() {
|
||||
// Set bias[0] = 5, bias[4] = -5, others = 0 with zero weights.
|
||||
// sigmoid(5) ≈ 0.993, sigmoid(0) = 0.5, sigmoid(-5) ≈ 0.0067.
|
||||
let dev = test_device();
|
||||
let w = HeadsWeights {
|
||||
w: vec![0.0; N_HORIZONS * HIDDEN_DIM],
|
||||
b: vec![5.0, 0.0, 0.0, 0.0, -5.0],
|
||||
};
|
||||
let h = vec![0.0; HIDDEN_DIM];
|
||||
let probs = multi_horizon_heads_gpu(&dev, &w, &h).expect("gpu");
|
||||
assert!(probs[0] > 0.99 && probs[0] < 1.0);
|
||||
assert_relative_eq!(probs[1], 0.5, epsilon = 1e-6);
|
||||
assert_relative_eq!(probs[2], 0.5, epsilon = 1e-6);
|
||||
assert_relative_eq!(probs[3], 0.5, epsilon = 1e-6);
|
||||
assert!(probs[4] > 0.0 && probs[4] < 0.01);
|
||||
}
|
||||
@@ -3,8 +3,31 @@
|
||||
> Companion to `2026-05-16-ml-alpha-phase-a.md`. When the plan's binding-code
|
||||
> blocks reference `dev.htod_copy`, `dev.alloc_zeros`, `dev.load_cubin_function`,
|
||||
> `func.launch_async`, or `dev.default_stream`, **use the patterns in this
|
||||
> addendum instead.** The plan's kernel `.cu` source, CPU oracles, tests, and
|
||||
> assertion thresholds are unchanged.
|
||||
> addendum instead.** The plan's kernel `.cu` source and assertion thresholds
|
||||
> are unchanged.
|
||||
|
||||
## Validation discipline override: no CPU mirrors
|
||||
|
||||
Plan 1 incorrectly listed a "CPU oracle + bit-equiv test" step for every kernel
|
||||
task. Per `feedback_no_cpu_test_fallbacks.md` — kernels must be validated via
|
||||
**invariants on the GPU output**, not CPU mirrors. The user re-confirmed this
|
||||
mid-execution.
|
||||
|
||||
**Canonical validation pattern for each kernel:**
|
||||
|
||||
1. Build analytically-known synthetic inputs (zeros, identities, magic constants
|
||||
with known kernel outputs).
|
||||
2. Run the GPU kernel.
|
||||
3. Assert algebraic invariants on the output: ranges (e.g. sigmoid outputs
|
||||
∈ [0, 1]), fixed-points (dt=0 → identity), monotonicity (input ↑ → output ↑),
|
||||
conservation (sum of softmax probs ≈ 1), bounded growth (output ≤ deterministic
|
||||
bound from input).
|
||||
4. Where forward + backward are paired (BCE, PPO loss), validate the gradient
|
||||
via **finite-difference on the GPU forward** at a small set of perturbed
|
||||
inputs — NOT against a CPU autograd.
|
||||
|
||||
Whenever Plan 1's task body says "CPU oracle" or "bit-equiv test", read it as
|
||||
"GPU-output invariant test."
|
||||
|
||||
This addendum exists because Plan 1 was written against an older cudarc
|
||||
device-centric API. The vendored cudarc 0.19 (`vendor/cudarc/`) moved
|
||||
|
||||
Reference in New Issue
Block a user