gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: CublasForward → CublasGemmSet with shared handle

Browse Source 
Replace per-instance cuBLAS/cublasLt handle ownership in CublasForward
with Arc<SharedCublasHandle>. The struct is renamed to CublasGemmSet
and a type alias preserves backward compatibility for external callers.

Removed from CublasGemmSet: handle, workspace, lt_handle, lt_workspace,
bias kernels (all now in SharedCublasHandle). Kept: branch streams,
branch workspaces, GEMM caches, events, network dimensions.

External callers (gpu_dqn_trainer, gpu_experience_collector) will be
updated in Tasks 4-5 to pass Arc<SharedCublasHandle> to the constructor.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-13 11:45:21 +02:00
parent 334ac22d30
commit ac0ae5b5ae
1 changed files with 45 additions and 207 deletions
Download Patch File Download Diff File Expand all files Collapse all files

252
crates/ml/src/cuda_pipeline/batched_forward.rs
View File

@@ -59,11 +59,9 @@ use std::collections::HashMap;
use std::sync::Arc;
use std::mem::ManuallyDrop;
use cudarc::cublas::result as cublas_result;
use cudarc::cublas::sys as cublas_sys;
use cudarc::cublaslt::sys as cublaslt_sys;
use cudarc::cublaslt::result as cublaslt_result;
use cudarc::driver::{CudaEvent, CudaFunction, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg};
use cudarc::driver::{CudaEvent, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg};
use crate::MLError;
@@ -88,7 +86,7 @@ struct CachedGemmDesc {
}
/// SAFETY: cublasLt descriptors are opaque pointers owned exclusively by
/// `CublasForward`, which is single-threaded in practice (owned by `GpuDqnTrainer`).
/// `CublasGemmSet`, which is single-threaded in practice (owned by `GpuDqnTrainer`).
unsafe impl Send for CachedGemmDesc {}
unsafe impl Sync for CachedGemmDesc {}
@@ -110,79 +108,20 @@ impl Drop for CachedGemmDesc {
/// All forward GEMMs use TRANSA=T, TRANSB=N so transpose flags are implicit.
type FwdGemmKey = (usize, usize, usize, usize);
// ── cuBLAS handle wrapper ─────────────────────────────────────────────────────
/// Wrapper to make `cublasHandle_t` Send + Sync.
///
/// cuBLAS handles are not thread-safe by default, but `CublasForward` is
/// owned exclusively by `GpuDqnTrainer` which is single-threaded in practice.
struct SendSyncCublasHandle(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);
}
}
}
/// Wrapper to make `cublasLtHandle_t` Send + Sync.
///
/// Same ownership model as `SendSyncCublasHandle` — exclusively owned by
/// `CublasForward` which is single-threaded in practice.
struct SendSyncCublasLtHandle(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);
}
}
}
// ── Batched forward context ───────────────────────────────────────────────────
/// cuBLAS-based batched forward pass context for the branching dueling DQN.
/// cuBLAS-based batched forward GEMM set for the branching dueling DQN.
///
/// Holds the cuBLAS handle and references to the trainer's pre-allocated
/// activation buffers. GEMM output lands directly in those buffers.
/// Holds an `Arc<SharedCublasHandle>` (cuBLAS + cublasLt handles, workspaces,
/// bias kernels) and the forward-specific resources: branch streams, per-branch
/// workspaces, cached GEMM descriptors, and network dimensions.
///
/// The handle is set to the trainer's stream so all GEMMs execute in order.
#[allow(missing_debug_implementations)] // handle is not Debug
pub struct CublasForward {
handle: SendSyncCublasHandle,
/// cuBLAS workspace buffer — must outlive the handle for CUDA Graph replay.
_workspace_buf: CudaSlice<u8>,
/// Raw device pointer + size for workspace restoration after cublasSetStream.
/// cublasSetStream resets workspace to the default pool, so we must re-set it.
workspace_ptr: u64,
workspace_size: usize,
// ── cublasLt handle + workspace (per-call stream, no handle state conflicts) ──
/// cublasLt handle for forward GEMMs — passes workspace+stream per-call,
/// eliminating the `cublasSetStream` workspace state conflict that hangs
/// CUDA Graph mega-capture on H100.
lt_handle: SendSyncCublasLtHandle,
/// cublasLt workspace buffer — must outlive handle for CUDA Graph replay.
_lt_workspace_buf: CudaSlice<u8>,
/// Raw device pointer for cublasLt workspace.
lt_workspace_ptr: u64,
/// Size of cublasLt workspace in bytes.
lt_workspace_size: usize,
// ── Compiled bias+activation kernels (loaded from precompiled cubin) ──
/// `add_bias_relu_f32_kernel(output, bias, out_dim, total)` — f32 output + f32 bias + ReLU.
add_bias_relu_f32_kernel: CudaFunction,
/// `add_bias_f32_kernel(output, bias, out_dim, total)` — f32 output + f32 bias (no activation).
add_bias_f32_kernel: CudaFunction,
pub struct CublasGemmSet {
/// Shared cuBLAS + cublasLt handle (owned via Arc, shared with backward pass
/// and experience collector for deterministic single-handle training).
handle: Arc<super::shared_cublas_handle::SharedCublasHandle>,
// ── Network dimensions (baked at construction) ──
batch_size: usize,
@@ -232,13 +171,20 @@ pub struct CublasForward {
gemm_cache_relu_bias: HashMap<FwdGemmKey, CachedGemmDesc>,
}
impl CublasForward {
/// Create a new cuBLAS forward context.
/// Backward-compatible type alias — existing callers (`GpuDqnTrainer`,
/// `GpuExperienceCollector`, etc.) keep using `CublasForward` until Task 4
/// wires the shared handle through all constructors.
pub type CublasForward = CublasGemmSet;
impl CublasGemmSet {
/// Create a new cuBLAS forward GEMM set backed by a shared handle.
///
/// Compiles the `add_bias_relu` and `add_bias` kernels from inline NVRTC source.
/// Binds the cuBLAS handle to the provided stream.
/// The shared handle owns the cuBLAS + cublasLt handles, workspaces, and
/// bias kernels. This constructor only creates forward-specific resources:
/// branch streams, per-branch workspaces, cached GEMM descriptors, events,
/// and stores the network dimensions.
pub fn new(
stream: &Arc<CudaStream>,
shared: Arc<super::shared_cublas_handle::SharedCublasHandle>,
batch_size: usize,
state_dim: usize,
shared_h1: usize,
@@ -252,70 +198,7 @@ impl CublasForward {
branch_3_size: usize,
s1_input_dim: usize,
) -> Result<Self, MLError> {
// ── Create cuBLAS handle ────────────────────────────────────────
let raw_handle = cublas_result::create_handle()
.map_err(|e| MLError::ModelError(format!("cublasCreate: {e:?}")))?;
// Bind handle to stream — ALL GEMMs will execute on this stream.
// SAFETY: cudarc::driver::sys and cudarc::cublas::sys both typedef
// CUstream to *mut CUstream_st, which is the same underlying type.
// The cast is safe as long as the pointer value is preserved.
unsafe {
let cu_stream = stream.cu_stream() as *mut cublas_sys::CUstream_st;
cublas_result::set_stream(raw_handle, cu_stream)
.map_err(|e| MLError::ModelError(format!("cublasSetStream: {e:?}")))?;
}
// Deterministic FP32 GEMM — ensures reproducible training across runs.
// TF32 tensor ops (CUBLAS_TF32_TENSOR_OP_MATH) use non-deterministic
// reduction order, causing ~1e-4 per-GEMM variance that compounds over
// thousands of forward+backward passes. DEFAULT_MATH forces IEEE FP32
// accumulation which is deterministic. Cost: ~1.5× slower GEMM (but GEMM
// is <20% of total step time — most time is in C51 loss/grad kernels).
unsafe {
cublas_sys::cublasSetMathMode(
raw_handle,
cublas_sys::cublasMath_t::CUBLAS_DEFAULT_MATH,
);
}
// Allocate explicit cuBLAS workspace for CUDA Graph compatibility.
// CRITICAL: cublasSetStream() resets workspace to the default pool,
// which uses stream-ordered allocation (incompatible with CUDA Graph).
// We must re-set this workspace after every cublasSetStream call.
// 32 MB: TF32 HMMA algorithms on H100 need more workspace than f32.
let workspace_size: usize = 32 * 1024 * 1024; // 32 MB — H100 CUDA 13.0 requires this for TF32
let workspace_buf = stream.alloc_zeros::<u8>(workspace_size)
.map_err(|e| MLError::ModelError(format!("cuBLAS workspace alloc: {e}")))?;
let workspace_ptr = {
let (dev_ptr, _guard) = workspace_buf.device_ptr(stream);
dev_ptr
};
unsafe {
let ws_ptr = workspace_ptr as *mut std::ffi::c_void;
let status = cublas_sys::cublasSetWorkspace_v2(raw_handle, ws_ptr, workspace_size);
if status != cublas_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS {
return Err(MLError::ModelError(format!("cublasSetWorkspace_v2 failed: {status:?}")));
}
}
// ── Create cublasLt handle + workspace ────────────────────────────
// cublasLtMatmul passes workspace+stream per-call (no handle state),
// eliminating the cublasSetStream workspace conflict that hangs CUDA
// Graph mega-capture on H100.
let lt_raw_handle = cublaslt_result::create_handle()
.map_err(|e| MLError::ModelError(format!("cublasLtCreate: {e:?}")))?;
// 32 MB workspace for cublasLtMatmul.
// cudarc now uses cuMemAlloc_v2 (sync, globally accessible) — see vendor/cudarc comment.
let lt_ws_size: usize = 32 * 1024 * 1024;
let lt_workspace_buf = stream.alloc_zeros::<u8>(lt_ws_size)
.map_err(|e| MLError::ModelError(format!("cublasLt workspace alloc: {e}")))?;
let lt_ws_ptr = lt_workspace_buf.raw_ptr();
// ── Load f32 bias kernels from precompiled cubin ──
let (add_bias_relu_f32_kernel, add_bias_f32_kernel) =
compile_bias_kernels(stream)?;
let stream = &shared.stream;
// ── Fork 4 branch streams for parallel advantage head dispatch ──
let branch_streams = [
@@ -364,6 +247,9 @@ impl CublasForward {
let s1_ldb = if s1_input_dim < state_dim { s1_input_dim } else { state_dim_padded };
let branch_sizes = [branch_0_size, branch_1_size, branch_2_size, branch_3_size];
let lt_raw_handle = shared.lt_handle.0;
let lt_ws_size = shared.lt_workspace_size;
let mut unique_shapes: Vec<FwdGemmKey> = Vec::new();
// h_s1: M=shared_h1, N=batch, K=s1_input_dim, ldb=s1_ldb
unique_shapes.push((shared_h1, batch_size, s1_input_dim, s1_ldb));
@@ -408,16 +294,7 @@ impl CublasForward {
}
Ok(Self {
handle: SendSyncCublasHandle(raw_handle),
_workspace_buf: workspace_buf,
workspace_ptr,
workspace_size,
lt_handle: SendSyncCublasLtHandle(lt_raw_handle),
_lt_workspace_buf: lt_workspace_buf,
lt_workspace_ptr: lt_ws_ptr,
lt_workspace_size: lt_ws_size,
add_bias_relu_f32_kernel,
add_bias_f32_kernel,
handle: shared,
// Dimensions
batch_size,
state_dim,
@@ -458,29 +335,14 @@ impl CublasForward {
w_ptr: u64, a_ptr: u64, c_ptr: u64,
m: usize, n: usize, k: usize,
) -> Result<(), MLError> {
self.lt_matmul(stream, w_ptr, a_ptr, c_ptr, m, n, k, k, self.lt_workspace_ptr, self.lt_workspace_size, "DIAG_raw")
self.lt_matmul(stream, w_ptr, a_ptr, c_ptr, m, n, k, k, self.handle.lt_workspace_ptr, self.handle.lt_workspace_size, "DIAG_raw")
}
/// CRITICAL: cublasSetStream resets workspace to the default pool.
/// We must re-set our user-owned workspace immediately after.
/// Delegates to `SharedCublasHandle::set_stream` which atomically
/// rebinds the handle and restores the user-owned workspace.
pub 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:?}")))?;
self.restore_workspace();
}
Ok(())
}
/// Restore user-owned workspace after cublasSetStream reset.
/// Must be called after every cublasSetStream to prevent CUDA Graph hang.
unsafe fn restore_workspace(&self) {
cublas_sys::cublasSetWorkspace_v2(
self.handle.0,
self.workspace_ptr as *mut std::ffi::c_void,
self.workspace_size,
);
self.handle.set_stream(stream)
}
// ══════════════════════════════════════════════════════════════════════════
@@ -518,8 +380,8 @@ impl CublasForward {
// First layer: ldb = state_dim_padded (CUTLASS K-tile alignment).
// Try fused GEMM+bias+ReLU epilogue, fall back to separate kernels.
let ws = self.lt_workspace_ptr;
let wss = self.lt_workspace_size;
let ws = self.handle.lt_workspace_ptr;
let wss = self.handle.lt_workspace_size;
if self.sgemm_f32_fused_relu_bias(stream, w_ptrs[0], states_ptr, h_s1_ptr, w_ptrs[1],
self.shared_h1, b, self.s1_input_dim, self.s1_ldb, ws, wss, "h_s1").is_err()
{
@@ -587,7 +449,7 @@ impl CublasForward {
// Per-branch workspace: eliminates contention between parallel branch streams.
// Try fused GEMM+bias+ReLU epilogue, fall back to separate kernels.
let bws = self.branch_workspace_ptrs[d];
let bwss = self.lt_workspace_size;
let bwss = self.handle.lt_workspace_size;
if self.sgemm_f32_fused_relu_bias(bs, w_ptrs[w_fc_idx], h_s2_ptr, branch_h_ptrs[d], w_ptrs[w_fc_idx + 1],
self.adv_h, b, self.shared_h2, self.shared_h2, bws, bwss, "h_bd").is_err()
{
@@ -945,7 +807,7 @@ impl CublasForward {
k: usize, // in_dim
_label: &str,
) -> Result<(), MLError> {
self.lt_matmul(stream, w_f32_ptr, a_f32_ptr, c_f32_ptr, n, b, k, k, self.lt_workspace_ptr, self.lt_workspace_size, _label)
self.lt_matmul(stream, w_f32_ptr, a_f32_ptr, c_f32_ptr, n, b, k, k, self.handle.lt_workspace_ptr, self.handle.lt_workspace_size, _label)
}
/// F32 GEMM on a branch stream with per-branch workspace.
@@ -963,7 +825,7 @@ impl CublasForward {
branch_idx: usize,
_label: &str,
) -> Result<(), MLError> {
self.lt_matmul(stream, w_f32_ptr, a_f32_ptr, c_f32_ptr, n, b, k, k, self.branch_workspace_ptrs[branch_idx], self.lt_workspace_size, _label)
self.lt_matmul(stream, w_f32_ptr, a_f32_ptr, c_f32_ptr, n, b, k, k, self.branch_workspace_ptrs[branch_idx], self.handle.lt_workspace_size, _label)
}
/// F32 GEMM with TF32 + explicit `ldb` (for padded states buffer).
@@ -980,7 +842,7 @@ impl CublasForward {
ldb: usize,
_label: &str,
) -> Result<(), MLError> {
self.lt_matmul(stream, w_f32_ptr, a_f32_ptr, c_f32_ptr, n, b, k, ldb, self.lt_workspace_ptr, self.lt_workspace_size, _label)
self.lt_matmul(stream, w_f32_ptr, a_f32_ptr, c_f32_ptr, n, b, k, ldb, self.handle.lt_workspace_ptr, self.handle.lt_workspace_size, _label)
}
/// Inner cublasLtMatmul dispatch: uses pre-cached descriptors when available,
@@ -1040,7 +902,7 @@ impl CublasForward {
let cu_stream = stream.cu_stream() as cublaslt_sys::cudaStream_t;
tracing::debug!(
handle = ?self.lt_handle.0,
handle = ?self.handle.lt_handle.0,
matmul_desc = ?cached.matmul_desc,
w_ptr = w_f32_ptr,
a_ptr = a_f32_ptr,
@@ -1057,7 +919,7 @@ impl CublasForward {
);
let matmul_status = cublaslt_sys::cublasLtMatmul(
self.lt_handle.0,
self.handle.lt_handle.0,
cached.matmul_desc,
&alpha as *const f32 as *const std::ffi::c_void,
w_f32_ptr as *const std::ffi::c_void,
@@ -1126,7 +988,7 @@ impl CublasForward {
unsafe {
let cu_stream = stream.cu_stream() as cublaslt_sys::cudaStream_t;
let status = cublaslt_sys::cublasLtMatmul(
self.lt_handle.0,
self.handle.lt_handle.0,
cached.matmul_desc,
&alpha as *const f32 as *const std::ffi::c_void,
w_ptr as *const std::ffi::c_void,
@@ -1220,7 +1082,7 @@ impl CublasForward {
).map_err(|e| MLError::ModelError(format!("set pref ws {_label}: {e:?}")))?;
let heuristic = cublaslt_result::get_matmul_algo_heuristic(
self.lt_handle.0,
self.handle.lt_handle.0,
matmul_desc,
a_layout,
b_layout,
@@ -1237,7 +1099,7 @@ impl CublasForward {
let cu_stream = stream.cu_stream() as cublaslt_sys::cudaStream_t;
tracing::debug!(
handle = ?self.lt_handle.0,
handle = ?self.handle.lt_handle.0,
matmul_desc = ?matmul_desc,
w_ptr = w_f32_ptr,
a_ptr = a_f32_ptr,
@@ -1254,7 +1116,7 @@ impl CublasForward {
);
let matmul_status = cublaslt_sys::cublasLtMatmul(
self.lt_handle.0,
self.handle.lt_handle.0,
matmul_desc,
&alpha as *const f32 as *const std::ffi::c_void,
w_f32_ptr as *const std::ffi::c_void,
@@ -1317,7 +1179,7 @@ impl CublasForward {
let blocks = ((batch * out_dim + 255) / 256) as u32;
unsafe {
stream
.launch_builder(&self.add_bias_relu_f32_kernel)
.launch_builder(&self.handle.add_bias_relu_f32_kernel)
.arg(&out_ptr)
.arg(&bias_ptr)
.arg(&out_dim_i32)
@@ -1347,7 +1209,7 @@ impl CublasForward {
unsafe {
stream
.launch_builder(&self.add_bias_f32_kernel)
.launch_builder(&self.handle.add_bias_f32_kernel)
.arg(&out_ptr)
.arg(&bias_ptr)
.arg(&out_dim_i32)
@@ -1381,30 +1243,6 @@ fn raw_f32_ptr(slice: &CudaSlice<f32>, stream: &Arc<CudaStream>) -> u64 {
ptr
}
// ── Kernel compilation ────────────────────────────────────────────────────────
/// Precompiled bias kernels cubin (build.rs — ZERO runtime nvcc).
static BIAS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/bias_kernels.cubin"));
/// Load f32 bias kernels from precompiled cubin.
fn compile_bias_kernels(
stream: &Arc<CudaStream>,
) -> Result<(CudaFunction, CudaFunction), MLError> {
let context = stream.context();
let module = context
.load_cubin(BIAS_CUBIN.to_vec())
.map_err(|e| MLError::ModelError(format!("bias cubin load: {e}")))?;
let add_bias_relu_f32 = module
.load_function("add_bias_relu_f32_kernel")
.map_err(|e| MLError::ModelError(format!("add_bias_relu_f32_kernel load: {e}")))?;
let add_bias_f32 = module
.load_function("add_bias_f32_kernel")
.map_err(|e| MLError::ModelError(format!("add_bias_f32_kernel load: {e}")))?;
Ok((add_bias_relu_f32, add_bias_f32))
}
// ── Cached GEMM descriptor factory ──────────────────────────────────────────
/// Create a `CachedGemmDesc` for a forward GEMM shape (TRANSA=T, TRANSB=N, f32, FAST_TF32).