Files
foxhunt/docs/superpowers/plans/2026-05-16-ml-alpha-phase-a-api-addendum.md
jgrusewski d4e46aba94 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>
2026-05-16 21:51:48 +02:00

12 KiB

Plan 1 API Addendum — cudarc 0.19 + foxhunt canonical patterns

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 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 allocation, copy, and kernel-launch ownership from the device to the stream — partly to make CUDA Graph capture clean (everything that becomes a graph node has to happen on the captured stream).


Imports

use std::sync::Arc;
use cudarc::driver::{
    CudaContext, CudaStream, CudaModule, CudaFunction, CudaSlice,
    DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg, DriverError,
};
use crate::pinned_mem::MappedF32Buffer;  // local copy — see pinned_mem.rs header
use ml_core::device::MlDevice;

cudarc::driver::result is the raw-call namespace (used for memcpy_dtod_async).


Pattern 1 — Device + stream construction

Production / runtime:

let dev = MlDevice::cuda(0)?;                 // panics if no GPU
let ctx: &Arc<CudaContext> = dev.cuda_context()?;
let stream: &Arc<CudaStream> = dev.cuda_stream()?;

Tests:

fn test_device() -> MlDevice {
    MlDevice::cuda(0).expect("CUDA 0 required for ml-alpha tests")
}

The plan's references to MlDevice::default_for_test() should be replaced with a local test_device() helper inside the test module (don't add a public shim to ml-core for this).


Pattern 2 — Cubin load + function handle

The build script emits cubins to OUT_DIR. At runtime:

const CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/snap_feature_assemble.cubin"));

fn load_fn(ctx: &Arc<CudaContext>, name: &str) -> Result<CudaFunction, DriverError> {
    let module: Arc<CudaModule> = ctx.load_cubin(CUBIN.to_vec())?;
    module.load_function(name)
}

Cache the CudaFunction and the Arc<CudaModule> as struct fields where the kernel is used repeatedly (CfcTrunk, MultiHorizonHeads, BCE, AdamW holders). Don't reload on every forward call — that defeats the whole point of the pre-compiled cubin discipline.


Pattern 3 — Host data → mapped-pinned → device (the canonical CPU→GPU path)

For Phase A bit-equiv tests and the training inner loop, the standard pattern is mapped-pinned staging → DtoD-async into a CudaSlice<f32> → pass the slice to launch. This keeps kernel signatures uniform (every kernel takes CudaSlice args; mapped-pinned is just an upload optimization that is invisible to the kernel).

fn upload(stream: &Arc<CudaStream>, host: &[f32]) -> Result<CudaSlice<f32>, anyhow::Error> {
    let n = host.len();
    let staging = unsafe { MappedF32Buffer::new(n) }
        .map_err(|e| anyhow::anyhow!("mapped-pinned alloc: {e}"))?;
    staging.write_from_slice(host);

    let mut dst = stream.alloc_zeros::<f32>(n)?;
    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(),
            )?;
        }
    }
    Ok(dst)
}

Use upload() everywhere the plan's bindings say dev.htod_copy(vec). Same shape for i32 etc. via MappedI32Buffer.

No plain cuMemcpyHtoD anywhere per feedback_no_htod_htoh_only_mapped_pinned.md.


Pattern 4 — Device → host (slow-path readback only)

fn download(stream: &Arc<CudaStream>, src: &CudaSlice<f32>) -> Result<Vec<f32>, anyhow::Error> {
    let n = src.len();
    let staging = unsafe { MappedF32Buffer::new(n) }
        .map_err(|e| anyhow::anyhow!("readback alloc: {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(),
        )?;
    }
    stream.synchronize()?;
    Ok(staging.read_all())
}

Use download() only for slow-path test readback and ISV snapshots. Never on the hot path.


Pattern 5 — Kernel launch (the builder)

