/** * 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; }