feat(ml-alpha): BCE multi-horizon + AdamW kernels

bce_loss_multi_horizon: fused forward+backward, block tree-reduce (no
atomicAdd), NaN labels masked (drop). Loss = mean over valid; grad =
(p-y)/(p(1-p)) scaled by 1/N_valid.

adamw_step: element-wise AdamW with weight decay; one thread per param.

Tests pass on sm_86:
  BCE (4/4): positive+finite loss, near-zero loss when probs match
    labels, analytic grad matches GPU-computed finite-difference at
    eps=1e-3 / max_relative=5e-2 across 5 perturbation points, NaN
    labels mask grad and contribute zero to loss/N_valid.
  AdamW (4/4): zero-grad moves param only by weight-decay, positive
    grad decreases param, step counter increments, repeated descent
    on grad=theta drives ‖θ‖ from 40 to <1 in 200 steps.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

crates/ml-alpha/cuda/adamw_step.cu

@@ -1,2 +1,40 @@
-// placeholder — real implementation in the task adding this kernel
-extern "C" __global__ void adamw_step_stub() {}
+// adamw_step.cu — element-wise AdamW update.
+//
+//   m_t = beta1 * m_{t-1} + (1 - beta1) * g
+//   v_t = beta2 * v_{t-1} + (1 - beta2) * g^2
+//   m_hat = m_t / (1 - beta1^step)
+//   v_hat = v_t / (1 - beta2^step)
+//   theta -= lr * (m_hat / (sqrt(v_hat) + eps) + wd * theta)
+//
+// One thread per parameter. Block tree-reduce not needed; this is pure
+// element-wise. No atomicAdd.
+
+extern "C" __global__ void adamw_step(
+    float* __restrict__ theta,
+    const float* __restrict__ grad,
+    float* __restrict__ m,
+    float* __restrict__ v,
+    int n_params,
+    float lr,
+    float beta1,
+    float beta2,
+    float eps,
+    float wd,
+    int step
+) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i >= n_params) return;
+
+    const float g = grad[i];
+    const float m_n = beta1 * m[i] + (1.0f - beta1) * g;
+    const float v_n = beta2 * v[i] + (1.0f - beta2) * g * g;
+    m[i] = m_n;
+    v[i] = v_n;
+
+    const float bc1 = 1.0f - powf(beta1, (float) step);
+    const float bc2 = 1.0f - powf(beta2, (float) step);
+    const float m_hat = m_n / fmaxf(bc1, 1e-12f);
+    const float v_hat = v_n / fmaxf(bc2, 1e-12f);
+
+    theta[i] -= lr * (m_hat / (sqrtf(v_hat) + eps) + wd * theta[i]);
+}

crates/ml-alpha/cuda/bce_loss_multi_horizon.cu

