# 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
```
