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

fix(critical): pad backtest evaluator states buffer to state_dim_padded

Browse Source 
Root cause of cublasLtMatmul CUBLAS_STATUS_NOT_SUPPORTED:
- states_buf allocated [batch * state_dim] (unpadded)
- cuBLAS forward GEMM reads with stride state_dim_padded (pad128)
- Buffer overflow: GEMM reads past buffer end

cublasLtMatmul validates buffer sizes against layout descriptors and
returns NOT_SUPPORTED for undersized buffers. cublasSgemm silently
read garbage — this was the source of the "parallel test congestion
errors" seen previously.

Fix:
- Allocate states_buf with state_dim_padded stride (3 allocation sites)
- gather_states kernel: add padded_sd parameter, write with padded stride
- DtoD copy: use padded row bytes
- Both gather_states call sites updated with padded_sd arg

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-10 11:31:35 +02:00
parent b456691de5
commit eda50b4eb5
4 changed files with 268 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
crates/ml/src/cuda_pipeline/backtest_gather_kernel.cu
View File

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

20
crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs
View File

@@ -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::<f32>(n_windows * state_dim)
.alloc_zeros::<f32>(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<f32>.
let n_elems = self.n_windows * state_dim;
// Return a copy of the states buffer as CudaSlice<f32> (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::<f32>(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::<f32>();
let state_dim_padded = (state_dim + 127) & !127;
let state_row_bytes = n * state_dim_padded * std::mem::size_of::<f32>();
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::<f32>(cn * total_actions)
.map_err(|e| MLError::ModelError(format!("alloc chunked q_values: {e}")))?;
let ch_states_buf = self.stream.alloc_zeros::<f32>(cn * self.state_dim)
let state_dim_padded = (self.state_dim + 127) & !127;
let ch_states_buf = self.stream.alloc_zeros::<f32>(cn * state_dim_padded)
.map_err(|e| MLError::ModelError(format!("alloc chunked states_buf: {e}")))?;
let ch_actions_buf = self.stream.alloc_zeros::<i32>(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}"))

BIN
testing/cublaslt_debug Executable file
View File

Binary file not shown.

246
testing/cublaslt_debug.cu Normal file
View File

@@ -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 <cublasLt.h>
#include <cublas_v2.h>
#include <cuda_runtime.h>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <vector>
#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(&ltHandle));
// 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;
}