From c3076e9288cb1942cbdddb15c653a87716c56b31 Mon Sep 17 00:00:00 2001 From: oceanio Date: Tue, 20 Oct 2026 08:45:27 +0800 Subject: [PATCH 1/3] metal: the request's sampling (the CUDA sampler's chain, on the CPU over the window's head logits) The Mac engine picked the window's argmax for every request, so temperature, top_p, top_k, min_p, seed, the penalties and logit_bias changed nothing. Now Verifier::run's Metal branch does what the CUDA path does after its window: a sampled request, one with penalties or with a logit_bias, draws row t from the window's logits (unified memory, Window::logits_row) with Philox(seed, pos0 + t) and its own penalty row; drafts are kept while the sampled token equals them. Greedy without penalties or bias keeps the window's argmax. metal/src/engine/host_sampler.cpp: logit_bias -> penalties -> top_k (one NEON-filtered pass, the kernels' order: score descending, ids ascending) -> top_p -> min_p -> temperature -> draw, the kernels' float/double arithmetic. 9-46 us per row on 248,320 logits (M5 Max, host_sampler_parity --bench), so no GPU kernel. Verifier::set_logit_bias returned before storing the bias on Metal: the Metal init returned before the line that sets le_ to n_layers, and the -1 it left read as an earlier stage of a layer split. The Metal init now sets it as the others do, so the C++ that CUDA, HIP and SYCL compile is unchanged (unifdef -USTRATA_GPU_METAL gives v0.2.6's file). Co-Authored-By: Claude Opus 5.5 --- metal/src/engine/host_sampler.cpp | 178 ++++++++++++++++++++++++++++++ metal/src/engine/host_sampler.h | 28 +++++ metal/src/engine/window.cpp | 7 ++ src/core/verify.cpp | 25 ++++- 4 files changed, 237 insertions(+), 1 deletion(-) create mode 100644 metal/src/engine/host_sampler.cpp create mode 100644 metal/src/engine/host_sampler.h diff --git a/metal/src/engine/host_sampler.cpp b/metal/src/engine/host_sampler.cpp new file mode 100644 index 00000000000..6992bdf8fd2 --- /dev/null +++ b/metal/src/engine/host_sampler.cpp @@ -0,0 +1,178 @@ +// metal/src/engine/host_sampler.cpp - see host_sampler.h. +// +// The selection is the CUDA kernels' list: candidates ranked by (score descending, id ascending), a strict total +// order, so the list does not depend on the order candidates are offered in. NaN and -inf are never candidates; a +// list with fewer than k candidates is filled with (id 0, -inf), as a kernel round that finds nothing stores. The +// scan offers a block of 16 logits only when one of them reaches the list's last score (NEON compare), so a row costs +// about one pass over its 248,320 floats. The penalised ids (at most the window's length) are taken out of the scan +// and offered with their penalised scores. The tail (top_p, min_p, temperature, the draw) is the kernels' arithmetic +// in the same order: float scores, double sums. +#include "host_sampler.h" + +#include +#include +#include +#include +#include + +#if defined(__ARM_NEON) +#include +#endif + +namespace strata::metal_engine { +namespace { + +constexpr int kSelMax = 64; // the widest top_k list (sampler.cu's kSelMax) +constexpr float kNegInf = -std::numeric_limits::infinity(); + +struct TopList { + int k = 1, n = 0; + float val[kSelMax]; + int id[kSelMax]; + static bool before(float v, int i, float w, int j) { return v > w || (v == w && i < j); } + // the list's last score once full (a candidate below it cannot enter), -inf before + float floor() const { return n == k ? val[k - 1] : kNegInf; } + void offer(float v, int i) { + if (!(v > kNegInf)) return; // NaN and -inf are never picked + if (n == k) { + if (!before(v, i, val[k - 1], id[k - 1])) return; + --n; + } + int at = n; + for (; at > 0 && before(v, i, val[at - 1], id[at - 1]); --at) { + val[at] = val[at - 1]; + id[at] = id[at - 1]; + } + val[at] = v; + id[at] = i; + ++n; + } +}; + +// sampler.cu's apply_penalties (llama_sampler_penalties_apply): the repeat penalty multiplies non-positive scores and +// divides positive ones; the presence penalty is boolean +float penalised(float s, int count, const strata::kernels::SamplerParams& p) { + if (count <= 0) return s; + if (s <= 0.0f) s *= p.penalty_repeat; + else s /= p.penalty_repeat; + s -= (float) count * p.penalty_freq + (count > 0 ? 1.0f : 0.0f) * p.penalty_present; + return s; +} + +// offers every v in [0, nv) of `s` that is not marked in `skip` (null: none) +void scan(const float* s, int nv, const std::vector* skip, TopList& L) { + auto one = [&](int v) { + if (skip != nullptr && (((*skip)[(size_t) v >> 6] >> (v & 63)) & 1u)) return; + L.offer(s[v], v); + }; + int v = 0; + for (; v < nv && L.n < L.k; ++v) one(v); +#if defined(__ARM_NEON) + for (; v + 16 <= nv; v += 16) { + const float32x4_t f = vdupq_n_f32(L.floor()); + // a score equal to the floor can still enter (a lower id than the last entry's, from the penalised set) + const uint32x4_t m = vorrq_u32(vorrq_u32(vcgeq_f32(vld1q_f32(s + v), f), vcgeq_f32(vld1q_f32(s + v + 4), f)), + vorrq_u32(vcgeq_f32(vld1q_f32(s + v + 8), f), vcgeq_f32(vld1q_f32(s + v + 12), f))); + if (vmaxvq_u32(m) == 0) continue; + for (int j = 0; j < 16; ++j) one(v + j); + } +#endif + for (; v < nv; ++v) + if (s[v] >= L.floor()) one(v); +} + +} // namespace + +float philox_uniform(uint64_t seed, uint64_t counter) { + // sampler.cu's philox_uniform: Philox 4x32-10, key (round, 0), the counter in (c0, c1), the seed in (c2, c3) + uint32_t c0 = (uint32_t) counter, c1 = (uint32_t) (counter >> 32); + uint32_t c2 = (uint32_t) seed, c3 = (uint32_t) (seed >> 32); + for (uint32_t i = 0; i < 10; ++i) { + const uint64_t m0 = (uint64_t) 0x9E3779B9u * c0, m1 = (uint64_t) 0xBB67AE85u * c2; + const uint32_t n0 = (uint32_t) (m1 >> 32) ^ c1 ^ i; + const uint32_t n2 = (uint32_t) (m0 >> 32) ^ c3; + c1 = (uint32_t) m1; + c3 = (uint32_t) m0; + c0 = n0; + c2 = n2; + } + return (float) (c0 >> 8) * (1.0f / 16777216.0f); // 24 bits: uniform in [0, 1), never 1.0 +} + +int sample_row(const float* logits, int nv, const int32_t* history, int history_len, + const strata::kernels::SamplerParams& p, const float* logit_bias, uint64_t counter) { + if (nv <= 0) return 0; + const bool greedy = p.greedy || p.temperature <= 0.0f; + // the scores before the penalties: the logits, or logits + bias (one pass into a scratch row) + const float* s = logits; + thread_local std::vector biased; + if (logit_bias != nullptr) { + biased.resize((size_t) nv); + for (int v = 0; v < nv; ++v) biased[(size_t) v] = logits[v] + logit_bias[v]; + s = biased.data(); + } + // the counted window: the last min(penalty_last_n, history_len) entries; ids outside the vocabulary are ignored + thread_local std::vector> pen; // (id, count) + thread_local std::vector mark; + pen.clear(); + const int hlen = (history != nullptr && p.penalty_last_n > 0) ? std::min(p.penalty_last_n, history_len) : 0; + if (hlen > 0) { + for (int i = history_len - hlen; i < history_len; ++i) + if (history[i] >= 0 && history[i] < nv) pen.push_back({history[i], 0}); + std::sort(pen.begin(), pen.end()); + size_t w = 0; + for (size_t r = 0; r < pen.size();) { + size_t e = r; + while (e < pen.size() && pen[e].first == pen[r].first) ++e; + pen[w++] = {pen[r].first, (int) (e - r)}; + r = e; + } + pen.resize(w); + mark.assign(((size_t) nv + 63) / 64, 0); + for (const auto& [id, c] : pen) mark[(size_t) id >> 6] |= 1ull << (id & 63); + } + TopList L; + L.k = greedy ? 1 : std::min(nv, (p.top_k > 0 && p.top_k < kSelMax) ? p.top_k : kSelMax); + scan(s, nv, pen.empty() ? nullptr : &mark, L); + for (const auto& [id, c] : pen) L.offer(penalised(s[id], c, p), id); + if (greedy) return L.n > 0 ? L.id[0] : 0; // sampler_greedy_kernel: 0 when nothing is above -inf + const int k = L.k; + for (int i = L.n; i < k; ++i) { L.id[i] = 0; L.val[i] = kNegInf; } // the rounds that found nothing + const float* lg = L.val; + // top_p over the list, then min_p on its survivors, then the temperature and one draw (sampler.cu's tail) + int n_keep = k; + float mx = lg[0]; + for (int i = 1; i < k; ++i) mx = std::fmax(mx, lg[i]); + if (p.top_p < 1.0f) { + double sum = 0.0; + for (int i = 0; i < k; ++i) sum += std::exp((double) lg[i] - (double) mx); + double cum = 0.0; + int cut = k; + for (int i = 0; i < k; ++i) { + cum += std::exp((double) lg[i] - (double) mx) / sum; + if (cum >= (double) p.top_p) { cut = i + 1; break; } + } + if (cut < p.min_keep) cut = p.min_keep < k ? p.min_keep : k; + n_keep = cut; + } + if (p.min_p > 0.0f) { + const float thresh = lg[0] + std::log(p.min_p); + for (int i = 0; i < n_keep; ++i) + if (lg[i] < thresh) { n_keep = i; break; } + } + const float inv_t = 1.0f / p.temperature; + float smx = lg[0] * inv_t; + for (int i = 1; i < n_keep; ++i) smx = std::fmax(smx, lg[i] * inv_t); + double sum = 0.0; + for (int i = 0; i < n_keep; ++i) sum += std::exp((double) (lg[i] * inv_t) - (double) smx); + int pick = L.id[n_keep > 0 ? n_keep - 1 : 0]; + const float u = philox_uniform(p.seed, counter); + double cum = 0.0; + for (int i = 0; i < n_keep; ++i) { + cum += std::exp((double) (lg[i] * inv_t) - (double) smx) / sum; + if ((double) u < cum) { pick = L.id[i]; break; } + } + return pick; +} + +} // namespace strata::metal_engine diff --git a/metal/src/engine/host_sampler.h b/metal/src/engine/host_sampler.h new file mode 100644 index 00000000000..7280beca11c --- /dev/null +++ b/metal/src/engine/host_sampler.h @@ -0,0 +1,28 @@ +// metal/src/engine/host_sampler.h - the request's sampling on the Mac, on the CPU over the window's head logits. +// +// The CUDA engine's sampler (src/kernels/cuda/sampler.cu, sample_tokens) row for row: on unified memory the head +// logits are host memory once the window has finished, so the Mac picks the token where it reads it. Same chain, +// same order (llama.cpp's): logit_bias, penalties, top_k (1..64; 0 or wider = 64), top_p, min_p, temperature, one +// Philox 4x32-10 draw at (seed, counter). The greedy pick (temperature <= 0) is the argmax after the bias and the +// penalties, ties to the lowest id. metal/tests/engine/host_sampler_parity.cpp pins it against a serial +// transcription of the CUDA kernel, draw for draw, and times it. +#pragma once + +#include "strata/kernels/sampler.hpp" + +#include + +namespace strata::metal_engine { + +/// One row's token. `logits`: n_vocab floats (host memory). `history`: the row's penalty history, `history_len` +/// ids, most recent last, -1 for unused front slots (strata::kernels::penalty_rows); only the last +/// min(p.penalty_last_n, history_len) are counted; null or 0 = no penalties. `logit_bias`: n_vocab floats or null +/// (-inf bans). `counter`: the draw's Philox counter, the position the row predicts from (the solo window's +/// pos0 + t, a slot row's position). p.logit_bias and p.counter are not read (device pointer / per-call). +int sample_row(const float* logits, int n_vocab, const int32_t* history, int history_len, + const strata::kernels::SamplerParams& p, const float* logit_bias, uint64_t counter); + +/// The Philox 4x32-10 uniform in [0, 1) of the CUDA sampler (24 bits), for the tests. +float philox_uniform(uint64_t seed, uint64_t counter); + +} // namespace strata::metal_engine diff --git a/metal/src/engine/window.cpp b/metal/src/engine/window.cpp index add35c0e9f3..59bb6228ecf 100644 --- a/metal/src/engine/window.cpp +++ b/metal/src/engine/window.cpp @@ -124,6 +124,10 @@ class Window { bool commit(int n_keep, std::string& err); bool finish(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t* out, double t0); bool copy_logits(int t, float* host) const; + // row t of the last window's head logits, in place; null past the window's rows + const float* logits_row(int t) const { + return t >= 0 && t < head_rows_ ? lm_.as() + (size_t) t * (size_t) n_vocab_ : nullptr; + } int64_t vocab() const { return n_vocab_; } bool ready() const { return cx_ != nullptr; } bool streaming() const { return stream_; } @@ -309,6 +313,7 @@ class Window { std::vector ple_hist_; // NG_HIST x hc_dim normalized rows, row-fastest bool ple_on_ = false; int last_T_ = 0; + int head_rows_ = 0; // rows of lm_ the last window's head wrote (logits_row) double last_gpu_ms_ = 0; int last_n_disp_ = 0; int32_t last_tokens_[MAXT] = {}; @@ -1459,6 +1464,7 @@ bool Window::finish(int Tn, const int32_t* tokens, int64_t pos0, bool head, int3 last_head_ = head; for (int t = 0; t < Tn; ++t) out[t] = head ? tokout_.as()[t] : -1; last_T_ = Tn; + head_rows_ = head ? Tn : 0; for (int t = 0; t < Tn; ++t) last_tokens_[t] = tokens[t]; static const bool trace = std::getenv("STRATA_METAL_TRACE") != nullptr; if (trace && head) { @@ -2646,6 +2652,7 @@ bool window_commit(void* w, int n_keep, std::string& err) { return ok; } bool window_copy_logits(void* w, int t, float* host) { return ((Window*) w)->copy_logits(t, host); } +const float* window_logits(void* w, int t) { return ((Window*) w)->logits_row(t); } // the MTP hook (src/core/mtp.cpp under STRATA_GPU_METAL): the drafter lives beside the decode window it reads namespace { diff --git a/src/core/verify.cpp b/src/core/verify.cpp index c8eef3f30e9..318ac57fbd5 100644 --- a/src/core/verify.cpp +++ b/src/core/verify.cpp @@ -91,6 +91,11 @@ bool window_init(void* w, const strata::core::ModelGeometry& g, strata::core::Se bool window_run(void* w, int T, const int32_t* tokens, int64_t pos0, bool head, int32_t* out, std::string& err); bool window_commit(void* w, int n_keep, std::string& err); bool window_copy_logits(void* w, int t, float* host); +// the head logits of row t of the last window, host memory (unified); null past the window +const float* window_logits(void* w, int t); +// the request's sampling over one row of them, on the CPU (metal/src/engine/host_sampler.h: the CUDA sampler's chain) +int sample_row(const float* logits, int n_vocab, const int32_t* history, int history_len, + const strata::kernels::SamplerParams& p, const float* logit_bias, uint64_t counter); } // namespace strata::metal_engine #endif @@ -821,6 +826,7 @@ bool Verifier::init(const WeightTable& wt, const ModelGeometry& g, SessionState& sampling_.temperature = 0.0f; #if defined(STRATA_GPU_METAL) if (metal_ == nullptr) metal_ = strata::metal_engine::window_create(); + if (le_ < 0) le_ = g.n_layers; // the whole model, as below (set_logit_bias reads le_ < n_layers as an earlier stage) return strata::metal_engine::window_init(metal_, g, ss, max_t, ss.qsa_states[ss.qsa_primary()].max_cells, &n_vocab_, err); #endif @@ -3077,7 +3083,24 @@ bool Verifier::run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool, if (measure) { err = "verify: proposal unsupported on Metal"; return false; } (void) pool; (void) user; - return strata::metal_engine::window_run(metal_, T, tokens, pos0, head_sampling_, out, err); + if (!strata::metal_engine::window_run(metal_, T, tokens, pos0, head_sampling_, out, err)) return false; + { // a sampled or penalized request (or one with a logit_bias): the CUDA path's host-side sampling below, on the + // CPU over the window's logits (unified memory). Row t draws Philox(seed, pos0 + t) with its own penalty row. + // Greedy without either keeps the window's argmax: the tokens are the ones it always gave. + const bool sampled = !sampling_.greedy && sampling_.temperature > 0.0f; + if (head_sampling_ && (sampled || hist_d_ != nullptr || !logit_bias_host_.empty())) { + const int32_t* hist = hist_d_ != nullptr ? (const int32_t*) strata::gpu::host_view(hist_d_) : nullptr; + for (int t = 0; t < T; ++t) { + const float* row = strata::metal_engine::window_logits(metal_, t); + if (row == nullptr) { err = "verify: the window's head logits are missing"; return false; } + out[t] = strata::metal_engine::sample_row(row, (int) n_vocab_, hist ? hist + (size_t) t * hist_len_ : nullptr, + hist ? hist_len_ : 0, sampling_, + logit_bias_host_.empty() ? nullptr : logit_bias_host_.data(), + (uint64_t) (pos0 + t)); + } + } + } + return true; #endif using namespace strata::kernels; const OnDevice on_device(device_); From 072170609ad0dec60cce04857a071f8e89439e4b Mon Sep 17 00:00:00 2001 From: oceanio Date: Tue, 20 Oct 2026 08:45:27 +0800 Subject: [PATCH 2/3] tests: tools/metal_sampling_test.py (the request's sampling through --serve: reproducible seeds, picks draw for draw against a reference sampler on the dumped logits, penalties, logit_bias) and metal/tests/engine/host_sampler_parity.cpp metal_sampling_test.py: a reference sampler written from src/kernels/cuda/sampler.cu (Philox 4x32-10 at (seed, position), logit_bias, penalties, top_k, top_p, min_p, temperature) picks from each one-token request's first-window logits (STRATA_DUMP_FIRST_LOGITS); every engine pick must be its pick. 13 checks; on 0.2.0 without the engine change 11 fail (every seed gives the greedy text). --bench: decode tok/s greedy vs sampled. host_sampler_parity: the Mac host sampler (host_sampler.h) against a serial transcription of the CUDA kernels (sampler_parity.cpp's mirror_select/mirror_pick), 592 picks on 248,320-logit rows with ties, -inf, NaN, bans and penalties; --bench times a row. CPU only. Co-Authored-By: Claude Opus 5.5 --- metal/CMakeLists.txt | 7 + metal/tests/engine/host_sampler_parity.cpp | 222 ++++++++++++++ tools/metal_sampling_test.py | 322 +++++++++++++++++++++ 3 files changed, 551 insertions(+) create mode 100644 metal/tests/engine/host_sampler_parity.cpp create mode 100644 tools/metal_sampling_test.py diff --git a/metal/CMakeLists.txt b/metal/CMakeLists.txt index 4a8bdd8cd0c..932144e3294 100644 --- a/metal/CMakeLists.txt +++ b/metal/CMakeLists.txt @@ -122,3 +122,10 @@ target_link_libraries(strata_metal_kernels PUBLIC strata_metal_host mtl_rt) add_executable(strata ${S}/program/generate.cpp) target_link_libraries(strata PRIVATE strata_metal_kernels strata_metal_host strata_metal_kernels) + +# The host sampler (src/engine/host_sampler.cpp: the request's sampling over the window's head logits) against a serial +# transcription of the CUDA sampler, draw for draw; --bench times it per row. CPU only: no GPU, no model. In its own +# folder: tests/*.cpp above are the kernel tree's mtl_rt-only programs. +add_executable(host_sampler_parity tests/engine/host_sampler_parity.cpp) +target_link_libraries(host_sampler_parity PRIVATE strata_metal_kernels) +add_test(NAME host_sampler_parity COMMAND host_sampler_parity) diff --git a/metal/tests/engine/host_sampler_parity.cpp b/metal/tests/engine/host_sampler_parity.cpp new file mode 100644 index 00000000000..853a7a43fe9 --- /dev/null +++ b/metal/tests/engine/host_sampler_parity.cpp @@ -0,0 +1,222 @@ +// metal/tests/engine/host_sampler_parity.cpp - the Mac's host sampler (src/engine/host_sampler.cpp) against a serial +// transcription of the CUDA sampler, draw for draw, and its time per row. CPU only: no GPU, no model. +// +// host_sampler_parity [--bench] exit 0 = every pick equals the reference's +// +// The reference below is src/kernels/sampler_parity.cpp's `mirror_select` + `mirror_pick` (the CUDA kernels' +// semantics, which that test pins on the GPU), with the logit_bias added before the penalties as the kernels do, and +// the greedy pick of `sampler_greedy_kernel` (argmax after bias and penalties, ties to the lowest id, 0 when every +// value is -inf or NaN). Rows: 248,320 logits (the model's vocabulary) with planted ties at the top, -inf and NaN +// entries, a logit_bias with bans, and a penalty history; configurations from greedy to top_k 64 with top_p and min_p. +#include "../../src/engine/host_sampler.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +using strata::kernels::SamplerParams; + +namespace { + +float ref_philox(uint64_t seed, uint64_t counter) { // sampler_parity.cpp's host_philox_uniform + uint32_t c0 = (uint32_t) counter, c1 = (uint32_t) (counter >> 32); + uint32_t c2 = (uint32_t) seed, c3 = (uint32_t) (seed >> 32); + for (uint32_t i = 0; i < 10; ++i) { + const uint32_t hi0 = (uint32_t) (((uint64_t) 0x9E3779B9u * c0) >> 32); + const uint32_t hi1 = (uint32_t) (((uint64_t) 0xBB67AE85u * c2) >> 32); + const uint32_t lo0 = 0x9E3779B9u * c0; + const uint32_t lo1 = 0xBB67AE85u * c2; + const uint32_t n0 = hi1 ^ c1 ^ i, n1 = lo1, n2 = hi0 ^ c3 ^ 0u, n3 = lo0; + c0 = n0; c1 = n1; c2 = n2; c3 = n3; + } + return (float) (c0 >> 8) * (1.0f / 16777216.0f); +} + +// the scores the kernels rank: logit (+ bias), then the penalties over the counted window +std::vector ref_scores(const float* l, int nv, const float* bias, const std::vector& window, + const SamplerParams& p) { + std::vector s((size_t) nv); + for (int v = 0; v < nv; ++v) s[(size_t) v] = bias ? l[v] + bias[v] : l[v]; + std::vector cnt((size_t) nv, 0); + for (int h : window) if (h >= 0 && h < nv) ++cnt[(size_t) h]; + for (int v = 0; v < nv; ++v) { + const int c = cnt[(size_t) v]; + if (c <= 0) continue; + float x = s[(size_t) v]; + if (x <= 0.0f) x *= p.penalty_repeat; else x /= p.penalty_repeat; + x -= (float) c * p.penalty_freq + (c > 0 ? 1.0f : 0.0f) * p.penalty_present; + s[(size_t) v] = x; + } + return s; +} + +int ref_pick(const float* l, int nv, const float* bias, const std::vector& window, const SamplerParams& p, + uint64_t counter) { + const std::vector s = ref_scores(l, nv, bias, window, p); + if (p.greedy || p.temperature <= 0.0f) { // sampler_greedy_kernel + float bv = -std::numeric_limits::infinity(); + int best = nv; + for (int v = 0; v < nv; ++v) if (s[(size_t) v] > bv) { bv = s[(size_t) v]; best = v; } + return best < nv ? best : 0; + } + const int k = std::min(nv, (p.top_k > 0 && p.top_k < 64) ? p.top_k : 64); + std::vector ids; // mirror_select: k argmax rounds over the not-yet-taken + std::vector lg; + std::vector taken((size_t) nv, 0); + for (int i = 0; i < k; ++i) { + int best = nv; + float bv = -std::numeric_limits::infinity(); + for (int v = 0; v < nv; ++v) + if (!taken[(size_t) v] && s[(size_t) v] > bv) { bv = s[(size_t) v]; best = v; } + const int id = best < nv ? best : 0; + taken[(size_t) id] = 1; + ids.push_back(id); + lg.push_back(bv); + } + // mirror_pick + int n_keep = k; + float mx = lg[0]; + for (int i = 1; i < k; ++i) mx = std::fmax(mx, lg[(size_t) i]); + if (p.top_p < 1.0f) { + double sum = 0.0; + for (int i = 0; i < k; ++i) sum += std::exp((double) lg[(size_t) i] - (double) mx); + double cum = 0.0; + int cut = k; + for (int i = 0; i < k; ++i) { + cum += std::exp((double) lg[(size_t) i] - (double) mx) / sum; + if (cum >= (double) p.top_p) { cut = i + 1; break; } + } + if (cut < p.min_keep) cut = p.min_keep < k ? p.min_keep : k; + n_keep = cut; + } + if (p.min_p > 0.0f) { + const float thresh = lg[0] + std::log(p.min_p); + for (int i = 0; i < n_keep; ++i) + if (lg[(size_t) i] < thresh) { n_keep = i; break; } + } + const float inv_t = p.temperature > 0.0f ? 1.0f / p.temperature : 0.0f; + auto scaled = [&](int i) { return lg[(size_t) i] * inv_t; }; + float smx = scaled(0); + for (int i = 1; i < n_keep; ++i) smx = std::fmax(smx, scaled(i)); + double sum = 0.0; + for (int i = 0; i < n_keep; ++i) sum += std::exp((double) scaled(i) - (double) smx); + int pick = ids[(size_t) (n_keep > 0 ? n_keep - 1 : 0)]; + const float u = ref_philox(p.seed, counter); + double cum = 0.0; + for (int i = 0; i < n_keep; ++i) { + cum += std::exp((double) scaled(i) - (double) smx) / sum; + if ((double) u < cum) { pick = ids[(size_t) i]; break; } + } + return pick; +} + +// the last min(penalty_last_n, len) entries of a history row (the kernels' window) +std::vector window_of(const std::vector& h, const SamplerParams& p) { + const int n = std::min(p.penalty_last_n, (int) h.size()); + return n > 0 ? std::vector(h.end() - n, h.end()) : std::vector{}; +} + +} // namespace + +int main(int argc, char** argv) { + const bool bench = argc > 1 && std::strcmp(argv[1], "--bench") == 0; + constexpr int NV = 248320; + std::mt19937_64 rng(11); + std::normal_distribution nd(0.0f, 2.5f); + int bad = 0, checks = 0; + + // Philox: the engine's draw is the reference's + for (uint64_t c = 0; c < 4096; ++c) { + const uint64_t seed = rng(); + if (strata::metal_engine::philox_uniform(seed, c * 977) != ref_philox(seed, c * 977)) ++bad; + } + std::printf(" philox: %s\n", bad ? "*** WRONG ***" : "matches 4096 draws"); + + struct Cfg { const char* name; SamplerParams p; bool bias, pen; }; + auto mk = [](float t, int k, float tp, float mp, int pen_n, float rep, float fr, float pr) { + SamplerParams p; + p.greedy = t <= 0.0f; p.temperature = t; p.top_k = k; p.top_p = tp; p.min_p = mp; + p.penalty_last_n = pen_n; p.penalty_repeat = rep; p.penalty_freq = fr; p.penalty_present = pr; + return p; + }; + const Cfg cfgs[] = { + {"greedy + logit_bias", mk(0.0f, 20, 1.0f, 0.0f, 0, 1.0f, 0.0f, 0.0f), true, false}, + {"greedy + penalties", mk(0.0f, 20, 1.0f, 0.0f, 64, 1.3f, 0.0f, 0.0f), false, true}, + {"t0.7 k20 p0.95", mk(0.7f, 20, 0.95f, 0.0f, 0, 1.0f, 0.0f, 0.0f), false, false}, + {"t1.0 k40 p1.0", mk(1.0f, 40, 1.0f, 0.0f, 0, 1.0f, 0.0f, 0.0f), false, false}, + {"t1.0 k0 (64) p0.9 min_p0.05", mk(1.0f, 0, 0.9f, 0.05f, 0, 1.0f, 0.0f, 0.0f), false, false}, + {"t0.6 k1", mk(0.6f, 1, 0.95f, 0.0f, 0, 1.0f, 0.0f, 0.0f), false, false}, + {"t1.2 k64 + bias + penalties", mk(1.2f, 64, 0.95f, 0.02f, 256, 1.5f, 0.1f, 0.2f), true, true}, + {"t3.0 k50 p1.0 (flat)", mk(3.0f, 50, 1.0f, 0.0f, 0, 1.0f, 0.0f, 0.0f), false, false}, + }; + std::vector l((size_t) NV), bias((size_t) NV); + std::vector hist(512); + for (int trial = 0; trial < 12; ++trial) { + for (float& x : l) x = nd(rng); + // a head the chain cuts into: a few strong logits, two ties at the top, some -inf and a NaN + for (int j = 0; j < 30; ++j) l[(size_t) (rng() % NV)] = 9.0f + (float) (rng() % 40) * 0.1f; + const int a = (int) (rng() % NV), b = (int) (rng() % NV); + l[(size_t) a] = l[(size_t) b] = 14.0f; + for (int j = 0; j < 50; ++j) l[(size_t) (rng() % NV)] = -std::numeric_limits::infinity(); + l[(size_t) (rng() % NV)] = std::numeric_limits::quiet_NaN(); + std::fill(bias.begin(), bias.end(), 0.0f); + bias[(size_t) std::min(a, b)] = -std::numeric_limits::infinity(); // ban one of the tie + for (int j = 0; j < 20; ++j) bias[(size_t) (rng() % NV)] = (float) ((int) (rng() % 40) - 20) * 0.5f; + for (size_t j = 0; j < hist.size(); ++j) hist[j] = j < 20 ? -1 : (int32_t) (rng() % 64 == 0 ? a : rng() % NV); + hist[hist.size() - 1] = b; // the other of the tie is penalised + for (const Cfg& c : cfgs) { + int wrong = 0; + for (int d = 0; d < 6; ++d) { + SamplerParams p = c.p; + p.seed = rng(); + const uint64_t counter = rng() % 300000; + const float* bp = c.bias ? bias.data() : nullptr; + const int32_t* hp = c.pen ? hist.data() : nullptr; + const int got = strata::metal_engine::sample_row(l.data(), NV, hp, c.pen ? (int) hist.size() : 0, p, bp, + counter); + const int want = ref_pick(l.data(), NV, bp, c.pen ? window_of(hist, p) : std::vector{}, p, counter); + wrong += got != want; + ++checks; + } + if (wrong) std::printf(" trial %d %-30s *** %d WRONG ***\n", trial, c.name, wrong); + bad += wrong; + } + } + // edge rows: all -inf (greedy and sampled answer 0), fewer finite logits than top_k + { + std::vector e((size_t) NV, -std::numeric_limits::infinity()); + for (const Cfg& c : cfgs) { + SamplerParams p = c.p; + p.seed = 5; + bad += strata::metal_engine::sample_row(e.data(), NV, nullptr, 0, p, nullptr, 9) != ref_pick(e.data(), NV, nullptr, {}, p, 9); + e[100] = 1.0f; e[7] = 1.0f; e[200000] = 3.0f; + bad += strata::metal_engine::sample_row(e.data(), NV, nullptr, 0, p, nullptr, 9) != ref_pick(e.data(), NV, nullptr, {}, p, 9); + checks += 2; + } + } + std::printf("host_sampler_parity: %d picks, %s\n", checks, bad ? "*** WRONG ***" : "all equal the reference"); + + if (bench) { // time per row on a model-sized row (the window's head: 248,320 logits) + for (float& x : l) x = nd(rng); + for (const Cfg& c : cfgs) { + SamplerParams p = c.p; + const int R = 400; + volatile int sink = 0; + const auto t0 = std::chrono::steady_clock::now(); + for (int r = 0; r < R; ++r) { + p.seed = (uint64_t) r; + sink = sink + strata::metal_engine::sample_row(l.data(), NV, c.pen ? hist.data() : nullptr, + c.pen ? (int) hist.size() : 0, p, c.bias ? bias.data() : nullptr, + (uint64_t) r); + } + const double us = std::chrono::duration(std::chrono::steady_clock::now() - t0).count() / R; + std::printf(" bench %-30s %7.1f us/row\n", c.name, us); + } + } + return bad ? 1 : 0; +} diff --git a/tools/metal_sampling_test.py b/tools/metal_sampling_test.py new file mode 100644 index 00000000000..aa7e9ab3d8f --- /dev/null +++ b/tools/metal_sampling_test.py @@ -0,0 +1,322 @@ +#!/usr/bin/env python3 +"""Sampling test through the --serve protocol (the Mac engine; any engine that has STRATA_DUMP_FIRST_LOGITS). + +The request's sampling keys (temperature, top_p, top_k, min_p, seed, penalty_*, logit_bias) must act as on the CUDA +engine: the chain penalties -> top_k -> top_p -> min_p -> temperature, one Philox 4x32-10 draw at (seed, position), +and greedy when temperature is 0 or absent. + +Engine 1 (no --batch; logit_bias needs it): + 1. reproducible: the same sampled request twice gives the same tokens; another seed gives others; both differ from + the greedy tokens. + 2. draw for draw: many one-token requests with seeds 1..N; each request's first-window logits are dumped + (STRATA_DUMP_FIRST_LOGITS) and a reference sampler written here from the CUDA kernels (src/kernels/cuda/sampler.cu) + picks from them with the same seed and position. Every pick must be the reference's; the empirical distribution + is compared with the reference probabilities (total variation distance, reported). + 3. penalties: greedy with penalty_repeat changes the tokens; the first pick (greedy and sampled) equals the + reference's, which counts the prompt's last penalty_last_n tokens. + 4. logit_bias: banning the greedy first token moves the pick to the reference's; sampled picks never take a banned + token and equal the reference's. +Engine 2: + 5. the sampled request of 1. gives the same tokens in a second engine process. + + python3 tools/metal_sampling_test.py --exe /abs/build-metal/strata --config strata-.json [--draws 100] + python3 tools/metal_sampling_test.py ... --bench # decode speed: greedy vs sampled + +Exit 0 when every check passes, 1 otherwise. STRATA_IQ_MT_MIN=1 (as batch_test.py) for exact tokens. +""" +import argparse, json, math, os, sys, tempfile, time +from collections import Counter +from pathlib import Path + +import numpy as np + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +import batch_test as BT # noqa: E402 + +M32 = 0xFFFFFFFF +F32 = np.float32 + + +# ---------------------------------------------------------------- the reference sampler (src/kernels/cuda/sampler.cu) +def philox_uniform(seed, counter): + """philox_uniform: Philox 4x32-10, counter in (c0, c1), seed in (c2, c3), key (round, 0); 24 bits of c0.""" + c0, c1, c2, c3 = counter & M32, (counter >> 32) & M32, seed & M32, (seed >> 32) & M32 + for i in range(10): + p0, p1 = 0x9E3779B9 * c0, 0xBB67AE85 * c2 + c0, c1, c2, c3 = ((p1 >> 32) ^ c1 ^ i) & M32, p1 & M32, ((p0 >> 32) ^ c3) & M32, p0 & M32 + return (c0 >> 8) * (1.0 / 16777216.0) + + +class Params: + """A request's keys as the engine parses them (strtof: float32), with the engine's defaults.""" + + def __init__(self, keys): + kv = dict(t.split("=", 1) for t in keys.split()) + self.temperature = F32(kv.get("temperature", "0")) + self.top_p = F32(kv.get("top_p", "1")) + self.top_k = int(kv.get("top_k", "20")) + self.min_p = F32(min(max(float(F32(kv.get("min_p", "0"))), 0.0), 1.0)) + self.seed = int(kv.get("seed", "0")) + self.pen_n = max(int(kv.get("penalty_last_n", "0")), 0) + self.rep = F32(kv.get("penalty_repeat", "1")) + self.freq = F32(kv.get("penalty_freq", "0")) + self.pres = F32(kv.get("penalty_present", "0")) + self.bias = {} + if "logit_bias" in kv: + for e in kv["logit_bias"].split(","): + i, v = e.split(":") + self.bias[int(i)] = -np.inf if float(v) <= -100 else F32(v) + + +def scores(logits, p, history): + """logit (+ bias), then the penalties over the last min(penalty_last_n, 4096) history tokens.""" + s = logits.astype(F32).copy() + for i, v in p.bias.items(): + s[i] = F32(s[i] + F32(v)) + if p.pen_n > 0 and history: + for i, c in Counter(history[-min(p.pen_n, 4096):]).items(): + if not 0 <= i < len(s): + continue + x = s[i] + x = F32(x * p.rep) if x <= 0 else F32(x / p.rep) + s[i] = F32(x - F32(F32(F32(c) * p.freq) + F32(p.pres))) + return s + + +def chain(logits, p, history=()): + """(ids, probs) the pick is drawn from: the sampled chain's survivors and their final probabilities, or the + greedy pick with probability 1.""" + s = scores(logits, p, list(history)) + if p.temperature <= 0: + with np.errstate(invalid="ignore"): + best = int(np.argmax(np.where(np.isnan(s), -np.inf, s))) + return [best if s[best] > -np.inf else 0], [1.0] + k = p.top_k if 0 < p.top_k < 64 else 64 + k = min(k, len(s)) + order = np.argsort(-s, kind="stable") # score descending, ties to the lower id + ids = [int(i) for i in order[:k] if s[i] > -np.inf] + lg = [F32(s[i]) for i in ids] + while len(ids) < k: # a round that finds nothing stores (0, -inf) + ids.append(0) + lg.append(F32(-np.inf)) + n_keep = k + mx = max(lg) + if p.top_p < 1: + tot = sum(math.exp(float(v) - float(mx)) for v in lg) + cum, cut = 0.0, k + for i, v in enumerate(lg): + cum += math.exp(float(v) - float(mx)) / tot + if cum >= float(p.top_p): + cut = i + 1 + break + n_keep = max(cut, 1) + if p.min_p > 0: + thresh = F32(lg[0] + F32(np.log(p.min_p))) + for i in range(n_keep): + if lg[i] < thresh: + n_keep = i + break + inv_t = F32(F32(1) / p.temperature) + sc = [F32(v * inv_t) for v in lg[:max(n_keep, 1)]] + smx = max(sc[:max(n_keep, 1)]) + e = [math.exp(float(v) - float(smx)) for v in sc[:n_keep]] + tot = sum(e) + return ids[:n_keep], [x / tot for x in e] + + +def ref_pick(logits, p, counter, history=()): + ids, probs = chain(logits, p, history) + if p.temperature <= 0: + return ids[0] + u = philox_uniform(p.seed, counter) + cum = 0.0 + for i, q in zip(ids, probs): + cum += q + if u < cum: + return i + return ids[-1] + + +# ---------------------------------------------------------------- the engine +class Proto: + def __init__(self, exe, cfg, batch, dump_dir, extra): + env = dict(cfg.get("env") or {}) + env.update({"STRATA_IQ_MT_MIN": "1", "STRATA_DECODE_TIMING": "1"}) + if dump_dir: + env["STRATA_DUMP_FIRST_LOGITS"] = str(Path(dump_dir) / "fl") + self.dump_dir = dump_dir + base = os.environ.get("SAMPLING_TEST_LOG") or os.path.join(tempfile.gettempdir(), "sampling_test_engine") + os.environ["BATCH_TEST_LOG"] = f"{base}.batch{batch}.log" # (BT.Engine's log; one file per engine) + t0 = time.time() + self.eng = BT.Engine(exe, cfg, batch, env, extra.split()) + print(f"engine up in {time.time() - t0:.0f} s (batch {batch})", flush=True) + self.out = self.eng.lines() + + def line(self): + line = next(self.out, None) + if line is None: + raise SystemExit(f"the engine ended - see {self.eng.log_path}") + return line + + def gen(self, ids, max_new, keys=""): + """GEN: (tokens, DONE fields as a dict)""" + self.eng.send(" ".join(x for x in ("GEN", str(max_new), keys, ",".join(map(str, ids))) if x)) + got = [] + while True: + line = self.line() + if line.startswith("T "): + got.append(int(line.split()[1])) + elif line.startswith("ERR"): + raise SystemExit(f"engine: {line}") + elif line.startswith("DONE"): + f = line.split() + return got, {"gen": int(f[1]), "decode_ms": float(f[4]), "acc": int(f[6]), "off": int(f[7])} + + def last_logits(self): + files = sorted(Path(self.dump_dir).glob("fl.*"), key=lambda q: int(q.suffix[1:])) + return np.fromfile(files[-1], dtype=np.float32) + + def quit(self): + self.eng.send("QUIT") + self.eng.p.wait(timeout=180) + + +def wrap(user): + return f"<|im_start|>user\n{user}<|im_end|>\n<|im_start|>assistant\n\n\n\n\n" + + +FAILS = [] + + +def check(ok, what): + print(f" {'PASS' if ok else 'FAIL'} {what}", flush=True) + if not ok: + FAILS.append(what) + + +def draws(P, ids, keys_of, n, history=()): + """n one-token requests (seed s = 1..n): engine picks, reference picks, reference probabilities""" + got, want, mass = [], [], Counter() + pos = len(ids) - 1 # the first generated token is drawn at the last prompt position + for s in range(1, n + 1): + keys = keys_of(s) + toks, _ = P.gen(ids, 1, keys) + lg = P.last_logits() + p = Params(keys) + got.append(toks[0]) + want.append(ref_pick(lg, p, pos, history)) + for i, q in zip(*chain(lg, p, history)): + mass[i] += q / n + emp = Counter(got) + tv = 0.5 * sum(abs(emp[i] / n - mass[i]) for i in set(emp) | set(mass)) + wrong = sum(a != b for a, b in zip(got, want)) + return got, want, wrong, tv, mass + + +def run_tests(a, cfg, tok): + enc = lambda t: tok.encode(t, parse_special=True) + A = enc(wrap("Write a short poem about the sea.")) + P_ = enc(wrap("Pick a random everyday object and describe it in one sentence.")) + R = enc(wrap("Write the word 'hello' twenty times, separated by spaces.")) + C = enc(wrap(BT.QUESTIONS[0])) + D = enc(wrap(BT.QUESTIONS[1])) + K = "temperature=0.8 top_p=0.95 top_k=20" + dump = tempfile.mkdtemp(prefix="sampling_test_") + + print("engine 1: no --batch", flush=True) + P = Proto(a.exe, cfg, 0, dump, a.extra) + g, _ = P.gen(A, a.max_new) + s1, _ = P.gen(A, a.max_new, K + " seed=1") + s1b, _ = P.gen(A, a.max_new, K + " seed=1") + s2, _ = P.gen(A, a.max_new, K + " seed=2") + print(" greedy :", ascii(tok.decode(g)[:100])) + print(" seed 1 :", ascii(tok.decode(s1)[:100])) + print(" seed 2 :", ascii(tok.decode(s2)[:100])) + check(s1 == s1b, "1. the same seed twice gives the same tokens") + check(s1 != s2, "1. another seed gives other tokens") + check(s1 != g and s2 != g, "1. sampled tokens differ from the greedy ones") + + for keys_of, name, n in ((lambda s: f"temperature=1.0 top_k=40 top_p=0.95 seed={s}", "t1.0 k40 p0.95", a.draws), + (lambda s: f"temperature=0.7 top_k=0 top_p=1.0 min_p=0.05 seed={s}", "t0.7 k0 min_p0.05", + a.draws // 2)): + got, want, wrong, tv, mass = draws(P, P_, keys_of, n) + top = sorted(mass.items(), key=lambda x: -x[1])[:5] + print(f" {name}: {len(set(got))} distinct first tokens in {n} draws; reference top: " + + ", ".join(f"{ascii(tok.decode([i]))} {q:.2f}" for i, q in top) + f"; TV distance {tv:.3f}") + check(wrong == 0, f"2. {name}: every pick equals the reference sampler's on the same logits ({wrong} of {n} differ)") + check(len(set(got)) > 1, f"2. {name}: more than one token drawn") + + gp, _ = P.gen(R, a.max_new) + pen = "penalty_last_n=64 penalty_repeat=1.5" + rp, _ = P.gen(R, a.max_new, pen) + lg = P.last_logits() + print(" greedy :", ascii(tok.decode(gp)[:100])) + print(" greedy + penalties:", ascii(tok.decode(rp)[:100])) + check(rp != gp, "3. greedy with penalty_repeat 1.5 changes the tokens") + want = ref_pick(lg, Params(pen), len(R) - 1, R) + check(rp[0] == want, f"3. greedy + penalties: the first pick is the reference's ({rp[0]} vs {want})") + got, want, wrong, tv, _ = draws(P, R, lambda s: f"temperature=1.0 top_k=40 {pen} seed={s}", a.draws // 3, R) + check(wrong == 0, f"3. sampled + penalties: every pick equals the reference's ({wrong} of {len(got)} differ; TV {tv:.3f})") + + ban = g[0] + bg, _ = P.gen(A, 4, f"logit_bias={ban}:-100") + lg = P.last_logits() + check(bg[0] != ban and bg[0] == ref_pick(lg, Params(f"logit_bias={ban}:-100"), len(A) - 1), + f"4. logit_bias bans the greedy first token {ban}: the pick is the reference's ({bg[0]})") + P.gen(P_, 1, "temperature=1.0 top_k=40 seed=1") + top1 = chain(P.last_logits(), Params("temperature=1.0 top_k=40 seed=1"))[0][0] # the most likely first token + got, want, wrong, _, _ = draws(P, P_, lambda s: f"temperature=1.0 top_k=40 seed={s} logit_bias={top1}:-100", a.draws // 4) + check(wrong == 0 and top1 not in got, + f"4. sampled + logit_bias: the banned token {top1} is never drawn, picks equal the reference's ({wrong} differ)") + P.quit() + + print("engine 2", flush=True) + P = Proto(a.exe, cfg, 0, dump, a.extra) + s1c, _ = P.gen(A, a.max_new, K + " seed=1") + check(s1c == s1, "5. the sampled request gives the same tokens in another engine process") + P.quit() + +def bench(a, cfg, tok): + """decode tok/s, greedy vs sampled (the 4 short prompts, MTP drafts on)""" + enc = lambda t: tok.encode(t, parse_special=True) + prompts = [enc(wrap(q)) for q in BT.QUESTIONS[:4]] + S = "temperature=0.7 top_p=0.95 top_k=20 seed=1" + P = Proto(a.exe, cfg, 0, None, a.extra) + res = {} + for mode, keys in (("greedy", ""), ("sampled", S)): + rates, acc, off = [], 0, 0 + for ids in prompts: + toks, d = P.gen(ids, a.bench_tokens, keys) + rates.append(d["gen"] / (d["decode_ms"] / 1000.0)) + acc, off = acc + d["acc"], off + d["off"] + res[f"solo {mode}"] = rates + print(f"solo {mode:7s}: " + " / ".join(f"{r:.1f}" for r in rates) + f" tok/s (mean {sum(rates) / len(rates):.1f}); " + f"drafts accepted {acc}/{off} = {acc / max(off, 1):.2f}", flush=True) + P.quit() + print(json.dumps(res)) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--exe", required=True) + ap.add_argument("--config", required=True) + ap.add_argument("--max-new", type=int, default=48) + ap.add_argument("--draws", type=int, default=100, help="one-token requests per draw-for-draw check") + ap.add_argument("--bench", action="store_true") + ap.add_argument("--bench-tokens", type=int, default=256) + ap.add_argument("--extra", default="") + a = ap.parse_args() + if not os.path.isabs(a.exe): + raise SystemExit("--exe must be an absolute path (the engine runs in the config's cwd)") + cfg = json.loads(Path(a.config).read_text()) + tok = BT.tokenizer(cfg["tokenizer"]) + if a.bench: + bench(a, cfg, tok) + return 0 + run_tests(a, cfg, tok) + print("ALL PASS" if not FAILS else f"{len(FAILS)} FAILED") + return 0 if not FAILS else 1 + + +if __name__ == "__main__": + sys.exit(main()) From f8fa79061be5d7d6dcf315f8d4e950e88e1ed893 Mon Sep 17 00:00:00 2001 From: oceanio Date: Tue, 20 Oct 2026 08:45:27 +0800 Subject: [PATCH 3/3] docs: Mac sampling (MAC.md Sampling: the keys, the order, seeds, the host pick's cost, greedy vs sampled decode) Co-Authored-By: Claude Opus 5.5 --- docs/MAC.md | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/docs/MAC.md b/docs/MAC.md index 56d16318ca2..88b69f494ce 100644 --- a/docs/MAC.md +++ b/docs/MAC.md @@ -145,6 +145,34 @@ Gyro-M runs on the same kernels (checked tensor by tensor) but no Mac tested her checks the decoder against the format definition and the expert kernels against a double-precision reference. `STRATA_METAL_GYRO_TILES=0` turns off the dense-weight tiles. +### Sampling + +A request's sampling settings work as on the PC: `temperature`, `top_p`, `top_k` (1-64; 0 or more than 64 keeps 64), +`min_p`, `seed`, the penalties and `logit_bias`, in the PC engine's order (logit_bias, penalties, top_k, top_p, min_p, +temperature) with one Philox draw per token keyed by the seed and the token's position, so a request with a fixed seed +gives the same text again. Temperature 0, or none, is greedy and gives the same tokens as before. The server's default +sampling (the model card's, when neither the request nor the config sets one) reaches the Mac engine too; send +`temperature: 0` for greedy answers (and `presence_penalty: 0`: the card's presence penalty is added to a +temperature-0 request too, and on the Mac it now applies). + +- The MTP drafts stay greedy, and a draft is kept when the sampled token equals it (as on the PC), so fewer drafts are + kept than with greedy decoding. +- The Mac picks the token on the CPU from the window's logits (unified memory, no copy): 9-46 µs per token row, + 11 µs at temperature 0.7, top_p 0.95, top_k 20 (`metal/tests/engine/host_sampler_parity --bench`, 248,320 logits). + +Measured on the M5 Max (128 GB) with Qwen3.8-Flash-Next IQ3_S at 262K context, 2026-10-20 (`tools/metal_sampling_test.py +--bench`: 4 short prompts with MTP, 256 tokens each; sampled = temperature 0.7, top_p 0.95, top_k 20; 3 runs on a Mac +that had been running tests for about an hour; a cool Mac decodes faster): greedy 76.3 / 83.3 / +78.3 tok/s, sampled 74.8 / 72.9 / 68.8 tok/s (mean of the 4 prompts); drafts kept 0.72 greedy, 0.70 sampled. The +sampled rate is lower because fewer drafts are kept and the sampled texts are other texts. + +Checked with `tools/metal_sampling_test.py` (13 checks, all pass; 0.2.0 without this change fails 11 of them): the same +seed gives the same text twice and in a second engine process, another seed another text; 208 one-token requests drew +exactly the token a reference sampler (written in the test from the PC engine's sampler code) draws from the same +logits, with top_k, top_p, min_p, penalties and a logit_bias. With temperature 0 the tokens equal 0.2.0's (a 28-, a +2,257- and a 16,180-token prompt). Not ported: the opt-in draft schemes `STRATA_SPEC_COUPLED`, `STRATA_SPEC_PROB` and +`STRATA_SPEC_GUMBEL` (the Mac drafts greedily). Not measured: other Macs and models. + ## Measured (M5 Max, 128 GB) **Measured on:** Apple M5 Max, 128 GB, macOS 27.2, Qwen3.8-Flash-Next packs, greedy, 200 generated tokens, Strata engine version