fn launch<P: PushKernelArg>(stream: &Arc<CudaStream>, func: &CudaFunction, cfg: LaunchConfig, args: impl FnOnce(&mut cudarc::driver::LaunchArgs)) -> Result<(), DriverError> {
    let mut builder = stream.launch_builder(func);
    args(&mut builder);
    unsafe { builder.launch(cfg) }
}

But idiomatic usage is inline — no need for the helper:

let cfg = LaunchConfig { grid_dim: (1,1,1), block_dim: (128,1,1), shared_mem_bytes: 0 };
let mut launch = stream.launch_builder(&func);
launch
    .arg(&w_in_d)        // CudaSlice<f32>
    .arg(&w_rec_d)
    .arg(&b_d)
    .arg(&x_d)
    .arg(&h_old_d)
    .arg(&dt_s)          // scalar f32
    .arg(&(n_in as i32))
    .arg(&(n_hid as i32))
    .arg(&mut h_new_d);  // mutable arg uses .arg() too — &mut CudaSlice
unsafe { launch.launch(cfg)?; }

Scalar args (f32, i32, u64) pass by &value — they implement DeviceRepr. Device slices pass by &slice or &mut slice.

The plan's func.launch_async(cfg, (args...)) should be replaced with the above builder.


Pattern 6 — IsvBus (Task 3)

The plan's IsvBus uses dev.alloc_zeros and dev.default_stream — both nonexistent. The corrected impl:

pub struct IsvBus {
    stream: Arc<CudaStream>,
    buffer: CudaSlice<f32>,
}

impl IsvBus {
    pub fn new(dev: &MlDevice) -> Result<Self> {
        let stream = dev.cuda_stream()?.clone();
        let buffer = stream.alloc_zeros::<f32>(SLOTS)?;
        Ok(Self { stream, buffer })
    }

    pub fn write_slot(&mut self, idx: usize, value: f32) -> Result<()> {
        assert!(idx < SLOTS);
        let staging = unsafe { MappedF32Buffer::new(1) }.map_err(|e| anyhow!("isv staging: {e}"))?;
        staging.write_from_slice(&[value]);
        unsafe {
            let (dst_ptr, _g) = self.buffer.device_ptr_mut(&self.stream);
            cudarc::driver::result::memcpy_dtod_async(
                dst_ptr + (idx * 4) as u64,
                staging.dev_ptr,
                4,
                self.stream.cu_stream(),
            )?;
        }
        self.stream.synchronize()?;
        Ok(())
    }

    pub fn snapshot(&self) -> Result<[f32; SLOTS]> {
        let v = download(&self.stream, &self.buffer)?;
        let mut out = [0f32; SLOTS];
        out.copy_from_slice(&v);
        Ok(out)
    }

    pub fn buffer(&self) -> &CudaSlice<f32> { &self.buffer }
}

Pattern 7 — MappedPinnedSnapshotSlot (Task 4)

The plan invented as_slice_mut(), as_slice(), device_ptr() accessors that don't exist on MappedF32Buffer. The real accessors are host_slice_mut() -> &mut [f32], read_all() -> Vec<f32>, and the public dev_ptr: CUdeviceptr field.

Corrected MappedPinnedSnapshotSlot:

pub struct MappedPinnedSnapshotSlot { backing: MappedF32Buffer }

impl MappedPinnedSnapshotSlot {
    pub fn new() -> Result<Self> {
        let n = (std::mem::size_of::<SnapshotPayload>() + 3) / 4;
        let backing = unsafe { MappedF32Buffer::new(n) }
            .map_err(|e| anyhow!("snapshot slot alloc: {e}"))?;
        Ok(Self { backing })
    }

    pub fn write_volatile(&mut self, payload: &SnapshotPayload) {
        let host = self.backing.host_slice_mut();
        let dst = host.as_mut_ptr() as *mut SnapshotPayload;
        unsafe { std::ptr::write_volatile(dst, *payload); }
        std::sync::atomic::compiler_fence(std::sync::atomic::Ordering::Release);
    }

