antirez
1884 lines
84 KiB
C
1884 lines
84 KiB
C
/* 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 <cuda_runtime.h>
|
|
#include <math.h>
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
|
|
#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;
|
|
}
|