File size: 3,358 Bytes
c51a3e7 | 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 | // Full-vocabulary sampling for Top K = 0. Ordinary settings sample entirely on GPU.
export function sampleVocabulary(logits, settings, recent, position) {
const values = Float32Array.from(logits),
seen = new Set(recent.slice(-64));
for (const id of seen)
if (id < values.length)
values[id] =
values[id] < 0
? values[id] * settings.repeatPenalty
: values[id] / settings.repeatPenalty;
const order = Uint32Array.from({ length: values.length }, (_, i) => i).sort(
(a, b) => values[b] - values[a] || a - b,
);
if (settings.temperature <= 0) return order[0];
const weights = new Float64Array(order.length),
max = values[order[0]];
let sum = 0;
for (let i = 0; i < order.length; i++) {
weights[i] = Math.exp((values[order[i]] - max) / settings.temperature);
sum += weights[i];
}
let mass = 0,
count = 0;
do {
mass += weights[count++];
} while (count < order.length && mass < settings.topP * sum);
let hash =
(Math.imul(
(settings.seed ^ Math.imul(position, 2654435761)) >>> 0,
747796405,
) +
2891336453) >>>
0;
hash = Math.imul((hash >>> ((hash >>> 28) + 4)) ^ hash, 277803737) >>> 0;
hash = ((hash >>> 22) ^ hash) >>> 0;
const threshold = ((hash + 0.5) / 4294967296) * mass;
let cumulative = 0;
for (let i = 0; i < count; i++) {
cumulative += weights[i];
if (cumulative >= threshold) return order[i];
}
return order[count - 1];
}
// Exact speculative rejection sampling (Leviathan et al., 2023).
export function candidateDistribution(words, settings) {
const floats = new Float32Array(words.buffer, words.byteOffset, words.length),
k = Math.min(settings.topK || 1, words.length / 2);
const values = [];
let total = 0;
for (let i = 0; i < k; i++) {
const mass = Math.exp((floats[i * 2] - floats[0]) / settings.temperature);
values.push([words[i * 2 + 1], mass]);
total += mass;
}
let retained = 0,
count = 0;
do {
retained += values[count++][1];
} while (count < k && retained < settings.topP * total);
return new Map(
values.slice(0, count).map(([id, mass]) => [id, mass / retained]),
);
}
export function drawDistribution(distribution, random) {
let sum = 0,
last;
for (const [id, p] of distribution) {
sum += p;
last = id;
if (random < sum) return id;
}
return last;
}
export function verifySampledDraft(
target,
draft,
token,
acceptRandom,
resampleRandom,
) {
const q = draft.get(token) || 0,
p = target.get(token) || 0;
if (q > 0 && acceptRandom < Math.min(1, p / q))
return { accepted: true, token };
const residual = new Map();
let total = 0;
for (const [id, p] of target) {
const mass = Math.max(0, p - (draft.get(id) || 0));
residual.set(id, mass);
total += mass;
}
return {
accepted: false,
token: drawDistribution(
total > 1e-12
? new Map([...residual].map(([id, p]) => [id, p / total]))
: target,
resampleRandom,
),
};
}
export function randomStream(seed, position, stream) {
let x =
(seed ^ Math.imul(position, 2654435761) ^ Math.imul(stream, 2246822519)) >>>
0;
x = (Math.imul(x, 747796405) + 2891336453) >>> 0;
x = Math.imul((x >>> ((x >>> 28) + 4)) ^ x, 277803737) >>> 0;
return ((((x >>> 22) ^ x) >>> 0) + 0.5) / 4294967296;
}
|