From 98964c1108b5b08dbb215278f208a5d8ea19e7d4 Mon Sep 17 00:00:00 2001 From: Kevin Gilbertson <880013+gilby@users.noreply.github.com> Date: Sun, 1 Nov 2026 04:48:09 +0100 Subject: [PATCH 1/2] dsv41 metal: indexer scores, top-512, candidate blocks and the MoE route index.metal scores FP4 queries against FP4 keys in one stated order with causal and block masks, topk512.metal keeps each row's top 512, cand.metal picks the candidate blocks as a byte mask, and route.metal takes the top-k of sqrt(softplus(logit)) + bias for the backbone and the drafter. Each is registered with the contract. --- zig/kernels/metal/deepseek_v41/cand.metal | 129 ++++++ zig/kernels/metal/deepseek_v41/index.metal | 207 +++++++++ zig/kernels/metal/deepseek_v41/route.metal | 95 +++++ zig/kernels/metal/deepseek_v41/sources.zig | 8 + zig/kernels/metal/deepseek_v41/topk512.metal | 109 +++++ zig/src/core/contract/dsv41_cand_records.zig | 11 + zig/src/core/contract/dsv41_index_records.zig | 12 + zig/src/core/contract/dsv41_metal_records.zig | 9 +- zig/src/core/contract/dsv41_route_records.zig | 7 + .../core/contract/dsv41_topk512_records.zig | 11 + zig/src/families/deepseek_v41/metal/cand.zig | 28 ++ .../deepseek_v41/metal/contract_exports.zig | 1 + zig/src/families/deepseek_v41/metal/index.zig | 123 ++++++ zig/tests/contract/dsv41_index_ref.zig | 258 +++++++++++ zig/tests/contract/metal_dsv41_cand.zig | 366 ++++++++++++++++ zig/tests/contract/metal_dsv41_index.zig | 286 +++++++++++++ zig/tests/contract/metal_dsv41_limits.zig | 9 + zig/tests/contract/metal_dsv41_main.zig | 73 +++- zig/tests/contract/metal_dsv41_route.zig | 399 ++++++++++++++++++ zig/tests/contract/metal_dsv41_topk512.zig | 298 +++++++++++++ zig/tests/contract/metal_main.zig | 2 + 21 files changed, 2434 insertions(+), 7 deletions(-) create mode 100644 zig/kernels/metal/deepseek_v41/cand.metal create mode 100644 zig/kernels/metal/deepseek_v41/index.metal create mode 100644 zig/kernels/metal/deepseek_v41/route.metal create mode 100644 zig/kernels/metal/deepseek_v41/topk512.metal create mode 100644 zig/src/core/contract/dsv41_cand_records.zig create mode 100644 zig/src/core/contract/dsv41_index_records.zig create mode 100644 zig/src/core/contract/dsv41_route_records.zig create mode 100644 zig/src/core/contract/dsv41_topk512_records.zig create mode 100644 zig/src/families/deepseek_v41/metal/cand.zig create mode 100644 zig/src/families/deepseek_v41/metal/index.zig create mode 100644 zig/tests/contract/dsv41_index_ref.zig create mode 100644 zig/tests/contract/metal_dsv41_cand.zig create mode 100644 zig/tests/contract/metal_dsv41_index.zig create mode 100644 zig/tests/contract/metal_dsv41_route.zig create mode 100644 zig/tests/contract/metal_dsv41_topk512.zig diff --git a/zig/kernels/metal/deepseek_v41/cand.metal b/zig/kernels/metal/deepseek_v41/cand.metal new file mode 100644 index 0000000000..0cc1d371f4 --- /dev/null +++ b/zig/kernels/metal/deepseek_v41/cand.metal @@ -0,0 +1,129 @@ +// DeepSeek-V4.1 candidate blocks, a threadgroup a row: block maxima, the newest block pinned, the top blocks as a mask. +#include +using namespace metal; + +#ifndef TF_CTL +#define TF_CTL 0 // planted faults: 1 NaN ranks highest, 2 ties to the higher block, 3 no pin, 4 input write +#endif +#define TF_TG 1024u +#define TF_BINS 2048u +#define TF_NINF 0x007fffffu // the order key of -inf: a block at or below it is never a candidate +#define TF_PINF 0xff800000u // the order key of +inf, which the pinned block takes + +struct TfCandArgs { + uint keys; // scores a row (the row stride of the scores) + uint blocks; // the row stride of the keyed scratch and of the mask + uint topk; // blocks kept, at most 2048 +}; + +// Order-preserving key of a float's bits: NaN as -inf (control 1: +inf), -0 as +0. +inline uint tf_cand_key(uint u) { + if ((u & 0x7fffffffu) > 0x7f800000u) u = TF_CTL == 1 ? 0x7f800000u : 0xff800000u; + if (u == 0x80000000u) u = 0u; + return (u >> 31) ? ~u : (u | 0x80000000u); +} + +// The bin holding the k-th largest key (bins scanned from the top) and the keys in the bins above it. +inline uint2 tf_cand_pick(threadgroup atomic_uint* bins, uint k, threadgroup uint* sc, uint tid, uint lane, uint simd) { + const uint a = atomic_load_explicit(bins + (TF_BINS - 1u - 2u * tid), memory_order_relaxed); + const uint b = atomic_load_explicit(bins + (TF_BINS - 2u - 2u * tid), memory_order_relaxed); + const uint local = a + b, ex = simd_prefix_exclusive_sum(local); + if (lane == 31u) sc[simd] = ex + local; + threadgroup_barrier(mem_flags::mem_threadgroup); + uint before = ex; + for (uint s = 0; s < simd; ++s) before += sc[s]; + if (before < k && before + local >= k) { + const bool first = before + a >= k; + sc[32] = TF_BINS - (first ? 1u : 2u) - 2u * tid; + sc[33] = first ? before : before + a; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint2 r = uint2(sc[32], sc[33]); + threadgroup_barrier(mem_flags::mem_threadgroup); + return r; +} + +// scores [rows][keys], visible [rows], keyed [rows][blocks] scratch, mask [rows][blocks] bytes. +[[max_total_threads_per_threadgroup(TF_TG)]] kernel void tf_cand( + constant TfCandArgs& args [[buffer(0)]], + device const uint* scores [[buffer(1)]], + device const int* visible [[buffer(2)]], + device uint* keyed [[buffer(3)]], + device uchar* mask [[buffer(4)]], + uint row [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], + uint simd [[simdgroup_index_in_threadgroup]]) { + threadgroup atomic_uint bins[TF_BINS]; + threadgroup uint sc[34], part[64]; + const uint n = min((args.keys + 7u) / 8u, args.blocks); + const uint k = min(args.topk, 2048u); + const int vis = visible[row]; + const uint last = vis >= 1 ? uint(vis - 1) / 8u : 0xffffffffu; + device const uint* s = scores + ulong(row) * args.keys; + device uint* kd = keyed + ulong(row) * args.blocks; + device uchar* out = mask + ulong(row) * args.blocks; +#if TF_CTL == 4 + if (tid == 0u && args.keys > 0u) ((device uint*)s)[0] ^= 1u; +#endif + for (uint j = tid; j < args.blocks; j += TF_TG) { + uint m = TF_NINF; + if (j < n) { + const uint end = min(8u * j + 8u, args.keys); + for (uint i = 8u * j; i < end; ++i) m = max(m, tf_cand_key(s[i])); + } + kd[j] = (TF_CTL != 3 && j == last && j < n) ? TF_PINF : m; + } + uint threshold = 0u, above = 0u, ties = 0u; + if (n > k && k > 0u) { + uint prefix = 0u, left = k; + for (uint pass = 0; pass < 3u; ++pass) { + const uint shift = pass == 0u ? 21u : (pass == 1u ? 10u : 0u); + const uint bmask = pass == 2u ? 0x3ffu : 0x7ffu; + const uint up = pass == 1u ? 21u : 10u; + for (uint i = tid; i < TF_BINS; i += TF_TG) atomic_store_explicit(bins + i, 0u, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint i = tid; i < n; i += TF_TG) { + const uint key = kd[i]; + if (pass == 0u || (key >> up) == prefix) atomic_fetch_add_explicit(bins + ((key >> shift) & bmask), 1u, memory_order_relaxed); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint2 r = tf_cand_pick(bins, left, sc, tid, lane, simd); + prefix = (prefix << (pass == 2u ? 10u : 11u)) | r.x; + left -= r.y; + above += r.y; + } + threshold = prefix; + ties = atomic_load_explicit(bins + (prefix & 0x3ffu), memory_order_relaxed); + } + const bool all = n <= k; + const uint wanted = k - above; // keys equal to the threshold that make the set: the lowest blocks + const bool scan = !all && (TF_CTL == 2 || wanted < ties); + const uint skip = TF_CTL == 2 ? ties - wanted : 0u; + uint eq_run = 0u; + for (uint base = 0; base < args.blocks; base += TF_TG) { + const uint j = base + tid; + const uint key = j < n ? kd[j] : 0u; + const bool real = key > TF_NINF; + const bool eq = real && !all && key == threshold; + bool pick = real && k > 0u && (all || key > threshold); + if (scan) { + const uint ex = simd_prefix_exclusive_sum(eq ? 1u : 0u); + threadgroup uint* p = part + 32u * ((base / TF_TG) & 1u); + if (lane == 31u) p[simd] = ex + (eq ? 1u : 0u); + threadgroup_barrier(mem_flags::mem_threadgroup); + uint before = ex, total = 0u; + for (uint q = 0; q < 32u; ++q) { + const uint v = p[q]; + before += q < simd ? v : 0u; + total += v; + } + const uint eq_x = eq_run + before; + pick = pick || (eq && (TF_CTL == 2 ? eq_x >= skip : eq_x < wanted)); + eq_run += total; + } else { + pick = pick || (eq && k > 0u); + } + if (j < args.blocks) out[j] = pick ? 1 : 0; + } +} diff --git a/zig/kernels/metal/deepseek_v41/index.metal b/zig/kernels/metal/deepseek_v41/index.metal new file mode 100644 index 0000000000..a5b741574b --- /dev/null +++ b/zig/kernels/metal/deepseek_v41/index.metal @@ -0,0 +1,207 @@ +// DeepSeek V4.1's indexer scores on Metal: FP4xFP4 dots in one stated order, causal and candidate-block masks. +#include +using namespace metal; +#pragma METAL fp contract(off) + +#ifndef TF_CTL +#define TF_CTL 0 +#endif + +// The release's indexer geometry (release-config.json): 32 index heads of 128 dims, candidate blocks of 8 keys. +#define IH 32u +#define ID 128u +#define CBS 8u +#define QROW 2048u // one query row: 32 heads x 128 FP4 values, two a byte +#define QSROW 128u // one query row: 32 heads x 4 per-32-group scales, e8m0 +#define IXR 4u // simdgroups a threadgroup carries, each scoring a block of 8 rows +#define IXKT 2u // key fragments (8 keys each) a simdgroup scores together over one block's rows + +// Fully unrolled: the matrix arrays below must stay in registers, indexed at compile time. +#define UNR _Pragma("clang loop unroll(full)") + +struct TfDsv41IxArgs { + uint rows, keys, start, ratio, chunk, out_stride, w_stride, ik_stride, mask_stride; + float w_scale; // softmax_scale * n_heads^-0.5 (model.py:556, 513): scales each bf16 weight once, in fp32 +}; + +// One E2M1 code by its bit fields: 0, 0.5, 1, 1.5, 2, 3, 4, 6 and their negatives (the format has no NaN or infinity). +static inline float tf_ix_e2m1(const uint c) { + const uint e = (c >> 1) & 3u, m = c & 1u; + const float v = (e == 0u) ? float(m) * 0.5f : ldexp(float(m) + 2.0f, int(e) - 2); + return (c & 8u) ? -v : v; +} + +// One e8m0 scale: 2^(c-127) from its exponent bits; code 255 is that dtype's NaN, code 0 the subnormal 2^-127. +static inline float tf_ix_e8m0(const uint c) { + return c == 255u ? as_type(0x7FC00000u) : (c == 0u ? as_type(0x00400000u) : as_type(c << 23)); +} + +// ReLU keeping a NaN: fmax() folds one to 0, which would hide a broken dot (gate 1's planted control). +static inline float tf_ix_relu(const float x) { +#if TF_CTL == 1 + return fmax(x, 0.0f); +#else + return (x != x || x > 0.0f) ? x : 0.0f; +#endif +} + +// One FP4 code's value as an fp16 (exact; the format has no NaN or infinity), by its code. +constant constexpr half tf_ix_lut[16] = {0.0h, 0.5h, 1.0h, 1.5h, 2.0h, 3.0h, 4.0h, 6.0h, + -0.0h, -0.5h, -1.0h, -1.5h, -2.0h, -3.0h, -4.0h, -6.0h}; + +// Scores 8 rows on fp16 MMA, fp32 sums: FP4 codes are exact in fp16 and a group sum exact in any order. +[[kernel, max_total_threads_per_threadgroup(32 * IXR)]] void tf_dsv41_ix_scores( + const device uchar* Q [[buffer(0)]], const device uchar* QS [[buffer(1)]], const device bfloat* W [[buffer(2)]], + const device uchar* IK [[buffer(3)]], device float* S [[buffer(4)]], const device uchar* CM [[buffer(5)]], + constant TfDsv41IxArgs& a [[buffer(6)]], uint2 tg [[threadgroup_position_in_grid]], + uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]]) { + const uint r0 = (tg.y * IXR + sg) * 8u; + if (r0 >= a.rows) return; // uniform across the simdgroup, which has no barrier with the others + const uint qid = lane / 4; + const uint fm = (qid & 4u) + ((lane / 2u) % 4u); + const uint fn = (qid & 2u) * 2u + (lane % 2u) * 2u; + const bool row_ok = r0 + fm < a.rows; + const uint row = row_ok ? r0 + fm : a.rows - 1u; // the lane's query row (an MMA row); a tail's lanes repeat the last + // keys a query reaches: (start + row) / ratio (model.py:565 prefill with start 1, 568 decode with end_pos) + const uint visible = (a.start + row) / a.ratio; + const uint vmax = (a.start + min(r0 + 8u, a.rows) - 1u) / a.ratio; // the block's last row reaches the most + const uint first = tg.x * a.chunk; + const uint last = min(first + a.chunk, a.keys); + const device uchar* qrow = Q + size_t(row) * QROW + fn / 2u; + for (uint key0 = first; key0 < last; key0 += IXKT * 8u) { + float sacc[IXKT][2]; + UNR for (uint kf = 0; kf < IXKT; ++kf) sacc[kf][0] = sacc[kf][1] = 0.0f; + if (key0 < vmax) { + simdgroup_matrix bq[IXKT][16]; + float sk[IXKT][2][4]; + UNR for (uint kf = 0; kf < IXKT; ++kf) { + const uint ka = min(key0 + 8u * kf + fn, a.keys - 1u), kb = min(ka + 1u, a.keys - 1u); + const device uchar* ra = IK + size_t(ka) * a.ik_stride; + const device uchar* rb = IK + size_t(kb) * a.ik_stride; + UNR for (uint ks = 0; ks < 16u; ++ks) { + const uint d = 8u * ks + fm; + bq[kf][ks].thread_elements()[0] = tf_ix_lut[(ra[d >> 1] >> (4u * (d & 1u))) & 0xFu]; + bq[kf][ks].thread_elements()[1] = tf_ix_lut[(rb[d >> 1] >> (4u * (d & 1u))) & 0xFu]; + } + UNR for (uint g = 0; g < 4u; ++g) { // a key's row: 64 FP4 bytes then its 4 e8m0 scales + sk[kf][0][g] = tf_ix_e8m0(ra[ID / 2u + g]); + sk[kf][1][g] = tf_ix_e8m0(rb[ID / 2u + g]); + } + } +#if TF_CTL == 2 // gate 3's control: the head fold in reverse, and head 0 (the last one folded) dropped + for (uint h = IH - 1u; h >= 1u; --h) { +#else + for (uint h = 0u; h < IH; ++h) { +#endif + simdgroup_matrix aq[16]; + UNR for (uint ks = 0; ks < 16u; ++ks) { + const uchar b = qrow[h * (ID / 2u) + 4u * ks]; + aq[ks].thread_elements()[0] = tf_ix_lut[b & 0xFu]; + aq[ks].thread_elements()[1] = tf_ix_lut[b >> 4]; + } + float qs[4]; + UNR for (uint g = 0; g < 4u; ++g) qs[g] = tf_ix_e8m0(QS[size_t(row) * QSROW + h * 4u + g]); + // the row's weights times the softmax scale once, in fp32 (model.py:556), before any head meets a score + const float wv = float(W[size_t(row) * a.w_stride + h]) * a.w_scale; + simdgroup_matrix acc[IXKT][4]; + UNR for (uint kf = 0; kf < IXKT; ++kf) + UNR for (uint g = 0; g < 4u; ++g) acc[kf][g] = simdgroup_matrix(0.0f); + UNR for (uint kk = 0; kk < 4u; ++kk) + UNR for (uint kf = 0; kf < IXKT; ++kf) + UNR for (uint g = 0; g < 4u; ++g) { + simdgroup_multiply_accumulate(acc[kf][g], aq[4u * g + kk], bq[kf][4u * g + kk], acc[kf][g]); + } + UNR for (uint kf = 0; kf < IXKT; ++kf) + UNR for (uint e = 0; e < 2u; ++e) { + float p[4]; + UNR for (uint g = 0; g < 4u; ++g) { + p[g] = acc[kf][g].thread_elements()[e]; +#if TF_CTL == 3 // gate 4's control: the group sums accumulated in half + p[g] = float(half(p[g])); +#endif + } + const float prod = (((p[0] * qs[0]) * sk[kf][e][0]) + ((p[1] * qs[1]) * sk[kf][e][1])) + + (((p[2] * qs[2]) * sk[kf][e][2]) + ((p[3] * qs[3]) * sk[kf][e][3])); + sacc[kf][e] += tf_ix_relu(prod) * wv; + } + } + } + if (!row_ok) continue; + UNR for (uint kf = 0; kf < IXKT; ++kf) + UNR for (uint e = 0; e < 2u; ++e) { + const uint key = key0 + 8u * kf + fn + e; + if (key >= last) continue; + // keys a row reaches are still masked by the candidate blocks + const bool live = key < visible && (a.mask_stride == 0u || CM[size_t(row) * a.mask_stride + key / CBS] != 0u); +#if TF_CTL == 4 // gate 2's control: one float past the LAST output, into the guard bytes + if (row == a.rows - 1u && key == a.keys - 1u) S[size_t(a.rows) * a.out_stride] = 1.0f; +#endif + S[size_t(row) * a.out_stride + key] = live ? sacc[kf][e] : -INFINITY; // model.py:566, 576 mask unreachable keys + } + } +} + +#define TF_IX_QPAD 129u // a head's floats in threadgroup memory: one more than ID, so 32 lanes read 32 banks + +// A query row's head h, dim d: the nibble's value times its group's power of two (both exact in fp32). +static inline float tf_ix_q(const device uchar* Q, const device uchar* QS, const size_t at, const uint h, const uint d) { + return tf_ix_e2m1((Q[at + h * (ID / 2u) + (d >> 1)] >> (4u * (d & 1u))) & 0xFu) * tf_ix_e8m0(QS[at / QROW * QSROW + h * 4u + (d >> 5)]); +} + +// Decode's few rows: eight simdgroups share a row's decoded query and score a key each in turn, the same bits as above. +[[kernel, max_total_threads_per_threadgroup(256)]] void tf_dsv41_ix_scores_dec( + const device uchar* Q [[buffer(0)]], const device uchar* QS [[buffer(1)]], const device bfloat* W [[buffer(2)]], + const device uchar* IK [[buffer(3)]], device float* S [[buffer(4)]], const device uchar* CM [[buffer(5)]], + constant TfDsv41IxArgs& a [[buffer(6)]], uint2 tg [[threadgroup_position_in_grid]], + uint t [[thread_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], + uint sg [[simdgroup_index_in_threadgroup]], uint sgs [[simdgroups_per_threadgroup]]) { + const uint row = tg.y; + if (row >= a.rows) return; + threadgroup float qf[IH * TF_IX_QPAD]; // 16 KiB: this row's query, decoded once for the whole chunk + threadgroup float wf[IH]; + const size_t at = size_t(row) * QROW; + for (uint i = t; i < IH * ID; i += sgs * 32u) qf[(i / ID) * TF_IX_QPAD + i % ID] = tf_ix_q(Q, QS, at, i / ID, i % ID); + // the row's weights times the softmax scale once, in fp32 (model.py:556), before any head meets a score + if (t < IH) wf[t] = float(W[size_t(row) * a.w_stride + t]) * a.w_scale; + threadgroup_barrier(mem_flags::mem_threadgroup); + + const uint first = tg.x * a.chunk; + const uint last = min(first + a.chunk, a.keys); + // keys a query reaches: (start + row) / ratio (model.py:565 prefill with start 1, 568 decode with end_pos) + const uint visible = (a.start + row) / a.ratio; + const threadgroup float* qh = qf + size_t(lane) * TF_IX_QPAD; + const float wh = wf[lane]; + for (uint key = first + sg; key < last; key += sgs) { +#if TF_CTL == 4 // gate 2's control: one float past the LAST output, into the guard bytes (lane 0, before any mask) + if (lane == 0u && row == a.rows - 1u && key == a.keys - 1u) S[size_t(a.rows) * a.out_stride] = 1.0f; +#endif + const bool live = key < visible && (a.mask_stride == 0u || CM[size_t(row) * a.mask_stride + key / CBS] != 0u); + if (!live) { + if (lane == 0u) S[size_t(row) * a.out_stride + key] = -INFINITY; // model.py:566, 576 mask unreachable keys + continue; + } + const device uchar* ik = IK + size_t(key) * a.ik_stride; + const device uchar* ks = ik + ID / 2u; // a key's row: 64 FP4 bytes then its 4 e8m0 scales + float p[4]; + for (uint q = 0; q < 4u; ++q) { +#if TF_CTL == 3 // gate 4's control: the group sums accumulated in half + half h = 0.0h; + for (uint i = 32u * q; i < 32u * q + 32u; ++i) h += half(qh[i]) * half(tf_ix_e2m1((ik[i >> 1] >> (4u * (i & 1u))) & 0xFu)); + p[q] = float(h); +#else + float acc = 0.0f; + for (uint i = 32u * q; i < 32u * q + 32u; ++i) acc += qh[i] * tf_ix_e2m1((ik[i >> 1] >> (4u * (i & 1u))) & 0xFu); + p[q] = acc; +#endif + } + const float prod = ((p[0] * tf_ix_e8m0(ks[0])) + (p[1] * tf_ix_e8m0(ks[1]))) + ((p[2] * tf_ix_e8m0(ks[2])) + (p[3] * tf_ix_e8m0(ks[3]))); + const float term = tf_ix_relu(prod) * wh; + float acc = 0.0f; +#if TF_CTL == 2 // gate 3's control: the head fold in reverse, and head 0 (the last one folded) dropped + for (uint h = IH - 1u; h >= 1u; --h) acc += simd_shuffle(term, ushort(h)); +#else + for (uint h = 0u; h < IH; ++h) acc += simd_shuffle(term, ushort(h)); +#endif + if (lane == 0u) S[size_t(row) * a.out_stride + key] = acc; + } +} diff --git a/zig/kernels/metal/deepseek_v41/route.metal b/zig/kernels/metal/deepseek_v41/route.metal new file mode 100644 index 0000000000..da901548d4 --- /dev/null +++ b/zig/kernels/metal/deepseek_v41/route.metal @@ -0,0 +1,95 @@ +// DeepSeek-V4.1's MoE route: sqrt(softplus(logit)) scores, the top-k of score plus bias. Variants via TF_DSV41_E/TOPK. +#include +using namespace metal; + +#ifndef TF_DSV41_PER +#define TF_DSV41_PER ((TF_DSV41_E + 31) / 32) +#endif + +#ifndef TF_CTL +#define TF_CTL 0 // planted faults: 1 a higher id wins a tie, 2 a NaN key wins, 3 the bias in the weight, 4 a write past k +#endif + +// The reference's softplus: the linear tail above 20 and the quadratic one below -10, as the CUDA router computes it. +inline float tf_dsv41_softplus(float z) { + if (z != z) return z; + const float e = metal::precise::exp(fmin(z, 20.0f)); + if (z > 20.0f) return z; + if (z < -10.0f) return e - 0.5f * e * e; + return metal::precise::log(1.0f + e); +} + +// A candidate beats the incumbent when its key is larger, or the keys tie and its id is lower; NaN keys never win. +inline bool tf_dsv41_better(float as, int ai, float bs, int bi) { + if (ai >= TF_DSV41_E) return false; + if (bi >= TF_DSV41_E) return true; + if (as != as) return TF_CTL == 2; // control 2: a NaN key beats every other + if (bs != bs) return TF_CTL != 2; + if (as > bs) return true; + if (as < bs) return false; + return TF_CTL == 1 ? ai > bi : ai < bi; +} + +struct Dsv41RouteArgs { + int rows; + int topk; + float scale; + float eps; +}; + +// One simdgroup a row: the top-k of score plus bias (ties to the lower id), then each picked score over the sum. +[[kernel]] void tf_dsv41_route_select(const device float* LOGITS [[buffer(0)]], const device float* BIAS [[buffer(1)]], + const device float* BIAS_VL [[buffer(2)]], const device uchar* IMAGE [[buffer(3)]], + constant Dsv41RouteArgs& a [[buffer(4)]], device int* PICK [[buffer(5)]], + device float* WTS [[buffer(6)]], uint row [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_simdgroup]]) { + const int r = int(row); + if (r >= a.rows || a.topk > TF_DSV41_TOPK) return; + const device float* bias = IMAGE[r] != 0 ? BIAS_VL : BIAS; + float s[TF_DSV41_PER], ch[TF_DSV41_PER]; + int eid[TF_DSV41_PER]; + for (int j = 0; j < TF_DSV41_PER; j++) { + const int e = j * 32 + int(lane); + if (e < TF_DSV41_E) { + s[j] = metal::precise::sqrt(tf_dsv41_softplus(LOGITS[size_t(r) * TF_DSV41_E + e])); + ch[j] = s[j] + bias[e]; + eid[j] = e; + } else { + s[j] = 0.0f; + ch[j] = -INFINITY; + eid[j] = -1; + } + } + float w[TF_DSV41_TOPK]; + for (int k = 0; k < a.topk; k++) { + // The lane's own best: its ids ascend with j, so a strict > keeps the lower id on a tie. + float best = -INFINITY; + int bid = TF_DSV41_E; // nothing yet + for (int j = 0; j < TF_DSV41_PER; j++) { + const bool nan_wins = TF_CTL == 2 && ch[j] != ch[j]; // control 2: a NaN key beats every other + if (eid[j] >= 0 && (nan_wins || ch[j] > best || (ch[j] == best && eid[j] < bid))) { best = ch[j]; bid = eid[j]; } + } + // The butterfly: key descending then id ascending is a total order, so every lane ends with the same pick. + for (int off = 16; off > 0; off /= 2) { + const float ob = simd_shuffle_xor(best, off); + const int oi = simd_shuffle_xor(bid, off); + if (tf_dsv41_better(ob, oi, best, bid)) { best = ob; bid = oi; } + } + // The winner's own unbiased score: its lane holds it, and the other lanes take it from there. + float mine = 0.0f; + if (bid < TF_DSV41_E) { + for (int j = 0; j < TF_DSV41_PER; j++) { + if (eid[j] == bid) { mine = s[j]; eid[j] = -1; ch[j] = -INFINITY; } // the slot is spent + } + } + w[k] = bid < TF_DSV41_E ? simd_shuffle(mine, uint(bid & 31)) : 0.0f; + if (lane == 0) PICK[size_t(r) * a.topk + k] = bid < TF_DSV41_E ? bid : -1; + } + if (lane == 0) { + float total = w[0]; + for (int k = 1; k < a.topk; k++) total = total + w[k]; + for (int k = 0; k < a.topk; k++) WTS[size_t(r) * a.topk + k] = (w[k] / (total + a.eps)) * a.scale; + if (TF_CTL == 3) WTS[size_t(r) * a.topk] += bias[0]; // control 3: the bias added to the chosen weight + if (TF_CTL == 4 && r == a.rows - 1) WTS[size_t(r) * a.topk + a.topk] = 1.0f; // control 4: one weight past the last row's picks + } +} diff --git a/zig/kernels/metal/deepseek_v41/sources.zig b/zig/kernels/metal/deepseek_v41/sources.zig index 73252916fe..8ddd8b678b 100644 --- a/zig/kernels/metal/deepseek_v41/sources.zig +++ b/zig/kernels/metal/deepseek_v41/sources.zig @@ -11,8 +11,16 @@ pub const experts_kimi = @embedFile("experts_kimi.metal"); /// DeepSeek V4.1's routed-expert gate/up (clamped SwiGLU, the reference's activation rounding) over the copy. pub const experts = @embedFile("experts.metal"); pub const xmma = @embedFile("experts_mma.metal"); +/// DeepSeek-V4.1 sparse-attention indexer top-512 over fp32 score rows. +pub const topk512 = @embedFile("topk512.metal"); pub const glue = @embedFile("glue.metal"); /// DeepSeek-V4.1 router, indexer, compressor and head matvecs: lane chains over K, one butterfly, any row count. pub const mv = @embedFile("mv.metal"); +/// DeepSeek-V4.1 candidate blocks: block maxima, the newest block pinned, an exact top-2048 as a byte mask. +pub const cand = @embedFile("cand.metal"); +/// DeepSeek V4.1's indexer scores: FP4 query and key dots in one stated order, causal and block masks. +pub const index = @embedFile("index.metal"); +/// DeepSeek-V4.1's MoE route: the top-k of sqrt(softplus(logit)) + bias, per row (deepseek_v41/route.metal). +pub const route = @embedFile("route.metal"); /// DeepSeek-V4.1's phase-1 glue (deepseek_v41/glue_tmp.metal): the reference bits of the registered glue kernels. pub const glue_tmp = @embedFile("glue_tmp.metal"); diff --git a/zig/kernels/metal/deepseek_v41/topk512.metal b/zig/kernels/metal/deepseek_v41/topk512.metal new file mode 100644 index 0000000000..882f24c829 --- /dev/null +++ b/zig/kernels/metal/deepseek_v41/topk512.metal @@ -0,0 +1,109 @@ +// DeepSeek-V4.1 indexer top-512, a threadgroup a row: stable rank, NaN as -inf, -0 as +0, ascending indices, then -1. +#include +using namespace metal; + +#ifndef TF_CTL +#define TF_CTL 0 // planted faults: 1 NaN ranks highest, 2 ties go to the higher index, 3 last score dropped, 4 input write +#endif +#define TF_TG 1024u +#define TF_K 512u +#define TF_BINS 2048u + +struct TfTopk512Args { + uint width; // scores a row (the row stride) +}; + +inline uint tf_topk512_key(uint u) { + if ((u & 0x7fffffffu) > 0x7f800000u) u = TF_CTL == 1 ? 0x7f800000u : 0xff800000u; + if (u == 0x80000000u) u = 0u; + return (u >> 31) ? ~u : (u | 0x80000000u); +} + +// The bin holding the k-th largest key (bins scanned from the top) and the keys in the bins above it. +inline uint2 tf_topk512_pick(threadgroup atomic_uint* bins, uint k, threadgroup uint* sc, uint tid, uint lane, uint simd) { + const uint a = atomic_load_explicit(bins + (TF_BINS - 1u - 2u * tid), memory_order_relaxed); + const uint b = atomic_load_explicit(bins + (TF_BINS - 2u - 2u * tid), memory_order_relaxed); + const uint local = a + b, ex = simd_prefix_exclusive_sum(local); + if (lane == 31u) sc[simd] = ex + local; + threadgroup_barrier(mem_flags::mem_threadgroup); + uint before = ex; + for (uint s = 0; s < simd; ++s) before += sc[s]; + if (before < k && before + local >= k) { + const bool first = before + a >= k; + sc[32] = TF_BINS - (first ? 1u : 2u) - 2u * tid; + sc[33] = first ? before : before + a; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint2 r = uint2(sc[32], sc[33]); + threadgroup_barrier(mem_flags::mem_threadgroup); + return r; +} + +[[max_total_threads_per_threadgroup(TF_TG)]] kernel void tf_topk512( + constant TfTopk512Args& args [[buffer(0)]], + device const uint* scores [[buffer(1)]], + device const int* visible [[buffer(2)]], + device int* selected [[buffer(3)]], + uint row [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], + uint simd [[simdgroup_index_in_threadgroup]]) { + threadgroup atomic_uint bins[TF_BINS]; + threadgroup uint sc[34], part[64]; + const uint n = min(uint(max(visible[row], 0)), args.width); + device const uint* s = scores + ulong(row) * args.width; + device int* out = selected + ulong(row) * TF_K; + const uint keep = min(n, TF_K); + const uint m = (TF_CTL == 3 && n > 0u) ? n - 1u : n; +#if TF_CTL == 4 + if (tid == 0u && n > 0u) ((device uint*)s)[0] ^= 1u; +#endif + uint threshold = 0u, above = 0u, ties = 0u; + if (n > TF_K) { + uint prefix = 0u, k = TF_K; + for (uint pass = 0; pass < 3u; ++pass) { + const uint shift = pass == 0u ? 21u : (pass == 1u ? 10u : 0u); + const uint mask = pass == 2u ? 0x3ffu : 0x7ffu; + const uint up = pass == 1u ? 21u : 10u; + for (uint i = tid; i < TF_BINS; i += TF_TG) atomic_store_explicit(bins + i, 0u, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint i = tid; i < m; i += TF_TG) { + const uint key = tf_topk512_key(s[i]); + if (pass == 0u || (key >> up) == prefix) atomic_fetch_add_explicit(bins + ((key >> shift) & mask), 1u, memory_order_relaxed); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint2 r = tf_topk512_pick(bins, k, sc, tid, lane, simd); + prefix = (prefix << (pass == 2u ? 10u : 11u)) | r.x; + k -= r.y; + above += r.y; + } + threshold = prefix; + ties = atomic_load_explicit(bins + (prefix & 0x3ffu), memory_order_relaxed); + } + const uint wanted = keep - above; // keys equal to the threshold that make the set: the lowest indices + const uint skip = TF_CTL == 2 ? ties - wanted : 0u; + uint gt_run = 0u, eq_run = 0u; + for (uint base = 0; base < m && gt_run + min(eq_run, wanted) < keep; base += TF_TG) { + const uint i = base + tid; + const uint key = i < m ? tf_topk512_key(s[i]) : 0u; + const uint gt = (i < m && (n <= TF_K || key > threshold)) ? 1u : 0u; + const uint eq = (i < m && n > TF_K && key == threshold) ? 1u : 0u; + const uint packed = gt | (eq << 16); + const uint ex = simd_prefix_exclusive_sum(packed); + threadgroup uint* p = part + 32u * ((base / TF_TG) & 1u); + if (lane == 31u) p[simd] = ex + packed; + threadgroup_barrier(mem_flags::mem_threadgroup); + uint before = ex, total = 0u; + for (uint q = 0; q < 32u; ++q) { + const uint v = p[q]; + before += q < simd ? v : 0u; + total += v; + } + const uint gt_x = gt_run + (before & 0xffffu), eq_x = eq_run + (before >> 16); + const uint eq_taken = TF_CTL == 2 ? (eq_x > skip ? eq_x - skip : 0u) : min(eq_x, wanted); + if (gt != 0u || (eq != 0u && (TF_CTL == 2 ? eq_x >= skip : eq_x < wanted))) out[gt_x + eq_taken] = int(i); + gt_run += total & 0xffffu; + eq_run += total >> 16; + } + for (uint i = keep + tid; i < TF_K; i += TF_TG) out[i] = -1; +} diff --git a/zig/src/core/contract/dsv41_cand_records.zig b/zig/src/core/contract/dsv41_cand_records.zig new file mode 100644 index 0000000000..973d0ee795 --- /dev/null +++ b/zig/src/core/contract/dsv41_cand_records.zig @@ -0,0 +1,11 @@ +//! DeepSeek V4.1 candidate blocks: one 1,024-thread group a row, which the source reserves. +const std = @import("std"); +const M = @import("limits.zig").MetalLaunch; + +pub const records = [_]M{ + .{ .name = "dsv41 cand tf_cand", .threads = 1024, .tg_memory = 8592, .reserved = 1024, .m1 = 1024, .m2 = 1024 }, +}; + +test "the DeepSeek V4.1 candidate-block launch fits every Apple family" { + try std.testing.expectEqual(@as(usize, 0), try @import("registry.zig").checkMetal(std.testing.allocator, &records, true)); +} diff --git a/zig/src/core/contract/dsv41_index_records.zig b/zig/src/core/contract/dsv41_index_records.zig new file mode 100644 index 0000000000..72336652eb --- /dev/null +++ b/zig/src/core/contract/dsv41_index_records.zig @@ -0,0 +1,12 @@ +//! DeepSeek V4.1's indexer scores: prompt rows eight a simdgroup on MMA, decode rows eight simdgroups over one query. +const std = @import("std"); +const M = @import("limits.zig").MetalLaunch; + +pub const records = [_]M{ + .{ .name = "dsv41 ix scores tf_dsv41_ix_scores", .threads = 128, .tg_memory = 0, .reserved = 128 }, + .{ .name = "dsv41 ix scores tf_dsv41_ix_scores_dec", .threads = 256, .tg_memory = 16640, .reserved = 256 }, +}; + +test "the DeepSeek V4.1 indexer launches fit every Apple family" { + try std.testing.expectEqual(@as(usize, 0), try @import("registry.zig").checkMetal(std.testing.allocator, &records, true)); +} diff --git a/zig/src/core/contract/dsv41_metal_records.zig b/zig/src/core/contract/dsv41_metal_records.zig index b8a5250c15..2fd86c6f50 100644 --- a/zig/src/core/contract/dsv41_metal_records.zig +++ b/zig/src/core/contract/dsv41_metal_records.zig @@ -2,8 +2,9 @@ const M = @import("limits.zig").MetalLaunch; const head = @import("dsv41_fp8_records.zig").records ++ @import("dsv41_fp8_mma_records.zig").records ++ - @import("dsv41_attn_records.zig").records ++ @import("dsv41_experts_records.zig").records ++ - @import("dsv41_glue_records.zig").records ++ @import("dsv41_xmma_records.zig").records ++ - @import("dsv41_mv_records.zig").records; + @import("dsv41_topk512_records.zig").records ++ @import("dsv41_attn_records.zig").records ++ + @import("dsv41_experts_records.zig").records ++ @import("dsv41_index_records.zig").records ++ + @import("dsv41_glue_records.zig").records ++ @import("dsv41_cand_records.zig").records ++ + @import("dsv41_xmma_records.zig").records ++ @import("dsv41_mv_records.zig").records; -pub const records = head ++ @import("dsv41_glue_tmp_records.zig").records; +pub const records = head ++ @import("dsv41_glue_tmp_records.zig").records ++ @import("dsv41_route_records.zig").records; diff --git a/zig/src/core/contract/dsv41_route_records.zig b/zig/src/core/contract/dsv41_route_records.zig new file mode 100644 index 0000000000..e311c79739 --- /dev/null +++ b/zig/src/core/contract/dsv41_route_records.zig @@ -0,0 +1,7 @@ +//! DeepSeek V4.1's MoE route select launches (`tf_dsv41_route_select`), the backbone's and the drafter's. +const M = @import("limits.zig").MetalLaunch; + +pub const records = [_]M{ + .{ .name = "dsv41 route backbone tf_dsv41_route_select", .threads = 32, .tg_memory = 0, .m1 = 832, .m2 = 832 }, + .{ .name = "dsv41 route drafter tf_dsv41_route_select", .threads = 32, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, +}; diff --git a/zig/src/core/contract/dsv41_topk512_records.zig b/zig/src/core/contract/dsv41_topk512_records.zig new file mode 100644 index 0000000000..95df46da88 --- /dev/null +++ b/zig/src/core/contract/dsv41_topk512_records.zig @@ -0,0 +1,11 @@ +//! DeepSeek V4.1 top-512 over fp32 score rows: one 1,024-thread group a row, which the source reserves. +const std = @import("std"); +const M = @import("limits.zig").MetalLaunch; + +pub const records = [_]M{ + .{ .name = "dsv41 topk512 tf_topk512", .threads = 1024, .tg_memory = 8592, .reserved = 1024, .m1 = 1024, .m2 = 1024 }, +}; + +test "the DeepSeek V4.1 top-512 launch fits every Apple family" { + try std.testing.expectEqual(@as(usize, 0), try @import("registry.zig").checkMetal(std.testing.allocator, &records, true)); +} diff --git a/zig/src/families/deepseek_v41/metal/cand.zig b/zig/src/families/deepseek_v41/metal/cand.zig new file mode 100644 index 0000000000..57122d6caa --- /dev/null +++ b/zig/src/families/deepseek_v41/metal/cand.zig @@ -0,0 +1,28 @@ +//! The candidate-block selection of one launch: block maxima of the index scores, the exact top blocks as a byte mask. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +const glue = @import("glue.zig"); + +pub const Ref = glue.Ref; +/// Threads of the one group a row. +pub const threads = 1024; +/// The most blocks a row keeps. +pub const max_topk = 2048; + +pub fn load(device: mtl.Device) !mtl.Pipeline { + const lib = try mtl.Library.fromSource(device, sources.dsv41.cand, mtl.CompileOptions.mlx()); + defer lib.deinit(); + return mtl.Pipeline.init(device, lib, "tf_cand", false); +} + +/// Candidate blocks of scores [rows][keys] (vis[row] keys count): mask bytes, 1 = candidate; keyed u32 scratch. +pub fn encode(pipe: mtl.Pipeline, e: mtl.ComputeEncoder, scores: Ref, vis: Ref, keyed: Ref, mask: Ref, rows: u32, keys: u32, blocks: u32, topk: u32) void { + e.setPipeline(pipe); + e.setValue([3]u32{ keys, blocks, @min(topk, max_topk) }, 0); + e.setBuffer(scores.buf, scores.off, 1); + e.setBuffer(vis.buf, vis.off, 2); + e.setBuffer(keyed.buf, keyed.off, 3); + e.setBuffer(mask.buf, mask.off, 4); + e.dispatchGroups(mtl.Size.of(rows, 1, 1), mtl.Size.of(threads, 1, 1)); +} diff --git a/zig/src/families/deepseek_v41/metal/contract_exports.zig b/zig/src/families/deepseek_v41/metal/contract_exports.zig index 0cf4950e36..b361fe467d 100644 --- a/zig/src/families/deepseek_v41/metal/contract_exports.zig +++ b/zig/src/families/deepseek_v41/metal/contract_exports.zig @@ -1,5 +1,6 @@ //! DeepSeek-V4.1's Metal family behind one export: tf.dsv41_metal.NAME (families/deepseek_v41/metal). pub const fp8_mma = @import("fp8_mma.zig"); +pub const index = @import("index.zig"); pub const fp8_rows = @import("fp8_rows.zig"); pub const attn = @import("csa2_attn.zig"); pub const experts = @import("experts.zig"); diff --git a/zig/src/families/deepseek_v41/metal/index.zig b/zig/src/families/deepseek_v41/metal/index.zig new file mode 100644 index 0000000000..5b07b48e7f --- /dev/null +++ b/zig/src/families/deepseek_v41/metal/index.zig @@ -0,0 +1,123 @@ +//! DeepSeek V4.1's indexer scores on Metal: the pipeline, its arguments and the chunk a launch cuts its keys into. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); + +pub const heads = 32; // index_n_heads +pub const dim = 128; // index_head_dim +pub const block = 8; // candidate_block_size +/// Simdgroups a threadgroup carries (the source's IXR), each scoring a block of 8 rows. +pub const simdgroups = 4; +/// Rows a threadgroup scores. +pub const rows_group = 8 * simdgroups; +pub const max_rows = 8192; +pub const max_keys = 1048576; // max_position_embeddings + +/// A call: `rows` queries of `keys` keys; a weight row is w_stride bf16 elements apart, a key row ik_stride bytes. +pub const Call = struct { + rows: u32, + keys: u32, + start: u32 = 1, // a row r reaches (start + r) / ratio keys: 1 is prefill, keys is decode (all visible) + ratio: u32 = 1, + chunk: u32 = 0, // keys a simdgroup scores before moving on; 0: chunkOf. No bit depends on it. + w_stride: u32 = 32, + ik_stride: u32 = 68, // 64 FP4 bytes then 4 e8m0 scales + mask_stride: u32 = 0, // 0: no candidate mask; else one byte a block a row + w_scale: f32 = 0x1p-6, // softmax_scale * heads^-0.5 for 128 dims, 32 heads: exactly 2^-6, applied in the kernel +}; + +const Args = extern struct { + rows: u32, + keys: u32, + start: u32, + ratio: u32, + chunk: u32, + out_stride: u32, + w_stride: u32, + ik_stride: u32, + mask_stride: u32, + w_scale: f32, +}; + +pub const Ref = struct { buf: mtl.Buffer, off: usize = 0 }; + +pub const Pipelines = struct { + score: mtl.Pipeline, + /// decode's few rows: eight simdgroups a group over one decoded query (the same bits) + dec: mtl.Pipeline, + + pub fn load(device: mtl.Device) !Pipelines { + return loadSource(device, sources.dsv41.index); + } + + /// The same kernel from other text: the contract's planted controls prefix a TF_CTL define. + pub fn loadSource(device: mtl.Device, text: []const u8) !Pipelines { + const lib = try mtl.Library.fromSource(device, text, mtl.CompileOptions.mlx()); + defer lib.deinit(); + return .{ .score = try mtl.Pipeline.init(device, lib, "tf_dsv41_ix_scores", false), .dec = try mtl.Pipeline.init(device, lib, "tf_dsv41_ix_scores_dec", false) }; + } + + pub fn deinit(p: *Pipelines) void { + p.score.deinit(); + p.dec.deinit(); + } +}; + +pub fn check(c: Call) !void { + if (c.rows == 0 or c.rows > max_rows) return error.BadIxRows; + if (c.keys == 0 or c.keys > max_keys) return error.BadIxKeys; + if (c.ratio == 0 or c.chunk > c.keys) return error.BadIxChunk; + if (c.mask_stride != 0 and c.mask_stride * 8 < c.keys) return error.BadIxMask; +} + +/// Threadgroups a launch aims for, so a launch of few rows still fills the GPU. +const target_groups = 4000; + +/// The chunk one launch cuts a row's keys into: enough chunks for `target_groups`, between 32 and 1024 keys. +pub fn chunkOf(c: Call) u32 { + if (c.chunk != 0) return c.chunk; + const groups = (c.rows + rows_group - 1) / rows_group; + const chunks = target_groups / groups + 1; + return @min(c.keys, std.math.clamp((c.keys / chunks + 15) / 16 * 16, 32, 1024)); +} + +/// Q [rows][2048] FP4, QS [rows][128] e8m0, W bf16, IK [keys][ik_stride], S [rows][keys] fp32; CM if mask_stride != 0. +pub fn encode(p: *const Pipelines, e: mtl.ComputeEncoder, q: Ref, qs: Ref, w: Ref, ik: Ref, s: Ref, cm: Ref, c: Call) !void { + try check(c); + const args = Args{ + .rows = c.rows, + .keys = c.keys, + .start = c.start, + .ratio = c.ratio, + .chunk = chunkOf(c), + .out_stride = c.keys, + .w_stride = c.w_stride, + .ik_stride = c.ik_stride, + .mask_stride = c.mask_stride, + .w_scale = c.w_scale, + }; + e.setPipeline(p.score); + e.setBuffer(q.buf, q.off, 0); + e.setBuffer(qs.buf, qs.off, 1); + e.setBuffer(w.buf, w.off, 2); + e.setBuffer(ik.buf, ik.off, 3); + e.setBuffer(s.buf, s.off, 4); + e.setBuffer(cm.buf, cm.off, 5); + e.setValue(args, 6); + const chunks = (c.keys + args.chunk - 1) / args.chunk; + e.dispatchGroups(mtl.Size.of(chunks, (c.rows + rows_group - 1) / rows_group, 1), mtl.Size.of(32 * simdgroups, 1, 1)); +} + +/// Decode's threadgroup: eight simdgroups taking a chunk's keys in turn over one decoded query. +const dec_threads = 256; + +/// `encode` for decode's rows (a launch of one slot): groups of 32 keys a row, the same bits. +pub fn encodeDecode(p: *const Pipelines, e: mtl.ComputeEncoder, q: Ref, qs: Ref, w: Ref, ik: Ref, s: Ref, cm: Ref, c: Call) !void { + try check(c); + const chunk = if (c.chunk != 0) c.chunk else @min(c.keys, 32); + const args = Args{ .rows = c.rows, .keys = c.keys, .start = c.start, .ratio = c.ratio, .chunk = chunk, .out_stride = c.keys, .w_stride = c.w_stride, .ik_stride = c.ik_stride, .mask_stride = c.mask_stride, .w_scale = c.w_scale }; + e.setPipeline(p.dec); + for ([_]Ref{ q, qs, w, ik, s, cm }, 0..) |r, i| e.setBuffer(r.buf, r.off, i); + e.setValue(args, 6); + e.dispatchGroups(mtl.Size.of((c.keys + chunk - 1) / chunk, c.rows, 1), mtl.Size.of(dec_threads, 1, 1)); +} diff --git a/zig/tests/contract/dsv41_index_ref.zig b/zig/tests/contract/dsv41_index_ref.zig new file mode 100644 index 0000000000..291fca05c7 --- /dev/null +++ b/zig/tests/contract/dsv41_index_ref.zig @@ -0,0 +1,258 @@ +//! DeepSeek V4.1's indexer scores for the contract: inputs, the kernel's fp32 order as oracle, float64 reference. +const std = @import("std"); +const contract = @import("contract"); + +const c = contract.spec; +const rows = contract.rows; + +pub const heads = 32; +pub const dim = 128; +pub const block = 8; +pub const qrow = heads * dim / 2; // bytes of one query row: two FP4 codes a byte +pub const qsrow = heads * dim / 32; // e8m0 scales of one query row +pub const krow = dim / 2 + dim / 32; // bytes of one key: FP4 codes then four e8m0 scales +pub const w_scale: f32 = 0x1p-6; // softmax_scale * heads^-0.5 for 128 dims and 32 heads (model.py:513, 556) + +/// What bounds a query row's keys: all of them, a causal prompt, a candidate mask, or one decode step. +pub const Kind = enum { dense, causal, masked, decode }; + +pub const Geo = struct { start: u32, ratio: u32, mask_stride: u32 }; + +/// The call's arguments: row r reaches (start + r) / ratio keys (model.py:565 prefill, 568 decode), then the mask. +pub fn geometry(kind: Kind, nkeys: usize) Geo { + const k: u32 = @intCast(nkeys); + return switch (kind) { + .dense => .{ .start = k, .ratio = 1, .mask_stride = 0 }, + .causal => .{ .start = 1, .ratio = 4, .mask_stride = 0 }, + .masked => .{ .start = k, .ratio = 1, .mask_stride = (k + block - 1) / block + 3 }, + .decode => .{ .start = if (k * 128 > 130) k * 128 - 130 else 1, .ratio = 128, .mask_stride = 0 }, + }; +} + +pub fn e2m1(code: u8) f32 { + const e: u32 = (code >> 1) & 3; + const m: f32 = @floatFromInt(code & 1); + const v: f32 = if (e == 0) m * 0.5 else std.math.ldexp(m + 2, @as(i32, @intCast(e)) - 2); + return if (code & 8 != 0) -v else v; +} + +pub fn e8m0(code: u8) f32 { + if (code == 255) return std.math.nan(f32); + return std.math.ldexp(@as(f32, 1), @as(i32, code) - 127); +} + +/// ReLU keeping a NaN, as the kernel does (a max against zero would turn it into 0). +fn relu32(x: f32) f32 { + return if (x != x or x > 0) x else 0; +} + +fn relu64(x: f64) f64 { + return if (x != x or x > 0) x else 0; +} + +/// Longest chain of roundings: a 32-term group, the scale, the group tree, the product and the 32-head sum. +pub fn depth(_: c.Shape) usize { + return 32 + 1 + 2 + 1 + heads; +} + +pub const Out = struct { oracle: f32, exact: f64, mag: f64 }; + +/// Host inputs of one call: Q [rows][qrow], QS [rows][qsrow], W [rows][heads] bf16, IK [keys][krow], mask bytes. +pub const Ix = struct { + a: std.mem.Allocator, + shape: c.Shape, // rows, n = keys, k = dim + kind: Kind, + geo: Geo, + q: []u8, + qs: []u8, + w: []u16, + ik: []u8, + cm: []u8, + qf: []f32 = &.{}, // decoded query values times their group scale + wf: []f32 = &.{}, // weights times w_scale, as the kernel forms them + kv: []f32 = &.{}, // decoded key values + ksc: []f32 = &.{}, // key group scales + + pub fn init(a: std.mem.Allocator, shape: c.Shape, kind: Kind) !Ix { + const geo = geometry(kind, shape.n); + const q = try a.alloc(u8, shape.rows * qrow); + errdefer a.free(q); + const qs = try a.alloc(u8, shape.rows * qsrow); + errdefer a.free(qs); + const w = try a.alloc(u16, shape.rows * heads); + errdefer a.free(w); + const ik = try a.alloc(u8, shape.n * krow); + errdefer a.free(ik); + const cm = try a.alloc(u8, @max(16, shape.rows * geo.mask_stride)); + return .{ .a = a, .shape = shape, .kind = kind, .geo = geo, .q = q, .qs = qs, .w = w, .ik = ik, .cm = cm }; + } + + pub fn deinit(f: Ix) void { + for ([_][]const u8{ f.q, f.qs, f.ik, f.cm }) |s| f.a.free(s); + f.a.free(f.w); + for ([_][]const f32{ f.qf, f.wf, f.kv, f.ksc }) |s| f.a.free(s); + } + + pub fn outputs(f: Ix) usize { + return f.shape.rows * f.shape.n; + } + + pub fn fill(f: *Ix, mode: rows.Fill, seed: u64) void { + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (mode == .stress) return f.stress(rng); + for (f.q) |*b| b.* = rng.int(u8); + for (f.qs) |*b| b.* = 124 + rng.uintLessThan(u8, 7); + for (f.w) |*b| b.* = c.toBf16(rng.float(f32) * 2 - 0.5); + for (f.ik, 0..) |*b, i| b.* = if (i % krow >= dim / 2) 124 + rng.uintLessThan(u8, 7) else rng.int(u8); + f.fillMask(rng); + if (mode == .nonfinite) f.plant(); + } + + /// Row r % 7 == 3 has no live block, r % 11 == 5 all of them, the rest about a third; the padding reads as live. + fn fillMask(f: *Ix, rng: std.Random) void { + @memset(f.cm, 0xff); + const stride = f.geo.mask_stride; + if (stride == 0) return; + const blocks = (f.shape.n + block - 1) / block; + for (0..f.shape.rows) |r| for (0..blocks) |b| { + f.cm[r * stride + b] = if (r % 7 == 3) 0 else if (r % 11 == 5) 1 else @intFromBool(rng.uintLessThan(u8, 3) == 0); + }; + } + + /// Gate 1: a NaN query scale, an infinite and a NaN-making weight, a NaN key scale (every row meets that key). + fn plant(f: *Ix) void { + f.qs[3 * 4 + 1] = 255; // row 0, head 3, group 1 + f.w[1 * heads + 2] = 0x7f80; // row 1: +inf + f.w[2 * heads + 5] = 0xff80; // row 2: -inf + f.w[3 * heads + 7] = 0x7f80; // row 3: +inf and -inf together + f.w[3 * heads + 9] = 0xff80; + f.ik[(f.shape.n / 3) * krow + dim / 2 + 2] = 255; + } + + /// Gate 4's inputs: scales 1, FP4 magnitudes that make every sum exact in fp32 and half sums inexact, relu active. + fn stress(f: *Ix, rng: std.Random) void { + const big = [8]u8{ 1, 3, 5, 6, 7, 7, 7, 7 }; // 0.5 1.5 3 4 6 6 6 6 + @memset(f.qs, 127); + for (0..f.shape.rows) |r| for (0..heads) |h| for (0..dim / 2) |i| { + const sign: u8 = if (h % 5 == 4) 8 else 0; // a negative head: relu drops its score + const lo = big[rng.uintLessThan(usize, 8)] | sign; + const hi = big[rng.uintLessThan(usize, 8)] | sign; + f.q[r * qrow + h * (dim / 2) + i] = lo | (hi << 4); + }; + for (0..f.shape.n) |key| { + for (0..dim / 2) |i| f.ik[key * krow + i] = big[rng.uintLessThan(usize, 8)] | (big[rng.uintLessThan(usize, 8)] << 4); + @memset(f.ik[key * krow + dim / 2 ..][0 .. dim / 32], 127); + } + for (f.w) |*b| b.* = c.toBf16(@floatFromInt(1 + rng.uintLessThan(u32, 7))); + f.fillMask(rng); + } + + /// The decoded forms the oracle and the float64 reference read; built once a call. + pub fn decode(f: *Ix) !void { + if (f.qf.len != 0) return; + const nr = f.shape.rows; + f.qf = try f.a.alloc(f32, nr * heads * dim); + f.wf = try f.a.alloc(f32, nr * heads); + f.kv = try f.a.alloc(f32, f.shape.n * dim); + f.ksc = try f.a.alloc(f32, f.shape.n * 4); + for (0..nr) |r| for (0..heads * dim) |i| { + const h = i / dim; + const d = i % dim; + const byte = f.q[r * qrow + h * (dim / 2) + d / 2]; + f.qf[r * heads * dim + i] = e2m1((byte >> @intCast(4 * (d % 2))) & 15) * e8m0(f.qs[r * qsrow + h * 4 + d / 32]); + }; + for (f.wf, f.w) |*x, b| x.* = c.fromBf16(b) * w_scale; + for (0..f.shape.n) |key| { + for (0..dim) |d| f.kv[key * dim + d] = e2m1((f.ik[key * krow + d / 2] >> @intCast(4 * (d % 2))) & 15); + for (0..4) |g| f.ksc[key * 4 + g] = e8m0(f.ik[key * krow + dim / 2 + g]); + } + } + + pub fn live(f: Ix, r: usize, key: usize) bool { + if (key >= (@as(usize, f.geo.start) + r) / f.geo.ratio) return false; + return f.geo.mask_stride == 0 or f.cm[r * f.geo.mask_stride + key / block] != 0; + } + + /// Score (r, key) live: the kernel's fp32 order as oracle, the float64 sum, its terms' magnitude. + pub fn one(f: Ix, r: usize, key: usize) Out { + const kv = f.kv[key * dim ..][0..dim]; + const ks = f.ksc[key * 4 ..][0..4]; + var acc: f32 = 0; + var exact: f64 = 0; + var mag: f64 = 0; + for (0..heads) |h| { + const qh = f.qf[(r * heads + h) * dim ..][0..dim]; + var p: [4]f32 = @splat(0); + var pe: [4]f64 = @splat(0); + var pm: f64 = 0; + for (0..4) |g| { + var m: f64 = 0; + for (g * 32..g * 32 + 32) |i| { + p[g] += qh[i] * kv[i]; + pe[g] += @as(f64, qh[i]) * kv[i]; + m += @abs(@as(f64, qh[i]) * kv[i]); + } + pm += m * ks[g]; + } + const dot: f32 = ((p[0] * ks[0]) + (p[1] * ks[1])) + ((p[2] * ks[2]) + (p[3] * ks[3])); + const dote: f64 = ((pe[0] * ks[0]) + (pe[1] * ks[1])) + ((pe[2] * ks[2]) + (pe[3] * ks[3])); + const wh = f.wf[r * heads + h]; + acc += relu32(dot) * wh; + exact += relu64(dote) * wh; + mag += @abs(@as(f64, wh)) * pm; + } + return .{ .oracle = acc, .exact = exact, .mag = mag }; + } + + /// The oracle's scores as `want` (-inf where masked); each must sit inside the fp32 bound of float64. + pub fn reference(f: *Ix, want: []f64, mag: []f64) !void { + try f.decode(); + const d: f64 = @floatFromInt(depth(f.shape)); + const n = f.shape.n; + for (0..f.shape.rows) |r| for (0..n) |key| { + const at = r * n + key; + if (!f.live(r, key)) { + want[at] = -std.math.inf(f64); + mag[at] = 0; + continue; + } + const o = f.one(r, key); + if (std.math.isFinite(o.exact) and std.math.isFinite(o.mag) and @abs(@as(f64, o.oracle) - o.exact) > c.gamma(d, 0x1p-24) * o.mag + 0x1p-149) return error.OracleDrift; + want[at] = o.oracle; + mag[at] = o.mag; + }; + } + + /// On gate 4's inputs: each score as accumulators of `acc` leave it, every sum in `order` (heads, then dims). + pub fn simulate(f: *Ix, acc: c.Accum, order: rows.Order, out: []f64) !void { + try f.decode(); + const n = f.shape.n; + var t: [dim]f64 = undefined; + var hs: [heads]f64 = undefined; + for (0..f.shape.rows) |r| for (0..n) |key| { + if (!f.live(r, key)) { + out[r * n + key] = -std.math.inf(f64); + continue; + } + for (0..heads) |h| { + const qh = f.qf[(r * heads + h) * dim ..][0..dim]; + for (0..dim) |i| t[i] = @as(f64, qh[i]) * f.kv[key * dim + i] * f.ksc[key * 4 + i / 32]; + const dot = rows.sumIn(acc, order, &t); + hs[h] = relu64(dot) * f.wf[r * heads + h]; + } + out[r * n + key] = rows.sumIn(acc, order, &hs); + }; + } +}; + +test "FP4 and e8m0 codes decode to their values" { + const want = [8]f32{ 0, 0.5, 1, 1.5, 2, 3, 4, 6 }; + for (want, 0..) |v, i| { + try std.testing.expectEqual(v, e2m1(@intCast(i))); + try std.testing.expectEqual(-v, e2m1(@intCast(i + 8))); + } + try std.testing.expectEqual(@as(f32, 1), e8m0(127)); + try std.testing.expectEqual(@as(f32, 0x1p-6), e8m0(121)); + try std.testing.expect(std.math.isNan(e8m0(255))); +} diff --git a/zig/tests/contract/metal_dsv41_cand.zig b/zig/tests/contract/metal_dsv41_cand.zig new file mode 100644 index 0000000000..98d352facc --- /dev/null +++ b/zig/tests/contract/metal_dsv41_cand.zig @@ -0,0 +1,366 @@ +//! DeepSeek-V4.1's candidate blocks under the kernel contract on Metal: stable-sort check, controls, splits, timing. +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const sources = @import("kernel_sources"); +const rig_ = @import("rig.zig"); + +const c = contract.spec; +const g = contract.gates; +const base = "dsv41 cand "; +const inf = std.math.inf(f32); +const rule = "NaN ranks as -inf, -0 as +0; block maxima over 8 keys, the last visible block pinned first, ties to the lower block"; + +fn noSums(_: c.Shape) usize { + return 0; +} + +/// Shapes are (rows, topk, blocks); keys are 8 a block less 3 when blocks is odd; stride 4 over when blocks % 8 == 5. +const served = blk: { + const rows = [_]usize{ 1, 4, 64, 4096, 8192 }; + const blocks = [_]usize{ 64, 255, 2048, 2049, 2053, 16384 }; + var out: [rows.len * blocks.len + 12]c.Shape = undefined; + var n: usize = 0; + out[n] = .{ .rows = 8, .n = 2048, .k = 3000 }; + n += 1; + for (rows) |r| for (blocks) |b| { + if (r * b > (1 << 21) or (r > 64 and b > 255 and b < 16384)) continue; + out[n] = .{ .rows = r, .n = 2048, .k = b }; + n += 1; + }; + for ([_]c.Shape{ + .{ .rows = 16, .n = 1, .k = 300 }, + .{ .rows = 16, .n = 100, .k = 5000 }, + .{ .rows = 16, .n = 2047, .k = 2048 }, + .{ .rows = 5, .n = 2048, .k = 2049 }, + .{ .rows = 1, .n = 2048, .k = 131071 }, + .{ .rows = 1, .n = 2048, .k = 131072 }, + .{ .rows = 4, .n = 2048, .k = 131072 }, + }) |s| { + out[n] = s; + n += 1; + } + break :blk out[0..n].*; +}; + +pub fn spec(name: []const u8) c.Spec { + return .{ + .name = name, + .accum = .none, + .depth = noSums, + .out = .raw, + .nan = .{ .stated = rule }, + .tiles = .{ .{}, .{}, .{ .size = 4096, .step = 64, .min = 64 } }, + .served = &served, + }; +} + +fn rank(v: f32) f32 { + return if (std.math.isNan(v)) -inf else v; +} + +const Order = struct { + maxima: []const f32, + fn before(o: Order, a: u32, b: u32) bool { + return o.maxima[a] > o.maxima[b] or (o.maxima[a] == o.maxima[b] and a < b); + } +}; + +pub const Cand = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + ctl: u8, + full: bool = false, + shape: c.Shape = undefined, + keys: usize = 0, + stride: usize = 0, + scores: []f32 = &.{}, + vis: []i32 = &.{}, + plants: usize = 0, + s: rig_.Ref = undefined, + v: rig_.Ref = undefined, + kd: rig_.Ref = undefined, + out: rig_.Ref = undefined, + msg: [200]u8 = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, ctl: u8) !Cand { + const text = try std.fmt.allocPrint(a, "#define TF_CTL {d}\n{s}", .{ ctl, sources.dsv41.cand }); + defer a.free(text); + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .pipe = try rig_.pipeline(device, text, "tf_cand"), .ctl = ctl }; + } + + pub fn deinit(t: *Cand) void { + t.a.free(t.scores); + t.a.free(t.vis); + t.pipe.deinit(); + t.rig.deinit(); + } + + pub fn kernel(t: *Cand, name: []const u8) g.Kernel { + const control: ?g.Gate = switch (t.ctl) { + 1 => .nonfinite, + 2, 3 => .shapes, + 4 => .repeat, + else => null, + }; + return .{ .spec = spec(name), .ctx = t, .vt = &vt, .control = control }; + } + + const vt = g.hooks(Cand); + + pub fn prepare(t: *Cand, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.n < 1 or shape.n > 2048 or shape.k < 1) return false; + t.rig.reset(); + t.shape = shape; + const nb = shape.k; + t.keys = 8 * nb - (if (nb % 2 == 1) @as(usize, 3) else 0); + t.stride = nb + (if (nb % 8 == 5) @as(usize, 4) else 0); + const keys = t.keys; + t.a.free(t.scores); + t.a.free(t.vis); + t.scores = &.{}; + t.vis = &.{}; + t.scores = try t.a.alloc(f32, shape.rows * keys); + t.vis = try t.a.alloc(i32, shape.rows); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (0..shape.rows) |r| { + const row = t.scores[r * keys ..][0..keys]; + const vis: usize = if (t.full or r < 4) keys else switch (r % 16) { + 1 => keys / 2 + 1, + 6 => @min(keys, 8 * 2048), + 7 => @min(keys, 8 * 2048 + 1), + 9 => 0, + 13 => @min(keys, 1), + else => keys, + }; + t.vis[r] = @intCast(vis); + const style: usize = if (t.full) 0 else r % 5; + for (row, 0..) |*x, i| x.* = if (i >= vis) -inf else switch (style) { + 0 => @as(f32, @floatFromInt(rng.uintLessThan(u32, 512))) / 16 - 16, + 1 => 1.5, + 2 => (rng.float(f32) - 0.5) * 1e6, + 3 => -@as(f32, @floatFromInt(i / 3)), + else => switch (rng.uintLessThan(u32, 8)) { + 0 => 1.0, + 1...4 => 0.0, + else => -0.0, + }, + }; + } + t.plants = 0; + if (fill == .nonfinite) { + // an all-NaN block, a lone +inf, an all -inf block, a lone NaN and an all-NaN last block (which is pinned) + const at = [_][2]usize{ .{ 0, nb / 3 }, .{ 1, nb / 2 }, .{ 2, nb - 2 }, .{ 3, 0 }, .{ 4, nb - 1 }, .{ 3, nb - 1 } }; + const put = [_]f32{ std.math.nan(f32), inf, -inf, std.math.nan(f32), std.math.nan(f32), std.math.nan(f32) }; + for (at, put, 0..) |p, v, n| if (p[0] < shape.rows and p[1] < nb and t.vis[p[0]] > 0) { + const row = t.scores[p[0] * keys ..][0..keys]; + const whole = n != 1 and n != 3; + for (row[8 * p[1] .. @min(keys, 8 * p[1] + (if (whole) @as(usize, 8) else 1))]) |*x| x.* = v; + t.plants += 1; + }; + } + t.s = try t.rig.input(std.mem.sliceAsBytes(t.scores)); + t.v = try t.rig.input(std.mem.sliceAsBytes(t.vis)); + t.kd = try t.rig.scratch(shape.rows * t.stride * 4); + t.out = try t.rig.output(shape.rows * t.stride); + return true; + } + + /// Rows r0..r0+rows of the prepared call, as one dispatch. + pub fn encode(t: *Cand, e: anytype, r0: usize, rows: usize) void { + const args = [3]u32{ @intCast(t.keys), @intCast(t.stride), @intCast(t.shape.n) }; + e.setPipeline(t.pipe); + e.setBuffer(t.s.buf, t.s.off + r0 * t.keys * 4, 1); + e.setBuffer(t.v.buf, t.v.off + r0 * 4, 2); + e.setBuffer(t.kd.buf, t.kd.off + r0 * t.stride * 4, 3); + e.setBuffer(t.out.buf, t.out.off + r0 * t.stride, 4); + e.setValue(args, 0); + e.dispatchGroups(mtl.Size.of(rows, 1, 1), mtl.Size.of(1024, 1, 1)); + } + + pub fn run(t: *Cand) !void { + @memset(t.rig.bytes(t.out, t.count()), 0x55); + const pass = t.rig.begin(); + t.encode(pass.e, 0, t.shape.rows); + try t.rig.finish(pass); + } + + pub fn count(t: *Cand) usize { + return t.shape.rows * t.stride; + } + + pub fn reference(_: *Cand, want: []f64, mag: []f64) !void { + @memset(want, 0); + @memset(mag, 0); + } + + /// The CPU stable sort against each row: all rows up to 512K blocks, else a stride of rows and the last. + pub fn check(t: *Cand, at: g.Gate) !?[]const u8 { + const nb = t.shape.k; + const keys = t.keys; + const got = t.rig.bytes(t.out, t.count()); + const step = @max(1, t.shape.rows * nb / (1 << 19)); + const idx = try t.a.alloc(u32, nb); + defer t.a.free(idx); + const bm = try t.a.alloc(f32, nb); + defer t.a.free(bm); + const inset = try t.a.alloc(bool, nb); + defer t.a.free(inset); + var nonfinite: usize = 0; + for (0..t.shape.rows) |r| { + if (r >= 8 and r % step != 0 and r + 1 != t.shape.rows) continue; + const row = t.scores[r * keys ..][0..keys]; + for (row) |x| nonfinite += @intFromBool(!std.math.isFinite(x)); + for (bm, idx, 0..) |*m, *x, j| { + m.* = -inf; + for (row[8 * j .. @min(keys, 8 * j + 8)]) |v| m.* = @max(m.*, rank(v)); + x.* = @intCast(j); + } + const vis: usize = @intCast(t.vis[r]); + if (vis >= 1 and (vis - 1) / 8 < nb) bm[(vis - 1) / 8] = inf; + std.sort.pdq(u32, idx, Order{ .maxima = bm }, Order.before); + const keep = @min(t.shape.n, nb); + const out = got[r * t.stride ..][0..t.stride]; + @memset(inset, false); + for (idx[0..keep]) |j| inset[j] = bm[j] > -inf; + for (out, 0..) |have, j| { + const w: u8 = if (j < nb and inset[j]) 1 else 0; + if (have != w) return try std.fmt.bufPrint(&t.msg, "row {d} (visible {d} keys, {d} blocks, top {d}) block {d}: got {d}, stable rank gives {d}", .{ r, vis, nb, t.shape.n, j, have, w }); + } + } + if (at == .nonfinite and (t.plants < 5 or nonfinite < 5)) return try std.fmt.bufPrint(&t.msg, "{d} plants: no proof", .{t.plants}); + return null; + } +}; + +/// Every registered variant, the launch-split probe after the gates, the glue_tmp equality and the timing. +pub const Suite = struct { + a: std.mem.Allocator, + device: mtl.Device, + queue: mtl.Queue, + probes: std.ArrayList(*Cand) = .empty, + + const list = [_]struct { ctl: u8, name: []const u8 }{ + .{ .ctl = 0, .name = base ++ "tf_cand" }, + .{ .ctl = 1, .name = base ++ "control: NaN ranked above +inf" }, + .{ .ctl = 2, .name = base ++ "control: ties go to the higher block" }, + .{ .ctl = 3, .name = base ++ "control: the newest block not pinned" }, + .{ .ctl = 4, .name = base ++ "control: writes into its input" }, + }; + + pub fn create(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, kernels: *std.ArrayList(g.Kernel), only: ?[]const u8) !*Suite { + const s = try a.create(Suite); + s.* = .{ .a = a, .device = device, .queue = queue }; + for (list) |entry| { + if (only) |text| if (std.mem.indexOf(u8, entry.name, text) == null) continue; + const p = try a.create(Cand); + p.* = try Cand.init(a, device, queue, entry.ctl); + try s.probes.append(a, p); + try kernels.append(a, p.kernel(entry.name)); + } + return s; + } + + pub fn destroy(s: *Suite) void { + for (s.probes.items) |p| { + p.deinit(); + s.a.destroy(p); + } + s.probes.deinit(s.a); + s.a.destroy(s); + } + + /// Rows split into launches of 1, 7 and rows/3, and one launch a row, equal the single launch byte for byte. + pub fn invariance(s: *Suite, seed: u64) !usize { + const p = if (s.probes.items.len > 0) s.probes.items[0] else return 0; + if (p.ctl != 0) return 0; + var bad: usize = 0; + var launches: usize = 0; + for ([_]c.Shape{ .{ .rows = 64, .n = 2048, .k = 16384 }, .{ .rows = 33, .n = 2048, .k = 2053 }, .{ .rows = 5, .n = 2048, .k = 131071 } }) |shape| { + if (!try p.prepare(shape, .random, seed)) return error.Refused; + try p.run(); + const whole = try s.a.dupe(u8, p.rig.bytes(p.out, p.count())); + defer s.a.free(whole); + for ([_]usize{ 1, 7, shape.rows / 3 }) |step| { + @memset(p.rig.bytes(p.out, p.count()), 0x55); + const pass = p.rig.begin(); + var r0: usize = 0; + while (r0 < shape.rows) : (r0 += step) { + p.encode(pass.e, r0, @min(step, shape.rows - r0)); + launches += 1; + } + try p.rig.finish(pass); + if (!std.mem.eql(u8, whole, p.rig.bytes(p.out, p.count()))) { + bad += 1; + std.debug.print("FAIL launch split {d}x{d}: {d} rows a launch differs from one launch\n", .{ shape.rows, shape.k, step }); + } + } + } + const word = if (bad == 0) "PASS" else "FAIL"; + std.debug.print("{s} launch split {s}: 3 shapes at 1, 7 and rows/3 rows a launch ({d} launches) byte-equal to one launch, {d} differ\n", .{ word, base, launches, bad }); + return bad; + } + + /// The O(blocks^2) glue_tmp kernel and this one give the same mask bytes over glue_tmp's envelope. + pub fn equality(s: *Suite, seed: u64) !usize { + const p = if (s.probes.items.len > 0) s.probes.items[0] else return 0; + if (p.ctl != 0) return 0; + const old = try rig_.pipeline(s.device, sources.dsv41.glue_tmp, "tf_dsv41_g_cand"); + defer old.deinit(); + var bad: usize = 0; + var cases: usize = 0; + const blocks = [_]usize{ 1, 7, 64, 255, 2048, 2049, 2053, 3000, 4100 }; + const tops = [_]usize{ 2048, 100, 1 }; + for (blocks, 0..) |nb, bi| for (tops, 0..) |top, ti| { + const shape: c.Shape = .{ .rows = 1 + (bi * 3 + ti) % 9, .n = top, .k = nb }; + if (!try p.prepare(shape, .random, seed + 100 * bi + ti)) return error.Refused; + try p.run(); + const bs = try p.rig.scratch(shape.rows * p.stride * 4); + const want = try p.rig.output(p.count()); + const pass = p.rig.begin(); + const e = pass.e; + e.setPipeline(old); + e.setBuffer(p.s.buf, p.s.off, 0); + e.setBuffer(p.v.buf, p.v.off, 1); + e.setBuffer(bs.buf, bs.off, 2); + e.setBuffer(want.buf, want.off, 3); + e.setValue([3]u32{ @intCast(p.keys), @intCast(p.stride), @intCast(top) }, 4); + e.dispatchGroups(mtl.Size.of(shape.rows, 1, 1), mtl.Size.of(256, 1, 1)); + try p.rig.finish(pass); + cases += 1; + const at = std.mem.indexOfDiff(u8, p.rig.bytes(p.out, p.count()), p.rig.bytes(want, p.count())); + if (at) |i| { + bad += 1; + std.debug.print("FAIL glue_tmp equality {d} rows, {d} blocks, top {d}: first difference at byte {d}\n", .{ shape.rows, nb, top, i }); + } + }; + const word = if (bad == 0) "PASS" else "FAIL"; + std.debug.print("{s} equality {s}: {d} cases of 1..9 rows, 1..4100 blocks and top 1, 100, 2048 byte-equal to glue_tmp's tf_dsv41_g_cand, {d} differ\n", .{ word, base, cases, bad }); + return bad; + } + + /// Median microseconds a launch, over launches queued back to back, at full visible length. + pub fn timing(s: *Suite, io: std.Io) !void { + var p = try Cand.init(s.a, s.device, s.queue, 0); + defer p.deinit(); + p.full = true; + for ([_]usize{ 16384, 131072 }) |blocks| for ([_]usize{ 1, 16 }) |rows| { + if (!try p.prepare(.{ .rows = rows, .n = 2048, .k = blocks }, .random, 7)) return error.Refused; + const reps: usize = if (rows == 1) 32 else 4; + var samples: [9]f64 = undefined; + for (&samples) |*out| { + const t0 = std.Io.Clock.awake.now(io); + const pass = p.rig.begin(); + for (0..reps) |_| p.encode(pass.e, 0, rows); + try p.rig.finish(pass); + const ns: f64 = @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(io)).toNanoseconds()); + out.* = ns / 1000.0 / @as(f64, @floatFromInt(reps)); + } + std.mem.sort(f64, &samples, {}, std.sort.asc(f64)); + const med = samples[samples.len / 2]; + std.debug.print("TIME {s}: blocks {d} rows {d}: {d:.1} us a launch, {d:.3} us a row (median of {d} x {d} launches)\n", .{ base, blocks, rows, med, med / @as(f64, @floatFromInt(rows)), samples.len, reps }); + }; + } +}; diff --git a/zig/tests/contract/metal_dsv41_index.zig b/zig/tests/contract/metal_dsv41_index.zig new file mode 100644 index 0000000000..40fc1ec893 --- /dev/null +++ b/zig/tests/contract/metal_dsv41_index.zig @@ -0,0 +1,286 @@ +//! DeepSeek V4.1's indexer scores under the contract on Metal: exact fp32 order, masks, controls, invariance, timing. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const sources = @import("kernel_sources"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const ref = @import("dsv41_index_ref.zig"); + +const c = contract.spec; +const g = contract.gates; +const ix = tf.dsv41_metal.index; +const base = "dsv41 ix "; + +/// What a registration varies: how a row's keys are bounded, and a planted fault. +pub const Variant = struct { kind: ref.Kind = .dense, ctl: u8 = 0 }; + +// Pairs (rows x keys) stay near a million: the float64 reference costs ~4,000 products a pair. +const wide = [_]c.Shape{ + .{ .rows = 8, .n = 256, .k = 128 }, // first: the small gates use it + .{ .rows = 1, .n = 128, .k = 128 }, + .{ .rows = 1, .n = 129, .k = 128 }, + .{ .rows = 1, .n = 1000, .k = 128 }, + .{ .rows = 3, .n = 513, .k = 128 }, + .{ .rows = 16, .n = 4097, .k = 128 }, + .{ .rows = 33, .n = 1024, .k = 128 }, + .{ .rows = 64, .n = 4096, .k = 128 }, + .{ .rows = 1024, .n = 1024, .k = 128 }, + .{ .rows = 8192, .n = 128, .k = 128 }, + .{ .rows = 1, .n = 131072, .k = 128 }, + .{ .rows = 2, .n = 131072, .k = 128 }, +}; +const steps = [_]c.Shape{ + .{ .rows = 8, .n = 256, .k = 128 }, + .{ .rows = 1, .n = 128, .k = 128 }, + .{ .rows = 1, .n = 1000, .k = 128 }, + .{ .rows = 2, .n = 4097, .k = 128 }, + .{ .rows = 3, .n = 513, .k = 128 }, + .{ .rows = 8, .n = 8192, .k = 128 }, + .{ .rows = 4, .n = 65536, .k = 128 }, + .{ .rows = 1, .n = 131072, .k = 128 }, +}; + +pub fn spec(v: Variant, name: []const u8) c.Spec { + return .{ .name = name, .accum = .f32, .depth = ref.depth, .out = .f32, .served = if (v.kind == .decode) &steps else &wide }; +} + +pub const Proj = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + v: Variant, + name: []const u8, + pipes: ix.Pipelines, + f: ?ref.Ix = null, + q: rig_.Ref = undefined, + qs: rig_.Ref = undefined, + w: rig_.Ref = undefined, + ik: rig_.Ref = undefined, + cm: rig_.Ref = undefined, + s: rig_.Ref = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, v: Variant, name: []const u8) !Proj { + var p: Proj = .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .v = v, .name = name, .pipes = undefined }; + if (v.ctl == 0) { + p.pipes = try ix.Pipelines.load(device); + } else { + const text = try std.fmt.allocPrint(a, "#define TF_CTL {d}\n{s}", .{ v.ctl, sources.dsv41.index }); + defer a.free(text); + p.pipes = try ix.Pipelines.loadSource(device, text); + } + return p; + } + + pub fn deinit(p: *Proj) void { + p.drop(); + p.pipes.deinit(); + p.rig.deinit(); + } + + fn drop(p: *Proj) void { + if (p.f) |f| f.deinit(); + p.f = null; + p.rig.reset(); + } + + pub fn kernel(p: *Proj) g.Kernel { + const control: ?g.Gate = switch (p.v.ctl) { + 1 => .nonfinite, + 2 => .shapes, + 3 => .sums, + 4 => .repeat, + else => null, + }; + return .{ .spec = spec(p.v, p.name), .ctx = p, .vt = &vt, .control = control }; + } + + const vt = g.hooks(Proj); + + /// Rows r0.. of the prepared call, with chunk 0 the helper's own. + fn call(p: *const Proj, r0: usize, nrows: usize, chunk: u32) ix.Call { + const f = p.f.?; + return .{ + .rows = @intCast(nrows), + .keys = @intCast(f.shape.n), + .start = f.geo.start + @as(u32, @intCast(r0)), + .ratio = f.geo.ratio, + .chunk = chunk, + .mask_stride = f.geo.mask_stride, + .w_scale = ref.w_scale, + }; + } + + pub fn prepare(p: *Proj, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + p.drop(); + if (shape.k != ref.dim) return false; + ix.check(.{ .rows = @intCast(shape.rows), .keys = @intCast(shape.n) }) catch return false; + var f = try ref.Ix.init(p.a, shape, p.v.kind); + f.fill(fill, seed); + p.f = f; + p.q = try p.rig.input(f.q); + p.qs = try p.rig.input(f.qs); + p.w = try p.rig.input(std.mem.sliceAsBytes(f.w)); + p.ik = try p.rig.input(f.ik); + p.cm = try p.rig.input(f.cm); + p.s = try p.rig.output(f.outputs() * 4); + return true; + } + + fn at(r: rig_.Ref, extra: usize) ix.Ref { + return .{ .buf = r.buf, .off = r.off + extra }; + } + + /// Rows r0..r0+nrows as one launch: every input offset by its rows, the output the same. + pub fn encode(p: *Proj, e: anytype, r0: usize, nrows: usize, chunk: u32) !void { + const f = p.f.?; + const n = f.shape.n; + const cm = at(p.cm, r0 * f.geo.mask_stride); + try ix.encode(&p.pipes, e, at(p.q, r0 * ref.qrow), at(p.qs, r0 * ref.qsrow), at(p.w, r0 * ref.heads * 2), at(p.ik, 0), at(p.s, r0 * n * 4), cm, p.call(r0, nrows, chunk)); + } + + pub fn run(p: *Proj) !void { + const pass = p.rig.begin(); + try p.encode(pass.e, 0, p.f.?.shape.rows, 0); + try p.rig.finish(pass); + } + + pub fn count(p: *Proj) usize { + return p.f.?.outputs(); + } + + pub fn reference(p: *Proj, want: []f64, mag: []f64) !void { + return p.f.?.reference(want, mag); + } + + /// Exact order: every score must equal the oracle's bits, so no allowance at all. + pub fn bounds(_: *Proj, _: []const f64, _: []const f64, allow: []f64) !void { + @memset(allow, 0); + } + + pub fn simulate(p: *Proj, acc: c.Accum, order: g.Order, out: []f64) !void { + return p.f.?.simulate(acc, order, out); + } +}; + +/// Every registered variant, the invariance probes that run after the gates, and the timing. +pub const Suite = struct { + a: std.mem.Allocator, + device: mtl.Device, + queue: mtl.Queue, + probes: std.ArrayList(*Proj) = .empty, + + const list = [_]struct { v: Variant, name: []const u8 }{ + .{ .v = .{}, .name = base ++ "dense" }, + .{ .v = .{ .kind = .causal }, .name = base ++ "causal prompt" }, + .{ .v = .{ .kind = .masked }, .name = base ++ "candidate-masked" }, + .{ .v = .{ .kind = .decode }, .name = base ++ "decode" }, + .{ .v = .{ .ctl = 1 }, .name = base ++ "control: relu turns a NaN score into 0" }, + .{ .v = .{ .ctl = 2 }, .name = base ++ "control: heads folded in reverse, head 0 dropped" }, + .{ .v = .{ .ctl = 3 }, .name = base ++ "control: group sums in fp16" }, + .{ .v = .{ .ctl = 4 }, .name = base ++ "control: writes past its last output" }, + }; + + pub fn create(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, kernels: *std.ArrayList(g.Kernel), only: ?[]const u8) !*Suite { + const s = try a.create(Suite); + s.* = .{ .a = a, .device = device, .queue = queue }; + for (list) |entry| { + if (only) |text| if (std.mem.indexOf(u8, entry.name, text) == null) continue; + const p = try a.create(Proj); + p.* = try Proj.init(a, device, queue, entry.v, entry.name); + try s.probes.append(a, p); + try kernels.append(a, p.kernel()); + } + return s; + } + + pub fn destroy(s: *Suite) void { + for (s.probes.items) |p| { + p.deinit(); + s.a.destroy(p); + } + s.probes.deinit(s.a); + s.a.destroy(s); + } + + fn launch(p: *Proj, r0: usize, nrows: usize, chunk: u32) !void { + const pass = p.rig.begin(); + try p.encode(pass.e, r0, nrows, chunk); + try p.rig.finish(pass); + } + + fn poison(p: *Proj) void { + @memset(p.rig.bytes(p.s, p.count() * 4), 0xff); + } + + fn report(word: []const u8, what: []const u8, kind: ref.Kind, bad: usize, launches: usize) void { + std.debug.print("{s} {s} {s}{t}: {d} launches byte-equal to the reference launch, {d} differ\n", .{ word, what, base, kind, launches, bad }); + } + + /// Row counts, chunk sizes and prompt-vs-decode launches equal one 2048-row launch byte for byte; returns misses. + pub fn invariance(s: *Suite, seed: u64) !usize { + var bad_all: usize = 0; + for (s.probes.items) |p| { + if (p.v.ctl != 0 or p.v.kind == .decode) continue; + var bad: usize = 0; + var launches: usize = 0; + const shape: c.Shape = .{ .rows = 2048, .n = 1000, .k = 128 }; + if (!try p.prepare(shape, .random, seed)) return error.Refused; + const n = shape.n; + try launch(p, 0, shape.rows, 0); + const whole = try s.a.dupe(u8, p.rig.bytes(p.s, p.count() * 4)); + defer s.a.free(whole); + for ([_]usize{ 1, 3, 16, 2048 }) |r| { + poison(p); + try launch(p, 0, r, 0); + launches += 1; + if (!std.mem.eql(u8, whole[0 .. r * n * 4], p.rig.bytes(p.s, r * n * 4))) bad += 1; + } + report(if (bad == 0) "PASS" else "FAIL", "row invariance", p.v.kind, bad, launches); + bad_all += bad; + bad = 0; + launches = 0; + for ([_]u32{ 1, 7, 64, 255, 256, 1000 }) |chunk| { + poison(p); + try launch(p, 0, shape.rows, chunk); + launches += 1; + if (!std.mem.eql(u8, whole, p.rig.bytes(p.s, whole.len))) bad += 1; + } + report(if (bad == 0) "PASS" else "FAIL", "chunk invariance (1,7,64,255,256,1000 keys)", p.v.kind, bad, launches); + bad_all += bad; + bad = 0; + poison(p); + const pass = p.rig.begin(); + for (0..shape.rows) |r| try p.encode(pass.e, r, 1, 0); + try p.rig.finish(pass); + launches = 1; + if (!std.mem.eql(u8, whole, p.rig.bytes(p.s, whole.len))) bad += 1; + report(if (bad == 0) "PASS" else "FAIL", "prompt vs one decode launch a row", p.v.kind, bad, launches); + bad_all += bad; + } + return bad_all; + } + + /// Median microseconds a launch over launches queued back to back, every key live. + pub fn timing(s: *Suite, io: std.Io) !void { + var p = try Proj.init(s.a, s.device, s.queue, .{}, "timing"); + defer p.deinit(); + const cases = [_][2]usize{ .{ 1, 8192 }, .{ 1, 32768 }, .{ 1, 131072 }, .{ 8192, 32768 } }; + for (cases) |case| { + if (!try p.prepare(.{ .rows = case[0], .n = case[1], .k = 128 }, .random, 7)) return error.Refused; + const reps: usize = if (case[0] == 1) 32 else 2; + var samples: [9]f64 = undefined; + for (&samples) |*out| { + const t0 = std.Io.Clock.awake.now(io); + const pass = p.rig.begin(); + for (0..reps) |_| try p.encode(pass.e, 0, case[0], 0); + try p.rig.finish(pass); + const ns: f64 = @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(io)).toNanoseconds()); + out.* = ns / 1000.0 / @as(f64, @floatFromInt(reps)); + } + std.mem.sort(f64, &samples, {}, std.sort.asc(f64)); + const us = samples[samples.len / 2]; + const pairs: f64 = @floatFromInt(case[0] * case[1]); + std.debug.print("TIME {s}: rows {d} keys {d}: {d:.1} us a launch, {d:.2} ns a (row, key) pair (median of {d} x {d})\n", .{ base, case[0], case[1], us, us * 1000.0 / pairs, samples.len, reps }); + } + } +}; diff --git a/zig/tests/contract/metal_dsv41_limits.zig b/zig/tests/contract/metal_dsv41_limits.zig index 1063c5ea64..5fdd7eb0f7 100644 --- a/zig/tests/contract/metal_dsv41_limits.zig +++ b/zig/tests/contract/metal_dsv41_limits.zig @@ -15,12 +15,21 @@ pub fn launchOf(a: std.mem.Allocator, device: mtl.Device, d: limits.MetalLaunch, } if (std.mem.startsWith(u8, d.name, "dsv41 fp8 mma ")) return .{ .decl = d, .function = fun, .here = sources.dsv41.fp8_mma, .pre_m5 = sources.dsv41.fp8_mma }; if (std.mem.startsWith(u8, d.name, "dsv41 attn ")) return .{ .decl = d, .function = fun, .here = sources.dsv41.csa2_attn, .pre_m5 = sources.dsv41.csa2_attn }; + if (std.mem.startsWith(u8, d.name, "dsv41 topk512 ")) return .{ .decl = d, .function = fun, .here = sources.dsv41.topk512, .pre_m5 = sources.dsv41.topk512 }; + if (std.mem.startsWith(u8, d.name, "dsv41 ix ")) return .{ .decl = d, .function = fun, .here = sources.dsv41.index, .pre_m5 = sources.dsv41.index }; if (std.mem.startsWith(u8, d.name, "dsv41 experts ")) { const body = try tf.dsv41_metal.experts.source(a, "", &.{}); return .{ .decl = d, .function = fun, .here = body, .pre_m5 = body }; } + if (std.mem.startsWith(u8, d.name, "dsv41 route ")) { + const e: usize = if (std.mem.indexOf(u8, d.name, "backbone") != null) 384 else 128; + const k: usize = if (e == 384) 6 else 3; + const body = try std.fmt.allocPrint(a, "#define TF_DSV41_E {d}\n#define TF_DSV41_TOPK {d}\n{s}", .{ e, k, sources.dsv41.route }); + return .{ .decl = d, .function = fun, .here = body, .pre_m5 = body }; + } if (std.mem.startsWith(u8, d.name, "dsv41 glue_tmp ")) return .{ .decl = d, .function = fun, .here = sources.dsv41.glue_tmp, .pre_m5 = sources.dsv41.glue_tmp }; if (std.mem.startsWith(u8, d.name, "dsv41 glue ")) return .{ .decl = d, .function = fun, .here = sources.dsv41.glue, .pre_m5 = sources.dsv41.glue }; + if (std.mem.startsWith(u8, d.name, "dsv41 cand ")) return .{ .decl = d, .function = fun, .here = sources.dsv41.cand, .pre_m5 = sources.dsv41.cand }; if (std.mem.startsWith(u8, d.name, "dsv41 mv ")) return .{ .decl = d, .function = fun, .here = sources.dsv41.mv, .pre_m5 = sources.dsv41.mv }; if (std.mem.startsWith(u8, d.name, "dsv41 fp8 rows ")) return .{ .decl = d, .function = fun, .here = sources.dsv41.fp8_rows, .pre_m5 = sources.dsv41.fp8_rows }; return null; diff --git a/zig/tests/contract/metal_dsv41_main.zig b/zig/tests/contract/metal_dsv41_main.zig index 1f181e106b..9137675d21 100644 --- a/zig/tests/contract/metal_dsv41_main.zig +++ b/zig/tests/contract/metal_dsv41_main.zig @@ -1,26 +1,77 @@ -//! DeepSeek-V4.1's Metal contract behind one hook: suites and their extra gates. +//! DeepSeek-V4.1's whole Metal contract behind one hook: routes, suites and their extra gates. const std = @import("std"); const mtl = @import("metal"); const contract = @import("contract"); const g = contract.gates; -const dsv41_attn = @import("metal_dsv41_attn.zig"); +const dsv41_route = @import("metal_dsv41_route.zig"); const dsv41_fp8_rows = @import("metal_dsv41_fp8_rows.zig"); const dsv41_fp8_mma = @import("metal_dsv41_fp8_mma.zig"); +const dsv41_topk512 = @import("metal_dsv41_topk512.zig"); +const dsv41_cand = @import("metal_dsv41_cand.zig"); +const dsv41_attn = @import("metal_dsv41_attn.zig"); const dsv41_experts = @import("metal_dsv41_experts.zig"); const dsv41_xmma = @import("metal_dsv41_xmma.zig"); +const dsv41_index = @import("metal_dsv41_index.zig"); const dsv41_glue = @import("metal_dsv41_glue.zig"); const dsv41_mv = @import("metal_dsv41_mv.zig"); const dsv41_glue_tmp = @import("metal_dsv41_glue_tmp.zig"); const dsv41_glue_tmp2 = @import("metal_dsv41_glue_tmp2.zig"); +fn wanted(only: ?[]const u8, name: []const u8) bool { + const text = only orelse return true; + return std.mem.indexOf(u8, name, text) != null; +} + +/// The route probes, registered where the dsv41 registrations sat: their kernels join `kernels` in variant order. +pub const Routes = struct { + routes: []dsv41_route.Route, + made: usize = 0, + a: std.mem.Allocator, + + /// Every route registration in name order: the two served kernels, then a control for each planted fault. + const list = [_]struct { v: dsv41_route.Variant, name: []const u8 }{ + .{ .v = .{ .kind = .backbone }, .name = "metal dsv41 route backbone tf_dsv41_route_select" }, + .{ .v = .{ .kind = .drafter }, .name = "metal dsv41 route drafter tf_dsv41_route_select" }, + .{ .v = .{ .kind = .backbone, .ctl = 1 }, .name = "metal dsv41 route backbone control: a higher id wins a tie" }, + .{ .v = .{ .kind = .drafter, .ctl = 1 }, .name = "metal dsv41 route drafter control: a higher id wins a tie" }, + .{ .v = .{ .kind = .backbone, .ctl = 2 }, .name = "metal dsv41 route backbone control: a NaN key wins" }, + .{ .v = .{ .kind = .drafter, .ctl = 2 }, .name = "metal dsv41 route drafter control: a NaN key wins" }, + .{ .v = .{ .kind = .backbone, .ctl = 3 }, .name = "metal dsv41 route backbone control: the bias in the weight" }, + .{ .v = .{ .kind = .drafter, .ctl = 3 }, .name = "metal dsv41 route drafter control: the bias in the weight" }, + .{ .v = .{ .kind = .backbone, .ctl = 4 }, .name = "metal dsv41 route backbone control: a write past k" }, + .{ .v = .{ .kind = .drafter, .ctl = 4 }, .name = "metal dsv41 route drafter control: a write past k" }, + }; + + pub fn create(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, kernels: *std.ArrayList(g.Kernel), only: ?[]const u8) !Routes { + var r = Routes{ .routes = try a.alloc(dsv41_route.Route, list.len), .a = a }; + errdefer a.free(r.routes); + errdefer for (r.routes[0..r.made]) |*p| p.deinit(); + for (list) |entry| { + if (!wanted(only, entry.name)) continue; + r.routes[r.made] = try dsv41_route.Route.init(a, device, queue, entry.v, entry.name); + r.made += 1; + try kernels.append(a, r.routes[r.made - 1].kernel()); + } + return r; + } + + pub fn destroy(r: *Routes) void { + for (r.routes[0..r.made]) |*p| p.deinit(); + r.a.free(r.routes); + } +}; + /// Every dsv41 suite, created and destroyed in the order main held them. pub const Suites = struct { fp8: *dsv41_fp8_rows.Suite, fp8_mma: *dsv41_fp8_mma.Suite, + topk: *dsv41_topk512.Suite, + cand: *dsv41_cand.Suite, attn: *dsv41_attn.Suite, experts: *dsv41_experts.Suite, xmma: *dsv41_xmma.Suite, + ix: *dsv41_index.Suite, glue: *dsv41_glue.Suite, mv: *dsv41_mv.Suite, glue_tmp: *dsv41_glue_tmp.Suite, @@ -31,12 +82,18 @@ pub const Suites = struct { errdefer fp8.destroy(); const fp8_mma = try dsv41_fp8_mma.Suite.create(a, device, queue, kernels, only); errdefer fp8_mma.destroy(); + const topk = try dsv41_topk512.Suite.create(a, device, queue, kernels, only); + errdefer topk.destroy(); + const cand = try dsv41_cand.Suite.create(a, device, queue, kernels, only); + errdefer cand.destroy(); const attn = try dsv41_attn.Suite.create(a, device, queue, kernels, only); errdefer attn.destroy(); const experts = try dsv41_experts.Suite.create(a, device, queue, kernels, only); errdefer experts.destroy(); const xmma = try dsv41_xmma.Suite.create(a, device, queue, kernels, only); errdefer xmma.destroy(); + const ix = try dsv41_index.Suite.create(a, device, queue, kernels, only); + errdefer ix.destroy(); const glue = try dsv41_glue.Suite.create(a, device, queue, kernels, only); errdefer glue.destroy(); const mv = try dsv41_mv.Suite.create(a, device, queue, kernels, only); @@ -45,7 +102,7 @@ pub const Suites = struct { errdefer glue_tmp.destroy(); const glue_tmp2 = try dsv41_glue_tmp2.Suite.create(a, device, queue, kernels, only); errdefer glue_tmp2.destroy(); - return .{ .fp8 = fp8, .fp8_mma = fp8_mma, .attn = attn, .experts = experts, .xmma = xmma, .glue = glue, .mv = mv, .glue_tmp = glue_tmp, .glue_tmp2 = glue_tmp2 }; + return .{ .fp8 = fp8, .fp8_mma = fp8_mma, .topk = topk, .cand = cand, .attn = attn, .experts = experts, .xmma = xmma, .ix = ix, .glue = glue, .mv = mv, .glue_tmp = glue_tmp, .glue_tmp2 = glue_tmp2 }; } pub fn destroy(s: *Suites) void { @@ -53,9 +110,12 @@ pub const Suites = struct { s.glue_tmp.destroy(); s.mv.destroy(); s.glue.destroy(); + s.ix.destroy(); s.xmma.destroy(); s.experts.destroy(); s.attn.destroy(); + s.cand.destroy(); + s.topk.destroy(); s.fp8_mma.destroy(); s.fp8.destroy(); } @@ -71,8 +131,14 @@ pub const Suites = struct { t.failed += try s.fp8_mma.invariance(seed); t.failed += try s.fp8_mma.accuracy(seed); if (std.c.getenv("TF_DSV41_TIME") != null) try s.fp8_mma.timing(io); + t.failed += try s.topk.invariance(seed); + if (std.c.getenv("TF_DSV41_TIME") != null) try s.topk.timing(io); + t.failed += try s.cand.invariance(seed); + t.failed += try s.cand.equality(seed); + if (std.c.getenv("TF_DSV41_TIME") != null) try s.cand.timing(io); try s.attn.invariance(io, seed, t); if (std.c.getenv("TF_DSV41_TIME") != null) try s.attn.timing(io); + t.failed += try s.ix.invariance(seed); t.failed += try s.glue.invariance(seed); if (s.mv.probes.items.len > 0) { t.failed += try s.mv.invariance(seed); @@ -83,5 +149,6 @@ pub const Suites = struct { t.failed += try s.glue_tmp.invariance(seed); t.failed += try s.glue_tmp.bounds(seed); t.failed += try s.glue_tmp2.invariance(seed); + if (std.c.getenv("TF_DSV41_TIME") != null) try s.ix.timing(io); } }; diff --git a/zig/tests/contract/metal_dsv41_route.zig b/zig/tests/contract/metal_dsv41_route.zig new file mode 100644 index 0000000000..9a9d0f804b --- /dev/null +++ b/zig/tests/contract/metal_dsv41_route.zig @@ -0,0 +1,399 @@ +//! DeepSeek-V4.1's MoE route under the contract: the top-k of sqrt(softplus(logit)) + bias, weights scaled by 1.5. +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const c = contract.spec; +const g = contract.gates; +const rig_ = @import("rig.zig"); +const sources = @import("kernel_sources"); + +/// A registration: the variant, and the planted fault its kernel text carries (0: the kernel as served). +pub const Variant = struct { kind: Kind = .backbone, ctl: u8 = 0 }; + +pub const Kind = enum { backbone, drafter }; + +fn expertsOf(v: Kind) usize { + return if (v == .backbone) 384 else 128; +} + +fn picksOf(v: Kind) usize { + return if (v == .backbone) 6 else 3; +} + +const scale: f32 = 1.5; // the reference's routed_scaling_factor +/// The planted tie's logit: softplus's linear tail, so its score is an exact square root of an exact square. +const tie_logit: f32 = 40.5; +const eps: f32 = 1e-20; + +pub fn experts(v: Kind) usize { + return expertsOf(v); +} + +pub fn picks(v: Kind) usize { + return picksOf(v); +} + +fn depth(sh: c.Shape) usize { + return sh.k + 8; +} + +pub fn spec(v: Variant, name: []const u8) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = depth, + .ops = 4, + .out = .f32, + .nan = .{ .stated = "a NaN logit or bias is never picked; a +inf logit is picked at a NaN weight and its row's other weights 0" }, + .tiles = .{ .{ .size = 8, .step = 1, .min = 1 }, .{}, .{} }, + .served = if (v.kind == .backbone) &.{ + .{ .rows = 4, .n = 384, .k = 6 }, + .{ .rows = 1, .n = 384, .k = 6 }, + .{ .rows = 8, .n = 384, .k = 6 }, + .{ .rows = 8192, .n = 384, .k = 6 }, + } else &.{ + .{ .rows = 4, .n = 128, .k = 3 }, + .{ .rows = 1, .n = 128, .k = 3 }, + .{ .rows = 8, .n = 128, .k = 3 }, + .{ .rows = 8192, .n = 128, .k = 3 }, + }, + .sums = .{ .exact = true, .why = "the stress picks' scores are roots of exact squares, so their sum is exact in any fp32 order" }, + }; +} + +fn softplus64(z: f64) f64 { + if (std.math.isNan(z)) return z; + const e = @exp(@min(z, 20.0)); + if (z > 20.0) return z; + if (z < -10.0) return e - 0.5 * e * e; + return @log(1.0 + e); +} + +const Args = extern struct { rows: i32, topk: i32, scale: f32, eps: f32 }; + +/// The stress picks' scores: roots of exact squares (the logit is the square, above softplus's linear tail). +const stress_scores = [_]f64{ 32768, 9, 8, 7, 6, 5 }; + +pub const Route = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + v: Variant, + name: []const u8, + e: usize, + k: usize, + shape: c.Shape = undefined, + logits: []f32 = &.{}, + bias: []f32 = &.{}, + bias_vl: []f32 = &.{}, + image: []u8 = &.{}, + previous: []u8 = &.{}, + runs: usize = 0, + x: rig_.Ref = undefined, + bias_r: rig_.Ref = undefined, + bias_vl_r: rig_.Ref = undefined, + image_r: rig_.Ref = undefined, + picks: rig_.Ref = undefined, + weights: rig_.Ref = undefined, + msg: [240]u8 = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, v: Variant, name: []const u8) !Route { + const e = experts(v.kind); + const k = picks(v.kind); + const src = try std.fmt.allocPrint(a, "#define TF_DSV41_E {d}\n#define TF_DSV41_TOPK {d}\n#define TF_CTL {d}\n{s}", .{ e, k, v.ctl, sources.dsv41.route }); + defer a.free(src); + return .{ + .a = a, + .rig = .{ .a = a, .device = device, .queue = queue }, + .pipe = try rig_.pipeline(device, src, "tf_dsv41_route_select"), + .v = v, + .name = name, + .e = e, + .k = k, + }; + } + + pub fn deinit(t: *Route) void { + t.a.free(t.logits); + t.a.free(t.bias); + t.a.free(t.bias_vl); + t.a.free(t.image); + t.a.free(t.previous); + t.rig.deinit(); + t.pipe.deinit(); + } + + pub fn kernel(t: *Route) g.Kernel { + const control: ?g.Gate = switch (t.v.ctl) { + 1, 3 => .shapes, + 2 => .nonfinite, + 4 => .repeat, + else => null, + }; + return .{ .spec = spec(t.v, t.name), .ctx = t, .vt = &vt, .control = control }; + } + + const vt = g.hooks(Route); + + /// The two experts whose (logit, bias) the random fill makes exactly equal on every fourth row: a planted tie. + fn tiePair(t: *const Route) [2]usize { + return .{ 11, if (t.e == 384) 322 else 100 }; + } + + pub fn prepare(t: *Route, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.n != t.e or shape.k != t.k) return false; + t.rig.reset(); + t.shape = shape; + t.runs = 0; + t.a.free(t.logits); + t.logits = &.{}; + t.a.free(t.bias); + t.bias = &.{}; + t.a.free(t.bias_vl); + t.bias_vl = &.{}; + t.a.free(t.image); + t.image = &.{}; + t.a.free(t.previous); + t.previous = &.{}; + const rows = shape.rows; + t.logits = try t.a.alloc(f32, rows * t.e); + t.bias = try t.a.alloc(f32, t.e); + t.bias_vl = try t.a.alloc(f32, t.e); + t.image = try t.a.alloc(u8, rows); + t.previous = try t.a.alloc(u8, t.count() * 4); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (t.bias) |*v| v.* = rng.float(f32) - 0.5; + for (t.bias_vl) |*v| v.* = rng.float(f32) - 0.5; + const tie = t.tiePair(); + t.bias[tie[0]] = 0.25; + t.bias[tie[1]] = 0.25; + t.bias_vl[tie[0]] = 0.25; + t.bias_vl[tie[1]] = 0.25; + for (t.logits) |*v| v.* = rng.float(f32) * 12 - 6; + for (0..rows) |r| { + t.image[r] = @intCast(r % 2); + if (r % 4 == 1) { + t.logits[r * t.e + tie[0]] = tie_logit; // above every random logit, so the two are its first two picks + t.logits[r * t.e + tie[1]] = tie_logit; + } + } + switch (fill) { + .stress => { + // The k picks' logits are exact squares; everything else sits far below, and the bias is zero. + @memset(t.logits, -80.0); + @memset(t.bias, 0); + @memset(t.bias_vl, 0); + for (0..rows) |r| for (0..t.k) |j| { + const s = stress_scores[j]; + t.logits[r * t.e + j] = @floatCast(s * s); + }; + }, + .nonfinite => { + const values = [_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32) }; + const plus = if (t.e == 384) @as(usize, 100) else 60; + const minus = if (t.e == 384) @as(usize, 200) else 120; + t.logits[0 * t.e + 7] = values[0]; + t.logits[1 * t.e + plus] = values[1]; + t.logits[2 * t.e + minus] = values[2]; + t.logits[3 * t.e + 3] = values[1]; + t.logits[3 * t.e + 4] = values[2]; + t.bias[13] = values[0]; + t.bias_vl[17] = values[0]; + }, + .random => {}, + } + t.x = try t.rig.input(std.mem.sliceAsBytes(t.logits)); + t.bias_r = try t.rig.input(std.mem.sliceAsBytes(t.bias)); + t.bias_vl_r = try t.rig.input(std.mem.sliceAsBytes(t.bias_vl)); + t.image_r = try t.rig.input(t.image); + t.picks = try t.rig.scratch(t.count() * 4); + t.weights = try t.rig.output(t.count() * 4); + return true; + } + + pub fn run(t: *Route) !void { + const p = t.rig.begin(); + p.e.setPipeline(t.pipe); + p.e.setBuffer(t.x.buf, t.x.off, 0); + p.e.setBuffer(t.bias_r.buf, t.bias_r.off, 1); + p.e.setBuffer(t.bias_vl_r.buf, t.bias_vl_r.off, 2); + p.e.setBuffer(t.image_r.buf, t.image_r.off, 3); + p.e.setValue(Args{ + .rows = @intCast(t.shape.rows), + .topk = @intCast(t.k), + .scale = scale, + .eps = eps, + }, 4); + p.e.setBuffer(t.picks.buf, t.picks.off, 5); + p.e.setBuffer(t.weights.buf, t.weights.off, 6); + p.e.dispatchGroups(mtl.Size.of(t.shape.rows, 1, 1), mtl.Size.of(32, 1, 1)); + try t.rig.finish(p); + const picked = t.rig.bytes(t.picks, t.count() * 4); + if (t.runs > 0 and !std.mem.eql(u8, t.previous, picked)) return error.RoutePicksNotRepeatable; + @memcpy(t.previous, picked); + t.runs += 1; + } + + pub fn count(t: *Route) usize { + return t.shape.rows * t.k; + } + + fn rowScores(t: *Route, r: usize, s: *[384]f64, key: *[384]f64) void { + const bias = if (t.image[r] != 0) t.bias_vl else t.bias; + for (0..t.e) |i| { + s[i] = @sqrt(softplus64(t.logits[r * t.e + i])); + key[i] = s[i] + bias[i]; + } + } + + /// The k picks in fp64 (NaN keys never win, ties to the lower id), and their unbiased scores' sum in pick order. + fn select(t: *Route, key: *const [384]f64, s: *const [384]f64, chosen: *[6]usize, total: *f64) void { + var used: [384]bool = @splat(false); + total.* = 0; + for (0..t.k) |kk| { + var best: usize = t.e; // none + for (0..t.e) |i| { + if (used[i] or key[i] != key[i]) continue; + if (best == t.e or key[i] > key[best] or (key[i] == key[best] and i < best)) best = i; + } + if (best == t.e) { + chosen[kk] = t.e; + } else { + used[best] = true; + chosen[kk] = best; + total.* += s[best]; + } + } + } + + pub fn reference(t: *Route, want: []f64, mag: []f64) !void { + var s: [384]f64 = undefined; + var key: [384]f64 = undefined; + var chosen: [6]usize = undefined; + for (0..t.shape.rows) |r| { + t.rowScores(r, &s, &key); + var total: f64 = 0; + t.select(&key, &s, &chosen, &total); + for (0..t.k) |kk| { + const i = chosen[kk]; + const w = if (i == t.e) 0 else (s[i] / (total + 1e-20)) * scale; + want[r * t.k + kk] = w; + mag[r * t.k + kk] = @abs(w); + } + } + } + + /// Gate 4: the picks as fp64 gives them, the sum in the accumulator and order, then the kernel's fp32 weight chain. + pub fn simulate(t: *Route, acc: c.Accum, order: g.Order, out: []f64) !void { + var s: [384]f64 = undefined; + var key: [384]f64 = undefined; + var chosen: [6]usize = undefined; + var terms: [6]f64 = undefined; + for (0..t.shape.rows) |r| { + t.rowScores(r, &s, &key); + var total: f64 = 0; + t.select(&key, &s, &chosen, &total); + var nt: usize = 0; + for (chosen[0..t.k]) |i| if (i != t.e) { + terms[nt] = s[i]; + nt += 1; + }; + const sum = contract.rows.sumIn(acc, order, terms[0..nt]); + for (0..t.k) |kk| { + const i = chosen[kk]; + out[r * t.k + kk] = if (i == t.e) 0 else blk: { + const num: f32 = @floatCast(s[i]); + const den: f32 = @floatCast(sum + 1e-20); + break :blk @as(f64, num / den * scale); + }; + } + } + } + + pub fn check(t: *Route, at: g.Gate) !?[]const u8 { + const n = t.count(); + const picked = t.rig.bytes(t.picks, n * 4); + const weights = t.rig.bytes(t.weights, n * 4); + if (at == .nonfinite) { + var s: [384]f64 = undefined; + var key: [384]f64 = undefined; + for (0..t.shape.rows) |r| { + t.rowScores(r, &s, &key); + for (0..t.k) |kk| { + const pick = std.mem.readInt(i32, picked[(r * t.k + kk) * 4 ..][0..4], .little); + if (pick < 0 or pick >= t.e) continue; + if (key[@intCast(pick)] != key[@intCast(pick)]) + return try std.fmt.bufPrint(&t.msg, "row {d} pick {d} took expert {d} at a NaN key: a NaN is never picked", .{ r, kk, pick }); + } + } + const buf = try t.a.alloc(f64, 2 * n); + defer t.a.free(buf); + const want = buf[0..n]; + const mag = buf[n..]; + try t.reference(want, mag); + var nf: usize = 0; + var differ: usize = 0; + var first: usize = n; + for (want, 0..) |w, i| { + if (c.classOf(w) != .finite) nf += 1; + if (c.classOf(c.read(.f32, weights, i)) != c.classOf(w)) { + differ += 1; + if (first == n) first = i; + } + } + if (nf == 0 or nf == n) return try std.fmt.bufPrint(&t.msg, "the plants make {d} of {d} reference weights non-finite: no proof", .{ nf, n }); + if (differ > 0) return try std.fmt.bufPrint(&t.msg, "weight {d} is {t} where float64 gives {t}", .{ first, c.classOf(c.read(.f32, weights, first)), c.classOf(want[first]) }); + return null; + } + var s: [384]f64 = undefined; + var key: [384]f64 = undefined; + var chosen: [6]usize = undefined; + for (0..t.shape.rows) |r| { + t.rowScores(r, &s, &key); + var used: [384]bool = @splat(false); + for (0..t.k) |kk| { + const i = r * t.k + kk; + const pick = std.mem.readInt(i32, picked[i * 4 ..][0..4], .little); + if (pick < 0) return try std.fmt.bufPrint(&t.msg, "row {d} pick {d} took nothing, but every expert is a candidate", .{ r, kk }); + const p: usize = @intCast(pick); + if (p >= t.e) return try std.fmt.bufPrint(&t.msg, "row {d} pick {d} is expert {d} of {d}", .{ r, kk, p, t.e }); + if (used[p]) return try std.fmt.bufPrint(&t.msg, "row {d} pick {d} repeats expert {d}", .{ r, kk, p }); + used[p] = true; + // The best remaining key in fp64: the pick must reach it within the fp32 arithmetic's slack. + var bestk: f64 = -std.math.inf(f64); + var bestid: usize = t.e; + for (0..t.e) |j| { + if (used[j] or key[j] != key[j]) continue; + if (key[j] > bestk or (key[j] == bestk and j < bestid)) { + bestk = key[j]; + bestid = j; + } + } + const slack = 1e-5 + 1e-6 * @abs(bestk); + if (bestid != t.e and key[p] < bestk - slack) return try std.fmt.bufPrint(&t.msg, "row {d} pick {d} took expert {d} at key {d:.6} while expert {d} remains at {d:.6}", .{ r, kk, p, key[p], bestid, bestk }); + for (0..t.e) |j| { + if (used[j] or key[j] != key[p] or j >= p) continue; + return try std.fmt.bufPrint(&t.msg, "row {d} pick {d} took expert {d} while expert {d} ties it at key {d:.6}", .{ r, kk, p, j, key[p] }); + } + } + // The weights against the fp64 formula over the kernel's own picks. + var total: f64 = 0; + for (0..t.k) |kk| { + const pick = std.mem.readInt(i32, picked[(r * t.k + kk) * 4 ..][0..4], .little); + chosen[kk] = if (pick < 0) t.e else @intCast(pick); + if (chosen[kk] != t.e) total += s[chosen[kk]]; + } + for (0..t.k) |kk| { + const i = r * t.k + kk; + const bias = if (chosen[kk] == t.e) 0 else if (t.image[r] != 0) t.bias_vl[chosen[kk]] else t.bias[chosen[kk]]; + const want = if (chosen[kk] == t.e) 0 else (s[chosen[kk]] / (total + 1e-20)) * scale; + const got = c.read(.f32, weights, i); + if (c.classOf(got) != c.classOf(want)) return try std.fmt.bufPrint(&t.msg, "row {d} pick {d}: weight class {t} where the fp64 formula gives {t}", .{ r, kk, c.classOf(got), c.classOf(want) }); + if (at != .nonfinite and @abs(got - want) > c.bound(spec(t.v, t.name), t.shape, want, @abs(want))) + return try std.fmt.bufPrint(&t.msg, "row {d} pick {d}: weight {e:.6} where the fp64 formula from its own picks gives {e:.6} (its bias {e:.6} is in the key only)", .{ r, kk, got, want, bias }); + } + } + return null; + } +}; diff --git a/zig/tests/contract/metal_dsv41_topk512.zig b/zig/tests/contract/metal_dsv41_topk512.zig new file mode 100644 index 0000000000..2436c83459 --- /dev/null +++ b/zig/tests/contract/metal_dsv41_topk512.zig @@ -0,0 +1,298 @@ +//! DeepSeek-V4.1's indexer top-512 over fp32 score rows under the kernel contract on Metal. +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const sources = @import("kernel_sources"); +const rig_ = @import("rig.zig"); + +const c = contract.spec; +const g = contract.gates; +const base = "dsv41 topk512 "; +const rule = "NaN ranks as -inf, -0 as +0, ties go to the lower index; row r keeps min(512, visible[r]) indices, ascending, then -1"; + +fn noSums(_: c.Shape) usize { + return 0; +} + +const served = blk: { + const rows = [_]usize{ 1, 4, 64, 4096, 8192 }; + const widths = [_]usize{ 128, 511, 512, 513, 1024, 4096, 16384 }; + var out: [rows.len * widths.len + 4]c.Shape = undefined; + var n: usize = 0; + out[n] = .{ .rows = 8, .n = 512, .k = 1024 }; + n += 1; + for (rows) |r| for (widths) |w| { + if (r >= 4096 and w > 1024 and w < 16384) continue; + if (r == 8192 and w > 4096) continue; + out[n] = .{ .rows = r, .n = 512, .k = w }; + n += 1; + }; + out[n] = .{ .rows = 1, .n = 512, .k = 131072 }; + out[n + 1] = .{ .rows = 1, .n = 512, .k = 1048576 }; + out[n + 2] = .{ .rows = 4, .n = 512, .k = 1048576 }; + break :blk out[0 .. n + 3].*; +}; + +pub fn spec(name: []const u8) c.Spec { + return .{ + .name = name, + .accum = .none, + .depth = noSums, + .out = .raw, + .nan = .{ .stated = rule }, + .tiles = .{ .{ .size = 4096 }, .{}, .{ .size = 16384, .step = 128, .min = 128 } }, + .served = &served, + }; +} + +fn rank(v: f32) f32 { + return if (std.math.isNan(v)) -std.math.inf(f32) else v; +} + +const Order = struct { + row: []const f32, + fn before(o: Order, a: u32, b: u32) bool { + const x = rank(o.row[a]); + const y = rank(o.row[b]); + return x > y or (x == y and a < b); + } +}; + +pub const TopK = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + ctl: u8, + full: bool = false, + shape: c.Shape = undefined, + scores: []f32 = &.{}, + vis: []i32 = &.{}, + plants: usize = 0, + s: rig_.Ref = undefined, + v: rig_.Ref = undefined, + out: rig_.Ref = undefined, + msg: [200]u8 = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, ctl: u8) !TopK { + const text = try std.fmt.allocPrint(a, "#define TF_CTL {d}\n{s}", .{ ctl, sources.dsv41.topk512 }); + defer a.free(text); + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .pipe = try rig_.pipeline(device, text, "tf_topk512"), .ctl = ctl }; + } + + pub fn deinit(t: *TopK) void { + t.a.free(t.scores); + t.a.free(t.vis); + t.pipe.deinit(); + t.rig.deinit(); + } + + pub fn kernel(t: *TopK, name: []const u8) g.Kernel { + const control: ?g.Gate = switch (t.ctl) { + 1 => .nonfinite, + 2, 3 => .shapes, + 4 => .repeat, + else => null, + }; + return .{ .spec = spec(name), .ctx = t, .vt = &vt, .control = control }; + } + + const vt = g.hooks(TopK); + + pub fn prepare(t: *TopK, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.n != 512 or shape.k < 1) return false; + t.rig.reset(); + t.shape = shape; + const k = shape.k; + t.a.free(t.scores); + t.a.free(t.vis); + t.scores = &.{}; + t.vis = &.{}; + t.scores = try t.a.alloc(f32, shape.rows * k); + t.vis = try t.a.alloc(i32, shape.rows); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (0..shape.rows) |r| { + const row = t.scores[r * k ..][0..k]; + const vis: usize = if (t.full or r < 4) k else switch (r % 16) { + 1 => k / 2 + 1, + 6 => @min(k, 512), + 7 => @min(k, 513), + 9 => 0, + 13 => @min(k, 1), + else => k, + }; + t.vis[r] = @intCast(vis); + const style: usize = if (t.full) 0 else r % 5; + for (row, 0..) |*x, i| x.* = if (i >= vis) std.math.inf(f32) else switch (style) { + 0 => @as(f32, @floatFromInt(rng.uintLessThan(u32, 512))) / 16 - 16, + 1 => 1.5, + 2 => (rng.float(f32) - 0.5) * 1e6, + 3 => -@as(f32, @floatFromInt(i / 3)), + else => switch (rng.uintLessThan(u32, 8)) { + 0 => 1.0, + 1...4 => 0.0, + else => -0.0, + }, + }; + } + t.plants = 0; + if (fill == .nonfinite) { + const at = [_][2]usize{ .{ 0, k / 3 }, .{ 1, k / 2 }, .{ 2, k - 1 }, .{ 3, 0 }, .{ 3, k - 1 } }; + const put = [_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32), std.math.inf(f32), -std.math.inf(f32) }; + for (at, put) |p, v| if (p[0] < shape.rows) { + t.scores[p[0] * k + p[1]] = v; + t.plants += 1; + }; + } + t.s = try t.rig.input(std.mem.sliceAsBytes(t.scores)); + t.v = try t.rig.input(std.mem.sliceAsBytes(t.vis)); + t.out = try t.rig.output(shape.rows * 512 * @sizeOf(i32)); + return true; + } + + /// Rows r0..r0+rows of the prepared call, as one dispatch. + pub fn encode(t: *TopK, e: anytype, r0: usize, rows: usize) void { + const args = [1]u32{@intCast(t.shape.k)}; + e.setPipeline(t.pipe); + e.setBuffer(t.s.buf, t.s.off + r0 * t.shape.k * 4, 1); + e.setBuffer(t.v.buf, t.v.off + r0 * 4, 2); + e.setBuffer(t.out.buf, t.out.off + r0 * 512 * 4, 3); + e.setValue(args, 0); + e.dispatchGroups(mtl.Size.of(rows, 1, 1), mtl.Size.of(1024, 1, 1)); + } + + pub fn run(t: *TopK) !void { + @memset(t.rig.bytes(t.out, t.count()), 0x55); + const pass = t.rig.begin(); + t.encode(pass.e, 0, t.shape.rows); + try t.rig.finish(pass); + } + + pub fn count(t: *TopK) usize { + return t.shape.rows * 512 * @sizeOf(i32); + } + + pub fn reference(_: *TopK, want: []f64, mag: []f64) !void { + @memset(want, 0); + @memset(mag, 0); + } + + /// The CPU stable sort against each row: all rows up to 4M scores, else a stride of rows and the last. + pub fn check(t: *TopK, at: g.Gate) !?[]const u8 { + const k = t.shape.k; + const got = std.mem.bytesAsSlice(i32, t.rig.bytes(t.out, t.count())); + const stride = @max(1, t.shape.rows * k / (4 << 20)); + const idx = try t.a.alloc(u32, k); + defer t.a.free(idx); + var nonfinite: usize = 0; + for (0..t.shape.rows) |r| { + if (r >= 8 and r % stride != 0 and r + 1 != t.shape.rows) continue; + const vis: usize = @intCast(t.vis[r]); + const row = t.scores[r * k ..][0..vis]; + for (row) |x| nonfinite += @intFromBool(!std.math.isFinite(x)); + for (idx[0..vis], 0..) |*x, i| x.* = @intCast(i); + std.sort.pdq(u32, idx[0..vis], Order{ .row = row }, Order.before); + const keep = @min(512, vis); + std.mem.sort(u32, idx[0..keep], {}, std.sort.asc(u32)); + const out = got[r * 512 ..][0..512]; + for (out, 0..) |have, i| { + const want: i32 = if (i < keep) @intCast(idx[i]) else -1; + if (have != want) return try std.fmt.bufPrint(&t.msg, "row {d} (visible {d} of {d}) slot {d}: got {d}, stable rank gives {d}", .{ r, vis, k, i, have, want }); + } + } + if (at == .nonfinite and (t.plants < 5 or nonfinite < 5)) return try std.fmt.bufPrint(&t.msg, "{d} plants: no proof", .{t.plants}); + return null; + } +}; + +/// Every registered variant, the launch-split probe after the gates, and the timing. +pub const Suite = struct { + a: std.mem.Allocator, + device: mtl.Device, + queue: mtl.Queue, + probes: std.ArrayList(*TopK) = .empty, + + const list = [_]struct { ctl: u8, name: []const u8 }{ + .{ .ctl = 0, .name = base ++ "tf_topk512" }, + .{ .ctl = 1, .name = base ++ "control: NaN ranked above +inf" }, + .{ .ctl = 2, .name = base ++ "control: ties go to the higher index" }, + .{ .ctl = 3, .name = base ++ "control: the last score dropped" }, + .{ .ctl = 4, .name = base ++ "control: writes into its input" }, + }; + + pub fn create(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, kernels: *std.ArrayList(g.Kernel), only: ?[]const u8) !*Suite { + const s = try a.create(Suite); + s.* = .{ .a = a, .device = device, .queue = queue }; + for (list) |entry| { + if (only) |text| if (std.mem.indexOf(u8, entry.name, text) == null) continue; + const p = try a.create(TopK); + p.* = try TopK.init(a, device, queue, entry.ctl); + try s.probes.append(a, p); + try kernels.append(a, p.kernel(entry.name)); + } + return s; + } + + pub fn destroy(s: *Suite) void { + for (s.probes.items) |p| { + p.deinit(); + s.a.destroy(p); + } + s.probes.deinit(s.a); + s.a.destroy(s); + } + + /// Rows split into launches of 1, 7 and rows/3, and one launch a row, equal the single launch byte for byte. + pub fn invariance(s: *Suite, seed: u64) !usize { + const p = if (s.probes.items.len > 0) s.probes.items[0] else return 0; + if (p.ctl != 0) return 0; + var bad: usize = 0; + var launches: usize = 0; + for ([_]c.Shape{ .{ .rows = 64, .n = 512, .k = 16384 }, .{ .rows = 33, .n = 512, .k = 700 }, .{ .rows = 5, .n = 512, .k = 131072 } }) |shape| { + if (!try p.prepare(shape, .random, seed)) return error.Refused; + try p.run(); + const whole = try s.a.dupe(u8, p.rig.bytes(p.out, p.count())); + defer s.a.free(whole); + for ([_]usize{ 1, 7, shape.rows / 3 }) |step| { + @memset(p.rig.bytes(p.out, p.count()), 0x55); + const pass = p.rig.begin(); + var r0: usize = 0; + while (r0 < shape.rows) : (r0 += step) { + p.encode(pass.e, r0, @min(step, shape.rows - r0)); + launches += 1; + } + try p.rig.finish(pass); + if (!std.mem.eql(u8, whole, p.rig.bytes(p.out, p.count()))) { + bad += 1; + std.debug.print("FAIL launch split {d}x{d}: {d} rows a launch differs from one launch\n", .{ shape.rows, shape.k, step }); + } + } + } + const word = if (bad == 0) "PASS" else "FAIL"; + std.debug.print("{s} launch split {s}: 3 shapes at 1, 7 and rows/3 rows a launch ({d} launches) byte-equal to one launch, {d} differ\n", .{ word, base, launches, bad }); + return bad; + } + + /// Median microseconds a launch, over launches queued back to back, at full visible length. + pub fn timing(s: *Suite, io: std.Io) !void { + var p = try TopK.init(s.a, s.device, s.queue, 0); + defer p.deinit(); + p.full = true; + for ([_]usize{ 16384, 131072 }) |width| for ([_]usize{ 1, 4096 }) |rows| { + if (!try p.prepare(.{ .rows = rows, .n = 512, .k = width }, .random, 7)) return error.Refused; + const reps: usize = if (rows == 1) 64 else 4; + var samples: [9]f64 = undefined; + for (&samples) |*out| { + const t0 = std.Io.Clock.awake.now(io); + const pass = p.rig.begin(); + for (0..reps) |_| p.encode(pass.e, 0, rows); + try p.rig.finish(pass); + const ns: f64 = @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(io)).toNanoseconds()); + out.* = ns / 1000.0 / @as(f64, @floatFromInt(reps)); + } + std.mem.sort(f64, &samples, {}, std.sort.asc(f64)); + const med = samples[samples.len / 2]; + std.debug.print("TIME {s}: width {d} rows {d}: {d:.1} us a launch, {d:.3} us a row (median of {d} x {d} launches)\n", .{ base, width, rows, med, med / @as(f64, @floatFromInt(rows)), samples.len, reps }); + }; + } +}; diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index 1aa25764a8..e74219237d 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -307,6 +307,8 @@ pub fn main(init: std.process.Init) !u8 { d3router_count += 1; try kernels.append(a, d3routers[d3router_count - 1].kernel()); } + var dsv41_routes = try dsv41.Routes.create(a, device, queue, &kernels, o.only); + defer dsv41_routes.destroy(); var d3combines: [std.enums.values(drama_combine.Variant).len]drama_combine.Combine = undefined; var d3combine_count: usize = 0; defer for (d3combines[0..d3combine_count]) |*p| p.deinit(); From d9cf9de64282f346603536e40cffa9c5fd4f6372 Mon Sep 17 00:00:00 2001 From: Kevin Gilbertson <880013+gilby@users.noreply.github.com> Date: Sun, 1 Nov 2026 04:48:09 +0100 Subject: [PATCH 2/2] dsv41 metal: the KV compressor and Engram compress.metal pools the ratio-2 and ratio-4 compressed KV and round-trips it through FP4; engram.metal dequantises the FP8 Engram rows, projects them, gates them and adds them to the residual. Registered with the contract against float64 and bit-equal to the phase-1 glue where they replace it. --- zig/kernels/metal/deepseek_v41/compress.metal | 138 +++++ zig/kernels/metal/deepseek_v41/engram.metal | 129 ++++ zig/kernels/metal/deepseek_v41/sources.zig | 3 + .../core/contract/dsv41_compress_records.zig | 12 + .../core/contract/dsv41_engram_records.zig | 13 + zig/src/core/contract/dsv41_metal_records.zig | 3 +- .../families/deepseek_v41/metal/compress.zig | 68 +++ .../deepseek_v41/metal/contract_exports.zig | 2 + .../families/deepseek_v41/metal/engram.zig | 165 +++++ zig/tests/contract/dsv41_engram_ref.zig | 407 +++++++++++++ zig/tests/contract/metal_dsv41_compress.zig | 574 ++++++++++++++++++ zig/tests/contract/metal_dsv41_engram.zig | 435 +++++++++++++ zig/tests/contract/metal_dsv41_limits.zig | 2 + zig/tests/contract/metal_dsv41_main.zig | 15 +- 14 files changed, 1964 insertions(+), 2 deletions(-) create mode 100644 zig/kernels/metal/deepseek_v41/compress.metal create mode 100644 zig/kernels/metal/deepseek_v41/engram.metal create mode 100644 zig/src/core/contract/dsv41_compress_records.zig create mode 100644 zig/src/core/contract/dsv41_engram_records.zig create mode 100644 zig/src/families/deepseek_v41/metal/compress.zig create mode 100644 zig/src/families/deepseek_v41/metal/engram.zig create mode 100644 zig/tests/contract/dsv41_engram_ref.zig create mode 100644 zig/tests/contract/metal_dsv41_compress.zig create mode 100644 zig/tests/contract/metal_dsv41_engram.zig diff --git a/zig/kernels/metal/deepseek_v41/compress.metal b/zig/kernels/metal/deepseek_v41/compress.metal new file mode 100644 index 0000000000..bc38d948b3 --- /dev/null +++ b/zig/kernels/metal/deepseek_v41/compress.metal @@ -0,0 +1,138 @@ +// DeepSeek V4.1's compressor on Metal: the ratio-2 pool, the compressed KV's FP4 round trip and the index-key packing. +#include +using namespace metal; +#pragma METAL fp contract(off) + +#ifndef TF_CTL +#define TF_CTL 0 +#endif + +// Control 1 of every kernel: a non-finite input is read as 0, so the output class changes. +inline float tf_c_in(float v) { return (TF_CTL == 1 && !isfinite(v)) ? 0.0f : v; } + +inline float tf_c_pow2(int e) { return as_type(uint(e + 127) << 23); } + +// ceil(log2(v)) of a positive normal float from its bits, as the reference's fast_log2_ceil. +inline int tf_c_log2c(float v) { + const uint b = as_type(v); + return int((b >> 23) & 0xFFu) - 127 + (((b & 0x7FFFFFu) != 0u && TF_CTL != 2) ? 1 : 0); +} + +// A magnitude of at most 448 as its nearest E4M3 code, ties to even (codes 0..8: subnormals and the first normal). +inline uint tf_c_e4m3(float a) { + uint b = as_type(a); + b += 0x7FFFFu + ((b >> 20) & 1u); + return a < 0.015625f ? uint(rint(a * 512.0f)) : (((b >> 23) - 120u) << 3) | ((b >> 20) & 7u); +} + +inline float tf_c_e4m3v(uint c) { return float(as_type(ushort((c & 0x7Fu) << 7))) * 256.0f; } + +// A magnitude of at most 6 as its nearest E2M1 code, ties to the even code (control 3: ties away from zero). +inline uint tf_c_e2m1(float m) { + const bool t = TF_CTL == 3; + if (m < 0.25f || (!t && m == 0.25f)) return 0u; + if (m < 0.75f) return 1u; + if (m < 1.25f || (!t && m == 1.25f)) return 2u; + if (m < 1.75f) return 3u; + if (m < 2.5f || (!t && m == 2.5f)) return 4u; + if (m < 3.5f) return 5u; + return (m < 5.0f || (!t && m == 5.0f)) ? 6u : 7u; +} + +inline float tf_c_e2m1v(uint c) { + const uint e = (c >> 1) & 3u, m = c & 1u; + return e == 0u ? float(m) * 0.5f : ldexp(float(m) + 2.0f, int(e) - 2); +} + +struct TfCPoolArgs { + uint rows, dim; +}; + +// Compressor ratio 2 (model.py:429-487): a column a thread, rows in order; odd positions pool the two slots. +[[kernel, max_total_threads_per_threadgroup(32)]] void tf_dsv41_c_pool2( + const device float* KV [[buffer(0)]], const device float* SC [[buffer(1)]], device float* ST [[buffer(2)]], + const device int* P [[buffer(3)]], device bfloat* Y [[buffer(4)]], constant TfCPoolArgs& a [[buffer(5)]], + uint d [[thread_position_in_grid]]) { + if (d >= a.dim) return; + // the carry ST (kv[2][dim], score[2][dim]) keeps each slot's last value: any launch split gives the same bits + float kv[2] = {ST[d], ST[a.dim + d]}, sc[2] = {ST[2 * a.dim + d], ST[3 * a.dim + d]}; + for (uint r = 0; r < a.rows; ++r) { + const int pos = P[r]; + const int slot = TF_CTL == 3 ? int(r & 1u) : (pos & 1); + kv[slot] = tf_c_in(KV[ulong(r) * a.dim + d]); + sc[slot] = tf_c_in(SC[ulong(r) * a.dim + d]); + if (slot == 1) { + // pooled = kv0 * (e0 / s) + kv1 * (e1 / s) in fp32 with e = exp(score - max), rounded to bf16 once + const float m = max(sc[0], sc[1]); + const float e0 = precise::exp(sc[0] - m), e1 = precise::exp(sc[1] - m); + const float sum = e0 + e1; + const float pooled = kv[0] * precise::divide(e0, sum) + kv[1] * precise::divide(e1, sum); + Y[ulong(r) * a.dim + d] = bfloat(pooled); + } + } + if (TF_CTL == 4) Y[ulong(a.rows) * a.dim + d] = bfloat(1.0f); + if (TF_CTL == 2) return; + ST[d] = kv[0], ST[a.dim + d] = kv[1], ST[2 * a.dim + d] = sc[0], ST[3 * a.dim + d] = sc[1]; +} + +struct TfCRtArgs { + uint blocks, in_stride, out_row; +}; + +// The compressed KV's fp4_act_quant round trip (kernel.py:184), 16-value blocks, E4M3 scales, to row D[r] (< 0 skips). +[[kernel, max_total_threads_per_threadgroup(32)]] void tf_dsv41_c_fp4rt( + const device bfloat* X [[buffer(0)]], device bfloat* Y [[buffer(1)]], const device int* D [[buffer(2)]], + constant TfCRtArgs& a [[buffer(3)]], uint2 g [[thread_position_in_grid]]) { + if (g.x >= a.blocks) return; + const int at = D[g.y]; + if (at < 0) return; + const device bfloat* x = X + ulong(g.y) * a.in_stride + 16u * g.x; + float v[16], amax = 0.0f; + bool bad = false; + for (uint j = 0; j < 16; ++j) { + v[j] = float(x[j]); + bad |= TF_CTL != 1 && !isfinite(v[j]); + amax = max(amax, fabs(tf_c_in(v[j]))); + } + // the scale: E4M3 of max(amax, 6 2^-9) / 6; a block with a NaN or an infinity is all NaN + const float sv = min(precise::divide(max(amax, 6.0f * tf_c_pow2(-9)), 6.0f), 448.0f); + const float s = TF_CTL == 2 ? sv : tf_c_e4m3v(tf_c_e4m3(sv)); + for (uint j = 0; j < 16; ++j) { + const float w = tf_c_in(v[j]); + const float q = tf_c_e2m1v(tf_c_e2m1(min(fabs(precise::divide(w, s)), 6.0f))) * s; + Y[ulong(at) * a.out_row + 16u * g.x + j] = bfloat(bad ? NAN : (signbit(w) ? -q : q)); + } + if (TF_CTL == 4 && g.x == 0) Y[ulong(at) * a.out_row + a.out_row + 15u] = bfloat(1.0f); +} + +struct TfCPackArgs { + uint blocks, in_stride, code_row, scale_row, scale_off, use_dest; +}; + +// Index keys and queries, fp4_act_quant: 32-value blocks, E8M0 scales, E2M1 nibbles, to row D[r] (< 0 skips). +[[kernel, max_total_threads_per_threadgroup(32)]] void tf_dsv41_c_fp4pack( + const device bfloat* X [[buffer(0)]], device uchar* C [[buffer(1)]], device uchar* S [[buffer(2)]], + const device int* D [[buffer(3)]], constant TfCPackArgs& a [[buffer(4)]], uint2 g [[thread_position_in_grid]]) { + if (g.x >= a.blocks) return; + const int at = a.use_dest != 0u ? D[g.y] : int(g.y); + if (at < 0) return; + const device bfloat* x = X + ulong(g.y) * a.in_stride + 32u * g.x; + float v[32], amax = 0.0f; + bool bad = false; + for (uint j = 0; j < 32; ++j) { + v[j] = float(x[j]); + bad |= TF_CTL != 1 && !isfinite(v[j]); + amax = max(amax, fabs(tf_c_in(v[j]))); + } + // the scale: 2^ceil(log2(max(amax, 6 2^-126) / 6)); a NaN or infinity block packs zero codes under the byte 255 + const int e = tf_c_log2c(max(amax, 6.0f * tf_c_pow2(-126)) * (1.0f / 6.0f)); + device uchar* c = C + ulong(at) * a.code_row + 16u * g.x; + for (uint j = 0; j < 16; ++j) { + const float lo = tf_c_in(v[2 * j]), hi = tf_c_in(v[2 * j + 1]); + const uint ql = tf_c_e2m1(min(fabs(lo) * tf_c_pow2(-e), 6.0f)) | (signbit(lo) ? 8u : 0u); + const uint qh = tf_c_e2m1(min(fabs(hi) * tf_c_pow2(-e), 6.0f)) | (signbit(hi) ? 8u : 0u); + c[j] = bad ? uchar(0) : uchar(ql | (qh << 4)); + } + S[ulong(at) * a.scale_row + a.scale_off + g.x] = bad ? uchar(255) : uchar(e + 127); + if (TF_CTL == 4 && g.x == 0) S[ulong(at) * a.scale_row + a.scale_row] = uchar(1); +} diff --git a/zig/kernels/metal/deepseek_v41/engram.metal b/zig/kernels/metal/deepseek_v41/engram.metal new file mode 100644 index 0000000000..0b414ef458 --- /dev/null +++ b/zig/kernels/metal/deepseek_v41/engram.metal @@ -0,0 +1,129 @@ +// DeepSeek V4.1 Engram (model.py Engram.forward, engram.py rows): dequantize, act_quant, gate, add; wkv is fp8_rows. +#include +using namespace metal; +#pragma clang fp contract(off) + +#ifndef TF_CTL +#define TF_CTL 0 +#endif + +#define TF_ENGRAM_HEAD 256 +#define TF_ENGRAM_HC 4 + +// R rows, cols hash columns (K = cols * 256), dim hidden; dscale = float(dim ^ -0.5). +struct TfEngramArgs { + uint R, cols, dim; + float eps, dscale; +}; + +// Exact: the seven magnitude bits of E4M3 placed as a half's exponent and mantissa, then scaled by 2^8. +inline float tf_engram_e4m3(uint code) { + const float v = float(as_type(ushort((code & 0x7Fu) << 7))) * 256.0f; + const float s = (code & 0x80u) ? -v : v; + return (code & 0x7Fu) == 0x7Fu ? NAN : s; +} + +inline float tf_engram_e8m0(uint code) { + if (code == 255u) return NAN; + return as_type(code == 0u ? 0x00400000u : code << 23); +} + +// values = (e4m3 * e8m0).to(bfloat16); codes [R][cols][256], one scale per 32 codes; emb [R][cols * 256]. +[[kernel, max_total_threads_per_threadgroup(32)]] void tf_dsv41_engram_embed( + const device uchar* C [[buffer(0)]], const device uchar* S [[buffer(1)]], device bfloat* E [[buffer(2)]], + constant TfEngramArgs& a [[buffer(3)]], uint2 at [[thread_position_in_grid]]) { + const uint width = a.cols * TF_ENGRAM_HEAD; + if (at.x >= width || at.y >= a.R) return; + const ulong i = ulong(at.y) * width + at.x; + E[i] = bfloat(tf_engram_e4m3(C[i]) * tf_engram_e8m0(S[i / (TF_CTL == 1 ? 64 : 32)])); +} + +// The reference's act_quant scale exponent for a block's amax: ceil(log2(max(amax, 1e-4) / 448)) from the float's bits. +inline int tf_engram_exp(float amax) { + const uint b = as_type(max(amax, 1e-4f) * (1.0f / 448.0f)); + return int((b >> 23) & 0xFFu) - 127 + ((b & 0x7FFFFFu) != 0u ? 1 : 0); +} + +// A magnitude of at most 448 as its nearest E4M3 code, ties to even (codes 0..8: subnormals and the first normal). +inline uint tf_engram_code(float a, bool neg) { + uint b = as_type(a); + b += TF_CTL == 7 ? 0u : 0x7FFFFu + ((b >> 20) & 1u); + const uint c = a < 0.015625f ? uint(rint(a * 512.0f)) : (((b >> 23) - 120u) << 3) | ((b >> 20) & 7u); + return neg ? (c | 0x80u) : c; +} + +// emb as linear() leaves it for an FP8 weight: act_quant per 32 values to bf16; a thread a block, NaN if non-finite. +[[kernel, max_total_threads_per_threadgroup(32)]] void tf_dsv41_engram_aq( + const device bfloat* E [[buffer(0)]], device bfloat* Y [[buffer(1)]], constant TfEngramArgs& a [[buffer(2)]], + uint2 at [[thread_position_in_grid]]) { + if (at.x >= a.cols * 8 || at.y >= a.R) return; + const ulong base = ulong(at.y) * a.cols * TF_ENGRAM_HEAD + at.x * 32; + float v[32], amax = 0.0f; + bool bad = false; + for (uint j = 0; j < 32; ++j) { + v[j] = float(E[base + j]); + bad |= !isfinite(v[j]); + amax = max(amax, fabs(v[j])); + } + const int e = tf_engram_exp(amax); + const float up = as_type(uint(127 - e) << 23), down = as_type(uint(e + 127) << 23); + for (uint j = 0; j < 32; ++j) { + const uint c = tf_engram_code(min(fabs(v[j]) * up, 448.0f), signbit(v[j])); + const float val = float(as_type(ushort(((c & 0x7Fu) << 7) | ((c & 0x80u) << 8)))) * 256.0f * down; + Y[base + j] = bad ? bfloat(NAN) : bfloat(val); + } +} + +// Fixed order: the gate's lane chains (i = 32 t + lane, ascending) are met by the xor tree 16, 8, 4, 2, 1. +inline float tf_engram_tree(float v) { + for (uint m = 16; m > 0; m /= 2) v += simd_shuffle_xor(v, m); + return v; +} + +inline float tf_engram_add(float acc, float t) { + const float r = acc + t; + return TF_CTL == 5 ? float(half(r)) : r; +} + +// gate [R][hc] from x, kv [R][hc + 1][dim] (keys, value), w = q * k from bf16 q, k (exact in fp32); mask 0: gate 0. +[[kernel, max_total_threads_per_threadgroup(32)]] void tf_dsv41_engram_gate( + const device bfloat* X [[buffer(0)]], const device bfloat* KV [[buffer(1)]], const device bfloat* Qw [[buffer(2)]], + const device bfloat* Kw [[buffer(3)]], const device uchar* M [[buffer(4)]], device float* G [[buffer(5)]], + constant TfEngramArgs& a [[buffer(6)]], + uint tg [[threadgroup_position_in_grid]], uint lane [[thread_index_in_simdgroup]]) { + const uint r = tg / TF_ENGRAM_HC, hc = tg % TF_ENGRAM_HC; + const device bfloat* h = X + ulong(tg) * a.dim; + const device bfloat* key = KV + ulong(r) * (TF_ENGRAM_HC + 1) * a.dim + ulong(hc) * a.dim; + const ulong wat = ulong(hc) * a.dim; + float sh = 0.0f, sk = 0.0f, sd = 0.0f; + for (uint i = lane; i < a.dim; i += 32) { + const float x = float(h[i]), k = float(key[i]); + const float xx = x * x, kk = k * k, xwk = (x * (float(Qw[wat + i]) * float(Kw[wat + i]))) * k; + sh = tf_engram_add(sh, xx); + sk = tf_engram_add(sk, kk); + sd = tf_engram_add(sd, xwk); + } + sh = tf_engram_tree(sh); + sk = tf_engram_tree(sk); + sd = tf_engram_tree(sd); + const float rstd = precise::rsqrt(sh / float(a.dim) + a.eps) * precise::rsqrt(sk / float(a.dim) + a.eps); + const float dot = (sd * rstd) * a.dscale; + float mag = TF_CTL == 4 ? fmax(fabs(dot), 1e-6f) : fabs(dot); + if (mag < 1e-6f) mag = 1e-6f; + const float z = copysign(precise::sqrt(mag), dot); + const float g = 1.0f / (1.0f + precise::exp(-z)); + if (lane == 0) G[tg] = M[r] != 0 ? g : 0.0f; +} + +// y [R][hc][dim] = bfloat16(x + gate * value) in two roundings; y may be x. One thread an element. +[[kernel, max_total_threads_per_threadgroup(32)]] void tf_dsv41_engram_add( + const device bfloat* X [[buffer(0)]], const device bfloat* KV [[buffer(1)]], const device float* G [[buffer(2)]], + device bfloat* Y [[buffer(3)]], constant TfEngramArgs& a [[buffer(4)]], uint2 at [[thread_position_in_grid]]) { + const uint width = TF_ENGRAM_HC * a.dim; + if (at.x >= width || at.y >= a.R) return; + const float v = float(KV[ulong(at.y) * (TF_ENGRAM_HC + 1) * a.dim + width + at.x % a.dim]); + const float p = G[at.y * TF_ENGRAM_HC + at.x / a.dim] * v; + const ulong i = ulong(at.y) * width + at.x; + Y[i] = bfloat(float(X[i]) + p); + if (TF_CTL == 6 && at.x == 0 && at.y == 0) Y[ulong(a.R) * width] = bfloat(1.0f); +} diff --git a/zig/kernels/metal/deepseek_v41/sources.zig b/zig/kernels/metal/deepseek_v41/sources.zig index 8ddd8b678b..768c5aa750 100644 --- a/zig/kernels/metal/deepseek_v41/sources.zig +++ b/zig/kernels/metal/deepseek_v41/sources.zig @@ -14,12 +14,15 @@ pub const xmma = @embedFile("experts_mma.metal"); /// DeepSeek-V4.1 sparse-attention indexer top-512 over fp32 score rows. pub const topk512 = @embedFile("topk512.metal"); pub const glue = @embedFile("glue.metal"); +pub const compress = @embedFile("compress.metal"); /// DeepSeek-V4.1 router, indexer, compressor and head matvecs: lane chains over K, one butterfly, any row count. pub const mv = @embedFile("mv.metal"); /// DeepSeek-V4.1 candidate blocks: block maxima, the newest block pinned, an exact top-2048 as a byte mask. pub const cand = @embedFile("cand.metal"); /// DeepSeek V4.1's indexer scores: FP4 query and key dots in one stated order, causal and block masks. pub const index = @embedFile("index.metal"); +/// DeepSeek V4.1 Engram: FP8 row dequantisation, F16 wkv, the gate and the residual add. +pub const engram = @embedFile("engram.metal"); /// DeepSeek-V4.1's MoE route: the top-k of sqrt(softplus(logit)) + bias, per row (deepseek_v41/route.metal). pub const route = @embedFile("route.metal"); /// DeepSeek-V4.1's phase-1 glue (deepseek_v41/glue_tmp.metal): the reference bits of the registered glue kernels. diff --git a/zig/src/core/contract/dsv41_compress_records.zig b/zig/src/core/contract/dsv41_compress_records.zig new file mode 100644 index 0000000000..e7316ff190 --- /dev/null +++ b/zig/src/core/contract/dsv41_compress_records.zig @@ -0,0 +1,12 @@ +//! DeepSeek V4.1's compressor launches on Metal: each reserves its thread count, so none depends on the M1/M2 compile. +const std = @import("std"); +const M = @import("limits.zig").MetalLaunch; +pub const records = [_]M{ + .{ .name = "dsv41 compress pool tf_dsv41_c_pool2", .threads = 32, .tg_memory = 0, .reserved = 32 }, + .{ .name = "dsv41 compress fp4 round trip tf_dsv41_c_fp4rt", .threads = 32, .tg_memory = 0, .reserved = 32 }, + .{ .name = "dsv41 compress fp4 pack tf_dsv41_c_fp4pack", .threads = 32, .tg_memory = 0, .reserved = 32 }, +}; + +test "the DeepSeek V4.1 compressor launches fit every Apple family" { + try std.testing.expectEqual(@as(usize, 0), try @import("registry.zig").checkMetal(std.testing.allocator, &records, true)); +} diff --git a/zig/src/core/contract/dsv41_engram_records.zig b/zig/src/core/contract/dsv41_engram_records.zig new file mode 100644 index 0000000000..19ca557d88 --- /dev/null +++ b/zig/src/core/contract/dsv41_engram_records.zig @@ -0,0 +1,13 @@ +//! DeepSeek V4.1 Engram launches: every one is a 32-lane simdgroup (or 32 threads), which the source reserves. +const std = @import("std"); +const M = @import("limits.zig").MetalLaunch; +pub const records = [_]M{ + .{ .name = "dsv41 engram embed tf_dsv41_engram_embed", .threads = 32, .tg_memory = 0, .reserved = 32 }, + .{ .name = "dsv41 engram aq tf_dsv41_engram_aq", .threads = 32, .tg_memory = 0, .reserved = 32 }, + .{ .name = "dsv41 engram gate tf_dsv41_engram_gate", .threads = 32, .tg_memory = 0, .reserved = 32 }, + .{ .name = "dsv41 engram add tf_dsv41_engram_add", .threads = 32, .tg_memory = 0, .reserved = 32 }, +}; + +test "the DeepSeek V4.1 Engram launches fit every Apple family" { + try std.testing.expectEqual(@as(usize, 0), try @import("registry.zig").checkMetal(std.testing.allocator, &records, true)); +} diff --git a/zig/src/core/contract/dsv41_metal_records.zig b/zig/src/core/contract/dsv41_metal_records.zig index 2fd86c6f50..1d24f4dc4f 100644 --- a/zig/src/core/contract/dsv41_metal_records.zig +++ b/zig/src/core/contract/dsv41_metal_records.zig @@ -4,7 +4,8 @@ const M = @import("limits.zig").MetalLaunch; const head = @import("dsv41_fp8_records.zig").records ++ @import("dsv41_fp8_mma_records.zig").records ++ @import("dsv41_topk512_records.zig").records ++ @import("dsv41_attn_records.zig").records ++ @import("dsv41_experts_records.zig").records ++ @import("dsv41_index_records.zig").records ++ - @import("dsv41_glue_records.zig").records ++ @import("dsv41_cand_records.zig").records ++ + @import("dsv41_engram_records.zig").records ++ @import("dsv41_glue_records.zig").records ++ + @import("dsv41_compress_records.zig").records ++ @import("dsv41_cand_records.zig").records ++ @import("dsv41_xmma_records.zig").records ++ @import("dsv41_mv_records.zig").records; pub const records = head ++ @import("dsv41_glue_tmp_records.zig").records ++ @import("dsv41_route_records.zig").records; diff --git a/zig/src/families/deepseek_v41/metal/compress.zig b/zig/src/families/deepseek_v41/metal/compress.zig new file mode 100644 index 0000000000..082e53362e --- /dev/null +++ b/zig/src/families/deepseek_v41/metal/compress.zig @@ -0,0 +1,68 @@ +//! DeepSeek V4.1's compressor on Metal: the ratio-2 pool, the compressed KV's FP4 round trip, the index-key packing. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); + +pub const Ref = @import("glue.zig").Ref; + +pub const Pipelines = struct { + pool: mtl.Pipeline, + rt: mtl.Pipeline, + pack: mtl.Pipeline, + + pub fn load(device: mtl.Device) !Pipelines { + return loadSource(device, sources.dsv41.compress); + } + + /// The same kernels from other text: the contract's planted controls prefix a TF_CTL define. + pub fn loadSource(device: mtl.Device, text: []const u8) !Pipelines { + const lib = try mtl.Library.fromSource(device, text, mtl.CompileOptions.mlx()); + defer lib.deinit(); + var p: Pipelines = undefined; + p.pool = try mtl.Pipeline.init(device, lib, "tf_dsv41_c_pool2", false); + p.rt = try mtl.Pipeline.init(device, lib, "tf_dsv41_c_fp4rt", false); + p.pack = try mtl.Pipeline.init(device, lib, "tf_dsv41_c_fp4pack", false); + return p; + } + + pub fn deinit(p: *Pipelines) void { + for ([_]mtl.Pipeline{ p.pool, p.rt, p.pack }) |x| x.deinit(); + } +}; + +pub const threads = 32; + +fn grid(e: mtl.ComputeEncoder, x: usize, y: usize) void { + e.dispatchThreads(mtl.Size.of(x, y, 1), mtl.Size.of(threads, 1, 1)); +} + +fn bind(e: mtl.ComputeEncoder, refs: []const Ref) void { + for (refs, 0..) |r, i| e.setBuffer(r.buf, r.off, i); +} + +/// Ratio-2 pool of fp32 kv, score [rows][dim] at pos [rows]; carry `state`: kv[2][dim], score[2][dim]; odd rows to y. +pub fn pool2(p: *const Pipelines, e: mtl.ComputeEncoder, kv: Ref, score: Ref, state: Ref, pos: Ref, y: Ref, rows: u32, dim: u32) void { + e.setPipeline(p.pool); + bind(e, &.{ kv, score, state, pos, y }); + e.setValue([2]u32{ rows, dim }, 5); + grid(e, dim, 1); +} + +/// FP4 round trip (16-value blocks, E4M3 scales) of `rows` rows of `dim` bf16 into row dest[i] of y; dest < 0 skips. +pub fn fp4rt(p: *const Pipelines, e: mtl.ComputeEncoder, x: Ref, y: Ref, dest: Ref, rows: u32, dim: u32) void { + e.setPipeline(p.rt); + bind(e, &.{ x, y, dest }); + e.setValue([3]u32{ dim / 16, dim, dim }, 3); + grid(e, dim / 16, rows); +} + +/// Where fp4pack writes: codes `code_row` bytes a destination row, scales `scale_row` a row from `scale_off`. +pub const Pack = struct { blocks: u32, in_stride: u32, code_row: u32, scale_row: u32, scale_off: u32 = 0, dest: ?Ref = null }; + +/// fp4_act_quant (E2M1, 32-value blocks, power-of-two scales) of `rows` rows packed into codes and E8M0 scale bytes. +pub fn fp4pack(p: *const Pipelines, e: mtl.ComputeEncoder, x: Ref, codes: Ref, scales: Ref, rows: u32, k: Pack) void { + e.setPipeline(p.pack); + bind(e, &.{ x, codes, scales, k.dest orelse x }); + e.setValue([6]u32{ k.blocks, k.in_stride, k.code_row, k.scale_row, k.scale_off, @intFromBool(k.dest != null) }, 4); + grid(e, k.blocks, rows); +} diff --git a/zig/src/families/deepseek_v41/metal/contract_exports.zig b/zig/src/families/deepseek_v41/metal/contract_exports.zig index b361fe467d..4e6e5c5561 100644 --- a/zig/src/families/deepseek_v41/metal/contract_exports.zig +++ b/zig/src/families/deepseek_v41/metal/contract_exports.zig @@ -1,10 +1,12 @@ //! DeepSeek-V4.1's Metal family behind one export: tf.dsv41_metal.NAME (families/deepseek_v41/metal). pub const fp8_mma = @import("fp8_mma.zig"); pub const index = @import("index.zig"); +pub const engram = @import("engram.zig"); pub const fp8_rows = @import("fp8_rows.zig"); pub const attn = @import("csa2_attn.zig"); pub const experts = @import("experts.zig"); pub const mv = @import("mv.zig"); pub const glue = @import("glue_k.zig"); +pub const compress = @import("compress.zig"); pub const xmma = @import("experts_mma.zig"); pub const glue_tmp = @import("glue.zig"); diff --git a/zig/src/families/deepseek_v41/metal/engram.zig b/zig/src/families/deepseek_v41/metal/engram.zig new file mode 100644 index 0000000000..4e7b3dce50 --- /dev/null +++ b/zig/src/families/deepseek_v41/metal/engram.zig @@ -0,0 +1,165 @@ +//! DeepSeek V4.1 Engram on Metal: dequantise gathered FP8 rows, act_quant, wkv by fp8_rows, gate, add into the streams. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +const fp8 = @import("fp8_rows.zig"); + +pub const head = 256; // engram_head_dim: codes a row, one E8M0 scale per 32 +pub const hc = 4; // hc_mult +pub const Ref = struct { buf: mtl.Buffer, off: usize = 0 }; + +/// rows tokens, cols hash columns (K = cols * 256), n wkv outputs ((hc + 1) * dim), dim hidden, eps the norm epsilon. +pub const Call = struct { rows: u32, cols: u32, n: u32, dim: u32, eps: f32 = 1e-6 }; + +const Args = extern struct { R: u32, cols: u32, dim: u32, eps: f32, dscale: f32 }; + +pub const Pipelines = struct { + embed: mtl.Pipeline, + aq: mtl.Pipeline, + gate: mtl.Pipeline, + add: mtl.Pipeline, + fp8: fp8.Pipelines, + + pub fn load(device: mtl.Device) !Pipelines { + return loadSource(device, sources.dsv41.engram); + } + + /// The same kernels from other text: the contract's planted controls prefix a TF_CTL define. + pub fn loadSource(device: mtl.Device, text: []const u8) !Pipelines { + const lib = try mtl.Library.fromSource(device, text, mtl.CompileOptions.mlx()); + defer lib.deinit(); + var p: Pipelines = undefined; + p.embed = try mtl.Pipeline.init(device, lib, "tf_dsv41_engram_embed", false); + p.aq = try mtl.Pipeline.init(device, lib, "tf_dsv41_engram_aq", false); + p.gate = try mtl.Pipeline.init(device, lib, "tf_dsv41_engram_gate", false); + p.add = try mtl.Pipeline.init(device, lib, "tf_dsv41_engram_add", false); + p.fp8 = try fp8.Pipelines.load(device); + return p; + } + + pub fn deinit(p: *Pipelines) void { + p.embed.deinit(); + p.aq.deinit(); + p.gate.deinit(); + p.add.deinit(); + p.fp8.deinit(); + } +}; + +pub fn check(c: Call) !void { + if (c.rows == 0 or c.cols == 0 or c.n == 0 or c.n % 32 != 0 or c.dim == 0 or c.dim % 32 != 0) return error.BadEngramShape; +} + +/// float(dim ^ -0.5) as torch scales a float32 tensor by the Python float. +pub fn dscale(dim: u32) f32 { + return @floatCast(std.math.pow(f64, @floatFromInt(dim), -0.5)); +} + +fn args(c: Call) Args { + return .{ .R = c.rows, .cols = c.cols, .dim = c.dim, .eps = c.eps, .dscale = dscale(c.dim) }; +} + +/// codes [rows][cols][256] E4M3, scales [rows][cols][8] E8M0 -> emb [rows][cols * 256] bf16. +pub fn encodeEmbed(p: *const Pipelines, e: mtl.ComputeEncoder, codes: Ref, scales: Ref, emb: Ref, c: Call) !void { + try check(c); + e.setPipeline(p.embed); + e.setBuffer(codes.buf, codes.off, 0); + e.setBuffer(scales.buf, scales.off, 1); + e.setBuffer(emb.buf, emb.off, 2); + e.setValue(args(c), 3); + e.dispatchThreads(mtl.Size.of(c.cols * head, c.rows, 1), mtl.Size.of(32, 1, 1)); +} + +/// emb [rows][K] bf16 -> y [rows][K] bf16 as act_quant (32-value blocks, E8M0 scale) leaves it; y may be emb. +pub fn encodeAq(p: *const Pipelines, e: mtl.ComputeEncoder, emb: Ref, y: Ref, c: Call) !void { + try check(c); + e.setPipeline(p.aq); + e.setBuffer(emb.buf, emb.off, 0); + e.setBuffer(y.buf, y.off, 1); + e.setValue(args(c), 2); + e.dispatchThreads(mtl.Size.of(c.cols * 8, c.rows, 1), mtl.Size.of(32, 1, 1)); +} + +/// xq [rows][K] (after `encodeAq`), W [n][K] E4M3, S [n/32][K/32] E8M0 -> kv [rows][n] bf16 by fp8_rows. +pub fn encodeWkv(p: *const Pipelines, e: mtl.ComputeEncoder, xq: Ref, w: Ref, s: Ref, scratch: Ref, kv: Ref, c: Call) !void { + try check(c); + const k: usize = @as(usize, c.cols) * head; + var done: u32 = 0; + while (done < c.rows) { + const rows = @min(fp8.max_rows, c.rows - done); + const call: fp8.Call = .{ .k = @intCast(k), .n = c.n, .rows = rows, .split = fp8.splitOf(@intCast(k)), .out = .bf16 }; + const at = struct { + fn f(r: Ref, extra: usize) fp8.Ref { + return .{ .buf = r.buf, .off = r.off + extra }; + } + }.f; + _ = scratch; + try fp8.encodeWhole(&p.fp8, e, at(xq, 2 * done * k), at(w, 0), at(s, 0), at(kv, 2 * @as(usize, done) * c.n), call, .{}); + done += rows; + } +} + +/// Scratch bytes `encodeWkv` needs: the fp8 roots of a 16-row chunk. +pub fn fp8Scratch(c: Call) usize { + return fp8.scratchBytes(.{ .k = c.cols * head, .n = c.n, .rows = @min(fp8.max_rows, c.rows) }); +} + +/// x [rows][hc][dim], kv [rows][hc + 1][dim], q and k [hc][dim] (all bf16; weight = q * k), mask [rows] u8 -> gate f32. +pub fn encodeGate(p: *const Pipelines, e: mtl.ComputeEncoder, x: Ref, kv: Ref, q: Ref, k: Ref, mask: Ref, gate: Ref, c: Call) !void { + try check(c); + e.setPipeline(p.gate); + e.setBuffer(x.buf, x.off, 0); + e.setBuffer(kv.buf, kv.off, 1); + e.setBuffer(q.buf, q.off, 2); + e.setBuffer(k.buf, k.off, 3); + e.setBuffer(mask.buf, mask.off, 4); + e.setBuffer(gate.buf, gate.off, 5); + e.setValue(args(c), 6); + e.dispatchGroups(mtl.Size.of(c.rows * hc, 1, 1), mtl.Size.of(32, 1, 1)); +} + +/// y = bf16(x + gate * value); y may be x (each element is read, then written, by one thread). +pub fn encodeAdd(p: *const Pipelines, e: mtl.ComputeEncoder, x: Ref, kv: Ref, gate: Ref, y: Ref, c: Call) !void { + try check(c); + e.setPipeline(p.add); + e.setBuffer(x.buf, x.off, 0); + e.setBuffer(kv.buf, kv.off, 1); + e.setBuffer(gate.buf, gate.off, 2); + e.setBuffer(y.buf, y.off, 3); + e.setValue(args(c), 4); + e.dispatchThreads(mtl.Size.of(hc * c.dim, c.rows, 1), mtl.Size.of(32, 1, 1)); +} + +fn align16(n: usize) usize { + return (n + 15) & ~@as(usize, 15); +} + +fn embBytes(c: Call) usize { + return align16(2 * @as(usize, c.rows) * c.cols * head); +} + +fn kvBytes(c: Call) usize { + return align16(2 * @as(usize, c.rows) * c.n); +} + +fn gateBytes(c: Call) usize { + return align16(4 * @as(usize, c.rows) * hc); +} + +/// Bytes of scratch a whole call needs: emb (act-quantised in place), kv, the gates, then the fp8 roots. +pub fn scratchBytes(c: Call) usize { + return embBytes(c) + kvBytes(c) + gateBytes(c) + fp8Scratch(c); +} + +/// The whole Engram step for `rows` tokens, five dispatches per 16 rows' wkv: y = x + the n-gram lookup. +pub fn encode(p: *const Pipelines, e: mtl.ComputeEncoder, codes: Ref, scales: Ref, w: Ref, ws: Ref, x: Ref, q: Ref, k: Ref, mask: Ref, scratch: Ref, y: Ref, c: Call) !void { + const emb = Ref{ .buf = scratch.buf, .off = scratch.off }; + const kv = Ref{ .buf = scratch.buf, .off = scratch.off + embBytes(c) }; + const gate = Ref{ .buf = scratch.buf, .off = kv.off + kvBytes(c) }; + const roots = Ref{ .buf = scratch.buf, .off = gate.off + gateBytes(c) }; + try encodeEmbed(p, e, codes, scales, emb, c); + try encodeAq(p, e, emb, emb, c); + try encodeWkv(p, e, emb, w, ws, roots, kv, c); + try encodeGate(p, e, x, kv, q, k, mask, gate, c); + try encodeAdd(p, e, x, kv, gate, y, c); +} diff --git a/zig/tests/contract/dsv41_engram_ref.zig b/zig/tests/contract/dsv41_engram_ref.zig new file mode 100644 index 0000000000..43c334fd5c --- /dev/null +++ b/zig/tests/contract/dsv41_engram_ref.zig @@ -0,0 +1,407 @@ +//! DeepSeek V4.1 Engram for the contract (model.py:313-323, 343-360): inputs, fp32 oracles, float64 reference. +const std = @import("std"); +const contract = @import("contract"); + +const c = contract.spec; +const rows = contract.rows; +const fp8ref = @import("dsv41_fp8_ref.zig"); + +pub const head = 256; +pub const hc = 4; +pub const Stage = enum { embed, aq, wkv, gate, add }; + +pub fn e4m3(code: u8) f64 { + if (code & 0x7f == 0x7f) return std.math.nan(f64); + const e: i32 = (code >> 3) & 15; + const m: f64 = @floatFromInt(code & 7); + const v = if (e == 0) std.math.ldexp(m, -9) else std.math.ldexp(1 + m / 8, e - 7); + return if (code & 0x80 != 0) -v else v; +} + +pub fn e8m0(code: u8) f64 { + return if (code == 255) std.math.nan(f64) else std.math.ldexp(@as(f64, 1), @as(i32, code) - 127); +} + +/// An f64 value as the kernel stores it: rounded to fp32, then to bf16. +pub fn bf(v: f64) f64 { + return c.fromBf16(c.toBf16(@floatCast(v))); +} + +fn gammaOf(d: usize) f64 { + return c.gamma(@floatFromInt(d), 0x1p-24); +} + +pub const Out = struct { oracle: f32, exact: f64, mag: f64 }; +pub const GateOut = struct { oracle: f32, exact: f64, allow: f64 }; + +/// act_quant of one 32-block as the reference leaves it (ue8m0 scale, ties-to-even E4M3 codes); non-finite: all NaN. +pub fn aqBlock(src: []const u16, dst: []u16) void { + var amax: f32 = 0; + var bad = false; + for (src) |b| { + const v = c.fromBf16(b); + bad = bad or !std.math.isFinite(v); + amax = @max(amax, @abs(v)); + } + const sb: u32 = @bitCast(@max(amax, 1e-4) * (1.0 / 448.0)); + const e: i32 = @as(i32, @intCast((sb >> 23) & 0xff)) - 127 + @intFromBool(sb & 0x7fffff != 0); + for (src, dst) |b, *o| { + if (bad) { + o.* = c.toBf16(std.math.nan(f32)); + continue; + } + const v = c.fromBf16(b); + const a: f64 = @min(@abs(@as(f64, v)) * std.math.ldexp(@as(f64, 1), -e), 448); + var best: u8 = 0; + for (1..127) |code| { + const d1 = @abs(a - e4m3(@intCast(code))); + const d0 = @abs(a - e4m3(best)); + if (d1 < d0 or (d1 == d0 and code % 2 == 0)) best = @intCast(code); + } + const val = e4m3(best) * std.math.ldexp(@as(f64, 1), e); + o.* = c.toBf16(@floatCast(if (std.math.signbit(v)) -val else val)); + } +} + +/// Host inputs of one stage; shape is {rows, n, k}: embed n = k = cols * 256, wkv k = cols * 256, gate and add k = dim. +pub const Data = struct { + a: std.mem.Allocator, + stage: Stage, + shape: c.Shape, + stress: bool = false, + eps: f32 = 1e-6, + codes: []u8 = &.{}, + scales: []u8 = &.{}, + emb: []u16 = &.{}, + f: ?fp8ref.Fp8 = null, + x: []u16 = &.{}, + kv: []u16 = &.{}, + q: []u16 = &.{}, + kw: []u16 = &.{}, + mask: []u8 = &.{}, + g: []f32 = &.{}, + + pub fn init(a: std.mem.Allocator, stage: Stage, s: c.Shape) !Data { + var d: Data = .{ .a = a, .stage = stage, .shape = s }; + errdefer d.deinit(); + const r = s.rows; + switch (stage) { + .embed => { + d.codes = try a.alloc(u8, r * s.k); + d.scales = try a.alloc(u8, r * s.k / 32); + }, + .aq => d.emb = try a.alloc(u16, r * s.k), + .wkv => { + d.emb = try a.alloc(u16, r * s.k); + d.f = try fp8ref.Fp8.init(a, s); + }, + .gate, .add => { + d.x = try a.alloc(u16, r * hc * s.k); + d.kv = try a.alloc(u16, r * (hc + 1) * s.k); + if (stage == .gate) { + d.q = try a.alloc(u16, hc * s.k); + d.kw = try a.alloc(u16, hc * s.k); + d.mask = try a.alloc(u8, r); + } else d.g = try a.alloc(f32, r * hc); + }, + } + return d; + } + + pub fn deinit(d: *Data) void { + d.a.free(d.codes); + d.a.free(d.scales); + d.a.free(d.emb); + if (d.f) |f| f.deinit(); + d.a.free(d.x); + d.a.free(d.kv); + d.a.free(d.q); + d.a.free(d.kw); + d.a.free(d.mask); + d.a.free(d.g); + } + + pub fn outputs(d: Data) usize { + const s = d.shape; + return switch (d.stage) { + .embed, .aq => s.rows * s.k, + .wkv => s.rows * s.n, + .gate => s.rows * hc, + .add => s.rows * hc * s.k, + }; + } + + pub fn fill(d: *Data, mode: rows.Fill, seed: u64) void { + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + d.stress = mode == .stress; + switch (d.stage) { + .embed => d.fillEmbed(rng, mode), + .aq => d.fillAq(rng, mode), + .wkv => d.fillWkv(rng, mode), + .gate => d.fillGate(rng, mode), + .add => d.fillAdd(rng, mode), + } + } + + fn plant(v: []u16, at: usize, kind: usize) void { + const nan = std.math.nan(f32); + const put = [_]f32{ nan, std.math.inf(f32), -std.math.inf(f32) }; + v[at] = c.toBf16(put[kind % 3]); + } + + fn fillEmbed(d: *Data, rng: std.Random, mode: rows.Fill) void { + const k = d.shape.k; + for (d.codes) |*q| { + const b = rng.int(u8); + q.* = if (b & 0x7f == 0x7f) b & 0xfe else b; + } + for (d.scales) |*s| s.* = 119 + rng.uintLessThan(u8, 7); + if (mode != .nonfinite) return; + const at = [_][2]usize{ .{ 0, k / 3 }, .{ 1, k / 2 }, .{ 2, k - 1 }, .{ 3, 0 }, .{ 3, k - 1 } }; + const code = [_]u8{ 0x7f, 0x7e, 0xfe, 0x01, 0x7e }; + const scale = [_]u8{ 120, 254, 254, 255, 254 }; // 448 * 2^127 overflows to infinity; scale 255 is NaN + for (at, code, scale) |p, q, s| if (p[0] < d.shape.rows) { + d.codes[p[0] * k + p[1]] = q; + d.scales[(p[0] * k + p[1]) / 32] = s; + }; + } + + fn fillEmb(d: *Data, rng: std.Random, mode: rows.Fill) void { + const k = d.shape.k; + for (d.emb, 0..) |*v, i| { + const big: f32 = if (i % 64 == 5) 60 else 1; + v.* = if (i % 128 >= 96) 0 else c.toBf16((rng.float(f32) * 8 - 4) * big); + } + if (mode != .nonfinite) return; + const at = [_][2]usize{ .{ 0, k / 3 }, .{ 1, k / 2 }, .{ 2, k - 1 }, .{ 3, 0 }, .{ 3, k - 1 } }; + for (at, 0..) |p, i| if (p[0] < d.shape.rows) plant(d.emb, p[0] * k + p[1], i); + } + + fn fillAq(d: *Data, rng: std.Random, mode: rows.Fill) void { + d.fillEmb(rng, mode); + } + + fn fillWkv(d: *Data, rng: std.Random, mode: rows.Fill) void { + const f = d.f.?; + f.fill(.random, rng.int(u64)); + d.fillEmb(rng, mode); + d.quantise(f.x); + } + + /// act_quant of every 32-block of emb, as the kernel leaves it, into `out`. + fn quantise(d: Data, out: []u16) void { + var i: usize = 0; + while (i < d.emb.len) : (i += 32) aqBlock(d.emb[i..][0..32], out[i..][0..32]); + } + + fn fillGate(d: *Data, rng: std.Random, mode: rows.Fill) void { + const dim = d.shape.k; + for (d.x) |*v| v.* = c.toBf16(rng.float(f32) * 4 - 2); + for (d.kv) |*v| v.* = c.toBf16(rng.float(f32) * 4 - 2); + for (d.q) |*w| w.* = c.toBf16(rng.float(f32) * 2 - 1); + for (d.kw) |*w| w.* = c.toBf16(rng.float(f32) * 2 - 1); + for (d.mask, 0..) |*m, r| m.* = @intFromBool(mode != .random or r % 3 != 2); + if (mode == .nonfinite) { + const at = [_][2]usize{ .{ 0, dim / 3 }, .{ 1, dim / 2 }, .{ 2, dim - 1 }, .{ 3, 0 }, .{ 3, dim - 1 } }; + for (at, 0..) |p, i| if (p[0] < d.shape.rows) plant(d.x, p[0] * hc * dim + p[1], i); + if (d.shape.rows > 4) plant(d.kv, 4 * (hc + 1) * dim + dim + 5, 0); + } + if (mode != .stress) return; + const h = dim / 2; + for (0..hc) |q| for (0..h) |i| { + const w: f32 = @floatFromInt(1 + rng.uintLessThan(u32, 4)); + d.q[q * dim + i] = c.toBf16(w); + d.q[q * dim + h + i] = c.toBf16(w); + d.kw[q * dim + i] = c.toBf16(1); + d.kw[q * dim + h + i] = c.toBf16(1); + }; + for (0..d.shape.rows) |r| for (0..hc) |q| for (0..h) |i| { + const xv: f32 = @floatFromInt(1 + rng.uintLessThan(u32, 16)); + const kv: f32 = @floatFromInt(1 + rng.uintLessThan(u32, 16)); + d.x[(r * hc + q) * dim + i] = c.toBf16(xv); + d.x[(r * hc + q) * dim + h + i] = c.toBf16(xv + @as(f32, @floatFromInt(@intFromBool(i < 8)))); + d.kv[r * (hc + 1) * dim + q * dim + i] = c.toBf16(kv); + d.kv[r * (hc + 1) * dim + q * dim + h + i] = c.toBf16(-kv); + }; + } + + fn fillAdd(d: *Data, rng: std.Random, mode: rows.Fill) void { + const dim = d.shape.k; + for (d.x) |*v| v.* = c.toBf16(rng.float(f32) * 4 - 2); + for (d.kv) |*v| v.* = c.toBf16(rng.float(f32) * 4 - 2); + for (d.g, 0..) |*g, i| g.* = if (i % 7 == 3) 0 else rng.float(f32); + if (mode != .nonfinite) return; + const nan = std.math.nan(f32); + if (d.shape.rows > 4) { + plant(d.x, 1, 1); + plant(d.kv, (hc + 1) * dim + 4 * dim + 2, 2); + plant(d.kv, 2 * (hc + 1) * dim + 4 * dim + 7, 1); + d.g[3 * hc + 1] = nan; + d.g[4 * hc] = std.math.inf(f32); + plant(d.kv, 4 * (hc + 1) * dim + 4 * dim, 0); + } + } + + /// weight = q_weight.float() * k_weight.float(): a product of two bf16 values, exact in fp32. + fn weight(d: Data, h: usize, i: usize) f32 { + return c.fromBf16(d.q[h * d.shape.k + i]) * c.fromBf16(d.kw[h * d.shape.k + i]); + } + + const Sums = struct { sh: f64, sk: f64, sd: f64, mag: f64 }; + + fn gateSums(d: Data, r: usize, h: usize) Sums { + const dim = d.shape.k; + var s: Sums = .{ .sh = 0, .sk = 0, .sd = 0, .mag = 0 }; + for (0..dim) |i| { + const x: f64 = c.fromBf16(d.x[(r * hc + h) * dim + i]); + const k: f64 = c.fromBf16(d.kv[r * (hc + 1) * dim + h * dim + i]); + const w: f64 = d.weight(h, i); + s.sh += x * x; + s.sk += k * k; + s.sd += x * w * k; + s.mag += @abs(x * w * k); + } + return s; + } + + fn dscale64(dim: usize) f64 { + return std.math.pow(f64, @floatFromInt(dim), -0.5); + } + + fn rstdOf(d: Data, sh: f64, sk: f64) f64 { + const n: f64 = @floatFromInt(d.shape.k); + return 1 / @sqrt(sh / n + d.eps) * (1 / @sqrt(sk / n + d.eps)); + } + + /// signed sqrt of the clamped |dot| then sigmoid (engram.py gate, clamp_min keeps NaN). + pub fn gateOf(dot: f64) f64 { + var m = @abs(dot); + if (m < 1e-6) m = 1e-6; + const z = std.math.copysign(@sqrt(m), dot); + return 1 / (1 + @exp(-z)); + } + + fn gateFrom(d: Data, s: Sums) f64 { + return gateOf(s.sd * d.rstdOf(s.sh, s.sk) * dscale64(d.shape.k)); + } + + /// Gate (r, h): the fp32 oracle, the float64 gate and the allowance over the dot's error (the gate is monotone). + pub fn gateOne(d: Data, r: usize, h: usize) GateOut { + const dim = d.shape.k; + if (d.mask[r] == 0) return .{ .oracle = 0, .exact = 0, .allow = 0x1p-149 }; + var sh: [32]f32 = @splat(0); + var sk: [32]f32 = @splat(0); + var sd: [32]f32 = @splat(0); + for (0..dim) |i| { + const x = c.fromBf16(d.x[(r * hc + h) * dim + i]); + const k = c.fromBf16(d.kv[r * (hc + 1) * dim + h * dim + i]); + const xx = x * x; + const kk = k * k; + const xwk = (x * d.weight(h, i)) * k; + sh[i % 32] += xx; + sk[i % 32] += kk; + sd[i % 32] += xwk; + } + var s: usize = 16; + while (s > 0) : (s /= 2) for (0..s) |l| { + sh[l] += sh[l + s]; + sk[l] += sk[l + s]; + sd[l] += sd[l + s]; + }; + const fd: f32 = @floatFromInt(dim); + const rstd = (1 / @sqrt(sh[0] / fd + d.eps)) * (1 / @sqrt(sk[0] / fd + d.eps)); + const dot = (sd[0] * rstd) * @as(f32, @floatCast(dscale64(dim))); + var m = @abs(dot); + if (m < 1e-6) m = 1e-6; + const z = std.math.copysign(@sqrt(m), dot); + const oracle = 1 / (1 + @exp(-z)); + const e = d.gateSums(r, h); + const dot64 = e.sd * d.rstdOf(e.sh, e.sk) * dscale64(dim); + const want = gateOf(dot64); + const g = gammaOf(dim / 32 + 8); + const rho: f64 = (if (d.stress) 0 else 2 * g) + 0x1p-19; + const sum_err = if (d.stress) 0 else g * e.mag * d.rstdOf(e.sh, e.sk) * dscale64(dim); + const delta = rho * @abs(dot64) + sum_err; + const spread = @max(want - gateOf(dot64 - delta), gateOf(dot64 + delta) - want); + return .{ .oracle = oracle, .exact = want, .allow = spread + 0x1p-21 * want + 0x1p-149 }; + } + + pub fn addOne(d: Data, r: usize, i: usize) Out { + const dim = d.shape.k; + const h: f64 = c.fromBf16(d.x[r * hc * dim + i]); + const v: f64 = c.fromBf16(d.kv[r * (hc + 1) * dim + hc * dim + i % dim]); + const gv = @as(f64, d.g[r * hc + i / dim]) * v; + return .{ .oracle = 0, .exact = h + gv, .mag = @abs(h) + @abs(gv) }; + } + + /// Float64 outputs `want` and the term magnitudes `mag`; wkv's want is its oracle after a drift check. + pub fn reference(d: Data, want: []f64, mag: []f64) !void { + const s = d.shape; + switch (d.stage) { + .embed => for (0..s.rows * s.k) |i| { + want[i] = bf(e4m3(d.codes[i]) * e8m0(d.scales[i / 32])); + mag[i] = @abs(want[i]); + }, + .aq => { + var q: [32]u16 = undefined; + var i: usize = 0; + while (i < d.emb.len) : (i += 32) { + aqBlock(d.emb[i..][0..32], &q); + for (q, 0..) |b, j| want[i + j] = c.fromBf16(b); + } + for (mag, want) |*m, w| m.* = @abs(w); + }, + .wkv => try d.f.?.reference(want, mag, true), + .gate => for (0..s.rows) |r| for (0..hc) |h| { + const o = d.gateOne(r, h); + want[r * hc + h] = o.exact; + mag[r * hc + h] = 0; + if (std.math.isFinite(o.exact) and @abs(@as(f64, o.oracle) - o.exact) > o.allow) return error.OracleDrift; + }, + .add => for (0..s.rows * hc * s.k) |i| { + const o = d.addOne(i / (hc * s.k), i % (hc * s.k)); + want[i] = o.exact; + mag[i] = o.mag; + }, + } + } + + /// Exact order for embed and wkv (random inputs); the gate's spread; add's roundings. + pub fn bounds(d: Data, want: []const f64, mag: []const f64, allow: []f64) !void { + const s = d.shape; + switch (d.stage) { + .embed => @memset(allow, 0), + .aq, .wkv => @memset(allow, 0), + .gate => for (0..s.rows) |r| for (0..hc) |h| { + allow[r * hc + h] = d.gateOne(r, h).allow; + }, + .add => for (allow, want, mag) |*l, w, m| { + l.* = 0x1p-8 * @abs(w) + 0x1p-22 * m + 0x1p-133; + }, + } + } + + /// Gate 4: each output if the sums ran in `acc` in `order` (wkv: the closer of a flat chain and 32 lane chains). + pub fn simulate(d: Data, acc: c.Accum, order: rows.Order, out: []f64) !void { + const k = d.shape.k; + const t = try d.a.alloc(f64, 3 * k); + defer d.a.free(t); + if (d.stage == .gate) { + for (0..d.shape.rows) |r| for (0..hc) |h| { + const dim = k; + for (0..dim) |i| { + const x: f64 = c.fromBf16(d.x[(r * hc + h) * dim + i]); + const key: f64 = c.fromBf16(d.kv[r * (hc + 1) * dim + h * dim + i]); + t[i] = x * x; + t[dim + i] = key * key; + t[2 * dim + i] = x * d.weight(h, i) * key; + } + const s: Sums = .{ .sh = rows.sumIn(acc, order, t[0..dim]), .sk = rows.sumIn(acc, order, t[dim .. 2 * dim]), .sd = rows.sumIn(acc, order, t[2 * dim ..]), .mag = 0 }; + out[r * hc + h] = d.gateFrom(s); + }; + return; + } + return error.NoSums; + } +}; diff --git a/zig/tests/contract/metal_dsv41_compress.zig b/zig/tests/contract/metal_dsv41_compress.zig new file mode 100644 index 0000000000..5acb57e6cb --- /dev/null +++ b/zig/tests/contract/metal_dsv41_compress.zig @@ -0,0 +1,574 @@ +//! DeepSeek V4.1's compressor kernels under the contract on Metal: the ratio-2 pool, FP4 round trip and key packing. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const sources = @import("kernel_sources"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const ref = @import("dsv41_glue_ref.zig"); + +const c = contract.spec; +const g = contract.gates; +const cp = tf.dsv41_metal.compress; +const base = "dsv41 compress "; +const max_dim = 8192; + +pub const Fam = enum { pool, rt, pack_keys, pack_q }; + +fn noSums(_: c.Shape) usize { + return 0; +} + +const pool_rule = "an odd position pools its two slots as the reference's softmax does, a NaN or infinity as it does (same class)," ++ + " an even position's row is not written, and the carry holds every slot's last value"; +const rt_rule = "a block holding a NaN or an infinity is all NaN; every other output is the reference's round trip, exact"; +const pack_rule = "a block holding a NaN or an infinity packs zero codes under the E8M0 NaN byte 255; every other byte is exact"; + +pub fn spec(f: Fam, name: []const u8) c.Spec { + const blocks: c.Tile = .{ .size = 256, .step = 32, .min = 32 }; + return switch (f) { + .pool => .{ .name = name, .accum = .none, .depth = noSums, .out = .bf16, .nan = .{ .stated = pool_rule }, .served = &.{ + .{ .rows = 16, .n = 512, .k = 1 }, .{ .rows = 1, .n = 512, .k = 1 }, .{ .rows = 4, .n = 512, .k = 1 }, .{ .rows = 700, .n = 512, .k = 1 }, + .{ .rows = 5, .n = 64, .k = 1 }, .{ .rows = 3, .n = 33, .k = 1 }, .{ .rows = 9, .n = 1, .k = 1 }, .{ .rows = 8, .n = 100, .k = 1 }, + } }, + .rt => .{ .name = name, .accum = .none, .depth = noSums, .out = .bf16, .nan = .{ .stated = rt_rule }, .tiles = .{ .{}, .{ .size = 256, .step = 16, .min = 16 }, .{} }, .served = &.{ + .{ .rows = 16, .n = 512, .k = 1 }, .{ .rows = 1, .n = 512, .k = 1 }, .{ .rows = 5, .n = 512, .k = 1 }, + .{ .rows = 64, .n = 512, .k = 1 }, .{ .rows = 8, .n = 64, .k = 1 }, .{ .rows = 3, .n = 16, .k = 1 }, + } }, + .pack_keys => .{ .name = name, .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = pack_rule }, .tiles = .{ .{}, blocks, .{} }, .served = &.{ + .{ .rows = 16, .n = 128, .k = 1 }, .{ .rows = 1, .n = 128, .k = 1 }, .{ .rows = 5, .n = 128, .k = 1 }, + .{ .rows = 64, .n = 128, .k = 1 }, .{ .rows = 4, .n = 256, .k = 1 }, .{ .rows = 3, .n = 32, .k = 1 }, + } }, + .pack_q => .{ .name = name, .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = pack_rule }, .tiles = .{ .{}, blocks, .{} }, .served = &.{ + .{ .rows = 16, .n = 8192, .k = 1 }, .{ .rows = 1, .n = 8192, .k = 1 }, .{ .rows = 5, .n = 4096, .k = 1 }, + .{ .rows = 8, .n = 128, .k = 1 }, .{ .rows = 3, .n = 32, .k = 1 }, + } }, + }; +} + +fn pow2(e: i32) f32 { + return std.math.pow(f32, 2, @floatFromInt(e)); +} + +fn spread(rng: std.Random, lo: i32, hi: i32) f32 { + return (rng.float(f32) * 2 - 1) * pow2(rng.intRangeAtMost(i32, lo, hi)); +} + +fn halfUlp(v: f64) f64 { + const a = @abs(v); + if (a < 0x1p-126) return 0x1p-134; + return std.math.ldexp(@as(f64, 1), std.math.frexp(a).exponent - 1 - 8); +} + +fn isNanBits(v: u16) bool { + return (v & 0x7fff) > 0x7f80; +} + +fn at(r: rig_.Ref, extra: usize) cp.Ref { + return .{ .buf = r.buf, .off = r.off + extra }; +} + +/// The E2M1 code nearest to `a` (at least 0), a tie to the even code. +fn nearestCode(a: f64) u8 { + var best: usize = 0; + for (ref.e2m1, 0..) |t, i| { + const d = @abs(a - t); + const bd = @abs(a - ref.e2m1[best]); + if (d < bd or (d == bd and i % 2 == 0 and best % 2 == 1)) best = i; + } + return @intCast(best); +} + +/// One 32-value block's packed bytes from the math: sixteen code bytes and the E8M0 scale byte. +fn packBlock(x: []const u16, codes: *[16]u8) u8 { + var amax: f32 = 0; + for (x) |h| { + const v = c.fromBf16(h); + if (!std.math.isFinite(v)) { + codes.* = @splat(0); + return 255; + } + amax = @max(amax, @abs(v)); + } + const inv6: f32 = 1.0 / 6.0; + const e = ref.ceilLog2(@max(amax, 6 * 0x1p-126) * inv6); + const s = std.math.pow(f64, 2, @floatFromInt(e)); + for (0..16) |j| { + var byte: u8 = 0; + for (0..2) |k| { + const v = c.fromBf16(x[2 * j + k]); + const code = nearestCode(@min(@as(f64, @abs(v)) / s, 6)) | @as(u8, if (std.math.signbit(v)) 8 else 0); + byte |= code << @intCast(4 * k); + } + codes[j] = byte; + } + return @intCast(e + 127); +} + +const vt = g.hooks(Proj); + +pub const Proj = struct { + a: std.mem.Allocator, + arena: std.heap.ArenaAllocator, + rig: rig_.Rig, + fam: Fam, + ctl: u8, + name: []const u8, + pipes: cp.Pipelines, + /// the pipelines of glue_tmp, when the equality probe launches them in place of ours + tmp: ?cp.Pipelines = null, + shape: c.Shape = undefined, + x: []u16 = &.{}, + kv: []f32 = &.{}, + sc: []f32 = &.{}, + pos: []i32 = &.{}, + dest: []i32 = &.{}, + state0: []f32 = &.{}, + plants: usize = 0, + in: [4]rig_.Ref = undefined, + o: [2]rig_.Ref = undefined, + ol: [2]usize = .{ 0, 0 }, + no: usize = 0, + msg: [240]u8 = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, f: Fam, ctl: u8, name: []const u8) !Proj { + var p: Proj = .{ .a = a, .arena = std.heap.ArenaAllocator.init(a), .rig = .{ .a = a, .device = device, .queue = queue }, .fam = f, .ctl = ctl, .name = name, .pipes = undefined }; + if (ctl == 0) { + p.pipes = try cp.Pipelines.load(device); + } else { + const text = try std.fmt.allocPrint(a, "#define TF_CTL {d}\n{s}", .{ ctl, sources.dsv41.compress }); + defer a.free(text); + p.pipes = try cp.Pipelines.loadSource(device, text); + } + return p; + } + + pub fn deinit(p: *Proj) void { + p.arena.deinit(); + p.pipes.deinit(); + p.rig.deinit(); + } + + pub fn kernel(p: *Proj) g.Kernel { + const control: ?g.Gate = switch (p.ctl) { + 0 => null, + 1 => .nonfinite, + 2, 3 => .shapes, + else => .repeat, + }; + return .{ .spec = spec(p.fam, p.name), .ctx = p, .vt = &vt, .control = control }; + } + + pub fn count(p: *Proj) usize { + return p.shape.rows * p.shape.n; + } + + fn dim(p: *const Proj) usize { + return p.shape.n; + } + + fn output(p: *Proj, len: usize) !rig_.Ref { + p.o[p.no] = try p.rig.output(len); + p.ol[p.no] = len; + p.no += 1; + return p.o[p.no - 1]; + } + + /// Rows of destination: a permutation of the rows with every fifth one skipped (-1). + fn destinations(p: *Proj, a: std.mem.Allocator) ![]i32 { + const rows = p.shape.rows; + const d = try a.alloc(i32, rows); + for (d, 0..) |*o, r| o.* = if (r % 5 == 4) -1 else @intCast(rows - 1 - r); + return d; + } + + fn preparePool(p: *Proj, a: std.mem.Allocator, rng: std.Random, nf: bool) !void { + const n = p.dim(); + const rows = p.shape.rows; + p.kv = try a.alloc(f32, rows * n); + p.sc = try a.alloc(f32, rows * n); + p.pos = try a.alloc(i32, rows); + p.state0 = try a.alloc(f32, 4 * n); + var p0: i32 = @intCast(rng.uintLessThan(u32, 1 << 20)); + if (nf) p0 &= ~@as(i32, 1); + for (p.pos, 0..) |*o, r| o.* = p0 + @as(i32, @intCast(r)); + const fresh = !nf and rng.uintLessThan(u32, 3) == 0; + for (p.state0[0 .. 2 * n]) |*o| o.* = if (fresh) 0 else spread(rng, -4, 4); + for (p.state0[2 * n ..]) |*o| o.* = if (fresh) -std.math.inf(f32) else spread(rng, -3, 3); + for (0..rows) |r| { + const mode = rng.uintLessThan(u32, 5); + for (0..n) |d| { + p.kv[r * n + d] = spread(rng, -4, 4); + p.sc[r * n + d] = switch (mode) { + 0, 1 => spread(rng, -3, 3), + 2 => (rng.float(f32) * 2 - 1) * 80, + 3 => if (r > 0) p.sc[(r - 1) * n + d] else 0.5, + else => @as(f32, @floatFromInt(rng.intRangeAtMost(i32, -2, 2))), + }; + } + } + if (nf and rows > 3) { + p.kv[1 * n + n / 3] = std.math.nan(f32); + p.sc[3 * n + n / 2] = std.math.inf(f32); + p.kv[0 * n + 0] = std.math.inf(f32); + p.sc[0 * n + 0] = 0.5; + p.sc[1 * n + 0] = 0.5; + p.kv[1 * n + 0] = 1; + p.kv[2 * n + n - 1] = std.math.inf(f32); + p.kv[3 * n + n - 1] = -std.math.inf(f32); + p.sc[2 * n + n - 1] = 0.25; + p.sc[3 * n + n - 1] = 0.25; + p.plants = 4; + } + p.in[0] = try p.rig.input(std.mem.sliceAsBytes(p.kv)); + p.in[1] = try p.rig.input(std.mem.sliceAsBytes(p.sc)); + p.in[2] = try p.rig.input(std.mem.sliceAsBytes(p.pos)); + _ = try p.output(rows * n * 2); + _ = try p.output(4 * n * 4); + } + + pub fn prepare(p: *Proj, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.n == 0 or shape.n > max_dim) return false; + const blk: usize = switch (p.fam) { + .pool => 1, + .rt => 16, + else => 32, + }; + if (shape.n % blk != 0) return false; + p.rig.reset(); + _ = p.arena.reset(.retain_capacity); + const a = p.arena.allocator(); + p.shape = shape; + p.plants = 0; + p.no = 0; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const nf = fill == .nonfinite; + const n = shape.n; + const rows = shape.rows; + if (p.fam == .pool) { + try p.preparePool(a, rng, nf); + return true; + } + p.x = try a.alloc(u16, rows * n); + ref.fillQuant(if (p.fam == .rt) .fp4_e4m3 else .fp4_e8m0, p.x, blk, seed); + if (nf) p.plants = ref.plant(p.x, rows, n, 0, n); + p.in[0] = try p.rig.input(std.mem.sliceAsBytes(p.x)); + switch (p.fam) { + .rt => { + p.dest = try p.destinations(a); + p.in[1] = try p.rig.input(std.mem.sliceAsBytes(p.dest)); + _ = try p.output(rows * n * 2); + }, + .pack_keys => { + p.dest = try p.destinations(a); + p.in[1] = try p.rig.input(std.mem.sliceAsBytes(p.dest)); + _ = try p.output(rows * (n / 2 + n / 32)); + }, + else => { + _ = try p.output(rows * n / 2); + _ = try p.output(rows * n / 32); + }, + } + return true; + } + + /// Outputs back to the poison and the carry to its start: where a launch sequence begins. + pub fn reset(p: *Proj) void { + for (0..p.no) |i| @memset(p.rig.bytes(p.o[i], p.ol[i]), 0xff); + if (p.fam == .pool) @memcpy(p.rig.bytes(p.o[1], p.ol[1]), std.mem.sliceAsBytes(p.state0)); + } + + /// Rows r0 to r0 + m of the prepared call as one launch, our kernels or (tmp) glue_tmp's. + pub fn encode(p: *Proj, e: anytype, r0: usize, m: usize, tmp: bool) void { + const pl = if (tmp) &p.tmp.? else &p.pipes; + const n = p.shape.n; + const rows: u32 = @intCast(m); + const nn: u32 = @intCast(n); + switch (p.fam) { + .pool => cp.pool2(pl, e, at(p.in[0], r0 * n * 4), at(p.in[1], r0 * n * 4), at(p.o[1], 0), at(p.in[2], r0 * 4), at(p.o[0], r0 * n * 2), rows, nn), + .rt => cp.fp4rt(pl, e, at(p.in[0], r0 * n * 2), at(p.o[0], 0), at(p.in[1], r0 * 4), rows, nn), + .pack_keys => { + const row: u32 = nn / 2 + nn / 32; + cp.fp4pack(pl, e, at(p.in[0], r0 * n * 2), at(p.o[0], 0), at(p.o[0], 0), rows, .{ .blocks = nn / 32, .in_stride = nn, .code_row = row, .scale_row = row, .scale_off = nn / 2, .dest = at(p.in[1], r0 * 4) }); + }, + .pack_q => cp.fp4pack(pl, e, at(p.in[0], r0 * n * 2), at(p.o[0], r0 * n / 2), at(p.o[1], r0 * n / 32), rows, .{ .blocks = nn / 32, .in_stride = nn, .code_row = nn / 2, .scale_row = nn / 32 }), + } + } + + pub fn run(p: *Proj) !void { + p.reset(); + const pass = p.rig.begin(); + p.encode(pass.e, 0, p.shape.rows, false); + try p.rig.finish(pass); + } + + pub fn reference(_: *Proj, want: []f64, mag: []f64) !void { + @memset(want, 0); + @memset(mag, 0); + } + + fn fail(p: *Proj, comptime fmt: []const u8, args: anytype) ![]const u8 { + return try std.fmt.bufPrint(&p.msg, fmt, args); + } + + /// The pool against float64 by column: bounds on pooled rows, even rows untouched, the carry. + fn checkPool(p: *Proj, gate: g.Gate) !?[]const u8 { + const n = p.dim(); + const rows = p.shape.rows; + const y = std.mem.bytesAsSlice(u16, p.rig.bytes(p.o[0], rows * n * 2)); + const st = std.mem.bytesAsSlice(u32, p.rig.bytes(p.o[1], 4 * n * 4)); + var nonfinite: usize = 0; + var pooled: usize = 0; + for (0..n) |d| { + var k2 = [2]f32{ p.state0[d], p.state0[n + d] }; + var s2 = [2]f32{ p.state0[2 * n + d], p.state0[3 * n + d] }; + for (0..rows) |r| { + const slot: usize = @intCast(p.pos[r] & 1); + k2[slot] = p.kv[r * n + d]; + s2[slot] = p.sc[r * n + d]; + const got = y[r * n + d]; + if (slot == 0) { + if (got != 0xffff) return try p.fail("row {d} column {d}: an even position's row was written (0x{x:0>4})", .{ r, d, got }); + continue; + } + const m = @max(@as(f64, s2[0]), @as(f64, s2[1])); + const e0 = @exp(@as(f64, s2[0]) - m); + const e1 = @exp(@as(f64, s2[1]) - m); + const k0: f64 = k2[0]; + const k1: f64 = k2[1]; + const want = (k0 * e0 + k1 * e1) / (e0 + e1); + const mag = (@abs(k0 * e0) + @abs(k1 * e1)) / (e0 + e1); + const have: f64 = c.fromBf16(got); + pooled += 1; + if (!std.math.isFinite(want)) { + nonfinite += 1; + if (c.classOf(have) != c.classOf(want)) return try p.fail("row {d} column {d} of {d}x{d}: {t} where float64 gives {t}", .{ r, d, rows, n, c.classOf(have), c.classOf(want) }); + continue; + } + const work = 10 * 0x1p-24 * mag; + if (!(@abs(have - want) <= halfUlp(@abs(want) + work) + work)) return try p.fail("row {d} column {d} of {d}x{d}: {e:.6} where float64 gives {e:.6}", .{ r, d, rows, n, have, want }); + } + for ([_]f32{ k2[0], k2[1], s2[0], s2[1] }, 0..) |v, i| { + const h: f32 = @bitCast(st[i * n + d]); + if (!(h == v or (std.math.isNan(h) and std.math.isNan(v)))) return try p.fail("carry word {d} of column {d}: {e:.6} where the last slot value is {e:.6}", .{ i, d, h, v }); + } + } + if (gate == .nonfinite and (p.plants < 4 or nonfinite < 3 or nonfinite == pooled)) return try p.fail("{d} plants, {d} of {d} pooled outputs non-finite: no proof", .{ p.plants, nonfinite, pooled }); + return null; + } + + fn checkRt(p: *Proj, gate: g.Gate) !?[]const u8 { + const n = p.dim(); + const rows = p.shape.rows; + const got = std.mem.bytesAsSlice(u16, p.rig.bytes(p.o[0], rows * n * 2)); + const want = try p.a.alloc(u16, rows * n); + defer p.a.free(want); + const written = try p.a.alloc(bool, rows * n); + defer p.a.free(written); + @memset(want, 0xffff); + @memset(written, false); + var xs: [16]f32 = undefined; + var bad_blocks: usize = 0; + for (0..rows) |r| { + if (p.dest[r] < 0) continue; + const to: usize = @intCast(p.dest[r]); + for (0..n / 16) |b| { + for (&xs, p.x[r * n + 16 * b ..][0..16]) |*o, v| o.* = c.fromBf16(v); + ref.roundTrip(.fp4_e4m3, &xs, want[to * n + 16 * b ..][0..16]); + @memset(written[to * n + 16 * b ..][0..16], true); + bad_blocks += @intFromBool(isNanBits(want[to * n + 16 * b])); + } + } + for (want, got, written, 0..) |w, h, wr, i| { + const ok = if (wr) (w == h or (isNanBits(w) and isNanBits(h))) else h == 0xffff; + if (!ok) return try p.fail("row {d} element {d} of {d}x{d}: got 0x{x:0>4}, the round trip gives 0x{x:0>4}", .{ i / n, i % n, rows, n, h, w }); + } + if (gate == .nonfinite and (p.plants < 5 or bad_blocks < 4)) return try p.fail("{d} plants, {d} NaN blocks: no proof", .{ p.plants, bad_blocks }); + return null; + } + + fn checkPack(p: *Proj, gate: g.Gate) !?[]const u8 { + const n = p.dim(); + const rows = p.shape.rows; + const keys = p.fam == .pack_keys; + const code_row = if (keys) n / 2 + n / 32 else n / 2; + const scale_row = if (keys) code_row else n / 32; + const scale_off = if (keys) n / 2 else 0; + const cw = try p.a.alloc(u8, rows * code_row); + defer p.a.free(cw); + const sw = if (keys) cw else try p.a.alloc(u8, rows * scale_row); + defer if (!keys) p.a.free(sw); + @memset(cw, 0xff); + @memset(sw, 0xff); + var bad_blocks: usize = 0; + for (0..rows) |r| { + const to: usize = if (keys) (if (p.dest[r] < 0) continue else @intCast(p.dest[r])) else r; + for (0..n / 32) |b| { + const sb = packBlock(p.x[r * n + 32 * b ..][0..32], cw[to * code_row + 16 * b ..][0..16]); + sw[to * scale_row + scale_off + b] = sb; + bad_blocks += @intFromBool(sb == 255); + } + } + const got_c = p.rig.bytes(p.o[0], rows * code_row); + if (std.mem.indexOfDiff(u8, cw, got_c)) |i| return try p.fail("{s} byte {d} of {d}x{d}: got 0x{x:0>2}, the packing gives 0x{x:0>2}", .{ if (keys) "key" else "code", i, rows, n, got_c[i], cw[i] }); + if (!keys) { + const got_s = p.rig.bytes(p.o[1], rows * scale_row); + if (std.mem.indexOfDiff(u8, sw, got_s)) |i| return try p.fail("scale byte {d} of {d}x{d}: got 0x{x:0>2}, the packing gives 0x{x:0>2}", .{ i, rows, n, got_s[i], sw[i] }); + } + if (gate == .nonfinite and (p.plants < 5 or bad_blocks < 4)) return try p.fail("{d} plants, {d} NaN blocks: no proof", .{ p.plants, bad_blocks }); + return null; + } + + pub fn check(p: *Proj, gate: g.Gate) !?[]const u8 { + return switch (p.fam) { + .pool => p.checkPool(gate), + .rt => p.checkRt(gate), + else => p.checkPack(gate), + }; + } +}; + +/// Every registered variant, the launch-split probe after the gates and the equality with glue_tmp's kernels. +pub const Suite = struct { + a: std.mem.Allocator, + device: mtl.Device, + probes: std.ArrayList(*Proj) = .empty, + + const E = struct { f: Fam, ctl: u8, name: []const u8 }; + const pool_name = base ++ "pool tf_dsv41_c_pool2"; + const rt_name = base ++ "fp4 round trip tf_dsv41_c_fp4rt"; + const keys_name = base ++ "fp4 pack keys tf_dsv41_c_fp4pack"; + const list = [_]E{ + .{ .f = .pool, .ctl = 0, .name = pool_name }, + .{ .f = .pool, .ctl = 1, .name = base ++ "control: pool non-finite inputs flushed" }, + .{ .f = .pool, .ctl = 2, .name = base ++ "control: pool carry not written back" }, + .{ .f = .pool, .ctl = 3, .name = base ++ "control: pool slot taken from the row, not the position" }, + .{ .f = .pool, .ctl = 4, .name = base ++ "control: pool writes past its output" }, + .{ .f = .rt, .ctl = 0, .name = rt_name }, + .{ .f = .rt, .ctl = 1, .name = base ++ "control: fp4 round trip NaN not flagged" }, + .{ .f = .rt, .ctl = 2, .name = base ++ "control: fp4 round trip scale left unrounded to E4M3" }, + .{ .f = .rt, .ctl = 4, .name = base ++ "control: fp4 round trip writes past its output" }, + .{ .f = .pack_keys, .ctl = 0, .name = keys_name }, + .{ .f = .pack_keys, .ctl = 1, .name = base ++ "control: fp4 pack NaN not flagged" }, + .{ .f = .pack_keys, .ctl = 2, .name = base ++ "control: fp4 pack scale exponent floored" }, + .{ .f = .pack_keys, .ctl = 3, .name = base ++ "control: fp4 pack ties rounded away from zero" }, + .{ .f = .pack_keys, .ctl = 4, .name = base ++ "control: fp4 pack writes past its output" }, + .{ .f = .pack_q, .ctl = 0, .name = base ++ "fp4 pack queries tf_dsv41_c_fp4pack" }, + }; + + pub fn create(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, kernels: *std.ArrayList(g.Kernel), only: ?[]const u8) !*Suite { + const s = try a.create(Suite); + s.* = .{ .a = a, .device = device }; + for (list) |entry| { + if (only) |text| if (std.mem.indexOf(u8, entry.name, text) == null) continue; + const p = try a.create(Proj); + p.* = try Proj.init(a, device, queue, entry.f, entry.ctl, entry.name); + try s.probes.append(a, p); + try kernels.append(a, p.kernel()); + } + return s; + } + + pub fn destroy(s: *Suite) void { + for (s.probes.items) |p| { + p.deinit(); + s.a.destroy(p); + } + s.probes.deinit(s.a); + s.a.destroy(s); + } + + fn shapesOf(f: Fam) []const c.Shape { + return switch (f) { + .pool => &.{ .{ .rows = 33, .n = 512, .k = 1 }, .{ .rows = 200, .n = 100, .k = 1 }, .{ .rows = 9, .n = 33, .k = 1 } }, + .rt => &.{ .{ .rows = 33, .n = 512, .k = 1 }, .{ .rows = 9, .n = 64, .k = 1 } }, + .pack_keys => &.{ .{ .rows = 33, .n = 128, .k = 1 }, .{ .rows = 9, .n = 256, .k = 1 } }, + .pack_q => &.{ .{ .rows = 20, .n = 8192, .k = 1 }, .{ .rows = 9, .n = 512, .k = 1 } }, + }; + } + + /// Launches of 1, 7 and rows/3 rows in order (the carry crossing every launch) equal the whole launch's bytes. + fn probe(s: *Suite, p: *Proj, shape: c.Shape, seed: u64, bad: *usize, launches: *usize) !void { + if (!try p.prepare(shape, .random, seed)) return error.Refused; + try p.run(); + const whole = try s.a.dupe(u8, try p.rig.outputBytes()); + defer s.a.free(whole); + for ([_]usize{ 1, 7, @max(1, shape.rows / 3) }) |step| { + p.reset(); + const pass = p.rig.begin(); + var r0: usize = 0; + while (r0 < shape.rows) : (r0 += step) { + p.encode(pass.e, r0, @min(step, shape.rows - r0), false); + launches.* += 1; + } + try p.rig.finish(pass); + if (!std.mem.eql(u8, whole, try p.rig.outputBytes())) { + bad.* += 1; + std.debug.print("FAIL row invariance {s} {d}x{d}: {d} rows a launch differs from one launch\n", .{ p.name, shape.rows, shape.n, step }); + } + } + } + + pub fn invariance(s: *Suite, seed: u64) !usize { + var bad: usize = 0; + var launches: usize = 0; + var kernels: usize = 0; + for (s.probes.items) |p| { + if (p.ctl != 0) continue; + kernels += 1; + for (shapesOf(p.fam), 0..) |shape, i| try s.probe(p, shape, seed + i, &bad, &launches); + } + if (kernels == 0) return 0; + std.debug.print("{s} row invariance {s}: {d} kernels at 1, 7 and rows/3 rows a launch, the pool's carry across launches ({d} launches) byte-equal to one launch, {d} differ\n", .{ if (bad == 0) "PASS" else "FAIL", base, kernels, launches, bad }); + return bad; + } + + /// Our kernels and glue_tmp's give the same bytes on finite inputs, outputs and carry alike. + pub fn equality(s: *Suite, seed: u64) !usize { + var any = false; + for (s.probes.items) |p| any = any or p.ctl == 0; + if (!any) return 0; + const lib = try mtl.Library.fromSource(s.device, sources.dsv41.glue_tmp, mtl.CompileOptions.mlx()); + defer lib.deinit(); + var old: cp.Pipelines = .{ + .pool = try mtl.Pipeline.init(s.device, lib, "tf_dsv41_g_pool2", false), + .rt = try mtl.Pipeline.init(s.device, lib, "tf_dsv41_g_fp4rt", false), + .pack = try mtl.Pipeline.init(s.device, lib, "tf_dsv41_g_fp4pack", false), + }; + defer old.deinit(); + var cases: usize = 0; + var differ: usize = 0; + for (s.probes.items) |p| { + if (p.ctl != 0) continue; + p.tmp = old; + defer p.tmp = null; + const extra = [_]c.Shape{ .{ .rows = 700, .n = 512, .k = 1 }, .{ .rows = 64, .n = 8192, .k = 1 } }; + const lists = [_][]const c.Shape{ spec(p.fam, p.name).served, shapesOf(p.fam), &extra }; + var i: usize = 0; + for (lists) |group| for (group) |shape| { + i += 1; + for (0..3) |k| { + if (!try p.prepare(shape, .random, seed + 100 * i + k)) continue; + p.reset(); + var pass = p.rig.begin(); + p.encode(pass.e, 0, shape.rows, false); + try p.rig.finish(pass); + const mine = try s.a.dupe(u8, try p.rig.outputBytes()); + defer s.a.free(mine); + p.reset(); + pass = p.rig.begin(); + p.encode(pass.e, 0, shape.rows, true); + try p.rig.finish(pass); + cases += 1; + if (std.mem.indexOfDiff(u8, mine, try p.rig.outputBytes())) |at_byte| { + differ += 1; + std.debug.print("FAIL glue_tmp equality {s} {d}x{d}: first difference at byte {d}\n", .{ p.name, shape.rows, shape.n, at_byte }); + } + } + }; + } + std.debug.print("{s} equality {s}: {d} cases of the pool, the round trip and both packings byte-equal to glue_tmp's kernels on finite inputs, {d} differ\n", .{ if (differ == 0) "PASS" else "FAIL", base, cases, differ }); + return differ; + } +}; diff --git a/zig/tests/contract/metal_dsv41_engram.zig b/zig/tests/contract/metal_dsv41_engram.zig new file mode 100644 index 0000000000..5c6815e190 --- /dev/null +++ b/zig/tests/contract/metal_dsv41_engram.zig @@ -0,0 +1,435 @@ +//! DeepSeek V4.1's Engram under the contract on Metal: four stages against float64, controls, whole-step invariance. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const sources = @import("kernel_sources"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const ref = @import("dsv41_engram_ref.zig"); + +const c = contract.spec; +const g = contract.gates; +const eg = tf.dsv41_metal.engram; +const base = "dsv41 engram "; + +/// What a registration varies: the stage and a planted fault (TF_CTL in the source). +pub const Variant = struct { stage: ref.Stage, ctl: u8 = 0 }; + +fn s3(rows: usize, n: usize, k: usize) c.Shape { + return .{ .rows = rows, .n = n, .k = k }; +} + +const rowsK = blk: { + var out: [17]usize = undefined; + for (0..16) |i| out[i] = i + 1; + out[16] = 33; + break :blk out; +}; + +fn served(comptime stage: ref.Stage) []const c.Shape { + const list = comptime blk: { + var out: [rowsK.len + 8]c.Shape = undefined; + var n: usize = 0; + switch (stage) { + .embed => { + out[n] = s3(8, 6144, 6144); + n += 1; + for (rowsK) |r| { + out[n] = s3(r, 6144, 6144); + n += 1; + } + out[n] = s3(3, 512, 512); + n += 1; + }, + .aq => { + out[n] = s3(8, 6144, 6144); + n += 1; + for (rowsK) |r| { + out[n] = s3(r, 6144, 6144); + n += 1; + } + out[n] = s3(3, 512, 512); + n += 1; + }, + .wkv => { + out[n] = s3(8, 192, 6144); + n += 1; + for (rowsK) |r| { + out[n] = s3(r, 64, 1536); + n += 1; + } + for ([_]c.Shape{ s3(1, 512, 6144), s3(5, 96, 2560), s3(16, 128, 256), s3(2, 32, 6144), s3(4, 5120, 6144) }) |x| { + out[n] = x; + n += 1; + } + }, + .gate => { + out[n] = s3(8, 4, 5120); + n += 1; + for (rowsK) |r| { + out[n] = s3(r, 4, 5120); + n += 1; + } + out[n] = s3(2, 4, 32); + n += 1; + out[n] = s3(5, 4, 1056); + n += 1; + }, + .add => { + out[n] = s3(8, 20480, 5120); + n += 1; + for (rowsK) |r| { + out[n] = s3(r, 20480, 5120); + n += 1; + } + out[n] = s3(2, 128, 32); + n += 1; + }, + } + var fin: [n]c.Shape = undefined; + @memcpy(&fin, out[0..n]); + break :blk fin; + }; + return &list; +} + +fn depthOne(_: c.Shape) usize { + return 1; +} + +fn depthGate(s: c.Shape) usize { + return s.k / 32 + 8; +} + +fn depthAdd(_: c.Shape) usize { + return 2; +} + +pub fn spec(v: Variant, name: []const u8) c.Spec { + return switch (v.stage) { + .embed => .{ .name = name, .accum = .none, .depth = depthOne, .out = .bf16, .served = served(.embed) }, + .aq => .{ .name = name, .accum = .none, .depth = depthOne, .out = .bf16, .served = served(.aq) }, + .wkv => .{ .name = name, .accum = .none, .depth = depthOne, .out = .bf16, .served = served(.wkv) }, + .gate => .{ + .name = name, + .accum = .f32, + .depth = depthGate, + .out = .f32, + .tiles = .{ .{ .size = 16 }, .{}, .{ .size = 32, .step = 32, .min = 32 } }, + .served = served(.gate), + }, + .add => .{ .name = name, .accum = .none, .depth = depthAdd, .out = .bf16, .served = served(.add) }, + }; +} + +pub const Proj = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + v: Variant, + name: []const u8, + pipes: eg.Pipelines, + d: ?ref.Data = null, + ins: [5]rig_.Ref = undefined, + scr: rig_.Ref = undefined, + out: rig_.Ref = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, v: Variant, name: []const u8) !Proj { + var p: Proj = .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .v = v, .name = name, .pipes = undefined }; + if (v.ctl == 0) { + p.pipes = try eg.Pipelines.load(device); + } else { + const text = try std.fmt.allocPrint(a, "#define TF_CTL {d}\n{s}", .{ v.ctl, sources.dsv41.engram }); + defer a.free(text); + p.pipes = try eg.Pipelines.loadSource(device, text); + } + return p; + } + + pub fn deinit(p: *Proj) void { + p.drop(); + p.pipes.deinit(); + p.rig.deinit(); + } + + fn drop(p: *Proj) void { + if (p.d) |*d| d.deinit(); + p.d = null; + p.rig.reset(); + } + + pub fn kernel(p: *Proj) g.Kernel { + const control: ?g.Gate = switch (p.v.ctl) { + 1, 7 => .shapes, + 5 => .sums, + 4 => .nonfinite, + 6 => .repeat, + else => null, + }; + return .{ .spec = spec(p.v, p.name), .ctx = p, .vt = &vt, .control = control }; + } + + const vt = g.hooks(Proj); + + fn call(p: *const Proj) eg.Call { + const s = p.d.?.shape; + const rows: u32 = @intCast(s.rows); + return switch (p.v.stage) { + .embed, .aq => .{ .rows = rows, .cols = @intCast(s.k / ref.head), .n = 32, .dim = 32 }, + .wkv => .{ .rows = rows, .cols = @intCast(s.k / ref.head), .n = @intCast(s.n), .dim = 32 }, + .gate, .add => .{ .rows = rows, .cols = 1, .n = 32, .dim = @intCast(s.k) }, + }; + } + + pub fn prepare(p: *Proj, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + p.drop(); + const wide = p.v.stage != .gate and p.v.stage != .add; + if (shape.k % 32 != 0 or (wide and shape.k % ref.head != 0) or shape.rows == 0 or shape.n == 0 or (p.v.stage == .wkv and shape.n % 32 != 0)) return false; + var d = try ref.Data.init(p.a, p.v.stage, shape); + d.fill(fill, seed); + p.d = d; + try p.upload(); + return true; + } + + fn upload(p: *Proj) !void { + const d = p.d.?; + const bytes = d.outputs() * c.size(spec(p.v, p.name).out); + const b = std.mem.sliceAsBytes; + switch (p.v.stage) { + .embed => { + p.ins[0] = try p.rig.input(d.codes); + p.ins[1] = try p.rig.input(d.scales); + }, + .aq => p.ins[0] = try p.rig.input(b(d.emb)), + .wkv => { + p.ins[0] = try p.rig.input(b(d.emb)); + p.ins[1] = try p.rig.input(d.f.?.codes); + p.ins[2] = try p.rig.input(d.f.?.scales); + const c0 = p.call(); + p.scr = try p.rig.scratch(2 * @as(usize, c0.rows) * d.shape.k + eg.fp8Scratch(c0)); + }, + .gate => { + p.ins[0] = try p.rig.input(b(d.x)); + p.ins[1] = try p.rig.input(b(d.kv)); + p.ins[2] = try p.rig.input(b(d.q)); + p.ins[3] = try p.rig.input(b(d.kw)); + p.ins[4] = try p.rig.input(d.mask); + }, + .add => { + p.ins[0] = try p.rig.input(b(d.x)); + p.ins[1] = try p.rig.input(b(d.kv)); + p.ins[2] = try p.rig.input(b(d.g)); + }, + } + p.out = try p.rig.output(bytes); + } + + fn at(r: rig_.Ref) eg.Ref { + return .{ .buf = r.buf, .off = r.off }; + } + + pub fn run(p: *Proj) !void { + const pass = p.rig.begin(); + const e = pass.e; + const i = p.ins; + switch (p.v.stage) { + .embed => try eg.encodeEmbed(&p.pipes, e, at(i[0]), at(i[1]), at(p.out), p.call()), + .aq => try eg.encodeAq(&p.pipes, e, at(i[0]), at(p.out), p.call()), + .wkv => { + const cl = p.call(); + const roots = eg.Ref{ .buf = p.scr.buf, .off = p.scr.off + 2 * cl.rows * cl.cols * ref.head }; + try eg.encodeAq(&p.pipes, e, at(i[0]), at(p.scr), cl); + try eg.encodeWkv(&p.pipes, e, at(p.scr), at(i[1]), at(i[2]), roots, at(p.out), cl); + }, + .gate => try eg.encodeGate(&p.pipes, e, at(i[0]), at(i[1]), at(i[2]), at(i[3]), at(i[4]), at(p.out), p.call()), + .add => try eg.encodeAdd(&p.pipes, e, at(i[0]), at(i[1]), at(i[2]), at(p.out), p.call()), + } + try p.rig.finish(pass); + } + + pub fn count(p: *Proj) usize { + return p.d.?.outputs(); + } + + pub fn reference(p: *Proj, want: []f64, mag: []f64) !void { + return p.d.?.reference(want, mag); + } + + pub fn bounds(p: *Proj, want: []const f64, mag: []const f64, allow: []f64) !void { + return p.d.?.bounds(want, mag, allow); + } + + pub fn simulate(p: *Proj, acc: c.Accum, order: g.Order, out: []f64) !void { + return p.d.?.simulate(acc, order, out); + } +}; + +/// Every registered variant, and the whole-step probes that run after the gates. +pub const Suite = struct { + a: std.mem.Allocator, + device: mtl.Device, + queue: mtl.Queue, + probes: std.ArrayList(*Proj) = .empty, + + const list = [_]struct { v: Variant, name: []const u8 }{ + .{ .v = .{ .stage = .embed }, .name = base ++ "embed" }, + .{ .v = .{ .stage = .aq }, .name = base ++ "aq" }, + .{ .v = .{ .stage = .wkv }, .name = base ++ "wkv (act_quant, fp8_rows)" }, + .{ .v = .{ .stage = .gate }, .name = base ++ "gate" }, + .{ .v = .{ .stage = .add }, .name = base ++ "add" }, + .{ .v = .{ .stage = .embed, .ctl = 1 }, .name = base ++ "control: embed reads one scale for 64 codes" }, + .{ .v = .{ .stage = .aq, .ctl = 7 }, .name = base ++ "control: act_quant truncates its codes" }, + .{ .v = .{ .stage = .gate, .ctl = 4 }, .name = base ++ "control: gate clamp drops a NaN" }, + .{ .v = .{ .stage = .gate, .ctl = 5 }, .name = base ++ "control: gate sums rounded to fp16" }, + .{ .v = .{ .stage = .add, .ctl = 6 }, .name = base ++ "control: add writes past its output" }, + }; + + pub fn create(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, kernels: *std.ArrayList(g.Kernel), only: ?[]const u8) !*Suite { + const s = try a.create(Suite); + s.* = .{ .a = a, .device = device, .queue = queue }; + for (list) |entry| { + if (only) |text| if (std.mem.indexOf(u8, entry.name, text) == null) continue; + const p = try a.create(Proj); + p.* = try Proj.init(a, device, queue, entry.v, entry.name); + try s.probes.append(a, p); + try kernels.append(a, p.kernel()); + } + return s; + } + + pub fn destroy(s: *Suite) void { + for (s.probes.items) |p| { + p.deinit(); + s.a.destroy(p); + } + s.probes.deinit(s.a); + s.a.destroy(s); + } + + /// The whole step on 40 rows (some masked): any chunk or a launch a row gives the same bytes, also add in place. + pub fn invariance(s: *Suite, seed: u64) !usize { + if (s.probes.items.len == 0 or s.probes.items[0].v.ctl != 0) return 0; + const pipes = &s.probes.items[0].pipes; + const cols = 24; + const dim = 640; + const k = cols * ref.head; + const n = (ref.hc + 1) * dim; + const total = 40; + const xrow = ref.hc * dim; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const codes = try s.a.alloc(u8, total * k); + defer s.a.free(codes); + for (codes) |*q| { + const b = rng.int(u8); + q.* = if (b & 0x7f == 0x7f) b & 0xfe else b; + } + const scales = try s.a.alloc(u8, total * k / 32); + defer s.a.free(scales); + for (scales) |*q| q.* = 119 + rng.uintLessThan(u8, 7); + const w8 = try s.a.alloc(u8, n * k); + defer s.a.free(w8); + for (w8) |*q| { + const b = rng.int(u8); + q.* = if (b & 0x7f == 0x7f) b & 0xfe else b; + } + const ws = try s.a.alloc(u8, n / 32 * (k / 32)); + defer s.a.free(ws); + for (ws) |*q| q.* = 108 + rng.uintLessThan(u8, 7); + const x = try s.a.alloc(u16, total * xrow); + defer s.a.free(x); + for (x) |*v| v.* = c.toBf16(rng.float(f32) * 4 - 2); + const qk = try s.a.alloc(u16, 2 * xrow); + defer s.a.free(qk); + for (qk) |*w| w.* = c.toBf16(rng.float(f32) * 2 - 1); + var mask: [total]u8 = undefined; + for (&mask, 0..) |*m, r| m.* = @intFromBool(r % 5 != 3); + const ones: [total]u8 = @splat(1); + + var rig: rig_.Rig = .{ .a = s.a, .device = s.device, .queue = s.queue }; + defer rig.deinit(); + const b = std.mem.sliceAsBytes; + const rc = try rig.input(codes); + const rs = try rig.input(scales); + const rw = try rig.input(w8); + const rws = try rig.input(ws); + const rx = try rig.input(b(x)); + const rq = try rig.input(b(qk[0..xrow])); + const rk = try rig.input(b(qk[xrow..])); + const rm = try rig.input(&mask); + const r1 = try rig.input(&ones); + const rxio = try rig.input(b(x)); + const ry = try rig.output(2 * total * xrow); + const call: eg.Call = .{ .rows = total, .cols = cols, .n = n, .dim = dim }; + const rscratch = try rig.scratch(eg.scratchBytes(call)); + const whole = try s.a.alloc(u8, 2 * total * xrow); + defer s.a.free(whole); + const Run = struct { + fn go(pp: *const eg.Pipelines, rg: *rig_.Rig, from: usize, rows: u32, m: rig_.Ref, y: rig_.Ref, xin: rig_.Ref, ins: [7]rig_.Ref, sc: rig_.Ref) !void { + const pass = rg.begin(); + const cl: eg.Call = .{ .rows = rows, .cols = cols, .n = n, .dim = dim }; + const o = struct { + fn f(r: rig_.Ref, extra: usize) eg.Ref { + return .{ .buf = r.buf, .off = r.off + extra }; + } + }.f; + try eg.encode(pp, pass.e, o(ins[0], from * k), o(ins[1], from * k / 32), o(ins[2], 0), o(ins[3], 0), o(xin, 2 * from * xrow), o(ins[4], 0), o(ins[5], 0), o(m, from), o(sc, 0), o(y, 2 * from * xrow), cl); + try rg.finish(pass); + } + }; + const ins = [7]rig_.Ref{ rc, rs, rw, rws, rq, rk, rx }; + var bad: usize = 0; + var launches: usize = 0; + @memset(rig.bytes(ry, whole.len), 0xff); + try Run.go(pipes, &rig, 0, total, rm, ry, rx, ins, rscratch); + @memcpy(whole, rig.bytes(ry, whole.len)); + for ([_]u32{ 1, 2, 3, 4, 8, 15, 16, 17, 33 }) |r| { + @memset(rig.bytes(ry, whole.len), 0xff); + try Run.go(pipes, &rig, 0, r, rm, ry, rx, ins, rscratch); + launches += 1; + if (!std.mem.eql(u8, whole[0 .. 2 * r * xrow], rig.bytes(ry, 2 * r * xrow))) { + bad += 1; + std.debug.print("FAIL row invariance {s}: {d} rows differ from the {d}-row launch\n", .{ base, r, total }); + } + } + for (0..total) |r| { + @memset(rig.bytes(ry, whole.len), 0xff); + try Run.go(pipes, &rig, r, 1, rm, ry, rx, ins, rscratch); + launches += 1; + const got = rig.bytes(ry, 2 * (r + 1) * xrow)[2 * r * xrow ..]; + if (!std.mem.eql(u8, whole[2 * r * xrow ..][0 .. 2 * xrow], got)) { + bad += 1; + std.debug.print("FAIL row invariance {s}: row {d} alone differs from the {d}-row launch\n", .{ base, r, total }); + } + } + // In place: the add reads each element and writes it from one thread. + try Run.go(pipes, &rig, 0, total, rm, rxio, rxio, ins, rscratch); + launches += 1; + if (!std.mem.eql(u8, whole, rig.bytes(rxio, whole.len))) { + bad += 1; + std.debug.print("FAIL in place {s}: y == x differs from the out-of-place bytes\n", .{base}); + } + // Image rows: a masked row's streams come out as they went in, and no row depends on another's mask. + var through: usize = 0; + for (0..total) |r| if (mask[r] == 0) { + through += 1; + if (!std.mem.eql(u8, b(x)[2 * r * xrow ..][0 .. 2 * xrow], whole[2 * r * xrow ..][0 .. 2 * xrow])) { + bad += 1; + std.debug.print("FAIL image row {s}: masked row {d} changed\n", .{ base, r }); + } + }; + @memset(rig.bytes(ry, whole.len), 0xff); + try Run.go(pipes, &rig, 0, total, r1, ry, rx, ins, rscratch); + launches += 1; + for (0..total) |r| { + const same = std.mem.eql(u8, whole[2 * r * xrow ..][0 .. 2 * xrow], rig.bytes(ry, whole.len)[2 * r * xrow ..][0 .. 2 * xrow]); + if (mask[r] != 0 and !same) { + bad += 1; + std.debug.print("FAIL image row {s}: unmasked row {d} depends on the mask of the others\n", .{ base, r }); + } + } + const word = if (bad == 0) "PASS" else "FAIL"; + std.debug.print("{s} row invariance {s}: {d} launches ({d}-row, chunks 1..33, one a row, in place, {d} masked rows pass through), {d} differ\n", .{ word, base, launches, total, through, bad }); + return bad; + } +}; diff --git a/zig/tests/contract/metal_dsv41_limits.zig b/zig/tests/contract/metal_dsv41_limits.zig index 5fdd7eb0f7..69da8990c9 100644 --- a/zig/tests/contract/metal_dsv41_limits.zig +++ b/zig/tests/contract/metal_dsv41_limits.zig @@ -27,9 +27,11 @@ pub fn launchOf(a: std.mem.Allocator, device: mtl.Device, d: limits.MetalLaunch, const body = try std.fmt.allocPrint(a, "#define TF_DSV41_E {d}\n#define TF_DSV41_TOPK {d}\n{s}", .{ e, k, sources.dsv41.route }); return .{ .decl = d, .function = fun, .here = body, .pre_m5 = body }; } + if (std.mem.startsWith(u8, d.name, "dsv41 engram ")) return .{ .decl = d, .function = fun, .here = sources.dsv41.engram, .pre_m5 = sources.dsv41.engram }; if (std.mem.startsWith(u8, d.name, "dsv41 glue_tmp ")) return .{ .decl = d, .function = fun, .here = sources.dsv41.glue_tmp, .pre_m5 = sources.dsv41.glue_tmp }; if (std.mem.startsWith(u8, d.name, "dsv41 glue ")) return .{ .decl = d, .function = fun, .here = sources.dsv41.glue, .pre_m5 = sources.dsv41.glue }; if (std.mem.startsWith(u8, d.name, "dsv41 cand ")) return .{ .decl = d, .function = fun, .here = sources.dsv41.cand, .pre_m5 = sources.dsv41.cand }; + if (std.mem.startsWith(u8, d.name, "dsv41 compress ")) return .{ .decl = d, .function = fun, .here = sources.dsv41.compress, .pre_m5 = sources.dsv41.compress }; if (std.mem.startsWith(u8, d.name, "dsv41 mv ")) return .{ .decl = d, .function = fun, .here = sources.dsv41.mv, .pre_m5 = sources.dsv41.mv }; if (std.mem.startsWith(u8, d.name, "dsv41 fp8 rows ")) return .{ .decl = d, .function = fun, .here = sources.dsv41.fp8_rows, .pre_m5 = sources.dsv41.fp8_rows }; return null; diff --git a/zig/tests/contract/metal_dsv41_main.zig b/zig/tests/contract/metal_dsv41_main.zig index 9137675d21..029a148b46 100644 --- a/zig/tests/contract/metal_dsv41_main.zig +++ b/zig/tests/contract/metal_dsv41_main.zig @@ -13,7 +13,9 @@ const dsv41_attn = @import("metal_dsv41_attn.zig"); const dsv41_experts = @import("metal_dsv41_experts.zig"); const dsv41_xmma = @import("metal_dsv41_xmma.zig"); const dsv41_index = @import("metal_dsv41_index.zig"); +const dsv41_engram = @import("metal_dsv41_engram.zig"); const dsv41_glue = @import("metal_dsv41_glue.zig"); +const dsv41_compress = @import("metal_dsv41_compress.zig"); const dsv41_mv = @import("metal_dsv41_mv.zig"); const dsv41_glue_tmp = @import("metal_dsv41_glue_tmp.zig"); const dsv41_glue_tmp2 = @import("metal_dsv41_glue_tmp2.zig"); @@ -72,7 +74,9 @@ pub const Suites = struct { experts: *dsv41_experts.Suite, xmma: *dsv41_xmma.Suite, ix: *dsv41_index.Suite, + engram: *dsv41_engram.Suite, glue: *dsv41_glue.Suite, + compress: *dsv41_compress.Suite, mv: *dsv41_mv.Suite, glue_tmp: *dsv41_glue_tmp.Suite, glue_tmp2: *dsv41_glue_tmp2.Suite, @@ -94,22 +98,28 @@ pub const Suites = struct { errdefer xmma.destroy(); const ix = try dsv41_index.Suite.create(a, device, queue, kernels, only); errdefer ix.destroy(); + const engram = try dsv41_engram.Suite.create(a, device, queue, kernels, only); + errdefer engram.destroy(); const glue = try dsv41_glue.Suite.create(a, device, queue, kernels, only); errdefer glue.destroy(); + const compress = try dsv41_compress.Suite.create(a, device, queue, kernels, only); + errdefer compress.destroy(); const mv = try dsv41_mv.Suite.create(a, device, queue, kernels, only); errdefer mv.destroy(); const glue_tmp = try dsv41_glue_tmp.Suite.create(a, device, queue, kernels, only); errdefer glue_tmp.destroy(); const glue_tmp2 = try dsv41_glue_tmp2.Suite.create(a, device, queue, kernels, only); errdefer glue_tmp2.destroy(); - return .{ .fp8 = fp8, .fp8_mma = fp8_mma, .topk = topk, .cand = cand, .attn = attn, .experts = experts, .xmma = xmma, .ix = ix, .glue = glue, .mv = mv, .glue_tmp = glue_tmp, .glue_tmp2 = glue_tmp2 }; + return .{ .fp8 = fp8, .fp8_mma = fp8_mma, .topk = topk, .cand = cand, .attn = attn, .experts = experts, .xmma = xmma, .ix = ix, .engram = engram, .glue = glue, .compress = compress, .mv = mv, .glue_tmp = glue_tmp, .glue_tmp2 = glue_tmp2 }; } pub fn destroy(s: *Suites) void { s.glue_tmp2.destroy(); s.glue_tmp.destroy(); s.mv.destroy(); + s.compress.destroy(); s.glue.destroy(); + s.engram.destroy(); s.ix.destroy(); s.xmma.destroy(); s.experts.destroy(); @@ -139,6 +149,7 @@ pub const Suites = struct { try s.attn.invariance(io, seed, t); if (std.c.getenv("TF_DSV41_TIME") != null) try s.attn.timing(io); t.failed += try s.ix.invariance(seed); + t.failed += try s.engram.invariance(seed); t.failed += try s.glue.invariance(seed); if (s.mv.probes.items.len > 0) { t.failed += try s.mv.invariance(seed); @@ -146,6 +157,8 @@ pub const Suites = struct { t.failed += try s.mv.head(io, seed, std.c.getenv("TF_DSV41_TIME") != null); } t.failed += try s.glue.tables(); + t.failed += try s.compress.invariance(seed); + t.failed += try s.compress.equality(seed); t.failed += try s.glue_tmp.invariance(seed); t.failed += try s.glue_tmp.bounds(seed); t.failed += try s.glue_tmp2.invariance(seed);