eda50b4eb5
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>
247 lines
10 KiB
Plaintext
247 lines
10 KiB
Plaintext
/**
|
|
* 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(<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;
|
|
}
|