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

diag: standalone cublasLtMatmul Rust test for H100 debugging

Browse Source 
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-10 14:06:05 +02:00
parent 45ff940e4b
commit 171b9a241c
3 changed files with 196 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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}")))?;

192
crates/ml/src/cuda_pipeline/cublaslt_debug.rs Normal file
View File

@@ -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::<i32>(),
).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::<i32>(),
).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::<usize>(),
).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"));
}
}
}

2
crates/ml/src/cuda_pipeline/mod.rs
View File

@@ -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;