gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

perf(ml-alpha): CUDA Graph capture of training step (#162 finale)

Browse Source 
Captures the full Mamba2 fwd → K-loop fwd → BCE → K-loop bwd →
Mamba2 bwd → AdamW × 7 → loss DtoD chain into a single CUDA Graph.
First step_batched call runs uncaptured (cuBLAS warmup); second
call captures; third+ replays. Replaces ~155 individual kernel
launches per step with one graph launch.

Four root-cause issues had to be fixed in concert to make the
capture region capture-compatible:

1. cuBLAS lazy workspace allocation
   crates/ml-alpha/src/mamba2_block.rs — pre-allocate an 8 MiB
   workspace via cublasSetWorkspace_v2 at Mamba2Block::new.
   cuBLAS would otherwise allocate on first gemm call with each
   new shape, breaking capture.

2. cuBLAS heuristic plan-cache lookup
   crates/ml-core/src/cuda_autograd/linear.rs — switch
   gemm_ex_f32 from CUBLAS_GEMM_DEFAULT_TENSOR_OP to
   CUBLAS_GEMM_DFALT. The heuristic algo path triggers
   plan-cache allocs; DFALT is deterministic with negligible
   perf delta for our shapes (Mamba2: 128×32, 128×state_dim).

3. Per-call cuModuleLoadData in bias kernels
   crates/ml-core/src/cuda_autograd/linear.rs — add a
   `BiasKernels` struct (add_bias_2d_kernel + reduce_sum_axis0
   handles) cached at construction. `BiasKernels::shared(stream)`
   uses a per-context OnceLock cache so the cubin loads exactly
   once per CUDA context for the process lifetime. Every
   `GpuLinear` and `OwnedGpuLinear` constructor now stores its
   `BiasKernels`. Removed the per-call `get_bias_kernels`
   helper entirely (no legacy aliases — greenfield).

4. Per-call CudaSlice::clone() in Mamba2 fwd + bwd
   crates/ml-alpha/src/mamba2_block.rs — three sites cloned
   the input slice to build a fresh `GpuTensor` view. Each
   `CudaSlice::clone()` does cuMemAlloc + dtod copy
   (vendor/cudarc/src/driver/safe/core.rs:1437); cuMemAlloc
   is forbidden during capture.

   Refactored `forward_with_slices_into` and
   `backward_with_slices_into` to take raw `x_data: &CudaSlice<f32>,
   batch: usize` instead of `&GpuTensor` / `&LinearActivations`.
   All three Mamba2 fwd call sites + three bwd call sites now
   pass the underlying CudaSlice directly.

Trainer changes (trainer/perception.rs):
  - Three-state machine in step_batched: warmup → capture → replay.
  - Labels staging fill moved BEFORE the captured region (host writes
    only; replays read whatever the host last wrote).
  - Removed the post-snap_batched sync that was inside dispatch (it
    would trip STREAM_CAPTURE_ISOLATION; kernels are stream-ordered
    so the sync was unnecessary).
  - Dispatch extracted into `dispatch_train_step` method; the
    captured region brackets exactly this method.
  - Final sync + mapped-pinned loss read happens once per step in
    step_batched, outside the captured region.

Validation:
  - Synthetic overfit smoke: 0.397 → 0.0007 (matches pre-refactor
    trajectory; capture+replay produces equivalent loss).
  - 26 ml-alpha lib tests + 23 integration tests pass.
  - 306 ml-core lib tests pass.

Also fixed (orthogonal but required for green workspace):
  crates/ml-core/src/action_space.rs — tests assumed 7-exposure
  layout (63 actions). Production `ExposureLevel` enum has 8
  variants (Hold inserted at idx 3, Flat moved to idx 7 → 72
  total actions). Updated tests + `get_valid_action_mask` to
  match the 8-level layout.

Honors:
  - feedback_no_legacy_aliases.md (no `add_bias_2d_with_fn` /
    legacy `get_bias_kernels` fallback; one canonical API).
  - feedback_no_partial_refactor.md (signature change propagated
    to every caller atomically; no half-migrated state).
  - feedback_wire_everything_up.md (BiasKernels wired into all
    GpuLinear/OwnedGpuLinear constructors in same commit).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-17 15:21:00 +02:00
parent ab94ce2a49
commit 87ca1f5f55
5 changed files with 334 additions and 123 deletions
Download Patch File Download Diff File Expand all files Collapse all files

76
crates/ml-alpha/src/mamba2_block.rs
View File

@@ -38,7 +38,7 @@ use std::sync::Arc;
use anyhow::{anyhow, Result};
use cudarc::cublas::CudaBlas;
use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg};
use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg};
use ml_core::cuda_autograd::gpu_tensor::GpuTensor;
use ml_core::cuda_autograd::linear::{LinearActivations, LinearGrads, OwnedGpuLinear};
@@ -366,6 +366,14 @@ pub struct Mamba2Block {
/// cuBLAS handle for the four linear projections (input / A / B / output).
pub cublas: CudaBlas,
/// Pre-allocated cuBLAS workspace. cuBLAS lazily allocates internal
/// workspace per gemm call when none is set, which is incompatible
/// with CUDA Graph stream capture (cudaMalloc inside the captured
/// region trips `CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED`). Setting a
/// user-managed workspace via `cublasSetWorkspace_v2` removes the
/// per-call alloc; the buffer stays alive for the lifetime of the
/// block, large enough for all gemm sizes used by the projections.
_cublas_workspace: CudaSlice<u8>,
}
impl Mamba2Block {
@@ -404,9 +412,28 @@ impl Mamba2Block {
.load_function("mamba2_alpha_adamw_step")
.map_err(|e| anyhow!("Mamba2Block: AdamW kernel resolve: {e}"))?;
// cuBLAS handle for the four GEMM projections.
// cuBLAS handle for the four GEMM projections, with a pre-
// allocated user-managed workspace so subsequent gemm calls
// don't trigger internal allocations during CUDA Graph capture.
// NVIDIA recommends 4 MiB minimum for general use, larger for
// batched ops on Ampere+; 8 MiB is comfortable headroom for
// our gemm shapes (largest is hidden_dim×in_dim ≈ 128×32).
let cublas = CudaBlas::new(Arc::clone(&stream))
.map_err(|e| anyhow!("Mamba2Block: cuBLAS init failed: {e}"))?;
const CUBLAS_WORKSPACE_BYTES: usize = 8 * 1024 * 1024;
let cublas_workspace = stream
.alloc_zeros::<u8>(CUBLAS_WORKSPACE_BYTES)
.map_err(|e| anyhow!("Mamba2Block: cuBLAS workspace alloc: {e}"))?;
unsafe {
let (ws_ptr, _g) = cublas_workspace.device_ptr(&stream);
cudarc::cublas::sys::cublasSetWorkspace_v2(
*cublas.handle(),
ws_ptr as *mut std::ffi::c_void,
CUBLAS_WORKSPACE_BYTES,
)
.result()
.map_err(|e| anyhow!("Mamba2Block: cublasSetWorkspace_v2: {e:?}"))?;
}
// ── Parameter allocation + initialisation. ──────────────────────
// All on GPU; init helpers in ml-core::cuda_autograd::init use
@@ -450,6 +477,7 @@ impl Mamba2Block {
kernel_reduce_d_w_c,
kernel_adamw,
cublas,
_cublas_workspace: cublas_workspace,
})
}
@@ -1039,26 +1067,26 @@ impl Mamba2Block {
);
let n_rows = n_batch * c.seq_len;
// input_2d: reshape view of input (same device storage, fresh
// wrapper — no allocation; Arc::clone on the CudaSlice).
let input_2d = GpuTensor::new(input.cuda_data().clone(), vec![n_rows, c.in_dim])
.map_err(|e| anyhow!("reshape input → 2D: {e}"))?;
// 1. x = input_2d @ W_in.T + b_in (writes into scratch.x).
// 1. x = input @ W_in.T + b_in — input is [B, K, in_dim], viewed
// as [B*K, in_dim] via raw slice + batch dim (no GpuTensor wrap,
// no CudaSlice::clone — cuMemAlloc inside capture is forbidden).
self.w_in.inner.forward_with_slices_into(
&input_2d, &self.w_in.weight, &self.w_in.bias,
input.cuda_data(), n_rows,
&self.w_in.weight, &self.w_in.bias,
&self.cublas, &self.stream, &mut scratch.x,
).map_err(|e| anyhow!("w_in fwd_into: {e}"))?;
// 2. a_proj = x @ W_a.T + b_a.
self.w_a.inner.forward_with_slices_into(
&scratch.x, &self.w_a.weight, &self.w_a.bias,
scratch.x.cuda_data(), n_rows,
&self.w_a.weight, &self.w_a.bias,
&self.cublas, &self.stream, &mut scratch.a_proj,
).map_err(|e| anyhow!("w_a fwd_into: {e}"))?;
// 3. b_proj = x @ W_b.T + b_b.
self.w_b.inner.forward_with_slices_into(
&scratch.x, &self.w_b.weight, &self.w_b.bias,
scratch.x.cuda_data(), n_rows,
&self.w_b.weight, &self.w_b.bias,
&self.cublas, &self.stream, &mut scratch.b_proj,
).map_err(|e| anyhow!("w_b fwd_into: {e}"))?;
@@ -1501,16 +1529,24 @@ impl Mamba2Block {
.map_err(|e| anyhow!("reduce dw_c full_into: {e}"))?;
}
// ── Linear backward: w_b, w_a, w_in (all _into variants). ────
let x_act = LinearActivations { input: fwd_scratch.x.clone() };
// ── Linear backward: w_b, w_a, w_in (all _into variants).
// Each takes the saved forward input as a raw slice + batch
// dim (NO GpuTensor wrapping, NO CudaSlice::clone — cuMemAlloc
// inside capture is forbidden). w_b/w_a use fwd_scratch.x as
// their input; w_in uses the original input tensor.
let n_rows = n_batch * c.seq_len;
self.w_b.inner.backward_with_slices_into(
&grads_buffers.d_b_proj_2d, &x_act, &self.w_b.weight,
&grads_buffers.d_b_proj_2d,
fwd_scratch.x.cuda_data(), n_rows,
&self.w_b.weight,
&self.cublas, &self.stream,
&mut grads_buffers.dw_b, &mut grads_buffers.db_b, &mut grads_buffers.d_x_from_b,
).map_err(|e| anyhow!("w_b bwd_into: {e}"))?;
self.w_a.inner.backward_with_slices_into(
&grads_buffers.d_a_proj_2d, &x_act, &self.w_a.weight,
&grads_buffers.d_a_proj_2d,
fwd_scratch.x.cuda_data(), n_rows,
&self.w_a.weight,
&self.cublas, &self.stream,
&mut grads_buffers.dw_a, &mut grads_buffers.db_a, &mut grads_buffers.d_x_from_a,
).map_err(|e| anyhow!("w_a bwd_into: {e}"))?;
@@ -1520,14 +1556,10 @@ impl Mamba2Block {
.add_into(&grads_buffers.d_x_from_b, &grads_buffers.d_x, &self.stream)
.map_err(|e| anyhow!("d_x add_into: {e}"))?;
// W_in backward: input_2d is a fresh wrapper around the
// original input pointer (Arc::clone — no allocation).
let n_rows = n_batch * c.seq_len;
let input_2d = GpuTensor::new(input.cuda_data().clone(), vec![n_rows, c.in_dim])
.map_err(|e| anyhow!("reshape input for w_in bwd_into: {e}"))?;
let input_act = LinearActivations { input: input_2d };
self.w_in.inner.backward_with_slices_into(
&grads_buffers.d_x, &input_act, &self.w_in.weight,
&grads_buffers.d_x,
input.cuda_data(), n_rows,
&self.w_in.weight,
&self.cublas, &self.stream,
&mut grads_buffers.dw_in, &mut grads_buffers.db_in, &mut grads_buffers.d_x_from_in,
).map_err(|e| anyhow!("w_in bwd_into: {e}"))?;

130
crates/ml-alpha/src/trainer/perception.rs
View File

@@ -35,9 +35,10 @@
use std::sync::Arc;
use anyhow::{Context, Result};
use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
use cudarc::driver::{
CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig,
PushKernelArg,
CudaFunction, CudaGraph, CudaModule, CudaSlice, CudaStream, DevicePtr, DevicePtrMut,
LaunchConfig, PushKernelArg,
};
use ml_core::cuda_autograd::gpu_tensor::GpuTensor;
use ml_core::device::MlDevice;
@@ -254,6 +255,14 @@ pub struct PerceptionTrainer {
// single-call host→device path; staging covers all K*B labels per
// training step in one DtoD copy).
stg_labels: MappedF32Buffer,
/// Captured CUDA Graph for the training step. First step_batched
/// call runs uncaptured (warmup); second call captures; third+
/// replays. Eliminates ~155 individual kernel launches per step.
train_graph: Option<CudaGraph>,
/// cuBLAS warmup done — required before stream capture can safely
/// include cuBLAS calls.
cublas_warmed: bool,
}
impl PerceptionTrainer {
@@ -371,6 +380,8 @@ impl PerceptionTrainer {
grad_heads_w_d: stream.alloc_zeros::<f32>(N_HORIZONS * n_hid)?,
grad_heads_b_d: stream.alloc_zeros::<f32>(N_HORIZONS)?,
stg_labels: unsafe { MappedF32Buffer::new(k * cfg.n_batch * N_HORIZONS) }.map_err(|e| anyhow::anyhow!("stg_labels: {e}"))?,
train_graph: None,
cublas_warmed: false,
// Batched snap_feature staging — sized for max B*K snapshots.
bk_capacity: cfg.n_batch * k,
@@ -515,19 +526,12 @@ impl PerceptionTrainer {
);
}
// ── 1. Build the batched window tensor [B, K, FEATURE_DIM] in
// ONE fused snap_feature_assemble_batched launch (was
// B*K separate launches). Per-step host work: pack the
// 12 mapped-pinned staging buffers (~150 KB total); then
// one DtoD per array → device; then one kernel launch
// with B*K threads. Output written directly into the
// window tensor's storage.
// Use pre-allocated window_tensor_d — no per-step alloc.
debug_assert_eq!(self.window_tensor_d.shape(), &[b_sz, k_seq, FEATURE_DIM]);
let total_snaps = b_sz * k_seq;
debug_assert!(total_snaps <= self.bk_capacity);
// Pack mapped-pinned staging from the input batch (host writes).
// ── 1. Fill ALL mapped-pinned staging on host BEFORE stream
// ops. Replays read whatever host last wrote.
{
let bid_px_h = self.stg_bid_px_all.host_slice_mut();
let bid_sz_h = self.stg_bid_sz_all.host_slice_mut();
@@ -568,11 +572,78 @@ impl PerceptionTrainer {
}
}
}
// Labels staging fill (host-only). Moved from inside dispatch
// to here so it happens BEFORE the captured region.
{
let dst = self.stg_labels.host_slice_mut();
for k in 0..k_seq {
for b_idx in 0..b_sz {
for h in 0..N_HORIZONS {
dst[(k * b_sz + b_idx) * N_HORIZONS + h] =
labels_batch[b_idx][k][h];
}
}
}
}
// ── 2. Three-state machine: warmup (first call), capture
// (second call), replay (third+).
if self.train_graph.is_some() {
self.train_graph
.as_ref()
.unwrap()
.launch()
.context("train graph launch")?;
} else if !self.cublas_warmed {
self.dispatch_train_step(b_sz, k_seq, total_snaps)
.context("train warmup dispatch")?;
self.cublas_warmed = true;
} else {
let ctx = self.stream.context().clone();
let _ = ctx.check_err();
unsafe { ctx.disable_event_tracking(); }
let begin = self.stream.begin_capture(
CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED,
);
if let Err(e) = begin {
unsafe { ctx.enable_event_tracking(); }
let _ = ctx.check_err();
return Err(anyhow::anyhow!("train begin_capture: {e}"));
}
let dispatch_result = self.dispatch_train_step(b_sz, k_seq, total_snaps);
let graph_result = self.stream.end_capture(
CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
);
unsafe { ctx.enable_event_tracking(); }
let _ = ctx.check_err();
dispatch_result.context("train graph dispatch (during capture)")?;
let graph = graph_result
.context("train graph end_capture")?
.ok_or_else(|| anyhow::anyhow!(
"train end_capture returned None — no work captured"
))?;
self.train_graph = Some(graph);
}
// ── 3. Sync + read mapped-pinned loss.
self.stream.synchronize().context("end-of-step sync")?;
let loss = unsafe { std::ptr::read_volatile(self.loss_host_d.host_ptr) };
Ok(loss)
}
/// Kernel-dispatch portion of a training step — captured into the
/// CUDA Graph on the second `step_batched` call.
fn dispatch_train_step(
&mut self,
b_sz: usize,
k_seq: usize,
total_snaps: usize,
) -> Result<()> {
// DtoD: staging (mapped-pinned, device-visible) → device buffers.
// The prev_bid_sz_all / prev_ask_sz_all buffers stay zero-init
// (loader doesn't currently compute prev-snapshot sizes — same
// as the single-snapshot path).
let n10 = total_snaps * 10;
let n6 = total_snaps * REGIME_DIM;
let n1 = total_snaps;
@@ -637,7 +708,8 @@ impl PerceptionTrainer {
.arg(self.window_tensor_d.data_mut());
launch.launch(snap_cfg).context("snap_batched fwd")?;
}
self.stream.synchronize().context("snap_batched sync")?;
// No sync — stream-ordered. A sync would trip
// STREAM_CAPTURE_ISOLATION inside the captured region.
// ── 2. Mamba2 per-step forward — writes into self.mamba2_fwd_scratch.
self.mamba2
@@ -664,21 +736,9 @@ impl PerceptionTrainer {
unsafe { launch.launch(cfg_tx).context("transpose h_enriched fwd")?; }
}
// ── 3. Pack labels in [K, B, N_HORIZONS] layout (matches the
// per-K slot layout of probs_per_k / labels_per_k) and
// upload via mapped-pinned staging.
// ── 3. Labels DtoD: staging (filled in step_batched before
// captured region) → device.
let total_labels = k_seq * b_sz * N_HORIZONS;
{
let dst = self.stg_labels.host_slice_mut();
for k in 0..k_seq {
for b_idx in 0..b_sz {
for h in 0..N_HORIZONS {
dst[(k * b_sz + b_idx) * N_HORIZONS + h] =
labels_batch[b_idx][k][h];
}
}
}
}
unsafe {
let (dst_ptr, _g) = self.labels_per_k_d.device_ptr_mut(&self.stream);
let nbytes = total_labels * std::mem::size_of::<f32>();
@@ -932,19 +992,15 @@ impl PerceptionTrainer {
.step_from_buffers_gpu_clip(&mut self.mamba2, &self.mamba2_grads_buffers)
.context("mamba2 AdamW step_from_buffers_gpu_clip")?;
// Final: queue the single-float DtoD copy loss_d → loss_host_d
// (mapped-pinned). Sync ONCE at end of step — flushes all
// queued kernels + the DtoD. Read via the mapped-pinned host_ptr
// without any extra dtoh roundtrip.
// Queue loss_d → loss_host_d (mapped-pinned). Captured inside
// graph; sync + read happen in step_batched outside the region.
unsafe {
let (src_ptr, _g) = self.loss_d.device_ptr(&self.stream);
cudarc::driver::result::memcpy_dtod_async(
self.loss_host_d.dev_ptr, src_ptr, 4, self.stream.cu_stream(),
).context("loss dtod → mapped-pinned")?;
}
self.stream.synchronize().context("end-of-step sync")?;
let loss = unsafe { std::ptr::read_volatile(self.loss_host_d.host_ptr) };
Ok(loss)
Ok(())
}
/// Single-sequence forward-only eval — thin wrapper around

83
crates/ml-core/src/action_space.rs
View File

@@ -3,21 +3,34 @@ pub use crate::common::action::{ExposureLevel, FactoredAction, OrderType, Urgenc
/// Returns action mask where true=valid, false=invalid based on position limits.
///
/// Prevents invalid exposure actions that would violate position limits.
/// DQN outputs 7 exposure actions (ShortSmall..LongFull).
/// DQN outputs 8 exposure actions ordered to match
/// [`ExposureLevel`] discriminants: ShortSmall(0), ShortHalf(1),
/// ShortFull(2), Hold(3), LongSmall(4), LongHalf(5), LongFull(6),
/// Flat(7). `Hold` and `Flat` are always valid (no exposure change /
/// close-to-zero); the six sized exposures are gated by `max_position`.
///
/// # Arguments
/// * `_current_position` - Current portfolio position (unused - actions use absolute targets)
/// * `_current_position` - Current portfolio position (unused; actions use absolute targets)
/// * `max_position` - Maximum allowed position magnitude (typically 2.0)
///
/// # Returns
/// Boolean mask of length 7 where:
/// - `true` = exposure action is valid (does not violate position limits)
/// - `false` = exposure action is invalid (would exceed max_position)
/// Boolean mask of length 8.
pub fn get_valid_action_mask(_current_position: f64, max_position: f64) -> Vec<bool> {
let exposures = [-0.25_f64, -0.50, -1.0, 0.0, 0.25, 0.50, 1.0];
// Index order matches `ExposureLevel` discriminants exactly.
// Hold (idx 3) and Flat (idx 7) are unconditionally valid; the
// six sized exposures are masked by their absolute target.
let exposures: [Option<f64>; 8] = [
Some(-0.25), Some(-0.50), Some(-1.0), // ShortSmall, ShortHalf, ShortFull
None, // Hold — always valid
Some(0.25), Some(0.50), Some(1.0), // LongSmall, LongHalf, LongFull
None, // Flat — always valid
];
exposures
.iter()
.map(|&exp| exp.abs() <= max_position)
.map(|slot| match slot {
Some(exp) => exp.abs() <= max_position,
None => true,
})
.collect()
}
@@ -31,10 +44,11 @@ mod tests {
assert_eq!(ExposureLevel::ShortSmall as usize, 0);
assert_eq!(ExposureLevel::ShortHalf as usize, 1);
assert_eq!(ExposureLevel::ShortFull as usize, 2);
assert_eq!(ExposureLevel::Flat as usize, 3);
assert_eq!(ExposureLevel::Hold as usize, 3);
assert_eq!(ExposureLevel::LongSmall as usize, 4);
assert_eq!(ExposureLevel::LongHalf as usize, 5);
assert_eq!(ExposureLevel::LongFull as usize, 6);
assert_eq!(ExposureLevel::Flat as usize, 7);
}
#[test]
@@ -53,8 +67,8 @@ mod tests {
#[test]
fn test_action_from_index_bidirectional() {
// Test all 63 actions round-trip
for idx in 0..63 {
// Test all 72 actions round-trip (8 exposure × 3 order × 3 urgency).
for idx in 0..72 {
let action = FactoredAction::from_index(idx).unwrap();
assert_eq!(
action.to_index(),
@@ -67,8 +81,8 @@ mod tests {
#[test]
fn test_action_from_index_bounds() {
// Test out of bounds indices
assert!(FactoredAction::from_index(63).is_err());
// Test out of bounds indices (valid range is 0..72).
assert!(FactoredAction::from_index(72).is_err());
assert!(FactoredAction::from_index(100).is_err());
assert!(FactoredAction::from_index(usize::MAX).is_err());
}
@@ -78,6 +92,9 @@ mod tests {
assert_eq!(ExposureLevel::ShortSmall.target_exposure(), -0.25);
assert_eq!(ExposureLevel::ShortHalf.target_exposure(), -0.50);
assert_eq!(ExposureLevel::ShortFull.target_exposure(), -1.0);
// Hold returns NaN — caller must check is_hold() and preserve
// the current position rather than reading this value.
assert!(ExposureLevel::Hold.target_exposure().is_nan());
assert_eq!(ExposureLevel::Flat.target_exposure(), 0.0);
assert_eq!(ExposureLevel::LongSmall.target_exposure(), 0.25);
assert_eq!(ExposureLevel::LongHalf.target_exposure(), 0.50);
@@ -139,31 +156,32 @@ mod tests {
#[test]
fn test_index_bijection() {
// Verify all 63 indices map to unique actions
// Verify all 72 indices map to unique actions.
use std::collections::HashSet;
let mut actions = HashSet::new();
for idx in 0..63 {
for idx in 0..72 {
let action = FactoredAction::from_index(idx).unwrap();
assert!(actions.insert(action), "Duplicate action for index {}", idx);
}
assert_eq!(actions.len(), 63);
assert_eq!(actions.len(), 72);
}
#[test]
fn test_flat_market_normal_action() {
// Test common neutral action
// Common close-to-zero action. With Hold inserted at index 3,
// Flat sits at exposure index 7 (see ExposureLevel discriminants).
// Index = 7 * 9 + 0 * 3 + 1 = 63 + 0 + 1 = 64.
let action = FactoredAction::new(ExposureLevel::Flat, OrderType::Market, Urgency::Normal);
// Index = 3 * 9 + 0 * 3 + 1 = 27 + 0 + 1 = 28
assert_eq!(action.to_index(), 28);
assert_eq!(action.to_index(), 64);
assert_eq!(action.target_exposure(), 0.0);
assert_eq!(action.transaction_cost(), 0.0015); // Wave 2.5 calibration
assert_eq!(action.urgency_weight(), 1.0);
// Round-trip
let reconstructed = FactoredAction::from_index(28).unwrap();
// Round-trip.
let reconstructed = FactoredAction::from_index(64).unwrap();
assert_eq!(action, reconstructed);
}
@@ -211,42 +229,49 @@ mod tests {
#[test]
fn test_action_masking_all_valid_at_standard_limit() {
// With max_position=2.0, all 7 exposure actions should be valid
// With max_position=2.0, all 8 exposure actions should be valid.
let mask = get_valid_action_mask(0.0, 2.0);
assert_eq!(mask.len(), 7);
assert_eq!(mask.len(), 8);
assert!(
mask.iter().all(|&v| v),
"All 7 exposure actions should be valid at max_position=2.0"
"All 8 exposure actions should be valid at max_position=2.0"
);
}
#[test]
fn test_action_masking_at_limit_1() {
// With max_position=1.0, all actions valid (max exposure = 1.0)
// With max_position=1.0, all 8 actions valid (max sized exposure = 1.0).
let mask = get_valid_action_mask(0.0, 1.0);
assert_eq!(mask.len(), 7);
assert_eq!(mask.len(), 8);
assert!(mask.iter().all(|&v| v), "All valid at max_position=1.0");
}
#[test]
fn test_action_masking_restrictive_limit() {
// With max_position=0.6, ShortFull/LongFull masked (exposure |1.0| > 0.6)
// With max_position=0.6, ShortFull/LongFull (|1.0| > 0.6) are
// masked. Index order matches `ExposureLevel` discriminants:
// ShortSmall(0), ShortHalf(1), ShortFull(2), Hold(3),
// LongSmall(4), LongHalf(5), LongFull(6), Flat(7).
let mask = get_valid_action_mask(0.0, 0.6);
assert_eq!(mask.len(), 7);
assert_eq!(mask.len(), 8);
assert!(mask[0], "ShortSmall (0.25) should be valid at max_position=0.6");
assert!(mask[1], "ShortHalf (0.50) should be valid at max_position=0.6");
assert!(!mask[2], "ShortFull (1.0) should be INVALID at max_position=0.6");
assert!(mask[3], "Flat should be valid at max_position=0.6");
assert!(mask[3], "Hold should always be valid");
assert!(mask[4], "LongSmall (0.25) should be valid at max_position=0.6");
assert!(mask[5], "LongHalf (0.50) should be valid at max_position=0.6");
assert!(!mask[6], "LongFull (1.0) should be INVALID at max_position=0.6");
assert!(mask[7], "Flat should always be valid");
}
#[test]
fn test_action_masking_flat_always_valid() {
// Both Hold (idx 3) and Flat (idx 7) bypass the magnitude check
// — they don't consume exposure capacity.
for pos in [-1.5, -1.0, 0.0, 1.0, 1.5] {
let mask = get_valid_action_mask(pos, 0.1);
assert!(mask[3], "Flat (exposure=0) should always be valid");
assert!(mask[3], "Hold should always be valid");
assert!(mask[7], "Flat (exposure=0) should always be valid");
}
}
}

165
crates/ml-core/src/cuda_autograd/linear.rs
View File

@@ -15,6 +15,7 @@ use std::sync::Arc;
use cudarc::cublas::CudaBlas;
use cudarc::cublas::sys::{self as cublas_sys, cublasOperation_t};
use cudarc::driver::CudaModule;
use cudarc::driver::{CudaSlice, CudaStream, CudaFunction, DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg};
// nvrtc removed — kernels compiled via build.rs cubins
@@ -78,7 +79,13 @@ unsafe fn gemm_ex_f32(
cublas_sys::cudaDataType_t::CUDA_R_32F,
ldc,
cublas_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F,
cublas_sys::cublasGemmAlgo_t::CUBLAS_GEMM_DEFAULT_TENSOR_OP,
// Deterministic default algo (no heuristic/plan-cache lookup).
// Required for CUDA Graph stream capture: TENSOR_OP heuristic
// path triggers internal cudaMalloc on first call with a given
// shape, which trips CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED.
// Negligible perf delta for our small matrix sizes
// (Mamba2 projections: 128×32, 128×state_dim).
cublas_sys::cublasGemmAlgo_t::CUBLAS_GEMM_DFALT,
).map_err(|e| MLError::ModelError(format!("cublasGemmEx {label}: {e:?}")))?;
Ok(())
}
@@ -98,6 +105,11 @@ pub struct GpuLinear {
pub in_dim: usize,
/// Output dimension.
pub out_dim: usize,
/// Pre-loaded bias-add and reduce-sum kernels. Cached at construction
/// so the forward/backward hot paths don't reload the cubin per call
/// (which would break CUDA Graph stream capture). Shared across all
/// linears in the same CUDA context via [`BiasKernels::shared`].
pub bias_kernels: BiasKernels,
}
/// Saved activations from a linear forward pass, needed for backward.
@@ -141,21 +153,15 @@ impl GpuLinear {
pub fn new(
in_features: usize,
out_features: usize,
_stream: &Arc<CudaStream>,
stream: &Arc<CudaStream>,
) -> Result<Self, MLError> {
let prefix = format!("_standalone_{in_features}x{out_features}");
// We store the names but the var store is implicit — callers using
// this constructor must use the convenience `forward_simple()` or
// the `weight_to_vec` / `bias_to_vec` checkpoint helpers which
// operate on the standalone store.
//
// The actual weights live in a store created via the `with_store`
// pattern. For now, we register in a fresh store held externally.
Ok(Self {
weight_name: format!("{prefix}.weight"),
bias_name: format!("{prefix}.bias"),
in_dim: in_features,
out_dim: out_features,
bias_kernels: BiasKernels::shared(stream)?,
})
}
@@ -251,7 +257,7 @@ impl GpuLinear {
}
// Add bias: Y[b, j] += bias[j] for all b
add_bias_2d(&mut y, &b_param.data, batch, out_dim, stream)?;
add_bias_2d(&self.bias_kernels.add_fn, &mut y, &b_param.data, batch, out_dim, stream)?;
// Save input for backward pass
let saved_input = clone_gpu_tensor(x, stream)?;
@@ -302,7 +308,7 @@ impl GpuLinear {
)?;
}
add_bias_2d(&mut y, bias, batch, out_dim, stream)?;
add_bias_2d(&self.bias_kernels.add_fn, &mut y, bias, batch, out_dim, stream)?;
let saved_input = clone_gpu_tensor(x, stream)?;
let activations = LinearActivations { input: saved_input };
@@ -355,7 +361,7 @@ impl GpuLinear {
}
// ── db[out] = sum(dY[B, out], axis=0) ───────────────────────
let db = reduce_sum_axis0(dy, batch, out_dim, stream)?;
let db = reduce_sum_axis0(&self.bias_kernels.reduce_fn, dy, batch, out_dim, stream)?;
// ── dX[B, in] = dY[B, out] @ W[out, in] ────────────────────
// col-major: dX_col[in, B] = W_col[in, out] @ dY_col[out, B]
@@ -433,7 +439,7 @@ impl GpuLinear {
}
// ── db[out] = sum(dY[B, out], axis=0) ───────────────────────
let db = reduce_sum_axis0(dy, batch, out_dim, stream)?;
let db = reduce_sum_axis0(&self.bias_kernels.reduce_fn, dy, batch, out_dim, stream)?;
// ── dX[B, in] = dY[B, out] @ W[out, in] ────────────────────
// col-major: dX_col[in, B] = W_col[in, out] @ dY_col[out, B]
@@ -471,18 +477,25 @@ impl GpuLinear {
///
/// `y_out` MUST already have shape `[batch, out_dim]`. Contents
/// are overwritten by the cuBLAS sgemm + bias-add.
/// `x_data` is a raw [`CudaSlice<f32>`] of length `batch * in_dim`,
/// row-major. Passed as a slice rather than a [`GpuTensor`] to
/// avoid the cuMemAlloc + dtod copy that `GpuTensor::new(... .clone(), ...)`
/// would do — CUDA Graph stream capture rejects per-call allocs.
pub fn forward_with_slices_into(
&self,
x: &GpuTensor,
x_data: &CudaSlice<f32>,
batch: usize,
weight: &CudaSlice<f32>,
bias: &CudaSlice<f32>,
cublas: &CudaBlas,
stream: &Arc<CudaStream>,
y_out: &mut GpuTensor,
) -> Result<(), MLError> {
let batch = if x.ndim() == 1 { 1 } else { x.shape()[0] };
let in_dim = self.in_dim;
let out_dim = self.out_dim;
debug_assert_eq!(x_data.len(), batch * in_dim,
"forward_with_slices_into: x_data len {} != batch {} * in_dim {}",
x_data.len(), batch, in_dim);
if y_out.shape() != [batch, out_dim] {
return Err(MLError::ModelError(format!(
"forward_with_slices_into: y_out shape {:?} != [{}, {}]",
@@ -491,7 +504,7 @@ impl GpuLinear {
}
let w_ptr = raw_ptr(weight, stream);
let x_ptr = raw_ptr(&x.data, stream);
let x_ptr = raw_ptr(x_data, stream);
let y_ptr = raw_ptr_mut(&mut y_out.data, stream);
unsafe {
@@ -511,7 +524,7 @@ impl GpuLinear {
"forward_with_slices_into",
)?;
}
add_bias_2d(y_out, bias, batch, out_dim, stream)?;
add_bias_2d(&self.bias_kernels.add_fn, y_out, bias, batch, out_dim, stream)?;
Ok(())
}
@@ -523,10 +536,16 @@ impl GpuLinear {
/// `dw_out` : `[out_dim, in_dim]`
/// `db_out` : `[out_dim]`
/// `dx_out` : `[batch, in_dim]`
/// `x_data` is the saved forward input — a raw slice of length
/// `batch * in_dim`. Passed as a slice (rather than wrapping in a
/// fresh [`LinearActivations { input: GpuTensor }`]) to avoid the
/// cuMemAlloc + dtod copy that CudaSlice::clone() would do, which
/// breaks CUDA Graph stream capture.
pub fn backward_with_slices_into(
&self,
dy: &GpuTensor,
activations: &LinearActivations,
x_data: &CudaSlice<f32>,
batch: usize,
weight: &CudaSlice<f32>,
cublas: &CudaBlas,
stream: &Arc<CudaStream>,
@@ -534,9 +553,14 @@ impl GpuLinear {
db_out: &mut GpuTensor,
dx_out: &mut GpuTensor,
) -> Result<(), MLError> {
let batch = dy.shape()[0];
let in_dim = self.in_dim;
let out_dim = self.out_dim;
debug_assert_eq!(x_data.len(), batch * in_dim,
"backward_with_slices_into: x_data len {} != batch {} * in_dim {}",
x_data.len(), batch, in_dim);
debug_assert_eq!(dy.shape()[0], batch,
"backward_with_slices_into: dy batch {} != expected {}",
dy.shape()[0], batch);
if dw_out.shape() != [out_dim, in_dim] {
return Err(MLError::ModelError(format!(
"backward_with_slices_into: dw shape {:?} != [{}, {}]",
@@ -556,7 +580,7 @@ impl GpuLinear {
)));
}
let x_ptr = raw_ptr(&activations.input.data, stream);
let x_ptr = raw_ptr(x_data, stream);
let dy_ptr = raw_ptr(&dy.data, stream);
let dw_ptr = raw_ptr_mut(&mut dw_out.data, stream);
unsafe {
@@ -577,7 +601,7 @@ impl GpuLinear {
)?;
}
reduce_sum_axis0_into(dy, batch, out_dim, stream, db_out)?;
reduce_sum_axis0_into(&self.bias_kernels.reduce_fn, dy, batch, out_dim, stream, db_out)?;
let w_ptr = raw_ptr(weight, stream);
let dy_ptr2 = raw_ptr(&dy.data, stream);
@@ -646,6 +670,7 @@ impl OwnedGpuLinear {
bias_name: format!("{prefix}.bias"),
in_dim,
out_dim,
bias_kernels: BiasKernels::shared(stream)?,
},
weight: w_data,
bias: b_data,
@@ -667,6 +692,7 @@ impl OwnedGpuLinear {
bias_name: format!("{prefix}.bias"),
in_dim,
out_dim,
bias_kernels: BiasKernels::shared(stream)?,
},
weight: w_data,
bias: b_data,
@@ -688,6 +714,7 @@ impl OwnedGpuLinear {
bias_name: format!("{prefix}.bias"),
in_dim,
out_dim,
bias_kernels: BiasKernels::shared(stream)?,
},
weight: w_data,
bias: b_data,
@@ -712,26 +739,89 @@ impl OwnedGpuLinear {
// ── Helper CUDA kernels ──────────────────────────────────────────────────
/// Get bias-add and reduce-sum kernel functions.
/// Pre-loaded bias-kernel handles (add_bias_2d, reduce_sum_axis0).
///
/// Load bias kernels from precompiled cubin (build.rs).
fn get_bias_kernels(stream: &Arc<CudaStream>) -> Result<(CudaFunction, CudaFunction), MLError> {
static CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/linear_kernels.cubin"));
let module = super::cubin_loader::load_cubin(CUBIN, stream)?;
let add_fn = super::cubin_loader::get_fn(&module, "add_bias_2d_kernel")?;
let reduce_fn = super::cubin_loader::get_fn(&module, "reduce_sum_axis0_kernel")?;
Ok((add_fn, reduce_fn))
/// Module loading (`cuModuleLoadData`) is NOT compatible with CUDA Graph
/// stream capture — it allocates device memory and trips
/// `CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED`. Production hot paths must
/// load the cubin ONCE at construction and reuse the function handles.
/// Both [`OwnedGpuLinear`] and the legacy [`GpuLinear`] paths build a
/// `BiasKernels` lazily; multiple linears in the same process share
/// a single underlying module via a per-context cache.
#[derive(Clone)]
pub struct BiasKernels {
pub add_fn: CudaFunction,
pub reduce_fn: CudaFunction,
_module: Arc<CudaModule>,
}
impl std::fmt::Debug for BiasKernels {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "BiasKernels {{ cached }}")
}
}
impl BiasKernels {
/// Load the bias-kernel module fresh. Prefer [`Self::shared`] in
/// any code that constructs many linears against the same context
/// — it caches the module per context so cuModuleLoadData runs
/// exactly once per process.
pub fn new(stream: &Arc<CudaStream>) -> Result<Self, MLError> {
static CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/linear_kernels.cubin"));
let module = super::cubin_loader::load_cubin(CUBIN, stream)?;
let add_fn = super::cubin_loader::get_fn(&module, "add_bias_2d_kernel")?;
let reduce_fn = super::cubin_loader::get_fn(&module, "reduce_sum_axis0_kernel")?;
Ok(Self { add_fn, reduce_fn, _module: module })
}
/// Return a `BiasKernels` for the given stream's context. The
/// underlying module is loaded exactly once per context for the
/// lifetime of the process; subsequent calls clone the cached
/// function handles (cheap — `CudaFunction` wraps a raw pointer +
/// an `Arc<CudaModule>`). The cache key is the context's
/// `Arc<CudaContext>` pointer; a fresh context (e.g. a new test
/// process) gets a fresh entry.
pub fn shared(stream: &Arc<CudaStream>) -> Result<Self, MLError> {
use std::sync::{Mutex, OnceLock};
static CACHE: OnceLock<Mutex<Vec<(usize, BiasKernels)>>> = OnceLock::new();
let cache = CACHE.get_or_init(|| Mutex::new(Vec::new()));
let ctx_key = Arc::as_ptr(stream.context()) as usize;
{
let guard = cache.lock().expect("BiasKernels cache poisoned");
for (key, kernels) in guard.iter() {
if *key == ctx_key {
return Ok(kernels.clone());
}
}
}
let kernels = Self::new(stream)?;
let mut guard = cache.lock().expect("BiasKernels cache poisoned");
// Re-check after taking the write lock — another thread may
// have raced us. If so, return the entry they inserted so the
// module Arc stays canonical.
for (key, k) in guard.iter() {
if *key == ctx_key {
return Ok(k.clone());
}
}
guard.push((ctx_key, kernels.clone()));
Ok(kernels)
}
}
/// Add bias to a 2D tensor: Y[b, j] += bias[j].
/// Launch `add_bias_2d_kernel` using a pre-loaded kernel handle.
/// The handle MUST come from a [`BiasKernels`] cached at construction
/// time — loading inside the call would re-enter `cuModuleLoadData`
/// which is incompatible with CUDA Graph stream capture.
fn add_bias_2d(
add_fn: &CudaFunction,
y: &mut GpuTensor,
bias: &CudaSlice<f32>,
rows: usize,
cols: usize,
stream: &Arc<CudaStream>,
) -> Result<(), MLError> {
let (add_fn, _) = get_bias_kernels(stream)?;
let total = rows * cols;
let threads = 256_u32;
let blocks = ((total as u32) + threads - 1) / threads;
@@ -746,7 +836,7 @@ fn add_bias_2d(
unsafe {
stream
.launch_builder(&add_fn)
.launch_builder(add_fn)
.arg(&y.data)
.arg(bias)
.arg(&rows_i32)
@@ -757,15 +847,20 @@ fn add_bias_2d(
Ok(())
}
/// Reduce a [rows, cols] tensor to [cols] by summing over axis 0.
/// Reduce a [rows, cols] tensor to [cols] by summing over axis 0,
/// using a pre-loaded kernel handle (see [`BiasKernels`]). The
/// handle MUST be cached at construction time — loading inside
/// the call would re-enter `cuModuleLoadData` and break CUDA Graph
/// stream capture.
fn reduce_sum_axis0(
reduce_fn: &CudaFunction,
x: &GpuTensor,
rows: usize,
cols: usize,
stream: &Arc<CudaStream>,
) -> Result<GpuTensor, MLError> {
let mut out = GpuTensor::zeros(&[cols], stream)?;
reduce_sum_axis0_into(x, rows, cols, stream, &mut out)?;
reduce_sum_axis0_into(reduce_fn, x, rows, cols, stream, &mut out)?;
Ok(out)
}
@@ -773,6 +868,7 @@ fn reduce_sum_axis0(
/// Used by `OwnedGpuLinear::backward_with_slices_into` to eliminate
/// per-call allocation on the training hot path.
pub fn reduce_sum_axis0_into(
reduce_fn: &CudaFunction,
x: &GpuTensor,
rows: usize,
cols: usize,
@@ -785,7 +881,6 @@ pub fn reduce_sum_axis0_into(
out.shape(), cols
)));
}
let (_, reduce_fn) = get_bias_kernels(stream)?;
let threads = 256_u32;
let blocks = ((cols as u32) + threads - 1) / threads;
let rows_i32 = rows as i32;
@@ -799,7 +894,7 @@ pub fn reduce_sum_axis0_into(
unsafe {
stream
.launch_builder(&reduce_fn)
.launch_builder(reduce_fn)
.arg(&x.data)
.arg(&out.data)
.arg(&rows_i32)

3
crates/ml-core/src/cuda_autograd/var_store.rs
View File

@@ -200,6 +200,7 @@ impl GpuVarStore {
bias_name: b_name,
in_dim,
out_dim,
bias_kernels: super::linear::BiasKernels::shared(&self.stream)?,
})
}
@@ -238,6 +239,7 @@ impl GpuVarStore {
bias_name: b_name,
in_dim,
out_dim,
bias_kernels: super::linear::BiasKernels::shared(&self.stream)?,
})
}
@@ -262,6 +264,7 @@ impl GpuVarStore {
bias_name: b_name,
in_dim,
out_dim,
bias_kernels: super::linear::BiasKernels::shared(&self.stream)?,
})
}