File size: 6,799 Bytes
9425aed | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 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 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 | /* sampler.c — zero-libc top-p nucleus sampler with temperature
* No malloc, no qsort, no stdlib. xorshift64* RNG.
*/
typedef unsigned long long uint64_t;
typedef unsigned int uint32_t;
typedef int int32_t;
typedef unsigned long long CUdeviceptr;
typedef int CUresult;
#define CUDA_SUCCESS 0
extern CUresult (*g_cuMemcpyDtoH_v2)(void*, CUdeviceptr, unsigned long long);
/* ── intrinsic expf (no libm) ───────────────────────────────────
* exp(x) = 2^(x * log2e) via hardware ex2.approx equivalent
* We use the compiler's built-in when available; on MSVC/clang
* this will lower to the FPU instruction.
*/
static float sov_expf(float x) {
/* Clamp to avoid overflow */
if (x > 88.0f) return 3.402823466e+38f;
if (x < -88.0f) return 0.0f;
/* Use __builtin_expf if available, otherwise polynomial */
#if defined(__GNUC__) || defined(__clang__)
return __builtin_expf(x);
#else
/* Horner's method: e^x ~ 1 + x + x^2/2 + x^3/6 + x^4/24 + x^5/120 */
float t = 1.0f + x * (1.0f + x * (0.5f + x * (0.16666667f + x * (0.041666668f + x * 0.008333334f))));
return t;
#endif
}
/* ── candidate type ─────────────────────────────────────────── */
#define SOV_MAX_TOPK 64
typedef struct {
float prob;
int32_t index;
} sov_cand_t;
/* ── partial top-k selection via insertion sort ─────────────── */
static int sov_partial_topk(const float* logits, int vocab_size,
int topk, sov_cand_t* out) {
int n = (vocab_size < topk) ? vocab_size : topk;
for (int i = 0; i < n; ++i) {
out[i].prob = logits[i];
out[i].index = i;
}
/* insertion sort: descending by logit */
for (int i = 1; i < n; ++i) {
sov_cand_t tmp = out[i];
int j = i - 1;
while (j >= 0 && out[j].prob < tmp.prob) {
out[j + 1] = out[j];
--j;
}
out[j + 1] = tmp;
}
/* scan remaining elements */
for (int i = n; i < vocab_size; ++i) {
float v = logits[i];
if (v <= out[n - 1].prob) continue;
sov_cand_t tmp = { v, i };
int j = n - 2;
while (j >= 0 && out[j].prob < tmp.prob) {
out[j + 1] = out[j];
--j;
}
out[j + 1] = tmp;
}
return n;
}
/* ── top-p (nucleus) filter ─────────────────────────────────── */
static int sov_top_p_filter(sov_cand_t* cands, int n, float top_p) {
if (top_p >= 1.0f || n <= 1) return n;
/* softmax over top-k logits */
float max_l = cands[0].prob;
float sum = 0.0f;
for (int i = 0; i < n; ++i) {
float p = sov_expf(cands[i].prob - max_l);
cands[i].prob = p;
sum += p;
}
for (int i = 0; i < n; ++i) cands[i].prob /= sum;
/* nucleus truncation */
float cumsum = 0.0f;
int keep = n;
for (int i = 0; i < n; ++i) {
cumsum += cands[i].prob;
if (cumsum >= top_p) { keep = i + 1; break; }
}
return keep;
}
/* ── categorical sample from (already normalised) cands ─────── */
static int sov_sample_categorical(sov_cand_t* cands, int n,
uint64_t* rng) {
/* xorshift64* */
uint64_t x = *rng;
x ^= x >> 12; x ^= x << 25; x ^= x >> 27;
*rng = x;
float u = (float)((x * 0x2545F4914F6CDD1Dull) >> 32) / 4294967296.0f;
float cumsum = 0.0f;
for (int i = 0; i < n; ++i) {
cumsum += cands[i].prob;
if (u < cumsum) return cands[i].index;
}
return cands[n - 1].index;
}
/* ── sov_sample_token ───────────────────────────────────────────
* d_logits: device pointer to vocab_size float32 logits
* Returns sampled token id, or negative on error.
* ---------------------------------------------------------------- */
int sov_sample_token(CUdeviceptr d_logits, int vocab_size,
float temperature, float top_p, int top_k,
uint64_t* rng_state) {
static float h_logits[128000];
static sov_cand_t cands[SOV_MAX_TOPK];
if (vocab_size <= 0 || vocab_size > 128000) return -1;
if (!g_cuMemcpyDtoH_v2) return -2;
CUresult err = g_cuMemcpyDtoH_v2(h_logits, d_logits,
(unsigned long long)(vocab_size * (int)sizeof(float)));
if (err != CUDA_SUCCESS) return -3;
/* greedy fast-path */
if (temperature == 0.0f) {
float best = h_logits[0]; int best_i = 0;
for (int i = 1; i < vocab_size; ++i) {
if (h_logits[i] > best) { best = h_logits[i]; best_i = i; }
}
return best_i;
}
/* temperature scaling */
if (temperature != 1.0f && temperature > 0.0f) {
float inv_t = 1.0f / temperature;
for (int i = 0; i < vocab_size; ++i) h_logits[i] *= inv_t;
}
int eff_topk = top_k;
if (eff_topk <= 0 || eff_topk > vocab_size) eff_topk = vocab_size;
if (eff_topk > SOV_MAX_TOPK) eff_topk = SOV_MAX_TOPK;
int n = sov_partial_topk(h_logits, vocab_size, eff_topk, cands);
/* if top_p < 1 we need probabilities for nucleus filter */
if (top_p < 1.0f) {
n = sov_top_p_filter(cands, n, top_p);
} else {
/* just softmax for sampling */
float max_l = cands[0].prob, sum = 0.0f;
for (int i = 0; i < n; ++i) {
float p = sov_expf(cands[i].prob - max_l);
cands[i].prob = p; sum += p;
}
for (int i = 0; i < n; ++i) cands[i].prob /= sum;
}
return sov_sample_categorical(cands, n, rng_state);
}
/* ── sov_sample_greedy ──────────────────────────────────────────── */
int sov_sample_greedy(CUdeviceptr d_logits, int vocab_size) {
static float h_logits[128000];
if (vocab_size <= 0 || vocab_size > 128000) return -1;
if (!g_cuMemcpyDtoH_v2) return -2;
CUresult err = g_cuMemcpyDtoH_v2(h_logits, d_logits,
(unsigned long long)(vocab_size * (int)sizeof(float)));
if (err != CUDA_SUCCESS) return -3;
float best = h_logits[0]; int best_i = 0;
for (int i = 1; i < vocab_size; ++i) {
if (h_logits[i] > best) { best = h_logits[i]; best_i = i; }
}
return best_i;
}
|