File size: 3,434 Bytes
9425aed | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | #include <stdint.h>
#include "cuda_validation_chain.h"
/* cuLaunchKernel signature — resolved at runtime from nvcuda.dll via cuda_driver_loader */
typedef int CUresult;
typedef struct CUstream_st* CUstream;
typedef CUresult (*cuLaunchKernel_fn)(void*, unsigned,unsigned,unsigned,
unsigned,unsigned,unsigned,
unsigned, CUstream, void**, void**);
static sov_rowm_record_t* g_rowm_rec = 0;
static sov_cuda_auth_t* g_auth = 0;
static cuLaunchKernel_fn g_launch_fn = 0;
/* Called by cuda_driver_loader after it resolves cuLaunchKernel from nvcuda.dll */
void sov_cuda_kernels_set_auth(sov_rowm_record_t* rec, sov_cuda_auth_t* auth) {
g_rowm_rec = rec;
g_auth = auth;
}
void sov_cuda_kernels_set_launch_fn(void* fn) {
g_launch_fn = (cuLaunchKernel_fn)fn;
}
static int find_handle(const sov_cuda_auth_t* auth, void* fn) {
for (uint32_t i = 0; i < auth->count; i++) {
if (auth->handles[i] == fn) return (int)i;
}
return -1;
}
int sov_cuda_launch_kernel(void* fn,
unsigned gx, unsigned gy, unsigned gz,
unsigned bx, unsigned by, unsigned bz,
unsigned shared_mem, void* stream, void** args) {
if (!g_rowm_rec || !g_auth) return -3;
if (sov_rowm_check_authorized(g_rowm_rec, g_auth) != 0) return -3;
if (find_handle(g_auth, fn) < 0) return -5;
if (!g_launch_fn) return -6;
CUresult r = g_launch_fn(fn, gx, gy, gz, bx, by, bz,
shared_mem, (CUstream)stream, args, 0);
return r == 0 ? 0 : -7;
}
int sov_cuda_flash_attention(int seqs, int heads,
float* q, float* k, float* v, float* out,
int* block_table, int* seq_lens,
int head_dim, int block_size) {
if (!g_rowm_rec || !g_auth) return -3;
if (sov_rowm_check_authorized(g_rowm_rec, g_auth) != 0) return -3;
/* Flash attention kernel handle is auth->handles[0] by convention */
if (!g_auth->count || !g_launch_fn) return -6;
/* Grid: ceil(seqs/32) x heads x 1 — one warp per sequence per head */
unsigned grid_x = (unsigned)((seqs + 31) / 32);
unsigned grid_y = (unsigned)heads;
void* kargs[] = { &seqs, &heads, &q, &k, &v, &out,
&block_table, &seq_lens, &head_dim, &block_size };
return sov_cuda_launch_kernel(g_auth->handles[0],
grid_x, grid_y, 1,
32, 1, 1,
0, 0, kargs);
}
int sov_cuda_gemm(float* a, float* b, float* c, int m, int n, int k_dim) {
if (!g_rowm_rec || !g_auth) return -3;
if (sov_rowm_check_authorized(g_rowm_rec, g_auth) != 0) return -3;
/* GEMM kernel handle is auth->handles[1] by convention */
if (g_auth->count < 2 || !g_launch_fn) return -6;
/* Grid: ceil(m/16) x ceil(n/16) — 16x16 tiles */
unsigned grid_x = (unsigned)((m + 15) / 16);
unsigned grid_y = (unsigned)((n + 15) / 16);
void* kargs[] = { &a, &b, &c, &m, &n, &k_dim };
return sov_cuda_launch_kernel(g_auth->handles[1],
grid_x, grid_y, 1,
16, 16, 1,
0, 0, kargs);
}
|