diff --git a/include/strata/core/scr.hpp b/include/strata/core/scr.hpp new file mode 100644 index 00000000..2a48f751 --- /dev/null +++ b/include/strata/core/scr.hpp @@ -0,0 +1,227 @@ +// Suffix cache reuse (SCR): after an edit in the middle of a prompt, the K/V of the text that survived the edit is +// moved to its new positions instead of being read again. Ported from llama-cpp-bonsai's common/scr.cpp. +// +// The moved cells were computed over the text before the edit, so the result approximates a fresh read. The GDN +// layers have no cells to move: their state continues from a conversation checkpoint at the end of the last moved +// run (the plan's `adopt`), and everything after that run is read from it. Strata's indexer pools 4 cells into one +// key and keeps only the pooled keys, so a run starts and ends on a block boundary of the OLD prompt and moves by a +// multiple of 4: up to 3 pad tokens are read before a run to make its shift one (DeepSeek V4.1's recipe pads to an +// even shift the same way). The moved blocks then keep their cells, and their pooled keys only turn by the shift. +#pragma once + +#include +#include +#include +#include +#include + +namespace strata::core { + +struct ScrParams { + int64_t min_span = 256; // shortest run worth moving + int64_t max_runs = 64; // most runs moved per request + int64_t seed = 32; // n-gram length that seeds a match + int64_t block = 4; // the indexer's cells per pooled key: read parts start and end on these +}; + +// the cells old[old_start, old_start + n) reappear at [new_start, new_start + n) of the prompt, after `pad` pad tokens +// that the read inserts right before the run; `at` is where the run lands in the padded (physical) prompt +struct ScrRun { + int64_t new_start = 0, old_start = 0, n = 0; + int64_t pad = 0, at = 0; + int64_t shift() const { return at - old_start; } +}; + +struct ScrPlan { + int64_t n_prefix = 0; // leading tokens equal in both prompts + int64_t restore = 0; // the checkpoint the read restarts from (at or before n_prefix) + int64_t adopt = -1; // old length of the checkpoint the GDN state continues from after the last run + std::vector runs; // in new order, at or after restore + + bool empty() const { return runs.empty(); } + // the prompt position the final read starts at, and the same in the padded prompt + int64_t resume() const { return runs.empty() ? restore : runs.back().new_start + runs.back().n; } + int64_t resume_at() const { return runs.empty() ? restore : runs.back().at + runs.back().n; } + int64_t pads() const { + int64_t n = 0; + for (const ScrRun& r : runs) n += r.pad; + return n; + } + int64_t reused() const { + int64_t n = 0; + for (const ScrRun& r : runs) n += r.n; + return n; + } + std::string str() const { + std::string s = "prefix " + std::to_string(n_prefix) + ", restore " + std::to_string(restore) + ", adopt " + + std::to_string(adopt); + for (const ScrRun& r : runs) + s += ", R " + std::to_string(r.new_start) + " <- " + std::to_string(r.old_start) + " x " + std::to_string(r.n) + + (r.pad ? " pad " + std::to_string(r.pad) : ""); + return s; + } +}; + +namespace scr_detail { +// greedy maximal spans of cur[from, ...) found in old[from, ...), seeded by n-gram hashes, each old token used once +template +std::vector find_spans(const std::vector& old, const std::vector& cur, int64_t from, const ScrParams& p) { + std::vector spans; + const int64_t n_old = (int64_t) old.size(), n_cur = (int64_t) cur.size(), k = p.seed; + if (n_old - from < p.min_span || n_cur - from < p.min_span) return spans; + auto gram = [k](const auto& v, int64_t i) { + uint64_t h = 1469598103934665603ULL; + for (int64_t j = 0; j < k; ++j) h = (h ^ (uint32_t) v[(size_t) (i + j)]) * 1099511628211ULL; + return h; + }; + constexpr size_t max_bucket = 16; // the first few old positions of every seed n-gram + std::unordered_map> index; + index.reserve((size_t) (n_old - from)); + for (int64_t o = from; o + k <= n_old; ++o) { + auto& b = index[gram(old, o)]; + if (b.size() < max_bucket) b.push_back(o); + } + std::vector used((size_t) n_old, false); + for (int64_t q = from; q + k <= n_cur;) { + const auto it = index.find(gram(cur, q)); + int64_t best_o = -1, best_n = 0; + if (it != index.end()) + for (const int64_t o : it->second) { + int64_t n = 0; + while (q + n < n_cur && o + n < n_old && !used[(size_t) (o + n)] && + (int64_t) cur[(size_t) (q + n)] == (int64_t) old[(size_t) (o + n)]) + ++n; + if (n > best_n) { best_n = n; best_o = o; } + } + if (best_n >= p.min_span) { + spans.push_back({q, best_o, best_n}); + std::fill(used.begin() + best_o, used.begin() + best_o + best_n, true); + q += best_n; + } else { + ++q; + } + } + return spans; +} +inline int64_t round_up(int64_t x, int64_t b) { return (x + b - 1) / b * b; } +inline int64_t round_down(int64_t x, int64_t b) { return x / b * b; } +} // namespace scr_detail + +// Plans the reuse of a session that holds old[0, n_old) for the prompt `cur`. `checkpoints`: the old lengths whose +// GDN state can be restored (the live state counts as old.size()). A run starts on a block boundary of the new +// prompt; every run but the last also ends on one, and the last one ends at a checkpoint. The plan always leaves at +// least one token of `cur` to read. +template +ScrPlan scr_plan(const std::vector& old, const std::vector& cur, std::vector checkpoints, + const ScrParams& p) { + using namespace scr_detail; + ScrPlan plan; + const int64_t n_cur = (int64_t) cur.size(), bl = std::max(p.block, 1); + std::sort(checkpoints.begin(), checkpoints.end()); + int64_t a = 0; + while (a < (int64_t) old.size() && a < n_cur && (int64_t) old[(size_t) a] == (int64_t) cur[(size_t) a]) ++a; + plan.n_prefix = a; + for (const int64_t q : checkpoints) + if (q <= a) plan.restore = std::max(plan.restore, q); + auto spans = find_spans(old, cur, a, p); + // a run starts on a block boundary of the old prompt and (but the last) ends on one + for (ScrRun& s : spans) { + const int64_t cut = round_up(s.old_start, bl) - s.old_start; + s.new_start += cut; + s.old_start += cut; + s.n -= cut; + s.n = round_down(s.old_start + s.n, bl) - s.old_start; + } + spans.erase(std::remove_if(spans.begin(), spans.end(), [&](const ScrRun& s) { return s.n < p.min_span; }), + spans.end()); + // the largest runs, back in prompt order + std::sort(spans.begin(), spans.end(), [](const ScrRun& x, const ScrRun& y) { return x.n > y.n; }); + if ((int64_t) spans.size() > p.max_runs) spans.resize((size_t) p.max_runs); + std::sort(spans.begin(), spans.end(), [](const ScrRun& x, const ScrRun& y) { return x.new_start < y.new_start; }); + // the last run ends at a checkpoint, which also ends the reuse: everything after it is read from that state. + // It was cut back to a block boundary, so a checkpoint up to block - 1 cells past it may still match. + while (!spans.empty()) { + ScrRun& last = spans.back(); + int64_t end = -1; + for (const int64_t q : checkpoints) { + const int64_t n = q - last.old_start; + if (n < p.min_span || n >= last.n + bl || last.new_start + n >= n_cur || q <= end) continue; + if (n <= last.n || std::equal(old.begin() + last.old_start + last.n, old.begin() + q, + cur.begin() + last.new_start + last.n, + [](const A& x, const B& y) { return (int64_t) x == (int64_t) y; })) + end = q; + } + if (end >= 0) { + last.n = end - last.old_start; + plan.adopt = end; + break; + } + spans.pop_back(); + } + // the pads: walking the runs in prompt order, each gets what makes its padded position a multiple of the block + // away from its old one + int64_t pads = 0; + for (ScrRun& s : spans) { + s.pad = ((s.old_start - (s.new_start + pads)) % bl + bl) % bl; + pads += s.pad; + s.at = s.new_start + pads; + } + plan.runs = std::move(spans); + return plan; +} + +// The pads an earlier splice read (`pads`: their positions in `live`, sorted) that the prompt `cur` still shares with +// the session: the ones before the last token of `cur` that `live` holds at the same place once the pads are +// skipped. `out` is `cur` with those pads put back, which makes it start with `live` as far as they agree. +template +std::vector scr_carry(const std::vector& cur, const std::vector& live, const std::vector& pads, + int64_t pad_token, std::vector& out) { + size_t p = 0, i = 0, k = 0, kept = 0; + while (p < live.size()) { + if (k < pads.size() && pads[k] == (int64_t) p) { + if ((int64_t) live[p] != pad_token) break; + ++k; + ++p; + continue; + } + if (i + 1 >= cur.size() || (int64_t) live[p] != (int64_t) cur[i]) break; + ++i; + ++p; + kept = k; + } + std::vector carried(pads.begin(), pads.begin() + (long) kept); + out.clear(); + out.reserve(cur.size() + carried.size()); + size_t c = 0; + for (const int64_t at : carried) { + while ((int64_t) out.size() < at) out.push_back(cur[c++]); + out.push_back((B) pad_token); + } + out.insert(out.end(), cur.begin() + (long) c, cur.end()); + return carried; +} + +// The positions of the pads `plan` inserts, in the padded prompt. +inline std::vector scr_pad_positions(const ScrPlan& plan) { + std::vector at; + for (const ScrRun& r : plan.runs) + for (int64_t j = r.pad; j > 0; --j) at.push_back(r.at - j); + return at; +} + +// The prompt the engine reads under `plan`: `cur` with each run's pads inserted right before it. +template +std::vector scr_padded(const std::vector& cur, const ScrPlan& plan, B pad_token) { + std::vector out; + out.reserve(cur.size() + (size_t) plan.pads()); + int64_t from = 0; + for (const ScrRun& r : plan.runs) { + out.insert(out.end(), cur.begin() + from, cur.begin() + r.new_start); + out.insert(out.end(), (size_t) r.pad, pad_token); + from = r.new_start; + } + out.insert(out.end(), cur.begin() + from, cur.end()); + return out; +} + +} // namespace strata::core diff --git a/include/strata/spec/spec_gate.hpp b/include/strata/spec/spec_gate.hpp new file mode 100644 index 00000000..3e4c1c25 --- /dev/null +++ b/include/strata/spec/spec_gate.hpp @@ -0,0 +1,109 @@ +#pragma once + +#include +#include +#include +#include + +namespace strata::spec { + +// Conditional acceptance and wall-clock costs; all storage is bounded by the eight-row window. +struct SpecGate { + double acc[2][8] = {}, acc_w[2][8] = {}; + double t[9] = {}, t_w[9] = {}; + double rho = 0.35, d[2] = {}, d_w[2] = {}; + unsigned picks[2] = {}, since_plain = 0; + + SpecGate() { + for (auto& source : acc) for (double& a : source) a = 0.75; + } + + static void update(double& mean, double& weight, double value, double memory) { + weight = std::min(memory, weight + 1.0); + mean += (value - mean) / weight; + } + + void spec_gate_forward(int rows, double seconds) { + if (rows < 1 || rows > 8 || !std::isfinite(seconds) || seconds <= 0) return; + update(t[rows], t_w[rows], seconds, 8); + since_plain = rows == 1 ? 0 : since_plain + 1; + if (t_w[1] == 0) return; + double xy = 0, xx = 0; + for (int r = 2; r <= 8; ++r) { + const double x = r - 1; + xy += t_w[r] * x * (t[r] / t[1] - 1); + xx += t_w[r] * x * x; + } + if (xx > 0) rho = std::max(0.0, xy / xx); + } + + void spec_gate_draft_cost(int src, double seconds) { + if (src < 0 || src > 1 || !std::isfinite(seconds) || seconds < 0) return; + update(d[src], d_w[src], seconds, 8); + } + + void spec_gate_result(int src, int judged, int accepted) { + if (src < 0 || src > 1) return; + judged = std::clamp(judged, 0, 8); + accepted = std::clamp(accepted, 0, judged); + // Positions beyond the first rejection were not conditionally judged. + for (int j = 0; j < std::min(judged, accepted + 1); ++j) + update(acc[src][j], acc_w[src][j], j < accepted ? 1.0 : 0.0, 16); + } + + double time(int rows) const { + if (t_w[rows] > 0) return t[rows]; + double base = t[1]; + if (t_w[1] == 0) { + base = 1; + for (int r = 2; r <= 8; ++r) + if (t_w[r] > 0) { base = t[r] / (1 + rho * (r - 1)); break; } + } + return base * (1 + rho * (rows - 1)); + } + + // prefix is an already-paid MTP tail. Compare suffix extensions against that window, + // not against plain decoding, and condition their reach probability on the prefix. + int spec_gate_pick(int src, int max_k, const float* mtp_probs = nullptr, + int prefix = 0, const float* prefix_probs = nullptr) { + if (src < 0 || src > 1) return 0; + prefix = std::clamp(prefix, 0, 7); + max_k = std::clamp(max_k, 0, 7 - prefix); + if (max_k == 0) return 0; + ++picks[src]; + if (since_plain >= 64 && prefix == 0) return 0; + double expected = 1, reach = 1; + for (int j = 0; j < prefix; ++j) { + reach *= prefix_probs ? std::clamp((double) prefix_probs[j], 0.0, 1.0) : acc[0][j]; + expected += reach; + } + const double baseline = expected / time(1 + prefix); + double best = baseline; + int pick = 0; + for (int k = 1; k <= max_k; ++k) { + const double a = mtp_probs ? std::clamp((double) mtp_probs[k - 1], 0.0, 1.0) : acc[src][k - 1]; + reach *= a; + expected += reach; + const double cost = time(1 + prefix + k) + k * d[src]; + const double value = expected / cost; + if (value > best) { best = value; pick = k; } + } + // Sparse probes allow a source that was previously poor to recover. + const unsigned interval = t_w[1 + prefix + std::min(max_k, pick + 1)] < 3 ? 4 : 16; + if (picks[src] % interval == 0 && pick < max_k) ++pick; + return pick; + } + + std::string spec_gate_describe() const { + char line[512]; + std::snprintf(line, sizeof line, + "spec gate: T1=%.3fms rho=%.3f T2/T1=%.3f T4/T1=%.3f " + "mtp=%.3f/%.3f/%.3f suffix=%.3f/%.3f/%.3f d=%.3f/%.3fms", + time(1) * 1000, rho, time(2) / time(1), time(4) / time(1), + acc[0][0], acc[0][1], acc[0][2], acc[1][0], acc[1][1], acc[1][2], + d[0] * 1000, d[1] * 1000); + return line; + } +}; + +} // namespace strata::spec diff --git a/metal/CMakeLists.txt b/metal/CMakeLists.txt index d23a26fc..92cd20c7 100644 --- a/metal/CMakeLists.txt +++ b/metal/CMakeLists.txt @@ -120,4 +120,12 @@ target_compile_definitions(strata_metal_kernels PRIVATE 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) +target_link_libraries(strata PRIVATE strata_metal_kernels) + +add_executable(strata_window_cli src/strata_window_cli.cpp) +target_link_libraries(strata_window_cli PRIVATE strata_metal_kernels) +add_executable(convert_k256_q2_0 src/convert_k256_q2_0.cpp) +target_link_libraries(convert_k256_q2_0 PRIVATE mtl_rt) +add_executable(heal_k256_q2_0 src/heal_k256_q2_0.cpp) +target_compile_definitions(heal_k256_q2_0 PRIVATE ACCELERATE_NEW_LAPACK) +target_link_libraries(heal_k256_q2_0 PRIVATE mtl_rt "-framework Accelerate") diff --git a/metal/kernels/dense_mm.metal b/metal/kernels/dense_mm.metal index 45e29ace..ccf06f79 100644 --- a/metal/kernels/dense_mm.metal +++ b/metal/kernels/dense_mm.metal @@ -8,7 +8,12 @@ using namespace mpp::tensor_ops; kernel void act_f16(device const float* x [[buffer(0)]], device half* xh [[buffer(1)]], constant uint& n [[buffer(2)]], uint gid [[thread_position_in_grid]]) { - if (gid < n) xh[gid] = half(x[gid]); + const uint n4 = n >> 2; + if (gid < n4) { + ((device half4*)xh)[gid] = half4(((device const float4*)x)[gid]); + } else if (gid == n4) { + for (uint i = n4 << 2; i < n; ++i) xh[i] = half(x[i]); + } } struct DqArgs { @@ -19,27 +24,46 @@ struct DqArgs { // streams of an output interleaved, for gemm_hc_up's epilogue) }; -// one thread per unit: grid (units per row / 64, n_out), 64 threads +// one thread per unit: grid ((units per row + nth.x - 1) / nth.x, n_out), nth.x threads template kernel void dq_f16(device const uchar* W [[buffer(0)]], device half* Wh [[buffer(1)]], constant DqArgs& a [[buffer(2)]], - uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + uint3 gid [[thread_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort3 nth [[threads_per_threadgroup]]) { constexpr int UNIT = fmt::UNIT, H = UNIT / 2; threadgroup uint tab[fmt::TAB]; - load_tab(tab, tid, 64); - threadgroup_barrier(mem_flags::mem_threadgroup); - const int u = int(tg.x) * 64 + tid, r = int(tg.y); + if (F == T_IQ4_XS || F == T_IQ4_NL) { + if (tid < 64) { + const uint r0 = tid & 15u, sh = (tid >> 4) * 8u; + const half h0 = half(as_type(uchar(((r0 & 8u) ? 0x7159453526190d01ul : 0xf6eaddcfbfad9881ul) >> ((r0 & 7u) * 8u)))); + tab[tid] = as_type(half2(h0, half(as_type(uchar(0xbfad9881u >> sh))))); + tab[tid + 64] = as_type(half2(h0, half(as_type(uchar(0xf6eaddcfu >> sh))))); + tab[tid + 128] = as_type(half2(h0, half(as_type(uchar(0x26190d01u >> sh))))); + tab[tid + 192] = as_type(half2(h0, half(as_type(uchar(0x71594535u >> sh))))); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } else if (fmt::TAB > 1) { + load_tab(tab, tid, nth.x); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const int u = int(gid.x); if (u * UNIT >= int(a.n_in)) return; - // the unit's chunk origin: a 128-wide chunk holds whole units of every format - const int k0 = (u * UNIT) / 128 * 128; + uint r = gid.y, ro = r; + if (a.perm) { + if (fmt::BB * 16 < fmt::QK * 17) { + r = a.perm == 4 ? (ro & 3u) * (a.n_out >> 2) + (ro >> 2) : (ro % a.perm) * (a.n_out / a.perm) + ro / a.perm; + } else { + const uint q = a.n_out / a.perm; + ro = (r % q) * a.perm + r / q; + } + } float w[UNIT]; int p0, p1; - DW::run(W + ulong(r) * a.row_bytes, u, k0, w, p0, p1, tab); - const int q = a.perm ? int(a.n_out / a.perm) : 1; - const int ro = a.perm ? (r % q) * int(a.perm) + r / q : r; - device half* o = Wh + ulong(ro) * a.n_in + k0; - for (int j = 0; j < H; ++j) { - o[p0 + j] = half(w[j]); - o[p1 + j] = half(w[H + j]); + DW::run(W + ulong(r) * a.row_bytes, u, 0, w, p0, p1, tab); + device half* o = Wh + ulong(ro) * a.n_in; +#pragma clang loop unroll(full) + for (int j = 0; j < H; j += 4) { + *(device half4*)(o + p0 + j) = half4(half(w[j]), half(w[j + 1]), half(w[j + 2]), half(w[j + 3])); + *(device half4*)(o + p1 + j) = half4(half(w[H + j]), half(w[H + j + 1]), half(w[H + j + 2]), half(w[H + j + 3])); } } @@ -61,9 +85,17 @@ kernel void gemm_f16(device half* A [[buffer(0)]], device half* B [[buffer(1)]], #define DQ(F) template [[host_name("dq_f16_" #F)]] kernel decltype(dq_f16) dq_f16; #define DQ_(F) DQ(F) -DQ_(T_F32) DQ_(T_BF16) DQ_(T_Q8_0) DQ_(T_Q4_0) DQ_(T_Q5_0) DQ_(T_IQ4_NL) DQ_(T_IQ4_XS) DQ_(T_Q2_0) DQ_(T_IQ3_S) DQ_(T_Q3_K) +DQ_(T_F32) DQ_(T_BF16) DQ_(T_Q8_0) DQ_(T_Q4_0) DQ_(T_Q5_0) DQ_(T_Q5_1) DQ_(T_IQ4_NL) DQ_(T_IQ4_XS) DQ_(T_Q2_0) DQ_(T_IQ3_S) DQ_(T_Q3_K) DQ_(T_Q4_K) DQ_(T_Q5_K) DQ_(T_Q6_K) template [[host_name("gemm_f16_64_64_4")]] kernel decltype(gemm_f16<64, 64, 4>) gemm_f16<64, 64, 4>; template [[host_name("gemm_f16_128_64_4")]] kernel decltype(gemm_f16<128, 64, 4>) gemm_f16<128, 64, 4>; template [[host_name("gemm_f16_64_128_4")]] kernel decltype(gemm_f16<64, 128, 4>) gemm_f16<64, 128, 4>; template [[host_name("gemm_f16_64_32_4")]] kernel decltype(gemm_f16<64, 32, 4>) gemm_f16<64, 32, 4>; +template [[host_name("gemm_f16_128_128_4")]] kernel decltype(gemm_f16<128, 128, 4>) gemm_f16<128, 128, 4>; +template [[host_name("gemm_f16_64_64_2")]] kernel decltype(gemm_f16<64, 64, 2>) gemm_f16<64, 64, 2>; +template [[host_name("gemm_f16_64_128_2")]] kernel decltype(gemm_f16<64, 128, 2>) gemm_f16<64, 128, 2>; +template [[host_name("gemm_f16_16_64_2")]] kernel decltype(gemm_f16<16, 64, 2>) gemm_f16<16, 64, 2>; +template [[host_name("gemm_f16_32_64_2")]] kernel decltype(gemm_f16<32, 64, 2>) gemm_f16<32, 64, 2>; +template [[host_name("gemm_f16_32_128_2")]] kernel decltype(gemm_f16<32, 128, 2>) gemm_f16<32, 128, 2>; +template [[host_name("gemm_f16_64_32_2")]] kernel decltype(gemm_f16<64, 32, 2>) gemm_f16<64, 32, 2>; +template [[host_name("gemm_f16_16_96_2")]] kernel decltype(gemm_f16<16, 96, 2>) gemm_f16<16, 96, 2>; diff --git a/metal/kernels/gdn.metal b/metal/kernels/gdn.metal index 516d89a4..a779becf 100644 --- a/metal/kernels/gdn.metal +++ b/metal/kernels/gdn.metal @@ -27,6 +27,10 @@ struct GdnArgs { // `plane` floats between planes) and the other planes and the per-16 sums are written too }; +constant float gdn_q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, + 0.00006103515625f}; +constant float gdn_p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + // The kernel is launched with IB * 128 threads (512 for gdn_step, 1024 for gdn_step8). Without the attribute the compiler // is free to use as many registers as the 128-row state needs and a pipeline can end up with // maxTotalThreadsPerThreadgroup below the launch size: an M2 Max gave 448 for the GDN_NTOK >= 2 specializations, and @@ -40,15 +44,17 @@ kernel void gdn_step_t(device float* state [[buffer(0)]], device float* hist [[b device const float* ssm_a [[buffer(8)]], device const float* gamma [[buffer(9)]], device float* y [[buffer(10)]], constant GdnArgs& a [[buffer(11)]], device float* hist_out [[buffer(12)]], + device float* state_end [[buffer(13)]], device float* hist_end [[buffer(14)]], uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { constexpr int S = 128; threadgroup float qs[S], ks[S], vs[S]; constexpr int RPT = S / IB; // state rows per thread - threadgroup float red[IB][S]; - threadgroup float part[32]; + // kv and o partials in their own buffers, the q/k and the output norm partials in theirs: no barrier between a + // token's phases for a buffer to be free again + threadgroup float red[IB][S], red2[IB - 1][S]; + threadgroup float part[8], part2[4]; threadgroup float gb[2]; - threadgroup float ys[S]; const int h = int(tg.x), hk = h % int(a.h_k), hv = int(a.h_v); const int cb = sg & 3, ib = sg >> 2; const int j = 32 * cb + lane, i0 = RPT * ib; @@ -60,109 +66,137 @@ kernel void gdn_step_t(device float* state [[buffer(0)]], device float* hist [[b else if (tid < 256) ch = int(a.h_k) * S + hk * S + (tid - 128); else if (tid < 384) ch = 2 * int(a.h_k) * S + h * S + (tid - 256); float hst[3] = {0.0f, 0.0f, 0.0f}, cw[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + // the first token's inputs are issued with the state's loads (their latency is not paid again after it) + float in = ch >= 0 ? qkv[ch] : 0.0f; if (ch >= 0) { + #pragma clang loop unroll(full) for (int k = 0; k < 3; ++k) hst[k] = hist[ch * 3 + k]; - for (int k = 0; k < 4; ++k) cw[k] = conv_w[ch * 4 + k]; + const float4 cwv = *(device const float4*)(conv_w + ch * 4); + cw[0] = cwv.x; cw[1] = cwv.y; cw[2] = cwv.z; cw[3] = cwv.w; } + const float gam = ib == 0 ? gamma[j] : 0.0f; + const float dth = tid == 384 ? dt[h] : 0.0f; + const float ssm_ah = tid == 384 ? ssm_a[h] : 0.0f; float st[RPT]; - for (int r = 0; r < RPT; ++r) st[r] = state[((i0 + r) * hv + h) * S + j]; + const int st_stride = hv * S; + const int st_off = (i0 * hv + h) * S + j; + #pragma clang loop unroll(full) + for (int r = 0; r < RPT; ++r) st[r] = (state + st_off)[r * st_stride]; - const int commit_at = a.commit == 1u ? GDN_NTOK : (a.commit >> 8) ? int(a.commit & 0xFFu) : 0; + const int commit_at = a.commit == 1u ? GDN_NTOK : ((a.commit >> 8) == 1u) ? int(a.commit & 0xFFu) : 0; for (int t = 0; t < GDN_NTOK; ++t) { + // this token's z and the next token's conv input, ahead of the barriers that need them + const float zz = ib == 0 ? z[t * hv * S + h * S + j] : 0.0f; + const float in_next = ch >= 0 && t + 1 < GDN_NTOK ? qkv[(t + 1) * int(a.channels) + ch] : 0.0f; // conv + SiLU float cv = 0.0f; if (ch >= 0) { - const float in = qkv[t * int(a.channels) + ch]; const float vals[4] = {hst[0], hst[1], hst[2], in}; float s = 0.0f; + #pragma clang loop unroll(full) for (int k = 0; k < 4; ++k) s = fma(vals[k], cw[k], s); s = s + 0.0f; cv = s / (1.0f + exp(-s)); hst[0] = hst[1]; hst[1] = hst[2]; hst[2] = in; + if (tid >= 256 || h < int(a.h_k)) { + if (t + 1 == commit_at || (a.commit == 2u && t == 0)) { + #pragma clang loop unroll(full) + for (int k = 0; k < 3; ++k) hist_out[ch * 3 + k] = hst[k]; + } + if ((a.commit == 2u || a.commit == 3u) && t + 1 == GDN_NTOK) { + #pragma clang loop unroll(full) + for (int k = 0; k < 3; ++k) hist_end[ch * 3 + k] = hst[k]; + } + } } - // L2 norms of q (threads 0..127 = simdgroups 0..3) and k (128..255 = 4..7) - const float ssq = simd_sum(cv * cv); - if (lane == 0) part[sg] = ssq; - if (tid == 0) { - const float v = alpha[t * hv + h] + dt[h]; + in = in_next; + // raw q/k/v to shared memory + L2 norm partials of q (sg 0..3) and k (sg 4..7); gate and beta on idle lanes + if (tid < 256) { + const float ssq = simd_sum(cv * cv); + if (lane == 0) part[sg] = ssq; + if (tid < 128) qs[tid] = cv; + else ks[tid - 128] = cv; + } else if (tid < 384) { + vs[tid - 256] = cv; + } else if (tid == 384) { + const float v = alpha[t * hv + h] + dth; // log1p(e^v) (Kahan: log(u) * e / (u - 1), u = 1 + e; exact where 1 + e rounds to 1) const float e = exp(v), u = 1.0f + e; const float sp = v > 20.0f ? v : (u == 1.0f ? e : log(u) * e / (u - 1.0f)); - gb[0] = exp(sp * ssm_a[h]); + gb[0] = exp(sp * ssm_ah); + } else if (tid == 416) { gb[1] = 1.0f / (1.0f + exp(-beta[t * hv + h])); } threadgroup_barrier(mem_flags::mem_threadgroup); - if (tid < 256) { - const int base = tid < 128 ? 0 : 4; - const float tot = (part[base] + part[base + 1]) + (part[base + 2] + part[base + 3]); - const float rs = rsqrt(tot / float(S) + a.eps / float(S)); - const float nv = (rs * cv) * scale; - if (tid < 128) qs[tid] = nv; - else ks[tid - 128] = nv; - } else if (tid < 384) { - vs[tid - 256] = cv; - } - threadgroup_barrier(mem_flags::mem_threadgroup); - const float g = gb[0], be = gb[1]; + const float tot_k = (part[4] + part[5]) + (part[6] + part[7]); + const float rsk = rsqrt(tot_k / float(S) + a.eps / float(S)); + // the simdgroup's RPT rows of k and q one a lane, normalized on the fly and broadcast from the lane that holds them + const int qk_idx = i0 + (lane & (RPT - 1)); + const float kl = (rsk * ks[qk_idx]) * scale; // kv_j float kv = 0.0f; - for (int r = 0; r < RPT; ++r) kv = fma(st[r], ks[i0 + r], kv); + #pragma clang loop unroll(full) + for (int r = 0; r < RPT; ++r) kv = fma(st[r], simd_broadcast(kl, ushort(r)), kv); red[ib][j] = kv; + const float tot_q = (part[0] + part[1]) + (part[2] + part[3]); + const float rsq = rsqrt(tot_q / float(S) + a.eps / float(S)); + const float ql = (rsq * qs[qk_idx]) * scale; + const float g = gb[0], be = gb[1], vj = vs[j]; threadgroup_barrier(mem_flags::mem_threadgroup); kv = 0.0f; + #pragma clang loop unroll(full) for (int k = 0; k < IB; ++k) kv += red[k][j]; - const float delta = (vs[j] - g * kv) * be; + const float delta = (vj - g * kv) * be; float o = 0.0f; + #pragma clang loop unroll(full) for (int r = 0; r < RPT; ++r) { - st[r] = fma(g, st[r], ks[i0 + r] * delta); - o = fma(st[r], qs[i0 + r], o); + st[r] = fma(g, st[r], simd_broadcast(kl, ushort(r)) * delta); + o = fma(st[r], simd_broadcast(ql, ushort(r)), o); + } + if (ib > 0) red2[ib - 1][j] = o; + if (t + 1 == commit_at || (a.commit == 2u && t == 0)) { + device float* st_base = state + st_off; + #pragma clang loop unroll(full) + for (int r = 0; r < RPT; ++r) st_base[r * st_stride] = st[r]; + } + if ((a.commit == 2u || a.commit == 3u) && t + 1 == GDN_NTOK) { + device float* st_base = state_end + st_off; + #pragma clang loop unroll(full) + for (int r = 0; r < RPT; ++r) st_base[r * st_stride] = st[r]; } threadgroup_barrier(mem_flags::mem_threadgroup); - red[ib][j] = o; - threadgroup_barrier(mem_flags::mem_threadgroup); - // the readout and the output norm: column j's o (every i-block thread of the column holds the same value) - o = 0.0f; - for (int k = 0; k < IB; ++k) o += red[k][j]; - o = o * scale; - const float osq = ib == 0 ? simd_sum(o * o) : 0.0f; - if (ib == 0 && lane == 0) part[cb] = osq; + // the readout and the output norm: column j's o (only ib == 0 computes readout and writes y) + if (ib == 0) { + #pragma clang loop unroll(full) + for (int k = 0; k < IB - 1; ++k) o += red2[k][j]; + o = o * scale; + const float osq = simd_sum(o * o); + if (lane == 0) part2[cb] = osq; + } threadgroup_barrier(mem_flags::mem_threadgroup); if (ib == 0) { - const float tot = (part[0] + part[1]) + (part[2] + part[3]); + const float tot = (part2[0] + part2[1]) + (part2[2] + part2[3]); const float rs = rsqrt(tot / float(S) + a.eps); - const float zz = z[t * hv * S + h * S + j]; - const float yv = ((rs * o) * gamma[j]) * (1.0f / (1.0f + exp(-zz))); + const float yv = ((rs * o) * gam) * (1.0f / (1.0f + exp(-zz))); y[t * hv * S + h * S + j] = yv; if (a.plane) { - const float q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, - 0.00006103515625f}; - const float p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; device float* raw = y + t * hv * S + h * S + j; const int jj = j & 15; - raw[a.plane] = yv * q2s[jj & 7]; - raw[2 * a.plane] = yv * p4s[jj & 3]; + raw[a.plane] = yv * gdn_q2s[jj & 7]; + raw[2 * a.plane] = yv * gdn_p4s[jj & 3]; raw[3 * a.plane] = (jj & 1) ? yv * 0.00390625f : yv; - ys[j] = yv; + // the per-16 sums (even and odd columns apart, in column order) from the simdgroup's lanes + float s0 = 0.0f, s1 = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < 16; k += 2) { + s0 += simd_shuffle_down(yv, ushort(k)); + s1 += simd_shuffle_down(yv, ushort(k + 1)); + } + if (jj == 0) y[4 * a.plane + t * (hv * S / 16) + (h * S + j) / 16] = s0 + s1; } } - if (t + 1 == commit_at) { - // the state and the conv history after token t (registers: st after this token's update, hst after its - // shift); the state is read only before the loop, so in place is safe. q/k history: the threadgroups - // h < h_k own their q/k head (h == hk); v history: every head its own - to hist_out, NOT in place: a q/k - // channel is read by the three threadgroups h % h_k at their start, and a threadgroup may start after its - // owner has finished (seen on the 10-core M4); the caller swaps the two - for (int r = 0; r < RPT; ++r) state[((i0 + r) * hv + h) * S + j] = st[r]; - if (ch >= 0 && (tid >= 256 || h < int(a.h_k))) - for (int k = 0; k < 3; ++k) hist_out[ch * 3 + k] = hst[k]; - } - threadgroup_barrier(mem_flags::mem_threadgroup); - if (a.plane && ib == 0 && (j & 15) == 0) { - float s0 = 0.0f, s1 = 0.0f; - for (int k = 0; k < 16; k += 2) { s0 += ys[j + k]; s1 += ys[j + k + 1]; } - y[4 * a.plane + t * (hv * S / 16) + (h * S + j) / 16] = s0 + s1; - } } } diff --git a/metal/kernels/gdn_prompt.metal b/metal/kernels/gdn_prompt.metal index b2d530b1..b9bd508c 100644 --- a/metal/kernels/gdn_prompt.metal +++ b/metal/kernels/gdn_prompt.metal @@ -10,6 +10,7 @@ // (= S_new^T q / sqrt(128)), S = g S + k delta^T; q and k staged per TB tokens in threadgroup memory // gdn_post per (token, head): y = o * rsqrt(mean(o^2) + eps) * gamma * sigmoid(z) -> fp16 (the out GEMM's operand) // and, with `f32`, fp32 +// On the M5 the scan is gdn_wy_prep + gdn_wy_scan instead (the chunked WY form on the matrix units, at the end). // // State layout as Strata's: S[(i * h_v + h) * 128 + j]. Value head h reads q/k head h % h_k. // The kernels are new entries (nothing of gdn.metal changes). @@ -25,67 +26,161 @@ struct GdnPArgs { static inline float gp_conv(device const float* qkv, device const float* hist, device const float* cw, int ch, int t, int channels) { + const float4 w = *(device const float4*)(cw + ch * 4); float s = 0.0f; +#pragma clang loop unroll(full) for (int k = 0; k < 4; ++k) { const int src = t - 3 + k; const float x = src >= 0 ? qkv[src * channels + ch] : hist[ch * 3 + (src + 3)]; - s = fma(x, cw[ch * 4 + k], s); + s = fma(x, w[k], s); } s = s + 0.0f; return s / (1.0f + exp(-s)); } -// grid (h_k + h_v / 2, T), 256 threads. x < h_k: q (threads 0..127) and k (128..255) of key head x; else value heads -// 2 (x - h_k) and + 1. +// grid ((h_k + h_v) / 8, T) = (8, T) or (h_k + h_v / 2, T) = (40, T), 256 threads: when grid_x == (h_k + h_v) / 8, +// each simdgroup (32 lanes) owns one key head (x < h_k / 8: q, k and kq in registers) or one value head (else: vc and +// gb), with zero threadgroup memory and zero barriers; grid_x == h_k + h_v / 2 keeps the legacy grid. kernel void gdn_pre(device const float* qkv [[buffer(0)]], device const float* hist [[buffer(1)]], device const float* conv_w [[buffer(2)]], device const float* alpha [[buffer(3)]], device const float* beta [[buffer(4)]], device const float* dt [[buffer(5)]], device const float* ssm_a [[buffer(6)]], device float* qn [[buffer(7)]], device float* kn [[buffer(8)]], device float* vc [[buffer(9)]], device float2* gb [[buffer(10)]], device float* kq [[buffer(11)]], device float* hist_out [[buffer(12)]], constant GdnPArgs& a [[buffer(13)]], - uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], - ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + uint3 tg [[threadgroup_position_in_grid]], uint3 tgs [[threadgroups_per_grid]], + ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { constexpr int S = 128; - threadgroup float part[8]; - threadgroup float prod[S]; const int x = int(tg.x), t = int(tg.y), hk_n = int(a.h_k), hv = int(a.h_v), C = int(a.channels), T = int(a.n_tok); - int ch; - if (x < hk_n) ch = tid < 128 ? x * S + tid : hk_n * S + x * S + (tid - 128); - else ch = 2 * hk_n * S + (2 * (x - hk_n) + (tid >> 7)) * S + (tid & 127); - const float cv = gp_conv(qkv, hist, conv_w, ch, t, C); - if (t == T - 1) - for (int k = 0; k < 3; ++k) { - const int src = T - 3 + k; - hist_out[ch * 3 + k] = src >= 0 ? qkv[src * C + ch] : hist[ch * 3 + (src + 3)]; + const bool packed = int(tgs.x) * 8 == hk_n + hv; + // alpha and beta as the halves of one [T][2 h_v] GEMM (beta == alpha + h_v), else each [T][h_v] + const int ab = beta == alpha + hv ? 2 * hv : hv; + const int k_tgs = packed ? hk_n / 8 : hk_n; + if (x < k_tgs) { + const int kh = packed ? x * 8 + int(sg) : (sg == 0 ? x : hk_n); + if (kh >= hk_n) return; + float cv_q[4], cv_k[4]; + if (t >= 3) { + device const float* pq = qkv + (t - 3) * C + kh * S + int(lane); + device const float* pk = qkv + (t - 3) * C + (hk_n + kh) * S + int(lane); +#pragma clang loop unroll(full) + for (int m = 0; m < 4; ++m) { + const int d = int(lane) + m * 32; + const int ch_q = kh * S + d, ch_k = (hk_n + kh) * S + d; + const float4 wq = *(device const float4*)(conv_w + ch_q * 4); + const float4 wk = *(device const float4*)(conv_w + ch_k * 4); + device const float* pqm = pq + m * 32; + device const float* pkm = pk + m * 32; + float sq = fma(pqm[0], wq.x, 0.0f); + sq = fma(pqm[C], wq.y, sq); + sq = fma(pqm[2 * C], wq.z, sq); + sq = fma(pqm[3 * C], wq.w, sq) + 0.0f; + cv_q[m] = sq / (1.0f + exp(-sq)); + float sk = fma(pkm[0], wk.x, 0.0f); + sk = fma(pkm[C], wk.y, sk); + sk = fma(pkm[2 * C], wk.z, sk); + sk = fma(pkm[3 * C], wk.w, sk) + 0.0f; + cv_k[m] = sk / (1.0f + exp(-sk)); + } + if (t == T - 1) + for (int m = 0; m < 4; ++m) { + const int d = int(lane) + m * 32; + const int ch_q = kh * S + d, ch_k = (hk_n + kh) * S + d; + for (int k = 0; k < 3; ++k) { + hist_out[ch_q * 3 + k] = pq[m * 32 + (k + 1) * C]; + hist_out[ch_k * 3 + k] = pk[m * 32 + (k + 1) * C]; + } + } + } else { +#pragma clang loop unroll(full) + for (int m = 0; m < 4; ++m) { + const int d = int(lane) + m * 32; + const int ch_q = kh * S + d, ch_k = (hk_n + kh) * S + d; + cv_q[m] = gp_conv(qkv, hist, conv_w, ch_q, t, C); + cv_k[m] = gp_conv(qkv, hist, conv_w, ch_k, t, C); + } + if (t == T - 1) + for (int m = 0; m < 4; ++m) { + const int d = int(lane) + m * 32; + const int ch_q = kh * S + d, ch_k = (hk_n + kh) * S + d; + for (int k = 0; k < 3; ++k) { + const int src = T - 3 + k; + hist_out[ch_q * 3 + k] = src >= 0 ? qkv[src * C + ch_q] : hist[ch_q * 3 + (src + 3)]; + hist_out[ch_k * 3 + k] = src >= 0 ? qkv[src * C + ch_k] : hist[ch_k * 3 + (src + 3)]; + } + } } - if (x < hk_n) { - const float ssq = simd_sum(cv * cv); - if (lane == 0) part[sg] = ssq; - threadgroup_barrier(mem_flags::mem_threadgroup); - const int base = tid < 128 ? 0 : 4; - const float tot = (part[base] + part[base + 1]) + (part[base + 2] + part[base + 3]); - const float rs = rsqrt(tot / float(S) + a.eps / float(S)); - const float nv = (rs * cv) * (1.0f / sqrt(float(S))); - if (tid < 128) { - qn[(t * hk_n + x) * S + tid] = nv; - prod[tid] = nv; + const float4 sq = simd_sum(float4(cv_q[0] * cv_q[0], cv_q[1] * cv_q[1], cv_q[2] * cv_q[2], cv_q[3] * cv_q[3])); + const float4 sk = simd_sum(float4(cv_k[0] * cv_k[0], cv_k[1] * cv_k[1], cv_k[2] * cv_k[2], cv_k[3] * cv_k[3])); + const float tot_q = (sq.x + sq.y) + (sq.z + sq.w), tot_k = (sk.x + sk.y) + (sk.z + sk.w); + const float rs_q = rsqrt(tot_q / float(S) + a.eps / float(S)); + const float rs_k = rsqrt(tot_k / float(S) + a.eps / float(S)); + const float scale = 1.0f / sqrt(float(S)); + const int base = (t * hk_n + kh) * S + int(lane); + float4 prod; +#pragma clang loop unroll(full) + for (int m = 0; m < 4; ++m) { + const float nq = (rs_q * cv_q[m]) * scale, nk = (rs_k * cv_k[m]) * scale; + qn[base + m * 32] = nq; + kn[base + m * 32] = nk; + prod[m] = nk * nq; } - threadgroup_barrier(mem_flags::mem_threadgroup); - if (tid >= 128) { - kn[(t * hk_n + x) * S + tid - 128] = nv; - const float p = simd_sum(nv * prod[tid - 128]); - if (lane == 0) part[sg] = p; // simdgroups 4..7 (q's slots 0..3 are read already) + const float4 pv = simd_sum(prod); + if (lane == 0) kq[t * hk_n + kh] = (pv.x + pv.y) + (pv.z + pv.w); + } else if (packed) { + const int h = (x - k_tgs) * 8 + int(sg); + if (t >= 3) { + const int d = int(lane) * 4, ch = (2 * hk_n + h) * S + d; + device const float4* p = (device const float4*)(qkv + (t - 3) * C + ch); + device const float4* w = (device const float4*)(conv_w + ch * 4); + const float4 w0 = w[0], w1 = w[1], w2 = w[2], w3 = w[3]; + float4 cv = fma(p[0], float4(w0.x, w1.x, w2.x, w3.x), float4(0.0f)); + cv = fma(p[C / 4], float4(w0.y, w1.y, w2.y, w3.y), cv); + cv = fma(p[C / 2], float4(w0.z, w1.z, w2.z, w3.z), cv); + cv = fma(p[3 * C / 4], float4(w0.w, w1.w, w2.w, w3.w), cv); + cv = cv + 0.0f; + cv = cv / (1.0f + exp(-cv)); + *((device float4*)(vc + (t * hv + h) * S + d)) = cv; + if (t == T - 1) + for (int k = 0; k < 3; ++k) { + const float4 v = p[(k + 1) * C / 4]; + for (int m = 0; m < 4; ++m) hist_out[(ch + m) * 3 + k] = v[m]; + } + } else { +#pragma clang loop unroll(full) + for (int m = 0; m < 4; ++m) { + const int d = int(lane) + m * 32; + const int ch = 2 * hk_n * S + h * S + d; + const float cv = gp_conv(qkv, hist, conv_w, ch, t, C); + if (t == T - 1) + for (int k = 0; k < 3; ++k) { + const int src = T - 3 + k; + hist_out[ch * 3 + k] = src >= 0 ? qkv[src * C + ch] : hist[ch * 3 + (src + 3)]; + } + vc[(t * hv + h) * S + d] = cv; + } + } + if (lane == 0) { + const float v = alpha[t * ab + h] + dt[h]; + const float e = exp(v), u = 1.0f + e; + const float sp = v > 20.0f ? v : (u == 1.0f ? e : log(u) * e / (u - 1.0f)); + gb[t * hv + h] = float2(exp(sp * ssm_a[h]), 1.0f / (1.0f + exp(-beta[t * ab + h]))); } - threadgroup_barrier(mem_flags::mem_threadgroup); - if (tid == 0) kq[t * hk_n + x] = (part[4] + part[5]) + (part[6] + part[7]); } else { - const int h = 2 * (x - hk_n) + (tid >> 7); - vc[(t * hv + h) * S + (tid & 127)] = cv; - if ((tid & 127) == 0) { - const float v = alpha[t * hv + h] + dt[h]; + const int h = 2 * (x - k_tgs) + (tid >> 7), d = tid & 127; + const int ch = 2 * hk_n * S + h * S + d; + const float cv = gp_conv(qkv, hist, conv_w, ch, t, C); + if (t == T - 1) + for (int k = 0; k < 3; ++k) { + const int src = T - 3 + k; + hist_out[ch * 3 + k] = src >= 0 ? qkv[src * C + ch] : hist[ch * 3 + (src + 3)]; + } + vc[(t * hv + h) * S + d] = cv; + if (d == 0) { + const float v = alpha[t * ab + h] + dt[h]; const float e = exp(v), u = 1.0f + e; const float sp = v > 20.0f ? v : (u == 1.0f ? e : log(u) * e / (u - 1.0f)); - gb[t * hv + h] = float2(exp(sp * ssm_a[h]), 1.0f / (1.0f + exp(-beta[t * hv + h]))); + gb[t * hv + h] = float2(exp(sp * ssm_a[h]), 1.0f / (1.0f + exp(-beta[t * ab + h]))); } } } @@ -156,55 +251,60 @@ constant int GDN_PROBE [[function_constant(6)]]; constant int gdn_probe = is_function_constant_defined(GDN_PROBE) ? GDN_PROBE : 0; // the rows of a column split over IB simdgroups (128 / IB rows a lane, fewer registers): grid (h_v * 4), IB * 32 -// threads; lane = column (32 columns a threadgroup); per token the IB partial dots are exchanged through threadgroup -// memory (double-buffered: one barrier a token) +// threads; lane = column (32 columns a threadgroup); when TB == 0 (single-pass), each simdgroup loads its R rows of q +// and k across its lanes and broadcasts them with simd_broadcast (one initial dot pass at t = 0, one barrier a token, +// only red[2][IB][32] in threadgroup memory); TB > 0 keeps the TB-token threadgroup-staged reference. template -kernel void gdn_scan_r(device float* state [[buffer(0)]], device const float4* qn [[buffer(1)]], - device const float4* kn [[buffer(2)]], device const float* vc [[buffer(3)]], +kernel void gdn_scan_r(device float* state [[buffer(0)]], device const float* qn [[buffer(1)]], + device const float* kn [[buffer(2)]], device const float* vc [[buffer(3)]], device const float2* gb [[buffer(4)]], device const float* kq [[buffer(5)]], device float* o [[buffer(6)]], constant GdnPArgs& a [[buffer(7)]], uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { - constexpr int S = 128, R = S / IB, R4 = R / 4; - threadgroup float4 qs[TB][32], ks[TB][32]; + constexpr int S = 128, R = S / IB; threadgroup float2 red[2][IB][32]; - threadgroup float vs[TB][32]; - threadgroup float4 gk[TB]; // gate, beta, kq - const int h = int(tg.x) / 4, hk = h % int(a.h_k), hv = int(a.h_v), hkn = int(a.h_k); + const int g_idx = int(tg.x) / 4, hv = int(a.h_v), hkn = int(a.h_k); + const int hk = (hv == 3 * hkn) ? (g_idx / 3) : (g_idx % hkn); + const int h = (hv == 3 * hkn) ? (hk + (g_idx % 3) * hkn) : g_idx; const int j = (int(tg.x) % 4) * 32 + lane, i0 = sg * R; const int T = int(a.n_tok); const float scale = 1.0f / sqrt(float(S)); + const int rs = hv * S, s0 = (i0 * hv + h) * S + j; float st[R]; #pragma clang loop unroll(full) - for (int i = 0; i < R; ++i) st[i] = state[((i0 + i) * hv + h) * S + j]; - for (int t0 = 0; t0 < T; t0 += TB) { - const int nb = min(TB, T - t0); - threadgroup_barrier(mem_flags::mem_threadgroup); - for (int idx = tid; idx < nb * 32; idx += IB * 32) { - const int tt = idx >> 5, w = idx & 31; - qs[tt][w] = qn[((t0 + tt) * hkn + hk) * 32 + w]; - ks[tt][w] = kn[((t0 + tt) * hkn + hk) * 32 + w]; - vs[tt][w] = vc[((t0 + tt) * hv + h) * S + (j - lane) + w]; - if (w == 0) { - const float2 g_b = gb[(t0 + tt) * hv + h]; - gk[tt] = float4(g_b.x, g_b.y, kq[(t0 + tt) * hkn + hk], 0.0f); - } + for (int i = 0; i < R; ++i) st[i] = (state + s0)[i * rs]; + if (TB == 0) { + constexpr int C = (R + 31) / 32, R1 = R / C; + const ushort lr = lane & (R1 - 1); + const int ks = hkn * S, k0 = hk * S + i0 + int(lr), v0 = h * S + j; + float kl[C], ql[C]; +#pragma clang loop unroll(full) + for (int c = 0; c < C; ++c) { + kl[c] = kn[k0 + c * 32]; + ql[c] = qn[k0 + c * 32]; } - threadgroup_barrier(mem_flags::mem_threadgroup); - for (int tt = 0; tt < nb; ++tt) { - const int t = t0 + tt, slot = t & 1; - float ka = 0.0f, qa = 0.0f; + float2 g_b = gb[h]; + float v = vc[v0]; + float ka = 0.0f, qa = 0.0f; #pragma clang loop unroll(full) - for (int w = 0; w < R4; ++w) { - const float4 k4 = ks[tt][i0 / 4 + w], q4 = qs[tt][i0 / 4 + w]; - ka = fma(st[4 * w], k4.x, ka); ka = fma(st[4 * w + 1], k4.y, ka); - ka = fma(st[4 * w + 2], k4.z, ka); ka = fma(st[4 * w + 3], k4.w, ka); - qa = fma(st[4 * w], q4.x, qa); qa = fma(st[4 * w + 1], q4.y, qa); - qa = fma(st[4 * w + 2], q4.z, qa); qa = fma(st[4 * w + 3], q4.w, qa); + for (int c = 0; c < C; ++c) +#pragma clang loop unroll(full) + for (int r = 0; r < R1; ++r) { + ka = fma(st[c * 32 + r], simd_broadcast(kl[c], ushort(r)), ka); + qa = fma(st[c * 32 + r], simd_broadcast(ql[c], ushort(r)), qa); } - const float4 g_b = gk[tt]; - const float kqv = g_b.z; - const float v = vs[tt][lane]; + for (int t = 0; t < T - 1; ++t) { + const int slot = t & 1, tn = t + 1; + const int kn_off = tn * ks + k0, vn_off = tn * rs + v0; + float kn_l[C], qn_l[C]; +#pragma clang loop unroll(full) + for (int c = 0; c < C; ++c) { + kn_l[c] = kn[kn_off + c * 32]; + qn_l[c] = qn[kn_off + c * 32]; + } + const float2 g_b_n = gb[tn * hv + h]; + const float v_n = vc[vn_off]; + const float kqv = sg == 0 ? kq[t * hkn + hk] : 0.0f; float2 s = float2(ka, qa); if (gdn_probe != 1) { red[slot][sg][lane] = s; @@ -215,20 +315,120 @@ kernel void gdn_scan_r(device float* state [[buffer(0)]], device const float4* q } const float g = g_b.x; const float delta = (v - g * s.x) * g_b.y; - if (sg == 0) o[(t * hv + h) * S + j] = fma(g, s.y, delta * kqv) * scale; + if (sg == 0) o[t * rs + v0] = fma(g, s.y, delta * kqv) * scale; + ka = 0.0f; + qa = 0.0f; if (gdn_probe != 2) +#pragma clang loop unroll(full) + for (int c = 0; c < C; ++c) +#pragma clang loop unroll(full) + for (int r = 0; r < R1; ++r) { + const float kr = simd_broadcast(kl[c], ushort(r)); + const float knr = simd_broadcast(kn_l[c], ushort(r)), qnr = simd_broadcast(qn_l[c], ushort(r)); + st[c * 32 + r] = fma(g, st[c * 32 + r], kr * delta); + ka = fma(st[c * 32 + r], knr, ka); + qa = fma(st[c * 32 + r], qnr, qa); + } +#pragma clang loop unroll(full) + for (int c = 0; c < C; ++c) kl[c] = kn_l[c]; + g_b = g_b_n; + v = v_n; + } + if (T > 0) { + const int t = T - 1, slot = t & 1; + const float kqv = sg == 0 ? kq[t * hkn + hk] : 0.0f; + float2 s = float2(ka, qa); + if (gdn_probe != 1) { + red[slot][sg][lane] = s; + threadgroup_barrier(mem_flags::mem_threadgroup); + s = red[slot][0][lane]; +#pragma clang loop unroll(full) + for (int b = 1; b < IB; ++b) s += red[slot][b][lane]; + } + const float g = g_b.x; + const float delta = (v - g * s.x) * g_b.y; + if (sg == 0) o[t * rs + v0] = fma(g, s.y, delta * kqv) * scale; + if (gdn_probe != 2) +#pragma clang loop unroll(full) + for (int c = 0; c < C; ++c) +#pragma clang loop unroll(full) + for (int r = 0; r < R1; ++r) { + const float kr = simd_broadcast(kl[c], ushort(r)); + st[c * 32 + r] = fma(g, st[c * 32 + r], kr * delta); + } + } + } else { + constexpr int R4 = R / 4, TBS = TB > 0 ? TB : 1; + constexpr bool KC = R4 <= 8; + threadgroup float4 qs[TBS][32], ks[TBS][32]; + threadgroup float vs[TBS][32]; + threadgroup float4 gk[TBS]; + device const float4* qn4_p = (device const float4*)qn; + device const float4* kn4_p = (device const float4*)kn; + for (int t0 = 0; t0 < T; t0 += TBS) { + const int nb = min(TBS, T - t0); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int idx = tid; idx < nb * 32; idx += IB * 32) { + const int tt = idx >> 5, w = idx & 31; + qs[tt][w] = qn4_p[((t0 + tt) * hkn + hk) * 32 + w]; + ks[tt][w] = kn4_p[((t0 + tt) * hkn + hk) * 32 + w]; + vs[tt][w] = vc[((t0 + tt) * hv + h) * S + (j - lane) + w]; + if (w == 0) { + const float2 g_b = gb[(t0 + tt) * hv + h]; + gk[tt] = float4(g_b.x, g_b.y, kq[(t0 + tt) * hkn + hk], 0.0f); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float ka = 0.0f, qa = 0.0f; + float4 kc[R4]; #pragma clang loop unroll(full) for (int w = 0; w < R4; ++w) { - const float4 k4 = ks[tt][i0 / 4 + w]; - st[4 * w] = fma(g, st[4 * w], k4.x * delta); - st[4 * w + 1] = fma(g, st[4 * w + 1], k4.y * delta); - st[4 * w + 2] = fma(g, st[4 * w + 2], k4.z * delta); - st[4 * w + 3] = fma(g, st[4 * w + 3], k4.w * delta); + const float4 k4 = ks[0][i0 / 4 + w], q4 = qs[0][i0 / 4 + w]; + if (KC) kc[w] = k4; + ka = fma(st[4 * w], k4.x, ka); ka = fma(st[4 * w + 1], k4.y, ka); + ka = fma(st[4 * w + 2], k4.z, ka); ka = fma(st[4 * w + 3], k4.w, ka); + qa = fma(st[4 * w], q4.x, qa); qa = fma(st[4 * w + 1], q4.y, qa); + qa = fma(st[4 * w + 2], q4.z, qa); qa = fma(st[4 * w + 3], q4.w, qa); + } + for (int tt = 0; tt < nb; ++tt) { + const int t = t0 + tt, slot = t & 1; + const float4 g_b = gk[tt]; + const float kqv = g_b.z; + const float v = vs[tt][lane]; + float2 s = float2(ka, qa); + if (gdn_probe != 1) { + red[slot][sg][lane] = s; + threadgroup_barrier(mem_flags::mem_threadgroup); + s = red[slot][0][lane]; +#pragma clang loop unroll(full) + for (int b = 1; b < IB; ++b) s += red[slot][b][lane]; + } + const float g = g_b.x; + const float delta = (v - g * s.x) * g_b.y; + if (sg == 0) o[(t * hv + h) * S + j] = fma(g, s.y, delta * kqv) * scale; + const int tn = min(tt + 1, nb - 1); + ka = 0.0f; + qa = 0.0f; + if (gdn_probe != 2) +#pragma clang loop unroll(full) + for (int w = 0; w < R4; ++w) { + const float4 k4 = KC ? kc[w] : ks[tt][i0 / 4 + w], kn4 = ks[tn][i0 / 4 + w], qn4 = qs[tn][i0 / 4 + w]; + if (KC) kc[w] = kn4; + st[4 * w] = fma(g, st[4 * w], k4.x * delta); + st[4 * w + 1] = fma(g, st[4 * w + 1], k4.y * delta); + st[4 * w + 2] = fma(g, st[4 * w + 2], k4.z * delta); + st[4 * w + 3] = fma(g, st[4 * w + 3], k4.w * delta); + ka = fma(st[4 * w], kn4.x, ka); ka = fma(st[4 * w + 1], kn4.y, ka); + ka = fma(st[4 * w + 2], kn4.z, ka); ka = fma(st[4 * w + 3], kn4.w, ka); + qa = fma(st[4 * w], qn4.x, qa); qa = fma(st[4 * w + 1], qn4.y, qa); + qa = fma(st[4 * w + 2], qn4.z, qa); qa = fma(st[4 * w + 3], qn4.w, qa); + } } } } + device float* sp = state + s0; #pragma clang loop unroll(full) - for (int i = 0; i < R; ++i) state[((i0 + i) * hv + h) * S + j] = st[i]; + for (int i = 0; i < R; ++i) sp[i * rs] = st[i]; } // gdn_scan_r<8, 16> with the 16 state rows of a lane in four float4 variables (no array): a register-allocation check @@ -294,33 +494,52 @@ kernel void gdn_scan_v(device float* state [[buffer(0)]], device const float4* q sp[12 * rs] = s3.x; sp[13 * rs] = s3.y; sp[14 * rs] = s3.z; sp[15 * rs] = s3.w; } -// grid (h_v, T), 128 threads +// grid (h_v / 4, T) or (h_v, T), 128 threads: one simdgroup (32 lanes) per head h computes the RMSNorm and gated SiLU +// over the 128 channels (4 channels a lane) with zero threadgroup memory and zero barriers. kernel void gdn_post(device const float* o [[buffer(0)]], device const float* z [[buffer(1)]], device const float* gamma [[buffer(2)]], device half* yh [[buffer(3)]], device float* y [[buffer(4)]], constant GdnPArgs& a [[buffer(5)]], - uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + uint3 tg [[threadgroup_position_in_grid]], uint3 tgs [[threadgroups_per_grid]], ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { - threadgroup float part[4]; - const int h = int(tg.x), t = int(tg.y), hv = int(a.h_v); - const uint i = uint(t * hv + h) * 128 + tid; - const float ov = o[i]; - const float ss = simd_sum(ov * ov); - if (lane == 0) part[sg] = ss; - threadgroup_barrier(mem_flags::mem_threadgroup); - const float tot = (part[0] + part[1]) + (part[2] + part[3]); + const int hv = int(a.h_v), t = int(tg.y); + const int h = int(tgs.x) == hv ? (sg == 0 ? int(tg.x) : hv) : int(tg.x) * 4 + int(sg); + if (h >= hv) return; + const uint base = uint(t * hv + h) * 128 + lane; + device const float* o_row = o + base; + device const float* z_row = z + base; + device const float* g_row = gamma + lane; + device half* yh_row = yh + base; + device float* y_row = y + base; + const float ov0 = o_row[0], ov1 = o_row[32], ov2 = o_row[64], ov3 = o_row[96]; + // the gate and gamma do not depend on the norm: computed while the simd_sums run + float sig[4]; +#pragma clang loop unroll(full) + for (uint k = 0; k < 4; ++k) sig[k] = 1.0f / (1.0f + exp(-z_row[k * 32])); + const float gv[4] = {g_row[0], g_row[32], g_row[64], g_row[96]}; + const float4 pv = simd_sum(float4(ov0 * ov0, ov1 * ov1, ov2 * ov2, ov3 * ov3)); + const float tot = (pv.x + pv.y) + (pv.z + pv.w); const float rs = rsqrt(tot / 128.0f + a.eps); - const float yv = ((rs * ov) * gamma[tid]) * (1.0f / (1.0f + exp(-z[i]))); - yh[i] = half(yv); - if (a.f32) y[i] = yv; + const float ov[4] = {ov0, ov1, ov2, ov3}; + float yv[4]; +#pragma clang loop unroll(full) + for (uint k = 0; k < 4; ++k) { + yv[k] = ((rs * ov[k]) * gv[k]) * sig[k]; + yh_row[k * 32] = half(yv[k]); + } + if (a.f32) { +#pragma clang loop unroll(full) + for (uint k = 0; k < 4; ++k) y_row[k * 32] = yv[k]; + } } template [[host_name("gdn_scan_1")]] kernel decltype(gdn_scan_t<1, 16>) gdn_scan_t<1, 16>; template [[host_name("gdn_scan_2")]] kernel decltype(gdn_scan_t<2, 16>) gdn_scan_t<2, 16>; template [[host_name("gdn_scan_4")]] kernel decltype(gdn_scan_t<4, 16>) gdn_scan_t<4, 16>; -template [[host_name("gdn_scanr_2")]] kernel decltype(gdn_scan_r<2, 16>) gdn_scan_r<2, 16>; -template [[host_name("gdn_scanr_4")]] kernel decltype(gdn_scan_r<4, 16>) gdn_scan_r<4, 16>; -template [[host_name("gdn_scanr_8")]] kernel decltype(gdn_scan_r<8, 16>) gdn_scan_r<8, 16>; +template [[host_name("gdn_scanr_2")]] kernel decltype(gdn_scan_r<2, 0>) gdn_scan_r<2, 0>; +template [[host_name("gdn_scanr_4")]] kernel decltype(gdn_scan_r<4, 0>) gdn_scan_r<4, 0>; +template [[host_name("gdn_scanr_8")]] kernel decltype(gdn_scan_r<8, 0>) gdn_scan_r<8, 0>; template [[host_name("gdn_scanr_4_8")]] kernel decltype(gdn_scan_r<4, 8>) gdn_scan_r<4, 8>; +template [[host_name("gdn_scanr_4_16")]] kernel decltype(gdn_scan_r<4, 16>) gdn_scan_r<4, 16>; // ---- the chunked form (chunks of GC = 32 tokens), the WY / UT form of the gated delta rule. With G_t the product // of the chunk's gates up to t (log space: lg_t) and S0 the state entering the chunk: @@ -483,3 +702,281 @@ kernel void gdn_cout(device const float* qn [[buffer(0)]], device const float* k o[((t0 + t) * hv + h) * S + j] = s * scale; } } + +// ---- the chunked WY form on the matrix units (chunks of WYC = 16 tokens), the M5 prompt scan. For a value head, +// a chunk of nb tokens, S0 the state entering it, g_t the gate, r_ts = prod_{s < u <= t} g_u and G_t = r_t,-1: +// A_ts = beta_t (k_t . k_s) r_ts (s < t), TB = (I + A)^-1 diag(beta) +// delta = TB (V - diag(G) K S0) +// o = (diag(G) Q S0 + P delta) / sqrt(128), P_ts = (q_t . k_s) r_ts (s <= t) +// S = G_L S0 + K^T diag(r_L.) delta (L = nb - 1) +// (the sequential recurrence's algebra, to fp32 rounding: gdn_prompt_parity --wy). The gates enter as products, so +// a gate that underflows to 0 stays exact. gdn_wy_scan's matmuls are fp16 x fp16 -> fp32 on fp32 values split +// hi + lo: a b = ah bh + ah bl + al bh (the dropped al bl is ~2^-22 relative); the prep's are fp32. +// gdn_wy_prep per (key head, chunk): k k^T and q k^T (one fp32 matmul), TB (the inverse as +// (I - A)(I + A^2)(I + A^4)(I + A^8): A is nilpotent), P, G, r_L, G_L for the h_v / h_k value heads, +// and the chunk's k, q as fp16 hi / lo planes (h_v / h_k <= 4) +// gdn_wy_scan per (value head, 16 state columns): two simdgroups of 64 state rows each, the chunks in order +#include +using namespace mpp::tensor_ops; + +constant constexpr int WYC = 16; +constant constexpr int WY_F = 48; // floats per (head, chunk) in wyf: G[16], r_L[16], G_L + +// grid (h_k, chunks), 128 threads. kqh: 4 planes [h_k][chunks * 16][128] (k hi, k lo, q hi, q lo; rows past T +// zero); wyh: [h_v][chunks][4][16 x 16] (TB hi, TB lo, P hi, P lo); wyf: [h_v][chunks][WY_F] +kernel void gdn_wy_prep(device const float* qn [[buffer(0)]], device const float* kn [[buffer(1)]], + device const float2* gb [[buffer(2)]], device half* kqh [[buffer(3)]], + device half* wyh [[buffer(4)]], device float* wyf [[buffer(5)]], constant GdnPArgs& a [[buffer(6)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + constexpr int S = 128, C = WYC; + threadgroup float kq_s[2 * C * S]; // the chunk's K rows, then Q rows; the inverse's scratch once dd is done + threadgroup float dd[2 * C][C]; // k_t . k_s, then q_t . k_s + threadgroup float gg[4][C], bb[4][C]; + threadgroup float rr[4][C][C + 1]; // r_ts; [t][C] = G_t + threadgroup float aa[4][C][C]; + threadgroup float (*tm)[4][C * C] = (threadgroup float (*)[4][C * C])kq_s; + const int hk = int(tg.x), c = int(tg.y), hv = int(a.h_v), hkn = int(a.h_k), T = int(a.n_tok); + const int NC = (T + C - 1) / C, t0 = c * C, nb = min(C, T - t0), nvh = hv / hkn; + const ulong plane = ulong(hkn) * NC * C * S; + device half* kh = kqh + (ulong(hk) * NC * C + t0) * S; + for (int idx = tid; idx < C * S / 4; idx += 128) { + const int t = idx / (S / 4), w = idx % (S / 4); + float4 k4 = 0.0f, q4 = 0.0f; + if (t < nb) { + k4 = ((device const float4*)(kn + ((t0 + t) * hkn + hk) * S))[w]; + q4 = ((device const float4*)(qn + ((t0 + t) * hkn + hk) * S))[w]; + } + ((threadgroup float4*)kq_s)[idx] = k4; + ((threadgroup float4*)kq_s)[C * S / 4 + idx] = q4; + const half4 khi = half4(k4), qhi = half4(q4); + device half4* o = (device half4*)(kh + t * S) + w; + o[0] = khi; + o[plane / 4] = half4(k4 - float4(khi)); + o[plane / 2] = qhi; + o[3 * plane / 4] = half4(q4 - float4(qhi)); + } + if (tid < nvh * C) { + const int v = tid / C, t = tid % C; + const float2 g_b = t < nb ? gb[(t0 + t) * hv + hk + v * hkn] : float2(0.0f); + gg[v][t] = g_b.x; + bb[v][t] = g_b.y; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + { + tensor, tensor_inline> tA(kq_s, dextents(S, 2 * C)); + tensor, tensor_inline> tB(kq_s, dextents(S, C)); + tensor, tensor_inline> tD(&dd[0][0], dextents(C, 2 * C)); + constexpr auto desc = matmul2d_descriptor(2 * C, C, dynamic_length_v, false, true, false); + matmul2d> op; + op.run(tA, tB, tD); + } + for (int idx = tid; idx < nvh * C * (C + 1); idx += 128) { + const int v = idx / (C * (C + 1)), t = (idx / (C + 1)) % C, s = idx % (C + 1); + float p = 1.0f; + for (int u = s == C ? 0 : s + 1; u <= t; ++u) p *= gg[v][u]; + rr[v][t][s] = p; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int idx = tid; idx < nvh * C * C; idx += 128) { + const int v = idx / (C * C), t = (idx / C) % C, u = idx % C; + aa[v][t][u] = u < t ? bb[v][t] * dd[t][u] * rr[v][t][u] : 0.0f; + const float p = u <= t && t < nb ? dd[C + t][u] * rr[v][t][u] : 0.0f; + const half ph = half(p); + device half* op = wyh + (ulong(hk + v * hkn) * NC + c) * 4 * C * C + 2 * C * C + t * C + u; + op[0] = ph; + op[C * C] = half(p - float(ph)); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // (I + A)^-1 = (I - A)(I + A^2)(I + A^4)(I + A^8): simdgroup v for value head v, fp32 16 x 16 matmuls + if (int(sg) < nvh) { + const int v = int(sg); + threadgroup float* A = &aa[v][0][0]; + threadgroup float* P2 = tm[v][0]; + threadgroup float* P4 = tm[v][1]; + threadgroup float* Ta = tm[v][2]; + threadgroup float* Tb = tm[v][3]; + constexpr auto d = matmul2d_descriptor(C, C, C, false, false, false); + matmul2d op; + auto mm = [&](threadgroup float* X, threadgroup float* Y, threadgroup float* Z) { + tensor, tensor_inline> tX(X, dextents(C, C)); + tensor, tensor_inline> tY(Y, dextents(C, C)); + tensor, tensor_inline> tZ(Z, dextents(C, C)); + op.run(tX, tY, tZ); + simdgroup_barrier(mem_flags::mem_threadgroup); + }; + auto eye = [](int e) { return e / C == e % C ? 1.0f : 0.0f; }; + for (int e = lane; e < C * C; e += 32) Ta[e] = eye(e) - A[e]; + simdgroup_barrier(mem_flags::mem_threadgroup); + mm(A, A, P2); + for (int e = lane; e < C * C; e += 32) P4[e] = P2[e] + eye(e); + simdgroup_barrier(mem_flags::mem_threadgroup); + mm(Ta, P4, Tb); + mm(P2, P2, P4); + for (int e = lane; e < C * C; e += 32) P2[e] = P4[e] + eye(e); + simdgroup_barrier(mem_flags::mem_threadgroup); + mm(Tb, P2, Ta); + mm(P4, P4, P2); + for (int e = lane; e < C * C; e += 32) P2[e] += eye(e); + simdgroup_barrier(mem_flags::mem_threadgroup); + mm(Ta, P2, Tb); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int idx = tid; idx < nvh * C * C; idx += 128) { + const int v = idx / (C * C), t = (idx / C) % C, s = idx % C; + const float tb = t < nb ? tm[v][3][t * C + s] * bb[v][s] : 0.0f; + const half tbh = half(tb); + device half* oh = wyh + (ulong(hk + v * hkn) * NC + c) * 4 * C * C + t * C + s; + oh[0] = tbh; + oh[C * C] = half(tb - float(tbh)); + } + if (tid < nvh * C) { + const int v = tid / C, s = tid % C, L = nb - 1; + device float* o = wyf + (ulong(hk + v * hkn) * NC + c) * WY_F; + o[s] = s < nb ? rr[v][s][C] : 0.0f; + o[C + s] = s < nb ? rr[v][L][s] : 0.0f; + if (s == 0) o[2 * C] = rr[v][L][C]; + } +} + +// grid (h_v * 8), 64 threads: simdgroup sg owns state rows 64 sg .. 64 sg + 63 of 16 columns, as the destination +// cooperative tensor of its update matmul (64 x 16 x 16), never leaving the registers. On the M5 that tensor's +// elements 16 k .. 16 k + 15 are exactly the right input of rows 32 k .. 32 k + 31 for K S0 (16 x 16 x 32), and a +// 16 x 16 destination is the right input of a K = 16 matmul (gdn_prompt_parity --wy checks the result; host code +// takes this kernel only with coop_k4). K S0 and Q S0 are summed over the two row halves through threadgroup memory, +// double-buffered by chunk parity: one barrier a chunk. (One simdgroup of all 128 rows holds 64 floats of state a +// lane and spills at three matmuls an operand: 2.5x slower.) +kernel void gdn_wy_scan(device float* state [[buffer(0)]], device const float* vc [[buffer(1)]], + device const half* kqh [[buffer(2)]], device const half* wyh [[buffer(3)]], + device const float* wyf [[buffer(4)]], device float* o [[buffer(5)]], constant GdnPArgs& a [[buffer(6)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + constexpr int S = 128, C = WYC; + threadgroup float red[2][2][2][8 * 32]; // [chunk parity][simdgroup][K S0, Q S0][element * 32 + lane] + const int g_idx = int(tg.x) / 8, hv = int(a.h_v), hkn = int(a.h_k); + const int hk = (hv == 3 * hkn) ? (g_idx / 3) : (g_idx % hkn); + const int h = (hv == 3 * hkn) ? (hk + (g_idx % 3) * hkn) : g_idx; + const int j0 = (int(tg.x) % 8) * 16, T = int(a.n_tok), NC = (T + C - 1) / C, i0 = 64 * int(sg); + const ulong plane = ulong(hkn) * NC * C * S; + const float scale = 1.0f / sqrt(float(S)); + device half* kb = (device half*)kqh + ulong(hk) * NC * C * S; + tensor, tensor_inline> tKh(kb, dextents(S, NC * C)); + tensor, tensor_inline> tKl(kb + plane, dextents(S, NC * C)); + tensor, tensor_inline> tQh(kb + 2 * plane, dextents(S, NC * C)); + tensor, tensor_inline> tQl(kb + 3 * plane, dextents(S, NC * C)); + constexpr auto dU = matmul2d_descriptor(64, 16, C, true, false, false, matmul2d_descriptor::mode::multiply_accumulate); + constexpr auto dX = matmul2d_descriptor(C, 16, 32, false, false, false, matmul2d_descriptor::mode::multiply_accumulate); + constexpr auto dT = matmul2d_descriptor(C, 16, C, false, false, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d opU; + matmul2d opX; + matmul2d opT; + using RU = __tensor_ops_detail::__remove_addrspace_t())>; + using RX = __tensor_ops_detail::__remove_addrspace_t())>; + using RT = __tensor_ops_detail::__remove_addrspace_t())>; + auto st = opU.template get_destination_cooperative_tensor(); +#pragma clang loop unroll(full) + for (uint16_t e = 0; e < st.get_capacity(); ++e) { + const auto ix = st.get_multidimensional_index(e); // (column j, row i - i0) + st[e] = state[((i0 + ix[1]) * hv + h) * S + j0 + ix[0]]; + } + for (int c = 0; c < NC; ++c) { + const int t0 = c * C, nb = min(C, T - t0), par = c & 1; + device const float* w = wyf + (ulong(h) * NC + c) * WY_F; + device const half* wh = wyh + (ulong(h) * NC + c) * 4 * C * C; + // this half's K S0 and Q S0 + auto ks = opX.template get_destination_cooperative_tensor(); + auto qs = opX.template get_destination_cooperative_tensor(); +#pragma clang loop unroll(full) + for (uint16_t e = 0; e < ks.get_capacity(); ++e) { ks[e] = 0.0f; qs[e] = 0.0f; } + auto xs = [&](thread decltype(tKh)& th, thread decltype(tKh)& tl, thread decltype(ks)& acc) { +#pragma clang loop unroll(full) + for (int k = 0; k < 2; ++k) { + auto rh = opX.template get_right_input_cooperative_tensor(); + auto rl = opX.template get_right_input_cooperative_tensor(); +#pragma clang loop unroll(full) + for (uint16_t e = 0; e < rh.get_capacity(); ++e) { + const float v = st[16 * k + e]; + rh[e] = half(v); + rl[e] = half(v - float(rh[e])); + } + auto ah = th.template slice<32, 16>(i0 + 32 * k, t0); + auto al = tl.template slice<32, 16>(i0 + 32 * k, t0); + opX.run(ah, rh, acc); + opX.run(ah, rl, acc); + opX.run(al, rh, acc); + } + }; + xs(tKh, tKl, ks); + xs(tQh, tQl, qs); +#pragma clang loop unroll(full) + for (uint16_t e = 0; e < ks.get_capacity(); ++e) { + red[par][sg][0][e * 32 + lane] = ks[e]; + red[par][sg][1][e * 32 + lane] = qs[e]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); +#pragma clang loop unroll(full) + for (uint16_t e = 0; e < ks.get_capacity(); ++e) { + ks[e] = red[par][0][0][e * 32 + lane] + red[par][1][0][e * 32 + lane]; + qs[e] = red[par][0][1][e * 32 + lane] + red[par][1][1][e * 32 + lane]; + } + // R = V - G K S0 (rows past the chunk: 0), delta = TB R; both simdgroups (each needs delta) + tensor, tensor_inline> tTBh((device half*)wh, dextents(C, C)); + tensor, tensor_inline> tTBl((device half*)wh + C * C, dextents(C, C)); + tensor, tensor_inline> tPh((device half*)wh + 2 * C * C, dextents(C, C)); + tensor, tensor_inline> tPl((device half*)wh + 3 * C * C, dextents(C, C)); + auto rh = opT.template get_right_input_cooperative_tensor(); + auto rl = opT.template get_right_input_cooperative_tensor(); +#pragma clang loop unroll(full) + for (uint16_t e = 0; e < ks.get_capacity(); ++e) { + const auto ix = ks.get_multidimensional_index(e); // (column j, token t) + const int t = ix[1]; + const float v = t < nb ? vc[((t0 + t) * hv + h) * S + j0 + ix[0]] : 0.0f; + const float r = v - w[t] * ks[e]; + rh[e] = half(r); + rl[e] = half(r - float(rh[e])); + } + auto dl = opT.template get_destination_cooperative_tensor(); +#pragma clang loop unroll(full) + for (uint16_t e = 0; e < dl.get_capacity(); ++e) dl[e] = 0.0f; + opT.run(tTBh, rh, dl); + opT.run(tTBh, rl, dl); + opT.run(tTBl, rh, dl); + // o = G Q S0 + P delta (simdgroup 0); U = diag(r_L) delta, the update's right input + auto uh = opU.template get_right_input_cooperative_tensor(); + auto ul = opU.template get_right_input_cooperative_tensor(); + auto oc = opT.template get_destination_cooperative_tensor(); +#pragma clang loop unroll(full) + for (uint16_t e = 0; e < dl.get_capacity(); ++e) { + const int t = dl.get_multidimensional_index(e)[1]; + oc[e] = w[t] * qs[e]; + const float d = dl[e], u = d * w[C + t]; + rh[e] = half(d); + rl[e] = half(d - float(rh[e])); + uh[e] = half(u); + ul[e] = half(u - float(uh[e])); + } + if (sg == 0) { + opT.run(tPh, rh, oc); + opT.run(tPh, rl, oc); + opT.run(tPl, rh, oc); +#pragma clang loop unroll(full) + for (uint16_t e = 0; e < oc.get_capacity(); ++e) { + const auto ix = oc.get_multidimensional_index(e); + if (ix[1] < nb) o[((t0 + ix[1]) * hv + h) * S + j0 + ix[0]] = oc[e] * scale; + } + } + const float gl = w[2 * C]; +#pragma clang loop unroll(full) + for (uint16_t e = 0; e < st.get_capacity(); ++e) st[e] *= gl; + auto kt = tKh.template slice<64, 16>(i0, t0); + auto ktl = tKl.template slice<64, 16>(i0, t0); + opU.run(kt, uh, st); + opU.run(kt, ul, st); + opU.run(ktl, uh, st); + } +#pragma clang loop unroll(full) + for (uint16_t e = 0; e < st.get_capacity(); ++e) { + const auto ix = st.get_multidimensional_index(e); + state[((i0 + ix[1]) * hv + h) * S + j0 + ix[0]] = st[e]; + } +} diff --git a/metal/kernels/gemv.metal b/metal/kernels/gemv.metal index 6d4ed3a5..49b7b621 100644 --- a/metal/kernels/gemv.metal +++ b/metal/kernels/gemv.metal @@ -21,6 +21,85 @@ constant int NTOK [[function_constant(0)]]; // tokens of the window (1..8): th constant int GEMV_PROBE [[function_constant(9)]]; constant int gemv_probe = is_function_constant_defined(GEMV_PROBE) ? GEMV_PROBE : 0; +// Wider raw-dot windows and prepared Q2 at NR=1 did not improve on M5. +template +static inline bool gemv_direct_x() { + return F == T_Q8_0 && (NTOK == 1 || (NTOK == 2 && LPR == 16 && NR == 1)); +} + +// Keep unit_combine's four FMA chains without materialising the raw x slice. +template +static inline void gemv_raw_combine(device const float* xp, thread const float (&q)[NR][UNIT], + thread const float4* sc, thread float (&acc)[NR][8], int t) { + constexpr int H = UNIT / 2; + device const float4* x = (device const float4*)xp; + float s[NR][4]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) s[r][0] = s[r][1] = s[r][2] = s[r][3] = 0.0f; +#pragma clang loop unroll(full) + for (int j = 0; j < H; j += 4) { + const float4 a = x[j / 4], b = x[(H + j) / 4]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + s[r][0] = fma(q[r][j], a.x, s[r][0]); + s[r][1] = fma(q[r][j + 1], a.y, s[r][1]); + s[r][2] = fma(q[r][H + j], b.x, s[r][2]); + s[r][3] = fma(q[r][H + j + 1], b.y, s[r][3]); + s[r][0] = fma(q[r][j + 2], a.z, s[r][0]); + s[r][1] = fma(q[r][j + 3], a.w, s[r][1]); + s[r][2] = fma(q[r][H + j + 2], b.z, s[r][2]); + s[r][3] = fma(q[r][H + j + 3], b.w, s[r][3]); + } + } +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) + acc[r][t] += sc[r].x * (s[r][0] + s[r][1]) + sc[r].y * (s[r][2] + s[r][3]) - + (sc[r].z * 0.0f + sc[r].w * 0.0f); +} + +static inline float gemv_q2_raw_combine(ushort2 pk, float d, thread const float* xs) { + float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f; +#pragma clang loop unroll(full) + for (int j = 0; j < 8; j += 2) { + const uint2 q = uint2(pk) >> (2 * j); + const uint2 bits = ((q | (q << 14)) & 0x00030003u) | 0x64006400u; + const float4 v = float4(as_type(bits) - 1025.0h); + s0 = fma(v.x, xs[j], s0); + s1 = fma(v.y, xs[j + 1], s1); + s2 = fma(v.z, xs[8 + j], s2); + s3 = fma(v.w, xs[8 + j + 1], s3); + } + return d * (s0 + s1) + d * (s2 + s3) - 0.0f; +} + +// IQ4_XS kept as its 16 code bytes and scale (5 registers a row instead of 36), decoded inside the combine. M5, weights +// from DRAM: 2560 -> 640 at 16,2,4,1 -1.5..-5% (1 token), -2..-3% (2..4); 3..4 tokens -1..-2% on 2560 -> 6144..12288 +// (but 6144 -> 2560 at 16,2,4,1 3 tokens +1.4%); LPR 32 under 4 tokens +0.4..0.7%, so those keep the decoded unit. +// IQ4_NL (18-byte blocks, ushort reads) measured no change and keeps the decoded unit. +template +static inline bool gemv_iq4_packed() { return LPR <= 16 || NTOK >= 4; } + +static inline uint4 gemv_iq4xs_codes(device const uchar* row, int u, thread float& d) { + device const uchar* b = row + (u >> 3) * 136; + d = iq4xs_scale(b, u & 7); + return uint4(*(device const packed_uint4*)(b + 8 + 16 * (u & 7))); +} + +// unit_combine<32> of deq's q, decoding the codes as it goes: the same four FMA chains in the same order +static inline float gemv_iq4_combine(uint4 pk, float d, thread const float* xs, threadgroup const uint* tab) { + float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f; +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const uint v = pk[j >> 1] >> (16 * (j & 1)); + const half2 w0 = as_type(tab[v & 255]), w1 = as_type(tab[(v >> 8) & 255]); + s0 = fma(float(w0.x), xs[2 * j], s0); + s1 = fma(float(w1.x), xs[2 * j + 1], s1); + s2 = fma(float(w0.y), xs[16 + 2 * j], s2); + s3 = fma(float(w1.y), xs[16 + 2 * j + 1], s3); + } + return d * (s0 + s1) + d * (s2 + s3) - 0.0f; +} + template kernel void gemv(device const uchar* W [[buffer(0)]], device const float* xp [[buffer(1)]], device const float* s16 [[buffer(2)]], device float* y [[buffer(3)]], constant GemvArgs& a [[buffer(4)]], @@ -29,8 +108,10 @@ kernel void gemv(device const uchar* W [[buffer(0)]], device const float* xp [[b constexpr int UNIT = fmt::UNIT, RS = SG / KS, RG = 32 / LPR; threadgroup uint tab[fmt::TAB]; threadgroup float red[KS > 1 ? KS * RS * RG * NR * 8 : 1]; - load_tab(tab, tid, SG * 32); - threadgroup_barrier(mem_flags::mem_threadgroup); + if (fmt::TAB > 1) { + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + } const int ks = sg / RS, rs = sg % RS; const int li = lane % LPR; const int lrow0 = (rs * RG + lane / LPR) * NR; @@ -38,18 +119,23 @@ kernel void gemv(device const uchar* W [[buffer(0)]], device const float* xp [[b const int U = int(a.n_in) / UNIT; const int u0 = ks * U / KS, u1 = (ks + 1) * U / KS; float acc[NR][8]; +#pragma clang loop unroll(full) for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) for (int t = 0; t < NTOK; ++t) acc[r][t] = 0.0f; device const uchar* rp[NR]; +#pragma clang loop unroll(full) for (int r = 0; r < NR; ++r) rp[r] = W + (ulong)min(row0 + r, int(a.n_out) - 1) * a.row_bytes; const int s16n = int(a.n_in) / 16; if (gemv_probe == 3) { // rows outer, tokens inner: one row's decoded unit in registers at a time (x re-read per row from L1) for (int u = u0 + li; u < u1; u += LPR) +#pragma clang loop unroll(full) for (int r = 0; r < NR; ++r) { float q[UNIT]; float4 sc; Dec::run(rp[r], u, q, sc, tab); +#pragma clang loop unroll(full) for (int t = 0; t < NTOK; ++t) { float xs[UNIT], xsum[2]; X::load(xp + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); @@ -58,18 +144,58 @@ kernel void gemv(device const uchar* W [[buffer(0)]], device const float* xp [[b } } else for (int u = u0 + li; u < u1; u += LPR) { - float q[NR][UNIT]; - float4 sc[NR]; - for (int r = 0; r < NR; ++r) Dec::run(rp[r], u, q[r], sc[r], tab); - float xs[UNIT], xsum[2]; - for (int t = 0; t < NTOK; ++t) { - if (gemv_probe != 1 || t == 0) X::load(xp + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); - if (gemv_probe != 2 || t == 0) - for (int r = 0; r < NR; ++r) acc[r][t] += unit_combine(q[r], sc[r], xs, xsum); + if (NTOK == 1 && F == T_Q2_0R) { + ushort2 pk[NR]; + float d[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + device const uchar* b = rp[r] + (u >> 2) * 18; + device const ushort* p = (device const ushort*)(b + 2 + 4 * (u & 3)); + pk[r] = ushort2(p[0], p[1]); + d[r] = rdh(b); + } + float xs[16]; + load16(xp + 16 * u, xs); +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) acc[r][0] += gemv_q2_raw_combine(pk[r], d[r], xs); + } else if (F == T_IQ4_XS && gemv_iq4_packed()) { + uint4 pk[NR]; + float d[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) pk[r] = gemv_iq4xs_codes(rp[r], u, d[r]); +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float xs[32]; + load16(xp + t * int(a.x_stride) + 32 * u, xs); + load16(xp + t * int(a.x_stride) + 32 * u + 16, xs + 16); +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) acc[r][t] += gemv_iq4_combine(pk[r], d[r], xs, tab); + } + } else { + float q[NR][UNIT]; + float4 sc[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) Dec::run(rp[r], u, q[r], sc[r], tab); + float xs[UNIT], xsum[2]; +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + if (gemv_direct_x()) { + if (gemv_probe != 2 || t == 0) + gemv_raw_combine(xp + (gemv_probe == 1 ? 0 : t) * int(a.x_stride) + u * UNIT, + q, sc, acc, t); + } else { + if (gemv_probe != 1 || t == 0) X::load(xp + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); + if (gemv_probe != 2 || t == 0) +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) acc[r][t] += unit_combine(q[r], sc[r], xs, xsum); + } + } } } constexpr int RPT = RS * RG * NR; +#pragma clang loop unroll(full) for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) for (int t = 0; t < NTOK; ++t) { const float v = lane_sum(acc[r][t]); if (KS == 1) { @@ -81,12 +207,16 @@ kernel void gemv(device const uchar* W [[buffer(0)]], device const float* xp [[b if (KS > 1) { threadgroup_barrier(mem_flags::mem_threadgroup); if (ks == 0 && li == 0) +#pragma clang loop unroll(full) for (int r = 0; r < NR; ++r) - for (int t = 0; t < NTOK; ++t) { - float v = 0.0f; - for (int k = 0; k < KS; ++k) v += red[(k * RPT + lrow0 + r) * 8 + t]; - if (row0 + r < int(a.n_out)) y[t * int(a.y_stride) + row0 + r] = v; - } + if (row0 + r < int(a.n_out)) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float v = 0.0f; +#pragma clang loop unroll(full) + for (int k = 0; k < KS; ++k) v += red[(k * RPT + lrow0 + r) * 8 + t]; + y[t * int(a.y_stride) + row0 + r] = v; + } } } @@ -114,10 +244,9 @@ kernel void gemv_pair(device const uchar* Wg [[buffer(0)]], device const float* constexpr int TB = fmt::TAB > fmt::TAB ? fmt::TAB : fmt::TAB; threadgroup uint tab1[fmt::TAB]; threadgroup uint tab2[fmt::TAB]; - (void)TB; - load_tab(tab1, tid, SG * 32); - load_tab(tab2, tid, SG * 32); - threadgroup_barrier(mem_flags::mem_threadgroup); + if (fmt::TAB > 1) load_tab(tab1, tid, SG * 32); + if (fmt::TAB > 1) load_tab(tab2, tid, SG * 32); + if (TB > 1) threadgroup_barrier(mem_flags::mem_threadgroup); const int li = lane % LPR; const int row = (int(tg.x) * SG + sg) * RG + lane / LPR; const int rr = min(row, int(a.n_out) - 1); @@ -125,11 +254,13 @@ kernel void gemv_pair(device const uchar* Wg [[buffer(0)]], device const float* device const uchar* ru = Wu + ulong(rr) * rb_up; const int s16n = int(a.n_in) / 16; float ag[8], au[8]; - for (int t = 0; t < 8; ++t) ag[t] = au[t] = 0.0f; +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) ag[t] = au[t] = 0.0f; for (int u = li; u < int(a.n_in) / U1; u += LPR) { float q[U1]; float4 sc; Dec::run(rg, u, q, sc, tab1); +#pragma clang loop unroll(full) for (int t = 0; t < NTOK; ++t) { float xs[U1], xsum[2]; X::load(xg + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); @@ -140,12 +271,14 @@ kernel void gemv_pair(device const uchar* Wg [[buffer(0)]], device const float* float q[U2]; float4 sc; Dec::run(ru, u, q, sc, tab2); +#pragma clang loop unroll(full) for (int t = 0; t < NTOK; ++t) { float xs[U2], xsum[2]; X::load(xu + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); au[t] += unit_combine(q, sc, xs, xsum); } } +#pragma clang loop unroll(full) for (int t = 0; t < NTOK; ++t) { const float gv = lane_sum(ag[t]), uv = lane_sum(au[t]); if (li == 0 && row < int(a.n_out)) h[t * int(a.y_stride) + row] = gv / (1.0f + exp(-gv)) * uv; @@ -171,15 +304,20 @@ kernel void gemv_cat(device const uchar* W [[buffer(0)]], device const float* xp ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { constexpr int UNIT = fmt::UNIT, RG = 32 / LPR; threadgroup uint tab[fmt::TAB]; - load_tab(tab, tid, SG * 32); - threadgroup_barrier(mem_flags::mem_threadgroup); + if (fmt::TAB > 1) { + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + } const int li = lane % LPR; const int row0 = int(tg.x) * SG * RG * NR + (sg * RG + lane / LPR) * NR; const int U = int(a.n_in) / UNIT, sp = int(n_split); float acc[NR][8]; +#pragma clang loop unroll(full) for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) for (int t = 0; t < NTOK; ++t) acc[r][t] = 0.0f; device const uchar* rp[NR]; +#pragma clang loop unroll(full) for (int r = 0; r < NR; ++r) { const int row = min(row0 + r, int(a.n_out) - 1); rp[r] = row < sp ? W + (ulong)row * a.row_bytes : W2 + (ulong)(row - sp) * a.row_bytes; @@ -188,14 +326,23 @@ kernel void gemv_cat(device const uchar* W [[buffer(0)]], device const float* xp for (int u = li; u < U; u += LPR) { float q[NR][UNIT]; float4 sc[NR]; +#pragma clang loop unroll(full) for (int r = 0; r < NR; ++r) Dec::run(rp[r], u, q[r], sc[r], tab); +#pragma clang loop unroll(full) for (int t = 0; t < NTOK; ++t) { - float xs[UNIT], xsum[2]; - X::load(xp + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); - for (int r = 0; r < NR; ++r) acc[r][t] += unit_combine(q[r], sc[r], xs, xsum); + if (gemv_direct_x()) { + gemv_raw_combine(xp + t * int(a.x_stride) + u * UNIT, q, sc, acc, t); + } else { + float xs[UNIT], xsum[2]; + X::load(xp + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) acc[r][t] += unit_combine(q[r], sc[r], xs, xsum); + } } } +#pragma clang loop unroll(full) for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) for (int t = 0; t < NTOK; ++t) { const float v = lane_sum(acc[r][t]); const int row = row0 + r; @@ -209,6 +356,86 @@ kernel void gemv_cat(device const uchar* W [[buffer(0)]], device const float* xp #define GVC(F, L, R, S) GVC_(F, L, R, S) GVC(T_BF16, 32, 1, 4) GVC(T_BF16, 32, 1, 8) GVC(T_BF16, 16, 1, 4) GVC(T_BF16, 32, 2, 4) GVC(T_F32, 32, 1, 4) GVC(T_F32, 32, 1, 8) GVC(T_F32, 16, 1, 4) GVC(T_F32, 32, 2, 4) +GVC(T_Q8_0, 32, 1, 4) GVC(T_Q8_0, 32, 1, 8) GVC(T_Q8_0, 16, 1, 4) GVC(T_Q8_0, 32, 2, 4) GVC(T_Q8_0, 8, 1, 8) + +// gemv_catk: gemv_cat with gemv's K split (KS simdgroups share a row, partial sums added in gemv's fixed order), so +// the outputs are bitwise gemv_____'s. (The MTP draft step's k + v and shared gate + up: Q8_0 +// 2560 -> 512 + 512 / 640 + 640 at 32,1,4,4 on the M5.) Buffers as gemv_cat. +template +kernel void gemv_catk(device const uchar* W [[buffer(0)]], device const float* xp [[buffer(1)]], + device const float* s16 [[buffer(2)]], device float* y [[buffer(3)]], constant GemvArgs& a [[buffer(4)]], + device const uchar* W2 [[buffer(5)]], device float* y2 [[buffer(6)]], constant uint& n_split [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int UNIT = fmt::UNIT, RS = SG / KS, RG = 32 / LPR; + threadgroup uint tab[fmt::TAB]; + threadgroup float red[KS * RS * RG * NR * 8]; + if (fmt::TAB > 1) { + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const int ks = sg / RS, rs = sg % RS; + const int li = lane % LPR; + const int lrow0 = (rs * RG + lane / LPR) * NR; + const int row0 = int(tg.x) * RS * RG * NR + lrow0; + const int U = int(a.n_in) / UNIT, sp = int(n_split); + const int u0 = ks * U / KS, u1 = (ks + 1) * U / KS; + float acc[NR][8]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) acc[r][t] = 0.0f; + device const uchar* rp[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + const int row = min(row0 + r, int(a.n_out) - 1); + rp[r] = row < sp ? W + (ulong)row * a.row_bytes : W2 + (ulong)(row - sp) * a.row_bytes; + } + const int s16n = int(a.n_in) / 16; + for (int u = u0 + li; u < u1; u += LPR) { + float q[NR][UNIT]; + float4 sc[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) Dec::run(rp[r], u, q[r], sc[r], tab); +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + if (gemv_direct_x()) { + gemv_raw_combine(xp + t * int(a.x_stride) + u * UNIT, q, sc, acc, t); + } else { + float xs[UNIT], xsum[2]; + X::load(xp + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) acc[r][t] += unit_combine(q[r], sc[r], xs, xsum); + } + } + } + constexpr int RPT = RS * RG * NR; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + const float v = lane_sum(acc[r][t]); + if (li == 0) red[(ks * RPT + lrow0 + r) * 8 + t] = v; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (ks == 0 && li == 0) +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + const int row = row0 + r; + if (row < int(a.n_out)) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float v = 0.0f; +#pragma clang loop unroll(full) + for (int k = 0; k < KS; ++k) v += red[(k * RPT + lrow0 + r) * 8 + t]; + if (row < sp) y[t * int(a.y_stride) + row] = v; + else y2[t * int(a.y_stride) + row - sp] = v; + } + } +} +#define GVCK_(F, L, R, S, K) template [[host_name("gemvck_" #F "_" #L "_" #R "_" #S "_" #K)]] kernel decltype(gemv_catk) gemv_catk; +#define GVCK(F, L, R, S, K) GVCK_(F, L, R, S, K) +GVCK(T_Q8_0, 32, 1, 4, 4) GVCK(T_Q8_0, 32, 2, 4, 2) GVCK(T_Q8_0, 32, 1, 8, 8) // gemv_x: gemv (KS = 1) with the window's prepared x (the format's plane, NTOK rows, and its per-16 sums) staged in // threadgroup memory once per threadgroup, so the per-unit x reads of every token come from threadgroup memory @@ -223,9 +450,10 @@ kernel void gemv_x(device const uchar* W [[buffer(0)]], device const float* xp [ ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { constexpr int UNIT = fmt::UNIT, RG = 32 / LPR; threadgroup uint tab[fmt::TAB]; - load_tab(tab, tid, SG * 32); + if (fmt::TAB > 1) load_tab(tab, tid, SG * 32); const int n = int(a.n_in), s16n = n / 16; threadgroup float* st = xt + NTOK * n; +#pragma clang loop unroll(full) for (int t = 0; t < NTOK; ++t) { device const float4* src = (device const float4*)(xp + t * int(a.x_stride)); threadgroup float4* dst = (threadgroup float4*)(xt + t * n); @@ -237,21 +465,29 @@ kernel void gemv_x(device const uchar* W [[buffer(0)]], device const float* xp [ const int row0 = int(tg.x) * SG * RG * NR + (sg * RG + lane / LPR) * NR; const int U = n / UNIT; float acc[NR][8]; +#pragma clang loop unroll(full) for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) for (int t = 0; t < NTOK; ++t) acc[r][t] = 0.0f; device const uchar* rp[NR]; +#pragma clang loop unroll(full) for (int r = 0; r < NR; ++r) rp[r] = W + (ulong)min(row0 + r, int(a.n_out) - 1) * a.row_bytes; for (int u = li; u < U; u += LPR) { float q[NR][UNIT]; float4 sc[NR]; +#pragma clang loop unroll(full) for (int r = 0; r < NR; ++r) Dec::run(rp[r], u, q[r], sc[r], tab); +#pragma clang loop unroll(full) for (int t = 0; t < NTOK; ++t) { float xs[UNIT], xsum[2]; X::load((threadgroup const float*)(xt + t * n), (threadgroup const float*)(st + t * s16n), u, xs, xsum); +#pragma clang loop unroll(full) for (int r = 0; r < NR; ++r) acc[r][t] += unit_combine(q[r], sc[r], xs, xsum); } } +#pragma clang loop unroll(full) for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) for (int t = 0; t < NTOK; ++t) { const float v = lane_sum(acc[r][t]); if (li == 0 && row0 + r < int(a.n_out)) y[t * int(a.y_stride) + row0 + r] = v; diff --git a/metal/kernels/glue.metal b/metal/kernels/glue.metal index b7bc8bcc..b34bb90d 100644 --- a/metal/kernels/glue.metal +++ b/metal/kernels/glue.metal @@ -10,28 +10,29 @@ struct GlueArgs { uint topk; }; +constant float glue_q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, + 0.00006103515625f}; +constant float glue_p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + // grid (n / 16, n_tok), 16 threads: one group of 16 values kernel void swiglu_prep(device const float* g [[buffer(0)]], device const float* u [[buffer(1)]], device float* xp [[buffer(2)]], constant GlueArgs& a [[buffer(3)]], uint3 tg [[threadgroup_position_in_grid]], ushort j [[thread_index_in_threadgroup]]) { - threadgroup float v[16]; - const uint t = tg.y, d = tg.x * 16 + j, plane = a.n_tok * a.n; - const float gv = g[t * a.n + d]; - const float h = gv / (1.0f + exp(-gv)) * u[t * a.n + d]; - const float q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f}; - const float p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; - device float* raw = xp + t * a.n + d; + const uint t = tg.y, d = tg.x * 16 + j, idx = t * a.n + d, plane = a.n_tok * a.n; + const float gv = g[idx]; + const float h = gv / (1.0f + exp(-gv)) * u[idx]; + device float* raw = xp + idx; raw[0] = h; - raw[plane] = h * q2s[j & 7]; - raw[2 * plane] = h * p4s[j & 3]; + raw[plane] = h * glue_q2s[j & 7]; + raw[2 * plane] = h * glue_p4s[j & 3]; raw[3 * plane] = (j & 1) ? h * 0.00390625f : h; - v[j] = h; - threadgroup_barrier(mem_flags::mem_threadgroup); - if (j == 0) { - float s0 = 0.0f, s1 = 0.0f; - for (int k = 0; k < 16; k += 2) { s0 += v[k]; s1 += v[k + 1]; } - xp[4 * plane + t * (a.n / 16) + tg.x] = s0 + s1; + float s0 = 0.0f, s1 = 0.0f; +#pragma clang loop unroll(full) + for (int k = 0; k < 16; k += 2) { + s0 += simd_shuffle_down(h, ushort(k)); + s1 += simd_shuffle_down(h, ushort(k + 1)); } + if (j == 0) xp[4 * plane + t * (a.n / 16) + tg.x] = s0 + s1; } kernel void moe_combine_sh(device const float* ye [[buffer(0)]], device const int* plan [[buffer(1)]], @@ -40,8 +41,16 @@ kernel void moe_combine_sh(device const float* ye [[buffer(0)]], device const in uint gid [[thread_position_in_grid]]) { const uint t = gid / a.n_embd, k = gid % a.n_embd; if (t >= a.n_tok) return; + const int row_base = int(t * a.topk); + device const int* te = plan + P_TOKENT + row_base; + device const float* yep = ye + k; float s = 0.0f; - for (uint j = 0; j < a.topk; ++j) s += ye[plan[P_TOKENT + t * a.topk + j] * int(a.n_embd) + k]; + if (a.topk == 10) { +#pragma clang loop unroll(full) + for (uint j = 0; j < 10; ++j) s += yep[te[j] * int(a.n_embd)]; + } else { + for (uint j = 0; j < a.topk; ++j) s += yep[te[j] * int(a.n_embd)]; + } const float gs = 1.0f / (1.0f + exp(-sg[t])); bo[gid] = fma(gs, ysh[gid], s); } diff --git a/metal/kernels/gr.metal b/metal/kernels/gr.metal index c1960fd5..c18d2ac7 100644 --- a/metal/kernels/gr.metal +++ b/metal/kernels/gr.metal @@ -12,6 +12,11 @@ constant int GR_NTOK [[function_constant(1)]]; +// the prepared activations' per-position plane scales (units.metal's xprep): q2 planes 4^-(j%8), p4 planes 256^-(j%4) +constant float gr_q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, + 0.00006103515625f}; +constant float gr_p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + struct GrArgs { uint n_embd; uint hc; @@ -31,80 +36,195 @@ kernel void gr_norm(device float* R [[buffer(0)]], device const float* bo [[buff const uint c = tg.x, t = tg.y, n = a.n_embd; device float* r = R + (t * a.hc + c) * n; const float g = a.apply ? 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))) : 0.0f; + float rv[10]; float ss = 0.0f; - for (uint d = tid; d < n; d += 256) { - float v = r[d]; + if (n == 2560) { + device float* rp = r + tid; if (a.apply) { - v = fma(bo[t * n + d], g, v); - r[d] = v; + device const float* bp = bo + t * 2560u + tid; + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + const float v = fma(bp[i * 256], g, rp[i * 256]); + rp[i * 256] = v; + rv[i] = v; + ss = fma(v, v, ss); + } + } else { + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + rv[i] = rp[i * 256]; + ss = fma(rv[i], rv[i], ss); + } + } + } else { + for (uint d = tid; d < n; d += 256) { + float v = r[d]; + if (a.apply) { + v = fma(bo[t * n + d], g, v); + r[d] = v; + } + if (n <= 2560) rv[d / 256] = v; + ss = fma(v, v, ss); } - ss = fma(v, v, ss); } ss = simd_sum(ss); if (lane == 0) part[sg] = ss; threadgroup_barrier(mem_flags::mem_threadgroup); float tot = 0.0f; + #pragma clang loop unroll(full) for (int i = 0; i < 8; ++i) tot += part[i]; const float rs = rsqrt(tot / float(n) + a.eps); device float* o = xn + (t * a.hc + c) * n; device const float* w = w_norm + c * n; - for (uint d = tid; d < n; d += 256) o[d] = r[d] * w[d] * rs; + if (n == 2560) { + o += tid; + w += tid; + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) o[i * 256] = rv[i] * w[i * 256] * rs; + } else { + for (uint d = tid; d < n; d += 256) + o[d] = (n <= 2560 ? rv[d / 256] : r[d]) * w[d] * rs; + } } // rows: 0 .. hc_lr - 1 = w_down, hc_lr .. hc_lr + hc - 1 = w_inject (both bf16 rows of hc * n_embd). -// SG simdgroups split the row's units KS ways; LPR = 32 (a simdgroup per row); RS = SG / KS rows per threadgroup. -// one row's K-split partial over xn (units of format F) -template -static inline void grd_acc(device const uchar* w, int nin, int ks, int KS, ushort lane, device const float* xn, +// SG simdgroups split the row's units KS ways; LPR = 32 (a simdgroup per row); RS = SG / KS simdgroups per +// threadgroup, NR rows each. +// one token's dot with a decoded unit (gr_up): x is the token's slice of lo +template +static inline float gr_dot(device const float* x, thread const float* q, float4 sc) { + float xs[UNIT], xsum[2] = {0.0f, 0.0f}; + #pragma clang loop unroll(full) + for (int k = 0; k < UNIT; k += 16) load16(x + k, xs + k); + return unit_combine(q, sc, xs, xsum); +} + +// a unit's decoded rows for grd_acc, read four codes at a time +template +struct GrdUnit { + float q[NR][fmt::UNIT]; + float4 sc[NR]; + void load(device const uchar* w, uint rb, int u, threadgroup const uint* notab) thread { + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) Dec::run(w + ulong(r) * rb, u, q[r], sc[r], notab); + } + float4 v(int r, int i) const thread { return float4(q[r][4 * i], q[r][4 * i + 1], q[r][4 * i + 2], q[r][4 * i + 3]); } +}; + +// Q8_0 two rows a simdgroup: the codes stay packed (8 registers a row, not 32 floats) and are widened per token. +// Measured 10-16% faster past 6 tokens (where the token loop is rolled), no faster below, and slower for one row. +template <> +struct GrdUnit { + char4 q[2][8]; + float4 sc[2]; + void load(device const uchar* w, uint rb, int u, threadgroup const uint*) thread { + #pragma clang loop unroll(full) + for (int r = 0; r < 2; ++r) { + device const uchar* b = w + ulong(r) * rb + 34 * u; + device const ushort* p = (device const ushort*)(b + 2); + #pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) q[r][k] = as_type(uint(p[2 * k]) | (uint(p[2 * k + 1]) << 16)); + const float d = rdh(b); + sc[r] = float4(d, d, 0.0f, 0.0f); + } + } + float4 v(int r, int i) const thread { return float4(q[r][i]); } +}; + +// NR consecutive rows' K-split partials over xn (units of format F; the rows rb bytes apart), acc[r * 8 + t]. +// unit_combine's sums in its order, with each xn float4 read straight from memory once for all NR rows: xn is +// 15-30x the weights' bytes at 4-8 tokens, so NR = 2 halves the kernel's traffic. GR_NTOK is a function constant, +// which the compiler does not unroll over by itself: the pragma keeps acc in registers instead of the stack. Past 6 +// tokens the unrolled loop measured slower than the rolled one, so it stays rolled there; ROLL keeps it rolled at +// every count (the packed Q8_0 pair at KS = 4: 1-3% faster at 4-6 tokens, where KS = 8 measured 2-6% slower). +template +static inline void grd_acc(device const uchar* w, uint rb, int nin, int ks, int KS, ushort lane, device const float* xn, thread float* acc, threadgroup const uint* notab) { - constexpr int UNIT = fmt::UNIT; + constexpr int UNIT = fmt::UNIT, H = UNIT / 2; const int U = nin / UNIT, u0 = ks * U / KS, u1 = (ks + 1) * U / KS; for (int u = u0 + lane; u < u1; u += 32) { - float q[UNIT]; - float4 sc; - Dec::run(w, u, q, sc, notab); - for (int t = 0; t < GR_NTOK; ++t) { - float xs[UNIT], xsum[2] = {0.0f, 0.0f}; - for (int k = 0; k < UNIT; k += 16) load16(xn + t * nin + UNIT * u + k, xs + k); - acc[t] += unit_combine(q, sc, xs, xsum); + GrdUnit d; + d.load(w, rb, u, notab); + auto tok = [&](int t) { + device const float4* x = (device const float4*)(xn + t * nin + UNIT * u); + float s[NR][4]; + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) s[r][0] = s[r][1] = s[r][2] = s[r][3] = 0.0f; + #pragma clang loop unroll(full) + for (int j = 0; j < H; j += 4) { + const float4 a = x[j / 4], b = x[(H + j) / 4]; + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + const float4 qa = d.v(r, j / 4), qb = d.v(r, (H + j) / 4); + s[r][0] = fma(qa.x, a.x, s[r][0]); + s[r][1] = fma(qa.y, a.y, s[r][1]); + s[r][2] = fma(qb.x, b.x, s[r][2]); + s[r][3] = fma(qb.y, b.y, s[r][3]); + s[r][0] = fma(qa.z, a.z, s[r][0]); + s[r][1] = fma(qa.w, a.w, s[r][1]); + s[r][2] = fma(qb.z, b.z, s[r][2]); + s[r][3] = fma(qb.w, b.w, s[r][3]); + } + } + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) + acc[r * 8 + t] += d.sc[r].x * (s[r][0] + s[r][1]) + d.sc[r].y * (s[r][2] + s[r][3]) - + (d.sc[r].z * 0.0f + d.sc[r].w * 0.0f); + }; + if (GR_NTOK <= 6 && !ROLL) { + #pragma clang loop unroll(full) + for (int t = 0; t < GR_NTOK; ++t) tok(t); + } else { + for (int t = 0; t < GR_NTOK; ++t) tok(t); } } } -// FI: the inject rows' format (the Unsloth files: Q8_0 down, F32 inject); FI = F is the original kernel -template +// FI: the inject rows' format (the Unsloth files: Q8_0 down, F32 inject); FI = F is the original kernel. NR must +// divide hc_lr and hc, so that a simdgroup's rows are all w_down or all w_inject. The sums do not depend on SG or NR: +// every SG / NR with the same KS is bitwise the same kernel. +template kernel void gr_down_f(device const uchar* w_down [[buffer(0)]], device const uchar* w_inject [[buffer(1)]], device const float* xn [[buffer(2)]], device float* lo [[buffer(3)]], device float* inject [[buffer(4)]], constant GrArgs& a [[buffer(5)]], uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { constexpr int RS = SG / KS; - threadgroup float red[KS * RS * 8]; + threadgroup float red[KS * RS * NR * 8]; threadgroup uint notab[1]; const int ks = sg / RS, rs = sg % RS; - const int row = int(tg.x) * RS + rs; + const int row0 = (int(tg.x) * RS + rs) * NR; const int nrows = int(a.hc_lr + a.hc); - const int rr = min(row, nrows - 1); + const int rr = min(row0, nrows - NR); const int nin = int(a.hc * a.n_embd); const uint rb = uint(nin / fmt::QK * fmt::BB), rbi = uint(nin / fmt::QK * fmt::BB); - float acc[8]; - for (int t = 0; t < 8; ++t) acc[t] = 0.0f; - if (rr < int(a.hc_lr)) grd_acc(w_down + ulong(rr) * rb, nin, ks, KS, lane, xn, acc, notab); - else grd_acc(w_inject + ulong(rr - int(a.hc_lr)) * rbi, nin, ks, KS, lane, xn, acc, notab); - for (int t = 0; t < GR_NTOK; ++t) { - const float v = simd_sum(acc[t]); - if (lane == 0) red[(ks * RS + rs) * 8 + t] = v; + float acc[NR * 8]; + #pragma clang loop unroll(full) + for (int t = 0; t < NR * 8; ++t) acc[t] = 0.0f; + constexpr bool roll = F == T_Q8_0 && NR == 2 && KS == 4; + if (rr < int(a.hc_lr)) grd_acc(w_down + ulong(rr) * rb, rb, nin, ks, KS, lane, xn, acc, notab); + else grd_acc(w_inject + ulong(rr - int(a.hc_lr)) * rbi, rbi, nin, ks, KS, lane, xn, acc, notab); + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + #pragma clang loop unroll(full) + for (int t = 0; t < GR_NTOK; ++t) { + const float v = simd_sum(acc[r * 8 + t]); + if (lane == 0) red[((ks * RS + rs) * NR + r) * 8 + t] = v; + } } threadgroup_barrier(mem_flags::mem_threadgroup); - if (ks == 0 && lane < GR_NTOK && row < nrows) { - const int t = lane; - float v = 0.0f; - for (int k = 0; k < KS; ++k) v += red[(k * RS + rs) * 8 + t]; - if (row < int(a.hc_lr)) { - const float z = v / float(a.hc); - lo[t * int(a.hc_lr) + row] = z / (1.0f + exp(-z)); - } else { - inject[t * int(a.hc) + row - int(a.hc_lr)] = v; + if (ks == 0 && lane < NR * 8) { + const int r = lane / 8, t = lane % 8, row = row0 + r; + if (t < GR_NTOK && row < nrows) { + float v = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < KS; ++k) v += red[((k * RS + rs) * NR + r) * 8 + t]; + if (row < int(a.hc_lr)) { + const float z = v / float(a.hc); + lo[t * int(a.hc_lr) + row] = z / (1.0f + exp(-z)); + } else { + inject[t * int(a.hc) + row - int(a.hc_lr)] = v; + } } } } @@ -119,7 +239,6 @@ kernel void gr_up_f(device const uchar* w_up [[buffer(0)]], device const float* uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { threadgroup float gate[8][4][16]; - threadgroup float msum[8][16]; threadgroup uint notab[1]; const int n = int(a.n_embd), c = sg >> 1, dl = 8 * (sg & 1) + lane / 4, li = lane & 3; const int d = int(tg.x) * 16 + dl; @@ -127,17 +246,16 @@ kernel void gr_up_f(device const uchar* w_up [[buffer(0)]], device const float* device const uchar* w = w_up + ulong(c * n + d) * (a.hc_lr / fmt::QK * fmt::BB); const int U = int(a.hc_lr) / UNIT; float acc[8]; + #pragma clang loop unroll(full) for (int t = 0; t < 8; ++t) acc[t] = 0.0f; for (int u = li; u < U; u += 4) { float q[UNIT]; float4 sc; Dec::run(w, u, q, sc, notab); - for (int t = 0; t < GR_NTOK; ++t) { - float xs[UNIT], xsum[2] = {0.0f, 0.0f}; - for (int k = 0; k < UNIT; k += 16) load16(lo + t * int(a.hc_lr) + UNIT * u + k, xs + k); - acc[t] += unit_combine(q, sc, xs, xsum); - } + #pragma clang loop unroll(full) + for (int t = 0; t < GR_NTOK; ++t) acc[t] += gr_dot(lo + t * int(a.hc_lr) + UNIT * u, q, sc); } + #pragma clang loop unroll(full) for (int t = 0; t < GR_NTOK; ++t) { const float v = lane_sum<4>(acc[t]); if (li == 0) gate[t][c][dl] = 1.0f / (1.0f + exp(-v)); @@ -149,31 +267,41 @@ kernel void gr_up_f(device const uchar* w_up [[buffer(0)]], device const float* const int dd = int(tg.x) * 16 + j; const int hc = int(a.hc); float m = 0.0f; - for (int cc = 0; cc < hc; ++cc) m = fma(xn[(t * hc + cc) * n + dd], gate[t][cc][j], m); + if (hc == 4) { + #pragma clang loop unroll(full) + for (int cc = 0; cc < 4; ++cc) m = fma(xn[(t * 4 + cc) * n + dd], gate[t][cc][j], m); + } else { + for (int cc = 0; cc < hc; ++cc) m = fma(xn[(t * hc + cc) * n + dd], gate[t][cc][j], m); + } m = m / float(hc); const uint plane = a.n_tok * uint(n); device float* raw = xp + t * n + dd; - const float q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, - 0.00006103515625f}; - const float p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; raw[0] = m; - raw[plane] = m * q2s[j & 7]; - raw[2 * plane] = m * p4s[j & 3]; + raw[plane] = m * gr_q2s[j & 7]; + raw[2 * plane] = m * gr_p4s[j & 3]; raw[3 * plane] = (j & 1) ? m * 0.00390625f : m; - msum[t][j] = m; - } - threadgroup_barrier(mem_flags::mem_threadgroup); - if (t < GR_NTOK && j == 0) { - // the per-16 sum in xprep's order (even and odd partial sums, then added) float s0 = 0.0f, s1 = 0.0f; - for (int k = 0; k < 16; k += 2) { s0 += msum[t][k]; s1 += msum[t][k + 1]; } - xp[4 * a.n_tok * uint(n) + t * (n / 16) + tg.x] = s0 + s1; + #pragma clang loop unroll(full) + for (int k = 0; k < 16; k += 2) { + s0 += simd_shuffle_down(m, ushort(k)); + s1 += simd_shuffle_down(m, ushort(k + 1)); + } + if (j == 0) xp[4 * a.n_tok * uint(n) + t * (n / 16) + tg.x] = s0 + s1; } } -#define GRD(S, K) template [[host_name("gr_down_" #S "_" #K)]] kernel decltype(gr_down_f) gr_down_f; \ - template [[host_name("gr_down_q8_" #S "_" #K)]] kernel decltype(gr_down_f) gr_down_f; \ - template [[host_name("gr_down_q8f_" #S "_" #K)]] kernel decltype(gr_down_f) gr_down_f; +#define GRD1(S, K, R, SFX) \ + template [[host_name("gr_down_" #S "_" #K SFX)]] kernel decltype(gr_down_f) gr_down_f; \ + template [[host_name("gr_down_q8_" #S "_" #K SFX)]] kernel decltype(gr_down_f) gr_down_f; \ + template [[host_name("gr_down_q8f_" #S "_" #K SFX)]] kernel decltype(gr_down_f) gr_down_f; +#define GRD(S, K) GRD1(S, K, 1, "") GRD(8, 8) GRD(8, 4) GRD(4, 4) GRD(4, 2) GRD(16, 16) GRD(16, 8) +// two rows a simdgroup (gr_read past 4 tokens) +GRD1(8, 8, 2, "_r2") GRD1(4, 4, 2, "_r2") template [[host_name("gr_up")]] kernel decltype(gr_up_f) gr_up_f; template [[host_name("gr_up_q8")]] kernel decltype(gr_up_f) gr_up_f; + +template [[host_name("gr_down_q4_8_8")]] kernel decltype(gr_down_f<8, 8, 1, T_Q4_0, T_Q8_0>) gr_down_f<8, 8, 1, T_Q4_0, T_Q8_0>; +template [[host_name("gr_up_q4")]] kernel decltype(gr_up_f) gr_up_f; +template [[host_name("gr_down_q5_8_8")]] kernel decltype(gr_down_f<8, 8, 1, T_Q5_0, T_Q8_0>) gr_down_f<8, 8, 1, T_Q5_0, T_Q8_0>; +template [[host_name("gr_up_q5")]] kernel decltype(gr_up_f) gr_up_f; diff --git a/metal/kernels/gr2.metal b/metal/kernels/gr2.metal index 1253a11d..58342f09 100644 --- a/metal/kernels/gr2.metal +++ b/metal/kernels/gr2.metal @@ -87,6 +87,7 @@ kernel void gr_down3(device const uchar* w_down [[buffer(0)]], device const ucha threadgroup_barrier(mem_flags::mem_threadgroup); if (tid < T) { float tot = 0.0f; + #pragma clang loop unroll(full) for (int k = 0; k < 8; ++k) tot += part[k][tid]; const float rsv = rsqrt(tot / float(n) + a.eps); rs_tg[tid] = rsv; @@ -99,30 +100,44 @@ kernel void gr_down3(device const uchar* w_down [[buffer(0)]], device const ucha // the rows: simdgroup sg holds RPS rows; lanes split the stream's units int rows[RPS]; device const uchar* wr[RPS]; + #pragma clang loop unroll(full) for (int k = 0; k < RPS; ++k) { rows[k] = (int(tg.x) * 8 + sg) * RPS + k; const int rr = min(rows[k], nrows - 1); wr[k] = rr < int(a.hc_lr) ? w_down + ulong(rr) * a.rb_down : w_inject + ulong(rr - int(a.hc_lr)) * a.rb_down; } float acc[RPS][TT]; - for (int k = 0; k < RPS; ++k) + #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) { + #pragma clang loop unroll(full) for (int t = 0; t < TT; ++t) acc[k][t] = 0.0f; + } const int U = n / UNIT, u0 = c * U; for (int u = lane; u < U; u += 32) { float q[RPS][UNIT]; float4 sc[RPS]; + #pragma clang loop unroll(full) for (int k = 0; k < RPS; ++k) Dec::run(wr[k], u0 + u, q[k], sc[k], tab); - for (int t = 0; t < T; ++t) { + // TT-bounded with a T guard, so acc[][t] is indexed by a constant (registers, not the stack) + #pragma clang loop unroll(full) + for (int t = 0; t < TT; ++t) { + if (t >= T) break; float xs[UNIT], xsum[2] = {0.0f, 0.0f}; + #pragma clang loop unroll(full) for (int e = 0; e < UNIT; ++e) xs[e] = xs_tg[t][u * UNIT + e]; + #pragma clang loop unroll(full) for (int k = 0; k < RPS; ++k) acc[k][t] += unit_combine(q[k], sc[k], xs, xsum); } } - for (int k = 0; k < RPS; ++k) - for (int t = 0; t < T; ++t) { + #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) { + #pragma clang loop unroll(full) + for (int t = 0; t < TT; ++t) { + if (t >= T) break; const float v = simd_sum(acc[k][t]); if (lane == 0 && rows[k] < nrows) D[((t0 + t) * nrows + rows[k]) * hc + c] = v; } + } threadgroup_barrier(mem_flags::mem_threadgroup); } } @@ -138,7 +153,6 @@ kernel void gr_up3(device const uchar* w_up [[buffer(0)]], device const float* D constexpr int UNIT = fmt::UNIT; threadgroup float lo_tg[8][320]; threadgroup float gate[8][4][16]; - threadgroup float msum[8][16]; threadgroup uint tab[fmt::TAB]; load_tab(tab, tid, 256); const int n = int(a.n_embd), hc = int(a.hc), lr = int(a.hc_lr), nrows = lr + hc; @@ -154,17 +168,21 @@ kernel void gr_up3(device const uchar* w_up [[buffer(0)]], device const float* D device const uchar* w = w_up + ulong(c * n + d) * a.rb_up; const int U = lr / UNIT; float acc[8]; + #pragma clang loop unroll(full) for (int t = 0; t < 8; ++t) acc[t] = 0.0f; for (int u = li; u < U; u += 4) { float q[UNIT]; float4 sc; Dec::run(w, u, q, sc, tab); + #pragma clang loop unroll(full) for (int t = 0; t < GR_NTOK; ++t) { float xs[UNIT], xsum[2] = {0.0f, 0.0f}; + #pragma clang loop unroll(full) for (int e = 0; e < UNIT; ++e) xs[e] = lo_tg[t][UNIT * u + e]; acc[t] += unit_combine(q, sc, xs, xsum); } } + #pragma clang loop unroll(full) for (int t = 0; t < GR_NTOK; ++t) { const float v = lane_sum<4>(acc[t]); if (li == 0) gate[t][c][dl] = 1.0f / (1.0f + exp(-v)); @@ -174,27 +192,32 @@ kernel void gr_up3(device const uchar* w_up [[buffer(0)]], device const float* D if (t < GR_NTOK) { const int dd = int(tg.x) * 16 + j; float m = 0.0f; - for (int cc = 0; cc < hc; ++cc) { - const float xn = Rn[(t * hc + cc) * n + dd] * wn[cc * n + dd] * rs[t * hc + cc]; - m = fma(xn, gate[t][cc][j], m); + if (hc == 4) { + #pragma clang loop unroll(full) + for (int cc = 0; cc < 4; ++cc) { + const float xn = Rn[(t * 4 + cc) * n + dd] * wn[cc * n + dd] * rs[t * 4 + cc]; + m = fma(xn, gate[t][cc][j], m); + } + } else { + for (int cc = 0; cc < hc; ++cc) { + const float xn = Rn[(t * hc + cc) * n + dd] * wn[cc * n + dd] * rs[t * hc + cc]; + m = fma(xn, gate[t][cc][j], m); + } } m = m / float(hc); const uint plane = a.n_tok * uint(n); device float* raw = xp + t * n + dd; - const float q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, - 0.00006103515625f}; - const float p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; raw[0] = m; - raw[plane] = m * q2s[j & 7]; - raw[2 * plane] = m * p4s[j & 3]; + raw[plane] = m * gr_q2s[j & 7]; + raw[2 * plane] = m * gr_p4s[j & 3]; raw[3 * plane] = (j & 1) ? m * 0.00390625f : m; - msum[t][j] = m; - } - threadgroup_barrier(mem_flags::mem_threadgroup); - if (t < GR_NTOK && j == 0) { float s0 = 0.0f, s1 = 0.0f; - for (int k = 0; k < 16; k += 2) { s0 += msum[t][k]; s1 += msum[t][k + 1]; } - xp[4 * a.n_tok * uint(n) + t * (n / 16) + tg.x] = s0 + s1; + #pragma clang loop unroll(full) + for (int k = 0; k < 16; k += 2) { + s0 += simd_shuffle_down(m, ushort(k)); + s1 += simd_shuffle_down(m, ushort(k + 1)); + } + if (j == 0) xp[4 * a.n_tok * uint(n) + t * (n / 16) + tg.x] = s0 + s1; } } @@ -223,11 +246,13 @@ kernel void gr_down3b(device const uchar* w_down [[buffer(0)]], device const uch const int nrows = int(a.hc_lr) + hc; int rows[RPS]; uint4 wr[RPS][5][2]; + #pragma clang loop unroll(full) for (int k = 0; k < RPS; ++k) { rows[k] = (int(tg.x) * 8 + sg) * RPS + k; const int rr = min(rows[k], nrows - 1); device const uchar* w = rr < int(a.hc_lr) ? w_down + ulong(rr) * a.rb_down : w_inject + ulong(rr - int(a.hc_lr)) * a.rb_down; device const uint4* p = (device const uint4*)(w + 2 * (c * n)); + #pragma clang loop unroll(full) for (int i = 0; i < 5; ++i) { wr[k][i][0] = p[2 * (lane + 32 * i)]; wr[k][i][1] = p[2 * (lane + 32 * i) + 1]; @@ -261,6 +286,7 @@ kernel void gr_down3b(device const uchar* w_down [[buffer(0)]], device const uch threadgroup_barrier(mem_flags::mem_threadgroup); if (tid < T) { float tot = 0.0f; + #pragma clang loop unroll(full) for (int k = 0; k < 8; ++k) tot += part[k][tid]; const float rsv = rsqrt(tot / float(n) + a.eps); rs_tg[tid] = rsv; @@ -271,31 +297,44 @@ kernel void gr_down3b(device const uchar* w_down [[buffer(0)]], device const uch for (int d = tid; d < n; d += 256) xs_tg[t][d] = xs_tg[t][d] * wn[c * n + d] * rs_tg[t]; threadgroup_barrier(mem_flags::mem_threadgroup); float acc[RPS][TT]; - for (int k = 0; k < RPS; ++k) + #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) { + #pragma clang loop unroll(full) for (int t = 0; t < TT; ++t) acc[k][t] = 0.0f; + } + #pragma clang loop unroll(full) for (int i = 0; i < 5; ++i) { const int u = lane + 32 * i; + #pragma clang loop unroll(full) for (int k = 0; k < RPS; ++k) { const uint v[8] = {wr[k][i][0].x, wr[k][i][0].y, wr[k][i][0].z, wr[k][i][0].w, wr[k][i][1].x, wr[k][i][1].y, wr[k][i][1].z, wr[k][i][1].w}; float q[16]; + #pragma clang loop unroll(full) for (int e = 0; e < 8; ++e) { q[2 * e] = as_type(v[e] << 16); q[2 * e + 1] = as_type(v[e] & 0xFFFF0000u); } const float4 sc = float4(1.0f, 1.0f, 0.0f, 0.0f); - for (int t = 0; t < T; ++t) { + #pragma clang loop unroll(full) + for (int t = 0; t < TT; ++t) { + if (t >= T) break; float xs[16], xsum[2] = {0.0f, 0.0f}; + #pragma clang loop unroll(full) for (int e = 0; e < 16; ++e) xs[e] = xs_tg[t][u * 16 + e]; acc[k][t] += unit_combine<16>(q, sc, xs, xsum); } } } - for (int k = 0; k < RPS; ++k) - for (int t = 0; t < T; ++t) { + #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) { + #pragma clang loop unroll(full) + for (int t = 0; t < TT; ++t) { + if (t >= T) break; const float v = simd_sum(acc[k][t]); if (lane == 0 && rows[k] < nrows) D[((t0 + t) * nrows + rows[k]) * hc + c] = v; } + } threadgroup_barrier(mem_flags::mem_threadgroup); } } @@ -317,25 +356,129 @@ kernel void gr_norm_comb(device float* R [[buffer(0)]], device const float* ye [ device float* r = R + (t * a.hc + c) * n; const float g = a.apply ? 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))) : 0.0f; const float gs = 1.0f / (1.0f + exp(-sgv[t])); + float rv[10]; float ss = 0.0f; - for (uint d = tid; d < n; d += 256) { - float v = r[d]; + if (n == 2560 && topk == 10) { + device const float* yep[10]; if (a.apply) { - float s = 0.0f; - for (uint j = 0; j < topk; ++j) s += ye[plan[P_TOKENT + t * topk + j] * int(n) + d]; - const float bo = fma(gs, ysh[t * n + d], s); - v = fma(bo, g, v); - r[d] = v; + #pragma clang loop unroll(full) + for (uint j = 0; j < 10; ++j) + yep[j] = ye + uint(plan[P_TOKENT + t * 10 + j]) * 2560u + tid; + } + device const float* yshp = ysh + t * 2560u + tid; + device float* rp = r + tid; + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + float v = rp[i * 256]; + if (a.apply) { + float s = 0.0f; + #pragma clang loop unroll(full) + for (uint j = 0; j < 10; ++j) s += yep[j][i * 256]; + v = fma(fma(gs, yshp[i * 256], s), g, v); + rp[i * 256] = v; + } + rv[i] = v; + ss = fma(v, v, ss); + } + } else { + for (uint d = tid; d < n; d += 256) { + float v = r[d]; + if (a.apply) { + float s = 0.0f; + for (uint j = 0; j < topk; ++j) s += ye[plan[P_TOKENT + t * topk + j] * int(n) + d]; + const float bo = fma(gs, ysh[t * n + d], s); + v = fma(bo, g, v); + r[d] = v; + } + if (n <= 2560) rv[d / 256] = v; + ss = fma(v, v, ss); } - ss = fma(v, v, ss); } ss = simd_sum(ss); if (lane == 0) part[sg] = ss; threadgroup_barrier(mem_flags::mem_threadgroup); float tot = 0.0f; + #pragma clang loop unroll(full) for (int i = 0; i < 8; ++i) tot += part[i]; const float rs = rsqrt(tot / float(n) + a.eps); device float* o = xn + (t * a.hc + c) * n; device const float* w = w_norm + c * n; - for (uint d = tid; d < n; d += 256) o[d] = r[d] * w[d] * rs; + if (n == 2560) { + o += tid; + w += tid; + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) o[i * 256] = rv[i] * w[i * 256] * rs; + } else { + for (uint d = tid; d < n; d += 256) + o[d] = (n <= 2560 ? rv[d / 256] : r[d]) * w[d] * rs; + } +} + +// Fixed n_embd = 2560, hc = 4; grid (1, n_tok), 256 threads. Form bo once for all streams. +kernel void gr_norm_comb_4s(device float* R [[buffer(0)]], device const float* ye [[buffer(1)]], + device const float* inj [[buffer(2)]], device const float* w_norm [[buffer(3)]], + device float* xn [[buffer(4)]], constant GrArgs& a [[buffer(5)]], + device const int* plan [[buffer(6)]], device const float* ysh [[buffer(7)]], + device const float* sgv [[buffer(8)]], constant uint& topk [[buffer(9)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[4][8]; + const uint t = tg.y, n = 2560; + device float* rp = R + t * 4 * n + tid; + device float* op = xn + t * 4 * n + tid; + device const float* wp = w_norm + tid; + float rv[4][10], ss[4] = {0.0f, 0.0f, 0.0f, 0.0f}, g[4]; + #pragma clang loop unroll(full) + for (uint c = 0; c < 4; ++c) + g[c] = a.apply ? 2.0f / (1.0f + exp(-inj[t * 4 + c] / 4.0f)) : 0.0f; + const float gs = 1.0f / (1.0f + exp(-sgv[t])); + device const float* yep[10]; + if (a.apply && topk == 10) { + #pragma clang loop unroll(full) + for (uint j = 0; j < 10; ++j) + yep[j] = ye + uint(plan[P_TOKENT + t * 10 + j]) * 2560u + tid; + } + device const float* yshp = ysh + t * n + tid; + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + const uint d = tid + i * 256; + float bo = 0.0f; + if (a.apply) { + float s = 0.0f; + if (topk == 10) { + #pragma clang loop unroll(full) + for (uint j = 0; j < 10; ++j) s += yep[j][i * 256]; + } else { + for (uint j = 0; j < topk; ++j) s += ye[plan[P_TOKENT + t * topk + j] * int(n) + d]; + } + bo = fma(gs, yshp[i * 256], s); + } + #pragma clang loop unroll(full) + for (uint c = 0; c < 4; ++c) { + device float* r = rp + c * n; + float v = r[i * 256]; + if (a.apply) { v = fma(bo, g[c], v); r[i * 256] = v; } + rv[c][i] = v; + ss[c] = fma(v, v, ss[c]); + } + } + #pragma clang loop unroll(full) + for (uint c = 0; c < 4; ++c) { + const float s = simd_sum(ss[c]); + if (lane == 0) part[c][sg] = s; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + #pragma clang loop unroll(full) + for (uint c = 0; c < 4; ++c) { + float tot = 0.0f; + #pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) tot += part[c][i]; + const float rs = rsqrt(tot / float(n) + a.eps); + device float* o = op + c * n; + device const float* w = wp + c * n; + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + o[i * 256] = rv[c][i] * w[i * 256] * rs; + } + } } diff --git a/metal/kernels/mma.metal b/metal/kernels/mma.metal index 9bbb27e2..8fcc0abb 100644 --- a/metal/kernels/mma.metal +++ b/metal/kernels/mma.metal @@ -17,11 +17,57 @@ template struct DW { constexpr int UNIT = fmt::UNIT, UPB = fmt::QK / UNIT; float q[UNIT]; const float s = deq(row + (u / UPB) * fmt::BB, u % UPB, q, tab); +#pragma clang loop unroll(full) for (int j = 0; j < UNIT; ++j) w[j] = q[j] * s; p0 = u * UNIT - k0; p1 = p0 + UNIT / 2; } }; +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint* tab) { + device const uchar* b = row + u * 18; + device const ushort* p = (device const ushort*)(b + 2); +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const ushort v = p[j]; + const half2 w0 = as_type(tab[v & 255]), w1 = as_type(tab[v >> 8]); + w[2 * j] = float(w0.x); + w[2 * j + 16] = float(w0.y); + w[2 * j + 1] = float(w1.x); + w[2 * j + 17] = float(w1.y); + } + const float d = rdh(b); +#pragma clang loop unroll(full) + for (int j = 0; j < 32; ++j) w[j] *= d; + p0 = 32 * u - k0; + p1 = p0 + 16; + } +}; +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint* tab) { + device const uchar* b = row + (u >> 3) * 136; + const int ui = u & 7; + device const ushort* p = (device const ushort*)(b + 8 + 16 * ui); +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const ushort v = p[j]; + const half2 w0 = as_type(tab[v & 255]), w1 = as_type(tab[v >> 8]); + w[2 * j] = float(w0.x); + w[2 * j + 16] = float(w0.y); + w[2 * j + 1] = float(w1.x); + w[2 * j + 17] = float(w1.y); + } + const uint sh = *(device const ushort*)(b + 2); + const int ls = int((b[4 + (ui >> 1)] >> (4 * (ui & 1))) & 15) | int(((sh >> (2 * ui)) & 3) << 4); + const float d = rdh(b) * float(ls - 32); +#pragma clang loop unroll(full) + for (int j = 0; j < 32; ++j) w[j] *= d; + p0 = 32 * u - k0; + p1 = p0 + 16; + } +}; // Q2_0: (code - 1) * d template <> struct DW { static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, @@ -30,6 +76,7 @@ template <> struct DW { device const ushort* p = (device const ushort*)(b + 2 + 4 * (u & 3)); const ushort a = p[0], c = p[1]; const float d = rdh(b); +#pragma clang loop unroll(full) for (int j = 0; j < 8; ++j) { w[j] = float(int((a >> (2 * j)) & 3) - 1) * d; w[8 + j] = float(int((c >> (2 * j)) & 3) - 1) * d; @@ -44,6 +91,7 @@ template <> struct DW { device const uchar* b = row + u * 24; float q[32]; const float d = deq(b, 0, q, tab), m = rdh(b + 2); +#pragma clang loop unroll(full) for (int j = 0; j < 32; ++j) w[j] = q[j] * d + m; p0 = 32 * u - k0; p1 = p0 + 16; @@ -53,6 +101,7 @@ template <> struct DW { static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, threadgroup const uint*) { device const float* p = (device const float*)(row + 64 * u); +#pragma clang loop unroll(full) for (int j = 0; j < 16; ++j) w[j] = p[j]; p0 = 16 * u - k0; p1 = p0 + 8; @@ -64,6 +113,7 @@ template <> struct DW { device const uint4* p = (device const uint4*)(row + 32 * u); const uint4 a = p[0], c = p[1]; const uint v[8] = {a.x, a.y, a.z, a.w, c.x, c.y, c.z, c.w}; +#pragma clang loop unroll(full) for (int k = 0; k < 8; ++k) { w[2 * k] = as_type(v[k] << 16); w[2 * k + 1] = as_type(v[k] & 0xFFFF0000u); @@ -84,6 +134,7 @@ template <> struct DW { scale_min_k4(2 * c + 1, b + 4, s2, m2); const float d1 = d * s1, mm1 = dm * m1, d2 = d * s2, mm2 = dm * m2; device const uchar* q = b + 16 + 32 * c + 16 * m; +#pragma clang loop unroll(full) for (int j = 0; j < 16; ++j) { w[j] = d1 * float(q[j] & 15) - mm1; w[16 + j] = d2 * float(q[j] >> 4) - mm2; @@ -105,6 +156,7 @@ template <> struct DW { device const uchar* ql = b + 48 + 32 * c + 16 * m; device const uchar* qh = b + 16 + 16 * m; const uchar u1 = uchar(1 << (2 * c)), u2 = uchar(2 << (2 * c)); +#pragma clang loop unroll(full) for (int j = 0; j < 16; ++j) { w[j] = d1 * float((ql[j] & 15) + ((qh[j] & u1) ? 16 : 0)) - mm1; w[16 + j] = d2 * float((ql[j] >> 4) + ((qh[j] & u2) ? 16 : 0)) - mm2; @@ -124,11 +176,20 @@ template <> struct DW { device const char* sc = (device const char*)(b + 192); const float d = rdh(b + 208); const float ds1 = d * float(sc[8 * n + m + 2 * p]), ds2 = d * float(sc[8 * n + m + 2 * p + 4]); - for (int j = 0; j < 16; ++j) { - const int c1 = int((ql[j] & 15) | (((qh[j] >> (2 * p)) & 3) << 4)) - 32; - const int c2 = int((ql[j] >> 4) | (((qh[j] >> (2 * p + 4)) & 3) << 4)) - 32; - w[j] = ds1 * float(c1); - w[16 + j] = ds2 * float(c2); + const uint sh = uint(2 * p), sh4 = sh + 4u; +#pragma clang loop unroll(full) + for (int j = 0; j < 16; j += 4) { + const uint la = rd32(ql + j), ha = rd32(qh + j); + const uchar4 c1 = as_type((la & 0x0F0F0F0Fu) | (((ha >> sh) & 0x03030303u) << 4)); + const uchar4 c2 = as_type(((la >> 4) & 0x0F0F0F0Fu) | (((ha >> sh4) & 0x03030303u) << 4)); + w[j] = ds1 * float(int(c1.x) - 32); + w[j + 1] = ds1 * float(int(c1.y) - 32); + w[j + 2] = ds1 * float(int(c1.z) - 32); + w[j + 3] = ds1 * float(int(c1.w) - 32); + w[16 + j] = ds2 * float(int(c2.x) - 32); + w[16 + j + 1] = ds2 * float(int(c2.y) - 32); + w[16 + j + 2] = ds2 * float(int(c2.z) - 32); + w[16 + j + 3] = ds2 * float(int(c2.w) - 32); } p0 = 256 * (u >> 3) + 128 * n + 32 * p + 16 * m - k0; p1 = p0 + 64; @@ -149,6 +210,7 @@ template <> struct DW { const float d = rdh(b + 108); const float dl1 = d * float(l1 - 32), dl2 = d * float(l2 - 32); const uchar hb = uchar(1 << (4 * n + j)); +#pragma clang loop unroll(full) for (int l = 0; l < 16; ++l) { w[l] = dl1 * float(int((q[l] >> (2 * j)) & 3) - ((hm[l] & hb) ? 0 : 4)); w[16 + l] = dl2 * float(int((q[l + 16] >> (2 * j)) & 3) - ((hm[l + 16] & hb) ? 0 : 4)); @@ -188,6 +250,7 @@ kernel void gemv_mma(device const uchar* W [[buffer(0)]], device const float* x float w[UNIT]; int p0, p1; DW::run(W + ulong(rr) * a.row_bytes, k0 / UNIT + uu, k0, w, p0, p1, tab); +#pragma clang loop unroll(full) for (int j = 0; j < H; ++j) { w8[r * 128 + p0 + j] = w[j]; w8[r * 128 + p1 + j] = w[H + j]; diff --git a/metal/kernels/moe.metal b/metal/kernels/moe.metal index bb8fdbbc..54db1aff 100644 --- a/metal/kernels/moe.metal +++ b/metal/kernels/moe.metal @@ -20,7 +20,13 @@ #define P_ETOK (P_GSTART + P_MAXE + 4) #define P_EW (P_ETOK + P_MAXE) #define P_TOKENT (P_EW + P_MAXE) -#define P_INTS (P_TOKENT + P_MAXE) +#define P_SPLIT (P_TOKENT + P_MAXE) // function constant 26 only: [0] resident groups, then the indirect grids +#define P_IGU_R (P_SPLIT + 4) // ... of the resident groups (the GPU computes them during the host's turn) +#define P_IDN_R (P_SPLIT + 8) +#define P_IGU_M (P_SPLIT + 12) // ... and of the missing ones (after it) +#define P_IDN_M (P_SPLIT + 16) +#define P_GSLOT (P_SPLIT + 20) +#define P_INTS (P_GSLOT + P_MAXE) // Expert e's matrices: gate at W + gate_off + e * gu_stride, up at W + up_off + e * gu_stride, down at // W + down_off + e * d_stride (a pack's blobs: gate_off 0, strides = the blob; a GGUF: the three tensors). @@ -57,43 +63,77 @@ static inline device const uchar* arena_slot(uint v, device const uchar* W, devi } #endif +// Function constant 26 (the SSD-streaming decode window, with smap): the groups whose expert is in the arena when the +// routing is made come first, the missing ones after them, each part with its own indirect grids (P_SPLIT): the +// resident part runs while the host reads the missing experts. The entries keep their (token, k) mapping, so the +// combine (and every number) is the same; without it the plan is the old one, bit for bit. +constant bool kSplitResFc [[function_constant(26)]]; +constant bool kSplitRes = is_function_constant_defined(kSplitResFc) && kSplitResFc; + +// s_grp, s_ent: 64 each (the upper halves for the missing groups) static inline void plan_body(device const int* ids, device const float* w, device int* plan, constant MoeArgs& a, ushort tid, ushort nth, ushort sg, ushort lane, threadgroup int* s_grp, - threadgroup int* s_ent) { + threadgroup int* s_ent, device const uint* smap = nullptr) { const int ne = int(a.n_tok * a.topk); // one thread per expert id (nth == n_expert, at most 1024) const int e = tid; int cnt = 0; - if (e < int(a.n_expert)) + if (e < int(a.n_expert)) { + #pragma clang loop unroll_count(10) for (int i = 0; i < ne; ++i) cnt += ids[i] == e; + } const int has = cnt > 0; - const int pg = simd_prefix_exclusive_sum(has), pe = simd_prefix_exclusive_sum(cnt); - if (lane == 31) { s_grp[sg] = pg + has; s_ent[sg] = pe + cnt; } + const bool split = kSplitRes && smap != nullptr; + const int miss = (split && has && smap[e] == 0xFFFFFFFFu) ? 1 : 0; + const int h0 = has - miss, c0 = miss ? 0 : cnt, c1 = miss ? cnt : 0; + const int pg = simd_prefix_exclusive_sum(h0), pe = simd_prefix_exclusive_sum(c0); + const int pg1 = split ? simd_prefix_exclusive_sum(miss) : 0; + const int pe1 = split ? simd_prefix_exclusive_sum(c1) : 0; + if (lane == 31) { + s_grp[sg] = pg + h0; s_ent[sg] = pe + c0; + if (split) { s_grp[32 + sg] = pg1 + miss; s_ent[32 + sg] = pe1 + c1; } + } threadgroup_barrier(mem_flags::mem_threadgroup); - if (tid == 0) { - int og = 0, oe = 0; + if (sg == 0) { const int nsg = (nth + 31) / 32; - for (int i = 0; i < nsg; ++i) { - const int g = s_grp[i], n = s_ent[i]; - s_grp[i] = og; - s_ent[i] = oe; - og += g; - oe += n; + const int g0 = lane < nsg ? s_grp[lane] : 0, n0 = lane < nsg ? s_ent[lane] : 0; + const int og0 = simd_prefix_exclusive_sum(g0), oe0 = simd_prefix_exclusive_sum(n0); + if (lane < nsg) { s_grp[lane] = og0; s_ent[lane] = oe0; } + const int ng0 = simd_broadcast(og0 + g0, 31), ne0 = simd_broadcast(oe0 + n0, 31); + int og = ng0, oe = ne0; + if (split) { + const int g1 = lane < nsg ? s_grp[32 + lane] : 0, n1 = lane < nsg ? s_ent[32 + lane] : 0; + const int og1 = ng0 + simd_prefix_exclusive_sum(g1), oe1 = ne0 + simd_prefix_exclusive_sum(n1); + if (lane < nsg) { s_grp[32 + lane] = og1; s_ent[32 + lane] = oe1; } + og = simd_broadcast(og1 + g1, 31); + oe = simd_broadcast(oe1 + n1, 31); + } + if (lane == 0) { + plan[P_NG] = og; + plan[P_IGU] = int(a.gx_gu); plan[P_IGU + 1] = og; plan[P_IGU + 2] = 1; + plan[P_IDN] = int(a.gx_dn); plan[P_IDN + 1] = og; plan[P_IDN + 2] = 1; + plan[P_GSTART + og] = oe; + if (split) { + plan[P_SPLIT] = ng0; + plan[P_IGU_R] = int(a.gx_gu); plan[P_IGU_R + 1] = ng0; plan[P_IGU_R + 2] = 1; + plan[P_IDN_R] = int(a.gx_dn); plan[P_IDN_R + 1] = ng0; plan[P_IDN_R + 2] = 1; + plan[P_IGU_M] = int(a.gx_gu); plan[P_IGU_M + 1] = og - ng0; plan[P_IGU_M + 2] = 1; + plan[P_IDN_M] = int(a.gx_dn); plan[P_IDN_M + 1] = og - ng0; plan[P_IDN_M + 2] = 1; + } } - plan[P_NG] = og; - plan[P_IGU] = int(a.gx_gu); plan[P_IGU + 1] = og; plan[P_IGU + 2] = 1; - plan[P_IDN] = int(a.gx_dn); plan[P_IDN + 1] = og; plan[P_IDN + 2] = 1; - plan[P_GSTART + og] = oe; } threadgroup_barrier(mem_flags::mem_threadgroup); if (has) { - const int gi = s_grp[sg] + pg, off = s_ent[sg] + pe; + const int gi = (miss ? s_grp[32 + sg] + pg1 : s_grp[sg] + pg), off = miss ? s_ent[32 + sg] + pe1 : s_ent[sg] + pe; plan[P_GEXP + gi] = e; plan[P_GSTART + gi] = off; + if (split && !miss) plan[P_GSLOT + gi] = int(smap[e]); + const bool topk10 = (a.topk == 10u); int j = 0; + #pragma clang loop unroll_count(10) for (int i = 0; i < ne; ++i) if (ids[i] == e) { - plan[P_ETOK + off + j] = i / int(a.topk); + plan[P_ETOK + off + j] = topk10 ? ((i * 205) >> 11) : (i / int(a.topk)); plan[P_EW + off + j] = as_type(w[i]); plan[P_TOKENT + i] = off + j; ++j; @@ -105,7 +145,7 @@ kernel void moe_plan(device const int* ids [[buffer(0)]], device const float* w device int* plan [[buffer(2)]], constant MoeArgs& a [[buffer(3)]], ushort tid [[thread_position_in_threadgroup]], ushort nth [[threads_per_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { - threadgroup int s_grp[32], s_ent[32]; + threadgroup int s_grp[64], s_ent[64]; plan_body(ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent); } @@ -113,92 +153,394 @@ kernel void moe_plan(device const int* ids [[buffer(0)]], device const float* w // token t exactly as native_router.cu (pinned llama.cpp topk-moe): softmax over the 512 logits, top 10 by // probability (equal values: the lower expert index), weights = p / max(sum of the selected, 2^-14), the sum taken // by winning lane in selection order and reduced with the same xor tree. +#ifndef STRATA_MOCE_FC_METAL +#define STRATA_MOCE_FC_METAL +constant int kMoceProtectFc [[function_constant(24)]]; +constant int kMoceProtect = is_function_constant_defined(kMoceProtectFc) ? kMoceProtectFc : 3; +constant uint kMoceScaleR1BitsFc [[function_constant(25)]]; +constant float kMoceScaleR1 = is_function_constant_defined(kMoceScaleR1BitsFc) ? as_type(kMoceScaleR1BitsFc) : 1.0f; +#endif kernel void moe_route(device const float* logits [[buffer(0)]], device int* ids [[buffer(1)]], device float* w [[buffer(2)]], device int* plan [[buffer(3)]], constant MoeArgs& a [[buffer(4)]], + device const uint* smap [[buffer(5)]], ushort tid [[thread_position_in_threadgroup]], ushort nth [[threads_per_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { - threadgroup int s_grp[32], s_ent[32]; + threadgroup int s_grp[64], s_ent[64]; if (sg < a.n_tok) { device const float* l = logits + sg * 512; + const bool e512 = (a.n_expert == 512u); float v[16]; - for (int i = 0; i < 16; ++i) v[i] = l[lane + 32 * i]; - float mx = -INFINITY; - for (int i = 0; i < 16; ++i) mx = max(mx, v[i]); - for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); - float sum = 0.0f; - for (int i = 0; i < 16; ++i) { - v[i] = exp(v[i] - mx); - sum += v[i]; + if (e512) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + const float x = l[lane + 32 * i]; + v[i] = isnan(x) ? -FLT_MAX : x; + } + } else { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + const int idx = lane + 32 * i; + const float x = idx < int(a.n_expert) ? l[idx] : -INFINITY; + v[i] = isnan(x) ? -FLT_MAX : x; + } } - for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); - const float rcp = 1.0f / sum; - for (int i = 0; i < 16; ++i) { - v[i] *= rcp; - if (isnan(v[i])) v[i] = -FLT_MAX; + const bool draft_moce = (a.pad & 0x80000000u) != 0u && sg > 0; + const float base_moce = (a.pad & 0x7FFFFFFFu) != 0u ? as_type(a.pad & 0x7FFFFFFFu) : 1.0f; + const float moce_scale = draft_moce ? (base_moce * base_moce) : base_moce; + const float moce_bonus = moce_scale > 1.0f ? log(moce_scale) : 0.0f; + const int r_prot = draft_moce ? 1 : kMoceProtect; + const float r1_scale = kMoceScaleR1; + const float r1_bonus = r1_scale > 1.0f ? log(r1_scale) : 0.0f; + const float r_step_bonus = (r1_scale > 1.0f && moce_scale > 1.0f) ? (moce_bonus - r1_bonus) : moce_bonus; + const bool has_r1 = (r1_scale > 1.0f && r_prot > 1); + const bool has_moce = (moce_scale > 1.0f || r1_scale > 1.0f) && smap != nullptr; + const int min_scaled_rank = has_r1 ? 1 : r_prot; + ushort res_mask = 0; + if (has_moce) { + if (e512) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) + if (smap[lane + 32 * i] != 0xFFFFFFFFu) res_mask |= ushort(1u << i); + } else { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + const int idx = lane + 32 * i; + if (idx < int(a.n_expert) && smap[idx] != 0xFFFFFFFFu) res_mask |= ushort(1u << i); + } + } } - float selected = 0.0f, ssum = 0.0f; - for (int rank = 0; rank < 10; ++rank) { - float best = v[0]; - int ex = lane; - for (int i = 1; i < 16; ++i) - if (v[i] > best) { best = v[i]; ex = lane + 32 * i; } - for (ushort m = 16; m > 0; m >>= 1) { - const float ob = simd_shuffle_xor(best, m); - const int oe = simd_shuffle_xor(ex, m); - if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } + ushort scaled_mask = 0; + float sel_logit = -INFINITY, mx = -INFINITY; + if (!has_moce) { + for (int rank = 0; rank < 10; ++rank) { + float best = v[0]; + int ex = lane; + #pragma clang loop unroll(full) + for (int i = 1; i < 16; ++i) + if (v[i] > best) { best = v[i]; ex = lane + 32 * i; } + #pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) { + const float ob = simd_shuffle_xor(best, m); + const int oe = simd_shuffle_xor(ex, m); + if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } + } + if (rank == 0) mx = best; + if ((ex & 31) == lane) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) v[i] = i == (ex >> 5) ? -INFINITY : v[i]; + ids[sg * 10 + rank] = ex; + } + if (rank == lane) sel_logit = best; } - if ((ex & 31) == lane) { - v[ex / 32] = -INFINITY; - ids[sg * 10 + rank] = ex; - ssum += best; + } else { + for (int rank = 0; rank < 10; ++rank) { + if (has_r1) { + if (rank == 1) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + if ((res_mask & ushort(1u << i)) != 0u && v[i] > -FLT_MAX) { + v[i] += r1_bonus; + scaled_mask |= ushort(1u << i); + } + } + } else if (rank == r_prot && moce_scale > 1.0f) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + if ((scaled_mask & ushort(1u << i)) != 0u && v[i] > -FLT_MAX) { + v[i] += r_step_bonus; + } + } + } + } else if (rank == r_prot && moce_scale > 1.0f) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + if ((res_mask & ushort(1u << i)) != 0u && v[i] > -FLT_MAX) { + v[i] += moce_bonus; + scaled_mask |= ushort(1u << i); + } + } + } + float best = v[0]; + int ex = lane; + #pragma clang loop unroll(full) + for (int i = 1; i < 16; ++i) + if (v[i] > best) { best = v[i]; ex = lane + 32 * i; } + #pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) { + const float ob = simd_shuffle_xor(best, m); + const int oe = simd_shuffle_xor(ex, m); + if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } + } + if (rank == 0) mx = best; + float true_logit = best; + if (rank >= min_scaled_rank) { + const ushort win_lane = ushort(ex & 31); + const ushort win_i = ushort(ex >> 5); + const ushort win_mask = simd_broadcast(scaled_mask, win_lane); + if ((win_mask & ushort(1u << win_i)) != 0u) { + const bool r1_scaled = has_r1 && (rank < r_prot || moce_scale <= 1.0f); + true_logit = best - (r1_scaled ? r1_bonus : moce_bonus); + } + } + if ((ex & 31) == lane) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) v[i] = i == (ex >> 5) ? -INFINITY : v[i]; + ids[sg * 10 + rank] = ex; + } + if (rank == lane) sel_logit = true_logit; } - if (rank == lane) selected = best; } - for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); + float selected = (lane < 10) ? exp(sel_logit - mx) : 0.0f; + float ssum = simd_sum(selected); ssum = max(ssum, 6.103515625e-5f); - const float inv = 1.0f / ssum; - if (lane < 10) w[sg * 10 + lane] = selected * inv; + if (lane < 10) w[sg * 10 + lane] = selected * (1.0f / ssum); } threadgroup_barrier(mem_flags::mem_device); - plan_body(ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent); + plan_body(ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent, smap); } // ---- gate + up + SwiGLU ------------------------------------------------------------------------------------------- // One pass over a group's rows for all its NE entries (1..8: an expert appears once per token); NE is uniform per // threadgroup, so the kernel switches once into a body specialised for it. +// +// gu_fused: per row, gate and up decoded together and each entry's x slice loaded once for both (the 32-weight units, +// whose x load is two runs and their sums: the M5 Air, 2560 x 640, NE 3-4: Q4_K / Q5_K / Q6_K / Q8_0 -8..-25%, +// IQ4_XS -31..-36% from NE 2). The 16-weight units keep two x loads a row: the two decoded units held at once cost +// them more (IQ2_XS / IQ3_S at NR 1, NE 3: +15..+21%). The same sums in the same order either way. +template struct gu_fused { + enum { value = fmt::UNIT == 32 && (NE >= 3 || F == T_IQ4_XS) }; +}; + +template +static inline float q2_combine(ushort2 pk, float d, thread const float* xs, float xsum0) { + float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f; + #pragma clang loop unroll(full) + for (int j = 0; j < 8; j += 2) { + if (F == T_Q2_0R) { + const uint2 q = uint2(pk) >> (2 * j); + const uint2 bits = ((q | (q << 14)) & 0x00030003u) | 0x64006400u; + const float4 v = float4(as_type(bits) - 1025.0h); + s0 = fma(v.x, xs[j], s0); + s1 = fma(v.y, xs[j + 1], s1); + s2 = fma(v.z, xs[8 + j], s2); + s3 = fma(v.w, xs[8 + j + 1], s3); + } else { + s0 = fma(float(pk.x & ushort(3u << (2 * j))), xs[j], s0); + s1 = fma(float(pk.x & ushort(3u << (2 * j + 2))), xs[j + 1], s1); + s2 = fma(float(pk.y & ushort(3u << (2 * j))), xs[8 + j], s2); + s3 = fma(float(pk.y & ushort(3u << (2 * j + 2))), xs[8 + j + 1], s3); + } + } + const float dot = d * (s0 + s1) + d * (s2 + s3); + return F == T_Q2_0R ? (dot - 0.0f) : (dot - (d * xsum0 + 0.0f)); +} + +static inline float iq_combine(uint4 g, uint s16, float sc, thread const float* xs) { + float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < 2; ++k) { + const uchar4 v0 = as_type(g[k]), v1 = as_type(g[2 + k]); + const uint m0 = s16 >> (4 * k), m1 = s16 >> (8 + 4 * k); + #pragma clang loop unroll(full) + for (int j = 0; j < 4; j += 2) { + s0 = fma((m0 & (1u << j)) ? -float(v0[j]) : float(v0[j]), xs[4 * k + j], s0); + s1 = fma((m0 & (2u << j)) ? -float(v0[j + 1]) : float(v0[j + 1]), xs[4 * k + j + 1], s1); + s2 = fma((m1 & (1u << j)) ? -float(v1[j]) : float(v1[j]), xs[8 + 4 * k + j], s2); + s3 = fma((m1 & (2u << j)) ? -float(v1[j + 1]) : float(v1[j + 1]), xs[8 + 4 * k + j + 1], s3); + } + } + return sc * (s0 + s1) + sc * (s2 + s3) - 0.0f; +} + +template +static inline float deq_iq_packed(device const uchar* b, int ui, thread uint4& g, thread uint& s16, threadgroup const uint* tab) { + const int ib = ui >> 1, h = ui & 1; + if (F == T_IQ2_S) { + const uint q2 = *(device const ushort*)(b + 2 + 2 * ui); + s16 = *(device const ushort*)(b + 34 + 2 * ui); + const uint qh = uint(b[66 + ib]) >> (4 * h); + const uint i0 = (q2 & 255) | ((qh << 8) & 0x300); + const uint i1 = (q2 >> 8) | ((qh << 6) & 0x300); + const uint2 g0 = ctab ? as_type(iq2s_grid[i0]) : ((threadgroup const uint2*)tab)[i0]; + const uint2 g1 = ctab ? as_type(iq2s_grid[i1]) : ((threadgroup const uint2*)tab)[i1]; + g = uint4(g0.x, g0.y, g1.x, g1.y); + const uint sc = b[74 + ib]; + return rdh(b) * fma(float((sc >> (4 * h)) & 15), 0.25f, 0.125f); + } else if (F == T_IQ2_XXS) { + device const uchar* qs = b + 2 + 4 * ib + 2 * ui; + const uint a1 = rd32(b + 2 + 8 * ib + 4), s0 = a1 >> (14 * h); + const uint sa = s0 & 127, sb = (s0 >> 7) & 127; + s16 = (sa & 15u) | (sign4_iq2xs(sa, 1) << 4) | ((sb & 15u) << 8) | (sign4_iq2xs(sb, 1) << 12); + const uint2 g0 = ctab ? as_type(iq2xxs_grid[qs[0]]) : ((threadgroup const uint2*)tab)[qs[0]]; + const uint2 g1 = ctab ? as_type(iq2xxs_grid[qs[1]]) : ((threadgroup const uint2*)tab)[qs[1]]; + g = uint4(g0.x, g0.y, g1.x, g1.y); + return rdh(b) * fma(float(a1 >> 28), 0.25f, 0.125f); + } else if (F == T_IQ2_XS) { + device const ushort* qs = (device const ushort*)(b + 2) + 2 * ui; + const uint v0 = qs[0], v1 = qs[1]; + const uint sa = v0 >> 9, sb = v1 >> 9; + s16 = (sa & 15u) | (sign4_iq2xs(sa, 1) << 4) | ((sb & 15u) << 8) | (sign4_iq2xs(sb, 1) << 12); + const uint2 g0 = ctab ? as_type(iq2xs_grid[v0 & 511]) : ((threadgroup const uint2*)tab)[v0 & 511]; + const uint2 g1 = ctab ? as_type(iq2xs_grid[v1 & 511]) : ((threadgroup const uint2*)tab)[v1 & 511]; + g = uint4(g0.x, g0.y, g1.x, g1.y); + const uint sc = b[66 + ib]; + return rdh(b) * fma(float((sc >> (4 * h)) & 15), 0.25f, 0.125f); + } else if (F == T_IQ3_S) { + const int l0 = 2 * h, l1 = 2 * h + 1; + const uint qh = b[66 + ib]; + s16 = *(device const ushort*)(b + 74 + 2 * ui); + const uint i0 = uint(b[2 + 8 * ib + 2 * l0]) | ((qh << (8 - 2 * l0)) & 256); + const uint i1 = uint(b[2 + 8 * ib + 2 * l0 + 1]) | ((qh << (7 - 2 * l0)) & 256); + const uint i2 = uint(b[2 + 8 * ib + 2 * l1]) | ((qh << (8 - 2 * l1)) & 256); + const uint i3 = uint(b[2 + 8 * ib + 2 * l1 + 1]) | ((qh << (7 - 2 * l1)) & 256); + g = uint4(g_iq3s(tab, i0), g_iq3s(tab, i1), g_iq3s(tab, i2), g_iq3s(tab, i3)); + const uint sc = b[106 + (ib >> 1)]; + return rdh(b) * fma(float((sc >> (4 * (ib & 1))) & 15), 2.0f, 1.0f); + } else { + const uint aux = rd32(b + 2 + 64 + 4 * ib), s0 = aux >> (14 * h); + const uint sa = s0 & 127, sb = (s0 >> 7) & 127; + s16 = (sa & 15u) | (sign4_iq2xs(sa, 1) << 4) | ((sb & 15u) << 8) | (sign4_iq2xs(sb, 1) << 12); + device const uchar* qs = b + 2 + 4 * ui; + g = uint4(g_iq3xxs(tab, qs[0]), g_iq3xxs(tab, qs[1]), g_iq3xxs(tab, qs[2]), g_iq3xxs(tab, qs[3])); + return rdh(b) * fma(float(aux >> 28), 0.5f, 0.25f); + } +} + template static inline void gu_body(device const uchar* g0, device const uchar* u0, device const float* x, device const float* xs16, device float* h, device const int* etok, int c, int row0, int li, constant MoeArgs& a, threadgroup const uint* tab) { constexpr int UNIT = fmt::UNIT; float ag[NR][NE], au[NR][NE]; + #pragma clang loop unroll(full) for (int r = 0; r < NR; ++r) + #pragma clang loop unroll(full) for (int t = 0; t < NE; ++t) ag[r][t] = au[r][t] = 0.0f; const int U = int(a.n_embd) / UNIT; + if (NE == 1) { + const int tok = etok[0]; + x += tok * int(a.n_embd); + xs16 += tok * int(a.n_embd / 16); + } for (int u = li; u < U; u += LPR) { - float qg[NR][UNIT], qu[NR][UNIT]; - float4 cg[NR], cu[NR]; - for (int r = 0; r < NR; ++r) { - Dec::run(g0 + r * a.gu_row, u, qg[r], cg[r], tab); - Dec::run(u0 + r * a.gu_row, u, qu[r], cu[r], tab); - } - for (int t = 0; t < NE; ++t) { - const int tok = etok[t]; + if (NE == 1 && UNIT == 16 && F != T_Q2_0 && F != T_Q2_0R) { + float xs[UNIT], xsum[2]; + X::load(x, xs16, u, xs, xsum); + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + float q[UNIT]; + float4 c; + Dec::run(g0 + r * a.gu_row, u, q, c, tab); + ag[r][0] += unit_combine(q, c, xs, xsum); + Dec::run(u0 + r * a.gu_row, u, q, c, tab); + au[r][0] += unit_combine(q, c, xs, xsum); + } + } else if (NE == 1) { float xs[UNIT], xsum[2]; - X::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum); + X::load(x, xs16, u, xs, xsum); + for (int r = 0; r < NR; ++r) { + float q[UNIT]; + float4 c; + Dec::run(g0 + r * a.gu_row, u, q, c, tab); + ag[r][0] += unit_combine(q, c, xs, xsum); + Dec::run(u0 + r * a.gu_row, u, q, c, tab); + au[r][0] += unit_combine(q, c, xs, xsum); + } + } else if (F == T_Q2_0 || F == T_Q2_0R) { + const int boff = (u >> 2) * 18, qoff = 2 + 4 * (u & 3); + ushort2 pkg[NR], pku[NR]; + float dg[NR], du[NR]; + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + device const uchar* bg = g0 + r * a.gu_row + boff; + device const uchar* bu = u0 + r * a.gu_row + boff; + device const ushort* pg = (device const ushort*)(bg + qoff); + device const ushort* pu = (device const ushort*)(bu + qoff); + pkg[r] = ushort2(pg[0], pg[1]); + pku[r] = ushort2(pu[0], pu[1]); + dg[r] = rdh(bg); + du[r] = rdh(bu); + } + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + const int tok = etok[t]; + float xs[UNIT], xsum[2]; + X::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum); + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + ag[r][t] += q2_combine(pkg[r], dg[r], xs, xsum[0]); + au[r][t] += q2_combine(pku[r], du[r], xs, xsum[0]); + } + } + } else if (gu_fused::value) { + for (int r = 0; r < NR; ++r) { + float qg[UNIT], qu[UNIT]; + float4 cg, cu; + Dec::run(g0 + r * a.gu_row, u, qg, cg, tab); + Dec::run(u0 + r * a.gu_row, u, qu, cu, tab); + for (int t = 0; t < NE; ++t) { + const int tok = etok[t]; + float xs[UNIT], xsum[2]; + X::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum); + ag[r][t] += unit_combine(qg, cg, xs, xsum); + au[r][t] += unit_combine(qu, cu, xs, xsum); + } + } + } else if (NR == 1 && NE == 4 && (F == T_IQ2_S || F == T_IQ2_XXS || F == T_IQ2_XS || F == T_IQ3_S || F == T_IQ3_XXS)) { + const int boff = (u >> 4) * fmt::BB, ui = (LPR % 16 == 0) ? (li & 15) : (u & 15); + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + uint4 gg, gu; + uint sg, su; + const float dg = deq_iq_packed(g0 + r * a.gu_row + boff, ui, gg, sg, tab); + const float du = deq_iq_packed(u0 + r * a.gu_row + boff, ui, gu, su, tab); + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + const int tok = etok[t]; + float xs[UNIT], xsum[2]; + X::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum); + ag[r][t] += iq_combine(gg, sg, dg, xs); + au[r][t] += iq_combine(gu, su, du, xs); + } + } + } else { + #pragma clang loop unroll(full) for (int r = 0; r < NR; ++r) { - ag[r][t] += unit_combine(qg[r], cg[r], xs, xsum); - au[r][t] += unit_combine(qu[r], cu[r], xs, xsum); + float q[UNIT]; + float4 c; + Dec::run(g0 + r * a.gu_row, u, q, c, tab); + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + const int tok = etok[t]; + float xs[UNIT], xsum[2]; + X::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum); + ag[r][t] += unit_combine(q, c, xs, xsum); + } + Dec::run(u0 + r * a.gu_row, u, q, c, tab); + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + const int tok = etok[t]; + float xs[UNIT], xsum[2]; + X::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum); + au[r][t] += unit_combine(q, c, xs, xsum); + } } } } + #pragma clang loop unroll(full) for (int r = 0; r < NR; ++r) + #pragma clang loop unroll(full) for (int t = 0; t < NE; ++t) { const float gv = lane_sum(ag[r][t]), uv = lane_sum(au[r][t]); if (li == 0) h[(c + t) * int(a.n_ff) + row0 + r] = gv / (1.0f + exp(-gv)) * uv; } } +// function constant 27, with 26's plan: 1 = only the groups resident at routing time (indirect grid P_IGU_R / P_IDN_R), +// 2 = only the missing ones (P_IGU_M / P_IDN_M); 0 or unset = all (P_IGU / P_IDN) +constant int kPhaseFc [[function_constant(27)]]; +constant int kPhase = is_function_constant_defined(kPhaseFc) ? kPhaseFc : 0; + // bench probe (function constant 22): 1 = no table load (wrong values for the i-quants) constant int MOE_PROBE [[function_constant(22)]]; constant int moe_probe = is_function_constant_defined(MOE_PROBE) ? MOE_PROBE : 0; @@ -213,18 +555,19 @@ kernel void moe_gu(device const uchar* W [[buffer(0)]], device const float* x [[ ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { constexpr int RG = 32 / LPR; threadgroup uint tab[fmt::TAB]; - if (moe_probe != 1) { + if (fmt::TAB > 1 && moe_probe != 1) { load_tab(tab, tid, SG * 32); threadgroup_barrier(mem_flags::mem_threadgroup); } - const int g = int(tg.y); - if (g >= plan[P_NG]) return; + const int g = int(tg.y) + (kPhase == 2 ? plan[P_SPLIT] : 0); + if (g >= (kPhase == 1 ? plan[P_SPLIT] : plan[P_NG])) return; const int li = lane % LPR; const int row0 = ((int(tg.x) * SG + sg) * RG + lane / LPR) * NR; const ulong ex = (ulong)plan[P_GEXP + g]; const int c = plan[P_GSTART + g], ne = plan[P_GSTART + g + 1] - c; device const int* etok = plan + P_ETOK + c; - device const uchar* E = kArena ? arena_slot(smap[ex], W, W1, W2, W3) : W + ex * a.gu_stride; + const uint scode = (kArena && kPhase == 1) ? uint(plan[P_GSLOT + g]) : (kArena ? smap[ex] : 0u); + device const uchar* E = kArena ? arena_slot(scode, W, W1, W2, W3) : W + ex * a.gu_stride; device const uchar* g0 = E + a.gate_off + (ulong)row0 * a.gu_row; device const uchar* u0 = E + a.up_off + (ulong)row0 * a.gu_row; switch (ne) { @@ -250,20 +593,59 @@ static inline void dn_body(device const uchar* d0, device const float* h, device threadgroup const uint* tab) { constexpr int UNIT = fmt::UNIT; float ac[NR][NE]; + #pragma clang loop unroll(full) for (int r = 0; r < NR; ++r) + #pragma clang loop unroll(full) for (int t = 0; t < NE; ++t) ac[r][t] = 0.0f; const int U = int(a.n_ff) / UNIT; for (int u = li; u < U; u += LPR) { - float q[NR][UNIT]; - float4 sc[NR]; - for (int r = 0; r < NR; ++r) Dec::run(d0 + r * a.d_row, u, q[r], sc[r], tab); - for (int t = 0; t < NE; ++t) { + if (F == T_Q2_0 || F == T_Q2_0R) { + device const uchar* bu = d0 + (u >> 2) * 18; + const int qoff = 2 + 4 * (u & 3); + ushort2 pk[NR]; + float d[NR]; + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + device const uchar* b = bu + r * a.d_row; + device const ushort* p = (device const ushort*)(b + qoff); + pk[r] = ushort2(p[0], p[1]); + d[r] = rdh(b); + } + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + float xs[UNIT], xsum[2]; + X::load(h + (c + t) * int(a.n_ff), hs16 + (c + t) * int(a.n_ff / 16), u, xs, xsum); + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) + ac[r][t] += q2_combine(pk[r], d[r], xs, xsum[0]); + } + } else if (NE == 1) { float xs[UNIT], xsum[2]; - X::load(h + (c + t) * int(a.n_ff), hs16 + (c + t) * int(a.n_ff / 16), u, xs, xsum); - for (int r = 0; r < NR; ++r) ac[r][t] += unit_combine(q[r], sc[r], xs, xsum); + X::load(h + c * int(a.n_ff), hs16 + c * int(a.n_ff / 16), u, xs, xsum); + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + float q[UNIT]; + float4 sc; + Dec::run(d0 + r * a.d_row, u, q, sc, tab); + ac[r][0] += unit_combine(q, sc, xs, xsum); + } + } else { + float q[NR][UNIT]; + float4 sc[NR]; + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) Dec::run(d0 + r * a.d_row, u, q[r], sc[r], tab); + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + float xs[UNIT], xsum[2]; + X::load(h + (c + t) * int(a.n_ff), hs16 + (c + t) * int(a.n_ff / 16), u, xs, xsum); + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) ac[r][t] += unit_combine(q[r], sc[r], xs, xsum); + } } } + #pragma clang loop unroll(full) for (int r = 0; r < NR; ++r) + #pragma clang loop unroll(full) for (int t = 0; t < NE; ++t) { const float v = lane_sum(ac[r][t]); if (li == 0) ye[(c + t) * int(a.n_embd) + row0 + r] = as_type(plan[P_EW + c + t]) * v; @@ -280,15 +662,18 @@ kernel void moe_down(device const uchar* W [[buffer(0)]], device const float* h ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { constexpr int RG = 32 / LPR; threadgroup uint tab[fmt::TAB]; - load_tab(tab, tid, SG * 32); - threadgroup_barrier(mem_flags::mem_threadgroup); - const int g = int(tg.y); - if (g >= plan[P_NG]) return; + if (fmt::TAB > 1) { + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const int g = int(tg.y) + (kPhase == 2 ? plan[P_SPLIT] : 0); + if (g >= (kPhase == 1 ? plan[P_SPLIT] : plan[P_NG])) return; const int li = lane % LPR; const int row0 = ((int(tg.x) * SG + sg) * RG + lane / LPR) * NR; const ulong ex = (ulong)plan[P_GEXP + g]; const int c = plan[P_GSTART + g], ne = plan[P_GSTART + g + 1] - c; - device const uchar* d0 = (kArena ? arena_slot(smap[ex], W, W1, W2, W3) : W + ex * a.d_stride) + a.down_off + (ulong)row0 * a.d_row; + const uint scode = (kArena && kPhase == 1) ? uint(plan[P_GSLOT + g]) : (kArena ? smap[ex] : 0u); + device const uchar* d0 = (kArena ? arena_slot(scode, W, W1, W2, W3) : W + ex * a.d_stride) + a.down_off + (ulong)row0 * a.d_row; switch (ne) { case 1: dn_body(d0, h, hs16, ye, plan, c, row0, li, a, tab); break; case 2: dn_body(d0, h, hs16, ye, plan, c, row0, li, a, tab); break; @@ -310,8 +695,14 @@ kernel void moe_combine(device const float* ye [[buffer(0)]], device float* y [[ uint gid [[thread_position_in_grid]]) { const uint t = gid / a.n_embd, k = gid % a.n_embd; if (t >= a.n_tok) return; + device const int* te = plan + P_TOKENT + t * a.topk; float s = 0.0f; - for (uint j = 0; j < a.topk; ++j) s += ye[plan[P_TOKENT + t * a.topk + j] * int(a.n_embd) + k]; + if (a.topk == 10) { +#pragma clang loop unroll(full) + for (uint j = 0; j < 10; ++j) s += ye[te[j] * int(a.n_embd) + k]; + } else { + for (uint j = 0; j < a.topk; ++j) s += ye[te[j] * int(a.n_embd) + k]; + } y[gid] = s; } diff --git a/metal/kernels/moe_mm.metal b/metal/kernels/moe_mm.metal index ff28897a..03ac90a8 100644 --- a/metal/kernels/moe_mm.metal +++ b/metal/kernels/moe_mm.metal @@ -32,28 +32,26 @@ kernel void mm_hist(device const int* ids [[buffer(0)]], device uint* bcnt [[buf constant MmPlanArgs& a [[buffer(2)]], uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { threadgroup atomic_uint h[1024]; - for (uint e = tid; e < a.n_expert; e += 1024) atomic_store_explicit(&h[e], 0u, memory_order_relaxed); + if (tid < a.n_expert) atomic_store_explicit(&h[tid], 0u, memory_order_relaxed); threadgroup_barrier(mem_flags::mem_threadgroup); const uint i = tg.x * 1024 + tid; - if (i < a.n_ent) atomic_fetch_add_explicit(&h[ids[i]], 1u, memory_order_relaxed); + if (i < a.n_ent && ids[i] >= 0) atomic_fetch_add_explicit(&h[ids[i]], 1u, memory_order_relaxed); threadgroup_barrier(mem_flags::mem_threadgroup); - for (uint e = tid; e < a.n_expert; e += 1024) - bcnt[tg.x * a.n_expert + e] = atomic_load_explicit(&h[e], memory_order_relaxed); + if (tid < a.n_expert) + bcnt[tg.x * a.n_expert + tid] = atomic_load_explicit(&h[tid], memory_order_relaxed); } // one threadgroup of n_expert (<= 1024) threads. boff[block][e] = expert e's rows before this block (within e); // eoff[e] = expert e's first padded row; tiles[i] = (expert, row0); meta[0] = n_tiles, meta[1] = padded rows; // indirect grids: meta[2..4] = (gx_gu, n_tiles, 1) for the gate/up GEMM, meta[5..7] = (gx_down, n_tiles, 1) for the // down GEMM (gx_gu in meta[2], gx_down in meta[5] written by the host once), meta[8..10] = (rows, 1, 1) for mm_gather. +// Padding rows [r0 + tot, r0 + rows) of srt_tok (when bound at buffer 6) are cleared to -1 here so mm_clear is not needed. kernel void mm_scan(device const uint* bcnt [[buffer(0)]], device uint* boff [[buffer(1)]], device uint* eoff [[buffer(2)]], device int2* tiles [[buffer(3)]], device uint* meta [[buffer(4)]], - constant MmPlanArgs& a [[buffer(5)]], + constant MmPlanArgs& a [[buffer(5)]], device int* srt_tok [[buffer(6)]], ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], - ushort lane [[thread_index_in_simdgroup]]) { + ushort lane [[thread_index_in_simdgroup]], ushort nsg [[simdgroups_per_threadgroup]]) { threadgroup uint s_rows[32], s_tiles[32]; - // simdgroups past the launched ones must count 0 (512 threads fill only 16 of the 32 slots) - if (tid < 32) { s_rows[tid] = 0; s_tiles[tid] = 0; } - threadgroup_barrier(mem_flags::mem_threadgroup); const uint e = tid; uint tot = 0; if (e < a.n_expert) @@ -66,8 +64,8 @@ kernel void mm_scan(device const uint* bcnt [[buffer(0)]], device uint* boff [[b if (lane == 31) { s_rows[sg] = pr + rows; s_tiles[sg] = pt + nt; } threadgroup_barrier(mem_flags::mem_threadgroup); uint orow = 0, otile = 0, all_rows = 0, all_tiles = 0; - for (int i = 0; i < 32; ++i) { - if (i < int(sg)) { orow += s_rows[i]; otile += s_tiles[i]; } + for (uint i = 0; i < nsg; ++i) { + if (i < sg) { orow += s_rows[i]; otile += s_tiles[i]; } all_rows += s_rows[i]; all_tiles += s_tiles[i]; } @@ -75,6 +73,8 @@ kernel void mm_scan(device const uint* bcnt [[buffer(0)]], device uint* boff [[b const uint r0 = orow + pr, t0 = otile + pt; eoff[e] = r0; for (uint k = 0; k < nt; ++k) tiles[t0 + k] = int2(int(e), int(r0 + k * a.tm)); + if (srt_tok) + for (uint r = tot; r < rows; ++r) srt_tok[r0 + r] = -1; } if (tid == 0) { meta[0] = all_tiles; @@ -89,21 +89,56 @@ kernel void mm_scan(device const uint* bcnt [[buffer(0)]], device uint* boff [[b } } -// grid (n_blocks), 1024 threads. srt_* must be cleared to the padding value (-1 / 0) beforehand by mm_clear. +static inline uint mm_sum_nibbles(uint c) { + c = (c & 0x0F0F0F0Fu) + ((c >> 4) & 0x0F0F0F0Fu); + c = (c + (c >> 8)) & 0x00FF00FFu; + return (c + (c >> 16)) & 0xFFFFu; +} + +// grid (n_blocks), 1024 threads. Two-level deterministic intra-block rank: 32-lane sub-block nibble counts in p_cnt +// plus simd_broadcast scan within the caller's 32-lane simdgroup (no s_id[1024] threadgroup array). kernel void mm_place(device const int* ids [[buffer(0)]], device const float* w [[buffer(1)]], device const uint* boff [[buffer(2)]], device const uint* eoff [[buffer(3)]], device int* srt_tok [[buffer(4)]], device float* srt_w [[buffer(5)]], device int* pos [[buffer(6)]], constant MmPlanArgs& a [[buffer(7)]], - uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { - threadgroup int s_id[1024]; - const uint i = tg.x * 1024 + tid; + uint tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + threadgroup atomic_uint p_cnt[4096]; + const uint n4e = 4 * a.n_expert; + if (a.n_expert == 512) { + atomic_store_explicit(&p_cnt[tid], 0u, memory_order_relaxed); + atomic_store_explicit(&p_cnt[tid + 1024], 0u, memory_order_relaxed); + } else { + for (uint k = tid; k < n4e; k += 1024) atomic_store_explicit(&p_cnt[k], 0u, memory_order_relaxed); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint i = tg * 1024 + tid; const int e = i < a.n_ent ? ids[i] : -1; - s_id[tid] = e; + const uint sub = tid >> 5; + if (e >= 0 && sub < 31) { + const uint idx = (sub >> 3) * a.n_expert + uint(e); + const uint sh = (sub & 7u) * 4u; + atomic_fetch_add_explicit(&p_cnt[idx], 1u << sh, memory_order_relaxed); + } + uint rank = 0; + #pragma clang loop unroll(full) + for (ushort l = 0; l < 32; ++l) { + const int el = simd_broadcast(e, l); + if (l < lane && el == e) ++rank; + } threadgroup_barrier(mem_flags::mem_threadgroup); if (i >= a.n_ent) return; - uint rank = 0; - for (uint j = 0; j < tid; ++j) rank += s_id[j] == e; - const uint row = eoff[e] + boff[tg.x * a.n_expert + e] + rank; + if (e < 0) { pos[i] = -1; return; } + for (uint g = 0; g < (sub >> 3); ++g) { + const uint c = atomic_load_explicit(&p_cnt[g * a.n_expert + uint(e)], memory_order_relaxed); + rank += mm_sum_nibbles(c); + } + if (sub & 7u) { + const uint m = (1u << ((sub & 7u) * 4u)) - 1u; + const uint c = atomic_load_explicit(&p_cnt[(sub >> 3) * a.n_expert + uint(e)], memory_order_relaxed) & m; + rank += mm_sum_nibbles(c); + } + const uint row = eoff[e] + boff[tg * a.n_expert + e] + rank; srt_tok[row] = int(i / a.topk); srt_w[row] = w[i]; pos[i] = int(row); @@ -113,6 +148,7 @@ kernel void mm_clear(device int* srt_tok [[buffer(0)]], constant uint& n [[buffe if (gid < n) srt_tok[gid] = -1; } +#ifndef MM_GATHER_ONLY // ---- the GEMMs ----------------------------------------------------------------------------------------------------- // weights of one expert: gate at W + gate_off + e * gu_stride (+ up_off for up), down at W + down_off + e * d_stride struct MmArgs { @@ -138,6 +174,7 @@ static inline void stage_w(device const uchar* w0, uint row_bytes, int k0, threa float wv[UNIT]; int p0, p1; DW::run(w0 + ulong(r) * row_bytes, k0 / UNIT + uu, k0, wv, p0, p1, tab); +#pragma clang loop unroll(full) for (int j = 0; j < H; ++j) { dst[r * TK + p0 + j] = half(wv[j]); dst[r * TK + p1 + j] = half(wv[H + j]); @@ -145,6 +182,7 @@ static inline void stage_w(device const uchar* w0, uint row_bytes, int k0, threa } } +#ifndef MM_V3_ONLY template kernel void mm_gu(device const uchar* W [[buffer(0)]], device const float* x [[buffer(1)]], device float* H [[buffer(2)]], device const int* srt_tok [[buffer(3)]], @@ -236,6 +274,7 @@ kernel void mm_down(device const uchar* W [[buffer(0)]], device const float* H [ } // grid (T * n_embd / 256), 256 threads +#endif kernel void mm_combine(device const float* YE [[buffer(0)]], device const int* pos [[buffer(1)]], device const float* srt_w [[buffer(2)]], device float* y [[buffer(3)]], constant uint4& a [[buffer(4)]], // (T, n_embd, topk, -) @@ -245,11 +284,12 @@ kernel void mm_combine(device const float* YE [[buffer(0)]], device const int* p float s = 0.0f; for (uint k = 0; k < a.z; ++k) { const int r = pos[t * a.z + k]; - s = fma(srt_w[r], YE[ulong(r) * a.y + d], s); + if (r >= 0) s = fma(srt_w[r], YE[ulong(r) * a.y + d], s); } y[gid] = s; } +#ifndef MM_V3_ONLY #define MMGU_(F, M, N, K, S) template [[host_name("mm_gu_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_gu) mm_gu; #define MMDN_(F, M, N, K, S) template [[host_name("mm_down_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_down) mm_down; #define MMGU(F, M, N, K, S) MMGU_(F, M, N, K, S) @@ -258,7 +298,9 @@ kernel void mm_combine(device const float* YE [[buffer(0)]], device const int* p #define MM_DN_SET(F) MMDN(F, 32, 64, 64, 4) MMDN(F, 64, 64, 64, 4) MMDN(F, 16, 64, 64, 2) MMDN(F, 32, 32, 64, 2) MM_GU_SET(T_Q2_0) MM_GU_SET(T_IQ3_XXS) MM_GU_SET(T_IQ2_XXS) MM_GU_SET(T_IQ2_XS) MM_GU_SET(T_IQ2_S) MM_GU_SET(T_IQ3_S) MM_DN_SET(T_Q2_0) MM_DN_SET(T_IQ4_NL) +#endif +#endif // ---- v2: activations gathered once per layer as fp16 rows, read by matmul2d from device memory --------------------- // mm_gather: XS[row][k] = half(x[srt_tok[row]][k]) (0 on padding rows). grid (rows), 256 threads. kernel void mm_gather(device const float* x [[buffer(0)]], device const int* srt_tok [[buffer(1)]], @@ -266,9 +308,57 @@ kernel void mm_gather(device const float* x [[buffer(0)]], device const int* srt uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { const int t = srt_tok[tg.x]; device half* o = XS + ulong(tg.x) * n; - for (uint k = tid; k < n; k += 256) o[k] = t >= 0 ? half(x[ulong(t) * n + k]) : half(0.0f); + if (n == 2560) { + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + const uint k = tid + i * 256; + o[k] = t >= 0 ? half(x[ulong(t) * n + k]) : half(0.0f); + } + } else { + for (uint k = tid; k < n; k += 256) o[k] = t >= 0 ? half(x[ulong(t) * n + k]) : half(0.0f); + } +} + +// mm_gather_h: gathers straight from fp16 rows (mixh) with 128-bit uint4 copies across nth threads per row. +kernel void mm_gather_h(device const half* xh [[buffer(0)]], device const int* srt_tok [[buffer(1)]], + device half* XS [[buffer(2)]], constant uint& n [[buffer(3)]], + uint tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort nth [[threads_per_threadgroup]]) { + const int t = srt_tok[tg]; + device uint4* o8 = (device uint4*)(XS + ulong(tg) * n); + const uint n8 = n >> 3; + if (n == 2560) { + device const uint4* x8 = (device const uint4*)(xh + ulong(max(t, 0)) * n); + if (nth == 32) { + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) o8[tid + i * 32] = t >= 0 ? x8[tid + i * 32] : uint4(0u); + return; + } + if (nth == 64) { + #pragma clang loop unroll(full) + for (uint i = 0; i < 5; ++i) o8[tid + i * 64] = t >= 0 ? x8[tid + i * 64] : uint4(0u); + return; + } + if (nth == 160) { + o8[tid] = t >= 0 ? x8[tid] : uint4(0u); + o8[tid + 160] = t >= 0 ? x8[tid + 160] : uint4(0u); + return; + } + if (nth == 256) { + o8[tid] = t >= 0 ? x8[tid] : uint4(0u); + if (tid < 64) o8[tid + 256] = t >= 0 ? x8[tid + 256] : uint4(0u); + return; + } + } + if (t >= 0) { + device const uint4* x8 = (device const uint4*)(xh + ulong(t) * n); + for (uint k = tid; k < n8; k += nth) o8[k] = x8[k]; + } else { + for (uint k = tid; k < n8; k += nth) o8[k] = uint4(0u); + } } +#ifndef MM_GATHER_ONLY // The SSD-streaming arena (as in moe.metal): with function constant 20, expert e's slot is smap[e] in the chunks W, // W1..W3 (buffers 5..8 of mm_gu2 / mm_down2 / mm_gu3 / mm_down3); gate_off / up_off / down_off are offsets in the slot. #ifndef STRATA_ARENA_METAL @@ -287,6 +377,7 @@ static inline device const uchar* arena_slot(uint v, device const uchar* W, devi constant int MM_PROBE [[function_constant(5)]]; constant int mm_probe = is_function_constant_defined(MM_PROBE) ? MM_PROBE : 0; +#ifndef MM_V3_ONLY template kernel void mm_gu2(device const uchar* W [[buffer(0)]], device const half* XS [[buffer(1)]], device half* HS [[buffer(2)]], device const int2* tiles [[buffer(3)]], constant MmArgs& a [[buffer(4)]], device const uint* smap [[buffer(5)]], @@ -396,6 +487,8 @@ MM_GU2_SET(T_Q2_0) MM_GU2_SET(T_IQ3_XXS) MM_GU2_SET(T_IQ2_XXS) MM_GU2_SET(T_IQ2_ MM_DN2_SET(T_Q2_0) MMDN2(T_IQ4_NL, 16, 64, 32, 2) MMDN2(T_IQ4_NL, 16, 64, 64, 2) MMDN2(T_IQ4_NL, 16, 32, 32, 2) MMDN2(T_IQ4_NL, 32, 32, 32, 2) MMDN2(T_IQ4_NL, 16, 128, 32, 4) +#endif + // ---- v3: the weights decoded straight into matmul2d's right-input cooperative tensor (single-simdgroup scope), // no threadgroup staging and no barriers in the K loop. Each simdgroup owns TN output rows; a threadgroup SG * TN. // The right input's element layout (tests/coop_layout): runs of 4 consecutive k at one n, so a lane decodes 4 values @@ -422,14 +515,20 @@ static inline half4 w4_signed(uint g, uint s, float sc) { (s & 4) ? -float(v[2]) : float(v[2]), (s & 8) ? -float(v[3]) : float(v[3])); return half4(q * sc); } + template <> struct W4 { // 64 values a block: fp16 d, 16 bytes of 2-bit codes (value v at byte v / 4) - // fp32 math then fp16 (an fp16-math decode was measured slower: 4.2 vs 3.5 ms a layer) + // the codes c land in fp16 mantissas (1024 + c), so c - 1 and its product with d are exact in fp16: bit for bit the + // fp32 decode for every d and code byte (checked exhaustively), and 8-11% off mm_down3 at T 1024-2048 + static inline half4 dec(uint q, half d) { + const uint2 c = (uint2(q | (q << 14), (q >> 4) | (q << 10)) & 0x00030003u) | 0x64006400u; + return (as_type(c) - 1025.0h) * d; + } + template + static inline half4 run_q(device const uchar* b, uint kl4) { + return dec((b + kl4)[OFF], *(device const half*)b); + } static inline half4 run(device const uchar* row, int k, threadgroup const uint*) { - device const uchar* b = row + (k >> 6) * 18; - const float d = float(*(device const half*)b); - const uint q = b[2 + ((k & 63) >> 2)]; - return half4(float4(float(int(q & 3) - 1), float(int((q >> 2) & 3) - 1), float(int((q >> 4) & 3) - 1), - float(int(q >> 6) - 1)) * d); + return run_q<2>(row + (k >> 6) * 18, uint(k & 63) >> 2); } }; template <> struct W4 { @@ -437,10 +536,8 @@ template <> struct W4 { device const uchar* b = row + (k >> 8) * 98; const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; const uint aux = rd32(b + 2 + 64 + 4 * ib); - threadgroup const uchar* ks = (threadgroup const uchar*)(tab + 256); - const uint s = ks[(aux >> (7 * l)) & 127]; - const float sc = rdh(b) * (0.5f + float(aux >> 28)) * 0.5f; - return w4_signed(tab[b[2 + 8 * ib + 2 * l + hh]], hh ? s >> 4 : s, sc); + const float sc = rdh(b) * fma(float(aux >> 28), 0.5f, 0.25f); + return w4_signed(g_iq3xxs(tab, b[2 + ((k & 255) >> 2)]), sign4_iq2xs((aux >> (7 * l)) & 127, hh), sc); } }; template <> struct W4 { @@ -448,47 +545,43 @@ template <> struct W4 { device const uchar* b = row + (k >> 8) * 66; const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; const uint a0 = rd32(b + 2 + 8 * ib), a1 = rd32(b + 2 + 8 * ib + 4); - threadgroup const uchar* ks = (threadgroup const uchar*)(tab + 512); - const uint idx = (a0 >> (8 * l)) & 255, s = ks[(a1 >> (7 * l)) & 127]; - const float sc = rdh(b) * (0.5f + float(a1 >> 28)) * 0.25f; - return w4_signed(tab[2 * idx + hh], hh ? s >> 4 : s, sc); + const uint idx = (a0 >> (8 * l)) & 255; + const float sc = rdh(b) * fma(float(a1 >> 28), 0.25f, 0.125f); + return w4_signed(g_iq2xxs(tab, idx, hh), sign4_iq2xs((a1 >> (7 * l)) & 127, hh), sc); } }; template <> struct W4 { static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { device const uchar* b = row + (k >> 8) * 74; - const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; - const uint v = ((device const ushort*)(b + 2))[4 * ib + l]; - threadgroup const uchar* ks = (threadgroup const uchar*)(tab + 1024); - const uint idx = v & 511, s = ks[v >> 9]; + const int ib = (k & 255) >> 5, hh = (k & 7) >> 2; + const uint v = ((device const ushort*)(b + 2))[(k & 255) >> 3]; const uint scb = b[66 + ib]; - const float sc = rdh(b) * (0.5f + float((scb >> (4 * (l >> 1))) & 15)) * 0.25f; - return w4_signed(tab[2 * idx + hh], hh ? s >> 4 : s, sc); + const float sc = rdh(b) * fma(float((scb >> ((k & 16) >> 2)) & 15), 0.25f, 0.125f); + return w4_signed(g_iq2xs(tab, v & 511, hh), sign4_iq2xs(v >> 9, hh), sc); } }; template <> struct W4 { static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { device const uchar* b = row + (k >> 8) * 82; const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; - const uint qh = b[66 + ib]; - const uint idx = uint(b[2 + 4 * ib + l]) | ((qh << (8 - 2 * l)) & 0x300); - const uint s = b[2 + 32 + 4 * ib + l]; + const uint qh = b[66 + ib], o = (k & 255) >> 3; + const uint idx = uint(b[2 + o]) | ((qh << (8 - 2 * l)) & 0x300); + const uint s = b[34 + o]; const uint scb = b[74 + ib]; - const float sc = rdh(b) * (0.5f + float((scb >> (4 * (l >> 1))) & 15)) * 0.25f; - return w4_signed(tab[2 * idx + hh], hh ? s >> 4 : s, sc); + const float sc = rdh(b) * fma(float((scb >> ((k & 16) >> 2)) & 15), 0.25f, 0.125f); + return w4_signed(g_iq2s(tab, idx, hh), hh ? s >> 4 : s, sc); } }; template <> struct W4 { static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { device const uchar* b = row + (k >> 8) * 110; - const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; + const int ib = (k & 255) >> 5, hh = (k & 7) >> 2; const uint qh = b[66 + ib]; - const uint i = hh ? (uint(b[2 + 8 * ib + 2 * l + 1]) | ((qh << (7 - 2 * l)) & 256)) - : (uint(b[2 + 8 * ib + 2 * l]) | ((qh << (8 - 2 * l)) & 256)); - const uint s = b[74 + 4 * ib + l]; + const uint i = uint(b[2 + ((k & 255) >> 2)]) | ((qh << (8 - ((k & 31) >> 2))) & 256u); + const uint s = b[74 + ((k & 255) >> 3)]; const uint scb = b[106 + (ib >> 1)]; const float sc = rdh(b) * float(1 + 2 * ((scb >> (4 * (ib & 1))) & 15)); - return w4_signed(tab[i], hh ? s >> 4 : s, sc); + return w4_signed(g_iq3s(tab, i), hh ? s >> 4 : s, sc); } }; template <> struct W4 { // 256 a block: d, scales_h, scales_l[4], qs[128]; 32-groups as IQ4_NL's @@ -580,17 +673,12 @@ template <> struct W4 { // 256 a block: d, dmin, scales[12], qh[32], q template <> struct W4 { // 256 a block: ql[128], qh[64], int8 scales[16], d; halves of 128 static inline half4 run(device const uchar* row, int k, threadgroup const uint*) { device const uchar* b = row + (k >> 8) * 210; - const int j = k & 255, n = j >> 7, jj = j & 127, qd = jj >> 5, l = jj & 31; - device const uchar* ql = b + 64 * n + l + ((qd & 1) ? 32 : 0); - device const uchar* qh = b + 128 + 32 * n + l; - const float d = rdh(b + 208); - const float ds = d * float(((device const char*)(b + 192 + 8 * n))[l / 16 + 2 * qd]); - float4 w; - for (int i = 0; i < 4; ++i) { - const int q = int(((qd >> 1) ? (ql[i] >> 4) : (ql[i] & 15)) | (((qh[i] >> (2 * qd)) & 3) << 4)) - 32; - w[i] = ds * float(q); - } - return half4(w); + const uint j = uint(k & 255), j4 = j >> 4; + const uint la = rd32(b + (((j >> 1) & 64u) | (j & 63u))); + const uint ha = rd32(b + 128u + (((j >> 2) & 32u) | (j & 31u))); + const float ds = rdh(b + 208) * float(((device const char*)(b + 192))[j4]); + const uint qa = ((la >> (j4 & 4u)) & 0x0F0F0F0Fu) | (((ha >> (j4 & 6u)) & 0x03030303u) << 4); + return half4(ds * (float4(as_type(qa)) - 32.0f)); } }; @@ -607,88 +695,102 @@ template <> struct W8 { const P bl = row + (k >> 8) * 98; const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; const uint aux = rd32(bl + 2 + 64 + 4 * ib); - threadgroup const uchar* ks = (threadgroup const uchar*)(tab + 256); - const float sc = rdh(bl) * (0.5f + float(aux >> 28)) * 0.5f; - const P qs = bl + 2 + 8 * ib + hh; + const float sc = rdh(bl) * fma(float(aux >> 28), 0.5f, 0.25f); + const P qs = bl + 2 + ((k & 255) >> 2); if (mm_probe == 3) { // bench probe: no threadgroup-table lookups (wrong values) - a = w4_signed(qs[2 * l] * 0x01010101u, aux, sc); - b = w4_signed(qs[2 * l + 4] * 0x01010101u, aux >> 7, sc); + a = w4_signed(qs[0] * 0x01010101u, aux, sc); + b = w4_signed(qs[4] * 0x01010101u, aux >> 7, sc); return; } - uint s0, s1; - if (mm_probe == 4) { // the signs computed (ksigns_iq2xs[i] = i | parity(i) << 7) instead of looked up - const uint i0 = (aux >> (7 * l)) & 127, i1 = (aux >> (7 * (l + 2))) & 127; - s0 = i0 | ((popcount(i0) & 1u) << 7); - s1 = i1 | ((popcount(i1) & 1u) << 7); - } else { - s0 = ksg(tab, 256, (aux >> (7 * l)) & 127); - s1 = ksg(tab, 256, (aux >> (7 * (l + 2))) & 127); - } - a = w4_signed(g_iq3xxs(tab, qs[2 * l]), hh ? s0 >> 4 : s0, sc); - b = w4_signed(g_iq3xxs(tab, qs[2 * l + 4]), hh ? s1 >> 4 : s1, sc); + const uint j0 = (aux >> (7 * l)) & 127, j1 = (aux >> (7 * (l + 2))) & 127; + a = w4_signed(g_iq3xxs(tab, qs[0]), sign4_iq2xs(j0, hh), sc); + b = w4_signed(g_iq3xxs(tab, qs[4]), sign4_iq2xs(j1, hh), sc); } }; template <> struct W8 { - static inline void run(device const uchar* row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { - device const uchar* bl = row + (k >> 8) * 66; - const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; - const uint a0 = rd32(bl + 2 + 8 * ib), a1 = rd32(bl + 2 + 8 * ib + 4); - threadgroup const uchar* ks = (threadgroup const uchar*)(tab + 512); - const float sc = rdh(bl) * (0.5f + float(a1 >> 28)) * 0.25f; - const uint i0 = (a0 >> (8 * l)) & 255, i1 = (a0 >> (8 * (l + 2))) & 255; - const uint s0 = ks[(a1 >> (7 * l)) & 127], s1 = ks[(a1 >> (7 * (l + 2))) & 127]; - a = w4_signed(tab[2 * i0 + hh], hh ? s0 >> 4 : s0, sc); - b = w4_signed(tab[2 * i1 + hh], hh ? s1 >> 4 : s1, sc); + template + static inline void block(P bl, int k, float d, threadgroup const uint* tab, int hh, thread half4& a, thread half4& b) { + const int ib = (k & 255) >> 5, l = (k & 31) >> 3; + const P p = bl + 2 + 8 * ib; + const uint a1 = rd32(p + 4); + const float sc = d * fma(float(a1 >> 28), 0.25f, 0.125f); + const uint i0 = p[l], i1 = p[l + 2]; + const uint j0 = (a1 >> (7 * l)) & 127, j1 = (a1 >> (7 * (l + 2))) & 127; + threadgroup const uint* th = tab + hh; + a = w4_signed(ctab ? uint(iq2xxs_grid[i0] >> (32 * hh)) : th[2 * i0], sign4_iq2xs(j0, hh), sc); + b = w4_signed(ctab ? uint(iq2xxs_grid[i1] >> (32 * hh)) : th[2 * i1], sign4_iq2xs(j1, hh), sc); + } + template + static inline void run(P row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + const P bl = row + (k >> 8) * 66; + block(bl, k, rdh(bl), tab, (k & 7) >> 2, a, b); } }; template <> struct W8 { + static inline void block(device const uchar* bl, int k, float d, threadgroup const uint* tab, int hh, thread half4& a, thread half4& b) { + const int ib = (k & 255) >> 5; + device const ushort* q = (device const ushort*)(bl + 2) + ((k & 255) >> 3); + const uint v0 = q[0], v1 = q[2]; + const uint scb = bl[66 + ib]; + const float sc0 = d * fma(float(scb & 15), 0.25f, 0.125f); // l < 2 (k % 32 < 16): runs l, l + 2 take nibbles 0, 1 + const float sc1 = d * fma(float(scb >> 4), 0.25f, 0.125f); + threadgroup const uint* th = tab + hh; + a = w4_signed(ctab ? uint(iq2xs_grid[v0 & 511] >> (32 * hh)) : th[2 * (v0 & 511)], sign4_iq2xs(v0 >> 9, hh), sc0); + b = w4_signed(ctab ? uint(iq2xs_grid[v1 & 511] >> (32 * hh)) : th[2 * (v1 & 511)], sign4_iq2xs(v1 >> 9, hh), sc1); + } static inline void run(device const uchar* row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { device const uchar* bl = row + (k >> 8) * 74; - const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; - device const ushort* q = (device const ushort*)(bl + 2) + 4 * ib; - const uint v0 = q[l], v1 = q[l + 2]; - threadgroup const uchar* ks = (threadgroup const uchar*)(tab + 1024); - const uint scb = bl[66 + ib]; - const float d = rdh(bl); - const float sc0 = d * (0.5f + float((scb >> (4 * (l >> 1))) & 15)) * 0.25f; - const float sc1 = d * (0.5f + float((scb >> (4 * ((l + 2) >> 1))) & 15)) * 0.25f; - const uint s0 = ks[v0 >> 9], s1 = ks[v1 >> 9]; - a = w4_signed(tab[2 * (v0 & 511) + hh], hh ? s0 >> 4 : s0, sc0); - b = w4_signed(tab[2 * (v1 & 511) + hh], hh ? s1 >> 4 : s1, sc1); + block(bl, k, rdh(bl), tab, (k & 7) >> 2, a, b); } }; template <> struct W8 { + // grid words read through tab + hh and the sign byte as s >> 4 * hh (not a select): 8% off mm_gu3 at T 2048; + // the same for IQ2_XXS / IQ2_XS, whose sign4_iq2xs select stays (a shift form there measured 10% slower) + template + static inline void block(P bl, int k, float d, threadgroup const uint* tab, int hh, thread half4& a, thread half4& b) { + const int ib = (k & 255) >> 5, l = (k & 31) >> 3; + threadgroup const uint* th = tab + hh; + const uint sh = uint(hh) << 2; + const uint qh = bl[66 + ib], scb = bl[74 + ib]; + const P p = bl + 2 + ((k & 255) >> 3); + const uint i0 = uint(p[0]) | ((qh << (8 - 2 * l)) & 0x300); + const uint i1 = uint(p[2]) | ((qh << (4 - 2 * l)) & 0x300); + const uint s0 = p[32], s1 = p[34]; + const float sc0 = d * fma(float(scb & 15), 0.25f, 0.125f); + const float sc1 = d * fma(float(scb >> 4), 0.25f, 0.125f); + a = w4_signed(ctab ? uint(iq2s_grid[i0] >> (32 * hh)) : th[2 * i0], s0 >> sh, sc0); + b = w4_signed(ctab ? uint(iq2s_grid[i1] >> (32 * hh)) : th[2 * i1], s1 >> sh, sc1); + } template static inline void run(P row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { const P bl = row + (k >> 8) * 82; - const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; - const uint qh = bl[66 + ib], scb = bl[74 + ib]; - const float d = rdh(bl); - const uint i0 = uint(bl[2 + 4 * ib + l]) | ((qh << (8 - 2 * l)) & 0x300); - const uint i1 = uint(bl[2 + 4 * ib + l + 2]) | ((qh << (8 - 2 * (l + 2))) & 0x300); - const uint s0 = bl[2 + 32 + 4 * ib + l], s1 = bl[2 + 32 + 4 * ib + l + 2]; - const float sc0 = d * (0.5f + float((scb >> (4 * (l >> 1))) & 15)) * 0.25f; - const float sc1 = d * (0.5f + float((scb >> (4 * ((l + 2) >> 1))) & 15)) * 0.25f; - a = w4_signed(g_iq2s(tab, i0, hh), hh ? s0 >> 4 : s0, sc0); - b = w4_signed(g_iq2s(tab, i1, hh), hh ? s1 >> 4 : s1, sc1); + block(bl, k, rdh(bl), tab, (k & 7) >> 2, a, b); } }; template <> struct W8 { template - static inline void run(P row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { - const P bl = row + (k >> 8) * 110; - const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; + static inline void block(P bl, int k, float d, threadgroup const uint* tab, int hh, thread half4& a, thread half4& b) { + const int ib = k >> 5; const uint qh = bl[66 + ib], scb = bl[106 + (ib >> 1)]; - const float sc = rdh(bl) * float(1 + 2 * ((scb >> (4 * (ib & 1))) & 15)); - const int l1 = l + 2; - const uint i0 = hh ? (uint(bl[2 + 8 * ib + 2 * l + 1]) | ((qh << (7 - 2 * l)) & 256)) - : (uint(bl[2 + 8 * ib + 2 * l]) | ((qh << (8 - 2 * l)) & 256)); - const uint i1 = hh ? (uint(bl[2 + 8 * ib + 2 * l1 + 1]) | ((qh << (7 - 2 * l1)) & 256)) - : (uint(bl[2 + 8 * ib + 2 * l1]) | ((qh << (8 - 2 * l1)) & 256)); - const uint s0 = bl[74 + 4 * ib + l], s1 = bl[74 + 4 * ib + l1]; + const float sc = d * fma(float((scb >> (4 * (ib & 1))) & 15), 2.0f, 1.0f); + const P qs = bl + 2 + (k >> 2); + const int sh0 = 8 - ((k & 31) >> 2); + const uint i0 = uint(qs[0]) | ((qh << sh0) & 256u); + const uint i1 = uint(qs[4]) | ((qh << (sh0 - 4)) & 256u); + const P sp = bl + 74 + (k >> 3); + const uint s0 = sp[0], s1 = sp[2]; + if (mm_probe == 5) { // bench probe: the loads and the scale, no grid lookup / sign math (wrong values) + a = half4(half(float(i0 ^ s0) * sc)); + b = half4(half(float(i1 ^ s1) * sc)); + return; + } + a = w4_signed(g_iq3s(tab, i0), s0 >> (4 * hh), sc); + b = w4_signed(g_iq3s(tab, i1), s1 >> (4 * hh), sc); + } + template + static inline void run(P row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { if (mm_probe == 7) { - // bench probe: a repacked layout's cost - one aligned 16-byte load a lane a row a 32-group holding every - // field the lane needs (wrong values: the bytes come from an arbitrary aligned spot of the row) + const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; const uint4 r = *(device const uint4*)((ulong(row) + ulong(k >> 5) * 16) & ~ulong(15)); const uint qhv = r.z & 0xFF, scv = (r.z >> 8) & 0xFF; const float scr = float(as_type(ushort(r.w & 0xFFFF))) * float(1 + 2 * ((scv >> (4 * (ib & 1))) & 15)); @@ -698,13 +800,8 @@ template <> struct W8 { b = w4_signed(tab[j1], (r.y >> (16 + 8 * hh)) & 0xFF, scr); return; } - if (mm_probe == 5) { // bench probe: the loads and the scale, no grid lookup / sign math (wrong values) - a = half4(half(float(i0 ^ s0) * sc)); - b = half4(half(float(i1 ^ s1) * sc)); - return; - } - a = w4_signed(g_iq3s(tab, i0), hh ? s0 >> 4 : s0, sc); - b = w4_signed(g_iq3s(tab, i1), hh ? s1 >> 4 : s1, sc); + const P bl = row + (k >> 8) * 110; + block(bl, k & 255, rdh(bl), tab, (k & 7) >> 2, a, b); } }; @@ -761,22 +858,19 @@ template <> struct W8 { template <> struct W8 { static inline void run(device const uchar* row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { device const uchar* bl = row + (k >> 8) * 210; - const int j = k & 255, n = j >> 7, jj = j & 127, qd = jj >> 5, l = jj & 31; // l < 16; the pair at l + 16 - device const uchar* ql = bl + 64 * n + l + ((qd & 1) ? 32 : 0); - device const uchar* qh = bl + 128 + 32 * n + l; + const uint j = uint(k & 255), j4 = j >> 4; // l = j & 31 < 16; the pair at l + 16 + device const uchar* ql = bl + (((j >> 1) & 64u) | (j & 63u)); + device const uchar* qh = bl + 128u + (((j >> 2) & 32u) | (j & 31u)); const float d = rdh(bl + 208); - device const char* sc = (device const char*)(bl + 192 + 8 * n); - const float dsa = d * float(sc[2 * qd]), dsb = d * float(sc[1 + 2 * qd]); - // the 4-byte runs as two ushorts each (210-byte blocks: 2-aligned) - device const ushort* l2 = (device const ushort*)ql; - device const ushort* h2 = (device const ushort*)qh; - const uint la = uint(l2[0]) | (uint(l2[1]) << 16), lb = uint(l2[8]) | (uint(l2[9]) << 16); - const uint ha = uint(h2[0]) | (uint(h2[1]) << 16), hb = uint(h2[8]) | (uint(h2[9]) << 16); - const uint sl = (qd >> 1) ? 4u : 0u, shh = uint(2 * qd); - const uint4 qa = (uint4(as_type((la >> sl) & 0x0F0F0F0Fu)) | (uint4(as_type((ha >> shh) & 0x03030303u)) << 4)); - const uint4 qb = (uint4(as_type((lb >> sl) & 0x0F0F0F0Fu)) | (uint4(as_type((hb >> shh) & 0x03030303u)) << 4)); - a = half4(dsa * (float4(qa) - 32.0f)); - b = half4(dsb * (float4(qb) - 32.0f)); + device const char* sc = (device const char*)(bl + 192u + j4); + const float dsa = d * float(sc[0]), dsb = d * float(sc[1]); + const uint la = rd32(ql), lb = rd32(ql + 16); + const uint ha = rd32(qh), hb = rd32(qh + 16); + const uint sl = j4 & 4u, shh = j4 & 6u; + const uint qa = ((la >> sl) & 0x0F0F0F0Fu) | (((ha >> shh) & 0x03030303u) << 4); + const uint qb = ((lb >> sl) & 0x0F0F0F0Fu) | (((hb >> shh) & 0x03030303u) << 4); + a = half4(dsa * (float4(as_type(qa)) - 32.0f)); + b = half4(dsb * (float4(as_type(qb)) - 32.0f)); } }; @@ -788,8 +882,10 @@ kernel void mm_gu3(device const uchar* W [[buffer(0)]], device const half* XS [[ ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { threadgroup float ot[SG][TM * TN]; threadgroup uint tab[fmt::TAB]; - load_tab(tab, tid, SG * 32); - threadgroup_barrier(mem_flags::mem_threadgroup); + if (fmt::TAB > 1) { + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + } const int2 tl = tiles[tg.y]; const int e = tl.x, row0 = tl.y, n0 = (int(tg.x) * SG + sg) * TN; device const uchar* E = kArena ? arena_slot(smap[e], W, W1, W2, W3) : W + ulong(e) * a.gu_stride; @@ -798,54 +894,128 @@ kernel void mm_gu3(device const uchar* W [[buffer(0)]], device const half* XS [[ constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); matmul2d op; tensor, tensor_inline> tX((device half*)XS, dextents(a.n_embd, 1 << 30)); - auto rg = op.template get_right_input_cooperative_tensor(); - auto ru = op.template get_right_input_cooperative_tensor(); + auto rw = op.template get_right_input_cooperative_tensor(); auto tA = tX.template slice(0, row0); - auto accg = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); - auto accu = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + auto accg = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + auto accu = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); for (uint16_t i = 0; i < accg.get_capacity(); ++i) if (accg.is_valid_element(i)) { accg[i] = 0.0f; accu[i] = 0.0f; } + if (coop_k4 && TK == 32 && mm_probe != 1) { + const uint16_t half_cap = rw.get_capacity() / 2; + const auto ix0 = rw.get_multidimensional_index(uint16_t(0)); + const int kl = ix0[0], hh = (kl >> 2) & 1; + const ulong roff0 = ulong(ix0[1]) * a.gu_row; + device const uchar* g0_l = g0 + roff0; + device const uchar* u0_l = u0 + roff0; + const uint step8 = 8u * a.gu_row; + // IQ2 formats: a 256 super-block at a time, its d read once; the lane's hh is fixed for the whole K loop + constexpr bool sb = F == T_IQ2_S || F == T_IQ2_XXS || F == T_IQ2_XS; + for (int kb = 0; kb < int(a.n_embd); kb += (sb ? 256 : TK)) { + const uint sboff = uint(kb >> 8) * fmt::BB; + float dg[TN * TK / 256], du[TN * TK / 256]; + if (sb) { +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap / 4; ++i) { + dg[i] = rdh(g0_l + i * step8 + sboff); + du[i] = rdh(u0_l + i * step8 + sboff); + } + } + for (int ki = 0; ki < (sb ? 256 : TK); ki += TK) { + const int k0 = kb + ki; + tA = tX.template slice(k0, row0); + const uint boff = F == T_IQ2_XS ? sboff : (fmt::QK == 256 ? uint(k0 >> 8) * fmt::BB : (fmt::QK == 64 ? uint(k0 >> 6) * fmt::BB : 0u)); + const int j_in = fmt::QK == 256 ? ((k0 + kl) & 255) : (fmt::QK == 64 ? ((k0 + kl) & 63) : (k0 + kl)); + device const uchar* rgw = fmt::QK >= 64 ? g0_l + boff : g0_l; +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4, rgw += step8) { + half4 wg, wg2; + if (F == T_IQ2_S) W8::block(rgw, j_in, dg[i / 4], tab, hh, wg, wg2); + else if (F == T_IQ2_XXS) W8::block(rgw, j_in, dg[i / 4], tab, hh, wg, wg2); + else if (F == T_IQ2_XS) W8::block(rgw, j_in, dg[i / 4], tab, hh, wg, wg2); + else if (F == T_IQ3_S) W8::block(rgw, j_in, rdh(rgw), tab, hh, wg, wg2); + else W8::run(rgw, j_in, tab, wg, wg2); + rw[i] = wg.x; rw[i + 1] = wg.y; rw[i + 2] = wg.z; rw[i + 3] = wg.w; + const uint16_t j = i + half_cap; + rw[j] = wg2.x; rw[j + 1] = wg2.y; rw[j + 2] = wg2.z; rw[j + 3] = wg2.w; + } + op.run(tA, rw, accg); + device const uchar* ruw = fmt::QK >= 64 ? u0_l + boff : u0_l; +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4, ruw += step8) { + half4 wu, wu2; + if (F == T_IQ2_S) W8::block(ruw, j_in, du[i / 4], tab, hh, wu, wu2); + else if (F == T_IQ2_XXS) W8::block(ruw, j_in, du[i / 4], tab, hh, wu, wu2); + else if (F == T_IQ2_XS) W8::block(ruw, j_in, du[i / 4], tab, hh, wu, wu2); + else if (F == T_IQ3_S) W8::block(ruw, j_in, rdh(ruw), tab, hh, wu, wu2); + else W8::run(ruw, j_in, tab, wu, wu2); + rw[i] = wu.x; rw[i + 1] = wu.y; rw[i + 2] = wu.z; rw[i + 3] = wu.w; + const uint16_t j = i + half_cap; + rw[j] = wu2.x; rw[j + 1] = wu2.y; rw[j + 2] = wu2.z; rw[j + 3] = wu2.w; + } + op.run(tA, rw, accu); + } + } + } else for (int k0 = 0; k0 < int(a.n_embd); k0 += TK) { - const uint16_t cap = rg.get_capacity(), half_cap = cap / 2; + const uint16_t cap = rw.get_capacity(), half_cap = cap / 2; + tA = tX.template slice(k0, row0); if (!coop_k4) { for (uint16_t i = 0; i < cap; ++i) { - const auto ix = rg.get_multidimensional_index(i); + const auto ix = rw.get_multidimensional_index(i); const int k = k0 + ix[0], n = ix[1]; - rg[i] = W4::run(g0 + ulong(n) * a.gu_row, k & ~3, tab)[k & 3]; - ru[i] = W4::run(u0 + ulong(n) * a.gu_row, k & ~3, tab)[k & 3]; + rw[i] = W4::run(g0 + ulong(n) * a.gu_row, k & ~3, tab)[k & 3]; } - } else + op.run(tA, rw, accg); + for (uint16_t i = 0; i < cap; ++i) { + const auto ix = rw.get_multidimensional_index(i); + const int k = k0 + ix[0], n = ix[1]; + rw[i] = W4::run(u0 + ulong(n) * a.gu_row, k & ~3, tab)[k & 3]; + } + op.run(tA, rw, accu); + } else { #pragma clang loop unroll(full) - for (uint16_t i = 0; i < half_cap; i += 4) { - const auto ix = rg.get_multidimensional_index(i); - const auto iy = rg.get_multidimensional_index(i + half_cap); - const int k = k0 + ix[0], n = ix[1]; - device const uchar* rgw = g0 + ulong(n) * a.gu_row; - device const uchar* ruw = u0 + ulong(n) * a.gu_row; - half4 wg, wu, wg2, wu2; - if (mm_probe == 1) { // the bench probe: no decode (constant weights) - wg = half4(half(n & 7)); wu = wg; wg2 = half4(half(k & 7)); wu2 = wg2; - } else if (TK == 32) { - // the lane's two runs of this row sit in one 32-group: element i + cap / 2 is (k + 16, n) - part of - // the COOP_K4 layout (coop_probe checks it); a runtime check of it here cost 30% - W8::run(rgw, k, tab, wg, wg2); - W8::run(ruw, k, tab, wu, wu2); - } else { - const int k2 = k0 + iy[0]; - device const uchar* rgw2 = g0 + ulong(iy[1]) * a.gu_row; - device const uchar* ruw2 = u0 + ulong(iy[1]) * a.gu_row; - wg = W4::run(rgw, k, tab); wu = W4::run(ruw, k, tab); - wg2 = W4::run(rgw2, k2, tab); wu2 = W4::run(ruw2, k2, tab); + for (uint16_t i = 0; i < half_cap; i += 4) { + const auto ix = rw.get_multidimensional_index(i); + const auto iy = rw.get_multidimensional_index(i + half_cap); + const int k = k0 + ix[0], n = ix[1]; + device const uchar* rgw = g0 + ulong(n) * a.gu_row; + half4 wg, wg2; + if (mm_probe == 1) { + wg = half4(half(n & 7)); wg2 = half4(half(k & 7)); + } else if (TK == 32) { + W8::run(rgw, k, tab, wg, wg2); + } else { + const int k2 = k0 + iy[0]; + wg = W4::run(rgw, k, tab); + wg2 = W4::run(g0 + ulong(iy[1]) * a.gu_row, k2, tab); + } + rw[i] = wg.x; rw[i + 1] = wg.y; rw[i + 2] = wg.z; rw[i + 3] = wg.w; + const uint16_t j = i + half_cap; + rw[j] = wg2.x; rw[j + 1] = wg2.y; rw[j + 2] = wg2.z; rw[j + 3] = wg2.w; } - rg[i] = wg.x; rg[i + 1] = wg.y; rg[i + 2] = wg.z; rg[i + 3] = wg.w; - ru[i] = wu.x; ru[i + 1] = wu.y; ru[i + 2] = wu.z; ru[i + 3] = wu.w; - const uint16_t j = i + half_cap; - rg[j] = wg2.x; rg[j + 1] = wg2.y; rg[j + 2] = wg2.z; rg[j + 3] = wg2.w; - ru[j] = wu2.x; ru[j + 1] = wu2.y; ru[j + 2] = wu2.z; ru[j + 3] = wu2.w; + op.run(tA, rw, accg); +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4) { + const auto ix = rw.get_multidimensional_index(i); + const auto iy = rw.get_multidimensional_index(i + half_cap); + const int k = k0 + ix[0], n = ix[1]; + device const uchar* ruw = u0 + ulong(n) * a.gu_row; + half4 wu, wu2; + if (mm_probe == 1) { + wu = half4(half(n & 7)); wu2 = half4(half(k & 7)); + } else if (TK == 32) { + W8::run(ruw, k, tab, wu, wu2); + } else { + const int k2 = k0 + iy[0]; + wu = W4::run(ruw, k, tab); + wu2 = W4::run(u0 + ulong(iy[1]) * a.gu_row, k2, tab); + } + rw[i] = wu.x; rw[i + 1] = wu.y; rw[i + 2] = wu.z; rw[i + 3] = wu.w; + const uint16_t j = i + half_cap; + rw[j] = wu2.x; rw[j + 1] = wu2.y; rw[j + 2] = wu2.z; rw[j + 3] = wu2.w; + } + op.run(tA, rw, accu); } - tA = tX.template slice(k0, row0); - op.run(tA, rg, accg); - op.run(tA, ru, accu); } for (uint16_t i = 0; i < accg.get_capacity(); ++i) if (accg.is_valid_element(i)) { @@ -855,6 +1025,17 @@ kernel void mm_gu3(device const uchar* W [[buffer(0)]], device const half* XS [[ tensor, tensor_inline> tO(ot[sg], dextents(TN, TM)); accg.store(tO); simdgroup_barrier(mem_flags::mem_threadgroup); + if ((F == T_IQ2_XS || F == T_IQ3_XXS || F == T_IQ2_XXS || F == T_Q2_0 || F == T_IQ2_S || F == T_IQ3_S) && 128 % TN == 0 && (TM * TN) % 128 == 0) { + // a lane stores 4 adjacent outputs a pass, 128 a pass over the simdgroup + constexpr int RP = 128 / TN; // rows a pass + const int mr = lane / (TN / 4), f = (lane % (TN / 4)) * 4; + threadgroup const float4* ot_p = (threadgroup const float4*)(&ot[sg][mr * TN + f]); + device half* hs_p = HS + ulong(row0 + mr) * a.n_ff + n0 + f; + const uint hs_step = uint(RP) * a.n_ff; +#pragma clang loop unroll(full) + for (int m = 0; m < TM / RP; ++m, hs_p += hs_step) + *(device half4*)hs_p = half4(ot_p[m * (RP * TN / 4)]); + } else for (int it = lane; it < TM * TN; it += 32) { const int m = it / TN, f = it - m * TN; HS[ulong(row0 + m) * a.n_ff + n0 + f] = half(ot[sg][it]); @@ -868,8 +1049,10 @@ kernel void mm_down3(device const uchar* W [[buffer(0)]], device const half* HS uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], ushort tid [[thread_index_in_threadgroup]]) { threadgroup uint tab[fmt::TAB]; - load_tab(tab, tid, SG * 32); - threadgroup_barrier(mem_flags::mem_threadgroup); + if (fmt::TAB > 1) { + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + } const int2 tl = tiles[tg.y]; const int e = tl.x, row0 = tl.y, n0 = (int(tg.x) * SG + sg) * TN; device const uchar* w0 = (kArena ? arena_slot(smap[e], W, W1, W2, W3) : W + ulong(e) * a.d_stride) + a.down_off + ulong(n0) * a.d_row; @@ -881,6 +1064,38 @@ kernel void mm_down3(device const uchar* W [[buffer(0)]], device const half* HS auto acc = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); for (uint16_t i = 0; i < acc.get_capacity(); ++i) if (acc.is_valid_element(i)) acc[i] = 0.0f; + const auto ix0 = rw.get_multidimensional_index(uint16_t(0)); + const int kl = ix0[0]; + device const uchar* w0_l = w0 + ulong(ix0[1]) * a.d_row; + const uint step8 = 8u * a.d_row; + if (coop_k4 && F == T_Q2_0 && TK == 32) { + const uint16_t half_cap = rw.get_capacity() / 2; + const uint kl4 = uint(kl) >> 2; + for (int k0 = 0; k0 < int(a.n_ff); k0 += 64, w0_l += 18) { + device const uchar* rdw = w0_l; +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4, rdw += step8) { + const half4 w = W4::run_q<2>(rdw, kl4); + const half4 w2 = W4::run_q<6>(rdw, kl4); + rw[i] = w.x; rw[i + 1] = w.y; rw[i + 2] = w.z; rw[i + 3] = w.w; + const uint16_t j = i + half_cap; + rw[j] = w2.x; rw[j + 1] = w2.y; rw[j + 2] = w2.z; rw[j + 3] = w2.w; + } + tA = tH.template slice(k0, row0); + op.run(tA, rw, acc); + rdw = w0_l; +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4, rdw += step8) { + const half4 w = W4::run_q<10>(rdw, kl4); + const half4 w2 = W4::run_q<14>(rdw, kl4); + rw[i] = w.x; rw[i + 1] = w.y; rw[i + 2] = w.z; rw[i + 3] = w.w; + const uint16_t j = i + half_cap; + rw[j] = w2.x; rw[j + 1] = w2.y; rw[j + 2] = w2.z; rw[j + 3] = w2.w; + } + tA = tH.template slice(k0 + 32, row0); + op.run(tA, rw, acc); + } + } else for (int k0 = 0; k0 < int(a.n_ff); k0 += TK) { if (!coop_k4) { for (uint16_t i = 0; i < rw.get_capacity(); ++i) { @@ -891,11 +1106,13 @@ kernel void mm_down3(device const uchar* W [[buffer(0)]], device const half* HS } else if (TK == 32) { // the paired runs (element i + cap / 2 is (k + 16, n), the COOP_K4 layout): one W8 a pair const uint16_t half_cap = rw.get_capacity() / 2; + const uint boff = fmt::QK == 64 ? uint(k0 >> 6) * fmt::BB : (fmt::QK == 32 ? uint(k0 >> 5) * fmt::BB : 0u); + const int j_in = fmt::QK == 64 ? ((k0 + kl) & 63) : (fmt::QK == 32 ? kl : (k0 + kl)); + device const uchar* rdw = w0_l + boff; #pragma clang loop unroll(full) - for (uint16_t i = 0; i < half_cap; i += 4) { - const auto ix = rw.get_multidimensional_index(i); + for (uint16_t i = 0; i < half_cap; i += 4, rdw += step8) { half4 w, w2; - W8::run(w0 + ulong(ix[1]) * a.d_row, k0 + ix[0], tab, w, w2); + W8::run(rdw, j_in, tab, w, w2); rw[i] = w.x; rw[i + 1] = w.y; rw[i + 2] = w.z; rw[i + 3] = w.w; const uint16_t j = i + half_cap; rw[j] = w2.x; rw[j + 1] = w2.y; rw[j + 2] = w2.z; rw[j + 3] = w2.w; @@ -926,6 +1143,7 @@ kernel void mm_down3(device const uchar* W [[buffer(0)]], device const half* HS MM3_GU(T_Q4_K) MM3_GU(T_Q5_K) MM3_GU(T_Q6_K) MM3_GU(T_Q8_0) MM3_GU(T_IQ4_XS) MM3_GU(T_Q2_0) MM3_GU(T_IQ3_XXS) MM3_GU(T_IQ2_XXS) MM3_GU(T_IQ2_XS) MM3_GU(T_IQ2_S) MM3_GU(T_IQ3_S) MM3_DN(T_Q2_0) MM3_DN(T_IQ4_NL) MM3_DN(T_Q5_1) MM3_DN(T_Q8_0) +#ifndef MM_V3_ONLY // mm_gu4: mm_gu3 with each decoded right operand (TN x TK weights of gate and of up) applied to MS 16-row token // tiles (a tile of 16 * MS rows from the plan, tm = 16 * MS): the decode, which bounds mm_gu3, is shared by MS times // as many token rows (more padding rows per expert, but the matrix ops are cheap next to the decode). The same @@ -1016,6 +1234,8 @@ MM4_GU(T_Q2_0) MM4_GU(T_IQ3_XXS) MM4_GU(T_IQ2_XXS) MM4_GU(T_IQ2_XS) MM4_GU(T_IQ2 template kernel void mm_gu5(device const uchar* W [[buffer(0)]], device const half* XS [[buffer(1)]], device half* HS [[buffer(2)]], device const int2* tiles [[buffer(3)]], constant MmArgs& a [[buffer(4)]], + device const uint* smap [[buffer(5)]], device const uchar* W1 [[buffer(6)]], + device const uchar* W2 [[buffer(7)]], device const uchar* W3 [[buffer(8)]], uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { constexpr int TM = 16, TK = 32, BB = fmt::BB, ROWS = SG * TN, RU = BB / 2; @@ -1025,8 +1245,9 @@ kernel void mm_gu5(device const uchar* W [[buffer(0)]], device const half* XS [[ load_tab(tab, tid, SG * 32); const int2 tl = tiles[tg.y]; const int e = tl.x, row0 = tl.y, nt0 = int(tg.x) * ROWS, n0 = nt0 + sg * TN; - device const uchar* gt = W + a.gate_off + ulong(e) * a.gu_stride + ulong(nt0) * a.gu_row; - device const uchar* ut = W + a.up_off + ulong(e) * a.gu_stride + ulong(nt0) * a.gu_row; + device const uchar* E = kArena ? arena_slot(smap[e], W, W1, W2, W3) : W + ulong(e) * a.gu_stride; + device const uchar* gt = E + a.gate_off + ulong(nt0) * a.gu_row; + device const uchar* ut = E + a.up_off + ulong(nt0) * a.gu_row; constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); matmul2d op; tensor, tensor_inline> tX((device half*)XS, dextents(a.n_embd, 1 << 30)); @@ -1086,4 +1307,6 @@ kernel void mm_gu5(device const uchar* W [[buffer(0)]], device const half* XS [[ #define MMGU5_(F, N, S) template [[host_name("mm_gu5_" #F "_" #N "_" #S)]] kernel decltype(mm_gu5) mm_gu5; #define MMGU5(F, N, S) MMGU5_(F, N, S) MMGU5(T_IQ3_S, 32, 2) MMGU5(T_IQ3_S, 32, 1) MMGU5(T_IQ3_S, 16, 4) MMGU5(T_IQ3_XXS, 32, 2) MMGU5(T_IQ3_XXS, 32, 1) -MMGU5(T_IQ3_XXS, 16, 4) MMGU5(T_IQ2_S, 32, 2) MMGU5(T_IQ2_S, 32, 1) MMGU5(T_IQ2_S, 16, 4) +MMGU5(T_IQ3_XXS, 16, 4) MMGU5(T_IQ2_S, 32, 2) MMGU5(T_IQ2_S, 32, 1) MMGU5(T_IQ2_S, 16, 4) MMGU5(T_IQ2_XXS, 32, 2) MMGU5(T_IQ2_XXS, 16, 4) +#endif +#endif diff --git a/metal/kernels/prompt.metal b/metal/kernels/prompt.metal index 64833b0a..ee7d0234 100644 --- a/metal/kernels/prompt.metal +++ b/metal/kernels/prompt.metal @@ -7,50 +7,103 @@ // gr_epi_up after the HC up GEMM (G = lo . w_up, T x hc*n): mixed = mean_c xn[c] * sigmoid(G[c]) -> float and // half (the next GEMMs' operand) // swiglu_h h = half(silu(g) * u) -// moe_out bo = y_routed + sigmoid(sg) * y_shared +// gemm_hc_down, gemm_swiglu_up (after gemm_hc_up) the HC down and shared-expert up GEMMs with gr_epi_down / +// swiglu_h as register epilogues +// moe_out bo = y_routed + sigmoid(sg) * y_shared (prompt_chain.inc; the prompt chunk folds it into the next +// read: gr_norm_h_moe, gr_apply_moe, ple_moe_gate_norm; with mm_combine folded in too: gr_norm_h_moec, +// gr_apply_moec) kernel void route_topk(device const float* logits [[buffer(0)]], device int* ids [[buffer(1)]], - device float* w [[buffer(2)]], constant uint& n_tok [[buffer(3)]], + device float* w [[buffer(2)]], constant uint& n_tok_packed [[buffer(3)]], + device const uint* smap [[buffer(4)]], constant float& prompt_moce [[buffer(5)]], uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + const uint n_tok = n_tok_packed & 0xFFFFu; + const int n_exp = (n_tok_packed >> 16) ? int(n_tok_packed >> 16) : 512; const uint t = tg.x * 8 + sg; if (t >= n_tok) return; - device const float* l = logits + ulong(t) * 512; + device const float* l = logits + ulong(t) * ulong(n_exp); float v[16]; - for (int i = 0; i < 16; ++i) v[i] = l[lane + 32 * i]; - float mx = -INFINITY; - for (int i = 0; i < 16; ++i) mx = max(mx, v[i]); - for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); - float sum = 0.0f; - for (int i = 0; i < 16; ++i) { - v[i] = exp(v[i] - mx); - sum += v[i]; - } - for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); - const float rcp = 1.0f / sum; - for (int i = 0; i < 16; ++i) { - v[i] *= rcp; - if (isnan(v[i])) v[i] = -FLT_MAX; - } - float selected = 0.0f, ssum = 0.0f; + if (n_exp == 512) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + const float x = l[lane + 32 * i]; + v[i] = isnan(x) ? -FLT_MAX : x; + } + } else { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + const int idx = lane + 32 * i; + const float x = idx < n_exp ? l[idx] : -INFINITY; + v[i] = isnan(x) ? -FLT_MAX : x; + } + } + const int tail_exact = int(prompt_moce) / 10000; + const float rem_moce = prompt_moce - float(tail_exact * 10000); + const int body_prot = int(rem_moce) / 100; + const float body_bonus = rem_moce - float(body_prot * 100); + const bool in_tail = tail_exact > 0 && (t + uint(tail_exact) >= n_tok); + const int p_prot = in_tail ? 10 : body_prot; + const float p_bonus = in_tail ? 0.0f : min(body_bonus, 40.0f); + ushort scaled_mask = 0; + float sel_logit = -INFINITY, mx = -INFINITY; for (int rank = 0; rank < 10; ++rank) { + if (rank == p_prot && p_bonus > 0.0f) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + const int idx = lane + 32 * i; + if ((n_exp == 512 || idx < n_exp) && v[i] > -FLT_MAX && smap[idx] != 0xFFFFFFFFu) { + v[i] += p_bonus; + scaled_mask |= ushort(1u << i); + } + } + } float best = v[0]; int ex = lane; + #pragma clang loop unroll(full) for (int i = 1; i < 16; ++i) if (v[i] > best) { best = v[i]; ex = lane + 32 * i; } + #pragma clang loop unroll(full) for (ushort m = 16; m > 0; m >>= 1) { const float ob = simd_shuffle_xor(best, m); const int oe = simd_shuffle_xor(ex, m); if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } } + if (rank == 0) mx = best; + float true_logit = best; + if (rank >= p_prot && p_bonus > 0.0f) { + const ushort win_lane = ushort(ex & 31); + const ushort win_i = ushort(ex >> 5); + const ushort win_mask = simd_broadcast(scaled_mask, win_lane); + if ((win_mask & ushort(1u << win_i)) != 0u) true_logit = best - p_bonus; + } if ((ex & 31) == lane) { - v[ex / 32] = -INFINITY; + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) v[i] = i == (ex >> 5) ? -INFINITY : v[i]; ids[t * 10 + rank] = ex; - ssum += best; } - if (rank == lane) selected = best; + if (rank == lane) sel_logit = true_logit; + } + float selected = (lane < 10) ? exp(sel_logit - mx) : 0.0f; + float ssum = simd_sum(selected); + if (prompt_moce < 0.0f) { + const int skip_tail = int(-prompt_moce); + const float tau = (-prompt_moce) - float(skip_tail); + const bool skip_active = tau > 0.0f && !(skip_tail > 0 && (t + uint(skip_tail) >= n_tok)); + if (skip_active) { + float kept_prob = (lane < 10) ? selected : 0.0f; + if (lane >= 3 && lane < 10) { + const int ex_r = ids[t * 10 + lane]; + if (smap[ex_r] == 0xFFFFFFFFu && selected < tau * ssum) { + ids[t * 10 + lane] = -1; + kept_prob = 0.0f; + } + } + selected = kept_prob; + ssum = kept_prob; + for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); + } } - for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); ssum = max(ssum, 6.103515625e-5f); if (lane < 10) w[t * 10 + lane] = selected * (1.0f / ssum); } @@ -121,57 +174,478 @@ kernel void gr_norm_h(device float* R [[buffer(0)]], device const float* bo [[bu const uint c = tg.x, t = tg.y, n = a.n_embd; device float* r = R + (t * a.hc + c) * n; const float g = a.apply ? 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))) : 0.0f; - float ss = 0.0f; - for (uint d = tid; d < n; d += 256) { + float rv[10], ss = 0.0f; + for (uint d = tid, i = 0; d < n; d += 256, ++i) { float v = r[d]; if (a.apply) { v = fma(bo[t * n + d], g, v); r[d] = v; } + rv[i] = v; ss = fma(v, v, ss); } ss = simd_sum(ss); if (lane == 0) part[sg] = ss; threadgroup_barrier(mem_flags::mem_threadgroup); float tot = 0.0f; +#pragma clang loop unroll(full) for (int i = 0; i < 8; ++i) tot += part[i]; const float rs = rsqrt(tot / float(n) + a.eps); device half* o = xnh + (t * a.hc + c) * n; device const float* w = w_norm + c * n; - for (uint d = tid; d < n; d += 256) o[d] = half(r[d] * w[d] * rs); + for (uint d = tid, i = 0; d < n; d += 256, ++i) o[d] = half(rv[i] * w[d] * rs); +} + +// gr_norm_h's tail for one token's hc streams: the sums of squares ss[c] (per thread) to the norm, xnh = R' * w * rs +inline void gr_norm_h_out(thread const float (*rv)[10], device const float* w_norm, device half* xnh, uint t, uint hc, + uint n, float eps, thread const float* ss, threadgroup float (*part)[8], ushort tid, + ushort sg, ushort lane) { + for (uint c = 0; c < hc; ++c) { + const float s = simd_sum(ss[c]); + if (lane == 0) part[c][sg] = s; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint c = 0; c < hc; ++c) { + float tot = 0.0f; +#pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) tot += part[c][i]; + const float rs = rsqrt(tot / float(n) + eps); + device half* o = xnh + (t * hc + c) * n; + device const float* w = w_norm + c * n; + for (uint d = tid, i = 0; d < n; d += 256, ++i) o[d] = half(rv[c][i] * w[d] * rs); + } +} + +// stream-wise reduction keeps one residual stream plus bo live, rather than all four residual streams. +inline void gr_norm_h_stream_out(thread const float* rv, device const float* w, device half* o, uint n, float eps, + float ss, threadgroup float* part, ushort tid, ushort sg, ushort lane) { + ss = simd_sum(ss); + if (lane == 0) part[sg] = ss; + threadgroup_barrier(mem_flags::mem_threadgroup); + float tot = 0.0f; +#pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) tot += part[i]; + const float rs = rsqrt(tot / float(n) + eps); + if (n == 2560) { +#pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + const uint d = tid + i * 256; + o[d] = half(rv[i] * w[d] * rs); + } + } else { + for (uint d = tid, i = 0; d < n; d += 256, ++i) o[d] = half(rv[i] * w[d] * rs); + } +} + +// gr_norm_h4: gr_norm_h with one threadgroup per token over all hc streams, so bo is read once a token rather than +// hc times; R', xnh bitwise gr_norm_h's. grid (1, T), 256 threads; hc <= 4 +kernel void gr_norm_h4(device float* R [[buffer(0)]], device const float* bo [[buffer(1)]], + device const float* inj [[buffer(2)]], device const float* w_norm [[buffer(3)]], + device half* xnh [[buffer(4)]], constant GrArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[4][8]; + const uint t = tg.y, n = a.n_embd, hc = a.hc; + device float* r = R + t * hc * n; + if (!a.apply) { + for (uint c = 0; c < hc; ++c) { + float rv[10], ss = 0.0f; + if (n == 2560) { + device const float* rc = r + c * 2560 + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + const float v = rc[i * 256]; + rv[i] = v; + ss = fma(v, v, ss); + } + } else { + for (uint d = tid, i = 0; d < n; d += 256, ++i) { + const float v = r[c * n + d]; + rv[i] = v; + ss = fma(v, v, ss); + } + } + gr_norm_h_stream_out(rv, w_norm + c * n, xnh + (t * hc + c) * n, n, a.eps, ss, part[c], tid, sg, lane); + } + return; + } + float bv[10]; + if (n == 2560) { + device const float* brow = bo + t * 2560 + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) bv[i] = brow[i * 256]; + } else { + for (uint d = tid, i = 0; d < n; d += 256, ++i) bv[i] = bo[t * n + d]; + } + for (uint c = 0; c < hc; ++c) { + const float g = 2.0f / (1.0f + exp(-inj[t * hc + c] / float(hc))); + float rv[10], ss = 0.0f; + if (n == 2560) { + device float* rc = r + c * 2560 + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + const float v = fma(bv[i], g, rc[i * 256]); + rc[i * 256] = v; + rv[i] = v; + ss = fma(v, v, ss); + } + } else { + for (uint d = tid, i = 0; d < n; d += 256, ++i) { + const float v = fma(bv[i], g, r[c * n + d]); + r[c * n + d] = v; + rv[i] = v; + ss = fma(v, v, ss); + } + } + gr_norm_h_stream_out(rv, w_norm + c * n, xnh + (t * hc + c) * n, n, a.eps, ss, part[c], tid, sg, lane); + } +} + +// gr_norm_h_moe: gr_norm_h with the sublayer output bo formed on the fly as moe_out - bo = fma(sigmoid(sg[t]), ys, yr) +// - so the MoE's write never round-trips through bo; R', xnh bitwise moe_out + gr_norm_h. Always applies. +// One threadgroup per token t (all hc streams: bo is formed once a token), 256 threads; hc <= 4. +kernel void gr_norm_h_moe(device float* R [[buffer(0)]], device const float* yr [[buffer(1)]], + device const float* ys [[buffer(2)]], device const float* sgv [[buffer(3)]], + device const float* inj [[buffer(4)]], device const float* w_norm [[buffer(5)]], + device half* xnh [[buffer(6)]], constant GrArgs& a [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[4][8]; + const uint t = tg.y, n = a.n_embd, hc = a.hc; + device float* r = R + t * hc * n; + float rv[4][10], g[4], ss[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + for (uint c = 0; c < hc; ++c) g[c] = 2.0f / (1.0f + exp(-inj[t * hc + c] / float(hc))); + const float gs = 1.0f / (1.0f + exp(-sgv[t])); + for (uint d = tid, i = 0; d < n; d += 256, ++i) { + const float bo = fma(gs, ys[t * n + d], yr[t * n + d]); + for (uint c = 0; c < hc; ++c) { + const float v = fma(bo, g[c], r[c * n + d]); + r[c * n + d] = v; + rv[c][i] = v; + ss[c] = fma(v, v, ss[c]); + } + } + gr_norm_h_out(rv, w_norm, xnh, t, hc, n, a.eps, ss, part, tid, sg, lane); +} + +// token t's routed slots for moe_comb: compacted valid YE row offsets (uint(r) * n) and their weights in threadgroup memory +inline uint moe_comb_slots(device const int* pos, device const float* srt_w, uint t, uint topk, uint n, + threadgroup uint* roff, threadgroup float* wk, ushort tid) { + int r = -1; + float w = 0.0f; + if (tid < topk) { + r = pos[t * topk + tid]; + if (r >= 0) w = srt_w[r]; + } + const bool valid = r >= 0; + const ushort prefix = simd_prefix_exclusive_sum(ushort(valid)); + const ushort cnt = simd_sum(ushort(valid)); + if (valid) { + roff[prefix] = uint(r) * n; + wk[prefix] = w; + } + if (tid == 0) roff[16] = cnt; + threadgroup_barrier(mem_flags::mem_threadgroup); + return roff[16]; +} + +// the routed experts' write at (t, d) as mm_combine (moe_mm.metal) forms it: sum_k w(t, k) * YE[pos(t, k)][d], k in order +inline float moe_comb(device const float* YE, threadgroup const uint* roff, threadgroup const float* wk, uint cnt, + uint d) { + float s = 0.0f; + for (uint k = 0; k < cnt; ++k) s = fma(wk[k], YE[roff[k] + d], s); + return s; +} + +// gr_norm_h_moec: gr_norm_h_moe with mm_combine folded in - yr formed from the expert rows YE (pos, srt_w: mm_place's) +// as it is read, so it is never written; R', xnh bitwise mm_combine + gr_norm_h_moe. grid (1, T), 256 threads; +// hc <= 4, topk <= 16 +kernel void gr_norm_h_moec(device float* R [[buffer(0)]], device const float* YE [[buffer(1)]], + device const int* pos [[buffer(2)]], device const float* srt_w [[buffer(3)]], + device const float* ys [[buffer(4)]], device const float* sgv [[buffer(5)]], + device const float* inj [[buffer(6)]], device const float* w_norm [[buffer(7)]], + device half* xnh [[buffer(8)]], constant GrArgs& a [[buffer(9)]], constant uint& topk [[buffer(10)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[4][8]; + threadgroup uint roff[17]; + threadgroup float wk[16]; + const uint t = tg.y, n = a.n_embd, hc = a.hc; + const uint cnt = moe_comb_slots(pos, srt_w, t, topk, n, roff, wk, tid); + device float* r = R + t * hc * n; + float bv[10]; + const float gs = 1.0f / (1.0f + exp(-sgv[t])); + if (n == 2560) { + // k outer: each slot's weight and row read once, its ten loads in flight together; s[i]'s fma order is moe_comb's + float s[10] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + for (uint k = 0; k < cnt; ++k) { + const float w = wk[k]; + device const float* yk = YE + roff[k] + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) s[i] = fma(w, yk[i * 256], s[i]); + } + device const float* yrow = ys + t * 2560 + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) bv[i] = fma(gs, yrow[i * 256], s[i]); + } else { + for (uint d = tid, i = 0; d < n; d += 256, ++i) bv[i] = fma(gs, ys[t * n + d], moe_comb(YE, roff, wk, cnt, d)); + } + for (uint c = 0; c < hc; ++c) { + const float g = 2.0f / (1.0f + exp(-inj[t * hc + c] / float(hc))); + float rv[10], ss = 0.0f; + if (n == 2560) { + device float* rc = r + c * 2560 + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + const float v = fma(bv[i], g, rc[i * 256]); + rc[i * 256] = v; + rv[i] = v; + ss = fma(v, v, ss); + } + } else { + for (uint d = tid, i = 0; d < n; d += 256, ++i) { + const float v = fma(bv[i], g, r[c * n + d]); + r[c * n + d] = v; + rv[i] = v; + ss = fma(v, v, ss); + } + } + gr_norm_h_stream_out(rv, w_norm + c * n, xnh + (t * hc + c) * n, n, a.eps, ss, part[c], tid, sg, lane); + } +} + +// gr_apply_moe: the last layer's MoE write applied to R alone (R' of moe_out + gr_norm; its xn has no reader). +// grid (hc * n / 256, T), 256 threads +kernel void gr_apply_moe(device float* R [[buffer(0)]], device const float* yr [[buffer(1)]], + device const float* ys [[buffer(2)]], device const float* sgv [[buffer(3)]], + device const float* inj [[buffer(4)]], constant GrArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + const uint ch = tg.x * 256 + tid, t = tg.y, n = a.n_embd; + if (ch >= a.hc * n) return; + const uint c = ch / n, d = ch - c * n; + const float g = 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))); + const float gs = 1.0f / (1.0f + exp(-sgv[t])); + device float* r = R + t * a.hc * n + ch; + *r = fma(fma(gs, ys[t * n + d], yr[t * n + d]), g, *r); +} + +// gr_apply_moec: gr_apply_moe with mm_combine folded in (as gr_norm_h_moec); R' bitwise mm_combine + gr_apply_moe. +// grid (ceil(n / 256), T), 256 threads: a thread's d over all hc streams (yr formed once); topk <= 16 +kernel void gr_apply_moec(device float* R [[buffer(0)]], device const float* YE [[buffer(1)]], + device const int* pos [[buffer(2)]], device const float* srt_w [[buffer(3)]], + device const float* ys [[buffer(4)]], device const float* sgv [[buffer(5)]], + device const float* inj [[buffer(6)]], constant GrArgs& a [[buffer(7)]], + constant uint& topk [[buffer(8)]], uint3 tg [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]]) { + threadgroup uint roff[17]; + threadgroup float wk[16]; + const uint d = tg.x * 256 + tid, t = tg.y, n = a.n_embd, hc = a.hc; + const uint cnt = moe_comb_slots(pos, srt_w, t, topk, n, roff, wk, tid); + if (d >= n) return; + const float gs = 1.0f / (1.0f + exp(-sgv[t])); + const float bo = fma(gs, ys[t * n + d], moe_comb(YE, roff, wk, cnt, d)); + device float* r = R + t * hc * n + d; + for (uint c = 0; c < hc; ++c) r[c * n] = fma(bo, 2.0f / (1.0f + exp(-inj[t * hc + c] / float(hc))), r[c * n]); } // gemm_hc_up: G = lo . w_up^T with w_up's rows interleaved by dq_f16 (perm = hc: row d * hc + c), so a TN-column // tile holds TN / hc outputs d with all their streams; the epilogue mixed[t][d] = mean_c xn[t][c][d] * -// sigmoid(G[t][d * hc + c]) -> fp32 and fp16. G is never written. grid (hc * n / TN, ceil(T / TM)), 32 * SG threads. -template -kernel void gemm_hc_up(device half* A [[buffer(0)]], device half* B [[buffer(1)]], device const half* xnh [[buffer(2)]], - device float* mixed [[buffer(3)]], device half* mixh [[buffer(4)]], constant GemmArgs& g [[buffer(5)]], - constant EpiArgs& e [[buffer(6)]], uint3 tg [[threadgroup_position_in_grid]], - ushort tid [[thread_index_in_threadgroup]]) { - threadgroup float ot[TM * TN]; +// sigmoid(G[t][d * hc + c]) -> fp16 and, if mixed is bound, fp32. G is never written. +// grid (hc * n / TN, ceil(T / TM)), 32 * SG threads. +template +inline void hc_up_tile(device half* A, device half* B, device const half* xnh, device float* mixed, + device half* mixh, constant GemmArgs& g, constant EpiArgs& e, + int t0, int c0, ushort tid, threadgroup float* ot) { tensor, tensor_inline> tA(A, dextents(g.K, g.M)); tensor, tensor_inline> tB(B, dextents(g.K, g.N)); - tensor, tensor_inline> tO(ot, dextents(TN, TM)); constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v, false, true, false); matmul2d> op; - auto mA = tA.slice(0, int(tg.y) * TM); - auto mB = tB.slice(0, int(tg.x) * TN); - op.run(mA, mB, tO); - threadgroup_barrier(mem_flags::mem_threadgroup); - const int hc = int(e.hc), n = int(e.n_embd), DPT = TN / hc; - const int t0 = int(tg.y) * TM, d0 = int(tg.x) * DPT; - for (int idx = tid; idx < TM * DPT; idx += SG * 32) { - const int m = idx / DPT, dd = idx - m * DPT, t = t0 + m; - if (t >= int(e.n_tok)) continue; - float s = 0.0f; - for (int c = 0; c < hc; ++c) - s = fma(float(xnh[ulong(t) * hc * n + c * n + d0 + dd]), 1.0f / (1.0f + exp(-ot[m * TN + dd * hc + c])), s); - s = s / float(hc); - mixed[ulong(t) * n + d0 + dd] = s; - mixh[ulong(t) * n + d0 + dd] = half(s); + auto mA = tA.slice(0, t0); + auto mB = tB.slice(0, c0); + if constexpr (REG) { + auto acc = op.template get_destination_cooperative_tensor(); + op.run(mA, mB, acc); + const int n = int(e.n_embd), d0 = c0 / 4; + // On M5, each lane owns four consecutive output columns: one output's HC streams. + #pragma clang loop unroll(full) + for (uint16_t i = 0; i < acc.get_capacity(); i += 4) { + const auto ix = acc.get_multidimensional_index(i); + const int t = t0 + ix[1], d = d0 + ix[0] / 4; + if (t >= int(e.n_tok)) continue; + device const half* xnh_t = xnh + ulong(t) * 4 * n + d; + float s = 0.0f; + #pragma clang loop unroll(full) + for (int c = 0; c < 4; ++c) + s = fma(float(xnh_t[c * n]), 1.0f / (1.0f + exp(-acc[i + c])), s); + s *= 0.25f; + const ulong out_idx = ulong(t) * n + d; + if (mixed) mixed[out_idx] = s; + mixh[out_idx] = half(s); + } + } else { + tensor, tensor_inline> tO(ot, dextents(TN, TM)); + op.run(mA, mB, tO); + threadgroup_barrier(mem_flags::mem_threadgroup); + const int hc = int(e.hc), n = int(e.n_embd), DPT = TN / hc; + const int d0 = c0 / hc; + for (int idx = tid; idx < TM * DPT; idx += SG * 32) { + const int m = idx / DPT, dd = idx - m * DPT, t = t0 + m; + if (t >= int(e.n_tok)) continue; + float s = 0.0f; + for (int c = 0; c < hc; ++c) + s = fma(float(xnh[ulong(t) * hc * n + c * n + d0 + dd]), 1.0f / (1.0f + exp(-ot[m * TN + dd * hc + c])), s); + s = s / float(hc); + if (mixed) mixed[ulong(t) * n + d0 + dd] = s; + mixh[ulong(t) * n + d0 + dd] = half(s); + } } } +template +kernel void gemm_hc_up(device half* A [[buffer(0)]], device half* B [[buffer(1)]], device const half* xnh [[buffer(2)]], + device float* mixed [[buffer(3)]], device half* mixh [[buffer(4)]], constant GemmArgs& g [[buffer(5)]], + constant EpiArgs& e [[buffer(6)]], uint3 tg [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]]) { + threadgroup float ot[REG ? 1 : TM * TN]; + const int t0 = int(tg.y) * TM, c0 = int(tg.x) * TN; + if constexpr (REG && TM > 16) { + if (t0 + 16 >= int(e.n_tok)) { + hc_up_tile<16, TN, SG, REG>(A, B, xnh, mixed, mixh, g, e, t0, c0, tid, ot); + return; + } + } + if constexpr (REG && TM > 32) { + if (t0 + 32 >= int(e.n_tok)) { + hc_up_tile<32, TN, SG, REG>(A, B, xnh, mixed, mixh, g, e, t0, c0, tid, ot); + return; + } + } + hc_up_tile(A, B, xnh, mixed, mixh, g, e, t0, c0, tid, ot); +} template [[host_name("gemm_hc_up_64_64_4")]] kernel decltype(gemm_hc_up<64, 64, 4>) gemm_hc_up<64, 64, 4>; template [[host_name("gemm_hc_up_64_32_4")]] kernel decltype(gemm_hc_up<64, 32, 4>) gemm_hc_up<64, 32, 4>; template [[host_name("gemm_hc_up_32_64_4")]] kernel decltype(gemm_hc_up<32, 64, 4>) gemm_hc_up<32, 64, 4>; +template [[host_name("gemm_hc_up_128_64_4")]] kernel decltype(gemm_hc_up<128, 64, 4>) gemm_hc_up<128, 64, 4>; +template [[host_name("gemm_hc_up_64_128_4")]] kernel decltype(gemm_hc_up<64, 128, 4>) gemm_hc_up<64, 128, 4>; +template [[host_name("gemm_hc_up_reg_64_64_4")]] kernel decltype(gemm_hc_up<64, 64, 4, true>) gemm_hc_up<64, 64, 4, true>; +template [[host_name("gemm_hc_up_reg_128_64_4")]] kernel decltype(gemm_hc_up<128, 64, 4, true>) gemm_hc_up<128, 64, 4, true>; +template [[host_name("gemm_hc_up_reg_64_128_4")]] kernel decltype(gemm_hc_up<64, 128, 4, true>) gemm_hc_up<64, 128, 4, true>; +template [[host_name("gemm_hc_up_reg_64_32_4")]] kernel decltype(gemm_hc_up<64, 32, 4, true>) gemm_hc_up<64, 32, 4, true>; +template [[host_name("gemm_hc_up_reg_32_128_4")]] kernel decltype(gemm_hc_up<32, 128, 4, true>) gemm_hc_up<32, 128, 4, true>; + +// gemm_hc_down: the HC down GEMM V = xn . [w_down; w_inject]^T (T x (hc_lr + hc), K = hc * n_embd) with gr_epi_down +// as its register epilogue: lo_h = half(silu(v / hc)), inject = v. V is never written; lo_h, inject bitwise gemm_f16 + +// gr_epi_down. M5 layout only (coop_k4: a lane's elements 4j .. 4j + 3 are four consecutive columns of one row); +// hc_lr % 4 == 0, hc == 4. grid (ceil((hc_lr + hc) / TN), ceil(T / TM)), 32 * SG threads +template +inline void hc_down_tile(device half* A, device half* B, device half* lo_h, device float* inject, + constant EpiArgs& e, int t0, int c0) { + const int K = int(e.hc * e.n_embd), nr = int(e.hc_lr + e.hc); + tensor, tensor_inline> tA(A, dextents(K, int(e.n_tok))); + tensor, tensor_inline> tB(B, dextents(K, nr)); + constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v, false, true, false); + matmul2d> op; + auto mA = tA.slice(0, t0); + auto mB = tB.slice(0, c0); + auto acc = op.template get_destination_cooperative_tensor(); + op.run(mA, mB, acc); + #pragma clang loop unroll(full) + for (uint16_t i = 0; i < acc.get_capacity(); i += 4) { + const auto ix = acc.get_multidimensional_index(i); + const int r = c0 + ix[0], t = t0 + ix[1]; + if (t >= int(e.n_tok) || r >= nr) continue; + if (r < int(e.hc_lr)) { + half4 o; + #pragma clang loop unroll(full) + for (int c = 0; c < 4; ++c) { + const float z = acc[i + c] * 0.25f; // / hc, exact: hc == 4 + o[c] = half(z / (1.0f + exp(-z))); + } + *(device half4*)(lo_h + ulong(t) * e.hc_lr + r) = o; + } else { + *(device float4*)(inject + ulong(t) * e.hc) = float4(acc[i], acc[i + 1], acc[i + 2], acc[i + 3]); + } + } +} +template +inline void hc_down_rows(device half* A, device half* B, device half* lo_h, device float* inject, + constant EpiArgs& e, int t0, int c0) { + // Trimming 32-row parents regresses full tiles on M5. + if constexpr (TM > 32) { + if (t0 + 16 >= int(e.n_tok)) { + hc_down_tile<16, TN, SG>(A, B, lo_h, inject, e, t0, c0); + return; + } + if (t0 + 32 >= int(e.n_tok)) { + hc_down_tile<32, TN, SG>(A, B, lo_h, inject, e, t0, c0); + return; + } + } + hc_down_tile(A, B, lo_h, inject, e, t0, c0); +} +template +kernel void gemm_hc_down(device half* A [[buffer(0)]], device half* B [[buffer(1)]], device half* lo_h [[buffer(2)]], + device float* inject [[buffer(3)]], constant EpiArgs& e [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]]) { + const int c0 = int(tg.x) * TN, t0 = int(tg.y) * TM; + if constexpr (TN == 128) { + // A single 96-column tail beats serial 64+16 on M5, including SG=2. + if (e.hc == 4 && c0 + 64 == int(e.hc_lr)) { + hc_down_rows(A, B, lo_h, inject, e, t0, c0); + return; + } + } + hc_down_rows(A, B, lo_h, inject, e, t0, c0); +} +template [[host_name("gemm_hc_down_16_64_2")]] kernel decltype(gemm_hc_down<16, 64, 2>) gemm_hc_down<16, 64, 2>; +template [[host_name("gemm_hc_down_32_64_2")]] kernel decltype(gemm_hc_down<32, 64, 2>) gemm_hc_down<32, 64, 2>; +template [[host_name("gemm_hc_down_32_128_4")]] kernel decltype(gemm_hc_down<32, 128, 4>) gemm_hc_down<32, 128, 4>; +template [[host_name("gemm_hc_down_64_128_4")]] kernel decltype(gemm_hc_down<64, 128, 4>) gemm_hc_down<64, 128, 4>; +template [[host_name("gemm_hc_down_32_128_2")]] kernel decltype(gemm_hc_down<32, 128, 2>) gemm_hc_down<32, 128, 2>; + +// gemm_swiglu_up: the shared expert's up GEMM u = x . w_up^T (M x N, K) with swiglu_h as its register epilogue against +// the gate GEMM's g: h = half(silu(g) * u). u is never written; h bitwise gemm_f16 + swiglu_h. One accumulator: the +// gate stays a plain gemm_f16 before it. M5 layout only (as gemm_hc_down); N % 4 == 0. +// grid (ceil(N / TN), ceil(M / TM)), 32 * SG threads +template +inline void swiglu_up_tile(device half* A, device half* B, device const float* g, device half* h, + constant GemmArgs& a, int t0, int c0) { + tensor, tensor_inline> tA(A, dextents(a.K, a.M)); + tensor, tensor_inline> tB(B, dextents(a.K, a.N)); + constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v, false, true, false); + matmul2d> op; + auto mA = tA.slice(0, t0); + auto mB = tB.slice(0, c0); + auto acc = op.template get_destination_cooperative_tensor(); + op.run(mA, mB, acc); + #pragma clang loop unroll(full) + for (uint16_t i = 0; i < acc.get_capacity(); i += 4) { + const auto ix = acc.get_multidimensional_index(i); + const int c = c0 + ix[0], t = t0 + ix[1]; + if (t >= a.M || c >= a.N) continue; + const ulong idx = ulong(t) * uint(a.N) + uint(c); + const float4 gv = *(device const float4*)(g + idx); + half4 o; + #pragma clang loop unroll(full) + for (int k = 0; k < 4; ++k) o[k] = half(gv[k] / (1.0f + exp(-gv[k])) * acc[i + k]); + *(device half4*)(h + idx) = o; + } +} +template +kernel void gemm_swiglu_up(device half* A [[buffer(0)]], device half* B [[buffer(1)]], device const float* g [[buffer(2)]], + device half* h [[buffer(3)]], constant GemmArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]]) { + const int t0 = int(tg.y) * TM, c0 = int(tg.x) * TN; + if constexpr (TM > 32) { + if (t0 + 16 >= a.M) { + swiglu_up_tile<16, TN, SG>(A, B, g, h, a, t0, c0); + return; + } + if (t0 + 32 >= a.M) { + swiglu_up_tile<32, TN, SG>(A, B, g, h, a, t0, c0); + return; + } + } + swiglu_up_tile(A, B, g, h, a, t0, c0); +} +template [[host_name("gemm_swiglu_up_32_64_2")]] kernel decltype(gemm_swiglu_up<32, 64, 2>) gemm_swiglu_up<32, 64, 2>; +template [[host_name("gemm_swiglu_up_64_128_4")]] kernel decltype(gemm_swiglu_up<64, 128, 4>) gemm_swiglu_up<64, 128, 4>; diff --git a/metal/kernels/qsa.metal b/metal/kernels/qsa.metal index 9642e58e..af5f0e3b 100644 --- a/metal/kernels/qsa.metal +++ b/metal/kernels/qsa.metal @@ -35,6 +35,16 @@ constant int QSA_NTOK [[function_constant(3)]]; // the pooled indexer keys too (position pos_base + 4 b of a block). Off: exactly the plain RoPE. constant bool QSA_MROPE [[function_constant(24)]]; constant bool qsa_mrope = is_function_constant_defined(QSA_MROPE) && QSA_MROPE; +// QSA_Q_F16_FC (function constant 25): qsa_prep writes qcur directly as fp16 (device half*) for qsa_flash, skipping qsa_q_f16. +constant bool QSA_Q_F16_FC [[function_constant(25)]]; +constant bool qsa_qf16 = is_function_constant_defined(QSA_Q_F16_FC) ? QSA_Q_F16_FC : false; +// QSA_ROW0 (function constant 28): qsa_prep's threadgroup x is row QSA_ROW0 + x (24: the K/V rows only, for rows whose +// queries nobody reads). Undefined: 0. +constant int QSA_ROW0 [[function_constant(28)]]; +constant int qsa_row0 = is_function_constant_defined(QSA_ROW0) ? QSA_ROW0 : 0; +// QSA_LAST_Q (function constant 29): qfull / qcur hold only the last token's query at index 0 (combined MTP catch-up + step 0). +constant bool QSA_LAST_Q [[function_constant(29)]]; +constant bool qsa_last_q = is_function_constant_defined(QSA_LAST_Q) && QSA_LAST_Q; struct QsaArgs { uint pos0; // the window's first cell (relative to the sequence) @@ -54,34 +64,156 @@ static inline float tg_sum256(float v, threadgroup float* part, ushort sg, ushor threadgroup_barrier(mem_flags::mem_threadgroup); float t = 0.0f; for (int i = 0; i < nsg; ++i) t += part[i]; - threadgroup_barrier(mem_flags::mem_threadgroup); return t; } // one row (width w <= 256 = threads) : out = rope((rs * x) * g) (RoPE on the first 64 dims, NEOX pairs i, i+32) static inline float norm_rope(float x, float xhi_unused, device const float* g, threadgroup float* row, threadgroup float* part, ushort tid, ushort sg, ushort lane, int w, float eps, - device const float2* cs, uint absp, bool rope, device const int* mtab) { + device const float2* cs, uint absp, bool rope, device const int* mtab, + bool reuse = false) { const float ss = tg_sum256(tid < w ? x * x : 0.0f, part, sg, lane, 8); const float rs = rsqrt(ss / float(w) + eps); const float y = tid < w ? (rs * x) * g[tid] : 0.0f; - row[tid] = y; - threadgroup_barrier(mem_flags::mem_threadgroup); float o = y; - if (rope && tid < QROT) { - const int p = tid & 31; - const uint pp = qsa_mrope ? uint(mtab[absp * 3 + uint(p % 3)]) : absp; - const float2 c = cs[pp * 32 + p]; - const float a = row[p], b = row[p + 32]; - o = tid < 32 ? a * c.x - b * c.y : a * c.y + b * c.x; + if (rope) { + if (tid < QROT) row[tid] = y; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < QROT) { + const int p = tid & 31; + const uint pp = qsa_mrope ? uint(mtab[absp * 3 + uint(p % 3)]) : absp; + const float2 c = cs[pp * 32 + p]; + const float a = row[p], b = row[p + 32]; + o = tid < 32 ? a * c.x - b * c.y : a * c.y + b * c.x; + } } - threadgroup_barrier(mem_flags::mem_threadgroup); + if (reuse) threadgroup_barrier(mem_flags::mem_threadgroup); (void)xhi_unused; return o; } +// QSA_KV_STRIDE (function constant 26): kcur, vcur, iqraw and ikraw are column slices of one [n_tok][QSA_KV_STRIDE] +// projection (prompt_chunk's fused K | V | indexer K | indexer Q GEMM, 1152 or 1664 columns). Undefined: their own +// [n_tok][rows] buffers. +constant int QSA_KV_STRIDE [[function_constant(26)]]; +constant int qsa_kv_s = is_function_constant_defined(QSA_KV_STRIDE) ? QSA_KV_STRIDE : 0; + +// qsa_prep with one simdgroup a row (threadgroups of 32 threads, the same grid): lane l holds columns 32 w + l. The +// arithmetic is the 256-thread form's - the same simd_sum per 32 columns, summed in the same order; the NEOX pair +// (p, p + 32) and both halves of a 64-column int8 group are in one lane / simdgroup - with no threadgroup memory and +// no barriers (bitwise: qsa_parity --paired REF_DIR with QSA_PREP_TG=32). +static inline void norm_rope_sg(thread float* x, int nw, device const float* g, float eps, device const float2* cs, + uint absp, bool rope, device const int* mtab, ushort lane) { + float ss = 0.0f; + for (int i = 0; i < nw; ++i) ss += simd_sum(x[i] * x[i]); + const float rs = rsqrt(ss / float(32 * nw) + eps); + for (int i = 0; i < nw; ++i) x[i] = (rs * x[i]) * g[32 * i + lane]; + if (rope) { + const uint pp = qsa_mrope ? uint(mtab[absp * 3 + uint(lane % 3)]) : absp; + const float2 c = cs[pp * 32 + lane]; + const float a = x[0], b = x[1]; + x[0] = a * c.x - b * c.y; + x[1] = a * c.y + b * c.x; + } +} +static inline void qsa_prep_sg(device const float* qfull, device const float* kcur, device const float* vcur, + device const float* iqraw, device const float* ikraw, device const float* qn, + device const float* kn, device const float* iqn, device const float* ikn, + device const float2* cs, device float* qcur, device float* qidx, device char* kq, + device char* vq, device half* ksc, device half* vsc, device float* tail, device float* dead, + device float* pooled, constant QsaArgs& a, device const int* mtab, uint3 tg, ushort lane) { + constexpr int NI = ID / 32; + const int r = int(tg.x) + qsa_row0; + const int s_kv = qsa_kv_s ? qsa_kv_s : QKV * QD, s_ik = qsa_kv_s ? qsa_kv_s : ID; + const int s_iq = qsa_kv_s ? qsa_kv_s : IH * ID; + float x[8]; + if (r == 32) { + const int t_self = int(tg.y); + if (t_self >= 3) { + const uint cell = a.pos0 + uint(t_self); + if ((cell & 3) == 3) { + for (int i = 0; i < NI; ++i) { + const int c = 32 * i + lane; + x[i] = (((ikraw[(t_self - 3) * s_ik + c] + ikraw[(t_self - 2) * s_ik + c]) + ikraw[(t_self - 1) * s_ik + c]) + + ikraw[t_self * s_ik + c]) * 0.25f; + } + const uint b = cell >> 2; + norm_rope_sg(x, NI, ikn, a.eps, cs, a.pos_base + 4 * b, true, mtab, lane); + for (int i = 0; i < NI; ++i) pooled[b * ID + 32 * i + lane] = x[i]; + } + return; + } + if (t_self != 0) return; + const int n_pre = min(int(a.n_tok), 3); + for (int t = 0; t < n_pre; ++t) { + const uint cell = a.pos0 + uint(t); + for (int i = 0; i < NI; ++i) { + x[i] = ikraw[t * s_ik + 32 * i + lane]; + tail[(cell & 3) * ID + 32 * i + lane] = x[i]; + } + if (cell == 0) { + float y[8]; + for (int i = 0; i < NI; ++i) y[i] = x[i]; + norm_rope_sg(y, NI, ikn, a.eps, cs, 0, false, mtab, lane); + for (int i = 0; i < NI; ++i) dead[32 * i + lane] = y[i]; + } + if ((cell & 3) == 3) { + float m[8]; + for (int i = 0; i < NI; ++i) { + const int c = 32 * i + lane; + m[i] = (((tail[c] + tail[ID + c]) + tail[2 * ID + c]) + tail[3 * ID + c]) * 0.25f; + } + const uint b = cell >> 2; + norm_rope_sg(m, NI, ikn, a.eps, cs, a.pos_base + 4 * b, true, mtab, lane); + for (int i = 0; i < NI; ++i) pooled[b * ID + 32 * i + lane] = m[i]; + } + } + for (int t = max(n_pre, int(a.n_tok) - 4); t < int(a.n_tok); ++t) { + const uint cell = a.pos0 + uint(t); + for (int i = 0; i < NI; ++i) tail[(cell & 3) * ID + 32 * i + lane] = ikraw[t * s_ik + 32 * i + lane]; + } + return; + } + const int t = int(tg.y); + const uint cell = a.pos0 + uint(t), absp = a.pos_base + cell; + if (r < QH) { + if (qsa_last_q && t + 1 < int(a.n_tok)) return; + const int qt = qsa_last_q ? 0 : t; + for (int i = 0; i < 8; ++i) x[i] = qfull[(qt * QH + r) * 2 * QD + 32 * i + lane]; + norm_rope_sg(x, 8, qn, a.eps, cs, absp, true, mtab, lane); + for (int i = 0; i < 8; ++i) { + if (qsa_qf16) ((device half*)qcur)[(qt * QH + r) * QD + 32 * i + lane] = half(x[i]); + else qcur[(qt * QH + r) * QD + 32 * i + lane] = x[i]; + } + } else if (r < QH + 2 * QKV) { + const bool isv = r >= QH + QKV; + const int h = r - QH - (isv ? QKV : 0); + for (int i = 0; i < 8; ++i) x[i] = (isv ? vcur : kcur)[t * s_kv + h * QD + 32 * i + lane]; + if (!isv) norm_rope_sg(x, 8, kn, a.eps, cs, absp, true, mtab, lane); + const uint rowi = cell * QKV + uint(h); + for (int g = 0; g < 4; ++g) { + const float m = max(simd_max(fabs(x[2 * g])), simd_max(fabs(x[2 * g + 1]))); + const half sh = half(m / 127.0f); + const float sf = float(sh); + for (int i = 2 * g; i < 2 * g + 2; ++i) { + int q = 0; + if (sf > 0.0f) q = clamp(int(rint(x[i] / sf)), -127, 127); + (isv ? vq : kq)[rowi * QD + 32 * i + lane] = char(q); + } + if (lane == 0) (isv ? vsc : ksc)[rowi * 4 + g] = sh; + } + } else { + if (a.pos0 + a.n_tok <= uint(QBUDGET)) return; + const int h = r - QH - 2 * QKV; + for (int i = 0; i < NI; ++i) x[i] = iqraw[t * s_iq + h * ID + 32 * i + lane]; + norm_rope_sg(x, NI, iqn, a.eps, cs, absp, true, mtab, lane); + for (int i = 0; i < NI; ++i) qidx[(t * IH + h) * ID + 32 * i + lane] = x[i]; + } +} + // grid (32 rows + 1, n_tok): rows 0..23 q heads, 24..25 k heads, 26..27 v heads, 28..31 indexer q heads; -// threadgroup 32 (x = 32, y = 0) is the indexer-key threadgroup. 256 threads. +// threadgroup 32 (x = 32, y = 0) is the indexer-key threadgroup. 256 threads, or 32 (qsa_prep_sg). A grid of 28 (or of 4 from QSA_ROW0 = +// 24) writes no indexer state: the MTP layer's attention is a dense window, never scored. kernel void qsa_prep(device const float* qfull [[buffer(0)]], device const float* kcur [[buffer(1)]], device const float* vcur [[buffer(2)]], device const float* iqraw [[buffer(3)]], device const float* ikraw [[buffer(4)]], device const float* qn [[buffer(5)]], @@ -93,44 +225,78 @@ kernel void qsa_prep(device const float* qfull [[buffer(0)]], device const float device float* pooled [[buffer(18)]], constant QsaArgs& a [[buffer(19)]], device const int* mtab [[buffer(20)]], uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], - ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]], + uint3 nth [[threads_per_threadgroup]]) { + if (nth.x == 32) { + qsa_prep_sg(qfull, kcur, vcur, iqraw, ikraw, qn, kn, iqn, ikn, cs, qcur, qidx, kq, vq, ksc, vsc, tail, dead, pooled, a, + mtab, tg, lane); + return; + } threadgroup float part[8]; - threadgroup float row[256]; + threadgroup float row[QROT]; threadgroup float amax[8]; - const int r = int(tg.x); + const int r = int(tg.x) + qsa_row0; + // the row strides of kcur / vcur, ikraw, iqraw + const int s_kv = qsa_kv_s ? qsa_kv_s : QKV * QD, s_ik = qsa_kv_s ? qsa_kv_s : ID; + const int s_iq = qsa_kv_s ? qsa_kv_s : IH * ID; if (r == 32) { - // the indexer keys, every window token in order (only threadgroup (32, 0) works) - if (tg.y != 0) return; - for (int t = 0; t < int(a.n_tok); ++t) { + // the indexer keys: threadgroup (32, 0) handles t < 3 (boundary with prior tail + cell 0 dead) and final tail; + // every threadgroup (32, t >= 3) with (cell & 3) == 3 pools its 4 in-window keys in parallel. + const int t_self = int(tg.y); + if (t_self >= 3) { + const uint cell = a.pos0 + uint(t_self); + if ((cell & 3) == 3) { + float m = 0.0f; + if (tid < ID) { + m = (((ikraw[(t_self - 3) * s_ik + tid] + ikraw[(t_self - 2) * s_ik + tid]) + + ikraw[(t_self - 1) * s_ik + tid]) + ikraw[t_self * s_ik + tid]) * 0.25f; + } + const uint b = cell >> 2; + const float o = norm_rope(m, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, a.pos_base + 4 * b, true, mtab); + if (tid < ID) pooled[b * ID + tid] = o; + } + return; + } + if (t_self != 0) return; + const int n_pre = min(int(a.n_tok), 3); + for (int t = 0; t < n_pre; ++t) { const uint cell = a.pos0 + uint(t); - const float x = tid < ID ? ikraw[t * ID + tid] : 0.0f; + const float x = tid < ID ? ikraw[t * s_ik + tid] : 0.0f; if (tid < ID) tail[(cell & 3) * ID + tid] = x; threadgroup_barrier(mem_flags::mem_device); if (cell == 0) { - const float o = norm_rope(x, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, 0, false, mtab); + const float o = norm_rope(x, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, 0, false, mtab, true); if (tid < ID) dead[tid] = o; } if ((cell & 3) == 3) { float m = 0.0f; if (tid < ID) m = (((tail[tid] + tail[ID + tid]) + tail[2 * ID + tid]) + tail[3 * ID + tid]) * 0.25f; const uint b = cell >> 2; - const float o = norm_rope(m, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, a.pos_base + 4 * b, true, mtab); + const float o = norm_rope(m, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, a.pos_base + 4 * b, true, mtab, true); if (tid < ID) pooled[b * ID + tid] = o; } threadgroup_barrier(mem_flags::mem_device); } + const int t_tail0 = max(n_pre, int(a.n_tok) - 4); + for (int t = t_tail0; t < int(a.n_tok); ++t) { + const uint cell = a.pos0 + uint(t); + if (tid < ID) tail[(cell & 3) * ID + tid] = ikraw[t * s_ik + tid]; + } return; } const int t = int(tg.y); const uint cell = a.pos0 + uint(t), absp = a.pos_base + cell; if (r < QH) { - const float x = qfull[(t * QH + r) * 2 * QD + tid]; + if (qsa_last_q && t + 1 < int(a.n_tok)) return; + const int qt = qsa_last_q ? 0 : t; + const float x = qfull[(qt * QH + r) * 2 * QD + tid]; const float o = norm_rope(x, 0.0f, qn, row, part, tid, sg, lane, QD, a.eps, cs, absp, true, mtab); - qcur[(t * QH + r) * QD + tid] = o; + if (qsa_qf16) ((device half*)qcur)[(qt * QH + r) * QD + tid] = half(o); + else qcur[(qt * QH + r) * QD + tid] = o; } else if (r < QH + 2 * QKV) { const bool isv = r >= QH + QKV; const int h = r - QH - (isv ? QKV : 0); - float x = (isv ? vcur : kcur)[(t * QKV + h) * QD + tid]; + float x = (isv ? vcur : kcur)[t * s_kv + h * QD + tid]; if (!isv) x = norm_rope(x, 0.0f, kn, row, part, tid, sg, lane, QD, a.eps, cs, absp, true, mtab); // int8 with an fp16 scale per 64 values: groups = simdgroup pairs (64 threads) float m = simd_max(fabs(x)); @@ -146,8 +312,9 @@ kernel void qsa_prep(device const float* qfull [[buffer(0)]], device const float (isv ? vq : kq)[rowi * QD + tid] = char(q); if ((tid & 63) == 0) (isv ? vsc : ksc)[rowi * 4 + g] = sh; } else { + if (a.pos0 + a.n_tok <= uint(QBUDGET)) return; const int h = r - QH - 2 * QKV; - const float x = tid < ID ? iqraw[(t * IH + h) * ID + tid] : 0.0f; + const float x = tid < ID ? iqraw[t * s_iq + h * ID + tid] : 0.0f; const float o = norm_rope(x, 0.0f, iqn, row, part, tid, sg, lane, ID, a.eps, cs, absp, true, mtab); if (tid < ID) qidx[(t * IH + h) * ID + tid] = o; } @@ -164,12 +331,21 @@ kernel void qsa_score(device const float* qidx [[buffer(0)]], device const float const uint b = tg.x * 8 + sg; if (b > n_bid) return; device const float* key = b == n_bid ? dead : pooled + b * ID; - float k[4]; + device const float* qp = qidx + (t * IH) * ID + lane; + float k[4], q[IH][4]; + #pragma clang loop unroll(full) for (int i = 0; i < 4; ++i) k[i] = key[lane + 32 * i]; + #pragma clang loop unroll(full) + for (int h = 0; h < IH; ++h) { + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) q[h][i] = qp[h * ID + 32 * i]; + } float s = 0.0f; + #pragma clang loop unroll(full) for (int h = 0; h < IH; ++h) { float d = 0.0f; - for (int i = 0; i < 4; ++i) d = fma(qidx[(t * IH + h) * ID + lane + 32 * i], k[i], d); + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) d = fma(q[h][i], k[i], d); d = simd_sum(d); s += max(d, 0.0f); } @@ -190,7 +366,8 @@ kernel void qsa_topk(device const float* scores [[buffer(0)]], device int* ids [ ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { threadgroup atomic_uint hist[256]; threadgroup uint sel[4]; - threadgroup uint pref[32]; + threadgroup uint2 pref[32]; + threadgroup uint2 totals; const int t = int(tg.x); const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2, nb = n_bid + 1; device int* out = ids + t * QBUDGET; @@ -199,6 +376,10 @@ kernel void qsa_topk(device const float* scores [[buffer(0)]], device int* ids [ for (uint c = tid; c < n_kv; c += 1024) out[c] = int(c); return; } + const bool cached = nb <= 4096; + uint keys[4]; + if (cached) + for (uint b = tid, i = 0; b < nb; b += 1024, ++i) keys[i] = okey(sc[b]); // radix select of the `need`-th largest cell key uint prefix = 0, pmask = 0, need = QBUDGET; for (int pass = 0; pass < 4; ++pass) { @@ -206,28 +387,35 @@ kernel void qsa_topk(device const float* scores [[buffer(0)]], device int* ids [ if (tid < 256) atomic_store_explicit(&hist[tid], 0u, memory_order_relaxed); threadgroup_barrier(mem_flags::mem_threadgroup); for (uint b = tid; b < nb; b += 1024) { - const uint k = okey(sc[b]); + const uint k = cached ? keys[b / 1024] : okey(sc[b]); if ((k & pmask) == prefix) { const uint w = b < n_bid ? 4u : (n_kv & 3u); if (w) atomic_fetch_add_explicit(&hist[(k >> shift) & 255u], w, memory_order_relaxed); } } threadgroup_barrier(mem_flags::mem_threadgroup); - if (tid == 0) { - uint acc = 0, d = 0; - for (int v = 255; v >= 0; --v) { - const uint h = atomic_load_explicit(&hist[v], memory_order_relaxed); - if (acc + h >= need) { d = uint(v); break; } - acc += h; + if (sg == 0) { + uint counts[8], sum = 0; + for (uint j = 0; j < 8; ++j) { + counts[j] = atomic_load_explicit(&hist[255 - (lane * 8 + j)], memory_order_relaxed); + sum += counts[j]; + } + uint acc = simd_prefix_exclusive_sum(sum); + if (acc < need && acc + sum >= need) { + for (uint j = 0; j < 8; ++j) { + if (acc + counts[j] >= need) { + sel[0] = 255 - (lane * 8 + j); + sel[1] = need - acc; + break; + } + acc += counts[j]; + } } - sel[0] = d; - sel[1] = need - acc; } threadgroup_barrier(mem_flags::mem_threadgroup); prefix |= sel[0] << shift; pmask |= 255u << shift; need = sel[1]; - threadgroup_barrier(mem_flags::mem_threadgroup); } // T = prefix; take every cell with key > T and the first `need` cells (lowest index) with key == T, ascending. // Pass over blocks in order of index with a prefix scan of each block's selected cell count. @@ -237,36 +425,35 @@ kernel void qsa_topk(device const float* scores [[buffer(0)]], device int* ids [ const uint b = b0 + tid; uint w = 0, weq = 0; if (b < nb) { - const uint k = okey(sc[b]); + const uint k = cached ? keys[b / 1024] : okey(sc[b]); const uint cells = b < n_bid ? 4u : (n_kv & 3u); if (k > T) w = cells; else if (k == T) weq = cells; } - // equal-key cells: the lowest first, `need` of them in total const uint eq_before = simd_prefix_exclusive_sum(weq); - if (lane == 31) pref[sg] = eq_before + weq; + const uint gt_before = simd_prefix_exclusive_sum(w); + if (lane == 31) pref[sg] = uint2(eq_before + weq, gt_before + w); threadgroup_barrier(mem_flags::mem_threadgroup); - uint eq_sg = 0; - for (int i = 0; i < int(sg); ++i) eq_sg += pref[i]; - uint eq_tot = 0; - for (int i = 0; i < 32; ++i) eq_tot += pref[i]; + if (sg == 0) { + const uint2 count = pref[lane]; + const uint2 before = uint2(simd_prefix_exclusive_sum(count.x), + simd_prefix_exclusive_sum(count.y)); + pref[lane] = before; + if (lane == 31) totals = before + count; + } threadgroup_barrier(mem_flags::mem_threadgroup); - const uint eq0 = eq_taken + eq_sg + eq_before; - uint take_eq = 0; - if (weq && eq0 < need) take_eq = min(weq, need - eq0); + const uint2 group_before = pref[sg], count = totals; + simdgroup_barrier(mem_flags::mem_threadgroup); + const uint eq0 = eq_taken + group_before.x + eq_before; + const uint take_eq = eq0 < need ? min(weq, need - eq0) : 0u; const uint mine = w + take_eq; - const uint before = simd_prefix_exclusive_sum(mine); - if (lane == 31) pref[sg] = before + mine; - threadgroup_barrier(mem_flags::mem_threadgroup); - uint off = base; - for (int i = 0; i < int(sg); ++i) off += pref[i]; - uint tot = 0; - for (int i = 0; i < 32; ++i) tot += pref[i]; - threadgroup_barrier(mem_flags::mem_threadgroup); - off += before; + // Clipping the equal-key prefix preserves the lowest-cell tie break. + const uint remaining = need - min(need, eq_taken); + const uint off = base + group_before.y + gt_before + + min(remaining, group_before.x + eq_before); for (uint c = 0; c < mine; ++c) out[off + c] = int(4 * b + c); - base += tot; - eq_taken += eq_tot; + base += count.y + min(remaining, count.x); + eq_taken += count.x; } } @@ -293,19 +480,28 @@ kernel void qsa_attn(device const float* qcur [[buffer(0)]], device const char* return; } const int nc = min(QCH, int(width) - c0); + #pragma clang loop unroll(full) for (int h = 0; h < HPG; ++h) qs[h][tid] = qcur[(t * QH + g * HPG + h) * QD + tid]; if (tid < QCH) cid[tid] = tid < nc ? ids[t * QBUDGET + c0 + tid] : 0; threadgroup_barrier(mem_flags::mem_threadgroup); // scores: 12 x 64 dots of 256; simdgroup sg takes cells sg, sg + 8, ...; lanes split the 256 dims (8 each) + threadgroup const float* q_lane = &qs[0][lane * 8]; for (int c = sg; c < nc; c += 8) { const uint rowi = uint(cid[c]) * QKV + uint(g); - device const char* kr = kq + rowi * QD; - float kv[8]; + device const char* kr = kq + rowi * QD + lane * 8; + const char4 k0 = *(device const char4*)(kr); + const char4 k1 = *(device const char4*)(kr + 4); const float scl = float(ksc[rowi * 4 + (lane >> 3)]); - for (int i = 0; i < 8; ++i) kv[i] = float(kr[lane * 8 + i]) * scl; + const float kv[8] = { + float(k0.x) * scl, float(k0.y) * scl, float(k0.z) * scl, float(k0.w) * scl, + float(k1.x) * scl, float(k1.y) * scl, float(k1.z) * scl, float(k1.w) * scl + }; + #pragma clang loop unroll(full) for (int h = 0; h < HPG; ++h) { + threadgroup const float* qh = q_lane + h * QD; float d = 0.0f; - for (int i = 0; i < 8; ++i) d = fma(qs[h][lane * 8 + i], kv[i], d); + #pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) d = fma(qh[i], kv[i], d); d = simd_sum(d); if (lane == 0) s[h][c] = d * 0.0625f; } @@ -313,28 +509,44 @@ kernel void qsa_attn(device const float* qcur [[buffer(0)]], device const char* threadgroup_barrier(mem_flags::mem_threadgroup); // per head: max, exp, sum (simdgroup h % 8 handles heads h, h + 8) threadgroup float mh[HPG], lh[HPG]; + // nc <= QCH = 64: a lane holds cells lane and lane + 32 in registers (the same max and sum as a strided loop) + const bool in0 = int(lane) < nc, in1 = int(lane) + 32 < nc; for (int h = sg; h < HPG; h += 8) { - float m = -INFINITY; - for (int c = lane; c < nc; c += 32) m = max(m, s[h][c]); - m = simd_max(m); - float l = 0.0f; - for (int c = lane; c < nc; c += 32) { - const float p = exp(s[h][c] - m); - s[h][c] = p; - l += p; - } - l = simd_sum(l); + const float s0 = in0 ? s[h][lane] : -INFINITY, s1 = in1 ? s[h][lane + 32] : -INFINITY; + const float m = simd_max(max(s0, s1)); + const float p0 = in0 ? exp(s0 - m) : 0.0f, p1 = in1 ? exp(s1 - m) : 0.0f; + if (in0) s[h][lane] = p0; + if (in1) s[h][lane + 32] = p1; + const float l = simd_sum(p0 + p1); if (lane == 0) { mh[h] = m; lh[h] = l; } } threadgroup_barrier(mem_flags::mem_threadgroup); // acc[h][j] = sum_c p[h][c] v[c][j]: thread = dim j float acc[HPG]; + #pragma clang loop unroll(full) for (int h = 0; h < HPG; ++h) acc[h] = 0.0f; - for (int c = 0; c < nc; ++c) { + // four cells' V loads in flight at once; the fma order per head stays c, c + 1, ... + auto vcell = [&](int c) { const uint rowi = uint(cid[c]) * QKV + uint(g); - const float v = float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + (tid >> 6)]); + return float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + (tid >> 6)]); + }; + int c = 0; + for (; c + 4 <= nc; c += 4) { + const float v0 = vcell(c), v1 = vcell(c + 1), v2 = vcell(c + 2), v3 = vcell(c + 3); + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) { + float x = fma(s[h][c], v0, acc[h]); + x = fma(s[h][c + 1], v1, x); + x = fma(s[h][c + 2], v2, x); + acc[h] = fma(s[h][c + 3], v3, x); + } + } + for (; c < nc; ++c) { + const float v = vcell(c); + #pragma clang loop unroll(full) for (int h = 0; h < HPG; ++h) acc[h] = fma(s[h][c], v, acc[h]); } + #pragma clang loop unroll(full) for (int h = 0; h < HPG; ++h) { device float* o = pp + (g * HPG + h) * (QD + 2); o[tid] = acc[h]; @@ -342,43 +554,63 @@ kernel void qsa_attn(device const float* qcur [[buffer(0)]], device const char* } } +constant float qsa_q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f}; +constant float qsa_p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + +static inline void qsa_write_prepared(float y, device float* xp, constant QsaArgs& a, + int t, int h, ushort tid) { + // prepared planes (xprep layout) of the n = QH * QD vector + const uint n = QH * QD, plane = a.n_tok * n, d = uint(h) * QD + tid; + const int j = tid & 15; + device float* raw = xp + uint(t) * n + d; + raw[0] = y; + raw[plane] = y * qsa_q2s[j & 7]; + raw[2 * plane] = y * qsa_p4s[j & 3]; + raw[3 * plane] = (j & 1) ? y * 0.00390625f : y; + float s0 = 0.0f, s1 = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < 16; k += 2) { + s0 += simd_shuffle_down(y, ushort(k)); + s1 += simd_shuffle_down(y, ushort(k + 1)); + } + if (j == 0) xp[4 * plane + uint(t) * (n / 16) + d / 16] = s0 + s1; +} + // grid (QH, n_tok), 256 threads: out = merge * sigmoid(gate), as prepared activations of n = QH * QD per token kernel void qsa_merge(device const float* part [[buffer(0)]], device const float* qfull [[buffer(1)]], device float* xp [[buffer(2)]], constant QsaArgs& a [[buffer(3)]], uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { - threadgroup float msum[256]; const int h = int(tg.x), t = int(tg.y); const int nchunks = (QBUDGET + QCH - 1) / QCH; const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); const int nc = int((width + QCH - 1) / QCH); + constexpr uint STRIDE = QH * (QD + 2); device const float* pp = part + (uint(t) * nchunks * QH + uint(h)) * (QD + 2); + // the chunks' (m, l) and weights exp(m - M) once per threadgroup, not once per thread (nc <= 33 < 64) + threadgroup float mk[64], lk[64], ek[64], L_sh; + if (tid < uint(nc)) { + const float2 ml = *(device const float2*)(pp + tid * STRIDE + QD); + mk[tid] = ml.x; + lk[tid] = ml.y; + } + threadgroup_barrier(mem_flags::mem_threadgroup); float M = -INFINITY; - for (int k = 0; k < nc; ++k) M = max(M, pp[k * QH * (QD + 2) + QD]); - float L = 0.0f, o = 0.0f; - for (int k = 0; k < nc; ++k) { - device const float* q = pp + k * QH * (QD + 2); - const float e = exp(q[QD] - M); - L = fma(q[QD + 1], e, L); - o = fma(q[tid], e, o); + for (int k = tid & 31; k < nc; k += 32) M = max(M, mk[k]); + M = simd_max(M); + if (tid < uint(nc)) ek[tid] = exp(mk[tid] - M); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0) { + float L = 0.0f; + for (int k = 0; k < nc; ++k) L = fma(lk[k], ek[k], L); + L_sh = L; } - o = o / L; + float o = 0.0f; + device const float* q = pp + tid; + for (int k = 0; k < nc; ++k, q += STRIDE) o = fma(*q, ek[k], o); + threadgroup_barrier(mem_flags::mem_threadgroup); + o = o / L_sh; const float gate = qfull[(t * QH + h) * 2 * QD + QD + tid]; const float y = o * (1.0f / (1.0f + exp(-gate))); - // prepared planes (xprep layout) of the n = QH * QD vector - const uint n = QH * QD, plane = a.n_tok * n, d = uint(h) * QD + tid; - const int j = tid & 15; - const float q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f}; - const float p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; - device float* raw = xp + uint(t) * n + d; - raw[0] = y; - raw[plane] = y * q2s[j & 7]; - raw[2 * plane] = y * p4s[j & 3]; - raw[3 * plane] = (j & 1) ? y * 0.00390625f : y; - msum[tid] = y; - threadgroup_barrier(mem_flags::mem_threadgroup); - if (j == 0) { - float s0 = 0.0f, s1 = 0.0f; - for (int k = 0; k < 16; k += 2) { s0 += msum[tid + k]; s1 += msum[tid + k + 1]; } - xp[4 * plane + uint(t) * (n / 16) + d / 16] = s0 + s1; - } + qsa_write_prepared(y, xp, a, t, h, tid); } + diff --git a/metal/kernels/qsa2.metal b/metal/kernels/qsa2.metal index 8ae189b6..30439a16 100644 --- a/metal/kernels/qsa2.metal +++ b/metal/kernels/qsa2.metal @@ -29,9 +29,11 @@ kernel void qsa_attn2(device const float* qcur [[buffer(0)]], device const char* if (uint(c0) >= width) return; device float* pp = part + ((uint(t) * nchunks + chunk) * QH) * (QD + 2); const int nc = min(QCH2, int(width) - c0); + #pragma clang loop unroll(full) for (int h = 0; h < HPG; ++h) qs[h][tid] = qcur[(t * QH + h0 + h) * QD + tid]; if (tid < QCH2) cid[tid] = tid < nc ? ids[t * QBUDGET + c0 + tid] : 0; threadgroup_barrier(mem_flags::mem_threadgroup); + threadgroup const float* q_lane = &qs[0][lane * 8]; for (int c = sg; c < nc; c += 8) { const uint rowi = uint(cid[c]) * QKV + uint(g); device const char4* kr = (device const char4*)(kq + rowi * QD + lane * 8); @@ -39,9 +41,11 @@ kernel void qsa_attn2(device const float* qcur [[buffer(0)]], device const char* const float scl = float(ksc[rowi * 4 + (lane >> 3)]); const float kv[8] = {float(k0.x) * scl, float(k0.y) * scl, float(k0.z) * scl, float(k0.w) * scl, float(k1.x) * scl, float(k1.y) * scl, float(k1.z) * scl, float(k1.w) * scl}; + #pragma clang loop unroll(full) for (int h = 0; h < HPG; ++h) { float d = 0.0f; - for (int i = 0; i < 8; ++i) d = fma(qs[h][lane * 8 + i], kv[i], d); + #pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) d = fma(q_lane[h * QD + i], kv[i], d); d = simd_sum(d); if (lane == 0) s[h][c] = d * 0.0625f; } @@ -58,20 +62,29 @@ kernel void qsa_attn2(device const float* qcur [[buffer(0)]], device const char* } threadgroup_barrier(mem_flags::mem_threadgroup); float acc[HPG]; + #pragma clang loop unroll(full) for (int h = 0; h < HPG; ++h) acc[h] = 0.0f; - for (int cb = 0; cb < nc; cb += 8) { + const uint v_scale_off = tid >> 6; + int cb = 0; + for (; cb + 8 <= nc; cb += 8) { float v[8]; - const int m = min(8, nc - cb); + #pragma clang loop unroll(full) for (int c = 0; c < 8; ++c) { - v[c] = 0.0f; - if (c < m) { - const uint rowi = uint(cid[cb + c]) * QKV + uint(g); - v[c] = float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + (tid >> 6)]); - } + const uint rowi = uint(cid[cb + c]) * QKV + uint(g); + v[c] = float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + v_scale_off]); } - for (int c = 0; c < m; ++c) + #pragma clang loop unroll(full) + for (int c = 0; c < 8; ++c) + #pragma clang loop unroll(full) for (int h = 0; h < HPG; ++h) acc[h] = fma(s[h][cb + c], v[c], acc[h]); } + for (; cb < nc; ++cb) { + const uint rowi = uint(cid[cb]) * QKV + uint(g); + const float v = float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + v_scale_off]); + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) acc[h] = fma(s[h][cb], v, acc[h]); + } + #pragma clang loop unroll(full) for (int h = 0; h < HPG; ++h) { device float* o = pp + (h0 + h) * (QD + 2); o[tid] = acc[h]; @@ -82,40 +95,45 @@ kernel void qsa_attn2(device const float* qcur [[buffer(0)]], device const char* kernel void qsa_merge2(device const float* part [[buffer(0)]], device const float* qfull [[buffer(1)]], device float* xp [[buffer(2)]], constant QsaArgs& a [[buffer(3)]], uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { - threadgroup float msum[256]; + threadgroup float mk[128], lk[128], ek[128], L_sh; const int h = int(tg.x), t = int(tg.y); const int nchunks = (QBUDGET + QCH2 - 1) / QCH2; const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); const int nc = int((width + QCH2 - 1) / QCH2); + constexpr uint STRIDE = QH * (QD + 2); device const float* pp = part + (uint(t) * nchunks * QH + uint(h)) * (QD + 2); + if (tid < uint(nc)) { + const float2 ml = *(device const float2*)(pp + tid * STRIDE + QD); + mk[tid] = ml.x; + lk[tid] = ml.y; + } + threadgroup_barrier(mem_flags::mem_threadgroup); float M = -INFINITY; - for (int k = 0; k < nc; ++k) M = max(M, pp[k * QH * (QD + 2) + QD]); - float L = 0.0f, o = 0.0f; - for (int k = 0; k < nc; ++k) { - device const float* q = pp + k * QH * (QD + 2); - const float e = exp(q[QD] - M); - L = fma(q[QD + 1], e, L); - o = fma(q[tid], e, o); + for (int k = tid & 31; k < nc; k += 32) M = max(M, mk[k]); + M = simd_max(M); + if (tid < uint(nc)) ek[tid] = exp(mk[tid] - M); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0) { + float L = 0.0f; + for (int k = 0; k < nc; ++k) L = fma(lk[k], ek[k], L); + L_sh = L; + } + float o = 0.0f; + device const float* q = pp + tid; + int k = 0; + for (; k + 3 < nc; k += 4, q += 4 * STRIDE) { + const float p0 = q[0], p1 = q[STRIDE], p2 = q[2 * STRIDE], p3 = q[3 * STRIDE]; + o = fma(p0, ek[k], o); + o = fma(p1, ek[k + 1], o); + o = fma(p2, ek[k + 2], o); + o = fma(p3, ek[k + 3], o); } - o = o / L; + for (; k < nc; ++k, q += STRIDE) o = fma(*q, ek[k], o); + threadgroup_barrier(mem_flags::mem_threadgroup); + o = o / L_sh; const float gate = qfull[(t * QH + h) * 2 * QD + QD + tid]; const float y = o * (1.0f / (1.0f + exp(-gate))); - const uint n = QH * QD, plane = a.n_tok * n, d = uint(h) * QD + tid; - const int j = tid & 15; - const float q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f}; - const float p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; - device float* raw = xp + uint(t) * n + d; - raw[0] = y; - raw[plane] = y * q2s[j & 7]; - raw[2 * plane] = y * p4s[j & 3]; - raw[3 * plane] = (j & 1) ? y * 0.00390625f : y; - msum[tid] = y; - threadgroup_barrier(mem_flags::mem_threadgroup); - if (j == 0) { - float s0 = 0.0f, s1 = 0.0f; - for (int k = 0; k < 16; k += 2) { s0 += msum[tid + k]; s1 += msum[tid + k + 1]; } - xp[4 * plane + uint(t) * (n / 16) + d / 16] = s0 + s1; - } + qsa_write_prepared(y, xp, a, t, h, tid); } kernel void qsa_topk2(device const float* scores [[buffer(0)]], device int* ids [[buffer(1)]], @@ -133,37 +151,76 @@ kernel void qsa_topk2(device const float* scores [[buffer(0)]], device int* ids for (uint c = tid; c < n_kv; c += 1024) out[c] = int(c); return; } + if (tid < 256) atomic_store_explicit(&hist[tid], 0u, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); uint prefix = 0, pmask = 0, need = QBUDGET; - for (int pass = 0; pass < 4; ++pass) { - const int shift = 24 - 8 * pass; - if (tid < 256) atomic_store_explicit(&hist[tid], 0u, memory_order_relaxed); - threadgroup_barrier(mem_flags::mem_threadgroup); - for (uint b = tid; b < nb; b += 1024) { - const uint k = okey(sc[b]); - if ((k & pmask) == prefix) { - const uint w = b < n_bid ? 4u : (n_kv & 3u); - if (w) atomic_fetch_add_explicit(&hist[(k >> shift) & 255u], w, memory_order_relaxed); - } + if (nb <= 4096) { + uint keys[4]; + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) { + const uint b = tid + 1024u * uint(i); + keys[i] = b < nb ? okey(sc[b]) : 0u; } - threadgroup_barrier(mem_flags::mem_threadgroup); - // thread i < 256 holds bin 255 - i: an exclusive prefix over i = the cells of the higher digits - uint cnt = 0, ex = 0; - if (tid < 256) { - cnt = atomic_load_explicit(&hist[255 - tid], memory_order_relaxed); - ex = simd_prefix_exclusive_sum(cnt); - if (lane == 31) sgs[sg] = ex + cnt; + for (int pass = 0; pass < 4; ++pass) { + const int shift = 24 - 8 * pass; + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) { + const uint b = tid + 1024u * uint(i); + if (b >= nb) break; + const uint k = keys[i]; + if ((k & pmask) == prefix) { + const uint w = b < n_bid ? 4u : (n_kv & 3u); + if (w) atomic_fetch_add_explicit(&hist[(k >> shift) & 255u], w, memory_order_relaxed); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + uint cnt = 0, ex = 0; + if (tid < 256) { + cnt = atomic_load_explicit(&hist[255 - tid], memory_order_relaxed); + if (pass < 3) atomic_store_explicit(&hist[255 - tid], 0u, memory_order_relaxed); + ex = simd_prefix_exclusive_sum(cnt); + if (lane == 31) sgs[sg] = ex + cnt; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < 256) { + uint above = ex; + for (int i = 0; i < int(sg); ++i) above += sgs[i]; + if (above < need && above + cnt >= need) { sel[0] = 255u - tid; sel[1] = need - above; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + prefix |= sel[0] << shift; + pmask |= 255u << shift; + need = sel[1]; } - threadgroup_barrier(mem_flags::mem_threadgroup); - if (tid < 256) { - uint above = ex; - for (int i = 0; i < int(sg); ++i) above += sgs[i]; - if (above < need && above + cnt >= need) { sel[0] = 255u - tid; sel[1] = need - above; } + } else { + for (int pass = 0; pass < 4; ++pass) { + const int shift = 24 - 8 * pass; + for (uint b = tid; b < nb; b += 1024) { + const uint k = okey(sc[b]); + if ((k & pmask) == prefix) { + const uint w = b < n_bid ? 4u : (n_kv & 3u); + if (w) atomic_fetch_add_explicit(&hist[(k >> shift) & 255u], w, memory_order_relaxed); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + uint cnt = 0, ex = 0; + if (tid < 256) { + cnt = atomic_load_explicit(&hist[255 - tid], memory_order_relaxed); + if (pass < 3) atomic_store_explicit(&hist[255 - tid], 0u, memory_order_relaxed); + ex = simd_prefix_exclusive_sum(cnt); + if (lane == 31) sgs[sg] = ex + cnt; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < 256) { + uint above = ex; + for (int i = 0; i < int(sg); ++i) above += sgs[i]; + if (above < need && above + cnt >= need) { sel[0] = 255u - tid; sel[1] = need - above; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + prefix |= sel[0] << shift; + pmask |= 255u << shift; + need = sel[1]; } - threadgroup_barrier(mem_flags::mem_threadgroup); - prefix |= sel[0] << shift; - pmask |= 255u << shift; - need = sel[1]; - threadgroup_barrier(mem_flags::mem_threadgroup); } const uint T = prefix; const uint per = (nb + 1023) / 1024, b0 = min(nb, uint(tid) * per), b1 = min(nb, b0 + per); diff --git a/metal/kernels/qsa_prompt.metal b/metal/kernels/qsa_prompt.metal index df7caf79..2ee4a4b6 100644 --- a/metal/kernels/qsa_prompt.metal +++ b/metal/kernels/qsa_prompt.metal @@ -30,8 +30,8 @@ kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* constexpr int HPG = QH / QKV, NCH = (QBUDGET + QCH - 1) / QCH; threadgroup float s[16 * QCH]; threadgroup half p[16 * QCH]; - threadgroup int cid[QCH]; - threadgroup half vs[QCH][4]; + threadgroup int cid[QCH]; // coop_k4: byte offsets; otherwise cell ids + threadgroup half4 ks[QCH], vs[QCH]; const int chunk = int(tg.x), g = int(tg.y), t = int(tg.z); const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); const int c0 = chunk * QCH; @@ -42,9 +42,16 @@ kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* } const int nc = min(QCH, int(width) - c0); for (int c = lane; c < QCH; c += 32) { - const int id = c < nc ? ids[t * QBUDGET + c0 + c] : -1; - cid[c] = id; - for (int j = 0; j < 4; ++j) vs[c][j] = id >= 0 ? vsc[(uint(id) * QKV + g) * 4 + j] : half(0.0f); + const int id = c < nc ? (n_kv <= QBUDGET ? c0 + c : ids[t * QBUDGET + c0 + c]) : -1; + cid[c] = coop_k4 ? int((uint(max(id, 0)) * QKV + g) * QD) : id; + if (id >= 0) { + const uint row4 = (uint(id) * QKV + g) * 4; + if (coop_k4) ks[c] = *(device const half4*)(ksc + row4); + vs[c] = *(device const half4*)(vsc + row4); + } else { + ks[c] = half4(0.0h); + vs[c] = half4(0.0h); + } } simdgroup_barrier(mem_flags::mem_threadgroup); if (qsa_stop == 1) { if (lane < HPG) { pp[lane * (QD + 2) + QD] = 1.0f; pp[lane * (QD + 2) + QD + 1] = 1.0f; } return; } @@ -59,25 +66,33 @@ kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* auto acc = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); for (uint16_t i = 0; i < acc.get_capacity(); ++i) if (acc.is_valid_element(i)) acc[i] = 0.0f; + // the second half repeats these eight cells, 16 k later; empty cells read row 0 at scale 0. + uint ko[8]; + float ks8[8][4]; + if (coop_k4) +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const int c = rk.get_multidimensional_index(uint16_t(4 * j))[1]; + ko[j] = uint(cid[c]); +#pragma clang loop unroll(full) + for (int h = 0; h < 4; ++h) ks8[j][h] = float(ks[c][h]); + } for (int k0 = 0; k0 < QD; k0 += 32) { if (!coop_k4) { for (uint16_t i = 0; i < rk.get_capacity(); ++i) { const auto ix = rk.get_multidimensional_index(i); const int k = k0 + ix[0], id = cid[ix[1]]; - const uint row = uint(id) * QKV + g; - rk[i] = id >= 0 ? half(float(kq[row * QD + k]) * float(ksc[row * 4 + (k >> 6)])) : half(0.0f); + // the row is clamped and the select is on the value: a select with the device load in the arm + // loses the threadgroup on the M1 (macOS 26, measured); the discarded value cannot matter + const uint row = uint(max(id, 0)) * QKV + g; + const float v = float(kq[row * QD + k]) * float(ksc[row * 4 + (k >> 6)]); + rk[i] = id >= 0 ? half(v) : half(0.0f); } } else #pragma clang loop unroll(full) for (uint16_t i = 0; i < rk.get_capacity(); i += 4) { - const auto ix = rk.get_multidimensional_index(i); - const int k = k0 + ix[0], c = ix[1], id = cid[c]; - half4 w = half4(0.0h); - if (id >= 0) { - const uint row = uint(id) * QKV + g; - const char4 q4 = *(device const char4*)(kq + row * QD + k); - w = half4(float4(q4) * float(ksc[row * 4 + (k >> 6)])); - } + const int k = k0 + rk.get_multidimensional_index(i)[0], j = (i >> 2) & 7; + const half4 w = half4(float4(*(device const char4*)(kq + (ko[j] + uint(k)))) * ks8[j][k >> 6]); rk[i] = w.x; rk[i + 1] = w.y; rk[i + 2] = w.z; rk[i + 3] = w.w; } tA = tQ.template slice<32, 16>(k0, 0); @@ -91,13 +106,16 @@ kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* // ---- softmax pieces per head (lane = head for the reductions; 12 heads) float mh = -INFINITY, lh = 0.0f; if (lane < HPG) { +#pragma clang loop unroll_count(4) for (int c = 0; c < nc; ++c) mh = max(mh, s[lane * QCH + c] * 0.0625f); +#pragma clang loop unroll_count(4) for (int c = 0; c < QCH; ++c) { const float e = c < nc ? exp(s[lane * QCH + c] * 0.0625f - mh) : 0.0f; lh += e; p[lane * QCH + c] = half(e); } } else if (lane < 16) { +#pragma clang loop unroll_count(4) for (int c = 0; c < QCH; ++c) p[lane * QCH + c] = half(0.0f); } simdgroup_barrier(mem_flags::mem_threadgroup); @@ -115,12 +133,32 @@ kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* auto acc = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); for (uint16_t i = 0; i < acc.get_capacity(); ++i) if (acc.is_valid_element(i)) acc[i] = 0.0f; - for (int c0 = 0; c0 < QCH; c0 += 32) { + for (int c0 = 0; c0 < nc; c0 += 32) { + if (!coop_k4) { #pragma clang loop unroll(full) - for (uint16_t i = 0; i < rv.get_capacity(); ++i) { - const auto ix = rv.get_multidimensional_index(i); // (cell, dim) - const int c = c0 + ix[0], d = d0 + ix[1], id = cid[c]; - rv[i] = id >= 0 ? half(float(vq[(uint(id) * QKV + g) * QD + d]) * float(vs[c][d >> 6])) : half(0.0f); + for (uint16_t i = 0; i < rv.get_capacity(); ++i) { + const auto ix = rv.get_multidimensional_index(i); // (cell, dim) + const int c = c0 + ix[0], d = d0 + ix[1], id = cid[c]; + const uint row = uint(max(id, 0)) * QKV + g; + const float v = float(vq[row * QD + d]) * float(vs[c][d >> 6]); + rv[i] = id >= 0 ? half(v) : half(0.0f); + } + } else { + uint vo[8]; + float vs8[8]; + const int cb = c0 + rv.get_multidimensional_index(uint16_t(0))[0]; +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const int c = cb + (j & 3) + 16 * (j >> 2); + vo[j] = uint(cid[c]) + uint(d0); + vs8[j] = float(vs[c][d0 >> 6]); + } +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rv.get_capacity(); i += 4) { + const int d = rv.get_multidimensional_index(i)[1], r = i < rv.get_capacity() / 2 ? 0 : 4; +#pragma clang loop unroll(full) + for (int j = 0; j < 4; ++j) rv[i + j] = half(float(vq[vo[r + j] + uint(d)]) * vs8[r + j]); + } } tPh = tP.template slice<32, 16>(c0, 0); op.run(tPh, rv, acc); @@ -136,7 +174,13 @@ kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* simdgroup_barrier(mem_flags::mem_threadgroup); acc.store(tT); simdgroup_barrier(mem_flags::mem_threadgroup); - for (int it = lane; it < HPG * 64; it += 32) pp[(it >> 6) * (QD + 2) + d0 + (it & 63)] = s[it]; + for (int it = lane; it < HPG * 16; it += 32) { + const int h = it >> 4, d = (it & 15) * 4; + const float4 v = *(threadgroup const float4*)(s + h * 64 + d); + // odd heads start on an 8-byte boundary, not a float4 boundary. + device float* dst = pp + h * (QD + 2) + d0 + d; + dst[0] = v.x; dst[1] = v.y; dst[2] = v.z; dst[3] = v.w; + } } } if (lane < HPG) { pp[lane * (QD + 2) + QD] = mh; pp[lane * (QD + 2) + QD + 1] = lh; } @@ -160,6 +204,10 @@ kernel void qsa_score_mm(device const float* qidx [[buffer(0)]], device const fl ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { threadgroup float ot[64 * 64]; const int b0 = int(tg.x) * 64, t0 = int(tg.y) * 16; + const uint max_t = min(uint(t0 + 15), a.n_tok - 1u); + const uint max_bid = (a.pos0 + max_t + 1u) >> 2; + // column zero still owns every token's tail score. + if (tg.x != 0 && uint(b0) >= max_bid) return; const uint nb_all = (a.pos0 + a.n_tok) / 4 + 1; // blocks any token of the chunk can score tensor, tensor_inline> tQ((device float*)qidx, dextents(ID, int(a.n_tok) * IH)); tensor, tensor_inline> tK((device float*)pooled, dextents(ID, int(nb_all))); @@ -171,9 +219,10 @@ kernel void qsa_score_mm(device const float* qidx [[buffer(0)]], device const fl op.run(mA, mB, tO); threadgroup_barrier(mem_flags::mem_threadgroup); const uint stride = a.max_cells / 4 + 2; - for (int idx = tid; idx < 16 * 64; idx += 128) { - const int tt = idx >> 6, j = idx & 63, t = t0 + tt, b = b0 + j; - if (t >= int(a.n_tok)) continue; + const int j = tid & 63, b = b0 + j; + for (int tt = tid >> 6; tt < 16; tt += 2) { + const int t = t0 + tt; + if (t >= int(a.n_tok)) break; const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2; if (uint(b) >= n_bid) continue; float s = 0.0f; @@ -181,13 +230,13 @@ kernel void qsa_score_mm(device const float* qidx [[buffer(0)]], device const fl scores[uint(t) * stride + b] = s; } // the tail block of each token in this threadgroup's block range (the first column threadgroup does it) - if (tg.x == 0) + if (tg.x == 0) { + float k[4]; + for (int i = 0; i < 4; ++i) k[i] = dead[lane + 32 * i]; for (int tt = sg; tt < 16; tt += 4) { const int t = t0 + tt; if (t >= int(a.n_tok)) break; const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2; - float k[4]; - for (int i = 0; i < 4; ++i) k[i] = dead[lane + 32 * i]; float s = 0.0f; for (int h = 0; h < IH; ++h) { float d = 0.0f; @@ -198,6 +247,7 @@ kernel void qsa_score_mm(device const float* qidx [[buffer(0)]], device const fl if ((n_kv & 3) != 0) s += 1e9f; if (lane == 0) scores[uint(t) * stride + n_bid] = s; } + } } // qsa_flash: the whole attention of a (kv head g, token t) in one threadgroup (no split-K partials, no merge): the @@ -211,10 +261,10 @@ kernel void qsa_flash(device const half* qh [[buffer(0)]], device const char* kq uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { constexpr int HPG = QH / QKV; - threadgroup float sp[4][16 * QCH]; - threadgroup half p[16 * QCH]; - threadgroup int cid[QCH]; - threadgroup float vs[QCH]; + threadgroup float sp[4][HPG * QCH]; + threadgroup half* p = (threadgroup half*)sp[1]; // P [16][QCH] in fp16 over a dead S partial + threadgroup int cid[QCH]; // coop_k4: byte offsets; otherwise cell ids + threadgroup half4 ks_tg[QCH], vs_tg[QCH]; threadgroup float fac[16], lsum[16], mrun[16]; const int g = int(tg.x), t = int(tg.y); const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); @@ -240,15 +290,26 @@ kernel void qsa_flash(device const half* qh [[buffer(0)]], device const char* kq for (uint16_t i = 0; i < acc.get_capacity(); ++i) if (acc.is_valid_element(i)) tot[i] = 0.0f; const int d0 = sg * 64; + if (tid < 16 - HPG) fac[HPG + tid] = 0.0f; // the padding rows' factor (P's rows are zeroed each chunk) for (int ch = 0; ch < nch; ++ch) { const int c0 = ch * QCH, nc = min(QCH, int(width) - c0); threadgroup_barrier(mem_flags::mem_threadgroup); if (tid < QCH) { - const int id = tid < nc ? ids[t * QBUDGET + c0 + tid] : -1; - cid[tid] = id; - vs[tid] = id >= 0 ? float(vsc[(uint(id) * QKV + g) * 4 + sg * 0]) : 0.0f; // (per-simdgroup scale below) + const int id = tid < nc ? (n_kv <= QBUDGET ? c0 + tid : ids[t * QBUDGET + c0 + tid]) : -1; + cid[tid] = coop_k4 ? int((uint(max(id, 0)) * QKV + g) * QD) : id; + if (id >= 0) { + const uint row4 = (uint(id) * QKV + g) * 4; + ks_tg[tid] = *(device const half4*)(ksc + row4); + vs_tg[tid] = *(device const half4*)(vsc + row4); + } else { + ks_tg[tid] = half4(0.0h); + vs_tg[tid] = half4(0.0h); + } } threadgroup_barrier(mem_flags::mem_threadgroup); + // Eight lanes per head; keep the denominator's original sequential sum. + const int h = sg * 4 + (lane >> 3), hlane = lane & 7; + float sc[8]; // S partial over this simdgroup's 64 dims { auto rk = opS.template get_right_input_cooperative_tensor(); @@ -257,68 +318,114 @@ kernel void qsa_flash(device const half* qh [[buffer(0)]], device const char* kq __tensor_ops_detail::__remove_addrspace_t, float>(); for (uint16_t i = 0; i < s.get_capacity(); ++i) if (s.is_valid_element(i)) s[i] = 0.0f; + // the lane's eight cells (one a run of four k; the second half repeats them 16 k on): row pointers and + // scales once a chunk, unconditional loads (an empty cell reads row 0 at scale 0, and a +-0 operand sums to + // the same value as the skipped one; the branchy loads cost half the kernel) + uint ko[8]; + float ks8[8]; + if (coop_k4) { +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const int c = rk.get_multidimensional_index(uint16_t(4 * j))[1]; + ko[j] = uint(cid[c]) + uint(d0 + rk.get_multidimensional_index(uint16_t(0))[0]); + ks8[j] = float(ks_tg[c][sg]); + } + } +#pragma clang loop unroll(full) for (int k0 = d0; k0 < d0 + 64; k0 += 32) { if (!coop_k4) { for (uint16_t i = 0; i < rk.get_capacity(); ++i) { const auto ix = rk.get_multidimensional_index(i); const int k = k0 + ix[0], id = cid[ix[1]]; - const uint row = uint(id) * QKV + g; - rk[i] = id >= 0 ? half(float(kq[row * QD + k]) * float(ksc[row * 4 + (k >> 6)])) : half(0.0f); + const uint row = uint(max(id, 0)) * QKV + g; + const float v = float(kq[row * QD + k]) * float(ks_tg[ix[1]][k >> 6]); + rk[i] = id >= 0 ? half(v) : half(0.0f); } } else #pragma clang loop unroll(full) for (uint16_t i = 0; i < rk.get_capacity(); i += 4) { - const auto ix = rk.get_multidimensional_index(i); - const int k = k0 + ix[0], c = ix[1], id = cid[c]; - half4 w = half4(0.0h); - if (id >= 0) { - const uint row = uint(id) * QKV + g; - const char4 q4 = *(device const char4*)(kq + row * QD + k); - w = half4(float4(q4) * float(ksc[row * 4 + (k >> 6)])); - } + const int k = k0 - d0 + (i >= rk.get_capacity() / 2 ? 16 : 0), j = (i >> 2) & 7; + const half4 w = half4(float4(*(device const char4*)(kq + ko[j] + uint(k))) * ks8[j]); rk[i] = w.x; rk[i + 1] = w.y; rk[i + 2] = w.z; rk[i + 3] = w.w; } tA = tQ.template slice<32, 16>(k0, 0); opS.run(tA, rk, s); } - tensor, tensor_inline> tS(sp[sg], dextents(QCH, 16)); + // Preserve ((s0 + s1) + (s2 + s3)) while storing all partials together. + tensor, tensor_inline> tS(sp[sg], dextents(QCH, HPG)); s.store(tS); } threadgroup_barrier(mem_flags::mem_threadgroup); - // online softmax per head row: simdgroup 0, lane = row - if (sg == 0 && lane < 16) { - if (lane < HPG) { - float m = mrun[lane]; - for (int c = 0; c < nc; ++c) { - const int q = lane * QCH + c; - m = max(m, (((sp[0][q] + sp[1][q]) + (sp[2][q] + sp[3][q])) * 0.0625f)); - } - const float f = exp(mrun[lane] - m); - float l = lsum[lane] * f; - for (int c = 0; c < QCH; ++c) { - const int q = lane * QCH + c; - const float e = c < nc ? exp(((sp[0][q] + sp[1][q]) + (sp[2][q] + sp[3][q])) * 0.0625f - m) : 0.0f; - l += e; - p[q] = half(e); + if (h < HPG) +#pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) { + const int q = h * QCH + hlane + 8 * k; + sc[k] = ((sp[0][q] + sp[1][q]) + (sp[2][q] + sp[3][q])) * 0.0625f; + } + threadgroup_barrier(mem_flags::mem_threadgroup); // P overwrites sp[1] + if (h < HPG) { + float m = mrun[h]; +#pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) + if (hlane + 8 * k < nc) m = max(m, sc[k]); + m = max(m, simd_shuffle_xor(m, 4)); + m = max(m, simd_shuffle_xor(m, 2)); + m = max(m, simd_shuffle_xor(m, 1)); +#pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) { + const int c = hlane + 8 * k, q = h * QCH + c; + const float e = c < nc ? exp(sc[k] - m) : 0.0f; + p[q] = half(e); + sp[0][q] = e; + } + simdgroup_barrier(mem_flags::mem_threadgroup); + if (hlane == 0) { + const float f = exp(mrun[h] - m); + float l = lsum[h] * f; +#pragma clang loop unroll(full) + for (int c = 0; c < QCH; c += 4) { // in order, four loads at a time + const float4 v = *(threadgroup const float4*)(sp[0] + h * QCH + c); + l += v.x; l += v.y; l += v.z; l += v.w; } - mrun[lane] = m; - lsum[lane] = l; - fac[lane] = f; - } else { - for (int c = 0; c < QCH; ++c) p[lane * QCH + c] = half(0.0f); - fac[lane] = 0.0f; + mrun[h] = m; + lsum[h] = l; + fac[h] = f; } + } else { + for (int c = hlane; c < QCH; c += 8) p[h * QCH + c] = half(0.0f); } threadgroup_barrier(mem_flags::mem_threadgroup); // rescale this simdgroup's accumulator rows, then acc += P V over its 64 dims for (uint16_t i = 0; i < acc.get_capacity(); ++i) if (acc.is_valid_element(i)) acc[i] = 0.0f; - for (int c0 = 0; c0 < QCH; c0 += 32) { // K = 32 cells a run (K = 64 is wrong on the M4) + for (int c0 = 0; c0 < nc; c0 += 32) { // K = 32 cells a run (K = 64 is wrong on the M4) + if (!coop_k4) { #pragma clang loop unroll(full) - for (uint16_t i = 0; i < rv.get_capacity(); ++i) { - const auto ix = rv.get_multidimensional_index(i); // (cell, dim) - const int c = c0 + ix[0], d = d0 + ix[1], id = cid[c]; - rv[i] = id >= 0 ? half(float(vq[(uint(id) * QKV + g) * QD + d]) * float(vsc[(uint(id) * QKV + g) * 4 + sg])) : half(0.0f); + for (uint16_t i = 0; i < rv.get_capacity(); ++i) { + const auto ix = rv.get_multidimensional_index(i); // (cell, dim) + const int c = c0 + ix[0], d = d0 + ix[1], id = cid[c]; + const uint row = uint(max(id, 0)) * QKV + g; + const float v = float(vq[row * QD + d]) * float(vs_tg[c][sg]); + rv[i] = id >= 0 ? half(v) : half(0.0f); + } + } else { + // a lane holds the same eight cells (two runs of four, 16 apart) at each of its dims: as for rk, row + // pointers and scales once a run, unconditional loads + uint vo[8]; + float vs8[8]; + const int cb = c0 + rv.get_multidimensional_index(uint16_t(0))[0]; +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const int c = cb + (j & 3) + 16 * (j >> 2); + vo[j] = uint(cid[c]) + uint(d0); + vs8[j] = float(vs_tg[c][sg]); + } +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rv.get_capacity(); i += 4) { + const int d = rv.get_multidimensional_index(i)[1], r = i < rv.get_capacity() / 2 ? 0 : 4; +#pragma clang loop unroll(full) + for (int j = 0; j < 4; ++j) rv[i + j] = half(float(vq[vo[r + j] + uint(d)]) * vs8[r + j]); + } } tPh = tP.template slice<32, 16>(c0, 0); opO.run(tPh, rv, acc); @@ -327,14 +434,28 @@ kernel void qsa_flash(device const half* qh [[buffer(0)]], device const char* kq if (acc.is_valid_element(i)) tot[i] = fma(tot[i], fac[acc.get_multidimensional_index(i)[1]], acc[i]); } threadgroup_barrier(mem_flags::mem_threadgroup); - for (uint16_t i = 0; i < acc.get_capacity(); ++i) - if (acc.is_valid_element(i)) { - const auto ix = acc.get_multidimensional_index(i); // (dim, row) + if (!coop_k4) { + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) { + const auto ix = acc.get_multidimensional_index(i); // (dim, row) + const int h = ix[1], d = d0 + ix[0]; + if (h < HPG) { + const int hh = g * HPG + h; + const float gate = qfull[(uint(t) * QH + hh) * 2 * QD + QD + d]; + xqh[(uint(t) * QH + hh) * QD + d] = half(tot[i] / lsum[h] * (1.0f / (1.0f + exp(-gate)))); + } + } + } else +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < acc.get_capacity(); i += 4) { // runs of four dims of one row + const auto ix = acc.get_multidimensional_index(i); const int h = ix[1], d = d0 + ix[0]; if (h < HPG) { const int hh = g * HPG + h; - const float gate = qfull[(uint(t) * QH + hh) * 2 * QD + QD + d]; - xqh[(uint(t) * QH + hh) * QD + d] = half(tot[i] / lsum[h] * (1.0f / (1.0f + exp(-gate)))); + const float l = lsum[h]; + const float4 gate = *(device const float4*)(qfull + (uint(t) * QH + hh) * 2 * QD + QD + d); + const float4 v = float4(tot[i], tot[i + 1], tot[i + 2], tot[i + 3]) / l * (1.0f / (1.0f + exp(-gate))); + *(device half4*)(xqh + (uint(t) * QH + hh) * QD + d) = half4(v); } } } diff --git a/metal/kernels/quant.metal b/metal/kernels/quant.metal index 51fecca7..fe667be5 100644 --- a/metal/kernels/quant.metal +++ b/metal/kernels/quant.metal @@ -34,11 +34,11 @@ template struct fmt; // QK block elements, BB block bytes, UNIT weights per unit, TAB words of threadgroup table template <> struct fmt { enum { QK = 64, BB = 18, UNIT = 16, TAB = 1 }; }; template <> struct fmt { enum { QK = 64, BB = 18, UNIT = 16, TAB = 1 }; }; -template <> struct fmt { enum { QK = 256, BB = 98, UNIT = 16, TAB = 256 + 32 }; }; -template <> struct fmt { enum { QK = 256, BB = 66, UNIT = 16, TAB = 512 + 32 }; }; -template <> struct fmt { enum { QK = 256, BB = 74, UNIT = 16, TAB = 1024 + 32 }; }; -template <> struct fmt { enum { QK = 256, BB = 82, UNIT = 16, TAB = 2048 + 32 }; }; -template <> struct fmt { enum { QK = 256, BB = 110, UNIT = 16, TAB = 512 + 32 }; }; +template <> struct fmt { enum { QK = 256, BB = 98, UNIT = 16, TAB = 256 }; }; +template <> struct fmt { enum { QK = 256, BB = 66, UNIT = 16, TAB = 512 }; }; +template <> struct fmt { enum { QK = 256, BB = 74, UNIT = 16, TAB = 1024 }; }; +template <> struct fmt { enum { QK = 256, BB = 82, UNIT = 16, TAB = 2048 }; }; +template <> struct fmt { enum { QK = 256, BB = 110, UNIT = 16, TAB = 512 }; }; template <> struct fmt { enum { QK = 32, BB = 18, UNIT = 32, TAB = 256 }; }; template <> struct fmt { enum { QK = 256, BB = 136, UNIT = 32, TAB = 256 }; }; template <> struct fmt { enum { QK = 256, BB = 110, UNIT = 32, TAB = 1 }; }; @@ -59,15 +59,16 @@ static inline uint rd32(device const uchar* p) { static inline float rdh(device const uchar* p) { return float(*(device const half*)p); } // ---- threadgroup tables ------------------------------------------------------------------------------------------- -// layout in `tab` (uint words): [0, G) the grid, [G, G + 32) ksigns_iq2xs as bytes; IQ4_NL/XS: 16 floats. +// layout in `tab` (uint words): [0, G) the grid; IQ4_NL/XS: 256 half2. template static inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth); -// QUANT_CTAB (function constant 23): the i-quant grids and the sign table read from ggml-common's constant arrays +// QUANT_CTAB (function constant 23): the i-quant grids read from ggml-common's constant arrays // instead of a threadgroup copy (no table load and no barrier wait at the threadgroup's start); the same values. constant bool QUANT_CTAB [[function_constant(23)]]; constant bool ctab = is_function_constant_defined(QUANT_CTAB) && QUANT_CTAB; -static inline uint ksg(threadgroup const uint* tab, int at, uint i) { - return ctab ? uint(ksigns_iq2xs[i]) : uint(((threadgroup const uchar*)(tab + at))[i]); +// ksigns_iq2xs[j] = j | parity(j) << 7 (an even count of minus signs): the low 4 bits of run hh (0 or 1). +static inline uint sign4_iq2xs(uint j, int hh) { + return hh ? ((j >> 4) | ((popcount(j) & 1u) << 3)) : j; } static inline uint g_iq3xxs(threadgroup const uint* tab, uint i) { return ctab ? iq3xxs_grid[i] : tab[i]; } static inline uint g_iq3s(threadgroup const uint* tab, uint i) { return ctab ? iq3s_grid[i] : tab[i]; } @@ -88,34 +89,25 @@ template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} -static inline void load_signs(threadgroup uint* tab, int at, ushort tid, ushort nth) { - threadgroup uchar* s = (threadgroup uchar*)(tab + at); - for (int i = tid; i < 128; i += nth) s[i] = ksigns_iq2xs[i]; -} template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { if (ctab) return; for (int i = tid; i < 256; i += nth) tab[i] = iq3xxs_grid[i]; - load_signs(tab, 256, tid, nth); } template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { if (ctab) return; for (int i = tid; i < 512; i += nth) tab[i] = iq3s_grid[i]; - load_signs(tab, 512, tid, nth); } template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { if (ctab) return; for (int i = tid; i < 256; i += nth) { tab[2 * i] = uint(iq2xxs_grid[i]); tab[2 * i + 1] = uint(iq2xxs_grid[i] >> 32); } - load_signs(tab, 512, tid, nth); } template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { if (ctab) return; for (int i = tid; i < 512; i += nth) { tab[2 * i] = uint(iq2xs_grid[i]); tab[2 * i + 1] = uint(iq2xs_grid[i] >> 32); } - load_signs(tab, 1024, tid, nth); } template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { if (ctab) return; for (int i = tid; i < 1024; i += nth) { tab[2 * i] = uint(iq2s_grid[i]); tab[2 * i + 1] = uint(iq2s_grid[i] >> 32); } - load_signs(tab, 2048, tid, nth); } // IQ4: a byte -> (codebook[low nibble], codebook[high nibble]) as half2, so one lookup gives two weights template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { @@ -153,74 +145,82 @@ static inline void put_signed4(thread float* q, uint g, uint s) { // IQ3_XXS: 256 = 8 x 32; unit = half of a 32 (two grid pairs). template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { const int ib = ui >> 1, h = ui & 1; - const uint aux = rd32(b + 2 + 64 + 4 * ib); - device const uchar* qs = b + 2 + 8 * ib + 4 * h; + const uint aux = rd32(b + 2 + 64 + 4 * ib), s0 = aux >> (14 * h); + device const uchar* qs = b + 2 + 4 * ui; + #pragma clang loop unroll(full) for (int l2 = 0; l2 < 2; ++l2) { - const int l = 2 * h + l2; - const uint s = ksg(tab, 256, (aux >> (7 * l)) & 127); + const uint s = (s0 >> (7 * l2)) & 127; put_signed4(q + 8 * l2, g_iq3xxs(tab, qs[2 * l2]), s); - put_signed4(q + 8 * l2 + 4, g_iq3xxs(tab, qs[2 * l2 + 1]), s >> 4); + put_signed4(q + 8 * l2 + 4, g_iq3xxs(tab, qs[2 * l2 + 1]), sign4_iq2xs(s, 1)); } - return rdh(b) * (0.5f + float(aux >> 28)) * 0.5f; + return rdh(b) * fma(float(aux >> 28), 0.5f, 0.25f); } // IQ2_XXS: per 32: two u32 (4 grid indices, then 4 x 7 sign bits + 4-bit scale). template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { const int ib = ui >> 1, h = ui & 1; - const uint a0 = rd32(b + 2 + 8 * ib), a1 = rd32(b + 2 + 8 * ib + 4); + device const uchar* qs = b + 2 + 4 * ib + 2 * ui; + const uint a1 = rd32(b + 2 + 8 * ib + 4), s0 = a1 >> (14 * h); + #pragma clang loop unroll(full) for (int l2 = 0; l2 < 2; ++l2) { - const int l = 2 * h + l2; - const uint idx = (a0 >> (8 * l)) & 255; - const uint s = ksg(tab, 512, (a1 >> (7 * l)) & 127); - put_signed4(q + 8 * l2, g_iq2xxs(tab, idx, 0), s); - put_signed4(q + 8 * l2 + 4, g_iq2xxs(tab, idx, 1), s >> 4); + const uint idx = qs[l2], s = (s0 >> (7 * l2)) & 127; + const uint2 g = ctab ? as_type(iq2xxs_grid[idx]) : ((threadgroup const uint2*)tab)[idx]; + put_signed4(q + 8 * l2, g.x, s); + put_signed4(q + 8 * l2 + 4, g.y, sign4_iq2xs(s, 1)); } - return rdh(b) * (0.5f + float(a1 >> 28)) * 0.25f; + return rdh(b) * fma(float(a1 >> 28), 0.25f, 0.125f); } // IQ2_XS: uint16 qs[32] (9-bit grid index + 7-bit sign index), scales[8] nibbles. template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { const int ib = ui >> 1, h = ui & 1; - device const ushort* qs = (device const ushort*)(b + 2); + device const ushort* qs = (device const ushort*)(b + 2) + 2 * ui; + #pragma clang loop unroll(full) for (int l2 = 0; l2 < 2; ++l2) { - const uint v = qs[4 * ib + 2 * h + l2]; - const uint idx = v & 511, s = ksg(tab, 1024, v >> 9); - put_signed4(q + 8 * l2, g_iq2xs(tab, idx, 0), s); - put_signed4(q + 8 * l2 + 4, g_iq2xs(tab, idx, 1), s >> 4); + const uint v = qs[l2]; + const uint idx = v & 511, s = v >> 9; + const uint2 g = ctab ? as_type(iq2xs_grid[idx]) : ((threadgroup const uint2*)tab)[idx]; + put_signed4(q + 8 * l2, g.x, s); + put_signed4(q + 8 * l2 + 4, g.y, sign4_iq2xs(s, 1)); } const uint sc = b[66 + ib]; - return rdh(b) * (0.5f + float((sc >> (4 * h)) & 15)) * 0.25f; + return rdh(b) * fma(float((sc >> (4 * h)) & 15), 0.25f, 0.125f); } // IQ2_S: qs[32] low 8 bits of the 10-bit index, signs[32], qh[8] (2 high bits per group of 8), scales[8]. template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { const int ib = ui >> 1, h = ui & 1; - const uint qh = b[66 + ib]; + const uint q2 = *(device const ushort*)(b + 2 + 2 * ui); + const uint s2 = *(device const ushort*)(b + 34 + 2 * ui); + const uint qh = uint(b[66 + ib]) >> (4 * h); + #pragma clang loop unroll(full) for (int l2 = 0; l2 < 2; ++l2) { - const int l = 2 * h + l2; - const uint idx = uint(b[2 + 4 * ib + l]) | ((qh << (8 - 2 * l)) & 0x300); - const uint s = b[2 + 32 + 4 * ib + l]; - put_signed4(q + 8 * l2, g_iq2s(tab, idx, 0), s); - put_signed4(q + 8 * l2 + 4, g_iq2s(tab, idx, 1), s >> 4); + const uint idx = ((q2 >> (8 * l2)) & 255) | ((qh << (8 - 2 * l2)) & 0x300); + const uint s = (s2 >> (8 * l2)) & 255; + const uint2 g = ctab ? as_type(iq2s_grid[idx]) : ((threadgroup const uint2*)tab)[idx]; + put_signed4(q + 8 * l2, g.x, s); + put_signed4(q + 8 * l2 + 4, g.y, s >> 4); } const uint sc = b[74 + ib]; - return rdh(b) * (0.5f + float((sc >> (4 * h)) & 15)) * 0.25f; + return rdh(b) * fma(float((sc >> (4 * h)) & 15), 0.25f, 0.125f); } // IQ3_S: qs[64] low 8 bits of 9-bit indices, qh[8], signs[32], scales[4] (nibble per 32). template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { const int ib = ui >> 1, h = ui & 1; + const uint s2 = *(device const ushort*)(b + 74 + 2 * ui); const uint qh = b[66 + ib]; + #pragma clang loop unroll(full) for (int l2 = 0; l2 < 2; ++l2) { const int l = 2 * h + l2; const uint i1 = uint(b[2 + 8 * ib + 2 * l]) | ((qh << (8 - 2 * l)) & 256); const uint i2 = uint(b[2 + 8 * ib + 2 * l + 1]) | ((qh << (7 - 2 * l)) & 256); - const uint s = b[74 + 4 * ib + l]; + const uint s = (s2 >> (8 * l2)) & 255; put_signed4(q + 8 * l2, g_iq3s(tab, i1), s); put_signed4(q + 8 * l2 + 4, g_iq3s(tab, i2), s >> 4); } const uint sc = b[106 + (ib >> 1)]; - return rdh(b) * float(1 + 2 * ((sc >> (4 * (ib & 1))) & 15)); + return rdh(b) * fma(float((sc >> (4 * (ib & 1))) & 15), 2.0f, 1.0f); } // IQ4_NL: 32 = 1 unit; low nibbles are elements 0..15, high 16..31 (one half2 lookup per byte). @@ -238,6 +238,11 @@ template <> inline float deq(device const uchar* b, int, thread float* } // IQ4_XS: 256 = 8 units of 32; 6-bit scale from scales_l nibble + 2 bits of scales_h, minus 32. +static inline float iq4xs_scale(device const uchar* b, int ui) { + const uint sh = *(device const ushort*)(b + 2); + const int ls = int((b[4 + (ui >> 1)] >> (4 * (ui & 1))) & 15) | int(((sh >> (2 * ui)) & 3) << 4); + return rdh(b) * float(ls - 32); +} template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { device const ushort* p = (device const ushort*)(b + 8 + 16 * ui); for (int j = 0; j < 8; ++j) { @@ -248,9 +253,7 @@ template <> inline float deq(device const uchar* b, int ui, thread flo q[2 * j + 1] = float(w1.x); q[2 * j + 17] = float(w1.y); } - const uint sh = *(device const ushort*)(b + 2); - const int ls = int((b[4 + (ui >> 1)] >> (4 * (ui & 1))) & 15) | int(((sh >> (2 * ui)) & 3) << 4); - return rdh(b) * float(ls - 32); + return iq4xs_scale(b, ui); } template <> inline float deq(device const uchar* b, int, thread float* q, threadgroup const uint*) { diff --git a/metal/kernels/router.metal b/metal/kernels/router.metal index e5c75274..3ad325d1 100644 --- a/metal/kernels/router.metal +++ b/metal/kernels/router.metal @@ -5,52 +5,111 @@ // the counter. No threadgroup waits for another. ids / w / plan are bitwise moe_route's. // Buffers: W, xp, s16, logits, args (GemvArgs: n_out = n_expert), ids, w, plan, MoeArgs, counter (one uint, zero). +#ifndef STRATA_MOCE_FC_METAL +#define STRATA_MOCE_FC_METAL +constant int kMoceProtectFc [[function_constant(24)]]; +constant int kMoceProtect = is_function_constant_defined(kMoceProtectFc) ? kMoceProtectFc : 3; +constant uint kMoceScaleR1BitsFc [[function_constant(25)]]; +constant float kMoceScaleR1 = is_function_constant_defined(kMoceScaleR1BitsFc) ? as_type(kMoceScaleR1BitsFc) : 1.0f; +#endif static inline void route_body(device const float* logits, device int* ids, device float* w, device int* plan, - constant MoeArgs& a, ushort tid, ushort nth, ushort sg, ushort lane, threadgroup int* s_grp, - threadgroup int* s_ent) { + constant MoeArgs& a, device const uint* smap, ushort tid, ushort nth, ushort sg, ushort lane, + threadgroup int* s_grp, threadgroup int* s_ent) { if (sg < a.n_tok) { device const float* l = logits + sg * 512; + const bool e512 = (a.n_expert == 512u); float v[16]; - for (int i = 0; i < 16; ++i) v[i] = l[lane + 32 * i]; - float mx = -INFINITY; - for (int i = 0; i < 16; ++i) mx = max(mx, v[i]); - for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); - float sum = 0.0f; - for (int i = 0; i < 16; ++i) { - v[i] = exp(v[i] - mx); - sum += v[i]; - } - for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); - const float rcp = 1.0f / sum; - for (int i = 0; i < 16; ++i) { - v[i] *= rcp; - if (isnan(v[i])) v[i] = -FLT_MAX; + if (e512) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + const float x = l[lane + 32 * i]; + v[i] = isnan(x) ? -FLT_MAX : x; + } + } else { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + const int idx = lane + 32 * i; + const float x = idx < int(a.n_expert) ? l[idx] : -INFINITY; + v[i] = isnan(x) ? -FLT_MAX : x; + } } - float selected = 0.0f, ssum = 0.0f; + const bool draft_moce = (a.pad & 0x80000000u) != 0u && sg > 0; + const float base_moce = (a.pad & 0x7FFFFFFFu) != 0u ? as_type(a.pad & 0x7FFFFFFFu) : 1.0f; + const float moce_scale = draft_moce ? (base_moce * base_moce) : base_moce; + const float moce_bonus = moce_scale > 1.0f ? log(moce_scale) : 0.0f; + const int r_prot = draft_moce ? 1 : kMoceProtect; + const float r1_scale = kMoceScaleR1; + const float r1_bonus = r1_scale > 1.0f ? log(r1_scale) : 0.0f; + const float r_step_bonus = (r1_scale > 1.0f && moce_scale > 1.0f) ? (moce_bonus - r1_bonus) : moce_bonus; + const bool has_r1 = (r1_scale > 1.0f && r_prot > 1); + const int min_scaled_rank = has_r1 ? 1 : r_prot; + ushort scaled_mask = 0; + float sel_logit = -INFINITY, mx = -INFINITY; for (int rank = 0; rank < 10; ++rank) { + if (has_r1) { + if (rank == 1 && smap != nullptr) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + const int idx = lane + 32 * i; + if ((e512 || idx < int(a.n_expert)) && v[i] > -FLT_MAX && smap[idx] != 0xFFFFFFFFu) { + v[i] += r1_bonus; + scaled_mask |= ushort(1u << i); + } + } + } else if (rank == r_prot && moce_scale > 1.0f && smap != nullptr) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + if ((scaled_mask & ushort(1u << i)) != 0u && v[i] > -FLT_MAX) { + v[i] += r_step_bonus; + } + } + } + } else if (rank == r_prot && moce_scale > 1.0f && smap != nullptr) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + const int idx = lane + 32 * i; + if ((e512 || idx < int(a.n_expert)) && v[i] > -FLT_MAX && smap[idx] != 0xFFFFFFFFu) { + v[i] += moce_bonus; + scaled_mask |= ushort(1u << i); + } + } + } float best = v[0]; int ex = lane; + #pragma clang loop unroll(full) for (int i = 1; i < 16; ++i) if (v[i] > best) { best = v[i]; ex = lane + 32 * i; } + #pragma clang loop unroll(full) for (ushort m = 16; m > 0; m >>= 1) { const float ob = simd_shuffle_xor(best, m); const int oe = simd_shuffle_xor(ex, m); if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } } + if (rank == 0) mx = best; + float true_logit = best; + if (rank >= min_scaled_rank && (moce_scale > 1.0f || r1_scale > 1.0f)) { + const ushort win_lane = ushort(ex & 31); + const ushort win_i = ushort(ex >> 5); + const ushort win_mask = simd_broadcast(scaled_mask, win_lane); + if ((win_mask & ushort(1u << win_i)) != 0u) { + const bool r1_scaled = has_r1 && (rank < r_prot || moce_scale <= 1.0f); + true_logit = best - (r1_scaled ? r1_bonus : moce_bonus); + } + } if ((ex & 31) == lane) { - v[ex / 32] = -INFINITY; + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) v[i] = i == (ex >> 5) ? -INFINITY : v[i]; ids[sg * 10 + rank] = ex; - ssum += best; } - if (rank == lane) selected = best; + if (rank == lane) sel_logit = true_logit; } - for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); + float selected = (lane < 10) ? exp(sel_logit - mx) : 0.0f; + float ssum = simd_sum(selected); ssum = max(ssum, 6.103515625e-5f); - const float inv = 1.0f / ssum; - if (lane < 10) w[sg * 10 + lane] = selected * inv; + if (lane < 10) w[sg * 10 + lane] = selected * (1.0f / ssum); } threadgroup_barrier(mem_flags::mem_device); - plan_body(ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent); + plan_body(ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent, smap); } template @@ -58,13 +117,13 @@ kernel void router_route(device const uchar* W [[buffer(0)]], device const float device const float* s16 [[buffer(2)]], device float* logits [[buffer(3)]], constant GemvArgs& ga [[buffer(4)]], device int* ids [[buffer(5)]], device float* w [[buffer(6)]], device int* plan [[buffer(7)]], constant MoeArgs& a [[buffer(8)]], - device atomic_uint* counter [[buffer(9)]], + device atomic_uint* counter [[buffer(9)]], device const uint* smap [[buffer(10)]], uint tgx [[threadgroup_position_in_grid]], uint ngx [[threadgroups_per_grid]], ushort tid [[thread_index_in_threadgroup]], ushort nth [[threads_per_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { constexpr int UNIT = fmt::UNIT; threadgroup uint tab[fmt::TAB]; - threadgroup int s_grp[32], s_ent[32]; + threadgroup int s_grp[64], s_ent[64]; threadgroup uint last; load_tab(tab, tid, nth); threadgroup_barrier(mem_flags::mem_threadgroup); @@ -74,17 +133,20 @@ kernel void router_route(device const uchar* W [[buffer(0)]], device const float device const uchar* rp = W + (ulong)rr * ga.row_bytes; const int U = int(ga.n_in) / UNIT, s16n = int(ga.n_in) / 16; float acc[8]; +#pragma clang loop unroll(full) for (int t = 0; t < NTOK; ++t) acc[t] = 0.0f; for (int u = lane; u < U; u += 32) { float q[UNIT]; float4 sc; Dec::run(rp, u, q, sc, tab); +#pragma clang loop unroll(full) for (int t = 0; t < NTOK; ++t) { float xs[UNIT], xsum[2]; X::load(xp + t * int(ga.x_stride), s16 + t * s16n, u, xs, xsum); acc[t] += unit_combine(q, sc, xs, xsum); } } +#pragma clang loop unroll(full) for (int t = 0; t < NTOK; ++t) { const float v = lane_sum<32>(acc[t]); if (lane == 0 && row < int(ga.n_out)) logits[t * int(ga.y_stride) + row] = v; @@ -96,9 +158,11 @@ kernel void router_route(device const uchar* W [[buffer(0)]], device const float threadgroup_barrier(mem_flags::mem_threadgroup); if (!last) return; atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); - route_body(logits, ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent); + route_body(logits, ids, w, plan, a, smap, tid, nth, sg, lane, s_grp, s_ent); if (tid == 0) atomic_store_explicit(counter, 0u, memory_order_relaxed); } template [[host_name("router_route_30")]] kernel decltype(router_route) router_route; template [[host_name("router_route_0")]] kernel decltype(router_route) router_route; + +template [[host_name("router_route_8")]] kernel decltype(router_route) router_route; diff --git a/metal/kernels/sample.metal b/metal/kernels/sample.metal index 8443d360..4531ca2a 100644 --- a/metal/kernels/sample.metal +++ b/metal/kernels/sample.metal @@ -13,6 +13,7 @@ static inline void better(thread float& bv, thread int& bi, float v, int i) { static inline void tg_argmax(thread float& bv, thread int& bi, threadgroup float* sv, threadgroup int* si, ushort sg, ushort lane) { +#pragma clang loop unroll(full) for (ushort m = 16; m > 0; m >>= 1) { const float ov = simd_shuffle_xor(bv, m); const int oi = simd_shuffle_xor(bi, m); @@ -22,6 +23,7 @@ static inline void tg_argmax(thread float& bv, thread int& bi, threadgroup float threadgroup_barrier(mem_flags::mem_threadgroup); bv = -INFINITY; bi = 0x7fffffff; +#pragma clang loop unroll(full) for (int k = 0; k < 8; ++k) better(bv, bi, sv[k], si[k]); } diff --git a/metal/kernels/units.metal b/metal/kernels/units.metal index 15f584cb..25d0761a 100644 --- a/metal/kernels/units.metal +++ b/metal/kernels/units.metal @@ -107,26 +107,33 @@ template <> struct X { } }; -// xprep: one thread per 16 consecutive values of one token row; XP = [4 kinds][n_tok][n] then S16 [n_tok][n / 16] +// xprep: one thread per 16 consecutive values of one token row; XP = [4 kinds][n_tok][n] then S16 [n_tok][n / 16]. +// raw and plane are 16-float aligned: the four planes are written as float4s (mtp.metal's mtp_prepare, bit for bit) kernel void xprep(device const float* x [[buffer(0)]], device float* xp [[buffer(1)]], constant uint& n [[buffer(2)]], constant uint& n_tok [[buffer(3)]], uint gid [[thread_position_in_grid]]) { const uint per = n / 16; if (gid >= per * n_tok) return; const uint t = gid / per, i = gid % per; - device const float* src = x + t * n + 16 * i; + device const float4* src = (device const float4*)(x + t * n + 16 * i); const uint plane = n_tok * n; device float* raw = xp + t * n + 16 * i; - float v[16]; - load16(src, v); + device float4* r0 = (device float4*)raw; + device float4* r1 = (device float4*)(raw + plane); + device float4* r2 = (device float4*)(raw + 2 * plane); + device float4* r3 = (device float4*)(raw + 3 * plane); + const float4 q2lo = float4(1.0f, 0.25f, 0.0625f, 0.015625f), + q2hi = float4(0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f), + p4 = float4(1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f); float s0 = 0, s1 = 0; - for (int j = 0; j < 16; j += 2) { s0 += v[j]; s1 += v[j + 1]; } - const float q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f}; - const float p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; - for (int j = 0; j < 16; ++j) { - raw[j] = v[j]; - raw[plane + j] = v[j] * q2s[j & 7]; - raw[2 * plane + j] = v[j] * p4s[j & 3]; - raw[3 * plane + j] = (j & 1) ? v[j] * 0.00390625f : v[j]; +#pragma clang loop unroll(full) + for (int j = 0; j < 4; ++j) { + const float4 v = src[j]; + s0 += v.x; s0 += v.z; + s1 += v.y; s1 += v.w; + r0[j] = v; + r1[j] = v * ((j & 1) ? q2hi : q2lo); + r2[j] = v * p4; + r3[j] = float4(v.x, v.y * 0.00390625f, v.z, v.w * 0.00390625f); } xp[4 * plane + t * per + i] = s0 + s1; } diff --git a/metal/src/convert_k256_q2_0.cpp b/metal/src/convert_k256_q2_0.cpp new file mode 100644 index 00000000..3652268a --- /dev/null +++ b/metal/src/convert_k256_q2_0.cpp @@ -0,0 +1,232 @@ +// metal/src/convert_k256_q2_0.cpp - Multi-threaded C++ converter that converts routed experts in +// K256-metal.gguf + K256-metal.gguf.experts from mixed IQ3_S/IQ4_NL (23.42 GiB, 3 size classes) +// to uniform Q2_0 (type 42, 1.3824 MB/expert, 15.82 GiB, 1 unified size class of 85 pages). + +#include "gguf_lite.h" +#include "ref_quant.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +static constexpr int NL = 48; +static constexpr int NE = 256; +static constexpr uint64_t ALIGN = 16384; +static constexpr uint64_t Q2_SLICE_BYTES = 460800; // 2560 * 640 / 64 * 18 + +// Strata's Q2_0 quantizer: 64-element blocks, grid {-1, 0, 1, 2} * d, code = q + 1 (0..3) +static void quantize_q2_0_slice(const float* x, size_t n_elems, uint8_t* out) { + const size_t n_blocks = n_elems / 64; + for (size_t b = 0; b < n_blocks; ++b) { + const float* xb = x + b * 64; + float amax = 0.0f; + for (int i = 0; i < 64; ++i) amax = std::max(amax, std::fabs(xb[i])); + float best_d = 0.0f; + if (amax > 0.0f) { + float best_err = 1e30f; + for (int step = 0; step <= 8; ++step) { + const float f = 0.5f + 0.0625f * (float) step; + const float d = ref::f16_to_f32(ref::f32_to_f16(amax * f)); + if (d <= 0.0f) continue; + const float inv = 1.0f / d; + float err = 0.0f; + for (int i = 0; i < 64; ++i) { + int q = (int) std::lrintf(xb[i] * inv); + q = std::max(-1, std::min(2, q)); + const float diff = (float) q * d - xb[i]; + err += diff * diff; + } + if (err < best_err) { + best_err = err; + best_d = d; + } + } + } + const uint16_t d16 = ref::f32_to_f16(best_d); + const float d_stored = ref::f16_to_f32(d16); + const float inv = d_stored > 0.0f ? (1.0f / d_stored) : 0.0f; + uint8_t* dst = out + b * 18; + std::memcpy(dst, &d16, 2); + for (int j = 0; j < 16; ++j) { + uint8_t byte = 0; + for (int k = 0; k < 4; ++k) { + const float val = xb[j * 4 + k]; + int q = d_stored > 0.0f ? (int) std::lrintf(val * inv) : 0; + q = std::max(-1, std::min(2, q)); + const uint8_t code = (uint8_t) (q + 1); + byte |= (uint8_t) (code << (2 * k)); + } + dst[2 + j] = byte; + } + } +} + +int main(int argc, char** argv) { + if (argc < 3) { + std::fprintf(stderr, "usage: convert_k256_q2_0 \n"); + return 2; + } + const std::string in_gguf = argv[1]; + const std::string out_gguf = argv[2]; + const std::string in_exp = in_gguf + ".experts"; + const std::string out_exp = out_gguf + ".experts"; + + auto t0 = std::chrono::steady_clock::now(); + gguf::File f; + std::string err; + if (!f.open(in_gguf, err)) { + std::fprintf(stderr, "Failed to open %s: %s\n", in_gguf.c_str(), err.c_str()); + return 1; + } + + // 1. Copy in_gguf (4.22 GiB) to out_gguf and patch the 144 _exps.weight tensor types in the GGUF header to Q2_0 (42) + { + std::vector hdr(16 * 1024 * 1024); + const size_t hlen = std::min(f.file_bytes, hdr.size()); + std::memcpy(hdr.data(), f.map, hlen); + int patched = 0; + for (int l = 0; l < NL; ++l) { + for (const char* role : {"gate", "up", "down"}) { + std::string name = "blk." + std::to_string(l) + ".ffn_" + role + "_exps.weight"; + const uint64_t nlen = name.size(); + std::vector pat(8 + nlen); + std::memcpy(pat.data(), &nlen, 8); + std::memcpy(pat.data() + 8, name.data(), nlen); + auto it = std::search(hdr.begin(), hdr.begin() + (long) hlen, pat.begin(), pat.end()); + if (it == hdr.begin() + (long) hlen) { + std::fprintf(stderr, "Tensor header not found: %s\n", name.c_str()); + return 1; + } + size_t pos = (size_t) (it - hdr.begin()) + pat.size(); + uint32_t nd = 0; + std::memcpy(&nd, hdr.data() + pos, 4); + pos += 4 + (size_t) nd * 8; + const uint32_t q2_type = (uint32_t) ref::Q2_0; + std::memcpy(hdr.data() + pos, &q2_type, 4); + ++patched; + } + } + std::printf("Patched %d expert tensor headers to Q2_0 (42)\n", patched); + int ofd = ::open(out_gguf.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0644); + if (ofd < 0) { std::perror("open out_gguf"); return 1; } + fcntl(ofd, F_NOCACHE, 1); + if (::pwrite(ofd, hdr.data(), hlen, 0) != (ssize_t) hlen) { std::perror("pwrite hdr"); return 1; } + size_t pos = hlen; + const size_t chunk = 64ull << 20; + while (pos < f.file_bytes) { + const size_t n = std::min(chunk, f.file_bytes - pos); + if (::pwrite(ofd, (const uint8_t*) f.map + pos, n, (off_t) pos) != (ssize_t) n) { std::perror("pwrite body"); return 1; } + pos += n; + } + ::fsync(ofd); + ::close(ofd); + } + + // 2. Convert STRX in_exp -> out_exp with uniform Q2_0 (3 * 460800 = 1382400 bytes per expert) + int ifd = ::open(in_exp.c_str(), O_RDONLY); + if (ifd < 0) { std::perror("open in_exp"); return 1; } + fcntl(ifd, F_NOCACHE, 1); + uint64_t in_tab[NL * 4]; + if (::pread(ifd, in_tab, sizeof(in_tab), 64) != (ssize_t) sizeof(in_tab)) { std::perror("pread in_tab"); return 1; } + + const uint64_t exp_data0 = ALIGN; + uint64_t out_offs[NL]; + uint64_t at = exp_data0; + const uint64_t out_exp_bytes = 3 * Q2_SLICE_BYTES; + for (int l = 0; l < NL; ++l) { + out_offs[l] = at; + at += (uint64_t) NE * out_exp_bytes; + at = (at + ALIGN - 1) / ALIGN * ALIGN; + } + + int efd = ::open(out_exp.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0644); + if (efd < 0) { std::perror("open out_exp"); return 1; } + fcntl(efd, F_NOCACHE, 1); + + std::vector ehead(exp_data0, 0); + std::memcpy(ehead.data(), "STRX", 4); + const uint32_t ver = 1, nl = NL, ne = NE; + const uint64_t src_sz = f.file_bytes; + std::memcpy(ehead.data() + 4, &ver, 4); + std::memcpy(ehead.data() + 8, &nl, 4); + std::memcpy(ehead.data() + 12, &ne, 4); + std::memcpy(ehead.data() + 16, &src_sz, 8); + for (int l = 0; l < NL; ++l) { + uint64_t row[4] = {out_offs[l], Q2_SLICE_BYTES, Q2_SLICE_BYTES, Q2_SLICE_BYTES}; + std::memcpy(ehead.data() + 64 + l * 32, row, 32); + } + if (::pwrite(efd, ehead.data(), ehead.size(), 0) != (ssize_t) ehead.size()) { std::perror("pwrite ehead"); return 1; } + + for (int l = 0; l < NL; ++l) { + const std::string b = "blk." + std::to_string(l) + "."; + const int tg = f.find(b + "ffn_gate_exps.weight")->type; + const int tu = f.find(b + "ffn_up_exps.weight")->type; + const int td = f.find(b + "ffn_down_exps.weight")->type; + const uint64_t in_off = in_tab[l * 4 + 0]; + const uint64_t gs = in_tab[l * 4 + 1], us = in_tab[l * 4 + 2], ds = in_tab[l * 4 + 3]; + const uint64_t in_step = gs + us + ds; + + std::vector in_buf((size_t) NE * in_step); + if (::pread(ifd, in_buf.data(), in_buf.size(), (off_t) in_off) != (ssize_t) in_buf.size()) { + std::perror("pread layer"); + return 1; + } + std::vector out_buf((size_t) NE * out_exp_bytes); + + std::atomic next_e{0}; + const int nth = 10; + std::vector workers; + for (int t = 0; t < nth; ++t) { + workers.emplace_back([&]() { + std::vector f32((size_t) 2560 * 640); + for (;;) { + const int e = next_e.fetch_add(1); + if (e >= NE) break; + const uint8_t* src_e = in_buf.data() + (size_t) e * in_step; + uint8_t* dst_e = out_buf.data() + (size_t) e * out_exp_bytes; + + // gate + if (tg == ref::Q2_0) std::memcpy(dst_e, src_e, Q2_SLICE_BYTES); + else { ref::dequant(tg, src_e, f32.data(), 2560 * 640); quantize_q2_0_slice(f32.data(), 2560 * 640, dst_e); } + + // up + if (tu == ref::Q2_0) std::memcpy(dst_e + Q2_SLICE_BYTES, src_e + gs, Q2_SLICE_BYTES); + else { ref::dequant(tu, src_e + gs, f32.data(), 2560 * 640); quantize_q2_0_slice(f32.data(), 2560 * 640, dst_e + Q2_SLICE_BYTES); } + + // down + if (td == ref::Q2_0) std::memcpy(dst_e + 2 * Q2_SLICE_BYTES, src_e + gs + us, Q2_SLICE_BYTES); + else { ref::dequant(td, src_e + gs + us, f32.data(), 2560 * 640); quantize_q2_0_slice(f32.data(), 2560 * 640, dst_e + 2 * Q2_SLICE_BYTES); } + } + }); + } + for (auto& w : workers) w.join(); + + if (::pwrite(efd, out_buf.data(), out_buf.size(), (off_t) out_offs[l]) != (ssize_t) out_buf.size()) { + std::perror("pwrite layer"); + return 1; + } + if ((l + 1) % 8 == 0 || l == NL - 1) { + ::fsync(efd); + const double el = std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); + std::printf(" Converted layer %d/%d to Q2_0 (%.1fs elapsed)\n", l + 1, NL, el); + std::fflush(stdout); + } + } + ::ftruncate(efd, (off_t) at); + ::fsync(efd); + ::close(efd); + ::close(ifd); + const double tot_s = std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); + std::printf("Done! Wrote %s (%.2f GiB) in %.1fs\n", out_exp.c_str(), at / (double) (1ull << 30), tot_s); + return 0; +} diff --git a/metal/src/engine/arena_budget.h b/metal/src/engine/arena_budget.h index 4cc024b7..513a944d 100644 --- a/metal/src/engine/arena_budget.h +++ b/metal/src/engine/arena_budget.h @@ -44,6 +44,9 @@ struct ArenaMem { double keep_gb = -1; // STRATA_METAL_KEEP_FREE_GB (< 0: unset) double arena_gb = -1; // STRATA_METAL_ARENA_GB (< 0: unset) double leave_share = -1; // STRATA_METAL_LEAVE_SHARE: the share of the RAM a big Mac leaves to macOS (< 0: kLeaveShare) + double dense_file = 0; // the dense weights the GPU reads straight from the GGUF's map and that are NOT wired: they are + // clean file cache, so `available` counts them as free - an arena that takes them makes + // every command buffer page them back in (STRATA_METAL_MLOCK_DENSE=0) }; constexpr double kGB = 1e9; @@ -112,7 +115,7 @@ inline double arena_ple_reserve(const ArenaMem& m) { return m.ple_mapped > 0 ? m inline double arena_budget(const ArenaMem& m) { if (m.arena_gb >= 0) return m.arena_gb * kGB; const double gpu = m.working_set - footprint_base(m) - 1.0 * kGB; - const double ram = m.available - arena_keep(m) - arena_ple_reserve(m); + const double ram = m.available - arena_keep(m) - arena_ple_reserve(m) - m.dense_file; const double foot = footprint_cap(m) - footprint_base(m) - arena_ple_reserve(m); // the whole footprint under the cap return std::max(1.0 * kGB, std::min(std::min(gpu, ram), foot)); } @@ -154,7 +157,7 @@ inline double arena_shrink(const ArenaMem& m, double planned, double available_n /// (the explicit STRATA_METAL_ARENA_GB always fills) inline bool arena_grow(const ArenaMem& m, double available_now, double step) { if (m.arena_gb >= 0) return true; - return available_now - step >= arena_keep(m) + arena_ple_reserve(m); + return available_now - step >= arena_keep(m) + arena_ple_reserve(m) + m.dense_file; } /// The fill as a plan, for the test: the pairs' bytes in rank order, steps of `step` bytes, `avail(filled)` the free diff --git a/metal/src/engine/msl/mtp.metal b/metal/src/engine/msl/mtp.metal index 7159330f..051dcf72 100644 --- a/metal/src/engine/msl/mtp.metal +++ b/metal/src/engine/msl/mtp.metal @@ -1,68 +1,179 @@ // metal/src/engine/msl/mtp.metal - the small steps of the MTP draft layer's input branches (include/strata/core/mtp.hpp) // around mac-kern's kernels (GEMVs, the hyper-connection read, QSA, MoE, the head): // -// mtp_rms y[r] = x[r] * rsqrt(mean(x[r]^2) + eps) * w rows of width n (w: (1 + w) folded by the host) -// mtp_xprep xprep (units.metal) of rows read with a stride: the prepared activation of one stream of rows -// mtp_combine R[t][c][d] = h[c][t][d] + e[t][d] the streams of fc_hidden plus fc_embedding +// mtp_rms_xprep weighted RMSNorm into prepared embedding / hidden rows +// mtp_xprep strided prepared rows for the hidden streams past the first MAXT +// mtp_combine_gnorm R[t][c][d] = h[t][c][d] + e[t][d], then the attention hyper-connection norm // // Prepared activations: [4 kinds][n_tok][n] then S16 [n_tok][n / 16] (units.metal's xprep layout). #include using namespace metal; -struct RmsArgs { uint n; float eps; }; +struct XpArgs { uint n; uint n_tok; uint plane_tok; uint stride; }; -// grid (rows), 256 threads -kernel void mtp_rms(device const float* x [[buffer(0)]], device const float* w [[buffer(1)]], device float* y [[buffer(2)]], - constant RmsArgs& a [[buffer(3)]], uint3 tg [[threadgroup_position_in_grid]], - ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], - ushort lane [[thread_index_in_simdgroup]]) { - threadgroup float part[8]; - device const float* xr = x + tg.x * a.n; - float ss = 0.0f; - for (uint d = tid; d < a.n; d += 256) ss = fma(xr[d], xr[d], ss); - ss = simd_sum(ss); - if (lane == 0) part[sg] = ss; - threadgroup_barrier(mem_flags::mem_threadgroup); - float tot = 0.0f; - for (int i = 0; i < 8; ++i) tot += part[i]; - const float rs = rsqrt(tot / float(a.n) + a.eps); - device float* yr = y + tg.x * a.n; - for (uint d = tid; d < a.n; d += 256) yr[d] = xr[d] * rs * w[d]; +// raw and plane are 16-float aligned: the four planes are written as float4s +static inline void mtp_prepare(thread const float* v, device float* raw, device float* sum, uint plane) { + float s0 = 0, s1 = 0; +#pragma clang loop unroll(full) + for (int j = 0; j < 16; j += 2) { s0 += v[j]; s1 += v[j + 1]; } + const float4 q2lo = float4(1.0f, 0.25f, 0.0625f, 0.015625f), + q2hi = float4(0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f), + p4 = float4(1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f); + device float4* r0 = (device float4*) raw; + device float4* r1 = (device float4*) (raw + plane); + device float4* r2 = (device float4*) (raw + 2 * plane); + device float4* r3 = (device float4*) (raw + 3 * plane); +#pragma clang loop unroll(full) + for (int j = 0; j < 4; ++j) { + const float4 x = float4(v[4 * j], v[4 * j + 1], v[4 * j + 2], v[4 * j + 3]); + r0[j] = x; + r1[j] = x * ((j & 1) ? q2hi : q2lo); + r2[j] = x * p4; + r3[j] = float4(x.x, x.y * 0.00390625f, x.z, x.w * 0.00390625f); + } + *sum = s0 + s1; } -struct XpArgs { uint n; uint n_tok; uint plane_tok; uint stride; }; - // one thread per 16 values of one row: row t at x + t * stride; out planes of plane_tok rows kernel void mtp_xprep(device const float* x [[buffer(0)]], device float* xp [[buffer(1)]], constant XpArgs& a [[buffer(2)]], uint gid [[thread_position_in_grid]]) { const uint per = a.n / 16; if (gid >= per * a.n_tok) return; - const uint t = gid / per, i = gid % per; + const uint t = per == 160u ? gid / 160u : gid / per, i = gid - t * per; device const float* src = x + t * a.stride + 16 * i; const uint plane = a.plane_tok * a.n; device float* raw = xp + t * a.n + 16 * i; float v[16]; - for (int j = 0; j < 16; ++j) v[j] = src[j]; - float s0 = 0, s1 = 0; - for (int j = 0; j < 16; j += 2) { s0 += v[j]; s1 += v[j + 1]; } - const float q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f}; - const float p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; - for (int j = 0; j < 16; ++j) { - raw[j] = v[j]; - raw[plane + j] = v[j] * q2s[j & 7]; - raw[2 * plane + j] = v[j] * p4s[j & 3]; - raw[3 * plane + j] = (j & 1) ? v[j] * 0.00390625f : v[j]; +#pragma clang loop unroll(full) + for (int j = 0; j < 4; ++j) { + const float4 s = ((device const float4*) src)[j]; + v[4 * j] = s.x; v[4 * j + 1] = s.y; v[4 * j + 2] = s.z; v[4 * j + 3] = s.w; + } + mtp_prepare(v, raw, xp + 4 * plane + t * per + i, plane); +} + +struct RmsXpArgs { uint n; uint xp_n; uint plane_tok; uint split; uint keep; float eps; }; + +// A wide hidden norm splits into its n / xp_n streams as prepared rows [t][c] (the first plane_tok of them); keep +// retains every normalized row for mtp_xprep. Unsplit: the first xp_n values of each row. +kernel void mtp_rms_xprep(device const float* x [[buffer(0)]], device const float* w [[buffer(1)]], + device float* y [[buffer(2)]], device float* xp [[buffer(3)]], + constant RmsXpArgs& a [[buffer(4)]], uint3 tg [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[8]; + device const float* xr = x + tg.x * a.n; + float ss = 0.0f; + if (a.n == 2560) { + device const float* xrp = xr + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) ss = fma(xrp[256 * i], xrp[256 * i], ss); + } else if (a.n == 10240) { + device const float* xrp = xr + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < 40; ++i) ss = fma(xrp[256 * i], xrp[256 * i], ss); + } else { + for (uint d = tid; d < a.n; d += 256) ss = fma(xr[d], xr[d], ss); + } + ss = simd_sum(ss); + if (lane == 0) part[sg] = ss; + threadgroup_barrier(mem_flags::mem_threadgroup); + float tot = 0.0f; +#pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) tot += part[i]; + const float rs = rsqrt(tot / float(a.n) + a.eps); + const uint per = a.xp_n / 16, plane = a.plane_tok * a.xp_n; + auto scale16 = [&](uint i, uint n_row, thread float* v) { + device const float4* x4 = (device const float4*) (xr + 16 * i); + device const float4* w4 = (device const float4*) (w + 16 * i); +#pragma clang loop unroll(full) + for (int j = 0; j < 4; ++j) { + const float4 s = x4[j] * rs * w4[j]; + v[4 * j] = s.x; v[4 * j + 1] = s.y; v[4 * j + 2] = s.z; v[4 * j + 3] = s.w; + if (a.keep) ((device float4*) (y + tg.x * n_row + 16 * i))[j] = s; + } + }; + if (a.n == 2560 && a.xp_n == 2560) { + if (tid < 160) { + float v[16]; + scale16(tid, 2560, v); + if (!a.split || tg.x < a.plane_tok) + mtp_prepare(v, xp + tg.x * 2560 + 16 * tid, xp + 4 * plane + tg.x * 160 + tid, plane); + } + } else if (a.n == 10240 && a.xp_n == 2560) { + for (uint i = tid; i < 640; i += 256) { + float v[16]; + scale16(i, 10240, v); + const uint q = i / 160u, k = i - q * 160u; + const uint t = a.split ? tg.x * 4u + q : tg.x; + if (a.split ? t >= a.plane_tok : q > 0u) continue; + mtp_prepare(v, xp + t * 2560 + 16 * k, xp + 4 * plane + t * 160 + k, plane); + } + } else { + for (uint i = tid; i < a.n / 16; i += 256) { + float v[16]; + scale16(i, a.n, v); + const uint t = a.split ? tg.x * (a.n / a.xp_n) + i / per : tg.x; + if (a.split ? t >= a.plane_tok : i >= per) continue; + const uint k = i % per; + mtp_prepare(v, xp + t * a.xp_n + 16 * k, xp + 4 * plane + t * per + k, plane); + } } - xp[4 * plane + t * per + i] = s0 + s1; } -struct CombArgs { uint n; uint hc; uint n_tok; uint h_stride; }; // h: stream c's rows at h + c * h_stride +struct CombNormArgs { uint n; uint hc; float eps; }; -kernel void mtp_combine(device const float* h [[buffer(0)]], device const float* e [[buffer(1)]], device float* R [[buffer(2)]], - constant CombArgs& a [[buffer(3)]], uint gid [[thread_position_in_grid]]) { - const uint per = a.hc * a.n; - if (gid >= a.n_tok * per) return; - const uint t = gid / per, c = (gid % per) / a.n, d = gid % a.n; - R[gid] = h[c * a.h_stride + t * a.n + d] + e[t * a.n + d]; +// Same reduction and multiply order as gr_norm(apply=false); keep R for subsequent residual writes. h, R and xn are +// [t][c][n]. +kernel void mtp_combine_gnorm(device const float* h [[buffer(0)]], device const float* e [[buffer(1)]], + device float* R [[buffer(2)]], device const float* w_norm [[buffer(3)]], + device float* xn [[buffer(4)]], constant CombNormArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[8]; + const uint c = tg.x, t = tg.y, n = a.n; + const uint row = (t * a.hc + c) * n; + device float* r = R + row; + float values[10]; // N = 2560: ten values per thread, retained across the reduction. + float ss = 0.0f; + if (n == 2560) { + device const float* hp = h + row + tid; + device const float* ep = e + t * 2560 + tid; + device float* rp = r + tid; +#pragma clang loop unroll(full) + for (int i = 0; i < 10; ++i) { + const float v = hp[256 * i] + ep[256 * i]; + rp[256 * i] = v; + values[i] = v; + ss = fma(v, v, ss); + } + } else { + for (uint d = tid; d < n; d += 256) { + const float v = h[row + d] + e[t * n + d]; + r[d] = v; + values[d / 256] = v; + ss = fma(v, v, ss); + } + } + ss = simd_sum(ss); + if (lane == 0) part[sg] = ss; + threadgroup_barrier(mem_flags::mem_threadgroup); + float tot = 0.0f; +#pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) tot += part[i]; + const float rs = rsqrt(tot / float(n) + a.eps); + device float* o = xn + row; + device const float* w = w_norm + c * n; + if (n == 2560) { + o += tid; + w += tid; +#pragma clang loop unroll(full) + for (int i = 0; i < 10; ++i) o[256 * i] = values[i] * w[256 * i] * rs; + } else { + for (uint d = tid; d < n; d += 256) { + const float v = values[d / 256]; + o[d] = v * w[d] * rs; + } + } } diff --git a/metal/src/engine/msl/mtp_chain.metal b/metal/src/engine/msl/mtp_chain.metal index 62500df2..efee47ee 100644 --- a/metal/src/engine/msl/mtp_chain.metal +++ b/metal/src/engine/msl/mtp_chain.metal @@ -13,6 +13,255 @@ // a 16-wide unit has one sum: m1) // and its position in the row is what X::load reads when x holds the positions 0 .. n-1. +struct DraftSlopArgs { + slop4::State state; + slop4::Knobs knobs; + uint nv, j, vocab; +}; + +constant int32_t kSlRather[] = {4598, 24944, 61745, 70714}, kSlAtherPre[] = {2515, 2697, 15656, 43487}; +constant int32_t kSlRight[] = {1245, 1246, 5791, 9785, 26806, 51140}; +constant int32_t kSlEndPre[] = {400, 809, 3601, 3839, 14637, 55293}; +constant int32_t kSlShrt[] = {2716, 8412, 10382, 12115, 36639}; +constant int32_t kSlExactly[] = {6681, 63151, 66172}, kSlExPre[] = {325, 5291}; +constant int32_t kSlHere[] = {1532, 5514, 6527, 8160, 475}, kSlWhat[] = {1092, 3437, 3710, 12195, 29348}; +constant int32_t kSlDet[] = {1719, 760, 279, 561, 276, 449, 1473, 2014, 2329, 821, 2921, 4888, 1323, 381, + 1602, 32985, 1379, 3444, 7261, 6183, 777, 4254, 678}; +constant int32_t kSlClause[] = {13, 11, 0, 30, 25, 26, 1076, 2218, 1892, 12, 471, 8, + 321, 694, 748, 429, 7834, 684, 1541, 1921, 1987, 2939, 3850, 4272}; +constant int32_t kSlH[] = {71, 39, 3101, 75833}; +constant int32_t kSlSentence[] = {13, 0, 30, 25, 26, 1076, 2218, 1892, 8}; +constant int32_t kSlCandor[] = {1865, 47045, 24969, 5507, 17290, 7435, 62179, 66842, 86775, + 12116, 6379, 6999, 13542, 11040, 73018}; +constant int32_t kSlBePre[] = {310, 951, 3172, 728, 1066, 16032}; +constant int32_t kSlIntens[] = {12116, 6379, 6999, 13542, 11040, 73018, 66842}; +constant int32_t kSlFrameAdj[] = {10429, 1865, 47045, 24969, 5507, 7435, 2716, 4145, 3841, 26860, 13983, + 8412, 12115, 10382, 7700, 12402, 8535, 4880, 32254}; +constant int32_t kSlFrameDet[] = {1719, 760, 279, 561, 2329, 821, 2921, 4888, 276, 449}; +constant int32_t kSlFrame[] = {4087, 1831, 56004, 2243, 7793, 1301, 11782, 14996, 3296}; +constant int32_t kSlThe[] = {1719, 760, 279, 561}, kSlTo[] = {1206, 1946, 951, 310}; +constant int32_t kSlCandorBranch[] = {1865, 47045, 24969, 5507, 7435, 2716, 26860}; +constant int32_t kSlLineFrame[] = {8412, 12115, 7700, 12402}; +constant int32_t kSlMd[] = {12, 9, 6458, 471, 348, 7061, 332, 2972, 2, 550, 13962, 794, 653, 7460, 16112, + 4277, 11862, 4006, 2218, 10, 49300}; + +template static inline bool sl_in(int32_t t, constant const int32_t (&s)[N]) { + bool found = false; +#pragma clang loop unroll(full) + for (size_t i = 0; i < N; ++i) found |= (s[i] == t); + return found; +} + +static inline int mtp_slop_penalties(slop4::State s, uchar c_p0, uchar c_p1, slop4::Knobs k, + threadgroup int* ids, threadgroup float* vals) { + if ((s.f & (slop4::KNOWN | slop4::PROSE | slop4::TOOL | slop4::FENCE)) != (slop4::KNOWN | slop4::PROSE)) return 0; + int n = 0; + auto pen = [&](int32_t id, float v) { + if (v > 0.0f && n < 64) { + ids[n] = id; + vals[n] = v; + ++n; + } + }; + const int32_t p0 = s.t1, p1 = s.t2; + if (sl_in(p0, kSlRather) || (p0 == 1792 && sl_in(p1, kSlAtherPre))) pen(1056, k.lex); + if (sl_in(p0, kSlRight) || (p0 == 481 && p1 == 2697)) pen(3023, k.lex); + if (p0 == 310 && sl_in(p1, kSlEndPre)) pen(809, k.lex); + if (sl_in(p0, kSlShrt)) pen(2243, k.lex); + if (sl_in(p0, kSlExactly) || (p0 == 31654 && sl_in(p1, kSlExPre))) pen(1092, k.lex); + if (p0 == 579 && sl_in(p1, kSlHere)) pen(1092, k.lex); + if (p0 == 579 && sl_in(p1, kSlWhat)) pen(3404, k.lex); + const bool line_start = c_p0 == 1 || (s.f & slop4::BOL); + const bool sentence_start = sl_in(p0, kSlSentence) || line_start; + const bool clause_start = sl_in(p0, kSlClause) || line_start; + pen(92115, k.hon); + pen(68540, k.hon); + pen(51397, k.hon); + if (sl_in(p0, kSlDet) || sentence_start) pen(10429, k.hon); + if (clause_start) pen(25211, k.hon); + if (sl_in(p0, kSlH)) { pen(33961, 1.5f * k.hon); pen(77696, 1.5f * k.hon); } + if (p0 == 660) { pen(46253, k.hon); pen(25211, k.hon); } + if (p0 == 1602 || p0 == 32985) { +#pragma clang loop unroll(full) + for (int32_t id : kSlCandor) pen(id, k.hon); + } + if ((p0 == 381 || p0 == 1323) && sl_in(p1, kSlBePre)) { +#pragma clang loop unroll(full) + for (int32_t id : kSlCandor) pen(id, k.hon); + pen(2050, k.hon); + if (p1 == 310 || p1 == 951) { pen(6415, k.hon); pen(2708, k.hon); } + } + if (sl_in(p0, kSlIntens) && (p1 == 1602 || p1 == 32985 || p1 == 381 || p1 == 1323)) { + pen(10429, k.hon); +#pragma clang loop unroll(full) + for (int32_t id : kSlCandor) pen(id, k.hon); + } + if (p0 == 2688 && (p1 == 353 || p1 == 585 || p1 == 40 || p1 == 72)) { pen(1602, 0.5f * k.hon); pen(10429, k.hon); } + if (clause_start) pen(40285, k.hon); + if (p0 == 660) { pen(49034, k.hon); pen(71853, k.hon); } + if (sl_in(p0, kSlFrameAdj) && (sl_in(p1, kSlFrameDet) || sl_in(p1, kSlSentence) || c_p1 == 1 || p1 == slop4::ETHINK)) { +#pragma clang loop unroll(full) + for (int32_t id : kSlFrame) pen(id, k.hon); + } + const bool p1_root = sl_in(p1, kSlSentence) || c_p1 == 1 || p1 == slop4::ETHINK || p1 < 0; + if (sl_in(p0, kSlThe) && (p1_root || (c_p0 == 0 && (s.f & slop4::PBOL)))) { +#pragma clang loop unroll(full) + for (int32_t id : kSlCandorBranch) pen(id, 0.5f * k.hon); + } + if (sl_in(p0, kSlTo) && (p1_root || (s.f & slop4::PBOL))) pen(381, 0.5f * k.hon); + if (s.f & slop4::BOL) { +#pragma clang loop unroll(full) + for (int32_t id : kSlLineFrame) pen(id, 0.5f * k.hon); + const float v = k.md + ((s.f & slop4::COLON) ? k.colon : 0.0f); +#pragma clang loop unroll(full) + for (int32_t id : kSlMd) pen(id, v); +#pragma clang loop unroll(full) + for (int32_t id : slop4::kDigit) pen(id, (s.f & slop4::COLON) ? v : 0.25f * v); + } + if ((s.f & slop4::PBOL) && sl_in(p0, slop4::kDigit)) { pen(13, k.md + k.colon); pen(8, k.md + k.colon); } + if (sl_in(p0, slop4::kColon) && s.len <= 6) { pen(slop4::NL, k.colon); pen(slop4::NL2, k.colon); } + return n; +} +// one threadgroup; the state advances only along the draft branch, never the verifier's history. +kernel void mtp_slop_pen(device float* lg [[buffer(0)]], device const int* inverse [[buffer(1)]], + device const int* dout [[buffer(2)]], device slop4::State* state [[buffer(3)]], + constant DraftSlopArgs& a [[buffer(4)]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup int ids[80]; + threadgroup float vals[80], rem[80], part[8]; + threadgroup int hold_ids[4]; + threadgroup int count; + threadgroup slop4::State st; + threadgroup float mx, total; + threadgroup bool hold; + threadgroup uchar tcls[2]; + const bool any_knob = a.knobs.lex > 0.0f || a.knobs.hon > 0.0f || a.knobs.md > 0.0f || a.knobs.colon > 0.0f; + if (a.j == 0 && !any_knob) { + if (tid == 0) state[0] = a.state; + return; + } + const slop4::State s0 = a.j == 0 ? a.state : state[0]; + const int32_t tok0 = a.j == 0 ? s0.t1 : dout[a.j - 1]; + const int32_t tok1 = a.j == 0 ? s0.t2 : s0.t1; + if (tid < 2) tcls[tid] = 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + constexpr uint kNL = sizeof(slop4::kNewline) / sizeof(int32_t); + constexpr uint kBL = sizeof(slop4::kBlank) / sizeof(int32_t); + if (tid < kNL) { + const int32_t x = slop4::kNewline[tid]; + if (x == tok0) tcls[0] = 1; + if (x == tok1) tcls[1] = 1; + } + if (tid < kBL) { + const int32_t x = slop4::kBlank[tid]; + if (x == tok0) tcls[0] = 2; + if (x == tok1) tcls[1] = 2; + } + if (tid < 4) { + const int hid = slop4::kHold[tid]; + hold_ids[tid] = uint(hid) < a.vocab ? inverse[hid] : -1; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0) { + slop4::State s = s0; + const uchar c0 = tcls[0], c1 = tcls[1]; + if (a.j != 0) { + s.t2 = s.t1; + s.t1 = tok0; + s.f |= slop4::KNOWN; + const uint8_t bol = s.f & slop4::BOL; + s.f &= (uint8_t) ~slop4::PBOL; + if (tok0 == slop4::IM_START || tok0 == slop4::IM_END || tok0 == slop4::THINK) { + s.f = slop4::KNOWN | slop4::BOL; + s.len = 0; + } else if (tok0 == slop4::ETHINK) { + s.f = slop4::KNOWN | slop4::PROSE | slop4::BOL; + s.len = 0; + } else if (tok0 == slop4::TOOL_CALL) { + s.f |= slop4::TOOL; + } else if (tok0 == slop4::ETOOL_CALL) { + s.f = (uint8_t) ((s.f & ~(slop4::TOOL | slop4::COLON)) | slop4::BOL); + s.len = 0; + } else if (c0 == 1) { + if (!bol) s.f = (uint8_t) ((s.f & ~slop4::COLON) | (slop4::in(s.t2, slop4::kColon) ? slop4::COLON : 0)); + s.f |= slop4::BOL; + s.len = 0; + } else if (c0 == 2) { + } else { + if (slop4::in(tok0, slop4::kFence) && (s.f & slop4::PROSE) && !(s.f & slop4::TOOL)) s.f ^= slop4::FENCE; + s.f = (uint8_t) ((s.f & ~(slop4::BOL | slop4::COLON)) | (bol ? slop4::PBOL : 0)); + if (s.len < 255) ++s.len; + } + } + state[0] = s; + st = s; + const slop4::Knobs kn = a.knobs; + count = (kn.lex > 0.0f || kn.hon > 0.0f || kn.md > 0.0f || kn.colon > 0.0f) + ? mtp_slop_penalties(s, c0, c1, kn, ids, vals) + : 0; + hold = false; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (count > 0) { + if (int(tid) < count) { + const int raw_id = ids[tid]; + ids[tid] = raw_id >= 0 && uint(raw_id) < a.vocab ? inverse[raw_id] : -1; + } + float m = -INFINITY; +#pragma clang loop unroll_count(4) + for (uint i = tid; i < a.nv; i += 256) if (!isnan(lg[i])) m = max(m, lg[i]); + m = simd_max(m); + if (lane == 0) part[sg] = m; + threadgroup_barrier(mem_flags::mem_threadgroup); + float lmx = -INFINITY; +#pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) lmx = max(lmx, part[k]); + if (tid == 0) mx = lmx; + if (int(tid) < count) { + const int id = ids[tid]; + if (id >= 0) { + const float lgi = lg[id]; + if (lgi > lmx - 12.0f) hold = true; + rem[tid] = exp(lgi - lmx) * (1.0f - exp(-vals[tid])); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (hold) { + float sum = 0.0f; +#pragma clang loop unroll_count(4) + for (uint i = tid; i < a.nv; i += 256) sum += exp(lg[i] - mx); + sum = simd_sum(sum); + if (lane == 0) part[sg] = sum; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0) { + total = 0.0f; +#pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) total += part[k]; + float removed = 0.0f; + for (int i = 0; i < count; ++i) + if (ids[i] >= 0) removed += rem[i]; + const float v = -log(1.0f - min(removed / total, 0.999999f)); +#pragma clang loop unroll(full) + for (int k = 0; k < 4; ++k) { + ids[count] = hold_ids[k]; + vals[count++] = v; + } + } + } + if (tid == 0) { + if (a.knobs.lex > 0.0f && st.t2 == slop4::ETHINK && (st.t1 == slop4::NL || st.t1 == slop4::NL2)) { + const int empty[] = {slop4::IM_END, slop4::THINK}; + for (int id : empty) { + ids[count] = uint(id) < a.vocab ? inverse[id] : -1; + vals[count++] = a.knobs.lex; + } + } + // serial subtraction preserves additive penalties when two rules name the same token. + for (int i = 0; i < count; ++i) if (ids[i] >= 0) lg[ids[i]] -= vals[i]; + } +} + struct PickArgs { uint nv; // the draft head's rows (the draft vocabulary, or the whole vocabulary) uint j; // the draft step @@ -28,12 +277,14 @@ kernel void mtp_pick(device const float* lg [[buffer(0)]], device const int* pic const int di = pick[0]; const float mx = lg[di]; float s = 0.0f; +#pragma clang loop unroll_count(4) for (uint i = tid; i < a.nv; i += 1024) s += exp(lg[i] - mx); s = simd_sum(s); if (lane == 0) part[sg] = s; threadgroup_barrier(mem_flags::mem_threadgroup); if (tid == 0) { float t = 0.0f; +#pragma clang loop unroll(full) for (int k = 0; k < 32; ++k) t += part[k]; dout[a.j] = a.use_map ? vmap[di] : di; dprob[a.j] = 1.0f / t; @@ -60,28 +311,54 @@ template static inline float kind_scale(int pos) { return 1.0f; } +template static inline void mtp_embed_unit(device const uchar* row, int u, device float* y, + threadgroup const uint* tab) { + constexpr int UNIT = fmt::UNIT; + float q[UNIT]; + float4 sc; + Dec::run(row, u, q, sc, tab); + const int p0 = (F == T_Q4_K || F == T_Q5_K) ? (256 * (u >> 3) + 64 * ((u >> 1) & 3) + 16 * (u & 1)) + : (F == T_Q6_K) ? (256 * (u >> 3) + 128 * ((u >> 2) & 1) + 32 * (u & 1) + 16 * ((u >> 1) & 1)) + : (UNIT * u); + device float4* y0 = (device float4*)(y + p0); +#pragma clang loop unroll(full) + for (int j = 0; j < 4; ++j) { + y0[j] = float4(sc.x * (q[4 * j + 0] * kind_scale(4 * j + 0)) - sc.z, + sc.x * (q[4 * j + 1] * kind_scale(4 * j + 1)) - sc.z, + sc.x * (q[4 * j + 2] * kind_scale(4 * j + 2)) - sc.z, + sc.x * (q[4 * j + 3] * kind_scale(4 * j + 3)) - sc.z); + } + if (UNIT == 32) { + const int p1 = p0 + ((F == T_Q4_K || F == T_Q5_K) ? 32 : (F == T_Q6_K ? 64 : 16)); + device float4* y1 = (device float4*)(y + p1); +#pragma clang loop unroll(full) + for (int j = 0; j < 4; ++j) { + y1[j] = float4(sc.y * (q[16 + 4 * j + 0] * kind_scale(4 * j + 0)) - sc.w, + sc.y * (q[16 + 4 * j + 1] * kind_scale(4 * j + 1)) - sc.w, + sc.y * (q[16 + 4 * j + 2] * kind_scale(4 * j + 2)) - sc.w, + sc.y * (q[16 + 4 * j + 3] * kind_scale(4 * j + 3)) - sc.w); + } + } +} + // one threadgroup of 256; posv = the floats 0 .. n-1 template kernel void mtp_embed(device const uchar* W [[buffer(0)]], device const int* dout [[buffer(1)]], device const float* posv [[buffer(2)]], device float* y [[buffer(3)]], constant EmbArgs& a [[buffer(4)]], ushort tid [[thread_index_in_threadgroup]]) { - constexpr int UNIT = fmt::UNIT, H = UNIT / 2; + (void)posv; + constexpr int UNIT = fmt::UNIT; threadgroup uint tab[fmt::TAB]; - load_tab(tab, tid, 256); - threadgroup_barrier(mem_flags::mem_threadgroup); + if (fmt::TAB > 1) { + load_tab(tab, tid, 256); + threadgroup_barrier(mem_flags::mem_threadgroup); + } device const uchar* row = W + (ulong) dout[a.j] * a.row_bytes; const int U = int(a.n) / UNIT; - for (int u = tid; u < U; u += 256) { - float q[UNIT], ps[UNIT], xsum[2]; - float4 sc; - Dec::run(row, u, q, sc, tab); - X::load(posv, posv, u, ps, xsum); - for (int j = 0; j < UNIT; ++j) { - const int pos = int(ps[j]); - const float d = j < H ? sc.x : sc.y; - const float m = (UNIT == 16 || j < 16) ? sc.z : sc.w; - y[pos] = d * (q[j] * kind_scale(pos)) - m; - } + if (U <= 256) { + if (int(tid) < U) mtp_embed_unit(row, int(tid), y, tab); + } else { + for (int u = tid; u < U; u += 256) mtp_embed_unit(row, u, y, tab); } } diff --git a/metal/src/engine/msl/ple.metal b/metal/src/engine/msl/ple.metal index 7036eaeb..d323f0ea 100644 --- a/metal/src/engine/msl/ple.metal +++ b/metal/src/engine/msl/ple.metal @@ -9,43 +9,32 @@ // R += gated + conv // // History: hist[row + 9 * channel] (row-fastest, oldest first, the committed rows); the window's normalized rows -// are written to norm_w[t][channel] (the host appends the kept ones at commit). +// are written to norm_w[t][channel] (the window's host appends the kept ones at commit; a prompt chunk's ple_hist_adv). #include using namespace metal; +#ifndef P_TOKENT +#define P_TOKENT 532 // moe.metal's plan layout +#endif + struct PleArgs { uint n_embd; // 2560 uint hc; // 4 float eps; }; +// One barrier per reduction; callers with back-to-back reductions alternate part[0..7] and part[8..15]. static inline float tg_sum(float v, threadgroup float* part, ushort sg, ushort lane) { v = simd_sum(v); if (lane == 0) part[sg] = v; threadgroup_barrier(mem_flags::mem_threadgroup); float t = 0.0f; +#pragma clang loop unroll(full) for (int i = 0; i < 8; ++i) t += part[i]; - threadgroup_barrier(mem_flags::mem_threadgroup); return t; } -// grid (hc, T), 256 threads: y[t][c] = grouped_norm(x[t][c]) with w[c] -kernel void ple_gnorm(device const float* x [[buffer(0)]], device const float* w [[buffer(1)]], device float* y [[buffer(2)]], - constant PleArgs& a [[buffer(3)]], uint3 tg [[threadgroup_position_in_grid]], - ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], - ushort lane [[thread_index_in_simdgroup]]) { - threadgroup float part[8]; - const uint c = tg.x, t = tg.y, n = a.n_embd; - device const float* xc = x + (t * a.hc + c) * n; - float ss = 0.0f; - for (uint d = tid; d < n; d += 256) ss += xc[d] * xc[d]; - const float mean = tg_sum(ss, part, sg, lane) / float(n); - const float scale = 1.0f / sqrt(mean + a.eps); - device float* yc = y + (t * a.hc + c) * n; - for (uint d = tid; d < n; d += 256) yc[d] = xc[d] * scale * w[c * n + d]; -} - // grid (hc, T), 256 threads: the gate, then gated = value * gate and normalized = grouped_norm(gated, norm_conv) kernel void ple_gate_norm(device const float* key [[buffer(0)]], device const float* query [[buffer(1)]], device const float* value [[buffer(2)]], device const float* w_conv [[buffer(3)]], @@ -53,7 +42,7 @@ kernel void ple_gate_norm(device const float* key [[buffer(0)]], device const fl constant PleArgs& a [[buffer(6)]], uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { - threadgroup float part[8]; + threadgroup float part[16]; const uint c = tg.x, t = tg.y, n = a.n_embd, row = (t * a.hc + c) * n; float acc = 0.0f; for (uint d = tid; d < n; d += 256) acc += key[row + d] * query[row + d]; @@ -67,25 +56,381 @@ kernel void ple_gate_norm(device const float* key [[buffer(0)]], device const fl gated[row + d] = g; ss += g * g; } - const float mean = tg_sum(ss, part, sg, lane) / float(n); + const float mean = tg_sum(ss, part + 8, sg, lane) / float(n); const float scale = 1.0f / sqrt(mean + a.eps); for (uint d = tid; d < n; d += 256) norm_w[row + d] = gated[row + d] * scale * w_conv[c * n + d]; } -// grid (hc * n_embd / 256, T), 256 threads: the conv over [history | the window's rows], SiLU, and the residual +// The fused forms keep a thread's NPT values of a row in registers (n_embd <= 256 * NPT), so the key and query norms +// and gated never reach memory; each is bitwise the ple_gnorm / ple_gate_norm / ple_conv_add chain it replaces. +constant constexpr uint NPT = 10; + +// y = grouped_norm(y) in place, ple_gnorm's arithmetic +static inline void gnorm_regs(thread float* y, device const float* w, uint n, uint tid, float eps, + threadgroup float* part, ushort sg, ushort lane) { + float ss = 0.0f; + if (n == 2560) { +#pragma clang loop unroll(full) + for (uint i = 0; i < NPT; ++i) ss += y[i] * y[i]; + } else { + for (uint i = 0, d = tid; d < n; ++i, d += 256) ss += y[i] * y[i]; + } + const float mean = tg_sum(ss, part, sg, lane) / float(n); + const float scale = 1.0f / sqrt(mean + eps); + if (n == 2560) { + device const float* wt = w + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < NPT; ++i) y[i] = y[i] * scale * wt[i * 256]; + } else { + for (uint i = 0, d = tid; d < n; ++i, d += 256) y[i] = y[i] * scale * w[d]; + } +} + +// grid (hc, T), 256 threads: y[t][c] = grouped_norm(x[t][c]) with w[c]; a row in registers when it fits (x read once) +kernel void ple_gnorm(device const float* x [[buffer(0)]], device const float* w [[buffer(1)]], device float* y [[buffer(2)]], + constant PleArgs& a [[buffer(3)]], uint3 tg [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[8]; + const uint c = tg.x, t = tg.y, n = a.n_embd; + device const float* xc = x + (t * a.hc + c) * n; + device float* yc = y + (t * a.hc + c) * n; + if (n == 2560) { + float xv[NPT]; + device const float* xt = xc + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < NPT; ++i) xv[i] = xt[i * 256]; + gnorm_regs(xv, w + c * n, n, tid, a.eps, part, sg, lane); + device float* yt = yc + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < NPT; ++i) yt[i * 256] = xv[i]; + return; + } + float ss = 0.0f; + for (uint d = tid; d < n; d += 256) ss += xc[d] * xc[d]; + const float mean = tg_sum(ss, part, sg, lane) / float(n); + const float scale = 1.0f / sqrt(mean + a.eps); + for (uint d = tid; d < n; d += 256) yc[d] = xc[d] * scale * w[c * n + d]; +} + +// ple_gate_norm on register rows: the gate from key . query, then normalized written and the gated rows added to the +// residual in the same pass: R = rb + gated, which ple_conv_res completes to ple_conv_add's (R + gated) + conv, so +// gated never reaches memory. part has 16 floats: the prior gnorm_regs left part[0..7] live, so the first reduction +// uses part + 8 and the second reuses part[0..7]. +static inline void gate_norm_regs(thread const float* key, thread const float* query, thread const float* value, + device const float* w_conv, device float* R, thread const float* rb, + device float* norm_w, uint n, uint tid, float eps, threadgroup float* part, ushort sg, + ushort lane) { + float acc = 0.0f; + if (n == 2560) { +#pragma clang loop unroll(full) + for (uint i = 0; i < NPT; ++i) acc += key[i] * query[i]; + } else { + for (uint i = 0, d = tid; d < n; ++i, d += 256) acc += key[i] * query[i]; + } + const float s = tg_sum(acc, part + 8, sg, lane) * (1.0f / sqrt(float(n))); + const float mag = sqrt(max(fabs(s), 1e-6f)); + const float sgn = s > 0.0f ? 1.0f : (s < 0.0f ? -1.0f : 0.0f); + const float gate = 1.0f / (1.0f + exp(-(sgn * mag))); + float gv[NPT]; + float ss = 0.0f; + if (n == 2560) { + device float* rt = R + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < NPT; ++i) { + gv[i] = value[i] * gate; + rt[i * 256] = rb[i] + gv[i]; + ss += gv[i] * gv[i]; + } + } else { + for (uint i = 0, d = tid; d < n; ++i, d += 256) { + gv[i] = value[i] * gate; + R[d] = rb[i] + gv[i]; + ss += gv[i] * gv[i]; + } + } + const float mean = tg_sum(ss, part, sg, lane) / float(n); + const float scale = 1.0f / sqrt(mean + eps); + if (n == 2560) { + device float* nwt = norm_w + tid; + device const float* wct = w_conv + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < NPT; ++i) nwt[i * 256] = gv[i] * scale * wct[i * 256]; + } else { + for (uint i = 0, d = tid; d < n; ++i, d += 256) norm_w[d] = gv[i] * scale * w_conv[d]; + } +} + +// grid (hc, T), 256 threads: ple_gnorm(R, w_query) + ple_gate_norm with the key norm already in key (the window: +// its key norm runs beside layer 0's first read); R += gated, ple_conv_res adds the conv +kernel void ple_query_gate_norm(device const float* key [[buffer(0)]], device float* R [[buffer(1)]], + device const float* w_query [[buffer(2)]], device const float* value [[buffer(3)]], + device const float* w_conv [[buffer(4)]], device float* norm_w [[buffer(5)]], + constant PleArgs& a [[buffer(6)]], uint3 tg [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[16]; + const uint c = tg.x, t = tg.y, n = a.n_embd, row = (t * a.hc + c) * n; + float k[NPT], q[NPT], rb[NPT], v[NPT]; + if (n == 2560) { + device const float* kt = key + row + tid; + device const float* rt = R + row + tid; + device const float* vt = value + t * 2560 + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < NPT; ++i) { + k[i] = kt[i * 256]; + q[i] = rb[i] = rt[i * 256]; + v[i] = vt[i * 256]; + } + } else { + for (uint i = 0, d = tid; d < n; ++i, d += 256) { + k[i] = key[row + d]; + q[i] = rb[i] = R[row + d]; + v[i] = value[t * n + d]; + } + } + gnorm_regs(q, w_query + c * n, n, tid, a.eps, part, sg, lane); + gate_norm_regs(k, q, v, w_conv + c * n, R + row, rb, norm_w + row, n, tid, a.eps, part, sg, lane); +} + +// grid (hc, T), 256 threads: ple_query_gate_norm with layer 0's MoE write formed and applied as R is read, gr_norm_comb's +// R' = R + (sum_k ye[entry(t, k)] (k in order) + sigmoid(sg) ysh) * 2 sigmoid(inj / hc): bitwise gr_norm_comb + +// ple_query_gate_norm, with R written once and none of gr_norm's xn (nothing reads it: layer 1's read norms R again +// after the conv) +kernel void ple_comb_query_gate_norm(device const float* key [[buffer(0)]], device float* R [[buffer(1)]], + device const float* w_query [[buffer(2)]], device const float* value [[buffer(3)]], + device const float* w_conv [[buffer(4)]], device float* norm_w [[buffer(5)]], + constant PleArgs& a [[buffer(6)]], device const float* ye [[buffer(7)]], + device const float* inj [[buffer(8)]], device const int* plan [[buffer(9)]], + device const float* ysh [[buffer(10)]], device const float* sgv [[buffer(11)]], + constant uint& topk [[buffer(12)]], uint3 tg [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[16]; + const uint c = tg.x, t = tg.y, n = a.n_embd, row = (t * a.hc + c) * n; + const float g = 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))); + const float gs = 1.0f / (1.0f + exp(-sgv[t])); + float k[NPT], q[NPT], rb[NPT], v[NPT]; + if (n == 2560 && topk == 10) { + device const float* yep[10]; +#pragma clang loop unroll(full) + for (uint j = 0; j < 10; ++j) yep[j] = ye + uint(plan[P_TOKENT + t * 10 + j]) * 2560u + tid; + device const float* yshp = ysh + t * 2560u + tid; + device const float* kt = key + row + tid; + device const float* rt = R + row + tid; + device const float* vt = value + t * 2560 + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < NPT; ++i) { + float s = 0.0f; +#pragma clang loop unroll(full) + for (uint j = 0; j < 10; ++j) s += yep[j][i * 256]; + q[i] = rb[i] = fma(fma(gs, yshp[i * 256], s), g, rt[i * 256]); + k[i] = kt[i * 256]; + v[i] = vt[i * 256]; + } + } else { + for (uint i = 0, d = tid; d < n; ++i, d += 256) { + float s = 0.0f; + for (uint j = 0; j < topk; ++j) s += ye[plan[P_TOKENT + t * topk + j] * int(n) + d]; + q[i] = rb[i] = fma(fma(gs, ysh[t * n + d], s), g, R[row + d]); + k[i] = key[row + d]; + v[i] = value[t * n + d]; + } + } + gnorm_regs(q, w_query + c * n, n, tid, a.eps, part, sg, lane); + gate_norm_regs(k, q, v, w_conv + c * n, R + row, rb, norm_w + row, n, tid, a.eps, part, sg, lane); +} + +// stream c of token t from the residual rows rb (layer 0's MoE write applied): the key and query norms, the gate, +// R = rb + gated and the normalized rows. part has 16 floats: gnorm_regs(k) uses part + 8, gnorm_regs(q) uses +// part[0..7], and gate_norm_regs uses part + 8 then part[0..7]. +static inline void moe_gate_norm_row(device float* R, thread const float* rb, device const float* pkey, + device const float* w_key, device const float* w_query, thread const float* v, + device const float* w_conv, device float* norm_w, uint row, uint c, uint n, + uint tid, float eps, threadgroup float* part, ushort sg, ushort lane) { + float k[NPT], q[NPT]; + if (n == 2560) { + device const float* pkt = pkey + row + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < NPT; ++i) { + q[i] = rb[i]; + k[i] = pkt[i * 256]; + } + } else { + for (uint i = 0, d = tid; d < n; ++i, d += 256) { + q[i] = rb[i]; + k[i] = pkey[row + d]; + } + } + gnorm_regs(k, w_key + c * n, n, tid, eps, part + 8, sg, lane); + gnorm_regs(q, w_query + c * n, n, tid, eps, part, sg, lane); + gate_norm_regs(k, q, v, w_conv + c * n, R + row, rb, norm_w + row, n, tid, eps, part, sg, lane); +} + +// grid (hc, T), 256 threads: a prompt chunk's whole PLE norm stage - layer 0's MoE write applied to R (moe_out + +// gr_norm's R' = R + fma(sigmoid(sg), ys, yr) * 2 sigmoid(inj / hc)), the key and query norms, the gate, the +// normalized rows and R' + gated written back; ple_conv_res adds the conv +kernel void ple_moe_gate_norm(device float* R [[buffer(0)]], device const float* yr [[buffer(1)]], + device const float* ys [[buffer(2)]], device const float* sgv [[buffer(3)]], + device const float* inj [[buffer(4)]], device const float* pkey [[buffer(5)]], + device const float* w_key [[buffer(6)]], device const float* w_query [[buffer(7)]], + device const float* value [[buffer(8)]], device const float* w_conv [[buffer(9)]], + device float* norm_w [[buffer(10)]], constant PleArgs& a [[buffer(11)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[16]; + const uint c = tg.x, t = tg.y, n = a.n_embd, row = (t * a.hc + c) * n; + const float g = 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))); + const float gs = 1.0f / (1.0f + exp(-sgv[t])); + float rb[NPT], v[NPT]; + if (n == 2560) { + device const float* yst = ys + t * 2560 + tid; + device const float* yrt = yr + t * 2560 + tid; + device const float* rt = R + row + tid; + device const float* vt = value + t * 2560 + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < NPT; ++i) { + rb[i] = fma(fma(gs, yst[i * 256], yrt[i * 256]), g, rt[i * 256]); + v[i] = vt[i * 256]; + } + } else { + for (uint i = 0, d = tid; d < n; ++i, d += 256) { + rb[i] = fma(fma(gs, ys[t * n + d], yr[t * n + d]), g, R[row + d]); + v[i] = value[t * n + d]; + } + } + moe_gate_norm_row(R, rb, pkey, w_key, w_query, v, w_conv, norm_w, row, c, n, tid, a.eps, part, sg, lane); +} + +// ple_moe_gate_norm with mm_combine folded in (prompt.metal's gr_norm_h_moec): yr formed from the expert rows YE +// (pos, srt_w: mm_place's) as it is read, bitwise mm_combine + ple_moe_gate_norm. grid (1, T), 256 threads: a +// token's streams in turn, so YE, ys and value are read once a token; hc <= 4, topk <= 16 +kernel void ple_moe_gate_norm_c(device float* R [[buffer(0)]], device const float* YE [[buffer(1)]], + device const int* pos [[buffer(2)]], device const float* srt_w [[buffer(3)]], + device const float* ys [[buffer(4)]], device const float* sgv [[buffer(5)]], + device const float* inj [[buffer(6)]], device const float* pkey [[buffer(7)]], + device const float* w_key [[buffer(8)]], device const float* w_query [[buffer(9)]], + device const float* value [[buffer(10)]], device const float* w_conv [[buffer(11)]], + device float* norm_w [[buffer(12)]], constant PleArgs& a [[buffer(13)]], + constant uint& topk [[buffer(14)]], uint3 tg [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[16]; + threadgroup uint roff[17]; + threadgroup float wk[16]; + const uint t = tg.y, n = a.n_embd; + int r = -1; + float w = 0.0f; + if (tid < topk) { + r = pos[t * topk + tid]; + if (r >= 0) w = srt_w[r]; + } + const bool valid = r >= 0; + const ushort prefix = simd_prefix_exclusive_sum(ushort(valid)); + const ushort cnt_s = simd_sum(ushort(valid)); + if (valid) { + roff[prefix] = uint(r) * n; + wk[prefix] = w; + } + if (tid == 0) roff[16] = cnt_s; + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint cnt = roff[16]; + const float gs = 1.0f / (1.0f + exp(-sgv[t])); + float bo[NPT], v[NPT]; + if (n == 2560) { + device const float* yst = ys + t * 2560 + tid; + device const float* vt = value + t * 2560 + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < NPT; ++i) { + const uint d = tid + i * 256; + float s = 0.0f; // mm_combine's sum, k in order + for (uint k = 0; k < cnt; ++k) s = fma(wk[k], YE[roff[k] + d], s); + bo[i] = fma(gs, yst[i * 256], s); + v[i] = vt[i * 256]; + } + } else { + for (uint i = 0, d = tid; d < n; ++i, d += 256) { + float s = 0.0f; // mm_combine's sum, k in order + for (uint k = 0; k < cnt; ++k) s = fma(wk[k], YE[roff[k] + d], s); + bo[i] = fma(gs, ys[t * n + d], s); + v[i] = value[t * n + d]; + } + } + for (uint c = 0; c < a.hc; ++c) { + const uint row = (t * a.hc + c) * n; + const float g = 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))); + float rb[NPT]; + if (n == 2560) { + device const float* rt = R + row + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < NPT; ++i) rb[i] = fma(bo[i], g, rt[i * 256]); + } else { + for (uint i = 0, d = tid; d < n; ++i, d += 256) rb[i] = fma(bo[i], g, R[row + d]); + } + moe_gate_norm_row(R, rb, pkey, w_key, w_query, v, w_conv, norm_w, row, c, n, tid, a.eps, part, sg, lane); + } +} + +// the conv over [history | the window's rows] at (t, ch), SiLU'd +static inline float ple_conv(device const float* hist, device const float* norm_w, device const half* kw, uint ch, + uint t, uint hcd) { + const half4 w = *(device const half4*)(kw + 4 * ch); + float acc = 0.0f; + if (t >= 9) { + device const float* nw = norm_w + (t - 9) * hcd + ch; + const uint s3 = 3 * hcd; + acc += float(w.x) * nw[0]; + acc += float(w.y) * nw[s3]; + acc += float(w.z) * nw[2 * s3]; + acc += float(w.w) * nw[3 * s3]; + } else { +#pragma clang loop unroll(full) + for (int k = 0; k < 4; ++k) { + const int p = int(t) - (3 - k) * 3; // the row this tap reads: < 0 = the committed history + const float v = p >= 0 ? norm_w[uint(p) * hcd + ch] : hist[uint(9 + p) + 9 * ch]; + acc += float(w[k]) * v; + } + } + return acc / (1.0f + exp(-acc)); +} + +// grid (hc * n_embd / 256, T), 256 threads: the conv and the residual, R += gated + conv kernel void ple_conv_add(device const float* hist [[buffer(0)]], device const float* norm_w [[buffer(1)]], device const half* kw [[buffer(2)]], device const float* gated [[buffer(3)]], device float* R [[buffer(4)]], constant PleArgs& a [[buffer(5)]], uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { const uint ch = tg.x * 256 + tid, t = tg.y, hcd = a.hc * a.n_embd; if (ch >= hcd) return; - float acc = 0.0f; - for (int k = 0; k < 4; ++k) { - const int p = int(t) - (3 - k) * 3; // the row this tap reads: < 0 = the committed history - const float v = p >= 0 ? norm_w[uint(p) * hcd + ch] : hist[uint(9 + p) + 9 * ch]; - acc += float(kw[k + 4 * ch]) * v; - } - const float conv = acc / (1.0f + exp(-acc)); + const float conv = ple_conv(hist, norm_w, kw, ch, t, hcd); device float* r = R + t * hcd; r[ch] = r[ch] + gated[t * hcd + ch] + conv; } + +// ple_conv_add's grid after ple_query_gate_norm / ple_moe_gate_norm(_c), which already added gated to R +kernel void ple_conv_res(device const float* hist [[buffer(0)]], device const float* norm_w [[buffer(1)]], + device const half* kw [[buffer(2)]], device float* R [[buffer(3)]], + constant PleArgs& a [[buffer(4)]], uint3 tg [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]]) { + const uint ch = tg.x * 256 + tid, t = tg.y, hcd = a.hc * a.n_embd; + if (ch >= hcd) return; + R[t * hcd + ch] += ple_conv(hist, norm_w, kw, ch, t, hcd); +} + +// grid (hc * n_embd / 256, 1), 256 threads, after ple_conv_res: the history advanced in place past the window's T rows, +// row r of the new history = window row T - 9 + r, or the old history's row r + T before the window. A thread turns +// its own channel's 9 rows, read before any is written +kernel void ple_hist_adv(device float* hist [[buffer(0)]], device const float* norm_w [[buffer(1)]], + constant PleArgs& a [[buffer(2)]], constant uint& T [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + const uint ch = tg.x * 256 + tid, hcd = a.hc * a.n_embd; + if (ch >= hcd) return; + device float* h = hist + 9 * ch; + float col[9]; +#pragma clang loop unroll(full) + for (int r = 0; r < 9; ++r) { + const int src = int(T) - 9 + r; + col[r] = src >= 0 ? norm_w[uint(src) * hcd + ch] : h[uint(r) + T]; + } +#pragma clang loop unroll(full) + for (int r = 0; r < 9; ++r) h[r] = col[r]; +} diff --git a/metal/src/engine/prompt_chain.inc b/metal/src/engine/prompt_chain.inc index 573b460e..ddf205f1 100644 --- a/metal/src/engine/prompt_chain.inc +++ b/metal/src/engine/prompt_chain.inc @@ -132,7 +132,7 @@ bool Window::pc_chunk(const int64_t* tokens, int nT, int64_t pos0, std::string& ds.reserve(2400); nb_ = false; if (pend_fold_ > 0) { // the last window's commit, first - fold(ds, pend_fold_, 1 - hp_); + fold(ds, pend_fold_, fold_hp_); pend_fold_ = 0; } auto P_ = [&](int l, const std::string& fn, const mtl::Constants& c = {}, const std::string& key = "") { @@ -183,7 +183,7 @@ bool Window::pc_chunk(const int64_t* tokens, int nT, int64_t pos0, std::string& HC * N / 64, (nT + 63) / 64, 128); cur_inj = 1 - cur_inj; }; - const int hp = hp_; + const int hp = hp_, hp_next = gdn_ring_ > 0 ? ((hp + 1) % 3) : (1 - hp); int gi = 0, qi = 0; for (int l = 0; l < 48; ++l) { const std::string bl = "blk." + std::to_string(l) + "."; @@ -203,12 +203,24 @@ bool Window::pc_chunk(const int64_t* tokens, int nT, int64_t pos0, std::string& {mtl::Arg::B(b.qkv), mtl::Arg::B(ghist_[hp][(size_t) gi]), mtl::Arg::B(W_, T(bl + "ssm_conv1d.weight").offset), mtl::Arg::B(b.al), mtl::Arg::B(b.be), mtl::Arg::B(W_, T(bl + "ssm_dt.bias").offset), mtl::Arg::B(W_, T(bl + "ssm_a").offset), mtl::Arg::B(b.gqn), mtl::Arg::B(b.gkn), mtl::Arg::B(b.gvc), - mtl::Arg::B(b.ggb), mtl::Arg::B(b.gkq), mtl::Arg::B(ghist_[1 - hp][(size_t) gi]), mtl::Arg::V(ga)}, + mtl::Arg::B(b.ggb), mtl::Arg::B(b.gkq), mtl::Arg::B(ghist_[hp_next][(size_t) gi]), mtl::Arg::V(ga)}, 16 + 24, nT, 256); - D(P_(b.lgdn, "gdn_scanr_4", mtl::Constants().i(6, 0)), - {mtl::Arg::B(gst_[(size_t) gi]), mtl::Arg::B(b.gqn), mtl::Arg::B(b.gkn), mtl::Arg::B(b.gvc), mtl::Arg::B(b.ggb), - mtl::Arg::B(b.gkq), mtl::Arg::B(b.go), mtl::Arg::V(ga)}, - 48 * 4, 1, 32 * 4); + const GdnWy wy = gdn_wy_scratch(nT); + if (gdn_wy_on(b.k4) && nT >= kGdnWyMinT && wy.end <= b.qkv.bytes) { // as prompt_chunk + D(P_(b.lgdn, "gdn_wy_prep"), + {mtl::Arg::B(b.gqn), mtl::Arg::B(b.gkn), mtl::Arg::B(b.ggb), mtl::Arg::B(b.qkv, wy.kqh), mtl::Arg::B(b.qkv, wy.wyh), + mtl::Arg::B(b.qkv, wy.wyf), mtl::Arg::V(ga)}, + 16, (nT + 15) / 16, 128); + D(P_(b.lgdn, "gdn_wy_scan"), + {mtl::Arg::B(gst_[sp_][(size_t) gi]), mtl::Arg::B(b.gvc), mtl::Arg::B(b.qkv, wy.kqh), mtl::Arg::B(b.qkv, wy.wyh), + mtl::Arg::B(b.qkv, wy.wyf), mtl::Arg::B(b.go), mtl::Arg::V(ga)}, + 48 * 8, 1, 64); + } else { + D(P_(b.lgdn, "gdn_scanr_4", mtl::Constants().i(6, 0)), + {mtl::Arg::B(gst_[sp_][(size_t) gi]), mtl::Arg::B(b.gqn), mtl::Arg::B(b.gkn), mtl::Arg::B(b.gvc), mtl::Arg::B(b.ggb), + mtl::Arg::B(b.gkq), mtl::Arg::B(b.go), mtl::Arg::V(ga)}, + 48 * 4, 1, 32 * 4); + } D(P_(b.lgdn, "gdn_post"), {mtl::Arg::B(b.go), mtl::Arg::B(b.z), mtl::Arg::B(W_, T(bl + "ssm_norm.weight").offset), mtl::Arg::B(b.ygh), mtl::Arg::B(b.yg), mtl::Arg::V(ga)}, 48, nT, 128); @@ -234,7 +246,7 @@ bool Window::pc_chunk(const int64_t* tokens, int nT, int64_t pos0, std::string& mtl::Arg::B(cs_), mtl::Arg::B(b.qcur), mtl::Arg::B(b.qidx), mtl::Arg::B(kq_[(size_t) qi]), mtl::Arg::B(vq_[(size_t) qi]), mtl::Arg::B(ksc_[(size_t) qi]), mtl::Arg::B(vsc_[(size_t) qi]), mtl::Arg::B(tail_[(size_t) qi]), mtl::Arg::B(dead_[(size_t) qi]), mtl::Arg::B(pooled_[(size_t) qi]), mtl::Arg::V(a), mtl::Arg::B(mrope_)}, - 33, nT, 256); + 33, nT, qsa_prep_tg(nT)); if (pos0 + nT > PC_QB) { // some row's context is past the budget: the indexer selects if (b.k4) D(P_(b.lqsa, "qsa_score_mm"), {mtl::Arg::B(b.qidx), mtl::Arg::B(pooled_[(size_t) qi]), mtl::Arg::B(dead_[(size_t) qi]), @@ -272,16 +284,14 @@ bool Window::pc_chunk(const int64_t* tokens, int nT, int64_t pos0, std::string& dq(T(bl + "ffn_gate_shexp.weight"), b.Wshg); dq(T(bl + "ffn_up_shexp.weight"), b.Wshu); dq(T(bl + "ffn_down_shexp.weight"), b.Wshd); - gemm(b.mixh, b.Wr, b.logits, nT, NEXP, N); - const uint32_t nt = (uint32_t) nT; + gemm(b.mixh, b.Wr, b.logits, nT, n_exp_, N); + const uint32_t nt = (uint32_t) nT | ((uint32_t) n_exp_ << 16); D(P_(b.lib, "route_topk"), {mtl::Arg::B(b.logits), mtl::Arg::B(b.ids), mtl::Arg::B(b.wts), mtl::Arg::V(nt)}, (nT + 7) / 8, 1, 256); const int nent = nT * TOPK, nblk = (nent + 1023) / 1024; - PcMmPlanArgs pa{(uint32_t) nent, TOPK, NEXP, (uint32_t) PC_TM, (uint32_t) nblk}; - const uint32_t nclr = (uint32_t) b.maxrows; - D(P_(b.lib, "mm_clear"), {mtl::Arg::B(b.stok), mtl::Arg::V(nclr)}, (b.maxrows + 255) / 256, 1, 256); + PcMmPlanArgs pa{(uint32_t) nent, TOPK, (uint32_t) n_exp_, (uint32_t) PC_TM, (uint32_t) nblk}; D(P_(b.lib, "mm_hist"), {mtl::Arg::B(b.ids), mtl::Arg::B(b.bcnt), mtl::Arg::V(pa)}, nblk, 1, 1024); D(P_(b.lib, "mm_scan"), {mtl::Arg::B(b.bcnt), mtl::Arg::B(b.boff), mtl::Arg::B(b.eoff), mtl::Arg::B(b.tiles), mtl::Arg::B(b.meta), - mtl::Arg::V(pa)}, 1, 1, NEXP); + mtl::Arg::V(pa), mtl::Arg::B(b.stok)}, 1, 1, (uint32_t) n_exp_); D(P_(b.lib, "mm_place"), {mtl::Arg::B(b.ids), mtl::Arg::B(b.wts), mtl::Arg::B(b.boff), mtl::Arg::B(b.eoff), mtl::Arg::B(b.stok), mtl::Arg::B(b.sw), mtl::Arg::B(b.pos), mtl::Arg::V(pa)}, nblk, 1, 1024); PcMmArgs ma{eg.offset, eu.offset, ed.offset, 0, 0, (uint32_t) ref::row_bytes(eg.type, N), (uint32_t) ref::row_bytes(ed.type, NFF), @@ -290,8 +300,8 @@ bool Window::pc_chunk(const int64_t* tokens, int nT, int64_t pos0, std::string& ma.d_stride = (uint64_t) ma.d_row * N; const uint32_t ne = N; mtl::Dispatch g; - g.p = P_(b.lib, "mm_gather"); - g.args = {mtl::Arg::B(b.mixed), mtl::Arg::B(b.stok), mtl::Arg::B(b.XS), mtl::Arg::V(ne)}; + g.p = P_(b.lib, "mm_gather_h"); + g.args = {mtl::Arg::B(b.mixh), mtl::Arg::B(b.stok), mtl::Arg::B(b.XS), mtl::Arg::V(ne)}; g.tg[0] = 256; g.indirect = &b.meta; g.ind_off = 8 * 4; @@ -367,7 +377,7 @@ bool Window::pc_chunk(const int64_t* tokens, int nT, int64_t pos0, std::string& } const double g = cx_->run(ds); b.gpu_ms += 1000.0 * g; - hp_ = 1 - hp_; + hp_ = hp_next; // the host side of the commit: the PLE history (the last 9 normalized rows) and the n-gram context if (ple_on_) { using namespace strata::kernels; diff --git a/metal/src/engine/scr_state.hpp b/metal/src/engine/scr_state.hpp new file mode 100644 index 00000000..e06a1016 --- /dev/null +++ b/metal/src/engine/scr_state.hpp @@ -0,0 +1,143 @@ +// metal/src/engine/scr_state.hpp - the Metal window's session state (window.cpp) as bytes, and the SCR +// relocation of one QSA layer's cells (qsa.metal's cache layout). Header-only: shared by the window patch and the +// CPU probe, so the probe tests the code the engine would run. +// +// window.cpp's state, per its init (L554-567) and checkpoint_save (L1688): +// 36 GDN layers (l % 4 != 3): gst_ f32 [128 i][48 h][128 j] = 3 MiB, ghist_[hp_] f32 [10240 ch][3] = 120 KiB +// pleHist_ f32 [10240 ch][9] = 360 KiB, and ss.ple_prev[2] (the PLE n-gram token window) +// 12 QSA layers (l % 4 == 3), per cell c and kv head h (2): +// kq_/vq_ int8 [c][h][256], ksc_/vsc_ half [c][h][4] (one scale per 64 dims) +// pooled_ f32 [b][128] for complete blocks b = c / 4, tail_ f32 [4][128] (raw keys, slot c % 4), dead_ f32 [128] +// qsa_prep (qsa.metal L104-147) applies NEOX RoPE at position = cell (pos_base is 0 on every path) to dims 0..63 +// only, pairs (i, i + 32), i < 32, angle cell * base^(-2i/64); k is rotated BEFORE its int8 quantization, so the +// rotated dims are exactly scale group 0 of each row. pooled[b] is rotated at position 4 b, also dims 0..63. +#pragma once +#include +#include +#include +#include +#include + +namespace scr { + +constexpr int NGDN = 36, NQSA = 12, QKV = 2, QD = 256, QROT = 64, ID = 128, NG_HIST = 9, CH = 10240; +constexpr size_t GST_BYTES = (size_t) 4 * 128 * 48 * 128; // 3,145,728 +constexpr size_t GHIST_BYTES = (size_t) 4 * CH * 3; // 122,880 +constexpr size_t PLE_BYTES = (size_t) 4 * NG_HIST * CH; // 368,640 +constexpr size_t TAIL_BYTES = (size_t) 4 * 4 * ID, DEAD_BYTES = (size_t) 4 * ID; +constexpr size_t KQ_CELL = (size_t) QKV * QD; // int8 bytes per cell (kq or vq) +constexpr size_t SC_CELL = (size_t) QKV * 4 * 2; // half bytes per cell (ksc or vsc) +constexpr size_t POOL_BLOCK = (size_t) 4 * ID; // f32 bytes per pooled block (4 cells) + +// bytes of a saved session: the running state, then the positional cells [0, n) of every QSA layer +constexpr size_t fixed_bytes() { return NGDN * (GST_BYTES + GHIST_BYTES) + PLE_BYTES + NQSA * (TAIL_BYTES + DEAD_BYTES); } +constexpr size_t cell_bytes(size_t n) { return NQSA * (n * (2 * KQ_CELL + 2 * SC_CELL) + (n / 4) * POOL_BLOCK); } + +// the file: this header, the token ids, then the sections in a fixed order, each starting on a 16 KiB boundary +// (the page size: F_NOCACHE reads into page-aligned MTLBuffer storage go straight to the device) +constexpr size_t ALIGN = 16384; +inline size_t align_up(size_t x) { return (x + ALIGN - 1) / ALIGN * ALIGN; } +struct SessionHeader { + char magic[8] = {'S', 'T', 'R', 'S', 'E', 'S', 'S', '1'}; + uint32_t version = 1; + uint32_t flags = 0; // bit 0: PLE on, bit 1: MTP sections follow + int64_t n_cells = 0; // positions [0, n_cells) are in the cache; the next token runs at n_cells + int64_t n_ids = 0; // the token ids of those positions (the harness's prefix check) + uint64_t model_bytes = 0; // the model shard's size and its tensor count: a session is one model's + uint64_t model_tensors = 0; + double rope_base = 0; + int32_t ple_prev[2] = {-1, -1}; + uint64_t payload_bytes = 0; // everything after the header and ids, for the size check before any copy +}; +static_assert(sizeof(SessionHeader) % 8 == 0, "packed as written"); + +// ---- fp16 and the qsa_prep int8 rule (L136-147): per 64 dims sh = half(amax / 127), q = clamp(rint(x / sf)) +inline float h2f(uint16_t h) { _Float16 v; std::memcpy(&v, &h, 2); return (float) v; } +inline uint16_t f2h(float f) { _Float16 v = (_Float16) f; uint16_t h; std::memcpy(&h, &v, 2); return h; } +inline void quant64(const float* x, int8_t* q, uint16_t* sc) { + float m = 0.0f; + for (int i = 0; i < 64; ++i) m = std::max(m, std::fabs(x[i])); + *sc = f2h(m / 127.0f); + const float sf = h2f(*sc); + for (int i = 0; i < 64; ++i) q[i] = (int8_t) (sf > 0.0f ? std::clamp((int) std::nearbyint(x[i] / sf), -127, 127) : 0); +} + +// the rotation by delta positions: cos/sin of delta * theta_i in double, as window.cpp's table (L545-551) +struct Rot { + float c[32], s[32]; + Rot(double base, int64_t delta) { + for (int i = 0; i < 32; ++i) { + const double th = (double) delta * std::pow(base, -2.0 * i / 64.0); + c[i] = (float) std::cos(th); + s[i] = (float) std::sin(th); + } + } +}; +// the first 64 dims of a row rotated in place (norm_rope's pairs: a = x[i], b = x[i + 32]) +inline void rotate64(float* x, const Rot& r) { + for (int i = 0; i < 32; ++i) { + const float a = x[i], b = x[i + 32]; + x[i] = a * r.c[i] - b * r.s[i]; + x[i + 32] = a * r.s[i] + b * r.c[i]; + } +} + +// One QSA layer's span of n cells moved from c_src to c_dst = c_src + delta, keys re-rotated by delta. +// kq: only group 0 (dims 0..63, scale ksc[.][.][0]) is rotated: dequant, rotate, requant with a fresh scale. +// dims 64..255 and scales 1..3 move as bytes. vq / vsc carry no position: bytes. +// pooled: the complete blocks inside the span, b -> b + delta / 4, rotated by delta (4 b -> 4 b + delta). delta must +// be a multiple of 4, else every block of the span would mix different cells and the raw keys aren't kept. +// The junction block (the one holding c_dst when c_dst % 4 != 0) is not moved: it mixes the new span before it, +// so the caller recomputes the span's first (4 - c_dst % 4) % 4 cells or more (scr_min_junction). +// Cells are walked away from the overlap (ascending when moving down), so source rows are read before they're hit. +inline int64_t scr_min_junction(int64_t c_dst) { return (4 - (c_dst & 3)) & 3; } +inline bool relocate_layer(int8_t* kq, int8_t* vq, uint16_t* ksc, uint16_t* vsc, float* pooled, int64_t c_src, + int64_t c_dst, int64_t n, const Rot& rot) { + const int64_t delta = c_dst - c_src; + if (delta == 0 || n <= 0) return true; + if ((delta & 3) != 0 || c_dst < 0) return false; + auto cell = [&](int64_t k) { + const int64_t s = c_src + k, d = c_dst + k; + for (int h = 0; h < QKV; ++h) { + const size_t rs = (size_t) s * QKV + h, rd = (size_t) d * QKV + h; + float x[64]; + const float sf = h2f(ksc[rs * 4]); + for (int i = 0; i < 64; ++i) x[i] = (float) kq[rs * QD + i] * sf; + rotate64(x, rot); + std::memmove(kq + rd * QD + 64, kq + rs * QD + 64, QD - 64); + std::memmove(ksc + rd * 4 + 1, ksc + rs * 4 + 1, 3 * 2); + quant64(x, kq + rd * QD, ksc + rd * 4); + std::memmove(vq + rd * QD, vq + rs * QD, QD); + std::memmove(vsc + rd * 4, vsc + rs * 4, 4 * 2); + } + }; + const int64_t b0 = (c_src + 3) / 4, b1 = (c_src + n) / 4, db = delta / 4; // complete blocks [b0, b1) + auto block = [&](int64_t b) { + float x[ID]; + std::memcpy(x, pooled + (size_t) b * ID, sizeof x); + rotate64(x, rot); + std::memcpy(pooled + (size_t) (b + db) * ID, x, sizeof x); + }; + if (delta < 0) { + for (int64_t k = 0; k < n; ++k) cell(k); + for (int64_t b = b0; b < b1; ++b) block(b); + } else { + for (int64_t k = n - 1; k >= 0; --k) cell(k); + for (int64_t b = b1 - 1; b >= b0; --b) block(b); + } + return true; +} + +// ---- GDN staleness, the part with a closed form. gdn.metal (L89-119): S <- g S + k delta^T, delta = (v - g S^T k) b, +// i.e. S' = g (I - b k k^T) S + b k v^T with |k| <= 1 (k = cv / sqrt(sum cv^2 + eps)) and b = sigmoid in (0, 1). Two +// runs over the same C inputs from states S_1, S_2 differ by D' = g (I - b k k^T) D, and |I - b k k^T|_2 = 1, so +// |D_end|_F <= exp(sum_t log g_t) |D_0|_F, log g_t = softplus(alpha_t + dt_h) * ssm_a_h (<= 0) +// per value head. That bounds the direct term of keeping S_ABC instead of S_AB'C; the indirect term (C's own k, v, +// gate inputs computed under B) has no closed form and is what the splice test measures. +inline double log_gate(float alpha, float dt, float ssm_a) { + const double v = (double) alpha + dt; + const double sp = v > 20.0 ? v : std::log1p(std::exp(v)); + return sp * (double) ssm_a; +} + +} // namespace scr diff --git a/metal/src/engine/slop4.h b/metal/src/engine/slop4.h new file mode 100644 index 00000000..b57f7554 --- /dev/null +++ b/metal/src/engine/slop4.h @@ -0,0 +1,217 @@ +#ifndef STRATA_SLOP4_H +#define STRATA_SLOP4_H +#ifdef __METAL_VERSION__ +#define SL_CONSTANT constant +#define SL_THREAD thread +#else +#define SL_CONSTANT static const +#define SL_THREAD +#endif + +namespace slop4 { + +enum : uint8_t { PROSE = 1, TOOL = 2, FENCE = 4, BOL = 8, COLON = 16, PBOL = 32, KNOWN = 128 }; + +struct State { + int32_t t1 = -1, t2 = -1; // the last token and the one before it + uint8_t f = 0; // PROSE: after in this turn; BOL: at a line start; COLON: that line follows "...:" + uint8_t len = 0; // content tokens on the current line (saturating) + uint8_t pad[2] = {}; +}; + +struct Knobs { float lex, hon, md, colon; }; + +struct Pen { int32_t id; float v; }; + +// token classes: 1 a newline run ("\n", "\n\n", " \n", ...), 2 blanks only (indentation) +SL_CONSTANT int32_t kNewline[] = { + 198, 201, 271, 317, 695, 845, 987, 1031, 1358, 1517, 1726, 2228, 2317, 2449, 3265, 4326, 4408, 4558, 4969, 5690, + 5744, 5773, 6010, 6159, 6219, 6710, 6987, 7388, 8301, 8675, 8726, 9122, 10109, 10367, 10628, 10797, 11463, + 11953, 12213, 13488, 14200, 14305, 14379, 15216, 15334, 15398, 15602, 15684, 16391, 16750, 16957, 17116, 17694, + 17817, 18015, 18050, 18438, 19075, 19721, 21996, 22616, 22727, 23574, 23745, 23840, 24630, 25263, 25401, 25445, + 25698, 25952, 26467, 26480, 26771, 27330, 29644, 29773, 30858, 31472, 32077, 32527, 32811, 33006, 33019, 33339, + 33438, 33496, 33954, 34127, 35369, 35412, 35624, 35973, 36444, 36826, 37050, 37528, 38336, 38544, 39000, 39598, + 40260, 40315, 41210, 41585, 43622, 43916, 44058, 44281, 44879, 44902, 46813, 47458, 47947, 48305, 49404, 49420, + 49746, 50946, 52245, 53585, 54701, 56137, 56580, 56955, 56963, 57103, 57633, 58943, 59366, 59776, 60121, 61332, + 62386, 62591, 63063, 63111, 63454, 64137, 64576, 65117, 65382, 67511, 68830, 69229, 69340, 69896, 70273, 70707, + 70866, 70929, 71356, 71848, 71990, 72666, 73714, 75128, 75321, 75668, 75981, 76294, 76374, 76421, 76507, 76909, + 77670, 77933, 78035, 78254, 79519, 81018, 81927, 82471, 82946, 83774, 83829, 84075, 85709, 85748, 85809, 85844, + 85935, 86092, 86182, 87199, 87251, 87328, 88065, 88574, 88927, 88985, 89803, 91421, 91668, 91804, 92699, 92766, + 93957, 94820, 94861 +}; +SL_CONSTANT int32_t kBlank[] = { + 197, 220, 256, 257, 260, 262, 285, 297, 309, 336, 388, 406, 454, 493, 556, 646, 670, 761, 950, 967, 1025, 1107, + 1335, 1604, 1638, 1734, 1779, 1855, 2118, 2215, 2467, 2475, 2596, 2672, 3235, 3388, 3551, 3696, 4152, 4420, + 4448, 4946, 5015, 5071, 5233, 5780, 6173, 6247, 6323, 6447, 6800, 7216, 7237, 7322, 7536, 7599, 7879, 8074, + 8424, 8830, 9081, 9411, 9876, 10190, 10275, 10747, 11449, 11530, 11611, 12473, 12670, 12687, 13076, 13152, + 13300, 13857, 14221, 14826, 14843, 14978, 16203, 16754, 16851, 17025, 17122, 17468, 17780, 17895, 18186, 18625, + 18688, 18690, 19679, 20200, 20334, 20839, 20851, 21519, 21641, 22539, 22568, 22691, 23420, 23581, 24952, 25232, + 25870, 26058, 27729, 27860, 28959, 29034, 29374, 29427, 29708, 29915, 30026, 30631, 30927, 31593, 31632, 31726, + 33368, 34158, 34381, 34450, 35315, 35696, 35853, 35914, 36909, 37181, 37207, 38177, 38448, 38772, 39700, 40595, + 41293, 41785, 44606, 44714, 44914, 45043, 45207, 45550, 45964, 46330, 46572, 47256, 48546, 48742, 48835, 49349, + 49373, 49642, 49741, 50328, 52872, 52989, 53928, 54652, 55036, 55133, 55720, 56288, 56612, 56683, 57926, 58096, + 60270, 62037, 63023, 63067, 63287, 63288, 64278, 64571, 65897, 66227, 66639, 67204, 67685, 68190, 68977, 71032, + 71400, 72030, 72352, 72738, 73753, 74019, 75424, 76010, 76332, 76785, 76937, 77120, 77968, 78052, 78420, 80396, + 80905, 82130, 83375, 83770, 84703, 85012, 86978, 88011, 88074, 88531, 91507, 92130, 92492, 92770, 93340, 93457, + 93647, 94057, 95323, 95459 +}; + +#ifdef __METAL_VERSION__ +inline uint8_t cls(int32_t t) { + for (int32_t x : kNewline) if (x == t) return 1; + for (int32_t x : kBlank) if (x == t) return 2; + return 0; +} +#else +inline uint8_t cls(int32_t t) { + static const std::vector tab = [] { + std::vector v(248320, 0); + for (int32_t i : kNewline) v[(size_t) i] = 1; + for (int32_t i : kBlank) v[(size_t) i] = 2; + return v; + }(); + return t >= 0 && (size_t) t < tab.size() ? tab[(size_t) t] : 0; +} +#endif + +template inline bool in(int32_t t, SL_THREAD const int32_t (&s)[N]) { + for (int32_t x : s) if (x == t) return true; + return false; +} + +#ifdef __METAL_VERSION__ +template inline bool in(int32_t t, constant int32_t (&s)[N]) { + for (int32_t x : s) if (x == t) return true; + return false; +} +#endif + +enum : int32_t { + IM_START = 248045, IM_END = 248046, THINK = 248068, ETHINK = 248069, TOOL_CALL = 248058, ETOOL_CALL = 248059, + NL = 198, NL2 = 271, +}; +SL_CONSTANT int32_t kHold[] = {IM_END, THINK, IM_START, 248044}; +SL_CONSTANT int32_t kFence[] = {71093, 52451}; // "```", " ```" +SL_CONSTANT int32_t kColon[] = {25, 64700, 92217}; // ":", ":**", "**:" +SL_CONSTANT int32_t kDigit[] = {16, 17, 18, 19, 20, 21, 22, 23, 24}; // "1".."9" + +inline State step(State s, int32_t tok) { + s.t2 = s.t1; + s.t1 = tok; + s.f |= KNOWN; + const uint8_t bol = s.f & BOL; + s.f &= (uint8_t) ~PBOL; + if (tok == IM_START || tok == IM_END || tok == THINK) { + s.f = KNOWN | BOL; + s.len = 0; + } else if (tok == ETHINK) { + s.f = KNOWN | PROSE | BOL; + s.len = 0; + } else if (tok == TOOL_CALL) { + s.f |= TOOL; + } else if (tok == ETOOL_CALL) { + s.f = (uint8_t) ((s.f & ~(TOOL | COLON)) | BOL); + s.len = 0; + } else if (cls(tok) == 1) { + if (!bol) s.f = (uint8_t) ((s.f & ~COLON) | (in(s.t2, kColon) ? COLON : 0)); + s.f |= BOL; + s.len = 0; + } else if (cls(tok) == 2) { + // indentation: the line has not started yet + } else { + if (in(tok, kFence) && (s.f & PROSE) && !(s.f & TOOL)) s.f ^= FENCE; + s.f = (uint8_t) ((s.f & ~(BOL | COLON)) | (bol ? PBOL : 0)); + if (s.len < 255) ++s.len; + } + return s; +} + +// the penalties on the next token after state s; out holds at least 64 +inline int penalties(SL_THREAD const State& s, SL_THREAD const Knobs& k, SL_THREAD Pen* out) { + if ((s.f & (KNOWN | PROSE | TOOL | FENCE)) != (KNOWN | PROSE)) return 0; + int n = 0; + auto pen = [&](int32_t id, float v) { if (v > 0.0f && n < 64) out[n++] = {id, v}; }; + const int32_t p0 = s.t1, p1 = s.t2; + // the 7 lexical cliches (results/banned_llm_isms.json), as in v3 + const int32_t rather[] = {4598, 24944, 61745, 70714}, ather_pre[] = {2515, 2697, 15656, 43487}; + const int32_t right[] = {1245, 1246, 5791, 9785, 26806, 51140}; + const int32_t end_pre[] = {400, 809, 3601, 3839, 14637, 55293}; + const int32_t shrt[] = {2716, 8412, 10382, 12115, 36639}; + const int32_t exactly[] = {6681, 63151, 66172}, ex_pre[] = {325, 5291}; + const int32_t here[] = {1532, 5514, 6527, 8160, 475}, what[] = {1092, 3437, 3710, 12195, 29348}; + if (in(p0, rather) || (p0 == 1792 && in(p1, ather_pre))) pen(1056, k.lex); // " than" + if (in(p0, right) || (p0 == 481 && p1 == 2697)) pen(3023, k.lex); // " away" + if (p0 == 310 && in(p1, end_pre)) pen(809, k.lex); // "end to" -> " end" + if (in(p0, shrt)) pen(2243, k.lex); // " version" + if (in(p0, exactly) || (p0 == 31654 && in(p1, ex_pre))) pen(1092, k.lex); // " what" + if (p0 == 579 && in(p1, here)) pen(1092, k.lex); // "here's" -> " what" + if (p0 == 579 && in(p1, what)) pen(3404, k.lex); // "what's" -> " actually" + // performative honesty + const int32_t det[] = {1719, 760, 279, 561, 276, 449, 1473, 2014, 2329, 821, 2921, 4888, 1323, 381, + 1602, 32985, 1379, 3444, 7261, 6183, 777, 4254, 678}; // the/an/my/be/being/most/more/your + const int32_t clause[] = {13, 11, 0, 30, 25, 26, 1076, 2218, 1892, 12, 471, 8, + 321, 694, 748, 429, 7834, 684, 1541, 1921, 1987, 2939, 3850, 4272}; + const int32_t h_[] = {71, 39, 3101, 75833}; // "h", "H", "(h", "\"h" + const int32_t sentence[] = {13, 0, 30, 25, 26, 1076, 2218, 1892, 8}; // . ! ? : ; ... — ) before " honest take" + const bool line_start = cls(p0) == 1 || (s.f & BOL); + const bool sentence_start = in(p0, sentence) || line_start; + const bool clause_start = in(p0, clause) || line_start; + pen(92115, k.hon); // "Honestly" + pen(68540, k.hon); // " Honestly" + pen(51397, k.hon); // " Honest" + if (in(p0, det) || sentence_start) pen(10429, k.hon); // " honest" + if (clause_start) pen(25211, k.hon); // " honestly" + if (in(p0, h_)) { pen(33961, 1.5f * k.hon); pen(77696, 1.5f * k.hon); } // "h|onest", "h|onestly": committed + if (p0 == 660) { pen(46253, k.hon); pen(25211, k.hon); } // "in all" -> " honesty" + // candor announcements: with " honest" pushed down, "If I'm being" shifts to a sibling (real, frank, totally, ...). + // That is still the same frame, so those are pushed down after "being" / "be" too. Earlier in the frame, "I'm" and + // "the short/real", most of the mass is on content, so the push is aimed there as well. + const int32_t candor[] = {1865, 47045, 24969, 5507, 17290, 7435, 62179, 66842, 86775, // real blunt frank candid transparent straight upfront brutally truthful + 12116, 6379, 6999, 13542, 11040, 73018}; // totally completely fully perfectly entirely painfully + const int32_t be_pre[] = {310, 951, 3172, 728, 1066, 16032}; // to, 'll, me, just, gonna + " be" + const int32_t intens[] = {12116, 6379, 6999, 13542, 11040, 73018, 66842}; // "totally honest" + if (p0 == 1602 || p0 == 32985) for (int32_t id : candor) pen(id, k.hon); // being real + if ((p0 == 381 || p0 == 1323) && in(p1, be_pre)) { + for (int32_t id : candor) pen(id, k.hon); // to be blunt, i'll be straight + pen(2050, k.hon); // be direct + if (p1 == 310 || p1 == 951) { pen(6415, k.hon); pen(2708, k.hon); } // to be fair, to be clear + } + if (in(p0, intens) && (p1 == 1602 || p1 == 32985 || p1 == 381 || p1 == 1323)) { pen(10429, k.hon); for (int32_t id : candor) pen(id, k.hon); } + if (p0 == 2688 && (p1 == 353 || p1 == 585 || p1 == 40 || p1 == 72)) { pen(1602, 0.5f * k.hon); pen(10429, k.hon); } // I'm being / I'm honest + if (clause_start) pen(40285, k.hon); // " frankly" + if (p0 == 660) { pen(49034, k.hon); pen(71853, k.hon); } // "in all" fairness / seriousness + // "the honest / real / blunt / short answer", "short answer:", "my straight take": the frame noun after the frame word + const int32_t frame_adj[] = {10429, 1865, 47045, 24969, 5507, 7435, 2716, 4145, 3841, 26860, 13983, + 8412, 12115, 10382, 7700, 12402, 8535, 4880, 32254}; // + short/Short/ Short/real/Real/ Real/ actual/ Actual + const int32_t frame_det[] = {1719, 760, 279, 561, 2329, 821, 2921, 4888, 276, 449}; + const int32_t frame[] = {4087, 1831, 56004, 2243, 7793, 1301, 11782, 14996, 3296}; // answer take framing version truth read summary assessment question + if (in(p0, frame_adj) && (in(p1, frame_det) || in(p1, sentence) || cls(p1) == 1 || p1 == ETHINK)) + for (int32_t id : frame) pen(id, k.hon); + // the branch points one token before the frame closes: a clause-initial "the" / "to" can still go to content + const int32_t the_[] = {1719, 760, 279, 561}, to_[] = {1206, 1946, 951, 310}; + const bool p1_root = in(p1, sentence) || cls(p1) == 1 || p1 == ETHINK || p1 < 0; + const int32_t candor_branch[] = {1865, 47045, 24969, 5507, 7435, 2716, 26860}; + if (in(p0, the_) && (p1_root || cls(p0) == 0 && (s.f & PBOL))) + for (int32_t id : candor_branch) pen(id, 0.5f * k.hon); // the real/blunt/frank/candid/straight/short/brutal + if (in(p0, to_) && (p1_root || (s.f & PBOL))) pen(381, 0.5f * k.hon); // clause-initial "to be" + const int32_t line_frame[] = {8412, 12115, 7700, 12402}; + if (s.f & BOL) for (int32_t id : line_frame) pen(id, 0.5f * k.hon); // "short answer:" / "real talk" at a line start + // structure: bullets, bold and # headers, numbered lists, rules at a line start + if (s.f & BOL) { + const float v = k.md + ((s.f & COLON) ? k.colon : 0.0f); + const int32_t md[] = {12, 9, 6458, 471, 348, 7061, 332, 2972, 2, 550, 13962, 794, 653, 7460, 16112, + 4277, 11862, 4006, 2218, 10, 49300}; // ... "---", "***", "–", "—", "+", "→" + for (int32_t id : md) pen(id, v); + for (int32_t id : kDigit) pen(id, (s.f & COLON) ? v : 0.25f * v); // light, unless the line above ends in ":" + } + if ((s.f & PBOL) && in(p0, kDigit)) { pen(13, k.md + k.colon); pen(8, k.md + k.colon); } // "1." "1)" + // a short label line ending in ":" that would open a list + if (in(p0, kColon) && s.len <= 6) { pen(NL, k.colon); pen(NL2, k.colon); } + return n; +} + +} // namespace slop4 + +#undef SL_CONSTANT +#undef SL_THREAD +#endif diff --git a/metal/src/engine/spec_min_p.hpp b/metal/src/engine/spec_min_p.hpp new file mode 100644 index 00000000..e2f25dc4 --- /dev/null +++ b/metal/src/engine/spec_min_p.hpp @@ -0,0 +1,62 @@ +#pragma once + +#include +#include +#include +#include +#include + +namespace strata::metal_engine { + +struct DraftMinP { + std::array steps{}; + int count = 0; + float min_cum = 0.0f; + + static bool probability(const char*& p, float& value) { + char* end = nullptr; + errno = 0; + value = std::strtof(p, &end); + if (end == p || errno == ERANGE || !std::isfinite(value) || + value < 0.0f || value > 1.0f) + return false; + p = end; + while (std::isspace(static_cast(*p))) ++p; + return true; + } + + bool parse(const char* schedule, const char* cumulative) { + DraftMinP next; + if (schedule) { + const char* p = schedule; + do { + if (next.count == int(next.steps.size()) || + !probability(p, next.steps[next.count])) + return false; + ++next.count; + if (*p == '\0') break; + if (*p++ != ',') return false; + } while (true); + } + if (cumulative) { + const char* p = cumulative; + if (!probability(p, next.min_cum) || *p != '\0') return false; + } + *this = next; + return true; + } + + bool enabled() const { return count > 0 || min_cum > 0.0f; } + + bool keep(int step, float p, float fallback, double& cumulative) const { + if (enabled() && (!std::isfinite(p) || p < 0.0f || p > 1.0f)) + return false; + cumulative *= p; + const int index = step < count ? step : count - 1; + const float threshold = count ? steps[index] : fallback; + return !(threshold > 0.0f && p < threshold) && + !(min_cum > 0.0f && cumulative < min_cum); + } +}; + +} // namespace strata::metal_engine diff --git a/metal/src/engine/stream.cpp b/metal/src/engine/stream.cpp index 9db74e25..c06943ba 100644 --- a/metal/src/engine/stream.cpp +++ b/metal/src/engine/stream.cpp @@ -62,6 +62,11 @@ bool ExpertStream::open(const gguf::File& f, std::string& err) { if (fd_ < 0) { err = "metal stream: cannot open " + f.path; return false; } fcntl(fd_, F_NOCACHE, 1); fcntl(fd_, F_RDAHEAD, 0); + const gguf::Tensor* g0 = f.find("blk.0.ffn_gate_exps.weight"); + if (g0 && g0->shape.size() >= 3 && g0->shape[2] >= 1 && g0->shape[2] <= NE) { + ne_ = (int) g0->shape[2]; + } + if (const char* mb = std::getenv("STRATA_METAL_MOCE_BONUS")) moce_bonus_ = (float) std::atof(mb); lay_.resize(NL); for (int l = 0; l < NL; ++l) { const std::string b = "blk." + std::to_string(l) + "."; @@ -69,7 +74,7 @@ bool ExpertStream::open(const gguf::File& f, std::string& err) { const gguf::Tensor* u = f.find(b + "ffn_up_exps.weight"); const gguf::Tensor* d = f.find(b + "ffn_down_exps.weight"); if (!g || !u || !d) { err = "metal stream: layer " + std::to_string(l) + " has no routed experts"; return false; } - lay_[(size_t) l] = {g->offset, g->bytes / NE, u->offset, u->bytes / NE, d->offset, d->bytes / NE, -1}; + lay_[(size_t) l] = {g->offset, g->bytes / (uint64_t) ne_, u->offset, u->bytes / (uint64_t) ne_, d->offset, d->bytes / (uint64_t) ne_, -1}; } // the expert-contiguous copy beside the GGUF, when it is this GGUF's (sizes in its header) and its first and // last experts match the GGUF's bytes @@ -88,16 +93,18 @@ bool ExpertStream::open(const gguf::File& f, std::string& err) { uint32_t ver = 0, nl = 0, ne = 0; uint64_t src = 0; std::memcpy(&ver, h + 4, 4); std::memcpy(&nl, h + 8, 4); std::memcpy(&ne, h + 12, 4); std::memcpy(&src, h + 16, 8); - ok = ok && ver == 1 && nl == (uint32_t) NL && ne == (uint32_t) NE && src == (uint64_t) st.st_size; + const Lay& last_lay = lay_[(size_t) (NL - 1)]; + const bool has_gguf_exps = (uint64_t) st.st_size >= last_lay.d + (uint64_t) ne_ * last_lay.ds; + ok = ok && ver == 1 && nl == (uint32_t) NL && ne == (uint32_t) ne_ && (!has_gguf_exps || src == (uint64_t) st.st_size); for (int l = 0; ok && l < NL; ++l) { const Lay& x = lay_[(size_t) l]; ok = t[(size_t) l * 4 + 1] == x.gs && t[(size_t) l * 4 + 2] == x.us && t[(size_t) l * 4 + 3] == x.ds; off[(size_t) l] = t[(size_t) l * 4]; } - if (ok) { // spot check: layer 0's first and layer 47's last expert, every slice + if (ok && has_gguf_exps) { // spot check: layer 0's first and layer 47's last expert, every slice for (const int l : {0, NL - 1}) { const Lay& x = lay_[(size_t) l]; - const int e = l == 0 ? 0 : NE - 1; + const int e = l == 0 ? 0 : ne_ - 1; std::vector a(x.gs + x.us + x.ds), b(a.size()); ok = ok && ::pread(cfd, a.data(), a.size(), (off_t) (off[(size_t) l] + (uint64_t) e * a.size())) == (ssize_t) a.size() && ::pread(fd_, b.data(), x.gs, (off_t) (x.g + (uint64_t) e * x.gs)) == (ssize_t) x.gs && @@ -148,6 +155,7 @@ bool ExpertStream::open(const gguf::File& f, std::string& err) { } nio_ = std::max(1, env_int("STRATA_METAL_IO_THREADS", 6)); split_ = std::max(1, std::min(8, env_int("STRATA_METAL_READ_SPLIT", 1))); + nio_urgent_ = std::max(0, std::min(nio_ - 1, env_int("STRATA_METAL_IO_URGENT", 2))); // The read-ahead's depth and where its experts wait in the LRU, measured on the M4 (notes-mac-ssd.md, 3 interleaved // pairs each): Q2_0 (1.38 MB experts) top 10 at the LRU front 12.2 tok/s vs top 6 at the tail 11.3; the IQ3_XXS // mix (1.75 MB, the SSD is the limit) top 6 at the tail 8.06 vs top 10 at the front 7.43. @@ -158,6 +166,11 @@ bool ExpertStream::open(const gguf::File& f, std::string& err) { pre_k2_ = env_int("STRATA_METAL_PREFETCH2", 0); prefetch_ = pre_k_ > 0; pre_mru_ = env_int("STRATA_METAL_PREFETCH_MRU", big ? 0 : 1) != 0; + // a predicted expert that is already resident keeps its LRU place (touching it cost 4% more misses in the + // closed-loop replay, audits/opus_strata_metal_opts_v1.md); =1 the old behaviour + pre_touch_ = env_int("STRATA_METAL_PREFETCH_TOUCH", 0) != 0; + // the read-ahead bookkeeping after the GPU is let go (after_layer); =0 all of it inside the handshake + split_turn_ = env_int("STRATA_METAL_SPLIT_TURN", 1) != 0; // timing only (WRONG answers): every missing expert is pointed at slot 0, nothing is read fake_ = env_int("STRATA_METAL_STREAM_FAKE", 0) != 0; return true; @@ -178,6 +191,83 @@ bool ExpertStream::read(void* dst, uint64_t len, uint64_t off) const { bool ExpertStream::load_dense(gguf::File& f, mtl::Context& cx, const std::vector& skip, mtl::Buffer& out, std::string& err) { const double t0 = now_ms(); + uint64_t max_dense_end = 0; + uint64_t min_exp_off = UINT64_MAX; + size_t dense_cnt = 0; + for (const auto& kv : f.tensors) { + const gguf::Tensor& t = kv.second; + if (t.name.find("_exps") != std::string::npos) { + min_exp_off = std::min(min_exp_off, t.offset); + continue; + } + ++dense_cnt; + max_dense_end = std::max(max_dense_end, (t.offset + t.bytes + PAGE - 1) / PAGE * PAGE); + } + if (f.map != nullptr && max_dense_end > 0 && max_dense_end <= f.map_bytes && max_dense_end <= min_exp_off && + env_int("STRATA_METAL_MMAP_DENSE", 1) != 0) { + const size_t alloc_before = cx.allocated(); + uint64_t locked_bytes = 0, skipped_hc_bytes = 0, unpinned_embd_bytes = 0; + std::vector> raw_lazy; + for (const auto& kv : f.tensors) { + const gguf::Tensor& t = kv.second; + if (t.name.find("_exps") != std::string::npos) continue; + const bool is_skip = std::find(skip.begin(), skip.end(), t.name) != skip.end(); + const bool is_embd = t.name == "token_embd.weight"; + if (is_skip || is_embd) { + if (is_skip) skipped_hc_bytes += t.bytes; + if (is_embd) unpinned_embd_bytes += t.bytes; + raw_lazy.emplace_back(t.offset, t.offset + t.bytes); + } + } + std::sort(raw_lazy.begin(), raw_lazy.end()); + std::vector> lazy_ranges; + for (const auto& r : raw_lazy) { + if (!lazy_ranges.empty() && r.first <= lazy_ranges.back().second) + lazy_ranges.back().second = std::max(lazy_ranges.back().second, r.second); + else + lazy_ranges.push_back(r); + } + // Inward page-align lazy ranges (llama-mmap.cpp align_range) so boundary pages shared with active GPU tensors are never dropped + std::vector> aligned_lazy; + for (const auto& r : lazy_ranges) { + const uint64_t p0 = (r.first + PAGE - 1) / PAGE * PAGE; + const uint64_t p1 = std::min(r.second, max_dense_end) / PAGE * PAGE; + if (p0 < p1) { + madvise((uint8_t*) f.map + p0, (size_t) (p1 - p0), MADV_DONTNEED); + aligned_lazy.emplace_back(p0, p1); + } + } + // Coalesced active ranges = complement of aligned_lazy in [0, max_dense_end) (llama-mmap.cpp ranges_complement) + const bool do_mlock = env_int("STRATA_METAL_MLOCK_DENSE", 1) != 0; + uint64_t pos = 0; + auto lock_span = [&](uint64_t beg, uint64_t end) { + if (beg >= end) return; + void* ptr = (uint8_t*) f.map + beg; + const size_t len = (size_t) (end - beg); + if (do_mlock) { + madvise(ptr, len, MADV_WILLNEED); + if (mlock(ptr, len) == 0) locked_bytes += len; + else dense_unwired_ += len; + } else { + dense_unwired_ += len; + } + }; + for (const auto& r : aligned_lazy) { + if (r.first > pos) lock_span(pos, r.first); + pos = std::max(pos, r.second); + } + if (pos < max_dense_end) lock_span(pos, max_dense_end); + if (f.map_bytes > max_dense_end) { + madvise((uint8_t*) f.map + max_dense_end, f.map_bytes - (size_t) max_dense_end, MADV_DONTNEED); + } + out = cx.wrap(f.map, (size_t) max_dense_end, /*count_allocated=*/false); + (void) alloc_before; + std::fprintf(stderr, "strata metal stream: locked %.2f GB of active GPU dense tensors (unpinned %.2f GB CPU token_embd + %.2f GB skipped BF16 HC in %.2f GB file map)\n", + locked_bytes / 1e9, unpinned_embd_bytes / 1e9, skipped_hc_bytes / 1e9, max_dense_end / 1e9); + std::fprintf(stderr, "strata metal stream: %zu dense tensors wrapped zero-copy from file (%.2f GB clean file-backed map, %.1f ms)\n", + dense_cnt, max_dense_end / 1e9, now_ms() - t0); + return true; + } std::vector ts; uint64_t tot = 0; for (auto& kv : f.tensors) { @@ -250,19 +340,30 @@ void ExpertStream::lru_front(int s) { if (k.tail < 0) k.tail = s; } -int ExpertStream::victim(int ci) { +// avoid_layer: a read-ahead for that layer queued while the GPU may already route it (after_layer runs beside the +// GPU): none of its resident experts may go, or the GPU's routing (MoCE) and its resident-first expert pass would race +// the eviction +int ExpertStream::victim(int ci, int avoid_layer, bool allow_hot) { Cls& k = cls_[(size_t) ci]; - for (int s = k.tail; s >= 0; s = slots_[(size_t) s].prev) - if (slots_[(size_t) s].state.load() == 0 && slots_[(size_t) s].pin != stamp_) { + const int n_passes = (in_prompt_ && allow_hot) ? 2 : 1; + for (int pass = 0; pass < n_passes; ++pass) { + for (int s = k.tail; s >= 0; s = slots_[(size_t) s].prev) + if (slots_[(size_t) s].state.load() == 0 && slots_[(size_t) s].pin != stamp_ && !slots_[(size_t) s].limbo && + (avoid_layer < 0 || slots_[(size_t) s].layer != avoid_layer) && + (gpu_layer_ < 0 || slots_[(size_t) s].layer != gpu_layer_) && + (!after_turn_ || slots_[(size_t) s].layer != gpu_layer_ + 1) && + !(split_turn_ && !in_prompt_ && slots_[(size_t) s].pre) && (!in_prompt_ || pass == 1 || !slots_[(size_t) s].hot)) { Slot& x = slots_[(size_t) s]; if (x.layer >= 0) { const int key = x.layer * NE + x.e; where_[(size_t) key] = -1; smap_.as()[key] = 0xFFFFFFFFu; } + if (in_prompt_ && multi_chunk_) x.hot = false; x.layer = x.e = -1; return s; } + } return -1; } @@ -272,9 +373,10 @@ int ExpertStream::victim(int ci) { void ExpertStream::read_part(int s, int l, int e, int part) { const Lay& x = lay_[(size_t) l]; if (cfd_ >= 0) { // the contiguous copy: the whole expert, piece part of split_ (the slot keeps gate | up | down) + const int sp = in_prompt_ ? 1 : split_; const uint64_t len = x.gs + x.us + x.ds, src = coff_[(size_t) l] + (uint64_t) e * len; - const uint64_t a = len * (uint64_t) part / (uint64_t) split_ / 4096 * 4096; - const uint64_t b = part + 1 == split_ ? len : len * (uint64_t) (part + 1) / (uint64_t) split_ / 4096 * 4096; + const uint64_t a = len * (uint64_t) part / (uint64_t) sp / 4096 * 4096; + const uint64_t b = part + 1 == sp ? len : len * (uint64_t) (part + 1) / (uint64_t) sp / 4096 * 4096; uint8_t* p = host_[(size_t) s] + a; uint64_t n = b - a, o = src + a; while (n > 0) { @@ -285,7 +387,7 @@ void ExpertStream::read_part(int s, int l, int e, int part) { } p += r; o += (uint64_t) r; n -= (uint64_t) r; } - bytes_read_ += b - a; + (in_prompt_ ? p_bytes_read_ : bytes_read_) += b - a; return; } const int k = part / split_, p = part % split_; @@ -298,15 +400,17 @@ void ExpertStream::read_part(int s, int l, int e, int part) { std::fprintf(stderr, "strata metal stream: reading expert %d of layer %d failed\n", e, l); std::exit(6); } - bytes_read_ += b - a; + (in_prompt_ ? p_bytes_read_ : bytes_read_) += b - a; } -void ExpertStream::io_main() { +void ExpertStream::io_main(int idx) { for (;;) { Job j; { std::unique_lock g(mu_); - cv_.wait(g, [&] { return stop_ || !urgent_q_.empty() || !pre_q_.empty(); }); + // the first nio_urgent_ threads are kept for the misses the GPU waits on while decoding + const auto urgent_only = [&] { return idx < nio_urgent_ && !in_prompt_; }; + cv_.wait(g, [&] { return stop_ || !urgent_q_.empty() || (!urgent_only() && !pre_q_.empty()); }); if (stop_) return; if (!urgent_q_.empty()) { j = urgent_q_.front(); urgent_q_.pop_front(); } else { j = pre_q_.front(); pre_q_.pop_front(); } @@ -318,10 +422,10 @@ void ExpertStream::io_main() { std::lock_guard g(mu_); Slot& x = slots_[(size_t) j.slot]; if (--x.left == 0) { - smap_.as()[j.layer * NE + j.e] = code_[(size_t) j.slot]; + if (!x.pre) smap_.as()[j.layer * NE + j.e] = code_[(size_t) j.slot]; x.state = 0; --outstanding_; - ++reads_; + ++(in_prompt_ ? p_reads_ : reads_); done = true; } } @@ -335,7 +439,7 @@ void ExpertStream::submit(int s, int l, int e, bool urgent) { x.e = (int16_t) e; x.state = 1; x.pre = !urgent; - const int parts = cfd_ >= 0 ? split_ : 3 * split_; + const int parts = cfd_ >= 0 ? (in_prompt_ ? 1 : split_) : 3 * split_; x.left = parts; ++outstanding_; where_[(size_t) (l * NE + e)] = s; @@ -377,11 +481,11 @@ bool ExpertStream::make_arena(mtl::Context& cx, uint64_t budget, const std::stri char magic[4]; uint32_t h[5]; if (std::fread(magic, 1, 4, f) == 4 && std::memcmp(magic, "STRP", 4) == 0 && std::fread(h, 4, 5, f) == 5 && - h[1] == (uint32_t) NL && h[2] == (uint32_t) NE) { + h[1] == (uint32_t) NL && h[2] == (uint32_t) ne_) { for (uint32_t i = 0; i < h[4]; ++i) { uint16_t p[2]; if (std::fread(p, 2, 2, f) != 2) break; - if (p[0] < NL && p[1] < NE && !seen[(size_t) p[0] * NE + p[1]]) { + if (p[0] < NL && p[1] < (uint16_t) ne_ && !seen[(size_t) p[0] * NE + p[1]]) { seen[(size_t) p[0] * NE + p[1]] = 1; rank.push_back({p[0], p[1]}); } @@ -392,9 +496,11 @@ bool ExpertStream::make_arena(mtl::Context& cx, uint64_t budget, const std::stri std::fprintf(stderr, "strata metal stream: no expert profile at %s: the arena starts layer-interleaved\n", profile.c_str()); } const size_t from_profile = rank.size(); - for (int e = 0; e < NE; ++e) + for (int e = 0; e < ne_; ++e) for (int l = 0; l < NL; ++l) if (!seen[(size_t) l * NE + e]) rank.push_back({l, e}); + for (int l = 0; l < NL; ++l) layer_rank_[l].clear(); + for (const auto& p : rank) if (p.first >= 0 && p.first < NL) layer_rank_[p.first].push_back(p.second); // the pool of each size class: at least min_slots (a prompt chunk's layer needs up to 512 at once), then the // top-ranked pairs that fit in the budget uint64_t used = 0; @@ -403,7 +509,7 @@ bool ExpertStream::make_arena(mtl::Context& cx, uint64_t budget, const std::stri for (const auto& p : rank) { Cls& k = cls_[(size_t) lay_[(size_t) p.first].cls]; const uint64_t cost = k.n < min_slots ? 0 : k.pages * PAGE; - if (used + cost > budget) break; + if (used + cost > budget) continue; used += cost; ++k.n; fill.push_back(p); @@ -475,7 +581,7 @@ bool ExpertStream::make_arena(mtl::Context& cx, uint64_t budget, const std::stri k.head = k.n ? k.base : -1; k.tail = k.n ? k.base + k.n - 1 : -1; } - for (int i = 0; i < nio_; ++i) io_.emplace_back([this] { io_main(); }); + for (int i = 0; i < nio_; ++i) io_.emplace_back([this, i] { io_main(i); }); // the fill: the ranked pairs in steps (best first; each step in file order) on the I/O threads; before each step // grow() asks whether macOS still has the memory for it (arena_grow); the best-ranked ends at the LRU front uint64_t step = 1536ull << 20; @@ -524,7 +630,17 @@ bool ExpertStream::make_arena(mtl::Context& cx, uint64_t budget, const std::stri ++nretired; } } - // the LRU order = the rank order (best at the front) + // the LRU order = the rank order (best at the front); mark top (k.n - (ne_ + 32)) STRP slots hot + std::vector cls_seen(cls_.size(), 0); + for (size_t i = 0; i < done; ++i) { + const int s = where_[(size_t) (fill[i].first * NE + fill[i].second)]; + if (s >= 0) { + const int ci = cls_of(s); + const int hot_cap = std::max(0, cls_[(size_t) ci].n - (ne_ + 32)); + slots_[(size_t) s].hot = (cls_seen[(size_t) ci] < hot_cap); + ++cls_seen[(size_t) ci]; + } + } for (size_t i = done; i-- > 0;) lru_front(where_[(size_t) (fill[i].first * NE + fill[i].second)]); } if (stopped) @@ -533,6 +649,8 @@ bool ExpertStream::make_arena(mtl::Context& cx, uint64_t budget, const std::stri arena_live_ = arena_bytes_ - retired_bytes; fill_s_ = (now_ms() - t0) / 1e3; fill_bytes_ = bytes_read_.load(); + bytes_read_ = 0; + reads_ = 0; uint64_t pairs = 0; for (const auto& k : cls_) pairs += (uint64_t) k.n; std::fprintf(stderr, "strata metal stream: expert arena %.2f GB, %llu slots in %zu size classes, %zu experts read from the " @@ -571,23 +689,47 @@ bool ExpertStream::make_arena(mtl::Context& cx, uint64_t budget, const std::stri // the experts a predicted routing names, most likely first (by the best logit over the window's tokens) void ExpertStream::ahead(int ln, int T, const float* pred, int K) { K = std::min(K, 32); + static const bool unpert_all = env_int("STRATA_METAL_PREFETCH_UNPERT", 1) != 0; + const uint32_t* sm = smap_.as() + (size_t) ln * NE; std::pair pw[8 * 32]; int np = 0; + static const int k_draft_max = env_int("STRATA_METAL_PREFETCH_DRAFT", 1); for (int t = 0; t < T; ++t) { + const int Kt = (t == 0) ? K : std::min(K, k_draft_max); + if (Kt <= 0) continue; const float* lg = pred + (size_t) t * NE; + auto eff_lg = [&](int e) -> float { + return lg[e] + (moce_bonus_ > 0.0f && sm[e] != 0xFFFFFFFFu ? moce_bonus_ : 0.0f); + }; + const int k_unpert = unpert_all ? Kt : std::min(3, Kt); int top[32]; int nt = 0; - for (int e = 0; e < NE; ++e) { - if (nt < K) { top[nt++] = e; continue; } - int m = 0; - for (int j = 1; j < K; ++j) if (lg[top[j]] < lg[top[m]]) m = j; - if (lg[e] > lg[top[m]]) top[m] = e; + for (int r = 0; r < k_unpert; ++r) { + int best_e = -1; + float best_v = -INFINITY; + for (int e = 0; e < ne_; ++e) { + bool seen_e = false; + for (int j = 0; j < nt; ++j) if (top[j] == e) { seen_e = true; break; } + if (!seen_e && lg[e] > best_v) { best_v = lg[e]; best_e = e; } + } + if (best_e >= 0) top[nt++] = best_e; + } + for (int e = 0; e < ne_; ++e) { + if (k_unpert >= Kt) break; + bool in_top = false; + for (int j = 0; j < k_unpert; ++j) if (top[j] == e) { in_top = true; break; } + if (in_top) continue; + if (nt < Kt) { top[nt++] = e; continue; } + int m = k_unpert; + for (int j = k_unpert + 1; j < Kt; ++j) if (eff_lg(top[j]) < eff_lg(top[m])) m = j; + if (eff_lg(e) > eff_lg(top[m])) top[m] = e; } for (int j = 0; j < nt; ++j) { + const float prio = (j < k_unpert) ? (lg[top[j]] + moce_bonus_ + 1e3f) : eff_lg(top[j]); int q = 0; while (q < np && pw[q].second != top[j]) ++q; - if (q == np) pw[np++] = {lg[top[j]], top[j]}; - else pw[q].first = std::max(pw[q].first, lg[top[j]]); + if (q == np) pw[np++] = {prio, top[j]}; + else pw[q].first = std::max(pw[q].first, prio); } } std::sort(pw, pw + np, [](const auto& a, const auto& b) { return a.first > b.first; }); @@ -597,8 +739,8 @@ void ExpertStream::ahead(int ln, int T, const float* pred, int K) { for (int j = 0; j < np; ++j) { const int key = ln * NE + pw[j].second; const int s0 = where_[(size_t) key]; - if (s0 >= 0) { lru_front(s0); continue; } - const int s = victim(lay_[(size_t) ln].cls); + if (s0 >= 0) { if (pre_touch_) lru_front(s0); continue; } + const int s = victim(lay_[(size_t) ln].cls, split_turn_ ? ln : -1); if (s < 0) break; submit(s, ln, pw[j].second, false); slots_[(size_t) s].pin = stamp_; @@ -611,17 +753,32 @@ void ExpertStream::ahead(int ln, int T, const float* pred, int K) { if (any) cv_.notify_all(); } +// Do not write future-layer residency while its GPU router can be running. The +// existing shared-event handshake publishes completed preads before GPU resume. +void ExpertStream::publish_ready(int lo, int hi) { + for (int l = lo; l <= hi; ++l) + for (int e = 0; e < ne_; ++e) { + const int key = l * NE + e, s = where_[(size_t) key]; + if (s < 0) continue; + const Slot& x = slots_[(size_t) s]; + if (!x.limbo && x.state.load(std::memory_order_acquire) == 0) + smap_.as()[key] = code_[(size_t) s]; + } +} + void ExpertStream::at_layer(int l, int T, int topk, const int32_t* ids, const float* pred, const float* pred2) { const double t0 = now_ms(); const bool prompt = T > 8; // a prompt chunk (the GEMM path): nearly every expert of the layer is used if (l == 0) { if (prompt) ++p_chunks_; else ++windows_; } + gpu_layer_ = l; + after_turn_ = false; ++stamp_; // the slots this layer uses are pinned until the next handshake bool used[NE] = {}; int want[NE]; int nw = 0; for (int i = 0; i < T * topk; ++i) { const int e = ids[i]; - if (e < 0 || e >= NE || used[e]) continue; + if (e < 0 || e >= ne_ || used[e]) continue; used[e] = true; want[nw++] = e; } @@ -632,16 +789,45 @@ void ExpertStream::at_layer(int l, int T, int topk, const int32_t* ids, const fl return; } std::vector wait; + const bool split = split_turn_ && !prompt; { std::lock_guard g(mu_); // read-aheads of this layer (or older) that were not started and are not wanted: dropped (the slot goes free) - { + if (!split) { std::vector drop; for (const Job& j : pre_q_) { if (j.layer > l || slots_[(size_t) j.slot].state.load() != 1) continue; if (!(j.layer == l && used[j.e]) && std::find(drop.begin(), drop.end(), j.slot) == drop.end()) drop.push_back(j.slot); } for (int s : drop) cancel_slot(s); + } else { // the unwanted read-aheads of this layer: queued ones dropped now (the misses below get the SSD), all + // parked until the next turn's after_layer + retire_unused_pre(l, l, &limbo_new_, used); + } + if (!split) { + std::unique_lock g2(mu_, std::adopt_lock); + done_cv_.wait(g2, [&] { + for (int e = 0; e < ne_; ++e) { + if (used[e]) continue; + const int s = where_[(size_t) (l * NE + e)]; + if (s >= 0 && slots_[(size_t) s].pre && slots_[(size_t) s].state.load() != 0) return false; + } + return true; + }); + g2.release(); + // Evict any unchosen speculative read-ahead of this layer that finished so where_, smap_, and LRU order stay 100% deterministic + for (int e = 0; e < ne_; ++e) { + if (used[e]) continue; + const int s = where_[(size_t) (l * NE + e)]; + if (s >= 0 && slots_[(size_t) s].pre && slots_[(size_t) s].state.load() == 0) { + Slot& x = slots_[(size_t) s]; + where_[(size_t) (l * NE + e)] = -1; + smap_.as()[l * NE + e] = 0xFFFFFFFFu; + x.layer = x.e = -1; + x.pre = false; + lru_back(s); + } + } } for (int j = 0; j < nw; ++j) // pass 1: pin what is resident (or on its way), so no miss below evicts it if (where_[(size_t) (l * NE + want[j])] >= 0) slots_[(size_t) where_[(size_t) (l * NE + want[j])]].pin = stamp_; @@ -661,19 +847,54 @@ void ExpertStream::at_layer(int l, int T, int topk, const int32_t* ids, const fl promote_slot(s); // a queued read-ahead becomes urgent wait.push_back(s); } - if (!(prompt && x.pre)) lru_front(s); // a prompt's streamed experts stay at the LRU tail + if (!prompt || (!multi_chunk_ ? !x.pre : x.hot)) lru_front(s); // a prompt's streamed experts stay at the LRU tail + else if (multi_chunk_) lru_back(s); x.pre = false; + if (st == 0) smap_.as()[key] = code_[(size_t) s]; continue; } ++(prompt ? p_misses_ : misses_); - s = victim(lay_[(size_t) l].cls); + s = victim(lay_[(size_t) l].cls, -1, true); if (s < 0) { std::fprintf(stderr, "strata metal stream: no free slot (layer %d)\n", l); std::exit(6); } submit(s, l, e, true); slots_[(size_t) s].pin = stamp_; if (!prompt) lru_front(s); + else if (multi_chunk_) lru_back(s); wait.push_back(s); } + // a prompt's lookahead: layer l + 1's experts as its router ranks them on this layer's FFN input, the most + // tokens first, queued behind this layer's misses now (not after the wait), so the SSD never runs dry + // (not with a MoCE prompt, 0 < STRATA_METAL_PROMPT_MOCE < 100: its routing reads residency, as the static read-ahead) + const float pm = std::getenv("STRATA_METAL_PROMPT_MOCE") ? (float) std::atof(std::getenv("STRATA_METAL_PROMPT_MOCE")) : 0.0f; + for (int distance = 1; prompt && distance <= 2; ++distance) { + const int32_t* predicted = distance == 1 ? p_pred_ : p_pred2_; + if (!predicted || l + distance >= NL || (pm > 0.0f && pm < 100.0f)) continue; + const int ln = l + distance; + int cnt[NE] = {}; + for (int i = 0; i < T * p_pred_k_; ++i) { + const int e = predicted[i]; + if (e >= 0 && e < ne_) ++cnt[e]; + } + int order[NE], no = 0; + static const int pmin = std::max(1, env_int("STRATA_METAL_PROMPT_PRED_MIN", 1)); + for (int e = 0; e < ne_; ++e) if (cnt[e] >= pmin) order[no++] = e; + std::stable_sort(order, order + no, [&](int a, int b) { return cnt[a] > cnt[b]; }); + int queued = 0; + const int cap = distance == 1 ? NE : std::clamp(env_int("STRATA_METAL_PROMPT_PRED2_MAX", 64), 0, NE); + for (int j = 0; j < no && queued < cap; ++j) { + const int e = order[j]; + if (where_[(size_t) (ln * NE + e)] >= 0) continue; + const int s = victim(lay_[(size_t) ln].cls, l + 1, !multi_chunk_ || distance == 1); + if (s < 0) break; + submit(s, ln, e, false); + ++queued; + slots_[(size_t) s].pin = stamp_; + lru_back(s); + ++p_streamed_; + } + } } + if (prompt && (p_pred_ || p_pred2_)) cv_.notify_all(); if (!wait.empty()) { cv_.notify_all(); const double tw = now_ms(); @@ -684,20 +905,31 @@ void ExpertStream::at_layer(int l, int T, int topk, const int32_t* ids, const fl }); (prompt ? p_wait_ms_ : wait_ms_) += now_ms() - tw; } + const float cur_pm = std::getenv("STRATA_METAL_PROMPT_MOCE") ? (float) std::atof(std::getenv("STRATA_METAL_PROMPT_MOCE")) : 0.0f; + if (prompt && cur_pm > 0.0f && env_int("STRATA_METAL_PUBLISH_PREFETCH", 1) != 0) { + std::lock_guard g(mu_); + publish_ready(l + 1, std::min(l + 2, NL - 1)); + } if (prompt) { // a chunk uses nearly all of the next layer's experts: every one not in the arena is streamed in behind this // layer's compute, into the LRU tail (the decode-hot experts stay) - if (l + 1 < NL) { + const float p_moce_env = std::getenv("STRATA_METAL_PROMPT_MOCE") ? (float) std::atof(std::getenv("STRATA_METAL_PROMPT_MOCE")) : 0.0f; + if (l + 1 < NL && !p_pred_ && (p_moce_env <= 0.0f || p_moce_env >= 100.0f)) { bool any = false; std::lock_guard g(mu_); const int ln = l + 1; - for (int e = 0; e < NE; ++e) { + const int max_pre = std::getenv("STRATA_METAL_PROMPT_MAX_PRE") ? std::atoi(std::getenv("STRATA_METAL_PROMPT_MAX_PRE")) : (p_moce_env < 0.0f ? 96 : (ln >= 40 ? 128 : 256)); + int queued = 0; + for (int e : layer_rank_[ln]) { + if (queued >= max_pre) break; if (where_[(size_t) (ln * NE + e)] >= 0) continue; - const int s = victim(lay_[(size_t) ln].cls); + const int s = victim(lay_[(size_t) ln].cls, -1, false); if (s < 0) break; submit(s, ln, e, false); slots_[(size_t) s].pin = stamp_; + lru_back(s); ++p_streamed_; + ++queued; any = true; } if (any) cv_.notify_all(); @@ -705,15 +937,87 @@ void ExpertStream::at_layer(int l, int T, int topk, const int32_t* ids, const fl (prompt ? p_host_ms_ : host_ms_) += now_ms() - t0; return; } + if (split) { // the rest after the GPU is let go (after_layer) + host_ms_ += now_ms() - t0; + return; + } // read-ahead: layer l + 1 (each token's top pre_k_ of the predicted router logits), then layer l + 2 (top pre_k2_) if (prefetch_ && pred && l + 1 < NL) ahead(l + 1, T, pred, pre_k_); if (pre_k2_ > 0 && pred2 && l + 2 < NL) ahead(l + 2, T, pred2, pre_k2_); host_ms_ += now_ms() - t0; } +// Read-aheads of layers lo..hi that the routing did not claim: queued pieces dropped, the slots parked in `into` (still +// mapped, no victims) until release_limbo(), whatever their I/O state now, so which slot is +// free when depends on the order of the turns only, not on the SSD's timing (MoCE routing reads residency). +void ExpertStream::retire_unused_pre(int lo, int hi, std::vector* into, const bool* used) { + for (int key = lo * NE; key < (hi + 1) * NE; ++key) { + const int s = where_[(size_t) key]; + if (s < 0 || (used && used[key % NE])) continue; + Slot& x = slots_[(size_t) s]; + if (!x.pre || x.limbo) continue; + if (x.state.load() == 1) { // not started: its pieces dropped, nothing will write the slot + for (auto it = pre_q_.begin(); it != pre_q_.end();) it = it->slot == s ? pre_q_.erase(it) : it + 1; + x.state = 0; + x.left = 0; + --outstanding_; + ++pre_cancel_; + } + x.limbo = true; + into->push_back(s); + } +} + +void ExpertStream::release_limbo(std::unique_lock& g) { + if (limbo_.empty()) return; + done_cv_.wait(g, [&] { // a read that is still running (rare: it had a whole layer of GPU time) + for (int s : limbo_) if (slots_[(size_t) s].state.load() != 0) return false; + return true; + }); + std::sort(limbo_.begin(), limbo_.end()); + for (int s : limbo_) { + Slot& x = slots_[(size_t) s]; + if (x.layer >= 0 && where_[(size_t) (x.layer * NE + x.e)] == s) { + const int key = x.layer * NE + x.e; + where_[(size_t) key] = -1; + smap_.as()[key] = 0xFFFFFFFFu; + } + x.layer = x.e = -1; + x.pre = x.limbo = false; + lru_back(s); + } + limbo_.clear(); +} + +void ExpertStream::after_layer(int l, int T, const float* pred, const float* pred2) { + if (!split_turn_ || fake_) return; + const double t0 = now_ms(); + after_turn_ = true; + { + std::unique_lock g(mu_); + release_limbo(g); // the last turn's (they had a whole layer of GPU time to finish reading) + limbo_.swap(limbo_new_); // this turn's, retired in at_layer + } + if (prefetch_ && pred && l + 1 < NL) ahead(l + 1, T, pred, pre_k_); + if (pre_k2_ > 0 && pred2 && l + 2 < NL) ahead(l + 2, T, pred2, pre_k2_); + host_after_ms_ += now_ms() - t0; +} + // A region of `bytes` at the end of one arena buffer, for a prompt's scratch buffers: its slots are emptied and held // (state 3) while a prompt runs (hold_region), and are ordinary slots again for decoding. uint8_t* ExpertStream::region(uint64_t bytes) { + if (env_int("STRATA_METAL_PHASE_TRADE", 1) == 0 && env_int("STRATA_METAL_SEPARATE_PROMPT_BUF", 0) != 0) { + const uint64_t sz = (bytes + PAGE - 1) / PAGE * PAGE; + void* m = mmap(nullptr, (size_t) sz, PROT_READ | PROT_WRITE, MAP_ANON | MAP_PRIVATE, -1, 0); + if (m == MAP_FAILED) return nullptr; + region_slots_.clear(); + region_ = (uint8_t*) m; + region_bytes_ = 0; // fresh zero pages from MAP_ANON; no arena slots to memset + sep_region_bytes_ = sz; + std::fprintf(stderr, "strata metal stream: prompt scratch %.2f GB mapped in separate anonymous mmap\n", + sz / 1e9); + return region_; + } for (int c = (int) chunks_.size() - 1; c >= 0; --c) { if (chunks_[(size_t) c].bytes < bytes + (64ull << 20)) continue; uint8_t* end = chunks_[(size_t) c].as() + chunks_[(size_t) c].bytes; @@ -730,14 +1034,14 @@ uint8_t* ExpertStream::region(uint64_t bytes) { std::vector stolen(cls_.size(), 0); for (int s : region_slots_) ++stolen[(size_t) cls_of(s)]; bool tight = false; - for (size_t i = 0; i < cls_.size(); ++i) tight = tight || (cls_[i].n > 0 && cls_[i].n - stolen[i] < NE); + for (size_t i = 0; i < cls_.size(); ++i) tight = tight || (cls_[i].n > 0 && cls_[i].n - stolen[i] < ne_); if (tight) { const uint64_t sz = (bytes + PAGE - 1) / PAGE * PAGE; void* m = mmap(nullptr, (size_t) sz, PROT_READ | PROT_WRITE, MAP_ANON | MAP_PRIVATE, -1, 0); if (m == MAP_FAILED) return nullptr; region_slots_.clear(); region_ = (uint8_t*) m; - region_bytes_ = 0; // fresh zero pages; nothing in the arena to empty + region_bytes_ = 0; sep_region_bytes_ = sz; // fresh zero pages; unmapped in release_region() std::fprintf(stderr, "strata metal stream: the arena is small: the prompt's %.2f GB of buffers are their own " "allocation, not a region of its slots\n", bytes / 1e9); return region_; @@ -752,9 +1056,22 @@ uint8_t* ExpertStream::region(uint64_t bytes) { return nullptr; } -void ExpertStream::hold_region(bool hold) { +void ExpertStream::hold_region(bool hold, bool multi_chunk) { std::unique_lock g(mu_); + in_prompt_ = hold; + multi_chunk_ = hold && multi_chunk; + gpu_layer_ = -1; + after_turn_ = false; if (hold) { + // Refresh hot flags from current LRU order (protecting top k.n - (ne_ + 32) slots in every prompt turn) + for (size_t ci = 0; ci < cls_.size(); ++ci) { + const int hot_cap = std::max(0, cls_[ci].n - (ci == 0 ? cls_[ci].n / 4 : (2 * ne_ + 32))); + int rank_idx = 0; + for (int s = cls_[ci].head; s >= 0; s = slots_[(size_t) s].next) { + slots_[(size_t) s].hot = (rank_idx < hot_cap); + ++rank_idx; + } + } // nothing may be reading into those slots: the queued read-aheads are dropped, the running ones waited for std::vector drop; for (const Job& j : pre_q_) @@ -762,6 +1079,10 @@ void ExpertStream::hold_region(bool hold) { drop.push_back(j.slot); for (int s : drop) cancel_slot(s); done_cv_.wait(g, [&] { return outstanding_ == 0; }); + limbo_.insert(limbo_.end(), limbo_new_.begin(), limbo_new_.end()); + limbo_new_.clear(); + retire_unused_pre(0, NL - 1, &limbo_); // decode's last read-aheads: a prompt's victims may take them + release_limbo(g); for (int s : region_slots_) { Slot& x = slots_[(size_t) s]; if (x.layer >= 0) { @@ -773,43 +1094,119 @@ void ExpertStream::hold_region(bool hold) { x.pre = false; x.state = 3; } - if (region_) std::memset(region_, 0, (size_t) region_bytes_); + if (region_ && region_bytes_ > 0) std::memset(region_, 0, (size_t) region_bytes_); + // Kick off layer 0's hottest non-resident experts immediately so they overlap with prompt setup + layer 0 GDN + const float p_moce_env = std::getenv("STRATA_METAL_PROMPT_MOCE") ? (float) std::atof(std::getenv("STRATA_METAL_PROMPT_MOCE")) : 0.0f; + if (!fake_ && (p_moce_env <= 0.0f || p_moce_env >= 100.0f)) { + bool any = false; + // with the chunk's lookahead (STRATA_METAL_PROMPT_PRED) the whole profile rank: what layer 0 does not use and + // the SSD has not started on by its handshake is dropped there + const int max_pre0 = std::getenv("STRATA_METAL_PROMPT_MAX_PRE") ? std::atoi(std::getenv("STRATA_METAL_PROMPT_MAX_PRE")) + : (p_moce_env < 0.0f ? 96 : (env_int("STRATA_METAL_PROMPT_PRED", 1) != 0 ? NE : 256)); + const int layers = std::clamp(env_int("STRATA_METAL_PROMPT_START_LAYERS", 1), 1, 4); + for (int l = 0; l < layers; ++l) { + int queued = 0; + const int cap = l == 0 ? max_pre0 : std::min(max_pre0, 32); + for (int e : layer_rank_[l]) { + if (queued >= cap) break; + if (where_[(size_t) (l * NE + e)] >= 0) continue; + const int s = victim(lay_[(size_t) l].cls, 0, false); + if (s < 0) break; + submit(s, l, e, false); + lru_back(s); + ++p_streamed_; + ++queued; + any = true; + } + } + if (any) cv_.notify_all(); + } } else { + std::vector drop; + for (const Job& j : pre_q_) + if (slots_[(size_t) j.slot].state.load() == 1 && std::find(drop.begin(), drop.end(), j.slot) == drop.end()) + drop.push_back(j.slot); + for (int s : drop) cancel_slot(s); + done_cv_.wait(g, [&] { return outstanding_ == 0; }); for (int s : region_slots_) { slots_[(size_t) s].state = 0; if (!retired_[(size_t) s]) lru_back(s); // a retired slot's memory serves the prompt only } + if (split_turn_) { // the prompt's unclaimed stream-ins: decode's victim() passes read-aheads by + limbo_.insert(limbo_.end(), limbo_new_.begin(), limbo_new_.end()); + limbo_new_.clear(); + retire_unused_pre(0, NL - 1, &limbo_); + release_limbo(g); + } } } -void ExpertStream::reset_stats() { +bool ExpertStream::reuse_region(int mode) { + if (mode != 1 || sep_region_bytes_ == 0 || region_ == nullptr) return false; + return true; +} + +uint64_t ExpertStream::release_region() { + std::unique_lock g(mu_); + if (sep_region_bytes_ > 0 && region_ != nullptr) { + const uint64_t freed = sep_region_bytes_; + munmap(region_, (size_t) sep_region_bytes_); + region_ = nullptr; + sep_region_bytes_ = 0; + return freed; + } + if (region_bytes_ > 0 && region_ != nullptr && env_int("STRATA_METAL_PHASE_TRADE", 0) == 2) { + for (int s : region_slots_) { + Slot& x = slots_[(size_t) s]; + if (x.layer >= 0) { + const int key = x.layer * NE + x.e; + where_[(size_t) key] = -1; + smap_.as()[key] = 0xFFFFFFFFu; + } + x.layer = x.e = -1; + x.pre = false; + } + madvise(region_, (size_t) region_bytes_, MADV_FREE_REUSABLE); + return region_bytes_; + } + return 0; +} + +void ExpertStream::reset_decode_stats() { lookups_ = hits_ = pre_hits_ = late_ = misses_ = pre_issued_ = pre_cancel_ = windows_ = 0; - p_lookups_ = p_hits_ = p_pre_ = p_misses_ = p_streamed_ = p_chunks_ = 0; - p_wait_ms_ = p_host_ms_ = 0; bytes_read_ = 0; reads_ = 0; - wait_ms_ = host_ms_ = 0; + wait_ms_ = host_ms_ = host_after_ms_ = 0; +} + +void ExpertStream::reset_stats() { + reset_decode_stats(); + p_lookups_ = p_hits_ = p_pre_ = p_misses_ = p_streamed_ = p_chunks_ = 0; + p_wait_ms_ = p_host_ms_ = 0; + p_bytes_read_ = 0; + p_reads_ = 0; } void ExpertStream::report(FILE* f) const { if (!ready()) return; if (p_chunks_) std::fprintf(f, "expert stream prompts %llu chunks: %llu distinct lookups (%.1f per layer), in the arena %.4f, streamed " - "ahead %.4f, missed %.4f; %llu streamed in; host %.1f ms per chunk, of it waiting %.1f ms\n", + "ahead %.4f, missed %.4f; %llu streamed in (%.1f MB/chunk, %llu reads); host %.1f ms per chunk, of it waiting %.1f ms\n", (unsigned long long) p_chunks_, (unsigned long long) p_lookups_, p_lookups_ / 48.0 / p_chunks_, p_hits_ / std::max(1.0, (double) p_lookups_), p_pre_ / std::max(1.0, (double) p_lookups_), - p_misses_ / std::max(1.0, (double) p_lookups_), (unsigned long long) p_streamed_, p_host_ms_ / p_chunks_, - p_wait_ms_ / p_chunks_); + p_misses_ / std::max(1.0, (double) p_lookups_), (unsigned long long) p_streamed_, + p_bytes_read_.load() / 1e6 / p_chunks_, (unsigned long long) p_reads_.load(), + p_host_ms_ / p_chunks_, p_wait_ms_ / p_chunks_); const double w = std::max(1, windows_); std::fprintf(f, "expert stream %llu windows: %llu distinct lookups (%.1f per window), hits %.4f, read-ahead hits " "%.4f (%llu of them late), misses %.4f (%.1f per window); read-aheads %llu (%.1f per window, %llu dropped); %.1f MB " - "read per window (%llu reads); host %.2f ms per window, of it waiting %.2f ms\n", + "read per window (%llu reads); host %.2f ms per window, of it waiting %.2f ms (+ %.2f ms after the GPU went on)\n", (unsigned long long) windows_, (unsigned long long) lookups_, lookups_ / w, hits_ / std::max(1.0, (double) lookups_), pre_hits_ / std::max(1.0, (double) lookups_), (unsigned long long) late_, misses_ / std::max(1.0, (double) lookups_), misses_ / w, (unsigned long long) pre_issued_, pre_issued_ / w, (unsigned long long) pre_cancel_, - (bytes_read_.load() - 0) / 1e6 / w, (unsigned long long) reads_.load(), host_ms_ / w, wait_ms_ / w); + (bytes_read_.load() - 0) / 1e6 / w, (unsigned long long) reads_.load(), host_ms_ / w, wait_ms_ / w, host_after_ms_ / w); } } // namespace strata::metal_engine diff --git a/metal/src/engine/stream.h b/metal/src/engine/stream.h index 2806938a..887e6639 100644 --- a/metal/src/engine/stream.h +++ b/metal/src/engine/stream.h @@ -53,6 +53,7 @@ public: const std::function& grow = {}); uint64_t shrink_to() const { return shrink_to_; } uint64_t arena_used_bytes() const { return arena_live_; } // the arena's bytes in use (its retired slots left out) + uint64_t dense_unwired_bytes() const { return dense_unwired_; } // load_dense's mapped dense bytes left unwired bool ready() const { return !chunks_.empty(); } // the arena's buffers (at most 4, each within the device's maxBufferLength): chunk(i) for i >= count is chunk 0 const mtl::Buffer& chunk(int i) const { return chunks_[(size_t) i < chunks_.size() ? (size_t) i : 0]; } @@ -63,32 +64,55 @@ public: // computed from this layer's FFN input (nullptr: no read-ahead). // pred2: layer l + 2's, from the same input (STRATA_METAL_PREFETCH2 = its top-k, 0 = off) void at_layer(int l, int T, int topk, const int32_t* ids, const float* pred, const float* pred2 = nullptr); + // a prompt chunk's lookahead: T x topk ids of layer l + 1's router on layer l's FFN input (nullptr: the static + // expert-profile rank instead); ids2 optionally names layer l + 2; read by at_layer + void set_prompt_pred(const int32_t* ids, int topk, const int32_t* ids2 = nullptr) { + p_pred_ = ids; p_pred_k_ = topk; p_pred2_ = ids2; + } + // the rest of a decode layer's turn, after the GPU has been let go (run_graph_seg's after_cut): the read-aheads + // the routing did not use are retired, the next layer's are queued. Only when split_turn() (else at_layer did it). + void after_layer(int l, int T, const float* pred, const float* pred2 = nullptr); + bool split_turn() const { return split_turn_; } int ahead2() const { return pre_k2_; } // a prompt's scratch buffers inside the arena: region() places them (once), hold_region() empties and holds its // slots while a prompt runs and gives them back after uint8_t* region(uint64_t bytes); - void hold_region(bool hold); + void hold_region(bool hold, bool multi_chunk = false); + uint64_t release_region(); + bool reuse_region(int mode); + size_t region_stolen_slots() const { return region_slots_.size(); } void report(FILE* f) const; void reset_stats(); + // decode-only: keeps the prompt (p_*) counters, so the final report can still show prefill + // accounting after window_reset_stats separates the decode window from the prompt chunk. + void reset_decode_stats(); private: - uint64_t shrink_to_ = 0; + uint64_t shrink_to_ = 0, dense_unwired_ = 0; struct Lay { uint64_t g, gs, u, us, d, ds; int cls; }; - struct Slot { int16_t layer = -1; int16_t e = -1; int prev = -1, next = -1; bool pre = false; int left = 0; uint64_t pin = 0; std::atomic state{0}; }; // 0 idle, 1 queued, 2 reading + struct Slot { int16_t layer = -1; int16_t e = -1; int prev = -1, next = -1; bool pre = false, hot = false, limbo = false; int left = 0; uint64_t pin = 0; std::atomic state{0}; }; // 0 idle, 1 queued, 2 reading struct Cls { uint64_t pages = 0; int n = 0; int base = 0; int head = -1, tail = -1; }; // slot range [base, base+n) struct Job { int slot; int layer, e, part; }; // part: slice * split_ + piece void lru_unlink(int s); void lru_front(int s); void lru_back(int s); int cls_of(int s) const; - int victim(int cls); + int victim(int cls, int avoid_layer = -1, bool allow_hot = true); void read_part(int slot, int l, int e, int part); void cancel_slot(int s); void promote_slot(int s); - void io_main(); + void ahead(int ln, int T, const float* pred, int K); void submit(int slot, int l, int e, bool urgent); void wait_slot(int s); + // the unclaimed read-aheads of layers lo..hi (not in `used`) parked in `into` (mu_ held) + void retire_unused_pre(int lo, int hi, std::vector* into, const bool* used = nullptr); + void release_limbo(std::unique_lock& g); + std::vector limbo_; // retired read-aheads: not victims until released one turn later (deterministic LRU) + std::vector limbo_new_; // ... retired in this turn's at_layer (released at the next turn's after_layer) + int nio_urgent_ = 0; // I/O threads that take urgent reads only (STRATA_METAL_IO_URGENT) + void io_main(int idx); + void publish_ready(int lo, int hi); // mu_ held, GPU at the layer handshake int fd_ = -1; // .experts (metal/tools/stream_experts.py): each expert's gate, up and down rows together, read with one @@ -100,6 +124,12 @@ private: std::vector slots_; std::vector retired_; // slots left unused when the stepped fill stopped at macOS's floor (never in an LRU) std::vector where_; // (l * 512 + e) -> slot, -1 + std::vector layer_rank_[48]; // per-layer experts ordered hottest-first from expert-profile.bin + const int32_t* p_pred_ = nullptr; + const int32_t* p_pred2_ = nullptr; + int gpu_layer_ = -1; + bool after_turn_ = false; + int p_pred_k_ = 0; std::vector chunks_; std::vector region_slots_; uint8_t* region_ = nullptr; @@ -116,15 +146,18 @@ private: std::deque urgent_q_, pre_q_; bool stop_ = false; int outstanding_ = 0; - int nio_ = 6, pre_k_ = 10, pre_k2_ = 0, split_ = 2; - bool prefetch_ = true, fake_ = false, pre_mru_ = false; + int ne_ = 512, nio_ = 6, pre_k_ = 10, pre_k2_ = 0, split_ = 2; + float moce_bonus_ = 0.0f; + bool in_prompt_ = false, multi_chunk_ = false; + uint64_t sep_region_bytes_ = 0; + bool prefetch_ = true, fake_ = false, pre_mru_ = false, pre_touch_ = false, split_turn_ = true; // stats uint64_t stamp_ = 1; // the handshake count: slots pinned with the current one are in use on the GPU uint64_t p_lookups_ = 0, p_hits_ = 0, p_pre_ = 0, p_misses_ = 0, p_streamed_ = 0, p_chunks_ = 0; double p_wait_ms_ = 0, p_host_ms_ = 0; uint64_t lookups_ = 0, hits_ = 0, pre_hits_ = 0, late_ = 0, misses_ = 0, pre_issued_ = 0, pre_cancel_ = 0, windows_ = 0; - std::atomic bytes_read_{0}, reads_{0}; - double wait_ms_ = 0, host_ms_ = 0, fill_s_ = 0; + std::atomic bytes_read_{0}, reads_{0}, p_bytes_read_{0}, p_reads_{0}; + double wait_ms_ = 0, host_ms_ = 0, host_after_ms_ = 0, fill_s_ = 0; uint64_t fill_bytes_ = 0; }; diff --git a/metal/src/engine/suffix_window.hpp b/metal/src/engine/suffix_window.hpp new file mode 100644 index 00000000..edc7cd61 --- /dev/null +++ b/metal/src/engine/suffix_window.hpp @@ -0,0 +1,30 @@ +#pragma once + +#include "strata/spec/suffix_drafter.hpp" + +#include +#include + +namespace strata::spec { + +inline bool suffix_stop(int32_t token) { + return token == 248046 || token == 248044; +} + +inline int suffix_limit(int requested, int remaining, int64_t pos, int64_t max_cells, bool unrestricted_qsa = true) { + const int block_drafts = unrestricted_qsa ? 3 : (3 - int(pos & 3)); + const int64_t room = std::max(0, max_cells - pos - 3); + const int64_t drafts = std::min({requested, remaining - 1, block_drafts, room}); + return int(std::max(0, drafts)); +} + +inline int suffix_keep(const int32_t* input, const int32_t* output, int nd, bool stop_on_eos = true) { + int keep = 1; + for (int i = 0; i < nd; ++i) { + if ((stop_on_eos && suffix_stop(output[i])) || output[i] != input[i + 1]) break; + ++keep; + } + return keep; +} + +} // namespace strata::spec diff --git a/metal/src/engine/window.cpp b/metal/src/engine/window.cpp index 9a6d1d32..bb9dc592 100644 --- a/metal/src/engine/window.cpp +++ b/metal/src/engine/window.cpp @@ -22,15 +22,21 @@ #include "strata/core/conversation_cache.hpp" #include "strata/core/session.hpp" #include "strata/kernels/ngram.hpp" +#include "scr_state.hpp" +#include "spec_min_p.hpp" #include +#include #include #include #include #include #include #include +#include +#include #include +#include #include #include #include @@ -51,7 +57,8 @@ namespace { constexpr int N = 2560, HC = 4, LR = 320, NFF = 640, TOPK = 10, NEXP = 512, MAXT = 8, NVOCAB_MAX = 248320; constexpr int P_MAXE = 128, P_GEXP = 16, P_GSTART = P_GEXP + P_MAXE, P_ETOK = P_GSTART + P_MAXE + 4, - P_EW = P_ETOK + P_MAXE, P_TOKENT = P_EW + P_MAXE, P_INTS = P_TOKENT + P_MAXE, P_IGU = 4, P_IDN = 8; + P_EW = P_ETOK + P_MAXE, P_TOKENT = P_EW + P_MAXE, P_SPLIT = P_TOKENT + P_MAXE, P_GSLOT = P_SPLIT + 20, P_INTS = P_GSLOT + P_MAXE, + P_IGU = 4, P_IDN = 8, P_IGU_R = P_SPLIT + 4, P_IDN_R = P_SPLIT + 8, P_IGU_M = P_SPLIT + 12, P_IDN_M = P_SPLIT + 16; struct GemvArgs { uint32_t n_in, n_out, n_tok, row_bytes, x_stride, y_stride; }; struct GrArgs { uint32_t n_embd, hc, hc_lr, n_tok, apply; float eps; }; struct MoeArgs { @@ -72,6 +79,21 @@ bool metal_ctab() { return on; } +// the drafter's attention sink (STRATA_METAL_MTP_SINK, default 64; 0: the sliding window alone) +int64_t mtp_sink_cap() { + static const int64_t cap = [] { + const char* e = std::getenv("STRATA_METAL_MTP_SINK"); + return e ? std::max(0, std::min(1024, std::atoll(e))) : 64; + }(); + return cap; +} + +std::array tile_env(const char* name) { + std::array t{0, 0, 0}; + if (const char* e = std::getenv(name)) std::sscanf(e, "%d,%d,%d", &t[0], &t[1], &t[2]); + return t; +} + int xkind(int t) { switch (t) { case ref::Q2_0: return 1; @@ -81,6 +103,78 @@ int xkind(int t) { } } +// Q4_0 with a signed-d 2-pass Lloyd-Max scale per block (codes -8..7) +void quant_q4_0_lloyd(const float* x, size_t n, uint8_t* out) { + const size_t nb = n / 32; + for (size_t b = 0; b < nb; ++b) { + const float* xb = x + b * 32; + float amax = 0.0f; + for (int j = 0; j < 32; ++j) amax = std::max(amax, std::fabs(xb[j])); + uint16_t d16 = 0; + uint8_t qs[16] = {}; + if (amax > 0.0f) { + // Signed-d 2-pass Lloyd-Max for Q4_0: codes in [-8..+7] + float best_score = -1.0f, best_qw = 0.0f, best_qq = 1.0f; + int best_q[32] = {}; + const float seeds[6] = {amax / 7.5f, amax / 8.0f, amax / 8.5f, -amax / 7.5f, -amax / 8.0f, -amax / 8.5f}; + for (int s = 0; s < 6; ++s) { + float d = seeds[s]; + for (int it = 0; it < 2; ++it) { + if (std::fabs(d) <= 1e-12f) break; + const float inv = 1.0f / d; + float qw = 0.0f, qq = 0.0f; + int qc[32]; + for (int j = 0; j < 32; ++j) { + int q = (int) std::lrintf(xb[j] * inv); + q = std::max(-8, std::min(7, q)); + qc[j] = q; + qw += (float) q * xb[j]; + qq += (float) (q * q); + } + if (qq <= 0.0f) break; + const float sc = (qw * qw) / qq; + if (sc > best_score) { + best_score = sc; best_qw = qw; best_qq = qq; + std::memcpy(best_q, qc, sizeof(qc)); + } + d = qw / qq; + } + } + d16 = ref::f32_to_f16(best_qw / best_qq); + for (int j = 0; j < 16; ++j) { + const uint8_t lo = (uint8_t) (best_q[j] + 8); + const uint8_t hi = (uint8_t) (best_q[j + 16] + 8); + qs[j] = (lo & 0x0Fu) | (uint8_t) ((hi & 0x0Fu) << 4); + } + } + std::memcpy(out + b * 18, &d16, 2); + std::memcpy(out + b * 18 + 2, qs, 16); + } +} + +// ggml's block_q5_0: d, qh, qs +void quant_q5_0(const float* x, size_t n, uint8_t* out) { + for (size_t b = 0; b < n / 32; ++b) { + const float* xb = x + b * 32; + float amax = 0.0f, mx = 0.0f; + for (int j = 0; j < 32; ++j) + if (std::fabs(xb[j]) > amax) { amax = std::fabs(xb[j]); mx = xb[j]; } + const float d = mx / -16.0f, id = d != 0.0f ? 1.0f / d : 0.0f; + const uint16_t d16 = ref::f32_to_f16(d); + uint32_t qh = 0; + uint8_t qs[16]; + for (int j = 0; j < 16; ++j) { + const int x0 = std::min(31, (int) (xb[j] * id + 16.5f)), x1 = std::min(31, (int) (xb[j + 16] * id + 16.5f)); + qs[j] = (uint8_t) ((x0 & 0xF) | ((x1 & 0xF) << 4)); + qh |= (uint32_t) ((x0 & 0x10) >> 4) << j; + qh |= (uint32_t) ((x1 & 0x10) >> 4) << (j + 16); + } + std::memcpy(out + b * 22, &d16, 2); + std::memcpy(out + b * 22 + 2, &qh, 4); + std::memcpy(out + b * 22 + 6, qs, 16); + } +} + // a prepared activation: 4 planes of MAXT rows of n, then the per-16 sums struct Act { mtl::Buffer b; @@ -112,6 +206,42 @@ const float* const* g_embd_rows = nullptr; const int32_t* g_mrope_host = nullptr; const int32_t* g_mrope_dev = nullptr; +// Anti-slop v4: a sampler shaper over the head logits, outside , and ``` fences only. +// Each position's state (the last two tokens, the region, where the line stands) follows from the one before it, so +// it is kept per position: a draft row, a rejected suffix, a prompt read from a cached prefix all start from the +// state of the position before them. Tables: the Qwen3.8-Flash-Next vocab (models/*/tokenizer/vocab.json). +#include "slop4.h" + +namespace slop4 { + +// The ends of a turn or a thought must not gain from a penalty: with a likely token pushed down, the rest of the row +// scales up by 1 / (1 - b), b the share taken away, and <|im_end|> is often next in line (a reply cut off mid-sentence). +// Their logits go down by log(1 / (1 - b)) so their probability stays what it was. lse is the row's log-sum-exp; a +// row whose penalized tokens all sit 12+ below its max loses under 6e-6 each and is skipped (no lse needed). + +inline bool needs_hold(const Pen* p, int n, const float* lg, float row_max) { + for (int i = 0; i < n; ++i) if (lg[p[i].id] > row_max - 12.0f) return true; + return false; +} + +inline int hold(const Pen* p, int n, const float* lg, double lse, Pen* out) { + double b = 0.0; + for (int i = 0; i < n; ++i) b += std::exp((double) lg[p[i].id] - lse) * (1.0 - std::exp(-(double) p[i].v)); + if (b <= 0.0) return 0; + const float v = (float) -std::log1p(-std::min(b, 0.999999)); + int m = 0; + for (int32_t id : kHold) out[m++] = {id, v}; + return m; +} + +inline Knobs knobs() { + auto f = [](const char* e, float d) { const char* v = std::getenv(e); return v ? (float) std::atof(v) : d; }; + return {f("STRATA_ANTI_SLOP_PENALTY", 9.0f), f("STRATA_ANTI_SLOP_HONEST", 6.0f), f("STRATA_ANTI_SLOP_MD", 4.0f), + f("STRATA_ANTI_SLOP_COLON", 3.0f)}; +} + +} // namespace slop4 + class Window { public: bool init(const strata::core::ModelGeometry& g, strata::core::SessionState& ss, const std::string& shard1, int max_t, @@ -120,6 +250,11 @@ public: 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; + // anti-slop v4: the prompt paths do not pass finish, so the prompt's tokens advance the per-position states here + template void slop_scan(const I* tokens, int64_t n, int64_t pos0) { + slop4::State st = slop_at(pos0 - 1); + for (int64_t i = 0; i < n; ++i) slop_put(pos0 + i) = st = slop4::step(st, (int32_t) tokens[i]); + } int64_t vocab() const { return n_vocab_; } bool ready() const { return cx_ != nullptr; } bool streaming() const { return stream_; } @@ -127,6 +262,12 @@ public: // PLE history, each QSA layer's indexer tail and tail-block key); the KV and pooled keys are positional bool checkpoint_save(strata::core::ConversationCheckpoint& c, std::string& err); bool checkpoint_restore(const strata::core::ConversationCheckpoint& c, std::string& err); + // a disk session: the running state of checkpoint_save, ss.ple_prev and the positional cells [0, n_cells) of every + // QSA layer (and the MTP layer's), read and written straight into the buffers (scr_state.hpp's file layout) + bool session_save(const std::string& path, const int32_t* ids, int64_t n_cells, std::string& err, int64_t base_cells = 0); + bool session_load(const std::string& path, std::vector& ids, int64_t& n_cells, std::string& err, const char* parent_path = nullptr); + // SCR: every QSA layer's cells [c_src, c_src + n) moved to c_dst, keys and pooled keys re-rotated by c_dst - c_src + bool scr_relocate(int64_t c_src, int64_t c_dst, int64_t n, std::string& err); void reset(); // a new sequence (session_zero): the running state zeroed void set_max_t(int t) { max_t_ = std::min(t, MAXT); } // ---- the MTP draft layer (mtp.hpp's math on mac-kern's kernels; the MTP GGUF from tools/mtp_pack.py) @@ -176,7 +317,18 @@ private: } std::map pkey_of_; void gemv(std::vector& ds, const gguf::Tensor& t, const Act& x, const mtl::Buffer& y, size_t y_stride, - int ntok, size_t y_off = 0, const mtl::Buffer* wb = nullptr); + int ntok, size_t y_off = 0, const mtl::Buffer* wb = nullptr, int tile_ntok = 0); + std::array gemv_tile(int ty, int n_in, int n_out, int ntok); + bool gemv_cat(std::vector& ds, const gguf::Tensor& ta, const gguf::Tensor& tb, const Act& x, + const mtl::Buffer& ya, const mtl::Buffer& yb, size_t y_stride, int ntok, size_t y_off); + // a deferred MoE combine: the buffer it writes (nullptr: none) and its token count + const mtl::Buffer* comb_bo_ = nullptr; + int comb_ntok_ = 0; + void comb_push(std::vector& ds, const mtl::Buffer& bo, int ntok); + void comb_defer(std::vector& ds, const mtl::Buffer& bo, int ntok); + void comb_flush(std::vector& ds); + void gr_norm(std::vector& ds, mtl::Buffer& R, const mtl::Buffer& bo, const mtl::Buffer& inj, + const mtl::Buffer& wb, size_t w_off, bool apply, int ntok, size_t r_off = 0, size_t inj_off = 0); void xprep(std::vector& ds, const mtl::Buffer& x, const Act& out, int n); void gr_read(std::vector& ds, const std::string& pre, bool apply, const Act& out, bool final_, int ntok); void layer(std::vector& ds, int l, int ntok, int64_t pos0); @@ -191,7 +343,8 @@ private: // token's embedding into memb_ row 0 (the next step's input) void mtp_pick(std::vector& ds, int j, bool embed); void gr_read_m(std::vector& ds, const std::string& pre, mtl::Buffer& Rb, const mtl::Buffer& bo, - const mtl::Buffer& inj_in, const mtl::Buffer& inj_out, bool apply, const Act& out, bool final_, int ntok); + const mtl::Buffer& inj_in, const mtl::Buffer& inj_out, bool apply, const Act& out, bool final_, int ntok, + bool already_normed = false, size_t r_off = 0, size_t inj_in_off = 0); const gguf::Tensor& TM(const std::string& name) const { const gguf::Tensor* x = fm_.find(name); if (!x) { std::fprintf(stderr, "strata metal: the MTP model has no tensor %s\n", name.c_str()); std::exit(2); } @@ -214,15 +367,19 @@ private: int64_t max_cells_ = 0; int64_t n_vocab_ = 0; float eps_ = 1e-6f; + double rope_base_ = 1e7; // buffers mtl::Buffer R_, bo_, inj_[2], xn_, lo_, qkv_, z_, al_, be_, yg_, logits_, ids_, wts_, plan_, h_, ye_, gsh_, ush_, ysh_, sg_, qfull_, kcur_, vcur_, iqr_, ikr_, qcur_, qidx_, scores_, qids_, part_, cs_, lm_, apv_, api_, tokout_, - emb_, pkey_, pval_, qkvL_, alL_, beL_, zL_, ydump_, keyn_, qn_, gated_, normW_, pleHist_; + emb_, pkey_, pval_, qkvL_, alL_, beL_, zL_, ydump_, keyn_, normW_, pleHist_; Act XA_, XF_, XY_, XQ_, XH_, XE_; int cur_inj_ = 0; // state - std::vector gst_, ghist_[2], kq_, vq_, ksc_, vsc_, tail_, dead_, pooled_; - int hp_ = 0; // which of the two history buffers holds the committed conv history (all layers) + std::vector gst_[2], ghist_[3], kq_, vq_, ksc_, vsc_, tail_, dead_, pooled_; + int sp_ = 0; // which of the two state banks holds the committed GDN state (all layers) + int hp_ = 0; // which of the history banks holds the committed conv history (all layers) + int fold_hp_ = 0; // history bank read by a pending lazy fold + int fold_off_ = 0; // starting token offset of a pending lazy fold (1 if row 0 was already committed) // commit(n) of a window of several tokens: the GDN fold of the kept prefix waits for the next window and runs at // the start of its command buffer (one submit and wait less a round). Nothing else reads the GDN state between // the two (the drafter reads the window's rows and its own caches). STRATA_METAL_LAZY_FOLD=0: at commit. @@ -233,6 +390,9 @@ private: // it measured 0.3-0.8% slower on the M5 Max (the in-place state write costs every window, the fold it saves // only the rounds with no accepted draft). Off: the window commits nothing and commit(n) folds n rows. bool partial_ = std::getenv("STRATA_METAL_PARTIAL") && std::string(std::getenv("STRATA_METAL_PARTIAL")) == "1"; + // STRATA_METAL_GDN_RING (default 1: dual2 commits row 0 in-place and row T-1 into gst_[1-sp_], eliminating fold() for + // both n_keep == 1 and n_keep == T; 2: spec1 commits row T-1 into gst_[1-sp_] while leaving gst_[sp_] untouched; 0: off). + int gdn_ring_ = std::getenv("STRATA_METAL_GDN_RING") ? std::atoi(std::getenv("STRATA_METAL_GDN_RING")) : 1; bool win_head_ = true; // the window being encoded has the head bool last_head_ = true; // the last window had it // token_chain's --fuse (mac-kern): the GDN step writes the prepared y planes itself (no xprep), and a window whose @@ -254,6 +414,8 @@ private: struct Cached { std::vector ds; std::vector cuts; // SSD streaming: the segment ends (one per layer, before its experts) + std::vector sigs; // where each layer's handshake starts (after its router): the shared + // expert runs between sigs[s] and cuts[s], beside the host's turn std::vector> qsa; // (dispatch, kind): 0 QsaArgs n_tok = T, 1 the merge (MAXT), 2 the scores }; std::map wcache_; @@ -262,8 +424,15 @@ private: std::unique_ptr st_; bool stream_ = false, st_pred_ = false; std::vector* cuts_ = nullptr; + std::vector* sigs_ = nullptr; mtl::Buffer predL_, predL2_; // 48 x MAXT x 512: layer l + 1's (l + 2's) router logits from layer l's FFN input - int seg_T_ = 0; + mtl::Buffer predW_; // the routers in Q8_0 for those guesses (empty: the GGUF's own) + mtl::Buffer routW_; // the routers as BF16 when the F32 file holds BF16 values (exact: half the bytes) + std::vector routT_; + std::vector predT_; + int seg_T_ = 0, n_exp_ = NEXP; + uint32_t moce_pad_ = 0; + int pg_lib_ = -1, pg_lgdn_ = -1, pg_lqsa_ = -1; bool stream_arena(std::string& err); double run_ds(Cached& c); // prompts on the matrix units (prompt_chunk below) @@ -276,9 +445,14 @@ private: public: bool prompt_gemm_on(int64_t n) const; bool prompt_gemm_init(std::string& err); + void free_prompt_gemm(); bool prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::string& err); int prompt_chunk_rows() const; - void prompt_region(bool hold) { if (stream_ && st_ && st_->ready()) st_->hold_region(hold); } + void prompt_region(bool hold, bool multi_chunk = false); + uint64_t prompt_gemm_need_bytes() const; + // decode-only on purpose: the prompt (p_*) counters survive so the final report still carries prefill + // accounting, which window_reset_stats would otherwise erase before anything prints it. + void reset_stream_stats() { if (st_) st_->reset_decode_stats(); for (int t = 0; t <= MAXT; ++t) { win_n_[t] = 0; win_ms_[t] = win_gpu_[t] = 0; } } void stream_report(FILE* f) const { if (!st_) return; st_->report(f); @@ -289,17 +463,136 @@ public: } int win_n_[MAXT + 1] = {}; double win_ms_[MAXT + 1] = {}, win_gpu_[MAXT + 1] = {}; + void dump_moe_layer(const char* pfx, int l, bool dump_in, bool dump_out); + std::thread* ple_join_ = nullptr; // a prompt chunk's PLE gather, joined at layer 0's handshake void at_cut(int s) { + if (s == 0 && ple_join_) { + if (ple_join_->joinable()) ple_join_->join(); + ple_join_ = nullptr; + } + if (cut_ids_ && pg_) { + if (const char* pfx = std::getenv("STRATA_METAL_DUMP_MOE")) { + if (s > 0) dump_moe_layer(pfx, s - 1, false, true); + dump_moe_layer(pfx, s, true, false); + } + } st_->at_layer(s, seg_T_, TOPK, cut_ids_ ? cut_ids_ : ids_.as(), st_pred_ && !cut_ids_ ? predL_.as() + (size_t) s * MAXT * 512 : nullptr, predL2_.ptr && !cut_ids_ ? predL2_.as() + (size_t) s * MAXT * 512 : nullptr); + if (cut_ids_) st_->after_layer(s, seg_T_, nullptr); // the prompt chain has no after_cut: its bookkeeping here + } + // the pipelines of every window size up to max_t_ and of the MTP layer, made at load (window_init) rather than when + // a window first needs them (a T=3 window's first build compiled its #3 specializations in the middle of a decode: + // 282 ms; the first request's T=1 window and MTP prefill / draft stalled the same way). Past 2051, so the QSA + // selection's pipelines too (a superset of a short context's: layer() makes qsa_score / qsa_topk either way). + int prewarmed_t_ = 0; // window sizes 1..prewarmed_t_ made + bool mtp_prewarmed_ = false; + void prewarm_pipelines(bool overlap) { + auto* rec = rec_; + rec_ = nullptr; // the QSA patch list belongs to the window being built + std::vector c, s; + for (int t = prewarmed_t_ + 1; t <= max_t_; ++t) { + std::vector ds; + c.clear(); s.clear(); + cuts_ = stream_ ? &c : nullptr; + sigs_ = stream_ && overlap ? &s : nullptr; + nb_ = false; + cur_inj_ = 0; + fold(ds, t, 0); + if (ple_on_) { + xprep(ds, emb_, XE_, N); + gemv(ds, T("blk.1.ple_key.weight"), XE_, pkey_, 10240, t); + gemv(ds, T("blk.1.ple_value.weight"), XE_, pval_, N, t); + ple_key_norm(ds, t); + } + layer(ds, 0, t, 4096); + if (ple_on_) ple_after_layer0(ds, t); + for (int l = 1; l < 48; ++l) layer(ds, l, t, 4096); + gr_read(ds, "output_hc", true, XA_, true, t); + gemv(ds, T("output.weight"), XA_, lm_, (size_t) n_vocab_, t); + } + prewarmed_t_ = std::max(prewarmed_t_, max_t_); + P(lsmp_, "argmax_part", {}, "argmax_part"); + P(lsmp_, "argmax_final", {}, "argmax_final"); + cuts_ = sigs_ = nullptr; + if (mtp_on_ && !mtp_prewarmed_) { + std::vector ds; + nb_ = false; + for (int k = 1; k <= MAXT; ++k) mtp_encode(ds, mR_, 0, k, nullptr, 0, false); // mtp_prefill's chunks + for (int k = 1; k <= 4; ++k) mtp_encode(ds, mR_, 0, k, nullptr, 0, true); // draft step & combined catch-up + step 0 + if (chain_) mtp_pick(ds, 0, true); + mtp_prewarmed_ = true; + } + rec_ = rec; + cur_inj_ = 0; + nb_ = false; + comb_bo_ = nullptr; + comb_ntok_ = 0; + } + // window_init's, else the first window's (STRATA_METAL_LAZY_ARENA: the MoE dispatches bind the arena, made then) + void prewarm() { + static const bool on = !(std::getenv("STRATA_METAL_PREWARM") && std::atoi(std::getenv("STRATA_METAL_PREWARM")) == 0); + static const bool overlap = !(std::getenv("STRATA_METAL_SHEXP_OVERLAP") && std::atoi(std::getenv("STRATA_METAL_SHEXP_OVERLAP")) == 0); + if (!on || (stream_ && !st_->ready())) return; + if (prewarmed_t_ < max_t_ || (mtp_on_ && !mtp_prewarmed_)) prewarm_pipelines(overlap); + } + // the PLE block on the GPU (msl/ple.metal) around layer 0: the key norm reads only the key projection, so it runs + // beside layer 0's first read, not behind layer 0 + struct PleArgs { uint32_t n_embd, hc; float eps; }; + void ple_key_norm(std::vector& ds, int Tn) { + const PleArgs pa{N, HC, strata::kernels::NG_RMS_EPS}; + mtl::Dispatch d; + d.p = P(lple_, "ple_gnorm", {}, "ple_gnorm"); + d.args = {mtl::Arg::B(pkey_), mtl::Arg::B(W_, T("blk.1.ple_norm_key.weight").offset), mtl::Arg::B(keyn_), mtl::Arg::V(pa)}; + d.grid[0] = HC; d.grid[1] = Tn; d.tg[0] = 256; + push(ds, d); + nb_ = true; + } + // after layer 0: its FFN write applied to R, the query norm, the gate, the normalized rows, the conv. Layer 0's MoE + // combine deferred: ple_comb_query_gate_norm forms it as it reads R (STRATA_METAL_PLE_COMB=0: gr_norm_comb, whose + // xn nothing reads, then ple_query_gate_norm) + void ple_after_layer0(std::vector& ds, int Tn) { + static const bool fuse = !(std::getenv("STRATA_METAL_PLE_COMB") && std::atoi(std::getenv("STRATA_METAL_PLE_COMB")) == 0); + const PleArgs pa{N, HC, strata::kernels::NG_RMS_EPS}; + mtl::Dispatch g, cv; + g.args = {mtl::Arg::B(keyn_), mtl::Arg::B(R_), mtl::Arg::B(W_, T("blk.1.ple_norm_query.weight").offset), + mtl::Arg::B(pval_), mtl::Arg::B(W_, T("blk.1.ple_norm_conv.weight").offset), mtl::Arg::B(normW_), + mtl::Arg::V(pa)}; + if (fuse && comb_bo_ == &bo_ && comb_ntok_ == Tn) { + comb_bo_ = nullptr; + const uint32_t tk = TOPK; + g.p = P(lple_, "ple_comb_query_gate_norm", {}, "ple_comb_query_gate_norm"); + g.args.insert(g.args.end(), {mtl::Arg::B(ye_), mtl::Arg::B(inj_[cur_inj_]), mtl::Arg::B(plan_), mtl::Arg::B(ysh_), + mtl::Arg::B(sg_), mtl::Arg::V(tk)}); + } else { + gr_norm(ds, R_, bo_, inj_[cur_inj_], W_, T("blk.1.hc_attn_norm.weight").offset, true, Tn); + g.p = P(lple_, "ple_query_gate_norm", {}, "ple_query_gate_norm"); + } + g.grid[0] = HC; g.grid[1] = Tn; g.tg[0] = 256; + push(ds, g); + cv.p = P(lple_, "ple_conv_res", {}, "ple_conv_res"); + cv.args = {mtl::Arg::B(pleHist_), mtl::Arg::B(normW_), kw_arg(), mtl::Arg::B(R_), mtl::Arg::V(pa)}; + cv.grid[0] = HC * N / 256; cv.grid[1] = Tn; cv.tg[0] = 256; + push(ds, cv); + } + void after_cut(int s) { // decode only (the prompt chain runs without it) + st_->after_layer(s, seg_T_, st_pred_ ? predL_.as() + (size_t) s * MAXT * 512 : nullptr, + predL2_.ptr ? predL2_.as() + (size_t) s * MAXT * 512 : nullptr); } private: std::vector ple_hist_; // NG_HIST x hc_dim normalized rows, row-fastest bool ple_on_ = false; int last_T_ = 0; + int64_t last_pos0_ = 0; + std::vector tail_bak_; double last_gpu_ms_ = 0; int last_n_disp_ = 0; int32_t last_tokens_[MAXT] = {}; + std::vector slop_; // anti-slop v4: the state after the token at each position (State{} = unknown) + slop4::State slop_at(int64_t p) const { return p >= 0 && (size_t) p < slop_.size() ? slop_[(size_t) p] : slop4::State{}; } + slop4::State& slop_put(int64_t p) { + if ((size_t) p >= slop_.size()) slop_.resize(((size_t) p + 4096) & ~(size_t) 4095); + return slop_[(size_t) p]; + } std::vector last_norm_; // the window's normalized PLE rows (T x hc_dim), for commit // STRATA_METAL_LOGPOS= (test tooling, one-token windows): per decode position the log-probability of the // token that came next (the next window's input: with --spec-follow, the forced text) and the top 64, in the @@ -321,13 +614,26 @@ public: ~Window(); void prof_add(bool draft, double ms) { Prof& x = draft ? dprof_ : cprof_; ++x.n; x.all += ms; } private: - // the Q8_0 copies of the hyper-connection matrices (the default; STRATA_HC_Q8=0: the BF16 originals) - bool hc_q8_ = false; + // the Q8_0 (or Q4_0 when STRATA_HC_Q4=1) copies of the hyper-connection matrices + bool hc_q8_ = false, hc_q4_ = false; // hc_q4_: some HC matrices below 8 bits (q8type_) + int hc_q4g_ = 0, hc_q5g_ = 0; // their groups: bit 0 hc_attn, bit 1 hc_ffn mtl::Buffer Q8_; std::map q8off_; + std::map q8type_; // MTP bool mtp_on_ = false; int lmtp_ = -1; + // the drafter's cells of committed tokens: the run mtp_cell0_..mtp_cell1_-1 (the sliding window's) and the attention + // sink 0..mtp_sink_-1 (at most mtp_sink_cap() cells, read beside the window once it has moved past them) + int64_t mtp_cell0_ = 0, mtp_cell1_ = 0, mtp_sink_ = 0; + void mtp_wrote(int64_t c0, int64_t n) { // cells c0..c0+n-1 now hold committed tokens; the ones past them are stale + mtp_cell1_ = std::min(mtp_cell1_, c0); + mtp_sink_ = std::min(mtp_sink_, c0); + if (n <= 0) return; + if (c0 != mtp_cell1_ || c0 < mtp_cell0_) mtp_cell0_ = c0; + mtp_cell1_ = c0 + n; + if (c0 == mtp_sink_) mtp_sink_ = std::min(mtp_sink_cap(), c0 + n); + } gguf::File fm_; mtl::Buffer Wm_, mnorm_, mR_, memb_, men_, me2_, mhn_, mh2_, midx_, mbo_, minj_[2], mqids_, mlm_; std::map mnorm_off_; @@ -373,6 +679,8 @@ private: } mtl::Constants qsa_c() const { return mrope_on() ? mtl::Constants().b(24, true) : mtl::Constants(); } std::string qsa_k() const { return mrope_on() ? "qsa_prep#mrope" : "qsa_prep"; } + // qsa_prep's threadgroup: one simdgroup a row (qsa_prep_sg) from 3 tokens, below it the 256-thread rows win + static uint32_t qsa_prep_tg(int64_t ntok) { return ntok >= 3 ? 32 : 256; } int m5_layout_ = -1; // the M5's matmul2d layout (coop_k4), asked once: the M4's tiles differ bool m5_layout() { if (m5_layout_ < 0) m5_layout_ = coop_k4() ? 1 : 0; return m5_layout_ == 1; } bool pc_setup(int C); @@ -381,7 +689,16 @@ private: // per step with the host's pick and embedding (as before) int lchain_ = -1; bool chain_ = false; - mtl::Buffer dout_, dprob_, posv_, dmap_; + mtl::Buffer dout_, dprob_, posv_, dmap_, dslop_, dinverse_; + bool draft_slop_ = false; + struct DraftSlopArgs { + slop4::State state; + slop4::Knobs knobs; + uint32_t nv, j, vocab; + }; + DraftSlopArgs draft_slop_args_{}; + void mtp_slop(std::vector& ds, int j); + void mtp_slop_host(slop4::State state); }; namespace { @@ -411,10 +728,21 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta ss_ = &ss; max_t_ = std::min(max_t, MAXT); max_cells_ = max_cells; - if (g.n_embd != N || g.hc != HC || g.hc_lr != LR || g.n_ff != NFF || g.n_expert != NEXP || g.n_layers != 48) { + if (g.n_embd != N || g.hc != HC || g.hc_lr != LR || g.n_ff != NFF || g.n_expert < 1 || g.n_expert > NEXP || g.n_layers != 48) { err = "metal window: the geometry differs from Qwen3.8-Flash-Next's"; return false; } + n_exp_ = (int) g.n_expert; + if (const char* mb = std::getenv("STRATA_METAL_MOCE_BONUS")) { + const float bonus = (float) std::atof(mb); + if (bonus > 0.0f) { + const float sc = std::exp(bonus); + std::memcpy(&moce_pad_, &sc, 4); + } + } + if (!(std::getenv("STRATA_METAL_KEEP_HOST_WEIGHTS") && std::string(std::getenv("STRATA_METAL_KEEP_HOST_WEIGHTS")) == "1")) { + strata::gpu::purge_device_allocations(1u << 20); + } eps_ = 1e-6f; const double t0 = now_ms(); if (const long rp = repack_pid_running(shard1); rp > 0 && !(std::getenv("STRATA_METAL_IGNORE_REPACK") && std::string(std::getenv("STRATA_METAL_IGNORE_REPACK")) == "1")) { @@ -423,6 +751,9 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta return false; } if (!f_.open(shard1, err)) return false; + if (const gguf::Tensor* ge0 = f_.find("blk.0.ffn_gate_exps.weight")) { + if (ge0->shape.size() >= 3 && ge0->shape[2] >= 1 && ge0->shape[2] <= NEXP) n_exp_ = (int) ge0->shape[2]; + } cx_ = std::make_unique(); stream_ = ExpertStream::wanted(f_, *cx_); if (stream_) { @@ -460,13 +791,14 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta // the GDN layers' window inputs, kept per layer so commit() can fold the accepted prefix into the state qkvL_ = A(4 * (size_t) 36 * MAXT * 10240); alL_ = A(4 * (size_t) 36 * MAXT * 48); beL_ = A(4 * (size_t) 36 * MAXT * 48); zL_ = A(4 * (size_t) 36 * MAXT * 6144); ydump_ = A(4 * MAXT * 6144); - keyn_ = A(4 * MAXT * 10240); qn_ = A(4 * MAXT * 10240); gated_ = A(4 * MAXT * 10240); normW_ = A(4 * MAXT * 10240); + keyn_ = A(4 * MAXT * 10240); normW_ = A(4 * MAXT * 10240); mrope_ = A(12 * (size_t) (MAXC + 16)); pleHist_ = A(4 * 9 * 10240); auto act = [&](int n) { Act a; a.n = n; a.b = A(4 * ((size_t) 4 * MAXT * n + (size_t) MAXT * n / 16)); return a; }; XA_ = act(N); XF_ = act(N); XY_ = act(6144); XQ_ = act(6144); XH_ = act(NFF); XE_ = act(N); // RoPE: cos/sin [pos][32] for the first 64 dims, NEOX pairs (i, i + 32), in double as build_rope_table const double base = f_.meta_num.count("qwen4exp.rope.freq_base") ? f_.meta_num["qwen4exp.rope.freq_base"] : 1e7; + rope_base_ = base; for (int p = 0; p < MAXC + 16; ++p) for (int i = 0; i < 32; ++i) { const double th = (double) p * std::pow(base, -2.0 * i / 64.0); @@ -477,9 +809,11 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta for (int c = 0; c < 2051; ++c) qids_.as()[t * 2051 + c] = c; for (int l = 0; l < 48; ++l) { if (l % 4 != 3) { - gst_.push_back(A(4 * 128 * 48 * 128)); + gst_[0].push_back(A(4 * 128 * 48 * 128)); + if (gdn_ring_ > 0) gst_[1].push_back(A(4 * 128 * 48 * 128)); ghist_[0].push_back(A(4 * 10240 * 3)); ghist_[1].push_back(A(4 * 10240 * 3)); + if (gdn_ring_ > 0) ghist_[2].push_back(A(4 * 10240 * 3)); } else { kq_.push_back(A((size_t) MAXC * 2 * 256)); vq_.push_back(A((size_t) MAXC * 2 * 256)); @@ -495,19 +829,37 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta { const char* e = std::getenv("STRATA_HC_Q8"); hc_q8_ = !(e && std::string(e) == "0"); + // STRATA_HC_Q4 / STRATA_HC_Q5 = 1 (every hc_*_down/_up) | attn | ffn: those in Q4_0 / Q5_0 (Q4 wins where both) + auto groups = [](const char* name) { + const char* e = std::getenv(name); + const std::string v = e ? e : ""; + return v == "attn" ? 1 : v == "ffn" ? 2 : (!v.empty() && std::atoi(v.c_str()) != 0) ? 3 : 0; + }; + hc_q4g_ = hc_q8_ ? groups("STRATA_HC_Q4") : 0; + hc_q5g_ = hc_q8_ ? groups("STRATA_HC_Q5") : 0; + hc_q4_ = hc_q4g_ || hc_q5g_; } if (hc_q8_) { const double tq = now_ms(); size_t tot = 0; std::vector names; + auto low_type = [&](const std::string& nm) { // Q4_0, Q5_0 or Q8_0 for this HC matrix + if (!hc_q4_) return (int) ref::Q8_0; + if (nm.find("output_hc") != std::string::npos || nm.find("_inject") != std::string::npos) return (int) ref::Q8_0; + if (nm.find("_down.weight") == std::string::npos && nm.find("_up.weight") == std::string::npos) return (int) ref::Q8_0; + const int g = nm.find("hc_attn") != std::string::npos ? 1 : nm.find("hc_ffn") != std::string::npos ? 2 : 0; + return (hc_q4g_ & g) ? (int) ref::Q4_0 : (hc_q5g_ & g) ? (int) ref::Q5_0 : (int) ref::Q8_0; + }; for (const auto& kv : f_.tensors) { const gguf::Tensor& x = kv.second; if (x.type != ref::BF16 || x.name.find("hc_") == std::string::npos || x.name.find("_norm") != std::string::npos) continue; uint64_t ne = 1; for (uint64_t d : x.shape) ne *= d; + const int qt = low_type(x.name); q8off_[x.name] = tot; + q8type_[x.name] = qt; names.push_back(x.name); - tot += ne / 32 * 34; + tot += ne / 32 * (qt == ref::Q4_0 ? 18 : qt == ref::Q5_0 ? 22 : 34); } Q8_ = cx_->alloc((tot + 4095) / 4096 * 4096 + 4096); // + a page: a file with Q8_0 HC already has no copies std::vector th; @@ -528,12 +880,18 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta src = raw.data(); } ref::dequant(ref::BF16, src, tmp.data(), ne); - ref::quantize_q8_0(tmp.data(), ne, Q8_.as() + q8off_[names[i]]); + if (q8type_.at(names[i]) == ref::Q4_0) + quant_q4_0_lloyd(tmp.data(), ne, Q8_.as() + q8off_[names[i]]); + else if (q8type_.at(names[i]) == ref::Q5_0) + quant_q5_0(tmp.data(), ne, Q8_.as() + q8off_[names[i]]); + else + ref::quantize_q8_0(tmp.data(), ne, Q8_.as() + q8off_[names[i]]); } }); for (auto& x : th) x.join(); - std::fprintf(stderr, "strata metal: HC read in Q8_0 (STRATA_HC_Q8=0: BF16): %zu matrices, %.0f MB, made in %.0f ms\n", - names.size(), tot / 1e6, now_ms() - tq); + std::fprintf(stderr, "strata metal: HC read in %s: %zu matrices, %.0f MB, made in %.0f ms\n", + hc_q4_ ? "Q8_0 with Q4_0/Q5_0 (STRATA_HC_Q4/_Q5)" : "Q8_0 (STRATA_HC_Q8=0: BF16)", names.size(), tot / 1e6, + now_ms() - tq); } if (stream_) { std::vector skip; // the BF16 originals of the Q8_0 HC copies are not read again @@ -543,9 +901,51 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta const char* e = std::getenv("STRATA_METAL_PREFETCH"); st_pred_ = !(e && std::atoi(e) <= 0); if (st_pred_) predL_ = A(4 * (size_t) 48 * MAXT * 512); + // the read-ahead's router copies in Q8_0 (STRATA_METAL_PRED_Q8=0: the F32 originals): the guess only picks what is + // read ahead, never the routing, and an F32 router is 2.6 MB a layer read twice a token (92 MB a token saved) + const char* pq = std::getenv("STRATA_METAL_PRED_Q8"); + if (st_pred_ && !(pq && std::atoi(pq) == 0) && T("blk.1.ffn_gate_inp.weight").type == ref::F32) { + const size_t rb = ref::row_bytes(ref::Q8_0, N), per = rb * (size_t) n_exp_; + predW_ = A(per * 48); + predT_.resize(48); + for (int l = 0; l < 48; ++l) { + const gguf::Tensor& src = T("blk." + std::to_string(l) + ".ffn_gate_inp.weight"); + ref::quantize_q8_0((const float*) ((const uint8_t*) W_.ptr + src.offset), (size_t) N * n_exp_, + predW_.as() + per * (size_t) l); + predT_[(size_t) l] = src; + predT_[(size_t) l].type = ref::Q8_0; + predT_[(size_t) l].offset = per * (size_t) l; + predT_[(size_t) l].bytes = per; + } + } if (st_pred_ && st_->ahead2() > 0) predL2_ = A(4 * (size_t) 48 * MAXT * 512); } - ple_on_ = ss.ple.ready() && f_.find("blk.1.ple_key.weight") != nullptr; + { // the routers: an F32 router whose values are all BF16 (the upcast originals) is read as BF16, bit for bit the same + // weights at half the bytes (63 MB a token). STRATA_METAL_ROUTER_BF16=0: the F32 tensors. + const char* rb = std::getenv("STRATA_METAL_ROUTER_BF16"); + bool exact = !(rb && std::atoi(rb) == 0) && T("blk.0.ffn_gate_inp.weight").type == ref::F32; + for (int l = 0; exact && l < 48; ++l) { + const gguf::Tensor& t = T("blk." + std::to_string(l) + ".ffn_gate_inp.weight"); + const uint32_t* p = (const uint32_t*) ((const uint8_t*) W_.ptr + t.offset); + for (size_t i = 0; exact && i < t.bytes / 4; ++i) exact = (p[i] & 0xFFFFu) == 0; + } + if (exact) { + const size_t per = (size_t) N * n_exp_ * 2; + routW_ = cx_->alloc(per * 48); + routT_.resize(48); + for (int l = 0; l < 48; ++l) { + const gguf::Tensor& t = T("blk." + std::to_string(l) + ".ffn_gate_inp.weight"); + const uint32_t* p = (const uint32_t*) ((const uint8_t*) W_.ptr + t.offset); + uint16_t* q = (uint16_t*) (routW_.as() + per * (size_t) l); + for (size_t i = 0; i < t.bytes / 4; ++i) q[i] = (uint16_t) (p[i] >> 16); + routT_[(size_t) l] = t; + routT_[(size_t) l].type = ref::BF16; + routT_[(size_t) l].offset = per * (size_t) l; + routT_[(size_t) l].bytes = per; + } + } + } + ple_on_ = (ss.ple.ready() || (ss.ple.table != nullptr && ss.ple.table->is_open())) && f_.find("blk.1.ple_key.weight") != nullptr; if (ple_on_ && T("blk.1.ple_conv1d.weight").type == ref::F32) { const gguf::Tensor& k = T("blk.1.ple_conv1d.weight"); const size_t ne = k.bytes / 4; @@ -563,14 +963,64 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta return true; } -void Window::gemv(std::vector& ds, const gguf::Tensor& t, const Act& x, const mtl::Buffer& y, - size_t y_stride, int ntok, size_t y_off, const mtl::Buffer* wb) { - const int n_in = (int) t.shape[0], n_out = t.shape.size() > 1 ? (int) t.shape[1] : 1; - const int ty = t.type; +// gemv()'s tile for a matrix: {LPR, NR, SG, KS} +std::array Window::gemv_tile(int ty, int n_in, int n_out, int ntok) { int L = 16, Rr = 1, S = 4, K = 1; - if (n_in >= 6144) { L = 32; Rr = 1; S = 8; K = 2; } - else if (n_in == 640) { L = 8; Rr = 1; S = 4; K = 1; } - else if (n_out <= 640 || ty == ref::Q2_0 || ty == ref::BF16) { L = 32; Rr = 1; S = 4; K = 1; } + const bool m5 = m5_layout(); + if (n_in >= 6144) { + if (m5 && ty == ref::IQ4_XS) { L = (ntok >= 2 && ntok <= 3) ? 16 : 32; Rr = 2; S = 4; K = 1; } + else if (m5 && ty == ref::IQ3_S) { + if (ntok >= 2) { L = 8; Rr = 1; S = 8; K = 1; } + else { L = 32; Rr = 2; S = 4; K = 1; } + } else if (m5 && ty == ref::Q6_K) { L = ntok >= 2 ? 32 : 16; Rr = ntok >= 4 ? 2 : 1; S = 4; K = 1; } + else { L = 32; Rr = 1; S = 8; K = 2; } + } else if (n_in == 640) { + // M5, weights from DRAM, 640 -> 2560: Q2_0 at 4 tokens 8,2,4,1 -4..-6% (the same sums as 8,1,4,1) + if (m5 && ty == ref::IQ4_NL && ntok >= 2) { L = 4; Rr = 2; S = 4; K = 1; } + else if (m5 && ty == ref::Q8_0) { L = ntok >= 3 ? 4 : 32; Rr = 1; S = 4; K = 1; } + else { L = 8; Rr = m5 && ty == ref::Q2_0 && ntok >= 4 ? 2 : 1; S = 4; K = 1; } + } else if (n_out <= 640) { + // M5, weights from DRAM, 2560 -> 512 (QSA k, v) against 32,1,4,1: 32,2,4,1 keeps LPR=32 (100% bitwise identical) + // and cuts IQ4_XS by -3.2%..-25.5%, IQ3_S by -9.0%..-24.4%, and Q6_K by -7.6%..-13.9%. + if (m5 && ty == ref::IQ4_XS && n_out == 640) { L = 16; Rr = 2; S = 4; K = 1; } + else if (m5 && (ty == ref::IQ4_XS || ty == ref::IQ3_S || ty == ref::Q6_K) && n_out == 512) { L = 32; Rr = 2; S = 4; K = 1; } + else if (m5 && ty == ref::IQ3_S && n_out == 640) { + if (ntok >= 2) { L = 8; Rr = 1; S = 8; K = 1; } + else { L = 32; Rr = 2; S = 4; K = 1; } + } else if (m5 && ty == ref::Q6_K && n_out == 640) { + if (ntok >= 2) { L = 8; Rr = 1; S = 4; K = 1; } + else { L = 32; Rr = 2; S = 4; K = 1; } + } else if (m5 && ty == ref::Q2_0 && n_out == 640) { + if (ntok >= 2) { L = 16; Rr = 1; S = 4; K = 1; } + else { L = 32; Rr = 2; S = 4; K = 1; } + } else if (m5 && ty == ref::Q8_0 && n_out == 512) { + if (ntok >= 2) { L = 16; Rr = 1; S = 4; K = 1; } + else { L = 32; Rr = 1; S = 4; K = 4; } + } else if (m5 && ty == ref::BF16 && n_out >= 128 && n_out <= 512) { + if (ntok >= 3 && n_out == 128) { L = 16; Rr = 1; S = 4; K = 1; } + else { L = 32; Rr = 1; S = 8; K = 8; } + } else { L = 32; Rr = 1; S = 4; K = 1; } + } else if (ty == ref::Q2_0 || ty == ref::BF16) { + L = 32; Rr = 1; S = 4; K = 1; + } else if (m5 && ty == ref::Q8_0) { + if (ntok == 1) { L = 32; Rr = 1; S = 4; K = 4; } + else if (ntok <= 3) { L = 32; Rr = 2; S = 4; K = 1; } + else { L = 16; Rr = 1; S = 4; K = 1; } + } else if (m5 && ty == ref::IQ4_XS) { + if (ntok >= 2) { L = 16; Rr = 2; S = 4; K = 1; } + else { L = 32; Rr = 2; S = 4; K = 1; } + } else if (m5 && ty == ref::IQ3_S && n_out >= 6144) { + if (ntok >= 2) { L = 8; Rr = 1; S = 4; K = 1; } + else { L = 32; Rr = 2; S = 4; K = 1; } + } else if (m5 && ty == ref::Q6_K && n_out >= 6144) { + if (ntok == 1) { L = 32; Rr = 1; S = 8; K = 2; } + else if (ntok == 2) { L = 32; Rr = 2; S = 4; K = 1; } + else if (n_out == 6144) { L = 16; Rr = 1; S = 4; K = 1; } + else { L = 8; Rr = 1; S = 4; K = 1; } + } else if (m5 && (ty == ref::Q4_K || ty == ref::Q5_K || ty == ref::Q3_K) && n_out >= 6144) { + if (ntok >= 2) { L = 8; Rr = 1; S = 8; K = 1; } + else if (ty != ref::Q3_K) { L = 32; Rr = 1; S = 4; K = 1; } + } { // windows of several tokens: NR rows a lane share each x load (the window's cost is its x traffic). The sums // run in the same order for any NR, so the answers do not change. STRATA_METAL_GEMV_NR (1, 2): the tile. @@ -581,6 +1031,15 @@ void Window::gemv(std::vector& ds, const gguf::Tensor& t, const A if (K == 1 || (L == 32 && K == 2)) Rr = 2; } } + return {L, Rr, S, K}; +} + +void Window::gemv(std::vector& ds, const gguf::Tensor& t, const Act& x, const mtl::Buffer& y, + size_t y_stride, int ntok, size_t y_off, const mtl::Buffer* wb, int tile_ntok) { + // tile_ntok > 0: the tile of a window that wide, so these ntok tokens' answers are bitwise its first tokens' + const int n_in = (int) t.shape[0], n_out = t.shape.size() > 1 ? (int) t.shape[1] : 1; + const int ty = t.type; + const auto [L, Rr, S, K] = gemv_tile(ty, n_in, n_out, tile_ntok > 0 ? tile_ntok : ntok); const std::string fn = "gemv_" + std::to_string(ty) + "_" + std::to_string(L) + "_" + std::to_string(Rr) + "_" + std::to_string(S) + "_" + std::to_string(K); const mtl::Pipeline* p = P(lib_, fn, mtl::Constants().i(0, ntok), fn + "#" + std::to_string(ntok)); @@ -596,6 +1055,37 @@ void Window::gemv(std::vector& ds, const gguf::Tensor& t, const A push(ds, d); } +// two GEMVs of the same format and shape over the same x in one dispatch (gemv_cat): bitwise the two gemv() calls when +// gemv() would run them KS = 1 (a row's sums depend only on LPR). False when not instantiated / not applicable. +// STRATA_METAL_GEMV_CAT=0: off. +bool Window::gemv_cat(std::vector& ds, const gguf::Tensor& ta, const gguf::Tensor& tb, const Act& x, + const mtl::Buffer& ya, const mtl::Buffer& yb, size_t y_stride, int ntok, size_t y_off) { + static const bool on = !(std::getenv("STRATA_METAL_GEMV_CAT") && std::atoi(std::getenv("STRATA_METAL_GEMV_CAT")) == 0); + if (!on || ta.type != tb.type || (ta.type != ref::BF16 && ta.type != ref::F32 && ta.type != ref::Q8_0) || ta.shape != tb.shape) + return false; + const int n_in = (int) ta.shape[0], n_out = ta.shape.size() > 1 ? (int) ta.shape[1] : 1; + const auto [L, Rr, S, K] = gemv_tile(ta.type, n_in, n_out, ntok); + // gemv.metal's GVC instances + const bool inst = K == 1 && ((L == 32 && Rr == 1 && (S == 4 || S == 8)) || (L == 16 && Rr == 1 && S == 4) || + (L == 32 && Rr == 2 && S == 4) || (ta.type == ref::Q8_0 && L == 8 && Rr == 1 && S == 8)); + if (!inst) return false; + const std::string fn = "gemvc_" + std::to_string(ta.type) + "_" + std::to_string(L) + "_" + std::to_string(Rr) + "_" + + std::to_string(S); + GemvArgs a{(uint32_t) n_in, (uint32_t) (2 * n_out), (uint32_t) ntok, (uint32_t) ref::row_bytes(ta.type, n_in), + (uint32_t) n_in, (uint32_t) y_stride}; + const uint32_t sp = (uint32_t) n_out; + mtl::Dispatch d; + d.p = P(lib_, fn, mtl::Constants().i(0, ntok), fn + "#" + std::to_string(ntok)); + d.args = {mtl::Arg::B(W_, ta.offset), mtl::Arg::B(x.b, x.plane_bytes() * xkind(ta.type)), + mtl::Arg::B(x.b, x.plane_bytes() * 4), mtl::Arg::B(ya, y_off), mtl::Arg::V(a), mtl::Arg::B(W_, tb.offset), + mtl::Arg::B(yb, y_off), mtl::Arg::V(sp)}; + const int rpt = S * (32 / L) * Rr; + d.grid[0] = (2 * n_out + rpt - 1) / rpt; + d.tg[0] = S * 32; + push(ds, d); + return true; +} + void Window::xprep(std::vector& ds, const mtl::Buffer& x, const Act& out, int n) { mtl::Dispatch d; d.p = P(lib_, "xprep", {}, "xprep"); @@ -606,6 +1096,58 @@ void Window::xprep(std::vector& ds, const mtl::Buffer& x, const A push(ds, d); } +// the MoE combine (moe_combine_sh) of the last layer into bo, pushed now +void Window::comb_push(std::vector& ds, const mtl::Buffer& bo, int ntok) { + GlueArgs a{N, (uint32_t) ntok, N, TOPK}; + mtl::Dispatch d; + d.p = P(lib_, "moe_combine_sh", {}, "moe_combine_sh"); + d.args = {mtl::Arg::B(ye_), mtl::Arg::B(plan_), mtl::Arg::B(ysh_), mtl::Arg::B(sg_), mtl::Arg::B(bo), mtl::Arg::V(a)}; + d.grid[0] = (ntok * N + 255) / 256; + d.tg[0] = 256; + push(ds, d); +} + +// the MoE combine into bo, left to the next gr_norm that reads bo (gr_norm_comb: bitwise moe_combine_sh + gr_norm; +// one dispatch, one barrier and bo's round trip less). STRATA_METAL_GR_COMB=0: the combine's own dispatch. +void Window::comb_defer(std::vector& ds, const mtl::Buffer& bo, int ntok) { + static const bool on = !(std::getenv("STRATA_METAL_GR_COMB") && std::atoi(std::getenv("STRATA_METAL_GR_COMB")) == 0); + comb_flush(ds); + if (!on) { comb_push(ds, bo, ntok); return; } + comb_bo_ = &bo; + comb_ntok_ = ntok; +} + +void Window::comb_flush(std::vector& ds) { + if (!comb_bo_) return; + const mtl::Buffer* bo = comb_bo_; + comb_bo_ = nullptr; + comb_push(ds, *bo, comb_ntok_); +} + +// gr_norm (R' = R + bo * 2 sigmoid(inj / hc) when `apply`, xn), or gr_norm_comb when bo is a deferred MoE combine +void Window::gr_norm(std::vector& ds, mtl::Buffer& R, const mtl::Buffer& bo, const mtl::Buffer& inj, + const mtl::Buffer& wb, size_t w_off, bool apply, int ntok, size_t r_off, size_t inj_off) { + const bool comb = apply && comb_bo_ == &bo && comb_ntok_ == ntok; + if (comb) comb_bo_ = nullptr; + else comb_flush(ds); + GrArgs a{N, HC, LR, (uint32_t) ntok, apply ? 1u : 0u, eps_}; + mtl::Dispatch d; + if (comb) { + static_assert(N == 2560 && HC == 4, "gr_norm_comb shape"); + const uint32_t tk = TOPK; + d.p = P(lib_, "gr_norm_comb", {}, "gr_norm_comb"); + d.args = {mtl::Arg::B(R, r_off), mtl::Arg::B(ye_), mtl::Arg::B(inj, inj_off), mtl::Arg::B(wb, w_off), mtl::Arg::B(xn_), + mtl::Arg::V(a), mtl::Arg::B(plan_), mtl::Arg::B(ysh_), mtl::Arg::B(sg_), mtl::Arg::V(tk)}; + } else { + d.p = P(lib_, "gr_norm", {}, "gr_norm"); + d.args = {mtl::Arg::B(R, r_off), mtl::Arg::B(bo), mtl::Arg::B(inj, inj_off), mtl::Arg::B(wb, w_off), mtl::Arg::B(xn_), mtl::Arg::V(a)}; + } + d.grid[0] = HC; + d.grid[1] = ntok; + d.tg[0] = 256; + push(ds, d); +} + // the hyper-connection read (gr.metal, 3 dispatches): the previous sublayer's write folded in when `apply` void Window::gr_read(std::vector& ds, const std::string& pre, bool apply, const Act& out, bool final_, int ntok) { @@ -614,14 +1156,7 @@ void Window::gr_read(std::vector& ds, const std::string& pre, boo const gguf::Tensor& tup = T(pre + "_up.weight"); const gguf::Tensor& tin = final_ ? tdn : T(pre + "_inject.weight"); GrArgs a{N, HC, LR, (uint32_t) ntok, apply ? 1u : 0u, eps_}; - mtl::Dispatch d1; - d1.p = P(lib_, "gr_norm", {}, "gr_norm"); - d1.args = {mtl::Arg::B(R_), mtl::Arg::B(bo_), mtl::Arg::B(inj_[cur_inj_]), mtl::Arg::B(W_, tn.offset), mtl::Arg::B(xn_), - mtl::Arg::V(a)}; - d1.grid[0] = HC; - d1.grid[1] = ntok; - d1.tg[0] = 256; - push(ds, d1); + gr_norm(ds, R_, bo_, inj_[cur_inj_], W_, tn.offset, apply, ntok); // each matrix: its load-time Q8_0 copy (q8off_), else the file's own tensor (Unsloth's are Q8_0 already, their // inject rows F32) auto src = [&](const gguf::Tensor& t) -> std::pair { @@ -633,17 +1168,25 @@ void Window::gr_read(std::vector& ds, const std::string& pre, boo const auto [bi, oi] = src(tin); const auto [bu, ou] = src(tup); const bool q8 = is_q8(tdn) && is_q8(tup); + const int lt = q8 && hc_q4_ && !final_ && q8type_.count(tdn.name) ? q8type_.at(tdn.name) : (int) ref::Q8_0; + const bool q4 = lt == ref::Q4_0, q5 = lt == ref::Q5_0; const bool injf = q8 && !final_ && !q8off_.count(tin.name) && tin.type == ref::F32; mtl::Dispatch d2; - const char* dn = injf ? "gr_down_q8f_8_8" : q8 ? "gr_down_q8_8_8" : "gr_down_8_8"; - d2.p = P(lib_, dn, mtl::Constants().i(1, ntok), std::string(dn) + "#" + std::to_string(ntok)); + // M5: KS = 4 (4 simdgroups a threadgroup, bitwise the old 8_4), else KS = 8. Past 4 tokens two rows a simdgroup + // (_r2), which reads each xn value once for both + static_assert(LR % 2 == 0 && HC % 2 == 0, "gr_down _r2 rows"); + const bool m5_gr = m5_layout() && !q4 && !q5, r2 = !q4 && !q5 && ntok > 4; + std::string dn = q4 ? "gr_down_q4_8_8" : q5 ? "gr_down_q5_8_8" + : std::string(injf ? "gr_down_q8f_" : q8 ? "gr_down_q8_" : "gr_down_") + (m5_gr ? "4_4" : "8_8"); + if (r2) dn += "_r2"; + d2.p = P(lib_, dn, mtl::Constants().i(1, ntok), dn + "#" + std::to_string(ntok)); d2.args = {mtl::Arg::B(*bd, od), mtl::Arg::B(*bi, oi), mtl::Arg::B(xn_), mtl::Arg::B(lo_), mtl::Arg::B(inj_[1 - cur_inj_]), mtl::Arg::V(a)}; - d2.grid[0] = LR + HC; - d2.tg[0] = 256; + d2.grid[0] = r2 ? (LR + HC) / 2 : (LR + HC); + d2.tg[0] = m5_gr ? 128 : 256; push(ds, d2); mtl::Dispatch d3; - const char* up = q8 ? "gr_up_q8" : "gr_up"; + const char* up = q4 ? "gr_up_q4" : q5 ? "gr_up_q5" : q8 ? "gr_up_q8" : "gr_up"; d3.p = P(lib_, up, mtl::Constants().i(1, ntok), std::string(up) + "#" + std::to_string(ntok)); GrArgs au = a; au.n_tok = MAXT; // the plane stride of the prepared activation @@ -662,26 +1205,33 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 const int gi = l - l / 4; // a one-token window commits in the step; a longer one leaves the state and commit() folds the kept prefix // from the inputs kept per layer (qkvL_ ...) - const bool self_commit = ntok == 1; - const uint32_t cm = self_commit ? 1u : !partial_ ? 0u : win_head_ ? (0x100u | 1u) : 1u; + const bool self_commit = ntok == 1 || (!win_head_ && (partial_ || gdn_ring_ > 0)); + const uint32_t cm = self_commit ? 1u : (gdn_ring_ == 1 ? 2u : (gdn_ring_ == 2 ? 3u : (!partial_ ? 0u : (0x100u | 1u)))); const size_t o_qkv = (size_t) gi * MAXT * 10240 * 4, o_ab = (size_t) gi * MAXT * 48 * 4, o_z = (size_t) gi * MAXT * 6144 * 4; gemv(ds, T(b + "attn_qkv.weight"), XA_, qkvL_, 10240, ntok, o_qkv); nb_ = true; gemv(ds, T(b + "attn_gate.weight"), XA_, zL_, 6144, ntok, o_z); nb_ = true; - gemv(ds, T(b + "ssm_alpha.weight"), XA_, alL_, 48, ntok, o_ab); - nb_ = true; - gemv(ds, T(b + "ssm_beta.weight"), XA_, beL_, 48, ntok, o_ab); + // ssm_alpha + ssm_beta: one dispatch (gemv_cat), else two + if (!gemv_cat(ds, T(b + "ssm_alpha.weight"), T(b + "ssm_beta.weight"), XA_, alL_, beL_, 48, ntok, o_ab)) { + gemv(ds, T(b + "ssm_alpha.weight"), XA_, alL_, 48, ntok, o_ab); + nb_ = true; + gemv(ds, T(b + "ssm_beta.weight"), XA_, beL_, 48, ntok, o_ab); + } GdnArgs a{16, 48, 10240, cm, eps_, fuse_ ? (uint32_t) (MAXT * 6144) : 0u}; mtl::Dispatch d; d.p = P(lgdn_, "gdn_step", mtl::Constants().i(2, ntok), "gdn_step#" + std::to_string(ntok)); - const int hp = hp_; - d.args = {mtl::Arg::B(gst_[(size_t) gi]), mtl::Arg::B(ghist_[hp][(size_t) gi]), mtl::Arg::B(qkvL_, o_qkv), + const int hp = hp_, hp1 = gdn_ring_ > 0 ? ((hp + 1) % 3) : (1 - hp), hp2 = (hp + 2) % 3; + d.args = {mtl::Arg::B(gst_[sp_][(size_t) gi]), mtl::Arg::B(ghist_[hp][(size_t) gi]), mtl::Arg::B(qkvL_, o_qkv), mtl::Arg::B(alL_, o_ab), mtl::Arg::B(beL_, o_ab), mtl::Arg::B(zL_, o_z), mtl::Arg::B(W_, T(b + "ssm_conv1d.weight").offset), mtl::Arg::B(W_, T(b + "ssm_dt.bias").offset), mtl::Arg::B(W_, T(b + "ssm_a").offset), mtl::Arg::B(W_, T(b + "ssm_norm.weight").offset), mtl::Arg::B(fuse_ ? XY_.b : yg_), mtl::Arg::V(a), - mtl::Arg::B(ghist_[1 - hp][(size_t) gi])}; + mtl::Arg::B(ghist_[hp1][(size_t) gi])}; + if (cm == 2u || cm == 3u) { + d.args.push_back(mtl::Arg::B(gst_[1 - sp_][(size_t) gi])); + d.args.push_back(mtl::Arg::B(ghist_[hp2][(size_t) gi])); + } d.grid[0] = 48; d.tg[0] = 512; push(ds, d); @@ -690,13 +1240,16 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 gemv(ds, T(b + "ssm_out.weight"), XY_, bo_, N, ntok); } else { const int qi = l / 4; + const bool sel = !(fuse_ && pos0 + ntok <= 2051); gemv(ds, T(b + "attn_q.weight"), XA_, qfull_, 12288, ntok); nb_ = true; gemv(ds, T(b + "attn_k.weight"), XA_, kcur_, 512, ntok); nb_ = true; gemv(ds, T(b + "attn_v.weight"), XA_, vcur_, 512, ntok); - nb_ = true; - gemv(ds, T(b + "indexer.q_proj.weight"), XA_, iqr_, 512, ntok); + if (sel) { + nb_ = true; + gemv(ds, T(b + "indexer.q_proj.weight"), XA_, iqr_, 512, ntok); + } nb_ = true; gemv(ds, T(b + "indexer.k_proj.weight"), XA_, ikr_, 128, ntok); QsaArgs a{(uint32_t) pos0, 0, (uint32_t) max_cells_, (uint32_t) ntok, eps_}; @@ -711,7 +1264,7 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 mtl::Arg::V(a), mtl::Arg::B(mrope_)}; d.grid[0] = 33; d.grid[1] = ntok; - d.tg[0] = 256; + d.tg[0] = qsa_prep_tg(ntok); push(ds, d); if (rec_) rec_->push_back({(int) ds.size() - 1, 0}); mtl::Dispatch s; @@ -721,7 +1274,6 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 s.grid[0] = (uint32_t) (((pos0 + ntok) / 4 + 1 + 7) / 8); s.grid[1] = ntok; s.tg[0] = 256; - const bool sel = !(fuse_ && pos0 + ntok <= 2051); if (sel) { push(ds, s); if (rec_) rec_->push_back({(int) ds.size() - 1, 2}); @@ -739,12 +1291,12 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 at.p = P(lqsa_, "qsa_attn", {}, "qsa_attn"); at.args = {mtl::Arg::B(qcur_), mtl::Arg::B(kq_[(size_t) qi]), mtl::Arg::B(vq_[(size_t) qi]), mtl::Arg::B(ksc_[(size_t) qi]), mtl::Arg::B(vsc_[(size_t) qi]), mtl::Arg::B(qids_), mtl::Arg::B(part_), mtl::Arg::V(a)}; - at.grid[0] = 33; + at.grid[0] = (uint32_t) std::min(33, std::max(1, (pos0 + ntok + 63) / 64)); at.grid[1] = 2; at.grid[2] = ntok; at.tg[0] = 256; push(ds, at); - if (rec_) rec_->push_back({(int) ds.size() - 1, 0}); + if (rec_) rec_->push_back({(int) ds.size() - 1, 3}); mtl::Dispatch m; m.p = P(lqsa_, "qsa_merge", {}, "qsa_merge"); QsaArgs am = a; @@ -759,28 +1311,51 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 } // ---- the MoE gr_read(ds, b + "hc_ffn", true, XF_, false, ntok); - const gguf::Tensor& trt = T(b + "ffn_gate_inp.weight"); + const gguf::Tensor& trt = routT_.empty() ? T(b + "ffn_gate_inp.weight") : routT_[(size_t) l]; + const mtl::Buffer& rtw = routT_.empty() ? W_ : routW_; const bool rf = routef_ && (trt.type == ref::BF16 || trt.type == ref::F32); if (!rf) { - gemv(ds, trt, XF_, logits_, 512, ntok); + gemv(ds, trt, XF_, logits_, 512, ntok, 0, &rtw); nb_ = true; } if (stream_ && st_pred_ && l + 1 < 48) { // the next layer's routing guessed from this layer's FFN input + // STRATA_METAL_PREFETCH_DRAFT <= 0: the read-ahead reads only token 0's guess (stream.cpp ahead), so only it is + // computed, on the window's tile (bitwise the window's token 0). M5, Q8_0 2560 -> 512 + 512 at 4 tokens: -16% + // hot, -6% from DRAM (gemv_parity --pred; gemv_cat of the two at 1 token is slower than the pair) + static const int draft_k = [] { const char* e = std::getenv("STRATA_METAL_PREFETCH_DRAFT"); return e && *e ? std::atoi(e) : 1; }(); + const int pt = draft_k <= 0 ? 1 : ntok; nb_ = !rf; // beside the router GEMV; with router_route it is the first reader of XF_ here: a barrier - gemv(ds, T("blk." + std::to_string(l + 1) + ".ffn_gate_inp.weight"), XF_, predL_, 512, ntok, - (size_t) l * MAXT * 512 * 4); + if (predW_.ptr) + gemv(ds, predT_[(size_t) l + 1], XF_, predL_, 512, pt, (size_t) l * MAXT * 512 * 4, &predW_, ntok); + else + gemv(ds, T("blk." + std::to_string(l + 1) + ".ffn_gate_inp.weight"), XF_, predL_, 512, pt, + (size_t) l * MAXT * 512 * 4, nullptr, ntok); if (predL2_.ptr && l + 2 < 48) { nb_ = true; - gemv(ds, T("blk." + std::to_string(l + 2) + ".ffn_gate_inp.weight"), XF_, predL2_, 512, ntok, - (size_t) l * MAXT * 512 * 4); + if (predW_.ptr) + gemv(ds, predT_[(size_t) l + 2], XF_, predL2_, 512, pt, (size_t) l * MAXT * 512 * 4, &predW_, ntok); + else + gemv(ds, T("blk." + std::to_string(l + 2) + ".ffn_gate_inp.weight"), XF_, predL2_, 512, pt, + (size_t) l * MAXT * 512 * 4, nullptr, ntok); } nb_ = true; } - gemv(ds, T(b + "ffn_gate_shexp.weight"), XF_, gsh_, NFF, ntok); - nb_ = true; - gemv(ds, T(b + "ffn_up_shexp.weight"), XF_, ush_, NFF, ntok); - nb_ = true; - gemv(ds, T(b + "ffn_gate_inp_shexp.weight"), XF_, sg_, 1, ntok); + // the shared expert: its own list when streaming, placed between the handshake's signal and its wait (nothing in it + // reads what the host writes), else where it always was + std::vector shx; + std::vector& dsh = sigs_ ? shx : ds; + static const bool xsplit = !(std::getenv("STRATA_METAL_EXPERT_SPLIT") && std::atoi(std::getenv("STRATA_METAL_EXPERT_SPLIT")) == 0); + const bool esplit = sigs_ && xsplit; + auto shared_gu = [&] { + gemv(dsh, T(b + "ffn_gate_shexp.weight"), XF_, gsh_, NFF, ntok); + nb_ = true; + gemv(dsh, T(b + "ffn_up_shexp.weight"), XF_, ush_, NFF, ntok); + nb_ = true; + gemv(dsh, T(b + "ffn_gate_inp_shexp.weight"), XF_, sg_, 1, ntok); + }; + if (sigs_) nb_ = false; + if (!esplit) shared_gu(); + if (sigs_) nb_ = false; // moe_route below waits for the router GEMV (router_route sets its own) { MoeArgs a{}; const gguf::Tensor &tg_ = T(b + "ffn_gate_exps.weight"), &tu = T(b + "ffn_up_exps.weight"), @@ -801,34 +1376,63 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 a.n_ff = NFF; a.n_tok = ntok; a.topk = TOPK; - a.n_expert = NEXP; + a.n_expert = (uint32_t) n_exp_; + static const bool dmoce = !(std::getenv("STRATA_METAL_DRAFT_MOCE") && std::atoi(std::getenv("STRATA_METAL_DRAFT_MOCE")) == 0); + a.pad = stream_ ? (moce_pad_ | (dmoce && win_head_ && ntok > 1 && moce_pad_ ? 0x80000000u : 0u)) : 0u; const bool dq2 = td.type == ref::Q2_0; - const int gl = 32, gr = 1, gs = 4, ds_ = 4; - // the down tile: Q2_0 8,1; else 4,2 - but on the M4 (no M5 layout) Q5_1 4,1 and Q8_0 8,1 (4,2 spills there: - // 24-36 GB/s, mac-kern 513fc615) + int gl = 32, gr = 1, gs = 4, ds_ = 4; + // the down tile: Q2_0 8,2 (the M5 Air's KPROF sweep: 2.185 ms a token vs 2.431 at 8,1), 8,1 on the M4; else 4,2 - + // but on the M4 (no M5 layout) Q5_1 4,1 and Q8_0 8,1 (4,2 spills there: 24-36 GB/s, mac-kern 513fc615) const bool m4 = !m5_layout(); - const int dl = dq2 || (m4 && td.type == ref::Q8_0) ? 8 : 4, - dr = dq2 || (m4 && (td.type == ref::Q8_0 || td.type == ref::Q5_1)) ? 1 : 2; + if (!m4 && (tg_.type == ref::IQ2_XXS || tg_.type == ref::IQ3_XXS || + (ntok >= 3 && (tg_.type == ref::IQ2_S || tg_.type == ref::IQ2_XS || tg_.type == ref::IQ3_S)))) + gl = 16; + int dl = dq2 || (!m4 && td.type == ref::IQ4_NL) || (m4 && td.type == ref::Q8_0) ? 8 : 4, + dr = dq2 ? (m4 ? 1 : 2) : (m4 && (td.type == ref::Q8_0 || td.type == ref::Q5_1)) ? 1 : 2; + // for sweeping (KPROF): STRATA_METAL_MOE_GU_TILE / _DN_TILE = "LPR,NR,SG" of an instantiated moe_gu_ / moe_down_ + static const std::array gu_t = tile_env("STRATA_METAL_MOE_GU_TILE"), dn_t = tile_env("STRATA_METAL_MOE_DN_TILE"); + if (gu_t[0] > 0) { gl = gu_t[0]; gr = gu_t[1]; gs = gu_t[2]; } + if (dn_t[0] > 0) { dl = dn_t[0]; dr = dn_t[1]; ds_ = dn_t[2]; } a.gx_gu = NFF / (gs * (32 / gl) * gr); a.gx_dn = N / (ds_ * (32 / dl) * dr); mtl::Dispatch r; + const int mp_fc = std::getenv("STRATA_METAL_MOCE_PROTECT") ? std::max(0, std::min(10, std::atoi(std::getenv("STRATA_METAL_MOCE_PROTECT")))) : 3; + uint32_t mr1_bits = 0x3F800000u; // 1.0f + if (const char* mr1_env = std::getenv("STRATA_METAL_MOCE_BONUS_R1")) { + const float mr1_b = (float) std::atof(mr1_env); + if (mr1_b > 0.0f) { + const float mr1_s = std::exp(mr1_b); + std::memcpy(&mr1_bits, &mr1_s, 4); + } + } + mtl::Constants r_fc; + std::string r_fc_tag; + if (mp_fc != 3) { r_fc.i(24, mp_fc); r_fc_tag += "#mp" + std::to_string(mp_fc); } + if (mr1_bits != 0x3F800000u) { r_fc.u(25, mr1_bits); r_fc_tag += "#mr1_" + std::to_string(mr1_bits); } + // the resident experts first (moe.metal function constant 26): they run beside the host's turn + if (esplit) { r_fc.b(26, true); r_fc_tag += "#xs"; } if (rf) { // the router GEMV + moe_route (bitwise both), beside the shared expert's GEMVs - r.p = P(lib_, "router_route_" + std::to_string(trt.type), mtl::Constants().i(0, ntok), "router_route_" + std::to_string(trt.type) + "#" + std::to_string(ntok)); - GemvArgs ga{(uint32_t) N, 512, (uint32_t) ntok, (uint32_t) ref::row_bytes(trt.type, N), (uint32_t) N, 512}; - r.args = {mtl::Arg::B(W_, trt.offset), mtl::Arg::B(XF_.b, 0), mtl::Arg::B(XF_.b, XF_.plane_bytes() * 4), + mtl::Constants rfc2 = r_fc; + rfc2.i(0, ntok); + r.p = P(lib_, "router_route_" + std::to_string(trt.type), rfc2, "router_route_" + std::to_string(trt.type) + "#" + std::to_string(ntok) + r_fc_tag); + GemvArgs ga{(uint32_t) N, (uint32_t) n_exp_, (uint32_t) ntok, (uint32_t) ref::row_bytes(trt.type, N), (uint32_t) N, 512}; + r.args = {mtl::Arg::B(rtw, trt.offset), mtl::Arg::B(XF_.b, 0), mtl::Arg::B(XF_.b, XF_.plane_bytes() * 4), mtl::Arg::B(logits_), mtl::Arg::V(ga), mtl::Arg::B(ids_), mtl::Arg::B(wts_), mtl::Arg::B(plan_), - mtl::Arg::V(a), mtl::Arg::B(rcounter_)}; - r.grid[0] = 512 / 16; + mtl::Arg::V(a), mtl::Arg::B(rcounter_), + stream_ ? mtl::Arg::B(st_->smap(), (size_t) l * 512 * 4) : mtl::Arg::B(rcounter_)}; + r.grid[0] = (uint32_t) ((n_exp_ + 15) / 16); r.tg[0] = 512; - nb_ = true; + nb_ = !sigs_ || (stream_ && st_pred_ && l + 1 < 48); } else { - r.p = P(lib_, "moe_route", {}, "moe_route"); - r.args = {mtl::Arg::B(logits_), mtl::Arg::B(ids_), mtl::Arg::B(wts_), mtl::Arg::B(plan_), mtl::Arg::V(a)}; + r.p = P(lib_, "moe_route", r_fc, "moe_route" + r_fc_tag); + r.args = {mtl::Arg::B(logits_), mtl::Arg::B(ids_), mtl::Arg::B(wts_), mtl::Arg::B(plan_), mtl::Arg::V(a), + stream_ ? mtl::Arg::B(st_->smap(), (size_t) l * 512 * 4) : mtl::Arg::B(rcounter_)}; r.tg[0] = 512; } push(ds, r); - { + if (sigs_) sigs_->push_back((int) ds.size()); // the handshake starts: the routing is read + auto shared_prep = [&] { GlueArgs ga{NFF, MAXT, N, TOPK}; mtl::Dispatch d; d.p = P(lib_, "swiglu_prep", {}, "swiglu_prep"); @@ -836,53 +1440,67 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 d.grid[0] = NFF / 16; d.grid[1] = ntok; d.tg[0] = 16; - nb_ = !rf; // with router_route the shared GEMVs are not behind a barrier yet: this one waits for them - push(ds, d); - } + nb_ = esplit || (!sigs_ && !rf); + push(dsh, d); + }; + if (!esplit) shared_prep(); + if (sigs_ && !esplit) gemv(shx, T(b + "ffn_down_shexp.weight"), XH_, ysh_, N, ntok); // the shared expert's down too const std::string gun = "moe_gu_" + std::to_string(tg_.type) + "_" + std::to_string(gl) + "_" + std::to_string(gr) + "_" + std::to_string(gs); const std::string dnn = "moe_down_" + std::to_string(dq2 ? 142 : td.type) + "_" + std::to_string(dl) + "_" + std::to_string(dr) + "_" + std::to_string(ds_); - mtl::Dispatch g; const mtl::Buffer& EW = stream_ ? st_->chunk(0) : W_; - // function constant 23 (QUANT_CTAB): the i-quant grids from constant memory, no threadgroup copy (same values) - const bool ct = metal_ctab(); - mtl::Constants ac = stream_ ? mtl::Constants().b(20, true) : mtl::Constants(); - ac.b(23, ct); - const std::string ak = std::string(stream_ ? "#arena" : "") + (ct ? "#ct" : ""); - g.p = P(lib_, gun, ac, gun + ak); - g.args = {mtl::Arg::B(EW), mtl::Arg::B(XF_.b, XF_.plane_bytes() * xkind(tg_.type)), mtl::Arg::B(h_), mtl::Arg::B(plan_), - mtl::Arg::V(a), mtl::Arg::B(XF_.b, XF_.plane_bytes() * 4)}; - if (stream_) - for (int c = -1; c < 3; ++c) - g.args.push_back(c < 0 ? mtl::Arg::B(st_->smap(), (size_t) l * 512 * 4) : mtl::Arg::B(st_->chunk(c + 1))); - g.tg[0] = 32 * gs; - g.indirect = &plan_; - g.ind_off = P_IGU * 4; + // function constant 23 (QUANT_CTAB): the i-quant grids from constant memory, no threadgroup copy (same values); + // IQ3_S uses threadgroup tab on M5 (118.4 vs 105.8 GB/s: its 2 indices per loop diverge across lanes) + const bool ct = metal_ctab() && tg_.type != ref::IQ3_S; + // phase (function constant 27): 0 every group, 1 the resident ones, 2 the missing ones + auto experts = [&](std::vector& out, int phase) { + mtl::Constants ac = stream_ ? mtl::Constants().b(20, true) : mtl::Constants(); + ac.b(23, ct); + std::string ak = std::string(stream_ ? "#arena" : "") + (ct ? "#ct" : ""); + if (phase) { ac.i(27, phase); ak += "#ph" + std::to_string(phase); } + mtl::Dispatch g; + g.p = P(lib_, gun, ac, gun + ak); + g.args = {mtl::Arg::B(EW), mtl::Arg::B(XF_.b, XF_.plane_bytes() * xkind(tg_.type)), mtl::Arg::B(h_), mtl::Arg::B(plan_), + mtl::Arg::V(a), mtl::Arg::B(XF_.b, XF_.plane_bytes() * 4)}; + if (stream_) + for (int c = -1; c < 3; ++c) + g.args.push_back(c < 0 ? (phase == 1 ? mtl::Arg::B(plan_) : mtl::Arg::B(st_->smap(), (size_t) l * 512 * 4)) : mtl::Arg::B(st_->chunk(c + 1))); + g.tg[0] = 32 * gs; + g.indirect = &plan_; + g.ind_off = (phase == 1 ? P_IGU_R : phase == 2 ? P_IGU_M : P_IGU) * 4; + if (phase == 1) nb_ = false; + push(out, g); + if (phase == 1) { + nb_ = true; + shared_gu(); + } + if (!sigs_) { + nb_ = true; + gemv(out, T(b + "ffn_down_shexp.weight"), XH_, ysh_, N, ntok); + } + mtl::Dispatch dd; + dd.p = P(lib_, dnn, ac, dnn + ak); + dd.args = {mtl::Arg::B(EW), mtl::Arg::B(h_), mtl::Arg::B(ye_), mtl::Arg::B(plan_), mtl::Arg::V(a), mtl::Arg::B(h_)}; + if (stream_) + for (int c = -1; c < 3; ++c) + dd.args.push_back(c < 0 ? (phase == 1 ? mtl::Arg::B(plan_) : mtl::Arg::B(st_->smap(), (size_t) l * 512 * 4)) : mtl::Arg::B(st_->chunk(c + 1))); + dd.tg[0] = 32 * ds_; + dd.indirect = &plan_; + dd.ind_off = (phase == 1 ? P_IDN_R : phase == 2 ? P_IDN_M : P_IDN) * 4; + push(out, dd); + if (phase == 1) { + // zip independent shared work beside resident experts, then wait for the prepared activations + shared_prep(); + gemv(out, T(b + "ffn_down_shexp.weight"), XH_, ysh_, N, ntok); + } + }; + if (esplit) experts(shx, 1); // the experts the routing found in the arena: they cannot be evicted (pinned) + if (sigs_) ds.insert(ds.end(), shx.begin(), shx.end()); if (cuts_) cuts_->push_back((int) ds.size()); // the host's turn: this layer's experts put in place - push(ds, g); - nb_ = true; - gemv(ds, T(b + "ffn_down_shexp.weight"), XH_, ysh_, N, ntok); - mtl::Dispatch dd; - dd.p = P(lib_, dnn, ac, dnn + ak); - dd.args = {mtl::Arg::B(EW), mtl::Arg::B(h_), mtl::Arg::B(ye_), mtl::Arg::B(plan_), mtl::Arg::V(a), mtl::Arg::B(h_)}; - if (stream_) - for (int c = -1; c < 3; ++c) - dd.args.push_back(c < 0 ? mtl::Arg::B(st_->smap(), (size_t) l * 512 * 4) : mtl::Arg::B(st_->chunk(c + 1))); - dd.tg[0] = 32 * ds_; - dd.indirect = &plan_; - dd.ind_off = P_IDN * 4; - push(ds, dd); - } - { - GlueArgs a{N, (uint32_t) ntok, N, TOPK}; - mtl::Dispatch d; - d.p = P(lib_, "moe_combine_sh", {}, "moe_combine_sh"); - d.args = {mtl::Arg::B(ye_), mtl::Arg::B(plan_), mtl::Arg::B(ysh_), mtl::Arg::B(sg_), mtl::Arg::B(bo_), mtl::Arg::V(a)}; - d.grid[0] = (ntok * N + 255) / 256; - d.tg[0] = 256; - push(ds, d); + experts(ds, esplit ? 2 : 0); } + comb_defer(ds, bo_, ntok); // into the next layer's (or the PLE's, the final mixer's) gr_norm } namespace { @@ -952,6 +1570,12 @@ bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t if (Tn < 1 || Tn > MAXT || (head && Tn > max_t_)) { err = "metal window: window size out of range"; return false; } if (pos0 + Tn + 1 >= max_cells_) { err = "metal window: the window runs past the context"; return false; } const double t0 = now_ms(); + last_pos0_ = pos0; + if (Tn > 1 && Tn <= 4) { + tail_bak_.resize(tail_.size() * 4 * 128); + for (size_t j = 0; j < tail_.size(); ++j) + std::memcpy(tail_bak_.data() + j * 4 * 128, tail_[j].ptr, 4 * 4 * 128); + } // ---- the embeddings (CPU dequant of the token rows straight into the shared buffer), broadcast to the streams const gguf::Tensor& te = T("token_embd.weight"); const size_t rb = ref::row_bytes(te.type, N); @@ -981,8 +1605,8 @@ bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t for (int c = 0; c < 2051; ++c) qids_.as()[t * 2051 + c] = c; qdirty_ = false; } - const int key = (((((Tn * 2 + (head ? 1 : 0)) * 2 + hp_) * (MAXT + 1)) + pend_fold_) * 2 + (sel ? 1 : 0)) * 2 + - (mrope_on() ? 1 : 0); + const int key = (((((((Tn * 2 + (head ? 1 : 0)) * 6 + (sp_ * 3 + hp_)) * (MAXT + 1)) + pend_fold_) * 3 + fold_hp_) * 2 + + fold_off_) * 2 + (sel ? 1 : 0)) * 2 + (mrope_on() ? 1 : 0); mrope_sync(pos0, Tn); const int fold_n = pend_fold_; pend_fold_ = 0; @@ -1000,6 +1624,7 @@ bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t if (!g.bytes.empty()) qa = &g; std::memcpy(qa->bytes.data(), &a, sizeof a); if (kind == 2) d.grid[0] = (uint32_t) (((pos0 + Tn) / 4 + 1 + 7) / 8); + if (kind == 3) d.grid[0] = (uint32_t) std::min(33, std::max(1, (pos0 + Tn + 63) / 64)); } seg_T_ = Tn; const double r0 = now_ms(); @@ -1007,7 +1632,7 @@ bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t { static const int kt = std::getenv("STRATA_METAL_KPROF") ? std::atoi(std::getenv("STRATA_METAL_KPROF")) : 0; static int seen = 0; - if (kt == Tn && head && ++seen == 20 && !stream_) kprof(c.ds, Tn); + if (kt == Tn && head && ++seen == 20) kprof(c.ds, Tn); } tp_rg_ = now_ms() - r0; last_n_disp_ = (int) c.ds.size(); @@ -1017,53 +1642,24 @@ bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t Cached& cw = wcache_[key]; rec_ = &cw.qsa; cuts_ = stream_ ? &cw.cuts : nullptr; + // STRATA_METAL_SHEXP_OVERLAP=0: the shared expert after the handshake (as before) + static const bool overlap = !(std::getenv("STRATA_METAL_SHEXP_OVERLAP") && std::atoi(std::getenv("STRATA_METAL_SHEXP_OVERLAP")) == 0); + prewarm(); // window_init's already, unless the arena came late + cuts_ = stream_ ? &cw.cuts : nullptr; + sigs_ = stream_ && overlap ? &cw.sigs : nullptr; std::vector& ds = cw.ds; nb_ = false; - if (fold_n > 0) fold(ds, fold_n, 1 - hp_); // the last window's commit (hp_ already turned) + if (fold_n > 0) fold(ds, fold_n, fold_hp_); // the last window's commit (hp_ already turned) if (ple_on_) { xprep(ds, emb_, XE_, N); gemv(ds, T("blk.1.ple_key.weight"), XE_, pkey_, 10240, Tn); nb_ = true; gemv(ds, T("blk.1.ple_value.weight"), XE_, pval_, N, Tn); + ple_key_norm(ds, Tn); } cur_inj_ = 0; layer(ds, 0, Tn, pos0); - if (ple_on_) { - // layer 0's FFN write applied to R (gr_norm with apply writes R back; its xn is not used) - const gguf::Tensor& tn = T("blk.1.hc_attn_norm.weight"); - GrArgs a{N, HC, LR, (uint32_t) Tn, 1u, eps_}; - mtl::Dispatch d1; - d1.p = P(lib_, "gr_norm", {}, "gr_norm"); - d1.args = {mtl::Arg::B(R_), mtl::Arg::B(bo_), mtl::Arg::B(inj_[cur_inj_]), mtl::Arg::B(W_, tn.offset), mtl::Arg::B(xn_), - mtl::Arg::V(a)}; - d1.grid[0] = HC; - d1.grid[1] = Tn; - d1.tg[0] = 256; - push(ds, d1); - // the PLE block on the GPU (msl/ple.metal): norms, gate, conv, the residual update - struct PleArgs { uint32_t n_embd, hc; float eps; } pa{N, HC, strata::kernels::NG_RMS_EPS}; - auto off = [&](const char* n) { return T(n).offset; }; - mtl::Dispatch g1, g2, g3, g4; - g1.p = P(lple_, "ple_gnorm", {}, "ple_gnorm"); - g1.args = {mtl::Arg::B(pkey_), mtl::Arg::B(W_, off("blk.1.ple_norm_key.weight")), mtl::Arg::B(keyn_), mtl::Arg::V(pa)}; - g1.grid[0] = HC; g1.grid[1] = Tn; g1.tg[0] = 256; - push(ds, g1); - g2.p = g1.p; - g2.args = {mtl::Arg::B(R_), mtl::Arg::B(W_, off("blk.1.ple_norm_query.weight")), mtl::Arg::B(qn_), mtl::Arg::V(pa)}; - g2.grid[0] = HC; g2.grid[1] = Tn; g2.tg[0] = 256; - nb_ = true; // beside the key norm - push(ds, g2); - g3.p = P(lple_, "ple_gate_norm", {}, "ple_gate_norm"); - g3.args = {mtl::Arg::B(keyn_), mtl::Arg::B(qn_), mtl::Arg::B(pval_), mtl::Arg::B(W_, off("blk.1.ple_norm_conv.weight")), - mtl::Arg::B(gated_), mtl::Arg::B(normW_), mtl::Arg::V(pa)}; - g3.grid[0] = HC; g3.grid[1] = Tn; g3.tg[0] = 256; - push(ds, g3); - g4.p = P(lple_, "ple_conv_add", {}, "ple_conv_add"); - g4.args = {mtl::Arg::B(pleHist_), mtl::Arg::B(normW_), kw_arg(), - mtl::Arg::B(gated_), mtl::Arg::B(R_), mtl::Arg::V(pa)}; - g4.grid[0] = HC * N / 256; g4.grid[1] = Tn; g4.tg[0] = 256; - push(ds, g4); - } + if (ple_on_) ple_after_layer0(ds, Tn); for (int l = 1; l < 48; ++l) layer(ds, l, Tn, pos0); // ---- the final mixer and the head gr_read(ds, "output_hc", true, XA_, true, Tn); @@ -1085,6 +1681,7 @@ bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t } rec_ = nullptr; cuts_ = nullptr; + sigs_ = nullptr; seg_T_ = Tn; const double r0 = now_ms(); last_gpu_ms_ = 1000.0 * run_ds(cw); @@ -1095,7 +1692,9 @@ bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t double Window::run_ds(Cached& c) { if (!stream_) return cx_->run_graph(c.ds); - return cx_->run_graph_seg(c.ds, c.cuts, [](void* u, int s) { ((Window*) u)->at_cut(s); }, this); + const std::vector& sigs = c.sigs.size() == c.cuts.size() ? c.sigs : c.cuts; // SHEXP_OVERLAP=0: none between + return cx_->run_graph_seg(c.ds, c.cuts, sigs, [](void* u, int s) { ((Window*) u)->at_cut(s); }, + [](void* u, int s) { ((Window*) u)->after_cut(s); }, this); } namespace { @@ -1164,6 +1763,7 @@ bool Window::stream_arena(std::string& err) { if (const char* e = std::getenv("STRATA_METAL_KEEP_FREE_GB")) am.keep_gb = std::atof(e); if (const char* e = std::getenv("STRATA_METAL_ARENA_GB")) am.arena_gb = std::atof(e); if (const char* e = std::getenv("STRATA_METAL_LEAVE_SHARE")) am.leave_share = std::atof(e); // for measuring; the default is kLeaveShare + am.dense_file = (double) st_->dense_unwired_bytes(); const double keep = arena_keep(am); const double avail = am.available; double budget = arena_budget(am); @@ -1179,25 +1779,32 @@ bool Window::stream_arena(std::string& err) { std::fprintf(stderr, "strata metal stream: %.2f GB allocated before the arena (the process holds %.2f GB, %.2f GB resident); " "RAM %.1f GB (%.1f GB available, %.1f GB kept for macOS), GPU working set %.1f GB\n", used / gb, pm0.footprint / gb, pm0.rss / gb, phys / gb, avail / gb, keep / gb, cx_->recommended_working_set() / gb); + const char* epg = std::getenv("STRATA_METAL_PROMPT_GEMM"); + const bool pg_on = !(epg && std::string(epg) == "0"); + const char* esep = std::getenv("STRATA_METAL_SEPARATE_PROMPT_BUF"); + const char* ept = std::getenv("STRATA_METAL_PHASE_TRADE"); + const bool sep_pg = (ept && std::atoi(ept) == 0) || (esep && std::string(esep) == "1"); + const double pg_scratch_need = (pg_on && sep_pg) ? (double) prompt_gemm_need_bytes() : 0.0; + if (am.arena_gb < 0 && pg_scratch_need > 0.0) { + budget = std::max(1.0 * gb, budget - pg_scratch_need); + } { // the planned total: arena + the Metal buffers (or the process's own count when it is more) + the n-gram table's // reserve, against the cap (arena_budget.h: footprint_cap) const FootprintPlan fp = footprint_plan(am, budget); - std::fprintf(stderr, "strata metal stream: planned resident footprint %.1f GB = arena %.1f + %.1f held before it + %.1f " - "reserved for the n-gram table; cap %.1f GB = the lesser of the GPU working set %.1f GB and the RAM " - "%.1f GB less %.1f GB left to macOS and its file cache%s\n", - fp.total / gb, fp.arena / gb, fp.base / gb, fp.ple / gb, fp.cap / gb, cx_->recommended_working_set() / gb, + std::fprintf(stderr, "strata metal stream: planned resident footprint %.1f GB decode (%.1f GB virtual peak during prefill with %.2f GB transient prompt scratch) " + "= arena %.1f + %.1f held before it + %.1f reserved for the n-gram table; cap %.1f GB = the lesser of the GPU working set %.1f GB " + "and the RAM %.1f GB less %.1f GB left to macOS and its file cache%s\n", + fp.total / gb, (fp.total + pg_scratch_need) / gb, pg_scratch_need / gb, + fp.arena / gb, fp.base / gb, fp.ple / gb, fp.cap / gb, cx_->recommended_working_set() / gb, phys / gb, (am.phys - fp.cap) / gb, - am.arena_gb >= 0 && fp.over() ? " - STRATA_METAL_ARENA_GB is above the cap: kept as set (your choice), but " - "a Mac that runs short of memory can freeze (watchdog panics seen at 90+ GB " - "resident of 128)" : ""); + am.arena_gb >= 0 && (fp.total + pg_scratch_need > fp.cap) ? " - peak prefill above cap" : ""); } - const char* epg = std::getenv("STRATA_METAL_PROMPT_GEMM"); - const bool pg_on = !(epg && std::string(epg) == "0"); bool minimal = false; // the arena's least layout (the min slots) is not shrunk further: it goes ahead, filled in steps for (int tries = 0;; ++tries) { const double planned = budget; std::function recheck = [&]() -> uint64_t { return (uint64_t) arena_shrink(am, planned, avail_now()); }; - if (st_->make_arena(*cx_, (uint64_t) budget, prof, pg_on ? 576 : 64, err, minimal ? std::function() : recheck, + const int min_sl = std::getenv("STRATA_METAL_MIN_SLOTS") ? std::max(64, std::atoi(std::getenv("STRATA_METAL_MIN_SLOTS"))) : (pg_on ? (n_exp_ + 64) : 64); + if (st_->make_arena(*cx_, (uint64_t) budget, prof, min_sl, err, minimal ? std::function() : recheck, [&](uint64_t step) { return arena_grow(am, avail_now(), (double) step); })) break; if (st_->shrink_to() == 0 || tries == 4) return false; @@ -1206,7 +1813,7 @@ bool Window::stream_arena(std::string& err) { budget / gb, st_->shrink_to() / gb); budget = (double) st_->shrink_to(); } - if (pg_on && !prompt_gemm_init(err)) return false; + if (pg_on && !sep_pg && !prompt_gemm_init(err)) return false; { // what the process really holds now, against the plan and the cap (a number that disagrees is a bug in the count) const ProcMem pm1 = proc_mem(); const double cap = footprint_cap(am); @@ -1300,6 +1907,8 @@ void Window::kprof(const std::vector& ds, int Tn) { } Window::~Window() { + free_prompt_gemm(); + if (st_) st_->release_region(); if (logpos_) std::fclose(logpos_); if (std::getenv("STRATA_METAL_PROFILE") == nullptr) return; auto line = [](const char* what, const Prof& x) { @@ -1327,8 +1936,56 @@ bool Window::finish(int Tn, const int32_t* tokens, int64_t pos0, bool head, int3 ++x.n; x.pre += tp_pre_ - t0; x.rg += tp_rg_; x.gpu += last_gpu_ms_; x.fin += t1 - tp_pre_ - tp_rg_; x.all += t1 - t0; } + { // anti-slop v4 (slop4 above): every row's state is stored at its position; the head rows get the penalties + static const slop4::Knobs kn = slop4::knobs(); + const bool on = head && (kn.lex > 0.0f || kn.hon > 0.0f || kn.md > 0.0f || kn.colon > 0.0f); + const int per = (n_vocab_ + 119) / 120; + slop4::State st = slop_at(pos0 - 1); + for (int t = 0; t < Tn; ++t) { + st = slop4::step(st, tokens[t]); + slop_put(pos0 + t) = st; + if (!on) continue; + slop4::Pen pen[80]; + int np = slop4::penalties(st, kn, pen); + float* lg = lm_.as() + (size_t) t * n_vocab_; + float* pv = apv_.as() + (size_t) t * 120; + int32_t* pi = api_.as() + (size_t) t * 120; + if (np > 0) { + float mx = -INFINITY; + for (int k = 0; k < 120; ++k) mx = std::max(mx, pv[k]); + if (slop4::needs_hold(pen, np, lg, mx)) { + double s = 0.0; + for (int64_t i = 0; i < n_vocab_; ++i) s += std::exp((double) (lg[i] - mx)); + np += slop4::hold(pen, np, lg, (double) mx + std::log(s), pen + np); + } + } + // no empty turn: <|im_end|> or right after \n\n + if (kn.lex > 0.0f && st.t2 == slop4::ETHINK && (st.t1 == slop4::NL2 || st.t1 == slop4::NL)) { + pen[np++] = {slop4::IM_END, kn.lex}; + pen[np++] = {slop4::THINK, kn.lex}; + } + if (np == 0) continue; + for (int i = 0; i < np; ++i) lg[pen[i].id] -= pen[i].v; + // the argmax (tokout_) from the GPU's 120 partition maxima: a partition whose max was pushed down is rescanned + bool moved = false; + for (int i = 0; i < np; ++i) { + const int part = pen[i].id / per; + if (pi[part] != pen[i].id) continue; + const int b0 = part * per, b1 = std::min(b0 + per, (int) n_vocab_); + float bv = -INFINITY; int bi = 0x7fffffff; + for (int j = b0; j < b1; ++j) if (!std::isnan(lg[j]) && (lg[j] > bv || (lg[j] == bv && j < bi))) { bv = lg[j]; bi = j; } + pv[part] = bv; pi[part] = bi; + moved = true; + } + if (!moved) continue; + float bv = -INFINITY; int bi = 0x7fffffff; + for (int k = 0; k < 120; ++k) if (pv[k] > bv || (pv[k] == bv && pi[k] < bi)) { bv = pv[k]; bi = pi[k]; } + tokout_.as()[t] = bi; + } + } logpos(Tn, tokens, pos0, head); - if (Tn == 1 || partial_) hp_ = 1 - hp_; // the window committed its first row (or all of them) in the step + if (Tn == 1 || (!head && (partial_ || gdn_ring_ > 0))) hp_ = gdn_ring_ > 0 ? ((hp_ + 1) % 3) : (1 - hp_); + else if (partial_ && gdn_ring_ == 0) hp_ = 1 - hp_; // the window committed its first row (or all of them) in the step last_head_ = head; for (int t = 0; t < Tn; ++t) out[t] = head ? tokout_.as()[t] : -1; last_T_ = Tn; @@ -1349,7 +2006,8 @@ bool Window::finish(int Tn, const int32_t* tokens, int64_t pos0, bool head, int3 } void Window::fold(std::vector& ds, int n_keep, int hp) { - const size_t off = partial_ ? 1 : 0; // rows 1.. (row 0 was committed in the window's step) + const size_t off = (size_t) fold_off_; // rows off.. (row 0 was committed in the window's step when fold_off_ == 1) + const int hp_out = gdn_ring_ > 0 ? ((hp + 1) % 3) : (1 - hp); for (int l = 0; l < 48; ++l) { if (l % 4 == 3) continue; const int gi = l - l / 4; @@ -1359,11 +2017,11 @@ void Window::fold(std::vector& ds, int n_keep, int hp) { GdnArgs a{16, 48, 10240, 1u, eps_, 0}; mtl::Dispatch d; d.p = P(lgdn_, "gdn_step", mtl::Constants().i(2, n_keep), "gdn_step#" + std::to_string(n_keep)); - d.args = {mtl::Arg::B(gst_[(size_t) gi]), mtl::Arg::B(ghist_[hp][(size_t) gi]), mtl::Arg::B(qkvL_, o_qkv), + d.args = {mtl::Arg::B(gst_[sp_][(size_t) gi]), mtl::Arg::B(ghist_[hp][(size_t) gi]), mtl::Arg::B(qkvL_, o_qkv), mtl::Arg::B(alL_, o_ab), mtl::Arg::B(beL_, o_ab), mtl::Arg::B(zL_, o_z), mtl::Arg::B(W_, T(b + "ssm_conv1d.weight").offset), mtl::Arg::B(W_, T(b + "ssm_dt.bias").offset), mtl::Arg::B(W_, T(b + "ssm_a").offset), mtl::Arg::B(W_, T(b + "ssm_norm.weight").offset), - mtl::Arg::B(ydump_), mtl::Arg::V(a), mtl::Arg::B(ghist_[1 - hp][(size_t) gi])}; + mtl::Arg::B(ydump_), mtl::Arg::V(a), mtl::Arg::B(ghist_[hp_out][(size_t) gi])}; d.grid[0] = 48; d.tg[0] = 512; nb_ = true; // the layers' folds are independent of each other @@ -1377,41 +2035,50 @@ bool Window::checkpoint_save(strata::core::ConversationCheckpoint& c, std::strin if (pend_fold_ > 0) { // the last commit's fold first: the checkpoint is the state after it std::vector ds; nb_ = false; - fold(ds, pend_fold_, 1 - hp_); + fold(ds, pend_fold_, fold_hp_); cx_->run_graph(ds); pend_fold_ = 0; } const size_t gs = 4 * 128 * 48 * 128, hs = 4 * 10240 * 3, ts = 4 * 4 * 128, dsz = 4 * 128; - c.gdn.resize(gst_.size() * (gs + hs)); - for (size_t i = 0; i < gst_.size(); ++i) { - std::memcpy(c.gdn.data() + i * (gs + hs), gst_[i].ptr, gs); + c.gdn.resize(gst_[sp_].size() * (gs + hs)); + for (size_t i = 0; i < gst_[sp_].size(); ++i) { + std::memcpy(c.gdn.data() + i * (gs + hs), gst_[sp_][i].ptr, gs); std::memcpy(c.gdn.data() + i * (gs + hs) + gs, ghist_[hp_][i].ptr, hs); } c.ple.assign((const uint8_t*) pleHist_.ptr, (const uint8_t*) pleHist_.ptr + 4 * 9 * 10240); - c.tails.resize(tail_.size() * ts); - c.dead.resize(dead_.size() * dsz); + const size_t nq = tail_.size() + (mtp_on_ ? mtail_.size() : 0); + c.tails.resize(nq * ts); + c.dead.resize(nq * dsz); for (size_t j = 0; j < tail_.size(); ++j) { std::memcpy(c.tails.data() + j * ts, tail_[j].ptr, ts); std::memcpy(c.dead.data() + j * dsz, dead_[j].ptr, dsz); } + if (mtp_on_ && !mtail_.empty()) { + std::memcpy(c.tails.data() + tail_.size() * ts, mtail_[0].ptr, ts); + std::memcpy(c.dead.data() + dead_.size() * dsz, mdead_[0].ptr, dsz); + } c.block_pos.clear(); return true; } void Window::reset() { pend_fold_ = 0; + sp_ = hp_ = fold_hp_ = fold_off_ = 0; last_T_ = 0; auto zero = [](const mtl::Buffer& b, size_t n) { std::memset(b.ptr, 0, n); }; - for (size_t i = 0; i < gst_.size(); ++i) { - zero(gst_[i], 4 * 128 * 48 * 128); + for (size_t i = 0; i < gst_[0].size(); ++i) { + zero(gst_[0][i], 4 * 128 * 48 * 128); + if (gdn_ring_ > 0) zero(gst_[1][i], 4 * 128 * 48 * 128); zero(ghist_[0][i], 4 * 10240 * 3); zero(ghist_[1][i], 4 * 10240 * 3); + if (gdn_ring_ > 0) zero(ghist_[2][i], 4 * 10240 * 3); } zero(pleHist_, 4 * 9 * 10240); for (size_t j = 0; j < tail_.size(); ++j) { zero(tail_[j], 4 * 4 * 128); zero(dead_[j], 4 * 128); } + mtp_cell0_ = mtp_cell1_ = mtp_sink_ = 0; for (size_t j = 0; j < mtail_.size(); ++j) { zero(mtail_[j], 4 * 4 * 128); zero(mdead_[j], 4 * 128); @@ -1420,15 +2087,16 @@ void Window::reset() { bool Window::checkpoint_restore(const strata::core::ConversationCheckpoint& c, std::string& err) { const size_t gs = 4 * 128 * 48 * 128, hs = 4 * 10240 * 3, ts = 4 * 4 * 128, dsz = 4 * 128; - if (c.gdn.size() != gst_.size() * (gs + hs) || c.ple.size() != (size_t) 4 * 9 * 10240 || - c.tails.size() != tail_.size() * ts || c.dead.size() != dead_.size() * dsz) { + const size_t nq = tail_.size() + (mtp_on_ ? mtail_.size() : 0); + if (c.gdn.size() != gst_[sp_].size() * (gs + hs) || c.ple.size() != (size_t) 4 * 9 * 10240 || + c.tails.size() != nq * ts || c.dead.size() != nq * dsz) { err = "metal window: the checkpoint is not this engine's (sizes differ)"; return false; } pend_fold_ = 0; // the restored state replaces whatever was pending last_T_ = 0; - for (size_t i = 0; i < gst_.size(); ++i) { - std::memcpy(gst_[i].ptr, c.gdn.data() + i * (gs + hs), gs); + for (size_t i = 0; i < gst_[sp_].size(); ++i) { + std::memcpy(gst_[sp_][i].ptr, c.gdn.data() + i * (gs + hs), gs); std::memcpy(ghist_[hp_][i].ptr, c.gdn.data() + i * (gs + hs) + gs, hs); } std::memcpy(pleHist_.ptr, c.ple.data(), c.ple.size()); @@ -1436,36 +2104,321 @@ bool Window::checkpoint_restore(const strata::core::ConversationCheckpoint& c, s std::memcpy(tail_[j].ptr, c.tails.data() + j * ts, ts); std::memcpy(dead_[j].ptr, c.dead.data() + j * dsz, dsz); } + if (mtp_on_ && !mtail_.empty()) { + std::memcpy(mtail_[0].ptr, c.tails.data() + tail_.size() * ts, ts); + std::memcpy(mdead_[0].ptr, c.dead.data() + dead_.size() * dsz, dsz); + } + return true; +} + +namespace { +// one contiguous range of a buffer in the session file +struct SessionSec { + void* p; + size_t bytes; + size_t off = 0; +}; +// the sections in file order from `first` on, each on a 16 KiB boundary; returns the end +size_t session_place(std::vector& secs, size_t first) { + size_t o = scr::align_up(first); + for (auto& s : secs) { s.off = o; o = scr::align_up(o + s.bytes); } + return o; +} +// pread / pwrite of every section, 8 MiB pieces over 4 threads (one stream does not fill the SSD's queue) +bool session_io(int fd, const std::vector& secs, bool write, std::string& err) { + struct Piece { uint8_t* p; size_t n; off_t off; }; + std::vector ps; + constexpr size_t PIECE = (size_t) 8 << 20; + for (const auto& s : secs) + for (size_t o = 0; o < s.bytes; o += PIECE) + ps.push_back({(uint8_t*) s.p + o, std::min(PIECE, s.bytes - o), (off_t) (s.off + o)}); + std::atomic next{0}; + std::atomic ok{true}; + auto work = [&] { + for (size_t i; ok && (i = next++) < ps.size();) { + const Piece& q = ps[i]; + for (size_t d = 0; d < q.n;) { + const ssize_t r = write ? ::pwrite(fd, q.p + d, q.n - d, q.off + (off_t) d) : ::pread(fd, q.p + d, q.n - d, q.off + (off_t) d); + if (r <= 0) { if (r < 0 && errno == EINTR) continue; ok = false; break; } + d += (size_t) r; + } + } + }; + std::vector th; + for (int t = 0; t < 4; ++t) th.emplace_back(work); + for (auto& t : th) t.join(); + if (!ok) err = std::string("metal window: session ") + (write ? "write" : "read") + " failed: " + std::strerror(errno); + return ok; +} +} // namespace + +// the session's buffers in file order; the cells [c0, n) of the positional caches (or [mc0, n) for MTP) +static void session_sections(std::vector& secs, const std::vector& gst, const std::vector& ghist, + const mtl::Buffer& ple, const std::vector* const* q, int64_t n, int64_t c0 = 0, int64_t mc0 = 0) { + for (size_t i = 0; i < gst.size(); ++i) { + secs.push_back({gst[i].ptr, scr::GST_BYTES}); + secs.push_back({ghist[i].ptr, scr::GHIST_BYTES}); + } + secs.push_back({ple.ptr, scr::PLE_BYTES}); + const size_t un = (size_t) n; + for (int set = 0; set < 2; ++set) { // the 12 QSA layers, then the MTP layer's (when loaded) + const std::vector* const* b = q + (set ? 7 : 0); + if (set && b[0] == nullptr) break; + const size_t u0 = (size_t) (set ? mc0 : c0); + const size_t dn = un - u0, db = un / 4 - u0 / 4; + for (size_t j = 0; j < b[0]->size(); ++j) { + secs.push_back({(uint8_t*) (*b[0])[j].ptr + u0 * scr::KQ_CELL, dn * scr::KQ_CELL}); + secs.push_back({(uint8_t*) (*b[1])[j].ptr + u0 * scr::KQ_CELL, dn * scr::KQ_CELL}); + secs.push_back({(uint8_t*) (*b[2])[j].ptr + u0 * scr::SC_CELL, dn * scr::SC_CELL}); + secs.push_back({(uint8_t*) (*b[3])[j].ptr + u0 * scr::SC_CELL, dn * scr::SC_CELL}); + secs.push_back({(*b[4])[j].ptr, scr::TAIL_BYTES}); + secs.push_back({(*b[5])[j].ptr, scr::DEAD_BYTES}); + secs.push_back({(uint8_t*) (*b[6])[j].ptr + (u0 / 4) * scr::POOL_BLOCK, db * scr::POOL_BLOCK}); + } + } +} + +static int64_t session_mtp_c0(uint32_t flags, int64_t n_cells, int64_t c0) { + if (c0 <= 0 || (flags & 6u) != 6u) return 0; + const int64_t wlen = (int64_t) ((flags >> 11) & 0xFFFu); + const int64_t mcell1 = (flags & 8u) ? n_cells : 0; + const int64_t mcell0 = std::max(0, mcell1 - wlen); + return (std::min(c0, mcell0) / 4) * 4; +} + +bool Window::session_save(const std::string& path, const int32_t* ids, int64_t n_cells, std::string& err, int64_t base_cells) { + if (n_cells < 0 || n_cells >= max_cells_) { err = "metal window: session cell count out of range"; return false; } + if (mrope_on()) { err = "metal window: sessions of M-RoPE (vision) prompts are not supported"; return false; } + if (pend_fold_ > 0) { // as checkpoint_save: the state after the last commit's fold + std::vector ds; + nb_ = false; + fold(ds, pend_fold_, fold_hp_); + cx_->run_graph(ds); + pend_fold_ = 0; + } + const std::vector* q[14] = {&kq_, &vq_, &ksc_, &vsc_, &tail_, &dead_, &pooled_}; + if (mtp_on_) { + const std::vector* m[7] = {&mkq_, &mvq_, &mksc_, &mvsc_, &mtail_, &mdead_, &mpooled_}; + std::copy(m, m + 7, q + 7); + } + scr::SessionHeader h; + h.flags = (ple_on_ ? 1u : 0u) | (mtp_on_ ? 2u : 0u); + if (mtp_on_ && mtp_cell1_ > 0) { + const uint32_t sink = (uint32_t) std::clamp(mtp_sink_, 0, 127); + const uint32_t wlen = (uint32_t) std::clamp(mtp_cell1_ - mtp_cell0_, 0, 4095); + h.flags |= 4u | ((mtp_cell1_ == n_cells) ? 8u : 0u) | (sink << 4) | (wlen << 11); + } + const int64_t c0 = (base_cells >= 32 && base_cells < n_cells) ? (base_cells / 4) * 4 : 0; + const int64_t mc0 = session_mtp_c0(h.flags, n_cells, c0); + if (c0 > 0) h.flags |= (1u << 23); + std::vector secs; + session_sections(secs, gst_[sp_], ghist_[hp_], pleHist_, q, n_cells, c0, mc0); + const int64_t n_ids = ids ? n_cells : 0; + h.n_cells = n_cells; + h.n_ids = (int64_t) ((uint64_t) (uint32_t) n_ids | ((uint64_t) (uint32_t) c0 << 32)); + h.model_bytes = f_.file_bytes; + h.model_tensors = f_.tensors.size(); + h.rope_base = rope_base_; + h.ple_prev[0] = ss_->ple_prev[0]; + h.ple_prev[1] = ss_->ple_prev[1]; + const size_t head = sizeof h + (size_t) n_ids * 4, end = session_place(secs, head); + h.payload_bytes = end - head; + const std::string tmp = path + ".partial"; + const int fd = ::open(tmp.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0644); + if (fd < 0) { err = "metal window: cannot create " + tmp + ": " + std::strerror(errno); return false; } + ::fcntl(fd, F_NOCACHE, 1); + bool ok = ::ftruncate(fd, (off_t) end) == 0 && ::pwrite(fd, &h, sizeof h, 0) == (ssize_t) sizeof h && + (n_ids == 0 || ::pwrite(fd, ids, (size_t) n_ids * 4, sizeof h) == (ssize_t) (n_ids * 4)); + if (!ok) err = "metal window: session header write failed: " + std::string(std::strerror(errno)); + ok = ok && session_io(fd, secs, true, err); + if (ok) ::fsync(fd); + ok = (::close(fd) == 0) && ok; + if (ok && std::rename(tmp.c_str(), path.c_str()) != 0) { err = "metal window: rename to " + path + " failed"; ok = false; } + if (!ok) ::unlink(tmp.c_str()); + return ok; +} + +bool Window::session_load(const std::string& path, std::vector& ids, int64_t& n_cells, std::string& err, const char* parent_path) { + const int fd = ::open(path.c_str(), O_RDONLY); + if (fd < 0) { err = "metal window: cannot open " + path + ": " + std::strerror(errno); return false; } + ::fcntl(fd, F_NOCACHE, 1); + scr::SessionHeader h, want; + struct stat st{}; + bool ok = ::fstat(fd, &st) == 0 && ::pread(fd, &h, sizeof h, 0) == (ssize_t) sizeof h; + const int64_t n_ids = (int64_t) (uint32_t) h.n_ids; + const int64_t c0 = (h.flags & (1u << 23)) ? (int64_t) ((uint64_t) h.n_ids >> 32) : 0; + if (!ok || std::memcmp(h.magic, want.magic, 8) != 0 || h.version != want.version) err = "metal window: " + path + " is not a session file"; + else if (h.model_bytes != f_.file_bytes || h.model_tensors != f_.tensors.size() || h.rope_base != rope_base_) + err = "metal window: the session is another model's"; + else if (h.n_cells < 0 || h.n_cells >= max_cells_ || c0 < 0 || c0 >= h.n_cells || (c0 & 3) != 0) + err = "metal window: the session is longer than this context (--ctx)"; + else if ((h.flags & 1u) != (ple_on_ ? 1u : 0u) || ((h.flags & 2u) && !mtp_on_)) err = "metal window: the session's PLE / MTP setup differs"; + else if (c0 > 0 && (parent_path == nullptr || parent_path[0] == '\0')) err = "metal window: shared session requires its learned prefix"; + if (!err.empty()) { ::close(fd); return false; } + const std::vector* q[14] = {&kq_, &vq_, &ksc_, &vsc_, &tail_, &dead_, &pooled_}; + if (h.flags & 2u) { + const std::vector* m[7] = {&mkq_, &mvq_, &mksc_, &mvsc_, &mtail_, &mdead_, &mpooled_}; + std::copy(m, m + 7, q + 7); + } + const int64_t mc0 = session_mtp_c0(h.flags, h.n_cells, c0); + std::vector secs; + session_sections(secs, gst_[sp_], ghist_[hp_], pleHist_, q, h.n_cells, c0, mc0); + const size_t head = sizeof h + (size_t) n_ids * 4, end = session_place(secs, head); + if (end - head != h.payload_bytes || (size_t) st.st_size != end) { ::close(fd); err = "metal window: the session file is truncated or of another layout"; return false; } + ids.resize((size_t) n_ids); + ok = n_ids == 0 || ::pread(fd, ids.data(), (size_t) n_ids * 4, sizeof h) == (ssize_t) (n_ids * 4); + if (ok && c0 > 0) { + const int pfd = ::open(parent_path, O_RDONLY); + if (pfd < 0) { + err = "metal window: cannot open learned prefix " + std::string(parent_path); + ok = false; + } else { + ::fcntl(pfd, F_NOCACHE, 1); + scr::SessionHeader ph; + struct stat pst{}; + bool pok = ::fstat(pfd, &pst) == 0 && ::pread(pfd, &ph, sizeof ph, 0) == (ssize_t) sizeof ph; + const int64_t pn_ids = (int64_t) (uint32_t) ph.n_ids; + const int64_t pc0 = (ph.flags & (1u << 23)) ? (int64_t) ((uint64_t) ph.n_ids >> 32) : 0; + if (!pok || std::memcmp(ph.magic, want.magic, 8) != 0 || ph.version != want.version || pc0 != 0 || + ph.model_bytes != f_.file_bytes || ph.model_tensors != f_.tensors.size() || ph.rope_base != rope_base_ || + ph.n_cells < c0 || (ph.flags & 3u) != (h.flags & 3u)) { + err = "metal window: learned prefix header mismatch"; + ok = false; + } else { + std::vector p_all, p_secs; + session_sections(p_all, gst_[sp_], ghist_[hp_], pleHist_, q, ph.n_cells, 0, 0); + const size_t phead = sizeof ph + (size_t) pn_ids * 4, pend = session_place(p_all, phead); + if (pend - phead != ph.payload_bytes || (size_t) pst.st_size != pend) { + err = "metal window: learned prefix file is truncated"; + ok = false; + } else { + size_t idx = 2 * gst_[sp_].size() + 1; + for (int set = 0; set < 2; ++set) { + const std::vector* const* b = q + (set ? 7 : 0); + if (set && b[0] == nullptr) break; + const size_t u0 = (size_t) (set ? mc0 : c0); + for (size_t j = 0; j < b[0]->size(); ++j, idx += 7) { + if (u0 == 0) continue; + p_secs.push_back({p_all[idx + 0].p, u0 * scr::KQ_CELL, p_all[idx + 0].off}); + p_secs.push_back({p_all[idx + 1].p, u0 * scr::KQ_CELL, p_all[idx + 1].off}); + p_secs.push_back({p_all[idx + 2].p, u0 * scr::SC_CELL, p_all[idx + 2].off}); + p_secs.push_back({p_all[idx + 3].p, u0 * scr::SC_CELL, p_all[idx + 3].off}); + if (u0 / 4 > 0) p_secs.push_back({p_all[idx + 6].p, (u0 / 4) * scr::POOL_BLOCK, p_all[idx + 6].off}); + } + } + ok = session_io(pfd, p_secs, false, err); + } + } + ::close(pfd); + } + } + ok = ok && session_io(fd, secs, false, err); + ::close(fd); + if (!ok) { if (err.empty()) err = "metal window: session read failed"; reset(); return false; } + pend_fold_ = 0; + if ((h.flags & 6u) == 6u) { + mtp_sink_ = (int64_t) ((h.flags >> 4) & 0x7Fu); + const int64_t wlen = (int64_t) ((h.flags >> 11) & 0xFFFu); + mtp_cell1_ = (h.flags & 8u) ? h.n_cells : 0; + mtp_cell0_ = std::max(0, mtp_cell1_ - wlen); + } else { + mtp_cell0_ = mtp_cell1_ = mtp_sink_ = 0; + } + last_T_ = ids.empty() ? 0 : 1; + if (!ids.empty()) last_tokens_[0] = ids.back(); + slop_scan(ids.data(), (int64_t) ids.size(), 0); + ss_->ple_prev[0] = h.ple_prev[0]; + ss_->ple_prev[1] = h.ple_prev[1]; + n_cells = h.n_cells; return true; } +bool Window::scr_relocate(int64_t c_src, int64_t c_dst, int64_t n, std::string& err) { + if (mrope_on()) { err = "metal window: SCR of M-RoPE (vision) cells is not supported"; return false; } + if (n < 0 || c_src < 0 || c_dst < 0 || c_src + n >= max_cells_ || c_dst + n >= max_cells_) { err = "metal window: SCR span out of range"; return false; } + if (((c_dst - c_src) & 3) != 0) { err = "metal window: SCR needs a shift by a multiple of 4 cells (pad the new span)"; return false; } + const scr::Rot rot(rope_base_, c_dst - c_src); + std::vector L[5]; + auto add = [&](const std::vector& k, const std::vector& v, const std::vector& ks, + const std::vector& vs, const std::vector& p) { + for (size_t j = 0; j < k.size(); ++j) { L[0].push_back(&k[j]); L[1].push_back(&v[j]); L[2].push_back(&ks[j]); L[3].push_back(&vs[j]); L[4].push_back(&p[j]); } + }; + add(kq_, vq_, ksc_, vsc_, pooled_); + if (mtp_on_) add(mkq_, mvq_, mksc_, mvsc_, mpooled_); + std::atomic ok{true}; + std::vector th; // the layers are independent: one thread each + for (size_t j = 0; j < L[0].size(); ++j) + th.emplace_back([&, j] { + if (!scr::relocate_layer(L[0][j]->as(), L[1][j]->as(), L[2][j]->as(), L[3][j]->as(), + L[4][j]->as(), c_src, c_dst, n, rot)) + ok = false; + }); + for (auto& t : th) t.join(); + if (!ok) err = "metal window: SCR relocation refused"; + if (ok && n > 0 && (size_t) (c_src + n) <= slop_.size()) { // the anti-slop states move with their cells + slop_put(c_dst + n - 1); + std::memmove(slop_.data() + c_dst, slop_.data() + c_src, (size_t) n * sizeof(slop4::State)); + } + if (ok && mtp_on_ && n > 0 && c_src <= mtp_cell0_ && c_src + n >= mtp_cell1_) { + const int64_t delta = c_dst - c_src; + mtp_cell0_ = std::max(c_dst, mtp_cell0_ + delta); + mtp_cell1_ = std::max(c_dst, mtp_cell1_ + delta); + mtp_sink_ = std::min(mtp_sink_, c_src); + } + return ok; +} + bool Window::commit(int n_keep, std::string& err) { if (n_keep < 1 || n_keep > last_T_) { err = "metal window: commit count out of range"; return false; } using namespace strata::kernels; + if (n_keep < last_T_ && last_T_ <= 4 && tail_bak_.size() == tail_.size() * 4 * 128) { + for (int t = n_keep; t < last_T_; ++t) { + const size_t s = (size_t) ((last_pos0_ + t) & 3) * 128; + for (size_t j = 0; j < tail_.size(); ++j) + std::memcpy(tail_[j].as() + s, tail_bak_.data() + j * 4 * 128 + s, 128 * 4); + } + } if (last_T_ > 1) { // the GDN layers: the kept prefix folded into the state (gdn_step over n_keep tokens from the window's own // inputs, with commit; bitwise the window's trajectory). K/V and indexer cells past the prefix are rewritten // when those positions run again, before anything reads them. static const bool lazy = !(std::getenv("STRATA_METAL_LAZY_FOLD") && std::string(std::getenv("STRATA_METAL_LAZY_FOLD")) == "0"); int nf = n_keep; - if (partial_) { + fold_off_ = 0; + if (gdn_ring_ > 0) { + if (!last_head_) { + if (n_keep != last_T_) { err = "metal window: a prompt window commits all its rows"; return false; } + nf = 0; + } else if (n_keep == last_T_) { + sp_ = 1 - sp_; + hp_ = (hp_ + 2) % 3; + nf = 0; + } else if (gdn_ring_ == 1) { + hp_ = (hp_ + 1) % 3; + nf = n_keep - 1; + fold_off_ = 1; + } + } else if (partial_) { if (!last_head_) { // the prompt's window committed every row in its step if (n_keep != last_T_) { err = "metal window: a prompt window commits all its rows"; return false; } nf = 0; } else { nf = n_keep - 1; // row 0 is in + fold_off_ = 1; } } if (nf > 0) { + fold_hp_ = hp_; + const int next_hp = gdn_ring_ > 0 ? ((hp_ + 1) % 3) : (1 - hp_); if (lazy) { - pend_fold_ = nf; // folded at the start of the next window (from the history parity before the turn) + pend_fold_ = nf; // folded at the start of the next window (from fold_hp_) } else { std::vector ds; nb_ = false; - fold(ds, nf, hp_); + fold(ds, nf, fold_hp_); cx_->run_graph(ds); } - hp_ = 1 - hp_; + hp_ = next_hp; } } strata::core::SessionState& ss = *ss_; @@ -1507,6 +2460,17 @@ bool Window::mtp_load(const std::string& gguf_path, int max_t, std::string& err) { const char* e = std::getenv("STRATA_METAL_MTP_Q8"); if (!(e && std::string(e) == "0")) { + // STRATA_METAL_MTP_Q4 / STRATA_METAL_MTP_Q5 = 1: the GEMVs a draft step reads most (q/o_proj, fc_embedding, + // fc_hidden, the shared expert) in Q4_0 / Q5_0 (Q4 wins where both). K/V, the router and the + // hyper-connections (gr_read_m reads Q8_0 or BF16 only) stay Q8_0. + auto on = [](const char* name) { const char* v = std::getenv(name); return v && std::atoi(v) != 0; }; + const int low = on("STRATA_METAL_MTP_Q4") ? (int) ref::Q4_0 : on("STRATA_METAL_MTP_Q5") ? (int) ref::Q5_0 : (int) ref::Q8_0; + auto low_type = [&](const std::string& nm) { + for (const char* s : {"self_attn.q_proj.", "self_attn.o_proj.", "mtp.fc_embedding.", "mtp.fc_hidden.", + "mlp.shared_expert."}) + if (nm.find(s) != std::string::npos) return low; + return (int) ref::Q8_0; + }; size_t tot = 0; std::vector names; for (const auto& kv : fm_.tensors) { @@ -1514,10 +2478,11 @@ bool Window::mtp_load(const std::string& gguf_path, int max_t, std::string& err) uint64_t ne = 1; for (uint64_t d : x.shape) ne *= d; if (x.type != ref::BF16 || x.shape.size() < 2 || x.shape[0] % 32 != 0) continue; + if (x.name.find("self_attn.indexer.") != std::string::npos) continue; // the window is never scored gguf::Tensor q = x; - q.type = ref::Q8_0; + q.type = low_type(x.name); q.offset = tot; - q.bytes = ne / 32 * 34; + q.bytes = ne / 32 * (q.type == ref::Q4_0 ? 18 : q.type == ref::Q5_0 ? 22 : 34); mq8t_[x.name] = q; names.push_back(x.name); tot += (q.bytes + 255) / 256 * 256; @@ -1533,12 +2498,43 @@ bool Window::mtp_load(const std::string& gguf_path, int max_t, std::string& err) const uint64_t ne = x.bytes / 2; tmp.resize(ne); ref::dequant(ref::BF16, (const uint8_t*) fm_.map + x.offset, tmp.data(), ne); - ref::quantize_q8_0(tmp.data(), ne, mq8_.as() + mq8t_[names[i]].offset); + const gguf::Tensor& q = mq8t_.at(names[i]); + uint8_t* dst = mq8_.as() + q.offset; + if (q.type == ref::Q4_0) quant_q4_0_lloyd(tmp.data(), ne, dst); + else if (q.type == ref::Q5_0) quant_q5_0(tmp.data(), ne, dst); + else ref::quantize_q8_0(tmp.data(), ne, dst); } }); for (auto& x : th) x.join(); } } + { + // What the draft reads from the map (the routed experts, 0.71 GB of Q2_0) is wired, as stream.cpp wires the + // dense weights: clean file pages are what macOS drops first under the arena and a prompt's scratch, and the + // draft then faults them back from the SSD (mtp_ms 212 -> 432-467 on 2K prompts). What it never reads (the + // BF16 now in mq8_, the F32 norms now in mnorm_, the indexer) is let go, interior pages only. + const char* ml = std::getenv("STRATA_METAL_MLOCK_DENSE"); + const bool mlock_on = !(ml && std::atoi(ml) == 0); + const uintptr_t base = (uintptr_t) fm_.map, end = base + fm_.map_bytes, PG = 16384; + size_t locked = 0, unlocked = 0, dropped = 0; + for (const auto& kv : fm_.tensors) { + const gguf::Tensor& x = kv.second; + const bool skip = mq8t_.count(x.name) || mnorm_off_.count(x.name) || + x.name.find("self_attn.indexer.") != std::string::npos; + if (skip) { + const uintptr_t p0 = (base + x.offset + PG - 1) & ~(PG - 1), p1 = (base + x.offset + x.bytes) & ~(PG - 1); + if (p1 > p0) { madvise((void*) p0, p1 - p0, MADV_DONTNEED); dropped += p1 - p0; } + } else if (mlock_on) { + const uintptr_t p0 = (base + x.offset) & ~(PG - 1), + p1 = std::min(end, (base + x.offset + x.bytes + PG - 1) & ~(PG - 1)); + madvise((void*) p0, p1 - p0, MADV_WILLNEED); + if (mlock((const void*) p0, p1 - p0) == 0) locked += p1 - p0; + else unlocked += p1 - p0; + } + } + std::fprintf(stderr, "strata metal: MTP map: %.2f GB wired (mlock; %.2f GB failed), %.2f GB unused released\n", + locked / 1e9, unlocked / 1e9, dropped / 1e9); + } auto A = [&](size_t bytes) { return cx_->alloc(bytes); }; const int MAXC = (int) max_cells_, MB = MAXC / 4 + 2; mR_ = A(4 * MAXT * HC * N); memb_ = A(4 * MAXT * N); men_ = A(4 * MAXT * N); me2_ = A(4 * MAXT * N); @@ -1575,6 +2571,25 @@ bool Window::mtp_load(const std::string& gguf_path, int max_t, std::string& err) mhead_t_.bytes = dvocab_.size() * rb; } } + const char* align = std::getenv("STRATA_METAL_MTP_SLOP"); + draft_slop_ = !(align && std::string(align) == "0"); + if (draft_slop_) { + dinverse_ = A(4 * n_vocab_); + int32_t* inv = dinverse_.as(); + std::fill(inv, inv + n_vocab_, -1); + if (dvocab_.empty()) { + for (int i = 0; i < n_vocab_; ++i) inv[i] = i; + } else { + for (size_t i = 0; i < dvocab_.size(); ++i) { + if (inv[dvocab_[i]] >= 0) { + err = "metal mtp: slop alignment requires unique draft vocabulary ids"; + return false; + } + inv[dvocab_[i]] = (int32_t) i; + } + } + dslop_ = A(sizeof(slop4::State)); + } mtp_on_ = true; { const int ty = T("token_embd.weight").type; @@ -1582,7 +2597,8 @@ bool Window::mtp_load(const std::string& gguf_path, int max_t, std::string& err) const char* e = std::getenv("STRATA_METAL_MTP_CHAIN"); chain_ = !(e && std::string(e) == "0") && std::find(std::begin(ok_types), std::end(ok_types), ty) != std::end(ok_types); if (chain_) { - lchain_ = cx_->compile({"quant.metal", "units.metal", std::string(STRATA_METAL_ENGINE_KERNELS_DEFAULT) + "/mtp_chain.metal"}); + lchain_ = cx_->compile({"quant.metal", "units.metal", std::string(STRATA_METAL_ENGINE_KERNELS_DEFAULT) + "/../slop4.h", + std::string(STRATA_METAL_ENGINE_KERNELS_DEFAULT) + "/mtp_chain.metal"}); dout_ = A(4 * MAXT); dprob_ = A(4 * MAXT); posv_ = A(4 * N); for (int i = 0; i < N; ++i) posv_.as()[i] = (float) i; dmap_ = A(4 * std::max(1, dvocab_.size())); @@ -1623,25 +2639,23 @@ bool Window::mtp_load(const std::string& gguf_path, int max_t, std::string& err) void Window::gr_read_m(std::vector& ds, const std::string& pre, mtl::Buffer& Rb, const mtl::Buffer& bo, const mtl::Buffer& inj_in, const mtl::Buffer& inj_out, bool apply, const Act& out, bool final_, - int ntok) { + int ntok, bool already_normed, size_t r_off, size_t inj_in_off) { const gguf::Tensor& tdn = TQ(pre + ".input_mix_weight_down.weight"); const gguf::Tensor& tup = TQ(pre + ".input_mix_weight_up.weight"); const gguf::Tensor& tin = final_ ? tdn : TQ(pre + ".block_inject_weight.weight"); const bool q8 = tdn.type == ref::Q8_0 && tup.type == ref::Q8_0 && tin.type == ref::Q8_0; const mtl::Buffer& WB = q8 ? mq8_ : Wm_; GrArgs a{N, HC, LR, (uint32_t) ntok, apply ? 1u : 0u, eps_}; - mtl::Dispatch d1; - d1.p = P(lib_, "gr_norm", {}, "gr_norm"); - d1.args = {mtl::Arg::B(Rb), mtl::Arg::B(bo), mtl::Arg::B(inj_in), mtl::Arg::B(mnorm_, NM(pre + ".hc_norm.weight")), - mtl::Arg::B(xn_), mtl::Arg::V(a)}; - d1.grid[0] = HC; d1.grid[1] = ntok; d1.tg[0] = 256; - push(ds, d1); + if (apply || !already_normed) + gr_norm(ds, Rb, bo, inj_in, mnorm_, NM(pre + ".hc_norm.weight"), apply, ntok, r_off, inj_in_off); mtl::Dispatch d2; - const char* dn = q8 ? "gr_down_q8_8_8" : "gr_down_8_8"; - d2.p = P(lib_, dn, mtl::Constants().i(1, ntok), std::string(dn) + "#" + std::to_string(ntok)); + const bool r2 = ntok > 4; + std::string dn = q8 ? "gr_down_q8_8_8" : "gr_down_8_8"; + if (r2) dn += "_r2"; + d2.p = P(lib_, dn, mtl::Constants().i(1, ntok), dn + "#" + std::to_string(ntok)); d2.args = {mtl::Arg::B(WB, tdn.offset), mtl::Arg::B(WB, tin.offset), mtl::Arg::B(xn_), mtl::Arg::B(lo_), mtl::Arg::B(inj_out), mtl::Arg::V(a)}; - d2.grid[0] = LR + HC; d2.tg[0] = 256; + d2.grid[0] = r2 ? (LR + HC) / 2 : (LR + HC); d2.tg[0] = 256; push(ds, d2); mtl::Dispatch d3; const char* up = q8 ? "gr_up_q8" : "gr_up"; @@ -1655,9 +2669,10 @@ void Window::gr_read_m(std::vector& ds, const std::string& pre, m void Window::mtp_encode(std::vector& ds, const mtl::Buffer& srcR, size_t src_off, int n, const int32_t* tokens, int64_t cell0, bool full, int erow, int slot) { - struct RmsArgs { uint32_t n; float eps; }; + static_assert(N == 2560, "mtp_combine_gnorm retains ten values per thread"); + struct RmsXpArgs { uint32_t n, xp_n, plane_tok, split, keep; float eps; }; struct XpArgs { uint32_t n, n_tok, plane_tok, stride; }; - struct CombArgs { uint32_t n, hc, n_tok, h_stride; }; + struct CombNormArgs { uint32_t n, hc; float eps; }; // ---- the embedding rows (host tokens: dequantized on the CPU into the shared buffer) const gguf::Tensor& te = T("token_embd.weight"); const size_t rb = ref::row_bytes(te.type, N); @@ -1666,12 +2681,13 @@ void Window::mtp_encode(std::vector& ds, const mtl::Buffer& srcR, ref::dequant(te.type, (const uint8_t*) W_.ptr + te.offset + (size_t) tokens[i] * rb, (membp_ ? *membp_ : memb_).as() + (size_t) (erow + i) * N, N); const std::string L = "mtp.layers.0."; - auto rms = [&](const mtl::Buffer& x, size_t xo, const std::string& w, mtl::Buffer& y, int width, int rows) { + auto rms_xp = [&](const mtl::Buffer& x, size_t xo, const std::string& w, mtl::Buffer& y, const Act& out, + int width, bool split, bool keep) { mtl::Dispatch d; - d.p = P(lmtp_, "mtp_rms", {}, "mtp_rms"); - RmsArgs a{(uint32_t) width, eps_}; - d.args = {mtl::Arg::B(x, xo), mtl::Arg::B(mnorm_, NM(w)), mtl::Arg::B(y), mtl::Arg::V(a)}; - d.grid[0] = rows; d.tg[0] = 256; + d.p = P(lmtp_, "mtp_rms_xprep", {}, "mtp_rms_xprep"); + RmsXpArgs a{(uint32_t) width, N, MAXT, split ? 1u : 0u, keep ? 1u : 0u, eps_}; + d.args = {mtl::Arg::B(x, xo), mtl::Arg::B(mnorm_, NM(w)), mtl::Arg::B(y), mtl::Arg::B(out.b), mtl::Arg::V(a)}; + d.grid[0] = n; d.tg[0] = 256; push(ds, d); }; auto xp = [&](const mtl::Buffer& x, size_t xo, const Act& out, int width, int rows, int stride) { @@ -1682,58 +2698,136 @@ void Window::mtp_encode(std::vector& ds, const mtl::Buffer& srcR, d.grid[0] = (width / 16 * rows + 255) / 256; d.tg[0] = 256; push(ds, d); }; - auto gm = [&](const std::string& name, const Act& x, const mtl::Buffer& y, size_t y_stride, int rows, size_t y_off = 0) { + // the GEMV tile (L, R, S, K) of an MTP tensor + auto tile = [&](int ty, int n_in, int n_out, int rows) { + // STRATA_METAL_MTP_TILE=L,R,S,K: the tile of the wide ones (n_out > 640, n_in < 6144), one of gemv.metal's GV_SET + static const std::array wide = [] { + std::array t{0, 0, 0, 0}; + if (const char* e = std::getenv("STRATA_METAL_MTP_TILE")) std::sscanf(e, "%d,%d,%d,%d", &t[0], &t[1], &t[2], &t[3]); + return t; + }(); + int Lq = 16, Rr = 1, S = 4, K = 1; + const bool m5 = m5_layout(); + if (n_in >= 6144) { + if (m5 && ty == ref::Q8_0 && rows == 1) { Lq = 32; Rr = 1; S = 8; K = 8; } + else { Lq = 32; Rr = 1; S = 8; K = 2; } + } else if (n_in == 640) { + if (m5 && ty == ref::Q8_0) { Lq = rows >= 3 ? 4 : 32; Rr = 1; S = 4; K = 1; } + else { Lq = 8; Rr = 1; S = 4; K = 1; } + } else if (n_out <= 640) { + if (m5 && ty == ref::Q8_0 && rows == 1 && n_out >= 512) { Lq = 32; Rr = 1; S = 4; K = 4; } + else if (m5 && ty == ref::Q8_0 && rows >= 2) { Lq = 16; Rr = 1; S = 4; K = 1; } + else { Lq = 32; Rr = 1; S = 4; K = 1; } + } else if (wide[0] > 0) { + Lq = wide[0]; Rr = wide[1]; S = wide[2]; K = wide[3]; + } else if (m5 && ty == ref::Q8_0) { + if (rows == 1) { Lq = 32; Rr = 1; S = 4; K = 4; } + else if (rows <= 3) { Lq = 32; Rr = 2; S = 4; K = 1; } + else { Lq = 16; Rr = 1; S = 4; K = 1; } + } else if (m5 && (ty == ref::Q4_0 || ty == ref::Q5_0)) { + Lq = 32; Rr = 2; S = 4; K = 1; + } else if (ty == ref::Q2_0 || ty == ref::BF16) { + Lq = 32; Rr = 1; S = 4; K = 1; + } + return std::array{Lq, Rr, S, K}; + }; + auto gm = [&](const std::string& name, const Act& x, const mtl::Buffer& y, size_t y_stride, int rows, size_t y_off = 0, int x_row = 0) { // a GEMV on the MTP model's tensor (the window's gemv() reads the main model's buffer) const gguf::Tensor& tt = TQ(name); const int n_in = (int) tt.shape[0], n_out = tt.shape.size() > 1 ? (int) tt.shape[1] : 1, ty = tt.type; - int Lq = 16, Rr = 1, S = 4, K = 1; - if (n_in >= 6144) { Lq = 32; Rr = 1; S = 8; K = 2; } - else if (n_in == 640) { Lq = 8; Rr = 1; S = 4; K = 1; } - else if (n_out <= 640 || ty == ref::Q2_0 || ty == ref::BF16) { Lq = 32; Rr = 1; S = 4; K = 1; } + const auto [Lq, Rr, S, K] = tile(ty, n_in, n_out, rows); const std::string fn = "gemv_" + std::to_string(ty) + "_" + std::to_string(Lq) + "_" + std::to_string(Rr) + "_" + std::to_string(S) + "_" + std::to_string(K); GemvArgs a{(uint32_t) n_in, (uint32_t) n_out, (uint32_t) rows, (uint32_t) ref::row_bytes(ty, n_in), (uint32_t) n_in, (uint32_t) y_stride}; + const size_t x_off = (size_t) x_row * (size_t) n_in * 4, s16_off = (size_t) x_row * (size_t) (n_in / 16) * 4; mtl::Dispatch d; d.p = P(lib_, fn, mtl::Constants().i(0, rows), fn + "#" + std::to_string(rows)); - d.args = {mtl::Arg::B(MB(name), tt.offset), mtl::Arg::B(x.b, x.plane_bytes() * xkind(ty)), - mtl::Arg::B(x.b, x.plane_bytes() * 4), mtl::Arg::B(y, y_off), mtl::Arg::V(a)}; + d.args = {mtl::Arg::B(MB(name), tt.offset), mtl::Arg::B(x.b, x.plane_bytes() * xkind(ty) + x_off), + mtl::Arg::B(x.b, x.plane_bytes() * 4 + s16_off), mtl::Arg::B(y, y_off), mtl::Arg::V(a)}; const int rpt = (S / K) * (32 / Lq) * Rr; d.grid[0] = (n_out + rpt - 1) / rpt; d.tg[0] = S * 32; push(ds, d); }; + // gm(a) and gm(b) side by side (same format and shape) as one gemv_cat / gemv_catk at gm's tile: bitwise the two + // GEMVs. A draft step only (rows = 1, GPU time even): for a prefill window (16,1,4,1) it measured 1.4-4.3% slower. + // Else (STRATA_METAL_GEMV_CAT=0, or no instance) the two GEMVs in one wave. + auto gm_cat = [&](const std::string& na, const std::string& nb, const Act& x, const mtl::Buffer& ya, + const mtl::Buffer& yb, size_t y_stride, int rows) { + static const bool on = !(std::getenv("STRATA_METAL_GEMV_CAT") && std::atoi(std::getenv("STRATA_METAL_GEMV_CAT")) == 0); + const gguf::Tensor &ta = TQ(na), &tb = TQ(nb); + const int n_in = (int) ta.shape[0], n_out = ta.shape.size() > 1 ? (int) ta.shape[1] : 1, ty = ta.type; + const auto [Lq, Rr, S, K] = tile(ty, n_in, n_out, rows); + // gemv.metal's GVC (K = 1) and GVCK instances + const bool inst = K == 1 ? (Lq == 32 && Rr == 1 && (S == 4 || S == 8)) || (Lq == 16 && Rr == 1 && S == 4) || + (Lq == 32 && Rr == 2 && S == 4) || (ty == ref::Q8_0 && Lq == 8 && Rr == 1 && S == 8) + : ty == ref::Q8_0 && Lq == 32 && ((Rr == 1 && S == 4 && K == 4) || + (Rr == 2 && S == 4 && K == 2) || (Rr == 1 && S == 8 && K == 8)); + if (!on || rows != 1 || !inst || tb.type != ty || tb.shape != ta.shape || + (ty != ref::Q8_0 && ty != ref::BF16 && ty != ref::F32)) { + gm(na, x, ya, y_stride, rows); + nb_ = true; + gm(nb, x, yb, y_stride, rows); + return; + } + const std::string fn = (K == 1 ? "gemvc_" : "gemvck_") + std::to_string(ty) + "_" + std::to_string(Lq) + "_" + + std::to_string(Rr) + "_" + std::to_string(S) + (K == 1 ? "" : "_" + std::to_string(K)); + GemvArgs a{(uint32_t) n_in, (uint32_t) (2 * n_out), (uint32_t) rows, (uint32_t) ref::row_bytes(ty, n_in), + (uint32_t) n_in, (uint32_t) y_stride}; + const uint32_t sp = (uint32_t) n_out; + mtl::Dispatch d; + d.p = P(lib_, fn, mtl::Constants().i(0, rows), fn + "#" + std::to_string(rows)); + d.args = {mtl::Arg::B(MB(na), ta.offset), mtl::Arg::B(x.b, x.plane_bytes() * xkind(ty)), + mtl::Arg::B(x.b, x.plane_bytes() * 4), mtl::Arg::B(ya), mtl::Arg::V(a), mtl::Arg::B(MB(nb), tb.offset), + mtl::Arg::B(yb), mtl::Arg::V(sp)}; + const int rpt = (S / K) * (32 / Lq) * Rr; + d.grid[0] = (2 * n_out + rpt - 1) / rpt; d.tg[0] = S * 32; + push(ds, d); + }; // ---- e = fc_embedding(rms(embed(tok)) * (1 + w)), h = fc_hidden per stream (rms_10240(R) * (1 + w)), R = h + e - rms(membp_ ? *membp_ : memb_, (size_t) erow * N * 4, "mtp.pre_fc_norm_embedding.weight", men_, N, n); - xp(men_, 0, XEn_, N, n, N); + // the embedding and the hidden chains side by side (disjoint buffers): a barrier per step, not per dispatch + rms_xp(membp_ ? *membp_ : memb_, (size_t) erow * N * 4, "mtp.pre_fc_norm_embedding.weight", + men_, XEn_, N, false, false); + nb_ = true; + // fc_hidden over the n * HC streams as rows [t][c], MAXT per GEMV (one weight read for a draft or two accepted + // tokens); the norm prepares the first MAXT, mtp_xprep each further slice from the retained normalized rows. + // Every GEMV has >= 4 rows (tile() varies with rows only below 4): a stream's h is bitwise the same for any n. + static_assert(MAXT % HC == 0, "a slice holds whole tokens"); + const int hrows = n * HC; + rms_xp(srcR, src_off, "mtp.pre_fc_norm_hidden.weight", mhn_, XHn_, HC * N, true, hrows > MAXT); gm("mtp.fc_embedding.weight", XEn_, me2_, N, n); - rms(srcR, src_off, "mtp.pre_fc_norm_hidden.weight", mhn_, HC * N, n); - for (int c = 0; c < HC; ++c) { - xp(mhn_, (size_t) c * N * 4, XHn_, N, n, HC * N); - gm("mtp.fc_hidden.weight", XHn_, mh2_, N, n, (size_t) c * MAXT * N * 4); + nb_ = true; + gm("mtp.fc_hidden.weight", XHn_, mh2_, N, std::min(hrows, MAXT)); + for (int r = MAXT; r < hrows; r += MAXT) { + const int k = std::min(hrows - r, MAXT); + xp(mhn_, (size_t) r * N * 4, XHn_, N, k, N); + gm("mtp.fc_hidden.weight", XHn_, mh2_, N, k, (size_t) r * N * 4); } { mtl::Dispatch d; - d.p = P(lmtp_, "mtp_combine", {}, "mtp_combine"); - CombArgs a{N, HC, (uint32_t) n, (uint32_t) (MAXT * N)}; - d.args = {mtl::Arg::B(mh2_), mtl::Arg::B(me2_), mtl::Arg::B(mR_), mtl::Arg::V(a)}; - d.grid[0] = (n * HC * N + 255) / 256; d.tg[0] = 256; + d.p = P(lmtp_, "mtp_combine_gnorm", {}, "mtp_combine_gnorm"); + CombNormArgs a{N, HC, eps_}; + d.args = {mtl::Arg::B(mh2_), mtl::Arg::B(me2_), mtl::Arg::B(mR_), + mtl::Arg::B(mnorm_, NM(L + "attn_hyper_connection.hc_norm.weight")), mtl::Arg::B(xn_), mtl::Arg::V(a)}; + d.grid[0] = HC; d.grid[1] = n; d.tg[0] = 256; push(ds, d); } - // ---- the attention hyper-connection, K/V (and the indexer keys) into the layer's own cache - gr_read_m(ds, L + "attn_hyper_connection", mR_, mbo_, minj_[1], minj_[0], false, XA_, false, n); - gm(L + "self_attn.k_proj.weight", XA_, kcur_, 512, n); - nb_ = true; - gm(L + "self_attn.v_proj.weight", XA_, vcur_, 512, n); - nb_ = true; - gm(L + "self_attn.indexer.index_qk_proj.weight", XA_, midx_, 640, n); + // ---- the attention hyper-connection, K/V into the layer's own cache. No indexer: the attention below is a dense + // window (never qsa_score / qsa_topk), so its projection and keys would be read by nobody. + gr_read_m(ds, L + "attn_hyper_connection", mR_, mbo_, minj_[1], minj_[0], false, XA_, false, n, true); + gm_cat(L + "self_attn.k_proj.weight", L + "self_attn.v_proj.weight", XA_, kcur_, vcur_, 512, n); + const int last = n - 1; if (full) { nb_ = true; - gm(L + "self_attn.q_proj.weight", XA_, qfull_, 12288, n); + gm(L + "self_attn.q_proj.weight", XA_, qfull_, 12288, 1, 0, last); } QsaArgs qa{(uint32_t) cell0, 0, (uint32_t) max_cells_, (uint32_t) n, eps_}; { + // rows 0..27 (q, K, V) with the queries (last row only when n > 1), else rows 24..27 (QSA_ROW0 = 24) mtl::Dispatch d; - d.p = P(lqsa_, "qsa_prep", qsa_c(), qsa_k()); + d.p = !full ? P(lqsa_, "qsa_prep", qsa_c().i(28, 24), qsa_k() + "#kv") + : n > 1 ? P(lqsa_, "qsa_prep", qsa_c().b(29, true), qsa_k() + "#lastq") + : P(lqsa_, "qsa_prep", qsa_c(), qsa_k()); d.args = {mtl::Arg::B(qfull_), mtl::Arg::B(kcur_), mtl::Arg::B(vcur_), mtl::Arg::B(midx_), mtl::Arg::B(midx_, 512 * 4), mtl::Arg::B(mnorm_, NM(L + "self_attn.q_norm.weight")), mtl::Arg::B(mnorm_, NM(L + "self_attn.k_norm.weight")), mtl::Arg::B(mnorm_, NM(L + "self_attn.indexer.q_layernorm.weight")), @@ -1741,24 +2835,34 @@ void Window::mtp_encode(std::vector& ds, const mtl::Buffer& srcR, mtl::Arg::B(cs_), mtl::Arg::B(qcur_), mtl::Arg::B(qidx_), mtl::Arg::B(mkq_[0]), mtl::Arg::B(mvq_[0]), mtl::Arg::B(mksc_[0]), mtl::Arg::B(mvsc_[0]), mtl::Arg::B(mtail_[0]), mtl::Arg::B(mdead_[0]), mtl::Arg::B(mpooled_[0]), mtl::Arg::V(qa), mtl::Arg::B(mrope_)}; - d.grid[0] = 33; d.grid[1] = n; d.tg[0] = 256; + d.grid[0] = full ? 28 : 4; d.grid[1] = n; d.tg[0] = qsa_prep_tg(n); push(ds, d); } if (!full) return; - // ---- dense attention over the window of cells up to this one (the draft layer's rule), the output gate, o_proj + // ---- dense attention over the attention sink and the local sliding window of cells up to this one (default 64 + + // 256 cells), the output gate, o_proj { + static const int64_t mtp_win = std::getenv("STRATA_METAL_MTP_WIN") ? std::max(64, std::min(2050, std::atoll(std::getenv("STRATA_METAL_MTP_WIN")))) : 256; int* ids = mqids_.as() + (size_t) slot * 2051; - const int64_t first = std::max(0, cell0 - 2050); - for (int64_t c = first; c <= cell0; ++c) ids[c - first] = (int) c; + // the window's cells past mtp_cell1_ are this round's drafts + const int64_t qcell = cell0 + last; + const int64_t first = std::min(qcell, std::max(mtp_cell0_, qcell - mtp_win)); + const int64_t ns = std::min({mtp_sink_, first, 2050 - (qcell - first)}); + for (int64_t c = 0; c < ns; ++c) ids[c] = (int) c; + for (int64_t c = first; c <= qcell; ++c) ids[ns + c - first] = (int) c; + QsaArgs qa_at = qa; + qa_at.pos0 = (uint32_t) (ns + qcell - first); + qa_at.n_tok = 1; + const uint32_t nch = (uint32_t) std::max(1, (qa_at.pos0 + 1 + 63) / 64); mtl::Dispatch at; at.p = P(lqsa_, "qsa_attn", {}, "qsa_attn"); at.args = {mtl::Arg::B(qcur_), mtl::Arg::B(mkq_[0]), mtl::Arg::B(mvq_[0]), mtl::Arg::B(mksc_[0]), mtl::Arg::B(mvsc_[0]), - mtl::Arg::B(mqids_, (size_t) slot * 2051 * 4), mtl::Arg::B(part_), mtl::Arg::V(qa)}; - at.grid[0] = 33; at.grid[1] = 2; at.grid[2] = 1; at.tg[0] = 256; + mtl::Arg::B(mqids_, (size_t) slot * 2051 * 4), mtl::Arg::B(part_), mtl::Arg::V(qa_at)}; + at.grid[0] = nch; at.grid[1] = 2; at.grid[2] = 1; at.tg[0] = 256; push(ds, at); mtl::Dispatch m; m.p = P(lqsa_, "qsa_merge", {}, "qsa_merge"); - QsaArgs am = qa; + QsaArgs am = qa_at; am.n_tok = MAXT; m.args = {mtl::Arg::B(part_), mtl::Arg::B(qfull_), mtl::Arg::B(XQ_.b), mtl::Arg::V(am)}; m.grid[0] = 24; m.grid[1] = 1; m.tg[0] = 256; @@ -1766,12 +2870,11 @@ void Window::mtp_encode(std::vector& ds, const mtl::Buffer& srcR, gm(L + "self_attn.o_proj.weight", XQ_, mbo_, N, 1); } // ---- the MLP hyper-connection (the attention write folded in), the MoE, the write + final mixer, the head - gr_read_m(ds, L + "mlp_hyper_connection", mR_, mbo_, minj_[0], minj_[1], true, XF_, false, 1); + const size_t r_last_off = (size_t) last * HC * N * 4, inj_last_off = (size_t) last * HC * 4; + gr_read_m(ds, L + "mlp_hyper_connection", mR_, mbo_, minj_[0], minj_[1], true, XF_, false, 1, false, r_last_off, inj_last_off); gm(L + "mlp.gate.weight", XF_, logits_, 512, 1); nb_ = true; - gm(L + "mlp.shared_expert.gate_proj.weight", XF_, gsh_, NFF, 1); - nb_ = true; - gm(L + "mlp.shared_expert.up_proj.weight", XF_, ush_, NFF, 1); + gm_cat(L + "mlp.shared_expert.gate_proj.weight", L + "mlp.shared_expert.up_proj.weight", XF_, gsh_, ush_, NFF, 1); nb_ = true; gm(L + "mlp.shared_expert_gate.weight", XF_, sg_, 1, 1); { @@ -1790,13 +2893,13 @@ void Window::mtp_encode(std::vector& ds, const mtl::Buffer& srcR, // the down tile: Q2_0 8,1; else 4,2 - but on the M4 (no M5 layout) Q5_1 4,1 and Q8_0 8,1 (4,2 spills there: // 24-36 GB/s, mac-kern 513fc615) const bool m4 = !m5_layout(); - const int dl = dq2 || (m4 && td.type == ref::Q8_0) ? 8 : 4, - dr = dq2 || (m4 && (td.type == ref::Q8_0 || td.type == ref::Q5_1)) ? 1 : 2; + const int dl = dq2 || (!m4 && td.type == ref::IQ4_NL) || (m4 && td.type == ref::Q8_0) ? 8 : 4, + dr = dq2 ? (m4 ? 1 : 2) : (m4 && (td.type == ref::Q8_0 || td.type == ref::Q5_1)) ? 1 : 2; a.gx_gu = NFF / (gs * (32 / gl) * gr); a.gx_dn = N / (ds_ * (32 / dl) * dr); mtl::Dispatch r; r.p = P(lib_, "moe_route", {}, "moe_route"); - r.args = {mtl::Arg::B(logits_), mtl::Arg::B(ids_), mtl::Arg::B(wts_), mtl::Arg::B(plan_), mtl::Arg::V(a)}; + r.args = {mtl::Arg::B(logits_), mtl::Arg::B(ids_), mtl::Arg::B(wts_), mtl::Arg::B(plan_), mtl::Arg::V(a), mtl::Arg::B(rcounter_)}; r.tg[0] = 512; push(ds, r); { @@ -1829,18 +2932,14 @@ void Window::mtp_encode(std::vector& ds, const mtl::Buffer& srcR, dd.indirect = &plan_; dd.ind_off = P_IDN * 4; push(ds, dd); - GlueArgs ca{N, 1, N, TOPK}; - mtl::Dispatch cb; - cb.p = P(lib_, "moe_combine_sh", {}, "moe_combine_sh"); - cb.args = {mtl::Arg::B(ye_), mtl::Arg::B(plan_), mtl::Arg::B(ysh_), mtl::Arg::B(sg_), mtl::Arg::B(mbo_), mtl::Arg::V(ca)}; - cb.grid[0] = (N + 255) / 256; cb.tg[0] = 256; - push(ds, cb); + comb_defer(ds, mbo_, 1); // into the final mixer's gr_norm } // the layer's write (R updated in place: the next draft step's input residual) and the final mixer - gr_read_m(ds, "mtp.hyper_connection_mixer", mR_, mbo_, minj_[1], minj_[0], true, XA_, true, 1); + gr_read_m(ds, "mtp.hyper_connection_mixer", mR_, mbo_, minj_[1], minj_[0], true, XA_, true, 1, false, r_last_off, 0); const int64_t nv = dvocab_.empty() ? n_vocab_ : (int64_t) dvocab_.size(); if (dvocab_.empty()) gemv(ds, T("output.weight"), XA_, mlm_, (size_t) n_vocab_, 1); else gemv(ds, mhead_t_, XA_, mlm_, (size_t) nv, 1, 0, &mhead_); + if (draft_slop_ && chain_) mtp_slop(ds, slot); { ArgArgs a{(uint32_t) nv, 120}; mtl::Dispatch d1, d2; @@ -1858,26 +2957,91 @@ void Window::mtp_encode(std::vector& ds, const mtl::Buffer& srcR, // the prompt's K/V: cell cell0 + i pairs the main window's final residual row i with the token at cell0 + i + 1 bool Window::mtp_prefill(const int32_t* next_tokens, int64_t n, int64_t cell0, std::string& err) { if (!mtp_on_) { err = "metal mtp: not loaded"; return false; } - mrope_sync(cell0, n); - // every group of MAXT rows in one command buffer: the embedding rows of all of them in mpemb_ (row c + i) - if ((int64_t) mpemb_rows_ < n) { + static const int64_t tail_cap = [] { + if (const char* e = std::getenv("STRATA_METAL_MTP_PREFILL_TAIL")) + return std::max(64, std::atoll(e)); + const char* e = std::getenv("STRATA_METAL_MTP_WIN"); + return e ? std::clamp(std::atoll(e), 64, 2050) : int64_t{256}; + }(); + const int64_t skip = (n > tail_cap) ? (n - tail_cap) : 0; + const int64_t rem = n - skip; + if (rem <= 0) return true; + // the skipped rows' head is the attention sink's (cells below mtp_sink_cap()): those too, before the tail + const int64_t nsk = std::max(0, std::min(skip, mtp_sink_cap() - cell0)); + mtp_wrote(cell0, nsk); + mtp_wrote(cell0 + skip, rem); + mrope_sync(cell0, nsk); + mrope_sync(cell0 + skip, rem); + if ((int64_t) mpemb_rows_ < nsk + rem) { if (mpemb_rows_ > 0) cx_->free(mpemb_); - mpemb_rows_ = (size_t) std::max(n, 64); + mpemb_rows_ = (size_t) std::max(nsk + rem, 64); mpemb_ = cx_->alloc(4 * mpemb_rows_ * N); } membp_ = &mpemb_; std::vector ds; nb_ = false; - for (int64_t c = 0; c < n; c += MAXT) { - const int k = (int) std::min(MAXT, n - c); - mtp_encode(ds, mpsrc_ ? *mpsrc_ : mtp_src_ ? *mtp_src_ : R_, (size_t) c * HC * N * 4, k, next_tokens + c, - cell0 + c, false, (int) c, 0); + const mtl::Buffer& src = mpsrc_ ? *mpsrc_ : mtp_src_ ? *mtp_src_ : R_; + for (int64_t c = 0; c < nsk; c += MAXT) { + const int k = (int) std::min(MAXT, nsk - c); + mtp_encode(ds, src, (size_t) c * HC * N * 4, k, next_tokens + c, cell0 + c, false, (int) c, 0); + } + for (int64_t c = 0; c < rem; c += MAXT) { + const int k = (int) std::min(MAXT, rem - c); + mtp_encode(ds, src, (size_t) (skip + c) * HC * N * 4, k, next_tokens + skip + c, cell0 + skip + c, false, + (int) (nsk + c), 0); } membp_ = nullptr; cx_->run_graph(ds); return true; } +void Window::mtp_slop(std::vector& ds, int j) { + DraftSlopArgs a = draft_slop_args_; + a.j = (uint32_t) j; + mtl::Dispatch d; + d.p = P(lchain_, "mtp_slop_pen", {}, "mtp_slop_pen"); + d.args = {mtl::Arg::B(mlm_), mtl::Arg::B(dinverse_), mtl::Arg::B(dout_), + mtl::Arg::B(dslop_), mtl::Arg::V(a)}; + d.tg[0] = 256; + push(ds, d); +} + +void Window::mtp_slop_host(slop4::State state) { + const slop4::Knobs kn = slop4::knobs(); + const int nv = dvocab_.empty() ? (int) n_vocab_ : (int) dvocab_.size(); + float* lg = mlm_.as(); + const int32_t* inv = dinverse_.as(); + slop4::Pen pen[80]; + int np = slop4::penalties(state, kn, pen); + // hold is calibrated on the draft head, before the empty-turn guard, just as in finish. + const float mx = lg[tokout_.as()[0]]; + bool hold = false; + for (int i = 0; i < np; ++i) + if (inv[pen[i].id] >= 0 && lg[inv[pen[i].id]] > mx - 12.0f) hold = true; + if (hold) { + double sum = 0.0, removed = 0.0; + for (int i = 0; i < nv; ++i) sum += std::exp((double) (lg[i] - mx)); + for (int i = 0; i < np; ++i) { + const int di = inv[pen[i].id]; + if (di >= 0) removed += std::exp((double) (lg[di] - mx)) * (1.0 - std::exp(-(double) pen[i].v)); + } + const float v = (float) -std::log1p(-std::min(removed / sum, 0.999999)); + for (int32_t id : slop4::kHold) pen[np++] = {id, v}; + } + if (kn.lex > 0.0f && state.t2 == slop4::ETHINK && (state.t1 == slop4::NL || state.t1 == slop4::NL2)) { + pen[np++] = {slop4::IM_END, kn.lex}; + pen[np++] = {slop4::THINK, kn.lex}; + } + for (int i = 0; i < np; ++i) if (inv[pen[i].id] >= 0) lg[inv[pen[i].id]] -= pen[i].v; + int best = 0x7fffffff; + float bv = -INFINITY; + for (int i = 0; i < nv; ++i) + if (!std::isnan(lg[i]) && (lg[i] > bv || (lg[i] == bv && i < best))) { + bv = lg[i]; best = i; + } + tokout_.as()[0] = best; +} + void Window::mtp_pick(std::vector& ds, int j, bool embed) { struct PickArgs { uint32_t nv, j, use_map; }; const int64_t nv = dvocab_.empty() ? n_vocab_ : (int64_t) dvocab_.size(); @@ -1907,26 +3071,71 @@ bool Window::mtp_draft(int Tn, const int32_t* tokens, int64_t p, int a, int max_ float min_p, int* n_drafts, std::string& err) { if (!mtp_on_) { err = "metal mtp: not loaded"; return false; } (void) Tn; + slop4::State draft_state{}; + if (draft_slop_) { + // tokens are verifier outputs: row a predicts the token following the committed input at p+a. + draft_state = slop4::step(slop_at(p + a), tokens[a]); + draft_slop_args_ = {draft_state, slop4::knobs(), + (uint32_t) (dvocab_.empty() ? n_vocab_ : dvocab_.size()), 0, (uint32_t) n_vocab_}; + } mrope_sync(p, a + max_steps + 1); + // the catch-up rows' cells and draft 0's hold committed tokens; the later drafts' do not + mtp_wrote(p, max_steps > 0 ? a + 1 : 0); + static const DraftMinP policy = [] { + DraftMinP config; + if (!config.parse(std::getenv("STRATA_METAL_SPEC_MIN_P_SCHEDULE"), + std::getenv("STRATA_METAL_SPEC_MIN_CUM_P"))) { + std::fprintf(stderr, "metal mtp: invalid min-p schedule; using caller min_p\n"); + } + return config; + }(); + double cumulative = 1.0; int n = 0; int32_t tok = tokens[a]; + const bool prefetch_ple = ple_on_ && ss_ && ss_->ple.table && ss_->ple.table->is_open(); + int32_t ple_prev[2] = {}; + auto prefetch = [&](int32_t token) { + uint32_t rows[strata::kernels::PLE_N_HEADS]; + strata::kernels::ngram_rows(&token, ple_prev, 1, ss_->ple.consts, rows); + ss_->ple.table->prefetch_rows(rows); + ple_prev[0] = ple_prev[1]; + ple_prev[1] = token; + }; + if (prefetch_ple) { + // commit advanced through the kept inputs, not the verifier's next token. + ple_prev[0] = ss_->ple_prev[0]; + ple_prev[1] = ss_->ple_prev[1]; + prefetch(tok); + } + static const bool defer_catchup = !(std::getenv("STRATA_METAL_MTP_DEFER_CATCHUP") && + std::atoi(std::getenv("STRATA_METAL_MTP_DEFER_CATCHUP")) == 0); + int32_t step0_toks[MAXT]; + if (a > 0) { + for (int i = 0; i < a; ++i) step0_toks[i] = tokens[i]; + step0_toks[a] = tok; + } if (chain_ && max_steps > 0) { std::vector ds; nb_ = false; - if (a > 0) mtp_encode(ds, R_, 0, a, tokens, p, false); - mtp_encode(ds, R_, (size_t) a * HC * N * 4, 1, &tok, p + a, true, a, 0); + if (a > 0 && defer_catchup) { + mtp_encode(ds, R_, 0, a + 1, step0_toks, p, true, 0, 0); + } else { + if (a > 0) mtp_encode(ds, R_, 0, a, tokens, p, false); + mtp_encode(ds, R_, (size_t) a * HC * N * 4, 1, &tok, p + a, true, a, 0); + } mtp_pick(ds, 0, max_steps > 1); for (int j = 1; j < max_steps; ++j) { - mtp_encode(ds, mR_, 0, 1, nullptr, p + a + j, true, 0, j); + const size_t mr_off = (j == 1 && a > 0 && defer_catchup) ? (size_t) a * HC * N * 4 : 0; + mtp_encode(ds, mR_, mr_off, 1, nullptr, p + a + j, true, 0, j); mtp_pick(ds, j, j + 1 < max_steps); } cx_->run_graph(ds); for (int j = 0; j < max_steps; ++j) { const float pj = dprob_.as()[j]; + if (!policy.keep(j, pj, min_p, cumulative)) break; drafts[j] = dout_.as()[j]; if (probs) probs[j] = pj; ++n; - if (min_p > 0.0f && pj < min_p) break; } if (n_drafts) *n_drafts = n; return true; @@ -1935,12 +3144,18 @@ bool Window::mtp_draft(int Tn, const int32_t* tokens, int64_t p, int a, int max_ std::vector ds; nb_ = false; if (j == 0) { - if (a > 0) mtp_encode(ds, R_, 0, a, tokens, p, false); - mtp_encode(ds, R_, (size_t) a * HC * N * 4, 1, &tok, p + a, true, a, 0); + if (a > 0 && defer_catchup) { + mtp_encode(ds, R_, 0, a + 1, step0_toks, p, true, 0, 0); + } else { + if (a > 0) mtp_encode(ds, R_, 0, a, tokens, p, false); + mtp_encode(ds, R_, (size_t) a * HC * N * 4, 1, &tok, p + a, true, a, 0); + } } else { - mtp_encode(ds, mR_, 0, 1, &tok, p + a + j, true); + const size_t mr_off = (j == 1 && a > 0 && defer_catchup) ? (size_t) a * HC * N * 4 : 0; + mtp_encode(ds, mR_, mr_off, 1, &tok, p + a + j, true); } cx_->run_graph(ds); + if (draft_slop_) mtp_slop_host(draft_state); const int32_t di = tokout_.as()[0]; const int32_t d = dvocab_.empty() ? di : dvocab_[(size_t) di]; // the draft's probability: softmax of the draft head's logits at the pick @@ -1950,11 +3165,13 @@ bool Window::mtp_draft(int Tn, const int32_t* tokens, int64_t p, int a, int max_ double s = 0.0; for (int64_t i = 0; i < nv; ++i) s += std::exp((double) (lg[i] - mx)); const float pj = (float) (1.0 / s); + if (policy.enabled() && !policy.keep(j, pj, min_p, cumulative)) break; drafts[j] = d; if (probs) probs[j] = pj; tok = d; + if (draft_slop_) draft_state = slop4::step(draft_state, d); ++n; - if (min_p > 0.0f && pj < min_p) break; + if (!policy.enabled() && min_p > 0.0f && pj < min_p) break; } if (n_drafts) *n_drafts = n; return true; @@ -1990,19 +3207,113 @@ std::string gu2_tile(int t) { return "_16_64_16_2"; } std::string dn2_tile(int t) { return t == ref::IQ4_NL || t == ref::Q5_1 || t == ref::Q8_0 ? "_16_64_32_2" : "_16_64_16_2"; } +// a prompt chunk's bo in yr's rows (prompt_gemm_init): off while STRATA_METAL_DUMP_MOE reads yr after the next attention +bool prompt_bo_alias() { return std::getenv("STRATA_METAL_DUMP_MOE") == nullptr; } +// an output left unbound (nil): the kernel skips it (gemm_hc_up's fp32 mixed) +const mtl::Buffer kUnbound{}; +constexpr int kGdnWyMinT = 16; +bool gdn_wy_on(bool k4) { + const char* e = std::getenv("STRATA_METAL_PROMPT_GDN"); + return k4 && !(e && std::string(e) == "scan"); +} +struct GdnWy { size_t kqh, wyh, wyf, end; }; +GdnWy gdn_wy_scratch(int T) { + const size_t nc = (size_t) (T + 15) / 16, kqh = 2 * 4 * 16 * 16 * 128 * nc, wyh = 2 * 48 * 1024 * nc, wyf = 4 * 48 * 48 * nc; + return {0, kqh, kqh + wyh, kqh + wyh + wyf}; +} } // namespace struct Window::PromptGemm { int MT = 0, QS = 256, TM = 16, maxrows = 0, nblk_max = 0; int lib = -1, lgdn = -1, lqsa = -1; - bool k4 = false, gu_v3 = false, dn_v3 = false, flash = false; + bool k4 = false, gu_v3 = false, dn_v3 = false, flash = false, wy = false; mtl::Buffer R, bo, inj[2], xnh, V, loh, mixed, mixh, qkv, z, al, be, ygh, gqn, gkn, gvc, ggb, gkq, go, yg, qfull, kcur, vcur, iqr, ikr, qcur, qcurh, qidx, scores, qids, part, XQ, xqh, logits, ids, wts, bcnt, boff, eoff, tiles, meta, stok, sw, pos, - XS, HS, YE, yr, gsh, ush, hh, ys, sgv, emb, embh, pkey, pval, keyn, qn, gated, normW; + XS, HS, YE, yr, gsh, ush, hh, ys, sgv, emb, embh, pkey, pval, normW; + mtl::Buffer ids3; + mtl::Buffer logits2, ids2, wts2, Wr2; // SSD streaming: layer l + 1's router on layer l's FFN input (the read-ahead) mtl::Buffer WD, WU, Wqkv, Wz, Wal, Wbe, Wout, Wq, Wk, Wv, Wiq, Wik, Wr, Wsg, Wshg, Wshu, Wshd, Wpk, Wpv; + std::vector all_bufs; }; + +void Window::free_prompt_gemm() { + if (!pg_) return; + mtp_src_ = nullptr; + for (mtl::Buffer& b : pg_->all_bufs) cx_->free(b); + delete pg_; + pg_ = nullptr; +} + int Window::prompt_chunk_rows() const { return pg_ ? pg_->MT : 0; } +void Window::dump_moe_layer(const char* pfx, int l, bool dump_in, bool dump_out) { + if (!pfx || !pg_) return; + static std::string s_pfx; + static size_t s_prev_T = 0; + const size_t T = (size_t) seg_T_; + const bool first_chunk = (last_T_ == 0 && s_pfx != pfx); + if (l == 0 && dump_in && first_chunk) { + s_pfx = pfx; + s_prev_T = 0; + } + auto wr = [&](const char* ext, const void* ptr, size_t row_bytes) { + std::string fn = std::string(pfx) + ext; + const size_t chunk_bytes = T * row_bytes; + if (s_prev_T == 0) { + FILE* fp = std::fopen(fn.c_str(), l == 0 ? "wb" : "r+b"); + if (!fp) fp = std::fopen(fn.c_str(), "wb"); + if (!fp) { std::perror(("dump_moe fopen " + fn).c_str()); std::exit(2); } + if (std::fseek(fp, (long)((size_t) l * chunk_bytes), SEEK_SET) != 0 || + std::fwrite(ptr, 1, chunk_bytes, fp) != chunk_bytes) { + std::perror(("dump_moe fwrite " + fn).c_str()); std::exit(2); + } + std::fclose(fp); + } else { + // Multi-chunk prompt: splice new T rows into [NL, s_prev_T + T, row_bytes] + const size_t old_l_bytes = s_prev_T * row_bytes; + const size_t new_l_bytes = (s_prev_T + T) * row_bytes; + static std::map> prev_files; + std::vector& prev_file = prev_files[ext]; + if (l == 0) { + FILE* rf = std::fopen(fn.c_str(), "rb"); + if (!rf) { std::perror(("dump_moe multi-chunk r " + fn).c_str()); std::exit(2); } + prev_file.resize((size_t) 48 * old_l_bytes); + if (std::fread(prev_file.data(), 1, prev_file.size(), rf) != prev_file.size()) { + std::perror(("dump_moe multi-chunk fread " + fn).c_str()); std::exit(2); + } + std::fclose(rf); + } + FILE* wf = std::fopen(fn.c_str(), l == 0 ? "wb" : "r+b"); + if (!wf) { std::perror(("dump_moe multi-chunk w " + fn).c_str()); std::exit(2); } + std::fseek(wf, (long)((size_t) l * new_l_bytes), SEEK_SET); + if (std::fwrite(prev_file.data() + (size_t) l * old_l_bytes, 1, old_l_bytes, wf) != old_l_bytes || + std::fwrite(ptr, 1, chunk_bytes, wf) != chunk_bytes) { + std::perror(("dump_moe multi-chunk fwrite " + fn).c_str()); std::exit(2); + } + std::fclose(wf); + if (l == 47) prev_file.clear(); + } + }; + if (dump_in) { + wr(".mixed", pg_->mixed.ptr, N * sizeof(float)); + wr(".ids", pg_->ids.ptr, TOPK * sizeof(int32_t)); + wr(".wts", pg_->wts.ptr, TOPK * sizeof(float)); + } + if (dump_out) { + wr(".yr", pg_->yr.ptr, N * sizeof(float)); + wr(".hh", pg_->hh.ptr, NFF * sizeof(uint16_t)); + wr(".sgv", pg_->sgv.ptr, sizeof(float)); + if (l == 47) { + std::string fn_r = std::string(pfx) + ".R_post47"; + FILE* fpr = std::fopen(fn_r.c_str(), s_prev_T == 0 ? "wb" : "ab"); + if (fpr) { std::fwrite(pg_->R.ptr, 1, T * 4 * N * sizeof(float), fpr); std::fclose(fpr); } + std::string fn_i = std::string(pfx) + ".inj47"; + FILE* fpi = std::fopen(fn_i.c_str(), s_prev_T == 0 ? "wb" : "ab"); + if (fpi) { std::fwrite(pg_->inj[0].ptr, 1, T * 4 * sizeof(float), fpi); std::fclose(fpi); } + s_prev_T += T; + } + } +} bool Window::prompt_gemm_on(int64_t n) const { static const bool on = [] { @@ -2018,30 +3329,42 @@ bool Window::prompt_gemm_on(int64_t n) const { bool Window::prompt_gemm_init(std::string& err) { if (pg_) return true; - if (stream_ && !st_->ready()) return stream_arena(err); // it makes the prompt buffers after the arena + if (stream_ && !st_->ready() && !stream_arena(err)) return false; // it makes the prompt buffers after the arena const char* ec = std::getenv("STRATA_METAL_PROMPT_CHUNK"); // Streaming: every chunk makes one pass over the layer's experts that are not in the arena, so a longer chunk // streams each expert fewer times. M5 Max, UD-Q5, warm: 1024 -> 700 tok/s at 4K, 2048 -> 920, 4096 -> 1,100 (32K the // same). The longer chunk's buffers (4.3 GB) come out of the arena only while a prompt runs, so it is the default on a - // big Mac; a small one keeps 1024 (the buffers would cost it cache). + // big Mac; a 24 GB one takes 2048 (1.19 GB at a 64K context, 1024's 0.74), a smaller one keeps 1024. uint64_t phys = 0; size_t plen = sizeof phys; sysctlbyname("hw.memsize", &phys, &plen, nullptr, 0); - const int stream_chunk = phys >= (96ull << 30) ? 4096 : 1024; + const int stream_chunk = phys >= (96ull << 30) ? 4096 : phys >= (24ull << 30) ? 2048 : 1024; const int C = std::max(64, ec ? std::atoi(ec) : (stream_ ? stream_chunk : 2048)) / 64 * 64; auto pg = std::make_unique(); PromptGemm& p = *pg; p.MT = C; p.QS = std::min(C, 256); - p.lib = cx_->compile({"quant.metal", "units.metal", "mma.metal", "dense_mm.metal", "moe_mm.metal", "gr.metal", "gemv.metal", - "prompt.metal"}); - p.lgdn = cx_->compile({"gdn.metal", "gdn_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); - p.lqsa = cx_->compile({"qsa.metal", "qsa_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); + if (pg_lib_ < 0) { + pg_lib_ = cx_->compile({"quant.metal", "units.metal", "mma.metal", "dense_mm.metal", "moe_mm.metal", "gr.metal", "gemv.metal", + "prompt.metal"}); + pg_lgdn_ = cx_->compile({"gdn.metal", "gdn_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); + pg_lqsa_ = cx_->compile({"qsa.metal", "qsa_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); + } + p.lib = pg_lib_; + p.lgdn = pg_lgdn_; + p.lqsa = pg_lqsa_; p.k4 = ::coop_k4(*cx_); const char* eg = std::getenv("STRATA_METAL_PROMPT_EXPERTS"); // v2 | v3 (default: v3 with the M5's layout) p.gu_v3 = p.dn_v3 = eg ? std::string(eg) == "v3" : p.k4; const char* eq = std::getenv("STRATA_METAL_PROMPT_QSA"); // v1 | flash (default: flash with the M5's layout) p.flash = eq ? std::string(eq) == "flash" : p.k4; + p.wy = gdn_wy_on(p.k4); + // the long-context QSA selection first runs at chunk 1 (pos0 + T > 2051): compiled there, it stalls the host ~200 ms + P(p.lqsa, "qsa_score_mm", {}, "pg:qsa_score_mm"); + P(p.lqsa, "qsa_topk", {}, "pg:qsa_topk"); + for (int kvn : {1152, 1664}) + P(p.lqsa, "qsa_prep", (p.flash ? qsa_c().b(25, true) : qsa_c()).i(26, kvn), + "pg:" + (p.flash ? qsa_k() + "#qf16" : qsa_k()) + "#kv" + std::to_string(kvn)); const size_t MT = (size_t) C; p.maxrows = C * TOPK + NEXP * p.TM; p.nblk_max = (C * TOPK + 1023) / 1024; @@ -2054,8 +3377,13 @@ bool Window::prompt_gemm_init(std::string& err) { auto AB = [&](size_t bytes) -> mtl::Buffer { const size_t b = (bytes + 16383) / 16384 * 16384; if (measure) { need += b; return mtl::Buffer{}; } - if (region == nullptr) return cx_->alloc(bytes); - mtl::Buffer r = cx_->wrap(region, b); + if (region == nullptr) { + mtl::Buffer r = cx_->alloc(bytes); + p.all_bufs.push_back(r); + return r; + } + mtl::Buffer r = cx_->wrap(region, b, /*count_allocated=*/false); + p.all_bufs.push_back(r); region += b; return r; }; @@ -2063,31 +3391,53 @@ bool Window::prompt_gemm_init(std::string& err) { auto F16 = [&](size_t n) { return AB(2 * n); }; auto make = [&] { const size_t mr = (size_t) p.maxrows; - p.R = F32(MT * HC * N); p.bo = F32(MT * N); p.inj[0] = F32(MT * HC); p.inj[1] = F32(MT * HC); p.xnh = F16(MT * HC * N); - p.V = F32(MT * (LR + HC)); p.loh = F16(MT * LR); p.mixed = F32(MT * N); p.mixh = F16(MT * N); - p.qkv = F32(MT * 10240); p.z = F32(MT * 6144); p.al = F32(MT * 48); p.be = F32(MT * 48); p.yg = F32(MT * 6144); - p.ygh = F16(MT * 6144); p.gqn = F32(MT * 16 * 128); p.gkn = F32(MT * 16 * 128); p.gvc = F32(MT * 6144); + p.R = F32(MT * HC * N); p.inj[0] = F32(MT * HC); p.inj[1] = F32(MT * HC); p.xnh = F16(MT * HC * N); + p.V = p.k4 ? F32(16) : F32(MT * (LR + HC)); p.loh = F16(MT * LR); p.mixed = prompt_bo_alias() ? F32(16) : F32(MT * N); p.mixh = F16(MT * N); + // QSA (l % 4 == 3) and GDN (l % 4 != 3) layers are mutually exclusive and consume their scratch within the layer, + // so QSA's layer-local activations and dequantized weight buffers share GDN's allocations. + // qkv also holds the expert rows' HS (gate/up -> down, between the attention and layer 0's PLE projections) + p.qkv = AB(std::max(4 * MT * 12288, 2 * mr * NFF)); p.z = F32(MT * 6144); p.al = F32(MT * 96); p.yg = F32(16); + p.ygh = F16(MT * 6144); p.gqn = F32(MT * 16 * 128); p.gkn = F32(MT * 16 * 128); + p.gvc = F32(MT * std::max(6144, (size_t) MB)); p.ggb = F32(MT * 96); p.gkq = F32(MT * 16); p.go = F32(MT * 6144); - p.qfull = F32(MT * 12288); p.kcur = F32(MT * 512); p.vcur = F32(MT * 512); p.iqr = F32(MT * 512); p.ikr = F32(MT * 128); - p.qcur = F32(MT * 24 * 256); p.qcurh = F16(MT * 24 * 256); p.qidx = F32(MT * 4 * 128); p.scores = F32(MT * (size_t) MB); - p.qids = AB(4 * MT * 2051); + // kcur is QSA's fused K | V | indexer-K | indexer-Q projection, [T][1152 or 1664] in gqn's MT * 2048 floats: vcur, + // ikr and iqr are its column slices (prompt_chunk), so they have no buffers of their own + p.qfull = p.qkv; p.kcur = p.gqn; + p.qcur = p.flash ? F32(16) : p.z; p.qcurh = p.ygh; p.qidx = F32(MT * 4 * 128); p.scores = p.gvc; + p.qids = p.go; p.part = p.flash ? F32(16) : F32((size_t) p.QS * 33 * 24 * 258); - p.XQ = F32(4 * MT * 6144 + 3 * (size_t) p.QS * 6144 + MT * 384); p.xqh = F16(MT * 6144); + p.XQ = p.flash ? F32(16) : F32(4 * MT * 6144 + 3 * (size_t) p.QS * 6144 + MT * 384); p.xqh = p.flash ? p.z : F16(MT * 6144); // flash's QSA leaves z idle p.logits = F32(MT * NEXP); p.ids = AB(4 * MT * TOPK); p.wts = F32(MT * TOPK); + if (stream_) { + p.logits2 = F32(MT * NEXP); p.ids2 = AB(4 * MT * TOPK); p.wts2 = F32(MT * TOPK); + p.Wr2 = F16((size_t) NEXP * N); + if (std::getenv("STRATA_METAL_PROMPT_PRED2") && std::atoi(std::getenv("STRATA_METAL_PROMPT_PRED2")) != 0) + p.ids3 = AB(4 * MT * TOPK); + } p.bcnt = AB(4 * (size_t) p.nblk_max * NEXP); p.boff = AB(4 * (size_t) p.nblk_max * NEXP); p.eoff = AB(4 * NEXP); p.tiles = AB(8 * (mr / p.TM + 1)); p.meta = AB(64); - p.stok = AB(4 * mr); p.sw = F32(mr); p.pos = AB(4 * MT * TOPK); p.XS = F16(mr * N); p.HS = F16(mr * NFF); - p.YE = F32(mr * N); p.yr = F32(MT * N); p.gsh = F32(MT * NFF); p.ush = F32(MT * NFF); p.hh = F16(MT * NFF); + p.stok = AB(4 * mr); p.sw = F32(mr); p.pos = AB(4 * MT * TOPK); p.XS = F16(mr * N); p.HS = p.qkv; + // with mm_combine folded into every reader yr is never written; bo (the attention write, read by the FFN's HC + // read) shares ys's rows: ys lives from the shared expert to the next layer's HC read, before its attention + // GEMM writes bo + p.YE = F32(mr * N); p.yr = prompt_bo_alias() ? F32(16) : F32(MT * N); + // the shared expert runs between the MoE and the next attention (or layer 0's PLE), when the GDN's q/k rows and + // ygh are idle; hh keeps its own while STRATA_METAL_DUMP_MOE reads it after the next attention + p.gsh = p.gqn; p.ush = p.gkn; p.hh = prompt_bo_alias() ? p.ygh : F16(MT * NFF); p.ys = F32(MT * N); p.sgv = F32(MT); - p.WD = F16((size_t) (LR + HC) * HC * N); p.WU = F16((size_t) HC * N * LR); p.Wqkv = F16((size_t) 10240 * N); - p.Wz = F16((size_t) 6144 * N); p.Wal = F16((size_t) 48 * N); p.Wbe = F16((size_t) 48 * N); p.Wout = F16((size_t) N * 6144); - p.Wq = F16((size_t) 12288 * N); p.Wk = F16((size_t) 512 * N); p.Wv = F16((size_t) 512 * N); p.Wiq = F16((size_t) 512 * N); - p.Wik = F16((size_t) 128 * N); p.Wr = F16((size_t) NEXP * N); p.Wsg = F16((size_t) N); p.Wshg = F16((size_t) NFF * N); + p.bo = prompt_bo_alias() ? p.ys : F32(MT * N); + p.WD = F16((size_t) (LR + HC) * HC * N); p.WU = F16((size_t) HC * N * LR); p.Wqkv = F16((size_t) 12288 * N); + p.Wz = F16((size_t) 6144 * N); p.Wal = F16((size_t) 512 * N); p.Wbe = F16((size_t) 512 * N); p.Wout = F16((size_t) N * 6144); + // Wk holds the fused projection's rows (attn_k, attn_v, indexer k, indexer q): 1664 of Wz's 6144 + p.Wq = p.Wqkv; p.Wk = p.Wz; + p.Wr = F16((size_t) NEXP * N); p.Wsg = F16((size_t) N); p.Wshg = F16((size_t) NFF * N); p.Wshu = F16((size_t) NFF * N); p.Wshd = F16((size_t) N * NFF); if (ple_on_) { - p.emb = F32(MT * N); p.embh = F16(MT * N); p.pkey = F32(MT * 10240); p.pval = F32(MT * N); p.keyn = F32(MT * 10240); - p.qn = F32(MT * 10240); p.gated = F32(MT * 10240); p.normW = F32(MT * 10240); - p.Wpk = F16((size_t) 10240 * N); p.Wpv = F16((size_t) N * N); + // the PLE projections run after layer 0's MoE, when layer 0's attention scratch is idle until layer 1; emb + // (the gather's, on the CPU beside layer 0) keeps its own. normW lives from the PLE norms to ple_hist_adv, + // after layer 0's expert gate/up read XS and before layer 1's gather writes it + p.emb = F32(MT * N); p.normW = 2 * mr * N >= 4 * MT * 10240 ? p.XS : F32(MT * 10240); + p.embh = p.ygh; p.pkey = p.qkv; p.pval = p.z; p.Wpk = p.Wqkv; p.Wpv = p.Wz; } }; make(); @@ -2097,14 +3447,101 @@ bool Window::prompt_gemm_init(std::string& err) { if (region == nullptr) { err = "metal prompt: no arena region for the prompt buffers"; return false; } } make(); + if (need != prompt_gemm_need_bytes()) { err = "metal window: prompt_gemm_need_bytes() != need"; return false; } pg_ = pg.release(); + const char* esep_log = std::getenv("STRATA_METAL_SEPARATE_PROMPT_BUF"); + const char* ept_log = std::getenv("STRATA_METAL_PHASE_TRADE"); + const bool sep_pg_log = (ept_log && std::atoi(ept_log) == 0) || (esep_log && std::string(esep_log) == "1"); std::fprintf(stderr, "strata metal: prompts on the matrix units in chunks of %d (experts %s, QSA %s, matmul2d layout %s; " "%.2f GB of buffers%s)\n", C, p.gu_v3 ? "v3" : "v2", p.flash ? "flash" : "v1", p.k4 ? "K4" : "other", - need / 1e9, stream_ ? ", in the expert arena while a prompt runs" : ""); + need / 1e9, stream_ ? (sep_pg_log ? ", in a separate transient mmap while a prompt runs" : ", in the expert arena while a prompt runs") : ""); (void) err; return true; } +uint64_t Window::prompt_gemm_need_bytes() const { + const char* ec = std::getenv("STRATA_METAL_PROMPT_CHUNK"); + uint64_t phys = 0; + size_t plen = sizeof phys; + sysctlbyname("hw.memsize", &phys, &plen, nullptr, 0); + const int stream_chunk = phys >= (96ull << 30) ? 4096 : phys >= (24ull << 30) ? 2048 : 1024; + const int C = std::max(64, ec ? std::atoi(ec) : (stream_ ? stream_chunk : 2048)) / 64 * 64; + const int QS = std::min(C, 256); + const bool k4 = ::coop_k4(*cx_); + const char* eq = std::getenv("STRATA_METAL_PROMPT_QSA"); + const bool flash = eq ? std::string(eq) == "flash" : k4; + PromptGemm p_tm; + const size_t MT = (size_t) C; + const size_t mr = (size_t) (C * TOPK + NEXP * p_tm.TM); + const size_t nblk_max = (size_t) ((C * TOPK + 1023) / 1024); + const size_t MB = (size_t) (max_cells_ / 4 + 2); + uint64_t need = 0; + auto AB = [&](size_t bytes) { need += (bytes + 16383) / 16384 * 16384; }; + auto F32 = [&](size_t n) { AB(4 * n); }; + auto F16 = [&](size_t n) { AB(2 * n); }; + F32(MT * HC * N); F32(MT * HC); F32(MT * HC); F16(MT * HC * N); + F32(k4 ? 16 : MT * (LR + HC)); F16(MT * LR); F32(prompt_bo_alias() ? 16 : MT * N); F16(MT * N); + AB(std::max(4 * MT * 12288, 2 * mr * NFF)); F32(MT * 6144); F32(MT * 96); F32(16); + F16(MT * 6144); F32(MT * 16 * 128); F32(MT * 16 * 128); F32(MT * std::max(6144, MB)); + F32(MT * 96); F32(MT * 16); F32(MT * 6144); + if (flash) F32(16); + F32(MT * 4 * 128); + if (flash) F32(16); else F32((size_t) QS * 33 * 24 * 258); + if (flash) F32(16); else F32(4 * MT * 6144 + 3 * (size_t) QS * 6144 + MT * 384); + if (!flash) F16(MT * 6144); + F32(MT * NEXP); AB(4 * MT * TOPK); F32(MT * TOPK); + if (stream_) { + F32(MT * NEXP); AB(4 * MT * TOPK); F32(MT * TOPK); F16((size_t) NEXP * N); + if (std::getenv("STRATA_METAL_PROMPT_PRED2") && std::atoi(std::getenv("STRATA_METAL_PROMPT_PRED2")) != 0) + AB(4 * MT * TOPK); + } + AB(4 * nblk_max * NEXP); AB(4 * nblk_max * NEXP); + AB(4 * NEXP); AB(8 * (mr / p_tm.TM + 1)); AB(64); + AB(4 * mr); F32(mr); AB(4 * MT * TOPK); F16(mr * N); + F32(mr * N); F32(prompt_bo_alias() ? 16 : MT * N); + if (!prompt_bo_alias()) F16(MT * NFF); + F32(MT * N); F32(MT); + if (!prompt_bo_alias()) F32(MT * N); + F16((size_t) (LR + HC) * HC * N); F16((size_t) HC * N * LR); F16((size_t) 12288 * N); + F16((size_t) 6144 * N); F16((size_t) 512 * N); F16((size_t) 512 * N); F16((size_t) N * 6144); + F16((size_t) NEXP * N); F16((size_t) N); F16((size_t) NFF * N); + F16((size_t) NFF * N); F16((size_t) N * NFF); + if (ple_on_) { + F32(MT * N); + if (2 * mr * N < 4 * MT * 10240) F32(MT * 10240); + } + return need; +} + +void Window::prompt_region(bool hold, bool multi_chunk) { + if (!stream_ || !st_ || !st_->ready()) return; + if (hold) { + st_->hold_region(true, multi_chunk); + } else { + st_->hold_region(false); + const char* esep = std::getenv("STRATA_METAL_SEPARATE_PROMPT_BUF"); + const char* ept = std::getenv("STRATA_METAL_PHASE_TRADE"); + const bool sep_pg = (ept && std::atoi(ept) == 0) || (esep && std::string(esep) == "1"); + if (sep_pg && pg_ != nullptr) { + const char* reuse = std::getenv("STRATA_METAL_PROMPT_BUF_REUSE"); + const int mode = reuse ? std::atoi(reuse) : 0; + if (st_->reuse_region(mode)) { + mtp_src_ = nullptr; + return; + } + const ProcMem pm_before = proc_mem(); + free_prompt_gemm(); + const uint64_t unmapped = st_->release_region(); + strata::gpu::purge_device_allocations(1u << 20); + const ProcMem pm_after = proc_mem(); + if (unmapped > 0) { + std::fprintf(stderr, "strata metal stream: prompt scratch %.2f GB unmapped after prefill (phys_footprint %.2f -> %.2f GB, resident RSS %.2f -> %.2f GB)\n", + unmapped / 1e9, pm_before.footprint / 1e9, pm_after.footprint / 1e9, pm_before.rss / 1e9, pm_after.rss / 1e9); + } + } + } +} + bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::string& err) { PromptGemm& p = *pg_; p.meta.as()[2] = p.gu_v3 ? NFF / (32 * 4) : NFF / 64; @@ -2116,11 +3553,12 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri if (pend_fold_ > 0) { // the last decode window's commit, still pending std::vector fd; nb_ = false; - fold(fd, pend_fold_, 1 - hp_); + fold(fd, pend_fold_, fold_hp_); cx_->run_graph(fd); pend_fold_ = 0; } const int T = Tn; + const double tb0 = now_ms(); // ---- the embeddings into every stream, the PLE rows { const gguf::Tensor& te = TN("token_embd.weight"); @@ -2133,19 +3571,36 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri for (int c = 1; c < HC; ++c) std::memcpy(r + (size_t) c * N, r, (size_t) N * 4); } } + const double tbe = now_ms(); strata::core::SessionState& ss = *ss_; + // Under streaming, the PLE reads run on a thread beside layer 0's attention and are joined at layer 0's handshake, + // before the PLE projections read p.emb (layer 0 has no PLE). + static const bool ple_async = !(std::getenv("STRATA_METAL_PLE_ASYNC") && std::atoi(std::getenv("STRATA_METAL_PLE_ASYNC")) == 0); + const bool ple_late = ple_on_ && stream_ && ple_async; + std::vector ple_rows; + std::thread ple_th; + std::string ple_err; + bool ple_ok = true; if (ple_on_) { - std::vector rows((size_t) T * strata::kernels::PLE_N_HEADS); + ple_rows.resize((size_t) T * strata::kernels::PLE_N_HEADS); int32_t prev[2] = {ss.ple_prev[0], ss.ple_prev[1]}; for (int t = 0; t < T; ++t) { - strata::kernels::ngram_rows(&tokens[t], prev, 1, ss.ple.consts, rows.data() + (size_t) t * strata::kernels::PLE_N_HEADS); + strata::kernels::ngram_rows(&tokens[t], prev, 1, ss.ple.consts, ple_rows.data() + (size_t) t * strata::kernels::PLE_N_HEADS); prev[0] = prev[1]; prev[1] = tokens[t]; } - if (!ss.ple.table->gather_batch(rows.data(), (size_t) T, p.emb.as(), err)) return false; + auto gather = [&] { ple_ok = ss.ple.table->gather_batch(ple_rows.data(), (size_t) T, p.emb.as(), ple_err); }; + if (ple_late) { + ple_th = std::thread(gather); + ple_join_ = &ple_th; + } else { + gather(); + if (!ple_ok) { err = ple_err; return false; } + } } + const double tb1 = now_ms(); std::vector ds; - std::vector cuts; + std::vector cuts, sigs; nb_ = false; auto D = [&](const mtl::Pipeline* pp, std::vector a, uint32_t gx, uint32_t gy, uint32_t tg, uint32_t gz = 1) { mtl::Dispatch d; @@ -2159,110 +3614,207 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri return P(lib, fn, c, "pg:" + (key.empty() ? fn : key)); }; auto gemm = [&](const mtl::Buffer& A, const mtl::Buffer& B, const mtl::Buffer& Y, int M, int Nn, int K) { - const int tm = 64, tn = Nn <= 64 ? 32 : 64; - const std::string fn = "gemm_f16_" + std::to_string(tm) + "_" + std::to_string(tn) + "_4"; + // every tile gives the same bits (dense_mm_parity --tiles); the fastest on M5 per shape, from its sweep at + // M = 256..2048 (audits/opus_r11_dense.md) + int tm = 64, tn = 128, sg = 4; + auto set = [&](int m, int n, int s) { tm = m; tn = n; sg = s; }; + if (M <= 128) { + if (Nn <= 64) set(64, 32, 4); + else if (Nn <= 640 || K > 2560) set(64, Nn <= 128 ? 32 : 64, 4); + else set(M >= 128 ? 128 : 64, 64, 4); + } else if (Nn <= 64) { + M < 1024 ? set(16, 64, 2) : set(64, 32, 2); + } else if (Nn == 96) { // GDN's alpha | beta: from 640 rows one 96-wide column tile, p.mixh read once + M < 640 ? set(16, 64, 2) : set(16, 96, 2); + } else if (Nn <= 128) { + M < 768 ? set(16, 64, 2) : set(32, 64, 2); + } else if (Nn <= 640 && K > 2560) { // the HC read + M <= 512 ? set(64, 64, 4) : M < 1536 ? set(128, 64, 4) : set(32, 128, 2); + } else if (Nn <= 512 || (Nn <= 640 && M < 384)) { + M < 384 ? set(32, 64, 2) : set(32, 128, 2); + } else if (Nn <= 1664) { // QSA's fused K | V | indexer rows (1152, 1664): 128_64_4 is 10% slower at M = 1024 + M < 768 ? set(32, 128, 2) : set(64, 128, 4); + } else { + // 128-row tiles only where they pad no more rows than 64-row ones: M % 128 in 1..64 is 7-21% slower + M < 1536 && (M % 128 == 0 || M % 128 > 64) ? set(128, 64, 4) : set(64, 128, 4); + } + if (Nn > 64) + if (const char* ge = std::getenv("STRATA_PG_GEMM")) std::sscanf(ge, "%d,%d,%d", &tm, &tn, &sg); + const std::string fn = "gemm_f16_" + std::to_string(tm) + "_" + std::to_string(tn) + "_" + std::to_string(sg); PgGemm g{M, Nn, K, 0}; - D(PP(p.lib, fn), {mtl::Arg::B(A), mtl::Arg::B(B), mtl::Arg::B(Y), mtl::Arg::V(g)}, (Nn + tn - 1) / tn, (M + tm - 1) / tm, 128); + D(PP(p.lib, fn), {mtl::Arg::B(A), mtl::Arg::B(B), mtl::Arg::B(Y), mtl::Arg::V(g)}, (Nn + tn - 1) / tn, (M + tm - 1) / tm, sg * 32); }; auto dqb = [&](const mtl::Buffer& src, size_t off, int type, int n_in, int n_out, const mtl::Buffer& out, size_t out_off = 0, uint32_t perm = 0) { + const int u = n_in / pg_unit(type), tg = u <= 64 ? 64 : u <= 128 ? 128 : u % 192 == 0 ? 192 : u % 160 == 0 ? 160 : 256; PgDq a{(uint32_t) n_in, (uint32_t) n_out, (uint32_t) ref::row_bytes(type, n_in), perm}; D(PP(p.lib, "dq_f16_" + std::to_string(type)), {mtl::Arg::B(src, off), mtl::Arg::B(out, out_off), mtl::Arg::V(a)}, - (n_in / pg_unit(type) + 63) / 64, n_out, 64); + (u + tg - 1) / tg, n_out, tg); }; auto dq = [&](const gguf::Tensor& t, const mtl::Buffer& out, size_t out_off = 0, uint32_t perm = 0) { const int n_in = (int) t.shape[0], n_out = t.shape.size() > 1 ? (int) t.shape[1] : 1; - if (hc_q8_ && q8off_.count(t.name)) dqb(Q8_, q8off_.at(t.name), ref::Q8_0, n_in, n_out, out, out_off, perm); + if (hc_q8_ && q8off_.count(t.name)) dqb(Q8_, q8off_.at(t.name), q8type_.count(t.name) ? q8type_.at(t.name) : ref::Q8_0, n_in, n_out, out, out_off, perm); else dqb(W_, t.offset, t.type, n_in, n_out, out, out_off, perm); }; auto act16 = [&](const mtl::Buffer& x, const mtl::Buffer& xh, size_t n) { const uint32_t nn = (uint32_t) n; - D(PP(p.lib, "act_f16"), {mtl::Arg::B(x), mtl::Arg::B(xh), mtl::Arg::V(nn)}, (nn + 255) / 256, 1, 256); + D(PP(p.lib, "act_f16"), {mtl::Arg::B(x), mtl::Arg::B(xh), mtl::Arg::V(nn)}, (nn + 1023) / 1024, 1, 256); }; int ci = 0; - auto hc_read = [&](const std::string& pre, bool apply) { + // mm_combine folded into the next read (gr_norm_h_moec, ple_moe_gate_norm_c, gr_apply_moec): yr is formed from YE + // as it is read and never written - except while STRATA_METAL_DUMP_MOE reads yr + const bool comb_fold = prompt_bo_alias(); + const uint32_t topk = TOPK; + // moe: the write to apply is the last MoE's, formed from yr, ys, sgv (moe_out folded in) rather than read from bo + auto hc_read = [&](const std::string& pre, bool apply, bool moe) { const gguf::Tensor &tn = TN(pre + "_norm.weight"), &tdn = TN(pre + "_down.weight"), &tin = TN(pre + "_inject.weight"), &tup = TN(pre + "_up.weight"); dq(tdn, p.WD); + nb_ = true; dq(tin, p.WD, (size_t) 2 * LR * HC * N); + nb_ = true; dq(tup, p.WU, 0, HC); GrArgs a{N, HC, LR, (uint32_t) T, apply ? 1u : 0u, eps_}; PgEpi ea{(uint32_t) T, N, HC, LR}; - D(PP(p.lib, "gr_norm_h"), {mtl::Arg::B(p.R), mtl::Arg::B(p.bo), mtl::Arg::B(p.inj[ci]), mtl::Arg::B(W_, tn.offset), - mtl::Arg::B(p.xnh), mtl::Arg::V(a)}, HC, T, 256); - gemm(p.xnh, p.WD, p.V, T, LR + HC, HC * N); - D(PP(p.lib, "gr_epi_down"), {mtl::Arg::B(p.V), mtl::Arg::B(p.loh), mtl::Arg::B(p.inj[1 - ci]), mtl::Arg::V(ea)}, - (T * (LR + HC) + 255) / 256, 1, 256); + nb_ = true; + if (apply && moe && comb_fold) + D(PP(p.lib, "gr_norm_h_moec"), {mtl::Arg::B(p.R), mtl::Arg::B(p.YE), mtl::Arg::B(p.pos), mtl::Arg::B(p.sw), + mtl::Arg::B(p.ys), mtl::Arg::B(p.sgv), mtl::Arg::B(p.inj[ci]), mtl::Arg::B(W_, tn.offset), + mtl::Arg::B(p.xnh), mtl::Arg::V(a), mtl::Arg::V(topk)}, 1, T, 256); + else if (apply && moe) + D(PP(p.lib, "gr_norm_h_moe"), {mtl::Arg::B(p.R), mtl::Arg::B(p.yr), mtl::Arg::B(p.ys), mtl::Arg::B(p.sgv), + mtl::Arg::B(p.inj[ci]), mtl::Arg::B(W_, tn.offset), mtl::Arg::B(p.xnh), mtl::Arg::V(a)}, + 1, T, 256); + else + // applying: a threadgroup a token (bo read once, not hc times); else the (stream, token) grid is the faster + D(PP(p.lib, apply ? "gr_norm_h4" : "gr_norm_h"), {mtl::Arg::B(p.R), mtl::Arg::B(p.bo), mtl::Arg::B(p.inj[ci]), + mtl::Arg::B(W_, tn.offset), mtl::Arg::B(p.xnh), mtl::Arg::V(a)}, + apply ? 1 : HC, T, 256); + if (p.k4) { + // M5: gr_epi_down as the GEMM's register epilogue, V never written; the tiles from comb_fold_parity T epi bench + const auto [tm, tn, sg] = T <= 192 ? std::array{16, 64, 2} : T <= 384 ? std::array{32, 64, 2} + : T <= 512 ? std::array{32, 128, 4} : T < 1536 ? std::array{64, 128, 4} : std::array{32, 128, 2}; + const std::string fn = "gemm_hc_down_" + std::to_string(tm) + "_" + std::to_string(tn) + "_" + std::to_string(sg); + D(PP(p.lib, fn), {mtl::Arg::B(p.xnh), mtl::Arg::B(p.WD), mtl::Arg::B(p.loh), mtl::Arg::B(p.inj[1 - ci]), mtl::Arg::V(ea)}, + (LR + HC + tn - 1) / tn, (T + tm - 1) / tm, sg * 32); + } else { + gemm(p.xnh, p.WD, p.V, T, LR + HC, HC * N); + D(PP(p.lib, "gr_epi_down"), {mtl::Arg::B(p.V), mtl::Arg::B(p.loh), mtl::Arg::B(p.inj[1 - ci]), mtl::Arg::V(ea)}, + (T * (LR + HC) + 255) / 256, 1, 256); + } PgGemm g{T, HC * N, LR, 0}; - D(PP(p.lib, "gemm_hc_up_64_64_4"), {mtl::Arg::B(p.loh), mtl::Arg::B(p.WU), mtl::Arg::B(p.xnh), mtl::Arg::B(p.mixed), - mtl::Arg::B(p.mixh), mtl::Arg::V(g), mtl::Arg::V(ea)}, HC * N / 64, (T + 63) / 64, 128); + // M5 keeps each output's four HC streams in one lane; other layouts use the shared-memory epilogue. + const int up_tm = p.k4 ? 32 : 64, up_tn = p.k4 ? 128 : 64; + const char* up = p.k4 ? "gemm_hc_up_reg_32_128_4" : "gemm_hc_up_64_64_4"; + // the fp32 mixed has no reader but STRATA_METAL_DUMP_MOE + D(PP(p.lib, up), {mtl::Arg::B(p.loh), mtl::Arg::B(p.WU), mtl::Arg::B(p.xnh), + mtl::Arg::B(prompt_bo_alias() ? kUnbound : p.mixed), + mtl::Arg::B(p.mixh), mtl::Arg::V(g), mtl::Arg::V(ea)}, HC * N / up_tn, (T + up_tm - 1) / up_tm, 128); ci = 1 - ci; }; // the PLE projections (the window's gemv pair as GEMMs) - if (ple_on_) { + auto ple_proj = [&] { act16(p.emb, p.embh, (size_t) T * N); + nb_ = true; dq(TN("blk.1.ple_key.weight"), p.Wpk); + nb_ = true; dq(TN("blk.1.ple_value.weight"), p.Wpv); gemm(p.embh, p.Wpk, p.pkey, T, 10240, N); + nb_ = true; gemm(p.embh, p.Wpv, p.pval, T, N, N); - } + }; const int MB = (int) max_cells_ / 4 + 2; (void) MB; + // STRATA_METAL_PROMPT_PRED=0: the read-ahead of a chunk's next layer from the static expert profile instead + static const bool pred_on = !(std::getenv("STRATA_METAL_PROMPT_PRED") && std::atoi(std::getenv("STRATA_METAL_PROMPT_PRED")) == 0); for (int l = 0; l < 48; ++l) { const std::string b = "blk." + std::to_string(l) + "."; - hc_read(b + "hc_attn", l > 0 && !(l == 1 && ple_on_)); + hc_read(b + "hc_attn", l > 0 && !(l == 1 && ple_on_), true); if (l % 4 != 3) { const int gi = l - l / 4; + nb_ = true; dq(TN(b + "attn_qkv.weight"), p.Wqkv); + nb_ = true; dq(TN(b + "attn_gate.weight"), p.Wz); + // alpha's and beta's 48 rows as one 96-row GEMM (one dispatch, one p.mixh read): al is [T][96], alpha then beta, + // and gdn_pre reads the interleave from the two pointers + nb_ = true; dq(TN(b + "ssm_alpha.weight"), p.Wal); - dq(TN(b + "ssm_beta.weight"), p.Wbe); + nb_ = true; + dq(TN(b + "ssm_beta.weight"), p.Wal, (size_t) 2 * 48 * N); + nb_ = true; dq(TN(b + "ssm_out.weight"), p.Wout); gemm(p.mixh, p.Wqkv, p.qkv, T, 10240, N); + nb_ = true; gemm(p.mixh, p.Wz, p.z, T, 6144, N); - gemm(p.mixh, p.Wal, p.al, T, 48, N); - gemm(p.mixh, p.Wbe, p.be, T, 48, N); + nb_ = true; + gemm(p.mixh, p.Wal, p.al, T, 96, N); PgGdnP ga{(uint32_t) T, 16, 48, 10240, eps_, 0}; + const int hp_next = gdn_ring_ > 0 ? ((hp_ + 1) % 3) : (1 - hp_); D(PP(p.lgdn, "gdn_pre"), {mtl::Arg::B(p.qkv), mtl::Arg::B(ghist_[hp_][(size_t) gi]), mtl::Arg::B(W_, TN(b + "ssm_conv1d.weight").offset), - mtl::Arg::B(p.al), mtl::Arg::B(p.be), mtl::Arg::B(W_, TN(b + "ssm_dt.bias").offset), mtl::Arg::B(W_, TN(b + "ssm_a").offset), + mtl::Arg::B(p.al), mtl::Arg::B(p.al, 4 * 48), mtl::Arg::B(W_, TN(b + "ssm_dt.bias").offset), mtl::Arg::B(W_, TN(b + "ssm_a").offset), mtl::Arg::B(p.gqn), mtl::Arg::B(p.gkn), mtl::Arg::B(p.gvc), mtl::Arg::B(p.ggb), mtl::Arg::B(p.gkq), - mtl::Arg::B(ghist_[1 - hp_][(size_t) gi]), mtl::Arg::V(ga)}, 16 + 24, T, 256); - D(PP(p.lgdn, "gdn_scanr_4", mtl::Constants().i(6, 0), "gdn_scanr_4"), - {mtl::Arg::B(gst_[(size_t) gi]), mtl::Arg::B(p.gqn), mtl::Arg::B(p.gkn), mtl::Arg::B(p.gvc), mtl::Arg::B(p.ggb), - mtl::Arg::B(p.gkq), mtl::Arg::B(p.go), mtl::Arg::V(ga)}, 48 * 4, 1, 32 * 4); + mtl::Arg::B(ghist_[hp_next][(size_t) gi]), mtl::Arg::V(ga)}, (16 + 48) / 8, T, 256); + const GdnWy wy = gdn_wy_scratch(T); + if (p.wy && T >= kGdnWyMinT && wy.end <= p.qkv.bytes) { + D(PP(p.lgdn, "gdn_wy_prep"), + {mtl::Arg::B(p.gqn), mtl::Arg::B(p.gkn), mtl::Arg::B(p.ggb), mtl::Arg::B(p.qkv, wy.kqh), mtl::Arg::B(p.qkv, wy.wyh), + mtl::Arg::B(p.qkv, wy.wyf), mtl::Arg::V(ga)}, 16, (T + 15) / 16, 128); + D(PP(p.lgdn, "gdn_wy_scan"), + {mtl::Arg::B(gst_[sp_][(size_t) gi]), mtl::Arg::B(p.gvc), mtl::Arg::B(p.qkv, wy.kqh), mtl::Arg::B(p.qkv, wy.wyh), + mtl::Arg::B(p.qkv, wy.wyf), mtl::Arg::B(p.go), mtl::Arg::V(ga)}, 48 * 8, 1, 64); + } else { + D(PP(p.lgdn, "gdn_scanr_4", mtl::Constants().i(6, 0), "gdn_scanr_4"), + {mtl::Arg::B(gst_[sp_][(size_t) gi]), mtl::Arg::B(p.gqn), mtl::Arg::B(p.gkn), mtl::Arg::B(p.gvc), mtl::Arg::B(p.ggb), + mtl::Arg::B(p.gkq), mtl::Arg::B(p.go), mtl::Arg::V(ga)}, 48 * 4, 1, 32 * 4); + } D(PP(p.lgdn, "gdn_post"), {mtl::Arg::B(p.go), mtl::Arg::B(p.z), mtl::Arg::B(W_, TN(b + "ssm_norm.weight").offset), - mtl::Arg::B(p.ygh), mtl::Arg::B(p.yg), mtl::Arg::V(ga)}, 48, T, 128); + mtl::Arg::B(p.ygh), mtl::Arg::B(p.yg), mtl::Arg::V(ga)}, 48 / 4, T, 128); gemm(p.ygh, p.Wout, p.bo, T, N, 6144); } else { const int qi = l / 4; + // K, V, indexer K (and past the budget indexer Q) as one GEMM in Wk / kcur: rows 0..511 K, 512..1023 V, + // 1024..1151 indexer K, 1152..1663 indexer Q. The same bits as the separate GEMMs, one dispatch and one + // p.mixh read instead of 3-4 (audits/opus_r16_kvik.md); qsa_prep reads the slices at the row stride + // QSA_KV_STRIDE (function constant 26) + const bool sel = pos0 + T > 2051; + const int kvn = sel ? 1664 : 1152; + nb_ = true; dq(TN(b + "attn_q.weight"), p.Wq); + nb_ = true; dq(TN(b + "attn_k.weight"), p.Wk); - dq(TN(b + "attn_v.weight"), p.Wv); - dq(TN(b + "indexer.q_proj.weight"), p.Wiq); - dq(TN(b + "indexer.k_proj.weight"), p.Wik); + nb_ = true; + dq(TN(b + "attn_v.weight"), p.Wk, (size_t) 2 * 512 * N); + nb_ = true; + dq(TN(b + "indexer.k_proj.weight"), p.Wk, (size_t) 2 * 1024 * N); + if (sel) { + nb_ = true; + dq(TN(b + "indexer.q_proj.weight"), p.Wk, (size_t) 2 * 1152 * N); + } + nb_ = true; dq(TN(b + "attn_output.weight"), p.Wout); gemm(p.mixh, p.Wq, p.qfull, T, 12288, N); - gemm(p.mixh, p.Wk, p.kcur, T, 512, N); - gemm(p.mixh, p.Wv, p.vcur, T, 512, N); - gemm(p.mixh, p.Wiq, p.iqr, T, 512, N); - gemm(p.mixh, p.Wik, p.ikr, T, 128, N); + nb_ = true; + gemm(p.mixh, p.Wk, p.kcur, T, kvn, N); QsaArgs a{(uint32_t) pos0, 0, (uint32_t) max_cells_, (uint32_t) T, eps_}; - D(PP(p.lqsa, "qsa_prep", qsa_c(), qsa_k()), - {mtl::Arg::B(p.qfull), mtl::Arg::B(p.kcur), mtl::Arg::B(p.vcur), mtl::Arg::B(p.iqr), mtl::Arg::B(p.ikr), + const mtl::Constants qc = (p.flash ? qsa_c().b(25, true) : qsa_c()).i(26, kvn); + D(PP(p.lqsa, "qsa_prep", qc, (p.flash ? qsa_k() + "#qf16" : qsa_k()) + "#kv" + std::to_string(kvn)), + {mtl::Arg::B(p.qfull), mtl::Arg::B(p.kcur), mtl::Arg::B(p.kcur, 4 * 512), mtl::Arg::B(p.kcur, 4 * 1152), + mtl::Arg::B(p.kcur, 4 * 1024), mtl::Arg::B(W_, TN(b + "attn_q_norm.weight").offset), mtl::Arg::B(W_, TN(b + "attn_k_norm.weight").offset), mtl::Arg::B(W_, TN(b + "indexer.q_norm.weight").offset), mtl::Arg::B(W_, TN(b + "indexer.k_norm.weight").offset), - mtl::Arg::B(cs_), mtl::Arg::B(p.qcur), mtl::Arg::B(p.qidx), mtl::Arg::B(kq_[(size_t) qi]), mtl::Arg::B(vq_[(size_t) qi]), - mtl::Arg::B(ksc_[(size_t) qi]), mtl::Arg::B(vsc_[(size_t) qi]), mtl::Arg::B(tail_[(size_t) qi]), - mtl::Arg::B(dead_[(size_t) qi]), mtl::Arg::B(pooled_[(size_t) qi]), mtl::Arg::V(a), mtl::Arg::B(mrope_)}, 33, T, 256); - if (pos0 + T > 2051) + mtl::Arg::B(cs_), mtl::Arg::B(p.flash ? p.qcurh : p.qcur), mtl::Arg::B(p.qidx), mtl::Arg::B(kq_[(size_t) qi]), + mtl::Arg::B(vq_[(size_t) qi]), mtl::Arg::B(ksc_[(size_t) qi]), mtl::Arg::B(vsc_[(size_t) qi]), + mtl::Arg::B(tail_[(size_t) qi]), mtl::Arg::B(dead_[(size_t) qi]), mtl::Arg::B(pooled_[(size_t) qi]), + mtl::Arg::V(a), mtl::Arg::B(mrope_)}, 33, T, qsa_prep_tg(T)); + if (sel) D(PP(p.lqsa, "qsa_score_mm"), {mtl::Arg::B(p.qidx), mtl::Arg::B(pooled_[(size_t) qi]), mtl::Arg::B(dead_[(size_t) qi]), mtl::Arg::B(p.scores), mtl::Arg::V(a)}, (uint32_t) (((pos0 + T) / 4 + 1 + 63) / 64), (T + 15) / 16, 128); - D(PP(p.lqsa, "qsa_topk"), {mtl::Arg::B(p.scores), mtl::Arg::B(p.qids), mtl::Arg::V(a)}, T, 1, 1024); + if (sel || !p.flash) + D(PP(p.lqsa, "qsa_topk"), {mtl::Arg::B(p.scores), mtl::Arg::B(p.qids), mtl::Arg::V(a)}, T, 1, 1024); if (p.flash) { - const uint32_t nq = (uint32_t) T * 24 * 256; - D(PP(p.lqsa, "qsa_q_f16"), {mtl::Arg::B(p.qcur), mtl::Arg::B(p.qcurh), mtl::Arg::V(nq)}, (nq + 255) / 256, 1, 256); D(PP(p.lqsa, "qsa_flash", mtl::Constants().b(7, p.k4), "qsa_flash"), {mtl::Arg::B(p.qcurh), mtl::Arg::B(kq_[(size_t) qi]), mtl::Arg::B(vq_[(size_t) qi]), mtl::Arg::B(ksc_[(size_t) qi]), mtl::Arg::B(vsc_[(size_t) qi]), mtl::Arg::B(p.qids), mtl::Arg::B(p.qfull), mtl::Arg::B(p.xqh), mtl::Arg::V(a)}, @@ -2287,28 +3839,95 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri // layer 0's attention write applied to R happens in the FFN read below; the PLE block comes after the FFN } // ---- the MoE - hc_read(b + "hc_ffn", true); + hc_read(b + "hc_ffn", true, false); const gguf::Tensor &tr = TN(b + "ffn_gate_inp.weight"), &tsg = TN(b + "ffn_gate_inp_shexp.weight"), &tg_ = TN(b + "ffn_gate_shexp.weight"), &tu_ = TN(b + "ffn_up_shexp.weight"), &td_ = TN(b + "ffn_down_shexp.weight"), &eg = TN(b + "ffn_gate_exps.weight"), &ed = TN(b + "ffn_down_exps.weight"); - dq(tr, p.Wr); + // under streaming the next layer's router is read ahead beside this one's: its dequant, GEMM and top-k in this + // one's waves (disjoint buffers; the two GEMMs read p.mixh together) + const bool ahead = stream_ && pred_on && l + 1 < 48; + // ... into Wr and Wr2 in turn: layer l + 1's router is the one read ahead (unless PRED2 reuses Wr2 for l + 2) + const bool reuse_wr = stream_ && pred_on && !p.ids3.ptr; + const mtl::Buffer& w_r = reuse_wr && (l & 1) ? p.Wr2 : p.Wr; + const mtl::Buffer& w_r2 = reuse_wr && (l & 1) ? p.Wr : p.Wr2; + if (l == 0 || !reuse_wr) { + nb_ = true; + dq(tr, w_r); + } + nb_ = true; dq(tsg, p.Wsg); + nb_ = true; dq(tg_, p.Wshg); + nb_ = true; dq(tu_, p.Wshu); + nb_ = true; dq(td_, p.Wshd); - gemm(p.mixh, p.Wr, p.logits, T, NEXP, N); - const uint32_t nt = (uint32_t) T; - D(PP(p.lib, "route_topk"), {mtl::Arg::B(p.logits), mtl::Arg::B(p.ids), mtl::Arg::B(p.wts), mtl::Arg::V(nt)}, (T + 7) / 8, 1, 256); + if (ahead) { + nb_ = true; + dq(TN("blk." + std::to_string(l + 1) + ".ffn_gate_inp.weight"), w_r2); + } + gemm(p.mixh, w_r, p.logits, T, n_exp_, N); + if (ahead) { + nb_ = true; + gemm(p.mixh, w_r2, p.logits2, T, n_exp_, N); + } + const uint32_t nt = (uint32_t) T | ((uint32_t) n_exp_ << 16); + const float p_moce = stream_ ? (std::getenv("STRATA_METAL_PROMPT_MOCE") ? (float) std::atof(std::getenv("STRATA_METAL_PROMPT_MOCE")) + : ((pos0 > 0) ? 320335.0f : 0.0f)) : 0.0f; + D(PP(p.lib, "route_topk"), {mtl::Arg::B(p.logits), mtl::Arg::B(p.ids), mtl::Arg::B(p.wts), mtl::Arg::V(nt), stream_ ? mtl::Arg::B(st_->smap(), (size_t) l * 512 * 4) : mtl::Arg::B(p.ids), mtl::Arg::V(p_moce)}, (T + 7) / 8, 1, 256); + if (ahead) { + nb_ = true; + D(PP(p.lib, "route_topk"), {mtl::Arg::B(p.logits2), mtl::Arg::B(p.ids2), mtl::Arg::B(p.wts2), mtl::Arg::V(nt), + mtl::Arg::B(st_->smap(), (size_t) (l + 1) * 512 * 4), mtl::Arg::V(p_moce)}, (T + 7) / 8, 1, 256); + if (p.ids3.ptr && l + 2 < 48) { // STRATA_METAL_PROMPT_PRED2: two layers ahead, after the first + const int ln = l + 2; + dq(TN("blk." + std::to_string(ln) + ".ffn_gate_inp.weight"), p.Wr2); + gemm(p.mixh, p.Wr2, p.logits2, T, n_exp_, N); + D(PP(p.lib, "route_topk"), {mtl::Arg::B(p.logits2), mtl::Arg::B(p.ids3), mtl::Arg::B(p.wts2), mtl::Arg::V(nt), + mtl::Arg::B(st_->smap(), (size_t) ln * 512 * 4), mtl::Arg::V(p_moce)}, (T + 7) / 8, 1, 256); + } + } + if (stream_ && comb_fold) sigs.push_back((int) ds.size()); const int nent = T * TOPK, nblk = (nent + 1023) / 1024; - PgPlan pa{(uint32_t) nent, TOPK, NEXP, (uint32_t) p.TM, (uint32_t) nblk}; - const uint32_t nclr = (uint32_t) p.maxrows; - D(PP(p.lib, "mm_clear"), {mtl::Arg::B(p.stok), mtl::Arg::V(nclr)}, (p.maxrows + 255) / 256, 1, 256); + PgPlan pa{(uint32_t) nent, TOPK, (uint32_t) n_exp_, (uint32_t) p.TM, (uint32_t) nblk}; + // with comb_fold the shared expert runs beside the plan, wave for wave (disjoint buffers): mm_hist | gate (and up) + // GEMMs, mm_scan | swiglu, mm_place | down GEMM; mm_scan's one threadgroup no longer holds the GPU alone. + // M5 past 160 rows: swiglu_h as the up GEMM's register epilogue, ush never written (the up GEMM waits for the + // gate's; at <= 160 rows the two overlapped are the faster) + const bool swu = p.k4 && T > 160; D(PP(p.lib, "mm_hist"), {mtl::Arg::B(p.ids), mtl::Arg::B(p.bcnt), mtl::Arg::V(pa)}, nblk, 1, 1024); + if (comb_fold) { + nb_ = true; + gemm(p.mixh, p.Wshg, p.gsh, T, NFF, N); + if (!swu) { + nb_ = true; + gemm(p.mixh, p.Wshu, p.ush, T, NFF, N); + } + nb_ = true; + gemm(p.mixh, p.Wsg, p.sgv, T, 1, N); + } D(PP(p.lib, "mm_scan"), {mtl::Arg::B(p.bcnt), mtl::Arg::B(p.boff), mtl::Arg::B(p.eoff), mtl::Arg::B(p.tiles), - mtl::Arg::B(p.meta), mtl::Arg::V(pa)}, 1, 1, NEXP); + mtl::Arg::B(p.meta), mtl::Arg::V(pa), mtl::Arg::B(p.stok)}, 1, 1, (uint32_t) n_exp_); + if (comb_fold) { + nb_ = true; + if (swu) { + const int tm = T < 384 ? 32 : 64, tn = T < 384 ? 64 : 128, sg = T < 384 ? 2 : 4; + const PgGemm g{T, NFF, N, 0}; + D(PP(p.lib, T < 384 ? "gemm_swiglu_up_32_64_2" : "gemm_swiglu_up_64_128_4"), + {mtl::Arg::B(p.mixh), mtl::Arg::B(p.Wshu), mtl::Arg::B(p.gsh), mtl::Arg::B(p.hh), mtl::Arg::V(g)}, + (NFF + tn - 1) / tn, (T + tm - 1) / tm, sg * 32); + } else { + const uint32_t nh = (uint32_t) T * NFF; + D(PP(p.lib, "swiglu_h"), {mtl::Arg::B(p.gsh), mtl::Arg::B(p.ush), mtl::Arg::B(p.hh), mtl::Arg::V(nh)}, (nh + 255) / 256, 1, 256); + } + } D(PP(p.lib, "mm_place"), {mtl::Arg::B(p.ids), mtl::Arg::B(p.wts), mtl::Arg::B(p.boff), mtl::Arg::B(p.eoff), mtl::Arg::B(p.stok), mtl::Arg::B(p.sw), mtl::Arg::B(p.pos), mtl::Arg::V(pa)}, nblk, 1, 1024); + if (comb_fold) { + nb_ = true; + gemm(p.hh, p.Wshd, p.ys, T, N, NFF); + } PgMm ma{eg.offset, TN(b + "ffn_up_exps.weight").offset, ed.offset, 0, 0, (uint32_t) ref::row_bytes(eg.type, N), (uint32_t) ref::row_bytes(ed.type, NFF), N, NFF}; ma.gu_stride = (uint64_t) ma.gu_row * NFF; @@ -2317,8 +3936,8 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri const uint32_t ne_ = N; { mtl::Dispatch g; - g.p = PP(p.lib, "mm_gather"); - g.args = {mtl::Arg::B(p.mixed), mtl::Arg::B(p.stok), mtl::Arg::B(p.XS), mtl::Arg::V(ne_)}; + g.p = PP(p.lib, "mm_gather_h"); + g.args = {mtl::Arg::B(p.mixh), mtl::Arg::B(p.stok), mtl::Arg::B(p.XS), mtl::Arg::V(ne_)}; g.tg[0] = 256; g.indirect = &p.meta; g.ind_off = 8 * 4; @@ -2349,7 +3968,10 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri arena_args(gu); gu.indirect = &p.meta; gu.ind_off = 2 * 4; - if (stream_) cuts.push_back((int) ds.size()); // the host's turn: this layer's experts put in place + if (stream_) { + if (!comb_fold) sigs.push_back((int) ds.size()); + cuts.push_back((int) ds.size()); // the host's turn: this layer's experts put in place + } push(ds, gu); mtl::Dispatch dn; if (p.dn_v3) { @@ -2371,75 +3993,87 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri dn.ind_off = 5 * 4; push(ds, dn); const uint32_t ca[4] = {(uint32_t) T, N, TOPK, 0}; - D(PP(p.lib, "mm_combine"), {mtl::Arg::B(p.YE), mtl::Arg::B(p.pos), mtl::Arg::B(p.sw), mtl::Arg::B(p.yr), mtl::Arg::V(ca)}, - (T * N + 255) / 256, 1, 256); - gemm(p.mixh, p.Wshg, p.gsh, T, NFF, N); - gemm(p.mixh, p.Wshu, p.ush, T, NFF, N); - const uint32_t nh = (uint32_t) T * NFF; - D(PP(p.lib, "swiglu_h"), {mtl::Arg::B(p.gsh), mtl::Arg::B(p.ush), mtl::Arg::B(p.hh), mtl::Arg::V(nh)}, (nh + 255) / 256, 1, 256); - gemm(p.hh, p.Wshd, p.ys, T, N, NFF); - gemm(p.mixh, p.Wsg, p.sgv, T, 1, N); - const uint32_t oa[2] = {(uint32_t) T, N}; - D(PP(p.lib, "moe_out"), {mtl::Arg::B(p.yr), mtl::Arg::B(p.ys), mtl::Arg::B(p.sgv), mtl::Arg::B(p.bo), mtl::Arg::V(oa)}, - (T * N + 255) / 256, 1, 256); + if (!comb_fold) { + D(PP(p.lib, "mm_combine"), {mtl::Arg::B(p.YE), mtl::Arg::B(p.pos), mtl::Arg::B(p.sw), mtl::Arg::B(p.yr), mtl::Arg::V(ca)}, + (T * N + 255) / 256, 1, 256); + gemm(p.mixh, p.Wshg, p.gsh, T, NFF, N); + nb_ = true; + gemm(p.mixh, p.Wshu, p.ush, T, NFF, N); + nb_ = true; + gemm(p.mixh, p.Wsg, p.sgv, T, 1, N); + const uint32_t nh = (uint32_t) T * NFF; + D(PP(p.lib, "swiglu_h"), {mtl::Arg::B(p.gsh), mtl::Arg::B(p.ush), mtl::Arg::B(p.hh), mtl::Arg::V(nh)}, (nh + 255) / 256, 1, 256); + gemm(p.hh, p.Wshd, p.ys, T, N, NFF); + } + // no moe_out: the next read forms the MoE write from yr (or YE), ys, sgv (gr_norm_h_moe(c), ple_moe_gate_norm(_c), + // gr_apply_moe(c)) if (l == 0 && ple_on_) { - // layer 0's FFN write applied to R, then the PLE block (as the window, for T rows) - const gguf::Tensor& tn = TN("blk.1.hc_attn_norm.weight"); - GrArgs a{N, HC, LR, (uint32_t) T, 1u, eps_}; - D(P(lib_, "gr_norm", {}, "gr_norm"), {mtl::Arg::B(p.R), mtl::Arg::B(p.bo), mtl::Arg::B(p.inj[ci]), mtl::Arg::B(W_, tn.offset), - mtl::Arg::B(p.YE), mtl::Arg::V(a)}, HC, T, 256); + // the projections into layer 0's idle attention scratch (prompt_gemm_init), after layer 0's handshake: the + // gather has been joined + ple_proj(); + // layer 0's FFN write applied to R (with comb_fold its yr formed from YE as it is read), the PLE norms and + // R + gated in one dispatch, then the conv (as the window, for T rows) struct PleArgs { uint32_t n_embd, hc; float eps; } pl{N, HC, strata::kernels::NG_RMS_EPS}; auto off = [&](const char* n) { return TN(n).offset; }; - D(P(lple_, "ple_gnorm", {}, "ple_gnorm"), {mtl::Arg::B(p.pkey), mtl::Arg::B(W_, off("blk.1.ple_norm_key.weight")), - mtl::Arg::B(p.keyn), mtl::Arg::V(pl)}, HC, T, 256); - nb_ = true; - D(P(lple_, "ple_gnorm", {}, "ple_gnorm"), {mtl::Arg::B(p.R), mtl::Arg::B(W_, off("blk.1.ple_norm_query.weight")), - mtl::Arg::B(p.qn), mtl::Arg::V(pl)}, HC, T, 256); - D(P(lple_, "ple_gate_norm", {}, "ple_gate_norm"), - {mtl::Arg::B(p.keyn), mtl::Arg::B(p.qn), mtl::Arg::B(p.pval), mtl::Arg::B(W_, off("blk.1.ple_norm_conv.weight")), - mtl::Arg::B(p.gated), mtl::Arg::B(p.normW), mtl::Arg::V(pl)}, HC, T, 256); - D(P(lple_, "ple_conv_add", {}, "ple_conv_add"), - {mtl::Arg::B(pleHist_), mtl::Arg::B(p.normW), kw_arg(), mtl::Arg::B(p.gated), - mtl::Arg::B(p.R), mtl::Arg::V(pl)}, HC * N / 256, T, 256); + const char* fn = comb_fold ? "ple_moe_gate_norm_c" : "ple_moe_gate_norm"; + std::vector a = {mtl::Arg::B(p.R)}; + if (comb_fold) a.insert(a.end(), {mtl::Arg::B(p.YE), mtl::Arg::B(p.pos), mtl::Arg::B(p.sw)}); + else a.push_back(mtl::Arg::B(p.yr)); + a.insert(a.end(), {mtl::Arg::B(p.ys), mtl::Arg::B(p.sgv), mtl::Arg::B(p.inj[ci]), mtl::Arg::B(p.pkey), + mtl::Arg::B(W_, off("blk.1.ple_norm_key.weight")), + mtl::Arg::B(W_, off("blk.1.ple_norm_query.weight")), mtl::Arg::B(p.pval), + mtl::Arg::B(W_, off("blk.1.ple_norm_conv.weight")), mtl::Arg::B(p.normW), mtl::Arg::V(pl)}); + if (comb_fold) a.push_back(mtl::Arg::V(topk)); + D(P(lple_, fn, {}, fn), std::move(a), comb_fold ? 1 : HC, T, 256); + D(P(lple_, "ple_conv_res", {}, "ple_conv_res"), + {mtl::Arg::B(pleHist_), mtl::Arg::B(p.normW), kw_arg(), mtl::Arg::B(p.R), mtl::Arg::V(pl)}, HC * N / 256, T, 256); + // the PLE history advanced past the chunk once the conv has read it, so normW is dead after layer 0 + const uint32_t tu = (uint32_t) T; + D(P(lple_, "ple_hist_adv", {}, "ple_hist_adv"), {mtl::Arg::B(pleHist_), mtl::Arg::B(p.normW), mtl::Arg::V(pl), mtl::Arg::V(tu)}, + HC * N / 256, 1, 256); } } // the last layer's write applied to R (the draft layer's K/V read the final residual rows) { - const gguf::Tensor& tn = TN("output_hc_norm.weight"); GrArgs a{N, HC, LR, (uint32_t) T, 1u, eps_}; - D(P(lib_, "gr_norm", {}, "gr_norm"), {mtl::Arg::B(p.R), mtl::Arg::B(p.bo), mtl::Arg::B(p.inj[ci]), mtl::Arg::B(W_, tn.offset), - mtl::Arg::B(p.YE), mtl::Arg::V(a)}, HC, T, 256); + if (comb_fold) + D(PP(p.lib, "gr_apply_moec"), {mtl::Arg::B(p.R), mtl::Arg::B(p.YE), mtl::Arg::B(p.pos), mtl::Arg::B(p.sw), + mtl::Arg::B(p.ys), mtl::Arg::B(p.sgv), mtl::Arg::B(p.inj[ci]), mtl::Arg::V(a), + mtl::Arg::V(topk)}, (N + 255) / 256, T, 256); + else + D(PP(p.lib, "gr_apply_moe"), {mtl::Arg::B(p.R), mtl::Arg::B(p.yr), mtl::Arg::B(p.ys), mtl::Arg::B(p.sgv), + mtl::Arg::B(p.inj[ci]), mtl::Arg::V(a)}, HC * N / 256, T, 256); } const double t0 = now_ms(); if (stream_) { seg_T_ = T; cut_ids_ = p.ids.as(); - cx_->run_graph_seg(ds, cuts, [](void* u, int s) { ((Window*) u)->at_cut(s); }, this); + st_->set_prompt_pred(pred_on ? p.ids2.as() : nullptr, TOPK, + pred_on && p.ids3.ptr ? p.ids3.as() : nullptr); + cx_->run_graph_seg(ds, cuts, sigs, [](void* u, int s) { ((Window*) u)->at_cut(s); }, nullptr, this); + st_->set_prompt_pred(nullptr, 0); + if (ple_th.joinable()) ple_th.join(); // only if the graph stopped before layer 0's handshake + ple_join_ = nullptr; + if (!ple_ok) { err = ple_err; return false; } + if (const char* pfx = std::getenv("STRATA_METAL_DUMP_MOE")) { + dump_moe_layer(pfx, 47, false, true); + } cut_ids_ = nullptr; } else { cx_->run_graph(ds); } pg_ms_ += now_ms() - t0; pg_tok_ += T; + if (std::getenv("STRATA_METAL_PKPROF")) kprof(ds, T); static const bool trace = std::getenv("STRATA_METAL_TRACE") != nullptr; - if (trace) std::fprintf(stderr, "strata metal: prompt chunk pos %lld, %d rows: %.0f ms\n", (long long) pos0, T, now_ms() - t0); - // the window's state after the chunk: the conv history turned, the PLE history and token window advanced - hp_ = 1 - hp_; - if (ple_on_) { - using namespace strata::kernels; - float* hist = pleHist_.as(); - const float* nw = p.normW.as(); - for (int c = 0; c < NG_HC_DIM; ++c) { - float col[NG_HIST]; - for (int r = 0; r < NG_HIST; ++r) { - const int src = T - NG_HIST + r; // row r of the new history: chunk row src, or the old history - col[r] = src >= 0 ? nw[(size_t) src * NG_HC_DIM + c] : hist[(size_t) c * NG_HIST + (size_t) (r + T)]; - } - std::memcpy(hist + (size_t) c * NG_HIST, col, sizeof col); - } - } + if (trace) std::fprintf(stderr, "strata metal: prompt chunk pos %lld, %d rows: %.0f ms (embd %.0f ms, ple %.0f ms, graph build %.0f ms)\n", (long long) pos0, T, now_ms() - t0, tbe - tb0, tb1 - tbe, t0 - tb1); + // the window's state after the chunk: the conv history turned (the PLE history on the GPU: ple_hist_adv), the token + // window advanced + hp_ = gdn_ring_ > 0 ? ((hp_ + 1) % 3) : (1 - hp_); if (T >= 2) { ss.ple_prev[0] = tokens[T - 2]; ss.ple_prev[1] = tokens[T - 1]; } else { ss.ple_prev[0] = ss.ple_prev[1]; ss.ple_prev[1] = tokens[0]; } + last_T_ = 1; + last_tokens_[0] = tokens[T - 1]; mtp_src_ = &p.R; // the draft layer's prompt K/V read this chunk's final rows return true; } @@ -2479,8 +4113,16 @@ bool window_init(void* w, const strata::core::ModelGeometry& g, strata::core::Se if (!x->init(g, ss, g_shard1, MAXT, max_cells, err)) return false; window_register(w); if (!mtp_load_pending(x, err)) return false; + if (x->streaming() && !(std::getenv("STRATA_METAL_LAZY_ARENA") && std::atoi(std::getenv("STRATA_METAL_LAZY_ARENA")) != 0)) { + const char* esep = std::getenv("STRATA_METAL_SEPARATE_PROMPT_BUF"); + const char* ept = std::getenv("STRATA_METAL_PHASE_TRADE"); + const bool sep_pg = (ept && std::atoi(ept) == 0) || (esep && std::string(esep) == "1"); + if (!x->prompt_gemm_init(err)) return false; + if (!sep_pg) x->prompt_region(false); + } } x->set_max_t(max_t); + x->prewarm(); *n_vocab = x->vocab(); return true; } @@ -2494,6 +4136,7 @@ bool window_run(void* w, int T, const int32_t* tokens, int64_t pos0, bool head, static const bool logpos = std::getenv("STRATA_LOGPOS") != nullptr; return ((Window*) w)->run(T, tokens, pos0, head || logpos, out, err); } +void window_reset_stats(void* w) { if (w) { if (std::getenv("STRATA_METAL_PKPROF") || std::getenv("STRATA_METAL_TRACE")) ((Window*) w)->stream_report(stderr); ((Window*) w)->reset_stream_stats(); } } void window_reset() { if (g_the_window != nullptr && g_the_window->ready()) g_the_window->reset(); } @@ -2505,6 +4148,18 @@ bool window_checkpoint_restore(const strata::core::ConversationCheckpoint& c, st if (g_the_window == nullptr || !g_the_window->ready()) { err = "metal window: not up"; return false; } return g_the_window->checkpoint_restore(c, err); } +bool window_session_save(const char* path, const int32_t* ids, int64_t n_cells, std::string& err, int64_t base_cells = 0) { + if (g_the_window == nullptr || !g_the_window->ready()) { err = "metal window: not up"; return false; } + return g_the_window->session_save(path, ids, n_cells, err, base_cells); +} +bool window_session_load(const char* path, std::vector& ids, int64_t* n_cells, std::string& err, const char* parent_path = nullptr) { + if (g_the_window == nullptr || !g_the_window->ready()) { err = "metal window: not up"; return false; } + return g_the_window->session_load(path, ids, *n_cells, err, parent_path); +} +bool window_scr_relocate(int64_t c_src, int64_t c_dst, int64_t n, std::string& err) { + if (g_the_window == nullptr || !g_the_window->ready()) { err = "metal window: not up"; return false; } + return g_the_window->scr_relocate(c_src, c_dst, n, err); +} bool window_commit(void* w, int n_keep, std::string& err) { const double t0 = now_ms(); const bool ok = ((Window*) w)->commit(n_keep, err); @@ -2544,6 +4199,9 @@ bool mtp_load_pending(Window* w, std::string& err) { // committed; on_chunk (the MTP's K/V) after each, with the window's final residual rows in place. bool prompt_run_rows(const int64_t* tokens, int64_t n, int64_t pos0, const std::function& on_chunk, std::string& err); +void window_mrope_host_set(const int32_t* host_table) { + g_mrope_host = host_table; +} bool prompt_run(const int64_t* tokens, int64_t n, int64_t pos0, const float* const* embd_rows, const std::function& on_chunk, std::string& err) { struct Rows { @@ -2555,23 +4213,26 @@ bool prompt_run(const int64_t* tokens, int64_t n, int64_t pos0, const float* con bool prompt_run_rows(const int64_t* tokens, int64_t n, int64_t pos0, const std::function& on_chunk, std::string& err) { if (g_window == nullptr) { err = "metal prompt: the decode window is not up yet"; return false; } + g_window->slop_scan(tokens, n, pos0); static const bool windows = std::getenv("STRATA_METAL_PROMPT") && std::string(std::getenv("STRATA_METAL_PROMPT")) == "window"; static const int64_t min_rows = std::getenv("STRATA_METAL_PROMPT_MIN") ? std::atoll(std::getenv("STRATA_METAL_PROMPT_MIN")) : 32; if (!windows && !g_window->streaming() && n >= min_rows) return g_window->prompt(tokens, n, pos0, on_chunk, err); if (!windows && g_window->streaming()) { const char* e = std::getenv("STRATA_METAL_PROMPT_GEMM"); const char* em = std::getenv("STRATA_METAL_PROMPT_GEMM_MIN"); - const int64_t mn = em ? std::atoll(em) : 64; + const int64_t mn = em ? std::atoll(em) : 128; if (!(e && std::string(e) == "0") && n >= mn) { if (!g_window->prompt_gemm_init(err)) return false; const int C = g_window->prompt_chunk_rows(); std::vector tk32((size_t) C); - g_window->prompt_region(true); + g_window->prompt_region(true, n > C); bool ok = true; for (int64_t c = 0; c < n && ok; c += C) { const int k = (int) std::min(C, n - c); for (int i = 0; i < k; ++i) tk32[(size_t) i] = (int32_t) tokens[c + i]; - ok = g_window->prompt_chunk(tk32.data(), k, pos0 + c, err) && (!on_chunk || on_chunk(nullptr, k, pos0 + c, err)); + // the drafter reads the last chunk's rows, and the first's below the attention sink + const int64_t sink = c + k < n ? std::min(k, mtp_sink_cap() - (pos0 + c)) : k; + ok = g_window->prompt_chunk(tk32.data(), k, pos0 + c, err) && (!on_chunk || sink <= 0 || on_chunk(nullptr, sink, pos0 + c, err)); } g_window->prompt_region(false); return ok; diff --git a/metal/src/heal_k256_q2_0.cpp b/metal/src/heal_k256_q2_0.cpp new file mode 100644 index 00000000..123e3c5c --- /dev/null +++ b/metal/src/heal_k256_q2_0.cpp @@ -0,0 +1,834 @@ +// vendor/strata/metal/src/heal_k256_q2_0.cpp - Phase 3 Signed-d AWQ/Lloyd-Max Q2_0 Expert Healer + +// Real-Trace Layerwise Scale Calibration + Monotonic Shared-Expert Ridge Solver for K256_reap_l2_sum +// +// Modes: +// 1. heal_k256_q2_0 --quantize [gain_gu=0.50] [gain_d=0.30] [cal_iq3s_pfx] +// Quantizes all 48 x 256 routed experts to Signed-d AWQ/Lloyd-Max Q2_0 (unlocking the sign bit of fp16 d +// for a 0-bit polarity switch between {-|d|, 0, +|d|, +2|d|} and {-2|d|, -|d|, 0, +|d|}), weighted by +// tempered real activation RMS (E[mixed_d^2])^0.5 (beta=0.50) on W_gate/W_up and analytical +// fourth-root importance (sigma_g * sigma_u)^0.5 ~ (E[h_c^2])^0.25 (beta=0.25) on W_down +// when cal_iq3s_pfx is provided (falling back to hc_ffn_norm RMS if omitted), and saves .experts.d_bak.bin. +// 2. heal_k256_q2_0 --calibrate-and-heal [gamma=0.75] [eta_shd=0.15] [ridge_rel=6.0] +// Calibrates per-layer W_down fp16 d scale alpha_l = cos_l^(1-gamma) / norm_ratio_l directly from the real +// multi-domain teacher trace (all t_cal tokens x all 256 routed experts, eliminating Gaussian depth overshoot) +// and solves the ridge regression (full-rank 640/640 on 46/48 layers; L0=430, L2=617) for blk.0..47.ffn_down_shexp.weight with +// monotonic Q8_0 / IQ4_NL block quantization, writing the 48 modified ffn_down_shexp tensors to . + +#ifndef ACCELERATE_NEW_LAPACK +#define ACCELERATE_NEW_LAPACK +#endif +#include "gguf_lite.h" +#include "ref_quant.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +static constexpr int NL = 48; +static constexpr int NE = 256; +static constexpr int N = 2560; +static constexpr int D = 640; +static constexpr uint64_t ALIGN = 16384; +static constexpr uint64_t Q2_SLICE_BYTES = 460800; // 2560 * 640 / 64 * 18 +static constexpr size_t N_BLOCKS = Q2_SLICE_BYTES / 18; + +static constexpr int8_t KVALUES_IQ4NL[16] = { + -127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113 +}; + +static void require_io(bool ok, const char* msg) { + if (!ok) { + std::perror(msg); + std::exit(2); + } +} + +static void pread_exact(int fd, void* buf, size_t n, off_t off, const char* msg) { + uint8_t* p = static_cast(buf); + size_t got = 0; + while (got < n) { + ssize_t r = ::pread(fd, p + got, n - got, off + (off_t)got); + require_io(r > 0, msg); + got += (size_t)r; + } +} + +static void pwrite_exact(int fd, const void* buf, size_t n, off_t off, const char* msg) { + const uint8_t* p = static_cast(buf); + size_t put = 0; + while (put < n) { + ssize_t w = ::pwrite(fd, p + put, n - put, off + (off_t)put); + require_io(w > 0, msg); + put += (size_t)w; + } +} + +static uint64_t fnv1a_64(const void* data, size_t nbytes) { + const uint8_t* p = static_cast(data); + uint64_t h = 14695981039346656037ULL; + for (size_t i = 0; i < nbytes; ++i) { + h ^= (uint64_t)p[i]; + h *= 1099511628211ULL; + } + return h; +} + +static void quantize_q2_0_signed_awq_fast( + const float* x, size_t n_rows, size_t n_cols, const float* col_imp, uint8_t* out, float gain_exp +) { + const size_t blocks_per_row = n_cols / 64; + for (size_t r = 0; r < n_rows; ++r) { + for (size_t b = 0; b < blocks_per_row; ++b) { + const float* xb = x + r * n_cols + b * 64; + const float* wb = col_imp ? (col_imp + b * 64) : nullptr; + float wxb[64], wi_arr[64]; + float amax = 0.0f, norm_x2 = 0.0f; + for (int i = 0; i < 64; ++i) { + const float xi = xb[i]; + const float wi = wb ? wb[i] : 1.0f; + wi_arr[i] = wi; + wxb[i] = wi * xi; + amax = std::max(amax, std::fabs(xi)); + norm_x2 += wxb[i] * xi; + } + int best_q[64] = {0}; + uint16_t d16 = 0; + if (amax > 0.0f && norm_x2 > 0.0f) { + float best_score = -1.0f, best_qw = 0.0f, best_qq = 1.0f; + static constexpr float seeds[10] = { + 0.42f, 0.52f, 0.64f, 0.78f, 0.94f, + -0.42f, -0.52f, -0.64f, -0.78f, -0.94f + }; + for (int s = 0; s < 10; ++s) { + float d = amax * seeds[s]; + for (int iter = 0; iter < 2; ++iter) { + if (std::fabs(d) <= 1e-12f) break; + const float inv = 1.0f / d; + float qw = 0.0f, qq = 0.0f; + int q_cur[64]; + for (int i = 0; i < 64; ++i) { + int q = (int)std::lrintf(xb[i] * inv); + q = std::max(-1, std::min(2, q)); + q_cur[i] = q; + const float fq = (float)q; + qw += fq * wxb[i]; + qq += wi_arr[i] * (fq * fq); + } + if (qq <= 0.0f) break; + const float d_new = qw / qq; + if ( (d > 0.0f ? d_new <= 0.0f : d_new >= 0.0f) ) break; + const float score = (qw * qw) / qq; + if (score > best_score) { + best_score = score; + best_qw = qw; + best_qq = qq; + std::memcpy(best_q, q_cur, sizeof(best_q)); + } + d = d_new; + } + } + float d_ls = best_qw / best_qq; + float r2 = best_score / std::max(norm_x2, 1e-20f); + float gain = (gain_exp == 0.0f) ? 1.0f : std::pow(std::max(r2, 0.25f), -0.5f * gain_exp); + d16 = ref::f32_to_f16(d_ls * gain); + } + uint8_t* dst = out + (r * blocks_per_row + b) * 18; + std::memcpy(dst, &d16, 2); + for (int j = 0; j < 16; ++j) { + uint8_t byte = 0; + for (int k = 0; k < 4; ++k) { + const uint8_t code = (uint8_t)(best_q[j * 4 + k] + 1); + byte |= (uint8_t)(code << (2 * k)); + } + dst[2 + j] = byte; + } + } + } +} + +static void quantize_iq4nl_monotonic(const uint8_t* raw_old, const float* w_target, size_t n_elems, uint8_t* raw_new) { + const size_t n_blks = n_elems / 32; + std::vector w_old_deq(32); + for (size_t b = 0; b < n_blks; ++b) { + const uint8_t* bo = raw_old + b * 18; + uint8_t* bn = raw_new + b * 18; + std::memcpy(bn, bo, 18); + ref::dequant(ref::IQ4_NL, bo, w_old_deq.data(), 32); + + const float* wt = w_target + b * 32; + float err_old = 0.0f; + for (int i = 0; i < 32; ++i) { + float diff = w_old_deq[i] - wt[i]; + err_old += diff * diff; + } + uint16_t d_old_u; + std::memcpy(&d_old_u, bo, 2); + float d_old = ref::f16_to_f32(d_old_u); + if (std::fabs(d_old) < 1e-12f) continue; + + float best_err = err_old; + uint8_t best_blk[18]; + std::memcpy(best_blk, bo, 18); + + // Search 25 fine scale multipliers around d_old so IQ4_NL preserves small ridge updates + for (int s = -12; s <= 12; ++s) { + float mult = 1.0f + 0.015f * (float)s; + uint16_t d_cand_u = ref::f32_to_f16(d_old * mult); + float d_cand = ref::f16_to_f32(d_cand_u); + uint8_t idx[32]; + float err_c = 0.0f; + for (int i = 0; i < 32; ++i) { + float best_w_err = 1e30f; + uint8_t best_q = 0; + for (uint8_t q = 0; q < 16; ++q) { + float v = d_cand * (float)KVALUES_IQ4NL[q]; + float de = (v - wt[i]) * (v - wt[i]); + if (de < best_w_err) { + best_w_err = de; + best_q = q; + } + } + idx[i] = best_q; + err_c += best_w_err; + } + if (err_c + 1e-11f < best_err) { + best_err = err_c; + std::memcpy(best_blk, &d_cand_u, 2); + for (int j = 0; j < 16; ++j) { + best_blk[2 + j] = (uint8_t)((idx[j] & 0x0F) | ((idx[j + 16] & 0x0F) << 4)); + } + } + } + std::memcpy(bn, best_blk, 18); + } +} + +static void quantize_q8_0_monotonic(const uint8_t* raw_old, const float* w_target, size_t n_elems, uint8_t* raw_new) { + const size_t n_blks = n_elems / 32; + std::vector w_old_deq(32); + for (size_t b = 0; b < n_blks; ++b) { + const uint8_t* bo = raw_old + b * 34; + uint8_t* bn = raw_new + b * 34; + std::memcpy(bn, bo, 34); + ref::dequant(ref::Q8_0, bo, w_old_deq.data(), 32); + + const float* wt = w_target + b * 32; + float err_old = 0.0f, amax = 0.0f; + for (int i = 0; i < 32; ++i) { + float diff = w_old_deq[i] - wt[i]; + err_old += diff * diff; + amax = std::max(amax, std::fabs(wt[i])); + } + uint16_t d_cand_u = ref::f32_to_f16(amax / 127.0f); + float d_cand = ref::f16_to_f32(d_cand_u); + if (d_cand <= 0.0f) continue; + int8_t qs[32]; + float err_c = 0.0f; + for (int i = 0; i < 32; ++i) { + int q = (int)std::lrintf(wt[i] / d_cand); + q = std::max(-127, std::min(127, q)); + qs[i] = (int8_t)q; + float diff = d_cand * (float)q - wt[i]; + err_c += diff * diff; + } + if (err_c + 1e-11f < err_old) { + std::memcpy(bn, &d_cand_u, 2); + std::memcpy(bn + 2, qs, 32); + } + } +} + +static int run_quantize_mode(int argc, char** argv) { + if (argc < 4) { + std::fprintf(stderr, "usage: heal_k256_q2_0 --quantize [gain_gu=0.50] [gain_d=0.30] [cal_iq3s_pfx]\n"); + return 2; + } + const std::string in_gguf = argv[2]; + const std::string out_gguf = argv[3]; + const float gain_gu = (argc > 4) ? std::atof(argv[4]) : 0.50f; + const float gain_d = (argc > 5) ? std::atof(argv[5]) : 0.30f; + const std::string cal_pfx = (argc > 6) ? argv[6] : ""; + const std::string in_exp = in_gguf + ".experts"; + const std::string out_exp = out_gguf + ".experts"; + const std::string bak_d_path = out_exp + ".d_bak.bin"; + + gguf::File f; + std::string err; + if (!f.open(in_gguf, err)) { + std::fprintf(stderr, "gguf open failed: %s\n", err.c_str()); + return 1; + } + + // Check if real teacher activation dump is available for tempered RMS importance + bool use_real_act = false; + size_t t_cal = 0; + int fd_cal_x = -1; + if (!cal_pfx.empty()) { + struct stat st{}; + std::string mix_path = cal_pfx + ".mixed"; + if (::stat(mix_path.c_str(), &st) == 0 && st.st_size > 0 && st.st_size % ((size_t)NL * N * sizeof(float)) == 0) { + t_cal = (size_t)st.st_size / ((size_t)NL * N * sizeof(float)); + fd_cal_x = ::open(mix_path.c_str(), O_RDONLY); + require_io(fd_cal_x >= 0, "open cal_pfx.mixed"); + use_real_act = (t_cal >= 32); + } + } + + std::vector> layer_s_in(NL, std::vector(N, 1.0f)); + std::vector> layer_s2_raw(NL, std::vector(N, 1.0f)); + if (use_real_act) { + std::vector x_l(t_cal * N); + for (int l = 0; l < NL; ++l) { + pread_exact(fd_cal_x, x_l.data(), t_cal * N * sizeof(float), (off_t)((size_t)l * t_cal * N * sizeof(float)), "pread cal_pfx.mixed"); + for (int d = 0; d < N; ++d) { + double ms2 = 0.0; + for (size_t t = 0; t < t_cal; ++t) { + double v = x_l[t * N + d]; + ms2 += v * v; + } + double e_x2 = std::max(ms2 / (double)t_cal, 1e-8); + layer_s2_raw[l][d] = (float)e_x2; + // Tempered real activation RMS (beta = 0.5: sqrt(E[mixed_d^2])) + layer_s_in[l][d] = (float)std::sqrt(e_x2); + } + } + ::close(fd_cal_x); + std::printf("Using tempered real activation RMS (T=%zu, beta=0.50) + per-expert (sigma_g*sigma_u)^0.5 (beta=0.25) for Signed-d Q2_0\n", t_cal); + } else { + for (int l = 0; l < NL; ++l) { + const std::string tname = "blk." + std::to_string(l) + ".hc_ffn_norm.weight"; + const gguf::Tensor* tn = f.find(tname); + if (tn && tn->type == 0) { + const float* wn = reinterpret_cast(static_cast(f.map) + tn->offset); + for (int d = 0; d < N; ++d) { + float ms2 = 0.0f; + for (int c = 0; c < 4; ++c) { + const float v = wn[c * N + d]; + ms2 += v * v; + } + layer_s_in[l][d] = std::sqrt(std::max(ms2 * 0.25f, 1e-12f)); + } + } + } + } + + const size_t hlen = std::min(f.file_bytes, 16 * 1024 * 1024); + std::vector hdr(hlen); + std::memcpy(hdr.data(), f.map, hlen); + for (int l = 0; l < NL; ++l) { + for (const char* role : {"gate", "up", "down"}) { + std::string name = "blk." + std::to_string(l) + ".ffn_" + role + "_exps.weight"; + const uint64_t nlen = name.size(); + std::vector pat(8 + nlen); + std::memcpy(pat.data(), &nlen, 8); + std::memcpy(pat.data() + 8, name.data(), nlen); + auto it = std::search(hdr.begin(), hdr.end(), pat.begin(), pat.end()); + require_io(it != hdr.end(), ("find tensor in GGUF header: " + name).c_str()); + size_t pos = (size_t)(it - hdr.begin()) + pat.size(); + uint32_t nd = 0; + std::memcpy(&nd, hdr.data() + pos, 4); + pos += 4 + (size_t)nd * 8; + uint32_t new_type = 42; + std::memcpy(hdr.data() + pos, &new_type, 4); + } + } + + int ofd = ::open(out_gguf.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0644); + require_io(ofd >= 0, "open out_gguf"); + fcntl(ofd, F_NOCACHE, 1); + pwrite_exact(ofd, hdr.data(), hlen, 0, "pwrite out_gguf hdr"); + size_t pos = hlen; + const size_t chunk = 64ull << 20; + while (pos < f.file_bytes) { + const size_t n = std::min(chunk, f.file_bytes - pos); + pwrite_exact(ofd, (const uint8_t*)f.map + pos, n, (off_t)pos, "pwrite out_gguf body"); + pos += n; + } + ::fsync(ofd); + ::close(ofd); + + int ifd = ::open(in_exp.c_str(), O_RDONLY); + require_io(ifd >= 0, "open in_exp"); + fcntl(ifd, F_NOCACHE, 1); + uint64_t in_tab[NL * 4]; + pread_exact(ifd, in_tab, sizeof(in_tab), 64, "pread in_tab"); + + const uint64_t exp_data0 = ALIGN; + uint64_t out_offs[NL]; + uint64_t at = exp_data0; + const uint64_t out_exp_bytes = 3 * Q2_SLICE_BYTES; + const uint64_t out_layer_bytes = (uint64_t)NE * out_exp_bytes; + for (int l = 0; l < NL; ++l) { + at = (at + ALIGN - 1) & ~(ALIGN - 1); + out_offs[l] = at; + at += out_layer_bytes; + } + + int efd = ::open(out_exp.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0644); + require_io(efd >= 0, "open out_exp"); + fcntl(efd, F_NOCACHE, 1); + + uint8_t strx_hdr[ALIGN]; + std::memset(strx_hdr, 0, sizeof(strx_hdr)); + std::memcpy(strx_hdr, "STRX", 4); + uint32_t v1 = 1, nl_u = NL, ne_u = NE; + uint64_t src_sz = f.file_bytes; + std::memcpy(strx_hdr + 4, &v1, 4); + std::memcpy(strx_hdr + 8, &nl_u, 4); + std::memcpy(strx_hdr + 12, &ne_u, 4); + std::memcpy(strx_hdr + 16, &src_sz, 8); + for (int l = 0; l < NL; ++l) { + uint64_t row[4] = {out_offs[l], Q2_SLICE_BYTES, Q2_SLICE_BYTES, Q2_SLICE_BYTES}; + std::memcpy(strx_hdr + 64 + l * 32, row, 32); + } + pwrite_exact(efd, strx_hdr, ALIGN, 0, "pwrite strx_hdr"); + + std::vector ds_all((size_t)NL * NE * N_BLOCKS); + const size_t n_elems = (size_t)N * D; + auto t0 = std::chrono::steady_clock::now(); + + for (int l = 0; l < NL; ++l) { + const std::string b = "blk." + std::to_string(l) + "."; + int tg = f.find(b + "ffn_gate_exps.weight")->type; + int tu = f.find(b + "ffn_up_exps.weight")->type; + int td = f.find(b + "ffn_down_exps.weight")->type; + + const uint64_t in_off = in_tab[l * 4 + 0]; + const uint64_t gs = in_tab[l * 4 + 1]; + const uint64_t us = in_tab[l * 4 + 2]; + const uint64_t ds = in_tab[l * 4 + 3]; + const uint64_t in_exp_step = gs + us + ds; + const size_t in_layer_bytes = (size_t)NE * in_exp_step; + + std::vector in_buf(in_layer_bytes); + pread_exact(ifd, in_buf.data(), in_layer_bytes, (off_t)in_off, "pread in_layer"); + + std::vector out_buf(out_layer_bytes); + std::atomic next_e{0}; + const int n_threads = 10; + std::vector threads; + for (int tid = 0; tid < n_threads; ++tid) { + threads.emplace_back([&]() { + std::vector f_gate(n_elems), f_up(n_elems), f_down(n_elems), s_down(D, 1.0f); + const float* s2_l = layer_s2_raw[l].data(); + for (;;) { + int e = next_e.fetch_add(1); + if (e >= NE) break; + const uint8_t* src = in_buf.data() + (size_t)e * in_exp_step; + uint8_t* dst = out_buf.data() + (size_t)e * out_exp_bytes; + + if (tg == ref::Q2_0) { + std::memcpy(dst, src, Q2_SLICE_BYTES); + if (use_real_act && td != ref::Q2_0) ref::dequant(tg, src, f_gate.data(), n_elems); + } else { + ref::dequant(tg, src, f_gate.data(), n_elems); + quantize_q2_0_signed_awq_fast(f_gate.data(), D, N, layer_s_in[l].data(), dst, gain_gu); + } + if (tu == ref::Q2_0) { + std::memcpy(dst + Q2_SLICE_BYTES, src + gs, Q2_SLICE_BYTES); + if (use_real_act && td != ref::Q2_0) ref::dequant(tu, src + gs, f_up.data(), n_elems); + } else { + ref::dequant(tu, src + gs, f_up.data(), n_elems); + quantize_q2_0_signed_awq_fast(f_up.data(), D, N, layer_s_in[l].data(), dst + Q2_SLICE_BYTES, gain_gu); + } + uint8_t* qd_dst = dst + 2 * Q2_SLICE_BYTES; + if (td == ref::Q2_0) { + std::memcpy(qd_dst, src + gs + us, Q2_SLICE_BYTES); + } else { + ref::dequant(td, src + gs + us, f_down.data(), n_elems); + const float* d_imp = nullptr; + if (use_real_act) { + for (int c = 0; c < D; ++c) { + const float* rg = f_gate.data() + (size_t)c * N; + const float* ru = f_up.data() + (size_t)c * N; + double vg = 0.0, vu = 0.0; + for (int d = 0; d < N; ++d) { + const double s2d = s2_l[d]; + vg += (double)rg[d] * rg[d] * s2d; + vu += (double)ru[d] * ru[d] * s2d; + } + s_down[c] = (float)std::pow(std::max(vg * vu, 1e-16), 0.25); + } + d_imp = s_down.data(); + } + quantize_q2_0_signed_awq_fast(f_down.data(), N, D, d_imp, qd_dst, gain_d); + } + uint16_t* d_bak = ds_all.data() + ((size_t)l * NE + e) * N_BLOCKS; + for (size_t blk = 0; blk < N_BLOCKS; ++blk) { + std::memcpy(&d_bak[blk], qd_dst + blk * 18, 2); + } + } + }); + } + for (auto& th : threads) th.join(); + pwrite_exact(efd, out_buf.data(), out_buf.size(), (off_t)out_offs[l], "pwrite out_layer"); + if ((l + 1) % 8 == 0 || l == NL - 1) { + double el = std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); + std::printf(" Quantized layer %d/%d to Signed-d AWQ Q2_0 (%.1fs elapsed)\n", l + 1, NL, el); + std::fflush(stdout); + } + } + + FILE* fb = std::fopen(bak_d_path.c_str(), "wb"); + require_io(fb != nullptr, "fopen d_bak.bin"); + require_io(std::fwrite(ds_all.data(), sizeof(uint16_t), ds_all.size(), fb) == ds_all.size(), "fwrite d_bak.bin"); + std::fclose(fb); + + // Save .quant.meta fingerprint over (gain_gu, gain_d, cal_pfx, t_cal) + const std::string qmeta_path = out_exp + ".quant.meta"; + FILE* fqm = std::fopen(qmeta_path.c_str(), "wb"); + require_io(fqm != nullptr, ("fopen " + qmeta_path).c_str()); + uint32_t gu_bits = 0, gd_bits = 0; + std::memcpy(&gu_bits, &gain_gu, 4); + std::memcpy(&gd_bits, &gain_d, 4); + const uint64_t pfx_hash = fnv1a_64(cal_pfx.data(), cal_pfx.size()); + const uint64_t qmeta[3] = {((uint64_t)gu_bits << 32) | gd_bits, pfx_hash, (uint64_t)t_cal}; + require_io(std::fwrite(qmeta, sizeof(uint64_t), 3, fqm) == 3, "fwrite quant.meta"); + std::fclose(fqm); + + ::fsync(efd); + ::close(efd); + ::close(ifd); + return 0; +} + +static int run_calibrate_and_heal_mode(int argc, char** argv) { + if (argc < 7) { + std::fprintf(stderr, "usage: heal_k256_q2_0 --calibrate-and-heal [gamma=0.75] [eta_shd=0.15] [ridge_rel=6.0]\n"); + return 2; + } + const std::string orig_gguf = argv[2]; + const std::string healed_gguf = argv[3]; + const std::string cal_pfx = argv[4]; + const std::string yr_q2_bin = argv[5]; + const std::string out_writers_bin = argv[6]; + const float gamma = (argc > 7) ? std::atof(argv[7]) : 0.75f; + const float eta_shd = (argc > 8) ? std::atof(argv[8]) : 0.15f; + const float ridge_rel = (argc > 9) ? std::atof(argv[9]) : 6.0f; + + const std::string healed_exp = healed_gguf + ".experts"; + const std::string bak_d_path = healed_exp + ".d_bak.bin"; + + gguf::File gf; + std::string err; + if (!gf.open(orig_gguf, err)) { + std::fprintf(stderr, "gguf open failed: %s\n", err.c_str()); + return 1; + } + + int fd_heal = ::open(healed_exp.c_str(), O_RDWR); + require_io(fd_heal >= 0, "open healed_exp"); + uint64_t tab_heal[NL * 4]; + pread_exact(fd_heal, tab_heal, sizeof(tab_heal), 64, "pread tab_heal"); + + FILE* fb = std::fopen(bak_d_path.c_str(), "rb"); + require_io(fb != nullptr, ("fopen " + bak_d_path).c_str()); + std::vector ds_all((size_t)NL * NE * N_BLOCKS); + require_io(std::fread(ds_all.data(), sizeof(uint16_t), ds_all.size(), fb) == ds_all.size(), "fread d_bak.bin"); + std::fclose(fb); + + // Infer t_cal dynamically from (cal_pfx + ".mixed") file size and cross-check .yr, .ids, .wts, .hh sizes + struct stat st_mix{}, st_yr{}, st_ids{}, st_wts{}, st_hh{}; + require_io(::stat((cal_pfx + ".mixed").c_str(), &st_mix) == 0 && st_mix.st_size > 0, "stat cal_pfx.mixed"); + require_io(st_mix.st_size % ((size_t)NL * N * sizeof(float)) == 0, "cal_pfx.mixed size divisible by NL*N*4"); + const size_t t_cal = (size_t)st_mix.st_size / ((size_t)NL * N * sizeof(float)); + require_io(::stat((cal_pfx + ".yr").c_str(), &st_yr) == 0 && (size_t)st_yr.st_size == (size_t)NL * t_cal * N * sizeof(float), "cal_pfx.yr size match"); + require_io(::stat((cal_pfx + ".ids").c_str(), &st_ids) == 0 && (size_t)st_ids.st_size == (size_t)NL * t_cal * 10 * sizeof(int32_t), "cal_pfx.ids size match"); + require_io(::stat((cal_pfx + ".wts").c_str(), &st_wts) == 0 && (size_t)st_wts.st_size == (size_t)NL * t_cal * 10 * sizeof(float), "cal_pfx.wts size match"); + require_io(::stat((cal_pfx + ".hh").c_str(), &st_hh) == 0 && (size_t)st_hh.st_size == (size_t)NL * t_cal * D * sizeof(uint16_t), "cal_pfx.hh size match"); + + int fd_mix = ::open((cal_pfx + ".mixed").c_str(), O_RDONLY); + int fd_yr_ref = ::open((cal_pfx + ".yr").c_str(), O_RDONLY); + require_io(fd_mix >= 0 && fd_yr_ref >= 0, "open cal_pfx .mixed/.yr"); + + // Compute cache fingerprint over d_bak.bin + L0 & L47 of .mixed + L0/L24/L47 gate/up slices of healed_exp + t_cal + std::vector mix_l0_l47(2 * t_cal * N); + pread_exact(fd_mix, mix_l0_l47.data(), t_cal * N * sizeof(float), 0, "pread mix_l0"); + pread_exact(fd_mix, mix_l0_l47.data() + t_cal * N, t_cal * N * sizeof(float), (off_t)((size_t)(NL - 1) * t_cal * N * sizeof(float)), "pread mix_l47"); + std::vector gu_sample(3 * 2 * Q2_SLICE_BYTES); + pread_exact(fd_heal, gu_sample.data() + 0 * 2 * Q2_SLICE_BYTES, 2 * Q2_SLICE_BYTES, (off_t)tab_heal[0 * 4 + 0], "pread gu_l0"); + pread_exact(fd_heal, gu_sample.data() + 1 * 2 * Q2_SLICE_BYTES, 2 * Q2_SLICE_BYTES, (off_t)tab_heal[24 * 4 + 0], "pread gu_l24"); + pread_exact(fd_heal, gu_sample.data() + 2 * 2 * Q2_SLICE_BYTES, 2 * Q2_SLICE_BYTES, (off_t)tab_heal[47 * 4 + 0], "pread gu_l47"); + const uint64_t dbak_hash = fnv1a_64(ds_all.data(), ds_all.size() * sizeof(uint16_t)); + const uint64_t mix_hash = fnv1a_64(mix_l0_l47.data(), mix_l0_l47.size() * sizeof(float)); + const uint64_t gu_hash = fnv1a_64(gu_sample.data(), gu_sample.size()); + const uint64_t want_meta[4] = {dbak_hash, mix_hash, gu_hash, (uint64_t)t_cal}; + const std::string meta_path = yr_q2_bin + ".meta"; + + bool cache_valid = false; + struct stat st_q2{}; + if (::stat(yr_q2_bin.c_str(), &st_q2) == 0 && (size_t)st_q2.st_size == (size_t)NL * t_cal * N * sizeof(float)) { + FILE* fm = std::fopen(meta_path.c_str(), "rb"); + if (fm) { + uint64_t got_meta[4] = {0, 0, 0, 0}; + if (std::fread(got_meta, sizeof(uint64_t), 4, fm) == 4 && + std::memcmp(got_meta, want_meta, sizeof(want_meta)) == 0) { + cache_valid = true; + } + std::fclose(fm); + } + } + + if (!cache_valid) { + int fd_ids = ::open((cal_pfx + ".ids").c_str(), O_RDONLY); + int fd_wts = ::open((cal_pfx + ".wts").c_str(), O_RDONLY); + int fd_q2w = ::open(yr_q2_bin.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0644); + require_io(fd_ids >= 0 && fd_wts >= 0 && fd_q2w >= 0, "open cal_pfx .ids/.wts or yr_q2_bin"); + + std::atomic nl_q2{0}; + std::vector q2_ths; + for (int tid = 0; tid < 8; ++tid) { + q2_ths.emplace_back([&]() { + const size_t n_elems = (size_t)N * D; + std::vector x(t_cal * N), yr_q2(t_cal * N, 0.0f); + std::vector ids(t_cal * 10); + std::vector wts(t_cal * 10); + std::vector in_h((size_t)NE * 3 * Q2_SLICE_BYTES); + std::vector wg(n_elems), wu(n_elems), wd(n_elems); + for (;;) { + int l = nl_q2.fetch_add(1); + if (l >= NL) break; + pread_exact(fd_mix, x.data(), t_cal * N * sizeof(float), (off_t)((size_t)l * t_cal * N * sizeof(float)), "pread mix"); + pread_exact(fd_ids, ids.data(), t_cal * 10 * sizeof(int32_t), (off_t)((size_t)l * t_cal * 10 * sizeof(int32_t)), "pread ids"); + pread_exact(fd_wts, wts.data(), t_cal * 10 * sizeof(float), (off_t)((size_t)l * t_cal * 10 * sizeof(float)), "pread wts"); + pread_exact(fd_heal, in_h.data(), in_h.size(), (off_t)tab_heal[l * 4 + 0], "pread in_h"); + for (int e = 0; e < NE; ++e) { + uint8_t* qd_s = in_h.data() + ((size_t)e * 3 + 2) * Q2_SLICE_BYTES; + const uint16_t* d_b = ds_all.data() + ((size_t)l * NE + e) * N_BLOCKS; + for (size_t blk = 0; blk < N_BLOCKS; ++blk) std::memcpy(qd_s + blk * 18, &d_b[blk], 2); + } + std::fill(yr_q2.begin(), yr_q2.end(), 0.0f); + std::vector>> exp_toks(NE); + for (size_t t = 0; t < t_cal; ++t) { + for (int k = 0; k < 10; ++k) { + int e = ids[t * 10 + k]; + float w = wts[t * 10 + k]; + if (e >= 0 && e < NE && w > 0.0f) exp_toks[e].push_back({(int)t, w}); + } + } + std::vector x_e, g_e, u_e, y_e; + for (int e = 0; e < NE; ++e) { + const auto& hits = exp_toks[e]; + if (hits.empty()) continue; + int m = (int)hits.size(); + x_e.resize((size_t)m * N); g_e.resize((size_t)m * D); u_e.resize((size_t)m * D); y_e.resize((size_t)m * N); + for (int i = 0; i < m; ++i) std::memcpy(x_e.data() + (size_t)i * N, x.data() + (size_t)hits[i].first * N, N * sizeof(float)); + const uint8_t* rh = in_h.data() + (size_t)e * 3 * Q2_SLICE_BYTES; + ref::dequant(ref::Q2_0, rh, wg.data(), n_elems); + ref::dequant(ref::Q2_0, rh + Q2_SLICE_BYTES, wu.data(), n_elems); + ref::dequant(ref::Q2_0, rh + 2 * Q2_SLICE_BYTES, wd.data(), n_elems); + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, m, D, N, 1.0f, x_e.data(), N, wg.data(), N, 0.0f, g_e.data(), D); + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, m, D, N, 1.0f, x_e.data(), N, wu.data(), N, 0.0f, u_e.data(), D); + for (int i = 0; i < m * D; ++i) { float gv = g_e[i]; g_e[i] = (gv / (1.0f + std::exp(-gv))) * u_e[i]; } + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, m, N, D, 1.0f, g_e.data(), D, wd.data(), D, 0.0f, y_e.data(), N); + for (int i = 0; i < m; ++i) { + float* dst = yr_q2.data() + (size_t)hits[i].first * N; + const float* src = y_e.data() + (size_t)i * N; + float w = hits[i].second; + for (int c = 0; c < N; ++c) dst[c] += w * src[c]; + } + } + pwrite_exact(fd_q2w, yr_q2.data(), t_cal * N * sizeof(float), (off_t)((size_t)l * t_cal * N * sizeof(float)), "pwrite yr_q2"); + } + }); + } + for (auto& th : q2_ths) th.join(); + ::close(fd_q2w); ::close(fd_wts); ::close(fd_ids); + + FILE* fm = std::fopen(meta_path.c_str(), "wb"); + require_io(fm != nullptr, "fopen yr_q2_bin.meta"); + require_io(std::fwrite(want_meta, sizeof(uint64_t), 4, fm) == 4, "fwrite yr_q2_bin.meta"); + std::fclose(fm); + } + + int fd_yr_q2 = ::open(yr_q2_bin.c_str(), O_RDONLY); + require_io(fd_yr_q2 >= 0, "open yr_q2_bin"); + + // Part A (Real-Trace Layerwise Scale Calibration) + Part B (Full-Rank Shared-Expert Ridge Regression) + FILE* fout = std::fopen(out_writers_bin.c_str(), "wb"); + require_io(fout != nullptr, "fopen out_writers_bin"); + uint32_t n_tensors = (eta_shd > 0.0f) ? NL : 0; + require_io(std::fwrite(&n_tensors, sizeof(uint32_t), 1, fout) == 1, "fwrite n_tensors"); + + auto write_record = [&](const std::string& name, uint32_t gtype, const void* data, uint32_t nbytes) { + uint32_t nlen = (uint32_t)name.size(); + require_io(std::fwrite(&nlen, 4, 1, fout) == 1, "fwrite nlen"); + require_io(std::fwrite(name.data(), 1, nlen, fout) == nlen, "fwrite name"); + require_io(std::fwrite(>ype, 4, 1, fout) == 1, "fwrite gtype"); + require_io(std::fwrite(&nbytes, 4, 1, fout) == 1, "fwrite nbytes"); + require_io(std::fwrite(data, 1, nbytes, fout) == nbytes, "fwrite data"); + }; + + std::vector layer_alpha(NL, 1.0f); + std::vector slice(Q2_SLICE_BYTES); + double tot_ref2 = 0.0, tot_dy2_unscaled = 0.0, tot_dy2_before = 0.0, tot_dy2_after = 0.0; + + for (int l = 0; l < NL; ++l) { + std::vector x(t_cal * N), yr_r(t_cal * N), yr_q(t_cal * N), dy(t_cal * N); + pread_exact(fd_mix, x.data(), t_cal * N * sizeof(float), (off_t)((size_t)l * t_cal * N * sizeof(float)), "pread mix"); + pread_exact(fd_yr_ref, yr_r.data(), t_cal * N * sizeof(float), (off_t)((size_t)l * t_cal * N * sizeof(float)), "pread yr_ref"); + pread_exact(fd_yr_q2, yr_q.data(), t_cal * N * sizeof(float), (off_t)((size_t)l * t_cal * N * sizeof(float)), "pread yr_q2"); + + double sum_ref2 = 0.0, sum_q2 = 0.0, sum_dot = 0.0; + for (size_t i = 0; i < t_cal * N; ++i) { + double r = yr_r[i], q = yr_q[i]; + sum_ref2 += r * r; + sum_q2 += q * q; + sum_dot += r * q; + double du = r - q; + tot_dy2_unscaled += du * du; + } + tot_ref2 += sum_ref2; + + float alpha_l = 1.0f; + if (gamma >= 0.0f && sum_ref2 > 0.0 && sum_q2 > 0.0) { + double norm_ratio = std::sqrt(sum_q2 / sum_ref2); + double cos_sim = std::max(sum_dot / (std::sqrt(sum_ref2) * std::sqrt(sum_q2)), 0.1); + alpha_l = (float)(std::pow(cos_sim, 1.0 - (double)gamma) / norm_ratio); + } + layer_alpha[l] = alpha_l; + + // Apply alpha_l to W_down fp16 d headers in healed_exp + uint64_t off_h = tab_heal[l * 4 + 0]; + for (int e = 0; e < NE; ++e) { + const uint16_t* d_bak = ds_all.data() + ((size_t)l * NE + e) * N_BLOCKS; + uint64_t d_off = off_h + (uint64_t)e * (3 * Q2_SLICE_BYTES) + 2 * Q2_SLICE_BYTES; + pread_exact(fd_heal, slice.data(), Q2_SLICE_BYTES, (off_t)d_off, "pread slice"); + for (size_t b = 0; b < N_BLOCKS; ++b) { + float d = ref::f16_to_f32(d_bak[b]); + uint16_t d_new = ref::f32_to_f16(d * alpha_l); + std::memcpy(slice.data() + b * 18, &d_new, 2); + } + pwrite_exact(fd_heal, slice.data(), Q2_SLICE_BYTES, (off_t)d_off, "pwrite slice"); + } + + for (size_t i = 0; i < t_cal * N; ++i) { + dy[i] = yr_r[i] - alpha_l * yr_q[i]; + tot_dy2_before += (double)dy[i] * dy[i]; + } + + if (eta_shd <= 0.0f) { + tot_dy2_after = tot_dy2_before; + continue; + } + + const std::string b = "blk." + std::to_string(l) + "."; + const auto* tn_shg = gf.find(b + "ffn_gate_shexp.weight"); + const auto* tn_shu = gf.find(b + "ffn_up_shexp.weight"); + const auto* tn_shd = gf.find(b + "ffn_down_shexp.weight"); + const auto* tn_sg = gf.find(b + "ffn_gate_inp_shexp.weight"); + + std::vector w_shg(D * N), w_shu(D * N), w_shd(N * D), w_sg(N); + ref::dequant(tn_shg->type, (const uint8_t*)gf.map + tn_shg->offset, w_shg.data(), D * N); + ref::dequant(tn_shu->type, (const uint8_t*)gf.map + tn_shu->offset, w_shu.data(), D * N); + ref::dequant(tn_shd->type, (const uint8_t*)gf.map + tn_shd->offset, w_shd.data(), N * D); + std::memcpy(w_sg.data(), (const uint8_t*)gf.map + tn_sg->offset, N * sizeof(float)); + + std::vector g_sh(t_cal * D), u_sh(t_cal * D), z_sh(t_cal * D), sg_l(t_cal); + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, (int)t_cal, D, N, 1.0f, x.data(), N, w_shg.data(), N, 0.0f, g_sh.data(), D); + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, (int)t_cal, D, N, 1.0f, x.data(), N, w_shu.data(), N, 0.0f, u_sh.data(), D); + cblas_sgemv(CblasRowMajor, CblasNoTrans, (int)t_cal, N, 1.0f, x.data(), N, w_sg.data(), 1, 0.0f, sg_l.data(), 1); + + for (size_t t = 0; t < t_cal; ++t) { + float sig = 1.0f / (1.0f + std::exp(-sg_l[t])); + for (int c = 0; c < D; ++c) { + float gv = g_sh[t * D + c]; + z_sh[t * D + c] = sig * (gv / (1.0f + std::exp(-gv))) * u_sh[t * D + c]; + } + } + + std::vector G(D * D, 0.0f), B(D * N, 0.0f); + cblas_sgemm(CblasRowMajor, CblasTrans, CblasNoTrans, D, D, (int)t_cal, 1.0f, z_sh.data(), D, z_sh.data(), D, 0.0f, G.data(), D); + cblas_sgemm(CblasRowMajor, CblasTrans, CblasNoTrans, D, N, (int)t_cal, 1.0f, z_sh.data(), D, dy.data(), N, 0.0f, B.data(), N); + + double trace_G = 0.0; + for (int c = 0; c < D; ++c) trace_G += G[c * D + c]; + float lam = (float)(ridge_rel * (trace_G / D)); + for (int c = 0; c < D; ++c) G[c * D + c] += lam; + + std::vector Gd(D * D), L_chol(D * D, 0.0); + for (int i = 0; i < D * D; ++i) Gd[i] = G[i]; + for (int i = 0; i < D; ++i) { + for (int j = 0; j <= i; ++j) { + double s = Gd[i * D + j]; + for (int k = 0; k < j; ++k) s -= L_chol[i * D + k] * L_chol[j * D + k]; + if (i == j) L_chol[i * D + j] = std::sqrt(std::max(s, 1e-12)); + else L_chol[i * D + j] = s / L_chol[j * D + j]; + } + } + std::vector w_shd_target = w_shd; + for (int r = 0; r < N; ++r) { + double y_fwd[D], x_bwd[D]; + for (int i = 0; i < D; ++i) { + double s = B[i * N + r]; + for (int k = 0; k < i; ++k) s -= L_chol[i * D + k] * y_fwd[k]; + y_fwd[i] = s / L_chol[i * D + i]; + } + for (int i = D - 1; i >= 0; --i) { + double s = y_fwd[i]; + for (int k = i + 1; k < D; ++k) s -= L_chol[k * D + i] * x_bwd[k]; + x_bwd[i] = s / L_chol[i * D + i]; + } + for (int c = 0; c < D; ++c) { + w_shd_target[r * D + c] += eta_shd * (float)x_bwd[c]; + } + } + + const uint8_t* raw_shd_old = (const uint8_t*)gf.map + tn_shd->offset; + std::vector raw_shd_new(tn_shd->bytes); + if (tn_shd->type == ref::Q8_0) { + quantize_q8_0_monotonic(raw_shd_old, w_shd_target.data(), N * D, raw_shd_new.data()); + } else { + quantize_iq4nl_monotonic(raw_shd_old, w_shd_target.data(), N * D, raw_shd_new.data()); + } + + std::vector w_shd_new_deq(N * D), delta_w(N * D), pred_dy(t_cal * N); + ref::dequant(tn_shd->type, raw_shd_new.data(), w_shd_new_deq.data(), N * D); + for (int i = 0; i < N * D; ++i) delta_w[i] = w_shd_new_deq[i] - w_shd[i]; + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, (int)t_cal, N, D, 1.0f, z_sh.data(), D, delta_w.data(), D, 0.0f, pred_dy.data(), N); + for (size_t i = 0; i < t_cal * N; ++i) { + double rem = (double)dy[i] - (double)pred_dy[i]; + tot_dy2_after += rem * rem; + } + + write_record(b + "ffn_down_shexp.weight", (uint32_t)tn_shd->type, raw_shd_new.data(), (uint32_t)raw_shd_new.size()); + } + + std::fclose(fout); + ::fsync(fd_heal); + ::close(fd_heal); + ::close(fd_mix); + ::close(fd_yr_ref); + ::close(fd_yr_q2); + + std::printf("Calibrated 48 layers on real trace (T=%zu, gamma=%.2f, L0=%.3f L24=%.3f L47=%.3f) & solved %u ffn_down_shexp tensors (MoE MSE %.2f -> %.2f -> %.2f, %.2f%%)\n", + t_cal, gamma, layer_alpha[0], layer_alpha[24], layer_alpha[47], n_tensors, + tot_dy2_unscaled, tot_dy2_before, tot_dy2_after, (tot_dy2_after / tot_dy2_unscaled - 1.0) * 100.0); + return 0; +} + +int main(int argc, char** argv) { + if (argc >= 2 && std::string(argv[1]) == "--quantize") { + return run_quantize_mode(argc, argv); + } + if (argc >= 2 && std::string(argv[1]) == "--calibrate-and-heal") { + return run_calibrate_and_heal_mode(argc, argv); + } + if (argc >= 3) { + std::vector args = {argv[0], (char*)"--quantize"}; + for (int i = 1; i < argc; ++i) args.push_back(argv[i]); + return run_quantize_mode((int)args.size(), args.data()); + } + std::fprintf(stderr, "usage: heal_k256_q2_0 [--quantize | --calibrate-and-heal] ...\n"); + return 2; +} diff --git a/metal/src/mtl_rt.h b/metal/src/mtl_rt.h index e5d72592..9e12c80a 100644 --- a/metal/src/mtl_rt.h +++ b/metal/src/mtl_rt.h @@ -88,7 +88,7 @@ public: Pipeline pipeline(int lib, const std::string& fn, const Constants& c = {}); Buffer alloc(size_t bytes); // shared storage, zeroed - Buffer wrap(void* page_aligned, size_t bytes); // no-copy wrap (page-aligned, size multiple of page) + Buffer wrap(void* page_aligned, size_t bytes, bool count_allocated = true); // no-copy wrap (page-aligned, size multiple of page) void free(Buffer& b); // Encodes `ds` into ONE compute encoder (serial or concurrent dispatch type) and ONE command buffer, `reps` @@ -103,6 +103,11 @@ public: // missing experts are put in place). Returns the GPU time in seconds. double run_graph_seg(const std::vector& ds, const std::vector& cuts, void (*at_cut)(void*, int), void* user); + // sigs (same length as cuts, sigs[s] <= cuts[s]): the GPU signals at sigs[s] and waits at cuts[s]; the dispatches + // between run while the host takes its turn (they must not read what the host writes). after_cut(s) runs on the + // host once the GPU has been let go at cut s (work the GPU does not wait for). + double run_graph_seg(const std::vector& ds, const std::vector& cuts, const std::vector& sigs, + void (*at_cut)(void*, int), void (*after_cut)(void*, int), void* user); // One command buffer per `run_cb` call, empty or not: measures commit+wait round trips. double run_wall(const std::vector& ds, int reps, double* gpu_seconds = nullptr); diff --git a/metal/src/mtl_rt.mm b/metal/src/mtl_rt.mm index 4c6045b2..106657d5 100644 --- a/metal/src/mtl_rt.mm +++ b/metal/src/mtl_rt.mm @@ -183,13 +183,13 @@ Buffer Context::alloc(size_t bytes) { return r; } -Buffer Context::wrap(void* p, size_t bytes) { +Buffer Context::wrap(void* p, size_t bytes, bool count_allocated) { id b = [impl_->dev newBufferWithBytesNoCopy:p length:bytes options:MTLResourceStorageModeShared deallocator:nil]; if (!b) { std::fprintf(stderr, "mtl: wrap %zu failed\n", bytes); std::exit(4); } - allocated_ += bytes; + if (count_allocated) allocated_ += bytes; Buffer r; r.handle = (__bridge_retained void*)b; r.ptr = p; @@ -267,14 +267,19 @@ double Context::run_graph(const std::vector& ds) { double Context::run_graph_seg(const std::vector& ds, const std::vector& cuts, void (*at_cut)(void*, int), void* user) { + return run_graph_seg(ds, cuts, cuts, at_cut, nullptr, user); +} + +double Context::run_graph_seg(const std::vector& ds, const std::vector& cuts, const std::vector& sigs, + void (*at_cut)(void*, int), void (*after_cut)(void*, int), void* user) { @autoreleasepool { if (impl_->ev == nil) impl_->ev = [impl_->dev newSharedEvent]; id ev = impl_->ev; const uint64_t base = impl_->ev_v; id cb = [impl_->q commandBufferWithUnretainedReferences]; size_t i = 0; - for (size_t s = 0; s <= cuts.size(); ++s) { - const size_t end = s < cuts.size() ? (size_t) cuts[s] : ds.size(); + auto run_to = [&](size_t end) { // [i, end) as one concurrent encoder + if (i >= end) return; id enc = [cb computeCommandEncoderWithDispatchType:MTLDispatchTypeConcurrent]; const size_t first = i; for (; i < end; ++i) { @@ -282,11 +287,14 @@ double Context::run_graph_seg(const std::vector& ds, const std::vector encode(enc, ds[i]); } [enc endEncoding]; - if (s < cuts.size()) { - [cb encodeSignalEvent:ev value:base + 2 * s + 1]; - [cb encodeWaitForEvent:ev value:base + 2 * s + 2]; - } + }; + for (size_t s = 0; s < cuts.size(); ++s) { + run_to((size_t) sigs[s]); + [cb encodeSignalEvent:ev value:base + 2 * s + 1]; + run_to((size_t) cuts[s]); // beside the host's turn + [cb encodeWaitForEvent:ev value:base + 2 * s + 2]; } + run_to(ds.size()); [cb commit]; for (size_t s = 0; s < cuts.size(); ++s) { const uint64_t want = base + 2 * s + 1; @@ -295,6 +303,7 @@ double Context::run_graph_seg(const std::vector& ds, const std::vector } at_cut(user, (int) s); ev.signaledValue = base + 2 * s + 2; + if (after_cut) after_cut(user, (int) s); } impl_->ev_v = base + 2 * cuts.size() + 2; [cb waitUntilCompleted]; diff --git a/metal/src/strata_window_cli.cpp b/metal/src/strata_window_cli.cpp new file mode 100644 index 00000000..dd57e839 --- /dev/null +++ b/metal/src/strata_window_cli.cpp @@ -0,0 +1,1947 @@ +// metal/src/strata_window_cli.cpp - Direct zero-bloat Metal runner for Strata's Window + ExpertStream + PleTable. +// Supports both: +// 1. Prefill (PromptGemm) + decode (Window::run + ExpertStream F_NOCACHE pread + speculative router lookahead) +// 2. Teacher-forced full-vocab (248,320) FP64 KLD, PPL, and top-1 agreement evaluation (--logits-out / --logits-ref) + +#include "strata/core/conversation_cache.hpp" +#include "strata/core/layer.hpp" +#include "strata/core/session.hpp" +#include "strata/kernels/ngram.hpp" +#include "strata/spec/suffix_drafter.hpp" +#include "strata/spec/spec_gate.hpp" +#include "engine/scr_state.hpp" +#include "engine/suffix_window.hpp" +#include "strata/core/scr.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace strata::metal_engine { +void set_model_shard(const char* shard1); +void* window_create(); +void window_destroy(void* w); +bool window_init(void* w, const strata::core::ModelGeometry& g, strata::core::SessionState& ss, int max_t, int64_t max_cells, + int64_t* n_vocab, std::string& err); +void window_mrope_host_set(const int32_t* host_table); +bool prompt_run(const int64_t* tokens, int64_t n, int64_t pos0, const float* const* embd_rows, + const std::function& on_chunk, std::string& err); +bool prompt_run_rows(const int64_t* tokens, int64_t n, int64_t pos0, + const std::function& on_chunk, std::string& err); +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); +void window_reset_stats(void* w); +void window_reset(); +bool window_checkpoint_save(strata::core::ConversationCheckpoint& c, std::string& err); +bool window_checkpoint_restore(const strata::core::ConversationCheckpoint& c, std::string& err); +bool window_session_save(const char* path, const int32_t* ids, int64_t n_cells, std::string& err, int64_t base_cells = 0); +bool window_session_load(const char* path, std::vector& ids, int64_t* n_cells, std::string& err, const char* parent_path = nullptr); +bool window_scr_relocate(int64_t c_src, int64_t c_dst, int64_t n, std::string& err); +bool mtp_load(const std::string& rt_dir, int max_t, std::string& err); +bool mtp_prefill(const int32_t* next_tokens, int64_t n, int64_t cell0, std::string& err); +bool mtp_draft(int T, const int32_t* tokens, int64_t p, int a, int max_steps, int32_t* drafts, float* probs, float min_p, + int* n_drafts, std::string& err); +} // namespace strata::metal_engine + +static double now_ms() { + return std::chrono::duration(std::chrono::steady_clock::now().time_since_epoch()).count(); +} + +// One policy instance per decode: a new prompt does not inherit another workload's timings. +struct SpecControl { + strata::spec::SpecGate gate; + bool enabled = true, full = false; + int mtp_used = 0; + + SpecControl() { + const char* mode = std::getenv("STRATA_SPEC_GATE"); + enabled = mode && std::strcmp(mode, "0") != 0; + full = mode && std::strcmp(mode, "1") == 0; + } + + int propose(strata::spec::SuffixDrafter& suffix, int32_t* input, + const int32_t* mtp, const float* probs, int n_mtp, int suffix_max, + int rem, int64_t pos, int64_t cells) { + int nd = strata::spec::suffix_limit(n_mtp, rem, pos, cells); + if (full) nd = gate.spec_gate_pick(0, nd, probs); + mtp_used = nd; + for (int j = 0; j < nd; ++j) input[1 + j] = mtp[j]; + int limit = strata::spec::suffix_limit(std::max(suffix_max, nd), rem, pos, cells) - nd; + if (suffix_max <= 0 || limit <= 0) return nd; + if (enabled) limit = gate.spec_gate_pick(1, limit, nullptr, nd, probs); + if (limit <= 0) return nd; + const double begin = now_ms(); + const int sk = nd > 0 ? suffix.propose_after(input + 1, nd, limit, input + 1 + nd) + : suffix.propose(limit, input + 1); + if (sk > 0) gate.spec_gate_draft_cost(1, (now_ms() - begin) / (1000 * sk)); + return nd + sk; + } + + void result(int nd, int keep) { + const int accepted = keep - 1; + gate.spec_gate_result(0, std::min(mtp_used, keep), std::min(mtp_used, accepted)); + if (accepted >= mtp_used) + gate.spec_gate_result(1, std::min(nd - mtp_used, keep - mtp_used), accepted - mtp_used); + } + + void describe() const { + if (enabled) std::fprintf(stderr, "%s\n", gate.spec_gate_describe().c_str()); + } +}; + +static double logsumexp_f64(const float* lg, int64_t n, int32_t* argmax_out) { + float mx = lg[0]; + int32_t am = 0; + for (int64_t i = 1; i < n; ++i) { + if (lg[i] > mx) { + mx = lg[i]; + am = (int32_t) i; + } + } + if (argmax_out) *argmax_out = am; + double sum = 0.0; + for (int64_t i = 0; i < n; ++i) { + sum += std::exp((double) lg[i] - (double) mx); + } + return (double) mx + std::log(sum); +} + +static std::vector read_ids(const std::string& path) { + std::vector v; + std::ifstream f(path); + int64_t x; + while (f >> x) v.push_back((int32_t) x); + return v; +} + +// teacher-forced next-token distributions, as log-probabilities: inputs[i] runs at pos + i and predicts targets[i] +struct TfRun { + std::vector lp; // [n][n_vocab] + double nll = 0; + double ms = 0; +}; +static bool tf_run(void* w, const std::vector& inputs, const std::vector& targets, int64_t pos, int64_t n_vocab, + TfRun& out, std::string& err) { + namespace me = strata::metal_engine; + const size_t n = inputs.size(); + out.lp.assign(n * (size_t) n_vocab, 0.0f); + out.nll = 0; + const double t0 = now_ms(); + for (size_t i = 0; i < n; ++i) { + int32_t tok = inputs[i], next = 0; + if (!me::window_run(w, 1, &tok, pos + (int64_t) i, true, &next, err) || !me::window_commit(w, 1, err)) return false; + float* lg = out.lp.data() + i * (size_t) n_vocab; + if (!me::window_copy_logits(w, 0, lg)) { err = "copy_logits failed"; return false; } + const double lse = logsumexp_f64(lg, n_vocab, nullptr); + for (int64_t v = 0; v < n_vocab; ++v) lg[v] = (float) ((double) lg[v] - lse); + out.nll -= lg[targets[i]]; + } + out.ms = now_ms() - t0; + return true; +} +// KL(ref || cand) per position, nats +static void tf_compare(const TfRun& ref, const TfRun& c, int64_t n_vocab, double& mean, double& p90, double& mx, double& top1) { + const size_t n = ref.lp.size() / (size_t) n_vocab; + std::vector k(n); + int agree = 0; + for (size_t i = 0; i < n; ++i) { + const float* a = ref.lp.data() + i * (size_t) n_vocab; + const float* b = c.lp.data() + i * (size_t) n_vocab; + double s = 0; + int64_t ia = 0, ib = 0; + for (int64_t v = 0; v < n_vocab; ++v) { + s += std::exp((double) a[v]) * ((double) a[v] - (double) b[v]); + if (a[v] > a[ia]) ia = v; + if (b[v] > b[ib]) ib = v; + } + k[i] = std::max(0.0, s); + agree += ia == ib; + } + mean = 0; + for (double x : k) mean += x; + mean /= (double) n; + std::vector s = k; + std::sort(s.begin(), s.end()); + p90 = s[(size_t) (0.9 * (double) (n - 1))]; + mx = s.back(); + top1 = 100.0 * agree / (double) n; +} + +// --scr-test: the SCR splice [A B C] -> [A B' C] against a fresh prefill of [A B' C], scored teacher-forced on D. +// C's last token is D's first input (so every variant's state is "A, the middle span, C minus its last token"). +// Variants, each from the saved [A B C] session: +// orig no edit: the old context (how much B -> B' matters at all) +// keep J C's cells relocated, B' + C[:J] prefilled from the state at |A|, then the running state of [A B C] +// restored (the CLM rule: the post-edit GDN keeps having read B) +// gdnex J as keep, but the running state of the fresh [A B' C] restored: only the KV staleness is left +// lossy J as keep, without the restore: the GDN state has read A B' C[:J] only +// J is raised to the junction block's minimum; lossy with J = |C| is the fresh prefill again (a sanity row: KL 0). +static int scr_test(void* w, strata::core::SessionState& ss, int64_t n_vocab, const std::string& pa, const std::string& pb, + const std::string& pbp, const std::string& pc, const std::string& pd, const std::string& js, int32_t pad, + int n_eval, const std::string& tmp) { + namespace me = strata::metal_engine; + std::string err; + const std::vector A = read_ids(pa), B = read_ids(pb), D = read_ids(pd); + std::vector Bp = read_ids(pbp), C = read_ids(pc); + if (A.size() < 4 || B.empty() || Bp.empty() || C.size() < 2 || D.size() < 2) { std::fprintf(stderr, "scr-test: empty span\n"); return 2; } + while ((Bp.size() - B.size()) % 4 != 0) Bp.push_back(pad); // the shift must be a multiple of 4 cells + const int32_t c_last = C.back(); + C.pop_back(); + const int64_t nA = (int64_t) A.size(), nB = (int64_t) B.size(), nBp = (int64_t) Bp.size(), nC = (int64_t) C.size(); + std::vector inputs{c_last}, targets; + for (size_t i = 0; i + 1 < D.size() && (int) inputs.size() < n_eval; ++i) inputs.push_back(D[i]); + targets.assign(D.begin(), D.begin() + (ptrdiff_t) inputs.size()); + auto prefill = [&](const std::vector& ids, int64_t pos) { + std::vector v(ids.begin(), ids.end()); + return v.empty() || me::prompt_run_rows(v.data(), (int64_t) v.size(), pos, nullptr, err); + }; + auto cat = [](std::vector a, const std::vector& b, size_t nb) { a.insert(a.end(), b.begin(), b.begin() + (ptrdiff_t) nb); return a; }; + auto fresh = [&] { me::window_reset(); ss.ple_prev[0] = -1; ss.ple_prev[1] = -1; }; + const int32_t prev_c[2] = {C.size() >= 2 ? C[C.size() - 2] : Bp.back(), C.back()}; + auto fail = [&](const char* what) { std::fprintf(stderr, "scr-test: %s: %s\n", what, err.c_str()); return 1; }; + + // the reference: [A B' C] prefilled fresh + strata::core::ConversationCheckpoint ckA, ckABC, ckRef; + int32_t prevA[2]; + fresh(); + if (!prefill(A, 0) || !me::window_checkpoint_save(ckA, err)) return fail("prefill A"); + prevA[0] = ss.ple_prev[0]; prevA[1] = ss.ple_prev[1]; + double t0 = now_ms(); + if (!prefill(cat(Bp, C, C.size()), nA)) return fail("prefill B' C"); + const double ref_ms = now_ms() - t0; + if (!me::window_checkpoint_save(ckRef, err)) return fail("checkpoint"); + TfRun ref; + if (!tf_run(w, inputs, targets, nA + nBp + nC, n_vocab, ref, err)) return fail("ref eval"); + // the old context [A B C], saved as a session + fresh(); + if (!prefill(cat(A, B, B.size()), 0) || !prefill(C, nA + nB) || !me::window_checkpoint_save(ckABC, err)) return fail("prefill A B C"); + const std::vector abc = cat(cat(A, B, B.size()), C, C.size()); + t0 = now_ms(); + if (!me::window_session_save(tmp.c_str(), abc.data(), (int64_t) abc.size(), err)) return fail("session save"); + const double save_ms = now_ms() - t0; + TfRun orig; + if (!tf_run(w, inputs, targets, nA + nB + nC, n_vocab, orig, err)) return fail("orig eval"); + const double n = (double) inputs.size(); + std::printf("SCR_SPANS A=%lld B=%lld B'=%lld (padded) C=%lld eval=%zu delta=%lld session_save_ms=%.1f ref_prefill_ms(B'+C)=%.1f ref_ppl=%.4f\n", + (long long) nA, (long long) nB, (long long) nBp, (long long) nC, inputs.size(), (long long) (nBp - nB), save_ms, ref_ms, + std::exp(ref.nll / n)); + double km, kp, kx, t1; + tf_compare(ref, orig, n_vocab, km, kp, kx, t1); + std::printf("SCR_ROW variant=orig J=- mean_kld=%.6f p90_kld=%.6f max_kld=%.6f top1=%.2f ppl=%.4f\n", km, kp, kx, t1, std::exp(orig.nll / n)); + + std::vector Js; + { + std::stringstream s(js); + std::string x; + while (std::getline(s, x, ',')) Js.push_back(std::atoll(x.c_str())); + Js.push_back(nC); + } + for (const char* mode : {"keep", "gdnex", "lossy"}) + for (int64_t J : Js) { + if (J == nC && std::string(mode) != "lossy") continue; + const int64_t Je = std::min(nC, std::max(J, scr::scr_min_junction(nA + nBp))); + std::vector ids; + int64_t cells = 0; + t0 = now_ms(); + if (!me::window_session_load(tmp.c_str(), ids, &cells, err)) return fail("session load"); + const double load_ms = now_ms() - t0; + t0 = now_ms(); + if (!me::window_scr_relocate(nA + nB, nA + nBp, nC, err)) return fail("relocate"); + const double rel_ms = now_ms() - t0; + t0 = now_ms(); + if (!me::window_checkpoint_restore(ckA, err)) return fail("restore A"); + ss.ple_prev[0] = prevA[0]; ss.ple_prev[1] = prevA[1]; + if (!prefill(cat(Bp, C, (size_t) Je), nA)) return fail("prefill B' C[:J]"); + if (std::string(mode) == "keep" && !me::window_checkpoint_restore(ckABC, err)) return fail("restore ABC"); + if (std::string(mode) == "gdnex" && !me::window_checkpoint_restore(ckRef, err)) return fail("restore ref"); + ss.ple_prev[0] = prev_c[0]; ss.ple_prev[1] = prev_c[1]; // token ids are never stale + const double splice_ms = now_ms() - t0; + TfRun r; + if (!tf_run(w, inputs, targets, nA + nBp + nC, n_vocab, r, err)) return fail("eval"); + tf_compare(ref, r, n_vocab, km, kp, kx, t1); + std::printf("SCR_ROW variant=%s J=%lld mean_kld=%.6f p90_kld=%.6f max_kld=%.6f top1=%.2f ppl=%.4f load_ms=%.1f relocate_ms=%.1f " + "splice_prefill_ms=%.1f speedup_vs_ref=%.2f\n", mode, (long long) Je, km, kp, kx, t1, std::exp(r.nll / n), load_ms, + rel_ms, splice_ms, ref_ms / std::max(1e-3, rel_ms + splice_ms)); + std::fflush(stdout); + if (std::string(mode) == "keep" && (J == 0 || J == 8)) { + // Connectome loop check (PLAN.md §2.9): restore spliced state and free-generate 128 tokens greedy + if (!me::window_session_load(tmp.c_str(), ids, &cells, err)) return fail("loop session load"); + if (!me::window_scr_relocate(nA + nB, nA + nBp, nC, err)) return fail("loop relocate"); + if (!me::window_checkpoint_restore(ckA, err)) return fail("loop restore A"); + ss.ple_prev[0] = prevA[0]; ss.ple_prev[1] = prevA[1]; + if (!prefill(cat(Bp, C, (size_t) Je), nA)) return fail("loop prefill"); + if (!me::window_checkpoint_restore(ckABC, err)) return fail("loop restore ABC"); + ss.ple_prev[0] = prev_c[0]; ss.ple_prev[1] = prev_c[1]; + std::vector gen_toks; + int32_t cur = inputs[0], nxt = 0; + int64_t gpos = nA + nBp + nC; + for (int step = 0; step < 128; ++step) { + if (!me::window_run(w, 1, &cur, gpos, true, &nxt, err) || !me::window_commit(w, 1, err)) break; + gen_toks.push_back(nxt); + cur = nxt; + ++gpos; + } + int rep4 = 0, tot4 = std::max(0, (int) gen_toks.size() - 3); + for (int i = 0; i < tot4; ++i) { + for (int k = 0; k < i; ++k) { + if (gen_toks[i] == gen_toks[k] && gen_toks[i+1] == gen_toks[k+1] && + gen_toks[i+2] == gen_toks[k+2] && gen_toks[i+3] == gen_toks[k+3]) { + ++rep4; + break; + } + } + } + std::printf("SCR_LOOP_CHECK variant=keep J=%lld gen_tokens=%zu repeated_4grams=%d/%d (%.2f%%)\n", + (long long) Je, gen_toks.size(), rep4, tot4, tot4 > 0 ? 100.0 * rep4 / tot4 : 0.0); + std::fflush(stdout); + } + } + std::remove(tmp.c_str()); + return 0; +} + + +#include +#include + +// sample_logits' distribution: the top-k logits, the repetition / presence penalties over the last 128 tokens of +// `recent`, the temperature, min-p and top-p. tok / p: its tokens and unnormalized probabilities (sum: theirs); empty +// for greedy decoding. +struct Dist { + std::vector tok; + std::vector p; + double sum = 0.0; +}; +static void sample_dist(const float* lg, int64_t n_vocab, float temp, int top_k, float top_p, float min_p, float pen_rep, + float pen_pres, const std::vector& recent, Dist& d) { + d.tok.clear(); + d.p.clear(); + d.sum = 0.0; + if (temp <= 0.0f || top_k == 1) return; + const int k = top_k > 0 ? std::min((int) n_vocab, top_k) : 64; + std::vector> cand; + cand.reserve(256); + float cutoff = -INFINITY; + for (int64_t i = 0; i < n_vocab; ++i) { + float v = lg[i]; + if (std::isnan(v) || v <= cutoff) continue; + cand.push_back({v, (int32_t) i}); + if ((int) cand.size() >= k * 4) { + std::nth_element(cand.begin(), cand.begin() + k - 1, cand.end(), + [](const auto& a, const auto& b) { return a.first > b.first; }); + cand.resize((size_t) k); + cutoff = cand.back().first; + } + } + if (cand.empty()) return; + if (pen_rep != 1.0f || pen_pres != 0.0f) { + const size_t r0 = recent.size() > 128 ? recent.size() - 128 : 0; + for (auto& kv : cand) { + bool seen = false; + for (size_t ri = r0; ri < recent.size(); ++ri) if (recent[ri] == kv.second) { seen = true; break; } + if (seen) { + if (pen_rep > 1.0f) kv.first = kv.first > 0.0f ? kv.first / pen_rep : kv.first * pen_rep; + kv.first -= pen_pres; + } + } + } + std::partial_sort(cand.begin(), cand.begin() + std::min(cand.size(), (size_t) k), cand.end(), + [](const auto& a, const auto& b) { return a.first > b.first; }); + if ((int) cand.size() > k) cand.resize((size_t) k); + const float mx = cand[0].first; + double sum = 0.0; + std::vector probs(cand.size()); + for (size_t i = 0; i < cand.size(); ++i) { + probs[i] = std::exp((double) (cand[i].first - mx) / (double) temp); + sum += probs[i]; + } + if (sum <= 0.0) { + d.tok = {cand[0].second}; + d.p = {1.0}; + d.sum = 1.0; + return; + } + if (min_p > 0.0f && cand.size() > 1) { + const double p_best = probs[0] / sum; + size_t keep = 1; + while (keep < cand.size() && (probs[keep] / sum) >= p_best * (double) min_p) ++keep; + cand.resize(keep); + probs.resize(keep); + sum = 0.0; + for (double p : probs) sum += p; + } + if (top_p > 0.0f && top_p < 1.0f) { + double cum = 0.0; + size_t keep = cand.size(); + for (size_t i = 0; i < cand.size(); ++i) { + cum += probs[i] / sum; + if (cum >= (double) top_p) { + keep = i + 1; + break; + } + } + cand.resize(keep); + probs.resize(keep); + sum = 0.0; + for (double p : probs) sum += p; + } + d.tok.resize(cand.size()); + for (size_t i = 0; i < cand.size(); ++i) d.tok[i] = cand[i].second; + d.p = std::move(probs); + d.sum = sum; +} + +// a draw from d, token `skip` left out (-1: none) +static int32_t dist_draw(const Dist& d, std::mt19937& rng, int32_t skip = -1) { + double tot = d.sum; + for (size_t i = 0; i < d.tok.size(); ++i) if (d.tok[i] == skip) tot -= d.p[i]; + if (tot <= 0.0) return skip; + std::uniform_real_distribution dist(0.0, tot); + double r = dist(rng); + int32_t last = d.tok[0]; + for (size_t i = 0; i < d.tok.size(); ++i) { + if (d.tok[i] == skip) continue; + last = d.tok[i]; + r -= d.p[i]; + if (r <= 0.0) return d.tok[i]; + } + return last; +} + +static int32_t sample_logits(const float* lg, int64_t n_vocab, int32_t greedy_tok, float temp, int top_k, float top_p, float min_p, + float pen_rep, float pen_pres, const std::vector& recent, std::mt19937& rng) { + Dist d; + sample_dist(lg, n_vocab, temp, top_k, top_p, min_p, pen_rep, pen_pres, recent, d); + return d.tok.empty() ? greedy_tok : dist_draw(d, rng); +} + +// Speculative sampling (Leviathan et al. 2023, sec. 2.3) of a verified window: input[0] the last token, input[1..nd] the +// drafts, output[i] the window's greedy pick after input[0..i]. The drafts are deterministic (q = one-hot), so draft +// x after row i is accepted with probability p_i(x) (p_i: sample_dist of row i's logits, history `hist` + input[0..i]); +// a rejected row samples p_i without x (the residual max(p - q, 0) normalized), and a fully accepted window samples its +// last row. The tokens emitted are distributed as sample_logits one at a time. Returns keep (output[0..keep-1] the +// tokens, input[0..keep-1] appended to hist). +static int spec_sample(void* w, const int32_t* input, int32_t* output, int nd, int64_t n_vocab, float temp, int top_k, float top_p, + float min_p, float pen_rep, float pen_pres, std::vector& hist, std::vector& lg, + std::mt19937& rng) { + namespace me = strata::metal_engine; + Dist d; + std::uniform_real_distribution uni(0.0, 1.0); + for (int i = 0;; ++i) { + hist.push_back(input[i]); + me::window_copy_logits(w, i, lg.data()); + sample_dist(lg.data(), n_vocab, temp, top_k, top_p, min_p, pen_rep, pen_pres, hist, d); + if (d.tok.empty()) { // greedy (top_k 1): the drafts must be the picks + if (i == nd || output[i] != input[i + 1] || strata::spec::suffix_stop(output[i])) return i + 1; + continue; + } + if (i == nd) { + output[i] = dist_draw(d, rng); + return i + 1; + } + const int32_t x = input[i + 1]; + static const bool typical = std::getenv("STRATA_SPEC_SAMPLE") && std::string(std::getenv("STRATA_SPEC_SAMPLE")) == "typical"; + double px = 0.0; + for (size_t k = 0; k < d.tok.size(); ++k) if (d.tok[k] == x) px = d.p[k]; + const bool accept = typical ? (!d.tok.empty() && (d.tok[0] == x || px >= std::max(0.05 * d.sum, 0.25 * d.p[0]))) + : (uni(rng) * d.sum < px); + if (accept) { + output[i] = x; + if (strata::spec::suffix_stop(x)) return i + 1; + continue; + } + output[i] = typical ? dist_draw(d, rng) : dist_draw(d, rng, x); + return i + 1; + } +} + +struct DiskCache { + struct Entry { + std::string path; + std::string parent_path; + std::vector ids; + std::vector pads; + int64_t base_cells = 0; + size_t bytes = 0; + bool pinned = false; + }; + + static constexpr size_t kLearnMinTokens = 32; + static constexpr size_t kLearnMaxTokens = 65536; + static constexpr size_t kMaxPins = 16; + static constexpr size_t kSeenMax = 32; + + std::string dir; + size_t limit_bytes = 0; + std::vector entries; + std::deque> seen_prompts; + + bool enabled() const { return !dir.empty(); } + + static uint64_t fnv1a(const int64_t* data, size_t n) { + uint64_t h = 1469598103934665603ull; + const auto* p = reinterpret_cast(data); + const size_t nb = n * sizeof(int64_t); + for (size_t i = 0; i < nb; ++i) h = (h ^ p[i]) * 1099511628211ull; + return h; + } + + static bool is_prefix(const std::vector& a, const std::vector& b) { + return a.size() <= b.size() && std::equal(a.begin(), a.end(), b.begin()); + } + + static size_t common_prefix(const std::vector& a, const std::vector& b, size_t cap) { + const size_t lim = std::min({a.size(), b.size(), cap}); + size_t i = 0; + while (i < lim && a[i] == b[i]) ++i; + return i; + } + + static void read_pads(const std::string& ppath, size_t n_ids, int32_t pad_tok, const std::vector& ids, + std::vector& out_pads) { + out_pads.clear(); + FILE* pf = std::fopen(ppath.c_str(), "rb"); + if (!pf) return; + uint64_t np = 0; + if (std::fread(&np, sizeof(np), 1, pf) == 1 && np <= n_ids) { + std::vector p((size_t) np); + if (std::fread(p.data(), sizeof(int64_t), (size_t) np, pf) == (size_t) np) { + bool ok = true; + for (int64_t at : p) { + if (at < 0 || (size_t) at >= n_ids || ids[(size_t) at] != pad_tok) ok = false; + } + if (ok) out_pads = std::move(p); + } + } + std::fclose(pf); + } + + static void write_pads(const std::string& ppath, const std::vector& pads) { + if (pads.empty()) { + std::remove(ppath.c_str()); + return; + } + if (FILE* pf = std::fopen(ppath.c_str(), "wb")) { + const uint64_t np = (uint64_t) pads.size(); + std::fwrite(&np, sizeof(np), 1, pf); + std::fwrite(pads.data(), sizeof(int64_t), pads.size(), pf); + std::fclose(pf); + } + } + + static std::string read_parent(const std::string& ppath) { + std::ifstream in(ppath, std::ios::binary); + if (!in) return ""; + std::string s; + std::getline(in, s); + while (!s.empty() && (s.back() == '\r' || s.back() == '\n')) s.pop_back(); + return s; + } + + static void write_parent(const std::string& ppath, const std::string& parent_path) { + if (parent_path.empty()) { + std::remove(ppath.c_str()); + return; + } + std::ofstream out(ppath, std::ios::binary | std::ios::trunc); + if (out) out << parent_path << "\n"; + } + + static void remove_files(const std::string& p) { + namespace fs = std::filesystem; + std::error_code ec; + fs::remove(p, ec); + fs::remove(p + ".pads", ec); + fs::remove(p + ".pin", ec); + fs::remove(p + ".parent", ec); + } + + void load_seen() { + seen_prompts.clear(); + if (dir.empty()) return; + const std::string spath = dir + "/prompts-seen.bin"; + FILE* f = std::fopen(spath.c_str(), "rb"); + if (!f) return; + uint32_t magic = 0, count = 0; + if (std::fread(&magic, 4, 1, f) == 1 && magic == 0x5345454Eu && + std::fread(&count, 4, 1, f) == 1 && count <= kSeenMax) { + for (uint32_t i = 0; i < count; ++i) { + uint32_t len = 0; + if (std::fread(&len, 4, 1, f) != 1 || len > kLearnMaxTokens) break; + std::vector i32(len); + if (len > 0 && std::fread(i32.data(), 4, len, f) != len) break; + seen_prompts.emplace_back(i32.begin(), i32.end()); + } + } + std::fclose(f); + } + + void save_seen() const { + if (dir.empty()) return; + const std::string spath = dir + "/prompts-seen.bin"; + const std::string tmp = spath + ".tmp"; + FILE* f = std::fopen(tmp.c_str(), "wb"); + if (!f) return; + const uint32_t magic = 0x5345454Eu; + const uint32_t count = (uint32_t) seen_prompts.size(); + bool ok = std::fwrite(&magic, 4, 1, f) == 1 && std::fwrite(&count, 4, 1, f) == 1; + for (const auto& p : seen_prompts) { + const uint32_t len = (uint32_t) std::min(p.size(), kLearnMaxTokens); + std::vector i32(p.begin(), p.begin() + len); + ok = ok && std::fwrite(&len, 4, 1, f) == 1 && + (len == 0 || std::fwrite(i32.data(), 4, len, f) == len); + } + std::fclose(f); + if (ok) std::rename(tmp.c_str(), spath.c_str()); + else std::remove(tmp.c_str()); + } + + void init(const std::string& cache_dir, int64_t limit_mib, int64_t max_cells, int32_t pad_tok) { + namespace fs = std::filesystem; + if (cache_dir.empty()) return; + std::error_code ec; + fs::create_directories(cache_dir, ec); + if (ec) { + std::fprintf(stderr, "strata metal cache: cannot create %s: %s\n", cache_dir.c_str(), ec.message().c_str()); + return; + } + dir = cache_dir; + limit_bytes = limit_mib > 0 ? (size_t) limit_mib * 1024ull * 1024ull : 0; + std::vector> found; + for (const auto& de : fs::directory_iterator(dir, ec)) { + const auto p = de.path(); + const auto ext = p.extension().string(); + if (ext == ".partial" || ext == ".tmp") { + fs::remove(p, ec); + continue; + } + if (ext != ".kvc") continue; + FILE* f = std::fopen(p.string().c_str(), "rb"); + if (!f) continue; + scr::SessionHeader h, want; + const bool hdr_ok = std::fread(&h, sizeof(h), 1, f) == 1; + const int64_t n_ids = hdr_ok ? (int64_t) (uint32_t) h.n_ids : 0; + const int64_t base_cells = (hdr_ok && (h.flags & (1u << 23))) ? (int64_t) ((uint64_t) h.n_ids >> 32) : 0; + bool ok = hdr_ok && + std::memcmp(h.magic, want.magic, 8) == 0 && + h.version == want.version && + h.n_cells > 0 && h.n_cells < max_cells && + n_ids == h.n_cells && + base_cells >= 0 && base_cells < h.n_cells && (base_cells & 3) == 0; + std::vector i32; + if (ok) { + i32.resize((size_t) n_ids); + ok = std::fread(i32.data(), 4, (size_t) n_ids, f) == (size_t) n_ids; + } + std::fclose(f); + std::string parent_path = base_cells > 0 ? read_parent(p.string() + ".parent") : ""; + if (base_cells > 0 && parent_path.empty()) ok = false; + if (!ok) { + remove_files(p.string()); + continue; + } + Entry e; + e.path = p.string(); + e.parent_path = std::move(parent_path); + e.ids.assign(i32.begin(), i32.end()); + read_pads(e.path + ".pads", e.ids.size(), pad_tok, e.ids, e.pads); + e.base_cells = base_cells; + e.bytes = (size_t) de.file_size(ec); + e.pinned = base_cells == 0 && fs::exists(e.path + ".pin", ec); + found.push_back({de.last_write_time(ec), std::move(e)}); + } + std::sort(found.begin(), found.end(), [](const auto& a, const auto& b) { return a.first < b.first; }); + for (auto& kv : found) entries.push_back(std::move(kv.second)); + prune_orphans(); + enforce_limits(); + load_seen(); + size_t total = 0, n_pin = 0, n_shared = 0; + for (const auto& e : entries) { + total += e.bytes; + if (e.pinned) ++n_pin; + if (e.base_cells > 0) ++n_shared; + } + std::fprintf(stderr, "strata metal cache: disk tier %s, %zu sessions (%zu pinned, %zu shared), %.2f GiB (limit %.1f GiB)\n", + dir.c_str(), entries.size(), n_pin, n_shared, total / 1073741824.0, limit_bytes / 1073741824.0); + } + + void prune_orphans() { + for (auto it = entries.begin(); it != entries.end();) { + if (it->base_cells > 0) { + bool parent_ok = false; + for (const auto& p : entries) { + if (p.path == it->parent_path && p.base_cells == 0 && + (int64_t) common_prefix(p.ids, it->ids, p.ids.size()) >= it->base_cells) { + parent_ok = true; + break; + } + } + if (!parent_ok) { + remove_files(it->path); + it = entries.erase(it); + continue; + } + } + ++it; + } + } + + void enforce_limits() { + namespace fs = std::filesystem; + std::error_code ec; + size_t n_pin = 0; + for (const auto& e : entries) if (e.pinned) ++n_pin; + for (size_t i = 0; n_pin > kMaxPins && i < entries.size(); ++i) { + if (entries[i].pinned) { + entries[i].pinned = false; + fs::remove(entries[i].path + ".pin", ec); + --n_pin; + } + } + if (limit_bytes == 0) return; + size_t total = 0; + for (const auto& e : entries) total += e.bytes; + for (auto it = entries.begin(); total > limit_bytes && entries.size() > 1 && it != entries.end();) { + if (it->pinned) { + ++it; + continue; + } + const std::string gone = it->path; + total -= it->bytes; + remove_files(gone); + it = entries.erase(it); + for (auto cit = entries.begin(); cit != entries.end();) { + if (cit->parent_path == gone) { + if (cit < it) --it; + total -= cit->bytes; + remove_files(cit->path); + cit = entries.erase(cit); + } else { + ++cit; + } + } + } + } + + int64_t observe_prefix(const std::vector& ids, int64_t n_pre) { + if (!enabled() || n_pre < (int64_t) kLearnMinTokens) return -1; + const size_t cap = std::min((size_t) n_pre, kLearnMaxTokens); + std::vector cur(ids.begin(), ids.begin() + cap); + size_t best_lcp = 0; + bool added = false; + for (auto& s : seen_prompts) { + if (is_prefix(s, cur)) { + s = cur; + added = true; + break; + } + if (is_prefix(cur, s)) { + added = true; + break; + } + const size_t lcp = common_prefix(cur, s, kLearnMaxTokens); + const auto spans = strata::core::scr_detail::find_spans( + s, cur, (int64_t) lcp, strata::core::ScrParams{32, 64, 16, 4}); + if (!spans.empty()) { + // Same CLM conversation after a mid-prompt edit/compaction: update in place, do not pin the fold point + s = cur; + added = true; + break; + } + if (lcp >= kLearnMinTokens && lcp < cur.size() && lcp < s.size()) { + best_lcp = std::max(best_lcp, lcp); + } + } + if (!added) { + if (seen_prompts.size() >= kSeenMax) seen_prompts.pop_front(); + seen_prompts.push_back(std::move(cur)); + } + save_seen(); + if (best_lcp >= kLearnMinTokens) { + for (const auto& e : entries) { + if (e.ids.size() == best_lcp && is_prefix(e.ids, ids)) return -1; + } + return (int64_t) best_lcp; + } + return -1; + } + + void touch(size_t idx) { + if (idx >= entries.size()) return; + namespace fs = std::filesystem; + std::error_code ec; + fs::last_write_time(entries[idx].path, fs::file_time_type::clock::now(), ec); + if (idx + 1 < entries.size()) { + Entry e = std::move(entries[idx]); + entries.erase(entries.begin() + (long) idx); + entries.push_back(std::move(e)); + } + } + + void save_current(const std::vector& ids, const std::vector& pads, bool pinned, std::string& err, + std::string* active_path = nullptr) { + if (!enabled() || ids.size() < kLearnMinTokens) return; + namespace fs = std::filesystem; + std::error_code ec; + for (size_t i = 0; i < entries.size(); ++i) { + if (is_prefix(ids, entries[i].ids) && entries[i].pads == pads) { + if (!pinned || (entries[i].ids.size() == ids.size() && entries[i].base_cells == 0)) { + if (pinned && !entries[i].pinned) { + entries[i].pinned = true; + if (FILE* pf = std::fopen((entries[i].path + ".pin").c_str(), "wb")) std::fclose(pf); + } + if (!pinned && active_path != nullptr) *active_path = entries[i].path; + touch(i); + return; + } + } + } + std::string parent_path; + int64_t base_cells = 0; + if (!pinned) { + size_t best_len = 0; + for (const auto& e : entries) { + if (e.pinned && e.base_cells == 0 && e.pads.empty() && + e.ids.size() >= kLearnMinTokens && ids.size() >= kLearnMinTokens + 4) { + const size_t lcp = common_prefix(e.ids, ids, std::min(e.ids.size(), ids.size() - 4)); + const int64_t c0 = ((int64_t) lcp / 4) * 4; + if (c0 >= (int64_t) kLearnMinTokens && (size_t) c0 > best_len && + (pads.empty() || pads.front() >= c0)) { + best_len = (size_t) c0; + base_cells = c0; + parent_path = e.path; + } + } + } + } + const uint64_t h = fnv1a(ids.data(), ids.size()); + char fname[96]; + std::snprintf(fname, sizeof(fname), "/%016llx-%zu.kvc", (unsigned long long) h, ids.size()); + const std::string path = dir + fname; + std::vector i32(ids.begin(), ids.end()); + if (!strata::metal_engine::window_session_save(path.c_str(), i32.data(), (int64_t) i32.size(), err, base_cells)) return; + write_pads(path + ".pads", pads); + write_parent(path + ".parent", parent_path); + if (pinned) { + if (FILE* pf = std::fopen((path + ".pin").c_str(), "wb")) std::fclose(pf); + } else { + fs::remove(path + ".pin", ec); + } + const size_t sz = (size_t) fs::file_size(path, ec); + const std::string prev_active = (!pinned && active_path != nullptr) ? *active_path : ""; + for (auto it = entries.begin(); it != entries.end();) { + if (it->path == path || + (!it->pinned && (is_prefix(it->ids, ids) || (!prev_active.empty() && it->path == prev_active)))) { + if (it->path != path) { + remove_files(it->path); + } + it = entries.erase(it); + } else { + ++it; + } + } + if (!pinned && active_path != nullptr) *active_path = path; + entries.push_back({path, parent_path, ids, pads, base_cells, sz, pinned}); + prune_orphans(); + enforce_limits(); + } + + int find_best(const std::vector& raw_ids, int64_t max_new, int64_t max_cells, int32_t pad_tok, + int64_t& out_reused) const { + out_reused = 0; + if (!enabled()) return -1; + int best_idx = -1; + int64_t best_score = 0; + for (size_t i = 0; i < entries.size(); ++i) { + const auto& e = entries[i]; + if (e.ids.empty()) continue; + std::vector q_ids = raw_ids; + if (!e.pads.empty()) { + std::vector carried; + strata::core::scr_carry(raw_ids, e.ids, e.pads, pad_tok, carried); + if ((int64_t) carried.size() + max_new <= max_cells) q_ids = std::move(carried); + } + const int64_t q_n_pre = std::max(0, (int64_t) q_ids.size() - 1); + int64_t cand = 0; + if ((int64_t) e.ids.size() <= q_n_pre && is_prefix(e.ids, q_ids)) { + cand = (int64_t) e.ids.size(); + } else if (e.ids.size() >= 32 && q_n_pre >= 32) { + const strata::core::ScrPlan plan = strata::core::scr_plan( + e.ids, q_ids, {(int64_t) e.ids.size()}, strata::core::ScrParams{32, 64, 16, 4}); + if (!plan.empty() && (int64_t) q_ids.size() + plan.pads() + max_new <= max_cells) { + cand = plan.n_prefix + plan.reused(); + } + } + if (cand >= 32 && cand >= best_score) { + best_score = cand; + best_idx = (int) i; + } + } + out_reused = best_score; + return best_idx; + } +}; + +static int serve_loop(void* w, strata::core::SessionState& ss, int64_t n_vocab, int64_t max_cells, int suffix_draft, + bool use_mtp, int spec, float spec_min_p, const std::string& cache_disk_dir, int64_t cache_disk_mib) { + namespace me = strata::metal_engine; + constexpr int64_t kScrPadToken = 198; // '\n' pad for 4-cell QSA block alignment (ported from 01a11408 / scr.hpp) + DiskCache disk_cache; + disk_cache.init(cache_disk_dir, cache_disk_mib, max_cells, kScrPadToken); + + std::printf("INFO kv=%lld\n", (long long) max_cells); + std::printf("READY %lld\n", (long long) max_cells); + std::fflush(stdout); + + std::string line, err; + struct CkSlot { + std::vector ids; + strata::core::ConversationCheckpoint ck; + int32_t prev0 = -1, prev1 = -1; + }; + std::vector ck_ring; + auto save_ck_at = [&](const std::vector& pfx) { + if (pfx.empty()) return; + for (size_t ci = 0; ci < ck_ring.size(); ++ci) { + if (ck_ring[ci].ids == pfx) { + CkSlot slot = std::move(ck_ring[ci]); + ck_ring.erase(ck_ring.begin() + (long) ci); + if (me::window_checkpoint_save(slot.ck, err)) { + slot.prev0 = ss.ple_prev[0]; + slot.prev1 = ss.ple_prev[1]; + ck_ring.push_back(std::move(slot)); + } + return; + } + } + if (ck_ring.size() >= 2) ck_ring.erase(ck_ring.begin()); + CkSlot slot; + slot.ids = pfx; + if (me::window_checkpoint_save(slot.ck, err)) { + slot.prev0 = ss.ple_prev[0]; + slot.prev1 = ss.ple_prev[1]; + ck_ring.push_back(std::move(slot)); + } + }; + std::vector cached_ids; + std::vector cached_pads; + std::vector> cached_imgs; + std::vector mrope_host((size_t) (max_cells + 16) * 3, 0); + std::vector img_rows; + std::vector row_ptr; + std::string active_path; + bool ram_dirty = false; + std::vector lg_buf((size_t) n_vocab); + std::mt19937 rng(42); + // temp > 0 with drafts: speculative sampling (spec_sample), else STRATA_SPEC_SAMPLE=greedy: the drafts kept while + // they are the greedy picks and only the last token sampled (fast, but those tokens are greedy, unpenalized) + static const bool spec_exact = std::getenv("STRATA_SPEC_SAMPLE") && + (std::string(std::getenv("STRATA_SPEC_SAMPLE")) == "exact" || + std::string(std::getenv("STRATA_SPEC_SAMPLE")) == "typical"); + + while (std::getline(std::cin, line)) { + if (line.empty()) continue; + if (line == "QUIT" || line.rfind("QUIT ", 0) == 0) { + if (ram_dirty && cached_ids.size() >= 32) disk_cache.save_current(cached_ids, cached_pads, false, err, &active_path); + break; + } + if (line == "STOP") continue; + if (line == "VRAM" || line.rfind("VRAM ", 0) == 0) { + std::printf("VRAM expert_slots=4129 expert_cache_mib=6656 expert_cache_full_mib=6656 vram_free_mib=13312\n"); + std::fflush(stdout); + continue; + } + if (line.rfind("SAVE ", 0) == 0) { + std::string path = line.substr(5); + while (!path.empty() && (path.back() == '\r' || path.back() == ' ')) path.pop_back(); + const double t0 = now_ms(); + std::vector i32(cached_ids.begin(), cached_ids.end()); + if (!me::window_session_save(path.c_str(), i32.data(), (int64_t) i32.size(), err)) { + std::printf("SERR io 0 %s\n", err.c_str()); + } else { + DiskCache::write_pads(path + ".pads", cached_pads); + struct stat st{}; + const size_t fbytes = ::stat(path.c_str(), &st) == 0 ? (size_t) st.st_size + : scr::fixed_bytes() + scr::cell_bytes(i32.size()); + std::printf("SAVED %zu %zu %.1f\n", i32.size(), fbytes, now_ms() - t0); + } + std::fflush(stdout); + continue; + } + if (line.rfind("RESTORE ", 0) == 0) { + std::string path = line.substr(8); + while (!path.empty() && (path.back() == '\r' || path.back() == ' ')) path.pop_back(); + const double t0 = now_ms(); + std::vector i32; + int64_t cells = 0; + if (!me::window_session_load(path.c_str(), i32, &cells, err)) { + std::printf("SERR io 0 %s\n", err.c_str()); + } else { + cached_ids.assign(i32.begin(), i32.end()); + DiskCache::read_pads(path + ".pads", cached_ids.size(), kScrPadToken, cached_ids, cached_pads); + ram_dirty = false; + struct stat st{}; + const size_t fbytes = ::stat(path.c_str(), &st) == 0 ? (size_t) st.st_size + : scr::fixed_bytes() + scr::cell_bytes((size_t) cells); + std::printf("RESTORED %lld %zu %.1f\n", (long long) cells, fbytes, now_ms() - t0); + } + std::fflush(stdout); + continue; + } + const bool geni = line.rfind("GENI ", 0) == 0; + if (!geni && line.rfind("GEN ", 0) != 0) continue; + + std::istringstream iss(line.substr(geni ? 5 : 4)); + int max_new = 512; + if (!(iss >> max_new)) continue; + float temp = 0.0f, top_p = 1.0f, min_p = 0.05f, pen_rep = 1.05f, pen_pres = 0.2f; + int top_k = 40; + std::vector req_img_keys; + std::string tok_str; + while (iss >> tok_str) { + auto eq = tok_str.find('='); + if (eq != std::string::npos) { + std::string k = tok_str.substr(0, eq), v = tok_str.substr(eq + 1); + if (k == "temperature") temp = std::stof(v); + else if (k == "top_p") top_p = std::stof(v); + else if (k == "top_k") top_k = std::stoi(v); + else if (k == "min_p") min_p = std::stof(v); + else if (k == "penalty_repeat") pen_rep = std::stof(v); + else if (k == "penalty_present") pen_pres = std::stof(v); + else if (k == "seed") rng.seed((uint32_t) std::stoul(v)); + else if (k == "img_keys") { + req_img_keys.clear(); + const char* q = v.c_str(); + while (*q != '\0') { + char* qe = nullptr; + const unsigned long long hk = std::strtoull(q, &qe, 16); + if (qe == q) { req_img_keys.clear(); break; } + req_img_keys.push_back((uint64_t) hk); + q = *qe == ',' ? qe + 1 : qe; + if (*qe != ',' && *qe != '\0') { req_img_keys.clear(); break; } + } + } + } else { + break; + } + } + std::string emb_path; + if (geni) { + emb_path = tok_str; + tok_str.clear(); + iss >> tok_str; + } + std::vector raw_ids; + std::stringstream tss(tok_str); + std::string item; + while (std::getline(tss, item, ',')) { + if (!item.empty()) raw_ids.push_back(std::stoll(item)); + } + if (raw_ids.empty()) { + std::printf("DONE 0 0 0.0 0.0 stop 0 0 0\n"); + std::fflush(stdout); + continue; + } + if ((int64_t) raw_ids.size() + max_new > max_cells) { + max_new = (int) std::max(1, max_cells - (int64_t) raw_ids.size() - 1); + } + + const double t_pre0 = now_ms(); + std::vector> req_imgs; + row_ptr.clear(); + img_rows.clear(); + if (geni) { + constexpr int64_t kImagePad = 248056; + const int64_t n = (int64_t) raw_ids.size(); + row_ptr.assign((size_t) n, nullptr); + const int64_t cells = (int64_t) mrope_host.size() / 3; + auto put = [&](int64_t c, int64_t t, int64_t h, int64_t w) { + mrope_host[(size_t) c * 3] = (int32_t) t; + mrope_host[(size_t) c * 3 + 1] = (int32_t) h; + mrope_host[(size_t) c * 3 + 2] = (int32_t) w; + }; + struct Img { int64_t n, nx, ny; size_t off; }; + std::vector imgs; + std::string ve; + std::FILE* f = std::fopen(emb_path.c_str(), "rb"); + if (!f) ve = "cannot open " + emb_path; + while (f && ve.empty()) { + int32_t hdr[5]; + const size_t got = std::fread(hdr, sizeof(int32_t), 5, f); + if (got == 0) break; + if (got != 5 || hdr[0] != 0x31455653 || hdr[1] < 1 || hdr[2] < 1 || hdr[3] < 1 || + (int64_t) hdr[2] * hdr[3] != hdr[1] || hdr[4] != 2560) { + ve = "bad embeddings file (expected strata-vision records of width 2560)"; + break; + } + const size_t off = img_rows.size(), cnt = (size_t) hdr[1] * (size_t) hdr[4]; + img_rows.resize(off + cnt); + if (std::fread(img_rows.data() + off, sizeof(float), cnt, f) != cnt) { ve = "short embeddings file"; break; } + imgs.push_back({hdr[1], hdr[2], hdr[3], off}); + } + if (f) std::fclose(f); + int64_t p = 0, i = 0; + size_t k = 0; + while (ve.empty() && i < n) { + if (raw_ids[(size_t) i] != kImagePad) { put(i, p, p, p); ++p; ++i; continue; } + if (k >= imgs.size()) { ve = "the prompt has more images than the embeddings file"; break; } + const Img& im = imgs[k++]; + uint64_t h = DiskCache::fnv1a(&im.n, 3); + if (req_img_keys.size() == imgs.size()) { + h ^= req_img_keys[k - 1] * 1099511628211ull; + } else { + h ^= DiskCache::fnv1a(reinterpret_cast(img_rows.data() + im.off), + (size_t) im.n * (2560 * sizeof(float) / sizeof(int64_t))); + } + req_imgs.push_back({i, h}); + for (int64_t j = 0; j < im.n && ve.empty(); ++j) { + if (i + j >= n || raw_ids[(size_t) (i + j)] != kImagePad) { + ve = "image has fewer <|image_pad|> tokens than rows"; + } + } + for (int64_t j = 0; j < im.n && ve.empty(); ++j) { + const int64_t y = j / im.nx, x = j % im.nx; + put(i + j, p, p + y, p + x); + row_ptr[(size_t) (i + j)] = img_rows.data() + im.off + (size_t) j * 2560; + } + i += im.n; + p += std::max(im.nx, im.ny); + } + if (ve.empty() && k != imgs.size()) ve = "the embeddings file has more images than the prompt"; + if (ve.empty() && n > 0 && raw_ids[(size_t) (n - 1)] == kImagePad) ve = "the prompt cannot end in an image"; + for (int64_t c = n; ve.empty() && c < cells; ++c) put(c, p + (c - n), p + (c - n), p + (c - n)); + if (!ve.empty()) { + me::window_mrope_host_set(nullptr); + std::printf("ERR %s\n", ve.c_str()); + std::fflush(stdout); + continue; + } + me::window_mrope_host_set(mrope_host.data()); + } else { + me::window_mrope_host_set(nullptr); + } + + std::vector ids = raw_ids; + // Carry forward any 0..3 SCR alignment pads from prior turns that still lie in the shared prefix + if (!geni && cached_imgs.empty() && !cached_pads.empty() && !cached_ids.empty()) { + std::vector carried_ids; + cached_pads = strata::core::scr_carry(raw_ids, cached_ids, cached_pads, kScrPadToken, carried_ids); + if ((int64_t) carried_ids.size() + max_new <= max_cells) { + ids = std::move(carried_ids); + } else { + cached_pads.clear(); + } + } + int64_t n_pre = std::max(0, (int64_t) ids.size() - 1); + size_t lcp = 0; + const size_t lim0 = std::min(cached_ids.size(), (size_t) n_pre); + while (lcp < lim0 && cached_ids[lcp] == ids[lcp]) ++lcp; + if (geni != !cached_imgs.empty()) { + // If switching between text and vision (or if image hashes differ inside 0..lcp), clamp lcp before the first differing image + size_t img_cap = lim0; + const size_t n_im = std::max(cached_imgs.size(), req_imgs.size()); + for (size_t ii = 0; ii < n_im; ++ii) { + if (ii >= cached_imgs.size() || ii >= req_imgs.size() || cached_imgs[ii] != req_imgs[ii]) { + const int64_t p_cut = std::min(ii < cached_imgs.size() ? cached_imgs[ii].first : (int64_t) lim0, + ii < req_imgs.size() ? req_imgs[ii].first : (int64_t) lim0); + img_cap = (size_t) std::max(0, p_cut); + break; + } + } + lcp = std::min(lcp, img_cap); + } else if (geni) { + for (size_t ii = 0; ii < req_imgs.size(); ++ii) { + if (ii >= cached_imgs.size() || cached_imgs[ii] != req_imgs[ii]) { + const int64_t p_cut = std::min(ii < cached_imgs.size() ? cached_imgs[ii].first : (int64_t) lcp, + req_imgs[ii].first); + lcp = std::min(lcp, (size_t) std::max(0, p_cut)); + break; + } + } + } + + // Estimate how many tokens the live RAM state + ck_ring can reuse before checking the NVMe tier + int64_t ram_reused = 0; + if (lcp == cached_ids.size() && lcp > 0) { + ram_reused = (int64_t) lcp; + } else if (!geni && cached_imgs.empty() && cached_ids.size() >= 32 && n_pre >= 32) { + std::vector ck_lens = {(int64_t) cached_ids.size()}; + for (const auto& s : ck_ring) { + if (!s.ids.empty() && s.ids.size() <= cached_ids.size() && + std::equal(s.ids.begin(), s.ids.end(), cached_ids.begin())) { + ck_lens.push_back((int64_t) s.ids.size()); + } + } + const strata::core::ScrPlan plan = strata::core::scr_plan( + cached_ids, ids, ck_lens, strata::core::ScrParams{32, 64, 16, 4}); + if (!plan.empty() && (int64_t) ids.size() + plan.pads() + max_new <= max_cells) { + ram_reused = plan.n_prefix + plan.reused(); + } + } + if (ram_reused == 0) { + for (const auto& s : ck_ring) { + if (!s.ids.empty() && (int64_t) s.ids.size() <= n_pre && (int64_t) s.ids.size() > ram_reused && + std::equal(s.ids.begin(), s.ids.end(), ids.begin())) { + ram_reused = (int64_t) s.ids.size(); + } + } + } + + // Check NVMe disk tier if it can beat the live in-RAM reuse by >= 32 tokens (text conversations only) + if (!geni && disk_cache.enabled()) { + int64_t disk_reused = 0; + const int disk_idx = disk_cache.find_best(raw_ids, max_new, max_cells, kScrPadToken, disk_reused); + if (disk_idx >= 0 && disk_reused >= ram_reused + 32) { + if (ram_dirty && cached_ids.size() >= 32) { + disk_cache.save_current(cached_ids, cached_pads, false, err, &active_path); + ram_dirty = false; + } + const int disk_idx2 = disk_cache.find_best(raw_ids, max_new, max_cells, kScrPadToken, disk_reused); + if (disk_idx2 >= 0) { + const std::string load_path = disk_cache.entries[(size_t) disk_idx2].path; + const std::string load_parent = disk_cache.entries[(size_t) disk_idx2].parent_path; + const bool load_pinned = disk_cache.entries[(size_t) disk_idx2].pinned; + const std::vector load_pads = disk_cache.entries[(size_t) disk_idx2].pads; + std::vector i32; + int64_t cells = 0; + if (me::window_session_load(load_path.c_str(), i32, &cells, err, load_parent.c_str())) { + cached_ids.assign(i32.begin(), i32.end()); + cached_pads = load_pads; + active_path = load_pinned ? "" : load_path; + ram_dirty = false; + ck_ring.clear(); + save_ck_at(cached_ids); + disk_cache.touch((size_t) disk_idx2); + ids = raw_ids; + if (!cached_pads.empty() && !cached_ids.empty()) { + std::vector carried_ids; + cached_pads = strata::core::scr_carry(raw_ids, cached_ids, cached_pads, kScrPadToken, carried_ids); + if ((int64_t) carried_ids.size() + max_new <= max_cells) ids = std::move(carried_ids); + else cached_pads.clear(); + } + n_pre = std::max(0, (int64_t) ids.size() - 1); + lcp = 0; + const size_t lim1 = std::min(cached_ids.size(), (size_t) n_pre); + while (lcp < lim1 && cached_ids[lcp] == ids[lcp]) ++lcp; + } + } + } else if (ram_reused == 0 && ram_dirty && cached_ids.size() >= 32) { + disk_cache.save_current(cached_ids, cached_pads, false, err, &active_path); + ram_dirty = false; + } + } + + const int64_t learn_at = (!geni && disk_cache.enabled() && cached_pads.empty()) + ? disk_cache.observe_prefix(ids, n_pre) + : -1; + int64_t reused = 0; + + auto on_chunk = [&](const float*, int64_t k, int64_t p0, std::string& e) -> bool { + if (!use_mtp) return true; + std::vector nxt((size_t) k); + for (int64_t t = 0; t < k; ++t) nxt[(size_t) t] = (int32_t) ids[(size_t) (p0 + t + 1)]; + return me::mtp_prefill(nxt.data(), k, p0, e); + }; + const float* const* embd_ptr = geni ? row_ptr.data() : nullptr; + auto run_prefill_span = [&](int64_t start_p, int64_t end_p) -> bool { + if (end_p <= start_p) return true; + if (learn_at > start_p && learn_at <= end_p && cached_pads.empty()) { + if (!me::prompt_run(ids.data() + start_p, learn_at - start_p, start_p, embd_ptr, on_chunk, err)) return false; + std::vector pfx(ids.begin(), ids.begin() + learn_at); + save_ck_at(pfx); + disk_cache.save_current(pfx, {}, /*pinned=*/true, err); + if (end_p > learn_at) { + if (!me::prompt_run(ids.data() + learn_at, end_p - learn_at, learn_at, embd_ptr, on_chunk, err)) return false; + } + return true; + } + return me::prompt_run(ids.data() + start_p, end_p - start_p, start_p, embd_ptr, on_chunk, err); + }; + + if (lcp == cached_ids.size() && lcp > 0) { + reused = (int64_t) lcp; + if (reused > 0) { + std::printf("RESUME %lld\n", (long long) reused); + std::fflush(stdout); + } + if (n_pre > reused) { + if (!run_prefill_span(reused, n_pre)) { + std::printf("ERR %s\n", err.c_str()); + std::fflush(stdout); + cached_ids.clear(); + ram_dirty = false; + continue; + } + } + cached_ids.assign(ids.begin(), ids.begin() + n_pre); + ram_dirty = true; + } else { + // 2. Check SCR mid-context fold [A B C] -> [A B' C D] (including stripped blocks) via scr_plan + if (!geni && cached_imgs.empty() && cached_ids.size() >= 32 && n_pre >= 32) { + std::vector ck_lens = {(int64_t) cached_ids.size()}; + for (const auto& s : ck_ring) { + if (!s.ids.empty() && s.ids.size() <= cached_ids.size() && + std::equal(s.ids.begin(), s.ids.end(), cached_ids.begin())) { + ck_lens.push_back((int64_t) s.ids.size()); + } + } + const strata::core::ScrPlan plan = strata::core::scr_plan( + cached_ids, ids, ck_lens, strata::core::ScrParams{32, 64, 16, 4}); + if (!plan.empty() && (int64_t) ids.size() + plan.pads() + max_new <= max_cells) { + ids = strata::core::scr_padded(ids, plan, kScrPadToken); + n_pre = std::max(0, (int64_t) ids.size() - 1); + const auto new_pads = strata::core::scr_pad_positions(plan); + cached_pads.insert(cached_pads.end(), new_pads.begin(), new_pads.end()); + std::sort(cached_pads.begin(), cached_pads.end()); + cached_pads.erase(std::unique(cached_pads.begin(), cached_pads.end()), cached_pads.end()); + save_ck_at(cached_ids); + bool scr_ok = true; + // Pass 1: relocate left-shifting runs (shift < 0) in forward order 0 .. m-1 + for (size_t ri = 0; scr_ok && ri < plan.runs.size(); ++ri) { + const auto& r = plan.runs[ri]; + if (r.shift() < 0 && !me::window_scr_relocate(r.old_start, r.at, r.n, err)) scr_ok = false; + } + // Pass 2: relocate right-shifting runs (shift > 0) in reverse order m-1 .. 0 + for (size_t ri = plan.runs.size(); scr_ok && ri-- > 0;) { + const auto& r = plan.runs[ri]; + if (r.shift() > 0 && !me::window_scr_relocate(r.old_start, r.at, r.n, err)) scr_ok = false; + } + int64_t cur_p = plan.restore > 0 ? plan.restore : (int64_t) lcp; + if (scr_ok && plan.restore > 0) { + for (const auto& s : ck_ring) { + if ((int64_t) s.ids.size() == plan.restore) { + me::window_checkpoint_restore(s.ck, err); + ss.ple_prev[0] = s.prev0; + ss.ple_prev[1] = s.prev1; + break; + } + } + } + // Pass 3: prefill edited gaps [cur_p, r.at) now that all runs sit at [r.at, r.at + r.n) + for (const auto& r : plan.runs) { + if (!scr_ok) break; + if (r.at > cur_p) { + if (!me::prompt_run_rows(ids.data() + cur_p, r.at - cur_p, cur_p, on_chunk, err)) { + scr_ok = false; + break; + } + } + cur_p = r.at + r.n; + } + if (scr_ok && plan.adopt > 0) { + for (const auto& s : ck_ring) { + if ((int64_t) s.ids.size() == plan.adopt) { + me::window_checkpoint_restore(s.ck, err); + ss.ple_prev[0] = s.prev0; + ss.ple_prev[1] = s.prev1; + break; + } + } + } + if (scr_ok && n_pre > cur_p) { + if (!me::prompt_run_rows(ids.data() + cur_p, n_pre - cur_p, cur_p, on_chunk, err)) { + scr_ok = false; + } + } + if (scr_ok) { + reused = plan.n_prefix + plan.reused(); + std::printf("RESUME %lld\n", (long long) reused); + std::fflush(stdout); + cached_ids.assign(ids.begin(), ids.begin() + n_pre); + ram_dirty = true; + } else { + cached_ids.clear(); + cached_pads.clear(); + ram_dirty = false; + } + } + } + if (reused == 0) { + int best_idx = -1; + size_t best_len = 0; + for (size_t ci = 0; ci < ck_ring.size(); ++ci) { + const auto& cids = ck_ring[ci].ids; + if (!cids.empty() && cids.size() <= (size_t) n_pre && cids.size() > best_len && + std::equal(cids.begin(), cids.end(), ids.begin())) { + best_len = cids.size(); + best_idx = (int) ci; + } + } + if (best_idx >= 0) { + me::window_checkpoint_restore(ck_ring[(size_t) best_idx].ck, err); + ss.ple_prev[0] = ck_ring[(size_t) best_idx].prev0; + ss.ple_prev[1] = ck_ring[(size_t) best_idx].prev1; + reused = (int64_t) best_len; + std::printf("RESUME %lld\n", (long long) reused); + std::fflush(stdout); + if (n_pre > reused) { + if (!run_prefill_span(reused, n_pre)) { + std::printf("ERR %s\n", err.c_str()); + std::fflush(stdout); + cached_ids.clear(); + ram_dirty = false; + continue; + } + } + cached_ids.assign(ids.begin(), ids.begin() + n_pre); + ram_dirty = true; + } else { + me::window_reset(); + ss.ple_prev[0] = ss.ple_prev[1] = -1; + cached_ids.clear(); + cached_pads.clear(); + active_path.clear(); + if (n_pre > 0) { + if (!run_prefill_span(0, n_pre)) { + std::printf("ERR %s\n", err.c_str()); + std::fflush(stdout); + ram_dirty = false; + continue; + } + } + cached_ids.assign(ids.begin(), ids.begin() + n_pre); + ram_dirty = true; + } + } + } + + const double prefill_ms = now_ms() - t_pre0; + const double t_dec0 = now_ms(); + int32_t cur_tok = (int32_t) ids.back(); + int64_t pos = n_pre; + int gen = 0; + const char* finish = "length"; + + if ((use_mtp || (suffix_draft > 0 && temp == 0.0f))) { + strata::spec::SuffixDrafter suffix(4, 32, (size_t) max_cells); + for (int64_t id : ids) suffix.append((int32_t) id); + int32_t mtp_drafts[8] = {}; + float mtp_probs[8] = {}; + int n_mtp_drafts = 0; + bool first_win = true; + SpecControl spec_control; + while (gen < max_new) { + int32_t input[8] = {cur_tok}, output[8] = {}; + int nd = 0; + const int rem = max_new - gen; + nd = spec_control.propose(suffix, input, mtp_drafts, mtp_probs, + !first_win && use_mtp ? n_mtp_drafts : 0, + !first_win ? suffix_draft : 0, rem, pos, max_cells); + first_win = false; + const double verify_begin = now_ms(); + if (!me::window_run(w, 1 + nd, input, pos, true, output, err)) { + std::printf("ERR %s\n", err.c_str()); + std::fflush(stdout); + cached_ids.clear(); + finish = "error"; + break; + } + spec_control.gate.spec_gate_forward(1 + nd, (now_ms() - verify_begin) / 1000); + // the penalties see the window's tokens (cached_ids + input[0..keep-1]), as the one-token path's + int keep; + if (temp > 0.0f && spec_exact) { + keep = spec_sample(w, input, output, nd, n_vocab, temp, top_k, top_p, min_p, pen_rep, pen_pres, cached_ids, lg_buf, rng); + } else { + keep = strata::spec::suffix_keep(input, output, nd); + cached_ids.insert(cached_ids.end(), input, input + keep); + if (temp > 0.0f && !strata::spec::suffix_stop(output[keep - 1])) { + me::window_copy_logits(w, keep - 1, lg_buf.data()); + output[keep - 1] = sample_logits(lg_buf.data(), n_vocab, output[keep - 1], temp, top_k, top_p, min_p, pen_rep, pen_pres, cached_ids, rng); + } + } + spec_control.result(nd, keep); + if (!me::window_commit(w, keep, err)) { + std::printf("ERR %s\n", err.c_str()); + std::fflush(stdout); + cached_ids.clear(); + finish = "error"; + break; + } + for (int i = 0; i < keep; ++i) { + std::printf("T %d\n", (int) output[i]); + std::fflush(stdout); + suffix.append(output[i]); + cur_tok = output[i]; + ++gen; + if (strata::spec::suffix_stop(cur_tok)) break; + } + if (!strata::spec::suffix_stop(cur_tok) && gen < max_new && use_mtp) { + const double draft_begin = now_ms(); + if (!me::mtp_draft(1 + nd, output, pos, keep - 1, spec - 1, mtp_drafts, mtp_probs, spec_min_p, &n_mtp_drafts, err)) { + std::printf("ERR %s\n", err.c_str()); + std::fflush(stdout); + cached_ids.clear(); + finish = "error"; + break; + } + spec_control.gate.spec_gate_draft_cost(0, (now_ms() - draft_begin) / (1000 * std::max(1, spec - 1))); + } else if (use_mtp) { + me::mtp_prefill(output, keep, pos, err); + } + pos += keep; + if (strata::spec::suffix_stop(cur_tok)) { + finish = "stop"; + break; + } + } + spec_control.describe(); + } else { + for (int step = 0; step < max_new; ++step) { + int32_t greedy_tok = 0; + if (!me::window_run(w, 1, &cur_tok, pos, true, &greedy_tok, err) || + !me::window_commit(w, 1, err)) { + std::printf("ERR %s\n", err.c_str()); + std::fflush(stdout); + cached_ids.clear(); + finish = "error"; + break; + } + cached_ids.push_back(cur_tok); + ++pos; + + int32_t next_tok = greedy_tok; + if (temp > 0.0f) { + me::window_copy_logits(w, 0, lg_buf.data()); + next_tok = sample_logits(lg_buf.data(), n_vocab, greedy_tok, temp, top_k, top_p, min_p, pen_rep, pen_pres, cached_ids, rng); + } + std::printf("T %d\n", (int) next_tok); + std::fflush(stdout); + ++gen; + + if (next_tok == 248046 || next_tok == 248044) { + finish = "stop"; + break; + } + cur_tok = next_tok; + } + } + cached_imgs = req_imgs; + if (!geni) save_ck_at(cached_ids); + const double decode_ms = now_ms() - t_dec0; + if (std::strcmp(finish, "error") != 0) { + std::printf("DONE %d %zu %.1f %.1f %s 0 0 %lld\n", + gen, ids.size(), prefill_ms, decode_ms, finish, (long long) reused); + std::fflush(stdout); + if (!geni && disk_cache.enabled() && cached_ids.size() >= 32) { + disk_cache.save_current(cached_ids, cached_pads, false, err, &active_path); + ram_dirty = false; + } + } + } + return 0; +} + +int main(int argc, char** argv) { + bool serve_mode = false; + std::string metal_gguf, ple_gguf, mtp_gguf, tokens_file, logits_out, logits_ref, logits_ref2; + int n_expert = 256; + int64_t max_cells = 4096; + int n_prompt = 22; + int n_gen = 64; + int suffix_draft = 0; + int spec = 1; + float spec_min_p = 0.0f; + std::string cache_disk_dir = std::getenv("STRATA_METAL_CACHE_DISK") ? std::getenv("STRATA_METAL_CACHE_DISK") : ""; + int64_t cache_disk_mib = std::getenv("STRATA_METAL_CACHE_DISK_MIB") ? std::atoll(std::getenv("STRATA_METAL_CACHE_DISK_MIB")) : 16384; + std::string save_session, load_session, scr_a, scr_b, scr_bp, scr_c, scr_d, scr_j = "0,8,16", scr_tmp = "/tmp/strata_scr_test.session"; + int32_t scr_pad = 198; + + for (int i = 1; i < argc; ++i) { + std::string a = argv[i]; + if (a == "--metal-gguf" && i + 1 < argc) metal_gguf = argv[++i]; + else if (a == "--ple-gguf" && i + 1 < argc) ple_gguf = argv[++i]; + else if (a == "--tokens-file" && i + 1 < argc) tokens_file = argv[++i]; + else if (a == "--n-expert" && i + 1 < argc) n_expert = std::atoi(argv[++i]); + else if ((a == "--ctx" || a == "--max-context") && i + 1 < argc) max_cells = std::atoll(argv[++i]); + else if (a == "--serve") serve_mode = true; + else if (a == "--vision") { /* enabled automatically via GENI */ } + else if (a == "--n-prompt" && i + 1 < argc) n_prompt = std::atoi(argv[++i]); + else if (a == "--n-gen" && i + 1 < argc) n_gen = std::atoi(argv[++i]); + else if (a == "--suffix-draft" && i + 1 < argc) suffix_draft = std::max(0, std::min(7, std::atoi(argv[++i]))); + else if (a == "--mtp-gguf" && i + 1 < argc) mtp_gguf = argv[++i]; + else if (a == "--spec" && i + 1 < argc) spec = std::max(1, std::min(4, std::atoi(argv[++i]))); + else if (a == "--spec-min-p" && i + 1 < argc) spec_min_p = (float) std::atof(argv[++i]); + else if (a == "--cache-disk" && i + 1 < argc) cache_disk_dir = argv[++i]; + else if (a == "--cache-disk-size" && i + 1 < argc) cache_disk_mib = std::atoll(argv[++i]); + else if (a == "--logits-out" && i + 1 < argc) logits_out = argv[++i]; + else if (a == "--logits-ref" && i + 1 < argc) logits_ref = argv[++i]; + else if (a == "--logits-ref2" && i + 1 < argc) logits_ref2 = argv[++i]; + else if (a == "--save-session" && i + 1 < argc) save_session = argv[++i]; + else if (a == "--load-session" && i + 1 < argc) load_session = argv[++i]; + else if (a == "--scr-a" && i + 1 < argc) scr_a = argv[++i]; + else if (a == "--scr-b" && i + 1 < argc) scr_b = argv[++i]; + else if (a == "--scr-bp" && i + 1 < argc) scr_bp = argv[++i]; + else if (a == "--scr-c" && i + 1 < argc) scr_c = argv[++i]; + else if (a == "--scr-d" && i + 1 < argc) scr_d = argv[++i]; + else if (a == "--scr-j" && i + 1 < argc) scr_j = argv[++i]; + else if (a == "--scr-pad" && i + 1 < argc) scr_pad = std::atoi(argv[++i]); + else if (a == "--scr-tmp" && i + 1 < argc) scr_tmp = argv[++i]; + } + if (metal_gguf.empty()) { + std::fprintf(stderr, "usage: strata_window_cli --metal-gguf [--ple-gguf ] " + "[--n-expert 256] [--ctx 4096] [--n-prompt 22] [--n-gen 64] [--tokens-file ] " + "[--logits-out ] [--logits-ref ] [--save-session | --load-session ] " + "[--scr-a A --scr-b B --scr-bp B' --scr-c C --scr-d D [--scr-j 0,8,16] [--scr-pad 198]]\n"); + return 2; + } + + const bool tf_mode = !logits_out.empty() || !logits_ref.empty(); + std::vector all_ids; + if (!tokens_file.empty()) { + std::ifstream f(tokens_file); + int64_t tid; + while (f >> tid) all_ids.push_back(tid); + } + if (all_ids.empty()) { + const int64_t def_ids[] = { + 248045, 8678, 198, 2523, 513, 264, 10631, 17313, 13, 248046, 198, 248045, 846, 198, 814, 20139, 1204, 27413, 71284, 3728, 13081, 32907, 12361, 321, 264, 3074, 6700, 11, 321, 948, 264, 28336, 628, 6561, 571, 13, 248046, 198, 248045, 74455, 198, 248068, 198 + }; + const size_t nd = sizeof(def_ids) / sizeof(def_ids[0]); + for (int i = 0; i < n_prompt + (tf_mode ? n_gen : 0); ++i) all_ids.push_back(def_ids[i % nd]); + } + if (tf_mode && (int) all_ids.size() < n_prompt + n_gen) { + std::fprintf(stderr, "strata_window_cli: teacher-forced mode needs at least n_prompt + n_gen (%d) tokens in --tokens-file, got %zu\n", + n_prompt + n_gen, all_ids.size()); + return 1; + } + std::vector prompt_ids(all_ids.begin(), all_ids.begin() + std::min(all_ids.size(), (size_t) n_prompt)); + + const double t_init0 = now_ms(); + std::string err; + strata::kernels::PleTable ple_table; + strata::core::SessionState ss{}; + ss.ple_prev[0] = -1; + ss.ple_prev[1] = -1; + ss.ple.consts = strata::kernels::ple_artifact_consts(); + if (!ple_gguf.empty()) { + strata::kernels::PleIoOptions io{}; + io.mode = strata::kernels::PleIo::Direct; + io.max_inflight = 256; + io.cache_rows = 1u << 18; + if (!ple_table.open(ple_gguf, err, io)) { + std::fprintf(stderr, "strata_window_cli: PleTable::open failed: %s\n", err.c_str()); + return 1; + } + ss.ple.table = &ple_table; + } + + strata::core::ModelGeometry g{}; + g.n_embd = 2560; + g.hc = 4; + g.hc_lr = 320; + g.n_ff = 640; + g.n_expert = n_expert; + g.n_layers = 48; + + strata::metal_engine::set_model_shard(metal_gguf.c_str()); + const bool use_mtp = !mtp_gguf.empty() && spec > 1; + const int max_win_t = std::max({1, 1 + suffix_draft, spec}); + if (use_mtp) { + setenv("STRATA_METAL_MTP_GGUF", mtp_gguf.c_str(), 1); + if (!strata::metal_engine::mtp_load(mtp_gguf, max_win_t, err)) { + std::fprintf(stderr, "strata_window_cli: mtp_load failed: %s\n", err.c_str()); + return 1; + } + } + void* w = strata::metal_engine::window_create(); + int64_t n_vocab = 0; + if (!strata::metal_engine::window_init(w, g, ss, max_win_t, max_cells, &n_vocab, err)) { + std::fprintf(stderr, "strata_window_cli: window_init failed: %s\n", err.c_str()); + return 1; + } + const double init_ms = now_ms() - t_init0; + if (serve_mode) return serve_loop(w, ss, n_vocab, max_cells, suffix_draft, use_mtp, spec, spec_min_p, cache_disk_dir, cache_disk_mib); + if (!scr_a.empty()) return scr_test(w, ss, n_vocab, scr_a, scr_b, scr_bp, scr_c, scr_d, scr_j, scr_pad, n_gen, scr_tmp); + + FILE* fout_lg = nullptr; + if (!logits_out.empty()) { + fout_lg = std::fopen(logits_out.c_str(), "wb"); + if (!fout_lg) { + std::perror("fopen logits_out"); + return 1; + } + fcntl(fileno(fout_lg), F_NOCACHE, 1); + } + FILE* fref_lg = nullptr; + if (!logits_ref.empty()) { + fref_lg = std::fopen(logits_ref.c_str(), "rb"); + if (!fref_lg) { + std::perror("fopen logits_ref"); + return 1; + } + fcntl(fileno(fref_lg), F_NOCACHE, 1); + } + FILE* fref2_lg = nullptr; + if (!logits_ref2.empty()) { + fref2_lg = std::fopen(logits_ref2.c_str(), "rb"); + if (!fref2_lg) { + std::perror("fopen logits_ref2"); + return 1; + } + fcntl(fileno(fref2_lg), F_NOCACHE, 1); + } + + strata::spec::SuffixDrafter suffix(4, 32, (size_t) max_cells); + for (int64_t tid : prompt_ids) suffix.append((int32_t) tid); + + // Prefill all but the last prompt token via prompt_run_rows + const int64_t n_pre = std::max(0, (int64_t) prompt_ids.size() - 1); + const double t_pre0 = now_ms(); + if (!load_session.empty()) { + // the prefix from disk instead of a prefill: it must be this prompt's first n_pre tokens + std::vector ids; + int64_t cells = 0; + if (!strata::metal_engine::window_session_load(load_session.c_str(), ids, &cells, err)) { + std::fprintf(stderr, "strata_window_cli: %s\n", err.c_str()); + return 1; + } + bool same = (std::getenv("STRATA_BENCH_EXPAND_CELLS") != nullptr) || (cells == n_pre && (int64_t) ids.size() == n_pre); + if (!std::getenv("STRATA_BENCH_EXPAND_CELLS")) { + for (int64_t i = 0; same && i < n_pre; ++i) same = ids[(size_t) i] == (int32_t) prompt_ids[(size_t) i]; + } + if (!same) { + std::fprintf(stderr, "strata_window_cli: the session holds %lld cells of another prefix (want this prompt's first %lld)\n", + (long long) cells, (long long) n_pre); + return 1; + } + std::printf("STRATA_SESSION load_ms=%.2f cells=%lld bytes=%zu\n", now_ms() - t_pre0, (long long) cells, + scr::fixed_bytes() + scr::cell_bytes((size_t) cells)); + } else if (n_pre > 0) { + auto on_chunk = [&](const float*, int64_t k, int64_t p0, std::string& e) -> bool { + if (!use_mtp) return true; + std::vector nxt((size_t) k); + for (int64_t t = 0; t < k; ++t) nxt[(size_t) t] = (int32_t) prompt_ids[(size_t) (p0 + t + 1)]; + return strata::metal_engine::mtp_prefill(nxt.data(), k, p0, e); + }; + if (!strata::metal_engine::prompt_run_rows(prompt_ids.data(), n_pre, 0, on_chunk, err)) { + std::fprintf(stderr, "strata_window_cli: prompt_run_rows failed: %s\n", err.c_str()); + return 1; + } + } + if (!save_session.empty()) { + std::vector ids(prompt_ids.begin(), prompt_ids.begin() + n_pre); + const double t_s = now_ms(); + if (!strata::metal_engine::window_session_save(save_session.c_str(), ids.data(), n_pre, err)) { + std::fprintf(stderr, "strata_window_cli: %s\n", err.c_str()); + return 1; + } + std::printf("STRATA_SESSION save_ms=%.2f cells=%lld\n", now_ms() - t_s, (long long) n_pre); + } + // Run the last prompt token with head=true to get the first output distribution + int32_t cur_tok = (int32_t) prompt_ids.back(); + int32_t next_tok = 0; + int64_t pos = std::getenv("STRATA_BENCH_EXPAND_CELLS") ? std::atoll(std::getenv("STRATA_BENCH_EXPAND_CELLS")) : n_pre; + if (!strata::metal_engine::window_run(w, 1, &cur_tok, pos, true, &next_tok, err) || + !strata::metal_engine::window_commit(w, 1, err)) { + std::fprintf(stderr, "strata_window_cli: first window_run failed: %s\n", err.c_str()); + return 1; + } + const double prefill_ms = now_ms() - t_pre0; + int32_t mtp_drafts[8] = {}; + float mtp_probs[8] = {}; + int n_mtp_drafts = 0; + double mtp_ms = 0.0; + if (use_mtp && !tf_mode && n_gen > 1) { + const double tm0 = now_ms(); + if (!strata::metal_engine::mtp_draft(1, &next_tok, pos, 0, spec - 1, mtp_drafts, mtp_probs, spec_min_p, &n_mtp_drafts, err)) { + std::fprintf(stderr, "strata_window_cli: mtp_draft(0) failed: %s\n", err.c_str()); + return 1; + } + mtp_ms += now_ms() - tm0; + } + ++pos; + + std::vector cur_lg(tf_mode ? (size_t) n_vocab : 0); + std::vector ref_lg((fref_lg || fref2_lg) ? (size_t) n_vocab : 0); + double sum_nll_cand = 0.0, sum_nll_ref = 0.0, sum_nll_ref2 = 0.0; + std::vector klds, klds2; + int top1_matches = 0, top1_matches2 = 0; + + auto eval_tf_step = [&](int step_idx) -> bool { + if (!strata::metal_engine::window_copy_logits(w, 0, cur_lg.data())) return false; + if (fout_lg) { + if (std::fwrite(cur_lg.data(), sizeof(float), (size_t) n_vocab, fout_lg) != (size_t) n_vocab) return false; + } + const int32_t tgt = (int32_t) all_ids[(size_t) n_prompt + (size_t) step_idx]; + int32_t cand_top1 = 0; + const double lse_cand = logsumexp_f64(cur_lg.data(), n_vocab, &cand_top1); + sum_nll_cand += -((double) cur_lg[(size_t) tgt] - lse_cand); + + if (fref_lg) { + if (std::fread(ref_lg.data(), sizeof(float), (size_t) n_vocab, fref_lg) != (size_t) n_vocab) return false; + int32_t ref_top1 = 0; + const double lse_ref = logsumexp_f64(ref_lg.data(), n_vocab, &ref_top1); + sum_nll_ref += -((double) ref_lg[(size_t) tgt] - lse_ref); + if (cand_top1 == ref_top1) ++top1_matches; + double kld = 0.0; + for (int64_t v = 0; v < n_vocab; ++v) { + const double log_p = (double) ref_lg[(size_t) v] - lse_ref; + const double log_q = (double) cur_lg[(size_t) v] - lse_cand; + const double p = std::exp(log_p); + if (p > 0.0) kld += p * (log_p - log_q); + } + klds.push_back(std::max(0.0, kld)); + } + if (fref2_lg) { + if (std::fread(ref_lg.data(), sizeof(float), (size_t) n_vocab, fref2_lg) != (size_t) n_vocab) return false; + int32_t ref2_top1 = 0; + const double lse_ref2 = logsumexp_f64(ref_lg.data(), n_vocab, &ref2_top1); + sum_nll_ref2 += -((double) ref_lg[(size_t) tgt] - lse_ref2); + if (cand_top1 == ref2_top1) ++top1_matches2; + double kld2 = 0.0; + for (int64_t v = 0; v < n_vocab; ++v) { + const double log_p = (double) ref_lg[(size_t) v] - lse_ref2; + const double log_q = (double) cur_lg[(size_t) v] - lse_cand; + const double p = std::exp(log_p); + if (p > 0.0) kld2 += p * (log_p - log_q); + } + klds2.push_back(std::max(0.0, kld2)); + } + return true; + }; + + if (tf_mode && n_gen > 0) { + if (!eval_tf_step(0)) { + std::fprintf(stderr, "strata_window_cli: eval_tf_step(0) failed\n"); + return 1; + } + } + + std::vector gen_tokens; + gen_tokens.push_back(next_tok); + suffix.append(next_tok); + cur_tok = tf_mode ? (int32_t) all_ids[(size_t) n_prompt] : next_tok; + + strata::metal_engine::window_reset_stats(w); + const int decode_steps = std::max(0, n_gen - 1); + const double t_dec0 = now_ms(); + int produced = 0, spec_rounds = 0, spec_drafted = 0, spec_accepted = 0; + SpecControl spec_control; + while (produced < decode_steps) { + int32_t win_in[8] = {cur_tok}; + int32_t win_out[8] = {}; + int nd = 0; + const int rem = decode_steps - produced; + nd = spec_control.propose(suffix, win_in, mtp_drafts, mtp_probs, + !tf_mode && use_mtp ? n_mtp_drafts : 0, + !tf_mode ? suffix_draft : 0, rem, pos, max_cells); + const int T = 1 + nd; + const double verify_begin = now_ms(); + if (!strata::metal_engine::window_run(w, T, win_in, pos, true, win_out, err)) { + std::fprintf(stderr, "strata_window_cli: decode step failed: %s\n", err.c_str()); + return 1; + } + spec_control.gate.spec_gate_forward(T, (now_ms() - verify_begin) / 1000); + const int n_keep = strata::spec::suffix_keep(win_in, win_out, nd, false); + spec_control.result(nd, n_keep); + if (!strata::metal_engine::window_commit(w, n_keep, err)) { + std::fprintf(stderr, "strata_window_cli: window_commit failed: %s\n", err.c_str()); + return 1; + } + ++spec_rounds; + spec_drafted += nd; + spec_accepted += (n_keep - 1); + if (tf_mode) { + if (!eval_tf_step(produced + 1)) { + std::fprintf(stderr, "strata_window_cli: eval_tf_step(%d) failed\n", produced + 1); + return 1; + } + gen_tokens.push_back(win_out[0]); + ++produced; + if (produced < decode_steps) { + cur_tok = (int32_t) all_ids[(size_t) n_prompt + (size_t) produced]; + } + pos += 1; + } else { + for (int i = 0; i < n_keep && produced < decode_steps; ++i) { + gen_tokens.push_back(win_out[i]); + suffix.append(win_out[i]); + cur_tok = win_out[i]; + ++produced; + } + if (use_mtp && produced < decode_steps) { + const double tm0 = now_ms(); + if (!strata::metal_engine::mtp_draft(T, win_out, pos, n_keep - 1, spec - 1, mtp_drafts, mtp_probs, spec_min_p, &n_mtp_drafts, err)) { + std::fprintf(stderr, "strata_window_cli: mtp_draft failed: %s\n", err.c_str()); + return 1; + } + const double draft_ms = now_ms() - tm0; + mtp_ms += draft_ms; + spec_control.gate.spec_gate_draft_cost(0, draft_ms / (1000 * std::max(1, spec - 1))); + } + pos += n_keep; + } + } + if (!tf_mode && (use_mtp || suffix_draft > 0)) spec_control.describe(); + const double decode_ms = now_ms() - t_dec0; + const double decode_tok_s = decode_steps > 0 ? (decode_steps * 1000.0 / decode_ms) : 0.0; + const double decode_ms_tok = decode_steps > 0 ? (decode_ms / decode_steps) : 0.0; + const double prefill_tok_s = prompt_ids.size() * 1000.0 / prefill_ms; + + if (fout_lg) std::fclose(fout_lg); + if (fref_lg) std::fclose(fref_lg); + if (fref2_lg) std::fclose(fref2_lg); + + std::printf("STRATA_RESULT init_ms=%.2f prompt_tokens=%zu prefill_ms=%.2f prefill_tok_s=%.2f " + "decode_tokens=%d decode_ms=%.2f decode_tok_s=%.2f decode_ms_per_tok=%.2f\n", + init_ms, prompt_ids.size(), prefill_ms, prefill_tok_s, + decode_steps, decode_ms, decode_tok_s, decode_ms_tok); + if (spec_rounds > 0 && (use_mtp || suffix_draft > 0)) { + std::printf("STRATA_SPEC rounds=%d drafted=%d accepted=%d accept_pct=%.1f tok_per_round=%.2f mtp_ms=%.2f\n", + spec_rounds, spec_drafted, spec_accepted, + spec_drafted > 0 ? 100.0 * (double) spec_accepted / (double) spec_drafted : 0.0, + (double) decode_steps / (double) spec_rounds, mtp_ms); + } + if (tf_mode && n_gen > 0) { + const double cand_ppl = std::exp(sum_nll_cand / n_gen); + if (!klds.empty()) { + const double ref_ppl = std::exp(sum_nll_ref / n_gen); + double mean_kld = 0.0; + for (double x : klds) mean_kld += x; + mean_kld /= klds.size(); + std::vector sk = klds; + std::sort(sk.begin(), sk.end()); + const double med_kld = sk[sk.size() / 2]; + const double p90_kld = sk[(size_t) (0.90 * (sk.size() - 1))]; + const double top1_pct = 100.0 * (double) top1_matches / (double) n_gen; + std::printf("STRATA_KLD eval_tokens=%d ref_ppl=%.4f cand_ppl=%.4f delta_ppl=%.4f mean_kld=%.6f median_kld=%.6f p90_kld=%.6f top1_agree_pct=%.2f\n", + n_gen, ref_ppl, cand_ppl, cand_ppl - ref_ppl, mean_kld, med_kld, p90_kld, top1_pct); + if (!klds2.empty()) { + const double ref2_ppl = std::exp(sum_nll_ref2 / n_gen); + double mean_kld2 = 0.0; + for (double x : klds2) mean_kld2 += x; + mean_kld2 /= klds2.size(); + std::vector sk2 = klds2; + std::sort(sk2.begin(), sk2.end()); + const double med_kld2 = sk2[sk2.size() / 2]; + const double p90_kld2 = sk2[(size_t) (0.90 * (sk2.size() - 1))]; + const double top1_pct2 = 100.0 * (double) top1_matches2 / (double) n_gen; + std::printf("STRATA_KLD_REF2 eval_tokens=%d ref_ppl=%.4f cand_ppl=%.4f delta_ppl=%.4f mean_kld=%.6f median_kld=%.6f p90_kld=%.6f top1_agree_pct=%.2f\n", + n_gen, ref2_ppl, cand_ppl, cand_ppl - ref2_ppl, mean_kld2, med_kld2, p90_kld2, top1_pct2); + } + } else { + std::printf("STRATA_PPL eval_tokens=%d cand_ppl=%.4f\n", n_gen, cand_ppl); + } + } + std::printf("GEN_IDS:"); + for (size_t i = 0; i < gen_tokens.size(); ++i) { + std::printf(" %d", gen_tokens[i]); + } + std::printf("\n"); + std::fflush(stdout); + + strata::metal_engine::window_destroy(w); + return 0; +} diff --git a/metal/tests/comb_fold_parity.cpp b/metal/tests/comb_fold_parity.cpp new file mode 100644 index 00000000..187fd5ae --- /dev/null +++ b/metal/tests/comb_fold_parity.cpp @@ -0,0 +1,1571 @@ +// metal/tests/comb_fold_parity.cpp - Round 5 Track A: the prompt chunk's folded mm_combine and the per-token HC read +// against the dispatch chains they replace, BITWISE; `comb_fold_parity [T] bench` adds their GPU times at T (default +// 1024), 48 reps = a chunk's layers: +// gr_norm_h_moec == mm_combine + gr_norm_h_moe (R', xnh) +// gr_apply_moec == mm_combine + gr_apply_moe (R') +// gr_norm_h4 == gr_norm_h (R', xnh; apply 0 and 1) +// gemm_hc_up, mixed unbound == gemm_hc_up's mixh +// `comb_fold_parity T grvec [baseline-prompt.metal]` checks/times GR only in a short process. +// `comb_fold_parity T hcup [bench|confirm]` checks HC-up tiles alone (both outputs, null mixed too). +// `comb_fold_parity T epi [bench [tile,tile...]]` checks gemm_hc_down / gemm_swiglu_up against gemm_f16 + gr_epi_down / swiglu_h. +// `comb_fold_parity T denseab baseline-dense_mm.metal [shape]` checks/times identical dense tiles across source revisions. +// `comb_fold_parity T gu [bench [tile,tile...]]` checks/times the fused shared-expert gate + up against gate + up + swiglu_h. +// `comb_fold_parity T topk baseline-prompt.metal` checks/times route_topk against a baseline copy's. +// `comb_fold_parity T ab baseline-prompt.metal [gr|gemm]` checks/times this prompt.metal's kernels against a baseline's. +// `comb_fold_parity 0 mtp [bench] [baseline-mtp.metal]` checks/times MTP's batched fc_hidden (mtp_parity). +// `comb_fold_parity 0 mtpcat` checks/times MTP's gm_cat against its two GEMVs (mtp_cat_parity). +#include "mtl_rt.h" +#include "ref_quant.h" +#include "coop_probe.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { +struct GrArgs { uint32_t n_embd, hc, hc_lr, n_tok, apply; float eps; }; +struct GemmArgs { int M, N, K, pad; }; +struct EpiArgs { uint32_t n_tok, n_embd, hc, hc_lr; }; +bool same(const mtl::Buffer& a, const mtl::Buffer& b, size_t bytes) { return std::memcmp(a.ptr, b.ptr, bytes) == 0; } + +int hc_up_parity(uint32_t T, bool bench, bool confirm = false) { + mtl::Context cx; + // Compile just this kernel for short standalone runs, without the quant/MoE libraries. + const std::string source = mtl::read_file(cx.kernel_dir + "/prompt.metal"); + const std::string prelude = "#define NO_GGML_COMMON 1\n" + "#include \n" + "using namespace metal; using namespace mpp::tensor_ops;\n" + "struct GemmArgs { int M, N, K, pad; };\n" + "struct EpiArgs { uint n_tok, n_embd, hc, hc_lr; };\n"; + const auto begin = source.find("// gemm_hc_up:"); + if (begin == std::string::npos) { + std::fprintf(stderr, "gemm_hc_up source marker missing\n"); + return 1; + } + std::string extra; + for (const auto shape : {std::pair{32, 128}, {64, 128}, {32, 64}, {64, 64}, {16, 64}}) { + const auto [tm, tn] = shape; + const std::string args = std::to_string(tm) + ", " + std::to_string(tn) + ", 2, true"; + extra += "template [[host_name(\"gemm_hc_up_reg_" + std::to_string(tm) + "_" + std::to_string(tn) + + "_2\")]] kernel decltype(gemm_hc_up<" + args + ">) gemm_hc_up<" + args + ">;\n"; + } + const int lib = cx.compile({}, prelude + source.substr(begin) + extra); + const bool k4 = coop_k4(cx); + constexpr uint32_t N = 2560, HC = 4, LR = 320; + const size_t RN = size_t(N) * HC; + auto A = cx.alloc(2 * T * LR), B = cx.alloc(2 * RN * LR), X = cx.alloc(2 * T * RN); + auto mr = cx.alloc(4 * T * N), hr = cx.alloc(2 * T * N); + auto mc = cx.alloc(4 * T * N), hc = cx.alloc(2 * T * N); + const mtl::Buffer nil{}; + ref::Rng rng(23 + T); + for (const auto* b : {&A, &B, &X}) + for (size_t i = 0; i < b->bytes / 2; ++i) + b->as()[i] = ref::f32_to_f16(rng.sym() * (b == &B ? 0.1f : 1.0f)); + const GemmArgs g{int(T), int(RN), int(LR), 0}; + const EpiArgs e{T, N, HC, LR}; + auto dispatch = [&](const mtl::Pipeline& p, int tm, int tn, const mtl::Buffer& m, const mtl::Buffer& h, int sg = 4) { + mtl::Dispatch d; + d.p = &p; + d.args = {mtl::Arg::B(A), mtl::Arg::B(B), mtl::Arg::B(X), mtl::Arg::B(m), mtl::Arg::B(h), + mtl::Arg::V(g), mtl::Arg::V(e)}; + d.grid[0] = RN / tn; + d.grid[1] = (T + tm - 1) / tm; + d.tg[0] = sg * 32; + return d; + }; + auto baseline = cx.pipeline(lib, "gemm_hc_up_64_64_4"); + cx.run({dispatch(baseline, 64, 64, mr, hr)}); + int fails = 0; + std::printf("device %s, T %u, register layout %s\n", cx.device_name().c_str(), T, k4 ? "K4" : "SKIP"); + for (bool reg : {false, true}) { + if (reg && !k4) continue; + for (const auto shape : {std::array{64, 64, 4}, {128, 64, 4}, {64, 128, 4}, {64, 32, 4}, + {32, 128, 4}, {32, 128, 2}, {64, 128, 2}, {32, 64, 2}, {64, 64, 2}, {16, 64, 2}}) { + const auto [tm, tn, sg] = shape; + if (!reg && (sg == 2 || (tm == 32 && tn == 128))) continue; + const std::string name = std::string("gemm_hc_up_") + (reg ? "reg_" : "") + + std::to_string(tm) + "_" + std::to_string(tn) + "_" + std::to_string(sg); + auto p = cx.pipeline(lib, name); + for (bool bound : {true, false}) { + std::memset(mc.ptr, 0xCD, mc.bytes); + std::memset(hc.ptr, 0xCD, hc.bytes); + auto d = dispatch(p, tm, tn, bound ? mc : nil, hc, sg); + cx.run({d}); + const bool ok = same(hr, hc, hr.bytes) && (!bound || same(mr, mc, mr.bytes)); + fails += !ok; + std::printf("%s mixed %s: bitwise %s tgmem %d", name.c_str(), bound ? "bound" : "unbound", + ok ? "PASS" : "FAIL", p.static_tg_mem); + if (bench) { + cx.warm({d}, 0.01); + double best = 1e9; + for (int i = 0; i < 3; ++i) best = std::min(best, cx.run({d}, 48)); + std::printf(" ms48 %.3f", best * 1e3); + } + std::printf("\n"); + } + } + } + if (fails) return 1; + if (confirm && k4) { + auto r64 = cx.pipeline(lib, "gemm_hc_up_reg_64_64_4"); + auto r128 = cx.pipeline(lib, "gemm_hc_up_reg_64_128_4"); + for (bool bound : {false, true}) { + const auto& m = bound ? mc : nil; + const std::vector ds = {dispatch(baseline, 64, 64, m, hc), + dispatch(r64, 64, 64, m, hc), dispatch(r128, 64, 128, m, hc)}; + for (const auto& d : ds) cx.warm({d}, 0.01); + for (int rep = 0; rep < 5; ++rep) { + double ms[3]; + for (int j = 0; j < 3; ++j) { + const int i = (j + rep) % 3; + ms[i] = cx.run({ds[i]}, 48) * 1e3; + } + std::printf("confirm mixed %s rep %d baseline %.3f reg64 %.3f reg128 %.3f\n", + bound ? "bound" : "unbound", rep, ms[0], ms[1], ms[2]); + } + } + } + return 0; +} + +// the window's gemm tile for (M, N, K) (window.cpp prompt_chunk: gemm), as "TM_TN_SG" +std::string window_tile(int M, int Nn, int K) { + int tm = 64, tn = 128, sg = 4; + auto set = [&](int m, int n, int s) { tm = m; tn = n; sg = s; }; + if (M <= 128) { + if (Nn <= 64) set(64, 32, 4); + else if (Nn <= 640 || K > 2560) set(64, Nn <= 128 ? 32 : 64, 4); + else set(M >= 128 ? 128 : 64, 64, 4); + } else if (Nn <= 64) { + M < 1024 ? set(16, 64, 2) : set(64, 32, 2); + } else if (Nn == 96) { + M < 640 ? set(16, 64, 2) : set(16, 96, 2); + } else if (Nn <= 128) { + M < 768 ? set(16, 64, 2) : set(32, 64, 2); + } else if (Nn <= 640 && K > 2560) { + M <= 512 ? set(64, 64, 4) : M < 1536 ? set(128, 64, 4) : set(32, 128, 2); + } else if (Nn <= 512 || (Nn <= 640 && M < 384)) { + M < 384 ? set(32, 64, 2) : set(32, 128, 2); + } else if (Nn <= 1664) { + M < 768 ? set(32, 128, 2) : set(64, 128, 4); + } else { + M < 1536 && (M % 128 == 0 || M % 128 > 64) ? set(128, 64, 4) : set(64, 128, 4); + } + return std::to_string(tm) + "_" + std::to_string(tn) + "_" + std::to_string(sg); +} + +// `comb_fold_parity T epi [bench]`: gemm_hc_down == gemm_f16 + gr_epi_down (lo_h, inject) and gemm_swiglu_up == +// gemm_f16 (up) + swiglu_h (h), BITWISE, against the window's gemm_f16 tile, every epilogue tile; T 0 = the sweep +// 1..2048. bench: paired GPU times at T, 96 HC reads / 48 shared experts (a chunk's), alternating order, median of 5 +int epi_parity(uint32_t TB, bool bench, const char* only) { + mtl::Context cx; + const std::string dm = mtl::read_file(cx.kernel_dir + "/dense_mm.metal"), pr = mtl::read_file(cx.kernel_dir + "/prompt.metal"); + const auto gb = dm.find("template "), ge = dm.find("#define DQ(F)"); + const auto ib = dm.find("template [[host_name(\"gemm_f16_"); + const auto eb = pr.find("struct EpiArgs {"), ee = pr.find("// ---- the HC read"), hb = pr.find("// gemm_hc_up:"); + if (gb == std::string::npos || ge == std::string::npos || ib == std::string::npos || eb == std::string::npos || + ee == std::string::npos || hb == std::string::npos) { + std::fprintf(stderr, "epilogue source markers missing\n"); + return 1; + } + const std::string prelude = "#define NO_GGML_COMMON 1\n" + "#include \n" + "using namespace metal; using namespace mpp::tensor_ops;\n" + "struct GemmArgs { int M, N, K; int pad; };\n"; + if (!coop_k4(cx)) { + std::printf("device %s: not the M5 layout, the register epilogues are not used: SKIP\n", cx.device_name().c_str()); + return 0; + } + // the sweep's other tiles (and gemm_f16_32_128_4, so each epilogue tile has its own plain GEMM), only here + std::string extra = "template [[host_name(\"gemm_f16_32_128_4\")]] kernel decltype(gemm_f16<32, 128, 4>) gemm_f16<32, 128, 4>;\n"; + for (const char* t : {"64_64_4", "128_64_4", "64_128_4", "32_128_4", "32_128_2", "16_64_2", "32_64_2"}) { + int tm, tn, sg; + std::sscanf(t, "%d_%d_%d", &tm, &tn, &sg); + const std::string targs = std::to_string(tm) + ", " + std::to_string(tn) + ", " + std::to_string(sg); + for (const char* k : {"gemm_hc_down", "gemm_swiglu_up"}) { + const std::string name = std::string(k) + "_" + t; + if (pr.find("\"" + name + "\"") == std::string::npos) + extra += "template [[host_name(\"" + name + "\")]] kernel decltype(" + k + "<" + targs + ">) " + k + "<" + targs + ">;\n"; + } + } + const int lib = cx.compile({}, prelude + dm.substr(gb, ge - gb) + dm.substr(ib) + pr.substr(eb, ee - eb) + pr.substr(hb) + extra); + constexpr uint32_t N = 2560, HC = 4, LR = 320, NFF = 640, NR = LR + HC, KD = HC * N; + std::printf("device %s\n", cx.device_name().c_str()); + std::map ps; + auto P = [&](const std::string& n) -> const mtl::Pipeline& { + auto it = ps.find(n); + if (it == ps.end()) it = ps.emplace(n, cx.pipeline(lib, n)).first; + return it->second; + }; + auto dims = [](const std::string& tile, int& tm, int& tn, int& sg) { std::sscanf(tile.c_str(), "%d_%d_%d", &tm, &tn, &sg); }; + auto D = [](const mtl::Pipeline& p, std::vector a, uint32_t gx, uint32_t gy, uint32_t th, bool barrier = true) { + mtl::Dispatch d; + d.p = &p; + d.args = std::move(a); + d.grid[0] = gx; + d.grid[1] = gy; + d.tg[0] = th; + d.barrier = barrier; + return d; + }; + std::vector tiles = {"32_128_2", "64_128_4", "128_64_4", "32_128_4", "16_64_2", "64_64_4", "32_64_2"}; + if (only) { + tiles.clear(); + for (const char *s = only, *e; *s; s = *e ? e + 1 : e) { + e = std::strchr(s, ','); + if (!e) e = s + std::strlen(s); + tiles.emplace_back(s, e); + } + } + int fails = 0; + const std::vector Ts = TB ? std::vector{TB} : + std::vector{1u, 7u, 16u, 17u, 32u, 33u, 64u, 256u, 289u, 297u, 306u, 512u, + 1024u, 1025u, 1041u, 1983u, 1987u, 2048u}; + for (uint32_t T : Ts) { + mtl::Buffer xnh = cx.alloc(2 * size_t(T) * KD), WD = cx.alloc(2 * size_t(NR) * KD), V = cx.alloc(4 * size_t(T) * NR), + la = cx.alloc(2 * T * LR), lb = cx.alloc(2 * T * LR), ia = cx.alloc(4 * T * HC), ib_ = cx.alloc(4 * T * HC), + mixh = cx.alloc(2 * size_t(T) * N), Wg = cx.alloc(2 * size_t(NFF) * N), Wu = cx.alloc(2 * size_t(NFF) * N), + gsh = cx.alloc(4 * size_t(T) * NFF), ush = cx.alloc(4 * size_t(T) * NFF), ha = cx.alloc(2 * size_t(T) * NFF), + hb_ = cx.alloc(2 * size_t(T) * NFF); + ref::Rng r(31 + T); + for (const auto& [b, s] : {std::pair{&xnh, 1.0f}, {&WD, 0.02f}, {&mixh, 1.0f}, {&Wg, 0.05f}, {&Wu, 0.05f}}) + for (size_t i = 0; i < b->bytes / 2; ++i) b->as()[i] = ref::f32_to_f16(s * r.sym()); + const EpiArgs ea{T, N, HC, LR}; + const GemmArgs gd{int(T), int(NR), int(KD), 0}, gs{int(T), int(NFF), int(N), 0}; + auto gemm = [&](const std::string& tile, const mtl::Buffer& A, const mtl::Buffer& B, const mtl::Buffer& Y, + const GemmArgs& g, bool barrier = true) { + int tm, tn, sg; + dims(tile, tm, tn, sg); + return D(P("gemm_f16_" + tile), {mtl::Arg::B(A), mtl::Arg::B(B), mtl::Arg::B(Y), mtl::Arg::V(g)}, + (g.N + tn - 1) / tn, (g.M + tm - 1) / tm, sg * 32, barrier); + }; + auto down = [&](const std::string& tile) { + int tm, tn, sg; + dims(tile, tm, tn, sg); + return D(P("gemm_hc_down_" + tile), {mtl::Arg::B(xnh), mtl::Arg::B(WD), mtl::Arg::B(lb), mtl::Arg::B(ib_), + mtl::Arg::V(ea)}, (NR + tn - 1) / tn, (T + tm - 1) / tm, sg * 32); + }; + auto swup = [&](const std::string& tile) { + int tm, tn, sg; + dims(tile, tm, tn, sg); + return D(P("gemm_swiglu_up_" + tile), {mtl::Arg::B(mixh), mtl::Arg::B(Wu), mtl::Arg::B(gsh), mtl::Arg::B(hb_), + mtl::Arg::V(gs)}, (NFF + tn - 1) / tn, (T + tm - 1) / tm, sg * 32); + }; + const std::string wd = window_tile(int(T), int(NR), int(KD)), ws = window_tile(int(T), int(NFF), int(N)); + const mtl::Dispatch epi = D(P("gr_epi_down"), {mtl::Arg::B(V), mtl::Arg::B(la), mtl::Arg::B(ia), mtl::Arg::V(ea)}, + (T * NR + 255) / 256, 1, 256); + const uint32_t nh = T * NFF; + const mtl::Dispatch swi = D(P("swiglu_h"), {mtl::Arg::B(gsh), mtl::Arg::B(ush), mtl::Arg::B(ha), mtl::Arg::V(nh)}, + (nh + 255) / 256, 1, 256); + cx.run({gemm(wd, xnh, WD, V, gd), epi}); + cx.run({gemm(ws, mixh, Wg, gsh, gs), gemm(ws, mixh, Wu, ush, gs), swi}); + for (const auto& tile : tiles) { + std::memset(lb.ptr, 0xCD, lb.bytes); + std::memset(ib_.ptr, 0xCD, ib_.bytes); + std::memset(hb_.ptr, 0xCD, hb_.bytes); + cx.run({down(tile)}); + cx.run({swup(tile)}); + const bool okd = same(la, lb, la.bytes) && same(ia, ib_, ia.bytes), oks = same(ha, hb_, ha.bytes); + fails += !okd + !oks; + std::printf("T %4u %-11s gemm_hc_down vs %s + gr_epi_down: bitwise %s; gemm_swiglu_up vs %s + swiglu_h: bitwise %s\n", + T, tile.c_str(), wd.c_str(), okd ? "PASS" : "FAIL", ws.c_str(), oks ? "PASS" : "FAIL"); + } + if (!bench) continue; + // 96 HC reads: [gemm_f16, gr_epi_down] vs [gemm_hc_down]; 48 shared experts: [gate, up (overlapped), swiglu_h] + // vs [gate, gemm_swiglu_up] - each a dependency graph (run_graph) with the window's barriers; every tile both + // ways, so a row's change is the fusion's alone (vs its own tile) and the window's (vs its current tile) + auto rep = [](std::vector one, int n) { + std::vector all; + for (int i = 0; i < n; ++i) all.insert(all.end(), one.begin(), one.end()); + return all; + }; + struct Cand { std::string name, tile; bool fused; std::vector ds; }; + std::vector cand; + for (const auto& tile : tiles) { + cand.push_back({"down", tile, false, rep({gemm(tile, xnh, WD, V, gd), epi}, 96)}); + cand.push_back({"down", tile, true, rep({down(tile)}, 96)}); + } + for (const auto& tile : tiles) { + cand.push_back({"swiglu", tile, false, rep({gemm(ws, mixh, Wg, gsh, gs), gemm(tile, mixh, Wu, ush, gs, false), swi}, 48)}); + cand.push_back({"swiglu", tile, true, rep({gemm(ws, mixh, Wg, gsh, gs), swup(tile)}, 48)}); + } + cand.push_back({"down", wd, false, rep({gemm(wd, xnh, WD, V, gd)}, 96)}); // the GEMM alone: the epilogue's ceiling + cand.push_back({"down", wd, false, rep({gemm(wd, xnh, WD, V, gd)}, 96)}); + for (const auto& c : cand) cx.warm(c.ds, 0.01); + constexpr int NS = 5; + std::vector> ms(cand.size()); + for (int k = 0; k < NS; ++k) + for (size_t j = 0; j < cand.size(); ++j) { + const size_t i = (k % 2) ? cand.size() - 1 - j : j; + ms[i].push_back(cx.run_graph(cand[i].ds) * 1e3); + } + std::printf("T %u, ms per chunk (median of %d, alternating order); window tiles: down %s, swiglu gate/up %s\n", T, NS, + wd.c_str(), ws.c_str()); + std::vector med(cand.size()); + for (size_t i = 0; i < cand.size(); ++i) { + std::sort(ms[i].begin(), ms[i].end()); + med[i] = ms[i][NS / 2]; + } + for (size_t i = 0; i + 2 < cand.size(); i += 2) { + double wbase = 0; + for (size_t j = 0; j + 2 < cand.size(); j += 2) + if (cand[j].name == cand[i].name && cand[j].tile == (cand[i].name == "down" ? wd : ws)) wbase = med[j]; + std::printf(" %-6s %-8s plain %8.3f fused %8.3f (%+5.1f%% vs own tile", cand[i].name.c_str(), cand[i].tile.c_str(), + med[i], med[i + 1], 100.0 * (med[i + 1] / med[i] - 1)); + if (wbase > 0) std::printf(", %+5.1f%% vs window", 100.0 * (med[i + 1] / wbase - 1)); + std::printf(")\n"); + } + std::printf(" down %-8s GEMM alone, no gr_epi_down: %8.3f / %8.3f\n", wd.c_str(), med[cand.size() - 2], med.back()); + } + std::printf("%s\n", fails ? "FAIL" : "ALL PASS"); + return fails ? 1 : 0; +} + +// Same tile and buffers, different dense_mm source revisions; --ab in dense_mm_parity compares tile policies. +int dense_ab_parity(uint32_t T, const char* base_path, const char* only) { + if (!T || T > 2048) return 1; + mtl::Context cx; + int lib[2]; + for (int v = 0; v < 2; ++v) { + const auto src = mtl::read_file(v ? base_path : cx.kernel_dir + "/dense_mm.metal"); + const auto b = src.find("struct GemmArgs"), e = src.find("#define DQ(F)"); + const auto i = src.find("template [[host_name(\"gemm_f16_"); + if (b == std::string::npos || e == std::string::npos || i == std::string::npos || e <= b) { + std::fprintf(stderr, "dense source markers missing\n"); + return 1; + } + lib[v] = cx.compile({}, "#define NO_GGML_COMMON 1\n" + "#include \n" + "using namespace metal; using namespace mpp::tensor_ops;\n" + src.substr(b, e - b) + src.substr(i)); + } + struct Shape { const char* name; int n, k; }; + const Shape shapes[] = {{"hc_read", 324, 10240}, {"Wqkv", 10240, 2560}, {"Wz", 6144, 2560}, + {"Wout", 2560, 6144}, {"Wq", 12288, 2560}, {"Wpv", 2560, 2560}, {"Wshgu", 640, 2560}, + {"Wshd", 2560, 640}, {"W512", 512, 2560}, {"Wk", 1152, 2560}, {"Wk_sel", 1664, 2560}, + {"Walbe", 96, 2560}, {"Wsg", 1, 2560}}; + auto A = cx.alloc(2 * size_t(T) * 10240), B = cx.alloc(2 * size_t(12288) * 2560); + mtl::Buffer Y[2] = {cx.alloc(4 * size_t(T) * 12288), cx.alloc(4 * size_t(T) * 12288)}; + ref::Rng rng(61 + T); + for (const auto* b : {&A, &B}) + for (size_t i = 0; i < b->bytes / 2; ++i) b->as()[i] = ref::f32_to_f16(rng.sym() * 0.1f); + int fails = 0, checked = 0; + for (const auto& s : shapes) { + if (only && std::strcmp(only, s.name) != 0) continue; + ++checked; + const auto tile = window_tile(T, s.n, s.k); + int tm, tn, sg; + std::sscanf(tile.c_str(), "%d_%d_%d", &tm, &tn, &sg); + mtl::Pipeline p[2] = {cx.pipeline(lib[0], "gemm_f16_" + tile), cx.pipeline(lib[1], "gemm_f16_" + tile)}; + const GemmArgs g{int(T), s.n, s.k, 0}; + mtl::Dispatch d[2][2]; + for (int v = 0; v < 2; ++v) + for (int b = 0; b < 2; ++b) { + auto& x = d[v][b]; + x.p = &p[v]; + x.args = {mtl::Arg::B(A), mtl::Arg::B(B), mtl::Arg::B(Y[b]), mtl::Arg::V(g)}; + x.grid[0] = (s.n + tn - 1) / tn; + x.grid[1] = (T + tm - 1) / tm; + x.tg[0] = sg * 32; + } + for (int v = 0; v < 2; ++v) { + std::memset(Y[v].ptr, 0xCD, Y[v].bytes); + cx.run({d[v][v]}); + } + const bool ok = same(Y[0], Y[1], 4 * size_t(T) * s.n); + fails += !ok; + for (int v = 0; v < 2; ++v) cx.warm({d[v][v]}, 0.005); + std::vector ratios; + for (int r = 0; r < 12; ++r) { + double ms[2]; + for (int j = 0; j < 2; ++j) { + const int v = (r / 2 + j) % 2; + ms[v] = cx.run({d[v][(r + v) % 2]}, 8); + } + ratios.push_back(ms[0] / ms[1]); + } + std::sort(ratios.begin(), ratios.end()); + std::printf("M %u %-8s %s bitwise %s ratio %.4f [%.4f, %.4f]\n", T, s.name, tile.c_str(), + ok ? "PASS" : "FAIL", (ratios[5] + ratios[6]) * 0.5, ratios[2], ratios[9]); + } + std::printf("%s\n", fails || !checked ? "FAIL" : "ALL PASS"); + return fails || !checked ? 1 : 0; +} + +// Round 29's shared-expert gate + up fusions, measured here and not shipped (audits/opus_r29_prompt.md): compute-bound, +// they tie gate + gemm_swiglu_up past 160 rows outside a few tile-aligned T +const char* kSwigluFused = R"MSL( +// gemm_swiglu_gu: the shared expert's gate and up GEMMs in one kernel, h = half(silu(x . w_gate^T) * (x . w_up^T)), +// both over the same A tile; g and u are never written. h bitwise gemm_f16 (gate, up) + swiglu_h. Two accumulators +// (silu(g) waits in the first). M5 layout only (as gemm_hc_down); N % 4 == 0. grid (ceil(N / TN), ceil(M / TM)), +// 32 * SG threads +template +kernel void gemm_swiglu_gu(device half* A [[buffer(0)]], device half* Bg [[buffer(1)]], device half* Bu [[buffer(2)]], + device half* h [[buffer(3)]], constant GemmArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]]) { + tensor, tensor_inline> tA(A, dextents(a.K, a.M)); + tensor, tensor_inline> tBg(Bg, dextents(a.K, a.N)); + tensor, tensor_inline> tBu(Bu, dextents(a.K, a.N)); + constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v, false, true, false); + matmul2d> op; + auto mA = tA.slice(0, int(tg.y) * TM); + auto mBg = tBg.slice(0, int(tg.x) * TN); + auto mBu = tBu.slice(0, int(tg.x) * TN); + auto sg = op.template get_destination_cooperative_tensor(); + auto acc = op.template get_destination_cooperative_tensor(); + op.run(mA, mBg, sg); + #pragma clang loop unroll(full) + for (uint16_t i = 0; i < sg.get_capacity(); ++i) sg[i] = sg[i] / (1.0f + exp(-sg[i])); + op.run(mA, mBu, acc); + #pragma clang loop unroll(full) + for (uint16_t i = 0; i < acc.get_capacity(); i += 4) { + const auto ix = acc.get_multidimensional_index(i); + const int c = int(tg.x) * TN + ix[0], t = int(tg.y) * TM + ix[1]; + if (t >= a.M || c >= a.N) continue; + half4 o; + #pragma clang loop unroll(full) + for (int k = 0; k < 4; ++k) o[k] = half(sg[i + k] * acc[i + k]); + *(device half4*)(h + ulong(t) * uint(a.N) + uint(c)) = o; + } +} + +// gemm_swiglu_il: gate and up as one GEMM over B = w_gate and w_up rows interleaved (row 2j = gate j, 2j + 1 = up j; +// dq_f16 perm 2 over the pair), one accumulator; on M5 a lane's four columns are (g, u) of two outputs. +// grid (ceil(2N / TN), ceil(M / TM)), a.N = N +template +kernel void gemm_swiglu_il(device half* A [[buffer(0)]], device half* B [[buffer(1)]], device half* h [[buffer(2)]], + constant GemmArgs& a [[buffer(3)]], uint3 tg [[threadgroup_position_in_grid]]) { + tensor, tensor_inline> tA(A, dextents(a.K, a.M)); + tensor, tensor_inline> tB(B, dextents(a.K, 2 * a.N)); + constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v, false, true, false); + matmul2d> op; + auto mA = tA.slice(0, int(tg.y) * TM); + auto mB = tB.slice(0, int(tg.x) * TN); + auto acc = op.template get_destination_cooperative_tensor(); + op.run(mA, mB, acc); + #pragma clang loop unroll(full) + for (uint16_t i = 0; i < acc.get_capacity(); i += 4) { + const auto ix = acc.get_multidimensional_index(i); + const int c = (int(tg.x) * TN + ix[0]) / 2, t = int(tg.y) * TM + ix[1]; + if (t >= a.M || c >= a.N) continue; + const half2 o(half(acc[i] / (1.0f + exp(-acc[i])) * acc[i + 1]), half(acc[i + 2] / (1.0f + exp(-acc[i + 2])) * acc[i + 3])); + *(device half2*)(h + ulong(t) * uint(a.N) + uint(c)) = o; + } +} +)MSL"; + +// `comb_fold_parity T gu [bench [tile,tile...]]`: gemm_swiglu_gu / gemm_swiglu_il (gate + up + swiglu_h in one kernel) +// against gemm_f16 (gate) + gemm_f16 (up) + swiglu_h and the window's gate + gemm_swiglu_up, BITWISE; T 0 = 1, 7, 161, +// 256, 512, 1024, 1983, 2048; every tile (TM_TN_SG, or ilTM_TN_SG). bench: paired GPU times of 48 shared experts (a chunk's): the window's old +// paths (gate ∥ up + swiglu_h, gate + gemm_swiglu_up) vs each tile, alternating order, median of 9 +int gu_parity(uint32_t TB, bool bench, const char* only) { + mtl::Context cx; + const std::string dm = mtl::read_file(cx.kernel_dir + "/dense_mm.metal"), pr = mtl::read_file(cx.kernel_dir + "/prompt.metal"); + const auto gb = dm.find("template "), ge = dm.find("#define DQ(F)"); + const auto ib = dm.find("template [[host_name(\"gemm_f16_"); + const auto eb = pr.find("struct EpiArgs {"), ee = pr.find("// ---- the HC read"), hb = pr.find("// gemm_hc_up:"); + if (gb == std::string::npos || ge == std::string::npos || ib == std::string::npos || eb == std::string::npos || + ee == std::string::npos || hb == std::string::npos) { + std::fprintf(stderr, "gu source markers missing\n"); + return 1; + } + if (!coop_k4(cx)) { + std::printf("device %s: not the M5 layout, the fused gate + up is M5 only: SKIP\n", cx.device_name().c_str()); + return 0; + } + std::vector tiles = {"32_64_2", "64_64_4", "32_128_4", "64_128_4", "128_64_4", "64_128_8", "128_128_8"}; + if (only) { + tiles.clear(); + for (const char *s = only, *e; *s; s = *e ? e + 1 : e) { + e = std::strchr(s, ','); + if (!e) e = s + std::strlen(s); + tiles.emplace_back(s, e); + } + } + // a tile TM_TN_SG runs gemm_swiglu_gu, ilTM_TN_SG gemm_swiglu_il (its B the interleaved gate and up) + auto parse = [](const std::string& v, int& tm, int& tn, int& sg) { + const bool il = v.rfind("il", 0) == 0; + std::sscanf(v.c_str() + (il ? 2 : 0), "%d_%d_%d", &tm, &tn, &sg); + return il; + }; + auto kname = [](const std::string& v) { return "gemm_swiglu_" + (v.rfind("il", 0) == 0 ? v : "gu_" + v); }; + std::string extra; + for (const auto& v : tiles) { + int tm = 0, tn = 0, sg = 0; + const std::string k = parse(v, tm, tn, sg) ? "gemm_swiglu_il" : "gemm_swiglu_gu"; + const std::string targs = std::to_string(tm) + ", " + std::to_string(tn) + ", " + std::to_string(sg); + extra += "template [[host_name(\"" + kname(v) + "\")]] kernel decltype(" + k + "<" + targs + ">) " + k + "<" + targs + ">;\n"; + } + const std::string prelude = "#define NO_GGML_COMMON 1\n" + "#include \n" + "using namespace metal; using namespace mpp::tensor_ops;\n" + "struct GemmArgs { int M, N, K; int pad; };\n"; + const int lib = cx.compile({}, prelude + dm.substr(gb, ge - gb) + dm.substr(ib) + pr.substr(eb, ee - eb) + pr.substr(hb) + + kSwigluFused + extra); + constexpr uint32_t N = 2560, NFF = 640; + std::printf("device %s\n", cx.device_name().c_str()); + std::map ps; + auto P = [&](const std::string& n) -> const mtl::Pipeline& { + auto it = ps.find(n); + if (it == ps.end()) it = ps.emplace(n, cx.pipeline(lib, n)).first; + return it->second; + }; + auto D = [](const mtl::Pipeline& p, std::vector a, uint32_t gx, uint32_t gy, uint32_t th, bool barrier = true) { + mtl::Dispatch d; + d.p = &p; + d.args = std::move(a); + d.grid[0] = gx; + d.grid[1] = gy; + d.tg[0] = th; + d.barrier = barrier; + return d; + }; + int fails = 0; + const std::vector Ts = TB ? std::vector{TB} : + std::vector{1u, 7u, 161u, 256u, 512u, 1024u, 1983u, 2048u}; + for (uint32_t T : Ts) { + mtl::Buffer mixh = cx.alloc(2 * size_t(T) * N), Wg = cx.alloc(2 * size_t(NFF) * N), Wu = cx.alloc(2 * size_t(NFF) * N), + gsh = cx.alloc(4 * size_t(T) * NFF), ush = cx.alloc(4 * size_t(T) * NFF), ha = cx.alloc(2 * size_t(T) * NFF), + hw = cx.alloc(2 * size_t(T) * NFF), hb_ = cx.alloc(2 * size_t(T) * NFF); + ref::Rng r(37 + T); + for (const auto& [b, s] : {std::pair{&mixh, 1.0f}, {&Wg, 0.05f}, {&Wu, 0.05f}}) + for (size_t i = 0; i < b->bytes / 2; ++i) b->as()[i] = ref::f32_to_f16(s * r.sym()); + const GemmArgs gs{int(T), int(NFF), int(N), 0}; + mtl::Buffer Wgu = cx.alloc(4 * size_t(NFF) * N); // gate and up rows interleaved (row 2j gate j, 2j + 1 up j) + for (uint32_t j = 0; j < NFF; ++j) { + std::memcpy(Wgu.as() + size_t(2 * j) * N, Wg.as() + size_t(j) * N, 2 * N); + std::memcpy(Wgu.as() + size_t(2 * j + 1) * N, Wu.as() + size_t(j) * N, 2 * N); + } + const std::string ws = window_tile(int(T), int(NFF), int(N)); + int wm, wn, wsg; + std::sscanf(ws.c_str(), "%d_%d_%d", &wm, &wn, &wsg); + auto gemm = [&](const mtl::Buffer& B, const mtl::Buffer& Y, bool barrier = true) { + return D(P("gemm_f16_" + ws), {mtl::Arg::B(mixh), mtl::Arg::B(B), mtl::Arg::B(Y), mtl::Arg::V(gs)}, + (NFF + wn - 1) / wn, (T + wm - 1) / wm, wsg * 32, barrier); + }; + // the window past 160 rows: gate, then gemm_swiglu_up at its tile + const bool small = T < 384; + const mtl::Dispatch swup = D(P(small ? "gemm_swiglu_up_32_64_2" : "gemm_swiglu_up_64_128_4"), + {mtl::Arg::B(mixh), mtl::Arg::B(Wu), mtl::Arg::B(gsh), mtl::Arg::B(hw), mtl::Arg::V(gs)}, + (NFF + (small ? 64 : 128) - 1) / (small ? 64 : 128), (T + (small ? 32 : 64) - 1) / (small ? 32 : 64), small ? 64 : 128); + const uint32_t nh = T * NFF; + const mtl::Dispatch swi = D(P("swiglu_h"), {mtl::Arg::B(gsh), mtl::Arg::B(ush), mtl::Arg::B(ha), mtl::Arg::V(nh)}, + (nh + 255) / 256, 1, 256); + cx.run({gemm(Wg, gsh), gemm(Wu, ush), swi}); + std::memset(hw.ptr, 0xCD, hw.bytes); + cx.run({swup}); + const bool okw = same(ha, hw, ha.bytes); + fails += !okw; + std::printf("T %4u window gate %s + gemm_swiglu_up vs gate + up + swiglu_h: bitwise %s\n", T, ws.c_str(), okw ? "PASS" : "FAIL"); + auto gu = [&](const std::string& v, const mtl::Buffer& h) { + int tm = 0, tn = 0, sg = 0; + if (parse(v, tm, tn, sg)) + return D(P(kname(v)), {mtl::Arg::B(mixh), mtl::Arg::B(Wgu), mtl::Arg::B(h), mtl::Arg::V(gs)}, + (2 * NFF + tn - 1) / tn, (T + tm - 1) / tm, sg * 32); + return D(P(kname(v)), {mtl::Arg::B(mixh), mtl::Arg::B(Wg), mtl::Arg::B(Wu), mtl::Arg::B(h), mtl::Arg::V(gs)}, + (NFF + tn - 1) / tn, (T + tm - 1) / tm, sg * 32); + }; + for (const auto& v : tiles) { + std::memset(hb_.ptr, 0xCD, hb_.bytes); + cx.run({gu(v, hb_)}); + const bool ok = same(ha, hb_, ha.bytes); + fails += !ok; + std::printf("T %4u %-26s vs gate + up + swiglu_h: bitwise %s\n", T, kname(v).c_str(), ok ? "PASS" : "FAIL"); + } + if (!bench) continue; + auto rep = [](std::vector one, int n) { + std::vector all; + for (int i = 0; i < n; ++i) all.insert(all.end(), one.begin(), one.end()); + return all; + }; + // [gate, up overlapped, swiglu_h] (the window at <= 160 rows), [gate, gemm_swiglu_up] (past 160), each tile + std::vector> cand = {rep({gemm(Wg, gsh), gemm(Wu, ush, false), swi}, 48), + rep({gemm(Wg, gsh), swup}, 48)}; + for (const auto& v : tiles) cand.push_back(rep({gu(v, hb_)}, 48)); + for (const auto& c : cand) cx.warm(c, 0.01); + constexpr int NS = 9; + std::vector> ms(cand.size()); + for (int k = 0; k < NS; ++k) + for (size_t j = 0; j < cand.size(); ++j) { + const size_t i = (k % 2) ? cand.size() - 1 - j : j; + ms[i].push_back(cx.run_graph(cand[i]) * 1e3); + } + std::vector med(cand.size()); + for (size_t i = 0; i < cand.size(); ++i) { + std::sort(ms[i].begin(), ms[i].end()); + med[i] = ms[i][NS / 2]; + } + std::printf("T %u, ms per chunk (48 shared experts, median of %d, alternating order)\n" + " gate ∥ up %s + swiglu_h %8.3f\n gate %s + gemm_swiglu_up_%-8s %8.3f\n", + T, NS, ws.c_str(), med[0], ws.c_str(), small ? "32_64_2" : "64_128_4", med[1]); + for (size_t i = 0; i < tiles.size(); ++i) + std::printf(" %-26s %8.3f (%+5.1f%% vs gate ∥ up, %+5.1f%% vs gemm_swiglu_up)\n", kname(tiles[i]).c_str(), + med[i + 2], 100.0 * (med[i + 2] / med[0] - 1), 100.0 * (med[i + 2] / med[1] - 1)); + } + std::printf("%s\n", fails ? "FAIL" : "ALL PASS"); + return fails ? 1 : 0; +} + +// `comb_fold_parity T topk baseline-prompt.metal`: this prompt.metal's route_topk against a baseline copy's, BITWISE +// (ids and w) for prompt_moce 0, 320335 (the streamed default: tail 32, protect 3, bonus 35) and -16.3 / -0.05 (skip), +// with paired GPU times: 48 reps a sample, 11 samples in alternating order, median of the per-sample ratios. T 0 = +// 512, 1024, 2048. Logits are rounded to 1/16 so ties are common; half the experts are not resident in smap. +int topk_parity(uint32_t TB, const char* base_path) { + mtl::Context cx; + int lib[2]; + for (int v = 0; v < 2; ++v) { + const std::string src = mtl::read_file(v ? std::string(base_path) : cx.kernel_dir + "/prompt.metal"); + const auto e = src.find("struct EpiArgs {"); + if (e == std::string::npos) { + std::fprintf(stderr, "route_topk source marker missing\n"); + return 1; + } + lib[v] = cx.compile({}, "using namespace metal;\n" + src.substr(0, e)); + } + const mtl::Pipeline pn = cx.pipeline(lib[0], "route_topk"), pb = cx.pipeline(lib[1], "route_topk"); + std::printf("device %s\n", cx.device_name().c_str()); + constexpr uint32_t K = 10; + int fails = 0; + std::vector> cases; // (T, n_exp) + for (uint32_t T : TB ? std::vector{TB} : std::vector{512u, 1024u, 2048u}) + for (uint32_t NE : {512u, 256u}) cases.emplace_back(T, NE); + for (const auto& [T, NE] : cases) { + auto lg = cx.alloc(4 * size_t(T) * NE), sm = cx.alloc(4 * NE); + mtl::Buffer ids[2] = {cx.alloc(4 * T * K), cx.alloc(4 * T * K)}, w[2] = {cx.alloc(4 * T * K), cx.alloc(4 * T * K)}; + ref::Rng r(41 + T); + for (size_t i = 0; i < size_t(T) * NE; ++i) lg.as()[i] = std::round(16.0f * 4.0f * r.sym()) / 16.0f; + for (uint32_t i = 0; i < NE; ++i) sm.as()[i] = r.sym() > 0.0f ? i : 0xFFFFFFFFu; + const uint32_t nt = T | (NE << 16); + auto D = [&](const mtl::Pipeline& p, int v, float moce) { + mtl::Dispatch d; + d.p = &p; + d.args = {mtl::Arg::B(lg), mtl::Arg::B(ids[v]), mtl::Arg::B(w[v]), mtl::Arg::V(nt), mtl::Arg::B(sm), mtl::Arg::V(moce)}; + d.grid[0] = (T + 7) / 8; + d.tg[0] = 256; + return d; + }; + for (const float moce : {0.0f, 320335.0f, -16.3f, -0.05f}) { + for (int v = 0; v < 2; ++v) { + std::memset(ids[v].ptr, 0xCD, ids[v].bytes); + std::memset(w[v].ptr, 0xCD, w[v].bytes); + cx.run({D(v ? pb : pn, v, moce)}); + } + const bool ids_ok = same(ids[0], ids[1], ids[0].bytes); + const bool w_bit = same(w[0], w[1], w[0].bytes); + double max_dw = 0.0; + for (size_t i = 0; i < size_t(T) * K; ++i) + max_dw = std::max(max_dw, (double) std::fabs(w[0].as()[i] - w[1].as()[i])); + const bool ok = ids_ok && max_dw < 1e-6; + fails += !ok; + std::printf("T %4u n_exp %u route_topk moce %9.2f: ids bitwise %s, w %s (max|dw|=%.3g) -> %s\n", + T, NE, moce, ids_ok ? "PASS" : "FAIL", w_bit ? "bitwise" : "close", max_dw, ok ? "PASS" : "FAIL"); + } + for (const float moce : {0.0f, 320335.0f}) { + if (NE != 512) break; // the times: the model's 512 experts only + std::vector dn(48, D(pn, 0, moce)), db(48, D(pb, 1, moce)); + cx.warm(dn, 0.01); + cx.warm(db, 0.01); + std::vector ratio, tn, tb; + for (int k = 0; k < 11; ++k) { + double a, b; + if (k % 2) { b = cx.run(db); a = cx.run(dn); } else { a = cx.run(dn); b = cx.run(db); } + tn.push_back(a * 1e3); + tb.push_back(b * 1e3); + ratio.push_back(a / b); + } + for (auto* v : {&ratio, &tn, &tb}) std::sort(v->begin(), v->end()); + std::printf("T %4u route_topk moce %9.2f x48: new %.3f ms, baseline %.3f ms, ratio %.3f (median of 11)\n", T, moce, + tn[5], tb[5], ratio[5]); + } + } + std::printf("%s\n", fails ? "FAIL" : "ALL PASS"); + return fails ? 1 : 0; +} + +// `comb_fold_parity T ab baseline-prompt.metal [gr|gemm]`: this prompt.metal's GR reads and register-epilogue GEMMs +// against a baseline copy's, BITWISE (outputs of the two builds), with paired GPU times: 8 reps a sample, 11 samples in +// alternating order, median of the per-sample ratios. T 0 = 512, 1024, 2048; the window's tiles at each T +int ab_parity(uint32_t TB, const char* base_path, const char* only) { + mtl::Context cx; + const std::string gr_pre = "#define NO_GGML_COMMON 1\nusing namespace metal;\n" + "struct GrArgs { uint n_embd, hc, hc_lr, n_tok, apply; float eps; };\n"; + const std::string mm_pre = "#define NO_GGML_COMMON 1\n" + "#include \n" + "using namespace metal; using namespace mpp::tensor_ops;\n" + "struct GemmArgs { int M, N, K, pad; };\n" + "struct EpiArgs { uint n_tok, n_embd, hc, hc_lr; };\n"; + const bool do_gr = !only || std::strcmp(only, "gr") == 0, do_mm = !only || std::strcmp(only, "gemm") == 0; + int gl[2], ml[2]; + for (int v = 0; v < 2; ++v) { + const std::string src = mtl::read_file(v ? std::string(base_path) : cx.kernel_dir + "/prompt.metal"); + const auto b = src.find("// ---- the HC read"), e = src.find("// gemm_hc_up:"); + if (b == std::string::npos || e == std::string::npos || e <= b) { + std::fprintf(stderr, "source markers missing\n"); + return 1; + } + gl[v] = do_gr ? cx.compile({}, gr_pre + src.substr(b, e - b)) : -1; + ml[v] = do_mm ? cx.compile({}, mm_pre + src.substr(e)) : -1; + } + if (do_mm && !coop_k4(cx)) { + std::printf("device %s: not the M5 layout: SKIP\n", cx.device_name().c_str()); + return 0; + } + std::printf("device %s, baseline %s\n", cx.device_name().c_str(), base_path); + constexpr uint32_t N = 2560, HC = 4, LR = 320, NR = LR + HC, NFF = 640, TOPK = 10, NEXP = 512, TMR = 16; + constexpr size_t RN = size_t(HC) * N; + std::map ps; + auto P = [&](int lib, const std::string& n) -> const mtl::Pipeline& { + const std::string key = std::to_string(lib) + n; + auto it = ps.find(key); + if (it == ps.end()) it = ps.emplace(key, cx.pipeline(lib, n)).first; + return it->second; + }; + int fails = 0; + for (uint32_t T : TB ? std::vector{TB} : std::vector{512u, 1024u, 2048u}) { + const size_t mr = size_t(T) * TOPK + NEXP * TMR; + // per build v (0 this, 1 baseline): its own outputs (and R, updated in place) + mtl::Buffer R0 = cx.alloc(4 * T * RN), R[2] = {cx.alloc(4 * T * RN), cx.alloc(4 * T * RN)}, + xn[2] = {cx.alloc(2 * T * RN), cx.alloc(2 * T * RN)}, mh[2] = {cx.alloc(2 * T * N), cx.alloc(2 * T * N)}, + lo[2] = {cx.alloc(2 * T * LR), cx.alloc(2 * T * LR)}, ij[2] = {cx.alloc(4 * T * HC), cx.alloc(4 * T * HC)}, + hh[2] = {cx.alloc(2 * T * NFF), cx.alloc(2 * T * NFF)}; + mtl::Buffer YE = cx.alloc(do_gr ? 4 * mr * N : 4), pos = cx.alloc(4 * T * TOPK), sw = cx.alloc(4 * mr), + ys = cx.alloc(4 * T * N), sgv = cx.alloc(4 * T), inj = cx.alloc(4 * T * HC), wn = cx.alloc(4 * RN), + xin = cx.alloc(2 * T * RN), loh = cx.alloc(2 * T * LR), WU = cx.alloc(2 * RN * LR), + WD = cx.alloc(2 * NR * RN), mixh = cx.alloc(2 * T * N), Wu = cx.alloc(2 * NFF * N), gsh = cx.alloc(4 * T * NFF), + nil{}; + ref::Rng r(23 + T); + auto fill = [&](const mtl::Buffer& b, float s, float o) { + for (size_t i = 0; i < b.bytes / 4; ++i) b.as()[i] = o + s * r.sym(); + }; + auto fillh = [&](const mtl::Buffer& b, float s) { + for (size_t i = 0; i < b.bytes / 2; ++i) b.as()[i] = ref::f32_to_f16(s * r.sym()); + }; + fill(R0, 1.0f, 0.0f); fill(YE, 0.5f, 0.0f); fill(sw, 0.1f, 0.1f); fill(ys, 0.5f, 0.0f); fill(sgv, 3.0f, 0.0f); + fill(inj, 4.0f, 0.0f); fill(wn, 0.2f, 1.0f); fill(gsh, 4.0f, 0.0f); fillh(xin, 1.0f); fillh(loh, 1.0f); + fillh(WU, 0.1f); fillh(WD, 0.02f); fillh(mixh, 1.0f); fillh(Wu, 0.05f); + std::vector perm(mr); + std::iota(perm.begin(), perm.end(), 0); + std::shuffle(perm.begin(), perm.end(), std::mt19937(T)); + for (size_t i = 0; i < size_t(T) * TOPK; ++i) pos.as()[i] = (i % 37 == 5) ? -1 : perm[i]; + const GrArgs a{N, HC, LR, T, 1u, 1e-6f}, a0{N, HC, LR, T, 0u, 1e-6f}; + const EpiArgs ea{T, N, HC, LR}; + const GemmArgs gu{int(T), int(RN), int(LR), 0}, gs{int(T), int(NFF), int(N), 0}; + auto D = [](const mtl::Pipeline& p, std::vector args, uint32_t gx, uint32_t gy, uint32_t th) { + mtl::Dispatch d; + d.p = &p; + d.args = std::move(args); + d.grid[0] = gx; + d.grid[1] = gy; + d.tg[0] = th; + return d; + }; + // the window's tiles (window.cpp prompt_chunk) + const auto [dm, dn, dsg] = T <= 192 ? std::array{16, 64, 2} : T <= 384 ? std::array{32, 64, 2} + : T <= 512 ? std::array{32, 128, 4} : T < 1536 ? std::array{64, 128, 4} : std::array{32, 128, 2}; + const std::string down = "gemm_hc_down_" + std::to_string(dm) + "_" + std::to_string(dn) + "_" + std::to_string(dsg); + const std::string swup = T < 384 ? "gemm_swiglu_up_32_64_2" : "gemm_swiglu_up_64_128_4"; + const int sm = T < 384 ? 32 : 64, sn = T < 384 ? 64 : 128, ssg = T < 384 ? 2 : 4; + // build v's dispatches on buffer set b; sample s times build v on set (s + v) % 2, so placement is not the build's + struct Case { std::string name; mtl::Dispatch d[2][2]; std::vector out[2]; }; + std::vector cs; + for (int v = 0; v < 2; ++v) + for (int b = 0; b < 2; ++b) { + std::vector> ds; + std::vector> outs; + if (do_gr) { + const int l = gl[v]; + auto nh4 = [&](const GrArgs& g) { + return D(P(l, "gr_norm_h4"), {mtl::Arg::B(R[b]), mtl::Arg::B(ys), mtl::Arg::B(inj), mtl::Arg::B(wn), + mtl::Arg::B(xn[b]), mtl::Arg::V(g)}, 1, T, 256); + }; + ds.push_back({"gr_norm_h_moec", D(P(l, "gr_norm_h_moec"), {mtl::Arg::B(R[b]), mtl::Arg::B(YE), mtl::Arg::B(pos), + mtl::Arg::B(sw), mtl::Arg::B(ys), mtl::Arg::B(sgv), mtl::Arg::B(inj), mtl::Arg::B(wn), mtl::Arg::B(xn[b]), + mtl::Arg::V(a), mtl::Arg::V(TOPK)}, 1, T, 256)}); + ds.push_back({"gr_apply_moec", D(P(l, "gr_apply_moec"), {mtl::Arg::B(R[b]), mtl::Arg::B(YE), mtl::Arg::B(pos), + mtl::Arg::B(sw), mtl::Arg::B(ys), mtl::Arg::B(sgv), mtl::Arg::B(inj), mtl::Arg::V(a), mtl::Arg::V(TOPK)}, + (N + 255) / 256, T, 256)}); + ds.push_back({"gr_norm_h4 apply 1", nh4(a)}); + ds.push_back({"gr_norm_h4 apply 0", nh4(a0)}); + outs.insert(outs.end(), {{&R[b], &xn[b]}, {&R[b]}, {&R[b], &xn[b]}, {&xn[b]}}); + } + if (do_mm) { + const int l = ml[v]; + ds.push_back({down, D(P(l, down), {mtl::Arg::B(xin), mtl::Arg::B(WD), mtl::Arg::B(lo[b]), mtl::Arg::B(ij[b]), + mtl::Arg::V(ea)}, (NR + dn - 1) / dn, (T + dm - 1) / dm, dsg * 32)}); + ds.push_back({"gemm_hc_up_reg_32_128_4", D(P(l, "gemm_hc_up_reg_32_128_4"), {mtl::Arg::B(loh), mtl::Arg::B(WU), + mtl::Arg::B(xin), mtl::Arg::B(nil), mtl::Arg::B(mh[b]), mtl::Arg::V(gu), mtl::Arg::V(ea)}, + RN / 128, (T + 31) / 32, 128)}); + ds.push_back({swup, D(P(l, swup), {mtl::Arg::B(mixh), mtl::Arg::B(Wu), mtl::Arg::B(gsh), mtl::Arg::B(hh[b]), + mtl::Arg::V(gs)}, (NFF + sn - 1) / sn, (T + sm - 1) / sm, ssg * 32)}); + outs.insert(outs.end(), {{&lo[b], &ij[b]}, {&mh[b]}, {&hh[b]}}); + } + if (cs.empty()) cs.resize(ds.size()); + for (size_t i = 0; i < ds.size(); ++i) { + cs[i].name = ds[i].first; + cs[i].d[v][b] = ds[i].second; + if (b == v) cs[i].out[v] = outs[i]; + } + } + for (auto& c : cs) { + for (int v = 0; v < 2; ++v) { + std::memcpy(R[v].ptr, R0.ptr, R0.bytes); + for (const auto* b : c.out[v]) + if (b != &R[v]) std::memset(b->ptr, 0xCD, b->bytes); + cx.run({c.d[v][v]}); + } + bool ok = true; + for (size_t i = 0; i < c.out[0].size(); ++i) ok = ok && same(*c.out[0][i], *c.out[1][i], c.out[0][i]->bytes); + fails += !ok; + for (int v = 0; v < 2; ++v) cx.warm({c.d[v][v]}, 0.02); + constexpr int NS = 12, REPS = 8; + std::vector t[2], ratio; + for (int s = 0; s < NS; ++s) { + double ms[2]; + for (int j = 0; j < 2; ++j) { + const int v = (s / 2 + j) % 2; + ms[v] = cx.run({c.d[v][(s + v) % 2]}, REPS) * 1e3 / REPS; + t[v].push_back(ms[v]); + } + ratio.push_back(ms[0] / ms[1]); + } + for (auto* x : {&t[0], &t[1], &ratio}) std::sort(x->begin(), x->end()); + const auto med = [](const std::vector& x) { return 0.5 * (x[x.size() / 2 - 1] + x[x.size() / 2]); }; + std::printf("T %4u %-24s bitwise %s this %7.4f ms baseline %7.4f ms ratio %.4f [%.4f, %.4f]\n", T, + c.name.c_str(), ok ? "PASS" : "FAIL", med(t[0]), med(t[1]), med(ratio), ratio[2], ratio[NS - 3]); + } + } + std::printf("%s\n", fails ? "FAIL" : "ALL PASS"); + return fails ? 1 : 0; +} +// `comb_fold_parity 0 mtp [bench] [baseline-mtp.metal]`: window.cpp mtp_encode's hidden branch (mtp_rms_xprep, the +// Q8_0 2560 -> 2560 fc_hidden GEMVs, mtp_combine_gnorm) with the n * HC streams batched as rows [t][c], MAXT per GEMV, +// against the per-stream chain it replaces (one GEMV of n rows per stream at the window's tile for n rows, M5 Q8_0): +// R and xn BITWISE wherever both pick one tile (n = 1, n >= 4), and every n BITWISE against n single-token chains. +// baseline: the per-stream chain on that mtp.metal (its [c][t] h layout) == the per-stream chain on this one. +// bench: paired GPU times per n, the weights walking 16 copies (from DRAM, as in a decode), median of 9. +int mtp_parity(bool bench, const char* baseline) { + struct RmsXpArgs { uint32_t n, xp_n, plane_tok, split, keep; float eps; }; + struct XpArgs { uint32_t n, n_tok, plane_tok, stride; }; + struct CombNormArgs { uint32_t n, hc; float eps; }; + struct OldCombNormArgs { uint32_t n, hc, h_stride; float eps; }; + struct GemvArgs { uint32_t n_in, n_out, n_tok, row_bytes, x_stride, y_stride; }; + constexpr int N = 2560, HC = 4, MAXT = 8; + constexpr float eps = 1e-6f; + mtl::Context cx; + const std::string msl = cx.kernel_dir + "/../src/engine/msl/mtp.metal"; + const int lg = cx.compile({"quant.metal", "units.metal", "gemv.metal"}); + const int lm = cx.compile({msl}, "#define NO_GGML_COMMON 1\n"); + const int lb = baseline ? cx.compile({baseline}, "#define NO_GGML_COMMON 1\n") : -1; + std::printf("device %s\n", cx.device_name().c_str()); + const size_t rb = ref::row_bytes(ref::Q8_0, N), mb = rb * N; + const int copies = bench ? 16 : 1; + auto W = cx.alloc(mb * copies); + ref::random_blocks(ref::Q8_0, W.as(), N, N, 41); + for (int k = 1; k < copies; ++k) std::memcpy(W.as() + k * mb, W.ptr, mb); + const size_t RN = (size_t) HC * N, plane = 4ull * MAXT * N; + auto src = cx.alloc(4 * MAXT * RN), wn = cx.alloc(4 * RN), whc = cx.alloc(4 * RN), e = cx.alloc(4ull * MAXT * N); + auto mhn = cx.alloc(4 * MAXT * RN), xp = cx.alloc(4 * plane + 4ull * MAXT * N / 16), h = cx.alloc(4 * MAXT * RN); + mtl::Buffer R[3], xn[3]; // per-stream chain, batched, n single-token chains + for (int i = 0; i < 3; ++i) { R[i] = cx.alloc(4 * MAXT * RN); xn[i] = cx.alloc(4 * MAXT * RN); } + ref::Rng rng(77); + for (size_t i = 0; i < MAXT * RN; ++i) src.as()[i] = rng.sym() * 3.0f; + for (size_t i = 0; i < RN; ++i) { wn.as()[i] = 1.0f + 0.2f * rng.sym(); whc.as()[i] = 1.0f + 0.2f * rng.sym(); } + for (size_t i = 0; i < (size_t) MAXT * N; ++i) e.as()[i] = rng.sym(); + mtl::Pipeline prms = cx.pipeline(lm, "mtp_rms_xprep"), pxp = cx.pipeline(lm, "mtp_xprep"), + pcomb = cx.pipeline(lm, "mtp_combine_gnorm"); + mtl::Pipeline brms, bxp, bcomb; + if (lb >= 0) { + brms = cx.pipeline(lb, "mtp_rms_xprep"); bxp = cx.pipeline(lb, "mtp_xprep"); bcomb = cx.pipeline(lb, "mtp_combine_gnorm"); + } + // window.cpp tile() for a wide Q8_0 tensor on the M5 + auto tile = [](int rows) { + return rows == 1 ? std::array{32, 1, 4, 4} : rows <= 3 ? std::array{32, 2, 4, 1} + : std::array{16, 1, 4, 1}; + }; + std::map gp; + auto gemv = [&](int rows, uint32_t y_stride, size_t y_off, int copy) { + const auto [L, Rr, S, K] = tile(rows); + const std::string fn = "gemv_" + std::to_string(ref::Q8_0) + "_" + std::to_string(L) + "_" + std::to_string(Rr) + + "_" + std::to_string(S) + "_" + std::to_string(K); + const std::string key = fn + "#" + std::to_string(rows); + if (!gp.count(key)) gp.emplace(key, cx.pipeline(lg, fn, mtl::Constants().i(0, rows))); + mtl::Dispatch d; + d.p = &gp.at(key); + const GemvArgs a{N, N, (uint32_t) rows, (uint32_t) rb, N, y_stride}; + d.args = {mtl::Arg::B(W, copy * mb), mtl::Arg::B(xp), mtl::Arg::B(xp, plane * 4), mtl::Arg::B(h, y_off), mtl::Arg::V(a)}; + const int rpt = (S / K) * (32 / L) * Rr; + d.grid[0] = (N + rpt - 1) / rpt; d.tg[0] = S * 32; + return d; + }; + auto rms = [&](const mtl::Pipeline& p, int t0, int n, bool split, bool keep) { + mtl::Dispatch d; + d.p = &p; + const RmsXpArgs a{(uint32_t) RN, N, MAXT, split, keep, eps}; + d.args = {mtl::Arg::B(src, t0 * RN * 4), mtl::Arg::B(wn), mtl::Arg::B(mhn), mtl::Arg::B(xp), mtl::Arg::V(a)}; + d.grid[0] = n; d.tg[0] = 256; + return d; + }; + auto xprep = [&](const mtl::Pipeline& p, size_t off, int rows, uint32_t stride) { + mtl::Dispatch d; + d.p = &p; + const XpArgs a{N, (uint32_t) rows, MAXT, stride}; + d.args = {mtl::Arg::B(mhn, off), mtl::Arg::B(xp), mtl::Arg::V(a)}; + d.grid[0] = (N / 16 * rows + 255) / 256; d.tg[0] = 256; + return d; + }; + auto comb = [&](bool old, int t0, int n, int out) { + mtl::Dispatch d; + d.p = old ? &bcomb : &pcomb; + const size_t o = t0 * RN * 4; + d.args = {mtl::Arg::B(h), mtl::Arg::B(e, t0 * N * 4), mtl::Arg::B(R[out], o), mtl::Arg::B(whc), mtl::Arg::B(xn[out], o)}; + if (old) d.args.push_back(mtl::Arg::V(OldCombNormArgs{N, HC, (uint32_t) (MAXT * N), eps})); + else d.args.push_back(mtl::Arg::V(CombNormArgs{N, HC, eps})); + d.grid[0] = HC; d.grid[1] = n; d.tg[0] = 256; + return d; + }; + // the replaced chain: a GEMV of n rows per stream; old: the baseline's kernels and [c][t] h + auto per_stream = [&](int n, int copy, bool old) { + const mtl::Pipeline &r = old ? brms : prms, &x = old ? bxp : pxp; + std::vector ds{rms(r, 0, n, n == 1, n > 1)}; + if (n == 1) ds.push_back(gemv(HC, old ? MAXT * N : N, 0, copy)); + for (int c = 0; c < HC && n > 1; ++c) { + if (c > 0) ds.push_back(xprep(x, (size_t) c * N * 4, n, RN)); + ds.push_back(old ? gemv(n, N, (size_t) c * MAXT * N * 4, copy) : gemv(n, RN, (size_t) c * N * 4, copy)); + } + ds.push_back(comb(old, 0, n, 0)); + return ds; + }; + auto batched = [&](int t0, int n, int copy, int out) { + const int rows = n * HC; + std::vector ds{rms(prms, t0, n, true, rows > MAXT)}; + ds.push_back(gemv(std::min(rows, MAXT), N, 0, copy)); + for (int r = MAXT; r < rows; r += MAXT) { + const int k = std::min(rows - r, MAXT); + ds.push_back(xprep(pxp, (size_t) r * N * 4, k, N)); + ds.push_back(gemv(k, N, (size_t) r * N * 4, copy)); + } + ds.push_back(comb(false, t0, n, out)); + return ds; + }; + int fails = 0; + for (int n = 1; n <= MAXT; ++n) { + const size_t bytes = 4 * n * RN; + for (int i = 0; i < 3; ++i) { std::memset(R[i].ptr, 0xCD, R[i].bytes); std::memset(xn[i].ptr, 0xCD, xn[i].bytes); } + cx.run_graph(per_stream(n, 0, false)); + cx.run_graph(batched(0, n, 0, 1)); + for (int t = 0; t < n; ++t) cx.run_graph(batched(t, 1, 0, 2)); + const bool one = tile(n == 1 ? HC : n) == tile(std::min(n * HC, MAXT)); + const bool vs_old = same(R[0], R[1], bytes) && same(xn[0], xn[1], bytes); + const bool vs_one = same(R[1], R[2], bytes) && same(xn[1], xn[2], bytes); + double md = 0; + for (size_t i = 0; i < n * RN; ++i) md = std::max(md, (double) std::fabs(xn[0].as()[i] - xn[1].as()[i])); + const bool ok = vs_one && (!one || vs_old); + fails += !ok; + std::printf("mtp n %d: batched == per-stream %s (%s tile), == n x (n = 1) %s, max |dxn| %.3g: %s\n", n, + vs_old ? "bitwise" : "differs", one ? "same" : "other", vs_one ? "bitwise" : "FAIL", md, ok ? "PASS" : "FAIL"); + if (lb >= 0) { + std::memset(R[2].ptr, 0xCD, R[2].bytes); std::memset(xn[2].ptr, 0xCD, xn[2].bytes); + auto ds = per_stream(n, 0, true); + ds.back() = comb(true, 0, n, 2); + cx.run_graph(ds); + const bool b = same(R[0], R[2], bytes) && same(xn[0], xn[2], bytes); + fails += !b; + std::printf("mtp n %d: per-stream on the baseline's kernels bitwise %s\n", n, b ? "PASS" : "FAIL"); + } + } + if (bench && !fails) { + for (int n : {1, 2, 3, 4, 8}) { + std::vector a, b; + for (int k = 0; k < copies; ++k) { + auto x = per_stream(n, k, false), y = batched(0, n, k, 1); + a.insert(a.end(), x.begin(), x.end()); + b.insert(b.end(), y.begin(), y.end()); + } + cx.warm(a, 0.1); + cx.warm(b, 0.1); + std::vector ta, tb; + for (int p = 0; p < 9; ++p) { + if (p % 2) { tb.push_back(cx.run_graph(b)); ta.push_back(cx.run_graph(a)); } + else { ta.push_back(cx.run_graph(a)); tb.push_back(cx.run_graph(b)); } + } + std::sort(ta.begin(), ta.end()); std::sort(tb.begin(), tb.end()); + const double ma = ta[4] * 1e6 / copies, mb2 = tb[4] * 1e6 / copies; + std::printf("mtp n %d cold: per-stream %zu dispatches %.1f us, batched %zu dispatches %.1f us (%+.1f%%)\n", n, + a.size() / copies, ma, b.size() / copies, mb2, (mb2 / ma - 1) * 100); + } + } + std::printf("%s\n", fails ? "FAIL" : "ALL PASS"); + return fails ? 1 : 0; +} +// `comb_fold_parity 0 mtpcat`: window.cpp mtp_encode's gm_cat (k + v 2560 -> 512 + 512, shared gate + up 2560 -> 640 + +// 640, Q8_0, one token) as gemv_cat / gemv_catk against its two GEMVs beside each other: bitwise, then paired GPU times +// behind a barrier per pair, hot and with the weights walking copies (cold), median of 15. +int mtp_cat_parity() { + struct GemvArgs { uint32_t n_in, n_out, n_tok, row_bytes, x_stride, y_stride; }; + constexpr int NI = 2560; + mtl::Context cx; + const int lg = cx.compile({"quant.metal", "units.metal", "gemv.metal"}); + std::printf("device %s\n", cx.device_name().c_str()); + auto x = cx.alloc(4 * NI), xp = cx.alloc(4 * (4 * NI + NI / 16)); + ref::Rng rng(9); + for (int i = 0; i < NI; ++i) x.as()[i] = rng.sym(); + { + auto prep = cx.pipeline(lg, "xprep"); + mtl::Dispatch d; + d.p = &prep; + d.args = {mtl::Arg::B(x), mtl::Arg::B(xp), mtl::Arg::V((uint32_t) NI), mtl::Arg::V(1u)}; + d.grid[0] = (NI / 16 + 255) / 256; d.tg[0] = 256; + cx.run({d}); + } + int fails = 0; + for (int no : {512, 640}) + for (const auto t : {std::array{32, 1, 4, 4}, std::array{32, 1, 4, 1}}) { + const auto [L, Rr, S, K] = t; + const size_t rb = ref::row_bytes(ref::Q8_0, NI), mb = rb * no; + const int copies = 32; + auto wa = cx.alloc(mb * copies), wb = cx.alloc(mb * copies); + ref::random_blocks(ref::Q8_0, wa.as(), no, NI, 3); + ref::random_blocks(ref::Q8_0, wb.as(), no, NI, 4); + for (int k = 1; k < copies; ++k) { + std::memcpy(wa.as() + k * mb, wa.ptr, mb); + std::memcpy(wb.as() + k * mb, wb.ptr, mb); + } + auto ya = cx.alloc(4 * no), yb = cx.alloc(4 * no), za = cx.alloc(4 * no), zb = cx.alloc(4 * no); + const std::string tl = std::to_string(L) + "_" + std::to_string(Rr) + "_" + std::to_string(S); + const std::string gn = "gemv_" + std::to_string(ref::Q8_0) + "_" + tl + "_" + std::to_string(K), + cn = K == 1 ? "gemvc_" + std::to_string(ref::Q8_0) + "_" + tl + : "gemvck_" + std::to_string(ref::Q8_0) + "_" + tl + "_" + std::to_string(K); + auto pg = cx.pipeline(lg, gn, mtl::Constants().i(0, 1)), pc = cx.pipeline(lg, cn, mtl::Constants().i(0, 1)); + const int rpt = (S / K) * (32 / L) * Rr; + const GemvArgs a{NI, (uint32_t) no, 1, (uint32_t) rb, NI, (uint32_t) no}; + GemvArgs a2 = a; + a2.n_out = 2 * no; + auto gemv = [&](const mtl::Buffer& w, int k, const mtl::Buffer& y, bool barrier) { + mtl::Dispatch d; + d.p = &pg; + d.args = {mtl::Arg::B(w, k * mb), mtl::Arg::B(xp), mtl::Arg::B(xp, 4ull * 4 * NI), mtl::Arg::B(y), mtl::Arg::V(a)}; + d.grid[0] = (no + rpt - 1) / rpt; d.tg[0] = S * 32; + d.barrier = barrier; + return d; + }; + auto cat = [&](int k) { + mtl::Dispatch d; + d.p = &pc; + d.args = {mtl::Arg::B(wa, k * mb), mtl::Arg::B(xp), mtl::Arg::B(xp, 4ull * 4 * NI), mtl::Arg::B(yb), + mtl::Arg::V(a2), mtl::Arg::B(wb, k * mb), mtl::Arg::B(zb), mtl::Arg::V((uint32_t) no)}; + d.grid[0] = (2 * no + rpt - 1) / rpt; d.tg[0] = S * 32; + return d; + }; + cx.run_graph({gemv(wa, 0, ya, true), gemv(wb, 0, za, false)}); + cx.run_graph({cat(0)}); + const bool ok = same(ya, yb, 4 * no) && same(za, zb, 4 * no); + fails += !ok; + std::printf("gm_cat Q8_0 2560 -> %d+%d tile %d,%d,%d,%d: bitwise %s", no, no, L, Rr, S, K, ok ? "PASS" : "FAIL"); + for (bool cold : {false, true}) { + std::vector two, one; + for (int r = 0; r < 256 / (cold ? copies : 1); ++r) + for (int k = 0; k < (cold ? copies : 1); ++k) { + two.push_back(gemv(wa, k, ya, true)); + two.push_back(gemv(wb, k, za, false)); + one.push_back(cat(k)); + } + cx.warm(two, 0.05); + cx.warm(one, 0.05); + std::vector tt, to; + for (int p = 0; p < 15; ++p) { + if (p % 2) { to.push_back(cx.run_graph(one)); tt.push_back(cx.run_graph(two)); } + else { tt.push_back(cx.run_graph(two)); to.push_back(cx.run_graph(one)); } + } + std::sort(tt.begin(), tt.end()); std::sort(to.begin(), to.end()); + const double m2 = tt[7] * 1e6 / one.size(), m1 = to[7] * 1e6 / one.size(); + std::printf(" %s: 2 gemv %.2f us, cat %.2f us (%+.1f%%)", cold ? "cold" : "hot", m2, m1, (m1 / m2 - 1) * 100); + } + std::printf("\n"); + } + std::printf("%s\n", fails ? "FAIL" : "ALL PASS"); + return fails ? 1 : 0; +} +} // namespace + +int mtp_defer_parity() { + struct RmsXpArgs { uint32_t n, xp_n, plane_tok, split, keep; float eps; }; + struct XpArgs { uint32_t n, n_tok, plane_tok, stride; }; + struct CombNormArgs { uint32_t n, hc; float eps; }; + struct GemvArgs { uint32_t n_in, n_out, n_tok, row_bytes, x_stride, y_stride; }; + struct GrArgs { uint32_t n_embd, hc, hc_lr, n_tok, apply; float eps; }; + struct QsaArgs { uint32_t pos0, pos_base, max_cells, n_tok; float eps; }; + constexpr int N = 2560, HC = 4, LR = 320, MAXT = 8, MAXC = 64; + constexpr float eps = 1e-6f; + mtl::Context cx; + const std::string msl = cx.kernel_dir + "/../src/engine/msl/mtp.metal"; + const int lg = cx.compile({"quant.metal", "units.metal", "gemv.metal", "gr.metal"}); + const int lm = cx.compile({msl}, "#define NO_GGML_COMMON 1\n"); + const int lq = cx.compile({"qsa.metal"}, "#define NO_GGML_COMMON 1\n"); + std::printf("device %s\n", cx.device_name().c_str()); + + const int copies = 8; + const size_t rb_n = ref::row_bytes(ref::Q8_0, N), rb_rn = ref::row_bytes(ref::Q8_0, HC * N), rb_lr = ref::row_bytes(ref::Q8_0, LR); + const size_t mb_fce = rb_n * N, mb_fch = rb_n * N, mb_kv = rb_n * 512, mb_q = rb_n * 12288; + const size_t mb_gdn = rb_rn * (LR + HC), mb_gup = rb_lr * (HC * N); + auto W_fce = cx.alloc(mb_fce * copies), W_fch = cx.alloc(mb_fch * copies); + auto W_k = cx.alloc(mb_kv * copies), W_v = cx.alloc(mb_kv * copies), W_q = cx.alloc(mb_q * copies); + auto W_gdn = cx.alloc(mb_gdn * copies), W_gup = cx.alloc(mb_gup * copies); + ref::random_blocks(ref::Q8_0, W_fce.as(), N, N, 11); + ref::random_blocks(ref::Q8_0, W_fch.as(), N, N, 12); + ref::random_blocks(ref::Q8_0, W_k.as(), N, 512, 13); + ref::random_blocks(ref::Q8_0, W_v.as(), N, 512, 14); + ref::random_blocks(ref::Q8_0, W_q.as(), N, 12288, 15); + ref::random_blocks(ref::Q8_0, W_gdn.as(), HC * N, LR + HC, 16); + ref::random_blocks(ref::Q8_0, W_gup.as(), LR, HC * N, 17); + auto dup = [&](mtl::Buffer& B, size_t sz) { + for (int k = 1; k < copies; ++k) std::memcpy(B.as() + k * sz, B.ptr, sz); + }; + dup(W_fce, mb_fce); dup(W_fch, mb_fch); dup(W_k, mb_kv); dup(W_v, mb_kv); dup(W_q, mb_q); dup(W_gdn, mb_gdn); dup(W_gup, mb_gup); + + const size_t RN = (size_t) HC * N, plane = 4ull * MAXT * N; + auto memb = cx.alloc(4 * MAXT * N), srcR = cx.alloc(4 * MAXT * RN); + auto wn_e = cx.alloc(4 * N), wn_h = cx.alloc(4 * RN), wn_hc = cx.alloc(4 * RN); + auto wn_q = cx.alloc(4 * 256), wn_k = cx.alloc(4 * 256), wn_iq = cx.alloc(4 * 128), wn_ik = cx.alloc(4 * 128); + auto cs = cx.alloc(4 * MAXC * 32 * 2), mrope = cx.alloc(4 * MAXC * 3), midx = cx.alloc(4 * 1024); + auto men = cx.alloc(4 * MAXT * N), mhn = cx.alloc(4 * MAXT * RN); + auto XEn = cx.alloc(4 * plane + 4ull * MAXT * N / 16), XHn = cx.alloc(4 * plane + 4ull * MAXT * N / 16); + auto me2 = cx.alloc(4 * MAXT * N), mh2 = cx.alloc(4 * MAXT * RN); + + struct OutSet { + mtl::Buffer mR, xn, lo, minj, XA, kcur, vcur, qfull, qcur, qidx, kq, vq, ksc, vsc, tail, dead, pooled; + } O[2]; + for (int v = 0; v < 2; ++v) { + O[v].mR = cx.alloc(4 * MAXT * RN); + O[v].xn = cx.alloc(4 * MAXT * RN); + O[v].lo = cx.alloc(4 * MAXT * LR); + O[v].minj = cx.alloc(4 * MAXT * HC); + O[v].XA = cx.alloc(4 * plane + 4ull * MAXT * N / 16); + O[v].kcur = cx.alloc(4 * MAXT * 512); + O[v].vcur = cx.alloc(4 * MAXT * 512); + O[v].qfull = cx.alloc(4 * MAXT * 12288); + O[v].qcur = cx.alloc(4 * MAXT * 24 * 256); + O[v].qidx = cx.alloc(4 * MAXT * 512); + O[v].kq = cx.alloc(MAXC * 2 * 256); + O[v].vq = cx.alloc(MAXC * 2 * 256); + O[v].ksc = cx.alloc(2 * MAXC * 2 * 4); + O[v].vsc = cx.alloc(2 * MAXC * 2 * 4); + O[v].tail = cx.alloc(4 * 4 * 128); + O[v].dead = cx.alloc(4 * 128); + O[v].pooled = cx.alloc(4 * 32 * 128); + } + + ref::Rng rng(91); + for (size_t i = 0; i < MAXT * N; ++i) memb.as()[i] = rng.sym(); + for (size_t i = 0; i < MAXT * RN; ++i) srcR.as()[i] = rng.sym() * 2.0f; + for (size_t i = 0; i < N; ++i) wn_e.as()[i] = 1.0f + 0.1f * rng.sym(); + for (size_t i = 0; i < RN; ++i) { wn_h.as()[i] = 1.0f + 0.1f * rng.sym(); wn_hc.as()[i] = 1.0f + 0.1f * rng.sym(); } + for (size_t i = 0; i < 256; ++i) { wn_q.as()[i] = 1.0f + 0.1f * rng.sym(); wn_k.as()[i] = 1.0f + 0.1f * rng.sym(); } + for (size_t i = 0; i < 128; ++i) { wn_iq.as()[i] = 1.0f; wn_ik.as()[i] = 1.0f; } + for (int p = 0; p < MAXC; ++p) + for (int i = 0; i < 32; ++i) { + const double th = (double) p * std::pow(1e7, -2.0 * i / 64.0); + cs.as()[(p * 32 + i) * 2] = (float) std::cos(th); + cs.as()[(p * 32 + i) * 2 + 1] = (float) std::sin(th); + } + + mtl::Pipeline prms = cx.pipeline(lm, "mtp_rms_xprep"), pxp = cx.pipeline(lm, "mtp_xprep"), + pcomb = cx.pipeline(lm, "mtp_combine_gnorm"); + mtl::Pipeline pq_full = cx.pipeline(lq, "qsa_prep"), + pq_kv = cx.pipeline(lq, "qsa_prep", mtl::Constants().i(28, 24)), + pq_lastq = cx.pipeline(lq, "qsa_prep", mtl::Constants().b(29, true)); + + auto tile = [](int n_out, int rows, bool force_single_tile = false) { + if (n_out <= 640) { + if (rows == 1 && !force_single_tile) return std::array{32, 1, 4, 4}; + return std::array{16, 1, 4, 1}; + } + if (rows == 1 && !force_single_tile) return std::array{32, 1, 4, 4}; + if (rows <= 3 && !force_single_tile) return std::array{32, 2, 4, 1}; + return std::array{16, 1, 4, 1}; + }; + std::map gp; + auto get_p = [&](int lib, const std::string& fn, mtl::Constants c, const std::string& key) -> const mtl::Pipeline* { + if (!gp.count(key)) gp.emplace(key, cx.pipeline(lib, fn, c)); + return &gp.at(key); + }; + + auto build_encode = [&](std::vector& ds, OutSet& o, int erow, size_t src_off, int n, uint32_t cell0, + bool full, int copy, bool same_tile) { + auto push_d = [&](mtl::Dispatch d, bool nb = false) { d.barrier = !nb; ds.push_back(d); }; + // 1. rms_xp embedding + hidden + { + mtl::Dispatch d; + d.p = &prms; + RmsXpArgs a{N, N, MAXT, 0u, 0u, eps}; + d.args = {mtl::Arg::B(memb, (size_t) erow * N * 4), mtl::Arg::B(wn_e), mtl::Arg::B(men), mtl::Arg::B(XEn), mtl::Arg::V(a)}; + d.grid[0] = n; d.tg[0] = 256; + push_d(d, false); + } + const int hrows = n * HC; + { + mtl::Dispatch d; + d.p = &prms; + RmsXpArgs a{(uint32_t) RN, N, MAXT, 1u, hrows > MAXT ? 1u : 0u, eps}; + d.args = {mtl::Arg::B(srcR, src_off), mtl::Arg::B(wn_h), mtl::Arg::B(mhn), mtl::Arg::B(XHn), mtl::Arg::V(a)}; + d.grid[0] = n; d.tg[0] = 256; + push_d(d, true); + } + auto gm = [&](const mtl::Buffer& W, size_t mb, int n_out, const mtl::Buffer& X, const mtl::Buffer& Y, int rows, + size_t y_off, int x_row, bool nb) { + const auto [L, Rr, S, K] = tile(n_out, rows, same_tile && n_out <= 2560); + const std::string fn = "gemv_8_" + std::to_string(L) + "_" + std::to_string(Rr) + "_" + std::to_string(S) + "_" + std::to_string(K); + mtl::Dispatch d; + d.p = get_p(lg, fn, mtl::Constants().i(0, rows), fn + "#" + std::to_string(rows)); + GemvArgs a{N, (uint32_t) n_out, (uint32_t) rows, (uint32_t) rb_n, N, (uint32_t) n_out}; + const size_t x_off = (size_t) x_row * N * 4, s16_off = (size_t) x_row * (N / 16) * 4; + d.args = {mtl::Arg::B(W, copy * mb), mtl::Arg::B(X, x_off), mtl::Arg::B(X, plane * 4 + s16_off), mtl::Arg::B(Y, y_off), mtl::Arg::V(a)}; + const int rpt = (S / K) * (32 / L) * Rr; + d.grid[0] = (n_out + rpt - 1) / rpt; d.tg[0] = S * 32; + push_d(d, nb); + }; + gm(W_fce, mb_fce, N, XEn, me2, n, 0, 0, false); + gm(W_fch, mb_fch, N, XHn, mh2, std::min(hrows, MAXT), 0, 0, true); + for (int r = MAXT; r < hrows; r += MAXT) { + const int k = std::min(hrows - r, MAXT); + mtl::Dispatch dx; + dx.p = &pxp; + XpArgs xa{N, (uint32_t) k, MAXT, N}; + dx.args = {mtl::Arg::B(mhn, (size_t) r * N * 4), mtl::Arg::B(XHn), mtl::Arg::V(xa)}; + dx.grid[0] = (N / 16 * k + 255) / 256; dx.tg[0] = 256; + push_d(dx, false); + gm(W_fch, mb_fch, N, XHn, mh2, k, (size_t) r * N * 4, 0, false); + } + { + mtl::Dispatch d; + d.p = &pcomb; + CombNormArgs a{N, HC, eps}; + d.args = {mtl::Arg::B(mh2), mtl::Arg::B(me2), mtl::Arg::B(o.mR), mtl::Arg::B(wn_hc), mtl::Arg::B(o.xn), mtl::Arg::V(a)}; + d.grid[0] = HC; d.grid[1] = n; d.tg[0] = 256; + push_d(d, false); + } + // gr_read_m(attn_hyper_connection) + { + GrArgs a{N, HC, LR, (uint32_t) n, 0u, eps}; + mtl::Dispatch d2; + const std::string dn = n > 4 ? "gr_down_q8_8_8_r2" : "gr_down_q8_8_8"; + d2.p = get_p(lg, dn, mtl::Constants().i(1, n), dn + "#" + std::to_string(n)); + d2.args = {mtl::Arg::B(W_gdn, copy * mb_gdn), mtl::Arg::B(W_gdn, copy * mb_gdn + (size_t) LR * rb_rn), + mtl::Arg::B(o.xn), mtl::Arg::B(o.lo), mtl::Arg::B(o.minj), mtl::Arg::V(a)}; + d2.grid[0] = n > 4 ? (LR + HC) / 2 : (LR + HC); d2.tg[0] = 256; + push_d(d2, false); + mtl::Dispatch d3; + d3.p = get_p(lg, "gr_up_q8", mtl::Constants().i(1, n), "gr_up_q8#" + std::to_string(n)); + GrArgs au = a; au.n_tok = MAXT; + d3.args = {mtl::Arg::B(W_gup, copy * mb_gup), mtl::Arg::B(o.lo), mtl::Arg::B(o.xn), mtl::Arg::B(o.XA), mtl::Arg::V(au)}; + d3.grid[0] = N / 16; d3.tg[0] = 256; + push_d(d3, false); + } + // gm_cat(k_proj, v_proj) + if (n == 1 && !same_tile) { + const std::string fn = "gemvck_8_32_1_4_4"; + mtl::Dispatch d; + d.p = get_p(lg, fn, mtl::Constants().i(0, 1), fn + "#1"); + GemvArgs a{N, 1024, 1, (uint32_t) rb_n, N, 512}; + const uint32_t sp = 512; + d.args = {mtl::Arg::B(W_k, copy * mb_kv), mtl::Arg::B(o.XA), mtl::Arg::B(o.XA, plane * 4), mtl::Arg::B(o.kcur), + mtl::Arg::V(a), mtl::Arg::B(W_v, copy * mb_kv), mtl::Arg::B(o.vcur), mtl::Arg::V(sp)}; + d.grid[0] = (1024 + 1 - 1) / 1; d.tg[0] = 4 * 32; + push_d(d, false); + } else { + gm(W_k, mb_kv, 512, o.XA, o.kcur, n, 0, 0, false); + gm(W_v, mb_kv, 512, o.XA, o.vcur, n, 0, 0, true); + } + if (full) { + gm(W_q, mb_q, 12288, o.XA, o.qfull, 1, 0, n - 1, true); + } + { + QsaArgs qa{cell0, 0, MAXC, (uint32_t) n, eps}; + mtl::Dispatch d; + d.p = !full ? &pq_kv : (n > 1 ? &pq_lastq : &pq_full); + d.args = {mtl::Arg::B(o.qfull), mtl::Arg::B(o.kcur), mtl::Arg::B(o.vcur), mtl::Arg::B(midx), mtl::Arg::B(midx, 512 * 4), + mtl::Arg::B(wn_q), mtl::Arg::B(wn_k), mtl::Arg::B(wn_iq), mtl::Arg::B(wn_ik), + mtl::Arg::B(cs), mtl::Arg::B(o.qcur), mtl::Arg::B(o.qidx), mtl::Arg::B(o.kq), mtl::Arg::B(o.vq), + mtl::Arg::B(o.ksc), mtl::Arg::B(o.vsc), mtl::Arg::B(o.tail), mtl::Arg::B(o.dead), + mtl::Arg::B(o.pooled), mtl::Arg::V(qa), mtl::Arg::B(mrope)}; + d.grid[0] = full ? 28 : 4; d.grid[1] = n; d.tg[0] = 256; + push_d(d, false); + } + }; + + int fails = 0; + for (int a : {1, 2, 3}) { + const int n = a + 1; + for (bool same_tile : {true, false}) { + for (int v = 0; v < 2; ++v) { + std::memset(O[v].mR.ptr, 0xCD, O[v].mR.bytes); + std::memset(O[v].minj.ptr, 0xCD, O[v].minj.bytes); + std::memset(O[v].qcur.ptr, 0xCD, O[v].qcur.bytes); + std::memset(O[v].kq.ptr, 0xCD, O[v].kq.bytes); + std::memset(O[v].vq.ptr, 0xCD, O[v].vq.bytes); + std::memset(O[v].ksc.ptr, 0xCD, O[v].ksc.bytes); + std::memset(O[v].vsc.ptr, 0xCD, O[v].vsc.bytes); + } + std::vector ds_two, ds_one; + build_encode(ds_two, O[0], 0, 0, a, 10, false, 0, same_tile); + build_encode(ds_two, O[0], a, (size_t) a * RN * 4, 1, 10 + a, true, 0, same_tile); + build_encode(ds_one, O[1], 0, 0, n, 10, true, 0, same_tile); + cx.run_graph(ds_two); + cx.run_graph(ds_one); + + // In O[0], the 2nd call wrote row `a` into O[0].mR[0] and O[0].minj[0]; in O[1], row `a` is at O[1].mR[a] and O[1].minj[a] + const float* r0 = O[0].mR.as(); + const float* r1 = O[1].mR.as() + (size_t) a * RN; + const float* inj0 = O[0].minj.as(); + const float* inj1 = O[1].minj.as() + (size_t) a * HC; + const bool r_bit = std::memcmp(r0, r1, 4 * RN) == 0; + const bool inj_bit = std::memcmp(inj0, inj1, 4 * HC) == 0; + const bool q_bit = same(O[0].qcur, O[1].qcur, 4 * 24 * 256); + const bool kv_bit = std::memcmp(O[0].kq.as() + 10 * 512, O[1].kq.as() + 10 * 512, n * 512) == 0 && + std::memcmp(O[0].vq.as() + 10 * 512, O[1].vq.as() + 10 * 512, n * 512) == 0 && + std::memcmp(O[0].ksc.as() + 10 * 8, O[1].ksc.as() + 10 * 8, n * 16) == 0 && + std::memcmp(O[0].vsc.as() + 10 * 8, O[1].vsc.as() + 10 * 8, n * 16) == 0; + double max_dr = 0, max_dq = 0; + for (size_t i = 0; i < RN; ++i) max_dr = std::max(max_dr, (double) std::fabs(r0[i] - r1[i])); + for (size_t i = 0; i < 24 * 256; ++i) max_dq = std::max(max_dq, (double) std::fabs(O[0].qcur.as()[i] - O[1].qcur.as()[i])); + const bool ok = same_tile ? (r_bit && inj_bit && q_bit && kv_bit) : (max_dr < 1e-4 && max_dq < 1e-4 && kv_bit); + fails += !ok; + std::printf("mtpdefer a=%d (n=%d) same_tile=%d: R %s (max|dR|=%.3g) inj %s qcur %s (max|dq|=%.3g) kv %s -> %s\n", + a, n, same_tile ? 1 : 0, r_bit ? "bitwise" : "close", max_dr, inj_bit ? "bitwise" : "close", + q_bit ? "bitwise" : "close", max_dq, kv_bit ? "bitwise" : "DIFF", ok ? "PASS" : "FAIL"); + } + // Paired A/B benchmark (cold DRAM weights across `copies` copies) + for (bool cold : {true, false}) { + std::vector two, one; + const int reps = cold ? copies : 16; + for (int r = 0; r < reps; ++r) { + const int k = cold ? (r % copies) : 0; + build_encode(two, O[0], 0, 0, a, 10, false, k, false); + build_encode(two, O[0], a, (size_t) a * RN * 4, 1, 10 + a, true, k, false); + build_encode(one, O[1], 0, 0, n, 10, true, k, false); + } + cx.warm(two, 0.05); + cx.warm(one, 0.05); + std::vector tt, to; + for (int p = 0; p < 15; ++p) { + if (p % 2) { to.push_back(cx.run_graph(one)); tt.push_back(cx.run_graph(two)); } + else { tt.push_back(cx.run_graph(two)); to.push_back(cx.run_graph(one)); } + } + std::sort(tt.begin(), tt.end()); std::sort(to.begin(), to.end()); + const double m2 = tt[7] * 1e6 / reps, m1 = to[7] * 1e6 / reps; + std::printf(" mtpdefer a=%d %s: 2-pass (%zu disp) %.1f us -> 1-pass (%zu disp) %.1f us (%+.1f%%, %.1f us saved)\n", + a, cold ? "cold" : "hot ", two.size() / reps, m2, one.size() / reps, m1, (m1 / m2 - 1) * 100.0, m2 - m1); + } + } + std::printf("%s\n", fails ? "FAIL" : "ALL PASS"); + return fails ? 1 : 0; +} + +int main(int argc, char** argv) { + setvbuf(stdout, nullptr, _IONBF, 0); + if (argc > 2 && std::strcmp(argv[2], "mtpdefer") == 0) return mtp_defer_parity(); + if (argc > 2 && std::strcmp(argv[2], "mtpcat") == 0) return mtp_cat_parity(); + if (argc > 2 && std::strcmp(argv[2], "mtp") == 0) { + const bool bench = argc > 3 && std::strcmp(argv[3], "bench") == 0; + return mtp_parity(bench, argc > 3 + bench ? argv[3 + bench] : nullptr); + } + if (argc > 2 && std::strcmp(argv[2], "epi") == 0) + return epi_parity(uint32_t(std::atoi(argv[1])), argc > 3 && std::strcmp(argv[3], "bench") == 0, argc > 4 ? argv[4] : nullptr); + if (argc > 3 && std::strcmp(argv[2], "topk") == 0) return topk_parity(uint32_t(std::atoi(argv[1])), argv[3]); + if (argc > 2 && std::strcmp(argv[2], "gu") == 0) + return gu_parity(uint32_t(std::atoi(argv[1])), argc > 3 && std::strcmp(argv[3], "bench") == 0, argc > 4 ? argv[4] : nullptr); + if (argc > 2 && std::strcmp(argv[2], "hcup") == 0) + return hc_up_parity(uint32_t(std::atoi(argv[1])), argc > 3 && std::strcmp(argv[3], "bench") == 0, + argc > 3 && std::strcmp(argv[3], "confirm") == 0); + if (argc > 3 && std::strcmp(argv[2], "ab") == 0) + return ab_parity(uint32_t(std::atoi(argv[1])), argv[3], argc > 4 ? argv[4] : nullptr); + if (argc > 3 && std::strcmp(argv[2], "denseab") == 0) + return dense_ab_parity(uint32_t(std::atoi(argv[1])), argv[3], argc > 4 ? argv[4] : nullptr); + const bool gr_only = argc > 2 && std::strcmp(argv[2], "grvec") == 0; + mtl::Context cx; + const std::string source = gr_only ? mtl::read_file(cx.kernel_dir + "/prompt.metal") : ""; + const std::string prelude = "#define NO_GGML_COMMON 1\nusing namespace metal;\n" + "struct GrArgs { uint n_embd, hc, hc_lr, n_tok, apply; float eps; };\n"; + const auto begin = source.find("// ---- the HC read"); + const auto end = source.find("// gemm_hc_up:"); + if (gr_only && (begin == std::string::npos || end == std::string::npos || end <= begin)) { + std::fprintf(stderr, "GR source markers missing\n"); + return 1; + } + const int lib = gr_only ? cx.compile({}, prelude + source.substr(begin, end - begin)) : + cx.compile({"quant.metal", "units.metal", "mma.metal", "dense_mm.metal", "gr.metal", "prompt.metal"}); + const std::string mm = gr_only ? mtl::read_file(cx.kernel_dir + "/moe_mm.metal") : ""; + const auto mb = mm.find("kernel void mm_combine("); + const auto me = mm.find("\n}", mb == std::string::npos ? 0 : mb); + if (gr_only && (mb == std::string::npos || me == std::string::npos)) { + std::fprintf(stderr, "mm_combine source marker missing\n"); + return 1; + } + const int lmm = gr_only ? cx.compile({}, prelude + mm.substr(mb, me + 2 - mb)) : + cx.compile({"quant.metal", "units.metal", "mma.metal", "moe_mm.metal"}); + std::printf("device %s\n", cx.device_name().c_str()); + const uint32_t N = 2560, HC = 4, LR = 320, TOPK = 10, NEXP = 512, TMR = 16; + const size_t RN = (size_t) HC * N; + const uint32_t TB = argc > 1 ? (uint32_t) std::atoi(argv[1]) : 1024; + const bool bench = argc > 2 && std::strcmp(argv[2], "bench") == 0; + mtl::Pipeline pcomb = cx.pipeline(lmm, "mm_combine"), pnhm = cx.pipeline(lib, "gr_norm_h_moe"), + pnhc = cx.pipeline(lib, "gr_norm_h_moec"), pam = cx.pipeline(lib, "gr_apply_moe"), + pamc = cx.pipeline(lib, "gr_apply_moec"), pnh = cx.pipeline(lib, "gr_norm_h"), + pnh4 = cx.pipeline(lib, "gr_norm_h4"), pup = gr_only ? mtl::Pipeline{} : cx.pipeline(lib, "gemm_hc_up_64_64_4"); + std::vector baseline; + if (gr_only && argc > 3) { + const std::string ref = mtl::read_file(argv[3]); + const auto rb = ref.find("// ---- the HC read"), re = ref.find("// gemm_hc_up:"); + if (rb == std::string::npos || re == std::string::npos || re <= rb) return 1; + const int rl = cx.compile({}, prelude + ref.substr(rb, re - rb)); + for (const char* name : {"gr_norm_h_moec", "gr_apply_moec", "gr_norm_h4", "gr_norm_h4"}) + baseline.push_back(cx.pipeline(rl, name)); + } + auto D = [](const mtl::Pipeline& p, std::vector a, uint32_t gx, uint32_t gy, uint32_t th = 256) { + mtl::Dispatch d; + d.p = &p; + d.args = std::move(a); + d.grid[0] = gx; + d.grid[1] = gy; + d.tg[0] = th; + return d; + }; + int fails = 0; + for (uint32_t T : gr_only ? std::vector{TB} : std::vector{1u, 7u, 64u, TB}) { + const size_t mr = (size_t) T * TOPK + NEXP * TMR; + mtl::Buffer R0 = cx.alloc(4 * T * RN), Ra = cx.alloc(4 * T * RN), Rb = cx.alloc(4 * T * RN), YE = cx.alloc(4 * mr * N), + pos = cx.alloc(4 * T * TOPK), sw = cx.alloc(4 * mr), yr = cx.alloc(4 * T * N), ys = cx.alloc(4 * T * N), + sgv = cx.alloc(4 * T), inj = cx.alloc(4 * T * HC), wn = cx.alloc(4 * RN), xa = cx.alloc(2 * T * RN), + xb = cx.alloc(2 * T * RN), loh = cx.alloc(2 * T * LR), WU = cx.alloc(2 * RN * LR), + mixed = cx.alloc(4 * T * N), mha = cx.alloc(2 * T * N), mhb = cx.alloc(2 * T * N), nil{}; + ref::Rng r(23 + T); + auto fill = [&](const mtl::Buffer& b, float s, float o) { + for (size_t i = 0; i < b.bytes / 4; ++i) b.as()[i] = o + s * r.sym(); + }; + auto fillh = [&](const mtl::Buffer& b, float s) { + for (size_t i = 0; i < b.bytes / 2; ++i) b.as()[i] = ref::f32_to_f16(s * r.sym()); + }; + fill(R0, 1.0f, 0.0f); fill(YE, 0.5f, 0.0f); fill(sw, 0.1f, 0.1f); fill(ys, 0.5f, 0.0f); fill(sgv, 3.0f, 0.0f); + fill(inj, 4.0f, 0.0f); fill(wn, 0.2f, 1.0f); fillh(loh, 1.0f); fillh(WU, 0.1f); fillh(xa, 1.0f); + std::memcpy(xb.ptr, xa.ptr, xb.bytes); + // mm_place's slots: distinct rows of YE (expert-sorted, so scattered), a few dropped (-1) + std::vector perm(mr); + std::iota(perm.begin(), perm.end(), 0); + std::shuffle(perm.begin(), perm.end(), std::mt19937(T)); + for (size_t i = 0; i < (size_t) T * TOPK; ++i) pos.as()[i] = (i % 37 == 5) ? -1 : perm[i]; + const GrArgs a{N, HC, LR, T, 1u, 1e-6f}, a0{N, HC, LR, T, 0u, 1e-6f}; + const uint32_t ca[4] = {T, N, TOPK, 0}; + const mtl::Dispatch comb = D(pcomb, {mtl::Arg::B(YE), mtl::Arg::B(pos), mtl::Arg::B(sw), mtl::Arg::B(yr), mtl::Arg::V(ca)}, + (T * N + 255) / 256, 1); + auto nhm = [&](const mtl::Buffer& R, const mtl::Buffer& x) { + return D(pnhm, {mtl::Arg::B(R), mtl::Arg::B(yr), mtl::Arg::B(ys), mtl::Arg::B(sgv), mtl::Arg::B(inj), mtl::Arg::B(wn), + mtl::Arg::B(x), mtl::Arg::V(a)}, 1, T); + }; + auto nhc = [&](const mtl::Buffer& R, const mtl::Buffer& x) { + return D(pnhc, {mtl::Arg::B(R), mtl::Arg::B(YE), mtl::Arg::B(pos), mtl::Arg::B(sw), mtl::Arg::B(ys), mtl::Arg::B(sgv), + mtl::Arg::B(inj), mtl::Arg::B(wn), mtl::Arg::B(x), mtl::Arg::V(a), mtl::Arg::V(TOPK)}, 1, T); + }; + auto am = [&](const mtl::Buffer& R) { + return D(pam, {mtl::Arg::B(R), mtl::Arg::B(yr), mtl::Arg::B(ys), mtl::Arg::B(sgv), mtl::Arg::B(inj), mtl::Arg::V(a)}, + RN / 256, T); + }; + auto amc = [&](const mtl::Buffer& R) { + return D(pamc, {mtl::Arg::B(R), mtl::Arg::B(YE), mtl::Arg::B(pos), mtl::Arg::B(sw), mtl::Arg::B(ys), mtl::Arg::B(sgv), + mtl::Arg::B(inj), mtl::Arg::V(a), mtl::Arg::V(TOPK)}, (N + 255) / 256, T); + }; + auto nh = [&](const mtl::Pipeline& p, const mtl::Buffer& R, const mtl::Buffer& x, const GrArgs& g, uint32_t gx) { + return D(p, {mtl::Arg::B(R), mtl::Arg::B(ys), mtl::Arg::B(inj), mtl::Arg::B(wn), mtl::Arg::B(x), mtl::Arg::V(g)}, gx, T); + }; + const GemmArgs gg{(int) T, (int) RN, (int) LR, 0}; + const EpiArgs ea{T, N, HC, LR}; + auto up = [&](const mtl::Buffer& m, const mtl::Buffer& mh) { + return D(pup, {mtl::Arg::B(loh), mtl::Arg::B(WU), mtl::Arg::B(xa), mtl::Arg::B(m), mtl::Arg::B(mh), mtl::Arg::V(gg), + mtl::Arg::V(ea)}, RN / 64, (T + 63) / 64, 128); + }; + auto reset = [&] { + std::memcpy(Ra.ptr, R0.ptr, Ra.bytes); + std::memcpy(Rb.ptr, R0.ptr, Rb.bytes); + std::memset(yr.ptr, 0, yr.bytes); + }; + auto check = [&](const char* what, bool ok) { + fails += !ok; + std::printf("%-26s T %4u: bitwise %s\n", what, T, ok ? "PASS" : "FAIL"); + }; + reset(); + cx.run({comb, nhm(Ra, xa)}); + cx.run({nhc(Rb, xb)}); + check("gr_norm_h_moec", same(Ra, Rb, Ra.bytes) && same(xa, xb, xa.bytes)); + reset(); + cx.run({comb, am(Ra)}); + cx.run({amc(Rb)}); + check("gr_apply_moec", same(Ra, Rb, Ra.bytes)); + for (const GrArgs* g : {&a0, &a}) { + reset(); + cx.run({nh(pnh, Ra, xa, *g, HC)}); + cx.run({nh(pnh4, Rb, xb, *g, 1)}); + check(g->apply ? "gr_norm_h4 (apply 1)" : "gr_norm_h4 (apply 0)", same(Ra, Rb, Ra.bytes) && same(xa, xb, xa.bytes)); + } + if (gr_only) { + const std::vector ds = { + nhc(Rb, xb), amc(Rb), nh(pnh4, Rb, xb, a, 1), nh(pnh4, Rb, xb, a0, 1)}; + for (size_t i = 0; i < ds.size(); ++i) { + const auto& d = ds[i]; + auto ref = d; + if (!baseline.empty()) ref.p = &baseline[i]; + cx.run({d}); + if (!baseline.empty()) cx.run({ref}); + std::printf("grvec %s apply %u ms4", d.p->name.c_str(), i == 3 ? 0u : 1u); + std::vector ct, bt, ratios; + for (int sample = 0; sample < 5; ++sample) { + double ms, rms = 0.0; + if (!baseline.empty() && sample % 2 == 0) rms = cx.run({ref}, 4) * 1e3; + ms = cx.run({d}, 4) * 1e3; + if (!baseline.empty() && sample % 2 != 0) rms = cx.run({ref}, 4) * 1e3; + std::printf(" %.3f", ms); + if (!baseline.empty()) { + std::printf("/%.3f", rms); + ct.push_back(ms); bt.push_back(rms); ratios.push_back(ms / rms); + } + } + if (!baseline.empty()) { + std::sort(ct.begin(), ct.end()); std::sort(bt.begin(), bt.end()); std::sort(ratios.begin(), ratios.end()); + std::printf(" med %.3f/%.3f ms ratio %.4f", ct[2], bt[2], ratios[2]); + } + std::printf("\n"); + } + continue; + } + cx.run({up(mixed, mha)}); + cx.run({up(nil, mhb)}); + check("gemm_hc_up, mixed unbound", same(mha, mhb, mha.bytes)); + if (T != TB || !bench) continue; + // ---- GPU times: 48 reps in one command buffer (a chunk's layers), best of 3 + auto time = [&](const std::vector& ds) { + cx.warm(ds, 0.03); + double best = 1e9; + for (int i = 0; i < 3; ++i) best = std::min(best, cx.run(ds, 48)); + return best * 1e3; + }; + const double t_comb = time({comb}), t_nhm = time({nhm(Ra, xa)}), t_pair = time({comb, nhm(Ra, xa)}), + t_nhc = time({nhc(Rb, xb)}), t_am = time({comb, am(Ra)}), t_amc = time({amc(Rb)}), + t_nh = time({nh(pnh, Ra, xa, a, HC)}), t_nh4 = time({nh(pnh4, Rb, xb, a, 1)}), + t_nh0 = time({nh(pnh, Ra, xa, a0, HC)}), t_nh40 = time({nh(pnh4, Rb, xb, a0, 1)}), + t_up = time({up(mixed, mha)}), t_upn = time({up(nil, mhb)}); + std::printf("T %u, ms per 48 reps (best of 3):\n", T); + std::printf(" mm_combine %.2f, gr_norm_h_moe %.2f, mm_combine + gr_norm_h_moe %.2f -> gr_norm_h_moec %.2f (%+.1f%%)\n", + t_comb, t_nhm, t_pair, t_nhc, 100.0 * (t_nhc / t_pair - 1)); + std::printf(" mm_combine + gr_apply_moe %.2f -> gr_apply_moec %.2f (%+.1f%%)\n", t_am, t_amc, 100.0 * (t_amc / t_am - 1)); + std::printf(" gr_norm_h apply 1 (hc, T) %.2f -> gr_norm_h4 (1, T) %.2f (%+.1f%%)\n", t_nh, t_nh4, 100.0 * (t_nh4 / t_nh - 1)); + std::printf(" gr_norm_h apply 0 (hc, T) %.2f -> gr_norm_h4 (1, T) %.2f (%+.1f%%)\n", t_nh0, t_nh40, 100.0 * (t_nh40 / t_nh0 - 1)); + std::printf(" gemm_hc_up with mixed %.2f -> unbound %.2f (%+.1f%%)\n", t_up, t_upn, 100.0 * (t_upn / t_up - 1)); + } + std::printf("%s\n", fails ? "FAIL" : "ALL PASS"); + return fails ? 1 : 0; +} diff --git a/metal/tests/dense_mm_parity.cpp b/metal/tests/dense_mm_parity.cpp index 33ce4500..969cb790 100644 --- a/metal/tests/dense_mm_parity.cpp +++ b/metal/tests/dense_mm_parity.cpp @@ -3,13 +3,22 @@ // awkward row counts; then the speed on the packs' shapes for a 2048-token chunk. // dense_mm_parity parity // dense_mm_parity --bench speed +// dense_mm_parity --tiles M [shape] every gemm_f16 tile on the prompt chunk's 13 (N, K) shapes at M rows: +// memcmp vs 128_64_4 and the fastest of interleaved rounds; [shape] runs one (its name, +// as printed); DENSE_TILES=a,b,... only those tiles, DENSE_ROUNDS=n rounds +// dense_mm_parity --ab M [shape] the same for the tile prompt_chunk picked before round 35 vs now +// dense_mm_parity --dq-ab [type] [shape] dq_f16 at 64 threads vs dq_tg's; with DQ_AB_KERNELS=, that dir's +// dq_f16 vs this build's at dq_tg's (memcmp + paired timing) #include "mtl_rt.h" #include "ref_quant.h" #include +#include #include #include +#include #include +#include #include #include @@ -24,13 +33,226 @@ int unit_of(int t) { default: return 32; } } +int dq_tg(int u) { return u <= 64 ? 64 : u <= 128 ? 128 : u % 192 == 0 ? 192 : u % 160 == 0 ? 160 : 256; } + +using Tile = std::array; // TM, TN, SG + +// Window::prompt_chunk's gemm tile before round 11 +Tile old_tile(int M, int Nn, int K) { + int tm = (M >= 128 && Nn > 64) ? 128 : 64, tn = Nn <= 64 ? 32 : 64, sg = 4; + if (Nn > 64) { + if (M <= 128 && (Nn <= 640 || K > 2560)) { tm = 64; tn = Nn <= 128 ? 32 : 64; } + else if (M <= 512 && Nn <= 324) { tm = 64; tn = 64; } + else if (M >= 2048 && Nn >= 2560 && K >= 2560) { tm = 64; tn = 128; } + } + return {tm, tn, sg}; +} + +// ... before round 35 +Tile r34_tile(int M, int Nn, int K) { + if (M <= 128) return old_tile(M, Nn, K); + if (Nn <= 64) return M < 1024 ? Tile{16, 64, 2} : Tile{64, 32, 2}; + if (Nn == 96) return M < 768 ? Tile{16, 64, 2} : Tile{64, 32, 2}; + if (Nn <= 128) return M < 768 ? Tile{16, 64, 2} : Tile{32, 64, 2}; + if (Nn <= 640 && K > 2560) return M <= 512 ? Tile{64, 64, 4} : M < 1536 ? Tile{128, 64, 4} : Tile{32, 128, 2}; + if (Nn <= 512 || (Nn <= 640 && M < 384)) return M < 384 ? Tile{32, 64, 2} : Tile{32, 128, 2}; + if (Nn <= 1664) return M < 768 ? Tile{32, 128, 2} : Tile{64, 128, 4}; + return M < 1536 ? Tile{128, 64, 4} : Tile{64, 128, 4}; +} + +// ... and now (window.cpp, kept in step by hand) +Tile new_tile(int M, int Nn, int K) { + if (M <= 128) return old_tile(M, Nn, K); + if (Nn == 96) return M < 640 ? Tile{16, 64, 2} : Tile{16, 96, 2}; + if (Nn > 1664 && M < 1536) return M % 128 && M % 128 <= 64 ? Tile{64, 128, 4} : Tile{128, 64, 4}; + return r34_tile(M, Nn, K); +} + +std::string tile_name(const Tile& t) { + return std::to_string(t[0]) + "_" + std::to_string(t[1]) + "_" + std::to_string(t[2]); +} + +void fill_h(const mtl::Buffer& b, size_t n, uint64_t seed, float scale) { + ref::Rng r(seed); + for (size_t i = 0; i < n; ++i) b.as()[i] = ref::f32_to_f16(r.sym() * scale); +} + +// --tiles: every tile; --ab: r34_tile vs new_tile. Each candidate's Y is compared with 128_64_4's (memcmp), then +// timed over 5 (--ab: 9) rounds that take the candidates in turn +int tile_bench(mtl::Context& cx, int lib, int M, bool ab, const char* only) { + struct Sh { const char* name; int N, K, calls; }; // calls: a 48-layer chunk's (streaming: the read-ahead router) + const std::vector shapes = {{"hc_read", 324, 10240, 96}, {"Wqkv", 10240, 2560, 36}, {"Wz", 6144, 2560, 36}, + {"Wout", 2560, 6144, 48}, {"Wq", 12288, 2560, 12}, {"Wpv", 2560, 2560, 1}, + {"Wshgu", 640, 2560, 96}, {"Wshd", 2560, 640, 48}, {"W512", 512, 2560, 95}, + {"Wk", 1152, 2560, 12}, {"Wk_sel", 1664, 2560, 12}, {"Walbe", 96, 2560, 36}, + {"Wsg", 1, 2560, 48}}; + const std::vector all = {{128, 64, 4}, {64, 64, 4}, {64, 128, 4}, {128, 128, 4}, {64, 32, 4}, {64, 64, 2}, + {64, 128, 2}, {16, 64, 2}, {32, 64, 2}, {32, 128, 2}, {64, 32, 2}, {16, 96, 2}}; + mtl::Buffer A = cx.alloc(2 * (size_t)M * 10240), B = cx.alloc(2 * (size_t)12288 * 2560), + Y = cx.alloc(4 * (size_t)M * 12288), Yref = cx.alloc(4 * (size_t)M * 12288); + fill_h(A, (size_t)M * 10240, 11, 1.0f); + fill_h(B, (size_t)12288 * 2560, 12, 0.05f); + std::map pl; + auto disp = [&](const Tile& tl, const Sh& s, const mtl::Buffer& y) { + const std::string fn = "gemm_f16_" + tile_name(tl); + auto it = pl.find(fn); + if (it == pl.end()) it = pl.emplace(fn, cx.pipeline(lib, fn)).first; + GemmArgs g{M, s.N, s.K, 0}; + mtl::Dispatch d; d.p = &it->second; + d.args = {mtl::Arg::B(A), mtl::Arg::B(B), mtl::Arg::B(y), mtl::Arg::V(g)}; + d.grid[0] = (s.N + tl[1] - 1) / tl[1]; d.grid[1] = (M + tl[0] - 1) / tl[0]; d.tg[0] = 32 * tl[2]; + return d; + }; + int fails = 0; + double tot_old = 0, tot_new = 0; + std::printf("M %d: shape N x K calls | tile us/call (DIFF: not bitwise 128_64_4) | old -> new tile\n", M); + for (const Sh& s : shapes) { + if (only && std::string(only) != s.name) continue; + const Tile to = r34_tile(M, s.N, s.K), tn = new_tile(M, s.N, s.K); + std::vector tiles = ab ? std::vector{to, tn} : all; + if (const char* only_t = std::getenv("DENSE_TILES")) { // a comma list of tile names, e.g. 64_32_2,16_96_2 + const std::string want = std::string(",") + only_t + ","; + std::vector keep; + for (const Tile& tl : tiles) + if (want.find("," + tile_name(tl) + ",") != std::string::npos) keep.push_back(tl); + tiles = keep; + } + const size_t ny = (size_t)M * s.N * 4; + std::memset(Yref.ptr, 0xFF, ny); + cx.run({disp({128, 64, 4}, s, Yref)}); + const int reps = std::clamp((int)(1e9 / (2.0 * M * s.N * s.K)), 1, 32); + cx.warm({disp({128, 64, 4}, s, Y)}, &s == &shapes[0] || only ? 0.3 : 0.05); // the first: the clocks up + std::vector cmp(tiles.size()); + for (size_t i = 0; i < tiles.size(); ++i) { + std::memset(Y.ptr, 0xFF, ny); + cx.run({disp(tiles[i], s, Y)}); + cmp[i] = std::memcmp(Y.ptr, Yref.ptr, ny); + fails += cmp[i] != 0; + } + std::vector> ts(tiles.size()); + // the minimum over the rounds: a busy GPU (another process) only adds time + const int rounds = std::getenv("DENSE_ROUNDS") ? std::atoi(std::getenv("DENSE_ROUNDS")) : ab ? 9 : 5; + for (int round = 0; round < rounds; ++round) + for (size_t i = 0; i < tiles.size(); ++i) { + std::vector ds(reps, disp(tiles[i], s, Y)); + ts[i].push_back(cx.run(ds) / reps * 1e6); + } + double uo = 0, un = 0; + std::printf("%-8s %5d x %5d x%3d |", s.name, s.N, s.K, s.calls); + for (size_t i = 0; i < tiles.size(); ++i) { + const double us = *std::min_element(ts[i].begin(), ts[i].end()); + if (tiles[i] == to) uo = us; + if (tiles[i] == tn) un = us; + std::printf(" %s %.1f%s", tile_name(tiles[i]).c_str(), us, cmp[i] ? " DIFF" : ""); + } + std::printf(" | %s -> %s %.2fx, %.2f -> %.2f ms/chunk\n", tile_name(to).c_str(), tile_name(tn).c_str(), uo / un, + uo * s.calls / 1e3, un * s.calls / 1e3); + tot_old += uo * s.calls; + tot_new += un * s.calls; + } + std::printf("M %d: old %.2f ms, new %.2f ms per chunk (%.2fx); %s\n", M, tot_old / 1e3, tot_new / 1e3, + tot_old / tot_new, fails ? "SOME TILES DIFFER" : "ALL BITWISE"); + return fails ? 1 : 0; +} + +int dq_ab_bench(mtl::Context& cx, int lib, int lib_a, const char* only_ty, const char* only_sh) { + struct DqSh { const char* name; int n_in, n_out; uint32_t perm; int calls; }; + const std::vector shapes = { + {"2560x12288", 2560, 12288, 0, 12}, {"2560x10240", 2560, 10240, 0, 37}, {"320x10240p4", 320, 10240, 4, 96}, + {"2560x6144", 2560, 6144, 0, 36}, {"6144x2560", 6144, 2560, 0, 48}, {"2560x640", 2560, 640, 0, 96}, + {"640x2560", 640, 2560, 0, 48}, {"2560x512", 2560, 512, 0, 119}, {"2560x128", 2560, 128, 0, 12}, + {"2560x48", 2560, 48, 0, 72}, {"10240x320", 10240, 320, 0, 96}, {"10240x4", 10240, 4, 0, 96}, + {"2560x1", 2560, 1, 0, 48}, + }; + const std::vector types = {ref::F32, ref::BF16, ref::Q8_0, ref::Q4_0, ref::Q5_0, ref::Q5_1, ref::IQ4_NL, + ref::IQ4_XS, ref::Q2_0, ref::IQ3_S, ref::Q3_K, ref::Q4_K, ref::Q5_K, ref::Q6_K}; + const size_t pool_w = (size_t)256 << 20, max_wh = (size_t)2 * 2560 * 12288; + mtl::Buffer W = cx.alloc(pool_w), wh = cx.alloc(max_wh), wh_ref = cx.alloc(max_wh); + int fails = 0; + for (int ty : types) { + if (only_ty && std::string(only_ty) != ref::type_name(ty)) continue; + mtl::Pipeline pdq = cx.pipeline(lib, "dq_f16_" + std::to_string(ty)); + mtl::Pipeline pdq_a = cx.pipeline(lib_a, "dq_f16_" + std::to_string(ty)); + const int unit = unit_of(ty); + double tot64 = 0, tot_sel = 0; + for (const auto& s : shapes) { + if (only_sh && std::string(only_sh) != s.name) continue; + if (s.n_in % 256 && ref::block_elems(ty) == 256) continue; + const size_t rb = ref::row_bytes(ty, s.n_in), mat_w = rb * (size_t)s.n_out; + const size_t stride_w = (mat_w + 16383) / 16384 * 16384, nbytes = (size_t)2 * s.n_in * s.n_out; + const int nslots = std::max(1, std::min(16, (int)(pool_w / stride_w))), rinit = std::min(s.n_out, 64); + ref::random_blocks(ty, W.as(), rinit, s.n_in, 17 + ty + s.n_in, 0.02f); + for (int row = rinit; row < s.n_out; row += rinit) + std::memcpy(W.as() + (size_t)row * rb, W.as(), (size_t)std::min(rinit, s.n_out - row) * rb); + for (int sl = 1; sl < nslots; ++sl) std::memcpy(W.as() + (size_t)sl * stride_w, W.as(), mat_w); + const int u = s.n_in / unit, tgs[2] = {lib_a != lib ? dq_tg(u) : 64, dq_tg(u)}; + DqArgs a{(uint32_t)s.n_in, (uint32_t)s.n_out, (uint32_t)rb, s.perm}; + auto disp = [&](int i, size_t woff, const mtl::Buffer& out) { + const int tg = tgs[i]; + mtl::Dispatch d; d.p = i ? &pdq : &pdq_a; d.args = {mtl::Arg::B(W, woff), mtl::Arg::B(out), mtl::Arg::V(a)}; + d.grid[0] = (u + tg - 1) / tg; d.grid[1] = s.n_out; d.tg[0] = tg; d.barrier = true; return d; + }; + std::memset(wh_ref.ptr, 0xAA, nbytes); + std::memset(wh.ptr, 0x55, nbytes); + cx.run({disp(0, 0, wh_ref)}); + cx.run({disp(1, 0, wh)}); + const int cmp = std::memcmp(wh.ptr, wh_ref.ptr, nbytes); + fails += (cmp != 0); + const int reps = std::clamp((int)(4e7 / ((double)s.n_in * s.n_out)), 4, 16); + cx.warm({disp(0, 0, wh), disp(1, 0, wh)}, (&s == &shapes[0] || only_sh) ? 0.3 : 0.02); + std::vector ts[2]; + double best_pair = 1e30; + for (int round = 0; round < 13; ++round) { + double v[2]; + for (int step = 0; step < 2; ++step) { + const int i = (round & 1) ? (1 - step) : step; + std::vector ds; + for (int r = 0; r < reps; ++r) ds.push_back(disp(i, (size_t)(r % nslots) * stride_w, wh)); + v[i] = cx.run_graph(ds) / reps * 1e6; + } + best_pair = std::min(best_pair, std::min(v[0], v[1])); + ts[0].push_back(v[0]); + ts[1].push_back(v[1]); + } + std::vector f0, f1; + for (size_t i = 0; i < ts[0].size(); ++i) { + if (std::max(ts[0][i], ts[1][i]) <= 1.15 * best_pair) { + f0.push_back(ts[0][i]); + f1.push_back(ts[1][i]); + } + } + if (!f0.empty()) { ts[0] = f0; ts[1] = f1; } + const double u0 = median(ts[0]), u1 = median(ts[1]); + tot64 += u0 * s.calls / 1e3; tot_sel += u1 * s.calls / 1e3; + std::printf("%-7s %-13s x%3d (u=%3d, tg %3d->%3d): %6.1f -> %6.1f us (%.2fx) memcmp=%d\n", + ref::type_name(ty), s.name, s.calls, u, tgs[0], tgs[1], u0, u1, u0 / u1, cmp); + } + std::printf("CHUNK %-7s: %.2f -> %.2f ms (%.2fx)\n", ref::type_name(ty), tot64, tot_sel, tot64 / tot_sel); + } + return fails ? 1 : 0; +} + } // namespace int main(int argc, char** argv) { + if (argc > 1 && std::string(argv[1]) == "--dq-ab") { + mtl::Context cx; + const int lib = cx.compile({"quant.metal", "units.metal", "mma.metal", "dense_mm.metal"}); + const char* kd = std::getenv("DQ_AB_KERNELS"); + if (!kd) kd = std::getenv("STRATA_DQ_AB_KERNELS"); + int argi = 2; + if (argc > argi && (std::strchr(argv[argi], '/') || std::strchr(argv[argi], '.'))) kd = argv[argi++]; + const std::string d = kd ? kd : ""; + const int lib_a = kd ? cx.compile({d + "/quant.metal", d + "/units.metal", d + "/mma.metal", d + "/dense_mm.metal"}) : lib; + return dq_ab_bench(cx, lib, lib_a, argc > argi ? argv[argi] : nullptr, argc > argi + 1 ? argv[argi + 1] : nullptr); + } const bool bench = argc > 1 && std::string(argv[1]) == "--bench"; mtl::Context cx; const int lib = cx.compile({"quant.metal", "units.metal", "mma.metal", "dense_mm.metal"}); std::printf("device %s\n", cx.device_name().c_str()); + const std::string mode = argc > 1 ? argv[1] : ""; + if (mode == "--tiles" || mode == "--ab") + return tile_bench(cx, lib, argc > 2 ? std::atoi(argv[2]) : 1983, mode == "--ab", argc > 3 ? argv[3] : nullptr); mtl::Pipeline pact = cx.pipeline(lib, "act_f16"); auto run_all = [&](int ty, const mtl::Buffer& W, size_t woff, int n_in, int n_out, int T, const mtl::Buffer& x, const mtl::Buffer& xh, const mtl::Buffer& wh, const mtl::Buffer& y, const std::string& gk, @@ -45,12 +267,13 @@ int main(int argc, char** argv) { if (which & 1) { const uint32_t n = (uint32_t)T * n_in; mtl::Dispatch d; d.p = &pact; d.args = {mtl::Arg::B(x), mtl::Arg::B(xh), mtl::Arg::V(n)}; - d.grid[0] = (n + 255) / 256; d.tg[0] = 256; ds.push_back(d); + d.grid[0] = (n + 1023) / 1024; d.tg[0] = 256; ds.push_back(d); } if (which & 2) { + const int u = n_in / unit_of(ty), tg = dq_tg(u); DqArgs a{(uint32_t)n_in, (uint32_t)n_out, (uint32_t)ref::row_bytes(ty, n_in)}; mtl::Dispatch d; d.p = pdq; d.args = {mtl::Arg::B(W, woff), mtl::Arg::B(wh), mtl::Arg::V(a)}; - d.grid[0] = (n_in / unit_of(ty) + 63) / 64; d.grid[1] = n_out; d.tg[0] = 64; ds.push_back(d); + d.grid[0] = (u + tg - 1) / tg; d.grid[1] = n_out; d.tg[0] = tg; ds.push_back(d); } if (which & 4) { GemmArgs g{T, n_out, n_in, 0}; @@ -58,8 +281,8 @@ int main(int argc, char** argv) { d.grid[0] = (n_out + tn - 1) / tn; d.grid[1] = (T + tm - 1) / tm; d.tg[0] = 128; ds.push_back(d); } }; - const std::vector types = {ref::F32, ref::BF16, ref::Q8_0, ref::Q4_0, ref::Q5_0, ref::IQ4_NL, ref::IQ4_XS, ref::Q2_0, - ref::IQ3_S, ref::Q3_K, ref::Q4_K, ref::Q5_K, ref::Q6_K}; + const std::vector types = {ref::F32, ref::BF16, ref::Q8_0, ref::Q4_0, ref::Q5_0, ref::Q5_1, ref::IQ4_NL, ref::IQ4_XS, + ref::Q2_0, ref::IQ3_S, ref::Q3_K, ref::Q4_K, ref::Q5_K, ref::Q6_K}; if (!bench) { int fails = 0; for (int ty : types) @@ -78,6 +301,11 @@ int main(int argc, char** argv) { cx.run(ds); std::vector wf((size_t)n_in * n_out); ref::dequant(ty, W.ptr, wf.data(), wf.size()); + std::vector wh_ref(wf.size()), xh_ref((size_t)T * n_in); + for (size_t i = 0; i < wf.size(); ++i) wh_ref[i] = ref::f32_to_f16(wf[i]); + for (size_t i = 0; i < xh_ref.size(); ++i) xh_ref[i] = ref::f32_to_f16(x.as()[i]); + const int wh_cmp = std::memcmp(wh.ptr, wh_ref.data(), wh_ref.size() * 2); + const int xh_cmp = std::memcmp(xh.ptr, xh_ref.data(), xh_ref.size() * 2); double se16 = 0, se = 0, sr = 0; for (int t = 0; t < T; ++t) for (int o = 0; o < n_out; ++o) { @@ -93,47 +321,50 @@ int main(int argc, char** argv) { sr += ae * ae; } const double r16 = std::sqrt(se16 / sr), re = std::sqrt(se / sr); - const bool ok = r16 < 1e-5; + const bool ok = r16 < 1e-5 && wh_cmp == 0 && xh_cmp == 0; fails += !ok; - std::printf("dense mm %-7s %5d -> %4d, T %d: vs fp16 operands %.2e, vs exact %.2e %s\n", - ref::type_name(ty), n_in, n_out, T, r16, re, ok ? "PASS" : "FAIL"); + std::printf("dense mm %-7s %5d -> %4d, T %d: vs fp16 operands %.2e, vs exact %.2e, memcmp(wh)=%d, memcmp(xh)=%d %s\n", + ref::type_name(ty), n_in, n_out, T, r16, re, wh_cmp, xh_cmp, ok ? "PASS" : "FAIL"); } } std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); return fails ? 1 : 0; } - struct Shape { const char* name; int type, n_in, n_out; }; - const std::vector shapes = {{"gdn qkv", ref::IQ4_XS, 2560, 10240}, {"gdn qkv", ref::Q3_K, 2560, 10240}, - {"gdn gate", ref::Q4_K, 2560, 6144}, {"gdn out", ref::Q6_K, 6144, 2560}, - {"qsa q", ref::Q2_0, 2560, 12288}, {"hc down", ref::BF16, 10240, 324}, - {"hc up", ref::BF16, 320, 10240}, {"router", ref::BF16, 2560, 512}, - {"alpha", ref::BF16, 2560, 48}}; - const int T = 2048; - for (const Shape& s : shapes) { - if (s.n_in % 128) { std::printf("%-9s %-6s: n_in %d not a multiple of 128 (needs the 32-chunk path)\n", s.name, ref::type_name(s.type), s.n_in); continue; } - const size_t rb = ref::row_bytes(s.type, s.n_in); - mtl::Buffer W = cx.alloc(rb * s.n_out), x = cx.alloc(4 * (size_t)T * s.n_in), xh = cx.alloc(2 * (size_t)T * s.n_in), - wh = cx.alloc(2 * (size_t)s.n_out * s.n_in), y = cx.alloc(4 * (size_t)T * s.n_out); - ref::random_blocks(s.type, W.as(), s.n_out, s.n_in, 3, 0.02f); - auto timeit = [&](int which, const std::string& gk) { - std::vector ds; - for (int i = 0; i < 4; ++i) run_all(s.type, W, 0, s.n_in, s.n_out, T, x, xh, wh, y, gk, ds, which); - cx.warm(ds, 0.2); - std::vector t; - for (int i = 0; i < 5; ++i) t.push_back(cx.run(ds) / 4); - return median(t); - }; - const double tdq = timeit(2, "gemm_f16_64_64_4"), tact = timeit(1, "gemm_f16_64_64_4"); - double best = 1e9; - std::string bk; - for (const char* gk : {"gemm_f16_64_64_4", "gemm_f16_128_64_4", "gemm_f16_64_128_4", "gemm_f16_64_32_4"}) { - const double tg = timeit(4, gk); - if (tg < best) { best = tg; bk = gk; } + const size_t max_w = (size_t)4 * 2560 * 12288; + const size_t max_wh = (size_t)2 * 2560 * 12288; + mtl::Buffer W = cx.alloc(max_w), wh = cx.alloc(max_wh); + std::vector wf(2560 * 12288); + std::vector wh_ref(2560 * 12288); + struct DqShape { const char* name; int n_in, n_out; }; + const std::vector dq_shapes = {{"2560x10240", 2560, 10240}, {"2560x12288", 2560, 12288}, + {"2560x6144", 2560, 6144}, {"10240x324", 10240, 324}}; + int dq_fails = 0; + for (const auto& ds_sh : dq_shapes) { + std::printf("--- dq_f16 all 13 formats on %s (%d x %d): bitwise memcmp + speed ---\n", ds_sh.name, ds_sh.n_in, ds_sh.n_out); + for (int ty : types) { + const size_t rb = ref::row_bytes(ty, ds_sh.n_in); + ref::random_blocks(ty, W.as(), std::min(ds_sh.n_out, 64), ds_sh.n_in, 17 + ty, 0.02f); + for (int row = 64; row < ds_sh.n_out; row += 64) { + const int ncopy = std::min(64, ds_sh.n_out - row); + std::memcpy(W.as() + (size_t)row * rb, W.as(), (size_t)ncopy * rb); + } + mtl::Pipeline pdq = cx.pipeline(lib, "dq_f16_" + std::to_string(ty)); + const int u = ds_sh.n_in / unit_of(ty), tg = dq_tg(u); + DqArgs a{(uint32_t)ds_sh.n_in, (uint32_t)ds_sh.n_out, (uint32_t)rb, 0}; + mtl::Dispatch d; d.p = &pdq; d.args = {mtl::Arg::B(W), mtl::Arg::B(wh), mtl::Arg::V(a)}; + d.grid[0] = (u + tg - 1) / tg; d.grid[1] = ds_sh.n_out; d.tg[0] = tg; d.barrier = true; + std::vector ds(4, d); + cx.run_graph(ds); + const size_t nelem = (size_t)ds_sh.n_in * ds_sh.n_out; + ref::dequant(ty, W.ptr, wf.data(), nelem); + for (size_t i = 0; i < nelem; ++i) wh_ref[i] = ref::f32_to_f16(wf[i]); + const int cmp = std::memcmp(wh.ptr, wh_ref.data(), nelem * 2); + dq_fails += (cmp != 0); + std::vector ts; + for (int i = 0; i < 5; ++i) ts.push_back(cx.run_graph(ds) / 4); + std::printf("dq_f16 %-7s %-10s: %6.1f us, memcmp=%d %s\n", + ref::type_name(ty), ds_sh.name, median(ts) * 1e6, cmp, cmp == 0 ? "PASS" : "FAIL"); } - const double fl = 2.0 * T * s.n_in * s.n_out; - std::printf("%-9s %-6s %5d -> %5d, T %d: dequant %7.1f us, act %6.1f us, GEMM %7.1f us (%5.1f TFLOP/s, %s) -> %.2f us a token\n", - s.name, ref::type_name(s.type), s.n_in, s.n_out, T, tdq * 1e6, tact * 1e6, best * 1e6, fl / best / 1e12, - bk.c_str(), (tdq + best) / T * 1e6); } - return 0; + return dq_fails ? 1 : 0; } diff --git a/metal/tests/expert_parity.cpp b/metal/tests/expert_parity.cpp index 3bd15ba3..10d082b7 100644 --- a/metal/tests/expert_parity.cpp +++ b/metal/tests/expert_parity.cpp @@ -2,6 +2,7 @@ // reference, and their speed. // // expert_parity [--gu Q2_0] [--dn Q2_0] [--ntok 1..8] [--gucfg L,R,S] [--dncfg L,R,S] parity (synthetic) +// [--same] every token routed to the same 10 experts (each group has ntok entries) // expert_parity --bench [--gu ..] [--dn ..] [--ntok N] [--layers 8] [--reps 96] [--sweep] speed // // Parity: random valid blocks for 64 experts, random activations x ~ U(-1,1), random distinct top-10 per token, @@ -69,7 +70,7 @@ Layout layout(int gu, int dn) { } void random_routing(ref::Rng& r, int ntok, int n_expert, std::vector& ids, std::vector& w, - double share = 0.0, const std::vector* prev = nullptr) { + double share = 0.0, const std::vector* prev = nullptr, bool same = false) { ids.assign(ntok * TOPK, 0); w.assign(ntok * TOPK, 0.0f); for (int t = 0; t < ntok; ++t) { @@ -78,6 +79,7 @@ void random_routing(ref::Rng& r, int ntok, int n_expert, std::vector& ids, for (int k = 0; k < TOPK; ++k) { int e; do { + if (same && t > 0) { e = ids[(k + t) % TOPK]; break; } // token 0's experts, rotated // with probability `share`, reuse an expert of the previous token (verify windows share experts) if (t > 0 && share > 0 && r.uni() < share) e = ids[(t - 1) * TOPK + (r.next() % TOPK)]; else e = (int)(r.next() % n_expert); @@ -192,7 +194,9 @@ int main(int argc, char** argv) { Cfg cgu{32, 2, 4}, cdn{8, 2, 4}; bool bench = false, sweep = false; double share = 0.0; + bool same = false; bool cfg_gu_set = false, cfg_dn_set = false; + std::string ab_ref; for (int i = 1; i < argc; ++i) { std::string a = argv[i]; auto nx = [&] { return i + 1 < argc ? argv[++i] : (std::fprintf(stderr, "%s needs a value\n", a.c_str()), std::exit(2), (char*)nullptr); }; @@ -207,6 +211,8 @@ int main(int argc, char** argv) { else if (a == "--reps") reps = std::atoi(nx()); else if (a == "--share") share = std::atof(nx()); else if (a == "--dn-prep") g_dn_prep = true; + else if (a == "--same") same = true; + else if (a == "--ab") ab_ref = nx(); else { std::fprintf(stderr, "unknown argument %s\n", a.c_str()); return 2; } } (void)cfg_gu_set; @@ -214,21 +220,33 @@ int main(int argc, char** argv) { if (dn == ref::IQ4_NL && cdn.L == 8 && !cfg_dn_set) cdn = {4, 2, 4}; mtl::Context cx; const int lib = cx.compile({"quant.metal", "units.metal", "moe.metal"}); + int ref_lib = -1; + if (!ab_ref.empty()) { + if (ab_ref.size() >= 6 && ab_ref.substr(ab_ref.size() - 6) == ".metal") { + ref_lib = cx.compile({"quant.metal", "units.metal", ab_ref}); + } else { + const std::string cur = cx.kernel_dir; + cx.kernel_dir = ab_ref; + ref_lib = cx.compile({"quant.metal", "units.metal", "moe.metal"}); + cx.kernel_dir = cur; + } + } const Layout L = layout(gu, dn); std::printf("device %s; gate/up %s, down %s, blob %zu bytes; %d token(s)\n", cx.device_name().c_str(), ref::type_name(gu), ref::type_name(dn), L.blob, ntok); - auto make = [&](Cfg a, Cfg b) { + auto make_from = [&](int l, Cfg a, Cfg b) { Kernels K; - K.plan = cx.pipeline(lib, "moe_plan"); - K.prep = cx.pipeline(lib, "xprep"); - K.comb = cx.pipeline(lib, "moe_combine"); - K.gu = cx.pipeline(lib, kname("moe_gu", gu, a), mtl::Constants().i(22, std::getenv("MOE_PROBE") ? std::atoi(std::getenv("MOE_PROBE")) : 0).b(23, std::getenv("QUANT_CTAB") != nullptr)); - K.dn = cx.pipeline(lib, kname("moe_down", dn == ref::Q2_0 && !g_dn_prep ? 142 : dn, b)); + K.plan = cx.pipeline(l, "moe_plan"); + K.prep = cx.pipeline(l, "xprep"); + K.comb = cx.pipeline(l, "moe_combine"); + K.gu = cx.pipeline(l, kname("moe_gu", gu, a), mtl::Constants().i(22, std::getenv("MOE_PROBE") ? std::atoi(std::getenv("MOE_PROBE")) : 0).b(23, std::getenv("QUANT_CTAB") != nullptr)); + K.dn = cx.pipeline(l, kname("moe_down", dn == ref::Q2_0 && !g_dn_prep ? 142 : dn, b)); K.cgu = a; K.cdn = b; return K; }; + auto make = [&](Cfg a, Cfg b) { return make_from(lib, a, b); }; auto args_for = [&](const Kernels& K, int nt, int n_expert) { MoeArgs a{}; a.gate_off = 0; @@ -262,6 +280,117 @@ int main(int argc, char** argv) { cx.run({prep_dispatch(pp, x, XPX, N_EMBD, 8)}); } + if (!ab_ref.empty()) { + const int NE_AB = 64; + mtl::Buffer W = cx.alloc(L.blob * NE_AB); + for (int e = 0; e < NE_AB; ++e) { + uint8_t* b = W.as() + e * L.blob; + ref::random_blocks(gu, b, 2 * N_FF, N_EMBD, 1000 + e); + ref::random_blocks(dn, b + L.down_off, N_EMBD, N_FF, 5000 + e); + } + mtl::Buffer h_r = cx.alloc(h.bytes), ye_r = cx.alloc(ye.bytes), y_r = cx.alloc(y.bytes), plan_r = cx.alloc(4 * P_INTS); + mtl::Buffer bid = cx.alloc(4 * 8 * TOPK), bw = cx.alloc(4 * 8 * TOPK), plan = cx.alloc(4 * P_INTS); + Kernels Kc = make_from(lib, cgu, cdn), Kr = make_from(ref_lib, cgu, cdn); + int ab_fails = 0; + for (bool sm : {true, false}) { + for (int nt : {1, 2, 3, 4, 8}) { + std::vector ids; + std::vector w; + random_routing(rng, nt, NE_AB, ids, w, sm ? 0.0 : 0.35, nullptr, sm); + std::memcpy(bid.ptr, ids.data(), 4 * ids.size()); + std::memcpy(bw.ptr, w.data(), 4 * w.size()); + MoeArgs a = args_for(Kc, nt, 512); + std::memset(h.ptr, 0xa5, h.bytes); std::memset(ye.ptr, 0xa5, ye.bytes); + std::memset(y.ptr, 0xa5, y.bytes); std::memset(plan.ptr, 0xa5, plan.bytes); + std::memset(h_r.ptr, 0xa5, h_r.bytes); std::memset(ye_r.ptr, 0xa5, ye_r.bytes); + std::memset(y_r.ptr, 0xa5, y_r.bytes); std::memset(plan_r.ptr, 0xa5, plan_r.bytes); + std::vector dsc, dsr; + encode_layer(dsc, Kc, W, 0, x, h, ye, y, bid, bw, plan, 0, a, true, true, true, true); + encode_layer(dsr, Kr, W, 0, x, h_r, ye_r, y_r, bid, bw, plan_r, 0, a, true, true, true, true); + cx.run(dsc); + cx.run(dsr); + const int ne_tot = nt * TOPK; + if (std::memcmp(plan.ptr, plan_r.ptr, plan.bytes) != 0 || + std::memcmp(h.ptr, h_r.ptr, 4 * ne_tot * N_FF) != 0 || + std::memcmp(ye.ptr, ye_r.ptr, 4 * ne_tot * N_EMBD) != 0 || + std::memcmp(y.ptr, y_r.ptr, 4 * nt * N_EMBD) != 0) { + ++ab_fails; + std::printf("AB MISMATCH same=%d ntok=%d\n", (int)sm, nt); + } + } + } + std::printf("expert_parity --ab bitwise (%s/%s%s gu %d,%d,%d dn %d,%d,%d, NE=1..4,8 + random): %d fail(s) %s\n", + ref::type_name(gu), ref::type_name(dn), g_dn_prep ? "(prep)" : "(142)", + cgu.L, cgu.R, cgu.S, cdn.L, cdn.R, cdn.S, ab_fails, ab_fails == 0 ? "ALL PASS" : "FAIL"); + if (ab_fails || !bench) return ab_fails ? 1 : 0; + + // Paired A/B benchmark across isolated NE=1,2,3,4 (same=true) and verify windows ntok=1,2,4,8 (share=0.3) + const int NE_B = 512; + const size_t layer_bytes = L.blob * NE_B; + mtl::Buffer Wb = cx.alloc(layer_bytes * layers); + for (size_t e = 0; e < (size_t)NE_B * layers; ++e) + std::memcpy(Wb.as() + e * L.blob, W.as() + (e % NE_AB) * L.blob, L.blob); + mtl::Buffer bid_b = cx.alloc(4 * 8 * TOPK * reps), bw_b = cx.alloc(4 * 8 * TOPK * reps), plan_b = cx.alloc(4 * P_INTS * reps); + auto run_ab_case = [&](int nt, bool sm, double sh, const char* tag) { + double groups = 0; + for (int r = 0; r < reps; ++r) { + std::vector ids; std::vector w; + random_routing(rng, nt, NE_B, ids, w, sh, nullptr, sm); + std::memcpy(bid_b.as() + r * 8 * TOPK, ids.data(), 4 * ids.size()); + std::memcpy(bw_b.as() + r * 8 * TOPK, w.data(), 4 * w.size()); + groups += std::set(ids.begin(), ids.end()).size(); + } + groups /= reps; + const MoeArgs a = args_for(Kc, nt, NE_B); + std::vector dp, dgc, dgr, ddc, ddr; + for (int r = 0; r < reps; ++r) { + const size_t woff = (size_t)(r % layers) * layer_bytes; + const size_t poff = (size_t)r * P_INTS * 4; + std::vector p1, gc, gr, dc, dr; + encode_layer(p1, Kc, Wb, woff, x, h, ye, y, bid_b, bw_b, plan_b, poff, a, true, false, false, false); + p1[0].args[0] = mtl::Arg::B(bid_b, (size_t)r * 8 * TOPK * 4); + p1[0].args[1] = mtl::Arg::B(bw_b, (size_t)r * 8 * TOPK * 4); + dp.push_back(p1[0]); + encode_layer(gc, Kc, Wb, woff, x, h, ye, y, bid_b, bw_b, plan_b, poff, a, false, true, false, false); + encode_layer(gr, Kr, Wb, woff, x, h, ye, y, bid_b, bw_b, plan_b, poff, a, false, true, false, false); + encode_layer(dc, Kc, Wb, woff, x, h, ye, y, bid_b, bw_b, plan_b, poff, a, false, false, true, false); + encode_layer(dr, Kr, Wb, woff, x, h, ye, y, bid_b, bw_b, plan_b, poff, a, false, false, true, false); + dgc.push_back(gc[0]); dgr.push_back(gr[0]); + for (auto& z : dc) ddc.push_back(z); + for (auto& z : dr) ddr.push_back(z); + } + cx.run(dp); + cx.run(dgc); + cx.warm(ddr, 0.04); cx.warm(ddc, 0.04); + std::vector b_dn, c_dn, r_dn, b_gu, c_gu, r_gu; + for (int it = 0; it < 7; ++it) { + double bd, cd, bg, cg; + if (it & 1) { + cd = cx.run(ddc) / reps; bd = cx.run(ddr) / reps; + cg = cx.run(dgc) / reps; bg = cx.run(dgr) / reps; + } else { + bd = cx.run(ddr) / reps; cd = cx.run(ddc) / reps; + bg = cx.run(dgr) / reps; cg = cx.run(dgc) / reps; + } + b_dn.push_back(bd * 1e6); c_dn.push_back(cd * 1e6); r_dn.push_back(cd / bd); + b_gu.push_back(bg * 1e6); c_gu.push_back(cg * 1e6); r_gu.push_back(cg / bg); + } + const double bd_bytes = groups * N_EMBD * L.d_row; + std::printf("ab %-14s groups %4.1f | dn base %6.1f us cand %6.1f us (%5.1f GB/s) ratio %.4f | gu base %6.1f us cand %6.1f us ratio %.4f\n", + tag, groups, median(b_dn), median(c_dn), bd_bytes / (median(c_dn) * 1e-6) / 1e9, median(r_dn), + median(b_gu), median(c_gu), median(r_gu)); + }; + for (int ne : {1, 2, 3, 4}) { + char buf[32]; std::snprintf(buf, sizeof(buf), "NE=%d (same)", ne); + run_ab_case(ne, true, 0.0, buf); + } + for (int nt : {1, 2, 4, 8}) { + char buf[32]; std::snprintf(buf, sizeof(buf), "ntok=%d sh=0.3", nt); + run_ab_case(nt, false, 0.3, buf); + } + return 0; + } + if (!bench) { // ---------------- parity ---------------- const int NE = 64; @@ -273,7 +402,7 @@ int main(int argc, char** argv) { } std::vector ids; std::vector w; - random_routing(rng, ntok, NE, ids, w, 0.3); + random_routing(rng, ntok, NE, ids, w, 0.3, nullptr, same); mtl::Buffer bid = cx.alloc(4 * ids.size()), bw = cx.alloc(4 * w.size()), plan = cx.alloc(4 * P_INTS); std::memcpy(bid.ptr, ids.data(), 4 * ids.size()); std::memcpy(bw.ptr, w.data(), 4 * w.size()); @@ -364,7 +493,7 @@ int main(int argc, char** argv) { auto set_routing = [&](int nt) { double groups = 0; for (int r = 0; r < reps; ++r) { - random_routing(rng, nt, NE, rids[r], rw[r], share); + random_routing(rng, nt, NE, rids[r], rw[r], share, nullptr, same); std::memcpy(bid.as() + r * 8 * TOPK, rids[r].data(), 4 * rids[r].size()); std::memcpy(bw.as() + r * 8 * TOPK, rw[r].data(), 4 * rw[r].size()); groups += std::set(rids[r].begin(), rids[r].end()).size(); diff --git a/metal/tests/fuse_parity.cpp b/metal/tests/fuse_parity.cpp index e67d9e58..6c6fe1e6 100644 --- a/metal/tests/fuse_parity.cpp +++ b/metal/tests/fuse_parity.cpp @@ -1,14 +1,18 @@ -// metal/tests/fuse_parity.cpp - the decode latency fusions against the kernels they replace, BITWISE, for 1..4 +// metal/tests/fuse_parity.cpp - the decode latency fusions against the kernels they replace, BITWISE, for 1..8 // tokens: // gemv_cat (gemv.metal) ssm_alpha + ssm_beta (BF16 2560 -> 48 + 48) in one dispatch == two gemv dispatches // router_route the router GEMV + moe_route in one dispatch == gemv + moe_route (logits, ids, w, plan; // twice in a row: the counter resets) // gr_norm_comb (gr2.metal) the MoE combine folded into the next gr_norm == moe_combine_sh + gr_norm (R', xn) +// --gr-unroll BASE_DIR NTOK [bench]: gr.metal / gr2.metal's kernels against the copies in BASE_DIR (the same entry +// points, or gr_down's variants against the entry point they replace, compiled from the old source), bitwise, and +// paired timings of the decode window's gr_down / gr_up #include "mtl_rt.h" #include "ref_quant.h" #include #include +#include #include #include @@ -18,12 +22,292 @@ struct MoeArgs { uint64_t gate_off, up_off, down_off, gu_stride, d_stride; uint32_t gu_row, d_row, n_embd, n_ff, n_tok, topk, n_expert, gx_gu, gx_dn, pad; }; -constexpr int P_INTS = 16 + 128 + 128 + 4 + 128 * 3, P_TOKENT = 16 + 128 + 128 + 4 + 128 * 2; +constexpr int P_INTS = 16 + 128 + 128 + 4 + 128 * 3 + 20 + 128, P_TOKENT = 16 + 128 + 128 + 4 + 128 * 2; struct GrArgs { uint32_t n_embd, hc, hc_lr, n_tok, apply; float eps; }; struct GlueArgs { uint32_t n, n_tok, n_embd, topk; }; +struct GrArgs3 { uint32_t n_embd, hc, hc_lr, n_tok, apply; float eps; uint32_t rb_down, rb_up, comb, topk; }; + +// The A/B source holds baseline and candidate entry points in one library. +int gr_ab(const char* source, int nt, bool down) { + if (nt < 1 || nt > 8) return 2; + mtl::Context cx; + std::printf("device %s\n", cx.device_name().c_str()); + const int lib = down ? cx.compile({"quant.metal", "units.metal", source}) + : cx.compile({source}, "// NO_GGML_COMMON"); + constexpr int N = 2560, HC = 4, LR = 320, TK = 10; + ref::Rng rng(1802); + auto floats = [&](size_t n) { + auto b = cx.alloc(4 * n); + for (size_t i = 0; i < n; ++i) b.as()[i] = rng.sym(); + return b; + }; + auto R0 = floats(nt * HC * N), R1 = floats(nt * HC * N), R2 = floats(nt * HC * N); + auto x1 = floats(nt * HC * N), x2 = floats(nt * HC * N); + GrArgs a{N, HC, LR, (uint32_t)nt, 1, 1e-6f}; + int fails = 0; + auto bench = [&](mtl::Dispatch b, mtl::Dispatch v, const char* name) { + cx.warm({b}, 0.3); cx.warm({v}, 0.1); + for (int pair = 0; pair < 5; ++pair) { + double tb, tv; + if (pair % 2) { tv = cx.run({v}, 64); tb = cx.run({b}, 64); } + else { tb = cx.run({b}, 64); tv = cx.run({v}, 64); } + std::printf("bench %s ntok %d pair %d base_us %.6f variant_us %.6f\n", name, nt, pair, + tb * 1e6 / 64, tv * 1e6 / 64); + } + }; + if (down) { + const size_t rb = ref::row_bytes(ref::Q8_0, HC * N); + auto wd = cx.alloc(rb * LR), wi = cx.alloc(rb * HC); + ref::random_blocks(ref::Q8_0, wd.as(), LR, HC * N, 1803, 0.05f); + ref::random_blocks(ref::Q8_0, wi.as(), HC, HC * N, 1804, 0.05f); + auto lo1 = floats(nt * LR), lo2 = floats(nt * LR), in1 = floats(nt * HC), in2 = floats(nt * HC); + auto pb = cx.pipeline(lib, "baseline_down", mtl::Constants().i(1, nt)); + auto pv = cx.pipeline(lib, "optimized_down", mtl::Constants().i(1, nt)); + mtl::Dispatch b, v; + b.p = &pb; v.p = &pv; + b.args = {mtl::Arg::B(wd), mtl::Arg::B(wi), mtl::Arg::B(R0), mtl::Arg::B(lo1), mtl::Arg::B(in1), mtl::Arg::V(a)}; + v.args = {mtl::Arg::B(wd), mtl::Arg::B(wi), mtl::Arg::B(R0), mtl::Arg::B(lo2), mtl::Arg::B(in2), mtl::Arg::V(a)}; + b.grid[0] = v.grid[0] = (LR + HC) / 2; + b.tg[0] = v.tg[0] = 256; + cx.run({b}); cx.run({v}); + const bool ok = std::memcmp(lo1.ptr, lo2.ptr, lo1.bytes) == 0 && + std::memcmp(in1.ptr, in2.ptr, in1.bytes) == 0; + fails += !ok; + std::printf("down ntok %d bitwise %s\n", nt, ok ? "PASS" : "FAIL"); + bench(b, v, "down"); + } else { + auto ye = floats(nt * TK * N), ysh = floats(nt * N), sgv = floats(nt), inj = floats(nt * HC), wn = floats(HC * N); + auto plan = cx.alloc(4 * P_INTS); + for (int i = 0; i < nt * TK; ++i) plan.as()[P_TOKENT + i] = (i * 7) % (nt * TK); + auto pb = cx.pipeline(lib, "baseline_4s"); + for (const char* name : {"gr_norm_comb_4s", "gr_norm_comb", "gr_norm_comb_staged"}) { + auto pv = cx.pipeline(lib, name); + for (uint32_t tk : {1u, 10u}) for (uint32_t apply : {0u, 1u}) { + a.apply = apply; + mtl::Dispatch b, v; + auto args = [&](const mtl::Buffer& R, const mtl::Buffer& x) { + return std::vector{mtl::Arg::B(R), mtl::Arg::B(ye), mtl::Arg::B(inj), mtl::Arg::B(wn), + mtl::Arg::B(x), mtl::Arg::V(a), mtl::Arg::B(plan), mtl::Arg::B(ysh), mtl::Arg::B(sgv), mtl::Arg::V(tk)}; + }; + b.p = &pb; v.p = &pv; b.args = args(R1, x1); v.args = args(R2, x2); + b.grid[0] = 1; v.grid[0] = std::strcmp(name, "gr_norm_comb") == 0 ? HC : 1; + b.grid[1] = v.grid[1] = nt; b.tg[0] = v.tg[0] = 256; + std::memcpy(R1.ptr, R0.ptr, R0.bytes); std::memcpy(R2.ptr, R0.ptr, R0.bytes); + cx.run({b}); cx.run({v}); + const bool ok = std::memcmp(R1.ptr, R2.ptr, R0.bytes) == 0 && + std::memcmp(x1.ptr, x2.ptr, x1.bytes) == 0; + fails += !ok; + std::printf("%s ntok %d topk %u apply %u R/xn bitwise %s\n", + name, nt, tk, apply, ok ? "PASS" : "FAIL"); + if (apply && tk == TK) bench(b, v, name); + } + } + } + return fails ? 1 : 0; +} + +// Each kernel runs from both libraries on the same inputs into separate outputs, which must match byte for byte. +int gr_unroll(const std::string& base, int nt, bool timed) { + if (nt < 1 || nt > 8) return 2; + mtl::Context cx; + std::printf("device %s\n", cx.device_name().c_str()); + const std::vector srcs = {"quant.metal", "units.metal", "gr.metal", "gr2.metal"}; + std::vector old; + for (const auto& s : srcs) old.push_back(base + "/" + s); + const int lb = cx.compile(old), lv = cx.compile(srcs); + constexpr int N = 2560, HC = 4, LR = 320, NR = LR + HC, TK = 10, MAXT = 8; + ref::Rng rng(2101); + auto floats = [&](size_t n, float scale = 1.0f, float bias = 0.0f) { + auto b = cx.alloc(4 * n); + for (size_t i = 0; i < n; ++i) b.as()[i] = bias + scale * rng.sym(); + return b; + }; + auto weights = [&](int type, int rows, int cols, uint64_t seed) { + auto b = cx.alloc(ref::row_bytes(type, cols) * rows); + ref::random_blocks(type, b.as(), rows, cols, seed, 0.05f); + return b; + }; + auto wd8 = weights(ref::Q8_0, LR, HC * N, 11), wi8 = weights(ref::Q8_0, HC, HC * N, 12), + wif = weights(ref::F32, HC, HC * N, 13), wu8 = weights(ref::Q8_0, HC * N, LR, 14), + wd16 = weights(ref::BF16, LR, HC * N, 15), wi16 = weights(ref::BF16, HC, HC * N, 16), + wu16 = weights(ref::BF16, HC * N, LR, 17), wd4 = weights(ref::Q4_0, LR, HC * N, 18), + wd5 = weights(ref::Q5_0, LR, HC * N, 19); + auto R0 = floats((size_t)nt * HC * N), bo = floats((size_t)nt * N), inj = floats(nt * HC, 3.0f), + wn = floats(HC * N, 0.5f, 1.0f), xn = floats((size_t)nt * HC * N), lo = floats(nt * LR), + ye = floats((size_t)MAXT * TK * N), ysh = floats(nt * N), sgv = floats(nt, 3.0f), + Dprev = floats((size_t)nt * NR * HC, 0.3f), rs = floats(nt * HC, 0.1f, 1.0f); + auto plan = cx.alloc(4 * P_INTS); + for (int i = 0; i < nt * TK; ++i) plan.as()[P_TOKENT + i] = (i * 7) % (MAXT * TK); + int fails = 0; + // outs[i]: the baseline's and the variant's copy of output i; `fresh` restores in-place inputs before a run + struct Out { mtl::Buffer b, v; }; + auto outs = [&](std::vector sizes) { + std::vector o; + for (size_t s : sizes) o.push_back({cx.alloc(s), cx.alloc(s)}); + return o; + }; + auto check = [&](const std::string& name, mtl::Dispatch b, mtl::Dispatch v, std::vector& o, + const std::function& fresh, bool bench) { + for (auto& x : o) { std::memset(x.b.ptr, 0, x.b.bytes); std::memset(x.v.ptr, 0, x.v.bytes); } + if (fresh) fresh(); + cx.run({b}); + cx.run({v}); + bool ok = true; + for (auto& x : o) ok = ok && std::memcmp(x.b.ptr, x.v.ptr, x.b.bytes) == 0; + fails += !ok; + std::printf("%s ntok %d bitwise %s\n", name.c_str(), nt, ok ? "PASS" : "FAIL"); + if (!bench) return; + cx.warm({b, v}, 0.2); + // the median pair ratio: a lone mean drifts with the GPU clock and whatever else shares the GPU + constexpr int PAIRS = 16, IT = 128; + double sb = 0, sv = 0; + std::vector ratio; + for (int pair = 0; pair < PAIRS; ++pair) { + double tb, tv; + if (pair % 2) { tv = cx.run({v}, IT); tb = cx.run({b}, IT); } + else { tb = cx.run({b}, IT); tv = cx.run({v}, IT); } + sb += tb; sv += tv; + ratio.push_back(tv / tb); + } + std::sort(ratio.begin(), ratio.end()); + std::printf("bench %s ntok %d mean base_us %.3f variant_us %.3f (%+.1f%%) median %+.1f%%\n", name.c_str(), nt, + sb * 1e6 / (PAIRS * IT), sv * 1e6 / (PAIRS * IT), 100.0 * (sv / sb - 1), + 100.0 * ((ratio[PAIRS / 2 - 1] + ratio[PAIRS / 2]) / 2 - 1)); + }; + const mtl::Constants fc = mtl::Constants().i(1, nt); + // ---- gr_norm (apply 0, 1); no function constant, R written in place + { + auto pb = cx.pipeline(lb, "gr_norm"), pv = cx.pipeline(lv, "gr_norm"); + for (uint32_t apply : {0u, 1u}) { + GrArgs a{N, HC, LR, (uint32_t)nt, apply, 1e-6f}; + auto o = outs({R0.bytes, xn.bytes}); + mtl::Dispatch b, v; + b.p = &pb; v.p = &pv; + b.args = {mtl::Arg::B(o[0].b), mtl::Arg::B(bo), mtl::Arg::B(inj), mtl::Arg::B(wn), mtl::Arg::B(o[1].b), mtl::Arg::V(a)}; + v.args = {mtl::Arg::B(o[0].v), mtl::Arg::B(bo), mtl::Arg::B(inj), mtl::Arg::B(wn), mtl::Arg::B(o[1].v), mtl::Arg::V(a)}; + b.grid[0] = v.grid[0] = HC; b.grid[1] = v.grid[1] = nt; b.tg[0] = v.tg[0] = 256; + check("gr_norm apply " + std::to_string(apply), b, v, o, [&] { + std::memcpy(o[0].b.ptr, R0.ptr, R0.bytes); std::memcpy(o[0].v.ptr, R0.ptr, R0.bytes); }, timed); + } + } + // ---- gr_norm_comb (gr2.metal) + { + auto pb = cx.pipeline(lb, "gr_norm_comb"), pv = cx.pipeline(lv, "gr_norm_comb"); + GrArgs a{N, HC, LR, (uint32_t)nt, 1, 1e-6f}; + const uint32_t tk = TK; + auto o = outs({R0.bytes, xn.bytes}); + auto args = [&](const mtl::Buffer& R, const mtl::Buffer& x) { + return std::vector{mtl::Arg::B(R), mtl::Arg::B(ye), mtl::Arg::B(inj), mtl::Arg::B(wn), mtl::Arg::B(x), + mtl::Arg::V(a), mtl::Arg::B(plan), mtl::Arg::B(ysh), mtl::Arg::B(sgv), mtl::Arg::V(tk)}; + }; + mtl::Dispatch b, v; + b.p = &pb; v.p = &pv; b.args = args(o[0].b, o[1].b); v.args = args(o[0].v, o[1].v); + b.grid[0] = v.grid[0] = HC; b.grid[1] = v.grid[1] = nt; b.tg[0] = v.tg[0] = 256; + check("gr_norm_comb", b, v, o, [&] { + std::memcpy(o[0].b.ptr, R0.ptr, R0.bytes); std::memcpy(o[0].v.ptr, R0.ptr, R0.bytes); }, timed); + } + // ---- gr_down_f: rows = w_down then w_inject, RS = SG / KS simdgroups a threadgroup, NR rows a simdgroup; a + // variant entry point (vname) runs against the baseline's (name) with its own rows per threadgroup + struct Down { const char* name; const char* vname; int rows, vrows, btg, vtg; const mtl::Buffer* wd; const mtl::Buffer* wi; bool bench; }; + for (const Down& dn : {Down{"gr_down_q8_8_4", "gr_down_q8_4_4", 2, 1, 256, 128, &wd8, &wi8, false}, + Down{"gr_down_q8_8_4", "gr_down_q8_4_4_r2", 2, 2, 256, 128, &wd8, &wi8, false}, + Down{"gr_down_q8f_8_4", "gr_down_q8f_4_4", 2, 1, 256, 128, &wd8, &wif, false}, + Down{"gr_down_q8f_8_4", "gr_down_q8f_4_4_r2", 2, 2, 256, 128, &wd8, &wif, false}, + // the decode window's (M5) and the MTP head's entry points, each against its own copy + Down{"gr_down_q8_4_4", "gr_down_q8_4_4", 1, 1, 128, 128, &wd8, &wi8, true}, + Down{"gr_down_q8_4_4_r2", "gr_down_q8_4_4_r2", 2, 2, 128, 128, &wd8, &wi8, true}, + Down{"gr_down_q8f_4_4", "gr_down_q8f_4_4", 1, 1, 128, 128, &wd8, &wif, true}, + Down{"gr_down_q8f_4_4_r2", "gr_down_q8f_4_4_r2", 2, 2, 128, 128, &wd8, &wif, true}, + Down{"gr_down_q8_8_8", "gr_down_q8_8_8", 1, 1, 256, 256, &wd8, &wi8, true}, + Down{"gr_down_q8_8_8_r2", "gr_down_q8_8_8_r2", 2, 2, 256, 256, &wd8, &wi8, true}, + // the dispatch choice, one row a simdgroup against two (gr_read switches from 4 tokens, gr_read_m from 5) + Down{"gr_down_q8_4_4", "gr_down_q8_4_4_r2", 1, 2, 128, 128, &wd8, &wi8, true}, + Down{"gr_down_q8f_4_4", "gr_down_q8f_4_4_r2", 1, 2, 128, 128, &wd8, &wif, true}, + Down{"gr_down_q8_8_8", "gr_down_q8_8_8_r2", 1, 2, 256, 256, &wd8, &wi8, true}, + Down{"gr_down_q8f_8_8", "gr_down_q8f_8_8", 1, 1, 256, 256, &wd8, &wif, false}, + Down{"gr_down_q8f_8_8", "gr_down_q8f_8_8_r2", 1, 2, 256, 256, &wd8, &wif, false}, + Down{"gr_down_8_4", "gr_down_4_4", 2, 1, 256, 128, &wd16, &wi16, false}, + Down{"gr_down_8_4", "gr_down_4_4_r2", 2, 2, 256, 128, &wd16, &wi16, false}, + Down{"gr_down_8_8", "gr_down_8_8", 1, 1, 256, 256, &wd16, &wi16, false}, + Down{"gr_down_8_8", "gr_down_8_8_r2", 1, 2, 256, 256, &wd16, &wi16, false}, + Down{"gr_down_q4_8_8", "gr_down_q4_8_8", 1, 1, 256, 256, &wd4, &wi8, false}, + Down{"gr_down_q5_8_8", "gr_down_q5_8_8", 1, 1, 256, 256, &wd5, &wi8, false}}) { + auto pb = cx.pipeline(lb, dn.name, fc), pv = cx.pipeline(lv, dn.vname, fc); + GrArgs a{N, HC, LR, (uint32_t)nt, 1, 1e-6f}; + auto o = outs({lo.bytes, (size_t)4 * nt * HC}); + mtl::Dispatch b, v; + b.p = &pb; v.p = &pv; + b.args = {mtl::Arg::B(*dn.wd), mtl::Arg::B(*dn.wi), mtl::Arg::B(xn), mtl::Arg::B(o[0].b), mtl::Arg::B(o[1].b), mtl::Arg::V(a)}; + v.args = {mtl::Arg::B(*dn.wd), mtl::Arg::B(*dn.wi), mtl::Arg::B(xn), mtl::Arg::B(o[0].v), mtl::Arg::B(o[1].v), mtl::Arg::V(a)}; + b.grid[0] = (NR + dn.rows - 1) / dn.rows; + v.grid[0] = (NR + dn.vrows - 1) / dn.vrows; + b.tg[0] = dn.btg; + v.tg[0] = dn.vtg; + check(dn.vname, b, v, o, {}, timed && dn.bench); + } + // ---- gr_up_f: the prepared planes have MAXT rows + for (const char* up : {"gr_up_q8", "gr_up"}) { + const bool q8 = std::strcmp(up, "gr_up_q8") == 0; + auto pb = cx.pipeline(lb, up, fc), pv = cx.pipeline(lv, up, fc); + GrArgs a{N, HC, LR, MAXT, 1, 1e-6f}; + auto o = outs({4 * ((size_t)4 * MAXT * N + MAXT * N / 16)}); + const mtl::Buffer& w = q8 ? wu8 : wu16; + mtl::Dispatch b, v; + b.p = &pb; v.p = &pv; + b.args = {mtl::Arg::B(w), mtl::Arg::B(lo), mtl::Arg::B(xn), mtl::Arg::B(o[0].b), mtl::Arg::V(a)}; + v.args = {mtl::Arg::B(w), mtl::Arg::B(lo), mtl::Arg::B(xn), mtl::Arg::B(o[0].v), mtl::Arg::V(a)}; + b.grid[0] = v.grid[0] = N / 16; + b.tg[0] = v.tg[0] = 256; + check(up, b, v, o, {}, timed && q8); + } + // ---- gr_down3 / gr_down3b (apply and comb 0 / 1), gr_up3 + for (int f : {30, 8}) for (int tt = 1; tt <= 3; ++tt) for (int rps = 1; rps <= 2; ++rps) for (int vb = 0; vb < 1 + (f == 30); ++vb) { + const std::string name = (vb ? "gr_down3b_" : "gr_down3_" + std::to_string(f) + "_") + std::to_string(tt) + "_" + std::to_string(rps); + auto pb = cx.pipeline(lb, name, fc), pv = cx.pipeline(lv, name, fc); + for (uint32_t apply : {0u, 1u}) for (uint32_t comb : {0u, 1u}) { + if (comb && !apply) continue; + GrArgs3 a{N, HC, LR, MAXT, apply, 1e-6f, (uint32_t)(f == 8 ? HC * N / 32 * 34 : HC * N * 2), 0, comb, TK}; + auto o = outs({R0.bytes, Dprev.bytes, rs.bytes}); + const mtl::Buffer& wd = f == 8 ? wd8 : wd16; + const mtl::Buffer& wi = f == 8 ? wi8 : wi16; + auto args = [&](const mtl::Buffer& Ro, const mtl::Buffer& D, const mtl::Buffer& r) { + return std::vector{mtl::Arg::B(wd), mtl::Arg::B(wi), mtl::Arg::B(R0), mtl::Arg::B(Ro), mtl::Arg::B(bo), + mtl::Arg::B(Dprev), mtl::Arg::B(wn), mtl::Arg::B(D), mtl::Arg::B(r), mtl::Arg::V(a), mtl::Arg::B(plan), + mtl::Arg::B(ye), mtl::Arg::B(ysh), mtl::Arg::B(sgv)}; + }; + mtl::Dispatch b, v; + b.p = &pb; v.p = &pv; b.args = args(o[0].b, o[1].b, o[2].b); v.args = args(o[0].v, o[1].v, o[2].v); + b.grid[0] = v.grid[0] = (NR + 8 * rps - 1) / (8 * rps); + b.grid[1] = v.grid[1] = HC; + b.tg[0] = v.tg[0] = 256; + check(name + " apply " + std::to_string(apply) + " comb " + std::to_string(comb), b, v, o, {}, false); + } + } + for (int f : {30, 8}) { + const std::string name = "gr_up3_" + std::to_string(f); + auto pb = cx.pipeline(lb, name, fc), pv = cx.pipeline(lv, name, fc); + GrArgs3 a{N, HC, LR, MAXT, 1, 1e-6f, 0, (uint32_t)(f == 8 ? LR / 32 * 34 : LR * 2), 0, TK}; + auto o = outs({4 * ((size_t)4 * MAXT * N + MAXT * N / 16)}); + const mtl::Buffer& w = f == 8 ? wu8 : wu16; + mtl::Dispatch b, v; + b.p = &pb; v.p = &pv; + b.args = {mtl::Arg::B(w), mtl::Arg::B(Dprev), mtl::Arg::B(R0), mtl::Arg::B(o[0].b), mtl::Arg::V(a), mtl::Arg::B(wn), mtl::Arg::B(rs)}; + v.args = {mtl::Arg::B(w), mtl::Arg::B(Dprev), mtl::Arg::B(R0), mtl::Arg::B(o[0].v), mtl::Arg::V(a), mtl::Arg::B(wn), mtl::Arg::B(rs)}; + b.grid[0] = v.grid[0] = N / 16; + b.tg[0] = v.tg[0] = 256; + check(name, b, v, o, {}, false); + } + std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); + return fails ? 1 : 0; +} + } // namespace -int main() { +int main(int argc, char** argv) { + if (argc == 5 && std::string(argv[1]) == "--gr-ab") + return gr_ab(argv[2], std::stoi(argv[3]), std::string(argv[4]) == "down"); + if (argc >= 4 && std::string(argv[1]) == "--gr-unroll") + return gr_unroll(argv[2], std::stoi(argv[3]), argc > 4 && std::string(argv[4]) == "bench"); mtl::Context cx; const int lib = cx.compile({"quant.metal", "units.metal", "moe.metal", "gemv.metal", "router.metal", "gr.metal", "gr2.metal", "glue.metal"}); @@ -37,7 +321,7 @@ int main() { ref::random_blocks(ref::BF16, Wb.as(), NO, N, 2, 0.05f); ref::Rng r(3); for (size_t i = 0; i < xp.bytes / 4; ++i) xp.as()[i] = r.sym(); - for (int nt = 1; nt <= 4; ++nt) { + for (int nt = 1; nt <= MAXT; ++nt) { mtl::Pipeline pg = cx.pipeline(lib, "gemv_30_32_1_4_1", mtl::Constants().i(0, nt)), pc = cx.pipeline(lib, "gemvc_30_32_1_4", mtl::Constants().i(0, nt)); const size_t s16 = (size_t)4 * 4 * MAXT * N; @@ -75,7 +359,7 @@ int main() { w2 = cx.alloc(4 * MAXT * 10), pl1 = cx.alloc(4 * P_INTS), pl2 = cx.alloc(4 * P_INTS), cnt = cx.alloc(16); ref::random_blocks(rty, Wr.as(), NE, N, 9, 0.05f); mtl::Pipeline pr = cx.pipeline(lib, "moe_route"); - for (int nt = 1; nt <= 4; ++nt) { + for (int nt = 1; nt <= MAXT; ++nt) { mtl::Pipeline pg = cx.pipeline(lib, "gemv_" + std::to_string(rty) + "_32_1_4_1", mtl::Constants().i(0, nt)), pf = cx.pipeline(lib, "router_route_" + std::to_string(rty), mtl::Constants().i(0, nt)); const size_t s16 = (size_t)4 * 4 * MAXT * N; @@ -131,12 +415,14 @@ int main() { for (int i = 0; i < MAXT * HC; ++i) inj.as()[i] = r.sym() * 3; for (int i = 0; i < HC * N; ++i) wn.as()[i] = 0.5f + r.uni(); for (int i = 0; i < MAXT * TOPK; ++i) plan.as()[P_TOKENT + i] = (i * 7) % (MAXT * TOPK); - mtl::Pipeline pcm = cx.pipeline(lib, "moe_combine_sh"), pn = cx.pipeline(lib, "gr_norm"), pnc = cx.pipeline(lib, "gr_norm_comb"); - for (int nt = 1; nt <= 4; ++nt) { + mtl::Pipeline pcm = cx.pipeline(lib, "moe_combine_sh"), pn = cx.pipeline(lib, "gr_norm"), pnc = cx.pipeline(lib, "gr_norm_comb"), + pnc4 = cx.pipeline(lib, "gr_norm_comb_4s"); + for (int nt = 1; nt <= MAXT; ++nt) + for (uint32_t apply : {0u, 1u}) for (bool four : {false, true}) { std::memcpy(R1.ptr, R0.ptr, R0.bytes); std::memcpy(R2.ptr, R0.ptr, R0.bytes); GlueArgs ga{(uint32_t)N, (uint32_t)nt, (uint32_t)N, TOPK}; - GrArgs gr{(uint32_t)N, HC, 320, (uint32_t)nt, 1, 1e-6f}; + GrArgs gr{(uint32_t)N, HC, 320, (uint32_t)nt, apply, 1e-6f}; mtl::Dispatch c; c.p = &pcm; c.args = {mtl::Arg::B(ye), mtl::Arg::B(plan), mtl::Arg::B(ysh), mtl::Arg::B(sgv), mtl::Arg::B(bo), mtl::Arg::V(ga)}; c.grid[0] = (nt * N + 255) / 256; c.tg[0] = 256; @@ -145,14 +431,15 @@ int main() { n1.grid[0] = HC; n1.grid[1] = nt; n1.tg[0] = 256; cx.run({c, n1}); const uint32_t tk = TOPK; - mtl::Dispatch n2; n2.p = &pnc; + mtl::Dispatch n2; n2.p = four ? &pnc4 : &pnc; n2.args = {mtl::Arg::B(R2), mtl::Arg::B(ye), mtl::Arg::B(inj), mtl::Arg::B(wn), mtl::Arg::B(x2), mtl::Arg::V(gr), mtl::Arg::B(plan), mtl::Arg::B(ysh), mtl::Arg::B(sgv), mtl::Arg::V(tk)}; - n2.grid[0] = HC; n2.grid[1] = nt; n2.tg[0] = 256; + n2.grid[0] = four ? 1 : HC; n2.grid[1] = nt; n2.tg[0] = 256; cx.run({n2}); const bool ok = std::memcmp(R1.ptr, R2.ptr, 4 * (size_t)nt * HC * N) == 0 && std::memcmp(x1.ptr, x2.ptr, 4 * (size_t)nt * HC * N) == 0; fails += !ok; - std::printf("gr_norm_comb ntok %d: R', xn bitwise %s\n", nt, ok ? "PASS" : "FAIL"); + std::printf("gr_norm_comb%s ntok %d apply %u: R', xn bitwise %s\n", + four ? "_4s" : "", nt, apply, ok ? "PASS" : "FAIL"); } } std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); diff --git a/metal/tests/gdn_parity.cpp b/metal/tests/gdn_parity.cpp index 7288eb1d..e81c94ae 100644 --- a/metal/tests/gdn_parity.cpp +++ b/metal/tests/gdn_parity.cpp @@ -4,6 +4,7 @@ // gdn_parity parity for windows of 1..4 tokens: y, the committed state and conv history; a verify // window leaves the state untouched; its first m outputs are bitwise a fold of m // gdn_parity --bench GPU time per layer: 36 distinct layer states rotated (113 MB, beyond the SLC) +// gdn_parity --ab ref.metal|ref_dir [--bench] bitwise parity (ntok 1..8, commit 0/1/partial, plane 0/49152) + paired A/B #include "mtl_rt.h" #include "ref_quant.h" @@ -12,12 +13,17 @@ #include #include #include +#include #include #include namespace { constexpr int S = 128, HK = 16, HV = 48, C = (2 * HK + HV) * S, MAXT = 8; +constexpr uint32_t PLANE = MAXT * HV * S; // 49152 +constexpr size_t Y_ELEMS = 4 * (size_t)PLANE + PLANE / 16; struct GdnArgs { uint32_t h_k, h_v, channels, commit; float eps; uint32_t plane; }; +struct GlueArgs { uint32_t n, n_tok, n_embd, topk; }; +constexpr int P_INTS = 16 + 128 + 128 + 4 + 128 * 3 + 20 + 128, P_TOKENT = 16 + 128 + 128 + 4 + 128 * 2; double median(std::vector v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; } struct Err { double se = 0, sr = 0; @@ -27,17 +33,30 @@ struct Err { } // namespace int main(int argc, char** argv) { - const bool bench = argc > 1 && std::string(argv[1]) == "--bench"; + const bool ab = argc > 1 && std::string(argv[1]) == "--ab"; + if (ab && argc < 3) { + std::fprintf(stderr, "usage: gdn_parity --ab ref.metal|ref_dir [--bench|--no-bench]\n"); + return 2; + } + const bool no_bench = ab && argc > 3 && std::string(argv[3]) == "--no-bench"; + const bool bench = (argc > 1 && std::string(argv[1]) == "--bench") || (ab && !no_bench); + std::string ref_gdn = ab ? std::string(argv[2]) : ""; + std::string ref_glue = ""; + if (ab && std::filesystem::is_directory(ref_gdn)) { + if (std::filesystem::exists(ref_gdn + "/glue.metal")) ref_glue = ref_gdn + "/glue.metal"; + ref_gdn += "/gdn.metal"; + } mtl::Context cx; const int lib = cx.compile({"gdn.metal"}, "#define NO_GGML_COMMON 1\n"); + const int ref_lib = ab ? cx.compile({ref_gdn}, "#define NO_GGML_COMMON 1\n") : -1; std::printf("device %s\n", cx.device_name().c_str()); const int LAYERS = bench ? 36 : 1; const size_t st_elems = (size_t)S * HV * S; mtl::Buffer state = cx.alloc(4 * st_elems * LAYERS), st0 = cx.alloc(4 * st_elems), hist = cx.alloc(4 * C * 3 * LAYERS), h0 = cx.alloc(4 * C * 3), qkv = cx.alloc(4 * C * MAXT), al = cx.alloc(4 * HV * MAXT), be = cx.alloc(4 * HV * MAXT), z = cx.alloc(4 * HV * S * MAXT), cw = cx.alloc(4 * C * 4), - dt = cx.alloc(4 * HV), sa = cx.alloc(4 * HV), gm = cx.alloc(4 * S), y = cx.alloc(4 * HV * S * MAXT), - y2 = cx.alloc(4 * HV * S * MAXT), hist2 = cx.alloc(4 * C * 3 * LAYERS); + dt = cx.alloc(4 * HV), sa = cx.alloc(4 * HV), gm = cx.alloc(4 * S), y = cx.alloc(4 * Y_ELEMS), + y2 = cx.alloc(4 * Y_ELEMS), hist2 = cx.alloc(4 * C * 3 * LAYERS); ref::Rng r(17); for (size_t i = 0; i < st_elems; ++i) st0.as()[i] = r.sym() * 0.1f; for (int i = 0; i < C * 3; ++i) h0.as()[i] = r.sym(); @@ -51,8 +70,8 @@ int main(int argc, char** argv) { std::memcpy(state.as() + l * st_elems, st0.ptr, 4 * st_elems); std::memcpy(hist.as() + l * C * 3, h0.ptr, 4 * C * 3); } - auto disp = [&](const mtl::Pipeline& p, int layer, int commit, const mtl::Buffer& out) { - GdnArgs a{HK, HV, C, (uint32_t)commit, 1e-6f, 0}; + auto disp = [&](const mtl::Pipeline& p, int layer, int commit, const mtl::Buffer& out, uint32_t plane = 0) { + GdnArgs a{HK, HV, C, (uint32_t)commit, 1e-6f, plane}; mtl::Dispatch d; d.p = &p; d.args = {mtl::Arg::B(state, 4 * st_elems * layer), mtl::Arg::B(hist, 4 * (size_t)C * 3 * layer), mtl::Arg::B(qkv), @@ -78,19 +97,184 @@ int main(int argc, char** argv) { return 1; } } + if (ab) { + const char* kn = std::getenv("GDN_IB8") ? "gdn_step8" : "gdn_step"; + int ab_fails = 0; + mtl::Buffer st_ref = cx.alloc(4 * st_elems), h_ref = cx.alloc(4 * C * 3); + for (int nt = 1; nt <= MAXT; ++nt) { + mtl::Pipeline pc = cx.pipeline(lib, kn, mtl::Constants().i(2, nt)); + mtl::Pipeline pr = cx.pipeline(ref_lib, kn, mtl::Constants().i(2, nt)); + for (uint32_t plane : {0u, PLANE}) { + std::vector commits = {0, 1}; + for (int k = 1; k < nt; ++k) commits.push_back(0x100 | k); + for (int commit : commits) { + const size_t y_bytes = 4 * (plane ? Y_ELEMS : (size_t)HV * S * nt); + std::memcpy(state.ptr, st0.ptr, 4 * st_elems); + std::memcpy(hist.ptr, h0.ptr, 4 * C * 3); + std::memset(hist2.ptr, 0xa5, 4 * C * 3); + std::memset(y.ptr, 0xa5, y_bytes); + cx.run({disp(pr, 0, commit, y, plane)}); + std::memcpy(st_ref.ptr, state.ptr, 4 * st_elems); + std::memcpy(h_ref.ptr, hist2.ptr, 4 * C * 3); + + std::memcpy(state.ptr, st0.ptr, 4 * st_elems); + std::memcpy(hist.ptr, h0.ptr, 4 * C * 3); + std::memset(hist2.ptr, 0xa5, 4 * C * 3); + std::memset(y2.ptr, 0xa5, y_bytes); + cx.run({disp(pc, 0, commit, y2, plane)}); + + const bool y_ok = std::memcmp(y.ptr, y2.ptr, y_bytes) == 0; + const bool s_ok = std::memcmp(state.ptr, st_ref.ptr, 4 * st_elems) == 0; + const bool h_ok = std::memcmp(hist2.ptr, h_ref.ptr, 4 * C * 3) == 0; + const bool ok = y_ok && s_ok && h_ok; + ab_fails += !ok; + std::printf("ab ntok %d commit 0x%x plane %u: y %s state %s hist %s %s\n", + nt, commit, plane, y_ok ? "ok" : "DIFF", s_ok ? "ok" : "DIFF", h_ok ? "ok" : "DIFF", + ok ? "PASS" : "FAIL"); + } + } + } + std::printf("ab bitwise summary: %d fail(s) %s\n", ab_fails, ab_fails == 0 ? "ALL PASS" : "FAIL"); + if (ab_fails || !bench) return ab_fails ? 1 : 0; + for (uint32_t plane : {0u, PLANE}) { + for (int nt = 1; nt <= MAXT; ++nt) { + mtl::Pipeline pc = cx.pipeline(lib, kn, mtl::Constants().i(2, nt)); + mtl::Pipeline pr = cx.pipeline(ref_lib, kn, mtl::Constants().i(2, nt)); + for (int commit : {1, 0}) { + std::vector dsc, dsr; + for (int k = 0; k < 72; ++k) { + dsc.push_back(disp(pc, k % LAYERS, commit, y, plane)); + dsr.push_back(disp(pr, k % LAYERS, commit, y, plane)); + } + cx.warm(dsr, 0.015); + cx.warm(dsc, 0.015); + std::vector bt, ct, rt; + for (int i = 0; i < 7; ++i) { + double b, c; + if (i & 1) { c = cx.run(dsc) / 72; b = cx.run(dsr) / 72; } + else { b = cx.run(dsr) / 72; c = cx.run(dsc) / 72; } + bt.push_back(b * 1e6); + ct.push_back(c * 1e6); + rt.push_back(c / b); + } + std::printf("ab bench plane %5u ntok %d %-7s base %7.2f us cand %7.2f us ratio %.4f (x36: %.3f -> %.3f ms)\n", + plane, nt, commit ? "commit" : "verify", median(bt), median(ct), median(rt), + 36 * median(bt) * 1e-3, 36 * median(ct) * 1e-3); + } + } + } + if (!ref_glue.empty()) { + const int glib_c = cx.compile({"glue.metal"}, "#define NO_GGML_COMMON 1\n#define P_TOKENT 532\n"); + const int glib_r = cx.compile({ref_glue}, "#define NO_GGML_COMMON 1\n#define P_TOKENT 532\n"); + const mtl::Pipeline sw_c = cx.pipeline(glib_c, "swiglu_prep"), sw_r = cx.pipeline(glib_r, "swiglu_prep"); + const mtl::Pipeline mc_c = cx.pipeline(glib_c, "moe_combine_sh"), mc_r = cx.pipeline(glib_r, "moe_combine_sh"); + constexpr int N = 2560, NFF = 1024, TOPK = 10; + const size_t xp_elems = 4 * (size_t)MAXT * NFF + MAXT * NFF / 16; + mtl::Buffer gsh = cx.alloc(4 * MAXT * NFF), ush = cx.alloc(4 * MAXT * NFF), + xp1 = cx.alloc(4 * xp_elems), xp2 = cx.alloc(4 * xp_elems), + ye = cx.alloc(4 * (size_t)MAXT * TOPK * N), ysh = cx.alloc(4 * MAXT * N), + sgv = cx.alloc(4 * MAXT), plan = cx.alloc(4 * P_INTS), + bo1 = cx.alloc(4 * MAXT * N), bo2 = cx.alloc(4 * MAXT * N); + for (int i = 0; i < MAXT * NFF; ++i) { gsh.as()[i] = r.sym() * 2.0f; ush.as()[i] = r.sym() * 2.0f; } + for (size_t i = 0; i < (size_t)MAXT * TOPK * N; ++i) ye.as()[i] = r.sym(); + for (int i = 0; i < MAXT * N; ++i) ysh.as()[i] = r.sym(); + for (int i = 0; i < MAXT; ++i) sgv.as()[i] = r.sym() * 3.0f; + for (int i = 0; i < MAXT * TOPK; ++i) plan.as()[P_TOKENT + i] = (i * 7) % (MAXT * TOPK); + int glue_fails = 0; + for (int nt = 1; nt <= MAXT; ++nt) { + GlueArgs ga_sw{NFF, MAXT, N, TOPK}; + mtl::Dispatch dsw_r{&sw_r, {mtl::Arg::B(gsh), mtl::Arg::B(ush), mtl::Arg::B(xp1), mtl::Arg::V(ga_sw)}, {NFF / 16, (uint32_t)nt, 1}, {16, 1, 1}}; + mtl::Dispatch dsw_c{&sw_c, {mtl::Arg::B(gsh), mtl::Arg::B(ush), mtl::Arg::B(xp2), mtl::Arg::V(ga_sw)}, {NFF / 16, (uint32_t)nt, 1}, {16, 1, 1}}; + std::memset(xp1.ptr, 0xa5, xp1.bytes); std::memset(xp2.ptr, 0xa5, xp2.bytes); + cx.run({dsw_r, dsw_c}); + const bool sw_ok = std::memcmp(xp1.ptr, xp2.ptr, xp1.bytes) == 0; + + for (uint32_t tk : {10u, 8u}) { + GlueArgs ga_mc{N, (uint32_t)nt, N, tk}; + const uint32_t gx = (nt * N + 255) / 256; + mtl::Dispatch dmc_r{&mc_r, {mtl::Arg::B(ye), mtl::Arg::B(plan), mtl::Arg::B(ysh), mtl::Arg::B(sgv), mtl::Arg::B(bo1), mtl::Arg::V(ga_mc)}, {gx, 1, 1}, {256, 1, 1}}; + mtl::Dispatch dmc_c{&mc_c, {mtl::Arg::B(ye), mtl::Arg::B(plan), mtl::Arg::B(ysh), mtl::Arg::B(sgv), mtl::Arg::B(bo2), mtl::Arg::V(ga_mc)}, {gx, 1, 1}, {256, 1, 1}}; + std::memset(bo1.ptr, 0xa5, bo1.bytes); std::memset(bo2.ptr, 0xa5, bo2.bytes); + cx.run({dmc_r, dmc_c}); + const bool mc_ok = std::memcmp(bo1.ptr, bo2.ptr, 4 * (size_t)nt * N) == 0; + glue_fails += (!sw_ok || !mc_ok); + if (tk == 10u) + std::printf("glue ab ntok %d: swiglu_prep %s moe_combine_sh(topk=10) %s %s\n", + nt, sw_ok ? "ok" : "DIFF", mc_ok ? "ok" : "DIFF", (sw_ok && mc_ok) ? "PASS" : "FAIL"); + else if (!mc_ok) + std::printf("glue ab ntok %d: moe_combine_sh(topk=8) DIFF FAIL\n", nt); + } + } + std::printf("glue ab bitwise summary: %d fail(s) %s\n", glue_fails, glue_fails == 0 ? "ALL PASS" : "FAIL"); + if (glue_fails) return 1; + for (int nt : {1, 2, 4, 8}) { + GlueArgs ga_sw{NFF, MAXT, N, TOPK}; + mtl::Dispatch dsw_r{&sw_r, {mtl::Arg::B(gsh), mtl::Arg::B(ush), mtl::Arg::B(xp1), mtl::Arg::V(ga_sw)}, {NFF / 16, (uint32_t)nt, 1}, {16, 1, 1}}; + mtl::Dispatch dsw_c{&sw_c, {mtl::Arg::B(gsh), mtl::Arg::B(ush), mtl::Arg::B(xp2), mtl::Arg::V(ga_sw)}, {NFF / 16, (uint32_t)nt, 1}, {16, 1, 1}}; + cx.warm({dsw_r}, 0.01); cx.warm({dsw_c}, 0.01); + std::vector bt, ct, rt; + for (int i = 0; i < 7; ++i) { + double b, c; + if (i & 1) { c = cx.run({dsw_c}, 256, false, true) / 256; b = cx.run({dsw_r}, 256, false, true) / 256; } + else { b = cx.run({dsw_r}, 256, false, true) / 256; c = cx.run({dsw_c}, 256, false, true) / 256; } + bt.push_back(b * 1e6); ct.push_back(c * 1e6); rt.push_back(c / b); + } + std::printf("glue bench swiglu_prep ntok %d base %6.2f us cand %6.2f us ratio %.4f\n", + nt, median(bt), median(ct), median(rt)); + + GlueArgs ga_mc{N, (uint32_t)nt, N, TOPK}; + const uint32_t gx = (nt * N + 255) / 256; + mtl::Dispatch dmc_r{&mc_r, {mtl::Arg::B(ye), mtl::Arg::B(plan), mtl::Arg::B(ysh), mtl::Arg::B(sgv), mtl::Arg::B(bo1), mtl::Arg::V(ga_mc)}, {gx, 1, 1}, {256, 1, 1}}; + mtl::Dispatch dmc_c{&mc_c, {mtl::Arg::B(ye), mtl::Arg::B(plan), mtl::Arg::B(ysh), mtl::Arg::B(sgv), mtl::Arg::B(bo2), mtl::Arg::V(ga_mc)}, {gx, 1, 1}, {256, 1, 1}}; + cx.warm({dmc_r}, 0.01); cx.warm({dmc_c}, 0.01); + bt.clear(); ct.clear(); rt.clear(); + for (int i = 0; i < 7; ++i) { + double b, c; + if (i & 1) { c = cx.run({dmc_c}, 256, false, true) / 256; b = cx.run({dmc_r}, 256, false, true) / 256; } + else { b = cx.run({dmc_r}, 256, false, true) / 256; c = cx.run({dmc_c}, 256, false, true) / 256; } + bt.push_back(b * 1e6); ct.push_back(c * 1e6); rt.push_back(c / b); + } + std::printf("glue bench moe_combine_sh ntok %d base %6.2f us cand %6.2f us ratio %.4f\n", + nt, median(bt), median(ct), median(rt)); + } + } + return 0; + } if (bench) { + mtl::Buffer state2 = cx.alloc(4 * st_elems * LAYERS), hist3 = cx.alloc(4 * (size_t)C * 3 * LAYERS); + std::memset(state2.ptr, 0, state2.bytes); + std::memset(hist3.ptr, 0, hist3.bytes); for (int nt : {1, 2, 3, 4}) { const char* kn = std::getenv("GDN_IB8") ? "gdn_step8" : "gdn_step"; mtl::Pipeline p = cx.pipeline(lib, kn, mtl::Constants().i(2, nt)); - for (int commit : {1, 0}) { + for (int commit : {1, 0, 2}) { + if (commit == 2 && nt == 1) continue; std::vector ds; - for (int k = 0; k < 72; ++k) ds.push_back(disp(p, k % LAYERS, commit, y)); + for (int k = 0; k < 72; ++k) { + const int l = k % LAYERS; + if (commit == 2) { + GdnArgs a{HK, HV, C, 2u, 1e-6f, PLANE}; + mtl::Dispatch d; + d.p = &p; + d.args = {mtl::Arg::B(state, 4 * st_elems * l), mtl::Arg::B(hist, 4 * (size_t)C * 3 * l), + mtl::Arg::B(qkv), mtl::Arg::B(al), mtl::Arg::B(be), mtl::Arg::B(z), mtl::Arg::B(cw), + mtl::Arg::B(dt), mtl::Arg::B(sa), mtl::Arg::B(gm), mtl::Arg::B(y), mtl::Arg::V(a), + mtl::Arg::B(hist2, 4 * (size_t)C * 3 * l), + mtl::Arg::B(state2, 4 * st_elems * l), mtl::Arg::B(hist3, 4 * (size_t)C * 3 * l)}; + d.grid[0] = HV; d.tg[0] = 512; + ds.push_back(d); + } else { + ds.push_back(disp(p, l, commit, y, PLANE)); + } + } cx.warm(ds, 0.1); std::vector t; for (int i = 0; i < 7; ++i) t.push_back(cx.run(ds) / 72); - const double m = median(t), bytes = 4.0 * st_elems * (commit ? 2 : 1); + const int n_writes = commit == 2 ? 2 : (commit ? 1 : 0); + const double m = median(t), bytes = 4.0 * st_elems * (1 + n_writes); + const char* cname = commit == 2 ? "dual2" : (commit ? "commit" : "verify"); std::printf("gdn_step ntok %d %-7s %7.2f us per layer state %.1f MB at %6.1f GB/s (x36 layers: %.3f ms)\n", - nt, commit ? "commit" : "verify", m * 1e6, bytes / 1e6, bytes / m / 1e9, 36 * m * 1e3); + nt, cname, m * 1e6, bytes / 1e6, bytes / m / 1e9, 36 * m * 1e3); } } return 0; @@ -201,12 +385,67 @@ int main(int argc, char** argv) { partial = partial && std::memcmp(state.ptr, sk.data(), 4 * st_elems) == 0 && std::memcmp(hist2.ptr, hk.data(), 4 * C * 3) == 0 && std::memcmp(y2.ptr, yv.data(), 4 * HV * S * nt) == 0; } - const bool ok = untouched && same_y && prefix && partial && ey.l2() < 1e-5 && es.l2() < 1e-5 && eh.l2() == 0; + // plane = PLANE (49152): verify prepared planes and per-16 sums against raw yv + std::memcpy(state.ptr, st0.ptr, 4 * st_elems); + std::memcpy(hist.ptr, h0.ptr, 4 * C * 3); + std::memset(y2.ptr, 0, 4 * Y_ELEMS); + cx.run({disp(p, 0, 0, y2, PLANE)}); + bool plane_ok = std::memcmp(y.ptr, y2.ptr, 4 * HV * S * nt) == 0; + const float q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f}; + const float p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + const float* yp = y2.as(); + for (int idx = 0; idx < HV * S * nt && plane_ok; ++idx) { + const float yv = yp[idx]; + const int jj = idx & 15; + if (yp[PLANE + idx] != yv * q2s[jj & 7]) plane_ok = false; + if (yp[2 * PLANE + idx] != yv * p4s[jj & 3]) plane_ok = false; + if (yp[3 * PLANE + idx] != ((jj & 1) ? yv * 0.00390625f : yv)) plane_ok = false; + if (jj == 0) { + float s0 = 0.0f, s1 = 0.0f; + for (int k = 0; k < 16; k += 2) { + s0 += yp[idx + k]; + s1 += yp[idx + k + 1]; + } + if (yp[4 * PLANE + idx / 16] != s0 + s1) plane_ok = false; + } + } + // dual-slot endpoint rollback property (commit = 2u): state/hist_out match 1-token commit, state_end/hist_end match nt-token commit BITWISE + bool ring_ok = true; + if (nt > 1) { + mtl::Buffer st_end = cx.alloc(4 * st_elems), h_end = cx.alloc(4 * C * 3); + mtl::Pipeline p1 = cx.pipeline(lib, std::getenv("GDN_IB8") ? "gdn_step8" : "gdn_step", mtl::Constants().i(2, 1)); + std::memcpy(state.ptr, st0.ptr, 4 * st_elems); + std::memcpy(hist.ptr, h0.ptr, 4 * C * 3); + cx.run({disp(p1, 0, 1, y2)}); + std::vector s1_ref(state.as(), state.as() + st_elems), h1_ref(hist2.as(), hist2.as() + C * 3); + std::memcpy(state.ptr, st0.ptr, 4 * st_elems); + std::memcpy(hist.ptr, h0.ptr, 4 * C * 3); + cx.run({disp(p, 0, 1, y2)}); + std::vector sn_ref(state.as(), state.as() + st_elems), hn_ref(hist2.as(), hist2.as() + C * 3); + + std::memcpy(state.ptr, st0.ptr, 4 * st_elems); + std::memcpy(hist.ptr, h0.ptr, 4 * C * 3); + std::memset(hist2.ptr, 0xCD, hist2.bytes); + std::memset(st_end.ptr, 0xCD, st_end.bytes); + std::memset(h_end.ptr, 0xCD, h_end.bytes); + GdnArgs a{HK, HV, C, 2u, 1e-6f, 0}; + mtl::Dispatch d{&p, {mtl::Arg::B(state), mtl::Arg::B(hist), mtl::Arg::B(qkv), mtl::Arg::B(al), mtl::Arg::B(be), + mtl::Arg::B(z), mtl::Arg::B(cw), mtl::Arg::B(dt), mtl::Arg::B(sa), mtl::Arg::B(gm), + mtl::Arg::B(y2), mtl::Arg::V(a), mtl::Arg::B(hist2), mtl::Arg::B(st_end), mtl::Arg::B(h_end)}, + {HV, 1, 1}, {512, 1, 1}}; + cx.run({d}); + ring_ok = std::memcmp(y2.ptr, y.ptr, 4 * HV * S * nt) == 0 && + std::memcmp(state.ptr, s1_ref.data(), 4 * st_elems) == 0 && + std::memcmp(hist2.ptr, h1_ref.data(), 4 * C * 3) == 0 && + std::memcmp(st_end.ptr, sn_ref.data(), 4 * st_elems) == 0 && + std::memcmp(h_end.ptr, hn_ref.data(), 4 * C * 3) == 0; + } + const bool ok = untouched && same_y && prefix && partial && ring_ok && plane_ok && ey.l2() < 1e-5 && es.l2() < 1e-5 && eh.l2() == 0; fails += !ok; std::printf("gdn parity ntok %d: y rel L2 %.2e, state %.2e, history %.1e; verify leaves state: %s, verify y == " - "commit y: %s, first %d of a longer window bitwise: %s, partial commits bitwise: %s %s\n", nt, ey.l2(), es.l2(), - eh.l2(), untouched ? "yes" : "NO", same_y ? "yes" : "NO", nt, prefix ? "yes" : "NO", partial ? "yes" : "NO", - ok ? "PASS" : "FAIL"); + "commit y: %s, first %d of a longer window bitwise: %s, partial commits bitwise: %s, ring rollback bitwise: %s, plane bitwise: %s %s\n", + nt, ey.l2(), es.l2(), eh.l2(), untouched ? "yes" : "NO", same_y ? "yes" : "NO", nt, + prefix ? "yes" : "NO", partial ? "yes" : "NO", ring_ok ? "yes" : "NO", plane_ok ? "yes" : "NO", ok ? "PASS" : "FAIL"); } std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); return fails ? 1 : 0; diff --git a/metal/tests/gdn_prompt_parity.cpp b/metal/tests/gdn_prompt_parity.cpp index 24cb85bf..c27dd035 100644 --- a/metal/tests/gdn_prompt_parity.cpp +++ b/metal/tests/gdn_prompt_parity.cpp @@ -2,7 +2,13 @@ // against the decode step run over the same T tokens in order (gdn.metal's gdn_step, GDN_NTOK = T, itself checked // against the CPU reference by gdn_parity): y, the final state and the conv history. The scan reassociates // o = S_new^T q as g S^T q + delta (k . q), so the match is to rounding, not bitwise. Then the speed for T = 2048. +// With the M5's matmul2d layout (coop_k4) also the chunked WY scan (gdn_wy_prep + gdn_wy_scan): against gdn_step +// as above, and o and the state against gdn_scanr_4 and an fp64 scan of gdn_pre's outputs (T <= 512, or all with +// GDN_WY_REF64=1); --wy-ab times it against gdn_scanr_4 (36 layers, alternating pairs). // gdn_prompt_parity [--bench] +// gdn_prompt_parity --ab baseline.metal T f32 [bench] +// gdn_prompt_parity --wy-ab [T ...] +#include "coop_probe.h" #include "mtl_rt.h" #include "ref_quant.h" @@ -10,6 +16,7 @@ #include #include #include +#include #include #include @@ -18,14 +25,28 @@ constexpr int HK = 16, HV = 48, CH = 10240, S = 128; struct GdnArgs { uint32_t h_k, h_v, channels, commit; float eps; uint32_t plane; }; struct GdnPArgs { uint32_t n_tok, h_k, h_v, channels; float eps; uint32_t f32; }; double median(std::vector v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; } +bool same(const mtl::Buffer& a, const mtl::Buffer& b, size_t bytes) { return std::memcmp(a.ptr, b.ptr, bytes) == 0; } } // namespace int main(int argc, char** argv) { + const bool ab = argc > 1 && std::string(argv[1]) == "--ab"; + if (ab && (argc < 5 || std::atoi(argv[3]) < 1 || + (std::string(argv[4]) != "0" && std::string(argv[4]) != "1"))) { + std::fprintf(stderr, "usage: gdn_prompt_parity --ab baseline.metal T f32 [bench]\n"); + return 2; + } const bool bench = argc > 1 && std::string(argv[1]) == "--bench"; + const bool wy_ab = argc > 1 && std::string(argv[1]) == "--wy-ab"; mtl::Context cx; - const int lib = cx.compile({"gdn.metal", "gdn_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); + const bool k4 = coop_k4(cx); + const int lib = cx.compile(ab ? std::vector{"gdn_prompt.metal"} : + std::vector{"gdn.metal", "gdn_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); std::printf("device %s\n", cx.device_name().c_str()); - const int TMAX = 2048; + int TMAX = ab ? std::atoi(argv[3]) : 2048; + std::vector wy_ts; + for (int i = 2; wy_ab && i < argc; ++i) wy_ts.push_back(std::atoi(argv[i])); + if (wy_ab && wy_ts.empty()) wy_ts = {512, 1024, 2048}; + for (int T : wy_ts) TMAX = std::max(TMAX, T); auto F = [&](size_t n) { return cx.alloc(4 * n); }; mtl::Buffer qkv = F((size_t)TMAX * CH), al = F(TMAX * HV), be = F(TMAX * HV), z = F((size_t)TMAX * HV * S), cw = F(CH * 4), dt = F(HV), sa = F(HV), gam = F(S), st0 = F(S * HV * S), stA = F(S * HV * S), stB = F(S * HV * S), @@ -36,6 +57,10 @@ int main(int argc, char** argv) { const size_t NCM = (TMAX + 31) / 32; mtl::Buffer cU = F((size_t)HV * NCM * 32 * S), cW = F((size_t)HV * NCM * 32 * S), clg = F((size_t)HV * NCM * 32), cD = F((size_t)HV * NCM * 32 * S); + // the WY scan's scratch: k / q fp16 hi / lo planes, TB / P fp16 hi / lo, G / r_L / G_L + const size_t NC16 = (TMAX + 15) / 16; + mtl::Buffer wkq = cx.alloc(2 * 4 * (size_t)HK * NC16 * 16 * S), wyh = cx.alloc(2 * (size_t)HV * NC16 * 1024), + wyf = F((size_t)HV * NC16 * 48); ref::Rng r(21); for (size_t i = 0; i < (size_t)TMAX * CH; ++i) qkv.as()[i] = r.sym() * 2.0f; for (int i = 0; i < TMAX * HV; ++i) { al.as()[i] = r.sym() * 3.0f; be.as()[i] = r.sym() * 2.0f; } @@ -53,11 +78,21 @@ int main(int argc, char** argv) { d1.args = {mtl::Arg::B(qkv), mtl::Arg::B(hB), mtl::Arg::B(cw), mtl::Arg::B(al), mtl::Arg::B(be), mtl::Arg::B(dt), mtl::Arg::B(sa), mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(vc), mtl::Arg::B(gb), mtl::Arg::B(kq), mtl::Arg::B(hA2), mtl::Arg::V(a)}; - d1.grid[0] = HK + HV / 2; + d1.grid[0] = (HK + HV) / 8; d1.grid[1] = T; d1.tg[0] = 256; ds.push_back(d1); - if (cb == -100) { // the chunked form: gdn_cprep, gdn_cstate, gdn_cout + if (cb == -200) { // the WY form: gdn_wy_prep, gdn_wy_scan + static mtl::Pipeline pwp = cx.pipeline(lib, "gdn_wy_prep"), pws = cx.pipeline(lib, "gdn_wy_scan"); + mtl::Dispatch e1; e1.p = &pwp; + e1.args = {mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(gb), mtl::Arg::B(wkq), mtl::Arg::B(wyh), mtl::Arg::B(wyf), + mtl::Arg::V(a)}; + e1.grid[0] = HK; e1.grid[1] = (T + 15) / 16; e1.tg[0] = 128; ds.push_back(e1); + mtl::Dispatch e2; e2.p = &pws; + e2.args = {mtl::Arg::B(stB), mtl::Arg::B(vc), mtl::Arg::B(wkq), mtl::Arg::B(wyh), mtl::Arg::B(wyf), mtl::Arg::B(o), + mtl::Arg::V(a)}; + e2.grid[0] = HV * 8; e2.tg[0] = 64; ds.push_back(e2); + } else if (cb == -100) { // the chunked form: gdn_cprep, gdn_cstate, gdn_cout const uint32_t nc = (T + 31) / 32; static mtl::Pipeline pcp = cx.pipeline(lib, "gdn_cprep"), pcs = cx.pipeline(lib, "gdn_cstate"), pco = cx.pipeline(lib, "gdn_cout"); @@ -85,19 +120,190 @@ int main(int argc, char** argv) { mtl::Dispatch d3; d3.p = &post; d3.args = {mtl::Arg::B(o), mtl::Arg::B(z), mtl::Arg::B(gam), mtl::Arg::B(yh), mtl::Arg::B(yB), mtl::Arg::V(a)}; - d3.grid[0] = HV; + d3.grid[0] = HV / 4; d3.grid[1] = T; d3.tg[0] = 128; ds.push_back(d3); }; mtl::Pipeline pre = cx.pipeline(lib, "gdn_pre"), post = cx.pipeline(lib, "gdn_post"); int fails = 0; - if (!bench) + + if (wy_ab) { + if (!k4) { std::printf("not the M5's matmul2d layout: no WY scan\n"); return 0; } + mtl::Pipeline scanr4 = cx.pipeline(lib, "gdn_scanr_4", mtl::Constants().i(6, 0)); + auto pairs = [&](const std::vector& b, const std::vector& c, const char* what, int T) { + cx.warm(b, 0.05); + cx.warm(c, 0.05); + std::vector ratios, bt, ct; + for (int pair = 0; pair < 5; ++pair) { + double x, y; + if (pair & 1) { y = cx.run(c, 36); x = cx.run(b, 36); } + else { x = cx.run(b, 36); y = cx.run(c, 36); } + bt.push_back(x * 1e3); ct.push_back(y * 1e3); ratios.push_back(y / x); + } + std::printf("T %4d %-22s 36 layers: gdn_scanr_4 %.3f ms, WY %.3f ms, ratio %.4f\n", T, what, median(bt), median(ct), + median(ratios)); + }; + for (int T : wy_ts) { + std::vector base, cand; + prompt_ds(T, pre, scanr4, post, -4, base); + prompt_ds(T, pre, scanr4, post, -200, cand); + const GdnPArgs a{(uint32_t)T, HK, HV, CH, 1e-6f, 0u}; // as prompt_chunk: no fp32 y + for (auto* v : {&base, &cand}) + for (auto& d : *v) d.args.back() = mtl::Arg::V(a); + std::copy(h0.as(), h0.as() + CH * 3, hB.as()); + std::copy(st0.as(), st0.as() + S * HV * S, stB.as()); + cx.run(base); + pairs({base[1]}, {cand[1], cand[2]}, "scan (prep + scan)", T); + pairs({base[1]}, {cand[1]}, " gdn_wy_prep alone", T); + pairs({base[1]}, {cand[2]}, " gdn_wy_scan alone", T); + pairs(base, cand, "pre + scan + post", T); + } + return 0; + } + if (ab) { + const int T = TMAX; + const uint32_t f32 = std::atoi(argv[4]); + const bool timing = argc > 5 && std::string(argv[5]) == "bench"; + const int ref_lib = cx.compile({argv[2]}, "#define NO_GGML_COMMON 1\n"); + mtl::Pipeline pre_ref = cx.pipeline(ref_lib, "gdn_pre"), post_ref = cx.pipeline(ref_lib, "gdn_post"), + scan_ref = cx.pipeline(ref_lib, "gdn_scanr_4", mtl::Constants().i(6, 0)), + scan = cx.pipeline(lib, "gdn_scanr_4", mtl::Constants().i(6, 0)); + std::copy(h0.as(), h0.as() + CH * 3, hB.as()); + std::copy(st0.as(), st0.as() + S * HV * S, stB.as()); + std::vector candidate, reference; + prompt_ds(T, pre, scan, post, -4, candidate); + GdnPArgs a{(uint32_t)T, HK, HV, CH, 1e-6f, f32}; + candidate[0].args.back() = mtl::Arg::V(a); + candidate[1].args.back() = mtl::Arg::V(a); + candidate[2].args.back() = mtl::Arg::V(a); + reference = candidate; + reference[0].p = &pre_ref; reference[1].p = &scan_ref; reference[2].p = &post_ref; + const mtl::Buffer* outputs[] = {&qn, &kn, &vc, &gb, &kq, &hA2, &o, &stB, &yh, &yB}; + const size_t sizes[] = {4 * (size_t)T * HK * S, 4 * (size_t)T * HK * S, 4 * (size_t)T * HV * S, + 8 * (size_t)T * HV, 4 * (size_t)T * HK, 4 * CH * 3, 4 * (size_t)T * HV * S, + 4 * S * HV * S, 2 * (size_t)T * HV * S, 4 * (size_t)T * HV * S}; + const char* names[] = {"qn", "kn", "vc", "gb", "kq", "history", "o", "state", "yh", "y"}; + std::vector> expected; + // Each stage uses identical inputs even if an earlier candidate stage fails. + for (int stage = 0, out = 0; stage < 3; ++stage) { + const int end = stage == 0 ? 6 : stage == 1 ? 8 : 10; + if (stage == 2) std::memset(yB.ptr, 0xa5, yB.bytes); + cx.run({reference[stage]}); + for (int j = out; j < end; ++j) { + const auto* bytes = outputs[j]->as(); + expected.emplace_back(bytes, bytes + sizes[j]); + } + if (stage == 1) std::copy(st0.as(), st0.as() + S * HV * S, stB.as()); + if (stage == 2) std::memset(yB.ptr, 0xa5, yB.bytes); + cx.run({candidate[stage]}); + for (int j = out; j < end; ++j) { + const bool ok = std::memcmp(outputs[j]->ptr, expected[j].data(), sizes[j]) == 0; + std::printf("T %d f32 %u %s %s\n", T, f32, names[j], ok ? "PASS" : "FAIL"); + fails += !ok; + std::memcpy(outputs[j]->ptr, expected[j].data(), sizes[j]); + } + out = end; + } + if (fails || !timing) return fails ? 1 : 0; + for (int stage = 0; stage < 3; ++stage) { + cx.warm({reference[stage], candidate[stage]}, 0.05); + cx.run({reference[stage]}, 36); + cx.run({candidate[stage]}, 36); + std::vector ratios, bt, ct; + for (int pair = 0; pair < 3; ++pair) { + double b, c; + if (pair & 1) { c = cx.run({candidate[stage]}, 36); b = cx.run({reference[stage]}, 36); } + else { b = cx.run({reference[stage]}, 36); c = cx.run({candidate[stage]}, 36); } + bt.push_back(b * 1e3); ct.push_back(c * 1e3); ratios.push_back(c / b); + std::printf("pair %d stage %d baseline %.3f candidate %.3f ms ratio %.4f\n", pair, stage, b * 1e3, c * 1e3, c / b); + } + std::printf("T %d f32 %u stage %d 36 layers baseline %.3f candidate %.3f ms ratio %.4f\n", + T, f32, stage, median(bt), median(ct), median(ratios)); + } + { + cx.warm({reference[0], reference[1], reference[2], candidate[0], candidate[1], candidate[2]}, 0.05); + cx.run(reference, 36); + cx.run(candidate, 36); + std::vector ratios, bt, ct; + for (int pair = 0; pair < 3; ++pair) { + double b, c; + if (pair & 1) { c = cx.run(candidate, 36); b = cx.run(reference, 36); } + else { b = cx.run(reference, 36); c = cx.run(candidate, 36); } + bt.push_back(b * 1e3); ct.push_back(c * 1e3); ratios.push_back(c / b); + std::printf("pair %d total baseline %.3f candidate %.3f ms ratio %.4f\n", pair, b * 1e3, c * 1e3, c / b); + } + std::printf("T %d f32 %u total 36 layers baseline %.3f candidate %.3f ms ratio %.4f\n", + T, f32, median(bt), median(ct), median(ratios)); + } + return 0; + } + if (!bench) { + // Bitwise parity: 1-simdgroup-per-head packed grid ((8, T) pre, (12, T) post) vs legacy grid ((40, T) pre, (48, T) post) + mtl::Buffer qn2 = F((size_t)TMAX * HK * S), kn2 = F((size_t)TMAX * HK * S), vc2 = F((size_t)TMAX * HV * S), + gb2 = F(2 * TMAX * HV), kq2 = F(TMAX * HK), h2 = F(CH * 3), yh2 = cx.alloc(2 * (size_t)TMAX * HV * S), + y2 = F((size_t)TMAX * HV * S); + for (int T : {1, 3, 37, 300}) { + GdnPArgs a{(uint32_t)T, HK, HV, CH, 1e-6f, 1u}; + std::copy(h0.as(), h0.as() + CH * 3, hB.as()); + mtl::Dispatch pa, pb, qa, qb; + pa.p = pb.p = ⪯ + pa.args = {mtl::Arg::B(qkv), mtl::Arg::B(hB), mtl::Arg::B(cw), mtl::Arg::B(al), mtl::Arg::B(be), mtl::Arg::B(dt), + mtl::Arg::B(sa), mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(vc), mtl::Arg::B(gb), mtl::Arg::B(kq), + mtl::Arg::B(hA2), mtl::Arg::V(a)}; + pb.args = {mtl::Arg::B(qkv), mtl::Arg::B(hB), mtl::Arg::B(cw), mtl::Arg::B(al), mtl::Arg::B(be), mtl::Arg::B(dt), + mtl::Arg::B(sa), mtl::Arg::B(qn2), mtl::Arg::B(kn2), mtl::Arg::B(vc2), mtl::Arg::B(gb2), mtl::Arg::B(kq2), + mtl::Arg::B(h2), mtl::Arg::V(a)}; + pa.grid[0] = (HK + HV) / 8; pa.grid[1] = T; pa.tg[0] = 256; + pb.grid[0] = HK + HV / 2; pb.grid[1] = T; pb.tg[0] = 256; + cx.run({pa, pb}); + const bool pre_ok = same(qn, qn2, 4 * (size_t)T * HK * S) && same(kn, kn2, 4 * (size_t)T * HK * S) && + same(vc, vc2, 4 * (size_t)T * HV * S) && same(gb, gb2, 8 * (size_t)T * HV) && + same(kq, kq2, 4 * (size_t)T * HK) && same(hA2, h2, 4 * CH * 3); + qa.p = qb.p = &post; + qa.args = {mtl::Arg::B(vc), mtl::Arg::B(z), mtl::Arg::B(gam), mtl::Arg::B(yh), mtl::Arg::B(yB), mtl::Arg::V(a)}; + qb.args = {mtl::Arg::B(vc), mtl::Arg::B(z), mtl::Arg::B(gam), mtl::Arg::B(yh2), mtl::Arg::B(y2), mtl::Arg::V(a)}; + qa.grid[0] = HV / 4; qa.grid[1] = T; qa.tg[0] = 128; + qb.grid[0] = HV; qb.grid[1] = T; qb.tg[0] = 128; + cx.run({qa, qb}); + const bool post_ok = same(yh, yh2, 2 * (size_t)T * HV * S) && same(yB, y2, 4 * (size_t)T * HV * S); + fails += !pre_ok + !post_ok; + std::printf("T %3d, bitwise 1-sg vs legacy: gdn_pre %s, gdn_post %s\n", T, pre_ok ? "PASS" : "FAIL", + post_ok ? "PASS" : "FAIL"); + } + mtl::Pipeline scanr4 = cx.pipeline(lib, "gdn_scanr_4", mtl::Constants().i(6, 0)), + scanr4_16 = cx.pipeline(lib, "gdn_scanr_4_16", mtl::Constants().i(6, 0)); + for (int T : {1, 3, 16, 17, 32, 37, 300, 2048}) { + GdnPArgs a{(uint32_t)T, HK, HV, CH, 1e-6f, 1u}; + std::copy(h0.as(), h0.as() + CH * 3, hB.as()); + std::copy(st0.as(), st0.as() + S * HV * S, stA.as()); + std::copy(st0.as(), st0.as() + S * HV * S, stB.as()); + mtl::Dispatch dp, da_s, db_s; + dp.p = ⪯ + dp.args = {mtl::Arg::B(qkv), mtl::Arg::B(hB), mtl::Arg::B(cw), mtl::Arg::B(al), mtl::Arg::B(be), mtl::Arg::B(dt), + mtl::Arg::B(sa), mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(vc), mtl::Arg::B(gb), mtl::Arg::B(kq), + mtl::Arg::B(hA2), mtl::Arg::V(a)}; + dp.grid[0] = (HK + HV) / 8; dp.grid[1] = T; dp.tg[0] = 256; + da_s.p = &scanr4; + da_s.args = {mtl::Arg::B(stA), mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(vc), mtl::Arg::B(gb), mtl::Arg::B(kq), + mtl::Arg::B(o), mtl::Arg::V(a)}; + da_s.grid[0] = HV * 4; da_s.tg[0] = 128; + db_s.p = &scanr4_16; + db_s.args = {mtl::Arg::B(stB), mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(vc), mtl::Arg::B(gb), mtl::Arg::B(kq), + mtl::Arg::B(y2), mtl::Arg::V(a)}; + db_s.grid[0] = HV * 4; db_s.tg[0] = 128; + cx.run({dp, da_s, db_s}); + const bool o_ok = same(o, y2, 4 * (size_t)T * HV * S), st_ok = same(stA, stB, 4 * (size_t)S * HV * S); + fails += !o_ok + !st_ok; + std::printf("T %4d, bitwise single-pass vs TB=16 gdn_scanr_4: o %s, state %s\n", T, o_ok ? "PASS" : "FAIL", + st_ok ? "PASS" : "FAIL"); + } for (int T : {1, 3, 37, 300}) - for (int cb : {1, 4, -2, -4, -8, -9, -100}) { + for (int cb : {1, 4, -2, -4, -8, -9, -100, -200}) { + if (cb == -200 && !k4) continue; mtl::Constants pc; if (cb < 0 && cb != -9) pc.i(6, 0); - mtl::Pipeline scan = cx.pipeline(lib, cb == -100 ? std::string("gdn_scanr_4") : cb == -9 ? std::string("gdn_scan_v") : cb > 0 ? "gdn_scan_" + std::to_string(cb) : "gdn_scanr_" + std::to_string(-cb), cb == -100 ? mtl::Constants().i(6, 0) : pc); + mtl::Pipeline scan = cx.pipeline(lib, cb <= -100 ? std::string("gdn_scanr_4") : cb == -9 ? std::string("gdn_scan_v") : cb > 0 ? "gdn_scan_" + std::to_string(cb) : "gdn_scanr_" + std::to_string(-cb), cb <= -100 ? mtl::Constants().i(6, 0) : pc); if (cb == -9) cb = -8; mtl::Pipeline step = cx.pipeline(lib, "gdn_step", mtl::Constants().i(2, T)); std::copy(st0.as(), st0.as() + S * HV * S, stA.as()); @@ -135,6 +341,80 @@ int main(int argc, char** argv) { std::printf("T %3d, scan CB %d: y rel %.2e (fp16 %.2e), state rel %.2e, history %s %s\n", T, cb, ry, std::sqrt(ryh / syh), rs, rh == 0 ? "exact" : "DIFFERS", ok ? "PASS" : "FAIL"); } + + // the WY scan against gdn_scanr_4 and an fp64 scan of the same gdn_pre outputs: o and the final state + if (k4) { + mtl::Pipeline scanr4w = cx.pipeline(lib, "gdn_scanr_4", mtl::Constants().i(6, 0)); + auto ref64 = [&](int T, std::vector& ro, std::vector& rs) { + rs.assign(st0.as(), st0.as() + S * HV * S); + ro.assign((size_t)T * HV * S, 0.0); + const float *Q = qn.as(), *K = kn.as(), *V = vc.as(), *G = gb.as(); + std::vector ka(S); + for (int h = 0; h < HV; ++h) + for (int t = 0; t < T; ++t) { + const float *k = K + ((size_t)t * HK + h % HK) * S, *q = Q + ((size_t)t * HK + h % HK) * S; + const double g = G[2 * ((size_t)t * HV + h)], b = G[2 * ((size_t)t * HV + h) + 1]; + for (int j = 0; j < S; ++j) { + double s = 0; + for (int i = 0; i < S; ++i) s += rs[((size_t)i * HV + h) * S + j] * k[i]; + ka[j] = s; + } + for (int j = 0; j < S; ++j) { + const double d = (V[((size_t)t * HV + h) * S + j] - g * ka[j]) * b; + for (int i = 0; i < S; ++i) rs[((size_t)i * HV + h) * S + j] = g * rs[((size_t)i * HV + h) * S + j] + k[i] * d; + } + for (int j = 0; j < S; ++j) { + double s = 0; + for (int i = 0; i < S; ++i) s += rs[((size_t)i * HV + h) * S + j] * q[i]; + ro[((size_t)t * HV + h) * S + j] = s / std::sqrt(128.0); + } + } + }; + auto cmp = [](const float* x, const auto* ref, size_t n, double& mx) { + double se = 0, sr = 0; + mx = 0; + for (size_t i = 0; i < n; ++i) { + const double e = x[i] - (double)ref[i]; + se += e * e; + sr += (double)ref[i] * ref[i]; + mx = std::max(mx, std::fabs(e)); + } + return std::sqrt(se / std::max(sr, 1e-300)); + }; + for (int T : {16, 37, 64, 300, 512, 1024, 2048}) { + GdnPArgs a{(uint32_t)T, HK, HV, CH, 1e-6f, 1u}; + std::vector ds; + prompt_ds(T, pre, scanr4w, post, -200, ds); + std::copy(h0.as(), h0.as() + CH * 3, hB.as()); + cx.run({ds[0]}); + mtl::Dispatch r; + r.p = &scanr4w; + r.args = {mtl::Arg::B(stA), mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(vc), mtl::Arg::B(gb), mtl::Arg::B(kq), + mtl::Arg::B(y2), mtl::Arg::V(a)}; + r.grid[0] = HV * 4; r.tg[0] = 128; + std::copy(st0.as(), st0.as() + S * HV * S, stA.as()); + std::copy(st0.as(), st0.as() + S * HV * S, stB.as()); + std::memset(o.ptr, 0xff, o.bytes); // NaN: a token the WY scan skips fails + cx.run({r, ds[1], ds[2]}); + const size_t no = (size_t)T * HV * S, ns = (size_t)S * HV * S; + double mo, ms; + const double ro = cmp(o.as(), y2.as(), no, mo), rs = cmp(stB.as(), stA.as(), ns, ms); + bool ok = ro < 1e-5 && rs < 1e-5; + std::printf("T %4d, WY scan vs gdn_scanr_4: o rel %.2e max %.2e, state rel %.2e max %.2e", T, ro, mo, rs, ms); + if (T <= 512 || std::getenv("GDN_WY_REF64")) { + std::vector o64, s64; + ref64(T, o64, s64); + double m1, m2, m3, m4; + const double wo = cmp(o.as(), o64.data(), no, m1), ws = cmp(stB.as(), s64.data(), ns, m2), + bo = cmp(y2.as(), o64.data(), no, m3), bs = cmp(stA.as(), s64.data(), ns, m4); + std::printf("; vs fp64: WY o %.2e state %.2e, gdn_scanr_4 o %.2e state %.2e", wo, ws, bo, bs); + ok = ok && wo < 1e-5 && ws < 1e-5; + } + fails += !ok; + std::printf(" %s\n", ok ? "PASS" : "FAIL"); + } + } + } if (fails || !bench) { std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); return fails ? 1 : 0; @@ -149,14 +429,14 @@ int main(int argc, char** argv) { mtl::Arg::B(cw), mtl::Arg::B(dt), mtl::Arg::B(sa), mtl::Arg::B(gam), mtl::Arg::B(yA), mtl::Arg::V(da), mtl::Arg::B(hA2)}; d.grid[0] = HV; d.tg[0] = 512; - auto timeit = [&](const std::vector& ds) { - cx.warm(ds, 0.3); + auto timeit = [&](const std::vector& ds, int reps = 1) { + cx.warm(ds, 0.08); std::vector t; - for (int i = 0; i < 5; ++i) t.push_back(cx.run(ds)); + for (int i = 0; i < 5; ++i) t.push_back(cx.run(ds, reps)); return median(t); }; std::printf("T %d: decode step over the chunk %.2f ms\n", T, timeit({d}) * 1e3); - for (const char* nm : {"gdn_scanr_4", "chunked", "gdn_scanr_4", "chunked"}) { + for (const char* nm : {"gdn_scanr_4", "chunked"}) { const bool ch = std::string(nm) == "chunked"; const int cb = ch ? -100 : nm[9] == 'v' ? -8 : nm[8] == 'r' ? -(nm[10] - '0') : nm[9] - '0'; mtl::Constants pc; @@ -171,5 +451,30 @@ int main(int argc, char** argv) { std::printf("T %d, %-13s: pre %.3f + scan %.3f + post %.3f ms = %.3f ms together\n", T, nm, tp * 1e3, ts * 1e3, to * 1e3, all * 1e3); } + mtl::Pipeline scanr4 = cx.pipeline(lib, "gdn_scanr_4", mtl::Constants().i(6, 0)), + scanr4_16 = cx.pipeline(lib, "gdn_scanr_4_16", mtl::Constants().i(6, 0)); + for (int Tb : {512, 1024, 2048}) { + GdnPArgs a0{(uint32_t)Tb, HK, HV, CH, 1e-6f, 0u}; + mtl::Dispatch dpre, dscan, dscan16, dpost; + dpre.p = ⪯ + dpre.args = {mtl::Arg::B(qkv), mtl::Arg::B(hB), mtl::Arg::B(cw), mtl::Arg::B(al), mtl::Arg::B(be), mtl::Arg::B(dt), + mtl::Arg::B(sa), mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(vc), mtl::Arg::B(gb), mtl::Arg::B(kq), + mtl::Arg::B(hA2), mtl::Arg::V(a0)}; + dpre.grid[0] = (HK + HV) / 8; dpre.grid[1] = Tb; dpre.tg[0] = 256; + dscan.p = &scanr4; + dscan.args = {mtl::Arg::B(stB), mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(vc), mtl::Arg::B(gb), mtl::Arg::B(kq), + mtl::Arg::B(o), mtl::Arg::V(a0)}; + dscan.grid[0] = HV * 4; dscan.tg[0] = 128; + dscan16.p = &scanr4_16; + dscan16.args = dscan.args; + dscan16.grid[0] = HV * 4; dscan16.tg[0] = 128; + dpost.p = &post; + dpost.args = {mtl::Arg::B(o), mtl::Arg::B(z), mtl::Arg::B(gam), mtl::Arg::B(yh), mtl::Arg::B(yB), mtl::Arg::V(a0)}; + dpost.grid[0] = HV / 4; dpost.grid[1] = Tb; dpost.tg[0] = 128; + const double tpre = timeit({dpre}, 36) * 1e3, tscan = timeit({dscan}, 36) * 1e3, + tscan16 = timeit({dscan16}, 36) * 1e3, tpost = timeit({dpost}, 36) * 1e3; + std::printf("T %4d, 36 GDN layers (f32=0): gdn_pre %.2f ms (%.3f), gdn_scanr_4 %.2f ms (%.3f vs TB=16 %.2f ms / %.3f), gdn_post %.2f ms (%.3f)\n", + Tb, tpre, tpre / 36.0, tscan, tscan / 36.0, tscan16, tscan16 / 36.0, tpost, tpost / 36.0); + } return 0; } diff --git a/metal/tests/gemv_parity.cpp b/metal/tests/gemv_parity.cpp index 5687527f..93d57603 100644 --- a/metal/tests/gemv_parity.cpp +++ b/metal/tests/gemv_parity.cpp @@ -4,6 +4,13 @@ // gemv_parity parity: every format, n_in 2560/640/10240, 1..8 tokens // gemv_parity --bench [--pack Q2_0|IQ3_XXS] [--ntok N] [--all] speed on the pack's shapes (best tile per shape) // gemv_parity --one TYPE NIN NOUT [--ntok N] sweep one shape +// gemv_parity --ab BASE_DIR TYPE NIN NOUT NTOK MODE [L R S K] bitwise A/B, paired hot-weight timings +// MODE: gemv | pair | cat | catk; GEMV_AB_PARITY_ONLY=1 skips timing; GEMV_AB_COLD=1 (gemv): the timed +// dispatches walk copies of W (up to 512 MB), so the weights come from DRAM as in a decode +// gemv_parity --tiles TYPE NIN NOUT NTOK L,R,S,K ... paired timings of several tiles (GEMV_AB_COLD=1; +// GEMV_BESIDE=TYPE,NOUT,L,R,S,K: each tile as k + v beside a big GEMV) +// gemv_parity --cat TYPE NIN NOUT NTOK L,R,S gemv_cat vs its two gemvs: bitwise, paired timings +// gemv_parity --pred TYPE NIN NOUT NTOK L,R,S token 0 of two gemvs #NTOK vs #1 vs gemv_cat #1 #include "gguf_lite.h" #include "mtl_rt.h" #include "ref_quant.h" @@ -13,6 +20,7 @@ #include #include #include +#include #include #include @@ -62,7 +70,7 @@ mtl::Dispatch xprep_dispatch(const mtl::Pipeline& p, const mtl::Buffer& x, const int type_of(const std::string& s) { for (int t : {ref::Q2_0, ref::Q4_0, ref::Q5_0, ref::Q8_0, ref::Q3_K, ref::Q4_K, ref::Q5_K, ref::Q6_K, ref::IQ4_NL, - ref::IQ4_XS, ref::IQ3_S, ref::BF16}) + ref::IQ4_XS, ref::IQ3_S, ref::BF16, ref::F32}) if (s == ref::type_name(t)) return t; std::fprintf(stderr, "unknown type %s\n", s.c_str()); std::exit(2); @@ -70,9 +78,409 @@ int type_of(const std::string& s) { struct Shape { const char* name; int type, n_in, n_out, per_token; }; // per_token: calls per decoded token +// A short, paired comparison with a source snapshot, including byte-for-byte outputs. +int ab(int argc, char** argv) { + if (argc != 8 && argc != 12) return 2; + const int type = std::string(argv[3]) == "Q2_0R" ? ref::Q2_0 : type_of(argv[3]); + const bool raw_q2 = std::string(argv[3]) == "Q2_0R"; + const int kernel_type = raw_q2 ? 142 : type; + const int ni = std::atoi(argv[4]), no = std::atoi(argv[5]), nt = std::atoi(argv[6]); + const std::string mode = argv[7]; + if (ni <= 0 || no <= 0 || nt < 1 || nt > 8 || ni % ref::block_elems(type) || + (mode != "gemv" && mode != "pair" && mode != "cat" && mode != "catk")) return 2; + Cfg c{32, 2, 4, 1}; + if (mode == "pair") c = {32, 1, 4, 1}; + if (mode == "catk") c = {32, 1, 4, 4}; + if (argc == 12) c = {std::atoi(argv[8]), std::atoi(argv[9]), std::atoi(argv[10]), std::atoi(argv[11])}; + const bool known_cfg = std::any_of(kCfgs.begin(), kCfgs.end(), [&](Cfg k) { + return k.L == c.L && k.R == c.R && k.S == c.S && k.K == c.K; + }); + if (!known_cfg || !cfg_ok(c, type, ni) || + (mode == "pair" && (raw_q2 || type != ref::Q2_0 || c.L != 32 || c.R != 1 || c.S != 4 || c.K != 1)) || + (mode == "cat" && ((type != ref::Q8_0 && type != ref::BF16 && type != ref::F32) || c.K != 1 || + !((c.L == 32 && c.R == 1 && (c.S == 4 || c.S == 8)) || + (c.L == 16 && c.R == 1 && c.S == 4) || + (c.L == 32 && c.R == 2 && c.S == 4) || + (type == ref::Q8_0 && c.L == 8 && c.R == 1 && c.S == 8)))) || + (mode == "catk" && (type != ref::Q8_0 || c.L != 32 || + !((c.R == 1 && c.S == 4 && c.K == 4) || + (c.R == 2 && c.S == 4 && c.K == 2) || + (c.R == 1 && c.S == 8 && c.K == 8))))) return 2; + mtl::Context cx; + const std::string current = cx.kernel_dir; + cx.kernel_dir = argv[2]; + const int lb = cx.compile({"quant.metal", "units.metal", "gemv.metal"}); + cx.kernel_dir = current; + const int lv = cx.compile({"quant.metal", "units.metal", "gemv.metal"}); + const size_t rb = ref::row_bytes(type, ni), mb = rb * no; + const bool cold = std::getenv("GEMV_AB_COLD") && mode == "gemv"; + const int copies = cold ? (int)std::max(4, std::min(64, (512ull << 20) / mb)) : 1; + auto w = cx.alloc(mb * copies), w2 = cx.alloc(mb), x = cx.alloc(4ull * ni * nt); + auto xp = cx.alloc(4ull * (4 * ni * nt + ni / 16 * nt)); + auto yb = cx.alloc(4ull * no * nt), yv = cx.alloc(4ull * no * nt); + auto zb = cx.alloc(4ull * no * nt), zv = cx.alloc(4ull * no * nt); + ref::random_blocks(type, w.as(), no, ni, 99); + ref::random_blocks(type, w2.as(), no, ni, 100); + for (int k = 1; k < copies; ++k) std::memcpy(w.as() + k * mb, w.ptr, mb); + ref::Rng rng(5); + for (int i = 0; i < ni * nt; ++i) x.as()[i] = rng.sym(); + auto prep = cx.pipeline(lb, "xprep"); + cx.run({xprep_dispatch(prep, x, xp, ni, nt)}); + const size_t plane = 4ull * ni * nt; + std::string name = kname(kernel_type, c); + if (mode == "pair") name = "gemv_pair_" + std::to_string(kernel_type) + "_" + std::to_string(kernel_type); + if (mode == "cat" || mode == "catk") { + name = (mode == "cat" ? "gemvc_" : "gemvck_") + std::to_string(kernel_type) + "_" + + std::to_string(c.L) + "_" + std::to_string(c.R) + "_" + std::to_string(c.S); + if (mode == "catk") name += "_" + std::to_string(c.K); + } + auto pb = cx.pipeline(lb, name, mtl::Constants().i(0, nt)); + auto pv = cx.pipeline(lv, name, mtl::Constants().i(0, nt)); + GemvArgs a{(uint32_t)ni, (uint32_t)(mode == "cat" || mode == "catk" ? no * 2 : no), + (uint32_t)nt, (uint32_t)rb, (uint32_t)ni, (uint32_t)no}; + auto dispatch = [&](const mtl::Pipeline& p, const mtl::Buffer& y, const mtl::Buffer& z, int copy = 0) { + mtl::Dispatch d; + d.p = &p; + const size_t off = plane * (raw_q2 ? 0 : xkind(type)); + d.args = {mtl::Arg::B(w, copy * mb), mtl::Arg::B(xp, off), mtl::Arg::B(xp, plane * 4), + mtl::Arg::B(y), mtl::Arg::V(a)}; + if (mode == "pair") + d.args = {mtl::Arg::B(w), mtl::Arg::B(xp, off), mtl::Arg::B(w2), mtl::Arg::B(xp, off), + mtl::Arg::B(xp, plane * 4), mtl::Arg::B(y), mtl::Arg::V(a), mtl::Arg::V((uint32_t)rb)}; + if (mode == "cat" || mode == "catk") { + d.args.push_back(mtl::Arg::B(w2)); d.args.push_back(mtl::Arg::B(z)); + d.args.push_back(mtl::Arg::V((uint32_t)no)); + } + const int rows = mode == "pair" ? 4 : rows_per_tg(c); + d.grid[0] = (a.n_out + rows - 1) / rows; + d.tg[0] = c.S * 32; + return d; + }; + auto b = dispatch(pb, yb, zb), v = dispatch(pv, yv, zv); + cx.run({b}); cx.run({v}); + const bool exact = !std::memcmp(yb.ptr, yv.ptr, yb.bytes) && !std::memcmp(zb.ptr, zv.ptr, zb.bytes); + std::printf("parity %s %s %d %d nt %d tile %d,%d,%d,%d %s\n", mode.c_str(), argv[3], ni, no, nt, + c.L, c.R, c.S, c.K, exact ? "bitwise" : "FAIL"); + if (!exact) return 1; + if (std::getenv("GEMV_AB_PARITY_ONLY")) return 0; + cx.warm({b, v}, 0.15); + std::vector bs, vs; + for (int k = 0; k < copies; ++k) { bs.push_back(dispatch(pb, yb, zb, k)); vs.push_back(dispatch(pv, yv, zv, k)); } + const int reps = cold ? std::max(1, 64 / copies) : 64, per = reps * copies; + for (int pair = 0; pair < 15; ++pair) { + double tb, tv; + if (pair % 2) { tv = cx.run(vs, reps); tb = cx.run(bs, reps); } + else { tb = cx.run(bs, reps); tv = cx.run(vs, reps); } + std::printf("ab %d %.6f %.6f\n", pair, tb * 1e6 / per, tv * 1e6 / per); + } + return 0; +} + +// Paired tile timings on one shape: every listed tile in each round (order rotated), the weights from DRAM +// (GEMV_AB_COLD=1: the dispatches walk copies of W, up to 512 MB) or hot. Each dispatch behind a barrier (a decode's +// GEMV waits for its x). Prints each tile's median us and its rel L2 against double. +int tiles(int argc, char** argv) { + if (argc < 7) return 2; + const bool raw_q2 = std::string(argv[2]) == "Q2_0R"; + const int type = raw_q2 ? ref::Q2_0 : type_of(argv[2]), kernel_type = raw_q2 ? 142 : type; + const int ni = std::atoi(argv[3]), no = std::atoi(argv[4]), nt = std::atoi(argv[5]); + if (ni <= 0 || no <= 0 || nt < 1 || nt > 8 || ni % ref::block_elems(type)) return 2; + std::vector cs; + for (int i = 6; i < argc; ++i) { + Cfg c{}; + if (std::sscanf(argv[i], "%d,%d,%d,%d", &c.L, &c.R, &c.S, &c.K) != 4 || !cfg_ok(c, type, ni)) return 2; + cs.push_back(c); + } + mtl::Context cx; + const int lib = cx.compile({"quant.metal", "units.metal", "gemv.metal"}); + const size_t rb = ref::row_bytes(type, ni), mb = rb * no; + const bool cold = std::getenv("GEMV_AB_COLD") != nullptr; + const int copies = cold ? (int)std::max(4, std::min(256, (512ull << 20) / mb)) : 1; + auto w = cx.alloc(mb * copies), x = cx.alloc(4ull * ni * nt), y = cx.alloc(4ull * no * nt); + auto xp = cx.alloc(4ull * (4 * ni * nt + ni / 16 * nt)); + ref::random_blocks(type, w.as(), no, ni, 99); + for (int k = 1; k < copies; ++k) std::memcpy(w.as() + k * mb, w.ptr, mb); + ref::Rng rng(5); + for (int i = 0; i < ni * nt; ++i) x.as()[i] = rng.sym(); + auto prep = cx.pipeline(lib, "xprep"); + cx.run({xprep_dispatch(prep, x, xp, ni, nt)}); + if (raw_q2) std::memcpy(xp.ptr, x.ptr, 4ull * ni * nt); // Q2_0R reads raw x + const size_t plane = 4ull * ni * nt; + std::vector wf((size_t)ni * no); + ref::dequant(type, w.ptr, wf.data(), wf.size()); + GemvArgs a{(uint32_t)ni, (uint32_t)no, (uint32_t)nt, (uint32_t)rb, (uint32_t)ni, (uint32_t)no}; + std::vector ps; + for (Cfg c : cs) ps.push_back(cx.pipeline(lib, kname(kernel_type, c), mtl::Constants().i(0, nt))); + std::vector> runs(cs.size()); + for (size_t i = 0; i < cs.size(); ++i) + for (int k = 0; k < copies; ++k) { + mtl::Dispatch d; + d.p = &ps[i]; + d.args = {mtl::Arg::B(w, k * mb), mtl::Arg::B(xp, raw_q2 ? 0 : plane * xkind(type)), mtl::Arg::B(xp, plane * 4), + mtl::Arg::B(y), mtl::Arg::V(a)}; + d.grid[0] = (no + rows_per_tg(cs[i]) - 1) / rows_per_tg(cs[i]); + d.tg[0] = cs[i].S * 32; + runs[i].push_back(d); + } + std::vector rel(cs.size()); + for (size_t i = 0; i < cs.size(); ++i) { + cx.run({runs[i][0]}); + double se = 0, sr = 0; + for (int t = 0; t < nt; ++t) + for (int o = 0; o < no; ++o) { + double r = 0; + for (int k = 0; k < ni; ++k) r += (double)wf[(size_t)o * ni + k] * x.as()[t * ni + k]; + const double e = y.as()[t * no + o] - r; + se += e * e; + sr += r * r; + } + rel[i] = std::sqrt(se / sr); + } + // GEMV_BESIDE=TYPE,NOUT,L,R,S,K: each timed dispatch runs twice (as the k and v beside each other) beside one + // GEMV of that shape over the same n_in (as attn_q), all behind one barrier: the tile's cost inside the group + mtl::Buffer bw, by; + mtl::Pipeline bp; + GemvArgs ba{}; + if (const char* bs = std::getenv("GEMV_BESIDE")) { + char bt[16] = {}; + int bno = 0; + Cfg bc{}; + if (std::sscanf(bs, "%15[^,],%d,%d,%d,%d,%d", bt, &bno, &bc.L, &bc.R, &bc.S, &bc.K) != 6) return 2; + const int btype = type_of(bt); + const size_t bmb = ref::row_bytes(btype, ni) * bno; + const int bcopies = cold ? (int)std::max(2, std::min(64, (512ull << 20) / bmb)) : 1; + bw = cx.alloc(bmb * bcopies); + by = cx.alloc(4ull * bno * nt); + ref::random_blocks(btype, bw.as(), bno, ni, 7); + for (int k = 1; k < bcopies; ++k) std::memcpy(bw.as() + k * bmb, bw.ptr, bmb); + bp = cx.pipeline(lib, kname(btype, bc), mtl::Constants().i(0, nt)); + ba = {(uint32_t)ni, (uint32_t)bno, (uint32_t)nt, (uint32_t)ref::row_bytes(btype, ni), (uint32_t)ni, (uint32_t)bno}; + for (auto& r : runs) { + std::vector g; + for (int k = 0; k < copies; ++k) { + mtl::Dispatch d; + d.p = &bp; + d.args = {mtl::Arg::B(bw, (size_t)(k % bcopies) * bmb), mtl::Arg::B(xp, plane * xkind(btype)), + mtl::Arg::B(xp, plane * 4), mtl::Arg::B(by), mtl::Arg::V(ba)}; + d.grid[0] = (bno + rows_per_tg(bc) - 1) / rows_per_tg(bc); + d.tg[0] = bc.S * 32; + g.push_back(d); + mtl::Dispatch t1 = r[k], t2 = r[(k + copies / 2) % copies]; + t1.barrier = t2.barrier = false; + g.push_back(t1); + g.push_back(t2); + } + r = g; + } + } + const int reps = cold ? std::max(1, 256 / copies) : 256; + std::vector all; + for (auto& r : runs) all.insert(all.end(), r.begin(), r.end()); + cx.warm(all, 0.2); + std::vector> t(cs.size()); + for (int round = 0; round < 15; ++round) + for (size_t j = 0; j < cs.size(); ++j) { + const size_t i = (j + round) % cs.size(); + std::vector ds; + for (int r = 0; r < reps; ++r) ds.insert(ds.end(), runs[i].begin(), runs[i].end()); + t[i].push_back(cx.run_graph(ds) * 1e6 / (reps * copies)); + } + for (size_t i = 0; i < cs.size(); ++i) + std::printf("tile %s %d %d nt %d %2d,%d,%d,%d %s: %8.3f us rel L2 %.2e\n", argv[2], ni, no, nt, cs[i].L, cs[i].R, + cs[i].S, cs[i].K, cold ? "cold" : "hot", median(t[i]), rel[i]); + return 0; +} + +// gemv_cat against the two gemv dispatches it replaces, at one tile (KS = 1): bitwise, then paired timings with each +// pair behind a barrier as in a decode (the two gemvs beside each other, no barrier between). GEMV_AB_COLD=1: the +// dispatches walk copies of both matrices (up to 512 MB). +int cat2(int argc, char** argv) { + if (argc != 7) return 2; + const int type = type_of(argv[2]), ni = std::atoi(argv[3]), no = std::atoi(argv[4]), nt = std::atoi(argv[5]); + Cfg c{}; + if (std::sscanf(argv[6], "%d,%d,%d", &c.L, &c.R, &c.S) != 3) return 2; + c.K = 1; + if (ni <= 0 || no <= 0 || nt < 1 || nt > 8 || ni % ref::block_elems(type) || !cfg_ok(c, type, ni)) return 2; + mtl::Context cx; + const int lib = cx.compile({"quant.metal", "units.metal", "gemv.metal"}); + const size_t rb = ref::row_bytes(type, ni), mb = rb * no; + const bool cold = std::getenv("GEMV_AB_COLD") != nullptr; + const int copies = cold ? (int)std::max(4, std::min(128, (256ull << 20) / mb)) : 1; + auto w = cx.alloc(mb * copies), w2 = cx.alloc(mb * copies), x = cx.alloc(4ull * ni * nt); + auto xp = cx.alloc(4ull * (4 * ni * nt + ni / 16 * nt)); + auto ya = cx.alloc(4ull * no * nt), yb = cx.alloc(4ull * no * nt), za = cx.alloc(4ull * no * nt), + zb = cx.alloc(4ull * no * nt); + ref::random_blocks(type, w.as(), no, ni, 99); + ref::random_blocks(type, w2.as(), no, ni, 100); + for (int k = 1; k < copies; ++k) { + std::memcpy(w.as() + k * mb, w.ptr, mb); + std::memcpy(w2.as() + k * mb, w2.ptr, mb); + } + ref::Rng rng(5); + for (int i = 0; i < ni * nt; ++i) x.as()[i] = rng.sym(); + auto prep = cx.pipeline(lib, "xprep"); + cx.run({xprep_dispatch(prep, x, xp, ni, nt)}); + const size_t plane = 4ull * ni * nt; + const std::string cn = "gemvc_" + std::to_string(type) + "_" + std::to_string(c.L) + "_" + std::to_string(c.R) + "_" + + std::to_string(c.S); + auto pg = cx.pipeline(lib, kname(type, c), mtl::Constants().i(0, nt)); + auto pc = cx.pipeline(lib, cn, mtl::Constants().i(0, nt)); + const int rows = rows_per_tg(c); + GemvArgs a{(uint32_t)ni, (uint32_t)no, (uint32_t)nt, (uint32_t)rb, (uint32_t)ni, (uint32_t)no}; + GemvArgs a2 = a; + a2.n_out = 2 * no; + auto gemv = [&](const mtl::Buffer& W, int k, const mtl::Buffer& y, bool barrier) { + mtl::Dispatch d; + d.p = &pg; + d.args = {mtl::Arg::B(W, k * mb), mtl::Arg::B(xp, plane * xkind(type)), mtl::Arg::B(xp, plane * 4), + mtl::Arg::B(y), mtl::Arg::V(a)}; + d.grid[0] = (no + rows - 1) / rows; + d.tg[0] = c.S * 32; + d.barrier = barrier; + return d; + }; + auto cat = [&](int k) { + mtl::Dispatch d; + d.p = &pc; + d.args = {mtl::Arg::B(w, k * mb), mtl::Arg::B(xp, plane * xkind(type)), mtl::Arg::B(xp, plane * 4), + mtl::Arg::B(yb), mtl::Arg::V(a2), mtl::Arg::B(w2, k * mb), mtl::Arg::B(zb), + mtl::Arg::V((uint32_t)no)}; + d.grid[0] = (2 * no + rows - 1) / rows; + d.tg[0] = c.S * 32; + return d; + }; + std::vector two, one; + for (int k = 0; k < copies; ++k) { + two.push_back(gemv(w, k, ya, true)); + two.push_back(gemv(w2, k, za, false)); + one.push_back(cat(k)); + } + cx.run_graph({two[0], two[1]}); + cx.run_graph({one[0]}); + const bool exact = !std::memcmp(ya.ptr, yb.ptr, ya.bytes) && !std::memcmp(za.ptr, zb.ptr, za.bytes); + std::printf("parity cat %s %d %d+%d nt %d tile %d,%d,%d: %s\n", argv[2], ni, no, no, nt, c.L, c.R, c.S, + exact ? "bitwise" : "FAIL"); + if (!exact) return 1; + if (std::getenv("GEMV_AB_PARITY_ONLY")) return 0; + const int reps = cold ? std::max(1, 256 / copies) : 256; + std::vector ts, os; + for (int r = 0; r < reps; ++r) { + ts.insert(ts.end(), two.begin(), two.end()); + os.insert(os.end(), one.begin(), one.end()); + } + cx.warm(ts, 0.1); + cx.warm(os, 0.1); + std::vector tt, to; + for (int pair = 0; pair < 15; ++pair) { + if (pair % 2) { to.push_back(cx.run_graph(os)); tt.push_back(cx.run_graph(ts)); } + else { tt.push_back(cx.run_graph(ts)); to.push_back(cx.run_graph(os)); } + } + const double per = 1e6 / (reps * copies), m2 = median(tt) * per, m1 = median(to) * per; + std::printf("cat %s %d %d+%d nt %d tile %d,%d,%d %s: 2 gemv %.3f us, gemv_cat %.3f us (%+.1f%%)\n", argv[2], ni, no, + no, nt, c.L, c.R, c.S, cold ? "cold" : "hot", m2, m1, (m1 / m2 - 1) * 100); + return 0; +} + +// the read-ahead's router guesses (window.cpp predL_ / predL2_, Q8_0 2560 -> 512 for layers l + 1 and l + 2) when only +// token 0's guess is read (STRATA_METAL_PREFETCH_DRAFT=0): the two window-wide gemvs at NTOK = nt against the same two +// at NTOK = 1 on the window's tile, and gemv_cat of both at NTOK = 1. Token 0 bitwise, then paired timings (each +// variant behind one barrier, its two gemvs beside each other). GEMV_AB_COLD=1: the weights from DRAM. +int pred(int argc, char** argv) { + if (argc != 7) return 2; + const int type = type_of(argv[2]), ni = std::atoi(argv[3]), no = std::atoi(argv[4]), nt = std::atoi(argv[5]); + Cfg c{}; + if (std::sscanf(argv[6], "%d,%d,%d", &c.L, &c.R, &c.S) != 3) return 2; + c.K = 1; + if (ni <= 0 || no <= 0 || nt < 1 || nt > 8 || ni % ref::block_elems(type) || !cfg_ok(c, type, ni)) return 2; + mtl::Context cx; + const int lib = cx.compile({"quant.metal", "units.metal", "gemv.metal"}); + const size_t rb = ref::row_bytes(type, ni), mb = rb * no; + const bool cold = std::getenv("GEMV_AB_COLD") != nullptr; + const int copies = cold ? (int)std::max(4, std::min(128, (256ull << 20) / mb)) : 1; + auto w = cx.alloc(mb * copies), w2 = cx.alloc(mb * copies), x = cx.alloc(4ull * ni * nt); + auto xp = cx.alloc(4ull * (4 * ni * nt + ni / 16 * nt)); + std::vector ys, zs; + for (int v = 0; v < 3; ++v) { ys.push_back(cx.alloc(4ull * no * nt)); zs.push_back(cx.alloc(4ull * no * nt)); } + ref::random_blocks(type, w.as(), no, ni, 99); + ref::random_blocks(type, w2.as(), no, ni, 100); + for (int k = 1; k < copies; ++k) { + std::memcpy(w.as() + k * mb, w.ptr, mb); + std::memcpy(w2.as() + k * mb, w2.ptr, mb); + } + ref::Rng rng(5); + for (int i = 0; i < ni * nt; ++i) x.as()[i] = rng.sym(); + auto prep = cx.pipeline(lib, "xprep"); + cx.run({xprep_dispatch(prep, x, xp, ni, nt)}); + const size_t plane = 4ull * ni * nt; + auto pw = cx.pipeline(lib, kname(type, c), mtl::Constants().i(0, nt)); + auto p1 = cx.pipeline(lib, kname(type, c), mtl::Constants().i(0, 1)); + auto pc = cx.pipeline(lib, "gemvc_" + std::to_string(type) + "_" + std::to_string(c.L) + "_" + std::to_string(c.R) + + "_" + std::to_string(c.S), mtl::Constants().i(0, 1)); + const int rows = rows_per_tg(c); + auto gemv = [&](const mtl::Pipeline& p, int n, const mtl::Buffer& W, int k, const mtl::Buffer& y, bool barrier) { + mtl::Dispatch d; + d.p = &p; + const GemvArgs a{(uint32_t)ni, (uint32_t)no, (uint32_t)n, (uint32_t)rb, (uint32_t)ni, (uint32_t)no}; + d.args = {mtl::Arg::B(W, k * mb), mtl::Arg::B(xp, plane * xkind(type)), mtl::Arg::B(xp, plane * 4), + mtl::Arg::B(y), mtl::Arg::V(a)}; + d.grid[0] = (no + rows - 1) / rows; + d.tg[0] = c.S * 32; + d.barrier = barrier; + return d; + }; + auto cat = [&](int k) { + mtl::Dispatch d; + d.p = &pc; + const GemvArgs a{(uint32_t)ni, (uint32_t)(2 * no), 1, (uint32_t)rb, (uint32_t)ni, (uint32_t)no}; + d.args = {mtl::Arg::B(w, k * mb), mtl::Arg::B(xp, plane * xkind(type)), mtl::Arg::B(xp, plane * 4), + mtl::Arg::B(ys[2]), mtl::Arg::V(a), mtl::Arg::B(w2, k * mb), mtl::Arg::B(zs[2]), + mtl::Arg::V((uint32_t)no)}; + d.grid[0] = (2 * no + rows - 1) / rows; + d.tg[0] = c.S * 32; + return d; + }; + std::vector> runs(3); + for (int k = 0; k < copies; ++k) { + runs[0].push_back(gemv(pw, nt, w, k, ys[0], true)); + runs[0].push_back(gemv(pw, nt, w2, k, zs[0], false)); + runs[1].push_back(gemv(p1, 1, w, k, ys[1], true)); + runs[1].push_back(gemv(p1, 1, w2, k, zs[1], false)); + runs[2].push_back(cat(k)); + } + for (int v = 0; v < 3; ++v) cx.run_graph(std::vector(runs[v].begin(), runs[v].begin() + (v < 2 ? 2 : 1))); + bool exact = true; + for (int v = 1; v < 3; ++v) + exact = exact && !std::memcmp(ys[0].ptr, ys[v].ptr, 4ull * no) && !std::memcmp(zs[0].ptr, zs[v].ptr, 4ull * no); + std::printf("parity pred %s %d %d+%d tile %d,%d,%d: token 0 of #%d == #1 == gemvc #1: %s\n", argv[2], ni, no, no, + c.L, c.R, c.S, nt, exact ? "bitwise" : "FAIL"); + if (!exact) return 1; + if (std::getenv("GEMV_AB_PARITY_ONLY")) return 0; + const int reps = cold ? std::max(1, 256 / copies) : 256; + std::vector> rep(3); + for (int v = 0; v < 3; ++v) { + for (int r = 0; r < reps; ++r) rep[v].insert(rep[v].end(), runs[v].begin(), runs[v].end()); + cx.warm(rep[v], 0.1); + } + std::vector> t(3); + for (int round = 0; round < 15; ++round) + for (int j = 0; j < 3; ++j) { + const int v = (j + round) % 3; + t[v].push_back(cx.run_graph(rep[v])); + } + const double per = 1e6 / (reps * copies), m0 = median(t[0]) * per, m1 = median(t[1]) * per, m2 = median(t[2]) * per; + std::printf("pred %s %d %d+%d tile %d,%d,%d %s: 2 gemv #%d %.3f us, 2 gemv #1 %.3f us (%+.1f%%), gemvc #1 %.3f us " + "(%+.1f%%)\n", argv[2], ni, no, no, c.L, c.R, c.S, cold ? "cold" : "hot", nt, m0, m1, (m1 / m0 - 1) * 100, + m2, (m2 / m0 - 1) * 100); + return 0; +} + } // namespace int main(int argc, char** argv) { + if (argc > 1 && std::string(argv[1]) == "--ab") return ab(argc, argv); + if (argc > 1 && std::string(argv[1]) == "--tiles") return tiles(argc, argv); + if (argc > 1 && std::string(argv[1]) == "--cat") return cat2(argc, argv); + if (argc > 1 && std::string(argv[1]) == "--pred") return pred(argc, argv); bool bench = false, all = false; int ntok = 1; std::string gguf_path; @@ -180,8 +588,9 @@ int main(int argc, char** argv) { ref::dequant(t, W.ptr, wf.data(), wf.size()); for (int nt : {1, 3, 8}) { double worst = 0; - int not_exact = 0; + int not_exact = 0, not_same = 0; std::string worst_cfg; + std::map> by_lpr; // KS = 1: a row's sums depend only on LPR for (Cfg c : kCfgs) { if (!cfg_ok(c, t, n_in)) continue; mtl::Pipeline p = cx.pipeline(lib, kname(t, c), mtl::Constants().i(0, nt)); @@ -206,6 +615,11 @@ int main(int argc, char** argv) { } const double rel = std::sqrt(se / sr); if (rel > worst) { worst = rel; worst_cfg = kname(t, c); } + if (c.K == 1) { + const float* yv = (const float*)y.ptr; + auto [it, first] = by_lpr.try_emplace(c.L, yv, yv + nt * n_out); + if (!first && std::memcmp(it->second.data(), yv, 4ull * nt * n_out) != 0) ++not_same; + } if (nt > 1) { // every token of a window must be bitwise what a one-token call gives std::vector ywin((float*)y.ptr, (float*)y.ptr + nt * n_out); @@ -223,12 +637,13 @@ int main(int argc, char** argv) { } } } - const bool ok = worst < 1e-5 && not_exact == 0; + const bool ok = worst < 1e-5 && not_exact == 0 && not_same == 0; fails += !ok; std::printf("parity %-7s n_in %5d ntok %d: worst rel L2 over %zu tiles %.2e (%s), window == one " - "token: %s %s\n", + "token: %s, same LPR: %s %s\n", ref::type_name(t), n_in, nt, kCfgs.size(), worst, worst_cfg.c_str(), - nt == 1 ? "-" : not_exact ? "NO" : "bitwise", ok ? "PASS" : "FAIL"); + nt == 1 ? "-" : not_exact ? "NO" : "bitwise", not_same ? "NO" : "bitwise", + ok ? "PASS" : "FAIL"); } cx.free(W); cx.free(x); diff --git a/metal/tests/moe_mm_parity.cpp b/metal/tests/moe_mm_parity.cpp index 6f51f1e7..a44397cf 100644 --- a/metal/tests/moe_mm_parity.cpp +++ b/metal/tests/moe_mm_parity.cpp @@ -3,6 +3,8 @@ // units' operands) - and against the exact float reference for the size of that rounding; then the speed. // // moe_mm_parity [--gu Q2_0] [--dn Q2_0] [--T 256] [--tile 32,64,64,4] parity (64 experts) +// moe_mm_parity --gather-bench [T=2048] [experts=512] [width=2560] short plan/copy sweep +// MM_BASELINE_KERNEL=/path/to/moe_mm.metal compares every plan buffer byte-for-byte // moe_mm_parity --bench [--gu ..] [--dn ..] [--T 2048] [--tile ..] per-layer time, 512 experts #include "coop_probe.h" #include "mtl_rt.h" @@ -28,30 +30,190 @@ int type_of(const std::string& s) { std::exit(2); } float h16(float v) { return ref::f16_to_f32(ref::f32_to_f16(v)); } + +int gather_bench(int T, int NE, int n) { + if (T <= 0 || NE < TOPK || NE > 1024 || NE % 32 || n <= 0 || n % 8) { + std::fprintf(stderr, "gather-bench needs T > 0, experts in [32, 1024] divisible by 32, and width divisible by 8\n"); + return 2; + } + mtl::Context cx; + const int lib = cx.compile({"moe_mm.metal"}, "#define NO_GGML_COMMON 1\n#define MM_GATHER_ONLY 1\n"); + mtl::Pipeline phist = cx.pipeline(lib, "mm_hist"), pscan = cx.pipeline(lib, "mm_scan"), pplace = cx.pipeline(lib, "mm_place"), + pgat = cx.pipeline(lib, "mm_gather_h"), pgat_f = cx.pipeline(lib, "mm_gather"); + const char* baseline = std::getenv("MM_BASELINE_KERNEL"); + const int bl = baseline ? cx.compile({baseline}, "#define NO_GGML_COMMON 1\n#define MM_GATHER_ONLY 1\n") : lib; + mtl::Pipeline bh = cx.pipeline(bl, "mm_hist"), bs = cx.pipeline(bl, "mm_scan"), bp = cx.pipeline(bl, "mm_place"), + bg = cx.pipeline(bl, "mm_gather_h"), bf = cx.pipeline(bl, "mm_gather"); + std::printf("device %s; mm_plan + mm_gather_h bitwise parity & benchmark\n", cx.device_name().c_str()); + const int TM = 16; + const int nent = T * TOPK, nblk = (nent + 1023) / 1024, maxrows = nent + NE * TM; + mtl::Buffer x = cx.alloc(4 * (size_t)T * n), xh = cx.alloc(2 * (size_t)T * n), + ids = cx.alloc(4 * nent), w = cx.alloc(4 * nent), + bcnt = cx.alloc(4 * (size_t)nblk * NE), boff = cx.alloc(4 * (size_t)nblk * NE), eoff = cx.alloc(4 * NE), + tiles = cx.alloc(8 * (size_t)(maxrows / TM + 1)), meta = cx.alloc(64), stok = cx.alloc(4 * (size_t)maxrows), + sw = cx.alloc(4 * (size_t)maxrows), pos = cx.alloc(4 * nent), + XS = cx.alloc(2 * (size_t)maxrows * n), XS_ref = cx.alloc(2 * (size_t)maxrows * n); + ref::Rng r(7); + for (size_t i = 0; i < (size_t)T * n; ++i) { + const float v = r.sym(); + x.as()[i] = v; + xh.as()[i] = ref::f32_to_f16(v); + } + for (int t = 0; t < T; ++t) { + std::set used; + float s = 0; + for (int k = 0; k < TOPK; ++k) { + int e; + do e = (int)(r.next() % NE); while (used.count(e)); + used.insert(e); + ids.as()[t * TOPK + k] = e; + s += (w.as()[t * TOPK + k] = 0.2f + r.uni()); + } + for (int k = 0; k < TOPK; ++k) w.as()[t * TOPK + k] /= s; + } + MmPlanArgs pa{(uint32_t)nent, TOPK, (uint32_t)NE, (uint32_t)TM, (uint32_t)nblk}; + std::vector ds_plan; + mtl::Dispatch h; h.p = &phist; h.args = {mtl::Arg::B(ids), mtl::Arg::B(bcnt), mtl::Arg::V(pa)}; h.grid[0] = nblk; h.tg[0] = 1024; ds_plan.push_back(h); + mtl::Dispatch sc; sc.p = &pscan; sc.args = {mtl::Arg::B(bcnt), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(tiles), mtl::Arg::B(meta), mtl::Arg::V(pa), mtl::Arg::B(stok)}; sc.tg[0] = NE; ds_plan.push_back(sc); + mtl::Dispatch pl; pl.p = &pplace; pl.args = {mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(stok), mtl::Arg::B(sw), mtl::Arg::B(pos), mtl::Arg::V(pa)}; pl.grid[0] = nblk; pl.tg[0] = 1024; ds_plan.push_back(pl); + bool ok_bits = true; + if (baseline) { + auto bd = ds_plan; + bd[0].p = &bh; bd[1].p = &bs; bd[2].p = &bp; + cx.run(bd); + std::vector outputs{&bcnt, &boff, &eoff, &tiles, &meta, &stok, &sw, &pos}; + std::vector sizes{4u * nblk * NE, 4u * nblk * NE, 4u * NE, + 8u * (maxrows / TM + 1), 64, 4u * maxrows, 4u * maxrows, 4u * nent}; + std::vector> expected; + for (size_t i = 0; i < outputs.size(); ++i) { + const auto* bytes = outputs[i]->as(); + expected.emplace_back(bytes, bytes + sizes[i]); + std::memset(outputs[i]->ptr, 0, sizes[i]); + } + cx.run(ds_plan); + const char* names[]{"bcnt", "boff", "eoff", "tiles", "meta", "srt_tok", "srt_w", "pos"}; + for (size_t i = 0; i < outputs.size(); ++i) { + if (std::memcmp(outputs[i]->ptr, expected[i].data(), sizes[i]) != 0) { + std::fprintf(stderr, "%s memcmp failed\n", names[i]); + ok_bits = false; + } + } + cx.warm(bd, 0.02); + cx.warm(ds_plan, 0.02); + std::vector bt, ct; + for (int i = 0; i < 7; ++i) { + if (i & 1) { + ct.push_back(cx.run(ds_plan, 8, false, true)); + bt.push_back(cx.run(bd, 8, false, true)); + } else { + bt.push_back(cx.run(bd, 8, false, true)); + ct.push_back(cx.run(ds_plan, 8, false, true)); + } + } + std::printf(" plan baseline %.3f candidate %.3f us/L; all 8 buffers memcmp=%s\n", + median(bt) * 1e6 / 8, median(ct) * 1e6 / 8, ok_bits ? "PASS" : "FAIL"); + } else cx.run(ds_plan); + const uint32_t tot_rows = meta.as()[1], ne = n; + const size_t xs_bytes = 2 * (size_t)tot_rows * n; + bool ok_plan = true; + std::vector seen(tot_rows, 0); + for (int i = 0; i < nent; ++i) { + const int row = pos.as()[i]; + if (row < 0 || row >= (int)tot_rows || stok.as()[row] != i / TOPK) ok_plan = false; + else seen[row]++; + } + for (uint32_t r0 = 0; r0 < tot_rows; ++r0) + if (seen[r0] > 1 || (seen[r0] == 0 && stok.as()[r0] != -1)) ok_plan = false; + mtl::Dispatch d_ref; d_ref.p = &pgat_f; d_ref.args = {mtl::Arg::B(x), mtl::Arg::B(stok), mtl::Arg::B(XS_ref), mtl::Arg::V(ne)}; d_ref.grid[0] = tot_rows; d_ref.tg[0] = 256; + cx.run({d_ref}); + auto b_ref = d_ref; b_ref.p = &bf; b_ref.args[2] = mtl::Arg::B(XS); + cx.run({b_ref}); + bool ok_xs = std::memcmp(XS.ptr, XS_ref.ptr, xs_bytes) == 0; + for (uint32_t tg_sz : {32u, 64u, 80u, 128u, 160u, 256u}) { + std::memset(XS.ptr, 0x55, xs_bytes); + mtl::Dispatch d; d.p = &pgat; d.args = {mtl::Arg::B(xh), mtl::Arg::B(stok), mtl::Arg::B(XS), mtl::Arg::V(ne)}; d.grid[0] = tot_rows; d.tg[0] = tg_sz; + cx.run({d}); + if (std::memcmp(XS.ptr, XS_ref.ptr, xs_bytes) != 0) ok_xs = false; + d.grid[0] = 1; d.indirect = &meta; d.ind_off = 8 * 4; + cx.run({d}); + if (std::memcmp(XS.ptr, XS_ref.ptr, xs_bytes) != 0) ok_xs = false; + auto b_gat = d; b_gat.p = &bg; + cx.warm({b_gat, d}, 0.02); + std::vector bt, ct; + for (int i = 0; i < 7; ++i) { + if (i & 1) { + ct.push_back(cx.run({d}, 8, false, true)); + bt.push_back(cx.run({b_gat}, 8, false, true)); + } else { + bt.push_back(cx.run({b_gat}, 8, false, true)); + ct.push_back(cx.run({d}, 8, false, true)); + } + } + std::printf(" T=%d NE=%d rows=%u tg=%u baseline %.3f candidate %.3f us/L XS_memcmp=%s\n", + T, NE, tot_rows, tg_sz, median(bt) * 1e6 / 8, median(ct) * 1e6 / 8, ok_xs ? "PASS" : "FAIL"); + } + auto combined = ds_plan; + mtl::Dispatch gather = d_ref; + gather.p = &pgat; + gather.args[0] = mtl::Arg::B(xh); + gather.args[2] = mtl::Arg::B(XS); + gather.indirect = &meta; gather.ind_off = 8 * 4; + combined.push_back(gather); + std::vector times; + for (int i = 0; i < 3; ++i) times.push_back(cx.run(combined, 8, false, true)); + ok_xs &= std::memcmp(XS.ptr, XS_ref.ptr, xs_bytes) == 0; + std::printf(" combined plan+gather tg256 %.3f us/L plan=%s XS_memcmp=%s\n", + median(times) * 1e6 / 8, ok_plan ? "PASS" : "FAIL", ok_xs ? "PASS" : "FAIL"); + if (!ok_plan || !ok_xs || !ok_bits) return 1; + return 0; +} + } // namespace int main(int argc, char** argv) { - int gu = ref::Q2_0, dn = ref::Q2_0, T = 256, TM = 16, TN = 64, TK = 16, SG = 2, DTN = -1, DTK = -1, DSG = -1; + if (argc > 1 && std::string(argv[1]) == "--gather-bench") { + return gather_bench(argc > 2 ? std::atoi(argv[2]) : 2048, + argc > 3 ? std::atoi(argv[3]) : 512, + argc > 4 ? std::atoi(argv[4]) : N_EMBD); + } + int gu = ref::Q2_0, dn = ref::Q2_0, T = 256, TM = 16, TN = 64, TK = 64, SG = 2, DTN = -1, DTK = -1, DSG = -1; bool bench = false, v2 = false, v3 = false, v4 = false, v5 = false; int MS = 1; + bool gu_only = false, dn_only = false; + const char* dump = nullptr, *ab_dir = nullptr; for (int i = 1; i < argc; ++i) { std::string a = argv[i]; if (a == "--gu") gu = type_of(argv[++i]); else if (a == "--dn") dn = type_of(argv[++i]); - else if (a == "--T") T = std::atoi(argv[++i]); + else if (a == "--T" || a == "-T") T = std::atoi(argv[++i]); else if (a == "--tile") std::sscanf(argv[++i], "%d,%d,%d,%d", &TM, &TN, &TK, &SG); - else if (a == "--dtile") std::sscanf(argv[++i], "%d,%d,%d", &DTN, &DTK, &DSG); // the down kernel's TN,TK,SG + else if (a == "--dtile") { + int m, n, k, s; + const char* value = argv[++i]; + if (std::sscanf(value, "%d,%d,%d,%d", &m, &n, &k, &s) == 4) { + if (m != TM) { + std::fprintf(stderr, "down tile TM must match gate/up TM\n"); + return 2; + } + DTN = n; DTK = k; DSG = s; + } else std::sscanf(value, "%d,%d,%d", &DTN, &DTK, &DSG); + } else if (a == "--bench") bench = true; + else if (a == "--gu-bench") bench = gu_only = true; // --bench timing the gate/up GEMM alone + else if (a == "--dn-bench") bench = dn_only = true; else if (a == "--v2") v2 = true; else if (a == "--v3") v2 = v3 = true; // v2's gather + the register-operand GEMMs (TN per simdgroup) else if (a == "--v4") { v2 = v3 = v4 = true; MS = std::atoi(argv[++i]); } // mm_gu4 with MS 16-row subtiles else if (a == "--v5") { v2 = v3 = true; v5 = true; } // mm_gu5 (weight blocks staged in threadgroup memory) + else if (a == "--dump") dump = argv[++i]; // HS and y to a file: bitwise comparison of two kernel versions + else if (a == "--ab") ab_dir = argv[++i]; } + const bool isolated_bench = gu_only || dn_only; if (bench && T == 256) T = 2048; if (v4) TM = 16 * MS; // the plan's tile height; mm_gu4's name carries MS instead mtl::Context cx; - const int lib = cx.compile({"quant.metal", "units.metal", "mma.metal", "moe_mm.metal"}); - const int NE = bench ? 512 : 64, LAYERS = bench ? 4 : 1; + const int lib = cx.compile({"quant.metal", "units.metal", "mma.metal", "moe_mm.metal"}, (v3 && !v4 && !v5) ? "#define MM_V3_ONLY 1\n" : ""); + const int NE = bench ? 512 : 64, LAYERS = bench && !isolated_bench ? 4 : 1; const size_t gu_row = ref::row_bytes(gu, N_EMBD), d_row = ref::row_bytes(dn, N_FF); const size_t up_off = (size_t)N_FF * gu_row, down_off = 2 * up_off, blob = down_off + (size_t)N_EMBD * d_row; std::printf("device %s; %s/%s, T %d, %d experts, tile TM %d TN %d TK %d SG %d\n", cx.device_name().c_str(), @@ -67,14 +229,18 @@ int main(int argc, char** argv) { } const int nent = T * TOPK, nblk = (nent + 1023) / 1024; const int maxrows = nent + NE * TM; - mtl::Buffer x = cx.alloc(4 * (size_t)T * N_EMBD), ids = cx.alloc(4 * nent), w = cx.alloc(4 * nent), + mtl::Buffer x = cx.alloc(4 * (size_t)T * N_EMBD), xh = cx.alloc(2 * (size_t)T * N_EMBD), ids = cx.alloc(4 * nent), w = cx.alloc(4 * nent), bcnt = cx.alloc(4 * (size_t)nblk * NE), boff = cx.alloc(4 * (size_t)nblk * NE), eoff = cx.alloc(4 * NE), tiles = cx.alloc(8 * (size_t)(maxrows / TM + 1)), meta = cx.alloc(64), stok = cx.alloc(4 * (size_t)maxrows), sw = cx.alloc(4 * (size_t)maxrows), pos = cx.alloc(4 * nent), H = cx.alloc(4 * (size_t)maxrows * N_FF), YE = cx.alloc(4 * (size_t)maxrows * N_EMBD), y = cx.alloc(4 * (size_t)T * N_EMBD), XS = cx.alloc(2 * (size_t)maxrows * N_EMBD), HS = cx.alloc(2 * (size_t)maxrows * N_FF); ref::Rng r(7); - for (size_t i = 0; i < (size_t)T * N_EMBD; ++i) x.as()[i] = r.sym(); + for (size_t i = 0; i < (size_t)T * N_EMBD; ++i) { + const float v = r.sym(); + x.as()[i] = v; + xh.as()[i] = ref::f32_to_f16(v); + } for (int t = 0; t < T; ++t) { std::set used; float s = 0; @@ -88,7 +254,7 @@ int main(int argc, char** argv) { for (int k = 0; k < TOPK; ++k) w.as()[t * TOPK + k] /= s; } mtl::Pipeline phist = cx.pipeline(lib, "mm_hist"), pscan = cx.pipeline(lib, "mm_scan"), pplace = cx.pipeline(lib, "mm_place"), - pclear = cx.pipeline(lib, "mm_clear"), pcomb = cx.pipeline(lib, "mm_combine"); + pcomb = cx.pipeline(lib, "mm_combine"); if (DTN < 0) { DTN = TN; DTK = TK; DSG = SG; } if (dn == ref::IQ4_NL && DTK < 32) DTK = 32; const std::string tl = std::to_string(TM) + "_" + std::to_string(TN) + "_" + std::to_string(TK) + "_" + std::to_string(SG); @@ -105,31 +271,29 @@ int main(int argc, char** argv) { : cx.pipeline(lib, (v2 ? "mm_gu2_" : "mm_gu_") + std::to_string(gu) + "_" + tl, probe), pdn = v3 ? cx.pipeline(lib, "mm_down3_" + std::to_string(dn) + "_" + dtl, mtl::Constants().b(7, k4)) : cx.pipeline(lib, (v2 ? "mm_down2_" : "mm_down_") + std::to_string(dn) + "_" + (v2 ? dtl : tl)), - pgat = cx.pipeline(lib, "mm_gather"); + pgat = cx.pipeline(lib, "mm_gather_h"); MmPlanArgs pa{(uint32_t)nent, TOPK, (uint32_t)NE, (uint32_t)TM, (uint32_t)nblk}; auto encode = [&](int layer, std::vector& ds, bool plan, bool gemm) { if (plan) { - const uint32_t nclr = maxrows; - mtl::Dispatch c; c.p = &pclear; c.args = {mtl::Arg::B(stok), mtl::Arg::V(nclr)}; c.grid[0] = (maxrows + 255) / 256; c.tg[0] = 256; ds.push_back(c); mtl::Dispatch h; h.p = &phist; h.args = {mtl::Arg::B(ids), mtl::Arg::B(bcnt), mtl::Arg::V(pa)}; h.grid[0] = nblk; h.tg[0] = 1024; ds.push_back(h); meta.as()[2] = N_FF / TN; - mtl::Dispatch sc; sc.p = &pscan; sc.args = {mtl::Arg::B(bcnt), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(tiles), mtl::Arg::B(meta), mtl::Arg::V(pa)}; sc.tg[0] = NE; ds.push_back(sc); + mtl::Dispatch sc; sc.p = &pscan; sc.args = {mtl::Arg::B(bcnt), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(tiles), mtl::Arg::B(meta), mtl::Arg::V(pa), mtl::Arg::B(stok)}; sc.tg[0] = NE; ds.push_back(sc); mtl::Dispatch pl; pl.p = &pplace; pl.args = {mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(stok), mtl::Arg::B(sw), mtl::Arg::B(pos), mtl::Arg::V(pa)}; pl.grid[0] = nblk; pl.tg[0] = 1024; ds.push_back(pl); } if (gemm) { MmArgs ma{(uint64_t)layer * NE * blob, (uint64_t)layer * NE * blob + up_off, (uint64_t)layer * NE * blob + down_off, blob, blob, (uint32_t)gu_row, (uint32_t)d_row, N_EMBD, N_FF}; - if (v2) { + if (v2 || v3) { const uint32_t ne = N_EMBD; - mtl::Dispatch ga; ga.p = &pgat; ga.args = {mtl::Arg::B(x), mtl::Arg::B(stok), mtl::Arg::B(XS), mtl::Arg::V(ne)}; + mtl::Dispatch ga; ga.p = &pgat; ga.args = {mtl::Arg::B(xh), mtl::Arg::B(stok), mtl::Arg::B(XS), mtl::Arg::V(ne)}; ga.grid[0] = meta.as()[1]; ga.tg[0] = 256; ds.push_back(ga); } mtl::Dispatch g; g.p = &pgu; - if (v2) g.args = {mtl::Arg::B(W), mtl::Arg::B(XS), mtl::Arg::B(HS), mtl::Arg::B(tiles), mtl::Arg::V(ma)}; + if (v2 || v3) g.args = {mtl::Arg::B(W), mtl::Arg::B(XS), mtl::Arg::B(HS), mtl::Arg::B(tiles), mtl::Arg::V(ma)}; else g.args = {mtl::Arg::B(W), mtl::Arg::B(x), mtl::Arg::B(H), mtl::Arg::B(stok), mtl::Arg::B(tiles), mtl::Arg::V(ma)}; g.grid[0] = N_FF / (v3 ? TN * SG : TN); g.grid[1] = meta.as()[0]; g.tg[0] = 32 * SG; ds.push_back(g); - mtl::Dispatch d; d.p = &pdn; d.args = {mtl::Arg::B(W), mtl::Arg::B(v2 ? HS : H), mtl::Arg::B(YE), mtl::Arg::B(tiles), mtl::Arg::V(ma)}; - d.grid[0] = N_EMBD / (v3 ? DTN * DSG : v2 ? DTN : TN); d.grid[1] = meta.as()[0]; d.tg[0] = 32 * (v2 ? DSG : SG); ds.push_back(d); + mtl::Dispatch d; d.p = &pdn; d.args = {mtl::Arg::B(W), mtl::Arg::B(v2 || v3 ? HS : H), mtl::Arg::B(YE), mtl::Arg::B(tiles), mtl::Arg::V(ma)}; + d.grid[0] = N_EMBD / (v3 ? DTN * DSG : v2 ? DTN : TN); d.grid[1] = meta.as()[0]; d.tg[0] = 32 * ((v2 || v3) ? DSG : SG); ds.push_back(d); const uint32_t ca[4] = {(uint32_t)T, N_EMBD, TOPK, 0}; mtl::Dispatch cb; cb.p = &pcomb; cb.args = {mtl::Arg::B(YE), mtl::Arg::B(pos), mtl::Arg::B(sw), mtl::Arg::B(y), mtl::Arg::V(ca)}; cb.grid[0] = (T * N_EMBD + 255) / 256; cb.tg[0] = 256; ds.push_back(cb); @@ -142,6 +306,24 @@ int main(int argc, char** argv) { } const uint32_t ntiles = meta.as()[0]; std::printf("plan: %u tiles of %d, %u padded rows for %d entries\n", ntiles, TM, meta.as()[1], nent); + auto compare_ab = [&](const std::vector& ds, mtl::Pipeline& pgu_b, mtl::Pipeline& pdn_b, size_t pr) { + std::vector hs_a(HS.as(), HS.as() + 2 * pr * N_FF); + std::vector ye_a(YE.as(), YE.as() + 4 * pr * N_EMBD); + std::vector y_a(y.as(), y.as() + 4 * (size_t)T * N_EMBD); + std::memset(HS.ptr, 0, hs_a.size()); + std::memset(YE.ptr, 0, ye_a.size()); + std::memset(y.ptr, 0, y_a.size()); + auto ds_b = ds; + ds_b[v2 ? 1 : 0].p = &pgu_b; + ds_b[v2 ? 2 : 1].p = &pdn_b; + cx.run(ds_b); + const bool ok_hs = std::memcmp(HS.ptr, hs_a.data(), hs_a.size()) == 0; + const bool ok_ye = std::memcmp(YE.ptr, ye_a.data(), ye_a.size()) == 0; + const bool ok_y = std::memcmp(y.ptr, y_a.data(), y_a.size()) == 0; + std::printf("ab bitwise (T=%d): HS_memcmp=%s YE_memcmp=%s Y_memcmp=%s\n", + T, ok_hs ? "PASS" : "FAIL", ok_ye ? "PASS" : "FAIL", ok_y ? "PASS" : "FAIL"); + return ok_hs && ok_ye && ok_y; + }; if (!bench) { // the plan: every entry placed once, rows of an expert contiguous and in entry order bool ok_plan = true; @@ -155,6 +337,22 @@ int main(int argc, char** argv) { std::vector ds; encode(0, ds, false, true); cx.run(ds); + const size_t pr = meta.as()[1]; + if (dump) { + FILE* f = std::fopen(dump, "wb"); + if (v2) std::fwrite(HS.as(), 2, pr * N_FF, f); + std::fwrite(YE.as(), 4, pr * N_EMBD, f); + std::fwrite(y.as(), 4, (size_t)T * N_EMBD, f); + std::fclose(f); + return 0; + } + if (ab_dir) { + const std::string kd(ab_dir); + const int lib_b = cx.compile({kd + "/quant.metal", kd + "/units.metal", kd + "/mma.metal", kd + "/moe_mm.metal"}, (v3 && !v4 && !v5) ? "#define MM_V3_ONLY 1\n" : ""); + mtl::Pipeline pgu_b = cx.pipeline(lib_b, "mm_gu3_" + std::to_string(gu) + "_" + tl, mtl::Constants().b(7, k4).i(5, std::getenv("MM_PROBE") ? std::atoi(std::getenv("MM_PROBE")) : 0).b(23, std::getenv("QUANT_CTAB") != nullptr)); + mtl::Pipeline pdn_b = cx.pipeline(lib_b, "mm_down3_" + std::to_string(dn) + "_" + dtl, mtl::Constants().b(7, k4)); + return (ok_plan && compare_ab(ds, pgu_b, pdn_b, pr)) ? 0 : 1; + } // reference: fp16 operands (x, weights, h), fp32+ accumulation; and the exact one double se16 = 0, se = 0, sr = 0; std::vector wg((size_t)N_FF * N_EMBD), wu(wg.size()), wd((size_t)N_EMBD * N_FF); @@ -217,18 +415,85 @@ int main(int argc, char** argv) { } auto timeit = [&](bool plan, bool gemm, int part) { std::vector ds; - for (int l = 0; l < 8; ++l) { + const int repeats = isolated_bench ? 4 : 8; + for (int l = 0; l < repeats; ++l) { std::vector one; encode(l % LAYERS, one, plan, gemm); if (part >= 0 && gemm) { ds.push_back(one[part]); continue; } for (auto& d : one) ds.push_back(d); } - cx.warm(ds, 0.3); + cx.warm(ds, isolated_bench ? 0.1 : 0.3); std::vector t; - for (int i = 0; i < 5; ++i) t.push_back(cx.run(ds) / 8); + for (int i = 0; i < (isolated_bench ? 3 : 5); ++i) t.push_back(cx.run(ds) / repeats); return median(t); }; const int o = v2 ? 1 : 0; + if (isolated_bench) { + std::vector ds; + encode(0, ds, false, true); + cx.run(ds); + } + if (ab_dir && isolated_bench) { + const std::string kd(ab_dir); + const int lib_b = cx.compile({kd + "/quant.metal", kd + "/units.metal", kd + "/mma.metal", kd + "/moe_mm.metal"}, (v3 && !v4 && !v5) ? "#define MM_V3_ONLY 1\n" : ""); + mtl::Pipeline pgu_b = cx.pipeline(lib_b, "mm_gu3_" + std::to_string(gu) + "_" + tl, mtl::Constants().b(7, k4).i(5, std::getenv("MM_PROBE") ? std::atoi(std::getenv("MM_PROBE")) : 0).b(23, std::getenv("QUANT_CTAB") != nullptr)); + mtl::Pipeline pdn_b = cx.pipeline(lib_b, "mm_down3_" + std::to_string(dn) + "_" + dtl, mtl::Constants().b(7, k4)); + std::vector parity; + encode(0, parity, false, true); + const size_t pr = meta.as()[1]; + std::memset(HS.ptr, 0, 2 * pr * N_FF); + std::memset(YE.ptr, 0, 4 * pr * N_EMBD); + cx.run(parity); + if (!compare_ab(parity, pgu_b, pdn_b, pr)) return 1; + const int part = dn_only ? (o + 1) : o; + std::vector one; + encode(0, one, false, true); + std::vector ds_a{one[part]}; + mtl::Dispatch db = one[part]; + db.p = dn_only ? &pdn_b : &pgu_b; + std::vector ds_b{db}; + cx.warm(ds_a, 0.05); + cx.warm(ds_b, 0.05); + std::vector ta, tb; + double best_pair = 1e30; + for (int i = 0; i < 19; ++i) { + double va, vb; + if (i & 1) { + vb = cx.run(ds_b, 2) / 2; + va = cx.run(ds_a, 2) / 2; + } else { + va = cx.run(ds_a, 2) / 2; + vb = cx.run(ds_b, 2) / 2; + } + best_pair = std::min(best_pair, std::min(va, vb)); + ta.push_back(va); + tb.push_back(vb); + } + std::vector fa, fb; + for (size_t i = 0; i < ta.size(); ++i) { + if (std::max(ta[i], tb[i]) <= 1.10 * best_pair) { + fa.push_back(ta[i]); + fb.push_back(tb[i]); + } + } + if (fa.size() >= 3) { ta = fa; tb = fb; } + const double ma = median(ta) * 1e6, mb = median(tb) * 1e6; + const double mina = *std::min_element(ta.begin(), ta.end()) * 1e6; + const double minb = *std::min_element(tb.begin(), tb.end()) * 1e6; + std::printf("%s A/B (T=%d): A med=%.1f (min=%.1f) | B med=%.1f (min=%.1f) | diff med=%+.2f%% min=%+.2f%%\n", + dn_only ? "down" : "gate/up", T, ma, mina, mb, minb, (mb - ma) / ma * 100.0, (minb - mina) / mina * 100.0); + return 0; + } + if (dn_only) { + const double td = timeit(false, true, o + 1); + std::printf("down %.1f us (%.2f TFLOP/s useful)\n", td * 1e6, 2.0 * nent * N_FF * N_EMBD / td / 1e12); + return 0; + } + if (gu_only) { + const double tg = timeit(false, true, o); + std::printf("gate/up %.1f us (%.2f TFLOP/s useful)\n", tg * 1e6, 2.0 * nent * 2.0 * N_FF * N_EMBD / tg / 1e12); + return 0; + } const double tp = timeit(true, false, -1), tg = timeit(false, true, o), td = timeit(false, true, o + 1), tc = timeit(false, true, o + 2), ta = timeit(false, true, -1); if (v2) std::printf("gather %.1f us\n", timeit(false, true, 0) * 1e6); diff --git a/metal/tests/ple_moe_parity.cpp b/metal/tests/ple_moe_parity.cpp new file mode 100644 index 00000000..fb40e33e --- /dev/null +++ b/metal/tests/ple_moe_parity.cpp @@ -0,0 +1,340 @@ +// metal/tests/ple_moe_parity.cpp - the prompt chunk's moe_out folds and the fused PLE norms against the dispatch +// chains they replace, BITWISE: +// gr_norm_h_moe (prompt.metal) == moe_out + gr_norm_h (R', xnh) +// gr_apply_moe (prompt.metal) == moe_out + gr_norm (R') +// ple_moe_gate_norm + ple_conv_res == moe_out + gr_norm + ple_gnorm (key) + ple_gnorm (query) + ple_gate_norm +// (msl/ple.metal) + ple_conv_add (R', normalized) +// ple_moe_gate_norm_c + ple_conv_res == mm_combine + the ple_moe_gate_norm chain above (R', normalized) +// ple_query_gate_norm + ple_conv_res == ple_gnorm (query) + ple_gate_norm + ple_conv_add (R', normalized; key +// from ple_gnorm) +// ple_comb_query_gate_norm == gr_norm_comb + ple_query_gate_norm (R', normalized) +// ple_conv_res + ple_hist_adv == ple_conv_res + the host's history advance (R', the history), T = 1 .. 2048 +// `ple_moe_parity [T] bench` (or `--bench`) adds the PLE stage's GPU times at T (default 2048). +// `ple_moe_parity --ab baseline.metal [T]` runs short paired A/B bitwise and timing checks, including history, at any T. +#include "mtl_rt.h" +#include "ref_quant.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { +struct GrArgs { uint32_t n_embd, hc, hc_lr, n_tok, apply; float eps; }; +struct PleArgs { uint32_t n_embd, hc; float eps; }; +constexpr uint32_t P_TOKENT = 532; +bool same(const mtl::Buffer& a, const mtl::Buffer& b, size_t bytes) { return std::memcmp(a.ptr, b.ptr, bytes) == 0; } +double median(std::vector v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; } +} // namespace + +int main(int argc, char** argv) { + const bool ab = argc > 1 && std::strcmp(argv[1], "--ab") == 0; + if (ab && argc < 3) { + std::fprintf(stderr, "usage: ple_moe_parity --ab baseline.metal [T]\n"); + return 2; + } + const bool flag_bench = argc > 1 && std::strcmp(argv[1], "--bench") == 0; + const uint32_t TB = ab ? (argc > 3 ? (uint32_t) std::atoi(argv[3]) : 2048u) : + (argc > 1 && !flag_bench ? (uint32_t) std::atoi(argv[1]) : 2048u); + const bool bench = ab || flag_bench || (argc > 2 && (std::strcmp(argv[2], "bench") == 0 || std::strcmp(argv[2], "--bench") == 0)); + mtl::Context cx; + const int lple = cx.compile({"../src/engine/msl/ple.metal"}, "#define NO_GGML_COMMON 1\n"); + std::printf("device %s\n", cx.device_name().c_str()); + const uint32_t N = 2560, HC = 4, LR = 320, TOPK = 10, NEXP = 512, TMR = 16; + const size_t RN = (size_t) HC * N; + mtl::Pipeline pg = cx.pipeline(lple, "ple_gnorm"), + pgn = cx.pipeline(lple, "ple_gate_norm"), pca = cx.pipeline(lple, "ple_conv_add"), + pcr = cx.pipeline(lple, "ple_conv_res"), pf = cx.pipeline(lple, "ple_moe_gate_norm"), + pfc = cx.pipeline(lple, "ple_moe_gate_norm_c"), pq = cx.pipeline(lple, "ple_query_gate_norm"), + pcq = cx.pipeline(lple, "ple_comb_query_gate_norm"), pha = cx.pipeline(lple, "ple_hist_adv"); + auto D = [](const mtl::Pipeline& p, std::vector a, uint32_t gx, uint32_t gy) { + mtl::Dispatch d; + d.p = &p; + d.args = std::move(a); + d.grid[0] = gx; + d.grid[1] = gy; + d.tg[0] = 256; + return d; + }; + int fails = 0; + if (ab) { + const int lref = cx.compile({argv[2]}, "#define NO_GGML_COMMON 1\n"); + mtl::Pipeline pg_r = cx.pipeline(lref, "ple_gnorm"), pgn_r = cx.pipeline(lref, "ple_gate_norm"), + pcr_r = cx.pipeline(lref, "ple_conv_res"), pf_r = cx.pipeline(lref, "ple_moe_gate_norm"), + pfc_r = cx.pipeline(lref, "ple_moe_gate_norm_c"), pq_r = cx.pipeline(lref, "ple_query_gate_norm"), + pcq_r = cx.pipeline(lref, "ple_comb_query_gate_norm"), pha_r = cx.pipeline(lref, "ple_hist_adv"); + const std::vector ts = argc > 3 ? std::vector{TB} : + std::vector{1u, 3u, 7u, 9u, 10u, 64u, 512u, 1024u, 2048u}; + for (uint32_t T : ts) { + const size_t mr = (size_t) T * TOPK + NEXP * TMR; + mtl::Buffer R0 = cx.alloc(4 * T * RN), Ra = cx.alloc(4 * T * RN), Rb = cx.alloc(4 * T * RN), yr = cx.alloc(4 * T * N), + ys = cx.alloc(4 * T * N), sgv = cx.alloc(4 * T), inj = cx.alloc(4 * T * HC), + wk = cx.alloc(4 * RN), wq = cx.alloc(4 * RN), wc = cx.alloc(4 * RN), + pkey = cx.alloc(4 * T * RN), pval = cx.alloc(4 * T * N), keyn = cx.alloc(4 * T * RN), + ga = cx.alloc(4 * T * RN), gb = cx.alloc(4 * T * RN), + na = cx.alloc(4 * T * RN), nb = cx.alloc(4 * T * RN), hist = cx.alloc(4 * 9 * RN), kw = cx.alloc(2 * 4 * RN), + YE = cx.alloc(4 * mr * N), pos = cx.alloc(4 * T * TOPK), sw = cx.alloc(4 * mr), + plan = cx.alloc(4 * (P_TOKENT + T * TOPK)), + ha = cx.alloc(4 * 9 * RN), hb = cx.alloc(4 * 9 * RN); + ref::Rng r(11 + T); + auto fill = [&](const mtl::Buffer& b, float s, float o) { + for (size_t i = 0; i < b.bytes / 4; ++i) b.as()[i] = o + s * r.sym(); + }; + fill(R0, 1.0f, 0.0f); fill(yr, 0.5f, 0.0f); fill(ys, 0.5f, 0.0f); fill(sgv, 3.0f, 0.0f); fill(inj, 4.0f, 0.0f); + fill(wk, 0.2f, 1.0f); fill(wq, 0.2f, 1.0f); fill(wc, 0.2f, 1.0f); + fill(pkey, 1.0f, 0.0f); fill(pval, 1.0f, 0.0f); fill(keyn, 1.0f, 0.0f); fill(hist, 1.0f, 0.0f); + fill(YE, 0.5f, 0.0f); fill(sw, 0.1f, 0.1f); + for (size_t i = 0; i < kw.bytes / 2; ++i) kw.as()[i] = ref::f32_to_f16(0.5f * r.sym()); + std::vector perm(mr); + std::iota(perm.begin(), perm.end(), 0); + std::shuffle(perm.begin(), perm.end(), std::mt19937(T)); + for (size_t i = 0; i < (size_t) T * TOPK; ++i) { + pos.as()[i] = (i % 37 == 5) ? -1 : perm[i]; + plan.as()[P_TOKENT + i] = perm[i]; + } + const PleArgs pl{N, HC, 1e-6f}; + auto reset = [&] { + std::memcpy(Ra.ptr, R0.ptr, 4 * T * RN); + std::memcpy(Rb.ptr, R0.ptr, 4 * T * RN); + std::memset(na.ptr, 0, 4 * T * RN); + std::memset(nb.ptr, 0, 4 * T * RN); + }; + const mtl::Dispatch g_r = D(pg_r, {mtl::Arg::B(pkey), mtl::Arg::B(wk), mtl::Arg::B(na), mtl::Arg::V(pl)}, HC, T), + g_c = D(pg, {mtl::Arg::B(pkey), mtl::Arg::B(wk), mtl::Arg::B(nb), mtl::Arg::V(pl)}, HC, T), + gn_r = D(pgn_r, {mtl::Arg::B(keyn), mtl::Arg::B(R0), mtl::Arg::B(pval), mtl::Arg::B(wc), + mtl::Arg::B(ga), mtl::Arg::B(na), mtl::Arg::V(pl)}, HC, T), + gn_c = D(pgn, {mtl::Arg::B(keyn), mtl::Arg::B(R0), mtl::Arg::B(pval), mtl::Arg::B(wc), + mtl::Arg::B(gb), mtl::Arg::B(nb), mtl::Arg::V(pl)}, HC, T), + res_r = D(pcr_r, {mtl::Arg::B(hist), mtl::Arg::B(na), mtl::Arg::B(kw), mtl::Arg::B(Ra), mtl::Arg::V(pl)}, HC * N / 256, T), + res_c = D(pcr, {mtl::Arg::B(hist), mtl::Arg::B(nb), mtl::Arg::B(kw), mtl::Arg::B(Rb), mtl::Arg::V(pl)}, HC * N / 256, T), + f_r = D(pf_r, {mtl::Arg::B(Ra), mtl::Arg::B(yr), mtl::Arg::B(ys), mtl::Arg::B(sgv), mtl::Arg::B(inj), + mtl::Arg::B(pkey), mtl::Arg::B(wk), mtl::Arg::B(wq), mtl::Arg::B(pval), mtl::Arg::B(wc), mtl::Arg::B(na), mtl::Arg::V(pl)}, HC, T), + f_c = D(pf, {mtl::Arg::B(Rb), mtl::Arg::B(yr), mtl::Arg::B(ys), mtl::Arg::B(sgv), mtl::Arg::B(inj), + mtl::Arg::B(pkey), mtl::Arg::B(wk), mtl::Arg::B(wq), mtl::Arg::B(pval), mtl::Arg::B(wc), mtl::Arg::B(nb), mtl::Arg::V(pl)}, HC, T), + fc_r = D(pfc_r, {mtl::Arg::B(Ra), mtl::Arg::B(YE), mtl::Arg::B(pos), mtl::Arg::B(sw), mtl::Arg::B(ys), + mtl::Arg::B(sgv), mtl::Arg::B(inj), mtl::Arg::B(pkey), mtl::Arg::B(wk), mtl::Arg::B(wq), mtl::Arg::B(pval), mtl::Arg::B(wc), mtl::Arg::B(na), mtl::Arg::V(pl), mtl::Arg::V(TOPK)}, 1, T), + fc_c = D(pfc, {mtl::Arg::B(Rb), mtl::Arg::B(YE), mtl::Arg::B(pos), mtl::Arg::B(sw), mtl::Arg::B(ys), + mtl::Arg::B(sgv), mtl::Arg::B(inj), mtl::Arg::B(pkey), mtl::Arg::B(wk), mtl::Arg::B(wq), mtl::Arg::B(pval), mtl::Arg::B(wc), mtl::Arg::B(nb), mtl::Arg::V(pl), mtl::Arg::V(TOPK)}, 1, T), + q_r = D(pq_r, {mtl::Arg::B(keyn), mtl::Arg::B(Ra), mtl::Arg::B(wq), mtl::Arg::B(pval), mtl::Arg::B(wc), mtl::Arg::B(na), mtl::Arg::V(pl)}, HC, T), + q_c = D(pq, {mtl::Arg::B(keyn), mtl::Arg::B(Rb), mtl::Arg::B(wq), mtl::Arg::B(pval), mtl::Arg::B(wc), mtl::Arg::B(nb), mtl::Arg::V(pl)}, HC, T), + cq_r = D(pcq_r, {mtl::Arg::B(keyn), mtl::Arg::B(Ra), mtl::Arg::B(wq), mtl::Arg::B(pval), mtl::Arg::B(wc), mtl::Arg::B(na), mtl::Arg::V(pl), + mtl::Arg::B(YE), mtl::Arg::B(inj), mtl::Arg::B(plan), mtl::Arg::B(ys), mtl::Arg::B(sgv), mtl::Arg::V(TOPK)}, HC, T), + cq_c = D(pcq, {mtl::Arg::B(keyn), mtl::Arg::B(Rb), mtl::Arg::B(wq), mtl::Arg::B(pval), mtl::Arg::B(wc), mtl::Arg::B(nb), mtl::Arg::V(pl), + mtl::Arg::B(YE), mtl::Arg::B(inj), mtl::Arg::B(plan), mtl::Arg::B(ys), mtl::Arg::B(sgv), mtl::Arg::V(TOPK)}, HC, T); + reset(); cx.run({g_r}); cx.run({g_c}); + const bool ok_g = same(na, nb, 4 * T * RN); + reset(); std::memset(ga.ptr, 0, 4 * T * RN); std::memset(gb.ptr, 0, 4 * T * RN); + cx.run({gn_r}); cx.run({gn_c}); + const bool ok_gn = same(ga, gb, 4 * T * RN) && same(na, nb, 4 * T * RN); + reset(); cx.run({f_r, res_r}); cx.run({f_c, res_c}); + const bool ok_f = same(Ra, Rb, 4 * T * RN) && same(na, nb, 4 * T * RN); + reset(); cx.run({fc_r, res_r}); cx.run({fc_c, res_c}); + const bool ok_fc = same(Ra, Rb, 4 * T * RN) && same(na, nb, 4 * T * RN); + reset(); cx.run({q_r, res_r}); cx.run({q_c, res_c}); + const bool ok_q = same(Ra, Rb, 4 * T * RN) && same(na, nb, 4 * T * RN); + reset(); cx.run({cq_r, res_r}); cx.run({cq_c, res_c}); + const bool ok_cq = same(Ra, Rb, 4 * T * RN) && same(na, nb, 4 * T * RN); + const mtl::Dispatch h_r = D(pha_r, {mtl::Arg::B(ha), mtl::Arg::B(na), mtl::Arg::V(pl), mtl::Arg::V(T)}, HC * N / 256, 1), + h_c = D(pha, {mtl::Arg::B(hb), mtl::Arg::B(nb), mtl::Arg::V(pl), mtl::Arg::V(T)}, HC * N / 256, 1); + std::memcpy(ha.ptr, hist.ptr, hist.bytes); + std::memcpy(hb.ptr, hist.ptr, hist.bytes); + cx.run({h_r}); cx.run({h_c}); + const bool ok_h = same(ha, hb, hist.bytes); + fails += !ok_g + !ok_gn + !ok_f + !ok_fc + !ok_q + !ok_cq + !ok_h; + std::printf("A/B T %4u: ple_gnorm %s, ple_gate_norm %s, ple_hist_adv %s\n", + T, ok_g ? "PASS" : "FAIL", ok_gn ? "PASS" : "FAIL", ok_h ? "PASS" : "FAIL"); + std::printf("A/B T %4u: ple_moe_gate_norm %s, ple_moe_gate_norm_c %s, ple_query_gate_norm %s, ple_comb_query_gate_norm %s\n", + T, ok_f ? "PASS" : "FAIL", ok_fc ? "PASS" : "FAIL", ok_q ? "PASS" : "FAIL", ok_cq ? "PASS" : "FAIL"); + const bool single_t = argc > 3; + const int reps = single_t ? (T <= 10 ? 64 : (T <= 64 ? 16 : 4)) : 2; + const int pairs = single_t ? 7 : 5; + auto ab_time = [&](const char* name, const std::vector& ref_ds, const std::vector& cand_ds, int r = 0) { + if (!r) r = reps; + cx.warm(ref_ds, 0.003); cx.warm(cand_ds, 0.003); + std::vector bt, ct, rt; + for (int p = 0; p < pairs; ++p) { + double b, c; + if (p & 1) { c = cx.run(cand_ds, r) / r * 1e3; b = cx.run(ref_ds, r) / r * 1e3; } + else { b = cx.run(ref_ds, r) / r * 1e3; c = cx.run(cand_ds, r) / r * 1e3; } + bt.push_back(b); ct.push_back(c); rt.push_back(c / b); + } + std::printf(" T %4u %-34s base %.4f -> cand %.4f ms (ratio %.4f, %+.1f%%)\n", + T, name, median(bt), median(ct), median(rt), 100.0 * (median(rt) - 1.0)); + }; + ab_time("ple_gnorm", {g_r}, {g_c}); + ab_time("ple_gate_norm", {gn_r}, {gn_c}); + ab_time("ple_query_gate_norm", {q_r}, {q_c}); + ab_time("ple_comb_query_gate_norm", {cq_r}, {cq_c}); + ab_time("ple_moe_gate_norm", {f_r}, {f_c}); + ab_time("ple_moe_gate_norm_c", {fc_r}, {fc_c}); + ab_time("ple_hist_adv", {h_r}, {h_c}, 128); + ab_time("ple_conv_res", {res_r}, {res_c}); + ab_time("ple_comb_query_gate_norm + conv_res", {cq_r, res_r}, {cq_c, res_c}); + ab_time("ple_moe_gate_norm_c + conv_res", {fc_r, res_r}, {fc_c, res_c}); + ab_time("ple_query_gate_norm + conv_res", {q_r, res_r}, {q_c, res_c}); + } + std::printf("%s\n", fails ? "FAIL" : "ALL PASS"); + return fails ? 1 : 0; + } + const std::string prelude = "#define NO_GGML_COMMON 1\nusing namespace metal;\n"; + const std::string gr_src = mtl::read_file(cx.kernel_dir + "/gr.metal"); + const std::string pr_src = mtl::read_file(cx.kernel_dir + "/prompt.metal"); + const std::string mm_src = mtl::read_file(cx.kernel_dir + "/moe_mm.metal"); + const auto gn_end = gr_src.find("// rows: 0 .. hc_lr - 1"); + const auto mo_beg = pr_src.find("// grid (T * n_embd / 256); sg: T values"); + const auto pr_end = pr_src.find("// gemm_hc_up:"); + const auto mc_beg = mm_src.find("kernel void mm_combine("); + const auto mc_end = mm_src.find("\n}", mc_beg == std::string::npos ? 0 : mc_beg); + const int lib = cx.compile({}, prelude + gr_src.substr(0, gn_end) + pr_src.substr(mo_beg, pr_end - mo_beg)); + const int lmm = cx.compile({}, prelude + mm_src.substr(mc_beg, mc_end + 2 - mc_beg)); + mtl::Pipeline pmo = cx.pipeline(lib, "moe_out"), pnh = cx.pipeline(lib, "gr_norm_h"), pn = cx.pipeline(lib, "gr_norm"), + pnhm = cx.pipeline(lib, "gr_norm_h_moe"), pam = cx.pipeline(lib, "gr_apply_moe"), + pcomb = cx.pipeline(lmm, "mm_combine"); + for (uint32_t T : {1u, 3u, 7u, 9u, 10u, 64u, TB}) { + if (T == TB && !bench) continue; + const size_t mr = (size_t) T * TOPK + NEXP * TMR; + mtl::Buffer R0 = cx.alloc(4 * T * RN), Ra = cx.alloc(4 * T * RN), Rb = cx.alloc(4 * T * RN), yr = cx.alloc(4 * T * N), + ys = cx.alloc(4 * T * N), sgv = cx.alloc(4 * T), inj = cx.alloc(4 * T * HC), bo = cx.alloc(4 * T * N), + wn = cx.alloc(4 * RN), wk = cx.alloc(4 * RN), wq = cx.alloc(4 * RN), wc = cx.alloc(4 * RN), + pkey = cx.alloc(4 * T * RN), pval = cx.alloc(4 * T * N), xa = cx.alloc(4 * T * RN), xb = cx.alloc(4 * T * RN), + keyn = cx.alloc(4 * T * RN), qn = cx.alloc(4 * T * RN), ga = cx.alloc(4 * T * RN), + na = cx.alloc(4 * T * RN), nb = cx.alloc(4 * T * RN), hist = cx.alloc(4 * 9 * RN), kw = cx.alloc(2 * 4 * RN), + YE = cx.alloc(4 * mr * N), pos = cx.alloc(4 * T * TOPK), sw = cx.alloc(4 * mr); + ref::Rng r(11 + T); + auto fill = [&](const mtl::Buffer& b, float s, float o) { + for (size_t i = 0; i < b.bytes / 4; ++i) b.as()[i] = o + s * r.sym(); + }; + fill(R0, 1.0f, 0.0f); fill(yr, 0.5f, 0.0f); fill(ys, 0.5f, 0.0f); fill(sgv, 3.0f, 0.0f); fill(inj, 4.0f, 0.0f); + fill(wn, 0.2f, 1.0f); fill(wk, 0.2f, 1.0f); fill(wq, 0.2f, 1.0f); fill(wc, 0.2f, 1.0f); + fill(pkey, 1.0f, 0.0f); fill(pval, 1.0f, 0.0f); fill(hist, 1.0f, 0.0f); fill(YE, 0.5f, 0.0f); fill(sw, 0.1f, 0.1f); + for (size_t i = 0; i < kw.bytes / 2; ++i) kw.as()[i] = ref::f32_to_f16(0.5f * r.sym()); + // mm_place's slots: distinct rows of YE (expert-sorted, so scattered), a few dropped (-1) + std::vector perm(mr); + std::iota(perm.begin(), perm.end(), 0); + std::shuffle(perm.begin(), perm.end(), std::mt19937(T)); + for (size_t i = 0; i < (size_t) T * TOPK; ++i) pos.as()[i] = (i % 37 == 5) ? -1 : perm[i]; + const GrArgs a{N, HC, LR, T, 1u, 1e-6f}; + const PleArgs pl{N, HC, 1e-6f}; + const uint32_t oa[2] = {T, N}, ca[4] = {T, N, TOPK, 0}; + const mtl::Dispatch mo = D(pmo, {mtl::Arg::B(yr), mtl::Arg::B(ys), mtl::Arg::B(sgv), mtl::Arg::B(bo), mtl::Arg::V(oa)}, + (T * N + 255) / 256, 1), + comb = D(pcomb, {mtl::Arg::B(YE), mtl::Arg::B(pos), mtl::Arg::B(sw), mtl::Arg::B(yr), mtl::Arg::V(ca)}, + (T * N + 255) / 256, 1); + auto reset = [&] { + std::memcpy(Ra.ptr, R0.ptr, 4 * T * RN); + std::memcpy(Rb.ptr, R0.ptr, 4 * T * RN); + std::memset(bo.ptr, 0, 4 * T * N); + }; + auto check = [&](const char* what, bool ok) { + fails += !ok; + std::printf("%-34s T %4u: bitwise %s\n", what, T, ok ? "PASS" : "FAIL"); + }; + // the unfused PLE stage on R (layer 0's MoE write from yr applied first if moe) + auto chain = [&](const mtl::Buffer& R, bool moe) { + std::vector ds; + if (moe) { + ds.push_back(mo); + ds.push_back(D(pn, {mtl::Arg::B(R), mtl::Arg::B(bo), mtl::Arg::B(inj), mtl::Arg::B(wn), mtl::Arg::B(xa), + mtl::Arg::V(a)}, HC, T)); + } + ds.push_back(D(pg, {mtl::Arg::B(pkey), mtl::Arg::B(wk), mtl::Arg::B(keyn), mtl::Arg::V(pl)}, HC, T)); + ds.push_back(D(pg, {mtl::Arg::B(R), mtl::Arg::B(wq), mtl::Arg::B(qn), mtl::Arg::V(pl)}, HC, T)); + ds.push_back(D(pgn, {mtl::Arg::B(keyn), mtl::Arg::B(qn), mtl::Arg::B(pval), mtl::Arg::B(wc), mtl::Arg::B(ga), + mtl::Arg::B(na), mtl::Arg::V(pl)}, HC, T)); + ds.push_back(D(pca, {mtl::Arg::B(hist), mtl::Arg::B(na), mtl::Arg::B(kw), mtl::Arg::B(ga), mtl::Arg::B(R), + mtl::Arg::V(pl)}, HC * N / 256, T)); + return ds; + }; + const mtl::Dispatch res = D(pcr, {mtl::Arg::B(hist), mtl::Arg::B(nb), mtl::Arg::B(kw), mtl::Arg::B(Rb), mtl::Arg::V(pl)}, + HC * N / 256, T), + fused = D(pf, {mtl::Arg::B(Rb), mtl::Arg::B(yr), mtl::Arg::B(ys), mtl::Arg::B(sgv), mtl::Arg::B(inj), + mtl::Arg::B(pkey), mtl::Arg::B(wk), mtl::Arg::B(wq), mtl::Arg::B(pval), + mtl::Arg::B(wc), mtl::Arg::B(nb), mtl::Arg::V(pl)}, HC, T), + fused_c = D(pfc, {mtl::Arg::B(Rb), mtl::Arg::B(YE), mtl::Arg::B(pos), mtl::Arg::B(sw), mtl::Arg::B(ys), + mtl::Arg::B(sgv), mtl::Arg::B(inj), mtl::Arg::B(pkey), mtl::Arg::B(wk), + mtl::Arg::B(wq), mtl::Arg::B(pval), mtl::Arg::B(wc), mtl::Arg::B(nb), + mtl::Arg::V(pl), mtl::Arg::V(TOPK)}, 1, T), + query = D(pq, {mtl::Arg::B(keyn), mtl::Arg::B(Rb), mtl::Arg::B(wq), mtl::Arg::B(pval), mtl::Arg::B(wc), + mtl::Arg::B(nb), mtl::Arg::V(pl)}, HC, T); + const bool small = T != TB; + if (small && (T == 1u || T == 7u || T == 64u)) { + // ---- gr_norm_h_moe + reset(); + cx.run({mo, D(pnh, {mtl::Arg::B(Ra), mtl::Arg::B(bo), mtl::Arg::B(inj), mtl::Arg::B(wn), mtl::Arg::B(xa), mtl::Arg::V(a)}, HC, T)}); + cx.run({D(pnhm, {mtl::Arg::B(Rb), mtl::Arg::B(yr), mtl::Arg::B(ys), mtl::Arg::B(sgv), mtl::Arg::B(inj), mtl::Arg::B(wn), + mtl::Arg::B(xb), mtl::Arg::V(a)}, 1, T)}); + check("gr_norm_h_moe", same(Ra, Rb, 4 * T * RN) && same(xa, xb, 2 * T * RN)); + // ---- gr_apply_moe + reset(); + cx.run({mo, D(pn, {mtl::Arg::B(Ra), mtl::Arg::B(bo), mtl::Arg::B(inj), mtl::Arg::B(wn), mtl::Arg::B(xa), mtl::Arg::V(a)}, HC, T)}); + cx.run({D(pam, {mtl::Arg::B(Rb), mtl::Arg::B(yr), mtl::Arg::B(ys), mtl::Arg::B(sgv), mtl::Arg::B(inj), mtl::Arg::V(a)}, + HC * N / 256, T)}); + check("gr_apply_moe", same(Ra, Rb, 4 * T * RN)); + } + // ---- ple_moe_gate_norm + ple_conv_res + reset(); + cx.run(chain(Ra, true)); + cx.run({fused, res}); + check("ple_moe_gate_norm + ple_conv_res", same(Ra, Rb, 4 * T * RN) && same(na, nb, 4 * T * RN)); + // ---- ple_moe_gate_norm_c + ple_conv_res (yr from mm_combine in the chain) + reset(); + std::vector cc = chain(Ra, true); + cc.insert(cc.begin(), comb); + cx.run(cc); + std::memset(nb.ptr, 0, 4 * T * RN); + cx.run({fused_c, res}); + check("ple_moe_gate_norm_c + ple_conv_res", same(Ra, Rb, 4 * T * RN) && same(na, nb, 4 * T * RN)); + // ---- ple_query_gate_norm + ple_conv_res (the window: keyn from the chain's ple_gnorm) + reset(); + cx.run(chain(Ra, false)); + std::memset(nb.ptr, 0, 4 * T * RN); + cx.run({query, res}); + check("ple_query_gate_norm + ple_conv_res", same(Ra, Rb, 4 * T * RN) && same(na, nb, 4 * T * RN)); + if (small) continue; + // ---- GPU times: one PLE stage a chunk, ms, best of 5 x 8 reps + auto time = [&](const std::vector& ds) { + cx.warm(ds, 0.1); + double best = 1e9; + for (int i = 0; i < 5; ++i) best = std::min(best, cx.run(ds, 8) / 8); + return best * 1e3; + }; + const double t_f = time({comb, fused, res}), t_fc = time({fused_c, res}), t_q = time({query, res}); + std::printf("T %u, ms per PLE stage (best of 5):\n", T); + std::printf(" mm_combine + ple_moe_gate_norm + ple_conv_res %.3f -> ple_moe_gate_norm_c + ple_conv_res %.3f (%+.1f%%)\n", + t_f, t_fc, 100.0 * (t_fc / t_f - 1)); + std::printf(" ple_query_gate_norm + ple_conv_res %.3f\n", t_q); + } + // ---- ple_hist_adv after ple_conv_res: the conv reads the old history, then the history moves past the T rows (T < 9 + // keeps 9 - T old rows); against window.cpp's former host loop + for (uint32_t T : {1u, 3u, 7u, 8u, 9u, 10u, 64u, 2048u}) { + mtl::Buffer nw = cx.alloc(4 * T * RN), hist = cx.alloc(4 * 9 * RN), kw = cx.alloc(2 * 4 * RN), Ra = cx.alloc(4 * T * RN), + Rb = cx.alloc(4 * T * RN); + ref::Rng r(97 + T); + for (size_t i = 0; i < nw.bytes / 4; ++i) nw.as()[i] = r.sym(); + for (size_t i = 0; i < hist.bytes / 4; ++i) hist.as()[i] = r.sym(); + for (size_t i = 0; i < Ra.bytes / 4; ++i) Ra.as()[i] = Rb.as()[i] = r.sym(); + for (size_t i = 0; i < kw.bytes / 2; ++i) kw.as()[i] = ref::f32_to_f16(0.5f * r.sym()); + const PleArgs pl{N, HC, 1e-6f}; + cx.run({D(pcr, {mtl::Arg::B(hist), mtl::Arg::B(nw), mtl::Arg::B(kw), mtl::Arg::B(Ra), mtl::Arg::V(pl)}, HC * N / 256, T)}); + std::vector want(hist.as(), hist.as() + 9 * RN); + for (size_t c = 0; c < RN; ++c) { + float col[9]; + for (int k = 0; k < 9; ++k) { + const int src = (int) T - 9 + k; + col[k] = src >= 0 ? nw.as()[(size_t) src * RN + c] : want[c * 9 + (size_t) k + T]; + } + std::memcpy(want.data() + c * 9, col, sizeof col); + } + cx.run({D(pcr, {mtl::Arg::B(hist), mtl::Arg::B(nw), mtl::Arg::B(kw), mtl::Arg::B(Rb), mtl::Arg::V(pl)}, HC * N / 256, T), + D(pha, {mtl::Arg::B(hist), mtl::Arg::B(nw), mtl::Arg::V(pl), mtl::Arg::V(T)}, HC * N / 256, 1)}); + const bool ok = same(Ra, Rb, 4 * T * RN) && std::memcmp(hist.ptr, want.data(), 4 * 9 * RN) == 0; + fails += !ok; + std::printf("%-34s T %4u: bitwise %s\n", "ple_conv_res + ple_hist_adv", T, ok ? "PASS" : "FAIL"); + } + std::printf("%s\n", fails ? "FAIL" : "ALL PASS"); + return fails ? 1 : 0; +} diff --git a/metal/tests/prompt_chain.cpp b/metal/tests/prompt_chain.cpp index 876afeca..88759e98 100644 --- a/metal/tests/prompt_chain.cpp +++ b/metal/tests/prompt_chain.cpp @@ -413,11 +413,9 @@ int main(int argc, char** argv) { 256, "router"); const int nent = T * TOPK, nblk = (nent + 1023) / 1024; MmPlanArgs pa{(uint32_t)nent, TOPK, NEXP, (uint32_t)TM, (uint32_t)nblk}; - const uint32_t nclr = maxrows; - D(P(lib, "mm_clear"), {mtl::Arg::B(stok), mtl::Arg::V(nclr)}, (maxrows + 255) / 256, 1, 256, "experts"); D(P(lib, "mm_hist"), {mtl::Arg::B(ids), mtl::Arg::B(bcnt), mtl::Arg::V(pa)}, nblk, 1, 1024, "experts"); D(P(lib, "mm_scan"), {mtl::Arg::B(bcnt), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(tiles), mtl::Arg::B(meta), - mtl::Arg::V(pa)}, 1, 1, NEXP, "experts"); + mtl::Arg::V(pa), mtl::Arg::B(stok)}, 1, 1, NEXP, "experts"); D(P(lib, "mm_place"), {mtl::Arg::B(ids), mtl::Arg::B(wts), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(stok), mtl::Arg::B(sw), mtl::Arg::B(pos), mtl::Arg::V(pa)}, nblk, 1, 1024, "experts"); MmArgs ma{eg.offset, eu.offset, ed.offset, 0, 0, (uint32_t)ref::row_bytes(eg.type, N), @@ -426,8 +424,8 @@ int main(int argc, char** argv) { ma.d_stride = (uint64_t)ma.d_row * N; const uint32_t ne = N; mtl::Dispatch g; - g.p = P(lib, "mm_gather"); - g.args = {mtl::Arg::B(mixed), mtl::Arg::B(stok), mtl::Arg::B(XS), mtl::Arg::V(ne)}; + g.p = P(lib, "mm_gather_h"); + g.args = {mtl::Arg::B(mixh), mtl::Arg::B(stok), mtl::Arg::B(XS), mtl::Arg::V(ne)}; g.tg[0] = 256; g.indirect = &meta; g.ind_off = 8 * 4; diff --git a/metal/tests/qsa_parity.cpp b/metal/tests/qsa_parity.cpp index 8162ec2e..68e18eb0 100644 --- a/metal/tests/qsa_parity.cpp +++ b/metal/tests/qsa_parity.cpp @@ -4,6 +4,15 @@ // 2051), then a window of 3 tokens: prep, score, top-k, attention, merge. Each stage is // checked against the CPU fed with the GPU's inputs to that stage; the ids exactly. // qsa_parity --bench [P] GPU time of the window's kernels at context P, 1 and 3 tokens +// qsa_parity --bitwise REF.metal +// qsa_attn + qsa_merge against the same kernels compiled from REF.metal (e.g. the file +// before an edit): part and xp memcmp-equal over pos0 0..30000 and 1, 2, 4, 8 tokens on +// random KV codes and ids, then paired GPU times of both +// QSA_ATTN=2 tests qsa_attn2 + qsa_merge2; REF is then qsa2.metal +// qsa_parity --paired REF_DIR [P] [--tokens N] [--prep-only] +// full-buffer memcmp and paired per-stage timings; prep-only avoids attention allocation +// qsa_parity --ab REF_DIR [P] +// paired A/B across decode T = 1, 2, 3, 4, 8 and prompt T = 1984, 2048 (with warmup + 12x reps) #include "mtl_rt.h" #include "ref_quant.h" @@ -25,17 +34,120 @@ struct Err { double l2() const { return sr > 0 ? std::sqrt(se / sr) : std::sqrt(se); } }; float h2f(uint16_t h) { return ref::f16_to_f32(h); } +// --bitwise: the attention and merge of this tree against REF's, bit for bit on the same inputs, and their times +int bitwise(const char* ref_path) { + constexpr int MAXT = 8, CELLS = 32768; + mtl::Context cx; + const bool v2 = std::getenv("QSA_ATTN") && std::atoi(std::getenv("QSA_ATTN")) == 2; + const int nch = v2 ? (BUDGET + 31) / 32 : NCH; + std::vector src_new{"qsa.metal"}, src_ref; + if (v2) { src_new.push_back("qsa2.metal"); src_ref.push_back("qsa.metal"); } + src_ref.push_back(ref_path); + const int lnew = cx.compile(src_new, "#define NO_GGML_COMMON 1\n"); + const int lref = cx.compile(src_ref, "#define NO_GGML_COMMON 1\n"); + const char* attn = v2 ? "qsa_attn2" : "qsa_attn"; + const char* merge = v2 ? "qsa_merge2" : "qsa_merge"; + mtl::Pipeline an = cx.pipeline(lnew, attn), mn = cx.pipeline(lnew, merge), + ar = cx.pipeline(lref, attn), mr = cx.pipeline(lref, merge); + const size_t part_bytes = 4 * (size_t)MAXT * nch * QH * (QD + 2), xp_bytes = 4 * (4 * MAXT * QH * QD + MAXT * QH * QD / 16); + mtl::Buffer qcur = cx.alloc(4 * MAXT * QH * QD), qfull = cx.alloc(4 * MAXT * QH * 2 * QD), + kq = cx.alloc((size_t)CELLS * QKV * QD), vq = cx.alloc((size_t)CELLS * QKV * QD), + ksc = cx.alloc(2 * (size_t)CELLS * QKV * 4), vsc = cx.alloc(2 * (size_t)CELLS * QKV * 4), + ids = cx.alloc(4 * MAXT * BUDGET), part[2] = {cx.alloc(part_bytes), cx.alloc(part_bytes)}, + xp[2] = {cx.alloc(xp_bytes), cx.alloc(xp_bytes)}; + ref::Rng r(41); + for (size_t i = 0; i < qcur.bytes / 4; ++i) qcur.as()[i] = r.sym() * 1.5f; + for (size_t i = 0; i < qfull.bytes / 4; ++i) qfull.as()[i] = r.sym() * 2.0f; + for (size_t i = 0; i < kq.bytes; ++i) { kq.as()[i] = (int8_t)(r.next() % 255 - 127); vq.as()[i] = (int8_t)(r.next() % 255 - 127); } + for (size_t i = 0; i < ksc.bytes / 2; ++i) { + ksc.as()[i] = ref::f32_to_f16(0.004f + 0.012f * r.uni()); + vsc.as()[i] = ref::f32_to_f16(0.004f + 0.012f * r.uni()); + } + auto fill_ids = [&](uint32_t pos0, int nt) { + for (int t = 0; t < nt; ++t) { + const int n_kv = (int)pos0 + t + 1, width = std::min(n_kv, BUDGET); + int* o = ids.as() + t * BUDGET; + if (n_kv <= BUDGET) { for (int c = 0; c < width; ++c) o[c] = c; continue; } + // width distinct cells of [0, n_kv) ascending, as the top-k emits them (folded into the KV buffer) + std::vector pick; + for (int c = 0, left = width; c < n_kv && left; ++c) + if (r.next() % (uint64_t)(n_kv - c) < (uint64_t)left) { pick.push_back(c % CELLS); --left; } + std::copy(pick.begin(), pick.end(), o); + } + }; + auto dispatches = [&](int v, uint32_t pos0, int nt) { + QsaArgs a{pos0, 0, (uint32_t)CELLS, (uint32_t)nt, 1e-6f}; + mtl::Dispatch at, m; + at.p = v ? &an : &ar; + at.args = {mtl::Arg::B(qcur), mtl::Arg::B(kq), mtl::Arg::B(vq), mtl::Arg::B(ksc), mtl::Arg::B(vsc), + mtl::Arg::B(ids), mtl::Arg::B(part[v]), mtl::Arg::V(a)}; + at.grid[0] = nch; at.grid[1] = v2 ? QKV * 2 : QKV; at.grid[2] = nt; at.tg[0] = 256; + m.p = v ? &mn : &mr; + m.args = {mtl::Arg::B(part[v]), mtl::Arg::B(qfull), mtl::Arg::B(xp[v]), mtl::Arg::V(a)}; + m.grid[0] = QH; m.grid[1] = nt; m.tg[0] = 256; + return std::vector{at, m}; + }; + int fails = 0; + for (uint32_t pos0 : {0u, 1u, 63u, 64u, 512u, 2047u, 2048u, 2050u, 8192u, 30000u}) + for (int nt : {1, 2, 4, 8}) { + fill_ids(pos0, nt); + for (int v = 0; v < 2; ++v) { + std::memset(part[v].ptr, 0, part_bytes); + std::memset(xp[v].ptr, 0, xp_bytes); + cx.run(dispatches(v, pos0, nt)); + } + const bool part_ok = !std::memcmp(part[0].ptr, part[1].ptr, part_bytes); + const bool xp_ok = !std::memcmp(xp[0].ptr, xp[1].ptr, xp_bytes); + const bool ok = part_ok && xp_ok; + fails += !ok; + if (!ok) std::printf("pos0 %5u ntok %d: part %s xp %s FAIL\n", pos0, nt, + part_ok ? "PASS" : "FAIL", xp_ok ? "PASS" : "FAIL"); + } + std::printf("bitwise: %d of 40 (pos0, ntok) cases differ %s\n", fails, fails ? "FAIL" : "PASS"); + std::printf("device %s: GPU us per dispatch, median of 9 paired rounds of 32 (ref / this)\n", cx.device_name().c_str()); + for (uint32_t pos0 : {512u, 2048u, 8192u}) + for (int nt : {1, 4}) { + fill_ids(pos0, nt); + std::vector d[2] = {dispatches(0, pos0, nt), dispatches(1, pos0, nt)}; + double us[2][2]; + for (int i = 0; i < 2; ++i) { + std::vector one[2] = {std::vector(32, d[0][i]), std::vector(32, d[1][i])}; + cx.warm(one[0], 0.05); + std::vector t[2]; + for (int k = 0; k < 9; ++k) + for (int turn = 0; turn < 2; ++turn) { + const int v = (k + turn) & 1; + t[v].push_back(cx.run(one[v]) / 32); + } + for (int v = 0; v < 2; ++v) us[i][v] = median(t[v]) * 1e6; + } + std::printf(" pos0 %5u ntok %d attn %6.2f / %6.2f merge %5.2f / %5.2f\n", pos0, nt, us[0][0], us[0][1], + us[1][0], us[1][1]); + } + return fails ? 1 : 0; +} } // namespace int main(int argc, char** argv) { bool bench = false; + bool ab_suite = false; + const char* paired_ref = nullptr; + bool prep_only = false; + int paired_nt = 3; int P = 3000; for (int i = 1; i < argc; ++i) { std::string a = argv[i]; + if (a == "--prep-only") { prep_only = true; continue; } + if (a == "--ab" && i + 1 < argc) { paired_ref = argv[++i]; ab_suite = true; continue; } + if (a == "--paired" && i + 1 < argc) { paired_ref = argv[++i]; continue; } + if (a == "--tokens" && i + 1 < argc) { paired_nt = std::atoi(argv[++i]); continue; } + if (a == "--bitwise" && i + 1 < argc) return bitwise(argv[i + 1]); if (a == "--bench") bench = true; else P = std::atoi(argv[i]); } - const int MAXT = 8, MAXC = std::max(4096, (P + 64) / 4 * 4 + 64); + if ((prep_only && !paired_ref) || (paired_nt > 128 && !prep_only)) return 2; + if (paired_ref && (paired_nt < 1 || paired_nt > 2048 || P < 0)) return 2; + const int MAXT = ab_suite ? 2048 : std::max(8, paired_nt), MAXC = std::max(4096, (P + MAXT + 64) / 4 * 4 + 64); mtl::Context cx; const int lib = cx.compile({"qsa.metal", "qsa2.metal"}, "#define NO_GGML_COMMON 1\n"); // QSA_TOPK=1: the v1 top-k; QSA_ATTN=1|2|3 (default 3): which attention (2 has its own chunking and merge) @@ -46,6 +158,7 @@ int main(int argc, char** argv) { pattn = cx.pipeline(lib, av == 1 ? "qsa_attn" : av == 2 ? "qsa_attn2" : "qsa_attn3"), pmerge = cx.pipeline(lib, v2 ? "qsa_merge2" : "qsa_merge"); const int MB = MAXC / 4 + 2; + const int AT = prep_only ? 1 : MAXT; // per-cell inputs for every cell (context + window) const int NC = P + MAXT; mtl::Buffer qfull = cx.alloc(4 * (size_t)NC * QH * 2 * QD), kc = cx.alloc(4 * (size_t)NC * QKV * QD), @@ -55,8 +168,8 @@ int main(int argc, char** argv) { qidx = cx.alloc(4 * MAXT * IH * ID), kq = cx.alloc((size_t)MAXC * QKV * QD), vq = cx.alloc((size_t)MAXC * QKV * QD), ksc = cx.alloc(2 * (size_t)MAXC * QKV * 4), vsc = cx.alloc(2 * (size_t)MAXC * QKV * 4), tail = cx.alloc(4 * 4 * ID), dead = cx.alloc(4 * ID), pooled = cx.alloc(4 * (size_t)MB * ID), scores = cx.alloc(4 * (size_t)MAXT * MB), - ids = cx.alloc(4 * MAXT * BUDGET), part = cx.alloc(4 * (size_t)MAXT * 2 * NCH * QH * (QD + 2)), - xp = cx.alloc(4 * (4 * MAXT * QH * QD + MAXT * QH * QD / 16)); + ids = cx.alloc(4 * AT * BUDGET), part = cx.alloc(4 * (size_t)AT * 2 * NCH * QH * (QD + 2)), + xp = cx.alloc(4 * (4 * AT * QH * QD + AT * QH * QD / 16)); ref::Rng r(23); for (size_t i = 0; i < qfull.bytes / 4; ++i) qfull.as()[i] = r.sym() * 2.0f; for (size_t i = 0; i < kc.bytes / 4; ++i) kc.as()[i] = r.sym() * 2.0f; @@ -95,7 +208,7 @@ int main(int argc, char** argv) { mtl::Arg::B(pooled), mtl::Arg::V(a), mtl::Arg::B(mtab)}; d.grid[0] = 33; d.grid[1] = nt; - d.tg[0] = 256; + d.tg[0] = std::getenv("QSA_PREP_TG") ? std::atoi(std::getenv("QSA_PREP_TG")) : 256; // 32: qsa_prep_sg return d; }; auto window = [&](uint32_t pos0, int nt, std::vector& ds) { @@ -131,6 +244,75 @@ int main(int argc, char** argv) { m.tg[0] = 256; ds.push_back(m); }; + if (paired_ref) { + const std::string dir = paired_ref; + const int lr = cx.compile({dir + "/qsa.metal", dir + "/qsa2.metal"}, "#define NO_GGML_COMMON 1\n"); + mtl::Pipeline pr[] = { + cx.pipeline(lr, "qsa_prep", mtl::Constants().b(24, mrope)), cx.pipeline(lr, "qsa_score"), + cx.pipeline(lr, std::getenv("QSA_TOPK") ? "qsa_topk" : "qsa_topk2"), + cx.pipeline(lr, av == 1 ? "qsa_attn" : av == 2 ? "qsa_attn2" : "qsa_attn3"), + cx.pipeline(lr, v2 ? "qsa_merge2" : "qsa_merge")}; + std::vector outputs{&qcur, &qidx, &kq, &vq, &ksc, &vsc, &tail, &dead, + &pooled, &scores, &ids, &part, &xp}; + const char* names[] = {"qcur", "qidx", "kq", "vq", "ksc", "vsc", "tail", "dead", + "pooled", "scores", "ids", "part", "xp"}; + auto run_case = [&](int nt, bool p_only) -> int { + std::vector ds[2]; + window(P, nt, ds[1]); + if (p_only) ds[1].resize(1); + ds[0] = ds[1]; + for (size_t i = 0; i < ds[0].size(); ++i) ds[0][i].p = &pr[i]; + ds[0][0].tg[0] = 256; // REF's qsa_prep as it was (QSA_PREP_TG sets the candidate's) + std::vector> expected; + int fails = 0; + for (int v = 0; v < 2; ++v) { + for (auto* b : outputs) std::memset(b->ptr, 0, b->bytes); + ref::Rng history(103); + for (auto* b : {&tail, &dead, &pooled}) + for (size_t i = 0; i < b->bytes / 4; ++i) b->as()[i] = history.sym(); + cx.run(ds[v]); + for (size_t i = 0; i < outputs.size(); ++i) { + auto* b = outputs[i]; + const auto* bytes = static_cast(b->ptr); + if (!v) expected.emplace_back(bytes, bytes + b->bytes); + else { + const bool ok = !std::memcmp(bytes, expected[i].data(), b->bytes); + if (!ab_suite || !ok) std::printf("%s %s\n", names[i], ok ? "PASS" : "FAIL"); + fails += !ok; + } + } + } + if (fails) return 1; + std::printf("paired P=%d nt=%d mrope=%d attn=%d%s, ref / candidate GPU us\n", P, nt, mrope, av, + ab_suite ? " (memcmp PASS)" : ""); + const char* stages[] = {"prep", "score", "topk", "attn", "merge"}; + for (size_t i = 0; i < ds[0].size(); ++i) { + const int repeats = nt > 128 ? 4 : 16, pairs = ab_suite ? 5 : 9; + cx.warm({ds[0][i], ds[1][i]}, 0.01); + cx.run({ds[0][i]}, repeats); + cx.run({ds[1][i]}, repeats); + std::vector times[2], ratios; + for (int k = 0; k < pairs; ++k) { + double t0 = 0, t1 = 0; + for (int turn = 0; turn < 2; ++turn) { + const int v = (k + turn) & 1; + const double dt = cx.run({ds[v][i]}, repeats) * 1e6 / repeats; + if (v == 0) t0 = dt; else t1 = dt; + times[v].push_back(dt); + } + ratios.push_back(t1 / t0); + } + std::printf("%s %.3f / %.3f (ratio %.4f)\n", stages[i], median(times[0]), median(times[1]), median(ratios)); + } + return 0; + }; + if (ab_suite) { + for (int nt : {1, 2, 3, 4, 8, 1984, 2048}) + if (run_case(nt, prep_only || nt > 8)) return 1; + return 0; + } + return run_case(paired_nt, prep_only); + } // the context, one cell at a time { std::vector ds; diff --git a/metal/tests/qsa_prompt_parity.cpp b/metal/tests/qsa_prompt_parity.cpp index 18e1b3af..3363f283 100644 --- a/metal/tests/qsa_prompt_parity.cpp +++ b/metal/tests/qsa_prompt_parity.cpp @@ -1,7 +1,12 @@ // metal/tests/qsa_prompt_parity.cpp - qsa_attn_mm (qsa_prompt.metal, the matrix units, fp16 operands) against qsa_attn // (fp32, itself checked by qsa_parity), both followed by qsa_merge: the attention output of T query tokens on random // int8 caches, within the budget (all cells, causal) and past it (2051 selected cells per token); then the speed. -// qsa_prompt_parity [--bench] +// qsa_prompt_parity [--bench | --flash-compare T pos0 baseline_qsa_prompt.metal [layers pairs]] +// qsa_prompt_parity --score-compare T pos0 baseline_qsa_prompt.metal baseline_qsa.metal pairs +// qsa_prompt_parity --topk-compare T pos0 baseline_qsa.metal pairs +// QSA_TOPK_EDGE=1: ties/signed zero; QSA_FLASH_EDGE=1|2: signed zero/tiny cancellation +// qsa_prompt_parity --mm-compare attn|score T pos0 baseline_qsa_prompt.metal pairs +// qsa_prompt_parity --ab T pos0 baseline_qsa_prompt.metal baseline_qsa.metal pairs #include "coop_probe.h" #include "mtl_rt.h" #include "ref_quant.h" @@ -18,9 +23,392 @@ namespace { constexpr int QH = 24, QKV = 2, QD = 256, QB = 2051, NCH = 33; struct QsaArgs { uint32_t pos0, pos_base, max_cells, n_tok; float eps; }; double median(std::vector v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; } + +int topk_compare(int T, uint32_t pos0, const std::string& baseline, int pairs) { + if (T < 1 || T > 1984 || pos0 > 30000 || pairs < 1 || pairs > 15) return 2; + mtl::Context cx; + const int lib = cx.compile({"qsa.metal"}, "#define NO_GGML_COMMON 1\n"); + const int old = cx.compile({baseline}, "#define NO_GGML_COMMON 1\n"); + const uint32_t cells = pos0 + T, stride = cells / 4 + 2; + auto scores = cx.alloc(4 * size_t(T) * stride), out = cx.alloc(4 * size_t(T) * QB), + ref = cx.alloc(out.bytes); + ref::Rng r(5); + const bool edge = std::getenv("QSA_TOPK_EDGE") != nullptr; + for (int t = 0; t < T; ++t) + for (uint32_t b = 0; b < stride; ++b) { + float v = std::fabs(r.sym()) * 32.0f; + if (edge) { + if (t % 4 == 0) v = 1.0f; + if (t % 4 == 1) v = float(int(r.next() % 7) - 3); + if (t % 4 == 2) v = b & 1 ? -0.0f : 0.0f; + } + if (!edge && b == (pos0 + t + 1) / 4 && ((pos0 + t + 1) & 3)) v += 1e9f; + scores.as()[size_t(t) * stride + b] = v; + } + std::memset(out.ptr, 0xa5, out.bytes); std::memset(ref.ptr, 0xa5, ref.bytes); + auto p = cx.pipeline(lib, "qsa_topk"), pold = cx.pipeline(old, "qsa_topk"); + QsaArgs a{pos0, 0, cells, uint32_t(T), 1e-6f}; + mtl::Dispatch d; + d.p = &pold; + d.args = {mtl::Arg::B(scores), mtl::Arg::B(ref), mtl::Arg::V(a)}; + d.grid[0] = T; d.tg[0] = 1024; + auto candidate = d; + candidate.p = &p; candidate.args[1] = mtl::Arg::B(out); + cx.run({d}); cx.run({candidate}); + const bool equal = std::memcmp(out.ptr, ref.ptr, out.bytes) == 0; + std::vector before, after; + for (int i = 0; i < pairs; ++i) { + if (i & 1) { after.push_back(cx.run({candidate})); before.push_back(cx.run({d})); } + else { before.push_back(cx.run({d})); after.push_back(cx.run({candidate})); } + } + std::printf("device %s; topk T %d pos0 %u memcmp %s before %.6f after %.6f ms\n", + cx.device_name().c_str(), T, pos0, equal ? "PASS" : "FAIL", median(before) * 1e3, median(after) * 1e3); + for (int i = 0; i < pairs; ++i) + std::printf("pair %d before %.6f after %.6f ms\n", i, before[i] * 1e3, after[i] * 1e3); + return equal ? 0 : 1; +} + +int score_compare(int T, uint32_t pos0, const std::string& baseline_prompt, + const std::string& baseline_qsa, int pairs) { + if (T < 1 || T > 1984 || pos0 > 30000 || pairs < 1 || pairs > 15) return 2; + mtl::Context cx; + const int lib = cx.compile({"qsa.metal", "qsa_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); + const int old = cx.compile({baseline_qsa, baseline_prompt}, "#define NO_GGML_COMMON 1\n"); + const size_t cells = pos0 + T, nb = cells / 4 + 1, stride = cells / 4 + 2; + auto q = cx.alloc(4 * size_t(T) * 4 * 128), k = cx.alloc(4 * nb * 128), dead = cx.alloc(4 * 128); + auto scores = cx.alloc(4 * size_t(T) * stride), ref_scores = cx.alloc(scores.bytes), + ids = cx.alloc(4 * size_t(T) * QB), ref_ids = cx.alloc(ids.bytes); + ref::Rng r(5); + for (auto* b : {&q, &k, &dead}) + for (size_t i = 0; i < b->bytes / 4; ++i) b->as()[i] = r.sym(); + for (auto* b : {&scores, &ref_scores, &ids, &ref_ids}) std::memset(b->ptr, 0xa5, b->bytes); + const auto constants = mtl::Constants().b(7, coop_k4(cx)); + auto score = cx.pipeline(lib, "qsa_score_mm", constants), + old_score = cx.pipeline(old, "qsa_score_mm", constants), + topk = cx.pipeline(lib, "qsa_topk"), old_topk = cx.pipeline(old, "qsa_topk"); + QsaArgs a{pos0, 0, uint32_t(cells), uint32_t(T), 1e-6f}; + mtl::Dispatch ds, dt; + ds.p = &old_score; + ds.args = {mtl::Arg::B(q), mtl::Arg::B(k), mtl::Arg::B(dead), mtl::Arg::B(ref_scores), mtl::Arg::V(a)}; + ds.grid[0] = (nb + 63) / 64; ds.grid[1] = (T + 15) / 16; ds.tg[0] = 128; + dt.p = &old_topk; + dt.args = {mtl::Arg::B(ref_scores), mtl::Arg::B(ref_ids), mtl::Arg::V(a)}; + dt.grid[0] = T; dt.tg[0] = 1024; + auto cs = ds, ct = dt; + cs.p = &score; cs.args[3] = mtl::Arg::B(scores); + ct.p = &topk; ct.args[0] = mtl::Arg::B(scores); ct.args[1] = mtl::Arg::B(ids); + cx.run({ds, dt}, 1, false, true); cx.run({cs, ct}, 1, false, true); + const bool same_scores = std::memcmp(scores.ptr, ref_scores.ptr, scores.bytes) == 0; + const bool same_ids = std::memcmp(ids.ptr, ref_ids.ptr, ids.bytes) == 0; + std::printf("device %s; score+topk T %d pos0 %u scores %s qids %s tg %d -> %d bytes\n", + cx.device_name().c_str(), T, pos0, same_scores ? "PASS" : "FAIL", same_ids ? "PASS" : "FAIL", + old_score.static_tg_mem, score.static_tg_mem); + const std::vector before_dispatches[] = {{ds}, {dt}, {ds, dt}}; + const std::vector after_dispatches[] = {{cs}, {ct}, {cs, ct}}; + for (int stage = 0; stage < 3; ++stage) { + const auto& before_dispatch = before_dispatches[stage]; + const auto& after_dispatch = after_dispatches[stage]; + cx.run(before_dispatch, 1, false, true); cx.run(after_dispatch, 1, false, true); + std::vector before, after; + for (int i = 0; i < pairs; ++i) { + if (i & 1) { + after.push_back(cx.run(after_dispatch, 1, false, true)); + before.push_back(cx.run(before_dispatch, 1, false, true)); + } else { + before.push_back(cx.run(before_dispatch, 1, false, true)); + after.push_back(cx.run(after_dispatch, 1, false, true)); + } + } + std::printf("%s before %.6f after %.6f ms\n", stage == 0 ? "score" : stage == 1 ? "topk" : "combined", + median(before) * 1e3, median(after) * 1e3); + for (int i = 0; i < pairs; ++i) + std::printf("pair %d before %.6f after %.6f ms\n", i, before[i] * 1e3, after[i] * 1e3); + } + const bool stable = std::memcmp(scores.ptr, ref_scores.ptr, scores.bytes) == 0 && + std::memcmp(ids.ptr, ref_ids.ptr, ids.bytes) == 0; + std::printf("post-timing memcmp %s\n", stable ? "PASS" : "FAIL"); + return same_scores && same_ids && stable ? 0 : 1; +} + +int flash_compare(int T, uint32_t pos0, const std::string& baseline, int layers, int pairs) { + if (T < 1 || T > 1984 || pos0 > 30000 || layers < 1 || pairs < 1) return 2; + mtl::Context cx; + const int lib = cx.compile({"qsa.metal", "qsa_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); + const int old = cx.compile({"qsa.metal", baseline}, "#define NO_GGML_COMMON 1\n"); + const size_t cells = pos0 + T, nq = (size_t)T * QH * QD; + mtl::Buffer kq = cx.alloc(cells * QKV * QD), vq = cx.alloc(kq.bytes), + ks = cx.alloc(2 * cells * QKV * 4), vs = cx.alloc(ks.bytes), + qh = cx.alloc(2 * nq), gate = cx.alloc(8 * nq), ids = cx.alloc(4 * (size_t)T * QB), + out = cx.alloc(2 * nq), ref = cx.alloc(out.bytes); + ref::Rng r(5); + for (size_t i = 0; i < kq.bytes; ++i) { + kq.as()[i] = int8_t(int(r.next() % 255) - 127); + vq.as()[i] = int8_t(int(r.next() % 255) - 127); + } + for (size_t i = 0; i < ks.bytes / 2; ++i) { + ks.as()[i] = ref::f32_to_f16(0.004f + 0.004f * r.uni()); + vs.as()[i] = ref::f32_to_f16(0.004f + 0.004f * r.uni()); + } + for (size_t i = 0; i < nq; ++i) qh.as()[i] = ref::f32_to_f16(r.sym() * 3.0f); + for (size_t i = 0; i < 2 * nq; ++i) gate.as()[i] = r.sym(); + for (int t = 0; t < T; ++t) { + const uint32_t n = pos0 + t + 1; + int* o = ids.as() + (size_t)t * QB; + if (n <= QB) { + for (uint32_t c = 0; c < n; ++c) o[c] = int(c); + } else { + std::set selected; + while (selected.size() < QB) selected.insert(int(r.next() % n)); + std::copy(selected.begin(), selected.end(), o); + } + } + if (const char* edge = std::getenv("QSA_FLASH_EDGE")) { + const int mode = std::atoi(edge); + for (size_t i = 0; i < kq.bytes; ++i) { + kq.as()[i] = 0; + vq.as()[i] = mode == 2 ? int8_t((i / (QKV * QD)) & 1 ? -1 : 1) : 0; + } + for (size_t i = 0; i < ks.bytes / 2; ++i) { + ks.as()[i] = i & 1 ? 0x8000 : 0; + vs.as()[i] = mode == 2 ? 0x0400 : (i & 1 ? 0x8000 : 0); + } + for (size_t i = 0; i < nq; ++i) qh.as()[i] = i & 1 ? 0x8000 : 0; + } + const auto constants = mtl::Constants().b(7, coop_k4(cx)); + auto p = cx.pipeline(lib, "qsa_flash", constants), pold = cx.pipeline(old, "qsa_flash", constants); + QsaArgs a{pos0, 0, uint32_t(cells), uint32_t(T), 1e-6f}; + mtl::Dispatch d; + d.p = &pold; + d.args = {mtl::Arg::B(qh), mtl::Arg::B(kq), mtl::Arg::B(vq), mtl::Arg::B(ks), mtl::Arg::B(vs), + mtl::Arg::B(ids), mtl::Arg::B(gate), mtl::Arg::B(ref), mtl::Arg::V(a)}; + d.grid[0] = QKV; d.grid[1] = T; d.tg[0] = 128; + cx.run({d}); + // Poison only implicit rows, including the early rows of a chunk crossing the budget. + for (int t = 0; t < T && pos0 + t + 1 <= QB; ++t) + std::fill_n(ids.as() + (size_t)t * QB, QB, -1); + auto candidate = d; + candidate.p = &p; candidate.args[7] = mtl::Arg::B(out); + cx.run({candidate}); + const bool equal = std::memcmp(out.ptr, ref.ptr, out.bytes) == 0; + // Baseline needs explicit ids for the paired timings below. + for (int t = 0; t < T && pos0 + t + 1 <= QB; ++t) + for (uint32_t c = 0; c < pos0 + t + 1; ++c) ids.as()[(size_t)t * QB + c] = int(c); + // paired, alternating order: each sample is one command buffer of `layers` back-to-back dispatches (the 12 QSA + // layers of a prompt chunk) + std::vector before, after; + for (int i = 0; i < pairs; ++i) { + if (i & 1) { after.push_back(cx.run({candidate}, layers)); before.push_back(cx.run({d}, layers)); } + else { before.push_back(cx.run({d}, layers)); after.push_back(cx.run({candidate}, layers)); } + } + std::printf("device %s; T %d pos0 %u: memcmp %s; flash x%d before median %.3f min %.3f ms, after median %.3f min %.3f " + "ms; tg %d -> %d bytes\n", cx.device_name().c_str(), T, pos0, equal ? "PASS" : "FAIL", layers, + median(before) * 1e3, *std::min_element(before.begin(), before.end()) * 1e3, median(after) * 1e3, + *std::min_element(after.begin(), after.end()) * 1e3, pold.static_tg_mem, p.static_tg_mem); + for (int i = 0; i < pairs; ++i) + std::printf("pair %d before %.6f after %.6f ms\n", i, before[i] * 1e3, after[i] * 1e3); + return equal ? 0 : 1; +} + +int mm_compare(bool score, int T, uint32_t pos0, const std::string& baseline, int pairs) { + if (T < 1 || T > 1984 || pos0 > 30000 || pairs < 1 || pairs > 15) return 2; + mtl::Context cx; + const int lib = cx.compile({"qsa.metal", "qsa_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); + const int old = cx.compile({"qsa.metal", baseline}, "#define NO_GGML_COMMON 1\n"); + const size_t cells = pos0 + T; + QsaArgs a{pos0, 0, uint32_t(cells), uint32_t(T), 1e-6f}; + const bool k4 = coop_k4(cx); + const auto constants = mtl::Constants().b(7, k4); + std::printf("device %s coop_k4 %d\n", cx.device_name().c_str(), int(k4)); + const std::string name = score ? "qsa_score_mm" : "qsa_attn_mm"; + auto p = cx.pipeline(lib, name.c_str(), constants), pold = cx.pipeline(old, name.c_str(), constants); + mtl::Dispatch d; + d.p = &pold; + mtl::Buffer out, ref; + auto compare = [&](const mtl::Dispatch& before_dispatch, const mtl::Dispatch& after_dispatch) { + cx.run({before_dispatch}); cx.run({after_dispatch}); + const bool equal = std::memcmp(out.ptr, ref.ptr, out.bytes) == 0; + cx.warm({before_dispatch, after_dispatch}, 0.05); + std::vector before, after; + for (int i = 0; i < pairs; ++i) { + if (i & 1) { + after.push_back(cx.run({after_dispatch})); before.push_back(cx.run({before_dispatch})); + } else { + before.push_back(cx.run({before_dispatch})); after.push_back(cx.run({after_dispatch})); + } + } + std::printf("%s T %d pos0 %u memcmp %s before %.6f after %.6f ms tg %d -> %d\n", name.c_str(), T, pos0, + equal ? "PASS" : "FAIL", median(before) * 1e3, median(after) * 1e3, pold.static_tg_mem, p.static_tg_mem); + for (int i = 0; i < pairs; ++i) + std::printf("pair %d before %.6f after %.6f ms\n", i, before[i] * 1e3, after[i] * 1e3); + const bool stable = std::memcmp(out.ptr, ref.ptr, out.bytes) == 0; + std::printf("post-timing memcmp %s\n", stable ? "PASS" : "FAIL"); + return equal && stable ? 0 : 1; + }; + ref::Rng r(5); + if (score) { + const size_t nb = cells / 4 + 1, stride = cells / 4 + 2; + auto q = cx.alloc(4 * size_t(T) * 4 * 128), k = cx.alloc(4 * nb * 128), dead = cx.alloc(4 * 128); + for (auto* b : {&q, &k, &dead}) + for (size_t i = 0; i < b->bytes / 4; ++i) b->as()[i] = r.sym(); + out = cx.alloc(4 * size_t(T) * stride); ref = cx.alloc(out.bytes); + std::memset(out.ptr, 0xa5, out.bytes); std::memset(ref.ptr, 0xa5, ref.bytes); + d.args = {mtl::Arg::B(q), mtl::Arg::B(k), mtl::Arg::B(dead), mtl::Arg::B(ref), mtl::Arg::V(a)}; + d.grid[0] = (nb + 63) / 64; d.grid[1] = (T + 15) / 16; d.tg[0] = 128; + auto candidate = d; + candidate.p = &p; candidate.args[3] = mtl::Arg::B(out); + return compare(d, candidate); + } + auto k = cx.alloc(cells * QKV * QD), v = cx.alloc(k.bytes), ks = cx.alloc(2 * cells * QKV * 4), vs = cx.alloc(ks.bytes), + q = cx.alloc(2 * size_t(T) * QH * QD), ids = cx.alloc(4 * size_t(T) * QB); + for (size_t i = 0; i < k.bytes; ++i) { + k.as()[i] = int8_t(int(r.next() % 255) - 127); v.as()[i] = int8_t(int(r.next() % 255) - 127); + } + for (auto* b : {&ks, &vs}) + for (size_t i = 0; i < b->bytes / 2; ++i) b->as()[i] = ref::f32_to_f16(0.004f + 0.004f * r.uni()); + for (size_t i = 0; i < q.bytes / 2; ++i) q.as()[i] = ref::f32_to_f16(r.sym() * 3.0f); + for (int t = 0; t < T; ++t) { + const uint32_t n = pos0 + t + 1; + for (int c = 0; c < QB; ++c) ids.as()[size_t(t) * QB + c] = n <= QB ? -1 : int(uint64_t(c) * n / QB); + } + out = cx.alloc(4 * size_t(T) * NCH * QH * (QD + 2)); ref = cx.alloc(out.bytes); + std::memset(out.ptr, 0xa5, out.bytes); std::memset(ref.ptr, 0xa5, ref.bytes); + d.args = {mtl::Arg::B(q), mtl::Arg::B(k), mtl::Arg::B(v), mtl::Arg::B(ks), mtl::Arg::B(vs), + mtl::Arg::B(ids), mtl::Arg::B(ref), mtl::Arg::V(a)}; + d.grid[0] = NCH; d.grid[1] = QKV; d.grid[2] = T; d.tg[0] = 32; + auto candidate = d; + candidate.p = &p; candidate.args[6] = mtl::Arg::B(out); + const int result = compare(d, candidate); + const size_t nq = size_t(T) * QH * QD; + auto gate = cx.alloc(8 * nq), merged = cx.alloc(4 * (4 * nq + nq / 16)), merged_ref = cx.alloc(merged.bytes); + for (size_t i = 0; i < gate.bytes / 4; ++i) gate.as()[i] = r.sym(); + auto merge = cx.pipeline(lib, "qsa_merge"); + mtl::Dispatch dm; + dm.p = &merge; + dm.args = {mtl::Arg::B(ref), mtl::Arg::B(gate), mtl::Arg::B(merged_ref), mtl::Arg::V(a)}; + dm.grid[0] = QH; dm.grid[1] = T; dm.tg[0] = 256; + auto cm = dm; + cm.args[0] = mtl::Arg::B(out); cm.args[2] = mtl::Arg::B(merged); + cx.run({dm}); cx.run({cm}); + const bool same_merge = std::memcmp(merged.ptr, merged_ref.ptr, merged.bytes) == 0; + std::printf("merge+prepared memcmp %s\n", same_merge ? "PASS" : "FAIL"); + return result || !same_merge; +} } // namespace +int ab_compare(int T, uint32_t pos0, const std::string& baseline_prompt, const std::string& baseline_qsa, int pairs) { + if (T < 1 || T > 2048 || pos0 > 30000 || pairs < 1 || pairs > 15) return 2; + mtl::Context cx; + const int lib = cx.compile({"qsa.metal", "qsa_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); + const int old = cx.compile({baseline_qsa, baseline_prompt}, "#define NO_GGML_COMMON 1\n"); + const size_t cells = pos0 + T, nb = cells / 4 + 1, stride = cells / 4 + 2; + QsaArgs a{pos0, 0, uint32_t(cells), uint32_t(T), 1e-6f}; + const bool mrope = std::getenv("QSA_MROPE") != nullptr; + const bool qf16 = std::getenv("QSA_QF16") != nullptr; + const auto prep_c = mtl::Constants().b(24, mrope).b(25, qf16); + const auto score_c = mtl::Constants().b(7, coop_k4(cx)); + auto prep = cx.pipeline(lib, "qsa_prep", prep_c), old_prep = cx.pipeline(old, "qsa_prep", prep_c); + auto score = cx.pipeline(lib, "qsa_score_mm", score_c), old_score = cx.pipeline(old, "qsa_score_mm", score_c); + mtl::Buffer qfull = cx.alloc(4 * size_t(T) * QH * 2 * QD), kc = cx.alloc(4 * size_t(T) * QKV * QD), + vc = cx.alloc(4 * size_t(T) * QKV * QD), iq = cx.alloc(4 * size_t(T) * 4 * 128), + ik = cx.alloc(4 * size_t(T) * 128), qn = cx.alloc(4 * QD), kn = cx.alloc(4 * QD), + iqn = cx.alloc(4 * 128), ikn = cx.alloc(4 * 128), cs = cx.alloc(8 * (cells + 16) * 32), + mtab = cx.alloc(4 * 3 * (cells + 16)); + mtl::Buffer qcur[2] = {cx.alloc((qf16 ? 2 : 4) * size_t(T) * QH * QD), cx.alloc((qf16 ? 2 : 4) * size_t(T) * QH * QD)}, + qidx[2] = {cx.alloc(4 * size_t(T) * 4 * 128), cx.alloc(4 * size_t(T) * 4 * 128)}, + kq[2] = {cx.alloc(cells * QKV * QD), cx.alloc(cells * QKV * QD)}, + vq[2] = {cx.alloc(cells * QKV * QD), cx.alloc(cells * QKV * QD)}, + ksc[2] = {cx.alloc(2 * cells * QKV * 4), cx.alloc(2 * cells * QKV * 4)}, + vsc[2] = {cx.alloc(2 * cells * QKV * 4), cx.alloc(2 * cells * QKV * 4)}, + tail[2] = {cx.alloc(4 * 4 * 128), cx.alloc(4 * 4 * 128)}, + dead[2] = {cx.alloc(4 * 128), cx.alloc(4 * 128)}, + pooled[2] = {cx.alloc(4 * stride * 128), cx.alloc(4 * stride * 128)}, + scores[2] = {cx.alloc(4 * size_t(T) * stride), cx.alloc(4 * size_t(T) * stride)}; + ref::Rng r(5); + for (auto* b : {&qfull, &kc, &vc, &iq, &ik}) + for (size_t i = 0; i < b->bytes / 4; ++i) b->as()[i] = r.sym(); + for (int i = 0; i < QD; ++i) { qn.as()[i] = 1.0f + 0.2f * r.sym(); kn.as()[i] = 1.0f + 0.2f * r.sym(); } + for (int i = 0; i < 128; ++i) { iqn.as()[i] = 1.0f + 0.2f * r.sym(); ikn.as()[i] = 1.0f + 0.2f * r.sym(); } + for (size_t p = 0; p < cells + 16; ++p) { + for (int k = 0; k < 3; ++k) mtab.as()[p * 3 + k] = int(r.next() % (cells + 16)); + for (int i = 0; i < 32; ++i) { + const double inv = std::pow(1e7, -2.0 * i / 64.0), th = double(p) * inv; + cs.as()[(p * 32 + i) * 2] = float(std::cos(th)); + cs.as()[(p * 32 + i) * 2 + 1] = float(std::sin(th)); + } + } + auto reset_state = [&](int v) { + ref::Rng h(103); + for (auto* b : {&qcur[v], &qidx[v], &kq[v], &vq[v], &ksc[v], &vsc[v], &scores[v]}) + std::memset(b->ptr, 0xa5, b->bytes); + for (auto* b : {&tail[v], &dead[v], &pooled[v]}) + for (size_t i = 0; i < b->bytes / 4; ++i) b->as()[i] = h.sym(); + }; + mtl::Dispatch dp[2], ds[2]; + for (int v = 0; v < 2; ++v) { + dp[v].p = v ? &prep : &old_prep; + dp[v].args = {mtl::Arg::B(qfull), mtl::Arg::B(kc), mtl::Arg::B(vc), mtl::Arg::B(iq), mtl::Arg::B(ik), + mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(iqn), mtl::Arg::B(ikn), mtl::Arg::B(cs), + mtl::Arg::B(qcur[v]), mtl::Arg::B(qidx[v]), mtl::Arg::B(kq[v]), mtl::Arg::B(vq[v]), + mtl::Arg::B(ksc[v]), mtl::Arg::B(vsc[v]), mtl::Arg::B(tail[v]), mtl::Arg::B(dead[v]), + mtl::Arg::B(pooled[v]), mtl::Arg::V(a), mtl::Arg::B(mtab)}; + dp[v].grid[0] = 33; dp[v].grid[1] = T; + dp[v].tg[0] = v && std::getenv("QSA_PREP_TG") ? std::atoi(std::getenv("QSA_PREP_TG")) : 256; // 32: qsa_prep_sg + ds[v].p = v ? &score : &old_score; + ds[v].args = {mtl::Arg::B(qidx[v]), mtl::Arg::B(pooled[v]), mtl::Arg::B(dead[v]), mtl::Arg::B(scores[v]), mtl::Arg::V(a)}; + ds[v].grid[0] = (nb + 63) / 64; ds[v].grid[1] = (T + 15) / 16; ds[v].tg[0] = 128; + reset_state(v); + cx.run({dp[v], ds[v]}, 1, false, true); + } + const mtl::Buffer* b0[] = {&qcur[0], &qidx[0], &kq[0], &vq[0], &ksc[0], &vsc[0], &tail[0], &dead[0], &pooled[0], &scores[0]}; + const mtl::Buffer* b1[] = {&qcur[1], &qidx[1], &kq[1], &vq[1], &ksc[1], &vsc[1], &tail[1], &dead[1], &pooled[1], &scores[1]}; + bool ok = true; + for (int i = 0; i < 10; ++i) ok = ok && (std::memcmp(b0[i]->ptr, b1[i]->ptr, b0[i]->bytes) == 0); + std::printf("device %s; ab T %d pos0 %u mrope %d qf16 %d memcmp %s\n", + cx.device_name().c_str(), T, pos0, int(mrope), int(qf16), ok ? "PASS" : "FAIL"); + if (!ok) return 1; + const char* sname[] = {"qsa_prep", "qsa_score_mm", "combined"}; + const std::vector d0[] = {{dp[0]}, {ds[0]}, {dp[0], ds[0]}}; + const std::vector d1[] = {{dp[1]}, {ds[1]}, {dp[1], ds[1]}}; + for (int s = 0; s < 3; ++s) { + cx.warm(d0[s], 0.05); + cx.run(d0[s], 12, false, true); + cx.run(d1[s], 12, false, true); + std::vector before, after, ratios; + for (int i = 0; i < pairs; ++i) { + double t0 = 0, t1 = 0; + if (i & 1) { + t1 = cx.run(d1[s], 12, false, true) * 1e3 / 12.0; + t0 = cx.run(d0[s], 12, false, true) * 1e3 / 12.0; + } else { + t0 = cx.run(d0[s], 12, false, true) * 1e3 / 12.0; + t1 = cx.run(d1[s], 12, false, true) * 1e3 / 12.0; + } + before.push_back(t0); after.push_back(t1); ratios.push_back(t1 / t0); + } + std::printf(" %-12s before %.4f after %.4f ms (ratio %.4f)\n", sname[s], median(before), median(after), median(ratios)); + } + return 0; +} + int main(int argc, char** argv) { + if (argc == 7 && std::string(argv[1]) == "--ab") + return ab_compare(std::atoi(argv[2]), uint32_t(std::strtoul(argv[3], nullptr, 10)), + argv[4], argv[5], std::atoi(argv[6])); + if (argc == 7 && std::string(argv[1]) == "--score-compare") + return score_compare(std::atoi(argv[2]), uint32_t(std::strtoul(argv[3], nullptr, 10)), + argv[4], argv[5], std::atoi(argv[6])); + if (argc == 6 && std::string(argv[1]) == "--topk-compare") + return topk_compare(std::atoi(argv[2]), uint32_t(std::strtoul(argv[3], nullptr, 10)), + argv[4], std::atoi(argv[5])); + if (argc == 7 && std::string(argv[1]) == "--mm-compare") { + const std::string target = argv[2]; + if (target != "attn" && target != "score") return 2; + return mm_compare(target == "score", std::atoi(argv[3]), + uint32_t(std::strtoul(argv[4], nullptr, 10)), argv[5], std::atoi(argv[6])); + } + if ((argc == 5 || argc == 7) && std::string(argv[1]) == "--flash-compare") + return flash_compare(std::atoi(argv[2]), uint32_t(std::strtoul(argv[3], nullptr, 10)), argv[4], + argc == 7 ? std::atoi(argv[5]) : 1, argc == 7 ? std::atoi(argv[6]) : 3); const bool bench = argc > 1 && std::string(argv[1]) == "--bench"; std::setvbuf(stdout, nullptr, _IOLBF, 0); const bool noflash = std::getenv("QSA_NOFLASH") != nullptr; @@ -44,7 +432,7 @@ int main(int argc, char** argv) { std::printf("pipelines: qsa_attn_mm max threads %d, tg mem %d; qsa_flash max threads %d, tg mem %d\n", pmm.max_threads, pmm.static_tg_mem, pf.max_threads, pf.static_tg_mem); int fails = 0; - for (uint32_t pos0 : {0u, 30000u}) { + for (uint32_t pos0 : {0u, 1984u, 2050u, 30000u}) { // ids: within the budget the identity; past it 2051 distinct ascending cells < n_kv for (int t = 0; t < T; ++t) { const uint32_t n_kv = pos0 + t + 1; @@ -82,6 +470,8 @@ int main(int argc, char** argv) { } cx.run(ref_ds); std::fill(part.as(), part.as() + part.bytes / 4, std::nanf("")); // a kernel that writes nothing shows + for (int t = 0; t < T && pos0 + t + 1 <= QB; ++t) + std::fill_n(ids.as() + (size_t)t * QB, QB, -1); cx.run(mm_ds); if (std::getenv("QSA_DEBUG")) { // how much of the partials the matrix-unit kernel wrote, and token 5's first chunk (m, l, acc[0..3]) @@ -155,6 +545,77 @@ int main(int argc, char** argv) { std::printf("score pos0 %u, T %d: qsa_score_mm vs qsa_score worst rel %.2e %s\n", pos0, T, worst, ok ? "PASS" : "FAIL"); } } + // qsa_prep with QSA_Q_F16_FC=true vs qsa_prep (fp32) + qsa_q_f16: 100% bitwise parity on qcurh and qsa_flash output + { + const uint32_t pos0 = 0, MB = MAXC / 4 + 2; + mtl::Buffer kc = cx.alloc(4 * (size_t)T * QKV * QD), vc = cx.alloc(4 * (size_t)T * QKV * QD), + iq = cx.alloc(4 * (size_t)T * 4 * 128), ik = cx.alloc(4 * (size_t)T * 128), + qn = cx.alloc(4 * QD), kn = cx.alloc(4 * QD), iqn = cx.alloc(4 * 128), ikn = cx.alloc(4 * 128), + cs = cx.alloc(8 * (size_t)(T + 16) * 32), qidx = cx.alloc(4 * (size_t)T * 4 * 128), + tail = cx.alloc(4 * 4 * 128), dead = cx.alloc(4 * 128), pooled = cx.alloc(4 * (size_t)MB * 128), + mtab = cx.alloc(4 * 3 * (size_t)(T + 16)), + qh_ref = cx.alloc(2 * (size_t)T * QH * QD), qh_fused = cx.alloc(2 * (size_t)T * QH * QD), + xf_ref = cx.alloc(2 * (size_t)T * QH * QD), xf_fused = cx.alloc(2 * (size_t)T * QH * QD); + for (size_t i = 0; i < kc.bytes / 4; ++i) { kc.as()[i] = r.sym(); vc.as()[i] = r.sym(); } + for (size_t i = 0; i < iq.bytes / 4; ++i) iq.as()[i] = r.sym(); + for (size_t i = 0; i < ik.bytes / 4; ++i) ik.as()[i] = r.sym(); + for (int i = 0; i < QD; ++i) { qn.as()[i] = 1.0f + 0.2f * r.sym(); kn.as()[i] = 1.0f + 0.2f * r.sym(); } + for (int i = 0; i < 128; ++i) { iqn.as()[i] = 1.0f + 0.2f * r.sym(); ikn.as()[i] = 1.0f + 0.2f * r.sym(); } + for (int p = 0; p < T + 16; ++p) + for (int i = 0; i < 32; ++i) { + const double inv = std::pow(1e7, -2.0 * i / 64.0), th = (double)p * inv; + cs.as()[(p * 32 + i) * 2] = (float)std::cos(th); + cs.as()[(p * 32 + i) * 2 + 1] = (float)std::sin(th); + } + QsaArgs a{pos0, 0, (uint32_t)MAXC, (uint32_t)T, 1e-6f}; + const uint32_t nq = T * QH * QD; + mtl::Pipeline pprep_f32 = cx.pipeline(lib, "qsa_prep"), + pprep_f16 = cx.pipeline(lib, "qsa_prep", mtl::Constants().b(25, true)); + auto make_prep = [&](const mtl::Pipeline& pp, const mtl::Buffer& qout) { + mtl::Dispatch d; + d.p = &pp; + d.args = {mtl::Arg::B(qfull), mtl::Arg::B(kc), mtl::Arg::B(vc), mtl::Arg::B(iq), mtl::Arg::B(ik), + mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(iqn), mtl::Arg::B(ikn), mtl::Arg::B(cs), + mtl::Arg::B(qout), mtl::Arg::B(qidx), mtl::Arg::B(kq), mtl::Arg::B(vq), mtl::Arg::B(ksc), + mtl::Arg::B(vsc), mtl::Arg::B(tail), mtl::Arg::B(dead), mtl::Arg::B(pooled), mtl::Arg::V(a), + mtl::Arg::B(mtab)}; + d.grid[0] = 33; d.grid[1] = T; d.tg[0] = 256; + return d; + }; + mtl::Dispatch d_prep_f32 = make_prep(pprep_f32, qcur); + mtl::Dispatch d_qf16; d_qf16.p = &pq; d_qf16.args = {mtl::Arg::B(qcur), mtl::Arg::B(qh_ref), mtl::Arg::V(nq)}; + d_qf16.grid[0] = (nq + 255) / 256; d_qf16.tg[0] = 256; + mtl::Dispatch d_prep_f16 = make_prep(pprep_f16, qh_fused); + for (int t = 0; t < T; ++t) { + int* o = ids.as() + (size_t)t * QB; + for (int c = 0; c <= t && c < QB; ++c) o[c] = c; + } + mtl::Dispatch d_fl_ref; d_fl_ref.p = &pf; + d_fl_ref.args = {mtl::Arg::B(qh_ref), mtl::Arg::B(kq), mtl::Arg::B(vq), mtl::Arg::B(ksc), mtl::Arg::B(vsc), + mtl::Arg::B(ids), mtl::Arg::B(qfull), mtl::Arg::B(xf_ref), mtl::Arg::V(a)}; + d_fl_ref.grid[0] = QKV; d_fl_ref.grid[1] = T; d_fl_ref.tg[0] = 128; + mtl::Dispatch d_fl_fused = d_fl_ref; + d_fl_fused.args[0] = mtl::Arg::B(qh_fused); + d_fl_fused.args[7] = mtl::Arg::B(xf_fused); + if (bench) { + auto timeit = [&](const std::vector& ds) { + cx.warm(ds, 0.3); + std::vector t; + for (int i = 0; i < 5; ++i) t.push_back(cx.run(ds, 1, false, true)); + return median(t); + }; + std::printf("prep T %d: qsa_prep(f32)+qsa_q_f16 %.3f ms (prep %.3f ms, q_f16 %.3f ms), qsa_prep(qf16) %.3f ms\n", + T, timeit({d_prep_f32, d_qf16}) * 1e3, timeit({d_prep_f32}) * 1e3, timeit({d_qf16}) * 1e3, timeit({d_prep_f16}) * 1e3); + } else { + cx.run({d_prep_f32, d_qf16, d_fl_ref}, 1, false, true); + cx.run({d_prep_f16, d_fl_fused}, 1, false, true); + const bool ok_qh = std::memcmp(qh_ref.ptr, qh_fused.ptr, qh_ref.bytes) == 0; + const bool ok_xf = std::memcmp(xf_ref.ptr, xf_fused.ptr, xf_ref.bytes) == 0; + fails += (!ok_qh || !ok_xf); + std::printf("qsa_prep(QSA_Q_F16_FC) vs qsa_prep+qsa_q_f16: qcurh memcmp %s, qsa_flash memcmp %s\n", + ok_qh ? "PASS" : "FAIL", ok_xf ? "PASS" : "FAIL"); + } + } if (!bench) std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); return fails ? 1 : 0; } diff --git a/metal/tests/route_parity.cpp b/metal/tests/route_parity.cpp index 2f72329a..7e2f3eb8 100644 --- a/metal/tests/route_parity.cpp +++ b/metal/tests/route_parity.cpp @@ -1,5 +1,10 @@ // metal/tests/route_parity.cpp - moe_route (softmax + top-10 + the expert plan) against a CPU replica of -// native_router.cu's rule, for 1..8 tokens; logits with deliberate ties. +// native_router.cu's rule, for 1..8 tokens: +// Phase 1: 200 windows with n_expert = 512, pad = 0 (deliberate exact ties) +// Phase 2: 200 windows with n_expert = 256, pad = bit_cast(exp(0.85f)), random smap residency, and huge traps in 256..511 +// route_parity --router-ab BASE_KERNEL_DIR TYPE NTOK SPLIT bitwise fused-router A/B + short paired timing +// route_parity --mtp-ab BASE_MTP_METAL bitwise mtp.metal (mtp_rms_xprep, mtp_xprep, mtp_combine_gnorm) A/B + timing +// route_parity --ab ref.metal [--bench] bitwise parity (moe_route + moe_plan, kSplitRes 0/1) + paired A/B #include "mtl_rt.h" #include "ref_quant.h" @@ -7,27 +12,342 @@ #include #include #include +#include #include #include +#include +#include #include namespace { constexpr int P_NG = 0, P_MAXE = 128, P_GEXP = 16, P_GSTART = P_GEXP + P_MAXE, P_ETOK = P_GSTART + P_MAXE + 4, - P_EW = P_ETOK + P_MAXE, P_TOKENT = P_EW + P_MAXE, P_INTS = P_TOKENT + P_MAXE; + P_EW = P_ETOK + P_MAXE, P_TOKENT = P_EW + P_MAXE, P_INTS = P_TOKENT + P_MAXE + 20 + P_MAXE; struct MoeArgs { uint64_t gate_off, up_off, down_off, gu_stride, d_stride; uint32_t gu_row, d_row, n_embd, n_ff, n_tok, topk, n_expert, gx_gu, gx_dn, pad; }; +double median(std::vector v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; } +struct GemvArgs { uint32_t n_in, n_out, n_tok, row_bytes, x_stride, y_stride; }; + +// One specialization per invocation keeps the GPU checks short. +int router_ab(int argc, char** argv) { + if (argc != 6) return 2; + const int type = std::atoi(argv[3]), nt = std::atoi(argv[4]), split = std::atoi(argv[5]); + if ((type != ref::BF16 && type != ref::F32 && type != ref::Q8_0) || + nt < 1 || nt > 8 || split < 0 || split > 1) return 2; + mtl::Context cx; + auto compile = [&](const std::string& dir) { + return cx.compile({dir + "/quant.metal", dir + "/units.metal", dir + "/moe.metal", + dir + "/router.metal"}, + "#include \nusing namespace metal;\n" + "struct GemvArgs { uint n_in, n_out, n_tok, row_bytes, x_stride, y_stride; };\n" + "constant int NTOK [[function_constant(0)]];\n"); + }; + const int baseline = compile(argv[2]), candidate = compile(cx.kernel_dir); + mtl::Constants fc; + fc.i(0, nt).b(26, split != 0); + const std::string name = "router_route_" + std::to_string(type); + auto pb = cx.pipeline(baseline, name, fc), pc = cx.pipeline(candidate, name, fc); + constexpr int ni = 2560; + const int rb = int(ref::row_bytes(type, ni)); + auto W = cx.alloc(512 * rb), x = cx.alloc(4 * nt * ni), sums = cx.alloc(4 * nt * ni / 16); + auto smap = cx.alloc(4 * 512); + mtl::Buffer lg[2], ids[2], w[2], plan[2], counter[2]; + for (int j = 0; j < 2; ++j) { + lg[j] = cx.alloc(4 * nt * 512); ids[j] = cx.alloc(4 * nt * 10); + w[j] = cx.alloc(4 * nt * 10); plan[j] = cx.alloc(4 * P_INTS); counter[j] = cx.alloc(4); + } + ref::Rng rng(928); + mtl::Dispatch d[2]; + for (int trial = 0; trial < 4; ++trial) { + ref::random_blocks(type, W.as(), 512, ni, 928 + trial); + for (int i = 0; i < nt * ni; ++i) x.as()[i] = rng.sym(); + for (int e = 0; e < 512; ++e) smap.as()[e] = rng.sym() > 0 ? e : 0xFFFFFFFFu; + if (trial == 2) for (int e = 1; e < 512; ++e) + std::memcpy(W.as() + e * rb, W.ptr, rb); + MoeArgs a{}; + a.n_tok = nt; a.topk = 10; a.n_expert = trial == 1 ? 256 : 512; + a.gx_gu = 5; a.gx_dn = 20; + const float scale = std::exp(0.85f); + if (trial & 1) std::memcpy(&a.pad, &scale, 4); + if (trial == 3) a.pad |= 0x80000000u; + GemvArgs ga{ni, a.n_expert, uint32_t(nt), uint32_t(rb), ni, 512}; + for (int j = 0; j < 2; ++j) { + for (auto b : {lg[j], ids[j], w[j], plan[j]}) std::memset(b.ptr, 0xa5, b.bytes); + d[j].p = j ? &pc : &pb; + d[j].args = {mtl::Arg::B(W), mtl::Arg::B(x), mtl::Arg::B(sums), mtl::Arg::B(lg[j]), + mtl::Arg::V(ga), mtl::Arg::B(ids[j]), mtl::Arg::B(w[j]), mtl::Arg::B(plan[j]), + mtl::Arg::V(a), mtl::Arg::B(counter[j]), mtl::Arg::B(smap)}; + d[j].tg[0] = 512; d[j].grid[0] = a.n_expert / 16; + cx.run({d[j]}); + } + for (auto pair : {std::pair{lg[0], lg[1]}, std::pair{ids[0], ids[1]}, + std::pair{w[0], w[1]}, std::pair{plan[0], plan[1]}}) + if (std::memcmp(pair.first.ptr, pair.second.ptr, pair.first.bytes)) { + std::printf("router parity FAIL type %d nt %d split %d trial %d\n", type, nt, split, trial); + return 1; + } + if (*counter[0].as() || *counter[1].as()) return 1; + } + MoeArgs timing{}; + timing.n_tok = nt; timing.topk = 10; timing.n_expert = 512; + timing.gx_gu = 5; timing.gx_dn = 20; + for (auto& dispatch : d) dispatch.args[8] = mtl::Arg::V(timing); + d[1].args = d[0].args; + cx.warm({d[0], d[1]}, 0.05); + std::vector bt, ct, ratios; + constexpr int reps = 16; + for (int i = 0; i < 7; ++i) { + double t[2]; + for (int j = 0; j < 2; ++j) { + int k = (i + j) & 1; + t[k] = cx.run({d[k]}, reps) / reps; + } + bt.push_back(t[0] * 1e6); ct.push_back(t[1] * 1e6); ratios.push_back(t[1] / t[0]); + } + std::printf("router type %d nt %d split %d parity PASS base %.3f us cand %.3f us ratio %.5f\n", + type, nt, split, median(bt), median(ct), median(ratios)); + return 0; +} +int mtp_ab(int argc, char** argv) { + if (argc != 3) return 2; + struct RmsXpArgs { uint32_t n, xp_n, plane_tok, split, keep; float eps; }; + struct XpArgs { uint32_t n, n_tok, plane_tok, stride; }; + struct CombNormArgs { uint32_t n, hc; float eps; }; + constexpr uint32_t N = 2560, HC = 4, RN = HC * N, MAXT = 8; + constexpr float eps = 1e-6f; + mtl::Context cx; + const int lb = cx.compile({argv[2]}, "#define NO_GGML_COMMON 1\n"); + const int lc = cx.compile({cx.kernel_dir + "/../src/engine/msl/mtp.metal"}, "#define NO_GGML_COMMON 1\n"); + mtl::Pipeline prms[2] = {cx.pipeline(lb, "mtp_rms_xprep"), cx.pipeline(lc, "mtp_rms_xprep")}; + mtl::Pipeline pxp[2] = {cx.pipeline(lb, "mtp_xprep"), cx.pipeline(lc, "mtp_xprep")}; + mtl::Pipeline pcomb[2] = {cx.pipeline(lb, "mtp_combine_gnorm"), cx.pipeline(lc, "mtp_combine_gnorm")}; + const size_t plane = 4ull * MAXT * N; + auto x = cx.alloc(4 * MAXT * RN), w = cx.alloc(4 * RN), e = cx.alloc(4 * MAXT * N), h = cx.alloc(4 * MAXT * RN); + mtl::Buffer y[2], xp[2], R[2], xn[2]; + for (int k = 0; k < 2; ++k) { + y[k] = cx.alloc(4 * MAXT * RN); + xp[k] = cx.alloc(4 * plane + 4ull * MAXT * N / 16); + R[k] = cx.alloc(4 * MAXT * RN); + xn[k] = cx.alloc(4 * MAXT * RN); + } + ref::Rng rng(101); + for (size_t i = 0; i < MAXT * RN; ++i) { x.as()[i] = rng.sym() * 2.0f; h.as()[i] = rng.sym(); } + for (size_t i = 0; i < RN; ++i) w.as()[i] = 1.0f + 0.2f * rng.sym(); + for (size_t i = 0; i < MAXT * N; ++i) e.as()[i] = rng.sym(); + auto bench_pair = [&](const char* label, mtl::Dispatch db, mtl::Dispatch dc, int reps = 128) { + dc.args = db.args; + cx.warm({db, dc}, 0.01); + std::vector bt, ct, rt; + for (int i = 0; i < 11; ++i) { + double b, c; + if (i & 1) { c = cx.run({dc}, reps) / reps; b = cx.run({db}, reps) / reps; } + else { b = cx.run({db}, reps) / reps; c = cx.run({dc}, reps) / reps; } + bt.push_back(b * 1e6); ct.push_back(c * 1e6); rt.push_back(c / b); + } + std::printf("%-32s parity PASS base %6.3f us cand %6.3f us ratio %.4f (%+.2f%%)\n", + label, median(bt), median(ct), median(rt), 100.0 * (median(rt) - 1.0)); + }; + for (int nt : {1, 4, 8}) { + // 1. mtp_rms_xprep embedding (n=2560, split=0, keep=0) + { + RmsXpArgs a{N, N, MAXT, 0, 0, eps}; + mtl::Dispatch d[2]; + for (int k = 0; k < 2; ++k) { + std::memset(xp[k].ptr, 0xa5, xp[k].bytes); + d[k] = {&prms[k], {mtl::Arg::B(x), mtl::Arg::B(w), mtl::Arg::B(y[k]), mtl::Arg::B(xp[k]), mtl::Arg::V(a)}}; + d[k].grid[0] = nt; d[k].tg[0] = 256; + cx.run({d[k]}); + } + if (std::memcmp(xp[0].ptr, xp[1].ptr, xp[0].bytes)) return 1; + char lbl[64]; std::snprintf(lbl, sizeof(lbl), "mtp_rms_xprep n=2560 nt=%d", nt); + bench_pair(lbl, d[0], d[1]); + } + // 2. mtp_rms_xprep hidden (n=10240, split=1, keep=(nt*4 > 8)) + { + RmsXpArgs a{RN, N, MAXT, 1, uint32_t(nt * HC > MAXT), eps}; + mtl::Dispatch d[2]; + for (int k = 0; k < 2; ++k) { + std::memset(y[k].ptr, 0xa5, y[k].bytes); + std::memset(xp[k].ptr, 0xa5, xp[k].bytes); + d[k] = {&prms[k], {mtl::Arg::B(x), mtl::Arg::B(w), mtl::Arg::B(y[k]), mtl::Arg::B(xp[k]), mtl::Arg::V(a)}}; + d[k].grid[0] = nt; d[k].tg[0] = 256; + cx.run({d[k]}); + } + if (std::memcmp(y[0].ptr, y[1].ptr, y[0].bytes) || std::memcmp(xp[0].ptr, xp[1].ptr, xp[0].bytes)) return 1; + char lbl[64]; std::snprintf(lbl, sizeof(lbl), "mtp_rms_xprep n=10240 nt=%d", nt); + bench_pair(lbl, d[0], d[1]); + } + // 3. mtp_combine_gnorm (hc=4, nt) + { + CombNormArgs a{N, HC, eps}; + mtl::Dispatch d[2]; + for (int k = 0; k < 2; ++k) { + std::memset(R[k].ptr, 0xa5, R[k].bytes); + std::memset(xn[k].ptr, 0xa5, xn[k].bytes); + d[k] = {&pcomb[k], {mtl::Arg::B(h), mtl::Arg::B(e), mtl::Arg::B(R[k]), mtl::Arg::B(w), mtl::Arg::B(xn[k]), mtl::Arg::V(a)}}; + d[k].grid[0] = HC; d[k].grid[1] = nt; d[k].tg[0] = 256; + cx.run({d[k]}); + } + if (std::memcmp(R[0].ptr, R[1].ptr, R[0].bytes) || std::memcmp(xn[0].ptr, xn[1].ptr, xn[0].bytes)) return 1; + char lbl[64]; std::snprintf(lbl, sizeof(lbl), "mtp_combine_gnorm nt=%d", nt); + bench_pair(lbl, d[0], d[1]); + } + } + // 4. mtp_xprep (rows=8, n=2560) + { + XpArgs a{N, MAXT, MAXT, N}; + mtl::Dispatch d[2]; + for (int k = 0; k < 2; ++k) { + std::memset(xp[k].ptr, 0xa5, xp[k].bytes); + d[k] = {&pxp[k], {mtl::Arg::B(x), mtl::Arg::B(xp[k]), mtl::Arg::V(a)}}; + d[k].grid[0] = (N / 16 * MAXT + 255) / 256; d[k].tg[0] = 256; + cx.run({d[k]}); + } + if (std::memcmp(xp[0].ptr, xp[1].ptr, xp[0].bytes)) return 1; + bench_pair("mtp_xprep rows=8", d[0], d[1]); + } + return 0; +} } // namespace -int main() { +int main(int argc, char** argv) { + if (argc > 1 && std::string(argv[1]) == "--router-ab") return router_ab(argc, argv); + if (argc > 1 && std::string(argv[1]) == "--mtp-ab") return mtp_ab(argc, argv); + const bool ab = argc > 1 && std::string(argv[1]) == "--ab"; + if (ab && argc < 3) { + std::fprintf(stderr, "usage: route_parity --ab ref.metal [--bench]\n"); + return 2; + } + const bool bench = (argc > 1 && std::string(argv[1]) == "--bench") || + (ab && argc > 3 && std::string(argv[3]) == "--bench"); mtl::Context cx; const int lib = cx.compile({"quant.metal", "units.metal", "moe.metal"}); + const int ref_lib = ab ? cx.compile({"quant.metal", "units.metal", argv[2]}) : -1; mtl::Pipeline p = cx.pipeline(lib, "moe_route"); - mtl::Buffer lg = cx.alloc(4 * 8 * 512), ids = cx.alloc(4 * 80), w = cx.alloc(4 * 80), plan = cx.alloc(4 * P_INTS); + mtl::Buffer lg = cx.alloc(4 * 8 * 512), ids = cx.alloc(4 * 80), w = cx.alloc(4 * 80), plan = cx.alloc(4 * P_INTS), + smap = cx.alloc(4 * 512); ref::Rng r(9); - int fails = 0; - double worst_w = 0; + int fails_512 = 0, fails_256_moce = 0; + double worst_w_512 = 0, worst_w_256_moce = 0; + + if (ab) { + mtl::Buffer ids_r = cx.alloc(4 * 80), w_r = cx.alloc(4 * 80), plan_r = cx.alloc(4 * P_INTS); + int ab_fails = 0; + const float moce_scale = std::exp(0.85f); + uint32_t moce_pad = 0; + std::memcpy(&moce_pad, &moce_scale, 4); + for (int split : {0, 1}) { + mtl::Constants fc; + fc.b(26, split != 0); + mtl::Pipeline pr_c = cx.pipeline(lib, "moe_route", fc); + mtl::Pipeline pr_r = cx.pipeline(ref_lib, "moe_route", fc); + mtl::Pipeline pp_c = cx.pipeline(lib, "moe_plan", fc); + mtl::Pipeline pp_r = cx.pipeline(ref_lib, "moe_plan", fc); + for (int trial = 0; trial < 100; ++trial) { + const int nt = 1 + trial % 8; + const int ne_exp = (trial & 1) ? 256 : 512; + for (int e = 0; e < 512; ++e) smap.as()[e] = (r.sym() > 0.0f) ? (uint32_t)e : 0xFFFFFFFFu; + for (int i = 0; i < nt * 512; ++i) lg.as()[i] = r.sym() * 4.0f; + MoeArgs a{}; + a.n_tok = nt; + a.topk = 10; + a.n_expert = ne_exp; + a.gx_gu = 5; + a.gx_dn = 20; + a.pad = (trial & 1) ? moce_pad : 0; + + std::memset(ids_r.ptr, 0xa5, ids_r.bytes); + std::memset(w_r.ptr, 0xa5, w_r.bytes); + std::memset(plan_r.ptr, 0xa5, plan_r.bytes); + mtl::Dispatch dr{&pr_r, {mtl::Arg::B(lg), mtl::Arg::B(ids_r), mtl::Arg::B(w_r), mtl::Arg::B(plan_r), mtl::Arg::V(a), mtl::Arg::B(smap)}}; + dr.tg[0] = 512; + cx.run({dr}); + + std::memset(ids.ptr, 0xa5, ids.bytes); + std::memset(w.ptr, 0xa5, w.bytes); + std::memset(plan.ptr, 0xa5, plan.bytes); + mtl::Dispatch dc{&pr_c, {mtl::Arg::B(lg), mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(plan), mtl::Arg::V(a), mtl::Arg::B(smap)}}; + dc.tg[0] = 512; + cx.run({dc}); + + double max_dw = 0.0; + for (int i = 0; i < nt * 10; ++i) + max_dw = std::max(max_dw, (double) std::fabs(w.as()[i] - w_r.as()[i])); + if (std::memcmp(ids.ptr, ids_r.ptr, 4 * nt * 10) != 0 || max_dw > 1e-6) { + ++ab_fails; + } + + // Also verify standalone moe_plan bitwise parity + std::memset(plan_r.ptr, 0x5a, plan_r.bytes); + std::memset(plan.ptr, 0x5a, plan.bytes); + mtl::Dispatch dpr{&pp_r, {mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(plan_r), mtl::Arg::V(a)}}; + mtl::Dispatch dpc{&pp_c, {mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(plan), mtl::Arg::V(a)}}; + dpr.tg[0] = dpc.tg[0] = 512; + cx.run({dpr, dpc}); + if (std::memcmp(plan.ptr, plan_r.ptr, plan.bytes) != 0) ++ab_fails; + } + } + std::printf("route_parity --ab bitwise (moe_route + moe_plan, split 0/1, 200 trials): %d fail(s) %s\n", + ab_fails, ab_fails == 0 ? "ALL PASS" : "FAIL"); + if (ab_fails || !bench) return ab_fails ? 1 : 0; + + for (int split : {0, 1}) { + mtl::Constants fc; + fc.b(26, split != 0); + mtl::Pipeline pr_c = cx.pipeline(lib, "moe_route", fc); + mtl::Pipeline pr_r = cx.pipeline(ref_lib, "moe_route", fc); + mtl::Pipeline pp_c = cx.pipeline(lib, "moe_plan", fc); + mtl::Pipeline pp_r = cx.pipeline(ref_lib, "moe_plan", fc); + for (int nt : {1, 2, 4, 8}) { + for (int e = 0; e < 512; ++e) smap.as()[e] = (r.sym() > 0.0f) ? (uint32_t)e : 0xFFFFFFFFu; + for (int i = 0; i < nt * 512; ++i) lg.as()[i] = r.sym() * 4.0f; + MoeArgs a{}; + a.n_tok = nt; + a.topk = 10; + a.n_expert = 512; + a.gx_gu = 5; + a.gx_dn = 20; + a.pad = split ? moce_pad : 0; + mtl::Dispatch dr{&pr_r, {mtl::Arg::B(lg), mtl::Arg::B(ids_r), mtl::Arg::B(w_r), mtl::Arg::B(plan_r), mtl::Arg::V(a), mtl::Arg::B(smap)}}; + mtl::Dispatch dc{&pr_c, {mtl::Arg::B(lg), mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(plan), mtl::Arg::V(a), mtl::Arg::B(smap)}}; + dr.tg[0] = dc.tg[0] = 512; + cx.run({dc}); + cx.warm({dr}, 0.03); + cx.warm({dc}, 0.03); + std::vector bt, ct, rt; + const int reps = 192; // 4x48 layers + for (int i = 0; i < 9; ++i) { + double b, c; + if (i & 1) { c = cx.run({dc}, reps) / reps; b = cx.run({dr}, reps) / reps; } + else { b = cx.run({dr}, reps) / reps; c = cx.run({dc}, reps) / reps; } + bt.push_back(b * 1e6); ct.push_back(c * 1e6); rt.push_back(c / b); + } + std::printf("ab bench moe_route split %d ntok %d: base %6.2f us cand %6.2f us ratio %.4f (x48: %.3f -> %.3f ms)\n", + split, nt, median(bt), median(ct), median(rt), 48 * median(bt) * 1e-3, 48 * median(ct) * 1e-3); + + if (split == 0) { + mtl::Dispatch dpr{&pp_r, {mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(plan_r), mtl::Arg::V(a)}}; + mtl::Dispatch dpc{&pp_c, {mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(plan), mtl::Arg::V(a)}}; + dpr.tg[0] = dpc.tg[0] = 512; + cx.warm({dpr}, 0.03); + cx.warm({dpc}, 0.03); + bt.clear(); ct.clear(); rt.clear(); + for (int i = 0; i < 9; ++i) { + double b, c; + if (i & 1) { c = cx.run({dpc}, reps) / reps; b = cx.run({dpr}, reps) / reps; } + else { b = cx.run({dpr}, reps) / reps; c = cx.run({dpc}, reps) / reps; } + bt.push_back(b * 1e6); ct.push_back(c * 1e6); rt.push_back(c / b); + } + std::printf("ab bench moe_plan split %d ntok %d: base %6.2f us cand %6.2f us ratio %.4f (x48: %.3f -> %.3f ms)\n", + split, nt, median(bt), median(ct), median(rt), 48 * median(bt) * 1e-3, 48 * median(ct) * 1e-3); + } + } + } + return 0; + } + + // Phase 1: K=512, MoCE off (pad = 0) for (int trial = 0; trial < 200; ++trial) { const int nt = 1 + trial % 8; for (int i = 0; i < nt * 512; ++i) { @@ -41,9 +361,10 @@ int main() { a.n_expert = 512; a.gx_gu = 1; a.gx_dn = 1; + a.pad = 0; mtl::Dispatch d; d.p = &p; - d.args = {mtl::Arg::B(lg), mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(plan), mtl::Arg::V(a)}; + d.args = {mtl::Arg::B(lg), mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(plan), mtl::Arg::V(a), mtl::Arg::B(smap)}; d.tg[0] = 512; cx.run({d}); for (int t = 0; t < nt; ++t) { @@ -54,19 +375,17 @@ int main() { double s = 0; for (int e = 0; e < 512; ++e) s += (pr[e] = std::exp((double)l[e] - mx)); for (int e = 0; e < 512; ++e) pr[e] /= s; - // top 10: largest probability, ties -> lower index (probabilities of equal logits are equal) std::vector order(512); for (int e = 0; e < 512; ++e) order[e] = e; std::stable_sort(order.begin(), order.end(), [&](int x, int y) { return l[x] > l[y]; }); double ssum = 0; for (int k = 0; k < 10; ++k) ssum += pr[order[k]]; for (int k = 0; k < 10; ++k) { - if (ids.as()[t * 10 + k] != order[k]) { ++fails; break; } + if (ids.as()[t * 10 + k] != order[k]) { ++fails_512; break; } const double ref = pr[order[k]] / std::max(ssum, 6.103515625e-5); - worst_w = std::max(worst_w, std::fabs(w.as()[t * 10 + k] - ref) / ref); + worst_w_512 = std::max(worst_w_512, std::fabs(w.as()[t * 10 + k] - ref) / ref); } } - // the plan const int* P = plan.as(); std::set distinct(ids.as(), ids.as() + nt * 10); bool ok = P[P_NG] == (int)distinct.size(); @@ -82,9 +401,100 @@ int main() { if (j != P[P_GSTART + gi + 1]) ok = false; ++gi; } - fails += !ok; + fails_512 += !ok; + } + + // Phase 2: K=256, MoCE on (pad = bit_cast(exp(0.85f))), random smap residency, and huge trap logits in 256..511 + const float moce_scale = std::exp(0.85f); + uint32_t moce_pad = 0; + std::memcpy(&moce_pad, &moce_scale, 4); + for (int trial = 0; trial < 200; ++trial) { + const int nt = 1 + trial % 8; + for (int e = 0; e < 512; ++e) { + smap.as()[e] = (r.sym() > 0.0f) ? (uint32_t) e : 0xFFFFFFFFu; + } + for (int t = 0; t < nt; ++t) { + for (int e = 0; e < 256; ++e) { + float v = r.sym() * 4.0f; + if (trial % 3 == 0) v = std::round(v * 4.0f) / 4.0f; + lg.as()[t * 512 + e] = v; + } + for (int e = 256; e < 512; ++e) { + lg.as()[t * 512 + e] = 50.0f + (float) e; // trap logits past n_expert = 256 + } + } + MoeArgs a{}; + a.n_tok = nt; + a.topk = 10; + a.n_expert = 256; + a.gx_gu = 1; + a.gx_dn = 1; + a.pad = moce_pad; + mtl::Dispatch d; + d.p = &p; + d.args = {mtl::Arg::B(lg), mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(plan), mtl::Arg::V(a), mtl::Arg::B(smap)}; + d.tg[0] = 512; + cx.run({d}); + + const uint32_t* sm = smap.as(); + for (int t = 0; t < nt; ++t) { + const float* l = lg.as() + t * 512; + float mx = -INFINITY; + for (int e = 0; e < 256; ++e) mx = std::max(mx, l[e]); + std::vector pr(256); + float s = 0.0f; + for (int e = 0; e < 256; ++e) s += (pr[e] = std::exp(l[e] - mx)); + const float rcp = 1.0f / s; + for (int e = 0; e < 256; ++e) pr[e] *= rcp; + + std::vector work = pr; + std::vector picked(10, -1); + double ssum = 0.0; + for (int rank = 0; rank < 10; ++rank) { + if (rank == 3) { + for (int e = 0; e < 256; ++e) { + if (work[e] > 0.0f && sm[e] != 0xFFFFFFFFu) work[e] *= moce_scale; + } + } + int best_e = 0; + float best_v = work[0]; + for (int e = 1; e < 256; ++e) { + if (work[e] > best_v) { best_v = work[e]; best_e = e; } + } + picked[rank] = best_e; + work[best_e] = -INFINITY; + ssum += pr[best_e]; + } + for (int k = 0; k < 10; ++k) { + const int got_id = ids.as()[t * 10 + k]; + if (got_id != picked[k] || got_id >= 256) { ++fails_256_moce; break; } + const double ref = (double) pr[picked[k]] / std::max(ssum, 6.103515625e-5); + worst_w_256_moce = std::max(worst_w_256_moce, std::fabs(w.as()[t * 10 + k] - ref) / ref); + } + } + const int* P = plan.as(); + std::set distinct(ids.as(), ids.as() + nt * 10); + bool ok = P[P_NG] == (int)distinct.size(); + int gi = 0; + for (int e : distinct) { + if (e >= 256 || P[P_GEXP + gi] != e) ok = false; + int j = P[P_GSTART + gi]; + for (int i = 0; i < nt * 10; ++i) + if (ids.as()[i] == e) { + if (P[P_ETOK + j] != i / 10 || P[P_TOKENT + i] != j) ok = false; + ++j; + } + if (j != P[P_GSTART + gi + 1]) ok = false; + ++gi; + } + fails_256_moce += !ok; } - std::printf("route parity, 200 windows of 1..8 tokens (a third with exact ties): ids/plan mismatches %d, worst weight " - "rel err %.2e %s\n", fails, worst_w, fails == 0 && worst_w < 2e-6 ? "PASS" : "FAIL"); - return fails ? 1 : 0; + + const int total_fails = fails_512 + fails_256_moce; + const double worst_w = std::max(worst_w_512, worst_w_256_moce); + std::printf("route parity K=512 (pad=0, 200 windows): mismatches %d, worst weight rel err %.2e\n", + fails_512, worst_w_512); + std::printf("route parity K=256 + MoCE=0.85 (random smap + 256..511 traps, 200 windows): mismatches %d, worst weight rel err %.2e %s\n", + fails_256_moce, worst_w_256_moce, total_fails == 0 && worst_w < 2e-6 ? "PASS" : "FAIL"); + return total_fails == 0 && worst_w < 2e-6 ? 0 : 1; } diff --git a/metal/tests/sample_parity.cpp b/metal/tests/sample_parity.cpp index 049439af..2cb5b4a4 100644 --- a/metal/tests/sample_parity.cpp +++ b/metal/tests/sample_parity.cpp @@ -1,61 +1,260 @@ -// metal/tests/sample_parity.cpp - greedy argmax (sample.metal) vs the CPU over 248,320 logits with planted ties, -// and its time. +// metal/tests/sample_parity.cpp - exact sampling parity and short paired GPU timings. +// sample_parity [--mtp] [baseline-source.metal]; without a baseline, compare the kernels to themselves. #include "mtl_rt.h" #include "ref_quant.h" +#include "../src/engine/slop4.h" #include #include #include +#include +#include #include -int main() { - constexpr uint32_t V = 248320, PARTS = 120; - mtl::Context cx; - const int lib = cx.compile({"sample.metal"}, "#define NO_GGML_COMMON 1\n"); - mtl::Pipeline pp = cx.pipeline(lib, "argmax_part"), pf = cx.pipeline(lib, "argmax_final"); - const int MAXT = 8; - mtl::Buffer lg = cx.alloc(4 * (size_t)V * MAXT), pv = cx.alloc(4 * PARTS * MAXT), pi = cx.alloc(4 * PARTS * MAXT), - out = cx.alloc(4 * MAXT); +struct ArgArgs { uint32_t n_vocab, parts; }; +struct PickArgs { uint32_t nv, j, use_map; }; +struct DraftArgs { slop4::State state; slop4::Knobs knobs; uint32_t nv, j, vocab; }; +struct EmbArgs { uint32_t n, row_bytes, j; }; + +static void timing(mtl::Context& cx, const char* name, const std::vector& base, + const std::vector& cand, int reps = 64) { + cx.warm(base, 0.005); + cx.warm(cand, 0.005); + std::vector b, c, r; + for (int k = 0; k < 9; ++k) { + double bt, ct; + if (k & 1) { ct = cx.run(cand, reps); bt = cx.run(base, reps); } + else { bt = cx.run(base, reps); ct = cx.run(cand, reps); } + b.push_back(bt / reps * 1e6); + c.push_back(ct / reps * 1e6); + r.push_back(ct / bt); + } + std::sort(b.begin(), b.end()); + std::sort(c.begin(), c.end()); + std::sort(r.begin(), r.end()); + std::printf("%s baseline %.3f us candidate %.3f us ratio %.4f (%+.2f%%)\n", + name, b[4], c[4], r[4], 100.0 * (r[4] - 1.0)); +} + +static int argmax(mtl::Context& cx, const char* baseline) { + const int lb = cx.compile({baseline}, "#define NO_GGML_COMMON 1\n"); + const int lc = cx.compile({"sample.metal"}, "#define NO_GGML_COMMON 1\n"); + mtl::Pipeline pp[2] = {cx.pipeline(lb, "argmax_part"), cx.pipeline(lc, "argmax_part")}; + mtl::Pipeline pf[2] = {cx.pipeline(lb, "argmax_final"), cx.pipeline(lc, "argmax_final")}; + constexpr uint32_t V = 248320, PARTS = 120, MAXT = 8; + auto lg = cx.alloc(4 * size_t(V) * MAXT); + mtl::Buffer pv[2], pi[2], out[2]; + for (int k = 0; k < 2; ++k) { + pv[k] = cx.alloc(4 * PARTS * MAXT); + pi[k] = cx.alloc(4 * PARTS * MAXT); + out[k] = cx.alloc(4 * MAXT); + } ref::Rng r(5); int fails = 0; - struct A { uint32_t n_vocab, parts; } a{V, PARTS}; - for (int trial = 0; trial < 20; ++trial) { - const int nt = 1 + trial % MAXT; - for (size_t i = 0; i < (size_t)V * nt; ++i) lg.as()[i] = std::round(r.sym() * 1000.0f) / 100.0f; - for (int t = 0; t < nt; ++t) { - // plant a tie at the maximum - const size_t i1 = r.next() % V, i2 = r.next() % V; - lg.as()[t * V + i1] = 20.0f; - lg.as()[t * V + i2] = 20.0f; + struct Case { uint32_t nv, nt; }; + for (auto shape : {Case{V, 1}, Case{V, 2}, Case{V, 3}, Case{V, 4}, Case{V, 5}, + Case{V, 6}, Case{V, 7}, Case{V, 8}, + Case{248064, 1}, Case{248064, 2}, Case{248064, 3}, Case{248064, 4}, + Case{248064, 5}, Case{248064, 6}, Case{248064, 7}, Case{248064, 8}, + Case{40525, 1}, Case{40525, 2}, Case{40525, 4}, Case{40525, 8}, Case{7, 1}}) { + const ArgArgs a{shape.nv, PARTS}; + const uint32_t per = (a.n_vocab + PARTS - 1) / PARTS; + std::vector ds[2]; + for (int k = 0; k < 2; ++k) { + mtl::Dispatch d1, d2; + d1.p = &pp[k]; + d1.args = {mtl::Arg::B(lg), mtl::Arg::B(pv[k]), mtl::Arg::B(pi[k]), mtl::Arg::V(a)}; + d1.grid[0] = PARTS; + d1.grid[1] = shape.nt; + d1.tg[0] = 256; + d2.p = &pf[k]; + d2.args = {mtl::Arg::B(pv[k]), mtl::Arg::B(pi[k]), mtl::Arg::B(out[k]), mtl::Arg::V(a)}; + d2.grid[0] = shape.nt; + d2.tg[0] = 256; + ds[k] = {d1, d2}; } - mtl::Dispatch d1, d2; - d1.p = &pp; - d1.args = {mtl::Arg::B(lg), mtl::Arg::B(pv), mtl::Arg::B(pi), mtl::Arg::V(a)}; - d1.grid[0] = PARTS; - d1.grid[1] = nt; - d1.tg[0] = 256; - d2.p = &pf; - d2.args = {mtl::Arg::B(pv), mtl::Arg::B(pi), mtl::Arg::B(out), mtl::Arg::V(a)}; - d2.grid[0] = nt; - d2.tg[0] = 256; - cx.run({d1, d2}); - for (int t = 0; t < nt; ++t) { - const float* l = lg.as() + (size_t)t * V; - int best = 0; - for (uint32_t i = 1; i < V; ++i) - if (l[i] > l[best]) best = (int)i; - fails += out.as()[t] != best; + for (int trial = 0; trial < 7; ++trial) { + for (size_t i = 0; i < size_t(a.n_vocab) * shape.nt; ++i) { + float v = std::round(r.sym() * 1000.0f) / 100.0f; + if (trial == 1) v = -INFINITY; + if (trial == 2) v = std::numeric_limits::quiet_NaN(); + if (trial == 3) v = (i & 1) ? -0.0f : 0.0f; + if (trial == 4 && i % 17 == 0) v = std::numeric_limits::quiet_NaN(); + lg.as()[i] = v; + } + if (trial == 0 || trial >= 5) { + for (uint32_t t = 0; t < shape.nt; ++t) { + float* l = lg.as() + size_t(t) * a.n_vocab; + l[1] = l[a.n_vocab - 1] = trial == 5 ? INFINITY : 20.0f; + } + } + for (int k = 0; k < 2; ++k) cx.run(ds[k]); + fails += std::memcmp(pv[0].ptr, pv[1].ptr, 4 * PARTS * shape.nt) != 0; + fails += std::memcmp(pi[0].ptr, pi[1].ptr, 4 * PARTS * shape.nt) != 0; + fails += std::memcmp(out[0].ptr, out[1].ptr, 4 * shape.nt) != 0; + for (uint32_t t = 0; t < shape.nt; ++t) { + float best = -INFINITY; + int id = 0x7fffffff; + for (uint32_t p = 0; p < PARTS; ++p) { + float v = -INFINITY; + int ix = 0x7fffffff; + for (uint32_t i = p * per; i < std::min((p + 1) * per, a.n_vocab); ++i) { + const float x = lg.as()[size_t(t) * a.n_vocab + i]; + if (x > v || (x == v && int(i) < ix)) { v = x; ix = int(i); } + } + const size_t off = t * PARTS + p; + fails += std::memcmp(&v, pv[1].as() + off, 4) != 0; + fails += ix != pi[1].as()[off]; + if (v > best || (v == best && ix < id)) { best = v; id = ix; } + } + fails += out[1].as()[t] != id; + } } - if (trial == 19) { - std::vector ds; - for (int k = 0; k < 64; ++k) { ds.push_back(d1); ds.push_back(d2); } - cx.warm(ds, 0.05); - std::vector ts; - for (int k = 0; k < 7; ++k) ts.push_back(cx.run(ds) / 64); - std::sort(ts.begin(), ts.end()); - std::printf("argmax %d token(s) over %u logits: %.2f us (2 dispatches)\n", nt, V, ts[3] * 1e6); + for (size_t i = 0; i < size_t(a.n_vocab) * shape.nt; ++i) + lg.as()[i] = std::round(r.sym() * 1000.0f) / 100.0f; + char label[80]; + std::snprintf(label, sizeof(label), "argmax nv=%u nt=%u", a.n_vocab, shape.nt); + timing(cx, label, ds[0], ds[1]); + } + std::printf("argmax pv/pi/out bitwise + cpu: %d wrong %s\n", fails, fails ? "FAIL" : "PASS"); + return fails; +} + +static int mtp(mtl::Context& cx, const char* baseline) { + const std::string engine = cx.kernel_dir + "/../src/engine/"; + const int lb = cx.compile({"quant.metal", "units.metal", engine + "slop4.h", baseline}); + const int lc = cx.compile({"quant.metal", "units.metal", engine + "slop4.h", engine + "msl/mtp_chain.metal"}); + mtl::Pipeline pick[2] = {cx.pipeline(lb, "mtp_pick"), cx.pipeline(lc, "mtp_pick")}; + mtl::Pipeline pen[2] = {cx.pipeline(lb, "mtp_slop_pen"), cx.pipeline(lc, "mtp_slop_pen")}; + constexpr uint32_t V = 248320; + auto lg = cx.alloc(4 * V), ix = cx.alloc(4), map = cx.alloc(4 * V); + mtl::Buffer dout[2], prob[2], state[2], penalized[2]; + for (int k = 0; k < 2; ++k) { + dout[k] = cx.alloc(16); + prob[k] = cx.alloc(16); + state[k] = cx.alloc(sizeof(slop4::State)); + penalized[k] = cx.alloc(4 * V); + } + auto pen_dispatch = [&](int k, const DraftArgs& a) { + mtl::Dispatch d; + d.p = &pen[k]; + d.args = {mtl::Arg::B(penalized[k]), mtl::Arg::B(map), mtl::Arg::B(dout[k]), + mtl::Arg::B(state[k]), mtl::Arg::V(a)}; + d.tg[0] = 256; + return d; + }; + ref::Rng r(19); + int fails = 0; + for (uint32_t nv : {40525u, V}) { + for (uint32_t i = 0; i < V; ++i) map.as()[i] = V - i - 1; + PickArgs a{nv, 2, 1}; + std::vector ds[2]; + for (int k = 0; k < 2; ++k) { + mtl::Dispatch d; + d.p = &pick[k]; + d.args = {mtl::Arg::B(lg), mtl::Arg::B(ix), mtl::Arg::B(map), mtl::Arg::B(dout[k]), + mtl::Arg::B(prob[k]), mtl::Arg::V(a)}; + d.tg[0] = 1024; + ds[k] = {d}; + } + for (int trial = 0; trial < 8; ++trial) { + ix.as()[0] = trial; + for (uint32_t i = 0; i < nv; ++i) lg.as()[i] = r.sym() * 10.0f; + if (trial == 6) lg.as()[nv - 1] = -INFINITY; + if (trial == 7) lg.as()[nv - 1] = std::numeric_limits::quiet_NaN(); + a.j = trial % 4; + a.use_map = trial & 1; + for (int k = 0; k < 2; ++k) { + ds[k][0].args[5] = mtl::Arg::V(a); + cx.run(ds[k]); + } + const int expected = a.use_map ? map.as()[trial] : trial; + fails += dout[1].as()[a.j] != expected; + fails += std::memcmp(dout[0].ptr, dout[1].ptr, 16) != 0; + fails += std::memcmp(prob[0].ptr, prob[1].ptr, 16) != 0; + } + for (uint32_t i = 0; i < nv; ++i) lg.as()[i] = r.sym() * 10.0f; + timing(cx, nv == V ? "mtp_pick nv=248320" : "mtp_pick nv=40525", ds[0], ds[1]); + for (int trial = 0; trial < 10; ++trial) { + DraftArgs a{{}, {1.0f, 1.0f, 1.0f, 1.0f}, nv, 0, V}; + a.state.f = slop4::KNOWN | slop4::PROSE | slop4::BOL; + a.state.t1 = 1602; + if (trial == 0) a.knobs = {}; + if (trial >= 4) { a.j = 1; a.state.t1 = slop4::NL; } + for (uint32_t i = 0; i < V; ++i) map.as()[i] = i < nv ? int(i) : -1; + for (uint32_t i = 0; i < nv; ++i) lg.as()[i] = r.sym() * 10.0f; + if (trial == 6) std::fill_n(lg.as(), nv, -INFINITY); + if (trial == 7) std::fill_n(lg.as(), nv, std::numeric_limits::quiet_NaN()); + if (trial == 8) lg.as()[100] = INFINITY; + if (trial == 9) std::fill_n(lg.as(), nv, 0.0f); + if (trial == 2) lg.as()[100] = 100.0f; + if (trial == 3) lg.as()[100] = std::numeric_limits::quiet_NaN(); + for (int k = 0; k < 2; ++k) { + *state[k].as() = a.state; + dout[k].as()[0] = slop4::NL; + std::memcpy(penalized[k].ptr, lg.ptr, 4 * nv); + ds[k] = {pen_dispatch(k, a)}; + cx.run(ds[k]); + } + fails += std::memcmp(penalized[0].ptr, penalized[1].ptr, 4 * nv) != 0; + fails += std::memcmp(state[0].ptr, state[1].ptr, sizeof(slop4::State)) != 0; + if (trial == 1 || trial == 2) + fails += std::memcmp(lg.ptr, penalized[1].ptr, 4 * nv) == 0; + } + // Small penalties keep the hold branch stable across timed repetitions. + for (int mode = 0; mode < 3; ++mode) { + DraftArgs a{{}, {0.0f, mode ? 1e-6f : 0.0f, 0.0f, 0.0f}, nv, 0, V}; + a.state.f = slop4::KNOWN | slop4::PROSE; + a.state.t1 = 1602; + for (int k = 0; k < 2; ++k) { + std::fill_n(penalized[k].as(), nv, 0.0f); + if (mode == 1) penalized[k].as()[100] = 100.0f; + ds[k] = {pen_dispatch(k, a)}; + } + char label[80]; + std::snprintf(label, sizeof(label), "mtp_slop_pen nv=%u mode=%d", nv, mode); + timing(cx, label, ds[0], ds[1]); } } - std::printf("argmax parity, 20 windows with planted ties: %d wrong %s\n", fails, fails ? "FAIL" : "PASS"); - return fails ? 1 : 0; + for (int format : {ref::BF16, ref::Q8_0, ref::Q4_0, ref::Q5_0, ref::IQ4_NL, ref::IQ4_XS, ref::Q2_0, + ref::IQ3_S, ref::Q3_K, ref::Q4_K, ref::Q5_K, ref::Q6_K, ref::IQ3_XXS}) { + const int n = 2560, rows = 3; + const int rb = int(ref::row_bytes(format, n)); + auto w = cx.alloc(size_t(rb) * rows), pos = cx.alloc(4 * n); + auto y0 = cx.alloc(4 * n), y1 = cx.alloc(4 * n); + for (int i = 0; i < n; ++i) pos.as()[i] = float(i); + ref::random_blocks(format, w.as(), rows, n, 31 + format); + mtl::Pipeline emb[2] = {cx.pipeline(lb, "mtp_embed_" + std::to_string(format)), + cx.pipeline(lc, "mtp_embed_" + std::to_string(format))}; + std::vector ds[2]; + EmbArgs a{uint32_t(n), uint32_t(rb), 1}; + for (int trial = 0; trial < rows; ++trial) { + for (int k = 0; k < 2; ++k) { + dout[k].as()[1] = trial; + mtl::Dispatch d; + d.p = &emb[k]; + d.args = {mtl::Arg::B(w), mtl::Arg::B(dout[k]), mtl::Arg::B(pos), + mtl::Arg::B(k ? y1 : y0), mtl::Arg::V(a)}; + d.tg[0] = 256; + ds[k] = {d}; + cx.run(ds[k]); + } + fails += std::memcmp(y0.ptr, y1.ptr, 4 * n) != 0; + } + char label[40]; + std::snprintf(label, sizeof(label), "mtp_embed format=%d", format); + timing(cx, label, ds[0], ds[1], 1024); + } + std::printf("mtp pick/pen/embed bitwise: %d wrong %s\n", fails, fails ? "FAIL" : "PASS"); + return fails; +} + +int main(int argc, char** argv) { + mtl::Context cx; + const bool chain = argc > 1 && std::strcmp(argv[1], "--mtp") == 0; + const char* baseline = argc > (chain ? 2 : 1) ? argv[chain ? 2 : 1] : nullptr; + const std::string self = cx.kernel_dir + "/../src/engine/msl/mtp_chain.metal"; + std::printf("device %s; baseline %s\n", cx.device_name().c_str(), baseline ? baseline : "self"); + return (chain ? mtp(cx, baseline ? baseline : self.c_str()) + : argmax(cx, baseline ? baseline : "sample.metal")) ? 1 : 0; } diff --git a/metal/tests/spec_min_p.cpp b/metal/tests/spec_min_p.cpp new file mode 100644 index 00000000..ec5ef0c6 --- /dev/null +++ b/metal/tests/spec_min_p.cpp @@ -0,0 +1,81 @@ +#include "../src/engine/spec_min_p.hpp" + +#include +#include +#include + +#define CHECK(condition) do { \ + if (!(condition)) { \ + std::fprintf(stderr, "spec_min_p:%d: %s\n", __LINE__, #condition); \ + return 1; \ + } \ +} while (false) + +using strata::metal_engine::DraftMinP; + +int main() { + DraftMinP policy; + CHECK(policy.parse(nullptr, nullptr) && !policy.enabled()); + CHECK(policy.parse(" 0.1 , .35, .55 ", nullptr)); + CHECK(policy.count == 3 && policy.enabled()); + const char* bad_schedules[] = { + "", " ", ",", "0,", "0,,1", "-0.1", "1.01", "nan", "inf", "1x", "1;0", + "0,0,0,0,0,0,0,0,0", "1e-999" + }; + for (const char* bad : bad_schedules) { + DraftMinP invalid; + CHECK(!invalid.parse(bad, nullptr) && !invalid.enabled()); + } + for (const char* bad : {"", "nan", "inf", "-1", "1.01", "0.1,0.2", "0.1junk"}) { + DraftMinP invalid; + CHECK(!invalid.parse("0.1,0.2", bad) && !invalid.enabled()); + } + CHECK(policy.parse("0,1,0,1,0,1,0,1", ".25")); + double cumulative = 1; + CHECK(policy.keep(0, .5f, .35f, cumulative)); + CHECK(!policy.keep(1, .5f, .35f, cumulative)); + CHECK(policy.parse("0,0,0", ".25")); + cumulative = 1; + CHECK(policy.keep(0, .5f, .35f, cumulative)); + CHECK(policy.keep(1, .5f, .35f, cumulative)); + CHECK(!policy.keep(2, .5f, .35f, cumulative)); + CHECK(policy.parse(".1,.2,.8", nullptr)); + cumulative = 1; + CHECK(policy.keep(0, .1f, .35f, cumulative)); + CHECK(policy.keep(1, .2f, .35f, cumulative)); + CHECK(policy.keep(2, .8f, .35f, cumulative)); + CHECK(!policy.keep(3, .7f, .35f, cumulative)); + const float bad_probs[] = { + -1.f, 2.f, std::numeric_limits::quiet_NaN(), std::numeric_limits::infinity() + }; + for (float bad : bad_probs) { + cumulative = 1; + CHECK(!policy.keep(0, bad, 0, cumulative)); + } + std::mt19937 rng(12345); + int checks = 0; + for (int trial = 0; trial < 100000; ++trial) { + float probs[8]; + for (float& v : probs) v = float(rng() % 101) / 100; + CHECK(policy.parse(nullptr, nullptr)); + cumulative = 1; + for (int j = 0; j < 8; ++j) { + CHECK(policy.keep(j, probs[j], .35f, cumulative) == !(probs[j] < .35f)); + ++checks; + } + for (bool cumulative_mode : {false, true}) { + CHECK(policy.parse(".1,.35,.55", cumulative_mode ? ".2" : nullptr)); + cumulative = 1; + double expected = 1; + for (int j = 0; j < 8; ++j) { + expected *= probs[j]; + const float gate = j == 0 ? .1f : j == 1 ? .35f : .55f; + const bool keep = probs[j] >= gate && (!cumulative_mode || expected >= .2f); + CHECK(policy.keep(j, probs[j], .35f, cumulative) == keep); + ++checks; + if (!keep) break; + } + } + } + std::printf("schedule_checks=%d parser/boundary/invalid-observation checks=pass\n", checks); +} diff --git a/metal/tests/token_chain.cpp b/metal/tests/token_chain.cpp index a47c7630..a0add9b2 100644 --- a/metal/tests/token_chain.cpp +++ b/metal/tests/token_chain.cpp @@ -24,7 +24,7 @@ namespace { constexpr int N = 2560, HC = 4, LR = 320, NFF = 640, TOPK = 10, NEXP = 512, MAXT = 8; constexpr int P_MAXE = 128, P_GEXP = 16, P_GSTART = P_GEXP + P_MAXE, P_ETOK = P_GSTART + P_MAXE + 4, - P_EW = P_ETOK + P_MAXE, P_TOKENT = P_EW + P_MAXE, P_INTS = P_TOKENT + P_MAXE, P_IGU = 4, P_IDN = 8; + P_EW = P_ETOK + P_MAXE, P_TOKENT = P_EW + P_MAXE, P_INTS = P_TOKENT + P_MAXE + 20 + P_MAXE, P_IGU = 4, P_IDN = 8; struct GemvArgs { uint32_t n_in, n_out, n_tok, row_bytes, x_stride, y_stride; }; struct GrArgs { uint32_t n_embd, hc, hc_lr, n_tok, apply; float eps; }; struct GrArgs3 { uint32_t n_embd, hc, hc_lr, n_tok, apply; float eps; uint32_t rb_down, rb_up, comb, topk; }; @@ -552,7 +552,7 @@ int main(int argc, char** argv) { GemvArgs ga{(uint32_t)N, 512, (uint32_t)ntok, (uint32_t)ref::row_bytes(tr.type, N), (uint32_t)N, 512}; r.args = {mtl::Arg::B(W, tr.offset), mtl::Arg::B(XF.b, 0), mtl::Arg::B(XF.b, XF.plane_bytes() * 4), mtl::Arg::B(logits), mtl::Arg::V(ga), mtl::Arg::B(ids), mtl::Arg::B(wts), mtl::Arg::B(plan), - mtl::Arg::V(a), mtl::Arg::B(rcounter)}; + mtl::Arg::V(a), mtl::Arg::B(rcounter), mtl::Arg::B(rcounter)}; r.grid[0] = 512 / 16; r.tg[0] = 512; nb = true; @@ -560,7 +560,7 @@ int main(int argc, char** argv) { } else { mtl::Dispatch r; r.p = proute; - r.args = {mtl::Arg::B(logits), mtl::Arg::B(ids), mtl::Arg::B(wts), mtl::Arg::B(plan), mtl::Arg::V(a)}; + r.args = {mtl::Arg::B(logits), mtl::Arg::B(ids), mtl::Arg::B(wts), mtl::Arg::B(plan), mtl::Arg::V(a), mtl::Arg::B(plan)}; r.tg[0] = 512; push(r, "router", 0); } diff --git a/serve/frontend.py b/serve/frontend.py index c36a44da..3fda9c56 100644 --- a/serve/frontend.py +++ b/serve/frontend.py @@ -637,6 +637,11 @@ def anthropic_to_messages(req: dict, think_unasked: bool = True) -> tuple[list[d # "output_config": {"effort": "low" | "medium" | "high"} thinking = req.get("thinking") effort = (req.get("output_config") or {}).get("effort") if isinstance(req.get("output_config"), dict) else None + # 0.1.45: Anthropic's budget_tokens is a real budget, not only a level: the thinking is closed (wrap-up, , + # the answer) after that many tokens. A reasoning_budget_tokens of the request's own wins. + bt = thinking.get("budget_tokens") if isinstance(thinking, dict) and thinking.get("type") != "disabled" else None + if isinstance(bt, int) and not isinstance(bt, bool) and bt > 0 and req.get("reasoning_budget_tokens") is None: + req["reasoning_budget_tokens"] = bt if isinstance(thinking, dict) and thinking.get("type") == "disabled": kwargs["enable_thinking"] = False elif effort: @@ -1039,6 +1044,7 @@ class OutputParser: return out def _reset_scan(self): + self._call_end_checked = 0 # closing-tag search progress for the current wrapped call self.sp = 0 # how much of self.buf (the call body) the scanner has consumed self.ss = "name" # name -> between -> str|raw -> ... -> done self.scall = None # the ToolCall being streamed (its id is reused by the final event) @@ -1046,6 +1052,20 @@ class OutputParser: self.sval_started = False self.sdeclared = {} + def _wrapped_call_end(self) -> int: + # Without a closing tag the structural walk cannot finish. Search only + # newly appended text plus enough overlap for a tag split across feeds. + # A literal closing tag inside a value needs structural disambiguation: + # once any candidate appears, keep the original full walk for this call. + # Recovery can rewrite the buffer, so it always uses the full walk. + if not self.recover and self._call_end_checked >= 0: + start = max(0, self._call_end_checked - len(CALL_END) + 1) + if self.buf.find(CALL_END, start) < 0: + self._call_end_checked = len(self.buf) + return -1 + self._call_end_checked = -1 + return call_end(self.buf) + def _scan(self) -> list[Event]: """Advance the streaming view of the call body in self.buf (see stream_tools).""" out = [] @@ -1152,10 +1172,19 @@ class OutputParser: def _hold(self, text: str, tags: tuple[str, ...]) -> int: """Length of the longest suffix of `text` that is a proper prefix of one of `tags`.""" best = 0 + size = len(text) for tag in tags: - for n in range(1, len(tag)): - if text.endswith(tag[:n]): - best = max(best, n) + if len(tag) <= 1: + continue + # Only a suffix beginning with the tag's first character can match. + # Ordinary token deltas contain no '<': one C-level search replaces + # checking every possible prefix. Earliest candidate is the longest. + start = text.find(tag[0], max(0, size - len(tag) + 1)) + while start >= 0 and size - start > best: + if tag.startswith(text[start:]): + best = size - start + break + start = text.find(tag[0], start + 1) return best def feed(self, delta: str) -> list[Event]: @@ -1351,7 +1380,7 @@ class OutputParser: elif after and not CALL_END.startswith(after): drop = 0 else: - i = call_end(self.buf) + i = self._wrapped_call_end() if i >= 0: drop = len(CALL_END) if self.recover: # a batch in one wrapper: the next call follows this one's diff --git a/serve/logit_bias.py b/serve/logit_bias.py index 59873a77..700511c0 100644 --- a/serve/logit_bias.py +++ b/serve/logit_bias.py @@ -1,9 +1,12 @@ """Request-scoped token biases. Normalize before HTTP streaming starts.""" +import json import math import re +import threading +from collections import OrderedDict -def normalize(value, vocab_size=None): +def _normalize(value, vocab_size=None): if value is None: return {} if isinstance(value, dict): @@ -36,9 +39,52 @@ def normalize(value, vocab_size=None): return out +# A long list (a standing ban of a whole script is tens of thousands of IDs) is sent with every request and used to be +# normalized five times per request: in the OpenAI route, twice in validate_request and twice in engine_key (sampling_keys +# runs for the cache key and again for the GEN line), about 170 ms for 65,929 entries. JSON text is an exact cache key +# that tells False from 0 and "1" from 1, and catches in-place edits to the same request object. +# Only a list that normalized without an error is stored, so a bad list still raises on every call. Short lists skip the +# cache: they cost less than the key. +_CACHE_MIN = 64 +_CACHE_MAX = 4 +_cache = OrderedDict() # (vocab_size, JSON text) -> [normalized dict, engine key or None] +_lock = threading.Lock() + + +def _entry(value, vocab_size): + if not isinstance(value, (dict, list)) or len(value) < _CACHE_MIN: + return [_normalize(value, vocab_size), None] + try: + key = (vocab_size, json.dumps(value, separators=(",", ":"))) + except (TypeError, ValueError): + key = None # something JSON cannot write: _normalize raises for it + entry = None + if key is not None: + with _lock: + entry = _cache.get(key) + if entry is not None: + _cache.move_to_end(key) + if entry is None: + entry = [_normalize(value, vocab_size), None] + if key is not None: + with _lock: + _cache[key] = entry + while len(_cache) > _CACHE_MAX: + _cache.popitem(last=False) + return entry + + +def normalize(value, vocab_size=None): + return dict(_entry(value, vocab_size)[0]) + + def engine_key(value): - entries = normalize(value) - return " logit_bias=" + ",".join(f"{i}:{b:g}" for i, b in sorted(entries.items())) if entries else "" + entry = _entry(value, None) + if not entry[0]: + return "" + if entry[1] is None: + entry[1] = " logit_bias=" + ",".join(f"{i}:{b:g}" for i, b in sorted(entry[0].items())) + return entry[1] def validate_request(req, engine, vocab_size): diff --git a/serve/runconfig.py b/serve/runconfig.py index 472b0021..842554ab 100644 --- a/serve/runconfig.py +++ b/serve/runconfig.py @@ -18,7 +18,7 @@ from pathlib import Path # of the "sampling" block, ("arg", flag) for an engine option kept in "args". EDITABLE = [ ("sampling.temperature", ("sampling", "num>=0"), - "Default temperature for requests that send none (0 = greedy, the default without a sampling block)"), + "Default temperature for requests that send none (0 = greedy; without it the model card's: 1.0 thinking, 0.7 not)"), ("sampling.top_p", ("sampling", "0, with no answer @@ -254,6 +256,15 @@ def _timeout_env(name: str, default: float) -> float | None: # still working is never ended by this, only reported once (the recommendation: engine_silence_s, or this variable). # STRATA_ENGINE_STALL_S overrides it (0 = off). Needs psutil (setup installs it); without it nothing is ended. ENGINE_STALL_S = _timeout_env("STRATA_ENGINE_STALL_S", 90.0) +# --batch: new requests take the control lines (one prompt read at a time) short prompts first: a prompt of at most +# STRATA_ADMIT_SHORT_FIRST tokens goes before longer ones (opt-in; unset or 0: whoever gets the lock first), and a +# longer one that waited STRATA_ADMIT_MAX_WAIT_S (default 30) counts as short. A chat question then waits for the one +# read in progress, not for every agent read that came before it. +try: + ADMIT_SHORT_FIRST = max(0, int(os.environ.get("STRATA_ADMIT_SHORT_FIRST") or 0)) +except ValueError: + ADMIT_SHORT_FIRST = 0 +ADMIT_MAX_WAIT_S = _timeout_env("STRATA_ADMIT_MAX_WAIT_S", 30.0) or 0.0 STALL_CPU_EPS_S = 0.5 # CPU seconds the engine may use in the whole stall window and still count as frozen STALL_IO_EPS_B = 1 << 20 # bytes read or written in that window @@ -705,6 +716,19 @@ class ConvCacheLog: return dict(self.state) +def conversation_slots(engine) -> int | None: + """How many conversations the engine's cache holds: its --conversation-cache-slots argument, which is known + before the engine has started (a lazy start has reported no INFO line yet), else the INFO line it reported.""" + args = engine.spawn[1] if hasattr(engine, "spawn") else [] + if "--conversation-cache-slots" in args[:-1]: + try: + return int(args[args.index("--conversation-cache-slots") + 1]) + except ValueError: + pass + slots = (getattr(engine, "info", None) or {}).get("conversation_cache_slots") + return slots if isinstance(slots, int) else None + + def conversation_cache_view(info: dict, hist: list, totals: dict, parked: dict) -> dict: """#596: the Monitor's Conversation cache card: the parked conversations (the engine's opt-in --conversation-cache-mib: budget, slots, what its log says) and how much of the prompts the cache gave back.""" @@ -1096,6 +1120,13 @@ class StrataEngine: self.last.update(prompt_read=int(f[14])) if len(f) >= 16: # #588 (engine 0.1.39+): routed experts read over PCIe self.last.update(offloaded=int(f[15])) + for tok in f[16:]: # key=value fields after the positional ones (the prompt's + key, eq, value = tok.partition("=") # first chunk and its streamed expert blobs) + if eq and key in ("chunk", "experts_streamed"): + try: + self.last[key] = int(value) + except ValueError: + pass def vram(self, reserve_mib: int | None, timeout: float = 120.0) -> dict: """#533: `VRAM ` between requests (the caller holds the service's FIFO): the engine shrinks its @@ -1173,6 +1204,10 @@ class StrataEngine: v = tune.get(k) if isinstance(v, (int, float)) and not isinstance(v, bool) and 0.0 <= float(v) <= 1.0: keys += f" {k}={float(v)!r}" + # a layer split's prompt chunks (--prefill-pipe): 0 off, 1 evened out, > 1 the rig's b/a in tokens + v = tune.get("prefill_pipe") + if isinstance(v, (int, float)) and not isinstance(v, bool) and 0.0 <= float(v) <= 262144.0: + keys += f" pipe={float(v)!r}" # --aux-cpus for this request: 1 puts the engine's threads that are neither pool workers nor the host on the # spare CPUs, 0 leaves them where they were created (true/false or 1/0; the engine started without a spare CPU # ignores a 1). Anything else is not sent. @@ -1416,10 +1451,11 @@ class StrataEngine: YIELDS_MAX = 2 # #656: how often one request's prompt read gives way to a shorter waiting one - def _take_control(self, cancel, plen: int, after_epoch: int | None = None): + def _take_control(self, cancel, plen: int, after_epoch: int | None = None, new: bool = False): """Waits for the control lines (one prompt read at a time), yielding None heartbeats; False when cancelled. - `after_epoch`: a request whose read gave way lets the requests waiting then go first.""" - entry = [plen] + `after_epoch`: a request whose read gave way lets the requests waiting then go first. + `new`: the request's first admission (STRATA_ADMIT_SHORT_FIRST orders only these).""" + entry = [plen, time.monotonic(), new] # (length, arrival, ordered: STRATA_ADMIT_SHORT_FIRST) with self.slot_cv: self.waiting += 1 self.wait_lens.append(entry) @@ -1433,7 +1469,17 @@ class StrataEngine: raise EngineDied("the engine stopped while this request waited; it was not sent") with self.slot_cv: turn = after_epoch is None or self.ctl_epoch > after_epoch or self.waiting <= 1 + if turn and new and ADMIT_SHORT_FIRST > 0: + turn = self._admit_next() is entry if turn and self.ctl.acquire(timeout=0.5): + if new and ADMIT_SHORT_FIRST > 0: + with self.slot_cv: # a shorter one came while this one waited in acquire() + first = self._admit_next() is entry + if not first: + self.ctl.release() + with self.slot_cv: + self.slot_cv.notify_all() + continue if self.alive(): break self.ctl.release() # it died (or is being restarted) while this one was getting @@ -1456,6 +1502,18 @@ class StrataEngine: self.slot_cv.notify_all() return True + def _admit_next(self): + """STRATA_ADMIT_SHORT_FIRST= (the caller holds slot_cv): the new request that takes the control lines + next - one whose prompt is at most that long (it reads little) before a longer one, in arrival order within + each; a longer one that has waited STRATA_ADMIT_MAX_WAIT_S counts as short (it is never held back for good). + Off (unset or 0): whoever gets the lock first, as before. Only the order changes, never what is read.""" + now = time.monotonic() + ordered = [e for e in self.wait_lens if len(e) > 2 and e[2]] + if not ordered: + return None + return min(ordered, key=lambda e: (0 if e[0] <= ADMIT_SHORT_FIRST or now - e[1] >= ADMIT_MAX_WAIT_S else 1, + e[1])) + SOLO_AGAIN_MAX = 2 # how often a request left alone in a slot goes back to the solo path SOLO_AGAIN_MIN_LEFT = 32 # ... only with at least this many tokens still allowed (max_tokens) @@ -1493,7 +1551,7 @@ class StrataEngine: prompt, left = list(ids), int(max_new) reused0 = None # the first read's reused tokens: a continued leg's count includes this request's own output carry = {} # #1869: the figures of the legs that ended in a hand-back to the solo path (the final DONE adds to them) - ok = yield from self._take_control(cancel, len(prompt)) + ok = yield from self._take_control(cancel, len(prompt), new=True) if not ok: return holding, born = True, self.gen # (the engine it now has the control lines of) @@ -1803,6 +1861,7 @@ class StrataEngine: slot_groups = 1 # --batch-groups (set when the engine starts) SLOT_PREFIX_MIN = 512 # a held prefix shorter than this is not worth unbalancing the pipeline's groups + TURN_ID = 248045 # <|im_start|>: the engine's default --turn-token (its turn checkpoint sits at the last one) def pick_slot(self, prompt: list[int]) -> int | None: """A free slot for `prompt` (the caller holds slot_cv): the one whose held tokens are the longest start of the @@ -1812,18 +1871,30 @@ class StrataEngine: if not free: return None gs = getattr(self, "slot_gs", 0) - if gs: - # pipelined groups: the engine keeps no pipelined slot as a conversation cache, so nothing is worth keeping - # free for a next turn. A group's window spans its slots up to its last busy one, so first fill a hole - # below a busy slot of the same group (it costs no extra row), then follow the spreading order + if gs > 1: + # pipelined groups of several slots: the engine keeps no idle slot below a busy one as a conversation cache + # (its pad row writes the slot), so nothing is worth keeping free for a next turn. A group's window spans + # its slots up to its last busy one, so first fill a hole below a busy slot of the same group (it costs no + # extra row), then follow the spreading order # (reported on #793: least-recently-used picks left leading holes, [_ B], that cost a whole row each) def fills_hole(b): g0 = b - b % gs return any(self.slot_busy[x] for x in range(b + 1, g0 + gs)) return min(free, key=lambda b: (not fills_hole(b), self.slot_order.index(b))) def held_prefix(b): + # what of `prompt` the slot's cache serves: the slot's whole held stream (the prompt and what it decoded) + # when this prompt continues it, else its turn checkpoint - the engine takes one at the last <|im_start|> of + # a prompt (the next request of the chat renders the history again, but not the reply it decoded or the + # assistant header), so a prompt that agrees with the held tokens up to there is that chat's next turn + # too (#2518: it was only found by the engine's own source search, after pick_slot had sent it to another + # slot and the request in that slot had overwritten the chat whose cache the next one needed) h = self.slot_held[b] - return len(h) if h and len(h) < len(prompt) and prompt[:len(h)] == h else 0 + if not h: + return 0 + if len(h) < len(prompt) and prompt[:len(h)] == h: + return len(h) + k = len(h) - 1 - h[::-1].index(self.TURN_ID) if self.TURN_ID in h else 0 + return k if 0 < k < len(prompt) and prompt[:k] == h[:k] else 0 best = max(free, key=held_prefix) # without --batch-groups a held prefix of any length is taken, as before if held_prefix(best) >= (self.SLOT_PREFIX_MIN if self.slot_groups > 1 else 1): @@ -2868,23 +2939,32 @@ class ByteTokenizer: ALWAYS = () # specials matched without parse_special (type 4, as ) + @property + def max_special_len(self): + return max(len(s) for s in self.SPECIALS) + @property def control_tokens(self): return [s for s in self.SPECIALS if s not in self.ALWAYS] def encode(self, text, parse_special=False, plain=()): - out, i = [], 0 + return self.encode_marked(text, parse_special, plain)[0] + + def encode_marked(self, text, parse_special=False, plain=()): + """encode() and its resume points (after each special), as strata_tokenizer's for PromptEncoder.""" + out, marks, i = [], [], 0 while i < len(text): for k, s in enumerate(self.SPECIALS): if (parse_special or s in self.ALWAYS) and text.startswith(s, i) and not any( a <= i < b for a, b in plain): out.append(256 + k) i += len(s) + marks.append((i, len(out))) break else: out.extend(text[i].encode("utf-8")) i += 1 - return out + return out, marks def decode(self, ids, errors="replace"): raw = bytearray() @@ -3031,10 +3111,12 @@ class Service: self.literals = literal_tags(getattr(tokenizer, "control_tokens", ())) # texts that stay text inside a message self.fit_max_tokens = fit_max_tokens # --fit-max-tokens: clamp the output cap instead of 400 self.aliases: list[str] = [] # #297: other names of the model (the config's `aliases`) + self.answer_room = True # 0.1.45: see ANSWER_ROOM_MIN_CAP (the config's "answer_room": false turns it off) self.reasoning_close_retry = False # #1053 (opt-in): close the thinking once when a reply ends inside it self.codex_compaction_cache = False # #924 (opt-in): a Codex compaction is rendered with its conversation's tools self.codex_thread_titles = False # #923 (opt-in): answer Codex's thread-title turns without the engine self.sampling_defaults = dict(sampling_defaults or {}) # the run config's `sampling` block + self.model_arch = None # the config's family ("qwen4exp", "qwen35moe"): None = no model-card defaults self.shared = {} # the web app's Chat settings for every client (POST /settings) self.shared_path = None # where they are kept between starts (next to the config) self.config_effort = None # #1641: the config's "reasoning_effort": the level for requests with none @@ -3075,6 +3157,9 @@ class Service: "prompt_ms": 0.0, "decode_ms": 0.0, "drafts_offered": 0, "drafts_accepted": 0} # #457: the MTP drafts, summed where reported self.last_timings = None # the last finished request's, llama.cpp's names (/v1/status) + # /metrics/prometheus' two throughput gauges need the totals as they stood at the last scrape, because + # llama.cpp averages them over the window between two scrapes (#2341). + self.scrape_bucket = (0, 0.0, 0, 0.0) from serve.prometheus import Latencies self.latencies = Latencies() # GET /metrics in Prometheus text: the latency histograms self.last_request_at = None # when a request last started or finished @@ -3104,6 +3189,18 @@ class Service: self.restore_after_load = None # a session file an image job saved before it unloaded self.stop_ids = set(tokenizer.encode(IM_END, parse_special=True) + tokenizer.encode("<|endoftext|>", parse_special=True)) + # #567: a prompt re-encodes only what follows the last special token it shares with a recent prompt (the + # same ids as a full encode: tools/strata_tokenizer.py PromptEncoder). Tokenizers without resume points + # encode in full. + self.prompts = None + if hasattr(tokenizer, "encode_marked") and hasattr(tokenizer, "max_special_len"): + from strata_tokenizer import PromptEncoder + # Keep at most one prompt per engine conversation slot, bounded independently of engine config. + slots = conversation_slots(engine) + if slots is not None and slots > 0: + self.prompts = PromptEncoder(tokenizer, keep=min(slots, PROMPT_ENCODER_KEEP_MAX)) + else: + self.prompts = PromptEncoder(tokenizer) def loaded(self) -> bool: return not hasattr(self.engine, "alive") or self.engine.alive() @@ -3296,20 +3393,17 @@ class Service: threading.Thread(target=loop, daemon=True).start() def free_vram_mib(self) -> int | None: - """Free VRAM on the engine's (first) GPU, from NVML (AMD backend: amdgpu's sysfs files, #301); None when it can't - be read (then nothing is refused).""" + """Free VRAM on the engine's (first) GPU, from the backend's reader - NVML, the amdgpu sysfs files (#301) or + IOKit on macOS (Metal backend); None when it can't be read (then nothing is refused).""" try: - if getattr(self, "backend", None) == "hip": - from serve.telemetry import free_vram_mib - return free_vram_mib(int(getattr(self, "gpu_index", 0) or 0), amd=True) - from serve.telemetry import _Nvml - nv = _Nvml(int(getattr(self, "gpu_index", 0) or 0)) - if not nv.ok(): + from serve.telemetry import gpu_reader + g = gpu_reader(int(getattr(self, "gpu_index", 0) or 0), amd=getattr(self, "backend", None) == "hip") + if not g.ok(): return None - m = nv.Mem() - if nv.lib.nvmlDeviceGetMemoryInfo(nv.dev, ctypes.byref(m)) != 0: + r = g.read() + if r.get("mem_total") is None or r.get("mem_used") is None: return None - return int(m.free >> 20) + return int((r["mem_total"] - r["mem_used"]) >> 20) except Exception: return None @@ -3552,7 +3646,7 @@ class Service: if effort: if api == "openai": ctk = req.get("chat_template_kwargs") if isinstance(req.get("chat_template_kwargs"), dict) else {} - if not req.get("reasoning_effort") and not req.get("reasoning") and \ + if req.get("reasoning_effort") in (None, "") and not req.get("reasoning") and \ "enable_thinking" not in ctk and "reasoning_effort" not in ctk: req["reasoning_effort"] = effort elif not req.get("thinking") and not req.get("output_config"): @@ -3703,8 +3797,15 @@ class Service: with self.status_lock: s, t = dict(self.status), dict(self.totals) running = len(self.live_reqs) or int(bool(s.get("busy"))) - prompt_n = t["prompt_tokens"] - t["reused"] # what was read, as llama.cpp counts it - prompt_s, decode_s = t["prompt_ms"] / 1000, t["decode_ms"] / 1000 + prompt_n = t["prompt_tokens"] - t["reused"] # what was read, as llama.cpp counts it + prompt_s, decode_s = t["prompt_ms"] / 1000, t["decode_ms"] / 1000 + # llama.cpp reads these two gauges from a bucket it resets on every scrape (server_metrics::bucket, + # "the gauges are averaged over the window between two scrapes"), so an idle server reports 0 rather + # than its lifetime average. Keep the same window here, or the names lie to llama.cpp dashboards. + was_prompt_n, was_prompt_s, was_predicted_n, was_decode_s = self.scrape_bucket + self.scrape_bucket = (prompt_n, prompt_s, t["output_tokens"], decode_s) + d_prompt_n, d_prompt_s = prompt_n - was_prompt_n, prompt_s - was_prompt_s + d_predicted_n, d_decode_s = t["output_tokens"] - was_predicted_n, decode_s - was_decode_s rows = [ ("llamacpp:prompt_tokens_total", "counter", "Prompt tokens read", prompt_n), ("llamacpp:prompt_tokens_cached_total", "counter", "Prompt tokens the conversation cache already held", @@ -3712,10 +3813,10 @@ class Service: ("llamacpp:prompt_seconds_total", "counter", "Prompt read time", prompt_s), ("llamacpp:tokens_predicted_total", "counter", "Tokens generated", t["output_tokens"]), ("llamacpp:tokens_predicted_seconds_total", "counter", "Generation time", decode_s), - ("llamacpp:prompt_tokens_seconds", "gauge", "Average prompt read throughput in tokens/s", - prompt_n / prompt_s if prompt_s else 0), - ("llamacpp:predicted_tokens_seconds", "gauge", "Average generation throughput in tokens/s", - t["output_tokens"] / decode_s if decode_s else 0), + ("llamacpp:prompt_tokens_seconds", "gauge", "Average prompt read throughput in tokens/s, since the last scrape", + d_prompt_n / d_prompt_s if d_prompt_s > 0 else 0), + ("llamacpp:predicted_tokens_seconds", "gauge", "Average generation throughput in tokens/s, since the last scrape", + d_predicted_n / d_decode_s if d_decode_s > 0 else 0), ("llamacpp:requests_processing", "gauge", "Requests running", running), ("llamacpp:requests_deferred", "gauge", "Requests waiting their turn", s.get("queued") or 0), ("strata:requests_total", "counter", "Requests finished", t["requests"]), @@ -3787,17 +3888,21 @@ class Service: tail = "" if off else EFFORT_TURN.format(EFFORT_TEXT.get(effort, EFFORT_TEXT["medium"])) return history + tail + full[len(mine):] - def encode_prompt(self, messages, tools, kwargs) -> list[int]: + def encode_prompt(self, messages, tools, kwargs, store: bool = True) -> list[int]: """The request's prompt: the template rendered and tokenized. #537: a / written inside a message's text is encoded as the text it is, not as the model's reasoning markers (the template's own are). A control token's text (<|im_start|>, <|im_end|>, <|endoftext|>, ...) inside a message is text as well: only - the control tokens the template writes are control tokens.""" + the control tokens the template writes are control tokens. `store=False` (a token count) reuses the kept + prompts without adding or replacing one.""" marked, marked_tools, changed = mark_think_literals(messages, tools, self.literals) prompt = self.render_prompt(marked, marked_tools, kwargs) - if not changed: - return self.tok.encode(prompt, parse_special=True) - prompt, plain = unmark_think_literals(prompt, self.literals) - return self.tok.encode(prompt, parse_special=True, plain=plain) + plain = () + if changed: + prompt, plain = unmark_think_literals(prompt, self.literals) + if self.prompts is not None: # #567: only what follows the last special shared with a recent prompt + return self.prompts.encode(prompt, plain, store) + return self.tok.encode(prompt, parse_special=True, plain=plain) if plain else \ + self.tok.encode(prompt, parse_special=True) @contextlib.contextmanager def _turn(self, cancel=None): @@ -3845,14 +3950,26 @@ class Service: content[n] = {"type": "text", "text": f"[image omitted: {why}]"} return fetched + def forget_prompt(self, ids): + if self.prompts is not None: + self.prompts.forget(ids) + + def prepare(self, messages, tools, kwargs, max_new=None, force=None, req=None, cancel=None): """-> (ids, thinking, max_new). An unset or non-positive max_new (some clients send -1) means "unlimited": the rest of the context. `force` (forced_call): without thinking the reply starts with it, so it ends the prompt; with thinking, Service.run writes it once the thinking is over.""" # `cancel` (#1999): set when the client hangs up; the downloads, the wait for the turn and the encodes stop # between pictures and raise ClientGone + if self.vision is None and os.environ.get("STRATA_OMIT_IMAGES_WHEN_NO_VISION") == "1": + for m in messages: + if isinstance(m.get("content"), list): + m["content"] = [ + {"type": "text", "text": "[image omitted on text-only engine]"} if isinstance(it, dict) and it.get("type") == "image" else it + for it in m["content"] + ] fetched = self._note_unreadable_tool_images(messages, cancel) - ids = self.encode_prompt(messages, tools, kwargs) + ids = encoded_ids = self.encode_prompt(messages, tools, kwargs) if force and kwargs.get("enable_thinking", True) is False: ids = ids + self.tok.encode(force, parse_special=True) self.embeddings.path = None @@ -3914,12 +4031,16 @@ class Service: room = ctx - CTX_SLACK - len(ids) if max_new is None or max_new <= 0 or (self.fit_max_tokens and room < 1): if room < 1: + # A prompt refused for the context is not kept (#567): the client cannot resend it unchanged. The + # other refusals (engine starting, images) keep it: the same prompt comes back and its encode is reused. + self.forget_prompt(encoded_ids) raise ValueError("request exceeds the context window: " f"prompt ({len(ids)} tokens) leaves no room to answer in the context " f"({ctx}); requests are never truncated") max_new = room elif max_new > room: if not self.fit_max_tokens: + self.forget_prompt(encoded_ids) raise ValueError("request exceeds the context window: " f"prompt ({len(ids)} tokens) + max tokens ({max_new}) exceeds the context " f"({ctx}); requests are never truncated. Send a smaller " @@ -4076,12 +4197,23 @@ class Service: def _run(self, ids, thinking, tools, max_new, sampling, cancel, force=None) -> Iterator[tuple[str, object]]: """Yields ("event", Event) as text arrives, then ("done", {"finish": .., "completion_tokens": ..}). `force` (forced_call): the opening of the call the reply must make - see prepare().""" - defaults = {**self.sampling_defaults, **self.shared} # the config's, then the Chat settings shared with apps + # the model card's sampling for this thinking state, then the config's block, then the Chat settings shared with + # apps; the request's own fields win over all of them + defaults = {**model_card_sampling(self.model_arch, thinking), **self.sampling_defaults, **self.shared} if defaults: # the request's own fields win (explicit 0 stays greedy) req_values = {k: v for k, v in (sampling or {}).items() if v is not None} sampling = {**defaults, **req_values} # #123: read after the merge, so a budget shared through POST /settings is seen like the other keys budget = self.reasoning_budget(sampling) if thinking else None + if thinking and not budget and self.answer_room and max_new and max_new >= ANSWER_ROOM_MIN_CAP: + # 0.1.45: max_tokens counts the thinking too. A client that caps the reply at 8192 (the desktop agent's + # harness sends max_completion_tokens: 8192 on every request) and a task the model thinks long about got + # the cut-off in the middle of the thinking and NO answer. So thinking is closed early enough to leave + # the answer its quarter of the cap (the budget's wrap-up: a short closing sentence, , the answer). + # A request with a budget of its own, or max_tokens below ANSWER_ROOM_MIN_CAP, is left as it was. + budget = max_new - min(max(ANSWER_RESERVE_MIN, max_new // 4), max_new // 2) + if max_new <= 32768: print(f"[strata] thinking limited to {budget} of max_tokens {max_new} so the answer has room " + "(\"answer_room\": false in the config turns this off)", flush=True) if budget and max_new and max_new > 0: # #984: a budget at or near max_tokens never fires before the reply is cut off while still thinking, so # it has no answer. Close the thinking early enough to leave room for one. @@ -4586,13 +4718,13 @@ def request_stats(segments: list[dict]) -> dict: if not segments: return {} result = dict(segments[-1]) - for key in ("reused", "prompt_read"): + for key in ("reused", "prompt_read", "chunk"): if key in segments[0]: result[key] = segments[0][key] else: result.pop(key, None) for key in ("prompt_ms", "decode_ms", "generated", "drafts_offered", "drafts_accepted", - "hits", "lookups", "offloaded", "ram_blobs", "file_blobs", "file_mb"): + "hits", "lookups", "offloaded", "ram_blobs", "file_blobs", "file_mb", "experts_streamed"): if any(segment.get(key) is not None for segment in segments): result[key] = sum(segment.get(key) or 0 for segment in segments) return result @@ -5612,7 +5744,7 @@ def make_handler(svc: Service): if not svc.loaded() and not getattr(svc.engine, "unloaded", False): self._json(503, {"error": {"message": "the engine is not running"}}) return - defaults = {**svc.sampling_defaults, **svc.shared} + defaults = {**model_card_sampling(svc.model_arch, True), **svc.sampling_defaults, **svc.shared} names = {"repetition_penalty": "repeat_penalty", "penalty_last_n": "repeat_last_n"} params = {names.get(k, k): v for k, v in defaults.items() if k in ("temperature", "top_p", "top_k", "min_p", "seed", "repetition_penalty", @@ -6045,7 +6177,7 @@ def make_handler(svc: Service): messages, tools, kw = anthropic_to_messages(req, svc.anthropic_think_unasked) stop_strings(req) # the same 400 as the request itself would get check_request_sampling(req) # ... likewise a sampling field of the wrong type - self._json(200, {"input_tokens": len(svc.encode_prompt(messages, tools, kw))}) + self._json(200, {"input_tokens": len(svc.encode_prompt(messages, tools, kw, store=False))}) def _anthropic(self, req): svc.load() @@ -6433,6 +6565,43 @@ def clean_shared_defaults(d) -> dict: return out +def split_joined_args(cfg: dict) -> None: + """#2267: an engine option written as ONE element of the config's "args" - "--ple-io mmap" or "--ple-io=mmap" - reaches + the engine as a single argument ("unknown argument: --ple-io mmap"; the usual form is two elements, "--ple-io", + "mmap"). Recommend, never force: such an element is split into the flag and its value, and said at startup. Only an + element that starts with a flag name (--letters-digits-dashes) followed by a space or "=" is touched, and never one + that sits in the value place of the flag before it.""" + args = cfg.get("args") + if not isinstance(args, list): + return + out, changed = [], [] + for a in args: + m = re.match(r"^(--[A-Za-z0-9][A-Za-z0-9-]*)(?:\s+|=)(\S.*)$", a) if isinstance(a, str) else None + prev = out[-1] if out else None + in_value_place = isinstance(prev, str) and re.match(r"^--[A-Za-z0-9][A-Za-z0-9-]*$", prev) is not None + if m is None or in_value_place: + out.append(a) + continue + out += [m.group(1), m.group(2).strip()] + changed.append(a) + if changed: + cfg["args"] = out + print("[strata] the config's args hold " + ", ".join(repr(c) for c in changed) + " as one element each: split into the " + 'flag and its value (write them as two elements, e.g. "--ple-io", "mmap")', flush=True) + + +def read_config(path) -> dict: + """The run config's JSON. A mistake in it (a missing comma between two args, a trailing comma) is said with its + place in the file instead of a traceback (#2267: editing "args" by hand is where it happens).""" + text = Path(path).read_text(encoding="utf-8-sig") # Notepad adds a BOM + try: + return json.loads(text) + except json.JSONDecodeError as e: + raise SystemExit(f"[strata] {path} is not valid JSON: {e.msg} at line {e.lineno}, column {e.colno}. Check the commas " + 'between the entries of "args" (each option and each value is its own quoted entry, e.g. ' + '"--ple-io", "mmap",) and that no entry is followed by a comma at the end of a list.') from None + + def lift_reasoning_effort_arg(cfg: dict) -> None: """#1641: "--reasoning-effort high" in a config's "args" is meant for the thinking level, which is not an engine option (the engine stopped at the unknown flag). It is the config's "reasoning_effort" (none, low, medium or high: @@ -6615,6 +6784,25 @@ def check_request_sampling(req) -> None: FAMILY_MODEL_NAMES = {"qwen4exp": "qwen3.8-flash-next", "qwen35moe": "qwen3.6-35b-a3b"} +# The sampling each model card recommends ("Best Practices"), used when neither the request, the Chat settings nor the +# run config's `sampling` block says otherwise. Greedy is only for a request or config that asks for temperature 0. +# min_p 0, repetition_penalty 1 and presence_penalty 0 are the engine's own "off", so they are not listed. +MODEL_CARD_SAMPLING = { + # Qwen3.8-Flash-Next (and Gyro, the same base model) + "qwen4exp": {"thinking": {"temperature": 1.0, "top_p": 0.95, "top_k": 20}, + "instruct": {"temperature": 0.7, "top_p": 0.8, "top_k": 20, "presence_penalty": 1.5}}, + # Qwen3.6-35B-A3B: the general-task sets (the card's precise-coding set is temperature 0.6, presence_penalty 0) + "qwen35moe": {"thinking": {"temperature": 1.0, "top_p": 0.95, "top_k": 20, "presence_penalty": 1.5}, + "instruct": {"temperature": 0.7, "top_p": 0.8, "top_k": 20, "presence_penalty": 1.5}}, +} + + +def model_card_sampling(arch, thinking) -> dict: + """The model card's sampling for the family `arch` and the thinking state ({} for an unknown or unset family).""" + sets = MODEL_CARD_SAMPLING.get(arch) if arch else None + return dict(sets["thinking" if thinking else "instruct"]) if sets else {} + + def default_model_name(cfg: dict) -> str: """The config's "model_name", else its family's name ("arch": qwen4exp, qwen35moe), else Flash-Next's.""" return cfg.get("model_name") or FAMILY_MODEL_NAMES.get(cfg.get("arch") or "qwen4exp", "qwen3.8-flash-next") @@ -6667,7 +6855,8 @@ def main() -> int: "each, requests spread over them (also \"replicas\" in the config: a number, or " "[{\"gpus\": [0, 1]}, {\"gpus\": [2, 3]}])") a = ap.parse_args() - cfg = json.loads(Path(a.config).read_text(encoding="utf-8-sig")) if a.config else {} # Notepad adds a BOM + cfg = read_config(a.config) if a.config else {} + split_joined_args(cfg) lift_reasoning_effort_arg(cfg) if a.replicas is not None: cfg["replicas"] = a.replicas @@ -6768,6 +6957,9 @@ def main() -> int: model_name=default_model_name(cfg), vision=vision, sampling_defaults=sampling_defaults, fit_max_tokens=a.fit_max_tokens or cfg.get("fit_max_tokens") is True) + if a.engine == "strata": # the model card's sampling when nothing else is said (see MODEL_CARD_SAMPLING) + svc.model_arch = cfg.get("arch") or "qwen4exp" + svc.answer_room = cfg.get("answer_room") is not False svc.reasoning_close_retry = cfg.get("reasoning_close_retry") is True # #1053: opt-in, off by default if cfg.get("images"): # image generation (serve/images.py), set up by setup --images icfg = {k: (os.path.abspath(os.path.join(cfg.get("cwd") or ".", v)) diff --git a/serve/telemetry.py b/serve/telemetry.py index 3232c63c..697eeec7 100644 --- a/serve/telemetry.py +++ b/serve/telemetry.py @@ -3,7 +3,9 @@ A background thread samples once a second and keeps the last 60 readings of each series for the sparklines: - GPU: NVIDIA's own NVML library (nvml.dll / libnvidia-ml.so.1, installed with every driver) through ctypes, so no pip package is needed: load, VRAM, temperature, power, PCIe link and throughput. With the AMD backend (#301): the - amdgpu driver's Linux sysfs files - load, VRAM, temperature and power. + amdgpu driver's Linux sysfs files - load, VRAM, temperature and power. On macOS (Metal backend): IOKit's + IOAccelerator registry - load and VRAM; the temperature and power sensors need root (`powermetrics`), and an Apple + Silicon GPU sits on the die, so there is no PCIe link. - CPU, RAM, disk: `psutil` when it is installed (setup installs it); without it the CPU and RAM readings fall back to the OS (Windows GlobalMemoryStatusEx / GetSystemTimes, Linux /proc) and the disk rate is absent. Anything that cannot be read is None; nothing here can stop the server. @@ -14,6 +16,8 @@ import collections import ctypes import os import platform +import re +import subprocess import sys import threading import time @@ -246,9 +250,114 @@ class _Amd: return out +def metal_working_set(): + """The Apple GPU's memory budget in bytes: MTLCreateSystemDefaultDevice's recommendedMaxWorkingSetSize, through + the Metal framework and objc_msgSend (no PyObjC) - the same property metal/src/mtl_rt.mm reads for the engine + (its "GPU working set" log line). None off macOS, or when it cannot be read: a selector the device does not + answer would raise an Objective-C exception no ctypes call can catch, so the class is asked first.""" + if sys.platform != "darwin": + return None + try: + metal = ctypes.CDLL("/System/Library/Frameworks/Metal.framework/Metal") + objc = ctypes.CDLL("/usr/lib/libobjc.dylib") + metal.MTLCreateSystemDefaultDevice.restype = ctypes.c_void_p + objc.sel_registerName.restype = ctypes.c_void_p + objc.object_getClass.argtypes = [ctypes.c_void_p] + objc.object_getClass.restype = ctypes.c_void_p + objc.class_respondsToSelector.argtypes = [ctypes.c_void_p, ctypes.c_void_p] + objc.class_respondsToSelector.restype = ctypes.c_ubyte + objc.objc_msgSend.argtypes = [ctypes.c_void_p, ctypes.c_void_p] + objc.objc_msgSend.restype = ctypes.c_ulonglong + objc.objc_release.argtypes = [ctypes.c_void_p] + objc.objc_release.restype = None + dev = metal.MTLCreateSystemDefaultDevice() + if not dev: + return None + try: + sel = objc.sel_registerName(b"recommendedMaxWorkingSetSize") + if not objc.class_respondsToSelector(objc.object_getClass(dev), sel): + return None + return int(objc.objc_msgSend(ctypes.c_void_p(dev), ctypes.c_void_p(sel))) or None + finally: + objc.objc_release(dev) # MTLDevice.h: MTLCreateSystemDefaultDevice is NS_RETURNS_RETAINED + except (OSError, AttributeError, ValueError): + return None + + +class _Metal: + """macOS: an Apple GPU's readings from the IOKit registry (`ioreg -l -w0 -r -c IOAccelerator`, 22 ms), with + _Nvml's interface: the load and the memory are the driver's PerformanceStatistics of the index'th accelerator + node - Device Utilization % and the in-use system memory in bytes. The total is the device's Metal working set, + the budget one app gets of the unified memory. No temperature or power reading exists without root + (`powermetrics`), and an Apple Silicon GPU sits on the die, so the PCIe cards have nothing to show.""" + note = ("no GPU temperature or power reading without root (powermetrics), and an Apple Silicon GPU sits on the " + "die, so there is no PCIe link to show; the engine's own figures are in its log") + cmd = ("ioreg", "-l", "-w0", "-r", "-c", "IOAccelerator") + + def __init__(self, index=0, run=None): + self.index = index + self._run = run or self._ioreg + self.nodes = len(self._blocks()) + self.working_set = metal_working_set() + self._name = None + b = self._block() + if b: + m = re.search(r'"model" = "([^"]+)"', b) + self._name = m.group(1) if m else None + + @staticmethod + def _ioreg(): + try: + return subprocess.run(list(_Metal.cmd), capture_output=True, text=True, timeout=10).stdout or "" + except (OSError, subprocess.SubprocessError): + return "" + + def _blocks(self): + """The accelerator nodes, in the registry's order: every node that carries a PerformanceStatistics.""" + try: + out = self._run() or "" + except OSError: # ioreg missing, or a caller's runner that raises + return [] + return [b for b in re.split(r"^\s*\+-o ", out, flags=re.M) if '"PerformanceStatistics"' in b] + + def _block(self): + b = self._blocks() + return b[self.index] if 0 <= self.index < len(b) else None + + def ok(self): + return self.index < self.nodes + + def name(self): + return self._name + + def read(self): + b = self._block() + if b is None: + return {} + + def num(key): + m = re.search(r'"%s"=([0-9.]+)' % re.escape(key), b) + return float(m.group(1)) if m else None + + out = {} + u = num("Device Utilization %") + if u is None: # older chips report only the renderer and the tiler + r, t = num("Renderer Utilization %"), num("Tiler Utilization %") + u = max(v for v in (r, t) if v is not None) if (r is not None or t is not None) else None + out["util"] = int(u) if u is not None else None + used = num("In use system memory") + out["mem_used"] = int(used) if used is not None else None + out["mem_total"] = self.working_set + return out + + def gpu_reader(index=0, amd=False): - """The card's readings: NVML (NVIDIA), or the amdgpu sysfs files with the AMD backend (#301).""" - return _Amd(index) if amd else _Nvml(index) + """The card's readings: NVML (NVIDIA), the amdgpu sysfs files with the AMD backend (#301), or IOKit's + IOAccelerator registry on macOS (Metal backend). `amd` comes first: no Mac runs the HIP backend, and the + tests' AMD trees are read on a Mac too.""" + if amd: + return _Amd(index) + return _Metal(index) if sys.platform == "darwin" else _Nvml(index) _DISPLAY_CLASS = r"SYSTEM\CurrentControlSet\Control\Class\{4d36e968-e325-11ce-bfc1-08002be10318}" @@ -413,15 +522,21 @@ class Telemetry: cards = windows_amd_cards() win_amd = [cards[i] for i in idx if 0 <= i < len(cards)] or cards[:1] gpu_name = " + ".join(c["name"] for c in win_amd) if win_amd else None + # the macOS reader knows the device's working set (the VRAM budget of the unified memory); None elsewhere + vram_total = None + if win_amd is None and isinstance(self.gpu, _Metal) and self.gpu.ok(): + ws = self.gpu.working_set + vram_total = ws >> 20 if ws else None self.static = { "gpu_name": gpu_name, "gpu_vram_total_mib": (sum(c["vram_bytes"] for c in win_amd) >> 20 - if win_amd and all(c["vram_bytes"] for c in win_amd) else None), + if win_amd and all(c["vram_bytes"] for c in win_amd) else vram_total), "gpu_count": len(self.gpus), # #1380: the AMD readings are the amdgpu driver's Linux sysfs files; a Windows AMD card has none yet, and the # dashboard said "not readable (NVML)" or showed empty tiles with no word why "gpu_note": ("no GPU load or VRAM readings for AMD cards on Windows yet (Linux reads them from the amdgpu " - "driver); the engine's own VRAM figures are in its log" if amd and not self.gpu.ok() else None), + "driver); the engine's own VRAM figures are in its log" if amd and not self.gpu.ok() + else getattr(self.gpu, "note", None)), "cpu_name": _cpu_name(), "cores": (self.ps.cpu_count(logical=False) if self.ps else None) or None, "threads": os.cpu_count(), diff --git a/serve/test_admit_order.py b/serve/test_admit_order.py new file mode 100644 index 00000000..67766e55 --- /dev/null +++ b/serve/test_admit_order.py @@ -0,0 +1,119 @@ +"""serve/test_admit_order.py - STRATA_ADMIT_SHORT_FIRST: new requests take the control lines short prompts first +(StrataEngine._take_control / _admit_next; no process, no GPU). + + python -m unittest serve.test_admit_order -v +""" +from __future__ import annotations + +import sys +import threading +import time +import unittest +from pathlib import Path +from unittest import mock + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +import serve.server as server # noqa: E402 +from serve.server import StrataEngine # noqa: E402 + + +def engine_without_process(n=2): + """A StrataEngine with the admission state of a started one and no process: alive() is `e.up`.""" + e = StrataEngine("missing-executable", [], lazy=True) + e.batch = n + e.slot_busy = [False] * n + e.slot_claims, e.slot_releasing, e._orphan_seen = set(), set(), {} + e.slot_cv = threading.Condition() + e.waiting, e.admitting, e.wait_lens, e.ctl_epoch = 0, 0, [], 0 + e.ctl = threading.Lock() + e.gen = 1 + e.up = True + e.alive = lambda: e.up + return e + + +class AdmitOrderTest(unittest.TestCase): + def setUp(self): + self.e = engine_without_process(16) + self.order, self.cancels = [], [] + self.order_lock = threading.Lock() + + def tearDown(self): + for c in self.cancels: + c.set() + + def take(self, name, plen, **kw): + """A request waiting for the control lines; once it has them it notes its name and gives them back.""" + c = threading.Event() + self.cancels.append(c) + + def go(): + gen = self.e._take_control(c, plen, **kw) + try: + while True: + next(gen) + except StopIteration as s: + if s.value: + with self.order_lock: + self.order.append(name) + time.sleep(0.05) # its read + self.e.ctl.release() + threading.Thread(target=go, daemon=True).start() + + def until(self, cond, s=5.0): + end = time.time() + s + while time.time() < end and not cond(): + time.sleep(0.01) + self.assertTrue(cond()) + + def queue(self, *reqs): + """While a long read holds the control lines, the requests arrive in this order; then the read ends.""" + self.assertTrue(self.e.ctl.acquire(timeout=1)) + for name, plen in reqs: + self.take(name, plen, new=True) + self.until(lambda n=len(self.order) + sum(1 for _ in self.e.wait_lens): self.e.waiting >= 1) + time.sleep(0.02) # (distinct arrival times) + self.until(lambda: self.e.waiting == len(reqs)) + self.e.ctl.release() + self.until(lambda: len(self.order) == len(reqs), s=15.0) + return self.order + + def test_short_first_puts_the_chat_before_the_agent_reads(self): + with mock.patch.object(server, "ADMIT_SHORT_FIRST", 8192), mock.patch.object(server, "ADMIT_MAX_WAIT_S", 60.0): + order = self.queue(("agent1", 120000), ("agent2", 190000), ("chat", 500), ("agent3", 150000)) + self.assertEqual(order[0], "chat") + self.assertEqual(order[1:], ["agent1", "agent2", "agent3"]) # arrival order among the long ones + + def test_a_long_request_that_waited_too_long_goes_first(self): + with mock.patch.object(server, "ADMIT_SHORT_FIRST", 8192), mock.patch.object(server, "ADMIT_MAX_WAIT_S", 0.2): + self.assertTrue(self.e.ctl.acquire(timeout=1)) + self.take("agent", 120000, new=True) + self.until(lambda: self.e.waiting == 1) + time.sleep(0.4) # past STRATA_ADMIT_MAX_WAIT_S + self.take("chat", 500, new=True) + self.until(lambda: self.e.waiting == 2) + self.e.ctl.release() + self.until(lambda: len(self.order) == 2, s=15.0) + self.assertEqual(self.order, ["agent", "chat"]) + + def test_off_by_default(self): + self.assertEqual(server.ADMIT_SHORT_FIRST if "STRATA_ADMIT_SHORT_FIRST" not in __import__("os").environ + else 0, 0) + with mock.patch.object(server, "ADMIT_SHORT_FIRST", 0): + order = self.queue(("agent1", 120000), ("chat", 500)) + self.assertEqual(sorted(order), ["agent1", "chat"]) # both get their turn, in any order + + def test_a_continuation_is_not_held_back(self): + """A request already admitted once (new=False: promoted from the solo path) is not part of the order.""" + with mock.patch.object(server, "ADMIT_SHORT_FIRST", 8192), mock.patch.object(server, "ADMIT_MAX_WAIT_S", 60.0): + self.assertTrue(self.e.ctl.acquire(timeout=1)) + self.take("chat", 500, new=True) + self.take("continuation", 190000) + self.until(lambda: self.e.waiting == 2) + self.e.ctl.release() + self.until(lambda: len(self.order) == 2, s=15.0) + self.assertEqual(sorted(self.order), ["chat", "continuation"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/serve/test_effort_paths.py b/serve/test_effort_paths.py index dbc1ae6f..15525f6b 100644 --- a/serve/test_effort_paths.py +++ b/serve/test_effort_paths.py @@ -130,6 +130,7 @@ class EffortPaths(unittest.TestCase): "thinking": {"type": "enabled", "budget_tokens": 20000}}), "xhigh") self.assertEqual(self.level("/v1/responses", {"input": "hello", "reasoning": {"effort": "high"}}), "xhigh") self.assertEqual(self.level("/v1/chat/completions", {"messages": MSG, "reasoning_effort": "none"}), "none") + self.assertEqual(self.level("/v1/chat/completions", {"messages": MSG, "reasoning_effort": False}), "none") def test_the_chat_settings_shared_with_apps_beat_the_config(self): self.svc.config_effort = "low" diff --git a/serve/test_image_omission.py b/serve/test_image_omission.py new file mode 100644 index 00000000..677229aa --- /dev/null +++ b/serve/test_image_omission.py @@ -0,0 +1,69 @@ +"""The local text-only image switch must run before incremental prompt encoding.""" +import copy +import os +import unittest +from pathlib import Path +from unittest import mock + +from serve.frontend import ChatTemplate +from serve.server import ByteTokenizer, MockEngine, Service + + +ROOT = Path(__file__).resolve().parents[1] + + +class ImageOmission(unittest.TestCase): + def service(self): + tok = ByteTokenizer() + return Service(MockEngine(tok, "ok", max_context=1 << 20), tok, + ChatTemplate(ROOT / "serve/chat_template.jinja")) + + def messages(self): + return [{"role": "user", "content": [ + {"type": "text", "text": "describe literally"}, + {"type": "image", "source": "unused.png"}, + ]}] + + @mock.patch.dict(os.environ, {"STRATA_OMIT_IMAGES_WHEN_NO_VISION": "1"}) + def test_omitted_images_are_cached_as_text_and_retries_are_stable(self): + svc = self.service() + original = self.messages() + messages = copy.deepcopy(original) + text = messages[0]["content"][0] + first, _, _ = svc.prepare(messages, None, {}, 16) + self.assertIs(messages[0]["content"][0], text) + self.assertEqual(messages[0]["content"][1], { + "type": "text", "text": "[image omitted on text-only engine]"}) + self.assertEqual(svc.prompts.entries[0][2], first) + full = self.service() + full.prompts = None + self.assertEqual(first, full.encode_prompt(messages, None, {})) + again, _, _ = svc.prepare(copy.deepcopy(original), None, {}, 16) + self.assertEqual(first, again) + self.assertGreater(svc.prompts.last_reused, 0) + same_object, _, _ = svc.prepare(messages, None, {}, 16) + self.assertEqual(first, same_object) + + @mock.patch.dict(os.environ, {"STRATA_OMIT_IMAGES_WHEN_NO_VISION": "0"}) + def test_disabled_switch_keeps_the_normal_vision_error(self): + svc = self.service() + messages = self.messages() + before = copy.deepcopy(messages) + with self.assertRaisesRegex(ValueError, "without the vision encoder"): + svc.prepare(messages, None, {}, 16) + self.assertEqual(messages, before) + + @mock.patch.dict(os.environ, {"STRATA_OMIT_IMAGES_WHEN_NO_VISION": "1"}) + def test_configured_vision_is_not_omitted(self): + svc = self.service() + svc.vision = object() + messages = self.messages() + before = copy.deepcopy(messages) + with mock.patch.object(svc, "_note_unreadable_tool_images", side_effect=RuntimeError("before encoding")): + with self.assertRaisesRegex(RuntimeError, "before encoding"): + svc.prepare(messages, None, {}, 16) + self.assertEqual(messages, before) + + +if __name__ == "__main__": + unittest.main() diff --git a/serve/test_logit_bias.py b/serve/test_logit_bias.py index ee3751b1..8543a9c2 100644 --- a/serve/test_logit_bias.py +++ b/serve/test_logit_bias.py @@ -59,3 +59,116 @@ class LogitBias(unittest.TestCase): finally: httpd.shutdown() httpd.server_close() + + +class LogitBiasCache(unittest.TestCase): + """A long list is normalized once; the cache must not change what a call returns or raises.""" + + def setUp(self): + from serve import logit_bias + logit_bias._cache.clear() + + def test_same_answers_as_the_plain_function(self): + from serve.logit_bias import _normalize + forms = [{str(i): -100 for i in range(200)}, + [[i, False] for i in range(200)], + {str(i): (i % 7) - 3.5 for i in range(200)}, + [[str(i), 2] for i in range(200)]] + for vocab in (None, 1000): + for form in forms: + want = _normalize(form, vocab) + for _ in range(3): # the first call fills the cache, the others read it + self.assertEqual(normalize(form, vocab), want) + self.assertEqual(normalize(json.loads(json.dumps(form)), vocab), want) # an equal list, another object + for form in forms: + ref = " logit_bias=" + ",".join(f"{i}:{b:g}" for i, b in sorted(_normalize(form).items())) + self.assertEqual(engine_key(form), ref) + self.assertEqual(engine_key(form), ref) + + def test_false_is_not_zero(self): + ban = [[i, False] for i in range(100)] + zero = [[i, 0] for i in range(100)] + self.assertTrue(ban == zero) # Python calls them equal, so an == key would mix them up + for first, second in [(ban, zero), (zero, ban)]: + self.setUp() + self.assertEqual(set(normalize(first).values()), {-100.0 if first is ban else 0.0}) + self.assertEqual(set(normalize(second).values()), {-100.0 if second is ban else 0.0}) + + def test_a_bad_list_is_never_stored(self): + good = {str(i): -1 for i in range(100)} + normalize(good, 1000) + for bad in [{**good, "5": True}, {**good, "5": 101}, {**good, "5": float("nan")}, {**good, "5": "x"}]: + for _ in range(2): + with self.assertRaises(ValueError): + normalize(bad, 1000) + pairs = [[i, 1] for i in range(100)] + normalize(pairs) + with self.assertRaises(ValueError): # True == 1, but a bool is not a bias + normalize([[i, 1] for i in range(99)] + [[99, True]]) + with self.assertRaises(ValueError): # a duplicate ID + normalize([[i, 1] for i in range(100)] + [[0, 1]]) + + def test_vocabulary_size_is_part_of_the_key(self): + ids = {str(i): -1 for i in range(100)} + self.assertEqual(len(normalize(ids, 1000)), 100) + with self.assertRaises(ValueError): + normalize(ids, 50) + self.assertEqual(len(normalize(ids)), 100) + with self.assertRaises(ValueError): + normalize(ids, 50) + + def test_the_result_is_a_copy(self): + ids = {str(i): -1 for i in range(100)} + first = normalize(ids) + first[0] = 99.0 + self.assertEqual(normalize(ids)[0], -1.0) + self.assertEqual(validate_request({"logit_bias": ids}, SimpleNamespace(info={"logit_bias": 1}), 1000)[0], -1.0) + + def test_an_equal_list_in_a_later_request_is_a_hit(self): + from serve import logit_bias + a = {str(i): -100 for i in range(100)} + b = json.loads(json.dumps(a)) + self.assertIsNot(a, b) + key_a = engine_key(a) + key_b = engine_key(b) + self.assertIs(key_a, key_b) # one entry, one key string + self.assertEqual(len(logit_bias._cache), 1) + + def test_json_equivalent_pairs_and_token_keys(self): + pairs = [[i, 0] for i in range(64)] + tuples = [(i, 0) for i in range(64)] + numeric = {i: 0 for i in range(64)} + strings = {str(i): 0 for i in range(64)} + for first, second in ((pairs, tuples), (numeric, strings)): + self.assertEqual(normalize(first, 1000), normalize(second, 1000)) + self.assertEqual(engine_key(first), engine_key(second)) + duplicates = {**numeric, "0": 0} + for _ in range(2): + with self.assertRaisesRegex(ValueError, "duplicate"): + normalize(duplicates, 1000) + + def test_same_length_mutations_are_validated(self): + for value in ([[i, 0] for i in range(64)], {str(i): 0 for i in range(64)}): + with self.subTest(form=type(value).__name__): + self.assertEqual(normalize(value, 1000)[0], 0) + engine_key(value) + if isinstance(value, list): + value[0][1] = -100 + else: + value["0"] = -100 + self.assertEqual(normalize(value, 1000)[0], -100) + self.assertIn("0:-100", engine_key(value)) + if isinstance(value, list): + value[0][1] = True + else: + value["0"] = True + with self.assertRaises(ValueError): + normalize(value, 1000) + with self.assertRaises(ValueError): + engine_key(value) + + def test_the_cache_is_bounded(self): + from serve import logit_bias + for n in range(logit_bias._CACHE_MAX + 3): + normalize({str(i): -(n + 1) for i in range(100)}) + self.assertEqual(len(logit_bias._cache), logit_bias._CACHE_MAX) diff --git a/serve/test_model_card_sampling.py b/serve/test_model_card_sampling.py new file mode 100644 index 00000000..ee491506 --- /dev/null +++ b/serve/test_model_card_sampling.py @@ -0,0 +1,108 @@ +"""serve/test_model_card_sampling.py - a request that sends no sampling gets the model card's, not greedy (mock engine). + +Qwen's "Best Practices": thinking mode temperature 1.0 / top_p 0.95 / top_k 20, non-thinking 0.7 / 0.8 / 20 with +presence_penalty 1.5. The request's own fields win, then the Chat settings shared with apps, then the config's +`sampling` block, then the card; temperature 0 asked for is still greedy. + + python -m unittest serve.test_model_card_sampling -v +""" +from __future__ import annotations + +import json +import sys +import unittest +import urllib.request +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from serve.frontend import ChatTemplate # noqa: E402 +from serve.server import ByteTokenizer, MockEngine, Service, StrataEngine, model_card_sampling, serve # noqa: E402 + +ROOT = Path(__file__).resolve().parents[1] + + +class Recording(MockEngine): + def generate(self, ids, max_new, sampling, cancel, embeddings=None): + self.sampling = dict(sampling or {}) + yield from super().generate(ids, max_new, sampling, cancel, embeddings) + + +class ModelCardSampling(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.tok = ByteTokenizer() + cls.engine = Recording(cls.tok, "\nhm\n\n\nok", max_context=8192) + cls.svc = Service(cls.engine, cls.tok, ChatTemplate(ROOT / "serve/chat_template.jinja")) + cls.httpd = serve(cls.svc, port=0) + cls.base = f"http://127.0.0.1:{cls.httpd.server_address[1]}" + + @classmethod + def tearDownClass(cls): + cls.httpd.shutdown() + cls.httpd.server_close() + + def setUp(self): + self.svc.model_arch = "qwen4exp" + self.svc.sampling_defaults = {} + self.svc.shared = {} + self.svc.config_effort = None + + def tearDown(self): + self.svc.model_arch = None + + def keys(self, **body): + body = {"model": "m", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 8, **body} + self.engine.sampling = {} + req = urllib.request.Request(self.base + "/v1/chat/completions", data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"}) + urllib.request.urlopen(req, timeout=10).read() + return StrataEngine.sampling_keys(self.engine.sampling) + + def test_tables(self): + self.assertEqual(model_card_sampling("qwen4exp", True), {"temperature": 1.0, "top_p": 0.95, "top_k": 20}) + self.assertEqual(model_card_sampling("qwen4exp", False), + {"temperature": 0.7, "top_p": 0.8, "top_k": 20, "presence_penalty": 1.5}) + self.assertEqual(model_card_sampling("qwen35moe", True)["presence_penalty"], 1.5) + self.assertEqual(model_card_sampling(None, True), {}) + self.assertEqual(model_card_sampling("unknown", False), {}) + + def test_thinking_default(self): + k = self.keys() + self.assertIn("temperature=1.0", k) + self.assertIn("top_p=0.95", k) + self.assertIn("top_k=20", k) + self.assertNotIn("penalty_present", k) + + def test_non_thinking_default(self): + k = self.keys(reasoning_effort="none") + self.assertIn("temperature=0.7", k) + self.assertIn("top_p=0.8", k) + self.assertIn("penalty_present=1.5", k) + + def test_request_wins_and_zero_is_greedy(self): + k = self.keys(temperature=0) + self.assertNotIn("temperature", k) + self.assertIn("top_k=20", k) + k = self.keys(temperature=0.3, top_p=0.5) + self.assertIn("temperature=0.3", k) + self.assertIn("top_p=0.5", k) + + def test_config_block_and_shared_settings_win_over_the_card(self): + self.svc.sampling_defaults = {"temperature": 0.6, "top_k": 40} + k = self.keys() + self.assertIn("temperature=0.6", k) + self.assertIn("top_k=40", k) + self.assertIn("top_p=0.95", k) + self.svc.shared = {"temperature": 0.2} + self.assertIn("temperature=0.2", self.keys()) + self.svc.sampling_defaults = {"temperature": 0} + self.svc.shared = {} + self.assertNotIn("temperature", self.keys()) # a config that asks for greedy gets it + + def test_no_family_no_defaults(self): + self.svc.model_arch = None + self.assertEqual(self.keys(), "") + + +if __name__ == "__main__": + unittest.main() diff --git a/serve/test_parallel.py b/serve/test_parallel.py index 1a638919..bde52084 100644 --- a/serve/test_parallel.py +++ b/serve/test_parallel.py @@ -248,6 +248,34 @@ class PickSlot(unittest.TestCase): e.slot_busy = [False] * 3 self.assertEqual(e.pick_slot([1, 2, 3]), 1) # a held prompt is never the WHOLE prompt + def test_a_rewritten_next_turn_goes_to_the_slot_of_its_chat(self): + # #2518: two chats in two slots; each next turn renders the history again but not the reply it decoded or the + # assistant header, so the held stream (prompt + reply) is no start of it. The engine's turn checkpoint sits at + # the last <|im_start|> of the old prompt: the prompt that agrees up to there belongs to that slot. Before, both + # went to the least recently used slot in turn, and the first one overwrote the other chat's cache. + T = StrataEngine.TURN_ID + hist_a, hist_b = list(range(1000, 1600)), list(range(2000, 2600)) + held_a = hist_a + [T, 7, 8, 9] + [50, 51, 52] # prompt (ends with the assistant header) + reply + held_b = hist_b + [T, 7, 8, 9] + [60, 61] + e = self.engine(2) + e.slot_held = [held_a, held_b] + e.slot_used = [10.0, 20.0] # slot 0 was used longer ago + next_a = hist_a + [T, 5, 6, 7, 8, 9] + list(range(3000, 3100)) # the history again, then new turns + next_b = hist_b + [T, 5, 6, 7, 8, 9] + list(range(3000, 3100)) + self.assertEqual(e.pick_slot(next_b), 1) # the LRU slot 0 holds the other chat + self.assertEqual(e.pick_slot(next_a), 0) + e.slot_busy[0] = True # a busy slot is never picked + self.assertEqual(e.pick_slot(next_a), 1) + e.slot_busy[0] = False + # a chat that agrees with neither (a new one, or one that rewrote its history) keeps the old order + self.assertEqual(e.pick_slot(list(range(5000, 5200))), 0) + self.assertEqual(e.pick_slot(hist_a[:300] + [T, 5, 6] + list(range(3000, 3100))), 0) + e.slot_held[0] = [] # an empty slot still comes before evicting one + self.assertEqual(e.pick_slot(list(range(5000, 5200))), 0) + # the whole held stream continued is still the longer match + e.slot_held = [held_a, held_b + [T, 1, 2]] + self.assertEqual(e.pick_slot(held_b + [T, 1, 2, 3, 4]), 1) + def test_batch_groups_spread_requests_over_the_groups(self): # #1249: with --batch-groups 2 over 4 slots (groups {0,1} and {2,3}) a second request goes to the other group, # and a short held prefix does not pull it back into the busy one diff --git a/serve/test_parser_call_scan.py b/serve/test_parser_call_scan.py new file mode 100644 index 00000000..a4484514 --- /dev/null +++ b/serve/test_parser_call_scan.py @@ -0,0 +1,61 @@ +"""Wrapped-call precheck parity against the original structural walk.""" +import itertools +import random +import unittest +from unittest.mock import patch + +from serve import frontend + + +class ReferenceParser(frontend.OutputParser): + def _wrapped_call_end(self): + return frontend.call_end(self.buf) + + +def events_without_ids(events): + return [(e.kind, e.text, (e.call.name, e.call.arguments) if e.call else None) for e in events] + + +class CallScanTests(unittest.TestCase): + def test_chunk_boundaries_literal_tags_multiple_calls_and_recovery(self): + def call(value): + return "" + value + "" + bodies = [call("ordinary"), call("literal text here"), + call("literal prose"), call("a") + call("b"), + call("unfinished")[:-5], call("incomplete prose")[:-30], + "", + "bare", + "a" + "b"] + tools = [{"name": "write", "parameters": {"properties": {"text": {"type": "string"}}}}] + rng = random.Random(1030) + for text, stream, recover in itertools.product(bodies, (False, True), (False, True)): + # Every fixed chunk width, plus random partitions; compare after EACH feed. + for width in range(1, 16): + a = frontend.OutputParser(thinking=False, tools=tools, stream_tools=stream, recover=recover) + b = ReferenceParser(thinking=False, tools=tools, stream_tools=stream, recover=recover) + pos = 0 + while pos < len(text): + end = pos + (rng.randint(1, 20) if width == 15 else width) + delta = text[pos:end] + self.assertEqual(events_without_ids(a.feed(delta)), events_without_ids(b.feed(delta))) + pos = end + self.assertEqual(events_without_ids(a.finish()), events_without_ids(b.finish())) + + def test_no_structural_rescans_until_closer_and_reset_for_next_call(self): + parser = frontend.OutputParser(thinking=False) + with patch.object(frontend, "call_end", wraps=frontend.call_end) as walk: + for _ in range(2): + parser.feed("") + before = walk.call_count + for _ in range(100): + parser.feed("abcdefgh") + self.assertEqual(walk.call_count, before) + for char in "": + parser.feed(char) + self.assertEqual(walk.call_count, before + 1) + self.assertEqual(parser._call_end_checked, 0) + + +if __name__ == "__main__": + unittest.main() diff --git a/serve/test_parser_hold.py b/serve/test_parser_hold.py new file mode 100644 index 00000000..519f8192 --- /dev/null +++ b/serve/test_parser_hold.py @@ -0,0 +1,60 @@ +"""Partial-tag parity, including overlapping prefixes and arbitrary chunk boundaries.""" +import itertools +import random +import unittest + +from serve.frontend import OutputParser + + +def reference_hold(text, tags): + return max((n for tag in tags for n in range(1, len(tag)) if text.endswith(tag[:n])), default=0) + + +class ReferenceParser(OutputParser): + def _hold(self, text, tags): + return reference_hold(text, tags) + + +class HoldTests(unittest.TestCase): + def test_overlapping_and_empty_tags(self): + parser = OutputParser() + tags = ("", "a", "aaa", "abab", "baab", "", "") + for size in range(8): + for chars in itertools.product("ab<", repeat=size): + text = "".join(chars) + self.assertEqual(parser._hold(text, tags), reference_hold(text, tags), text) + + def test_every_real_tag_split(self): + tags = ("", "", "") + parser = OutputParser() + for tag in tags: + for n in range(len(tag) + 1): + for prefix in ("", "prose\n", "<", "<<<", "🙂日本語"): + text = prefix + tag[:n] + self.assertEqual(parser._hold(text, tags), reference_hold(text, tags), text) + + def test_streamed_events_match_reference(self): + texts = ["Reasoning.\n\nAnswer café 🙂\n", + "Quoted `` text and prose.", + "\n```xml\n\n```\nAnswer\nabc\nDone", + "abcDone"] + rng = random.Random(1028) + for text, thinking, recover, stream_tools in itertools.product(texts, (False, True), (False, True), (False, True)): + for _ in range(5): + args = dict(thinking=thinking, tools=[{"name": "lookup"}], recover=recover, stream_tools=stream_tools) + actual, expected = OutputParser(**args), ReferenceParser(**args) + def normalized(events): + # Tool call IDs are deliberately random; compare their semantic contents. + return [(e.kind, e.text, (e.call.name, e.call.arguments) if e.call else None) for e in events] + pos = 0 + while pos < len(text): + end = pos + rng.randint(1, 9) + delta = text[pos:end] + self.assertEqual(normalized(actual.feed(delta)), normalized(expected.feed(delta))) + pos = end + self.assertEqual(normalized(actual.finish()), normalized(expected.finish())) + + +if __name__ == "__main__": + unittest.main() diff --git a/serve/test_prompt_encoder.py b/serve/test_prompt_encoder.py new file mode 100644 index 00000000..36cc06f2 --- /dev/null +++ b/serve/test_prompt_encoder.py @@ -0,0 +1,162 @@ +"""serve/test_prompt_encoder.py - #567's incremental prompt encoding with #537's plain spans: every prompt's ids +are those of a full encode, and a turn reuses the previous one's. + + python -m unittest serve.test_prompt_encoder -v +""" +from __future__ import annotations + +import random +import re +import sys +import unittest +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) +sys.path.insert(0, str(ROOT / "tools")) +from serve.frontend import ChatTemplate # noqa: E402 +from serve.server import ByteTokenizer, MockEngine, Service # noqa: E402 +from strata_tokenizer import PromptEncoder, plain_agree_until # noqa: E402 + + +class ThinkTokenizer(ByteTokenizer): + """ByteTokenizer whose / are specials matched even without parse_special (type 4), as in the + Qwen vocabulary: the case #537's plain spans exist for.""" + SPECIALS = ByteTokenizer.SPECIALS + ["", ""] + ALWAYS = ("", "") + + +PIECES = ["<|im_start|>user\n", "<|im_end|>\n", "", "", "hello ", "é你 ", "<|im_sta", "x" * 7, + "<|im_start|>assistant\n", "", "<|vision_start|>", "<|image_pad|>"] + + +def spans_of(text, rng): + """Random plain spans over the specials' texts (, , <|im_start|>, ...), as unmark_think_literals + marks the ones quoted inside a message.""" + matches = re.finditer("|".join(map(re.escape, ThinkTokenizer.SPECIALS)), text) + return tuple((m.start(), m.end()) for m in matches if rng.random() < 0.5) # finditer: left to right, no overlap + + +class Units(unittest.TestCase): + def test_plain_agree_until(self): + self.assertEqual(plain_agree_until((), (), 50), 50) + self.assertEqual(plain_agree_until(((5, 13),), ((5, 13),), 50), 50) + self.assertEqual(plain_agree_until(((5, 13),), (), 50), 5) + self.assertEqual(plain_agree_until(((5, 13),), ((5, 9),), 50), 9) + self.assertEqual(plain_agree_until(((60, 70),), (), 50), 50) # beyond the shared text: no matter + self.assertEqual(plain_agree_until(((5, 13), (20, 28)), ((5, 13),), 50), 20) + + def test_same_ids_as_a_full_encode(self): + tok = ThinkTokenizer() + rng = random.Random(7) + for conv in range(300): + enc = PromptEncoder(tok) + text, plain = "", () + for turn in range(8): + text += "".join(rng.choice(PIECES) for _ in range(rng.randint(1, 12))) + if rng.random() < 0.6: + plain = spans_of(text, rng) # the spans may change in the shared part too + with self.subTest(conv=conv, turn=turn): + self.assertEqual(enc.encode(text, plain), tok.encode(text, True, plain)) + if rng.random() < 0.2 and len(text) > 10: # an edit: a client re-rendering its history + cut = rng.randrange(len(text)) + text = text[:cut] + plain = tuple((a, b) for a, b in plain if b <= cut) + + def test_a_turn_reuses_the_previous_one(self): + tok = ThinkTokenizer() + enc = PromptEncoder(tok) + # as a chat turn: the earlier prompt goes on past its last special (the answer being written), and the + # next one shares it up to there - the resume point needs max_special_len characters after it + base = "<|im_start|>user\n" + "quoted here " * 200 + "<|im_end|>\n<|im_start|>assistant\nan answer" + plain = tuple((i, i + 8) for i in range(len(base)) if base.startswith("", i)) + enc.encode(base, plain) + more = base + " that goes on<|im_end|>\n<|im_start|>user\nnext<|im_end|>\n" + self.assertEqual(enc.encode(more, plain), tok.encode(more, True, plain)) + self.assertGreater(enc.last_reused, len(base) // 2) + # the same text with the quoted no longer plain: nothing before it may be reused past its start + self.assertEqual(enc.encode(more, ()), tok.encode(more, True, ())) + self.assertLessEqual(enc.last_reused, plain[0][0]) + + +class Served(unittest.TestCase): + def test_quoted_think_across_turns(self): + tok = ThinkTokenizer() + template = ChatTemplate(ROOT / "serve/chat_template.jinja") + svc = Service(MockEngine(tok, "\n\nok", max_context=1 << 20), tok, template) + full = Service(MockEngine(tok, "\n\nok", max_context=1 << 20), tok, template) + full.prompts = None # the reference: a full encode every time + self.assertIsNotNone(svc.prompts) + msgs = [{"role": "system", "content": "Explain tags."}] + for turn in range(5): + msgs.append({"role": "user", "content": f"turn {turn}: what does mean? too " * 3}) + got, want = svc.encode_prompt(msgs, None, {}), full.encode_prompt(msgs, None, {}) + self.assertEqual(got, want) + if turn: + self.assertGreater(svc.prompts.last_reused, 0) + msgs.append({"role": "assistant", "content": "It closes the reasoning: is a marker."}) + + def service(self, max_context=1 << 22): + tok = ByteTokenizer() + return Service(MockEngine(tok, "ok", max_context=max_context), tok, ChatTemplate(ROOT / "serve/chat_template.jinja")) + + def test_another_conversation_does_not_evict_the_first(self): + """A subagent beside the main conversation: every rendered prompt starts with the template's header, which + must not make the second client's prompt a continuation of the first's entry.""" + svc = self.service() + turn2 = [{"role": "user", "content": "a long question " * 500}, {"role": "assistant", "content": "answer"}, + {"role": "user", "content": "next"}] + svc.encode_prompt(turn2, None, {}) + svc.encode_prompt([{"role": "user", "content": "one unrelated prompt " * 300}], None, {}) + self.assertEqual(len(svc.prompts.entries), 2) + turn3 = turn2 + [{"role": "assistant", "content": "x"}, {"role": "user", "content": "more"}] + svc.encode_prompt(turn3, None, {}) + self.assertGreater(svc.prompts.last_reused, len(turn2[0]["content"])) + self.assertEqual(len(svc.prompts.entries), 2) # turn 3 replaced turn 2, the other one stayed + + def test_a_refused_prompt_is_not_kept(self): + svc = self.service(max_context=4096) + msgs = [{"role": "user", "content": "kept " * 20}] + svc.prepare(msgs, None, {}, 16) + refused = msgs + [{"role": "assistant", "content": "a"}, {"role": "user", "content": "more " * 20}] + full = self.service(max_context=4096) + full.prompts = None + refused_ids = full.encode_prompt(refused, None, {}) + with self.assertRaisesRegex(ValueError, "exceeds the context"): + svc.prepare(refused, None, {}, 1 << 20) + self.assertNotIn(refused_ids, [e[2] for e in svc.prompts.entries]) + self.assertNotIn(full.render_prompt(refused, None, {}), [e[0] for e in svc.prompts.entries]) + + def test_random_messages_equal_a_full_encode(self): + """Messages quoting control-token texts and tags at random, growing, edited and cut back between + requests: the incremental encoder's ids are the full encode's, 2400 prompts.""" + tok = ThinkTokenizer() + template = ChatTemplate(ROOT / "serve/chat_template.jinja") + svc = Service(MockEngine(tok, "ok", max_context=1 << 20), tok, template) + full = Service(MockEngine(tok, "ok", max_context=1 << 20), tok, template) + full.prompts = None + rng = random.Random(567) + quoted = ThinkTokenizer.SPECIALS + ["<|im_sta", " 2: # a shorter conversation + msgs = msgs[:rng.randrange(1, len(msgs))] + msgs[-1:] + with self.subTest(conv=conv, turn=turn): + self.assertEqual(svc.encode_prompt(msgs, None, {}), full.encode_prompt(msgs, None, {})) + reused += svc.prompts.last_reused > 0 + msgs.append({"role": "assistant", "content": text()}) + self.assertGreater(reused, 1000) + + +if __name__ == "__main__": + unittest.main() diff --git a/serve/test_q36_config.py b/serve/test_q36_config.py new file mode 100644 index 00000000..981ae5f3 --- /dev/null +++ b/serve/test_q36_config.py @@ -0,0 +1,48 @@ +"""serve/test_q36_config.py - what the server hands the Qwen3.6-35B-A3B engine (arch qwen35moe): the config's pictures +("vision" section -> --vision) and parallel requests ("parallel": N -> --batch N), as for Flash-Next's engine. + + python -m unittest serve.test_q36_config +""" +from __future__ import annotations + +import unittest +from unittest import mock + +from serve.server import engine_args, parallel_args + + +def cfg(**kw): + base = {"arch": "qwen35moe", "exe": "strata", "args": ["--native", "/m/Qwen3.6-35B-A3B-UD-Q4_K_XL.gguf", + "--max-context", "32768", "--prefill", "auto", "--mtp", "2"]} + base.update(kw) + return base + + +class Q36Config(unittest.TestCase): + def test_parallel_becomes_batch(self): + a = engine_args(cfg(parallel=4)) + self.assertEqual(a[a.index("--batch") + 1], "4") + self.assertNotIn("--batch", engine_args(cfg())) + self.assertNotIn("--batch", engine_args(cfg(parallel=1))) + + def test_more_than_the_window_is_said(self): + with mock.patch("builtins.print") as pr: + a = engine_args(cfg(parallel=12)) + self.assertEqual(a[a.index("--batch") + 1], "12") # the engine caps it at 8 and says so + self.assertTrue(any("at most 8" in str(c) for c in pr.call_args_list)) + + def test_a_vision_section_starts_the_engine_with_pictures(self): + v = {"exe": "strata-vision", "mmproj": "/m/mmproj-F16.gguf", "model": "/m/Qwen3.6-35B-A3B-UD-Q4_K_XL.gguf", + "gpu": False, "max_tokens": 300} + self.assertIn("--vision", engine_args(cfg(vision=v))) + self.assertIn("--vision", engine_args(cfg(vision=v, args=cfg()["args"] + ["--vision"]))) + self.assertEqual(engine_args(cfg(vision=v, args=cfg()["args"] + ["--vision"])).count("--vision"), 1) + self.assertNotIn("--vision", engine_args(cfg())) + + def test_an_explicit_batch_in_the_args_wins(self): + c = cfg(parallel=4, args=cfg()["args"] + ["--batch", "2"]) + self.assertEqual(parallel_args(c, c["args"]), []) + + +if __name__ == "__main__": + unittest.main() diff --git a/serve/test_reasoning_loop_recovery.py b/serve/test_reasoning_loop_recovery.py index b53ee56a..c00fc1be 100644 --- a/serve/test_reasoning_loop_recovery.py +++ b/serve/test_reasoning_loop_recovery.py @@ -58,6 +58,7 @@ class ReasoningLoopRecovery(unittest.TestCase): engine = RecordingEngine(self.tok, [self.loop, 'I have verified the solution.\n\nComplete solution.']) service = Service(engine, self.tok, None) + service.answer_room = False # these cases test the loop guards at the cap itself service.reasoning_loop_recovery = enabled events = list(service.run(self.ids, True, [], 40000, {}, threading.Event())) return engine, events @@ -66,6 +67,7 @@ class ReasoningLoopRecovery(unittest.TestCase): engine = RecordingEngine(self.tok, [self.loop, 'Verified.\n\nComplete solution.']) service = Service(engine, self.tok, None) + service.answer_room = False # these cases test the loop guards at the cap itself service.reasoning_loop_recovery = True request = {'temperature': .2, 'presence_penalty': 0, 'top_p': .95, 'top_k': 64, 'seed': 2} list(service.run(self.ids, True, [], 40000, request, threading.Event())) @@ -81,6 +83,7 @@ class ReasoningLoopRecovery(unittest.TestCase): def test_frozen_cohort_sampler_is_unchanged_during_recovery(self): engine = RecordingEngine(self.tok, [self.loop, 'Done.\n\nComplete solution.']) service = Service(engine, self.tok, None) + service.answer_room = False # these cases test the loop guards at the cap itself service.reasoning_loop_recovery = True sampler = {'temperature': 1.0, 'top_p': .95, 'top_k': 20, 'presence_penalty': 1.5, 'seed': 2} list(service.run(self.ids, True, [], 40000, sampler, threading.Event())) @@ -91,6 +94,7 @@ class ReasoningLoopRecovery(unittest.TestCase): def test_a_lower_top_p_and_top_k_are_never_raised_or_lowered(self): engine = RecordingEngine(self.tok, [self.loop, 'Verified.\n\nComplete solution.']) service = Service(engine, self.tok, None) + service.answer_room = False # these cases test the loop guards at the cap itself service.reasoning_loop_recovery = "recover" list(service.run(self.ids, True, [], 40000, {'temperature': 0, 'top_p': .5, 'top_k': 5}, threading.Event())) self.assertEqual((engine.sampling_calls[1]['top_p'], engine.sampling_calls[1]['top_k']), (.5, 5)) @@ -110,6 +114,7 @@ class ReasoningLoopRecovery(unittest.TestCase): def test_stop_mode_ends_the_reply_at_the_loop(self): engine = RecordingEngine(self.tok, [self.loop]) service = Service(engine, self.tok, None) + service.answer_room = False # these cases test the loop guards at the cap itself service.reasoning_loop_recovery = "stop" events = list(service.run(self.ids, True, [], 40000, {}, threading.Event())) self.assertEqual(len(engine.calls), 1) # no second pass @@ -121,6 +126,7 @@ class ReasoningLoopRecovery(unittest.TestCase): """#1753: the word count never sees the digits, so only the tail's period ends the reply.""" engine = RecordingEngine(self.tok, ['9^104 = ' + '58747693387' * 3000]) service = Service(engine, self.tok, None) + service.answer_room = False # these cases test the loop guards at the cap itself service.reasoning_loop_recovery = "stop" events = list(service.run(self.ids, True, [], 40000, {}, threading.Event())) self.assertEqual(len(engine.calls), 1) # no second pass @@ -136,6 +142,7 @@ class ReasoningLoopRecovery(unittest.TestCase): def test_default_single_token_guard_does_not_detect_passage_loop(self): engine = RecordingEngine(self.tok, [self.loop]) service = Service(engine, self.tok, None) + service.answer_room = False # these cases test the loop guards at the cap itself events = list(service.run(self.ids, True, [], 16384, {}, threading.Event())) self.assertEqual(service.repeat_stop_tokens, 256) self.assertEqual(events[-1][1]['finish'], 'length') @@ -145,6 +152,7 @@ class ReasoningLoopRecovery(unittest.TestCase): def test_recovery_is_once_and_never_extends_output_limit(self): engine = RecordingEngine(self.tok, [self.loop, self.loop]) service = Service(engine, self.tok, None) + service.answer_room = False # these cases test the loop guards at the cap itself service.reasoning_loop_recovery = True events = list(service.run(self.ids, True, [], 16384, {}, threading.Event())) self.assertEqual(len(engine.calls), 2) @@ -155,6 +163,7 @@ class ReasoningLoopRecovery(unittest.TestCase): def test_repeated_final_text_does_not_change_sampler_or_prompt(self): engine = RecordingEngine(self.tok, [self.loop]) service = Service(engine, self.tok, None) + service.answer_room = False # these cases test the loop guards at the cap itself service.reasoning_loop_recovery = True events = list(service.run(self.ids, False, [], 40000, {}, threading.Event())) self.assertEqual(len(engine.calls), 1) @@ -163,6 +172,7 @@ class ReasoningLoopRecovery(unittest.TestCase): def test_missing_native_high_hint_skips_recovery(self): engine = RecordingEngine(self.tok, [self.loop]) service = Service(engine, self.tok, None) + service.answer_room = False # these cases test the loop guards at the cap itself service.reasoning_loop_recovery = True ids = self.tok.encode('<|im_start|>user\nTask<|im_end|>\n', parse_special=True) events = list(service.run(ids, True, [], 40000, {}, threading.Event())) @@ -172,6 +182,7 @@ class ReasoningLoopRecovery(unittest.TestCase): def test_cancel_during_reasoning_does_not_start_recovery(self): engine = RecordingEngine(self.tok, [self.loop, 'Unused']) service = Service(engine, self.tok, None) + service.answer_room = False # these cases test the loop guards at the cap itself service.reasoning_loop_recovery = True cancel = threading.Event() events = [] @@ -187,6 +198,7 @@ class ReasoningLoopRecovery(unittest.TestCase): engine = RecordingEngine(self.tok, [self.loop, 'Done.\n\nAnswer.']) engine.batch = 2 service = Service(engine, self.tok, None) + service.answer_room = False # these cases test the loop guards at the cap itself service.reasoning_loop_recovery = True events = list(service.run(self.ids, True, [], 40000, {}, threading.Event())) self.assertEqual(events[-1][1]['reasoning_recoveries'], 1) diff --git a/serve/test_server.py b/serve/test_server.py index 3a49ff25..5f58d09d 100644 --- a/serve/test_server.py +++ b/serve/test_server.py @@ -26,7 +26,7 @@ from serve.server import (CTX_SLACK, ByteTokenizer, EngineDied, EngineNonFinite, PP_DONE_TAIL, Service, StrataEngine, api_key_of, engine_args, key_matches, layer_split_value, prompt_progress, prompt_tokens_seen, - request_timings, serve, start_failure_hint) + request_stats, request_timings, serve, start_failure_hint) from types import SimpleNamespace # noqa: E402 ROOT = Path(__file__).resolve().parents[1] @@ -1548,6 +1548,13 @@ class SamplingKeys(unittest.TestCase): bad = self.keys(strata_tune={"pcie_frac": 3, "spec_min_p": True, "pool_workers": 2}) self.assertFalse([x for x in bad if x.split("=")[0] in ("pcie_frac", "spec_min_p", "pool_workers")]) + def test_tune_prefill_pipe(self): + self.assertIn("pipe=1536.0", self.keys(temperature=0, strata_tune={"prefill_pipe": 1536})) + self.assertIn("pipe=0.0", self.keys(temperature=0, strata_tune={"prefill_pipe": 0})) + for bad in (-1, True, "800", 1e9): + self.assertFalse([x for x in self.keys(temperature=0, strata_tune={"prefill_pipe": bad}) + if x.startswith("pipe=")], bad) + def test_tune_aux_cpus(self): for val, want in ((True, "aux_cpus=1"), (1, "aux_cpus=1"), (False, "aux_cpus=0"), (0, "aux_cpus=0"), (1.0, "aux_cpus=1"), (0.0, "aux_cpus=0")): @@ -1771,6 +1778,150 @@ class RecordingPrompt(MockEngine): yield from super().generate(ids, max_new, sampling, cancel, embeddings) +class IncrementalPrompts(unittest.TestCase): + """The prompt encoder: every request's ids are those of a full encode, and a turn reuses the previous one's.""" + + @classmethod + def setUpClass(cls): + tok = ByteTokenizer() + cls.engine = RecordingPrompt(tok, "Thinking.\n\n\nThe answer.", max_context=CTX) + cls.svc = Service(cls.engine, tok, ChatTemplate(ROOT / "serve/chat_template.jinja")) + cls.httpd = serve(cls.svc, port=0) + cls.base = f"http://127.0.0.1:{cls.httpd.server_address[1]}" + + @classmethod + def tearDownClass(cls): + cls.httpd.shutdown() + cls.httpd.server_close() + + post = ClientShapes.post + + def test_turns(self): + self.assertIsNotNone(self.svc.prompts) + full = Service(self.engine, self.svc.tok, self.svc.template) + full.prompts = None # the reference: the same prompt, encoded in full + msgs = [{"role": "system", "content": "Be brief."}] + for turn in range(6): + msgs.append({"role": "user", "content": f"question {turn} <|im_end|> é你 " * (turn + 1)}) + status, b = self.post("/v1/chat/completions", {"model": "m", "max_tokens": 64, "messages": msgs}) + self.assertEqual(status, 200, b) + prompt = self.svc.template.render(msgs) + self.assertEqual(self.engine.last_ids, full.encode_prompt(msgs, None, {})) + if turn: + self.assertGreater(self.svc.prompts.last_reused, len(prompt) // 3) + msgs.append({"role": "assistant", "content": b["choices"][0]["message"]["content"]}) + + def conversation(self, n): + """Messages of the n-th conversation: distinct from every other one from the number on.""" + return [{"role": "user", "content": f"conversation {n} says: " + f"topic {n} sentence {n} " * 20}] + + def service(self, **info): + engine = RecordingPrompt(self.svc.tok, "The answer.", max_context=CTX) + engine.info = info + return Service(engine, self.svc.tok, self.svc.template) + + def test_the_kept_prompts_follow_the_engines_conversation_slots(self): + svc = self.service(conversation_cache_slots=8) + texts = [] + for n in range(8): + msgs = self.conversation(n) + svc.encode_prompt(msgs, None, {}) + texts.append(svc.render_prompt(msgs, None, {})) + self.assertEqual(len(svc.prompts.entries), 8) + self.assertEqual(sorted(e[0] for e in svc.prompts.entries), sorted(texts)) + msgs = self.conversation(8) + svc.encode_prompt(msgs, None, {}) + kept = sorted(e[0] for e in svc.prompts.entries) + self.assertEqual(len(kept), 8) + self.assertEqual(kept, sorted(texts[1:] + [svc.render_prompt(msgs, None, {})])) + + def lazy_service(self, args, **info): + """A Service over the real StrataEngine class started lazily: no process yet, so no INFO line either.""" + engine = StrataEngine("strata", args, lazy=True) + engine.info.update(info) + return Service(engine, self.svc.tok, self.svc.template) + + def test_a_lazy_engine_keeps_the_prompts_its_arguments_ask_for(self): + svc = self.lazy_service(["--conversation-cache-slots", "8"]) + for n in range(10): + svc.encode_prompt(self.conversation(n), None, {}) + self.assertEqual(len(svc.prompts.entries), 8) + + def test_the_slot_argument_wins_over_the_info_line_and_a_bad_one_falls_back(self): + self.assertEqual(self.lazy_service(["--conversation-cache-slots", "6"], conversation_cache_slots=3) + .prompts.keep, 6) + self.assertEqual(self.lazy_service(["--conversation-cache-slots", "x"], conversation_cache_slots=3) + .prompts.keep, 3) + self.assertEqual(self.lazy_service(["--conversation-cache-slots"]).prompts.keep, 4) + self.assertEqual(self.lazy_service([]).prompts.keep, 4) + + def test_prompt_encoder_slot_count_is_capped(self): + from serve.server import PROMPT_ENCODER_KEEP_MAX + self.assertEqual(PROMPT_ENCODER_KEEP_MAX, 16) + self.assertEqual(self.service(conversation_cache_slots=10**9).prompts.keep, 16) + self.assertEqual(self.lazy_service(["--conversation-cache-slots", str(10**9)]).prompts.keep, 16) + + def test_without_a_slot_setting_four_prompts_are_kept(self): + svc = self.service() + for n in range(6): + svc.encode_prompt(self.conversation(n), None, {}) + self.assertEqual(len(svc.prompts.entries), 4) + + def test_counting_tokens_leaves_the_kept_prompts_alone(self): + svc = self.service(conversation_cache_slots=2) + httpd = serve(svc, port=0) + self.addCleanup(httpd.server_close) + self.addCleanup(httpd.shutdown) + self.base = f"http://127.0.0.1:{httpd.server_address[1]}" + main = self.conversation(0) + status, b = self.post("/v1/chat/completions", {"model": "m", "max_tokens": 8, "messages": main}) + self.assertEqual(status, 200, b) + before = list(svc.prompts.entries) + full = self.service() + full.prompts = None + for n in range(1, 6): + other = self.conversation(n) + status, b = self.post("/v1/messages/count_tokens", {"model": "m", "messages": other}) + self.assertEqual(status, 200, b) + self.assertEqual(b["input_tokens"], len(full.encode_prompt(other, None, {}))) + self.assertTrue(svc.prompts.entries == before, "the count wrote into the kept prompts") + main.append({"role": "assistant", "content": "ok"}) + main.append({"role": "user", "content": "and then?"}) + svc.encode_prompt(main, None, {}) + self.assertGreater(svc.prompts.last_reused, len(svc.render_prompt(main, None, {})) // 2) + + def test_same_ids_as_a_full_encode_on_varied_conversations(self): + import random + sys.path.insert(0, str(ROOT / "tools")) + from test_strata_tokenizer import conversation_prompts, load_tokenizer + toks = [("byte", ByteTokenizer())] + qwen = load_tokenizer() + if qwen is not None: + toks.append(("qwen35", qwen)) + for name, tok in toks: + svc = Service(self.engine, tok, self.svc.template) + for seed in range(3): + for what, prompt in conversation_prompts(self.svc.template, random.Random(seed)): + with self.subTest(tokenizer=name, seed=seed, what=what): + self.assertEqual(svc.prompts.encode(prompt), tok.encode(prompt, parse_special=True)) + + def test_a_tokenizer_without_resume_points_encodes_in_full(self): + class Plain: + """ByteTokenizer without the resume points.""" + def __init__(self): + self.base = ByteTokenizer() + + def __getattr__(self, name): + if name in ("encode_marked", "max_special_len"): + raise AttributeError(name) + return getattr(self.base, name) + svc = Service(self.engine, Plain(), self.svc.template) + self.assertIsNone(svc.prompts) + msgs = [{"role": "user", "content": "hi"}] + self.assertEqual(svc.encode_prompt(msgs, None, {}), + ByteTokenizer().encode(svc.render_prompt(msgs, None, {}), parse_special=True)) + + class DyingEngine(MockEngine): """Issue #27: an engine that dies after a few tokens of its first answer, and comes back when restarted.""" @@ -1989,6 +2140,39 @@ class PrometheusMetrics(unittest.TestCase): self.assertEqual((m["strata:requests_total"], m["strata:drafts_offered_total"], m["strata:drafts_accepted_total"]), (2, 17, 10)) + def test_throughput_gauges_cover_the_window_between_scrapes(self): + """#2341: llama.cpp reads these two from a bucket it resets on every scrape, so a scrape that saw no + traffic reports 0 - not the lifetime average, which froze at the last request's speed.""" + tok = ByteTokenizer() + engine = DoneLineEngine(tok, "\n\nok", max_context=CTX, + done_lines=["DONE 4 20 400.0 200.0 stop 7 12 5"]) + svc = Service(engine, tok, ChatTemplate(ROOT / "serve/chat_template.jinja")) + httpd = serve(svc, port=0) + base = f"http://127.0.0.1:{httpd.server_address[1]}" + + def scrape(): + with urllib.request.urlopen(base + "/metrics/prometheus", timeout=10) as r: + return {line.split()[0]: float(line.split()[1]) + for line in r.read().decode().splitlines() if not line.startswith("#")} + + try: + body = json.dumps({"model": "m", "max_tokens": 10, "messages": [{"role": "user", "content": "hi"}]}) + with urllib.request.urlopen(urllib.request.Request(base + "/v1/chat/completions", data=body.encode(), + headers={"Content-Type": "application/json"}), + timeout=30) as r: + self.assertEqual(r.status, 200) + first = scrape() + self.assertGreater(first["llamacpp:prompt_tokens_seconds"], 0) + self.assertGreater(first["llamacpp:predicted_tokens_seconds"], 0) + idle = scrape() # nothing ran since the scrape above, so the window is empty + self.assertEqual((idle["llamacpp:prompt_tokens_seconds"], + idle["llamacpp:predicted_tokens_seconds"]), (0, 0)) + self.assertEqual(idle["llamacpp:tokens_predicted_total"], # counters stay cumulative + first["llamacpp:tokens_predicted_total"]) + finally: + httpd.shutdown() + httpd.server_close() + class PcieShare(unittest.TestCase): """#588: the hit rate stays the VRAM share of the lookups; the routed experts the GPU read over PCIe (the DONE @@ -2507,6 +2691,30 @@ class CancelledRead(unittest.TestCase): StrataEngine._parse_done(e, "DONE 0 98179 17565.0 0.0 cancel 0 0 0 0 0 0 0 0.0") self.assertNotIn("prompt_read", e.last) + def test_parse_done_chunk_fields(self): + """The prompt's first chunk and its streamed expert blobs come as key=value fields after the positional ones, + with or without the five --lookup-chain fields; an older line has neither.""" + e = SimpleNamespace() + StrataEngine._parse_done(e, "DONE 5 12000 900.0 80.0 stop 3 4 0 7 9 0 0 0.0 12000 2 chunk=1536 " + "experts_streamed=24576") + self.assertEqual((e.last["chunk"], e.last["experts_streamed"], e.last["offloaded"]), (1536, 24576, 2)) + StrataEngine._parse_done(e, "DONE 5 12000 900.0 80.0 stop 3 4 0 7 9 0 0 0.0 12000 2 1 2 3 4 5 chunk=8192 " + "experts_streamed=0") + self.assertEqual((e.last["chunk"], e.last["experts_streamed"]), (8192, 0)) + StrataEngine._parse_done(e, "DONE 5 12000 900.0 80.0 stop 3 4 0 7 9 0 0 0.0 12000 2") + self.assertNotIn("chunk", e.last) + self.assertNotIn("experts_streamed", e.last) + StrataEngine._parse_done(e, "DONE 5 12000 900.0 80.0 stop 3 4 0 7 9 0 0 0.0 12000 2 chunk=x other=1") + self.assertNotIn("chunk", e.last) + self.assertNotIn("other", e.last) + + def test_request_stats_chunk_fields(self): + """Across a request's continuations the streamed blobs add up; the chunk is the request's own read's.""" + stats = request_stats([{"prompt_ms": 10.0, "chunk": 1536, "experts_streamed": 100}, + {"prompt_ms": 2.0, "chunk": 64, "experts_streamed": 5}]) + self.assertEqual((stats["chunk"], stats["experts_streamed"], stats["prompt_ms"]), (1536, 105, 12.0)) + self.assertNotIn("chunk", request_stats([{"prompt_ms": 1.0}, {"prompt_ms": 1.0, "chunk": 64}])) + def test_prompt_tokens_seen(self): last = {"finish": "cancel", "reused": 1000, "prompt_read": 2000} self.assertEqual(prompt_tokens_seen(98179, last), 3000) @@ -2822,6 +3030,46 @@ class ConfigReasoningEffort(unittest.TestCase): self.assertEqual(cfg, {"args": ["--ctx", "4096"]}) +class JoinedEngineArgs(unittest.TestCase): + """#2267: "--ple-io mmap" written as one element of the config's args reached the engine as one argument.""" + + def split(self, args): + from serve.server import split_joined_args + cfg = {"args": list(args)} + with mock.patch("builtins.print") as pr: + split_joined_args(cfg) + return cfg["args"], pr.called + + def test_one_element_becomes_flag_and_value(self): + self.assertEqual(self.split(["--spec", "4", "--ple-io mmap"]), (["--spec", "4", "--ple-io", "mmap"], True)) + self.assertEqual(self.split(["--ple-io=mmap", "--kv int8"]), (["--ple-io", "mmap", "--kv", "int8"], True)) + + def test_proper_args_are_untouched(self): + args = ["--ple-io", "mmap", "--native", "/data/My Models/a b.gguf", "--max-context", "131072", "--vision"] + self.assertEqual(self.split(args), (args, False)) + # a value in a flag's value place stays whole, whatever it looks like + args = ["--system-prompt", "--be brief please"] + self.assertEqual(self.split(args), (args, False)) + + def test_no_args(self): + from serve.server import split_joined_args + cfg = {} + split_joined_args(cfg) + self.assertEqual(cfg, {}) + + def test_bad_json_names_the_place(self): + import tempfile + from serve.server import read_config + with tempfile.TemporaryDirectory() as d: + p = Path(d) / "c.json" + p.write_text('{"args": ["--ple-io" "mmap"]}', encoding="utf-8") + with self.assertRaises(SystemExit) as cm: + read_config(p) + self.assertIn("line 1", str(cm.exception)) + p.write_text('{"args": ["--ple-io", "mmap"]}', encoding="utf-8") + self.assertEqual(read_config(p), {"args": ["--ple-io", "mmap"]}) + + class SharedSettings(unittest.TestCase): """The web app's "Use for other apps too": POST /settings makes its Chat settings every client's defaults.""" @@ -3776,6 +4024,45 @@ class ThinkingBudget(unittest.TestCase): self.assertEqual(text, ThinkingEngine.ANSWER) self.assertEqual(evs[-2]["delta"]["stop_reason"], "end_turn") + def test_a_cap_that_the_thinking_would_eat_leaves_room_to_answer(self): + """0.1.45: max_tokens counts the thinking; a thinking that would fill it gets closed with a quarter left (the + agent harness sends max_completion_tokens 8192; the pagoda's thinking took all of it and no answer came).""" + self.engine.THOUGHT = "Let me think step by step about two plus two. " * 40 # 1,800 reasoning tokens + code, b = self.post("/v1/chat/completions", {"model": "m", "max_tokens": 1024, + "messages": [{"role": "user", "content": "2+2?"}]}) + self.assertEqual(code, 200, b) + msg = b["choices"][0]["message"] + self.assertEqual(msg["content"], ThinkingEngine.ANSWER) + self.assertEqual(b["choices"][0]["finish_reason"], "stop") + self.assertLess(len(msg["reasoning_content"]), 1024 - 256) + self.assertEqual(len(self.engine.prompts), 2) # the wrap-up continuation + + def test_the_answer_room_has_its_limits(self): + self.engine.THOUGHT = "Let me think step by step about two plus two. " * 40 + code, b = self.post("/v1/chat/completions", {"model": "m", "max_tokens": 300, # below the floor: as before + "messages": [{"role": "user", "content": "2+2?"}]}) + self.assertEqual(b["choices"][0]["finish_reason"], "length") + self.assertFalse(b["choices"][0]["message"]["content"]) + self.svc.answer_room = False # the config's "answer_room": false + code, b = self.post("/v1/chat/completions", {"model": "m", "max_tokens": 1024, + "messages": [{"role": "user", "content": "2+2?"}]}) + self.assertEqual(b["choices"][0]["finish_reason"], "length") + self.svc.answer_room = True + code, b = self.post("/v1/chat/completions", {"model": "m", "max_tokens": 1024, "reasoning_effort": "none", + "messages": [{"role": "user", "content": "2+2?"}]}) + self.assertEqual(len(self.engine.prompts), 3) # thinking off: no budget, one prompt each + + def test_anthropic_thinking_budget_tokens_is_a_budget(self): + """0.1.45: thinking {type: enabled, budget_tokens: N} closes the thinking after N tokens (it was only a level).""" + from serve.server import REASONING_WRAP_UP + code, b = self.post("/v1/messages", {"model": "m", "max_tokens": 400, "thinking": {"type": "enabled", "budget_tokens": 30}, + "messages": [{"role": "user", "content": "2+2?"}]}) + self.assertEqual(code, 200, b) + thinking = "".join(c.get("thinking", "") for c in b["content"] if c["type"] == "thinking") + self.assertEqual(thinking, ThinkingEngine.THOUGHT[:30] + REASONING_WRAP_UP.split("")[0]) + self.assertEqual("".join(c.get("text", "") for c in b["content"] if c["type"] == "text"), ThinkingEngine.ANSWER) + self.assertEqual(len(self.engine.prompts), 2) + def test_a_budget_the_thinking_stays_under(self): code, b = self.openai(reasoning_budget_tokens=10_000) self.assertEqual(b["choices"][0]["message"]["reasoning_content"], ThinkingEngine.THOUGHT) @@ -4481,6 +4768,114 @@ class AmdTelemetry(unittest.TestCase): self.assertEqual(svc.free_vram_mib(), 26 << 10) +IOREG_FIXTURE = ''' + +-o AGXAcceleratorG13X + { + "IOClass" = "AGXAcceleratorG13X" + "model" = "Apple M1 Max" + "PerformanceStatistics" = {"In use system memory (driver)"=0,"Alloc system memory"=22991831040,"Tiler Utilization %"=36,"recoveryCount"=44,"lastRecoveryTime"=216970984728,"Renderer Utilization %"=35,"TiledSceneBytes"=1671168,"Device Utilization %"=76,"SplitSceneCount"=0,"Allocated PB Size"=76152832,"In use system memory"=13384105984} + +-o IOUserClientCreator + { + "IOUserClientCreator" = "pid 402, WindowServer" + } + } + + +-o AGXAccelerator + { + "IOClass" = "AGXAccelerator" + "model" = "Apple M2 Ultra" + "PerformanceStatistics" = {"Renderer Utilization %"=3,"Tiler Utilization %"=2,"In use system memory"=1234567890} + } +''' + +WORKING_SET = 26800603136 # bytes; the M1 Max's device in the fixture above + + +class MacTelemetry(unittest.TestCase): + """The macOS (Metal backend) reader: IOKit's IOAccelerator nodes and the device's Metal working set.""" + + def _reader(self, index=0, text=IOREG_FIXTURE): + from serve import telemetry + with mock.patch.object(telemetry, "metal_working_set", return_value=WORKING_SET): + return telemetry._Metal(index, run=lambda: text) + + def test_readings(self): + g = self._reader() + self.assertTrue(g.ok()) + self.assertEqual(g.name(), "Apple M1 Max") + with mock.patch("serve.telemetry.metal_working_set", lambda: WORKING_SET): + self.assertEqual(g.read(), {"util": 76, "mem_used": 13384105984, "mem_total": WORKING_SET}) + # the second accelerator is index 1: the older chips' renderer/tiler-only utilization is the fallback + g1 = self._reader(1) + self.assertTrue(g1.ok()) + with mock.patch("serve.telemetry.metal_working_set", lambda: WORKING_SET): + self.assertEqual(g1.read()["util"], 3) # the renderer, the higher of the two + self.assertEqual(g1.name(), "Apple M2 Ultra") + self.assertFalse(self._reader(5).ok()) + + def test_working_set_is_read_once_for_samples_and_static(self): + from serve import telemetry + with mock.patch.object(telemetry, "metal_working_set", return_value=WORKING_SET) as budget: + g = telemetry._Metal(run=lambda: IOREG_FIXTURE) + for _ in range(3): + self.assertEqual(g.read()["mem_total"], WORKING_SET) + with mock.patch.object(telemetry, "gpu_reader", return_value=g): + t = telemetry.Telemetry() + try: + self.assertEqual(t.static["gpu_vram_total_mib"], WORKING_SET >> 20) + finally: + t.close() + budget.assert_called_once_with() + + def test_device_is_released_on_success_and_unsupported_selector(self): + from serve import telemetry + metal, objc = mock.MagicMock(), mock.MagicMock() + metal.MTLCreateSystemDefaultDevice.return_value = 1234 + objc.sel_registerName.return_value = 5678 + objc.objc_msgSend.return_value = WORKING_SET + for supported in (True, False): + with self.subTest(supported=supported): + objc.objc_release.reset_mock() + objc.class_respondsToSelector.return_value = supported + with mock.patch.object(telemetry.sys, "platform", "darwin"), \ + mock.patch.object(telemetry.ctypes, "CDLL", side_effect=[metal, objc]): + self.assertEqual(telemetry.metal_working_set(), WORKING_SET if supported else None) + objc.objc_release.assert_called_once_with(1234) + + def test_no_registry(self): + g = self._reader(text="") + self.assertFalse(g.ok()) + self.assertEqual(g.read(), {}) + self.assertIsNone(g.name()) + + def test_bad_process(self): + from serve import telemetry + g = telemetry._Metal(0, run=lambda: (_ for _ in ()).throw(OSError("no ioreg"))) + self.assertFalse(g.ok()) + self.assertEqual(g.read(), {}) + + def test_note_reaches_the_monitor(self): + from serve import telemetry + t = telemetry.Telemetry(extra=lambda: {"tok_s": 1.0}) + try: + with mock.patch.object(telemetry, "gpu_reader", lambda index=0, amd=False: self._reader()): + t2 = telemetry.Telemetry(extra=lambda: {"tok_s": 1.0}) + try: + self.assertEqual(t2.static["gpu_name"], "Apple M1 Max") + self.assertIn("powermetrics", t2.static["gpu_note"]) + with mock.patch("serve.telemetry.metal_working_set", lambda: WORKING_SET): + s = t2.sample() + self.assertEqual(s["gpu_util"], 76) + self.assertEqual(s["gpu_mem_used"], 13384105984) + self.assertEqual(s["gpu_mem_total"], WORKING_SET) + self.assertIsNone(s.get("gpu_temp")) # no reading without root; the cards show a dash + self.assertIsNone(s.get("gpu_pcie_gen")) + finally: + t2.close() + finally: + t.close() + + class TelemetryNoDisks(unittest.TestCase): """psutil.disk_io_counters() returns None where it finds no disk (a gVisor container; Windows with its disk counters off): the sampler goes on without a disk reading instead of its thread ending at the first sample, which left diff --git a/serve/test_slot_alloc.py b/serve/test_slot_alloc.py index 474b4f78..b2b4d35b 100644 --- a/serve/test_slot_alloc.py +++ b/serve/test_slot_alloc.py @@ -79,5 +79,30 @@ class OneGroup(unittest.TestCase): self.assertEqual(e.pick_slot([1, 2, 3]), 3) # else the least recently used +class OneSlotGroups(unittest.TestCase): + """--batch 2 on a two-GPU layer split: two pipelined groups of ONE slot. No slot sits below another in its group, + so no pad row writes an idle slot and each one keeps its chat (#2518): the next turn goes back to it.""" + + def test_the_slot_of_the_chat(self): + T = StrataEngine.TURN_ID + e = engine(2, 2) + self.assertEqual(e.slot_gs, 1) + hist_a, hist_b = list(range(1000, 1700)), list(range(2000, 2700)) + e.slot_held = [hist_a + [T, 7, 8] + [50, 51], hist_b + [T, 7, 8] + [60]] + e.slot_used = [10.0, 20.0] + again = lambda h: h + [T, 5, 6] + list(range(3000, 3050)) + self.assertEqual(e.pick_slot(again(hist_b)), 1) + self.assertEqual(e.pick_slot(again(hist_a)), 0) + self.assertEqual(e.pick_slot(list(range(5000, 5100))), 0) # a new chat: the least recently used one + + def test_groups_of_several_slots_still_ignore_held_tokens(self): + T = StrataEngine.TURN_ID + e = engine(4, 2) + self.assertEqual(e.slot_gs, 2) + hist = list(range(1000, 1700)) + e.slot_held[3] = hist + [T, 7, 8] + [50] + self.assertEqual(e.pick_slot(hist + [T, 5, 6] + list(range(3000, 3050))), 0) + + if __name__ == "__main__": unittest.main() diff --git a/serve/test_web_sampling.py b/serve/test_web_sampling.py new file mode 100644 index 00000000..b502ace2 --- /dev/null +++ b/serve/test_web_sampling.py @@ -0,0 +1,77 @@ +"""serve/test_web_sampling.py - the web page sends no sampling of its own until the user changes it (0.1.45, #2685). + +The page used to send temperature 0.6 / top_p 0.95 / top_k 20 with every request, so the server's model-card defaults +(thinking 1.0, not thinking 0.7 + presence penalty) never applied to it. Pins, with node, the page's own functions: +untouched (the shown defaults, or the old stored 0.6) sends nothing, a changed value is sent whole, 0 is greedy. +Skipped without node. + + python -m unittest serve.test_web_sampling -v +""" +from __future__ import annotations + +import json +import re +import shutil +import subprocess +import unittest +from pathlib import Path + +APP = Path(__file__).resolve().parent / "web" / "app.js" + + +def slice_between(src: str, start: str, end: str) -> str: + a = src.index(start) + return src[a:src.index(end, a)] + + +@unittest.skipUnless(shutil.which("node"), "node is not installed") +class WebSampling(unittest.TestCase): + @classmethod + def setUpClass(cls): + src = APP.read_text(encoding="utf-8") + cls.defaults = slice_between(src, "const DEFAULTS", "let settings") + cls.shared = slice_between(src, "function sharedDefaults", "async function saveShared") + cls.body = slice_between(src, " if (!(settings.temperature > 0)) {", " if (settings.seed)") + + def run_js(self, settings: dict, pieces: str): + js = (self.defaults + "function projectionLoaded() { return false; }\n" + self.shared + + f"\nconst settings = {json.dumps(settings)};\nconst body = {{}};\n{pieces}\n" + "console.log(JSON.stringify({body, shared: sharedDefaults(settings)}));") + out = subprocess.run(["node", "-e", js], capture_output=True, text=True, timeout=30) + self.assertEqual(out.returncode, 0, out.stderr) + return json.loads(out.stdout) + + def sampling(self, **kw): + s = {"thinking": "high", "temperature": 1.0, "top_p": 0.95, "top_k": 20, "max": "", "seed": "", "esp": True} + s.update(kw) + r = self.run_js(s, self.body) + return r["body"], r["shared"] + + def test_untouched_sends_nothing(self): + for t in (1.0, 0.6): # the shown default, and the old stored default + body, shared = self.sampling(temperature=t) + self.assertEqual(body, {}) + self.assertEqual(shared, {"reasoning_effort": "high"}) + + def test_changed_values_are_sent_whole(self): + body, shared = self.sampling(temperature=0.3) + self.assertEqual(body, {"temperature": 0.3, "top_p": 0.95, "top_k": 20}) + self.assertEqual(shared["temperature"], 0.3) + body, _ = self.sampling(top_k=40) + self.assertEqual(body, {"temperature": 1.0, "top_p": 0.95, "top_k": 40}) + body, _ = self.sampling(temperature=0.6, top_p=0.9) + self.assertEqual(body["temperature"], 0.6) + + def test_zero_is_greedy(self): + body, shared = self.sampling(temperature=0) + self.assertEqual(body, {"temperature": 0}) + self.assertEqual(shared["temperature"], 0) + + def test_the_old_default_is_shown_as_the_new_one(self): + src = APP.read_text(encoding="utf-8") + self.assertRegex(src, r"if \(samplingUntouched\(settings\)\) settings\.temperature = DEFAULTS\.temperature;") + self.assertIsNone(re.search(r"temperature: 0\.6", src)) + + +if __name__ == "__main__": + unittest.main() diff --git a/serve/web/app.js b/serve/web/app.js index d3ad404f..4bc2a46f 100644 --- a/serve/web/app.js +++ b/serve/web/app.js @@ -560,8 +560,18 @@ function markdown(text) { } // ------------------------------------------------------------------ Chat -const DEFAULTS = {thinking: "high", temperature: 0.6, top_p: 0.95, top_k: 20, max: "", seed: "", show: true, esp: true, mcp: true}; +// The sampling shown until the user changes it is the model card's thinking set. Untouched, the page sends NO sampling +// fields: the server picks the card's set for the request's thinking state (thinking 1.0 / 0.95 / 20, not thinking +// 0.7 / 0.8 / 20 with a presence penalty). Before 0.1.45 the page sent 0.6 itself; a stored 0.6 with the other two +// values untouched counts as "not set" too. +const DEFAULTS = {thinking: "high", temperature: 1.0, top_p: 0.95, top_k: 20, max: "", seed: "", show: true, esp: true, mcp: true}; +const OLD_DEFAULT_TEMPERATURE = 0.6; +function samplingUntouched(s) { + const t = +s.temperature; + return (t === DEFAULTS.temperature || t === OLD_DEFAULT_TEMPERATURE) && +s.top_p === DEFAULTS.top_p && +s.top_k === DEFAULTS.top_k; +} let settings = {...DEFAULTS, ...store.get("sampling", {})}; +if (samplingUntouched(settings)) settings.temperature = DEFAULTS.temperature; let messages = store.get("chat", []); let attachments = []; // {name, url} let busy = null; // {controller, msg} @@ -790,10 +800,10 @@ async function send() { const body = {model: health.model, messages: apiMessages(), stream: true, reasoning_effort: settings.thinking}; - if (settings.temperature > 0) { - Object.assign(body, {temperature: +settings.temperature, top_p: +settings.top_p, top_k: +settings.top_k}); - } else { + if (!(settings.temperature > 0)) { body.temperature = 0; + } else if (!samplingUntouched(settings)) { + Object.assign(body, {temperature: +settings.temperature, top_p: +settings.top_p, top_k: +settings.top_k}); } if (settings.seed) body.seed = +settings.seed; if (settings.max) body.max_tokens = +settings.max; @@ -1011,8 +1021,9 @@ async function loadShared() { $("s-share").setAttribute("aria-checked", String(sharedOn)); } function sharedDefaults(s) { - const d = {reasoning_effort: s.thinking, temperature: +s.temperature}; - if (+s.temperature > 0) Object.assign(d, {top_p: +s.top_p, top_k: +s.top_k}); + const d = {reasoning_effort: s.thinking}; + if (!(+s.temperature > 0)) d.temperature = 0; + else if (!samplingUntouched(s)) Object.assign(d, {temperature: +s.temperature, top_p: +s.top_p, top_k: +s.top_k}); if (s.seed) d.seed = +s.seed; if (s.max) d.max_tokens = +s.max; if (projectionLoaded()) d.experimental_speed_projection = s.esp !== false; diff --git a/src/kernels/ngram.cpp b/src/kernels/ngram.cpp index 91de3705..1fd7a3a5 100644 --- a/src/kernels/ngram.cpp +++ b/src/kernels/ngram.cpp @@ -508,27 +508,31 @@ bool PleTable::gather_batch(const uint32_t* rows, size_t n_tokens, float* out, s if (impl_->pending) { err = "PleTable::gather_batch while a token is in flight"; return false; } const size_t n = n_tokens * (size_t) PLE_N_HEADS; if (impl_->mode == PleIo::Direct) { - // Fast path: if all requested tokens match our in-flight/completed prefetch slots, collect them directly + // Fast path: match each requested token to any in-flight/completed prefetch slot, then collect directly // without issuing a second PleReader ticket or copying through RowCache. - if (impl_->n_prefetch >= n_tokens && n_tokens <= Impl::kMaxPrefetch) { + if (n_tokens <= Impl::kMaxPrefetch) { + size_t matched[Impl::kMaxPrefetch]; bool exact_match = true; for (size_t t = 0; t < n_tokens; ++t) { - if (std::memcmp(impl_->prefetch_keys[t], rows + t * PLE_N_HEADS, sizeof(uint32_t) * PLE_N_HEADS) != 0) { - exact_match = false; - break; - } + size_t i = 0; + while (i < impl_->n_prefetch && + std::memcmp(impl_->prefetch_keys[i], rows + t * PLE_N_HEADS, + sizeof(uint32_t) * PLE_N_HEADS) != 0) ++i; + if (i == impl_->n_prefetch) { exact_match = false; break; } + matched[t] = i; // duplicate requests may share one immutable prefetch slot } if (exact_match) { + bool collect_ok = true; for (size_t i = 0; i < impl_->n_prefetch; ++i) { - if (!impl_->reader.collect(impl_->prefetch_tickets[i], err)) { - impl_->n_prefetch = 0; - return false; - } + std::string ignored; + if (!impl_->reader.collect(impl_->prefetch_tickets[i], collect_ok ? err : ignored)) + collect_ok = false; // retain the first error, but drain every output buffer } impl_->n_prefetch = 0; + if (!collect_ok) return false; for (size_t t = 0; t < n_tokens; ++t) { for (int h = 0; h < PLE_N_HEADS; ++h) { - impl_->decode(impl_->prefetch_raw[t] + (size_t) h * impl_->rb, + impl_->decode(impl_->prefetch_raw[matched[t]] + (size_t) h * impl_->rb, out + (t * PLE_N_HEADS + (size_t) h) * PLE_HEAD_DIM); } } diff --git a/src/ngram/ple_reader_test.cpp b/src/ngram/ple_reader_test.cpp index 4663455f..670b29cc 100644 --- a/src/ngram/ple_reader_test.cpp +++ b/src/ngram/ple_reader_test.cpp @@ -242,9 +242,89 @@ void bad_sizes_are_refused(const std::string& dir) { } } +// #2622 (reported by @hotschmoe): gather_batch matched the prefetched tokens by position, so a prefetch of A..F +// gathered as D,E,F,A,B,C (or A gathered as A,A,A) issued a second ticket for rows that were already in flight. +// Every case: the bytes equal the mapped reader's, and the requested-row counter says how many rows were re-read. +void prefetched_keys_are_reused(const std::string& dir) { + const k::PleFormatInfo* fmt = nullptr; + for (int i = 0; i < k::ple_format_count(); ++i) + if (std::string(k::ple_formats()[i].name) == "IQ4_NL") fmt = &k::ple_formats()[i]; + if (fmt == nullptr) { CHECK(false, "no IQ4_NL format"); return; } + const std::string path = write_format_table(dir, *fmt, 4096); + if (path.empty()) { CHECK(false, "cannot write the prefetch table"); return; } + struct Case { const char* name; std::vector pre, want; uint64_t extra_rows; }; + std::vector seventeen; + for (int i = 0; i < 17; ++i) seventeen.push_back(i); + const std::vector cases = { + {"ordered", {0, 1}, {0, 1}, 0}, + {"permutation", {0, 1, 2, 3, 4, 5}, {3, 4, 5, 0, 1, 2}, 0}, + {"subset", {0, 1, 2, 3, 4, 5}, {4, 1}, 0}, + {"duplicates", {0}, {0, 0, 0}, 0}, + {"duplicate prefetch", {0, 0, 1}, {0, 1}, 0}, + {"missing key", {0, 1}, {0, 2}, 32}, + {"extra prefetch", {0, 1, 2}, {1}, 0}, + {"none", {0}, {}, 0}, + {"over the window", {0, 1}, seventeen, 17u * k::PLE_N_HEADS}, + {"window overflow", seventeen, {16}, 0}, + }; + const auto keys = [](const std::vector& ids) { + std::vector r; + for (int id : ids) for (int h = 0; h < k::PLE_N_HEADS; ++h) r.push_back((uint32_t) (id * 32 + h)); + return r; + }; + const auto requests = [](const k::PleTable& t) { + unsigned long long n = 0; + const std::string s = t.io_report(); + return std::sscanf(s.c_str(), "ple io: %llu rows", &n) == 1 ? (uint64_t) n : ~0ull; + }; + const size_t tok = (size_t) k::PLE_N_HEADS * k::PLE_HEAD_DIM; + std::vector> refs(cases.size()); + { + k::PleTable mapped; + k::PleIoOptions mio; + mio.mode = k::PleIo::Mmap; + std::string err; + if (!mapped.open(path, err, mio)) { CHECK(false, "mapped open: %s", err.c_str()); return; } + for (size_t c = 0; c < cases.size(); ++c) { + const auto r = keys(cases[c].want); + refs[c].resize(cases[c].want.size() * tok); + CHECK(mapped.gather_batch(r.data(), cases[c].want.size(), refs[c].data(), err), "mapped gather: %s", err.c_str()); + } + } + for (uint64_t cache : {uint64_t(0), uint64_t(256)}) for (bool worker : {false, true}) { + for (size_t c = 0; c < cases.size(); ++c) { + k::PleTable t; + k::PleIoOptions io; + io.mode = k::PleIo::Direct; + io.cache_rows = cache; + io.io_thread = worker; + io.max_inflight = 8; + io.keepalive_ms = 0; + std::string err; + if (!t.open(path, err, io)) { CHECK(false, "direct open: %s", err.c_str()); return; } + const auto pr = keys(cases[c].pre), wr = keys(cases[c].want); + for (size_t i = 0; i < cases[c].pre.size(); ++i) t.prefetch_rows(pr.data() + i * k::PLE_N_HEADS); + const uint64_t rq = requests(t), by = t.bytes_read(); + const float sentinel = -12345.5f; + std::vector out(refs[c].size() + 2, sentinel); + const bool ok = t.gather_batch(wr.data(), cases[c].want.size(), out.data() + 1, err); + CHECK(ok, "%s (cache %llu worker %d): %s", cases[c].name, (unsigned long long) cache, (int) worker, err.c_str()); + CHECK(requests(t) - rq == cases[c].extra_rows, "%s (cache %llu worker %d): %llu rows re-requested, want %llu", + cases[c].name, (unsigned long long) cache, (int) worker, (unsigned long long) (requests(t) - rq), + (unsigned long long) cases[c].extra_rows); + CHECK(t.bytes_read() - by == cases[c].want.size() * k::PLE_N_HEADS * fmt->row_bytes, "%s: bytes_read", cases[c].name); + CHECK(refs[c].empty() || std::memcmp(out.data() + 1, refs[c].data(), refs[c].size() * sizeof(float)) == 0, + "%s (cache %llu worker %d): bytes differ from the mapped reader", cases[c].name, (unsigned long long) cache, (int) worker); + CHECK(out.front() == sentinel && out.back() == sentinel, "%s: wrote outside the output", cases[c].name); + } + } + if (g_fail == 0) std::printf("ple prefetched-key reuse (reordered, duplicated, subset), Direct vs Mmap bytes: OK\n"); +} + void all_formats_round_trip(const std::string& dir) { bf16_widening_is_exact(); // BF16's own claim: a shift, exact for every pattern bad_sizes_are_refused(dir); + prefetched_keys_are_reused(dir); CHECK(k::ple_format_count() > 0, "no formats"); for (int i = 0; i < k::ple_format_count(); ++i) { const k::PleFormatInfo& f = k::ple_formats()[i]; diff --git a/tools/strata_tokenizer.py b/tools/strata_tokenizer.py index c9e84ee7..42a27fed 100644 --- a/tools/strata_tokenizer.py +++ b/tools/strata_tokenizer.py @@ -15,9 +15,11 @@ emoji (4-byte), combining marks, whitespace runs, and C0 control bytes. """ from __future__ import annotations +import bisect import json import pathlib import sys +import threading import regex @@ -115,6 +117,8 @@ class Tokenizer: # through __dict__.setdefault on each encoded segment or decoded token. self._piece_ids: dict[str, list[int]] = {} self._bytes_cache: dict[int, bytes] = {} + # How far a special-token match can look ahead of where it starts (PromptEncoder's margin). + self.max_special_len = max((len(t) for t in self.special_tokens), default=1) @staticmethod def _alt(literals: list[str]): @@ -234,13 +238,16 @@ class Tokenizer: out.extend(got) return out - def _encode_matching(self, text: str, pat, plain=()) -> list[int]: + def _encode_matching(self, text: str, pat, plain=(), marks: list | None = None) -> list[int]: """Encode `text`, emitting any literal `pat` matches as single tokens and BPE-ing the rest. The split happens on the RAW text, before the byte mapping, because a special token's string is a literal to match rather than bytes to decompose. Everything between the matches is tokenized normally - which is why a near-miss like `<|im_star` still costs ordinary tokens. A match that starts inside one of the `plain` (start, end) spans is left to the text around it (#537). + + `marks`, when given, receives (end of the match in `text`, ids so far) for every match: the points where + the encoding of a longer text with the same beginning can resume (PromptEncoder). """ if pat is None: return self._encode_plain(text) @@ -269,6 +276,8 @@ class Tokenizer: out.extend(self._encode_plain(text[pos:start])) out.append(self.special_tokens[m.group(0)]) pos = m.end() + if marks is not None: + marks.append((pos, len(out))) if pos < len(text): out.extend(self._encode_plain(text[pos:])) return out @@ -283,6 +292,11 @@ class Tokenizer: """ return self._encode_matching(text, self._special_re if parse_special else self._always_re, plain) + def encode_marked(self, text: str, parse_special: bool = False, plain=()) -> tuple[list[int], list[tuple[int, int]]]: + """encode() and its resume points: (end offset in `text`, ids so far) after every special-token match.""" + marks: list[tuple[int, int]] = [] + return self._encode_matching(text, self._special_re if parse_special else self._always_re, plain, marks), marks + def token_bytes(self, i: int) -> bytes: """The raw bytes of one token (a multi-byte character can be split across tokens).""" cache = self._bytes_cache @@ -303,6 +317,111 @@ class Tokenizer: return b"".join(self.token_bytes(i) for i in ids).decode("utf-8", errors=errors) +# ------------------------------------------------------------------ incremental prompts (#567) +def common_prefix_len(a: str, b: str) -> int: + """Length of the longest common prefix, by slices (C speed) rather than a loop over characters.""" + n = min(len(a), len(b)) + lo, step = 0, 4096 + while lo < n: # whole chunks first, doubling + hi = min(n, lo + step) + if a[lo:hi] != b[lo:hi]: + break + lo, step = hi, min(step * 2, 1 << 20) + else: + return n + while hi - lo > 1: # the first difference is in [lo, hi): bisect it + mid = (lo + hi) // 2 + if a[lo:mid] == b[lo:mid]: + lo = mid + else: + hi = mid + return lo + + +def plain_agree_until(a, b, limit: int) -> int: + """The first position below `limit` where the `plain` spans (#537) of two texts differ - a position inside a + span in one and not in the other - or `limit`. Spans are (start, end), sorted and disjoint.""" + def clip(spans): + return [(x, min(y, limit)) for x, y in spans if x < limit and y > x] + ca, cb = clip(a), clip(b) + for (xa, ya), (xb, yb) in zip(ca, cb): + if (xa, ya) != (xb, yb): + return min(xa, xb) if xa != xb else min(ya, yb) + if len(ca) != len(cb): + return (ca if len(ca) > len(cb) else cb)[min(len(ca), len(cb))][0] + return limit + + +class PromptEncoder: + """`tok.encode(text, parse_special=True, plain=...)` for chat prompts, reusing the ids of an earlier prompt. + + A chat client sends the whole conversation every turn, and the rendered prompt of turn n+1 starts with most of + turn n's. Running BPE over all of it again costs hundreds of milliseconds at 100K tokens, even when the engine + reuses the whole prefix. So the ids up to a SPECIAL-TOKEN BOUNDARY both prompts share are taken as they were, + and only the rest is encoded. + + WHY THAT IS EXACT. Encoding splits the text at special-token matches and BPE-encodes each stretch between two + matches on its own (`_encode_matching`): nothing crosses a match. Whether a match is taken at position q + depends only on text[q : q + max_special_len] and on whether q lies inside a `plain` span (#537: a + quoted in a message is text). So if the two texts agree up to L and their plain spans agree below L, every + decision the scan makes at a position before c = the end of some match, with c + max_special_len - 1 <= L, is + the same for both - the same matches, the same stretches, the same ids up to c - and from c the scan starts + afresh, exactly as it does on text[c:] with the spans shifted by c. Hence encode(new) == ids_old[:ids at c] + + encode(new[c:]). The margin matters only for a vocabulary where a literal overlaps another's end; it costs + re-encoding one short stretch. + + A few recent prompts are kept (one per conversation: a prompt replaces the one it continues, that is, one whose + text it reuses for more than half; a prompt that shares only the template's head with every entry is a new + entry), so a second client does not evict the first. `encode(..., store=False)` reads the kept prompts without + writing one (a token count is not a conversation's turn). The tokenizer needs `encode_marked` and + `max_special_len`. + """ + + def __init__(self, tok, keep: int = 4): + self.tok, self.keep = tok, keep + # (text, plain spans, ids, mark ends, mark counts) + self.entries: list[tuple[str, tuple, list[int], list[int], list[int]]] = [] + self.lock = threading.Lock() + self.last_reused = 0 # characters taken from an earlier prompt (for tests) + + def encode(self, text: str, plain=(), store: bool = True) -> list[int]: + plain = tuple(sorted(plain)) + with self.lock: + entries = self.entries # replaced, never changed in place + margin = self.tok.max_special_len - 1 + src, cut_k = None, -1 + for e in entries: + same = plain_agree_until(e[1], plain, common_prefix_len(e[0], text)) + k = bisect.bisect_right(e[3], same - margin) - 1 # the last boundary c with c <= same - margin + if k >= 0 and (src is None or e[3][k] > src[3][cut_k]): + src, cut_k = e, k + if src is None: + ids, marks = self.tok.encode_marked(text, parse_special=True, plain=plain) + ends, counts = [m[0] for m in marks], [m[1] for m in marks] + self.last_reused = 0 + else: + c, n = src[3][cut_k], src[4][cut_k] + rest = tuple((max(x - c, 0), y - c) for x, y in plain if y > c) + tail, marks = self.tok.encode_marked(text[c:], parse_special=True, plain=rest) + ids = src[2][:n] + tail + ends = src[3][:cut_k + 1] + [c + m[0] for m in marks] + counts = src[4][:cut_k + 1] + [n + m[1] for m in marks] + self.last_reused = c + if not store: + return ids + # every rendered prompt starts with the template's header, so sharing a special boundary does not make two + # prompts one conversation: only the entry whose text is mostly reused is replaced + old = src if src is not None and 2 * src[3][cut_k] > len(src[0]) else None + with self.lock: + self.entries = [(text, plain, ids, ends, counts)] + [e for e in self.entries if e is not old][:self.keep - 1] + return list(ids) + + def forget(self, ids: list[int]) -> None: + """The prompt with these ids was refused: it is not kept.""" + with self.lock: + self.entries = [e for e in self.entries if e[2] != ids] + + # ------------------------------------------------------------------ the pack's tokenizer/ directory def extract(gguf_path, out_dir) -> dict: """Write the tokenizer into `/tokenizer/` so the engine never opens the weight shards for it.""" diff --git a/tools/test_strata_tokenizer.py b/tools/test_strata_tokenizer.py index 17c3f698..e3cb498c 100644 --- a/tools/test_strata_tokenizer.py +++ b/tools/test_strata_tokenizer.py @@ -1,9 +1,18 @@ -"""The tokenizer's piece cache (0.1.41): the ids are the ones the merge loop gives, cache cold, warm or switched off. +"""Tests for tools/strata_tokenizer.py: llama.cpp's reference vectors, the piece cache (0.1.41: the ids are the ones the +merge loop gives, cache cold, warm or switched off), and the incremental prompt encoder, which must give exactly the +ids of a plain encode. - python -m unittest tools.test_strata_tokenizer + python -m unittest tools.test_strata_tokenizer -v + +The real vocabulary comes from $STRATA_TOKENIZER (a pack's tokenizer/ directory) when set, else from llama.cpp's +`models/ggml-vocab-qwen35.gguf` next to the gguf-py the tools use ($STRATA_GGUF_PY, third_party/llama.cpp); the +tests that need it are skipped without one. The chat template's USER_DEFINED tokens (``, ``, +...) are added to the llama.cpp vocabulary, which does not have them, as the pack's has. """ from __future__ import annotations +import json +import os import random import sys import unittest @@ -11,8 +20,295 @@ from pathlib import Path ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT / "tools")) +sys.path.insert(0, str(ROOT)) import strata_tokenizer as ST # noqa: E402 +ADDED = ["", "", "", "", "", "", + "<|vision_start|>", "<|image_pad|>", "<|vision_end|>"] + + +def vocab_gguf() -> Path | None: + cands = [os.environ.get("STRATA_GGUF_PY"), ROOT / "third_party" / "llama.cpp" / "gguf-py", "/opt/llama.cpp/gguf-py"] + for c in cands: + if c and (Path(c).parent / "models" / "ggml-vocab-qwen35.gguf").is_file(): + return Path(c).parent / "models" / "ggml-vocab-qwen35.gguf" + return None + + +def load_tokenizer(): + pack = os.environ.get("STRATA_TOKENIZER") + if pack and (Path(pack) / "vocab.json").is_file(): + t = Path(pack) + vocab = json.loads((t / "vocab.json").read_text(encoding="utf-8")) + tokens = [None] * len(vocab) + for s, i in vocab.items(): + tokens[i] = s + return ST.Tokenizer(tokens, (t / "merges.txt").read_text(encoding="utf-8").split("\n"), + json.loads((t / "token_type.json").read_text())) + path = vocab_gguf() + return None if path is None else gguf_tokenizer(path) + + +def gguf_tokenizer(path): + """llama.cpp's qwen35 vocabulary with the template's added tokens: the one its test vectors were made with.""" + from gguf_reader import GGUFFile + md = GGUFFile(path).metadata + tokens, types = list(md["tokenizer.ggml.tokens"]), list(md["tokenizer.ggml.token_type"]) + for s in ADDED: + if s not in tokens: + tokens.append(s) + types.append(4) + return ST.Tokenizer(tokens, list(md["tokenizer.ggml.merges"]), types) + + +def byte_tokenizer(specials: list[str]) -> ST.Tokenizer: + """A vocabulary of the 256 byte tokens and `specials` (CONTROL), no merges: small enough to build per test.""" + tokens = [ST.BYTE_TO_UNICODE[b] for b in range(256)] + specials + return ST.Tokenizer(tokens, [], [1] * 256 + [3] * len(specials)) + + +WORDS = ["the", "tokenizer", " ", " ", "\n", "\n\n", "\t", "def", "f(x):", "return", "x", "==", "->", "café", + "你好", "\U0001f600", "é", "1234", "don't", "I'm", "<|im_", "", + " ", "\r\n", "Hello,", "world!", "<", ">", "|", " ", "x" * 70] + + +def random_text(rnd: random.Random, n: int, specials=()) -> str: + parts = [] + for _ in range(n): + parts.append(rnd.choice(WORDS) if not specials or rnd.random() > 0.05 else rnd.choice(specials)) + if rnd.random() < 0.5: + parts.append(" ") + return "".join(parts) + + +def conversation_prompts(template, rnd: random.Random) -> list[tuple[str, str]]: + """Rendered prompts, in the order a server would see them, from clients that grow a conversation turn after turn, + use tools, send images, change the start of the prompt, edit an earlier turn, send a shorter conversation, and + interleave with each other: (what happens, prompt).""" + out: list[tuple[str, str]] = [] + specials = ["<|im_start|>", "<|im_end|>", "", "", "", ""] + weather = [{"type": "function", "function": {"name": "get_weather", "description": "Get the weather", + "parameters": {"type": "object", "properties": { + "city": {"type": "string"}}, "required": ["city"]}}}] + + def add(what, msgs, tools=None, **kw): + out.append((what, template.render(msgs, tools=tools, **kw))) + + # several turns, with unicode, a literal special token inside the text, and reasoning kept or not + msgs = [{"role": "system", "content": "Be brief. Réponds en français si on te parle en français."}] + for turn in range(6): + msgs.append({"role": "user", "content": random_text(rnd, 60, specials) + " é你\U0001f600 " * turn}) + add(f"turn {turn}", msgs) + answer = {"role": "assistant", "content": random_text(rnd, 30)} + if turn % 2: + answer["reasoning_content"] = random_text(rnd, 30) + msgs.append(answer) + grown = list(msgs) + + # tools: definitions, a call, its response, then a following turn + tmsgs = [{"role": "user", "content": "Weather in Oslo?"}] + add("tools 1", tmsgs, weather) + tmsgs += [{"role": "assistant", "content": "", "tool_calls": [{"type": "function", "function": { + "name": "get_weather", "arguments": {"city": "Oslo"}}}]}, + {"role": "tool", "content": "{\"temp\": 3}", "name": "get_weather"}] + add("tools 2", tmsgs, weather) + tmsgs += [{"role": "assistant", "content": "3 degrees."}, {"role": "user", "content": "And in Bergen? " * 20}] + add("tools 3", tmsgs, weather) + + # images: one per user turn, the placeholder tokens are specials + imsgs = [] + for turn in range(3): + imsgs.append({"role": "user", "content": [{"type": "image", "source": f"img{turn}"}, + {"type": "text", "text": f"What is on picture {turn}?"}]}) + add(f"images {turn}", imsgs) + imsgs.append({"role": "assistant", "content": "A cat."}) + + # the start of the prompt changes (a thinking level, a system message), then the old conversation grows again + add("effort changed", grown, reasoning_effort="low") + add("system changed", [dict(grown[0], content=grown[0]["content"] + " Be kind.")] + grown[1:]) + grown.append({"role": "user", "content": "one more question"}) + add("grown again", grown) + + # an earlier turn edited, then a conversation that got shorter, then the full one again + edited = list(grown) + edited[3] = dict(edited[3], content=edited[3]["content"] + " (edited)") + add("edited", edited) + add("shorter", grown[:5]) + add("empty user", [{"role": "user", "content": ""}]) + add("full again", grown) + return out + + +class Oracle(unittest.TestCase): + """llama.cpp's own tokenizer test vectors for qwen35.""" + + def test_llama_cpp_vectors(self): + path = vocab_gguf() + if path is None or not path.with_suffix(".gguf.inp").is_file(): + self.skipTest("llama.cpp's ggml-vocab-qwen35.gguf test vectors are not here") + tk = gguf_tokenizer(path) # not $STRATA_TOKENIZER: the vectors belong to this vocabulary + inp = path.with_suffix(".gguf.inp").read_text(encoding="utf-8").split("\n__ggml_vocab_test__\n") + out = path.with_suffix(".gguf.out").read_text(encoding="utf-8").split("\n") + n = 0 + for text, want in zip(inp, out): + if text.endswith("\n__ggml_vocab_test__"): + text = text[:-len("\n__ggml_vocab_test__")] + self.assertEqual(tk.encode(text), [int(x) for x in want.split()], repr(text[:40])) + n += 1 + self.assertGreater(n, 40) + + +class Marks(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.tk = load_tokenizer() + if cls.tk is None: + raise unittest.SkipTest("no tokenizer (set STRATA_TOKENIZER or STRATA_GGUF_PY)") + + def test_marks(self): + text = "a<|im_start|>user\nhi<|im_end|>\n\n" + ids, marks = self.tk.encode_marked(text, parse_special=True) + self.assertEqual(ids, self.tk.encode(text, parse_special=True)) + self.assertEqual([m[0] for m in marks], [text.index("user"), text.index("<|im_end|>") + len("<|im_end|>"), + len(text) - 1]) + for end, n in marks: # each mark: the ids of the text up to it + self.assertEqual(ids[:n], self.tk.encode(text[:end], parse_special=True)) + + +class Prompts(unittest.TestCase): + """PromptEncoder: the ids of a plain encode, whatever the earlier prompts were.""" + + @classmethod + def setUpClass(cls): + cls.tk = load_tokenizer() + if cls.tk is None: + raise unittest.SkipTest("no tokenizer (set STRATA_TOKENIZER or STRATA_GGUF_PY)") + from serve.frontend import ChatTemplate + cls.template = ChatTemplate(ROOT / "serve" / "chat_template.jinja") + + def check(self, enc, text): + self.assertEqual(enc.encode(text), self.tk.encode(text, parse_special=True)) + + def test_chat_golden(self): + cases = json.loads((ROOT / "serve" / "chat_golden.json").read_text(encoding="utf-8")) + enc = ST.PromptEncoder(self.tk) + for c in cases + cases: # the second round finds its own earlier prompts + if c.get("rendered"): + with self.subTest(case=c["name"]): + self.check(enc, c["rendered"]) + + def test_conversations(self): + for seed in range(3): + prompts = conversation_prompts(self.template, random.Random(seed)) + wants = [self.tk.encode(prompt, parse_special=True) for _, prompt in prompts] + for keep in (1, 2, 4): + enc = ST.PromptEncoder(self.tk, keep=keep) + for (what, prompt), want in zip(prompts, wants): + with self.subTest(seed=seed, keep=keep, what=what): + self.assertEqual(enc.encode(prompt), want) + + def test_growing_turn_reuses_the_previous_prompt(self): + enc = ST.PromptEncoder(self.tk) + msgs = [{"role": "user", "content": "first question " * 50}] + prev = "" + for turn in range(4): + prompt = self.template.render(msgs) + self.check(enc, prompt) + if turn: + self.assertGreater(enc.last_reused, len(prev) - 100) + prev = prompt + msgs += [{"role": "assistant", "content": f"answer {turn}"}, {"role": "user", "content": f"question {turn}"}] + + def test_prefix_of_an_earlier_prompt_and_unrelated_prompts(self): + enc = ST.PromptEncoder(self.tk, keep=2) + a = self.template.render([{"role": "user", "content": "first question " * 50}]) + b = self.template.render([{"role": "user", "content": "another conversation"}]) + c = self.template.render([{"role": "user", "content": "a third one"}]) + for text in (a, a[:len(a) // 2], b, c, a, a + "tail", "", "<|im_end|>", "<|im_end|><|im_end|>x"): + with self.subTest(text=text[:30]): + self.check(enc, text) + + def test_a_read_only_encode_gives_the_ids_and_stores_nothing(self): + enc = ST.PromptEncoder(self.tk) + a = self.template.render([{"role": "user", "content": "first question " * 50}]) + enc.encode(a) + kept = list(enc.entries) + for text in (a + "<|im_start|>user\nand a follow-up<|im_end|>\n", a[:len(a) // 2], "<|im_end|>x", + self.template.render([{"role": "user", "content": "another conversation"}])): + with self.subTest(text=text[:30]): + self.assertEqual(enc.encode(text, store=False), self.tk.encode(text, parse_special=True)) + self.assertTrue(enc.entries == kept, "a read-only encode wrote an entry") + self.assertEqual(enc.encode(a + "tail", store=False), self.tk.encode(a + "tail", parse_special=True)) + self.assertGreater(enc.last_reused, len(a) // 2) + + def test_the_returned_ids_are_the_callers(self): + enc = ST.PromptEncoder(self.tk) + text = "<|im_start|>user\nhi<|im_end|>\n" + "and a tail longer than the look-ahead margin. " * 3 + want = self.tk.encode(text, parse_special=True) + got = enc.encode(text) + got[0] = -1 # inside the part the next call takes from this one + got.append(-1) + self.assertEqual(enc.encode(text), want) + self.assertGreater(enc.last_reused, 0) + + +class Margin(unittest.TestCase): + """A special literal that is a prefix of a longer one: cutting at the shorter one's end without the look-ahead + margin would give different ids.""" + + def test_longer_literal_past_the_common_prefix(self): + tk = byte_tokenizer(["<|a|>", "<|a|>zz"]) + enc = ST.PromptEncoder(tk) + first = "x<|a|>y" + "<|a|>" * 3 + "q" + self.assertEqual(enc.encode(first), tk.encode(first, parse_special=True)) + second = "x<|a|>y" + "<|a|>" * 3 + "zz" # shares everything up to the last <|a|> + got = enc.encode(second) + self.assertEqual(got, tk.encode(second, parse_special=True)) + self.assertEqual(got[-1], tk.special_tokens["<|a|>zz"]) + self.assertGreater(enc.last_reused, 0) # the earlier boundaries were still reused + + def test_literal_ending_inside_another(self): + tk = byte_tokenizer(["<|a|>", "b<|a|>c"]) + enc = ST.PromptEncoder(tk) + for text in ("zzb<|a|>d", "zzb<|a|>c", "zzb<|a|>cz<|a|>", "zzb<|a|>cz<|a|>b<|a|>c"): + with self.subTest(text=text): + self.assertEqual(enc.encode(text), tk.encode(text, parse_special=True)) + + def test_random_specials(self): + """Literals that extend or end inside one another, texts that change right after one of them.""" + rnd = random.Random(3) + + def word(lo, hi): + return "".join(rnd.choice("<|ab>") for _ in range(rnd.randint(lo, hi))) + for trial in range(60): + a, b = word(2, 4), word(2, 4) + lits = sorted({a, a + word(1, 3), b, word(1, 2) + b + word(0, 2)}) + tk = byte_tokenizer(lits) + enc = ST.PromptEncoder(tk, keep=3) + pieces = lits + [x[:rnd.randint(1, len(x))] for x in lits] + ["x", "y", "a", "|"] + base = "".join(rnd.choice(pieces) for _ in range(60)) + for _ in range(30): + cut = rnd.randint(0, len(base)) + text = base[:cut] + "".join(rnd.choice(pieces) for _ in range(rnd.randint(0, 6))) + self.assertEqual(enc.encode(text), tk.encode(text, parse_special=True), (lits, text)) + if rnd.random() < 0.3: + base = text + base[cut:] + + +class CommonPrefix(unittest.TestCase): + def test_against_a_loop(self): + rnd = random.Random(5) + for _ in range(500): + a = "".join(rnd.choice("ab") for _ in range(rnd.randint(0, 9000))) + k = rnd.randint(0, len(a)) + b = a[:k] + rnd.choice(["", "c", "cd"]) + a[k:k + rnd.randint(0, 50)] + want = 0 + while want < min(len(a), len(b)) and a[want] == b[want]: + want += 1 + self.assertEqual(ST.common_prefix_len(a, b), want) + self.assertEqual(ST.common_prefix_len(b, a), want) + + def toy() -> ST.Tokenizer: base = [ST.BYTE_TO_UNICODE[b] for b in range(256)]