# cublasLtMatmul Migration — TF32 Tensor Cores in CUDA Graph Mega-Capture

> **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:** Replace cublasGemmEx with cublasLtMatmul for TF32 tensor core acceleration inside CUDA Graph mega-capture on H100.

**Architecture:** The current cublasGemmEx with CUBLAS_COMPUTE_32F_FAST_TF32 hangs inside graph_mega on H100 due to workspace state conflicts during multi-op graph capture. cublasLtMatmul passes workspace and stream per-call (no handle state), eliminating this conflict. The migration replaces the internal GEMM implementation in `sgemm_f32` (forward) and `sgemm_f32` (backward) — all 29 call sites use these two methods, so only the method bodies change.

**Tech Stack:** cudarc vendored cublaslt (safe.rs + result.rs + sys), CUDA 12.9, H100 SM90

---

## Key Insight: Only 2 Methods Change

All 22 forward GEMM calls go through `CublasForward::sgemm_f32()` and `CublasForward::sgemm_f32_ldb()`.
All 7 backward GEMM calls go through `CublasBackward::sgemm_f32()`.

The call sites don't change. Only the internal implementation of these 3 methods changes from cublasGemmEx to cublasLtMatmul.

## Pre-created Descriptors

cublasLtMatmul requires matmul descriptors, matrix layouts, and algorithm heuristics. Creating these per-call is expensive. Instead:
- Create the matmul descriptor ONCE at construction (compute_type + transpose config)
- Create matrix layouts per-call (they're cheap and dimensions vary per layer)
- Cache algorithm heuristics per unique (m,n,k) shape via a HashMap

The workspace is already allocated by cudarc's `Workspace::new()` which auto-sizes to 32MB on Hopper.

---

### Task 1: Add cublasLt Handle + Workspace to CublasForward

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/batched_forward.rs:92-230`

- [ ] **Step 1: Add cublasLt fields to CublasForward struct**

Add to the struct (after the existing cuBLAS handle):

```rust
use cudarc::cublaslt::sys as cublaslt_sys;
use cudarc::cublaslt::result as cublaslt_result;

pub struct CublasForward {
    handle: SendSyncCublasHandle,
    // ... existing fields ...

    // cublasLt for TF32 tensor core GEMMs inside CUDA Graphs
    lt_handle: cublaslt_sys::cublasLtHandle_t,
    lt_workspace: CudaSlice<u8>,
    lt_workspace_ptr: u64,
    lt_workspace_size: usize,
}
```

Wrap lt_handle in a Send+Sync newtype (same pattern as SendSyncCublasHandle).

- [ ] **Step 2: Initialize cublasLt in constructor**

In `CublasForward::new()`, after the existing cuBLAS handle creation:

```rust
// cublasLt handle — stateless, no set_stream needed
let lt_handle = cublaslt_result::create_handle()
    .map_err(|e| MLError::ModelError(format!("cublasLtCreate: {e:?}")))?;

// Workspace: 32MB on Hopper (SM90+), 4MB otherwise
let major = stream.context()
    .attribute(cudarc::driver::sys::CUdevice_attribute::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR)
    .unwrap_or(8);
let lt_workspace_size = if major >= 9 { 32 * 1024 * 1024 } else { 4 * 1024 * 1024 };
let lt_workspace = stream.alloc_zeros::<u8>(lt_workspace_size)
    .map_err(|e| MLError::ModelError(format!("cublasLt workspace alloc: {e}")))?;
let lt_workspace_ptr = lt_workspace.raw_ptr();
```

- [ ] **Step 3: Add Drop impl for lt_handle**

```rust
// In the Drop for CublasForward (or in a wrapper struct):
unsafe { cublaslt_result::destroy_handle(self.lt_handle); }
```

- [ ] **Step 4: Compile and verify**

Run: `SQLX_OFFLINE=true cargo check -p ml`
Expected: Compiles clean

- [ ] **Step 5: Commit**

```bash
git add crates/ml/src/cuda_pipeline/batched_forward.rs
git commit -m "feat: add cublasLt handle + workspace to CublasForward"
```

---

### Task 2: Replace sgemm_f32 (forward) with cublasLtMatmul

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/batched_forward.rs` — `sgemm_f32()` method body

- [ ] **Step 1: Implement cublasLtMatmul in sgemm_f32**

Replace the cublasGemmEx body with cublasLtMatmul. The method signature stays identical (raw u64 pointers):

```rust
pub(crate) fn sgemm_f32(
    &self,
    w_f32_ptr: u64,
    a_f32_ptr: u64,
    c_f32_ptr: u64,
    n: usize,     // out_dim (rows of W^T = cols of result)
    b: usize,     // batch
    k: usize,     // in_dim
    _label: &str,
) -> Result<(), MLError> {
    let alpha = 1.0_f32;
    let beta = 0.0_f32;

    // cuBLAS column-major: our row-major C[B,N] = A[B,K] @ W[N,K]^T
    // becomes col-major: C'[N,B] = W[N,K] @ A'[K,B]
    // So: transa=N (W is N×K, no transpose), transb=N (A' is K×B, no transpose)
    // m=N, n=B, k=K, lda=K (W row stride), ldb=K (A row stride), ldc=N
    //
    // But our current cublasGemmEx uses transa=T, transb=N with W as A-matrix:
    //   GemmEx(OP_T, OP_N, N, B, K, W[K×N], K, A[K×B], K, C[N×B], N)
    // For cublasLtMatmul, set TRANSA=T on the matmul descriptor.

    unsafe {
        // Create matmul descriptor with TF32
        let matmul_desc = cublaslt_result::create_matmul_desc(
            cublaslt_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F_FAST_TF32,
            cublaslt_sys::cudaDataType_t::CUDA_R_32F,
        ).map_err(|e| MLError::ModelError(format!("cublasLtMatmulDescCreate {_label}: {e:?}")))?;

        // Set transpose A (weight matrix)
        let op_t = cublaslt_sys::cublasOperation_t::CUBLAS_OP_T as i32;
        let op_n = cublaslt_sys::cublasOperation_t::CUBLAS_OP_N as i32;
        cublaslt_result::set_matmul_desc_attribute(
            matmul_desc,
            cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSA,
            &op_t as *const _ as *const _,
            std::mem::size_of::<i32>(),
        ).map_err(|e| MLError::ModelError(format!("set TRANSA {_label}: {e:?}")))?;
        cublaslt_result::set_matmul_desc_attribute(
            matmul_desc,
            cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSB,
            &op_n as *const _ as *const _,
            std::mem::size_of::<i32>(),
        ).map_err(|e| MLError::ModelError(format!("set TRANSB {_label}: {e:?}")))?;

        // Matrix layouts: A=W[K,N] (before transpose), B=A[K,B], C=D=C[N,B]
        let f32_type = cublaslt_sys::cudaDataType_t::CUDA_R_32F;
        let a_layout = cublaslt_result::create_matrix_layout(f32_type, k as u64, n as u64, k as i64)
            .map_err(|e| MLError::ModelError(format!("A layout {_label}: {e:?}")))?;
        let b_layout = cublaslt_result::create_matrix_layout(f32_type, k as u64, b as u64, k as i64)
            .map_err(|e| MLError::ModelError(format!("B layout {_label}: {e:?}")))?;
        let c_layout = cublaslt_result::create_matrix_layout(f32_type, n as u64, b as u64, n as i64)
            .map_err(|e| MLError::ModelError(format!("C layout {_label}: {e:?}")))?;

        // Algorithm heuristic
        let pref = cublaslt_result::create_matmul_pref()
            .map_err(|e| MLError::ModelError(format!("matmul pref {_label}: {e:?}")))?;
        cublaslt_result::set_matmul_pref_attribute(
            pref,
            cublaslt_sys::cublasLtMatmulPreferenceAttributes_t::CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
            &self.lt_workspace_size as *const _ as *const _,
            std::mem::size_of::<usize>(),
        ).map_err(|e| MLError::ModelError(format!("set pref workspace {_label}: {e:?}")))?;

        let heuristic = cublaslt_result::get_matmul_algo_heuristic(
            self.lt_handle, matmul_desc, a_layout, b_layout, c_layout, c_layout, pref,
        ).map_err(|e| MLError::ModelError(format!("algo heuristic {_label}: {e:?}")))?;

        // Execute matmul — per-call workspace + stream, no handle state!
        let stream_ptr = self.stream().cu_stream() as cublaslt_sys::cudaStream_t;
        cublaslt_result::matmul(
            self.lt_handle,
            matmul_desc,
            &alpha as *const _ as *const _,
            &beta as *const _ as *const _,
            w_f32_ptr as *const _,    // A (weight)
            a_layout,
            a_f32_ptr as *const _,    // B (input)
            b_layout,
            c_f32_ptr as *const _,    // C (for beta accumulation)
            c_layout,
            c_f32_ptr as *mut _,      // D (output, same as C)
            c_layout,
            &heuristic.algo as *const _,
            self.lt_workspace_ptr as *mut _,
            self.lt_workspace_size,
            stream_ptr,
        ).map_err(|e| MLError::ModelError(format!("cublasLtMatmul {_label}: {e:?}")))?;

        // Cleanup descriptors
        cublaslt_result::destroy_matrix_layout(a_layout).ok();
        cublaslt_result::destroy_matrix_layout(b_layout).ok();
        cublaslt_result::destroy_matrix_layout(c_layout).ok();
        cublaslt_result::destroy_matmul_desc(matmul_desc).ok();
        cublaslt_result::destroy_matmul_pref(pref).ok();
    }
    Ok(())
}
```

- [ ] **Step 2: Implement sgemm_f32_ldb similarly**

Same as sgemm_f32 but the B-matrix layout uses `ldb` instead of `k` for its leading dimension:
```rust
let b_layout = cublaslt_result::create_matrix_layout(f32_type, k as u64, b as u64, ldb as i64)
```

- [ ] **Step 3: Remove the cublasSetStream workspace restoration**

Since cublasLtMatmul takes stream per-call, the multi-stream branch dispatch no longer needs cublasSetStream for the GEMM handle. Remove `restore_workspace()` and the workspace ptr/size fields from the struct. Keep the old cuBLAS handle for now (used by set_stream in non-GEMM paths).

Actually — the branch dispatch still uses `cublasSetStream` for the `sgemm_f32` calls. But now `sgemm_f32` uses cublasLtMatmul which doesn't use the cuBLAS handle at all. We need to pass the branch stream to sgemm_f32 instead.

Add a `stream` parameter to `sgemm_f32`:
```rust
pub(crate) fn sgemm_f32(
    &self,
    stream: &CudaStream,  // NEW: per-call stream
    w_f32_ptr: u64,
    ...
```

Update all 22 call sites to pass the appropriate stream (main stream or branch stream).

- [ ] **Step 4: Run smoke test locally**

Run: `SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests::regression::test_no_hang_single_epoch --ignored --nocapture`
Expected: PASS with healthy grad_norm

- [ ] **Step 5: Commit**

```bash
git add crates/ml/src/cuda_pipeline/batched_forward.rs
git commit -m "feat: replace cublasGemmEx with cublasLtMatmul in forward (TF32 + CUDA Graph safe)"
```

---

### Task 3: Replace sgemm_f32 (backward) with cublasLtMatmul

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/batched_backward.rs`

- [ ] **Step 1: Add cublasLt handle + workspace to CublasBackward**

Same pattern as forward. The backward sgemm_f32 takes flexible (transa, transb, m, n, k) arguments:

```rust
fn sgemm_f32(
    &self,
    transa: cublas_sys::cublasOperation_t,
    transb: cublas_sys::cublasOperation_t,
    m: i32, n: i32, k: i32,
    alpha: f32,
    a: u64, lda: i32,
    b_ptr: u64, ldb: i32,
    beta: f32,
    c: u64, ldc: i32,
    label: &str,
) -> Result<(), MLError>
```

Set TRANSA/TRANSB on the matmul descriptor based on the transa/transb parameters. Matrix layout dimensions depend on transpose flags:
- If transa=N: A is [m, k], lda
- If transa=T: A is [k, m], lda

- [ ] **Step 2: Remove old cuBLAS handle from backward**

The backward cuBLAS handle is no longer needed for GEMMs. But `set_stream` is called on it. Check if any non-GEMM operations use the backward cuBLAS handle. If not, remove it entirely.

- [ ] **Step 3: Run all smoke tests**

Run: `SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture`
Expected: 19/19 pass

- [ ] **Step 4: Commit**

```bash
git add crates/ml/src/cuda_pipeline/batched_backward.rs
git commit -m "feat: replace cublasGemmEx with cublasLtMatmul in backward (TF32 + CUDA Graph safe)"
```

---

### Task 4: Deploy to H100 and Verify Performance

- [ ] **Step 1: Push and deploy**

```bash
git push
./scripts/argo-train.sh dqn --baseline --epochs 5
```

- [ ] **Step 2: Verify epoch duration**

Expected: ~10s/epoch (TF32 tensor cores via cublasLtMatmul)
Previous: 77s/epoch (cublasSgemm CUDA cores), hang (cublasGemmEx in graph_mega)

- [ ] **Step 3: Verify gradient norms across C51 transition**

RECAPTURE_DIAG should show non-zero grad_norm with c51_alpha > 0.

- [ ] **Step 4: Commit final with performance note**

---

### Task 5: Cleanup Dead Code

- [ ] **Step 1: Remove old cuBLAS handle from CublasForward if no longer needed**

Check if any non-GEMM cuBLAS operations still use `self.handle`. If not, remove:
- `SendSyncCublasHandle` wrapper
- `handle` field
- `_workspace_buf`, `workspace_ptr`, `workspace_size` fields
- `set_stream`, `restore_workspace` methods
- `cublasSetMathMode` call

- [ ] **Step 2: Remove old cuBLAS handle from CublasBackward if no longer needed**

Same check.

- [ ] **Step 3: Run full test suite + deploy**

```bash
SQLX_OFFLINE=true cargo test -p ml --lib -- smoke_tests --ignored
git push && ./scripts/argo-train.sh dqn --baseline --epochs 20
```