@@ -1,2 +1,85 @@
-// placeholder — real implementation in the task adding this kernel
-extern "C" __global__ void bce_loss_multi_horizon_stub() {}
+// bce_loss_multi_horizon.cu
+//
+// Fused multi-horizon BCE forward + backward.
+//
+// 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)
+// Output:
+//   loss_out[1]              — mean BCE over valid entries
+//   grad_probs[n_pos, n_horizons] — d loss / d prob (scaled by 1/N_valid)
+//
+//   per element:
+//     p = clamp(probs[i], 1e-6, 1 - 1e-6)
+//     y = labels[i]
+//     L_i = -[y log p + (1-y) log(1-p)]
+//     dL/dp = (p - y) / (p (1-p))
+//
+//   total_loss = (1/N_valid) * sum_i L_i
+//   grad_probs[i] = (p - y) / (p (1-p) * N_valid)
+//
+// Block tree-reduce over a single block. n_pos * n_horizons must fit
+// within the block's per-thread iteration count (this kernel is for
+// minibatches up to a few thousand elements per call — well within
+// reach for n_pos = 256 * seq_len = 16384 elements). For larger
+// batches, the caller chunks.
+
+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]
+    int n_pos,
+    int n_horizons,
+    float* __restrict__ loss_out,       // [1]
+    float* __restrict__ grad_probs,     // [n_pos * n_horizons]
+    int* __restrict__ valid_count_out   // [1] — number of non-NaN labels
+) {
+    const int tid = threadIdx.x;
+    const int total = n_pos * n_horizons;
+
+    __shared__ float sloss[256];
+    __shared__ int   svalid[256];
+    sloss[tid] = 0.0f;
+    svalid[tid] = 0;
+
+    // First pass: count valid entries and accumulate loss numerator.
+    for (int i = tid; i < total; i += blockDim.x) {
+        const float y = labels[i];
+        if (isnan(y)) {
+            grad_probs[i] = 0.0f;
+            continue;
+        }
+        const float p = fminf(fmaxf(probs[i], 1e-6f), 1.0f - 1e-6f);
+        sloss[tid] -= y * logf(p) + (1.0f - y) * logf(1.0f - p);
+        svalid[tid] += 1;
+    }
+    __syncthreads();
+
+    // Block tree-reduce on sloss and svalid.
+    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
+        if (tid < s) {
+            sloss[tid] += sloss[tid + s];
+            svalid[tid] += svalid[tid + s];
+        }
+        __syncthreads();
+    }
+
+    if (tid == 0) {
+        const int n_valid = svalid[0];
+        loss_out[0] = (n_valid > 0) ? sloss[0] / (float) n_valid : 0.0f;
+        valid_count_out[0] = n_valid;
+    }
+    __syncthreads();
+
+    // Second pass: emit grad scaled by 1/n_valid. Uses the published
+    // valid_count_out[0] (already written by thread 0 above).
+    const int n_valid = valid_count_out[0];
+    const float scale = (n_valid > 0) ? (1.0f / (float) n_valid) : 0.0f;
+    for (int i = tid; i < total; i += blockDim.x) {
+        const float y = labels[i];
+        if (isnan(y)) {
+            continue;
+        }
+        const float p = fminf(fmaxf(probs[i], 1e-6f), 1.0f - 1e-6f);
+        grad_probs[i] = (p - y) / (p * (1.0f - p)) * scale;
+    }
+}

crates/ml-alpha/src/lib.rs

@@ -30,6 +30,7 @@ pub mod heads;
 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;

crates/ml-alpha/src/trainer/loss.rs

@@ -0,0 +1,122 @@
+//! Fused multi-horizon BCE loss + gradient w.r.t. probs.
+
+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 CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/bce_loss_multi_horizon.cubin"));
+
+pub struct BceInput {
+    pub probs: Vec<f32>,
+    pub labels: Vec<f32>, // NaN entries are masked out
+    pub n_horizons: usize,
+    pub n_pos: usize,
+}
+
+pub struct BceOutput {
+    pub loss: f32,
+    pub n_valid: i32,
+    pub grad_probs: Vec<f32>,
+}
+
+pub fn bce_multi_horizon_loss_and_grad_gpu(dev: &MlDevice, input: &BceInput) -> Result<BceOutput> {
+    let total = input.n_pos * input.n_horizons;
+    assert_eq!(input.probs.len(), total);
+    assert_eq!(input.labels.len(), total);
+
+    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 probs_d = upload(stream, &input.probs)?;
+    let labels_d = upload(stream, &input.labels)?;
+    let mut loss_d = stream.alloc_zeros::<f32>(1).context("loss alloc")?;
+    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 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(&func);
+    launch
+        .arg(&probs_d).arg(&labels_d)
+        .arg(&n_pos_i).arg(&n_horizons_i)
+        .arg(&mut loss_d).arg(&mut grad_d).arg(&mut valid_d);
+    unsafe { launch.launch(cfg).context("bce launch")?; }
+
+    stream.synchronize().context("bce sync")?;
+    let loss = download_f32(stream, &loss_d)?[0];
+    let grad = download_f32(stream, &grad_d)?;
+    let valid = download_i32(stream, &valid_d)?[0];
+
+    Ok(BceOutput {
+        loss,
+        n_valid: valid,
+        grad_probs: grad,
+    })
+}
+
+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 upload staging: {e}"))?;
+    staging.write_from_slice(host);
+    let mut dst = stream.alloc_zeros::<f32>(n).context("bce 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 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 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 download DtoD")?;
+    }
+    stream.synchronize().context("bce download sync")?;
+    Ok(staging.read_all())
+}
+
+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}"))?;
+    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 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 {
+        for i in 0..n {
+            out.push(std::ptr::read_volatile(host_ptr.add(i)));
+        }
+    }
+    Ok(out)
+}

crates/ml-alpha/src/trainer/mod.rs

