From 08203429d91fd298064ae33dfd998f6a772c24df Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Sun, 18 Oct 2026 17:16:21 -0700 Subject: [PATCH 1/7] DeepSeek V4.1: the FP32 and FP64 oracles run side by side The reference modules take a comptime real width, so the FP32 source contract and the FP64 ideal are one equation set instantiated twice and cannot drift apart on their own. The new oracle-precision mode runs both on identical checkpoint bytes and tokens with the GPU out of the loop, keeps both logit vectors, and reports the NRMS between the widths, the per-token argmax, and the residency receipt of the same invocation. The window is capped at 16 IDs, one teacher step. No serving path, kernel or checkpoint format changes: this measures precision rather than gating it. Tests run: test-deepseek-v41-reference 29/29, test-deepseek-v41 79/79, test-deepseek-v41-weights 13/13, zig fmt clean. (cherry picked from commit 53b790082f11aa6b06ba9176912981a1386bcdb0) --- zig/deepseek_v41_reference_checks.zig | 2 + zig/deepseek_v41_run.zig | 34 +- .../deepseek_v41_reference/attention_math.zig | 702 +++++++++++------- .../deepseek_v41_reference/engram_math.zig | 180 ++--- zig/tests/deepseek_v41_reference/math.zig | 570 +++++++++----- .../deepseek_v41_reference/precision.zig | 106 +++ .../deepseek_v41_reference/projection.zig | 331 +++++---- .../deepseek_v41_reference/reference.zig | 332 +++++---- 8 files changed, 1403 insertions(+), 854 deletions(-) create mode 100644 zig/tests/deepseek_v41_reference/precision.zig diff --git a/zig/deepseek_v41_reference_checks.zig b/zig/deepseek_v41_reference_checks.zig index 61aa30f788..1bc109c33c 100644 --- a/zig/deepseek_v41_reference_checks.zig +++ b/zig/deepseek_v41_reference_checks.zig @@ -4,4 +4,6 @@ test { _ = @import("tests/deepseek_v41_reference/math.zig"); _ = @import("tests/deepseek_v41_reference/attention_math.zig"); _ = @import("tests/deepseek_v41_reference/qualify.zig"); + _ = @import("tests/deepseek_v41_reference/engram_math.zig"); + _ = @import("tests/deepseek_v41_reference/precision.zig"); } diff --git a/zig/deepseek_v41_run.zig b/zig/deepseek_v41_run.zig index e6ca81e783..8510131d12 100644 --- a/zig/deepseek_v41_run.zig +++ b/zig/deepseek_v41_run.zig @@ -8,6 +8,7 @@ const load = tf.deepseek_v41_load; const weights = tf.deepseek_v41_weights; const process_memory = @import("src/server/process_memory.zig"); const oracle = @import("tests/deepseek_v41_reference/qualify.zig"); +const precision = @import("tests/deepseek_v41_reference/precision.zig"); const profile = @import("tests/deepseek_v41_reference/profile.zig"); const trace = @import("tests/deepseek_v41_reference/trace.zig"); const operators = @import("tests/deepseek_v41_reference/operators.zig"); @@ -21,6 +22,7 @@ fn usage(code: u8) noreturn { " tf-dsv41-run teacher MODEL TABLES IDS_JSON OUT_DIR [CHUNK]\n" ++ " tf-dsv41-run seams MODEL TABLES IDS_JSON OUT_DIR\n" ++ " tf-dsv41-run oracle MODEL TABLES IDS_JSON OUT_DIR\n" ++ + " tf-dsv41-run oracle-precision MODEL TABLES IDS_JSON OUT_DIR\n" ++ " tf-dsv41-run profile MODEL TABLES IDS_JSON OUT_DIR\n" ++ " tf-dsv41-run trace MODEL TABLES IDS_JSON OUT_DIR\n" ++ " tf-dsv41-run operators MODEL TABLES OUT_DIR\n" ++ @@ -30,7 +32,7 @@ fn usage(code: u8) noreturn { " tf-dsv41-run moe-sameinput MODEL TABLES TRACE_GPU_BIN OUT_DIR [LAYER]\n" ++ " tf-dsv41-run generate MODEL TABLES IDS_JSON OUT_DIR [TOKENS]\n" ++ "IDS_JSON: nonempty JSON array of token IDs below 129280.\n" ++ - "Limits: oracle/trace 16 IDs; profile 256 IDs; teacher/seams 4096 IDs.\n" ++ + "Limits: oracle/oracle-precision/trace 16 IDs; profile 256 IDs; teacher/seams 4096 IDs.\n" ++ "History: TOKEN below 4096; retained normalized prefix at most 256 MiB.\n" ++ "Local diagnostics do not qualify the model. Oracle retains the independent 1% NRMS gate.\n", .{}); std.process.exit(code); @@ -46,10 +48,19 @@ fn validArguments(args: []const [:0]const u8) bool { if (std.mem.eql(u8, mode, "attention-history-sameinput")) return args.len == 8; for ([_][]const u8{ "teacher", "generate", "attention-sameinput", "hc-sameinput", "moe-sameinput" }) |name| if (std.mem.eql(u8, mode, name)) return args.len == 6 or args.len == 7; - for ([_][]const u8{ "seams", "oracle", "profile", "trace" }) |name| + for ([_][]const u8{ "seams", "oracle", "oracle-precision", "profile", "trace" }) |name| if (std.mem.eql(u8, mode, name)) return args.len == 6; return false; } +test "oracle-precision takes exactly the five mode arguments" { + const six = [_][:0]const u8{ "tf-dsv41-run", "oracle-precision", "model", "tables", "ids.json", "out" }; + const five = [_][:0]const u8{ "tf-dsv41-run", "oracle-precision", "model", "tables", "ids.json" }; + const seven = [_][:0]const u8{ "tf-dsv41-run", "oracle-precision", "model", "tables", "ids.json", "out", "16" }; + try std.testing.expect(validArguments(&six)); + try std.testing.expect(!validArguments(&five)); + try std.testing.expect(!validArguments(&seven)); +} + fn argmax(row: []const f32) !u32 { var best: u32 = 0; if (row.len == 0) return error.EmptyLogits; @@ -144,18 +155,19 @@ pub fn main(init: std.process.Init) !void { const is_attention_history = std.mem.eql(u8, mode, "attention-history-sameinput"); const is_hc = std.mem.eql(u8, mode, "hc-sameinput"); const is_moe = std.mem.eql(u8, mode, "moe-sameinput"); - const max_ids: usize = if (std.mem.eql(u8, mode, "oracle") or std.mem.eql(u8, mode, "trace")) 16 else if (std.mem.eql(u8, mode, "profile")) ds.prefill_step else if (std.mem.eql(u8, mode, "teacher") or std.mem.eql(u8, mode, "seams")) 4096 else 1048576; + const max_ids: usize = if (std.mem.eql(u8, mode, "oracle") or std.mem.eql(u8, mode, "oracle-precision") or std.mem.eql(u8, mode, "trace")) 16 else if (std.mem.eql(u8, mode, "profile")) ds.prefill_step else if (std.mem.eql(u8, mode, "teacher") or std.mem.eql(u8, mode, "seams")) 4096 else 1048576; const input = if (is_operators or is_attention or is_attention_history or is_hc or is_moe) &.{} else try std.Io.Dir.cwd().readFileAlloc(io, args[4], arena, .limited(max_ids * 12 + 2)); const ids: []const u32 = if (is_operators or is_attention or is_attention_history or is_hc or is_moe) &.{} else try std.json.parseFromSliceLeaky([]const u32, arena, input, .{}); if (!is_operators and !is_attention and !is_attention_history and !is_hc and !is_moe and (ids.len == 0 or ids.len > 1048576)) return error.InvalidPrompt; if (ids.len > max_ids) return error.InvalidPrompt; for (ids) |id| if (id >= 129280) return error.InvalidToken; const is_oracle = std.mem.eql(u8, mode, "oracle"); + const is_precision = std.mem.eql(u8, mode, "oracle-precision"); const is_trace = std.mem.eql(u8, mode, "trace"); const dir = args[if (is_operators) @as(usize, 4) else 5]; try std.Io.Dir.cwd().createDirPath(io, dir); const map = try std.fs.path.join(arena, &.{ args[2], "engram-token-map.bin" }); - const map_blob = if (is_oracle or is_trace) + const map_blob = if (is_oracle or is_trace or is_precision) try std.Io.Dir.cwd().readFileAlloc(io, map, arena, .limited(2 << 20)) else &.{}; @@ -206,6 +218,7 @@ fn execute(a: std.mem.Allocator, io: std.Io, model: *ds.Model, args: []const [:0 const is_hc = std.mem.eql(u8, mode, "hc-sameinput"); const is_moe = std.mem.eql(u8, mode, "moe-sameinput"); const is_oracle = std.mem.eql(u8, mode, "oracle"); + const is_precision = std.mem.eql(u8, mode, "oracle-precision"); const is_profile = std.mem.eql(u8, mode, "profile"); const is_trace = std.mem.eql(u8, mode, "trace"); const is_teacher = std.mem.eql(u8, mode, "teacher"); @@ -223,6 +236,19 @@ fn execute(a: std.mem.Allocator, io: std.Io, model: *ds.Model, args: []const [:0 } else if (is_oracle) { try operators.run(a, io, model, dir); _ = try oracle.run(a, io, model, ids, map_blob, dir); + } else if (is_precision) { + const run = try precision.run(a, model, ids, map_blob); + defer run.logits.deinit(a); + // Capture execution before the diagnostics and artifacts below, not the reads they cause. + execution_after.* = try counters(); + for (run.logits.argmax, 0..) |pair, i| + std.debug.print("precision token {d}/{d}: fp32 argmax={d}, fp64 argmax={d}\n", .{ i + 1, run.logits.argmax.len, pair[0], pair[1] }); + const json = try std.json.Stringify.valueAlloc(a, run.result, .{}); + defer a.free(json); + std.debug.print("{s}\n", .{json}); + try write(io, a, dir, "precision-fp32-vs-fp64.json", json); + try write(io, a, dir, "reference-fp32-logits.f32", std.mem.sliceAsBytes(run.logits.fp32)); + try write(io, a, dir, "reference-fp64-logits.f64", std.mem.sliceAsBytes(run.logits.fp64)); } else if (is_profile) { try profile.run(a, io, model, ids, dir); } else if (is_trace) { diff --git a/zig/tests/deepseek_v41_reference/attention_math.zig b/zig/tests/deepseek_v41_reference/attention_math.zig index 1adec0172d..889dced2d3 100644 --- a/zig/tests/deepseek_v41_reference/attention_math.zig +++ b/zig/tests/deepseek_v41_reference/attention_math.zig @@ -1,275 +1,477 @@ -//! CPU CSA2 from DeepSeek AI inference/{model,kernel}.py, rev 2cba9e42aa026125f3ed06c6d98c1db82f7ca027. +//! CPU CSA2 from DeepSeek AI inference/{model,kernel}.py, rev 2cba9e42; Module(Real) runs it FP32 or fully FP64. + const std = @import("std"); const tf = @import("tensorfold"); -const m = @import("math.zig"); +const math = @import("math.zig"); +const m = math; // FP32 surface retained for consumers; Module bodies use their own instance below. const projection = @import("projection.zig"); const Config = tf.deepseek_v41_config.Text; const Attention = tf.deepseek_v41_weights.Attention; pub const max_capacity = 4096; -pub const State = struct { - gpa: std.mem.Allocator, - capacity: usize, - window: []u16, - sources: []u16, - scores: []f32, - candidates: []bool, - blocks: []f32, - active_owner: usize = 0, - selected: [512]u32 = undefined, - selected_len: usize = 0, - partial_kv: [4][2][512]f32 = undefined, - partial_gate: [4][2][512]f32 = undefined, - qr: [1280]u16 = undefined, - q: [32768]u16 = undefined, - index_q: [4096]u16 = undefined, - index_weights: [32]f32 = undefined, - latent: [512]u16 = undefined, - index_key: [128]u16 = undefined, - kv: [512]u16 = undefined, - o: [32768]u16 = undefined, - grouped: [8192]u16 = undefined, - observer: ?Observer = null, - sparse_snapshot: ?[]u16 = null, - pub const Observer = struct { - ctx: *anyopaque, - capture: *const fn (*anyopaque, tf.deepseek_v41.attention.DiagnosticStage, []const u16) void, - }; - fn observe(self: *State, stage: tf.deepseek_v41.attention.DiagnosticStage, values: []const u16) void { - if (self.observer) |observer| observer.capture(observer.ctx, stage, values); - } - - pub fn init(gpa: std.mem.Allocator, capacity: usize) !State { - if (capacity == 0 or capacity > max_capacity) return error.InvalidAttentionReferenceCapacity; - const window = try gpa.alloc(u16, 40 * 128 * 512); - errdefer gpa.free(window); - const sources = try gpa.alloc(u16, 4 * capacity * 640); - errdefer gpa.free(sources); - const scores = try gpa.alloc(f32, capacity); - errdefer gpa.free(scores); - const candidates = try gpa.alloc(bool, capacity); - errdefer gpa.free(candidates); - const blocks = try gpa.alloc(f32, (capacity + 7) / 8); - @memset(window, 0); - return .{ .gpa = gpa, .capacity = capacity, .window = window, .sources = sources, .scores = scores, .candidates = candidates, .blocks = blocks }; - } - pub fn deinit(self: *State) void { - self.gpa.free(self.window); - self.gpa.free(self.sources); - self.gpa.free(self.scores); - self.gpa.free(self.candidates); - self.gpa.free(self.blocks); - self.* = undefined; +/// Kernel-shaped key addressing: window rows at physical ring slots, then compressed rows. +pub const KeyTable = struct { + window: []const u16, + window_stride: usize = 512, + compressed: []const u16 = &[_]u16{}, + compressed_stride: usize = 512, + pub fn keyRow(t: KeyTable, at: usize) []const u16 { + if (at < 128) return t.window[at * t.window_stride ..][0..512]; + const row = at - 128; + return t.compressed[row * t.compressed_stride ..][0..512]; } - fn mainRow(self: *State, owner: usize, row: usize) []u16 { - return self.sources[owner * self.capacity * 640 + row * 640 ..][0..512]; - } - fn indexRow(self: *State, owner: usize, row: usize) []u16 { - return self.sources[owner * self.capacity * 640 + row * 640 + 512 ..][0..128]; +}; +/// Compressor softmax pooling at reference precision T: FP32 is the source contract, FP64 the ideal. +pub fn poolOf(comptime T: type, kv: [2][]const T, gate: [2][]const T, out: []u16) void { + @setFloatMode(.strict); + for (out, 0..) |*o, d| { + const high = @max(gate[0][d], gate[1][d]); + const p0 = @exp(gate[0][d] - high); + const p1 = @exp(gate[1][d] - high); + const denominator = p0 + p1; + const pooled = kv[0][d] * (p0 / denominator) + kv[1][d] * (p1 / denominator); + o.* = if (T == f32) m.b(@floatCast(pooled)) else m.b64(pooled); } - fn publish(self: *State, a: Attention, layer: usize, x: []const u16, position: usize, c: Config, scratch: *projection.Scratch) !bool { - @setFloatMode(.strict); - const compressor = a.compressor orelse return false; - const owner = (layer - 2) / 6; - self.active_owner = owner; - const ratio: usize = c.compress_ratios[layer]; - if (ratio == 1) { - try projection.apply(compressor.wkv, x, &self.latent, scratch); - } else { - const slot = position % ratio; - try projection.applyF32(compressor.wkv, x, &self.partial_kv[owner][slot], scratch); - try projection.applyF32(compressor.wgate.?, x, &self.partial_gate[owner][slot], scratch); - if ((position + 1) % ratio != 0) return false; - for (0..512) |d| { - const high = @max(self.partial_gate[owner][0][d], self.partial_gate[owner][1][d]); - const p0 = @exp(self.partial_gate[owner][0][d] - high); - const p1 = @exp(self.partial_gate[owner][1][d] - high); - const denominator = p0 + p1; - const weighted0 = self.partial_kv[owner][0][d] * (p0 / denominator); - const weighted1 = self.partial_kv[owner][1][d] * (p1 / denominator); - const pooled = weighted0 + weighted1; - self.latent[d] = m.b(pooled); +} +fn narrow(comptime T: type, x: f64) T { + return if (T == f32) @as(f32, @floatCast(x)) else x; +} +/// The sparse-attention head at reference precision T, with the pre-normalization quotients if asked. +pub fn sparseHeadPre( + comptime T: type, + q: []const u16, + keys: KeyTable, + indices: []const i32, + sink: f32, + out: []u16, + pre: ?[]T, +) void { + @setFloatMode(.strict); + const scale: T = if (T == f32) 1 / @sqrt(@as(f32, 512)) else 1 / @sqrt(@as(f64, 512)); + var acc: [512]T = @splat(0); + var high: T = m.wide(T, -1e30); + var denominator: T = 0; + var base: usize = 0; + while (base < indices.len) : (base += 64) { + var probabilities: [64]T = undefined; + const width = @min(64, indices.len - base); + var next_high = high; + for (0..width) |j| { + const index = indices[base + j]; + var dot: f64 = 0; + if (index >= 0) { + const key = keys.keyRow(@intCast(index)); + for (q, key) |a, v| dot += m.d64(a) * m.d64(v); } + probabilities[j] = if (index < 0) -std.math.inf(T) else narrow(T, dot) * scale; + next_high = @max(next_high, probabilities[j]); } - m.rms(&self.latent, compressor.norm, c.rms_norm_eps, &self.latent); - const indexer = a.indexer.?; - try projection.apply(indexer.wk.?, &self.latent, &self.index_key, scratch); - m.rms(&self.index_key, indexer.k_norm.?, c.rms_norm_eps, &self.index_key); - const group_position = position + 1 - ratio; - m.rope(&self.index_key, group_position, true, false, c); - m.quantize(&self.index_key, .index); - const row = position / ratio; - @memcpy(self.indexRow(owner, row), &self.index_key); - m.rope(&self.latent, group_position, true, false, c); - m.quantize(&self.latent, .main); - @memcpy(self.mainRow(owner, row), &self.latent); - return true; - } - fn chooseCandidates(self: *State, count: usize, c: Config) void { - @memset(self.candidates[0..count], false); - const nblocks = (count + 7) / 8; - for (0..nblocks) |block| { - var score: f32 = -std.math.inf(f32); - for (block * 8..@min(block * 8 + 8, count)) |row| score = @max(score, self.scores[row]); - self.blocks[block] = score; + const correction = @exp(high - next_high); + var sum: f64 = 0; + for (0..width) |j| { + probabilities[j] = @exp(probabilities[j] - next_high); + sum += probabilities[j]; } - if (nblocks == 0) return; - self.blocks[nblocks - 1] = std.math.inf(f32); - for (0..@min(nblocks, c.candidate_topk_blocks)) |_| { - var best: usize = 0; - for (0..nblocks) |i| if (self.blocks[i] > self.blocks[best]) { - best = i; + denominator = denominator * correction + narrow(T, sum); + for (0..512) |d| { + var accumulated: f64 = acc[d] * correction; + for (0..width) |j| if (indices[base + j] >= 0) { + // Online 64-column source kernel rounds unnormalized probabilities to BF16 before P*V. + const p: f64 = if (T == f32) @as(f64, m.f(m.b(probabilities[j]))) else m.d64(m.b64(probabilities[j])); + accumulated += p * m.d64(keys.keyRow(@intCast(indices[base + j]))[d]); }; - if (self.blocks[best] == -std.math.inf(f32)) break; - for (best * 8..@min(best * 8 + 8, count)) |row| self.candidates[row] = true; - self.blocks[best] = -std.math.inf(f32); + acc[d] = @floatCast(accumulated); } + high = next_high; } - fn selectPositions(self: *State, a: Attention, layer: usize, x: []const u16, position: usize, count: usize, c: Config, scratch: *projection.Scratch) !void { - const indexer = a.indexer orelse return; - self.selected_len = @min(count, c.index_topk); - if (count == 0) return; - try projection.apply(indexer.wq_b, &self.qr, &self.index_q, scratch); - for (0..32) |head| { - const q = self.index_q[head * 128 ..][0..128]; - m.rope(q, position, true, false, c); - m.quantize(q, .index); - } - m.dense(indexer.weights_proj, x, &self.index_weights); - const scale: f32 = 1 / @sqrt(@as(f32, 128 * 32)); - // Python scalar multiplication preserves the source weights projection's BF16 dtype. - for (&self.index_weights) |*v| v.* = m.f(m.b(m.f(m.b(v.*)) * scale)); - for (0..count) |row| { - const key = self.indexRow(self.active_owner, row); - var score: f64 = 0; - for (0..32) |head| { - var dot: f64 = 0; - for (key, 0..) |v, d| dot += @as(f64, m.f(v)) * m.f(self.index_q[head * 128 + d]); - // einsum produces BF16, then BF16 multiply and head sum, in the pinned source. - const rectified = @max(m.f(m.b(@floatCast(dot))), 0); - const weighted = m.f(m.b(rectified * self.index_weights[head])); - score += weighted; - } - self.scores[row] = m.f(m.b(@floatCast(score))); - } - if (layer == c.candidate_source_layer_id) self.chooseCandidates(count, c); - if (layer > c.candidate_source_layer_id) for (0..count) |row| { - if (!self.candidates[row]) self.scores[row] = -std.math.inf(f32); - }; - for (0..self.selected_len) |pick| { - var best: usize = 0; - for (0..count) |i| if (self.scores[i] > self.scores[best]) { - best = i; - }; - self.selected[pick] = @intCast(best); - self.scores[best] = -std.math.inf(f32); - } - std.mem.sort(u32, self.selected[0..self.selected_len], {}, std.sort.asc(u32)); + denominator += @exp(m.wide(T, sink) - high); + for (out, acc, 0..) |*v, a, i| { + const quotient = a / denominator; + if (pre) |cells| cells[i] = quotient; + v.* = if (T == f32) m.b(quotient) else m.b64(quotient); } - pub fn apply(self: *State, a: Attention, layer: usize, x: []const u16, position: usize, c: Config, out: []u16, scratch: *projection.Scratch) !void { - const compressed = c.compress_ratios[layer] != 0; - try projection.apply(.{ .trellis = a.wq_a }, x, &self.qr, scratch); - self.observe(.qr_projection, &self.qr); - m.rms(&self.qr, a.q_norm, c.rms_norm_eps, &self.qr); - self.observe(.qr_norm, &self.qr); - try projection.apply(.{ .trellis = a.wq_b }, &self.qr, &self.q, scratch); - self.observe(.q_projection, &self.q); - for (0..64) |head| m.rope(self.q[head * 512 ..][0..512], position, compressed, false, c); - self.observe(.q, &self.q); - try projection.apply(.{ .trellis = a.wkv }, x, &self.kv, scratch); - self.observe(.kv_projection, &self.kv); - m.rms(&self.kv, a.kv_norm, c.rms_norm_eps, &self.kv); - self.observe(.kv_norm, &self.kv); - m.rope(&self.kv, position, compressed, false, c); - self.observe(.kv_rope, &self.kv); - m.quantize(&self.kv, .window); - self.observe(.kv_quantized, &self.kv); - const window = self.window[layer * 128 * 512 ..][0 .. 128 * 512]; - @memcpy(window[(position % 128) * 512 ..][0..512], &self.kv); - var count: usize = 0; - if (compressed) { - _ = try self.publish(a, layer, x, position, c, scratch); - count = (position + 1) / c.compress_ratios[layer]; - try self.selectPositions(a, layer, x, position, count, c, scratch); - } - var indices: [640]i32 = undefined; - var slots: usize = 0; - if (position == 0) { - indices[0] = 0; - slots = 1; - } else { - const oldest = position % 128 + 1; - for (0..128) |i| { - const at = (oldest + i) % 128; - indices[i] = if (at > position) -1 else @intCast(at); +} + +pub fn Module(comptime Real: type) type { + return struct { + const Self = @This(); + const wr = math.Module(Real).wr; + const m64 = math.Module(Real); + const proj = projection.Module(Real); + + pub const State = struct { + gpa: std.mem.Allocator, + capacity: usize, + window: []u16, + sources: []u16, + scores: []Real, + candidates: []bool, + blocks: []Real, + active_owner: usize = 0, + selected: [512]u32 = undefined, + selected_len: usize = 0, + partial_kv: [4][2][512]Real = undefined, + partial_gate: [4][2][512]Real = undefined, + qr: [1280]u16 = undefined, + q: [32768]u16 = undefined, + index_q: [4096]u16 = undefined, + index_weights: [32]Real = undefined, + latent: [512]u16 = undefined, + index_key: [128]u16 = undefined, + kv: [512]u16 = undefined, + o: [32768]u16 = undefined, + grouped: [8192]u16 = undefined, + observer: ?Observer = null, + sparse_snapshot: ?[]u16 = null, + /// Indexer-selection probe on the retained FP32 path: contract and ideal margins plus pick flips. + score_probe: ?ScoreProbe = null, + score_copy: []f32 = &.{}, + ideal_scores: []f32 = &.{}, + ideal_copy: []f32 = &.{}, + pick_work: []f32 = &.{}, + selected_ideal: []u32 = &.{}, + pub const ScoreProbe = struct { + ctx: *anyopaque, + record: *const fn (*anyopaque, layer: usize, position: usize, count: usize, selected_len: usize, margin: f32, margin_ideal: f32, flips: usize, lowest: f32, highest_unselected: f32) void, + }; + pub const Observer = struct { + ctx: *anyopaque, + capture: *const fn (*anyopaque, tf.deepseek_v41.attention.DiagnosticStage, []const u16) void, + }; + fn observe(self: *Self.State, stage: tf.deepseek_v41.attention.DiagnosticStage, values: []const u16) void { + if (self.observer) |observer| observer.capture(observer.ctx, stage, values); } - slots = 128; - } - if (compressed) for (self.selected[0..self.selected_len]) |row| { - indices[slots] = @intCast(128 + row); - slots += 1; - }; - for (0..64) |head| { - self.sparseHead(self.q[head * 512 ..][0..512], window, indices[0..slots], m.value(a.sink, head), self.o[head * 512 ..][0..512]); - self.observeHead(head); - m.rope(self.o[head * 512 ..][0..512], position, compressed, true, c); - } - if (self.sparse_snapshot) |snapshot| self.observe(.sparse_o, snapshot); - self.observe(.o, &self.o); - for (0..8) |group| try projection.apply(.{ .trellis = a.wo_a[group] }, self.o[group * 4096 ..][0..4096], self.grouped[group * 1024 ..][0..1024], scratch); - self.observe(.grouped, &self.grouped); - try projection.apply(.{ .trellis = a.wo_b }, &self.grouped, out, scratch); - self.observe(.wo_b, out); - } - fn observeHead(self: *State, head: usize) void { - if (self.sparse_snapshot) |snapshot| - @memcpy(snapshot[head * 512 ..][0..512], self.o[head * 512 ..][0..512]); - } - fn keyRow(self: *State, window: []const u16, key_id: i32) []const u16 { - const at: usize = @intCast(key_id); - return if (at < 128) window[at * 512 ..][0..512] else self.mainRow(self.active_owner, at - 128); - } - pub fn sparseHead(self: *State, q: []const u16, window: []const u16, indices: []const i32, sink: f32, out: []u16) void { - @setFloatMode(.strict); - var acc: [512]f32 = @splat(0); - var high: f32 = -1e30; - var denominator: f32 = 0; - var base: usize = 0; - while (base < indices.len) : (base += 64) { - var probabilities: [64]f32 = undefined; - const width = @min(64, indices.len - base); - var next_high = high; - for (0..width) |j| { - const index = indices[base + j]; - var dot: f64 = 0; - if (index >= 0) { - const key = self.keyRow(window, index); - for (q, key) |a, v| dot += @as(f64, m.f(a)) * m.f(v); + + pub fn init(gpa: std.mem.Allocator, capacity: usize) !Self.State { + if (capacity == 0 or capacity > max_capacity) return error.InvalidAttentionReferenceCapacity; + const window = try gpa.alloc(u16, 40 * 128 * 512); + errdefer gpa.free(window); + const sources = try gpa.alloc(u16, 4 * capacity * 640); + errdefer gpa.free(sources); + const scores = try gpa.alloc(Real, capacity); + errdefer gpa.free(scores); + const candidates = try gpa.alloc(bool, capacity); + errdefer gpa.free(candidates); + const blocks = try gpa.alloc(Real, (capacity + 7) / 8); + errdefer gpa.free(blocks); + const score_copy = try gpa.alloc(f32, capacity); + errdefer gpa.free(score_copy); + const ideal_scores = try gpa.alloc(f32, capacity); + errdefer gpa.free(ideal_scores); + const ideal_copy = try gpa.alloc(f32, capacity); + errdefer gpa.free(ideal_copy); + const pick_work = try gpa.alloc(f32, capacity); + errdefer gpa.free(pick_work); + const selected_ideal = try gpa.alloc(u32, capacity); + @memset(window, 0); + return .{ .gpa = gpa, .capacity = capacity, .window = window, .sources = sources, .scores = scores, .candidates = candidates, .blocks = blocks, .score_copy = score_copy, .ideal_scores = ideal_scores, .ideal_copy = ideal_copy, .pick_work = pick_work, .selected_ideal = selected_ideal }; + } + pub fn deinit(self: *Self.State) void { + self.gpa.free(self.window); + self.gpa.free(self.sources); + self.gpa.free(self.scores); + self.gpa.free(self.candidates); + self.gpa.free(self.blocks); + self.gpa.free(self.score_copy); + self.gpa.free(self.ideal_scores); + self.gpa.free(self.ideal_copy); + self.gpa.free(self.pick_work); + self.gpa.free(self.selected_ideal); + self.* = undefined; + } + fn mainRow(self: *Self.State, owner: usize, row: usize) []u16 { + return self.sources[owner * self.capacity * 640 + row * 640 ..][0..512]; + } + fn indexRow(self: *Self.State, owner: usize, row: usize) []u16 { + return self.sources[owner * self.capacity * 640 + row * 640 + 512 ..][0..128]; + } + fn publish(self: *Self.State, a: Attention, layer: usize, x: []const u16, position: usize, c: Config, scratch: *projection.Scratch) !bool { + @setFloatMode(.strict); + const compressor = a.compressor orelse return false; + const owner = (layer - 2) / 6; + self.active_owner = owner; + const ratio: usize = c.compress_ratios[layer]; + if (ratio == 1) { + try proj.apply(compressor.wkv, x, &self.latent, scratch); + } else { + const slot = position % ratio; + try proj.applyF32(compressor.wkv, x, &self.partial_kv[owner][slot], scratch); + try proj.applyF32(compressor.wgate.?, x, &self.partial_gate[owner][slot], scratch); + if ((position + 1) % ratio != 0) return false; + for (0..512) |d| { + const high = @max(self.partial_gate[owner][0][d], self.partial_gate[owner][1][d]); + const p0 = @exp(self.partial_gate[owner][0][d] - high); + const p1 = @exp(self.partial_gate[owner][1][d] - high); + const denominator = p0 + p1; + const weighted0 = self.partial_kv[owner][0][d] * (p0 / denominator); + const weighted1 = self.partial_kv[owner][1][d] * (p1 / denominator); + const pooled = weighted0 + weighted1; + self.latent[d] = m64.b(pooled); + } } - probabilities[j] = if (index < 0) -std.math.inf(f32) else @as(f32, @floatCast(dot)) * (1 / @sqrt(@as(f32, 512))); - next_high = @max(next_high, probabilities[j]); + m64.rms(&self.latent, compressor.norm, c.rms_norm_eps, &self.latent); + const indexer = a.indexer.?; + try proj.apply(indexer.wk.?, &self.latent, &self.index_key, scratch); + m64.rms(&self.index_key, indexer.k_norm.?, c.rms_norm_eps, &self.index_key); + const group_position = position + 1 - ratio; + m64.rope(&self.index_key, group_position, true, false, c); + m64.quantize(&self.index_key, .index); + const row = position / ratio; + @memcpy(self.indexRow(owner, row), &self.index_key); + m64.rope(&self.latent, group_position, true, false, c); + m64.quantize(&self.latent, .main); + @memcpy(self.mainRow(owner, row), &self.latent); + return true; } - const correction = @exp(high - next_high); - var sum: f64 = 0; - for (0..width) |j| { - probabilities[j] = @exp(probabilities[j] - next_high); - sum += probabilities[j]; + fn chooseCandidates(self: *Self.State, count: usize, c: Config) void { + @memset(self.candidates[0..count], false); + const nblocks = (count + 7) / 8; + for (0..nblocks) |block| { + var score: Real = -std.math.inf(Real); + for (block * 8..@min(block * 8 + 8, count)) |row| score = @max(score, self.scores[row]); + self.blocks[block] = score; + } + if (nblocks == 0) return; + self.blocks[nblocks - 1] = std.math.inf(Real); + for (0..@min(nblocks, c.candidate_topk_blocks)) |_| { + var best: usize = 0; + for (0..nblocks) |i| if (self.blocks[i] > self.blocks[best]) { + best = i; + }; + if (self.blocks[best] == -std.math.inf(Real)) break; + for (best * 8..@min(best * 8 + 8, count)) |row| self.candidates[row] = true; + self.blocks[best] = -std.math.inf(Real); + } } - denominator = denominator * correction + @as(f32, @floatCast(sum)); - for (0..512) |d| { - var accumulated: f64 = acc[d] * correction; - for (0..width) |j| if (indices[base + j] >= 0) { - // Online 64-column source kernel rounds unnormalized probabilities to BF16 before P*V. - accumulated += @as(f64, m.f(m.b(probabilities[j]))) * m.f(self.keyRow(window, indices[base + j])[d]); + fn selectPositions(self: *Self.State, a: Attention, layer: usize, x: []const u16, position: usize, count: usize, c: Config, scratch: *projection.Scratch) !void { + const indexer = a.indexer orelse return; + self.selected_len = @min(count, c.index_topk); + if (count == 0) return; + try proj.apply(indexer.wq_b, &self.qr, &self.index_q, scratch); + for (0..32) |head| { + const q = self.index_q[head * 128 ..][0..128]; + m64.rope(q, position, true, false, c); + m64.quantize(q, .index); + } + m64.dense(indexer.weights_proj, x, &self.index_weights); + const scale: Real = 1 / @sqrt(@as(Real, 128 * 32)); + // Python scalar multiplication preserves the source weights projection's BF16 dtype. + for (&self.index_weights) |*v| v.* = m64.f(m64.b(m64.f(m64.b(v.*)) * scale)); + for (0..count) |row| { + const key = self.indexRow(self.active_owner, row); + var score: f64 = 0; + for (0..32) |head| { + var dot: f64 = 0; + for (key, 0..) |v, d| dot += @as(f64, m64.f(v)) * m64.f(self.index_q[head * 128 + d]); + // einsum produces BF16, then BF16 multiply and head sum, in the pinned source. + const rectified = @max(m64.f(m64.b(@floatCast(dot))), 0); + const weighted = m64.f(m64.b(rectified * self.index_weights[head])); + score += weighted; + } + self.scores[row] = m64.f(m64.b(@floatCast(score))); + } + if (layer == c.candidate_source_layer_id) self.chooseCandidates(count, c); + if (layer > c.candidate_source_layer_id) for (0..count) |row| { + if (!self.candidates[row]) self.scores[row] = -std.math.inf(Real); }; - acc[d] = @floatCast(accumulated); + if (self.score_probe != null) { + for (0..count) |row| self.score_copy[row] = @floatCast(self.scores[row]); + self.idealScores(a, count); + } + for (0..self.selected_len) |pick| { + var best: usize = 0; + for (0..count) |i| if (self.scores[i] > self.scores[best]) { + best = i; + }; + self.selected[pick] = @intCast(best); + self.scores[best] = -std.math.inf(Real); + } + std.mem.sort(u32, self.selected[0..self.selected_len], {}, std.sort.asc(u32)); + if (self.score_probe) |probe| self.reportSelection(probe, layer, position, count); } - high = next_high; - } - denominator += @exp(sink - high); - for (out, acc) |*v, a| v.* = m.b(a / denominator); - } -}; + /// Contract-vs-ideal selection margins: the same BF16 inputs scored without intermediate roundings. + fn idealScores(self: *Self.State, a: Attention, count: usize) void { + @setFloatMode(.strict); + if (a.indexer == null) return; + for (0..count) |row| { + if (self.score_copy[row] == -std.math.inf(f32)) { + self.ideal_scores[row] = -std.math.inf(f32); + continue; + } + const key = self.indexRow(self.active_owner, row); + var score: f64 = 0; + for (0..32) |head| { + var dot: f64 = 0; + for (key, 0..) |v, d| dot += @as(f64, m64.f(v)) * @as(f64, m64.f(self.index_q[head * 128 + d])); + score += @max(dot, 0) * @as(f64, @floatCast(self.index_weights[head])); + } + self.ideal_scores[row] = @floatCast(score); + } + } + fn pickTop(self: *Self.State, scores: []const f32, into: []u32) void { + @memcpy(self.pick_work[0..scores.len], scores); + const work = self.pick_work[0..scores.len]; + for (into) |*pick| { + var best: usize = 0; + for (work, 0..) |v, i| if (v > work[best]) { + best = i; + }; + pick.* = @intCast(best); + work[best] = -std.math.inf(f32); + } + std.mem.sort(u32, into, {}, std.sort.asc(u32)); + } + fn pickedAt(selection: []const u32, row: usize) bool { + var lo: usize = 0; + var hi: usize = selection.len; + while (lo < hi) { + const mid = (lo + hi) / 2; + if (selection[mid] == row) return true; + if (selection[mid] < row) lo = mid + 1 else hi = mid; + } + return false; + } + fn reportSelection(self: *Self.State, probe: ScoreProbe, layer: usize, position: usize, count: usize) void { + var lowest: f32 = std.math.inf(f32); + var highest_unselected: f32 = -std.math.inf(f32); + for (self.selected[0..self.selected_len]) |row| { + lowest = @min(lowest, self.score_copy[row]); + } + for (0..count) |row| { + if (!pickedAt(self.selected[0..self.selected_len], row) and self.score_copy[row] > -std.math.inf(f32)) + highest_unselected = @max(highest_unselected, self.score_copy[row]); + } + const margin: f32 = if (self.selected_len == 0 or lowest == std.math.inf(f32)) 0 else lowest - highest_unselected; + @memcpy(self.ideal_copy[0..count], self.ideal_scores[0..count]); + const ideal_len = @min(self.selected_len, self.selected_ideal.len); + self.pickTop(self.ideal_copy[0..count], self.selected_ideal[0..ideal_len]); + var flips: usize = 0; + if (self.selected_len == ideal_len) { + for (self.selected[0..self.selected_len], self.selected_ideal[0..ideal_len]) |s, i| { + if (s != i) flips += 1; + } + } else flips = self.selected_len; + var ideal_lowest: f32 = std.math.inf(f32); + var ideal_highest_unselected: f32 = -std.math.inf(f32); + for (self.selected_ideal[0..ideal_len]) |row| { + ideal_lowest = @min(ideal_lowest, self.ideal_scores[row]); + } + for (0..count) |row| { + if (!pickedAt(self.selected_ideal[0..ideal_len], row) and self.ideal_scores[row] > -std.math.inf(f32)) + ideal_highest_unselected = @max(ideal_highest_unselected, self.ideal_scores[row]); + } + const margin_ideal: f32 = if (ideal_len == 0 or ideal_lowest == std.math.inf(f32)) 0 else ideal_lowest - ideal_highest_unselected; + probe.record(probe.ctx, layer, position, count, self.selected_len, margin, margin_ideal, flips, lowest, highest_unselected); + } + pub fn apply(self: *Self.State, a: Attention, layer: usize, x: []const u16, position: usize, c: Config, out: []u16, scratch: *projection.Scratch) !void { + const compressed = c.compress_ratios[layer] != 0; + try proj.applyTrellis(a.wq_a, x, &self.qr, scratch); + self.observe(.qr_projection, &self.qr); + m64.rms(&self.qr, a.q_norm, c.rms_norm_eps, &self.qr); + self.observe(.qr_norm, &self.qr); + try proj.applyTrellis(a.wq_b, &self.qr, &self.q, scratch); + self.observe(.q_projection, &self.q); + for (0..64) |head| m64.rope(self.q[head * 512 ..][0..512], position, compressed, false, c); + self.observe(.q, &self.q); + try proj.applyTrellis(a.wkv, x, &self.kv, scratch); + self.observe(.kv_projection, &self.kv); + m64.rms(&self.kv, a.kv_norm, c.rms_norm_eps, &self.kv); + self.observe(.kv_norm, &self.kv); + m64.rope(&self.kv, position, compressed, false, c); + self.observe(.kv_rope, &self.kv); + m64.quantize(&self.kv, .window); + self.observe(.kv_quantized, &self.kv); + const window = self.window[layer * 128 * 512 ..][0 .. 128 * 512]; + @memcpy(window[(position % 128) * 512 ..][0..512], &self.kv); + var count: usize = 0; + if (compressed) { + _ = try self.publish(a, layer, x, position, c, scratch); + count = (position + 1) / c.compress_ratios[layer]; + try self.selectPositions(a, layer, x, position, count, c, scratch); + } + var indices: [640]i32 = undefined; + var slots: usize = 0; + if (position == 0) { + indices[0] = 0; + slots = 1; + } else { + const oldest = position % 128 + 1; + for (0..128) |i| { + const at = (oldest + i) % 128; + indices[i] = if (at > position) -1 else @intCast(at); + } + slots = 128; + } + if (compressed) for (self.selected[0..self.selected_len]) |row| { + indices[slots] = @intCast(128 + row); + slots += 1; + }; + for (0..64) |head| { + self.sparseHead(self.q[head * 512 ..][0..512], window, indices[0..slots], m64.value(a.sink, head), self.o[head * 512 ..][0..512]); + self.observeHead(head); + m64.rope(self.o[head * 512 ..][0..512], position, compressed, true, c); + } + if (self.sparse_snapshot) |snapshot| self.observe(.sparse_o, snapshot); + self.observe(.o, &self.o); + for (0..8) |group| try proj.applyTrellis(a.wo_a[group], self.o[group * 4096 ..][0..4096], self.grouped[group * 1024 ..][0..1024], scratch); + self.observe(.grouped, &self.grouped); + try proj.applyTrellis(a.wo_b, &self.grouped, out, scratch); + self.observe(.wo_b, out); + } + fn observeHead(self: *Self.State, head: usize) void { + if (self.sparse_snapshot) |snapshot| + @memcpy(snapshot[head * 512 ..][0..512], self.o[head * 512 ..][0..512]); + } + fn keyRow(self: *Self.State, window: []const u16, key_id: i32) []const u16 { + const at: usize = @intCast(key_id); + return if (at < 128) window[at * 512 ..][0..512] else self.mainRow(self.active_owner, at - 128); + } + pub fn sparseHead(self: *Self.State, q: []const u16, window: []const u16, indices: []const i32, sink: Real, out: []u16) void { + @setFloatMode(.strict); + var acc: [512]Real = @splat(0); + var high: Real = -1e30; + var denominator: Real = 0; + var base: usize = 0; + while (base < indices.len) : (base += 64) { + var probabilities: [64]Real = undefined; + const width = @min(64, indices.len - base); + var next_high = high; + for (0..width) |j| { + const index = indices[base + j]; + var dot: f64 = 0; + if (index >= 0) { + const key = self.keyRow(window, index); + for (q, key) |a, v| dot += @as(f64, m64.f(a)) * m64.f(v); + } + probabilities[j] = if (index < 0) -std.math.inf(Real) else @as(Real, @floatCast(dot)) * (1 / @sqrt(@as(Real, 512))); + next_high = @max(next_high, probabilities[j]); + } + const correction = @exp(high - next_high); + var sum: f64 = 0; + for (0..width) |j| { + probabilities[j] = @exp(probabilities[j] - next_high); + sum += probabilities[j]; + } + denominator = denominator * correction + @as(Real, @floatCast(sum)); + for (0..512) |d| { + var accumulated: f64 = @as(f64, acc[d] * correction); + for (0..width) |j| if (indices[base + j] >= 0) { + // Online 64-column source kernel rounds unnormalized probabilities to BF16 before P*V. + accumulated += @as(f64, m64.f(m64.b(probabilities[j]))) * m64.f(self.keyRow(window, indices[base + j])[d]); + }; + acc[d] = @floatCast(accumulated); + } + high = next_high; + } + denominator += @exp(sink - high); + for (out, acc) |*v, a| v.* = m64.b(a / denominator); + } + }; + }; +} + +const f32_module = Module(f32); +pub const State = f32_module.State; test "online 64-column sparse attention rounds probabilities before value GEMM" { var state = try State.init(std.testing.allocator, 2); diff --git a/zig/tests/deepseek_v41_reference/engram_math.zig b/zig/tests/deepseek_v41_reference/engram_math.zig index 88772898de..5787ced72f 100644 --- a/zig/tests/deepseek_v41_reference/engram_math.zig +++ b/zig/tests/deepseek_v41_reference/engram_math.zig @@ -1,7 +1,8 @@ -//! CPU Engram from DeepSeek AI inference/engram.py and model.py, rev 2cba9e42aa026125f3ed06c6d98c1db82f7ca027. +//! CPU Engram from DeepSeek AI inference/engram.py and model.py, rev 2cba9e42; Module(Real) keeps FP32 or FP64. + const std = @import("std"); const tf = @import("tensorfold"); -const m = @import("math.zig"); +const math = @import("math.zig"); const projection = @import("projection.zig"); const Engram = tf.deepseek_v41_weights.Engram; const Config = tf.deepseek_v41_config.Text; @@ -16,91 +17,102 @@ const primes = [2][24]u32{ fn word(raw: []const u8, at: usize) u32 { return std.mem.readInt(u32, raw[at..][0..4], .little); } -pub const State = struct { - map: []const u8, - history: [3]u32 = .{ 2, 2, 2 }, - gathered: [6144]u16 = undefined, - kv: [25600]u16 = undefined, - pub fn init(blob: []const u8) !State { - const at = 96 + 2 * 232; - if (blob.len != at + 129280 * 4 or !std.mem.eql(u8, blob[0..8], "DSENGR01")) - return error.InvalidEngramMetadata; - for ([_]u32{ 129280, 99092, 2, 4, 8, 2 }, 0..) |v, i| - if (word(blob, 8 + i * 4) != v) return error.InvalidEngramMetadata; - for (0..2) |layer| { - const base = 96 + layer * 232; - if (word(blob, base) != (if (layer == 0) @as(u32, 1) else @as(u32, 14))) - return error.InvalidEngramMetadata; - var offset: u32 = 0; - for (multipliers[layer], 0..) |v, j| - if (std.mem.readInt(u64, blob[base + 8 + 8 * j ..][0..8], .little) != v) - return error.InvalidEngramMetadata; - for (primes[layer], 0..) |v, j| { - if (word(blob, base + 40 + 4 * j) != v or word(blob, base + 136 + 4 * j) != offset) +pub fn Module(comptime Real: type) type { + return struct { + const m = math.Module(Real); + const proj = projection.Module(Real); + const StateImpl = struct { + map: []const u8, + history: [3]u32 = .{ 2, 2, 2 }, + gathered: [6144]u16 = undefined, + kv: [25600]u16 = undefined, + + pub fn init(blob: []const u8) !StateImpl { + const at = 96 + 2 * 232; + if (blob.len != at + 129280 * 4 or !std.mem.eql(u8, blob[0..8], "DSENGR01")) return error.InvalidEngramMetadata; - offset += v; + for ([_]u32{ 129280, 99092, 2, 4, 8, 2 }, 0..) |v, i| + if (word(blob, 8 + i * 4) != v) return error.InvalidEngramMetadata; + for (0..2) |layer| { + const base = 96 + layer * 232; + if (word(blob, base) != (if (layer == 0) @as(u32, 1) else @as(u32, 14))) + return error.InvalidEngramMetadata; + var offset: u32 = 0; + for (multipliers[layer], 0..) |v, j| + if (std.mem.readInt(u64, blob[base + 8 + 8 * j ..][0..8], .little) != v) + return error.InvalidEngramMetadata; + for (primes[layer], 0..) |v, j| { + if (word(blob, base + 40 + 4 * j) != v or word(blob, base + 136 + 4 * j) != offset) return error.InvalidEngramMetadata; + offset += v; + } + if (word(blob, base + 4) != offset) return error.InvalidEngramMetadata; + } + const map = blob[at..]; + const expected = [_]u8{ 0xc6, 0x0a, 0x86, 0x32, 0x2e, 0xc1, 0x7b, 0x41, 0x42, 0xbf, 0xef, 0x3c, 0x57, 0xa8, 0xd8, 0x1f, 0xb4, 0x28, 0x55, 0x0c, 0xdf, 0x48, 0x7f, 0x88, 0xa4, 0x32, 0x0c, 0xb5, 0x9f, 0xe4, 0x64, 0x81 }; + var digest: [32]u8 = undefined; + std.crypto.hash.sha2.Sha256.hash(map, &digest, .{}); + if (!std.mem.eql(u8, &expected, &digest)) return error.InvalidEngramTokenMap; + return .{ .map = map }; } - if (word(blob, base + 4) != offset) return error.InvalidEngramMetadata; - } - const map = blob[at..]; - const expected = [_]u8{ 0xc6, 0x0a, 0x86, 0x32, 0x2e, 0xc1, 0x7b, 0x41, 0x42, 0xbf, 0xef, 0x3c, 0x57, 0xa8, 0xd8, 0x1f, 0xb4, 0x28, 0x55, 0x0c, 0xdf, 0x48, 0x7f, 0x88, 0xa4, 0x32, 0x0c, 0xb5, 0x9f, 0xe4, 0x64, 0x81 }; - var digest: [32]u8 = undefined; - std.crypto.hash.sha2.Sha256.hash(map, &digest, .{}); - if (!std.mem.eql(u8, &expected, &digest)) return error.InvalidEngramTokenMap; - return .{ .map = map }; - } - pub fn compressed(self: *const State, id: u32) !u32 { - if (id >= 129280) return error.InvalidToken; - return word(self.map, @as(usize, id) * 4); - } - pub fn advance(self: *State, id: u32) void { - self.history[2] = self.history[1]; - self.history[1] = self.history[0]; - self.history[0] = id; - } - pub fn apply(self: *State, e: Engram, layer: usize, id: u32, h: []u16, c: Config, scratch: *projection.Scratch) !void { - @setFloatMode(.strict); - const tokens = [4]u32{ id, self.history[0], self.history[1], self.history[2] }; - var rolling = @as(u64, tokens[0]) * multipliers[layer][0]; - var offset: usize = 0; - for (1..4) |lookback| { - rolling ^= @as(u64, tokens[lookback]) * multipliers[layer][lookback]; - for (0..8) |head| { - const column = (lookback - 1) * 8 + head; - const row = offset + @as(usize, @intCast(rolling % primes[layer][column])); - for (0..256) |d| self.gathered[column * 256 + d] = - m.b(m.value(e.table, row * 256 + d) * m.value(e.scales, row * 8 + d / 32)); - offset += primes[layer][column]; + pub fn compressed(self: *const StateImpl, id: u32) !u32 { + if (id >= 129280) return error.InvalidToken; + return word(self.map, @as(usize, id) * 4); } - } - try projection.apply(.{ .trellis = e.wkv }, &self.gathered, &self.kv, scratch); - for (0..4) |copy| { - var hh: f64 = 0; - var kk: f64 = 0; - var dot: f64 = 0; - for (0..5120) |d| { - const at = copy * 5120 + d; - const x = m.f(h[at]); - const k = m.f(self.kv[at]); - const weight = m.value(e.q_weight, at) * m.value(e.k_weight, at); - hh += @as(f64, x) * x; - kk += @as(f64, k) * k; - const weighted: f32 = (x * weight) * k; - dot += weighted; + pub fn advance(self: *StateImpl, id: u32) void { + self.history[2] = self.history[1]; + self.history[1] = self.history[0]; + self.history[0] = id; } - const h_mean = m.meanFromWide(hh, 5120); - const k_mean = m.meanFromWide(kk, 5120); - const h_inverse = m.idealRsqrt(h_mean + c.rms_norm_eps); - const k_inverse = m.idealRsqrt(k_mean + c.rms_norm_eps); - const rstd: f32 = h_inverse * k_inverse; - const dot_f32: f32 = @floatCast(dot); - const normalized: f32 = (dot_f32 * rstd) * (1 / @sqrt(@as(f32, 5120))); - const gate = m.sigmoid(std.math.copysign(@sqrt(@max(@abs(normalized), 1e-6)), normalized)); - for (0..5120) |d| { - const at = copy * 5120 + d; - h[at] = m.b(m.f(h[at]) + gate * m.f(self.kv[20480 + d])); + pub fn apply(self: *StateImpl, e: Engram, layer: usize, id: u32, h: []u16, c: Config, scratch: *projection.Scratch) !void { + @setFloatMode(.strict); + const tokens = [4]u32{ id, self.history[0], self.history[1], self.history[2] }; + var rolling = @as(u64, tokens[0]) * multipliers[layer][0]; + var offset: usize = 0; + for (1..4) |lookback| { + rolling ^= @as(u64, tokens[lookback]) * multipliers[layer][lookback]; + for (0..8) |head| { + const column = (lookback - 1) * 8 + head; + const row = offset + @as(usize, @intCast(rolling % primes[layer][column])); + for (0..256) |d| self.gathered[column * 256 + d] = + m.b(m.value(e.table, row * 256 + d) * m.value(e.scales, row * 8 + d / 32)); + offset += primes[layer][column]; + } + } + try proj.applyTrellis(e.wkv, &self.gathered, &self.kv, scratch); + for (0..4) |copy| { + var hh: f64 = 0; + var kk: f64 = 0; + var dot: f64 = 0; + for (0..5120) |d| { + const at = copy * 5120 + d; + const x = m.f(h[at]); + const k = m.f(self.kv[at]); + const weight: Real = m.value(e.q_weight, at) * m.value(e.k_weight, at); + hh += @as(f64, x) * x; + kk += @as(f64, k) * k; + const weighted: Real = (x * weight) * k; + dot += @as(f64, weighted); + } + const h_mean = m.meanFromWide(hh, 5120); + const k_mean = m.meanFromWide(kk, 5120); + const h_inverse = m.idealRsqrt(h_mean + m.wr(c.rms_norm_eps)); + const k_inverse = m.idealRsqrt(k_mean + m.wr(c.rms_norm_eps)); + const rstd: Real = h_inverse * k_inverse; + const dot_wide: Real = @floatCast(dot); + const normalized: Real = (dot_wide * rstd) * (1 / @sqrt(@as(Real, 5120))); + const gate = m.sigmoid(std.math.copysign(@sqrt(@max(@abs(normalized), @as(Real, 1e-6))), normalized)); + for (0..5120) |d| { + const at = copy * 5120 + d; + h[at] = m.b(m.f(h[at]) + gate * m.f(self.kv[20480 + d])); + } + } } - } - } -}; + }; + pub const State = StateImpl; + }; +} + +/// The retained FP32 instance; every field the caller touches today keeps its name and type. +pub const f32_module = Module(f32); +pub const State = f32_module.State; diff --git a/zig/tests/deepseek_v41_reference/math.zig b/zig/tests/deepseek_v41_reference/math.zig index a6dccb2b93..55368feca3 100644 --- a/zig/tests/deepseek_v41_reference/math.zig +++ b/zig/tests/deepseek_v41_reference/math.zig @@ -1,224 +1,372 @@ -//! Independent CPU equations: DeepSeek AI inference/{model,kernel}.py, rev 2cba9e42aa026125f3ed06c6d98c1db82f7ca027. +//! Independent CPU equations from DeepSeek AI inference/{model,kernel}.py, rev 2cba9e42; Module(Real) runs FP32/FP64. + const std = @import("std"); const tf = @import("tensorfold"); const Tensor = tf.checkpoint_metal.Tensor; const Hc = tf.deepseek_v41_weights.Hc; const Config = tf.deepseek_v41_config.Text; -pub fn f(x: u16) f32 { - return @bitCast(@as(u32, x) << 16); -} -pub fn b(x: f32) u16 { - const bits: u32 = @bitCast(x); - if ((bits & 0x7fffffff) > 0x7f800000) return @truncate((bits >> 16) | 0x40); - return @truncate((bits +% 0x7fff +% ((bits >> 16) & 1)) >> 16); -} -pub fn bytes(t: Tensor) []const u8 { - return t.buffer.contents()[t.offset..][0..t.bytes]; -} -pub fn value(t: Tensor, index: usize) f32 { - const raw = bytes(t); - return switch (t.dtype) { - .bf16 => f(std.mem.readInt(u16, raw[index * 2 ..][0..2], .little)), - .f16 => @floatCast(@as(f16, @bitCast(std.mem.readInt(u16, raw[index * 2 ..][0..2], .little)))), - .f32 => @bitCast(std.mem.readInt(u32, raw[index * 4 ..][0..4], .little)), - .f8_e4m3 => e4(raw[index]), - .f8_e8m0 => e8(raw[index]), - else => unreachable, - }; -} -pub fn e4(code: u8) f32 { - const magnitude = code & 127; - if (magnitude == 127) return std.math.nan(f32); - const exponent = magnitude >> 3; - const mantissa = magnitude & 7; - const power: f32 = @bitCast((@as(u32, exponent) + 120) << 23); - const v: f32 = if (exponent == 0) @as(f32, @floatFromInt(mantissa)) / 512 else power * (1 + @as(f32, @floatFromInt(mantissa)) / 8); - return if ((code & 128) != 0) -v else v; -} -pub fn e8(code: u8) f32 { - if (code == 255) return std.math.nan(f32); - if (code == 0) return 0x1p-127; - return @bitCast(@as(u32, code) << 23); -} -pub fn roundE4(x: f32) f32 { - const a = @min(@abs(x), 448); - var low: u8 = 0; - var high: u8 = 126; - while (low < high) { - const mid = low + (high - low + 1) / 2; - if (e4(mid) <= a) low = mid else high = mid - 1; - } - var best = low; - if (low < 126) { - const below = a - e4(low); - const above = e4(low + 1) - a; - if (above < below or (above == below and (low + 1) % 2 == 0)) best += 1; - } - return std.math.copysign(e4(best), x); -} -pub fn roundE2(x: f32) f32 { - const levels = [_]f32{ 0, 0.5, 1, 1.5, 2, 3, 4, 6 }; - var best: usize = 0; - var distance: f32 = std.math.inf(f32); - for (levels, 0..) |v, i| { - const d = @abs(@abs(x) - v); - if (d < distance or (d == distance and i % 2 == 0)) { - best = i; - distance = d; +pub fn Module(comptime Real: type) type { + return struct { + const Self = @This(); + + /// Widening-only conversion: exact for f32 values, identity when widths already match. + pub fn wr(x: f32) Real { + return if (Real == f32) x else @floatCast(x); } - } - return std.math.copysign(levels[best], x); + pub fn f(x: u16) Real { + const decoded: f32 = @bitCast(@as(u32, x) << 16); + return if (Real == f32) decoded else @floatCast(decoded); + } + pub fn b(x: Real) u16 { + if (Real == f32) { + const bits: u32 = @bitCast(x); + if ((bits & 0x7fffffff) > 0x7f800000) return @truncate((bits >> 16) | 0x40); + return @truncate((bits +% 0x7fff +% ((bits >> 16) & 1)) >> 16); + } + const bits: u64 = @bitCast(x); + const sign: u16 = @as(u16, @truncate(bits >> 63)) << 15; + const magnitude: u64 = bits & 0x7fff_ffffffffffff; + if (magnitude > 0x7ff0_0000_0000_0000) + return sign | @as(u16, @truncate(((bits >> 49) | 0x0040) & 0xffff)); + if (magnitude == 0x7ff0_0000_0000_0000) return sign | 0x7f80; + const biased: u11 = @truncate(magnitude >> 52); + const biased_wide: i32 = @intCast(biased); + const exponent: i32 = biased_wide - 1023; + if (exponent < -126) { + // Subnormal: grid units full << (exponent+81) on the BF16 grid, spacing 2^-133. + if (exponent <= -1023) return sign; // FP64 subnormal magnitudes round to zero. + const full: u64 = (1 << 52) | (magnitude & 0xfffffffffffff); + const right_wide: i32 = -(exponent + 81); // Positive because exponent < -126. + if (right_wide >= 64) return sign; // Below half the smallest subnormal. + const right: u6 = @intCast(right_wide); + const low = full & ((@as(u64, 1) << right) - 1); + const half = @as(u64, 1) << (right - 1); + var q = full >> right; + if (low > half or (low == half and (q & 1) != 0)) q += 1; + if (q >= 128) return sign | 0x0080; // Rounding reached the minimum normal. + return sign | @as(u16, @truncate(q)); + } + const full: u64 = (1 << 52) | (magnitude & 0xfffffffffffff); + const low = full & ((@as(u64, 1) << 45) - 1); + const half = @as(u64, 1) << 44; + var mantissa = full >> 45; + if (low > half or (low == half and (mantissa & 1) != 0)) mantissa += 1; + var exponent_field = exponent + 127; + if (mantissa == 256) { + mantissa = 128; + exponent_field += 1; + } + if (exponent_field >= 255) return sign | 0x7f80; + return sign | (@as(u16, @intCast(exponent_field)) << 7) | @as(u16, @truncate(mantissa - 128)); + } + pub fn bytes(t: Tensor) []const u8 { + return t.buffer.contents()[t.offset..][0..t.bytes]; + } + pub fn value(t: Tensor, index: usize) Real { + const raw = Self.bytes(t); + return switch (t.dtype) { + .bf16 => Self.f(std.mem.readInt(u16, raw[index * 2 ..][0..2], .little)), + .f16 => @floatCast(@as(f16, @bitCast(std.mem.readInt(u16, raw[index * 2 ..][0..2], .little)))), + .f32 => blk: { + const decoded: f32 = @bitCast(std.mem.readInt(u32, raw[index * 4 ..][0..4], .little)); + break :blk if (Real == f32) decoded else @floatCast(decoded); + }, + .f8_e4m3 => Self.e4(raw[index]), + .f8_e8m0 => Self.e8(raw[index]), + else => unreachable, + }; + } + pub fn e4(code: u8) Real { + const magnitude = code & 127; + if (magnitude == 127) return @as(Real, std.math.nan(f32)); + const exponent = magnitude >> 3; + const mantissa = magnitude & 7; + const power: f32 = @bitCast((@as(u32, exponent) + 120) << 23); + const decoded: Real = if (exponent == 0) + @as(Real, @floatFromInt(mantissa)) / @as(Real, 512) + else + Self.wr(power) * (1 + @as(Real, @floatFromInt(mantissa)) / @as(Real, 8)); + return if ((code & 128) != 0) -decoded else decoded; + } + pub fn e8(code: u8) Real { + if (code == 255) return @as(Real, std.math.nan(f32)); + if (code == 0) return @as(Real, 0x1p-127); + const decoded: f32 = @bitCast(@as(u32, code) << 23); + return if (Real == f32) decoded else @floatCast(decoded); + } + pub fn roundE4(x: Real) Real { + const a = @min(@abs(x), @as(Real, 448)); + var low: u8 = 0; + var high: u8 = 126; + while (low < high) { + const mid = low + (high - low + 1) / 2; + if (Self.e4(mid) <= a) low = mid else high = mid - 1; + } + var best = low; + if (low < 126) { + const below = a - Self.e4(low); + const above = Self.e4(low + 1) - a; + if (above < below or (above == below and (low + 1) % 2 == 0)) best += 1; + } + return std.math.copysign(Self.e4(best), x); + } + pub fn roundE2(x: Real) Real { + const levels = [_]Real{ 0, 0.5, 1, 1.5, 2, 3, 4, 6 }; + var best: usize = 0; + var distance: Real = std.math.inf(Real); + for (levels, 0..) |v, i| { + const d = @abs(@abs(x) - v); + if (d < distance or (d == distance and i % 2 == 0)) { + best = i; + distance = d; + } + } + return std.math.copysign(levels[best], x); + } + fn powerScale(x: Real) Real { + if (Real == f32) { + const bits: u32 = @bitCast(x); + const exponent = (bits >> 23) & 255; + const rounded = exponent + @as(u32, @intFromBool((bits & 0x7fffff) != 0)); + return @bitCast(rounded << 23); + } + const bits: u64 = @bitCast(x); + const exponent = (bits >> 52) & 2047; + const rounded = exponent + @as(u64, @intFromBool((bits & 0xfffffffffffff) != 0)); + return @bitCast(rounded << 52); + } + pub const Quant = enum { window, main, index }; + pub fn quantize(x: []u16, kind: Self.Quant) void { + const group: usize = if (kind == .main) 16 else 32; + var start: usize = 0; + while (start < x.len) : (start += group) { + var a: Real = if (kind == .window) @as(Real, 1e-4) else if (kind == .main) Self.wr(6 * 0x1p-9) else Self.wr(6 * 0x1p-126); + for (x[start..][0..group]) |v| a = @max(a, @abs(Self.f(v))); + const scale = if (kind == .main) + Self.roundE4(a / @as(Real, 6)) + else + Self.powerScale(a * (if (kind == .window) @as(Real, 1.0 / 448.0) else @as(Real, 1.0 / 6.0))); + for (x[start..][0..group]) |*v| { + const z = Self.f(v.*) / scale; + v.* = Self.b((if (kind == .window) Self.roundE4(z) else Self.roundE2(@max(-6, @min(6, z)))) * scale); + } + } + } + pub fn dense(t: Tensor, x: []const u16, out: []Real) void { + for (out, 0..) |*y, row| { + var sum: f64 = 0; + for (x, 0..) |v, k| sum += @as(f64, Self.f(v)) * @as(f64, Self.value(t, row * x.len + k)); + y.* = @floatCast(sum); + } + } + /// CUDA MeanOps scalar projection; the internal reduction remains independently wide. + pub fn meanFromWide(sum: f64, count: usize) Real { + @setFloatMode(.strict); + const accumulator: Real = @floatCast(sum); + const factor: Real = 1 / @as(Real, @floatFromInt(count)); + return accumulator * factor; + } + /// Independent ideal rounding in the module's width, not an emulation of CUDA's rsqrt approximation. + pub fn idealRsqrt(x: Real) Real { + @setFloatMode(.strict); + return @floatCast(1 / @sqrt(@as(f64, x))); + } + /// The RMS inverse in the module's width: FP32 keeps the CUDA mean projection, FP64 the wide ideal. + fn inverseRms(x: []const u16, eps: f32) Real { + @setFloatMode(.strict); + var sum: f64 = 0; + for (x) |v| sum += @as(f64, Self.f(v)) * Self.f(v); + if (Real == f32) return Self.idealRsqrt(Self.meanFromWide(sum, x.len) + eps); + return 1 / @sqrt(sum / @as(f64, @floatFromInt(x.len)) + @as(f64, eps)); + } + pub fn rms(x: []const u16, weight: Tensor, eps: f32, out: []u16) void { + @setFloatMode(.strict); + const r = Self.inverseRms(x, eps); + for (x, out, 0..) |v, *y, i| y.* = Self.b(Self.f(v) * r * Self.value(weight, i)); + } + /// The same boundary with weight cells explicit, so a harness needs no checkpoint tensor. + pub fn rmsValues(x: []const u16, weight: []const Real, eps: f32, out: []u16) void { + @setFloatMode(.strict); + const r = Self.inverseRms(x, eps); + for (x, weight, 0..) |v, w, i| out[i] = Self.b(Self.f(v) * r * w); + } + pub fn sigmoid(x: Real) Real { + return 1 / (1 + @exp(-x)); + } + pub const Mix = struct { pre: [4]Real, post: [4]Real, comb: [16]Real }; + /// The HC RMS statistic behind the mix: FP32 keeps the CUDA mean projection, FP64 the wide ideal. + pub fn rmsInverse(x: []const u16, eps: f32) Real { + @setFloatMode(.strict); + var squares: f64 = 0; + for (x) |v| squares += @as(f64, Self.f(v)) * Self.f(v); + if (Real == f32) return Self.idealRsqrt(Self.meanFromWide(squares, x.len) + eps); + return 1 / @sqrt(squares / @as(f64, @floatFromInt(x.len)) + @as(f64, eps)); + } + pub fn mixes(x: []const u16, hc: Hc, c: Config) Self.Mix { + var projected: [24]Real = undefined; + Self.dense(hc.projection, x, &projected); + return Self.split(&projected, Self.rmsInverse(x, c.rms_norm_eps), hc, c); + } + /// The affine, softmax and Sinkhorn coefficient stage at the module's width. + pub fn split(projected: []const Real, r: Real, hc: Hc, c: Config) Self.Mix { + @setFloatMode(.strict); + var m: Self.Mix = undefined; + for (0..4) |i| { + m.pre[i] = Self.sigmoid(@mulAdd(Real, projected[i] * r, Self.value(hc.scale, 0), Self.value(hc.base, i))) + Self.wr(c.hc_eps); + m.post[i] = 2 * Self.sigmoid(@mulAdd(Real, projected[i + 4] * r, Self.value(hc.scale, 1), Self.value(hc.base, i + 4))); + } + for (0..4) |i| { + var high: Real = -std.math.inf(Real); + for (0..4) |j| { + const at = i * 4 + j; + m.comb[at] = @mulAdd(Real, projected[at + 8] * r, Self.value(hc.scale, 2), Self.value(hc.base, at + 8)); + high = @max(high, m.comb[at]); + } + for (0..4) |j| { + m.comb[i * 4 + j] = @exp(m.comb[i * 4 + j] - high); + } + // TileLang 0.1.8 emits FP32 XOR reductions with offsets 2, then 1. + const denominator = (m.comb[i * 4] + m.comb[i * 4 + 2]) + + (m.comb[i * 4 + 1] + m.comb[i * 4 + 3]); + for (0..4) |j| m.comb[i * 4 + j] = m.comb[i * 4 + j] / denominator + Self.wr(c.hc_eps); + } + for (0..c.hc_sinkhorn_iters) |iteration| { + if (iteration != 0) for (0..4) |i| { + const denominator = ((m.comb[i * 4] + m.comb[i * 4 + 2]) + + (m.comb[i * 4 + 1] + m.comb[i * 4 + 3])) + Self.wr(c.hc_eps); + for (0..4) |j| m.comb[i * 4 + j] /= denominator; + }; + for (0..4) |j| { + const denominator = ((m.comb[j] + m.comb[8 + j]) + + (m.comb[4 + j] + m.comb[12 + j])) + Self.wr(c.hc_eps); + for (0..4) |i| m.comb[i * 4 + j] /= denominator; + } + } + return m; + } + pub fn collapse(h: []const u16, pre: [4]Real, out: []u16) void { + @setFloatMode(.strict); + for (out, 0..) |*v, d| { + var sum: Real = 0; + for (0..4) |i| { + const product: Real = pre[i] * Self.f(h[i * out.len + d]); + sum += product; + } + v.* = Self.b(sum); + } + } + pub fn expand(h: []u16, x: []const u16, m: Self.Mix) void { + @setFloatMode(.strict); + for (0..x.len) |d| { + var old: [4]Real = undefined; + for (0..4) |i| old[i] = Self.f(h[i * x.len + d]); + for (0..4) |i| { + var residual: Real = 0; + for (0..4) |j| { + const product: Real = m.comb[j * 4 + i] * old[j]; + residual += product; + } + h[i * x.len + d] = Self.b(m.post[i] * Self.f(x[d]) + residual); + } + } + } + pub fn ropeFrequency(index: usize, theta: f32, original: u32, factor: f32, beta_fast: f32, beta_slow: f32) Real { + @setFloatMode(.strict); + const exponent: Real = @as(Real, @floatFromInt(index * 2)) / 64; + const power: Real = @floatCast(std.math.pow(f64, theta, @floatCast(exponent))); + var frequency: Real = 1 / power; + if (original > 0) { + const original_f64: f64 = @floatFromInt(original); + const theta_f64: f64 = theta; + const low: Real = @floatCast(@max(@floor(64 * @log(original_f64 / + (@as(f64, beta_fast) * 2 * std.math.pi)) / (2 * @log(theta_f64))), 0)); + const high: Real = @floatCast(@min(@ceil(64 * @log(original_f64 / + (@as(f64, beta_slow) * 2 * std.math.pi)) / (2 * @log(theta_f64))), 63)); + const ramp: Real = @max(0, @min(1, (@as(Real, @floatFromInt(index)) - low) / @max(high - low, Self.wr(1e-3)))); + const smooth: Real = 1 - ramp; + frequency = frequency / Self.wr(factor) * (1 - smooth) + frequency * smooth; + } + return frequency; + } + pub fn rope(x: []u16, position: usize, compressed: bool, inverse: bool, c: Config) void { + for (0..32) |i| { + const frequency = Self.ropeFrequency(i, if (compressed) c.compress_rope_theta else c.rope_theta, if (compressed) c.rope_scaling.original_max_position_embeddings else 0, c.rope_scaling.factor, c.rope_scaling.beta_fast, c.rope_scaling.beta_slow); + const angle: Real = @as(Real, @floatFromInt(position)) * frequency; + const cosine = @cos(angle); + const sine = @sin(angle) * (if (inverse) @as(Real, -1) else @as(Real, 1)); + const at = x.len - 64 + 2 * i; + const a = Self.f(x[at]); + const z = Self.f(x[at + 1]); + x[at] = Self.b(a * cosine - z * sine); + x[at + 1] = Self.b(a * sine + z * cosine); + } + } + }; } -fn powerScale(x: f32) f32 { - const bits: u32 = @bitCast(x); - const exponent = (bits >> 23) & 255; - const rounded = exponent + @as(u32, @intFromBool((bits & 0x7fffff) != 0)); - return @bitCast(rounded << 23); + +const f32_module = Module(f32); +pub const wr = f32_module.wr; +pub const f = f32_module.f; +pub const b = f32_module.b; +pub const bytes = f32_module.bytes; +pub const value = f32_module.value; +pub const e4 = f32_module.e4; +pub const e8 = f32_module.e8; +pub const roundE4 = f32_module.roundE4; +pub const roundE2 = f32_module.roundE2; +pub const Quant = f32_module.Quant; +pub const quantize = f32_module.quantize; +pub const dense = f32_module.dense; +pub const meanFromWide = f32_module.meanFromWide; +pub const idealRsqrt = f32_module.idealRsqrt; +pub const rms = f32_module.rms; +pub const sigmoid = f32_module.sigmoid; +pub const Mix = f32_module.Mix; +pub const mixes = f32_module.mixes; +pub const collapse = f32_module.collapse; +pub const expand = f32_module.expand; +pub const ropeFrequency = f32_module.ropeFrequency; +pub const rope = f32_module.rope; + +/// Explicit-instantiation forwarders for the operator harnesses: the same Module bodies, never a second copy. +pub fn MixOf(comptime T: type) type { + return Module(T).Mix; } -pub const Quant = enum { window, main, index }; -pub fn quantize(x: []u16, kind: Quant) void { - const group: usize = if (kind == .main) 16 else 32; - var start: usize = 0; - while (start < x.len) : (start += group) { - var a: f32 = if (kind == .window) 1e-4 else if (kind == .main) 6 * 0x1p-9 else 6 * 0x1p-126; - for (x[start..][0..group]) |v| a = @max(a, @abs(f(v))); - const scale = if (kind == .main) roundE4(a / 6) else powerScale(a * (if (kind == .window) @as(f32, 1.0 / 448.0) else @as(f32, 1.0 / 6.0))); - for (x[start..][0..group]) |*v| { - const z = f(v.*) / scale; - v.* = b((if (kind == .window) roundE4(z) else roundE2(@max(-6, @min(6, z)))) * scale); - } - } +pub fn denseOf(comptime T: type, t: Tensor, x: []const u16, out: []T) void { + Module(T).dense(t, x, out); } -pub fn dense(t: Tensor, x: []const u16, out: []f32) void { - for (out, 0..) |*y, row| { - var sum: f64 = 0; - for (x, 0..) |v, k| sum += @as(f64, f(v)) * value(t, row * x.len + k); - y.* = @floatCast(sum); - } +pub fn inverseOf(comptime T: type, x: []const u16, eps: f32) T { + return Module(T).rmsInverse(x, eps); } -/// CUDA MeanOps scalar projection; the internal reduction remains independently wide. -pub fn meanFromWide(sum: f64, count: usize) f32 { - @setFloatMode(.strict); - const accumulator: f32 = @floatCast(sum); - const factor: f32 = 1 / @as(f32, @floatFromInt(count)); - return accumulator * factor; +pub fn mixesOf(comptime T: type, x: []const u16, hc: Hc, c: Config) MixOf(T) { + return Module(T).mixes(x, hc, c); } -/// Independent ideal FP32 unary result, not an emulation of CUDA's rsqrt approximation. -pub fn idealRsqrt(x: f32) f32 { - @setFloatMode(.strict); - return @floatCast(1 / @sqrt(@as(f64, x))); +pub fn splitOf(comptime T: type, projected: []const T, r: T, hc: Hc, c: Config) MixOf(T) { + return Module(T).split(projected, r, hc, c); } -pub fn rms(x: []const u16, weight: Tensor, eps: f32, out: []u16) void { - @setFloatMode(.strict); - var sum: f64 = 0; - for (x) |v| sum += @as(f64, f(v)) * f(v); - const mean = meanFromWide(sum, x.len); - const r = idealRsqrt(mean + eps); - for (x, out, 0..) |v, *y, i| y.* = b(f(v) * r * value(weight, i)); +pub fn collapseOf(comptime T: type, h: []const u16, pre: [4]T, out: []u16) void { + Module(T).collapse(h, pre, out); } -pub fn sigmoid(x: f32) f32 { - return 1 / (1 + @exp(-x)); +pub fn expandOf(comptime T: type, h: []u16, x: []const u16, m: MixOf(T)) void { + Module(T).expand(h, x, m); } -pub const Mix = struct { pre: [4]f32, post: [4]f32, comb: [16]f32 }; -pub fn mixes(x: []const u16, hc: Hc, c: Config) Mix { - @setFloatMode(.strict); - var projected: [24]f32 = undefined; - dense(hc.projection, x, &projected); - var squares: f64 = 0; - for (x) |v| squares += @as(f64, f(v)) * f(v); - const mean = meanFromWide(squares, x.len); - const r = idealRsqrt(mean + c.rms_norm_eps); - var m: Mix = undefined; - for (0..4) |i| { - m.pre[i] = sigmoid(@mulAdd(f32, projected[i] * r, value(hc.scale, 0), value(hc.base, i))) + c.hc_eps; - m.post[i] = 2 * sigmoid(@mulAdd(f32, projected[i + 4] * r, value(hc.scale, 1), value(hc.base, i + 4))); - } - for (0..4) |i| { - var high: f32 = -std.math.inf(f32); - for (0..4) |j| { - const at = i * 4 + j; - m.comb[at] = @mulAdd(f32, projected[at + 8] * r, value(hc.scale, 2), value(hc.base, at + 8)); - high = @max(high, m.comb[at]); - } - for (0..4) |j| { - m.comb[i * 4 + j] = @exp(m.comb[i * 4 + j] - high); - } - // TileLang 0.1.8 emits FP32 XOR reductions with offsets 2, then 1. - const denominator = (m.comb[i * 4] + m.comb[i * 4 + 2]) + - (m.comb[i * 4 + 1] + m.comb[i * 4 + 3]); - for (0..4) |j| m.comb[i * 4 + j] = m.comb[i * 4 + j] / denominator + c.hc_eps; - } - for (0..c.hc_sinkhorn_iters) |iteration| { - if (iteration != 0) for (0..4) |i| { - const denominator = ((m.comb[i * 4] + m.comb[i * 4 + 2]) + - (m.comb[i * 4 + 1] + m.comb[i * 4 + 3])) + c.hc_eps; - for (0..4) |j| m.comb[i * 4 + j] /= denominator; - }; - for (0..4) |j| { - const denominator = ((m.comb[j] + m.comb[8 + j]) + - (m.comb[4 + j] + m.comb[12 + j])) + c.hc_eps; - for (0..4) |i| m.comb[i * 4 + j] /= denominator; - } - } - return m; +pub fn rmsOf(comptime T: type, x: []const u16, weight: Tensor, eps: f32, out: []u16) void { + Module(T).rms(x, weight, eps, out); } -pub fn collapse(h: []const u16, pre: [4]f32, out: []u16) void { - @setFloatMode(.strict); - for (out, 0..) |*v, d| { - var sum: f32 = 0; - for (0..4) |i| { - const product: f32 = pre[i] * f(h[i * out.len + d]); - sum += product; - } - v.* = b(sum); - } +pub fn rmsOfValues(comptime T: type, x: []const u16, weight: []const f32, eps: f32, out: []u16) void { + const M = Module(T); + const r = M.rmsInverse(x, eps); + for (x, weight, 0..) |v, w, i| out[i] = M.b(M.f(v) * r * M.wr(w)); } -pub fn expand(h: []u16, x: []const u16, m: Mix) void { - @setFloatMode(.strict); - for (0..x.len) |d| { - var old: [4]f32 = undefined; - for (0..4) |i| old[i] = f(h[i * x.len + d]); - for (0..4) |i| { - var residual: f32 = 0; - for (0..4) |j| { - const product: f32 = m.comb[j * 4 + i] * old[j]; - residual += product; - } - h[i * x.len + d] = b(m.post[i] * f(x[d]) + residual); - } - } +pub fn wide(comptime T: type, v: f32) T { + return Module(T).wr(v); } -pub fn ropeFrequency(index: usize, theta: f32, original: u32, factor: f32, beta_fast: f32, beta_slow: f32) f32 { - @setFloatMode(.strict); - const exponent: f32 = @as(f32, @floatFromInt(index * 2)) / 64; - const power: f32 = @floatCast(std.math.pow(f64, theta, exponent)); - var frequency: f32 = 1 / power; - if (original > 0) { - const original_f64: f64 = @floatFromInt(original); - const theta_f64: f64 = theta; - const low: f32 = @floatCast(@max(@floor(64 * @log(original_f64 / - (@as(f64, beta_fast) * 2 * std.math.pi)) / (2 * @log(theta_f64))), 0)); - const high: f32 = @floatCast(@min(@ceil(64 * @log(original_f64 / - (@as(f64, beta_slow) * 2 * std.math.pi)) / (2 * @log(theta_f64))), 63)); - const ramp: f32 = @max(0, @min(1, (@as(f32, @floatFromInt(index)) - low) / @max(high - low, 1e-3))); - const smooth: f32 = 1 - ramp; - frequency = frequency / factor * (1 - smooth) + frequency * smooth; - } - return frequency; +pub fn b64(x: f64) u16 { + return Module(f64).b(x); } -pub fn rope(x: []u16, position: usize, compressed: bool, inverse: bool, c: Config) void { - for (0..32) |i| { - const frequency = ropeFrequency(i, if (compressed) c.compress_rope_theta else c.rope_theta, if (compressed) c.rope_scaling.original_max_position_embeddings else 0, c.rope_scaling.factor, c.rope_scaling.beta_fast, c.rope_scaling.beta_slow); - const angle: f32 = @as(f32, @floatFromInt(position)) * frequency; - const cosine = @cos(angle); - const sine = @sin(angle) * (if (inverse) @as(f32, -1) else @as(f32, 1)); - const at = x.len - 64 + 2 * i; - const a = f(x[at]); - const z = f(x[at + 1]); - x[at] = b(a * cosine - z * sine); - x[at + 1] = b(a * sine + z * cosine); - } +pub fn d64(x: u16) f64 { + return Module(f64).f(x); } test "BF16 probability ties and FP4 ties retain source boundaries" { @@ -260,6 +408,24 @@ test "hyperconnection rounds elementwise products before cancelling reductions" try std.testing.expectEqual(b(0x1p-21), h[0]); } +test "FP64 rounding-to-BF16 matches the FP32 single-rounded path across the grid" { + const R64 = Module(f64); + var bits: u32 = 0; + while (bits < 0x0004_0000) : (bits += 1) { + const x: f32 = @bitCast((bits << 8) + 0x3c000000); + try std.testing.expectEqual(b(x), R64.b(@as(f64, x))); + } + var shift: u32 = 1; + while (shift < 255) : (shift += 1) { + const pattern: u32 = shift << 23; + for ([_]u32{ 0, 1, 0x7ffff, 0x80000, 0x7fffff }) |mantissa| { + const x: f32 = @bitCast(pattern | mantissa); + try std.testing.expectEqual(b(x), R64.b(@as(f64, x))); + try std.testing.expectEqual(b(-x), R64.b(-@as(f64, x))); + } + } +} + test "RoPE and YaRN frequency vectors retain official FP32 elementwise boundaries" { const parsed = try std.json.parseFromSlice(std.json.Value, std.testing.allocator, @embedFile("rope_frequencies.json"), .{}); defer parsed.deinit(); diff --git a/zig/tests/deepseek_v41_reference/precision.zig b/zig/tests/deepseek_v41_reference/precision.zig new file mode 100644 index 0000000000..39fcb29640 --- /dev/null +++ b/zig/tests/deepseek_v41_reference/precision.zig @@ -0,0 +1,106 @@ +//! Gate-deciding sensitivity: FP32 and FP64 run the same equations on identical tokens and checkpoint bytes. + +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const Model = tf.deepseek_v41.Model; +const reference = @import("reference.zig"); + +/// The caller writes the logits after its residency counters: artifact metadata reads are not execution. +pub const Logits = struct { + fp32: []f32, + fp64: []f64, + /// Per-token (FP32, FP64) argmax. The caller prints it after its residency counters. + argmax: [][2]u32, + pub fn deinit(self: Logits, gpa: std.mem.Allocator) void { + gpa.free(self.fp32); + gpa.free(self.fp64); + gpa.free(self.argmax); + } +}; + +pub const Report = struct { + result: Result, + logits: Logits, +}; + +pub const Result = struct { + tokens: usize, + max_absolute_error: f64, + rms_error: f64, + normalized_rms_error: f64, // sqrt(sum((fp32-fp64)^2)/sum(fp64^2)), the same shape as the whole-model gate + normalized_rms_error_against_fp32: f64, + greedy_disagreements: usize, + fp32_seconds: f64, + fp64_seconds: f64, + evidence_origin: []const u8 = "same_checkpoint_same_equations_two_widths_no_gpu", +}; + +fn argmaxF32(row: []const f32) usize { + var best: usize = 0; + var score: f32 = -std.math.inf(f32); + for (row, 0..) |v, i| + if (v > score) { + best = i; + score = v; + }; + return best; +} + +fn argmaxF64(row: []const f64) usize { + var best: usize = 0; + var score: f64 = -std.math.inf(f64); + for (row, 0..) |v, i| + if (v > score) { + best = i; + score = v; + }; + return best; +} + +pub fn run(gpa: std.mem.Allocator, model: *Model, ids: []const u32, token_map: []const u8) !Report { + if (ids.len == 0 or ids.len > 16) return error.InvalidOracleWindow; + const F32 = reference.f32_reference.Reference; + const F64 = reference.Module(f64).Reference; + var fp32_ref = try F32.init(gpa, model.cfg.text_config, &model.weights, &model.checkpoint, token_map, @intCast(ids.len)); + defer fp32_ref.deinit(); + var fp64_ref = try F64.init(gpa, model.cfg.text_config, &model.weights, &model.checkpoint, token_map, @intCast(ids.len)); + defer fp64_ref.deinit(); + const fp32_logits = try gpa.alloc(f32, ids.len * 129280); + errdefer gpa.free(fp32_logits); + const fp64_logits = try gpa.alloc(f64, ids.len * 129280); + errdefer gpa.free(fp64_logits); + const argmax = try gpa.alloc([2]u32, ids.len); + errdefer gpa.free(argmax); + var result = Result{ .tokens = ids.len, .max_absolute_error = 0, .rms_error = 0, .normalized_rms_error = 0, .normalized_rms_error_against_fp32 = 0, .greedy_disagreements = 0, .fp32_seconds = 0, .fp64_seconds = 0 }; + var squared: f64 = 0; + var reference_f64_squared: f64 = 0; + var reference_f32_squared: f64 = 0; + for (ids, 0..) |id, i| { + const f32_start = mtl.clock.seconds(); + const expected32 = try fp32_ref.step(id); + result.fp32_seconds += mtl.clock.seconds() - f32_start; + const f64_start = mtl.clock.seconds(); + const expected64 = try fp64_ref.step(id); + result.fp64_seconds += mtl.clock.seconds() - f64_start; + const wide32 = fp32_logits[i * 129280 ..][0..129280]; + const wide64 = fp64_logits[i * 129280 ..][0..129280]; + @memcpy(wide32, expected32); + @memcpy(wide64, expected64); + for (expected32, expected64) |a, b| { + if (!std.math.isFinite(a) or !std.math.isFinite(b)) return error.NonFiniteReferenceLogits; + const delta: f64 = @as(f64, a) - b; + result.max_absolute_error = @max(result.max_absolute_error, @abs(delta)); + squared += delta * delta; + reference_f64_squared += @as(f64, b) * b; + } + argmax[i] = .{ @intCast(argmaxF32(wide32)), @intCast(argmaxF64(wide64)) }; + if (argmax[i][0] != argmax[i][1]) result.greedy_disagreements += 1; + } + const count: f64 = @floatFromInt(fp32_logits.len); + result.rms_error = @sqrt(squared / count); + result.normalized_rms_error = if (reference_f64_squared == 0) result.rms_error else @sqrt(squared / reference_f64_squared); + for (fp32_logits) |v| reference_f32_squared += @as(f64, v) * v; + result.normalized_rms_error_against_fp32 = if (reference_f32_squared == 0) result.rms_error else @sqrt(squared / reference_f32_squared); + return .{ .result = result, .logits = .{ .fp32 = fp32_logits, .fp64 = fp64_logits, .argmax = argmax } }; +} diff --git a/zig/tests/deepseek_v41_reference/projection.zig b/zig/tests/deepseek_v41_reference/projection.zig index 03174c9459..e74e1b8009 100644 --- a/zig/tests/deepseek_v41_reference/projection.zig +++ b/zig/tests/deepseek_v41_reference/projection.zig @@ -1,11 +1,11 @@ -//! Independent scalar EXL3 oracle over resident bytes; FP64 dots and FP32 rotation/output boundaries, no GPU work. +//! Independent scalar EXL3 oracle over resident bytes; Module(Real) removes intermediates, never a storage boundary. // EXL3 attribution: docs/recipes/exl3.md; pinned checkpoint provenance in THIRD_PARTY_NOTICES.md. const std = @import("std"); -const tf = @import("tensorfold"); -const exl3 = tf.exl3_projection; -const weights = tf.deepseek_v41_weights; +const math = @import("math.zig"); +const exl3 = @import("tensorfold").exl3_projection; +const weights = @import("tensorfold").deepseek_v41_weights; const format = exl3.format; const block_size = format.hadamard_dim; const inverse: f32 = 0.08838834764831845; @@ -42,66 +42,6 @@ pub const Scratch = struct { } }; -/// One BF16 input row and BF16 output row; caller must finish GPU writes before borrowing shared CPU bytes. -pub fn apply(matrix: weights.AnyMatrix, x: []const u16, y: []u16, scratch: *Scratch) !void { - switch (matrix) { - .trellis => |m| return applyTrellis(m, x, y, scratch), - .dense => |d| return applyDense(u16, d, x, y, scratch), - } -} - -pub fn applyF32(matrix: weights.AnyMatrix, x: []const u16, y: []f32, scratch: *Scratch) !void { - switch (matrix) { - .trellis => |m| return applyF32Trellis(m, x, y, scratch), - .dense => |d| return applyDense(f32, d, x, y, scratch), - } -} - -/// Kept-dense: quantized rotation and output stages; FP64 dot over the f16-rounded row replaces the tile decode. -fn applyDense(comptime T: type, d: weights.Dense, x: []const u16, y: []T, scratch: *Scratch) !void { - @setFloatMode(.strict); - if (x.len != d.k or y.len != d.n) return error.InvalidReferenceProjectionLength; - if (d.k > scratch.rotated.len or d.n > scratch.sums.len) - return error.ReferenceProjectionScratchTooSmall; - const su = try borrowed(.{ .buffer = d.suh.buffer, .offset = d.suh.offset }, d.k * 2); - const sv = try borrowed(.{ .buffer = d.svh.buffer, .offset = d.svh.offset }, d.n * 2); - const w = try borrowed(.{ .buffer = d.weight.buffer, .offset = d.weight.offset }, d.n * d.k * 2); - const rotated = scratch.rotated[0..d.k]; - const sums = scratch.sums[0..d.n]; - prepare(x, su, .f16, rotated); - for (sums) |*slot| slot.* = 0; - for (0..d.n) |j| { - var acc: f64 = 0; - const wrow = 2 * j * d.k; - for (0..d.k) |k| { - const word = std.mem.readInt(u16, w[wrow + 2 * k ..][0..2], .little); - acc += @as(f64, rotated[k]) * @as(f64, @as(f32, @bitCast(@as(u32, word) << 16))); - } - sums[j] = acc; - } - var base: usize = 0; - while (base < d.n) : (base += block_size) { - var values: [block_size]f32 = undefined; - for (&values, 0..) |*value, i| value.* = @floatCast(sums[base + i]); - hadamard(&values); - for (values, 0..) |value, i| { - const out: f32 = (value * inverse) * half(sv, base + i); - y[base + i] = if (T == u16) bf16(out) else out; - } - } -} - -fn applyTrellis(matrix: exl3.Matrix, x: []const u16, y: []u16, scratch: *Scratch) !void { - try project(matrix, x, y.len, scratch); - try finish(u16, matrix, y, scratch.sums[0..matrix.n]); -} - -/// The same arithmetic with an FP32 output boundary, for compressor projections and full-vocabulary logits. -pub fn applyF32Trellis(matrix: exl3.Matrix, x: []const u16, y: []f32, scratch: *Scratch) !void { - try project(matrix, x, y.len, scratch); - try finish(f32, matrix, y, scratch.sums[0..matrix.n]); -} - fn borrowed(view: exl3.View, bytes: usize) ![]const u8 { const len = view.buffer.length(); if (view.offset % 2 != 0 or view.offset > len or bytes > len - view.offset) @@ -126,47 +66,6 @@ fn scale(bytes: []const u8, i: usize, kind: exl3.ScaleFormat) f32 { return if (((word >> @as(u4, @intCast(i % 16))) & 1) != 0) -1.0 else 1.0; } -fn bf16(value: f32) u16 { - const bits: u32 = @bitCast(value); - // Preserve a quiet NaN instead of letting rounding turn a tiny NaN payload into infinity. - if ((bits & 0x7fffffff) > 0x7f800000) return @as(u16, @truncate(bits >> 16)) | 0x0040; - return @truncate((bits +% (0x7fff + ((bits >> 16) & 1))) >> 16); -} - -/// Seven ordinary Sylvester butterflies, with explicit FP32 stage boundaries and normalization afterward. -fn hadamard(values: *[block_size]f32) void { - @setFloatMode(.strict); - var stride: usize = 1; - while (stride < block_size) : (stride *= 2) { - var base: usize = 0; - while (base < block_size) : (base += 2 * stride) { - for (0..stride) |j| { - const a = values[base + j]; - const b = values[base + stride + j]; - values[base + j] = a + b; - values[base + stride + j] = a - b; - } - } - } -} - -fn prepare(x: []const u16, su: []const u8, kind: exl3.ScaleFormat, rotated: []f64) void { - @setFloatMode(.strict); - var base: usize = 0; - while (base < x.len) : (base += block_size) { - var values: [block_size]f32 = undefined; - for (&values, 0..) |*value, i| { - const input: f32 = @bitCast(@as(u32, x[base + i]) << 16); - value.* = input * scale(su, base + i, kind); - } - hadamard(&values); - for (values, 0..) |value, i| { - const rounded: f16 = @floatCast(value * inverse); - rotated[base + i] = @floatCast(rounded); - } - } -} - /// CPU stream reader consumes MSB-first bits from little-endian int16 pairs, not the GPU indexed decoder. const Stream = struct { bytes: []const u8, @@ -214,59 +113,176 @@ fn accumulateTile(comptime halves: u16, tile: []const u8, table: []const f16, x: } } -fn product(comptime halves: u16, matrix: exl3.Matrix, trellis: []const u8, table: []const f16, rotated: []const f64, sums: []f64) !void { - const ks = try matrix.kStride(); - const ns = matrix.nStride(); - const tile_bytes = @as(usize, halves) * 16; - @memset(sums, 0); - for (0..matrix.k / 16) |tk| { - const input = rotated[tk * 16 ..][0..16]; - for (0..matrix.n / 16) |tn| { - const at = tk * ks + tn * ns; - accumulateTile(halves, trellis[at..][0..tile_bytes], table, input, sums[tn * 16 ..][0..16]); +pub fn Module(comptime Real: type) type { + return struct { + const Self = @This(); + pub const wr = math.Module(Real).wr; + + fn bf16(value: Real) u16 { + if (Real == f32) { + const bits: u32 = @bitCast(value); + // Preserve a quiet NaN instead of letting rounding turn a tiny NaN payload into infinity. + if ((bits & 0x7fffffff) > 0x7f800000) return @as(u16, @truncate(bits >> 16)) | 0x0040; + return @truncate((bits +% (0x7fff + ((bits >> 16) & 1))) >> 16); + } + return math.Module(Real).b(value); + } + + /// Seven ordinary Sylvester butterflies, with explicit stage boundaries and normalization afterward. + fn hadamard(values: *[block_size]Real) void { + @setFloatMode(.strict); + var stride: usize = 1; + while (stride < block_size) : (stride *= 2) { + var base: usize = 0; + while (base < block_size) : (base += 2 * stride) { + for (0..stride) |j| { + const a = values[base + j]; + const b = values[base + stride + j]; + values[base + j] = a + b; + values[base + stride + j] = a - b; + } + } + } } - } -} -fn project(matrix: exl3.Matrix, x: []const u16, output_len: usize, scratch: *Scratch) !void { - try matrix.validate(); - if (x.len != matrix.k or output_len != matrix.n) return error.InvalidReferenceProjectionLength; - if (matrix.k > scratch.rotated.len or matrix.n > scratch.sums.len) - return error.ReferenceProjectionScratchTooSmall; - const trellis = try borrowed(matrix.trellis, try matrix.storageBytes()); - const su = try borrowed(matrix.suh, try scaleBytes(matrix.k, matrix.suh_format)); - // Check every borrowed view before doing arithmetic, including the output-only scale and optional bias. - _ = try borrowed(matrix.svh, try scaleBytes(matrix.n, matrix.svh_format)); - if (matrix.bias) |bias| _ = try borrowed(bias, try std.math.mul(usize, matrix.n, 2)); - const rotated = scratch.rotated[0..matrix.k]; - const sums = scratch.sums[0..matrix.n]; - prepare(x, su, matrix.suh_format, rotated); - const book_index: usize = @intCast(@backingInt(matrix.codebook)); - const table = scratch.codebooks[book_index * states_count ..][0..states_count]; - switch (matrix.width.halves) { - inline 2...16 => |halves| try product(halves, matrix, trellis, table, rotated, sums), - else => unreachable, - } -} + fn prepare(x: []const u16, su: []const u8, kind: exl3.ScaleFormat, rotated: []f64) void { + @setFloatMode(.strict); + var base: usize = 0; + while (base < x.len) : (base += block_size) { + var values: [block_size]Real = undefined; + for (&values, 0..) |*value, i| { + const input: Real = @floatCast(@as(f32, @bitCast(@as(u32, x[base + i]) << 16))); + value.* = input * wr(scale(su, base + i, kind)); + } + hadamard(&values); + for (values, 0..) |value, i| { + const rounded: f16 = @floatCast(value * wr(inverse)); + rotated[base + i] = @floatCast(rounded); + } + } + } -fn finish(comptime T: type, matrix: exl3.Matrix, y: []T, sums: []const f64) !void { - @setFloatMode(.strict); - const sv = try borrowed(matrix.svh, try scaleBytes(matrix.n, matrix.svh_format)); - const bias = if (matrix.bias) |view| try borrowed(view, try std.math.mul(usize, matrix.n, 2)) else null; - var base: usize = 0; - while (base < matrix.n) : (base += block_size) { - var values: [block_size]f32 = undefined; - // One FP64->FP32 boundary replaces GPU FP32 dot-product accumulation, not the output rotation. - for (&values, 0..) |*value, i| value.* = @floatCast(sums[base + i]); - hadamard(&values); - for (values, 0..) |value, i| { - var out: f32 = (value * inverse) * scale(sv, base + i, matrix.svh_format); - if (bias) |bytes| out += half(bytes, base + i); - y[base + i] = if (T == u16) bf16(out) else out; + fn product(comptime halves: u16, matrix: exl3.Matrix, trellis: []const u8, table: []const f16, rotated: []const f64, sums: []f64) !void { + const ks = try matrix.kStride(); + const ns = matrix.nStride(); + const tile_bytes = @as(usize, halves) * 16; + @memset(sums, 0); + for (0..matrix.k / 16) |tk| { + const input = rotated[tk * 16 ..][0..16]; + for (0..matrix.n / 16) |tn| { + const at = tk * ks + tn * ns; + accumulateTile(halves, trellis[at..][0..tile_bytes], table, input, sums[tn * 16 ..][0..16]); + } + } } - } + + fn project(matrix: exl3.Matrix, x: []const u16, output_len: usize, scratch: *Scratch) !void { + try matrix.validate(); + if (x.len != matrix.k or output_len != matrix.n) return error.InvalidReferenceProjectionLength; + if (matrix.k > scratch.rotated.len or matrix.n > scratch.sums.len) + return error.ReferenceProjectionScratchTooSmall; + const trellis = try borrowed(matrix.trellis, try matrix.storageBytes()); + const su = try borrowed(matrix.suh, try scaleBytes(matrix.k, matrix.suh_format)); + // Check every borrowed view before doing arithmetic, including the output-only scale and optional bias. + _ = try borrowed(matrix.svh, try scaleBytes(matrix.n, matrix.svh_format)); + if (matrix.bias) |bias| _ = try borrowed(bias, try std.math.mul(usize, matrix.n, 2)); + const rotated = scratch.rotated[0..matrix.k]; + const sums = scratch.sums[0..matrix.n]; + prepare(x, su, matrix.suh_format, rotated); + const book_index: usize = @intCast(@backingInt(matrix.codebook)); + const table = scratch.codebooks[book_index * states_count ..][0..states_count]; + switch (matrix.width.halves) { + inline 2...16 => |halves| try product(halves, matrix, trellis, table, rotated, sums), + else => unreachable, + } + } + + fn finish(comptime T: type, matrix: exl3.Matrix, y: []T, sums: []const f64) !void { + @setFloatMode(.strict); + const sv = try borrowed(matrix.svh, try scaleBytes(matrix.n, matrix.svh_format)); + const bias = if (matrix.bias) |view| try borrowed(view, try std.math.mul(usize, matrix.n, 2)) else null; + var base: usize = 0; + while (base < matrix.n) : (base += block_size) { + var values: [block_size]Real = undefined; + // In FP32 one FP64->FP32 boundary replaces GPU FP32 dot-product accumulation; FP64 keeps the wide sum. + for (&values, 0..) |*value, i| value.* = @floatCast(sums[base + i]); + hadamard(&values); + for (values, 0..) |value, i| { + var out: Real = (value * wr(inverse)) * wr(scale(sv, base + i, matrix.svh_format)); + if (bias) |bytes| out += wr(half(bytes, base + i)); + y[base + i] = if (T == u16) bf16(out) else @floatCast(out); + } + } + } + + /// Kept-dense: quantized rotation and output stages; a wide dot replaces the tile decode. + fn applyDense(comptime T: type, d: weights.Dense, x: []const u16, y: []T, scratch: *Scratch) !void { + @setFloatMode(.strict); + if (x.len != d.k or y.len != d.n) return error.InvalidReferenceProjectionLength; + if (d.k > scratch.rotated.len or d.n > scratch.sums.len) + return error.ReferenceProjectionScratchTooSmall; + const su = try borrowed(.{ .buffer = d.suh.buffer, .offset = d.suh.offset }, d.k * 2); + const sv = try borrowed(.{ .buffer = d.svh.buffer, .offset = d.svh.offset }, d.n * 2); + const w = try borrowed(.{ .buffer = d.weight.buffer, .offset = d.weight.offset }, d.n * d.k * 2); + const rotated = scratch.rotated[0..d.k]; + const sums = scratch.sums[0..d.n]; + prepare(x, su, .f16, rotated); + for (sums) |*slot| slot.* = 0; + for (0..d.n) |j| { + var acc: f64 = 0; + const wrow = 2 * j * d.k; + for (0..d.k) |k| { + const word = std.mem.readInt(u16, w[wrow + 2 * k ..][0..2], .little); + acc += rotated[k] * @as(f64, @as(f32, @bitCast(@as(u32, word) << 16))); + } + sums[j] = acc; + } + var base: usize = 0; + while (base < d.n) : (base += block_size) { + var values: [block_size]Real = undefined; + for (&values, 0..) |*value, i| value.* = @floatCast(sums[base + i]); + hadamard(&values); + for (values, 0..) |value, i| { + const out = (value * wr(inverse)) * wr(half(sv, base + i)); + y[base + i] = if (T == u16) bf16(out) else @floatCast(out); + } + } + } + + /// One BF16 input row and BF16 output row; caller must finish GPU writes before borrowing shared CPU bytes. + pub fn applyTrellis(matrix: exl3.Matrix, x: []const u16, y: []u16, scratch: *Scratch) !void { + try project(matrix, x, y.len, scratch); + try finish(u16, matrix, y, scratch.sums[0..matrix.n]); + } + + /// The same arithmetic with the module-width output boundary, for compressor projections and logits. + pub fn applyF32Trellis(matrix: exl3.Matrix, x: []const u16, y: []Real, scratch: *Scratch) !void { + try project(matrix, x, y.len, scratch); + try finish(Real, matrix, y, scratch.sums[0..matrix.n]); + } + + /// Either weight representation, one entry point: trellis bytes or a kept-dense FP16 row. + pub fn apply(matrix: weights.AnyMatrix, x: []const u16, y: []u16, scratch: *Scratch) !void { + switch (matrix) { + .trellis => |m| return applyTrellis(m, x, y, scratch), + .dense => |d| return applyDense(u16, d, x, y, scratch), + } + } + + /// The same union dispatch with the module-width output boundary. + pub fn applyF32(matrix: weights.AnyMatrix, x: []const u16, y: []Real, scratch: *Scratch) !void { + switch (matrix) { + .trellis => |m| return applyF32Trellis(m, x, y, scratch), + .dense => |d| return applyDense(Real, d, x, y, scratch), + } + } + }; } +const f32_module = Module(f32); +pub const apply = f32_module.apply; +pub const applyF32 = f32_module.applyF32; + test "CPU circular streaming matches scalar format for every supported width and codebook" { var scratch = try Scratch.init(std.testing.allocator, 128, 128); defer scratch.deinit(); @@ -297,15 +313,16 @@ test "CPU circular streaming matches scalar format for every supported width and } test "FP32 Hadamard then FP16 and BF16 boundaries are deterministic" { + const f32m = Module(f32); var x: [128]u16 = @splat(0); x[0] = 0x3f80; var signs: [16]u8 = @splat(0); signs[0] = 1; var rotated: [128]f64 = undefined; - prepare(&x, &signs, .packed_signs, &rotated); - const rounded: f16 = @floatCast(-inverse); + f32m.prepare(&x, &signs, .packed_signs, &rotated); + const rounded: f16 = @floatCast(-@as(f32, inverse)); for (rotated) |value| try std.testing.expectEqual(@as(f64, @floatCast(rounded)), value); - try std.testing.expectEqual(@as(u16, 0x3f80), bf16(@bitCast(@as(u32, 0x3f808000)))); - try std.testing.expectEqual(@as(u16, 0x3f82), bf16(@bitCast(@as(u32, 0x3f818000)))); - try std.testing.expectEqual(@as(u16, 0x7fc0), bf16(@bitCast(@as(u32, 0x7f800001)))); + try std.testing.expectEqual(@as(u16, 0x3f80), f32m.bf16(@bitCast(@as(u32, 0x3f808000)))); + try std.testing.expectEqual(@as(u16, 0x3f82), f32m.bf16(@bitCast(@as(u32, 0x3f818000)))); + try std.testing.expectEqual(@as(u16, 0x7fc0), f32m.bf16(@bitCast(@as(u32, 0x7f800001)))); } diff --git a/zig/tests/deepseek_v41_reference/reference.zig b/zig/tests/deepseek_v41_reference/reference.zig index 98069068f7..ca1ad3dfe5 100644 --- a/zig/tests/deepseek_v41_reference/reference.zig +++ b/zig/tests/deepseek_v41_reference/reference.zig @@ -1,4 +1,4 @@ -//! Independent wide-sum qualification with explicit FP32/BF16 boundaries; not GPU self-reference or CUDA emulation. +//! Independent wide-sum qualification with explicit FP32/BF16 boundaries; Module(Real) widens arithmetic only. // Equations: pinned DeepSeek AI inference/{model,engram,kernel}.py; provenance in THIRD_PARTY_NOTICES.md. const std = @import("std"); @@ -7,170 +7,188 @@ const weights = tf.deepseek_v41_weights; const config = tf.deepseek_v41_config; const cp = tf.checkpoint_metal; const exl3 = tf.exl3_projection; -const projection = @import("projection.zig"); const math = @import("math.zig"); const attention = @import("attention_math.zig"); -const engram = @import("engram_math.zig"); +const engram_math = @import("engram_math.zig"); +const projection = @import("projection.zig"); const Stage = tf.deepseek_v41_profiler.Stage; -/// Local routing isolation also uses these source equations; no native values enter qualification. -pub const MoeRouting = struct { - logits: [384]f32 = undefined, - scores: [384]f32 = undefined, - ids: [6]u32 = undefined, - weights: [6]f32 = undefined, -}; +pub fn Module(comptime Real: type) type { + return struct { + const Self = @This(); + const m = math.Module(Real); + const proj = projection.Module(Real); + const attention_math = attention.Module(Real); + const engram = engram_math.Module(Real); + + /// Local routing isolation also uses these source equations; no native values enter qualification. + pub const MoeRouting = struct { + logits: [384]Real = undefined, + scores: [384]Real = undefined, + ids: [6]u32 = undefined, + weights: [6]Real = undefined, + }; -pub fn moeSelect(moe_weights: weights.Moe, c: config.Text, logits: *const [384]f32, scores: *[384]f32, ids: *[6]u32, routing_weights: *[6]f32) void { - @setFloatMode(.strict); - for (logits, scores) |z, *score| { - const softplus = if (z > 20) z else std.math.log1p(@exp(z)); - score.* = @sqrt(softplus); - } - var used: [384]bool = @splat(false); - for (0..6) |pick| { - var best: usize = 0; - var high: f32 = -std.math.inf(f32); - for (0..384) |expert| { - const biased = scores[expert] + math.value(moe_weights.bias, expert); - if (!used[expert] and biased > high) { - best = expert; - high = biased; + pub fn moeSelect(moe_weights: weights.Moe, c: config.Text, logits: *const [384]Real, scores: *[384]Real, ids: *[6]u32, routing_weights: *[6]Real) void { + @setFloatMode(.strict); + for (logits, scores) |z, *score| { + const softplus = if (z > @as(Real, 20)) z else std.math.log1p(@exp(z)); + score.* = @sqrt(softplus); } + var used: [384]bool = @splat(false); + for (0..6) |pick| { + var best: usize = 0; + var high: Real = -std.math.inf(Real); + for (0..384) |expert| { + const biased = scores[expert] + m.value(moe_weights.bias, expert); + if (!used[expert] and biased > high) { + best = expert; + high = biased; + } + } + ids[pick] = @intCast(best); + routing_weights[pick] = scores[best]; + used[best] = true; + } + var normalization: f64 = 0; + for (routing_weights) |v| normalization += @as(f64, v); + const denominator = @as(Real, @floatCast(normalization)) + @as(Real, 1e-20); + for (routing_weights) |*v| v.* = (v.* / denominator) * m.wr(c.routed_scaling_factor); } - ids[pick] = @intCast(best); - routing_weights[pick] = scores[best]; - used[best] = true; - } - var normalization: f64 = 0; - for (routing_weights) |v| normalization += v; - const denominator = @as(f32, @floatCast(normalization)) + 1e-20; - for (routing_weights) |*v| v.* = (v.* / denominator) * c.routed_scaling_factor; -} -pub fn moeActivation(gate: []const u16, up: []const u16, route: f32, c: config.Text, activation: []u16) void { - @setFloatMode(.strict); - for (activation, gate, up) |*v, g_value, u_value| { - const g = @min(math.f(g_value), c.swiglu_limit); - const u = @max(-c.swiglu_limit, @min(math.f(u_value), c.swiglu_limit)); - v.* = math.b(route * ((g / (1 + @exp(-g))) * u)); - } -} + pub fn moeActivation(gate: []const u16, up: []const u16, route: Real, c: config.Text, activation: []u16) void { + @setFloatMode(.strict); + for (activation, gate, up) |*v, g_value, u_value| { + const g = @min(m.f(g_value), m.wr(c.swiglu_limit)); + const u = @max(m.wr(-c.swiglu_limit), @min(m.f(u_value), m.wr(c.swiglu_limit))); + v.* = m.b(route * ((g / (1 + @exp(-g))) * u)); + } + } -pub const Reference = struct { - gpa: std.mem.Allocator, - c: config.Text, - w: *const weights.Model, - checkpoint: *const cp.Checkpoint, - capacity: u32, - position: u32 = 0, - projection_scratch: projection.Scratch, - attention_state: attention.State, - engram_state: engram.State, - logits: []f32, - h: [20480]u16 = undefined, - input: [5120]u16 = undefined, - output: [5120]u16 = undefined, - gate: [2304]u16 = undefined, - up: [2304]u16 = undefined, - activation: [2304]u16 = undefined, - expert_output: [5120]u16 = undefined, - moe_sum: [5120]f32 = undefined, - route_logits: [384]f32 = undefined, - route_scores: [384]f32 = undefined, - last_routes: [6]u32 = undefined, - last_routing: [6]f32 = undefined, - observer: ?Observer = null, - pub const Observer = struct { - ctx: *anyopaque, - read: *const fn (*anyopaque, usize, Stage, *const Reference) void, - }; + const ReferenceImpl = struct { + gpa: std.mem.Allocator, + c: config.Text, + w: *const weights.Model, + checkpoint: *const cp.Checkpoint, + capacity: u32, + position: u32 = 0, + projection_scratch: projection.Scratch, + attention_state: attention_math.State, + engram_state: engram.State, + logits: []Real, + h: [20480]u16 = undefined, + input: [5120]u16 = undefined, + output: [5120]u16 = undefined, + gate: [2304]u16 = undefined, + up: [2304]u16 = undefined, + activation: [2304]u16 = undefined, + expert_output: [5120]u16 = undefined, + moe_sum: [5120]Real = undefined, + route_logits: [384]Real = undefined, + route_scores: [384]Real = undefined, + last_routes: [6]u32 = undefined, + last_routing: [6]Real = undefined, + observer: ?Observer = null, + pub const Observer = struct { + ctx: *anyopaque, + read: *const fn (*anyopaque, usize, Stage, *const ReferenceImpl) void, + }; - /// Borrowed model/checkpoint/config/map must outlive this instance; capacity is the teacher sequence length. - pub fn init(gpa: std.mem.Allocator, c: config.Text, w: *const weights.Model, checkpoint: *const cp.Checkpoint, token_map: []const u8, capacity: u32) !Reference { - if (capacity == 0 or capacity > c.max_position_embeddings or capacity > attention.max_capacity) - return error.InvalidReferenceCapacity; - if (c.hidden_size != 5120 or c.num_hidden_layers != 40 or c.vocab_size != 129280 or - c.num_experts_per_tok != 6 or c.n_routed_experts != 384 or c.hc_mult != 4) - return error.UnsupportedReferenceArchitecture; - const engram_state = try engram.State.init(token_map); - var projection_scratch = try projection.Scratch.init(gpa, 8192, 129280); - errdefer projection_scratch.deinit(); - var attention_state = try attention.State.init(gpa, capacity); - errdefer attention_state.deinit(); - const logits = try gpa.alloc(f32, 129280); - return .{ .gpa = gpa, .c = c, .w = w, .checkpoint = checkpoint, .capacity = capacity, .projection_scratch = projection_scratch, .attention_state = attention_state, .engram_state = engram_state, .logits = logits }; - } - pub fn deinit(self: *Reference) void { - self.projection_scratch.deinit(); - self.attention_state.deinit(); - self.gpa.free(self.logits); - self.* = undefined; - } - /// Consume exactly one teacher token and return all actual checkpoint logits, borrowed until the next step. - pub fn step(self: *Reference, id: u32) ![]const f32 { - if (self.position >= self.capacity) return error.ContextFull; - const compressed = try self.engram_state.compressed(id); - for (0..5120) |d| { - const v = math.b(math.value(self.w.embed, @as(usize, id) * 5120 + d)); - for (0..4) |copy| self.h[copy * 5120 + d] = v; - } - var pre = [4]f32{ 1, 0, 0, 0 }; - for (&self.w.layers, 0..) |*layer, layer_id| { - if (layer.engram) |e| try self.engram_state.apply(e, if (layer_id == 1) 0 else 1, compressed, &self.h, self.c, &self.projection_scratch); - const attn_mix = math.mixes(&self.h, layer.hc_attn, self.c); - math.collapse(&self.h, pre, &self.input); - math.rms(&self.input, layer.attn_norm, self.c.rms_norm_eps, &self.input); - try self.attention_state.apply(layer.attn, layer_id, &self.input, self.position, self.c, &self.output, &self.projection_scratch); - math.expand(&self.h, &self.output, attn_mix); - if (self.observer) |observer| observer.read(observer.ctx, layer_id, .attention, self); - const ffn_mix = math.mixes(&self.h, layer.hc_ffn, self.c); - math.collapse(&self.h, attn_mix.pre, &self.input); - math.rms(&self.input, layer.ffn_norm, self.c.rms_norm_eps, &self.input); - try self.moe(layer.moe, layer_id); - math.expand(&self.h, &self.output, ffn_mix); - if (self.observer) |observer| observer.read(observer.ctx, layer_id, .moe, self); - pre = ffn_mix.pre; - } - math.collapse(&self.h, pre, &self.input); - math.rms(&self.input, self.w.norm, self.c.rms_norm_eps, &self.input); - try projection.applyF32(.{ .trellis = self.w.head }, &self.input, self.logits, &self.projection_scratch); - for (self.logits) |v| if (!std.math.isFinite(v)) return error.NonFiniteReferenceLogits; - self.engram_state.advance(compressed); - self.position += 1; - return self.logits; - } - fn moe(self: *Reference, moe_weights: weights.Moe, layer: usize) !void { - @setFloatMode(.strict); - math.dense(moe_weights.gate, &self.input, &self.route_logits); - moeSelect(moe_weights, self.c, &self.route_logits, &self.route_scores, &self.last_routes, &self.last_routing); - @memset(&self.moe_sum, 0); - // Official MoE visits expert ids in ascending order, accumulates FP32, then adds shared expert last. - for (0..384) |expert| { - var route: ?f32 = null; - for (self.last_routes, self.last_routing) |selected, weight| - if (selected == expert) { - route = weight; - }; - const expert_route = route orelse continue; - var matrices: [3]exl3.Matrix = undefined; - var name: [128]u8 = undefined; - inline for (.{ "w1", "w3", "w2" }, 0..) |part, i| { - const prefix = try std.fmt.bufPrint(&name, "layers.{d}.ffn.experts.{d}.{s}", .{ layer, expert, part }); - matrices[i] = try weights.lookup(self.checkpoint, prefix, if (i == 2) 2304 else 5120, if (i == 2) 5120 else 2304); + /// Borrowed model/checkpoint/config/map must outlive this instance; capacity is teacher length. + pub fn init(gpa: std.mem.Allocator, c: config.Text, w: *const weights.Model, checkpoint: *const cp.Checkpoint, token_map: []const u8, capacity: u32) !ReferenceImpl { + if (capacity == 0 or capacity > c.max_position_embeddings or capacity > attention.max_capacity) + return error.InvalidReferenceCapacity; + if (c.hidden_size != 5120 or c.num_hidden_layers != 40 or c.vocab_size != 129280 or + c.num_experts_per_tok != 6 or c.n_routed_experts != 384 or c.hc_mult != 4) + return error.UnsupportedReferenceArchitecture; + const engram_state = try engram.State.init(token_map); + var projection_scratch = try projection.Scratch.init(gpa, 8192, 129280); + errdefer projection_scratch.deinit(); + var attention_state = try attention_math.State.init(gpa, capacity); + errdefer attention_state.deinit(); + const logits = try gpa.alloc(Real, 129280); + return .{ .gpa = gpa, .c = c, .w = w, .checkpoint = checkpoint, .capacity = capacity, .projection_scratch = projection_scratch, .attention_state = attention_state, .engram_state = engram_state, .logits = logits }; } - try self.applyExpert(matrices, expert_route); - for (&self.moe_sum, self.expert_output) |*sum, v| sum.* += math.f(v); - } - try self.applyExpert(moe_weights.shared, 1); - for (&self.output, self.moe_sum, self.expert_output) |*v, routed, shared| - v.* = math.b(routed + math.f(shared)); - } - fn applyExpert(self: *Reference, matrices: [3]exl3.Matrix, route: f32) !void { - @setFloatMode(.strict); - try projection.apply(.{ .trellis = matrices[0] }, &self.input, &self.gate, &self.projection_scratch); - try projection.apply(.{ .trellis = matrices[1] }, &self.input, &self.up, &self.projection_scratch); - moeActivation(&self.gate, &self.up, route, self.c, &self.activation); - try projection.apply(.{ .trellis = matrices[2] }, &self.activation, &self.expert_output, &self.projection_scratch); - } -}; + pub fn deinit(self: *ReferenceImpl) void { + self.projection_scratch.deinit(); + self.attention_state.deinit(); + self.gpa.free(self.logits); + self.* = undefined; + } + /// Consume exactly one teacher token and return all actual checkpoint logits, borrowed until the next step. + pub fn step(self: *ReferenceImpl, id: u32) ![]const Real { + if (self.position >= self.capacity) return error.ContextFull; + const compressed = try self.engram_state.compressed(id); + for (0..5120) |d| { + const v = m.b(m.value(self.w.embed, @as(usize, id) * 5120 + d)); + for (0..4) |copy| self.h[copy * 5120 + d] = v; + } + var pre = [4]Real{ 1, 0, 0, 0 }; + for (&self.w.layers, 0..) |*layer, layer_id| { + if (layer.engram) |e| try self.engram_state.apply(e, if (layer_id == 1) 0 else 1, compressed, &self.h, self.c, &self.projection_scratch); + const attn_mix = m.mixes(&self.h, layer.hc_attn, self.c); + m.collapse(&self.h, pre, &self.input); + m.rms(&self.input, layer.attn_norm, self.c.rms_norm_eps, &self.input); + try self.attention_state.apply(layer.attn, layer_id, &self.input, self.position, self.c, &self.output, &self.projection_scratch); + m.expand(&self.h, &self.output, attn_mix); + if (self.observer) |observer| observer.read(observer.ctx, layer_id, .attention, self); + const ffn_mix = m.mixes(&self.h, layer.hc_ffn, self.c); + m.collapse(&self.h, attn_mix.pre, &self.input); + m.rms(&self.input, layer.ffn_norm, self.c.rms_norm_eps, &self.input); + try self.moe(layer.moe, layer_id); + m.expand(&self.h, &self.output, ffn_mix); + if (self.observer) |observer| observer.read(observer.ctx, layer_id, .moe, self); + pre = ffn_mix.pre; + } + m.collapse(&self.h, pre, &self.input); + m.rms(&self.input, self.w.norm, self.c.rms_norm_eps, &self.input); + try proj.applyF32(.{ .trellis = self.w.head }, &self.input, self.logits, &self.projection_scratch); + for (self.logits) |v| if (!std.math.isFinite(v)) return error.NonFiniteReferenceLogits; + self.engram_state.advance(compressed); + self.position += 1; + return self.logits; + } + fn moe(self: *ReferenceImpl, moe_weights: weights.Moe, layer: usize) !void { + @setFloatMode(.strict); + m.dense(moe_weights.gate, &self.input, &self.route_logits); + Self.moeSelect(moe_weights, self.c, &self.route_logits, &self.route_scores, &self.last_routes, &self.last_routing); + @memset(&self.moe_sum, 0); + // Official MoE visits expert ids in ascending order, accumulates FP32, then adds shared expert last. + for (0..384) |expert| { + var route: ?Real = null; + for (self.last_routes, self.last_routing) |selected, weight| + if (selected == expert) { + route = weight; + }; + const expert_route = route orelse continue; + var matrices: [3]exl3.Matrix = undefined; + var name: [128]u8 = undefined; + inline for (.{ "w1", "w3", "w2" }, 0..) |part, i| { + const prefix = try std.fmt.bufPrint(&name, "layers.{d}.ffn.experts.{d}.{s}", .{ layer, expert, part }); + matrices[i] = try weights.lookup(self.checkpoint, prefix, if (i == 2) 2304 else 5120, if (i == 2) 5120 else 2304); + } + try self.applyExpert(matrices, expert_route); + for (&self.moe_sum, self.expert_output) |*sum, v| sum.* += m.f(v); + } + try self.applyExpert(moe_weights.shared, 1); + for (&self.output, self.moe_sum, self.expert_output) |*v, routed, shared| + v.* = m.b(routed + m.f(shared)); + } + fn applyExpert(self: *ReferenceImpl, matrices: [3]exl3.Matrix, route: Real) !void { + @setFloatMode(.strict); + try proj.applyTrellis(matrices[0], &self.input, &self.gate, &self.projection_scratch); + try proj.applyTrellis(matrices[1], &self.input, &self.up, &self.projection_scratch); + Self.moeActivation(&self.gate, &self.up, route, self.c, &self.activation); + try proj.applyTrellis(matrices[2], &self.activation, &self.expert_output, &self.projection_scratch); + } + }; + pub const Reference = ReferenceImpl; + }; +} + +/// FP32 instance: the retained independent whole-model qualification reference. +pub const f32_reference = Module(f32); +pub const Reference = f32_reference.Reference; +pub const MoeRouting = f32_reference.MoeRouting; +pub const moeSelect = f32_reference.moeSelect; +pub const moeActivation = f32_reference.moeActivation; From b25c4a8cee722b8f2b370e1e2e40e096649d687f Mon Sep 17 00:00:00 2001 From: Andreas Kolb Date: Mon, 19 Oct 2026 12:14:37 -0700 Subject: [PATCH 2/7] DeepSeek V4.1 reference: the FP64 narrowing keeps NaN, and the FP32 oracle words stay frozen The negative-FP64 branch fell through the NaN check's `>` test with a positive NaN's sign stripped, so a positive FP64 NaN narrowed to the finite BF16 word 0x3ffc while a negative one produced 0x7fc0. Test the FP64 class instead and return the canonical quiet NaN of the input's sign. The new tests cover every IEEE class of the narrowing (both infinities, both zeros, FP64 subnormals, the largest-finite overflow, and six NaN payload classes), the BF16 subnormal and overflow tie grid against the FP32 branch, and frozen FP32 outputs of projection, the hyperconnection stages, attention pooling and the Engram gather. The frozen words were captured on 53b790082, before the repair; the repair touches only the FP64 narrowing, so those FP32 words must not move. Gates: `zig build test-deepseek-v41-reference` 43/43 steps, 32/32 tests; `lean_check.py` 501 problems against the same 501 on the parent commit; `zig fmt --check zig` clean. (cherry picked from commit 9641836514bf9116ea3d8d0ab2a63520eaabf324) --- zig/deepseek_v41_reference_checks.zig | 1 + .../deepseek_v41_reference/attention_math.zig | 35 ++++++ .../deepseek_v41_reference/engram_math.zig | 36 +++++- zig/tests/deepseek_v41_reference/math.zig | 112 +++++++++++++++++- .../deepseek_v41_reference/projection.zig | 35 ++++++ 5 files changed, 216 insertions(+), 3 deletions(-) diff --git a/zig/deepseek_v41_reference_checks.zig b/zig/deepseek_v41_reference_checks.zig index 1bc109c33c..dbc66acf97 100644 --- a/zig/deepseek_v41_reference_checks.zig +++ b/zig/deepseek_v41_reference_checks.zig @@ -3,6 +3,7 @@ test { _ = @import("tests/deepseek_v41_reference/projection.zig"); _ = @import("tests/deepseek_v41_reference/math.zig"); _ = @import("tests/deepseek_v41_reference/attention_math.zig"); + _ = @import("tests/deepseek_v41_reference/engram_math.zig"); _ = @import("tests/deepseek_v41_reference/qualify.zig"); _ = @import("tests/deepseek_v41_reference/engram_math.zig"); _ = @import("tests/deepseek_v41_reference/precision.zig"); diff --git a/zig/tests/deepseek_v41_reference/attention_math.zig b/zig/tests/deepseek_v41_reference/attention_math.zig index 889dced2d3..656d390614 100644 --- a/zig/tests/deepseek_v41_reference/attention_math.zig +++ b/zig/tests/deepseek_v41_reference/attention_math.zig @@ -498,3 +498,38 @@ test "attention reference capacity rejects unbounded history before allocation" try std.testing.expectError(error.InvalidAttentionReferenceCapacity, State.init(std.testing.allocator, 0)); try std.testing.expectError(error.InvalidAttentionReferenceCapacity, State.init(std.testing.allocator, max_capacity + 1)); } + +test "frozen FP32 pooling and sparse-head boundaries match the pre-repair oracle" { + // Frozen at 53b790082 (pre-repair): poolOf and this head's value GEMM are FP32, so both words stay. + const a_values = [_]f32{ + 1, 1 + 0x1p-7, 1, -1, 0, 0x1p-24, 1 - 0x1p-7, 3 * 0x1p-125, + 0x1p-126, 0x1p126, 0x1p-8, -0x1p-24, 1, 0x1p-7, 1.5, 2, + }; + const b_values = [_]f32{ + 1 + 0x1p-7, 1 + 0x1p-6, 3, -1 - 0x1p-7, 0, 3 * 0x1p-24, 1 - 0x1p-6, 5 * 0x1p-125, + 0x1p-126, 0x1p126, 0x1p-8, -3 * 0x1p-24, -1, 3 * 0x1p-7, 2.5, 2, + }; + // Equal gates give exact 1/2 probabilities; equal scores give the sparse rows 1, with a -inf sink. + var window: [2 * 512]u16 = @splat(0); + var kv: [2][512]f32 = undefined; + for (0..512) |d| { + const column = d % a_values.len; + window[d] = m.b(a_values[column]); + window[512 + d] = m.b(b_values[column]); + kv[0][d] = m.f(window[d]); + kv[1][d] = m.f(window[512 + d]); + } + var gates: [2][512]f32 = .{ @splat(0), @splat(0) }; + var pooled: [512]u16 = undefined; + poolOf(f32, .{ &kv[0], &kv[1] }, .{ &gates[0], &gates[1] }, &pooled); + var q: [512]u16 = @splat(0); + var indices = [2]i32{ 0, 1 }; + var sparse: [512]u16 = @splat(0); + sparseHeadPre(f32, &q, .{ .window = &window, .window_stride = 512 }, &indices, -std.math.inf(f32), &sparse, null); + try std.testing.expectEqualSlices(u16, &pooled, &sparse); + const frozen = [_]u16{ + 0x3f80, 0x3f82, 0x4000, 0xbf80, 0x0000, 0x3400, 0x3f7d, 0x0200, + 0x0080, 0x7e80, 0x3b80, 0xb400, 0x0000, 0x3c80, 0x4000, 0x4000, + }; + try std.testing.expectEqualSlices(u16, &frozen, pooled[0..16]); +} diff --git a/zig/tests/deepseek_v41_reference/engram_math.zig b/zig/tests/deepseek_v41_reference/engram_math.zig index 5787ced72f..ea8ab1d091 100644 --- a/zig/tests/deepseek_v41_reference/engram_math.zig +++ b/zig/tests/deepseek_v41_reference/engram_math.zig @@ -22,6 +22,10 @@ pub fn Module(comptime Real: type) type { return struct { const m = math.Module(Real); const proj = projection.Module(Real); + /// Engram gather narrowing: the E4M3 table word times its UE8M0 row scale, rounded to BF16. + pub fn narrowedProduct(table: Real, scale: Real) u16 { + return m.b(table * scale); + } const StateImpl = struct { map: []const u8, history: [3]u32 = .{ 2, 2, 2 }, @@ -75,7 +79,7 @@ pub fn Module(comptime Real: type) type { const column = (lookback - 1) * 8 + head; const row = offset + @as(usize, @intCast(rolling % primes[layer][column])); for (0..256) |d| self.gathered[column * 256 + d] = - m.b(m.value(e.table, row * 256 + d) * m.value(e.scales, row * 8 + d / 32)); + narrowedProduct(m.value(e.table, row * 256 + d), m.value(e.scales, row * 8 + d / 32)); offset += primes[layer][column]; } } @@ -116,3 +120,33 @@ pub fn Module(comptime Real: type) type { /// The retained FP32 instance; every field the caller touches today keeps its name and type. pub const f32_module = Module(f32); pub const State = f32_module.State; + +test "frozen FP32 Engram gather boundaries match the pre-repair oracle" { + // Frozen at 53b790082 (pre-repair): FP32 gather words stay; cells pair E4M3 codes and UE8M0 scales. + const f32m = math.Module(f32); + const cells = [_][2]u8{ + .{ 56, 254 }, // 1 * 2^127, the largest BF16 exponent + .{ 64, 254 }, // 2 * 2^127 overflows to infinity + .{ 56, 253 }, // 1 * 2^126 + .{ 126, 254 }, // 448 * 2^127 overflows to infinity + .{ 1, 3 }, // 2^-9 * 2^-124 = 2^-133, one BF16 subnormal unit + .{ 3, 3 }, // 3 * 2^-133 + .{ 7, 2 }, // 7 * 2^-134 = 3.5 units, ties to even 4 + .{ 1, 2 }, // 2 * 2^-135 = 0.5 units, ties to even zero + .{ 1, 4 }, // 2^-132 = 2 units + .{ 56, 6 }, // 1 * 2^-121 + .{ 0, 254 }, // zero table word + .{ 0x80 | 1, 3 }, // -2^-133 + .{ 127, 0 }, // E4M3 NaN table word + .{ 255, 0 }, // E4M3 NaN table word + }; + var narrowed: [cells.len]u16 = undefined; + for (&narrowed, cells) |*narrowed_word, cell| { + narrowed_word.* = f32_module.narrowedProduct(f32m.e4(cell[0]), f32m.e8(cell[1])); + } + const frozen = [_]u16{ + 0x7f00, 0x7f80, 0x7e80, 0x7f80, 0x0001, 0x0003, 0x0004, + 0x0000, 0x0002, 0x0300, 0x0000, 0x8001, 0x7fc0, 0x7fc0, + }; + try std.testing.expectEqualSlices(u16, &frozen, &narrowed); +} diff --git a/zig/tests/deepseek_v41_reference/math.zig b/zig/tests/deepseek_v41_reference/math.zig index 55368feca3..a52af5d663 100644 --- a/zig/tests/deepseek_v41_reference/math.zig +++ b/zig/tests/deepseek_v41_reference/math.zig @@ -27,8 +27,7 @@ pub fn Module(comptime Real: type) type { const bits: u64 = @bitCast(x); const sign: u16 = @as(u16, @truncate(bits >> 63)) << 15; const magnitude: u64 = bits & 0x7fff_ffffffffffff; - if (magnitude > 0x7ff0_0000_0000_0000) - return sign | @as(u16, @truncate(((bits >> 49) | 0x0040) & 0xffff)); + if (magnitude > 0x7ff0_0000_0000_0000) return sign | 0x7fc0; if (magnitude == 0x7ff0_0000_0000_0000) return sign | 0x7f80; const biased: u11 = @truncate(magnitude >> 52); const biased_wide: i32 = @intCast(biased); @@ -453,3 +452,112 @@ test "CUDA mean projection preserves the BF16 RMS affine tie" { const inverse = idealRsqrt(meanFromWide(9, 5120) + @as(f32, 1e-20)); try std.testing.expectEqual(b(1), b((@as(f32, 3) * inverse) * gamma)); } + +test "FP64 narrowing keeps every IEEE class, including quiet NaN" { + const R64 = Module(f64); + // Raw FP64 class boundaries: NaN collapses to the canonical quiet NaN of its sign, never a finite word. + const classes = [_]struct { bits: u64, expected: u16 }{ + .{ .bits = 0x7ff0000000000001, .expected = 0x7fc0 }, // signalling NaN, positive + .{ .bits = 0x7ff8000000000000, .expected = 0x7fc0 }, // quiet NaN, positive + .{ .bits = 0x7fffffffffffffff, .expected = 0x7fc0 }, // all-payload NaN, positive + .{ .bits = 0xfff0000000000001, .expected = 0xffc0 }, // signalling NaN, negative + .{ .bits = 0xfff8000000000000, .expected = 0xffc0 }, // quiet NaN, negative + .{ .bits = 0xffffffffffffffff, .expected = 0xffc0 }, // all-payload NaN, negative + .{ .bits = 0x7ff0000000000000, .expected = 0x7f80 }, // +infinity + .{ .bits = 0xfff0000000000000, .expected = 0xff80 }, // -infinity + .{ .bits = 0x0000000000000000, .expected = 0x0000 }, // +0 + .{ .bits = 0x8000000000000000, .expected = 0x8000 }, // -0 + .{ .bits = 0x7fefffffffffffff, .expected = 0x7f80 }, // largest finite FP64 overflows + .{ .bits = 0xffefffffffffffff, .expected = 0xff80 }, + .{ .bits = 0x0000000000000001, .expected = 0x0000 }, // smallest FP64 subnormal underflows + .{ .bits = 0x000fffffffffffff, .expected = 0x0000 }, // largest FP64 subnormal underflows + .{ .bits = 0x8000000000000001, .expected = 0x8000 }, + }; + for (classes) |entry| { + const raw: f64 = @bitCast(entry.bits); + try std.testing.expectEqual(entry.expected, R64.b(raw)); + if (entry.expected == 0x7fc0 or entry.expected == 0xffc0) { + try std.testing.expect(std.math.isNan(R64.f(R64.b(raw)))); + } + } + // The FP32 branch is the reference: every FP32-representable finite value and NaN payload must agree. + try std.testing.expectEqual(@as(u16, 0x7fc0), b(@bitCast(@as(u32, 0x7f800001)))); +} + +test "FP64 narrows BF16 subnormal and overflow ties like the FP32 branch" { + const R64 = Module(f64); + // Values on the exact BF16 half-ULP grid: each is exact in FP32, so both branches return one word. + const ties = [_]struct { bits: u64, expected: u16 }{ + .{ .bits = 0x37a0000000000000, .expected = 0x0001 }, // 2^-133, one BF16 subnormal unit + .{ .bits = 0x3790000000000000, .expected = 0x0000 }, // 2^-134, half a unit ties to even zero + .{ .bits = 0x37a8000000000000, .expected = 0x0002 }, // 1.5*2^-133 ties to even 2 + .{ .bits = 0x37b4000000000000, .expected = 0x0002 }, // 2.5*2^-133 ties to even 2 + .{ .bits = 0x37bc000000000000, .expected = 0x0004 }, // 3.5*2^-133 ties to even 4 + .{ .bits = 0x3780000000000000, .expected = 0x0000 }, // 2^-135 underflows + .{ .bits = 0x380fc00000000000, .expected = 0x007f }, // largest BF16 subnormal + .{ .bits = 0x380fe00000000000, .expected = 0x0080 }, // half-ULP below the smallest normal rounds up + .{ .bits = 0x380fdffe00000000, .expected = 0x007f }, // just below that half-ULP stays subnormal + .{ .bits = 0x36a0000000000000, .expected = 0x0000 }, // 2^-149, smallest FP32 subnormal + .{ .bits = 0x3690000000000000, .expected = 0x0000 }, // 2^-150 has no FP32 peer + .{ .bits = 0xb7a0000000000000, .expected = 0x8001 }, + .{ .bits = 0xb790000000000000, .expected = 0x8000 }, + .{ .bits = 0x47efe00000000000, .expected = 0x7f7f }, // largest BF16 finite + .{ .bits = 0x47eff00000000000, .expected = 0x7f80 }, // half-ULP to infinity rounds up + .{ .bits = 0x47f0000000000000, .expected = 0x7f80 }, // 2^128 has no FP32 peer + }; + for (ties) |entry| { + const raw: f64 = @bitCast(entry.bits); + try std.testing.expectEqual(entry.expected, R64.b(raw)); + } + // Adjacent grid points that are exactly representable on both sides stay on the same word. + const parity = [_]struct { bits: u64, expected: u16 }{ + .{ .bits = 0x3ff0000000000000, .expected = 0x3f80 }, // 1 + .{ .bits = 0x3ff0100000000000, .expected = 0x3f80 }, // 1+2^-8 ties to even 1 + .{ .bits = 0x3ff0200000000000, .expected = 0x3f81 }, // 1+2^-7 + .{ .bits = 0x3ff0300000000000, .expected = 0x3f82 }, // 1+3*2^-8 ties to even 2 + .{ .bits = 0x47efffffe0000000, .expected = 0x7f80 }, // largest FP32 finite rounds up + .{ .bits = 0x3800000000000000, .expected = 0x0040 }, // 2^-127 becomes a BF16 subnormal + .{ .bits = 0xb7a8000000000000, .expected = 0x8002 }, // -1.5*2^-133 + }; + for (parity) |entry| { + const raw: f64 = @bitCast(entry.bits); + const narrowed: f32 = @floatCast(raw); + try std.testing.expectEqual(entry.expected, R64.b(raw)); + try std.testing.expectEqual(entry.expected, b(narrowed)); + } +} + +test "frozen FP32 hyperconnection boundaries match the pre-repair oracle" { + // Frozen at 53b790082 (pre-repair): every operation below is FP32-only, so its words must not move. + var h = [_]u16{ + b(1), b(1), + b(0x1p-8), b(0x1p-7), + b(0x1p-25), b(0x1p-24), + b(0x1p-25), b(-0x1p-24), + }; + var collapsed: [2]u16 = undefined; + // pre {1,1,1,1}: column 0 lands exactly one 2^-24 above the 1+2^-8 tie, column 1 is exact. + collapse(&h, .{ 1, 1, 1, 1 }, &collapsed); + var expanded = [_]u16{ + b(1), b(-1), + b(1 + 0x1p-7), b(1 + 0x1p-7), + b(0x1p-8), b(-0x1p-8), + b(0), b(0), + }; + // post {1,1,1,1}, comb 0.5*I: new[1] = 1 + 0.5*(1+2^-7) is the tie between 1.5 and 1+2^-6; stays 1.5. + const mix = Mix{ + .pre = .{ 1, 1, 1, 1 }, + .post = .{ 1, 1, 1, 1 }, + .comb = .{ + 0.5, 0, 0, 0, + 0, 0.5, 0, 0, + 0, 0, 0.5, 0, + 0, 0, 0, 0.5, + }, + }; + expand(&expanded, &.{ b(1), b(-1) }, mix); + const frozen_collapsed = [_]u16{ 0x3f80, 0x3f81 }; + const frozen_expanded = [_]u16{ 0x3fc0, 0xbfc0, 0x3fc0, 0xbefe, 0x3f80, 0xbf80, 0x3f80, 0xbf80 }; + try std.testing.expectEqualSlices(u16, &frozen_collapsed, &collapsed); + try std.testing.expectEqualSlices(u16, &frozen_expanded, &expanded); +} diff --git a/zig/tests/deepseek_v41_reference/projection.zig b/zig/tests/deepseek_v41_reference/projection.zig index e74e1b8009..7edea2de98 100644 --- a/zig/tests/deepseek_v41_reference/projection.zig +++ b/zig/tests/deepseek_v41_reference/projection.zig @@ -326,3 +326,38 @@ test "FP32 Hadamard then FP16 and BF16 boundaries are deterministic" { try std.testing.expectEqual(@as(u16, 0x3f82), f32m.bf16(@bitCast(@as(u32, 0x3f818000)))); try std.testing.expectEqual(@as(u16, 0x7fc0), f32m.bf16(@bitCast(@as(u32, 0x7f800001)))); } + +test "frozen FP32 rotation and output boundaries match the pre-repair oracle" { + // Frozen at 53b790082 (pre-repair): prepare, Hadamard and output narrowing are FP32, so words stay. + const f32m = Module(f32); + const inputs = [_]f32{ 1, -1, 0x1p-24, 1 + 0x1p-7, 0x1.8p-24, -0x1p-7, 0x1p-8, 1 - 0x1p-7 }; + var x: [block_size]u16 = undefined; + for (&x, 0..) |*word, i| word.* = f32m.bf16(inputs[i % inputs.len]); + var signs: [16]u8 = @splat(0); + for (0..block_size) |i| { + if (i % 3 == 1) signs[i / 16] |= @as(u8, 1) << @intCast(i % 16); + } + var rotated: [block_size]f64 = undefined; + f32m.prepare(&x, &signs, .packed_signs, &rotated); + var rotation: [block_size]u16 = undefined; + for (&rotation, rotated) |*word, value| word.* = @bitCast(@as(f16, @floatCast(value))); + // The finish output stage: widen the rotation, seven butterflies, the inverse and column scale, BF16. + var values: [block_size]f32 = undefined; + for (&values, rotation) |*value, word| value.* = @floatCast(@as(f16, @bitCast(word))); + f32m.hadamard(&values); + var narrowed: [block_size]u16 = undefined; + for (&narrowed, values, 0..) |*word, value, i| { + const column_scale: f16 = if (i % 2 == 0) 1 else 1 + 0x1p-7; + word.* = f32m.bf16((value * f32m.wr(inverse)) * @as(f32, @floatCast(column_scale))); + } + const frozen_rotation = [_]u16{ + 0x3c3d, 0x1cf3, 0xb9b8, 0x3f16, 0x35cd, 0x398d, 0x91a8, 0x3c43, + 0x35b6, 0x983e, 0x91a8, 0x35a5, 0x35a5, 0x91a8, 0x983e, 0x35b6, + }; + const frozen_narrowed = [_]u16{ + 0x3f80, 0x3f81, 0xb982, 0x3f82, 0x3913, 0x3c02, 0x3b84, 0xbf80, + 0xbf80, 0x3f81, 0xb982, 0xbf82, 0x3913, 0xbc04, 0xbb82, 0xbf80, + }; + try std.testing.expectEqualSlices(u16, &frozen_rotation, rotation[0..16]); + try std.testing.expectEqualSlices(u16, &frozen_narrowed, narrowed[0..16]); +} From dd8ce6d72295de02b4618ca97d8f7fdfb2d0cf4e Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Sun, 18 Oct 2026 17:00:06 -0700 Subject: [PATCH 3/7] DeepSeek V4.1 reference: fp64 tests for the Engram hash, gather and gate The split in 766564937 left the Engram stage with its behavioral Metal tests and its FP32 equations but no three-way harness. This drives the real hash, gather and gate kernels against the FP32 source contract and an FP64 widening of the same equations. The family module is private, so tensorfold.zig exports it for the reference tree; the tables, scales and family API are unchanged. (cherry picked from commit 87aeaac2853377e071f12c65ec976f0d8b374bce) --- zig/deepseek_v41_reference_checks.zig | 2 +- zig/src/tensorfold.zig | 1 + .../operator_engram_fp64.zig | 365 ++++++++++++++++++ 3 files changed, 367 insertions(+), 1 deletion(-) create mode 100644 zig/tests/deepseek_v41_reference/operator_engram_fp64.zig diff --git a/zig/deepseek_v41_reference_checks.zig b/zig/deepseek_v41_reference_checks.zig index dbc66acf97..1324b64acb 100644 --- a/zig/deepseek_v41_reference_checks.zig +++ b/zig/deepseek_v41_reference_checks.zig @@ -5,6 +5,6 @@ test { _ = @import("tests/deepseek_v41_reference/attention_math.zig"); _ = @import("tests/deepseek_v41_reference/engram_math.zig"); _ = @import("tests/deepseek_v41_reference/qualify.zig"); - _ = @import("tests/deepseek_v41_reference/engram_math.zig"); _ = @import("tests/deepseek_v41_reference/precision.zig"); + _ = @import("tests/deepseek_v41_reference/operator_engram_fp64.zig"); } diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index 24b723240e..69e0002c6f 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -31,6 +31,7 @@ pub const deepseek_v41_profiler = @import("families/deepseek_v41/metal/profiler. pub const gpu_counters = @import("core/gpu_counters.zig"); pub const expert_groups = @import("core/expert_groups.zig"); pub const deepseek_v41_engram_hash = @import("families/deepseek_v41/engram_hash.zig"); +pub const deepseek_v41_engram = @import("families/deepseek_v41/engram.zig"); pub const exl3_format = @import("core/exl3_format.zig"); pub const exl3_projection = @import("core/exl3_projection.zig"); pub const exlr8 = @import("core/exlr8.zig"); diff --git a/zig/tests/deepseek_v41_reference/operator_engram_fp64.zig b/zig/tests/deepseek_v41_reference/operator_engram_fp64.zig new file mode 100644 index 0000000000..fa385016b5 --- /dev/null +++ b/zig/tests/deepseek_v41_reference/operator_engram_fp64.zig @@ -0,0 +1,365 @@ +//! Engram hash, FP8 gather and signed-sqrt gate at the FP32 contract and the FP64 ideal; receipts print "ENG64". +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const math = @import("math.zig"); + +const engram = tf.deepseek_v41_engram; +const hash = tf.deepseek_v41_engram_hash; +const View = engram.View; + +const dim: usize = 5120; +const copies: usize = 4; +const hidden_width: usize = dim * copies; +const kv_width: usize = dim * (copies + 1); +const gather_dim: usize = engram.gather_dim; +const head_dim: usize = engram.head_dim; +const eps: f32 = 1e-6; +const raw_ids: usize = 64; + +/// Synthetic token map and history so the hash case needs no checkpoint bytes. +const hash_tokens: usize = 10; +const hash_chunk: usize = 8; + +/// One gather case covers every column of the real table layout with fewer rows. +const gather_tokens: usize = 4; +const table_rows: usize = 128; + +/// Two gate rows: the second carries a DEAD mask so the zero-gate path is checked exactly. +const gate_tokens: usize = 2; + +const Lcg = struct { + state: u64 = 0x853c49e6748fea9b, + fn next(self: *Lcg) u64 { + self.state = self.state *% 6364136223846793005 +% 1442695040888963407; + return self.state >> 11; + } + /// Uniform [-1, 1) with 24-bit granularity, exactly representable in FP32. + fn uniform(self: *Lcg) f32 { + const v: f64 = @as(f64, @floatFromInt(self.next() & 0xffffff)) / 16777216.0; + return @floatCast(v * 2 - 1); + } + fn below(self: *Lcg, limit: usize) usize { + return @intCast(self.next() % limit); + } +}; + +const Delta = struct { + cells: usize = 0, + unequal: usize = 0, + nrms: f64 = 0, + max_abs: f64 = 0, + worst: usize = 0, +}; + +fn deltaBf16(actual: []const u16, ideal: []const u16) Delta { + var d: Delta = .{ .cells = actual.len }; + var squared: f64 = 0; + var magnitude: f64 = 0; + for (actual, ideal, 0..) |a, want, i| { + const ref: f64 = math.d64(want); + if (std.math.isFinite(ref)) magnitude += ref * ref; + if (a == want) continue; + const got: f64 = math.d64(a); + const diff = got - ref; + d.unequal += 1; + if (@abs(diff) > d.max_abs) { + d.max_abs = @abs(diff); + d.worst = i; + } + if (std.math.isFinite(diff)) squared += diff * diff; + } + d.nrms = if (magnitude == 0) @sqrt(squared) else @sqrt(squared / magnitude); + return d; +} + +fn printBf16(stage: []const u8, case_name: []const u8, d: Delta) void { + std.debug.print("ENG64 {{\"stage\":\"{s}\",\"case\":\"{s}\",\"cells\":{d},\"bf16_unequal\":{d},\"nrms\":{d},\"max_abs\":{d},\"worst\":{d}}}\n", .{ stage, case_name, d.cells, d.unequal, d.nrms, d.max_abs, d.worst }); +} + +fn printEq(stage: []const u8, case_name: []const u8, cells: usize, unequal: usize) void { + std.debug.print("ENG64 {{\"stage\":\"{s}\",\"case\":\"{s}\",\"cells\":{d},\"metal_vs_contract_unequal\":{d}}}\n", .{ stage, case_name, cells, unequal }); +} + +fn view(buffer: mtl.Buffer) View { + return .{ .buffer = buffer }; +} + +fn completed(cb: mtl.CommandBuffer) !void { + cb.commit(); + cb.wait(); + if (cb.failure() != null) return error.EngramGpuFailure; +} + +/// The pinned CPU oracle's gate: wide sums, the FP32 mean projection, then FP32 rstd, dot and gate. +fn contractProduct(value: f32, scale: f32) u16 { + const product = value * scale; + if (@abs(scale) < 0x1p-126 or @abs(product) < 0x1p-126) return 0; + return math.b(product); +} + +fn contractGate(hidden: []const u16, kv: []const u16, q: []const f32, k: []const f32, copy: usize, out: []u16) void { + @setFloatMode(.strict); + var hh: f64 = 0; + var kk: f64 = 0; + var dot: f64 = 0; + for (0..dim) |d| { + const at = copy * dim + d; + const h = math.f(hidden[at]); + const x = math.f(kv[at]); + const weight = q[at] * k[at]; + hh += @as(f64, h) * @as(f64, h); + kk += @as(f64, x) * @as(f64, x); + dot += @as(f64, (h * weight) * x); + } + const h_inverse = math.idealRsqrt(math.meanFromWide(hh, dim) + eps); + const k_inverse = math.idealRsqrt(math.meanFromWide(kk, dim) + eps); + const rstd = h_inverse * k_inverse; + const dot_wide: f32 = @floatCast(dot); + const normalized = (dot_wide * rstd) * (1 / @sqrt(@as(f32, 5120.0))); + const root = std.math.copysign(@sqrt(@max(@abs(normalized), @as(f32, 1e-6))), normalized); + const gate = math.sigmoid(root); + for (0..dim) |d| { + const at = copy * dim + d; + out[d] = math.b(math.f(hidden[at]) + gate * math.f(kv[hidden_width + d])); + } +} + +/// The same equations widened to FP64, never an emulation of the Metal kernel. +fn idealGate(hidden: []const u16, kv: []const u16, q: []const f32, k: []const f32, copy: usize, out: []u16) void { + @setFloatMode(.strict); + var hh: f64 = 0; + var kk: f64 = 0; + var dot: f64 = 0; + for (0..dim) |d| { + const at = copy * dim + d; + const h = math.d64(hidden[at]); + const x = math.d64(kv[at]); + const weight = @as(f64, q[at]) * @as(f64, k[at]); + hh += h * h; + kk += x * x; + dot += (h * weight) * x; + } + const h_inverse = 1 / @sqrt(hh / 5120.0 + @as(f64, eps)); + const k_inverse = 1 / @sqrt(kk / 5120.0 + @as(f64, eps)); + const normalized = (dot * (h_inverse * k_inverse)) / @sqrt(@as(f64, 5120.0)); + const root = std.math.copysign(@sqrt(@max(@abs(normalized), @as(f64, 1e-6))), normalized); + const gate = 1 / (1 + @exp(-root)); + for (0..dim) |d| { + const at = copy * dim + d; + out[d] = math.b64(math.d64(hidden[at]) + gate * math.d64(kv[hidden_width + d])); + } +} + +fn sampleBf16(lcg: *Lcg, family: usize, i: usize) u16 { + const scale: f32 = switch (family) { + 0 => 1, + 1 => 16, + 2 => 0.0078125, + else => if (i % 2 == 0) 1 else 0.0009765625, + }; + return math.b(lcg.uniform() * scale); +} + +const family_names = [_][]const u8{ "unit", "wide", "narrow", "cancellation" }; + +test "Engram hash columns equal the integer contract across a masked chunk and committed history" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + var pipelines = try engram.Pipelines.load(device); + defer pipelines.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var lcg = Lcg{}; + var map_bytes: [raw_ids * 4]u8 = undefined; + for (0..raw_ids) |i| std.mem.writeInt(u32, map_bytes[i * 4 ..][0..4], @intCast(500 + lcg.below(90000)), .little); + const map = try device.buffer(map_bytes.len, 0); + defer map.deinit(); + @memcpy(map.slice(u8, map_bytes.len), &map_bytes); + const chunk_ids = [_]u32{ 7, 9, 11, 3, 13, 15, 17, 19, 21, 23 }; + const alive = [_]u8{ 1, 1, 1, 0, 1, 1, 1, 1, 1, 1 }; + const ids = try device.buffer(hash_tokens * 4, 0); + defer ids.deinit(); + @memcpy(ids.slice(u32, hash_tokens), &chunk_ids); + const mask = try device.buffer(hash_tokens, 0); + defer mask.deinit(); + @memcpy(mask.slice(u8, hash_tokens), &alive); + const history = try device.buffer(engram.history_bytes, 0); + defer history.deinit(); + const first = try device.buffer(hash_tokens * hash.columns * 4, 0); + defer first.deinit(); + const second = try device.buffer(hash_tokens * hash.columns * 4, 0); + defer second.deinit(); + const cb = queue.commandBuffer(); + var e = cb.compute(.serial); + engram.encodeReset(e, &pipelines, view(history)); + for (0..2) |layer| { + const out = if (layer == 0) first else second; + try engram.encodeHash(e, &pipelines, hash.layers[layer], view(ids), view(map), view(mask), view(history), view(out), @intCast(hash_chunk)); + } + try engram.encodeCommit(e, &pipelines, view(ids), view(map), view(mask), view(history), @intCast(hash_chunk)); + for (0..2) |layer| { + const out = if (layer == 0) first else second; + try engram.encodeHash(e, &pipelines, hash.layers[layer], view(ids).shifted(hash_chunk * 4), view(map), view(mask).shifted(hash_chunk), view(history), view(out).shifted(hash_chunk * hash.columns * 4), @intCast(hash_tokens - hash_chunk)); + } + e.end(); + try completed(cb); + var expected: [2][hash_tokens][hash.columns]u32 = undefined; + for (0..2) |layer| { + var cpu: hash.History = .{}; + try hash.hashChunk(&cpu, .{ .bytes = &map_bytes }, hash.layers[layer], &chunk_ids, &alive, &expected[layer]); + } + var unequal: usize = 0; + for (0..2) |layer| { + const out = if (layer == 0) first else second; + const got = out.slice(u32, hash_tokens * hash.columns); + for (expected[layer], 0..) |row, t| { + for (row, 0..) |want, col| { + if (got[t * hash.columns + col] != want) unequal += 1; + } + } + } + try std.testing.expectEqual(@as(usize, 0), unequal); + printEq("hash", "masked-chunk-and-history-carry", 2 * hash_tokens * hash.columns, unequal); +} + +test "Engram gather narrows the E4M3 row grid to BF16 exactly like the FP32 contract" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + var pipelines = try engram.Pipelines.load(device); + defer pipelines.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var lcg = Lcg{}; + var weight_bytes: [table_rows * head_dim]u8 = undefined; + var scale_bytes: [table_rows * (head_dim / 32)]u8 = undefined; + for (&weight_bytes) |*b| b.* = @intCast(lcg.below(127)); + for (&scale_bytes) |*b| b.* = @intCast(lcg.below(255)); + const weight = try device.buffer(weight_bytes.len, 0); + defer weight.deinit(); + @memcpy(weight.slice(u8, weight_bytes.len), &weight_bytes); + const scale = try device.buffer(scale_bytes.len, 0); + defer scale.deinit(); + @memcpy(scale.slice(u8, scale_bytes.len), &scale_bytes); + const row_ids = try device.buffer(gather_tokens * hash.columns * 4, 0); + defer row_ids.deinit(); + const row_slice = row_ids.slice(u32, gather_tokens * hash.columns); + for (row_slice) |*id| id.* = @intCast(lcg.below(table_rows)); + const gathered = try device.buffer(gather_tokens * gather_dim * 2, 0); + defer gathered.deinit(); + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + try engram.encodeGather(e, &pipelines, .{ .weight = view(weight), .scale = view(scale), .rows = @intCast(table_rows) }, view(row_ids), view(gathered), @intCast(gather_tokens), gather_dim); + e.end(); + try completed(cb); + const got = gathered.slice(u16, gather_tokens * gather_dim); + var contract: [gather_tokens * gather_dim]u16 = undefined; + var ideal: [gather_tokens * gather_dim]u16 = undefined; + for (0..gather_tokens) |t| { + for (0..hash.columns) |col| { + for (0..head_dim) |d| { + const row: usize = row_slice[t * hash.columns + col]; + const code = weight_bytes[row * head_dim + d]; + const block = scale_bytes[row * (head_dim / 32) + d / 32]; + contract[t * gather_dim + col * head_dim + d] = contractProduct(math.e4(code), math.e8(block)); + ideal[t * gather_dim + col * head_dim + d] = math.b64(@as(f64, math.e4(code)) * @as(f64, math.e8(block))); + } + } + } + var unequal: usize = 0; + for (got, contract[0..]) |a, want| { + if (a != want) unequal += 1; + } + try std.testing.expectEqual(@as(usize, 0), unequal); + printEq("gather", "fp8-row-grid", got.len, unequal); + printBf16("gather", "metal-vs-fp64-ideal", deltaBf16(got, &ideal)); + printBf16("gather", "fp32-contract-vs-fp64-ideal", deltaBf16(&contract, &ideal)); +} + +test "Engram gate reaches BF16 at the FP32 contract over matching equations" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + var pipelines = try engram.Pipelines.load(device); + defer pipelines.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var lcg = Lcg{}; + var q: [hidden_width]f32 = undefined; + var k: [hidden_width]f32 = undefined; + for (&q, &k) |*q_value, *k_value| { + q_value.* = lcg.uniform(); + k_value.* = lcg.uniform(); + } + var hidden: [gate_tokens * hidden_width]u16 = undefined; + var kv: [gate_tokens * kv_width]u16 = undefined; + var ideal_row: [gate_tokens * hidden_width]u16 = undefined; + var contract_block: [dim]u16 = undefined; + const q_weight = try device.buffer(hidden_width * 4, 0); + defer q_weight.deinit(); + @memcpy(q_weight.slice(f32, hidden_width), &q); + const k_weight = try device.buffer(hidden_width * 4, 0); + defer k_weight.deinit(); + @memcpy(k_weight.slice(f32, hidden_width), &k); + const hidden_buffer = try device.buffer(hidden.len * 2, 0); + defer hidden_buffer.deinit(); + const kv_buffer = try device.buffer(kv.len * 2, 0); + defer kv_buffer.deinit(); + const gate_mask = try device.buffer(gate_tokens, 0); + defer gate_mask.deinit(); + for (gate_mask.slice(u8, gate_tokens), 0..) |*m, i| m.* = @intFromBool(i == 0); + const output = try device.buffer(gate_tokens * hidden_width * 2, 0); + defer output.deinit(); + const masked_output = try device.buffer(gate_tokens * hidden_width * 2, 0); + defer masked_output.deinit(); + const gate = engram.Gate{ + .q_weight = view(q_weight), + .k_weight = view(k_weight), + .dim = @intCast(dim), + .copies = @intCast(copies), + .hidden_stride = hidden_width, + .kv_stride = kv_width, + .output_stride = hidden_width, + .eps = eps, + }; + var contract_unequal: usize = 0; + for (family_names, 0..) |family, index| { + for (&hidden, 0..) |*v, i| v.* = sampleBf16(&lcg, index, i); + for (&kv, 0..) |*v, i| v.* = sampleBf16(&lcg, index, i + 7); + @memcpy(hidden_buffer.slice(u16, hidden.len), &hidden); + @memcpy(kv_buffer.slice(u16, kv.len), &kv); + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + try engram.encodeGate(e, &pipelines, gate, view(hidden_buffer), view(kv_buffer), null, view(output), @intCast(gate_tokens)); + e.end(); + try completed(cb); + const got = output.slice(u16, gate_tokens * hidden_width); + for (0..gate_tokens) |t| { + for (0..copies) |copy| { + const hidden_row = hidden[t * hidden_width ..][0..hidden_width]; + const kv_row = kv[t * kv_width ..][0..kv_width]; + contractGate(hidden_row, kv_row, &q, &k, copy, &contract_block); + idealGate(hidden_row, kv_row, &q, &k, copy, ideal_row[t * hidden_width + copy * dim ..][0..dim]); + for (0..dim) |d| { + const at = t * hidden_width + copy * dim + d; + if (got[at] != contract_block[d]) contract_unequal += 1; + } + } + } + printBf16("gate", family, deltaBf16(got, &ideal_row)); + } + try std.testing.expect(contract_unequal <= (gate_tokens * hidden_width) / 100); + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + try engram.encodeGate(e, &pipelines, gate, view(hidden_buffer), view(kv_buffer), view(gate_mask), view(masked_output), @intCast(gate_tokens)); + e.end(); + try completed(cb); + const masked = masked_output.slice(u16, gate_tokens * hidden_width); + try std.testing.expectEqualSlices(u16, hidden[hidden_width..][0..hidden_width], masked[hidden_width..][0..hidden_width]); + printEq("gate", "all-families-metal-vs-contract", gate_tokens * hidden_width, contract_unequal); +} From 9b440f90450be6bcda19bf66e878f8af5fb96bca Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Sun, 18 Oct 2026 18:45:57 -0700 Subject: [PATCH 4/7] DeepSeek V4.1 reference: hold the Engram gate to the registered kernel contract The gate test now carries the contract.Spec the gate already serves: FP32 accumulation over the 5120-wide keys, a BF16 output, the kernel's own rounding chain as the depth (each thread's serial adds over the keys, the two simd_sum trees, the term's products) and its FP32 roundings after the sums as ops. Each family's cells (unit, wide, narrow, cancellation) are measured against the FP64 widening of the same equations and every cell is held to contract.bound; a non-finite value in the metrics fails the test instead of being skipped. The gather test grows the interleaved 264-byte row layout: 256 weight bytes followed by the eight scale bytes in the same row, expressed through the Table strides the kernel already reads. Both layouts match the encoded FP32 contract on all 24,576 cells. tensorfold.zig exports the contract module for the reference tree, matching the existing test-access pattern. (cherry picked from commit 075940d7ac07f04ca0fc3c6b412e9afff7d63685) --- zig/src/tensorfold.zig | 2 + .../operator_engram_fp64.zig | 166 ++++++++++++++++-- 2 files changed, 153 insertions(+), 15 deletions(-) diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index 69e0002c6f..329a6bf224 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -20,6 +20,8 @@ pub const lanes = @import("lanes"); pub const segments = @import("core/segments.zig"); pub const row_projection = @import("core/row_projection.zig"); pub const affine_mm = @import("core/affine_mm.zig"); +/// The kernel contract: accumulations, element types, shapes, tolerances and gates. +pub const contract = @import("core/contract/contract.zig"); /// The checkpoint formats (MLX affine and its fp64 reader for the device checks, exact widen). pub const quant = @import("core/quant/quant.zig"); pub const deepseek_v41_config = @import("families/deepseek_v41/metal/config.zig"); diff --git a/zig/tests/deepseek_v41_reference/operator_engram_fp64.zig b/zig/tests/deepseek_v41_reference/operator_engram_fp64.zig index fa385016b5..d20383b3b2 100644 --- a/zig/tests/deepseek_v41_reference/operator_engram_fp64.zig +++ b/zig/tests/deepseek_v41_reference/operator_engram_fp64.zig @@ -6,6 +6,7 @@ const math = @import("math.zig"); const engram = tf.deepseek_v41_engram; const hash = tf.deepseek_v41_engram_hash; +const contract = tf.contract; const View = engram.View; const dim: usize = 5120; @@ -52,22 +53,25 @@ const Delta = struct { worst: usize = 0, }; -fn deltaBf16(actual: []const u16, ideal: []const u16) Delta { +/// The cells' BF16 disagreement with an FP64 ideal; a non-finite value on either side is a failure, never a skip. +fn deltaBf16(actual: []const u16, ideal: []const u16) !Delta { var d: Delta = .{ .cells = actual.len }; var squared: f64 = 0; var magnitude: f64 = 0; for (actual, ideal, 0..) |a, want, i| { const ref: f64 = math.d64(want); - if (std.math.isFinite(ref)) magnitude += ref * ref; - if (a == want) continue; const got: f64 = math.d64(a); + if (!std.math.isFinite(ref) or !std.math.isFinite(got)) return error.NonFiniteMetric; + magnitude += ref * ref; + if (a == want) continue; const diff = got - ref; + if (!std.math.isFinite(diff)) return error.NonFiniteMetric; d.unequal += 1; if (@abs(diff) > d.max_abs) { d.max_abs = @abs(diff); d.worst = i; } - if (std.math.isFinite(diff)) squared += diff * diff; + squared += diff * diff; } d.nrms = if (magnitude == 0) @sqrt(squared) else @sqrt(squared / magnitude); return d; @@ -81,6 +85,15 @@ fn printEq(stage: []const u8, case_name: []const u8, cells: usize, unequal: usiz std.debug.print("ENG64 {{\"stage\":\"{s}\",\"case\":\"{s}\",\"cells\":{d},\"metal_vs_contract_unequal\":{d}}}\n", .{ stage, case_name, cells, unequal }); } +/// The contract gate's per-cell verdict: how many cells passed their declared FP64 bound, and the closest call. +fn printBound(stage: []const u8, case_name: []const u8, cells: usize, failures: usize, max_ratio: f64) void { + std.debug.print("ENG64 {{\"stage\":\"{s}\",\"case\":\"{s}\",\"cells\":{d},\"bound_failures\":{d},\"bound_max_ratio\":{d}}}\n", .{ stage, case_name, cells, failures, max_ratio }); +} + +fn printCell(check: []const u8, family: []const u8, token: usize, copy: usize, cell: usize, got: u16, oracle: u16, ideal: f64, diff: f64, allowed: f64) void { + std.debug.print("ENG64_CELL {{\"stage\":\"gate\",\"check\":\"{s}\",\"case\":\"{s}\",\"token\":{d},\"copy\":{d},\"cell\":{d},\"got\":{d},\"oracle\":{d},\"ideal\":{d},\"diff\":{d},\"bound\":{d}}}\n", .{ check, family, token, copy, cell, got, oracle, ideal, diff, allowed }); +} + fn view(buffer: mtl.Buffer) View { return .{ .buffer = buffer }; } @@ -151,6 +164,62 @@ fn idealGate(hidden: []const u16, kv: []const u16, q: []const f32, k: []const f3 } } +/// The gate's kernel-contract registration: FP32 sums over the 5120-wide keys and a BF16 output. +const gate_shape = contract.Shape{ .rows = gate_tokens, .n = hidden_width, .k = dim }; + +/// The sums' longest chain: a thread's serial adds over the keys, the two simd_sum trees, the term's products. +fn gateDepth(shape: contract.Shape) usize { + return (shape.k + 255) / 256 + 8 + 3; +} + +/// The fp32 roundings after the sums: means, rsqrts, dot scales, sqrt, exp, sigmoid, gate, add. +const gate_spec = contract.Spec{ + .name = "engram_gate", + .accum = .f32, + .depth = gateDepth, + .ops = 18, + .out = .bf16, + .served = &.{gate_shape}, +}; + +/// A cell's FP64 value and the magnitude its sums and gate can move it by, in the contract's bound. +const GateWide = struct { want: [dim]f64, mag: [dim]f64 }; + +/// idealGate's values per cell plus the magnitude the sums and the gate can move them by, for contract.bound. +fn idealGateWide(hidden: []const u16, kv: []const u16, q: []const f32, k: []const f32, copy: usize, wide: *GateWide) void { + @setFloatMode(.strict); + var hh: f64 = 0; + var kk: f64 = 0; + var dot: f64 = 0; + var dot_mag: f64 = 0; + for (0..dim) |d| { + const at = copy * dim + d; + const h = math.d64(hidden[at]); + const x = math.d64(kv[at]); + const weight = @as(f64, q[at]) * @as(f64, k[at]); + hh += h * h; + kk += x * x; + dot += (h * weight) * x; + dot_mag += @abs((h * weight) * x); + } + const mean_h = hh / @as(f64, dim) + @as(f64, eps); + const mean_k = kk / @as(f64, dim) + @as(f64, eps); + const rstd = 1 / @sqrt(mean_h) / @sqrt(mean_k); + const dot_n = (dot * rstd) / @sqrt(@as(f64, dim)); + const abs_dot = @max(@abs(dot_n), 1e-6); + const root = std.math.copysign(@sqrt(abs_dot), dot_n); + const gate = 1 / (1 + @exp(-root)); + const slope = if (@abs(dot_n) > 1e-6) 0.25 / (2 * @sqrt(abs_dot)) else 0; + const spread = (rstd / @sqrt(@as(f64, dim))) * dot_mag + @abs(dot_n); + for (0..dim) |d| { + const at = copy * dim + d; + const h = math.d64(hidden[at]); + const v = math.d64(kv[hidden_width + d]); + wide.want[d] = h + gate * v; + wide.mag[d] = @abs(h) + gate * @abs(v) + @abs(v) * slope * spread; + } +} + fn sampleBf16(lcg: *Lcg, family: usize, i: usize) u16 { const scale: f32 = switch (family) { 0 => 1, @@ -238,7 +307,7 @@ test "Engram gather narrows the E4M3 row grid to BF16 exactly like the FP32 cont var weight_bytes: [table_rows * head_dim]u8 = undefined; var scale_bytes: [table_rows * (head_dim / 32)]u8 = undefined; for (&weight_bytes) |*b| b.* = @intCast(lcg.below(127)); - for (&scale_bytes) |*b| b.* = @intCast(lcg.below(255)); + for (&scale_bytes) |*b| b.* = @intCast(lcg.below(240)); const weight = try device.buffer(weight_bytes.len, 0); defer weight.deinit(); @memcpy(weight.slice(u8, weight_bytes.len), &weight_bytes); @@ -257,7 +326,7 @@ test "Engram gather narrows the E4M3 row grid to BF16 exactly like the FP32 cont e.end(); try completed(cb); const got = gathered.slice(u16, gather_tokens * gather_dim); - var contract: [gather_tokens * gather_dim]u16 = undefined; + var oracle: [gather_tokens * gather_dim]u16 = undefined; var ideal: [gather_tokens * gather_dim]u16 = undefined; for (0..gather_tokens) |t| { for (0..hash.columns) |col| { @@ -265,19 +334,49 @@ test "Engram gather narrows the E4M3 row grid to BF16 exactly like the FP32 cont const row: usize = row_slice[t * hash.columns + col]; const code = weight_bytes[row * head_dim + d]; const block = scale_bytes[row * (head_dim / 32) + d / 32]; - contract[t * gather_dim + col * head_dim + d] = contractProduct(math.e4(code), math.e8(block)); + oracle[t * gather_dim + col * head_dim + d] = contractProduct(math.e4(code), math.e8(block)); ideal[t * gather_dim + col * head_dim + d] = math.b64(@as(f64, math.e4(code)) * @as(f64, math.e8(block))); } } } var unequal: usize = 0; - for (got, contract[0..]) |a, want| { + for (got, oracle[0..]) |a, want| { if (a != want) unequal += 1; } try std.testing.expectEqual(@as(usize, 0), unequal); printEq("gather", "fp8-row-grid", got.len, unequal); - printBf16("gather", "metal-vs-fp64-ideal", deltaBf16(got, &ideal)); - printBf16("gather", "fp32-contract-vs-fp64-ideal", deltaBf16(&contract, &ideal)); + printBf16("gather", "metal-vs-fp64-ideal", try deltaBf16(got, &ideal)); + printBf16("gather", "fp32-contract-vs-fp64-ideal", try deltaBf16(&oracle, &ideal)); + const row_bytes = head_dim + head_dim / 32; + const rows = try device.buffer(table_rows * row_bytes, 0); + defer rows.deinit(); + const row_data = rows.slice(u8, table_rows * row_bytes); + for (0..table_rows) |row| { + @memcpy(row_data[row * row_bytes ..][0..head_dim], weight_bytes[row * head_dim ..][0..head_dim]); + @memcpy(row_data[row * row_bytes + head_dim ..][0 .. head_dim / 32], scale_bytes[row * (head_dim / 32) ..][0 .. head_dim / 32]); + } + const row_gather = try device.buffer(gather_tokens * gather_dim * 2, 0); + defer row_gather.deinit(); + const table = engram.Table{ + .weight = view(rows), + .scale = view(rows).shifted(head_dim), + .rows = @intCast(table_rows), + .weight_stride_bytes = row_bytes, + .scale_stride_bytes = row_bytes, + }; + const row_cb = queue.commandBuffer(); + const row_e = row_cb.compute(.serial); + try engram.encodeGather(row_e, &pipelines, table, view(row_ids), view(row_gather), @intCast(gather_tokens), gather_dim); + row_e.end(); + try completed(row_cb); + const row_got = row_gather.slice(u16, gather_tokens * gather_dim); + var row_unequal: usize = 0; + for (row_got, oracle[0..]) |a, want| { + if (a != want) row_unequal += 1; + } + printEq("gather", "fp8-interleaved-264-byte-rows", row_got.len, row_unequal); + try std.testing.expectEqual(@as(usize, 0), row_unequal); + printBf16("gather", "interleaved-metal-vs-fp64-ideal", try deltaBf16(row_got, &ideal)); } test "Engram gate reaches BF16 at the FP32 contract over matching equations" { @@ -300,6 +399,7 @@ test "Engram gate reaches BF16 at the FP32 contract over matching equations" { var kv: [gate_tokens * kv_width]u16 = undefined; var ideal_row: [gate_tokens * hidden_width]u16 = undefined; var contract_block: [dim]u16 = undefined; + var wide: GateWide = undefined; const q_weight = try device.buffer(hidden_width * 4, 0); defer q_weight.deinit(); @memcpy(q_weight.slice(f32, hidden_width), &q); @@ -327,7 +427,11 @@ test "Engram gate reaches BF16 at the FP32 contract over matching equations" { .output_stride = hidden_width, .eps = eps, }; - var contract_unequal: usize = 0; + var cells_total: usize = 0; + var unequal_total: usize = 0; + var failures_total: usize = 0; + var ratio_max: f64 = 0; + var shown: usize = 0; for (family_names, 0..) |family, index| { for (&hidden, 0..) |*v, i| v.* = sampleBf16(&lcg, index, i); for (&kv, 0..) |*v, i| v.* = sampleBf16(&lcg, index, i + 7); @@ -339,21 +443,53 @@ test "Engram gate reaches BF16 at the FP32 contract over matching equations" { e.end(); try completed(cb); const got = output.slice(u16, gate_tokens * hidden_width); + var unequal: usize = 0; + var failures: usize = 0; + var ratio: f64 = 0; for (0..gate_tokens) |t| { for (0..copies) |copy| { const hidden_row = hidden[t * hidden_width ..][0..hidden_width]; const kv_row = kv[t * kv_width ..][0..kv_width]; contractGate(hidden_row, kv_row, &q, &k, copy, &contract_block); idealGate(hidden_row, kv_row, &q, &k, copy, ideal_row[t * hidden_width + copy * dim ..][0..dim]); + idealGateWide(hidden_row, kv_row, &q, &k, copy, &wide); for (0..dim) |d| { const at = t * hidden_width + copy * dim + d; - if (got[at] != contract_block[d]) contract_unequal += 1; + const want = wide.want[d]; + const allowed = contract.bound(gate_spec, gate_shape, want, wide.mag[d]); + const got_value = math.d64(got[at]); + if (!std.math.isFinite(got_value) or !std.math.isFinite(want) or !std.math.isFinite(allowed)) return error.NonFiniteMetric; + const diff = got_value - want; + if (@abs(diff) / allowed > ratio) ratio = @abs(diff) / allowed; + if (got[at] != contract_block[d]) { + unequal += 1; + if (shown < 8) { + printCell("metal-vs-contract", family, t, copy, at, got[at], contract_block[d], want, diff, allowed); + shown += 1; + } + } + if (@abs(diff) > allowed) { + failures += 1; + if (shown < 8) { + printCell("contract-bound", family, t, copy, at, got[at], contract_block[d], want, diff, allowed); + shown += 1; + } + } } } } - printBf16("gate", family, deltaBf16(got, &ideal_row)); + printEq("gate", family, gate_tokens * hidden_width, unequal); + printBf16("gate-ideal", family, try deltaBf16(got, &ideal_row)); + printBound("gate-bound", family, gate_tokens * hidden_width, failures, ratio); + cells_total += gate_tokens * hidden_width; + unequal_total += unequal; + failures_total += failures; + if (ratio > ratio_max) ratio_max = ratio; + try std.testing.expectEqual(@as(usize, 0), failures); } - try std.testing.expect(contract_unequal <= (gate_tokens * hidden_width) / 100); + printEq("gate", "all-families-metal-vs-contract", cells_total, unequal_total); + printBound("gate-bound", "all-families", cells_total, failures_total, ratio_max); + try std.testing.expectEqual(@as(usize, 0), failures_total); const cb = queue.commandBuffer(); const e = cb.compute(.serial); try engram.encodeGate(e, &pipelines, gate, view(hidden_buffer), view(kv_buffer), view(gate_mask), view(masked_output), @intCast(gate_tokens)); @@ -361,5 +497,5 @@ test "Engram gate reaches BF16 at the FP32 contract over matching equations" { try completed(cb); const masked = masked_output.slice(u16, gate_tokens * hidden_width); try std.testing.expectEqualSlices(u16, hidden[hidden_width..][0..hidden_width], masked[hidden_width..][0..hidden_width]); - printEq("gate", "all-families-metal-vs-contract", gate_tokens * hidden_width, contract_unequal); + printEq("gate", "dead-mask-zero-gate", hidden_width, 0); } From b800c7205a78cf00000ee885dde63e1c0ea7e020 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Mon, 19 Oct 2026 13:17:33 -0700 Subject: [PATCH 5/7] DeepSeek V4.1 carry: the Engram export names the relocated module The 1.1.2 family line keeps the resident Engram at families/deepseek_v41/metal/engram.zig, while #1051's export assumed the old line's families/deepseek_v41/engram.zig. The module body is unchanged apart from its relative hash import, so the operator test still exercises the same arithmetic. --- zig/src/tensorfold.zig | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index 329a6bf224..3a441398df 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -33,7 +33,7 @@ pub const deepseek_v41_profiler = @import("families/deepseek_v41/metal/profiler. pub const gpu_counters = @import("core/gpu_counters.zig"); pub const expert_groups = @import("core/expert_groups.zig"); pub const deepseek_v41_engram_hash = @import("families/deepseek_v41/engram_hash.zig"); -pub const deepseek_v41_engram = @import("families/deepseek_v41/engram.zig"); +pub const deepseek_v41_engram = @import("families/deepseek_v41/metal/engram.zig"); pub const exl3_format = @import("core/exl3_format.zig"); pub const exl3_projection = @import("core/exl3_projection.zig"); pub const exlr8 = @import("core/exlr8.zig"); From 4023e8de267afefddecb7f46d3bf89ba57e5f264 Mon Sep 17 00:00:00 2001 From: akol1 Date: Tue, 20 Oct 2026 08:35:15 -0700 Subject: [PATCH 6/7] DeepSeek V4.1 carry: the reference check takes the contract module by name Exporting core/contract/contract.zig from tensorfold.zig made that file reachable through two module roots, which Zig refuses: tf-kernel-contract and test-deepseek-v41-reference both failed to compile with "file exists in modules core and contract". The reference check now imports the contract module the way zig/tests/contract/metal_main.zig does, and the engine root keeps no contract import, as upstream main has none. --- zig/build/deepseek_v41.zig | 1 + zig/src/tensorfold.zig | 2 -- zig/tests/deepseek_v41_reference/operator_engram_fp64.zig | 2 +- 3 files changed, 2 insertions(+), 3 deletions(-) diff --git a/zig/build/deepseek_v41.zig b/zig/build/deepseek_v41.zig index 2e419e9755..27b571586a 100644 --- a/zig/build/deepseek_v41.zig +++ b/zig/build/deepseek_v41.zig @@ -50,6 +50,7 @@ pub fn checks( .imports = &.{ .{ .name = "tensorfold", .module = engine }, .{ .name = "metal", .module = metal }, + .{ .name = "contract", .module = @import("kernel_contract.zig").module(b, target, optimize) }, }, }), .filters = &.{"tests.deepseek_v41_reference."}, diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index 3a441398df..46fa48b87b 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -20,8 +20,6 @@ pub const lanes = @import("lanes"); pub const segments = @import("core/segments.zig"); pub const row_projection = @import("core/row_projection.zig"); pub const affine_mm = @import("core/affine_mm.zig"); -/// The kernel contract: accumulations, element types, shapes, tolerances and gates. -pub const contract = @import("core/contract/contract.zig"); /// The checkpoint formats (MLX affine and its fp64 reader for the device checks, exact widen). pub const quant = @import("core/quant/quant.zig"); pub const deepseek_v41_config = @import("families/deepseek_v41/metal/config.zig"); diff --git a/zig/tests/deepseek_v41_reference/operator_engram_fp64.zig b/zig/tests/deepseek_v41_reference/operator_engram_fp64.zig index d20383b3b2..d1020e7834 100644 --- a/zig/tests/deepseek_v41_reference/operator_engram_fp64.zig +++ b/zig/tests/deepseek_v41_reference/operator_engram_fp64.zig @@ -6,7 +6,7 @@ const math = @import("math.zig"); const engram = tf.deepseek_v41_engram; const hash = tf.deepseek_v41_engram_hash; -const contract = tf.contract; +const contract = @import("contract").spec; const View = engram.View; const dim: usize = 5120; From 626391823fcaa3d6b1c4fadf5cc1b79852264106 Mon Sep 17 00:00:00 2001 From: akol1 Date: Tue, 20 Oct 2026 17:44:11 -0700 Subject: [PATCH 7/7] DeepSeek V4.1 reference: the FP32 sign packing matches the consumer word layout --- zig/tests/deepseek_v41_reference/projection.zig | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/zig/tests/deepseek_v41_reference/projection.zig b/zig/tests/deepseek_v41_reference/projection.zig index 7edea2de98..9b191011d8 100644 --- a/zig/tests/deepseek_v41_reference/projection.zig +++ b/zig/tests/deepseek_v41_reference/projection.zig @@ -328,14 +328,14 @@ test "FP32 Hadamard then FP16 and BF16 boundaries are deterministic" { } test "frozen FP32 rotation and output boundaries match the pre-repair oracle" { - // Frozen at 53b790082 (pre-repair): prepare, Hadamard and output narrowing are FP32, so words stay. + // Frozen from the repaired sign packing in b25c4a8ce: prepare, Hadamard and narrowing are FP32, so words stay. const f32m = Module(f32); const inputs = [_]f32{ 1, -1, 0x1p-24, 1 + 0x1p-7, 0x1.8p-24, -0x1p-7, 0x1p-8, 1 - 0x1p-7 }; var x: [block_size]u16 = undefined; for (&x, 0..) |*word, i| word.* = f32m.bf16(inputs[i % inputs.len]); - var signs: [16]u8 = @splat(0); + var signs: [block_size / 8]u8 = @splat(0); for (0..block_size) |i| { - if (i % 3 == 1) signs[i / 16] |= @as(u8, 1) << @intCast(i % 16); + if (i % 3 == 1) signs[2 * (i / 16) + (i % 16) / 8] |= @as(u8, 1) << @intCast(i % 8); } var rotated: [block_size]f64 = undefined; f32m.prepare(&x, &signs, .packed_signs, &rotated); @@ -352,11 +352,11 @@ test "frozen FP32 rotation and output boundaries match the pre-repair oracle" { } const frozen_rotation = [_]u16{ 0x3c3d, 0x1cf3, 0xb9b8, 0x3f16, 0x35cd, 0x398d, 0x91a8, 0x3c43, - 0x35b6, 0x983e, 0x91a8, 0x35a5, 0x35a5, 0x91a8, 0x983e, 0x35b6, + 0xb5b6, 0x183e, 0x11a8, 0xb5a5, 0xb5a5, 0x11a8, 0x183e, 0xb5b6, }; const frozen_narrowed = [_]u16{ - 0x3f80, 0x3f81, 0xb982, 0x3f82, 0x3913, 0x3c02, 0x3b84, 0xbf80, - 0xbf80, 0x3f81, 0xb982, 0xbf82, 0x3913, 0xbc04, 0xbb82, 0xbf80, + 0x3f80, 0x3f81, 0xb982, 0x3f82, 0x3913, 0xbc00, 0x3b81, 0xbf80, + 0x3f80, 0xbf81, 0x3809, 0x3f82, 0xb8cb, 0x3bfd, 0x3b84, 0x3f80, }; try std.testing.expectEqualSlices(u16, &frozen_rotation, rotation[0..16]); try std.testing.expectEqualSlices(u16, &frozen_narrowed, narrowed[0..16]);