    pub fn peek(&self) -> SnapshotPayload {
        let v = self.backing.read_all();
        let src = v.as_ptr() as *const SnapshotPayload;
        unsafe { std::ptr::read_volatile(src) }
    }

    pub fn device_ptr(&self) -> cudarc::driver::sys::CUdeviceptr { self.backing.dev_ptr }
}

Tests asserting slot.device_ptr() != 0 work unchanged (the field is CUdeviceptr, a u64, so != 0 is meaningful).

MappedPinnedFillSlot follows the same shape.


Pattern 8 — Kernel argument: passing a MappedF32Buffer's pinned region

If a kernel must read directly from a mapped-pinned slot (hot-path zero-copy), the kernel arg type on the Rust side is &u64 carrying the CUdeviceptr:

let dev_ptr_arg: u64 = slot.device_ptr();
let mut launch = stream.launch_builder(&func);
launch.arg(&dev_ptr_arg);  // passed as u64 -> CUdeviceptr in the kernel signature
unsafe { launch.launch(cfg)?; }

The CUDA kernel signature for that arg is const float* __restrict__ as usual — no kernel-side change.

For Phase A, prefer the DtoD-into-CudaSlice pattern (Pattern 3) — it keeps every kernel binding uniform. Save zero-copy pinned reads for the live hot path (Plan 3 work).


Pattern 9 — CUDA Graph A capture (Task 11)

cudarc 0.19's safe graph API:

use cudarc::driver::{CudaGraph, CudaGraphExec};
use cudarc::driver::sys::CUstreamCaptureMode;

impl CfcTrunk {
    pub fn capture_graph_a(&mut self) -> Result<()> {
        // Warm-up (kernels loaded, allocations stable)
        self.forward_snapshot(&Mbp10RawInput::default())?;

        let stream = self.stream.clone();
        unsafe {
            stream.begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL)?;
            self.dispatch_perception_kernels()?;
            let graph = stream.end_capture()?;
            let exec = graph.instantiate(cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH)?;
            self.graph_a_exec = Some(exec);
        }
        Ok(())
    }

    pub fn perception_forward_captured(&mut self, input: &Mbp10RawInput) -> Result<([f32; 5], [f32; 8])> {
        // Stage scalars + raw input device buffers (outside the graph for v1).
        self.upload_input(input)?;
        let exec = self.graph_a_exec.as_ref().expect("Graph A not captured");
        let stream = &self.stream;
        unsafe { exec.launch(stream)?; }
        stream.synchronize()?;
        Ok((download_array5(&stream, &self.probs_d)?, download_array8(&stream, &self.proj_out_d)?))
    }
}

Verify stream.begin_capture / stream.end_capture / graph.instantiate / exec.launch signatures from vendor/cudarc/src/driver/safe/graph.rs when landing the task — minor flags shape may differ.


Per-task summary of plan-vs-real changes

Plan 1 Task Change
3 (IsvBus) Use Pattern 6
4 (pinned slots) Use Pattern 7
5 (snap_feature_assemble) Pattern 1 + 2 + 3 + 5; replace func.launch_async((args...)) with launch_builder
6 (cfc_step) Same as Task 5
7 (multi_horizon_heads) Same
8 (projection) Same
9 (CfcTrunk) Cache Arc<CudaStream>, CudaFunctions, and CudaSlices in struct; remove device-method calls
10 (BCE + AdamW) Same
11 (Graph A capture) Use Pattern 9
13 (Phase A trainer) Inherits CfcTrunk's stream/handles
14 (Backprop kernels) Same as Task 5 for binding; kernels themselves are the real work

The plan's kernel .cu source, CPU oracles, finite-diff thresholds, smoke criteria, and gate verdict logic are unchanged. Only the Rust binding code uses these patterns.