@@ -0,0 +1,7 @@
+//! Phase A trainer scaffolding.
+//!
+//! Task 10 lands loss + optimizer kernels. Tasks 13-14 land the trainer
+//! loop + real CfC backward.
+
+pub mod loss;
+pub mod optim;

crates/ml-alpha/src/trainer/optim.rs

@@ -0,0 +1,75 @@
+//! AdamW optimizer state + kernel binding.
+
+use std::sync::Arc;
+
+use anyhow::{Context, Result};
+use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg};
+use ml_core::device::MlDevice;
+
+const CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/adamw_step.cubin"));
+
+pub struct AdamW {
+    stream: Arc<CudaStream>,
+    _module: Arc<CudaModule>,
+    func: CudaFunction,
+    m: CudaSlice<f32>,
+    v: CudaSlice<f32>,
+    step: i32,
+    pub lr: f32,
+    pub beta1: f32,
+    pub beta2: f32,
+    pub eps: f32,
+    pub wd: f32,
+}
+
+impl AdamW {
+    pub fn new(dev: &MlDevice, n_params: usize, lr: f32) -> Result<Self> {
+        let stream = dev.cuda_stream().context("adamw stream")?.clone();
+        let ctx = dev.cuda_context().context("adamw ctx")?;
+        let module = ctx.load_cubin(CUBIN.to_vec()).context("load adamw cubin")?;
+        let func = module.load_function("adamw_step").context("adamw fn")?;
+        let m = stream.alloc_zeros::<f32>(n_params).context("m alloc")?;
+        let v = stream.alloc_zeros::<f32>(n_params).context("v alloc")?;
+        Ok(Self {
+            stream,
+            _module: module,
+            func,
+            m,
+            v,
+            step: 0,
+            lr,
+            beta1: 0.9,
+            beta2: 0.999,
+            eps: 1e-8,
+            wd: 1e-2,
+        })
+    }
+
+    pub fn step_count(&self) -> i32 {
+        self.step
+    }
+
+    pub fn step(&mut self, theta: &mut CudaSlice<f32>, grad: &CudaSlice<f32>) -> Result<()> {
+        let n = theta.len();
+        assert_eq!(grad.len(), n, "grad/theta size mismatch");
+        assert_eq!(self.m.len(), n, "m size mismatch");
+        assert_eq!(self.v.len(), n, "v size mismatch");
+        self.step += 1;
+        let n_params_i = n as i32;
+        let step_i = self.step;
+        let cfg = LaunchConfig {
+            grid_dim: (((n as u32) + 255) / 256, 1, 1),
+            block_dim: (256, 1, 1),
+            shared_mem_bytes: 0,
+        };
+        let mut launch = self.stream.launch_builder(&self.func);
+        launch
+            .arg(theta).arg(grad).arg(&mut self.m).arg(&mut self.v)
+            .arg(&n_params_i)
+            .arg(&self.lr).arg(&self.beta1).arg(&self.beta2).arg(&self.eps).arg(&self.wd)
+            .arg(&step_i);
+        unsafe { launch.launch(cfg).context("adamw launch")?; }
+        self.stream.synchronize().context("adamw sync")?;
+        Ok(())
+    }
+}

crates/ml-alpha/tests/adamw_invariants.rs

