tomevault-io Tutorial: Adding a New Kernel to RL.cu
Goal
Implement high performance cuda kernel with target format:
- include files: include/kernels/[kernel_name].cuh
- source files: src/kernels/[kernel_name].cu
- testing: tests/test_[kernel_name].cu
- Makefile: Compile the testing implementation Same format for non-cuda function (e.g. cpp) After finsihed implementation, you should explain how the kernel works.
Example
include/kernels/attention.cuh — Flash Attention 2 Kernel
#pragma once
#include <cuda_fp16.h>
#include <cuda_runtime.h>
// ---------------------------------------------------------------------------
// Flash Attention 2 — Prefill
// ---------------------------------------------------------------------------
// Q, K, V layout: [B, S, H, D] (H = num_heads, D = head_dim)
// O layout: [B, S, H_q, D]
//
// GQA supported: H_q % H_kv == 0. kv_head = q_head * H_kv / H_q.
// Causal mask applied (token i attends only to positions <= i).
// FP16 I/O, FP32 accumulation. Currently specialised for head_dim = 128.
//
// Tile sizes: Br=16 (Q rows per block), Bc=64 (KV cols per tile).
// Shared memory per block: 2 × Bc × D × sizeof(half) = 32 KB.
// ---------------------------------------------------------------------------
void launch_flash_attention_prefill(
const half* Q,
const half* K,
const half* V,
half* O,
int B, int S, int H_q, int H_kv, int head_dim,
cudaStream_t stream = 0
);
// ---------------------------------------------------------------------------
// Paged Attention — Decode
// ---------------------------------------------------------------------------
// Single new token (S_q=1) attends to full context stored in a paged KV cache.
//
// q layout: [num_seqs, H_q, D]
// k_cache / v_cache: [num_blocks, H_kv, block_size, D]
// block_tables: [num_seqs, max_blocks_per_seq] (logical→physical)
// seq_lens: [num_seqs] — context length including the new token
//
// Specialised for head_dim=128, block_size=16.
// ---------------------------------------------------------------------------
void launch_paged_attention_decode(
const half* q,
const half* k_cache,
const half* v_cache,
half* out,
const int* block_tables,
const int* seq_lens,
int num_seqs, int H_q, int H_kv, int head_dim,
int max_blocks_per_seq, int block_size,
cudaStream_t stream = 0
);
src/kernels/attention.cu — Flash Attention 2 Kernel Implementation
#include "kernels/attention.cuh"
#include <cuda_fp16.h>
#include <cuda_runtime.h>
#include <math.h>
// =============================================================================
// FA2 Prefill Kernel
// =============================================================================
//
// Grid: (ceil(S/Br), H_q, B)
// Block: (Br) — one thread per Q row in the current tile
//
// Algorithm (Flash Attention 2 with online softmax):
// Outer loop: iterate over KV tiles (K and V loaded into shared memory)
// Inner loop: each thread computes scores Q[q_row] · K[kv_pos] for all
// positions in the tile, updating (row_max, row_sum, O_acc)
// with the standard online softmax rescaling step.
//
// Shared memory layout:
// K_smem[Bc][HEAD_DIM] — current K tile, FP16, 16 KB
// V_smem[Bc][HEAD_DIM] — current V tile, FP16, 16 KB
// Total: 32 KB per block
//
// Registers per thread (FP32):
// q_reg[HEAD_DIM] — query row (loaded once)
// o_acc[HEAD_DIM] — output accumulator
// row_max, row_sum — online softmax state
// =============================================================================
template<int HEAD_DIM, int Br, int Bc>
__global__ void flash_attention_prefill_kernel(
const half* __restrict__ Q, // [B, S, H_q, D]
const half* __restrict__ K, // [B, S, H_kv, D]
const half* __restrict__ V, // [B, S, H_kv, D]
half* __restrict__ O, // [B, S, H_q, D]
int B, int S, int H_q, int H_kv, float scale
) {
const int r = threadIdx.x; // Row within Q tile (0..Br-1)
const int q_tile = blockIdx.x;
const int h_q = blockIdx.y;
const int b = blockIdx.z;
const int h_kv = h_q * H_kv / H_q; // GQA: integer division
const int q_row = q_tile * Br + r; // Global row in Q
// NOTE: do NOT return early here — all Br threads must participate in the
// cooperative K/V tile load and reach every __syncthreads(). Use a flag
// instead and skip computation/writes for out-of-bounds rows.
const bool valid_q = (q_row < S);
// Shared memory — K and V tiles, FP16
__shared__ half K_smem[Bc][HEAD_DIM];
__shared__ half V_smem[Bc][HEAD_DIM];
// Load Q row into registers (FP32, converted once) — only for valid rows.
float q_reg[HEAD_DIM];
#pragma unroll
for (int d = 0; d < HEAD_DIM; d++) q_reg[d] = 0.0f;
if (valid_q) {
const half* q_ptr = Q + ((long)(b * S + q_row) * H_q + h_q) * HEAD_DIM;
#pragma unroll
for (int d = 0; d < HEAD_DIM; d++)
q_reg[d] = __half2float(q_ptr[d]);
}
// Online softmax state (per Q row, in registers)
float row_max = -INFINITY;
float row_sum = 0.0f;
float o_acc[HEAD_DIM];
#pragma unroll
for (int d = 0; d < HEAD_DIM; d++) o_acc[d] = 0.0f;
const int num_kv_tiles = (S + Bc - 1) / Bc;
for (int kv_tile = 0; kv_tile < num_kv_tiles; kv_tile++) {
const int tile_start = kv_tile * Bc;
// ------------------------------------------------------------------
// Cooperative load: ALL Br threads fill K_smem and V_smem (Bc rows).
// Each thread is responsible for rows: r, r+Br, r+2*Br, r+3*Br.
// All threads must participate — including those with q_row >= S —
// so that every K/V row needed by valid threads is populated before
// the __syncthreads() below.
// ------------------------------------------------------------------
for (int row = r; row < Bc; row += Br) {
const int kv_row = tile_start + row;
if (kv_row < S) {
const half* k_ptr = K + ((long)(b * S + kv_row) * H_kv + h_kv) * HEAD_DIM;
const half* v_ptr = V + ((long)(b * S + kv_row) * H_kv + h_kv) * HEAD_DIM;
#pragma unroll
for (int d = 0; d < HEAD_DIM; d++) {
K_smem[row][d] = k_ptr[d];
V_smem[row][d] = v_ptr[d];
}
}
}
__syncthreads();
// ------------------------------------------------------------------
// Compute attention scores and update online softmax (valid rows only)
// ------------------------------------------------------------------
if (valid_q) {
const int tile_end = min(tile_start + Bc, S);
for (int c = 0; c < tile_end - tile_start; c++) {
const int kv_pos = tile_start + c;
// Causal mask: later positions in this row are also masked (break ok)
if (kv_pos > q_row) break;
// Dot product: Q[q_row] · K[kv_pos]
float dot = 0.0f;
#pragma unroll
for (int d = 0; d < HEAD_DIM; d++)
dot += q_reg[d] * __half2float(K_smem[c][d]);
dot *= scale;
// Online softmax update (numerically stable)
const float new_max = fmaxf(row_max, dot);
const float alpha = expf(row_max - new_max); // rescale old accum
const float p = expf(dot - new_max); // weight for new token
row_sum = row_sum * alpha + p;
#pragma unroll
for (int d = 0; d < HEAD_DIM; d++)
o_acc[d] = o_acc[d] * alpha + p * __half2float(V_smem[c][d]);
row_max = new_max;
}
}
__syncthreads();
}
// Normalise and write output (FP32 → FP16) — valid rows only
if (valid_q) {
half* o_ptr = O + ((long)(b * S + q_row) * H_q + h_q) * HEAD_DIM;
const float inv_sum = 1.0f / row_sum;
#pragma unroll
for (int d = 0; d < HEAD_DIM; d++)
o_ptr[d] = __float2half(o_acc[d] * inv_sum);
}
}
// =============================================================================
// Paged Attention Decode Kernel
// =============================================================================
//
// Grid: (num_seqs, H_q)
// Block: (1) — one thread per (sequence, q_head)
//
// The single query token attends to all context_len tokens stored in the
// paged KV cache. Block tables map logical KV blocks to physical locations.
//
// Algorithm: same online softmax as prefill, but without shared memory tiling
// (Q has only one row; no tiling over Q needed). Walks block_table to handle
// non-contiguous physical storage.
// =============================================================================
template<int HEAD_DIM, int BLOCK_SIZE>
__global__ void paged_attention_decode_kernel(
const half* __restrict__ q, // [num_seqs, H_q, D]
const half* __restrict__ k_cache, // [num_blocks, H_kv, BLOCK_SIZE, D]
const half* __restrict__ v_cache, // [num_blocks, H_kv, BLOCK_SIZE, D]
half* __restrict__ out, // [num_seqs, H_q, D]
const int* __restrict__ block_tables, // [num_seqs, max_blocks_per_seq]
const int* __restrict__ seq_lens, // [num_seqs]
float scale,
int max_blocks_per_seq,
int H_q, int H_kv
) {
const int seq_idx = blockIdx.x;
const int h_q = blockIdx.y;
const int h_kv = h_q * H_kv / H_q; // GQA
const int context_len = seq_lens[seq_idx];
const int num_blocks = (context_len + BLOCK_SIZE - 1) / BLOCK_SIZE;
// Load Q into registers (FP32)
float q_reg[HEAD_DIM];
{
const half* q_ptr = q + (seq_idx * H_q + h_q) * HEAD_DIM;
#pragma unroll
for (int d = 0; d < HEAD_DIM; d++)
q_reg[d] = __half2float(q_ptr[d]);
}
float max_score = -INFINITY;
float sum_exp = 0.0f;
float acc[HEAD_DIM];
#pragma unroll
for (int d = 0; d < HEAD_DIM; d++) acc[d] = 0.0f;
const int* seq_block_table = block_tables + seq_idx * max_blocks_per_seq;
for (int blk = 0; blk < num_blocks; blk++) {
// Indirection: logical block index → physical block in the pool
const int physical_block = seq_block_table[blk];
const int tokens_in_block = min(BLOCK_SIZE, context_len - blk * BLOCK_SIZE);
// Base offset for this physical block + KV head
// k_cache layout: [num_blocks, H_kv, BLOCK_SIZE, D]
const long block_base =
((long)physical_block * H_kv + h_kv) * BLOCK_SIZE * HEAD_DIM;
for (int tok = 0; tok < tokens_in_block; tok++) {
const half* k_ptr = k_cache + block_base + (long)tok * HEAD_DIM;
const half* v_ptr = v_cache + block_base + (long)tok * HEAD_DIM;
float score = 0.0f;
#pragma unroll
for (int d = 0; d < HEAD_DIM; d++)
score += q_reg[d] * __half2float(k_ptr[d]);
score *= scale;
// Online softmax update
const float new_max = fmaxf(max_score, score);
const float alpha = expf(max_score - new_max);
const float p = expf(score - new_max);
sum_exp = sum_exp * alpha + p;
#pragma unroll
for (int d = 0; d < HEAD_DIM; d++)
acc[d] = acc[d] * alpha + p * __half2float(v_ptr[d]);
max_score = new_max;
}
}
// Write normalised output
half* out_ptr = out + (seq_idx * H_q + h_q) * HEAD_DIM;
const float inv = 1.0f / sum_exp;
#pragma unroll
for (int d = 0; d < HEAD_DIM; d++)
out_ptr[d] = __float2half(acc[d] * inv);
}
// =============================================================================
// Launch wrappers
// =============================================================================
void launch_flash_attention_prefill(
const half* Q,
const half* K,
const half* V,
half* O,
int B, int S, int H_q, int H_kv, int head_dim,
cudaStream_t stream
) {
constexpr int Br = 16;
constexpr int Bc = 64;
const float scale = 1.0f / sqrtf((float)head_dim);
dim3 grid((S + Br - 1) / Br, H_q, B);
dim3 block(Br);
flash_attention_prefill_kernel<128, Br, Bc><<<grid, block, 0, stream>>>(
Q, K, V, O, B, S, H_q, H_kv, scale
);
}
void launch_paged_attention_decode(
const half* q,
const half* k_cache,
const half* v_cache,
half* out,
const int* block_tables,
const int* seq_lens,
int num_seqs, int H_q, int H_kv, int head_dim,
int max_blocks_per_seq, int block_size,
cudaStream_t stream
) {
const float scale = 1.0f / sqrtf((float)head_dim);
dim3 grid(num_seqs, H_q);
dim3 block(1);
paged_attention_decode_kernel<128, 16><<<grid, block, 0, stream>>>(
q, k_cache, v_cache, out,
block_tables, seq_lens,
scale, max_blocks_per_seq, H_q, H_kv
);
}
tests/test_attention.cu — Flash Attention 2 Kernel Testing
// test_attention.cu
// Validates Flash Attention 2 (prefill) and Paged Attention (decode) kernels
// against a CPU FP32 naive reference.
//
// Pass criterion (from DESIGN.md): max absolute error < 5e-3 vs naive attention.
//
// Prefill test cases:
// 1. B=1, S=16, H_q=2, H_kv=1 — small, GQA ratio 2:1
// 2. B=1, S=64, H_q=4, H_kv=2 — single KV tile (Bc=64)
// 3. B=1, S=128, H_q=16, H_kv=8 — Qwen3 head config, two KV tiles
// 4. B=2, S=96, H_q=4, H_kv=2 — batched
//
// Decode test cases:
// 5. num_seqs=1, context_len=32, H_q=2, H_kv=1
// 6. num_seqs=2, context_len=128, H_q=16, H_kv=8
#include <cstdio>
#include <cstdlib>
#include <cmath>
#include <cuda_fp16.h>
#include <cuda_runtime.h>
#include "kernels/attention.cuh"
// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
#define CUDA_CHECK(call) \
do { \
cudaError_t err = (call); \
if (err != cudaSuccess) { \
fprintf(stderr, "CUDA error at %s:%d — %s\n", \
__FILE__, __LINE__, cudaGetErrorString(err)); \
exit(EXIT_FAILURE); \
} \
} while (0)
// Deterministic LCG (no srand dependency)
static unsigned int lcg_state = 42u;
static float lcg_randf() {
lcg_state = lcg_state * 1664525u + 1013904223u;
return ((float)(lcg_state >> 1) / (float)0x7fffffffu) * 2.0f - 1.0f;
}
// ---------------------------------------------------------------------------
// CPU reference: naive causal attention
// Q, K, V layout: [B, S, H, D] (H = num_heads)
// O layout: [B, S, H_q, D]
// FP32 throughout.
// ---------------------------------------------------------------------------
static void ref_attention(
float* O,
const float* Q,
const float* K,
const float* V,
int B, int S, int H_q, int H_kv, int D
) {
const float scale = 1.0f / sqrtf((float)D);
float* scores = new float[S];
for (int b = 0; b < B; b++) {
for (int h = 0; h < H_q; h++) {
const int hkv = h * H_kv / H_q;
for (int i = 0; i < S; i++) {
const float* qi = Q + ((b * S + i) * H_q + h) * D;
// Compute scores[0..i]
float max_s = -1e30f;
for (int j = 0; j <= i; j++) {
const float* kj = K + ((b * S + j) * H_kv + hkv) * D;
float dot = 0.0f;
for (int d = 0; d < D; d++) dot += qi[d] * kj[d];
scores[j] = dot * scale;
if (scores[j] > max_s) max_s = scores[j];
}
// Softmax over [0..i]
float sum_e = 0.0f;
for (int j = 0; j <= i; j++) {
scores[j] = expf(scores[j] - max_s);
sum_e += scores[j];
}
for (int j = 0; j <= i; j++) scores[j] /= sum_e;
// O[i] = softmax(S) · V[0..i]
float* oi = O + ((b * S + i) * H_q + h) * D;
for (int d = 0; d < D; d++) {
float acc = 0.0f;
for (int j = 0; j <= i; j++) {
const float* vj = V + ((b * S + j) * H_kv + hkv) * D;
acc += scores[j] * vj[d];
}
oi[d] = acc;
}
}
}
}
delete[] scores;
}
// ---------------------------------------------------------------------------
// CPU reference: non-causal attention (for decode, S_q=1)
// q layout: [num_seqs, H_q, D]
// K, V layout: [num_seqs, S_kv, H_kv, D] — contiguous KV context
// O layout: [num_seqs, H_q, D]
// ---------------------------------------------------------------------------
static void ref_decode_attention(
float* O,
const float* q,
const float* K,
const float* V,
int num_seqs, int S_kv, int H_q, int H_kv, int D
) {
const float scale = 1.0f / sqrtf((float)D);
float* scores = new float[S_kv];
for (int s = 0; s < num_seqs; s++) {
for (int h = 0; h < H_q; h++) {
const int hkv = h * H_kv / H_q;
const float* qi = q + (s * H_q + h) * D;
float max_s = -1e30f;
for (int j = 0; j < S_kv; j++) {
const float* kj = K + ((s * S_kv + j) * H_kv + hkv) * D;
float dot = 0.0f;
for (int d = 0; d < D; d++) dot += qi[d] * kj[d];
scores[j] = dot * scale;
if (scores[j] > max_s) max_s = scores[j];
}
float sum_e = 0.0f;
for (int j = 0; j < S_kv; j++) {
scores[j] = expf(scores[j] - max_s);
sum_e += scores[j];
}
for (int j = 0; j < S_kv; j++) scores[j] /= sum_e;
float* oi = O + (s * H_q + h) * D;
for (int d = 0; d < D; d++) {
float acc = 0.0f;
for (int j = 0; j < S_kv; j++) {
const float* vj = V + ((s * S_kv + j) * H_kv + hkv) * D;
acc += scores[j] * vj[d];
}
oi[d] = acc;
}
}
}
delete[] scores;
}
// ---------------------------------------------------------------------------
// Prefill test runner
// ---------------------------------------------------------------------------
static bool run_prefill_test(
const char* name, int B, int S, int H_q, int H_kv,
int D = 128, float tol = 5e-3f
) {
const long N_q = (long)B * S * H_q * D;
const long N_kv = (long)B * S * H_kv * D;
// Host buffers (FP32 for reference, FP16 for kernel)
float* h_Q_f32 = new float[N_q];
float* h_K_f32 = new float[N_kv];
float* h_V_f32 = new float[N_kv];
float* h_ref = new float[N_q];
half* h_Q = new half[N_q];
half* h_K = new half[N_kv];
half* h_V = new half[N_kv];
half* h_out = new half[N_q];
// Fill with random values, scale to prevent softmax saturation
for (long i = 0; i < N_q; i++) h_Q[i] = __float2half(lcg_randf() * 0.5f);
for (long i = 0; i < N_kv; i++) h_K[i] = __float2half(lcg_randf() * 0.5f);
for (long i = 0; i < N_kv; i++) h_V[i] = __float2half(lcg_randf() * 0.5f);
// FP32 reference uses the FP16-quantised values (round-trip)
for (long i = 0; i < N_q; i++) h_Q_f32[i] = __half2float(h_Q[i]);
for (long i = 0; i < N_kv; i++) h_K_f32[i] = __half2float(h_K[i]);
for (long i = 0; i < N_kv; i++) h_V_f32[i] = __half2float(h_V[i]);
ref_attention(h_ref, h_Q_f32, h_K_f32, h_V_f32, B, S, H_q, H_kv, D);
// Device buffers
half *d_Q, *d_K, *d_V, *d_O;
CUDA_CHECK(cudaMalloc(&d_Q, N_q * sizeof(half)));
CUDA_CHECK(cudaMalloc(&d_K, N_kv * sizeof(half)));
CUDA_CHECK(cudaMalloc(&d_V, N_kv * sizeof(half)));
CUDA_CHECK(cudaMalloc(&d_O, N_q * sizeof(half)));
CUDA_CHECK(cudaMemcpy(d_Q, h_Q, N_q * sizeof(half), cudaMemcpyHostToDevice));
CUDA_CHECK(cudaMemcpy(d_K, h_K, N_kv * sizeof(half), cudaMemcpyHostToDevice));
CUDA_CHECK(cudaMemcpy(d_V, h_V, N_kv * sizeof(half), cudaMemcpyHostToDevice));
launch_flash_attention_prefill(d_Q, d_K, d_V, d_O, B, S, H_q, H_kv, D);
CUDA_CHECK(cudaDeviceSynchronize());
CUDA_CHECK(cudaMemcpy(h_out, d_O, N_q * sizeof(half), cudaMemcpyDeviceToHost));
// Compare
float max_err = 0.0f;
for (long i = 0; i < N_q; i++) {
float diff = fabsf(__half2float(h_out[i]) - h_ref[i]);
if (diff > max_err) max_err = diff;
}
bool passed = (max_err < tol);
printf("[%s] %-50s max_err=%.6f %s\n",
passed ? "PASS" : "FAIL", name, max_err,
passed ? "" : "<-- EXCEEDS 5e-3");
cudaFree(d_Q); cudaFree(d_K); cudaFree(d_V); cudaFree(d_O);
delete[] h_Q_f32; delete[] h_K_f32; delete[] h_V_f32; delete[] h_ref;
delete[] h_Q; delete[] h_K; delete[] h_V; delete[] h_out;
return passed;
}
// ---------------------------------------------------------------------------
// Paged attention decode test runner
// Uses a contiguous block mapping (block i → physical block i) so the paged
// result is equivalent to standard attention on contiguous K/V.
// ---------------------------------------------------------------------------
static bool run_decode_test(
const char* name,
int num_seqs, int context_len, int H_q, int H_kv,
int D = 128, int BLOCK_SIZE = 16, float tol = 5e-3f
) {
// Total KV tokens per sequence (all seqs same length for simplicity)
const long Q_total = (long)num_seqs * H_q * D;
const long O_total = Q_total;
// Host buffers (FP32 reference)
float* h_q_f32 = new float[Q_total];
float* h_K_f32 = new float[(long)num_seqs * context_len * H_kv * D];
float* h_V_f32 = new float[(long)num_seqs * context_len * H_kv * D];
float* h_ref = new float[O_total];
half* h_q = new half[Q_total];
half* h_out = new half[O_total];
for (long i = 0; i < Q_total; i++) h_q[i] = __float2half(lcg_randf() * 0.5f);
for (long i = 0; i < (long)num_seqs * context_len * H_kv * D; i++) {
h_K_f32[i] = lcg_randf() * 0.5f;
h_V_f32[i] = lcg_randf() * 0.5f;
}
for (long i = 0; i < Q_total; i++) h_q_f32[i] = __half2float(h_q[i]);
// Quantise K/V round-trip
for (long i = 0; i < (long)num_seqs * context_len * H_kv * D; i++) {
h_K_f32[i] = __half2float(__float2half(h_K_f32[i]));
h_V_f32[i] = __half2float(__float2half(h_V_f32[i]));
}
ref_decode_attention(h_ref, h_q_f32, h_K_f32, h_V_f32,
num_seqs, context_len, H_q, H_kv, D);
// Build the paged KV cache with an identity block table
// Each sequence uses context_len/BLOCK_SIZE (rounded up) blocks.
const int blocks_per_seq = (context_len + BLOCK_SIZE - 1) / BLOCK_SIZE;
const int total_phys_blocks = num_seqs * blocks_per_seq;
// k_cache layout: [num_blocks, H_kv, BLOCK_SIZE, D]
const long cache_elems = (long)total_phys_blocks * H_kv * BLOCK_SIZE * D;
half* h_k_cache = new half[cache_elems]();
half* h_v_cache = new half[cache_elems]();
// Fill paged cache from contiguous K/V
// seq s, token t, kv head h → logical_block = t/BLOCK_SIZE, offset = t%BLOCK_SIZE
// physical_block = s * blocks_per_seq + logical_block
for (int s = 0; s < num_seqs; s++) {
for (int t = 0; t < context_len; t++) {
for (int h = 0; h < H_kv; h++) {
int logical_block = t / BLOCK_SIZE;
int tok_offset = t % BLOCK_SIZE;
int phys_block = s * blocks_per_seq + logical_block;
const float* ksrc = h_K_f32 + ((s * context_len + t) * H_kv + h) * D;
const float* vsrc = h_V_f32 + ((s * context_len + t) * H_kv + h) * D;
half* kdst = h_k_cache
+ ((long)phys_block * H_kv + h) * BLOCK_SIZE * D
+ tok_offset * D;
half* vdst = h_v_cache
+ ((long)phys_block * H_kv + h) * BLOCK_SIZE * D
+ tok_offset * D;
for (int d = 0; d < D; d++) {
kdst[d] = __float2half(ksrc[d]);
vdst[d] = __float2half(vsrc[d]);
}
}
}
}
// Block tables: identity mapping (seq s, logical block b → physical block s*bps+b)
int* h_block_tables = new int[num_seqs * blocks_per_seq];
int* h_seq_lens = new int[num_seqs];
for (int s = 0; s < num_seqs; s++) {
h_seq_lens[s] = context_len;
for (int b = 0; b < blocks_per_seq; b++)
h_block_tables[s * blocks_per_seq + b] = s * blocks_per_seq + b;
}
// Device allocations
half *d_q, *d_k_cache, *d_v_cache, *d_out;
int *d_block_tables, *d_seq_lens;
CUDA_CHECK(cudaMalloc(&d_q, Q_total * sizeof(half)));
CUDA_CHECK(cudaMalloc(&d_k_cache, cache_elems * sizeof(half)));
CUDA_CHECK(cudaMalloc(&d_v_cache, cache_elems * sizeof(half)));
CUDA_CHECK(cudaMalloc(&d_out, O_total * sizeof(half)));
CUDA_CHECK(cudaMalloc(&d_block_tables, num_seqs * blocks_per_seq * sizeof(int)));
CUDA_CHECK(cudaMalloc(&d_seq_lens, num_seqs * sizeof(int)));
CUDA_CHECK(cudaMemcpy(d_q, h_q, Q_total * sizeof(half), cudaMemcpyHostToDevice));
CUDA_CHECK(cudaMemcpy(d_k_cache, h_k_cache, cache_elems * sizeof(half), cudaMemcpyHostToDevice));
CUDA_CHECK(cudaMemcpy(d_v_cache, h_v_cache, cache_elems * sizeof(half), cudaMemcpyHostToDevice));
CUDA_CHECK(cudaMemcpy(d_block_tables, h_block_tables, num_seqs * blocks_per_seq * sizeof(int), cudaMemcpyHostToDevice));
CUDA_CHECK(cudaMemcpy(d_seq_lens, h_seq_lens, num_seqs * sizeof(int), cudaMemcpyHostToDevice));
launch_paged_attention_decode(
d_q, d_k_cache, d_v_cache, d_out,
d_block_tables, d_seq_lens,
num_seqs, H_q, H_kv, D,
blocks_per_seq, BLOCK_SIZE
);
CUDA_CHECK(cudaDeviceSynchronize());
CUDA_CHECK(cudaMemcpy(h_out, d_out, O_total * sizeof(half), cudaMemcpyDeviceToHost));
float max_err = 0.0f;
for (long i = 0; i < O_total; i++) {
float diff = fabsf(__half2float(h_out[i]) - h_ref[i]);
if (diff > max_err) max_err = diff;
}
bool passed = (max_err < tol);
printf("[%s] %-50s max_err=%.6f %s\n",
passed ? "PASS" : "FAIL", name, max_err,
passed ? "" : "<-- EXCEEDS 5e-3");
cudaFree(d_q); cudaFree(d_k_cache); cudaFree(d_v_cache);
cudaFree(d_out); cudaFree(d_block_tables); cudaFree(d_seq_lens);
delete[] h_q_f32; delete[] h_K_f32; delete[] h_V_f32; delete[] h_ref;
delete[] h_q; delete[] h_out;
delete[] h_k_cache; delete[] h_v_cache;
delete[] h_block_tables; delete[] h_seq_lens;
return passed;
}
// ---------------------------------------------------------------------------
// Benchmarks
// ---------------------------------------------------------------------------
static void run_prefill_benchmark(
const char* name, int B, int S, int H_q, int H_kv,
int D = 128, int warmup = 10, int iters = 200
) {
const long N_q = (long)B * S * H_q * D;
const long N_kv = (long)B * S * H_kv * D;
half *d_Q, *d_K, *d_V, *d_O;
CUDA_CHECK(cudaMalloc(&d_Q, N_q * sizeof(half)));
CUDA_CHECK(cudaMalloc(&d_K, N_kv * sizeof(half)));
CUDA_CHECK(cudaMalloc(&d_V, N_kv * sizeof(half)));
CUDA_CHECK(cudaMalloc(&d_O, N_q * sizeof(half)));
cudaEvent_t ev0, ev1;
CUDA_CHECK(cudaEventCreate(&ev0));
CUDA_CHECK(cudaEventCreate(&ev1));
for (int i = 0; i < warmup; i++)
launch_flash_attention_prefill(d_Q, d_K, d_V, d_O, B, S, H_q, H_kv, D);
CUDA_CHECK(cudaDeviceSynchronize());
CUDA_CHECK(cudaEventRecord(ev0));
for (int i = 0; i < iters; i++)
launch_flash_attention_prefill(d_Q, d_K, d_V, d_O, B, S, H_q, H_kv, D);
CUDA_CHECK(cudaEventRecord(ev1));
CUDA_CHECK(cudaEventSynchronize(ev1));
float ms = 0.0f;
CUDA_CHECK(cudaEventElapsedTime(&ms, ev0, ev1));
const float us = ms * 1000.0f / iters;
// FLOPs: 2 * B * H_q * S^2 * D (Q@K^T and softmax(S)@V, approximately)
const double flops = 2.0 * B * H_q * (double)S * S * D * 2;
const double tflops = flops / (us * 1e-6) / 1e12;
printf("[BENCH] %-50s %7.2f us %5.2f TFLOPS\n", name, us, tflops);
CUDA_CHECK(cudaEventDestroy(ev0));
CUDA_CHECK(cudaEventDestroy(ev1));
cudaFree(d_Q); cudaFree(d_K); cudaFree(d_V); cudaFree(d_O);
}
// ---------------------------------------------------------------------------
// main
// ---------------------------------------------------------------------------
int main() {
printf("=== Flash Attention kernel tests ===\n\n");
// ── Prefill tests ────────────────────────────────────────────────────────
printf("--- Prefill (FA2) ---\n");
bool all_pass = true;
all_pass &= run_prefill_test(
"Prefill B=1 S=16 H_q=2 H_kv=1 (small GQA 2:1)",
1, 16, 2, 1);
all_pass &= run_prefill_test(
"Prefill B=1 S=64 H_q=4 H_kv=2 (single KV tile)",
1, 64, 4, 2);
all_pass &= run_prefill_test(
"Prefill B=1 S=128 H_q=16 H_kv=8 (Qwen3 heads, 2 KV tiles)",
1, 128, 16, 8);
all_pass &= run_prefill_test(
"Prefill B=2 S=96 H_q=4 H_kv=2 (batched)",
2, 96, 4, 2);
all_pass &= run_prefill_test(
"Prefill B=1 S=256 H_q=16 H_kv=8 (4 KV tiles)",
1, 256, 16, 8);
// ── Decode tests ─────────────────────────────────────────────────────────
printf("\n--- Decode (Paged Attention) ---\n");
all_pass &= run_decode_test(
"Decode num_seqs=1 ctx=32 H_q=2 H_kv=1",
1, 32, 2, 1);
all_pass &= run_decode_test(
"Decode num_seqs=1 ctx=128 H_q=16 H_kv=8 (Qwen3)",
1, 128, 16, 8);
all_pass &= run_decode_test(
"Decode num_seqs=2 ctx=64 H_q=4 H_kv=2 (batched seqs)",
2, 64, 4, 2);
all_pass &= run_decode_test(
"Decode num_seqs=4 ctx=128 H_q=16 H_kv=8 (GRPO-style batch)",
4, 128, 16, 8);
// ── Summary ──────────────────────────────────────────────────────────────
printf("\n%s\n", all_pass ? "All tests PASSED." : "Some tests FAILED.");
// ── Benchmarks ───────────────────────────────────────────────────────────
printf("\n=== Prefill benchmarks (warmup=10, iters=200) ===\n");
run_prefill_benchmark("Prefill B=1 S=128 H_q=16 H_kv=8", 1, 128, 16, 8);
run_prefill_benchmark("Prefill B=1 S=512 H_q=16 H_kv=8", 1, 512, 16, 8);
run_prefill_benchmark("Prefill B=1 S=2048 H_q=16 H_kv=8", 1, 2048, 16, 8);
return all_pass ? 0 : 1;
}
# ──────────────────────────────────────────────────────────────────────────────
# Makefile — GRPO-CUDA kernel tests
#
# Quick-start:
# make tests # build + run all three kernel tests
# make test_rmsnorm # build + run rmsnorm only
# make generate_refs # generate PyTorch reference data (needs torch)
# make clean # remove build/
#
# To target a specific GPU architecture instead of auto-detect:
# make tests ARCH=sm_120 # Blackwell (RTX PRO 6000 / RTX 50xx) ← this machine
# make tests ARCH=sm_89 # Ada (RTX 40xx)
# make tests ARCH=sm_86 # Ampere (RTX 30xx)
# make tests ARCH=sm_80 # A100
# ──────────────────────────────────────────────────────────────────────────────
CUDA_HOME := /usr/local/cuda-12.8
NVCC := $(CUDA_HOME)/bin/nvcc
ARCH := sm_120
INCLUDES := -I include
NVCCFLAGS := -O2 -std=c++17 $(INCLUDES) --gpu-architecture=$(ARCH)
BUILDDIR := build
PYTHON := python3
# ── Kernel test binaries ───────────────────────────────────────────────────────
$(BUILDDIR)/test_rmsnorm: src/kernels/rmsnorm.cu tests/test_rmsnorm.cu | $(BUILDDIR)
$(NVCC) $(NVCCFLAGS) $^ -o $@
$(BUILDDIR)/test_softmax: src/kernels/softmax.cu tests/test_softmax.cu | $(BUILDDIR)
$(NVCC) $(NVCCFLAGS) $^ -o $@
$(BUILDDIR)/test_swiglu: src/kernels/swiglu.cu tests/test_swiglu.cu | $(BUILDDIR)
$(NVCC) $(NVCCFLAGS) $^ -o $@
$(BUILDDIR)/test_attention: src/kernels/attention.cu tests/test_attention.cu | $(BUILDDIR)
$(NVCC) $(NVCCFLAGS) $^ -o $@
$(BUILDDIR)/test_kv_cache: src/model/kv_cache.cu tests/test_kv_cache.cu | $(BUILDDIR)
$(NVCC) $(NVCCFLAGS) $^ -o $@
$(BUILDDIR)/test_rope: src/kernels/rope.cu tests/test_rope.cu | $(BUILDDIR)
$(NVCC) $(NVCCFLAGS) $^ -o $@
$(BUILDDIR)/test_embedding: src/kernels/embedding.cu tests/test_embedding.cu | $(BUILDDIR)
$(NVCC) $(NVCCFLAGS) $^ -o $@
$(BUILDDIR)/test_linear: src/kernels/linear.cu tests/test_linear.cu | $(BUILDDIR)
$(NVCC) $(NVCCFLAGS) $^ -o $@ -lcublas
QWEN3_SRCS := src/model/qwen3.cu src/model/kv_cache.cu \
src/kernels/rmsnorm.cu src/kernels/rope.cu src/kernels/attention.cu \
src/kernels/swiglu.cu src/kernels/embedding.cu src/kernels/linear.cu \
src/kernels/config.cpp src/kernels/weights.cpp
$(BUILDDIR)/test_qwen3: $(QWEN3_SRCS) tests/test_qwen3.cu | $(BUILDDIR)
$(NVCC) $(NVCCFLAGS) $^ -o $@ -lcublas
$(BUILDDIR)/bench_decode: $(QWEN3_SRCS) tests/bench_decode.cu | $(BUILDDIR)
$(NVCC) $(NVCCFLAGS) $^ -o $@ -lcublas
$(BUILDDIR)/test_loading_weights: src/kernels/config.cpp src/kernels/weights.cpp tests/test_loading_weights.cpp | $(BUILDDIR)
$(NVCC) $(NVCCFLAGS) $^ -o $@
$(BUILDDIR):
mkdir -p $(BUILDDIR)
# ── Run targets ────────────────────────────────────────────────────────────────
.PHONY: test_rmsnorm test_softmax test_swiglu test_attention test_kv_cache test_rope test_embedding test_linear test_qwen3 test_loading_weights bench_decode tests generate_refs clean
test_rmsnorm: $(BUILDDIR)/test_rmsnorm
./$(BUILDDIR)/test_rmsnorm
test_softmax: $(BUILDDIR)/test_softmax
./$(BUILDDIR)/test_softmax
test_swiglu: $(BUILDDIR)/test_swiglu
./$(BUILDDIR)/test_swiglu
test_attention: $(BUILDDIR)/test_attention
./$(BUILDDIR)/test_attention
test_kv_cache: $(BUILDDIR)/test_kv_cache
./$(BUILDDIR)/test_kv_cache
test_rope: $(BUILDDIR)/test_rope
./$(BUILDDIR)/test_rope
test_embedding: $(BUILDDIR)/test_embedding
./$(BUILDDIR)/test_embedding
test_linear: $(BUILDDIR)/test_linear
./$(BUILDDIR)/test_linear
test_qwen3: $(BUILDDIR)/test_qwen3
./$(BUILDDIR)/test_qwen3
test_loading_weights: $(BUILDDIR)/test_loading_weights
./$(BUILDDIR)/test_loading_weights
bench_decode: $(BUILDDIR)/bench_decode
./$(BUILDDIR)/bench_decode
tests: test_rmsnorm test_softmax test_swiglu test_attention test_kv_cache test_rope test_embedding test_linear test_qwen3 test_loading_weights
# ── PyTorch reference generator ───────────────────────────────────────────────
generate_refs:
$(PYTHON) tests/generate_references.py --outdir tests/reference_data
# ── Cleanup ────────────────────────────────────────────────────────────────────
clean:
rm -rf $(BUILDDIR)
Key points:
- This fully implementation with CUDA/C++ and cublas, do not contain torch api.
- The model will load in fp16, the kernel should optimization for fp16 computation.
- The kernel implementation should consider for LLM RL training.
References
docs/DESIGN.mdinclude/kernels/attention.cuhsrc/kernels/attention.cuinclude/kernels/rmsnorm.cuhsrc/kernels/rmsnorm.cutests/test_attention.cutests/test_rmsnorm.cu
Source: KJLdefeated/RL.cu — distributed by TomeVault.