diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index 018947342d..5b34839762 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -217,6 +217,13 @@ TensorFold ships no model weights. The `z-lab/Qwen3.8-27B-DFlash2` model card st The optional `incoai/GLM-5.3-Flash-DFlash2` model card states CC BY-NC-ND 4.0, for non-commercial use without derivatives; only the Python engine reads that drafter. Each checkpoint keeps its own license. +The native `zig/src/families/deepseek_v41/` implementation and `zig/tests/deepseek_v41_reference/` +follow [DeepSeek-V4.1-Flash's official inference](https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash/tree/2cba9e42aa026125f3ed06c6d98c1db82f7ca027/inference), +MIT License, Copyright (c) 2023 DeepSeek: the reference's equations, the Engram normalization, +PCG64 multipliers and prime bucket layout, and the pinned chat template come from that revision. +The EXL3 checkpoint reads through the ExLlamaV3 format covered above; no captured CUDA/Triton +kernel assets, full model shards or binary traces are bundled. + ## Structured output: libtfgrammar `lib/libtfgrammar` in the release archives is `zig/src/grammar/tf_grammar.cc` compiled with the C++ sources of diff --git a/build.zig b/build.zig index b12fb3c40f..b8884fe5f0 100644 --- a/build.zig +++ b/build.zig @@ -208,11 +208,10 @@ pub fn build(b: *std.Build) void { } /// The native Metal server build is owned by the shared Metal target module. +const metalTargets = metal_build.metalTargets; const nativeServer = metal_build.nativeServer; /// macOS: the metallib, the Metal runtime's test programs and the engine over Metal. -const metalTargets = metal_build.metalTargets; - /// The server and distribution use the same embedded-source Metal engine. const metalEngineModules = metal_build.metalEngineModules; diff --git a/tools/zig/lean_check.py b/tools/zig/lean_check.py index 6b85198058..1e1e2cd8d1 100644 --- a/tools/zig/lean_check.py +++ b/tools/zig/lean_check.py @@ -12,7 +12,7 @@ ".sh": re.compile(r"^\s*#(?!!)")} BLOCK = re.compile(r"/\*") DOCSTRING = re.compile(r'^\s*(?:[rubf]{0,2})("""|\'\'\')') -GENERATED = ("unicode_data.zig",) +GENERATED = ("unicode_data.zig", "nfkd_data.zig") HEX_METAL = re.compile(r"_[0-9a-f]{16}\.metal$") # Generated outputs are skipped entirely. unicode_data.zig has only the length exemption. GENERATOR_OUTPUT = ( diff --git a/zig/build/deepseek_v41.zig b/zig/build/deepseek_v41.zig new file mode 100644 index 0000000000..a1ae4ad2d1 --- /dev/null +++ b/zig/build/deepseek_v41.zig @@ -0,0 +1,51 @@ +//! Experimental DS4.1 weight-loading qualification targets. +const std = @import("std"); + +fn withCore(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, root: []const u8, core: *std.Build.Module, imports: []const std.Build.Module.Import) *std.Build.Module { + var list: std.ArrayList(std.Build.Module.Import) = .empty; + list.append(b.allocator, .{ .name = "core", .module = core }) catch @panic("oom"); + list.appendSlice(b.allocator, imports) catch @panic("oom"); + return b.createModule(.{ .root_source_file = b.path(root), .target = target, .optimize = optimize, .link_libc = true, .imports = list.items }); +} + +pub fn weights( + b: *std.Build, + target: std.Build.ResolvedTarget, + optimize: std.builtin.OptimizeMode, + metal: *std.Build.Module, + core: *std.Build.Module, +) *std.Build.Step { + const host = b.addTest(.{ .root_module = withCore(b, target, optimize, "zig/deepseek_v41_weights_checks.zig", core, &.{ + .{ .name = "metal", .module = metal }, + }) }); + const step = b.step("test-deepseek-v41-weights", "V4.1 configuration, Engram table plans and resident weight staging"); + step.dependOn(&b.addRunArtifact(host).step); + return step; +} + +pub fn checks( + b: *std.Build, + target: std.Build.ResolvedTarget, + optimize: std.builtin.OptimizeMode, + metal: *std.Build.Module, + sources: *std.Build.Module, + engine: *std.Build.Module, + core: *std.Build.Module, +) struct { device: *std.Build.Step, reference: *std.Build.Step } { + const device = b.addTest(.{ .root_module = withCore(b, target, optimize, "zig/deepseek_v41_checks.zig", core, &.{ + .{ .name = "metal", .module = metal }, + .{ .name = "kernel_sources", .module = sources }, + }) }); + const device_step = b.step("test-deepseek-v41", "Experimental DS4.1 EXL3, Engram, mHC and sparse attention on Metal"); + device_step.dependOn(&b.addRunArtifact(device).step); + const reference = b.addTest(.{ + .root_module = withCore(b, target, optimize, "zig/deepseek_v41_reference_checks.zig", core, &.{ + .{ .name = "tensorfold", .module = engine }, + .{ .name = "metal", .module = metal }, + }), + .filters = &.{"tests.deepseek_v41_reference."}, + }); + const reference_step = b.step("test-deepseek-v41-reference", "Independent CPU teacher-oracle arithmetic boundaries"); + reference_step.dependOn(&b.addRunArtifact(reference).step); + return .{ .device = device_step, .reference = reference_step }; +} diff --git a/zig/build/dsv41_shared.zig b/zig/build/dsv41_shared.zig index 0751a1100f..a9b5a18f72 100644 --- a/zig/build/dsv41_shared.zig +++ b/zig/build/dsv41_shared.zig @@ -2,12 +2,16 @@ const std = @import("std"); pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, file: []const u8) *std.Build.Module { return moduleAt(b, target, .Debug, file); } -pub fn moduleAt(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, file: []const u8) *std.Build.Module { +pub fn coreModule(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *std.Build.Module { const tok = b.createModule(.{ .root_source_file = b.path("zig/src/core/tokenizer/tokenizer.zig"), .target = target, .optimize = optimize, .link_libc = true }); - const core = b.createModule(.{ .root_source_file = b.path("zig/src/core/root.zig"), .target = target, .optimize = optimize, .link_libc = true }); + const core = b.createModule(.{ .root_source_file = b.path("zig/dsv41_core_family.zig"), .target = target, .optimize = optimize, .link_libc = true }); + core.addImport("metal", b.createModule(.{ .root_source_file = b.path("zig/src/metal/metal.zig"), .target = target, .optimize = optimize, .link_libc = true })); core.addImport("gpu_memory", @import("gpu_memory.zig").module(b, target)); core.addImport("tokenizer", tok); - return b.createModule(.{ .root_source_file = b.path(b.fmt("zig/src/families/deepseek_v41/{s}", .{file})), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "core", .module = core }} }); + return core; +} +pub fn moduleAt(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, file: []const u8) *std.Build.Module { + return b.createModule(.{ .root_source_file = b.path(b.fmt("zig/src/families/deepseek_v41/{s}", .{file})), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "core", .module = coreModule(b, target, optimize) }} }); } pub fn steps(b: *std.Build) void { const host = b.addTest(.{ .root_module = module(b, b.graph.host, "shared_test.zig") }); diff --git a/zig/build/metal_targets.zig b/zig/build/metal_targets.zig index d5ab369dd7..9174c01215 100644 --- a/zig/build/metal_targets.zig +++ b/zig/build/metal_targets.zig @@ -1,5 +1,6 @@ //! Metal build targets and their shared source/library/module construction. const std = @import("std"); +const dsv41_build = @import("deepseek_v41.zig"); const builtin = @import("builtin"); const qwen27_build = @import("qwen27.zig"); const core_checks = @import("core_checks.zig"); @@ -170,6 +171,7 @@ pub fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s const programs = [_]struct { name: []const u8, path: []const u8, about: []const u8 }{ .{ .name = "tf-metal-tests", .path = "zig/tests/gpu_tests.zig", .about = "GPU runtime tests" }, .{ .name = "tf-dispatch-bench", .path = "zig/tests/dispatch_bench.zig", .about = "Dispatch boundary benchmark" }, + .{ .name = "tf-dsv41-run", .path = "zig/deepseek_v41_run.zig", .about = "Full-resident DeepSeek V4.1 admission, teacher-forced logits, chunk seams and decode timing" }, .{ .name = "tf-kernel-exact", .path = "zig/tests/kernel_exact.zig", .about = "Bit-exact check of a kernel against an MLX oracle case" }, }; for (programs) |p| { @@ -276,6 +278,12 @@ pub fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s const load_admission_run = b.addRunArtifact(load_admission_cases); b.step("test-flashnext-load-admission", "CPU-only Flash Next runtime quantization admission").dependOn(&load_admission_run.step); test_step.dependOn(&load_admission_run.step); + // The family reaches core through `@import("core")`; this facade owns exactly the files it names. + const core_root = b.createModule(.{ .root_source_file = b.path("zig/dsv41_core_family.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "gpu_memory", .module = @import("gpu_memory.zig").module(b, target) }, .{ .name = "tokenizer", .module = b.createModule(.{ .root_source_file = b.path("zig/src/core/tokenizer/tokenizer.zig"), .target = target, .optimize = optimize, .link_libc = true }) }, .{ .name = "metal", .module = metal } } }); + test_step.dependOn(dsv41_build.weights(b, target, optimize, metal, core_root)); + const dsv41 = dsv41_build.checks(b, target, optimize, metal, mods.sources, mods.engine, core_root); + test_step.dependOn(dsv41.device); + test_step.dependOn(dsv41.reference); const pack_cache_cases = b.addTest(.{ .root_module = b.createModule(.{ .root_source_file = b.path("zig/flashnext_pack_cache_cases.zig"), .target = target, diff --git a/zig/deepseek_v41_checks.zig b/zig/deepseek_v41_checks.zig new file mode 100644 index 0000000000..d80b25bbc0 --- /dev/null +++ b/zig/deepseek_v41_checks.zig @@ -0,0 +1,14 @@ +//! Source-built V4.1 operator qualification; real Metal execution, no checkpoint needed. +test { + _ = @import("core").exl3_projection_gpu_test; + _ = @import("src/families/deepseek_v41/engram.zig"); + _ = @import("src/families/deepseek_v41/csa2.zig"); + _ = @import("src/families/deepseek_v41/ops_check.zig"); + _ = @import("src/families/deepseek_v41/load.zig"); + _ = @import("src/families/deepseek_v41/metal_model.zig"); + _ = @import("src/families/deepseek_v41/index_precision_test.zig"); + _ = @import("src/families/deepseek_v41/config.zig"); + _ = @import("src/families/deepseek_v41/profiler.zig"); + _ = @import("core").gpu_counters; + _ = @import("core").expert_groups; +} diff --git a/zig/deepseek_v41_reference_checks.zig b/zig/deepseek_v41_reference_checks.zig new file mode 100644 index 0000000000..61aa30f788 --- /dev/null +++ b/zig/deepseek_v41_reference_checks.zig @@ -0,0 +1,7 @@ +//! Independent CPU arithmetic boundaries used by the resident-checkpoint teacher oracle. +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/qualify.zig"); +} diff --git a/zig/deepseek_v41_run.zig b/zig/deepseek_v41_run.zig new file mode 100644 index 0000000000..0b57148839 --- /dev/null +++ b/zig/deepseek_v41_run.zig @@ -0,0 +1,313 @@ +//! Full-checkpoint admission and execution receipts; no synthetic weights or replay substitute. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const ds = tf.deepseek_v41; +const config = tf.deepseek_v41_config; +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 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"); +const attention_operators = @import("tests/deepseek_v41_reference/operator_attention.zig"); +const history_operators = @import("tests/deepseek_v41_reference/operator_history.zig"); +const hc_operators = @import("tests/deepseek_v41_reference/operator_hc.zig"); +const moe_operators = @import("tests/deepseek_v41_reference/operator_moe.zig"); + +fn usage(code: u8) noreturn { + std.debug.print("usage: tf-dsv41-run plan MODEL TABLES\n" ++ + " 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 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" ++ + " tf-dsv41-run attention-sameinput MODEL TABLES TRACE_CPU_BIN OUT_DIR [LAYER]\n" ++ + " tf-dsv41-run attention-history-sameinput MODEL TABLES TRACE_DIR OUT_DIR TOKEN LAYER\n" ++ + " tf-dsv41-run hc-sameinput MODEL TABLES TRACE_CPU_BIN OUT_DIR [LAYER]\n" ++ + " 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" ++ + "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); +} +fn number(s: []const u8) !u32 { + return std.fmt.parseInt(u32, s, 10); +} +fn validArguments(args: []const [:0]const u8) bool { + if (args.len < 2) return false; + const mode = args[1]; + for ([_][]const u8{ "plan", "operators" }, [_]usize{ 4, 5 }) |name, count| + if (std.mem.eql(u8, mode, name)) return args.len == count; + 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| + if (std.mem.eql(u8, mode, name)) return args.len == 6; + return false; +} +fn argmax(row: []const f32) !u32 { + var best: u32 = 0; + if (row.len == 0) return error.EmptyLogits; + var score = -std.math.inf(f32); + for (row, 0..) |v, i| { + if (!std.math.isFinite(v)) return error.NonfiniteLogits; + if (v > score) { + best = @intCast(i); + score = v; + } + } + return best; +} +fn write(io: std.Io, a: std.mem.Allocator, dir: []const u8, name: []const u8, data: []const u8) !void { + const path = try std.fs.path.join(a, &.{ dir, name }); + defer a.free(path); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = path, .data = data }); +} +const Counters = struct { disk_read: u64, pageins: u64, footprint: u64 }; +/// Darwin's rusage_info_v4 in full: flavor 4 writes 296 bytes (uuid plus 35 u64 fields), so a shorter buffer overruns. +const RUsage = extern struct { + uuid: [16]u8, + user_time: u64, + system_time: u64, + pkg_idle_wkups: u64, + interrupt_wkups: u64, + pageins: u64, + wired_size: u64, + resident_size: u64, + phys_footprint: u64, + proc_start_abstime: u64, + proc_exit_abstime: u64, + child_user_time: u64, + child_system_time: u64, + child_pkg_idle_wkups: u64, + child_interrupt_wkups: u64, + child_pageins: u64, + child_elapsed_abstime: u64, + diskio_bytesread: u64, + /// The 18 fields after `diskio_bytesread`, carried so the struct spans all 296 bytes Darwin writes. + reserved: [18]u64, + + comptime { + // The C interface fixes the size, so a shorter declaration is a build error rather than a silent overrun. + if (@sizeOf(RUsage) != 296) @compileError("RUsage must match Darwin's rusage_info_v4 (296 bytes)"); + } +}; +fn counters() !Counters { + const proc_pid_rusage = struct { + extern "c" fn proc_pid_rusage(pid: c_int, flavor: c_int, buffer: *anyopaque) c_int; + }.proc_pid_rusage; + const getpid = struct { + extern "c" fn getpid() c_int; + }.getpid; + var r: RUsage = undefined; + if (proc_pid_rusage(getpid(), 4, &r) != 0) return error.MemoryCountersUnavailable; + return .{ .disk_read = r.diskio_bytesread, .pageins = r.pageins, .footprint = r.phys_footprint }; +} +fn teacher(m: *ds.Model, ids: []const u32, chunk: u32, out: []f32) !f64 { + if (chunk == 0 or chunk > m.workspace.plan.rows or ids.len > m.cache.plan.capacity) + return error.InvalidTeacherWindow; + if (out.len != ids.len * 129280) return error.InvalidTeacherOutputLength; + try m.reset(); + var elapsed: f64 = 0; + var start: usize = 0; + while (start < ids.len) { + const end = @min(start + chunk, ids.len); + const row = try m.forward(ids[start..end], .all); + @memcpy(out[start * 129280 .. end * 129280], row); + for (row) |value| if (!std.math.isFinite(value)) return error.NonfiniteLogits; + elapsed += m.last_gpu_seconds; + start = end; + } + return elapsed; +} + +pub fn main(init: std.process.Init) !void { + const a = init.gpa; + const io = init.io; + const args = try init.minimal.args.toSlice(init.arena.allocator()); + if (args.len == 2 and (std.mem.eql(u8, args[1], "--help") or std.mem.eql(u8, args[1], "-h"))) usage(0); + if (!validArguments(args)) usage(2); + const arena = init.arena.allocator(); + const mode = args[1]; + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + if (std.mem.eql(u8, mode, "plan")) { + const cfg = try config.Config.load(arena, io, args[2]); + var plan = try load.Plan.open(a, io, args[2], args[3]); + defer plan.deinit(); + try weights.validate(&plan); + const device = try mtl.Device.init(); + defer device.deinit(); + std.debug.print("{{\"mode\":\"plan\",\"weight_bytes\":{d},\"shards\":{d},\"tensors\":{d},\"context\":{d},\"metal_working_set_bytes\":{d}}}\n", .{ plan.bytes, plan.shards.items.len, plan.tensors.count(), cfg.text_config.max_position_embeddings, device.maxWorkingSet() }); + return; + } + const is_operators = std.mem.eql(u8, mode, "operators"); + const is_attention = std.mem.eql(u8, mode, "attention-sameinput"); + 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 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_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) + try std.Io.Dir.cwd().readFileAlloc(io, map, arena, .limited(2 << 20)) + else + &.{}; + const model = try a.create(ds.Model); + defer a.destroy(model); + const begin = mtl.clock.seconds(); + model.* = try ds.Model.open(a, io, .{ .dir = args[2], .table_dir = args[3], .token_map = map }); + defer model.deinit(); + const ready = mtl.clock.seconds(); + const before = try counters(); + var execution_after: ?Counters = null; + var execution_error: ?anyerror = null; + execute(a, io, model, args, ids, map_blob, dir, before, &execution_after) catch |err| { + execution_error = err; + }; + const after = execution_after orelse try counters(); + if (after.disk_read < before.disk_read or after.pageins < before.pageins) return error.NonmonotonicCounters; + const resident_gate_passed = after.disk_read == before.disk_read; + const receipt = try std.json.Stringify.valueAlloc(a, .{ + .mode = mode, + .evidence_origin = "current_invocation_not_historical_qualification", + .context_capacity = model.cache.plan.capacity, + .consumed_tokens = model.cache.position, + .weight_bytes = model.checkpoint.residentBytes(), + .resident_allocation_bytes = model.residency.allocatedSize(), + .ready_footprint_bytes = before.footprint, + .final_footprint_bytes = after.footprint, + .load_seconds = ready - begin, + .measurement_scope = if (execution_after != null) "model_execution" else "execution_and_diagnostic_io", + .execution_disk_read_bytes = after.disk_read - before.disk_read, + .execution_pageins = after.pageins - before.pageins, + .resident_gate_passed = resident_gate_passed, + .execution_error = if (execution_error) |err| @as(?[]const u8, @errorName(err)) else null, + }, .{}); + defer a.free(receipt); + std.debug.print("{s}\n", .{receipt}); + try write(io, a, dir, "residency.json", receipt); + if (!resident_gate_passed) std.debug.print("residency gate failed: ResidentExecutionReadFromDisk\n", .{}); + if (execution_error) |err| return err; + if (!resident_gate_passed) return error.ResidentExecutionReadFromDisk; +} + +fn execute(a: std.mem.Allocator, io: std.Io, model: *ds.Model, args: []const [:0]const u8, ids: []const u32, map_blob: []const u8, dir: []const u8, before: Counters, execution_after: *?Counters) !void { + const mode = args[1]; + const is_operators = std.mem.eql(u8, mode, "operators"); + const is_attention = std.mem.eql(u8, mode, "attention-sameinput"); + 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 is_oracle = std.mem.eql(u8, mode, "oracle"); + 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"); + const is_seams = std.mem.eql(u8, mode, "seams"); + if (is_operators) { + try operators.run(a, io, model, dir); + } else if (is_attention) { + try attention_operators.runAttention(a, io, model, args[4], dir, if (args.len > 6) try number(args[6]) else 0); + } else if (is_attention_history) { + try history_operators.runAttentionHistory(a, io, model, args[4], dir, try number(args[6]), try number(args[7])); + } else if (is_hc) { + try hc_operators.runHyperconnections(a, io, model, args[4], dir, if (args.len > 6) try number(args[6]) else 0); + } else if (is_moe) { + try moe_operators.runMoe(a, io, model, args[4], dir, if (args.len > 6) try number(args[6]) else 21); + } else if (is_oracle) { + try operators.run(a, io, model, dir); + _ = try oracle.run(a, io, model, ids, map_blob, dir); + } else if (is_profile) { + try profile.run(a, io, model, ids, dir); + } else if (is_trace) { + try trace.run(a, io, model, ids, map_blob, dir); + } else if (is_teacher or is_seams) { + const chunk = if (args.len > 6) try number(args[6]) else model.workspace.plan.rows; + if (chunk == 0 or chunk > model.workspace.plan.rows) return error.InvalidPrefillChunk; + const outputs = try a.alloc(f32, ids.len * 129280); + defer a.free(outputs); + const elapsed = try teacher(model, ids, chunk, outputs); + // Capture execution before writing the receipt, not filesystem metadata reads from the receipt itself. + if (is_seams) { + const alternative = try a.alloc(f32, outputs.len); + defer a.free(alternative); + for ([_]u32{ 1, 2, 3, 7, 31, 127, 128, 129, 255, 256 }) |split| { + if (split > model.workspace.plan.rows) continue; + _ = try teacher(model, ids, split, alternative); + var error_max: f64 = 0; + var disagreements: usize = 0; + var bit_disagreements: usize = 0; + for (outputs, alternative) |x, y| { + error_max = @max(error_max, @abs(@as(f64, x) - y)); + bit_disagreements += @intFromBool(@as(u32, @bitCast(x)) != @as(u32, @bitCast(y))); + } + for (0..ids.len) |i| { + if (try argmax(outputs[i * 129280 ..][0..129280]) != try argmax(alternative[i * 129280 ..][0..129280])) + disagreements += 1; + } + std.debug.print("{{\"mode\":\"seam\",\"chunk\":{d},\"max_logit_error\":{d},\"bit_disagreements\":{d},\"greedy_disagreements\":{d}}}\n", .{ split, error_max, bit_disagreements, disagreements }); + if (disagreements != 0 or bit_disagreements != 0) { + var name: [48]u8 = undefined; + const file = try std.fmt.bufPrint(&name, "seam-{d}-logits.f32", .{split}); + execution_after.* = try counters(); + try write(io, a, dir, file, std.mem.sliceAsBytes(alternative)); + return error.ChunkSeamMismatch; + } + } + } + std.debug.print("{{\"mode\":\"teacher\",\"tokens\":{d},\"gpu_seconds\":{d}}}\n", .{ ids.len, elapsed }); + execution_after.* = try counters(); + try write(io, a, dir, "logits.f32", std.mem.sliceAsBytes(outputs)); + } else if (std.mem.eql(u8, mode, "generate")) { + const n = if (args.len > 6) try number(args[6]) else 64; + if (ids.len > model.cache.plan.capacity or n == 0 or n > model.cache.plan.capacity - ids.len) + return error.ContextFull; + try model.reset(); + const prefill_start = mtl.clock.seconds(); + var row: []const f32 = &.{}; + var start: usize = 0; + var progress_at = prefill_start; + while (start < ids.len) { + const end = @min(start + model.workspace.plan.rows, ids.len); + row = try model.forward(ids[start..end], if (end == ids.len) .last else .none); + start = end; + const now = mtl.clock.seconds(); + if (now - progress_at >= 60 or start == ids.len) { + const memory = try counters(); + std.debug.print("{{\"mode\":\"prefill_progress\",\"completed_tokens\":{d},\"prompt_tokens\":{d},\"elapsed_seconds\":{d},\"footprint_bytes\":{d},\"disk_read_bytes\":{d},\"pageins\":{d}}}\n", .{ start, ids.len, now - prefill_start, memory.footprint, memory.disk_read - before.disk_read, memory.pageins - before.pageins }); + progress_at = now; + } + } + const prefill_end = mtl.clock.seconds(); + const generated = try a.alloc(u32, n); + defer a.free(generated); + var gpu_seconds: f64 = 0; + const decode_start = mtl.clock.seconds(); + for (generated, 0..) |*token, i| { + token.* = try argmax(row); + if (i + 1 < generated.len) { + row = try model.forward(&.{token.*}, .last); + gpu_seconds += model.last_gpu_seconds; + } + } + const decode_end = mtl.clock.seconds(); + execution_after.* = try counters(); + const json = try std.json.Stringify.valueAlloc(a, generated, .{}); + defer a.free(json); + try write(io, a, dir, "tokens.json", json); + std.debug.print("{{\"mode\":\"generate\",\"prompt_tokens\":{d},\"completion_tokens\":{d},\"prefill_seconds\":{d},\"decode_seconds\":{d},\"decode_gpu_seconds\":{d}}}\n", .{ ids.len, n, prefill_end - prefill_start, decode_end - decode_start, gpu_seconds }); + } else usage(2); +} diff --git a/zig/deepseek_v41_weights_checks.zig b/zig/deepseek_v41_weights_checks.zig new file mode 100644 index 0000000000..b3d4fedcac --- /dev/null +++ b/zig/deepseek_v41_weights_checks.zig @@ -0,0 +1,8 @@ +//! Host checks for V4.1 configuration, Engram table plans and resident weight staging; no kernels compiled. +test { + _ = @import("src/families/deepseek_v41/config.zig"); + _ = @import("src/families/deepseek_v41/view.zig"); + _ = @import("src/families/deepseek_v41/engram_hash.zig"); + _ = @import("src/families/deepseek_v41/load.zig"); + _ = @import("src/families/deepseek_v41/weights.zig"); +} diff --git a/zig/dsv41_core_family.zig b/zig/dsv41_core_family.zig new file mode 100644 index 0000000000..8992aeef76 --- /dev/null +++ b/zig/dsv41_core_family.zig @@ -0,0 +1,13 @@ +//! The core modules the DeepSeek V4.1 family names as `core`: the family is its own module, so it cannot climb out. +pub const checkpoint_metal = @import("src/core/checkpoint_metal.zig"); +pub const contract = @import("src/core/contract.zig"); +pub const exl3_layout = @import("src/core/exl3_layout.zig"); +pub const exl3_load = @import("src/core/exl3_load.zig"); +pub const exl3_prepare = @import("src/core/exl3_prepare.zig"); +pub const exl3_projection = @import("src/core/exl3_projection.zig"); +pub const exl3_projection_gpu_test = @import("src/core/exl3_projection_gpu_test.zig"); +pub const expert_groups = @import("src/core/expert_groups.zig"); +pub const gpu_counters = @import("src/core/gpu_counters.zig"); +pub const paced_read = @import("src/core/paced_read.zig"); +pub const safetensors = @import("src/core/safetensors.zig"); +pub const safetensors_security_test = @import("src/core/safetensors_security_test.zig"); diff --git a/zig/kernels/metal/core/exl3_projection.metal b/zig/kernels/metal/core/exl3_projection.metal index 0eca0d80c5..93ef9cfdf3 100644 --- a/zig/kernels/metal/core/exl3_projection.metal +++ b/zig/kernels/metal/core/exl3_projection.metal @@ -241,6 +241,20 @@ kernel void tf_exl3_state_window_check(const device uint4* inputs [[buffer(0)]], outputs[tid] = tf_exl3_mapped_state(tile, (off << 6) | (1u << 11)); } +// Kept-dense: rotated rows x bf16 weight, f32 in rotate_out sums; sequential k, one thread an output, in fp64 order. +kernel void tf_exl3_dense(const device half* rotated [[buffer(0)]], + const device bfloat* weight [[buffer(1)]], + device float* sums [[buffer(2)]], + constant TfExl3Args& a [[buffer(3)]], + const device uint* input_ids [[buffer(4)]], + uint2 g [[thread_position_in_grid]]) { + const uint row = a.input_ids ? input_ids[g.y] : g.y; + float acc = 0; + for (uint k = 0; k < a.k; ++k) + acc += float(rotated[ulong(g.y) * a.k + k]) * float(weight[ulong(g.x) * a.k + k]); + sums[ulong(g.y) * a.n + g.x] = acc; +} + // Four SIMDgroups share decoded weights while each retains its original eight-row compensated dots. kernel void tf_exl3_matmul_prefill(const device half* x [[buffer(0)]], const device uchar* trellis [[buffer(1)]], diff --git a/zig/kernels/metal/core/expert_groups.metal b/zig/kernels/metal/core/expert_groups.metal new file mode 100644 index 0000000000..be543d5780 --- /dev/null +++ b/zig/kernels/metal/core/expert_groups.metal @@ -0,0 +1,36 @@ +// Group routed slots into padded eight-row expert tiles; scatter preserves original token and output IDs. +#include +using namespace metal; +struct TfExpertGroups { uint rows, experts, repeat, padded; }; +#define TF_GROUP_ARGS \ + const device uint* ids [[buffer(0)]], device atomic_uint* counts [[buffer(1)]], \ + device uint* offsets [[buffer(2)]], device uint* ranks [[buffer(3)]], \ + device uint* expert_ids [[buffer(4)]], device uint* token_ids [[buffer(5)]], \ + device uint* slot_ids [[buffer(6)]], constant TfExpertGroups& d [[buffer(7)]], \ + uint i [[thread_position_in_grid]] + +kernel void tf_expert_group_clear(TF_GROUP_ARGS) { + if (i < d.experts) atomic_store_explicit(counts + i, 0u, memory_order_relaxed); + if (i < d.padded) expert_ids[i] = 0xffffffffu; +} +kernel void tf_expert_group_count(TF_GROUP_ARGS) { + if (i >= d.rows || ids[i] >= d.experts) return; + ranks[i] = atomic_fetch_add_explicit(counts + ids[i], 1u, memory_order_relaxed); +} +kernel void tf_expert_group_prefix(TF_GROUP_ARGS) { + if (i != 0) return; + uint offset = 0; + for (uint expert = 0; expert < d.experts; ++expert) { + offsets[expert] = offset; + const uint count = atomic_load_explicit(counts + expert, memory_order_relaxed); + offset += (count + 7u) & ~7u; + } +} +kernel void tf_expert_group_scatter(TF_GROUP_ARGS) { + if (i >= d.rows || ids[i] >= d.experts) return; + const uint destination = offsets[ids[i]] + ranks[i]; + if (destination >= d.padded) return; + expert_ids[destination] = ids[i]; + token_ids[destination] = i / d.repeat; + slot_ids[destination] = i; +} diff --git a/zig/kernels/metal/deepseek_v41/csa2.metal b/zig/kernels/metal/deepseek_v41/csa2.metal new file mode 100644 index 0000000000..9d16cfcce5 --- /dev/null +++ b/zig/kernels/metal/deepseek_v41/csa2.metal @@ -0,0 +1,102 @@ +// Source is compiled after kv.metal, which provides exact packed-cache decoders. +struct Dsv41Attention { + uint rows, heads, width, q_stride, position, chunk_rows, window, ring_stride; + uint chunk_stride, main_stride, main_block, ratio, compressed_rows, topk, index_stride; + float scale; +}; +// BF16 query/key products are exact in FP32; compensate accumulation and lane reduction. +inline void dsv41_dot_add(thread float &sum, thread float &error, float value) { + float next = sum + value; + float recovered = next - sum; + error += (sum - (next - recovered)) + (value - recovered); + sum = next; +} +inline float dsv41_dot_reduce(float sum, float error) { + for (uint offset = 16; offset != 0; offset >>= 1) { + float other = simd_shuffle_xor(sum, offset); + float other_error = simd_shuffle_xor(error, offset); + dsv41_dot_add(sum, error, other); + dsv41_dot_add(sum, error, other_error); + } + return sum + error; +} +// BF16 probability rounding uses fixed query-based tiles, never a chunk-dependent physical key layout. +inline uint dsv41_window_columns(constant Dsv41Attention &p, uint row) { + return p.position + row == 0 ? 1 : p.window; +} +inline device const uchar *dsv41_attn_key(uint slot, uint row, device const uchar *ring, + device const uchar *chunk, device const uchar *main, device const int *indices, + constant Dsv41Attention &p, thread uint &kind) { + uint query = p.position + row; + uint columns = dsv41_window_columns(p, row); + if (slot < columns) { + uint pos; + if (query == 0) { + pos = 0; + } else if (query + 1 < p.window) { + pos = (query + 1 + slot) % p.window; + } else { + uint begin = query + 1 > p.window ? query + 1 - p.window : 0; + pos = begin + slot; + } + if (pos > query) return nullptr; + kind = 0; + return pos >= p.position ? chunk + (pos - p.position) * p.chunk_stride + : ring + (pos % p.window) * p.ring_stride; + } + if (!p.ratio) return nullptr; + int pos = indices[row * p.index_stride + slot - columns]; + if (pos < 0 || uint(pos) >= min(p.compressed_rows, (query + 1) / p.ratio)) return nullptr; + kind = 1; + return main + uint(pos) * p.main_stride; +} +kernel void dsv41_sparse_attention(device const ushort *q [[buffer(0)]], device const uchar *ring [[buffer(1)]], + device const uchar *chunk [[buffer(2)]], device const uchar *main [[buffer(3)]], + device const int *indices [[buffer(4)]], device const float *sink [[buffer(5)]], + device ushort *out [[buffer(6)]], constant Dsv41Attention &p [[buffer(7)]], + uint2 g [[threadgroup_position_in_grid]], uint lane [[thread_index_in_threadgroup]]) { + uint head = g.x, row = g.y; + if (head >= p.heads || row >= p.rows) return; + threadgroup float scores[64]; + float query[16], acc[16]; + for (uint j = 0; j < p.width / 32; ++j) { + query[j] = dsv41_bf(q[row * p.q_stride + head * p.width + lane + j * 32]); acc[j] = 0; + } + float mx = -1e30f, sum = 0; + uint query_position = p.position + row; + uint compressed_columns = p.ratio ? min(p.topk, min(p.compressed_rows, (query_position + 1) / p.ratio)) : 0; + uint count = dsv41_window_columns(p, row) + compressed_columns; + for (uint base = 0; base < count; base += 64) { + float next_max = mx; + for (uint j = 0; j < 64; ++j) { + uint kind = 0; + device const uchar *key = base + j < count ? dsv41_attn_key(base + j, row, ring, chunk, main, indices, p, kind) : nullptr; + float dot = 0, dot_error = 0; + if (key) for (uint d = 0; d < p.width / 32; ++d) + dsv41_dot_add(dot, dot_error, query[d] * dsv41_unpack(key, lane + d * 32, p.width, kind ? p.main_block : 32, kind)); + dot = key ? dsv41_dot_reduce(dot, dot_error) * p.scale : -INFINITY; + next_max = max(next_max, dot); + if (lane == 0) scores[j] = dot; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float correction = exp(mx - next_max); + float tile_sum = 0; + for (uint d = 0; d < p.width / 32; ++d) acc[d] *= correction; + for (uint j = 0; j < 64; ++j) { + float probability = exp(scores[j] - next_max); + tile_sum += probability; + float rounded = dsv41_bf(dsv41_round(probability)); + uint kind = 0; + device const uchar *key = base + j < count ? dsv41_attn_key(base + j, row, ring, chunk, main, indices, p, kind) : nullptr; + if (key) for (uint d = 0; d < p.width / 32; ++d) + acc[d] += rounded * dsv41_unpack(key, lane + d * 32, p.width, kind ? p.main_block : 32, kind); + } + sum = sum * correction + tile_sum; + mx = next_max; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + // The source adds sink mass to the completed online denominator before the FP32 division. + float denom = sum + exp(sink[head] - mx); + for (uint d = 0; d < p.width / 32; ++d) + out[row * p.q_stride + head * p.width + lane + d * 32] = dsv41_round(acc[d] / denom); +} diff --git a/zig/kernels/metal/deepseek_v41/engram.metal b/zig/kernels/metal/deepseek_v41/engram.metal new file mode 100644 index 0000000000..b3d454cf59 --- /dev/null +++ b/zig/kernels/metal/deepseek_v41/engram.metal @@ -0,0 +1,155 @@ +// Math/layout attribution: DeepSeek AI DeepSeek-V4.1-Flash inference/{engram,model}.py (MIT). +#include +using namespace metal; +inline float rd(const device uchar* p, ulong i, uint dtype) { + if (dtype == 1u) return ((const device float*)p)[i]; + if (dtype == 2u) return float(((const device half*)p)[i]); + return as_type(uint(((const device ushort*)p)[i]) << 16); +} + +struct EngramLayer { + uint layer_id, rows; + ulong multipliers[4]; + uint primes[24], offsets[24]; +}; +struct EngramHistory { int tail[3]; uint available; }; +struct EngramHashArgs { uint tokens, pad, masked; }; + +inline int engram_token(const device uint* ids, const device uint* map, + const device uchar* mask, uint pos, uint masked) { + return masked && mask[pos] == 0 ? -1 : int(map[ids[pos]]); +} + +// One lane per token; history is immutable until all layers have hashed this chunk. +kernel void tf_v41_engram_hash(const device uint* ids [[buffer(0)]], + const device uint* map [[buffer(1)]], + const device uchar* mask [[buffer(2)]], + const device EngramHistory& history [[buffer(3)]], + device uint* out [[buffer(4)]], + constant EngramLayer& layer [[buffer(5)]], + constant EngramHashArgs& args [[buffer(6)]], + uint pos [[thread_position_in_grid]]) { + if (pos >= args.tokens) return; + bool blocked = false; + ulong rolling = 0; + for (uint shift = 0; shift < 4; shift++) { + int token; + if (shift <= pos) token = engram_token(ids, map, mask, pos - shift, args.masked); + else { + const uint back = shift - pos - 1; + token = back < history.available ? history.tail[back] : -1; + } + blocked = blocked || token == -1; + const ulong compressed = blocked ? args.pad : uint(token); + rolling ^= compressed * layer.multipliers[shift]; + if (shift != 0) for (uint head = 0; head < 8; head++) { + const uint col = (shift - 1) * 8 + head; + out[size_t(pos) * 24 + col] = uint(rolling % layer.primes[col]) + layer.offsets[col]; + } + } +} + +// Called exactly once after all Engram layers hash a chunk, never once per layer. +kernel void tf_v41_engram_commit(const device uint* ids [[buffer(0)]], + const device uint* map [[buffer(1)]], + const device uchar* mask [[buffer(2)]], + device EngramHistory& history [[buffer(3)]], + constant EngramHashArgs& args [[buffer(4)]]) { + int tail[3]; + for (uint i = 0; i < 3; i++) { + tail[i] = i < args.tokens ? engram_token(ids, map, mask, args.tokens - i - 1, args.masked) + : i - args.tokens < history.available ? history.tail[i - args.tokens] : -1; + } + for (uint i = 0; i < 3; i++) history.tail[i] = tail[i]; + history.available = min(3u, history.available + args.tokens); +} + +kernel void tf_v41_engram_reset(device EngramHistory& history [[buffer(0)]]) { + for (uint i = 0; i < 3; i++) history.tail[i] = -1; + history.available = 0; +} + +inline float engram_e4m3(uchar bits) { + const uint magnitude = bits & 127, exponent = magnitude >> 3, mantissa = magnitude & 7; + float value = magnitude == 127 ? NAN : exponent == 0 ? float(mantissa) * 0x1p-9f + : ldexp(1.0f + float(mantissa) * 0.125f, int(exponent) - 7); + return bits & 128 ? -value : value; +} +inline float engram_ue8m0(uchar bits) { + return bits == 255 ? NAN : as_type(bits == 0 ? 0x00400000u : uint(bits) << 23); +} + +struct EngramGatherArgs { + ulong weight_stride, scale_stride, output_stride; + uint tokens, columns, head_dim, scale_block; +}; + +// Read only the selected E4M3FN rows; one UE8M0 scale per 32 values of each row. +kernel void tf_v41_engram_gather(const device uint* ids [[buffer(0)]], + const device uchar* table [[buffer(1)]], + const device uchar* scales [[buffer(2)]], + device bfloat* out [[buffer(3)]], + constant EngramGatherArgs& args [[buffer(4)]], + uint2 pos [[thread_position_in_grid]]) { + if (pos.y >= args.tokens || pos.x >= args.columns * args.head_dim) return; + const uint col = pos.x / args.head_dim, d = pos.x % args.head_dim; + const size_t row = ids[size_t(pos.y) * args.columns + col]; + const float value = engram_e4m3(table[row * args.weight_stride + d]); + const float scale = engram_ue8m0(scales[row * args.scale_stride + d / args.scale_block]); + out[size_t(pos.y) * args.output_stride + pos.x] = bfloat(value * scale); +} + +struct EngramGateArgs { + ulong hidden_stride, kv_stride, output_stride; + uint dim, copies, masked, weight_dtype; + float eps; +}; + +// Each copy gets its own RMS dot gate; all copies share the value following the flattened keys. +kernel void tf_v41_engram_gate(const device bfloat* hidden [[buffer(0)]], + const device bfloat* kv [[buffer(1)]], + const device uchar* q_weight [[buffer(2)]], + const device uchar* k_weight [[buffer(3)]], + const device uchar* mask [[buffer(4)]], + device bfloat* out [[buffer(5)]], + constant EngramGateArgs& args [[buffer(6)]], + uint2 group [[threadgroup_position_in_grid]], + uint t [[thread_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], + uint sg [[simdgroup_index_in_threadgroup]]) { + threadgroup float3 sums[8]; + threadgroup float gate; + const size_t hbase = size_t(group.x) * args.hidden_stride + size_t(group.y) * args.dim; + const size_t kbase = size_t(group.x) * args.kv_stride + size_t(group.y) * args.dim; + const size_t wbase = size_t(group.y) * args.dim; + float3 acc = 0; + for (uint d = t; d < args.dim; d += 256) { + const float h = float(hidden[hbase + d]), k = float(kv[kbase + d]); + const float weight = rd(q_weight, wbase + d, args.weight_dtype) * rd(k_weight, wbase + d, args.weight_dtype); + acc += float3(h * h, k * k, (h * weight) * k); + } + acc = float3(simd_sum(acc.x), simd_sum(acc.y), simd_sum(acc.z)); + if (lane == 0) sums[sg] = acc; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + acc = lane < 8 ? sums[lane] : float3(0); + acc = float3(simd_sum(acc.x), simd_sum(acc.y), simd_sum(acc.z)); + if (lane == 0) { + float rstd; + { +#pragma clang fp contract(off) + const float factor = 1.0f / float(args.dim); + rstd = metal::precise::rsqrt(acc.x * factor + args.eps) + * metal::precise::rsqrt(acc.y * factor + args.eps); + } + const float dot = (acc.z * rstd) * metal::precise::rsqrt(float(args.dim)); + const float signed_root = copysign(metal::precise::sqrt(max(abs(dot), 1e-6f)), dot); + gate = args.masked && mask[group.x] == 0 ? 0.0f : 1.0f / (1.0f + metal::precise::exp(-signed_root)); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const size_t vbase = size_t(group.x) * args.kv_stride + size_t(args.copies) * args.dim; + const size_t obase = size_t(group.x) * args.output_stride + size_t(group.y) * args.dim; + for (uint d = t; d < args.dim; d += 256) + out[obase + d] = bfloat(float(hidden[hbase + d]) + gate * float(kv[vbase + d])); +} diff --git a/zig/kernels/metal/deepseek_v41/index.metal b/zig/kernels/metal/deepseek_v41/index.metal new file mode 100644 index 0000000000..8d62337624 --- /dev/null +++ b/zig/kernels/metal/deepseek_v41/index.metal @@ -0,0 +1,103 @@ +// Source is compiled after kv.metal; no query-by-key-by-head score tensor is materialized. +struct Dsv41Index { + uint rows, heads, width, q_stride, weight_stride, position, ratio, keys, key_stride; + uint tile_start, tile_width, topk, candidate_topk, block_size, use_candidates, make_candidates; + float scale; +}; +struct Dsv41Pick { float score; int position; }; +inline bool dsv41_less(Dsv41Pick a, Dsv41Pick b) { + return a.score < b.score || (a.score == b.score && a.position > b.position); +} +inline void dsv41_heap(device Dsv41Pick *h, uint size, Dsv41Pick v) { + if (!dsv41_less(h[0], v)) return; + uint at = 0; + while (2 * at + 1 < size) { + uint child = 2 * at + 1; + if (child + 1 < size && dsv41_less(h[child + 1], h[child])) ++child; + if (!dsv41_less(h[child], v)) break; + h[at] = h[child]; at = child; + } + h[at] = v; +} +inline int dsv41_index_position(uint local, uint row, device const Dsv41Pick *candidates, + constant Dsv41Index &p) { + uint slot = p.tile_start + local; + if (!p.use_candidates) return int(slot); + uint block = slot / p.block_size; + if (block >= p.candidate_topk) return -1; + int id = candidates[row * p.candidate_topk + block].position; + return id < 0 ? -1 : id * int(p.block_size) + int(slot % p.block_size); +} +kernel void dsv41_index_score(device const ushort *q [[buffer(0)]], device const ushort *weight [[buffer(1)]], + device const uchar *key [[buffer(2)]], device const Dsv41Pick *candidates [[buffer(3)]], + device float *scores [[buffer(4)]], constant Dsv41Index &p [[buffer(5)]], + uint2 g [[threadgroup_position_in_grid]], uint lane [[thread_index_in_threadgroup]]) { + uint local = g.x, row = g.y; + if (local >= p.tile_width || row >= p.rows) return; + int pos = dsv41_index_position(local, row, candidates, p); + uint visible = min(p.keys, (p.position + row + 1) / p.ratio); + if (pos < 0 || uint(pos) >= visible) { + if (lane == 0) scores[row * p.tile_width + local] = -INFINITY; + return; + } + device const uchar *k = key + uint(pos) * p.key_stride; + float key_values[4]; + for (uint d = 0; d < p.width / 32; ++d) key_values[d] = dsv41_unpack(k, lane + 32 * d, p.width, 32, 2); + float total = 0; + for (uint head = 0; head < p.heads; ++head) { + float dot = 0; + for (uint d = 0; d < p.width / 32; ++d) + dot += dsv41_bf(q[row * p.q_stride + head * p.width + lane + 32 * d]) * key_values[d]; + dot = dsv41_bf(dsv41_round(simd_sum(dot))); + const float scaled_weight = dsv41_bf(dsv41_round(dsv41_bf(weight[row * p.weight_stride + head]) * p.scale)); + total += dsv41_bf(dsv41_round(max(dot, 0.0f) * scaled_weight)); + } + if (lane == 0) scores[row * p.tile_width + local] = dsv41_bf(dsv41_round(total)); +} +kernel void dsv41_index_clear(device Dsv41Pick *top [[buffer(0)]], device Dsv41Pick *candidate [[buffer(1)]], + constant Dsv41Index &p [[buffer(2)]], uint2 i [[thread_position_in_grid]]) { + if (i.y >= p.rows) return; + if (i.x < p.topk) top[i.y * p.topk + i.x] = {-INFINITY, -1}; + if (p.make_candidates && i.x < p.candidate_topk) candidate[i.y * p.candidate_topk + i.x] = {-INFINITY, -1}; +} +kernel void dsv41_index_merge(device const float *scores [[buffer(0)]], device Dsv41Pick *top [[buffer(1)]], + device Dsv41Pick *candidate [[buffer(2)]], constant Dsv41Index &p [[buffer(3)]], + uint row [[thread_position_in_grid]]) { + if (row >= p.rows) return; + device Dsv41Pick *heap = top + row * p.topk; + for (uint i = 0; i < p.tile_width; ++i) { + float score = scores[row * p.tile_width + i]; + int pos = dsv41_index_position(i, row, candidate, p); + if (isfinite(score)) dsv41_heap(heap, p.topk, {score, pos}); + } + if (!p.make_candidates) return; + uint visible = min(p.keys, (p.position + row + 1) / p.ratio); + device Dsv41Pick *blocks = candidate + row * p.candidate_topk; + for (uint base = 0; base < p.tile_width; base += p.block_size) { + float score = -INFINITY; + for (uint j = 0; j < p.block_size; ++j) score = max(score, scores[row * p.tile_width + base + j]); + uint block = (p.tile_start + base) / p.block_size; + if (visible && block == (visible - 1) / p.block_size) score = INFINITY; + if (score > -INFINITY) dsv41_heap(blocks, p.candidate_topk, {score, int(block)}); + } +} +inline void dsv41_sort_positions(device Dsv41Pick *p, uint size) { + // Selection is score-ordered; attention sees the reference's ascending position order. + for (uint i = 1; i < size; ++i) { + Dsv41Pick value = p[i]; uint j = i; + while (j > 0 && (p[j - 1].position < 0 || + (value.position >= 0 && p[j - 1].position > value.position))) { + if (value.position < 0) break; + p[j] = p[j - 1]; --j; + } + p[j] = value; + } +} +kernel void dsv41_index_finish(device Dsv41Pick *top [[buffer(0)]], device Dsv41Pick *candidate [[buffer(1)]], + device int *positions [[buffer(2)]], constant Dsv41Index &p [[buffer(3)]], uint row [[thread_position_in_grid]]) { + if (row >= p.rows) return; + device Dsv41Pick *heap = top + row * p.topk; + dsv41_sort_positions(heap, p.topk); + for (uint i = 0; i < p.topk; ++i) positions[row * p.topk + i] = heap[i].position; + if (p.make_candidates) dsv41_sort_positions(candidate + row * p.candidate_topk, p.candidate_topk); +} diff --git a/zig/kernels/metal/deepseek_v41/kv.metal b/zig/kernels/metal/deepseek_v41/kv.metal new file mode 100644 index 0000000000..6ac2017efb --- /dev/null +++ b/zig/kernels/metal/deepseek_v41/kv.metal @@ -0,0 +1,141 @@ +// DeepSeek V4.1 cache quantization and pooling, following the official inference kernels. +#include +using namespace metal; +inline float dsv41_bf(ushort x) { return as_type(uint(x) << 16); } +inline ushort dsv41_round(float x) { + uint u = as_type(x); + return ushort((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} +inline float dsv41_e4(uchar c) { + uint e = (c >> 3) & 15, m = c & 7; + float v = e == 0 ? float(m) * 0x1p-9f : ldexp(1.0f + float(m) / 8.0f, int(e) - 7); + return (c & 128) ? -v : v; +} +inline uchar dsv41_e4code(float x) { + float a = min(abs(x), 448.0f); + if (a == 0) return signbit(x) ? 128 : 0; + int e = max(-6, int((as_type(a) >> 23) & 255) - 127); + int q = int(rint(ldexp(a, 3 - e))); + uint c = e == -6 && q < 8 ? uint(q) : uint((e + 7) * 8 + q - 8); + return uchar(min(c, 126u) | (signbit(x) ? 128u : 0u)); +} +inline float dsv41_e2(uchar c) { + const float v[8] = {0, .5f, 1, 1.5f, 2, 3, 4, 6}; + return (c & 8) ? -v[c & 7] : v[c & 7]; +} +inline uchar dsv41_e2code(float x) { + float a = min(abs(x), 6.0f); + uint pick = a <= .25f ? 0 : a < .75f ? 1 : a <= 1.25f ? 2 : a < 1.75f ? 3 + : a <= 2.5f ? 4 : a < 3.5f ? 5 : a <= 5.0f ? 6 : 7; + return uchar(pick | (signbit(x) ? 8 : 0)); +} +inline float dsv41_unpack(device const uchar *p, uint d, uint width, uint block, uint kind) { + if (kind == 0) { + float s = as_type(uint(p[width + d / 32]) << 23); + return dsv41_bf(dsv41_round(dsv41_e4(p[d]) * s)); + } + uchar c = p[d / 2]; + float s = kind == 1 ? dsv41_e4(p[width / 2 + d / block]) + : as_type(uint(p[width / 2 + d / block]) << 23); + return dsv41_bf(dsv41_round(dsv41_e2((d & 1) ? c >> 4 : c & 15) * s)); +} +struct Dsv41Pack { uint rows, width, stride, packed_stride, block, kind, first, ring, inplace; }; +kernel void dsv41_pack(device ushort *x [[buffer(0)]], device uchar *out [[buffer(1)]], + constant Dsv41Pack &p [[buffer(2)]], uint2 g [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_threadgroup]]) { + uint row = g.y, group = g.x, d = group * p.block + lane; + if (row >= p.rows) return; + float v = lane < p.block ? dsv41_bf(x[row * p.stride + d]) : 0; + float a = simd_max(abs(v)); + float scale; uchar sc; + if (p.kind == 1) { + sc = dsv41_e4code(max(a, 6.0f * 0x1p-9f) / 6.0f); scale = dsv41_e4(sc); + } else { + float minimum = p.kind == 0 ? 1e-4f : 6.0f * 0x1p-126f; + float z = max(a, minimum) / (p.kind == 0 ? 448.0f : 6.0f); + uint bits = as_type(z); + int exp = int((bits >> 23) & 255) - 127 + ((bits & 0x7fffff) != 0); + sc = uchar(exp + 127); scale = ldexp(1.0f, exp); + } + uint dst = p.ring ? (p.first + row) % p.ring : p.first + row; + device uchar *packed = out + dst * p.packed_stride; + if (lane == 0) packed[(p.kind == 0 ? p.width : p.width / 2) + group] = sc; + if (lane >= p.block) return; + if (p.kind == 0) { + uchar c = dsv41_e4code(v / scale); packed[d] = c; + if (p.inplace) x[row * p.stride + d] = dsv41_round(dsv41_e4(c) * scale); + } else { + uchar c = dsv41_e2code(v / scale); + uchar other = uchar(simd_shuffle_xor(uint(c), 1)); + if (!(lane & 1)) packed[d / 2] = c | (other << 4); + if (p.inplace) x[row * p.stride + d] = dsv41_round(dsv41_e2(c) * scale); + } +} +struct Dsv41Pool { uint rows, width, ratio, position, stride, input_f32; float eps; }; +kernel void dsv41_pool(device const uchar *values [[buffer(0)]], device const float *gate [[buffer(1)]], + device float *partial [[buffer(2)]], device const ushort *weight [[buffer(3)]], + device ushort *latent [[buffer(4)]], constant Dsv41Pool &p [[buffer(5)]], + uint group [[threadgroup_position_in_grid]], uint lane [[thread_index_in_threadgroup]]) { + // A group reads the prior chunk's tail before the independent tail-update dispatch. + uint start = (p.position / p.ratio + group) * p.ratio; + threadgroup float pooled[512]; + float square = 0; + for (uint d = lane; d < p.width; d += 128) { + float mx = -INFINITY, sum = 0, total = 0; + for (uint j = 0; j < p.ratio; ++j) { + uint pos = start + j; + float score = p.ratio == 1 ? 0 : (pos < p.position ? partial[p.ratio * p.width + (pos % p.ratio) * p.width + d] + : gate[(pos - p.position) * p.stride + d]); + mx = max(mx, score); + } + float slots[2], exponent[2]; + for (uint j = 0; j < p.ratio; ++j) { + uint pos = start + j; + slots[j] = pos < p.position ? partial[(pos % p.ratio) * p.width + d] + : (p.input_f32 ? reinterpret_cast(values)[(pos - p.position) * p.stride + d] + : dsv41_bf(reinterpret_cast(values)[(pos - p.position) * p.stride + d])); + float score = p.ratio == 1 ? 0 : (pos < p.position ? partial[p.ratio * p.width + (pos % p.ratio) * p.width + d] + : gate[(pos - p.position) * p.stride + d]); + exponent[j] = exp(score - mx); sum += exponent[j]; + } + for (uint j = 0; j < p.ratio; ++j) total += slots[j] * (exponent[j] / sum); + float v = dsv41_bf(dsv41_round(total)); pooled[d] = v; square += v * v; + } + threadgroup float sums[4]; + float ss = simd_sum(square); + if (!(lane % 32)) sums[lane / 32] = ss; + threadgroup_barrier(mem_flags::mem_threadgroup); + float inv = rsqrt((sums[0] + sums[1] + sums[2] + sums[3]) / p.width + p.eps); + for (uint d = lane; d < p.width; d += 128) + latent[group * p.width + d] = dsv41_round(pooled[d] * inv * dsv41_bf(weight[d])); +} +kernel void dsv41_pool_tail(device const float *values [[buffer(0)]], device const float *gate [[buffer(1)]], + device float *partial [[buffer(2)]], constant Dsv41Pool &p [[buffer(3)]], + uint d [[thread_position_in_grid]]) { + if (d >= p.width || p.ratio == 1) return; + uint end = p.position + p.rows, tail = end % p.ratio; + for (uint j = 0; j < tail; ++j) { + uint pos = end - tail + j; + if (pos < p.position) continue; + partial[j * p.width + d] = values[(pos - p.position) * p.stride + d]; + partial[p.ratio * p.width + j * p.width + d] = gate[(pos - p.position) * p.stride + d]; + } +} +kernel void dsv41_copy(device const uchar *src [[buffer(0)]], device uchar *dst [[buffer(1)]], + constant uint &bytes [[buffer(2)]], uint i [[thread_position_in_grid]]) { + if (i < bytes) dst[i] = src[i]; +} +kernel void dsv41_publish_window(device const uchar *chunk [[buffer(0)]], device uchar *ring [[buffer(1)]], + constant uint4 &p [[buffer(2)]], uint2 i [[thread_position_in_grid]]) { + uint count = min(p.x, 128u), first = p.x - count; + if (i.x >= p.z || i.y >= count) return; + uint row = first + i.y, slot = (p.y + row) % 128; + ring[slot * p.w + i.x] = chunk[row * p.z + i.x]; +} +kernel void dsv41_unpack_rows(device const uchar *packed [[buffer(0)]], device ushort *out [[buffer(1)]], + constant Dsv41Pack &p [[buffer(2)]], uint2 i [[thread_position_in_grid]]) { + if (i.x >= p.width || i.y >= p.rows) return; + uint row = p.ring ? (p.first + i.y) % p.ring : p.first + i.y; + out[i.y * p.stride + i.x] = dsv41_round(dsv41_unpack(packed + row * p.packed_stride, + i.x, p.width, p.block, p.kind)); +} diff --git a/zig/kernels/metal/deepseek_v41/mhc.metal b/zig/kernels/metal/deepseek_v41/mhc.metal new file mode 100644 index 0000000000..8b6203a4cf --- /dev/null +++ b/zig/kernels/metal/deepseek_v41/mhc.metal @@ -0,0 +1,120 @@ +// Reference: DeepSeek-V4.1-Flash inference/model.py and kernel.py, revision 2cba9e42aa026125f3ed06c6d98c1db82f7ca027. +#include +using namespace metal; +inline float bf(ushort x) { return as_type(uint(x)<<16); } +inline ushort rb(float x) { + uint u=as_type(x); + if ((u&0x7fffffff)>0x7f800000) return ushort((u>>16)|0x40); + return ushort((u+0x7fff+((u>>16)&1))>>16); +} +// Torch CUDA materializes FP32 squares before MeanOps accumulation. +#pragma clang fp contract(off) +struct Rms { uint rows,width,stream_stride; float eps; }; +kernel void dsv41_hc_rms(const device ushort* x [[buffer(0)]], device float* inv [[buffer(1)]], + constant Rms& d [[buffer(2)]], uint row [[threadgroup_position_in_grid]], + uint t [[thread_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], + uint sg [[simdgroup_index_in_threadgroup]]) { + // Fixed single-row CUDA tree: four vector accumulators, stride 512, independent of prefill chunks. + threadgroup float sums[512]; + float4 accum=0.0f; + for (uint i=4*t;i<4*d.width;i+=4*512) { + float4 a=float4(bf(x[ulong(row)*d.stream_stride+i]), + bf(x[ulong(row)*d.stream_stride+i+1]), + bf(x[ulong(row)*d.stream_stride+i+2]), + bf(x[ulong(row)*d.stream_stride+i+3])); + float4 square=a*a; + accum+=square; + } + float s=((accum.x+accum.y)+accum.z)+accum.w; + sums[t]=s; + for (uint offset=256;offset>=32;offset>>=1) { + threadgroup_barrier(mem_flags::mem_threadgroup); + if (t0;offset>>=1) s+=simd_shuffle_down(s,offset); + // MeanOps projects with a rounded FP32 factor, not a division of the sum. + float factor=1.0f/float(4*d.width); + if (lane==0) inv[row]=rsqrt(s*factor+d.eps); + } +} +#pragma clang fp contract(on) +struct Split { uint rows,mix_stride,coefficient_stride,comb_stride; float eps; }; +kernel void dsv41_hc_split(const device float* mixes [[buffer(0)]], const device float* inv [[buffer(1)]], + const device float* scale [[buffer(2)]], const device float* base [[buffer(3)]], + device float* pre [[buffer(4)]], device float* post [[buffer(5)]], + device float* comb [[buffer(6)]], constant Split& d [[buffer(7)]], + uint row [[thread_position_in_grid]]) { + if (row>=d.rows) return; + float m[16]; + for (uint j=0;j<4;j++) { + float a=fma(mixes[ulong(row)*d.mix_stride+j]*inv[row],scale[0],base[j]); + float b=fma(mixes[ulong(row)*d.mix_stride+4+j]*inv[row],scale[1],base[4+j]); + pre[ulong(row)*d.coefficient_stride+j]=1.0f/(1.0f+exp(-a))+d.eps; + post[ulong(row)*d.coefficient_stride+j]=2.0f/(1.0f+exp(-b)); + for (uint k=0;k<4;k++) { + uint q=4*j+k; + m[q]=fma(mixes[ulong(row)*d.mix_stride+8+q]*inv[row],scale[2],base[8+q]); + } + } + // Initial row softmax + eps, then a column normalization. This counts as iteration one. + for (uint j=0;j<4;j++) { + float mx=max(max(m[4*j],m[4*j+1]),max(m[4*j+2],m[4*j+3])); + for (uint k=0;k<4;k++) m[4*j+k]=exp(m[4*j+k]-mx); + // Pinned TileLang HC row/column sums use FP32 butterfly pairs. + float s=(m[4*j]+m[4*j+2])+(m[4*j+1]+m[4*j+3]); + for (uint k=0;k<4;k++) m[4*j+k]=m[4*j+k]/s+d.eps; + } + for (uint it=0;it<20;it++) { + if (it>0) { + for (uint j=0;j<4;j++) { + float s=((m[4*j]+m[4*j+2])+(m[4*j+1]+m[4*j+3]))+d.eps; + for (uint k=0;k<4;k++) m[4*j+k]/=s; + } + } + for (uint k=0;k<4;k++) { + float s=((m[k]+m[8+k])+(m[4+k]+m[12+k]))+d.eps; + for (uint j=0;j<4;j++) m[4*j+k]/=s; + } + } + for (uint q=0;q<16;q++) comb[ulong(row)*d.comb_stride+q]=m[q]; +} +// PyTorch materializes these FP32 products before reduction and branch addition. +#pragma clang fp contract(off) +struct Collapse { uint rows,width,stream_stride,y_stride,pre_stride; }; +kernel void dsv41_hc_collapse(const device ushort* x [[buffer(0)]], const device float* previous_pre [[buffer(1)]], + device ushort* y [[buffer(2)]], constant Collapse& d [[buffer(3)]], + uint2 i [[thread_position_in_grid]]) { + if (i.x>=d.width || i.y>=d.rows) return; + float s=0; + for (uint c=0;c<4;c++) s+=previous_pre[ulong(i.y)*d.pre_stride+c]*bf(x[ulong(i.y)*d.stream_stride+ulong(c)*d.width+i.x]); + y[ulong(i.y)*d.y_stride+i.x]=rb(s); +} +struct Post { uint rows,width,branch_stride,residual_stride,y_stride,coefficient_stride,comb_stride; }; +kernel void dsv41_hc_post(const device ushort* branch [[buffer(0)]], const device ushort* residual [[buffer(1)]], + const device float* post [[buffer(2)]], const device float* comb [[buffer(3)]], + device ushort* y [[buffer(4)]], constant Post& d [[buffer(5)]], + uint2 i [[thread_position_in_grid]]) { + if (i.x>=d.width || i.y>=d.rows) return; + float r[4]; + for (uint j=0;j<4;j++) r[j]=bf(residual[ulong(i.y)*d.residual_stride+ulong(j)*d.width+i.x]); + float b=bf(branch[ulong(i.y)*d.branch_stride+i.x]); + for (uint k=0;k<4;k++) { + float s=0; + for (uint j=0;j<4;j++) s+=comb[ulong(i.y)*d.comb_stride+4*j+k]*r[j]; + y[ulong(i.y)*d.y_stride+ulong(k)*d.width+i.x]=rb(post[ulong(i.y)*d.coefficient_stride+k]*b+s); + } +} +#pragma clang fp contract(on) +struct Expand { uint rows,width,x_stride,stream_stride,pre_stride; }; +kernel void dsv41_hc_expand(const device ushort* x [[buffer(0)]], device ushort* stream [[buffer(1)]], + device float* pre [[buffer(2)]], constant Expand& d [[buffer(3)]], + uint2 i [[thread_position_in_grid]]) { + if (i.x>=d.width || i.y>=d.rows) return; + for (uint c=0;c<4;c++) stream[ulong(i.y)*d.stream_stride+ulong(c)*d.width+i.x]=x[ulong(i.y)*d.x_stride+i.x]; + if (i.x==0) for (uint c=0;c<4;c++) pre[ulong(i.y)*d.pre_stride+c]=c==0?1.0f:0.0f; +} diff --git a/zig/kernels/metal/deepseek_v41/ops.metal b/zig/kernels/metal/deepseek_v41/ops.metal new file mode 100644 index 0000000000..d831901dd5 --- /dev/null +++ b/zig/kernels/metal/deepseek_v41/ops.metal @@ -0,0 +1,124 @@ +// Math: deepseek-ai/DeepSeek-V4.1-Flash inference/model.py, revision 2cba9e42aa026125f3ed06c6d98c1db82f7ca027. +#include +using namespace metal; +inline float bf(ushort x) { return as_type(uint(x) << 16); } +inline ushort rb(float x) { + uint u = as_type(x); + if ((u & 0x7fffffff) > 0x7f800000) return ushort((u >> 16) | 0x40); + return ushort((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} +inline float rd(const device uchar* p, ulong i, uint dtype) { + if (dtype==1) return ((const device float*)p)[i]; + if (dtype==2) return float(((const device half*)p)[i]); + return bf(((const device ushort*)p)[i]); +} +struct Rms { uint rows, width, x_stride, y_stride; float eps; uint weight_f32; }; +#pragma clang fp contract(off) +kernel void dsv41_rms(const device ushort* x [[buffer(0)]], const device uchar* w [[buffer(1)]], + device ushort* y [[buffer(2)]], constant Rms& d [[buffer(3)]], + uint row [[threadgroup_position_in_grid]], uint t [[thread_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], uint sg [[simdgroup_index_in_threadgroup]]) { + threadgroup float sums[8]; + float s = 0; + for (uint i = t; i < d.width; i += 256) { float a = bf(x[ulong(row)*d.x_stride+i]); s += a*a; } + s = simd_sum(s); + if (lane == 0) sums[sg] = s; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + float total = simd_sum(lane < 8 ? sums[lane] : 0.0f); + const float factor = 1.0f / float(d.width); + if (lane == 0) sums[0] = rsqrt(total * factor + d.eps); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint i = t; i < d.width; i += 256) + y[ulong(row)*d.y_stride+i] = rb((bf(x[ulong(row)*d.x_stride+i])*sums[0])*rd(w,i,d.weight_f32)); +} +#pragma clang fp contract(on) +// Each SIMD group emits four columns; compensated dots do not depend on token count. + +// Precision: 0=BF16, 1=FP32, 2=FP16; FP16 weights never round through BF16. +inline void dense_add(thread float& high, thread float& low, float value) { + float next = high + value; + float borrowed = next - high; + low += (high - (next - borrowed)) + (value - borrowed); + high = next; +} +struct Dense { uint rows, k, n, groups, x_stride, w_stride, y_stride, x_dtype, w_dtype, y_dtype; }; +kernel void dsv41_dense(const device uchar* x [[buffer(0)]], const device uchar* w [[buffer(1)]], + device uchar* y [[buffer(2)]], constant Dense& d [[buffer(3)]], + uint3 g [[threadgroup_position_in_grid]], uint lane [[thread_index_in_simdgroup]], + uint sg [[simdgroup_index_in_threadgroup]]) { + uint col = g.x*16 + sg*4; + float a[4] = {0,0,0,0}, residual[4] = {0,0,0,0}; + for (uint k = lane; k < d.k; k += 32) { + float v = rd(x,ulong(g.y)*d.x_stride+ulong(g.z)*d.k+k,d.x_dtype); + for (uint j=0;j<4;j++) if (col+j>= 1) { + float other = simd_shuffle_xor(a[j], offset); + float other_residual = simd_shuffle_xor(residual[j], offset); + dense_add(a[j], residual[j], other); + residual[j] += other_residual; + } + float v = a[j] + residual[j]; + if (lane==0 && col+j=d.width || i.y>=d.rows) return; + ushort v=table[ulong(ids[i.y])*d.table_stride+i.x]; + for (uint c=0;c=d.width || i.y>=d.rows) return; + ulong k=ulong(i.y)*d.input_stride+i.x; + float a=min(rd(gate,k,d.input_f32),10.0f), b=clamp(rd(up,k,d.input_f32),-10.0f,10.0f); + float v=(a/(1.0f+exp(-a)))*b; + if (d.weighted) v=weights[i.y]*v; + y[ulong(i.y)*d.output_stride+i.x]=rb(v); +} +struct Combine { uint rows,width,routed_row_stride,slot_stride,shared_stride,y_stride,input_f32,id_stride; }; +kernel void dsv41_combine(const device uchar* routed [[buffer(0)]], const device uchar* shared [[buffer(1)]], + const device uint* ids [[buffer(2)]], device ushort* y [[buffer(3)]], + constant Combine& d [[buffer(4)]], + uint2 i [[thread_position_in_grid]], uint t [[thread_index_in_threadgroup]]) { + threadgroup uint slots[6]; + if (t==0) { + for (uint a=0;a<6;a++) slots[a]=a; + for (uint a=1;a<6;a++) { + uint slot=slots[a], j=a; + while (j>0 && ids[ulong(i.y)*d.id_stride+slots[j-1]]>ids[ulong(i.y)*d.id_stride+slot]) { + slots[j]=slots[j-1]; j--; + } + slots[j]=slot; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (i.x>=d.width || i.y>=d.rows) return; + for (uint slot=0;slot<6;slot++) if (ids[ulong(i.y)*d.id_stride+slot]>=384) { + y[ulong(i.y)*d.y_stride+i.x]=rb(NAN); + return; + } + float s=0; + for (uint rank=0;rank<6;rank++) + s+=rd(routed,ulong(i.y)*d.routed_row_stride+ulong(slots[rank])*d.slot_stride+i.x,d.input_f32); + s+=rd(shared,ulong(i.y)*d.shared_stride+i.x,d.input_f32); + y[ulong(i.y)*d.y_stride+i.x]=rb(s); +} diff --git a/zig/kernels/metal/deepseek_v41/rope.metal b/zig/kernels/metal/deepseek_v41/rope.metal new file mode 100644 index 0000000000..09c4d1cb2e --- /dev/null +++ b/zig/kernels/metal/deepseek_v41/rope.metal @@ -0,0 +1,17 @@ +// Adjacent-pair RoPE and YaRN match DeepSeek V4.1's in-place BF16 rotation. +#include +using namespace metal; +struct Dsv41Rope { uint rows, heads, width, rotary, stride, position, step, inverse, spectrum; }; +kernel void dsv41_rope(device ushort *x [[buffer(0)]], constant Dsv41Rope &p [[buffer(1)]], + device const float2 *factors [[buffer(2)]], + uint3 i [[thread_position_in_grid]]) { + if (i.x >= p.rotary / 2 || i.y >= p.heads || i.z >= p.rows) return; + // The encoder binds the selected spectrum; rows are absolute source positions. + float2 rotation = factors[(ulong(p.position) + ulong(i.z) * p.step) * (p.rotary / 2) + i.x]; + float c = rotation.x, s = p.inverse ? -rotation.y : rotation.y; + uint at = i.z * p.stride + i.y * p.width + p.width - p.rotary + 2 * i.x; + float a = as_type(uint(x[at]) << 16), b = as_type(uint(x[at + 1]) << 16); + uint u = as_type(a * c - b * s), v = as_type(a * s + b * c); + x[at] = ushort((u + 0x7fff + ((u >> 16) & 1)) >> 16); + x[at + 1] = ushort((v + 0x7fff + ((v >> 16) & 1)) >> 16); +} diff --git a/zig/kernels/metal/deepseek_v41/route.metal b/zig/kernels/metal/deepseek_v41/route.metal new file mode 100644 index 0000000000..51f5454dd3 --- /dev/null +++ b/zig/kernels/metal/deepseek_v41/route.metal @@ -0,0 +1,53 @@ +// Reference: DeepSeek-V4.1-Flash inference/model.py Gate, revision 2cba9e42aa026125f3ed06c6d98c1db82f7ca027. +#include +using namespace metal; +inline float lp(float x) { float u=1.0f+x; return u==1.0f?x:x*(log(u)/(u-1.0f)); } +inline float score(float x) { return sqrt(x>20.0f?x:lp(exp(x))); } +struct Shape { uint rows,logit_stride,output_stride; float temperature; uint bias_dtype; }; +kernel void dsv41_route(const device float* logits [[buffer(0)]], const device uchar* bias [[buffer(1)]], + device uint* ids [[buffer(2)]], device float* weights [[buffer(3)]], + constant Shape& d [[buffer(4)]], uint row [[threadgroup_position_in_grid]], + uint t [[thread_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], + uint sg [[simdgroup_index_in_threadgroup]]) { + threadgroup float scores[384], select[384], best[4], picked[6]; + threadgroup uint best_id[4], winner; + for (uint e=t;e<384;e+=128) { + float b; + if (d.bias_dtype==1) b=((const device float*)bias)[e]; + else if (d.bias_dtype==2) b=float(((const device half*)bias)[e]); + else b=as_type(uint(((const device ushort*)bias)[e])<<16); + scores[e]=score(logits[ulong(row)*d.logit_stride+e]/d.temperature); + select[e]=scores[e]+b; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // Every expert participates in every rank selection; ties use the lower expert id deterministically. + for (uint slot=0;slot<6;slot++) { + float v=-INFINITY; + uint id=0xffffffffu; + for (uint e=t;e<384;e+=128) if (select[e]>v || (select[e]==v && e std.math.maxInt(u32)) return error.InvalidExl3Experts; diff --git a/zig/src/core/exl3_projection_gpu_test.zig b/zig/src/core/exl3_projection_gpu_test.zig index 6f6445c433..1beb5ee218 100644 --- a/zig/src/core/exl3_projection_gpu_test.zig +++ b/zig/src/core/exl3_projection_gpu_test.zig @@ -285,6 +285,34 @@ test "packed signs and GPU expert filtering preserve untouched output slots" { try std.testing.expectEqualSlices(u16, plain.slice(u16, 6 * n), filtered.slice(u16, 6 * n)); } +test "the dense projection checks its borrowed buffers against the output it writes" { + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try ex.Pipelines.load(device); + defer p.deinit(); + const owned = try device.buffer(n * k * 2 + (k + n) * 2 + rows * (k + n) * 2, mtl.ResourceOptions.shared); + defer owned.deinit(); + const scratch = try device.buffer(try ex.scratchBytes(k, n, rows), mtl.ResourceOptions.shared); + defer scratch.deinit(); + const spans = @as(usize, n * k * 2); + const weight = ex.View{ .buffer = owned, .offset = 0 }; + const suh = ex.View{ .buffer = owned, .offset = spans }; + const svh = ex.View{ .buffer = owned, .offset = spans + k * 2 }; + const xv = ex.View{ .buffer = owned, .offset = spans + (k + n) * 2 }; + const yv = ex.View{ .buffer = owned, .offset = xv.offset + rows * k * 2 }; + const sv = ex.View{ .buffer = scratch, .offset = 0 }; + const r = ex.Rows{ .count = rows, .x_stride = k, .y_stride = n }; + const cb = queue.commandBuffer(); + const e = cb.compute(.concurrent); + // A weight or scale sharing the output's storage is read while `y` is written, so both are rejected. + const aliased = ex.DenseMatrix{ .weight = yv, .suh = suh, .svh = svh, .k = k, .n = n }; + try std.testing.expectError(error.AliasedExl3View, ex.encodeDense(&p, e, aliased, xv, yv, sv, r)); + try std.testing.expectError(error.AliasedExl3View, ex.encodeDense(&p, e, .{ .weight = weight, .suh = yv, .svh = svh, .k = k, .n = n }, xv, yv, sv, r)); + try ex.encodeDense(&p, e, .{ .weight = weight, .suh = suh, .svh = svh, .k = k, .n = n }, xv, yv, sv, r); +} + test "rejected views leave an encoder reusable and successful output retains its guards" { const device = try mtl.Device.init(); defer device.deinit(); diff --git a/zig/src/core/expert_groups.zig b/zig/src/core/expert_groups.zig new file mode 100644 index 0000000000..5b672ab6fd --- /dev/null +++ b/zig/src/core/expert_groups.zig @@ -0,0 +1,101 @@ +//! Borrowed bounded GPU route grouping; padded eight-row tiles preserve each original output slot. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +const exl3 = @import("exl3_projection.zig"); +pub const Field = enum { counts, offsets, ranks, expert_ids, token_ids, slot_ids }; +const field_count = std.enums.values(Field).len; + +pub const Plan = struct { + rows: u32, + experts: u32, + padded: u32, + offsets: [field_count]usize, + bytes: usize, + + pub fn init(rows: usize, experts: usize) !Plan { + if (rows == 0 or experts == 0 or rows > std.math.maxInt(u32) or experts > std.math.maxInt(u32)) + return error.InvalidExpertGrouping; + const active = @min(rows, experts); + const tiles = try std.math.add(usize, active, (rows - active) / 8); + const padded = try std.math.mul(usize, tiles, 8); + if (padded > std.math.maxInt(u32)) return error.ExpertGroupingTooLarge; + const sizes = [_]usize{ experts * 4, experts * 4, rows * 4, padded * 4, padded * 4, padded * 4 }; + var plan = Plan{ .rows = @intCast(rows), .experts = @intCast(experts), .padded = @intCast(padded), .offsets = undefined, .bytes = 0 }; + for (sizes, 0..) |size, i| { + plan.offsets[i] = std.mem.alignForward(usize, plan.bytes, 16); + plan.bytes = try std.math.add(usize, plan.offsets[i], size); + } + plan.bytes = std.mem.alignForward(usize, plan.bytes, 16); + return plan; + } +}; + +pub const Scratch = struct { + buffer: mtl.Buffer, + offset: usize = 0, + plan: Plan, + + pub fn view(self: Scratch, field: Field) exl3.View { + return .{ .buffer = self.buffer, .offset = self.offset + self.plan.offsets[@backingInt(field)] }; + } +}; + +pub const Pipelines = struct { + stages: [4]mtl.Pipeline, + + pub fn load(device: mtl.Device) !Pipelines { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const lib = try mtl.Library.fromSource(device, sources.core_expert_groups, mtl.CompileOptions.mlx()); + defer lib.deinit(); + var stages: [4]mtl.Pipeline = undefined; + var loaded: usize = 0; + errdefer for (stages[0..loaded]) |pipeline| pipeline.deinit(); + inline for (.{ "tf_expert_group_clear", "tf_expert_group_count", "tf_expert_group_prefix", "tf_expert_group_scatter" }, 0..) |name, i| { + stages[i] = try mtl.Pipeline.init(device, lib, name, false); + loaded += 1; + } + return .{ .stages = stages }; + } + pub fn deinit(self: Pipelines) void { + for (self.stages) |pipeline| pipeline.deinit(); + } +}; +const Args = extern struct { rows: u32, experts: u32, repeat: u32, padded: u32 }; + +pub fn encode(enc: mtl.ComputeEncoder, pipelines: *const Pipelines, ids: exl3.View, scratch: Scratch, repeat: u32) !void { + // Plan must come from Plan.init; route IDs remain GPU-resident and are bounds-checked by the kernels. + if (repeat == 0 or scratch.plan.rows == 0 or scratch.plan.experts == 0 or + scratch.plan.rows % repeat != 0) return error.InvalidExpertGrouping; + if (ids.offset % 4 != 0 or scratch.offset % 4 != 0 or + ids.offset > ids.buffer.length() or @as(usize, scratch.plan.rows) * 4 > ids.buffer.length() - ids.offset or + scratch.offset > scratch.buffer.length() or scratch.plan.bytes > scratch.buffer.length() - scratch.offset) + return error.InvalidExpertGroupingStorage; + const args = Args{ .rows = scratch.plan.rows, .experts = scratch.plan.experts, .repeat = repeat, .padded = scratch.plan.padded }; + for (pipelines.stages, 0..) |pipeline, phase| { + enc.setPipeline(pipeline); + enc.setBuffer(ids.buffer, ids.offset, 0); + inline for (std.enums.values(Field), 1..) |field, i| { + const view = scratch.view(field); + enc.setBuffer(view.buffer, view.offset, i); + } + enc.setValue(args, 7); + enc.dispatchThreads(.{ .width = switch (phase) { + 0 => @max(args.experts, args.padded), + 2 => 1, + else => args.rows, + } }, .{ .width = 256 }); + enc.barrier(); + } +} + +test "expert grouping bounds padding without context-sized allocation" { + const small = try Plan.init(6, 384); + try std.testing.expectEqual(@as(u32, 48), small.padded); + const full = try Plan.init(1536, 384); + try std.testing.expectEqual(@as(u32, 4224), full.padded); + try std.testing.expect(full.bytes < 64 * 1024); + try std.testing.expectError(error.InvalidExpertGrouping, Plan.init(0, 384)); + try std.testing.expectError(error.ExpertGroupingTooLarge, Plan.init(std.math.maxInt(u32), 384)); +} diff --git a/zig/src/core/gpu_counters.zig b/zig/src/core/gpu_counters.zig index 2e0b216165..72edee1343 100644 --- a/zig/src/core/gpu_counters.zig +++ b/zig/src/core/gpu_counters.zig @@ -1,4 +1,4 @@ -//! Same-command-buffer GPU timestamps, bounded resources and no per-span CPU waits. +//! Bounded same-command GPU timestamps; disabled init allocates nothing, no per-span CPU waits or CPU timing fallback. const std = @import("std"); const mtl = @import("metal"); const counters = mtl.counters; @@ -9,6 +9,8 @@ pub const max_span_capacity: usize = 65536; pub const Options = struct { enabled: bool = false, max_spans: usize = 512 }; pub const Span = struct { kind: Class = .dense, + /// Caller-owned diagnostic label; does not affect sampling or generic class aggregation. + metadata: u32 = 0, begin_timestamp: u64 = 0, end_timestamp: u64 = 0, gpu_seconds: f64 = 0, @@ -16,7 +18,7 @@ pub const Span = struct { pub const Aggregate = struct { calls: u64 = 0, gpu_seconds: f64 = 0 }; const State = enum { idle, recording, resolved, read, broken }; -/// Two reference samples convert GPU ticks to nanoseconds; never use CPU span time. +/// Apple's paired reference samples convert GPU ticks to CPU nanoseconds, not elapsed operator time. pub const Calibration = struct { before: counters.ClockSample, after: counters.ClockSample, @@ -29,7 +31,8 @@ pub const Calibration = struct { } pub fn duration(self: Calibration, begin: u64, end: u64) !f64 { - if (begin == 0 or end == 0 or begin == counters.error_value or end == counters.error_value) return error.InvalidCounterSample; + if (begin == 0 or end == 0 or begin == counters.error_value or end == counters.error_value) + return error.InvalidCounterSample; if (end < begin) return error.ReversedCounterSample; if (begin < self.before.gpu or end > self.after.gpu) return error.CounterOutsideCalibration; const ticks: f64 = @floatFromInt(end - begin); @@ -79,7 +82,7 @@ pub const Recorder = struct { calibration: ?Calibration = null, totals: [class_count]Aggregate = @splat(.{}), - /// Init-only bounded storage; disabled init allocates nothing. + /// Init-only storage, descriptors and fence; 1..65536 spans, at most 16 CPU-resolve bytes per span. pub fn init(allocator: std.mem.Allocator, device: mtl.Device, options: Options) !Recorder { var self = Recorder{ .allocator = allocator }; if (!options.enabled) return self; @@ -120,8 +123,8 @@ pub const Recorder = struct { self.state = .recording; } - /// Non-nested spans; stage mode rotates fenced encoders, requiring resource and pipeline rebinding. - pub fn beginSpan(self: *Recorder, encoder: *mtl.ComputeEncoder, kind: Class) !void { + /// No nesting; stage sampling rotates fenced serial encoders: callers must rebind and redeclare resources. + pub fn beginSpan(self: *Recorder, encoder: *mtl.ComputeEncoder, kind: Class, metadata: u32) !void { if (!self.enabled()) return; if (self.state != .recording or !self.encoder_open) return error.CounterNotRecording; if (self.active_span != null) return error.NestedCounterSpan; @@ -141,7 +144,7 @@ pub const Recorder = struct { self.encoder_open = true; resources.fence.waitCompute(encoder.*); } - self.spans[index] = .{ .kind = kind }; + self.spans[index] = .{ .kind = kind, .metadata = metadata }; self.span_count += 1; self.active_span = index; } @@ -172,7 +175,7 @@ pub const Recorder = struct { self.active_span = null; } - /// Ends the final encoder before the caller's existing commit/wait; resolution requires completion. + /// End the final encoder before existing commit/wait; resolve counters only after command completion. pub fn finishEncoding(self: *Recorder, encoder: *mtl.ComputeEncoder) !void { if (!self.enabled()) { encoder.end(); @@ -186,7 +189,7 @@ pub const Recorder = struct { self.state = .resolved; } - /// Cancels unsubmitted work, closing any surviving encoder, including after encoder creation failure. + /// Close and discard failed, unsubmitted encoding without committing; safe after encoder creation failure. pub fn cancelEncoding(self: *Recorder, encoder: *mtl.ComputeEncoder) !void { if (self.command_buffer) |cb| { if (cb.status() != .not_enqueued) return error.CounterCommandAlreadySubmitted; @@ -198,7 +201,7 @@ pub const Recorder = struct { try self.discardAfterCompletion(); } - /// No wait or fallback; results are borrowed until beginCommand, discardAfterCompletion or deinit. + /// No wait or fallback; results are borrowed until reuse/discard/deinit, errors and invalid calibration propagate. pub fn readAfterCompletion(self: *Recorder) ![]const Span { if (!self.enabled()) return &.{}; if (self.state == .read) return self.spans[0..self.span_count]; @@ -244,7 +247,7 @@ pub const Recorder = struct { return self.totals; } - /// Discards closed, unsubmitted or completed recordings; in-flight resources are never reused. + /// Discard completed/unsubmitted recordings after closing encoders; in-flight reuse is rejected, never waited. pub fn discardAfterCompletion(self: *Recorder) !void { if (self.encoder_open) return error.CounterEncoderStillOpen; if (self.command_buffer) |cb| { @@ -288,14 +291,14 @@ test "disabled recorder allocates nothing and never touches a device or encoder var recorder = try Recorder.init(std.testing.failing_allocator, undefined, .{ .enabled = false }); defer recorder.deinit() catch unreachable; var encoder: mtl.ComputeEncoder = undefined; - try recorder.beginSpan(&encoder, .attention); + try recorder.beginSpan(&encoder, .attention, 0); try recorder.endSpan(&encoder); try std.testing.expect(!recorder.enabled()); try std.testing.expectEqual(@as(usize, 0), (try recorder.readAfterCompletion()).len); for (recorder.aggregates()) |aggregate| try std.testing.expectEqual(@as(u64, 0), aggregate.calls); } -test "calibration rejects missing samples, both reversed clocks and both range boundaries" { +test "calibration rejects missing samples and both clock and range boundaries" { const c = Calibration{ .before = .{ .cpu = 1000, .gpu = 100 }, .after = .{ .cpu = 5000, .gpu = 300 } }; try std.testing.expectError(error.InvalidCounterSample, c.duration(0, 200)); try std.testing.expectError(error.InvalidCounterSample, c.duration(100, counters.error_value)); @@ -308,12 +311,45 @@ test "calibration rejects missing samples, both reversed clocks and both range b } test "enabled capacity validation occurs before any device access" { - try std.testing.expectError(error.InvalidSpanCapacity, Recorder.init(std.testing.failing_allocator, undefined, .{ - .enabled = true, - .max_spans = 0, - })); - try std.testing.expectError(error.InvalidSpanCapacity, Recorder.init(std.testing.failing_allocator, undefined, .{ - .enabled = true, - .max_spans = max_span_capacity + 1, - })); + try std.testing.expectError(error.InvalidSpanCapacity, Recorder.init(std.testing.failing_allocator, undefined, .{ .enabled = true, .max_spans = 0 })); + try std.testing.expectError(error.InvalidSpanCapacity, Recorder.init(std.testing.failing_allocator, undefined, .{ .enabled = true, .max_spans = max_span_capacity + 1 })); +} + +test "counter ownership rejects open and in-flight reuse without a CPU span wait" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + if (counters.query(device).mode() == .unsupported) return error.SkipZigTest; + const queue = try device.queue(); + defer queue.deinit(); + const gate = try device.sharedEvent(); + defer gate.deinit(); + var recorder = try Recorder.init(std.testing.allocator, device, .{ .enabled = true, .max_spans = 1 }); + defer recorder.deinit() catch unreachable; + const cb = queue.commandBuffer(); + cb.waitFor(gate, 1); + var encoder = cb.compute(.serial); + try recorder.beginCommand(cb); + var submitted = false; + defer { + if (submitted) { + gate.set(1); + cb.wait(); + recorder.discardAfterCompletion() catch unreachable; + } else recorder.cancelEncoding(&encoder) catch unreachable; + } + try std.testing.expectError(error.CounterEncoderStillOpen, recorder.deinit()); + try recorder.beginSpan(&encoder, .dense, 42); + try std.testing.expectError(error.NestedCounterSpan, recorder.beginSpan(&encoder, .copy, 43)); + try std.testing.expectError(error.UnclosedCounterSpan, recorder.finishEncoding(&encoder)); + try recorder.endSpan(&encoder); + try std.testing.expectError(error.CounterSpanCapacityExceeded, recorder.beginSpan(&encoder, .copy, 43)); + try recorder.finishEncoding(&encoder); + cb.commit(); + submitted = true; + try std.testing.expectError(error.CounterCommandNotCompleted, recorder.readAfterCompletion()); + try std.testing.expectError(error.CounterRecordingInFlight, recorder.beginCommand(cb)); + try std.testing.expectError(error.CounterCommandNotCompleted, recorder.discardAfterCompletion()); + try std.testing.expectError(error.CounterCommandNotCompleted, recorder.deinit()); } diff --git a/zig/src/families/deepseek_v41/attention.zig b/zig/src/families/deepseek_v41/attention.zig new file mode 100644 index 0000000000..84864c972f --- /dev/null +++ b/zig/src/families/deepseek_v41/attention.zig @@ -0,0 +1,273 @@ +//! Pinned DeepSeek V4.1 attention, resident EXL3 projections and causal packed-cache chunk execution. +const mtl = @import("metal"); +const exl3 = @import("core").exl3_projection; +const checkpoint = @import("core").checkpoint_metal; +const weights = @import("weights.zig"); +const config = @import("config.zig"); +const ops = @import("ops.zig"); +const csa2 = @import("csa2.zig"); +const kv = @import("kv.zig"); +const index = @import("csa2_index.zig"); +const rope = @import("rope.zig"); +pub const View = @import("view.zig").View; +pub const State = kv.cache.State; + +/// Opt-in bounded intermediate capture; ordinary execution does not install an observer. +pub const DiagnosticStage = enum { + qr_projection, + qr_norm, + q_projection, + q, + kv_projection, + kv_norm, + kv_rope, + kv_quantized, + sparse_o, + o, + grouped, + wo_b, +}; +pub const DiagnosticObserver = struct { + ctx: *anyopaque, + capture: *const fn (*anyopaque, mtl.ComputeEncoder, DiagnosticStage, View) anyerror!void, + + fn observe(comptime enabled: bool, self: ?DiagnosticObserver, e: mtl.ComputeEncoder, stage: DiagnosticStage, source: View) !void { + if (enabled) if (self) |observer| try observer.capture(observer.ctx, e, stage, source); + } +}; + +pub const Pipelines = struct { + plain: ops.Pipelines, + cache: kv.Pipelines, + select: index.Pipelines, + rotary: rope.Pipelines, + sparse: csa2.Pipelines, + + pub fn load(device: mtl.Device, capacity: u32, c: config.Text) !Pipelines { + var plain = try ops.Pipelines.load(device); + errdefer plain.deinit(); + var cache = try kv.Pipelines.load(device); + errdefer cache.deinit(); + var select = try index.Pipelines.load(device); + errdefer select.deinit(); + var rotary = try rope.Pipelines.load(device, capacity, c.qk_rope_head_dim, .{ .theta = c.rope_theta, .original = 0 }, .{ .theta = c.compress_rope_theta, .factor = c.rope_scaling.factor, .original = c.rope_scaling.original_max_position_embeddings, .beta_fast = c.rope_scaling.beta_fast, .beta_slow = c.rope_scaling.beta_slow }); + errdefer rotary.deinit(); + const sparse = try csa2.Pipelines.load(device); + return .{ .plain = plain, .cache = cache, .select = select, .rotary = rotary, .sparse = sparse }; + } + pub fn residentBytes(capacity: u32, c: config.Text) usize { + return rope.Pipelines.residentBytes(capacity, c.qk_rope_head_dim); + } + pub fn registerResidency(p: *const Pipelines, set: mtl.ResidencySet) void { + p.rotary.registerResidency(set); + } + pub fn deinit(p: *Pipelines) void { + p.plain.deinit(); + p.cache.deinit(); + p.select.deinit(); + p.rotary.deinit(); + p.sparse.deinit(); + } +}; + +// Row-sized regions never contain dequantized weights or a query-by-context tensor. +const regions = [_]usize{ 1280 * 2, 32768 * 2, 512 * 2, 512 * 4, 512 * 4, 512 * 2, 128 * 2, 4096 * 2, 32 * 2, 8192 * 2 }; +const Region = enum(usize) { qr, query, raw_kv, values, gates, latent, index_k, index_q, index_weights, grouped }; + +pub const Workspace = struct { + buffer: mtl.Buffer, + selection: index.Scratch, + max_rows: u32, + selection_layer: ?u32 = null, + selection_position: u32 = 0, + selection_rows: u32 = 0, + candidate_position: u32 = 0, + candidate_rows: u32 = 0, + + pub fn bytes(max_rows: u32) usize { + var per_row: usize = 0; + for (regions) |size| per_row += size; + return @as(usize, max_rows) * per_row + index.Scratch.bytes(max_rows); + } + pub fn init(device: mtl.Device, max_rows: u32) !Workspace { + if (max_rows == 0 or max_rows > 4096) return error.InvalidAttentionWorkspace; + const buffer = try device.buffer(bytes(max_rows) - index.Scratch.bytes(max_rows), mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + errdefer buffer.deinit(); + const selection = try index.Scratch.init(device, max_rows); + return .{ .buffer = buffer, .selection = selection, .max_rows = max_rows }; + } + pub fn deinit(w: *Workspace) void { + w.buffer.deinit(); + w.selection.deinit(); + } + pub fn residentBytes(w: *const Workspace) usize { + return bytes(w.max_rows); + } + pub fn registerResidency(w: *const Workspace, set: mtl.ResidencySet) void { + set.add(w.buffer); + set.add(w.selection.buffer); + } + fn view(w: *const Workspace, region: Region) View { + var offset: usize = 0; + for (regions[0..@backingInt(region)]) |size| offset += size * w.max_rows; + return .{ .buffer = w.buffer, .offset = offset }; + } +}; + +fn tensor(t: checkpoint.Tensor) View { + return .{ .buffer = t.buffer, .offset = t.offset }; +} +fn core(v: View) exl3.View { + return .{ .buffer = v.buffer, .offset = v.offset }; +} +fn projection(p: *const exl3.Pipelines, e: mtl.ComputeEncoder, m: weights.AnyMatrix, x: View, y: View, scratch: exl3.View, rows: u32, xs: u32, ys: u32, output: exl3.Output) !void { + switch (m) { + .trellis => |t| try exl3.encode(p, e, t, core(x), core(y), scratch, .{ .count = rows, .x_stride = xs, .y_stride = ys, .output = output }), + // A kept-dense repack: the same rotate stages, the middle a plain dense product. + .dense => |d| try exl3.encodeDense(p, e, .{ + .weight = core(tensor(d.weight)), + .suh = core(tensor(d.suh)), + .svh = core(tensor(d.svh)), + .k = d.k, + .n = d.n, + }, core(x), core(y), scratch, .{ .count = rows, .x_stride = xs, .y_stride = ys, .output = output }), + } +} +fn norm(p: *const Pipelines, e: mtl.ComputeEncoder, x: View, weight: checkpoint.Tensor, rows: u32, width: u32, eps: f32) void { + ops.encodeRms(e, &p.plain, x, tensor(weight), x, .{ .rows = rows, .width = width, .x_stride = width, .y_stride = width, .eps = eps, .weight_f32 = weight.dtype == .f32 }); + e.barrier(); +} +fn rotate(p: *const Pipelines, e: mtl.ComputeEncoder, x: View, rows: u32, heads: u32, width: u32, position: u32, ratio: u32, step: u32, inverse: bool, c: config.Text) !void { + try p.rotary.encode(e, x, .{ .rows = rows, .heads = heads, .width = width, .rotary = c.qk_rope_head_dim, .stride = heads * width, .position = position, .step = step, .inverse = @intFromBool(inverse), .spectrum = if (ratio == 0) .normal else .compressed }); +} +fn precedingIndex(layer: u32) u32 { + var source: u32 = 2; + for (kv.cache.index_sources) |id| { + if (id > layer) break; + source = id; + } + return source; +} + +/// BF16 x/y [rows,5120]; layers 0..39 share the chunk-start cursor, appended by the caller after all layers. + +// No extra RMS after wq_b or index-query Hadamard outside EXL3. +pub fn encode(p: *const Pipelines, e: mtl.ComputeEncoder, w: *const weights.Attention, state: *State, work: *Workspace, x: View, y: View, ep: *const exl3.Pipelines, scratch: exl3.View, layer: u32, rows: u32, c: config.Text) !void { + return encodeDiagnostic(p, e, w, state, work, x, y, ep, scratch, layer, rows, c, false, null); +} + +/// Same encoder and arithmetic as encode, with optional copies at operator boundaries. +pub fn encodeDiagnostic(p: *const Pipelines, e: mtl.ComputeEncoder, w: *const weights.Attention, state: *State, work: *Workspace, x: View, y: View, ep: *const exl3.Pipelines, scratch: exl3.View, layer: u32, rows: u32, c: config.Text, comptime diagnostic: bool, observer: ?DiagnosticObserver) !void { + if (layer >= 40 or rows == 0 or rows > work.max_rows or + state.position > state.plan.capacity or rows > state.plan.capacity - state.position) + return error.InvalidAttentionChunk; + const position = state.position; + const ratio = c.compress_ratios[layer]; + const end = position + rows; + const q_width = c.num_attention_heads * c.head_dim; + const qr = work.view(.qr); + const q = work.view(.query); + const raw = work.view(.raw_kv); + const ring = try state.ring(layer); + const chunk = work.selection.windowChunk(); + e.barrier(); + try projection(ep, e, .{ .trellis = w.wq_a }, x, qr, scratch, rows, c.hidden_size, c.q_lora_rank, .bf16); + if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .qr_projection, qr); + norm(p, e, qr, w.q_norm, rows, c.q_lora_rank, c.rms_norm_eps); + if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .qr_norm, qr); + try projection(ep, e, .{ .trellis = w.wq_b }, qr, q, scratch, rows, c.q_lora_rank, q_width, .bf16); + if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .q_projection, q); + try rotate(p, e, q, rows, c.num_attention_heads, c.head_dim, position, ratio, 1, false, c); + if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .q, q); + try projection(ep, e, .{ .trellis = w.wkv }, x, raw, scratch, rows, c.hidden_size, c.head_dim, .bf16); + if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .kv_projection, raw); + norm(p, e, raw, w.kv_norm, rows, c.head_dim, c.rms_norm_eps); + if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .kv_norm, raw); + try rotate(p, e, raw, rows, 1, c.head_dim, position, ratio, 1, false, c); + if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .kv_rope, raw); + try p.cache.encodePack(e, raw, chunk, .{ .rows = rows, .width = c.head_dim, .stride = c.head_dim, .packed_stride = kv.cache.ring_stride, .block = 32, .kind = .window_fp8 }); + if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .kv_quantized, chunk); + + var main = chunk; + var compressed_rows: u32 = 0; + if (ratio != 0) { + const source = try state.source(layer, ratio); + main = source.main; + compressed_rows = end / ratio; + const complete = compressed_rows - position / ratio; + const latent = work.view(.latent); + if (source.owner == layer) { + const compressor = if (w.compressor) |*value| value else return error.MissingAttentionCompressor; + const output: exl3.Output = if (ratio == 1) .bf16 else .f32; + const values = work.view(.values); + const gates = work.view(.gates); + try projection(ep, e, compressor.wkv, x, values, scratch, rows, c.hidden_size, c.head_dim, output); + if (ratio > 1) { + const gate = compressor.wgate orelse return error.MissingAttentionCompressorGate; + try projection(ep, e, gate, x, gates, scratch, rows, c.hidden_size, c.head_dim, .f32); + } + try p.cache.encodePool(e, values, gates, source.partial, tensor(compressor.norm), latent, .{ .rows = rows, .ratio = ratio, .position = position, .stride = c.head_dim, .input_f32 = @intFromBool(ratio > 1), .eps = c.rms_norm_eps }); + if (complete != 0) { + const selector = if (w.indexer) |*value| value else return error.MissingAttentionIndexer; + const wk = selector.wk orelse return error.MissingAttentionIndexKeyProjection; + const k_norm = selector.k_norm orelse return error.MissingAttentionIndexKeyNorm; + const key = work.view(.index_k); + try projection(ep, e, wk, latent, key, scratch, complete, c.head_dim, c.index_head_dim, .bf16); + norm(p, e, key, k_norm, complete, c.index_head_dim, c.rms_norm_eps); + try rotate(p, e, key, complete, 1, c.index_head_dim, position / ratio * ratio, ratio, ratio, false, c); + try p.cache.encodePack(e, key, source.index, .{ .rows = complete, .width = c.index_head_dim, .stride = c.index_head_dim, .packed_stride = kv.cache.index_stride, .block = 32, .kind = .index_fp4, .first = position / ratio }); + } + } + if (kv.cache.isIndexSource(layer)) { + const selector = if (w.indexer) |*value| value else return error.MissingAttentionIndexer; + const iq = work.view(.index_q); + const iw = work.view(.index_weights); + const iq_width = c.index_n_heads * c.index_head_dim; + if (compressed_rows != 0) { + try projection(ep, e, selector.wq_b, qr, iq, scratch, rows, c.q_lora_rank, iq_width, .bf16); + try rotate(p, e, iq, rows, c.index_n_heads, c.index_head_dim, position, ratio, 1, false, c); + try p.cache.encodePack(e, iq, work.selection.queryPacked(), .{ .rows = rows * c.index_n_heads, .width = c.index_head_dim, .stride = c.index_head_dim, .packed_stride = kv.cache.index_stride, .block = 32, .kind = .index_fp4, .inplace = 1 }); + ops.encodeDense(e, &p.plain, x, tensor(selector.weights_proj), iw, .{ .rows = rows, .k = c.hidden_size, .n = c.index_n_heads, .x_stride = c.hidden_size, .w_stride = c.hidden_size, .y_stride = c.index_n_heads, .w_dtype = switch (selector.weights_proj.dtype) { + .bf16 => .bf16, + .f16 => .f16, + .f32 => .f32, + else => unreachable, + } }); + e.barrier(); + } + const make_candidates = layer == c.candidate_source_layer_id; + const use_candidates = layer > c.candidate_source_layer_id; + if (use_candidates and (work.candidate_position != position or work.candidate_rows != rows)) + return error.MissingAttentionCandidates; + try p.select.encode(e, iq, iw, source.index, work.selection, .{ .rows = rows, .q_stride = iq_width, .weight_stride = c.index_n_heads, .position = position, .ratio = ratio, .keys = compressed_rows, .make_candidates = @intFromBool(make_candidates), .use_candidates = @intFromBool(use_candidates) }); + work.selection_layer = layer; + work.selection_position = position; + work.selection_rows = rows; + if (make_candidates) { + work.candidate_position = position; + work.candidate_rows = rows; + } + } else if (work.selection_layer != precedingIndex(layer) or + work.selection_position != position or work.selection_rows != rows) return error.MissingAttentionSelection; + // Index K consumes the normalized latent before this in-place RoPE mutates it. + if (source.owner == layer and complete != 0) { + try rotate(p, e, latent, complete, 1, c.head_dim, position / ratio * ratio, ratio, ratio, false, c); + try p.cache.encodePack(e, latent, source.main, .{ .rows = complete, .width = c.head_dim, .stride = c.head_dim, .packed_stride = kv.cache.main_stride, .block = 16, .kind = .main_fp4, .first = position / ratio }); + } + } + try p.sparse.encode(e, q, ring, chunk, main, work.selection.positions(), tensor(w.sink), q, .{ .rows = rows, .q_stride = q_width, .position = position, .chunk_rows = rows, .ratio = ratio, .compressed_rows = compressed_rows, .topk = @min(c.index_topk, compressed_rows) }); + e.barrier(); + if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .sparse_o, q); + try p.cache.publishWindow(e, chunk, ring, rows, position); + try rotate(p, e, q, rows, c.num_attention_heads, c.head_dim, position, ratio, 1, true, c); + if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .o, q); + const grouped = work.view(.grouped); + const group_input = q_width / c.o_groups; + const group_output = c.o_groups * c.o_lora_rank; + for (w.wo_a, 0..) |matrix, group| { + try projection(ep, e, .{ .trellis = matrix }, q.shifted(group * group_input * 2), grouped.shifted(group * c.o_lora_rank * 2), scratch, rows, q_width, group_output, .bf16); + } + if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .grouped, grouped); + try projection(ep, e, .{ .trellis = w.wo_b }, grouped, y, scratch, rows, group_output, c.hidden_size, .bf16); + if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .wo_b, y); +} diff --git a/zig/src/families/deepseek_v41/csa2.zig b/zig/src/families/deepseek_v41/csa2.zig new file mode 100644 index 0000000000..395de9107a --- /dev/null +++ b/zig/src/families/deepseek_v41/csa2.zig @@ -0,0 +1,77 @@ +//! Exact causal SWA plus shared compressed KV; fixed-buffer chunked prefill and decode use one path. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +const View = @import("view.zig").View; +pub const index = @import("csa2_index.zig"); +pub const kv = @import("kv.zig"); +pub const rope = @import("rope.zig"); +pub const Params = extern struct { + rows: u32, + heads: u32 = 64, + width: u32 = 512, + q_stride: u32, + position: u32, + chunk_rows: u32, + window: u32 = 128, + ring_stride: u32 = 528, + chunk_stride: u32 = 528, + main_stride: u32 = 288, + main_block: u32 = 16, + ratio: u32, + compressed_rows: u32, + topk: u32 = 512, + index_stride: u32 = 512, + scale: f32 = 0.04419417382415922, + pub fn validate(p: Params) !void { + if (p.rows == 0 or p.heads == 0 or p.width == 0 or p.width > 512 or p.width % 32 != 0 or + @as(u64, p.q_stride) < @as(u64, p.heads) * p.width or + p.chunk_rows < p.rows or p.window != 128 or + p.ring_stride < p.width + p.width / 32 or p.chunk_stride < p.width + p.width / 32 or + p.main_stride < p.width / 2 + p.width / 16 or p.main_block != 16 or + (p.ratio != 0 and p.ratio != 1 and p.ratio != 2) or p.topk > 512 or p.index_stride < p.topk or + @as(u64, p.position) + p.rows > kv.cache.max_context or + p.compressed_rows > kv.cache.max_context or !std.math.isFinite(p.scale) or p.scale <= 0 or + @as(u64, p.rows -| 1) * p.q_stride + @as(u64, p.heads) * p.width > std.math.maxInt(u32) or + @as(u64, p.rows) * p.index_stride > std.math.maxInt(u32) or + @as(u64, p.rows) * p.chunk_stride > std.math.maxInt(u32) or + @as(u64, p.window) * p.ring_stride > std.math.maxInt(u32) or + @as(u64, p.compressed_rows) * p.main_stride > std.math.maxInt(u32)) + return error.InvalidSparseAttentionShape; + } +}; +pub const Pipelines = struct { + attention: mtl.Pipeline, + pub fn load(device: mtl.Device) !Pipelines { + var options = mtl.CompileOptions.mlx(); + options.math_functions = .precise; + const lib = try mtl.Library.fromSource(device, sources.deepseek_v41_kv ++ sources.deepseek_v41_csa2, options); + defer lib.deinit(); + return .{ .attention = try mtl.Pipeline.init(device, lib, "dsv41_sparse_attention", false) }; + } + pub fn deinit(p: *Pipelines) void { + p.attention.deinit(); + } + /// BF16 Q/output [rows,heads,width], FP32 sinks, I32 indices; publish this chunk to the ring only after dispatch. + pub fn encode(p: *const Pipelines, enc: mtl.ComputeEncoder, q: View, ring: View, chunk: View, compressed: View, positions: View, sink: View, out: View, args: Params) !void { + try args.validate(); + try q.requireRows(args.rows, args.q_stride, @as(usize, args.heads) * args.width, 2); + try out.requireRows(args.rows, args.q_stride, @as(usize, args.heads) * args.width, 2); + try sink.require(@as(usize, args.heads) * 4, 4); + if (args.position != 0) try ring.requireRows(@min(args.position, args.window), args.ring_stride, args.width + args.width / 32, 1); + try chunk.requireRows(args.rows, args.chunk_stride, args.width + args.width / 32, 1); + if (args.ratio != 0 and args.compressed_rows != 0 and args.topk != 0) { + try compressed.requireRows(args.compressed_rows, args.main_stride, args.width / 2 + args.width / 16, 1); + try positions.requireRows(args.rows, args.index_stride, args.topk, 4); + } + enc.setPipeline(p.attention); + for ([_]View{ q, ring, chunk, compressed, positions, sink, out }, 0..) |v, i| + enc.setBuffer(v.buffer, v.offset, i); + enc.setValue(args, 7); + enc.dispatchGroups(mtl.Size.of(args.heads, args.rows, 1), mtl.Size.of(32, 1, 1)); + } +}; + +test { + _ = @import("kv_test.zig"); +} diff --git a/zig/src/families/deepseek_v41/csa2_index.zig b/zig/src/families/deepseek_v41/csa2_index.zig new file mode 100644 index 0000000000..ed2ee3a833 --- /dev/null +++ b/zig/src/families/deepseek_v41/csa2_index.zig @@ -0,0 +1,139 @@ +//! Tiled GPU scores and bounded online heaps; the memory footprint does not grow with context length. +const std = @import("std"); +const max_context = @import("metal_kv.zig").max_context; +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +const View = @import("view.zig").View; +pub const score_tile_width = 4096; +pub const Params = extern struct { + rows: u32, + heads: u32 = 32, + width: u32 = 128, + q_stride: u32, + weight_stride: u32, + position: u32, + ratio: u32, + keys: u32, + key_stride: u32 = 68, + tile_start: u32 = 0, + tile_width: u32 = score_tile_width, + topk: u32 = 512, + candidate_topk: u32 = 2048, + block_size: u32 = 8, + use_candidates: u32 = 0, + make_candidates: u32 = 0, + scale: f32 = 0.015625, + pub fn validate(p: Params) !void { + if (p.rows == 0 or p.heads == 0 or p.heads > 32 or p.width == 0 or p.width > 128 or p.width % 32 != 0 or + p.q_stride < p.heads * p.width or p.weight_stride < p.heads or + (p.ratio != 1 and p.ratio != 2) or p.key_stride < p.width / 2 + p.width / 32 or + p.topk != 512 or p.candidate_topk != 2048 or p.block_size != 8 or + p.tile_width != score_tile_width or p.use_candidates > 1 or p.make_candidates > 1 or + (p.use_candidates != 0 and p.make_candidates != 0) or + @as(u64, p.position) + p.rows > max_context or p.keys > max_context or + !std.math.isFinite(p.scale) or p.scale <= 0 or + @as(u64, p.rows) * p.q_stride > std.math.maxInt(u32) or + @as(u64, p.rows) * p.weight_stride > std.math.maxInt(u32) or + @as(u64, p.keys) * p.key_stride > std.math.maxInt(u32)) + return error.InvalidIndexerShape; + } +}; +pub const Scratch = struct { + buffer: mtl.Buffer, + rows: u32, + pub fn bytes(rows: u32) usize { + return @as(usize, rows) * (score_tile_width * 4 + 512 * 8 + 2048 * 8 + 512 * 4 + 528 + 32 * 68); + } + pub fn init(device: mtl.Device, rows: u32) !Scratch { + if (rows == 0 or rows > 4096) return error.InvalidIndexerScratch; + return .{ .buffer = try device.buffer(bytes(rows), mtl.ResourceOptions.shared), .rows = rows }; + } + pub fn deinit(s: *Scratch) void { + s.buffer.deinit(); + } + pub fn scores(s: Scratch) View { + return .{ .buffer = s.buffer }; + } + pub fn heap(s: Scratch) View { + return s.scores().shifted(@as(usize, s.rows) * score_tile_width * 4); + } + pub fn candidates(s: Scratch) View { + return s.heap().shifted(@as(usize, s.rows) * 512 * 8); + } + pub fn positions(s: Scratch) View { + return s.candidates().shifted(@as(usize, s.rows) * 2048 * 8); + } + pub fn windowChunk(s: Scratch) View { + return s.positions().shifted(@as(usize, s.rows) * 512 * 4); + } + pub fn queryPacked(s: Scratch) View { + return s.windowChunk().shifted(@as(usize, s.rows) * 528); + } +}; +pub const Pipelines = struct { + score: mtl.Pipeline, + clear: mtl.Pipeline, + merge: mtl.Pipeline, + finish: mtl.Pipeline, + pub fn load(device: mtl.Device) !Pipelines { + var options = mtl.CompileOptions.mlx(); + options.math_functions = .precise; + const lib = try mtl.Library.fromSource(device, sources.deepseek_v41_kv ++ sources.deepseek_v41_index, options); + defer lib.deinit(); + var p: Pipelines = undefined; + p.score = try mtl.Pipeline.init(device, lib, "dsv41_index_score", false); + errdefer p.score.deinit(); + p.clear = try mtl.Pipeline.init(device, lib, "dsv41_index_clear", false); + errdefer p.clear.deinit(); + p.merge = try mtl.Pipeline.init(device, lib, "dsv41_index_merge", false); + errdefer p.merge.deinit(); + p.finish = try mtl.Pipeline.init(device, lib, "dsv41_index_finish", false); + return p; + } + pub fn deinit(p: *Pipelines) void { + p.score.deinit(); + p.clear.deinit(); + p.merge.deinit(); + p.finish.deinit(); + } + /// Q has already undergone BF16 RoPE and in-place FP4/E8M0 quantization; weights are BF16 projection output. + pub fn encode(p: *const Pipelines, enc: mtl.ComputeEncoder, q: View, weights: View, keys: View, scratch: Scratch, args: Params) !void { + try args.validate(); + if (args.rows > scratch.rows) return error.IndexerScratchExceeded; + try (View{ .buffer = scratch.buffer }).require(Scratch.bytes(scratch.rows), 4); + if (args.keys != 0) { + try q.requireRows(args.rows, args.q_stride, @as(usize, args.heads) * args.width, 2); + try weights.requireRows(args.rows, args.weight_stride, args.heads, 2); + try keys.requireRows(args.keys, args.key_stride, args.width / 2 + args.width / 32, 1); + } + enc.setPipeline(p.clear); + for ([_]View{ scratch.heap(), scratch.candidates() }, 0..) |v, i| enc.setBuffer(v.buffer, v.offset, i); + enc.setValue(args, 2); + enc.dispatchThreads(mtl.Size.of(@max(args.topk, args.candidate_topk), args.rows, 1), mtl.Size.of(128, 1, 1)); + enc.barrier(); + const count = if (args.use_candidates != 0) args.candidate_topk * args.block_size else args.keys; + var at: u32 = 0; + while (at < count) : (at += score_tile_width) { + var tile = args; + tile.tile_start = at; + enc.setPipeline(p.score); + for ([_]View{ q, weights, keys, scratch.candidates(), scratch.scores() }, 0..) |v, i| + enc.setBuffer(v.buffer, v.offset, i); + enc.setValue(tile, 5); + enc.dispatchGroups(mtl.Size.of(score_tile_width, args.rows, 1), mtl.Size.of(32, 1, 1)); + enc.barrier(); + enc.setPipeline(p.merge); + for ([_]View{ scratch.scores(), scratch.heap(), scratch.candidates() }, 0..) |v, i| + enc.setBuffer(v.buffer, v.offset, i); + enc.setValue(tile, 3); + enc.dispatchThreads(mtl.Size.of(args.rows, 1, 1), mtl.Size.of(1, 1, 1)); + enc.barrier(); + } + enc.setPipeline(p.finish); + for ([_]View{ scratch.heap(), scratch.candidates(), scratch.positions() }, 0..) |v, i| + enc.setBuffer(v.buffer, v.offset, i); + enc.setValue(args, 3); + enc.dispatchThreads(mtl.Size.of(args.rows, 1, 1), mtl.Size.of(1, 1, 1)); + enc.barrier(); + } +}; diff --git a/zig/src/families/deepseek_v41/csa2_test.zig b/zig/src/families/deepseek_v41/csa2_test.zig new file mode 100644 index 0000000000..0ac81cacca --- /dev/null +++ b/zig/src/families/deepseek_v41/csa2_test.zig @@ -0,0 +1,309 @@ +const std = @import("std"); +const mtl = @import("metal"); +const csa2 = @import("csa2.zig"); +const index = @import("csa2_index.zig"); +const rope = @import("rope.zig"); +const View = @import("view.zig").View; +fn bf(x: f32) u16 { + const bits: u32 = @bitCast(x); + return @intCast((bits +% 0x7fff +% ((bits >> 16) & 1)) >> 16); +} +fn view(b: mtl.Buffer) View { + return .{ .buffer = b }; +} +fn submit(cb: mtl.CommandBuffer, enc: mtl.ComputeEncoder) !void { + enc.end(); + cb.commit(); + cb.wait(); + try std.testing.expect(cb.failure() == null); +} + +test "sparse attention preserves a causal window across ring wrap and masks future compressed keys" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try csa2.Pipelines.load(device); + defer p.deinit(); + const q = try device.buffer(3 * 512 * 2, mtl.ResourceOptions.shared); + defer q.deinit(); + const ring = try device.buffer(128 * 528, mtl.ResourceOptions.shared); + defer ring.deinit(); + const chunk = try device.buffer(3 * 528, mtl.ResourceOptions.shared); + defer chunk.deinit(); + const main = try device.buffer(201 * 288, mtl.ResourceOptions.shared); + defer main.deinit(); + const positions = try device.buffer(6 * 4, mtl.ResourceOptions.shared); + defer positions.deinit(); + const sink = try device.buffer(4, mtl.ResourceOptions.shared); + defer sink.deinit(); + const out = try device.buffer(3 * 512 * 2, mtl.ResourceOptions.shared); + defer out.deinit(); + @memset(q.contents()[0 .. 3 * 512 * 2], 0); + sink.slice(f32, 1)[0] = 0; + for (0..128) |r| { + @memset(ring.contents()[r * 528 .. r * 528 + 512], 0x38); + @memset(ring.contents()[r * 528 + 512 .. (r + 1) * 528], 127); + } + for (0..3) |r| { + @memset(chunk.contents()[r * 528 .. r * 528 + 512], 0x40); + @memset(chunk.contents()[r * 528 + 512 .. (r + 1) * 528], 127); + positions.slice(i32, 6)[2 * r] = 0; + positions.slice(i32, 6)[2 * r + 1] = 200; + } + @memset(main.contents()[0..256], 0x66); + @memset(main.contents()[256..288], 0x38); + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + try p.encode(enc, view(q), view(ring), view(chunk), view(main), view(positions), view(sink), view(out), .{ .rows = 3, .heads = 1, .q_stride = 512, .position = 129, .chunk_rows = 3, .ratio = 2, .compressed_rows = 201, .topk = 2, .index_stride = 2 }); + try submit(cb, enc); + for (0..3) |r| { + const expected = bf((133 + @as(f32, @floatFromInt(r))) / 130); + for (out.slice(u16, 3 * 512)[r * 512 .. (r + 1) * 512]) |v| try std.testing.expectEqual(expected, v); + } +} + +test "sparse QK cancellation retains a small score before sink normalization" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try csa2.Pipelines.load(device); + defer p.deinit(); + const q = try device.buffer(512 * 2, mtl.ResourceOptions.shared); + defer q.deinit(); + const chunk = try device.buffer(528, mtl.ResourceOptions.shared); + defer chunk.deinit(); + const sink = try device.buffer(4, mtl.ResourceOptions.shared); + defer sink.deinit(); + const out = try device.buffer(512 * 2, mtl.ResourceOptions.shared); + defer out.deinit(); + @memset(q.slice(u16, 512), 0); + q.slice(u16, 512)[0] = bf(0x1p24); + q.slice(u16, 512)[1] = bf(-0x1p24); + q.slice(u16, 512)[32] = bf(0x1p-6); + @memset(chunk.contents()[0..512], 0x38); + @memset(chunk.contents()[512..528], 127); + sink.slice(f32, 1)[0] = 0; + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + // No prior window or compressed keys: their unused bindings share the immutable chunk. + try p.encode(enc, view(q), view(chunk), view(chunk), view(chunk), view(chunk), view(sink), view(out), .{ .rows = 1, .heads = 1, .q_stride = 512, .position = 0, .chunk_rows = 1, .ratio = 0, .compressed_rows = 0, .scale = 1 }); + try submit(cb, enc); + // Exact dot = 2^-6; 1 / (1 + exp(-2^-6)) rounds to this BF16 value. + for (out.slice(u16, 512)) |v| try std.testing.expectEqual(bf(0.50390625), v); +} + +test "the first query retains BF16 probability arithmetic regardless of future prefill rows" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try csa2.Pipelines.load(device); + defer p.deinit(); + const q = try device.buffer(32 * 2, mtl.ResourceOptions.shared); + defer q.deinit(); + const ring = try device.buffer(128 * 33, mtl.ResourceOptions.shared); + defer ring.deinit(); + const chunk = try device.buffer(128 * 33, mtl.ResourceOptions.shared); + defer chunk.deinit(); + const main = try device.buffer(18, mtl.ResourceOptions.shared); + defer main.deinit(); + const positions = try device.buffer(4, mtl.ResourceOptions.shared); + defer positions.deinit(); + const sink = try device.buffer(4, mtl.ResourceOptions.shared); + defer sink.deinit(); + const out = try device.buffer(32 * 2, mtl.ResourceOptions.shared); + defer out.deinit(); + @memset(q.slice(u16, 32), 0); + q.slice(u16, 32)[0] = bf(1); + @memset(chunk.contents()[0..32], 0x38); + chunk.contents()[32] = 127; + @memset(main.contents()[0..16], 0x44); + @memset(main.contents()[16..18], 0x38); + positions.slice(i32, 1)[0] = 0; + sink.slice(f32, 1)[0] = 0.005; + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + try p.encode(enc, view(q), view(ring), view(chunk), view(main), view(positions), view(sink), view(out), .{ .rows = 1, .heads = 1, .width = 32, .q_stride = 32, .position = 0, .chunk_rows = 1, .ring_stride = 33, .chunk_stride = 33, .main_stride = 18, .ratio = 1, .compressed_rows = 1, .topk = 1, .index_stride = 1, .scale = 1 }); + try submit(cb, enc); + // BF16(exp(-1))*1 + 1*2; padded windows would move key 2 into a later online-softmax tile. + for (out.slice(u16, 32)) |v| try std.testing.expectEqual(bf(1.5703125), v); + const cb2 = queue.commandBuffer(); + const enc2 = cb2.compute(.serial); + try p.encode(enc2, view(q), view(ring), view(chunk), view(main), view(positions), view(sink), view(out), .{ .rows = 1, .heads = 1, .width = 32, .q_stride = 32, .position = 0, .chunk_rows = 128, .ring_stride = 33, .chunk_stride = 33, .main_stride = 18, .ratio = 1, .compressed_rows = 1, .topk = 1, .index_stride = 1, .scale = 1 }); + try submit(cb2, enc2); + for (out.slice(u16, 32)) |v| try std.testing.expectEqual(bf(1.5703125), v); +} + +test "hierarchical candidate selection pins the newest partial block and subsequent reindex uses those blocks" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try index.Pipelines.load(device); + defer p.deinit(); + var scratch = try index.Scratch.init(device, 1); + defer scratch.deinit(); + const q = try device.buffer(32 * 2, mtl.ResourceOptions.shared); + defer q.deinit(); + const weights = try device.buffer(2, mtl.ResourceOptions.shared); + defer weights.deinit(); + const keys_count = 2048 * 8 + 3; + const keys = try device.buffer(keys_count * 17, mtl.ResourceOptions.shared); + defer keys.deinit(); + @memset(q.slice(u16, 32), bf(1)); + weights.slice(u16, 1)[0] = bf(1); + for (0..keys_count) |r| { + @memset(keys.contents()[r * 17 .. r * 17 + 16], 0x22); + keys.contents()[r * 17 + 16] = if (r >= 2048 * 8) 126 else 127; + } + var args = index.Params{ .rows = 1, .heads = 1, .width = 32, .q_stride = 32, .weight_stride = 1, .position = keys_count - 1, .ratio = 1, .keys = keys_count, .key_stride = 17, .make_candidates = 1, .scale = 1 }; + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + try p.encode(enc, view(q), view(weights), view(keys), scratch, args); + try submit(cb, enc); + const Pick = extern struct { score: f32, position: i32 }; + const candidates = @as([*]Pick, @ptrCast(@alignCast(scratch.candidates().buffer.contents() + scratch.candidates().offset)))[0..2048]; + try std.testing.expectEqual(@as(i32, 2048), candidates[2047].position); + for (candidates[0..2047], 0..) |pick, i| try std.testing.expectEqual(@as(i32, @intCast(i)), pick.position); + const cb2 = queue.commandBuffer(); + const enc2 = cb2.compute(.serial); + args.make_candidates = 0; + args.use_candidates = 1; + try p.encode(enc2, view(q), view(weights), view(keys), scratch, args); + try submit(cb2, enc2); + const selected = @as([*]i32, @ptrCast(@alignCast(scratch.positions().buffer.contents() + scratch.positions().offset)))[0..512]; + for (selected, 0..) |pos, i| try std.testing.expectEqual(@as(i32, @intCast(i)), pos); +} + +test "RoPE rotates adjacent tail pairs with BF16 assignment and compressed positions use group starts" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try rope.Pipelines.load(device, 3, 32, .{ .theta = 10000 }, .{ .theta = 160000 }); + defer p.deinit(); + const x = try device.buffer(2 * 64 * 2, mtl.ResourceOptions.shared); + defer x.deinit(); + @memset(x.slice(u16, 128), bf(1)); + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + try std.testing.expectError(error.InvalidRotaryShape, p.encode(enc, view(x), .{ .rows = 2, .heads = 1, .width = 64, .rotary = 32, .stride = 64, .position = 1, .step = 2 })); + try std.testing.expectError(error.InvalidRotaryShape, p.encode(enc, view(x), .{ .rows = 1, .heads = 1, .width = 64, .rotary = 16, .stride = 64, .position = 0 })); + try std.testing.expectError(error.InvalidRotaryShape, p.encode(enc, view(x), .{ .rows = 2, .heads = 1, .width = 64, .rotary = 32, .stride = 64, .position = std.math.maxInt(u32), .step = std.math.maxInt(u32) })); + try p.encode(enc, view(x), .{ .rows = 2, .heads = 1, .width = 64, .rotary = 32, .stride = 64, .position = 0, .step = 2 }); + try submit(cb, enc); + for (x.slice(u16, 128)[0..96]) |v| try std.testing.expectEqual(bf(1), v); + for (0..16) |pair| { + const power: f32 = @floatCast(std.math.pow(f64, 10000, @as(f32, @floatFromInt(pair * 2)) / 32)); + const frequency: f32 = 1 / power; + const angle: f32 = 2 * frequency; + const c: f32 = @floatCast(@cos(@as(f64, angle))); + const s: f32 = @floatCast(@sin(@as(f64, angle))); + const a = bf(c - s); + const b = bf(s + c); + try std.testing.expectEqual(a, x.slice(u16, 128)[96 + pair * 2]); + try std.testing.expectEqual(b, x.slice(u16, 128)[97 + pair * 2]); + } +} + +test "indexer visibility follows completed compressor groups for every query in a chunk" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try index.Pipelines.load(device); + defer p.deinit(); + var scratch = try index.Scratch.init(device, 3); + defer scratch.deinit(); + const q = try device.buffer(3 * 32 * 2, mtl.ResourceOptions.shared); + defer q.deinit(); + const weights = try device.buffer(3 * 2, mtl.ResourceOptions.shared); + defer weights.deinit(); + const keys = try device.buffer(3 * 17, mtl.ResourceOptions.shared); + defer keys.deinit(); + @memset(q.slice(u16, 3 * 32), bf(1)); + @memset(weights.slice(u16, 3), bf(1)); + for (0..3) |r| { + @memset(keys.contents()[r * 17 .. r * 17 + 16], 0x22); + keys.contents()[r * 17 + 16] = 127; + } + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + try p.encode(enc, view(q), view(weights), view(keys), scratch, .{ .rows = 3, .heads = 1, .width = 32, .q_stride = 32, .weight_stride = 1, .position = 0, .ratio = 2, .keys = 3, .key_stride = 17, .scale = 1 }); + try submit(cb, enc); + const selected = @as([*]i32, @ptrCast(@alignCast(scratch.positions().buffer.contents() + + scratch.positions().offset)))[0 .. 3 * 512]; + for (selected[0..512]) |pos| try std.testing.expectEqual(@as(i32, -1), pos); + for (1..3) |r| { + try std.testing.expectEqual(@as(i32, 0), selected[r * 512]); + for (selected[r * 512 + 1 .. (r + 1) * 512]) |pos| try std.testing.expectEqual(@as(i32, -1), pos); + } +} + +test "YaRN interpolates its correction band and inverse rotation conjugates the frequency" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try rope.Pipelines.load(device, 65537, 32, .{ .theta = 10000 }, .{ .theta = 160000, .original = 65536 }); + defer p.deinit(); + const x = try device.buffer(64 * 2, mtl.ResourceOptions.shared); + defer x.deinit(); + @memset(x.slice(u16, 64), bf(1)); + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + try p.encode(enc, view(x), .{ .rows = 1, .heads = 1, .width = 64, .rotary = 32, .stride = 64, .position = 65536, .inverse = 1, .spectrum = .compressed }); + try submit(cb, enc); + const low: f32 = @floatCast(@max(@floor(32 * @log(@as(f64, 65536) / (32 * 2 * std.math.pi)) / (2 * @log(@as(f64, 160000)))), 0)); + const high: f32 = @floatCast(@min(@ceil(32 * @log(@as(f64, 65536) / (2 * std.math.pi)) / (2 * @log(@as(f64, 160000)))), 31)); + for ([_]usize{ 0, 10, 15 }) |pair| { + const power: f32 = @floatCast(std.math.pow(f64, 160000, @as(f32, @floatFromInt(pair * 2)) / 32)); + const frequency: f32 = 1 / power; + const smooth: f32 = 1 - std.math.clamp((@as(f32, @floatFromInt(pair)) - low) / @max(high - low, 1e-3), 0, 1); + const slow: f32 = frequency / 16 * (1 - smooth); + const fast: f32 = frequency * smooth; + const angle: f32 = 65536 * (slow + fast); + const c: f32 = @floatCast(@cos(@as(f64, angle))); + const s: f32 = -@as(f32, @floatCast(@sin(@as(f64, angle)))); + try std.testing.expectEqual(bf(c - s), x.slice(u16, 64)[32 + pair * 2]); + try std.testing.expectEqual(bf(s + c), x.slice(u16, 64)[33 + pair * 2]); + } +} + +test "RoPE retained source pair rounds the position-five imaginary consumer to pinned CUDA BF16" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try rope.Pipelines.load(device, 6, 64, .{ .theta = 10000 }, .{ .theta = 160000, .original = 65536 }); + defer p.deinit(); + const x = try device.buffer(512 * 2, mtl.ResourceOptions.shared); + defer x.deinit(); + @memset(x.slice(u16, 512), bf(0)); + // Token-5/layer-0 rotary pair 21 crosses a BF16 tie under per-head cos; pinned source factors give .546875. + x.slice(u16, 512)[490] = bf(1.15625); + x.slice(u16, 512)[491] = bf(0.53125); + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + try p.encode(enc, view(x), .{ .rows = 1, .heads = 1, .width = 512, .rotary = 64, .stride = 512, .position = 5 }); + try submit(cb, enc); + try std.testing.expectEqual(bf(0.546875), x.slice(u16, 512)[491]); +} diff --git a/zig/src/families/deepseek_v41/engram.zig b/zig/src/families/deepseek_v41/engram.zig new file mode 100644 index 0000000000..867cc78d39 --- /dev/null +++ b/zig/src/families/deepseek_v41/engram.zig @@ -0,0 +1,207 @@ +//! Resident Engram FP8 gathers and signed-sqrt gates; WKV is projected separately by the caller. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +pub const View = @import("view.zig").View; +const Dtype = @import("ops.zig").Dtype; +pub const hash = @import("engram_hash.zig"); + +pub const head_dim = 256; +pub const gather_dim = hash.columns * head_dim; +pub const history_bytes = 16; + +/// Borrowed full-resident E4M3FN rows and UE8M0 rowwise scale blocks, never a 32x32 scale matrix. +pub const Table = struct { + weight: View, + scale: View, + rows: u32, + weight_stride_bytes: usize = head_dim, + scale_stride_bytes: usize = head_dim / 32, + scale_block: u32 = 32, + + pub fn validate(self: Table) !void { + if (self.rows == 0 or self.weight_stride_bytes < head_dim or self.scale_stride_bytes < head_dim / 32 or self.scale_block != 32) return error.InvalidEngramTable; + try self.weight.requireRows(self.rows, self.weight_stride_bytes, head_dim, 1); + try self.scale.requireRows(self.rows, self.scale_stride_bytes, head_dim / 32, 1); + } +}; + +/// BF16 hidden [tokens,copies,dim], WKV [tokens,(copies+1)*dim]; norm weights contiguous FP32 [copies,dim]. +pub const Gate = struct { + q_weight: View, + k_weight: View, + /// The released checkpoints store these f32; rank-split repacks store bf16. + weight_dtype: Dtype = .f32, + dim: u32, + copies: u32, + hidden_stride: usize, + kv_stride: usize, + output_stride: usize, + eps: f32, + + pub fn validate(self: Gate) !void { + const width = try std.math.mul(usize, self.dim, self.copies); + const projected_width = try std.math.add(usize, width, self.dim); + if (self.dim == 0 or self.copies == 0 or self.hidden_stride < width or self.output_stride < width or self.kv_stride < projected_width or !std.math.isFinite(self.eps) or self.eps <= 0 or + self.dim > std.math.maxInt(u32) - 255) return error.InvalidEngramGate; + const weight_bytes = try std.math.mul(usize, width, 4); + try self.q_weight.require(weight_bytes, 4); + try self.k_weight.require(weight_bytes, 4); + } +}; + +pub const Pipelines = struct { + hash_ids: mtl.Pipeline, + commit: mtl.Pipeline, + reset: mtl.Pipeline, + gather: mtl.Pipeline, + gate: mtl.Pipeline, + + pub fn load(device: mtl.Device) !Pipelines { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const lib = try mtl.Library.fromSource(device, sources.deepseek_v41_engram, mtl.CompileOptions.mlx()); + defer lib.deinit(); + const h = try mtl.Pipeline.init(device, lib, "tf_v41_engram_hash", false); + errdefer h.deinit(); + const c = try mtl.Pipeline.init(device, lib, "tf_v41_engram_commit", false); + errdefer c.deinit(); + const r = try mtl.Pipeline.init(device, lib, "tf_v41_engram_reset", false); + errdefer r.deinit(); + const g = try mtl.Pipeline.init(device, lib, "tf_v41_engram_gather", false); + errdefer g.deinit(); + const gate = try mtl.Pipeline.init(device, lib, "tf_v41_engram_gate", false); + return .{ .hash_ids = h, .commit = c, .reset = r, .gather = g, .gate = gate }; + } + + pub fn deinit(self: *Pipelines) void { + self.hash_ids.deinit(); + self.commit.deinit(); + self.reset.deinit(); + self.gather.deinit(); + self.gate.deinit(); + } +}; + +const HashArgs = extern struct { tokens: u32, pad: u32, masked: u32 }; +const GatherArgs = extern struct { + weight_stride: u64, + scale_stride: u64, + output_stride: u64, + tokens: u32, + columns: u32, + head_dim: u32, + scale_block: u32, +}; +const GateArgs = extern struct { + hidden_stride: u64, + kv_stride: u64, + output_stride: u64, + dim: u32, + copies: u32, + masked: u32, + weight_dtype: u32, + eps: f32, +}; + +fn bind(e: mtl.ComputeEncoder, view: View, at: usize) void { + e.setBuffer(view.buffer, view.offset, at); +} + +/// Each history View is one sequence's 16-byte GPU record; initialize before its first hash. +pub fn encodeReset(e: mtl.ComputeEncoder, p: *const Pipelines, history: View) void { + e.setPipeline(p.reset); + bind(e, history, 0); + e.dispatchGroups(.{ .width = 1 }, .{ .width = 1 }); + e.barrier(); +} + +/// Snapshot/rollback are GPU copies; caller ends compute encoding before beginning a blit encoder. +pub fn encodeSnapshot(e: mtl.BlitEncoder, history: View, saved: View) void { + e.copy(history.buffer, history.offset, saved.buffer, saved.offset, history_bytes); +} + +pub fn encodeRollback(e: mtl.BlitEncoder, saved: View, history: View) void { + encodeSnapshot(e, saved, history); +} + +/// Raw u32 IDs -> [tokens,24] hashes; map covers valid IDs, optional 0=DEAD mask, history is unchanged. +pub fn encodeHash(e: mtl.ComputeEncoder, p: *const Pipelines, meta: hash.Layer, ids: View, token_map: View, mask: ?View, history: View, out: View, tokens: u32) !void { + if (tokens == 0) return error.EmptyEngramChunk; + if (!std.mem.eql(u8, std.mem.asBytes(&meta), std.mem.asBytes(&hash.layers[0])) and + !std.mem.eql(u8, std.mem.asBytes(&meta), std.mem.asBytes(&hash.layers[1]))) return error.InvalidEngramMetadata; + try ids.requireRows(tokens, 1, 1, 4); + try token_map.require(4, 4); + try history.require(history_bytes, 4); + try out.requireRows(tokens, hash.columns, hash.columns, 4); + if (mask) |m| try m.require(tokens, 1); + e.setPipeline(p.hash_ids); + bind(e, ids, 0); + bind(e, token_map, 1); + bind(e, mask orelse ids, 2); + bind(e, history, 3); + bind(e, out, 4); + e.setValue(meta, 5); + e.setValue(HashArgs{ .tokens = tokens, .pad = hash.pad_id, .masked = @intFromBool(mask != null) }, 6); + e.dispatchThreads(.{ .width = tokens }, .{ .width = 64 }); + e.barrier(); +} + +/// Commit exactly once per chunk after both layers have hashed; remains bounded at three past tokens. +pub fn encodeCommit(e: mtl.ComputeEncoder, p: *const Pipelines, ids: View, token_map: View, mask: ?View, history: View, tokens: u32) !void { + if (tokens == 0) return error.EmptyEngramChunk; + try ids.requireRows(tokens, 1, 1, 4); + try token_map.require(4, 4); + try history.require(history_bytes, 4); + if (mask) |m| try m.require(tokens, 1); + e.barrier(); + e.setPipeline(p.commit); + bind(e, ids, 0); + bind(e, token_map, 1); + bind(e, mask orelse ids, 2); + bind(e, history, 3); + e.setValue(HashArgs{ .tokens = tokens, .pad = hash.pad_id, .masked = @intFromBool(mask != null) }, 4); + e.dispatchGroups(.{ .width = 1 }, .{ .width = 1 }); + e.barrier(); +} + +/// Gather uint32 [tokens,24] row IDs into BF16 [tokens,6144]; IDs must be within the borrowed table. +pub fn encodeGather(e: mtl.ComputeEncoder, p: *const Pipelines, table: Table, ids: View, out: View, tokens: u32, output_stride: usize) !void { + try table.validate(); + if (tokens == 0 or output_stride < gather_dim) return error.InvalidEngramGather; + try ids.requireRows(tokens, hash.columns, hash.columns, 4); + try out.requireRows(tokens, output_stride, gather_dim, 2); + e.setPipeline(p.gather); + bind(e, ids, 0); + bind(e, table.weight, 1); + bind(e, table.scale, 2); + bind(e, out, 3); + e.setValue(GatherArgs{ .weight_stride = table.weight_stride_bytes, .scale_stride = table.scale_stride_bytes, .output_stride = output_stride, .tokens = tokens, .columns = hash.columns, .head_dim = head_dim, .scale_block = table.scale_block }, 4); + e.dispatchThreads(.{ .width = gather_dim, .height = tokens }, .{ .width = 256 }); + e.barrier(); +} + +/// Add per-copy gate times shared value to BF16 hidden; out may alias hidden, never the projected KV. +pub fn encodeGate(e: mtl.ComputeEncoder, p: *const Pipelines, gate: Gate, hidden: View, projected_kv: View, mask: ?View, out: View, tokens: u32) !void { + try gate.validate(); + if (tokens == 0) return error.EmptyEngramChunk; + const width = @as(usize, gate.dim) * gate.copies; + try hidden.requireRows(tokens, gate.hidden_stride, width, 2); + try projected_kv.requireRows(tokens, gate.kv_stride, width + gate.dim, 2); + try out.requireRows(tokens, gate.output_stride, width, 2); + if (mask) |m| try m.require(tokens, 1); + e.setPipeline(p.gate); + bind(e, hidden, 0); + bind(e, projected_kv, 1); + bind(e, gate.q_weight, 2); + bind(e, gate.k_weight, 3); + bind(e, mask orelse hidden, 4); + bind(e, out, 5); + e.setValue(GateArgs{ .hidden_stride = gate.hidden_stride, .kv_stride = gate.kv_stride, .output_stride = gate.output_stride, .dim = gate.dim, .copies = gate.copies, .masked = @intFromBool(mask != null), .weight_dtype = @backingInt(gate.weight_dtype), .eps = gate.eps }, 6); + e.dispatchGroups(.{ .width = tokens, .height = gate.copies }, .{ .width = 256 }); + e.barrier(); +} + +test { + _ = @import("engram_test.zig"); +} diff --git a/zig/src/families/deepseek_v41/engram_test.zig b/zig/src/families/deepseek_v41/engram_test.zig new file mode 100644 index 0000000000..dd085a8744 --- /dev/null +++ b/zig/src/families/deepseek_v41/engram_test.zig @@ -0,0 +1,247 @@ +//! Behavioral hash, rollback, selected-row dequantization and per-copy signed-sqrt gate regressions. +const std = @import("std"); +const hash = @import("engram_hash.zig"); +const engram = @import("engram.zig"); +const mtl = @import("metal"); +const View = engram.View; + +// Generated offline with the official NumPy multipliers and prime-bucket algorithm, not the native code. +const first = [_]u32{ + 8939594, 28321402, 40265357, 55817216, 71147020, 86924254, 110813108, 121427189, + 128454377, 156247504, 165979524, 188168291, 198564788, 211763979, 239563936, 241631034, + 263018334, 286373017, 297618161, 319241270, 333656732, 339279600, 364226324, 373889199, +}; +const third = [_]u32{ + 13176822, 30752650, 42532473, 56771564, 70133176, 85254602, 108815600, 115334181, + 138863141, 149049926, 167314083, 191090346, 204680718, 211476972, 237672992, 245604111, + 269210034, 282954889, 298312790, 313992267, 329778842, 345458689, 359433926, 374584261, +}; +const after_dead = [_]u32{ + 762556, 18016262, 38494153, 58317688, 64056565, 81970637, 106927994, 122902996, + 129060773, 153088767, 160206490, 178949865, 205667218, 211027556, 235070602, 251994813, + 263310652, 278975004, 292595948, 315407209, 325281398, 348093105, 353903080, 377405725, +}; + +test "official hash IDs preserve order, padding and globally unique layer buckets" { + var history: hash.History = .{}; + const a = try history.push(7, hash.layers[0]); + try std.testing.expectEqualSlices(u32, &first, &a); + _ = try history.push(9, hash.layers[0]); + const c = try history.push(11, hash.layers[0]); + try std.testing.expectEqualSlices(u32, &third, &c); + for (hash.layers) |meta| for (meta.primes, meta.offsets) |p, offset| { + try std.testing.expect(p >= 16000000); + try std.testing.expect(offset + p <= meta.rows); + }; + var high: hash.History = .{ .tail = .{ 1234, 5678, 42 }, .position = 3 }; + const result = try high.push(99091, hash.layers[1]); + const expected = [_]u32{ + 3315102, 30010869, 33829090, 51274514, 77267621, 84401643, 106522093, 127296527, + 137402994, 151593930, 162661867, 180852958, 195297507, 213361283, 228864003, 255457679, + 262256667, 282063670, 300635716, 304966933, 321424067, 350721784, 358839309, 370654406, + }; + try std.testing.expectEqualSlices(u32, &expected, &result); +} + +test "dead image positions block all older lookbacks and snapshot rollback restores the exact boundary" { + var history: hash.History = .{}; + _ = try history.push(7, hash.layers[0]); + _ = try history.push(9, hash.layers[0]); + const saved = history.snapshot(); + _ = try history.push(hash.dead, hash.layers[0]); + const got = try history.push(13, hash.layers[0]); + try std.testing.expectEqualSlices(u32, &after_dead, &got); + history.rollback(saved); + const replay = try history.push(11, hash.layers[0]); + try std.testing.expectEqualSlices(u32, &third, &replay); + history.reset(); + const reset = try history.push(7, hash.layers[0]); + try std.testing.expectEqualSlices(u32, &first, &reset); + try std.testing.expectError(error.InvalidEngramCompressedId, history.push(-2, hash.layers[0])); + try std.testing.expectEqual(@as(u64, 1), history.position); +} + +test "all prefill/decode partitions preserve hashes including a dead token at the chunk boundary" { + var bytes: [20 * 4]u8 = undefined; + for (0..20) |i| std.mem.writeInt(u32, bytes[i * 4 ..][0..4], @intCast(i), .little); + const map = hash.TokenMap{ .bytes = &bytes }; + const ids = [_]u32{ 7, 9, 11, 3, 13, 15, 17 }; + const mask = [_]u8{ 1, 1, 1, 0, 1, 1, 1 }; + for (hash.layers) |meta| { + var whole: hash.History = .{}; + var expected: [ids.len][hash.columns]u32 = undefined; + try hash.hashChunk(&whole, map, meta, &ids, &mask, &expected); + for (0..ids.len + 1) |split| { + var chunked: hash.History = .{}; + var actual: [ids.len][hash.columns]u32 = undefined; + try hash.hashChunk(&chunked, map, meta, ids[0..split], mask[0..split], actual[0..split]); + try hash.hashChunk(&chunked, map, meta, ids[split..], mask[split..], actual[split..]); + try std.testing.expectEqualDeep(expected, actual); + try std.testing.expectEqualDeep(whole, chunked); + } + const saved = whole.snapshot(); + var invalid: [2][hash.columns]u32 = undefined; + try std.testing.expectError(error.InvalidEngramToken, hash.hashChunk(&whole, map, meta, &.{ 1, 20 }, null, &invalid)); + try std.testing.expectEqualDeep(saved, whole); + } +} + +test "raw tokenizer IDs are compressed before hashing and normalization-equivalent tokens share hashes" { + var bytes: [4 * 4]u8 = undefined; + for ([_]u32{ 7, 9, 11, 9 }, 0..) |id, i| std.mem.writeInt(u32, bytes[i * 4 ..][0..4], id, .little); + const map = hash.TokenMap{ .bytes = &bytes }; + var a: hash.History = .{}; + var b: hash.History = .{}; + var out_a: [3][hash.columns]u32 = undefined; + var out_b: [3][hash.columns]u32 = undefined; + try hash.hashChunk(&a, map, hash.layers[0], &.{ 0, 1, 2 }, null, &out_a); + try hash.hashChunk(&b, map, hash.layers[0], &.{ 0, 3, 2 }, null, &out_b); + try std.testing.expectEqualDeep(out_a, out_b); + try std.testing.expectEqualSlices(u32, &first, &out_a[0]); + try std.testing.expectEqualSlices(u32, &third, &out_a[2]); +} + +fn view(buffer: mtl.Buffer) View { + return .{ .buffer = buffer }; +} +fn bf(value: f32) u16 { + const bits: u32 = @bitCast(value); + return @truncate((bits + 0x7fff + ((bits >> 16) & 1)) >> 16); +} +fn float(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} +fn completed(cb: mtl.CommandBuffer) !void { + cb.commit(); + cb.wait(); + if (cb.failure() != null) return error.EngramGPUFailure; +} + +test "Metal hashes chunks with image breaks and rolls back bounded GPU history" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + var p = try engram.Pipelines.load(device); + defer p.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + const ids = try device.buffer(7 * 4, 0); + defer ids.deinit(); + @memcpy(ids.slice(u32, 7), &[_]u32{ 7, 9, 11, 3, 13, 15, 17 }); + const map = try device.buffer(20 * 4, 0); + defer map.deinit(); + for (map.slice(u32, 20), 0..) |*id, i| id.* = @intCast(i); + const mask = try device.buffer(7, 0); + defer mask.deinit(); + @memcpy(mask.slice(u8, 7), &[_]u8{ 1, 1, 1, 0, 1, 1, 1 }); + const history = try device.buffer(engram.history_bytes, 0); + defer history.deinit(); + const saved = try device.buffer(engram.history_bytes, 0); + defer saved.deinit(); + const out = try device.buffer(7 * 24 * 4, 0); + defer out.deinit(); + const replay = try device.buffer(24 * 4, 0); + defer replay.deinit(); + const cb = queue.commandBuffer(); + var e = cb.compute(.serial); + engram.encodeReset(e, &p, view(history)); + try engram.encodeHash(e, &p, hash.layers[0], view(ids), view(map), view(mask), view(history), view(out), 3); + try engram.encodeCommit(e, &p, view(ids), view(map), view(mask), view(history), 3); + e.end(); + var blit = cb.blit(); + engram.encodeSnapshot(blit, view(history), view(saved)); + blit.end(); + e = cb.compute(.serial); + try engram.encodeHash(e, &p, hash.layers[0], view(ids).shifted(12), view(map), view(mask).shifted(3), view(history), view(out).shifted(3 * 24 * 4), 4); + try engram.encodeCommit(e, &p, view(ids).shifted(12), view(map), view(mask).shifted(3), view(history), 4); + e.end(); + blit = cb.blit(); + engram.encodeRollback(blit, view(saved), view(history)); + blit.end(); + e = cb.compute(.serial); + try engram.encodeHash(e, &p, hash.layers[1], view(ids).shifted(16), view(map), null, view(history), view(replay), 1); + e.end(); + try completed(cb); + var cpu: hash.History = .{}; + for (ids.slice(u32, 7), mask.slice(u8, 7), 0..) |id, alive, i| { + const want = try cpu.push(if (alive == 0) hash.dead else @intCast(id), hash.layers[0]); + try std.testing.expectEqualSlices(u32, &want, out.slice(u32, 7 * 24)[i * 24 ..][0..24]); + } + cpu = .{ .tail = .{ 11, 9, 7 }, .position = 3 }; + const want = try cpu.push(13, hash.layers[1]); + try std.testing.expectEqualSlices(u32, &want, replay.slice(u32, 24)); + try std.testing.expectEqualSlices(i32, &.{ 11, 9, 7 }, history.slice(i32, 4)[0..3]); +} + +test "Metal gathers only selected FP8 rows with rowwise scales and gates each mHC copy independently" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + var p = try engram.Pipelines.load(device); + defer p.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + const table = try device.buffer(3 * 272, 0); + defer table.deinit(); + @memset(table.slice(u8, 3 * 272), 0x7f); + const scales = try device.buffer(3 * 12, 0); + defer scales.deinit(); + @memset(scales.slice(u8, 3 * 12), 255); + for (0..3) |r| { + @memset(table.slice(u8, 3 * 272)[r * 272 ..][0..256], if (r == 0) 0x38 else if (r == 1) 0xb8 else 0x7e); + for (0..8) |g| scales.slice(u8, 3 * 12)[r * 12 + g] = @intCast(127 + g); + } + table.slice(u8, 3 * 272)[2 * 272] = 0x01; + table.slice(u8, 3 * 272)[2 * 272 + 1] = 0x80; + scales.slice(u8, 3 * 12)[12 + 7] = 255; + const ids = try device.buffer(24 * 4, 0); + defer ids.deinit(); + for (ids.slice(u32, 24), 0..) |*id, i| id.* = @intCast((i * 2 + 1) % 3); + const gathered = try device.buffer((engram.gather_dim + 16) * 2, 0); + defer gathered.deinit(); + const hidden = try device.buffer(2 * 4 * 2, 0); + defer hidden.deinit(); + for (hidden.slice(u16, 8), 0..) |*v, i| v.* = bf(if (i < 4) 1 else 2); + const kv = try device.buffer(12 * 2, 0); + defer kv.deinit(); + for (kv.slice(u16, 12), 0..) |*v, i| v.* = bf(if (i < 4) 1 else if (i < 8) -2 else 3); + const q = try device.buffer(8 * 4, 0); + defer q.deinit(); + for (q.slice(f32, 8), 0..) |*v, i| v.* = if (i < 4) 1 else 0.5; + const k = try device.buffer(8 * 4, 0); + defer k.deinit(); + for (k.slice(f32, 8), 0..) |*v, i| v.* = if (i < 4) 1 else 2; + const output = try device.buffer(8 * 2, 0); + defer output.deinit(); + const dead_mask = try device.buffer(1, 0); + defer dead_mask.deinit(); + dead_mask.slice(u8, 1)[0] = 0; + const masked_output = try device.buffer(8 * 2, 0); + defer masked_output.deinit(); + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + try engram.encodeGather(e, &p, .{ .weight = view(table), .scale = view(scales), .rows = 3, .weight_stride_bytes = 272, .scale_stride_bytes = 12 }, view(ids), view(gathered), 1, engram.gather_dim + 16); + try engram.encodeGate(e, &p, .{ .q_weight = view(q), .k_weight = view(k), .dim = 4, .copies = 2, .hidden_stride = 8, .kv_stride = 12, .output_stride = 8, .eps = 1e-20 }, view(hidden), view(kv), null, view(output), 1); + try engram.encodeGate(e, &p, .{ .q_weight = view(q), .k_weight = view(k), .dim = 4, .copies = 2, .hidden_stride = 8, .kv_stride = 12, .output_stride = 8, .eps = 1e-20 }, view(hidden), view(kv), view(dead_mask), view(masked_output), 1); + e.end(); + try completed(cb); + for (ids.slice(u32, 24), 0..) |r, col| for (0..256) |d| { + const got = gathered.slice(u16, engram.gather_dim)[col * 256 + d]; + if (r == 1 and d >= 224) { + try std.testing.expect(std.math.isNan(float(got))); + continue; + } + const base: f32 = if (r == 2 and d == 0) 0x1p-9 else if (r == 2 and d == 1) -0.0 else if (r == 0) 1 else if (r == 1) -1 else 448; + const scale: f32 = @floatFromInt(@as(u32, 1) << @as(u5, @intCast(d / 32))); + try std.testing.expectEqual(bf(base * scale), got); + }; + const gate_positive = 1.0 / (1.0 + @exp(-@sqrt(@as(f32, 2)))); + const gate_negative = 1.0 / (1.0 + @exp(@sqrt(@as(f32, 2)))); + for (output.slice(u16, 8), 0..) |v, i| { + const want = bf(if (i < 4) 1 + 3 * gate_positive else 2 + 3 * gate_negative); + try std.testing.expectApproxEqAbs(float(want), float(v), 0.015625); + } + try std.testing.expectEqualSlices(u16, hidden.slice(u16, 8), masked_output.slice(u16, 8)); +} diff --git a/zig/src/families/deepseek_v41/fixtures/hc_attention_rms.json b/zig/src/families/deepseek_v41/fixtures/hc_attention_rms.json new file mode 100644 index 0000000000..d1f1c2b61d --- /dev/null +++ b/zig/src/families/deepseek_v41/fixtures/hc_attention_rms.json @@ -0,0 +1,25 @@ +{ + "embedding_bf16": "WT1bPAO857tZvTS9Cr2HvDY9MT2nvUI867vvPIG9Xr3wPMs8GLwLvCG91b1CPf87mTzuvC69rTzOvM28x7w+PSC6VL0YPOG8Fb1TvRE84zwNuw48pTwrvF89OTxEveU8SDybPBA9Yz0XPKU8Nj33PJK8jjtTvNW8C731PDk8Rz0SvAK9l7xnPB69ojvNuxM9y7yMvDC97Dx3PA89Ib2TO0k8j7orvDc7xbwwvCM9lLydvDS9OD1JPQC9JL2yvBu8ebyCvUE9Qz3+vDm8Bb0XPGG8gryUuAi8ZT1vvN+8wbyQPI47orz/vCE93jzVPHU90DwgvVe8rLzPvDe9Xr1tPZK8ojtCvcG84byJOzq9QD2APGE9Xzu3OwY90LzkPNY9iLsZPao8kz2Ivbs85LxKvZu9AjuiPQO9BL6RvAk9hL37u+08gjoavBw97rxDPXq5kz3vO5e7arp3vX+9hj2dO6m8B7sBPHS9jDywPGI907zYvAy8P7y9PDw8u7zrO1O8rjzLvD28Yz3yu5u8lLz0vCa9kjxuPE+9ED1XvJW7RbwfPDe9Mz1yvRy9Br0CPNc8OL2zPLI8tjwdPQw8Fb03vIM8kTzrvPU8DTwgvdM8YbwJPKY877vVPLo8Y7x6PAw9IbyOPEu9+TxuPfo8E7zVPO48ybzNPOe8EbtlvYi8Kb3VvCI8Wz2rvb88Wz1HvYe9oLzFPMY8pDyCPXO8k7yQPQq9IrxbvS89JjwKvQo7Jb0QvVa9oDxWPLM8rrgMvI46mjwWPdi8nzwevUG9Dz3eumy8K70bPZS9Jz2/vIO9xrxLvZi9W73GOsY9mL2OvCe8BL1DvOE8D70tvbU9mzxPvT082zwnvCy9Ar0+O1s9qTwbvIY9Sb2qvTY9KzvvvCA93TxAPeG83TtbPN07s7vXu2M8RjzyvAe9wb1fvUk9YrwwvE+8nbywPVG8WzsDPZI9XD2RvG09+DybPAo81zx8PdI8kr0JvJI7Xj1GPTS9irycPFS8gTzcvES9wLwvPY+9Bju2vRs9IDziPF09GD3dvHA7ML0MvEa9nroCPbo8yrtcvZg7QT3WvC47aLxmPKa7zbwVvGq7Dz30Oyc9Bz2FPBo9Ob14vVW8ELsDPUc8qL2kvCA8DbzlOz49EbwsvS8907wQO5w6sDzcuo085jt3Pdu6Hr0UPBm9OLzVvOq8LT0wO8y7Hz2/vFa9rLwePSC91LyXPa87rr0xveK8jDuMu8C54rx6uzU7HT3AvLE8Jj37O047/zz1PB+7Ab0QPDa9Oz2JvZW8Lb3guwQ9njtivey8vTyhOuE8Brx5vXC80jxdOw08i70OvcS8MD2qvbu8XbyZPLm9mz1aPV66iTx3vMQ7oTyyPIw9PbuOuzM9HjzgO4c7ubzivS69PDzTvLi8H712u767ZL2dPFk9gTzvvK48BTwOvYA9TbtlPIy8Vj0wPSu8D7x2vZQ88ruvPLy8tjwyPSk99rycOxk91DzgvL87pjyGPUi8Uz1ivBS9hb28OjM9izxevFq7Rb39vK07fL2QvcU8uD1vvPq6TLz2PCk9hLzhvIC7gzyBO2Q65DySux6967zbvA09Iz2cPVe8rrwmvbC8djwtPVe8Pz1qvUi9e71ZvYU9R7okPS49jbpIvaY8QbwiPee8rLwKvbE6WT2MPFa8YbxJvLo7Hr0iPVM94DzMu4s8Ojw5PRs9N730O568Fz27vFK8tL2uvEa8irrtO0K89TtAPVk8vbgmvTk9pzqDvLs7KTxtO6W8GjoWPCc9DjxBPZi9oryoPJS8O72jvBu9Fz3lPKA8bj1lPAS9UD0iPQ48KbxFPJy9J703Pbi8/rwqPVs7Er3Iu3q9qDsvvR28VrxzPVE8Y72UPB+9WbsfvS+8KL2BOu08B71ZPVK6Nbo0PZg9iT1pPYE8RD0IPB89Mr3NPEG8oz2wvLU8QjzavLK7qLwgve88ybxCPV68obxbvbu8f71xPGa87LwRPa08Pj1vvG89mz1WPdk8E711O5C89zwXPOy8hDzxPAc9ELwqvGI9I70SO7Q7Ab3AOzg8abtfvTe8zbzAvLS8ibyVPTu9tTyYPKy857wnPZi9JDx5uz07E70MPVq8ZLysOy67BT00Pda8bTx0PN082T0nvdS8Lztyuzm95DyoO1I9vLv5PBO9FD2kPWQ9kLuguQY8A73LPD27pryUvPq8Rb2Pvf68kjxzvGS9+LyDvOC8TDyRPAG9HDy0vLW89DyGPQg78LqMOza9ULt+POM44zzRvIU9xTyLPQc8zbxtvaY9bLuoPKw79zvMO7U8h72pPAa9ML0BvUK907xDvZC9Fbx3PfA8VDxsvW88Bz0TvDK91DwrPIk9ab2uu0a8iTtBvck8sjz7PCQ9D70JvQe94DrHOhG7hLzvvJ894DwxO2q7mb1jPDY8qTzPvEM9N7zZPHE9IboPPUM84zxbO0y9q7ylOh68JD3/u5c7ujwKPIu7yTzLvOO8Ej0YPdE9r726PR28C732PBW8wztrvZY8eLyVPLU8T7xhPea9i7qvPES7qbsvPI08jj2uPOM8T71bvOC8ijwrPY28G707PeO8JD0BvC85sbwPPA491bzNO5Y6J7wbPX88n71hPA+9wbyzPBi7kj2JPM+7j7sROgG8Zj2nPJ88uzwGvJE9j70PPay7yLz+vAC9Sj19vYe8hz3OvBW9DrswvYM8Eb2dvGy9/TuKPeI7JTwIvUk9mLuLvI68rrx2vUg9izxfPf08ez33vB+8L70lvdk8Sb0lvYY8Pry0POk7yDy/PCA8cDwNPQY8dT2DvYu6FT2dPWy8C7wcPUw9nzxhvKA9t7sWPci8Tr25vL47Vb02vBQ9j71vva+8tzyNPDe9lrm3PAK82jq1O2M7WbuXvCw9krsmvag7tDoIPXM86zzJPAU9aj2bvRG8uD3SOqa7XrxNuy09gDvHPF+93DwxPAu9Ej1FPU69hTuyPOu8L70avQ89zTqzvUk9tjvvPBw7u7zEO7k8Ez0hO0U7NL01PQM83zzJO3y8oDz8PO0797yOPIG9nLx9PLs8sbx2PZa7kT2NPCe8k7vUPJU6Gb0fPL48cDzDOzQ8RLxKPU89Wr36PJY8rDy+uya7E7wMPau6TryIO6O9kDyyvc87Er3hvAg8BLzWPEi8yzxqO5I9vLwWvRs9sryUPCa8trw6PA09EDuNvaa7QDxXvSq8U71fvHQ9M71QvU88TLuYvV89xrxFPXS8/jxAu2E9cLyZPL29+DvHPOq8QL0XvCi9TT0JPX+8prxPPSe9wr7FvG09OD2mPJM8g72NvLS80jytPLo8F73cOz+9CD2LO1Q9Mrx1vA+9WT2UO4K9MT1uPIk96L1Qu9i8cL1oPU69v7rePNg7Kz2ePLA7Dj0oPEU9Mj3/vDY9hL2SO488J7sGvIE8Kz1Xuxs9KDwCPpa8O725vGA8SrxAvSI8A73HvF47n7x7PEA9ejyWPUw8NT0TPXo9DT2ePGu9yjzXu4I9tTsAPee97rsEvSO9GDw/PEq96DvOPQS8Yz1dvSw9cz2vOxq71boZvVM8AD1vvYQ5UT08POm8zjuEvZy8JT0Puw466LyfPfg87zqFvMw6Ujzqulm6fz2muwM9t72pPc87+LzIPci8xLzzO2S95LtoPXu9bDx0vBG8/TzPPHC97zzWPDe9k7wXPJQ78rwavQW9gj2NOFw8/rxgvX28gjwRvei8F7ywvDW9z7s0PUE90rsePV2987yZO8o8Rj2TO8e7mT3oPM47zDs3vBS90jykPBa9zTyNvOM8kTwTvT87Wjz2PLi8rDs8ugS9CT11u3e8Wr0rvNs7ITz3vA09T73bvIu8qjyIvYC9oz21vIC9Z7uAPEs9Gr0MvYe84ztXPeq86TykvAW9FD0uvII8ELxAvK67Sj3GPIe9HrxcPAW9yjxdPSc9RL3kvC48gbysPBg9BL0SvUE83jzCPMI7Cryyu448+rr0PEM8Wj3mvAY8mLyGuyw8Mjp4vcS9AjwlPXw8Rz1BPq89HT2yu9i8Zbz4Ois9S701Oym9FrwIvaI9lz1nPMO8nb0OutA7rT2IPFa96ryrvTG9e73xvKS8ETzsvGs95zuzPC28dz2dPIc9/DxPvGG8tDw+PGo8h70HPCq93TwqPJG6erznO8C8qbxfPBK8iD0EvWw68bwpO2a9lbvSPHE73DoEPGo9YD3vOzY9Q7ouPCo9I7wjPeu65TysvBA7gD2VO5S9BD07vbo8hL2QPG28AD0vvWG9JL3EPFw9q72guw68qzwRPRy9rDybPPG8BL3svGK9nj17PWi7ETwEO+s7Czz8u/E80DyPvFA9fLyHPEI9vT1ovJe7Gb0aPRQ8br0IPVe8Wr04u8c7BTw8vT+9Kb3ZPJW9+DxQPFC8UrzBPDc9Jb2Dvbq8Sz2zPIa9aT17POO8lbzEPL88Ar66O4U8Zj0Xu5W8Vz0NPSQ9tr2dvIM8PL0VvMs8n7wBvXq9sTuCPbu9ezwBvQO9Ubw4PSA9ITyruwm70bwSPUA8hDwPvZ89F73tvOa7Ij3OuZU8AT1svb28P7xIOzu8jz1oPQo9HDwePcG8ozy+vSm9vjyJPEW9ezyOPO28tTzLu8E8Er1mvdK8TD25vCO9fz1VPJu9BbxxvGc8rbwJvYO8bTzKPIo9Rb0yPJo82LzAu9I8Bz1XvGI9RrwMvSE9IjvAvA08FzoUvbQ8PLxbvfm8Zz0zvIK7zTv/PF88Mz1PvJa9IT3JPKW8Sb0QvQW8cjyjPLs737z5O4+807w0vQs9ZD0LvUk9Gj20PPe8I7suvOm8l7xVvca8KrtzvQ89QzyHPWU8ejzUPA+9Aj2aPBY9ebwmPXy7bj1mPS09Z7sLPZ+8KL28PCS9z7xQvW47QTw3vdu7/DvhvPq8X7ySvTq6O72nPUY80LwJvR496TxXPIU7LrywPfq87zz5uwW9/DsEvUI8Qr2wPPE7pz1SvbS8Xjx5vTe8Qj0LPcM9JjwzPZa8HT0zvBw9KjyIPZQ9WLwTPIg8cTutvHm8+Lweveo8k72ivOM8Sz17vYO9VD0Zvek8wjzAvdg8cr0IPFo9mLygPZe9GTz+OYc9mD1NvBs97TyHvTq96bwovXE8zTtIPK08qz2XPCw8MT2/O0m9ezwEPE07lT1uPU29s7zHPAO9CD3Ru2g8K71yPUA8nD3Ku0I8YT0ePU48TbzYu1s8sDwtPY29BD0UPPI74joFPUs9rr0uPD88Cb3CvOa8rjx1vVe8xbzlO6S8yTyBPJq7YLpOvLe8IbyzPDW9C7vkO8s8izwtvQQ9gbwbvIW8HT5PPXo8ID0oPAu8dbstvRq7pzxkvMM87jymPcE6BDyOvSI9w7tiPXW8vbwJvAi8pryJvHa6oz3APPS8QD3wvIE94Ty8vNc8Qz2RPGs9LD3YO4+9r7yhvEO8nzutvDe88bxVPHc8tTwuvV68Ab0NO5w9ijz5O6W8S71UPQG87TyIvAA9KLyWuys9/buVPHG8dL3oPMg8oTy7vNO86DlLvQs94jwUPHK8rTsTvQ49SzwyvXo8XTzbut28aTwQvOS5KbzkvJs7Pbkcvdm8P7x9vek8V70KPb489zx8PVo8gz3lPAE9RD04vSs6Trx6u9M8uDwOPLq7OT0rvSm8Jz2svI08t7yiO5A9X73Au1i9zjxvPES9Qj3zvGw8Tb0+vSw9Vb1pOxu9zL00vYu9QL0fPcs8djy1vS88vD1PPcI8NzyFPIK97zxXvOE8hz3iOg49Gz3XvBA9sr0bvJ884jzkvNy8kbyzO++8U7s/vXo8hj2cvDQ9R72DPYI9P7p/u2m9z7w2vDK8GD02vbc97bwqPP28Yb0xvWY8RjsQvX47eDySvAU9nLwVPVG9Hz3SPIa8f723PEo8Gb28PIW7Nr1cO5M9+rxsPDA9ODz5vIE8Az3QPFu9jTw7PQs9Yj2IvBk9kL2GvAS9zLx6PI+8oLzMuwu9Jj3pvGu8Vr1SvOw7V7yJPLc8G72ovGA8dT1JvKu8KT2CPHU9ET2ZPKM84rn5ug49mD0sPfo8Uz1yu8u8T7kPvXs9qDyrPe27Lb2ePKo8/rswvPS7wb2pvLc7n72VvFe8BzyrPKC8CD2ovHA8oD0ivF07+ryLvTA8mTxtPBS9xTvXO/K8Qb2fPDu8ID2sPLO7Rj0Kve68ibzrus68jrxMOye9ybyFPF88TD2KvfU84Ty1vKM86Lw4vFm9iTzcPCq997wIPBo9TD3MPUg8W7xbPVA93jiwvGA9J70uPbu8Rz0hPM883LxNPR48Jb0HPCu9sDyNuwu8ED3RPP28gr0dPUo8Ej0pvM29tT3dvLy85rw6vZc9k7ynPM+7Q7zeu0u9Rr3tPLK8Fz2FvGs7kbxWvOS7LD2FvDC90jwjvIE8Kb2mvH48PDzlvNC5G7z7PBW9Kb1vPG48oTwkO6e9HT26OzW8QL0APEg9t7uRvA68l73zOaG87LwwvVy7ZLy+vNg7HT0sOyM8vDyevcg7xrx0vNE9wryNvYw7Nz0wvIy9W701vCi8qD2DPba83zx7PSU9+ju5PD69Lj2BvCU9hjlJPQO9sjwdPQS7Mb2ROzm9OjwYvWQ9Ur13O0I9gT3dPKU8xjuVPIk5gjyku1+97rwcvIq9dTsKvVQ88jvLPB894Tv0vAu8yL1XuqC9ibwivSY8RDw1vYw8oTwEvZI8C71xvEM9HztYPXY9FLwdu5G8dDxavDu8ir3ZvPo8tbyWvVu8Kr0pvZA8ZT3Euxe9e7youQ28ZjwHvVk9Dr1OPEQ8/ryfPD+9kb0RPFS91jurvL08pDsFPRS957wLPSa9QLw9PVu8J73+uFK89DxXve68mD2XO4Q9Bb1FvIg8Uz1pPWg9qbtsvI494TwPuzI9hLysu4k9QDvKvPM8bbsGPea7Jb30vBs8yzxavRa9STyWvOc8Sj2KvHy9ZTwqO4y8mLwhvWW8GTxZOzo9wzx+vVK8Srx6vJi8EL2nvCg9AjxwvSM8ZD0uuK88x7zuvDa9UL1SvIe8CD0hvZ088Lymu4I81D2TPIi6OD3aPGk857yOPaA8yTzdvJG8uTsmvcs8ijz8PNE8J70WvG68OL3lu9q8l732vY49Ab1sPG28bbw6vRm9zbxEvUA9+Lw7vH07Yr0QvAo8orxmvHi6FLyavCu9gDyNvBy8zjxAvQY9aT38ux49bTynvAY9ojycvQ080zy9uhi9Zz3bOxC9jT3tu0m8iLuRvYY8s7zQu5g9jT3nvM67LD0nPcG8fTwkuxc6tr11utk8l7uvPDq9Dr2kPVW9Lj2hPCO9Qz1TvY68cL21PD69mLw/O6q84T2JPYO9cz2mPZA8v7xeuXy8hz14vJ08ljxyPZY9Uz16Paq65TmKvGc7OLrou928Ib1mPcC8ub0tPSg7LTz3O6W8/LyIPUM7Trxkvbw9Dr3cvOO8C71aPeS8lrw4vDe8db07PXq5CLzDPMI96DywvAA9KrzrPJY8Dz2jPHG8k7xBvGi9OTyxPMi8QbyIPJu8k7t/POc6rbrWPCi9NTxOPZ49VT0UPSy9Sj1YPCO9JztUPQ68gjwwvLi8jDuRPEw9U70xvcW73btcPbE9/jw/PfW8ND0Kvay8W72wPFU7Iz2bPeG7OTyCPao8nzwUPTI8fLvdO/G6trt1vfy897tBvIa8JzwWPS49uzytPKM8k71VPN+8KT1bvdM8zrxaPD07VT2kPBg9TLyPPJa9xLsjPSe7GT1hvO88QjxSvds8bz2mPGw9xLx8PW89pb0Kuww9jzzAPW+7E7wxvLC8WDxnPAq8sLzNu0C9KD3jPNy8/jvqPKq9I70WPDs8UD27PNE8vjyAvWM9Xrx0vIM8Zjw3vAi9Bj5XvT09j7xvvDm9Rr15PJS8/DuXPKU9vrw0PYm8CzuYPOe8LT0zvQG9aL36uxI8RjxKu667Uj0fO6w8Lb29uqM8hrzmPMW4Xr2OPBm9/zyNPZ+7Lbw4PVi8j72/u4q7tjtJvDO9ED0QPHk9wzybPKI8i727PIe8RjwNPVo9ADwgu1U8ybxJPey8VztTOwc9KD3lO0+9sbwpvTM8gLwQvSQ8SL2ZPNC8ST2pPNC8aT0CPXe9dLuCPGw9YD1fPAY81TxIOye8wrwcPRe9SjwePdw8BT12vQe8Vb0JvYQ96zuTPS09PrzdPJE8rDzmvKe7/Dz9u/K8Zj2IPZ28Ez1SvCc8T7wEvb67R73mO/O8sTvXO3W8FL1Avfs7xTwavVY757x/PcW7GL30vBc91rz0uiY9i71kvIG8WzzQvMc9QrzHPDu9ULvwPGQ9XT26vPy6q7uGvQQ9DT07PCk9mbyHO0S7eD0BPSE9Aj3bOsQ8Fr1dvbO8IT0QPsG81L2UPN28JT0GPYG8lzwsPWy7kLyyPeE8Wru0PBq86btYu/68ibxsPGk8Jj13vMe8jLxqPGg90btGPIO8gT27PJK9P737vF28tro6vZO8Jr20O3A96zw3vZe8eT1yPYw8vDsqPRk7obxnPd87j73lPNs8Hzo2PGW82jxOvYo8vbxPPWU8HD1SPZm8Br3GvKG8ZTylPG89kD3xvH69Ir1+Ozm95DoEPIO8Pj1BPQo9hb05PHW96jxevaC8TT2BuoO807w8vNq8E71JPQy8gz0DvKS8lTyQvTC9vjyUvcG7Qr0WvK+8CL27PcK8Hj25PEo9KD2KvBo9ybv5u628abouvAW96bwpPVG8tzoOO+a7WrxKvcE8xLtQvWs95bqDPa889bsePaY8STkyPbI7FD2HO6W8u7yJvT89FjwOPcG84Tyhvey8arzGOtc81ryyvNs8fT0avYq7Kz26PNK9Szk5vfI8PT0ZvQc8vD1IvaK8Bj2mPAe94bvhPEA9TzyEPZE9zDsKvSI9s71APXS73TyBPKQ8Hr0APQs8bz0zu1e9qDx5vZq9fzybPXa9qr2OvKg9STz1vIs9cLzvPBQ96Tsmvam9/rsMvUo957o6PbA8tbslO3W7az0MPaw8Cr1gPMa8ab1ku828Gb23OwI9b7yYvCQ9aT0WvbG847wevFA7BjwbPRc8hLx5u4i8gDy4PCA9dD2MvFC8iD30PFc9SDyNPHi7xrwyvXW9kT2LPFy9+Tt+PKw8qLwUPe+55DtzvRo9HD0NupG9rjsIPCc7pjpDOwk7Nr2UvMO86zyHvTK8H70lvas9MTwvPZG8hL5UPXg7FzxrvCu9d71LvS69vbwKvVG9tzy2vZe7ED1kvFy8nDvgPIo9BL2DPQ89LL3avEA9Tj1TPYY99r0ZPR27w7tru7Q8O7sBPd+827wqPYM8WDxkPaO8Pjx1Pc26mrzXOyW9mj0Tvc29sjzWuzY9Arx4vUM81Twduhs9hjsCvS89qDudPDM6lj3NuxW7Gb1EPcK8xjzwPAi9Xz0SOw094rkXPdq8ab0/PfI7bbv7uh28GL1YPfw8K7xJPQS9mbxbvMa8xjt0PBK9a70qPSS9HDwHPTO9yTxNPMs8DbxkPCW9lL38O029nz1QvL08hb2wvJS8Zjx6vMC8JTs8PQ28aDwUu1A8Bb0GO8c8Dj1NPFi8G72EPZ48DzxrPbI7ljwYu9e5Nj0TvaC65zyLPU886j0OvZo907u1vCo9Jb0au1E7szwdPSc98Ls0Pdc8Bjz1PLK8FT0Jvmm8l7yzO1E8Uj1BPfu7MD2lvHK8i7wWvTu8TzzlvIm98zvZPBi5gb3ZO4w9xzz7O/C84Tv1vCE8PL1GvRS8Y7z+PMs8nbx3OxG9XrwsvRA927uQvJ+9U71zvDo9oTyVPTS8wTwRvbG8FD3CPYy9vb2QPUC9XD0iPC06PL2VvLa8C72VOz89ZL33vHi8Sr0KPQI9nrxUPQq9xDpOvaO7sjxWPUc8+jxbvco8jDwAPXs8TLyWvPe7vLz7u7y6Wz1Ju7Y8BD0uu2C8j7zbutM7IDz/Oki9YbsoPCC9F7ytvBO85Ls2vb28UDuWO5w8Br0TPSu9ZbwUvfe8dj2gvUA9pT2nPFC9bzy9vLa8+TqCPPY8ibx/PeY8kLw6vfK7vzzzOw88u7wnvcs6mzwwvbu8PLx2PVA6Zj2SvQu9G72rPIy93bxiPQA9Iz3KvJQ9t7zPPX49oTy3vDS7XDzZORY83rwRPeM8/zvVvC+9CLy2vBM9SD3gvB07qzzgvIS8qz2KvOs6iLtIPS+8hT1dvP88lTzUPJq6WL1KvY69WD3HvD+93rzbvAY8HLx3vXK8ujtMPNk5yrxqvKS8nLtevYO9or3lu5C97rwOOma8lT09Pci8iT2KvJI8Hr0Kvcw8Jz0OvPY8kT33O6m9mzubvOq8JD2XOje9Mj1FPZc8pbvYvCK93Tsxu9U8ELomPZo8wbyHPL87+LyHPJ+6c7x2va68i7wqvTa8mj1zvaS8lj23PDC9sryhPSi9sTy9u6C8t7ktvCI9VT6+O7o8vDvfPGG8ozyJPH68tDuQvE69AzwrvPA8YD3MveI74bwDPBO9YD0SO1e9Fj0qO2W8IbvEPAQ9WjwCvb+8s72uOQo6nL2YPKq8Zb2zPJk8Cr0WPGi9Jj2tuxW9izzGPTY81rwEvM28pbnjPbS8gDyIPQQ96bzwPFS9SrqWPJe90zugO+q7G714OkI8vDyeOzW9aLzCPcS8cj1NPci8/rycPLc8Qb2IPPg6irw1PQi9irzkPIk9brzuvAo9cD2pvA87vjyeunK8kL1RvUY8ebrEO/u83DwLvMM8jz2OPVI8FT3evB09PzwgvUs9JT3ovG+8yTvDuGq8urxfvJg9uT0uPY69YzxHvIo8AL0cPSw8trz9PMG9xL0QPdo8CL1bPNm9UjyzO4K907zKu5w8ID0yPBo8s7wwvYu8jrynu328TztyvAQ9YzygvAo9hDxMN9Y8DT2yPc0837yVuhU8kbw3PV07h7wPOw69XD2rPDk7azxXPRQ9Hz0LvVm8zzxkvX+9Gb1GvJE7oL0/uwS8WDxAvYO9JD3xu887Fj3LvEo94zwlPSi74rwnu/Y8ZD2/PAy9/zymOhQ9uDuJvEm8STwLvZS8jj1xPSK9oDygPTk8VT3bPFq7rDsWPJQ8iTwFPQa9ezztvCE9bLs1u/s8bz0fvQe4Q7sfvea8Db0dPca8MD2bPYM8Ir0OPaY6hjwkPSQ6fDsBu1C9zzxNPd65LLuQvVk9ibtMvDS81jxnPWK9SjzTvKi8Sjx3vSg9KT0NPV68Jr0Rvau8DL0IvIk83DxePYK8iz2RO9u8tb3/u6y8LT23PGY9x7wru5E9ojzlPJ28hLyDOx28JLzwvCs9zTzzPNO6gT1xO4m8nr0uu3I8FT1vPFu8Dry+vBG9tbxYvWu8Az1tPKK8HT3CPI49pzxUPQO8sTxmvXK9STxAPLQ9bTy2PMA8tj2GvKu8eLwWvSy9Kjz5uzQ8F733PLK8q7zuO0O9rDx5PYC9BL0kvRu827z8u5S9izz4O5A86Dv3PBE8Gz0dvKc7ozxWvHG9EL1DvZC9KTwNPci83jqevYE8mDzxPNq7Dr3bunA9Tr01vZc7UD2AOhU9dTp0PIE66LxUPFU8GLv4vA89YjxjvaQ8w7y+PAa8eDw1O049Ajw9vMS9R73ovB481zuTvPQ5az1lPcA8urzVu3c95zrSvA47ejzIvIO9WTuSu7Y8JD1sOza9FL2EPMO8nz3yuys8Yz2EOQG9f7zHPEO9LD3QPLA8Xr1MPU88+bxZPKK7dTwNPfS7kTyYu1Y8Sb3zPEA8NL0UPdc8nzqBvJ28gzwTPaY8y7z8vBM8+ryeOiI9vDwzPfc7zbsVvNu8QjxlPfw6ZLwwvJ28FbsCvWW9LT2rveE8HT2FPUg8v7xKPO29VL1dve08Qry3u1E8IbwOPaw86jyuO5g8nz3mvOO8jzy7vAO9BT1CPQq8H72hvNO7ijwvuz+8pL2QvVg9d7sCvtK8U70XPa07IL1rvea8PLugPRS9hDxqvP+8G7ywvHY7OztIvdc627sJvKg80DxFPEQ9dj0Zuwo9uzx2PG87kD1aPac7oj1xvHE7dDyfPMW8Ejucu5e7Ub3UOwc8qDwMPXu6uz3VvF+8j7zEPIk6tLyUPWY7LbypvPW8NbxGuuy6qD1TvWq707wWvfO8Q71yPTQ8Ej1sPU+7wzqdu9c4X7xKvQ88SrzTuxK9qjtZvQW9vjuJvMQ8Ir1gvRo9ojxfPSg9yjyFPIA8A71EO7a8ebzUvIi8FDwUPRe9yr1xPeK6JrsOveC8hzzmO3+8IDySPRa9Jz1MPKE8Dz2SvfQ8xjyWvMC85LyavB28lDwVvQC7TjuGvAK96zyxPIG8VDzkuzq6ZzqsPHY7Az2EvTU9B7zDvNC8mzzsPJE8Kj2JvCK8lDz/PLc8l7xJvFw9Bz1uPIA8sDzXvHA9Zz1Cvam90zyjvAM9qDmGvHo7nDsQPd28bLzZvKg9Cr0bvTE9qLdAPc86Hr2oPKu8drwNvVm8Hr32vKu8cb0ivZU6870HPJk9hbxsPaE8CD2AvF09sLyOPDG9tzz5vK+9nbpBPRc9jT1CPNW7i7xaPai8I72cvB07j7zOO0m9p7xvPYC7pLu3PEy93rwzPQ09EDzEPPM8XbwLO4w9JD1Hu2O9qDm1OxK9yTx8u7y8qryZvNA8FT0Ku5m9TT0LvOO8sbyMPTa9VTy2vKi7Cz2RPJK8JrsrvQ69iL3MPL68F70mPai9NT39OyI8xbyBPVo8yLuBvdm7gTxSO9E8JrwhPBE9RL2fvNe7Hzw+PJy9wTqzPGq7ET0mvYY8NL2TvHU9FD2WPO+88Dw/vWU9VrwavRo9h70RPbS7bz1HvQe727w7PQe7J70ZO0G9Kb4hPK681TswPDO8ozyBPB49Qj33vP47Pr2WvcW8I70UPfw8Vz0qPRW7vrwFOzu9YLmWPEG98bwUPS69Or1cvPc8lTz4OrG7Qb2IPIk9mjwVPLk8zryGvII8U7zdPIW8jLwxvHM8BbyOvPA6Dr3OOS+81rwNvF06tzzRvM08Vj3TPHU7jjsPPXC75rwlPQG88jzaPJS7Ij3ivKw9qjxTPeY8DT0OPTk9Lj1BvEq9RT1TPIy8lroSPdo7ez1AvYy8xzsbPYc9ND25vJQ9qz3mvBA99LtivMG8WLyePLc8xDyrvPC8AjzYPI49V7yuPJo9nT0mvUq8P70pvSi9hbwEvNS7Mz0pPQu7mbxWPIm8u7xrvXK7LzxkPDG9IL0aPS69g72jPD09ELwIPWw9ETytvEA8NTyYPM08xjwPPYe8hjyjO8G8bryGOhw8MTyIvPA7TD06vZy8Br3zPM48Dz1SvbM87b0wuwC9hrxvvTY9VroZPUq8WLyBuxw9sLyWPYi8Eb1SvRc9uzxWu628BL1GvHy9pDuMvAy9A7xPvIU8dbxVuwM9X7wevTC8jL2yvK29ObyYPbM82ryrPLQ6yrwVu068yD2evAu9Bz0qvYm7UjwNPTQ97TuhvFG9pLsHvl+6Yj00PQG817wsvf48jjwiPai7H72JvDM73DyovDI8o7yUvG69qLy0Pbg8VrzOPAs9qzw5vuM78jwRPTA9JLtZPAm8sL0PPRg9HjzoO2g9C7wJvea8VL3IvMa8Oz0Uvby89L28POa8Qz1KvZq8zbuPPDk97jk2vM88/zwnPRu9lLxCPTe9szyuPLQ6mz1vvGG9mb26PA==", + "mixes_f32": "5ertwO6atMABjNfAmaqHQE2ySsG204rBMOp7wJFAr8AAOxE/S0CVuhU2ssD8mqBA2IAxv1nkMT83yZu6uaNIuUHJwbqG2oG/SCKCPzAy3LloCZa7EL8+uv22bLroEMs7", + "scale_f32": "UyCBPWQ7RDxMEXw9", + "base_f32": "ChY2QHj3VcDr8U7AWOUNwHV/hcDK9HzApJWcwXxptb9nNFU+IZzLwZlLo8AJrcU9GSQBP0av4b4WAsrBt7OswV3ty8HInUA+BOflvQtYsMFBCZzB416twTi2s8G+BFdB", + "branch_bf16_bits": [ + 16053, + 16067 + ], + "residual_bf16_bits": [ + 48359, + 48377, + 48359, + 48377, + 48359, + 48377, + 48359, + 48377 + ], + "expected_source_tree_target_bf16_bits": [ + 14092, + 46834 + ], + "provenance": "Token 671, layer-0 attention HC, official model revision 2cba9e42aa026125f3ed06c6d98c1db82f7ca027. Four copies of the actual BF16 embedding; mixes are the independently executed native FP32 projection boundary. Target columns 3093 and 4687 of copy 3 distinguish native old RMS from pinned Torch CUDA single-row mean. Base64 payloads are little-endian BF16 or FP32." +} diff --git a/zig/src/families/deepseek_v41/index_precision_test.zig b/zig/src/families/deepseek_v41/index_precision_test.zig new file mode 100644 index 0000000000..1822f7ad0f --- /dev/null +++ b/zig/src/families/deepseek_v41/index_precision_test.zig @@ -0,0 +1,57 @@ +//! BF16 tensor boundaries are observable through which compressed positions enter sparse attention. +const std = @import("std"); +const mtl = @import("metal"); +const index = @import("csa2_index.zig"); +const kv = @import("kv.zig"); + +// The last key's extra 2^-9 dot term disappears at the source BF16 boundary; all 513 scores must tie. +test "index BF16 dot product and reduction boundaries preserve top-512 tie selection" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var ip = try index.Pipelines.load(device); + defer ip.deinit(); + var kp = try kv.Pipelines.load(device); + defer kp.deinit(); + var scratch = try index.Scratch.init(device, 1); + defer scratch.deinit(); + const q = try device.buffer(4096 * 2, mtl.ResourceOptions.shared); + defer q.deinit(); + const w = try device.buffer(32 * 2, mtl.ResourceOptions.shared); + defer w.deinit(); + const raw = try device.buffer(513 * 128 * 2, mtl.ResourceOptions.shared); + defer raw.deinit(); + const storage = try device.buffer(513 * 68, mtl.ResourceOptions.shared); + defer storage.deinit(); + @memset(q.slice(u16, 4096), 0); + @memset(w.slice(u16, 32), 0); + @memset(raw.slice(u16, 513 * 128), 0); + q.slice(u16, 4096)[0] = 0x3f80; // 1 + q.slice(u16, 4096)[32] = 0x3b00; // 2^-9, in a separate FP4 scale block + q.slice(u16, 4096)[128] = 0x3f80; + w.slice(u16, 32)[0] = 0x3f40; // 3/4 + w.slice(u16, 32)[1] = 0x3b00; + for (0..513) |row| raw.slice(u16, 513 * 128)[row * 128] = 0x3f80; + raw.slice(u16, 513 * 128)[512 * 128 + 32] = 0x3f80; + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + try kp.encodePack(enc, .{ .buffer = raw }, .{ .buffer = storage }, .{ .rows = 513, .width = 128, .stride = 128, .packed_stride = 68, .block = 32, .kind = .index_fp4 }); + try ip.encode(enc, .{ .buffer = q }, .{ .buffer = w }, .{ .buffer = storage }, scratch, .{ .rows = 1, .q_stride = 4096, .weight_stride = 32, .position = 512, .ratio = 1, .keys = 513 }); + enc.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |message| { + std.log.err("index precision qualification: {s}", .{message}); + return error.GpuExecutionFailed; + } + const positions = scratch.positions(); + const ids: [*]const i32 = @ptrCast(@alignCast(positions.buffer.contents() + positions.offset)); + for (0..512) |i| try std.testing.expectEqual(@as(i32, @intCast(i)), ids[i]); + const scores = scratch.scores(); + const values: [*]const f32 = @ptrCast(@alignCast(scores.buffer.contents() + scores.offset)); + try std.testing.expectEqual(@as(f32, 0.01171875), values[0]); + try std.testing.expectEqual(values[0], values[512]); +} diff --git a/zig/src/families/deepseek_v41/kv.zig b/zig/src/families/deepseek_v41/kv.zig new file mode 100644 index 0000000000..c8882ff590 --- /dev/null +++ b/zig/src/families/deepseek_v41/kv.zig @@ -0,0 +1,162 @@ +//! Cache operators accept projected inputs; no source projection or hidden allocations live here. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +const View = @import("view.zig").View; +pub const cache = @import("metal_kv.zig"); +pub const Kind = enum(u32) { window_fp8 = 0, main_fp4 = 1, index_fp4 = 2 }; +pub const Pack = extern struct { + rows: u32, + width: u32, + stride: u32, + packed_stride: u32, + block: u32, + kind: Kind, + first: u32 = 0, + ring: u32 = 0, + inplace: u32 = 0, + pub fn validate(p: Pack) !void { + const block: u32 = if (p.kind == .main_fp4) 16 else 32; + const bytes = if (p.kind == .window_fp8) p.width else p.width / 2; + if (p.rows == 0 or p.width == 0 or p.width % block != 0 or p.block != block or + p.stride < p.width or @as(u64, p.packed_stride) < @as(u64, bytes) + p.width / block or + p.inplace > 1 or @as(u64, p.first) + p.rows > std.math.maxInt(u32) or + @as(u64, p.rows -| 1) * p.stride + p.width > std.math.maxInt(u32) or + @as(u64, if (p.ring != 0) p.ring else p.first +| p.rows) * p.packed_stride > + std.math.maxInt(u32) or (p.ring != 0 and p.rows > p.ring)) + return error.InvalidCachePackShape; + } +}; +pub const Pool = extern struct { + rows: u32, + width: u32 = 512, + ratio: u32, + position: u32, + stride: u32, + input_f32: u32, + eps: f32 = 1e-20, + pub fn completed(p: Pool) u32 { + return (p.position + p.rows) / p.ratio - p.position / p.ratio; + } + pub fn validate(p: Pool) !void { + if (p.rows == 0 or p.width == 0 or p.width > 512 or (p.ratio != 1 and p.ratio != 2) or + p.stride < p.width or p.input_f32 != @intFromBool(p.ratio > 1) or + !std.math.isFinite(p.eps) or p.eps <= 0 or + @as(u64, p.position) + p.rows > cache.max_context or + @as(u64, p.rows -| 1) * p.stride + p.width > std.math.maxInt(u32)) + return error.InvalidCompressorShape; + } +}; +pub const Pipelines = struct { + pack: mtl.Pipeline, + pool: mtl.Pipeline, + tail: mtl.Pipeline, + copy: mtl.Pipeline, + publish: mtl.Pipeline, + unpack: mtl.Pipeline, + pub fn load(device: mtl.Device) !Pipelines { + var options = mtl.CompileOptions.mlx(); + options.math_functions = .precise; + const lib = try mtl.Library.fromSource(device, sources.deepseek_v41_kv, options); + defer lib.deinit(); + var p: Pipelines = undefined; + p.pack = try mtl.Pipeline.init(device, lib, "dsv41_pack", false); + errdefer p.pack.deinit(); + p.pool = try mtl.Pipeline.init(device, lib, "dsv41_pool", false); + errdefer p.pool.deinit(); + p.tail = try mtl.Pipeline.init(device, lib, "dsv41_pool_tail", false); + errdefer p.tail.deinit(); + p.copy = try mtl.Pipeline.init(device, lib, "dsv41_copy", false); + errdefer p.copy.deinit(); + p.publish = try mtl.Pipeline.init(device, lib, "dsv41_publish_window", false); + errdefer p.publish.deinit(); + p.unpack = try mtl.Pipeline.init(device, lib, "dsv41_unpack_rows", false); + return p; + } + pub fn deinit(p: *Pipelines) void { + p.pack.deinit(); + p.pool.deinit(); + p.tail.deinit(); + p.copy.deinit(); + p.publish.deinit(); + p.unpack.deinit(); + } + pub fn encodePack(p: *const Pipelines, enc: mtl.ComputeEncoder, x: View, storage: View, args: Pack) !void { + try args.validate(); + try x.requireRows(args.rows, args.stride, args.width, 2); + try storage.requireRows(if (args.ring != 0) args.ring else args.first + args.rows, args.packed_stride, args.width / (if (args.kind == .window_fp8) @as(u32, 1) else 2) + + args.width / args.block, 1); + enc.setPipeline(p.pack); + enc.setBuffer(x.buffer, x.offset, 0); + enc.setBuffer(storage.buffer, storage.offset, 1); + enc.setValue(args, 2); + enc.dispatchGroups(mtl.Size.of(args.width / args.block, args.rows, 1), mtl.Size.of(32, 1, 1)); + enc.barrier(); + } + pub fn encodeUnpack(p: *const Pipelines, enc: mtl.ComputeEncoder, storage: View, out: View, args: Pack) !void { + try args.validate(); + try out.requireRows(args.rows, args.stride, args.width, 2); + try storage.requireRows(if (args.ring != 0) args.ring else args.first + args.rows, args.packed_stride, args.width / (if (args.kind == .window_fp8) @as(u32, 1) else 2) + + args.width / args.block, 1); + enc.setPipeline(p.unpack); + enc.setBuffer(storage.buffer, storage.offset, 0); + enc.setBuffer(out.buffer, out.offset, 1); + enc.setValue(args, 2); + enc.dispatchThreads(mtl.Size.of(args.width, args.rows, 1), mtl.Size.of(32, 1, 1)); + enc.barrier(); + } + /// ratio2 values/gates are FP32; ratio1 values and RMS weight are BF16. Latent output is BF16. + pub fn encodePool(p: *const Pipelines, enc: mtl.ComputeEncoder, values: View, gates: View, partial: View, norm_weight: View, latent: View, args: Pool) !void { + try args.validate(); + try values.requireRows(args.rows, args.stride, args.width, if (args.input_f32 != 0) 4 else 2); + if (args.ratio > 1) { + try gates.requireRows(args.rows, args.stride, args.width, 4); + try partial.require(@as(usize, 2) * args.ratio * args.width * 4, 4); + } + if (args.completed() != 0) { + try norm_weight.require(@as(usize, args.width) * 2, 2); + try latent.requireRows(args.completed(), args.width, args.width, 2); + enc.setPipeline(p.pool); + for ([_]View{ values, gates, partial, norm_weight, latent }, 0..) |v, i| + enc.setBuffer(v.buffer, v.offset, i); + enc.setValue(args, 5); + enc.dispatchGroups(mtl.Size.of(args.completed(), 1, 1), mtl.Size.of(128, 1, 1)); + enc.barrier(); + } + if (args.ratio > 1) { + enc.setPipeline(p.tail); + for ([_]View{ values, gates, partial }, 0..) |v, i| enc.setBuffer(v.buffer, v.offset, i); + enc.setValue(args, 3); + enc.dispatchThreads(mtl.Size.of(args.width, 1, 1), mtl.Size.of(128, 1, 1)); + enc.barrier(); + } + } + pub fn encodeCopy(p: *const Pipelines, enc: mtl.ComputeEncoder, src: View, dst: View, bytes: u32) void { + enc.setPipeline(p.copy); + enc.setBuffer(src.buffer, src.offset, 0); + enc.setBuffer(dst.buffer, dst.offset, 1); + enc.setValue(bytes, 2); + enc.dispatchThreads(mtl.Size.of(bytes, 1, 1), mtl.Size.of(256, 1, 1)); + enc.barrier(); + } + /// Copy only the final window; each destination slot has one writer even for large prefill chunks. + pub fn publishWindow(p: *const Pipelines, enc: mtl.ComputeEncoder, chunk: View, ring: View, rows: u32, position: u32) !void { + if (rows == 0 or @as(u64, position) + rows > cache.max_context) return error.NoCacheRows; + try chunk.requireRows(rows, 528, 528, 1); + try ring.require(128 * 528, 1); + enc.setPipeline(p.publish); + enc.setBuffer(chunk.buffer, chunk.offset, 0); + enc.setBuffer(ring.buffer, ring.offset, 1); + enc.setValue([4]u32{ rows, position, 528, 528 }, 2); + enc.dispatchThreads(mtl.Size.of(528, @min(rows, 128), 1), mtl.Size.of(32, 1, 1)); + enc.barrier(); + } + pub fn snapshot(p: *const Pipelines, enc: mtl.ComputeEncoder, state: *cache.State) !void { + try state.begin(); + p.encodeCopy(enc, .{ .buffer = state.live }, .{ .buffer = state.backup }, @intCast(state.plan.mutableBytes())); + } + pub fn rollback(p: *const Pipelines, enc: mtl.ComputeEncoder, state: *cache.State) !void { + try state.rollback(); + p.encodeCopy(enc, .{ .buffer = state.backup }, .{ .buffer = state.live }, @intCast(state.plan.mutableBytes())); + } +}; diff --git a/zig/src/families/deepseek_v41/kv_test.zig b/zig/src/families/deepseek_v41/kv_test.zig new file mode 100644 index 0000000000..088f36a3fe --- /dev/null +++ b/zig/src/families/deepseek_v41/kv_test.zig @@ -0,0 +1,234 @@ +const std = @import("std"); +const mtl = @import("metal"); +const kv = @import("kv.zig"); +const cache = @import("metal_kv.zig"); +const View = @import("view.zig").View; +fn bf(x: f32) u16 { + const bits: u32 = @bitCast(x); + return @intCast((bits +% 0x7fff +% ((bits >> 16) & 1)) >> 16); +} +fn view(b: mtl.Buffer) View { + return .{ .buffer = b }; +} +fn submit(cb: mtl.CommandBuffer, enc: mtl.ComputeEncoder) !void { + enc.end(); + cb.commit(); + cb.wait(); + try std.testing.expect(cb.failure() == null); +} + +test "cache owners and one million token storage match the released layout" { + try std.testing.expectEqual(@as(usize, 0), try cache.ownerFor(7, 2)); + try std.testing.expectEqual(@as(usize, 1), try cache.ownerFor(13, 2)); + try std.testing.expectEqual(@as(usize, 2), try cache.ownerFor(19, 2)); + try std.testing.expectEqual(@as(usize, 3), try cache.ownerFor(39, 1)); + try std.testing.expectError(error.NoCompressedCache, cache.ownerFor(0, 0)); + const plan = cache.Plan{ .capacity = 1048576, .layers = 40 }; + try std.testing.expectEqual(@as(usize, 1048576 * 890), plan.packedBytes()); + try std.testing.expectEqual(@as(usize, 2732032), plan.mutableBytes()); + try plan.validate(); + try std.testing.expectError(error.InvalidCacheCapacity, (cache.Plan{ .capacity = 1, .layers = 41 }).validate()); + try std.testing.expectError(error.InvalidCacheLayer, cache.ownerFor(40, 1)); + try std.testing.expectError(error.InvalidCacheOwner, cache.ownerFor(39, 2)); +} + +test "FP4 E2M1 ties, signs, zero scales and BF16 in-place values" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var pipes = try kv.Pipelines.load(device); + defer pipes.deinit(); + const x = try device.buffer(64, mtl.ResourceOptions.shared); + defer x.deinit(); + const storage = try device.buffer(36, mtl.ResourceOptions.shared); + defer storage.deinit(); + const values = [_]f32{ 0, 0.5, 1, 1.5, 2, 3, 4, 6, -0.5, -1, -1.5, -2, 0.25, 0.75, 1.25, 1.75 }; + for (values, 0..) |v, i| x.slice(u16, 32)[i] = bf(v); + @memset(x.slice(u16, 32)[16..32], 0); + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + try pipes.encodePack(enc, view(x), view(storage), .{ .rows = 2, .width = 16, .stride = 16, .packed_stride = 9, .block = 16, .kind = .main_fp4, .inplace = 1 }); + try submit(cb, enc); + const expected = [_]u8{ 0x10, 0x32, 0x54, 0x76, 0xa9, 0xcb, 0x20, 0x42, 0x38 }; + try std.testing.expectEqualSlices(u8, &expected, storage.contents()[0..9]); + try std.testing.expectEqual(@as(u8, 1), storage.contents()[17]); + const rounded = [_]f32{ 0, 0.5, 1, 1.5, 2, 3, 4, 6, -0.5, -1, -1.5, -2, 0, 1, 1, 2 }; + for (rounded, 0..) |v, i| try std.testing.expectEqual(bf(v), x.slice(u16, 32)[i]); +} + +test "compressor carries a partial group across chunks without changing the pooled latent" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var pipes = try kv.Pipelines.load(device); + defer pipes.deinit(); + const values = try device.buffer(3 * 512 * 4, mtl.ResourceOptions.shared); + defer values.deinit(); + const gates = try device.buffer(3 * 512 * 4, mtl.ResourceOptions.shared); + defer gates.deinit(); + const partial = try device.buffer(4 * 512 * 4, mtl.ResourceOptions.shared); + defer partial.deinit(); + const one_shot = try device.buffer(4 * 512 * 4, mtl.ResourceOptions.shared); + defer one_shot.deinit(); + const weight = try device.buffer(512 * 2, mtl.ResourceOptions.shared); + defer weight.deinit(); + const out = try device.buffer(2 * 512 * 2, mtl.ResourceOptions.shared); + defer out.deinit(); + @memset(partial.contents()[0 .. 4 * 512 * 4], 0); + @memset(one_shot.contents()[0 .. 4 * 512 * 4], 0); + @memset(weight.slice(u16, 512), bf(1)); + for (0..3 * 512) |i| { + values.slice(f32, 3 * 512)[i] = @as(f32, @floatFromInt((i % 512) % 11)) - 5 + @as(f32, @floatFromInt(i / 512)); + gates.slice(f32, 3 * 512)[i] = @as(f32, @floatFromInt(i % 7)) * 0.25; + } + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + try pipes.encodePool(enc, view(values), view(gates), view(partial), view(weight), view(out), .{ .rows = 1, .ratio = 2, .position = 0, .stride = 512, .input_f32 = 1 }); + try pipes.encodePool(enc, view(values).shifted(512 * 4), view(gates).shifted(512 * 4), view(partial), view(weight), view(out), .{ .rows = 2, .ratio = 2, .position = 1, .stride = 512, .input_f32 = 1 }); + try pipes.encodePool(enc, view(values), view(gates), view(one_shot), view(weight), view(out).shifted(512 * 2), .{ .rows = 3, .ratio = 2, .position = 0, .stride = 512, .input_f32 = 1 }); + try submit(cb, enc); + try std.testing.expectEqualSlices(u16, out.slice(u16, 1024)[0..512], out.slice(u16, 1024)[512..1024]); + try std.testing.expectEqualSlices(f32, partial.slice(f32, 2048)[0..512], one_shot.slice(f32, 2048)[0..512]); + try std.testing.expectEqualSlices(f32, partial.slice(f32, 2048)[1024..1536], one_shot.slice(f32, 2048)[1024..1536]); +} + +test "compressor rounds the pooled FP32 latent to BF16 before RMS normalization" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var pipes = try kv.Pipelines.load(device); + defer pipes.deinit(); + const values = try device.buffer(4 * 4, mtl.ResourceOptions.shared); + defer values.deinit(); + const gates = try device.buffer(4 * 4, mtl.ResourceOptions.shared); + defer gates.deinit(); + const partial = try device.buffer(4 * 2 * 4, mtl.ResourceOptions.shared); + defer partial.deinit(); + const weight = try device.buffer(2 * 2, mtl.ResourceOptions.shared); + defer weight.deinit(); + const out = try device.buffer(2 * 2, mtl.ResourceOptions.shared); + defer out.deinit(); + @memcpy(values.slice(f32, 4), &[_]f32{ 1.00390625, 0.5, 1.00390625, 0.5 }); + @memset(gates.slice(f32, 4), 0); + @memset(partial.contents()[0..32], 0); + @memset(weight.slice(u16, 2), bf(1)); + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + try pipes.encodePool(enc, view(values), view(gates), view(partial), view(weight), view(out), .{ .rows = 2, .width = 2, .ratio = 2, .position = 0, .stride = 2, .input_f32 = 1 }); + try submit(cb, enc); + // 1.00390625 ties to 1 in BF16; normalizing the unrounded value changes the second result. + try std.testing.expectEqual(bf(1.265625), out.slice(u16, 2)[0]); + try std.testing.expectEqual(bf(0.6328125), out.slice(u16, 2)[1]); +} + +test "ring publish handles chunks beyond the window and transaction rollback restores every mutable byte" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var pipes = try kv.Pipelines.load(device); + defer pipes.deinit(); + var state = try cache.State.init(device, .{ .capacity = 1024, .layers = 40 }); + defer state.deinit(); + const chunk = try device.buffer(131 * 528, mtl.ResourceOptions.shared); + defer chunk.deinit(); + for (0..131) |r| @memset(chunk.contents()[r * 528 .. (r + 1) * 528], @intCast(r)); + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + try pipes.publishWindow(enc, view(chunk), try state.ring(0), 131, 0); + try state.append(131); + try pipes.snapshot(enc, &state); + try pipes.publishWindow(enc, view(chunk), try state.ring(0), 131, 131); + try state.append(131); + try pipes.rollback(enc, &state); + try submit(cb, enc); + try std.testing.expectEqual(@as(u32, 131), state.position); + try std.testing.expectEqualSlices(u8, state.backup.contents()[0..state.plan.mutableBytes()], state.live.contents()[0..state.plan.mutableBytes()]); + for (3..131) |r| { + const slot = r % 128; + for (state.live.contents()[slot * 528 .. (slot + 1) * 528]) |byte| + try std.testing.expectEqual(@as(u8, @intCast(r)), byte); + } + const cb2 = queue.commandBuffer(); + const enc2 = cb2.compute(.serial); + try pipes.snapshot(enc2, &state); + try pipes.publishWindow(enc2, view(chunk), try state.ring(0), 131, 131); + try state.append(131); + try state.commit(); + try submit(cb2, enc2); + try std.testing.expectEqual(@as(u32, 262), state.position); + try std.testing.expect(state.saved_position == null); +} + +test "window FP8 quantization includes the RoPE tail and index FP4 uses E8M0 scale bytes" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var pipes = try kv.Pipelines.load(device); + defer pipes.deinit(); + const x = try device.buffer(64 * 2, mtl.ResourceOptions.shared); + defer x.deinit(); + const storage = try device.buffer(66, mtl.ResourceOptions.shared); + defer storage.deinit(); + const decoded = try device.buffer(64 * 2, mtl.ResourceOptions.shared); + defer decoded.deinit(); + @memset(x.slice(u16, 64), 0); + @memset(x.slice(u16, 64)[32..64], bf(0.3)); + x.slice(u16, 64)[63] = bf(448); + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + const args = kv.Pack{ .rows = 1, .width = 64, .stride = 64, .packed_stride = 66, .block = 32, .kind = .window_fp8, .inplace = 1 }; + try pipes.encodePack(enc, view(x), view(storage), args); + try pipes.encodeUnpack(enc, view(storage), view(decoded), args); + try submit(cb, enc); + try std.testing.expectEqual(@as(u8, 127), storage.contents()[65]); + try std.testing.expectEqual(@as(u8, 126), storage.contents()[63]); + try std.testing.expectEqual(bf(0.3125), decoded.slice(u16, 64)[62]); + try std.testing.expectEqualSlices(u16, x.slice(u16, 64), decoded.slice(u16, 64)); + @memset(x.slice(u16, 64)[0..32], bf(6)); + const cb2 = queue.commandBuffer(); + const enc2 = cb2.compute(.serial); + try pipes.encodePack(enc2, view(x), view(storage), .{ .rows = 1, .width = 32, .stride = 32, .packed_stride = 17, .block = 32, .kind = .index_fp4, .inplace = 1 }); + try submit(cb2, enc2); + for (storage.contents()[0..16]) |byte| try std.testing.expectEqual(@as(u8, 0x77), byte); + try std.testing.expectEqual(@as(u8, 127), storage.contents()[16]); +} + +test { + _ = @import("csa2_test.zig"); +} + +test "cache pack rejects undersized borrowed input and output before encoding" { + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const small = try device.buffer(16, mtl.ResourceOptions.shared); + defer small.deinit(); + const input = try device.buffer(64, mtl.ResourceOptions.shared); + defer input.deinit(); + const args = kv.Pack{ .rows = 1, .width = 32, .stride = 32, .packed_stride = 33, .block = 32, .kind = .window_fp8 }; + const pipelines: *const kv.Pipelines = undefined; + try std.testing.expectError(error.InvalidTensorView, pipelines.encodePack(undefined, view(small), view(input), args)); + try std.testing.expectError(error.InvalidTensorView, pipelines.encodePack(undefined, view(input), view(small), args)); + try std.testing.expectError(error.InvalidTensorView, pipelines.encodePack(undefined, view(input).shifted(1), view(input), args)); +} + +test "cache shapes reject overflow and nonfinite compressor epsilon" { + try std.testing.expectError(error.InvalidCachePackShape, (kv.Pack{ .rows = 2, .width = 32, .stride = 32, .packed_stride = 33, .block = 32, .kind = .window_fp8, .first = std.math.maxInt(u32) }).validate()); + try std.testing.expectError(error.InvalidCompressorShape, (kv.Pool{ .rows = 1, .ratio = 1, .position = 0, .stride = 512, .input_f32 = 0, .eps = std.math.nan(f32) }).validate()); +} diff --git a/zig/src/families/deepseek_v41/metal_kv.zig b/zig/src/families/deepseek_v41/metal_kv.zig new file mode 100644 index 0000000000..f2b70175a7 --- /dev/null +++ b/zig/src/families/deepseek_v41/metal_kv.zig @@ -0,0 +1,111 @@ +//! Shared cache ownership and exact transaction snapshots; appended rows are cursor-delimited. +const std = @import("std"); +const mtl = @import("metal"); +const View = @import("view.zig").View; +pub const owners = [_]u32{ 2, 8, 14, 20 }; +pub const ratios = [_]u32{ 2, 2, 2, 1 }; +pub const index_sources = [_]u32{ 2, 8, 14, 20, 24, 28, 32, 36 }; +pub const main_stride = 288; +pub const index_stride = 68; +pub const ring_stride = 528; +pub const window = 128; +pub const max_context = 1048576; + +pub fn ownerFor(layer: u32, ratio: u32) !usize { + if (ratio == 0) return error.NoCompressedCache; + if (layer >= 40) return error.InvalidCacheLayer; + if (layer < 2 or ratio != (if (layer < 20) @as(u32, 2) else 1)) return error.InvalidCacheOwner; + var found: ?usize = null; + for (owners, 0..) |owner, i| { + if (owner <= layer and ratios[i] == ratio) found = i; + } + return found orelse error.InvalidCacheOwner; +} +pub fn isIndexSource(layer: u32) bool { + return std.mem.indexOfScalar(u32, &index_sources, layer) != null; +} +pub const Source = struct { owner: u32, ratio: u32, rows: u32, main: View, index: View, partial: View }; +pub const Plan = struct { + capacity: u32, + layers: u32, + pub fn validate(p: Plan) !void { + if (p.capacity == 0 or p.capacity > max_context or p.layers != 40) + return error.InvalidCacheCapacity; + } + pub fn packedBytes(p: Plan) usize { + var bytes: usize = 0; + for (ratios) |r| bytes += @as(usize, p.capacity / r) * (main_stride + index_stride); + return bytes; + } + pub fn mutableBytes(p: Plan) usize { + return @as(usize, p.layers) * window * ring_stride + 7 * 2 * 512 * 4; + } + pub fn residentBytes(p: Plan) usize { + return p.packedBytes() + 2 * p.mutableBytes(); + } +}; +pub const State = struct { + plan: Plan, + storage: mtl.Buffer, + live: mtl.Buffer, + backup: mtl.Buffer, + position: u32 = 0, + saved_position: ?u32 = null, + + pub fn init(device: mtl.Device, plan: Plan) !State { + try plan.validate(); + const storage = try device.buffer(plan.packedBytes(), mtl.ResourceOptions.shared); + errdefer storage.deinit(); + const live = try device.buffer(plan.mutableBytes(), mtl.ResourceOptions.shared); + errdefer live.deinit(); + const backup = try device.buffer(plan.mutableBytes(), mtl.ResourceOptions.shared); + @memset(live.contents()[0..plan.mutableBytes()], 0); + return .{ .plan = plan, .storage = storage, .live = live, .backup = backup }; + } + pub fn deinit(s: *State) void { + s.storage.deinit(); + s.live.deinit(); + s.backup.deinit(); + } + pub fn residentBytes(s: *const State) usize { + return s.plan.residentBytes(); + } + pub fn registerResidency(s: *const State, set: mtl.ResidencySet) void { + set.add(s.storage); + set.add(s.live); + set.add(s.backup); + } + pub fn ring(s: *const State, layer: u32) !View { + if (layer >= s.plan.layers) return error.InvalidCacheLayer; + return .{ .buffer = s.live, .offset = @as(usize, layer) * window * ring_stride }; + } + pub fn source(s: *const State, layer: u32, ratio: u32) !Source { + if (layer >= s.plan.layers) return error.InvalidCacheLayer; + if (s.position > s.plan.capacity) return error.CacheCapacityExceeded; + const id = try ownerFor(layer, ratio); + var packed_off: usize = 0; + var partial_off: usize = @as(usize, s.plan.layers) * window * ring_stride; + for (ratios[0..id]) |r| { + packed_off += @as(usize, s.plan.capacity / r) * (main_stride + index_stride); + partial_off += 2 * r * 512 * 4; + } + return .{ .owner = owners[id], .ratio = ratio, .rows = s.position / ratio, .main = .{ .buffer = s.storage, .offset = packed_off }, .index = .{ .buffer = s.storage, .offset = packed_off + @as(usize, s.plan.capacity / ratio) * main_stride }, .partial = .{ .buffer = s.live, .offset = partial_off } }; + } + pub fn append(s: *State, rows: u32) !void { + if (s.position > s.plan.capacity or rows == 0 or rows > s.plan.capacity - s.position) + return error.CacheCapacityExceeded; + s.position += rows; + } + pub fn begin(s: *State) !void { + if (s.saved_position != null) return error.CacheTransactionActive; + s.saved_position = s.position; + } + pub fn commit(s: *State) !void { + if (s.saved_position == null) return error.NoCacheTransaction; + s.saved_position = null; + } + pub fn rollback(s: *State) !void { + s.position = s.saved_position orelse return error.NoCacheTransaction; + s.saved_position = null; + } +}; diff --git a/zig/src/families/deepseek_v41/metal_model.zig b/zig/src/families/deepseek_v41/metal_model.zig new file mode 100644 index 0000000000..9823a9a019 --- /dev/null +++ b/zig/src/families/deepseek_v41/metal_model.zig @@ -0,0 +1,376 @@ +//! Plain V4.1 execution over resident EXL3 weights and source-owned packed caches. +const std = @import("std"); +const mtl = @import("metal"); +const checkpoint = @import("core").checkpoint_metal; +const exl3 = @import("core").exl3_projection; +const config = @import("config.zig"); +const load = @import("resident_load.zig"); +const weights = @import("weights.zig"); +const work = @import("workspace.zig"); +pub const attention = @import("attention.zig"); +const kv = @import("kv.zig"); +pub const ops = @import("ops.zig"); +pub const mhc = @import("mhc.zig"); +const route = @import("route.zig"); +const engram = @import("engram.zig"); +const View = @import("view.zig").View; +const gpu_counters = @import("profiler.zig"); +const expert_groups = @import("core").expert_groups; + +pub const prefill_step = 256; +pub const Logits = enum { none, last, all }; +pub const Options = struct { + dir: []const u8, + table_dir: ?[]const u8 = null, + token_map: []const u8, + context: u32 = 0, + chunk_rows: u32 = prefill_step, +}; + +const Pipelines = struct { + projection: exl3.Pipelines, + ops: ops.Pipelines, + mhc: mhc.Pipelines, + route: route.Pipelines, + engram: engram.Pipelines, + attention: attention.Pipelines, + groups: expert_groups.Pipelines, + + fn init(device: mtl.Device, capacity: u32, c: config.Text) !Pipelines { + var ep = try exl3.Pipelines.load(device); + errdefer ep.deinit(); + var op = try ops.Pipelines.load(device); + errdefer op.deinit(); + var hp = try mhc.Pipelines.load(device); + errdefer hp.deinit(); + var rp = try route.Pipelines.load(device); + errdefer rp.deinit(); + var gp = try engram.Pipelines.load(device); + errdefer gp.deinit(); + var groups = try expert_groups.Pipelines.load(device); + errdefer groups.deinit(); + const ap = try attention.Pipelines.load(device, capacity, c); + return .{ .projection = ep, .ops = op, .mhc = hp, .route = rp, .engram = gp, .attention = ap, .groups = groups }; + } + + fn deinit(p: *Pipelines) void { + p.attention.deinit(); + p.engram.deinit(); + p.route.deinit(); + p.groups.deinit(); + p.mhc.deinit(); + p.ops.deinit(); + p.projection.deinit(); + } +}; + +/// Serialized unique owner; submissions finish before return/release, cancellation is between bounded chunks. +pub const Model = struct { + gpa: std.mem.Allocator, + device: mtl.Device, + queue: mtl.Queue, + residency: mtl.ResidencySet, + checkpoint: checkpoint.Checkpoint, + weights: weights.Model, + cfg: config.Config, + cfg_arena: std.heap.ArenaAllocator, + cache: kv.cache.State, + workspace: work.Workspace, + attention_workspace: attention.Workspace, + pipelines: Pipelines, + admitted_bytes: usize, + last_gpu_seconds: f64 = 0, + logits_rows: usize = 0, + valid: bool = true, + profile: ?*gpu_counters.Recorder = null, + grouped_prefill: bool = true, + observer: ?Observer = null, + pub const Observer = struct { + ctx: *anyopaque, + encode: *const fn (*anyopaque, mtl.ComputeEncoder, u32, gpu_counters.Stage, *const Model) void, + }; + + pub fn open(gpa: std.mem.Allocator, io: std.Io, o: Options) !Model { + var arena = std.heap.ArenaAllocator.init(gpa); + errdefer arena.deinit(); + const c = try config.Config.load(arena.allocator(), io, o.dir); + const capacity = try c.context(o.context); + const table_dir = o.table_dir orelse c.engram_table_dir orelse return error.MissingEngramTables; + var plan = try load.Plan.open(gpa, io, o.dir, table_dir); + defer plan.deinit(); + try weights.validate(&plan); + const map_blob = try std.Io.Dir.cwd().readFileAlloc(io, o.token_map, arena.allocator(), .limited(2 << 20)); + const map = try engram.hash.TokenMap.parse(map_blob); + const wp = try work.Plan.init(o.chunk_rows); + const kp = kv.cache.Plan{ .capacity = capacity, .layers = 40 }; + const workspace_bytes = try std.math.add(usize, wp.bytes, attention.Workspace.bytes(o.chunk_rows)); + const cache_bytes = try std.math.add(usize, kp.residentBytes(), weights.sizeofDescriptors()); + const execution_bytes = try std.math.add(usize, workspace_bytes, cache_bytes); + const rotary_bytes = attention.Pipelines.residentBytes(capacity, c.text_config); + const needed = try std.math.add(usize, plan.bytes, try std.math.add(usize, execution_bytes, rotary_bytes)); + const device = try mtl.Device.init(); + errdefer device.deinit(); + const margin: u64 = 8 << 30; + const available = device.maxWorkingSet() -| device.allocated() -| margin; + if (needed > available) { + std.log.err("full residency needs {d} bytes plus {d} margin; Metal admission allows {d}", .{ needed, margin, available }); + return error.InsufficientResidentMemory; + } + const max_buffer = mtl.objc.msg(u64, device.id, "maxBufferLength", .{}); + for (plan.shards.items) |shard| if (shard.end - shard.begin > max_buffer) return error.ShardExceedsMetalBufferLimit; + std.log.info("full resident admission: weights={d}, total={d}, context={d}, headroom={d}", .{ plan.bytes, needed, capacity, available - needed }); + var pipelines = try Pipelines.init(device, capacity, c.text_config); + errdefer pipelines.deinit(); + const queue = try device.queue(); + errdefer queue.deinit(); + const residency = try device.residencySet(plan.shards.items.len + 128); + errdefer residency.deinit(); + var cp = try plan.materialize(gpa, device, residency, queue); + errdefer cp.deinit(); + var bound = try weights.bind(device, &cp); + errdefer bound.deinit(); + var cache = try kv.cache.State.init(device, kp); + errdefer cache.deinit(); + var workspace = try work.Workspace.init(device, wp, map.bytes); + errdefer workspace.deinit(); + var aw = try attention.Workspace.init(device, o.chunk_rows); + errdefer aw.deinit(); + bound.registerResidency(residency); + cache.registerResidency(residency); + residency.add(workspace.buffer); + aw.registerResidency(residency); + pipelines.attention.registerResidency(residency); + residency.commit(); + residency.requestResidency(); + queue.addResidencySet(residency); + return .{ .gpa = gpa, .device = device, .queue = queue, .residency = residency, .checkpoint = cp, .weights = bound, .cfg = c, .cfg_arena = arena, .cache = cache, .workspace = workspace, .attention_workspace = aw, .pipelines = pipelines, .admitted_bytes = needed }; + } + + pub fn deinit(m: *Model) void { + m.queue.removeResidencySet(m.residency); + m.residency.deinit(); + m.attention_workspace.deinit(); + m.workspace.deinit(); + m.cache.deinit(); + m.weights.deinit(); + m.checkpoint.deinit(); + m.pipelines.deinit(); + m.queue.deinit(); + m.device.deinit(); + m.cfg_arena.deinit(); + } + + pub fn reset(m: *Model) !void { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + m.valid = false; + m.logits_rows = 0; + m.cache.position = 0; + m.attention_workspace.selection_layer = null; + m.attention_workspace.selection_rows = 0; + m.attention_workspace.candidate_rows = 0; + m.last_gpu_seconds = 0; + m.cache.saved_position = null; + @memset(m.cache.live.contents()[0..m.cache.plan.mutableBytes()], 0); + const cb = m.queue.commandBufferUnretained(); + const enc = cb.compute(.serial); + engram.encodeReset(enc, &m.pipelines.engram, m.workspace.view(.history)); + enc.end(); + cb.commit(); + cb.wait(); + try completed(cb); + m.valid = true; + m.logits_rows = 0; + } + + /// One command buffer per chunk; no operator waits or reads routes back to the CPU. + pub fn forward(m: *Model, ids: []const u32, logits: Logits) ![]const f32 { + if (!m.valid) return error.InvalidModelState; + if (ids.len == 0 or ids.len > m.workspace.plan.rows) return error.InvalidPrefillChunk; + if (m.cache.position > m.cache.plan.capacity or + ids.len > m.cache.plan.capacity - m.cache.position) return error.ContextFull; + for (ids) |id| if (id >= 129280) return error.UnsupportedToken; + m.logits_rows = 0; + @memcpy(m.workspace.tokenIds(ids.len), ids); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const rows: u32 = @intCast(ids.len); + const e = m.queue.commandBufferUnretained(); + var enc = e.compute(.serial); + for (m.checkpoint.shards.items) |shard| enc.useResource(shard.buffer, mtl.ResourceUsage.read); + var encoding = true; + var record_started = false; + errdefer { + if (m.profile) |recorder| { + if (record_started) { + if (encoding) recorder.cancelEncoding(&enc) catch {} else recorder.discardAfterCompletion() catch {}; + } else if (encoding) enc.end(); + } else if (encoding) enc.end(); + } + if (m.profile) |recorder| { + try recorder.beginCommand(e); + record_started = true; + } + try m.profileSpan(&enc, .other); + const p = &m.pipelines; + const ws = &m.workspace; + for (engram.hash.layers, 0..) |meta, index| { + try engram.encodeHash(enc, &p.engram, meta, ws.view(.ids), ws.view(.token_map), null, ws.view(.history), ws.view(if (index == 0) .engram_ids_0 else .engram_ids_1), rows); + } + try engram.encodeCommit(enc, &p.engram, ws.view(.ids), ws.view(.token_map), null, ws.view(.history), rows); + ops.encodeEmbedding(enc, &p.ops, ws.view(.ids), tensor(m.weights.embed), ws.view(.collapsed), .{ .rows = rows, .width = 5120, .table_stride = 5120, .y_stride = 5120, .copies = 1 }); + enc.barrier(); + mhc.encodeExpand(enc, &p.mhc, ws.view(.collapsed), ws.view(.stream), ws.view(.previous_pre), .{ .rows = rows, .width = 5120, .x_stride = 5120, .stream_stride = 20480 }); + enc.barrier(); + for (&m.weights.layers, 0..) |*layer, index| { + if (layer.engram) |g| { + try m.profileSpan(&enc, .engram); + try engram.encodeGather(enc, &p.engram, .{ .weight = tensor(g.table), .scale = tensor(g.scales), .rows = @intCast(g.table.shape[0]) }, ws.view(if (index == 1) .engram_ids_0 else .engram_ids_1), ws.view(.engram_gather), rows, 6144); + try m.project(enc, g.wkv, ws.view(.engram_gather), ws.view(.engram_kv), rows, 6144, 25600, .bf16); + try engram.encodeGate(enc, &p.engram, .{ .q_weight = tensor(g.q_weight), .k_weight = tensor(g.k_weight), .dim = 5120, .copies = 4, .weight_dtype = switch (g.q_weight.dtype) { + .bf16 => .bf16, + .f16 => .f16, + .f32 => .f32, + else => unreachable, + }, .hidden_stride = 20480, .kv_stride = 25600, .output_stride = 20480, .eps = m.cfg.text_config.rms_norm_eps }, ws.view(.stream), ws.view(.engram_kv), null, ws.view(.stream), rows); + } + try m.profileSpan(&enc, .mhc); + m.mixes(enc, layer.hc_attn, ws.view(.attention_pre), rows); + mhc.encodeCollapse(enc, &p.mhc, ws.view(.stream), ws.view(.previous_pre), ws.view(.collapsed), .{ .rows = rows, .width = 5120, .stream_stride = 20480, .y_stride = 5120 }); + enc.barrier(); + m.normalize(enc, ws.view(.collapsed), layer.attn_norm, ws.view(.normalized), rows); + try m.profileSpan(&enc, .attention); + try attention.encode(&p.attention, enc, &layer.attn, &m.cache, &m.attention_workspace, ws.view(.normalized), ws.view(.branch), &p.projection, ws.projection(), @intCast(index), rows, m.cfg.text_config); + enc.barrier(); + try m.profileSpan(&enc, .mhc); + m.post(enc, rows); + if (m.observer) |observer| observer.encode(observer.ctx, enc, @intCast(index), .attention, m); + m.mixes(enc, layer.hc_ffn, ws.view(.previous_pre), rows); + mhc.encodeCollapse(enc, &p.mhc, ws.view(.stream), ws.view(.attention_pre), ws.view(.collapsed), .{ .rows = rows, .width = 5120, .stream_stride = 20480, .y_stride = 5120 }); + enc.barrier(); + m.normalize(enc, ws.view(.collapsed), layer.ffn_norm, ws.view(.normalized), rows); + try m.profileSpan(&enc, .moe); + try m.moe(enc, &layer.moe, rows); + try m.profileSpan(&enc, .mhc); + m.post(enc, rows); + if (m.observer) |observer| observer.encode(observer.ctx, enc, @intCast(index), .moe, m); + } + if (logits != .none) { + try m.profileSpan(&enc, .head); + const output_rows: u32 = if (logits == .all) rows else 1; + const first: usize = if (logits == .all) 0 else rows - 1; + mhc.encodeCollapse(enc, &p.mhc, ws.view(.stream).shifted(first * 20480 * 2), ws.view(.previous_pre).shifted(first * 4 * 4), ws.view(.collapsed), .{ .rows = output_rows, .width = 5120, .stream_stride = 20480, .y_stride = 5120 }); + enc.barrier(); + m.normalize(enc, ws.view(.collapsed), m.weights.norm, ws.view(.normalized), output_rows); + try m.project(enc, m.weights.head, ws.view(.normalized), ws.view(.logits), output_rows, 5120, 129280, .f32); + } + if (m.profile) |recorder| { + if (recorder.activeSpan()) try recorder.endSpan(&enc); + try recorder.finishEncoding(&enc); + } else enc.end(); + encoding = false; + m.valid = false; + e.commit(); + e.wait(); + try completed(e); + if (m.profile) |recorder| _ = try recorder.readAfterCompletion(); + try m.cache.append(rows); + m.last_gpu_seconds = e.gpuSeconds(); + m.valid = true; + m.logits_rows = if (logits == .all) rows else if (logits == .last) 1 else 0; + return ws.logits(m.logits_rows); + } + + fn profileSpan(m: *Model, encoder: *mtl.ComputeEncoder, stage: gpu_counters.Stage) !void { + const recorder = m.profile orelse return; + if (!recorder.enabled()) return; + if (recorder.activeSpan()) try recorder.endSpan(encoder); + try recorder.beginSpan(encoder, stage); + for (m.checkpoint.shards.items) |shard| encoder.useResource(shard.buffer, mtl.ResourceUsage.read); + } + + fn normalize(m: *Model, enc: mtl.ComputeEncoder, x: View, w: checkpoint.Tensor, y: View, rows: u32) void { + ops.encodeRms(enc, &m.pipelines.ops, x, tensor(w), y, .{ .rows = rows, .width = 5120, .x_stride = 5120, .y_stride = 5120, .eps = m.cfg.text_config.rms_norm_eps }); + enc.barrier(); + } + + fn mixes(m: *Model, enc: mtl.ComputeEncoder, w: weights.Hc, pre: View, rows: u32) void { + const ws = &m.workspace; + mhc.encodeRms(enc, &m.pipelines.mhc, ws.view(.stream), ws.view(.inv), .{ .rows = rows, .width = 5120, .stream_stride = 20480, .eps = m.cfg.text_config.rms_norm_eps }); + ops.encodeDense(enc, &m.pipelines.ops, ws.view(.stream), tensor(w.projection), ws.view(.mixes), .{ .rows = rows, .k = 20480, .n = 24, .x_stride = 20480, .w_stride = 20480, .y_stride = 24, .w_dtype = .f32, .y_dtype = .f32 }); + enc.barrier(); + mhc.encodeSplit(enc, &m.pipelines.mhc, ws.view(.mixes), ws.view(.inv), tensor(w.scale), tensor(w.base), pre, ws.view(.post), ws.view(.comb), .{ .rows = rows, .eps = m.cfg.text_config.hc_eps }); + enc.barrier(); + } + + fn post(m: *Model, enc: mtl.ComputeEncoder, rows: u32) void { + const ws = &m.workspace; + mhc.encodePost(enc, &m.pipelines.mhc, ws.view(.branch), ws.view(.stream), ws.view(.post), ws.view(.comb), ws.view(.stream), .{ .rows = rows, .width = 5120, .branch_stride = 5120, .residual_stride = 20480, .y_stride = 20480 }); + enc.barrier(); + } + + fn project(m: *Model, enc: mtl.ComputeEncoder, matrix: exl3.Matrix, x: View, y: View, rows: u32, xs: usize, ys: usize, output: exl3.Output) !void { + try exl3.encode(&m.pipelines.projection, enc, matrix, convert(x), convert(y), m.workspace.projection(), .{ .count = rows, .x_stride = xs, .y_stride = ys, .output = output }); + enc.barrier(); + } + + fn moe(m: *Model, enc: mtl.ComputeEncoder, w: *const weights.Moe, rows: u32) !void { + const ws = &m.workspace; + const p = &m.pipelines; + ops.encodeDense(enc, &p.ops, ws.view(.normalized), tensor(w.gate), ws.view(.router_logits), .{ .rows = rows, .k = 5120, .n = 384, .x_stride = 5120, .w_stride = 5120, .y_stride = 384, .w_dtype = switch (w.gate.dtype) { + .bf16 => .bf16, + else => .f16, + }, .y_dtype = .f32 }); + enc.barrier(); + route.encode(enc, &p.route, ws.view(.router_logits), tensor(w.bias), ws.view(.route_ids), ws.view(.route_weights), .{ .rows = rows, .bias_dtype = switch (w.bias.dtype) { + .f32 => .f32, + .f16 => .f16, + .bf16 => .bf16, + else => .f16, + } }); + enc.barrier(); + const grouped: ?expert_groups.Scratch = if (m.grouped_prefill and rows >= 8) try ws.grouping(rows) else null; + var routed_rows = exl3.Rows{ .count = @as(usize, rows) * 6, .x_stride = 5120, .y_stride = 2304, .input_rows = rows, .output_rows = @as(usize, rows) * 6, .routes = .{ .ids = convert(ws.view(.route_ids)), .repeat = 6 } }; + if (grouped) |g| { + try expert_groups.encode(enc, &p.groups, convert(ws.view(.route_ids)), g, 6); + routed_rows.count = g.plan.padded; + routed_rows.input_ids = g.view(.token_ids); + routed_rows.output_ids = g.view(.slot_ids); + routed_rows.routes = .{ .ids = g.view(.expert_ids) }; + routed_rows.grouped = true; + } + for (0..2) |i| { + try exl3.encodeIndexed(&p.projection, enc, .{ .descriptors = .{ .buffer = w.descriptors[i] }, .count = 384, .k = 5120, .n = 2304 }, convert(ws.view(.normalized)), convert(ws.view(if (i == 0) .gate else .up)), ws.projection(), routed_rows); + enc.barrier(); + } + ops.encodeSwiGLU(enc, &p.ops, ws.view(.gate), ws.view(.up), ws.view(.route_weights), ws.view(.activation), .{ .rows = rows * 6, .width = 2304, .input_stride = 2304, .output_stride = 2304, .weighted = true }); + enc.barrier(); + routed_rows.x_stride = 2304; + routed_rows.y_stride = 5120; + routed_rows.routes.?.repeat = 1; + routed_rows.input_rows = @as(usize, rows) * 6; + if (grouped) |g| routed_rows.input_ids = g.view(.slot_ids); + try exl3.encodeIndexed(&p.projection, enc, .{ .descriptors = .{ .buffer = w.descriptors[2] }, .count = 384, .k = 2304, .n = 5120 }, convert(ws.view(.activation)), convert(ws.view(.routed)), ws.projection(), routed_rows); + enc.barrier(); + try m.project(enc, w.shared[0], ws.view(.normalized), ws.view(.shared_gate), rows, 5120, 2304, .bf16); + try m.project(enc, w.shared[1], ws.view(.normalized), ws.view(.shared_up), rows, 5120, 2304, .bf16); + ops.encodeSwiGLU(enc, &p.ops, ws.view(.shared_gate), ws.view(.shared_up), ws.view(.route_weights), ws.view(.shared_activation), .{ .rows = rows, .width = 2304, .input_stride = 2304, .output_stride = 2304 }); + enc.barrier(); + try m.project(enc, w.shared[2], ws.view(.shared_activation), ws.view(.shared_down), rows, 2304, 5120, .bf16); + ops.encodeCombine(enc, &p.ops, ws.view(.routed), ws.view(.shared_down), ws.view(.route_ids), ws.view(.branch), .{ .rows = rows, .width = 5120, .routed_row_stride = 6 * 5120, .slot_stride = 5120, .shared_stride = 5120, .y_stride = 5120 }); + enc.barrier(); + } +}; + +fn tensor(t: checkpoint.Tensor) View { + return .{ .buffer = t.buffer, .offset = t.offset }; +} +fn convert(v: View) exl3.View { + return .{ .buffer = v.buffer, .offset = v.offset }; +} +fn completed(cb: mtl.CommandBuffer) !void { + if (cb.failure()) |message| { + std.log.err("DeepSeek V4.1 GPU execution failed: {s}", .{message}); + return error.GpuExecutionFailed; + } +} diff --git a/zig/src/families/deepseek_v41/mhc.zig b/zig/src/families/deepseek_v41/mhc.zig new file mode 100644 index 0000000000..3ac499423f --- /dev/null +++ b/zig/src/families/deepseek_v41/mhc.zig @@ -0,0 +1,95 @@ +//! Four-copy mHC; callers validate shapes/storage and own previous-pre coefficients. +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +pub const View = @import("view.zig").View; +pub const Pipelines = struct { + rms: mtl.Pipeline, + split: mtl.Pipeline, + collapse: mtl.Pipeline, + post: mtl.Pipeline, + expand: mtl.Pipeline, + pub fn load(device: mtl.Device) !Pipelines { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + var options = mtl.CompileOptions.mlx(); + options.math_functions = .precise; + const lib = try mtl.Library.fromSource(device, sources.deepseek_v41_mhc, options); + defer lib.deinit(); + const rms = try mtl.Pipeline.init(device, lib, "dsv41_hc_rms", false); + errdefer rms.deinit(); + const split = try mtl.Pipeline.init(device, lib, "dsv41_hc_split", false); + errdefer split.deinit(); + const collapse = try mtl.Pipeline.init(device, lib, "dsv41_hc_collapse", false); + errdefer collapse.deinit(); + const post = try mtl.Pipeline.init(device, lib, "dsv41_hc_post", false); + errdefer post.deinit(); + const expand = try mtl.Pipeline.init(device, lib, "dsv41_hc_expand", false); + return .{ .rms = rms, .split = split, .collapse = collapse, .post = post, .expand = expand }; + } + pub fn deinit(p: *Pipelines) void { + p.rms.deinit(); + p.split.deinit(); + p.collapse.deinit(); + p.post.deinit(); + p.expand.deinit(); + } +}; +fn bind(e: mtl.ComputeEncoder, refs: anytype) void { + inline for (refs, 0..) |r, i| e.setBuffer(r.buffer, r.offset, i); +} +pub const Rms = extern struct { rows: u32, width: u32, stream_stride: u32, eps: f32 }; +/// BF16 flattened [rows,4*width] -> FP32 inv[rows]; fixed CUDA single-row tree, not every batched CUDA grouping. +pub fn encodeRms(e: mtl.ComputeEncoder, p: *const Pipelines, stream: View, inv: View, d: Rms) void { + e.setPipeline(p.rms); + bind(e, .{ stream, inv }); + e.setValue(d, 2); + e.dispatchGroups(mtl.Size.of(d.rows, 1, 1), mtl.Size.of(512, 1, 1)); +} +pub const Split = extern struct { + rows: u32, + mix_stride: u32 = 24, + coefficient_stride: u32 = 4, + comb_stride: u32 = 16, + eps: f32 = 1e-6, +}; +/// FP32 mixes[rows,24], inv, scale[3], base[24] -> pre/post[rows,4], comb[rows,4,4]; 20 iterations. +pub fn encodeSplit(e: mtl.ComputeEncoder, p: *const Pipelines, mixes: View, inv: View, scale: View, base: View, pre: View, post: View, comb: View, d: Split) void { + e.setPipeline(p.split); + bind(e, .{ mixes, inv, scale, base, pre, post, comb }); + e.setValue(d, 7); + e.dispatchThreads(mtl.Size.of(d.rows, 1, 1), mtl.Size.of(32, 1, 1)); +} +pub const Collapse = extern struct { rows: u32, width: u32, stream_stride: u32, y_stride: u32, pre_stride: u32 = 4 }; +/// Collapse BF16 [rows,4,width] with previous sublayer pre[rows,4]; head uses final FFN pre, initially one-hot. +pub fn encodeCollapse(e: mtl.ComputeEncoder, p: *const Pipelines, stream: View, previous_pre: View, y: View, d: Collapse) void { + e.setPipeline(p.collapse); + bind(e, .{ stream, previous_pre, y }); + e.setValue(d, 3); + e.dispatchThreads(mtl.Size.of(d.width, d.rows, 1), mtl.Size.of(256, 1, 1)); +} +pub const Post = extern struct { + rows: u32, + width: u32, + branch_stride: u32, + residual_stride: u32, + y_stride: u32, + coefficient_stride: u32 = 4, + comb_stride: u32 = 16, +}; +/// y[k,d]=post[k]*branch[d]+sum_j comb[j,k]*residual[j,d]; BF16 inputs/output, FP32 coefficients/math. + +// Element pitches; exact residual alias is safe, but output must not alias branch/post/comb. +pub fn encodePost(e: mtl.ComputeEncoder, p: *const Pipelines, branch: View, residual: View, post: View, comb: View, y: View, d: Post) void { + e.setPipeline(p.post); + bind(e, .{ branch, residual, post, comb, y }); + e.setValue(d, 5); + e.dispatchThreads(mtl.Size.of(d.width, d.rows, 1), mtl.Size.of(256, 1, 1)); +} +pub const Expand = extern struct { rows: u32, width: u32, x_stride: u32, stream_stride: u32, pre_stride: u32 = 4 }; +/// BF16 x[rows,width] copied into four streams and FP32 initial pre[rows,4] set to [1,0,0,0]. +pub fn encodeExpand(e: mtl.ComputeEncoder, p: *const Pipelines, x: View, stream: View, pre: View, d: Expand) void { + e.setPipeline(p.expand); + bind(e, .{ x, stream, pre }); + e.setValue(d, 3); + e.dispatchThreads(mtl.Size.of(d.width, d.rows, 1), mtl.Size.of(256, 1, 1)); +} diff --git a/zig/src/families/deepseek_v41/mhc_test.zig b/zig/src/families/deepseek_v41/mhc_test.zig new file mode 100644 index 0000000000..ee0d888e28 --- /dev/null +++ b/zig/src/families/deepseek_v41/mhc_test.zig @@ -0,0 +1,209 @@ +//! Consumer-visible mHC coefficient, reduction and BF16 boundary regressions. +const std = @import("std"); +const mtl = @import("metal"); +const mhc = @import("mhc.zig"); +const support = @import("test_support.zig"); +const bf = support.bf; +const fp = support.fp; +const view = support.view; +const Buffers = support.Buffers; +const done = support.done; +const close = support.close; +fn splitReference(mixes: []const f32, inv: f64, scales: []const f32, base: []const f32) [24]f64 { + var out: [24]f64 = undefined; + for (0..4) |j| { + out[j] = 1 / (1 + @exp(-(@as(f64, mixes[j]) * inv * scales[0] + base[j]))) + 1e-6; + out[4 + j] = 2 / (1 + @exp(-(@as(f64, mixes[4 + j]) * inv * scales[1] + base[4 + j]))); + } + for (0..4) |j| { + var max: f64 = -std.math.inf(f64); + for (0..4) |k| { + const q = 8 + 4 * j + k; + out[q] = @as(f64, mixes[q]) * inv * scales[2] + base[q]; + max = @max(max, out[q]); + } + var sum: f64 = 0; + for (0..4) |k| { + const q = 8 + 4 * j + k; + out[q] = @exp(out[q] - max); + sum += out[q]; + } + for (0..4) |k| out[8 + 4 * j + k] = out[8 + 4 * j + k] / sum + 1e-6; + } + for (0..20) |it| { + if (it > 0) for (0..4) |j| { + var s: f64 = 0; + for (0..4) |k| s += out[8 + 4 * j + k]; + for (0..4) |k| out[8 + 4 * j + k] /= s + 1e-6; + }; + for (0..4) |k| { + var s: f64 = 0; + for (0..4) |j| s += out[8 + 4 * j + k]; + for (0..4) |j| out[8 + 4 * j + k] /= s + 1e-6; + } + } + return out; +} + +test "V4.1 flattened mHC RMS, 20 Sinkhorn iterations, previous collapse, source-destination post and identity expand" { + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try mhc.Pipelines.load(device); + defer p.deinit(); + var b = Buffers{ .device = device }; + defer b.deinit(); + const stream = try b.get(u16, 16); + for (stream.slice(u16, 16), 0..) |*v, i| v.* = bf(@as(f32, @floatFromInt(i + 1)) * 0.25); + const inv = try b.get(f32, 1); + const mixes = try b.get(f32, 24); + for (mixes.slice(f32, 24), 0..) |*v, i| v.* = (@as(f32, @floatFromInt(i % 9)) - 4) * 3; + const scale = try b.get(f32, 3); + @memcpy(scale.slice(f32, 3), &[_]f32{ 0.3, -0.7, 1.2 }); + const base = try b.get(f32, 24); + for (base.slice(f32, 24), 0..) |*v, i| v.* = (@as(f32, @floatFromInt(i % 5)) - 2) * 0.3; + const pre = try b.get(f32, 4); + const post = try b.get(f32, 4); + const comb = try b.get(f32, 16); + const previous = try b.get(f32, 4); + @memcpy(previous.slice(f32, 4), &[_]f32{ 0.1, 0.2, 0.3, 0.4 }); + const collapse = try b.get(u16, 4); + const branch = try b.get(u16, 4); + @memset(branch.slice(u16, 4), bf(2)); + const y = try b.get(u16, 16); + const expanded = try b.get(u16, 16); + const initial = try b.get(f32, 4); + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + mhc.encodeRms(e, &p, view(stream), view(inv), .{ .rows = 1, .width = 4, .stream_stride = 16, .eps = 1e-20 }); + mhc.encodeSplit(e, &p, view(mixes), view(inv), view(scale), view(base), view(pre), view(post), view(comb), .{ .rows = 1 }); + mhc.encodeCollapse(e, &p, view(stream), view(previous), view(collapse), .{ .rows = 1, .width = 4, .stream_stride = 16, .y_stride = 4 }); + mhc.encodePost(e, &p, view(branch), view(stream), view(post), view(comb), view(y), .{ .rows = 1, .width = 4, .branch_stride = 4, .residual_stride = 16, .y_stride = 16 }); + mhc.encodeExpand(e, &p, view(branch), view(expanded), view(initial), .{ .rows = 1, .width = 4, .x_stride = 4, .stream_stride = 16 }); + try done(cb, e); + var squares: f64 = 0; + for (stream.slice(u16, 16)) |v| squares += @as(f64, fp(v)) * fp(v); + const ir = 1 / @sqrt(squares / 16 + 1e-20); + try close(ir, inv.slice(f32, 1)[0], 1e-6); + const ref = splitReference(mixes.slice(f32, 24), ir, scale.slice(f32, 3), base.slice(f32, 24)); + for (0..4) |j| { + try close(ref[j], pre.slice(f32, 4)[j], 2e-6); + try close(ref[4 + j], post.slice(f32, 4)[j], 2e-6); + } + for (0..16) |q| try close(ref[8 + q], comb.slice(f32, 16)[q], 3e-6); + for (0..4) |i| { + var collapsed: f64 = 0; + for (0..4) |j| collapsed += @as(f64, previous.slice(f32, 4)[j]) * fp(stream.slice(u16, 16)[j * 4 + i]); + try close(collapsed, fp(collapse.slice(u16, 4)[i]), 0.02); + for (0..4) |k| { + var expected: f64 = 2 * ref[4 + k]; + for (0..4) |j| expected += ref[8 + 4 * j + k] * fp(stream.slice(u16, 16)[j * 4 + i]); + try close(expected, fp(y.slice(u16, 16)[k * 4 + i]), 0.04); + try std.testing.expectEqual(bf(2), expanded.slice(u16, 16)[k * 4 + i]); + } + try std.testing.expectEqual(@as(f32, if (i == 0) 1 else 0), initial.slice(f32, 4)[i]); + } +} + +test "V4.1 source single-row HC RMS reaches the retained attention BF16 consumer" { + const Fixture = struct { + embedding_bf16: []const u8, + mixes_f32: []const u8, + scale_f32: []const u8, + base_f32: []const u8, + branch_bf16_bits: [2]u16, + residual_bf16_bits: [8]u16, + expected_source_tree_target_bf16_bits: [2]u16, + }; + const fixture = try std.json.parseFromSlice(Fixture, std.testing.allocator, @embedFile("fixtures/hc_attention_rms.json"), .{ .ignore_unknown_fields = true }); + defer fixture.deinit(); + const f = fixture.value; + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try mhc.Pipelines.load(device); + defer p.deinit(); + var b = Buffers{ .device = device }; + defer b.deinit(); + const width = 5120; + const stream = try b.get(u16, 4 * width); + const xs = stream.slice(u16, 4 * width); + try std.base64.standard.Decoder.decode(std.mem.sliceAsBytes(xs[0..width]), f.embedding_bf16); + for (1..4) |copy| @memcpy(xs[copy * width ..][0..width], xs[0..width]); + const inv = try b.get(f32, 1); + const mixes = try b.get(f32, 24); + const scale = try b.get(f32, 3); + const base = try b.get(f32, 24); + try std.base64.standard.Decoder.decode(std.mem.sliceAsBytes(mixes.slice(f32, 24)), f.mixes_f32); + try std.base64.standard.Decoder.decode(std.mem.sliceAsBytes(scale.slice(f32, 3)), f.scale_f32); + try std.base64.standard.Decoder.decode(std.mem.sliceAsBytes(base.slice(f32, 24)), f.base_f32); + const pre = try b.get(f32, 4); + const post = try b.get(f32, 4); + const comb = try b.get(f32, 16); + const branch = try b.get(u16, 2); + const residual = try b.get(u16, 8); + const output = try b.get(u16, 8); + @memcpy(branch.slice(u16, 2), &f.branch_bf16_bits); + @memcpy(residual.slice(u16, 8), &f.residual_bf16_bits); + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + // Recompute statistic, coefficients and post products from the retained projection; never inject an oracle inverse. + mhc.encodeRms(e, &p, view(stream), view(inv), .{ .rows = 1, .width = width, .stream_stride = 4 * width, .eps = 1e-20 }); + mhc.encodeSplit(e, &p, view(mixes), view(inv), view(scale), view(base), view(pre), view(post), view(comb), .{ .rows = 1 }); + mhc.encodePost(e, &p, view(branch), view(residual), view(post), view(comb), view(output), .{ .rows = 1, .width = 2, .branch_stride = 2, .residual_stride = 8, .y_stride = 8 }); + try done(cb, e); + // The former reduction produces [14091,46835], despite unchanged branch BF16 bits. + for (0..2) |i| try std.testing.expectEqual(f.expected_source_tree_target_bf16_bits[i], output.slice(u16, 8)[6 + i]); +} + +test "V4.1 hyperconnections preserve FP32 products and residual boundaries under cancellation" { + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try mhc.Pipelines.load(device); + defer p.deinit(); + var b = Buffers{ .device = device }; + defer b.deinit(); + const stream = try b.get(u16, 8); + @memcpy(stream.slice(u16, 8), &[_]u16{ bf(3), bf(3), bf(0), bf(0), bf(1), bf(1), bf(1), bf(1) }); + const previous = try b.get(f32, 8); + @memcpy(previous.slice(f32, 8), &[_]f32{ -1, 1 + 0x1p-23, 0, 0, 1, 0x1p-24, 0x1p-24, -1 }); + const collapsed = try b.get(u16, 2); + const branch = try b.get(u16, 2); + @memcpy(branch.slice(u16, 2), &[_]u16{ bf(0), bf(1) }); + const post = try b.get(f32, 8); + post.slice(f32, 8)[4] = -1; + const comb = try b.get(f32, 32); + comb.slice(f32, 32)[0] = -1; + comb.slice(f32, 32)[4] = 1 + 0x1p-23; + comb.slice(f32, 32)[16] = 1; + comb.slice(f32, 32)[20] = 0x1p-24; + comb.slice(f32, 32)[24] = 0x1p-24; + const output = try b.get(u16, 8); + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + mhc.encodeCollapse(e, &p, view(stream), view(previous), view(collapsed), .{ .rows = 2, .width = 1, .stream_stride = 4, .y_stride = 1 }); + mhc.encodePost(e, &p, view(branch), view(stream), view(post), view(comb), view(output), .{ .rows = 2, .width = 1, .branch_stride = 1, .residual_stride = 4, .y_stride = 4 }); + try done(cb, e); + try std.testing.expectEqual(bf(0x1p-21), collapsed.slice(u16, 2)[0]); + try std.testing.expectEqual(bf(0), collapsed.slice(u16, 2)[1]); + try std.testing.expectEqual(bf(0x1p-21), output.slice(u16, 8)[0]); + try std.testing.expectEqual(bf(0), output.slice(u16, 8)[4]); + branch.slice(u16, 2)[0] = bf(3); + post.slice(f32, 8)[0] = 1 + 0x1p-23; + comb.slice(f32, 32)[4] = 0; + const branch_cb = queue.commandBuffer(); + const branch_enc = branch_cb.compute(.serial); + mhc.encodePost(branch_enc, &p, view(branch), view(stream), view(post), view(comb), view(output), .{ .rows = 1, .width = 1, .branch_stride = 1, .residual_stride = 4, .y_stride = 4 }); + try done(branch_cb, branch_enc); + try std.testing.expectEqual(bf(0x1p-21), output.slice(u16, 8)[0]); +} diff --git a/zig/src/families/deepseek_v41/ops.zig b/zig/src/families/deepseek_v41/ops.zig new file mode 100644 index 0000000000..68fd169857 --- /dev/null +++ b/zig/src/families/deepseek_v41/ops.zig @@ -0,0 +1,110 @@ +//! Trusted plaintext encoders: callers validate nonzero shapes/storage; strides count elements, EXL3 is separate. +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +pub const View = @import("view.zig").View; + +pub const Pipelines = struct { + rms: mtl.Pipeline, + dense: mtl.Pipeline, + embedding: mtl.Pipeline, + swiglu: mtl.Pipeline, + combine: mtl.Pipeline, + + pub fn load(device: mtl.Device) !Pipelines { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + var options = mtl.CompileOptions.mlx(); + options.math_functions = .precise; + const lib = try mtl.Library.fromSource(device, sources.deepseek_v41_ops, options); + defer lib.deinit(); + const rms = try mtl.Pipeline.init(device, lib, "dsv41_rms", false); + errdefer rms.deinit(); + const dense = try mtl.Pipeline.init(device, lib, "dsv41_dense", false); + errdefer dense.deinit(); + const embedding = try mtl.Pipeline.init(device, lib, "dsv41_embedding", false); + errdefer embedding.deinit(); + const swiglu = try mtl.Pipeline.init(device, lib, "dsv41_swiglu", false); + errdefer swiglu.deinit(); + const combine = try mtl.Pipeline.init(device, lib, "dsv41_combine", false); + return .{ .rms = rms, .dense = dense, .embedding = embedding, .swiglu = swiglu, .combine = combine }; + } + pub fn deinit(p: *Pipelines) void { + p.rms.deinit(); + p.dense.deinit(); + p.embedding.deinit(); + p.swiglu.deinit(); + p.combine.deinit(); + } +}; +fn bind(e: mtl.ComputeEncoder, refs: anytype) void { + inline for (refs, 0..) |r, i| e.setBuffer(r.buffer, r.offset, i); +} +pub const Rms = struct { rows: u32, width: u32, x_stride: u32, y_stride: u32, eps: f32, weight_f32: bool = false }; +/// BF16 x/y [rows,width], BF16 or FP32 weight [width]; sums, mean and affine product are FP32. +pub fn encodeRms(e: mtl.ComputeEncoder, p: *const Pipelines, x: View, weight: View, y: View, d: Rms) void { + e.setPipeline(p.rms); + bind(e, .{ x, weight, y }); + const D = extern struct { rows: u32, width: u32, x_stride: u32, y_stride: u32, eps: f32, weight_f32: u32 }; + e.setValue(D{ .rows = d.rows, .width = d.width, .x_stride = d.x_stride, .y_stride = d.y_stride, .eps = d.eps, .weight_f32 = @intFromBool(d.weight_f32) }, 3); + e.dispatchGroups(mtl.Size.of(d.rows, 1, 1), mtl.Size.of(256, 1, 1)); +} +pub const Dtype = enum(u32) { bf16 = 0, f32 = 1, f16 = 2 }; +pub const Dense = struct { + rows: u32, + k: u32, + n: u32, + groups: u32 = 1, + x_stride: u32, + w_stride: u32, + y_stride: u32, + x_dtype: Dtype = .bf16, + w_dtype: Dtype = .bf16, + y_dtype: Dtype = .bf16, +}; +/// x[rows,groups,k], w[groups,n,k], y[rows,groups,n]; packed groups, padded token pitches, FP32 sums, final BF16. +pub fn encodeDense(e: mtl.ComputeEncoder, p: *const Pipelines, x: View, w: View, y: View, d: Dense) void { + e.setPipeline(p.dense); + bind(e, .{ x, w, y }); + e.setValue([10]u32{ d.rows, d.k, d.n, d.groups, d.x_stride, d.w_stride, d.y_stride, @backingInt(d.x_dtype), @backingInt(d.w_dtype), @backingInt(d.y_dtype) }, 3); + e.dispatchGroups(mtl.Size.of((d.n + 15) / 16, d.rows, d.groups), mtl.Size.of(128, 1, 1)); +} +pub const Embedding = extern struct { rows: u32, width: u32, table_stride: u32, y_stride: u32, copies: u32 = 4 }; +/// ids[rows] u32 must be in vocabulary; table[vocab,width] BF16 -> y[rows,copies,width] BF16. +pub fn encodeEmbedding(e: mtl.ComputeEncoder, p: *const Pipelines, ids: View, table: View, y: View, d: Embedding) void { + e.setPipeline(p.embedding); + bind(e, .{ ids, table, y }); + e.setValue(d, 3); + e.dispatchThreads(mtl.Size.of(d.width, d.rows, 1), mtl.Size.of(256, 1, 1)); +} +pub const Activation = struct { + rows: u32, + width: u32, + input_stride: u32, + output_stride: u32, + weighted: bool = false, + input_f32: bool = false, +}; +/// Separate gate/up; apply FP32 weights before BF16 cast/down. Unweighted calls borrow but never read weights. +pub fn encodeSwiGLU(e: mtl.ComputeEncoder, p: *const Pipelines, gate: View, up: View, weights: View, y: View, d: Activation) void { + e.setPipeline(p.swiglu); + bind(e, .{ gate, up, weights, y }); + e.setValue([6]u32{ d.rows, d.width, d.input_stride, d.output_stride, @intFromBool(d.weighted), @intFromBool(d.input_f32) }, 4); + e.dispatchThreads(mtl.Size.of(d.width, d.rows, 1), mtl.Size.of(256, 1, 1)); +} +pub const Combine = struct { + rows: u32, + width: u32, + routed_row_stride: u32, + slot_stride: u32, + shared_stride: u32, + y_stride: u32, + input_f32: bool = false, + id_stride: u32 = 6, +}; +/// Sum six weighted down slots by ascending expert ID in FP32, then shared; final BF16, invalid IDs yield NaN. +pub fn encodeCombine(e: mtl.ComputeEncoder, p: *const Pipelines, routed: View, shared: View, ids: View, y: View, d: Combine) void { + e.setPipeline(p.combine); + bind(e, .{ routed, shared, ids, y }); + e.setValue([8]u32{ d.rows, d.width, d.routed_row_stride, d.slot_stride, d.shared_stride, d.y_stride, @intFromBool(d.input_f32), d.id_stride }, 4); + e.dispatchThreads(mtl.Size.of(d.width, d.rows, 1), mtl.Size.of(256, 1, 1)); +} diff --git a/zig/src/families/deepseek_v41/ops_check.zig b/zig/src/families/deepseek_v41/ops_check.zig new file mode 100644 index 0000000000..bfe08f4dd8 --- /dev/null +++ b/zig/src/families/deepseek_v41/ops_check.zig @@ -0,0 +1,191 @@ +//! Independent GPU consumer fixtures; pinned official inference equations and provenance in THIRD_PARTY_NOTICES.md. +const std = @import("std"); +const mtl = @import("metal"); +const ops = @import("ops.zig"); +const support = @import("test_support.zig"); +const bf = support.bf; +const fp = support.fp; +const view = support.view; +const Buffers = support.Buffers; +const done = support.done; +const close = support.close; + +test "V4.1 BF16 RMS, F16 dense row independence, grouped wo_a, embedding, clamped weighted SwiGLU and shared combine" { + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try ops.Pipelines.load(device); + defer p.deinit(); + var b = Buffers{ .device = device }; + defer b.deinit(); + const x = try b.get(u16, 16); + const w = try b.get(u16, 4); + const norm = try b.get(u16, 16); + const xs = x.slice(u16, 16); + const ws = w.slice(u16, 4); + for (0..2) |r| for (0..4) |i| { + xs[r * 8 + i] = bf(@as(f32, @floatFromInt(1 + i + r))); + }; + for (ws, 0..) |*v, i| v.* = bf(@as(f32, @floatFromInt(i + 1)) * 0.5); + const f16w = try b.get(f16, 8); + const fws = f16w.slice(f16, 8); + for (fws, 0..) |*v, i| v.* = @floatCast(@as(f32, @floatFromInt(i + 1)) * 0.125 + 0.0009765625); + const dense = try b.get(f32, 8); + const dense_one = try b.get(f32, 2); + const grouped = try b.get(u16, 8); + const gw = try b.get(u16, 8); + for (gw.slice(u16, 8), 0..) |*v, i| v.* = bf(@as(f32, @floatFromInt(i + 1)) * 0.25); + const ids = try b.get(u32, 2); + ids.slice(u32, 2)[0] = 1; + ids.slice(u32, 2)[1] = 0; + const embed = try b.get(u16, 32); + const gate = try b.get(f32, 8); + const up = try b.get(f32, 8); + @memcpy(gate.slice(f32, 8), &[_]f32{ -100, -12, 1, 100, -2, 0, 8, 12 }); + @memcpy(up.slice(f32, 8), &[_]f32{ -100, -12, 1, 100, 12, -20, 0.5, 1 }); + const weights = try b.get(f32, 2); + @memcpy(weights.slice(f32, 2), &[_]f32{ 0.25, 1.25 }); + const act = try b.get(u16, 8); + const routed = try b.get(u16, 48); + const combine_ids = try b.get(u32, 12); + @memcpy(combine_ids.slice(u32, 12), &[_]u32{ 5, 3, 1, 4, 2, 0, 5, 3, 1, 4, 2, 0 }); + const shared = try b.get(u16, 8); + const combined = try b.get(u16, 8); + const values = [_]f32{ 1, 33554432, -33554432, 2, 4, 8 }; + for (routed.slice(u16, 48), 0..) |*v, i| v.* = bf(values[(i % 24) / 4]); + @memset(shared.slice(u16, 8), bf(2.5)); + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + ops.encodeRms(e, &p, view(x), view(w), view(norm), .{ .rows = 2, .width = 4, .x_stride = 8, .y_stride = 8, .eps = 1e-20 }); + ops.encodeDense(e, &p, view(x), view(f16w), view(dense), .{ .rows = 2, .k = 4, .n = 2, .x_stride = 8, .w_stride = 4, .y_stride = 4, .w_dtype = .f16, .y_dtype = .f32 }); + ops.encodeDense(e, &p, view(x).shifted(16), view(f16w), view(dense_one), .{ .rows = 1, .k = 4, .n = 2, .x_stride = 8, .w_stride = 4, .y_stride = 2, .w_dtype = .f16, .y_dtype = .f32 }); + ops.encodeDense(e, &p, view(x), view(gw), view(grouped), .{ .rows = 2, .k = 2, .n = 2, .groups = 2, .x_stride = 8, .w_stride = 2, .y_stride = 4 }); + ops.encodeEmbedding(e, &p, view(ids), view(x), view(embed), .{ .rows = 2, .width = 4, .table_stride = 8, .y_stride = 16 }); + ops.encodeSwiGLU(e, &p, view(gate), view(up), view(weights), view(act), .{ .rows = 2, .width = 4, .input_stride = 4, .output_stride = 4, .weighted = true, .input_f32 = true }); + ops.encodeCombine(e, &p, view(routed), view(shared), view(combine_ids), view(combined), .{ .rows = 2, .width = 4, .routed_row_stride = 24, .slot_stride = 4, .shared_stride = 4, .y_stride = 4 }); + try done(cb, e); + for (0..2) |r| { + var sum: f64 = 0; + for (0..4) |i| sum += @as(f64, fp(xs[r * 8 + i])) * fp(xs[r * 8 + i]); + for (0..4) |i| try close(@as(f64, fp(xs[r * 8 + i])) / @sqrt(sum / 4 + 1e-20) * fp(ws[i]), fp(norm.slice(u16, 16)[r * 8 + i]), 0.04); + for (0..2) |j| { + var ref: f64 = 0; + for (0..4) |k| ref += @as(f64, fp(xs[r * 8 + k])) * @as(f64, fws[j * 4 + k]); + try close(ref, dense.slice(f32, 8)[r * 4 + j], 1e-6); + } + for (0..2) |g| for (0..2) |j| { + var ref: f32 = 0; + for (0..2) |k| ref += fp(xs[r * 8 + g * 2 + k]) * fp(gw.slice(u16, 8)[g * 4 + j * 2 + k]); + try std.testing.expectEqual(bf(ref), grouped.slice(u16, 8)[r * 4 + g * 2 + j]); + }; + for (0..4) |copy| for (0..4) |i| try std.testing.expectEqual(xs[(1 - r) * 8 + i], embed.slice(u16, 32)[r * 16 + copy * 4 + i]); + for (0..4) |i| { + const k = r * 4 + i; + const a: f64 = @min(gate.slice(f32, 8)[k], 10); + const u: f64 = @max(-10, @min(up.slice(f32, 8)[k], 10)); + const activated = a / (1 + @exp(-a)) * u * weights.slice(f32, 2)[r]; + try close(activated, fp(act.slice(u16, 8)[k]), @abs(activated) * 0.004 + 1e-35); + var ref: f32 = 0; + for ([_]usize{ 5, 2, 4, 1, 3, 0 }) |slot| ref += fp(routed.slice(u16, 48)[r * 24 + slot * 4 + i]); + try std.testing.expectEqual(bf(ref + 2.5), combined.slice(u16, 8)[k]); + } + } + for (0..2) |j| try std.testing.expectEqual(dense.slice(f32, 8)[4 + j], dense_one.slice(f32, 2)[j]); +} + +test "V4.1 CUDA RMS mean projection preserves the BF16 affine tie" { + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try ops.Pipelines.load(device); + defer p.deinit(); + var buffers = Buffers{ .device = device }; + defer buffers.deinit(); + const width = 5120; + const x = try buffers.get(u16, width); + const weight = try buffers.get(f32, width); + const output = try buffers.get(u16, width); + x.slice(u16, width)[0] = bf(3); + @memset(weight.slice(f32, width), @bitCast(@as(u32, 0x3c65de29))); + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + ops.encodeRms(e, &p, view(x), view(weight), view(output), .{ .rows = 1, .width = width, .x_stride = width, .y_stride = width, .eps = 1e-20, .weight_f32 = true }); + try done(cb, e); + // Pinned CUDA RMSNorm.forward returns 1; sum / width crosses to the next BF16 value. + try std.testing.expectEqual(bf(1), output.slice(u16, width)[0]); + for (output.slice(u16, width)[1..]) |v| try std.testing.expectEqual(@as(u16, 0), v); +} + +test "V4.1 FP32 dense retains product and reduction cancellation" { + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try ops.Pipelines.load(device); + defer p.deinit(); + var buffers = Buffers{ .device = device }; + defer buffers.deinit(); + const k = 96; + const n = 3; + const x = try buffers.get(u16, k); + const w = try buffers.get(f32, k * n); + const y = try buffers.get(f32, n); + @memset(x.slice(u16, k), bf(3)); + const weights = w.slice(f32, k * n); + weights[0] = 0x1p25; + weights[32] = 1; + weights[64] = -0x1p25; + weights[k] = 1 + 0x1p-23; + weights[k + 32] = -1; + weights[2 * k] = 0x1p25; + weights[2 * k + 1] = 1; + weights[2 * k + 2] = -0x1p25; + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + ops.encodeDense(enc, &p, view(x), view(w), view(y), .{ .rows = 1, .k = k, .n = n, .x_stride = k, .w_stride = k, .y_stride = n, .w_dtype = .f32, .y_dtype = .f32 }); + try done(cb, enc); + for (0..n) |column| { + var expected: f64 = 0; + for (0..k) |index| expected += @as(f64, fp(x.slice(u16, k)[index])) * weights[column * k + index]; + try std.testing.expectEqual(@as(f32, @floatCast(expected)), y.slice(f32, n)[column]); + } +} + +test "V4.1 weighted SiLU preserves division rounding before BF16 activation" { + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try ops.Pipelines.load(device); + defer p.deinit(); + var buffers = Buffers{ .device = device }; + defer buffers.deinit(); + const gate = try buffers.get(u16, 1); + const up = try buffers.get(u16, 1); + const weights = try buffers.get(f32, 1); + const output = try buffers.get(u16, 1); + gate.slice(u16, 1)[0] = bf(0.78515625); + up.slice(u16, 1)[0] = bf(-6.96875); + weights.slice(f32, 1)[0] = 0.25; + const cb = queue.commandBuffer(); + const encoder = cb.compute(.serial); + ops.encodeSwiGLU(encoder, &p, view(gate), view(up), view(weights), view(output), .{ .rows = 1, .width = 1, .input_stride = 1, .output_stride = 1, .weighted = true }); + try done(cb, encoder); + // Source SiLU divides before up/route multiplication; reciprocal multiplication crosses this BF16 tie. + try std.testing.expectEqual(bf(-0.94140625), output.slice(u16, 1)[0]); +} + +test { + _ = @import("mhc_test.zig"); + _ = @import("route_test.zig"); +} diff --git a/zig/src/families/deepseek_v41/profiler.zig b/zig/src/families/deepseek_v41/profiler.zig new file mode 100644 index 0000000000..239be77e40 --- /dev/null +++ b/zig/src/families/deepseek_v41/profiler.zig @@ -0,0 +1,93 @@ +//! Family stage labels over the generic same-command-buffer counter recorder; no extra storage or GPU waits. +const std = @import("std"); +const mtl = @import("metal"); +const counters = @import("core").gpu_counters; + +pub const Stage = enum { engram, mhc, attention, moe, head, other }; +pub const stage_count = @typeInfo(Stage).@"enum".field_names.len; +pub const Options = counters.Options; +pub const Aggregate = counters.Aggregate; +pub const Span = counters.Span; + +pub fn class(stage: Stage) counters.Class { + return switch (stage) { + .engram => .embedding, + .mhc => .norm, + .attention => .attention, + .moe, .head => .dense, + .other => .copy, + }; +} +pub fn stageOf(span: Span) !Stage { + if (span.metadata >= stage_count) return error.InvalidProfileStage; + return @fromBackingInt(@intCast(span.metadata)); +} + +/// Unique owner of the recorder; stage metadata and class totals share the recorder's bounded span storage. +pub const Recorder = struct { + recorder: counters.Recorder, + totals: [stage_count]Aggregate = @splat(.{}), + + pub fn init(allocator: std.mem.Allocator, device: mtl.Device, options: Options) !Recorder { + return .{ .recorder = try counters.Recorder.init(allocator, device, options) }; + } + pub fn enabled(self: *const Recorder) bool { + return self.recorder.enabled(); + } + pub fn capability(self: *const Recorder) mtl.counters.Capability { + return self.recorder.capability(); + } + pub fn activeSpan(self: *const Recorder) bool { + return self.recorder.active_span != null; + } + pub fn beginCommand(self: *Recorder, cb: mtl.CommandBuffer) !void { + try self.recorder.beginCommand(cb); + self.totals = @splat(.{}); + } + pub fn beginSpan(self: *Recorder, encoder: *mtl.ComputeEncoder, stage: Stage) !void { + try self.recorder.beginSpan(encoder, class(stage), @intCast(@backingInt(stage))); + } + pub fn endSpan(self: *Recorder, encoder: *mtl.ComputeEncoder) !void { + try self.recorder.endSpan(encoder); + } + pub fn finishEncoding(self: *Recorder, encoder: *mtl.ComputeEncoder) !void { + try self.recorder.finishEncoding(encoder); + } + pub fn cancelEncoding(self: *Recorder, encoder: *mtl.ComputeEncoder) !void { + try self.recorder.cancelEncoding(encoder); + self.totals = @splat(.{}); + } + pub fn discardAfterCompletion(self: *Recorder) !void { + try self.recorder.discardAfterCompletion(); + self.totals = @splat(.{}); + } + pub fn readAfterCompletion(self: *Recorder) ![]const Span { + const spans = try self.recorder.readAfterCompletion(); + var totals: [stage_count]Aggregate = @splat(.{}); + for (spans) |span| { + const total = &totals[@backingInt(try stageOf(span))]; + total.calls += 1; + total.gpu_seconds += span.gpu_seconds; + } + self.totals = totals; + return spans; + } + pub fn aggregates(self: *const Recorder) [stage_count]Aggregate { + return self.totals; + } + pub fn deinit(self: *Recorder) !void { + try self.recorder.deinit(); + self.totals = @splat(.{}); + } +}; + +test "disabled family profiling allocates nothing and leaves no stage receipts" { + var recorder = try Recorder.init(std.testing.failing_allocator, undefined, .{}); + defer recorder.deinit() catch unreachable; + var encoder: mtl.ComputeEncoder = undefined; + try recorder.beginSpan(&encoder, .mhc); + try recorder.endSpan(&encoder); + try std.testing.expect(!recorder.activeSpan()); + try std.testing.expectEqual(@as(usize, 0), (try recorder.readAfterCompletion()).len); + for (recorder.aggregates()) |total| try std.testing.expectEqual(@as(u64, 0), total.calls); +} diff --git a/zig/src/families/deepseek_v41/resident_load.zig b/zig/src/families/deepseek_v41/resident_load.zig new file mode 100644 index 0000000000..807a6debe8 --- /dev/null +++ b/zig/src/families/deepseek_v41/resident_load.zig @@ -0,0 +1,266 @@ +//! Header-first resident loading: no mapped weights, deferred disk lookup or duplicate Engram projections. +const std = @import("std"); +const mtl = @import("metal"); +const st = @import("core").safetensors; +const cp = @import("core").checkpoint_metal; +const paced = @import("core").paced_read; + +pub const Reference = struct { shard: usize, entry: st.Entry }; +pub const Shard = struct { + path: [:0]const u8, + file_bytes: usize, + data_start: usize, + begin: usize, + end: usize, + names: st.Header, +}; + +pub const Plan = struct { + arena: std.heap.ArenaAllocator, + shards: std.ArrayList(Shard) = .empty, + tensors: std.StringHashMapUnmanaged(Reference) = .empty, + bytes: usize = 0, + + pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, table_dir: []const u8) !Plan { + var self = Plan{ .arena = std.heap.ArenaAllocator.init(gpa) }; + errdefer self.deinit(); + const a = self.arena.allocator(); + const index_path = try std.fs.path.join(a, &.{ dir, "model.safetensors.index.json" }); + const index_bytes = try std.Io.Dir.cwd().readFileAlloc(io, index_path, a, .limited(64 << 20)); + const index = try std.json.parseFromSliceLeaky(struct { weight_map: std.json.Value }, a, index_bytes, .{ + .ignore_unknown_fields = true, + }); + if (index.weight_map != .object or index.weight_map.object.count() == 0) return error.BadCheckpointIndex; + var files: std.StringHashMapUnmanaged(void) = .empty; + var entries = index.weight_map.object.iterator(); + while (entries.next()) |item| { + if (item.value_ptr.* != .string) return error.BadCheckpointIndex; + const file = item.value_ptr.string; + if (!safeBasename(file)) return error.BadCheckpointIndex; + try files.put(a, file, {}); + } + const ordered = try a.alloc([]const u8, files.count()); + var file_it = files.keyIterator(); + var n: usize = 0; + while (file_it.next()) |file| : (n += 1) ordered[n] = file.*; + std.mem.sort([]const u8, ordered, {}, less); + var owned: std.StringHashMapUnmanaged([]const u8) = .empty; + entries.reset(); + while (entries.next()) |item| try owned.put(a, item.key_ptr.*, item.value_ptr.string); + for (ordered) |file| try self.add(try std.fs.path.join(a, &.{ dir, file }), null, &owned); + if (self.tensors.count() != index.weight_map.object.count()) return error.CheckpointIndexMismatch; + entries = index.weight_map.object.iterator(); + while (entries.next()) |item| { + const ref = self.tensors.get(item.key_ptr.*) orelse return error.CheckpointIndexMismatch; + if (!std.mem.eql(u8, std.fs.path.basename(self.shards.items[ref.shard].path), item.value_ptr.string)) + return error.CheckpointIndexMismatch; + } + for ([_]struct { layer: u32, file: []const u8, rows: usize }{ + .{ .layer = 1, .file = "model-00047-of-00048.safetensors", .rows = 384006168 }, + .{ .layer = 14, .file = "model-00048-of-00048.safetensors", .rows = 384016682 }, + }) |table| { + const prefix = try std.fmt.allocPrint(a, "layers.{d}.engram.embed.", .{table.layer}); + try self.add(try std.fs.path.join(a, &.{ table_dir, table.file }), prefix, null); + const weight = try std.fmt.allocPrint(a, "{s}weight", .{prefix}); + const scale = try std.fmt.allocPrint(a, "{s}scale", .{prefix}); + try self.require(weight, .f8_e4m3, &.{ table.rows, 256 }); + try self.require(scale, .f8_e8m0, &.{ table.rows, 8 }); + if (self.shards.items[self.shards.items.len - 1].names.count() != 2) return error.UnexpectedEngramTensor; + } + return self; + } + + pub fn deinit(self: *Plan) void { + self.arena.deinit(); + self.* = undefined; + } + + pub fn get(self: *const Plan, name: []const u8) !st.Entry { + return (self.tensors.get(name) orelse return error.MissingTensor).entry; + } + + pub fn has(self: *const Plan, name: []const u8) bool { + return self.tensors.contains(name); + } + + pub fn require(self: *const Plan, name: []const u8, dtype: st.DType, shape: []const usize) !void { + const e = try self.get(name); + if (e.dtype != dtype or e.rank != shape.len or !std.mem.eql(usize, e.shape[0..e.rank], shape)) { + std.log.err("unsupported tensor {s}: dtype {s}, shape {any}", .{ name, @tagName(e.dtype), e.shape[0..e.rank] }); + return error.UnsupportedTensor; + } + } + + fn add(self: *Plan, path: []const u8, prefix: ?[]const u8, owned: ?*const std.StringHashMapUnmanaged([]const u8)) !void { + const a = self.arena.allocator(); + const own_path = try a.dupeSentinel(u8, path, 0); + const fd = std.c.open(own_path, .{ .ACCMODE = .RDONLY }); + if (fd < 0) return error.OpenFailed; + defer _ = std.c.close(fd); + const length = std.c.lseek(fd, 0, std.c.SEEK.END); + if (length < 8) return error.BadSafetensors; + const file_bytes: usize = @intCast(length); + var head: [8]u8 = undefined; + try readAll(fd, &head, 0); + const header_bytes = std.mem.readInt(u64, &head, .little); + if (header_bytes > @min(file_bytes - 8, 64 << 20)) return error.BadSafetensors; + const json = try a.alloc(u8, @intCast(header_bytes)); + try readAll(fd, json, 8); + const data_start: usize = 8 + @as(usize, @intCast(header_bytes)); + const names = try st.parseHeaderPrefix(a, json, file_bytes - data_start, prefix); + if (names.count() == 0) return error.MissingTensor; + var begin: usize = if (prefix == null) 0 else file_bytes - data_start; + var end: usize = if (prefix == null) file_bytes - data_start else 0; + for (names.keys(), names.values()) |name, entry| { + // A rank-split pack duplicates whole-load tensors beside the shards; the index decides the owner file. + if (owned) |map| { + const owner = map.get(name) orelse continue; + if (!std.mem.eql(u8, owner, std.fs.path.basename(path))) continue; + } + begin = @min(begin, entry.begin); + end = @max(end, entry.end); + const slot = try self.tensors.getOrPut(a, name); + if (slot.found_existing) { + return error.DuplicateTensor; + } + slot.value_ptr.* = .{ .shard = self.shards.items.len, .entry = entry }; + } + if (prefix != null) { + var selected_bytes: usize = 0; + for (names.values()) |entry| selected_bytes = try std.math.add(usize, selected_bytes, entry.end - entry.begin); + if (selected_bytes != end - begin) return error.NoncontiguousEngramTables; + } + self.bytes = try std.math.add(usize, self.bytes, end - begin); + try self.shards.append(a, .{ .path = own_path, .file_bytes = file_bytes, .data_start = data_start, .begin = begin, .end = end, .names = names }); + } + + /// Only called after architecture validation and the complete memory admission calculation. + pub fn materialize(self: *const Plan, gpa: std.mem.Allocator, device: mtl.Device, residency: mtl.ResidencySet, queue: mtl.Queue) !cp.Checkpoint { + var out = cp.Checkpoint.init(gpa); + errdefer out.deinit(); + const stage = try paced.stage(gpa, 32 << 20); + defer gpa.free(stage); + var warm = mtl.keepalive.Target{ .queue = queue, .sets = @as(*const [1]mtl.ResidencySet, &residency) }; + var last_tick: f64 = -std.math.inf(f64); + var short: paced.Memory = undefined; + for (self.shards.items, 0..) |shard, i| { + const fd = std.c.open(shard.path, .{ .ACCMODE = .RDONLY }); + if (fd < 0) return error.OpenFailed; + defer _ = std.c.close(fd); + if (std.c.fcntl(fd, 48, @as(c_int, 1)) != 0) return error.UncachedReadUnavailable; + if (std.c.lseek(fd, 0, std.c.SEEK.END) != shard.file_bytes) return error.CheckpointChanged; + const bytes = shard.end - shard.begin; + const buffer = try device.buffer(@max(bytes, 16), mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + var owned = false; + errdefer if (!owned) buffer.deinit(); + try out.own(buffer, bytes); + owned = true; + residency.add(buffer); + // Unwinding must deregister before `out.deinit` frees the buffer, or the set points at released storage. + var registered = false; + errdefer if (registered) { + residency.remove(buffer); + residency.commit(); + }; + registered = true; + residency.commit(); + residency.requestResidency(); + var copied: usize = 0; + while (copied < bytes) { + const count = @min(bytes - copied, stage.len - paced.alignment); + paced.read(fd, buffer.contents()[copied..][0..count], shard.data_start + shard.begin + copied, stage, &short) catch |err| { + if (err == error.MemoryShort) { + var msg: [512]u8 = undefined; + std.log.err("deepseek-v41: {s}", .{paced.refusal(&msg, short, paced.floor())}); + } + return err; + }; + copied += count; + const now = mtl.clock.seconds(); + if (now - last_tick >= 0.4) { + mtl.keepalive.Target.tick(&warm); + last_tick = now; + } + } + for (shard.names.keys(), shard.names.values()) |name, e| { + try out.set(name, .{ .buffer = buffer, .offset = e.begin - shard.begin, .bytes = e.end - e.begin, .dtype = e.dtype, .shape = e.shape, .rank = e.rank }); + } + std.log.info("resident checkpoint shard {d}/{d}: {d} bytes, residency set {d} bytes", .{ i + 1, self.shards.items.len, bytes, residency.allocatedSize() }); + } + residency.commit(); + residency.requestResidency(); + return out; + } +}; + +fn readAll(fd: std.c.fd_t, dest: []u8, at: usize) !void { + var done: usize = 0; + while (done < dest.len) { + const n = std.c.pread(fd, dest.ptr + done, dest.len - done, @intCast(at + done)); + if (n <= 0) return error.ShortRead; + done += @intCast(n); + } +} + +fn safeBasename(s: []const u8) bool { + return s.len > 0 and !std.mem.eql(u8, s, ".") and !std.mem.eql(u8, s, "..") and + std.mem.indexOfAny(u8, s, "/\\") == null and std.mem.indexOfScalar(u8, s, 0) == null and + std.mem.endsWith(u8, s, ".safetensors"); +} + +fn less(_: void, a: []const u8, b: []const u8) bool { + return std.mem.lessThan(u8, a, b); +} + +test "selected table plan excludes projections and preserves tensor byte offsets" { + const a = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const header = + \\{"padding":{"dtype":"U8","shape":[3],"data_offsets":[0,3]}, + \\ "engram.embed.weight":{"dtype":"U8","shape":[4],"data_offsets":[3,7]}, + \\ "engram.embed.scale":{"dtype":"U8","shape":[2],"data_offsets":[7,9]}, + \\ "engram.wkv":{"dtype":"U8","shape":[1],"data_offsets":[9,10]}} + ; + var data: [8 + header.len + 10]u8 = undefined; + std.mem.writeInt(u64, data[0..8], header.len, .little); + @memcpy(data[8..][0..header.len], header); + @memset(data[8 + header.len ..], 0); + try tmp.dir.writeFile(io, .{ .sub_path = "table.safetensors", .data = &data }); + const path = try std.fmt.allocPrint(a, ".zig-cache/tmp/{s}/table.safetensors", .{tmp.sub_path}); + defer a.free(path); + var plan = Plan{ .arena = std.heap.ArenaAllocator.init(a) }; + defer plan.deinit(); + try plan.add(path, "engram.embed.", null); + try std.testing.expectEqual(@as(usize, 6), plan.bytes); + try std.testing.expectEqual(@as(usize, 3), plan.shards.items[0].begin); + try std.testing.expectEqual(@as(usize, 9), plan.shards.items[0].end); + try std.testing.expectEqual(@as(usize, 7), (try plan.get("engram.embed.scale")).begin); + try std.testing.expectError(error.MissingTensor, plan.get("engram.wkv")); + try std.testing.expectError(error.DuplicateTensor, plan.add(path, "engram.embed.", null)); +} + +test "selected table payload must be contiguous and indexed files cannot escape checkpoint directory" { + const a = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const header = + \\{"engram.embed.weight":{"dtype":"U8","shape":[4],"data_offsets":[0,4]}, + \\ "projection":{"dtype":"U8","shape":[3],"data_offsets":[4,7]}, + \\ "engram.embed.scale":{"dtype":"U8","shape":[2],"data_offsets":[7,9]}} + ; + var data: [8 + header.len + 9]u8 = undefined; + std.mem.writeInt(u64, data[0..8], header.len, .little); + @memcpy(data[8..][0..header.len], header); + @memset(data[8 + header.len ..], 0); + try tmp.dir.writeFile(io, .{ .sub_path = "gap.safetensors", .data = &data }); + const path = try std.fmt.allocPrint(a, ".zig-cache/tmp/{s}/gap.safetensors", .{tmp.sub_path}); + defer a.free(path); + var plan = Plan{ .arena = std.heap.ArenaAllocator.init(a) }; + defer plan.deinit(); + try std.testing.expectError(error.NoncontiguousEngramTables, plan.add(path, "engram.embed.", null)); + for ([_][]const u8{ "../x.safetensors", "/x.safetensors", "a\\x.safetensors", ".", "..", "x.safetensors\x00", "x.json", "" }) |name| try std.testing.expect(!safeBasename(name)); + try std.testing.expect(safeBasename("model-00001-of-00039.safetensors")); +} diff --git a/zig/src/families/deepseek_v41/rope.zig b/zig/src/families/deepseek_v41/rope.zig new file mode 100644 index 0000000000..cc3419f7c3 --- /dev/null +++ b/zig/src/families/deepseek_v41/rope.zig @@ -0,0 +1,134 @@ +//! Adjacent-pair BF16 RoPE using resident, source-style complex rotation factors. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +const View = @import("view.zig").View; + +pub const Spec = struct { + theta: f32, + factor: f32 = 16, + original: u32 = 0, + beta_fast: f32 = 32, + beta_slow: f32 = 1, + + fn validate(s: Spec) !void { + if (!std.math.isFinite(s.theta) or s.theta <= 1 or + !std.math.isFinite(s.factor) or s.factor <= 0 or + !std.math.isFinite(s.beta_fast) or s.beta_fast <= 0 or + !std.math.isFinite(s.beta_slow) or s.beta_slow <= 0) + return error.InvalidRotarySpec; + } +}; +pub const Spectrum = enum(u32) { normal, compressed }; +pub const Params = extern struct { + rows: u32, + heads: u32, + width: u32, + rotary: u32, + stride: u32, + position: u32, + step: u32 = 1, + inverse: u32 = 0, + spectrum: Spectrum = .normal, + + pub fn validate(p: Params) !void { + if (p.rows == 0 or p.heads == 0 or p.rotary == 0 or p.rotary % 2 != 0 or + p.rotary > p.width or @as(u64, p.stride) < @as(u64, p.heads) * p.width or + p.step == 0 or p.inverse > 1 or + @as(u64, p.rows - 1) * p.stride + @as(u64, p.heads) * p.width > std.math.maxInt(u32)) + return error.InvalidRotaryShape; + } +}; + +// Keep FP32 tensor/angle boundaries, FP64 corrected_dim logs, and FP32-rounded polar unaries; never per-head recompute. +fn fillFactors(out: []f32, capacity: u32, rotary: u32, spec: Spec) void { + @setFloatMode(.strict); + const dim: f32 = @floatFromInt(rotary); + const pairs: usize = rotary / 2; + var low: f32 = 0; + var high: f32 = 0; + if (spec.original > 0) { + const original: f64 = @floatFromInt(spec.original); + const dim64: f64 = @floatFromInt(rotary); + const denominator = 2 * @log(@as(f64, spec.theta)); + low = @floatCast(@max(@floor(dim64 * @log(original / + (@as(f64, spec.beta_fast) * 2 * std.math.pi)) / denominator), 0)); + high = @floatCast(@min(@ceil(dim64 * @log(original / + (@as(f64, spec.beta_slow) * 2 * std.math.pi)) / denominator), dim64 - 1)); + } + for (0..pairs) |pair| { + const exponent: f32 = @as(f32, @floatFromInt(pair * 2)) / dim; + const power: f32 = @floatCast(std.math.pow(f64, spec.theta, exponent)); + var frequency: f32 = 1 / power; + if (spec.original > 0) { + const ramp: f32 = std.math.clamp((@as(f32, @floatFromInt(pair)) - low) / + @max(high - low, @as(f32, 1e-3)), 0, 1); + const smooth: f32 = 1 - ramp; + const scaled: f32 = frequency / spec.factor; + const slow: f32 = scaled * (1 - smooth); + const fast: f32 = frequency * smooth; + frequency = slow + fast; + } + for (0..capacity) |position| { + const angle: f32 = @as(f32, @floatFromInt(position)) * frequency; + const at = (position * pairs + pair) * 2; + out[at] = @floatCast(@cos(@as(f64, angle))); + out[at + 1] = @floatCast(@sin(@as(f64, angle))); + } + } +} + +pub const Pipelines = struct { + rotate: mtl.Pipeline, + factors: mtl.Buffer, + capacity: u32, + rotary: u32, + + pub fn residentBytes(capacity: u32, rotary: u32) usize { + return 2 * @as(usize, capacity) * (rotary / 2) * 2 * @sizeOf(f32); + } + + pub fn load(device: mtl.Device, capacity: u32, rotary: u32, normal: Spec, compressed: Spec) !Pipelines { + if (capacity == 0 or rotary == 0 or rotary % 2 != 0 or + @as(u64, capacity) * rotary > std.math.maxInt(usize) / (2 * @sizeOf(f32))) + return error.InvalidRotaryShape; + try normal.validate(); + try compressed.validate(); + var options = mtl.CompileOptions.mlx(); + options.math_functions = .precise; + const lib = try mtl.Library.fromSource(device, sources.deepseek_v41_rope, options); + defer lib.deinit(); + const rotate = try mtl.Pipeline.init(device, lib, "dsv41_rope", false); + errdefer rotate.deinit(); + const factors = try device.buffer(residentBytes(capacity, rotary), mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + errdefer factors.deinit(); + const count = @as(usize, capacity) * rotary; + const values = factors.slice(f32, 2 * count); + fillFactors(values[0..count], capacity, rotary, normal); + fillFactors(values[count..], capacity, rotary, compressed); + return .{ .rotate = rotate, .factors = factors, .capacity = capacity, .rotary = rotary }; + } + pub fn deinit(p: *Pipelines) void { + p.factors.deinit(); + p.rotate.deinit(); + } + pub fn registerResidency(p: *const Pipelines, set: mtl.ResidencySet) void { + set.add(p.factors); + } + pub fn encode(p: *const Pipelines, enc: mtl.ComputeEncoder, x: View, args: Params) !void { + try args.validate(); + const last = @as(u64, args.position) + @as(u64, args.rows - 1) * args.step; + if (args.rotary != p.rotary or last >= p.capacity) return error.InvalidRotaryShape; + try x.requireRows(args.rows, args.stride, @as(usize, args.heads) * args.width, 2); + const offset = switch (args.spectrum) { + .normal => 0, + .compressed => @as(usize, p.capacity) * p.rotary * @sizeOf(f32), + }; + enc.setPipeline(p.rotate); + enc.setBuffer(x.buffer, x.offset, 0); + enc.setValue(args, 1); + enc.setBuffer(p.factors, offset, 2); + enc.dispatchThreads(mtl.Size.of(args.rotary / 2, args.heads, args.rows), mtl.Size.of(32, 1, 1)); + enc.barrier(); + } +}; diff --git a/zig/src/families/deepseek_v41/route.zig b/zig/src/families/deepseek_v41/route.zig new file mode 100644 index 0000000000..00c1a1ff10 --- /dev/null +++ b/zig/src/families/deepseek_v41/route.zig @@ -0,0 +1,39 @@ +//! Exhaustive 384-expert sqrt-softplus top-six routing. Bias affects selection, never returned weights. +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +pub const View = @import("view.zig").View; +pub const Dtype = @import("ops.zig").Dtype; +pub const Pipelines = struct { + select: mtl.Pipeline, + pub fn load(device: mtl.Device) !Pipelines { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + var options = mtl.CompileOptions.mlx(); + options.math_functions = .precise; + const lib = try mtl.Library.fromSource(device, sources.deepseek_v41_route, options); + defer lib.deinit(); + return .{ .select = try mtl.Pipeline.init(device, lib, "dsv41_route", false) }; + } + pub fn deinit(p: *Pipelines) void { + p.select.deinit(); + } +}; +pub const Shape = struct { + rows: u32, + logit_stride: u32 = 384, + output_stride: u32 = 6, + temperature: f32 = 1, + bias_dtype: Dtype = .f32, +}; +/// FP32 logits[rows,384], bias F16/BF16/FP32 -> u32 IDs and FP32 weights[rows,6]; lower ID wins biased-score ties. + +// Nonzero rows, logit/output pitches >=384/6, finite positive temperature, complete storage are caller invariants. + +// All six execute; weights use unbiased scores / (sum+1e-20) * 1.5. Invalid comparisons emit invalid IDs/NaNs. +pub fn encode(e: mtl.ComputeEncoder, p: *const Pipelines, logits: View, bias: View, ids: View, weights: View, d: Shape) void { + e.setPipeline(p.select); + inline for (.{ logits, bias, ids, weights }, 0..) |r, i| e.setBuffer(r.buffer, r.offset, i); + const D = extern struct { rows: u32, logit_stride: u32, output_stride: u32, temperature: f32, bias_dtype: u32 }; + e.setValue(D{ .rows = d.rows, .logit_stride = d.logit_stride, .output_stride = d.output_stride, .temperature = d.temperature, .bias_dtype = @backingInt(d.bias_dtype) }, 4); + e.dispatchGroups(mtl.Size.of(d.rows, 1, 1), mtl.Size.of(128, 1, 1)); +} diff --git a/zig/src/families/deepseek_v41/route_test.zig b/zig/src/families/deepseek_v41/route_test.zig new file mode 100644 index 0000000000..f653fa9241 --- /dev/null +++ b/zig/src/families/deepseek_v41/route_test.zig @@ -0,0 +1,152 @@ +//! Consumer-visible expert IDs, weights and nonfinite-input routing boundaries. +const std = @import("std"); +const mtl = @import("metal"); +const route = @import("route.zig"); +const ops = @import("ops.zig"); +const support = @import("test_support.zig"); +const view = support.view; +const Buffers = support.Buffers; +const done = support.done; +const close = support.close; +fn scoreReference(x: f64) f64 { + return @sqrt(if (x > 20) x else std.math.log1p(@exp(x))); +} +test "V4.1 exhaustive router F16 correction bias, unbiased normalization, extreme logits and row independence" { + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try route.Pipelines.load(device); + defer p.deinit(); + var b = Buffers{ .device = device }; + defer b.deinit(); + const logits = try b.get(f32, 2 * 389); + const bias = try b.get(f16, 384); + for (0..384) |i| { + logits.slice(f32, 778)[i] = (@as(f32, @floatFromInt(i)) - 190) * 0.5; + logits.slice(f32, 778)[389 + i] = 0; + bias.slice(f16, 384)[i] = if (i >= 17 and i < 23) 100 else 0; + } + const ids = try b.get(u32, 16); + const weights = try b.get(f32, 16); + const ids_one = try b.get(u32, 6); + const weights_one = try b.get(f32, 6); + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + route.encode(e, &p, view(logits), view(bias), view(ids), view(weights), .{ .rows = 2, .logit_stride = 389, .output_stride = 8, .bias_dtype = .f16 }); + route.encode(e, &p, view(logits).shifted(389 * 4), view(bias), view(ids_one), view(weights_one), .{ .rows = 1, .bias_dtype = .f16 }); + try done(cb, e); + for (0..2) |r| { + var used: [384]bool = @splat(false); + var picks: [6]usize = undefined; + var raw: [6]f64 = undefined; + var sum: f64 = 0; + for (0..6) |slot| { + var best: f64 = -std.math.inf(f64); + var pick: usize = 0; + for (0..384) |i| { + const sc = scoreReference(logits.slice(f32, 778)[r * 389 + i]) + @as(f64, bias.slice(f16, 384)[i]); + if (!used[i] and sc > best) { + best = sc; + pick = i; + } + } + used[pick] = true; + picks[slot] = pick; + raw[slot] = scoreReference(logits.slice(f32, 778)[r * 389 + pick]); + sum += raw[slot]; + } + for (0..6) |slot| { + try std.testing.expectEqual(@as(u32, @intCast(picks[slot])), ids.slice(u32, 16)[r * 8 + slot]); + try close(raw[slot] / (sum + 1e-20) * 1.5, weights.slice(f32, 16)[r * 8 + slot], 3e-6); + } + } + for (0..6) |i| { + try std.testing.expectEqual(ids.slice(u32, 16)[8 + i], ids_one.slice(u32, 6)[i]); + try std.testing.expectEqual(weights.slice(f32, 16)[8 + i], weights_one.slice(f32, 6)[i]); + } +} + +test "V4.1 router normalization matches source CUDA six-score FP32 reduction boundaries" { + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try route.Pipelines.load(device); + defer p.deinit(); + var b = Buffers{ .device = device }; + defer b.deinit(); + // Power-of-two sqrt(softplus(score^2)) with logits >20 eliminates unary-rounding ambiguity. + const scores = [2][6]f32{ + .{ 268435456, 16, 16, 16, 16, 16 }, + .{ 2097152, 8, 16, 134217728, 8388608, 8192 }, + }; + // Independent pinned Torch CUDA Gate expected bits; the second row also distinguishes wide normalization. + const expected_bits = [2][6]u32{ + .{ 1069547517, 868220925, 868220925, 868220925, 868220925, 868220925 }, + .{ 1018303405, 867308461, 875697069, 1068635053, 1035080621, 951194541 }, + }; + const logits = try b.get(f32, scores.len * 384); + const bias = try b.get(f32, 384); + const ids = try b.get(u32, scores.len * 6); + const weights = try b.get(f32, scores.len * 6); + @memset(logits.slice(f32, scores.len * 384), -100); + for (scores, 0..) |row, r| { + for (row, 0..) |s, slot| logits.slice(f32, scores.len * 384)[r * 384 + slot] = s * s; + } + // Distinct biases force slot order without topk ties; weights still use unbiased scores. + for (0..6) |slot| bias.slice(f32, 384)[slot] = 1099511627776 - @as(f32, @floatFromInt(slot)) * 68719476736; + const cb = queue.commandBuffer(); + const encoder = cb.compute(.serial); + route.encode(encoder, &p, view(logits), view(bias), view(ids), view(weights), .{ .rows = scores.len, .bias_dtype = .f32 }); + try done(cb, encoder); + for (expected_bits, 0..) |row, r| { + for (row, 0..) |expected, slot| { + try std.testing.expectEqual(@as(u32, @intCast(slot)), ids.slice(u32, scores.len * 6)[r * 6 + slot]); + try std.testing.expectEqual(expected, @as(u32, @bitCast(weights.slice(f32, scores.len * 6)[r * 6 + slot]))); + } + } +} + +test "nonfinite router inputs propagate invalid routes without reading outside the expert arrays" { + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try route.Pipelines.load(device); + defer p.deinit(); + var buffers = Buffers{ .device = device }; + defer buffers.deinit(); + const logits = try buffers.get(f32, 384); + const bias = try buffers.get(f32, 384); + const ids = try buffers.get(u32, 6); + const weights = try buffers.get(f32, 6); + @memset(logits.slice(f32, 384), std.math.nan(f32)); + const cb = queue.commandBuffer(); + const encoder = cb.compute(.serial); + route.encode(encoder, &p, view(logits), view(bias), view(ids), view(weights), .{ .rows = 1 }); + try done(cb, encoder); + for (ids.slice(u32, 6), weights.slice(f32, 6)) |id, weight| { + try std.testing.expectEqual(std.math.maxInt(u32), id); + try std.testing.expect(std.math.isNan(weight)); + } + var combine = try ops.Pipelines.load(device); + defer combine.deinit(); + const routed = try buffers.get(u16, 6 * 4); + const shared = try buffers.get(u16, 4); + const output = try buffers.get(u16, 4); + // Invalid routes skip expert writes; finite slots from a prior chunk must not hide that failure. + @memset(routed.slice(u16, 24), support.bf(2)); + @memset(shared.slice(u16, 4), support.bf(3)); + const combine_cb = queue.commandBuffer(); + const combine_encoder = combine_cb.compute(.serial); + ops.encodeCombine(combine_encoder, &combine, view(routed), view(shared), view(ids), view(output), .{ .rows = 1, .width = 4, .routed_row_stride = 24, .slot_stride = 4, .shared_stride = 4, .y_stride = 4 }); + try done(combine_cb, combine_encoder); + for (output.slice(u16, 4)) |value| try std.testing.expect(std.math.isNan(support.fp(value))); +} diff --git a/zig/src/families/deepseek_v41/test_support.zig b/zig/src/families/deepseek_v41/test_support.zig new file mode 100644 index 0000000000..80bf87e8d2 --- /dev/null +++ b/zig/src/families/deepseek_v41/test_support.zig @@ -0,0 +1,40 @@ +const std = @import("std"); +const mtl = @import("metal"); +const View = @import("view.zig").View; + +pub fn bf(x: f32) u16 { + const u: u32 = @bitCast(x); + return @intCast((u +% 0x7fff +% ((u >> 16) & 1)) >> 16); +} +pub fn fp(x: u16) f32 { + return @bitCast(@as(u32, x) << 16); +} +pub fn view(b: mtl.Buffer) View { + return .{ .buffer = b }; +} +pub const Buffers = struct { + device: mtl.Device, + items: [32]mtl.Buffer = undefined, + count: usize = 0, + + pub fn get(b: *Buffers, comptime T: type, n: usize) !mtl.Buffer { + if (b.count == b.items.len) return error.TooManyTestBuffers; + const x = try b.device.buffer(try std.math.mul(usize, @sizeOf(T), n), mtl.ResourceOptions.shared); + b.items[b.count] = x; + b.count += 1; + @memset(x.slice(T, n), 0); + return x; + } + pub fn deinit(b: *Buffers) void { + for (b.items[0..b.count]) |x| x.deinit(); + } +}; +pub fn done(cb: mtl.CommandBuffer, e: mtl.ComputeEncoder) !void { + e.end(); + cb.commit(); + cb.wait(); + try std.testing.expect(cb.failure() == null); +} +pub fn close(expected: f64, actual: f32, tolerance: f64) !void { + try std.testing.expect(@abs(expected - @as(f64, actual)) <= tolerance); +} diff --git a/zig/src/families/deepseek_v41/view.zig b/zig/src/families/deepseek_v41/view.zig new file mode 100644 index 0000000000..fbaa08ce39 --- /dev/null +++ b/zig/src/families/deepseek_v41/view.zig @@ -0,0 +1,25 @@ +//! A borrowed Metal buffer region; element precision is carried by each operator's explicit descriptor. +const std = @import("std"); +const mtl = @import("metal"); + +pub const View = struct { + buffer: mtl.Buffer, + offset: usize = 0, + + pub fn shifted(self: View, bytes: usize) View { + return .{ .buffer = self.buffer, .offset = self.offset + bytes }; + } + + pub fn require(self: View, bytes: usize, alignment: usize) !void { + if (self.offset % alignment != 0 or self.offset > self.buffer.length() or + bytes > self.buffer.length() - self.offset) return error.InvalidTensorView; + } + + /// Rows are strided elements; only the final row's used width is required. + pub fn requireRows(self: View, rows: usize, stride: usize, width: usize, element_bytes: usize) !void { + if (rows == 0 or width == 0 or stride < width) return error.InvalidTensorView; + const preceding = try std.math.mul(usize, rows - 1, stride); + const elements = try std.math.add(usize, preceding, width); + try self.require(try std.math.mul(usize, elements, element_bytes), element_bytes); + } +}; diff --git a/zig/src/families/deepseek_v41/weights.zig b/zig/src/families/deepseek_v41/weights.zig new file mode 100644 index 0000000000..3be000c53c --- /dev/null +++ b/zig/src/families/deepseek_v41/weights.zig @@ -0,0 +1,348 @@ +//! Strict pinned backbone schema and borrowed resident bindings; only expert descriptor tables are owned. +const std = @import("std"); +const mtl = @import("metal"); +const cp = @import("core").checkpoint_metal; +const st = @import("core").safetensors; +const exl3 = @import("core").exl3_projection; +const load = @import("resident_load.zig"); + +pub const Tensor = cp.Tensor; +pub const Matrix = exl3.Matrix; + +/// A kept-dense repack stores small matrices as a bf16 weight beside the rotation scales. +pub const Dense = struct { weight: Tensor, suh: Tensor, svh: Tensor, k: usize, n: usize }; +pub const AnyMatrix = union(enum) { trellis: Matrix, dense: Dense }; +pub const Hc = struct { projection: Tensor, scale: Tensor, base: Tensor }; +pub const Compressor = struct { wkv: AnyMatrix, wgate: ?AnyMatrix, norm: Tensor }; +pub const Indexer = struct { wq_b: AnyMatrix, wk: ?AnyMatrix, weights_proj: Tensor, k_norm: ?Tensor }; +pub const Attention = struct { + wq_a: Matrix, + wq_b: Matrix, + wkv: Matrix, + wo_b: Matrix, + wo_a: [8]Matrix, + sink: Tensor, + q_norm: Tensor, + kv_norm: Tensor, + compressor: ?Compressor, + indexer: ?Indexer, +}; +/// Projection order is w1 (gate), w3 (up), w2 (down), for shared and routed experts alike. +pub const Moe = struct { + gate: Tensor, + bias: Tensor, + shared: [3]Matrix, + descriptors: [3]mtl.Buffer, +}; +pub const Engram = struct { table: Tensor, scales: Tensor, q_weight: Tensor, k_weight: Tensor, wkv: Matrix }; +pub const Layer = struct { + attn_norm: Tensor, + ffn_norm: Tensor, + hc_attn: Hc, + hc_ffn: Hc, + attn: Attention, + moe: Moe, + engram: ?Engram, +}; +pub const Model = struct { + embed: Tensor, + head: Matrix, + norm: Tensor, + layers: [40]Layer, + + /// Checkpoint buffers must outlive this borrowed model and every submitted command using it. + pub fn deinit(self: *Model) void { + for (&self.layers) |*layer| for (layer.moe.descriptors) |buffer| buffer.deinit(); + self.* = undefined; + } + pub fn registerResidency(self: *const Model, set: mtl.ResidencySet) void { + for (&self.layers) |*layer| for (layer.moe.descriptors) |buffer| set.add(buffer); + } + pub fn residentBytes(_: *const Model) usize { + return sizeofDescriptors(); + } +}; + +/// Admission adds these bytes to load.Plan.bytes; no dequantized weights or duplicate projections are allocated. +pub fn sizeofDescriptors() usize { + return 40 * 3 * 384 * @sizeOf(exl3.ExpertDescriptor); +} + +/// Every executed tensor is checked using headers only, before checkpoint materialization or model allocation. +pub fn validate(plan: *const load.Plan) !void { + try populate(false, plan, null); +} + +/// Bind resident tensors and marker payloads; route descriptors replace redundant host matrix metadata. +pub fn bind(device: mtl.Device, checkpoint: *const cp.Checkpoint) !Model { + var model: Model = undefined; + try populate(true, checkpoint, &model); + var owned: usize = 0; + errdefer for (0..owned) |i| model.layers[i / 3].moe.descriptors[i % 3].deinit(); + var prefix: [128]u8 = undefined; + for (&model.layers, 0..) |*layer, layer_id| { + inline for (.{ "w1", "w3", "w2" }, 0..) |part, projection| { + const bytes = 384 * @sizeOf(exl3.ExpertDescriptor); + const buffer = try device.buffer(bytes, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + layer.moe.descriptors[projection] = buffer; + owned += 1; + const out: [*]exl3.ExpertDescriptor = @ptrCast(@alignCast(buffer.contents())); + for (0..384) |expert| { + var m: Matrix = undefined; + const k: usize = if (projection == 2) 2304 else 5120; + const n: usize = if (projection == 2) 5120 else 2304; + try matrix(true, checkpoint, try named(&prefix, "layers.{d}.ffn.experts.{d}.{s}", .{ layer_id, expert, part }), k, n, &m); + out[expert] = try exl3.descriptor(m); + } + } + } + return model; +} + +/// A borrowed CPU/GPU matrix view for checkpoint qualification, using the same strict storage contract. +pub fn lookup(checkpoint: *const cp.Checkpoint, prefix: []const u8, k: usize, n: usize) !Matrix { + var out: Matrix = undefined; + try matrix(true, checkpoint, prefix, k, n, &out); + return out; +} + +fn named(buffer: []u8, comptime fmt: []const u8, args: anytype) ![]const u8 { + return std.fmt.bufPrint(buffer, fmt, args); +} + +/// A plain tensor the encoders decode as f32, f16 or bf16; the shape must match exactly. +fn plainEither(comptime binding: bool, source: anytype, name: []const u8, shape: []const usize, destination: ?*Tensor) !void { + const tensor = try source.get(name); + if ((tensor.dtype != .f32 and tensor.dtype != .f16 and tensor.dtype != .bf16) or tensor.rank != shape.len or + !std.mem.eql(usize, tensor.shape[0..tensor.rank], shape)) + { + std.log.err("unsupported tensor {s}: {s} {any}", .{ name, @tagName(tensor.dtype), tensor.shape[0..tensor.rank] }); + return error.UnsupportedTensor; + } + if (binding) { + var count: usize = 1; + for (shape) |dim| count = try std.math.mul(usize, count, dim); + if (tensor.bytes != try std.math.mul(usize, count, tensor.dtype.size())) return error.UnsupportedTensor; + if (tensor.offset % tensor.dtype.size() != 0 or tensor.offset > tensor.buffer.length() or + tensor.bytes > tensor.buffer.length() - tensor.offset) return error.InvalidTensorView; + destination.?.* = tensor; + } else if (tensor.begin % 4 != 0) return error.InvalidTensorView; +} + +fn plain(comptime binding: bool, source: anytype, name: []const u8, dtype: st.DType, shape: []const usize, destination: ?*Tensor) !void { + const tensor = try source.get(name); + if (tensor.dtype != dtype or tensor.rank != shape.len or + !std.mem.eql(usize, tensor.shape[0..tensor.rank], shape)) + { + std.log.err("unsupported backbone tensor {s}: {s} {any}", .{ name, @tagName(tensor.dtype), tensor.shape[0..tensor.rank] }); + return error.UnsupportedTensor; + } + const offset = if (binding) tensor.offset else tensor.begin; + if (offset % dtype.size() != 0) return error.InvalidTensorView; + if (binding) { + var count: usize = 1; + for (shape) |dim| count = try std.math.mul(usize, count, dim); + if (tensor.bytes != try std.math.mul(usize, count, dtype.size())) return error.UnsupportedTensor; + if (tensor.offset % dtype.size() != 0 or tensor.offset > tensor.buffer.length() or + tensor.bytes > tensor.buffer.length() - tensor.offset) return error.InvalidTensorView; + destination.?.* = tensor; + } +} + +fn has(source: anytype, name: []const u8) bool { + return source.has(name); +} + +const Scale = struct { name: []const u8, format: exl3.ScaleFormat }; +fn scale(comptime binding: bool, source: anytype, prefix: []const u8, size: usize, comptime fp: []const u8, comptime signs: []const u8, storage: []u8) !Scale { + var alternative: [160]u8 = undefined; + const f = try named(storage, "{s}.{s}", .{ prefix, fp }); + const s = try named(&alternative, "{s}.{s}", .{ prefix, signs }); + const have_fp = has(source, f); + if (have_fp == has(source, s)) return error.MissingOrAmbiguousExl3Scale; + const name = if (have_fp) f else try named(storage, "{s}.{s}", .{ prefix, signs }); + const t = try source.get(name); + if (t.rank != 1 or t.shape[0] != (if (have_fp) size else size / 16) or + t.dtype != (if (have_fp) st.DType.f16 else st.DType.i16)) return error.UnsupportedExl3Scale; + const offset = if (binding) t.offset else t.begin; + if (offset % 2 != 0) return error.InvalidTensorView; + if (binding) { + const bytes = (if (have_fp) size else size / 16) * 2; + if (t.bytes != bytes or t.offset % 2 != 0 or t.offset > t.buffer.length() or + bytes > t.buffer.length() - t.offset) return error.InvalidTensorView; + } + return .{ .name = name, .format = if (have_fp) .f16 else .packed_signs }; +} + +/// Compressor/indexer matrices ride either storage: packs quantize; rank-split repacks stay dense with rotation scales. +fn matrixAny(comptime binding: bool, source: anytype, prefix: []const u8, k: usize, n: usize, destination: ?*AnyMatrix) !void { + var name_buf: [160]u8 = undefined; + if (has(source, try named(&name_buf, "{s}.weight", .{prefix}))) { + const weight_name = try named(&name_buf, "{s}.weight", .{prefix}); + const w = try source.get(weight_name); + if (w.dtype != .bf16 or w.rank != 2 or w.shape[0] != n or w.shape[1] != k) { + std.log.err("unsupported dense tensor {s}: {s} {any}", .{ weight_name, @tagName(w.dtype), w.shape[0..w.rank] }); + return error.UnsupportedTensor; + } + const su = try source.get(try named(&name_buf, "{s}.suh", .{prefix})); + const sv = try source.get(try named(&name_buf, "{s}.svh", .{prefix})); + if (su.dtype != .f16 or su.rank != 1 or su.shape[0] != k or + sv.dtype != .f16 or sv.rank != 1 or sv.shape[0] != n) return error.UnsupportedExl3Scale; + if (binding) { + if (w.offset % 2 != 0) return error.InvalidTensorView; + destination.?.* = .{ .dense = .{ .weight = w, .suh = su, .svh = sv, .k = k, .n = n } }; + } + return; + } + var m: Matrix = undefined; + try matrix(binding, source, prefix, k, n, if (binding) &m else null); + if (binding) destination.?.* = .{ .trellis = m }; +} + +fn matrix(comptime binding: bool, source: anytype, prefix: []const u8, k: usize, n: usize, destination: ?*Matrix) !void { + var name: [160]u8 = undefined; + const trellis = try source.get(try named(&name, "{s}.trellis", .{prefix})); + if (trellis.dtype != .i16 or trellis.rank != 3 or trellis.shape[0] != k / 16 or + trellis.shape[1] != n / 16 or trellis.shape[2] % 8 != 0) return error.UnsupportedExl3Shape; + if ((if (binding) trellis.offset else trellis.begin) % 2 != 0) return error.InvalidTensorView; + const halves = trellis.shape[2] / 8; + if (halves < 2 or halves > 16 or (halves % 2 != 0 and halves > 7)) return error.UnsupportedExl3Width; + // Pinned checkpoint uses mul1 throughout; missing markers cannot silently select the legacy codebook. + const marker = try source.get(try named(&name, "{s}.mul1", .{prefix})); + if (marker.dtype != .i32 or marker.rank != 0) return error.UnsupportedExl3Marker; + if ((if (binding) marker.offset else marker.begin) % 4 != 0) return error.InvalidTensorView; + if (has(source, try named(&name, "{s}.mcg", .{prefix})) or + has(source, try named(&name, "{s}.3inst", .{prefix})) or + has(source, try named(&name, "{s}.bias", .{prefix}))) + { + std.log.err("unexpected exl3 parts beside {s}", .{prefix}); + return error.UnsupportedExl3Parts; + } + var su_name: [160]u8 = undefined; + var sv_name: [160]u8 = undefined; + const su = try scale(binding, source, prefix, k, "suh", "su", &su_name); + const sv = try scale(binding, source, prefix, n, "svh", "sv", &sv_name); + if (binding) { + const expected = k / 16 * (n / 16) * halves * 16; + if (trellis.bytes != expected or trellis.offset % 2 != 0 or + trellis.offset > trellis.buffer.length() or expected > trellis.buffer.length() - trellis.offset) + return error.InvalidTensorView; + if (marker.bytes != 4 or marker.offset % 4 != 0 or marker.offset > marker.buffer.length() or + 4 > marker.buffer.length() - marker.offset) return error.InvalidTensorView; + const marker_bytes = marker.buffer.contents()[marker.offset..][0..4]; + if (std.mem.readInt(u32, marker_bytes, .little) != exl3.format.Mul1Mul) return error.UnsupportedExl3Marker; + const suh = try source.get(su.name); + const svh = try source.get(sv.name); + destination.?.* = .{ + .trellis = .{ .buffer = trellis.buffer, .offset = trellis.offset }, + .suh = .{ .buffer = suh.buffer, .offset = suh.offset }, + .svh = .{ .buffer = svh.buffer, .offset = svh.offset }, + .k = k, + .n = n, + .width = .{ .halves = @intCast(halves) }, + .codebook = .mul1, + .suh_format = su.format, + .svh_format = sv.format, + }; + try destination.?.validate(); + } +} + +fn populate(comptime binding: bool, source: anytype, model: ?*Model) !void { + try plain(binding, source, "embed.weight", .bf16, &.{ 129280, 5120 }, if (binding) &model.?.embed else null); + try plain(binding, source, "norm.weight", .bf16, &.{5120}, if (binding) &model.?.norm else null); + try matrix(binding, source, "head", 5120, 129280, if (binding) &model.?.head else null); + var name: [160]u8 = undefined; + var prefix: [128]u8 = undefined; + for (0..40) |layer| { + const out: ?*Layer = if (binding) &model.?.layers[layer] else null; + try plain(binding, source, try named(&name, "layers.{d}.attn_norm.weight", .{layer}), .bf16, &.{5120}, if (binding) &out.?.attn_norm else null); + try plain(binding, source, try named(&name, "layers.{d}.ffn_norm.weight", .{layer}), .bf16, &.{5120}, if (binding) &out.?.ffn_norm else null); + inline for (.{ "attn", "ffn" }) |kind| { + const hc: ?*Hc = if (binding) (if (std.mem.eql(u8, kind, "attn")) &out.?.hc_attn else &out.?.hc_ffn) else null; + try plain(binding, source, try named(&name, "layers.{d}.hc_{s}_fn", .{ layer, kind }), .f32, &.{ 24, 20480 }, if (binding) &hc.?.projection else null); + try plain(binding, source, try named(&name, "layers.{d}.hc_{s}_scale", .{ layer, kind }), .f32, &.{3}, if (binding) &hc.?.scale else null); + try plain(binding, source, try named(&name, "layers.{d}.hc_{s}_base", .{ layer, kind }), .f32, &.{24}, if (binding) &hc.?.base else null); + } + const attn: ?*Attention = if (binding) &out.?.attn else null; + inline for (.{ "wq_a", "wq_b", "wkv", "wo_b" }, .{ 5120, 1280, 5120, 8192 }, .{ 1280, 32768, 512, 5120 }) |part, k, n| { + try matrix(binding, source, try named(&prefix, "layers.{d}.attn.{s}", .{ layer, part }), k, n, if (binding) &@field(attn.?, part) else null); + } + for (0..8) |slice| try matrix(binding, source, try named(&prefix, "layers.{d}.attn.wo_a.slice.{d}", .{ layer, slice }), 4096, 1024, if (binding) &attn.?.wo_a[slice] else null); + try plain(binding, source, try named(&name, "layers.{d}.attn.attn_sink", .{layer}), .f32, &.{64}, if (binding) &attn.?.sink else null); + try plain(binding, source, try named(&name, "layers.{d}.attn.q_norm.weight", .{layer}), .bf16, &.{1280}, if (binding) &attn.?.q_norm else null); + try plain(binding, source, try named(&name, "layers.{d}.attn.kv_norm.weight", .{layer}), .bf16, &.{512}, if (binding) &attn.?.kv_norm else null); + const owner = layer == 2 or layer == 8 or layer == 14 or layer == 20; + if (binding) { + attn.?.compressor = if (owner) @as(Compressor, undefined) else null; + attn.?.indexer = if (owner or layer == 24 or layer == 28 or layer == 32 or layer == 36) + @as(Indexer, undefined) + else + null; + } + if (owner) { + const compressor: ?*Compressor = if (binding) &attn.?.compressor.? else null; + try matrixAny(binding, source, try named(&prefix, "layers.{d}.attn.compressor.wkv", .{layer}), 5120, 512, if (binding) &compressor.?.wkv else null); + try plain(binding, source, try named(&name, "layers.{d}.attn.compressor.norm.weight", .{layer}), .bf16, &.{512}, if (binding) &compressor.?.norm else null); + if (binding) compressor.?.wgate = if (layer < 20) @as(AnyMatrix, undefined) else null; + if (layer < 20) try matrixAny(binding, source, try named(&prefix, "layers.{d}.attn.compressor.wgate", .{layer}), 5120, 512, if (binding) &compressor.?.wgate.? else null); + } + if (owner or layer == 24 or layer == 28 or layer == 32 or layer == 36) { + const indexer: ?*Indexer = if (binding) &attn.?.indexer.? else null; + try matrixAny(binding, source, try named(&prefix, "layers.{d}.attn.indexer.wq_b", .{layer}), 1280, 4096, if (binding) &indexer.?.wq_b else null); + try plainEither(binding, source, try named(&name, "layers.{d}.attn.indexer.weights_proj.weight", .{layer}), &.{ 32, 5120 }, if (binding) &indexer.?.weights_proj else null); + if (binding) { + indexer.?.wk = if (owner) @as(AnyMatrix, undefined) else null; + indexer.?.k_norm = if (owner) @as(Tensor, undefined) else null; + } + if (owner) { + try matrixAny(binding, source, try named(&prefix, "layers.{d}.attn.indexer.wk", .{layer}), 512, 128, if (binding) &indexer.?.wk.? else null); + try plain(binding, source, try named(&name, "layers.{d}.attn.indexer.k_norm.weight", .{layer}), .bf16, &.{128}, if (binding) &indexer.?.k_norm.? else null); + } + } + const moe: ?*Moe = if (binding) &out.?.moe else null; + // The released checkpoints store the router gate f16; rank-split repacks store it bf16. + if (binding) { + const tensor = try source.get(try named(&name, "layers.{d}.ffn.gate.weight", .{layer})); + if ((tensor.dtype != .f16 and tensor.dtype != .bf16) or tensor.rank != 2 or + tensor.shape[0] != 384 or tensor.shape[1] != 5120) + { + std.log.err("unsupported router gate {s}: {s} {any}", .{ name, @tagName(tensor.dtype), tensor.shape[0..tensor.rank] }); + return error.UnsupportedTensor; + } + moe.?.gate = tensor; + } else { + const tensor = try source.get(try named(&name, "layers.{d}.ffn.gate.weight", .{layer})); + if ((tensor.dtype != .f16 and tensor.dtype != .bf16) or tensor.rank != 2 or + tensor.shape[0] != 384 or tensor.shape[1] != 5120) return error.UnsupportedTensor; + } + // The bias rides the same three dtypes the route kernel decodes (f32, f16, bf16). + if (binding) { + const tensor = try source.get(try named(&name, "layers.{d}.ffn.gate.bias", .{layer})); + if ((tensor.dtype != .f16 and tensor.dtype != .bf16 and tensor.dtype != .f32) or tensor.rank != 1 or tensor.shape[0] != 384) + return error.UnsupportedTensor; + moe.?.bias = tensor; + } else { + const tensor = try source.get(try named(&name, "layers.{d}.ffn.gate.bias", .{layer})); + if ((tensor.dtype != .f16 and tensor.dtype != .bf16 and tensor.dtype != .f32) or tensor.rank != 1 or tensor.shape[0] != 384) + return error.UnsupportedTensor; + } + inline for (.{ "w1", "w3", "w2" }, 0..) |part, projection| { + const k: usize = if (projection == 2) 2304 else 5120; + const n: usize = if (projection == 2) 5120 else 2304; + try matrix(binding, source, try named(&prefix, "layers.{d}.ffn.shared_experts.{s}", .{ layer, part }), k, n, if (binding) &moe.?.shared[projection] else null); + if (!binding) for (0..384) |expert| try matrix(false, source, try named(&prefix, "layers.{d}.ffn.experts.{d}.{s}", .{ layer, expert, part }), k, n, null); + } + const engram = layer == 1 or layer == 14; + if (binding) out.?.engram = if (engram) @as(Engram, undefined) else null; + if (engram) { + const e: ?*Engram = if (binding) &out.?.engram.? else null; + const rows: usize = if (layer == 1) 384006168 else 384016682; + try plain(binding, source, try named(&name, "layers.{d}.engram.embed.weight", .{layer}), .f8_e4m3, &.{ rows, 256 }, if (binding) &e.?.table else null); + try plain(binding, source, try named(&name, "layers.{d}.engram.embed.scale", .{layer}), .f8_e8m0, &.{ rows, 8 }, if (binding) &e.?.scales else null); + try plainEither(binding, source, try named(&name, "layers.{d}.engram.q_weight", .{layer}), &.{ 4, 5120 }, if (binding) &e.?.q_weight else null); + try plainEither(binding, source, try named(&name, "layers.{d}.engram.k_weight", .{layer}), &.{ 4, 5120 }, if (binding) &e.?.k_weight else null); + try matrix(binding, source, try named(&prefix, "layers.{d}.engram.wkv", .{layer}), 6144, 25600, if (binding) &e.?.wkv else null); + } + } +} diff --git a/zig/src/families/deepseek_v41/workspace.zig b/zig/src/families/deepseek_v41/workspace.zig new file mode 100644 index 0000000000..a144bff477 --- /dev/null +++ b/zig/src/families/deepseek_v41/workspace.zig @@ -0,0 +1,107 @@ +//! Fixed chunk-sized activations: no layer history or context-by-query score matrices. +const std = @import("std"); +const mtl = @import("metal"); +const exl3 = @import("core").exl3_projection; +const View = @import("view.zig").View; +const hash = @import("engram_hash.zig"); +const groups = @import("core").expert_groups; + +pub const Field = enum { + ids, + token_map, + history, + stream, + collapsed, + normalized, + branch, + previous_pre, + attention_pre, + post, + comb, + mixes, + inv, + engram_ids_0, + engram_ids_1, + engram_gather, + engram_kv, + router_logits, + route_ids, + route_weights, + gate, + up, + activation, + routed, + shared_gate, + shared_up, + shared_activation, + shared_down, + logits, + projection, + grouping, +}; +const count = std.enums.values(Field).len; + +pub const Plan = struct { + rows: u32, + offsets: [count]usize, + bytes: usize, + + pub fn init(rows: u32) !Plan { + if (rows == 0 or rows > 4096) return error.InvalidPrefillChunk; + const t: usize = rows; + const grouped = if (rows > 1) try groups.Plan.init(t * 6, 384) else null; + const projection_bytes = @max(try exl3.scratchBytes(5120, 129280, t), if (grouped) |g| @max(try exl3.scratchBytes(5120, 2304, g.padded), try exl3.scratchBytes(2304, 5120, g.padded)) else try exl3.scratchBytes(2304, 5120, t * 6)); + const sizes = [_]usize{ + t * 4, hash.raw_vocab * 4, 16, t * 20480 * 2, t * 5120 * 2, t * 5120 * 2, t * 5120 * 2, + t * 4 * 4, t * 4 * 4, t * 4 * 4, t * 16 * 4, t * 24 * 4, t * 4, t * 24 * 4, + t * 24 * 4, t * 6144 * 2, t * 25600 * 2, t * 384 * 4, t * 6 * 4, t * 6 * 4, t * 6 * 2304 * 2, + t * 6 * 2304 * 2, t * 6 * 2304 * 2, t * 6 * 5120 * 2, t * 2304 * 2, t * 2304 * 2, t * 2304 * 2, t * 5120 * 2, + t * 129280 * 4, projection_bytes, if (grouped) |g| g.bytes else 0, + }; + var out = Plan{ .rows = rows, .offsets = undefined, .bytes = 0 }; + for (sizes, 0..) |size, i| { + out.offsets[i] = std.mem.alignForward(usize, out.bytes, 16); + out.bytes = try std.math.add(usize, out.offsets[i], size); + } + out.bytes = std.mem.alignForward(usize, out.bytes, 16384); + return out; + } +}; + +pub const Workspace = struct { + buffer: mtl.Buffer, + plan: Plan, + + pub fn init(device: mtl.Device, plan: Plan, token_map: []const u8) !Workspace { + if (token_map.len != hash.raw_vocab * 4) return error.InvalidTokenMap; + const buffer = try device.buffer(plan.bytes, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + const w = Workspace{ .buffer = buffer, .plan = plan }; + @memcpy(buffer.contents()[plan.offsets[@backingInt(Field.token_map)]..][0..token_map.len], token_map); + @memset(buffer.contents()[plan.offsets[@backingInt(Field.history)]..][0..16], 0); + return w; + } + + pub fn deinit(w: *Workspace) void { + w.buffer.deinit(); + } + pub fn view(w: *const Workspace, field: Field) View { + return .{ .buffer = w.buffer, .offset = w.plan.offsets[@backingInt(field)] }; + } + pub fn projection(w: *const Workspace) exl3.View { + const v = w.view(.projection); + return .{ .buffer = v.buffer, .offset = v.offset }; + } + pub fn grouping(w: *const Workspace, rows: u32) !groups.Scratch { + if (rows < 2 or rows > w.plan.rows) return error.InvalidPrefillChunk; + const v = w.view(.grouping); + return .{ .buffer = v.buffer, .offset = v.offset, .plan = try groups.Plan.init(@as(usize, rows) * 6, 384) }; + } + pub fn tokenIds(w: *const Workspace, rows: usize) []u32 { + const v = w.view(.ids); + return @as([*]u32, @ptrCast(@alignCast(v.buffer.contents() + v.offset)))[0..rows]; + } + pub fn logits(w: *const Workspace, rows: usize) []const f32 { + const v = w.view(.logits); + return @as([*]const f32, @ptrCast(@alignCast(v.buffer.contents() + v.offset)))[0 .. rows * 129280]; + } +}; diff --git a/zig/src/metal/counters.zig b/zig/src/metal/counters.zig index 610e22086b..0bbd65f186 100644 --- a/zig/src/metal/counters.zig +++ b/zig/src/metal/counters.zig @@ -1,4 +1,4 @@ -//! Metal counter sampling; owned resources must outlive every command buffer using them. +//! Metal counter sampling; owned descriptors, samples and fences must outlive their command buffers. const std = @import("std"); const objc = @import("objc.zig"); const types = @import("types.zig"); @@ -107,7 +107,7 @@ pub const SampleBuffer = struct { objc.msg(void, encoder.id, "sampleCountersInBuffer:atSampleIndex:withBarrier:", .{ self.id, index, true }); } - /// Caller must ensure GPU completion; borrows a bounded NSData result without a GPU blit or second copy. + /// Borrow bounded NSData bytes after completion; same-command blit resolution is unsupported, no result copy. pub fn resolveCompleted(self: SampleBuffer, count: usize) !Resolved { if (count == 0 or count > self.count) return error.InvalidSampleCount; const data = objc.msg(?objc.Id, self.id, "resolveCounterRange:", .{ diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index 26b68d7925..e41bd7d02e 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -24,6 +24,15 @@ pub const frags = @import("core/frags.zig"); pub const affine_mm = @import("core/affine_mm.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/config.zig"); +pub const deepseek_v41_load = @import("families/deepseek_v41/load.zig"); +pub const deepseek_v41_weights = @import("families/deepseek_v41/weights.zig"); +pub const deepseek_v41_view = @import("families/deepseek_v41/view.zig"); +pub const deepseek_v41 = @import("families/deepseek_v41/metal_model.zig"); +pub const deepseek_v41_profiler = @import("families/deepseek_v41/profiler.zig"); +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 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/attention_math.zig b/zig/tests/deepseek_v41_reference/attention_math.zig new file mode 100644 index 0000000000..1adec0172d --- /dev/null +++ b/zig/tests/deepseek_v41_reference/attention_math.zig @@ -0,0 +1,298 @@ +//! CPU CSA2 from DeepSeek AI inference/{model,kernel}.py, rev 2cba9e42aa026125f3ed06c6d98c1db82f7ca027. +const std = @import("std"); +const tf = @import("tensorfold"); +const m = @import("math.zig"); +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; + } + 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]; + } + 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); + } + } + 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; + } + 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; + }; + 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); + } + } + 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)); + } + 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); + } + 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); + } + 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]); + } + 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(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]); + }; + acc[d] = @floatCast(accumulated); + } + high = next_high; + } + denominator += @exp(sink - high); + for (out, acc) |*v, a| v.* = m.b(a / denominator); + } +}; + +test "online 64-column sparse attention rounds probabilities before value GEMM" { + var state = try State.init(std.testing.allocator, 2); + defer state.deinit(); + var q: [512]u16 = @splat(0); + q[0] = m.b(8); + state.window[1] = m.b(1); + state.window[512] = m.b(1); + var indices: [65]i32 = @splat(-1); + indices[0] = 0; + indices[1] = 1; + indices[64] = 0; + var out: [512]u16 = undefined; + state.sparseHead(&q, state.window[0..1024], &indices, 0, &out); + const probability = @exp(-@as(f32, 8) * (1 / @sqrt(@as(f32, 512)))); + var denominator = @as(f32, @floatCast(@as(f64, probability) + 1)); + denominator += probability; + denominator += probability; + const expected = m.b((2 * m.f(m.b(probability))) / denominator); + try std.testing.expectEqual(expected, out[1]); +} + +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)); +} diff --git a/zig/tests/deepseek_v41_reference/engram_math.zig b/zig/tests/deepseek_v41_reference/engram_math.zig new file mode 100644 index 0000000000..88772898de --- /dev/null +++ b/zig/tests/deepseek_v41_reference/engram_math.zig @@ -0,0 +1,106 @@ +//! CPU Engram from DeepSeek AI inference/engram.py and model.py, rev 2cba9e42aa026125f3ed06c6d98c1db82f7ca027. +const std = @import("std"); +const tf = @import("tensorfold"); +const m = @import("math.zig"); +const projection = @import("projection.zig"); +const Engram = tf.deepseek_v41_weights.Engram; +const Config = tf.deepseek_v41_config.Text; +const multipliers = [2][4]u64{ + .{ 76632096046245, 4839876093313, 35959672319349, 73987337458391 }, + .{ 67716810739261, 51510806800915, 30921347202721, 82619226485591 }, +}; +const primes = [2][24]u32{ + .{ 16000057, 16000079, 16000081, 16000097, 16000121, 16000129, 16000133, 16000183, 16000189, 16000207, 16000211, 16000253, 16000277, 16000289, 16000307, 16000321, 16000339, 16000381, 16000393, 16000399, 16000403, 16000409, 16000447, 16000463 }, + .{ 16000477, 16000487, 16000499, 16000507, 16000511, 16000573, 16000609, 16000627, 16000667, 16000669, 16000693, 16000697, 16000711, 16000729, 16000759, 16000769, 16000781, 16000799, 16000813, 16000819, 16000841, 16000877, 16000879, 16000889 }, +}; +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) + 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 }; + } + 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]; + } + } + 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; + } + 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])); + } + } + } +}; diff --git a/zig/tests/deepseek_v41_reference/math.zig b/zig/tests/deepseek_v41_reference/math.zig new file mode 100644 index 0000000000..a6dccb2b93 --- /dev/null +++ b/zig/tests/deepseek_v41_reference/math.zig @@ -0,0 +1,289 @@ +//! Independent CPU equations: DeepSeek AI inference/{model,kernel}.py, rev 2cba9e42aa026125f3ed06c6d98c1db82f7ca027. +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; + } + } + return std.math.copysign(levels[best], x); +} +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); +} +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 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); + } +} +/// 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; +} +/// 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 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 sigmoid(x: f32) f32 { + return 1 / (1 + @exp(-x)); +} +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 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 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 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 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); + } +} + +test "BF16 probability ties and FP4 ties retain source boundaries" { + try std.testing.expectEqual(@as(u16, 0x3f80), b(@bitCast(@as(u32, 0x3f808000)))); + try std.testing.expectEqual(@as(f32, 2), roundE2(2.5)); + try std.testing.expectEqual(@as(f32, 4), roundE2(5)); + try std.testing.expectEqual(@as(f32, 448), roundE4(500)); +} + +test "delayed hyperconnection residual uses comb columns, not rows" { + var h = [_]u16{ b(1), b(2), b(3), b(4) }; + const mix = Mix{ .pre = .{ 1, 0, 0, 0 }, .post = .{ 0, 0, 0, 0 }, .comb = .{ 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 1, 0, 0, 0 } }; + expand(&h, &.{b(0)}, mix); + try std.testing.expectEqualSlices(u16, &.{ b(4), b(1), b(2), b(3) }, &h); +} + +test "hyperconnection FP32 accumulation preserves collapse and branch cancellation" { + var h = [_]u16{ b(1), b(1), b(1), b(1) }; + var collapsed: [1]u16 = undefined; + collapse(&h, .{ 1, 0x1p-24, 0x1p-24, -1 }, &collapsed); + try std.testing.expectEqual(@as(u16, 0), collapsed[0]); + const mix = Mix{ .pre = .{ 1, 0, 0, 0 }, .post = .{ -1, 0, 0, 0 }, .comb = .{ 1, 0, 0, 0, 0x1p-24, 0, 0, 0, 0x1p-24, 0, 0, 0, 0, 0, 0, 0 } }; + expand(&h, &.{b(1)}, mix); + try std.testing.expectEqual(@as(u16, 0), h[0]); +} + +test "hyperconnection rounds elementwise products before cancelling reductions" { + var h = [_]u16{ b(3), b(3), b(0), b(0) }; + const coefficient: f32 = 1 + 0x1p-23; + var collapsed: [1]u16 = undefined; + collapse(&h, .{ -1, coefficient, 0, 0 }, &collapsed); + try std.testing.expectEqual(b(0x1p-21), collapsed[0]); + const mix = Mix{ .pre = .{ 1, 0, 0, 0 }, .post = .{ 0, 0, 0, 0 }, .comb = .{ -1, 0, 0, 0, coefficient, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 } }; + expand(&h, &.{b(0)}, mix); + try std.testing.expectEqual(b(0x1p-21), h[0]); + h = .{ b(3), b(0), b(0), b(0) }; + const branch_mix = Mix{ .pre = .{ 1, 0, 0, 0 }, .post = .{ coefficient, 0, 0, 0 }, .comb = .{ -1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 } }; + expand(&h, &.{b(3)}, branch_mix); + try std.testing.expectEqual(b(0x1p-21), h[0]); +} + +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(); + for (parsed.value.object.get("vectors").?.array.items) |vector| { + const theta: f32 = @floatFromInt(vector.object.get("theta").?.integer); + const original: u32 = @intCast(vector.object.get("original").?.integer); + for (vector.object.get("bits").?.array.items, 0..) |expected, index| { + const actual = ropeFrequency(index, theta, original, 16, 32, 1); + const bits: u32 = @intCast(expected.integer); + const source: f32 = @bitCast(bits); + try std.testing.expectApproxEqAbs(source, actual, source * std.math.floatEps(f32)); + } + } +} + +test "ideal FP32 inverse square root avoids an intermediate square-root rounding" { + // Actual CUDA agrees on this HC RMS argument; CPU/MPS's two-round result is one ULP higher. + try std.testing.expectEqual(@as(f32, 0.3651483356952667), idealRsqrt(7.500000953674316)); +} + +test "CUDA mean projection preserves the BF16 RMS affine tie" { + @setFloatMode(.strict); + // Constructed exact-sum row isolates canonical CUDA RMSNorm mean projection through the BF16 affine tie. + const gamma: f32 = @bitCast(@as(u32, 0x3c65de29)); + const inverse = idealRsqrt(meanFromWide(9, 5120) + @as(f32, 1e-20)); + try std.testing.expectEqual(b(1), b((@as(f32, 3) * inverse) * gamma)); +} diff --git a/zig/tests/deepseek_v41_reference/operator_attention.zig b/zig/tests/deepseek_v41_reference/operator_attention.zig new file mode 100644 index 0000000000..73e94641c3 --- /dev/null +++ b/zig/tests/deepseek_v41_reference/operator_attention.zig @@ -0,0 +1,155 @@ +//! Single-position retained-input attention isolation. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const projection = @import("projection.zig"); +const math = @import("math.zig"); +const attention = tf.deepseek_v41.attention; +const Model = tf.deepseek_v41.Model; +const diagnostic = @import("operator_diagnostic.zig"); +const diagnosticRow = diagnostic.diagnosticRow; +const attentionMetric = diagnostic.attentionMetric; +const attention_math = @import("attention_math.zig"); +const AttentionStage = diagnostic.AttentionStage; +const AttentionCapture = diagnostic.AttentionCapture; +const attention_elements = diagnostic.attention_elements; +const attentionProjection = diagnostic.attentionProjection; +const hyperconnections = @import("operator_hc.zig").hyperconnections; + +/// Local isolation receipt, not a replacement for the full eight-token oracle or its 1% gate. +pub fn runAttention(gpa: std.mem.Allocator, io: std.Io, model: *Model, snapshot: []const u8, dir: []const u8, layer: usize) !void { + const blob = try std.Io.Dir.cwd().readFileAlloc(io, snapshot, gpa, .limited(8 << 20)); + defer gpa.free(blob); + const cfg = model.cfg.text_config; + if (layer >= model.weights.layers.len) return error.InvalidAttentionLayer; + const snapshot_offset = layer * 2 * 61488; + if (blob.len < snapshot_offset + 40960 + 10240) return error.InvalidAttentionSnapshot; + try model.reset(); + const input = model.workspace.view(.normalized); + @memcpy(input.buffer.contents()[input.offset..][0..10240], blob[snapshot_offset + 40960 ..][0..10240]); + const input_row = @as([*]const u16, @ptrCast(@alignCast(input.buffer.contents() + input.offset)))[0..5120]; + const native = try model.device.buffer(attention_elements * 2, mtl.ResourceOptions.shared); + defer native.deinit(); + const reference_rows = try gpa.alloc(u16, attention_elements); + defer gpa.free(reference_rows); + var capture = AttentionCapture{ .model = model, .native = native, .reference = reference_rows }; + model.residency.add(native); + model.residency.commit(); + model.residency.requestResidency(); + defer { + model.residency.remove(native); + model.residency.commit(); + } + var scratch = try projection.Scratch.init(gpa, 8192, 32768); + defer scratch.deinit(); + var state = try attention_math.State.init(gpa, 1); + defer state.deinit(); + var output: [5120]u16 = undefined; + // Replay prerequisite normalized layers; complete each GPU command before reusing the CPU input buffer. + for (0..layer) |prior| { + const prior_offset = prior * 2 * 61488 + 40960; + @memcpy(input.buffer.contents()[input.offset..][0..10240], blob[prior_offset..][0..10240]); + const warm_cb = model.queue.commandBufferUnretained(); + const warm_enc = warm_cb.compute(.serial); + var warm_encoding = true; + errdefer if (warm_encoding) warm_enc.end(); + for (model.checkpoint.shards.items) |shard| warm_enc.useResource(shard.buffer, mtl.ResourceUsage.read); + try attention.encode(&model.pipelines.attention, warm_enc, &model.weights.layers[prior].attn, &model.cache, &model.attention_workspace, input, model.workspace.view(.branch), &model.pipelines.projection, model.workspace.projection(), @intCast(prior), 1, cfg); + warm_enc.end(); + warm_encoding = false; + warm_cb.commit(); + warm_cb.wait(); + if (warm_cb.failure() != null) return error.OperatorGpuFailure; + try state.apply(model.weights.layers[prior].attn, prior, input_row, 0, cfg, &output, &scratch); + } + @memcpy(input.buffer.contents()[input.offset..][0..10240], blob[snapshot_offset + 40960 ..][0..10240]); + const weights = &model.weights.layers[layer].attn; + const cb = model.queue.commandBufferUnretained(); + const enc = cb.compute(.serial); + var encoding = true; + errdefer if (encoding) enc.end(); + for (model.checkpoint.shards.items) |shard| enc.useResource(shard.buffer, mtl.ResourceUsage.read); + try attention.encodeDiagnostic(&model.pipelines.attention, enc, weights, &model.cache, &model.attention_workspace, input, model.workspace.view(.branch), &model.pipelines.projection, model.workspace.projection(), @intCast(layer), 1, cfg, true, .{ .ctx = &capture, .capture = AttentionCapture.gpu }); + enc.end(); + encoding = false; + cb.commit(); + cb.wait(); + if (cb.failure() != null) return error.OperatorGpuFailure; + const sparse_snapshot = try gpa.alloc(u16, 32768); + defer gpa.free(sparse_snapshot); + state.sparse_snapshot = sparse_snapshot; + state.observer = .{ .ctx = &capture, .capture = AttentionCapture.cpu }; + try state.apply(weights.*, layer, input_row, 0, cfg, &output, &scratch); + state.observer = null; + state.sparse_snapshot = null; + + var report: std.ArrayList(u8) = .empty; + defer report.deinit(gpa); + var count: usize = 0; + var header: [512]u8 = undefined; + try report.appendSlice(gpa, try std.fmt.bufPrint(&header, "{{\"mode\":\"attention-sameinput\",\"token\":0,\"layer\":{d},\"input_offset_bytes\":{d},\"input_length_bytes\":10240,\"identical_normalized_input\":true,\"retained_normalized_prefix_layers\":{d},\"compression_ratio\":{d},\"reference_dot\":\"independent_EXL3_FP64\",\"qualification_gate_applied\":false,\"cases\":[\n", .{ layer, snapshot_offset + 40960, layer, cfg.compress_ratios[layer] })); + inline for (std.enums.values(AttentionStage)) |stage| + try attentionMetric(u16, gpa, &report, &count, @tagName(stage), "same_normalized_input_chain", capture.row(stage), capture.cpuRow(stage)); + + const expected = try gpa.alloc(f32, 32768); + defer gpa.free(expected); + try attentionProjection(gpa, model, weights.wq_a, input, "wq_a", &scratch, expected, &report, &count); + try attentionProjection(gpa, model, weights.wq_b, capture.view(.qr_norm), "wq_b", &scratch, expected, &report, &count); + try attentionProjection(gpa, model, weights.wkv, input, "wkv", &scratch, expected, &report, &count); + for (weights.wo_a, 0..) |matrix, group| { + var name: [32]u8 = undefined; + try attentionProjection(gpa, model, matrix, capture.view(.o).shifted(group * 4096 * 2), try std.fmt.bufPrint(&name, "wo_a_group{d}", .{group}), &scratch, expected, &report, &count); + } + try attentionProjection(gpa, model, weights.wo_b, capture.view(.grouped), "wo_b", &scratch, expected, &report, &count); + + var same_input: [32768]u16 = undefined; + math.rms(capture.row(.qr_projection), weights.q_norm, cfg.rms_norm_eps, same_input[0..1280]); + try attentionMetric(u16, gpa, &report, &count, "qr_norm", "same_native_input_reference_rms", capture.row(.qr_norm), same_input[0..1280]); + math.rms(capture.row(.kv_projection), weights.kv_norm, cfg.rms_norm_eps, same_input[0..512]); + try attentionMetric(u16, gpa, &report, &count, "kv_norm", "same_native_input_reference_rms", capture.row(.kv_norm), same_input[0..512]); + @memcpy(same_input[0..512], capture.row(.kv_rope)); + math.quantize(same_input[0..512], .window); + try attentionMetric(u16, gpa, &report, &count, "kv_quantized", "same_native_input_reference_quantize", capture.row(.kv_quantized), same_input[0..512]); + if (cfg.compress_ratios[layer] == 1) { + const source = try model.cache.source(@intCast(layer), 1); + const compressed = model.workspace.view(.collapsed); + const unpack_cb = model.queue.commandBufferUnretained(); + const unpack_enc = unpack_cb.compute(.serial); + var unpack_encoding = true; + errdefer if (unpack_encoding) unpack_enc.end(); + try model.pipelines.attention.cache.encodeUnpack(unpack_enc, source.main, compressed, .{ .rows = 1, .width = 512, .stride = 512, .packed_stride = 288, .block = 16, .kind = .main_fp4 }); + unpack_enc.end(); + unpack_encoding = false; + unpack_cb.commit(); + unpack_cb.wait(); + if (unpack_cb.failure() != null) return error.OperatorGpuFailure; + const actual_compressed = diagnosticRow(u16, compressed, 512); + const reference_compressed = state.sources[state.active_owner * state.capacity * 640 ..][0..512]; + try attentionMetric(u16, gpa, &report, &count, "compressed_kv", "same_retained_normalized_prefix", actual_compressed, reference_compressed); + var sinks: [64]f32 = undefined; + for (&sinks, 0..) |*v, head| v.* = math.value(weights.sink, head); + inline for (.{ "attention-compressed-native.bf16", "attention-compressed-reference.bf16", "attention-sinks.f32" }, .{ std.mem.sliceAsBytes(actual_compressed), std.mem.sliceAsBytes(reference_compressed), std.mem.sliceAsBytes(&sinks) }) |name, data| { + const path = try std.fs.path.join(gpa, &.{ dir, name }); + defer gpa.free(path); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = path, .data = data }); + } + @memcpy(reference_compressed, actual_compressed); + const indices = [_]i32{ 0, 128 }; + for (0..64) |head| + state.sparseHead(capture.row(.q)[head * 512 ..][0..512], capture.row(.kv_quantized), &indices, sinks[head], same_input[head * 512 ..][0..512]); + try attentionMetric(u16, gpa, &report, &count, "sparse_o", "same_native_q_window_and_compressed_kv_reference_sparse", capture.row(.sparse_o), &same_input); + } + if (cfg.compress_ratios[layer] == 0) { + const indices = [_]i32{0}; + for (0..64) |head| + state.sparseHead(capture.row(.q)[head * 512 ..][0..512], capture.row(.kv_quantized), &indices, math.value(weights.sink, head), same_input[head * 512 ..][0..512]); + try attentionMetric(u16, gpa, &report, &count, "sparse_o", "same_native_q_and_quantized_kv_reference_sparse", capture.row(.sparse_o), &same_input); + } + try hyperconnections(gpa, model, blob[snapshot_offset..], capture.view(.wo_b), &report, &count, layer); + try report.appendSlice(gpa, "\n]}\n"); + inline for (.{ "attention-sameinput.json", "attention-native.bf16", "attention-reference.bf16", "attention-input.bf16" }, .{ report.items, std.mem.sliceAsBytes(native.slice(u16, attention_elements)), std.mem.sliceAsBytes(reference_rows), std.mem.sliceAsBytes(input_row) }) |name, data| { + const path = try std.fs.path.join(gpa, &.{ dir, name }); + defer gpa.free(path); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = path, .data = data }); + } +} diff --git a/zig/tests/deepseek_v41_reference/operator_diagnostic.zig b/zig/tests/deepseek_v41_reference/operator_diagnostic.zig new file mode 100644 index 0000000000..66dc039bcf --- /dev/null +++ b/zig/tests/deepseek_v41_reference/operator_diagnostic.zig @@ -0,0 +1,124 @@ +//! Shared retained-input diagnostic capture and metrics. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const projection = @import("projection.zig"); +const math = @import("math.zig"); +const attention = tf.deepseek_v41.attention; +const Model = tf.deepseek_v41.Model; +const ex = tf.exl3_projection; + +pub const AttentionStage = attention.DiagnosticStage; +pub const attention_widths = [_]usize{ 1280, 1280, 32768, 32768, 512, 512, 512, 512, 32768, 32768, 8192, 5120 }; +pub const attention_elements = blk: { + var total: usize = 0; + for (attention_widths) |width| total += width; + break :blk total; +}; +pub fn attentionOffset(stage: AttentionStage) usize { + var offset: usize = 0; + for (attention_widths[0..@backingInt(stage)]) |width| offset += width; + return offset; +} +pub const AttentionCapture = struct { + model: *Model, + native: mtl.Buffer, + reference: []u16, + + pub fn view(self: *const AttentionCapture, stage: AttentionStage) attention.View { + return .{ .buffer = self.native, .offset = attentionOffset(stage) * 2 }; + } + pub fn row(self: *const AttentionCapture, stage: AttentionStage) []u16 { + return self.native.slice(u16, attention_elements)[attentionOffset(stage)..][0..attention_widths[@backingInt(stage)]]; + } + pub fn cpuRow(self: *const AttentionCapture, stage: AttentionStage) []u16 { + return self.reference[attentionOffset(stage)..][0..attention_widths[@backingInt(stage)]]; + } + pub fn gpu(ctx: *anyopaque, enc: mtl.ComputeEncoder, stage: AttentionStage, source: attention.View) !void { + const self: *AttentionCapture = @ptrCast(@alignCast(ctx)); + const cache = &self.model.pipelines.attention.cache; + if (stage == .kv_quantized) { + try cache.encodeUnpack(enc, source, self.view(stage), .{ .rows = 1, .width = 512, .stride = 512, .packed_stride = 528, .block = 32, .kind = .window_fp8 }); + } else { + cache.encodeCopy(enc, source, self.view(stage), @intCast(attention_widths[@backingInt(stage)] * 2)); + } + } + pub fn cpu(ctx: *anyopaque, stage: AttentionStage, values: []const u16) void { + const self: *AttentionCapture = @ptrCast(@alignCast(ctx)); + @memcpy(self.cpuRow(stage), values); + } +}; + +pub fn attentionMetric(comptime T: type, gpa: std.mem.Allocator, report: *std.ArrayList(u8), count: *usize, name: []const u8, comparison: []const u8, actual: []const T, expected: []const T) !void { + if (actual.len == 0 or actual.len != expected.len) return error.InvalidAttentionMetricShape; + var squared: f64 = 0; + var magnitude: f64 = 0; + var max_abs: f64 = 0; + var worst: usize = 0; + var unequal: usize = 0; + var first: ?usize = null; + for (actual, expected, 0..) |a, r, i| { + const x = if (T == u16) math.f(a) else a; + const y = if (T == u16) math.f(r) else r; + if (!std.math.isFinite(x) or !std.math.isFinite(y)) return error.NonfiniteAttentionDiagnostic; + const delta = @as(f64, x) - y; + squared += delta * delta; + magnitude += @as(f64, y) * y; + if (@abs(delta) > max_abs) { + max_abs = @abs(delta); + worst = i; + } + const ab = if (T == u16) a else math.b(a); + const rb = if (T == u16) r else math.b(r); + if (ab != rb) { + unequal += 1; + if (first == null) first = i; + } + } + const line = try std.json.Stringify.valueAlloc(gpa, .{ + .name = name, + .comparison = comparison, + .elements = actual.len, + .output_dtype = if (T == u16) "bf16" else "f32", + .normalized_rms_error = if (magnitude == 0) @sqrt(squared) else @sqrt(squared / magnitude), + .max_absolute_error = max_abs, + .worst_index = worst, + .bf16_unequal_count = unequal, + .bf16_first_unequal_index = first, + .worst_native = if (T == u16) math.f(actual[worst]) else actual[worst], + .worst_reference = if (T == u16) math.f(expected[worst]) else expected[worst], + }, .{}); + defer gpa.free(line); + if (count.* != 0) try report.appendSlice(gpa, ",\n"); + try report.appendSlice(gpa, line); + count.* += 1; + std.debug.print("{s}\n", .{line}); +} + +pub fn attentionProjection(gpa: std.mem.Allocator, model: *Model, matrix: ex.Matrix, input: attention.View, name: []const u8, scratch: *projection.Scratch, expected: []f32, report: *std.ArrayList(u8), count: *usize) !void { + const input_row = @as([*]const u16, @ptrCast(@alignCast(input.buffer.contents() + input.offset)))[0..matrix.k]; + try projection.applyF32(.{ .trellis = matrix }, input_row, expected[0..matrix.n], scratch); + const output = model.workspace.view(.logits); + const cb = model.queue.commandBufferUnretained(); + const enc = cb.compute(.serial); + var encoding = true; + errdefer if (encoding) enc.end(); + ex.useMatrix(enc, matrix); + try ex.encode(&model.pipelines.projection, enc, matrix, .{ .buffer = input.buffer, .offset = input.offset }, .{ .buffer = output.buffer, .offset = output.offset }, model.workspace.projection(), .{ .count = 1, .x_stride = matrix.k, .y_stride = matrix.n, .output = .f32 }); + enc.end(); + encoding = false; + cb.commit(); + cb.wait(); + if (cb.failure() != null) return error.OperatorGpuFailure; + const actual = @as([*]const f32, @ptrCast(@alignCast(output.buffer.contents() + output.offset)))[0..matrix.n]; + try attentionMetric(f32, gpa, report, count, name, "same_native_input_fp64_dot", actual, expected[0..matrix.n]); +} + +pub fn diagnosticRow(comptime T: type, view: attention.View, len: usize) []T { + return @as([*]T, @ptrCast(@alignCast(view.buffer.contents() + view.offset)))[0..len]; +} +pub fn attentionWrite(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, name: []const u8, data: []const u8) !void { + const path = try std.fs.path.join(gpa, &.{ dir, name }); + defer gpa.free(path); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = path, .data = data }); +} diff --git a/zig/tests/deepseek_v41_reference/operator_hc.zig b/zig/tests/deepseek_v41_reference/operator_hc.zig new file mode 100644 index 0000000000..4aec23362a --- /dev/null +++ b/zig/tests/deepseek_v41_reference/operator_hc.zig @@ -0,0 +1,98 @@ +//! Retained-input hyperconnection isolation, separate from qualification. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const math = @import("math.zig"); +const attention = tf.deepseek_v41.attention; +const Model = tf.deepseek_v41.Model; +const diagnostic = @import("operator_diagnostic.zig"); +const diagnosticRow = diagnostic.diagnosticRow; +const attentionMetric = diagnostic.attentionMetric; + +pub fn hyperconnections(gpa: std.mem.Allocator, model: *Model, stream_bytes: []const u8, branch: attention.View, report: *std.ArrayList(u8), count: *usize, layer: usize) !void { + const ws = &model.workspace; + const hc = model.weights.layers[layer].hc_ffn; + const cfg = model.cfg.text_config; + const stream = ws.view(.stream); + @memcpy(stream.buffer.contents()[stream.offset..][0..40960], stream_bytes[0..40960]); + var input: [20480]u16 = undefined; + @memcpy(&input, diagnosticRow(u16, stream, input.len)); + var projected: [24]f32 = undefined; + math.dense(hc.projection, &input, &projected); + const expected_mix = math.mixes(&input, hc, cfg); + var squares: f64 = 0; + for (input) |v| squares += @as(f64, math.f(v)) * math.f(v); + const mean = math.meanFromWide(squares, input.len); + const expected_inv = [_]f32{math.idealRsqrt(mean + cfg.rms_norm_eps)}; + const cb = model.queue.commandBufferUnretained(); + const enc = cb.compute(.serial); + for (model.checkpoint.shards.items) |shard| enc.useResource(shard.buffer, mtl.ResourceUsage.read); + tf.deepseek_v41.mhc.encodeRms(enc, &model.pipelines.mhc, stream, ws.view(.inv), .{ .rows = 1, .width = 5120, .stream_stride = 20480, .eps = cfg.rms_norm_eps }); + tf.deepseek_v41.ops.encodeDense(enc, &model.pipelines.ops, stream, .{ .buffer = hc.projection.buffer, .offset = hc.projection.offset }, ws.view(.mixes), .{ .rows = 1, .k = 20480, .n = 24, .x_stride = 20480, .w_stride = 20480, .y_stride = 24, .w_dtype = .f32, .y_dtype = .f32 }); + enc.barrier(); + tf.deepseek_v41.mhc.encodeSplit(enc, &model.pipelines.mhc, ws.view(.mixes), ws.view(.inv), .{ .buffer = hc.scale.buffer, .offset = hc.scale.offset }, .{ .buffer = hc.base.buffer, .offset = hc.base.offset }, ws.view(.attention_pre), ws.view(.post), ws.view(.comb), .{ .rows = 1, .eps = cfg.hc_eps }); + enc.barrier(); + tf.deepseek_v41.mhc.encodeCollapse(enc, &model.pipelines.mhc, stream, ws.view(.attention_pre), ws.view(.collapsed), .{ .rows = 1, .width = 5120, .stream_stride = 20480, .y_stride = 5120 }); + enc.barrier(); + tf.deepseek_v41.mhc.encodePost(enc, &model.pipelines.mhc, branch, stream, ws.view(.post), ws.view(.comb), stream, .{ .rows = 1, .width = 5120, .branch_stride = 5120, .residual_stride = 20480, .y_stride = 20480 }); + enc.end(); + cb.commit(); + cb.wait(); + if (cb.failure() != null) return error.OperatorGpuFailure; + try attentionMetric(f32, gpa, report, count, "hc_projection", "same_retained_stream_fp64_dot", diagnosticRow(f32, ws.view(.mixes), 24), &projected); + try attentionMetric(f32, gpa, report, count, "hc_inverse", "same_retained_stream_reference_statistic", diagnosticRow(f32, ws.view(.inv), 1), &expected_inv); + try attentionMetric(f32, gpa, report, count, "hc_pre", "same_retained_stream_chain", diagnosticRow(f32, ws.view(.attention_pre), 4), &expected_mix.pre); + try attentionMetric(f32, gpa, report, count, "hc_post", "same_retained_stream_chain", diagnosticRow(f32, ws.view(.post), 4), &expected_mix.post); + try attentionMetric(f32, gpa, report, count, "hc_comb", "same_retained_stream_chain", diagnosticRow(f32, ws.view(.comb), 16), &expected_mix.comb); + var native_mix: math.Mix = undefined; + @memcpy(&native_mix.pre, diagnosticRow(f32, ws.view(.attention_pre), 4)); + @memcpy(&native_mix.post, diagnosticRow(f32, ws.view(.post), 4)); + @memcpy(&native_mix.comb, diagnosticRow(f32, ws.view(.comb), 16)); + var collapsed: [5120]u16 = undefined; + math.collapse(&input, native_mix.pre, &collapsed); + try attentionMetric(u16, gpa, report, count, "hc_collapse", "same_retained_stream_and_native_coefficients", diagnosticRow(u16, ws.view(.collapsed), 5120), &collapsed); + math.expand(&input, diagnosticRow(u16, branch, 5120), native_mix); + try attentionMetric(u16, gpa, report, count, "hc_expand", "same_retained_stream_branch_and_native_coefficients", diagnosticRow(u16, stream, 20480), &input); + if (layer + 1 < model.weights.layers.len and stream_bytes.len >= 61488 + 40960) { + @memcpy(stream.buffer.contents()[stream.offset..][0..40960], stream_bytes[61488..][0..40960]); + @memcpy(&input, diagnosticRow(u16, stream, input.len)); + const norm = model.weights.layers[layer + 1].attn_norm; + const next_cb = model.queue.commandBufferUnretained(); + const next_enc = next_cb.compute(.serial); + next_enc.useResource(norm.buffer, mtl.ResourceUsage.read); + tf.deepseek_v41.mhc.encodeCollapse(next_enc, &model.pipelines.mhc, stream, ws.view(.attention_pre), ws.view(.collapsed), .{ .rows = 1, .width = 5120, .stream_stride = 20480, .y_stride = 5120 }); + next_enc.barrier(); + tf.deepseek_v41.ops.encodeRms(next_enc, &model.pipelines.ops, ws.view(.collapsed), .{ .buffer = norm.buffer, .offset = norm.offset }, ws.view(.branch), .{ .rows = 1, .width = 5120, .x_stride = 5120, .y_stride = 5120, .eps = cfg.rms_norm_eps }); + next_enc.end(); + next_cb.commit(); + next_cb.wait(); + if (next_cb.failure() != null) return error.OperatorGpuFailure; + math.collapse(&input, native_mix.pre, &collapsed); + try attentionMetric(u16, gpa, report, count, "next_collapse", "same_retained_ffn_stream_and_native_pre", diagnosticRow(u16, ws.view(.collapsed), 5120), &collapsed); + math.collapse(&input, expected_mix.pre, &collapsed); + try attentionMetric(u16, gpa, report, count, "next_collapse", "same_retained_ffn_stream_reference_pre", diagnosticRow(u16, ws.view(.collapsed), 5120), &collapsed); + math.rms(&collapsed, norm, cfg.rms_norm_eps, &collapsed); + try attentionMetric(u16, gpa, report, count, "next_normalized", "same_retained_ffn_stream_reference_pre_and_norm", diagnosticRow(u16, ws.view(.branch), 5120), &collapsed); + } +} + +pub fn runHyperconnections(gpa: std.mem.Allocator, io: std.Io, model: *Model, snapshot: []const u8, dir: []const u8, layer: usize) !void { + if (layer >= model.weights.layers.len) return error.InvalidHyperconnectionLayer; + const blob = try std.Io.Dir.cwd().readFileAlloc(io, snapshot, gpa, .limited(8 << 20)); + defer gpa.free(blob); + const offset = layer * 2 * 61488; + if (blob.len < offset + 2 * 61488) return error.InvalidHyperconnectionSnapshot; + try model.reset(); + const branch = model.workspace.view(.branch); + @memcpy(branch.buffer.contents()[branch.offset..][0..10240], blob[offset + 61488 + 51200 ..][0..10240]); + var report: std.ArrayList(u8) = .empty; + defer report.deinit(gpa); + var count: usize = 0; + var header: [256]u8 = undefined; + try report.appendSlice(gpa, try std.fmt.bufPrint(&header, "{{\"mode\":\"hc-sameinput\",\"token\":null,\"token_identity_source\":\"unknown_headerless_snapshot\",\"layer\":{d},\"qualification_gate_applied\":false,\"cases\":[\n", .{layer})); + try hyperconnections(gpa, model, blob[offset..], branch, &report, &count, layer); + try report.appendSlice(gpa, "\n]}\n"); + const path = try std.fs.path.join(gpa, &.{ dir, "hc-sameinput.json" }); + defer gpa.free(path); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = path, .data = report.items }); +} diff --git a/zig/tests/deepseek_v41_reference/operator_history.zig b/zig/tests/deepseek_v41_reference/operator_history.zig new file mode 100644 index 0000000000..499a1cd463 --- /dev/null +++ b/zig/tests/deepseek_v41_reference/operator_history.zig @@ -0,0 +1,298 @@ +//! Retained-input attention history and source-layout receipts. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const projection = @import("projection.zig"); +const math = @import("math.zig"); +const attention = tf.deepseek_v41.attention; +const Model = tf.deepseek_v41.Model; +const diagnostic = @import("operator_diagnostic.zig"); +const diagnosticRow = diagnostic.diagnosticRow; +const attentionMetric = diagnostic.attentionMetric; +const attention_math = @import("attention_math.zig"); +const AttentionStage = diagnostic.AttentionStage; +const AttentionCapture = diagnostic.AttentionCapture; +const attention_widths = diagnostic.attention_widths; +const attention_elements = diagnostic.attention_elements; +const attentionOffset = diagnostic.attentionOffset; +const attentionProjection = diagnostic.attentionProjection; +const attentionWrite = diagnostic.attentionWrite; + +// Source indices address physical ring slots then compressed rows; reference reserves 128 window rows even at zero. +fn attentionHistoryIndices(position: usize, selected: []const i32, compressed_offset: usize, out: *[640]i32) []const i32 { + const window_columns: usize = if (position == 0) 1 else 128; + if (position == 0) { + out[0] = 0; + } else { + for (0..128) |slot| { + const at = (position % 128 + 1 + slot) % 128; + out[slot] = if (at > position) -1 else @intCast(at); + } + } + for (selected, 0..) |row, slot| + out[window_columns + slot] = if (row < 0) -1 else @intCast(compressed_offset + @as(usize, @intCast(row))); + return out[0 .. window_columns + selected.len]; +} + +/// Replay prerequisite layers at each position; native/reference cursors advance without borrowed intermediates. +pub fn runAttentionHistory(gpa: std.mem.Allocator, io: std.Io, model: *Model, trace_dir: []const u8, dir: []const u8, token: usize, layer: usize) !void { + if (layer >= model.weights.layers.len) return error.InvalidAttentionLayer; + if (token >= model.cache.plan.capacity) return error.InvalidAttentionHistory; + if (token >= attention_math.max_capacity) return error.DiagnosticHistoryTooLarge; + const prefix_elements = try std.math.mul(usize, token + 1, (layer + 1) * 5120); + if (prefix_elements > (256 << 20) / @sizeOf(u16)) return error.DiagnosticHistoryTooLarge; + const cfg = model.cfg.text_config; + const history_count = token + 1; + const ratio = cfg.compress_ratios[layer]; + const compressed_rows = if (ratio == 0) 0 else history_count / ratio; + const topk = @min(compressed_rows, cfg.index_topk); + const window_rows: usize = if (token == 0) 1 else 128; + const snapshot_bytes = 61488; + const snapshot_offset = layer * 2 * snapshot_bytes; + const weights = &model.weights.layers[layer].attn; + const input = model.workspace.view(.normalized); + const native = try model.device.buffer(attention_elements * 2, mtl.ResourceOptions.shared); + defer native.deinit(); + const decoded = try model.device.buffer((128 * 512 + compressed_rows * 640) * 2, mtl.ResourceOptions.shared); + defer decoded.deinit(); + const reference_rows = try gpa.alloc(u16, attention_elements); + defer gpa.free(reference_rows); + const sparse_snapshot = try gpa.alloc(u16, 32768); + defer gpa.free(sparse_snapshot); + const normalized_prefix = try gpa.alloc(u16, prefix_elements); + defer gpa.free(normalized_prefix); + var retained_branch: [5120]u16 = undefined; + var capture = AttentionCapture{ .model = model, .native = native, .reference = reference_rows }; + for ([_]mtl.Buffer{ native, decoded }) |buffer| model.residency.add(buffer); + model.residency.commit(); + model.residency.requestResidency(); + defer { + for ([_]mtl.Buffer{ native, decoded }) |buffer| model.residency.remove(buffer); + model.residency.commit(); + } + var scratch = try projection.Scratch.init(gpa, 8192, 32768); + defer scratch.deinit(); + var state = try attention_math.State.init(gpa, history_count); + defer state.deinit(); + var output: [5120]u16 = undefined; + try model.reset(); + for (0..history_count) |position| { + var name: [64]u8 = undefined; + const trace_name = try std.fmt.bufPrint(&name, "trace-{d}-gpu.bin", .{position}); + const path = try std.fs.path.join(gpa, &.{ trace_dir, trace_name }); + defer gpa.free(path); + const blob = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(8 << 20)); + defer gpa.free(blob); + if (blob.len != model.weights.layers.len * 2 * snapshot_bytes) return error.InvalidAttentionSnapshot; + if (model.cache.position != position) return error.InvalidAttentionHistoryCursor; + for (0..layer + 1) |prior| { + const offset = prior * 2 * snapshot_bytes + 40960; + const row = normalized_prefix[(position * (layer + 1) + prior) * 5120 ..][0..5120]; + @memcpy(std.mem.sliceAsBytes(row), blob[offset..][0..10240]); + @memcpy(input.buffer.contents()[input.offset..][0..10240], std.mem.sliceAsBytes(row)); + const target = position == token and prior == layer; + const cb = model.queue.commandBufferUnretained(); + const enc = cb.compute(.serial); + var encoding = true; + errdefer if (encoding) enc.end(); + for (model.checkpoint.shards.items) |shard| enc.useResource(shard.buffer, mtl.ResourceUsage.read); + if (target) { + try attention.encodeDiagnostic(&model.pipelines.attention, enc, &model.weights.layers[prior].attn, &model.cache, &model.attention_workspace, input, model.workspace.view(.branch), &model.pipelines.projection, model.workspace.projection(), @intCast(prior), 1, cfg, true, .{ .ctx = &capture, .capture = AttentionCapture.gpu }); + @memcpy(std.mem.sliceAsBytes(&retained_branch), blob[snapshot_offset + 51200 ..][0..10240]); + } else { + try attention.encode(&model.pipelines.attention, enc, &model.weights.layers[prior].attn, &model.cache, &model.attention_workspace, input, model.workspace.view(.branch), &model.pipelines.projection, model.workspace.projection(), @intCast(prior), 1, cfg); + } + enc.end(); + encoding = false; + cb.commit(); + cb.wait(); + if (cb.failure() != null) return error.OperatorGpuFailure; + // Wait before the next retained row replaces the native workspace input. + if (target) { + state.observer = .{ .ctx = &capture, .capture = AttentionCapture.cpu }; + state.sparse_snapshot = sparse_snapshot; + } + try state.apply(model.weights.layers[prior].attn, prior, row, position, cfg, &output, &scratch); + state.observer = null; + state.sparse_snapshot = null; + } + try model.cache.append(1); + } + + // Export post-insertion physical ring slots, masking never-written slots in source indices. + const decoded_view = attention.View{ .buffer = decoded }; + const native_window = decoded.slice(u16, 128 * 512 + compressed_rows * 640)[0 .. 128 * 512]; + const native_compressed = decoded.slice(u16, 128 * 512 + compressed_rows * 640)[128 * 512 ..][0 .. compressed_rows * 512]; + const native_index_keys = decoded.slice(u16, 128 * 512 + compressed_rows * 640)[128 * 512 + compressed_rows * 512 ..][0 .. compressed_rows * 128]; + const unpack_cb = model.queue.commandBufferUnretained(); + const unpack_enc = unpack_cb.compute(.serial); + var unpack_encoding = true; + errdefer if (unpack_encoding) unpack_enc.end(); + const ring = try model.cache.ring(@intCast(layer)); + try model.pipelines.attention.cache.encodeUnpack(unpack_enc, ring, decoded_view, .{ .rows = 128, .width = 512, .stride = 512, .packed_stride = 528, .block = 32, .kind = .window_fp8 }); + if (compressed_rows != 0) { + const source = try model.cache.source(@intCast(layer), ratio); + try model.pipelines.attention.cache.encodeUnpack(unpack_enc, source.main, decoded_view.shifted(128 * 512 * 2), .{ .rows = @intCast(compressed_rows), .width = 512, .stride = 512, .packed_stride = 288, .block = 16, .kind = .main_fp4 }); + try model.pipelines.attention.cache.encodeUnpack(unpack_enc, source.index, decoded_view.shifted((128 * 512 + compressed_rows * 512) * 2), .{ .rows = @intCast(compressed_rows), .width = 128, .stride = 128, .packed_stride = 68, .block = 32, .kind = .index_fp4 }); + } + unpack_enc.end(); + unpack_encoding = false; + unpack_cb.commit(); + unpack_cb.wait(); + if (unpack_cb.failure() != null) return error.OperatorGpuFailure; + const reference_window = state.window[layer * 128 * 512 ..][0 .. 128 * 512]; + const reference_compressed = try gpa.alloc(u16, compressed_rows * 640); + defer gpa.free(reference_compressed); + const reference_main = reference_compressed[0 .. compressed_rows * 512]; + const reference_index_keys = reference_compressed[compressed_rows * 512 ..]; + for (0..compressed_rows) |row| { + const source_row = state.sources[state.active_owner * state.capacity * 640 + row * 640 ..][0..640]; + @memcpy(reference_main[row * 512 ..][0..512], source_row[0..512]); + @memcpy(reference_index_keys[row * 128 ..][0..128], source_row[512..640]); + } + const native_selected = diagnosticRow(i32, model.attention_workspace.selection.positions(), 512)[0..topk]; + var reference_selected: [512]i32 = undefined; + for (state.selected[0..topk], 0..) |row, i| reference_selected[i] = @intCast(row); + var native_indices_storage: [640]i32 = undefined; + var reference_indices_storage: [640]i32 = undefined; + const native_indices = attentionHistoryIndices(token, native_selected, window_rows, &native_indices_storage); + const reference_indices = attentionHistoryIndices(token, reference_selected[0..topk], window_rows, &reference_indices_storage); + var sinks: [64]f32 = undefined; + for (&sinks, 0..) |*value, head| value.* = math.value(weights.sink, head); + + var first_stage: ?[]const u8 = null; + inline for (std.enums.values(AttentionStage)) |stage| + if (first_stage == null and !std.mem.eql(u16, capture.row(stage), capture.cpuRow(stage))) { + first_stage = @tagName(stage); + }; + const header = try std.json.Stringify.valueAlloc(gpa, .{ + .mode = "attention-history-sameinput", + .token = token, + .layer = layer, + .history_count = history_count, + .token_identity_source = "caller_supplied_TOKEN_and_trace_POSITION_gpu_filenames_raw_snapshots_have_no_token_header", + .normalized_prefix_replay = .{ .file_pattern = "trace-{position}-gpu.bin", .origin = "retained_native_GPU_normalized_attention_inputs", .first_position = 0, .last_position = token, .first_layer = 0, .last_layer = layer, .order = "position_major_then_layer_ascending_inclusive", .snapshot_count_per_file = model.weights.layers.len * 2, .snapshot_length_bytes = snapshot_bytes, .attention_snapshot_index = "2*layer", .normalized_input_offset_within_snapshot_bytes = 40960, .normalized_input_length_bytes = 10240, .native_cache_reset = true, .native_cache_cursor_matches_each_position = true, .native_cache_append_rows_per_position = 1, .reference_position_passed_explicitly = true, .reference_capacity = state.capacity }, + .identical_retained_normalized_inputs = true, + .independent_reference_history = true, + .native_history_replay_branch_matches_retained_trace = std.mem.eql(u16, capture.row(.wo_b), &retained_branch), + .first_differing_same_history_input_stage = first_stage, + .window_history_equal = std.mem.eql(u16, native_window, reference_window), + .compressed_main_equal = std.mem.eql(u16, native_compressed, reference_main), + .compressed_index_keys_equal = std.mem.eql(u16, native_index_keys, reference_index_keys), + .selection_equal = std.mem.eql(i32, native_selected, reference_selected[0..topk]), + .compression_ratio = ratio, + .compressed_rows = compressed_rows, + .selection_count = topk, + .reference_dot = "independent_EXL3_FP64", + .reference_sparse_dot = "independent_FP64", + .qualification_gate_applied = false, + .qualification_oracle_modified_to_reuse_native_values = false, + .same_native_query_cache_diagnostic = "separate_reference_State_with_captured_native_window_main_selection_and_q_not_used_in_independent_chain", + .official_revision = "2cba9e42aa026125f3ed06c6d98c1db82f7ca027", + .receipt_dtype_endianness = "little_endian", + .receipt_strides = "contiguous_element_strides", + .files = .{ + .stage_chain = .{ .native = "attention-native.bf16", .reference = "attention-reference.bf16", .dtype = "bf16", .shape = .{attention_elements}, .strides = .{1}, .layout = "stages_table_below" }, + .normalized_prefix = .{ .file = "attention-normalized-prefix.bf16", .dtype = "bf16", .shape = .{ history_count, layer + 1, 5120 }, .strides = .{ (layer + 1) * 5120, 5120, 1 } }, + .target_input = .{ .file = "attention-input.bf16", .dtype = "bf16", .shape = .{5120}, .strides = .{1} }, + .retained_branch = .{ .file = "attention-retained-branch.bf16", .dtype = "bf16", .shape = .{5120}, .strides = .{1} }, + .query = .{ .native = "attention-query-native.bf16", .reference = "attention-query-reference.bf16", .dtype = "bf16", .shape = .{ 1, 1, 64, 512 }, .strides = .{ 32768, 32768, 512, 1 } }, + .window = .{ .native = "attention-window-native.bf16", .reference = "attention-window-reference.bf16", .dtype = "bf16", .shape = .{ 1, 128, 512 }, .strides = .{ 65536, 512, 1 }, .layout = "physical_ring_slot_position_modulo_128_includes_current_quantized_KV" }, + .compressed_main = .{ .native = "attention-compressed-native.bf16", .reference = "attention-compressed-reference.bf16", .dtype = "bf16", .shape = .{ 1, compressed_rows, 512 }, .strides = .{ compressed_rows * 512, 512, 1 } }, + .compressed_index_keys = .{ .native = "attention-index-keys-native.bf16", .reference = "attention-index-keys-reference.bf16", .dtype = "bf16", .shape = .{ 1, compressed_rows, 128 }, .strides = .{ compressed_rows * 128, 128, 1 } }, + .selection = .{ .native = "attention-selection-native.i32", .reference = "attention-selection-reference.i32", .dtype = "i32", .shape = .{ 1, topk }, .strides = .{ topk, 1 }, .layout = "actual_compressed_row_ids_compact_export_native_workspace_row_stride_is_512" }, + .source_kv = .{ .native = "attention-source-kv-native.bf16", .reference = "attention-source-kv-reference.bf16", .dtype = "bf16", .shape = .{ 1, window_rows + compressed_rows, 512 }, .strides = .{ (window_rows + compressed_rows) * 512, 512, 1 }, .layout = "source_sparse_attn_KV_window_then_compressed_rows_no_requantization" }, + .source_indices = .{ .native = "attention-source-indices-native.i32", .reference = "attention-source-indices-reference.i32", .dtype = "i32", .shape = .{ 1, 1, native_indices.len }, .strides = .{ native_indices.len, native_indices.len, 1 }, .layout = "source_get_window_topk_idxs_then_actual_selection_plus_window_rows_minus_one_remains_masked" }, + .sinks = .{ .file = "attention-sinks.f32", .dtype = "f32", .shape = .{64}, .strides = .{1} }, + .sparse_output = .{ .native = "attention-sparse-native.bf16", .reference = "attention-sparse-reference.bf16", .same_native_input_reference = "attention-sparse-same-native-input-reference.bf16", .dtype = "bf16", .shape = .{ 1, 1, 64, 512 }, .strides = .{ 32768, 32768, 512, 1 } }, + }, + .source_sparse_call = "sparse_attn(query,source_kv,sinks,source_indices,1/sqrt(512)); output_before_inverse_RoPE", + .source_sparse_replay_is_operator_evidence_only = true, + }, .{}); + defer gpa.free(header); + var report: std.ArrayList(u8) = .empty; + defer report.deinit(gpa); + try report.appendSlice(gpa, header[0 .. header.len - 1]); + try report.appendSlice(gpa, ",\"cases\":[\n"); + var count: usize = 0; + inline for (std.enums.values(AttentionStage)) |stage| + try attentionMetric(u16, gpa, &report, &count, @tagName(stage), "same_retained_normalized_history_input_chain", capture.row(stage), capture.cpuRow(stage)); + try attentionMetric(u16, gpa, &report, &count, "retained_branch", "native_history_replay_vs_retained_native_trace", capture.row(.wo_b), &retained_branch); + try attentionMetric(u16, gpa, &report, &count, "window_history", "same_retained_normalized_history_input_chain", native_window, reference_window); + if (compressed_rows != 0) { + try attentionMetric(u16, gpa, &report, &count, "compressed_kv", "same_retained_normalized_history_input_chain", native_compressed, reference_main); + try attentionMetric(u16, gpa, &report, &count, "compressed_index_keys", "same_retained_normalized_history_input_chain", native_index_keys, reference_index_keys); + } + + // This separate diagnostic State never overwrites independent retained-history reference data. + var native_input_state = try attention_math.State.init(gpa, history_count); + defer native_input_state.deinit(); + for (0..compressed_rows) |row| + @memcpy(native_input_state.sources[row * 640 ..][0..512], native_compressed[row * 512 ..][0..512]); + var helper_indices_storage: [640]i32 = undefined; + const helper_indices = attentionHistoryIndices(token, native_selected, 128, &helper_indices_storage); + var same_input: [32768]u16 = undefined; + for (0..64) |head| + native_input_state.sparseHead(capture.row(.q)[head * 512 ..][0..512], native_window, helper_indices, sinks[head], same_input[head * 512 ..][0..512]); + try attentionMetric(u16, gpa, &report, &count, "sparse_o", "same_native_query_window_compressed_cache_selection_reference_sparse_diagnostic", capture.row(.sparse_o), &same_input); + try attentionWrite(gpa, io, dir, "attention-sparse-same-native-input-reference.bf16", std.mem.sliceAsBytes(&same_input)); + math.rms(capture.row(.qr_projection), weights.q_norm, cfg.rms_norm_eps, same_input[0..1280]); + try attentionMetric(u16, gpa, &report, &count, "qr_norm", "same_native_input_reference_rms_diagnostic", capture.row(.qr_norm), same_input[0..1280]); + math.rms(capture.row(.kv_projection), weights.kv_norm, cfg.rms_norm_eps, same_input[0..512]); + try attentionMetric(u16, gpa, &report, &count, "kv_norm", "same_native_input_reference_rms_diagnostic", capture.row(.kv_norm), same_input[0..512]); + @memcpy(&same_input, capture.row(.q_projection)); + for (0..64) |head| math.rope(same_input[head * 512 ..][0..512], token, ratio != 0, false, cfg); + try attentionMetric(u16, gpa, &report, &count, "q", "same_native_input_reference_rope_diagnostic", capture.row(.q), &same_input); + @memcpy(same_input[0..512], capture.row(.kv_norm)); + math.rope(same_input[0..512], token, ratio != 0, false, cfg); + try attentionMetric(u16, gpa, &report, &count, "kv_rope", "same_native_input_reference_rope_diagnostic", capture.row(.kv_rope), same_input[0..512]); + @memcpy(same_input[0..512], capture.row(.kv_rope)); + math.quantize(same_input[0..512], .window); + try attentionMetric(u16, gpa, &report, &count, "kv_quantized", "same_native_input_reference_quantize_diagnostic", capture.row(.kv_quantized), same_input[0..512]); + const expected = try gpa.alloc(f32, 32768); + defer gpa.free(expected); + try attentionProjection(gpa, model, weights.wq_a, input, "wq_a", &scratch, expected, &report, &count); + try attentionProjection(gpa, model, weights.wq_b, capture.view(.qr_norm), "wq_b", &scratch, expected, &report, &count); + try attentionProjection(gpa, model, weights.wkv, input, "wkv", &scratch, expected, &report, &count); + for (weights.wo_a, 0..) |matrix, group| { + var name: [32]u8 = undefined; + try attentionProjection(gpa, model, matrix, capture.view(.o).shifted(group * 4096 * 2), try std.fmt.bufPrint(&name, "wo_a_group{d}", .{group}), &scratch, expected, &report, &count); + } + try attentionProjection(gpa, model, weights.wo_b, capture.view(.grouped), "wo_b", &scratch, expected, &report, &count); + try report.appendSlice(gpa, "\n],\"stages\":["); + inline for (std.enums.values(AttentionStage), 0..) |stage, i| { + const descriptor = try std.json.Stringify.valueAlloc(gpa, .{ .name = @tagName(stage), .offset_bytes = attentionOffset(stage) * 2, .elements = attention_widths[i] }, .{}); + defer gpa.free(descriptor); + if (i != 0) try report.appendSlice(gpa, ","); + try report.appendSlice(gpa, descriptor); + } + try report.appendSlice(gpa, "]}\n"); + const source_kv = try gpa.alloc(u16, 2 * (window_rows + compressed_rows) * 512); + defer gpa.free(source_kv); + const source_native = source_kv[0 .. (window_rows + compressed_rows) * 512]; + const source_reference = source_kv[(window_rows + compressed_rows) * 512 ..]; + @memcpy(source_native[0 .. window_rows * 512], native_window[0 .. window_rows * 512]); + @memcpy(source_native[window_rows * 512 ..], native_compressed); + @memcpy(source_reference[0 .. window_rows * 512], reference_window[0 .. window_rows * 512]); + @memcpy(source_reference[window_rows * 512 ..], reference_main); + const input_row = normalized_prefix[(history_count * (layer + 1) - 1) * 5120 ..][0..5120]; + inline for (.{ + "attention-history-sameinput.json", "attention-native.bf16", "attention-reference.bf16", + "attention-normalized-prefix.bf16", "attention-input.bf16", "attention-retained-branch.bf16", + "attention-query-native.bf16", "attention-query-reference.bf16", "attention-window-native.bf16", + "attention-window-reference.bf16", "attention-compressed-native.bf16", "attention-compressed-reference.bf16", + "attention-index-keys-native.bf16", "attention-index-keys-reference.bf16", "attention-selection-native.i32", + "attention-selection-reference.i32", "attention-source-kv-native.bf16", "attention-source-kv-reference.bf16", + "attention-source-indices-native.i32", "attention-source-indices-reference.i32", "attention-sinks.f32", + "attention-sparse-native.bf16", "attention-sparse-reference.bf16", + }, .{ + report.items, std.mem.sliceAsBytes(native.slice(u16, attention_elements)), std.mem.sliceAsBytes(reference_rows), + std.mem.sliceAsBytes(normalized_prefix), std.mem.sliceAsBytes(input_row), std.mem.sliceAsBytes(&retained_branch), + std.mem.sliceAsBytes(capture.row(.q)), std.mem.sliceAsBytes(capture.cpuRow(.q)), std.mem.sliceAsBytes(native_window), + std.mem.sliceAsBytes(reference_window), std.mem.sliceAsBytes(native_compressed), std.mem.sliceAsBytes(reference_main), + std.mem.sliceAsBytes(native_index_keys), std.mem.sliceAsBytes(reference_index_keys), std.mem.sliceAsBytes(native_selected), + std.mem.sliceAsBytes(reference_selected[0..topk]), std.mem.sliceAsBytes(source_native), std.mem.sliceAsBytes(source_reference), + std.mem.sliceAsBytes(native_indices), std.mem.sliceAsBytes(reference_indices), std.mem.sliceAsBytes(&sinks), + std.mem.sliceAsBytes(capture.row(.sparse_o)), std.mem.sliceAsBytes(capture.cpuRow(.sparse_o)), + }) |name, data| try attentionWrite(gpa, io, dir, name, data); +} diff --git a/zig/tests/deepseek_v41_reference/operator_moe.zig b/zig/tests/deepseek_v41_reference/operator_moe.zig new file mode 100644 index 0000000000..558234569f --- /dev/null +++ b/zig/tests/deepseek_v41_reference/operator_moe.zig @@ -0,0 +1,409 @@ +//! Retained-input MoE routing and expert isolation. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const projection = @import("projection.zig"); +const math = @import("math.zig"); +const attention = tf.deepseek_v41.attention; +const Model = tf.deepseek_v41.Model; +const ex = tf.exl3_projection; +const diagnosticRow = @import("operator_diagnostic.zig").diagnosticRow; + +const moe_reference = @import("reference.zig"); +const MoeRows = struct { + gate: [7][2304]u16 = undefined, + up: [7][2304]u16 = undefined, + activation: [7][2304]u16 = undefined, + down: [7][5120]u16 = undefined, + branch: [5120]u16 = undefined, +}; +const MoeFirst = struct { + source_chain: ?[]const u8 = null, + source_expert_stage: ?[]const u8 = null, + native_input_diagnostic: ?[]const u8 = null, + native_weights_reference_chain_diagnostic: ?[]const u8 = null, + retained_weights_reference_chain_diagnostic: ?[]const u8 = null, + native_down_reference_combine_diagnostic_differs: bool = false, +}; + +fn moeMetric(comptime T: type, gpa: std.mem.Allocator, report: *std.ArrayList(u8), count: *usize, name: []const u8, comparison: []const u8, diagnostic_reuse: bool, actual: []const T, expected: []const T) !bool { + if (actual.len == 0 or actual.len != expected.len) return error.InvalidMoeMetricShape; + var squared: f64 = 0; + var magnitude: f64 = 0; + var max_abs: f64 = 0; + var worst: usize = 0; + var unequal: usize = 0; + var first: ?usize = null; + for (actual, expected, 0..) |a, r, i| { + const x = if (T == u16) math.f(a) else a; + const y = if (T == u16) math.f(r) else r; + if (!std.math.isFinite(x) or !std.math.isFinite(y)) return error.NonfiniteMoeDiagnostic; + const delta = @as(f64, x) - y; + squared += delta * delta; + magnitude += @as(f64, y) * y; + if (@abs(delta) > max_abs) { + max_abs = @abs(delta); + worst = i; + } + const ab = if (T == u16) a else @as(u32, @bitCast(a)); + const rb = if (T == u16) r else @as(u32, @bitCast(r)); + if (ab != rb) { + unequal += 1; + if (first == null) first = i; + } + } + const line = try std.json.Stringify.valueAlloc(gpa, .{ + .name = name, + .comparison = comparison, + .diagnostic_native_value_reuse = diagnostic_reuse, + .qualification_gate_applied = false, + .elements = actual.len, + .output_dtype = if (T == u16) "bf16" else "f32", + .normalized_rms_error = if (magnitude == 0) @sqrt(squared) else @sqrt(squared / magnitude), + .max_absolute_error = max_abs, + .worst_index = worst, + .bit_unequal_count = unequal, + .first_bit_unequal_index = first, + .worst_native = if (T == u16) math.f(actual[worst]) else actual[worst], + .worst_reference = if (T == u16) math.f(expected[worst]) else expected[worst], + }, .{}); + defer gpa.free(line); + if (count.* != 0) try report.appendSlice(gpa, ",\n"); + try report.appendSlice(gpa, line); + count.* += 1; + std.debug.print("{s}\n", .{line}); + return unequal != 0; +} + +fn moeWrite(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, name: []const u8, bytes: []const u8) !void { + const path = try std.fs.path.join(gpa, &.{ dir, name }); + defer gpa.free(path); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = path, .data = bytes }); +} + +fn moeWriteRows(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, prefix: []const u8, rows: *const MoeRows) !void { + inline for (.{ "gate", "up", "activation", "down", "branch" }) |stage| { + const name = try std.fmt.allocPrint(gpa, "{s}-{s}.bf16", .{ prefix, stage }); + defer gpa.free(name); + try moeWrite(gpa, io, dir, name, std.mem.asBytes(&@field(rows.*, stage))); + } +} + +fn moeMatrices(model: *Model, layer: usize, expert: ?u32) ![3]ex.Matrix { + if (expert == null) return model.weights.layers[layer].moe.shared; + var matrices: [3]ex.Matrix = undefined; + var name: [128]u8 = undefined; + inline for (.{ "w1", "w3", "w2" }, 0..) |part, i| { + const key = try std.fmt.bufPrint(&name, "layers.{d}.ffn.experts.{d}.{s}", .{ layer, expert.?, part }); + matrices[i] = try tf.deepseek_v41_weights.lookup(&model.checkpoint, key, if (i == 2) 2304 else 5120, if (i == 2) 5120 else 2304); + } + return matrices; +} + +fn moeCombineDown(ids: *const [6]u32, down: *const [7][5120]u16, branch: *[5120]u16) void { + @setFloatMode(.strict); + var sum: [5120]f32 = @splat(0); + // Keep slots in top-six order in receipts, but match the official ascending expert traversal. + for (0..384) |expert| { + for (ids, 0..) |id, slot| { + if (id == expert) + for (&sum, down[slot]) |*s, value| { + s.* += math.f(value); + }; + } + } + for (branch, sum, down[6]) |*v, routed, shared| v.* = math.b(routed + math.f(shared)); +} + +fn moeCombine(ids: *const [6]u32, rows: *MoeRows) void { + moeCombineDown(ids, &rows.down, &rows.branch); +} + +fn moeReferenceRows(model: *Model, layer: usize, input: []const u16, ids: *const [6]u32, routing: *const [6]f32, rows: *MoeRows, scratch: *projection.Scratch, previous: ?*const MoeRows, previous_ids: ?*const [6]u32) !void { + for (0..7) |slot| { + const matrices = try moeMatrices(model, layer, if (slot == 6) null else ids[slot]); + var reuse: ?usize = if (slot == 6 and previous != null) 6 else null; + if (slot != 6 and previous_ids != null) + for (previous_ids.?, 0..) |id, prior| { + if (id == ids[slot]) { + reuse = prior; + break; + } + }; + if (reuse) |prior| { + rows.gate[slot] = previous.?.gate[prior]; + rows.up[slot] = previous.?.up[prior]; + } else { + try projection.apply(.{ .trellis = matrices[0] }, input, &rows.gate[slot], scratch); + try projection.apply(.{ .trellis = matrices[1] }, input, &rows.up[slot], scratch); + } + moe_reference.moeActivation(&rows.gate[slot], &rows.up[slot], if (slot == 6) 1 else routing[slot], model.cfg.text_config, &rows.activation[slot]); + // A shared or identical-weight expert result can be reused only when the BF16 input is identical. + if (reuse != null and std.mem.eql(u16, &rows.activation[slot], &previous.?.activation[reuse.?])) { + rows.down[slot] = previous.?.down[reuse.?]; + } else { + try projection.apply(.{ .trellis = matrices[2] }, &rows.activation[slot], &rows.down[slot], scratch); + } + } + moeCombine(ids, rows); +} + +fn moeEncodeNative(model: *Model, layer: usize, select_routes: bool) !void { + const ws = &model.workspace; + const p = &model.pipelines; + const w = &model.weights.layers[layer].moe; + const cb = model.queue.commandBufferUnretained(); + const enc = cb.compute(.serial); + var encoding = true; + errdefer if (encoding) enc.end(); + for (model.checkpoint.shards.items) |shard| enc.useResource(shard.buffer, mtl.ResourceUsage.read); + for (w.descriptors) |buffer| enc.useResource(buffer, mtl.ResourceUsage.read); + if (select_routes) { + tf.deepseek_v41.ops.encodeDense(enc, &p.ops, ws.view(.normalized), .{ .buffer = w.gate.buffer, .offset = w.gate.offset }, ws.view(.router_logits), .{ .rows = 1, .k = 5120, .n = 384, .x_stride = 5120, .w_stride = 5120, .y_stride = 384, .w_dtype = .f16, .y_dtype = .f32 }); + enc.barrier(); + // Bind the existing production router; no new native interface or diagnostic math kernel. + enc.setPipeline(p.route.select); + inline for (.{ ws.view(.router_logits), attention.View{ .buffer = w.bias.buffer, .offset = w.bias.offset }, ws.view(.route_ids), ws.view(.route_weights) }, 0..) |view, i| + enc.setBuffer(view.buffer, view.offset, i); + const Shape = extern struct { rows: u32, logit_stride: u32, output_stride: u32, temperature: f32, bias_dtype: u32 }; + enc.setValue(Shape{ .rows = 1, .logit_stride = 384, .output_stride = 6, .temperature = 1, .bias_dtype = 2 }, 4); + enc.dispatchGroups(mtl.Size.of(1, 1, 1), mtl.Size.of(128, 1, 1)); + enc.barrier(); + } + var routed_rows = ex.Rows{ .count = 6, .x_stride = 5120, .y_stride = 2304, .routes = .{ .ids = .{ .buffer = ws.view(.route_ids).buffer, .offset = ws.view(.route_ids).offset }, .repeat = 6 } }; + for (0..2) |i| { + const target = ws.view(if (i == 0) .gate else .up); + const input = ws.view(.normalized); + try ex.encodeIndexed(&p.projection, enc, .{ .descriptors = .{ .buffer = w.descriptors[i] }, .count = 384, .k = 5120, .n = 2304 }, .{ .buffer = input.buffer, .offset = input.offset }, .{ .buffer = target.buffer, .offset = target.offset }, ws.projection(), routed_rows); + enc.barrier(); + } + tf.deepseek_v41.ops.encodeSwiGLU(enc, &p.ops, ws.view(.gate), ws.view(.up), ws.view(.route_weights), ws.view(.activation), .{ .rows = 6, .width = 2304, .input_stride = 2304, .output_stride = 2304, .weighted = true }); + enc.barrier(); + routed_rows.x_stride = 2304; + routed_rows.y_stride = 5120; + routed_rows.routes.?.repeat = 1; + const activation = ws.view(.activation); + const down = ws.view(.routed); + try ex.encodeIndexed(&p.projection, enc, .{ .descriptors = .{ .buffer = w.descriptors[2] }, .count = 384, .k = 2304, .n = 5120 }, .{ .buffer = activation.buffer, .offset = activation.offset }, .{ .buffer = down.buffer, .offset = down.offset }, ws.projection(), routed_rows); + enc.barrier(); + for (0..2) |i| { + const target = ws.view(if (i == 0) .shared_gate else .shared_up); + const input = ws.view(.normalized); + ex.useMatrix(enc, w.shared[i]); + try ex.encode(&p.projection, enc, w.shared[i], .{ .buffer = input.buffer, .offset = input.offset }, .{ .buffer = target.buffer, .offset = target.offset }, ws.projection(), .{ .count = 1, .x_stride = 5120, .y_stride = 2304 }); + enc.barrier(); + } + tf.deepseek_v41.ops.encodeSwiGLU(enc, &p.ops, ws.view(.shared_gate), ws.view(.shared_up), ws.view(.route_weights), ws.view(.shared_activation), .{ .rows = 1, .width = 2304, .input_stride = 2304, .output_stride = 2304 }); + enc.barrier(); + const shared_activation = ws.view(.shared_activation); + const shared_down = ws.view(.shared_down); + ex.useMatrix(enc, w.shared[2]); + try ex.encode(&p.projection, enc, w.shared[2], .{ .buffer = shared_activation.buffer, .offset = shared_activation.offset }, .{ .buffer = shared_down.buffer, .offset = shared_down.offset }, ws.projection(), .{ .count = 1, .x_stride = 2304, .y_stride = 5120 }); + enc.barrier(); + tf.deepseek_v41.ops.encodeCombine(enc, &p.ops, ws.view(.routed), ws.view(.shared_down), ws.view(.route_ids), ws.view(.branch), .{ .rows = 1, .width = 5120, .routed_row_stride = 6 * 5120, .slot_stride = 5120, .shared_stride = 5120, .y_stride = 5120 }); + enc.end(); + encoding = false; + cb.commit(); + cb.wait(); + if (cb.failure() != null) return error.OperatorGpuFailure; +} + +fn moeCaptureNative(model: *Model, rows: *MoeRows) void { + const ws = &model.workspace; + inline for (.{ "gate", "up", "activation", "down" }, .{ .gate, .up, .activation, .routed }, .{ .shared_gate, .shared_up, .shared_activation, .shared_down }) |stage, routed, shared| { + const width = if (comptime std.mem.eql(u8, stage, "down")) 5120 else 2304; + const source = diagnosticRow(u16, ws.view(routed), 6 * width); + for (0..6) |slot| @memcpy(&@field(rows.*, stage)[slot], source[slot * width ..][0..width]); + @memcpy(&@field(rows.*, stage)[6], diagnosticRow(u16, ws.view(shared), width)); + } + @memcpy(&rows.branch, diagnosticRow(u16, ws.view(.branch), 5120)); +} + +fn moeCompareRows(gpa: std.mem.Allocator, report: *std.ArrayList(u8), count: *usize, comparison: []const u8, reuse: bool, actual: *const MoeRows, actual_ids: *const [6]u32, expected: *const MoeRows, expected_ids: *const [6]u32, first: *?[]const u8) !void { + // Compare like expert IDs rather than assuming independent top-six orders are identical. + inline for (.{ "gate", "up", "activation", "down" }) |stage| { + var differs = false; + for (0..7) |slot| { + var reference_slot: ?usize = if (slot == 6) 6 else null; + if (slot != 6) + for (expected_ids, 0..) |id, i| { + if (id == actual_ids[slot]) { + reference_slot = i; + break; + } + }; + if (reference_slot == null) continue; + var name: [64]u8 = undefined; + const label = if (slot == 6) try std.fmt.bufPrint(&name, "shared_{s}", .{stage}) else try std.fmt.bufPrint(&name, "expert_{d}_{s}", .{ actual_ids[slot], stage }); + differs = (try moeMetric(u16, gpa, report, count, label, comparison, reuse, &@field(actual.*, stage)[slot], &@field(expected.*, stage)[reference_slot.?])) or differs; + } + if (differs and first.* == null) first.* = stage; + } + if (try moeMetric(u16, gpa, report, count, "combined_branch", comparison, reuse, &actual.branch, &expected.branch)) + if (first.* == null) { + first.* = "combined_branch"; + }; +} + +/// Retained local snapshot isolation. MoE has no history; qualification computes every value independently. +pub fn runMoe(gpa: std.mem.Allocator, io: std.Io, model: *Model, snapshot: []const u8, dir: []const u8, layer: usize) !void { + if (layer >= model.weights.layers.len) return error.InvalidMoeLayer; + const blob = try std.Io.Dir.cwd().readFileAlloc(io, snapshot, gpa, .limited(8 << 20)); + defer gpa.free(blob); + const snapshot_offset = (layer * 2 + 1) * 61488; + if (blob.len != 80 * 61488 or blob.len < snapshot_offset + 61488) return error.InvalidMoeSnapshot; + try model.reset(); + const input = model.workspace.view(.normalized); + @memcpy(input.buffer.contents()[input.offset..][0..10240], blob[snapshot_offset + 40960 ..][0..10240]); + const input_row = diagnosticRow(u16, input, 5120); + var retained_ids: [6]u32 = undefined; + var retained_weights: [6]f32 = undefined; + var retained_branch: [5120]u16 = undefined; + for (&retained_ids, 0..) |*id, i| id.* = std.mem.readInt(u32, blob[snapshot_offset + 61440 + 4 * i ..][0..4], .little); + for (&retained_weights, 0..) |*weight, i| + weight.* = @bitCast(std.mem.readInt(u32, blob[snapshot_offset + 61464 + 4 * i ..][0..4], .little)); + @memcpy(std.mem.asBytes(&retained_branch), blob[snapshot_offset + 51200 ..][0..10240]); + for (retained_ids, retained_weights, 0..) |id, weight, i| { + if (id >= 384 or !std.math.isFinite(weight)) return error.InvalidMoeRetainedRouting; + for (retained_ids[0..i]) |prior| if (id == prior) return error.InvalidMoeRetainedRouting; + } + const native = try gpa.create(MoeRows); + defer gpa.destroy(native); + const source = try gpa.create(MoeRows); + defer gpa.destroy(source); + const native_weights_reference = try gpa.create(MoeRows); + defer gpa.destroy(native_weights_reference); + const retained_reference = try gpa.create(MoeRows); + defer gpa.destroy(retained_reference); + const native_inputs_reference = try gpa.create(MoeRows); + defer gpa.destroy(native_inputs_reference); + const retained_native = try gpa.create(MoeRows); + defer gpa.destroy(retained_native); + var scratch = try projection.Scratch.init(gpa, 5120, 5120); + defer scratch.deinit(); + + try moeEncodeNative(model, layer, true); + moeCaptureNative(model, native); + var native_ids: [6]u32 = undefined; + var native_weights: [6]f32 = undefined; + var native_logits_routing: moe_reference.MoeRouting = .{}; + @memcpy(&native_ids, diagnosticRow(u32, model.workspace.view(.route_ids), 6)); + @memcpy(&native_weights, diagnosticRow(f32, model.workspace.view(.route_weights), 6)); + @memcpy(&native_logits_routing.logits, diagnosticRow(f32, model.workspace.view(.router_logits), 384)); + moe_reference.moeSelect(model.weights.layers[layer].moe, model.cfg.text_config, &native_logits_routing.logits, &native_logits_routing.scores, &native_logits_routing.ids, &native_logits_routing.weights); + var source_routing: moe_reference.MoeRouting = .{}; + math.dense(model.weights.layers[layer].moe.gate, input_row, &source_routing.logits); + moe_reference.moeSelect(model.weights.layers[layer].moe, model.cfg.text_config, &source_routing.logits, &source_routing.scores, &source_routing.ids, &source_routing.weights); + try moeReferenceRows(model, layer, input_row, &source_routing.ids, &source_routing.weights, source, &scratch, null, null); + try moeReferenceRows(model, layer, input_row, &native_ids, &native_weights, native_weights_reference, &scratch, source, &source_routing.ids); + try moeReferenceRows(model, layer, input_row, &retained_ids, &retained_weights, retained_reference, &scratch, native_weights_reference, &native_ids); + for (0..7) |slot| { + native_inputs_reference.gate[slot] = native.gate[slot]; + native_inputs_reference.up[slot] = native.up[slot]; + moe_reference.moeActivation(&native.gate[slot], &native.up[slot], if (slot == 6) 1 else native_weights[slot], model.cfg.text_config, &native_inputs_reference.activation[slot]); + // Deliberately use the captured native activation, not the preceding reference activation. + if (std.mem.eql(u16, &native.activation[slot], &native_weights_reference.activation[slot])) { + native_inputs_reference.down[slot] = native_weights_reference.down[slot]; + } else { + const matrices = try moeMatrices(model, layer, if (slot == 6) null else native_ids[slot]); + try projection.apply(.{ .trellis = matrices[2] }, &native.activation[slot], &native_inputs_reference.down[slot], &scratch); + } + } + moeCombine(&native_ids, native_inputs_reference); + @memcpy(diagnosticRow(u32, model.workspace.view(.route_ids), 6), &retained_ids); + @memcpy(diagnosticRow(f32, model.workspace.view(.route_weights), 6), &retained_weights); + try moeEncodeNative(model, layer, false); + moeCaptureNative(model, retained_native); + + var report: std.ArrayList(u8) = .empty; + defer report.deinit(gpa); + const header = try std.json.Stringify.valueAlloc(gpa, .{ + .mode = "moe-sameinput", + .token = @as(?usize, null), + .layer = layer, + .token_identity_source = "caller_supplied_snapshot_argument_raw_snapshot_has_no_token_header", + .retained_route_origin = "native_GPU_snapshot_command_contract", + .snapshot_count = 80, + .snapshot_length_bytes = 61488, + .snapshot_index = layer * 2 + 1, + .snapshot_stage = "moe", + .input_offset_bytes = snapshot_offset + 40960, + .input_length_bytes = 10240, + .retained_ids_offset_bytes = snapshot_offset + 61440, + .retained_weights_offset_bytes = snapshot_offset + 61464, + .identical_normalized_input = true, + .reference_dot = "independent_EXL3_FP64", + .qualification_gate_applied = false, + .qualification_oracle_modified_to_reuse_native_values = false, + .official_revision = "2cba9e42aa026125f3ed06c6d98c1db82f7ca027", + .source_ids = source_routing.ids, + .source_weights = source_routing.weights, + .native_ids = native_ids, + .native_weights = native_weights, + .retained_ids = retained_ids, + .retained_weights = retained_weights, + .source_ids_equal_native_order = std.mem.eql(u32, &source_routing.ids, &native_ids), + .native_ids_equal_retained_order = std.mem.eql(u32, &native_ids, &retained_ids), + .source_select_on_native_logits_ids = native_logits_routing.ids, + .source_select_on_native_logits_weights = native_logits_routing.weights, + .score_receipt_origin = "source_sqrt_softplus_applied_to_captured_native_logits_not_GPU_score_capture", + .expert_receipt_layout = "slot_major_top6_ids_order_then_shared_slot6", + .expert_sum_order = "ascending_expert_id_FP32_shared_last_final_BF16", + .receipt_prefixes = .{ "moe-native", "moe-source", "moe-native-weights-reference-diagnostic", "moe-native-inputs-reference-diagnostic", "moe-retained-weights-reference-diagnostic", "moe-retained-route-native-diagnostic" }, + .receipt_shapes = .{ .gate = .{ 7, 2304 }, .up = .{ 7, 2304 }, .activation = .{ 7, 2304 }, .down = .{ 7, 5120 }, .branch = .{5120} }, + .native_inputs_reference_contract = "activation:reference_on_native_gate_up_and_weight;down:FP64_on_native_activation;branch:reference_combine_of_these_down_outputs", + }, .{}); + defer gpa.free(header); + try report.appendSlice(gpa, header[0 .. header.len - 1]); + try report.appendSlice(gpa, ",\"cases\":[\n"); + var count: usize = 0; + var first: MoeFirst = .{}; + if (try moeMetric(f32, gpa, &report, &count, "route_logits", "source_chain", false, &native_logits_routing.logits, &source_routing.logits)) first.source_chain = "route_logits"; + if (try moeMetric(f32, gpa, &report, &count, "route_scores_source_transform", "source_score_transform_on_captured_native_logits_diagnostic", true, &native_logits_routing.scores, &source_routing.scores)) + if (first.source_chain == null) { + first.source_chain = "route_scores_source_transform"; + }; + if (!std.mem.eql(u32, &source_routing.ids, &native_ids) and first.source_chain == null) first.source_chain = "route_ids"; + if (std.mem.eql(u32, &source_routing.ids, &native_ids)) { + if (try moeMetric(f32, gpa, &report, &count, "route_weights", "source_chain", false, &native_weights, &source_routing.weights)) + if (first.source_chain == null) { + first.source_chain = "route_weights"; + }; + } + if (std.mem.eql(u32, &native_logits_routing.ids, &native_ids)) + _ = try moeMetric(f32, gpa, &report, &count, "route_weights", "source_select_on_same_native_logits_diagnostic", true, &native_weights, &native_logits_routing.weights); + if (std.mem.eql(u32, &retained_ids, &native_ids)) + _ = try moeMetric(f32, gpa, &report, &count, "route_weights", "retained_native_weights_vs_reexecuted_native", true, &native_weights, &retained_weights); + try moeCompareRows(gpa, &report, &count, "source_chain", false, native, &native_ids, source, &source_routing.ids, &first.source_expert_stage); + if (first.source_chain == null) first.source_chain = first.source_expert_stage; + try moeCompareRows(gpa, &report, &count, "same_native_ids_weights_reference_chain_diagnostic", true, native, &native_ids, native_weights_reference, &native_ids, &first.native_weights_reference_chain_diagnostic); + inline for (.{ "activation", "down" }) |stage| { + var differs = false; + for (0..7) |slot| { + var name: [64]u8 = undefined; + const label = if (slot == 6) try std.fmt.bufPrint(&name, "shared_{s}", .{stage}) else try std.fmt.bufPrint(&name, "expert_{d}_{s}", .{ native_ids[slot], stage }); + differs = (try moeMetric(u16, gpa, &report, &count, label, if (comptime std.mem.eql(u8, stage, "activation")) "same_native_gate_up_weight_reference_activation_diagnostic" else "same_native_activation_FP64_down_diagnostic", true, &@field(native.*, stage)[slot], &@field(native_inputs_reference.*, stage)[slot])) or differs; + } + if (differs and first.native_input_diagnostic == null) first.native_input_diagnostic = stage; + } + _ = try moeMetric(u16, gpa, &report, &count, "combined_branch", "same_native_activation_FP64_down_then_reference_combine_diagnostic", true, &native.branch, &native_inputs_reference.branch); + // Isolate combination arithmetic independently of every projection and route-weight difference. + var native_down_reference_branch: [5120]u16 = undefined; + moeCombineDown(&native_ids, &native.down, &native_down_reference_branch); + first.native_down_reference_combine_diagnostic_differs = + try moeMetric(u16, gpa, &report, &count, "combined_branch", "same_native_down_reference_combine_diagnostic", true, &native.branch, &native_down_reference_branch); + try moeWrite(gpa, io, dir, "moe-native-down-reference-combine-diagnostic.bf16", std.mem.asBytes(&native_down_reference_branch)); + try moeCompareRows(gpa, &report, &count, "same_retained_native_ids_weights_reference_chain_diagnostic", true, retained_native, &retained_ids, retained_reference, &retained_ids, &first.retained_weights_reference_chain_diagnostic); + _ = try moeMetric(u16, gpa, &report, &count, "combined_branch", "retained_native_branch_vs_reexecuted_native_routing", true, &native.branch, &retained_branch); + _ = try moeMetric(u16, gpa, &report, &count, "combined_branch", "retained_native_branch_vs_reexecuted_retained_routes_diagnostic", true, &retained_native.branch, &retained_branch); + const finding = try std.json.Stringify.valueAlloc(gpa, first, .{}); + defer gpa.free(finding); + try report.appendSlice(gpa, "\n],\"first_differing_intermediate\":"); + try report.appendSlice(gpa, finding); + try report.appendSlice(gpa, "}\n"); + try moeWrite(gpa, io, dir, "moe-sameinput.json", report.items); + inline for (.{ "moe-native", "moe-source", "moe-native-weights-reference-diagnostic", "moe-native-inputs-reference-diagnostic", "moe-retained-weights-reference-diagnostic", "moe-retained-route-native-diagnostic" }, .{ native, source, native_weights_reference, native_inputs_reference, retained_reference, retained_native }) |prefix, rows| + try moeWriteRows(gpa, io, dir, prefix, rows); + inline for (.{ "moe-input.bf16", "moe-retained-branch.bf16", "moe-native-logits.f32", "moe-source-logits.f32", "moe-native-logits-source-scores.f32", "moe-source-scores.f32", "moe-native-ids.u32", "moe-source-ids.u32", "moe-retained-ids.u32", "moe-native-weights.f32", "moe-source-weights.f32", "moe-retained-weights.f32", "moe-native-logits-source-select-ids.u32", "moe-native-logits-source-select-weights.f32" }, .{ std.mem.sliceAsBytes(input_row), std.mem.asBytes(&retained_branch), std.mem.asBytes(&native_logits_routing.logits), std.mem.asBytes(&source_routing.logits), std.mem.asBytes(&native_logits_routing.scores), std.mem.asBytes(&source_routing.scores), std.mem.asBytes(&native_ids), std.mem.asBytes(&source_routing.ids), std.mem.asBytes(&retained_ids), std.mem.asBytes(&native_weights), std.mem.asBytes(&source_routing.weights), std.mem.asBytes(&retained_weights), std.mem.asBytes(&native_logits_routing.ids), std.mem.asBytes(&native_logits_routing.weights) }) |name, bytes| + try moeWrite(gpa, io, dir, name, bytes); +} diff --git a/zig/tests/deepseek_v41_reference/operators.zig b/zig/tests/deepseek_v41_reference/operators.zig new file mode 100644 index 0000000000..d7c8cb65de --- /dev/null +++ b/zig/tests/deepseek_v41_reference/operators.zig @@ -0,0 +1,92 @@ +//! Actual checkpoint projections on identical deterministic BF16 rows, against independent FP64 dot products. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const projection = @import("projection.zig"); +const math = @import("math.zig"); +const Model = tf.deepseek_v41.Model; +const ex = tf.exl3_projection; +const Case = struct { name: []const u8, matrix: ex.Matrix }; + +pub fn run(gpa: std.mem.Allocator, io: std.Io, model: *Model, dir: []const u8) !void { + const first = &model.weights.layers[0]; + const cases = [_]Case{ + .{ .name = "wq_a", .matrix = first.attn.wq_a }, + .{ .name = "wq_b", .matrix = first.attn.wq_b }, + .{ .name = "wkv", .matrix = first.attn.wkv }, + .{ .name = "wo_a_group0", .matrix = first.attn.wo_a[0] }, + .{ .name = "wo_b", .matrix = first.attn.wo_b }, + .{ .name = "engram_wkv", .matrix = model.weights.layers[1].engram.?.wkv }, + .{ .name = "routed_gate", .matrix = try tf.deepseek_v41_weights.lookup(&model.checkpoint, "layers.0.ffn.experts.0.w1", 5120, 2304) }, + .{ .name = "routed_up", .matrix = try tf.deepseek_v41_weights.lookup(&model.checkpoint, "layers.0.ffn.experts.0.w3", 5120, 2304) }, + .{ .name = "routed_down", .matrix = try tf.deepseek_v41_weights.lookup(&model.checkpoint, "layers.0.ffn.experts.0.w2", 2304, 5120) }, + .{ .name = "head", .matrix = model.weights.head }, + }; + const x = try model.device.buffer(8192 * 2, mtl.ResourceOptions.shared); + defer x.deinit(); + const y = try model.device.buffer(2 * 129280 * 4, mtl.ResourceOptions.shared); + defer y.deinit(); + const scratch = try model.device.buffer(try ex.scratchBytes(8192, 129280, 1), mtl.ResourceOptions.shared); + defer scratch.deinit(); + for (x.slice(u16, 8192), 0..) |*value, i| + value.* = math.b(@as(f32, @floatFromInt(@as(i32, @intCast((i * 17) % 97)) - 48)) / 32); + for ([_]mtl.Buffer{ x, y, scratch }) |buffer| model.residency.add(buffer); + model.residency.commit(); + model.residency.requestResidency(); + defer { + for ([_]mtl.Buffer{ x, y, scratch }) |buffer| model.residency.remove(buffer); + model.residency.commit(); + } + var reference = try projection.Scratch.init(gpa, 8192, 129280); + defer reference.deinit(); + const expected = try gpa.alloc(f32, 129280); + defer gpa.free(expected); + var report: std.ArrayList(u8) = .empty; + defer report.deinit(gpa); + try report.appendSlice(gpa, "{\"input\":\"bf16((17*i%97-48)/32)\",\"reference_dot\":\"FP64\",\"cases\":["); + var all_pass = true; + for (cases, 0..) |case, index| { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const matrix = case.matrix; + const cb = model.queue.commandBufferUnretained(); + const enc = cb.compute(.serial); + var encoding = true; + errdefer if (encoding) enc.end(); + ex.useMatrix(enc, matrix); + try ex.encode(&model.pipelines.projection, enc, matrix, .{ .buffer = x }, .{ .buffer = y }, .{ .buffer = scratch }, .{ .count = 1, .x_stride = matrix.k, .y_stride = matrix.n, .output = .f32 }); + enc.end(); + encoding = false; + cb.commit(); + cb.wait(); + if (cb.failure() != null) return error.OperatorGpuFailure; + const actual = y.slice(f32, 129280); + try projection.applyF32(.{ .trellis = matrix }, x.slice(u16, matrix.k), expected[0..matrix.n], &reference); + var max_abs: f64 = 0; + var worst: usize = 0; + var squared: f64 = 0; + var magnitude: f64 = 0; + for (actual[0..matrix.n], expected[0..matrix.n], 0..) |a, b, column| { + if (!std.math.isFinite(a) or !std.math.isFinite(b)) return error.NonfiniteOperatorOutput; + const delta: f64 = @as(f64, a) - b; + if (@abs(delta) > max_abs) { + max_abs = @abs(delta); + worst = column; + } + squared += delta * delta; + magnitude += @as(f64, b) * b; + } + const nrms = if (magnitude == 0) @sqrt(squared) else @sqrt(squared / magnitude); + all_pass = all_pass and nrms <= 1e-4; + const line = try std.fmt.allocPrint(gpa, "{{\"name\":\"{s}\",\"k\":{d},\"n\":{d},\"normalized_rms_error\":{d},\"max_absolute_error\":{d},\"worst_column\":{d},\"worst_gpu\":{d},\"worst_reference\":{d}}}", .{ case.name, matrix.k, matrix.n, nrms, max_abs, worst, actual[worst], expected[worst] }); + defer gpa.free(line); + if (index > 0) try report.appendSlice(gpa, ","); + try report.appendSlice(gpa, line); + std.debug.print("{s}\n", .{line}); + } + try report.appendSlice(gpa, "]}\n"); + const path = try std.fs.path.join(gpa, &.{ dir, "operators.json" }); + defer gpa.free(path); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = path, .data = report.items }); + if (!all_pass) return error.OperatorReferenceMismatch; +} diff --git a/zig/tests/deepseek_v41_reference/profile.zig b/zig/tests/deepseek_v41_reference/profile.zig new file mode 100644 index 0000000000..0fb15a8b90 --- /dev/null +++ b/zig/tests/deepseek_v41_reference/profile.zig @@ -0,0 +1,52 @@ +//! Real-checkpoint calibrated GPU stage timings, with profile-on/off logits checked before reporting. +const std = @import("std"); +const tf = @import("tensorfold"); +const counters = tf.deepseek_v41_profiler; +const Model = tf.deepseek_v41.Model; + +pub fn run(gpa: std.mem.Allocator, io: std.Io, model: *Model, ids: []const u32, dir: []const u8) !void { + if (ids.len == 0 or ids.len > model.workspace.plan.rows) return error.InvalidProfileWindow; + const previous_grouping = model.grouped_prefill; + defer model.grouped_prefill = previous_grouping; + model.grouped_prefill = false; + try model.reset(); + const baseline = try model.forward(ids, .all); + const baseline_gpu_seconds = model.last_gpu_seconds; + const expected = try gpa.dupe(f32, baseline); + defer gpa.free(expected); + model.grouped_prefill = true; + try model.reset(); + const candidate = try model.forward(ids, .all); + const candidate_gpu_seconds = model.last_gpu_seconds; + try check(expected, candidate); + var recorder = try counters.Recorder.init(gpa, model.device, .{ .enabled = true, .max_spans = 512 }); + defer recorder.deinit() catch unreachable; + const previous_profile = model.profile; + model.profile = &recorder; + defer model.profile = previous_profile; + try model.reset(); + const actual = try model.forward(ids, .last); + try check(expected[expected.len - 129280 ..], actual); + const totals = recorder.aggregates(); + var sampled_gpu_seconds: f64 = 0; + for (totals) |value| sampled_gpu_seconds += value.gpu_seconds; + var text: std.ArrayList(u8) = .empty; + defer text.deinit(gpa); + try text.print(gpa, "{{\"tokens\":{d},\"comparison_logits\":\"all\",\"candidate_grouping\":true,\"same_logits\":true,\"baseline_command_gpu_seconds\":{d},\"candidate_command_gpu_seconds\":{d},\"profiled_logits\":\"last\",\"sampled_span_gpu_seconds\":{d},\"stages\":[", .{ ids.len, baseline_gpu_seconds, candidate_gpu_seconds, sampled_gpu_seconds }); + for (std.enums.values(counters.Stage), totals, 0..) |stage, value, i| { + try text.print(gpa, "{s}{{\"stage\":\"{s}\",\"calls\":{d},\"gpu_seconds\":{d}}}", .{ if (i == 0) "" else ",", @tagName(stage), value.calls, value.gpu_seconds }); + } + try text.appendSlice(gpa, "]}\n"); + const path = try std.fs.path.join(gpa, &.{ dir, "gpu-profile.json" }); + defer gpa.free(path); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = path, .data = text.items }); + std.debug.print("{s}", .{text.items}); +} + +fn check(expected: []const f32, actual: []const f32) !void { + if (expected.len != actual.len) return error.ProfileLogitShapeMismatch; + for (actual, expected) |a, b| { + if (!std.math.isFinite(a) or !std.math.isFinite(b)) return error.NonfiniteProfileLogits; + if (@as(u32, @bitCast(a)) != @as(u32, @bitCast(b))) return error.ProfileChangedLogits; + } +} diff --git a/zig/tests/deepseek_v41_reference/projection.zig b/zig/tests/deepseek_v41_reference/projection.zig new file mode 100644 index 0000000000..03174c9459 --- /dev/null +++ b/zig/tests/deepseek_v41_reference/projection.zig @@ -0,0 +1,311 @@ +//! Independent scalar EXL3 oracle over resident bytes; FP64 dots and FP32 rotation/output boundaries, no GPU work. + +// 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 format = exl3.format; +const block_size = format.hadamard_dim; +const inverse: f32 = 0.08838834764831845; +const states_count = 65536; + +/// Reuse bounded vector/codebook storage for heterogeneous matrices; never allocate decoded W[K,N]. +pub const Scratch = struct { + gpa: std.mem.Allocator, + rotated: []f64, + sums: []f64, + codebooks: []f16, + + pub fn init(gpa: std.mem.Allocator, max_k: usize, max_n: usize) !Scratch { + if (max_k == 0 or max_n == 0) return error.InvalidReferenceProjectionCapacity; + const rotated = try gpa.alloc(f64, max_k); + errdefer gpa.free(rotated); + const sums = try gpa.alloc(f64, max_n); + errdefer gpa.free(sums); + const codebooks = try gpa.alloc(f16, 3 * states_count); + errdefer gpa.free(codebooks); + for ([_]format.Codebook{ .inst3, .mcg, .mul1 }, 0..) |book, bi| { + for (0..states_count) |state| { + codebooks[bi * states_count + state] = format.codebookValue(book, @intCast(state)); + } + } + return .{ .gpa = gpa, .rotated = rotated, .sums = sums, .codebooks = codebooks }; + } + + pub fn deinit(s: *Scratch) void { + s.gpa.free(s.rotated); + s.gpa.free(s.sums); + s.gpa.free(s.codebooks); + s.* = undefined; + } +}; + +/// 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) + return error.InvalidReferenceProjectionView; + const address = view.buffer.contents(); + if (@intFromPtr(address) == 0) return error.NonCpuAccessibleExl3Buffer; + return address[view.offset..][0..bytes]; +} + +fn scaleBytes(count: usize, kind: exl3.ScaleFormat) !usize { + return std.math.mul(usize, if (kind == .packed_signs) count / 16 else count, 2); +} + +fn half(bytes: []const u8, i: usize) f32 { + const value: f16 = @bitCast(std.mem.readInt(u16, bytes[2 * i ..][0..2], .little)); + return @floatCast(value); +} + +fn scale(bytes: []const u8, i: usize, kind: exl3.ScaleFormat) f32 { + if (kind == .f16) return half(bytes, i); + const word = std.mem.readInt(u16, bytes[2 * (i / 16) ..][0..2], .little); + 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, + cursor: usize = 0, + pending: u32 = 0, + available: u6 = 0, + + fn take(s: *Stream, count: u5) u16 { + if (s.available == 0) { + s.pending = std.mem.readInt(u32, s.bytes[s.cursor..][0..4], .little); + s.cursor += 4; + s.available = 32; + } + const first: u5 = @intCast(@min(s.available, count)); + const head: u32 = s.pending >> @as(u5, @intCast(32 - @as(u6, first))); + s.pending <<= first; + s.available -= first; + if (first == count) return @intCast(head); + const rest = count - first; + s.pending = std.mem.readInt(u32, s.bytes[s.cursor..][0..4], .little); + s.cursor += 4; + const tail = s.pending >> @as(u5, @intCast(32 - @as(u6, rest))); + s.pending <<= rest; + s.available = 32 - @as(u6, rest); + return @intCast((head << rest) | tail); + } +}; + +/// Circular initial state uses final 16 stream bits; decode only a sixteen-column accumulator and live stream. +fn accumulateTile(comptime halves: u16, tile: []const u8, table: []const f16, x: []const f64, acc: *[16]f64) void { + @setFloatMode(.strict); + var stream: Stream = .{ .bytes = tile }; + var state: u16 = @truncate(std.mem.readInt(u32, tile[tile.len - 4 ..][0..4], .little)); + var p: usize = 0; + while (p < 256) : (p += 1) { + // Odd half-widths alternate floor(bits),ceil(bits), exactly as format.streamEnds specifies. + const advance: u4 = @intCast(halves / 2 + (if (halves % 2 == 1) p % 2 else 0)); + state = (state << advance) | stream.take(advance); + const lane = p / 8; + const j = p % 8; + const row = 2 * (lane % 4) + (j & 1) + 8 * ((j >> 1) & 1); + const col = lane / 4 + 8 * (j >> 2); + const weight: f64 = @floatCast(table[state]); + acc[col] += x[row] * weight; + } +} + +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]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; + } + } +} + +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(); + var words: [128]i16 = undefined; + for (&words, 0..) |*word, i| word.* = @bitCast(@as(u16, @truncate(i * 40503 + 17329))); + var x: [16]f64 = undefined; + for (&x, 0..) |*value, i| value.* = @as(f64, @floatFromInt(i + 1)) * 0.0625; + var placement: [256]format.Placement = undefined; + format.tilePlacement(&placement); + inline for (2..17) |halves| { + const width: format.Width = .{ .halves = halves }; + const tile = words[0..width.tileWords()]; + var states: [256]u32 = undefined; + format.tileStates(width, tile, &states); + for ([_]format.Codebook{ .inst3, .mcg, .mul1 }, 0..) |book, bi| { + if (width.isHalf() and (book != .mul1 or halves > 7)) continue; + const table = scratch.codebooks[bi * states_count ..][0..states_count]; + var expected: [16]f64 = @splat(0); + for (states, placement) |state, where| { + const weight: f64 = @floatCast(format.codebookValue(book, @intCast(state))); + expected[where.col] += x[where.row] * weight; + } + var actual: [16]f64 = @splat(0); + accumulateTile(halves, std.mem.sliceAsBytes(tile), table, &x, &actual); + try std.testing.expectEqualSlices(f64, &expected, &actual); + } + } +} + +test "FP32 Hadamard then FP16 and BF16 boundaries are deterministic" { + 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); + 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)))); +} diff --git a/zig/tests/deepseek_v41_reference/qualify.zig b/zig/tests/deepseek_v41_reference/qualify.zig new file mode 100644 index 0000000000..e5190ee546 --- /dev/null +++ b/zig/tests/deepseek_v41_reference/qualify.zig @@ -0,0 +1,97 @@ +//! Actual checkpoint teacher comparison: independent CPU arithmetic versus the resident GPU decode path. +const std = @import("std"); +const mtl = @import("metal"); +const Model = @import("tensorfold").deepseek_v41.Model; +const Reference = @import("reference.zig").Reference; +const nrms_limit: f64 = 0.01; + +pub const Result = struct { + tokens: usize, + max_absolute_error: f64, + rms_error: f64, + normalized_rms_error: f64, + greedy_disagreements: usize, + gpu_seconds: f64, + cpu_seconds: f64, + normalized_rms_limit: f64 = nrms_limit, + numeric_gate_passed: bool = false, + greedy_gate_passed: bool = false, +}; + +fn maximum(row: []const f32) !u32 { + var id: u32 = 0; + var best = -std.math.inf(f32); + for (row, 0..) |v, i| { + if (!std.math.isFinite(v)) return error.NonfiniteTeacherLogits; + if (v > best) { + best = v; + id = @intCast(i); + } + } + return id; +} + +pub fn run(gpa: std.mem.Allocator, io: std.Io, model: *Model, ids: []const u32, token_map: []const u8, dir: []const u8) !Result { + if (ids.len == 0 or ids.len > 16) return error.InvalidOracleWindow; + var reference = try Reference.init(gpa, model.cfg.text_config, &model.weights, &model.checkpoint, token_map, @intCast(ids.len)); + defer reference.deinit(); + const cpu = try gpa.alloc(f32, ids.len * 129280); + defer gpa.free(cpu); + const gpu = try gpa.alloc(f32, cpu.len); + defer gpa.free(gpu); + try model.reset(); + var result = Result{ .tokens = ids.len, .max_absolute_error = 0, .rms_error = 0, .normalized_rms_error = 0, .greedy_disagreements = 0, .gpu_seconds = 0, .cpu_seconds = 0 }; + var squared_error: f64 = 0; + var squared_reference: f64 = 0; + for (ids, 0..) |id, i| { + const actual = try model.forward(&.{id}, .all); + result.gpu_seconds += model.last_gpu_seconds; + @memcpy(gpu[i * 129280 ..][0..129280], actual); + const start = mtl.clock.seconds(); + const expected = try reference.step(id); + result.cpu_seconds += mtl.clock.seconds() - start; + @memcpy(cpu[i * 129280 ..][0..129280], expected); + if (try maximum(actual) != try maximum(expected)) result.greedy_disagreements += 1; + for (actual, expected) |a, b| { + if (!std.math.isFinite(a) or !std.math.isFinite(b)) return error.NonfiniteTeacherLogits; + const error_: f64 = @as(f64, a) - b; + result.max_absolute_error = @max(result.max_absolute_error, @abs(error_)); + squared_error += error_ * error_; + squared_reference += @as(f64, b) * b; + } + std.debug.print("oracle teacher {d}/{d}: GPU token={d}, CPU token={d}\n", .{ i + 1, ids.len, try maximum(actual), try maximum(expected) }); + } + result.rms_error = @sqrt(squared_error / @as(f64, @floatFromInt(cpu.len))); + result.normalized_rms_error = if (squared_reference == 0) result.rms_error else @sqrt(squared_error / squared_reference); + result.numeric_gate_passed = result.normalized_rms_error <= result.normalized_rms_limit; + result.greedy_gate_passed = result.greedy_disagreements == 0; + const cp = try std.fs.path.join(gpa, &.{ dir, "oracle-logits.f32" }); + defer gpa.free(cp); + const gp = try std.fs.path.join(gpa, &.{ dir, "gpu-logits.f32" }); + defer gpa.free(gp); + const report = try std.fs.path.join(gpa, &.{ dir, "oracle.json" }); + defer gpa.free(report); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = cp, .data = std.mem.sliceAsBytes(cpu) }); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = gp, .data = std.mem.sliceAsBytes(gpu) }); + const json = try std.json.Stringify.valueAlloc(gpa, result, .{}); + defer gpa.free(json); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = report, .data = json }); + std.debug.print("{s}\n", .{json}); + try enforceGate(result.normalized_rms_error, result.greedy_disagreements); + return result; +} + +fn enforceGate(normalized_rms_error: f64, greedy_disagreements: usize) !void { + if (!std.math.isFinite(normalized_rms_error)) return error.NonfiniteTeacherLogits; + if (normalized_rms_error > nrms_limit) return error.TeacherReferenceMismatch; + if (greedy_disagreements != 0) return error.TeacherGreedyMismatch; +} + +test "teacher numeric gate remains independent of greedy agreement" { + try std.testing.expectError(error.TeacherReferenceMismatch, enforceGate(0.04, 0)); + try std.testing.expectError(error.TeacherReferenceMismatch, enforceGate(0.010000001, 1)); + try std.testing.expectError(error.TeacherGreedyMismatch, enforceGate(0.01, 1)); + try enforceGate(0.01, 0); + try std.testing.expectError(error.NonfiniteTeacherLogits, enforceGate(std.math.nan(f64), 0)); + try std.testing.expectError(error.NonfiniteTeacherLogits, enforceGate(std.math.inf(f64), 0)); +} diff --git a/zig/tests/deepseek_v41_reference/reference.zig b/zig/tests/deepseek_v41_reference/reference.zig new file mode 100644 index 0000000000..98069068f7 --- /dev/null +++ b/zig/tests/deepseek_v41_reference/reference.zig @@ -0,0 +1,176 @@ +//! Independent wide-sum qualification with explicit FP32/BF16 boundaries; not GPU self-reference or CUDA emulation. + +// Equations: pinned DeepSeek AI inference/{model,engram,kernel}.py; provenance in THIRD_PARTY_NOTICES.md. +const std = @import("std"); +const tf = @import("tensorfold"); +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 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 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; + } + } + 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 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, + }; + + /// 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); + } + 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); + } +}; diff --git a/zig/tests/deepseek_v41_reference/rope_frequencies.json b/zig/tests/deepseek_v41_reference/rope_frequencies.json new file mode 100644 index 0000000000..2ffe1d1556 --- /dev/null +++ b/zig/tests/deepseek_v41_reference/rope_frequencies.json @@ -0,0 +1,14 @@ +{ + "reference": "deepseek-ai/DeepSeek-V4.1-Flash", + "revision": "2cba9e42aa026125f3ed06c6d98c1db82f7ca027", + "source_sha256": "4e9ae23620edc8028ccc5d5fef552ab7fdc7dcd6f79608754fe9f67644056f65", + "torch_version": "2.14.1", + "extraction": "precompute_freqs_cis FP32 frequency vector before outer/polar; CPU", + "factor": 16, + "beta_fast": 32, + "beta_slow": 1, + "vectors": [ + {"theta": 10000, "original": 0, "bits": [1065353216,1061157137,1058010522,1054337179,1050798235,1047712803,1043732615,1040747895,1036831949,1033475085,1030116802,1026341397,1023510243,1019364175,1016180025,1012562930,1008981770,1005959600,1001931932,998911607,995049016,991652109,988353832,984533390,981668463,977573004,974350870,970790693,967166816,964208661,960132947,957076594]}, + {"theta": 160000, "original": 65536, "bits": [1065353216,1060112954,1056054303,1051098369,1046804782,1042117749,1037601972,1033169251,1028443341,1024251131,1019326489,1015361742,1010249151,1006366088,1001209182,997081655,991197921,985076960,979287478,973418702,966505391,959469944,952012877,943627103,933484300,917783504,913828774,908753896,904558591,899759068,895336240,890797133]} + ] +} diff --git a/zig/tests/deepseek_v41_reference/trace.zig b/zig/tests/deepseek_v41_reference/trace.zig new file mode 100644 index 0000000000..79cdef0fa4 --- /dev/null +++ b/zig/tests/deepseek_v41_reference/trace.zig @@ -0,0 +1,123 @@ +//! Bounded investigation snapshots; file I/O and independent CPU work are not performance measurements. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const Model = tf.deepseek_v41.Model; +const Reference = @import("reference.zig").Reference; +const Stage = tf.deepseek_v41_profiler.Stage; +const widths = [_]usize{ 20480 * 2, 5120 * 2, 5120 * 2, 6 * 4, 6 * 4 }; +const stride = 61488; +const bytes = stride * 80; + +const Capture = struct { + buffer: mtl.Buffer, + pipeline: mtl.Pipeline, + cpu: []u8, + + fn gpu(ctx: *anyopaque, enc: mtl.ComputeEncoder, layer: u32, stage: Stage, model: *const Model) void { + const self: *Capture = @ptrCast(@alignCast(ctx)); + var destination: usize = (@as(usize, layer) * 2 + @intFromBool(stage == .moe)) * stride; + const ws = &model.workspace; + inline for (.{ ws.view(.stream), ws.view(.normalized), ws.view(.branch), ws.view(.route_ids), ws.view(.route_weights) }, widths, 0..) |source, count, field| { + if (stage == .moe or field < 3) { + enc.setPipeline(self.pipeline); + enc.setBuffer(source.buffer, source.offset, 0); + enc.setBuffer(self.buffer, destination, 1); + enc.setValue(@as(u32, @intCast(count / 4)), 2); + enc.dispatchThreads(.{ .width = count / 4 }, .{ .width = 256 }); + destination += count; + } + } + enc.barrier(); + } + fn reference(ctx: *anyopaque, layer: usize, stage: Stage, ref: *const Reference) void { + const self: *Capture = @ptrCast(@alignCast(ctx)); + const start = (layer * 2 + @intFromBool(stage == .moe)) * stride; + var destination: []u8 = self.cpu[start..][0..stride]; + inline for (.{ std.mem.sliceAsBytes(&ref.h), std.mem.sliceAsBytes(&ref.input), std.mem.sliceAsBytes(&ref.output), std.mem.sliceAsBytes(&ref.last_routes), std.mem.sliceAsBytes(&ref.last_routing) }, 0..) |source, field| { + if (stage == .moe or field < 3) { + @memcpy(destination[0..source.len], source); + destination = destination[source.len..]; + } + } + } +}; +fn normalizedError(a: []const u8, b: []const u8) !f64 { + var squared: f64 = 0; + var denominator: f64 = 0; + var i: usize = 0; + while (i < a.len) : (i += 2) { + const x: f32 = @bitCast(@as(u32, std.mem.readInt(u16, a[i..][0..2], .little)) << 16); + const y: f32 = @bitCast(@as(u32, std.mem.readInt(u16, b[i..][0..2], .little)) << 16); + if (!std.math.isFinite(x) or !std.math.isFinite(y)) return error.NonfiniteTraceValues; + squared += (@as(f64, x) - y) * (@as(f64, x) - y); + denominator += @as(f64, y) * y; + } + return if (denominator == 0) @sqrt(squared) else @sqrt(squared / denominator); +} + +pub fn run(gpa: std.mem.Allocator, io: std.Io, model: *Model, ids: []const u32, token_map: []const u8, dir: []const u8) !void { + if (ids.len == 0 or ids.len > 16) return error.InvalidOracleWindow; + const lib = try mtl.Library.fromSource(model.device, "#include \nusing namespace metal;\nkernel void trace_copy(const device uint* x [[buffer(0)]], device uint* y [[buffer(1)]], constant uint& n [[buffer(2)]], uint i [[thread_position_in_grid]]) { if(i 0.01 or !routes_match) std.debug.print("trace token={d} layer={d} stage={s} stream={d} input={d} branch={d} routes_match={}\n", .{ token, point / 2, if (point % 2 == 0) "attention" else "moe", stream_error, input_error, branch_error, routes_match }); + } + inline for (.{ "gpu", "cpu" }, .{ gpu, cpu }) |name, data| { + const path = try std.fmt.allocPrint(gpa, "{s}/trace-{d}-{s}.bin", .{ dir, token, name }); + defer gpa.free(path); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = path, .data = data }); + } + std.debug.print("trace captured token {d}/{d}, same_logits=true\n", .{ token + 1, ids.len }); + } + try report.appendSlice(gpa, "\n]\n"); + const path = try std.fs.path.join(gpa, &.{ dir, "layer-trace.json" }); + defer gpa.free(path); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = path, .data = report.items }); +}