@@ -0,0 +1,106 @@
+//! AdamW invariants — element-wise correctness without a CPU mirror.
+
+use approx::assert_relative_eq;
+use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut};
+use ml_alpha::pinned_mem::MappedF32Buffer;
+use ml_alpha::trainer::optim::AdamW;
+use ml_core::device::MlDevice;
+use std::sync::Arc;
+
+fn test_device() -> MlDevice {
+    MlDevice::cuda(0).expect("CUDA 0 required for ml-alpha tests")
+}
+
+fn upload(stream: &Arc<CudaStream>, host: &[f32]) -> CudaSlice<f32> {
+    let n = host.len();
+    let staging = unsafe { MappedF32Buffer::new(n) }.unwrap();
+    staging.write_from_slice(host);
+    let mut dst = stream.alloc_zeros::<f32>(n).unwrap();
+    let nbytes = n * 4;
+    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()).unwrap();
+    }
+    dst
+}
+
+fn download(stream: &Arc<CudaStream>, src: &CudaSlice<f32>) -> Vec<f32> {
+    let n = src.len();
+    let staging = unsafe { MappedF32Buffer::new(n) }.unwrap();
+    let nbytes = n * 4;
+    unsafe {
+        let (src_ptr, _g) = src.device_ptr(stream);
+        cudarc::driver::result::memcpy_dtod_async(staging.dev_ptr, src_ptr, nbytes, stream.cu_stream()).unwrap();
+    }
+    stream.synchronize().unwrap();
+    staging.read_all()
+}
+
+#[test]
+fn zero_grad_moves_parameter_only_by_weight_decay() {
+    let dev = test_device();
+    let stream = dev.cuda_stream().unwrap().clone();
+    let theta_init: Vec<f32> = vec![1.0, 2.0, -3.0, 0.5];
+    let mut theta = upload(&stream, &theta_init);
+    let grad = upload(&stream, &vec![0.0; 4]);
+    let mut opt = AdamW::new(&dev, 4, 1e-3).unwrap();
+    opt.wd = 0.1;
+    opt.step(&mut theta, &grad).unwrap();
+    let theta_after = download(&stream, &theta);
+    // With g=0, m=v=0 too, so the Adam update is 0; only weight decay
+    // applies: theta -= lr * wd * theta = theta * (1 - lr*wd)
+    let factor = 1.0_f32 - 1e-3 * 0.1;
+    for (a, b) in theta_after.iter().zip(&theta_init) {
+        assert_relative_eq!(*a, *b * factor, epsilon = 1e-6);
+    }
+}
+
+#[test]
+fn positive_grad_decreases_parameter() {
+    let dev = test_device();
+    let stream = dev.cuda_stream().unwrap().clone();
+    let theta_init: Vec<f32> = vec![5.0, 3.0, 1.0, 0.1];
+    let mut theta = upload(&stream, &theta_init);
+    let grad = upload(&stream, &vec![1.0, 1.0, 1.0, 1.0]);
+    let mut opt = AdamW::new(&dev, 4, 1e-2).unwrap();
+    opt.wd = 0.0; // isolate the Adam update
+    opt.step(&mut theta, &grad).unwrap();
+    let theta_after = download(&stream, &theta);
+    for (a, b) in theta_after.iter().zip(&theta_init) {
+        assert!(*a < *b, "positive grad must reduce parameter: {b} -> {a}");
+    }
+}
+
+#[test]
+fn adam_step_count_increments() {
+    let dev = test_device();
+    let stream = dev.cuda_stream().unwrap().clone();
+    let mut theta = upload(&stream, &vec![1.0; 8]);
+    let grad = upload(&stream, &vec![0.1; 8]);
+    let mut opt = AdamW::new(&dev, 8, 1e-3).unwrap();
+    assert_eq!(opt.step_count(), 0);
+    opt.step(&mut theta, &grad).unwrap();
+    assert_eq!(opt.step_count(), 1);
+    opt.step(&mut theta, &grad).unwrap();
+    assert_eq!(opt.step_count(), 2);
+}
+
+#[test]
+fn repeated_grad_descent_drives_parameter_toward_zero() {
+    // grad = theta (gradient of 0.5 * theta^2 wrt theta), repeated AdamW steps
+    // must reduce |theta| toward 0 over many steps.
+    let dev = test_device();
+    let stream = dev.cuda_stream().unwrap().clone();
+    let initial: Vec<f32> = vec![10.0; 16];
+    let mut theta = upload(&stream, &initial);
+    let mut opt = AdamW::new(&dev, 16, 1e-1).unwrap();
+    opt.wd = 0.0;
+    for _ in 0..200 {
+        let current = download(&stream, &theta);
+        let grad = upload(&stream, &current);
+        opt.step(&mut theta, &grad).unwrap();
+    }
+    let final_theta = download(&stream, &theta);
+    let norm: f32 = final_theta.iter().map(|v| v * v).sum::<f32>().sqrt();
+    assert!(norm < 1.0, "‖θ‖ should drop well below initial 40 after 200 steps, got {norm}");
+}

crates/ml-alpha/tests/bce_grad_finite_diff.rs

