diff --git a/CHANGELOG.md b/CHANGELOG.md index 876346618c..c16e8c2ad6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,18 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 release notes](RELEASE-NOTES-1.0.0.md) for the native binary's supported models and migration details. +## Unreleased: experimental draft + +- Add an experimental DeepSeek V4.1 Flash resident Metal backend and independent + qualification runner. Whole-model numerical qualification remains failed: + retained eight-token NRMS 4.0960584769% against an unchanged 1% gate, despite + zero greedy disagreements. This is draft-only, not a qualified-model addition. + See the [implementation and qualification guide](docs/deepseek-v41.md). +- Harden resident buffer/view admission, invalid-route error propagation, mixed DSML/XML + whole/SSE containment and tokenizer failure reporting. Bounded actual serving, tool + calls, unsupported-feature refusal, disconnect recovery and zero-I/O request intervals + are documented separately from the still-failed whole-model numerical gate. + ## 1.0.7 - Living Weights learns each fact into about ten new neurons inside the model's own shards, on Macs and on a GB10. Every MoE layer's shared expert gets a pocket of 64 extra neurons; a fact takes one at each of about ten layers, fired at the first row of the reply (a fact about you) or at the subject's last token where a question names it (a fact about someone or something), with outputs fitted through the model so the answers come back, thinking on too, while near misses, free writing about the subject and other answers stay as they were. Nothing the model had is changed, no adapter is loaded, and any MLX loader reads the folder as a slightly wider model. A fact is kept only when held-out wordings bring it back and nothing else moved; otherwise it is taken out and leaves nothing behind. Documents are learned fact by fact: the model finds the facts, names the organisation on each, and keeps only those whose names, numbers and words the text holds (a list holding others is read again), then writes questions that name their subject (rewriting ones that ask for it), and leaves a question another fact answers to that fact. Three personal facts taught one after another each came back on both test questions on an M5 Ultra, an M3 Ultra and a GB10, with no leaks and no drift, live and after a plain restart; a six-fact handbook gave 3 to 6 of 6 answers across six runs, and long documents about one subject are still weak (6 of 16 from a 21-fact handbook). Without learned facts, replies and speed equal 1.0.6's on both backends; a folder with 21 facts decodes about 16% slower on a Mac. `--slide` refuses a folder taught by 1.0.3 to 1.0.6, whose lessons sit in the shared experts the pocket widens; such folders still serve without it. It builds on the CUDA learner and the take-out rule of #591, #599 and #632. Thanks to @ZackSample and @ThinkOffApp. diff --git a/README.md b/README.md index 4b45c06bcc..97ff5cfab8 100644 --- a/README.md +++ b/README.md @@ -30,6 +30,13 @@ curl -fsS http://127.0.0.1:8080/v1/chat/completions \ -d '{"model":"local-model","messages":[{"role":"user","content":"Say hello in one sentence."}],"max_tokens":128,"temperature":0}' ``` +## DeepSeek V4.1 contributor draft + +The [experimental resident Metal implementation](docs/deepseek-v41.md) is draft-only, +not a qualified model. Its retained independent whole-model comparison fails the +unchanged 1% numerical gate. See the guide for prerequisites, exact reproduction +inputs, limitations and the conditions required before readiness or merge. + ## Living Weights (experimental) Living Weights teaches a running model new facts. Start a Nemotron 3.5 Lightning server with `--slide`, tell it diff --git a/RUNBOOK.md b/RUNBOOK.md index ec04482c45..821bdde831 100644 --- a/RUNBOOK.md +++ b/RUNBOOK.md @@ -142,6 +142,20 @@ On one Mac, GLM loads only when its weights and caches fit in 70% of RAM, 179.2 `GLM_LOAD_LIMIT_GB=N` sets that limit in GB, up to the GPU's recommended working set; a value that isn't a number above 0 stops the server. Past 70% of RAM the server warns at startup: a load that wires most of the RAM can stall macOS until its watchdog restarts the Mac, so run nothing else large beside it. +## DeepSeek V4.1 Flash: experimental draft only + +The resident Metal implementation is not release-qualified. The retained independent +eight-token whole-model comparison fails the unchanged 1% NRMS gate at 4.0960584769%; +zero greedy disagreements do not clear it. Keep the proposal in draft. + +The [DS4.1 draft guide](docs/deepseek-v41.md) provides pinned checkpoint revisions, +resident-memory admission, supported/unsupported features, exact teacher IDs, build, +diagnostic and serving commands, and the readiness gates. +The [historical evidence record](docs/deepseek-v41-evidence.md) preserves numerical +counterexamples, rejected experiments, residency failures and limited speed receipts. +Historical receipts are not verification of the current-main integration. + + ## Access, context and updates To accept remote clients, choose `--host 0.0.0.0` and configure `--api-key-file` with a file readable only by its owner. diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index cb54012303..84395982f1 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -76,6 +76,35 @@ checked bit for bit against ExLlamaV3's dequantization. Flash Next's optional int8 and int4 KV caches (`families/qwen4_exp/cuda/kvcache.py`) follow the cache quantization scheme of [ExLlamaV3](https://github.com/turboderp-org/exllamav3) `-cq 8` and `-cq 4` (MIT License, Copyright (c) 2025 Turboderp, text below): groups of 32, one fp16 absmax scale per group, the group rotated by a 32-point Hadamard, midpoint-grid codes, `compand_a == 0`. 8-bit stores each code as a signed int8 (`q - 128`). 4-bit stores two unsigned codes per byte, low nibble first (the same bits as ExLlamaV3's little-endian packing, a uint8 tensor rather than their uint32 words). Their dequantizer folds another `1/sqrt(32)` into the scale and applies the unnormalized butterfly on the way out; this cache applies the normalized H32 to the query and to the merged output instead, and leaves the stored codes rotated. Scales match their quantizer bit for bit. Reconstructed values agree within fp16/bf16 rounding (under 0.01 on random groups), not bit for bit. The quantizer and the attention dequant are written for TensorFold and checked against an independent reference of that arithmetic. +## Native DeepSeek V4.1 and Metal EXL3 + +The native `zig/src/families/deepseek_v41/` implementation and source Metal kernels 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 offline Engram token-map builder +`tools/zig/dsv41_token_map.py` implements that revision's normalization, PCG64 multipliers and prime bucket layout. +The independent CPU qualification arithmetic is in `zig/tests/deepseek_v41_reference/`. + +`zig/src/server/deepseek_v41_official.jinja` is the unchanged official chat template from the same revision; +its text, token-ID and tool-decoding fixtures are generated from DeepSeek's official encoder. +The MIT license text below applies to these adaptations and the template. + +The bounded `zig/src/families/deepseek_v41/fixtures/hc_attention_rms.json` regression +contains model-derived numeric rows and independently captured operator boundaries from +that pinned MIT-licensed checkpoint. Its embedded provenance identifies the token, layer, +source revision, encodings and target columns; it is not a full model shard or trace. +`zig/tests/deepseek_v41_reference/rope_frequencies.json` contains 64 formula-derived +frequency values with extraction provenance. The draft's historical evidence is labeled +separately from current-branch verification. +The RMS mean-boundary test uses a constructed BF16 row `[3, 0, ...]` of width 5,120 +and a uniform constructed FP32 weight `0x3c65de29`, not a captured model coefficient. +Its expected BF16 tie was exercised through the pinned official `RMSNorm.forward` +on PyTorch `2.13.0a0+9186a08b2c.nv26.07` / CUDA 13.3. It tests an operator boundary, +not source-coefficient injection into whole-model qualification. + +The native `zig/src/core/exl3_projection*.zig` decoder and `zig/kernels/metal/core/exl3_projection.metal` +separately implement ExLlamaV3's format and codebooks under the Turboderp attribution above. +No captured CUDA/Triton binary kernel assets, full model shards or binary traces are bundled. + ## Vendored code and weights `src/tensorfold/drafters/vendor/z_lab_dflash/model_mlx.py` is the unmodified `dflash/model_mlx.py` from diff --git a/build.zig b/build.zig index 054e688280..aa22bbdfa6 100644 --- a/build.zig +++ b/build.zig @@ -5,6 +5,7 @@ const dist_build = @import("zig/build/dist.zig"); const qwen27_build = @import("zig/build/qwen27.zig"); const hip_build = @import("zig/build/hip.zig"); const grammar_build = @import("zig/build/grammar.zig"); +const dsv41_build = @import("zig/build/deepseek_v41.zig"); comptime { const required = std.mem.trim(u8, @embedFile(".zig-version"), "\r\n"); @@ -207,7 +208,10 @@ fn nativeServer(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b .link_libc = true, .imports = &.{ .{ .name = "engine_api", .module = api }, .{ .name = "tokenizer", .module = tokenizer }, .{ .name = "template", .module = template } }, }); - test_step.?.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests })).step); + const server_run = b.addRunArtifact(b.addTest(.{ .root_module = server_tests })); + test_step.?.dependOn(&server_run.step); + b.step("test-server", "HTTP routes, official templates and tool parsing without model weights") + .dependOn(&server_run.step); // the parse half's tests, as their own test binary (zig/src/server/tool_parse_test.zig) const tool_parse_tests = b.addTest(.{ .root_module = b.createModule(.{ .root_source_file = b.path("zig/src/server/tool_parse_test.zig"), .target = target, .optimize = .Debug, .link_libc = true, .imports = &.{ .{ .name = "engine_api", .module = api }, .{ .name = "tokenizer", .module = tokenizer }, .{ .name = "template", .module = template } } }) }); test_step.?.dependOn(&b.addRunArtifact(tool_parse_tests).step); @@ -381,6 +385,7 @@ fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b 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); + dsv41_build.checks(b, target, optimize, metal, sources, engine); const pack_cache_cases = b.addTest(.{ .root_module = b.createModule(.{ .root_source_file = b.path("zig/flashnext_pack_cache_cases.zig"), .target = target, @@ -392,6 +397,8 @@ fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b test_step.dependOn(&pack_cache_run.step); const engine_programs = [_]struct { name: []const u8, path: []const u8, about: []const u8, c_source: ?[]const u8 = null }{ .{ .name = "tf-qwen35-check", .path = "zig/tests/qwen35_check.zig", .about = "Qwen3.5-2B checkpoint and native operations" }, + .{ .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-exl3-bench", .path = "zig/exl3_projection_bench.zig", .about = "Actual GPU codebook parity and real-shape EXL3 decode performance" }, .{ .name = "tf-qwen35-forward", .path = "zig/tests/qwen35_forward.zig", .about = "Qwen3.5-2B native teacher-forced logits" }, .{ .name = "tf-qwen35-exact", .path = "zig/tests/qwen35_exact.zig", .about = "Qwen3.5-2B forward and committed-state exactness" }, .{ .name = "tensorfold", .path = "zig/src/main.zig", .about = "The native engine's command line" }, diff --git a/docs/deepseek-v41-evidence.md b/docs/deepseek-v41-evidence.md new file mode 100644 index 0000000000..605839ce4a --- /dev/null +++ b/docs/deepseek-v41-evidence.md @@ -0,0 +1,288 @@ +# DeepSeek V4.1 historical numerical and runtime evidence + +These are retained investigation receipts, not verification of the current-main draft. +The measured source was `bec00ae6961a43a9ab73c4fa9338c167b3f82851` plus the local +DS4.1 candidate changes; it was not a clean committed release. The draft integrates +that work onto main `e7c523ebfa16b8094cf9e35e3773a7d25befcc5b`. +Do not cite intermediate test counts or lower historical errors as the current gate. +The latest retained gate is **4.0960584769% NRMS against an unchanged 1% limit: FAIL**. +The [draft guide](deepseek-v41.md) states current boundaries and reproduction commands. + +The following chronology preserves rejected experiments and negative results. It is +not a speedup claim, a source-CUDA parity claim, or a passing residency receipt. +Large local traces, checkpoint weights, and private filesystem paths are not included. + +## Retained investigation + +The eight-token teacher comparison initially measured 7.6812% normalized RMS error. Intermediate reference +corrections measured 2.8089% and 2.7056%, but those are not the current qualification result. Preserving the +source compressor's FP32 softmax-before-weighted-sum boundary measured 5.4496%; aligning native attention's +tile-denominator and sink arithmetic with the pinned source measured 6.1794%. Correcting native and reference +SiLU division ordering then measured 3.4807%. Revisiting compensated native EXL3 accumulation after +same-input attention isolation eliminated all first-layer attention-branch differences and measured 3.2922% +whole-model NRMS. A scalar compensated-dot candidate then matched all ten projection checks' FP32 outputs +from the independent FP64-dot oracle, but measured 3.6140% whole-model NRMS. Correcting the reference's +FP32 accumulator-before-mean-division boundary measured 3.8218%. Compensating native FP32-weight dense +products and reductions then removed the retained layer-4 hyperconnection projection error and its +downstream layer-5 normalized-input BF16 mismatch, measuring 3.5010% whole-model NRMS. Running the +pinned TileLang 0.1.8 HC kernel on CUDA and inspecting its generated PTX then established fused affine +multiply-adds and FP32 butterfly denominator reductions. Applying those contracts to native HC and +the independent reference measured 3.7623% eight-token whole-model NRMS. Computing the oracle's unary +inverse square root independently in FP64 and rounding its result once to FP32 then measured 4.1783%, +with maximum absolute logit error 1.583322 and zero greedy disagreements. Correcting the reference's +HC collapse/post FP32 accumulation boundary and unary power rounding then measured 4.4685% NRMS, +maximum absolute error 2.321231, and zero greedy disagreements across all eight tokens. Disabling +implicit FMA contraction in native HC collapse/post then removed the isolated layer-11 BF16 expansion +error, but measured 4.7181% whole-model NRMS, maximum absolute error 2.052191, and zero greedy +disagreements. Compensating native sparse QK accumulation and lane reductions then measured +2.9347% NRMS, maximum absolute error 1.153051, and zero greedy disagreements across all eight tokens. +Correcting native six-score routing normalization to its independently established source-CUDA +FP32 reduction then measured 1.2571% NRMS, maximum absolute error 0.488874, and zero greedy +disagreements. Precomputing resident source-style complex RoPE factors then fixes the retained +token-5/layer-0 attention case, but raises whole-model NRMS to 3.9876%, maximum absolute error +1.563581, with zero greedy disagreements. Correcting native HC's statistic to the independently +executed canonical single-row CUDA reduction then measures 3.7143% NRMS, maximum absolute +error 1.373373, and zero greedy disagreements. Correcting the CUDA mean scalar projection +in ordinary RMS, Engram and the independent reference then measures 4.0961% NRMS, maximum +absolute error 1.512364, and zero greedy disagreements. This is the current candidate, not +the lowest historical metric. The unchanged admission gate is 1%; `TeacherReferenceMismatch` +remains an upstream blocker. The ideal unary oracle is not CUDA hardware emulation. +The latest GPU teacher replay took 5.351 seconds; this is not served-decode throughput +or a comparative speedup. +The earlier compensated-accumulation experiment had been reverted after a worse controlled whole-model +result with the then-current reference. A proposed single-round inverse-square-root reference change was +also reverted: PyTorch 2.14.1 CPU and MPS matched FP32 square-root followed by reciprocal on all 10,001 +positive probe inputs, not the proposed FP64 inverse-square-root rounded once. That probe does not establish +CUDA intrinsic rounding. These numerical runs are not new API speed measurements. + +The mean-boundary probe matches PyTorch 2.14.1 CPU and MPS on the retained layer-4 statistic. +The reference still uses independent wide sums internally, not either backend's FP32 reduction tree. +Actual PyTorch 2.13.0a0 / CUDA 13.3 probing on an NVIDIA GB10 shows that CUDA `rsqrt` differs from +CPU/MPS square-root-then-reciprocal rounding on 3,372 of 10,003 positive inputs, and from a correctly +rounded FP64 inverse on 2,255. Both retained inverse examples agree with the original Metal intrinsic. +Compensating native HC variance measured 3.9120% whole-model NRMS without removing the retained +inverse discrepancy; a CPU/MPS-only native inverse rewrite measured 3.4688%. Both experiments were +reverted: the first did not address the causal discrepancy, and the second is not justified by the +actual CUDA operation, irrespective of its lower metric. CUDA mean also uses a rounded reciprocal +factor rather than CPU/MPS mean division. A complete source-backend reference remains unqualified. +The FP64 projection oracle and admission threshold have not been relaxed. Dense cancellation +regressions fail before compensation and pass after it. +The earlier inverse-reference experiment was based only on CPU/MPS behavior. Actual CUDA evidence now +supports retaining a distinct ideal unary FP32 result instead of adding an intermediate FP32 square-root +rounding to the mathematical oracle. On the retained layer-5 statistic that removes the one-ULP inverse +discrepancy; pre/post/comb coefficient errors remain below 2.8e-9. The 10,003-input CUDA probe above +still rules out claiming bitwise source-CUDA parity for this ideal reference. + +The TileLang probe uses independently computed FP64 projection outputs rounded to FP32 as inputs; +it is an isolated source-kernel check, not a whole-model CUDA comparison. Generated PTX confirms +`fma.rn.f32` affine operations and descending XOR row/column reductions. The corrected retained +layer-4 HC comparison has maximum coefficient errors of 4.66e-10 (pre), 3.73e-9 (post), and +1.86e-9 (comb), with identical next collapsed and normalized BF16 rows. Those checks do not clear the +whole-model gate; remaining CPU/reference versus CUDA primitive differences still require resolution. + +Executing the pinned `Block.hc_pre` and `Block.hc_post` methods on CUDA also establishes FP32 +accumulation of their materialized FP32 products. Two half-ULP terms are lost when accumulated +after 1.0; summing them in FP64 instead changes a cancelling BF16 result from zero to 2^-23. +The reference now retains the source FP32 accumulation boundary. Its test target previously +omitted the newer HC, RoPE and inverse-square-root regressions; enabling them reproduced the +residual failure and a two-ULP frequency error. The reference power operation now computes an +independent ideal unary result rounded to FP32 before reciprocal. All 64 retained CPU +frequency-vector values pass the unchanged test tolerance; this is not source-CUDA bitwise qualification. + +The native collapse/post product boundaries also require disabling implicit FMA contraction. +With identical retained layer-11 input, branch and native coefficients, the old expansion rounded +one BF16 cell to 26.625 instead of 26.5. Reversing the cancellation fixture reproduced the contraction +error: BF16 0x34c0 instead of 0x3500. Actual pinned CUDA collapse, residual and branch-product +cancellations all produce 2^-21. The corrected native retained expansion matches all 20,480 BF16 +cells; all 56 Metal checks and 10 independent reference checks pass. Explicit HC affine FMAs remain. + +The retained layer-21 compressed-attention isolation replays the normalized prefix to populate +the real compressed cache. Queries, window keys and compressed keys match, but uncompensated +native QK accumulation changed one sparse-output BF16 cell from 0.38671875 to 0.388671875. +The pinned CUDA sparse kernel's 16-head specialization agrees with the independent reference. +A separate cancellation fixture loses a 2^-6 term beside cancelling 2^24 terms in native lane +accumulation, changing sink-normalized output from source BF16 0.50390625 to 0.5. Compensated +per-lane accumulation and lane reduction reproduce the source result; the regression fails before +and passes after the correction. All retained layer-21 attention-stage BF16 outputs now match, +and all 57 Metal plus 10 reference checks pass. + +These CUDA probes use TileLang 0.1.8 with its required Apache TVM FFI 0.1.6. Re-running the +HC source check with that supported dependency preserves every previously observed coefficient +bit. The full 64-head sparse kernel compiles but cannot launch on the GB10: it requests 141,312 +dynamic shared-memory bytes, exceeding the device's 101,376-byte opt-in limit. The executed +16-head specialization is isolated operator evidence, not full-head or whole-model qualification. + +The source-CUDA routing probe establishes the contiguous six-score sum as +`((s0+s4)+s2)+((s1+s5)+s3)`: four lane accumulators followed by descending shuffle offsets 2 and 1. +Actual CUDA execution matches this reduction on 8,918 positive, signed and permuted boundary +rows; sequential and wide reductions disagree on 2,815 and 2,468 rows respectively. Two +consumer regression fixtures reconstruct exact scores through above-threshold softplus and +square root, avoiding a unary-rounding confound. Native normalization now follows the source +contract; the independent mathematical oracle retains wide normalization rather than copying +the backend reduction tree. + +`tf-dsv41-run moe-sameinput MODEL TABLES TRACE_GPU_BIN OUT_DIR [LAYER]` isolates +routing, six routed experts and the shared expert on a retained normalized input; layer defaults +to 21. It records independent-chain comparisons separately from native-input/native-weight +diagnostic reuse. Raw snapshots have no token header, so this history-free MoE receipt reports +token identity as unknown and preserves the supplied snapshot path rather than claiming token zero. +The corrected retained layer-21 logits, weights, gate/up projections, +activations, down projections and all 5,120 combined-branch BF16 cells match the independent +reference. Replaying the retained old weights reproduces all 11 old branch differences. +All 58 Metal and 10 reference checks pass; those operator checks do not clear the full-model gate. + +`tf-dsv41-run attention-history-sameinput MODEL TABLES TRACE_DIR OUT_DIR TOKEN LAYER` +replays retained normalized attention inputs at every position through the selected layer, +advancing real native and independent reference caches together. It exports queries, physical +window rows, compressed keys, selected indices and sinks for separate source-kernel replay; +native-query/cache reuse is explicitly diagnostic and never enters qualification. +At token 2, layer 1, all attention-stage BF16 outputs and all 13 isolated FP32 projections match; +the replayed native branch matches the retained full trace, with zero disk reads and page-ins. +Thus the original layer-1 mismatch at that token depends on earlier differing history inputs, +not a demonstrated operator error on identical history. This does not clear the eight-token gate. + +The earlier ideal-unary first-token diagnostic preserved native logits, but its residency check +failed: process disk reads increased by 12,288 bytes with zero page-ins. The first layer-11 isolation +also failed its residency check with 4,096 read bytes. Neither receipt establishes weight streaming. +The corrected-reference first-token trace and corrected native layer-11 isolation both measured +zero process disk-read bytes and zero page-ins; those limited runs do not qualify every workload. +The compressed-attention isolation runs measured 8,192, 4,096 and, after QK compensation, +16,384 process disk-read bytes respectively, each with zero page-ins. Their numerical evidence +does not make them passing zero-I/O qualification receipts or establish weight streaming. +The QK-compensated full eight-token layer trace preserves every native logit bit, but also fails +residency with 16,384 read bytes and zero page-ins. Its first persistent first-token divergence +was layer-21 MoE with identical normalized input; the routing correction removes that isolated +discrepancy. The corrected MoE isolation also fails residency with 4,096 read bytes and zero +page-ins. Neither numerical isolation is a passing zero-I/O receipt. +The routing-corrected eight-token trace also preserves all native logit bits, but fails residency +with 20,480 read bytes and zero page-ins. The first persistent first-token difference is now one +layer-24 HC post-expansion cell; later teacher tokens also diverge in history-dependent attention. +Retained layer-24 HC projection and collapse/post with identical native coefficients match exactly; +the inverse differs by one FP32 ULP from the ideal reference. Actual intermediate-statistic capture +then shows native and pinned Torch CUDA agree: sum 561,016,537,088, mean/argument 27,393,386, +inverse 0.00019106324180029333. The wide oracle rounds its sum to 561,016,602,624 and argument +to 27,393,388; changing that argument alone explains the inverse difference. Metal `rsqrt` agrees +with ideal unary rounding on both captured arguments, so variance compensation or an inverse +rewrite is not a source-backed correction. The isolation measures zero reads and page-ins, +but this retained-case CUDA parity does not clear the whole-model gate. + +A separate token-1/layer-0 attention probe exposes a real native HC statistic mismatch: squared +sum 28.332973480224609 versus actual pinned CUDA 28.332975387573242. CUDA matches the wide +oracle on that input. Native now follows CUDA's canonical single-row contiguous FP32 mean: +512 threads, four independent vector accumulators, descending block/warp reductions, then +FP32 mean-factor multiplication. Independent CUDA probes establish the tree on both retained +inputs and ten generated inputs; standalone native execution matches the source single-row +inverse for each input at all 14 tested row counts. The older HC24 source result is preserved. +Changing only the token-1 inverse fixes both retained expansion cells; production-tree replay +matches all 20,480 reference BF16 cells. A permanent regression recomputes RMS, coefficients +and post-expansion from the retained input rather than injecting the oracle inverse. +Native fixes the same per-row tree across chunks: CUDA can choose another reduction grouping +for batched rows, so this is decode-tree parity, not every batched-CUDA rounding behavior. +The integrated build passes all 60 Metal and 10 reference checks. The actual full 257-token +history preserves every FP32 logit bit and greedy token across chunks 1, 2, 3, 7, 31, 127, 128, +129, 255 and 256, with zero execution disk reads and page-ins. The eight-token whole-model +numerical gate nevertheless remains failed at 3.7143% NRMS. + +Pinned Torch CUDA exposes a separate mean-projection boundary: with BF16 input `[3, 0, ...]` +of width 5,120, exact squared sum 9 and FP32 gamma bits `0x3c65de29`, canonical +`RMSNorm.forward` returns BF16 1.0 on CUDA but 1.0078125 on CPU. CUDA materializes the +FP32 reciprocal count and multiplies the FP32 accumulator; division crosses this affine +rounding tie. The native ordinary RMS reproduces the CPU result before correction and +the CUDA result afterward. Native Engram uses the same corrected mean projection. +The reference now preserves that scalar boundary around its independent FP64 sums for +ordinary RMS, HC and Engram; it does not reuse native statistics or reduction trees. +All 61 Metal and 11 reference checks pass. The eight actual initial teacher-embedding +RMS cases already matched at every BF16 cell, so this constructed consumer does not +establish the whole-model root cause. The corrected candidate's actual full 257-token +history remains bit-exact across chunks 1, 2, 3, 7, 31, 127, 128, 129, 255 and 256, with +zero execution disk reads and page-ins. The full eight-token gate nevertheless fails +at 4.0961% NRMS. Its separate trace preserves every native logit bit but records 20,480 +execution disk-read bytes and zero page-ins, so it is not a zero-I/O receipt. + +Later ordinary RMS isolation also rules out a reduction rewrite for token-1/layer-2: +production replay reproduces all 5,120 retained normalized values, and exact pinned +`RMSNorm.forward` CUDA matches native on the identical real collapsed input and BF16 +norm weights. CUDA collapse with fixed native pre-coefficients is bit-exact too. +The delayed discrepancy instead comes from layer-1 FFN HC pre-coefficients: replacing +only those with independently computed wide coefficients changes collapsed cell 1,580 +and exactly reproduces the retained reference normalization. Native HC projection and +inverse match the wide calculation; canonical HC coefficient execution must establish +whether the remaining unary difference is a native error. + +The earliest token-0/layer-16 MoE branch discrepancy is source-conformant, not a justified +routing rewrite. Gate/up projections are bit-exact; wide normalization changes two BF16 +expert activations and then one combined-branch cell. Actual pinned Torch CPU and CUDA +Gate/Expert tails on the retained projections match all 16,128 native activation cells. +The diagnostic using native weights matches every expert stage and branch, but never +replaces qualification. This isolation records 12,288 read bytes and zero page-ins. + +Canonical TileLang execution separates the token-1/layer-0 FFN's three differing +expansion cells. Native's inverse matches actual pinned CUDA, not the wide statistic; +CUDA therefore agrees with native at cells 9,998 and 16,906. The third cell, 18,987, +remains a genuine unary-consumer mismatch: controlling exponent input 1.4393715858459473 +produces Metal FP32 bits `0x4086fa39`, versus source CUDA and ideal-rounded `0x4086fa38`. +That difference changes post-coefficient 3 and BF16 expansion 14,832 versus source 14,831. +Substituting only the actual source post-coefficient fixes that cell; substituting source +comb coefficients does not. No source coefficient is injected into qualification. +The delayed layer-2 normalization discrepancy also uses a pre-coefficient where actual +TileLang agrees with native, not the wide reference. + +Six actual Metal exponent variants were checked on the 16 retained HC arguments and +65,537 FP32 grid points from -32 through 32, against pinned Torch CUDA and ideal rounding. +Explicit `precise::exp` does not change production output; single- and split-log2 `exp2` +substitutions do not fix the controlling cell and increase source disagreements, reaching +29 and 22 ULP respectively versus production's 2-ULP maximum on this campaign. +All substitutes are rejected. A portable correction to this HC unary boundary and the +whole-model discrepancy remain unresolved; the unchanged numerical gate still blocks +ready-for-review or merge of the family implementation. + +The retained token-5/layer-0 query differs in one BF16 cell before this RoPE correction: native +0.54296875 versus source/reference 0.546875. Exact pinned Torch CPU and CUDA operations agree +with every independent-reference query cell across all 64 heads. Intermediate capture isolates +a one-ULP Metal cosine difference; frequency, sine and FMA counterfactuals do not fix that cell. +Native normal and compressed spectra now use resident FP32 complex-factor tables, computed once +from FP32 source-style angles, rather than per-head GPU transcendental calls. The tables add +512 MiB at maximum context, included in admission and residency, not reported as weight bytes. +There is no per-request factor allocation, transfer or CPU trigonometry. Retained-history replay +now matches every attention-stage BF16 cell and all 13 isolated FP32 projections, with zero +execution disk reads and page-ins. All 59 Metal and 10 reference checks pass. This retained-case +proof does not establish CUDA factor parity at every long-context position or clear the failing +eight-token whole-model gate; its oracle logits remain bit-identical to the previous candidate. + +The resident-factor full eight-token trace preserves every native logit bit, but fails zero-I/O +qualification with 61,440 process disk-read bytes and zero page-ins. The corrected token-5 +layer-0 attention now matches. Token-6 layer-0 MoE replay isolates two one-ULP routing-weight +differences with identical normalized inputs, logits and expert IDs. Expert-146 activation cell +477 is native 0.0810546875 versus wide-oracle 0.08056640625; exact pinned Torch CPU and CUDA +activation tails agree with native at every tested BF16 cell. That single cell propagates to +1,614 expert down-projection cells and 246 combined-branch cells. The diagnostic using native +weights matches every expert and branch stage, but never replaces the qualification oracle. +This replay measures zero disk reads and page-ins; no source-backed native correction is justified. +A separate complex-product cancellation probe agrees between +native Metal and exact pinned Torch CPU/CUDA: the real result is 2^-24, not zero. Disabling +complex-product contraction would therefore be an unjustified change. + +A broader exact-source factor probe covers positions 0–7, all 32 pairs, both configured spectra. +Reference FP32 unaries and native host-FP64 unaries rounded to FP32 agree across all 1,024 +components, so changing the reference unary width is not justified. Canonical CUDA differs +in 57 factor components; captured power, YaRN ramp and same-angle polar operations expose +device rounding differences. Both factor sets still give identical BF16 outputs on the retained +token-5 query. The table correction is therefore a proven retained-consumer fix, not a claim +of universal source-CUDA factor identity. + +The minimal native API campaign used one uncached greedy request per case, thinking off, 16 output tokens, +on an M3 Ultra with 512 GiB RAM. Every stream completed with the exact requested prompt/output token counts +and zero cached tokens. Decode excludes the first token: 15 divided by the first-to-last-token interval. +These are single-run smoke measurements, not statistical estimates or a before/after speedup claim. + +| Workload | Prompt tokens | TTFT (s) | Decode (tok/s) | +| --- | ---: | ---: | ---: | +| Prose | 128 | 10.57 | 1.46 | +| Prose | 512 | 10.71 | 1.25 | +| Prose | 1,024 | 23.23 | 1.17 | +| Prose | 2,048 | 50.00 | 1.15 | +| Prose | 4,096 | 103.93 | 1.12 | +| Code | 4,096 | 105.68 | 1.11 | + +A separate long-prefill run passed 128k tokens with zero observed process disk-read and page-in deltas. +It was intentionally cancelled before generation; million-token execution remains unqualified. diff --git a/docs/deepseek-v41.md b/docs/deepseek-v41.md new file mode 100644 index 0000000000..b48fbab3e0 --- /dev/null +++ b/docs/deepseek-v41.md @@ -0,0 +1,253 @@ +# DeepSeek V4.1 Flash: experimental resident Metal draft + +**Not release-qualified. Keep the proposal in draft until the numerical and runtime gates pass.** +[Proposal #685](https://github.com/ashhart/TensorFold/issues/685) covers this implementation. +The backend executes text inference; it is not a placeholder. Its independent whole-model +comparison still fails. Greedy-token agreement does not establish numerical correctness. + +## Qualification boundary + +The following receipts are historical, from main +`bec00ae6961a43a9ab73c4fa9338c167b3f82851` plus the local DS4.1 candidate changes. +This draft integrates that work onto main `e7c523ebfa16b8094cf9e35e3773a7d25befcc5b`. +They are not new verification of the integrated draft. Current-branch checks are reported in the PR. + +| Retained measurement | Result | Qualification | +| --- | --- | --- | +| Eight-token whole-model NRMS | 0.040960584768951355 (4.0960584769%) | FAIL; unchanged limit 0.01 (1%) | +| Maximum absolute logit error | 1.5123639106750488 | Not a substitute for NRMS | +| Greedy disagreements | 0 of 8 | Does not clear numerical gate | +| Teacher replay / independent CPU oracle | 5.350994958 s / 346.856063250 s | Not served-decode throughput or speedup | +| Operator checks | 61 Metal + 11 independent reference | Limited operator evidence | +| 257-token chunk/history replay | Identical FP32 logits and greedy tokens | Chunks 1, 2, 3, 7, 31, 127, 128, 129, 255, 256 | +| Same history replay residency | 0 read bytes; 0 page-ins | Passing measured execution scope only | +| Separate full trace residency | 20,480 read bytes; 0 page-ins | FAIL; not a zero-I/O receipt | + +`TeacherReferenceMismatch` remains a real failure, not a warning to suppress. +The [historical evidence record](deepseek-v41-evidence.md) preserves intermediate results, +rejected experiments, CUDA counterexamples and earlier speed receipts. + +### Unresolved numerical causes + +- A retained HC exponent input `1.4393715858459473` gives Metal FP32 `0x4086fa39` + versus source CUDA and ideal-rounded `0x4086fa38`. Post-coefficient 3 changes a BF16 + expansion cell from source 14831 to native 14832. Replacing only the actual source + coefficient fixes that retained cell; qualification never injects source coefficients. +- Six Metal exponent variants were exercised on 16 retained arguments and 65,537 grid + points from -32 through 32. None fixed the controlling cell. Single/split-log2 `exp2` + increased worst differences to 29/22 ULP versus production's 2 ULP; all were rejected. +- Independent wide reductions and ideal FP32 unary rounding are not a CUDA hardware + emulator. Actual source-CUDA execution also agrees with native on some cells where + the independent oracle differs. Lower NRMS alone does not justify changing arithmetic. +- The canonical 64-head CUDA sparse kernel needs 141,312 bytes of dynamic shared memory; + the available GB10 permits 101,376. Compilation succeeded but launch did not. + Executed 16-head checks are operator evidence, not full-head or whole-model qualification. + +## Implementation and limits + +`zig/src/families/deepseek_v41/` owns configuration, resident weights, workspace, +hyperconnections, Engram, sparse/compressed attention, routed/shared experts, RoPE and +history state. `zig/kernels/metal/deepseek_v41/` implements device operators. +The native host uses the existing serialized worker and HTTP transports. Independent +qualification lives in `zig/tests/deepseek_v41_reference/`; it does not import native +statistics or replace its wide reduction with the backend's reduction tree. +Cancellation is observed between bounded prompt/decode chunks; it does not interrupt +an in-flight Metal command buffer. The serialized host owns model/history state until +that command completes before reset, reuse or teardown. + +Shared EXL3 APIs check descriptor arithmetic, aligned view extents and writable-view +nonaliasing. Callers must supply actual input/output capacities for GPU ID maps, +guarantee IDs are in range and scatter output IDs are unique, register indirect weights +resident, and keep indirect descriptor-referenced weights disjoint from writable buffers. +Raw family encoders are trusted fixed-shape internals, not unchecked external APIs. +Generic counters classify `gpu_profile.Class`; the family profiler owns DS stage metadata +and totals without allocation or copied spans. Counter cleanup is fallible while encoding +or execution is in flight. There are no per-span CPU waits or fabricated timing fallbacks. +Shared EXL3 projections and hardware counters are separate from family arithmetic. + +The complete EXL3 checkpoint and original Engram tables are owned resident Metal buffers. +There is no weight mmap or per-request SSD table streaming. Startup reads are expected; +steady-state resident claims require measured process read/page-in deltas. Admission +leaves an 8 GiB reserve below Metal's recommended working set and rejects incomplete fit. + +Only one active sequence is supported; concurrent clients queue. Vision execution, +draft heads, thinking budgets, Sliding/Living Weights, learned states and retained prompt +caches are unsupported. Request admission must reject unsupported combinations. +`--context` includes prompt plus reply, from 1 through 1,048,576; prompt chunks are at most +256 rows. The packed cache uses 890 bytes per token plus bounded window/compression state. +Allocated or advertised context is not evidence that every position was exercised. +A historical prefill passed 128k tokens, then was intentionally cancelled before generation; +full million-token execution remains unqualified. + +On the measured M3 Ultra with 512 GiB RAM, the weight payload was 413,338,642,120 bytes +(384.95 GiB). Full-window cache, scratch and resident RoPE factors brought admission to +415,199,770,312 bytes (386.69 GiB). RoPE factors account for 512 MiB at maximum context; +they are not weight bytes and incur no per-request CPU trigonometry or factor transfer. +Use `/memory`, process counters and `vmmap` together; aliased virtual mappings are not +additive physical memory, and a read delta alone does not prove weight streaming. + +## Pinned inputs and build + +Keep snapshots outside the checkout. This PR contains no model shards or binary traces. + +- EXL3: [`Mia-AiLab/DeepSeek-V4.1-Flash-EXL3-2.9bpw`](https://huggingface.co/Mia-AiLab/DeepSeek-V4.1-Flash-EXL3-2.9bpw/tree/64ba41b6c916a587db06eae2e19b7845f7be6e6b), + revision `64ba41b6c916a587db06eae2e19b7845f7be6e6b`, complete snapshot. +- Source: [`deepseek-ai/DeepSeek-V4.1-Flash`](https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash/tree/2cba9e42aa026125f3ed06c6d98c1db82f7ca027), + revision `2cba9e42aa026125f3ed06c6d98c1db82f7ca027`. Supply + `model-00047-of-00048.safetensors`, `model-00048-of-00048.safetensors`, `config.json`, + `model.safetensors.index.json`, `tokenizer.json`, `tokenizer_config.json`, + `chat_template.jinja`, `LICENSE`, and `inference/config.json` for the offline map builder. + Only `.engram.embed.weight` and `.engram.embed.scale` come from these shards; + projections come from EXL3. + +Use Zig 0.17.0 and Xcode's macOS SDK/Metal toolchain: + +```sh +DEVELOPER_DIR=/Applications/Xcode.app/Contents/Developer \ + zig build native tf-dsv41-run -Dcpu=apple_m1 -Doptimize=fast +python3 -m venv .venv-weights +.venv-weights/bin/pip install numpy==2.5.3 tokenizers==0.23.3 +.venv-weights/bin/python tools/zig/dsv41_token_map.py \ + --tokenizer /models/dsv41-exl3/tokenizer.json \ + --config /models/dsv41-original/inference/config.json \ + --output /models/dsv41-exl3/engram-token-map.bin +``` + +The offline map builder preserves official normalization, PCG64 multipliers and prime +bucket layout. Native inference requires no Python. Use the pinned original template +explicitly; do not silently substitute the EXL3 snapshot's older template. +The map builder requires the original inference config's field names, not the converted +model config schema. Both pinned configs use 16,000,000 Engram buckets and 129,280 token IDs. + +## Reproduce diagnostics + +Run one full-model process at a time. `MODEL` means `/models/dsv41-exl3` and `TABLES` +means `/models/dsv41-original`; these paths are examples, not required directory names. +Use fresh output directories and preserve exit status plus JSON receipts. + +Exact retained teacher input, saved as `teacher.json`: + +```json +[0, 671, 6102, 294, 8760, 344, 11111, 16] +``` + +The retained 257-token seam input is exactly `(teacher_ids * 33)[:257]`: + +```sh +python3 -c 'import json; ids=json.load(open("teacher.json")); json.dump((ids*33)[:257],open("seams-257.json","w"))' +zig-out/bin/tf-dsv41-run plan /models/dsv41-exl3 /models/dsv41-original +zig-out/bin/tf-dsv41-run teacher /models/dsv41-exl3 /models/dsv41-original teacher.json /tmp/dsv41-teacher 1 +zig-out/bin/tf-dsv41-run oracle /models/dsv41-exl3 /models/dsv41-original teacher.json /tmp/dsv41-oracle +zig-out/bin/tf-dsv41-run seams /models/dsv41-exl3 /models/dsv41-original seams-257.json /tmp/dsv41-seams +zig build test-deepseek-v41 test-deepseek-v41-reference -Dcpu=apple_m1 -Doptimize=fast +``` + +A failed oracle exits with `TeacherReferenceMismatch`; do not rerun merely to relabel it. +The oracle's independent-chain output is qualification. Native-input or native-coefficient +reuse is an explicitly diagnostic counterfactual and cannot replace that output. +Teacher/seams accept at most 4,096 IDs; oracle/trace at most 16; profile at most 256. +IDs are nonempty JSON arrays of integer vocabulary IDs below 129,280. + +The runner also provides `generate`, `profile`, `trace`, `operators`, `attention-sameinput`, +`attention-history-sameinput`, `hc-sameinput` and `moe-sameinput`; invoke it without +arguments for syntax. History replay needs `trace-{position}-gpu.bin` for every position +through the requested token, not an invented cache or a single late snapshot. +`moe-sameinput` defaults to layer 21 and cannot infer token identity from a headerless snapshot. +Counter capture may rotate encoders and change submission overhead; profile timing is +not an uninstrumented throughput measurement. Receipts distinguish `model_execution` +from `execution_and_diagnostic_io`; never compare scope differences as weight-streaming proof. + +## Exercise actual serving + +Start with a bounded context; the complete resident weight payload still requires the +large-memory machine. The example is intentionally loopback-only and thinking off: + +```sh +zig-out/native/bin/tensorfold-native serve /models/dsv41-exl3 \ + --engram-table-dir /models/dsv41-original \ + --chat-template /models/dsv41-original/chat_template.jinja \ + --context 4096 --parallel 1 --no-drafts --prompt-cache-gib 0 \ + --host 127.0.0.1 --port 8080 --name DeepSeek-V4.1-Flash \ + --keep-warm 0 --no-update-check +curl -fsS http://127.0.0.1:8080/v1/chat/completions \ + -H 'Content-Type: application/json' \ + -d '{"model":"DeepSeek-V4.1-Flash","messages":[{"role":"user","content":"Respond with only the word READY."}],"temperature":0,"max_tokens":24,"chat_template_kwargs":{"enable_thinking":false}}' +``` + +For throughput receipts, preserve public prompt/token IDs, uncached counts, sampling +settings, output count, first/last token times, exact executable commit and hardware. +Decode rate excludes the first token. Do not turn single-run smoke receipts into a +statistical comparison or a whole-model numerical qualification claim. + +## Integrated draft smoke receipt + +These checks exercised the integrated implementation, not the older numerical campaign. +Environment: Apple M3 Ultra, 512 GiB RAM, macOS 26.6.2, Zig 0.17.0 and Xcode Metal; +the PR records the exact published commit. Checkpoints/revisions are the pinned inputs above. +The known failed whole-model comparison was not rerun merely to reconfirm failure. + +```sh +DEVELOPER_DIR=/Applications/Xcode.app/Contents/Developer \ +TENSORFOLD_V41_TOKENIZER_DIR=/models/dsv41-original \ + zig build native tf-dsv41-run tf-exl3-bench test-deepseek-v41 \ + test-deepseek-v41-reference test-server test-golden test \ + -Dcpu=apple_m1 -Doptimize=fast -j2 --summary all +zig-out/bin/tf-exl3-bench --shape 4 --warmup 1 --repeats 1 --prefill 2 +``` + +- Final build: 158/158 steps succeeded; 178 newly executed tests passed, with remaining + successful test steps reused from cache. Dedicated DS checks passed 75 device/host + cases and 13 independent reference cases during integration; these are not 88 + whole-model comparisons. The server consumer suite passed 155/155. +- Frozen HTTP goldens: 309/319 exact matches, 10 explicitly known differences, + zero unexpected differences. No golden answers were rewritten. +- Actual EXL3 CLI: all 196,608 codebook states matched FP16 bits; the selected + `K=4096,N=1024` one-row, routed-six and two-row prefill paths had zero comparator + bit mismatches and passed the sampled independent FP64 oracle gate. One repeat + is smoke evidence, not a statistically supported speed claim. +- Actual offline map packing produced 517,680 bytes, raw/compressed vocabularies + 129,280/99,092 and compressed pad 2. Metadata/lookup matched the retained map; + the only 32 differing bytes were the correctly refreshed input-config SHA-256. +- Actual `plan` admitted 41 shards, 192,456 tensors and 413,338,642,120 weight bytes. + Actual serving at context 4,096 admitted 413,735,409,352 device bytes and printed + `EXPERIMENTAL / UNQUALIFIED`. Startup took 113.9 seconds; this is not decode latency. +- Greedy, thinking-off whole and SSE requests for “Respond with only the word READY.” + returned `READY`, with 12 prompt tokens, 3 completion tokens and zero cached tokens. + An intervening “What is 2 + 3? Respond with only the answer.” returned `5`; + repeating the first request preserved content/counts and its logged token-hash prefix. +- Whole and SSE requests for “Call lookup with q set to Paris. Do not answer directly.” + with one `lookup(q:string)` tool and `parallel_tool_calls:false` returned exactly one + call, `{"q":"Paris"}`, index 0 in SSE, and `finish_reason:"tool_calls"`. + Both used 275 prompt tokens and 38 completion tokens, zero cached tokens; + this actual prompt crossed the 256-row prefill boundary. +- Actual requests rejected positive thinking budgets while thinking was off, explicit + drafts, image input, and a 4,096-token prompt leaving no reply room with HTTP 400. + Closing a counting stream after its first content delta logged `reason=client-left`; + the next request returned `READY` normally. SIGTERM shutdown exited 0. +- Process `proc_pid_rusage` deltas were zero read bytes and zero page-ins across the + three text/reset requests, across disconnect/recovery, and across the two tool requests. + These bounded request intervals do not reclassify the historical failed full trace, + qualify long-context execution, or clear the whole-model numerical gate. +- Source-rule scan: 1,241 files, 532 inherited findings versus main's recorded + 1,169 files/534 findings; zero added findings. The checker remains nonzero, + not “lint clean.” Outgoing text contained no private home paths or credential markers; + weights, binary traces, build outputs and local receipts were excluded. + +## Publication and ready gate + +Shared contributions overlap #648, #649, #650, #680, #681 and #684. This draft is based +on main, not on unmerged PR heads; overlap must be reconciled before merge. +The qualified-model table remains unchanged. The PR stays draft while any gate below fails: + +- Unchanged independent whole-model NRMS gate passes on the retained teacher IDs, + without coefficient injection, oracle/backend conflation or omitted history. +- Remaining source-CUDA unary discrepancy is resolved with consumer-visible evidence; + unsupported full-head/source-backend checks are explicitly completed or scoped by reviewers. +- Affected operator, reference, server and golden checks pass on the integrated commit. +- Actual serving, reset/cancellation, admission and relevant chunk/history paths are exercised. +- Zero-I/O receipts identify their scope; long-context claims match exercised positions. +- Outgoing files/history contain no credentials, private hostnames, local paths, weights, + binary traces or generated build artifacts. Public provenance and license notices remain. + +No automated action in this work should mark the PR ready, merge it, or advertise the +family as release-qualified. Further DS4.1 numerical investigation continues locally. diff --git a/tools/zig/dsv41_token_map.py b/tools/zig/dsv41_token_map.py new file mode 100755 index 0000000000..31006cdf70 --- /dev/null +++ b/tools/zig/dsv41_token_map.py @@ -0,0 +1,134 @@ +#!/usr/bin/env python3 +"""Offline Engram metadata packing; official algorithm and MIT attribution in THIRD_PARTY_NOTICES.md.""" +from __future__ import annotations + +import argparse +import hashlib +import json +import math +from pathlib import Path +import struct + +import numpy as np +from tokenizers import Regex, Tokenizer, normalizers + +SOURCE = "https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash/blob/2cba9e42aa026125f3ed06c6d98c1db82f7ca027/inference/engram.py" +MAGIC = b"DSENGR01" +RAW_VOCAB = 129280 +COMPRESSED_VOCAB = 99092 + + +def compressed_map(backend: Tokenizer) -> tuple[list[int], int]: + sentinel = "\ue000" + normalize = normalizers.Sequence([ + normalizers.NFKC(), normalizers.NFD(), normalizers.StripAccents(), + normalizers.Lowercase(), normalizers.Replace(Regex(r"[ \t\r\n]+"), " "), + normalizers.Replace(Regex(r"^ $"), sentinel), normalizers.Strip(), + normalizers.Replace(sentinel, " "), + ]) + keys: dict[str, int] = {} + lookup: list[int] = [] + if backend.get_vocab_size(with_added_tokens=True) != RAW_VOCAB: + raise ValueError("This native format supports the released vocabulary only") + for token_id in range(backend.get_vocab_size(with_added_tokens=True)): + text = backend.decode([token_id], skip_special_tokens=False) + key = backend.id_to_token(token_id) if "\ufffd" in text else normalize.normalize_str(text) or text + if key is None: + raise ValueError(f"Missing tokenizer entry {token_id}") + if key not in keys: + keys[key] = len(keys) + lookup.append(keys[key]) + return lookup, len(keys) + + +def prime(n: int) -> bool: + if n < 2 or n % 2 == 0: + return n == 2 + return all(n % divisor for divisor in range(3, math.isqrt(n) + 1, 2)) + + +def layout(config: dict, compressed: int) -> list[dict]: + seen: set[int] = set() + layers = [] + order, heads = config["engram_max_ngram_size"], config["engram_n_heads"] + bound = max(1, ((2**63 - 1) // compressed) // 2) + for layer_id, rows in zip(config["engram_layer_ids"], config["engram_num_embeddings"], strict=True): + primes, offsets = [], [] + for _ in range(order - 1): + current = config["engram_vocab_size"] - 1 + for _ in range(heads): + current += 1 + while current in seen or not prime(current): + current += 1 + seen.add(current) + offsets.append(sum(primes)) + primes.append(current) + if sum(primes) != rows: + raise ValueError(f"Layer {layer_id}: prime buckets sum {sum(primes)} != table rows {rows}") + multipliers = np.random.default_rng(10007 * layer_id).integers(0, bound, size=order, dtype=np.int64) * 2 + 1 + layers.append(dict(layer_id=layer_id, rows=rows, primes=primes, offsets=offsets, + multipliers=[int(x) for x in multipliers])) + return layers + + +def released_config(config: dict) -> None: + expected = dict(vocab_size=RAW_VOCAB, engram_compressed_vocab_size=COMPRESSED_VOCAB, + engram_max_ngram_size=4, engram_n_heads=8, engram_vocab_size=16000000) + for name, value in expected.items(): + if type(config.get(name)) is not int or config[name] != value: + raise ValueError(f"Unsupported released configuration: {name}") + if config.get("engram_layer_ids") != [1, 14]: + raise ValueError("Unsupported Engram layers") + rows = config.get("engram_num_embeddings") + if not isinstance(rows, list) or len(rows) != 2 or any( + type(row) is not int or not 0 < row <= 24 * 16010000 for row in rows + ): + raise ValueError("Invalid Engram table dimensions") + + +def bounded_bytes(path: Path, limit: int) -> bytes: + with path.open("rb") as source: + data = source.read(limit + 1) + if len(data) > limit: + raise ValueError(f"Metadata exceeds {limit} bytes") + return data + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--tokenizer", type=Path, required=True) + parser.add_argument("--config", type=Path, required=True, help="Pinned official inference/config.json") + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + config_bytes = bounded_bytes(args.config, 1 << 20) + outer = json.loads(config_bytes) + config = outer.get("text_config", outer) + released_config(config) + tokenizer_bytes = bounded_bytes(args.tokenizer, 64 << 20) + lookup, compressed = compressed_map(Tokenizer.from_str(tokenizer_bytes.decode("utf-8"))) + if compressed != config["engram_compressed_vocab_size"] or len(lookup) != config["vocab_size"]: + raise ValueError(f"Tokenizer mismatch: raw={len(lookup)}, compressed={compressed}") + order, heads = config["engram_max_ngram_size"], config["engram_n_heads"] + pad_id = config.get("engram_pad_id", config.get("engram_pad_token_id")) + if type(pad_id) is not int or not 0 <= pad_id < len(lookup): + raise ValueError("Invalid Engram pad token") + pad = lookup[pad_id] + layers = layout(config, compressed) + tokenizer_sha = hashlib.sha256(tokenizer_bytes).digest() + config_sha = hashlib.sha256(config_bytes).digest() + # LE header: magic, raw/compressed vocab, compressed pad, order, heads, layers, provenance hashes. + packed = bytearray(MAGIC + struct.pack("<6I", len(lookup), compressed, pad, order, heads, len(layers))) + packed.extend(tokenizer_sha + config_sha) + for layer in layers: + packed.extend(struct.pack("<2I4Q24I24I", layer["layer_id"], layer["rows"], + *layer["multipliers"], *layer["primes"], *layer["offsets"])) + packed.extend(struct.pack(f"<{len(lookup)}I", *lookup)) + args.output.write_bytes(packed) + manifest = dict(source=SOURCE, tokenizer_sha256=tokenizer_sha.hex(), config_sha256=config_sha.hex(), + packed_sha256=hashlib.sha256(packed).hexdigest(), raw_vocab=len(lookup), + compressed_vocab=compressed, pad=pad, order=order, heads=heads, layers=layers) + args.output.with_suffix(args.output.suffix + ".json").write_text(json.dumps(manifest, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/zig/build/deepseek_v41.zig b/zig/build/deepseek_v41.zig new file mode 100644 index 0000000000..a6676b0bfb --- /dev/null +++ b/zig/build/deepseek_v41.zig @@ -0,0 +1,39 @@ +//! Experimental DS4.1 operator and independent-reference qualification targets. +const std = @import("std"); + +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, +) void { + const device = b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/deepseek_v41_checks.zig"), + .target = target, + .optimize = optimize, + .link_libc = true, + .imports = &.{ + .{ .name = "metal", .module = metal }, + .{ .name = "kernel_sources", .module = sources }, + }, + }) }); + b.step("test-deepseek-v41", "Experimental DS4.1 EXL3, Engram, mHC and sparse attention on Metal") + .dependOn(&b.addRunArtifact(device).step); + const reference = b.addTest(.{ + .root_module = b.createModule(.{ + .root_source_file = b.path("zig/deepseek_v41_reference_checks.zig"), + .target = target, + .optimize = optimize, + .link_libc = true, + .imports = &.{ + .{ .name = "tensorfold", .module = engine }, + .{ .name = "metal", .module = metal }, + }, + }), + .filters = &.{"tests.deepseek_v41_reference."}, + }); + b.step("test-deepseek-v41-reference", "Independent CPU teacher-oracle arithmetic boundaries") + .dependOn(&b.addRunArtifact(reference).step); +} diff --git a/zig/deepseek_v41_checks.zig b/zig/deepseek_v41_checks.zig new file mode 100644 index 0000000000..755a35c618 --- /dev/null +++ b/zig/deepseek_v41_checks.zig @@ -0,0 +1,13 @@ +//! Source-built V4.1 operator qualification; real Metal execution, no checkpoint needed. +test { + _ = @import("src/core/exl3_projection_gpu_test.zig"); + _ = @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/index_precision_test.zig"); + _ = @import("src/families/deepseek_v41/config.zig"); + _ = @import("src/families/deepseek_v41/profiler.zig"); + _ = @import("src/core/gpu_counters.zig"); + _ = @import("src/core/expert_groups.zig"); +} 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..871eaf20ab --- /dev/null +++ b/zig/deepseek_v41_run.zig @@ -0,0 +1,279 @@ +//! 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 }; +fn counters() !Counters { + var r: process_memory.RUsage = undefined; + if (!process_memory.readUsage(&r)) 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/exl3_projection_bench.zig b/zig/exl3_projection_bench.zig new file mode 100644 index 0000000000..3dc2233b28 --- /dev/null +++ b/zig/exl3_projection_bench.zig @@ -0,0 +1,302 @@ +//! Source-built, bounded EXL3 qualification: real shapes, GPU timestamps and exact output parity. +const std = @import("std"); +const mtl = @import("metal"); +const ex = @import("tensorfold").exl3_projection; +const sources = @import("kernel_sources"); +const oracle = @import("exl3_projection_oracle.zig"); + +const Shape = struct { k: usize, n: usize }; +const shapes = [_]Shape{ + .{ .k = 5120, .n = 2304 }, .{ .k = 2304, .n = 5120 }, + .{ .k = 1280, .n = 32768 }, .{ .k = 5120, .n = 129280 }, + .{ .k = 4096, .n = 1024 }, .{ .k = 8192, .n = 5120 }, + .{ .k = 5120, .n = 1280 }, .{ .k = 5120, .n = 512 }, + .{ .k = 6144, .n = 25600 }, +}; +const Mode = enum { one, routed6, prefill }; +const Settings = struct { + warmup: usize = 2, + repeats: usize = 5, + prefill: usize = 32, + candidate: ex.Decode = .hoisted_decode, +}; +const limit = 2 * 1024 * 1024 * 1024; + +const Owned = struct { + matrix: ex.Matrix, + buffers: [4]mtl.Buffer, + + fn bytes(s: Shape, halves: u16) usize { + const ns = @as(usize, halves) * 16 + 2; + return 16 + (s.k / 16) * ((s.n / 16) * ns + 16) + (s.k + 2 * s.n) * 2 + 48; + } + fn init(device: mtl.Device, s: Shape, halves: u16, book: ex.format.Codebook, seed: u32) !Owned { + const ns = @as(usize, halves) * 16 + 2; + const ks = (s.n / 16) * ns + 16; + const t = try device.buffer(16 + (s.k / 16) * ks, mtl.ResourceOptions.shared); + errdefer t.deinit(); + const su = try device.buffer(s.k * 2 + 16, mtl.ResourceOptions.shared); + errdefer su.deinit(); + const sv = try device.buffer(s.n * 2 + 16, mtl.ResourceOptions.shared); + errdefer sv.deinit(); + const bias = try device.buffer(s.n * 2 + 16, mtl.ResourceOptions.shared); + errdefer bias.deinit(); + // Any circular packed bitstream is valid; independent deterministic bytes also fill padding. + var state = seed; + for (t.slice(u16, t.length() / 2)) |*v| { + state ^= state << 13; + state ^= state >> 17; + state ^= state << 5; + v.* = @truncate(state); + } + for (su.slice(f16, s.k + 8)[8..], 0..) |*v, i| v.* = if (i % 3 == 0) -0.75 else 1.125; + for (sv.slice(f16, s.n + 8)[8..], 0..) |*v, i| v.* = if (i % 5 == 0) -0.25 else 0.875; + for (bias.slice(f16, s.n + 8)[8..], 0..) |*v, i| + v.* = @floatCast(@as(f32, @floatFromInt(i % 7)) / 16); + var matrix: ex.Matrix = .{ + .trellis = .{ .buffer = t, .offset = 16 }, + .suh = .{ .buffer = su, .offset = 16 }, + .svh = .{ .buffer = sv, .offset = 16 }, + .bias = .{ .buffer = bias, .offset = 16 }, + .k = s.k, + .n = s.n, + .width = .{ .halves = halves }, + .codebook = book, + .tile_k_stride = ks, + .tile_n_stride = ns, + }; + if (seed % 2 == 0) { + matrix.suh_format = .packed_signs; + matrix.svh_format = .packed_signs; + @memset(su.slice(u16, s.k + 8)[8..][0 .. s.k / 16], 0x9635); + @memset(sv.slice(u16, s.n + 8)[8..][0 .. s.n / 16], 0x4a17); + } + return .{ .matrix = matrix, .buffers = .{ t, su, sv, bias } }; + } + fn deinit(w: Owned) void { + for (w.buffers) |b| b.deinit(); + } +}; + +fn finish(cb: mtl.CommandBuffer, e: mtl.ComputeEncoder) !f64 { + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.log.err("EXL3 GPU failure: {s}", .{msg}); + return error.Exl3GpuFailure; + } + const seconds = cb.gpuSeconds(); + if (!std.math.isFinite(seconds) or seconds <= 0) return error.GpuTimestampUnavailable; + return seconds; +} + +fn checkCodebooks(device: mtl.Device, queue: mtl.Queue) !void { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const lib = try mtl.Library.fromSource(device, sources.core_exl3_projection, mtl.CompileOptions.mlx()); + defer lib.deinit(); + const check = try mtl.Pipeline.init(device, lib, "tf_exl3_codebook_check", false); + defer check.deinit(); + const arithmetic = try device.buffer(ex.codebook_bytes, mtl.ResourceOptions.shared); + defer arithmetic.deinit(); + const cb = queue.commandBuffer(); + const e = cb.compute(.concurrent); + e.setPipeline(check); + e.setBuffer(arithmetic, 0, 0); + e.dispatchThreads(.{ .width = ex.codebook_entries }, .{ .width = 256 }); + _ = try finish(cb, e); + var wrong: usize = 0; + const books = [_]ex.format.Codebook{ .inst3, .mcg, .mul1 }; + for (arithmetic.slice(u16, ex.codebook_entries), 0..) |actual, i| { + const expected = ex.format.codebookValue(books[i / 65536], @intCast(i % 65536)); + if (actual != @as(u16, @bitCast(expected))) wrong += 1; + } + std.debug.print("EXL3 codebook_states={d} fp16_bit_mismatches={d} sample_bytes={d}\n", .{ ex.codebook_entries, wrong, arithmetic.length() }); + if (wrong != 0) return error.CodebookParityFailure; +} + +fn bf16(x: f32) u16 { + const bits: u32 = @bitCast(x); + return @truncate((bits +% (0x7fff + ((bits >> 16) & 1))) >> 16); +} + +const Work = struct { + matrix: ex.Matrix, + experts: ex.Experts, + indexed: bool, + x: ex.View, + ys: [2]ex.View, + scratch: ex.View, + rows: ex.Rows, + + fn run(w: Work, queue: mtl.Queue, p: *ex.Pipelines, decode: ex.Decode) !f64 { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + p.decode = decode; + const y = w.ys[if (decode == .arithmetic) @as(usize, 0) else 1]; + const cb = queue.commandBuffer(); + const e = cb.compute(.concurrent); + const encoded = if (w.indexed) ex.encodeIndexed(p, e, w.experts, w.x, y, w.scratch, w.rows) else ex.encode(p, e, w.matrix, w.x, y, w.scratch, w.rows); + encoded catch |err| { + e.end(); + return err; + }; + return finish(cb, e); + } + fn parity(w: Work, n: usize) !void { + var wrong: usize = 0; + var max_abs: f32 = 0; + const size: usize = if (w.rows.output == .f32) 4 else 2; + const length = (w.rows.count - 1) * w.rows.y_stride + n; + const a = w.ys[0].buffer.contents() + w.ys[0].offset; + const b = w.ys[1].buffer.contents() + w.ys[1].offset; + for (0..w.rows.count) |row| for (0..n) |col| { + const i = row * w.rows.y_stride + col; + if (size == 4) { + const av = @as([*]const f32, @ptrCast(@alignCast(a)))[0..length][i]; + const bv = @as([*]const f32, @ptrCast(@alignCast(b)))[0..length][i]; + if (!std.math.isFinite(av) or !std.math.isFinite(bv)) return error.NonFiniteOutput; + if (@as(u32, @bitCast(av)) != @as(u32, @bitCast(bv))) wrong += 1; + max_abs = @max(max_abs, @abs(av - bv)); + } else { + const av = @as([*]const u16, @ptrCast(@alignCast(a)))[0..length][i]; + const bv = @as([*]const u16, @ptrCast(@alignCast(b)))[0..length][i]; + const af: f32 = @bitCast(@as(u32, av) << 16); + const bf: f32 = @bitCast(@as(u32, bv) << 16); + if (!std.math.isFinite(af) or !std.math.isFinite(bf)) return error.NonFiniteOutput; + if (av != bv) wrong += 1; + max_abs = @max(max_abs, @abs(af - bf)); + } + }; + std.debug.print(" parity_bit_mismatches={d} max_abs={d:.9}\n", .{ wrong, max_abs }); + if (wrong != 0) return error.OutputParityFailure; + } +}; + +fn benchmark(gpa: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, p: *ex.Pipelines, s: Shape, mode: Mode, settings: Settings) !void { + const count: usize = switch (mode) { + .one => 1, + .routed6 => 6, + .prefill => settings.prefill, + }; + const books = [_]ex.format.Codebook{ .mul1, .inst3, .mul1, .mcg, .mul1, .inst3 }; + const halves = [_]u16{ 6, 4, 3, 6, 5, 6 }; + const expert_count: usize = if (mode == .routed6) 6 else 1; + const output: ex.Output = if (s.n == 129280) .f32 else .bf16; + const element_bytes: usize = if (output == .f32) 4 else 2; + const x_stride = s.k + 16; + const y_stride = s.n + 8; + const scratch_bytes = try ex.scratchBytes(s.k, s.n, count); + var total = count * x_stride * 2 + 32 + 2 * (count * y_stride * element_bytes + 32) + + scratch_bytes + expert_count * @sizeOf(ex.ExpertDescriptor) + count * 12; + for (0..expert_count) |i| total += Owned.bytes(s, halves[i]); + if (total >= limit) return error.BenchmarkAllocationLimit; + var owned: [6]Owned = undefined; + var initialized: usize = 0; + defer for (owned[0..initialized]) |w| w.deinit(); + for (0..expert_count) |i| { + owned[i] = try Owned.init(device, s, halves[i], books[i], @intCast(0x173819 + i)); + initialized += 1; + } + const x = try device.buffer(count * x_stride * 2 + 32, mtl.ResourceOptions.shared); + defer x.deinit(); + for (x.slice(u16, count * x_stride + 16), 0..) |*v, i| + v.* = bf16(@as(f32, @floatFromInt(@as(i32, @intCast(i % 23)) - 11)) / 32); + var ys: [2]mtl.Buffer = undefined; + ys[0] = try device.buffer(count * y_stride * element_bytes + 32, mtl.ResourceOptions.shared); + defer ys[0].deinit(); + ys[1] = try device.buffer(count * y_stride * element_bytes + 32, mtl.ResourceOptions.shared); + defer ys[1].deinit(); + for (ys) |y| @memset(y.slice(u8, y.length()), 0xa5); + const scratch = try device.buffer(scratch_bytes, mtl.ResourceOptions.shared); + defer scratch.deinit(); + const descriptors = try device.buffer(expert_count * @sizeOf(ex.ExpertDescriptor), mtl.ResourceOptions.shared); + defer descriptors.deinit(); + for (owned[0..expert_count], descriptors.slice(ex.ExpertDescriptor, expert_count)) |w, *d| + d.* = try ex.descriptor(w.matrix); + const routes = try device.buffer(count * 4, mtl.ResourceOptions.shared); + defer routes.deinit(); + const gather = try device.buffer(count * 4, mtl.ResourceOptions.shared); + defer gather.deinit(); + const scatter = try device.buffer(count * 4, mtl.ResourceOptions.shared); + defer scatter.deinit(); + for (0..count) |i| { + routes.slice(u32, count)[i] = @intCast(i % expert_count); + gather.slice(u32, count)[i] = 0; + scatter.slice(u32, count)[i] = @intCast(count - 1 - i); + } + const set = try device.residencySet(40); + defer set.deinit(); + for (owned[0..expert_count]) |w| for (w.buffers) |b| set.add(b); + for ([_]mtl.Buffer{ x, ys[0], ys[1], scratch, descriptors, routes, gather, scatter }) |b| set.add(b); + set.commit(); + set.requestResidency(); + queue.addResidencySet(set); + defer queue.removeResidencySet(set); + const w: Work = .{ + .matrix = owned[0].matrix, + .experts = .{ .descriptors = .{ .buffer = descriptors }, .count = expert_count, .k = s.k, .n = s.n }, + .indexed = mode == .routed6, + .x = .{ .buffer = x, .offset = 32 }, + .ys = .{ .{ .buffer = ys[0], .offset = 32 }, .{ .buffer = ys[1], .offset = 32 } }, + .scratch = .{ .buffer = scratch }, + .rows = .{ .count = count, .x_stride = x_stride, .y_stride = y_stride, .output = output, .routes = if (mode == .routed6) .{ .ids = .{ .buffer = routes }, .repeat = 6 } else null, .input_ids = if (mode == .routed6) .{ .buffer = gather } else null, .output_ids = if (mode == .routed6) .{ .buffer = scatter } else null, .input_rows = 1, .output_rows = count }, + }; + for (0..settings.warmup) |_| { + _ = try w.run(queue, p, .arithmetic); + _ = try w.run(queue, p, settings.candidate); + } + var totals = [_]f64{ 0, 0 }; + var best = [_]f64{ std.math.inf(f64), std.math.inf(f64) }; + for (0..settings.repeats) |r| { + // Alternate order to avoid consistently favoring either path's cache/thermal state. + for (0..2) |j| { + const index = (r + j) % 2; + const seconds = try w.run(queue, p, if (index == 0) .arithmetic else settings.candidate); + totals[index] += seconds; + best[index] = @min(best[index], seconds); + } + try w.parity(s.n); + } + var matrices: [6]ex.Matrix = undefined; + for (owned[0..expert_count], 0..) |item, index| matrices[index] = item.matrix; + try oracle.verify(gpa, matrices[0..expert_count], w.x, w.ys, w.rows); + const repeats: f64 = @floatFromInt(settings.repeats); + std.debug.print("EXL3 k={d} n={d} mode={s} rows={d} output={s} candidate={s} allocated_bytes={d} " ++ + "warmup={d} repeats={d} baseline_mean_gpu_s={d:.9} candidate_mean_gpu_s={d:.9} " ++ + "baseline_best_gpu_s={d:.9} candidate_best_gpu_s={d:.9} mean_speedup={d:.4}\n", .{ s.k, s.n, @tagName(mode), count, @tagName(output), @tagName(settings.candidate), total, settings.warmup, settings.repeats, totals[0] / repeats, totals[1] / repeats, best[0], best[1], totals[0] / totals[1] }); +} + +pub fn main(init: std.process.Init) !void { + const args = try init.minimal.args.toSlice(init.arena.allocator()); + var settings: Settings = .{}; + var shape_index: ?usize = null; + var i: usize = 1; + while (i < args.len) : (i += 2) { + if (i + 1 >= args.len) return error.OptionValueRequired; + if (std.mem.eql(u8, args[i], "--candidate")) { + settings.candidate = std.meta.stringToEnum(ex.Decode, args[i + 1]) orelse return error.UnknownCandidate; + if (settings.candidate == .arithmetic) return error.InvalidBenchmarkSettings; + continue; + } + const v = try std.fmt.parseInt(usize, args[i + 1], 10); + if (std.mem.eql(u8, args[i], "--repeats")) settings.repeats = v else if (std.mem.eql(u8, args[i], "--warmup")) settings.warmup = v else if (std.mem.eql(u8, args[i], "--prefill")) settings.prefill = v else if (std.mem.eql(u8, args[i], "--shape")) shape_index = v else return error.UnknownOption; + } + if (settings.repeats == 0 or settings.prefill == 0 or settings.prefill > 128 or + (shape_index != null and shape_index.? >= shapes.len)) return error.InvalidBenchmarkSettings; + 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 ex.Pipelines.load(device); + defer p.deinit(); + try checkCodebooks(device, queue); + for (shapes, 0..) |s, si| { + if (shape_index != null and shape_index.? != si) continue; + for ([_]Mode{ .one, .routed6, .prefill }) |mode| + try benchmark(init.gpa, device, queue, &p, s, mode, settings); + } +} diff --git a/zig/exl3_projection_oracle.zig b/zig/exl3_projection_oracle.zig new file mode 100644 index 0000000000..b00927ba05 --- /dev/null +++ b/zig/exl3_projection_oracle.zig @@ -0,0 +1,144 @@ +//! Bounded scalar-format decoding and FP64 dots; no Metal decoder, butterfly mirror or full weight matrix. +const std = @import("std"); +const ex = @import("tensorfold").exl3_projection; + +fn pointer(comptime T: type, v: ex.View) [*]const T { + return @ptrCast(@alignCast(v.buffer.contents() + v.offset)); +} +fn scale(v: ex.View, kind: ex.ScaleFormat, i: usize) f64 { + if (kind == .packed_signs) return if ((pointer(u16, v)[i / 16] >> @intCast(i % 16)) & 1 == 1) -1 else 1; + return @floatCast(pointer(f16, v)[i]); +} +fn output(v: ex.View, kind: ex.Output, i: usize) f64 { + return if (kind == .f32) pointer(f32, v)[i] else @as(f32, @bitCast(@as(u32, pointer(u16, v)[i]) << 16)); +} + +fn rotateInput(m: ex.Matrix, x: ex.View, row: usize, stride: usize, rotated: []f16) void { + for (0..m.k / 128) |b| { + var block: [128]f64 = undefined; + for (&block, 0..) |*v, i| { + const j = b * 128 + i; + const input: f32 = @bitCast(@as(u32, pointer(u16, x)[row * stride + j]) << 16); + v.* = @as(f64, input) * scale(m.suh, m.suh_format, j); + } + ex.format.rotate(&block, 1, 128, 1); + for (block, 0..) |v, i| rotated[b * 128 + i] = @floatCast(v); + } +} + +fn referenceBlock(m: ex.Matrix, rotated: []const f16, base: usize, out: *[128]f64) !void { + @memset(out, 0); + var placement: [256]ex.format.Placement = undefined; + ex.format.tilePlacement(&placement); + const k_stride = try m.kStride(); + for (0..m.k / 16) |tk| { + for (0..8) |tn| { + const bytes = tk * k_stride + (base / 16 + tn) * m.nStride(); + const tile: [*]const i16 = @ptrCast(@alignCast(m.trellis.buffer.contents() + m.trellis.offset + bytes)); + var states: [256]u32 = undefined; + ex.format.tileStates(m.width, tile[0..m.width.tileWords()], &states); + for (states, placement) |state, at| { + const weight: f64 = @floatCast(ex.format.codebookValue(m.codebook, @intCast(state))); + out[tn * 16 + at.col] += @as(f64, rotated[tk * 16 + at.row]) * weight; + } + } + } + ex.format.rotate(out, 1, 128, 1); + for (out, 0..) |*v, i| { + v.* *= scale(m.svh, m.svh_format, base + i); + if (m.bias) |bias| v.* += @as(f64, pointer(f16, bias)[base + i]); + } +} + +const Stats = struct { + squared_error: f64 = 0, + squared_reference: f64 = 0, + max_abs: f64 = 0, + count: usize = 0, + + fn add(s: *Stats, got: f64, want: f64) !void { + if (!std.math.isFinite(got) or !std.math.isFinite(want)) return error.NonFiniteOracleOutput; + const difference = got - want; + s.squared_error += difference * difference; + s.squared_reference += want * want; + s.max_abs = @max(s.max_abs, @abs(difference)); + s.count += 1; + } + fn nrms(s: Stats) f64 { + if (s.squared_reference == 0) return if (s.squared_error == 0) 0 else std.math.inf(f64); + return std.math.sqrt(s.squared_error / s.squared_reference); + } +}; + +pub fn verify(gpa: std.mem.Allocator, matrices: []const ex.Matrix, x: ex.View, ys: [2]ex.View, rows: ex.Rows) !void { + const n = matrices[0].n; + const rotated = try gpa.alloc(f16, matrices[0].k); + defer gpa.free(rotated); + var aggregate = [_]Stats{ .{}, .{} }; + var worst_abs: f64 = 0; + var worst_column: usize = 0; + var worst_nrms: f64 = 0; + var worst_nrms_column: usize = 0; + var baseline_worst_abs: f64 = 0; + var baseline_worst_nrms: f64 = 0; + // Compact fixtures check every row and column; large shapes sample three rotation blocks and two rows. + const blocks: usize = @min(n / 128, 3); + const checked_rows = if (n == 128) rows.count else @min(rows.count, 2); + for (0..blocks) |sample| { + const base = (if (blocks == 1) 0 else sample * (n / 128 - 1) / (blocks - 1)) * 128; + var columns: [2][128]Stats = .{ @splat(.{}), @splat(.{}) }; + for (0..checked_rows) |r| { + const expert = if (rows.routes) |routes| pointer(u32, routes.ids)[r] else 0; + if (expert >= matrices.len) return error.InvalidOracleExpert; + const m = matrices[expert]; + const source = if (rows.input_ids) |ids| pointer(u32, ids)[r] else r / (if (rows.routes) |routes| routes.repeat else @as(usize, 1)); + const dest = if (rows.output_ids) |ids| pointer(u32, ids)[r] else r; + rotateInput(m, x, source, rows.x_stride, rotated); + var expected: [128]f64 = undefined; + try referenceBlock(m, rotated, base, &expected); + for (expected, 0..) |want, c| { + for (0..2) |variant| { + const got = output(ys[variant], rows.output, dest * rows.y_stride + base + c); + try aggregate[variant].add(got, want); + try columns[variant][c].add(got, want); + } + } + } + for (0..128) |c| { + const before = columns[0][c]; + const after = columns[1][c]; + if (after.squared_error > before.squared_error or after.max_abs > before.max_abs) + return error.WorseOracleColumn; + baseline_worst_abs = @max(baseline_worst_abs, before.max_abs); + baseline_worst_nrms = @max(baseline_worst_nrms, before.nrms()); + if (after.max_abs > worst_abs) { + worst_abs = after.max_abs; + worst_column = base + c; + } + if (after.nrms() > worst_nrms) { + worst_nrms = after.nrms(); + worst_nrms_column = base + c; + } + } + } + std.debug.print("EXL3 oracle=scalar_format_fp64_dot_fp16_rotated_input k={d} n={d} " ++ + "book={s} halves={d} output={s} sampled_rows={d} sampled_columns={d} values={d} " ++ + "baseline_nrms={d:.9} candidate_nrms={d:.9} worst_sampled_column={d} worst_sampled_abs={d:.9} " ++ + "worst_sampled_column_nrms_index={d} worst_sampled_column_nrms={d:.9} " ++ + "baseline_worst_sampled_abs={d:.9} baseline_worst_sampled_column_nrms={d:.9}\n", .{ + matrices[0].k, n, @tagName(matrices[0].codebook), matrices[0].width.halves, @tagName(rows.output), + checked_rows, blocks * 128, aggregate[1].count, aggregate[0].nrms(), aggregate[1].nrms(), + worst_column, worst_abs, worst_nrms_column, worst_nrms, baseline_worst_abs, + baseline_worst_nrms, + }); + if (aggregate[1].nrms() > 0.01) return error.OracleNrmsFailure; +} + +test "oracle error statistics retain cancellation columns instead of averaging them away" { + var stats: Stats = .{}; + try stats.add(1.5, 1); + try stats.add(-0.5, -1); + try std.testing.expectApproxEqAbs(@as(f64, 0.5), stats.nrms(), 1e-15); + try std.testing.expectEqual(@as(f64, 0.5), stats.max_abs); + try std.testing.expectError(error.NonFiniteOracleOutput, stats.add(std.math.nan(f64), 1)); +} diff --git a/zig/kernels/metal/core/exl3_projection.metal b/zig/kernels/metal/core/exl3_projection.metal new file mode 100644 index 0000000000..a76a912ab3 --- /dev/null +++ b/zig/kernels/metal/core/exl3_projection.metal @@ -0,0 +1,230 @@ +// Native EXL3: stored trellis tiles, row-independent compensated dots, and normalized Sylvester rotations. +#include +using namespace metal; + +struct TfExl3Args { + uint k, n, rows, halves; + uint codebook, su_signs, sv_signs, has_bias; + uint output_f32, input_ids, output_ids, routed; + uint expert_id, repeat, indexed, experts; + ulong x_stride, y_stride, tile_k_stride, tile_n_stride; +}; + +// Metal 3 / argument-buffer tier 2 descriptors are populated once from resident buffers' GPU addresses. +struct TfExl3Expert { + const device uchar* trellis; + const device ushort* su; + const device ushort* sv; + const device half* bias; + ulong tile_k_stride, tile_n_stride; + uint halves, codebook, su_signs, sv_signs, has_bias, reserved; +}; + +inline bool tf_exl3_active(const device uint* routes, constant TfExl3Args& a, uint row) { + return row < a.rows && (a.indexed ? routes[row] < a.experts : (!a.routed || routes[row] == a.expert_id)); +} + +inline float tf_exl3_scale(const device ushort* p, uint i, uint signs) { + if (signs) return ((p[i / 16] >> (i % 16)) & 1u) ? -1.0f : 1.0f; + return float(as_type(p[i])); +} + +// Two little-endian int16 words form one MSB-first stream word, including unaligned borrowed tiles. +inline uint tf_exl3_word(const device ushort* tile, uint i) { + return uint(tile[2 * i]) | (uint(tile[2 * i + 1]) << 16); +} + +inline ushort tf_exl3_state(const device ushort* tile, uint halves, uint p) { + const uint ring = 128 * halves; + const uint end = ((p + 1) * halves - ((halves & 1u) ? ((p + 1) & 1u) : 0u)) / 2; + const uint start = (end + ring - 16) % ring; + const uint wi = start / 32, off = start % 32; + const uint lo = tf_exl3_word(tile, wi); + if (off <= 16) return ushort(lo >> (16 - off)); + const uint hi = tf_exl3_word(tile, (wi + 1) % (ring / 32)); + return ushort((lo << (off - 16)) | (hi >> (48 - off))); +} + +inline ushort tf_exl3_mapped_state(const device ushort* tile, uint map) { + const uint wi = map & 63u, off = (map >> 6) & 31u; + const uint lo = tf_exl3_word(tile, wi); + if (off <= 16) return ushort(lo >> (16 - off)); + const uint hi = tf_exl3_word(tile, map >> 11); + return ushort((lo << (off - 16)) | (hi >> (48 - off))); +} + +// Both operations are exact in FP32 before the single FP16 round; do not split the mul1 FMA. +inline half tf_exl3_value(ushort state, uint codebook) { + uint x = uint(state) * (codebook == 2 ? 0x83DCD12Du : codebook == 1 ? 0xCBAC1FEDu : 89226354u); + if (codebook == 2) { + const uint h = 1024 + (x & 255) + ((x >> 8) & 255) + ((x >> 16) & 255) + (x >> 24); + return half(fma(float(h), float(as_type(ushort(0x1EEE))), float(as_type(ushort(0xC931))))); + } + if (codebook == 0) x += 64248484u; + x = (x & 0x8FFF8FFFu) ^ 0x3B603B60u; + return half(float(as_type(ushort(x))) + float(as_type(ushort(x >> 16)))); +} + +// Lane l holds indices l, l+32, l+64, l+96; all seven butterfly stages have a fixed FP32 order. +inline void tf_exl3_hadamard(thread float (&v)[4], uint lane) { + for (uint d = 1; d < 32; d *= 2) { + for (uint j = 0; j < 4; ++j) { + const float peer = simd_shuffle_xor(v[j], ushort(d)); + v[j] = (lane & d) ? peer - v[j] : v[j] + peer; + } + } + const float a = v[0], b = v[1], c = v[2], d = v[3]; + v[0] = a + b; v[1] = a - b; v[2] = c + d; v[3] = c - d; + const float e = v[0], f = v[1], g = v[2], h = v[3]; + v[0] = e + g; v[1] = f + h; v[2] = e - g; v[3] = f - h; +} + +kernel void tf_exl3_rotate_in(const device ushort* x [[buffer(0)]], + const device ushort* su [[buffer(1)]], + device half* rotated [[buffer(2)]], + constant TfExl3Args& a [[buffer(3)]], + const device uint* input_ids [[buffer(4)]], + const device uint* routes [[buffer(5)]], + const device TfExl3Expert* descriptors [[buffer(6)]], + uint lane [[thread_index_in_simdgroup]], + uint3 tg [[threadgroup_position_in_grid]]) { + const uint row = tg.y; + if (!tf_exl3_active(routes, a, row)) return; + const device ushort* scales = a.indexed ? descriptors[routes[row]].su : su; + const uint signs = a.indexed ? descriptors[routes[row]].su_signs : a.su_signs; + const uint source = a.input_ids ? input_ids[row] : row / a.repeat; + const uint base = tg.x * 128; + float v[4]; + for (uint j = 0; j < 4; ++j) { + const uint k = base + lane + 32 * j; + v[j] = as_type(uint(x[ulong(source) * a.x_stride + k]) << 16) * tf_exl3_scale(scales, k, signs); + } + tf_exl3_hadamard(v, lane); + for (uint j = 0; j < 4; ++j) + rotated[ulong(row) * a.k + base + lane + 32 * j] = half(v[j] * 0.08838834764831845f); +} + +// One SIMDgroup owns eight rows and sixteen columns. Only its current 16x16 weight tile is decoded. +template +inline void tf_exl3_matmul_impl(const device half* x, const device uchar* trellis, + device float* sums, constant TfExl3Args& a, + const device uint* routes, const device TfExl3Expert* descriptors, + uint lane, uint3 tg, + threadgroup float* wx, threadgroup float* xx) { + const uint rb = a.indexed == 1 ? tg.y : tg.y * 8, nb = tg.x * 16; + if (!simd_any(tf_exl3_active(routes, a, rb + (a.indexed == 1 ? 0 : lane % 8)))) return; + const uint id = a.indexed ? routes[rb] : 0; + const device uchar* data = a.indexed ? descriptors[id].trellis : trellis; + const ulong ks = a.indexed ? descriptors[id].tile_k_stride : a.tile_k_stride; + const ulong ns = a.indexed ? descriptors[id].tile_n_stride : a.tile_n_stride; + const uint halves = a.indexed ? descriptors[id].halves : a.halves; + const uint codebook = a.indexed ? descriptors[id].codebook : a.codebook; + uint maps[8]; + if (Hoist) { + for (uint j = 0; j < 8; ++j) { + const uint p = lane + 32 * j, ring = 128 * halves; + const uint end = ((p + 1) * halves - ((halves & 1u) ? ((p + 1) & 1u) : 0u)) / 2; + const uint start = (end + ring - 16) % ring; + const uint wi = start / 32, off = start % 32, next = wi + 1 == 4 * halves ? 0 : wi + 1; + maps[j] = wi | (off << 6) | (next << 11); + } + } + float high[4]={0,0,0,0}, low[4]={0,0,0,0}; + for (uint kb = 0; kb < a.k; kb += 16) { + const device ushort* tile = reinterpret_cast( + data + ulong(kb / 16) * ks + ulong(nb / 16) * ns); + for (uint j = 0; j < 8; ++j) { + const uint p = lane + 32 * j, l = p / 8, e = p % 8; + const uint r = 2 * (l % 4) + (e & 1u) + 8 * ((e >> 1) & 1u); + const uint c = l / 4 + 8 * (e >> 2); + const ushort state = Hoist ? tf_exl3_mapped_state(tile, maps[j]) : tf_exl3_state(tile, halves, p); + wx[r * 16 + c] = float(tf_exl3_value(state, codebook)); + } + for (uint j = lane; j < 128; j += 32) { + const uint row = rb + j / 16; + xx[j] = (a.indexed != 1 || j / 16 == 0) && tf_exl3_active(routes, a, row) + ? float(x[ulong(row) * a.k + kb + j % 16]) : 0.0f; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint i=0;i<4;++i) { + const uint at=lane+32*i, row=rb+at/16, col=at%16; + if ((a.indexed==1 && at/16!=0) || !tf_exl3_active(routes,a,row)) continue; + for (uint k=0;k<16;++k) { + // FP16 products are exact in FP32. TwoSum retains each addition's rounding error. + const float product=xx[(at/16)*16+k]*wx[k*16+col]; + const float next=high[i]+product; + const float borrowed=next-high[i]; + low[i]+=(high[i]-(next-borrowed))+(product-borrowed); + high[i]=next; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + for (uint i=0;i<4;++i) { + const uint at=lane+32*i, row=rb+at/16; + if ((a.indexed!=1 || at/16==0) && tf_exl3_active(routes,a,row)) + sums[ulong(row)*a.n+nb+at%16]=high[i]+low[i]; + } +} + +kernel void tf_exl3_matmul(const device half* x [[buffer(0)]], + const device uchar* trellis [[buffer(1)]], + device float* sums [[buffer(2)]], + constant TfExl3Args& a [[buffer(3)]], + const device uint* routes [[buffer(4)]], + const device TfExl3Expert* descriptors [[buffer(5)]], + uint lane [[thread_index_in_simdgroup]], + uint3 tg [[threadgroup_position_in_grid]]) { + threadgroup float wx[256], xx[128]; + tf_exl3_matmul_impl(x, trellis, sums, a, routes, descriptors, lane, tg, wx, xx); +} + + +kernel void tf_exl3_matmul_hoisted(const device half* x [[buffer(0)]], + const device uchar* trellis [[buffer(1)]], + device float* sums [[buffer(2)]], + constant TfExl3Args& a [[buffer(3)]], + const device uint* routes [[buffer(4)]], + const device TfExl3Expert* descriptors [[buffer(5)]], + uint lane [[thread_index_in_simdgroup]], + uint3 tg [[threadgroup_position_in_grid]]) { + threadgroup float wx[256], xx[128]; + tf_exl3_matmul_impl(x, trellis, sums, a, routes, descriptors, lane, tg, wx, xx); +} + +// Qualification only: compare every arithmetic state against the independent CPU-generated FP16 table. +kernel void tf_exl3_codebook_check(device half* arithmetic [[buffer(0)]], + uint i [[thread_position_in_grid]]) { + if (i < 3 * 65536) arithmetic[i] = tf_exl3_value(ushort(i % 65536), i / 65536); +} + +kernel void tf_exl3_rotate_out(const device float* sums [[buffer(0)]], + const device ushort* sv [[buffer(1)]], + const device half* bias [[buffer(2)]], + device uchar* y [[buffer(3)]], + constant TfExl3Args& a [[buffer(4)]], + const device uint* output_ids [[buffer(5)]], + const device uint* routes [[buffer(6)]], + const device TfExl3Expert* descriptors [[buffer(7)]], + uint lane [[thread_index_in_simdgroup]], + uint3 tg [[threadgroup_position_in_grid]]) { + const uint row = tg.y, base = tg.x * 128; + if (!tf_exl3_active(routes, a, row)) return; + const uint id = a.indexed ? routes[row] : 0; + const device ushort* scales = a.indexed ? descriptors[id].sv : sv; + const device half* biases = a.indexed ? descriptors[id].bias : bias; + const uint signs = a.indexed ? descriptors[id].sv_signs : a.sv_signs; + const uint has_bias = a.indexed ? descriptors[id].has_bias : a.has_bias; + float v[4]; + for (uint j = 0; j < 4; ++j) v[j] = sums[ulong(row) * a.n + base + lane + 32 * j]; + tf_exl3_hadamard(v, lane); + const uint dest = a.output_ids ? output_ids[row] : row; + for (uint j = 0; j < 4; ++j) { + const uint n = base + lane + 32 * j; + float out = (v[j] * 0.08838834764831845f) * tf_exl3_scale(scales, n, signs); + if (has_bias) out += float(biases[n]); + const ulong at = ulong(dest) * a.y_stride + n; + if (a.output_f32) reinterpret_cast(y)[at] = out; + else reinterpret_cast(y)[at] = bfloat(out); + } +} 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..532bbd01ee --- /dev/null +++ b/zig/kernels/metal/deepseek_v41/engram.metal @@ -0,0 +1,150 @@ +// Math/layout attribution: DeepSeek AI DeepSeek-V4.1-Flash inference/{engram,model}.py (MIT). +#include +using namespace metal; + +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; + 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 float* q_weight [[buffer(2)]], + const device float* 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 = q_weight[wbase + d] * k_weight[wbase + d]; + 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 16 or + (m.width.isHalf() and (m.codebook != .mul1 or m.width.halves > 7))) return error.UnsupportedExl3Width; + const row_bytes = try std.math.mul(usize, m.n / 16, m.nStride()); + const k_stride = if (m.tile_k_stride == 0) row_bytes else m.tile_k_stride; + if (m.nStride() < m.tileBytes() or m.nStride() % 2 != 0 or + k_stride < row_bytes or k_stride % 2 != 0) + return error.InvalidExl3Stride; + } + + pub fn tileBytes(m: Matrix) usize { + return m.width.tileWords() * 2; + } + pub fn nStride(m: Matrix) usize { + return if (m.tile_n_stride == 0) m.tileBytes() else m.tile_n_stride; + } + pub fn kStride(m: Matrix) !usize { + return if (m.tile_k_stride == 0) std.math.mul(usize, m.n / 16, m.nStride()) else m.tile_k_stride; + } + pub fn storageBytes(m: Matrix) !usize { + try m.validate(); + const last_k = try std.math.mul(usize, m.k / 16 - 1, try m.kStride()); + const last_n = try std.math.mul(usize, m.n / 16 - 1, m.nStride()); + return std.math.add(usize, try std.math.add(usize, last_k, last_n), m.tileBytes()); + } + fn checkStorage(m: Matrix) !usize { + const bytes = try m.storageBytes(); + try check(m.trellis, bytes, 2); + try check(m.suh, m.k * 2 / (if (m.suh_format == .packed_signs) @as(usize, 16) else 1), 2); + try check(m.svh, m.n * 2 / (if (m.svh_format == .packed_signs) @as(usize, 16) else 1), 2); + if (m.bias) |b| try check(b, m.n * 2, 2); + return bytes; + } +}; + +/// Individual expert filtering reads u32 IDs on GPU; repeat top-k slots share one input row. +pub const Routes = struct { ids: View, expert_id: u32 = 0, repeat: usize = 1 }; + +/// Strides are elements, X is BF16, IDs are u32[count]. Gather/scatter IDs must be in bounds and outputs unique. +pub const Rows = struct { + count: usize, + x_stride: usize, + y_stride: usize, + output: Output = .bf16, + input_ids: ?View = null, + output_ids: ?View = null, + /// Required for ID maps; caller guarantees IDs below capacities and unique output IDs. + input_rows: usize = 0, + output_rows: usize = 0, + routes: ?Routes = null, + /// Caller-prepared eight-row tiles, same expert per tile; maxInt(u32) padding must follow valid rows. + grouped: bool = false, +}; + +/// Metal 3 / tier-2 argument-buffer ABI; borrowed pointers remain resident until the encoded work completes. +pub const ExpertDescriptor = extern struct { + trellis: u64, + suh: u64, + svh: u64, + bias: u64, + tile_k_stride: u64, + tile_n_stride: u64, + halves: u32, + codebook: u32, + suh_signs: u32, + svh_signs: u32, + has_bias: u32, + reserved: u32 = 0, +}; + +/// Populate a caller-owned shared descriptor table once; no weight copying, packing or dequantization. +pub fn descriptor(m: Matrix) !ExpertDescriptor { + _ = try m.checkStorage(); + return .{ + .trellis = try address(m.trellis), + .suh = try address(m.suh), + .svh = try address(m.svh), + .bias = if (m.bias) |v| try address(v) else 0, + .tile_k_stride = try m.kStride(), + .tile_n_stride = m.nStride(), + .halves = m.width.halves, + .codebook = @intCast(@backingInt(m.codebook)), + .suh_signs = @intFromBool(m.suh_format == .packed_signs), + .svh_signs = @intFromBool(m.svh_format == .packed_signs), + .has_bias = @intFromBool(m.bias != null), + }; +} + +fn address(v: View) !u64 { + return std.math.add(u64, v.buffer.gpuAddress(), v.offset); +} + +/// Call once per encoder for indirectly addressed matrices, unless the caller declares them via residency sets. +pub fn useMatrix(e: mtl.ComputeEncoder, m: Matrix) void { + e.useResource(m.trellis.buffer, mtl.ResourceUsage.read); + e.useResource(m.suh.buffer, mtl.ResourceUsage.read); + e.useResource(m.svh.buffer, mtl.ResourceUsage.read); + if (m.bias) |b| e.useResource(b.buffer, mtl.ResourceUsage.read); +} + +/// Common K,N; descriptor weights must remain resident and disjoint from writable views through completion. +pub const Experts = struct { descriptors: View, count: usize, k: usize, n: usize }; + +pub const Pipelines = struct { + rotate_in: mtl.Pipeline, + matmul: mtl.Pipeline, + matmul_hoisted: mtl.Pipeline, + decode: Decode = .hoisted_decode, + rotate_out: 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_exl3_projection, mtl.CompileOptions.mlx()); + defer lib.deinit(); + const ri = try mtl.Pipeline.init(device, lib, "tf_exl3_rotate_in", false); + errdefer ri.deinit(); + const mm = try mtl.Pipeline.init(device, lib, "tf_exl3_matmul", false); + errdefer mm.deinit(); + const hoisted = try mtl.Pipeline.init(device, lib, "tf_exl3_matmul_hoisted", false); + errdefer hoisted.deinit(); + const ro = try mtl.Pipeline.init(device, lib, "tf_exl3_rotate_out", false); + errdefer ro.deinit(); + for ([_]mtl.Pipeline{ ri, mm, hoisted, ro }) |p| + if (p.simdWidth() != 32 or p.maxThreads() < 32) return error.UnsupportedExl3Simd; + return .{ .rotate_in = ri, .matmul = mm, .matmul_hoisted = hoisted, .rotate_out = ro }; + } + pub fn deinit(p: *Pipelines) void { + p.rotate_in.deinit(); + p.matmul.deinit(); + p.matmul_hoisted.deinit(); + p.rotate_out.deinit(); + } +}; + +const Args = extern struct { + k: u32, + n: u32, + rows: u32, + halves: u32, + codebook: u32, + su_signs: u32, + sv_signs: u32, + has_bias: u32, + output_f32: u32, + input_ids: u32, + output_ids: u32, + routed: u32, + expert_id: u32, + repeat: u32, + indexed: u32 = 0, + experts: u32 = 0, + x_stride: u64, + y_stride: u64, + tile_k_stride: u64, + tile_n_stride: u64, +}; + +fn shape(k: usize, n: usize) !void { + if (k == 0 or n == 0 or k % 128 != 0 or n % 128 != 0 or + k > std.math.maxInt(u32) or n > std.math.maxInt(u32)) return error.UnsupportedExl3Shape; +} +fn rotatedBytes(k: usize, rows: usize) !usize { + const bytes = try std.math.mul(usize, try std.math.mul(usize, k, rows), 2); + return (try std.math.add(usize, bytes, 15)) & ~@as(usize, 15); +} + +/// Caller allocates this once for its largest window; scratch never holds a dequantized full weight matrix. +pub fn scratchBytes(k: usize, n: usize, rows: usize) !usize { + try shape(k, n); + if (rows == 0 or rows > std.math.maxInt(u32) - 7) return error.InvalidExl3Rows; + return std.math.add(usize, try rotatedBytes(k, rows), try std.math.mul(usize, try std.math.mul(usize, n, rows), 4)); +} +fn check(v: View, bytes: usize, alignment: usize) !void { + const len = v.buffer.length(); + if (v.offset % alignment != 0 or v.offset > len or bytes > len - v.offset) return error.InvalidExl3View; +} +fn disjoint(a: View, a_bytes: usize, b: View, b_bytes: usize) !void { + if (a.buffer.id == b.buffer.id and a.offset < b.offset + b_bytes and b.offset < a.offset + a_bytes) + return error.AliasedExl3View; +} +fn readonly(v: View, bytes: usize, y: View, y_bytes: usize, scratch: View, scratch_bytes: usize) !void { + try disjoint(v, bytes, y, y_bytes); + try disjoint(v, bytes, scratch, scratch_bytes); +} +fn outputBytes(n: usize, r: Rows) !usize { + const rows = if (r.output_ids != null) r.output_rows else r.count; + const end = try std.math.add(usize, try std.math.mul(usize, rows - 1, r.y_stride), n); + return std.math.mul(usize, end, if (r.output == .f32) @as(usize, 4) else 2); +} +fn rowsArgs(k: usize, n: usize, r: Rows, x: View, y: View, scratch: View) !Args { + const scratch_size = try scratchBytes(k, n, r.count); + if (r.x_stride < k or r.y_stride < n) return error.InvalidExl3Stride; + const repeat = if (r.routes) |v| v.repeat else 1; + if (repeat == 0 or repeat > std.math.maxInt(u32) or r.count % repeat != 0) return error.InvalidExl3Repeat; + if ((r.input_ids != null and r.input_rows == 0) or + (r.output_ids != null and r.output_rows == 0)) return error.InvalidExl3Capacity; + const xrows = if (r.input_ids != null) r.input_rows else r.count / repeat; + const xend = try std.math.add(usize, try std.math.mul(usize, xrows - 1, r.x_stride), k); + const xbytes = try std.math.mul(usize, xend, 2); + const ybytes = try outputBytes(n, r); + try check(x, xbytes, 2); + try check(y, ybytes, if (r.output == .f32) @as(usize, 4) else 2); + try check(scratch, scratch_size, 16); + try readonly(x, xbytes, y, ybytes, scratch, scratch_size); + try disjoint(y, ybytes, scratch, scratch_size); + const id_bytes = try std.math.mul(usize, r.count, 4); + for ([_]?View{ r.input_ids, r.output_ids, if (r.routes) |v| v.ids else null }) |optional| { + if (optional) |v| { + try check(v, id_bytes, 4); + try readonly(v, id_bytes, y, ybytes, scratch, scratch_size); + } + } + return .{ + .k = @intCast(k), + .n = @intCast(n), + .rows = @intCast(r.count), + .halves = 0, + .codebook = 0, + .su_signs = 0, + .sv_signs = 0, + .has_bias = 0, + .output_f32 = @intFromBool(r.output == .f32), + .input_ids = @intFromBool(r.input_ids != null), + .output_ids = @intFromBool(r.output_ids != null), + .routed = @intFromBool(r.routes != null), + .expert_id = if (r.routes) |v| v.expert_id else 0, + .repeat = @intCast(repeat), + .x_stride = r.x_stride, + .y_stride = r.y_stride, + .tile_k_stride = 0, + .tile_n_stride = 0, + }; +} +fn bind(e: mtl.ComputeEncoder, v: View, index: usize) void { + e.setBuffer(v.buffer, v.offset, index); +} + +/// No allocation or commit; retain all views until completion, disjoint from writes, with valid GPU IDs. +pub fn encode(p: *const Pipelines, e: mtl.ComputeEncoder, m: Matrix, x: View, y: View, scratch: View, r: Rows) !void { + const storage_bytes = try m.checkStorage(); + var a = try rowsArgs(m.k, m.n, r, x, y, scratch); + if (r.grouped) return error.InvalidExl3GroupedRows; + const ybytes = try outputBytes(m.n, r); + const scratch_bytes = try scratchBytes(m.k, m.n, r.count); + try readonly(m.trellis, storage_bytes, y, ybytes, scratch, scratch_bytes); + try readonly(m.suh, m.k * 2 / (if (m.suh_format == .packed_signs) @as(usize, 16) else 1), y, ybytes, scratch, scratch_bytes); + try readonly(m.svh, m.n * 2 / (if (m.svh_format == .packed_signs) @as(usize, 16) else 1), y, ybytes, scratch, scratch_bytes); + if (m.bias) |b| try readonly(b, m.n * 2, y, ybytes, scratch, scratch_bytes); + a.halves = m.width.halves; + a.codebook = @intCast(@backingInt(m.codebook)); + a.su_signs = @intFromBool(m.suh_format == .packed_signs); + a.sv_signs = @intFromBool(m.svh_format == .packed_signs); + a.has_bias = @intFromBool(m.bias != null); + a.tile_k_stride = try m.kStride(); + a.tile_n_stride = m.nStride(); + stages(p, e, a, m.trellis, m.suh, m.svh, m.bias orelse m.svh, scratch, x, y, scratch, r); +} + +/// GPU u32 routes select descriptors; IDs must be valid, unused eight-row tile slots are zero padded. +pub fn encodeIndexed(p: *const Pipelines, e: mtl.ComputeEncoder, experts: Experts, x: View, y: View, scratch: View, r: Rows) !void { + if (r.routes == null or experts.count == 0 or experts.count > std.math.maxInt(u32)) return error.InvalidExl3Experts; + if (r.grouped and (r.input_ids == null or r.output_ids == null or r.count % 8 != 0)) + return error.InvalidExl3GroupedRows; + try shape(experts.k, experts.n); + try check(experts.descriptors, try std.math.mul(usize, experts.count, @sizeOf(ExpertDescriptor)), 8); + var a = try rowsArgs(experts.k, experts.n, r, x, y, scratch); + try readonly(experts.descriptors, experts.count * @sizeOf(ExpertDescriptor), y, try outputBytes(experts.n, r), scratch, try scratchBytes(experts.k, experts.n, r.count)); + a.indexed = if (r.grouped) 2 else 1; + a.experts = @intCast(experts.count); + stages(p, e, a, scratch, scratch, scratch, scratch, experts.descriptors, x, y, scratch, r); +} + +fn stages(p: *const Pipelines, e: mtl.ComputeEncoder, a: Args, trellis: View, su: View, sv: View, bias: View, descriptors: View, x: View, y: View, scratch: View, r: Rows) void { + const sums = View{ .buffer = scratch.buffer, .offset = scratch.offset + ((@as(usize, a.rows) * a.k * 2 + 15) & ~@as(usize, 15)) }; + const routes = if (r.routes) |v| v.ids else scratch; + e.setPipeline(p.rotate_in); + bind(e, x, 0); + bind(e, su, 1); + bind(e, scratch, 2); + e.setValue(a, 3); + bind(e, r.input_ids orelse scratch, 4); + bind(e, routes, 5); + bind(e, descriptors, 6); + e.dispatchGroups(.{ .width = a.k / 128, .height = a.rows }, .{ .width = 32 }); + e.barrier(); + e.setPipeline(switch (p.decode) { + .arithmetic => p.matmul, + .hoisted_decode => p.matmul_hoisted, + }); + bind(e, scratch, 0); + bind(e, trellis, 1); + bind(e, sums, 2); + e.setValue(a, 3); + bind(e, routes, 4); + bind(e, descriptors, 5); + e.dispatchGroups(.{ .width = a.n / 16, .height = if (a.indexed == 1) a.rows else (@as(usize, a.rows) + 7) / 8 }, .{ .width = 32 }); + e.barrier(); + e.setPipeline(p.rotate_out); + bind(e, sums, 0); + bind(e, sv, 1); + bind(e, bias, 2); + bind(e, y, 3); + e.setValue(a, 4); + bind(e, r.output_ids orelse scratch, 5); + bind(e, routes, 6); + bind(e, descriptors, 7); + e.dispatchGroups(.{ .width = a.n / 128, .height = a.rows }, .{ .width = 32 }); + e.barrier(); +} + +test "EXL3 scratch is row compact and bounded independently of full matrices" { + try std.testing.expectEqual(@as(usize, 214016), try scratchBytes(5120, 2304, 11)); + try std.testing.expectError(error.InvalidExl3Rows, scratchBytes(128, 128, 0)); + try std.testing.expectError(error.UnsupportedExl3Shape, scratchBytes(129, 128, 1)); + try std.testing.expectEqual(@as(usize, 72), @sizeOf(ExpertDescriptor)); + try std.testing.expectEqual(@as(usize, 96), @sizeOf(Args)); +} + +test "stored per-matrix width and rectangular strides include public target shapes" { + var m = Matrix{ .trellis = undefined, .suh = undefined, .svh = undefined, .k = 5120, .n = 2304, .width = .{ .halves = 6 } }; + try m.validate(); + try std.testing.expectEqual(@as(usize, 4423680), try m.storageBytes()); + m.width.halves = 5; + try m.validate(); + m.codebook = .mcg; + try std.testing.expectError(error.UnsupportedExl3Width, m.validate()); + m.width.halves = 16; + m.k = 1280; + m.n = 4096; + try m.validate(); + m.tile_n_stride = 258; + try m.validate(); + m.tile_k_stride = 256; + try std.testing.expectError(error.InvalidExl3Stride, m.validate()); +} + +test "invalid and overflowing extents fail before storage arithmetic can wrap" { + var m = Matrix{ .trellis = undefined, .suh = undefined, .svh = undefined, .k = 0, .n = 128, .width = .{ .halves = 6 } }; + try std.testing.expectError(error.UnsupportedExl3Shape, m.storageBytes()); + m.k = 128; + m.n = 256; + m.tile_n_stride = std.math.maxInt(usize) - 1; + try std.testing.expectError(error.Overflow, m.validate()); + m.tile_n_stride = m.tileBytes(); + m.tile_k_stride = std.math.maxInt(usize) - 1; + try std.testing.expectError(error.Overflow, m.storageBytes()); + try std.testing.expectError(error.InvalidExl3Rows, scratchBytes(128, 128, std.math.maxInt(u32))); + try std.testing.expectError(error.Overflow, scratchBytes(0xffffff80, 0xffffff80, 0xffffff80)); +} + +test { + _ = decoder; +} diff --git a/zig/src/core/exl3_projection_decode.zig b/zig/src/core/exl3_projection_decode.zig new file mode 100644 index 0000000000..fb4f8a90d5 --- /dev/null +++ b/zig/src/core/exl3_projection_decode.zig @@ -0,0 +1,121 @@ +//! CPU mirror of the Metal two-word circular decoder, tested against the deliberately scalar format oracle. +const std = @import("std"); +const fmt = @import("exl3_format.zig"); + +fn word(tile: []const i16, i: usize) u32 { + return @as(u32, @as(u16, @bitCast(tile[2 * i]))) | (@as(u32, @as(u16, @bitCast(tile[2 * i + 1]))) << 16); +} +pub fn state(width: fmt.Width, tile: []const i16, p: u32) u16 { + const ring: u32 = 128 * @as(u32, width.halves); + const end = ((p + 1) * width.halves - (if (width.isHalf()) (p + 1) & 1 else @as(u32, 0))) / 2; + const start = (end + ring - 16) % ring; + const wi = start / 32; + const off = start % 32; + const lo = word(tile, wi); + if (off <= 16) return @truncate(lo >> @as(u5, @intCast(16 - off))); + const hi = word(tile, (wi + 1) % (ring / 32)); + return @truncate((lo << @as(u5, @intCast(off - 16))) | (hi >> @as(u5, @intCast(48 - off)))); +} + +/// FP32 arithmetic used by Metal is exact here, followed by exactly one FP16 rounding. +pub fn value(codebook: fmt.Codebook, s: u16) f16 { + var x: u32 = @as(u32, s) *% switch (codebook) { + .mul1 => fmt.Mul1Mul, + .mcg => fmt.McgMul, + .inst3 => fmt.Inst3Mul, + }; + if (codebook == .mul1) { + const h: f32 = @floatFromInt(1024 + (x & 255) + ((x >> 8) & 255) + ((x >> 16) & 255) + (x >> 24)); + const scale: f16 = @bitCast(fmt.Mul1ScaleBits); + const bias: f16 = @bitCast(fmt.Mul1BiasBits); + return @floatCast(@mulAdd(f32, h, @floatCast(scale), @floatCast(bias))); + } + if (codebook == .inst3) x +%= fmt.Inst3Add; + x = (x & fmt.Mask) ^ fmt.Flip; + const lo: f16 = @bitCast(@as(u16, @truncate(x))); + const hi: f16 = @bitCast(@as(u16, @truncate(x >> 16))); + return @floatCast(@as(f32, @floatCast(lo)) + @as(f32, @floatCast(hi))); +} + +/// The row/column inverse of the CUDA fragment permutation lets the MMA decoder place each stream value. +pub fn placement(p: u32) fmt.Placement { + const l = p / 8; + const j = p % 8; + return .{ .row = @intCast(2 * (l % 4) + (j & 1) + 8 * ((j >> 1) & 1)), .col = @intCast(l / 4 + 8 * (j >> 2)) }; +} + +/// Same seven FP32 butterfly stages as the SIMDgroup transform, normalized only at the end. +pub fn hadamard(block: *[128]f32) void { + var d: usize = 1; + while (d < 128) : (d *= 2) { + var base: usize = 0; + while (base < 128) : (base += 2 * d) { + for (0..d) |j| { + const a = block[base + j]; + const b = block[base + j + d]; + block[base + j] = a + b; + block[base + j + d] = a - b; + } + } + } + for (block) |*x| x.* *= 0.08838834764831845; +} + +test "two-word circular extraction agrees at every stream boundary and supported half width" { + var prng = std.Random.DefaultPrng.init(0x45584c3353544154); + const random = prng.random(); + var tile: [128]i16 = undefined; + for ([_]u16{ 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16 }) |halves| { + const width = fmt.Width{ .halves = halves }; + for (0..12) |pattern| { + for (tile[0..width.tileWords()], 0..) |*v, j| v.* = switch (pattern) { + 0 => 0, + 1 => -1, + 2 => @bitCast(@as(u16, if (j % 2 == 0) 0x8001 else 0x1234)), + else => random.int(i16), + }; + var expected: [256]u32 = undefined; + fmt.tileStates(width, tile[0..width.tileWords()], &expected); + for (expected, 0..) |want, p| try std.testing.expectEqual(@as(u16, @intCast(want)), state(width, tile[0..width.tileWords()], @intCast(p))); + } + } +} + +test "Metal FP32 codebook arithmetic preserves FP16 bits for all 65536 states and all codebooks" { + for ([_]fmt.Codebook{ .inst3, .mcg, .mul1 }) |codebook| { + for (0..65536) |s| { + const got: u16 = @bitCast(value(codebook, @intCast(s))); + const want: u16 = @bitCast(fmt.codebookValue(codebook, @intCast(s))); + try std.testing.expectEqual(want, got); + } + } +} + +test "fragment placement is a bijection agreeing with the scalar oracle" { + var expected: [256]fmt.Placement = undefined; + fmt.tilePlacement(&expected); + var visited: [256]bool = @splat(false); + for (expected, 0..) |want, p| { + const got = placement(@intCast(p)); + try std.testing.expectEqual(want, got); + const at = @as(usize, got.row) * 16 + got.col; + try std.testing.expect(!visited[at]); + visited[at] = true; + } +} + +test "fixed butterfly transform preserves block signs and is its own inverse" { + var x: [128]f32 = undefined; + for (&x, 0..) |*v, i| v.* = @as(f32, @floatFromInt(@as(i32, @intCast(i)) - 64)) / 32; + const before = x; + hadamard(&x); + hadamard(&x); + for (x, before) |got, want| try std.testing.expectApproxEqAbs(want, got, 0.000002); + var impulse: [128]f32 = @splat(0); + impulse[67] = 1; + hadamard(&impulse); + for (impulse, 0..) |got, i| { + const want: f32 = if (@popCount(i & 67) & 1 == 0) 0.08838834764831845 else -0.08838834764831845; + try std.testing.expectEqual(want, got); + } +} diff --git a/zig/src/core/exl3_projection_gpu_test.zig b/zig/src/core/exl3_projection_gpu_test.zig new file mode 100644 index 0000000000..5abe91d3c8 --- /dev/null +++ b/zig/src/core/exl3_projection_gpu_test.zig @@ -0,0 +1,420 @@ +//! Explicit GPU qualification tests; run through the native test module on a Metal SIMDgroup-matrix device. +const std = @import("std"); +const mtl = @import("metal"); +const ex = @import("exl3_projection.zig"); +const k = 256; +const n = 128; +const rows = 9; + +const Owned = struct { + matrix: ex.Matrix, + buffers: [4]mtl.Buffer, + + fn init(device: mtl.Device, halves: u16, book: ex.format.Codebook, seed: u64) !Owned { + const words = @as(usize, halves) * 8 * (k / 16) * (n / 16); + const t = try device.buffer(words * 2 + 16, mtl.ResourceOptions.shared); + errdefer t.deinit(); + const su = try device.buffer(k * 2 + 16, mtl.ResourceOptions.shared); + errdefer su.deinit(); + const sv = try device.buffer(n * 2 + 16, mtl.ResourceOptions.shared); + errdefer sv.deinit(); + const bias = try device.buffer(n * 2 + 16, mtl.ResourceOptions.shared); + errdefer bias.deinit(); + var prng = std.Random.DefaultPrng.init(seed); + for (t.slice(i16, words + 8)[8..]) |*x| x.* = prng.random().int(i16); + for (su.slice(f16, k + 8)[8..], 0..) |*x, i| x.* = if (i % 3 == 0) -0.75 else 1.125; + for (sv.slice(f16, n + 8)[8..], 0..) |*x, i| x.* = if (i % 5 == 0) -0.25 else 0.875; + for (bias.slice(f16, n + 8)[8..], 0..) |*x, i| x.* = @floatCast(@as(f32, @floatFromInt(i % 7)) / 16); + return .{ .matrix = .{ .trellis = .{ .buffer = t, .offset = 16 }, .suh = .{ .buffer = su, .offset = 16 }, .svh = .{ .buffer = sv, .offset = 16 }, .bias = .{ .buffer = bias, .offset = 16 }, .k = k, .n = n, .width = .{ .halves = halves }, .codebook = book }, .buffers = .{ t, su, sv, bias } }; + } + fn deinit(w: Owned) void { + for (w.buffers) |b| b.deinit(); + } +}; + +fn bf16(x: f32) u16 { + const bits: u32 = @bitCast(x); + return @truncate((bits +% (0x7fff + ((bits >> 16) & 1))) >> 16); +} +fn input(device: mtl.Device, count: usize) !mtl.Buffer { + const x = try device.buffer(count * (k + 16) * 2 + 32, mtl.ResourceOptions.shared); + for (x.slice(u16, count * (k + 16) + 16), 0..) |*v, i| + v.* = bf16(@as(f32, @floatFromInt(@as(i32, @intCast(i % 23)) - 11)) / 32); + return x; +} +fn complete(cb: mtl.CommandBuffer, e: mtl.ComputeEncoder) !void { + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure() != null) return error.Exl3GpuFailure; +} + +fn reference(w: Owned, x: []const u16, out: *[n]f32) void { + var rotated: [k]f16 = undefined; + const su = w.buffers[1].slice(f16, k + 8)[8..]; + for (0..k / 128) |b| { + var block: [128]f32 = undefined; + for (&block, 0..) |*v, i| v.* = @as(f32, @bitCast(@as(u32, x[b * 128 + i]) << 16)) * @as(f32, su[b * 128 + i]); + ex.decoder.hadamard(&block); + for (block, 0..) |v, i| rotated[b * 128 + i] = @floatCast(v); + } + const word_count = w.matrix.width.tileWords() * (k / 16) * (n / 16); + const words = w.buffers[0].slice(i16, word_count + 8)[8..]; + var unpacked: [k * n]f16 = undefined; + ex.format.unpack(w.matrix.width, w.matrix.codebook, words, k / 16, n / 16, &unpacked); + for (out, 0..) |*v, col| { + var sum: f64 = 0; + for (0..k) |j| sum += @as(f64, rotated[j]) * @as(f64, unpacked[j * n + col]); + v.* = @floatCast(sum); + } + ex.decoder.hadamard(out); + const sv = w.buffers[2].slice(f16, n + 8)[8..]; + const bias = w.buffers[3].slice(f16, n + 8)[8..]; + for (out, 0..) |*v, i| v.* = v.* * @as(f32, sv[i]) + @as(f32, bias[i]); +} + +test "EXL3 long cancellation retains small terms in BF16 projection output" { + const long_k = 4096; + const blocks = long_k / 128; + const device = try mtl.Device.init(); + defer device.deinit(); + var pipelines = try ex.Pipelines.load(device); + defer pipelines.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + const base = try Owned.init(device, 6, .mul1, 0x71717); + defer base.deinit(); + const block_words = base.matrix.width.tileWords() * (128 / 16) * (n / 16); + const block_bytes = base.buffers[0].contents()[16..][0 .. block_words * 2]; + const trellis = try device.buffer(blocks * block_bytes.len, mtl.ResourceOptions.shared); + defer trellis.deinit(); + for (0..blocks) |i| @memcpy(trellis.contents()[i * block_bytes.len ..][0..block_bytes.len], block_bytes); + const scales = try device.buffer(long_k * 2, mtl.ResourceOptions.shared); + defer scales.deinit(); + @memset(scales.slice(f16, long_k), 1); + var matrix = base.matrix; + matrix.k = long_k; + matrix.trellis = .{ .buffer = trellis }; + matrix.suh = .{ .buffer = scales }; + matrix.bias = null; + const x = try device.buffer(long_k * 2, mtl.ResourceOptions.shared); + defer x.deinit(); + var rotated: [long_k]f16 = undefined; + for (0..blocks) |i| { + const v: f32 = if (i == 0) 4096 else if (i == blocks - 1) -4096 else 0x1p-13; + @memset(x.slice(u16, long_k)[i * 128 ..][0..128], bf16(v)); + var block: [128]f32 = @splat(v); + ex.decoder.hadamard(&block); + for (block, 0..) |value, j| rotated[i * 128 + j] = @floatCast(value); + } + // Equal large endpoint blocks cancel; intervening terms must survive the long reduction. + var decoded: [128 * n]f16 = undefined; + ex.format.unpack(base.matrix.width, base.matrix.codebook, base.buffers[0].slice(i16, block_words + 8)[8..], 128 / 16, n / 16, &decoded); + var expected: [n]f32 = undefined; + for (&expected, 0..) |*value, col| { + var sum: f64 = 0; + for (rotated, 0..) |v, j| sum += @as(f64, v) * @as(f64, decoded[(j % 128) * n + col]); + value.* = @floatCast(sum); + } + ex.decoder.hadamard(&expected); + const sv = base.buffers[2].slice(f16, n + 8)[8..]; + var expected_bf16: [n]u16 = undefined; + for (expected, sv, &expected_bf16) |v, scale, *out| out.* = bf16(v * @as(f32, scale)); + const y = try device.buffer(n * 2, mtl.ResourceOptions.shared); + defer y.deinit(); + const scratch = try device.buffer(try ex.scratchBytes(long_k, n, 1), mtl.ResourceOptions.shared); + defer scratch.deinit(); + const cb = queue.commandBuffer(); + const encoder = cb.compute(.concurrent); + try ex.encode(&pipelines, encoder, matrix, .{ .buffer = x }, .{ .buffer = y }, .{ .buffer = scratch }, .{ .count = 1, .x_stride = long_k, .y_stride = n }); + try complete(cb, encoder); + try std.testing.expectEqualSlices(u16, &expected_bf16, y.slice(u16, n)); +} + +test "Metal EXL3 codebooks and half widths match the rotated scalar projection" { + const device = try mtl.Device.init(); + defer device.deinit(); + var p = try ex.Pipelines.load(device); + defer p.deinit(); + for ([_]ex.Decode{ .arithmetic, .hoisted_decode }) |decode| { + p.decode = decode; + const queue = try device.queue(); + defer queue.deinit(); + const x = try input(device, rows); + defer x.deinit(); + const y = try device.buffer(rows * (n + 8) * 4 + 32, mtl.ResourceOptions.shared); + defer y.deinit(); + const scratch = try device.buffer(try ex.scratchBytes(k, n, rows), mtl.ResourceOptions.shared); + defer scratch.deinit(); + for ([_]ex.format.Codebook{ .inst3, .mcg, .mul1 }) |book| { + for ([_]u16{ 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16 }) |halves| { + if (halves % 2 == 1 and book != .mul1) continue; + const w = try Owned.init(device, halves, book, 0xdec0de + @as(u64, halves)); + defer w.deinit(); + const cb = queue.commandBuffer(); + const e = cb.compute(.concurrent); + try ex.encode(&p, e, w.matrix, .{ .buffer = x, .offset = 32 }, .{ .buffer = y, .offset = 32 }, .{ .buffer = scratch }, .{ .count = rows, .x_stride = k + 16, .y_stride = n + 8, .output = .f32 }); + try complete(cb, e); + for (0..rows) |r| { + var want: [n]f32 = undefined; + reference(w, x.slice(u16, rows * (k + 16) + 16)[16 + r * (k + 16) ..][0..k], &want); + const got = y.slice(f32, rows * (n + 8) + 8)[8 + r * (n + 8) ..][0..n]; + var mag: f32 = 1; + for (want) |v| mag = @max(mag, @abs(v)); + for (got, want) |a, b| try std.testing.expectApproxEqAbs(b, a, 0.00005 * mag); + } + } + } + } +} + +test "Metal EXL3 rows retain exact BF16 bits across full and independently resumed windows" { + const device = try mtl.Device.init(); + defer device.deinit(); + var p = try ex.Pipelines.load(device); + defer p.deinit(); + for ([_]ex.Decode{ .arithmetic, .hoisted_decode }) |decode| { + p.decode = decode; + const queue = try device.queue(); + defer queue.deinit(); + const w = try Owned.init(device, 6, .mul1, 0x71717); + defer w.deinit(); + const x = try input(device, rows); + defer x.deinit(); + const all = try device.buffer(rows * n * 2, mtl.ResourceOptions.shared); + defer all.deinit(); + const separate = try device.buffer(rows * n * 2, mtl.ResourceOptions.shared); + defer separate.deinit(); + const scratch = try device.buffer(try ex.scratchBytes(k, n, rows), mtl.ResourceOptions.shared); + defer scratch.deinit(); + const cb = queue.commandBuffer(); + const e = cb.compute(.concurrent); + try ex.encode(&p, e, w.matrix, .{ .buffer = x, .offset = 32 }, .{ .buffer = all }, .{ .buffer = scratch }, .{ .count = rows, .x_stride = k + 16, .y_stride = n }); + for (0..rows) |r| try ex.encode(&p, e, w.matrix, .{ .buffer = x, .offset = 32 + r * (k + 16) * 2 }, .{ .buffer = separate, .offset = r * n * 2 }, .{ .buffer = scratch }, .{ .count = 1, .x_stride = k + 16, .y_stride = n }); + try complete(cb, e); + try std.testing.expectEqualSlices(u16, all.slice(u16, rows * n), separate.slice(u16, rows * n)); + } +} + +test "GPU router descriptors select mixed-width experts with gather and scatter without host routing" { + const device = try mtl.Device.init(); + defer device.deinit(); + var p = try ex.Pipelines.load(device); + defer p.deinit(); + for ([_]ex.Decode{ .arithmetic, .hoisted_decode }) |decode| { + p.decode = decode; + const queue = try device.queue(); + defer queue.deinit(); + const w0 = try Owned.init(device, 3, .mul1, 11); + defer w0.deinit(); + const w1 = try Owned.init(device, 16, .mcg, 12); + defer w1.deinit(); + const descriptors = try device.buffer(2 * @sizeOf(ex.ExpertDescriptor), mtl.ResourceOptions.shared); + defer descriptors.deinit(); + descriptors.slice(ex.ExpertDescriptor, 2)[0] = try ex.descriptor(w0.matrix); + descriptors.slice(ex.ExpertDescriptor, 2)[1] = try ex.descriptor(w1.matrix); + const ids = try device.buffer(6 * 4, mtl.ResourceOptions.shared); + defer ids.deinit(); + @memcpy(ids.slice(u32, 6), &[_]u32{ 0, 1, 1, 0, 1, 0 }); + const gather = try device.buffer(6 * 4, mtl.ResourceOptions.shared); + defer gather.deinit(); + @memcpy(gather.slice(u32, 6), &[_]u32{ 2, 0, 1, 2, 1, 0 }); + const scatter = try device.buffer(6 * 4, mtl.ResourceOptions.shared); + defer scatter.deinit(); + @memcpy(scatter.slice(u32, 6), &[_]u32{ 5, 4, 3, 2, 1, 0 }); + const x = try input(device, 3); + defer x.deinit(); + const indexed = try device.buffer(6 * n * 4, mtl.ResourceOptions.shared); + defer indexed.deinit(); + const plain = try device.buffer(6 * n * 4, mtl.ResourceOptions.shared); + defer plain.deinit(); + const scratch = try device.buffer(try ex.scratchBytes(k, n, 6), mtl.ResourceOptions.shared); + defer scratch.deinit(); + const cb = queue.commandBuffer(); + const e = cb.compute(.concurrent); + ex.useMatrix(e, w0.matrix); + ex.useMatrix(e, w1.matrix); + try ex.encodeIndexed(&p, e, .{ .descriptors = .{ .buffer = descriptors }, .count = 2, .k = k, .n = n }, .{ .buffer = x, .offset = 32 }, .{ .buffer = indexed }, .{ .buffer = scratch }, .{ .count = 6, .x_stride = k + 16, .y_stride = n, .output = .f32, .routes = .{ .ids = .{ .buffer = ids }, .repeat = 2 }, .input_ids = .{ .buffer = gather }, .output_ids = .{ .buffer = scatter }, .input_rows = 3, .output_rows = 6 }); + for (0..6) |slot| { + const w = if (ids.slice(u32, 6)[slot] == 0) w0 else w1; + const source = gather.slice(u32, 6)[slot]; + const dest = scatter.slice(u32, 6)[slot]; + try ex.encode(&p, e, w.matrix, .{ .buffer = x, .offset = 32 + source * (k + 16) * 2 }, .{ .buffer = plain, .offset = dest * n * 4 }, .{ .buffer = scratch }, .{ .count = 1, .x_stride = k + 16, .y_stride = n, .output = .f32 }); + } + try complete(cb, e); + try std.testing.expectEqualSlices(u32, plain.slice(u32, 6 * n), indexed.slice(u32, 6 * n)); + } +} + +test "packed signs and GPU expert filtering preserve untouched output slots" { + const device = try mtl.Device.init(); + defer device.deinit(); + var p = try ex.Pipelines.load(device); + defer p.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + for ([_]ex.Decode{ .arithmetic, .hoisted_decode }) |decode| { + p.decode = decode; + const w = try Owned.init(device, 4, .inst3, 0xfeed); + defer w.deinit(); + const su = try device.buffer(k / 8, mtl.ResourceOptions.shared); + defer su.deinit(); + const sv = try device.buffer(n / 8, mtl.ResourceOptions.shared); + defer sv.deinit(); + @memset(su.slice(i16, k / 16), @bitCast(@as(u16, 0x9635))); + @memset(sv.slice(i16, n / 16), @bitCast(@as(u16, 0x4a17))); + ex.format.unpackSigns(su.slice(i16, k / 16), w.buffers[1].slice(f16, k + 8)[8..]); + ex.format.unpackSigns(sv.slice(i16, n / 16), w.buffers[2].slice(f16, n + 8)[8..]); + var signed = w.matrix; + signed.suh = .{ .buffer = su }; + signed.svh = .{ .buffer = sv }; + signed.suh_format = .packed_signs; + signed.svh_format = .packed_signs; + const ids = try device.buffer(6 * 4, mtl.ResourceOptions.shared); + defer ids.deinit(); + @memcpy(ids.slice(u32, 6), &[_]u32{ 2, 1, 1, 2, 2, 1 }); + const x = try input(device, 3); + defer x.deinit(); + const filtered = try device.buffer(6 * n * 2, mtl.ResourceOptions.shared); + defer filtered.deinit(); + const plain = try device.buffer(6 * n * 2, mtl.ResourceOptions.shared); + defer plain.deinit(); + @memset(filtered.slice(u16, 6 * n), 0x7f7f); + @memset(plain.slice(u16, 6 * n), 0x7f7f); + const scratch = try device.buffer(try ex.scratchBytes(k, n, 6), mtl.ResourceOptions.shared); + defer scratch.deinit(); + const cb = queue.commandBuffer(); + const e = cb.compute(.concurrent); + try ex.encode(&p, e, signed, .{ .buffer = x, .offset = 32 }, .{ .buffer = filtered }, .{ .buffer = scratch }, .{ .count = 6, .x_stride = k + 16, .y_stride = n, .routes = .{ .ids = .{ .buffer = ids }, .expert_id = 1, .repeat = 2 } }); + for (0..6) |slot| { + if (ids.slice(u32, 6)[slot] != 1) continue; + try ex.encode(&p, e, w.matrix, .{ .buffer = x, .offset = 32 + (slot / 2) * (k + 16) * 2 }, .{ .buffer = plain, .offset = slot * n * 2 }, .{ .buffer = scratch }, .{ .count = 1, .x_stride = k + 16, .y_stride = n }); + } + try complete(cb, e); + try std.testing.expectEqualSlices(u16, plain.slice(u16, 6 * n), filtered.slice(u16, 6 * n)); + } +} + +test "grouped mixed-width routes preserve every output slot across eight-row tails and both input maps" { + const groups = @import("expert_groups.zig"); + 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(); + var grouping = try groups.Pipelines.load(device); + defer grouping.deinit(); + var matrices: [8]Owned = undefined; + var loaded: usize = 0; + defer for (matrices[0..loaded]) |matrix| matrix.deinit(); + const widths = [_]u16{ 3, 16, 6 }; + const books = [_]ex.format.Codebook{ .mul1, .mcg, .inst3 }; + for (&matrices, 0..) |*matrix, i| { + matrix.* = try Owned.init(device, widths[i % 3], books[i % 3], 713 + i); + loaded += 1; + } + const descriptors = try device.buffer(8 * @sizeOf(ex.ExpertDescriptor), mtl.ResourceOptions.shared); + defer descriptors.deinit(); + for (matrices, descriptors.slice(ex.ExpertDescriptor, 8)) |matrix, *entry| + entry.* = try ex.descriptor(matrix.matrix); + const slots = 17 * 6; + const ids = try device.buffer(slots * 4, mtl.ResourceOptions.shared); + defer ids.deinit(); + for (ids.slice(u32, slots), 0..) |*id, i| id.* = @intCast((i / 6 + i % 6) % 8); + const plan = try groups.Plan.init(slots, 8); + const metadata = try device.buffer(plan.bytes, mtl.ResourceOptions.shared); + defer metadata.deinit(); + const grouped = groups.Scratch{ .buffer = metadata, .plan = plan }; + const x = try input(device, slots); + defer x.deinit(); + const old = try device.buffer(slots * n * 4 + 16, mtl.ResourceOptions.shared); + defer old.deinit(); + const candidate = try device.buffer(slots * n * 4 + 16, mtl.ResourceOptions.shared); + defer candidate.deinit(); + const scratch = try device.buffer(try ex.scratchBytes(k, n, plan.padded), mtl.ResourceOptions.shared); + defer scratch.deinit(); + for ([_]usize{ 6, 1 }) |repeat| { + @memset(old.slice(u32, slots * n + 4), 0x7fbebebe); + @memset(candidate.slice(u32, slots * n + 4), 0x7fbebebe); + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + for (matrices) |matrix| ex.useMatrix(e, matrix.matrix); + try groups.encode(e, &grouping, .{ .buffer = ids }, grouped, 6); + p.decode = .arithmetic; + try ex.encodeIndexed(&p, e, .{ .descriptors = .{ .buffer = descriptors }, .count = 8, .k = k, .n = n }, .{ .buffer = x, .offset = 32 }, .{ .buffer = old }, .{ .buffer = scratch }, .{ .count = slots, .x_stride = k + 16, .y_stride = n, .output = .f32, .routes = .{ .ids = .{ .buffer = ids }, .repeat = repeat } }); + p.decode = .hoisted_decode; + try ex.encodeIndexed(&p, e, .{ .descriptors = .{ .buffer = descriptors }, .count = 8, .k = k, .n = n }, .{ .buffer = x, .offset = 32 }, .{ .buffer = candidate }, .{ .buffer = scratch }, .{ .count = plan.padded, .x_stride = k + 16, .y_stride = n, .output = .f32, .routes = .{ .ids = grouped.view(.expert_ids) }, .grouped = true, .input_ids = grouped.view(if (repeat == 6) .token_ids else .slot_ids), .output_ids = grouped.view(.slot_ids), .input_rows = slots / repeat, .output_rows = slots }); + try complete(cb, e); + try std.testing.expectEqualSlices(u32, old.slice(u32, slots * n + 4), candidate.slice(u32, slots * n + 4)); + const expert_view = grouped.view(.expert_ids); + const token_view = grouped.view(.token_ids); + const slot_view = grouped.view(.slot_ids); + const experts: [*]const u32 = @ptrCast(@alignCast(metadata.contents() + expert_view.offset)); + const tokens: [*]const u32 = @ptrCast(@alignCast(metadata.contents() + token_view.offset)); + const original: [*]const u32 = @ptrCast(@alignCast(metadata.contents() + slot_view.offset)); + var seen: [slots]bool = @splat(false); + for (0..plan.padded) |row| { + if (experts[row] == std.math.maxInt(u32)) continue; + try std.testing.expect(original[row] < slots); + try std.testing.expect(!seen[original[row]]); + seen[original[row]] = true; + try std.testing.expectEqual(ids.slice(u32, slots)[original[row]], experts[row]); + try std.testing.expectEqual(original[row] / 6, tokens[row]); + try std.testing.expectEqual(experts[row / 8 * 8], experts[row]); + } + for (seen) |present| try std.testing.expect(present); + } +} + +test "rejected views leave an encoder reusable and successful output retains its guards" { + 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 w = try Owned.init(device, 6, .mul1, 81); + defer w.deinit(); + const x = try input(device, 1); + defer x.deinit(); + const y = try device.buffer(n * 2 + 32, mtl.ResourceOptions.shared); + defer y.deinit(); + @memset(y.slice(u16, n + 16), 0x7f7f); + const scratch_bytes = try ex.scratchBytes(k, n, 1); + const scratch = try device.buffer(scratch_bytes + 32, mtl.ResourceOptions.shared); + defer scratch.deinit(); + @memset(scratch.slice(u8, scratch.length()), 0xa5); + const cb = queue.commandBuffer(); + const e = cb.compute(.concurrent); + const xv = ex.View{ .buffer = x, .offset = 32 }; + const yv = ex.View{ .buffer = y, .offset = 16 }; + const sv = ex.View{ .buffer = scratch, .offset = 16 }; + const r = ex.Rows{ .count = 1, .x_stride = k, .y_stride = n }; + try std.testing.expectError(error.InvalidExl3View, ex.encode(&p, e, w.matrix, .{ .buffer = x, .offset = 33 }, yv, sv, r)); + try std.testing.expectError(error.InvalidExl3View, ex.encode(&p, e, w.matrix, xv, yv, .{ .buffer = scratch, .offset = 48 }, r)); + try std.testing.expectError(error.AliasedExl3View, ex.encode(&p, e, w.matrix, xv, xv, sv, r)); + var bad_rows = r; + bad_rows.x_stride = std.math.maxInt(usize); + bad_rows.count = 2; + try std.testing.expectError(error.Overflow, ex.encode(&p, e, w.matrix, xv, yv, sv, bad_rows)); + bad_rows = r; + bad_rows.input_ids = .{ .buffer = x }; + try std.testing.expectError(error.InvalidExl3Capacity, ex.encode(&p, e, w.matrix, xv, yv, sv, bad_rows)); + bad_rows.input_rows = 2; + try std.testing.expectError(error.InvalidExl3View, ex.encode(&p, e, w.matrix, xv, yv, sv, bad_rows)); + var bad_matrix = w.matrix; + bad_matrix.trellis.offset = bad_matrix.trellis.buffer.length(); + try std.testing.expectError(error.InvalidExl3View, ex.encode(&p, e, bad_matrix, xv, yv, sv, r)); + try ex.encode(&p, e, w.matrix, xv, yv, sv, r); + try complete(cb, e); + for (y.slice(u16, n + 16)[0..8]) |guard| try std.testing.expectEqual(@as(u16, 0x7f7f), guard); + for (y.slice(u16, n + 16)[n + 8 ..]) |guard| try std.testing.expectEqual(@as(u16, 0x7f7f), guard); + for (scratch.slice(u8, scratch.length())[0..16]) |guard| try std.testing.expectEqual(@as(u8, 0xa5), guard); + for (scratch.slice(u8, scratch.length())[scratch_bytes + 16 ..]) |guard| + try std.testing.expectEqual(@as(u8, 0xa5), guard); + var expected: [n]f32 = undefined; + reference(w, x.slice(u16, k + 32)[16..][0..k], &expected); + for (expected, y.slice(u16, n + 16)[8..][0..n]) |want, got| + try std.testing.expectApproxEqAbs(want, @as(f32, @bitCast(@as(u32, got) << 16)), 0.008 * @max(1, @abs(want))); +} 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 new file mode 100644 index 0000000000..72edee1343 --- /dev/null +++ b/zig/src/core/gpu_counters.zig @@ -0,0 +1,355 @@ +//! 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; + +pub const Class = @import("gpu_profile.zig").Class; +pub const class_count = @typeInfo(Class).@"enum".field_names.len; +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, +}; +pub const Aggregate = struct { calls: u64 = 0, gpu_seconds: f64 = 0 }; +const State = enum { idle, recording, resolved, read, broken }; + +/// 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, + + pub fn secondsPerTick(self: Calibration) !f64 { + if (self.after.cpu <= self.before.cpu or self.after.gpu <= self.before.gpu) return error.InvalidClockCalibration; + const cpu_delta: f64 = @floatFromInt(self.after.cpu - self.before.cpu); + const gpu_delta: f64 = @floatFromInt(self.after.gpu - self.before.gpu); + return cpu_delta / gpu_delta * 1e-9; + } + + 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 (end < begin) return error.ReversedCounterSample; + if (begin < self.before.gpu or end > self.after.gpu) return error.CounterOutsideCalibration; + const ticks: f64 = @floatFromInt(end - begin); + return ticks * try self.secondsPerTick(); + } +}; + +const Resources = struct { + device: mtl.Device, + samples: counters.SampleBuffer, + pass: ?counters.ComputePass, + fence: counters.Fence, + + fn init(device: mtl.Device, mode: counters.Mode, sample_count: usize) !Resources { + const samples = try counters.SampleBuffer.init(device, sample_count); + errdefer samples.deinit(); + const fence = try counters.Fence.init(device); + errdefer fence.deinit(); + const pass = if (mode == .stage) try counters.ComputePass.init() else null; + return .{ + .device = .{ .id = mtl.objc.retain(device.id) }, + .samples = samples, + .pass = pass, + .fence = fence, + }; + } + + fn deinit(self: Resources) void { + if (self.pass) |pass| pass.deinit(); + self.fence.deinit(); + self.samples.deinit(); + self.device.deinit(); + } +}; + +pub const Recorder = struct { + allocator: std.mem.Allocator, + resources: ?Resources = null, + hardware: counters.Capability = .{}, + spans: []Span = &.{}, + span_count: usize = 0, + active_span: ?usize = null, + command_buffer: ?mtl.CommandBuffer = null, + encoder_open: bool = false, + state: State = .idle, + clock_before: counters.ClockSample = .{}, + calibration: ?Calibration = null, + totals: [class_count]Aggregate = @splat(.{}), + + /// 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; + if (options.max_spans == 0 or options.max_spans > max_span_capacity) return error.InvalidSpanCapacity; + self.hardware = counters.query(device); + const mode = self.hardware.mode(); + if (mode == .unsupported) return error.UnsupportedGpuCounters; + self.spans = try allocator.alloc(Span, options.max_spans); + errdefer allocator.free(self.spans); + self.resources = try Resources.init(device, mode, options.max_spans * 2); + return self; + } + + pub fn enabled(self: *const Recorder) bool { + return self.resources != null; + } + + /// For disabled Recorders this is intentionally unqueried; use metal.counters.query(device) to inspect hardware. + pub fn capability(self: *const Recorder) counters.Capability { + return self.hardware; + } + + /// Bind the caller's uncommitted command buffer with an open compute encoder; do not commit inside this API. + pub fn beginCommand(self: *Recorder, cb: mtl.CommandBuffer) !void { + if (!self.enabled()) { + self.encoder_open = true; + return; + } + if (self.state != .idle and self.state != .read) return error.CounterRecordingInFlight; + if (cb.status() != .not_enqueued) return error.CounterCommandAlreadySubmitted; + self.command_buffer = cb.retain(); + self.clock_before = counters.sampleClocks(self.resources.?.device); + self.calibration = null; + self.span_count = 0; + self.active_span = null; + self.totals = @splat(.{}); + self.encoder_open = true; + self.state = .recording; + } + + /// 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; + if (self.span_count == self.spans.len) return error.CounterSpanCapacityExceeded; + const resources = self.resources.?; + const index = self.span_count; + if (self.hardware.mode() == .dispatch) { + try resources.samples.sample(encoder.*, index * 2); + } else { + resources.fence.update(encoder.*); + encoder.end(); + self.encoder_open = false; + encoder.* = resources.pass.?.compute(self.command_buffer.?, resources.samples, index * 2, index * 2 + 1) catch |err| { + self.state = .broken; + return err; + }; + self.encoder_open = true; + resources.fence.waitCompute(encoder.*); + } + self.spans[index] = .{ .kind = kind, .metadata = metadata }; + self.span_count += 1; + self.active_span = index; + } + + /// On stage-only devices the returned uninstrumented encoder still waits for all writes of the previous span. + pub fn endSpan(self: *Recorder, encoder: *mtl.ComputeEncoder) !void { + if (!self.enabled()) return; + if (self.state != .recording or !self.encoder_open) return error.CounterNotRecording; + const index = self.active_span orelse return error.NoActiveCounterSpan; + const resources = self.resources.?; + if (self.hardware.mode() == .dispatch) { + try resources.samples.sample(encoder.*, index * 2 + 1); + } else { + resources.fence.update(encoder.*); + encoder.end(); + self.encoder_open = false; + const cb = self.command_buffer.?; + const id = mtl.objc.msg(?mtl.objc.Id, cb.id, "computeCommandEncoderWithDispatchType:", .{ + @backingInt(mtl.DispatchType.serial), + }) orelse { + self.state = .broken; + return error.CounterEncoderCreationFailed; + }; + encoder.* = .{ .id = id }; + self.encoder_open = true; + resources.fence.waitCompute(encoder.*); + } + self.active_span = null; + } + + /// 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(); + self.encoder_open = false; + return; + } + if (self.state != .recording or !self.encoder_open) return error.CounterNotRecording; + if (self.active_span != null) return error.UnclosedCounterSpan; + encoder.end(); + self.encoder_open = false; + self.state = .resolved; + } + + /// 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; + } + if (self.encoder_open) { + encoder.end(); + self.encoder_open = false; + } + try self.discardAfterCompletion(); + } + + /// 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]; + if (self.state != .resolved) return error.CounterResultsNotEncoded; + const cb = self.command_buffer.?; + const status = cb.status(); + if (status == .@"error") return error.CounterCommandFailed; + if (status != .completed) return error.CounterCommandNotCompleted; + const resources = self.resources.?; + const calibration = Calibration{ + .before = self.clock_before, + .after = counters.sampleClocks(resources.device), + }; + if (self.span_count == 0) { + cb.release(); + self.command_buffer = null; + self.state = .read; + return &.{}; + } + const resolved = try resources.samples.resolveCompleted(self.span_count * 2); + defer resolved.deinit(); + const results = resolved.values(); + // Validate every sample before publishing partial totals; retrying cannot duplicate accumulated calls. + for (self.spans[0..self.span_count], 0..) |_, i| { + _ = try calibration.duration(results[i * 2].timestamp, results[i * 2 + 1].timestamp); + } + for (self.spans[0..self.span_count], 0..) |*span, i| { + span.begin_timestamp = results[i * 2].timestamp; + span.end_timestamp = results[i * 2 + 1].timestamp; + span.gpu_seconds = try calibration.duration(span.begin_timestamp, span.end_timestamp); + const aggregate = &self.totals[@backingInt(span.kind)]; + aggregate.calls += 1; + aggregate.gpu_seconds += span.gpu_seconds; + } + self.calibration = calibration; + cb.release(); + self.command_buffer = null; + self.state = .read; + return self.spans[0..self.span_count]; + } + + pub fn aggregates(self: *const Recorder) [class_count]Aggregate { + return self.totals; + } + + /// 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| { + switch (cb.status()) { + .not_enqueued, .completed, .@"error" => {}, + else => return error.CounterCommandNotCompleted, + } + cb.release(); + self.command_buffer = null; + } + self.state = .idle; + self.span_count = 0; + self.active_span = null; + self.totals = @splat(.{}); + self.calibration = null; + } + + /// Unique ownership: do not copy; cancel unsubmitted work or await completion before destruction. + pub fn deinit(self: *Recorder) !void { + try self.discardAfterCompletion(); + if (self.resources) |resources| { + resources.deinit(); + self.allocator.free(self.spans); + } + self.* = .{ .allocator = self.allocator }; + } +}; + +test "calibration converts GPU ticks rather than CPU span time" { + const calibration = Calibration{ .before = .{ .cpu = 1000, .gpu = 100 }, .after = .{ .cpu = 5000, .gpu = 300 } }; + try std.testing.expectApproxEqAbs(@as(f64, 1e-6), try calibration.duration(150, 200), 1e-15); + try std.testing.expectApproxEqAbs(@as(f64, 0), try calibration.duration(175, 175), 1e-15); + try std.testing.expectError(error.InvalidCounterSample, calibration.duration(counters.error_value, 200)); + try std.testing.expectError(error.ReversedCounterSample, calibration.duration(200, 150)); + try std.testing.expectError(error.CounterOutsideCalibration, calibration.duration(99, 200)); + const invalid = Calibration{ .before = .{ .cpu = 1000, .gpu = 100 }, .after = .{ .cpu = 1000, .gpu = 300 } }; + try std.testing.expectError(error.InvalidClockCalibration, invalid.secondsPerTick()); +} + +test "disabled recorder allocates nothing and never touches a device or encoder for spans" { + 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, 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 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)); + try std.testing.expectError(error.CounterOutsideCalibration, c.duration(150, 301)); + try std.testing.expectApproxEqAbs(@as(f64, 4e-6), try c.duration(100, 300), 1e-15); + const bad_gpu = Calibration{ .before = c.before, .after = .{ .cpu = 5000, .gpu = 99 } }; + const bad_cpu = Calibration{ .before = c.before, .after = .{ .cpu = 999, .gpu = 300 } }; + try std.testing.expectError(error.InvalidClockCalibration, bad_gpu.secondsPerTick()); + try std.testing.expectError(error.InvalidClockCalibration, bad_cpu.secondsPerTick()); +} + +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 })); +} + +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/core/lanes/stream.zig b/zig/src/core/lanes/stream.zig index d43d01b103..96c89af201 100644 --- a/zig/src/core/lanes/stream.zig +++ b/zig/src/core/lanes/stream.zig @@ -64,10 +64,12 @@ pub fn detectLoop(emitted: []const u32) ?u32 { return cycleAt(emitted, &.{}, emitted.len); } -/// A stop-string check over the emitted tokens (the server's StopPolicy decodes their tail). +/// Check emitted-token stop strings; decoder/allocation errors fail the host instead of missing a stop. pub const StopCheck = struct { ptr: *anyopaque, - check: *const fn (ptr: *anyopaque, emitted: []const u32) bool, + check: *const fn (ptr: *anyopaque, emitted: []const u32) Error!bool, + + pub const Error = error{ Template, OutOfMemory }; }; /// The host's answer to whether the stream's request was cancelled, asked between prompt chunks. @@ -352,7 +354,7 @@ pub const Stream = struct { } } const fire = if (s.loop_guard and s.think_open and s.think_end >= 0 and s.loop_period == null) detectLoop(s.emitted()) else null; - if (s.isEos(t) or s.stopped()) { + if (s.isEos(t) or try s.stopped()) { s.finished = true; s.reason = .stop; } else if (fire) |period| { @@ -377,7 +379,7 @@ pub const Stream = struct { return landed; } - fn stopped(s: *const Stream) bool { + fn stopped(s: *const Stream) StopCheck.Error!bool { const c = s.stop_check orelse return false; return c.check(c.ptr, s.emitted()); } @@ -397,6 +399,29 @@ test "commit stops at eos and length" { try std.testing.expect(s.finished and s.reason == .stop); } +test "stop callback errors propagate instead of becoming a miss or a normal stop" { + const Check = struct { + fn fail(ptr: *anyopaque, _: []const u32) StopCheck.Error!bool { + const failure: *StopCheck.Error = @ptrCast(@alignCast(ptr)); + return failure.*; + } + }; + const gpa = std.testing.allocator; + for ([_]StopCheck.Error{ error.OutOfMemory, error.Template }) |failure| { + var stored = failure; + var s = try Stream.init(gpa, .{ + .id = "stop-error", + .prompt = &.{1}, + .max_new = 3, + .stop_check = .{ .ptr = &stored, .check = Check.fail }, + }); + defer s.deinit(gpa); + try std.testing.expectError(failure, s.commit(gpa, &.{ 5, 6 }, &.{})); + try std.testing.expectEqualSlices(u32, &.{5}, s.emitted()); + try std.testing.expect(!s.finished and s.reason == .none); + } +} + test "a logprob stream lands a row with each token it lands, and refuses tokens without rows" { const gpa = std.testing.allocator; var s = try Stream.init(gpa, .{ .id = "a", .prompt = &.{ 1, 2 }, .max_new = 3, .eos = &.{9}, .logprobs = 0 }); diff --git a/zig/src/core/paced_read.zig b/zig/src/core/paced_read.zig index b23dd15c4a..54fea1ee0d 100644 --- a/zig/src/core/paced_read.zig +++ b/zig/src/core/paced_read.zig @@ -24,17 +24,25 @@ pub fn freeBytes() u64 { var pages: u32 = 0; var len: usize = @sizeOf(u32); if (std.c.sysctlbyname("vm.page_free_count", &pages, &len, null, 0) != 0) return std.math.maxInt(u64); + var speculative: u32 = 0; + var spec_len: usize = @sizeOf(u32); + // sysctl's free count excludes reclaimable speculative pages; Mach's free count includes them. + if (std.c.sysctlbyname("vm.page_speculative_count", &speculative, &spec_len, null, 0) != 0) speculative = 0; var size: u64 = 0; var len2: usize = @sizeOf(u64); if (std.c.sysctlbyname("hw.pagesize", &size, &len2, null, 0) != 0) size = alignment; - return @as(u64, pages) * size; + return (@as(u64, pages) + speculative) * size; } -/// Wait (up to `limit_ms`) while free memory is under the floor, so the kernel reclaims before the next read commits pages. +/// Wait under memory pressure and a low free-page floor; normal pressure permits the kernel to reclaim file cache. pub fn waitForRoom(limit_ms: u32) void { const want = floor(); var waited: u32 = 0; while (waited < limit_ms and freeBytes() < want) : (waited += 2) { + var pressure: c_int = 0; + var pressure_len: usize = @sizeOf(c_int); + if (std.c.sysctlbyname("kern.memorystatus_vm_pressure_level", &pressure, &pressure_len, null, 0) == 0 and + pressure == 1) return; // DISPATCH_MEMORYPRESSURE_NORMAL; cached pages are reclaimable, not exhausted RAM. const ts: std.c.timespec = .{ .sec = 0, .nsec = 2 * std.time.ns_per_ms }; _ = std.c.nanosleep(&ts, null); } diff --git a/zig/src/core/serial_batch_test.zig b/zig/src/core/serial_batch_test.zig index f619635719..670f206b9e 100644 --- a/zig/src/core/serial_batch_test.zig +++ b/zig/src/core/serial_batch_test.zig @@ -95,7 +95,7 @@ const Box = struct { return error.HostTimeout; } }; -fn stop(_: *anyopaque, tokens: []const u32) bool { +fn stop(_: *anyopaque, tokens: []const u32) api.Stop.Error!bool { return tokens.len == 4; } test "draft batch routing, bounded delivery and EOS preserve shared host events" { @@ -130,6 +130,33 @@ test "draft batch routing, bounded delivery and EOS preserve shared host events" } } +test "stop callback failure finishes the serial request with its actual error" { + const Check = struct { + fn fail(ptr: *anyopaque, _: []const u32) api.Stop.Error!bool { + const failure: *api.Stop.Error = @ptrCast(@alignCast(ptr)); + return failure.*; + } + }; + for ([_]api.Stop.Error{ error.OutOfMemory, error.Template }) |failure| { + var stored = failure; + var f = Fixture{}; + var b = Box{}; + const h = try serial.Host.init(std.testing.allocator, io, f.driver()); + defer h.deinit(); + const request = api.Request{ + .prompt = &.{1}, + .max_tokens = 20, + .stop = .{ .ctx = &stored, .check = Check.fail }, + }; + try h.engine().submit(7, &request, .{ .ctx = &b, .event = Box.event }); + try b.wait(); + try std.testing.expectEqual(api.Reason.failed, b.reason.?); + try std.testing.expectEqualStrings(@errorName(failure), b.message); + try std.testing.expectEqual(@as(usize, 1), b.count); + try std.testing.expectEqual(@as(usize, 0), f.batches); + } +} + test "draft cancellation publishes a complete returned round and resets the next request" { var f = Fixture{}; const h = try serial.Host.init(std.testing.allocator, io, f.driver()); diff --git a/zig/src/core/serial_host.zig b/zig/src/core/serial_host.zig index 7127ffc6b8..0f1c13ca3d 100644 --- a/zig/src/core/serial_host.zig +++ b/zig/src/core/serial_host.zig @@ -26,6 +26,8 @@ pub const Driver = struct { fed: ?*const fn (*anyopaque) u64 = null, /// The family's queue as the server's idle keepalive target (Metal); null keeps the ticker off. keepalive: ?api.keepalive.Target = null, + /// Immutable or synchronized backend memory counters; HTTP reads must not touch the GPU. + memory: ?api.MemorySource = null, }; const Job = struct { id: api.Id, request: *const api.Request, sink: api.Sink }; @@ -81,6 +83,7 @@ pub const Host = struct { const h = self(ptr); var i = h.driver.info; i.plain_only = h.driver.draft_batch == null; + i.thinking_budget = false; return i; } fn submit(ptr: *anyopaque, id: api.Id, request: *const api.Request, sink: api.Sink) api.SubmitError!void { @@ -128,8 +131,9 @@ pub const Host = struct { if (live and stream_tokens.len > 0) stream_tokens[0] = h.prompt_count + h.generated_count; out.* = .{ .running = @intFromBool(live), .waiting = @intCast(h.queued.items.len), .streams = if (live and stream_tokens.len > 0) 1 else 0, .generation_tokens = h.generated_count, .prefill_tokens_per_second = h.prefill_rate, .decode_tokens_per_second = h.decode_rate }; } - fn memory(_: *anyopaque, _: bool) ?api.Memory { - return null; + fn memory(ptr: *anyopaque, reset_peak: bool) ?api.Memory { + const source = self(ptr).driver.memory orelse return null; + return source.read(source.ctx, reset_peak); } fn emit(job: *Job, event: api.Event) void { job.sink.event(job.sink.ctx, job.id, &event); @@ -252,7 +256,8 @@ pub const Host = struct { h.lock(); h.generated_count += 1; h.unlock(); - const stop = std.mem.indexOfScalar(u32, request.eos, token) != null or (if (request.stop) |check| check.check(check.ctx, emitted.items) else false); + const stop = std.mem.indexOfScalar(u32, request.eos, token) != null or + (if (request.stop) |check| try check.check(check.ctx, emitted.items) else false); if (stop or i + 1 == request.max_tokens) { const ns = draw_start.durationTo(std.Io.Clock.awake.now(h.io)).toNanoseconds(); h.lock(); diff --git a/zig/src/core/template/ast.zig b/zig/src/core/template/ast.zig index c97c2665fd..7b353b0aaf 100644 --- a/zig/src/core/template/ast.zig +++ b/zig/src/core/template/ast.zig @@ -14,7 +14,7 @@ pub const Name = struct { pub const Kwarg = struct { name: []const u8, value: *Expr }; -pub const Args = struct { pos: []*Expr = &.{}, kw: []Kwarg = &.{} }; +pub const Args = struct { pos: []*Expr = &.{}, kw: []Kwarg = &.{}, dyn_kw: ?*Expr = null }; pub const CmpOp = enum { eq, ne, lt, lteq, gt, gteq, in, notin }; @@ -154,4 +154,5 @@ pub fn eachChild(comptime E: type, e: *const Expr, ctx: anytype, comptime visit: fn eachArg(comptime E: type, args: Args, ctx: anytype, comptime visit: fn (@TypeOf(ctx), *Expr) E!void) E!void { for (args.pos) |x| try visit(ctx, x); for (args.kw) |k| try visit(ctx, k.value); + if (args.dyn_kw) |x| try visit(ctx, x); } diff --git a/zig/src/core/template/eval.zig b/zig/src/core/template/eval.zig index 5186d0ddda..52409f6637 100644 --- a/zig/src/core/template/eval.zig +++ b/zig/src/core/template/eval.zig @@ -157,9 +157,26 @@ fn execFor(c: *Ctx, act: *Activation, f: *const ast.For) Error!void { fn args(c: *Ctx, act: *Activation, a: ast.Args) Error!struct { []Value, []Kw } { const pos = try c.rt.a.alloc(Value, a.pos.len); for (a.pos, pos) |e, *x| x.* = try eval(c, act, e); - const kw = try c.rt.a.alloc(Kw, a.kw.len); - for (a.kw, kw) |k, *x| x.* = .{ .name = k.name, .value = try eval(c, act, k.value) }; - return .{ pos, kw }; + if (a.dyn_kw == null) { + const kw = try c.rt.a.alloc(Kw, a.kw.len); + for (a.kw, kw) |k, *x| x.* = .{ .name = k.name, .value = try eval(c, act, k.value) }; + return .{ pos, kw }; + } + var kw: std.ArrayList(Kw) = .empty; + try kw.ensureTotalCapacity(c.rt.a, a.kw.len); + for (a.kw) |k| kw.appendAssumeCapacity(.{ .name = k.name, .value = try eval(c, act, k.value) }); + if (a.dyn_kw) |expr| { + const expanded = try eval(c, act, expr); + if (expanded != .dict) return c.rt.fail("argument after ** must be a mapping", .{}); + try kw.ensureUnusedCapacity(c.rt.a, expanded.dict.map.count()); + for (expanded.dict.map.keys(), expanded.dict.map.values()) |name, value| { + for (kw.items[0..a.kw.len]) |existing| { + if (std.mem.eql(u8, name, existing.name)) return c.rt.fail("got multiple values for keyword argument '{s}'", .{name}); + } + kw.appendAssumeCapacity(.{ .name = name, .value = value }); + } + } + return .{ pos, kw.items }; } fn applyFilters(c: *Ctx, act: *Activation, e: *const ast.Expr, hole: Value) Error!Value { diff --git a/zig/src/core/template/parser.zig b/zig/src/core/template/parser.zig index 78193dfe34..c43ec479a5 100644 --- a/zig/src/core/template/parser.zig +++ b/zig/src/core/template/parser.zig @@ -496,16 +496,24 @@ pub const Parser = struct { try p.expectOp(.lparen); var pos: std.ArrayList(*ast.Expr) = .empty; var kw: std.ArrayList(ast.Kwarg) = .empty; + var dyn_kw: ?*ast.Expr = null; var need_comma = false; while (!p.isOp(.rparen)) { if (need_comma) { try p.expectOp(.comma); if (p.isOp(.rparen)) break; } - if (p.isOp(.mul) or p.isOp(.pow)) return p.fail("*args and **kwargs calls are not supported", .{}); - if (p.cur().tag == .name and p.look().tag == .op and p.look().op == .assign) { + if (dyn_kw != null) return p.fail("invalid syntax for function call expression", .{}); + if (p.skipOp(.pow)) { + dyn_kw = try p.expression(true); + } else if (p.isOp(.mul)) { + return p.fail("*args calls are not supported", .{}); + } else if (p.cur().tag == .name and p.look().tag == .op and p.look().op == .assign) { const key = p.advance().text; _ = p.advance(); + for (kw.items) |existing| { + if (std.mem.eql(u8, existing.name, key)) return p.fail("keyword argument repeated: {s}", .{key}); + } try kw.append(p.a, .{ .name = key, .value = try p.expression(true) }); } else { if (kw.items.len > 0) return p.fail("invalid syntax for function call expression", .{}); @@ -514,7 +522,7 @@ pub const Parser = struct { need_comma = true; } try p.expectOp(.rparen); - return .{ .pos = pos.items, .kw = kw.items }; + return .{ .pos = pos.items, .kw = kw.items, .dyn_kw = dyn_kw }; } fn dottedName(p: *Parser) ![]const u8 { diff --git a/zig/src/core/template/tests.zig b/zig/src/core/template/tests.zig index 6d4d0ae1ab..1d11e6074a 100644 --- a/zig/src/core/template/tests.zig +++ b/zig/src/core/template/tests.zig @@ -119,3 +119,14 @@ test "syntax and unknown filters fail at compile time" { try expect("{% if %}", shared, .{ .text = null, .raised = null }); try expect("{% for x in y %}", shared, .{ .text = null, .raised = null }); } + +test "dynamic keyword mapping expands into dict constructors and macros" { + try expect("{{ dict({'a': 1}, **{'b': 2}) | tojson }}", "{}", .{ .text = "{\"a\": 1, \"b\": 2}", .raised = null }); + try expect("{% macro m(a, b) %}{{ a }}:{{ b }}{% endmacro %}{{ m(1, **{'b': 2}) }}", "{}", .{ .text = "1:2", .raised = null }); + try expect("{{ dict(x=1, **{'x': 2}) }}", "{}", .{ .text = null, .raised = null }); + try expect("{{ dict(**[1, 2]) }}", "{}", .{ .text = null, .raised = null }); + try expect("{{ dict(x=1, x=2) }}", "{}", .{ .text = null, .raised = null }); + try expect("{{ dict(**{'x': 1}, y=2) }}", "{}", .{ .text = null, .raised = null }); + try expect("{{ dict(**{'x': 1},) | tojson }}", "{}", .{ .text = "{\"x\": 1}", .raised = null }); + try expect("{{ dict(**{'x': value | nofilter}) }}", "{}", .{ .text = null, .raised = null }); +} diff --git a/zig/src/families/deepseek_v41/attention.zig b/zig/src/families/deepseek_v41/attention.zig new file mode 100644 index 0000000000..6d7cfe5da7 --- /dev/null +++ b/zig/src/families/deepseek_v41/attention.zig @@ -0,0 +1,258 @@ +//! Pinned DeepSeek V4.1 attention, resident EXL3 projections and causal packed-cache chunk execution. +const mtl = @import("metal"); +const exl3 = @import("../../core/exl3_projection.zig"); +const checkpoint = @import("../../core/checkpoint_metal.zig"); +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: exl3.Matrix, x: View, y: View, scratch: exl3.View, rows: u32, xs: u32, ys: u32, output: exl3.Output) !void { + try exl3.encode(p, e, m, 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, 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, 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, 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 = .f16 }); + 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, 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, 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/config.zig b/zig/src/families/deepseek_v41/config.zig new file mode 100644 index 0000000000..67658fb4c8 --- /dev/null +++ b/zig/src/families/deepseek_v41/config.zig @@ -0,0 +1,343 @@ +//! The admitted DeepSeek V4.1 EXL3 architecture and its source-owned global caches. +const std = @import("std"); + +pub const Rope = struct { + rope_type: []const u8, + factor: f32, + beta_fast: f32, + beta_slow: f32, + original_max_position_embeddings: u32, +}; + +pub const Text = struct { + vocab_size: u32, + hidden_size: u32, + moe_intermediate_size: u32, + num_hidden_layers: u32, + num_attention_heads: u32, + num_key_value_heads: u32, + head_dim: u32, + qk_rope_head_dim: u32, + q_lora_rank: u32, + o_lora_rank: u32, + o_groups: u32, + swiglu_limit: f32, + rms_norm_eps: f32, + max_position_embeddings: u32, + rope_theta: f32, + rope_scaling: Rope, + n_routed_experts: u32, + n_shared_experts: u32, + num_experts_per_tok: u32, + scoring_func: []const u8, + topk_method: []const u8, + norm_topk_prob: bool, + routed_scaling_factor: f32, + sliding_window: u32, + compress_ratios: []const u32, + compress_rope_theta: f32, + kv_source_layer_ids: []const u32, + index_source_layer_ids: []const u32, + index_n_heads: u32, + index_head_dim: u32, + index_topk: u32, + candidate_source_layer_id: u32, + candidate_topk_blocks: u32, + candidate_block_size: u32, + hc_mult: u32, + hc_sinkhorn_iters: u32, + hc_eps: f32, + engram_layer_ids: []const u32, + engram_num_embeddings: []const u64, + engram_max_ngram_size: u32, + engram_vocab_size: u32, + engram_n_heads: u32, + engram_head_dim: u32, + engram_pad_token_id: u32, + engram_compressed_vocab_size: u32, + num_nextn_predict_layers: u32, +}; + +pub const Config = struct { + model_type: []const u8, + text_config: Text, + quantization_config: struct { quant_method: []const u8, version: []const u8 }, + engram_table_dir: ?[]const u8 = null, + + pub fn load(arena: std.mem.Allocator, io: std.Io, dir: []const u8) !Config { + const path = try std.fs.path.join(arena, &.{ dir, "config.json" }); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, arena, .limited(1 << 20)); + return parse(arena, bytes); + } + + pub fn parse(arena: std.mem.Allocator, bytes: []const u8) !Config { + const root = try std.json.parseFromSliceLeaky(std.json.Value, arena, bytes, .{ .parse_numbers = false }); + const text = try field(root, "text_config"); + const rope = try field(text, "rope_scaling"); + const quant = try field(root, "quantization_config"); + const c: Config = .{ + .model_type = try string(try field(root, "model_type")), + .text_config = .{ + .vocab_size = try unsigned(u32, try field(text, "vocab_size")), + .hidden_size = try unsigned(u32, try field(text, "hidden_size")), + .moe_intermediate_size = try unsigned(u32, try field(text, "moe_intermediate_size")), + .num_hidden_layers = try unsigned(u32, try field(text, "num_hidden_layers")), + .num_attention_heads = try unsigned(u32, try field(text, "num_attention_heads")), + .num_key_value_heads = try unsigned(u32, try field(text, "num_key_value_heads")), + .head_dim = try unsigned(u32, try field(text, "head_dim")), + .qk_rope_head_dim = try unsigned(u32, try field(text, "qk_rope_head_dim")), + .q_lora_rank = try unsigned(u32, try field(text, "q_lora_rank")), + .o_lora_rank = try unsigned(u32, try field(text, "o_lora_rank")), + .o_groups = try unsigned(u32, try field(text, "o_groups")), + .swiglu_limit = try float(try field(text, "swiglu_limit")), + .rms_norm_eps = try float(try field(text, "rms_norm_eps")), + .max_position_embeddings = try unsigned(u32, try field(text, "max_position_embeddings")), + .rope_theta = try float(try field(text, "rope_theta")), + .rope_scaling = .{ + .rope_type = try string(try field(rope, "rope_type")), + .factor = try float(try field(rope, "factor")), + .beta_fast = try float(try field(rope, "beta_fast")), + .beta_slow = try float(try field(rope, "beta_slow")), + .original_max_position_embeddings = try unsigned(u32, try field(rope, "original_max_position_embeddings")), + }, + .n_routed_experts = try unsigned(u32, try field(text, "n_routed_experts")), + .n_shared_experts = try unsigned(u32, try field(text, "n_shared_experts")), + .num_experts_per_tok = try unsigned(u32, try field(text, "num_experts_per_tok")), + .scoring_func = try string(try field(text, "scoring_func")), + .topk_method = try string(try field(text, "topk_method")), + .norm_topk_prob = try boolean(try field(text, "norm_topk_prob")), + .routed_scaling_factor = try float(try field(text, "routed_scaling_factor")), + .sliding_window = try unsigned(u32, try field(text, "sliding_window")), + .compress_ratios = try integers(u32, arena, try field(text, "compress_ratios")), + .compress_rope_theta = try float(try field(text, "compress_rope_theta")), + .kv_source_layer_ids = try integers(u32, arena, try field(text, "kv_source_layer_ids")), + .index_source_layer_ids = try integers(u32, arena, try field(text, "index_source_layer_ids")), + .index_n_heads = try unsigned(u32, try field(text, "index_n_heads")), + .index_head_dim = try unsigned(u32, try field(text, "index_head_dim")), + .index_topk = try unsigned(u32, try field(text, "index_topk")), + .candidate_source_layer_id = try unsigned(u32, try field(text, "candidate_source_layer_id")), + .candidate_topk_blocks = try unsigned(u32, try field(text, "candidate_topk_blocks")), + .candidate_block_size = try unsigned(u32, try field(text, "candidate_block_size")), + .hc_mult = try unsigned(u32, try field(text, "hc_mult")), + .hc_sinkhorn_iters = try unsigned(u32, try field(text, "hc_sinkhorn_iters")), + .hc_eps = try float(try field(text, "hc_eps")), + .engram_layer_ids = try integers(u32, arena, try field(text, "engram_layer_ids")), + .engram_num_embeddings = try integers(u64, arena, try field(text, "engram_num_embeddings")), + .engram_max_ngram_size = try unsigned(u32, try field(text, "engram_max_ngram_size")), + .engram_vocab_size = try unsigned(u32, try field(text, "engram_vocab_size")), + .engram_n_heads = try unsigned(u32, try field(text, "engram_n_heads")), + .engram_head_dim = try unsigned(u32, try field(text, "engram_head_dim")), + .engram_pad_token_id = try unsigned(u32, try field(text, "engram_pad_token_id")), + .engram_compressed_vocab_size = try unsigned(u32, try field(text, "engram_compressed_vocab_size")), + .num_nextn_predict_layers = try unsigned(u32, try field(text, "num_nextn_predict_layers")), + }, + .quantization_config = .{ + .quant_method = try string(try field(quant, "quant_method")), + .version = try string(try field(quant, "version")), + }, + .engram_table_dir = if (root.object.get("engram_table_dir")) |v| + (if (v == .null) null else try string(v)) + else + null, + }; + try c.validate(); + return c; + } + + pub fn validate(c: Config) !void { + const t = c.text_config; + if (!std.mem.eql(u8, c.model_type, "deepseek_v41") or + !std.mem.eql(u8, c.quantization_config.quant_method, "exl3")) return error.UnsupportedCheckpoint; + if (t.hidden_size != 5120 or t.moe_intermediate_size != 2304 or t.num_hidden_layers != 40 or + t.num_attention_heads != 64 or t.num_key_value_heads != 1 or t.head_dim != 512 or + t.qk_rope_head_dim != 64 or t.q_lora_rank != 1280 or t.o_lora_rank != 1024 or t.o_groups != 8 or + t.vocab_size != 129280 or t.max_position_embeddings != 1048576) return error.UnsupportedArchitecture; + if (t.n_routed_experts != 384 or t.n_shared_experts != 1 or t.num_experts_per_tok != 6 or + !std.mem.eql(u8, t.scoring_func, "sqrtsoftplus") or !std.mem.eql(u8, t.topk_method, "noaux_tc") or + !t.norm_topk_prob or t.routed_scaling_factor != 1.5 or t.swiglu_limit != 10 or + t.sliding_window != 128 or t.hc_mult != 4 or t.hc_sinkhorn_iters != 20 or + t.index_n_heads != 32 or t.index_head_dim != 128 or t.index_topk != 512 or + t.candidate_source_layer_id != 20 or t.candidate_topk_blocks != 2048 or + t.candidate_block_size != 8) return error.UnsupportedArchitecture; + if (!std.mem.eql(u32, t.kv_source_layer_ids, &.{ 2, 8, 14, 20 }) or + !std.mem.eql(u32, t.index_source_layer_ids, &.{ 2, 8, 14, 20, 24, 28, 32, 36 }) or + !std.mem.eql(u32, t.engram_layer_ids, &.{ 1, 14 }) or t.compress_ratios.len != 43 or + !std.mem.eql(u64, t.engram_num_embeddings, &.{ 384006168, 384016682 })) return error.UnsupportedCacheLayout; + for (t.compress_ratios, 0..) |ratio, layer| { + const expected: u32 = if (layer < 2 or layer >= 40) 0 else if (layer < 20) 2 else 1; + if (ratio != expected) return error.UnsupportedCacheLayout; + } + if (t.engram_vocab_size != 16000000 or + t.engram_max_ngram_size != 4 or t.engram_n_heads != 8 or t.engram_head_dim != 256 or + t.engram_pad_token_id != 2 or t.engram_compressed_vocab_size != 99092 or + t.num_nextn_predict_layers != 3 or !std.mem.eql(u8, t.rope_scaling.rope_type, "yarn") or + t.rope_scaling.factor != 16 or t.rope_scaling.original_max_position_embeddings != 65536 or + !std.math.isFinite(t.rms_norm_eps) or t.rms_norm_eps <= 0 or + !std.math.isFinite(t.hc_eps) or t.hc_eps <= 0) return error.UnsupportedArchitecture; + if (!std.math.isFinite(t.rope_theta) or t.rope_theta <= 1 or + !std.math.isFinite(t.compress_rope_theta) or t.compress_rope_theta <= 1 or + !std.math.isFinite(t.rope_scaling.beta_fast) or t.rope_scaling.beta_fast <= 0 or + !std.math.isFinite(t.rope_scaling.beta_slow) or t.rope_scaling.beta_slow <= 0) + return error.UnsupportedArchitecture; + } + + pub fn context(c: Config, requested: u32) !u32 { + const limit = c.text_config.max_position_embeddings; + if (requested > limit) return error.ContextFull; + return if (requested == 0) limit else requested; + } +}; + +fn field(value: std.json.Value, name: []const u8) !std.json.Value { + if (value != .object) return error.UnexpectedToken; + return value.object.get(name) orelse error.MissingField; +} + +fn string(value: std.json.Value) ![]const u8 { + if (value != .string) return error.UnexpectedToken; + return value.string; +} + +fn boolean(value: std.json.Value) !bool { + if (value != .bool) return error.UnexpectedToken; + return value.bool; +} + +fn numberText(value: std.json.Value) ![]const u8 { + return switch (value) { + .number_string, .string => |s| s, + else => error.UnexpectedToken, + }; +} + +fn float(value: std.json.Value) !f32 { + const result = try std.fmt.parseFloat(f32, try numberText(value)); + if (!std.math.isFinite(result)) return error.Overflow; + return result; +} + +/// Decode integral decimal/exponent notation from digits exactly, without floating-point coercion. +fn unsigned(comptime T: type, value: std.json.Value) !T { + const text = try numberText(value); + if (text.len == 0) return error.InvalidCharacter; + if (std.fmt.parseInt(T, text, 10)) |n| return n else |_| {} + var start: usize = 0; + const negative = text[0] == '-'; + if (negative or text[0] == '+') start = 1; + const exp_pos = std.mem.indexOfAnyPos(u8, text, start, "eE") orelse text.len; + const exponent: i64 = if (exp_pos < text.len) + try std.fmt.parseInt(i32, text[exp_pos + 1 ..], 10) + else + 0; + var digits: usize = 0; + var fractional: usize = 0; + var dot = false; + var nonzero = false; + for (text[start..exp_pos]) |ch| { + if (ch == '.' and !dot) { + dot = true; + continue; + } + if (ch < '0' or ch > '9') return error.InvalidCharacter; + digits += 1; + if (dot) fractional += 1; + nonzero = nonzero or ch != '0'; + } + if (digits == 0) return error.InvalidCharacter; + if (!nonzero) return 0; + if (negative) return error.Overflow; + const scale = exponent - @as(i64, @intCast(fractional)); + const discarded: usize = if (scale < 0) @intCast(-scale) else 0; + if (discarded >= digits) return error.InvalidNumber; + const kept = digits - discarded; + var n: T = 0; + var index: usize = 0; + for (text[start..exp_pos]) |ch| { + if (ch == '.') continue; + if (index < kept) { + const digit: T = @intCast(ch - '0'); + if (n > (std.math.maxInt(T) - digit) / 10) return error.Overflow; + n = n * 10 + digit; + } else if (ch != '0') return error.InvalidNumber; + index += 1; + } + // Any nonzero ordinary integer overflows after at most 20 decimal shifts. + if (scale > 20) return error.Overflow; + var shift: i64 = 0; + while (shift < scale) : (shift += 1) { + if (n > std.math.maxInt(T) / 10) return error.Overflow; + n *= 10; + } + return n; +} + +fn integers(comptime T: type, arena: std.mem.Allocator, value: std.json.Value) ![]const T { + if (value != .array) return error.UnexpectedToken; + const result = try arena.alloc(T, value.array.items.len); + for (value.array.items, result) |item, *out| out.* = try unsigned(T, item); + return result; +} + +test "config integers retain exact unsigned boundaries and integral notation" { + try std.testing.expectEqual(@as(u32, 4294967295), try unsigned(u32, .{ .number_string = "4294967295" })); + try std.testing.expectEqual(@as(u64, 18446744073709551615), try unsigned(u64, .{ .number_string = "18446744073709551615.0" })); + try std.testing.expectEqual(@as(u32, 5120), try unsigned(u32, .{ .number_string = "5.120e3" })); + try std.testing.expectEqual(@as(u32, 5120), try unsigned(u32, .{ .string = "5120" })); + try std.testing.expectEqual(@as(u32, 0), try unsigned(u32, .{ .number_string = "-0.0" })); + try std.testing.expectError(error.Overflow, unsigned(u32, .{ .number_string = "4294967296" })); + try std.testing.expectError(error.Overflow, unsigned(u64, .{ .number_string = "18446744073709551616" })); + try std.testing.expectError(error.Overflow, unsigned(u64, .{ .number_string = "18446744073709551616.0" })); + try std.testing.expectError(error.InvalidNumber, unsigned(u64, .{ .number_string = "18446744073709551615.1" })); + try std.testing.expectError(error.InvalidNumber, unsigned(u32, .{ .number_string = "5120.0000000000000000001" })); + try std.testing.expectError(error.InvalidNumber, unsigned(u32, .{ .number_string = "1e-100" })); + try std.testing.expectError(error.Overflow, unsigned(u32, .{ .number_string = "-1.0" })); + try std.testing.expectError(error.Overflow, unsigned(u32, .{ .number_string = "1e100" })); +} + +test "config decoding rejects missing fields and wrong shapes without defaults" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + try std.testing.expectError(error.MissingField, Config.parse(arena.allocator(), "{}")); + try std.testing.expectError(error.UnexpectedToken, Config.parse(arena.allocator(), "[]")); + try std.testing.expectError(error.UnexpectedToken, Config.parse(arena.allocator(), "{\"text_config\":false}")); + try std.testing.expectError(error.UnexpectedToken, boolean(.{ .string = "true" })); + try std.testing.expectError(error.UnexpectedToken, unsigned(u32, .{ .bool = true })); + try std.testing.expectError(error.Overflow, float(.{ .number_string = "1e100" })); +} + +test "released architecture rejects inconsistent Engram buckets and nonfinite rotary parameters" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const checkpoint = + \\{"model_type":"deepseek_v41","quantization_config":{"quant_method":"exl3","version":"1"}, + \\"text_config":{"vocab_size":129280,"hidden_size":5120,"moe_intermediate_size":2304, + \\"num_hidden_layers":40,"num_attention_heads":64,"num_key_value_heads":1,"head_dim":512, + \\"qk_rope_head_dim":64,"q_lora_rank":1280,"o_lora_rank":1024,"o_groups":8, + \\"swiglu_limit":10,"rms_norm_eps":0.000001,"max_position_embeddings":1048576, + \\"rope_theta":10000,"rope_scaling":{"rope_type":"yarn","factor":16,"beta_fast":32, + \\"beta_slow":1,"original_max_position_embeddings":65536},"n_routed_experts":384, + \\"n_shared_experts":1,"num_experts_per_tok":6,"scoring_func":"sqrtsoftplus", + \\"topk_method":"noaux_tc","norm_topk_prob":true,"routed_scaling_factor":1.5, + \\"sliding_window":128,"compress_ratios":[0,0,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2, + \\1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,0,0,0],"compress_rope_theta":10000, + \\"kv_source_layer_ids":[2,8,14,20],"index_source_layer_ids":[2,8,14,20,24,28,32,36], + \\"index_n_heads":32,"index_head_dim":128,"index_topk":512,"candidate_source_layer_id":20, + \\"candidate_topk_blocks":2048,"candidate_block_size":8,"hc_mult":4,"hc_sinkhorn_iters":20, + \\"hc_eps":0.000001,"engram_layer_ids":[1,14],"engram_num_embeddings":[384006168,384016682], + \\"engram_max_ngram_size":4,"engram_vocab_size":16000000,"engram_n_heads":8,"engram_head_dim":256, + \\"engram_pad_token_id":2,"engram_compressed_vocab_size":99092,"num_nextn_predict_layers":3} + ; + const absent = try Config.parse(arena.allocator(), try std.fmt.allocPrint(arena.allocator(), "{s}}}", .{checkpoint})); + var invalid = absent; + invalid.text_config.engram_vocab_size = 1; + try std.testing.expectError(error.UnsupportedArchitecture, invalid.validate()); + invalid = absent; + invalid.text_config.engram_vocab_size = invalid.text_config.vocab_size; + try std.testing.expectError(error.UnsupportedArchitecture, invalid.validate()); + invalid = absent; + invalid.text_config.rope_theta = 1; + try std.testing.expectError(error.UnsupportedArchitecture, invalid.validate()); + invalid = absent; + invalid.text_config.rope_scaling.beta_slow = std.math.nan(f32); + try std.testing.expectError(error.UnsupportedArchitecture, invalid.validate()); +} 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..2ac845ef07 --- /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("kv_state.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..0f5a0287c2 --- /dev/null +++ b/zig/src/families/deepseek_v41/engram.zig @@ -0,0 +1,203 @@ +//! 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; +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, + 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, + 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), .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_hash.zig b/zig/src/families/deepseek_v41/engram_hash.zig new file mode 100644 index 0000000000..609aff7ca8 --- /dev/null +++ b/zig/src/families/deepseek_v41/engram_hash.zig @@ -0,0 +1,148 @@ +//! Exact DeepSeek AI inference/engram.py hashes; metadata from the released Flash configuration. +const std = @import("std"); + +pub const order = 4; +pub const heads = 8; +pub const columns = 24; +pub const dead: i32 = -1; +pub const raw_vocab = 129280; +pub const compressed_vocab = 99092; +pub const pad_id = 2; + +pub const Layer = extern struct { + layer_id: u32, + rows: u32, + multipliers: [order]u64, + primes: [columns]u32, + offsets: [columns]u32, +}; + +fn layer(id: u32, rows: u32, multipliers: [order]u64, primes: [columns]u32) Layer { + var offsets: [columns]u32 = undefined; + var offset: u32 = 0; + for (primes, 0..) |p, i| { + offsets[i] = offset; + offset += p; + } + std.debug.assert(offset == rows); + return .{ .layer_id = id, .rows = rows, .multipliers = multipliers, .primes = primes, .offsets = offsets }; +} + +pub const layers = [_]Layer{ + layer(1, 384006168, .{ 76632096046245, 4839876093313, 35959672319349, 73987337458391 }, .{ + 16000057, 16000079, 16000081, 16000097, 16000121, 16000129, 16000133, 16000183, + 16000189, 16000207, 16000211, 16000253, 16000277, 16000289, 16000307, 16000321, + 16000339, 16000381, 16000393, 16000399, 16000403, 16000409, 16000447, 16000463, + }), + layer(14, 384016682, .{ 67716810739261, 51510806800915, 30921347202721, 82619226485591 }, .{ + 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(bytes: []const u8, at: usize) u32 { + return std.mem.readInt(u32, bytes[at..][0..4], .little); +} + +/// Borrowed packed metadata; parent owns both the blob and its resident Metal upload. +pub const TokenMap = struct { + bytes: []const u8, + + pub fn parse(blob: []const u8) !TokenMap { + const header = 96; + const layer_size = 232; + const map_offset = header + layers.len * layer_size; + if (blob.len != map_offset + raw_vocab * 4 or !std.mem.eql(u8, blob[0..8], "DSENGR01")) + return error.InvalidEngramMetadata; + const expected = [_]u32{ raw_vocab, compressed_vocab, pad_id, order, heads, layers.len }; + for (expected, 0..) |value, i| if (word(blob, 8 + i * 4) != value) return error.InvalidEngramMetadata; + // Reject a tokenizer which happened to have the right size but normalized differently. + const tokenizer_sha = [_]u8{ + 0xc9, 0x0d, 0xfa, 0x01, 0x24, 0x9d, 0xb1, 0xbe, 0x42, 0x45, 0x78, 0x0a, 0x05, 0x2e, 0xde, 0x75, + 0x2e, 0x13, 0x61, 0xc6, 0x12, 0xac, 0x6d, 0x08, 0xe2, 0xbd, 0xad, 0xa7, 0xd5, 0x99, 0x47, 0x6b, + }; + if (!std.mem.eql(u8, blob[32..64], &tokenizer_sha)) return error.InvalidEngramTokenizer; + for (layers, 0..) |meta, i| { + const at = header + i * layer_size; + if (word(blob, at) != meta.layer_id or word(blob, at + 4) != meta.rows) return error.InvalidEngramMetadata; + for (meta.multipliers, 0..) |v, j| { + if (std.mem.readInt(u64, blob[at + 8 + j * 8 ..][0..8], .little) != v) + return error.InvalidEngramMetadata; + } + for (meta.primes, 0..) |v, j| if (word(blob, at + 40 + j * 4) != v) return error.InvalidEngramMetadata; + for (meta.offsets, 0..) |v, j| if (word(blob, at + 136 + j * 4) != v) return error.InvalidEngramMetadata; + } + const result = TokenMap{ .bytes = blob[map_offset..] }; + for (0..raw_vocab) |id| if (word(result.bytes, id * 4) >= compressed_vocab) return error.InvalidEngramTokenMap; + if (try result.get(2) != pad_id) return error.InvalidEngramTokenMap; + const map_sha = [_]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(result.bytes, &digest, .{}); + if (!std.mem.eql(u8, &digest, &map_sha)) return error.InvalidEngramTokenMap; + return result; + } + + pub fn get(self: TokenMap, token: u32) !u32 { + if (token >= self.bytes.len / 4) return error.InvalidEngramToken; + return word(self.bytes, @as(usize, token) * 4); + } +}; + +/// Only the previous three positions are retained, including DEAD positions; newest is first. +pub const History = struct { + tail: [order - 1]i32 = .{ dead, dead, dead }, + position: u64 = 0, + + pub fn reset(self: *History) void { + self.* = .{}; + } + + pub fn snapshot(self: History) History { + return self; + } + + /// Restore before replaying the accepted prefix of a speculative chunk. + pub fn rollback(self: *History, saved: History) void { + self.* = saved; + } + + pub fn push(self: *History, compressed: i32, meta: Layer) ![columns]u32 { + if (compressed < dead or compressed >= compressed_vocab) return error.InvalidEngramCompressedId; + var tokens: [order]u32 = @splat(pad_id); + var blocked = compressed == dead; + if (!blocked) tokens[0] = @intCast(compressed); + for (0..order - 1) |i| { + blocked = blocked or self.position <= i or self.tail[i] == dead; + if (!blocked) tokens[i + 1] = @intCast(self.tail[i]); + } + var result: [columns]u32 = undefined; + var rolling = @as(u64, tokens[0]) * meta.multipliers[0]; + for (1..order) |shift| { + rolling ^= @as(u64, tokens[shift]) * meta.multipliers[shift]; + for (0..heads) |head| { + const col = (shift - 1) * heads + head; + result[col] = @as(u32, @intCast(rolling % meta.primes[col])) + meta.offsets[col]; + } + } + self.tail[2] = self.tail[1]; + self.tail[1] = self.tail[0]; + self.tail[0] = compressed; + self.position += 1; + return result; + } +}; + +/// Single sequence, token mask bytes (0=DEAD), output [tokens,24], no allocations or lost chunk history. +pub fn hashChunk(history: *History, map: TokenMap, meta: Layer, ids: []const u32, mask: ?[]const u8, out: [][columns]u32) !void { + if (out.len != ids.len or (mask != null and mask.?.len != ids.len)) return error.InvalidEngramChunk; + // Validate the entire chunk before mutating history. + for (ids) |token| _ = try map.get(token); + for (ids, 0..) |token, i| { + const compressed: i32 = if (mask != null and mask.?[i] == 0) dead else @intCast(try map.get(token)); + out[i] = try history.push(compressed, meta); + } +} 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..ac9c4839b9 --- /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("kv_state.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_state.zig b/zig/src/families/deepseek_v41/kv_state.zig new file mode 100644 index 0000000000..f2b70175a7 --- /dev/null +++ b/zig/src/families/deepseek_v41/kv_state.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/kv_test.zig b/zig/src/families/deepseek_v41/kv_test.zig new file mode 100644 index 0000000000..0d1ae738bc --- /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("kv_state.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/load.zig b/zig/src/families/deepseek_v41/load.zig new file mode 100644 index 0000000000..cc6bdc19fe --- /dev/null +++ b/zig/src/families/deepseek_v41/load.zig @@ -0,0 +1,242 @@ +//! 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.zig"); +const cp = @import("../../core/checkpoint_metal.zig"); +const paced = @import("../../core/paced_read.zig"); + +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); + for (ordered) |file| try self.add(try std.fs.path.join(a, &.{ dir, file }), null); + 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); + 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) !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| { + 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); + 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); + residency.commit(); + residency.requestResidency(); + var copied: usize = 0; + while (copied < bytes) { + const count = @min(bytes - copied, stage.len - paced.alignment); + try paced.read(fd, buffer.contents()[copied..][0..count], shard.data_start + shard.begin + copied, stage); + 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."); + 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.")); +} + +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.")); + 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/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/model.zig b/zig/src/families/deepseek_v41/model.zig new file mode 100644 index 0000000000..1889e00913 --- /dev/null +++ b/zig/src/families/deepseek_v41/model.zig @@ -0,0 +1,363 @@ +//! 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.zig"); +const exl3 = @import("../../core/exl3_projection.zig"); +const config = @import("config.zig"); +const load = @import("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.zig"); + +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, .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 = .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 = .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/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..fe3a85b5fc --- /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.zig"); + +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/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..10afd22356 --- /dev/null +++ b/zig/src/families/deepseek_v41/weights.zig @@ -0,0 +1,273 @@ +//! 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.zig"); +const st = @import("../../core/safetensors.zig"); +const exl3 = @import("../../core/exl3_projection.zig"); +const load = @import("load.zig"); + +pub const Tensor = cp.Tensor; +pub const Matrix = exl3.Matrix; +pub const Hc = struct { projection: Tensor, scale: Tensor, base: Tensor }; +pub const Compressor = struct { wkv: Matrix, wgate: ?Matrix, norm: Tensor }; +pub const Indexer = struct { wq_b: Matrix, wk: ?Matrix, 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); +} + +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 }; +} + +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})) or + has(source, try named(&name, "{s}.weight", .{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 matrix(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(Matrix, undefined) else null; + if (layer < 20) try matrix(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 matrix(binding, source, try named(&prefix, "layers.{d}.attn.indexer.wq_b", .{layer}), 1280, 4096, if (binding) &indexer.?.wq_b else null); + try plain(binding, source, try named(&name, "layers.{d}.attn.indexer.weights_proj.weight", .{layer}), .f16, &.{ 32, 5120 }, if (binding) &indexer.?.weights_proj else null); + if (binding) { + indexer.?.wk = if (owner) @as(Matrix, undefined) else null; + indexer.?.k_norm = if (owner) @as(Tensor, undefined) else null; + } + if (owner) { + try matrix(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; + try plain(binding, source, try named(&name, "layers.{d}.ffn.gate.weight", .{layer}), .f16, &.{ 384, 5120 }, if (binding) &moe.?.gate else null); + try plain(binding, source, try named(&name, "layers.{d}.ffn.gate.bias", .{layer}), .f16, &.{384}, if (binding) &moe.?.bias else null); + 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 plain(binding, source, try named(&name, "layers.{d}.engram.q_weight", .{layer}), .f32, &.{ 4, 5120 }, if (binding) &e.?.q_weight else null); + try plain(binding, source, try named(&name, "layers.{d}.engram.k_weight", .{layer}), .f32, &.{ 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..fcef57cd6e --- /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.zig"); +const View = @import("view.zig").View; +const hash = @import("engram_hash.zig"); +const groups = @import("../../core/expert_groups.zig"); + +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 new file mode 100644 index 0000000000..0bbd65f186 --- /dev/null +++ b/zig/src/metal/counters.zig @@ -0,0 +1,207 @@ +//! 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"); +const device = @import("device.zig"); +const command = @import("command.zig"); + +extern "c" var MTLCommonCounterSetTimestamp: objc.Id; +extern "c" var MTLCommonCounterTimestamp: objc.Id; + +pub const error_value: u64 = std.math.maxInt(u64); +pub const max_samples: usize = 131072; +const max_counter_sets: usize = 64; +pub const TimestampResult = extern struct { timestamp: u64 }; +pub const SamplingPoint = enum(usize) { stage = 0, draw = 1, dispatch = 2, tile_dispatch = 3, blit = 4 }; +pub const Mode = enum { unsupported, dispatch, stage }; +pub const ClockSample = struct { cpu: u64 = 0, gpu: u64 = 0 }; + +pub const Capability = struct { + timestamp_set: bool = false, + timestamp_counter: bool = false, + clock_calibration: bool = false, + stage_boundary: bool = false, + dispatch_boundary: bool = false, + + pub fn mode(self: Capability) Mode { + if (!self.timestamp_set or !self.timestamp_counter or !self.clock_calibration) return .unsupported; + if (self.dispatch_boundary) return .dispatch; + if (self.stage_boundary) return .stage; + return .unsupported; + } +}; + +fn timestampSet(dev: device.Device) ?objc.Id { + const sets = objc.msg(?objc.Id, dev.id, "counterSets", .{}) orelse return null; + const count = objc.msg(usize, sets, "count", .{}); + if (count > max_counter_sets) return null; + for (0..count) |i| { + const set = objc.msg(objc.Id, sets, "objectAtIndex:", .{i}); + const name = objc.msg(objc.Id, set, "name", .{}); + if (objc.msg(bool, name, "isEqualToString:", .{MTLCommonCounterSetTimestamp})) return set; + } + return null; +} + +fn hasTimestamp(set: objc.Id) bool { + const counters = objc.msg(objc.Id, set, "counters", .{}); + const count = objc.msg(usize, counters, "count", .{}); + if (count > max_counter_sets) return false; + for (0..count) |i| { + const counter = objc.msg(objc.Id, counters, "objectAtIndex:", .{i}); + const name = objc.msg(objc.Id, counter, "name", .{}); + if (objc.msg(bool, name, "isEqualToString:", .{MTLCommonCounterTimestamp})) return true; + } + return false; +} + +/// Device queries, never architecture-name assumptions; no GPU allocations or command submissions. +pub fn query(dev: device.Device) Capability { + var result: Capability = .{}; + if (!dev.responds("counterSets") or !dev.responds("newCounterSampleBufferWithDescriptor:error:")) return result; + if (timestampSet(dev)) |set| { + result.timestamp_set = true; + result.timestamp_counter = hasTimestamp(set); + } + result.clock_calibration = dev.responds("sampleTimestamps:gpuTimestamp:"); + if (dev.responds("supportsCounterSampling:")) { + result.stage_boundary = objc.msg(bool, dev.id, "supportsCounterSampling:", .{@backingInt(SamplingPoint.stage)}); + result.dispatch_boundary = objc.msg(bool, dev.id, "supportsCounterSampling:", .{@backingInt(SamplingPoint.dispatch)}); + } + return result; +} + +/// Apple's CPU clock samples are nanoseconds; GPU ticks require two-point calibration, not a guessed tick rate. +pub fn sampleClocks(dev: device.Device) ClockSample { + var result: ClockSample = .{}; + objc.msg(void, dev.id, "sampleTimestamps:gpuTimestamp:", .{ &result.cpu, &result.gpu }); + return result; +} + +pub const SampleBuffer = struct { + id: objc.Id, + count: usize, + + /// Shared opaque samples support CPU resolution after command completion, including stage-only devices. + pub fn init(dev: device.Device, count: usize) !SampleBuffer { + if (count == 0 or count > max_samples) return error.InvalidSampleCount; + const capability = query(dev); + if (capability.mode() == .unsupported) return error.UnsupportedGpuCounters; + const allocated = objc.msg(?objc.Id, objc.class("MTLCounterSampleBufferDescriptor"), "alloc", .{}) orelse + return error.CounterSampleBufferCreationFailed; + const desc = objc.msg(?objc.Id, allocated, "init", .{}) orelse return error.CounterSampleBufferCreationFailed; + defer objc.release(desc); + objc.msg(void, desc, "setCounterSet:", .{timestampSet(dev).?}); + objc.msg(void, desc, "setStorageMode:", .{@as(usize, 0)}); + objc.msg(void, desc, "setSampleCount:", .{count}); + var err: ?objc.Id = null; + const id = objc.msg(?objc.Id, dev.id, "newCounterSampleBufferWithDescriptor:error:", .{ desc, &err }) orelse { + std.log.err("GPU counter sample buffer: {s}", .{objc.errorText(err)}); + return error.CounterSampleBufferCreationFailed; + }; + return .{ .id = id, .count = count }; + } + + pub fn sample(self: SampleBuffer, encoder: command.ComputeEncoder, index: usize) !void { + if (index >= self.count) return error.CounterSampleIndexOutOfBounds; + objc.msg(void, encoder.id, "sampleCountersInBuffer:atSampleIndex:withBarrier:", .{ self.id, index, true }); + } + + /// 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:", .{ + types.Range{ .location = 0, .length = count }, + }) orelse return error.CounterResolveFailed; + if (objc.msg(usize, data, "length", .{}) != count * @sizeOf(TimestampResult)) + return error.InvalidCounterResultLength; + return .{ .id = objc.retain(data), .count = count }; + } + + pub fn deinit(self: SampleBuffer) void { + objc.release(self.id); + } +}; + +pub const Resolved = struct { + id: objc.Id, + count: usize, + + pub fn values(self: Resolved) []const TimestampResult { + const bytes = objc.msg([*]const u8, self.id, "bytes", .{}); + const values_: [*]const TimestampResult = @ptrCast(@alignCast(bytes)); + return values_[0..self.count]; + } + + pub fn deinit(self: Resolved) void { + objc.release(self.id); + } +}; + +pub const ComputePass = struct { + id: objc.Id, + attachment: objc.Id, + + /// Reusable descriptor; Metal snapshots it when creating an encoder. All timing encoders are serial. + pub fn init() !ComputePass { + const allocated = objc.msg(?objc.Id, objc.class("MTLComputePassDescriptor"), "alloc", .{}) orelse + return error.CounterPassCreationFailed; + const id = objc.msg(?objc.Id, allocated, "init", .{}) orelse return error.CounterPassCreationFailed; + objc.msg(void, id, "setDispatchType:", .{@backingInt(types.DispatchType.serial)}); + const array = objc.msg(objc.Id, id, "sampleBufferAttachments", .{}); + const attachment = objc.msg(objc.Id, array, "objectAtIndexedSubscript:", .{@as(usize, 0)}); + return .{ .id = id, .attachment = attachment }; + } + + pub fn compute(self: ComputePass, cb: command.CommandBuffer, samples: SampleBuffer, start: usize, end: usize) !command.ComputeEncoder { + if (start >= samples.count or end >= samples.count) return error.CounterSampleIndexOutOfBounds; + objc.msg(void, self.attachment, "setSampleBuffer:", .{samples.id}); + objc.msg(void, self.attachment, "setStartOfEncoderSampleIndex:", .{start}); + objc.msg(void, self.attachment, "setEndOfEncoderSampleIndex:", .{end}); + return .{ .id = objc.msg(?objc.Id, cb.id, "computeCommandEncoderWithDescriptor:", .{self.id}) orelse + return error.CounterEncoderCreationFailed }; + } + + pub fn deinit(self: ComputePass) void { + objc.release(self.id); + } +}; + +pub const Fence = struct { + id: objc.Id, + + pub fn init(dev: device.Device) !Fence { + return .{ .id = objc.msg(?objc.Id, dev.id, "newFence", .{}) orelse return error.CounterFenceCreationFailed }; + } + + pub fn update(self: Fence, encoder: command.ComputeEncoder) void { + objc.msg(void, encoder.id, "updateFence:", .{self.id}); + } + + pub fn waitCompute(self: Fence, encoder: command.ComputeEncoder) void { + objc.msg(void, encoder.id, "waitForFence:", .{self.id}); + } + + pub fn deinit(self: Fence) void { + objc.release(self.id); + } +}; + +test "counter capability prefers dispatch and refuses missing timestamp data" { + var capability = Capability{ .timestamp_set = true, .timestamp_counter = true, .clock_calibration = true }; + try std.testing.expectEqual(Mode.unsupported, capability.mode()); + capability.stage_boundary = true; + try std.testing.expectEqual(Mode.stage, capability.mode()); + capability.dispatch_boundary = true; + try std.testing.expectEqual(Mode.dispatch, capability.mode()); + capability.timestamp_counter = false; + try std.testing.expectEqual(Mode.unsupported, capability.mode()); +} + +test "sample buffer capacity rejects invalid counts before device access" { + try std.testing.expectError(error.InvalidSampleCount, SampleBuffer.init(undefined, 0)); + try std.testing.expectError(error.InvalidSampleCount, SampleBuffer.init(undefined, max_samples + 1)); + const samples = SampleBuffer{ .id = undefined, .count = 2 }; + try std.testing.expectError(error.CounterSampleIndexOutOfBounds, samples.sample(undefined, 2)); + try std.testing.expectError(error.InvalidSampleCount, samples.resolveCompleted(3)); +} diff --git a/zig/src/metal/metal.zig b/zig/src/metal/metal.zig index a8e47f78e0..c7a6c835f9 100644 --- a/zig/src/metal/metal.zig +++ b/zig/src/metal/metal.zig @@ -26,6 +26,7 @@ pub const ResidencySet = @import("residency.zig").ResidencySet; pub const keepalive = @import("keepalive.zig"); pub const clock = @import("clock.zig"); pub const ane = @import("ane.zig"); +pub const counters = @import("counters.zig"); test { std.testing.refAllDecls(@This()); diff --git a/zig/src/native/deepseek_v41_host.zig b/zig/src/native/deepseek_v41_host.zig new file mode 100644 index 0000000000..c0962aac31 --- /dev/null +++ b/zig/src/native/deepseek_v41_host.zig @@ -0,0 +1,140 @@ +//! V4.1's resident model behind the existing serialized request worker. +const std = @import("std"); +const mtl = @import("metal"); +const api = @import("engine_api"); +const tf = @import("tensorfold"); +const ds = tf.deepseek_v41; +const Serial = api.serial_host; +const Allocator = std.mem.Allocator; + +pub const Host = struct { + gpa: Allocator, + model: *ds.Model, + serial: *Serial.Host, + sampler_storage: []u8, + memory_info: api.Memory, + warm: mtl.keepalive.Target, + + fn self(ptr: *anyopaque) *Host { + return @ptrCast(@alignCast(ptr)); + } + fn reset(ptr: *anyopaque) !void { + try self(ptr).model.reset(); + } + fn chunk(ptr: *anyopaque, ids: []const u32, last: bool) !void { + _ = try self(ptr).model.forward(ids, if (last) .last else .none); + } + fn advance(ptr: *anyopaque, token: u32) !void { + _ = try self(ptr).model.forward(&.{token}, .last); + } + fn fed(ptr: *anyopaque) u64 { + return self(ptr).model.cache.position; + } + fn draw(ptr: *anyopaque, sampling: ?api.Sampling) !u32 { + const h = self(ptr); + if (!h.model.valid or h.model.logits_rows != 1) return error.RoundNotReady; + const row = h.model.workspace.logits(1); + if (sampling) |s| if (s.temperature > 0) { + for (row) |value| if (!std.math.isFinite(value)) return error.NonfiniteLogits; + var arena = std.heap.FixedBufferAllocator.init(h.sampler_storage); + return Serial.choose(arena.allocator(), .{ .f32 = row }, h.model.cache.position, s); + }; + var best: u32 = 0; + var maximum = -std.math.inf(f32); + for (row, 0..) |value, id| { + if (!std.math.isFinite(value)) return error.NonfiniteLogits; + if (value > maximum) { + maximum = value; + best = @intCast(id); + } + } + return best; + } + fn memory(ptr: ?*anyopaque, _: bool) ?api.Memory { + const h: *Host = @ptrCast(@alignCast(ptr.?)); + return h.memory_info; + } +}; + +pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !?api.Opened { + if (o.speed_up != null or o.slide or o.learn != null or + (o.prompt_cache_gib != null and o.prompt_cache_gib.? != 0) or o.prompt_cache_over_cap) + { + problem.* = "V4.1 resident serving does not support speed-up, Sliding Weights or retained prompt-cache flags"; + return null; + } + if (o.drafter != null or o.device != null or o.segments != null or + o.mtp_drafts != null or o.mtp_confidence != null or + (o.keep != null and o.keep.? != 0) or o.tp != 1 or o.rank != 0 or + o.master.len != 0 or o.policy.len != 0) + { + problem.* = "experimental V4.1 serving does not support drafter, MTP, cache slots, device, TP or policy flags"; + return null; + } + if (o.lanes_fixed and o.lanes != 1) { + problem.* = "V4.1's plain resident worker serves one active sequence; use --parallel 1 or auto"; + return null; + } + const requested = o.context orelse 0; + if ((o.context != null and requested <= 0) or requested > 1048576) { + problem.* = "DeepSeek V4.1 context must be from 1 to 1,048,576 tokens (omitted: full model limit)"; + return null; + } + const h = try gpa.create(Host); + errdefer gpa.destroy(h); + const model = try gpa.create(ds.Model); + errdefer gpa.destroy(model); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const token_map = try std.fs.path.join(a, &.{ o.dir, "engram-token-map.bin" }); + model.* = ds.Model.open(gpa, io, .{ .dir = o.dir, .table_dir = o.engram_table_dir, .token_map = token_map, .context = @intCast(requested) }) catch |err| { + problem.* = try std.fmt.allocPrint(a, "the full-resident DeepSeek V4.1 engine cannot load ({s})", .{@errorName(err)}); + gpa.destroy(model); + gpa.destroy(h); + return null; + }; + errdefer model.deinit(); + const sampler_storage = try gpa.alloc(u8, 8 << 20); + errdefer gpa.free(sampler_storage); + const bytes = model.admitted_bytes; + h.* = .{ .gpa = gpa, .model = model, .serial = undefined, .sampler_storage = sampler_storage, .memory_info = .{ .active = bytes, .peak = bytes, .pinned = model.residency.allocatedSize() }, .warm = .{ .queue = model.queue, .sets = @as(*const [1]mtl.ResidencySet, &model.residency) } }; + const startup = try std.fmt.allocPrint(a, "EXPERIMENTAL / UNQUALIFIED DeepSeek V4.1: whole-model accuracy gate has not passed; " ++ + "all weights resident ({d} bytes), packed caches + bounded scratch ({d} bytes), context {d}; no SSD lookups", .{ model.checkpoint.residentBytes(), bytes - model.checkpoint.residentBytes(), model.cache.plan.capacity }); + h.serial = try Serial.Host.init(gpa, io, .{ + .ctx = h, + .info = .{ .name = "deepseek-v41-metal-exl3-resident", .lanes = 1, .context_window = model.cache.plan.capacity, .prefill_step = model.workspace.plan.rows, .plain_only = true, .startup = startup }, + .reset = Host.reset, + .prompt_chunk = Host.chunk, + .decode_chunk = Host.chunk, + .draw = Host.draw, + .advance = Host.advance, + .fed = Host.fed, + .memory = .{ .ctx = h, .read = Host.memory }, + .keepalive = .{ .ctx = &h.warm, .tick = mtl.keepalive.Target.tick }, + }); + return .{ .engine = h.serial.engine(), .close = close, .ctx = h }; +} + +pub fn close(ptr: *anyopaque) void { + const h = Host.self(ptr); + h.serial.deinit(); + h.model.deinit(); + h.gpa.destroy(h.model); + h.gpa.free(h.sampler_storage); + h.gpa.destroy(h); +} + +test "unsupported DeepSeek serving flags are refused before allocating or loading a checkpoint" { + var problem: []const u8 = ""; + const bad_options = [_]api.Open{ + .{ .dir = "", .model_type = "deepseek_v41", .tp = 2 }, + .{ .dir = "", .model_type = "deepseek_v41", .policy = "unknown=1" }, + .{ .dir = "", .model_type = "deepseek_v41", .drafter = "draft" }, + .{ .dir = "", .model_type = "deepseek_v41", .mtp_drafts = 1 }, + .{ .dir = "", .model_type = "deepseek_v41", .device = 0 }, + .{ .dir = "", .model_type = "deepseek_v41", .keep = 1 }, + }; + for (bad_options) |options| { + try std.testing.expect(try open(std.testing.failing_allocator, std.testing.failing_allocator, std.testing.io, options, &problem) == null); + } +} diff --git a/zig/src/native/flashnext_host.zig b/zig/src/native/flashnext_host.zig index e81619a8f0..c5e1ba2674 100644 --- a/zig/src/native/flashnext_host.zig +++ b/zig/src/native/flashnext_host.zig @@ -237,6 +237,7 @@ pub const Host = struct { const Ctx = struct { h: *Host, job: *Job, + failure: ?api.Stop.Error = null, fn prefilled(ctx: *anyopaque) void { const c: *Ctx = @ptrCast(@alignCast(ctx)); @@ -264,9 +265,15 @@ pub const Host = struct { var matched = false; var n: usize = 0; for (toks) |t| { // stop strings are checked after each token, as the lane core does - job.emitted.append(c.h.gpa, t) catch return true; + job.emitted.append(c.h.gpa, t) catch |err| { + c.failure = err; + return true; + }; n += 1; - if (job.request.stop) |s| if (s.check(s.ctx, job.emitted.items)) { + if (job.request.stop) |s| if (s.check(s.ctx, job.emitted.items) catch |err| { + c.failure = err; + return true; + }) { matched = true; break; }; @@ -349,6 +356,10 @@ pub const Host = struct { .cancelled => .cancelled, }; fin.stats = .{ .rounds = res.rounds, .drafted = res.drafted, .accepted = res.accepted, .min_rows = res.min_rows, .prefill_seconds = if (job.prefilled) |done| @as(f64, @floatFromInt(@as(i64, @intCast(@max(0, done - job.began))))) / 1e9 else null }; + if (c.failure) |err| { + fin.reason = .failed; + fin.message = @errorName(err); + } } }; diff --git a/zig/src/native/metal.zig b/zig/src/native/metal.zig index eb02076946..07e8559775 100644 --- a/zig/src/native/metal.zig +++ b/zig/src/native/metal.zig @@ -19,6 +19,7 @@ pub const families: []const api.Family = &.{ .{ .model_type = "qwen4_exp", .formats = &.{"mlx-q6g32"} }, .{ .model_type = "glm5_next", .formats = &.{"mlx-q4g64"} }, .{ .model_type = "qwen3_5", .formats = &.{"mlx-q4g64"} }, + .{ .model_type = "deepseek_v41", .formats = &.{"exl3"} }, }; /// The chip class gate entries name ("apple-m5" for an Apple M5 Max); null without an Apple GPU. @@ -99,6 +100,10 @@ const Host = struct { /// The engine for `o.dir`, or null with `problem` set when no Metal engine reads the checkpoint. pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !?api.Opened { + if (o.engram_table_dir != null and !std.mem.eql(u8, o.model_type, "deepseek_v41")) { + problem.* = "--engram-table-dir applies only to DeepSeek V4.1 resident checkpoints"; + return null; + } const dense = std.mem.eql(u8, o.model_type, "qwen3_5") and qwen27Engine(a, io, o.dir); if (o.drafter != null and !dense) { problem.* = "--drafter is supported for Qwen3.8-27B-class checkpoints (qwen3_5 with its own output head)"; @@ -108,6 +113,8 @@ pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]c if (std.mem.eql(u8, o.model_type, "qwen4_exp")) return openFlashNext(a, gpa, io, o, problem); if (std.mem.eql(u8, o.model_type, "glm5_next")) return openGlm(a, gpa, io, o, problem); if (std.mem.eql(u8, o.model_type, "qwen3_5")) return @import("qwen35.zig").open(a, gpa, io, o, problem); + if (std.mem.eql(u8, o.model_type, "deepseek_v41")) + return @import("deepseek_v41_host.zig").open(a, gpa, io, o, problem); if (!std.mem.eql(u8, o.model_type, "nemotron_h")) { problem.* = try std.fmt.allocPrint(a, "the native engine has no backend for {s} checkpoints yet; the Python engine 0.6.6 may serve them: python -m pip install git+https://github.com/ashhart/TensorFold.git@v0.6.6", .{o.model_type}); return null; diff --git a/zig/src/server/admission_consumer_regressions.zig b/zig/src/server/admission_consumer_regressions.zig new file mode 100644 index 0000000000..3445852dce --- /dev/null +++ b/zig/src/server/admission_consumer_regressions.zig @@ -0,0 +1,34 @@ +//! Consumer-visible refusals that must precede opening the event stream. +const std = @import("std"); +const api = @import("engine_api"); + +pub fn budgetWithoutThinking(engine: api.Engine, comptime sendWith: anytype) !void { + const body = "{\"model\":\"test-model\",\"messages\":[{\"role\":\"user\",\"content\":\"x\"}]," ++ + "\"max_tokens\":2,\"temperature\":0,\"thinking_budget\":5," ++ + "\"chat_template_kwargs\":{\"enable_thinking\":false}"; + for ([_][]const u8{ ",\"stream\":false}", ",\"stream\":true}" }) |suffix| { + const raw = try std.mem.concat(std.testing.allocator, u8, &.{ body, suffix }); + defer std.testing.allocator.free(raw); + const seen = try sendWith(engine, .chat, raw, .{}); + try std.testing.expectEqual(@as(?u16, 400), seen.status); + try std.testing.expect(!seen.opened); + try std.testing.expectEqual(@as(i64, 0), seen.preparing); + } +} + +pub fn media(comptime send: anytype) !void { + const bodies = [_][]const u8{ + "{\"model\":\"test-model\",\"messages\":[{\"role\":\"user\",\"content\":[" ++ + "{\"type\":\"image_url\",\"image_url\":{\"url\":\"https://example.com/image.png\"}}]}],", + "{\"model\":\"test-model\",\"messages\":[{\"role\":\"user\",\"content\":\"x\"}]," ++ + "\"modalities\":[\"text\",\"audio\"],", + }; + for (bodies) |body| for ([_][]const u8{ "\"stream\":false}", "\"stream\":true}" }) |suffix| { + const raw = try std.mem.concat(std.testing.allocator, u8, &.{ body, suffix }); + defer std.testing.allocator.free(raw); + const seen = try send(.chat, raw); + try std.testing.expectEqual(@as(?u16, 400), seen.status); + try std.testing.expect(!seen.opened); + try std.testing.expectEqual(@as(i64, 0), seen.preparing); + }; +} diff --git a/zig/src/server/anthropic.zig b/zig/src/server/anthropic.zig index 4821e0f3eb..65cda89982 100644 --- a/zig/src/server/anthropic.zig +++ b/zig/src/server/anthropic.zig @@ -4,7 +4,6 @@ const json = @import("json.zig"); const errors = @import("errors.zig"); const http_body = @import("http_body.zig"); const fields = @import("fields.zig"); -const messages_mod = @import("messages.zig"); const openai = @import("openai.zig"); const prompt = @import("prompt.zig"); const routes = @import("routes.zig"); @@ -47,11 +46,11 @@ fn sendError(conn: *Conn, a: Allocator, message: []const u8, status: u16) void { } /// Rendered like the chat route renders, without generating: the prompt's token count. -fn countTokens(srv: *Server, cx: *Cx, chat: Value) errors.Refused!usize { +pub fn countTokens(srv: *Server, cx: *Cx, chat: Value) errors.Refused!usize { const thinking = try fields.thinkingFields(cx, chat, srv.effort_levels); const on = thinking.enable orelse srv.config.enable_thinking; const tools = try tool_specs.active(cx, chat.get("tools"), chat.get("tool_choice")); - const msgs = try messages_mod.normalize(cx, chat.get("messages"), "system", srv.needs_user_after_tool); + const msgs = try openai.validatedMessages(srv, cx, chat.get("messages")); const ids_ = prompt.renderIds(srv, cx, msgs, tools, on, srv.effortFor(thinking.effort), true) catch |e| switch (e) { error.Refused => return cx.other(cx.message), else => |x| return x, @@ -275,7 +274,8 @@ const Reply = struct { target.id = id.string; }; const function = call.get("function") orelse Value.null; - target.name = std.mem.concat(r.a, u8, &.{ target.name, json.strOr(function.get("name")) }) catch return error.Closed; + const name = tool_specs.toolName(r.a, call) catch return error.Closed; + target.name = std.mem.concat(r.a, u8, &.{ target.name, name }) catch return error.Closed; target.arguments = std.mem.concat(r.a, u8, &.{ target.arguments, json.strOr(function.get("arguments")) }) catch return error.Closed; }; }; @@ -361,7 +361,7 @@ const Reply = struct { const b = try json.newObject(a); try b.put(a, "type", .{ .string = "tool_use" }); try b.put(a, "id", call.get("id").?); - try b.put(a, "name", function.get("name").?); + try b.put(a, "name", .{ .string = try tool_specs.toolName(a, call) }); try b.put(a, "input", args); try content.append(a, .{ .object = b }); }; diff --git a/zig/src/server/chat.zig b/zig/src/server/chat.zig index 6315e32991..d18876ce3d 100644 --- a/zig/src/server/chat.zig +++ b/zig/src/server/chat.zig @@ -57,66 +57,12 @@ pub const Reply = struct { effort: ?[]const u8 = null, /// ``choices[0].logprobs`` (``{"content": [...]}``) for a request that asked; null otherwise. logprobs: ?Value = null, + parsed_calls: ?tool_parse.Parsed = null, }; pub const Failure = error{ Refused, Cancelled, Failed, OutOfMemory }; -/// Events the engine thread hands this request, read by the request's own thread. -const Mailbox = struct { - io: std.Io, - gpa: Allocator, - mutex: std.Io.Mutex = .init, - cond: std.Io.Condition = .init, - tokens: std.ArrayList(u32) = .empty, - chunks: std.ArrayList(usize) = .empty, // where each round's tokens end - rows: std.ArrayList(api.LogprobRow) = .empty, // one a token of `tokens`, for requests with logprobs - cached: ?u32 = null, - prefilled_ns: ?i96 = null, - finished: bool = false, - reason: api.Reason = .stop, - stats: api.Stats = .{}, - widths: []u32 = &.{}, - raised: []bool = &.{}, - telemetry: []u8 = "", - message: []u8 = "", - - fn onEvent(ctx: *anyopaque, _: api.Id, event: *const api.Event) void { - const m: *Mailbox = @ptrCast(@alignCast(ctx)); - m.mutex.lockUncancelable(m.io); - defer m.mutex.unlock(m.io); - switch (event.*) { - .prefilled => |cached| { - m.cached = cached; - m.prefilled_ns = std.Io.Clock.awake.now(m.io).toNanoseconds(); - }, - .logprobs => |r| m.rows.appendSlice(m.gpa, r) catch {}, - .tokens => |t| { - m.tokens.appendSlice(m.gpa, t) catch {}; - m.chunks.append(m.gpa, m.tokens.items.len) catch {}; - }, - .finished => |f| { - m.finished = true; - m.reason = f.reason; - m.stats = f.stats; - m.widths = m.gpa.dupe(u32, f.stats.prefill_widths) catch &.{}; - m.raised = m.gpa.dupe(bool, f.stats.prefill_raised) catch &.{}; - m.telemetry = m.gpa.dupe(u8, f.stats.telemetry_json) catch ""; - m.message = m.gpa.dupe(u8, f.message) catch ""; - }, - } - m.cond.signal(m.io); - } - - fn deinit(m: *Mailbox) void { - m.tokens.deinit(m.gpa); - m.chunks.deinit(m.gpa); - m.rows.deinit(m.gpa); - m.gpa.free(m.widths); - m.gpa.free(m.raised); - m.gpa.free(m.telemetry); - m.gpa.free(m.message); - } -}; +const Mailbox = @import("chat_mailbox.zig").Mailbox; const nowNs = clock.nowNs; const seconds = clock.seconds; @@ -195,9 +141,11 @@ pub fn prepare(srv: *Server, cx: *Cx, input: Input, gone: anytype) Failure!Prepa .decide = if (input.prompt == null and !background) prompt_mod.decideOf(srv, rendered.ids) else null, }; try srv.checkFeatures(cx, f, input.tools.len > 0, thinking, rendered.ids, input.tools, &request); + const budget_field = f.get("thinking_budget"); + const budget: i64 = if (budget_field != null and budget_field.?.truthy()) budget_field.?.int64() orelse 0 else srv.config.thinking_budget; + if (budget > 0 and !srv.info.thinking_budget) + return cx.refuse("thinking_budget is not supported by this native engine"); if (thinking) { - const budget_field = f.get("thinking_budget"); - const budget: i64 = if (budget_field != null and budget_field.?.truthy()) budget_field.?.int64() orelse 0 else srv.config.thinking_budget; if (srv.think_close.len > 0) request.loop_guard = srv.config.loop_guard; if ((budget > 0 or request.loop_guard) and srv.think_close.len > 0) { if (budget > 0) request.think_budget = @intCast(@min(budget, std.math.maxInt(u32))); @@ -250,7 +198,13 @@ pub fn generate(srv: *Server, cx: *Cx, prepared: Prepared, sink: ?Sink, gone: an errdefer if (!gen.engine_done) gen.cancel(); // the engine writes to the mailbox until it says finished const result: Failure!Reply = blk: { if (sink != null and input.tools.len > 0) gen.calls = tool_stream.Streamer.init(a, input.tools) catch |e| break :blk e; - gen.loop(gone) catch |e| break :blk e; + gen.loop(gone) catch |e| break :blk switch (e) { + error.Template => cx.other("the tokenizer failed to decode the reply"), + error.Cancelled => error.Cancelled, + error.Failed => error.Failed, + error.OutOfMemory => error.OutOfMemory, + error.Refused => error.Refused, + }; break :blk gen.finish(cx, prepared.prompt_len, received, thinking, effort, prepared.sampled, prepared.drafts); }; if (result) |_| {} else |e| logEnded(input.id, e, cx.message, prepared.prompt_len, gen.collected.items.len, seconds(nowNs(io) - received)); @@ -348,7 +302,7 @@ const Generation = struct { } /// The engine's stop check, here: the newest tokens' text holds a stop string. - fn stopHit(g: *Generation) Allocator.Error!bool { + fn stopHit(g: *Generation) @import("model_text.zig").Error!bool { if (g.stops.strings.len == 0) return false; const t = g.collected.items; const tail = t[t.len -| g.stops.tail()..]; @@ -358,7 +312,7 @@ const Generation = struct { } /// Commits a round's tokens as ``LaneStream.commit`` would, then streams what they add. - fn commit(g: *Generation, chunk: []const u32) Failure!void { + fn commit(g: *Generation, chunk: []const u32) (Failure || @import("model_text.zig").Error)!void { var landed: usize = 0; for (chunk) |t| { if (g.reason != null) break; @@ -437,7 +391,7 @@ const Generation = struct { } /// Takes rounds until the reply ends; a client that leaves cancels it. - fn loop(g: *Generation, gone: anytype) Failure!void { + fn loop(g: *Generation, gone: anytype) (Failure || @import("model_text.zig").Error)!void { const m = g.box; while (true) { m.mutex.lockUncancelable(m.io); @@ -479,7 +433,10 @@ const Generation = struct { }; const finished_ns = nowNs(g.srv.io); const content_tokens = if (g.ignore_eos) g.collected.items else reply_text.stripTrailing(g.collected.items, g.srv.eos); - const raw = try g.srv.text.decode(a, content_tokens); + const raw = g.srv.text.decode(a, content_tokens) catch |e| switch (e) { + error.Template => return cx.other("the tokenizer failed to decode the reply"), + error.OutOfMemory => return error.OutOfMemory, + }; const visible = g.stops.visible(raw, false); // closed before the think split, so a call ending an unclosed think block is the answer const close = try g.closeCall(visible); @@ -496,19 +453,31 @@ const Generation = struct { content = h.content; reasoning = h.reasoning; } + var final_calls: std.ArrayList(Value) = .empty; + if (g.calls) |*c| { + if (tool_parse.dsmlOpening(content, 0) != null) + try c.feed(content, &final_calls); + if (c.invalidDsml(content)) + return cx.other("the model generated a malformed or unoffered DeepSeek V4.1 tool call"); + } + const parsed_calls = if (g.tools.len > 0 and tool_parse.dsmlOpening(content, 0) != null) + try tool_parse.parse(a, content, g.tools, null) + else + null; + if (parsed_calls) |parsed| if (reply_text.findTag(parsed.content, 0, "<｜DSML｜ calls>") != null) + return cx.other("the model generated a malformed or unoffered DeepSeek V4.1 tool call"); if (g.sink) |sink| { const thought = g.streamed_reasoning.text.items; if (reasoning) |r| if (std.mem.startsWith(u8, r, thought) and r.len > thought.len) try g.emit(sink, try deltaOf(a, "reasoning_content", r[thought.len..])); - const shown = if (g.calls != null) try reply_text.finishedProse(a, content) else content; + const shown = if (parsed_calls) |parsed| parsed.content else if (g.calls != null) try reply_text.finishedProse(a, content) else content; const sent = g.streamed.text.items; if (std.mem.startsWith(u8, shown, sent) and shown.len > sent.len) try g.emit(sink, .{ .string = shown[sent.len..] }); - if (g.calls) |*c| if (close.len > 0) { - // the streamer reads the closers as markup the model wrote, so the call ends with the same deltas - var out: std.ArrayList(Value) = .empty; - try c.feed(if (g.thinking) content else text, &out); - for (out.items) |d| try g.emit(sink, d); + if (g.calls) |*c| if (close.len > 0 and parsed_calls == null) { + // Preserve legacy streaming's synthesized-closer path and delta ordering. + try c.feed(if (g.thinking) content else text, &final_calls); }; + for (final_calls.items) |d| try g.emit(sink, d); } const prefilled = m.prefilled_ns; const decode_s: f64 = if (prefilled) |p| @max(0, seconds(finished_ns - p)) else 0; @@ -551,6 +520,7 @@ const Generation = struct { .stop_sequence = g.stops.matched(raw), .reasoning = reasoning, .tool_calls_streamed = g.calls != null and g.calls.?.streamed, + .parsed_calls = parsed_calls, .finish_reason = reason, .prompt_tokens = prompt_len, .cached_tokens = m.cached orelse 0, @@ -571,21 +541,7 @@ const Generation = struct { } }; -/// The request's stop strings as the engine checks them after each token: the newest tokens' text holds one. -const StopHook = struct { - srv: *Server, - stops: reply_text.Stops, - - fn check(ctx: *anyopaque, emitted: []const u32) bool { - const h: *StopHook = @ptrCast(@alignCast(ctx)); - var arena: std.heap.ArenaAllocator = .init(h.srv.gpa); - defer arena.deinit(); - const tail = emitted[emitted.len -| h.stops.tail()..]; - const text = h.srv.text.decode(arena.allocator(), tail) catch return false; - for (h.stops.strings) |stop| if (std.mem.indexOf(u8, text, stop) != null) return true; - return false; - } -}; +const StopHook = @import("chat_stop.zig").Hook; fn decodeText(t: anytype, a: Allocator, tokens: []const u32) ![]u8 { return t.decode(a, tokens); diff --git a/zig/src/server/chat_mailbox.zig b/zig/src/server/chat_mailbox.zig new file mode 100644 index 0000000000..b90af4ec83 --- /dev/null +++ b/zig/src/server/chat_mailbox.zig @@ -0,0 +1,60 @@ +//! Engine-thread events owned by one request until its terminal event. +const std = @import("std"); +const api = @import("engine_api"); +const Allocator = std.mem.Allocator; + +pub const Mailbox = struct { + io: std.Io, + gpa: Allocator, + mutex: std.Io.Mutex = .init, + cond: std.Io.Condition = .init, + tokens: std.ArrayList(u32) = .empty, + chunks: std.ArrayList(usize) = .empty, // where each round's tokens end + rows: std.ArrayList(api.LogprobRow) = .empty, // one a token of `tokens`, for requests with logprobs + cached: ?u32 = null, + prefilled_ns: ?i96 = null, + finished: bool = false, + reason: api.Reason = .stop, + stats: api.Stats = .{}, + widths: []u32 = &.{}, + raised: []bool = &.{}, + telemetry: []u8 = "", + message: []u8 = "", + + pub fn onEvent(ctx: *anyopaque, _: api.Id, event: *const api.Event) void { + const m: *Mailbox = @ptrCast(@alignCast(ctx)); + m.mutex.lockUncancelable(m.io); + defer m.mutex.unlock(m.io); + switch (event.*) { + .prefilled => |cached| { + m.cached = cached; + m.prefilled_ns = std.Io.Clock.awake.now(m.io).toNanoseconds(); + }, + .logprobs => |r| m.rows.appendSlice(m.gpa, r) catch {}, + .tokens => |t| { + m.tokens.appendSlice(m.gpa, t) catch {}; + m.chunks.append(m.gpa, m.tokens.items.len) catch {}; + }, + .finished => |f| { + m.finished = true; + m.reason = f.reason; + m.stats = f.stats; + m.widths = m.gpa.dupe(u32, f.stats.prefill_widths) catch &.{}; + m.raised = m.gpa.dupe(bool, f.stats.prefill_raised) catch &.{}; + m.telemetry = m.gpa.dupe(u8, f.stats.telemetry_json) catch ""; + m.message = m.gpa.dupe(u8, f.message) catch ""; + }, + } + m.cond.signal(m.io); + } + + pub fn deinit(m: *Mailbox) void { + m.tokens.deinit(m.gpa); + m.chunks.deinit(m.gpa); + m.rows.deinit(m.gpa); + m.gpa.free(m.widths); + m.gpa.free(m.raised); + m.gpa.free(m.telemetry); + m.gpa.free(m.message); + } +}; diff --git a/zig/src/server/chat_stop.zig b/zig/src/server/chat_stop.zig new file mode 100644 index 0000000000..613b534730 --- /dev/null +++ b/zig/src/server/chat_stop.zig @@ -0,0 +1,20 @@ +//! The request's stop strings as the engine checks each newly emitted token. +const std = @import("std"); +const Server = @import("server.zig").Server; +const reply_text = @import("reply_text.zig"); +const api = @import("engine_api"); + +pub const Hook = struct { + srv: *Server, + stops: reply_text.Stops, + + pub fn check(ctx: *anyopaque, emitted: []const u32) api.Stop.Error!bool { + const h: *Hook = @ptrCast(@alignCast(ctx)); + var arena: std.heap.ArenaAllocator = .init(h.srv.gpa); + defer arena.deinit(); + const tail = emitted[emitted.len -| h.stops.tail()..]; + const text = try h.srv.text.decode(arena.allocator(), tail); + for (h.stops.strings) |stop| if (std.mem.indexOf(u8, text, stop) != null) return true; + return false; + } +}; diff --git a/zig/src/server/cli.zig b/zig/src/server/cli.zig index 40ee1e0ee8..384aa1c1df 100644 --- a/zig/src/server/cli.zig +++ b/zig/src/server/cli.zig @@ -55,6 +55,7 @@ pub const flags = [_]Flag{ .{ .name = "--vision-max-images" }, .{ .name = "--vision-image-tokens" }, .{ .name = "--context", .native = true }, + .{ .name = "--engram-table-dir", .native = !cuda_build }, .{ .name = "--speed-up", .native = true }, .{ .name = "--max-tokens", .native = true }, .{ .name = "--temperature", .native = true }, @@ -129,6 +130,7 @@ pub const Args = struct { metrics_open: bool = false, dashboard: bool = false, context: ?i64 = null, + engram_table_dir: ?[]const u8 = null, speed_up: ?[]const u8 = null, // speed-up mode: this Mac's settings for the two-Mac link (Flash Next) prompt_cache_gib: ?f64 = null, // kept prompt states' budget (0: none; null: what 70% of RAM leaves past the loaded server) prompt_cache_over_cap: bool = false, // a --prompt-cache-gib past that is kept, not refused @@ -250,6 +252,11 @@ fn apply(a: Allocator, out: *Args, name: []const u8, value: ?[]const u8, u: *Usa out.chat_template = v; return; } + if (is(name, "--engram-table-dir")) { + if (v.len == 0) return fail(u, a, "argument --engram-table-dir: expected a directory", .{}); + out.engram_table_dir = v; + return; + } if (is(name, "--host")) out.host = v else if (is(name, "--port")) { const p = try int(u, a, name, v); if (p < 0 or p > 65535) return fail(u, a, "argument --port: invalid port: '{s}'", .{v}); @@ -342,6 +349,7 @@ pub fn request(a: Allocator, dir: []const u8, model_type: []const u8, args: Args .dir = dir, .model_type = model_type, .context = args.context, + .engram_table_dir = args.engram_table_dir, .lanes = parallel(args.parallel) orelse 8, .lanes_fixed = parallelFixed(args.parallel), .drafts = !args.no_drafts, diff --git a/zig/src/server/decoder_consumer_regressions.zig b/zig/src/server/decoder_consumer_regressions.zig new file mode 100644 index 0000000000..780826391c --- /dev/null +++ b/zig/src/server/decoder_consumer_regressions.zig @@ -0,0 +1,102 @@ +//! Real text-interface decoder failures must survive through the consumer error boundary. +const std = @import("std"); +const api = @import("engine_api"); +const json = @import("json.zig"); +const model_text = @import("model_text.zig"); +const openai = @import("openai.zig"); +const Server = @import("server.zig").Server; +const Conn = @import("http_conn.zig").Conn; + +fn failedDecode(_: *anyopaque, _: std.mem.Allocator, _: []const u32) model_text.Error![]u8 { + return error.Template; +} + +fn failedBytes(_: *anyopaque, _: std.mem.Allocator, _: u32) model_text.Error!?[]u8 { + return error.Template; +} + +fn info(_: *anyopaque) api.Info { + return .{ .context_window = 4096, .logprobs = true }; +} + +const Output = struct { + opened: bool = false, + status: ?u16 = null, + decoder_error: bool = false, + content: bool = false, + + fn out(o: *@This()) openai.Out { + return .{ .ctx = o, .vt = &.{ .open = open, .event = event, .reply = reply } }; + } + fn open(ctx: *anyopaque) error{Closed}!void { + const o: *Output = @ptrCast(@alignCast(ctx)); + o.opened = true; + } + fn observe(o: *Output, payload: json.Value) void { + if (payload.get("error")) |err| if (err.strField("message")) |message| { + o.decoder_error = std.mem.indexOf(u8, message, "tokenizer") != null and + std.mem.indexOf(u8, message, "decode") != null and std.mem.indexOf(u8, message, "memory") == null; + }; + if (payload.get("choices")) |choices| if (choices == .array) for (choices.array) |choice| { + const message = choice.get("delta") orelse choice.get("message") orelse continue; + if (message.strField("content")) |content| o.content = o.content or content.len > 0; + }; + } + fn event(ctx: *anyopaque, payload: ?json.Value) error{Closed}!void { + const o: *Output = @ptrCast(@alignCast(ctx)); + if (payload) |p| o.observe(p); + } + fn reply(ctx: *anyopaque, status: u16, payload: json.Value) void { + const o: *Output = @ptrCast(@alignCast(ctx)); + o.status = status; + o.observe(payload); + } +}; + +pub fn run(comptime TestText: type, comptime HiEngine: type) !void { + const ReplyDecoder = struct { + fn decode(ctx: *anyopaque, a: std.mem.Allocator, ids: []const u32) model_text.Error![]u8 { + if (ids.len > 0 and ids[0] == 'h') return error.Template; + const original: *TestText = @ptrCast(@alignCast(ctx)); + return original.text().decode(a, ids); + } + }; + for (0..3) |mode| for ([_]bool{ false, true }) |stream| { + var baseline: TestText = .{}; + const base = baseline.text(); + var text_vtable = base.vtable.*; + switch (mode) { + 0 => text_vtable.decode = failedDecode, + 1 => text_vtable.token_bytes = failedBytes, + else => text_vtable.decode = ReplyDecoder.decode, + } + var backend: HiEngine = .{}; + const engine = backend.engine(); + var engine_vtable = engine.vtable.*; + engine_vtable.info = info; + const srv = try Server.init(std.testing.allocator, std.testing.io, .{ .ctx = engine.ctx, .vtable = &engine_vtable }, .{ .ctx = base.ctx, .vtable = &text_vtable }, .{ .served_name = "test-model", .model_ids = &.{"test-model"}, .enable_thinking = false, .use_drafts = false }, null); + defer srv.deinit(); + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const fields = switch (mode) { + 0 => "\"prompt\":[[120]]", + 1 => "\"messages\":[{\"role\":\"user\",\"content\":\"x\"}],\"logprobs\":true", + else => "\"messages\":[{\"role\":\"user\",\"content\":\"x\"}]", + }; + const raw = try std.fmt.allocPrint(a, "{{\"model\":\"test-model\",\"max_tokens\":2,{s},\"stream\":{s}}}", .{ fields, if (stream) "true" else "false" }); + var seen: Output = .{}; + var conn: Conn = .{ .fd = -1, .peer = "", .buf = &.{}, .gpa = a }; + openai.run(srv, a, seen.out(), .{ .conn = &conn }, mode != 0, (try json.parse(a, raw)).ok); + try std.testing.expect(seen.decoder_error); + try std.testing.expect(!seen.content); + if (mode < 2 or !stream) { + try std.testing.expectEqual(@as(?u16, 500), seen.status); + try std.testing.expect(!seen.opened); + } else { + try std.testing.expect(seen.opened); + try std.testing.expectEqual(@as(?u16, null), seen.status); + } + try std.testing.expectEqual(@as(i64, 0), srv.preparing.load(.acquire)); + }; +} diff --git a/zig/src/server/deepseek_v41_official.jinja b/zig/src/server/deepseek_v41_official.jinja new file mode 100644 index 0000000000..eb12dd280d --- /dev/null +++ b/zig/src/server/deepseek_v41_official.jinja @@ -0,0 +1,252 @@ +{#- DeepSeek-V4.1 chat template, ported from encoding/encoding.py in the model repo. + Extra kwargs: thinking_mode ("thinking" | "chat", default "thinking"; enable_thinking=true/false is accepted + as an alias), drop_thinking (default true), reasoning_effort (int in [1, 100] or "low"/"medium"/"high"/"max", + default "high"). Assistant reasoning is read from reasoning_content, or from thinking as an alias. -#} +{%- set IMG = "<｜deepseek_image｜>" -%} +{%- set DSML = "｜DSML｜" -%} + +{%- if thinking_mode is not defined -%} + {%- set thinking_mode = "chat" if (enable_thinking is defined and enable_thinking is not none and not enable_thinking) else "thinking" -%} +{%- endif -%} +{%- if thinking_mode not in ["chat", "thinking"] -%} + {{- raise_exception("Invalid thinking_mode `" ~ thinking_mode ~ "`, should be \"chat\" or \"thinking\"") -}} +{%- endif -%} +{%- if drop_thinking is not defined -%}{%- set drop_thinking = true -%}{%- endif -%} +{%- if add_generation_prompt is not defined -%}{%- set add_generation_prompt = false -%}{%- endif -%} +{%- if tools is not defined -%}{%- set tools = none -%}{%- endif -%} + +{%- if reasoning_effort is not defined or reasoning_effort is none -%}{%- set reasoning_effort = "high" -%}{%- endif -%} +{%- set effort_map = {"low": 50, "medium": 62, "high": 75, "max": 100} -%} +{%- if reasoning_effort is string and reasoning_effort in effort_map -%} + {%- set budget = effort_map[reasoning_effort] -%} +{%- elif reasoning_effort is integer and 1 <= reasoning_effort <= 100 -%} + {%- set budget = reasoning_effort -%} +{%- else -%} + {{- raise_exception("Invalid reasoning effort: " ~ reasoning_effort ~ ", should be int within [1,100] or one of \"low\", \"medium\", \"high\", \"max\"") -}} +{%- endif -%} + +{#- Text of text/image blocks joined by blank lines; other block types are dropped. -#} +{%- macro joined_text(blocks) -%} + {%- set jt = namespace(parts=[]) -%} + {%- for b in blocks -%} + {%- if b.type == "text" -%} + {%- if IMG in (b.text or "") -%}{{- raise_exception("Text block contains image placeholder " ~ IMG ~ ". Images should be separate content blocks.") -}}{%- endif -%} + {%- set jt.parts = jt.parts + [b.text or ""] -%} + {%- elif b.type in ["image", "image_url"] -%} + {%- set jt.parts = jt.parts + [IMG] -%} + {%- endif -%} + {%- endfor -%} + {{- jt.parts | join("\n\n") -}} +{%- endmacro -%} + +{#- Splits/validates "ns::name" against an explicit namespace and returns the qualified name. -#} +{%- macro qualified_name(name, ns) -%} + {%- if ns is mapping -%}{%- set ns = ns.name -%}{%- endif -%} + {%- set parts = name.split("::", 1) -%} + {%- if parts | length == 2 -%} + {%- if ns is not none and ns != parts[0] -%}{{- raise_exception("Conflicting tool namespaces: " ~ ns ~ " != " ~ parts[0]) -}}{%- endif -%} + {%- set ns = parts[0] -%} + {%- set name = parts[1] -%} + {%- endif -%} + {%- if "::" in name -%}{{- raise_exception("Tool name must not contain '::': " ~ name) -}}{%- endif -%} + {%- if ns is not none and "::" in ns -%}{{- raise_exception("Tool namespace must not contain '::': " ~ ns) -}}{%- endif -%} + {{- name if ns is none else ns ~ "::" ~ name -}} +{%- endmacro -%} + +{%- macro render_tools(tool_list) -%} + {%- set ts = namespace(schemas=[]) -%} + {%- for tool in tool_list -%} + {%- set fn = tool.function if tool.function is defined else tool -%} + {%- set ns = tool.namespace if tool.namespace is defined and tool.namespace is not none else fn.get("namespace") -%} + {%- set schema = namespace(d={}) -%} + {%- for k, v in fn.items() if k != "namespace" -%}{%- set schema.d = dict(schema.d, **{k: v}) -%}{%- endfor -%} + {%- set schema.d = dict(schema.d, name=qualified_name(fn.name, ns)) -%} + {%- if ns is mapping and ns.description -%} + {%- set schema.d = dict(schema.d, description=ns.description ~ "\n" ~ (fn.description or "")) -%} + {%- endif -%} + {%- set ts.schemas = ts.schemas + [schema.d | tojson] -%} + {%- endfor -%} + {{- "## Tools\n\nYou have access to a set of tools to help answer the user's question. You can invoke tools by writing a \"<｜DSML｜ calls>\" block like the following:\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"$TOOL_NAME\">\n<｜DSML｜ parameter name=\"$PARAMETER_NAME\" string=\"true|false\">$PARAMETER_VALUE\n...\n\n<｜DSML｜ invoke name=\"$TOOL_NAME2\">\n...\n\n\n\nString parameters should be specified as is and set `string=\"true\"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string=\"false\"`.\n\nIf thinking_mode is enabled (triggered by ), you MUST output your complete reasoning inside ... BEFORE any tool calls or final response.\n\nOtherwise, output directly after with tool calls or final response.\n\n### Available Tool Schemas\n\n" -}} + {{- ts.schemas | join("\n") -}} + {{- "\n\nYou MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.\n" -}} +{%- endmacro -%} + +{#- 1. Merge tool results (and consecutive user turns) into user messages made of content blocks. -#} +{%- set mns = namespace(list=[]) -%} +{%- for msg in messages -%} + {%- set content = msg.get("content") -%} + {%- set is_blocks = content is iterable and content is not string and content is not mapping -%} + {%- if content is string and IMG in content -%} + {{- raise_exception("Message content contains image special token " ~ IMG ~ ". Images should be provided as image content blocks.") -}} + {%- endif -%} + {%- set rc = msg.reasoning_content if msg.get("reasoning_content") is not none else msg.get("thinking") -%} + {%- if rc is string and IMG in rc -%} + {{- raise_exception("reasoning_content contains image special token " ~ IMG) -}} + {%- endif -%} + {%- set prev = mns.list[-1] if mns.list else none -%} + {%- if msg.role == "tool" -%} + {%- set tblock = {"type": "tool_result", "tool_use_id": msg.get("tool_call_id", ""), "content": joined_text(content) if is_blocks else msg.get("content", "")} -%} + {%- if prev is not none and prev.role == "user" -%} + {%- set mns.list = mns.list[:-1] + [dict(prev, content_blocks=prev.content_blocks + [tblock])] -%} + {%- else -%} + {%- set mns.list = mns.list + [{"role": "user", "content_blocks": [tblock]}] -%} + {%- endif -%} + {%- elif msg.role == "user" -%} + {%- set blocks = content if is_blocks else [{"type": "text", "text": content or ""}] -%} + {%- if prev is not none and prev.role == "user" and prev.get("task") is none -%} + {%- set mns.list = mns.list[:-1] + [dict(prev, content_blocks=prev.content_blocks + blocks)] -%} + {%- else -%} + {%- set mns.list = mns.list + [dict(msg, content_blocks=blocks)] -%} + {%- endif -%} + {%- else -%} + {%- set mns.list = mns.list + [msg] -%} + {%- endif -%} +{%- endfor -%} + +{#- Top-level tools attach to a leading system message, or to a new empty one. -#} +{%- if tools -%} + {%- if mns.list and mns.list[0].role == "system" -%} + {%- if not mns.list[0].get("tools") -%}{%- set mns.list = [dict(mns.list[0], tools=tools)] + mns.list[1:] -%}{%- endif -%} + {%- else -%} + {%- set mns.list = [{"role": "system", "content": "", "tools": tools}] + mns.list -%} + {%- endif -%} +{%- endif -%} + +{#- 2. Reorder tool results to match the tool_calls order of the preceding assistant message. -#} +{%- set sns = namespace(list=[], ids=[]) -%} +{%- for m in mns.list -%} + {%- if m.role == "assistant" and m.get("tool_calls") -%} + {%- set sns.ids = [] -%} + {%- for tc in m.tool_calls -%} + {%- set sns.ids = sns.ids + [tc.get("id") or tc.get("function", {}).get("id", "")] -%} + {%- endfor -%} + {%- set sns.list = sns.list + [m] -%} + {%- elif m.role == "user" and m.content_blocks -%} + {%- set tb = namespace(blocks=[]) -%} + {%- for b in m.content_blocks if b.type == "tool_result" -%} + {%- set order = namespace(idx=0) -%} + {%- for tc_id in sns.ids -%} + {%- if tc_id and tc_id == b.get("tool_use_id", "") -%}{%- set order.idx = loop.index0 -%}{%- endif -%} + {%- endfor -%} + {%- set tb.blocks = tb.blocks + [dict(b, _order=order.idx)] -%} + {%- endfor -%} + {%- if tb.blocks | length > 1 -%} + {%- set sorted_blocks = tb.blocks | sort(attribute="_order") -%} + {%- set rb = namespace(blocks=[], i=0) -%} + {%- for b in m.content_blocks -%} + {%- if b.type == "tool_result" -%} + {%- set rb.blocks = rb.blocks + [sorted_blocks[rb.i]] -%} + {%- set rb.i = rb.i + 1 -%} + {%- else -%} + {%- set rb.blocks = rb.blocks + [b] -%} + {%- endif -%} + {%- endfor -%} + {%- set sns.list = sns.list + [dict(m, content_blocks=rb.blocks)] -%} + {%- else -%} + {%- set sns.list = sns.list + [m] -%} + {%- endif -%} + {%- else -%} + {%- set sns.list = sns.list + [m] -%} + {%- endif -%} +{%- endfor -%} +{%- set msgs = sns.list -%} + +{#- Mid-conversation system messages count as user turns. Tools anywhere disable drop_thinking. -#} +{%- set state = namespace(last_user=-1, has_tools=false) -%} +{%- for m in msgs -%} + {%- if m.role == "user" or (m.role == "system" and loop.index0 > 0) -%}{%- set state.last_user = loop.index0 -%}{%- endif -%} + {%- if m.get("tools") -%}{%- set state.has_tools = true -%}{%- endif -%} +{%- endfor -%} +{%- set effective_drop = drop_thinking and not state.has_tools -%} +{%- set thinking = thinking_mode == "thinking" -%} + +{{- "<｜begin▁of▁sentence｜>" -}} +{#- 3. Render each message. -#} +{%- for m in msgs -%} + {%- set idx = loop.index0 -%} + {%- set content = m.get("content") -%} + {%- if content is iterable and content is not string and content is not mapping -%}{%- set content = joined_text(content) -%}{%- endif -%} + + {%- if idx == 0 -%} + {%- if thinking or m.role == "system" -%}{{- "<｜System｜>" -}}{%- endif -%} + {%- if thinking -%}{{- "Reasoning Effort: " ~ budget ~ " (range 1-100, the higher the value, the more thorough the reasoning)\n\n" -}}{%- endif -%} + {%- endif -%} + + {%- if m.role == "system" -%} + {%- if idx > 0 -%}{{- "<｜System｜>" -}}{%- endif -%} + {{- content or "" -}} + {%- if m.get("tools") -%}{{- "\n\n" ~ render_tools(m.tools) -}}{%- endif -%} + {%- if m.get("response_format") -%} + {{- "\n\n## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n" ~ (m.response_format | tojson) -}} + {%- endif -%} + + {%- elif m.role == "user" -%} + {{- "<｜User｜>" -}} + {%- set up = namespace(parts=[]) -%} + {%- for b in m.content_blocks -%} + {%- if b.type == "text" -%} + {%- if IMG in (b.text or "") -%}{{- raise_exception("Text block contains image placeholder " ~ IMG ~ ". Images should be separate content blocks.") -}}{%- endif -%} + {%- set up.parts = up.parts + [b.text or ""] -%} + {%- elif b.type in ["image", "image_url"] -%} + {%- set up.parts = up.parts + [IMG] -%} + {%- elif b.type == "tool_result" -%} + {%- set tc = b.get("content", "") -%} + {%- if tc is iterable and tc is not string and tc is not mapping -%} + {%- set sub = namespace(parts=[]) -%} + {%- for s in tc -%} + {%- if s.type == "text" -%}{%- set sub.parts = sub.parts + [s.get("text", "")] -%} + {%- elif s.type in ["image", "image_url"] -%}{%- set sub.parts = sub.parts + [IMG] -%} + {%- else -%}{%- set sub.parts = sub.parts + ["[Unsupported " ~ s.type ~ "]"] -%} + {%- endif -%} + {%- endfor -%} + {%- set tc = sub.parts | join("\n\n") -%} + {%- endif -%} + {%- set up.parts = up.parts + ["" ~ tc ~ ""] -%} + {%- else -%} + {%- set up.parts = up.parts + ["[Unsupported " ~ b.type ~ "]"] -%} + {%- endif -%} + {%- endfor -%} + {{- up.parts | join("\n\n") -}} + + {%- elif m.role == "latest_reminder" -%} + {{- "<｜latest_reminder｜>" ~ content -}} + + {%- elif m.role == "assistant" -%} + {%- set prev_has_task = idx > 0 and msgs[idx - 1].get("task") is not none -%} + {%- if thinking and not prev_has_task and (not effective_drop or idx > state.last_user) -%} + {{- ((m.reasoning_content if m.get("reasoning_content") is not none else m.get("thinking")) or "") ~ "" -}} + {%- endif -%} + {{- content or "" -}} + {%- if m.get("tool_calls") -%} + {%- set tcs = namespace(lines=[]) -%} + {%- for tc in m.tool_calls -%} + {%- set fn = tc.function -%} + {%- set params = namespace(lines=[]) -%} + {%- set args = fn.arguments if fn.arguments is mapping else {"arguments": fn.arguments} -%} + {%- for k, v in args.items() -%} + {%- set params.lines = params.lines + ["<" ~ DSML ~ " parameter name=\"" ~ k ~ "\" string=\"" ~ ("true" if v is string else "false") ~ "\">" ~ (v if v is string else v | tojson) ~ ""] -%} + {%- endfor -%} + {%- set name = qualified_name(fn.name, tc.get("namespace") or fn.get("namespace")) -%} + {%- set tcs.lines = tcs.lines + ["<" ~ DSML ~ " invoke name=\"" ~ name ~ "\">\n" ~ params.lines | join("\n") ~ "\n"] -%} + {%- endfor -%} + {{- "\n\n<" ~ DSML ~ " calls>\n" ~ tcs.lines | join("\n") ~ "\n" -}} + {%- endif -%} + {%- if not m.get("wo_eos", false) -%}{{- "<｜end▁of▁sentence｜>" -}}{%- endif -%} + + {%- else -%} + {{- raise_exception("Unknown role: " ~ m.role) -}} + {%- endif -%} + + {#- Transition tokens, only when the next message is an assistant/latest_reminder turn or this is the last one. -#} + {%- if loop.last or msgs[idx + 1].role in ["assistant", "latest_reminder"] -%} + {%- set task = m.get("task") -%} + {%- if task is not none -%} + {%- set task_tokens = {"action": "<｜action｜>", "query": "<｜query｜>", "authority": "<｜authority｜>", "domain": "<｜domain｜>", "title": "<｜title｜>", "read_url": "<｜read_url｜>"} -%} + {%- if task not in task_tokens -%}{{- raise_exception("Invalid task: '" ~ task ~ "'") -}}{%- endif -%} + {%- if task == "action" -%}{{- "<｜Assistant｜>" ~ ("" if thinking else "") -}}{%- endif -%} + {{- task_tokens[task] -}} + {%- elif (m.role == "user" or (m.role == "system" and idx > 0)) and (add_generation_prompt or not loop.last) -%} + {{- "<｜Assistant｜>" ~ ("" if thinking and (not effective_drop or idx >= state.last_user) else "") -}} + {%- endif -%} + {%- endif -%} +{%- endfor -%} diff --git a/zig/src/server/deepseek_v41_text.zig b/zig/src/server/deepseek_v41_text.zig new file mode 100644 index 0000000000..b982847ef0 --- /dev/null +++ b/zig/src/server/deepseek_v41_text.zig @@ -0,0 +1,134 @@ +//! Text-only official DS4.1 template adaptation; pinned DeepSeek-AI revision and MIT notice in THIRD_PARTY_NOTICES.md. +const std = @import("std"); +const json = @import("json.zig"); +const Value = json.Value; +const Allocator = std.mem.Allocator; +const Error = error{ Template, OutOfMemory }; + +pub fn isTemplate(source: []const u8) bool { + return std.mem.indexOf(u8, source, "DeepSeek-V4.1 chat template") != null and + std.mem.indexOf(u8, source, " calls>") != null; +} + +fn rejectMedia(value: Value) Error!void { + switch (value) { + .string => |s| if (std.mem.indexOf(u8, s, "<｜deepseek_image｜>") != null or + std.mem.indexOf(u8, s, "") != null or std.mem.indexOf(u8, s, "") != null) return error.Template, + .array => |items| for (items) |item| try rejectMedia(item), + .object => |o| { + for ([_][]const u8{ "image", "images", "image_url", "input_image", "audio", "input_audio", "video", "video_url" }) |name| { + if (value.typeIs(name) or json.truthyField(value, name)) return error.Template; + } + for (o.values()) |item| try rejectMedia(item); + }, + else => {}, + } +} + +fn validText(value: ?Value) Error!void { + const v = value orelse return; + if (v == .null or v == .string) return; + if (v != .array) return error.Template; + for (v.array) |part| { + if (part != .object or part.strField("type") == null) return error.Template; + if (part.typeIs("text")) { + if (part.field("text")) |t| if (t != .string) return error.Template; + } else if (part.typeIs("tool_result")) try validText(part.get("content")); + } +} + +fn validateMessage(message: Value) Error!void { + if (message != .object) return error.Template; + for ([_][]const u8{ "image", "images", "image_url", "input_image", "audio", "input_audio", "video", "video_url" }) |name| { + if (json.truthyField(message, name)) return error.Template; + } + for ([_][]const u8{ "content", "content_blocks", "reasoning_content", "thinking" }) |name| { + if (message.get(name)) |value| try rejectMedia(value); + } + try validText(message.get("content")); + try validText(message.get("content_blocks")); + const role = message.strField("role") orelse return error.Template; + for ([_][]const u8{ "system", "developer", "user", "assistant", "tool", "latest_reminder" }) |accepted| { + if (std.mem.eql(u8, role, accepted)) return; + } + return error.Template; +} + +pub fn validate(messages: Value, problem: *[]const u8) Error!void { + if (messages != .array or messages.array.len == 0) { + problem.* = "DeepSeek V4.1 messages must be a nonempty array"; + return error.Template; + } + for (messages.array) |message| validateMessage(message) catch |err| { + problem.* = "DeepSeek V4.1 requires text-only messages with a supported role"; + return err; + }; +} + +/// Preserve multipart separators, leading/mid systems and metadata; only normalize API aliases and call arguments. +pub fn prepare(a: Allocator, messages: Value, problem: *[]const u8) Error!Value { + return prepareImpl(a, messages) catch |err| { + if (err == error.Template) + problem.* = "DeepSeek V4.1 requires text-only messages with supported roles; images, audio and video are unsupported"; + return err; + }; +} + +fn prepareImpl(a: Allocator, messages: Value) Error!Value { + if (messages != .array or messages.array.len == 0) return error.Template; + const list = try a.alloc(Value, messages.array.len); + for (messages.array, list) |message, *dst| { + try validateMessage(message); + const role = message.strField("role").?; + const copy = try json.copyObject(a, message.object); + if (std.mem.eql(u8, role, "developer")) try copy.put(a, "role", .{ .string = "system" }); + if (message.field("content_blocks")) |blocks| { + try validText(blocks); + try copy.put(a, "content", blocks); + } + if (message.field("tool_calls")) |calls| { + if (calls != .array) return error.Template; + const changed = try a.alloc(Value, calls.array.len); + for (calls.array, changed) |call, *slot| { + const function = call.get("function") orelse return error.Template; + if (function != .object or function.strField("name") == null) return error.Template; + const original = function.get("arguments") orelse return error.Template; + var arguments = original; + for (0..2) |_| { + if (arguments != .string) break; + arguments = switch (try json.parseText(a, arguments.string)) { + .ok => |v| v, + .err => break, + }; + } + if (arguments != .object) { + const wrapped = try json.newObject(a); + try wrapped.put(a, "arguments", original); + arguments = .{ .object = wrapped }; + } + const fn_copy = try json.copyObject(a, function.object); + try fn_copy.put(a, "arguments", arguments); + const call_copy = try json.copyObject(a, call.object); + try call_copy.put(a, "function", .{ .object = fn_copy }); + slot.* = .{ .object = call_copy }; + } + try copy.put(a, "tool_calls", .{ .array = changed }); + } + dst.* = .{ .object = copy }; + } + return .{ .array = list }; +} + +/// Keep official aliases; server xhigh means max, and templates without medium use high. +pub fn effort(a: Allocator, source: []const u8, input: []const u8) Error!std.json.Value { + if (std.fmt.parseInt(i64, input, 10)) |n| { + if (n < 1 or n > 100) return error.Template; + return .{ .integer = n }; + } else |_| {} + if (std.mem.eql(u8, input, "xhigh")) return .{ .string = "max" }; + if (std.mem.eql(u8, input, "minimal")) return .{ .string = "low" }; + if (std.mem.eql(u8, input, "medium") and std.mem.indexOf(u8, source, "\"medium\"") == null) return .{ .string = "high" }; + if (std.mem.eql(u8, input, "none")) return .{ .string = "high" }; + _ = a; + return .{ .string = input }; +} diff --git a/zig/src/server/deepseek_v41_text_goldens.json b/zig/src/server/deepseek_v41_text_goldens.json new file mode 100644 index 0000000000..b53a09de70 --- /dev/null +++ b/zig/src/server/deepseek_v41_text_goldens.json @@ -0,0 +1,2465 @@ +{ + "reference": "deepseek-ai/DeepSeek-V4.1-Flash", + "revision": "2cba9e42aa026125f3ed06c6d98c1db82f7ca027", + "template_sha256": "d959d804d5101b79a49b6ff1bf3c54cd5affa9a5a78c503d925f3090040b31c3", + "tokenizer_sha256": "c90dfa01249db1be4245780a052ede752e1361c612ac6d08e2bdada7d599476b", + "cases": [ + { + "name": "chat_user", + "messages": [ + { + "role": "user", + "content": "hello" + } + ], + "thinking": false, + "effort": null, + "tools": null, + "prompt": "<｜begin▁of▁sentence｜><｜User｜>hello<｜Assistant｜>", + "ids": [ + 0, + 128803, + 33310, + 128804, + 128822 + ] + }, + { + "name": "leading_system", + "messages": [ + { + "role": "system", + "content": "You are helpful." + }, + { + "role": "user", + "content": "你好" + } + ], + "thinking": false, + "effort": null, + "tools": null, + "prompt": "<｜begin▁of▁sentence｜><｜System｜>You are helpful.<｜User｜>你好<｜Assistant｜>", + "ids": [ + 0, + 128799, + 3476, + 477, + 11502, + 16, + 128803, + 30594, + 128804, + 128822 + ] + }, + { + "name": "thinking_None", + "messages": [ + { + "role": "user", + "content": "question" + } + ], + "thinking": true, + "effort": null, + "tools": null, + "prompt": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>question<｜Assistant｜>", + "ids": [ + 0, + 128799, + 77666, + 288, + 9504, + 482, + 28, + 223, + 2402, + 343, + 16801, + 223, + 19, + 15, + 1457, + 14, + 270, + 4077, + 270, + 1990, + 14, + 270, + 850, + 15146, + 270, + 22805, + 868, + 128803, + 20072, + 128804, + 128821 + ] + }, + { + "name": "thinking_low", + "messages": [ + { + "role": "user", + "content": "question" + } + ], + "thinking": true, + "effort": "low", + "tools": null, + "prompt": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 50 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>question<｜Assistant｜>", + "ids": [ + 0, + 128799, + 77666, + 288, + 9504, + 482, + 28, + 223, + 1328, + 343, + 16801, + 223, + 19, + 15, + 1457, + 14, + 270, + 4077, + 270, + 1990, + 14, + 270, + 850, + 15146, + 270, + 22805, + 868, + 128803, + 20072, + 128804, + 128821 + ] + }, + { + "name": "thinking_high", + "messages": [ + { + "role": "user", + "content": "question" + } + ], + "thinking": true, + "effort": "high", + "tools": null, + "prompt": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>question<｜Assistant｜>", + "ids": [ + 0, + 128799, + 77666, + 288, + 9504, + 482, + 28, + 223, + 2402, + 343, + 16801, + 223, + 19, + 15, + 1457, + 14, + 270, + 4077, + 270, + 1990, + 14, + 270, + 850, + 15146, + 270, + 22805, + 868, + 128803, + 20072, + 128804, + 128821 + ] + }, + { + "name": "thinking_max", + "messages": [ + { + "role": "user", + "content": "question" + } + ], + "thinking": true, + "effort": "max", + "tools": null, + "prompt": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 100 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>question<｜Assistant｜>", + "ids": [ + 0, + 128799, + 77666, + 288, + 9504, + 482, + 28, + 223, + 1457, + 343, + 16801, + 223, + 19, + 15, + 1457, + 14, + 270, + 4077, + 270, + 1990, + 14, + 270, + 850, + 15146, + 270, + 22805, + 868, + 128803, + 20072, + 128804, + 128821 + ] + }, + { + "name": "thinking_1", + "messages": [ + { + "role": "user", + "content": "question" + } + ], + "thinking": true, + "effort": 1, + "tools": null, + "prompt": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 1 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>question<｜Assistant｜>", + "ids": [ + 0, + 128799, + 77666, + 288, + 9504, + 482, + 28, + 223, + 19, + 343, + 16801, + 223, + 19, + 15, + 1457, + 14, + 270, + 4077, + 270, + 1990, + 14, + 270, + 850, + 15146, + 270, + 22805, + 868, + 128803, + 20072, + 128804, + 128821 + ] + }, + { + "name": "thinking_42", + "messages": [ + { + "role": "user", + "content": "question" + } + ], + "thinking": true, + "effort": 42, + "tools": null, + "prompt": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 42 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>question<｜Assistant｜>", + "ids": [ + 0, + 128799, + 77666, + 288, + 9504, + 482, + 28, + 223, + 3180, + 343, + 16801, + 223, + 19, + 15, + 1457, + 14, + 270, + 4077, + 270, + 1990, + 14, + 270, + 850, + 15146, + 270, + 22805, + 868, + 128803, + 20072, + 128804, + 128821 + ] + }, + { + "name": "thinking_100", + "messages": [ + { + "role": "user", + "content": "question" + } + ], + "thinking": true, + "effort": 100, + "tools": null, + "prompt": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 100 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>question<｜Assistant｜>", + "ids": [ + 0, + 128799, + 77666, + 288, + 9504, + 482, + 28, + 223, + 1457, + 343, + 16801, + 223, + 19, + 15, + 1457, + 14, + 270, + 4077, + 270, + 1990, + 14, + 270, + 850, + 15146, + 270, + 22805, + 868, + 128803, + 20072, + 128804, + 128821 + ] + }, + { + "name": "mid_system", + "messages": [ + { + "role": "system", + "content": "sys" + }, + { + "role": "user", + "content": "q1" + }, + { + "role": "assistant", + "content": "a1", + "reasoning_content": "r1" + }, + { + "role": "system", + "content": "mid sys" + } + ], + "thinking": true, + "effort": 88, + "tools": null, + "prompt": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 88 (range 1-100, the higher the value, the more thorough the reasoning)\n\nsys<｜User｜>q1<｜Assistant｜>a1<｜end▁of▁sentence｜><｜System｜>mid sys<｜Assistant｜>", + "ids": [ + 0, + 128799, + 77666, + 288, + 9504, + 482, + 28, + 223, + 2851, + 343, + 16801, + 223, + 19, + 15, + 1457, + 14, + 270, + 4077, + 270, + 1990, + 14, + 270, + 850, + 15146, + 270, + 22805, + 868, + 44489, + 128803, + 83, + 19, + 128804, + 128822, + 67, + 19, + 1, + 128799, + 14062, + 22832, + 128804, + 128821 + ] + }, + { + "name": "multipart_and_consecutive_users", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "one" + }, + { + "type": "text", + "text": "two" + } + ] + }, + { + "role": "user", + "content": "three" + } + ], + "thinking": false, + "effort": null, + "tools": null, + "prompt": "<｜begin▁of▁sentence｜><｜User｜>one\n\ntwo\n\nthree<｜Assistant｜>", + "ids": [ + 0, + 128803, + 791, + 271, + 23315, + 271, + 30230, + 128804, + 128822 + ] + }, + { + "name": "strip_old_reasoning", + "messages": [ + { + "role": "user", + "content": "q1" + }, + { + "role": "assistant", + "content": "a1", + "reasoning_content": "private reasoning" + }, + { + "role": "user", + "content": "q2" + }, + { + "role": "assistant", + "content": "a2", + "reasoning_content": "keep this" + } + ], + "thinking": true, + "effort": null, + "tools": null, + "prompt": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>q1<｜Assistant｜>a1<｜end▁of▁sentence｜><｜User｜>q2<｜Assistant｜>keep thisa2<｜end▁of▁sentence｜>", + "ids": [ + 0, + 128799, + 77666, + 288, + 9504, + 482, + 28, + 223, + 2402, + 343, + 16801, + 223, + 19, + 15, + 1457, + 14, + 270, + 4077, + 270, + 1990, + 14, + 270, + 850, + 15146, + 270, + 22805, + 868, + 128803, + 83, + 19, + 128804, + 128822, + 67, + 19, + 1, + 128803, + 83, + 20, + 128804, + 128821, + 38140, + 566, + 128822, + 67, + 20, + 1 + ] + }, + { + "name": "latest_reminder", + "messages": [ + { + "role": "user", + "content": "q" + }, + { + "role": "latest_reminder", + "content": "Today is Monday." + }, + { + "role": "assistant", + "content": "a" + } + ], + "thinking": false, + "effort": null, + "tools": null, + "prompt": "<｜begin▁of▁sentence｜><｜User｜>q<｜Assistant｜><｜latest_reminder｜>Today is Monday.a<｜end▁of▁sentence｜>", + "ids": [ + 0, + 128803, + 83, + 128804, + 128822, + 128828, + 19610, + 344, + 9219, + 16323, + 1 + ] + }, + { + "name": "task_action", + "messages": [ + { + "role": "user", + "content": "classify", + "task": "action" + }, + { + "role": "assistant", + "content": "result" + } + ], + "thinking": true, + "effort": null, + "tools": null, + "prompt": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>classify<｜Assistant｜><｜action｜>result<｜end▁of▁sentence｜>", + "ids": [ + 0, + 128799, + 77666, + 288, + 9504, + 482, + 28, + 223, + 2402, + 343, + 16801, + 223, + 19, + 15, + 1457, + 14, + 270, + 4077, + 270, + 1990, + 14, + 270, + 850, + 15146, + 270, + 22805, + 868, + 128803, + 3767, + 2196, + 128804, + 128821, + 128829, + 9001, + 1 + ] + }, + { + "name": "task_query", + "messages": [ + { + "role": "user", + "content": "classify", + "task": "query" + }, + { + "role": "assistant", + "content": "result" + } + ], + "thinking": true, + "effort": null, + "tools": null, + "prompt": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>classify<｜query｜>result<｜end▁of▁sentence｜>", + "ids": [ + 0, + 128799, + 77666, + 288, + 9504, + 482, + 28, + 223, + 2402, + 343, + 16801, + 223, + 19, + 15, + 1457, + 14, + 270, + 4077, + 270, + 1990, + 14, + 270, + 850, + 15146, + 270, + 22805, + 868, + 128803, + 3767, + 2196, + 128830, + 9001, + 1 + ] + }, + { + "name": "task_authority", + "messages": [ + { + "role": "user", + "content": "classify", + "task": "authority" + }, + { + "role": "assistant", + "content": "result" + } + ], + "thinking": true, + "effort": null, + "tools": null, + "prompt": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>classify<｜authority｜>result<｜end▁of▁sentence｜>", + "ids": [ + 0, + 128799, + 77666, + 288, + 9504, + 482, + 28, + 223, + 2402, + 343, + 16801, + 223, + 19, + 15, + 1457, + 14, + 270, + 4077, + 270, + 1990, + 14, + 270, + 850, + 15146, + 270, + 22805, + 868, + 128803, + 3767, + 2196, + 128831, + 9001, + 1 + ] + }, + { + "name": "task_domain", + "messages": [ + { + "role": "user", + "content": "classify", + "task": "domain" + }, + { + "role": "assistant", + "content": "result" + } + ], + "thinking": true, + "effort": null, + "tools": null, + "prompt": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>classify<｜domain｜>result<｜end▁of▁sentence｜>", + "ids": [ + 0, + 128799, + 77666, + 288, + 9504, + 482, + 28, + 223, + 2402, + 343, + 16801, + 223, + 19, + 15, + 1457, + 14, + 270, + 4077, + 270, + 1990, + 14, + 270, + 850, + 15146, + 270, + 22805, + 868, + 128803, + 3767, + 2196, + 128832, + 9001, + 1 + ] + }, + { + "name": "task_title", + "messages": [ + { + "role": "user", + "content": "classify", + "task": "title" + }, + { + "role": "assistant", + "content": "result" + } + ], + "thinking": true, + "effort": null, + "tools": null, + "prompt": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>classify<｜title｜>result<｜end▁of▁sentence｜>", + "ids": [ + 0, + 128799, + 77666, + 288, + 9504, + 482, + 28, + 223, + 2402, + 343, + 16801, + 223, + 19, + 15, + 1457, + 14, + 270, + 4077, + 270, + 1990, + 14, + 270, + 850, + 15146, + 270, + 22805, + 868, + 128803, + 3767, + 2196, + 128836, + 9001, + 1 + ] + }, + { + "name": "task_read_url", + "messages": [ + { + "role": "user", + "content": "classify", + "task": "read_url" + }, + { + "role": "assistant", + "content": "result" + } + ], + "thinking": true, + "effort": null, + "tools": null, + "prompt": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>classify<｜read_url｜>result<｜end▁of▁sentence｜>", + "ids": [ + 0, + 128799, + 77666, + 288, + 9504, + 482, + 28, + 223, + 2402, + 343, + 16801, + 223, + 19, + 15, + 1457, + 14, + 270, + 4077, + 270, + 1990, + 14, + 270, + 850, + 15146, + 270, + 22805, + 868, + 128803, + 3767, + 2196, + 128845, + 9001, + 1 + ] + }, + { + "name": "wo_eos", + "messages": [ + { + "role": "user", + "content": "q" + }, + { + "role": "assistant", + "content": "partial", + "reasoning_content": "r", + "wo_eos": true + } + ], + "thinking": true, + "effort": null, + "tools": null, + "prompt": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>q<｜Assistant｜>rpartial", + "ids": [ + 0, + 128799, + 77666, + 288, + 9504, + 482, + 28, + 223, + 2402, + 343, + 16801, + 223, + 19, + 15, + 1457, + 14, + 270, + 4077, + 270, + 1990, + 14, + 270, + 850, + 15146, + 270, + 22805, + 868, + 128803, + 83, + 128804, + 128821, + 84, + 128822, + 14004 + ] + }, + { + "name": "response_format", + "messages": [ + { + "role": "system", + "content": "sys", + "response_format": { + "type": "object", + "properties": { + "值": { + "type": "integer" + } + } + } + }, + { + "role": "user", + "content": "q" + } + ], + "thinking": false, + "effort": null, + "tools": null, + "prompt": "<｜begin▁of▁sentence｜><｜System｜>sys\n\n## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n{\"type\": \"object\", \"properties\": {\"值\": {\"type\": \"integer\"}}}<｜User｜>q<｜Assistant｜>", + "ids": [ + 0, + 128799, + 44489, + 271, + 372, + 17038, + 29478, + 979, + 3476, + 74366, + 29851, + 42917, + 304, + 270, + 2502, + 37913, + 304, + 20430, + 1137, + 24313, + 4611, + 3362, + 582, + 10325, + 1760, + 582, + 68838, + 3362, + 28612, + 2022, + 3362, + 28612, + 4611, + 3362, + 582, + 39648, + 4, + 41653, + 128803, + 83, + 128804, + 128822 + ] + }, + { + "name": "namespace_multiturn_tools", + "messages": [ + { + "role": "user", + "content": "question" + }, + { + "role": "assistant", + "reasoning_content": "retained reasoning", + "content": "summary", + "tool_calls": [ + { + "id": "first", + "type": "function", + "namespace": "search", + "function": { + "name": "lookup", + "arguments": "{\"query\":\"天气\\nParis\", \"limit\":2, \"data\":[1,true,null], \"obj\":{\"x\":1}}" + } + }, + { + "id": "second", + "type": "function", + "function": { + "name": "files::lookup", + "arguments": "\"{\\\"query\\\":\\\"value\\\"}\"" + } + } + ] + }, + { + "role": "tool", + "tool_call_id": "second", + "content": "second result" + }, + { + "role": "tool", + "tool_call_id": "first", + "content": "first result" + }, + { + "role": "assistant", + "reasoning_content": "continue reasoning", + "content": "done" + } + ], + "thinking": true, + "effort": null, + "tools": [ + { + "type": "function", + "namespace": { + "name": "search", + "description": "Search tools." + }, + "function": { + "name": "lookup", + "description": "Look up a value", + "parameters": { + "type": "object", + "properties": { + "query": { + "type": "string" + }, + "limit": { + "type": "integer" + } + } + } + } + }, + { + "type": "function", + "namespace": "files", + "function": { + "name": "lookup", + "parameters": { + "type": "object" + } + } + } + ], + "prompt": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)\n\n\n\n## Tools\n\nYou have access to a set of tools to help answer the user's question. You can invoke tools by writing a \"<｜DSML｜ calls>\" block like the following:\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"$TOOL_NAME\">\n<｜DSML｜ parameter name=\"$PARAMETER_NAME\" string=\"true|false\">$PARAMETER_VALUE\n...\n\n<｜DSML｜ invoke name=\"$TOOL_NAME2\">\n...\n\n\n\nString parameters should be specified as is and set `string=\"true\"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string=\"false\"`.\n\nIf thinking_mode is enabled (triggered by ), you MUST output your complete reasoning inside ... BEFORE any tool calls or final response.\n\nOtherwise, output directly after with tool calls or final response.\n\n### Available Tool Schemas\n\n{\"name\": \"search::lookup\", \"description\": \"Search tools.\\nLook up a value\", \"parameters\": {\"type\": \"object\", \"properties\": {\"query\": {\"type\": \"string\"}, \"limit\": {\"type\": \"integer\"}}}}\n{\"name\": \"files::lookup\", \"parameters\": {\"type\": \"object\"}}\n\nYou MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.\n<｜User｜>question<｜Assistant｜>retained reasoningsummary\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"search::lookup\">\n<｜DSML｜ parameter name=\"query\" string=\"true\">天气\nParis\n<｜DSML｜ parameter name=\"limit\" string=\"false\">2\n<｜DSML｜ parameter name=\"data\" string=\"false\">[1, true, null]\n<｜DSML｜ parameter name=\"obj\" string=\"false\">{\"x\": 1}\n\n<｜DSML｜ invoke name=\"files::lookup\">\n<｜DSML｜ parameter name=\"query\" string=\"true\">value\n\n<｜end▁of▁sentence｜><｜User｜>first result\n\nsecond result<｜Assistant｜>continue reasoningdone<｜end▁of▁sentence｜>", + "ids": [ + 0, + 128799, + 77666, + 288, + 9504, + 482, + 28, + 223, + 2402, + 343, + 16801, + 223, + 19, + 15, + 1457, + 14, + 270, + 4077, + 270, + 1990, + 14, + 270, + 850, + 15146, + 270, + 22805, + 86262, + 372, + 27193, + 271, + 3476, + 611, + 3278, + 304, + 260, + 1341, + 294, + 6704, + 304, + 1694, + 3287, + 270, + 3967, + 734, + 3417, + 16, + 2042, + 588, + 34756, + 6704, + 513, + 4985, + 260, + 38322, + 128825, + 10699, + 48902, + 5603, + 1277, + 270, + 2502, + 979, + 30, + 128825, + 10699, + 1018, + 30, + 128825, + 34756, + 2329, + 84476, + 6792, + 5896, + 35044, + 3816, + 30, + 128825, + 10767, + 2329, + 84476, + 46630, + 4130, + 61729, + 35044, + 4, + 3418, + 1281, + 11476, + 94, + 19836, + 3320, + 6, + 46630, + 4130, + 61729, + 56909, + 1718, + 128825, + 10767, + 1018, + 7835, + 1718, + 128825, + 34756, + 1018, + 30, + 128825, + 34756, + 2329, + 84476, + 6792, + 5896, + 35044, + 20, + 3816, + 7835, + 1718, + 128825, + 34756, + 1018, + 1718, + 128825, + 10699, + 4697, + 3524, + 8252, + 1531, + 366, + 12038, + 412, + 344, + 305, + 1341, + 3608, + 4463, + 1281, + 11476, + 4, + 37419, + 1884, + 710, + 915, + 4815, + 343, + 62896, + 14, + 2631, + 2065, + 634, + 14, + 31939, + 14, + 8435, + 754, + 2281, + 270, + 1990, + 295, + 26639, + 8786, + 305, + 1341, + 3608, + 4463, + 1281, + 19836, + 4, + 108526, + 3575, + 6892, + 71144, + 344, + 22104, + 343, + 89944, + 284, + 513, + 223, + 128821, + 754, + 440, + 74366, + 5238, + 782, + 5553, + 22805, + 6352, + 223, + 128821, + 1613, + 128822, + 119907, + 1117, + 4105, + 10699, + 469, + 4087, + 4256, + 339, + 13079, + 6922, + 14, + 5238, + 6578, + 1561, + 223, + 128822, + 418, + 4105, + 10699, + 469, + 4087, + 4256, + 339, + 795, + 17829, + 28249, + 13178, + 8380, + 271, + 24313, + 2852, + 3362, + 582, + 14346, + 2366, + 12533, + 1425, + 1760, + 582, + 20855, + 3362, + 582, + 12386, + 6704, + 10276, + 80, + 18948, + 890, + 260, + 1990, + 1760, + 582, + 46172, + 3362, + 28612, + 4611, + 3362, + 582, + 10325, + 1760, + 582, + 68838, + 3362, + 28612, + 12066, + 3362, + 28612, + 4611, + 3362, + 582, + 4463, + 87928, + 582, + 39673, + 3362, + 28612, + 4611, + 3362, + 582, + 39648, + 4, + 3831, + 30316, + 24313, + 2852, + 3362, + 582, + 44585, + 2366, + 12533, + 1425, + 1760, + 582, + 46172, + 3362, + 28612, + 4611, + 3362, + 582, + 10325, + 4, + 33236, + 3476, + 74366, + 29851, + 1605, + 270, + 3554, + 6428, + 4105, + 2329, + 305, + 10767, + 5815, + 8380, + 304, + 34756, + 4105, + 10699, + 603, + 128803, + 20072, + 128804, + 128821, + 2101, + 3659, + 22805, + 128822, + 11235, + 271, + 30, + 128825, + 10699, + 1018, + 30, + 128825, + 34756, + 2329, + 1281, + 14346, + 2366, + 12533, + 1425, + 3816, + 30, + 128825, + 10767, + 2329, + 1281, + 12066, + 4, + 3418, + 1281, + 11476, + 3320, + 16652, + 201, + 51119, + 1718, + 128825, + 10767, + 1018, + 30, + 128825, + 10767, + 2329, + 1281, + 39673, + 4, + 3418, + 1281, + 19836, + 3320, + 20, + 1718, + 128825, + 10767, + 1018, + 30, + 128825, + 10767, + 2329, + 1281, + 5263, + 4, + 3418, + 1281, + 19836, + 3320, + 61, + 19, + 14, + 3292, + 14, + 5100, + 63, + 1718, + 128825, + 10767, + 1018, + 30, + 128825, + 10767, + 2329, + 1281, + 17371, + 4, + 3418, + 1281, + 19836, + 3320, + 24313, + 90, + 3362, + 223, + 19, + 24568, + 128825, + 10767, + 1018, + 1718, + 128825, + 34756, + 1018, + 30, + 128825, + 34756, + 2329, + 1281, + 44585, + 2366, + 12533, + 1425, + 3816, + 30, + 128825, + 10767, + 2329, + 1281, + 12066, + 4, + 3418, + 1281, + 11476, + 3320, + 5197, + 1718, + 128825, + 10767, + 1018, + 1718, + 128825, + 34756, + 1018, + 1718, + 128825, + 10699, + 32, + 1, + 128803, + 30, + 72461, + 46148, + 32, + 13213, + 1529, + 1718, + 72461, + 46148, + 4697, + 30, + 72461, + 46148, + 32, + 10930, + 1529, + 1718, + 72461, + 46148, + 32, + 128804, + 128821, + 45911, + 22805, + 128822, + 44629, + 1 + ] + }, + { + "name": "fallback_arguments_not json", + "messages": [ + { + "role": "user", + "content": "q" + }, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "type": "function", + "function": { + "name": "lookup", + "arguments": "not json" + } + } + ] + } + ], + "thinking": false, + "effort": null, + "tools": null, + "prompt": "<｜begin▁of▁sentence｜><｜User｜>q<｜Assistant｜>\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"lookup\">\n<｜DSML｜ parameter name=\"arguments\" string=\"true\">not json\n\n<｜end▁of▁sentence｜>", + "ids": [ + 0, + 128803, + 83, + 128804, + 128822, + 271, + 30, + 128825, + 10699, + 1018, + 30, + 128825, + 34756, + 2329, + 1281, + 12533, + 1425, + 3816, + 30, + 128825, + 10767, + 2329, + 1281, + 83772, + 4, + 3418, + 1281, + 11476, + 3320, + 2869, + 18786, + 1718, + 128825, + 10767, + 1018, + 1718, + 128825, + 34756, + 1018, + 1718, + 128825, + 10699, + 32, + 1 + ] + }, + { + "name": "fallback_arguments_[1,2]", + "messages": [ + { + "role": "user", + "content": "q" + }, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "type": "function", + "function": { + "name": "lookup", + "arguments": "[1,2]" + } + } + ] + } + ], + "thinking": false, + "effort": null, + "tools": null, + "prompt": "<｜begin▁of▁sentence｜><｜User｜>q<｜Assistant｜>\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"lookup\">\n<｜DSML｜ parameter name=\"arguments\" string=\"true\">[1,2]\n\n<｜end▁of▁sentence｜>", + "ids": [ + 0, + 128803, + 83, + 128804, + 128822, + 271, + 30, + 128825, + 10699, + 1018, + 30, + 128825, + 34756, + 2329, + 1281, + 12533, + 1425, + 3816, + 30, + 128825, + 10767, + 2329, + 1281, + 83772, + 4, + 3418, + 1281, + 11476, + 3320, + 61, + 19, + 14, + 20, + 63, + 1718, + 128825, + 10767, + 1018, + 1718, + 128825, + 34756, + 1018, + 1718, + 128825, + 10699, + 32, + 1 + ] + }, + { + "name": "fallback_arguments_\"plain\"", + "messages": [ + { + "role": "user", + "content": "q" + }, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "type": "function", + "function": { + "name": "lookup", + "arguments": "\"plain\"" + } + } + ] + } + ], + "thinking": false, + "effort": null, + "tools": null, + "prompt": "<｜begin▁of▁sentence｜><｜User｜>q<｜Assistant｜>\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"lookup\">\n<｜DSML｜ parameter name=\"arguments\" string=\"true\">\"plain\"\n\n<｜end▁of▁sentence｜>", + "ids": [ + 0, + 128803, + 83, + 128804, + 128822, + 271, + 30, + 128825, + 10699, + 1018, + 30, + 128825, + 34756, + 2329, + 1281, + 12533, + 1425, + 3816, + 30, + 128825, + 10767, + 2329, + 1281, + 83772, + 4, + 3418, + 1281, + 11476, + 3320, + 4, + 38938, + 4, + 1718, + 128825, + 10767, + 1018, + 1718, + 128825, + 34756, + 1018, + 1718, + 128825, + 10699, + 32, + 1 + ] + }, + { + "name": "official_case_1", + "messages": [ + { + "role": "system", + "content": "You are a helpful assistant." + }, + { + "role": "user", + "content": "What's the weather like in Beijing?" + }, + { + "role": "assistant", + "reasoning_content": "The user wants the weather in Beijing. I should call get_weather.", + "content": "", + "tool_calls": [ + { + "type": "function", + "function": { + "name": "get_weather", + "arguments": "{\"location\": \"Beijing\", \"unit\": \"celsius\"}" + } + } + ] + }, + { + "role": "tool", + "tool_call_id": "call_0", + "content": "{\"temperature\": 22, \"condition\": \"sunny\", \"humidity\": 45}" + }, + { + "role": "assistant", + "reasoning_content": "Got the weather data. Let me format a nice response.", + "content": "The weather in Beijing is currently sunny with a temperature of 22°C and 45% humidity." + } + ], + "thinking": true, + "effort": null, + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather for a specific location", + "parameters": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "The city name" + }, + "unit": { + "type": "string", + "enum": [ + "celsius", + "fahrenheit" + ], + "description": "Temperature unit" + } + }, + "required": [ + "location" + ] + } + } + }, + { + "type": "function", + "function": { + "name": "search", + "description": "Search the web for information", + "parameters": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "Search query" + }, + "num_results": { + "type": "integer", + "description": "Number of results to return" + } + }, + "required": [ + "query" + ] + } + } + } + ], + "prompt": "<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)\n\nYou are a helpful assistant.\n\n## Tools\n\nYou have access to a set of tools to help answer the user's question. You can invoke tools by writing a \"<｜DSML｜ calls>\" block like the following:\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"$TOOL_NAME\">\n<｜DSML｜ parameter name=\"$PARAMETER_NAME\" string=\"true|false\">$PARAMETER_VALUE\n...\n\n<｜DSML｜ invoke name=\"$TOOL_NAME2\">\n...\n\n\n\nString parameters should be specified as is and set `string=\"true\"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string=\"false\"`.\n\nIf thinking_mode is enabled (triggered by ), you MUST output your complete reasoning inside ... BEFORE any tool calls or final response.\n\nOtherwise, output directly after with tool calls or final response.\n\n### Available Tool Schemas\n\n{\"name\": \"get_weather\", \"description\": \"Get the weather for a specific location\", \"parameters\": {\"type\": \"object\", \"properties\": {\"location\": {\"type\": \"string\", \"description\": \"The city name\"}, \"unit\": {\"type\": \"string\", \"enum\": [\"celsius\", \"fahrenheit\"], \"description\": \"Temperature unit\"}}, \"required\": [\"location\"]}}\n{\"name\": \"search\", \"description\": \"Search the web for information\", \"parameters\": {\"type\": \"object\", \"properties\": {\"query\": {\"type\": \"string\", \"description\": \"Search query\"}, \"num_results\": {\"type\": \"integer\", \"description\": \"Number of results to return\"}}, \"required\": [\"query\"]}}\n\nYou MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.\n<｜User｜>What's the weather like in Beijing?<｜Assistant｜>The user wants the weather in Beijing. I should call get_weather.\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"get_weather\">\n<｜DSML｜ parameter name=\"location\" string=\"true\">Beijing\n<｜DSML｜ parameter name=\"unit\" string=\"true\">celsius\n\n<｜end▁of▁sentence｜><｜User｜>{\"temperature\": 22, \"condition\": \"sunny\", \"humidity\": 45}<｜Assistant｜>Got the weather data. Let me format a nice response.The weather in Beijing is currently sunny with a temperature of 22°C and 45% humidity.<｜end▁of▁sentence｜>", + "ids": [ + 0, + 128799, + 77666, + 288, + 9504, + 482, + 28, + 223, + 2402, + 343, + 16801, + 223, + 19, + 15, + 1457, + 14, + 270, + 4077, + 270, + 1990, + 14, + 270, + 850, + 15146, + 270, + 22805, + 868, + 3476, + 477, + 260, + 11502, + 22896, + 339, + 372, + 27193, + 271, + 3476, + 611, + 3278, + 304, + 260, + 1341, + 294, + 6704, + 304, + 1694, + 3287, + 270, + 3967, + 734, + 3417, + 16, + 2042, + 588, + 34756, + 6704, + 513, + 4985, + 260, + 38322, + 128825, + 10699, + 48902, + 5603, + 1277, + 270, + 2502, + 979, + 30, + 128825, + 10699, + 1018, + 30, + 128825, + 34756, + 2329, + 84476, + 6792, + 5896, + 35044, + 3816, + 30, + 128825, + 10767, + 2329, + 84476, + 46630, + 4130, + 61729, + 35044, + 4, + 3418, + 1281, + 11476, + 94, + 19836, + 3320, + 6, + 46630, + 4130, + 61729, + 56909, + 1718, + 128825, + 10767, + 1018, + 7835, + 1718, + 128825, + 34756, + 1018, + 30, + 128825, + 34756, + 2329, + 84476, + 6792, + 5896, + 35044, + 20, + 3816, + 7835, + 1718, + 128825, + 34756, + 1018, + 1718, + 128825, + 10699, + 4697, + 3524, + 8252, + 1531, + 366, + 12038, + 412, + 344, + 305, + 1341, + 3608, + 4463, + 1281, + 11476, + 4, + 37419, + 1884, + 710, + 915, + 4815, + 343, + 62896, + 14, + 2631, + 2065, + 634, + 14, + 31939, + 14, + 8435, + 754, + 2281, + 270, + 1990, + 295, + 26639, + 8786, + 305, + 1341, + 3608, + 4463, + 1281, + 19836, + 4, + 108526, + 3575, + 6892, + 71144, + 344, + 22104, + 343, + 89944, + 284, + 513, + 223, + 128821, + 754, + 440, + 74366, + 5238, + 782, + 5553, + 22805, + 6352, + 223, + 128821, + 1613, + 128822, + 119907, + 1117, + 4105, + 10699, + 469, + 4087, + 4256, + 339, + 13079, + 6922, + 14, + 5238, + 6578, + 1561, + 223, + 128822, + 418, + 4105, + 10699, + 469, + 4087, + 4256, + 339, + 795, + 17829, + 28249, + 13178, + 8380, + 271, + 24313, + 2852, + 3362, + 582, + 1133, + 65, + 50219, + 1760, + 582, + 20855, + 3362, + 582, + 6287, + 270, + 9670, + 362, + 260, + 3549, + 6830, + 1760, + 582, + 46172, + 3362, + 28612, + 4611, + 3362, + 582, + 10325, + 1760, + 582, + 68838, + 3362, + 28612, + 33182, + 3362, + 28612, + 4611, + 3362, + 582, + 4463, + 1760, + 582, + 20855, + 3362, + 582, + 671, + 4593, + 2329, + 87928, + 582, + 15165, + 3362, + 28612, + 4611, + 3362, + 582, + 4463, + 1760, + 582, + 47002, + 3362, + 20584, + 69, + 33030, + 1760, + 582, + 72, + 37988, + 45390, + 582, + 20855, + 3362, + 582, + 78130, + 4761, + 4, + 55695, + 582, + 24486, + 3362, + 20584, + 33182, + 13747, + 30316, + 24313, + 2852, + 3362, + 582, + 14346, + 1760, + 582, + 20855, + 3362, + 582, + 12386, + 270, + 6545, + 362, + 1951, + 1760, + 582, + 46172, + 3362, + 28612, + 4611, + 3362, + 582, + 10325, + 1760, + 582, + 68838, + 3362, + 28612, + 12066, + 3362, + 28612, + 4611, + 3362, + 582, + 4463, + 1760, + 582, + 20855, + 3362, + 582, + 12386, + 13607, + 87928, + 582, + 10064, + 80757, + 3362, + 28612, + 4611, + 3362, + 582, + 39648, + 1760, + 582, + 20855, + 3362, + 582, + 7501, + 294, + 3083, + 304, + 1354, + 4, + 55695, + 582, + 24486, + 3362, + 20584, + 12066, + 13747, + 33236, + 3476, + 74366, + 29851, + 1605, + 270, + 3554, + 6428, + 4105, + 2329, + 305, + 10767, + 5815, + 8380, + 304, + 34756, + 4105, + 10699, + 603, + 128803, + 3085, + 734, + 270, + 9670, + 1277, + 295, + 29660, + 33, + 128804, + 128821, + 671, + 3967, + 10059, + 270, + 9670, + 295, + 29660, + 16, + 342, + 1531, + 1926, + 1178, + 65, + 50219, + 16, + 128822, + 271, + 30, + 128825, + 10699, + 1018, + 30, + 128825, + 34756, + 2329, + 1281, + 1133, + 65, + 50219, + 3816, + 30, + 128825, + 10767, + 2329, + 1281, + 33182, + 4, + 3418, + 1281, + 11476, + 3320, + 121458, + 1718, + 128825, + 10767, + 1018, + 30, + 128825, + 10767, + 2329, + 1281, + 15165, + 4, + 3418, + 1281, + 11476, + 3320, + 69, + 33030, + 1718, + 128825, + 10767, + 1018, + 1718, + 128825, + 34756, + 1018, + 1718, + 128825, + 10699, + 32, + 1, + 128803, + 30, + 72461, + 46148, + 32, + 24313, + 88634, + 3362, + 223, + 1302, + 14, + 582, + 22526, + 3362, + 582, + 69958, + 2928, + 1760, + 582, + 34859, + 13168, + 3362, + 223, + 1883, + 24568, + 72461, + 46148, + 32, + 128804, + 128821, + 70935, + 270, + 9670, + 1499, + 16, + 4480, + 678, + 8786, + 260, + 10722, + 4256, + 16, + 128822, + 671, + 9670, + 295, + 29660, + 344, + 7491, + 43905, + 418, + 260, + 5169, + 294, + 223, + 1302, + 2614, + 37, + 305, + 223, + 1883, + 7, + 27392, + 16, + 1 + ] + } + ] +} diff --git a/zig/src/server/deepseek_v41_text_test.zig b/zig/src/server/deepseek_v41_text_test.zig new file mode 100644 index 0000000000..0757c3a988 --- /dev/null +++ b/zig/src/server/deepseek_v41_text_test.zig @@ -0,0 +1,242 @@ +//! Deterministic official encoder/tokenizer goldens; pinned DeepSeek-AI provenance in THIRD_PARTY_NOTICES.md. +const std = @import("std"); +const json = @import("json.zig"); +const model_text = @import("model_text.zig"); +const HfText = @import("hf_text.zig").HfText; +const v41 = @import("deepseek_v41_text.zig"); +const tool_parse = @import("tool_parse.zig"); +const tool_stream = @import("tool_stream.zig"); +const reply_text = @import("reply_text.zig"); +const fields = @import("fields.zig"); +const errors = @import("errors.zig"); +const Value = json.Value; +const official = @embedFile("deepseek_v41_official.jinja"); +const golden_bytes = @embedFile("deepseek_v41_text_goldens.json"); + +fn parse(a: std.mem.Allocator, bytes: []const u8) !Value { + return (try json.parse(a, bytes)).ok; +} +fn options(c: Value) model_text.RenderOptions { + const effort = c.field("effort"); + return .{ + .enable_thinking = c.get("thinking").?.bool, + .reasoning_effort = if (effort) |e| (if (e == .int) e.int else e.string) else null, + .tools = c.field("tools"), + }; +} +fn tinyCheckpoint(dir: std.Io.Dir) !void { + const io = std.testing.io; + try dir.writeFile(io, .{ .sub_path = "tokenizer.json", .data = + \\{"model":{"type":"BPE","vocab":{"H":0},"merges":[]},"added_tokens":[]} + }); + try dir.writeFile(io, .{ .sub_path = "config.json", .data = "{\"model_type\":\"deepseek_v41\"}" }); + try dir.writeFile(io, .{ .sub_path = "chat_template.jinja", .data = official }); +} + +test "official V41 text reference goldens render through default checkpoint HfText" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try tinyCheckpoint(tmp.dir); + const dir = try std.fmt.allocPrint(a, ".zig-cache/tmp/{s}", .{tmp.sub_path}); + var problem: []const u8 = ""; + const text = try HfText.load(std.testing.allocator, std.testing.io, dir, null, a, &problem); + defer text.deinit(); + const cases = (try parse(a, golden_bytes)).get("cases").?.array; + for (cases) |c| { + const rendered = text.text().render(a, c.get("messages").?, options(c), &problem) catch |err| { + std.debug.print("V41 golden {s}: {s}\n", .{ c.strField("name").?, problem }); + return err; + }; + try std.testing.expectEqualStrings(c.strField("prompt").?, rendered); + } + const user = try parse(a, "[{\"role\":\"user\",\"content\":\"hello\"}]"); + try std.testing.expectEqualStrings("<｜begin▁of▁sentence｜><｜User｜>hello", try text.text().render(a, user, .{ .add_generation_prompt = false }, &problem)); + try tmp.dir.writeFile(std.testing.io, .{ .sub_path = "override.jinja", .data = "override{{ messages[0].content }}" }); + const path = try std.fmt.allocPrint(a, "{s}/override.jinja", .{dir}); + const override = try HfText.load(std.testing.allocator, std.testing.io, dir, path, a, &problem); + defer override.deinit(); + try std.testing.expectEqualStrings("overridehello", try override.text().render(a, user, .{}, &problem)); +} + +test "V41 real tokenizer IDs match official reference goldens when checkpoint is supplied" { + const dir = std.testing.environ.getPosix("TENSORFOLD_V41_TOKENIZER_DIR") orelse return error.SkipZigTest; + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var problem: []const u8 = ""; + const text = try HfText.load(std.testing.allocator, std.testing.io, dir, null, a, &problem); + defer text.deinit(); + for ((try parse(a, golden_bytes)).get("cases").?.array) |c| { + const ids = try text.text().renderIds(a, c.get("messages").?, options(c), &problem); + const expected = c.get("ids").?.array; + const wanted = try a.alloc(u32, expected.len); + for (expected, wanted) |id, *slot| slot.* = @intCast(id.int64().?); + try std.testing.expectEqualSlices(u32, wanted, ids); + } +} + +test "V41 numeric reasoning efforts are accepted without changing other templates" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var cx: errors.Cx = .{ .a = a }; + const levels = try fields.effortLevels(a, official); + const effort = try fields.thinkingFields(&cx, try parse(a, "{\"reasoning_effort\":42}"), levels); + try std.testing.expectEqualStrings("42", effort.effort.?); + try std.testing.expect(effort.enable.?); + for ([_][]const u8{ "0", "101", "true", "1.5" }) |value| { + const body = try parse(a, try std.mem.concat(a, u8, &.{ "{\"reasoning_effort\":", value, "}" })); + try std.testing.expectError(error.Refused, fields.thinkingFields(&cx, body, levels)); + } + try std.testing.expectError(error.Refused, fields.thinkingFields(&cx, try parse(a, "{\"reasoning_effort\":42}"), &.{ "low", "high" })); + const numeric = try v41.effort(a, official, "42"); + try std.testing.expectEqual(@as(i64, 42), numeric.integer); + try std.testing.expectEqualStrings("max", (try v41.effort(a, official, "xhigh")).string); + try std.testing.expectError(error.Template, v41.effort(a, official, "0")); + try std.testing.expectError(error.Template, v41.effort(a, official, "101")); +} + +test "V41 text preparation rejects media at every supported nesting boundary" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var problem: []const u8 = ""; + for ([_][]const u8{ + "[{\"role\":\"user\",\"content\":[{\"type\":\"image_url\",\"image_url\":{\"url\":\"image.png\"}}]}]", + "[{\"role\":\"tool\",\"content\":[{\"type\":\"image\",\"source\":{\"data\":\"x\"}}]}]", + "[{\"role\":\"user\",\"content\":\"image.png\"}]", + "[{\"role\":\"user\",\"content\":\"<｜deepseek_image｜>\"}]", + "[{\"role\":\"assistant\",\"content\":\"\",\"reasoning_content\":\"<｜deepseek_image｜>\"}]", + "[{\"role\":\"user\",\"content\":[{\"type\":\"input_audio\",\"input_audio\":{\"data\":\"x\"}}]}]", + "[{\"role\":\"user\",\"content_blocks\":[{\"type\":\"video_url\",\"video_url\":\"clip.mp4\"}]}]", + "[{\"role\":\"user\",\"content\":[{\"type\":\"tool_result\",\"content\":[{\"type\":\"image\",\"source\":\"x\"}]}]}]", + "[{\"role\":\"assistant\",\"content\":\"\",\"thinking\":\"image.png\"}]", + }) |input| try std.testing.expectError(error.Template, v41.prepare(a, try parse(a, input), &problem)); +} + +test "V41 image-named tool arguments remain ordinary typed tool data" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try tinyCheckpoint(tmp.dir); + const dir = try std.fmt.allocPrint(a, ".zig-cache/tmp/{s}", .{tmp.sub_path}); + var problem: []const u8 = ""; + const text = try HfText.load(std.testing.allocator, std.testing.io, dir, null, a, &problem); + defer text.deinit(); + const messages = try parse(a, "[{\"role\":\"user\",\"content\":\"inspect it\"}," ++ + "{\"role\":\"assistant\",\"content\":\"\",\"tool_calls\":[{\"function\":{\"name\":\"inspect\"," ++ + "\"arguments\":{\"image\":\"cat.jpg\",\"audio\":true}}}]}," ++ + "{\"role\":\"tool\",\"content\":\"done\"},{\"role\":\"user\",\"content\":\"continue\"}]"); + const rendered = try text.text().render(a, messages, .{ .enable_thinking = false }, &problem); + try std.testing.expect(std.mem.indexOf(u8, rendered, "<｜DSML｜ parameter name=\"image\" string=\"true\">cat.jpg") != null); + try std.testing.expect(std.mem.indexOf(u8, rendered, "<｜DSML｜ parameter name=\"audio\" string=\"false\">true") != null); +} + +const completion = "summary\n\n<｜DSML｜ calls>\n" ++ + "<｜DSML｜ invoke name=\"search::lookup\">\n" ++ + "<｜DSML｜ parameter name=\"query\" string=\"true\"> value\n\"quoted\" \n" ++ + "<｜DSML｜ parameter name=\"limit\" string=\"false\">2\n" ++ + "<｜DSML｜ parameter name=\"data\" string=\"false\">[true,null,{\"x\":1}]\n" ++ + "\n" ++ + "<｜DSML｜ invoke name=\"files::lookup\">\n\n" ++ + ""; +const offered = "[{\"type\":\"function\",\"namespace\":\"search\",\"function\":{\"name\":\"lookup\"}}," ++ + "{\"type\":\"function\",\"namespace\":{\"name\":\"files\"},\"function\":{\"name\":\"lookup\"}}]"; + +test "V41 leading-space DSML preserves namespace and typed parameter boundaries" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const tools = (try parse(a, offered)).array; + const parsed = try tool_parse.parse(a, completion, tools, null); + try std.testing.expectEqualStrings("summary", parsed.content); + try std.testing.expectEqual(@as(usize, 2), parsed.calls.?.len); + const first = parsed.calls.?[0]; + try std.testing.expectEqualStrings("search", first.strField("namespace").?); + try std.testing.expectEqualStrings("lookup", first.get("function").?.strField("name").?); + const arguments = try parse(a, first.get("function").?.strField("arguments").?); + try std.testing.expectEqualStrings(" value\n\"quoted\" ", arguments.strField("query").?); + try std.testing.expectEqual(@as(i64, 2), arguments.get("limit").?.int64().?); + try std.testing.expect(json.equal(arguments.get("data").?, try parse(a, "[true,null,{\"x\":1}]"))); + try std.testing.expectEqualStrings("files", parsed.calls.?[1].strField("namespace").?); + try std.testing.expectEqualStrings("{}", parsed.calls.?[1].get("function").?.strField("arguments").?); + const old = "<｜DSML｜tool_calls><｜DSML｜invoke name=\"lookup\"><｜DSML｜parameter name=\"q\" string=\"true\">v"; + const old_tools = (try parse(a, "[{\"function\":{\"name\":\"lookup\"}}]")).array; + try std.testing.expect((try tool_parse.parse(a, old, old_tools, null)).calls != null); +} + +test "V41 DSML streaming survives every byte boundary and emits each validated call once" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const tools = (try parse(a, offered)).array; + var streamer = try tool_stream.Streamer.init(a, tools); + var deltas: std.ArrayList(Value) = .empty; + for (0..completion.len + 1) |end| try streamer.feed(completion[0..end], &deltas); + try streamer.feed(completion, &deltas); + try std.testing.expect(streamer.streamed and !streamer.invalidDsml(completion)); + try std.testing.expectEqual(@as(usize, 4), deltas.items.len); + try std.testing.expectEqualStrings("search", deltas.items[0].get("tool_calls").?.array[0].strField("namespace").?); + try std.testing.expectEqual(@as(i64, 1), deltas.items[2].get("tool_calls").?.array[0].get("index").?.int64().?); + try std.testing.expectEqualStrings("files", deltas.items[2].get("tool_calls").?.array[0].strField("namespace").?); + try std.testing.expectEqualStrings("summary\n\n", try reply_text.hideToolCalls(a, completion, true)); +} + +test "V41 malformed and incomplete DSML never silently truncates into a valid call" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const tools = (try parse(a, offered)).array; + for ([_][]const u8{ + "<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"search::lookup\">\n<｜DSML｜ parameter name=\"x\" string=\"false\">bad json\n\n", + "<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"search::lookup\">\n<｜DSML｜ parameter name=\"x\" string=\"true\">a\n<｜DSML｜ parameter name=\"x\" string=\"true\">b\n\n", + "<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"search::lookup\">\n", + }) |bad| { + const parsed = try tool_parse.parse(a, bad, tools, 1); + try std.testing.expect(parsed.calls == null); + try std.testing.expectEqualStrings(reply_text.pyStrip(bad), reply_text.pyStrip(parsed.content)); + var streamer = try tool_stream.Streamer.init(a, tools); + var deltas: std.ArrayList(Value) = .empty; + try streamer.feed(bad, &deltas); + try std.testing.expectEqual(@as(usize, 0), deltas.items.len); + try std.testing.expect(streamer.invalidDsml(bad)); + } +} + +test "V41 completion parsing matches independently generated official decode goldens" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const tools = (try parse(a, "[" ++ + "{\"function\":{\"name\":\"lookup\"}}," ++ + "{\"namespace\":\"search\",\"function\":{\"name\":\"lookup\"}}," ++ + "{\"namespace\":\"files\",\"function\":{\"name\":\"lookup\"}}]")).array; + const goldens = try parse(a, @embedFile("deepseek_v41_tool_goldens.json")); + for (goldens.get("cases").?.array) |c| { + const raw = c.strField("completion").?; + const eos = "<｜end▁of▁sentence｜>"; + const complete = raw[0 .. raw.len - eos.len]; + const expected = c.get("parsed").?; + const split = if (c.get("thinking").?.bool) + try reply_text.splitThinking(a, complete, true, reply_text.think_markers) + else + reply_text.Split{ .reasoning = "", .answer = complete }; + try std.testing.expectEqualStrings(expected.strField("reasoning_content").?, split.reasoning); + const parsed = try tool_parse.parse(a, split.answer, tools, null); + try std.testing.expectEqualStrings(expected.strField("content").?, parsed.content); + const wanted = expected.get("tool_calls").?.array; + try std.testing.expectEqual(wanted.len, parsed.calls.?.len); + for (wanted, parsed.calls.?) |reference_call, actual| { + try std.testing.expectEqualStrings(json.strOr(reference_call.get("namespace")), json.strOr(actual.get("namespace"))); + const ref_fn = reference_call.get("function").?; + const actual_fn = actual.get("function").?; + try std.testing.expectEqualStrings(ref_fn.strField("name").?, actual_fn.strField("name").?); + try std.testing.expect(json.equal(try parse(a, ref_fn.strField("arguments").?), try parse(a, actual_fn.strField("arguments").?))); + } + } +} diff --git a/zig/src/server/deepseek_v41_tool_goldens.json b/zig/src/server/deepseek_v41_tool_goldens.json new file mode 100644 index 0000000000..94b86ce1d8 --- /dev/null +++ b/zig/src/server/deepseek_v41_tool_goldens.json @@ -0,0 +1,72 @@ +{ + "reference": "DeepSeek-AI encoding.py", + "revision": "2cba9e42aa026125f3ed06c6d98c1db82f7ca027", + "cases": [ + { + "name": "namespaces_and_parameter_boundaries", + "thinking": true, + "completion": " reason summary\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"search::lookup\">\n<｜DSML｜ parameter name=\"query\" string=\"true\"> value\n\"quoted\" \n<｜DSML｜ parameter name=\"limit\" string=\"false\">2\n<｜DSML｜ parameter name=\"data\" string=\"false\">[true, null, {\"x\": 1}]\n\n<｜DSML｜ invoke name=\"files::lookup\">\n\n\n<｜end▁of▁sentence｜>", + "parsed": { + "role": "assistant", + "content": "summary", + "reasoning_content": " reason ", + "tool_calls": [ + { + "type": "function", + "function": { + "name": "lookup", + "arguments": "{\"query\": \" value\\n\\\"quoted\\\" \", \"limit\": 2, \"data\": [true, null, {\"x\": 1}]}" + }, + "namespace": "search" + }, + { + "type": "function", + "function": { + "name": "lookup", + "arguments": "{}" + }, + "namespace": "files" + } + ] + } + }, + { + "name": "chat_empty_call", + "thinking": false, + "completion": "summary\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"lookup\">\n\n\n<｜end▁of▁sentence｜>", + "parsed": { + "role": "assistant", + "content": "summary", + "reasoning_content": "", + "tool_calls": [ + { + "type": "function", + "function": { + "name": "lookup", + "arguments": "{}" + } + } + ] + } + }, + { + "name": "escaped_keys_and_unicode", + "thinking": false, + "completion": "summary\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"lookup\">\n<｜DSML｜ parameter name=\"key\"quoted\" string=\"true\">中文\n<｜DSML｜ parameter name=\"false\" string=\"false\">false\n<｜DSML｜ parameter name=\"null\" string=\"false\">null\n<｜DSML｜ parameter name=\"object\" string=\"false\">{\"b\": [1, 2]}\n\n<｜end▁of▁sentence｜>", + "parsed": { + "role": "assistant", + "content": "summary", + "reasoning_content": "", + "tool_calls": [ + { + "type": "function", + "function": { + "name": "lookup", + "arguments": "{\"key\\\"quoted\": \"中文\", \"false\": false, \"null\": null, \"object\": {\"b\": [1, 2]}}" + } + } + ] + } + } + ] +} diff --git a/zig/src/server/fields.zig b/zig/src/server/fields.zig index 0de8c8abfa..5c9c52435b 100644 --- a/zig/src/server/fields.zig +++ b/zig/src/server/fields.zig @@ -211,9 +211,19 @@ pub fn thinkingFields(cx: *Cx, body: Value, levels: []const []const u8) errors.R effort = k.field("reasoning_effort"); }; if (effort) |e| { - if (e != .string or !isEffort(e.string)) return cx.refuse("reasoning_effort must be none, minimal, low, medium, high, xhigh or max"); - out.effort = coerceEffort(e.string, levels); - out.enable = !std.mem.eql(u8, e.string, "none"); + const numeric_v41 = for (levels) |level| { + if (std.mem.eql(u8, level, "numeric")) break true; + } else false; + if (e == .int and numeric_v41) { + const n = e.int64() orelse return cx.refuse("DeepSeek V4.1 reasoning_effort must be an integer from 1 to 100"); + if (n < 1 or n > 100) return cx.refuse("DeepSeek V4.1 reasoning_effort must be an integer from 1 to 100"); + out.effort = e.int; + out.enable = true; + } else { + if (e != .string or !isEffort(e.string)) return cx.refuse("reasoning_effort must be none, minimal, low, medium, high, xhigh or max"); + out.effort = coerceEffort(e.string, levels); + out.enable = !std.mem.eql(u8, e.string, "none"); + } } if (kw_object) if (enable) |on| { out.enable = on.truthy(); @@ -258,6 +268,7 @@ pub fn coerceEffort(effort: []const u8, levels: []const []const u8) []const u8 { /// The efforts a chat template names in quotes (``effort_levels``). pub fn effortLevels(a: Allocator, template: []const u8) Allocator.Error![]const []const u8 { var found: std.ArrayList([]const u8) = .empty; + if (@import("deepseek_v41_text.zig").isTemplate(template)) try found.append(a, "numeric"); for (effort_order) |name| { var at: usize = 0; while (std.mem.indexOfPos(u8, template, at, name)) |i| : (at = i + 1) { diff --git a/zig/src/server/hf_text.zig b/zig/src/server/hf_text.zig index 22e4b976dc..5f9380799b 100644 --- a/zig/src/server/hf_text.zig +++ b/zig/src/server/hf_text.zig @@ -4,6 +4,7 @@ const tokenizer = @import("tokenizer"); const template = @import("template"); const json = @import("json.zig"); const model_text = @import("model_text.zig"); +const deepseek_v41 = @import("deepseek_v41_text.zig"); const Value = json.Value; const Allocator = std.mem.Allocator; @@ -13,8 +14,11 @@ pub const HfText = struct { tok: tokenizer.Tokenizer, source: []const u8 = "", tool_source: []const u8 = "", + // Loaded once; HTTP readers never allocate into the checkpoint's arena. + combined_source: []const u8 = "", compiled: ?template.Compiled = null, tool_compiled: ?template.Compiled = null, + v41: bool = false, specials: std.json.Value = .null, eos: []const u32 = &.{}, prefix: []const u32 = &.{}, // the post-processor's tokens before a text (add_special_tokens) @@ -28,18 +32,25 @@ pub const HfText = struct { errdefer t.arena.deinit(); const a = t.arena.allocator(); t.tok = tokenizer.loadTokenizer(io, gpa, dir) catch |e| { + if (e == error.OutOfMemory) return error.OutOfMemory; problem.* = try std.fmt.allocPrint(pa, "cannot read {s}/tokenizer.json ({s})", .{ dir, @errorName(e) }); return error.Load; }; errdefer t.tok.deinit(); - const config = readJson(io, a, dir, "tokenizer_config.json"); - const model_config = readJson(io, a, dir, "config.json"); + const config = try readJson(io, a, dir, "tokenizer_config.json"); + const model_config = try readJson(io, a, dir, "config.json"); if (template_override) |path| { t.source = std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(16 << 20)) catch |e| { + if (e == error.OutOfMemory) return error.OutOfMemory; problem.* = try std.fmt.allocPrint(pa, "cannot read chat template override ({s})", .{@errorName(e)}); return error.Load; }; } else try t.readTemplates(io, a, dir, config); + t.v41 = deepseek_v41.isTemplate(t.source) or + (template_override == null and model_config == .object and if (model_config.object.get("model_type")) |kind| + kind == .string and std.mem.eql(u8, kind.string, "deepseek_v41") + else + false); t.specials = try specialsOf(a, config); t.eos = try t.eosIds(a, model_config, config); try t.readPostProcessor(io, a, dir); @@ -47,11 +58,26 @@ pub const HfText = struct { var diag: template.Diag = .{}; defer if (diag.msg.len > 0) gpa.free(diag.msg); t.compiled = template.compile(gpa, t.source, &diag) catch |e| { + if (e == error.OutOfMemory) return error.OutOfMemory; problem.* = try std.fmt.allocPrint(pa, "the chat template does not compile: {s} ({s})", .{ diag.msg, @errorName(e) }); return error.Load; }; } - if (t.tool_source.len > 0) t.tool_compiled = template.compile(gpa, t.tool_source, null) catch null; + errdefer if (t.compiled) |*c| c.deinit(); + if (t.tool_source.len > 0) { + var diag: template.Diag = .{}; + defer if (diag.msg.len > 0) gpa.free(diag.msg); + t.tool_compiled = template.compile(gpa, t.tool_source, &diag) catch |e| { + if (e == error.OutOfMemory) return error.OutOfMemory; + problem.* = try std.fmt.allocPrint(pa, "the tool chat template does not compile: {s} ({s})", .{ diag.msg, @errorName(e) }); + return error.Load; + }; + } + errdefer if (t.tool_compiled) |*c| c.deinit(); + t.combined_source = if (t.tool_source.len > 0) + try std.mem.concat(a, u8, &.{ t.source, " ", t.tool_source }) + else + t.source; return t; } @@ -63,10 +89,12 @@ pub const HfText = struct { t.gpa.destroy(t); } - fn readJson(io: std.Io, a: Allocator, dir: []const u8, name: []const u8) std.json.Value { - const path = std.fs.path.join(a, &.{ dir, name }) catch return .null; - const bytes = std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(64 << 20)) catch return .null; - return std.json.parseFromSliceLeaky(std.json.Value, a, bytes, .{}) catch .null; + fn readJson(io: std.Io, a: Allocator, dir: []const u8, name: []const u8) Allocator.Error!std.json.Value { + const path = try std.fs.path.join(a, &.{ dir, name }); + const bytes = std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(64 << 20)) catch |err| + return if (err == error.OutOfMemory) error.OutOfMemory else .null; + return std.json.parseFromSliceLeaky(std.json.Value, a, bytes, .{}) catch |err| + if (err == error.OutOfMemory) error.OutOfMemory else .null; } /// chat_template.jinja, else chat_template.json, else tokenizer_config.json's (named: "default", and "tool_use"). @@ -75,8 +103,10 @@ pub const HfText = struct { if (std.Io.Dir.cwd().readFileAlloc(io, jinja, a, .limited(16 << 20))) |s| { t.source = s; return; - } else |_| {} - const sidecar = readJson(io, a, dir, "chat_template.json"); + } else |err| { + if (err == error.OutOfMemory) return error.OutOfMemory; + } + const sidecar = try readJson(io, a, dir, "chat_template.json"); if (sidecar == .object) if (sidecar.object.get("chat_template")) |s| if (s == .string) { t.source = s.string; return; @@ -150,7 +180,7 @@ pub const HfText = struct { /// The tokens ``add_special_tokens`` adds: TemplateProcessing's single template, or Bert and Roberta's cls and sep. fn readPostProcessor(t: *HfText, io: std.Io, a: Allocator, dir: []const u8) !void { - const doc = readJson(io, a, dir, "tokenizer.json"); + const doc = try readJson(io, a, dir, "tokenizer.json"); if (doc != .object) return; const pp = doc.object.get("post_processor") orelse return; var prefix: std.ArrayList(u32) = .empty; @@ -226,11 +256,13 @@ pub const HfText = struct { const t = self(ctx); const ids = t.tok.encode(a, s) catch |e| return if (e == error.OutOfMemory) error.OutOfMemory else error.Template; if (!add_special or (t.prefix.len == 0 and t.suffix.len == 0)) return ids; + defer a.free(ids); return std.mem.concat(a, u32, &.{ t.prefix, ids, t.suffix }); } - fn decodeFn(ctx: *anyopaque, a: Allocator, ids: []const u32) Allocator.Error![]u8 { - return self(ctx).tok.decode(a, ids, false) catch |e| if (e == error.OutOfMemory) error.OutOfMemory else a.dupe(u8, "\u{fffd}"); + fn decodeFn(ctx: *anyopaque, a: Allocator, ids: []const u32) model_text.Error![]u8 { + return self(ctx).tok.decode(a, ids, false) catch |e| + return if (e == error.OutOfMemory) error.OutOfMemory else error.Template; } fn tokenIdFn(ctx: *anyopaque, piece: []const u8) ?u32 { @@ -242,8 +274,9 @@ pub const HfText = struct { return a.dupe(u8, piece); } - fn tokenBytesFn(ctx: *anyopaque, a: Allocator, id: u32) Allocator.Error!?[]u8 { - return self(ctx).tok.tokenBytes(a, id) catch |e| if (e == error.OutOfMemory) error.OutOfMemory else null; + fn tokenBytesFn(ctx: *anyopaque, a: Allocator, id: u32) model_text.Error!?[]u8 { + return self(ctx).tok.tokenBytes(a, id) catch |e| + return if (e == error.OutOfMemory) error.OutOfMemory else error.Template; } fn vocabFn(ctx: *anyopaque) u32 { @@ -255,27 +288,39 @@ pub const HfText = struct { } fn sourceFn(ctx: *anyopaque) []const u8 { - const t = self(ctx); - return if (t.tool_source.len > 0) std.mem.concat(t.arena.allocator(), u8, &.{ t.source, " ", t.tool_source }) catch t.source else t.source; + return self(ctx).combined_source; } fn renderFn(ctx: *anyopaque, a: Allocator, messages: Value, options: model_text.RenderOptions, problem: *[]const u8) model_text.Error![]u8 { const t = self(ctx); - const compiled = if (options.tools != null and t.tool_compiled != null) &t.tool_compiled.? else if (t.compiled) |*c| c else { + const use_tools = options.tools != null and t.tool_compiled != null; + const compiled = if (use_tools) &t.tool_compiled.? else if (t.compiled) |*c| c else { problem.* = "this checkpoint has no chat template: send a completion prompt instead"; return error.Template; }; + var scratch = std.heap.ArenaAllocator.init(a); + defer scratch.deinit(); + const work = scratch.allocator(); + const source = if (use_tools) t.tool_source else t.source; + const v41 = t.v41 or deepseek_v41.isTemplate(source); var context: std.json.ObjectMap = .empty; if (t.specials == .object) { var it = t.specials.object.iterator(); - while (it.next()) |e| try context.put(a, e.key_ptr.*, e.value_ptr.*); + while (it.next()) |e| try context.put(work, e.key_ptr.*, e.value_ptr.*); + } + try context.put(work, "enable_thinking", .{ .bool = options.enable_thinking }); + try context.put(work, "thinking_mode", .{ .string = if (options.enable_thinking) "thinking" else "chat" }); + if (options.reasoning_effort) |e| { + const effort_value: std.json.Value = if (v41) deepseek_v41.effort(work, source, e) catch |err| { + problem.* = "DeepSeek V4.1 reasoning effort must be within 1..100 or name a supported effort"; + return err; + } else .{ .string = e }; + try context.put(work, "reasoning_effort", effort_value); } - try context.put(a, "enable_thinking", .{ .bool = options.enable_thinking }); - try context.put(a, "thinking_mode", .{ .string = if (options.enable_thinking) "thinking" else "chat" }); - if (options.reasoning_effort) |e| try context.put(a, "reasoning_effort", .{ .string = e }); var diag: template.Diag = .{}; - const tools: std.json.Value = if (options.tools) |tl| try toStd(a, tl) else .null; - const out = template.render(a, compiled, try toStd(a, messages), tools, .{ .object = context }, options.add_generation_prompt, &diag) catch |e| { + const tools: std.json.Value = if (options.tools) |tl| try toStd(work, tl) else .null; + const prepared = if (v41) try deepseek_v41.prepare(work, messages, problem) else messages; + const out = template.render(a, compiled, try toStd(work, prepared), tools, .{ .object = context }, options.add_generation_prompt, &diag) catch |e| { if (e == error.OutOfMemory) return error.OutOfMemory; problem.* = diag.msg; return error.Template; @@ -344,6 +389,12 @@ test "explicit template overrides checkpoint templates including tool prompts an const messages = (try json.parse(a, "[{\"role\":\"user\",\"content\":\"H\"}]")).ok; const named = try HfText.load(gpa, io, dir, null, a, &problem); defer named.deinit(); + const source = named.text().templateSource(); + try std.testing.expectEqualStrings("default{{ messages[0].content }} tools{{ messages[0].content }}", source); + try std.testing.expectEqual(@intFromPtr(source.ptr), @intFromPtr(named.text().templateSource().ptr)); + const owned = try named.text().render(gpa, messages, .{}, &problem); + defer gpa.free(owned); + try std.testing.expectEqualStrings("defaultH", owned); try std.testing.expectEqualStrings("defaultH", try named.text().render(a, messages, .{}, &problem)); try std.testing.expectEqualStrings("toolsH", try named.text().render(a, messages, .{ .tools = .{ .array = &.{} } }, &problem)); try tmp.dir.writeFile(io, .{ .sub_path = "override.jinja", .data = "override{{ messages[0].content }}" }); @@ -366,4 +417,44 @@ test "explicit template overrides checkpoint templates including tool prompts an const empty = try HfText.load(gpa, io, dir, path, a, &problem); defer empty.deinit(); try std.testing.expectError(error.Template, empty.text().render(a, messages, .{}, &problem)); + try tmp.dir.deleteFile(io, "chat_template.jinja"); + try tmp.dir.writeFile(io, .{ .sub_path = "tokenizer_config.json", .data = + \\{"chat_template":{"default":"ok","tool_use":"{% unknown %}"}} + }); + try std.testing.expectError(error.Load, HfText.load(gpa, io, dir, null, a, &problem)); +} + +test "postprocessor encoding frees the unwrapped token allocation" { + const a = std.testing.allocator; + var t: HfText = .{ + .gpa = a, + .arena = .init(a), + .tok = try tokenizer.parse(a, + \\{"model":{"type":"BPE","vocab":{"H":0},"merges":[]},"added_tokens":[]} + ), + .prefix = &.{1}, + .suffix = &.{2}, + }; + defer t.tok.deinit(); + defer t.arena.deinit(); + const ids = try t.text().encode(a, "H", true); + defer a.free(ids); + try std.testing.expectEqualSlices(u32, &.{ 1, 0, 2 }, ids); +} + +test "decoder and token bytes preserve allocation failures" { + const a = std.testing.allocator; + var t: HfText = .{ + .gpa = a, + .arena = .init(a), + .tok = try tokenizer.parse(a, + \\{"model":{"type":"BPE","vocab":{"H":0},"merges":[]},"added_tokens":[],"decoder":{"type":"ByteLevel"}} + ), + }; + defer t.tok.deinit(); + defer t.arena.deinit(); + var decode_fail = std.testing.FailingAllocator.init(a, .{ .fail_index = 0 }); + try std.testing.expectError(error.OutOfMemory, t.text().decode(decode_fail.allocator(), &.{0})); + var bytes_fail = std.testing.FailingAllocator.init(a, .{ .fail_index = 0 }); + try std.testing.expectError(error.OutOfMemory, t.text().tokenBytes(bytes_fail.allocator(), 0)); } diff --git a/zig/src/server/legacy_consumer_regressions.zig b/zig/src/server/legacy_consumer_regressions.zig new file mode 100644 index 0000000000..0f7847b17a --- /dev/null +++ b/zig/src/server/legacy_consumer_regressions.zig @@ -0,0 +1,198 @@ +//! Containment final-review consumer regressions, exercised by the preflight suite. +const std = @import("std"); +const json = @import("json.zig"); + +const dsml_prefix = "<｜DSML｜tool_calls><｜DSML｜invoke name=\"lookup\">" ++ + "<｜DSML｜parameter name=\"q\" string=\"true\">"; +const dsml_suffix = ""; +const dsml_body = "{\"model\":\"test-model\",\"max_tokens\":2000,\"messages\":[{\"role\":\"user\",\"content\":\"x\"}]," ++ + "\"tools\":[{\"type\":\"function\",\"function\":{\"name\":\"lookup\",\"parameters\":{\"type\":\"object\"}}}]"; + +pub fn run(comptime TestText: type, comptime CallEngine: type, comptime sendWith: anytype) !void { + // DSML literal calls and adjacent XML calls agree in whole and streamed consumer replies. + { + const literals = [_][]const u8{ + "", + "", + }; + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const adjacent = "7"; + for (literals) |literal| for ([_]bool{ false, true }) |with_adjacent| { + const raw = try std.mem.concat(a, u8, &.{ dsml_prefix, literal, dsml_suffix, if (with_adjacent) adjacent else "" }); + const tokens = try TestText.encode(undefined, a, raw, false); + for ([_]bool{ false, true }) |single| { + const expected = try std.mem.concat(a, u8, &.{ + "{\"q\":", try json.quote(a, literal, .{}), "}", if (with_adjacent and !single) "{\"n\":\"7\"}" else "", + }); + const body = try std.mem.concat(a, u8, &.{ dsml_body, if (single) ",\"parallel_tool_calls\":false" else "" }); + for ([_]bool{ false, true }) |batched| { + var backend: CallEngine = .{ .tokens = tokens, .batched = batched }; + const whole = try sendWith(backend.engine(), .chat, try std.mem.concat(a, u8, &.{ body, ",\"stream\":false}" }), .{}); + const streamed = try sendWith(backend.engine(), .chat, try std.mem.concat(a, u8, &.{ body, ",\"stream\":true}" }), .{}); + try std.testing.expectEqual(@as(?u16, 200), whole.status); + try std.testing.expect(streamed.opened and streamed.status == null); + try std.testing.expectEqual(@as(usize, if (with_adjacent and !single) 2 else 1), whole.calls); + try std.testing.expectEqual(whole.calls, streamed.calls); + try std.testing.expectEqualStrings(expected, whole.arguments[0..whole.arguments_len]); + try std.testing.expectEqualStrings(expected, streamed.arguments[0..streamed.arguments_len]); + try std.testing.expectEqualStrings("", whole.content()); + try std.testing.expectEqualStrings("", streamed.content()); + try std.testing.expectEqual(@as(i64, 0), streamed.preparing); + } + } + }; + } + // unfinished and malformed DSML never emit their complete nested call in consumer replies. + { + const prefix = dsml_prefix ++ ""; + const replies = [_][]const u8{ + prefix, + prefix ++ "", + }; + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + for (replies) |raw| { + const tokens = try TestText.encode(undefined, a, raw, false); + for ([_]bool{ false, true }) |single| for ([_]bool{ false, true }) |batched| { + var backend: CallEngine = .{ .tokens = tokens, .batched = batched }; + const body = try std.mem.concat(a, u8, &.{ dsml_body, if (single) ",\"parallel_tool_calls\":false" else "" }); + const whole = try sendWith(backend.engine(), .chat, try std.mem.concat(a, u8, &.{ body, ",\"stream\":false}" }), .{}); + const streamed = try sendWith(backend.engine(), .chat, try std.mem.concat(a, u8, &.{ body, ",\"stream\":true}" }), .{}); + try std.testing.expectEqual(@as(?u16, 200), whole.status); + try std.testing.expect(streamed.opened and streamed.status == null); + try std.testing.expectEqual(@as(usize, 0), whole.calls); + try std.testing.expectEqual(@as(usize, 0), streamed.calls); + try std.testing.expectEqual(@as(usize, 0), streamed.arguments_len); + try std.testing.expectEqualStrings(raw, whole.content()); + try std.testing.expectEqualStrings(raw, streamed.content()); + try std.testing.expectEqual(@as(i64, 0), streamed.preparing); + }; + } + } + // DSML and adjacent existing non-XML envelopes preserve whole and SSE calls without duplicates. + { + const outer = dsml_prefix ++ "outer" ++ dsml_suffix; + const cases = [_][2][]const u8{ + .{ "{\"name\":\"lookup\",\"arguments\":{\"q\":\"adjacent-json\"}}", "{\"q\":\"adjacent-json\"}" }, + .{ "lookupqadjacent-glm", "{\"q\":\"adjacent-glm\"}" }, + .{ "<|tool_call>call:lookup{q:<|\"|>adjacent-gemma<|\"|>}", "{\"q\":\"adjacent-gemma\"}" }, + .{ "{\"name\":\"lookup\",\"arguments\":{\"q\":\"adjacent-namespaced\"}}", "{\"q\":\"adjacent-namespaced\"}" }, + }; + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + for (cases) |case| for ([_]bool{ false, true }) |outer_first| { + const raw = try std.mem.concat(a, u8, if (outer_first) &.{ outer, case[0] } else &.{ case[0], outer }); + const expected = try std.mem.concat(a, u8, if (outer_first) + &.{ "{\"q\":\"outer\"}", case[1] } + else + &.{ case[1], "{\"q\":\"outer\"}" }); + const tokens = try TestText.encode(undefined, a, raw, false); + for ([_]bool{ false, true }) |batched| { + var backend: CallEngine = .{ .tokens = tokens, .batched = batched }; + const whole = try sendWith(backend.engine(), .chat, dsml_body ++ ",\"stream\":false}", .{}); + const streamed = try sendWith(backend.engine(), .chat, dsml_body ++ ",\"stream\":true}", .{}); + try std.testing.expectEqual(@as(?u16, 200), whole.status); + try std.testing.expect(streamed.opened and streamed.status == null); + try std.testing.expectEqual(@as(usize, 2), whole.calls); + try std.testing.expectEqual(@as(usize, 2), streamed.calls); + try std.testing.expectEqualStrings(expected, whole.arguments[0..whole.arguments_len]); + try std.testing.expectEqualStrings(expected, streamed.arguments[0..streamed.arguments_len]); + try std.testing.expectEqual(@as(i64, 0), streamed.preparing); + } + }; + } + // a supported bare JSON call containing a literal DSML opener agrees in whole and SSE replies. + { + const call = "{\"name\":\"lookup\",\"arguments\":{\"q\":\"<｜DSML｜tool_calls>\"}}"; + const replies = [_][]const u8{ call, "[" ++ call ++ "]", "```json\n" ++ call ++ "\n```" }; + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + for (replies) |raw| { + const tokens = try TestText.encode(undefined, a, raw, false); + for ([_]bool{ false, true }) |batched| { + var backend: CallEngine = .{ .tokens = tokens, .batched = batched }; + const whole = try sendWith(backend.engine(), .chat, dsml_body ++ ",\"stream\":false}", .{}); + const streamed = try sendWith(backend.engine(), .chat, dsml_body ++ ",\"stream\":true}", .{}); + try std.testing.expectEqual(@as(?u16, 200), whole.status); + try std.testing.expect(streamed.opened and streamed.status == null); + try std.testing.expectEqual(@as(usize, 1), whole.calls); + try std.testing.expectEqual(@as(usize, 1), streamed.calls); + const expected = "{\"q\":\"<｜DSML｜tool_calls>\"}"; + try std.testing.expectEqualStrings(expected, whole.arguments[0..whole.arguments_len]); + try std.testing.expectEqualStrings(expected, streamed.arguments[0..streamed.arguments_len]); + try std.testing.expectEqual(@as(i64, 0), streamed.preparing); + } + } + } + // closed rejected calls and literal less-than prose preserve later consumer calls. + { + const outer = dsml_prefix ++ "outer" ++ dsml_suffix; + const adjacent = "adjacent"; + const cases = [_][2][]const u8{ + .{ outer ++ "{\"name\":\"other\",\"arguments\":{}}" ++ adjacent, "{\"q\":\"outer\"}{\"q\":\"adjacent\"}" }, + .{ outer ++ "" ++ adjacent, "{\"q\":\"outer\"}{\"q\":\"adjacent\"}" }, + .{ outer ++ "" ++ adjacent, "{\"q\":\"outer\"}{\"q\":\"adjacent\"}" }, + .{ outer ++ "" ++ adjacent, "{\"q\":\"outer\"}{\"q\":\"adjacent\"}" }, + .{ adjacent ++ "1 < 2 " ++ adjacent, "{\"q\":\"adjacent\"}{\"q\":\"adjacent\"}" }, + .{ "1 < 2 " ++ outer ++ adjacent, "{\"q\":\"outer\"}{\"q\":\"adjacent\"}" }, + }; + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + for (cases) |case| { + const tokens = try TestText.encode(undefined, a, case[0], false); + for ([_]bool{ false, true }) |batched| { + var backend: CallEngine = .{ .tokens = tokens, .batched = batched }; + const whole = try sendWith(backend.engine(), .chat, dsml_body ++ ",\"stream\":false}", .{}); + const streamed = try sendWith(backend.engine(), .chat, dsml_body ++ ",\"stream\":true}", .{}); + try std.testing.expectEqual(@as(?u16, 200), whole.status); + try std.testing.expect(streamed.opened and streamed.status == null); + try std.testing.expectEqual(@as(usize, 2), whole.calls); + try std.testing.expectEqual(@as(usize, 2), streamed.calls); + try std.testing.expectEqualStrings(case[1], whole.arguments[0..whole.arguments_len]); + try std.testing.expectEqualStrings(case[1], streamed.arguments[0..streamed.arguments_len]); + try std.testing.expectEqual(@as(i64, 0), streamed.preparing); + } + } + } +} + +pub fn mixedChain(comptime TestText: type, comptime CallEngine: type, comptime sendWith: anytype) !void { + const xml = "xml"; + const v41 = "<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"lookup\">\n" ++ + "<｜DSML｜ parameter name=\"q\" string=\"true\">v41\n" ++ + "\n"; + const cases = [_][]const u8{ + "{\"name\":\"lookup\",\"arguments\":{\"q\":\"middle\"}}", + "lookupqmiddle", + "<|tool_call>call:lookup{q:<|\"|>middle<|\"|>}", + "{\"name\":\"lookup\",\"arguments\":{\"q\":\"middle\"}}", + }; + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + for (cases) |middle| { + const raw = try std.mem.concat(a, u8, &.{ xml, middle, v41, xml }); + const tokens = try TestText.encode(undefined, a, raw, false); + for ([_]bool{ false, true }) |batched| { + var backend: CallEngine = .{ .tokens = tokens, .batched = batched }; + const whole = try sendWith(backend.engine(), .chat, dsml_body ++ ",\"stream\":false}", .{}); + const streamed = try sendWith(backend.engine(), .chat, dsml_body ++ ",\"stream\":true}", .{}); + try std.testing.expectEqual(@as(?u16, 200), whole.status); + try std.testing.expect(streamed.opened and streamed.status == null); + try std.testing.expect(!whole.got_error and !streamed.got_error); + try std.testing.expectEqual(@as(usize, 4), whole.calls); + try std.testing.expectEqual(whole.calls, streamed.calls); + const expected = "{\"q\":\"xml\"}{\"q\":\"middle\"}{\"q\":\"v41\"}{\"q\":\"xml\"}"; + try std.testing.expectEqualStrings(expected, whole.arguments[0..whole.arguments_len]); + try std.testing.expectEqualStrings(expected, streamed.arguments[0..streamed.arguments_len]); + try std.testing.expectEqualStrings(whole.content(), streamed.content()); + try std.testing.expectEqual(@as(i64, 0), streamed.preparing); + } + } +} diff --git a/zig/src/server/logprobs.zig b/zig/src/server/logprobs.zig index 0d0b37d0d6..45146d7ce6 100644 --- a/zig/src/server/logprobs.zig +++ b/zig/src/server/logprobs.zig @@ -4,6 +4,7 @@ const api = @import("engine_api"); const json = @import("json.zig"); const errors = @import("errors.zig"); const Server = @import("server.zig").Server; +const model_text = @import("model_text.zig"); const Value = json.Value; const Cx = errors.Cx; const Allocator = std.mem.Allocator; @@ -17,6 +18,14 @@ pub fn admit(srv: *Server, cx: *Cx, thinking: bool, think_budget: u32, count: u8 /// ``choices[0].logprobs`` for the answer's tokens: after the first end-of-thinking token and dropped newlines. pub fn value(srv: *Server, cx: *Cx, a: Allocator, thinking: bool, tokens: []const u32, rows: []const api.LogprobRow) errors.Refused!Value { + return rowsValue(srv, cx, a, thinking, tokens, rows) catch |err| switch (err) { + error.OutOfMemory => return error.OutOfMemory, + error.Template => return cx.other("the tokenizer could not decode logprob tokens"), + error.Refused => return error.Refused, + }; +} + +fn rowsValue(srv: *Server, cx: *Cx, a: Allocator, thinking: bool, tokens: []const u32, rows: []const api.LogprobRow) (model_text.Error || errors.Refused)!Value { if (rows.len < tokens.len) return cx.other("the engine gave fewer logprob rows than reply tokens"); var start: usize = 0; if (thinking) { @@ -43,7 +52,7 @@ pub fn value(srv: *Server, cx: *Cx, a: Allocator, thinking: bool, tokens: []cons } /// ``{token, logprob, bytes}``: the token's text, and its own bytes (a partial UTF-8 sequence kept whole). -fn entry(srv: *Server, a: Allocator, id: u32, logprob: f32) Allocator.Error!*json.Object { +fn entry(srv: *Server, a: Allocator, id: u32, logprob: f32) model_text.Error!*json.Object { const bytes = (try srv.text.tokenBytes(a, id)) orelse &.{}; const o = try json.newObject(a); try o.put(a, "token", .{ .string = try srv.text.decode(a, &.{id}) }); diff --git a/zig/src/server/model_text.zig b/zig/src/server/model_text.zig index 0b36055a7f..9273afa73f 100644 --- a/zig/src/server/model_text.zig +++ b/zig/src/server/model_text.zig @@ -21,8 +21,8 @@ pub const Text = struct { pub const VTable = struct { /// Ids of ``text``, special-token text included as its token; ``add_special`` adds what the tokenizer prepends. encode: *const fn (ctx: *anyopaque, a: Allocator, text: []const u8, add_special: bool) Error![]u32, - /// UTF-8 of ``ids`` with special tokens; bytes that end mid-character decode as U+FFFD, as Python's replace does. - decode: *const fn (ctx: *anyopaque, a: Allocator, ids: []const u32) Allocator.Error![]u8, + /// Decode UTF-8 with special tokens; incomplete bytes become U+FFFD, tokenizer failures return Template. + decode: *const fn (ctx: *anyopaque, a: Allocator, ids: []const u32) Error![]u8, /// The id of a token the vocabulary holds whole (``convert_tokens_to_ids``), else null. token_id: *const fn (ctx: *anyopaque, piece: []const u8) ?u32, /// A token's vocabulary string (``convert_ids_to_tokens``). @@ -36,15 +36,15 @@ pub const Text = struct { template_source: *const fn (ctx: *anyopaque) []const u8, /// Whether an id is a special token (an added token flagged special); null: none are known. special: ?*const fn (ctx: *anyopaque, id: u32) bool = null, - /// A token's own bytes (a partial UTF-8 sequence kept whole) for logprob rows; null: the tokenizer can't say. - token_bytes: ?*const fn (ctx: *anyopaque, a: Allocator, id: u32) Allocator.Error!?[]u8 = null, + /// A token's exact bytes; null means context-dependent decoding, while tokenizer failures return Template. + token_bytes: ?*const fn (ctx: *anyopaque, a: Allocator, id: u32) Error!?[]u8 = null, }; pub fn encode(t: Text, a: Allocator, text: []const u8, add_special: bool) Error![]u32 { return t.vtable.encode(t.ctx, a, text, add_special); } - pub fn decode(t: Text, a: Allocator, ids: []const u32) Allocator.Error![]u8 { + pub fn decode(t: Text, a: Allocator, ids: []const u32) Error![]u8 { return t.vtable.decode(t.ctx, a, ids); } @@ -61,7 +61,7 @@ pub const Text = struct { } /// A token's own bytes, or null when this text cannot give them. - pub fn tokenBytes(t: Text, a: Allocator, id: u32) Allocator.Error!?[]u8 { + pub fn tokenBytes(t: Text, a: Allocator, id: u32) Error!?[]u8 { const f = t.vtable.token_bytes orelse return null; return f(t.ctx, a, id); } @@ -77,6 +77,7 @@ pub const Text = struct { /// The chat template's ids, as ``apply_chat_template(tokenize=True)`` gives them. pub fn renderIds(t: Text, a: Allocator, messages: json.Value, options: RenderOptions, problem: *[]const u8) Error![]u32 { const text = try t.render(a, messages, options, problem); + defer a.free(text); return t.encode(a, text, false); } diff --git a/zig/src/server/openai.zig b/zig/src/server/openai.zig index 91d02cceb8..a24c79e713 100644 --- a/zig/src/server/openai.zig +++ b/zig/src/server/openai.zig @@ -119,16 +119,24 @@ const Plan = struct { fn plan(srv: *Server, cx: *Cx, is_chat: bool, raw: Value) errors.Refused!Plan { const body = try fields.parseNumbers(cx, raw); try fields.validateModalities(cx, body); + if (body.get("draft")) |draft| { + if (draft == .bool and draft.bool and (srv.info.plain_only or !srv.config.use_drafts)) + return cx.refuse("draft decoding is not available for this served model"); + } // rows need each token's exact bytes: a tokenizer whose decoding depends on context (null) can't give them - const exact_bytes = json.truthyField(body, "logprobs") and srv.info.logprobs and (srv.text.tokenBytes(cx.a, 0) catch return error.OutOfMemory) != null; + const exact_bytes = json.truthyField(body, "logprobs") and srv.info.logprobs and + (srv.text.tokenBytes(cx.a, 0) catch |err| switch (err) { + error.OutOfMemory => return error.OutOfMemory, + error.Template => return cx.fail(.server, "the tokenizer could not decode token bytes", .{}), + }) != null; const logprobs = try fields.probabilityOptions(cx, body, exact_bytes); if (srv.config.request_log) |path| request_log.append(cx.a, path, body); var input: chat.Input = .{ .fields = undefined }; if (is_chat) { - input.messages = try messages.normalize(cx, body.get("messages"), "system", srv.needs_user_after_tool); + input.messages = try validatedMessages(srv, cx, body.get("messages")); input.tools = try tool_specs.active(cx, body.get("tools"), body.get("tool_choice")); } else if (body.get("messages")) |m| if (m == .array and m.array.len > 0) { - input.messages = try messages.normalize(cx, m, "system", srv.needs_user_after_tool); + input.messages = try validatedMessages(srv, cx, m); }; if (!is_chat and input.messages.array.len == 0) input.prompt = try legacyPrompt(srv, cx, body.get("prompt")); // ``body.get("max_tokens") or body.get("max_completion_tokens")``: both are ints or None by now @@ -176,6 +184,19 @@ fn plan(srv: *Server, cx: *Cx, is_chat: bool, raw: Value) errors.Refused!Plan { }; } +pub fn validatedMessages(srv: *Server, cx: *Cx, input: ?Value) errors.Refused!Value { + const v41 = @import("deepseek_v41_text.zig"); + if (!v41.isTemplate(srv.text.templateSource())) + return messages.normalize(cx, input, "system", srv.needs_user_after_tool); + const raw = input orelse return cx.refuse("messages must be a nonempty array"); + var problem: []const u8 = "messages must be a nonempty array"; + v41.validate(raw, &problem) catch |err| { + if (err == error.OutOfMemory) return error.OutOfMemory; + return cx.refuse(problem); + }; + return raw; +} + /// A completion's prompt as the model reads it: token ids as given, anything else as text. fn legacyPrompt(srv: *Server, cx: *Cx, prompt: ?Value) errors.Refused!chat.Prompt { const p = prompt orelse return .{ .text = "" }; @@ -209,7 +230,10 @@ fn promptText(srv: *Server, cx: *Cx, p: Value) errors.Refused![]const u8 { const vocab = srv.text.vocabSize(); const toks = try cx.a.alloc(u32, items.len); for (items, toks) |t, *slot| slot.* = try promptId(cx, vocab, t); - return srv.text.decode(cx.a, toks) catch return error.OutOfMemory; + return srv.text.decode(cx.a, toks) catch |err| switch (err) { + error.OutOfMemory => return error.OutOfMemory, + error.Template => return cx.fail(.server, "the tokenizer could not decode the completion prompt", .{}), + }; } var parts: std.ArrayList([]const u8) = .empty; for (items) |item| try parts.append(cx.a, try promptText(srv, cx, item)); @@ -314,7 +338,7 @@ const Run = struct { fn attachCalls(r: *Run, reply: *chat.Reply) Allocator.Error!?[]Value { const tools = r.plan.input.tools; if (tools.len == 0) return null; - const parsed = try tool_parse.parse(r.a, reply.content, tools, r.plan.policy.maxCalls()); + const parsed = reply.parsed_calls orelse try tool_parse.parse(r.a, reply.content, tools, r.plan.policy.maxCalls()); const calls = parsed.calls orelse { if (r.plan.policy.single) reply.content = try r.plan.policy.parsedContent(r.a, parsed.content); return null; @@ -373,7 +397,7 @@ const Run = struct { if (!r.is_chat) return; const message = json.newObject(r.a) catch return; message.put(r.a, "role", .{ .string = "assistant" }) catch return; - message.put(r.a, "content", if (calls != null) .null else .{ .string = reply.content }) catch return; + message.put(r.a, "content", if (calls != null and reply.content.len == 0) .null else .{ .string = reply.content }) catch return; if (calls) |c| message.put(r.a, "tool_calls", .{ .array = c }) catch return; warm.reply(r.srv, r.plan.input, .{ .object = message }, reply.thinking, reply.effort); } @@ -389,7 +413,7 @@ const Run = struct { if (r.is_chat) { const message = try json.newObject(a); try message.put(a, "role", .{ .string = "assistant" }); - try message.put(a, "content", if (calls != null) .null else .{ .string = reply.content }); + try message.put(a, "content", if (calls != null and reply.content.len == 0) .null else .{ .string = reply.content }); if (reply.reasoning) |t| if (t.len > 0) try message.put(a, "reasoning_content", .{ .string = t }); if (calls) |c| try message.put(a, "tool_calls", .{ .array = c }); try choice.put(a, "message", .{ .object = message }); diff --git a/zig/src/server/process_memory.zig b/zig/src/server/process_memory.zig index 1e8fab92db..fb22215e1c 100644 --- a/zig/src/server/process_memory.zig +++ b/zig/src/server/process_memory.zig @@ -5,7 +5,7 @@ const builtin = @import("builtin"); pub const Snapshot = struct { physical_footprint_bytes: u64, lifetime_peak_physical_footprint_bytes: u64 }; // Darwin's rusage_info_v4 ABI, including the UUID and trailing fields the kernel writes but we never publish. -const RUsage = extern struct { +pub const RUsage = extern struct { uuid: [16]u8, user_time: u64, system_time: u64, @@ -47,11 +47,16 @@ const RUsage = extern struct { extern "c" fn proc_pid_rusage(pid: c_int, flavor: c_int, buffer: *anyopaque) c_int; extern "c" fn getpid() c_int; +/// The existing Darwin ABI, for qualification counters without C translation or extra allocation. +pub fn readUsage(info: *RUsage) bool { + if (builtin.os.tag != .macos) return false; + return proc_pid_rusage(getpid(), 4, info) == 0; +} /// One kernel sample, including peaks between HTTP polls; null on unsupported platforms or failed reads. pub fn read() ?Snapshot { if (builtin.os.tag != .macos) return null; var info: RUsage = undefined; - if (proc_pid_rusage(getpid(), 4, &info) != 0) return null; + if (!readUsage(&info)) return null; return .{ .physical_footprint_bytes = info.phys_footprint, .lifetime_peak_physical_footprint_bytes = info.lifetime_max_phys_footprint, diff --git a/zig/src/server/prompt.zig b/zig/src/server/prompt.zig index e68363bfe7..0a0c3ad676 100644 --- a/zig/src/server/prompt.zig +++ b/zig/src/server/prompt.zig @@ -27,7 +27,8 @@ pub const isTitle = messages_mod.isTitleRequest; /// ``render_prompt_ids``: messages normalized for this template, rendered, and a dangling ```` closed. pub fn renderIds(srv: *Server, cx: *Cx, messages: Value, tools: []const Value, thinking: bool, effort: ?[]const u8, generation: bool) errors.Refused![]const u32 { - const normalized = try messages_mod.toolArguments(cx, try messages_mod.normalize(cx, messages, srv.late_system, srv.needs_user_after_tool)); + const native_v41 = @import("deepseek_v41_text.zig").isTemplate(srv.text.templateSource()); + const normalized = if (native_v41) messages else try messages_mod.toolArguments(cx, try messages_mod.normalize(cx, messages, srv.late_system, srv.needs_user_after_tool)); var problem: []const u8 = ""; const options: model_text.RenderOptions = .{ .tools = if (tools.len > 0) Value{ .array = @constCast(tools) } else null, @@ -40,11 +41,14 @@ pub fn renderIds(srv: *Server, cx: *Cx, messages: Value, tools: []const Value, t error.Template => return cx.fail(.server, "{s}", .{problem}), }; if (!thinking and generation and ids.len > 0) { - const last = try srv.text.decode(cx.a, ids[ids.len - 1 ..]); + const last = srv.text.decode(cx.a, ids[ids.len - 1 ..]) catch |err| switch (err) { + error.OutOfMemory => return error.OutOfMemory, + error.Template => return cx.fail(.server, "the tokenizer cannot decode the generation prompt", .{}), + }; if (std.mem.eql(u8, @import("reply_text.zig").pyStrip(last), "")) { const close = srv.text.encode(cx.a, "", false) catch |e| switch (e) { error.OutOfMemory => return error.OutOfMemory, - error.Template => &[_]u32{}, + error.Template => return cx.fail(.server, "the tokenizer cannot encode the thinking close marker", .{}), }; if (close.len == 1) ids = try std.mem.concat(cx.a, u32, &.{ ids, close }); } diff --git a/zig/src/server/reply_text.zig b/zig/src/server/reply_text.zig index 648a1d8cd2..2bde1b12ad 100644 --- a/zig/src/server/reply_text.zig +++ b/zig/src/server/reply_text.zig @@ -13,6 +13,7 @@ pub const calls = [_][2][]const u8{ .{ "", "" }, .{ "<|tool_call>", "" }, .{ "<\u{ff5c}DSML\u{ff5c}tool_calls>", "" }, + .{ "<｜DSML｜ calls>", "" }, }; /// ``text.find(tag, from)`` led by a vector scan for the tag's first byte: markup is rare in prose, so it is fast. @@ -99,6 +100,47 @@ pub fn hideToolCalls(a: Allocator, text: []const u8, finished: bool) Allocator.E return hideInto(a, &out, text, finished); } +fn namespaceByte(ch: u8) bool { + return std.ascii.isAlphanumeric(ch) or ch == '_' or ch == '.' or ch == '-' or ch >= 0x80; +} + +fn namespacedCall(text: []const u8, from: usize) ?struct { start: usize, open_end: usize } { + var cursor = from; + while (std.mem.indexOfScalarPos(u8, text, cursor, '<')) |start| { + cursor = start + 1; + if (cursor == text.len or text[cursor] == '.' or text[cursor] == '-') continue; + var end = cursor; + while (end < text.len and namespaceByte(text[end])) end += 1; + if (end == cursor or !std.ascii.startsWithIgnoreCase(text[end..], ":tool_call>")) continue; + return .{ .start = start, .open_end = end + ":tool_call>".len }; + } + return null; +} + +fn partialNamespace(text: []const u8) usize { + const start = std.mem.lastIndexOfScalar(u8, text, '<') orelse return 0; + const tail = text[start + 1 ..]; + if (tail.len > 0 and (tail[0] == '.' or tail[0] == '-')) return 0; + var end: usize = 0; + while (end < tail.len and namespaceByte(tail[end])) end += 1; + if (end == tail.len) return text.len - start; + if (end == 0 or tail[end] != ':') return 0; + const suffix = tail[end + 1 ..]; + return if (std.ascii.startsWithIgnoreCase("tool_call>", suffix)) text.len - start else 0; +} + +fn namespacedClose(text: []const u8, open: anytype) ?usize { + const name = text[open.start + 1 .. open.open_end]; + var cursor = open.open_end; + while (std.mem.indexOfScalarPos(u8, text, cursor, '<')) |start| { + cursor = start + 1; + if (std.mem.startsWith(u8, text[start..], " 0 and !srv.info.thinking_budget) { + std.debug.print("tensorfold: --thinking-budget is not supported by this native engine\n", .{}); + return 2; + } if (raiseOpenFileLimit()) |r| log.line("open-file limit raised: soft {d} -> {d}", .{ r[0], r[1] }); const address = resolve(io, args.host, args.port) orelse { std.debug.print("tensorfold: cannot resolve --host {s}\n", .{args.host}); diff --git a/zig/src/server/server.zig b/zig/src/server/server.zig index 255b02a2b6..7579f5f2e3 100644 --- a/zig/src/server/server.zig +++ b/zig/src/server/server.zig @@ -240,7 +240,13 @@ pub const Server = struct { fn callForm(srv: *Server, cx: *Cx) errors.Refused!?Form { const probe = (try json.parse(cx.a, "[{\"role\": \"user\", \"content\": \"x\"}, {\"role\": \"assistant\", \"content\": \"\", \"tool_calls\": [{\"id\": \"call_0\", \"type\": \"function\", \"function\": {\"name\": \"tfprobe_fn\", \"arguments\": {}}}]}]")).ok; var scratch: Cx = .{ .a = cx.a }; - const text = if (prompt.renderIds(srv, &scratch, probe, &.{}, false, null, false)) |ids| try srv.text.decode(cx.a, ids) else |_| ""; + const text = if (prompt.renderIds(srv, &scratch, probe, &.{}, false, null, false)) |ids| + srv.text.decode(cx.a, ids) catch |err| switch (err) { + error.OutOfMemory => return error.OutOfMemory, + error.Template => return cx.other("the tokenizer could not decode the tool call probe"), + } + else |_| + ""; if (std.mem.lastIndexOf(u8, text, "tfprobe_fn")) |at| { var best: ?usize = null; var opener: []const u8 = ""; diff --git a/zig/src/server/slide_think.zig b/zig/src/server/slide_think.zig index dce2cb1ba6..9f32b643ee 100644 --- a/zig/src/server/slide_think.zig +++ b/zig/src/server/slide_think.zig @@ -35,7 +35,10 @@ pub fn thinkingHead(srv: *Server, cx: *Cx, question: []const u8) !?[]const u32 { const asked = try a.dupe(Value, &.{try message(a, "user", question)}); const head = try prompt.renderIds(srv, cx, .{ .array = asked }, &.{}, true, null, true); if (head.len == 0) return null; - const tail = try srv.text.decode(a, head[head.len - @min(head.len, 3) ..]); + const tail = srv.text.decode(a, head[head.len - @min(head.len, 3) ..]) catch |e| switch (e) { + error.Template => return cx.other("the tokenizer failed to decode the thinking prompt"), + error.OutOfMemory => return error.OutOfMemory, + }; if (!std.mem.endsWith(u8, std.mem.trimEnd(u8, tail, " \t\r\n"), "")) return null; return head; } diff --git a/zig/src/server/stream_preflight_test.zig b/zig/src/server/stream_preflight_test.zig index d3fb16c7e3..c2b550ab38 100644 --- a/zig/src/server/stream_preflight_test.zig +++ b/zig/src/server/stream_preflight_test.zig @@ -10,21 +10,33 @@ const responses_translate = @import("responses_translate.zig"); const errors = @import("errors.zig"); const Conn = @import("http_conn.zig").Conn; -/// Bytes as tokens; every prompt renders as "prompt" (6 tokens). +/// Unicode scalars as tokens; prompts are six tokens unless the official V41 renderer is selected. const TestText = struct { - fn text(t: *@This()) model_text.Text { + v41: bool = false, + pub fn text(t: *@This()) model_text.Text { return .{ .ctx = t, .vtable = &.{ .encode = encode, .decode = decode, .token_id = tokenId, .token_string = tokenString, .vocab_size = vocabSize, .eos_ids = eosIds, .render = render, .template_source = templateSource } }; } - fn encode(_: *anyopaque, a: std.mem.Allocator, input: []const u8, _: bool) model_text.Error![]u32 { - const ids = try a.alloc(u32, input.len); - for (input, ids) |byte, *id| id.* = byte; + pub fn encode(_: *anyopaque, a: std.mem.Allocator, input: []const u8, _: bool) model_text.Error![]u32 { + const view = std.unicode.Utf8View.init(input) catch return error.Template; + const count = std.unicode.utf8CountCodepoints(input) catch return error.Template; + const ids = try a.alloc(u32, count); + var it = view.iterator(); + for (ids) |*id| id.* = it.nextCodepoint().?; return ids; } fn decode(_: *anyopaque, a: std.mem.Allocator, ids: []const u32) std.mem.Allocator.Error![]u8 { - const decoded = try a.alloc(u8, ids.len); - for (ids, decoded) |id, *byte| byte.* = @intCast(id); + var count: usize = 0; + var scalar: [4]u8 = undefined; + for (ids) |id| count += std.unicode.utf8Encode(@intCast(id), &scalar) catch unreachable; + const decoded = try a.alloc(u8, count); + var offset: usize = 0; + for (ids) |id| { + const size = std.unicode.utf8Encode(@intCast(id), &scalar) catch unreachable; + @memcpy(decoded[offset..][0..size], scalar[0..size]); + offset += size; + } return decoded; } @@ -44,18 +56,29 @@ const TestText = struct { return &.{}; } - fn render(_: *anyopaque, a: std.mem.Allocator, _: json.Value, _: model_text.RenderOptions, _: *[]const u8) model_text.Error![]u8 { - return a.dupe(u8, "prompt"); + fn render(ctx: *anyopaque, a: std.mem.Allocator, messages: json.Value, options: model_text.RenderOptions, problem: *[]const u8) model_text.Error![]u8 { + const t: *TestText = @ptrCast(@alignCast(ctx)); + if (!t.v41) return a.dupe(u8, "prompt"); + const v41 = @import("deepseek_v41_text.zig"); + const hf = @import("hf_text.zig"); + var context: std.json.ObjectMap = .empty; + try context.put(a, "enable_thinking", .{ .bool = options.enable_thinking }); + const prepared = try v41.prepare(a, messages, problem); + return @import("template").renderChat(a, @embedFile("deepseek_v41_official.jinja"), try hf.toStd(a, prepared), if (options.tools) |tools| try hf.toStd(a, tools) else .null, .{ .object = context }, options.add_generation_prompt, null) catch |err| { + if (err == error.OutOfMemory) return error.OutOfMemory; + return error.Template; + }; } - fn templateSource(_: *anyopaque) []const u8 { - return ""; + fn templateSource(ctx: *anyopaque) []const u8 { + const t: *TestText = @ptrCast(@alignCast(ctx)); + return if (t.v41) @embedFile("deepseek_v41_official.jinja") else ""; } }; /// An 8-token window whose engine answers "hi" at once. const HiEngine = struct { - fn engine(e: *@This()) api.Engine { + pub fn engine(e: *@This()) api.Engine { return .{ .ctx = e, .vtable = &.{ .info = info, .submit = submit, .cancel = cancel, .status = status, .memory = memory } }; } @@ -89,8 +112,12 @@ const CallEngine = struct { for (reply, &out) |byte, *id| id.* = byte; break :blk out; }; + tokens: []const u32 = &ids, + batched: bool = false, + prompt: [4096]u32 = undefined, + prompt_len: usize = 0, - fn engine(e: *@This()) api.Engine { + pub fn engine(e: *@This()) api.Engine { return .{ .ctx = e, .vtable = &.{ .info = info, .submit = submit, .cancel = HiEngine.cancel, .status = HiEngine.status, .memory = HiEngine.memory } }; } @@ -98,9 +125,14 @@ const CallEngine = struct { return .{ .context_window = 4096 }; } - fn submit(_: *anyopaque, id: api.Id, _: *const api.Request, sink: api.Sink) api.SubmitError!void { + fn submit(ctx: *anyopaque, id: api.Id, request: *const api.Request, sink: api.Sink) api.SubmitError!void { + const e: *CallEngine = @ptrCast(@alignCast(ctx)); + e.prompt_len = request.prompt.len; + @memcpy(e.prompt[0..e.prompt_len], request.prompt); sink.event(sink.ctx, id, &.{ .prefilled = 0 }); - for (0..ids.len) |i| sink.event(sink.ctx, id, &.{ .tokens = ids[i..][0..1] }); + if (e.batched) { + sink.event(sink.ctx, id, &.{ .tokens = e.tokens }); + } else for (0..e.tokens.len) |i| sink.event(sink.ctx, id, &.{ .tokens = e.tokens[i..][0..1] }); sink.event(sink.ctx, id, &.{ .finished = .{ .reason = .stop } }); } }; @@ -115,17 +147,40 @@ const Seen = struct { context_code: bool = false, label: ?enum { exact, greedy } = null, preparing: i64 = 0, - text: [64]u8 = undefined, // the streamed delta content, or the whole reply's message content + text: [4096]u8 = undefined, // bounded complete content for legacy-envelope regressions text_len: usize = 0, calls: usize = 0, // the tool calls the stream started, or the whole reply's tool calls + got_error: bool = false, + arguments: [1024]u8 = undefined, + arguments_len: usize = 0, + argument_fragments: usize = 0, + reasoning: [256]u8 = undefined, + reasoning_len: usize = 0, + + fn keepArguments(c: *Seen, call: json.Value) void { + const f = call.get("function") orelse return; + const args = f.strField("arguments") orelse return; + if (args.len > 0) c.argument_fragments += 1; + const n = @min(args.len, c.arguments.len - c.arguments_len); + @memcpy(c.arguments[c.arguments_len..][0..n], args[0..n]); + c.arguments_len += n; + } - fn content(c: *const Seen) []const u8 { + fn keepReasoning(c: *Seen, value: json.Value) void { + if (value != .string) return; + const n = @min(value.string.len, c.reasoning.len - c.reasoning_len); + @memcpy(c.reasoning[c.reasoning_len..][0..n], value.string[0..n]); + c.reasoning_len += n; + } + + pub fn content(c: *const Seen) []const u8 { return c.text[0..c.text_len]; } fn keep(c: *Seen, part: json.Value) void { if (part != .string) return; - const n = @min(part.string.len, c.text.len - c.text_len); + std.debug.assert(part.string.len <= c.text.len - c.text_len); + const n = part.string.len; @memcpy(c.text[c.text_len..][0..n], part.string[0..n]); c.text_len += n; } @@ -145,14 +200,17 @@ const Seen = struct { c.events += 1; const p = payload orelse return; if (p.get("error")) |err| { + c.got_error = true; const code = err.get("code") orelse return; c.context_code = code == .string and std.mem.eql(u8, code.string, "context_length_exceeded"); return; } if (p.get("choices")) |ch| if (ch == .array and ch.array.len > 0) if (ch.array[0].get("delta")) |d| { if (d.get("content")) |t| c.keep(t); + if (d.get("reasoning_content")) |t| c.keepReasoning(t); if (d.get("tool_calls")) |t| if (t == .array) for (t.array) |call| { if (call.get("id") != null) c.calls += 1; + c.keepArguments(call); }; }; if (c.events != 1) return; @@ -170,10 +228,13 @@ const Seen = struct { }; if (payload.get("choices")) |ch| if (ch == .array and ch.array.len > 0) if (ch.array[0].get("message")) |m| { if (m.get("content")) |t| c.keep(t); + if (m.get("reasoning_content")) |t| c.keepReasoning(t); if (m.get("tool_calls")) |t| if (t == .array) { c.calls += t.array.len; + for (t.array) |call| c.keepArguments(call); }; }; + c.got_error = payload.get("error") != null; const code = (payload.get("error") orelse return).get("code") orelse return; c.context_code = code == .string and std.mem.eql(u8, code.string, "context_length_exceeded"); } @@ -303,6 +364,116 @@ test "a streamed Responses tool reply sends no whitespace that the whole reply d "\"tools\":[{\"type\":\"function\",\"name\":\"lookup\",\"parameters\":{\"type\":\"object\"}}]"); } +test "malformed and unoffered V41 DSML fails both whole and streamed chat" { + const inputs = [_][]const u8{ + "<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"lookup\">\n", + "<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"lookup\">\n<｜DSML｜ parameter name=\"x\" string=\"false\">bad\n\n", + "<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"not_offered\">\n\n", + }; + for (inputs) |raw| { + const tokens = try TestText.encode(undefined, std.testing.allocator, raw, false); + defer std.testing.allocator.free(tokens); + var backend: CallEngine = .{ .tokens = tokens }; + const body = "{\"model\":\"test-model\",\"max_tokens\":1000,\"messages\":[{\"role\":\"user\",\"content\":\"x\"}]," ++ + "\"tools\":[{\"type\":\"function\",\"function\":{\"name\":\"lookup\",\"parameters\":{\"type\":\"object\"}}}]"; + const whole = try sendWith(backend.engine(), .chat, body ++ ",\"stream\":false}", .{}); + const streamed = try sendWith(backend.engine(), .chat, body ++ ",\"stream\":true}", .{}); + try std.testing.expectEqual(@as(?u16, 500), whole.status); + try std.testing.expect(whole.got_error and streamed.got_error); + try std.testing.expectEqual(@as(usize, 0), whole.calls); + try std.testing.expectEqual(@as(usize, 0), streamed.calls); + } +} + +test "legacy calls inside V41 DSML parameters stay literal in whole and streamed replies" { + const literal = ""; + const raw = "<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"lookup\">\n" ++ + "<｜DSML｜ parameter name=\"q\" string=\"true\">" ++ literal ++ "\n" ++ + "\n"; + const tokens = try TestText.encode(undefined, std.testing.allocator, raw, false); + defer std.testing.allocator.free(tokens); + const body = "{\"model\":\"test-model\",\"max_tokens\":1000,\"messages\":[{\"role\":\"user\",\"content\":\"x\"}]," ++ + "\"tools\":[{\"type\":\"function\",\"function\":{\"name\":\"lookup\",\"parameters\":{\"type\":\"object\"}}}]"; + for ([_]bool{ false, true }) |batched| { + var backend: CallEngine = .{ .tokens = tokens, .batched = batched }; + const whole = try sendWith(backend.engine(), .chat, body ++ ",\"stream\":false}", .{}); + const streamed = try sendWith(backend.engine(), .chat, body ++ ",\"stream\":true}", .{}); + try std.testing.expectEqual(@as(?u16, 200), whole.status); + try std.testing.expect(!whole.got_error and !streamed.got_error); + try std.testing.expectEqual(@as(usize, 1), whole.calls); + try std.testing.expectEqual(@as(usize, 1), streamed.calls); + try std.testing.expectEqualStrings("{\"q\":\"" ++ literal ++ "\"}", whole.arguments[0..whole.arguments_len]); + try std.testing.expectEqualStrings(whole.arguments[0..whole.arguments_len], streamed.arguments[0..streamed.arguments_len]); + } +} + +test "complete V41 calls recovered from unclosed thinking agree in both reply modes" { + const raw = "reason\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"lookup\">\n\n"; + const tokens = try TestText.encode(undefined, std.testing.allocator, raw, false); + defer std.testing.allocator.free(tokens); + var backend: CallEngine = .{ .tokens = tokens }; + const body = "{\"model\":\"test-model\",\"max_tokens\":1000,\"messages\":[{\"role\":\"user\",\"content\":\"x\"}]," ++ + "\"chat_template_kwargs\":{\"enable_thinking\":true}," ++ + "\"tools\":[{\"type\":\"function\",\"function\":{\"name\":\"lookup\",\"parameters\":{\"type\":\"object\"}}}]"; + const whole = try sendWith(backend.engine(), .chat, body ++ ",\"stream\":false}", .{}); + const streamed = try sendWith(backend.engine(), .chat, body ++ ",\"stream\":true}", .{}); + try std.testing.expectEqual(@as(?u16, 200), whole.status); + try std.testing.expect(!whole.got_error and !streamed.got_error); + try std.testing.expectEqual(@as(usize, 1), whole.calls); + try std.testing.expectEqual(@as(usize, 1), streamed.calls); + try std.testing.expectEqualStrings("{}", whole.arguments[0..whole.arguments_len]); + try std.testing.expectEqualStrings("{}", streamed.arguments[0..streamed.arguments_len]); + try std.testing.expectEqualStrings("reason", whole.reasoning[0..whole.reasoning_len]); + try std.testing.expectEqualStrings("reason", @import("reply_text.zig").pyStrip(streamed.reasoning[0..streamed.reasoning_len])); + try std.testing.expectEqualStrings(whole.content(), streamed.content()); +} + +test "Anthropic V41 counting and generation preserve the same multipart prompt" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var text: TestText = .{ .v41 = true }; + var backend: CallEngine = .{ .tokens = &.{ 'h', 'i' } }; + var srv = try server.Server.init(std.testing.allocator, std.testing.io, backend.engine(), text.text(), .{ + .served_name = "test-model", + .model_ids = &.{"test-model"}, + .enable_thinking = false, + .use_drafts = false, + }, null); + defer srv.deinit(); + var cx: errors.Cx = .{ .a = a }; + const body = (try json.parse(a, "{\"model\":\"test-model\",\"max_tokens\":2,\"thinking\":{\"type\":\"disabled\"}," ++ + "\"messages\":[{\"role\":\"user\",\"content\":[{\"type\":\"text\",\"text\":\"one\"},{\"type\":\"text\",\"text\":\"two\"}]}]}")).ok; + const count_chat = try anthropic_translate.translate(&cx, body, true); + const count = try @import("anthropic.zig").countTokens(srv, &cx, count_chat); + var seen: Seen = .{}; + var conn: Conn = .{ .fd = -1, .peer = "", .buf = &.{}, .gpa = a }; + openai.run(srv, a, seen.out(), .{ .conn = &conn }, true, try anthropic_translate.translate(&cx, body, false)); + try std.testing.expectEqual(@as(?u16, 200), seen.status); + try std.testing.expectEqual(count, backend.prompt_len); + const expected = "<｜begin▁of▁sentence｜><｜User｜>one\n\ntwo<｜Assistant｜>"; + try std.testing.expectEqualStrings(expected, try TestText.decode(undefined, a, backend.prompt[0..backend.prompt_len])); +} + +test "valid V41 DSML has the same call and prose in whole and streamed chat" { + const raw = "summary\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"lookup\">\n" ++ + "<｜DSML｜ parameter name=\"image\" string=\"true\">cat.jpg\n" ++ + "\n"; + const tokens = try TestText.encode(undefined, std.testing.allocator, raw, false); + defer std.testing.allocator.free(tokens); + var backend: CallEngine = .{ .tokens = tokens }; + const body = "{\"model\":\"test-model\",\"max_tokens\":1000,\"messages\":[{\"role\":\"user\",\"content\":\"x\"}]," ++ + "\"tools\":[{\"type\":\"function\",\"function\":{\"name\":\"lookup\",\"parameters\":{\"type\":\"object\"}}}]"; + const whole = try sendWith(backend.engine(), .chat, body ++ ",\"stream\":false}", .{}); + const streamed = try sendWith(backend.engine(), .chat, body ++ ",\"stream\":true}", .{}); + try std.testing.expectEqual(@as(?u16, 200), whole.status); + try std.testing.expect(!whole.got_error and !streamed.got_error); + try std.testing.expectEqual(@as(usize, 1), whole.calls); + try std.testing.expectEqual(@as(usize, 1), streamed.calls); + try std.testing.expectEqualStrings("summary", whole.content()); + try std.testing.expectEqualStrings(whole.content(), streamed.content()); +} + /// HiEngine's replies from an engine that decodes greedily only. const GreedyEngine = struct { fn engine(e: *@This()) api.Engine { @@ -310,10 +481,23 @@ const GreedyEngine = struct { } fn info(_: *anyopaque) api.Info { - return .{ .context_window = 4096, .greedy_only = true }; + return .{ .context_window = 4096, .greedy_only = true, .thinking_budget = false }; } }; +test "unsupported thinking budgets are refused before opening a chat stream" { + var backend: GreedyEngine = .{}; + const body = "{\"model\":\"test-model\",\"messages\":[{\"role\":\"user\",\"content\":\"x\"}]," ++ + "\"max_tokens\":2,\"temperature\":0,\"thinking_budget\":5,\"chat_template_kwargs\":{\"enable_thinking\":true}"; + const whole = try sendWith(backend.engine(), .chat, body ++ ",\"stream\":false}", .{}); + const streamed = try sendWith(backend.engine(), .chat, body ++ ",\"stream\":true}", .{}); + try std.testing.expectEqual(@as(?u16, 400), whole.status); + try std.testing.expectEqual(@as(?u16, 400), streamed.status); + try std.testing.expect(!streamed.opened); + try std.testing.expectEqual(@as(i64, 0), streamed.preparing); + try @import("admission_consumer_regressions.zig").budgetWithoutThinking(backend.engine(), sendWith); +} + test "a greedy-only engine's sampled request is a 400 before the stream opens; greedy ones are served" { var backend: GreedyEngine = .{}; const sampled = try sendWith(backend.engine(), .chat, "{\"model\":\"test-model\",\"messages\":[{\"role\":\"user\",\"content\":\"x\"}],\"stream\":true,\"max_tokens\":2,\"temperature\":0.7}", .{}); @@ -343,3 +527,73 @@ test "a greedy-only engine's refusal names --temperature 0 when the server's own try std.testing.expect(try srv.resolveSampling(&cx, (try json.parse(a, "{\"temperature\":0}")).ok, 0, &.{}) == null); } } + +test "nested legacy markup in malformed V41 DSML is never executable" { + const prefix = "<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"lookup\">\n" ++ + "<｜DSML｜ parameter name=\"q\" string=\"true\">"; + const inputs = [_][]const u8{ + prefix, + prefix ++ "\n", + "reason\n" ++ prefix ++ "\n", + "reason\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"not_offered\">\n\n", + }; + for (inputs, 0..) |raw, i| { + const tokens = try TestText.encode(undefined, std.testing.allocator, raw, false); + defer std.testing.allocator.free(tokens); + var backend: CallEngine = .{ .tokens = tokens }; + const body = try std.fmt.allocPrint(std.testing.allocator, "{{\"model\":\"test-model\",\"max_tokens\":1000,\"messages\":[{{\"role\":\"user\",\"content\":\"x\"}}]," ++ + "\"chat_template_kwargs\":{{\"enable_thinking\":{s}}}," ++ + "\"tools\":[{{\"type\":\"function\",\"function\":{{\"name\":\"lookup\",\"parameters\":{{\"type\":\"object\"}}}}}}]", .{if (i >= 2) "true" else "false"}); + defer std.testing.allocator.free(body); + const whole_body = try std.mem.concat(std.testing.allocator, u8, &.{ body, ",\"stream\":false}" }); + defer std.testing.allocator.free(whole_body); + const streamed_body = try std.mem.concat(std.testing.allocator, u8, &.{ body, ",\"stream\":true}" }); + defer std.testing.allocator.free(streamed_body); + const whole = try sendWith(backend.engine(), .chat, whole_body, .{}); + const streamed = try sendWith(backend.engine(), .chat, streamed_body, .{}); + try std.testing.expectEqual(@as(?u16, 500), whole.status); + try std.testing.expect(whole.got_error and streamed.got_error); + try std.testing.expectEqual(@as(usize, 0), whole.calls); + try std.testing.expectEqual(@as(usize, 0), streamed.calls); + } +} + +test "legacy XML recovered from unclosed thinking retains its whole-arguments fallback" { + const raw = "reason\n\n\n\nfast cars\n\n" ++ + "\n5\n\n\n"; + const tokens = try TestText.encode(undefined, std.testing.allocator, raw, false); + defer std.testing.allocator.free(tokens); + var backend: CallEngine = .{ .tokens = tokens }; + const body = "{\"model\":\"test-model\",\"max_tokens\":1000,\"messages\":[{\"role\":\"user\",\"content\":\"x\"}]," ++ + "\"chat_template_kwargs\":{\"enable_thinking\":true}," ++ + "\"tools\":[{\"type\":\"function\",\"function\":{\"name\":\"lookup\",\"parameters\":{\"type\":\"object\"," ++ + "\"properties\":{\"limit\":{\"type\":\"integer\"}}}}}]"; + const whole = try sendWith(backend.engine(), .chat, body ++ ",\"stream\":false}", .{}); + const streamed = try sendWith(backend.engine(), .chat, body ++ ",\"stream\":true}", .{}); + try std.testing.expectEqual(@as(?u16, 200), whole.status); + try std.testing.expect(!whole.got_error and !streamed.got_error); + try std.testing.expectEqual(@as(usize, 1), whole.calls); + try std.testing.expectEqual(@as(usize, 1), streamed.calls); + try std.testing.expectEqualStrings("{\"query\":\"fast cars\",\"limit\":5}", whole.arguments[0..whole.arguments_len]); + try std.testing.expectEqualStrings(whole.arguments[0..whole.arguments_len], streamed.arguments[0..streamed.arguments_len]); + try std.testing.expectEqual(@as(usize, 1), streamed.argument_fragments); +} +test "containment final-review legacy consumer regressions remain supported" { + try @import("legacy_consumer_regressions.zig").run(TestText, CallEngine, sendWith); + try @import("legacy_consumer_regressions.zig").mixedChain(TestText, CallEngine, sendWith); + try @import("decoder_consumer_regressions.zig").run(TestText, HiEngine); +} +test "vision input and audio output are refused before opening a stream" { + try @import("admission_consumer_regressions.zig").media(send); +} +test "explicit drafts are refused rather than ignored when server drafts are disabled" { + const body = "{\"model\":\"test-model\",\"messages\":[{\"role\":\"user\",\"content\":\"x\"}]," ++ + "\"max_tokens\":2,\"draft\":true"; + for ([_][]const u8{ ",\"stream\":false}", ",\"stream\":true}" }) |suffix| { + const raw = try std.mem.concat(std.testing.allocator, u8, &.{ body, suffix }); + defer std.testing.allocator.free(raw); + const seen = try send(.chat, raw); + try std.testing.expectEqual(@as(?u16, 400), seen.status); + try std.testing.expect(!seen.opened); + } +} diff --git a/zig/src/server/tokens.zig b/zig/src/server/tokens.zig index 3a2e351921..888cf0680d 100644 --- a/zig/src/server/tokens.zig +++ b/zig/src/server/tokens.zig @@ -76,7 +76,10 @@ fn detokenize(srv: *Server, cx: *Cx, body: Value) errors.Refused!Value { if (n < 0 or n >= vocab) return cx.fail(.request, "tokens token ids must be in the vocabulary's range 0 to {d}", .{@as(i64, vocab) - 1}); slot.* = @intCast(n); } - const text = srv.text.decode(cx.a, out) catch return error.OutOfMemory; + const text = srv.text.decode(cx.a, out) catch |err| switch (err) { + error.OutOfMemory => return error.OutOfMemory, + error.Template => return cx.other("the tokenizer could not decode the requested tokens"), + }; const o = try json.newObject(cx.a); try o.put(cx.a, "prompt", .{ .string = text }); return .{ .object = o }; diff --git a/zig/src/server/tool_grammar.zig b/zig/src/server/tool_grammar.zig index 27b76df1c9..8c35630a27 100644 --- a/zig/src/server/tool_grammar.zig +++ b/zig/src/server/tool_grammar.zig @@ -35,7 +35,13 @@ pub fn style(rendered: []const u8) ?Style { pub fn structure(srv: *Server, cx: *Cx, tools: []const Value, parallel: bool, thinking: bool) errors.Refused!api.Structure { const probe_turn = (try json.parse(cx.a, "[{\"role\": \"user\", \"content\": \"x\"}, {\"role\": \"assistant\", \"content\": \"\", \"tool_calls\": [{\"id\": \"call_0\", \"type\": \"function\", \"function\": {\"name\": \"" ++ probe ++ "\", \"arguments\": {\"" ++ key ++ "\": \"v\"}}}]}]")).ok; var scratch: Cx = .{ .a = cx.a }; - const rendered = if (prompt.renderIds(srv, &scratch, probe_turn, &.{}, false, null, false)) |ids| try srv.text.decode(cx.a, ids) else |_| ""; + const rendered = if (prompt.renderIds(srv, &scratch, probe_turn, &.{}, false, null, false)) |ids| + srv.text.decode(cx.a, ids) catch |err| switch (err) { + error.OutOfMemory => return error.OutOfMemory, + error.Template => return cx.other("the tokenizer could not decode the tool grammar probe"), + } + else |_| + ""; const s = style(rendered) orelse return cx.refuse("tool_choice \"required\" or a named function needs a chat template that writes calls (XML parameters or a JSON object), and this one does not: send \"auto\""); const text = try build(cx.a, s, tools, parallel); return srv.compiledStructure(cx, .structural_tag, text, "tool_choice \"required\" or a named function", thinking); diff --git a/zig/src/server/tool_params.zig b/zig/src/server/tool_params.zig index c59d11cfbe..5985e187a4 100644 --- a/zig/src/server/tool_params.zig +++ b/zig/src/server/tool_params.zig @@ -30,7 +30,7 @@ pub fn schemas(a: Allocator, tools: []const Value) Allocator.Error!Schemas { if (properties == .object) for (properties.object.keys(), properties.object.values()) |k, v| { if (v == .object) try kept.put(a, k, v); }; - const name = if (function.get("name")) |n| try tool_specs.pyStr(a, n) else ""; + const name = try tool_specs.toolName(a, tool); try out.put(a, try std.ascii.allocLowerString(a, name), kept); } return out; diff --git a/zig/src/server/tool_parse.zig b/zig/src/server/tool_parse.zig index 58ce898dd7..c81303bd33 100644 --- a/zig/src/server/tool_parse.zig +++ b/zig/src/server/tool_parse.zig @@ -15,17 +15,27 @@ pub const deltas = tool_streaming.deltas; pub const wrapCalls = tool_streaming.wrapCalls; const Error = error{ Invalid, OutOfMemory }; -const Call = struct { name: []const u8, arguments: Value }; -const Envelope = struct { start: usize, end: usize, payload: []const u8 }; +const Call = struct { name: []const u8, arguments: Value, namespaced: bool = false }; +const Envelope = struct { start: usize, end: usize, payload: []const u8, complete: bool = true }; pub const Parsed = struct { content: []const u8, calls: ?[]Value }; -const dsml_open = "<\u{ff5c}DSML\u{ff5c}tool_calls>"; -const dsml_close = ""; -const invoke_open = "<\u{ff5c}DSML\u{ff5c}invoke name=\""; -const invoke_close = ""; -const param_open = "<\u{ff5c}DSML\u{ff5c}parameter name=\""; -const param_close = ""; +const Dsml = struct { open: []const u8, close: []const u8, invoke_open: []const u8, invoke_close: []const u8, param_open: []const u8, param_close: []const u8 }; +const dsml = [_]Dsml{ + .{ .open = "<｜DSML｜tool_calls>", .close = "", .invoke_open = "<｜DSML｜invoke name=\"", .invoke_close = "", .param_open = "<｜DSML｜parameter name=\"", .param_close = "" }, + .{ .open = "<｜DSML｜ calls>", .close = "", .invoke_open = "<｜DSML｜ invoke name=\"", .invoke_close = "", .param_open = "<｜DSML｜ parameter name=\"", .param_close = "" }, +}; + +const DsmlOpening = struct { start: usize, open: []const u8, close: []const u8 }; + +/// The next DSML envelope opener, before any nested markup is interpreted. +pub fn dsmlOpening(text: []const u8, from: usize) ?DsmlOpening { + var first: ?DsmlOpening = null; + for (dsml) |tags| if (find(text, tags.open, from)) |at| { + if (first == null or at < first.?.start) first = .{ .start = at, .open = tags.open, .close = tags.close }; + }; + return first; +} fn findCI(hay: []const u8, needle: []const u8, from: usize) ?usize { if (from > hay.len) return null; @@ -53,12 +63,20 @@ fn envelopes(a: Allocator, text: []const u8) Allocator.Error![]Envelope { var found: std.ArrayList(Envelope) = .empty; var pos: usize = 0; while (findCI(text, "", pos)) |s| { + if (dsmlOpening(text, pos)) |outer| if (outer.start < s) { + pos = (find(text, outer.close, outer.start + outer.open.len) orelse break) + outer.close.len; + continue; + }; const e = findCI(text, "", s + 11) orelse break; try found.append(a, .{ .start = s, .end = e + 12, .payload = strip(text[s + 11 .. e]) }); pos = e + 12; } pos = 0; while (std.mem.indexOfScalarPos(u8, text, pos, '<')) |s| { + if (dsmlOpening(text, pos)) |outer| if (outer.start <= s) { + pos = (find(text, outer.close, outer.start + outer.open.len) orelse break) + outer.close.len; + continue; + }; pos = s + 1; var j = s + 1; if (j >= text.len or !(std.ascii.isAlphabetic(text[j]) or text[j] == '_')) continue; @@ -72,15 +90,21 @@ fn envelopes(a: Allocator, text: []const u8) Allocator.Error![]Envelope { } pos = 0; while (find(text, "<|tool_call>", pos)) |s| { + if (dsmlOpening(text, pos)) |outer| if (outer.start < s) { + pos = (find(text, outer.close, outer.start + outer.open.len) orelse break) + outer.close.len; + continue; + }; const e = find(text, "", s + 12) orelse break; try found.append(a, .{ .start = s, .end = e + 12, .payload = strip(text[s + 12 .. e]) }); pos = e + 12; } pos = 0; - while (find(text, dsml_open, pos)) |s| { - const e = find(text, dsml_close, s + dsml_open.len) orelse break; - try found.append(a, .{ .start = s, .end = e + dsml_close.len, .payload = text[s .. e + dsml_close.len] }); - pos = e + dsml_close.len; + while (dsmlOpening(text, pos)) |outer| { + const s = outer.start; + const e = find(text, outer.close, s + outer.open.len); + const end = if (e) |at| at + outer.close.len else text.len; + try found.append(a, .{ .start = s, .end = end, .payload = text[s..end], .complete = e != null }); + pos = end; } std.mem.sort(Envelope, found.items, {}, struct { fn less(_: void, x: Envelope, y: Envelope) bool { @@ -102,23 +126,26 @@ pub fn parse(a: Allocator, text: []const u8, tools: []const Value, max_calls: ?u const name = try tool_specs.toolName(a, tool); try known.put(a, try std.ascii.allocLowerString(a, name), name); } + if (try bareCalls(a, text, &known, max_calls)) |calls| return .{ .content = "", .calls = calls }; const schemas = try tool_params.schemas(a, tools); const envs = try envelopes(a, text); - if (envs.len == 0) { - if (try bareCalls(a, text, &known, max_calls)) |calls| return .{ .content = "", .calls = calls }; - return .{ .content = text, .calls = null }; - } + if (envs.len == 0) return .{ .content = text, .calls = null }; var calls: std.ArrayList(Value) = .empty; var residue: std.ArrayList(u8) = .empty; var cursor: usize = 0; for (envs) |env| { try residue.appendSlice(a, text[cursor..env.start]); cursor = env.end; + if (!env.complete) { + try residue.appendSlice(a, text[env.start..env.end]); + continue; + } if (max_calls) |m| if (calls.items.len >= m) continue; const parsed: ?[]Call = blk: { - if (std.mem.startsWith(u8, env.payload, dsml_open)) { - const inner = env.payload[dsml_open.len .. env.payload.len - dsml_close.len]; - break :blk dsmlCalls(a, inner) catch |e| switch (e) { + for (dsml) |tags| { + if (!std.mem.startsWith(u8, env.payload, tags.open)) continue; + const inner = env.payload[tags.open.len .. env.payload.len - tags.close.len]; + break :blk dsmlCalls(a, inner, tags) catch |e| switch (e) { error.OutOfMemory => return error.OutOfMemory, error.Invalid => null, }; @@ -138,12 +165,13 @@ pub fn parse(a: Allocator, text: []const u8, tools: []const Value, max_calls: ?u if ((try callName(a, c, &known)) == null) break false; } else true else false; if (!good) { - // a malformed or unoffered call stays text when calls are unlimited, and drops out of a one-call reply - if (max_calls == null) try residue.appendSlice(a, text[env.start..env.end]); + // V4.1 malformed calls stay text even under a one-call policy. + if (max_calls == null or std.mem.startsWith(u8, env.payload, dsml[1].open)) + try residue.appendSlice(a, text[env.start..env.end]); continue; } for (parsed.?) |c| { - if (max_calls == null or calls.items.len < max_calls.?) try calls.append(a, try openaiCall(a, (try callName(a, c, &known)).?, c.arguments)); + if (max_calls == null or calls.items.len < max_calls.?) try calls.append(a, try openaiCall(a, (try callName(a, c, &known)).?, c.arguments, c.namespaced)); } } try residue.appendSlice(a, text[cursor..]); @@ -155,6 +183,10 @@ pub fn closeCall(a: Allocator, text: []const u8, tools: []const Value) Allocator var pos: usize = 0; // the opener ``envelopes`` leaves without a closer, so the one appended pairs with it const start = while (findCI(text, "", pos)) |s| { + if (dsmlOpening(text, pos)) |outer| if (outer.start < s) { + pos = (find(text, outer.close, outer.start + outer.open.len) orelse return "") + outer.close.len; + continue; + }; pos = (findCI(text, "", s + 11) orelse break s) + 12; } else return ""; const block = text[start + 11 ..]; @@ -180,14 +212,16 @@ fn callName(a: Allocator, c: Call, known: *const std.StringArrayHashMapUnmanaged } /// ``_openai_tool_call``: the call under ``name`` (``callName``), arguments as compact JSON. -fn openaiCall(a: Allocator, name: []const u8, arguments: Value) Allocator.Error!Value { +fn openaiCall(a: Allocator, name: []const u8, arguments: Value, namespaced: bool) Allocator.Error!Value { const function = try json.newObject(a); - try function.put(a, "name", .{ .string = name }); + const namespace_end = if (namespaced) std.mem.indexOf(u8, name, "::") else null; + try function.put(a, "name", .{ .string = if (namespace_end) |at| name[at + 2 ..] else name }); try function.put(a, "arguments", .{ .string = try json.stringify(a, arguments, .{ .ascii = false, .compact = true }) }); const call = try json.newObject(a); try call.put(a, "id", .{ .string = try ids.make(a, "call_", 24) }); try call.put(a, "type", .{ .string = "function" }); try call.put(a, "function", .{ .object = function }); + if (namespace_end) |at| try call.put(a, "namespace", .{ .string = name[0..at] }); return .{ .object = call }; } @@ -214,7 +248,7 @@ fn bareCalls(a: Allocator, text: []const u8, known: *const std.StringArrayHashMa const c = parsed orelse return null; // bare JSON naming a tool the request did not offer is a structured answer, the reply's content const name = known.get(try std.ascii.allocLowerString(a, c.name)) orelse return null; - try calls.append(a, try openaiCall(a, name, c.arguments)); + try calls.append(a, try openaiCall(a, name, c.arguments, false)); } return if (calls.items.len > 0) calls.items else null; } @@ -454,37 +488,44 @@ fn gemmaCall(a: Allocator, block: []const u8) Error!?Call { } /// Every invoke of a DSML block, or Invalid when anything in it is not a well-formed invoke. -fn dsmlCalls(a: Allocator, block: []const u8) Error![]Call { +fn dsmlCalls(a: Allocator, block: []const u8, tags: Dsml) Error![]Call { var calls: std.ArrayList(Call) = .empty; var at: usize = 0; var pos: usize = 0; - while (find(block, invoke_open, pos)) |s| { + while (find(block, tags.invoke_open, pos)) |s| { pos = s + 1; - const name_start = s + invoke_open.len; - const q = std.mem.indexOfScalarPos(u8, block, name_start, '"') orelse continue; + const name_start = s + tags.invoke_open.len; + const q = (if (std.mem.eql(u8, tags.open, dsml[1].open)) + find(block, "\">", name_start) + else + std.mem.indexOfScalarPos(u8, block, name_start, '"')) orelse continue; if (!std.mem.startsWith(u8, block[q..], "\">")) continue; - const e = find(block, invoke_close, q + 2) orelse continue; + const e = find(block, tags.invoke_close, q + 2) orelse continue; if (!allSpace(block[at..s])) return error.Invalid; - at = e + invoke_close.len; + at = e + tags.invoke_close.len; pos = at; const body = block[q + 2 .. e]; const arguments = try json.newObject(a); var p: usize = 0; var kept: usize = 0; var residue: std.ArrayList(u8) = .empty; - while (find(body, param_open, p)) |ps| { + while (find(body, tags.param_open, p)) |ps| { p = ps + 1; - const ns = ps + param_open.len; - const nq = std.mem.indexOfScalarPos(u8, body, ns, '"') orelse continue; + const ns = ps + tags.param_open.len; + const nq = (if (std.mem.eql(u8, tags.open, dsml[1].open)) + find(body, "\" string=\"", ns) + else + std.mem.indexOfScalarPos(u8, body, ns, '"')) orelse continue; const flag = body[nq..]; const is_string = std.mem.startsWith(u8, flag, "\" string=\"true\">"); if (!is_string and !std.mem.startsWith(u8, flag, "\" string=\"false\">")) continue; const v_start = nq + (if (is_string) "\" string=\"true\">".len else "\" string=\"false\">".len); - const pe = find(body, param_close, v_start) orelse continue; + const pe = find(body, tags.param_close, v_start) orelse continue; const value = body[v_start..pe]; try residue.appendSlice(a, body[kept..ps]); - kept = pe + param_close.len; + kept = pe + tags.param_close.len; p = kept; + if (std.mem.eql(u8, tags.open, dsml[1].open) and arguments.contains(body[ns..nq])) return error.Invalid; if (is_string) { try arguments.put(a, body[ns..nq], .{ .string = value }); } else switch (try json.parseText(a, value)) { @@ -494,7 +535,7 @@ fn dsmlCalls(a: Allocator, block: []const u8) Error![]Call { } try residue.appendSlice(a, body[kept..]); if (!allSpace(residue.items)) return error.Invalid; - try calls.append(a, .{ .name = strip(block[name_start..q]), .arguments = .{ .object = arguments } }); + try calls.append(a, .{ .name = strip(block[name_start..q]), .arguments = .{ .object = arguments }, .namespaced = std.mem.eql(u8, tags.open, dsml[1].open) }); } if (calls.items.len == 0 or !allSpace(block[at..])) return error.Invalid; return calls.items; diff --git a/zig/src/server/tool_parse_test.zig b/zig/src/server/tool_parse_test.zig index 2c43c3f17e..34d80ba064 100644 --- a/zig/src/server/tool_parse_test.zig +++ b/zig/src/server/tool_parse_test.zig @@ -149,3 +149,145 @@ test "GLM calls keep a declared name with spaces, dots, colons or unicode; an un const prose = "two\nlinesn1"; try std.testing.expect((try parse(a, prose, tools, null)).calls == null); } + +test "unfinished V41 DSML never exposes or completes a nested legacy call" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const tools = (try json.parse(a, "[{\"function\":{\"name\":\"lookup\"}}]")).ok.array; + const prefix = "<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"lookup\">\n<｜DSML｜ parameter name=\"q\" string=\"true\">"; + for ([_][]const u8{ + prefix ++ "", + prefix ++ "", + }) |text| { + try std.testing.expectEqualStrings("", try closeCall(a, text, tools)); + for ([_]?usize{ null, 1 }) |limit| { + const parsed = try parse(a, text, tools, limit); + try std.testing.expect(parsed.calls == null); + try std.testing.expectEqualStrings(text, parsed.content); + } + } +} + +const dsml_prefix = "<｜DSML｜tool_calls><｜DSML｜invoke name=\"lookup\">" ++ + "<｜DSML｜parameter name=\"q\" string=\"true\">"; +const dsml_suffix = ""; +const nested_calls = [_][]const u8{ + "", + "", + "{\"name\":\"lookup\",\"arguments\":{}}", + "{\"name\":\"lookup\",\"arguments\":{}", +}; + +test "DSML parameters contain complete and unfinished legacy calls without consuming adjacent calls" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const tools = (try json.parse(a, "[{\"type\":\"function\",\"function\":{\"name\":\"lookup\"}}]")).ok.array; + const adjacent = "{\"name\":\"lookup\",\"arguments\":{\"n\":7}}"; + for (nested_calls) |literal| { + const text = try std.mem.concat(a, u8, &.{ "Before.", dsml_prefix, literal, dsml_suffix, adjacent, "After." }); + for ([_]?usize{ null, 1 }) |quota| { + const parsed = try parse(a, text, tools, quota); + try std.testing.expectEqual(@as(usize, if (quota == null) 2 else 1), parsed.calls.?.len); + const outer = parsed.calls.?[0].get("function").?; + try std.testing.expectEqualStrings("lookup", outer.strField("name").?); + const arguments = (try json.parse(a, outer.strField("arguments").?)).ok; + try std.testing.expectEqualStrings(literal, arguments.strField("q").?); + if (quota == null) { + try std.testing.expectEqualStrings("{\"n\":7}", parsed.calls.?[1].get("function").?.strField("arguments").?); + } + try std.testing.expectEqualStrings("Before.After.", parsed.content); + try std.testing.expectEqualStrings("", try closeCall(a, text, tools)); + } + } +} + +test "an unfinished DSML envelope never parses or autocompletes its nested legacy call" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const tools = (try json.parse(a, "[{\"type\":\"function\",\"function\":{\"name\":\"lookup\"}}]")).ok.array; + for (nested_calls) |literal| { + const text = try std.mem.concat(a, u8, &.{ "Before.", dsml_prefix, literal }); + for ([_]?usize{ null, 1 }) |quota| { + const parsed = try parse(a, text, tools, quota); + try std.testing.expect(parsed.calls == null); + try std.testing.expectEqualStrings(text, parsed.content); + } + try std.testing.expectEqualStrings("", try closeCall(a, text, tools)); + } +} + +test "closeCall skips a closed DSML parameter's unfinished legacy opener and completes the adjacent call" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const tools = (try json.parse(a, "[{\"type\":\"function\",\"function\":{\"name\":\"lookup\"}}]")).ok.array; + const text = dsml_prefix ++ "" ++ dsml_suffix ++ + "{\"name\":\"lookup\",\"arguments\":{\"n\":7}"; + const tail = try closeCall(a, text, tools); + try std.testing.expectEqualStrings("}", tail); + const parsed = try parse(a, try std.mem.concat(a, u8, &.{ text, tail }), tools, null); + try std.testing.expectEqual(@as(usize, 2), parsed.calls.?.len); + try std.testing.expectEqualStrings("{\"n\":7}", parsed.calls.?[1].get("function").?.strField("arguments").?); +} + +test "bare JSON calls keep literal unfinished DSML markers as argument data" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const tools = (try json.parse(a, "[{\"type\":\"function\",\"function\":{\"name\":\"lookup\"}}]")).ok.array; + const call = "{\"name\":\"lookup\",\"arguments\":{\"q\":\"<｜DSML｜tool_calls>\"}}"; + const inputs = [_][]const u8{ call, "[" ++ call ++ "]", "```json\n" ++ call ++ "\n```" }; + for (inputs) |text| for ([_]?usize{ null, 1 }) |quota| { + const parsed = try parse(a, text, tools, quota); + try std.testing.expectEqual(@as(usize, 1), parsed.calls.?.len); + const arguments = (try json.parse(a, parsed.calls.?[0].get("function").?.strField("arguments").?)).ok; + try std.testing.expectEqualStrings("<｜DSML｜tool_calls>", arguments.strField("q").?); + try std.testing.expectEqualStrings("", parsed.content); + }; +} + +test "both DSML formats contain unfinished alternate protocols without swallowing adjacent calls" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const tools = (try json.parse(a, "[{\"function\":{\"name\":\"lookup\"}}]")).ok.array; + const formats = [_][2][]const u8{ + .{ dsml_prefix, dsml_suffix }, + .{ "<｜DSML｜ calls><｜DSML｜ invoke name=\"lookup\"><｜DSML｜ parameter name=\"q\" string=\"true\">", "" }, + }; + const protocols = [_][2][]const u8{ + .{ "{\"name\":\"lookup\"}", "{\"name\":\"lookup\",\"arguments\":{\"n\":7}}" }, + .{ "<|tool_call>call:lookup{q:<|\"|>literal<|\"|>}", "<|tool_call>call:lookup{n:7}" }, + .{ "<｜DSML｜ calls><｜DSML｜ invoke name=\"lookup\">", "<｜DSML｜ calls><｜DSML｜ invoke name=\"lookup\">" }, + }; + for (formats) |format| for (protocols) |protocol| { + const text = try std.mem.concat(a, u8, &.{ format[0], protocol[0], format[1], protocol[1] }); + for ([_]?usize{ null, 1 }) |quota| { + const parsed = try parse(a, text, tools, quota); + try std.testing.expectEqual(@as(usize, if (quota == null) 2 else 1), parsed.calls.?.len); + const arguments = (try json.parse(a, parsed.calls.?[0].get("function").?.strField("arguments").?)).ok; + try std.testing.expectEqualStrings(protocol[0], arguments.strField("q").?); + try std.testing.expectEqualStrings("", parsed.content); + } + }; +} + +test "bare JSON V41 markers remain literal arguments under either quota" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const tools = (try json.parse(a, "[{\"function\":{\"name\":\"lookup\"}}]")).ok.array; + const call = "{\"name\":\"lookup\",\"arguments\":{\"q\":\"<｜DSML｜ calls>\"}}"; + for ([_][]const u8{ call, "[" ++ call ++ "]", "```json\n" ++ call ++ "\n```" }) |text| { + for ([_]?usize{ null, 1 }) |quota| { + const parsed = try parse(a, text, tools, quota); + try std.testing.expectEqual(@as(usize, 1), parsed.calls.?.len); + const arguments = (try json.parse(a, parsed.calls.?[0].get("function").?.strField("arguments").?)).ok; + try std.testing.expectEqualStrings("<｜DSML｜ calls>", arguments.strField("q").?); + try std.testing.expectEqualStrings("", parsed.content); + } + } +} diff --git a/zig/src/server/tool_specs.zig b/zig/src/server/tool_specs.zig index 1a20f48025..c0c20abe64 100644 --- a/zig/src/server/tool_specs.zig +++ b/zig/src/server/tool_specs.zig @@ -30,8 +30,12 @@ fn nameText(a: Allocator, v: ?Value) Allocator.Error![]const u8 { pub fn toolName(a: Allocator, tool: Value) Allocator.Error![]const u8 { if (tool != .object) return ""; - if (tool.get("function")) |function| if (function == .object) return nameText(a, function.get("name")); - return nameText(a, tool.get("name")); + const function = if (tool.get("function")) |f| (if (f == .object) f else tool) else tool; + const name = try nameText(a, function.get("name")); + var namespace = tool.field("namespace") orelse function.field("namespace") orelse Value.null; + if (namespace == .object) namespace = namespace.get("name") orelse Value.null; + if (namespace != .string or std.mem.indexOf(u8, name, "::") != null) return name; + return std.mem.concat(a, u8, &.{ namespace.string, "::", name }); } fn lowerEq(text: []const u8, want: []const u8) bool { @@ -88,7 +92,7 @@ pub fn active(cx: *Cx, tools: ?Value, choice: ?Value) errors.Refused![]Value { if (!try isFunctionChoice(cx.a, c)) return specs; const function = c.get("function") orelse return cx.other("tool_choice function must include a function object"); if (function != .object) return cx.other("tool_choice function must include a function object"); - const requested = try nameText(cx.a, function.get("name")); + const requested = try toolName(cx.a, c); if (requested.len == 0) return cx.other("tool_choice function must include a name"); var named: std.ArrayList(Value) = .empty; for (specs) |tool| if (std.mem.eql(u8, try toolName(cx.a, tool), requested)) try named.append(cx.a, tool); diff --git a/zig/src/server/tool_stream.zig b/zig/src/server/tool_stream.zig index a02ee06d92..98105368b9 100644 --- a/zig/src/server/tool_stream.zig +++ b/zig/src/server/tool_stream.zig @@ -5,6 +5,7 @@ const ids = @import("ids.zig"); const reply_text = @import("reply_text.zig"); const tool_params = @import("tool_params.zig"); const tool_specs = @import("tool_specs.zig"); +const tool_parse = @import("tool_parse.zig"); const Value = json.Value; const Allocator = std.mem.Allocator; @@ -32,7 +33,7 @@ fn closing(rest: []const u8, text: []const u8) ?usize { return if (std.mem.startsWith(u8, rest[i..], text)) i + text.len else null; } -/// Streams Qwen's XML calls as JSON argument deltas matching the final parser. +/// Streams XML argument deltas and validated DSML envelopes, matching the final parser. pub const Streamer = struct { a: Allocator, schemas: tool_params.Schemas, @@ -42,20 +43,27 @@ pub const Streamer = struct { typed: bool = false, buffer: std.ArrayList(u8) = .empty, pos: usize = 0, - state: enum { outside, call, function, value, closing, abort } = .outside, + state: enum { outside, call, buffered, function, value, closing, abort } = .outside, index: i64 = -1, args_open: bool = false, lead: bool = false, held: []const u8 = "", streamed: bool = false, searched: usize = 0, // the answer is only ever extended, so prose already searched holds no opener + tools: []const Value, + dsml_start: ?usize = null, + dsml_close: []const u8 = "", + invalid_dsml: bool = false, + dsml_searched: usize = 0, + dsml_seen: bool = false, + buffered_start: usize = 0, + buffered_close: []const u8 = "", pub fn init(a: Allocator, tools: []const Value) Allocator.Error!Streamer { - var s: Streamer = .{ .a = a, .schemas = try tool_params.schemas(a, tools) }; + var s: Streamer = .{ .a = a, .schemas = try tool_params.schemas(a, tools), .tools = tools }; for (tools) |tool| { if (tool != .object) continue; - const spec = if (tool.get("function")) |f| (if (f == .object) f else tool) else tool; - const name = if (spec.get("name")) |n| (if (n.truthy()) try tool_specs.pyStr(a, n) else "") else ""; + const name = try tool_specs.toolName(a, tool); if (name.len == 0) continue; try s.known.put(a, try std.ascii.allocLowerString(a, name), name); } @@ -76,31 +84,132 @@ pub const Streamer = struct { return q[1 .. q.len - 1]; } + fn emitCalls(s: *Streamer, calls: []const Value, out: *std.ArrayList(Value)) Allocator.Error!void { + for (try tool_parse.deltas(s.a, calls)) |delta| { + const one = delta.get("tool_calls").?.array[0]; + const index = one.get("index").?.int64().?; + try one.object.put(s.a, "index", try json.intValue(s.a, index + s.index + 1)); + try out.append(s.a, delta); + } + s.index += @intCast(calls.len); + s.streamed = true; + } + /// The deltas the reply's ``text`` (its whole answer so far) adds. pub fn feed(s: *Streamer, text: []const u8, out: *std.ArrayList(Value)) Allocator.Error!void { while (s.state != .abort) { const rest = text[@min(s.pos, text.len)..]; switch (s.state) { .outside => { + if (s.pos == 0) { + const trimmed = reply_text.pyStrip(text); + if (trimmed.len > 0 and (trimmed[0] == '{' or trimmed[0] == '[' or + std.mem.startsWith(u8, trimmed, "```"))) + { + s.state = .abort; + return; + } + } + const dsml_open = reply_text.calls[2][0]; + if (s.dsml_start) |start| { + const from = @max(s.pos, s.dsml_searched -| (s.dsml_close.len - 1)); + const end = reply_text.findTag(text, from, s.dsml_close) orelse { + s.dsml_searched = text.len; + return; + }; + s.pos = end + s.dsml_close.len; + s.searched = s.pos; + s.dsml_start = null; + const parsed = try tool_parse.parse(s.a, text[start..s.pos], s.tools, null); + if (parsed.calls) |calls| { + try s.emitCalls(calls, out); + } else if (std.mem.eql(u8, s.dsml_close, "")) { + s.invalid_dsml = true; + } + } if (text.len < s.searched) s.searched = s.pos; - const from = @min(text.len, @max(s.pos, s.searched -| ("".len - 1))); - const j = reply_text.findTag(text, from, "") orelse { + const from = @min(text.len, @max(s.pos, s.searched -| (dsml_open.len - 1))); + var at = from; + while (std.mem.indexOfScalarPos(u8, text, at, '<')) |start| { + if (tool_parse.dsmlOpening(text, start)) |outer| if (outer.start == start) { + s.dsml_seen = true; + s.dsml_start = start; + s.dsml_close = outer.close; + s.pos = start + outer.open.len; + s.dsml_searched = s.pos; + break; + }; + if (std.mem.startsWith(u8, text[start..], "")) { + s.buffered_start = start; + s.pos = start + "".len; + s.state = .call; + break; + } + if (std.mem.startsWith(u8, text[start..], "<|tool_call>")) { + if (!s.dsml_seen and !s.streamed) { + s.state = .abort; + return; + } + s.buffered_start = start; + s.buffered_close = ""; + s.pos = start + "<|tool_call>".len; + s.dsml_searched = s.pos; + s.state = .buffered; + break; + } + const end = std.mem.indexOfAnyPos(u8, text, start + 1, "<>") orelse { + s.searched = start; + return; + }; + if (text[end] == '<') { + at = end; + continue; + } + if (callTag(text[start + 1 .. end])) { + if (!s.dsml_seen and !s.streamed) { + s.state = .abort; + return; + } + s.buffered_start = start; + s.buffered_close = try std.mem.concat(s.a, u8, &.{ "" }); + s.pos = end + 1; + s.dsml_searched = s.pos; + s.state = .buffered; + break; + } + at = end + 1; + } else { s.searched = text.len; return; - }; - s.pos = j + "".len; - s.state = .call; + } }, .call => { const m = opening(rest, " 256 or (trimmed.len > 0 and trimmed[0] != '<')) s.state = .abort; + if (reply_text.charCount(rest) > 256 or (trimmed.len > 0 and trimmed[0] != '<') or + ((s.dsml_seen or s.streamed) and std.ascii.findIgnoreCase(rest, "") != null)) + { + if (!s.dsml_seen and !s.streamed) { + s.state = .abort; + return; + } + s.buffered_close = ""; + s.dsml_searched = s.pos; + s.state = .buffered; + continue; + } return; }; const lower = try std.ascii.allocLowerString(s.a, reply_text.pyStrip(m.name)); const name = s.known.get(lower) orelse { - s.state = .abort; - return; + if (!s.dsml_seen and !s.streamed) { + s.state = .abort; + return; + } + s.buffered_close = ""; + s.dsml_searched = s.pos; + s.state = .buffered; + continue; }; s.current = name; s.index += 1; @@ -118,6 +227,19 @@ pub const Streamer = struct { s.pos += m.len; s.state = .function; }, + .buffered => { + const tail = s.buffered_close; + const from = @max(s.pos, s.dsml_searched -| (tail.len - 1)); + const end = std.ascii.findIgnoreCasePos(text, from, tail) orelse { + s.dsml_searched = text.len; + return; + }; + s.pos = end + tail.len; + s.searched = s.pos; + s.state = .outside; + const parsed = try tool_parse.parse(s.a, text[s.buffered_start..s.pos], s.tools, null); + if (parsed.calls) |calls| try s.emitCalls(calls, out); + }, .function => { if (opening(rest, "")) return true; + const outer = tool_parse.dsmlOpening(text, s.pos) orelse return false; + return std.mem.eql(u8, outer.close, ""); + } }; /// Bytes of ``rest`` that cannot be the start of a ``tail``-byte closer (Python slices characters). @@ -236,8 +366,11 @@ pub const TextFilter = struct { const end = std.mem.indexOfScalar(u8, rest, '>') orelse { const name = rest[1..]; const colon = std.mem.indexOfScalar(u8, name, ':'); - const possible = name.len == 0 or nameLike(name) or (colon != null and nameLike(name[0..colon.?]) and - std.ascii.startsWithIgnoreCase("tool_call", name[colon.? + 1 ..])); + const possible = name.len == 0 or nameLike(name) or + std.mem.startsWith(u8, "｜DSML｜ calls", name) or + std.mem.startsWith(u8, "｜DSML｜tool_calls", name) or + (colon != null and nameLike(name[0..colon.?]) and + std.ascii.startsWithIgnoreCase("tool_call", name[colon.? + 1 ..])); if (possible) { try f.reset(a, rest); break; @@ -271,6 +404,7 @@ pub const TextFilter = struct { /// ``(?:NAME:)?tool_call``, any case. fn callTag(name: []const u8) bool { if (std.ascii.eqlIgnoreCase(name, "tool_call")) return true; + if (std.mem.eql(u8, name, "｜DSML｜ calls") or std.mem.eql(u8, name, "｜DSML｜tool_calls")) return true; const colon = std.mem.indexOfScalar(u8, name, ':') orelse return false; return nameLike(name[0..colon]) and std.ascii.eqlIgnoreCase(name[colon + 1 ..], "tool_call"); } @@ -303,6 +437,8 @@ pub const Policy = struct { /// ``parsed_content``: a parsed reply's content, whole envelopes the parser left as text kept and the rest filtered, so an unclosed envelope stays hidden. pub fn parsedContent(p: *const Policy, a: Allocator, s: []const u8) Allocator.Error![]const u8 { if (!p.single) return s; + if (std.mem.indexOf(u8, s, "<｜DSML｜ calls>") != null or + std.mem.indexOf(u8, s, "<｜DSML｜tool_calls>") != null) return s; var out: std.ArrayList(u8) = .empty; var cursor: usize = 0; while (try wholeEnvelope(a, s, cursor)) |e| : (cursor = e.end) { @@ -336,69 +472,6 @@ fn wholeEnvelope(a: Allocator, s: []const u8, from: usize) Allocator.Error!?stru return null; } -test "a streamed call its end token left open completes as its closing markup completes it" { - var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); - defer arena.deinit(); - const a = arena.allocator(); - const tools = (try json.parse(a, "[{\"type\":\"function\",\"function\":{\"name\":\"lookup\",\"parameters\":{\"properties\":{\"n\":{\"type\":\"integer\"}}}}}]")).ok.array; - const open = "\n\n\nfast cars\n\n\n5\n\n"; - const sent = try sentArguments(a, try streamed(a, tools, &.{open})); - const ended = try streamed(a, tools, &.{ open, try std.mem.concat(a, u8, &.{ open, try tool_parse.closeCall(a, open, tools) }) }); - const marked = try streamed(a, tools, &.{ open, open ++ "\n" }); - // the arguments sent before the end, closed by closedJson, are the ones the client ends with - try std.testing.expectEqualStrings((try tool_params.closedJson(a, sent)).?, try sentArguments(a, ended)); - try std.testing.expectEqual(marked.len, ended.len); - for (marked[1..], ended[1..]) |x, y| try std.testing.expectEqualStrings(try json.stringify(a, x, .{}), try json.stringify(a, y, .{})); -} - -test "a streamed call cut inside a value stays open" { - var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); - defer arena.deinit(); - const a = arena.allocator(); - const tools = (try json.parse(a, "[{\"type\":\"function\",\"function\":{\"name\":\"lookup\"}}]")).ok.array; - const open = "\n\n\nfast cars and more"; - try std.testing.expectEqualStrings("", try tool_parse.closeCall(a, open, tools)); - // its opening delta went out and cannot be taken back; what it sent does not close under closedJson - const sent = try sentArguments(a, try streamed(a, tools, &.{open})); - try std.testing.expectEqualStrings("{\"q\":\"fast c", sent); - try std.testing.expect((try tool_params.closedJson(a, sent)) == null); -} - -test "a one-call reply drops an unoffered call and still hides an unclosed one" { - var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); - defer arena.deinit(); - const a = arena.allocator(); - const tools = (try json.parse(a, "[{\"type\":\"function\",\"function\":{\"name\":\"lookup\"}}]")).ok.array; - const single: Policy = .{ .single = true }; - const text = "Launching. now"; - // what attachCalls makes of the reply: the parse, then the policy's content - const r = try tool_parse.parse(a, text, tools, single.maxCalls()); - try std.testing.expect(r.calls == null); - try std.testing.expectEqualStrings("Launching.", r.content); - try std.testing.expectEqualStrings("Launching.", try single.parsedContent(a, r.content)); - try std.testing.expectEqualStrings("Launching. ", try single.parsedContent(a, "Launching. {\"name\"")); - // a stream that sent the prose before the block ends with the block - try std.testing.expectEqualStrings(text["Launching. ".len..], single.kept(text, "\nLaunching. ").?); - try std.testing.expect(single.kept("Launching.", "Launching. ") == null); - const parallel: Policy = .{}; - try std.testing.expect(parallel.kept(text, "Launching. ") == null); -} - -const tool_parse = @import("tool_parse.zig"); - -/// The deltas one streamer sends while the reply grows through ``texts``. -fn streamed(a: Allocator, tools: []const Value, texts: []const []const u8) Allocator.Error![]Value { - var s = try Streamer.init(a, tools); - var out: std.ArrayList(Value) = .empty; - for (texts) |t| try s.feed(t, &out); - return out.items; -} - -/// The arguments a client joins from streamed deltas. -fn sentArguments(a: Allocator, deltas: []const Value) Allocator.Error![]const u8 { - var out: std.ArrayList(u8) = .empty; - for (deltas) |d| for (d.get("tool_calls").?.array) |one| { - try out.appendSlice(a, one.get("function").?.get("arguments").?.string); - }; - return out.items; +test { + _ = @import("tool_stream_test.zig"); } diff --git a/zig/src/server/tool_stream_test.zig b/zig/src/server/tool_stream_test.zig new file mode 100644 index 0000000000..45f5408f31 --- /dev/null +++ b/zig/src/server/tool_stream_test.zig @@ -0,0 +1,244 @@ +const std = @import("std"); +const json = @import("json.zig"); +const tool_params = @import("tool_params.zig"); +const stream = @import("tool_stream.zig"); +const Streamer = stream.Streamer; +const Policy = stream.Policy; +const Value = json.Value; +const Allocator = std.mem.Allocator; + +test "a streamed call its end token left open completes as its closing markup completes it" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const tools = (try json.parse(a, "[{\"type\":\"function\",\"function\":{\"name\":\"lookup\",\"parameters\":{\"properties\":{\"n\":{\"type\":\"integer\"}}}}}]")).ok.array; + const open = "\n\n\nfast cars\n\n\n5\n\n"; + const sent = try sentArguments(a, try streamed(a, tools, &.{open})); + const ended = try streamed(a, tools, &.{ open, try std.mem.concat(a, u8, &.{ open, try tool_parse.closeCall(a, open, tools) }) }); + const marked = try streamed(a, tools, &.{ open, open ++ "\n" }); + // the arguments sent before the end, closed by closedJson, are the ones the client ends with + try std.testing.expectEqualStrings((try tool_params.closedJson(a, sent)).?, try sentArguments(a, ended)); + try std.testing.expectEqual(marked.len, ended.len); + for (marked[1..], ended[1..]) |x, y| try std.testing.expectEqualStrings(try json.stringify(a, x, .{}), try json.stringify(a, y, .{})); +} + +test "a streamed call cut inside a value stays open" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const tools = (try json.parse(a, "[{\"type\":\"function\",\"function\":{\"name\":\"lookup\"}}]")).ok.array; + const open = "\n\n\nfast cars and more"; + try std.testing.expectEqualStrings("", try tool_parse.closeCall(a, open, tools)); + // its opening delta went out and cannot be taken back; what it sent does not close under closedJson + const sent = try sentArguments(a, try streamed(a, tools, &.{open})); + try std.testing.expectEqualStrings("{\"q\":\"fast c", sent); + try std.testing.expect((try tool_params.closedJson(a, sent)) == null); +} + +test "a one-call reply drops an unoffered call and still hides an unclosed one" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const tools = (try json.parse(a, "[{\"type\":\"function\",\"function\":{\"name\":\"lookup\"}}]")).ok.array; + const single: Policy = .{ .single = true }; + const text = "Launching. now"; + // what attachCalls makes of the reply: the parse, then the policy's content + const r = try tool_parse.parse(a, text, tools, single.maxCalls()); + try std.testing.expect(r.calls == null); + try std.testing.expectEqualStrings("Launching.", r.content); + try std.testing.expectEqualStrings("Launching.", try single.parsedContent(a, r.content)); + try std.testing.expectEqualStrings("Launching. ", try single.parsedContent(a, "Launching. {\"name\"")); + // a stream that sent the prose before the block ends with the block + try std.testing.expectEqualStrings(text["Launching. ".len..], single.kept(text, "\nLaunching. ").?); + try std.testing.expect(single.kept("Launching.", "Launching. ") == null); + const parallel: Policy = .{}; + try std.testing.expect(parallel.kept(text, "Launching. ") == null); +} + +const tool_parse = @import("tool_parse.zig"); + +/// The deltas one streamer sends while the reply grows through ``texts``. +fn streamed(a: Allocator, tools: []const Value, texts: []const []const u8) Allocator.Error![]Value { + var s = try Streamer.init(a, tools); + var out: std.ArrayList(Value) = .empty; + for (texts) |t| try s.feed(t, &out); + return out.items; +} + +/// The arguments a client joins from streamed deltas. +fn sentArguments(a: Allocator, deltas: []const Value) Allocator.Error![]const u8 { + var out: std.ArrayList(u8) = .empty; + for (deltas) |d| for (d.get("tool_calls").?.array) |one| { + try out.appendSlice(a, one.get("function").?.get("arguments").?.string); + }; + return out.items; +} + +test "DSML containment survives every byte boundary and preserves adjacent call order" { + const prefix = "<｜DSML｜tool_calls><｜DSML｜invoke name=\"lookup\">" ++ + "<｜DSML｜parameter name=\"q\" string=\"true\">"; + const suffix = ""; + const legacy = ""; + for ([_][]const u8{ legacy, "" }) |literal| { + var outer_arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer outer_arena.deinit(); + const outer_a = outer_arena.allocator(); + const text = try std.mem.concat(outer_a, u8, &.{ legacy, prefix, literal, suffix, legacy }); + const expected = try std.mem.concat(outer_a, u8, &.{ "{}{\"q\":", try json.quote(outer_a, literal, .{}), "}{}" }); + for (0..text.len + 1) |boundary| { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const tools = (try json.parse(a, "[{\"type\":\"function\",\"function\":{\"name\":\"lookup\"}}]")).ok.array; + var s = try Streamer.init(a, tools); + var out: std.ArrayList(Value) = .empty; + try s.feed(text[0..boundary], &out); + try s.feed(text, &out); + try s.feed(text, &out); + var calls: i64 = 0; + for (out.items) |delta| for (delta.get("tool_calls").?.array) |call| { + if (call.get("id") == null) continue; + try std.testing.expectEqual(calls, call.get("index").?.int64().?); + try std.testing.expectEqualStrings("lookup", call.get("function").?.strField("name").?); + calls += 1; + }; + try std.testing.expectEqual(@as(i64, 3), calls); + try std.testing.expectEqualStrings(expected, try sentArguments(a, out.items)); + } + } +} + +test "unfinished DSML retains containment through byte feeds and final closeCall" { + const prefix = "<｜DSML｜tool_calls><｜DSML｜invoke name=\"lookup\">" ++ + "<｜DSML｜parameter name=\"q\" string=\"true\">"; + const literals = [_][]const u8{ + "", + "", + "{\"name\":\"lookup\",\"arguments\":{}", + }; + for (literals) |literal| { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const tools = (try json.parse(a, "[{\"type\":\"function\",\"function\":{\"name\":\"lookup\"}}]")).ok.array; + const text = try std.mem.concat(a, u8, &.{ prefix, literal }); + var s = try Streamer.init(a, tools); + var out: std.ArrayList(Value) = .empty; + for (0..text.len + 1) |n| { + try s.feed(text[0..n], &out); + try std.testing.expectEqual(@as(usize, 0), out.items.len); + } + const tail = try tool_parse.closeCall(a, text, tools); + try std.testing.expectEqualStrings("", tail); + try s.feed(try std.mem.concat(a, u8, &.{ text, tail }), &out); + try std.testing.expectEqual(@as(usize, 0), out.items.len); + try std.testing.expect(!s.streamed); + } +} + +test "malformed closed DSML masks its nested XML and resumes the next legitimate call" { + const malformed = "<｜DSML｜tool_calls><｜DSML｜invoke name=\"lookup\">" ++ + "<｜DSML｜parameter name=\"q\" string=\"true\">" ++ + "" ++ + ""; + const legacy = ""; + const text = malformed ++ legacy; + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const tools = (try json.parse(a, "[{\"type\":\"function\",\"function\":{\"name\":\"lookup\"}}]")).ok.array; + var s = try Streamer.init(a, tools); + var out: std.ArrayList(Value) = .empty; + for (0..malformed.len + 1) |n| { + try s.feed(text[0..n], &out); + try std.testing.expectEqual(@as(usize, 0), out.items.len); + } + for (malformed.len + 1..text.len + 1) |n| try s.feed(text[0..n], &out); + var calls: usize = 0; + for (out.items) |delta| for (delta.get("tool_calls").?.array) |call| { + if (call.get("id") != null) calls += 1; + }; + try std.testing.expectEqual(@as(usize, 1), calls); + try std.testing.expectEqualStrings("{}", try sentArguments(a, out.items)); +} + +test "V41 DSML containment and mixed protocols preserve indices at every byte boundary" { + const prefix = "<｜DSML｜ calls><｜DSML｜ invoke name=\"lookup\"><｜DSML｜ parameter name=\"q\" string=\"true\">"; + const suffix = ""; + const xml = ""; + const protocols = [_][2][]const u8{ + .{ "{\"name\":\"lookup\"}", "{\"name\":\"lookup\",\"arguments\":{\"n\":7}}" }, + .{ "<|tool_call>call:lookup{q:<|\"|>literal<|\"|>}", "<|tool_call>call:lookup{n:7}" }, + .{ "", xml }, + }; + var source_arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer source_arena.deinit(); + const source = source_arena.allocator(); + for (protocols) |protocol| { + const text = try std.mem.concat(source, u8, &.{ xml, prefix, protocol[0], suffix, protocol[1] }); + for (0..text.len + 1) |boundary| { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const tools = (try json.parse(a, "[{\"function\":{\"name\":\"lookup\"}}]")).ok.array; + const parsed = try tool_parse.parse(a, text, tools, null); + var s = try Streamer.init(a, tools); + var out: std.ArrayList(Value) = .empty; + try s.feed(text[0..boundary], &out); + try s.feed(text, &out); + try s.feed(text, &out); + var count: i64 = 0; + var expected: std.ArrayList(u8) = .empty; + for (parsed.calls.?) |call| try expected.appendSlice(a, call.get("function").?.strField("arguments").?); + for (out.items) |delta| for (delta.get("tool_calls").?.array) |call| { + if (call.get("id") == null) continue; + try std.testing.expectEqual(count, call.get("index").?.int64().?); + count += 1; + }; + try std.testing.expectEqual(@as(i64, 3), count); + try std.testing.expect(!s.invalidDsml(text)); + try std.testing.expectEqualStrings(expected.items, try sentArguments(a, out.items)); + } + } +} + +test "XML then non XML then DSML never loses later calls to final fallback" { + const xml = ""; + const outer = "<｜DSML｜ calls><｜DSML｜ invoke name=\"lookup\"><｜DSML｜ parameter name=\"q\" string=\"true\">outer" ++ + ""; + const middles = [_][]const u8{ + "{\"name\":\"lookup\",\"arguments\":{\"n\":7}}", + "lookupn7", + "<|tool_call>call:lookup{n:7}", + "{\"name\":\"lookup\",\"arguments\":{\"n\":7}}", + "", + }; + var source_arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer source_arena.deinit(); + const source = source_arena.allocator(); + for (middles) |middle| { + const text = try std.mem.concat(source, u8, &.{ xml, middle, "1 < 2 ", outer, xml }); + for (0..text.len + 1) |boundary| { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const tools = (try json.parse(a, "[{\"function\":{\"name\":\"lookup\"}}]")).ok.array; + const parsed = try tool_parse.parse(a, text, tools, null); + var s = try Streamer.init(a, tools); + var out: std.ArrayList(Value) = .empty; + try s.feed(text[0..boundary], &out); + try s.feed(text, &out); + try s.feed(text, &out); + var count: i64 = 0; + var expected: std.ArrayList(u8) = .empty; + for (parsed.calls.?) |call| try expected.appendSlice(a, call.get("function").?.strField("arguments").?); + for (out.items) |delta| for (delta.get("tool_calls").?.array) |call| { + if (call.get("id") == null) continue; + try std.testing.expectEqual(count, call.get("index").?.int64().?); + count += 1; + }; + try std.testing.expectEqual(@as(i64, @intCast(parsed.calls.?.len)), count); + try std.testing.expectEqualStrings(expected.items, try sentArguments(a, out.items)); + } + } +} diff --git a/zig/src/server/tool_streaming.zig b/zig/src/server/tool_streaming.zig index 700caedd7c..590f5b1087 100644 --- a/zig/src/server/tool_streaming.zig +++ b/zig/src/server/tool_streaming.zig @@ -21,6 +21,7 @@ pub fn deltas(a: Allocator, calls: []const Value) Allocator.Error![]Value { const kind = call.get("type"); try head.put(a, "type", .{ .string = if (kind != null and kind.?.truthy()) try tool_specs.pyStr(a, kind.?) else "function" }); try head.put(a, "function", .{ .object = head_fn }); + if (call.field("namespace")) |namespace| try head.put(a, "namespace", namespace); try out.append(a, try wrapCalls(a, .{ .object = head })); const args = function.get("arguments"); const text = if (args != null and args.?.truthy()) try tool_specs.pyStr(a, args.?) else ""; diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index eeaee2fc7a..dc0e8cb8a9 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -67,3 +67,19 @@ pub const bf16_topk = @import("bf16_topk"); pub const bf16_topk_gpu = @import("bf16_topk_gpu"); pub const tree_commit_gpu = @import("tree_commit_gpu"); + +pub const exl3_projection = @import("core/exl3_projection.zig"); +pub const gpu_counters = @import("core/gpu_counters.zig"); +pub const expert_groups = @import("core/expert_groups.zig"); +pub const deepseek_v41 = @import("families/deepseek_v41/model.zig"); +pub const deepseek_v41_profiler = @import("families/deepseek_v41/profiler.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"); + +test { + _ = @import("core/exl3_projection_gpu_test.zig"); + _ = @import("families/deepseek_v41/engram.zig"); + _ = @import("families/deepseek_v41/csa2.zig"); + _ = @import("families/deepseek_v41/ops_check.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..d3df4d0ac7 --- /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(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(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(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(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(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..d57567957f --- /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(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..c37682bf8e --- /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(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..eb9780db96 --- /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(matrices[0], input, &rows.gate[slot], scratch); + try projection.apply(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(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(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..ae3ba82fab --- /dev/null +++ b/zig/tests/deepseek_v41_reference/operators.zig @@ -0,0 +1,99 @@ +//! 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); + const previous = model.pipelines.projection.decode; + defer model.pipelines.projection.decode = previous; + 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); + for ([_]ex.Decode{ .arithmetic, .hoisted_decode }, 0..) |decode, output| { + model.pipelines.projection.decode = decode; + try ex.encode(&model.pipelines.projection, enc, matrix, .{ .buffer = x }, .{ .buffer = y, .offset = output * 129280 * 4 }, .{ .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, 2 * 129280); + const candidate = actual[129280..][0..matrix.n]; + try projection.applyF32(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], candidate, expected[0..matrix.n], 0..) |a, c, b, column| { + if (!std.math.isFinite(a) or !std.math.isFinite(b)) return error.NonfiniteOperatorOutput; + if (@as(u32, @bitCast(a)) != @as(u32, @bitCast(c))) return error.DecoderChangedOperator; + 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},\"same_decoder_bits\":true,\"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..d5748807d2 --- /dev/null +++ b/zig/tests/deepseek_v41_reference/profile.zig @@ -0,0 +1,56 @@ +//! 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_decode = model.pipelines.projection.decode; + defer model.pipelines.projection.decode = previous_decode; + const previous_grouping = model.grouped_prefill; + defer model.grouped_prefill = previous_grouping; + model.grouped_prefill = false; + model.pipelines.projection.decode = .arithmetic; + 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.pipelines.projection.decode = .hoisted_decode; + 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\",\"baseline_decode\":\"arithmetic\",\"candidate_decode\":\"hoisted_decode\",\"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..ca0f932c22 --- /dev/null +++ b/zig/tests/deepseek_v41_reference/projection.zig @@ -0,0 +1,261 @@ +//! 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 exl3 = @import("tensorfold").exl3_projection; +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: 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 applyF32(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..23d811a532 --- /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(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(matrices[0], &self.input, &self.gate, &self.projection_scratch); + try projection.apply(matrices[1], &self.input, &self.up, &self.projection_scratch); + moeActivation(&self.gate, &self.up, route, self.c, &self.activation); + try projection.apply(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 }); +} diff --git a/zig/tests/server/fake_engine.zig b/zig/tests/server/fake_engine.zig index 0e413f3286..2dfb94effa 100644 --- a/zig/tests/server/fake_engine.zig +++ b/zig/tests/server/fake_engine.zig @@ -90,6 +90,7 @@ pub const ScriptEngine = struct { const prefill_done = std.Io.Clock.awake.now(e.io).toNanoseconds(); var stats: api.Stats = .{ .prefill_seconds = @as(f64, @floatFromInt(@max(0, prefill_done - prefill_start))) / 1e9 }; var reason: api.Reason = .stop; + var message: []const u8 = ""; var done = false; for (script.object.get("chunks").?.array.items) |chunk| { if (delay > 0) std.Io.sleep(e.io, .fromMilliseconds(delay), .awake) catch {}; @@ -109,7 +110,14 @@ pub const ScriptEngine = struct { const start = emitted.items.len; for (tokens.items) |t| { emitted.append(a, t) catch {}; - if (std.mem.indexOfScalar(u32, r.eos, t) != null or (r.stop != null and r.stop.?.check(r.stop.?.ctx, emitted.items))) { + const stopped = std.mem.indexOfScalar(u32, r.eos, t) != null or + (if (r.stop) |stop| stop.check(stop.ctx, emitted.items) catch |err| { + reason = .failed; + message = @errorName(err); + done = true; + break; + } else false); + if (stopped) { reason = .stop; done = true; } else if (emitted.items.len >= r.max_tokens) { @@ -132,6 +140,6 @@ pub const ScriptEngine = struct { e.active -= 1; _ = e.cancelled.remove(id); e.mutex.unlock(e.io); - sink.event(sink.ctx, id, &.{ .finished = .{ .reason = reason, .stats = stats } }); + sink.event(sink.ctx, id, &.{ .finished = .{ .reason = reason, .stats = stats, .message = message } }); } };