From 87ca1f5f55bbc5078ab4c44585fa0243bdfa0e1d Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 17 May 2026 15:21:00 +0200 Subject: [PATCH] perf(ml-alpha): CUDA Graph capture of training step (#162 finale) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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, 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 --- crates/ml-alpha/src/mamba2_block.rs | 76 +++++--- crates/ml-alpha/src/trainer/perception.rs | 130 ++++++++++---- crates/ml-core/src/action_space.rs | 83 ++++++--- crates/ml-core/src/cuda_autograd/linear.rs | 165 ++++++++++++++---- crates/ml-core/src/cuda_autograd/var_store.rs | 3 + 5 files changed, 334 insertions(+), 123 deletions(-) diff --git a/crates/ml-alpha/src/mamba2_block.rs b/crates/ml-alpha/src/mamba2_block.rs index 0525fb356..81051dc38 100644 --- a/crates/ml-alpha/src/mamba2_block.rs +++ b/crates/ml-alpha/src/mamba2_block.rs @@ -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, } 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::(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}"))?; diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index b89b55754..39f861bd0 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -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, + /// 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::(N_HORIZONS * n_hid)?, grad_heads_b_d: stream.alloc_zeros::(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::(); @@ -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 diff --git a/crates/ml-core/src/action_space.rs b/crates/ml-core/src/action_space.rs index a8134d4bd..750ab94a9 100644 --- a/crates/ml-core/src/action_space.rs +++ b/crates/ml-core/src/action_space.rs @@ -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 { - 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; 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"); } } } diff --git a/crates/ml-core/src/cuda_autograd/linear.rs b/crates/ml-core/src/cuda_autograd/linear.rs index 3bbb00a2a..27ff14780 100644 --- a/crates/ml-core/src/cuda_autograd/linear.rs +++ b/crates/ml-core/src/cuda_autograd/linear.rs @@ -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, + stream: &Arc, ) -> Result { 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`] 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, + batch: usize, weight: &CudaSlice, bias: &CudaSlice, cublas: &CudaBlas, stream: &Arc, 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, + batch: usize, weight: &CudaSlice, cublas: &CudaBlas, stream: &Arc, @@ -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) -> 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, +} + +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) -> Result { + 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`). The cache key is the context's + /// `Arc` pointer; a fresh context (e.g. a new test + /// process) gets a fresh entry. + pub fn shared(stream: &Arc) -> Result { + use std::sync::{Mutex, OnceLock}; + static CACHE: OnceLock>> = 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, rows: usize, cols: usize, stream: &Arc, ) -> 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, ) -> Result { 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) diff --git a/crates/ml-core/src/cuda_autograd/var_store.rs b/crates/ml-core/src/cuda_autograd/var_store.rs index f6fb6cdd3..58514840c 100644 --- a/crates/ml-core/src/cuda_autograd/var_store.rs +++ b/crates/ml-core/src/cuda_autograd/var_store.rs @@ -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)?, }) }