From 5df50190fad0fb57174acd30c89dc090c037f315 Mon Sep 17 00:00:00 2001 From: Arthur Hanson Rasmusson Date: Sun, 18 Oct 2026 12:05:31 -0400 Subject: cuda-ab: transformer + transformer_mt clients (single-head attention + 4-layer FFN, fp32, CPU-reference checked; multi-stream and 5-launch loop variants) --- tools/cuda-ab/transformer_mt.cu | 468 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 468 insertions(+) create mode 100644 tools/cuda-ab/transformer_mt.cu (limited to 'tools/cuda-ab/transformer_mt.cu') diff --git a/tools/cuda-ab/transformer_mt.cu b/tools/cuda-ab/transformer_mt.cu new file mode 100644 index 0000000..ffd23aa --- /dev/null +++ b/tools/cuda-ab/transformer_mt.cu @@ -0,0 +1,468 @@ +// SPDX-License-Identifier: BSD-3-Clause +/* + * transformer_mt.cu - attention forward on 3 independent streams + a + * 5-launch sequential loop, in one binary. + * + * Same single-head attention op as transformer.cu (X[64x256], + * Wq/Wk/Wv/Wo[256x256], head dim 256, fp32, softmax(Q*K^T/16), Y=(A*V)*Wo), + * exercised two ways: + * + * 1. Multi-stream class: three independent streams each run the full op on + * their own input (X_s[i] = xval(i + s)). Every stream records an event + * after its chain; stream 2 waits on streams 0 and 1 via cudaStreamWaitEvent + * and the whole span is timed with cudaEventElapsedTime. + * 2. Repeated-launch class: the op is launched 5 times in a loop on one + * stream with per-iteration inputs and per-iteration outputs, one + * synchronization at the end. + * + * Both classes are verified against an internal double-precision CPU + * reference (tolerance 1e-3 relative, 1e-6 absolute floor). + * + * Output lines are stable and greppable: + * + * DEVICE cc=. sm= mem= + * TRANSFORMER_MT streams=3 bad= ysum= + * TRANSFORMER_LOOP launches=5 bad= ysum= + * TIMING ms= + * STATUS OK|FAIL + * + * Build with a static CUDA runtime (--cudart=static) so the guest only needs + * the driver's libcuda.so, not a CUDA toolkit installation: + * + * nvcc -O2 --cudart=static -arch=sm_75 -o transformer_mt transformer_mt.cu + * + * Then, from the guest: /tmp/transformer_mt + */ + +#include +#include +#include +#include +#include + +#define CHECK(x) \ + do { \ + cudaError_t e_ = (x); \ + if (e_ != cudaSuccess) { \ + printf("CUDA_ERROR %s at line %d\n", \ + cudaGetErrorString(e_), __LINE__); \ + printf("STATUS FAIL\n"); \ + return 1; \ + } \ + } while (0) + +#define SEQ 64 +#define DH 256 +#define TILE 16 +#define TOL 1e-3f +#define NST 3 +#define NLOOP 5 + +__global__ void gemm_nn(const float *A, const float *B, float *C, + int M, int N, int K) +{ + __shared__ float As[TILE][TILE]; + __shared__ float Bs[TILE][TILE]; + int row = blockIdx.y * TILE + threadIdx.y; + int col = blockIdx.x * TILE + threadIdx.x; + float acc = 0.0f; + + for (int t = 0; t < K / TILE; t++) { + As[threadIdx.y][threadIdx.x] = A[row * K + t * TILE + threadIdx.x]; + Bs[threadIdx.y][threadIdx.x] = + B[(t * TILE + threadIdx.y) * N + col]; + __syncthreads(); + for (int k = 0; k < TILE; k++) + acc += As[threadIdx.y][k] * Bs[k][threadIdx.x]; + __syncthreads(); + } + C[row * N + col] = acc; +} + +__global__ void gemm_nt(const float *A, const float *B, float *C, + int M, int N, int K) +{ + /* C[i][j] = sum_k A[i][k] * B[j][k] (B is [N][K], row-major). */ + __shared__ float As[TILE][TILE]; + __shared__ float Bs[TILE][TILE]; + int row = blockIdx.y * TILE + threadIdx.y; + int tcol = blockIdx.x * TILE + threadIdx.x; + float acc = 0.0f; + + for (int t = 0; t < K / TILE; t++) { + As[threadIdx.y][threadIdx.x] = A[row * K + t * TILE + threadIdx.x]; + Bs[threadIdx.y][threadIdx.x] = + B[(blockIdx.x * TILE + threadIdx.y) * K + + t * TILE + threadIdx.x]; + __syncthreads(); + for (int k = 0; k < TILE; k++) + acc += As[threadIdx.y][k] * Bs[threadIdx.x][k]; + __syncthreads(); + } + C[row * N + tcol] = acc; +} + +__global__ void softmax_rows(const float *S, float *P, int rows, int cols, + float scale) +{ + extern __shared__ float sm[]; + int row = blockIdx.x; + int t = threadIdx.x; + float v = S[row * cols + t] * scale; + int stride; + + sm[t] = v; + __syncthreads(); + for (stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (t < stride) + sm[t] = fmaxf(sm[t], sm[t + stride]); + __syncthreads(); + } + v = expf(v - sm[0]); + sm[t] = v; + __syncthreads(); + for (stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (t < stride) + sm[t] += sm[t + stride]; + __syncthreads(); + } + P[row * cols + t] = v / sm[0]; +} + +/* ------------------------------------------------------------------ host */ + +struct attn_bufs { + float *q, *k, *v, *s, *a, *av, *y; +}; + +static void launch_chain(cudaStream_t st, const float *x, struct attn_bufs b, + const float *wq, const float *wk, const float *wv, + const float *wo) +{ + dim3 g1(DH / TILE, SEQ / TILE); + dim3 g2(SEQ / TILE, SEQ / TILE); + dim3 tb(TILE, TILE); + + gemm_nn<<>>(x, wq, b.q, SEQ, DH, DH); + gemm_nn<<>>(x, wk, b.k, SEQ, DH, DH); + gemm_nn<<>>(x, wv, b.v, SEQ, DH, DH); + gemm_nt<<>>(b.q, b.k, b.s, SEQ, SEQ, DH); + softmax_rows<<>>(b.s, b.a, + SEQ, SEQ, + 1.0f / 16.0f); + gemm_nn<<>>(b.a, b.v, b.av, SEQ, DH, SEQ); + gemm_nn<<>>(b.av, wo, b.y, SEQ, DH, DH); +} + +static float xval(int i) +{ + return (float)(((i % 7) - 3)) * 0.25f; +} + +static float wval(int i, int seed) +{ + return (float)(((i * 13 + seed * 101 + 7) % 5) - 2) * 0.125f; +} + +static void ref_gemm(const double *A, const double *B, double *C, + int M, int N, int K) +{ + for (int i = 0; i < M; i++) + for (int j = 0; j < N; j++) { + double acc = 0.0; + + for (int k = 0; k < K; k++) + acc += A[i * K + k] * B[k * N + j]; + C[i * N + j] = acc; + } +} + +static void ref_gemm_nt(const double *A, const double *B, double *C, + int M, int N, int K) +{ + for (int i = 0; i < M; i++) + for (int j = 0; j < N; j++) { + double acc = 0.0; + + for (int k = 0; k < K; k++) + acc += A[i * K + k] * B[j * K + k]; + C[i * N + j] = acc; + } +} + +static void ref_softmax(const double *S, double *P, int rows, int cols, + double scale) +{ + for (int i = 0; i < rows; i++) { + const double *s = S + i * cols; + double *p = P + i * cols; + double m = s[0] * scale, sum = 0.0; + + for (int j = 1; j < cols; j++) + if (s[j] * scale > m) + m = s[j] * scale; + for (int j = 0; j < cols; j++) { + p[j] = exp(s[j] * scale - m); + sum += p[j]; + } + for (int j = 0; j < cols; j++) + p[j] /= sum; + } +} + +static int count_bad(const float *gpu, const double *ref, int n) +{ + int bad = 0; + + for (int i = 0; i < n; i++) { + double g = (double)gpu[i]; + double r = ref[i]; + + if (fabs(g - r) > (double)TOL * (fabs(r) + 1e-6)) + bad++; + } + return bad; +} + +static double dsum(const float *h, int n) +{ + double s = 0.0; + + for (int i = 0; i < n; i++) + s += (double)h[i]; + return s; +} + +/* Verify one attention result against the CPU reference for input offset o. */ +static int verify_one(const float *h_y, int o, double *ysum_out) +{ + const int seqdh = SEQ * DH, seq2 = SEQ * SEQ; + double *x, *q, *k, *v, *s, *a, *av, *y, *w; + int i, bad; + + x = (double *)malloc((size_t)seqdh * sizeof(double)); + q = (double *)malloc((size_t)seqdh * sizeof(double)); + k = (double *)malloc((size_t)seqdh * sizeof(double)); + v = (double *)malloc((size_t)seqdh * sizeof(double)); + s = (double *)malloc((size_t)seq2 * sizeof(double)); + a = (double *)malloc((size_t)seq2 * sizeof(double)); + av = (double *)malloc((size_t)seqdh * sizeof(double)); + y = (double *)malloc((size_t)seqdh * sizeof(double)); + w = (double *)malloc((size_t)DH * DH * sizeof(double)); + if (!x || !q || !k || !v || !s || !a || !av || !y || !w) + exit(1); + + for (i = 0; i < seqdh; i++) + x[i] = (double)xval(i + o); + for (i = 0; i < DH * DH; i++) + w[i] = (double)wval(i, 1); + ref_gemm(x, w, q, SEQ, DH, DH); + for (i = 0; i < DH * DH; i++) + w[i] = (double)wval(i, 2); + ref_gemm(x, w, k, SEQ, DH, DH); + for (i = 0; i < DH * DH; i++) + w[i] = (double)wval(i, 3); + ref_gemm(x, w, v, SEQ, DH, DH); + ref_gemm_nt(q, k, s, SEQ, SEQ, DH); + ref_softmax(s, a, SEQ, SEQ, 1.0 / 16.0); + ref_gemm(a, v, av, SEQ, DH, SEQ); + for (i = 0; i < DH * DH; i++) + w[i] = (double)wval(i, 4); + ref_gemm(av, w, y, SEQ, DH, DH); + + bad = count_bad(h_y, y, seqdh); + *ysum_out = dsum(h_y, seqdh); + free(x); + free(q); + free(k); + free(v); + free(s); + free(a); + free(av); + free(y); + free(w); + return bad; +} + +int main(void) +{ + cudaDeviceProp prop; + const int seqdh = SEQ * DH, dh2 = DH * DH; + const size_t b_seqdh = (size_t)seqdh * sizeof(float); + const size_t b_dh2 = (size_t)dh2 * sizeof(float); + const size_t b_seq2 = (size_t)SEQ * SEQ * sizeof(float); + cudaStream_t st[NST], st_loop; + cudaEvent_t ev[NST], t0, t1; + struct attn_bufs ab[NST], lb; + float *d_x[NST], *h_y[NST]; + float *d_xl[NLOOP], *d_yl[NLOOP], *h_yl[NLOOP]; + float *d_wq, *d_wk, *d_wv, *d_wo, *wl; + float ms = 0.0f; + int bad_mt = 0, bad_loop = 0, s, i; + + setvbuf(stdout, NULL, _IONBF, 0); + + CHECK(cudaGetDeviceProperties(&prop, 0)); + printf("DEVICE %s cc=%d.%d sm=%d mem=%llu\n", prop.name, + prop.major, prop.minor, prop.multiProcessorCount, + (unsigned long long)prop.totalGlobalMem); + + CHECK(cudaMalloc(&d_wq, b_dh2)); + CHECK(cudaMalloc(&d_wk, b_dh2)); + CHECK(cudaMalloc(&d_wv, b_dh2)); + CHECK(cudaMalloc(&d_wo, b_dh2)); + wl = (float *)malloc(b_dh2); + if (!wl) + return 1; + for (i = 0; i < dh2; i++) + wl[i] = wval(i, 1); + CHECK(cudaMemcpy(d_wq, wl, b_dh2, cudaMemcpyHostToDevice)); + for (i = 0; i < dh2; i++) + wl[i] = wval(i, 2); + CHECK(cudaMemcpy(d_wk, wl, b_dh2, cudaMemcpyHostToDevice)); + for (i = 0; i < dh2; i++) + wl[i] = wval(i, 3); + CHECK(cudaMemcpy(d_wv, wl, b_dh2, cudaMemcpyHostToDevice)); + for (i = 0; i < dh2; i++) + wl[i] = wval(i, 4); + CHECK(cudaMemcpy(d_wo, wl, b_dh2, cudaMemcpyHostToDevice)); + + /* ------------- class 1: 3 independent streams ------------- */ + for (s = 0; s < NST; s++) { + float *hx = (float *)malloc(b_seqdh); + + if (!hx) + return 1; + for (i = 0; i < seqdh; i++) + hx[i] = xval(i + s); + h_y[s] = (float *)malloc(b_seqdh); + if (!h_y[s]) + return 1; + + CHECK(cudaMalloc(&d_x[s], b_seqdh)); + CHECK(cudaMalloc(&ab[s].q, b_seqdh)); + CHECK(cudaMalloc(&ab[s].k, b_seqdh)); + CHECK(cudaMalloc(&ab[s].v, b_seqdh)); + CHECK(cudaMalloc(&ab[s].s, b_seq2)); + CHECK(cudaMalloc(&ab[s].a, b_seq2)); + CHECK(cudaMalloc(&ab[s].av, b_seqdh)); + CHECK(cudaMalloc(&ab[s].y, b_seqdh)); + CHECK(cudaStreamCreate(&st[s])); + CHECK(cudaEventCreate(&ev[s])); + CHECK(cudaMemcpyAsync(d_x[s], hx, b_seqdh, + cudaMemcpyHostToDevice, st[s])); + free(hx); + } + CHECK(cudaEventCreate(&t0)); + CHECK(cudaEventCreate(&t1)); + CHECK(cudaEventRecord(t0)); + + for (s = 0; s < NST; s++) { + launch_chain(st[s], d_x[s], ab[s], d_wq, d_wk, d_wv, d_wo); + CHECK(cudaGetLastError()); + CHECK(cudaEventRecord(ev[s], st[s])); + CHECK(cudaMemcpyAsync(h_y[s], ab[s].y, b_seqdh, + cudaMemcpyDeviceToHost, st[s])); + } + /* event-ordered join: stream 2 waits on streams 0 and 1 */ + CHECK(cudaStreamWaitEvent(st[2], ev[0], 0)); + CHECK(cudaStreamWaitEvent(st[2], ev[1], 0)); + CHECK(cudaEventRecord(t1, st[2])); + for (s = 0; s < NST; s++) + CHECK(cudaStreamSynchronize(st[s])); + CHECK(cudaEventElapsedTime(&ms, t0, t1)); + + for (s = 0; s < NST; s++) { + double ys; + + bad_mt += verify_one(h_y[s], s, &ys); + } + printf("TRANSFORMER_MT streams=%d bad=%d ysum=%.6f\n", NST, bad_mt, + dsum(h_y[0], seqdh)); + + /* ----------------- class 2: 5-launch loop ----------------- */ + CHECK(cudaStreamCreate(&st_loop)); + CHECK(cudaMalloc(&lb.q, b_seqdh)); + CHECK(cudaMalloc(&lb.k, b_seqdh)); + CHECK(cudaMalloc(&lb.v, b_seqdh)); + CHECK(cudaMalloc(&lb.s, b_seq2)); + CHECK(cudaMalloc(&lb.a, b_seq2)); + CHECK(cudaMalloc(&lb.av, b_seqdh)); + CHECK(cudaMalloc(&lb.y, b_seqdh)); /* unused; op writes per-iter y */ + for (i = 0; i < NLOOP; i++) { + float *hx = (float *)malloc(b_seqdh); + + if (!hx) + return 1; + for (s = 0; s < seqdh; s++) + hx[s] = xval(s + 11 * i); + CHECK(cudaMalloc(&d_xl[i], b_seqdh)); + CHECK(cudaMalloc(&d_yl[i], b_seqdh)); + CHECK(cudaMemcpy(d_xl[i], hx, b_seqdh, + cudaMemcpyHostToDevice)); + h_yl[i] = (float *)malloc(b_seqdh); + if (!h_yl[i]) + return 1; + free(hx); + } + + for (i = 0; i < NLOOP; i++) { + struct attn_bufs b = lb; + + b.y = d_yl[i]; + launch_chain(st_loop, d_xl[i], b, d_wq, d_wk, d_wv, d_wo); + CHECK(cudaGetLastError()); + } + CHECK(cudaStreamSynchronize(st_loop)); + for (i = 0; i < NLOOP; i++) { + double ys; + + CHECK(cudaMemcpy(h_yl[i], d_yl[i], b_seqdh, + cudaMemcpyDeviceToHost)); + bad_loop += verify_one(h_yl[i], 11 * i, &ys); + } + printf("TRANSFORMER_LOOP launches=%d bad=%d ysum=%.6f\n", NLOOP, + bad_loop, dsum(h_yl[NLOOP - 1], seqdh)); + printf("TIMING ms=%.3f\n", ms); + + if (bad_mt != 0 || bad_loop != 0) { + printf("STATUS FAIL\n"); + return 1; + } + + printf("STATUS OK\n"); + for (s = 0; s < NST; s++) { + CHECK(cudaEventDestroy(ev[s])); + CHECK(cudaStreamDestroy(st[s])); + CHECK(cudaFree(d_x[s])); + CHECK(cudaFree(ab[s].q)); + CHECK(cudaFree(ab[s].k)); + CHECK(cudaFree(ab[s].v)); + CHECK(cudaFree(ab[s].s)); + CHECK(cudaFree(ab[s].a)); + CHECK(cudaFree(ab[s].av)); + CHECK(cudaFree(ab[s].y)); + free(h_y[s]); + } + for (i = 0; i < NLOOP; i++) { + CHECK(cudaFree(d_xl[i])); + CHECK(cudaFree(d_yl[i])); + free(h_yl[i]); + } + CHECK(cudaEventDestroy(t0)); + CHECK(cudaEventDestroy(t1)); + CHECK(cudaStreamDestroy(st_loop)); + CHECK(cudaFree(lb.q)); + CHECK(cudaFree(lb.k)); + CHECK(cudaFree(lb.v)); + CHECK(cudaFree(lb.s)); + CHECK(cudaFree(lb.a)); + CHECK(cudaFree(lb.av)); + CHECK(cudaFree(lb.y)); + CHECK(cudaFree(d_wq)); + CHECK(cudaFree(d_wk)); + CHECK(cudaFree(d_wv)); + CHECK(cudaFree(d_wo)); + free(wl); + return 0; +} -- cgit v1.2.3