docs: implementation plan for single shared cuBLAS handle

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-13 11:32:23 +02:00
parent 44b07d1022
commit 6621177f78
1 changed files with 462 additions and 0 deletions
--- a/docs/superpowers/plans/2026-04-13-single-cublas-handle.md
+++ 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<u8>,
+    pub(crate) workspace_ptr: u64,
+    pub(crate) workspace_size: usize,
+    _lt_workspace_buf: CudaSlice<u8>,
+    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<CudaStream>,
+}
+```
+
+- [ ] **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<CudaStream>) -> Result<Self, MLError>`.
+
+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<SharedCublasHandle>` reference. The struct becomes:
+
+```rust
+pub struct CublasGemmSet {
+    handle: Arc<SharedCublasHandle>,
+    // 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<CudaStream>; 4],
+    _branch_workspace_bufs: [CudaSlice<u8>; 4],
+    pub(crate) branch_workspace_ptrs: [u64; 4],
+    trunk_done_event: CudaEvent,
+    branch_done_events: [CudaEvent; 4],
+    // Cached GEMM descriptors
+    gemm_cache: HashMap<FwdGemmKey, CachedGemmDesc>,
+    gemm_cache_relu_bias: HashMap<FwdGemmKey, CachedGemmDesc>,
+}
+```
+
+- [ ] **Step 2: Implement CublasGemmSet::new**
+
+Constructor takes `Arc<SharedCublasHandle>` + 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<SharedCublasHandle>,
+        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<Self, MLError> {
+        // ... 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<SharedCublasHandle>`. Keep relu_mask_kernel, bias_grad_kernel, gemm_cache, and network dimensions.
+
+```rust
+pub struct CublasBackwardSet {
+    handle: Arc<SharedCublasHandle>,
+    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<BwdGemmKey, CachedBwdGemmDesc>,
+}
+```
+
+- [ ] **Step 2: Update CublasBackwardSet::new**
+
+Takes `Arc<SharedCublasHandle>` + 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<SharedCublasHandle>,
+```
+
+- [ ] **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<SharedCublasHandle> {
+    &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<SharedCublasHandle>` and pass it to `CublasGemmSet::new` instead of creating its own `CublasForward`:
+
+```rust
+pub fn new(
+    stream: Arc<CudaStream>,
+    shared_cublas: Arc<SharedCublasHandle>,
+    // ... other params
+) -> Result<Self, MLError> {
+    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
+```