/* Cross-device GPU tensor copy test (mgpu-device-aware-cuda). * * Exercises: * - ds4_gpu_init_multi single-device and multi-device init paths. * - ds4_gpu_tensor_alloc_on / ds4_gpu_tensor_free_in_place. * - ds4_gpu_tensor_copy_xdev same-device fast path. * - ds4_gpu_tensor_copy_xdev peer-auto and DS4_FORCE_HOST_BOUNCE paths * (when 2+ GPUs are visible). */ /* ds4_gpu_mgpu.h is standalone-C-compatible — it now provides the * complete ds4_gpu_tensor struct + typedef so callers can use the * bare type name. ds4_gpu.h is also included here only for the * legacy ds4_gpu_init / _cleanup / _tensor_read / _tensor_write * prototypes the test uses; the mgpu header does not duplicate them. */ #include "ds4_gpu.h" #include "ds4_gpu_mgpu.h" #include #include #include #include #include #define CHECK(cond, msg) \ do { \ if (!(cond)) { \ fprintf(stderr, "FAIL: %s (line %d)\n", (msg), __LINE__); \ return 1; \ } \ } while (0) static int run_one(int n_gpus_wanted, int force_bounce) { int dev_count = 0; (void)cudaGetDeviceCount(&dev_count); if (dev_count < 1) { fprintf(stderr, "no CUDA devices visible\n"); return 0; } if (n_gpus_wanted > dev_count) { fprintf(stderr, "skip: wanted %d GPUs, have %d\n", n_gpus_wanted, dev_count); return 0; } ds4_gpu_config cfg; memset(&cfg, 0, sizeof(cfg)); cfg.n_gpus = n_gpus_wanted; for (int i = 0; i < n_gpus_wanted; i++) cfg.device_indices[i] = i; if (force_bounce) setenv("DS4_FORCE_HOST_BOUNCE", "1", 1); else unsetenv("DS4_FORCE_HOST_BOUNCE"); CHECK(ds4_gpu_init_multi(&cfg), "init_multi"); const size_t N = 256 * 1024; float *host_src = (float *)malloc(N * sizeof(float)); float *host_dst = (float *)malloc(N * sizeof(float)); CHECK(host_src && host_dst, "host alloc"); for (size_t i = 0; i < N; i++) host_src[i] = (float)(i % 997) * 0.5f; ds4_gpu_tensor a; memset(&a, 0, sizeof(a)); ds4_gpu_tensor b; memset(&b, 0, sizeof(b)); CHECK(ds4_gpu_tensor_alloc_on(&a, 0, N * sizeof(float)) == 0, "alloc_on dev 0"); int dst_dev = (n_gpus_wanted == 2) ? 1 : 0; CHECK(ds4_gpu_tensor_alloc_on(&b, dst_dev, N * sizeof(float)) == 0, "alloc_on dst dev"); CHECK(ds4_gpu_tensor_write(&a, 0, host_src, N * sizeof(float)), "tensor_write"); CHECK(ds4_gpu_tensor_copy_xdev(&b, &a, N * sizeof(float)), "tensor_copy_xdev"); CHECK(ds4_gpu_tensor_read(&b, 0, host_dst, N * sizeof(float)), "tensor_read"); CHECK(memcmp(host_src, host_dst, N * sizeof(float)) == 0, "data mismatch"); /* device_id round-trip. */ CHECK(ds4_gpu_tensor_device(&a) == 0, "device_id a"); CHECK(ds4_gpu_tensor_device(&b) == dst_dev, "device_id b"); ds4_gpu_tensor_free_in_place(&a); ds4_gpu_tensor_free_in_place(&b); free(host_src); free(host_dst); ds4_gpu_cleanup(); fprintf(stderr, " test_gpu_xdev OK (n_gpus=%d, force_bounce=%d)\n", n_gpus_wanted, force_bounce); return 0; } static int run_copy3(int n_gpus_wanted, int force_bounce) { int dev_count = 0; (void)cudaGetDeviceCount(&dev_count); if (dev_count < 1) { fprintf(stderr, "no CUDA devices visible\n"); return 0; } if (n_gpus_wanted > dev_count) { fprintf(stderr, "skip copy3: wanted %d GPUs, have %d\n", n_gpus_wanted, dev_count); return 0; } ds4_gpu_config cfg; memset(&cfg, 0, sizeof(cfg)); cfg.n_gpus = n_gpus_wanted; for (int i = 0; i < n_gpus_wanted; i++) cfg.device_indices[i] = i; if (force_bounce) setenv("DS4_FORCE_HOST_BOUNCE", "1", 1); else unsetenv("DS4_FORCE_HOST_BOUNCE"); CHECK(ds4_gpu_init_multi(&cfg), "copy3 init_multi"); const uint32_t n_embd = 4096u; const uint32_t n_expert = 3u; const size_t norm_bytes = (size_t)n_embd * sizeof(float); const size_t selected_bytes = (size_t)n_expert * sizeof(int32_t); const size_t weights_bytes = (size_t)n_expert * sizeof(float); int32_t host_selected[3] = {17, 99, 203}; float host_weights[3] = {0.25f, 0.5f, 0.75f}; int32_t got_selected[3] = {0, 0, 0}; float got_weights[3] = {0.0f, 0.0f, 0.0f}; float *host_norm = (float *)malloc(norm_bytes); float *got_norm = (float *)malloc(norm_bytes); CHECK(host_norm && got_norm, "copy3 host alloc"); for (uint32_t i = 0; i < n_embd; i++) host_norm[i] = (float)((int)i - 1024) * 0.015625f; const int dst_dev = (n_gpus_wanted == 2) ? 1 : 0; ds4_gpu_tensor norm_src; memset(&norm_src, 0, sizeof(norm_src)); ds4_gpu_tensor sel_src; memset(&sel_src, 0, sizeof(sel_src)); ds4_gpu_tensor w_src; memset(&w_src, 0, sizeof(w_src)); ds4_gpu_tensor norm_dst; memset(&norm_dst, 0, sizeof(norm_dst)); ds4_gpu_tensor sel_dst; memset(&sel_dst, 0, sizeof(sel_dst)); ds4_gpu_tensor w_dst; memset(&w_dst, 0, sizeof(w_dst)); CHECK(ds4_gpu_tensor_alloc_on(&norm_src, 0, norm_bytes) == 0, "copy3 alloc norm_src"); CHECK(ds4_gpu_tensor_alloc_on(&sel_src, 0, selected_bytes) == 0, "copy3 alloc sel_src"); CHECK(ds4_gpu_tensor_alloc_on(&w_src, 0, weights_bytes) == 0, "copy3 alloc w_src"); CHECK(ds4_gpu_tensor_alloc_on(&norm_dst, dst_dev, norm_bytes) == 0, "copy3 alloc norm_dst"); CHECK(ds4_gpu_tensor_alloc_on(&sel_dst, dst_dev, selected_bytes) == 0, "copy3 alloc sel_dst"); CHECK(ds4_gpu_tensor_alloc_on(&w_dst, dst_dev, weights_bytes) == 0, "copy3 alloc w_dst"); int ok = ds4_gpu_tensor_write(&norm_src, 0, host_norm, norm_bytes) && ds4_gpu_tensor_write(&sel_src, 0, host_selected, selected_bytes) && ds4_gpu_tensor_write(&w_src, 0, host_weights, weights_bytes) && ds4_gpu_tensor_copy_xdev3(&norm_dst, &norm_src, norm_bytes, &sel_dst, &sel_src, selected_bytes, &w_dst, &w_src, weights_bytes) && ds4_gpu_tensor_read(&norm_dst, 0, got_norm, norm_bytes) && ds4_gpu_tensor_read(&sel_dst, 0, got_selected, selected_bytes) && ds4_gpu_tensor_read(&w_dst, 0, got_weights, weights_bytes); CHECK(ok, "copy3 IO"); CHECK(memcmp(host_norm, got_norm, norm_bytes) == 0, "copy3 norm"); CHECK(memcmp(host_selected, got_selected, selected_bytes) == 0, "copy3 selected"); CHECK(memcmp(host_weights, got_weights, weights_bytes) == 0, "copy3 weights"); ds4_gpu_tensor_free_in_place(&norm_src); ds4_gpu_tensor_free_in_place(&sel_src); ds4_gpu_tensor_free_in_place(&w_src); ds4_gpu_tensor_free_in_place(&norm_dst); ds4_gpu_tensor_free_in_place(&sel_dst); ds4_gpu_tensor_free_in_place(&w_dst); free(host_norm); free(got_norm); ds4_gpu_cleanup(); fprintf(stderr, " copy_xdev3 OK (n_gpus=%d, force_bounce=%d)\n", n_gpus_wanted, force_bounce); return 0; } static int run_top1(void) { int dev_count = 0; (void)cudaGetDeviceCount(&dev_count); if (dev_count < 1) { fprintf(stderr, "no CUDA devices visible\n"); return 0; } ds4_gpu_config cfg; memset(&cfg, 0, sizeof(cfg)); cfg.n_gpus = 1; cfg.device_indices[0] = 0; CHECK(ds4_gpu_init_multi(&cfg), "top1 init_multi"); const uint32_t n_comp = 131073u; const uint32_t n_tokens = 3u; const size_t n_score = (size_t)n_comp * n_tokens; float *host_scores = (float *)malloc(n_score * sizeof(float)); uint32_t *host_sel = (uint32_t *)calloc(n_tokens, sizeof(uint32_t)); CHECK(host_scores && host_sel, "top1 host alloc"); for (size_t i = 0; i < n_score; i++) host_scores[i] = -1000.0f; host_scores[17] = 3.5f; host_scores[120000] = 3.5f; /* tie: lower id wins */ host_scores[1000] = NAN; /* NaN is ignored like CPU argmax */ host_scores[(size_t)n_comp + 5] = 8.0f; host_scores[(size_t)n_comp + n_comp - 1u] = 8.0f; for (uint32_t i = 0; i < n_comp; i++) { host_scores[(size_t)2u * n_comp + i] = -INFINITY; } host_scores[(size_t)2u * n_comp + 1u] = NAN; ds4_gpu_tensor scores; memset(&scores, 0, sizeof(scores)); ds4_gpu_tensor selected; memset(&selected, 0, sizeof(selected)); CHECK(ds4_gpu_tensor_alloc_on(&scores, 0, n_score * sizeof(float)) == 0, "top1 alloc scores"); CHECK(ds4_gpu_tensor_alloc_on(&selected, 0, n_tokens * sizeof(uint32_t)) == 0, "top1 alloc selected"); int ok = ds4_gpu_tensor_write(&scores, 0, host_scores, n_score * sizeof(float)) && ds4_gpu_indexer_topk_tensor(&selected, &scores, n_comp, n_tokens, 1) && ds4_gpu_tensor_read(&selected, 0, host_sel, n_tokens * sizeof(uint32_t)); CHECK(ok, "top1 compute"); CHECK(host_sel[0] == 17u, "top1 row0"); CHECK(host_sel[1] == 5u, "top1 row1"); CHECK(host_sel[2] == 0u, "top1 row2"); ds4_gpu_tensor_free_in_place(&scores); ds4_gpu_tensor_free_in_place(&selected); free(host_scores); free(host_sel); ds4_gpu_cleanup(); fprintf(stderr, " top1 OK\n"); return 0; } typedef struct topk_ref_entry { float score; uint32_t index; } topk_ref_entry; static int topk_ref_cmp(const void *ap, const void *bp) { const topk_ref_entry *a = (const topk_ref_entry *)ap; const topk_ref_entry *b = (const topk_ref_entry *)bp; if (a->score > b->score) return -1; if (a->score < b->score) return 1; return a->index < b->index ? -1 : (a->index > b->index ? 1 : 0); } static int run_topk2048(void) { int dev_count = 0; (void)cudaGetDeviceCount(&dev_count); if (dev_count < 1) return 0; ds4_gpu_config cfg; memset(&cfg, 0, sizeof(cfg)); cfg.n_gpus = 1; cfg.device_indices[0] = 0; CHECK(ds4_gpu_init_multi(&cfg), "topk2048 init_multi"); const uint32_t cases[] = {3355u, 5003u}; const uint32_t n_tokens = 2u; const uint32_t top_k = 2048u; for (uint32_t ci = 0; ci < 2u; ci++) { const uint32_t n_comp = cases[ci]; const uint64_t n_scores = (uint64_t)n_tokens * n_comp; float *host_scores = (float *)malloc((size_t)n_scores * sizeof(float)); uint32_t *host_selected = (uint32_t *)malloc((size_t)n_tokens * top_k * sizeof(uint32_t)); topk_ref_entry *ref = (topk_ref_entry *)malloc((size_t)n_comp * sizeof(topk_ref_entry)); CHECK(host_scores && host_selected && ref, "topk2048 host alloc"); for (uint32_t t = 0; t < n_tokens; t++) { for (uint32_t i = 0; i < n_comp; i++) { host_scores[(uint64_t)t * n_comp + i] = (float)(((uint64_t)i * 37u + t * 101u) % 1009u); } } ds4_gpu_tensor scores; memset(&scores, 0, sizeof(scores)); ds4_gpu_tensor selected; memset(&selected, 0, sizeof(selected)); CHECK(ds4_gpu_tensor_alloc_on( &scores, 0, n_scores * sizeof(float)) == 0, "topk2048 alloc scores"); CHECK(ds4_gpu_tensor_alloc_on( &selected, 0, (uint64_t)n_tokens * top_k * sizeof(uint32_t)) == 0, "topk2048 alloc selected"); int ok = ds4_gpu_tensor_write( &scores, 0, host_scores, n_scores * sizeof(float)) && ds4_gpu_indexer_topk_tensor( &selected, &scores, n_comp, n_tokens, top_k) && ds4_gpu_tensor_read( &selected, 0, host_selected, (uint64_t)n_tokens * top_k * sizeof(uint32_t)); CHECK(ok, "topk2048 compute"); for (uint32_t t = 0; t < n_tokens; t++) { for (uint32_t i = 0; i < n_comp; i++) { ref[i].score = host_scores[(uint64_t)t * n_comp + i]; ref[i].index = i; } qsort(ref, n_comp, sizeof(ref[0]), topk_ref_cmp); for (uint32_t i = 0; i < top_k; i++) { if (host_selected[(uint64_t)t * top_k + i] != ref[i].index) { fprintf(stderr, "FAIL: topk2048 n=%u token=%u rank=%u got=%u want=%u\n", n_comp, t, i, host_selected[(uint64_t)t * top_k + i], ref[i].index); return 1; } } } ds4_gpu_tensor_free_in_place(&scores); ds4_gpu_tensor_free_in_place(&selected); free(host_scores); free(host_selected); free(ref); } ds4_gpu_cleanup(); fprintf(stderr, " topk2048 OK\n"); return 0; } static int run_glm_selected_attention(void) { int dev_count = 0; (void)cudaGetDeviceCount(&dev_count); if (dev_count < 1) return 0; ds4_gpu_config cfg; memset(&cfg, 0, sizeof(cfg)); cfg.n_gpus = 1; cfg.device_indices[0] = 0; CHECK(ds4_gpu_init_multi(&cfg), "glm selected attention init_multi"); const uint32_t n_head = 1u; const uint32_t kv_dim = 512u; const uint32_t qk_nope = 192u; const uint32_t qk_rope = 64u; const uint32_t q_dim = qk_nope + qk_rope; const uint32_t cache_cap = 5u; const uint32_t n_selected = 3u; float host_q[q_dim]; float host_low[kv_dim]; float host_kv[cache_cap * kv_dim]; float host_rope[cache_cap * qk_rope]; float host_kv_compact[n_selected * kv_dim]; float host_rope_compact[n_selected * qk_rope]; uint32_t host_selected[3] = {4u, 1u, 3u}; float causal[kv_dim]; float indexed[kv_dim]; for (uint32_t i = 0; i < q_dim; i++) host_q[i] = (float)((int)(i % 17u) - 8) * 0.01f; /* Isolate selected-row KV gathering. RoPE depends on the original cache * row, so a physically compacted reference is equivalent only when the * query's RoPE tail contributes zero to the score. */ for (uint32_t i = qk_nope; i < q_dim; i++) host_q[i] = 0.0f; for (uint32_t i = 0; i < kv_dim; i++) host_low[i] = (float)((int)(i % 13u) - 6) * 0.02f; for (uint32_t i = 0; i < cache_cap * kv_dim; i++) host_kv[i] = (float)((int)((i * 7u) % 19u) - 9) * 0.03f; for (uint32_t i = 0; i < cache_cap * qk_rope; i++) host_rope[i] = (float)((int)((i * 5u) % 11u) - 5) * 0.015f; for (uint32_t i = 0; i < n_selected; i++) { memcpy(host_kv_compact + (size_t)i * kv_dim, host_kv + (size_t)host_selected[i] * kv_dim, (size_t)kv_dim * sizeof(float)); memcpy(host_rope_compact + (size_t)i * qk_rope, host_rope + (size_t)host_selected[i] * qk_rope, (size_t)qk_rope * sizeof(float)); } ds4_gpu_tensor q; memset(&q, 0, sizeof(q)); ds4_gpu_tensor low; memset(&low, 0, sizeof(low)); ds4_gpu_tensor kv; memset(&kv, 0, sizeof(kv)); ds4_gpu_tensor rope; memset(&rope, 0, sizeof(rope)); ds4_gpu_tensor kv_compact; memset(&kv_compact, 0, sizeof(kv_compact)); ds4_gpu_tensor rope_compact; memset(&rope_compact, 0, sizeof(rope_compact)); ds4_gpu_tensor selected; memset(&selected, 0, sizeof(selected)); ds4_gpu_tensor out_causal; memset(&out_causal, 0, sizeof(out_causal)); ds4_gpu_tensor out_indexed; memset(&out_indexed, 0, sizeof(out_indexed)); CHECK(ds4_gpu_tensor_alloc_on(&q, 0, sizeof(host_q)) == 0, "glm selected attention alloc q"); CHECK(ds4_gpu_tensor_alloc_on(&low, 0, sizeof(host_low)) == 0, "glm selected attention alloc low"); CHECK(ds4_gpu_tensor_alloc_on(&kv, 0, sizeof(host_kv)) == 0, "glm selected attention alloc kv"); CHECK(ds4_gpu_tensor_alloc_on(&rope, 0, sizeof(host_rope)) == 0, "glm selected attention alloc rope"); CHECK(ds4_gpu_tensor_alloc_on(&kv_compact, 0, sizeof(host_kv_compact)) == 0, "glm selected attention alloc compact kv"); CHECK(ds4_gpu_tensor_alloc_on(&rope_compact, 0, sizeof(host_rope_compact)) == 0, "glm selected attention alloc compact rope"); CHECK(ds4_gpu_tensor_alloc_on(&selected, 0, sizeof(host_selected)) == 0, "glm selected attention alloc selected"); CHECK(ds4_gpu_tensor_alloc_on(&out_causal, 0, sizeof(causal)) == 0, "glm selected attention alloc causal"); CHECK(ds4_gpu_tensor_alloc_on(&out_indexed, 0, sizeof(indexed)) == 0, "glm selected attention alloc indexed"); int ok = ds4_gpu_tensor_write(&q, 0, host_q, sizeof(host_q)) && ds4_gpu_tensor_write(&low, 0, host_low, sizeof(host_low)) && ds4_gpu_tensor_write(&kv, 0, host_kv, sizeof(host_kv)) && ds4_gpu_tensor_write(&rope, 0, host_rope, sizeof(host_rope)) && ds4_gpu_tensor_write(&kv_compact, 0, host_kv_compact, sizeof(host_kv_compact)) && ds4_gpu_tensor_write(&rope_compact, 0, host_rope_compact, sizeof(host_rope_compact)) && ds4_gpu_tensor_write(&selected, 0, host_selected, sizeof(host_selected)) && ds4_gpu_glm_attention_indexed_batch_lora_causal_tensor( &out_causal, &q, &low, &kv_compact, &rope_compact, 1u, 2u, n_selected, n_selected, false, n_head, kv_dim, qk_nope, qk_rope, 0u, 8000000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f) && ds4_gpu_glm_attention_indexed_batch_lora_valid_tensor( &out_indexed, &q, &low, &kv, &rope, &selected, 1u, n_selected, cache_cap, false, n_head, kv_dim, qk_nope, qk_rope, 0u, 8000000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f) && ds4_gpu_tensor_read(&out_causal, 0, causal, sizeof(causal)) && ds4_gpu_tensor_read(&out_indexed, 0, indexed, sizeof(indexed)); CHECK(ok, "glm selected attention compute"); CHECK(memcmp(causal, indexed, sizeof(causal)) == 0, "glm selected attention mismatch"); ds4_gpu_tensor_free_in_place(&q); ds4_gpu_tensor_free_in_place(&low); ds4_gpu_tensor_free_in_place(&kv); ds4_gpu_tensor_free_in_place(&rope); ds4_gpu_tensor_free_in_place(&kv_compact); ds4_gpu_tensor_free_in_place(&rope_compact); ds4_gpu_tensor_free_in_place(&selected); ds4_gpu_tensor_free_in_place(&out_causal); ds4_gpu_tensor_free_in_place(&out_indexed); ds4_gpu_cleanup(); fprintf(stderr, " glm_selected_attention OK\n"); return 0; } static int run_glm_indexer_scores(void) { int dev_count = 0; (void)cudaGetDeviceCount(&dev_count); if (dev_count < 1) return 0; ds4_gpu_config cfg; memset(&cfg, 0, sizeof(cfg)); cfg.n_gpus = 1; cfg.device_indices[0] = 0; CHECK(ds4_gpu_init_multi(&cfg), "glm indexer scores init_multi"); const uint32_t n_rows = 137u; const uint32_t n_tokens = 17u; const uint32_t pos0 = 5u; const uint32_t n_head = 32u; const uint32_t head_dim = 128u; const float scale = 1.0f / 64.0f; const size_t q_count = (size_t)n_tokens * n_head * head_dim; const size_t weight_count = (size_t)n_tokens * n_head; const size_t cache_count = (size_t)n_rows * head_dim; const size_t score_count = (size_t)n_tokens * n_rows; float *host_q = (float *)malloc(q_count * sizeof(float)); float *host_weights = (float *)malloc(weight_count * sizeof(float)); float *host_cache = (float *)malloc(cache_count * sizeof(float)); float *host_fast = (float *)malloc(score_count * sizeof(float)); float *host_quality = (float *)malloc(score_count * sizeof(float)); float *host_ref = (float *)malloc(score_count * sizeof(float)); CHECK(host_q && host_weights && host_cache && host_fast && host_quality && host_ref, "glm indexer scores host alloc"); for (size_t i = 0; i < q_count; i++) host_q[i] = (float)((int)(i % 15u) - 7) / 32.0f; for (size_t i = 0; i < weight_count; i++) host_weights[i] = (float)((int)(i % 7u) - 3) / 8.0f; for (size_t i = 0; i < cache_count; i++) host_cache[i] = (float)((int)((i * 5u) % 13u) - 6) / 32.0f; for (uint32_t t = 0; t < n_tokens; t++) { for (uint32_t row = 0; row < n_rows; row++) { float total = 0.0f; for (uint32_t h = 0; h < n_head; h++) { float dot = 0.0f; const float *qh = host_q + ((size_t)t * n_head + h) * head_dim; const float *kr = host_cache + (size_t)row * head_dim; for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kr[d]; if (dot > 0.0f) total += dot * host_weights[(size_t)t * n_head + h]; } host_ref[(size_t)t * n_rows + row] = row < pos0 + t + 1u ? total * scale : -INFINITY; } } ds4_gpu_tensor q; memset(&q, 0, sizeof(q)); ds4_gpu_tensor weights; memset(&weights, 0, sizeof(weights)); ds4_gpu_tensor cache; memset(&cache, 0, sizeof(cache)); ds4_gpu_tensor scores; memset(&scores, 0, sizeof(scores)); CHECK(ds4_gpu_tensor_alloc_on(&q, 0, q_count * sizeof(float)) == 0, "glm indexer scores alloc q"); CHECK(ds4_gpu_tensor_alloc_on(&weights, 0, weight_count * sizeof(float)) == 0, "glm indexer scores alloc weights"); CHECK(ds4_gpu_tensor_alloc_on(&cache, 0, cache_count * sizeof(float)) == 0, "glm indexer scores alloc cache"); CHECK(ds4_gpu_tensor_alloc_on(&scores, 0, score_count * sizeof(float)) == 0, "glm indexer scores alloc output"); int ok = ds4_gpu_tensor_write(&q, 0, host_q, q_count * sizeof(float)) && ds4_gpu_tensor_write(&weights, 0, host_weights, weight_count * sizeof(float)) && ds4_gpu_tensor_write(&cache, 0, host_cache, cache_count * sizeof(float)); CHECK(ok, "glm indexer scores upload"); ds4_gpu_set_quality(false); ok = ds4_gpu_glm_indexer_scores_batch_tensor( &scores, &q, &weights, &cache, n_rows, n_tokens, pos0, n_head, head_dim, scale, false) && ds4_gpu_tensor_read(&scores, 0, host_fast, score_count * sizeof(float)); CHECK(ok, "glm indexer scores fast compute"); ds4_gpu_set_quality(true); ok = ds4_gpu_glm_indexer_scores_batch_tensor( &scores, &q, &weights, &cache, n_rows, n_tokens, pos0, n_head, head_dim, scale, false) && ds4_gpu_tensor_read(&scores, 0, host_quality, score_count * sizeof(float)); CHECK(ok, "glm indexer scores quality compute"); ds4_gpu_set_quality(false); CHECK(memcmp(host_fast, host_ref, score_count * sizeof(float)) == 0, "glm indexer fast scores mismatch"); CHECK(memcmp(host_quality, host_ref, score_count * sizeof(float)) == 0, "glm indexer quality scores mismatch"); ds4_gpu_tensor_free_in_place(&q); ds4_gpu_tensor_free_in_place(&weights); ds4_gpu_tensor_free_in_place(&cache); ds4_gpu_tensor_free_in_place(&scores); free(host_q); free(host_weights); free(host_cache); free(host_fast); free(host_quality); free(host_ref); ds4_gpu_cleanup(); fprintf(stderr, " glm_indexer_scores OK\n"); return 0; } static int run_moe_handoff_pack(void) { int dev_count = 0; (void)cudaGetDeviceCount(&dev_count); if (dev_count < 1) { fprintf(stderr, "no CUDA devices visible\n"); return 0; } ds4_gpu_config cfg; memset(&cfg, 0, sizeof(cfg)); cfg.n_gpus = 1; cfg.device_indices[0] = 0; CHECK(ds4_gpu_init_multi(&cfg), "moe_pack init_multi"); const uint32_t n_embd = 4096u; const uint32_t n_expert = 3u; const size_t norm_bytes = (size_t)n_embd * sizeof(float); const size_t selected_bytes = (size_t)n_expert * sizeof(int32_t); const size_t weights_bytes = (size_t)n_expert * sizeof(float); const size_t packed_bytes = norm_bytes + selected_bytes + weights_bytes; float *host_norm = (float *)malloc(norm_bytes); int32_t *host_selected = (int32_t *)malloc(selected_bytes); float *host_weights = (float *)malloc(weights_bytes); unsigned char *host_packed = (unsigned char *)calloc(1, packed_bytes); CHECK(host_norm && host_selected && host_weights && host_packed, "moe_pack host alloc"); for (uint32_t i = 0; i < n_embd; i++) host_norm[i] = (float)((int)i - 2000) * 0.03125f; for (uint32_t i = 0; i < n_expert; i++) { host_selected[i] = (int32_t)(200 + i * 7); host_weights[i] = 0.125f * (float)(i + 1); } ds4_gpu_tensor norm; memset(&norm, 0, sizeof(norm)); ds4_gpu_tensor selected; memset(&selected, 0, sizeof(selected)); ds4_gpu_tensor weights; memset(&weights, 0, sizeof(weights)); ds4_gpu_tensor packed; memset(&packed, 0, sizeof(packed)); CHECK(ds4_gpu_tensor_alloc_on(&norm, 0, norm_bytes) == 0, "moe_pack alloc norm"); CHECK(ds4_gpu_tensor_alloc_on(&selected, 0, selected_bytes) == 0, "moe_pack alloc selected"); CHECK(ds4_gpu_tensor_alloc_on(&weights, 0, weights_bytes) == 0, "moe_pack alloc weights"); CHECK(ds4_gpu_tensor_alloc_on(&packed, 0, packed_bytes) == 0, "moe_pack alloc packed"); int ok = ds4_gpu_tensor_write(&norm, 0, host_norm, norm_bytes) && ds4_gpu_tensor_write(&selected, 0, host_selected, selected_bytes) && ds4_gpu_tensor_write(&weights, 0, host_weights, weights_bytes) && ds4_gpu_moe_handoff_pack_tensor(&packed, &norm, &selected, &weights, n_embd, n_expert) && ds4_gpu_tensor_read(&packed, 0, host_packed, packed_bytes); CHECK(ok, "moe_pack compute"); CHECK(memcmp(host_packed, host_norm, norm_bytes) == 0, "moe_pack norm"); CHECK(memcmp(host_packed + norm_bytes, host_selected, selected_bytes) == 0, "moe_pack selected"); CHECK(memcmp(host_packed + norm_bytes + selected_bytes, host_weights, weights_bytes) == 0, "moe_pack weights"); ds4_gpu_tensor_free_in_place(&norm); ds4_gpu_tensor_free_in_place(&selected); ds4_gpu_tensor_free_in_place(&weights); ds4_gpu_tensor_free_in_place(&packed); free(host_norm); free(host_selected); free(host_weights); free(host_packed); ds4_gpu_cleanup(); fprintf(stderr, " moe_handoff_pack OK\n"); return 0; } /* Stress sub-test: repeated cross-device round-trips at realistic sizes. * Designed to catch the RTX 6000 Ada / driver bug where peer copies pass * a single small probe but corrupt data at larger sizes or on repeat. */ static int run_stress(int force_bounce) { int dev_count = 0; (void)cudaGetDeviceCount(&dev_count); if (dev_count < 2) { fprintf(stderr, " skipping stress (need >= 2 devices)\n"); return 0; } ds4_gpu_config cfg; memset(&cfg, 0, sizeof(cfg)); cfg.n_gpus = 2; cfg.device_indices[0] = 0; cfg.device_indices[1] = 1; if (force_bounce) setenv("DS4_FORCE_HOST_BOUNCE", "1", 1); else unsetenv("DS4_FORCE_HOST_BOUNCE"); CHECK(ds4_gpu_init_multi(&cfg), "stress init_multi"); const size_t sizes[] = { 256u * 1024u, 1u * 1024u * 1024u, 16u * 1024u * 1024u, }; const int n_sizes = (int)(sizeof(sizes) / sizeof(sizes[0])); const int iters = 32; const size_t max_n = sizes[n_sizes - 1]; unsigned char *host_src = (unsigned char *)malloc(max_n); unsigned char *host_dst = (unsigned char *)malloc(max_n); CHECK(host_src && host_dst, "stress host alloc"); ds4_gpu_tensor a; memset(&a, 0, sizeof(a)); ds4_gpu_tensor b; memset(&b, 0, sizeof(b)); CHECK(ds4_gpu_tensor_alloc_on(&a, 0, max_n) == 0, "stress alloc_on 0"); CHECK(ds4_gpu_tensor_alloc_on(&b, 1, max_n) == 0, "stress alloc_on 1"); int total_ok = 1; for (int s = 0; s < n_sizes && total_ok; s++) { size_t n = sizes[s]; for (int it = 0; it < iters && total_ok; it++) { for (size_t k = 0; k < n; k++) { host_src[k] = (unsigned char) ((k * 31u + (size_t)it * 17u + (size_t)s * 53u + 11u) & 0xffu); } memset(host_dst, 0, n); int io_ok = ds4_gpu_tensor_write(&a, 0, host_src, n) && ds4_gpu_tensor_copy_xdev(&b, &a, n) && ds4_gpu_tensor_read(&b, 0, host_dst, n); if (!io_ok) { fprintf(stderr, "FAIL: stress IO error size=%zu iter=%d force_bounce=%d\n", n, it, force_bounce); total_ok = 0; break; } for (size_t k = 0; k < n; k++) { if (host_src[k] != host_dst[k]) { fprintf(stderr, "FAIL: stress data mismatch size=%zu iter=%d offset=%zu" " src=0x%02x dst=0x%02x force_bounce=%d\n", n, it, k, host_src[k], host_dst[k], force_bounce); total_ok = 0; break; } } } } ds4_gpu_tensor_free_in_place(&a); ds4_gpu_tensor_free_in_place(&b); free(host_src); free(host_dst); ds4_gpu_cleanup(); if (total_ok) { fprintf(stderr, " stress OK (force_bounce=%d, %d iters x %d sizes, max %zu MiB)\n", force_bounce, iters, n_sizes, max_n / (1024u * 1024u)); } return total_ok ? 0 : 1; } static int run_add_xdev(int force_bounce) { int dev_count = 0; (void)cudaGetDeviceCount(&dev_count); if (dev_count < 2) { fprintf(stderr, " skipping add_xdev (need >= 2 devices)\n"); return 0; } ds4_gpu_config cfg; memset(&cfg, 0, sizeof(cfg)); cfg.n_gpus = 2; cfg.device_indices[0] = 0; cfg.device_indices[1] = 1; if (force_bounce) setenv("DS4_FORCE_HOST_BOUNCE", "1", 1); else unsetenv("DS4_FORCE_HOST_BOUNCE"); CHECK(ds4_gpu_init_multi(&cfg), "add_xdev init_multi"); const size_t N = 512 * 1024; float *host_local = (float *)malloc(N * sizeof(float)); float *host_remote = (float *)malloc(N * sizeof(float)); float *host_out = (float *)malloc(N * sizeof(float)); CHECK(host_local && host_remote && host_out, "add_xdev host alloc"); for (size_t i = 0; i < N; i++) { host_local[i] = (float)((int)(i % 1009) - 500) * 0.125f; host_remote[i] = (float)((int)(i % 997) - 400) * 0.25f; } ds4_gpu_tensor local; memset(&local, 0, sizeof(local)); ds4_gpu_tensor remote; memset(&remote, 0, sizeof(remote)); ds4_gpu_tensor staging; memset(&staging, 0, sizeof(staging)); ds4_gpu_tensor out; memset(&out, 0, sizeof(out)); CHECK(ds4_gpu_tensor_alloc_on(&local, 0, N * sizeof(float)) == 0, "add_xdev alloc local"); CHECK(ds4_gpu_tensor_alloc_on(&remote, 1, N * sizeof(float)) == 0, "add_xdev alloc remote"); CHECK(ds4_gpu_tensor_alloc_on(&staging, 0, N * sizeof(float)) == 0, "add_xdev alloc staging"); CHECK(ds4_gpu_tensor_alloc_on(&out, 0, N * sizeof(float)) == 0, "add_xdev alloc out"); int ok = ds4_gpu_tensor_write(&local, 0, host_local, N * sizeof(float)) && ds4_gpu_tensor_write(&remote, 0, host_remote, N * sizeof(float)) && ds4_gpu_add_xdev_tensor(&out, &local, &remote, &staging, (uint32_t)N) && ds4_gpu_tensor_read(&out, 0, host_out, N * sizeof(float)); CHECK(ok, "add_xdev IO"); for (size_t i = 0; i < N; i++) { float want = host_local[i] + host_remote[i]; if (host_out[i] != want) { fprintf(stderr, "FAIL: add_xdev mismatch i=%zu got=%f want=%f force_bounce=%d\n", i, host_out[i], want, force_bounce); return 1; } } ds4_gpu_tensor_free_in_place(&local); ds4_gpu_tensor_free_in_place(&remote); ds4_gpu_tensor_free_in_place(&staging); ds4_gpu_tensor_free_in_place(&out); free(host_local); free(host_remote); free(host_out); ds4_gpu_cleanup(); fprintf(stderr, " add_xdev OK (force_bounce=%d)\n", force_bounce); return 0; } static void pack_q8_identity_scale(unsigned char *w, uint64_t in_dim, uint64_t out_dim) { const uint64_t blocks = (in_dim + 31u) / 32u; for (uint64_t r = 0; r < out_dim; r++) { for (uint64_t b = 0; b < blocks; b++) { unsigned char *blk = w + (r * blocks + b) * 34u; /* IEEE-754 half representation of 1.0 on little-endian hosts. */ blk[0] = 0x00u; blk[1] = 0x3cu; for (uint64_t j = 0; j < 32u; j++) { uint64_t i = b * 32u + j; int v = i < in_dim ? (int)((r * 17u + i * 5u + 11u) % 23u) - 11 : 0; blk[2u + j] = (unsigned char)(int8_t)v; } } } } static int run_glm_decode_attention_staged(void) { int dev_count = 0; (void)cudaGetDeviceCount(&dev_count); if (dev_count < 1) return 0; ds4_gpu_config cfg; memset(&cfg, 0, sizeof(cfg)); cfg.n_gpus = 1; cfg.device_indices[0] = 0; CHECK(ds4_gpu_init_multi(&cfg), "glm staged attention init_multi"); const uint32_t n_head = 4u; const uint32_t kv_dim = 512u; const uint32_t qk_nope = 192u; const uint32_t qk_rope = 64u; const uint32_t q_dim = qk_nope + qk_rope; const uint32_t value_dim = 128u; const uint32_t cache_cap = 2053u; const uint32_t cases[] = {513u, 2048u}; const uint64_t value_rows = (uint64_t)n_head * value_dim; const uint64_t model_size = value_rows * (kv_dim / 32u) * 34u; const size_t q_count = (size_t)n_head * q_dim; const size_t low_count = (size_t)n_head * kv_dim; const size_t kv_count = (size_t)cache_cap * kv_dim; const size_t rope_count = (size_t)cache_cap * qk_rope; const size_t out_count = (size_t)n_head * value_dim; unsigned char *model = (unsigned char *)malloc((size_t)model_size); float *host_q = (float *)malloc(2u * q_count * sizeof(float)); float *host_low = (float *)malloc(2u * low_count * sizeof(float)); float *host_kv = (float *)malloc(kv_count * sizeof(float)); float *host_rope = (float *)malloc(rope_count * sizeof(float)); uint32_t *host_selected = (uint32_t *)malloc((size_t)cache_cap * sizeof(uint32_t)); float *host_fused = (float *)malloc(2u * out_count * sizeof(float)); float *host_staged = (float *)malloc(2u * out_count * sizeof(float)); CHECK(model && host_q && host_low && host_kv && host_rope && host_selected && host_fused && host_staged, "glm staged attention host alloc"); pack_q8_identity_scale(model, kv_dim, value_rows); for (size_t i = 0; i < 2u * q_count; i++) host_q[i] = (float)((int)(i % 19u) - 9) * 0.015625f; for (size_t i = 0; i < 2u * low_count; i++) host_low[i] = (float)((int)((i * 3u) % 23u) - 11) * 0.0078125f; for (size_t i = 0; i < kv_count; i++) host_kv[i] = (float)((int)((i * 5u) % 29u) - 14) * 0.00390625f; for (size_t i = 0; i < rope_count; i++) host_rope[i] = (float)((int)((i * 7u) % 17u) - 8) * 0.0078125f; for (uint32_t i = 0; i < cases[1]; i++) host_selected[i] = (i * 29u + 7u) % cache_cap; CHECK(ds4_gpu_set_model_map(model, model_size), "glm staged attention set model map"); ds4_gpu_tensor q; memset(&q, 0, sizeof(q)); ds4_gpu_tensor low; memset(&low, 0, sizeof(low)); ds4_gpu_tensor kv; memset(&kv, 0, sizeof(kv)); ds4_gpu_tensor rope; memset(&rope, 0, sizeof(rope)); ds4_gpu_tensor selected; memset(&selected, 0, sizeof(selected)); ds4_gpu_tensor fused; memset(&fused, 0, sizeof(fused)); ds4_gpu_tensor staged; memset(&staged, 0, sizeof(staged)); CHECK(ds4_gpu_tensor_alloc_on(&q, 0, 2u * q_count * sizeof(float)) == 0, "glm staged attention alloc q"); CHECK(ds4_gpu_tensor_alloc_on(&low, 0, 2u * low_count * sizeof(float)) == 0, "glm staged attention alloc low"); CHECK(ds4_gpu_tensor_alloc_on(&kv, 0, kv_count * sizeof(float)) == 0, "glm staged attention alloc kv"); CHECK(ds4_gpu_tensor_alloc_on(&rope, 0, rope_count * sizeof(float)) == 0, "glm staged attention alloc rope"); CHECK(ds4_gpu_tensor_alloc_on(&selected, 0, (uint64_t)cache_cap * sizeof(uint32_t)) == 0, "glm staged attention alloc selected"); CHECK(ds4_gpu_tensor_alloc_on(&fused, 0, 2u * out_count * sizeof(float)) == 0, "glm staged attention alloc fused"); CHECK(ds4_gpu_tensor_alloc_on(&staged, 0, 2u * out_count * sizeof(float)) == 0, "glm staged attention alloc staged"); int ok = ds4_gpu_tensor_write(&q, 0, host_q, 2u * q_count * sizeof(float)) && ds4_gpu_tensor_write(&low, 0, host_low, 2u * low_count * sizeof(float)) && ds4_gpu_tensor_write(&kv, 0, host_kv, kv_count * sizeof(float)) && ds4_gpu_tensor_write(&rope, 0, host_rope, rope_count * sizeof(float)) && ds4_gpu_tensor_write(&selected, 0, host_selected, (uint64_t)cases[1] * sizeof(uint32_t)); CHECK(ok, "glm staged attention upload"); for (uint32_t ci = 0; ci < 2u; ci++) { const uint32_t n_selected = cases[ci]; setenv("DS4_GLM_ATTN_NO_STAGED_DECODE", "1", 1); ok = ds4_gpu_glm_attention_indexed_decode_typed_tensor( &fused, &q, &low, &kv, &rope, model, model_size, 0u, 8u, &selected, n_selected, cache_cap, false, n_head, kv_dim, qk_nope, qk_rope, value_dim, 0u, 8000000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f) && ds4_gpu_tensor_read(&fused, 0, host_fused, out_count * sizeof(float)); CHECK(ok, "glm staged attention fused compute"); unsetenv("DS4_GLM_ATTN_NO_STAGED_DECODE"); ok = ds4_gpu_glm_attention_indexed_decode_typed_tensor( &staged, &q, &low, &kv, &rope, model, model_size, 0u, 8u, &selected, n_selected, cache_cap, false, n_head, kv_dim, qk_nope, qk_rope, value_dim, 0u, 8000000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f) && ds4_gpu_tensor_read(&staged, 0, host_staged, out_count * sizeof(float)); CHECK(ok, "glm staged attention staged compute"); if (memcmp(host_fused, host_staged, out_count * sizeof(float)) != 0) { fprintf(stderr, "FAIL: glm staged attention mismatch selected=%u\n", n_selected); return 1; } } const uint32_t tok2_cases[] = {255u, 513u}; for (uint32_t ci = 0; ci < 2u; ci++) { const uint32_t n_selected = tok2_cases[ci]; for (uint32_t i = 0; i <= n_selected; i++) host_selected[i] = i; CHECK(ds4_gpu_tensor_write(&selected, 0, host_selected, (uint64_t)(n_selected + 1u) * sizeof(uint32_t)), "glm tok2 attention selected upload"); ds4_gpu_tensor *q0 = ds4_gpu_tensor_view( &q, 0, q_count * sizeof(float)); ds4_gpu_tensor *q1 = ds4_gpu_tensor_view( &q, q_count * sizeof(float), q_count * sizeof(float)); ds4_gpu_tensor *low0 = ds4_gpu_tensor_view( &low, 0, low_count * sizeof(float)); ds4_gpu_tensor *low1 = ds4_gpu_tensor_view( &low, low_count * sizeof(float), low_count * sizeof(float)); ds4_gpu_tensor *ref0 = ds4_gpu_tensor_view( &fused, 0, out_count * sizeof(float)); ds4_gpu_tensor *ref1 = ds4_gpu_tensor_view( &fused, out_count * sizeof(float), out_count * sizeof(float)); CHECK(q0 && q1 && low0 && low1 && ref0 && ref1, "glm tok2 attention views"); ok = ds4_gpu_glm_attention_indexed_decode_typed_tensor( ref0, q0, low0, &kv, &rope, model, model_size, 0u, 8u, &selected, n_selected, cache_cap, false, n_head, kv_dim, qk_nope, qk_rope, value_dim, 0u, 8000000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f) && ds4_gpu_glm_attention_indexed_decode_typed_tensor( ref1, q1, low1, &kv, &rope, model, model_size, 0u, 8u, &selected, n_selected + 1u, cache_cap, false, n_head, kv_dim, qk_nope, qk_rope, value_dim, 0u, 8000000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f); CHECK(ok, "glm tok2 attention reference compute"); ds4_gpu_set_glm_mtp_verify_mode(true); ok = ds4_gpu_glm_attention_indexed_decode_typed_tensor( &staged, &q, &low, &kv, &rope, model, model_size, 0u, 8u, &selected, n_selected, cache_cap, false, n_head, kv_dim, qk_nope, qk_rope, value_dim, 0u, 8000000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f); ds4_gpu_set_glm_mtp_verify_mode(false); ok = ok && ds4_gpu_tensor_read(&fused, 0, host_fused, 2u * out_count * sizeof(float)) && ds4_gpu_tensor_read(&staged, 0, host_staged, 2u * out_count * sizeof(float)); ds4_gpu_tensor_free(q0); ds4_gpu_tensor_free(q1); ds4_gpu_tensor_free(low0); ds4_gpu_tensor_free(low1); ds4_gpu_tensor_free(ref0); ds4_gpu_tensor_free(ref1); CHECK(ok, "glm tok2 attention batched compute"); CHECK(memcmp(host_fused, host_staged, 2u * out_count * sizeof(float)) == 0, "glm tok2 attention exact output"); } ds4_gpu_tensor_free_in_place(&q); ds4_gpu_tensor_free_in_place(&low); ds4_gpu_tensor_free_in_place(&kv); ds4_gpu_tensor_free_in_place(&rope); ds4_gpu_tensor_free_in_place(&selected); ds4_gpu_tensor_free_in_place(&fused); ds4_gpu_tensor_free_in_place(&staged); free(model); free(host_q); free(host_low); free(host_kv); free(host_rope); free(host_selected); free(host_fused); free(host_staged); ds4_gpu_cleanup(); fprintf(stderr, " glm_decode_attention_staged OK\n"); return 0; } static void pack_f16_small_mask(uint16_t *w, uint64_t in_dim, uint64_t out_dim) { for (uint64_t r = 0; r < out_dim; r++) { for (uint64_t i = 0; i < in_dim; i++) { int on = ((r * 13u + i * 7u + 5u) % 11u) < 3u; w[r * in_dim + i] = on ? 0x3c00u : 0x0000u; } } } static int run_q8_kslice(void) { int dev_count = 0; (void)cudaGetDeviceCount(&dev_count); if (dev_count < 1) { fprintf(stderr, " skipping q8_kslice (need CUDA device)\n"); return 0; } ds4_gpu_config cfg; memset(&cfg, 0, sizeof(cfg)); cfg.n_gpus = 1; cfg.device_indices[0] = 0; CHECK(ds4_gpu_init_multi(&cfg), "q8_kslice init_multi"); const uint64_t in_dim = 128; const uint64_t out_dim = 96; const uint64_t split = 64; const uint64_t blocks = (in_dim + 31u) / 32u; const uint64_t model_size = out_dim * blocks * 34u; unsigned char *model = (unsigned char *)malloc((size_t)model_size); float *host_x = (float *)malloc((size_t)in_dim * sizeof(float)); float *host_full = (float *)malloc((size_t)out_dim * sizeof(float)); float *host_sum = (float *)malloc((size_t)out_dim * sizeof(float)); CHECK(model && host_x && host_full && host_sum, "q8_kslice host alloc"); pack_q8_identity_scale(model, in_dim, out_dim); for (uint64_t i = 0; i < in_dim; i++) { host_x[i] = (float)((int)(i % 37u) - 18) * 0.03125f; } CHECK(ds4_gpu_set_model_map(model, model_size), "q8_kslice set model map"); ds4_gpu_tensor x; memset(&x, 0, sizeof(x)); ds4_gpu_tensor full; memset(&full, 0, sizeof(full)); ds4_gpu_tensor p0; memset(&p0, 0, sizeof(p0)); ds4_gpu_tensor p1; memset(&p1, 0, sizeof(p1)); ds4_gpu_tensor sum; memset(&sum, 0, sizeof(sum)); CHECK(ds4_gpu_tensor_alloc_on(&x, 0, in_dim * sizeof(float)) == 0, "q8_kslice alloc x"); CHECK(ds4_gpu_tensor_alloc_on(&full, 0, out_dim * sizeof(float)) == 0, "q8_kslice alloc full"); CHECK(ds4_gpu_tensor_alloc_on(&p0, 0, out_dim * sizeof(float)) == 0, "q8_kslice alloc p0"); CHECK(ds4_gpu_tensor_alloc_on(&p1, 0, out_dim * sizeof(float)) == 0, "q8_kslice alloc p1"); CHECK(ds4_gpu_tensor_alloc_on(&sum, 0, out_dim * sizeof(float)) == 0, "q8_kslice alloc sum"); CHECK(ds4_gpu_tensor_write(&x, 0, host_x, in_dim * sizeof(float)), "q8_kslice write x"); ds4_gpu_tensor *x0 = ds4_gpu_tensor_view(&x, 0, split * sizeof(float)); ds4_gpu_tensor *x1 = ds4_gpu_tensor_view(&x, split * sizeof(float), (in_dim - split) * sizeof(float)); CHECK(x0 && x1, "q8_kslice views"); int ok = ds4_gpu_matmul_q8_0_tensor(&full, model, model_size, 0, in_dim, out_dim, &x, 1) && ds4_gpu_matmul_q8_0_kslice_rows_tensor(&p0, model, model_size, 0, in_dim, out_dim, 0, split, x0, 1) && ds4_gpu_matmul_q8_0_kslice_rows_tensor(&p1, model, model_size, 0, in_dim, out_dim, split, in_dim - split, x1, 1) && ds4_gpu_add_tensor(&sum, &p0, &p1, (uint32_t)out_dim) && ds4_gpu_tensor_read(&full, 0, host_full, out_dim * sizeof(float)) && ds4_gpu_tensor_read(&sum, 0, host_sum, out_dim * sizeof(float)); CHECK(ok, "q8_kslice compute"); for (uint64_t i = 0; i < out_dim; i++) { float diff = fabsf(host_full[i] - host_sum[i]); if (diff > 1.0e-4f) { fprintf(stderr, "FAIL: q8_kslice mismatch row=%llu full=%f sum=%f diff=%g\n", (unsigned long long)i, host_full[i], host_sum[i], diff); return 1; } } ds4_gpu_tensor_free(x0); ds4_gpu_tensor_free(x1); ds4_gpu_tensor_free_in_place(&x); ds4_gpu_tensor_free_in_place(&full); ds4_gpu_tensor_free_in_place(&p0); ds4_gpu_tensor_free_in_place(&p1); ds4_gpu_tensor_free_in_place(&sum); free(model); free(host_x); free(host_full); free(host_sum); ds4_gpu_cleanup(); fprintf(stderr, " q8_kslice OK\n"); return 0; } static int run_q8_matmul_top1_fused(void) { int dev_count = 0; (void)cudaGetDeviceCount(&dev_count); if (dev_count < 1) { fprintf(stderr, " skipping q8_matmul_top1_fused (need CUDA device)\n"); return 0; } ds4_gpu_config cfg; memset(&cfg, 0, sizeof(cfg)); cfg.n_gpus = 1; cfg.device_indices[0] = 0; CHECK(ds4_gpu_init_multi(&cfg), "q8_top1_fused init_multi"); const uint64_t in_dim = 160; const uint64_t out_dim = 257; const uint32_t index_offset = 1234u; const uint64_t blocks = (in_dim + 31u) / 32u; const uint64_t model_size = out_dim * blocks * 34u; unsigned char *model = (unsigned char *)malloc((size_t)model_size); float *host_x = (float *)malloc((size_t)in_dim * sizeof(float)); uint32_t ref_id = 0, fused_id = 0; float ref_value = 0.0f, fused_value = 0.0f; CHECK(model && host_x, "q8_top1_fused host alloc"); pack_q8_identity_scale(model, in_dim, out_dim); for (uint64_t i = 0; i < in_dim; i++) { host_x[i] = (float)((int)(i % 43u) - 21) * 0.0234375f; } CHECK(ds4_gpu_set_model_map(model, model_size), "q8_top1_fused set model map"); ds4_gpu_tensor x; memset(&x, 0, sizeof(x)); ds4_gpu_tensor logits; memset(&logits, 0, sizeof(logits)); ds4_gpu_tensor ref_selected; memset(&ref_selected, 0, sizeof(ref_selected)); ds4_gpu_tensor ref_values; memset(&ref_values, 0, sizeof(ref_values)); ds4_gpu_tensor fused_selected; memset(&fused_selected, 0, sizeof(fused_selected)); ds4_gpu_tensor fused_values; memset(&fused_values, 0, sizeof(fused_values)); CHECK(ds4_gpu_tensor_alloc_on(&x, 0, in_dim * sizeof(float)) == 0, "q8_top1_fused alloc x"); CHECK(ds4_gpu_tensor_alloc_on(&logits, 0, out_dim * sizeof(float)) == 0, "q8_top1_fused alloc logits"); CHECK(ds4_gpu_tensor_alloc_on(&ref_selected, 0, sizeof(uint32_t)) == 0, "q8_top1_fused alloc ref selected"); CHECK(ds4_gpu_tensor_alloc_on(&ref_values, 0, sizeof(float)) == 0, "q8_top1_fused alloc ref values"); CHECK(ds4_gpu_tensor_alloc_on(&fused_selected, 0, sizeof(uint32_t)) == 0, "q8_top1_fused alloc fused selected"); CHECK(ds4_gpu_tensor_alloc_on(&fused_values, 0, sizeof(float)) == 0, "q8_top1_fused alloc fused values"); CHECK(ds4_gpu_tensor_write(&x, 0, host_x, in_dim * sizeof(float)), "q8_top1_fused write x"); int ok = ds4_gpu_matmul_q8_0_tensor(&logits, model, model_size, 0, in_dim, out_dim, &x, 1) && ds4_gpu_indexer_top1_value_tensor(&ref_selected, &ref_values, &logits, (uint32_t)out_dim, 1, index_offset) && ds4_gpu_matmul_q8_0_top1_tensor(&fused_selected, &fused_values, model, model_size, 0, in_dim, out_dim, &x, index_offset) && ds4_gpu_tensor_read(&ref_selected, 0, &ref_id, sizeof(ref_id)) && ds4_gpu_tensor_read(&fused_selected, 0, &fused_id, sizeof(fused_id)) && ds4_gpu_tensor_read(&ref_values, 0, &ref_value, sizeof(ref_value)) && ds4_gpu_tensor_read(&fused_values, 0, &fused_value, sizeof(fused_value)); CHECK(ok, "q8_top1_fused compute"); CHECK(ref_id == fused_id, "q8_top1_fused selected"); if (fabsf(ref_value - fused_value) > 1.0e-6f) { fprintf(stderr, "FAIL: q8_top1_fused value mismatch id=%u ref=%f fused=%f diff=%g\n", ref_id, ref_value, fused_value, fabsf(ref_value - fused_value)); return 1; } ds4_gpu_tensor_free_in_place(&x); ds4_gpu_tensor_free_in_place(&logits); ds4_gpu_tensor_free_in_place(&ref_selected); ds4_gpu_tensor_free_in_place(&ref_values); ds4_gpu_tensor_free_in_place(&fused_selected); ds4_gpu_tensor_free_in_place(&fused_values); free(model); free(host_x); ds4_gpu_cleanup(); fprintf(stderr, " q8_matmul_top1_fused OK\n"); return 0; } static int run_f16_small_out(void) { int dev_count = 0; (void)cudaGetDeviceCount(&dev_count); if (dev_count < 1) { fprintf(stderr, " skipping f16_small_out (need CUDA device)\n"); return 0; } setenv("DS4_CUDA_F16_SMALL_OUT", "1", 1); unsetenv("DS4_CUDA_NO_F16_SMALL_OUT"); ds4_gpu_config cfg; memset(&cfg, 0, sizeof(cfg)); cfg.n_gpus = 1; cfg.device_indices[0] = 0; CHECK(ds4_gpu_init_multi(&cfg), "f16_small_out init_multi"); const uint64_t in_dim = 32768; const uint64_t out_dim = 16; const uint64_t model_size = in_dim * out_dim * sizeof(uint16_t); uint16_t *model = (uint16_t *)malloc((size_t)model_size); float *host_x = (float *)malloc((size_t)in_dim * sizeof(float)); float *host_out = (float *)malloc((size_t)out_dim * sizeof(float)); float *host_ref = (float *)calloc((size_t)out_dim, sizeof(float)); CHECK(model && host_x && host_out && host_ref, "f16_small_out host alloc"); pack_f16_small_mask(model, in_dim, out_dim); for (uint64_t i = 0; i < in_dim; i++) { host_x[i] = (float)((int)(i % 31u) - 15) * 0.03125f; } for (uint64_t r = 0; r < out_dim; r++) { float sum = 0.0f; for (uint64_t i = 0; i < in_dim; i++) { if (model[r * in_dim + i] == 0x3c00u) sum += host_x[i]; } host_ref[r] = sum; } CHECK(ds4_gpu_set_model_map(model, model_size), "f16_small_out set model map"); ds4_gpu_tensor x; memset(&x, 0, sizeof(x)); ds4_gpu_tensor out; memset(&out, 0, sizeof(out)); CHECK(ds4_gpu_tensor_alloc_on(&x, 0, in_dim * sizeof(float)) == 0, "f16_small_out alloc x"); CHECK(ds4_gpu_tensor_alloc_on(&out, 0, out_dim * sizeof(float)) == 0, "f16_small_out alloc out"); CHECK(ds4_gpu_tensor_write(&x, 0, host_x, in_dim * sizeof(float)), "f16_small_out write x"); int ok = ds4_gpu_matmul_f16_tensor(&out, model, model_size, 0, in_dim, out_dim, &x, 1) && ds4_gpu_tensor_read(&out, 0, host_out, out_dim * sizeof(float)); CHECK(ok, "f16_small_out compute"); for (uint64_t r = 0; r < out_dim; r++) { float diff = fabsf(host_out[r] - host_ref[r]); if (diff > 1.0e-3f) { fprintf(stderr, "FAIL: f16_small_out mismatch row=%llu got=%f want=%f diff=%g\n", (unsigned long long)r, host_out[r], host_ref[r], diff); return 1; } } ds4_gpu_tensor_free_in_place(&x); ds4_gpu_tensor_free_in_place(&out); free(model); free(host_x); free(host_out); free(host_ref); ds4_gpu_cleanup(); fprintf(stderr, " f16_small_out OK\n"); return 0; } static int run_f16_small_batch(void) { int dev_count = 0; (void)cudaGetDeviceCount(&dev_count); if (dev_count < 1) { fprintf(stderr, " skipping f16_small_batch (need CUDA device)\n"); return 0; } setenv("DS4_CUDA_F16_SMALL_BATCH", "1", 1); unsetenv("DS4_CUDA_NO_F16_SMALL_BATCH"); unsetenv("DS4_CUDA_NO_F16_CUBLAS_BATCH"); ds4_gpu_config cfg; memset(&cfg, 0, sizeof(cfg)); cfg.n_gpus = 1; cfg.device_indices[0] = 0; CHECK(ds4_gpu_init_multi(&cfg), "f16_small_batch init_multi"); const uint64_t in_dim = 8192; const uint64_t out_dim = 24; const uint64_t n_tok = 9; const uint64_t model_size = in_dim * out_dim * sizeof(uint16_t); uint16_t *model = (uint16_t *)malloc((size_t)model_size); float *host_x = (float *)malloc((size_t)(n_tok * in_dim) * sizeof(float)); float *host_fast = (float *)malloc((size_t)(n_tok * out_dim) * sizeof(float)); float *host_ref = (float *)malloc((size_t)(n_tok * out_dim) * sizeof(float)); CHECK(model && host_x && host_fast && host_ref, "f16_small_batch host alloc"); const uint16_t vals[5] = {0x0000u, 0x3c00u, 0xbc00u, 0x3800u, 0xb800u}; for (uint64_t r = 0; r < out_dim; r++) { for (uint64_t i = 0; i < in_dim; i++) { model[r * in_dim + i] = vals[(r * 17u + i * 5u + 3u) % 5u]; } } for (uint64_t t = 0; t < n_tok; t++) { for (uint64_t i = 0; i < in_dim; i++) { host_x[t * in_dim + i] = (float)((int)((t * 19u + i * 7u) % 53u) - 26) * 0.015625f; } } CHECK(ds4_gpu_set_model_map(model, model_size), "f16_small_batch set model map"); ds4_gpu_tensor x; memset(&x, 0, sizeof(x)); ds4_gpu_tensor fast; memset(&fast, 0, sizeof(fast)); ds4_gpu_tensor ref; memset(&ref, 0, sizeof(ref)); CHECK(ds4_gpu_tensor_alloc_on(&x, 0, n_tok * in_dim * sizeof(float)) == 0, "f16_small_batch alloc x"); CHECK(ds4_gpu_tensor_alloc_on(&fast, 0, n_tok * out_dim * sizeof(float)) == 0, "f16_small_batch alloc fast"); CHECK(ds4_gpu_tensor_alloc_on(&ref, 0, n_tok * out_dim * sizeof(float)) == 0, "f16_small_batch alloc ref"); CHECK(ds4_gpu_tensor_write(&x, 0, host_x, n_tok * in_dim * sizeof(float)), "f16_small_batch write x"); int ok = ds4_gpu_matmul_f16_tensor(&fast, model, model_size, 0, in_dim, out_dim, &x, n_tok); setenv("DS4_CUDA_NO_F16_SMALL_BATCH", "1", 1); setenv("DS4_CUDA_NO_F16_CUBLAS_BATCH", "1", 1); ok = ok && ds4_gpu_matmul_f16_tensor(&ref, model, model_size, 0, in_dim, out_dim, &x, n_tok) && ds4_gpu_tensor_read(&fast, 0, host_fast, n_tok * out_dim * sizeof(float)) && ds4_gpu_tensor_read(&ref, 0, host_ref, n_tok * out_dim * sizeof(float)); CHECK(ok, "f16_small_batch compute"); for (uint64_t i = 0; i < n_tok * out_dim; i++) { if (memcmp(&host_fast[i], &host_ref[i], sizeof(float)) != 0) { fprintf(stderr, "FAIL: f16_small_batch mismatch i=%llu got=%f want=%f diff=%g\n", (unsigned long long)i, host_fast[i], host_ref[i], fabsf(host_fast[i] - host_ref[i])); return 1; } } ds4_gpu_tensor_free_in_place(&x); ds4_gpu_tensor_free_in_place(&fast); ds4_gpu_tensor_free_in_place(&ref); free(model); free(host_x); free(host_fast); free(host_ref); ds4_gpu_cleanup(); fprintf(stderr, " f16_small_batch OK\n"); return 0; } static int run_f16_pair_batch(void) { int dev_count = 0; (void)cudaGetDeviceCount(&dev_count); if (dev_count < 1) { fprintf(stderr, " skipping f16_pair_batch (need CUDA device)\n"); return 0; } unsetenv("DS4_CUDA_NO_F16_PAIR_MATMUL"); unsetenv("DS4_CUDA_NO_F16_CUBLAS_BATCH"); unsetenv("DS4_CUDA_F16_SMALL_BATCH"); unsetenv("DS4_CUDA_NO_F16_SMALL_BATCH"); ds4_gpu_config cfg; memset(&cfg, 0, sizeof(cfg)); cfg.n_gpus = 1; cfg.device_indices[0] = 0; CHECK(ds4_gpu_init_multi(&cfg), "f16_pair_batch init_multi"); const uint64_t in_dim = 512; const uint64_t out_dim = 80; const uint64_t n_tok = 11; const uint64_t weight_bytes = in_dim * out_dim * sizeof(uint16_t); const uint64_t model_size = weight_bytes * 2u; uint16_t *model = (uint16_t *)malloc((size_t)model_size); float *host_x = (float *)malloc((size_t)(n_tok * in_dim) * sizeof(float)); float *fast0 = (float *)malloc((size_t)(n_tok * out_dim) * sizeof(float)); float *fast1 = (float *)malloc((size_t)(n_tok * out_dim) * sizeof(float)); float *ref0 = (float *)malloc((size_t)(n_tok * out_dim) * sizeof(float)); float *ref1 = (float *)malloc((size_t)(n_tok * out_dim) * sizeof(float)); CHECK(model && host_x && fast0 && fast1 && ref0 && ref1, "f16_pair_batch host alloc"); const uint16_t vals[7] = {0x0000u, 0x3c00u, 0xbc00u, 0x3800u, 0xb800u, 0x3400u, 0xb400u}; for (uint64_t r = 0; r < out_dim; r++) { for (uint64_t i = 0; i < in_dim; i++) { model[r * in_dim + i] = vals[(r * 11u + i * 3u + 1u) % 7u]; model[(weight_bytes / sizeof(uint16_t)) + r * in_dim + i] = vals[(r * 13u + i * 5u + 4u) % 7u]; } } for (uint64_t t = 0; t < n_tok; t++) { for (uint64_t i = 0; i < in_dim; i++) { host_x[t * in_dim + i] = (float)((int)((t * 23u + i * 17u) % 71u) - 35) * 0.0078125f; } } CHECK(ds4_gpu_set_model_map(model, model_size), "f16_pair_batch set model map"); ds4_gpu_tensor x; memset(&x, 0, sizeof(x)); ds4_gpu_tensor out_fast0; memset(&out_fast0, 0, sizeof(out_fast0)); ds4_gpu_tensor out_fast1; memset(&out_fast1, 0, sizeof(out_fast1)); ds4_gpu_tensor out_ref0; memset(&out_ref0, 0, sizeof(out_ref0)); ds4_gpu_tensor out_ref1; memset(&out_ref1, 0, sizeof(out_ref1)); CHECK(ds4_gpu_tensor_alloc_on(&x, 0, n_tok * in_dim * sizeof(float)) == 0, "f16_pair_batch alloc x"); CHECK(ds4_gpu_tensor_alloc_on(&out_fast0, 0, n_tok * out_dim * sizeof(float)) == 0, "f16_pair_batch alloc fast0"); CHECK(ds4_gpu_tensor_alloc_on(&out_fast1, 0, n_tok * out_dim * sizeof(float)) == 0, "f16_pair_batch alloc fast1"); CHECK(ds4_gpu_tensor_alloc_on(&out_ref0, 0, n_tok * out_dim * sizeof(float)) == 0, "f16_pair_batch alloc ref0"); CHECK(ds4_gpu_tensor_alloc_on(&out_ref1, 0, n_tok * out_dim * sizeof(float)) == 0, "f16_pair_batch alloc ref1"); CHECK(ds4_gpu_tensor_write(&x, 0, host_x, n_tok * in_dim * sizeof(float)), "f16_pair_batch write x"); int ok = ds4_gpu_matmul_f16_pair_tensor(&out_fast0, &out_fast1, model, model_size, 0, weight_bytes, in_dim, out_dim, &x, n_tok); setenv("DS4_CUDA_NO_F16_PAIR_MATMUL", "1", 1); ok = ok && ds4_gpu_matmul_f16_pair_tensor(&out_ref0, &out_ref1, model, model_size, 0, weight_bytes, in_dim, out_dim, &x, n_tok) && ds4_gpu_tensor_read(&out_fast0, 0, fast0, n_tok * out_dim * sizeof(float)) && ds4_gpu_tensor_read(&out_fast1, 0, fast1, n_tok * out_dim * sizeof(float)) && ds4_gpu_tensor_read(&out_ref0, 0, ref0, n_tok * out_dim * sizeof(float)) && ds4_gpu_tensor_read(&out_ref1, 0, ref1, n_tok * out_dim * sizeof(float)); CHECK(ok, "f16_pair_batch compute"); for (uint64_t i = 0; i < n_tok * out_dim; i++) { if (memcmp(&fast0[i], &ref0[i], sizeof(float)) != 0 || memcmp(&fast1[i], &ref1[i], sizeof(float)) != 0) { fprintf(stderr, "FAIL: f16_pair_batch mismatch i=%llu fast0=%f ref0=%f fast1=%f ref1=%f\n", (unsigned long long)i, fast0[i], ref0[i], fast1[i], ref1[i]); return 1; } } ds4_gpu_tensor_free_in_place(&x); ds4_gpu_tensor_free_in_place(&out_fast0); ds4_gpu_tensor_free_in_place(&out_fast1); ds4_gpu_tensor_free_in_place(&out_ref0); ds4_gpu_tensor_free_in_place(&out_ref1); free(model); free(host_x); free(fast0); free(fast1); free(ref0); free(ref1); ds4_gpu_cleanup(); fprintf(stderr, " f16_pair_batch OK\n"); return 0; } static int run_q8_rowsplit(void) { int dev_count = 0; (void)cudaGetDeviceCount(&dev_count); if (dev_count < 1) { fprintf(stderr, " skipping q8_rowsplit (need CUDA device)\n"); return 0; } ds4_gpu_config cfg; memset(&cfg, 0, sizeof(cfg)); cfg.n_gpus = 1; cfg.device_indices[0] = 0; CHECK(ds4_gpu_init_multi(&cfg), "q8_rowsplit init_multi"); const uint64_t in_dim = 96; const uint64_t out_dim = 97; const uint64_t split = 41; const uint64_t blocks = (in_dim + 31u) / 32u; const uint64_t row_bytes = blocks * 34u; const uint64_t model_size = out_dim * row_bytes; unsigned char *model = (unsigned char *)malloc((size_t)model_size); float *host_x = (float *)malloc((size_t)in_dim * sizeof(float)); float *host_full = (float *)malloc((size_t)out_dim * sizeof(float)); float *host_split = (float *)malloc((size_t)out_dim * sizeof(float)); CHECK(model && host_x && host_full && host_split, "q8_rowsplit host alloc"); pack_q8_identity_scale(model, in_dim, out_dim); for (uint64_t i = 0; i < in_dim; i++) { host_x[i] = (float)((int)(i % 43u) - 21) * 0.015625f; } CHECK(ds4_gpu_set_model_map(model, model_size), "q8_rowsplit set model map"); ds4_gpu_tensor x; memset(&x, 0, sizeof(x)); ds4_gpu_tensor full; memset(&full, 0, sizeof(full)); ds4_gpu_tensor split_out; memset(&split_out, 0, sizeof(split_out)); CHECK(ds4_gpu_tensor_alloc_on(&x, 0, in_dim * sizeof(float)) == 0, "q8_rowsplit alloc x"); CHECK(ds4_gpu_tensor_alloc_on(&full, 0, out_dim * sizeof(float)) == 0, "q8_rowsplit alloc full"); CHECK(ds4_gpu_tensor_alloc_on(&split_out, 0, out_dim * sizeof(float)) == 0, "q8_rowsplit alloc split"); CHECK(ds4_gpu_tensor_write(&x, 0, host_x, in_dim * sizeof(float)), "q8_rowsplit write x"); ds4_gpu_tensor *split0 = ds4_gpu_tensor_view(&split_out, 0, split * sizeof(float)); ds4_gpu_tensor *split1 = ds4_gpu_tensor_view(&split_out, split * sizeof(float), (out_dim - split) * sizeof(float)); CHECK(split0 && split1, "q8_rowsplit views"); int ok = ds4_gpu_matmul_q8_0_tensor(&full, model, model_size, 0, in_dim, out_dim, &x, 1) && ds4_gpu_matmul_q8_0_tensor(split0, model, model_size, 0, in_dim, split, &x, 1) && ds4_gpu_matmul_q8_0_tensor(split1, model, model_size, split * row_bytes, in_dim, out_dim - split, &x, 1) && ds4_gpu_tensor_read(&full, 0, host_full, out_dim * sizeof(float)) && ds4_gpu_tensor_read(&split_out, 0, host_split, out_dim * sizeof(float)); CHECK(ok, "q8_rowsplit compute"); for (uint64_t i = 0; i < out_dim; i++) { float diff = fabsf(host_full[i] - host_split[i]); if (diff > 1.0e-6f) { fprintf(stderr, "FAIL: q8_rowsplit mismatch row=%llu full=%f split=%f diff=%g\n", (unsigned long long)i, host_full[i], host_split[i], diff); return 1; } } ds4_gpu_tensor_free(split0); ds4_gpu_tensor_free(split1); ds4_gpu_tensor_free_in_place(&x); ds4_gpu_tensor_free_in_place(&full); ds4_gpu_tensor_free_in_place(&split_out); free(model); free(host_x); free(host_full); free(host_split); ds4_gpu_cleanup(); fprintf(stderr, " q8_rowsplit OK\n"); return 0; } static int run_q8_pair_batch(void) { int dev_count = 0; (void)cudaGetDeviceCount(&dev_count); if (dev_count < 1) { fprintf(stderr, " skipping q8_pair_batch (need CUDA device)\n"); return 0; } setenv("DS4_CUDA_Q8_PAIR_BATCH", "1", 1); ds4_gpu_config cfg; memset(&cfg, 0, sizeof(cfg)); cfg.n_gpus = 1; cfg.device_indices[0] = 0; CHECK(ds4_gpu_init_multi(&cfg), "q8_pair_batch init_multi"); const uint64_t in_dim = 96; const uint64_t out0_dim = 40; const uint64_t out1_dim = 33; const uint64_t n_tok = 9; const uint64_t blocks = (in_dim + 31u) / 32u; const uint64_t row_bytes = blocks * 34u; const uint64_t off1 = out0_dim * row_bytes; const uint64_t model_size = off1 + out1_dim * row_bytes; unsigned char *model = (unsigned char *)malloc((size_t)model_size); float *host_x = (float *)malloc((size_t)(n_tok * in_dim) * sizeof(float)); float *host_sep0 = (float *)malloc((size_t)(n_tok * out0_dim) * sizeof(float)); float *host_sep1 = (float *)malloc((size_t)(n_tok * out1_dim) * sizeof(float)); float *host_pair0 = (float *)malloc((size_t)(n_tok * out0_dim) * sizeof(float)); float *host_pair1 = (float *)malloc((size_t)(n_tok * out1_dim) * sizeof(float)); CHECK(model && host_x && host_sep0 && host_sep1 && host_pair0 && host_pair1, "q8_pair_batch host alloc"); pack_q8_identity_scale(model, in_dim, out0_dim); pack_q8_identity_scale(model + off1, in_dim, out1_dim); for (uint64_t i = 0; i < n_tok * in_dim; i++) { host_x[i] = (float)((int)(i % 47u) - 23) * 0.03125f; } CHECK(ds4_gpu_set_model_map(model, model_size), "q8_pair_batch set model map"); ds4_gpu_tensor x; memset(&x, 0, sizeof(x)); ds4_gpu_tensor sep0; memset(&sep0, 0, sizeof(sep0)); ds4_gpu_tensor sep1; memset(&sep1, 0, sizeof(sep1)); ds4_gpu_tensor pair0; memset(&pair0, 0, sizeof(pair0)); ds4_gpu_tensor pair1; memset(&pair1, 0, sizeof(pair1)); CHECK(ds4_gpu_tensor_alloc_on(&x, 0, n_tok * in_dim * sizeof(float)) == 0, "q8_pair_batch alloc x"); CHECK(ds4_gpu_tensor_alloc_on(&sep0, 0, n_tok * out0_dim * sizeof(float)) == 0, "q8_pair_batch alloc sep0"); CHECK(ds4_gpu_tensor_alloc_on(&sep1, 0, n_tok * out1_dim * sizeof(float)) == 0, "q8_pair_batch alloc sep1"); CHECK(ds4_gpu_tensor_alloc_on(&pair0, 0, n_tok * out0_dim * sizeof(float)) == 0, "q8_pair_batch alloc pair0"); CHECK(ds4_gpu_tensor_alloc_on(&pair1, 0, n_tok * out1_dim * sizeof(float)) == 0, "q8_pair_batch alloc pair1"); CHECK(ds4_gpu_tensor_write(&x, 0, host_x, n_tok * in_dim * sizeof(float)), "q8_pair_batch write x"); int ok = ds4_gpu_matmul_q8_0_tensor(&sep0, model, model_size, 0, in_dim, out0_dim, &x, n_tok) && ds4_gpu_matmul_q8_0_tensor(&sep1, model, model_size, off1, in_dim, out1_dim, &x, n_tok) && ds4_gpu_matmul_q8_0_pair_tensor(&pair0, &pair1, model, model_size, 0, off1, in_dim, out0_dim, out1_dim, &x, n_tok) && ds4_gpu_tensor_read(&sep0, 0, host_sep0, n_tok * out0_dim * sizeof(float)) && ds4_gpu_tensor_read(&sep1, 0, host_sep1, n_tok * out1_dim * sizeof(float)) && ds4_gpu_tensor_read(&pair0, 0, host_pair0, n_tok * out0_dim * sizeof(float)) && ds4_gpu_tensor_read(&pair1, 0, host_pair1, n_tok * out1_dim * sizeof(float)); CHECK(ok, "q8_pair_batch compute"); CHECK(memcmp(host_sep0, host_pair0, (size_t)(n_tok * out0_dim) * sizeof(float)) == 0, "q8_pair_batch out0"); CHECK(memcmp(host_sep1, host_pair1, (size_t)(n_tok * out1_dim) * sizeof(float)) == 0, "q8_pair_batch out1"); ds4_gpu_tensor_free_in_place(&x); ds4_gpu_tensor_free_in_place(&sep0); ds4_gpu_tensor_free_in_place(&sep1); ds4_gpu_tensor_free_in_place(&pair0); ds4_gpu_tensor_free_in_place(&pair1); free(model); free(host_x); free(host_sep0); free(host_sep1); free(host_pair0); free(host_pair1); ds4_gpu_cleanup(); unsetenv("DS4_CUDA_Q8_PAIR_BATCH"); fprintf(stderr, " q8_pair_batch OK\n"); return 0; } static int run_attention_output_tp(void) { int dev_count = 0; (void)cudaGetDeviceCount(&dev_count); if (dev_count < 1) { fprintf(stderr, " skipping attention_output_tp (need CUDA device)\n"); return 0; } ds4_gpu_config cfg; memset(&cfg, 0, sizeof(cfg)); cfg.n_gpus = 1; cfg.device_indices[0] = 0; CHECK(ds4_gpu_init_multi(&cfg), "attention_tp init_multi"); const uint64_t group_dim = 64; const uint64_t rank = 64; const uint32_t n_groups = 4; const uint32_t group_cnt = 2; const uint64_t out_dim = 48; const uint64_t low_dim = (uint64_t)n_groups * rank; const uint64_t a_blocks = (group_dim + 31u) / 32u; const uint64_t b_blocks = (low_dim + 31u) / 32u; const uint64_t a_bytes = (uint64_t)n_groups * rank * a_blocks * 34u; const uint64_t b_off = a_bytes; const uint64_t b_bytes = out_dim * b_blocks * 34u; const uint64_t model_size = a_bytes + b_bytes; unsigned char *model = (unsigned char *)malloc((size_t)model_size); float *host_heads = (float *)malloc((size_t)n_groups * group_dim * sizeof(float)); float *host_full = (float *)malloc((size_t)out_dim * sizeof(float)); float *host_sum = (float *)malloc((size_t)out_dim * sizeof(float)); CHECK(model && host_heads && host_full && host_sum, "attention_tp host alloc"); pack_q8_identity_scale(model, group_dim, (uint64_t)n_groups * rank); pack_q8_identity_scale(model + b_off, low_dim, out_dim); for (uint64_t i = 0; i < (uint64_t)n_groups * group_dim; i++) { host_heads[i] = (float)((int)(i % 29u) - 14) * 0.0625f; } CHECK(ds4_gpu_set_model_map(model, model_size), "attention_tp set model map"); ds4_gpu_tensor heads; memset(&heads, 0, sizeof(heads)); ds4_gpu_tensor low; memset(&low, 0, sizeof(low)); ds4_gpu_tensor low0; memset(&low0, 0, sizeof(low0)); ds4_gpu_tensor low1; memset(&low1, 0, sizeof(low1)); ds4_gpu_tensor full; memset(&full, 0, sizeof(full)); ds4_gpu_tensor p0; memset(&p0, 0, sizeof(p0)); ds4_gpu_tensor p1; memset(&p1, 0, sizeof(p1)); ds4_gpu_tensor sum; memset(&sum, 0, sizeof(sum)); CHECK(ds4_gpu_tensor_alloc_on(&heads, 0, (uint64_t)n_groups * group_dim * sizeof(float)) == 0, "attention_tp alloc heads"); CHECK(ds4_gpu_tensor_alloc_on(&low, 0, low_dim * sizeof(float)) == 0, "attention_tp alloc low"); CHECK(ds4_gpu_tensor_alloc_on(&low0, 0, (uint64_t)group_cnt * rank * sizeof(float)) == 0, "attention_tp alloc low0"); CHECK(ds4_gpu_tensor_alloc_on(&low1, 0, (uint64_t)group_cnt * rank * sizeof(float)) == 0, "attention_tp alloc low1"); CHECK(ds4_gpu_tensor_alloc_on(&full, 0, out_dim * sizeof(float)) == 0, "attention_tp alloc full"); CHECK(ds4_gpu_tensor_alloc_on(&p0, 0, out_dim * sizeof(float)) == 0, "attention_tp alloc p0"); CHECK(ds4_gpu_tensor_alloc_on(&p1, 0, out_dim * sizeof(float)) == 0, "attention_tp alloc p1"); CHECK(ds4_gpu_tensor_alloc_on(&sum, 0, out_dim * sizeof(float)) == 0, "attention_tp alloc sum"); CHECK(ds4_gpu_tensor_write(&heads, 0, host_heads, (uint64_t)n_groups * group_dim * sizeof(float)), "attention_tp write heads"); int ok = ds4_gpu_attention_output_q8_batch_tensor(&full, &low, NULL, NULL, model, model_size, 0, b_off, group_dim, rank, n_groups, out_dim, &heads, 1) && ds4_gpu_attention_output_q8_tp_tensor(&p0, &low0, model, model_size, 0, b_off, group_dim, rank, n_groups, 0, group_cnt, out_dim, &heads) && ds4_gpu_attention_output_q8_tp_tensor(&p1, &low1, model, model_size, 0, b_off, group_dim, rank, n_groups, group_cnt, group_cnt, out_dim, &heads) && ds4_gpu_add_tensor(&sum, &p0, &p1, (uint32_t)out_dim) && ds4_gpu_tensor_read(&full, 0, host_full, out_dim * sizeof(float)) && ds4_gpu_tensor_read(&sum, 0, host_sum, out_dim * sizeof(float)); CHECK(ok, "attention_tp compute"); for (uint64_t i = 0; i < out_dim; i++) { float diff = fabsf(host_full[i] - host_sum[i]); if (diff > 1.0e-3f) { fprintf(stderr, "FAIL: attention_tp mismatch row=%llu full=%f sum=%f diff=%g\n", (unsigned long long)i, host_full[i], host_sum[i], diff); return 1; } } ds4_gpu_tensor_free_in_place(&heads); ds4_gpu_tensor_free_in_place(&low); ds4_gpu_tensor_free_in_place(&low0); ds4_gpu_tensor_free_in_place(&low1); ds4_gpu_tensor_free_in_place(&full); ds4_gpu_tensor_free_in_place(&p0); ds4_gpu_tensor_free_in_place(&p1); ds4_gpu_tensor_free_in_place(&sum); free(model); free(host_heads); free(host_full); free(host_sum); ds4_gpu_cleanup(); fprintf(stderr, " attention_output_tp OK\n"); return 0; } static int run_attention_output_tp_peer_read(void) { int dev_count = 0; (void)cudaGetDeviceCount(&dev_count); if (dev_count < 2) { fprintf(stderr, " skipping attention_output_tp_peer_read (need >= 2 devices)\n"); return 0; } unsetenv("DS4_FORCE_HOST_BOUNCE"); ds4_gpu_config cfg; memset(&cfg, 0, sizeof(cfg)); cfg.n_gpus = 2; cfg.device_indices[0] = 0; cfg.device_indices[1] = 1; CHECK(ds4_gpu_init_multi(&cfg), "attention_peer init_multi"); if (!g_gpu_peer_ok[1][0]) { ds4_gpu_cleanup(); fprintf(stderr, " skipping attention_output_tp_peer_read (no peer access 1<-0)\n"); return 0; } const uint64_t group_dim = 64; const uint64_t rank = 64; const uint32_t n_groups = 4; const uint32_t group_cnt = 2; const uint64_t out_dim = 48; const uint64_t low_dim = (uint64_t)n_groups * rank; const uint64_t a_blocks = (group_dim + 31u) / 32u; const uint64_t b_blocks = (low_dim + 31u) / 32u; const uint64_t a_bytes = (uint64_t)n_groups * rank * a_blocks * 34u; const uint64_t b_off = a_bytes; const uint64_t b_bytes = out_dim * b_blocks * 34u; const uint64_t model_size = a_bytes + b_bytes; unsigned char *model = (unsigned char *)malloc((size_t)model_size); float *host_heads = (float *)malloc((size_t)n_groups * group_dim * sizeof(float)); float *host_full = (float *)malloc((size_t)out_dim * sizeof(float)); float *host_sum = (float *)malloc((size_t)out_dim * sizeof(float)); CHECK(model && host_heads && host_full && host_sum, "attention_peer host alloc"); pack_q8_identity_scale(model, group_dim, (uint64_t)n_groups * rank); pack_q8_identity_scale(model + b_off, low_dim, out_dim); for (uint64_t i = 0; i < (uint64_t)n_groups * group_dim; i++) { host_heads[i] = (float)((int)(i % 31u) - 15) * 0.046875f; } CHECK(ds4_gpu_set_model_map(model, model_size), "attention_peer set model map"); ds4_tensor_range r0 = { 0, model_size, 0 }; ds4_tensor_range r1 = { 0, model_size, 1 }; CHECK(ds4_gpu_device_cache_tensors(0, &r0, 1) == 0, "attention_peer cache model dev0"); CHECK(ds4_gpu_device_cache_tensors(1, &r1, 1) == 0, "attention_peer cache model dev1"); CHECK(ds4_gpu_set_current_device(0) == 0, "attention_peer current dev0"); ds4_gpu_tensor heads; memset(&heads, 0, sizeof(heads)); ds4_gpu_tensor low; memset(&low, 0, sizeof(low)); ds4_gpu_tensor low0; memset(&low0, 0, sizeof(low0)); ds4_gpu_tensor low1; memset(&low1, 0, sizeof(low1)); ds4_gpu_tensor full; memset(&full, 0, sizeof(full)); ds4_gpu_tensor p0; memset(&p0, 0, sizeof(p0)); ds4_gpu_tensor p1; memset(&p1, 0, sizeof(p1)); ds4_gpu_tensor staging; memset(&staging, 0, sizeof(staging)); ds4_gpu_tensor sum; memset(&sum, 0, sizeof(sum)); CHECK(ds4_gpu_tensor_alloc_on(&heads, 0, (uint64_t)n_groups * group_dim * sizeof(float)) == 0, "attention_peer alloc heads"); CHECK(ds4_gpu_tensor_alloc_on(&low, 0, low_dim * sizeof(float)) == 0, "attention_peer alloc low"); CHECK(ds4_gpu_tensor_alloc_on(&low0, 0, (uint64_t)group_cnt * rank * sizeof(float)) == 0, "attention_peer alloc low0"); CHECK(ds4_gpu_tensor_alloc_on(&low1, 1, (uint64_t)group_cnt * rank * sizeof(float)) == 0, "attention_peer alloc low1"); CHECK(ds4_gpu_tensor_alloc_on(&full, 0, out_dim * sizeof(float)) == 0, "attention_peer alloc full"); CHECK(ds4_gpu_tensor_alloc_on(&p0, 0, out_dim * sizeof(float)) == 0, "attention_peer alloc p0"); CHECK(ds4_gpu_tensor_alloc_on(&p1, 1, out_dim * sizeof(float)) == 0, "attention_peer alloc p1"); CHECK(ds4_gpu_tensor_alloc_on(&staging, 0, out_dim * sizeof(float)) == 0, "attention_peer alloc staging"); CHECK(ds4_gpu_tensor_alloc_on(&sum, 0, out_dim * sizeof(float)) == 0, "attention_peer alloc sum"); CHECK(ds4_gpu_tensor_write(&heads, 0, host_heads, (uint64_t)n_groups * group_dim * sizeof(float)), "attention_peer write heads"); int ok = ds4_gpu_attention_output_q8_batch_tensor(&full, &low, NULL, NULL, model, model_size, 0, b_off, group_dim, rank, n_groups, out_dim, &heads, 1) && ds4_gpu_attention_output_q8_tp_tensor(&p0, &low0, model, model_size, 0, b_off, group_dim, rank, n_groups, 0, group_cnt, out_dim, &heads) && ds4_gpu_tensor_wait_xdev(&heads, 1) && ds4_gpu_set_current_device(1) == 0 && ds4_gpu_attention_output_q8_tp_tensor(&p1, &low1, model, model_size, 0, b_off, group_dim, rank, n_groups, group_cnt, group_cnt, out_dim, &heads) && ds4_gpu_set_current_device(0) == 0 && ds4_gpu_add_xdev_tensor(&sum, &p0, &p1, &staging, (uint32_t)out_dim) && ds4_gpu_tensor_read(&full, 0, host_full, out_dim * sizeof(float)) && ds4_gpu_tensor_read(&sum, 0, host_sum, out_dim * sizeof(float)); CHECK(ok, "attention_peer compute"); for (uint64_t i = 0; i < out_dim; i++) { float diff = fabsf(host_full[i] - host_sum[i]); if (diff > 1.0e-3f) { fprintf(stderr, "FAIL: attention_peer mismatch row=%llu full=%f sum=%f diff=%g\n", (unsigned long long)i, host_full[i], host_sum[i], diff); return 1; } } ds4_gpu_tensor_free_in_place(&heads); ds4_gpu_tensor_free_in_place(&low); ds4_gpu_tensor_free_in_place(&low0); ds4_gpu_tensor_free_in_place(&low1); ds4_gpu_tensor_free_in_place(&full); ds4_gpu_tensor_free_in_place(&p0); ds4_gpu_tensor_free_in_place(&p1); ds4_gpu_tensor_free_in_place(&staging); ds4_gpu_tensor_free_in_place(&sum); free(model); free(host_heads); free(host_full); free(host_sum); ds4_gpu_cleanup(); fprintf(stderr, " attention_output_tp_peer_read OK\n"); return 0; } int main(void) { int dev_count = 0; (void)cudaGetDeviceCount(&dev_count); fprintf(stderr, "test_gpu_xdev: %d CUDA devices visible\n", dev_count); /* N=1 same-device path. */ if (run_one(1, 0)) return 1; if (run_copy3(1, 0)) return 1; if (run_top1()) return 1; if (run_topk2048()) return 1; if (run_glm_selected_attention()) return 1; if (run_glm_indexer_scores()) return 1; if (run_glm_decode_attention_staged()) return 1; if (run_moe_handoff_pack()) return 1; if (run_q8_kslice()) return 1; if (run_q8_matmul_top1_fused()) return 1; if (run_f16_small_out()) return 1; if (run_f16_small_batch()) return 1; if (run_f16_pair_batch()) return 1; if (run_q8_rowsplit()) return 1; if (run_q8_pair_batch()) return 1; if (run_attention_output_tp()) return 1; /* If 2+ GPUs, exercise peer-auto and forced-bounce paths. */ if (dev_count >= 2) { if (run_one(2, 0)) return 1; if (run_copy3(2, 0)) return 1; if (run_attention_output_tp_peer_read()) return 1; if (run_one(2, 1)) return 1; if (run_copy3(2, 1)) return 1; /* Stress: catches the cudaMemcpyPeer driver-corruption pattern that * single-size single-iter tests can miss. Runs both peer-auto and * forced-bounce so we get coverage of both code paths under load. */ if (run_stress(0)) return 1; if (run_stress(1)) return 1; if (run_add_xdev(0)) return 1; if (run_add_xdev(1)) return 1; } else { fprintf(stderr, " skipping multi-GPU paths (need >= 2 devices)\n"); } fprintf(stderr, "test_gpu_xdev PASS\n"); return 0; }