dsv41 metal: the KV compressor and Engram

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 <metal_stdlib>
+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;
+  }
+}

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 <metal_stdlib>
+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<float>(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<uint>(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<uint>(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<half>(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);
+}

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 <metal_stdlib>
+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<half>(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<float>(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<uint>(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<uint>(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<float>(uint(127 - e) << 23), down = as_type<float>(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<half>(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);
+}