@@ -0,0 +1,85 @@
+//! BCE forward + backward validation.
+//!
+//! Per `feedback_no_cpu_test_fallbacks.md`: no CPU oracle. The gradient
+//! is validated against finite-difference computed via GPU forward calls
+//! at perturbed inputs (the kernel itself is the reference; we only
+//! verify that the analytic grad matches its own numerical derivative).
+
+use approx::assert_relative_eq;
+use ml_alpha::trainer::loss::{bce_multi_horizon_loss_and_grad_gpu, BceInput};
+use ml_core::device::MlDevice;
+
+fn test_device() -> MlDevice {
+    MlDevice::cuda(0).expect("CUDA 0 required for ml-alpha tests")
+}
+
+fn make_input(probs: Vec<f32>, labels: Vec<f32>) -> BceInput {
+    BceInput { probs, labels, n_horizons: 5, n_pos: 4 }
+}
+
+#[test]
+fn bce_loss_is_positive_and_finite() {
+    let dev = test_device();
+    let probs: Vec<f32> = (0..20).map(|i| (0.1 + i as f32 * 0.04).min(0.95)).collect();
+    let labels: Vec<f32> = (0..20).map(|i| (i % 2) as f32).collect();
+    let out = bce_multi_horizon_loss_and_grad_gpu(&dev, &make_input(probs, labels)).unwrap();
+    assert!(out.loss.is_finite());
+    assert!(out.loss > 0.0);
+    assert_eq!(out.n_valid, 20);
+}
+
+#[test]
+fn bce_grad_zero_when_probs_match_labels() {
+    let dev = test_device();
+    // probs ≈ labels (within clamp window): grad ≈ 0 wherever |p-y| ≈ 0
+    let mut probs = vec![0.5_f32; 20];
+    let labels: Vec<f32> = (0..20).map(|i| (i % 2) as f32).collect();
+    // Move probs very close to labels (but within (1e-6, 1-1e-6) clamp window).
+    for i in 0..20 {
+        probs[i] = if labels[i] > 0.5 { 1.0 - 1e-5 } else { 1e-5 };
+    }
+    let out = bce_multi_horizon_loss_and_grad_gpu(&dev, &make_input(probs, labels)).unwrap();
+    // Near-perfect predictions → near-zero loss
+    assert!(out.loss < 1e-3, "loss should be small for matching probs/labels, got {}", out.loss);
+}
+
+#[test]
+fn bce_grad_matches_finite_difference() {
+    let dev = test_device();
+    let probs: Vec<f32> = (0..20).map(|i| (0.1 + i as f32 * 0.04).min(0.95)).collect();
+    let labels: Vec<f32> = (0..20).map(|i| (i % 2) as f32).collect();
+    let base = bce_multi_horizon_loss_and_grad_gpu(&dev, &make_input(probs.clone(), labels.clone())).unwrap();
+
+    let eps = 1e-3_f32;
+    for i in [0usize, 5, 10, 15, 19] {
+        let mut up = probs.clone();
+        up[i] += eps;
+        let mut dn = probs.clone();
+        dn[i] -= eps;
+        let lu = bce_multi_horizon_loss_and_grad_gpu(&dev, &make_input(up, labels.clone())).unwrap().loss;
+        let ld = bce_multi_horizon_loss_and_grad_gpu(&dev, &make_input(dn, labels.clone())).unwrap().loss;
+        let fd = (lu - ld) / (2.0 * eps);
+        assert_relative_eq!(
+            base.grad_probs[i],
+            fd,
+            epsilon = 5e-3,
+            max_relative = 5e-2
+        );
+    }
+}
+
+#[test]
+fn bce_nan_labels_mask_grad_and_loss() {
+    let dev = test_device();
+    let mut labels: Vec<f32> = (0..20).map(|i| (i % 2) as f32).collect();
+    // Mask out indices 5, 10, 15
+    labels[5] = f32::NAN;
+    labels[10] = f32::NAN;
+    labels[15] = f32::NAN;
+    let probs: Vec<f32> = (0..20).map(|i| (0.1 + i as f32 * 0.04).min(0.95)).collect();
+    let out = bce_multi_horizon_loss_and_grad_gpu(&dev, &make_input(probs, labels)).unwrap();
+    assert_eq!(out.n_valid, 17, "should drop 3 NaN-masked entries");
+    for &i in &[5usize, 10, 15] {
+        assert_eq!(out.grad_probs[i], 0.0, "masked grad should be 0 at index {i}, got {}", out.grad_probs[i]);
+    }
+}