File size: 6,483 Bytes
8d0b310 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 | #include "ds4_gpu.h"
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <time.h>
static double monotonic_seconds(void) {
struct timespec ts;
clock_gettime(CLOCK_MONOTONIC, &ts);
return (double)ts.tv_sec + (double)ts.tv_nsec / 1000000000.0;
}
static double getenv_seconds(const char *name, double fallback) {
const char *s = getenv(name);
if (!s || !s[0]) return fallback;
char *end = NULL;
const double v = strtod(s, &end);
return end != s && v > 0.0 ? v : fallback;
}
static int check_large_topk(void) {
const uint32_t n_comp = 32768;
const uint32_t n_tokens = 32;
const uint32_t top_k = 512;
const uint64_t score_count = (uint64_t)n_comp * n_tokens;
float *scores_host = (float *)malloc((size_t)score_count * sizeof(float));
uint32_t *selected_host = (uint32_t *)malloc((size_t)n_tokens * top_k * sizeof(uint32_t));
if (!scores_host || !selected_host) return 1;
for (uint32_t t = 0; t < n_tokens; t++) {
for (uint32_t i = 0; i < n_comp; i++) {
scores_host[(uint64_t)t * n_comp + i] = (float)i;
}
}
ds4_gpu_tensor *scores = ds4_gpu_tensor_alloc(score_count * sizeof(float));
ds4_gpu_tensor *selected = ds4_gpu_tensor_alloc((uint64_t)n_tokens * top_k * sizeof(uint32_t));
int rc = 1;
double elapsed = 0.0;
if (scores && selected &&
ds4_gpu_tensor_write(scores, 0, scores_host, score_count * sizeof(float))) {
/* Exclude one-time CUDA module/kernel setup from the throughput guard. */
if (!ds4_gpu_indexer_topk_tensor(selected, scores, n_comp, n_tokens, top_k) ||
!ds4_gpu_synchronize()) {
rc = 1;
goto cleanup;
}
const double t0 = monotonic_seconds();
if (ds4_gpu_indexer_topk_tensor(selected, scores, n_comp, n_tokens, top_k) &&
ds4_gpu_synchronize()) {
elapsed = monotonic_seconds() - t0;
rc = ds4_gpu_tensor_read(selected, 0, selected_host,
(uint64_t)n_tokens * top_k * sizeof(uint32_t)) ? 0 : 1;
}
}
if (rc == 0) {
for (uint32_t t = 0; t < n_tokens && rc == 0; t++) {
for (uint32_t i = 0; i < top_k; i++) {
const uint32_t expected = n_comp - 1u - i;
const uint32_t got = selected_host[(uint64_t)t * top_k + i];
if (got != expected) {
fprintf(stderr, "top-k mismatch token=%u rank=%u got=%u expected=%u\n",
t, i, got, expected);
rc = 1;
break;
}
}
}
}
if (rc == 0) {
const double max_seconds = getenv_seconds("DS4_CUDA_TOPK_REGRESSION_SEC", 2.0);
fprintf(stderr, "cuda-regression: top-k n_comp=%u n_tokens=%u elapsed=%.3fs\n",
n_comp, n_tokens, elapsed);
if (elapsed > max_seconds) {
fprintf(stderr, "top-k regression: %.3fs exceeds %.3fs\n", elapsed, max_seconds);
rc = 1;
}
}
cleanup:
ds4_gpu_tensor_free(selected);
ds4_gpu_tensor_free(scores);
free(selected_host);
free(scores_host);
return rc;
}
static int check_decode_attention_overflow_path(void) {
const uint32_t n_head = 8;
const uint32_t head_dim = 512;
const uint32_t n_raw = 128;
const uint32_t n_comp = 8100;
const uint64_t q_count = (uint64_t)n_head * head_dim;
const uint64_t raw_count = (uint64_t)n_raw * head_dim;
const uint64_t comp_count = (uint64_t)n_comp * head_dim;
float *sinks = (float *)calloc(n_head, sizeof(float));
float *q_host = (float *)calloc((size_t)q_count, sizeof(float));
float *raw_host = (float *)calloc((size_t)raw_count, sizeof(float));
float *comp_host = (float *)calloc((size_t)comp_count, sizeof(float));
float *heads_host = (float *)calloc((size_t)q_count, sizeof(float));
if (!sinks || !q_host || !raw_host || !comp_host || !heads_host) return 1;
for (uint32_t c = 0; c < n_comp; c++) {
comp_host[(uint64_t)c * head_dim] = 1.0f;
}
ds4_gpu_tensor *heads = ds4_gpu_tensor_alloc(q_count * sizeof(float));
ds4_gpu_tensor *q = ds4_gpu_tensor_alloc(q_count * sizeof(float));
ds4_gpu_tensor *raw = ds4_gpu_tensor_alloc(raw_count * sizeof(float));
ds4_gpu_tensor *comp = ds4_gpu_tensor_alloc(comp_count * sizeof(float));
int rc = 1;
if (heads && q && raw && comp &&
ds4_gpu_tensor_write(q, 0, q_host, q_count * sizeof(float)) &&
ds4_gpu_tensor_write(raw, 0, raw_host, raw_count * sizeof(float)) &&
ds4_gpu_tensor_write(comp, 0, comp_host, comp_count * sizeof(float)) &&
ds4_gpu_attention_decode_heads_tensor(heads,
sinks,
n_head * sizeof(float),
0,
q,
raw,
n_raw,
n_raw,
0,
comp,
0,
n_comp,
NULL,
0,
n_head,
head_dim) &&
ds4_gpu_synchronize() &&
ds4_gpu_tensor_read(heads, 0, heads_host, q_count * sizeof(float))) {
rc = 0;
for (uint32_t h = 0; h < n_head; h++) {
const float v = heads_host[(uint64_t)h * head_dim];
if (v < 0.90f) {
fprintf(stderr, "attention fallback ignored compressed rows for head=%u value=%f\n",
h, (double)v);
rc = 1;
}
}
}
ds4_gpu_tensor_free(comp);
ds4_gpu_tensor_free(raw);
ds4_gpu_tensor_free(q);
ds4_gpu_tensor_free(heads);
free(heads_host);
free(comp_host);
free(raw_host);
free(q_host);
free(sinks);
return rc;
}
int main(void) {
if (!ds4_gpu_init()) return 1;
int rc = check_large_topk();
if (check_decode_attention_overflow_path() != 0) rc = 1;
ds4_gpu_cleanup();
if (rc == 0) puts("cuda long-context regression: OK");
return rc;
}
|