From 171b9a241c004138ef7efcab8dd4e8ceeedee1b9 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 10 Apr 2026 14:06:05 +0200 Subject: [PATCH] diag: standalone cublasLtMatmul Rust test for H100 debugging Co-Authored-By: Claude Opus 4.6 (1M context) --- .../ml/src/cuda_pipeline/batched_forward.rs | 5 +- crates/ml/src/cuda_pipeline/cublaslt_debug.rs | 192 ++++++++++++++++++ crates/ml/src/cuda_pipeline/mod.rs | 2 + 3 files changed, 196 insertions(+), 3 deletions(-) create mode 100644 crates/ml/src/cuda_pipeline/cublaslt_debug.rs diff --git a/crates/ml/src/cuda_pipeline/batched_forward.rs b/crates/ml/src/cuda_pipeline/batched_forward.rs index 4b5f9debc..247eb4e6d 100644 --- a/crates/ml/src/cuda_pipeline/batched_forward.rs +++ b/crates/ml/src/cuda_pipeline/batched_forward.rs @@ -501,9 +501,8 @@ impl CublasForward { let distinct_branches = h_b0_ptr != h_b1_ptr && h_b1_ptr != h_b2_ptr && h_b0_ptr != h_b2_ptr && h_b0_ptr != h_b3_ptr && h_b1_ptr != h_b3_ptr && h_b2_ptr != h_b3_ptr; - if false && distinct_branches { - // Multi-stream branch dispatch DISABLED for H100 cublasLtMatmul debugging. - // cublasLtMatmul returns NOT_SUPPORTED on branch streams on H100. + if distinct_branches { + // Multi-stream branch dispatch — cublasLtMatmul takes stream per-call. self.trunk_done_event.record(stream) .map_err(|e| MLError::ModelError(format!("f32 trunk event record: {e}")))?; diff --git a/crates/ml/src/cuda_pipeline/cublaslt_debug.rs b/crates/ml/src/cuda_pipeline/cublaslt_debug.rs new file mode 100644 index 000000000..404f876c0 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/cublaslt_debug.rs @@ -0,0 +1,192 @@ +//! Standalone cublasLtMatmul debug test. +//! Tests the EXACT same API call that fails in the experience collector on H100. +//! Run: SQLX_OFFLINE=true cargo test -p ml --lib -- cuda_pipeline::cublaslt_debug --nocapture + +#[cfg(test)] +mod tests { + use cudarc::cublaslt::sys as cublaslt_sys; + use cudarc::cublaslt::result as cublaslt_result; + use cudarc::driver::CudaContext; + + struct TestCase { + m: i32, + n: i32, + k: i32, + label: &'static str, + } + + fn test_cublaslt_matmul( + lt_handle: cublaslt_sys::cublasLtHandle_t, + stream: cudarc::driver::sys::CUstream, + ws_ptr: u64, + ws_size: usize, + tc: &TestCase, + ) -> Result<(), String> { + let m = tc.m as usize; + let n = tc.n as usize; + let k = tc.k as usize; + + // Allocate buffers via cuMemAlloc (NOT stream-ordered) + let mut d_w: cudarc::driver::sys::CUdeviceptr = 0; + let mut d_a: cudarc::driver::sys::CUdeviceptr = 0; + let mut d_c: cudarc::driver::sys::CUdeviceptr = 0; + unsafe { + cudarc::driver::sys::cuMemAlloc_v2(&mut d_w, k * m * 4); + cudarc::driver::sys::cuMemAlloc_v2(&mut d_a, k * n * 4); + cudarc::driver::sys::cuMemAlloc_v2(&mut d_c, m * n * 4); + cudarc::driver::sys::cuMemsetD8_v2(d_w, 0, k * m * 4); + cudarc::driver::sys::cuMemsetD8_v2(d_a, 0, k * n * 4); + cudarc::driver::sys::cuMemsetD8_v2(d_c, 0, m * n * 4); + } + + let alpha = 1.0_f32; + let beta = 0.0_f32; + let f32_type = cublaslt_sys::cudaDataType_t::CUDA_R_32F; + let compute_type = cublaslt_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F; + + let result = unsafe { + // Create matmul descriptor + let matmul_desc = cublaslt_result::create_matmul_desc(compute_type, f32_type) + .map_err(|e| format!("create_matmul_desc: {e:?}"))?; + + // TRANSA=T, TRANSB=N + let transa: i32 = 1; + let transb: i32 = 0; + cublaslt_result::set_matmul_desc_attribute( + matmul_desc, + cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSA, + &transa as *const i32 as *const std::ffi::c_void, + std::mem::size_of::(), + ).map_err(|e| format!("set TRANSA: {e:?}"))?; + cublaslt_result::set_matmul_desc_attribute( + matmul_desc, + cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSB, + &transb as *const i32 as *const std::ffi::c_void, + std::mem::size_of::(), + ).map_err(|e| format!("set TRANSB: {e:?}"))?; + + // Matrix layouts + let a_layout = cublaslt_result::create_matrix_layout(f32_type, k as u64, m as u64, k as i64) + .map_err(|e| format!("A layout: {e:?}"))?; + let b_layout = cublaslt_result::create_matrix_layout(f32_type, k as u64, n as u64, k as i64) + .map_err(|e| format!("B layout: {e:?}"))?; + let c_layout = cublaslt_result::create_matrix_layout(f32_type, m as u64, n as u64, m as i64) + .map_err(|e| format!("C layout: {e:?}"))?; + let d_layout = cublaslt_result::create_matrix_layout(f32_type, m as u64, n as u64, m as i64) + .map_err(|e| format!("D layout: {e:?}"))?; + + // Heuristic + let matmul_pref = cublaslt_result::create_matmul_pref() + .map_err(|e| format!("create_matmul_pref: {e:?}"))?; + cublaslt_result::set_matmul_pref_attribute( + matmul_pref, + cublaslt_sys::cublasLtMatmulPreferenceAttributes_t::CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, + &ws_size as *const usize as *const std::ffi::c_void, + std::mem::size_of::(), + ).map_err(|e| format!("set pref ws: {e:?}"))?; + + let heuristic = cublaslt_result::get_matmul_algo_heuristic( + lt_handle, matmul_desc, a_layout, b_layout, c_layout, d_layout, matmul_pref, + ).map_err(|e| format!("heuristic: {e:?}"))?; + + eprintln!(" heuristic OK (ws_needed={})", heuristic.workspaceSize); + + // Execute matmul + let cu_stream = stream as cublaslt_sys::cudaStream_t; + let matmul_status = cublaslt_sys::cublasLtMatmul( + lt_handle, + matmul_desc, + &alpha as *const f32 as *const std::ffi::c_void, + d_w as *const std::ffi::c_void, + a_layout, + d_a as *const std::ffi::c_void, + b_layout, + &beta as *const f32 as *const std::ffi::c_void, + d_c as *const std::ffi::c_void, + c_layout, + d_c as *mut std::ffi::c_void, + d_layout, + &heuristic.algo as *const cublaslt_sys::cublasLtMatmulAlgo_t, + ws_ptr as *mut std::ffi::c_void, + ws_size, + cu_stream, + ); + + // Sync + cudarc::driver::sys::cuStreamSynchronize(stream); + + // Cleanup + let _ = cublaslt_result::destroy_matmul_pref(matmul_pref); + let _ = cublaslt_result::destroy_matrix_layout(d_layout); + let _ = cublaslt_result::destroy_matrix_layout(c_layout); + let _ = cublaslt_result::destroy_matrix_layout(b_layout); + let _ = cublaslt_result::destroy_matrix_layout(a_layout); + let _ = cublaslt_result::destroy_matmul_desc(matmul_desc); + + if matmul_status == cublaslt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS { + Ok(()) + } else { + Err(format!("cublasLtMatmul status={matmul_status:?}")) + } + }; + + // Free buffers + unsafe { + cudarc::driver::sys::cuMemFree_v2(d_w); + cudarc::driver::sys::cuMemFree_v2(d_a); + cudarc::driver::sys::cuMemFree_v2(d_c); + } + + result + } + + #[test] + fn test_cublaslt_all_dimensions() { + let ctx = CudaContext::new(0).unwrap(); + let stream = ctx.default_stream(); + + // Create cublasLt handle + let lt_handle = cublaslt_result::create_handle().unwrap(); + + // Allocate workspace via cuMemAlloc + let ws_size: usize = 32 * 1024 * 1024; + let mut ws_ptr: cudarc::driver::sys::CUdeviceptr = 0; + unsafe { + cudarc::driver::sys::cuMemAlloc_v2(&mut ws_ptr, ws_size); + } + + let cases = [ + TestCase { m: 128, n: 64, k: 64, label: "train_small" }, + TestCase { m: 128, n: 512, k: 64, label: "eval_chunk" }, + TestCase { m: 128, n: 4096, k: 256, label: "h100_batch" }, + TestCase { m: 256, n: 4096, k: 256, label: "h100_h_s2" }, + TestCase { m: 128, n: 16384, k: 256, label: "h100_full" }, + TestCase { m: 51, n: 4096, k: 128, label: "h100_v_logits" }, + TestCase { m: 3, n: 5, k: 4, label: "tiny" }, + ]; + + let cu_stream = stream.cu_stream(); + let mut failures = Vec::new(); + + for tc in &cases { + eprint!(" [{}] ({},{},{}) ... ", tc.label, tc.m, tc.n, tc.k); + match test_cublaslt_matmul(lt_handle, cu_stream, ws_ptr, ws_size, tc) { + Ok(()) => eprintln!("OK ✓"), + Err(e) => { + eprintln!("FAIL ✗: {e}"); + failures.push(format!("[{}] ({},{},{}): {e}", tc.label, tc.m, tc.n, tc.k)); + } + } + } + + // Cleanup + unsafe { + cudarc::driver::sys::cuMemFree_v2(ws_ptr); + let _ = cublaslt_result::destroy_handle(lt_handle); + } + + if !failures.is_empty() { + panic!("cublasLtMatmul failures:\n{}", failures.join("\n")); + } + } +} diff --git a/crates/ml/src/cuda_pipeline/mod.rs b/crates/ml/src/cuda_pipeline/mod.rs index 3f38271ca..6f90d9358 100644 --- a/crates/ml/src/cuda_pipeline/mod.rs +++ b/crates/ml/src/cuda_pipeline/mod.rs @@ -30,6 +30,8 @@ pub mod gpu_dqn_trainer; pub mod batched_forward; pub mod batched_backward; pub mod gpu_her; +#[cfg(test)] +mod cublaslt_debug; pub mod gpu_iql_trainer; pub mod gpu_iqn_head; pub mod gpu_attention;