diff --git a/crates/ml/src/cuda_pipeline/backtest_gather_kernel.cu b/crates/ml/src/cuda_pipeline/backtest_gather_kernel.cu index fdc8989e7..c200bee88 100644 --- a/crates/ml/src/cuda_pipeline/backtest_gather_kernel.cu +++ b/crates/ml/src/cuda_pipeline/backtest_gather_kernel.cu @@ -34,7 +34,7 @@ extern "C" __global__ void gather_states( const __nv_bfloat16* __restrict__ features, // [n_windows * max_len * feat_dim] (bf16 from data pipeline) const float* __restrict__ portfolio, // [n_windows * 8] (f32 from env kernel) - float* states_out, // [n_windows * state_dim] (f32 for cublasSgemm) + float* states_out, // [n_windows * padded_sd] (f32 for cublasLtMatmul, padded stride) int n_windows, int max_len, int feat_dim, @@ -42,7 +42,8 @@ extern "C" __global__ void gather_states( int current_step, float initial_capital, // host scalar — convert on first use float spread_cost, // host scalar — convert on first use - int ofi_dim // 0 = no OFI, 8 = standard OFI features + int ofi_dim, // 0 = no OFI, 8 = standard OFI features + int padded_sd // pad128(state_dim) — output row stride for cuBLAS K-tile alignment ) { // Shared memory tile for coalesced portfolio reads — 8 F32 values per thread __shared__ float shmem_portfolio[256 * 8]; @@ -63,7 +64,7 @@ extern "C" __global__ void gather_states( if (w >= n_windows) return; int feat_base = (w * max_len + current_step) * feat_dim; - int out_base = w * state_dim; + int out_base = w * padded_sd; // Read portfolio state from shared memory tile — all f32 (output is f32 for cublasSgemm) float f_value = shmem_portfolio[local_tid * 8 + 0]; @@ -222,15 +223,16 @@ extern "C" __global__ void gather_states( } } - // ── 5. Zero-pad remaining [filled .. state_dim) ── // + // ── 5. Zero-pad remaining [filled .. padded_sd) ── // + // Pads both the state_dim gap AND the cuBLAS K-tile alignment columns. int filled = market_dim + 8 + 16 + ofi_dim; - for (i = filled; i + 3 < state_dim; i += 4) { + for (i = filled; i + 3 < padded_sd; i += 4) { states_out[out_base + i] = 0.0f; states_out[out_base + i + 1] = 0.0f; states_out[out_base + i + 2] = 0.0f; states_out[out_base + i + 3] = 0.0f; } - for (; i < state_dim; i++) { + for (; i < padded_sd; i++) { states_out[out_base + i] = 0.0f; } diff --git a/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs b/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs index 677651b46..6ad61a9a3 100644 --- a/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs +++ b/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs @@ -551,8 +551,9 @@ impl GpuBacktestEvaluator { // 8-aligned for H100 tensor core HMMA dispatch. const PORTFOLIO_AND_MTF_DIM: usize = 24; // 8 portfolio + 16 multi-timeframe let state_dim = (feature_dim + PORTFOLIO_AND_MTF_DIM + 7) & !7; + let state_dim_padded = (state_dim + 127) & !127; // pad128 for cuBLAS K-tile alignment let states_buf = stream - .alloc_zeros::(n_windows * state_dim) + .alloc_zeros::(n_windows * state_dim_padded) .map_err(|e| MLError::ModelError(format!("states_buf alloc: {e}")))?; // Pre-allocate forward kernel output buffer (used by evaluate_dqn path) @@ -713,7 +714,8 @@ impl GpuBacktestEvaluator { // Safety: argument order matches `gather_states` signature exactly: // features (bf16), portfolio (f32), states_out (f32), n_windows, max_len, feat_dim, - // state_dim, current_step, initial_capital, spread_cost, ofi_dim + // state_dim, current_step, initial_capital, spread_cost, ofi_dim, padded_sd + let padded_sd_i32 = ((state_dim + 127) & !127) as i32; unsafe { self.stream .launch_builder(&self.gather_kernel) @@ -728,12 +730,14 @@ impl GpuBacktestEvaluator { .arg(&initial_capital) .arg(&spread_cost) .arg(&ofi_dim_i32) + .arg(&padded_sd_i32) .launch(launch_cfg) .map_err(|e| MLError::ModelError(format!("gather_states launch step {step}: {e}")))?; } - // Return a copy of the states buffer as CudaSlice. - let n_elems = self.n_windows * state_dim; + // Return a copy of the states buffer as CudaSlice (padded stride). + let state_dim_padded = (state_dim + 127) & !127; + let n_elems = self.n_windows * state_dim_padded; let dst = self.stream.alloc_zeros::(n_elems) .map_err(|e| MLError::ModelError(format!("alloc states step {step}: {e}")))?; let src_view = self.states_buf.slice(..n_elems); @@ -1042,7 +1046,8 @@ impl GpuBacktestEvaluator { let b3 = dqn_cfg.branch_3_size as i32; let v_min_f = dqn_cfg.v_min; let v_max_f = dqn_cfg.v_max; - let state_row_bytes = n * state_dim * std::mem::size_of::(); + let state_dim_padded = (state_dim + 127) & !127; + let state_row_bytes = n * state_dim_padded * std::mem::size_of::(); let total_chunks = (self.max_len + DQN_BACKTEST_CHUNK_SIZE - 1) / DQN_BACKTEST_CHUNK_SIZE; for chunk_start in (0..self.max_len).step_by(DQN_BACKTEST_CHUNK_SIZE) { @@ -1511,7 +1516,8 @@ impl GpuBacktestEvaluator { let ch_q_values = self.stream.alloc_zeros::(cn * total_actions) .map_err(|e| MLError::ModelError(format!("alloc chunked q_values: {e}")))?; - let ch_states_buf = self.stream.alloc_zeros::(cn * self.state_dim) + let state_dim_padded = (self.state_dim + 127) & !127; + let ch_states_buf = self.stream.alloc_zeros::(cn * state_dim_padded) .map_err(|e| MLError::ModelError(format!("alloc chunked states_buf: {e}")))?; let ch_actions_buf = self.stream.alloc_zeros::(cn) .map_err(|e| MLError::ModelError(format!("alloc chunked actions_buf: {e}")))?; @@ -1711,6 +1717,7 @@ impl GpuBacktestEvaluator { let initial_capital = self.config.initial_capital; let spread_cost = self.config.spread_cost; let ofi_dim_i32 = self.config.ofi_dim as i32; + let padded_sd_i32 = ((state_dim + 127) & !127) as i32; unsafe { self.stream @@ -1726,6 +1733,7 @@ impl GpuBacktestEvaluator { .arg(&initial_capital) .arg(&spread_cost) .arg(&ofi_dim_i32) + .arg(&padded_sd_i32) .launch(launch_cfg) .map_err(|e| { MLError::ModelError(format!("gather_states launch step {step}: {e}")) diff --git a/testing/cublaslt_debug b/testing/cublaslt_debug new file mode 100755 index 000000000..7128bac0a Binary files /dev/null and b/testing/cublaslt_debug differ diff --git a/testing/cublaslt_debug.cu b/testing/cublaslt_debug.cu new file mode 100644 index 000000000..2890b46cd --- /dev/null +++ b/testing/cublaslt_debug.cu @@ -0,0 +1,246 @@ +/** + * Minimal cublasLtMatmul debug test. + * + * Tests the EXACT same configuration that fails in the Rust code: + * C[N, B] = op(W)[N, K] @ A[K, B] where op(W) = W^T + * + * Failing dimensions on H100: m=128, n=4096, k=256 + * Failing dimensions on 3050: m=128, n=512, k=64 + * + * Compile: + * nvcc -o cublaslt_debug cublaslt_debug.cu -lcublasLt -lcublas -lcudart + * + * Run: + * ./cublaslt_debug + */ + +#include +#include +#include +#include +#include +#include +#include + +#define CHECK_CUDA(call) \ + do { \ + cudaError_t err = (call); \ + if (err != cudaSuccess) { \ + fprintf(stderr, "CUDA error at %s:%d: %s\n", __FILE__, __LINE__, \ + cudaGetErrorString(err)); \ + exit(1); \ + } \ + } while (0) + +#define CHECK_CUBLAS(call) \ + do { \ + cublasStatus_t st = (call); \ + if (st != CUBLAS_STATUS_SUCCESS) { \ + fprintf(stderr, "cuBLAS error at %s:%d: status=%d\n", \ + __FILE__, __LINE__, (int)st); \ + exit(1); \ + } \ + } while (0) + +// ── Test configuration ───────────────────────────────────────────── +struct TestCase { + int m, n, k; + const char* label; +}; + +void test_cublaslt_matmul(cublasLtHandle_t ltHandle, + cudaStream_t stream, + void* workspace, size_t wsSize, + const TestCase& tc, + cublasComputeType_t computeType, + const char* computeLabel) { + int m = tc.m, n = tc.n, k = tc.k; + printf(" [%s] (%d, %d, %d) compute=%s ... ", tc.label, m, n, k, computeLabel); + fflush(stdout); + + // Allocate device memory + float *d_W, *d_A, *d_C; + CHECK_CUDA(cudaMalloc(&d_W, (size_t)k * m * sizeof(float))); // W[m, k] row-major = [k, m] col-major + CHECK_CUDA(cudaMalloc(&d_A, (size_t)k * n * sizeof(float))); // A[n, k] row-major = [k, n] col-major + CHECK_CUDA(cudaMalloc(&d_C, (size_t)m * n * sizeof(float))); // C[n, m] row-major = [m, n] col-major + CHECK_CUDA(cudaMemset(d_W, 0, (size_t)k * m * sizeof(float))); + CHECK_CUDA(cudaMemset(d_A, 0, (size_t)k * n * sizeof(float))); + CHECK_CUDA(cudaMemset(d_C, 0, (size_t)m * n * sizeof(float))); + + float alpha = 1.0f, beta = 0.0f; + + // ── Approach 1: TRANSA=T (matching our Rust code) ────────────── + // GEMM: C[m, n] = W^T[m, k] @ A[k, n] + // Physical: W is [k, m] col-major (row-major W[m, k]), TRANSA=T → [m, k] + // A is [k, n] col-major (row-major A[n, k]), TRANSB=N → [k, n] + // C is [m, n] col-major + { + cublasLtMatmulDesc_t matmulDesc; + CHECK_CUBLAS(cublasLtMatmulDescCreate(&matmulDesc, computeType, CUDA_R_32F)); + + cublasOperation_t opT = CUBLAS_OP_T; + cublasOperation_t opN = CUBLAS_OP_N; + CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(matmulDesc, + CUBLASLT_MATMUL_DESC_TRANSA, &opT, sizeof(opT))); + CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(matmulDesc, + CUBLASLT_MATMUL_DESC_TRANSB, &opN, sizeof(opN))); + + // A = W: physical [k, m], ld=k. TRANSA=T → logical [m, k]. + cublasLtMatrixLayout_t Adesc, Bdesc, Cdesc, Ddesc; + CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Adesc, CUDA_R_32F, k, m, k)); + CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Bdesc, CUDA_R_32F, k, n, k)); + CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Cdesc, CUDA_R_32F, m, n, m)); + CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Ddesc, CUDA_R_32F, m, n, m)); + + cublasLtMatmulPreference_t pref; + CHECK_CUBLAS(cublasLtMatmulPreferenceCreate(&pref)); + CHECK_CUBLAS(cublasLtMatmulPreferenceSetAttribute(pref, + CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &wsSize, sizeof(wsSize))); + + cublasLtMatmulHeuristicResult_t heurResult; + int returnedAlgoCount = 0; + cublasStatus_t heurStatus = cublasLtMatmulAlgoGetHeuristic( + ltHandle, matmulDesc, Adesc, Bdesc, Cdesc, Ddesc, + pref, 1, &heurResult, &returnedAlgoCount); + + if (heurStatus != CUBLAS_STATUS_SUCCESS || returnedAlgoCount == 0) { + printf("HEURISTIC FAILED (status=%d, count=%d)\n", (int)heurStatus, returnedAlgoCount); + } else { + printf("heuristic OK (ws=%zu) ", heurResult.workspaceSize); + fflush(stdout); + + cublasStatus_t matmulStatus = cublasLtMatmul( + ltHandle, matmulDesc, + &alpha, + d_W, Adesc, + d_A, Bdesc, + &beta, + d_C, Cdesc, + d_C, Ddesc, + &heurResult.algo, + workspace, wsSize, + stream); + + CHECK_CUDA(cudaStreamSynchronize(stream)); + + if (matmulStatus == CUBLAS_STATUS_SUCCESS) { + printf("MATMUL OK ✓\n"); + } else { + printf("MATMUL FAILED (status=%d) ✗\n", (int)matmulStatus); + } + } + + cublasLtMatmulPreferenceDestroy(pref); + cublasLtMatrixLayoutDestroy(Ddesc); + cublasLtMatrixLayoutDestroy(Cdesc); + cublasLtMatrixLayoutDestroy(Bdesc); + cublasLtMatrixLayoutDestroy(Adesc); + cublasLtMatmulDescDestroy(matmulDesc); + } + + // ── Approach 2: NO TRANSPOSE (swap A/B, row-major trick) ─────── + // GEMM: C'[m, n] = W'[m, k] @ A'[k, n] TRANSA=N, TRANSB=N + // But W stored row-major [m, k] → col-major [k, m]. So W' has rows=m, cols=k, ld=k. + // Wait — col-major [rows, cols] with ld: element (i,j) = data[j*ld + i]. + // Row-major W[m, k]: element (i,j) = data[i*k + j]. + // Col-major interpretation: rows=k, cols=m, ld=k. (j*k + i) maps to data[col*k + row]. + // This is [k, m] col-major. To get [m, k] col-major without transpose, + // we'd need ld=m, but ld must be >= rows=m. That works for an [m, k] layout with ld=m. + // But that's NOT how the data is stored. Row-major [m, k] IS col-major [k, m]. + // + // Correct no-transpose approach: swap the matrices. + // Row-major: C[n, m] = A[n, k] @ W[m, k]^T + // → col-major: C'[m, n] = W'[k, m]^T_col @ A'[k, n] + // But W'[k, m] transposed is [m, k]. That needs TRANSA=T again. + // + // Alternative: D = alpha * A * B + beta * C → m,n,k from D perspective. + // Row-major: out[batch, out_dim] = in[batch, in_dim] @ W[out_dim, in_dim]^T + // cuBLAS col-major: out'[out_dim, batch] = W'[in_dim, out_dim]^T[out_dim, in_dim] @ in'[in_dim, batch] + // So: M=out_dim, N=batch, K=in_dim, A=W with TRANSA=T, B=in with TRANSB=N. + // This is approach 1. There's no clean no-transpose version. + { + // Skip approach 2 — approach 1 is the correct formulation. + } + + // ── Approach 3: cublasSgemm for comparison ───────────────────── + { + cublasHandle_t handle; + CHECK_CUBLAS(cublasCreate(&handle)); + CHECK_CUBLAS(cublasSetStream(handle, stream)); + + cublasStatus_t sgemmStatus = cublasSgemm(handle, + CUBLAS_OP_T, CUBLAS_OP_N, + m, n, k, + &alpha, + d_W, k, // W[k, m] col-major, transpose → [m, k] + d_A, k, // A[k, n] col-major + &beta, + d_C, m); // C[m, n] col-major + + CHECK_CUDA(cudaStreamSynchronize(stream)); + printf(" [%s] (%d, %d, %d) cublasSgemm ... %s\n", tc.label, m, n, k, + sgemmStatus == CUBLAS_STATUS_SUCCESS ? "OK ✓" : "FAILED ✗"); + + cublasDestroy(handle); + } + + CHECK_CUDA(cudaFree(d_W)); + CHECK_CUDA(cudaFree(d_A)); + CHECK_CUDA(cudaFree(d_C)); +} + +int main() { + // Print GPU info + cudaDeviceProp prop; + CHECK_CUDA(cudaGetDeviceProperties(&prop, 0)); + printf("GPU: %s (SM %d.%d)\n", prop.name, prop.major, prop.minor); + printf("CUDA driver: "); + int driverVersion; + CHECK_CUDA(cudaDriverGetVersion(&driverVersion)); + printf("%d.%d\n", driverVersion / 1000, (driverVersion % 1000) / 10); + + // Create cublasLt handle + cublasLtHandle_t ltHandle; + CHECK_CUBLAS(cublasLtCreate(<Handle)); + + // Allocate workspace (32 MB) + size_t wsSize = 32 * 1024 * 1024; + void* workspace; + CHECK_CUDA(cudaMalloc(&workspace, wsSize)); + + // Create stream + cudaStream_t stream; + CHECK_CUDA(cudaStreamCreate(&stream)); + + // Test cases: the dimensions that fail in Rust + TestCase cases[] = { + {128, 64, 64, "training_small"}, + {128, 512, 64, "eval_chunk"}, + {128, 4096, 256, "h100_batch"}, + {256, 64, 256, "shared_h2"}, + {64, 64, 256, "shared_h1"}, + {51, 64, 128, "v_logits"}, + {3, 5, 4, "cudarc_test"}, + }; + int numCases = sizeof(cases) / sizeof(cases[0]); + + printf("\n=== CUBLAS_COMPUTE_32F_FAST_TF32 ===\n"); + for (int i = 0; i < numCases; i++) { + test_cublaslt_matmul(ltHandle, stream, workspace, wsSize, cases[i], + CUBLAS_COMPUTE_32F_FAST_TF32, "FAST_TF32"); + } + + printf("\n=== CUBLAS_COMPUTE_32F ===\n"); + for (int i = 0; i < numCases; i++) { + test_cublaslt_matmul(ltHandle, stream, workspace, wsSize, cases[i], + CUBLAS_COMPUTE_32F, "32F"); + } + + // Cleanup + CHECK_CUDA(cudaStreamDestroy(stream)); + CHECK_CUDA(cudaFree(workspace)); + CHECK_CUBLAS(cublasLtDestroy(ltHandle)); + + printf("\nDone.\n"); + return 0; +}