diff --git a/docs/superpowers/plans/2026-04-13-single-cublas-handle.md b/docs/superpowers/plans/2026-04-13-single-cublas-handle.md new file mode 100644 index 000000000..61eebebcc --- /dev/null +++ b/docs/superpowers/plans/2026-04-13-single-cublas-handle.md @@ -0,0 +1,462 @@ +# Single Shared cuBLAS Handle Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Eliminate cuBLAS non-determinism by sharing a single cuBLAS handle across all components on the same CUDA stream. + +**Architecture:** Split `CublasForward` into `SharedCublasHandle` (one per stream, `Arc`-shared) and `CublasGemmSet` (per-component cached descriptors). Same split for `CublasBackward` → `CublasBackwardSet`. The experience collector, training forward, DDQN forward, and backward all share the same handle. + +**Tech Stack:** Rust, cudarc (cuBLAS/cublasLt FFI), CUDA 12.9 + +--- + +### Task 1: Create SharedCublasHandle + +**Files:** +- Create: `crates/ml/src/cuda_pipeline/shared_cublas_handle.rs` +- Modify: `crates/ml/src/cuda_pipeline/mod.rs` + +- [ ] **Step 1: Create the shared handle struct** + +Create `crates/ml/src/cuda_pipeline/shared_cublas_handle.rs`: + +```rust +#![allow(unsafe_code)] + +//! Single shared cuBLAS handle for deterministic training. +//! +//! NVIDIA guarantees bit-wise reproducibility only with one cuBLAS handle +//! per CUDA stream. This module provides a shared handle used by all +//! components (training forward, backward, experience collector, DDQN). + +use std::sync::Arc; +use cudarc::driver::{CudaFunction, CudaSlice, CudaStream}; +use crate::MLError; + +// Re-export the cublas result/sys modules used by consumers +use cudarc::cublas::result as cublas_result; +use cudarc::cublas::sys as cublas_sys; +use cudarc::cublaslt::result as cublaslt_result; +use cudarc::cublaslt::sys as cublaslt_sys; + +/// Send+Sync wrapper for cuBLAS handle. +pub(crate) struct SendSyncCublasHandle(pub(crate) cublas_sys::cublasHandle_t); +unsafe impl Send for SendSyncCublasHandle {} +unsafe impl Sync for SendSyncCublasHandle {} +impl Drop for SendSyncCublasHandle { + fn drop(&mut self) { + unsafe { let _ = cublas_result::destroy_handle(self.0); } + } +} + +/// Send+Sync wrapper for cublasLt handle. +pub(crate) struct SendSyncCublasLtHandle(pub(crate) cublaslt_sys::cublasLtHandle_t); +unsafe impl Send for SendSyncCublasLtHandle {} +unsafe impl Sync for SendSyncCublasLtHandle {} +impl Drop for SendSyncCublasLtHandle { + fn drop(&mut self) { + unsafe { let _ = cublaslt_result::destroy_handle(self.0); } + } +} + +/// Single shared cuBLAS context — one per CUDA stream. +/// +/// Created once by the training system, shared via `Arc` with all components +/// that need cuBLAS GEMM operations (forward, backward, experience collector). +/// NVIDIA guarantees bit-wise reproducible results only with one handle per stream. +#[allow(missing_debug_implementations)] +pub struct SharedCublasHandle { + pub(crate) handle: SendSyncCublasHandle, + pub(crate) lt_handle: SendSyncCublasLtHandle, + // Workspace buffers — must outlive the handle for CUDA Graph replay + _workspace_buf: CudaSlice, + pub(crate) workspace_ptr: u64, + pub(crate) workspace_size: usize, + _lt_workspace_buf: CudaSlice, + pub(crate) lt_workspace_ptr: u64, + pub(crate) lt_workspace_size: usize, + // Compiled bias+activation kernels (shared across all GemmSets) + pub(crate) add_bias_relu_f32_kernel: CudaFunction, + pub(crate) add_bias_f32_kernel: CudaFunction, + // Stream reference + pub(crate) stream: Arc, +} +``` + +- [ ] **Step 2: Implement SharedCublasHandle::new** + +The constructor consolidates handle creation from `CublasForward::new` and `CublasBackward::new`. Extract lines 255-410 from `batched_forward.rs` (handle creation, workspace, math mode, Lt handle, kernel compilation) into `SharedCublasHandle::new(stream: &Arc) -> Result`. + +Key details: +- `cublasSetMathMode(handle, CUBLAS_PEDANTIC_MATH)` +- `cublasSetWorkspace_v2(handle, workspace, 32MB)` +- cublasLt handle creation + 32MB Lt workspace +- Compile `add_bias_relu_f32_kernel` and `add_bias_f32_kernel` from the existing cubin +- NO network dimensions, NO cached descriptors, NO branch streams + +- [ ] **Step 3: Add `set_stream` method** + +```rust +impl SharedCublasHandle { + pub(crate) fn set_stream(&self, stream: &CudaStream) -> Result<(), MLError> { + unsafe { + let cu_stream = stream.cu_stream() as *mut cublas_sys::CUstream_st; + cublas_result::set_stream(self.handle.0, cu_stream) + .map_err(|e| MLError::ModelError(format!("cublasSetStream: {e:?}")))?; + // Restore workspace (cublasSetStream resets it) + let ws_ptr = self.workspace_ptr as *mut std::ffi::c_void; + cublas_sys::cublasSetWorkspace_v2(self.handle.0, ws_ptr, self.workspace_size); + } + Ok(()) + } +} +``` + +- [ ] **Step 4: Register module** + +Add to `crates/ml/src/cuda_pipeline/mod.rs`: +```rust +pub mod shared_cublas_handle; +``` + +- [ ] **Step 5: Build check** + +Run: `SQLX_OFFLINE=true cargo check -p ml` +Expected: Clean build (new module, no consumers yet) + +- [ ] **Step 6: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/shared_cublas_handle.rs crates/ml/src/cuda_pipeline/mod.rs +git commit -m "feat: SharedCublasHandle — single cuBLAS handle per stream" +``` + +--- + +### Task 2: Create CublasGemmSet (forward descriptors) + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/batched_forward.rs` + +- [ ] **Step 1: Add CublasGemmSet struct** + +Keep `CachedGemmDesc`, `FwdGemmKey`, and all the GEMM dispatch methods. Replace the handle/workspace fields with an `Arc` reference. The struct becomes: + +```rust +pub struct CublasGemmSet { + handle: Arc, + // Network dimensions + pub(crate) batch_size: usize, + pub(crate) state_dim: usize, + pub(crate) state_dim_padded: usize, + pub(crate) s1_input_dim: usize, + pub(crate) s1_ldb: usize, + pub(crate) shared_h1: usize, + pub(crate) shared_h2: usize, + pub(crate) value_h: usize, + pub(crate) adv_h: usize, + pub(crate) num_atoms: usize, + pub(crate) branch_0_size: usize, + pub(crate) branch_1_size: usize, + pub(crate) branch_2_size: usize, + pub(crate) branch_3_size: usize, + // Branch dispatch (parallel advantage heads) + branch_streams: [Arc; 4], + _branch_workspace_bufs: [CudaSlice; 4], + pub(crate) branch_workspace_ptrs: [u64; 4], + trunk_done_event: CudaEvent, + branch_done_events: [CudaEvent; 4], + // Cached GEMM descriptors + gemm_cache: HashMap, + gemm_cache_relu_bias: HashMap, +} +``` + +- [ ] **Step 2: Implement CublasGemmSet::new** + +Constructor takes `Arc` + network dimensions. Creates branch streams, allocates branch workspaces, pre-caches GEMM descriptors using `handle.lt_handle` for `AlgoGetIds` + `AlgoCheck`. Does NOT create any cuBLAS handle. + +```rust +impl CublasGemmSet { + pub fn new( + handle: Arc, + batch_size: usize, + state_dim: usize, + shared_h1: usize, shared_h2: usize, + value_h: usize, adv_h: usize, + num_atoms: usize, + branch_0_size: usize, branch_1_size: usize, + branch_2_size: usize, branch_3_size: usize, + s1_input_dim: usize, + ) -> Result { + // ... fork branch streams, allocate branch workspaces, + // pre-cache GEMM descriptors using handle.lt_handle + } +} +``` + +- [ ] **Step 3: Move GEMM dispatch methods** + +Move `sgemm_f32`, `sgemm_f32_branch`, `sgemm_f32_ldb`, `sgemm_f32_fused_relu_bias`, `lt_matmul`, `lt_matmul_cached`, `lt_matmul_uncached`, `forward_online_f32`, `forward_online_raw`, `forward_value_head`, `launch_add_bias_f32_raw`, `launch_add_bias_relu_f32_raw`, `set_stream` from `CublasForward` to `CublasGemmSet`. + +Each method that previously used `self.handle` / `self.lt_handle` now uses `self.handle.handle` / `self.handle.lt_handle`. Each method that used `self.lt_workspace_ptr` now uses `self.handle.lt_workspace_ptr`. Bias/relu kernels: `self.handle.add_bias_relu_f32_kernel`. + +- [ ] **Step 4: Keep CublasForward as a type alias (temporary)** + +During migration, keep `CublasForward` as a type alias or thin wrapper around `CublasGemmSet` to avoid breaking all callers at once: + +```rust +pub type CublasForward = CublasGemmSet; +``` + +This lets existing code compile while we migrate callers one at a time. + +- [ ] **Step 5: Build check** + +Run: `SQLX_OFFLINE=true cargo check -p ml` +Expected: Clean build (CublasForward is now CublasGemmSet via alias) + +- [ ] **Step 6: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/batched_forward.rs +git commit -m "refactor: CublasForward → CublasGemmSet with shared handle" +``` + +--- + +### Task 3: Create CublasBackwardSet + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/batched_backward.rs` + +- [ ] **Step 1: Replace CublasBackward internals with shared handle** + +Same pattern as Task 2. Replace handle/lt_handle/workspace fields with `Arc`. Keep relu_mask_kernel, bias_grad_kernel, gemm_cache, and network dimensions. + +```rust +pub struct CublasBackwardSet { + handle: Arc, + relu_mask_kernel: CudaFunction, + bias_grad_kernel: CudaFunction, + // Network dimensions + batch_size: usize, + state_dim: usize, + state_dim_padded: usize, + s1_input_dim: usize, + shared_h1: usize, shared_h2: usize, + value_h: usize, adv_h: usize, + num_atoms: usize, + branch_0_size: usize, branch_1_size: usize, + branch_2_size: usize, branch_3_size: usize, + gemm_cache: HashMap, +} +``` + +- [ ] **Step 2: Update CublasBackwardSet::new** + +Takes `Arc` + config. Compiles relu_mask + bias_grad kernels. Pre-caches backward GEMM descriptors using `handle.lt_handle`. Does NOT create any cuBLAS handle. + +- [ ] **Step 3: Type alias for migration** + +```rust +pub type CublasBackward = CublasBackwardSet; +``` + +- [ ] **Step 4: Build check + commit** + +Run: `SQLX_OFFLINE=true cargo check -p ml` + +```bash +git add crates/ml/src/cuda_pipeline/batched_backward.rs +git commit -m "refactor: CublasBackward → CublasBackwardSet with shared handle" +``` + +--- + +### Task 4: Wire shared handle through GpuDqnTrainer + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` + +- [ ] **Step 1: Add shared handle field** + +Add to GpuDqnTrainer struct: +```rust +shared_cublas: Arc, +``` + +- [ ] **Step 2: Create shared handle in constructor** + +In `GpuDqnTrainer::new()`, create ONE `SharedCublasHandle` before any `CublasGemmSet`: + +```rust +let shared_cublas = Arc::new(SharedCublasHandle::new(&stream)?); +let cublas_forward = CublasGemmSet::new(Arc::clone(&shared_cublas), batch_size, ...)?; +let cublas_forward_ddqn = CublasGemmSet::new(Arc::clone(&shared_cublas), batch_size, ...)?; +let cublas_backward = CublasBackwardSet::new(Arc::clone(&shared_cublas), &config)?; +``` + +- [ ] **Step 3: Delete duplicate handle wrapper types** + +Delete `SendSyncCublasHandle` and `SendSyncCublasLtHandle` from `batched_forward.rs` and `batched_backward.rs` — they now live in `shared_cublas_handle.rs`. + +- [ ] **Step 4: Delete double_dqn_stream** + +Remove `double_dqn_stream` field, its fork in the constructor, the pass1_event / pass3_event sync logic in `replay_forward`. DDQN already runs on the main stream. + +- [ ] **Step 5: Expose shared handle for experience collector** + +Add accessor: +```rust +pub(crate) fn shared_cublas(&self) -> &Arc { + &self.shared_cublas +} +``` + +- [ ] **Step 6: Build check + commit** + +Run: `SQLX_OFFLINE=true cargo check -p ml` + +```bash +git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs crates/ml/src/cuda_pipeline/batched_forward.rs crates/ml/src/cuda_pipeline/batched_backward.rs +git commit -m "refactor: GpuDqnTrainer uses single SharedCublasHandle" +``` + +--- + +### Task 5: Wire shared handle through experience collector + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` +- Modify: `crates/ml/src/trainers/dqn/fused_training.rs` + +- [ ] **Step 1: Accept shared handle in constructor** + +Change `GpuExperienceCollector::new` to accept `Arc` and pass it to `CublasGemmSet::new` instead of creating its own `CublasForward`: + +```rust +pub fn new( + stream: Arc, + shared_cublas: Arc, + // ... other params +) -> Result { + let cublas_forward = CublasGemmSet::new( + shared_cublas, + alloc_episodes, state_dim, shared_h1, shared_h2, + value_h, adv_h, num_atoms, + branch_sizes[0], branch_sizes[1], branch_sizes[2], branch_sizes[3], + s1_input_dim, + )?; + // ... +} +``` + +- [ ] **Step 2: Wire through FusedTrainingCtx** + +In `fused_training.rs`, pass the shared handle from the trainer to the experience collector: + +```rust +// In FusedTrainingCtx::new or wherever the collector is constructed: +let shared_cublas = Arc::clone(trainer.shared_cublas()); +// Pass to experience collector constructor +``` + +- [ ] **Step 3: Update all GpuExperienceCollector::new call sites** + +Search for all `GpuExperienceCollector::new(` calls and add the `shared_cublas` parameter. + +- [ ] **Step 4: Build check + commit** + +Run: `SQLX_OFFLINE=true cargo check -p ml` + +```bash +git add crates/ml/src/cuda_pipeline/gpu_experience_collector.rs crates/ml/src/trainers/dqn/fused_training.rs +git commit -m "refactor: experience collector shares cuBLAS handle with trainer" +``` + +--- + +### Task 6: Delete dead code + type aliases + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/batched_forward.rs` +- Modify: `crates/ml/src/cuda_pipeline/batched_backward.rs` +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` + +- [ ] **Step 1: Remove CublasForward type alias** + +Delete `pub type CublasForward = CublasGemmSet;`. Update all imports/references across the codebase to use `CublasGemmSet` directly. + +- [ ] **Step 2: Remove CublasBackward type alias** + +Delete `pub type CublasBackward = CublasBackwardSet;`. Update all references. + +- [ ] **Step 3: Remove dead fields from GpuDqnTrainer** + +Delete: +- `double_dqn_stream` field +- `pass1_event`, `pass3_event` fields +- `eval_td_snapshot`, `eval_loss_snapshot` fields (if still present) +- Any `exp_fwd_graph` leak infrastructure + +- [ ] **Step 4: Remove dead code from experience collector** + +Delete `exp_fwd_graph` field and the graph capture/replay code (the shared handle + CUDA Graph on training side provides determinism). + +- [ ] **Step 5: Build check + smoke test** + +Run: `SQLX_OFFLINE=true cargo check -p ml` +Run: `SQLX_OFFLINE=true cargo build -p ml --lib` +Run smoke test: +```bash +FOXHUNT_TEST_DATA=test_data/futures-baseline SQLX_OFFLINE=true cargo test -p ml --lib -- smoke_tests::walk_forward::test_walk_forward_oos_metrics --ignored --nocapture 2>&1 | grep "val_Sharpe" | head -5 +``` + +Expected: Test passes, val_Sharpe values appear. + +- [ ] **Step 6: Commit** + +```bash +git add -A +git commit -m "cleanup: remove dead code, type aliases, duplicate handles" +``` + +--- + +### Task 7: Determinism verification + +**Files:** None (test only) + +- [ ] **Step 1: Run 4-way determinism test** + +```bash +for i in 1 2 3 4; do + FOXHUNT_TEST_DATA=test_data/futures-baseline SQLX_OFFLINE=true \ + cargo test -p ml --lib -- smoke_tests::walk_forward::test_walk_forward_oos_metrics \ + --ignored --nocapture 2>&1 | grep -oP 'val_Sharpe=[\-0-9.]+' > /tmp/run_shared_${i}.txt +done +diff /tmp/run_shared_1.txt /tmp/run_shared_2.txt && \ +diff /tmp/run_shared_2.txt /tmp/run_shared_3.txt && \ +diff /tmp/run_shared_3.txt /tmp/run_shared_4.txt && \ +echo "ALL 4 IDENTICAL — FULLY DETERMINISTIC" || echo "DIFFERENT" +``` + +Expected: `ALL 4 IDENTICAL — FULLY DETERMINISTIC` + +- [ ] **Step 2: Commit final state** + +```bash +git add -A +git commit -m "feat: fully deterministic training + evaluation via single cuBLAS handle + +NVIDIA guarantees bit-wise reproducibility with one cuBLAS handle per stream. +SharedCublasHandle provides a single handle shared by all components: +- Training forward (CublasGemmSet, batch=64) +- DDQN forward (CublasGemmSet, batch=64) +- Experience collector (CublasGemmSet, batch=32) +- Backward (CublasBackwardSet) + +Verified: 4 consecutive runs produce identical val_Sharpe across all 30 epochs." +git push +```