| |
| |
| |
|
|
| #include "cuda_driver_loader.h"
|
| #include "kv_allocator.h"
|
| #include "rowm_cuda_validation.h"
|
|
|
| #include <stdint.h>
|
| #include <stdio.h>
|
| #include <stdlib.h>
|
| #include <string.h>
|
|
|
| #define PTR_VALUE(type, value) ((type)(uintptr_t)(value))
|
|
|
| static void check_failed(const char* expression, const char* file, int line) {
|
| fprintf(stderr, "CHECK failed: %s (%s:%d)\n", expression, file, line);
|
| exit(1);
|
| }
|
| #define CHECK(expression) \
|
| ((expression) ? (void)0 : check_failed(#expression, __FILE__, __LINE__))
|
|
|
| const unsigned char flash_attention_ptx[] = {'.', 'v', 0};
|
| const unsigned int flash_attention_ptx_len = 3;
|
| const unsigned char gemm_ptx[] = {'.', 'g', 0};
|
| const unsigned int gemm_ptx_len = 3;
|
| float g_janet_kernel_config[8] = {
|
| 1.0f, 0.5f, 0.25f, 0.125f, 8.0f, 16.0f, 32.0f, 64.0f
|
| };
|
|
|
| int sov_cuda_flash_attention(int seqs, int query_heads,
|
| uint32_t layer_index,
|
| CUdeviceptr q, CUdeviceptr out,
|
| CUdeviceptr block_table,
|
| CUdeviceptr seq_lens);
|
|
|
| static CUmodule g_fake_module = PTR_VALUE(CUmodule, 0x1000);
|
| static CUfunction g_fake_flash = PTR_VALUE(CUfunction, 0x2000);
|
| static CUfunction g_fake_rmsnorm = PTR_VALUE(CUfunction, 0x3000);
|
| static CUfunction g_fake_silu = PTR_VALUE(CUfunction, 0x4000);
|
| static CUdeviceptr g_fake_power = UINT64_C(0x5000);
|
| static CUdeviceptr g_fake_config = UINT64_C(0x6000);
|
|
|
| static int g_driver_initialized;
|
| static int g_driver_init_calls;
|
| static int g_module_load_calls;
|
| static int g_module_unload_calls;
|
| static int g_lookup_calls;
|
| static int g_fail_lookup_call;
|
| static int g_register_calls;
|
| static CUdeviceptr g_registered_power;
|
| static int g_power_sync_calls;
|
| static int g_config_copy_calls;
|
| static int g_gemm_init_calls;
|
| static int g_gemm_init_result;
|
| static int g_gemm_shutdown_calls;
|
| static int g_expected_launch;
|
| static int g_launch_calls;
|
| static int g_rowm_commit_calls;
|
|
|
| static int g_allocator_ready;
|
| static sov_kv_allocator_config_t g_allocator_config;
|
| static CUdeviceptr g_allocator_k;
|
| static CUdeviceptr g_allocator_v;
|
| static size_t g_token_stride;
|
| static size_t g_page_stride;
|
| static size_t g_layer_stride;
|
| static size_t g_plane_stride;
|
|
|
| static void reset_allocator_contract(void) {
|
| g_allocator_ready = 1;
|
| g_allocator_config.physical_block_count = 3;
|
| g_allocator_config.layer_count = 2;
|
| g_allocator_config.kv_head_count = 2;
|
| g_allocator_config.head_dim = 64;
|
| g_allocator_config.element_bytes = 2;
|
| g_allocator_k = UINT64_C(0x20000);
|
| g_token_stride = 256;
|
| g_page_stride = 4096;
|
| g_layer_stride = 12288;
|
| g_plane_stride = 24576;
|
| g_allocator_v = g_allocator_k + (CUdeviceptr)g_plane_stride;
|
| }
|
|
|
| int sov_kv_get_config(sov_kv_allocator_config_t* out) {
|
| if (!g_allocator_ready || !out) return -1;
|
| *out = g_allocator_config;
|
| return 0;
|
| }
|
| CUdeviceptr sov_kv_get_k_base(void) { return g_allocator_k; }
|
| CUdeviceptr sov_kv_get_v_base(void) { return g_allocator_v; }
|
| size_t sov_kv_get_token_stride(void) { return g_token_stride; }
|
| size_t sov_kv_get_page_stride(void) { return g_page_stride; }
|
| size_t sov_kv_get_layer_stride(void) { return g_layer_stride; }
|
| size_t sov_kv_get_plane_stride(void) { return g_plane_stride; }
|
|
|
| int sov_cuda_init(void) { ++g_driver_init_calls; g_driver_initialized = 1; return CUDA_SUCCESS; }
|
| int sov_cuda_is_initialized(void) { return g_driver_initialized; }
|
| uint64_t sov_cuda_context_generation(void) { return g_driver_initialized ? UINT64_C(1) : UINT64_C(0); }
|
|
|
| CUresult sov_cuda_get_runtime_identity(sov_cuda_runtime_identity_t* id) {
|
| size_t i;
|
| CHECK(id != 0);
|
| memset(id, 0, sizeof(*id));
|
| id->driver_version = 12080;
|
| id->compute_capability_major = 8;
|
| id->compute_capability_minor = 9;
|
| id->context_generation = sov_cuda_context_generation();
|
| for (i = 0; i < sizeof(id->device_uuid); ++i)
|
| id->device_uuid[i] = (uint8_t)(i + 1u);
|
| return CUDA_SUCCESS;
|
| }
|
|
|
| CUresult sov_cuda_module_load_data(CUmodule* module, const void* image) {
|
| CHECK(module != 0);
|
| CHECK(image == flash_attention_ptx);
|
| ++g_module_load_calls;
|
| *module = g_fake_module;
|
| return CUDA_SUCCESS;
|
| }
|
|
|
| CUresult sov_cuda_module_get_function(CUfunction* function, CUmodule module,
|
| const char* name) {
|
| ++g_lookup_calls;
|
| CHECK(function != 0);
|
| CHECK(module == g_fake_module);
|
| if (g_lookup_calls == g_fail_lookup_call) return CUDA_ERROR_NOT_FOUND;
|
| if (strcmp(name, "flash_attention_paged") == 0) *function = g_fake_flash;
|
| else if (strcmp(name, "rmsnorm_fused") == 0) *function = g_fake_rmsnorm;
|
| else if (strcmp(name, "silu_fused") == 0) *function = g_fake_silu;
|
| else CHECK(!"unexpected kernel lookup");
|
| return CUDA_SUCCESS;
|
| }
|
|
|
| CUresult sov_cuda_module_get_global(CUdeviceptr* device_ptr, size_t* bytes,
|
| CUmodule module, const char* name) {
|
| CHECK(device_ptr != 0); CHECK(bytes != 0); CHECK(module == g_fake_module);
|
| if (strcmp(name, "power_state") == 0) {
|
| *device_ptr = g_fake_power; *bytes = sizeof(int32_t);
|
| } else if (strcmp(name, "janet_kernel_config") == 0) {
|
| *device_ptr = g_fake_config; *bytes = 256;
|
| } else CHECK(!"unexpected module-global lookup");
|
| return CUDA_SUCCESS;
|
| }
|
|
|
| CUresult sov_cuda_module_unload(CUmodule module) {
|
| CHECK(module == g_fake_module);
|
| ++g_module_unload_calls;
|
| return CUDA_SUCCESS;
|
| }
|
|
|
| int sov_cuda_memcpy_h2d(void* device_dst, const void* host_src, size_t bytes) {
|
| CHECK((CUdeviceptr)(uintptr_t)device_dst == g_fake_config);
|
| CHECK(host_src == g_janet_kernel_config);
|
| CHECK(bytes == sizeof(g_janet_kernel_config));
|
| ++g_config_copy_calls;
|
| return CUDA_SUCCESS;
|
| }
|
|
|
| CUresult sov_cuda_register_power_state_device(CUdeviceptr device_ptr) {
|
| ++g_register_calls;
|
| g_registered_power = device_ptr;
|
| return CUDA_SUCCESS;
|
| }
|
|
|
| CUresult sov_cuda_sync_power_state(void) {
|
| CHECK(g_registered_power == g_fake_power);
|
| ++g_power_sync_calls;
|
| return CUDA_SUCCESS;
|
| }
|
|
|
| CUresult sov_cuda_launch_kernel(CUfunction function,
|
| unsigned int grid_x, unsigned int grid_y,
|
| unsigned int grid_z,
|
| unsigned int block_x, unsigned int block_y,
|
| unsigned int block_z,
|
| unsigned int shared_mem_bytes,
|
| CUstream stream,
|
| void** kernel_params,
|
| void** extra) {
|
| ++g_launch_calls;
|
| CHECK(grid_z == 1); CHECK(block_z == 1);
|
| CHECK(shared_mem_bytes == 0); CHECK(stream == 0);
|
| CHECK(kernel_params != 0); CHECK(extra == 0);
|
|
|
| if (g_expected_launch == 1) {
|
| CHECK(function == g_fake_flash);
|
| CHECK(grid_x == 2 && grid_y == 4);
|
| CHECK(block_x == 1 && block_y == 1);
|
| CHECK(*(CUdeviceptr*)kernel_params[0] == UINT64_C(0x10000));
|
| CHECK(*(CUdeviceptr*)kernel_params[1] == g_allocator_k);
|
| CHECK(*(CUdeviceptr*)kernel_params[2] == g_allocator_v);
|
| CHECK(*(CUdeviceptr*)kernel_params[3] == UINT64_C(0x40000));
|
| CHECK(*(CUdeviceptr*)kernel_params[4] == UINT64_C(0x50000));
|
| CHECK(*(CUdeviceptr*)kernel_params[5] == UINT64_C(0x60000));
|
| CHECK(*(uint64_t*)kernel_params[6] == 256);
|
| CHECK(*(uint64_t*)kernel_params[7] == 4096);
|
| CHECK(*(uint64_t*)kernel_params[8] == 12288);
|
| CHECK(*(uint32_t*)kernel_params[9] == 1);
|
| CHECK(*(uint32_t*)kernel_params[10] == 4);
|
| CHECK(*(uint32_t*)kernel_params[11] == 2);
|
| CHECK(*(uint32_t*)kernel_params[12] == 64);
|
| CHECK(*(uint32_t*)kernel_params[13] == 3);
|
| CHECK(*(uint32_t*)kernel_params[14] == SOV_KV_MAX_BLOCKS_PER_SEQ);
|
| } else if (g_expected_launch == 2) {
|
| CHECK(function == g_fake_rmsnorm);
|
| CHECK(grid_x == 1 && grid_y == 1);
|
| CHECK(block_x == 1 && block_y == 1);
|
| CHECK(*(CUdeviceptr*)kernel_params[0] == UINT64_C(0x70000));
|
| CHECK(*(CUdeviceptr*)kernel_params[1] == UINT64_C(0x80000));
|
| CHECK(*(CUdeviceptr*)kernel_params[2] == UINT64_C(0x90000));
|
| CHECK(*(uint32_t*)kernel_params[3] == 129);
|
| } else if (g_expected_launch == 3) {
|
| CHECK(function == g_fake_silu);
|
| CHECK(grid_x == 2 && grid_y == 1);
|
| CHECK(block_x == 128 && block_y == 1);
|
| CHECK(*(CUdeviceptr*)kernel_params[0] == UINT64_C(0xa0000));
|
| CHECK(*(CUdeviceptr*)kernel_params[1] == UINT64_C(0xb0000));
|
| CHECK(*(uint32_t*)kernel_params[2] == 129);
|
| } else CHECK(!"unexpected launch");
|
| return CUDA_SUCCESS;
|
| }
|
|
|
| int sov_cuda_gemm_init(void) { ++g_gemm_init_calls; return g_gemm_init_result; }
|
| void sov_cuda_gemm_shutdown(void) { ++g_gemm_shutdown_calls; }
|
|
|
| static int commit_cuda_validation(
|
| void* rowm_context,
|
| const sov_cuda_validation_evidence_t* evidence,
|
| sov_rowm_cuda_validation_commit_t* commit_out) {
|
| size_t index;
|
| ++g_rowm_commit_calls;
|
| CHECK(rowm_context == (void*)(uintptr_t)0xcafeu);
|
| CHECK(evidence != 0);
|
| CHECK(evidence->ptx_target == 89);
|
| CHECK(evidence->flash_ptx.bytes == flash_attention_ptx);
|
| CHECK(evidence->flash_ptx.size == 2);
|
| CHECK(evidence->gemm_ptx.bytes == gemm_ptx);
|
| CHECK(evidence->gemm_ptx.size == 2);
|
| CHECK(evidence->resolved_kernel_mask == SOV_CUDA_REQUIRED_KERNEL_MASK);
|
| memset(commit_out, 0, sizeof(*commit_out));
|
| commit_out->version = SOV_ROWM_CUDA_COMMIT_VERSION;
|
| commit_out->validation_epoch = 1;
|
| for (index = 0; index < 32; ++index) {
|
| commit_out->record_id[index] = (uint8_t)(index + 1u);
|
| commit_out->committed_rowm_root[index] = (uint8_t)(index + 2u);
|
| commit_out->embedded_worm_receipt_hash[index] = (uint8_t)(index + 3u);
|
| }
|
| return 0;
|
| }
|
|
|
| static int call_attention(int seqs, int query_heads, uint32_t layer,
|
| CUdeviceptr q, CUdeviceptr out,
|
| CUdeviceptr table, CUdeviceptr lens) {
|
| return sov_cuda_flash_attention(seqs, query_heads, layer, q, out, table, lens);
|
| }
|
|
|
| static void test_lookup_failure_cleanup(void) {
|
| g_fail_lookup_call = 2;
|
| CHECK(sov_cuda_kernels_init() == CUDA_ERROR_NOT_FOUND);
|
| CHECK(g_driver_init_calls == 1);
|
| CHECK(g_module_load_calls == 1);
|
| CHECK(g_module_unload_calls == 1);
|
| CHECK(g_registered_power == 0);
|
| CHECK(g_gemm_init_calls == 0);
|
| CHECK(g_gemm_shutdown_calls == 1);
|
| g_fail_lookup_call = 0;
|
| }
|
|
|
| static void test_gemm_failure_cleanup(void) {
|
| g_gemm_init_result = -33;
|
| CHECK(sov_cuda_kernels_init() == -33);
|
| CHECK(g_module_load_calls == 2);
|
| CHECK(g_module_unload_calls == 2);
|
| CHECK(g_registered_power == 0);
|
| CHECK(g_gemm_init_calls == 1);
|
| CHECK(g_gemm_shutdown_calls == 2);
|
| g_gemm_init_result = 0;
|
| }
|
|
|
| static void test_dispatch_and_shutdown(void) {
|
| reset_allocator_contract();
|
| CHECK(sov_cuda_kernels_init() == CUDA_SUCCESS);
|
| CHECK(sov_cuda_kernels_init() == CUDA_SUCCESS);
|
| CHECK(g_module_load_calls == 3);
|
| CHECK(g_config_copy_calls == 2);
|
| CHECK(g_registered_power == g_fake_power);
|
| CHECK(g_gemm_init_calls == 3);
|
|
|
| g_expected_launch = 1;
|
| CHECK(call_attention(2, 4, 1,
|
| UINT64_C(0x10000), UINT64_C(0x40000),
|
| UINT64_C(0x50000), UINT64_C(0x60000))
|
| == SOV_CUDA_ROWM_UNAUTHORIZED);
|
| CHECK(g_launch_calls == 0);
|
|
|
| CHECK(sov_cuda_kernels_authorize_rowm(
|
| commit_cuda_validation, (void*)(uintptr_t)0xcafeu)
|
| == SOV_CUDA_ROWM_OK);
|
| CHECK(g_rowm_commit_calls == 1);
|
|
|
| g_expected_launch = 1;
|
| CHECK(call_attention(2, 4, 1,
|
| UINT64_C(0x10000), UINT64_C(0x40000),
|
| UINT64_C(0x50000), UINT64_C(0x60000))
|
| == CUDA_SUCCESS);
|
|
|
| g_expected_launch = 2;
|
| CHECK(sov_cuda_rmsnorm_fused(
|
| UINT64_C(0x70000), UINT64_C(0x80000),
|
| UINT64_C(0x90000), 129) == CUDA_SUCCESS);
|
|
|
| g_expected_launch = 3;
|
| CHECK(sov_cuda_silu_fused(
|
| UINT64_C(0xa0000), UINT64_C(0xb0000), 129) == CUDA_SUCCESS);
|
|
|
| sov_cuda_kernels_shutdown();
|
| CHECK(g_registered_power == 0);
|
| CHECK(g_module_unload_calls == 3);
|
| CHECK(g_gemm_shutdown_calls == 3);
|
| CHECK(g_power_sync_calls == 3);
|
| CHECK(g_launch_calls == 3);
|
| CHECK(sov_cuda_validation_require_authorized(1) == SOV_CUDA_ROWM_UNAUTHORIZED);
|
| }
|
|
|
| int main(void) {
|
| reset_allocator_contract();
|
| test_lookup_failure_cleanup();
|
| test_gemm_failure_cleanup();
|
| test_dispatch_and_shutdown();
|
| puts("cuda kernel dispatch host tests passed");
|
| return 0;
|
| }
|
|
|
