diff --git a/CHANGELOG.md b/CHANGELOG.md index 00c804d41c..a8247e5fca 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Docs: the optimization ledger's cells for #2011, #2162 and #2282 say they landed (#2431). - Lane host: a stream paid for at load (`free_streams`, the engine's own sequence) is a standing lease in the KV ledger, so a family that reserves per request admits as many full-window requests as the window plan, not one fewer, and its prompt cache sits beside the same streams; the job on it promises only what it adds, and the lease is not reported at stop. No family on main reserves per request yet, so every current load is unchanged (#1943 step 2). Thanks to @laudney for finding the gap. +- Nemotron Parse on HIP: NVIDIA's `nvidia/Nemotron-3.5-Parse` is served natively from its FP32 safetensors on a Radeon AI PRO R9700 (gfx1201), with its RADIO image tower, convolution neck and mBART decoder. The token arrays match the published example's on three pages, plain and drafted, solo and concurrent (#1941, part of #1927). Thanks to @laudney. - Developers: `test-glm-snapshots-host` runs under `test-cuda-host`, as GLM's other host steps do, so CI builds it and runs its 21 tests; its host-cache test sets `TF_GLM_HOST_CACHE` through `sys.env`, so it links on Windows too (#2426, refs #2031). Thanks to @Bizuayeu. - Conversation states on disk try one shorter compatible state after a refused read before prefilling from the start. Failed reads still name their error in the log. @@ -23,6 +24,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel streams and copy drafts; there is no drafter or prompt cache yet (#1492, part of #1488). Thanks to @jschmied. - HIP: capture a drafter's requested target-layer taps in each stream's ring, with graph-on/off and planted kernel controls (#2314). +- Docs: `/health`'s `memory` on Metal is the process footprint from the families with a reader, not `{}`, and `/memory`'s `prompt_cache_plan` row names every Mac family that reports it, with Qwen3.5-2B under `--prompt-cache-gib` (#2439). Thanks to @Whq02. + ## 1.2.4 - Docs: four subcommand-flag fixes checked against the code: `slide audit --threshold`'s range, `benchmark --context` and `--kv-cache`, `info` on Flash Next, and `run --device` and `--carveout` on CUDA builds (#2348). diff --git a/LICENSES/NVIDIA-OpenMDW-1.1.txt b/LICENSES/NVIDIA-OpenMDW-1.1.txt new file mode 100644 index 0000000000..3618e26f59 --- /dev/null +++ b/LICENSES/NVIDIA-OpenMDW-1.1.txt @@ -0,0 +1,57 @@ +Unless separately indicated, the model weights and NVIDIA-provided source code +and documentation are licensed under the OpenMDW-1.1 License. +The tokenizer included with this model is separately licensed under the +MIT license. See TOKENIZER_ATTRIBUTION.md for the license text, +original copyright notice, source, and NVIDIA modification history. + +------------ +Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + +OpenMDW License Agreement, version 1.1 (OpenMDW-1.1) + +By exercising rights granted to you under this agreement, you accept and agree +to its terms. + +As used in this agreement, "Model Materials" means the materials provided to +you under this agreement, consisting of: (1) one or more machine learning +models (including architecture and parameters); and (2) all related artifacts +(including associated data, documentation and software) that are provided to +you hereunder. + +Subject to your compliance with this agreement, permission is hereby granted, +free of charge, to deal in the Model Materials without restriction, including +under all copyright, patent, database, and trade secret rights included or +embodied therein. + +If you distribute any portion of the Model Materials, you shall retain in your +distribution (1) a copy of this agreement, and (2) all copyright notices and +other notices of origin included in the Model Materials that are applicable to +your distribution. +If you file, maintain, or voluntarily participate in a lawsuit against any +person or entity asserting that the Model Materials directly or indirectly +infringe any patent or copyright, then all rights and grants made to you +hereunder are terminated, unless that lawsuit was in response to a +corresponding lawsuit first brought against you. + +This agreement does not impose any restrictions or obligations with respect to +any use, modification, or sharing of any outputs generated by using the Model +Materials. + +THE MODEL MATERIALS ARE PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS +OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE, TITLE, NONINFRINGEMENT, ACCURACY, OR THE +ABSENCE OF LATENT OR OTHER DEFECTS OR ERRORS, WHETHER OR NOT DISCOVERABLE, ALL +TO THE GREATEST EXTENT PERMISSIBLE UNDER APPLICABLE LAW. + +YOU ARE SOLELY RESPONSIBLE FOR (1) CLEARING RIGHTS OF OTHER PERSONS THAT MAY +APPLY TO THE MODEL MATERIALS OR ANY USE THEREOF, INCLUDING WITHOUT LIMITATION +ANY PERSON'S COPYRIGHTS OR OTHER RIGHTS INCLUDED OR EMBODIED IN THE MODEL +MATERIALS; (2) OBTAINING ANY NECESSARY CONSENTS, PERMISSIONS OR OTHER RIGHTS +REQUIRED FOR ANY USE OF THE MODEL MATERIALS; OR (3) PERFORMING ANY DUE +DILIGENCE OR UNDERTAKING ANY OTHER INVESTIGATIONS INTO THE MODEL MATERIALS OR +ANYTHING INCORPORATED OR EMBODIED THEREIN. +IN NO EVENT SHALL THE PROVIDERS OF THE MODEL MATERIALS BE LIABLE FOR ANY CLAIM, +DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR +OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE MODEL MATERIALS, THE +USE THEREOF OR OTHER DEALINGS THEREIN. diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index d0d7070278..f2f3d6b8f5 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -35,6 +35,23 @@ The native engine installs and links neither. It reads checkpoints in their form `zig/kernels/metal/` repeat the arithmetic of MLX's kernels and mlx-lm's model math, many of them bit for bit. The Python engine, frozen at 0.6.6, installs MLX and mlx-lm as dependencies. +## Nemotron 3.5 Parse on HIP + +The reference tool `tools/zig/nemotron_parse_reference.py` adapts `run_transformers_inference.py` from NVIDIA's +[nvidia/Nemotron-3.5-Parse](https://huggingface.co/nvidia/Nemotron-3.5-Parse) at revision +`4035ffb1d8f8bafcdb867dfaf6f7514f101464f9`, and keeps that file's SPDX header. The native family follows the same +release's code for its math, rounding and resizing; each file is written for TensorFold, and the file named is the +reference: + +- `zig/src/families/nemotron_parse/processor.zig`: `NemotronParseImageProcessor` in `hf_nemotron_parse_processor.py` + (the aspect-preserving resize, the NaViT white padding, the normalization and the patch order). +- `zig/src/families/nemotron_parse/tower.zig` and `decoder.zig`: `RadioWithNeck`, `_VisionTransformer`, + `NemotronParseDecoder`, `_DecoderLayer`, `_GQAAttention` and `_ScaledEmbedding` in `hf_nemotron_parse_modeling.py`. +- `zig/kernels/hip/ops/vl_position.hpp`: `_normalized_2d_position_pages` in `hf_prefix_geometry.py`. + +These release files are Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES, under the NVIDIA OpenMDW License 1.1; see +[the license text](LICENSES/NVIDIA-OpenMDW-1.1.txt). + ## MLX and mlx-lm adaptations In the Python engine, frozen at 0.6.6, the DeltaNet implementations in diff --git a/build.zig b/build.zig index c070382123..bfff1002e2 100644 --- a/build.zig +++ b/build.zig @@ -205,6 +205,7 @@ pub fn build(b: *std.Build) void { const with_hip = b.option(bool, "hip", "Enable the opt-in HIP backend and its host/device targets") orelse false; const hip_options = hip_build.options(b); // declared on every target (#1165), used only with -Dhip=true hip_build.loaderTests(b, test_step); + @import("zig/build/nemotron_parse.zig").hostTests(b, test_step); if (with_hip) hip_build.steps(b, target, optimize, build_options, test_step, hip_options); const with_vulkan = b.option(bool, "vulkan", "Enable the opt-in Vulkan backend and its host/device targets") orelse false; const vulkan_options = vulkan_build.options(b); // declared on every target, used only with -Dvulkan=true diff --git a/docs/api.md b/docs/api.md index 108280b480..d8d22a7559 100644 --- a/docs/api.md +++ b/docs/api.md @@ -963,7 +963,7 @@ configured the body is only `{"status": "ok"}`. An engine that reports unhealthy | `max_batch_size` | integer | The lanes the engine serves at once | | `policy` | string | The engine's policy line; absent when the engine reports none | | `warming` | boolean | Always `false` on the native engines | -| `memory` | object | CUDA: `active`, `cache` and `peak` in bytes (`cache` is 0), plus `limit` when the engine reports one. Metal: `{}` | +| `memory` | object | `active`, `cache` and `peak` in bytes (`cache` is 0), plus `limit` when the engine reports one: CUDA's allocator counts; on Metal the process footprint, from Nemotron, Qwen3.5-2B, Qwen3.8-27B, Qwen3.5-MoE, Flash Next, GLM-5.3-Flash and GLM-5.3 (EXLR8); Flash Next and both GLMs add `limit`, their load limit. `{}` from the other Metal families | | `live` | object | `connections`, `waiting`, `decode_tokens_per_second`, `prefill_tokens_per_second` | | `version` | string | The release this server was built from | | `backend` | string | The backend built in: `metal`, `cuda`, `hip` or `xpu`; absent from a build without one | diff --git a/docs/memory-diagnostics.md b/docs/memory-diagnostics.md index 03aa3d9763..f8c71742e4 100644 --- a/docs/memory-diagnostics.md +++ b/docs/memory-diagnostics.md @@ -17,7 +17,7 @@ zero means a measured or configured zero. The endpoint publishes no paths, UUIDs | `device.active_bytes`, `device.cache_bytes` | Existing backend allocator counters, where implemented. | | `device.peak_bytes` | Backend peak since startup or the last `/health?reset_peak=1`. | | `device.peak_scope` | `since_start_or_health_reset`; distinct from the process lifetime peak. | -| `prompt_cache_plan` | Immutable retained-prefix budget applied at startup, reported by Nemotron, Flash Next and Qwen3.8-27B on Macs. | +| `prompt_cache_plan` | Immutable retained-prefix budget applied at startup, reported by Nemotron, Flash Next, Qwen3.8-27B, Qwen3.5-MoE, Gemma 4, LFM2.5-24B-A2B and LFM3-1.7B on Macs, and by Qwen3.5-2B with `--prompt-cache-gib`. | The cache plan contains `source`, `budget_bytes`, and `explicit_budget`. With OS readings, `source` is `physical_footprint` and the plan also includes the sizing inputs: `ram_bytes`, @@ -30,11 +30,11 @@ rather than merely the presence of an override flag. Without both OS readings, the existing fallback uses an explicit budget (`source: explicit`) or the spare Metal working set up to 16 GiB (`source: metal_working_set`). Its OS inputs and unknown `over_cap` remain omitted. -Nemotron and Flash Next report a disabled retained-prefix cache as a known `budget_bytes: 0`; Qwen3.8-27B omits -`prompt_cache_plan` when its budget is 0. GLM-5.3-Flash and the CUDA server omit it too: GLM keeps -`--prompt-cache-gib` or 16 GiB inside its load limit, and the CUDA server keeps `--prompt-cache-gib` or 16 GiB inside -what its streams leave and names that size in its startup line. The 2B keeps prompt states only with -`--prompt-cache-gib`, and then reports the plan. +Nemotron, Flash Next, LFM2.5-24B-A2B and LFM3-1.7B report a disabled retained-prefix cache as a known `budget_bytes: 0`; +Qwen3.8-27B, Qwen3.5-MoE and Gemma 4 omit `prompt_cache_plan` when their budget is 0. GLM-5.3-Flash and the CUDA server +omit it too: GLM keeps `--prompt-cache-gib` or 16 GiB inside its load limit, and the CUDA server keeps `--prompt-cache-gib` +or 16 GiB inside what its streams leave and names that size in its startup line. The 2B keeps prompt states only with +`--prompt-cache-gib`, and then reports the plan, `budget_bytes: 0` included. This plan is a cache allowance, not a process limit, current cache occupancy, or an admission guarantee. `/metrics` keeps `tensorfold:process_footprint_bytes` and adds `tensorfold:process_footprint_peak_bytes` diff --git a/docs/recipes/README.md b/docs/recipes/README.md index e1ae0136f7..6abb258dff 100644 --- a/docs/recipes/README.md +++ b/docs/recipes/README.md @@ -10,6 +10,7 @@ For agent-assisted contributions, read the [TensorFold contributor skill](SKILL. | dots.mocr | [Native HIP OCR, BF16, one GPU](dots-mocr-hip.md) | | Kimi K3 / K3.1 | [Native Metal bring-up, synthetic only, resident budget refusal](kimi-k3-metal.md) | | Nemotron 3.5 Lightning | [MLX](nemotron-3.5.md), [Intel Arc, experimental](xpu-arc.md) | +| Nemotron 3.5 Parse | [Native HIP OCR, validation and qualification limits](nemotron-parse.md) | | Nemotron 3.5 Super VL | [Native CUDA on one or two GB10s, images](nemotron-3.5-super-vl.md) | | Qwen3.8-27B | [MLX, quantization and CUDA](qwen3.8-27b.md) | | Qwen3.5-2B | [Native Metal preview, tied affine 4-bit](qwen3.5-2b-native.md) | diff --git a/docs/recipes/nemotron-parse.md b/docs/recipes/nemotron-parse.md new file mode 100644 index 0000000000..1020a638ed --- /dev/null +++ b/docs/recipes/nemotron-parse.md @@ -0,0 +1,141 @@ +# Nemotron 3.5 Parse on native HIP + +This port runs `nvidia/Nemotron-3.5-Parse` on one HIP rank. The +validated checkpoint revision is +`4035ffb1d8f8bafcdb867dfaf6f7514f101464f9`. The family is +`nemotron_parse`; it is separate from Nemotron-H and Super VL. + +The original FP32 safetensors stay unchanged. The loader validates all +510 tensors, then owns BF16 device copies of the 509 floating tensors. +Activations and KV use BF16 with FP32 sums. The image tower, horizontal +neck, causal mBART decoder and trained one-step proposal run in native HIP +kernels. Python is used only for independent development references. + +## Build and serve + +Use the repository Zig version. On ROCm installations with a different +layout, give the actual compiler and device-library paths: + +```sh +zig build native-hip -Dhip=true -Dhip-arch=gfx1201 \ + -Dhipcc=/opt/rocm/core/bin/hipcc \ + -Dhip-device-lib=/opt/rocm/core/lib/llvm/amdgcn/bitcode + +zig-out/native-hip/bin/tensorfold-native serve "$PARSE_MODEL" \ + --backend rocm --device 0 --vision --name parse --parallel 2 --context 4096 \ + --mtp-drafts 1 --prompt-cache-gib 0.25 --host 127.0.0.1 --port 8000 +``` + +`PARSE_MODEL` names the already cached snapshot. Select the intended +visible device and verify its PCI identity; management and HIP ordinals +can differ. The engine checks driver free memory at load and admission. +The developer checks require at least six GiB free before a bounded +two-lane run. + +Send one user message with one PNG `image_url` data URL to +`/v1/chat/completions`. For example, with an ordinary OpenAI client: + +```python +import base64 +from pathlib import Path +from openai import OpenAI + +image = base64.b64encode(Path("page.png").read_bytes()).decode() +client = OpenAI(base_url="http://127.0.0.1:8000/v1", api_key="local") +reply = client.chat.completions.create( + model="parse", + messages=[{"role": "user", "content": [ + {"type": "image_url", "image_url": {"url": "data:image/png;base64," + image}} + ]}], + temperature=0, + max_tokens=1024, +) +print(reply.choices[0].message.content) +``` + +Use the served model name reported by `/health`, or set `--name` explicitly. +The built-in media template puts the visual prefix before the example's +six task tokens. Optional text parts follow those tokens. Generated OCR, +including its location/class markers, is returned in `content`. A custom +`--chat-template` still overrides the built-in bridge. + +The image processor shrinks into 2048-by-1664 without enlargement, pads +bottom/right to 16/64, and preserves the visible content size for RADIO +positions. A full page uses 3329 visual context rows. Visual and text rows +share the 16384-row model context; `--context` may lower this limit. One +image is required. Multi-turn chat, video, tensor parallelism, lossy KV, +Sliding Weights and multiple proposal heads are not part of this port. + +Default generation is greedy with repetition penalty 1.1. The shared +sampler widens raw logits to FP32 before applying penalties, and forces +EOS at the configured final output position. `"draft": false` selects +plain target decoding; normal drafted requests use the trained head and +the existing target verifier. Shared rounds batch dense work across all +admitted streams. Prompt snapshots use image content and frame identities. + +## Reproduce the checks + +Create offline BF16 references from the checkpoint's own example workflow: + +```sh +python tools/zig/nemotron_parse_reference.py --model "$PARSE_MODEL" \ + --image page.png --output "$REFERENCE" +python tools/zig/nemotron_parse_reference.py positions --model "$PARSE_MODEL" \ + --output "$POSITION_REFERENCE" + +zig build test-nemotron-parse-host -Dparse-model-dir="$PARSE_MODEL" \ + -Dparse-reference-dir="$REFERENCE" + +TF_PARSE_DIR="$PARSE_MODEL" TF_PARSE_POSITION_DIR="$POSITION_REFERENCE" \ + zig build hip-parse-ops-test -Dhip=true -Dhip-arch=gfx1201 \ + -Dhipcc=/opt/rocm/core/bin/hipcc \ + -Dhip-device-lib=/opt/rocm/core/lib/llvm/amdgcn/bitcode +``` + +The reference environment used CPU Torch 2.11.0, TorchVision 0.26.0, +Transformers 5.14.1 and timm 1.0.22. It loads only local assets. The tool +retains NVIDIA's OpenMDW notice; the license is +[LICENSES/NVIDIA-OpenMDW-1.1.txt](../../LICENSES/NVIDIA-OpenMDW-1.1.txt). +Reference directories must be new; large tensors are not checked into Git. +`nemotron_parse_reference.py processor` prints the model-free +`zig/src/families/nemotron_parse/fixtures/processor.json`. + +`tf-nemotron-parse-test` builds a developer probe. Give it the model, +reference and output directories. It compares intermediate seams and +teacher-forced logits, then checks full/chunked and solo/shared KV/hidden +bytes and suffix rollback. `nemotron_parse_reference.py compare +--reference "$REFERENCE" --native "$PROBE_OUTPUT"` applies the original HF +logits processor to saved native rows. + +The kernel contract registers twelve new operator cases, which +`tf-hip-kernel-contract` runs with no argument, beside the rest of the +contract and its planted controls. `hip-kernel-contract-objects` builds +the code objects for gfx1030, gfx1100, gfx1151 and gfx1201, and +`hip-kernel-contract-limits` checks every launch record against them +without a GPU. + +## Evidence and limits + +Native HTTP token arrays matched three PNG references (two 512-by-256 +pages, one 1664-by-2048 page), plain/drafted and solo/concurrent. The +full page uses 3335 prompt rows including the six text tokens. Prompt KV +and final-hidden bytes agree across full/chunked and solo/shared forwards +and after suffix rollback. No speed target is claimed. + +All twelve new operator cases pass arithmetic gates 1–4, including five +repeat runs, guards, non-finites, tile edges and accumulator stress. +Their seven launch records match the code objects of all four recorded +chips and the live gfx1201 device. The complete contract reports only +the limit problems that main reports too: two stale records on the live +gfx1201 (`tf_sum_partials`, `tf_sum_partials_act`) and 13 across the four +chips' code objects. + +BF16 intermediate tensors differ from HF's CPU GEMM rounding; saved +seam errors are diagnostic, while actual output token arrays are checked +separately. Runtime model checks cover gfx1201 only. CUDA, Metal, other +AMD hardware have not been qualified here. Sampled plain/drafted equality +passes at two small-page seeds, one with presence and frequency penalties; +this does not compare TensorFold's keyed sampler with HF's random draws. +PNG is tested. JPEG uses the repository's pinned libjpeg-turbo decoder +and has no qualification receipt yet. Structured output needs the +existing grammar library and has no real-model receipt yet. diff --git a/tools/zig/nemotron_parse_reference.py b/tools/zig/nemotron_parse_reference.py new file mode 100644 index 0000000000..e6a952da77 --- /dev/null +++ b/tools/zig/nemotron_parse_reference.py @@ -0,0 +1,294 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Offline Nemotron Parse references for native tests: HF example, positions, processor fingerprints, compare.""" + +import argparse +import hashlib +import json +import sys +from pathlib import Path + +# Adapted from run_transformers_inference.py in nvidia/Nemotron-3.5-Parse at the revision below. + +REVISION = '4035ffb1d8f8bafcdb867dfaf6f7514f101464f9' +PROMPT = '' +COMMANDS = ('reference', 'sources', 'positions', 'processor', 'compare') +# The revision's non-weight sources: model, processor and example code, configuration, tokenizer and the model card. +SOURCES = { + 'README.md': 'a97942afdd486e6e413afd66b5f7d955a995ad78b8eaf4a73cd010ac829813ef', + 'chat_template.jinja': '86e90f392a2779c96c96b5848248981a05750e7a11e0380b092a8c35e2c3e32f', + 'config.json': 'bff3460ae530573ad2e276627be86cab1b6859754b76e8b196410f17e7f0b7d5', + 'generation_config.json': '0ec0669b5d3df6c4f2fcbf41334c3d64f44fdee1d1e5f158a33b112edf1646c3', + 'hf_logits_processor.py': '37bb9612db8840eeda67f3bc9a2d8837aeebc3b206315229671d6d56c570704c', + 'hf_nemotron_parse_config.py': '5f4d2530457e41dc100df1fa614485aec8c94715f25cd69ba6e6ebdaf4611899', + 'hf_nemotron_parse_modeling.py': 'ec43d337be2797e2ad137326a68578af67694763eac304fd85b604b6775b93cc', + 'hf_nemotron_parse_processor.py': '983fa8aeb0d872da01296035b30d60c84fd857879a37c939c5c7203bad3c8b81', + 'hf_prefix_geometry.py': '41406f33399fd39bd0601fd2acd320f87fec5f9865f13fbcb4e4577d94b97048', + 'postprocessing.py': 'cef5a7379720d479f7fd6fe47fb0e83097ff47bc053f751e747134c882d5a3fe', + 'preprocessor_config.json': '16d54d73fc265b31f3771caf39770dc85213cae7594823ba7a6ac3f085315ada', + 'processor_config.json': 'b3e2ff2cc94f3d9412f179acf06d8bc5543305ec27966617977b3bd47fee9e5e', + 'run_transformers_inference.py': 'fc3568bb3e9cf48dbcb8ce0cd4e25ee7360cec9fa9da331141497ecee5dcea41', + 'special_tokens_map.json': '7073bb56328d2b66724495a7cf53841fccdcc9efcafb1e1b7cc916baa80e3c8c', + 'tokenizer.json': '34f566e728995769d971fc2707596bfb92c83298d357bc6ec29a7a1d3d12a74a', + 'tokenizer_config.json': '701dad9bcc5f964a6fa020e7d5018c71efc93aad79d471cfca541b9bb4f6d012', +} +VERSIONS = {'torch': '2.11.0', 'torchvision': '0.26.0', 'transformers': '5.14.1', 'timm': '1.0.22'} + + +def pinned(model_dir): + """Return the checkpoint's source hashes, or refuse sources that differ from REVISION.""" + found = {name: hashlib.sha256((model_dir / name).read_bytes()).hexdigest() if (model_dir / name).is_file() else None + for name in SOURCES} + if wrong := sorted(name for name, digest in found.items() if digest != SOURCES[name]): + raise RuntimeError(f"checkpoint sources differ from {REVISION}: {', '.join(wrong)}") + return found + + +def versions(): + """Refuse a reference environment whose packages differ from the pinned versions.""" + import importlib.metadata + + if wrong := [f'{name} {importlib.metadata.version(name)}' for name, want in VERSIONS.items() + if importlib.metadata.version(name).split('+')[0] != want]: + raise RuntimeError(f"reference packages differ from {VERSIONS}: {', '.join(wrong)}") + + +def sources(args): + """Check a checkpoint's sources against the pins and print their hashes; no model or package is loaded.""" + print(json.dumps(dict(reference_revision=REVISION, versions=VERSIONS, source_sha256=pinned(args.model.resolve())), + indent=2)) + + +def reference(args): + model_dir = args.model.resolve() + sources_sha256 = pinned(model_dir) + versions() + import torch + from PIL import Image + from transformers import AutoModel, AutoProcessor, GenerationConfig + + torch.set_num_threads(4) + args.output.mkdir(parents=True, exist_ok=False) + (args.output / "input.image").write_bytes(args.image.read_bytes()) + source = model_dir / 'hf_nemotron_parse_modeling.py' + manifest = dict(reference_revision=REVISION, device=args.device, dtype=args.dtype, + torch_version=torch.__version__, source_sha256=hashlib.sha256(source.read_bytes()).hexdigest(), + image_sha256=hashlib.sha256(args.image.read_bytes()).hexdigest(), sources_sha256=sources_sha256, + versions=VERSIONS, tensors={}) + + def save(name, value): + value = value.detach().cpu().contiguous() + payload = value.view(torch.uint8).numpy().tobytes() + filename = name + '.bin' + (args.output / filename).write_bytes(payload) + manifest['tensors'][name] = dict(file=filename, shape=list(value.shape), dtype=str(value.dtype), + sha256=hashlib.sha256(payload).hexdigest()) + + model = AutoModel.from_pretrained(model_dir, trust_remote_code=True, local_files_only=True, + dtype=getattr(torch, args.dtype)).to(args.device).eval() + processor = AutoProcessor.from_pretrained(model_dir, trust_remote_code=True, local_files_only=True) + image = Image.open(args.image).convert('RGB') + inputs = processor(images=[image], text=args.prompt, return_tensors='pt', + add_special_tokens=False).to(args.device) + for name, value in inputs.items(): + if isinstance(value, torch.Tensor): + save('input.' + name, value) + + handles = [] + capture = True + + def tap(name): + def record(_module, _inputs, output): + if not capture: + return + if isinstance(output, torch.Tensor): + save(name, output) + elif isinstance(output, tuple): + for i, value in enumerate(output): + if isinstance(value, torch.Tensor): + save(f'{name}.{i}', value) + return record + + handles.append(model.encoder.register_forward_hook(tap('encoder.neck'))) + blocks = model.encoder.model_encoder.radio_model.model.blocks + for i, block in enumerate(blocks): + handles.append(block.register_forward_hook(tap(f'encoder.block.{i}'))) + handles.append(blocks[0].register_forward_pre_hook(lambda _module, values: save('encoder.positioned', values[0]) if capture else None)) + for i, layer in enumerate(model.decoder.layers): + handles.append(layer.register_forward_hook(tap(f'decoder.layer.{i}'))) + handles.append(model.decoder.layernorm_embedding.register_forward_hook(tap('decoder.embedding_norm'))) + handles.append(model.decoder.layer_norm.register_forward_hook(tap('decoder.final_norm'))) + config = GenerationConfig.from_pretrained(model_dir, local_files_only=True) + config.max_new_tokens = args.max_new_tokens + config.do_sample = False + config.num_beams = 1 + config.repetition_penalty = 1.1 + generation_step = 0 + + def save_logits(_module, _inputs, output): + nonlocal generation_step, capture + if generation_step == 0: + save('prefill.logits', output) + capture = False + save(f'generation.logits.{generation_step}', output[:, -1:, :]) + generation_step += 1 + + generation_hook = model.lm_head.register_forward_hook(save_logits) + with torch.inference_mode(): + output_ids = model.generate(**inputs, generation_config=config) + generation_hook.remove() + for handle in handles: + handle.remove() + manifest['prompt_ids'] = inputs['input_ids'].tolist()[0] + manifest['output_ids'] = output_ids.tolist()[0] + manifest['new_ids'] = manifest['output_ids'][len(manifest['prompt_ids']):] + manifest['text'] = processor.batch_decode(output_ids, skip_special_tokens=False)[0] + manifest['settings'] = dict(max_new_tokens=args.max_new_tokens, do_sample=False, num_beams=1, + repetition_penalty=1.1) + (args.output / 'manifest.json').write_text(json.dumps(manifest, indent=2) + '\n') + print(manifest['text']) + + +def positions(args): + """Content-normalized RADIO positions from the checkpoint's hf_prefix_geometry.py, one-patch axes included.""" + pinned(args.model.resolve()) + import importlib.util + + import torch + + torch.set_num_threads(4) + spec = importlib.util.spec_from_file_location('parse_geometry', args.model / 'hf_prefix_geometry.py') + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + args.output.mkdir(parents=True, exist_ok=False) + hidden, prefix = 8, 10 + table = (((torch.arange(128 * 128 * hidden) * 17) % 193).float() - 96).div(97).to(torch.bfloat16) + cls = torch.arange(prefix * hidden).float().div(83).to(torch.bfloat16) + + def save(name, value): + (args.output / name).write_bytes(value.contiguous().view(torch.uint8).numpy().tobytes()) + + save('table.bin', table) + save('cls.bin', cls) + cells = [] + for i, (height, width) in enumerate([(1, 1), (8, 73), (17, 65), (113, 97), (512, 256), (2047, 1663)]): + gh, gw = (height + 15) // 16, ((width + 63) // 64) * 4 + patch = (((torch.arange(gh * gw * hidden) * 13) % 101).float() - 50).div(257).to(torch.bfloat16) + with torch.inference_mode(): + positioned = module._normalized_2d_position_pages( + torch, patches=patch.reshape(1, gh * gw, hidden), pos_embed=table.reshape(1, 128 * 128, hidden), + source_rows=128, source_cols=128, image_sizes=[(gh * 16, gw * 16)], patch_size=16, + content_sizes=[(height, width)], content_normalized=True) + save(f'patches-{i}.bin', patch) + save(f'expected-{i}.bin', torch.cat([cls.reshape(prefix, hidden), positioned[0]], dim=0)) + cells.append(dict(height=height, width=width, gh=gh, gw=gw)) + (args.output / 'manifest.json').write_text(json.dumps(dict(hidden=hidden, prefix=prefix, cells=cells), indent=2)) + + +# Formulas of nvidia/Nemotron-3.5-Parse@4035ffb1d8f8bafcdb867dfaf6f7514f101464f9 processor; no native code. +PROCESSOR_CASES = [(1, 1), (17, 65), (31, 129), (113, 97), (2049, 101), (99, 2001), (2201, 1777), (2048, 1664)] + + +def processor_case(np, Image, height, width): + mean = np.array([0.48145466, 0.4578275, 0.40821073], dtype=np.float32) + std = np.array([0.26862954, 0.26130258, 0.27577711], dtype=np.float32) + y, x, ch = np.indices((height, width, 3), dtype=np.uint32) + rgb = ((y * 17 + x * 29 + ch * 83) % 256).astype(np.uint8) + ratio = width / height + h, w = height, width + if h > 2048: + h = 2048 + w = int(h * ratio) + if w > 1664: + w = 1664 + h = int(w / ratio) + image = Image.fromarray(rgb) + if (h, w) != (height, width): + image = image.resize((w, h), Image.Resampling.BILINEAR) + content = np.array(image) + padded = np.pad(content, ((0, -h % 16), (0, -w % 64), (0, 0)), constant_values=255) + # ToTensor divides float32 by 255; Normalize uses two separate float32 operations. + normalized = (padded.astype(np.float32) / np.float32(255.0) - mean) / std + ph, pw = padded.shape[:2] + patches = normalized.reshape(ph // 16, 16, pw // 16, 16, 3).transpose(0, 2, 4, 1, 3) + data = patches.astype('(h) * rows + r) * dim + c]); } + +// Dense BF16 image/text operators; compensated FP32 sums keep their order at every prompt or verify width. + +namespace { + +extern "C" __global__ __launch_bounds__(256) +void tf_vl_embed(const void* table, const int* ids, void* out, int rows, int width, int kind, float scale) { + const long long i = gid(); + if (i >= (long long)rows * width) return; + const int row = i / width, column = i % width; + store(out, kind, i, load(table, kind, (long long)ids[row] * width + column) * scale); +} + +extern "C" __global__ __launch_bounds__(256) +void tf_vl_linear(const void* x, const void* w, const void* bias, void* out, + int rows, int n, int k, int kind, int out_kind) { + __shared__ float xs[16][16]; + __shared__ float ws[64][16]; + const int tid = threadIdx.x, row = tid / 16, col = tid % 16; + const int r = blockIdx.y * 16 + row, base = blockIdx.x * 64; + float acc[4] = {0, 0, 0, 0}; + float correction[4] = {0, 0, 0, 0}; + for (int first = 0; first < k; first += 16) { + xs[row][col] = r < rows && first + col < k ? load(x, kind, (long long)r * k + first + col) : 0; + for (int j = tid; j < 64 * 16; j += 256) { + int nc = base + j / 16, kc = first + j % 16; + ws[j / 16][j % 16] = nc < n && kc < k ? load(w, kind, (long long)nc * k + kc) : 0; + } + __syncthreads(); + for (int j = 0; j < 16 && first + j < k; ++j) { + float xv = xs[row][j]; + for (int tile = 0; tile < 4; ++tile) { + float product = xv * ws[col + tile * 16][j]; + if (isfinite(acc[tile]) && isfinite(product)) { + float adjusted = product - correction[tile]; + float next = acc[tile] + adjusted; + correction[tile] = isfinite(next) ? (next - acc[tile]) - adjusted : 0; + acc[tile] = next; + } else { + acc[tile] += product; + correction[tile] = 0; + } + } + } + __syncthreads(); + } + if (r >= rows) return; + for (int tile = 0; tile < 4; ++tile) { + int nc = base + col + tile * 16; + if (nc < n) { + float value = acc[tile]; + if (bias) value += load(bias, kind, nc); + store(out, out_kind, (long long)r * n + nc, value); + } + } +} + +// Not tf_layer_norm_bf16: that one takes rsqrtf, this one 1 / sqrtf. The sums are the same, yet about one +// BF16 output in 100,000 differs, so dots.mocr and Nemotron Parse each keep the norm their receipts were made on. +extern "C" __global__ __launch_bounds__(256) +void tf_vl_layernorm(const void* x, const void* weight, const void* bias, void* out, + int rows, int width, int kind, float eps) { + __shared__ float sums[256]; + const int tid = threadIdx.x, row = blockIdx.x; + if (row >= rows) return; + const long long base = (long long)row * width; + float sum = 0; + for (int j = tid; j < width; j += 256) sum += load(x, kind, base + j); + sums[tid] = sum; + __syncthreads(); + for (int step = 128; step; step /= 2) { + if (tid < step) sums[tid] += sums[tid + step]; + __syncthreads(); + } + const float mean = sums[0] / width; + __syncthreads(); + sum = 0; + for (int j = tid; j < width; j += 256) { + float centered = load(x, kind, base + j) - mean; + sum += centered * centered; + } + sums[tid] = sum; + __syncthreads(); + for (int step = 128; step; step /= 2) { + if (tid < step) sums[tid] += sums[tid + step]; + __syncthreads(); + } + const float inv = 1.0f / sqrtf(sums[0] / width + eps); + for (int j = tid; j < width; j += 256) { + float value = (load(x, kind, base + j) - mean) * inv; + value = value * load(weight, kind, j) + load(bias, kind, j); + store(out, kind, base + j, value); + } +} + +extern "C" __global__ __launch_bounds__(256) +void tf_vl_pointwise(const void* x, const void* y, void* out, long long count, int kind, int mode, float scale) { + const long long i = gid(); + if (i >= count) return; + const float value = load(x, kind, i); + float result; + if (mode == 0) result = (0.5f * value) * (1.0f + erff(value * 0.7071067811865475f)); + else if (mode == 1) result = value + load(y, kind, i); + else result = value * scale; + store(out, kind, i, result); +} + +extern "C" __global__ __launch_bounds__(256) +void tf_vl_layout(const void* x, void* out, int rows, int heads, int dim, int in_kind, int out_kind, int mode) { + const long long i = gid(), count = (long long)rows * heads * dim; + if (i >= count) return; + long long source; + if (mode == 0) { + // Attention returns [head,row,dim]; projections consume [row,head,dim]. + const int d = i % dim, h = (i / dim) % heads, r = i / (heads * dim); + source = ((long long)h * rows + r) * dim + d; + } else { + // Conv2's input groups four horizontal patches, channel before group member. + const int member = i % 4, channel = (i / 4) % dim; + source = ((i / (4 * dim)) * 4 + member) * dim + channel; + } + store(out, out_kind, i, load(x, in_kind, source)); +} + +} // namespace + +#include "vl_position.hpp" diff --git a/zig/kernels/hip/ops/vl_position.hpp b/zig/kernels/hip/ops/vl_position.hpp new file mode 100644 index 0000000000..66a28a21b7 --- /dev/null +++ b/zig/kernels/hip/ops/vl_position.hpp @@ -0,0 +1,77 @@ +#pragma once + +#include "common.hpp" + +// The learned field's antialiased resize in FP32, then content-normalized border sampling into BF16 patch rows. + +namespace { + +extern "C" __global__ __launch_bounds__(256) +void tf_vl_resize(const void* src, float* out, int ih, int iw, int oh, int ow, int width, int kind, int axis) { + const long long i = gid(); + if (i >= (long long)oh * ow * width) return; + const int ch = i % width, x = (i / width) % ow, y = i / (width * ow); + if (axis == 2) { + const int y0 = (y * ih) / oh, y1 = ((y + 1) * ih + oh - 1) / oh; + const int x0 = (x * iw) / ow, x1 = ((x + 1) * iw + ow - 1) / ow; + float sum = 0; + for (int yy = y0; yy < y1; ++yy) for (int xx = x0; xx < x1; ++xx) + sum += load(src, kind, ((long long)yy * iw + xx) * width + ch); + out[i] = (sum / (y1 - y0)) / (x1 - x0); + return; + } + const int in = axis == 0 ? iw : ih, size = axis == 0 ? ow : oh, index = axis == 0 ? x : y; + if (in == size) { + out[i] = load(src, kind, ((long long)y * iw + x) * width + ch); + return; + } + const float scale = (float)in / size, support = scale > 1.0f ? scale : 1.0f; + const float center = (index + 0.5f) * scale, inv = 1.0f / support; + const int first = max((int)(center - support + 0.5f), 0); + const int end = min((int)(center + support + 0.5f), in); + float total = 0; + for (int j = first; j < end; ++j) { + float distance = fabsf((j - center + 0.5f) * inv); + total += distance < 1.0f ? 1.0f - distance : 0.0f; + } + float result = 0; + for (int j = first; j < end; ++j) { + float distance = fabsf((j - center + 0.5f) * inv); + float weight = (distance < 1.0f ? 1.0f - distance : 0.0f) / total; + long long at = axis == 0 ? ((long long)y * iw + j) * width + ch : ((long long)j * iw + x) * width + ch; + result += load(src, kind, at) * weight; + } + out[i] = result; +} + +__device__ float vl_sample_coordinate(int index, int pixels, int size) { + const float start = (float)(index * 16); + const float extent = fminf(fmaxf((float)pixels - start, 0.0f), 16.0f); + const float normalized = fminf(start + 0.5f * extent, (float)pixels) / pixels; + const float grid = 2.0f * normalized - 1.0f; + return fminf(fmaxf(((grid + 1.0f) * size - 1.0f) * 0.5f, 0.0f), (float)(size - 1)); +} + +extern "C" __global__ __launch_bounds__(256) +void tf_vl_position(const void* patches, const void* cls, const float* field, void* out, + int gh, int gw, int height, int width, int hidden, int prefix, int kind) { + const long long i = gid(); + const int count = gh * gw + prefix; + if (i >= (long long)count * hidden) return; + const int row = i / hidden, ch = i % hidden; + if (row < prefix) { + store(out, kind, i, load(cls, kind, i)); + return; + } + const int r = row - prefix, oh = (height + 15) / 16, ow = (width + 15) / 16; + const float y = vl_sample_coordinate(r / gw, height, oh), x = vl_sample_coordinate(r % gw, width, ow); + const int y0 = (int)floorf(y), x0 = (int)floorf(x), y1 = min(y0 + 1, oh - 1), x1 = min(x0 + 1, ow - 1); + const float sy = y - y0, sx = x - x0; + float value = field[((long long)y0 * ow + x0) * hidden + ch] * ((1 - sx) * (1 - sy)); + value += field[((long long)y0 * ow + x1) * hidden + ch] * (sx * (1 - sy)); + value += field[((long long)y1 * ow + x0) * hidden + ch] * ((1 - sx) * sy); + value += field[((long long)y1 * ow + x1) * hidden + ch] * (sx * sy); + store(out, kind, i, load(patches, kind, (long long)r * hidden + ch) + rounded(value, kind)); +} + +} // namespace diff --git a/zig/src/core/contract/gates.zig b/zig/src/core/contract/gates.zig index 6c0ca7d413..a815417a16 100644 --- a/zig/src/core/contract/gates.zig +++ b/zig/src/core/contract/gates.zig @@ -320,7 +320,7 @@ pub fn sums(a: std.mem.Allocator, k: Kernel, seed: u64) !Verdict { const q = misses(got, want, allow); if (!(q <= 1)) return say(.fail, .sums, s.name, "{d:.1}x over its exact-sum bound at k {d}: not fp32 sums (fp16 would miss {d:.0}x, bf16 {d:.0}x)", .{ q, shape.k, margin[0], margin[1] }); if (need[1] == 0 or need[0] == 0) return say(.pass, .sums, s.name, "within its bound at k {d}; fp16 sums would miss {d:.0}x, bf16 {d:.1}x ({s})", .{ shape.k, margin[0], margin[1], s.sums.why }); - return say(.pass, .sums, s.name, "exact sums at k {d}, where fp16 sums would miss by {d:.0}x and bf16 by {d:.0}x", .{ shape.k, margin[0], margin[1] }); + return say(.pass, .sums, s.name, "{s} at k {d}, where fp16 sums would miss by {d:.0}x and bf16 by {d:.0}x", .{ if (s.sums.exact) "exact sums" else "FP32 sums within their derived bound", shape.k, margin[0], margin[1] }); } /// A kernel's or control's gates 1-4; a control must fail its own gate and passes when it does. diff --git a/zig/src/core/contract/hip_records_families.zig b/zig/src/core/contract/hip_records_families.zig index 2d9d96315d..49d86fa1b5 100644 --- a/zig/src/core/contract/hip_records_families.zig +++ b/zig/src/core/contract/hip_records_families.zig @@ -16,6 +16,6 @@ pub const all = join: { test "the HIP records stay main's; the family lists only regroup GLM-OCR's and the joint schema's" { const every = @import("hip_records.zig").records; - try std.testing.expectEqual(@as(usize, 199), every.len); // 197 frozen records and two drafter-tap records - try std.testing.expectEqual(@as(usize, 177), every.len - all.len); + try std.testing.expectEqual(@as(usize, 206), every.len); // 199 main records and seven Parse vision records + try std.testing.expectEqual(@as(usize, 184), every.len - all.len); } diff --git a/zig/src/core/contract/hip_records_vision.zig b/zig/src/core/contract/hip_records_vision.zig index 89fd81b464..806f7c3edb 100644 --- a/zig/src/core/contract/hip_records_vision.zig +++ b/zig/src/core/contract/hip_records_vision.zig @@ -47,4 +47,67 @@ pub const records = [_]Record{ .{ .vgprs = 229, .sgprs = 48, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, }, }, + .{ + .launch = .{ .name = "vision tf_vl_embed", .symbol = "tf_vl_embed", .image = "ops", .threads = 256 }, + .usage = .{ + .{ .vgprs = 13, .sgprs = 23, .lds = 0, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 13, .sgprs = 23, .lds = 0, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 13, .sgprs = 23, .lds = 0, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 13, .sgprs = 25, .lds = 0, .private = 0, .bound = 256, .wave = 32 }, + }, + }, + .{ + .launch = .{ .name = "vision tf_vl_linear", .symbol = "tf_vl_linear", .image = "ops", .threads = 256 }, + .usage = .{ + .{ .vgprs = 26, .sgprs = 26, .lds = 5120, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 26, .sgprs = 26, .lds = 5120, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 26, .sgprs = 26, .lds = 5120, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 26, .sgprs = 26, .lds = 5120, .private = 0, .bound = 256, .wave = 32 }, + }, + }, + .{ + .launch = .{ .name = "vision tf_vl_layernorm", .symbol = "tf_vl_layernorm", .image = "ops", .threads = 256 }, + .usage = .{ + .{ .vgprs = 17, .sgprs = 29, .lds = 1024, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 16, .sgprs = 30, .lds = 1024, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 16, .sgprs = 30, .lds = 1024, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 16, .sgprs = 30, .lds = 1024, .private = 0, .bound = 256, .wave = 32 }, + }, + }, + .{ + .launch = .{ .name = "vision tf_vl_pointwise", .symbol = "tf_vl_pointwise", .image = "ops", .threads = 256 }, + .usage = .{ + .{ .vgprs = 10, .sgprs = 18, .lds = 0, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 10, .sgprs = 18, .lds = 0, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 10, .sgprs = 14, .lds = 0, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 10, .sgprs = 14, .lds = 0, .private = 0, .bound = 256, .wave = 32 }, + }, + }, + .{ + .launch = .{ .name = "vision tf_vl_layout", .symbol = "tf_vl_layout", .image = "ops", .threads = 256 }, + .usage = .{ + .{ .vgprs = 15, .sgprs = 27, .lds = 0, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 15, .sgprs = 27, .lds = 0, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 15, .sgprs = 27, .lds = 0, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 15, .sgprs = 32, .lds = 0, .private = 0, .bound = 256, .wave = 32 }, + }, + }, + .{ + .launch = .{ .name = "vision tf_vl_resize", .symbol = "tf_vl_resize", .image = "ops", .threads = 256 }, + .usage = .{ + .{ .vgprs = 27, .sgprs = 31, .lds = 0, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 27, .sgprs = 31, .lds = 0, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 27, .sgprs = 31, .lds = 0, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 27, .sgprs = 38, .lds = 0, .private = 0, .bound = 256, .wave = 32 }, + }, + }, + .{ + .launch = .{ .name = "vision tf_vl_position", .symbol = "tf_vl_position", .image = "ops", .threads = 256 }, + .usage = .{ + .{ .vgprs = 24, .sgprs = 26, .lds = 0, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 25, .sgprs = 26, .lds = 0, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 24, .sgprs = 26, .lds = 0, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 24, .sgprs = 30, .lds = 0, .private = 0, .bound = 256, .wave = 32 }, + }, + }, }; diff --git a/zig/src/core/engine_vision.zig b/zig/src/core/engine_vision.zig index 987b3c593b..a4d4e819da 100644 --- a/zig/src/core/engine_vision.zig +++ b/zig/src/core/engine_vision.zig @@ -8,6 +8,8 @@ pub const PreparedImage = struct { gh: u32, gw: u32, tokens: u32, + /// Visible pixel height and width before padding, for towers with content-normalized positions. + content_size: ?[2]u32 = null, hash: u64, /// The image as tower inputs in order, each image tokens between `before` and `after`; empty: `pixels` alone. pieces: []const Piece = &.{}, diff --git a/zig/src/core/lanes/stream.zig b/zig/src/core/lanes/stream.zig index d950d5640d..baff3bd310 100644 --- a/zig/src/core/lanes/stream.zig +++ b/zig/src/core/lanes/stream.zig @@ -108,6 +108,8 @@ pub const Image = struct { gw: u32, at: u32, tokens: u32, + /// Visible pixel height and width before padding, for towers with content-normalized positions. + content_size: ?[2]u32 = null, hash: u64 = 0, // core/picture/identity.zig's: the bytes and the grid, as the prompt cache's key names the image }; diff --git a/zig/src/families/nemotron_parse/backend.zig b/zig/src/families/nemotron_parse/backend.zig new file mode 100644 index 0000000000..cce97f24a9 --- /dev/null +++ b/zig/src/families/nemotron_parse/backend.zig @@ -0,0 +1,243 @@ +//! Parse's target and one-step head behind the existing shared lane engine; no family-owned scheduler. +const std = @import("std"); +const hip = @import("hip"); +const lanes = @import("lanes"); +const serve = @import("hip_serve"); +const glue = serve.backend; +const draw = serve.draw; +const be = lanes.backend; +const c = @import("config.zig"); +const decoder = @import("decoder.zig"); +const engine = @import("engine.zig"); +const Compute = @import("compute.zig").Compute; +const Tensor = hip.ops.Tensor; +const prefill = @import("prefill.zig"); + +pub const max_window = 32; +pub const Lane = struct { + state: decoder.State, + prefix: hip.DeviceBuffer, + image_rows: usize, + fill: usize = 0, + first: u64 = 0, + first_row: ?[]u16 = null, + held: [1]u32 = undefined, + held_n: usize = 0, + pending: ?struct { start: usize, rows: usize } = null, +}; + +pub const Backend = struct { + a: std.mem.Allocator, + e: *engine.Engine, + streams: glue.Streams(Lane), + sampler: draw.Drawer, + drafts: bool, + cache: @import("engine_api").prompt_cache.Store, + + pub fn init(a: std.mem.Allocator, e: *engine.Engine, drafts: bool) !*Backend { + const h = try a.create(Backend); + errdefer a.destroy(h); + h.* = .{ .a = a, .e = e, .streams = .init(a), .sampler = try draw.Drawer.init(a, &e.r, .bf16, c.vocab, e.streams * max_window), .drafts = drafts, .cache = undefined }; + h.cache = @import("engine_api").prompt_cache.Store.init(a, @import("cache.zig").snapshots(h), .{ .chunk = 1, .min_prompt = 0, .repeat_min = 1, .repeat_keyed = true, .warm = .on }, 0); + return h; + } + + pub fn deinit(h: *Backend) void { + h.e.stream.synchronize() catch {}; + h.cache.deinit(); + var it = h.streams.map.valueIterator(); + while (it.next()) |l| h.drop(l.*); + h.streams.deinit(); + h.sampler.deinit(); + const a = h.a; + a.destroy(h); + } + + pub fn drop(h: *Backend, l: *Lane) void { + h.e.stream.synchronize() catch {}; + if (l.first_row) |row| h.a.free(row); + l.prefix.free(); + l.state.deinit(); + h.a.destroy(l); + } + + pub fn facts(h: *const Backend) lanes.Model { + return .{ .exact_width = max_window, .mtp = h.drafts, .speculate = h.drafts, .speculate_early = false, .drafts = 1, .hidden_rows = true, .batch_rows = @intCast(h.e.streams * max_window), .max_streams = @intCast(h.e.streams), .draft_streams = h.drafts }; + } + + pub fn backend(h: *Backend) be.Backend { + return .{ .ptr = h, .vtable = &vtable, .penalties = true, .chunk_invariant = true }; + } + + const vtable: be.Backend.VTable = .{ .reports_fill = true, .prefill = fillAll, .prefill_step = fillStep, .first = first, .first_masked = firstMasked, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .tree = tree, .held = held, .release = release }; + + pub fn of(ptr: *anyopaque) *Backend { + const h: *Backend = @ptrCast(@alignCast(ptr)); + glue.bind(&h.e.ctx); + return h; + } + + fn fillAll(ptr: *anyopaque, s: *lanes.Stream) anyerror!void { + while (!try fillStep(ptr, s, 0)) {} + } + + fn fillStep(ptr: *anyopaque, s: *lanes.Stream, max_rows: u32) anyerror!bool { + return prefill.step(of(ptr), s, max_rows); + } + + fn first(ptr: *anyopaque, s: *lanes.Stream, position: u64) anyerror!u64 { + if (position != s.prompt_len) return error.PositionMismatch; + return (try of(ptr).streams.get(s)).first; + } + + fn firstMasked(ptr: *anyopaque, s: *lanes.Stream, position: u64, bits: []const u32) anyerror!u64 { + const h = of(ptr); + if (position != s.prompt_len) return error.PositionMismatch; + const l = try h.streams.get(s); + const row = try h.a.dupe(u16, l.first_row orelse return error.NoFirstRow); + defer h.a.free(row); + var r = try h.request(h.a, s, l.image_rows, s.prompt_len, &.{}, position); + defer h.a.free(r.history); + r.mask = bits; + return h.streams.drawn.take(try h.sampler.drawHost(row, r)); + } + + fn queue(_: *anyopaque, _: *lanes.Stream, _: be.Feed, _: u64) anyerror!u64 { + return error.NotPipelined; + } + + fn read(ptr: *anyopaque, handle: u64) anyerror!u32 { + return of(ptr).streams.drawn.read(handle); + } + + /// Sampling history includes the current window's input prefix, and excludes visual placeholder IDs. + pub fn request(_: *Backend, a: std.mem.Allocator, s: *lanes.Stream, image_rows: usize, at: usize, extra: []const u32, position: u64) !draw.Request { + if (at < image_rows or at > s.context.items.len or at < s.prompt_len) return error.PositionMismatch; + const base = s.context.items[image_rows..at]; + const history = try a.alloc(u32, base.len + extra.len); + @memcpy(history[0..base.len], base); + @memcpy(history[base.len..], extra); + return .{ .sampling = s.sampling, .position = position, .repetition = 1.1, .forced = if (s.max_new > 0 and position == s.prompt_len + s.max_new - 1) 2 else null, .history = history, .reply = history[s.prompt_len - image_rows ..], .penalties = s.penalties }; + } + + fn settle(l: *Lane, rows: ?usize) !void { + if (l.pending) |p| { + const n = rows orelse p.rows; + if (n > p.rows) return error.BadKeep; + try l.state.keep(p.start + n); + l.pending = null; + } + } + + fn verify(ptr: *anyopaque, ws: []const be.Window, out: []be.Verified) anyerror!void { + const h = of(ptr); + if (ws.len == 0 or ws.len > h.e.streams or out.len != ws.len) return error.WindowTooWide; + var arena: std.heap.ArenaAllocator = .init(h.a); + defer arena.deinit(); + const a = arena.allocator(); + const spans = try a.alloc(decoder.Span, ws.len); + const starts = try a.alloc(usize, ws.len); + const tokens = try a.alloc(u32, ws.len * max_window); + const reqs = try a.alloc(draw.Request, tokens.len); + var total: usize = 0; + for (ws, 0..) |w, i| { + const l = try h.streams.get(w.stream); + try settle(l, null); + const n = try glue.windowTokens(w, l.held[0..l.held_n], l.state.len, tokens[total..][0..max_window]); + if (w.positions.len != n or out[i].sampled.len != n or out[i].drafts.len != n - 1) return error.BadWindow; + spans[i] = .{ .state = &l.state, .rows = n }; + starts[i] = l.state.len; + for (0..n) |j| { + reqs[total + j] = try h.request(a, w.stream, l.image_rows, l.state.len, tokens[total..][0 .. j + 1], w.positions[j]); + const words = if (w.masks.len == 0) 0 else w.masks.len / n; + if (w.masks.len != words * n) return error.BadMask; + reqs[total + j].mask = w.masks[j * words ..][0..words]; + } + total += n; + } + for (tokens[0..total]) |id| if (id >= c.vocab) return error.BadInputToken; + _ = try decoder.rowsOf(spans); + h.e.scratch.reset(); + const compute = h.e.compute(); + var device = try hip.DeviceBuffer.fromHost(&h.e.r, std.mem.sliceAsBytes(tokens[0..total])); + defer device.free(); + const embedded = try decoder.embed(compute, device.base(), total); + const hidden = try decoder.forwardMany(compute, spans, embedded); + const logits = try decoder.logits(compute, hidden, total); + const choices = try a.alloc(u32, total); + try h.sampler.draw(compute.ops, h.e.stream, logits, reqs[0..total], choices, false); + var at: usize = 0; + for (ws, out, starts) |w, *o, start| { + const l = try h.streams.get(w.stream); + const n = w.rows(); + @memcpy(o.sampled, choices[at..][0..n]); + @memcpy(o.drafts, tokens[at + 1 ..][0 .. n - 1]); + l.pending = .{ .start = start, .rows = n }; + l.held_n = 0; + at += n; + } + } + + fn keep(ptr: *anyopaque, ws: []const be.Window, paths: []const []const u32) anyerror!void { + const h = of(ptr); + for (ws, paths) |w, path| { + for (path, 0..) |row, i| if (row != i) return error.TreesNotBuilt; + try settle(try h.streams.get(w.stream), path.len); + } + } + + /// The trained auxiliary head predicts one token; all streams' proposals use one packed head forward. + fn draft(ptr: *anyopaque, rs: []const be.DraftRequest) anyerror!void { + const h = of(ptr); + if (rs.len == 0) return; + if (!h.drafts or rs.len > h.e.streams) return error.NoDraftHead; + var arena: std.heap.ArenaAllocator = .init(h.a); + defer arena.deinit(); + const a = arena.allocator(); + const active = try a.alloc(*Lane, rs.len); + const tokens = try a.alloc(u32, rs.len); + const reqs = try a.alloc(draw.Request, rs.len); + h.e.scratch.reset(); + const compute = h.e.compute(); + const hidden = try compute.alloc(rs.len * c.hidden, .bf16); + var count: usize = 0; + for (rs) |r| { + const l = try h.streams.get(r.stream); + try settle(l, if (r.rows) |rows| rows.len else null); + l.held_n = 0; + if (r.depth == 0) continue; + if (r.position != l.state.len + 1 or r.depth > 1 or r.lanes != null) return error.BadDraftRequest; + const token = if (r.follow.len > 0) r.follow[r.follow.len - 1] else try glue.firstToken(&h.streams.drawn, r); + tokens[count] = token; + active[count] = l; + reqs[count] = try h.request(a, r.stream, l.image_rows, l.state.len, tokens[count..][0..1], r.position); + try compute.ops.copyCols(try l.state.hiddenAt(l.state.len - 1), c.hidden, 0, hidden.ptr + count * c.hidden * 2, 1, c.hidden); + count += 1; + } + if (count == 0) return; + var device = try hip.DeviceBuffer.fromHost(&h.e.r, std.mem.sliceAsBytes(tokens[0..count])); + defer device.free(); + const logits = try @import("mtp.zig").logits(compute, hidden, device.base(), count); + const choices = try a.alloc(u32, count); + try h.sampler.draw(compute.ops, h.e.stream, logits, reqs[0..count], choices, false); + for (active[0..count], choices) |l, token| { + l.held[0] = token; + l.held_n = 1; + } + } + + fn tree(ptr: *anyopaque, s: *lanes.Stream, _: std.mem.Allocator) anyerror!?lanes.stream.Held { + return .{ .count = @intCast((try of(ptr).streams.get(s)).held_n) }; + } + + fn held(ptr: *anyopaque, s: *lanes.Stream, out: []u32) anyerror!void { + const l = try of(ptr).streams.get(s); + if (out.len > l.held_n) return error.NoHeldDrafts; + @memcpy(out, l.held[0..out.len]); + } + + fn release(ptr: *anyopaque, s: *lanes.Stream) void { + const h = of(ptr); + h.drop(h.streams.forget(s) orelse return); + } +}; diff --git a/zig/src/families/nemotron_parse/cache.zig b/zig/src/families/nemotron_parse/cache.zig new file mode 100644 index 0000000000..f98e57bb6b --- /dev/null +++ b/zig/src/families/nemotron_parse/cache.zig @@ -0,0 +1,70 @@ +//! Owned exact KV/hidden snapshots for the shared prompt cache, keyed by the server's image-aware token identity. +const std = @import("std"); +const hip = @import("hip"); +const api = @import("engine_api"); +const lanes = @import("lanes"); +const c = @import("config.zig"); +const Backend = @import("backend.zig").Backend; +const State = @import("decoder.zig").State; +const engine = @import("engine.zig"); +const pc = api.prompt_cache; + +pub fn snapshots(h: *Backend) pc.Snapshots { + return .{ .ptr = h, .vtable = &.{ .bytes = bytes, .save = save, .restore = restore, .drop = drop } }; +} + +fn of(ptr: *anyopaque) *Backend { + return Backend.of(ptr); +} + +fn bytes(_: *anyopaque, at: u32) u64 { + return engine.stateBytes(at); +} + +fn lane(h: *Backend, owner: ?*anyopaque) !*@import("backend.zig").Lane { + const s: *lanes.Stream = @ptrCast(@alignCast(owner orelse return error.NoCacheOwner)); + return h.streams.get(s); +} + +fn copy(h: *Backend, x: *const State, y: *State, n: usize) !void { + if (n == 0 or n > x.len or n > y.total) return error.BadSnapshot; + try hip.raw.copy(y.hidden, 0, x.hidden.base(), n * c.hidden * 2, h.e.stream.handle); + for (x.layers, y.layers) |src, dst| for (0..c.kv_heads) |head| { + const xo = head * x.total * c.head_dim * 2; + const yo = head * y.total * c.head_dim * 2; + try hip.raw.copy(dst.k, yo, src.k.base() + xo, n * c.head_dim * 2, h.e.stream.handle); + try hip.raw.copy(dst.v, yo, src.v.base() + xo, n * c.head_dim * 2, h.e.stream.handle); + }; + try h.e.stream.synchronize(); + y.len = n; +} + +fn save(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!pc.Saved { + const h = of(ptr); + const l = try lane(h, owner); + if (at == 0 or at > l.state.len) return error.BadSnapshot; + const pool = try hip.admission.budget(h.a, h.e.io, h.e.ctx); + if (pool.room(hip.usage(false).device) < engine.stateBytes(at)) return error.LaneMemory; + const saved = try h.a.create(State); + errdefer h.a.destroy(saved); + saved.* = try State.init(&h.e.r, at); + errdefer saved.deinit(); + try copy(h, &l.state, saved, at); + return saved; +} + +pub fn restore(ptr: *anyopaque, owner: ?*anyopaque, saved: pc.Saved) anyerror!void { + const h = of(ptr); + const x: *State = @ptrCast(@alignCast(saved)); + const l = try lane(h, owner); + try copy(h, x, &l.state, x.len); + l.fill = x.len; +} + +fn drop(ptr: *anyopaque, saved: pc.Saved) void { + const h = of(ptr); + const x: *State = @ptrCast(@alignCast(saved)); + h.e.stream.synchronize() catch {}; + x.deinit(); + h.a.destroy(x); +} diff --git a/zig/src/families/nemotron_parse/checkpoint.zig b/zig/src/families/nemotron_parse/checkpoint.zig new file mode 100644 index 0000000000..7a51666448 --- /dev/null +++ b/zig/src/families/nemotron_parse/checkpoint.zig @@ -0,0 +1,42 @@ +//! A complete Parse checkpoint mapped read-only; opening this never initializes a GPU. +const std = @import("std"); +const st = @import("core").safetensors; +const config = @import("config.zig"); +const schema = @import("schema.zig"); + +pub const Checkpoint = struct { + file: st.File, + io: std.Io, + + pub fn open(a: std.mem.Allocator, io: std.Io, dir: []const u8) !Checkpoint { + try validateConfig(a, io, dir, "config.json", config.validate); + try validateConfig(a, io, dir, "preprocessor_config.json", config.validateProcessor); + const path = try std.fs.path.join(a, &.{ dir, "model.safetensors" }); + defer a.free(path); + var file = try st.File.open(a, io, path); + errdefer file.close(io); + try schema.validate(&file.names); + // Summary order is part of the neck, not merely an index tensor's shape. + const indices = file.get(schema.radio ++ "summary_idxs").?.bytes; + if (std.mem.readInt(i64, indices[0..8], .little) != 0 or + std.mem.readInt(i64, indices[8..16], .little) != 1) return error.BadParseSummary; + return .{ .file = file, .io = io }; + } + + pub fn close(s: *Checkpoint) void { + s.file.close(s.io); + s.* = undefined; + } + + pub fn get(s: *const Checkpoint, name: []const u8) !st.Tensor { + return s.file.get(name) orelse error.MissingParseTensor; + } +}; + +fn validateConfig(a: std.mem.Allocator, io: std.Io, dir: []const u8, name: []const u8, check: anytype) !void { + const path = try std.fs.path.join(a, &.{ dir, name }); + defer a.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(4 << 20)); + defer a.free(bytes); + try check(a, bytes); +} diff --git a/zig/src/families/nemotron_parse/compute.zig b/zig/src/families/nemotron_parse/compute.zig new file mode 100644 index 0000000000..39ea695174 --- /dev/null +++ b/zig/src/families/nemotron_parse/compute.zig @@ -0,0 +1,82 @@ +//! The family graph over shared HIP operators; scratch is reused in stream order, never read back for inference. +const std = @import("std"); +const hip = @import("hip"); +const Weights = @import("weights.zig").Weights; +const Tensor = hip.ops.Tensor; + +pub const Tap = struct { + ctx: *anyopaque, + record: *const fn (*anyopaque, []const u8, Tensor, usize, usize) anyerror!void, +}; + +pub const Compute = struct { + ops: hip.ops.Ops, + kernels: hip.vision.Kernels, + stream: hip.Stream, + weights: *const Weights, + tap: ?Tap = null, + cancel: ?struct { ptr: *anyopaque, check: *const fn (*anyopaque) bool } = null, + + pub fn poll(c: Compute) !void { + if (c.cancel) |x| { + try c.stream.synchronize(); + if (x.check(x.ptr)) return error.Cancelled; + } + } + + pub fn alloc(c: Compute, count: usize, kind: hip.ops.Kind) !Tensor { + return .{ .ptr = try c.ops.arena.take(count * kind.size()), .kind = kind }; + } + + pub fn named(c: Compute, stem: []const u8, suffix: []const u8) !Tensor { + var buf: [256]u8 = undefined; + return c.weights.get(try std.fmt.bufPrint(&buf, "{s}.{s}", .{ stem, suffix })); + } + + pub fn linear(c: Compute, x: Tensor, stem: []const u8, rows: usize, n: usize, k: usize, bias: bool) !Tensor { + const out = try c.alloc(rows * n, .bf16); + try c.kernels.linear(c.stream, x, (try c.named(stem, "weight")).ptr, if (bias) (try c.named(stem, "bias")).ptr else null, out, rows, n, k); + return out; + } + + pub fn norm(c: Compute, x: Tensor, stem: []const u8, rows: usize, width: usize, eps: f32) !Tensor { + const out = try c.alloc(rows * width, .bf16); + try c.kernels.norm(c.stream, x, (try c.named(stem, "weight")).ptr, (try c.named(stem, "bias")).ptr, out, rows, width, eps); + return out; + } + + pub fn residual(c: Compute, x: Tensor, y: Tensor, count: usize) !void { + try c.kernels.point(c.stream, x, y, x, count, .add, 1); + } + + pub fn gelu(c: Compute, x: Tensor, count: usize) !void { + try c.kernels.point(c.stream, x, null, x, count, .gelu, 1); + } + + pub fn record(c: Compute, name: []const u8, x: Tensor, rows: usize, width: usize) !void { + if (c.tap) |tap| try tap.record(tap.ctx, name, x, rows, width); + } + + /// Copy interleaved projection rows into [head,row,dim], without applying text rotary positions. + pub fn heads(c: Compute, x: Tensor, rows: usize, count: usize, dim: usize, wide: bool) !Tensor { + const out = try c.alloc(rows * count * dim, if (wide) .f32 else .bf16); + try c.ops.ropePrefill(x, out, rows * dim, dim, rows, count, dim, 0, 0, 1); + return out; + } + + /// Global attention is causal attention whose first query can already see the complete span. + pub fn global(c: Compute, qkv: Tensor, rows: usize, count: usize, dim: usize) !Tensor { + const width = count * dim; + var parts: [3]Tensor = undefined; + for (&parts, 0..) |*p, i| { + const slice = try c.alloc(rows * width, .bf16); + try c.ops.copyCols(qkv, 3 * width, i * width, slice.ptr, rows, width); + p.* = try c.heads(slice, rows, count, dim, i == 0); + } + const att = try c.alloc(rows * width, .f32); + try c.ops.causalPrefill(parts[0].ptr, .{ .k = parts[1].ptr, .v = parts[2].ptr, .kind = .bf16, .kv_heads = count, .total = rows, .d = dim }, att.ptr, rows, rows, count, 1 / @sqrt(@as(f32, @floatFromInt(dim))), rows - 1); + const out = try c.alloc(rows * width, .bf16); + try c.kernels.layout(c.stream, att, out, rows, count, dim, .heads_to_rows); + return out; + } +}; diff --git a/zig/src/families/nemotron_parse/config.zig b/zig/src/families/nemotron_parse/config.zig new file mode 100644 index 0000000000..1e1a9b5043 --- /dev/null +++ b/zig/src/families/nemotron_parse/config.zig @@ -0,0 +1,131 @@ +//! The pinned Parse architecture; this is not the Nemotron-H or Super VL family. +const std = @import("std"); +const Object = std.json.ObjectMap; +const Value = std.json.Value; + +pub const model_type = "nemotron_parse"; +pub const hidden = 1024; +pub const layers = 10; +pub const ffn = 4096; +pub const heads = 16; +pub const kv_heads = 8; +pub const head_dim = 64; +pub const vocab = 72512; +pub const context = 16384; +pub const vision_hidden = 1152; +pub const vision_layers = 27; +pub const vision_ffn = 4304; +pub const vision_heads = 16; +pub const vision_prefix = 10; +pub const patch = 16; +pub const max_height = 2048; +pub const max_width = 1664; +pub const norm_eps: f32 = 1e-5; +pub const vision_eps: f32 = 1e-6; +pub const mean = [3]f32{ 0.48145466, 0.4578275, 0.40821073 }; +pub const deviation = [3]f32{ 0.26862954, 0.26130258, 0.27577711 }; + +pub const Error = error{ BadParseConfig, UnsupportedParseConfig } || std.mem.Allocator.Error; + +/// Check fields that affect arithmetic or geometry; exported aliases must agree. +pub fn validate(a: std.mem.Allocator, bytes: []const u8) !void { + const doc = try std.json.parseFromSlice(Value, a, bytes, .{}); + defer doc.deinit(); + const root = try object(doc.value); + try string(root, "model_type", model_type); + try boolean(root, "legacy_causal_vision_prefix", true); + try boolean(root, "tie_word_embeddings", true); + try string(root, "navit_position_embedding_mode", "content_normalized_2d_interpolate"); + try string(root, "navit_coordinate_frame", "content_normalized"); + try boolean(root, "navit_packing", true); + try number(root, "max_sequence_length", context); + try number(root, "num_extra_heads", 1); + try number(root, "vision_spatial_dim", vision_hidden); + try number(root, "vision_summary_dim", vision_hidden * 2); + try number(root, "decoder_start_token_id", 2); + try number(root, "eos_token_id", 2); + try number(root, "pad_token_id", 1); + inline for (.{ .{ "hidden_size", hidden }, .{ "intermediate_size", ffn }, .{ "vocab_size", vocab }, .{ "num_hidden_layers", layers }, .{ "decoder_layers", layers }, .{ "num_attention_heads", heads }, .{ "num_key_value_heads", kv_heads }, .{ "self_attention_kv_heads", kv_heads }, .{ "cross_attention_kv_heads", 0 }, .{ "max_position_embeddings", context }, .{ "legacy_causal_prefix_context_limit", context } }) |pair| try alias(root, pair[0], pair[1]); + for ([_][]const u8{ "quantization", "quantization_config" }) |key| { + if (root.get(key)) |v| if (v != .null) return error.UnsupportedParseConfig; + } + const decoder = try object(try field(root, "decoder")); + try string(decoder, "model_type", "nemotron_parse_text"); + try string(decoder, "activation_function", "gelu"); + try boolean(decoder, "scale_embedding", true); + try boolean(decoder, "tie_word_embeddings", true); + try boolean(decoder, "causal_vision_prefix", true); + try boolean(decoder, "add_cross_attention", false); + inline for (.{ .{ "d_model", hidden }, .{ "decoder_layers", layers }, .{ "decoder_ffn_dim", ffn }, .{ "vocab_size", vocab }, .{ "decoder_attention_heads", heads }, .{ "self_attention_kv_heads", kv_heads }, .{ "cross_attention_kv_heads", 0 } }) |pair| + try number(decoder, pair[0], pair[1]); + inline for (.{ .{ "hidden_size", hidden }, .{ "num_hidden_layers", layers } }) |pair| + try alias(decoder, pair[0], pair[1]); + const encoder = try object(try field(root, "encoder")); + try string(encoder, "version", "c-radio_v4-so400m"); + try number(encoder, "patch_size", patch); + try number(encoder, "max_resolution", max_height); + const args = try object(try field(encoder, "args")); + try string(args, "model", "vit_so400m_patch16_224"); + try number(args, "cpe_max_size", max_height); + try number(args, "register_multiple", vision_prefix); +} + +pub fn validateProcessor(a: std.mem.Allocator, bytes: []const u8) !void { + const doc = try std.json.parseFromSlice(Value, a, bytes, .{}); + defer doc.deinit(); + const root = try object(doc.value); + try boolean(root, "navit_packing", true); + try boolean(root, "do_normalize", true); + try boolean(root, "do_rescale", true); + try string(root, "navit_padding_mode", "top_left_white"); + try string(root, "navit_coordinate_frame", "content_normalized"); + const size = try object(try field(root, "size")); + try number(size, "height", max_height); + try number(size, "width", max_width); + const final = try field(root, "final_size"); + if (final != .array or final.array.items.len != 2) return error.BadParseConfig; + for (final.array.items, [_]i64{ max_height, max_width }) |v, want| { + if (v != .integer or v.integer != want) return error.UnsupportedParseConfig; + } + try floats(root, "image_mean", &mean); + try floats(root, "image_std", &deviation); + const scale = try field(root, "rescale_factor"); + if (scale != .float or @abs(scale.float - 1.0 / 255.0) > 1e-15) return error.UnsupportedParseConfig; +} + +fn object(v: Value) Error!Object { + return if (v == .object) v.object else error.BadParseConfig; +} +fn field(o: Object, key: []const u8) Error!Value { + return o.get(key) orelse error.BadParseConfig; +} +fn number(o: Object, key: []const u8, want: i64) Error!void { + const v = try field(o, key); + if (v != .integer) return error.BadParseConfig; + if (v.integer != want) return error.UnsupportedParseConfig; +} +fn alias(o: Object, key: []const u8, want: i64) Error!void { + if (o.contains(key)) try number(o, key, want); +} +fn string(o: Object, key: []const u8, want: []const u8) Error!void { + const v = try field(o, key); + if (v != .string) return error.BadParseConfig; + if (!std.mem.eql(u8, v.string, want)) return error.UnsupportedParseConfig; +} +fn boolean(o: Object, key: []const u8, want: bool) Error!void { + const v = try field(o, key); + if (v != .bool) return error.BadParseConfig; + if (v.bool != want) return error.UnsupportedParseConfig; +} +fn floats(o: Object, key: []const u8, want: []const f32) Error!void { + const v = try field(o, key); + if (v != .array or v.array.items.len != want.len) return error.BadParseConfig; + for (v.array.items, want) |x, w| { + const got: f32 = switch (x) { + .float => @floatCast(x.float), + .integer => @floatFromInt(x.integer), + else => return error.BadParseConfig, + }; + if (got != w) return error.UnsupportedParseConfig; + } +} diff --git a/zig/src/families/nemotron_parse/decoder.zig b/zig/src/families/nemotron_parse/decoder.zig new file mode 100644 index 0000000000..3450cae47f --- /dev/null +++ b/zig/src/families/nemotron_parse/decoder.zig @@ -0,0 +1,151 @@ +//! Parse's causal mBART target: biased pre-norm GQA, exact GELU and tied BF16 logits, without text RoPE. +const std = @import("std"); +const hip = @import("hip"); +const c = @import("config.zig"); +const Compute = @import("compute.zig").Compute; +const Tensor = hip.ops.Tensor; + +pub const Layer = struct { k: hip.DeviceBuffer, v: hip.DeviceBuffer }; +pub const State = struct { + layers: [c.layers]Layer, + hidden: hip.DeviceBuffer, + total: usize, + len: usize = 0, + + pub fn init(r: *const hip.Runtime, total: usize) !State { + if (total == 0 or total > c.context) return error.BadContext; + var out: State = .{ .layers = undefined, .hidden = undefined, .total = total }; + var loaded: usize = 0; + errdefer for (out.layers[0..loaded]) |*l| { + l.k.free(); + l.v.free(); + }; + const bytes = c.kv_heads * c.head_dim * total * 2; + for (&out.layers) |*l| { + l.k = try hip.DeviceBuffer.alloc(r, bytes); + errdefer l.k.free(); + l.v = try hip.DeviceBuffer.alloc(r, bytes); + loaded += 1; + } + out.hidden = try hip.DeviceBuffer.alloc(r, total * c.hidden * 2); + return out; + } + + pub fn deinit(s: *State) void { + for (&s.layers) |*l| { + l.k.free(); + l.v.free(); + } + s.hidden.free(); + s.* = undefined; + } + + /// Truncation leaves accepted KV and final-hidden rows intact; later windows overwrite discarded suffix rows. + pub fn keep(s: *State, length: usize) !void { + if (length > s.len) return error.BadKeep; + s.len = length; + } + + pub fn hiddenAt(s: *const State, position: usize) !Tensor { + if (position >= s.len) return error.BadPosition; + return .{ .ptr = s.hidden.base() + position * c.hidden * 2, .kind = .bf16 }; + } +}; + +pub fn embed(compute: Compute, ids: u64, rows: usize) !Tensor { + const out = try compute.alloc(rows * c.hidden, .bf16); + try compute.kernels.embed(compute.stream, try compute.weights.get("decoder.embed_tokens.weight"), ids, out, rows, c.hidden, 32); + return out; +} + +/// One stream's contiguous span within a packed forward; its KV remains separate from every other stream. +pub const Span = struct { state: *State, rows: usize }; + +/// Append embedded rows to one stream; the returned final norm lives in owned stream state. +pub fn forward(compute: Compute, state: *State, embeddings: Tensor, rows: usize) !Tensor { + const at = state.len; + _ = try forwardMany(compute, &.{.{ .state = state, .rows = rows }}, embeddings); + return state.hiddenAt(at); +} + +/// One shared dense forward over all spans, with causal attention confined to each span's own KV. +pub fn forwardMany(compute: Compute, spans: []const Span, embeddings: Tensor) !Tensor { + const rows = try rowsOf(spans); + const x = try compute.norm(embeddings, "decoder.layernorm_embedding", rows, c.hidden, c.norm_eps); + const mark = compute.ops.arena.mark(); + try compute.record("decoder.embedding_norm", x, rows, c.hidden); + var buf: [160]u8 = undefined; + for (0..c.layers) |i| { + try compute.poll(); + try layer(compute, spans, i, x, rows); + try compute.record(try std.fmt.bufPrint(&buf, "decoder.layer.{d}", .{i}), x, rows, c.hidden); + compute.ops.arena.release(mark); + } + const out = try compute.norm(x, "decoder.layer_norm", rows, c.hidden, c.norm_eps); + try compute.record("decoder.final_norm", out, rows, c.hidden); + var offset: usize = 0; + for (spans) |span| { + try hip.raw.copy(span.state.hidden, span.state.len * c.hidden * 2, out.ptr + offset * c.hidden * 2, span.rows * c.hidden * 2, compute.stream.handle); + offset += span.rows; + } + // Publish lengths only after all commands have been queued successfully. + for (spans) |span| span.state.len += span.rows; + return out; +} + +/// Check every context and duplicate state before a batch can enqueue any GPU work. +pub fn rowsOf(spans: []const Span) !usize { + var rows: usize = 0; + for (spans, 0..) |span, i| { + const state = span.state; + if (span.rows == 0 or state.len > state.total or span.rows > state.total - state.len) return error.ContextExceeded; + for (spans[0..i]) |prior| if (prior.state == state) return error.DuplicateState; + rows = try std.math.add(usize, rows, span.rows); + } + if (rows == 0) return error.ContextExceeded; + return rows; +} + +pub fn logits(compute: Compute, hidden: Tensor, rows: usize) !Tensor { + return compute.linear(hidden, "decoder.embed_tokens", rows, c.vocab, c.hidden, false); +} + +fn layer(compute: Compute, spans: []const Span, index: usize, x: Tensor, rows: usize) !void { + var name: [192]u8 = undefined; + var root_buf: [64]u8 = undefined; + const root = try std.fmt.bufPrint(&root_buf, "decoder.layers.{d}", .{index}); + const norm = try compute.norm(x, try sub(&name, root, "self_attn_layer_norm"), rows, c.hidden, c.norm_eps); + const q = try compute.linear(norm, try sub(&name, root, "self_attn.q_proj"), rows, c.hidden, c.hidden, true); + const k = try compute.linear(norm, try sub(&name, root, "self_attn.k_proj"), rows, c.kv_heads * c.head_dim, c.hidden, true); + const v = try compute.linear(norm, try sub(&name, root, "self_attn.v_proj"), rows, c.kv_heads * c.head_dim, c.hidden, true); + const row_att = try compute.alloc(rows * c.hidden, .bf16); + const mark = compute.ops.arena.mark(); + var offset: usize = 0; + for (spans) |span| { + const state = span.state; + const qs: Tensor = .{ .ptr = q.ptr + offset * c.hidden * 2, .kind = .bf16 }; + const ks: Tensor = .{ .ptr = k.ptr + offset * c.kv_heads * c.head_dim * 2, .kind = .bf16 }; + const vs: Tensor = .{ .ptr = v.ptr + offset * c.kv_heads * c.head_dim * 2, .kind = .bf16 }; + const wide_q = try compute.heads(qs, span.rows, c.heads, c.head_dim, true); + const cache = state.layers[index]; + try compute.ops.kvWrite(ks, cache.k.base(), span.rows, c.kv_heads, c.head_dim, state.total, state.len); + try compute.ops.kvWrite(vs, cache.v.base(), span.rows, c.kv_heads, c.head_dim, state.total, state.len); + const attended = try compute.alloc(span.rows * c.hidden, .f32); + try compute.ops.causalPrefill(wide_q.ptr, .{ .k = cache.k.base(), .v = cache.v.base(), .kind = .bf16, .kv_heads = c.kv_heads, .total = state.total, .d = c.head_dim }, attended.ptr, span.rows, state.len + span.rows, c.heads, 1 / @sqrt(@as(f32, c.head_dim)), state.len); + const dst: Tensor = .{ .ptr = row_att.ptr + offset * c.hidden * 2, .kind = .bf16 }; + try compute.kernels.layout(compute.stream, attended, dst, span.rows, c.heads, c.head_dim, .heads_to_rows); + offset += span.rows; + compute.ops.arena.release(mark); + } + const projected = try compute.linear(row_att, try sub(&name, root, "self_attn.out_proj"), rows, c.hidden, c.hidden, true); + try compute.residual(x, projected, rows * c.hidden); + const ffn_norm = try compute.norm(x, try sub(&name, root, "final_layer_norm"), rows, c.hidden, c.norm_eps); + const up = try compute.linear(ffn_norm, try sub(&name, root, "fc1"), rows, c.ffn, c.hidden, true); + try compute.gelu(up, rows * c.ffn); + const down = try compute.linear(up, try sub(&name, root, "fc2"), rows, c.hidden, c.ffn, true); + try compute.residual(x, down, rows * c.hidden); +} + +fn sub(buf: []u8, root: []const u8, end: []const u8) ![]const u8 { + return std.fmt.bufPrint(buf, "{s}.{s}", .{ root, end }); +} diff --git a/zig/src/families/nemotron_parse/engine.zig b/zig/src/families/nemotron_parse/engine.zig new file mode 100644 index 0000000000..3b8350fbb7 --- /dev/null +++ b/zig/src/families/nemotron_parse/engine.zig @@ -0,0 +1,84 @@ +//! Owned single-rank HIP resources; the shared lane core owns scheduling and request lifetime. +const std = @import("std"); +const hip = @import("hip"); +const c = @import("config.zig"); +const Weights = @import("weights.zig").Weights; +const Compute = @import("compute.zig").Compute; + +pub const scratch_bytes = 1 << 30; +pub const weight_bytes = 1261601004; +pub const max_image_rows = 3329; +pub const max_raw_bytes = 13312 * 768 * 4; +pub const prefix_bytes = max_image_rows * c.hidden * 2; + +pub fn stateBytes(window: usize) usize { + return window * (c.hidden + c.layers * 2 * c.kv_heads * c.head_dim) * 2; +} + +pub const Engine = struct { + a: std.mem.Allocator, + io: std.Io, + r: hip.Runtime, + ctx: hip.Context, + stream: hip.Stream, + launcher: hip.Launcher, + kernels: hip.vision.Kernels, + weights: Weights, + scratch: hip.Arena, + window: usize, + streams: usize, + + pub fn open(a: std.mem.Allocator, io: std.Io, dir: []const u8, dev: hip.Device, window: usize, asked: usize, fixed: bool) !*Engine { + if (window == 0 or window > c.context or asked == 0 or asked > 64) return error.BadParseAdmission; + const e = try a.create(Engine); + errdefer a.destroy(e); + e.a = a; + e.io = io; + e.window = window; + e.r = try hip.Runtime.open(); + errdefer e.r.close(); + e.ctx = try hip.Context.init(&e.r, dev.index); + errdefer e.ctx.deinit(); + _ = try hip.Device.kernelCaps(&e.r, dev.index); + const pool = try hip.admission.budget(a, io, e.ctx); + const room = pool.room(hip.usage(false).device); + const base = weight_bytes + scratch_bytes + max_raw_bytes; + const per = stateBytes(window) + prefix_bytes; + if (room < base + per) return error.ParseDoesNotFit; + e.streams = @min(asked, (room - base) / per); + if (fixed and e.streams < asked) return error.ParseDoesNotFit; + e.stream = try hip.Stream.init(&e.r); + errdefer e.stream.deinit(); + e.launcher = try hip.Launcher.load(&e.r, dev.policy.choices(), hip.kernels.images); + errdefer e.launcher.unload(); + // Every row uses the same attention tile across prompt chunks and shared verification widths. + e.launcher.wide = false; + e.kernels = try hip.vision.Kernels.load(e.launcher.mods[0]); + e.weights = try Weights.load(a, &e.r, io, dir); + errdefer e.weights.deinit(); + e.scratch = try hip.Arena.init(&e.r, scratch_bytes); + return e; + } + + pub fn deinit(e: *Engine) void { + e.stream.synchronize() catch {}; + e.scratch.deinit(); + e.weights.deinit(); + e.launcher.unload(); + e.stream.deinit(); + e.ctx.deinit(); + e.r.close(); + const a = e.a; + a.destroy(e); + } + + pub fn compute(e: *Engine) Compute { + return .{ .ops = .{ .l = &e.launcher, .stream = e.stream.handle, .bf16 = true, .arena = &e.scratch }, .kernels = e.kernels, .stream = e.stream, .weights = &e.weights }; + } + + /// Refuse new lane storage if another GPU user consumed the unallocated part of our initial plan. + pub fn laneRoom(e: *Engine) !void { + const pool = try hip.admission.budget(e.a, e.io, e.ctx); + if (pool.room(hip.usage(false).device) < stateBytes(e.window) + prefix_bytes + max_raw_bytes) return error.LaneMemory; + } +}; diff --git a/zig/src/families/nemotron_parse/fixtures/processor.json b/zig/src/families/nemotron_parse/fixtures/processor.json new file mode 100644 index 0000000000..0a6f8c8c06 --- /dev/null +++ b/zig/src/families/nemotron_parse/fixtures/processor.json @@ -0,0 +1,85 @@ +{ + "source_revision": "4035ffb1d8f8bafcdb867dfaf6f7514f101464f9", + "cases": [ + { + "height": 1, + "width": 1, + "content_h": 1, + "content_w": 1, + "padded_h": 16, + "padded_w": 64, + "tokens": 2, + "sha256": "a4588849fb862743ef48b8b385d9525d8fefe3cd533d910e63a23eb28fb28566" + }, + { + "height": 17, + "width": 65, + "content_h": 17, + "content_w": 65, + "padded_h": 32, + "padded_w": 128, + "tokens": 5, + "sha256": "2f10fad8a77696095ffe16ff7ab61a20f262e3290e7f784cf4d0a4b2c3dcf7b9" + }, + { + "height": 31, + "width": 129, + "content_h": 31, + "content_w": 129, + "padded_h": 32, + "padded_w": 192, + "tokens": 7, + "sha256": "ade4e732a116c088a16248142c22f9876182d6bdb1c5f4756b7b3aa673692a49" + }, + { + "height": 113, + "width": 97, + "content_h": 113, + "content_w": 97, + "padded_h": 128, + "padded_w": 128, + "tokens": 17, + "sha256": "e7f0adc4fa13b6d738c3294471c6ab762d560c9b3b10bd173b82e184a8da6e39" + }, + { + "height": 2049, + "width": 101, + "content_h": 2048, + "content_w": 100, + "padded_h": 2048, + "padded_w": 128, + "tokens": 257, + "sha256": "e025eca32b2f9c8f362a0de9838a5ba91a7abb203af00105655d1413094c968d" + }, + { + "height": 99, + "width": 2001, + "content_h": 82, + "content_w": 1664, + "padded_h": 96, + "padded_w": 1664, + "tokens": 157, + "sha256": "48c2d9cc6b2c59372ae95e634ed6341547372f1dd3409dbac9bbba1f88b425bf" + }, + { + "height": 2201, + "width": 1777, + "content_h": 2048, + "content_w": 1653, + "padded_h": 2048, + "padded_w": 1664, + "tokens": 3329, + "sha256": "a1884e9025579569bb00835e7db1f94603646fef0b8b1dbb11752ceabd8b4bf4" + }, + { + "height": 2048, + "width": 1664, + "content_h": 2048, + "content_w": 1664, + "padded_h": 2048, + "padded_w": 1664, + "tokens": 3329, + "sha256": "d3fe3629bc22f8ee2f75ccbc16236777e1dad2406f118db0f2515be42af33c7d" + } + ] +} diff --git a/zig/src/families/nemotron_parse/hip.zig b/zig/src/families/nemotron_parse/hip.zig new file mode 100644 index 0000000000..b4100027e9 --- /dev/null +++ b/zig/src/families/nemotron_parse/hip.zig @@ -0,0 +1,116 @@ +//! The native HIP family registry entry; the shared lane engine serves image-prefix OCR and verifies proposals. +const std = @import("std"); +const hip = @import("hip"); +const api = @import("engine_api"); +const lanes = @import("lanes"); +const serve = @import("hip_serve"); +const config = @import("config.zig"); +const engine = @import("engine.zig"); +const Backend = @import("backend.zig").Backend; +const processor = @import("processor.zig"); + +pub const model_type = config.model_type; +pub const raw_text = false; +/// OCR output has no reasoning block: every reply token is answer content. +pub const reasoning = false; +pub const formats: []const []const u8 = &.{"fp32"}; +pub const default_context: i64 = config.context; +pub const prefill_step: u32 = 512; +pub const slide_refusal = "--slide: Nemotron Parse has no native learner"; +pub const confidence_refusal = "--mtp-confidence: Nemotron Parse uses its trained one-step proposal without a confidence cut"; + +/// Refuse unsupported combinations before GPU setup or a tensor-parallel join (hip_families.refused). +pub fn refusal(o: api.Open) ?[]const u8 { + if (o.tp != 1 or o.rank != 0) return "Nemotron Parse currently serves one HIP rank (--tp 1 --rank 0)"; + if (o.vision == null) return "Nemotron Parse requires --vision and one image per request"; + if (o.context) |n| if (n > config.context) return "Nemotron Parse's visual prefix and text share a 16384-token context"; + if (o.mtp_drafts) |n| if (n > 1) return "Nemotron Parse has one trained proposal head; use --mtp-drafts 0 or 1"; + return null; +} + +pub const Loaded = struct { + backend: lanes.backend.Backend, + facts: lanes.Model, + rows: u32, + ctx: *anyopaque, + deinit: *const fn (*anyopaque) void, + streams: usize, + startup: []const u8, + /// Its image tower, and its prompt cache when the budget and the slots allow one. + offer: serve.Offer, +}; + +const Owned = struct { + a: std.mem.Allocator, + e: *engine.Engine, + h: *Backend, + startup: []const u8, + image_tokens: u32, +}; + +fn release(ptr: *anyopaque) void { + const own: *Owned = @ptrCast(@alignCast(ptr)); + own.h.deinit(); + own.e.deinit(); + own.a.free(own.startup); + own.a.destroy(own); +} + +fn prepare(ptr: *anyopaque, a: std.mem.Allocator, bytes: []const u8, limit: u32) anyerror!api.PreparedImage { + const own: *Owned = @ptrCast(@alignCast(ptr)); + const p = try processor.prepare(own.e.io, a, bytes); + errdefer a.free(p.pixels); + const g = p.geometry; + if (g.tokens() > @min(limit, own.image_tokens)) return error.ImageTooSmall; + return .{ .pixels = p.pixels, .gh = @intCast(g.padded.h / 16), .gw = @intCast(g.padded.w / 16), .tokens = @intCast(g.tokens()), .content_size = .{ @intCast(g.content.h), @intCast(g.content.w) }, .hash = p.hash }; +} + +pub fn open(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, dev: hip.Device, dir: []const u8, o: serve.Options, problem: *[]const u8) !Loaded { + if (o.world != 1 or o.rank != 0 or o.vision == null) return error.Refused; + const e = engine.Engine.open(gpa, io, dir, dev, o.window, o.streams, o.fixed) catch |err| { + problem.* = try std.fmt.allocPrint(a, "cannot load native HIP Nemotron Parse ({t})", .{err}); + return error.Refused; + }; + errdefer e.deinit(); + const h = try Backend.init(gpa, e, dev.policy.mtp.drafts != 0); + errdefer h.deinit(); + const pool = try hip.admission.budget(gpa, io, e.ctx); + const planned = e.streams * (engine.stateBytes(e.window) + engine.prefix_bytes) + engine.max_raw_bytes; + const spare = pool.room(hip.usage(false).device) -| planned; + const budget: usize = if (o.cache_gib) |g| blk: { + if (!std.math.isFinite(g) or g < 0 or g * (1 << 30) > @as(f64, @floatFromInt(spare))) { + problem.* = "the Parse prompt cache exceeds free memory after its admitted lanes; lower --prompt-cache-gib"; + return error.Refused; + } + break :blk @intFromFloat(g * (1 << 30)); + } else @min(spare, 256 << 20); + h.cache.budget = if (dev.policy.slots == 0) 0 else budget; + const own = try gpa.create(Owned); + errdefer gpa.destroy(own); + const startup = try std.fmt.allocPrint(gpa, "HIP Nemotron Parse BF16: {d} streams, {d} context rows including image; native RADIO and mBART, one-step MTP", .{ e.streams, e.window }); + own.* = .{ .a = gpa, .e = e, .h = h, .startup = startup, .image_tokens = @min(o.vision.?.image_tokens, engine.max_image_rows) }; + return .{ .backend = h.backend(), .facts = h.facts(), .rows = @import("backend.zig").max_window, .ctx = own, .deinit = release, .streams = e.streams, .startup = startup, .offer = .{ + .vision = .{ .ctx = own, .prepare = prepare, .image_token = 1, .markers = &.{1}, .image_tokens = own.image_tokens, .max_images = 1 }, + .cache = if (h.cache.budget > 0) &h.cache else null, + } }; +} + +pub fn follow(_: std.mem.Allocator, _: std.mem.Allocator, _: std.Io, _: hip.Device, _: []const u8, _: serve.Options, problem: *[]const u8) !void { + problem.* = "Nemotron Parse supports one HIP rank"; + return error.Refused; +} + +test "Parse refuses missing image input, tensor parallelism and unsupported trained-head depth before GPU setup" { + try std.testing.expect(refusal(.{ .dir = "d", .model_type = model_type }) != null); + const base: api.Open = .{ .dir = "d", .model_type = model_type, .vision = .{} }; + try std.testing.expect(refusal(base) == null); + var other = base; + other.tp = 2; + try std.testing.expect(refusal(other) != null); + other = base; + other.mtp_drafts = 2; + try std.testing.expect(refusal(other) != null); + other = base; + other.context = 16385; + try std.testing.expect(refusal(other) != null); +} diff --git a/zig/src/families/nemotron_parse/hip_ops_test.zig b/zig/src/families/nemotron_parse/hip_ops_test.zig new file mode 100644 index 0000000000..ab989e4277 --- /dev/null +++ b/zig/src/families/nemotron_parse/hip_ops_test.zig @@ -0,0 +1,339 @@ +//! Focused shared HIP operator, position and optional checkpoint upload checks for Parse. +const std = @import("std"); +const hip = @import("hip"); +const core = @import("core"); +const a = std.testing.allocator; +const bf = core.bf16_round; +const Tensor = hip.ops.Tensor; + +const Rig = struct { + r: hip.Runtime, + ctx: hip.Context, + s: hip.Stream, + module: hip.Module, + k: hip.vision.Kernels, + + fn init(t: *Rig) !void { + t.r = try hip.Runtime.open(); + errdefer t.r.close(); + t.ctx = try hip.Context.init(&t.r, 0); + errdefer t.ctx.deinit(); + t.s = try hip.Stream.init(&t.r); + errdefer t.s.deinit(); + t.module = try hip.Module.load(&t.r, @import("parse_gpu").bytes); + errdefer t.module.unload(); + t.k = try hip.vision.Kernels.load(t.module); + } + fn deinit(t: *Rig) void { + t.s.synchronize() catch {}; + t.module.unload(); + t.s.deinit(); + t.ctx.deinit(); + t.r.close(); + } +}; + +const Buffer = struct { + dev: hip.DeviceBuffer, + host: []u16, + count: usize, + + fn init(r: *const hip.Runtime, count: usize, seed: usize) !Buffer { + const host = try a.alloc(u16, count + 64); + errdefer a.free(host); + @memset(host, 0xa5a5); + for (host[32..][0..count], 0..) |*v, i| { + const x: i32 = @intCast((i * 7 + seed) % 37); + v.* = bf.bits(@as(f32, @floatFromInt(x - 18)) / 16); + } + return .{ .dev = try hip.DeviceBuffer.fromHost(r, std.mem.sliceAsBytes(host)), .host = host, .count = count }; + } + fn free(b: *Buffer) void { + b.dev.free(); + a.free(b.host); + } + fn tensor(b: Buffer) Tensor { + return .{ .ptr = b.dev.base() + 64, .kind = .bf16 }; + } + fn values(b: Buffer) []u16 { + return b.host[32..][0..b.count]; + } + fn read(b: Buffer) ![]u16 { + const got = try a.alloc(u16, b.host.len); + errdefer a.free(got); + try b.dev.download(0, std.mem.sliceAsBytes(got)); + try std.testing.expectEqualSlices(u16, b.host[0..32], got[0..32]); + try std.testing.expectEqualSlices(u16, b.host[b.count + 32 ..], got[b.count + 32 ..]); + return got; + } + fn unchanged(b: Buffer) !void { + const got = try b.read(); + defer a.free(got); + try std.testing.expectEqualSlices(u16, b.host, got); + } +}; + +test "Parse HIP dense projections match FP64 at tails and served widths, with five repeats and guards" { + var t: Rig = undefined; + try t.init(); + defer t.deinit(); + for ([_][3]usize{ .{ 1, 1, 1 }, .{ 15, 63, 19 }, .{ 16, 64, 16 }, .{ 17, 65, 1152 }, .{ 2, 1024, 4304 }, .{ 1, 4304, 1152 } }) |shape| { + const m, const n, const k = .{ shape[0], shape[1], shape[2] }; + var x = try Buffer.init(&t.r, m * k, 3); + defer x.free(); + var w = try Buffer.init(&t.r, n * k, 5); + defer w.free(); + var bias = try Buffer.init(&t.r, n, 7); + defer bias.free(); + var out = try Buffer.init(&t.r, m * n, 0); + defer out.free(); + const want = try a.alloc(u16, m * n); + defer a.free(want); + for (0..m) |r| for (0..n) |col| { + var sum: f64 = bf.value(bias.values()[col]); + for (0..k) |j| sum += @as(f64, bf.value(x.values()[r * k + j])) * bf.value(w.values()[col * k + j]); + want[r * n + col] = bf.bits(@floatCast(sum)); // dyadic fixture sums are exactly representable in FP32 + }; + for (0..5) |_| { + try t.k.linear(t.s, x.tensor(), w.tensor().ptr, bias.tensor().ptr, out.tensor(), m, n, k); + try t.s.synchronize(); + const got = try out.read(); + defer a.free(got); + try std.testing.expectEqualSlices(u16, want, got[32..][0 .. m * n]); + } + try x.unchanged(); + try w.unchanged(); + try bias.unchanged(); + } +} + +test "Parse HIP LayerNorm matches FP64 with BF16 output and propagates non-finite rows" { + var t: Rig = undefined; + try t.init(); + defer t.deinit(); + for ([_]usize{ 1, 255, 256, 257, 1024, 1152 }) |width| { + var x = try Buffer.init(&t.r, 4 * width, 3); + defer x.free(); + var w = try Buffer.init(&t.r, width, 6); + defer w.free(); + var b = try Buffer.init(&t.r, width, 9); + defer b.free(); + var out = try Buffer.init(&t.r, 4 * width, 0); + defer out.free(); + for ([_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32) }, 1..) |v, row| + x.values()[row * width] = bf.bits(v); + try x.dev.upload(0, std.mem.sliceAsBytes(x.host)); + try t.k.norm(t.s, x.tensor(), w.tensor().ptr, b.tensor().ptr, out.tensor(), 4, width, 1e-6); + try t.s.synchronize(); + const got = try out.read(); + defer a.free(got); + var mean: f64 = 0; + for (x.values()[0..width]) |v| mean += bf.value(v); + mean /= @floatFromInt(width); + var variance: f64 = 0; + for (x.values()[0..width]) |v| variance += std.math.pow(f64, @as(f64, bf.value(v)) - mean, 2); + variance /= @floatFromInt(width); + for (0..width) |j| { + const want = (@as(f64, bf.value(x.values()[j])) - mean) / @sqrt(variance + 1e-6) * bf.value(w.values()[j]) + bf.value(b.values()[j]); + const actual: f64 = bf.value(got[32 + j]); + try std.testing.expect(@abs(actual - want) <= @abs(want) / 256 + 2e-6); + } + for (got[32 + width ..][0 .. 3 * width]) |v| try std.testing.expect(std.math.isNan(bf.value(v))); + try x.unchanged(); + try w.unchanged(); + try b.unchanged(); + } +} + +test "Parse HIP horizontal neck and attention layouts preserve every element" { + var t: Rig = undefined; + try t.init(); + defer t.deinit(); + const rows = 12; + const heads = 3; + const dim = 7; + var x = try Buffer.init(&t.r, rows * heads * dim, 3); + defer x.free(); + var out = try Buffer.init(&t.r, rows * heads * dim, 0); + defer out.free(); + for ([_]hip.vision.Kernels.Layout{ .heads_to_rows, .horizontal_four }) |mode| { + try t.k.layout(t.s, x.tensor(), out.tensor(), rows, if (mode == .heads_to_rows) heads else 1, if (mode == .heads_to_rows) dim else heads * dim, mode); + try t.s.synchronize(); + const got = try out.read(); + defer a.free(got); + for (0..x.count) |i| { + const from = if (mode == .heads_to_rows) + ((i / dim % heads) * rows + i / (heads * dim)) * dim + i % dim + else + (i / (4 * heads * dim) * 4 + i % 4) * heads * dim + i / 4 % (heads * dim); + try std.testing.expectEqual(x.values()[from], got[32 + i]); + } + } + try x.unchanged(); +} + +extern "c" fn erf(x: f64) f64; + +test "Parse HIP GELU, residual and embedding scale retain guards and repeat bytes" { + var t: Rig = undefined; + try t.init(); + defer t.deinit(); + const count = 1027; + var x = try Buffer.init(&t.r, count, 3); + defer x.free(); + var y = try Buffer.init(&t.r, count, 9); + defer y.free(); + var out = try Buffer.init(&t.r, count, 0); + defer out.free(); + for ([_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32) }, 0..) |v, i| + x.values()[i] = bf.bits(v); + try x.dev.upload(0, std.mem.sliceAsBytes(x.host)); + for ([_]hip.vision.Kernels.Point{ .gelu, .add, .scale }) |mode| { + var first: ?[]u16 = null; + defer if (first) |v| a.free(v); + for (0..5) |_| { + try t.k.point(t.s, x.tensor(), y.tensor(), out.tensor(), count, mode, 32); + try t.s.synchronize(); + const got = try out.read(); + defer a.free(got); + if (first) |v| try std.testing.expectEqualSlices(u16, v, got) else first = try a.dupe(u16, got); + for (x.values(), y.values(), got[32..][0..count]) |xv, yv, actual| { + const v: f64 = bf.value(xv); + const want = switch (mode) { + .gelu => (0.5 * v) * (1 + erf(v / @sqrt(@as(f64, 2)))), + .add => v + bf.value(yv), + .scale => v * 32, + }; + const value: f64 = bf.value(actual); + if (std.math.isNan(want)) { + try std.testing.expect(std.math.isNan(value)); + } else if (std.math.isInf(want)) { + try std.testing.expectEqual(want, value); + } else { + try std.testing.expect(@abs(value - want) <= @abs(want) / 256 + 2e-7); + } + } + } + } + try x.unchanged(); + try y.unchanged(); +} + +test "Parse HIP uploads the actual FP32 checkpoint as owned BF16 weights (TF_PARSE_DIR)" { + const dir = @import("sys").getenv("TF_PARSE_DIR") orelse return error.SkipZigTest; + var t: Rig = undefined; + try t.init(); + defer t.deinit(); + const before = hip.usage(false).device; + var w = try @import("weights.zig").Weights.load(a, &t.r, std.testing.io, dir); + var live = true; + defer if (live) w.deinit(); + try std.testing.expectEqual(@as(usize, 509), w.names.count()); + try std.testing.expectEqual(@as(usize, 1261601004), w.bytes); + var ck = try @import("checkpoint.zig").Checkpoint.open(a, std.testing.io, dir); + defer ck.close(); + const table = try ck.get("decoder.embed_tokens.weight"); + var got: [32]u16 = undefined; + try w.names.get("decoder.embed_tokens.weight").?.download(0, std.mem.sliceAsBytes(&got)); + for (got, 0..) |value, i| { + const bits = std.mem.readInt(u32, table.bytes[i * 4 ..][0..4], .little); + try std.testing.expectEqual(bf.bits(@bitCast(bits)), value); + } + try std.testing.expectEqual(before + w.bytes, hip.usage(false).device); + { + const ids = [_]i32{ 0, 1, 2, 72511 }; + var dev_ids = try hip.DeviceBuffer.fromHost(&t.r, std.mem.sliceAsBytes(&ids)); + defer dev_ids.free(); + var out = try hip.DeviceBuffer.alloc(&t.r, ids.len * 1024 * 2); + defer out.free(); + try t.k.embed(t.s, try w.get("decoder.embed_tokens.weight"), dev_ids.base(), .{ .ptr = out.base(), .kind = .bf16 }, ids.len, 1024, 32); + try t.s.synchronize(); + var values: [4 * 1024]u16 = undefined; + try out.download(0, std.mem.sliceAsBytes(&values)); + for (ids, 0..) |id, row| for (0..1024) |j| { + const at = (@as(usize, @intCast(id)) * 1024 + j) * 4; + const bits = std.mem.readInt(u32, table.bytes[at..][0..4], .little); + const rounded = bf.round32(@bitCast(bits)); + try std.testing.expectEqual(bf.bits(rounded * 32), values[row * 1024 + j]); + }; + } + + w.deinit(); + live = false; + try std.testing.expectEqual(before, hip.usage(false).device); +} + +test { + _ = @import("hip_position_test.zig"); +} + +test "Parse decoder refuses invalid contexts and suffix growth before GPU access" { + const decoder = @import("decoder.zig"); + const r: hip.Runtime = undefined; + try std.testing.expectError(error.BadContext, decoder.State.init(&r, 0)); + try std.testing.expectError(error.BadContext, decoder.State.init(&r, 16385)); + var state: decoder.State = .{ .layers = undefined, .hidden = undefined, .total = 128, .len = 65 }; + try state.keep(63); + try std.testing.expectEqual(@as(usize, 63), state.len); + try std.testing.expectError(error.BadKeep, state.keep(64)); + try std.testing.expectError(error.BadPosition, state.hiddenAt(63)); + try std.testing.expectError(error.ContextExceeded, decoder.rowsOf(&.{})); + try std.testing.expectError(error.DuplicateState, decoder.rowsOf(&.{ .{ .state = &state, .rows = 1 }, .{ .state = &state, .rows = 1 } })); + const compute: @import("compute.zig").Compute = undefined; + try std.testing.expectError(error.ContextExceeded, decoder.forward(compute, &state, undefined, 66)); + try std.testing.expectError(error.ContextExceeded, decoder.forward(compute, &state, undefined, 0)); +} + +test "Parse HIP linear matches FP64 classes for non-finite products and preserves finite dots" { + var t: Rig = undefined; + try t.init(); + defer t.deinit(); + const rows = 5; + const n = 65; + const width = 1025; + var x = try Buffer.init(&t.r, rows * width, 0); + defer x.free(); + var w = try Buffer.init(&t.r, n * width, 3); + defer w.free(); + var bias = try Buffer.init(&t.r, n, 7); + defer bias.free(); + var out = try Buffer.init(&t.r, rows * n, 0); + defer out.free(); + var rng: std.Random.DefaultPrng = .init(42); + const random = rng.random(); + for (x.values()) |*v| v.* = bf.bits(std.math.ldexp(random.float(f32) * 2 - 1, @as(i32, @intCast(random.uintLessThan(u8, 17))) - 8)); + for (w.values()) |*v| v.* = bf.bits(std.math.ldexp(random.float(f32) * 2 - 1, @as(i32, @intCast(random.uintLessThan(u8, 17))) - 8)); + x.values()[width + 3] = bf.bits(std.math.nan(f32)); + x.values()[2 * width + 5] = bf.bits(std.math.inf(f32)); + x.values()[3 * width + 7] = bf.bits(-std.math.inf(f32)); + x.values()[4 * width + 9] = bf.bits(std.math.inf(f32)); + x.values()[4 * width + 11] = bf.bits(-std.math.inf(f32)); + try x.dev.upload(0, std.mem.sliceAsBytes(x.host)); + try w.dev.upload(0, std.mem.sliceAsBytes(w.host)); + try t.k.linear(t.s, x.tensor(), w.tensor().ptr, bias.tensor().ptr, out.tensor(), rows, n, width); + try t.s.synchronize(); + const got = try out.read(); + defer a.free(got); + const contract = core.contract.spec; + const spec: contract.Spec = .{ .name = "linear probe", .accum = .f32, .depth = struct { + fn f(shape: contract.Shape) usize { + return shape.k * 4 + 1; + } + }.f, .out = .bf16, .served = &.{} }; + for (0..rows) |row| for (0..n) |col| { + var sum: f64 = bf.value(bias.values()[col]); + var magnitude = @abs(sum); + for (0..width) |j| { + const term = @as(f64, bf.value(x.values()[row * width + j])) * bf.value(w.values()[col * width + j]); + sum += term; + magnitude += @abs(term); + } + const actual: f64 = bf.value(got[32 + row * n + col]); + try std.testing.expectEqual(contract.classOf(sum), contract.classOf(actual)); + if (std.math.isFinite(sum)) try std.testing.expect(@abs(actual - sum) <= + contract.bound(spec, .{ .rows = rows, .n = n, .k = width }, sum, magnitude)); + }; + try x.unchanged(); + try w.unchanged(); + try bias.unchanged(); +} diff --git a/zig/src/families/nemotron_parse/hip_position_test.zig b/zig/src/families/nemotron_parse/hip_position_test.zig new file mode 100644 index 0000000000..af4c87007f --- /dev/null +++ b/zig/src/families/nemotron_parse/hip_position_test.zig @@ -0,0 +1,80 @@ +//! Position interpolation against the actual pinned HF helper, independently generated before native execution. +const std = @import("std"); +const hip = @import("hip"); +const bf = @import("core").bf16_round; +const a = std.testing.allocator; +const io = std.testing.io; +const Cell = struct { height: usize, width: usize, gh: usize, gw: usize }; +const Manifest = struct { hidden: usize, prefix: usize, cells: []Cell }; + +fn read(dir: []const u8, name: []const u8) ![]u8 { + const path = try std.fs.path.join(a, &.{ dir, name }); + defer a.free(path); + return std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(16 << 20)); +} + +fn upload(r: *const hip.Runtime, dir: []const u8, name: []const u8) !hip.DeviceBuffer { + const bytes = try read(dir, name); + defer a.free(bytes); + return hip.DeviceBuffer.fromHost(r, bytes); +} + +test "Parse HIP positions match HF content-normalized geometry (TF_PARSE_POSITION_DIR)" { + const dir = @import("sys").getenv("TF_PARSE_POSITION_DIR") orelse return error.SkipZigTest; + const text = try read(dir, "manifest.json"); + defer a.free(text); + const doc = try std.json.parseFromSlice(Manifest, a, text, .{}); + defer doc.deinit(); + const m = doc.value; + var r = try hip.Runtime.open(); + defer r.close(); + var ctx = try hip.Context.init(&r, 0); + defer ctx.deinit(); + var s = try hip.Stream.init(&r); + defer s.deinit(); + var module = try hip.Module.load(&r, @import("parse_gpu").bytes); + defer module.unload(); + const k = try hip.vision.Kernels.load(module); + var table = try upload(&r, dir, "table.bin"); + defer table.free(); + var cls = try upload(&r, dir, "cls.bin"); + defer cls.free(); + for (m.cells, 0..) |cell, index| { + var name: [64]u8 = undefined; + var patches = try upload(&r, dir, try std.fmt.bufPrint(&name, "patches-{d}.bin", .{index})); + defer patches.free(); + const oh = (cell.height + 15) / 16; + const ow = (cell.width + 15) / 16; + var field = try hip.DeviceBuffer.alloc(&r, oh * ow * m.hidden * 4); + defer field.free(); + var mid = try hip.DeviceBuffer.alloc(&r, 128 * ow * m.hidden * 4); + defer mid.free(); + const count = (cell.gh * cell.gw + m.prefix) * m.hidden; + var out = try hip.DeviceBuffer.alloc(&r, count * 2); + defer out.free(); + const table_tensor: hip.ops.Tensor = .{ .ptr = table.base(), .kind = .bf16 }; + if (@min(oh, ow) == 1) { + try k.resize(s, table_tensor, field.base(), .{ 128, 128 }, .{ oh, ow }, m.hidden, .area); + } else { + try k.resize(s, table_tensor, mid.base(), .{ 128, 128 }, .{ 128, ow }, m.hidden, .horizontal); + try k.resize(s, .{ .ptr = mid.base(), .kind = .f32 }, field.base(), .{ 128, ow }, .{ oh, ow }, m.hidden, .vertical); + } + try k.position(s, .{ .ptr = patches.base(), .kind = .bf16 }, cls.base(), field.base(), .{ .ptr = out.base(), .kind = .bf16 }, .{ cell.gh, cell.gw }, .{ cell.height, cell.width }, m.hidden, m.prefix); + try s.synchronize(); + const want = try read(dir, try std.fmt.bufPrint(&name, "expected-{d}.bin", .{index})); + defer a.free(want); + const got = try a.alloc(u16, count); + defer a.free(got); + try out.download(0, std.mem.sliceAsBytes(got)); + try std.testing.expectEqual(count * 2, want.len); + var differences: usize = 0; + var worst: f64 = 0; + for (got, 0..) |v, i| { + const expected = std.mem.readInt(u16, want[i * 2 ..][0..2], .little); + differences += @intFromBool(v != expected); + worst = @max(worst, @abs(@as(f64, bf.value(v)) - bf.value(expected))); + } + std.debug.print("Parse positions {d}x{d}: {d}/{d} byte differences, max abs {d:.9}\n", .{ cell.height, cell.width, differences, count, worst }); + try std.testing.expectEqual(@as(usize, 0), differences); + } +} diff --git a/zig/src/families/nemotron_parse/host_test.zig b/zig/src/families/nemotron_parse/host_test.zig new file mode 100644 index 0000000000..58db72787e --- /dev/null +++ b/zig/src/families/nemotron_parse/host_test.zig @@ -0,0 +1,255 @@ +//! Device-free Parse admission and processor checks; the optional checkpoint must already exist. +const std = @import("std"); +const core = @import("core"); +const family = @import("root.zig"); +const c = family.config; +const schema = family.schema; +const processor = @import("processor.zig"); +const a = std.testing.allocator; +const io = std.testing.io; + +const valid = + \\{"model_type":"nemotron_parse","legacy_causal_vision_prefix":true,"tie_word_embeddings":true, + \\"navit_position_embedding_mode":"content_normalized_2d_interpolate","navit_coordinate_frame":"content_normalized", + \\"navit_packing":true,"max_sequence_length":16384,"num_extra_heads":1,"vision_spatial_dim":1152, + \\"vision_summary_dim":2304,"decoder_start_token_id":2,"eos_token_id":2,"pad_token_id":1, + \\"decoder":{"model_type":"nemotron_parse_text","activation_function":"gelu","scale_embedding":true, + \\"tie_word_embeddings":true,"causal_vision_prefix":true,"add_cross_attention":false, + \\"d_model":1024,"decoder_layers":10,"decoder_ffn_dim":4096,"vocab_size":72512, + \\"decoder_attention_heads":16,"self_attention_kv_heads":8,"cross_attention_kv_heads":0}, + \\"encoder":{"version":"c-radio_v4-so400m","patch_size":16,"max_resolution":2048, + \\"args":{"model":"vit_so400m_patch16_224","cpe_max_size":2048,"register_multiple":10}}} +; + +test "Parse config accepts its graph and refuses wrong arithmetic and conflicting aliases" { + try c.validate(a, valid); + for ([_]struct { []const u8, []const u8 }{ + .{ "\"d_model\":1024", "\"d_model\":2048" }, + .{ "\"self_attention_kv_heads\":8", "\"self_attention_kv_heads\":16" }, + .{ "\"gelu\"", "\"gelu_new\"" }, + .{ "\"scale_embedding\":true", "\"scale_embedding\":false" }, + .{ "\"max_sequence_length\":16384", "\"max_sequence_length\":32768" }, + .{ "\"num_extra_heads\":1", "\"num_extra_heads\":0" }, + .{ "\"model_type\":\"nemotron_parse\"", "\"model_type\":\"nemotron_h\"" }, + .{ "\"pad_token_id\":1", "\"pad_token_id\":1,\"hidden_size\":2048" }, + .{ "\"pad_token_id\":1", "\"pad_token_id\":1,\"quantization\":{}" }, + }) |change| { + const bad = try replace(valid, change[0], change[1]); + defer a.free(bad); + try std.testing.expectError(error.UnsupportedParseConfig, c.validate(a, bad)); + } + try std.testing.expectError(error.BadParseConfig, c.validate(a, "[]")); +} + +fn replace(text: []const u8, old: []const u8, new: []const u8) ![]u8 { + const at = std.mem.indexOf(u8, text, old) orelse return error.MissingTestField; + return std.mem.concat(a, u8, &.{ text[0..at], new, text[at + old.len ..] }); +} + +const Header = struct { + arena: std.heap.ArenaAllocator, + names: core.safetensors.Header = .empty, + + pub fn tensor(h: *Header, name: []const u8, dtype: core.safetensors.DType, shape: []const usize) !void { + const ar = h.arena.allocator(); + var dims: [core.safetensors.max_rank]usize = @splat(1); + @memcpy(dims[0..shape.len], shape); + try h.names.put(ar, try ar.dupe(u8, name), .{ .dtype = dtype, .rank = @intCast(shape.len), .shape = dims, .begin = 0, .end = 0 }); + } +}; + +test "Parse tensor admission rejects missing, wrong-shape, mixed-dtype and extra tensors" { + var fixture: Header = .{ .arena = .init(a) }; + defer fixture.arena.deinit(); + try schema.each(&fixture); + try schema.validate(&fixture.names); + try std.testing.expectEqual(@as(usize, 510), fixture.names.count()); + const name = "encoder.conv2.weight"; + const original = fixture.names.get(name).?; + fixture.names.getPtr(name).?.shape[3] = 2; // a 2-by-2 shuffle is a different model + try std.testing.expectError(error.BadParseTensor, schema.validate(&fixture.names)); + fixture.names.getPtr(name).?.* = original; + fixture.names.getPtr(name).?.dtype = .bf16; + try std.testing.expectError(error.BadParseTensor, schema.validate(&fixture.names)); + fixture.names.getPtr(name).?.* = original; + try fixture.tensor("decoder.embed_positions.weight", .f32, &.{ 16384, 1024 }); + try std.testing.expectError(error.UnexpectedParseTensor, schema.validate(&fixture.names)); + _ = fixture.names.swapRemove("decoder.embed_positions.weight"); + _ = fixture.names.swapRemove(name); + try std.testing.expectError(error.MissingParseTensor, schema.validate(&fixture.names)); +} + +test "Parse real checkpoint validates all 510 tensors without GPU initialization" { + const dir = @import("parse_options").model_dir orelse return error.SkipZigTest; + var ck = try family.checkpoint.Checkpoint.open(a, io, dir); + defer ck.close(); + try std.testing.expectEqual(@as(usize, 510), ck.file.names.count()); + const embed = try ck.get("decoder.embed_tokens.weight"); + try std.testing.expect(embed.is(.f32, &.{ c.vocab, c.hidden })); + try std.testing.expectEqual(@as(usize, 4 * c.vocab * c.hidden), embed.bytes.len); +} + +test "Parse geometry keeps small pages and counts prefix slots and context edges" { + const small = try processor.geometry(17, 65); + try std.testing.expectEqualDeep(processor.Size{ .h = 17, .w = 65 }, small.content); + try std.testing.expectEqualDeep(processor.Size{ .h = 32, .w = 128 }, small.padded); + try std.testing.expectEqual(@as(usize, 5), small.tokens()); + const large = try processor.geometry(4096, 3328); + try std.testing.expectEqualDeep(processor.Size{ .h = 2048, .w = 1664 }, large.content); + try std.testing.expectEqual(@as(usize, 3329), large.tokens()); + try large.fits(c.context - 3329, c.context); + try std.testing.expectError(error.ContextExceeded, large.fits(c.context - 3328, c.context)); + try std.testing.expectError(error.ContextExceeded, large.fits(1, 3329)); + try std.testing.expectError(error.Invalid, processor.geometry(0, 7)); + try std.testing.expectError(error.ImageTooSmall, processor.geometry(8192, 1)); +} + +test "Parse patches preserve channel-row-column order and white padding" { + var pixels: [17 * 65 * 3]u8 = undefined; + for (&pixels, 0..) |*v, i| v.* = @truncate(i * 37); + var p = try processor.prepareRgb(a, .{ .height = 17, .width = 65, .pixels = &pixels }); + defer p.deinit(a); + try std.testing.expectEqual(@as(usize, 16 * 768), p.pixels.len); + // Each channel has its own contiguous 16x16 patch, not interleaved RGB pixels. + try std.testing.expectEqual((@as(f32, pixels[3]) / 255 - c.mean[0]) / c.deviation[0], p.pixels[1]); + try std.testing.expectEqual((@as(f32, pixels[1]) / 255 - c.mean[1]) / c.deviation[1], p.pixels[256]); + const pad = 7 * 768 + 256 + 15; // rightmost patch, row 0, channel 1, padded column + try std.testing.expectEqual((1 - c.mean[1]) / c.deviation[1], p.pixels[pad]); + try std.testing.expectError(error.BadPixels, processor.prepareRgb(a, .{ .height = 17, .width = 65, .pixels = pixels[0..3] })); +} + +test "Parse processor matches independent Pillow and NumPy fingerprints" { + const Case = struct { height: usize, width: usize, content_h: usize, content_w: usize, padded_h: usize, padded_w: usize, tokens: usize, sha256: []const u8 }; + const Fixture = struct { source_revision: []const u8, cases: []const Case }; + // Generated by `python tools/zig/nemotron_parse_reference.py processor`. + const doc = try std.json.parseFromSlice(Fixture, a, @embedFile("fixtures/processor.json"), .{}); + defer doc.deinit(); + for (doc.value.cases) |cell| { + const pixels = try a.alloc(u8, cell.height * cell.width * 3); + defer a.free(pixels); + for (0..cell.height) |y| for (0..cell.width) |x| for (0..3) |ch| { + pixels[(y * cell.width + x) * 3 + ch] = @truncate(y * 17 + x * 29 + ch * 83); + }; + var p = try processor.prepareRgb(a, .{ .height = cell.height, .width = cell.width, .pixels = pixels }); + defer p.deinit(a); + try std.testing.expectEqualDeep(processor.Size{ .h = cell.content_h, .w = cell.content_w }, p.geometry.content); + try std.testing.expectEqualDeep(processor.Size{ .h = cell.padded_h, .w = cell.padded_w }, p.geometry.padded); + try std.testing.expectEqual(cell.tokens, p.geometry.tokens()); + var hash: std.crypto.hash.sha2.Sha256 = .init(.{}); + for (p.pixels) |v| { + var bytes: [4]u8 = undefined; + std.mem.writeInt(u32, &bytes, @bitCast(v), .little); + hash.update(&bytes); + } + var want_bytes: [32]u8 = undefined; + const want = try std.fmt.hexToBytes(&want_bytes, cell.sha256); + try std.testing.expectEqualSlices(u8, want, &hash.finalResult()); + } +} + +test "Parse processor config rejects an effective canvas that conflicts with size metadata" { + const text = + \\{"navit_packing":true,"do_normalize":true,"do_rescale":true, + \\"navit_padding_mode":"top_left_white","navit_coordinate_frame":"content_normalized", + \\"size":{"height":2048,"width":1664},"final_size":[2048,1664], + \\"image_mean":[0.48145466,0.4578275,0.40821073],"image_std":[0.26862954,0.26130258,0.27577711], + \\"rescale_factor":0.00392156862745098} + ; + try c.validateProcessor(a, text); + for ([_]struct { []const u8, []const u8 }{ + .{ "[2048,1664]", "[512,512]" }, + .{ "top_left_white", "center_white" }, + .{ "0.48145466", "0.5" }, + .{ "0.00392156862745098", "1.0" }, + }) |change| { + const bad = try replace(text, change[0], change[1]); + defer a.free(bad); + try std.testing.expectError(error.UnsupportedParseConfig, c.validateProcessor(a, bad)); + } +} + +test "Parse image identity includes content dimensions when normalized padded pixels are equal" { + var first: [17 * 65 * 3]u8 = @splat(255); + var second: [18 * 66 * 3]u8 = @splat(255); + var p = try processor.prepareRgb(a, .{ .height = 17, .width = 65, .pixels = &first }); + defer p.deinit(a); + var q = try processor.prepareRgb(a, .{ .height = 18, .width = 66, .pixels = &second }); + defer q.deinit(a); + try std.testing.expectEqualSlices(f32, p.pixels, q.pixels); + try std.testing.expect(p.hash != q.hash); + try std.testing.expectError(error.Oversized, processor.geometry(std.math.maxInt(usize), 7)); +} + +test "Parse decoded image matches the actual Hugging Face example processor output" { + const dir = @import("parse_options").reference_dir orelse return error.SkipZigTest; + const image = try readReference(dir, "input.image", 16 << 20); + defer a.free(image); + const expected = try readReference(dir, "input.pixel_values.bin", 64 << 20); + defer a.free(expected); + var p = try processor.prepare(io, a, image); + defer p.deinit(a); + try std.testing.expectEqual(p.pixels.len * 4, expected.len); + for (p.pixels, 0..) |v, i| { + const bits = std.mem.readInt(u32, expected[i * 4 ..][0..4], .little); + try std.testing.expectEqual(bits, @as(u32, @bitCast(v))); + } + const sizes = try readReference(dir, "input.content_sizes.bin", 17); + defer a.free(sizes); + try std.testing.expectEqual(@as(usize, 16), sizes.len); + try std.testing.expectEqual(p.geometry.content.h, std.mem.readInt(u64, sizes[0..8], .little)); + try std.testing.expectEqual(p.geometry.content.w, std.mem.readInt(u64, sizes[8..16], .little)); +} + +fn readReference(dir: []const u8, name: []const u8, limit: usize) ![]u8 { + const path = try std.fs.path.join(a, &.{ dir, name }); + defer a.free(path); + return std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(limit)); +} + +test "Parse media admission preserves unpadded content and refuses malformed visual prefixes" { + const Image = struct { pixels: []const f32, gh: u32, gw: u32, at: u32, tokens: u32, content_size: ?[2]u32 }; + const pixels = try a.alloc(f32, 2 * 8 * 768); + defer a.free(pixels); + var image: Image = .{ .pixels = pixels, .gh = 2, .gw = 8, .at = 0, .tokens = 5, .content_size = .{ 17, 65 } }; + const media = @import("media.zig"); + const valid_geometry = try media.geometry(image); + try std.testing.expectEqualDeep(processor.Size{ .h = 17, .w = 65 }, valid_geometry.content); + image.at = 1; + try std.testing.expectError(error.ParseImageMustBePrefix, media.geometry(image)); + image.at = 0; + image.content_size = null; + try std.testing.expectError(error.ParseContentSizeRequired, media.geometry(image)); + for ([_][2]u32{ .{ 0, 65 }, .{ 33, 65 }, .{ 17, 129 }, .{ 17, 64 } }) |bad| { + image.content_size = bad; + try std.testing.expectError(error.BadParseImage, media.geometry(image)); + } + image.content_size = .{ 17, 65 }; + image.tokens = 4; + try std.testing.expectError(error.BadParseImage, media.geometry(image)); + image.tokens = 5; + image.pixels = pixels[0..3]; + try std.testing.expectError(error.BadParseImage, media.geometry(image)); + image.pixels = pixels; + image.gw = 7; + try std.testing.expectError(error.BadParseImage, media.geometry(image)); +} + +test "Parse native tokenizer matches HF task IDs and full serialized OCR output" { + const model = @import("parse_options").model_dir orelse return error.SkipZigTest; + const reference = @import("parse_options").reference_dir orelse return error.SkipZigTest; + const text = try readReference(reference, "manifest.json", 4 << 20); + defer a.free(text); + const Manifest = struct { prompt_ids: []u32, output_ids: []u32, text: []const u8 }; + const doc = try std.json.parseFromSlice(Manifest, a, text, .{ .ignore_unknown_fields = true }); + defer doc.deinit(); + var tokenizer = try core.tokenizer.loadTokenizer(io, a, model); + defer tokenizer.deinit(); + const prompt = ""; + const ids = try tokenizer.encode(a, prompt); + defer a.free(ids); + try std.testing.expectEqualSlices(u32, doc.value.prompt_ids, ids); + const decoded = try tokenizer.decode(a, doc.value.output_ids, false); + defer a.free(decoded); + try std.testing.expectEqualStrings(doc.value.text, decoded); + try std.testing.expect(tokenizer.isAdded(1)); // the visual prefix uses the checkpoint's added marker +} diff --git a/zig/src/families/nemotron_parse/media.zig b/zig/src/families/nemotron_parse/media.zig new file mode 100644 index 0000000000..15ef17b7f8 --- /dev/null +++ b/zig/src/families/nemotron_parse/media.zig @@ -0,0 +1,16 @@ +//! Validate a server image's content geometry before allocating or launching the Parse tower. +const c = @import("config.zig"); +const processor = @import("processor.zig"); + +pub fn geometry(image: anytype) !processor.Geometry { + if (image.at != 0) return error.ParseImageMustBePrefix; + const content = image.content_size orelse return error.ParseContentSizeRequired; + if (image.gh == 0 or image.gh > c.max_height / c.patch or image.gw == 0 or + image.gw > c.max_width / c.patch or image.gw % 4 != 0) return error.BadParseImage; + if (content[0] == 0 or content[1] == 0 or content[0] > image.gh * 16 or content[1] > image.gw * 16) + return error.BadParseImage; + const g: processor.Geometry = .{ .content = .{ .h = content[0], .w = content[1] }, .padded = .{ .h = @as(usize, image.gh) * 16, .w = @as(usize, image.gw) * 16 } }; + if ((g.content.h + 15) / 16 != image.gh or (g.content.w + 63) / 64 * 4 != image.gw or + g.tokens() != image.tokens or g.patches() * 768 != image.pixels.len) return error.BadParseImage; + return g; +} diff --git a/zig/src/families/nemotron_parse/model_check.zig b/zig/src/families/nemotron_parse/model_check.zig new file mode 100644 index 0000000000..ff81d67221 --- /dev/null +++ b/zig/src/families/nemotron_parse/model_check.zig @@ -0,0 +1,99 @@ +//! Development-only byte checks of actual model state at prompt cuts, shared forwards and suffix rollback. +const std = @import("std"); +const hip = @import("hip"); +const Compute = @import("compute.zig").Compute; +const decoder = @import("decoder.zig"); +const config = @import("config.zig"); +const Tensor = hip.ops.Tensor; + +pub fn run(a: std.mem.Allocator, c: Compute, r: *const hip.Runtime, prefix: Tensor, image_rows: usize, ids: []const i32, baseline: *decoder.State) !void { + const rows = image_rows + ids.len; + if (ids.len == 0) return error.BadReferenceIds; + var text = try hip.DeviceBuffer.fromHost(r, std.mem.sliceAsBytes(ids)); + defer text.free(); + var input = try hip.DeviceBuffer.alloc(r, rows * config.hidden * 2); + defer input.free(); + const tensor: Tensor = .{ .ptr = input.base(), .kind = .bf16 }; + try c.ops.copyCols(prefix, config.hidden, 0, tensor.ptr, image_rows, config.hidden); + try c.kernels.embed(c.stream, try c.weights.get("decoder.embed_tokens.weight"), text.base(), slice(tensor, image_rows), ids.len, config.hidden, 32); + try c.stream.synchronize(); + var cut = try decoder.State.init(r, rows + 4); + defer cut.deinit(); + var other = try decoder.State.init(r, rows + 4); + defer other.deinit(); + var shared = try decoder.State.init(r, rows + 7); + defer shared.deinit(); + var chunk: usize = 17; + while (cut.len < rows) { + c.ops.arena.reset(); + const n = @min(chunk, rows - cut.len); + _ = try decoder.forward(c, &cut, slice(tensor, cut.len), n); + chunk = if (chunk == 17) 511 else 256; + } + try c.stream.synchronize(); + try same(a, baseline, &cut, rows); + std.debug.print("STATE: full/chunked prompt KV and hidden bytes agree ({d} rows)\n", .{rows}); + + // A different final text token makes cross-stream KV contamination visible. + const different: i32 = if (ids[ids.len - 1] == 243) 809 else 243; + try text.upload(0, std.mem.asBytes(&different)); + var combined = try hip.DeviceBuffer.alloc(r, 2 * rows * config.hidden * 2); + defer combined.free(); + try hip.raw.copy(combined, 0, tensor.ptr, rows * config.hidden * 2, c.stream.handle); + try hip.raw.copy(combined, rows * config.hidden * 2, tensor.ptr, rows * config.hidden * 2, c.stream.handle); + const combined_tensor: Tensor = .{ .ptr = combined.base(), .kind = .bf16 }; + try c.kernels.embed(c.stream, try c.weights.get("decoder.embed_tokens.weight"), text.base(), slice(combined_tensor, 2 * rows - 1), 1, config.hidden, 32); + c.ops.arena.reset(); + _ = try decoder.forward(c, &other, slice(combined_tensor, rows), rows); + try cut.keep(0); + c.ops.arena.reset(); + _ = try decoder.forwardMany(c, &.{ .{ .state = &cut, .rows = rows }, .{ .state = &shared, .rows = rows } }, combined_tensor); + try c.stream.synchronize(); + try same(a, baseline, &cut, rows); + try same(a, &other, &shared, rows); + std.debug.print("STATE: two distinct prompts share one dense forward and match solo KV/hidden bytes\n", .{}); + + // Roll back inside a target window, then replay only its accepted suffix. + try shared.keep(rows - 3); + c.ops.arena.reset(); + _ = try decoder.forward(c, &shared, slice(combined_tensor, 2 * rows - 3), 3); + try c.stream.synchronize(); + try same(a, &other, &shared, rows); + std.debug.print("STATE: partial suffix rollback/replay preserves every KV and hidden byte\n", .{}); + c.ops.arena.reset(); +} + +fn slice(x: Tensor, start: usize) Tensor { + return .{ .ptr = x.ptr + start * config.hidden * 2, .kind = .bf16 }; +} + +fn bytes(a: std.mem.Allocator, x: hip.DeviceBuffer, xo: usize, y: hip.DeviceBuffer, yo: usize, n: usize) !void { + const left = try a.alloc(u8, n); + defer a.free(left); + const right = try a.alloc(u8, n); + defer a.free(right); + try x.download(xo, left); + try y.download(yo, right); + if (!std.mem.eql(u8, left, right)) { + var different: usize = 0; + for (left, right) |l, v| different += @intFromBool(l != v); + std.debug.print("STATE mismatch: {d}/{d} bytes, offsets {d}/{d}\n", .{ different, n, xo, yo }); + return error.ModelStateMismatch; + } +} + +fn same(a: std.mem.Allocator, x: *const decoder.State, y: *const decoder.State, rows: usize) !void { + if (x.len < rows or y.len < rows) return error.BadComparedState; + try bytes(a, x.hidden, 0, y.hidden, 0, rows * config.hidden * 2); + for (x.layers, y.layers, 0..) |xl, yl, layer| { + for (0..config.kv_heads) |head| { + const xo = head * x.total * config.head_dim * 2; + const yo = head * y.total * config.head_dim * 2; + bytes(a, xl.k, xo, yl.k, yo, rows * config.head_dim * 2) catch |err| { + std.debug.print("STATE: K layer {d}, head {d}\n", .{ layer, head }); + return err; + }; + try bytes(a, xl.v, xo, yl.v, yo, rows * config.head_dim * 2); + } + } +} diff --git a/zig/src/families/nemotron_parse/mtp.zig b/zig/src/families/nemotron_parse/mtp.zig new file mode 100644 index 0000000000..7d27d20d09 --- /dev/null +++ b/zig/src/families/nemotron_parse/mtp.zig @@ -0,0 +1,15 @@ +//! The checkpoint's trained one-step proposal: unscaled token embedding, biased projection, residual and head. +const hip = @import("hip"); +const c = @import("config.zig"); +const Compute = @import("compute.zig").Compute; + +/// Pack one state row and one pending token per chain; the target alone verifies the returned proposal logits. +pub fn logits(compute: Compute, hidden: hip.ops.Tensor, ids: u64, rows: usize) !hip.ops.Tensor { + if (rows == 0) return error.EmptyProposal; + const embedding = try compute.alloc(rows * c.hidden, .bf16); + try compute.kernels.embed(compute.stream, try compute.weights.get("decoder.embed_tokens.weight"), ids, embedding, rows, c.hidden, 1); + const projected = try compute.linear(embedding, "decoder.extra_proj.0", rows, c.hidden, c.hidden, true); + try compute.residual(projected, hidden, rows * c.hidden); + const out = try compute.linear(projected, "decoder.extra_heads.0", rows, c.hidden, c.hidden, true); + return @import("decoder.zig").logits(compute, out, rows); +} diff --git a/zig/src/families/nemotron_parse/prefill.zig b/zig/src/families/nemotron_parse/prefill.zig new file mode 100644 index 0000000000..59fd7b390e --- /dev/null +++ b/zig/src/families/nemotron_parse/prefill.zig @@ -0,0 +1,106 @@ +//! Native image-prefix preparation and cancellable prompt chunks driven by the shared lane host. +const std = @import("std"); +const hip = @import("hip"); +const lanes = @import("lanes"); +const c = @import("config.zig"); +const backend = @import("backend.zig"); +const decoder = @import("decoder.zig"); +const Compute = @import("compute.zig").Compute; +const Tensor = hip.ops.Tensor; + +fn cancelled(ptr: *anyopaque) bool { + const s: *lanes.Stream = @ptrCast(@alignCast(ptr)); + return s.isCancelled(); +} + +fn compute(h: *backend.Backend, s: *lanes.Stream) Compute { + var out = h.e.compute(); + out.cancel = .{ .ptr = s, .check = cancelled }; + return out; +} + +fn begin(h: *backend.Backend, s: *lanes.Stream) !*backend.Lane { + if (s.images.len != 1) return error.ParseNeedsOneImage; + if (h.streams.map.count() >= h.e.streams) return error.LaneMemory; + const image = s.images[0]; + const g = try @import("media.zig").geometry(image); + if (image.tokens >= s.prompt_len or s.prompt_len > h.e.window) return error.ContextExceeded; + for (s.prompt()) |id| if (id >= c.vocab) return error.BadInputToken; + try h.e.laneRoom(); + const l = try h.a.create(backend.Lane); + errdefer h.a.destroy(l); + l.* = .{ .state = try decoder.State.init(&h.e.r, h.e.window), .prefix = undefined, .image_rows = image.tokens }; + errdefer l.state.deinit(); + l.prefix = try hip.DeviceBuffer.alloc(&h.e.r, image.tokens * c.hidden * 2); + errdefer l.prefix.free(); + try h.streams.map.put(h.a, s, l); + errdefer _ = h.streams.map.remove(s); + errdefer h.e.stream.synchronize() catch {}; + if (s.reuse.saved) |saved| { + try @import("cache.zig").restore(h, s, saved); + if (l.fill != s.reuse.at) return error.BadSnapshot; + s.cached = s.reuse.at; + if (l.fill >= l.image_rows) return l; + } + h.e.scratch.reset(); + const comp = compute(h, s); + var raw = try hip.DeviceBuffer.fromHost(&h.e.r, std.mem.sliceAsBytes(image.pixels)); + defer raw.free(); + const input = try comp.alloc(image.pixels.len, .bf16); + try comp.ops.cast(.{ .ptr = raw.base(), .kind = .f32 }, input, image.pixels.len); + try @import("tower.zig").encode(comp, input, g, .{ .ptr = l.prefix.base(), .kind = .bf16 }); + try h.e.stream.synchronize(); + return l; +} + +pub fn step(h: *backend.Backend, s: *lanes.Stream, max_rows: u32) !bool { + if (s.isCancelled()) return error.Cancelled; + const l = h.streams.map.get(s) orelse try begin(h, s); + const from = l.fill; + if (from == s.prompt_len) { + h.e.scratch.reset(); + try finish(h, s, l, compute(h, s)); + s.fill_at.store(@intCast(from), .release); + return true; + } + var to = @min(s.prompt_len, from + @min(if (max_rows == 0) 512 else max_rows, 512)); + for (s.reuse.marks) |at| if (at > from and at < to) { + to = at; + }; + for (s.chunks) |at| if (at > from and at < to) { + to = at; + }; + if (to <= from) return error.EmptyPromptChunk; + h.e.scratch.reset(); + const comp = compute(h, s); + var ids = try hip.DeviceBuffer.fromHost(&h.e.r, std.mem.sliceAsBytes(s.prompt()[from..to])); + defer ids.free(); + const embedded = try decoder.embed(comp, ids.base(), to - from); + if (from < l.image_rows) { + const image_end = @min(to, l.image_rows); + const src: Tensor = .{ .ptr = l.prefix.base() + from * c.hidden * 2, .kind = .bf16 }; + try comp.ops.copyCols(src, c.hidden, 0, embedded.ptr, image_end - from, c.hidden); + } + const hidden = try decoder.forward(comp, &l.state, embedded, to - from); + _ = hidden; + if (to == s.prompt_len) try finish(h, s, l, comp); + try h.e.stream.synchronize(); + l.fill = to; + s.fill_at.store(@intCast(to), .release); + if (s.reuse.hook) |hook| for (s.reuse.marks) |at| if (at == to) { + hook.at(hook.ptr, s, @intCast(to)); + break; + }; + return to == s.prompt_len; +} + +fn finish(h: *backend.Backend, s: *lanes.Stream, l: *backend.Lane, comp: Compute) !void { + const logits = try decoder.logits(comp, try l.state.hiddenAt(s.prompt_len - 1), 1); + var r = try h.request(h.a, s, l.image_rows, s.prompt_len, &.{}, s.prompt_len); + defer h.a.free(r.history); + if (s.grammar != null and l.first_row == null) l.first_row = try h.a.alloc(u16, c.vocab); + r.keep = l.first_row; + var choice: [1]u32 = undefined; + try h.sampler.draw(comp.ops, h.e.stream, logits, &.{r}, &choice, false); + l.first = h.streams.drawn.take(choice[0]); +} diff --git a/zig/src/families/nemotron_parse/probe.zig b/zig/src/families/nemotron_parse/probe.zig new file mode 100644 index 0000000000..e3d26c9ae5 --- /dev/null +++ b/zig/src/families/nemotron_parse/probe.zig @@ -0,0 +1,161 @@ +//! Developer probe of Parse model seams against the pinned HF example; metrics are not a qualification gate. +const std = @import("std"); +const hip = @import("hip"); +const core = @import("core"); +const processor = @import("processor.zig"); +const Weights = @import("weights.zig").Weights; +const Compute = @import("compute.zig").Compute; + +pub fn main(init: std.process.Init) !void { + const args = try init.minimal.args.toSlice(init.arena.allocator()); + if (args.len != 4) { + std.debug.print("usage: tf-nemotron-parse-test MODEL REFERENCE OUTPUT\n", .{}); + return error.InvalidArguments; + } + const a = init.gpa; + const io = init.io; + var r = try hip.Runtime.open(); + defer r.close(); + var ctx = try hip.Context.init(&r, 0); + defer ctx.deinit(); + const room = try ctx.memInfo(); + if (room.free < 5 * (1 << 30)) return error.InsufficientProbeMemory; + std.debug.print("PROBE: free VRAM {d} MiB, allocations below 3 GiB\n", .{room.free >> 20}); + var s = try hip.Stream.init(&r); + defer s.deinit(); + var l = try hip.Launcher.load(&r, .{}, hip.kernels.images); + defer l.unload(); + l.wide = false; // use the same 16-row attention arithmetic for prefill, decode and verification + const kernels = try hip.vision.Kernels.load(l.mods[0]); + var w = try Weights.load(a, &r, io, args[1]); + defer w.deinit(); + const image_path = try std.fs.path.join(a, &.{ args[2], "input.image" }); + defer a.free(image_path); + const image = try std.Io.Dir.cwd().readFileAlloc(io, image_path, a, .limited(16 << 20)); + defer a.free(image); + var pixels = try processor.prepare(io, a, image); + defer pixels.deinit(a); + var raw = try hip.DeviceBuffer.fromHost(&r, std.mem.sliceAsBytes(pixels.pixels)); + defer raw.free(); + var scratch = try hip.Arena.init(&r, 1 << 30); + defer scratch.deinit(); + var prefix = try hip.DeviceBuffer.alloc(&r, pixels.geometry.tokens() * 1024 * 2); + defer prefix.free(); + try std.Io.Dir.cwd().createDirPath(io, args[3]); + var tap: Probe = .{ .a = a, .io = io, .r = &r, .s = s, .reference = args[2], .output = args[3] }; + const c: Compute = .{ .ops = .{ .l = &l, .stream = s.handle, .bf16 = true, .arena = &scratch }, .kernels = kernels, .stream = s, .weights = &w, .tap = .{ .ctx = &tap, .record = Probe.record } }; + const input = try c.alloc(pixels.pixels.len, .bf16); + try c.ops.cast(.{ .ptr = raw.base(), .kind = .f32 }, input, pixels.pixels.len); + try @import("tower.zig").encode(c, input, pixels.geometry, .{ .ptr = prefix.base(), .kind = .bf16 }); + try s.synchronize(); + scratch.reset(); + var state = try @import("decoder.zig").State.init(&r, 16384); + defer state.deinit(); + const ids_path = try std.fs.path.join(a, &.{ args[2], "input.input_ids.bin" }); + defer a.free(ids_path); + const id_bytes = try std.Io.Dir.cwd().readFileAlloc(io, ids_path, a, .limited(1 << 20)); + defer a.free(id_bytes); + if (id_bytes.len == 0 or id_bytes.len % 8 != 0) return error.BadReferenceIds; + const ids = try a.alloc(i32, id_bytes.len / 8); + defer a.free(ids); + for (ids, 0..) |*id, i| { + const value = std.mem.readInt(i64, id_bytes[i * 8 ..][0..8], .little); + if (value < 0 or value >= 72512) return error.BadReferenceIds; + id.* = @intCast(value); + } + var dev_ids = try hip.DeviceBuffer.fromHost(&r, std.mem.sliceAsBytes(ids)); + defer dev_ids.free(); + const total = pixels.geometry.tokens() + ids.len; + const embeddings = try c.alloc(total * 1024, .bf16); + try c.ops.copyCols(.{ .ptr = prefix.base(), .kind = .bf16 }, 1024, 0, embeddings.ptr, pixels.geometry.tokens(), 1024); + try kernels.embed(s, try w.get("decoder.embed_tokens.weight"), dev_ids.base(), .{ .ptr = embeddings.ptr + pixels.geometry.tokens() * 1024 * 2, .kind = .bf16 }, ids.len, 1024, 32); + _ = try @import("decoder.zig").forward(c, &state, embeddings, total); + const logits = try @import("decoder.zig").logits(c, try state.hiddenAt(total - 1), 1); + try c.record("prefill.logits", logits, 1, 72512); + var pick = try hip.DeviceBuffer.alloc(&r, 4); + defer pick.free(); + try c.ops.argmaxRows(logits, 1, 72512, pick.base()); + try s.synchronize(); + var token: i32 = undefined; + try pick.download(0, std.mem.asBytes(&token)); + std.debug.print("PREFILL first token {d}\n", .{token}); + const manifest_path = try std.fs.path.join(a, &.{ args[2], "manifest.json" }); + defer a.free(manifest_path); + const manifest_bytes = try std.Io.Dir.cwd().readFileAlloc(io, manifest_path, a, .limited(4 << 20)); + defer a.free(manifest_bytes); + const manifest = try std.json.parseFromSlice(std.json.Value, a, manifest_bytes, .{}); + defer manifest.deinit(); + const generated = manifest.value.object.get("new_ids").?.array.items; + try Probe.record(&tap, "generation.logits.0", logits, 1, 72512); + var dc = c; + dc.tap = null; + for (1..generated.len) |step| { + var previous: i32 = @intCast(generated[step - 1].integer); + if (previous < 0 or previous >= 72512) return error.BadReferenceIds; + try dev_ids.upload(0, std.mem.asBytes(&previous)); + scratch.reset(); + const embedded = try @import("decoder.zig").embed(dc, dev_ids.base(), 1); + const hidden = try @import("decoder.zig").forward(dc, &state, embedded, 1); + const next_logits = try @import("decoder.zig").logits(dc, hidden, 1); + var name: [64]u8 = undefined; + try Probe.record(&tap, try std.fmt.bufPrint(&name, "generation.logits.{d}", .{step}), next_logits, 1, 72512); + } + std.debug.print("PROBE: {d} teacher-forced generation rows saved for HF logits-processor comparison\n", .{generated.len}); + + try @import("model_check.zig").run(a, dc, &r, .{ .ptr = prefix.base(), .kind = .bf16 }, pixels.geometry.tokens(), ids, &state); + + std.debug.print("PROBE: image tokens {d}, scratch peak {d} MiB; model-seam metrics only\n", .{ pixels.geometry.tokens(), scratch.peak >> 20 }); +} + +const Probe = struct { + a: std.mem.Allocator, + io: std.Io, + r: *const hip.Runtime, + s: hip.Stream, + reference: []const u8, + output: []const u8, + + fn record(ctx: *anyopaque, name: []const u8, x: hip.ops.Tensor, rows: usize, width: usize) !void { + const t: *Probe = @ptrCast(@alignCast(ctx)); + if (x.kind != .bf16) return error.UnexpectedProbeDtype; + try t.s.synchronize(); + const length = rows * width; + const got = try t.a.alloc(u16, length); + defer t.a.free(got); + try hip.runtime.check(t.r.api.hipMemcpyDtoH(std.mem.sliceAsBytes(got).ptr, @ptrFromInt(x.ptr), length * 2)); + const file = try std.fmt.allocPrint(t.a, "{s}.bin", .{name}); + defer t.a.free(file); + const path = try std.fs.path.join(t.a, &.{ t.reference, file }); + defer t.a.free(path); + const expected_all = try std.Io.Dir.cwd().readFileAlloc(t.io, path, t.a, .limited(128 << 20)); + defer t.a.free(expected_all); + const expected = if (std.mem.eql(u8, name, "prefill.logits") and expected_all.len >= length * 2) + expected_all[expected_all.len - length * 2 ..] + else + expected_all; + if (expected.len != length * 2) return error.ProbeShapeMismatch; + var equal: usize = 0; + var worst: f64 = 0; + var absolute: f64 = 0; + for (0..rows) |row| { + var square: f64 = 0; + var norm: f64 = 0; + for (0..width) |col| { + const i = row * width + col; + const bits = std.mem.readInt(u16, expected[i * 2 ..][0..2], .little); + const want: f64 = core.bf16_round.value(bits); + const value: f64 = core.bf16_round.value(got[i]); + if (!std.math.isFinite(value) or !std.math.isFinite(want)) return error.NonFiniteProbe; + equal += @intFromBool(got[i] == bits); + square += (value - want) * (value - want); + norm += want * want; + absolute = @max(absolute, @abs(value - want)); + } + worst = @max(worst, @sqrt(square / @max(norm, 1e-30))); + } + std.debug.print("{s}: {d}x{d}, equal {d}/{d}, max row relative L2 {d:.7}, max abs {d:.7}\n", .{ name, rows, width, equal, length, worst, absolute }); + const out_path = try std.fs.path.join(t.a, &.{ t.output, file }); + defer t.a.free(out_path); + try std.Io.Dir.cwd().writeFile(t.io, .{ .sub_path = out_path, .data = std.mem.sliceAsBytes(got) }); + } +}; diff --git a/zig/src/families/nemotron_parse/processor.zig b/zig/src/families/nemotron_parse/processor.zig new file mode 100644 index 0000000000..40637cb066 --- /dev/null +++ b/zig/src/families/nemotron_parse/processor.zig @@ -0,0 +1,100 @@ +//! Parse's aspect-preserving Pillow resize, NaViT white padding, and raw RGB patch order. +const std = @import("std"); +const picture = @import("core").picture; +const c = @import("config.zig"); +const Allocator = std.mem.Allocator; +pub const Error = picture.Error || error{ ImageTooSmall, BadPixels, ContextExceeded }; + +pub const Size = struct { h: usize, w: usize }; +pub const Geometry = struct { + content: Size, + padded: Size, + + pub fn patches(g: Geometry) usize { + return (g.padded.h / c.patch) * (g.padded.w / c.patch); + } + pub fn tokens(g: Geometry) usize { + return g.patches() / 4 + 1; // four horizontal patches per row, plus the summary + } + pub fn fits(g: Geometry, text: usize, window: usize) Error!void { + if (window > c.context or g.tokens() > window or text > window - g.tokens()) return error.ContextExceeded; + } +}; + +/// Follow the processor's two sequential float64 truncations, including small-image preservation. +pub fn geometry(height: usize, width: usize) Error!Geometry { + if (height == 0 or width == 0) return error.Invalid; + if (height > picture.max_side or width > picture.max_side) return error.Oversized; + try picture.fits(width, height); + const ratio = @as(f64, @floatFromInt(width)) / @as(f64, @floatFromInt(height)); + var h = height; + var w = width; + if (h > c.max_height) { + h = c.max_height; + w = @intFromFloat(@as(f64, @floatFromInt(h)) * ratio); + } + if (w > c.max_width) { + w = c.max_width; + h = @intFromFloat(@as(f64, @floatFromInt(w)) / ratio); + } + if (h == 0 or w == 0) return error.ImageTooSmall; + return .{ .content = .{ .h = h, .w = w }, .padded = .{ + .h = std.mem.alignForward(usize, h, 16), + .w = std.mem.alignForward(usize, w, 64), + } }; +} + +pub const Prepared = struct { + geometry: Geometry, + pixels: []f32, // [patches, 3, 16, 16], RGB after normalization + hash: u64, + + pub fn deinit(p: *Prepared, a: Allocator) void { + a.free(p.pixels); + p.* = undefined; + } +}; + +pub fn prepare(io: std.Io, a: Allocator, bytes: []const u8) Error!Prepared { + var rgb = try picture.decode(io, a, bytes); + defer rgb.deinit(a); + return prepareRgb(a, rgb); +} + +pub fn prepareRgb(a: Allocator, rgb: picture.Rgb) Error!Prepared { + const g = try geometry(rgb.height, rgb.width); + if (rgb.pixels.len != rgb.height * rgb.width * 3) return error.BadPixels; + const resized = try picture.resample.resizeWith(a, .pillow, .bilinear, .of(rgb.pixels, rgb.width, rgb.height), g.content.w, g.content.h); + defer a.free(resized); + const out = try a.alloc(f32, g.patches() * 768); + const table = normalization(); + const cols = g.padded.w / c.patch; + for (0..g.padded.h) |y| for (0..g.padded.w) |x| { + const base = ((y / 16) * cols + x / 16) * 768 + (y % 16) * 16 + x % 16; + for (0..3) |ch| { + const value = if (y < g.content.h and x < g.content.w) + resized[(y * g.content.w + x) * 3 + ch] + else + 255; + out[base + ch * 256] = table[ch][value]; + } + }; + // Content dimensions affect positional sampling even when the padded grid is equal. + var h: std.hash.Wyhash = .init(0x7061727365); + h.update(std.mem.sliceAsBytes(out)); + for ([_]usize{ g.content.h, g.content.w, g.padded.h, g.padded.w }) |v| { + var bytes: [8]u8 = undefined; + std.mem.writeInt(u64, &bytes, @intCast(v), .little); + h.update(&bytes); + } + return .{ .geometry = g, .pixels = out, .hash = h.final() }; +} + +fn normalization() [3][256]f32 { + var out: [3][256]f32 = undefined; + for (&out, 0..) |*channel, ch| for (channel, 0..) |*v, i| { + const rescaled = @as(f32, @floatFromInt(i)) / @as(f32, 255); + v.* = (rescaled - c.mean[ch]) / c.deviation[ch]; + }; + return out; +} diff --git a/zig/src/families/nemotron_parse/root.zig b/zig/src/families/nemotron_parse/root.zig new file mode 100644 index 0000000000..eca0cef044 --- /dev/null +++ b/zig/src/families/nemotron_parse/root.zig @@ -0,0 +1,6 @@ +//! Nemotron 3.5 Parse: a RADIO encoder and causal mBART visual-prefix decoder. +pub const config = @import("config.zig"); +pub const schema = @import("schema.zig"); +pub const checkpoint = @import("checkpoint.zig"); +pub const processor = @import("processor.zig"); +pub const model_type = config.model_type; diff --git a/zig/src/families/nemotron_parse/schema.zig b/zig/src/families/nemotron_parse/schema.zig new file mode 100644 index 0000000000..fdbfba3d9b --- /dev/null +++ b/zig/src/families/nemotron_parse/schema.zig @@ -0,0 +1,89 @@ +//! Every tensor required by the pinned RADIO/mBART graph, before device allocation. +const std = @import("std"); +const st = @import("core").safetensors; +const c = @import("config.zig"); +pub const radio = "encoder.model_encoder.radio_model."; +pub const tower = radio ++ "model."; +pub const patch_generator = tower ++ "patch_generator."; +pub const Error = error{ MissingParseTensor, BadParseTensor, UnexpectedParseTensor }; + +/// Used for both the real header and disposable malformed header fixtures. +pub fn validate(header: *const st.Header) Error!void { + var checker: Checker = .{ .header = header }; + each(&checker) catch |err| switch (err) { + error.NoSpaceLeft => unreachable, // fixed tensor stems fit the local name buffers + else => |e| return e, + }; + if (checker.count != header.count()) return error.UnexpectedParseTensor; +} + +const Checker = struct { + header: *const st.Header, + count: usize = 0, + fn tensor(s: *Checker, name: []const u8, dtype: st.DType, shape: []const usize) Error!void { + const t = s.header.get(name) orelse return error.MissingParseTensor; + if (t.dtype != dtype or t.rank != shape.len or !std.mem.eql(usize, t.shape[0..t.rank], shape)) + return error.BadParseTensor; + s.count += 1; + } +}; + +/// Emit the complete expected schema. Names are borrowed only for the callback. +pub fn each(sink: anytype) !void { + try sink.tensor("decoder.embed_tokens.weight", .f32, &.{ c.vocab, c.hidden }); + for ([_][]const u8{ "decoder.layer_norm", "decoder.layernorm_embedding" }) |stem| + try norm(sink, stem, c.hidden); + try sink.tensor(radio ++ "summary_idxs", .i64, &.{2}); + try sink.tensor(radio ++ "input_conditioner.norm_mean", .f32, &.{ 3, 1, 1 }); + try sink.tensor(radio ++ "input_conditioner.norm_std", .f32, &.{ 3, 1, 1 }); + try sink.tensor(patch_generator ++ "cls_token.token", .f32, &.{ c.vision_prefix, c.vision_hidden }); + try sink.tensor(patch_generator ++ "embedder.weight", .f32, &.{ c.vision_hidden, 768 }); + try sink.tensor(patch_generator ++ "pos_embed", .f32, &.{ 1, 128 * 128, c.vision_hidden }); + try linear(sink, "encoder.conv1", &.{ c.hidden, c.vision_hidden, 1 }, c.hidden); + try sink.tensor("encoder.conv2.weight", .f32, &.{ c.hidden, c.hidden, 1, 4 }); + try linear(sink, "encoder.sum_proj", &.{ c.hidden, c.vision_hidden * 2 }, c.hidden); + for ([_][]const u8{ "encoder.layer_norm1", "encoder.layer_norm2", "encoder.layer_norm3" }) |stem| + try norm(sink, stem, c.hidden); + var buf: [192]u8 = undefined; + for (0..c.layers) |i| { + const stem = try std.fmt.bufPrint(&buf, "decoder.layers.{d}", .{i}); + try subNorm(sink, stem, "self_attn_layer_norm", c.hidden); + try subNorm(sink, stem, "final_layer_norm", c.hidden); + try subLinear(sink, stem, "self_attn.q_proj", c.hidden, c.hidden); + try subLinear(sink, stem, "self_attn.k_proj", c.kv_heads * c.head_dim, c.hidden); + try subLinear(sink, stem, "self_attn.v_proj", c.kv_heads * c.head_dim, c.hidden); + try subLinear(sink, stem, "self_attn.out_proj", c.hidden, c.hidden); + try subLinear(sink, stem, "fc1", c.ffn, c.hidden); + try subLinear(sink, stem, "fc2", c.hidden, c.ffn); + } + try subLinear(sink, "decoder", "extra_heads.0", c.hidden, c.hidden); + try subLinear(sink, "decoder", "extra_proj.0", c.hidden, c.hidden); + for (0..c.vision_layers) |i| { + const stem = try std.fmt.bufPrint(&buf, tower ++ "blocks.{d}", .{i}); + try subNorm(sink, stem, "norm1", c.vision_hidden); + try subNorm(sink, stem, "norm2", c.vision_hidden); + try subLinear(sink, stem, "attn.qkv", 3 * c.vision_hidden, c.vision_hidden); + try subLinear(sink, stem, "attn.proj", c.vision_hidden, c.vision_hidden); + try subLinear(sink, stem, "mlp.fc1", c.vision_ffn, c.vision_hidden); + try subLinear(sink, stem, "mlp.fc2", c.vision_hidden, c.vision_ffn); + } +} + +fn norm(sink: anytype, stem: []const u8, n: usize) !void { + var buf: [256]u8 = undefined; + for ([_][]const u8{ "weight", "bias" }) |end| + try sink.tensor(try std.fmt.bufPrint(&buf, "{s}.{s}", .{ stem, end }), .f32, &.{n}); +} +fn linear(sink: anytype, stem: []const u8, shape: []const usize, n: usize) !void { + var buf: [256]u8 = undefined; + try sink.tensor(try std.fmt.bufPrint(&buf, "{s}.weight", .{stem}), .f32, shape); + try sink.tensor(try std.fmt.bufPrint(&buf, "{s}.bias", .{stem}), .f32, &.{n}); +} +fn subNorm(sink: anytype, stem: []const u8, sub: []const u8, n: usize) !void { + var buf: [224]u8 = undefined; + try norm(sink, try std.fmt.bufPrint(&buf, "{s}.{s}", .{ stem, sub }), n); +} +fn subLinear(sink: anytype, stem: []const u8, sub: []const u8, n: usize, k: usize) !void { + var buf: [224]u8 = undefined; + try linear(sink, try std.fmt.bufPrint(&buf, "{s}.{s}", .{ stem, sub }), &.{ n, k }, n); +} diff --git a/zig/src/families/nemotron_parse/tower.zig b/zig/src/families/nemotron_parse/tower.zig new file mode 100644 index 0000000000..58da70f0ba --- /dev/null +++ b/zig/src/families/nemotron_parse/tower.zig @@ -0,0 +1,90 @@ +//! C-RADIOv4-SO400M and Parse's horizontal convolution neck on native HIP kernels. +const std = @import("std"); +const hip = @import("hip"); +const config = @import("config.zig"); +const names = @import("schema.zig"); +const Geometry = @import("processor.zig").Geometry; +const Compute = @import("compute.zig").Compute; +const Tensor = hip.ops.Tensor; + +const width = config.vision_hidden; +const heads = config.vision_heads; +const head_width = width / heads; +comptime { + std.debug.assert(head_width * heads == width); +} + +/// "." in `buf`, valid until the next call; the helpers append ".weight" or ".bias". +fn sub(buf: []u8, stem: []const u8, name: []const u8) ![]const u8 { + return std.fmt.bufPrint(buf, "{s}.{s}", .{ stem, name }); +} + +/// One pre-norm ViT layer (arXiv 2010.11929, Eq. 2-3), unmasked, biased projections, erf GELU. +fn block(c: Compute, x: Tensor, stem: []const u8, rows: usize) !void { + var buf: [224]u8 = undefined; + + // Eq. 2: attention sub-layer. + const a = try c.norm(x, try sub(&buf, stem, "norm1"), rows, width, config.vision_eps); + const qkv = try c.linear(a, try sub(&buf, stem, "attn.qkv"), rows, 3 * width, width, true); + const mixed = try c.global(qkv, rows, heads, head_width); + const attended = try c.linear(mixed, try sub(&buf, stem, "attn.proj"), rows, width, width, true); + try c.residual(x, attended, rows * width); + + // Eq. 3: feed-forward sub-layer. + const b = try c.norm(x, try sub(&buf, stem, "norm2"), rows, width, config.vision_eps); + const up = try c.linear(b, try sub(&buf, stem, "mlp.fc1"), rows, config.vision_ffn, width, true); + try c.gelu(up, rows * config.vision_ffn); + const down = try c.linear(up, try sub(&buf, stem, "mlp.fc2"), rows, width, config.vision_ffn, true); + try c.residual(x, down, rows * width); +} + +/// `out` holds the image's visual prefix [geometry.tokens(),1024] and survives scratch reuse. +pub fn encode(c: Compute, patches: Tensor, geometry: Geometry, out: Tensor) !void { + const count = geometry.patches(); + const rows = count + config.vision_prefix; + const hidden = config.vision_hidden; + const x = try c.alloc(rows * hidden, .bf16); + const mark = c.ops.arena.mark(); + try positioned(c, patches, geometry, x); + try c.record("encoder.positioned", x, rows, hidden); + c.ops.arena.release(mark); + var buf: [192]u8 = undefined; + for (0..config.vision_layers) |i| { + try c.poll(); + const stem = try std.fmt.bufPrint(&buf, names.tower ++ "blocks.{d}", .{i}); + try block(c, x, stem, rows); + try c.record(try std.fmt.bufPrint(&buf, "encoder.block.{d}", .{i}), x, rows, hidden); + c.ops.arena.release(mark); + } + try neck(c, x, count, out); + try c.record("encoder.neck.0", out, geometry.tokens(), config.hidden); +} + +fn positioned(c: Compute, patches: Tensor, geometry: Geometry, x: Tensor) !void { + const h = (geometry.content.h + 15) / 16; + const w = (geometry.content.w + 15) / 16; + const projected = try c.linear(patches, names.patch_generator ++ "embedder", geometry.patches(), config.vision_hidden, 768, false); + const field = try c.alloc(h * w * config.vision_hidden, .f32); + const table = try c.weights.get(names.patch_generator ++ "pos_embed"); + if (@min(h, w) == 1) { + try c.kernels.resize(c.stream, table, field.ptr, .{ 128, 128 }, .{ h, w }, config.vision_hidden, .area); + } else { + const horizontal = try c.alloc(128 * w * config.vision_hidden, .f32); + try c.kernels.resize(c.stream, table, horizontal.ptr, .{ 128, 128 }, .{ 128, w }, config.vision_hidden, .horizontal); + try c.kernels.resize(c.stream, horizontal, field.ptr, .{ 128, w }, .{ h, w }, config.vision_hidden, .vertical); + } + try c.kernels.position(c.stream, projected, (try c.weights.get(names.patch_generator ++ "cls_token.token")).ptr, field.ptr, x, .{ geometry.padded.h / 16, geometry.padded.w / 16 }, .{ geometry.content.h, geometry.content.w }, config.vision_hidden, config.vision_prefix); +} + +fn neck(c: Compute, x: Tensor, patches: usize, out: Tensor) !void { + const features: Tensor = .{ .ptr = x.ptr + config.vision_prefix * config.vision_hidden * 2, .kind = .bf16 }; + const first = try c.linear(features, "encoder.conv1", patches, config.hidden, config.vision_hidden, true); + const norm1 = try c.norm(first, "encoder.layer_norm1", patches, config.hidden, config.vision_eps); + const grouped = try c.alloc(patches * config.hidden, .bf16); + try c.kernels.layout(c.stream, norm1, grouped, patches, 1, config.hidden, .horizontal_four); + const reduced = try c.linear(grouped, "encoder.conv2", patches / 4, config.hidden, config.hidden * 4, false); + try c.kernels.norm(c.stream, reduced, (try c.named("encoder.layer_norm2", "weight")).ptr, (try c.named("encoder.layer_norm2", "bias")).ptr, out, patches / 4, config.hidden, config.vision_eps); + const summary = try c.linear(x, "encoder.sum_proj", 1, config.hidden, config.vision_hidden * 2, true); + const last: Tensor = .{ .ptr = out.ptr + patches / 4 * config.hidden * 2, .kind = .bf16 }; + try c.kernels.norm(c.stream, summary, (try c.named("encoder.layer_norm3", "weight")).ptr, (try c.named("encoder.layer_norm3", "bias")).ptr, last, 1, config.hidden, config.vision_eps); +} diff --git a/zig/src/families/nemotron_parse/weights.zig b/zig/src/families/nemotron_parse/weights.zig new file mode 100644 index 0000000000..8812d5c739 --- /dev/null +++ b/zig/src/families/nemotron_parse/weights.zig @@ -0,0 +1,51 @@ +//! Owned BF16 device weights; the FP32 checkpoint stays unchanged and all tensors are checked before upload. +const std = @import("std"); +const hip = @import("hip"); +const core = @import("core"); +const checkpoint = @import("checkpoint.zig"); + +pub const Weights = struct { + gpa: std.mem.Allocator, + names: std.StringHashMapUnmanaged(hip.DeviceBuffer) = .empty, + bytes: usize = 0, + + pub fn load(a: std.mem.Allocator, r: *const hip.Runtime, io: std.Io, dir: []const u8) !Weights { + var ck = try checkpoint.Checkpoint.open(a, io, dir); + defer ck.close(); + var w: Weights = .{ .gpa = a }; + errdefer w.deinit(); + for (ck.file.names.keys()) |name| { + const tensor = try ck.get(name); + if (tensor.dtype == .i64) continue; // the checked summary indices are the fixed first two class rows + const count = tensor.numel(); + const host = try a.alloc(u16, count); + defer a.free(host); + for (host, 0..) |*out, i| { + const bits = std.mem.readInt(u32, tensor.bytes[i * 4 ..][0..4], .little); + out.* = core.bf16_round.bits(@bitCast(bits)); + } + const key = try a.dupe(u8, name); + errdefer a.free(key); + var dev = try hip.DeviceBuffer.fromHost(r, std.mem.sliceAsBytes(host)); + errdefer dev.free(); + try w.names.put(a, key, dev); + w.bytes += dev.len; + } + return w; + } + + pub fn deinit(w: *Weights) void { + var it = w.names.iterator(); + while (it.next()) |entry| { + entry.value_ptr.free(); + w.gpa.free(entry.key_ptr.*); + } + w.names.deinit(w.gpa); + w.* = undefined; + } + + pub fn get(w: *const Weights, name: []const u8) !hip.ops.Tensor { + const dev = w.names.get(name) orelse return error.MissingParseTensor; + return .{ .ptr = dev.base(), .kind = .bf16 }; + } +}; diff --git a/zig/src/hip/ops/vision.zig b/zig/src/hip/ops/vision.zig index b3a14411b8..6666ff2c05 100644 --- a/zig/src/hip/ops/vision.zig +++ b/zig/src/hip/ops/vision.zig @@ -1,9 +1,15 @@ -//! Shape and address admission for shared BF16 vision operations. +//! Shared vision operations: BF16 ops through the launcher (dots.mocr), dense kernels loaded by name (Nemotron Parse). const std = @import("std"); const Ops = @import("ops.zig").Ops; const t = @import("types.zig"); const raw = @import("../launches/vision.zig"); const count = @import("../launches/columns.zig").count; +const Module = @import("../module.zig").Module; +const Function = @import("../module.zig").Function; +const Stream = @import("../stream.zig").Stream; +const launch = @import("../launch.zig"); +const Tensor = t.Tensor; +pub const Error = t.Error; const Span = struct { at: u64, bytes: usize }; @@ -72,6 +78,133 @@ pub fn fullAttention(o: Ops, q: u64, k: u64, v: u64, out: u64, rows: usize, head try raw.full(o.l, q, k, v, out, rows, heads, d, 1 / @sqrt(@as(f32, @floatFromInt(d))), o.stream); } +pub const Kernels = struct { + embed_fn: Function, + linear_fn: Function, + norm_fn: Function, + point_fn: Function, + layout_fn: Function, + resize_fn: Function, + position_fn: Function, + + pub fn load(module: Module) !Kernels { + return .{ + .embed_fn = try module.function("tf_vl_embed"), + .linear_fn = try module.function("tf_vl_linear"), + .norm_fn = try module.function("tf_vl_layernorm"), + .point_fn = try module.function("tf_vl_pointwise"), + .layout_fn = try module.function("tf_vl_layout"), + .resize_fn = try module.function("tf_vl_resize"), + .position_fn = try module.function("tf_vl_position"), + }; + } + + /// Device IDs must be in the table's admitted range; embedding scaling rounds in the output dtype. + pub fn embed(k: Kernels, s: Stream, table: Tensor, ids: u64, out: Tensor, rows: usize, width: usize, scale: f32) Error!void { + if (table.kind != out.kind or rows == 0 or rows > 16384 or width == 0 or width > 16384 or !std.math.isFinite(scale)) return error.BadShape; + var args: launch.Args = .{}; + for ([_]u64{ table.ptr, ids, out.ptr }) |v| try args.add(v); + try args.add(@as(c_int, @intCast(rows))); + try args.add(@as(c_int, @intCast(width))); + try args.add(@backingInt(out.kind)); + try args.add(scale); + try flat(k.embed_fn, s, rows * width, &args); + } + + /// Weights and bias have x's dtype; the sole output rounding is after the biased FP32 dot. + pub fn linear(k: Kernels, s: Stream, x: Tensor, w: u64, bias: ?u64, out: Tensor, rows: usize, n: usize, width: usize) Error!void { + if (rows == 0 or n == 0 or width == 0 or rows > 16384 or n > 1 << 20 or width > 16384) return error.BadShape; + var args: launch.Args = .{}; + try args.add(x.ptr); + try args.add(w); + try args.add(bias orelse @as(u64, 0)); + try args.add(out.ptr); + for ([_]usize{ rows, n, width }) |v| try args.add(@as(c_int, @intCast(v))); + try args.add(@backingInt(x.kind)); + try args.add(@backingInt(out.kind)); + try launch.launch(k.linear_fn, .{ .grid = .{ .x = @intCast((n + 63) / 64), .y = @intCast((rows + 15) / 16) }, .block = .{ .x = 256 } }, s, &args); + } + + /// The affine parameters have x's dtype, as torch LayerNorm after model.to(bfloat16). + pub fn norm(k: Kernels, s: Stream, x: Tensor, weight: u64, bias: u64, out: Tensor, rows: usize, width: usize, eps: f32) Error!void { + if (x.kind != out.kind or rows == 0 or rows > 16384 or width == 0 or width > 16384 or !(eps > 0) or !std.math.isFinite(eps)) return error.BadShape; + var args: launch.Args = .{}; + for ([_]u64{ x.ptr, weight, bias, out.ptr }) |p| try args.add(p); + try args.add(@as(c_int, @intCast(rows))); + try args.add(@as(c_int, @intCast(width))); + try args.add(@backingInt(x.kind)); + try args.add(eps); + try launch.launch(k.norm_fn, .{ .grid = .{ .x = @intCast(rows) }, .block = .{ .x = 256 } }, s, &args); + } + + pub const Point = enum(c_int) { gelu, add, scale }; + pub fn point(k: Kernels, s: Stream, x: Tensor, y: ?Tensor, out: Tensor, elements: usize, mode: Point, scale: f32) Error!void { + if (x.kind != out.kind or elements == 0 or elements > std.math.maxInt(c_int)) return error.BadShape; + if (mode == .add and (y == null or y.?.kind != x.kind)) return error.BadShape; + var args: launch.Args = .{}; + try args.add(x.ptr); + try args.add(if (y) |v| v.ptr else @as(u64, 0)); + try args.add(out.ptr); + try args.add(@as(c_longlong, @intCast(elements))); + try args.add(@backingInt(x.kind)); + try args.add(@backingInt(mode)); + try args.add(scale); + try flat(k.point_fn, s, elements, &args); + } + + /// One axis of antialiased bilinear resizing, or adaptive-average pooling of both axes. + pub const Resize = enum(c_int) { horizontal, vertical, area }; + pub fn resize(k: Kernels, s: Stream, src: Tensor, out: u64, input: [2]usize, output: [2]usize, width: usize, mode: Resize) Error!void { + for ([_][2]usize{ input, output }) |pair| for (pair) |v| { + if (v == 0 or v > 128) return error.BadShape; + }; + if (width == 0 or width > 4096) return error.BadShape; + if ((mode == .horizontal and input[0] != output[0]) or + (mode == .vertical and input[1] != output[1])) return error.BadShape; + var args: launch.Args = .{}; + try args.add(src.ptr); + try args.add(out); + for ([_][2]usize{ input, output }) |pair| for (pair) |v| { + try args.add(@as(c_int, @intCast(v))); + }; + try args.add(@as(c_int, @intCast(width))); + try args.add(@backingInt(src.kind)); + try args.add(@backingInt(mode)); + try flat(k.resize_fn, s, output[0] * output[1] * width, &args); + } + + /// Prefix class rows and projected patches plus BF16-rounded, content-normalized positions. + pub fn position(k: Kernels, s: Stream, patches: Tensor, cls: u64, field: u64, out: Tensor, grid: [2]usize, content: [2]usize, hidden: usize, prefix: usize) Error!void { + if (patches.kind != out.kind or patches.kind == .f32 or hidden == 0 or hidden > 4096 or prefix > 32) return error.BadShape; + for (grid, content) |g, c| if (g == 0 or g > 128 or c == 0 or c > g * 16) return error.BadShape; + var args: launch.Args = .{}; + for ([_]u64{ patches.ptr, cls, field, out.ptr }) |v| try args.add(v); + for ([_][2]usize{ grid, content, .{ hidden, prefix } }) |pair| for (pair) |v| { + try args.add(@as(c_int, @intCast(v))); + }; + try args.add(@backingInt(patches.kind)); + try flat(k.position_fn, s, (grid[0] * grid[1] + prefix) * hidden, &args); + } + + pub const Layout = enum(c_int) { heads_to_rows, horizontal_four }; + pub fn layout(k: Kernels, s: Stream, x: Tensor, out: Tensor, rows: usize, heads: usize, dim: usize, mode: Layout) Error!void { + if (rows == 0 or rows > 16384 or heads == 0 or heads > 256 or dim == 0 or dim > 16384) return error.BadShape; + if (mode == .horizontal_four and (heads != 1 or rows % 4 != 0)) return error.BadShape; + var args: launch.Args = .{}; + try args.add(x.ptr); + try args.add(out.ptr); + for ([_]usize{ rows, heads, dim }) |v| try args.add(@as(c_int, @intCast(v))); + try args.add(@backingInt(x.kind)); + try args.add(@backingInt(out.kind)); + try args.add(@backingInt(mode)); + try flat(k.layout_fn, s, rows * heads * dim, &args); + } +}; + +fn flat(f: Function, s: Stream, elements: usize, args: *launch.Args) Error!void { + try launch.launch(f, .{ .grid = .{ .x = @intCast((elements + 255) / 256) }, .block = .{ .x = 256 } }, s, args); +} + test "dots HIP vision: malformed geometry and buffer spans fail before a launch" { var o: Ops = undefined; o.l = undefined; diff --git a/zig/src/hip_serve/draw.zig b/zig/src/hip_serve/draw.zig index d520e64889..9ad5e67767 100644 --- a/zig/src/hip_serve/draw.zig +++ b/zig/src/hip_serve/draw.zig @@ -16,11 +16,19 @@ pub const Request = struct { mask: []const u32 = &.{}, /// The row's logits copied here (the prompt's last row, kept for a masked first draw). keep: ?[]u16 = null, + /// The checkpoint's multiplicative penalty over prompt and reply tokens, applied once per seen token in FP32. + repetition: f32 = 1, + /// A checkpoint-forced final EOS, after its ordinary score processors. + forced: ?u32 = null, + history: []const u32 = &.{}, + /// OpenAI penalties count reply tokens only; history may also contain prompt tokens. + reply: []const u32 = &.{}, + penalties: lanes.sampling.Penalties = .{}, }; /// Whether a row is drawn on the host from the whole row. fn whole(r: Request) bool { - return r.mask.len > 0 or r.keep != null; + return r.mask.len > 0 or r.keep != null or r.repetition != 1 or r.penalties.active() or r.forced != null; } fn greedy(s: ?lanes.Sampling) bool { @@ -154,7 +162,54 @@ pub const Drawer = struct { /// A host row's token per `r`: its mask applied in place, then greedy or the keyed draw. pub fn drawHost(w: *const Drawer, row: []u16, r: Request) !u32 { + if (r.forced) |id| { + if (id >= row.len) return error.BadForcedToken; + if (r.mask.len > 0 and (id / 32 >= r.mask.len or (r.mask[id / 32] >> @intCast(id % 32)) & 1 == 0)) return error.NoAllowedToken; + return id; + } if (r.mask.len > 0) sample.mask(row, w.dtype, r.mask); - return sample.draw(w.gpa, row, w.dtype, r.sampling, r.position); + if (r.repetition == 1 and !r.penalties.active()) return sample.draw(w.gpa, row, w.dtype, r.sampling, r.position); + const values = try w.gpa.alloc(f32, row.len); + defer w.gpa.free(values); + for (values, 0..) |*v, i| v.* = sample.widen(row, w.dtype, i); + try process(w.gpa, values, r); + return sample.drawWide(w.gpa, values, r.sampling, r.position); } }; + +/// HF repetition processing widens raw scores before FP32 multiply/divide; it must not round back to BF16. +pub fn process(a: std.mem.Allocator, values: []f32, r: Request) !void { + if (!(r.repetition > 0) or !std.math.isFinite(r.repetition)) return error.BadRepetitionPenalty; + const counts = try a.alloc(u32, values.len); + defer a.free(counts); + @memset(counts, 0); + if (r.repetition != 1) { + for (r.history) |id| { + if (id >= values.len) return error.BadHistoryToken; + counts[id] = 1; + } + for (values, counts) |*v, seen| if (seen != 0) { + v.* = if (v.* < 0) v.* * r.repetition else v.* / r.repetition; + }; + } + if (r.penalties.active()) { + @memset(counts, 0); + for (r.reply) |id| { + if (id >= values.len) return error.BadHistoryToken; + counts[id] += 1; + } + for (values, counts) |*v, count| if (count != 0) { + v.* -= r.penalties.of(count); + }; + } +} + +test "repetition keeps FP32 ranking and counts prompt tokens once, OpenAI penalties count only reply tokens" { + var values = [_]f32{ 1.5, 1.3671875, -2, 4 }; + try process(std.testing.allocator, &values, .{ .sampling = null, .position = 0, .repetition = 1.1, .history = &.{ 0, 0, 2 }, .reply = &.{ 3, 3 }, .penalties = .{ .presence = 1, .frequency = 0.5 } }); + try std.testing.expectEqual(@as(f32, 1.5) / @as(f32, 1.1), values[0]); + try std.testing.expectEqual(@as(f32, -2) * @as(f32, 1.1), values[2]); + try std.testing.expectEqual(@as(f32, 2), values[3]); + try std.testing.expectEqual(@as(u32, 1), try sample.drawWide(std.testing.allocator, values[0..2], null, 0)); + try std.testing.expectError(error.BadHistoryToken, process(std.testing.allocator, &values, .{ .sampling = null, .position = 0, .repetition = 1.1, .history = &.{4} })); +} diff --git a/zig/src/hip_serve/sample.zig b/zig/src/hip_serve/sample.zig index d0939841de..8fb049b112 100644 --- a/zig/src/hip_serve/sample.zig +++ b/zig/src/hip_serve/sample.zig @@ -52,8 +52,25 @@ fn larger(_: void, a: Candidate, b: Candidate) bool { /// draw(): greedy (no sampling or temperature <= 0) is the argmax; else `choose` at `position` over the candidates. pub fn draw(gpa: Allocator, row: []const u16, dtype: Dtype, sampling: ?lanes.Sampling, position: u64) !u32 { - const s = sampling orelse return argmax(row, dtype); - if (s.temperature <= 0.0) return argmax(row, dtype); + if (sampling == null or sampling.?.temperature <= 0) return argmax(row, dtype); + const values = try gpa.alloc(f32, row.len); + defer gpa.free(values); + for (values, 0..) |*v, i| v.* = widen(row, dtype, i); + return drawWide(gpa, values, sampling, position); +} + +/// Already processed FP32 scores retain their precision through ranking and the keyed sampler. +pub fn drawWide(gpa: Allocator, row: []const f32, sampling: ?lanes.Sampling, position: u64) !u32 { + if (row.len == 0) return error.NoAllowedToken; + if (sampling == null or sampling.?.temperature <= 0) { + var best: usize = 0; + for (row, 0..) |v, i| { + if (std.math.isNan(row[best])) break; + if (std.math.isNan(v) or v > row[best]) best = i; + } + return @intCast(best); + } + const s = sampling.?; const width = row.len; const count = if (s.top_k != 0) @min(width, @as(usize, s.top_k) + MARGIN) else width; const best = try gpa.alloc(Candidate, count); @@ -61,10 +78,10 @@ pub fn draw(gpa: Allocator, row: []const u16, dtype: Dtype, sampling: ?lanes.Sam var filled: usize = 0; // the whole row when top_k is 0 (choose orders it), else the top `count` by value then id in one pass if (count == width) { - for (best, 0..) |*c, i| c.* = .{ .value = widen(row, dtype, i), .id = @intCast(i) }; + for (best, 0..) |*c, i| c.* = .{ .value = row[i], .id = @intCast(i) }; filled = width; } else for (0..width) |i| { - const c: Candidate = .{ .value = widen(row, dtype, i), .id = @intCast(i) }; + const c: Candidate = .{ .value = row[i], .id = @intCast(i) }; if (filled == count and !larger({}, c, best[count - 1])) continue; var at = if (filled < count) filled else count - 1; if (filled < count) filled += 1; diff --git a/zig/src/native/hip.zig b/zig/src/native/hip.zig index 9457fde272..b291835de9 100644 --- a/zig/src/native/hip.zig +++ b/zig/src/native/hip.zig @@ -11,13 +11,14 @@ const serve = @import("hip_serve"); const glm_ocr = @import("glm_ocr"); const qwen3 = @import("qwen3"); const dots = @import("dots_mocr"); +const parse = @import("nemotron_parse"); const fam = @import("hip_families.zig"); const Allocator = std.mem.Allocator; const admission = hip.admission; const modelContext = admission.modelContext; /// The HIP families: namespaces with `model_type`, `formats`, `default_context`, `prefill_step`, `open` and `follow`. -const registry = .{ qwen35.native, qwen35.native_moe, flashnext.native, nemotron.native, glm_ocr.native, qwen3.native, dots.native }; +const registry = .{ qwen35.native, qwen35.native_moe, flashnext.native, nemotron.native, glm_ocr.native, qwen3.native, dots.native, parse }; pub const backends: []const []const u8 = &.{"hip"}; /// No per-thread alternate signal stack: its thread-local leaves RCCL's and HIP's threads too little stack to start. @@ -370,7 +371,7 @@ test "a HIP family declares replies without reasoning with `pub const reasoning })); } -test "the text families carry no offer, so the host gives them the empty one; GLM-OCR, Qwen3, Qwen3.5 and dots.mocr offer through their Loaded" { +test "the text families carry no offer, so the host gives them the empty one; GLM-OCR, Qwen3, Qwen3.5, dots.mocr and Parse offer through their Loaded" { inline for (.{ flashnext.native, nemotron.native }) |F| { const Loaded = @typeInfo(@typeInfo(@TypeOf(F.open)).@"fn".return_type.?).error_union.payload; try std.testing.expect(!@hasField(Loaded, "offer")); @@ -383,6 +384,8 @@ test "the text families carry no offer, so the host gives them the empty one; GL try std.testing.expect(@FieldType(qwen35.native_engine.Loaded, "offer") == serve.Offer and !@hasField(qwen35.native_engine.Loaded, "readouts")); try std.testing.expect(@FieldType(dots.native.Loaded, "offer") == serve.Offer); try std.testing.expect(!@hasDecl(dots.native, "attach") and !@hasDecl(dots.native, "openVision")); + try std.testing.expect(@FieldType(parse.Loaded, "offer") == serve.Offer and !@hasField(parse.Loaded, "vision") and !@hasField(parse.Loaded, "cache")); + try std.testing.expect(@hasDecl(parse, "refusal") and !@hasDecl(parse, "flagRefused")); } test { @@ -398,6 +401,8 @@ test "every registered family is listed for capabilities" { try std.testing.expectEqualStrings("glm_ocr", families[4].model_type); try std.testing.expectEqualStrings("qwen3", families[5].model_type); try std.testing.expectEqualStrings("dots_ocr", families[6].model_type); + try std.testing.expectEqualStrings("nemotron_parse", families[7].model_type); + try std.testing.expect(!families[7].raw_text); } test "GLM-OCR native flags: MTP and tensor ranks refuse before readiness; copy drafts remain" { @@ -408,6 +413,10 @@ test "GLM-OCR native flags: MTP and tensor ranks refuse before readiness; copy d try std.testing.expect(flagRefused(F, .{ .dir = "missing", .model_type = F.model_type, .mtp_drafts = 0 }) == null); } +test "Nemotron Parse declares no reasoning block, so its OCR text stays answer content" { + try std.testing.expect(!reasoning(parse)); +} + test "Nemotron refuses --slide and --mtp-confidence on HIP before it loads, and keeps a cap" { const o: api.Open = .{ .dir = "d", .model_type = "nemotron_h", .slide = true }; try std.testing.expect(std.mem.indexOf(u8, flagRefused(nemotron.native, o).?, "--slide") != null); diff --git a/zig/src/server/fixtures/nemotron_parse.jinja b/zig/src/server/fixtures/nemotron_parse.jinja new file mode 100644 index 0000000000..6b6966a21d --- /dev/null +++ b/zig/src/server/fixtures/nemotron_parse.jinja @@ -0,0 +1,14 @@ +{%- if messages|length != 1 or messages[0]['role'] != 'user' -%} +{{- raise_exception('Nemotron Parse requires one user message with one image') -}} +{%- endif -%} +{%- set content = messages[0]['content'] -%} +{%- if content is string -%} +{{- raise_exception('Nemotron Parse requires image content parts') -}} +{%- endif -%} +{%- for part in content -%} +{%- if part['type'] == 'image_url' -%}{{ '' }}{%- endif -%} +{%- endfor -%} +{{- '' -}} +{%- for part in content -%} +{%- if part['type'] == 'text' -%}{{ part['text'] }}{%- endif -%} +{%- endfor -%} diff --git a/zig/src/server/hf_text.zig b/zig/src/server/hf_text.zig index ce7bfe7cf3..7b2ce54768 100644 --- a/zig/src/server/hf_text.zig +++ b/zig/src/server/hf_text.zig @@ -111,6 +111,7 @@ pub const HfText = struct { }; t.overridden = true; } else try t.readTemplates(io, a, dir, config); + if (template_override == null) t.source = @import("hf_text_parse.zig").template(model_config) orelse t.source; // DeepSeek-V4.x releases ship encoding.py, not Jinja: a checkpoint's Jinja is a third party's (--chat-template) if (t.native_deepseek) { t.unread_template = t.source.len > 0 or t.tool_source.len > 0; diff --git a/zig/src/server/hf_text_parse.zig b/zig/src/server/hf_text_parse.zig new file mode 100644 index 0000000000..7544c16f0c --- /dev/null +++ b/zig/src/server/hf_text_parse.zig @@ -0,0 +1,10 @@ +//! Nemotron Parse's chat template, which hf_text.zig loads in place of the checkpoint's. +const std = @import("std"); + +/// The bridge from server media parts to Parse's visual prefix (its OCR example passes images separately), else null. +pub fn template(model_config: std.json.Value) ?[]const u8 { + if (model_config != .object) return null; + const kind = model_config.object.get("model_type") orelse return null; + if (kind != .string or !std.mem.eql(u8, kind.string, "nemotron_parse")) return null; + return @embedFile("fixtures/nemotron_parse.jinja"); +} diff --git a/zig/src/server/images.zig b/zig/src/server/images.zig index 6ebf4acb32..a96adc24e6 100644 --- a/zig/src/server/images.zig +++ b/zig/src/server/images.zig @@ -266,7 +266,7 @@ pub fn both(a: std.mem.Allocator, ids: *std.ArrayList(u32), key: *std.ArrayList( /// Appends one image's rows, its engine entry and its cache words; `named`: the word before them is its marker's. fn placeImage(a: std.mem.Allocator, ids: *std.ArrayList(u32), key: *std.ArrayList(u32), imgs: *std.ArrayList(api.Image), token: u32, img: api.PreparedImage, named: bool) !void { if (named) key.items[key.items.len - 1] = image_key.open(img.hash); - try imgs.append(a, .{ .pixels = img.pixels, .gh = img.gh, .gw = img.gw, .at = @intCast(ids.items.len), .tokens = img.tokens, .hash = img.hash }); + try imgs.append(a, .{ .pixels = img.pixels, .gh = img.gh, .gw = img.gw, .at = @intCast(ids.items.len), .tokens = img.tokens, .content_size = img.content_size, .hash = img.hash }); for (0..img.tokens) |j| { try ids.append(a, token); try key.append(a, image_key.row(img.hash, j)); @@ -295,7 +295,7 @@ fn fakePrepare(ctx: *anyopaque, a: std.mem.Allocator, bytes: []const u8, max_tok _ = ctx; _ = max_tokens; const tokens: u32 = @intCast(bytes.len - 3); // the test's images: a JPEG's first bytes, then one token a byte - return .{ .pixels = try a.alloc(f32, 0), .gh = 2, .gw = 2 * tokens, .tokens = tokens, .hash = bytes[3] }; + return .{ .pixels = try a.alloc(f32, 0), .gh = 2, .gw = 2 * tokens, .tokens = tokens, .content_size = .{ 16, 16 }, .hash = bytes[3] }; } test "placeholders expand to each image's tokens, the history and cache key with them" { @@ -320,6 +320,8 @@ test "placeholders expand to each image's tokens, the history and cache key with try std.testing.expectEqual(@as(usize, 5), x.history_len); // the first image is in the history, the second not try std.testing.expectEqual(@as(u32, 1), x.images[0].at); try std.testing.expectEqual(@as(u32, 5), x.images[1].at); + try std.testing.expectEqualDeep(@as(?[2]u32, .{ 16, 16 }), x.images[0].content_size); + try std.testing.expectEqualDeep(@as(?[2]u32, .{ 16, 16 }), x.images[1].content_size); try std.testing.expect(x.cache_key[1] >= 0x8000_0000 and x.cache_key[1] != x.cache_key[2]); try std.testing.expectEqual(@as(u32, 2), x.cache_key[3]); // a placeholder the images do not account for is refused diff --git a/zig/tests/hip/contract_main.zig b/zig/tests/hip/contract_main.zig index 7ed7e3012c..c0ed124c00 100644 --- a/zig/tests/hip/contract_main.zig +++ b/zig/tests/hip/contract_main.zig @@ -34,6 +34,8 @@ const columns_k = @import("contract_columns.zig"); const vision_k = @import("contract_vision.zig"); const full_k = @import("contract_vision_attention.zig"); const text_k = @import("contract_dense_text.zig"); +const vl_k = @import("contract_vl.zig"); +const pos_k = @import("contract_vl_position.zig"); const g = core.contract.gates; @@ -303,6 +305,15 @@ fn run(a: std.mem.Allocator, io: std.Io) !usize { for (&text, std.enums.values(text_k.Which)) |*v, w| v.* = .init(a, &r, s, &l, w); defer for (&text) |*v| v.deinit(); for (&text) |*v| try kernels.append(a, v.kernel()); + const vl_kernels = try hip.vision.Kernels.load(l.mods[0]); + var vl: [std.enums.values(vl_k.Which).len]vl_k.Vision = undefined; + for (&vl, std.enums.values(vl_k.Which)) |*v, which_vl| v.* = .init(a, &r, s, vl_kernels, which_vl); + defer for (&vl) |*v| v.deinit(); + for (&vl) |*v| try kernels.append(a, v.kernel()); + var pos: [std.enums.values(pos_k.Which).len]pos_k.Position = undefined; + for (&pos, std.enums.values(pos_k.Which)) |*p, which_pos| p.* = .init(a, &r, s, vl_kernels, which_pos); + defer for (&pos) |*p| p.deinit(); + for (&pos) |*p| try kernels.append(a, p.kernel()); var t: g.Tally = .{}; const t0 = ms(io); try g.runAll(a, kernels.items, 0x6b63, &t); diff --git a/zig/tests/hip/contract_vl.zig b/zig/tests/hip/contract_vl.zig new file mode 100644 index 0000000000..4d2f985bd9 --- /dev/null +++ b/zig/tests/hip/contract_vl.zig @@ -0,0 +1,234 @@ +//! Shared vision and prefix-decoder operators through the HIP contract: independent FP64 formulas and buffers. +const std = @import("std"); +const hip = @import("hip"); +const core = @import("core"); +const Rig = @import("contract_rig.zig").Rig; +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const Dense = contract.dense.Dense; + +pub const Which = enum { linear, norm, gelu, add, scale, heads, four, embed }; +fn depth(s: c.Shape) usize { + return 4 * s.k + 257; +} +fn noSums(_: c.Shape) usize { + return 0; +} + +pub fn spec(w: Which) c.Spec { + var s: c.Spec = .{ .name = @tagName(w), .accum = .none, .depth = noSums, .out = .bf16, .served = &.{.{ .rows = 17, .n = 1, .k = 1024 }}, .tiles = .{ .{ .size = 256 }, .{}, .{ .size = 256 } } }; + s.name = switch (w) { + .linear => "hip tf_vl_linear BF16 biased FP32 dense", + .norm => "hip tf_vl_layernorm BF16 biased two-pass norm", + .gelu => "hip tf_vl_pointwise exact GELU", + .add => "hip tf_vl_pointwise residual", + .scale => "hip tf_vl_pointwise scale", + .heads => "hip tf_vl_layout heads to rows", + .four => "hip tf_vl_layout horizontal four", + .embed => "hip tf_vl_embed scaled BF16 table", + }; + switch (w) { + .linear => { + s.accum = .f32; + s.depth = depth; + s.tiles = .{ .{ .size = 16 }, .{ .size = 64 }, .{ .size = 16 } }; + s.served = &.{ .{ .rows = 1, .n = 72512, .k = 1024 }, .{ .rows = 17, .n = 1152, .k = 1152 }, .{ .rows = 1, .n = 4304, .k = 1152 }, .{ .rows = 1, .n = 1024, .k = 4304 } }; + }, + .norm => { + s.accum = .f32; + s.depth = depth; + s.ops = 5; + s.served = &.{ .{ .rows = 17, .n = 1, .k = 1024 }, .{ .rows = 1, .n = 1, .k = 1152 }, .{ .rows = 1, .n = 1, .k = 2304 } }; + }, + .gelu => s.ops = 5, + .heads => { + s.served = &.{ .{ .rows = 17, .n = 16, .k = 72 }, .{ .rows = 1, .n = 16, .k = 64 } }; + s.tiles[1] = .{ .size = 16 }; + }, + .four => { + s.served = &.{.{ .rows = 416, .n = 1, .k = 1024 }}; + s.tiles[0] = .{ .size = 256, .min = 4, .step = 4 }; + }, + .embed, .scale => s.ops = 1, + .add => {}, + } + return s; +} + +pub const Vision = struct { + rig: Rig, + kernels: hip.vision.Kernels, + which: Which, + arena: std.heap.ArenaAllocator, + shape: c.Shape = undefined, + dense: ?Dense = null, + x: []u16 = &.{}, + y: []u16 = &.{}, + weight: []u16 = &.{}, + bias: []u16 = &.{}, + ids: []i32 = &.{}, + ptrs: [5]u64 = @splat(0), + stress: bool = false, + + pub fn init(a: std.mem.Allocator, r: *const hip.Runtime, stream: hip.Stream, kernels: hip.vision.Kernels, which: Which) Vision { + return .{ .rig = .{ .a = a, .r = r, .s = stream }, .kernels = kernels, .which = which, .arena = .init(a) }; + } + pub fn deinit(v: *Vision) void { + v.rig.deinit(); + v.arena.deinit(); + } + pub fn kernel(v: *Vision) g.Kernel { + return .{ .spec = spec(v.which), .ctx = v, .vt = &vt }; + } + const vt = g.hooks(Vision); + + pub fn count(v: *Vision) usize { + const s = v.shape; + return s.rows * switch (v.which) { + .linear => s.n, + .heads => s.n * s.k, + else => s.k, + }; + } + + pub fn prepare(v: *Vision, s: c.Shape, fill: g.Fill, seed: u64) !bool { + if (s.rows == 0 or s.k == 0 or s.k > 16384 or s.n == 0 or s.rows > 16384) return false; + if (v.which == .four and s.rows % 4 != 0) return false; + v.rig.reset(); + _ = v.arena.reset(.retain_capacity); + const a = v.arena.allocator(); + v.shape = s; + v.stress = fill == .stress; + var rng = std.Random.DefaultPrng.init(seed); + const random = rng.random(); + v.dense = null; + const xn = if (v.which == .embed) 8 * s.k else if (v.which == .heads) s.rows * s.n * s.k else s.rows * s.k; + v.x = try a.alloc(u16, xn); + v.y = try a.alloc(u16, xn); + v.weight = try a.alloc(u16, if (v.which == .linear) s.n * s.k else s.k); + v.bias = try a.alloc(u16, if (v.which == .linear) s.n else s.k); + for (v.x) |*x| x.* = c.toBf16(random.float(f32) * 4 - 2); + for (v.y) |*x| x.* = c.toBf16(random.float(f32) * 4 - 2); + for (v.weight) |*x| x.* = c.toBf16(0.5 + random.float(f32)); + for (v.bias) |*x| x.* = c.toBf16(random.float(f32) * 2 - 1); + if (v.which == .linear) { + const d = try Dense.init(a, s, 1); + d.fillX(fill, seed, 2); + v.x = d.x; + for (v.weight, d.w, 0..) |*bits, *w, i| { + const j = i % s.k; + w.* = if (fill == .stress) + (if (j < s.k / 2) @as(f64, 1) else -1) * @as(f64, @floatFromInt(1 + ((j % (s.k / 2)) * 7 + i / s.k) % 6)) + else + c.fromBf16(bits.*); + bits.* = c.toBf16(@floatCast(w.*)); + } + if (fill == .stress) @memset(v.bias, 0); + v.dense = d; + } else if (fill == .nonfinite) { + for ([_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32) }, 0..) |x, i| { + const at = if (v.which == .norm) i * s.k + s.k / 2 else i; + if (at < v.x.len) v.x[at] = c.toBf16(x); + } + } else if (fill == .stress and v.which == .norm) { + for (v.x, 0..) |*x, i| x.* = c.toBf16(256 + @as(f32, @floatFromInt(i % s.k % 7)) * 2); + @memset(v.weight, c.toBf16(1)); + @memset(v.bias, 0); + } + v.ids = try a.alloc(i32, s.rows); + for (v.ids, 0..) |*id, i| id.* = @intCast(i % 8); + v.ptrs = .{ + try v.rig.input(std.mem.sliceAsBytes(v.x)), + try v.rig.input(std.mem.sliceAsBytes(v.weight)), + try v.rig.input(std.mem.sliceAsBytes(v.bias)), + try v.rig.input(if (v.which == .embed) std.mem.sliceAsBytes(v.ids) else std.mem.sliceAsBytes(v.y)), + try v.rig.output(v.count() * 2), + }; + return true; + } + + pub fn run(v: *Vision) !void { + const s = v.shape; + const x: hip.ops.Tensor = .{ .ptr = v.ptrs[0], .kind = .bf16 }; + const y: hip.ops.Tensor = .{ .ptr = v.ptrs[3], .kind = .bf16 }; + const out: hip.ops.Tensor = .{ .ptr = v.ptrs[4], .kind = .bf16 }; + const k = v.kernels; + const stream = v.rig.s; + switch (v.which) { + .linear => try k.linear(stream, x, v.ptrs[1], v.ptrs[2], out, s.rows, s.n, s.k), + .norm => try k.norm(stream, x, v.ptrs[1], v.ptrs[2], out, s.rows, s.k, 1e-6), + .gelu, .add, .scale => try k.point(stream, x, y, out, v.count(), switch (v.which) { + .gelu => .gelu, + .add => .add, + else => .scale, + }, 32), + .heads => try k.layout(stream, x, out, s.rows, s.n, s.k, .heads_to_rows), + .four => try k.layout(stream, x, out, s.rows, 1, s.k, .horizontal_four), + .embed => try k.embed(stream, x, v.ptrs[3], out, s.rows, s.k, 32), + } + try stream.synchronize(); + } + + pub fn reference(v: *Vision, want: []f64, mag: []f64) !void { + const s = v.shape; + @memset(mag, 0); + if (v.dense) |d| { + try d.reference(want, mag); + for (want, mag, 0..) |*w, *m, i| { + const bias = c.fromBf16(v.bias[i % s.n]); + w.* += bias; + m.* += @abs(bias); + } + return; + } + for (want, mag, 0..) |*w, *m, i| { + const row = i / s.k; + const col = i % s.k; + const x: f64 = c.fromBf16(v.x[@min(i, v.x.len - 1)]); + w.* = switch (v.which) { + .norm => blk: { + const values = v.x[row * s.k ..][0..s.k]; + var mean: f64 = 0; + var absolute: f64 = 0; + for (values) |z| { + mean += c.fromBf16(z); + absolute += @abs(c.fromBf16(z)); + } + mean /= @floatFromInt(s.k); + var variance: f64 = 0; + for (values) |z| variance += (c.fromBf16(z) - mean) * (c.fromBf16(z) - mean); + variance /= @floatFromInt(s.k); + const inv = 1 / @sqrt(variance + 1e-6); + const weight: f64 = c.fromBf16(v.weight[col]); + m.* = @abs(weight) * inv * (absolute / @as(f64, @floatFromInt(s.k)) + @abs(mean)); + break :blk (x - mean) * inv * weight + c.fromBf16(v.bias[col]); + }, + .gelu => (0.5 * x) * (1 + erf(x / @sqrt(@as(f64, 2)))), + .add => x + c.fromBf16(v.y[i]), + .scale => x * 32, + .heads => c.fromBf16(v.x[((i / s.k % s.n) * s.rows + i / (s.n * s.k)) * s.k + i % s.k]), + .four => c.fromBf16(v.x[(i / (4 * s.k) * 4 + i % 4) * s.k + i / 4 % s.k]), + .embed => @as(f64, c.fromBf16(v.x[@as(usize, @intCast(v.ids[row])) * s.k + col])) * 32, + .linear => unreachable, + }; + } + } + + pub fn simulate(v: *Vision, acc: c.Accum, order: g.Order, out: []f64) !void { + if (v.dense) |d| return d.simulate(acc, order, out); + const s = v.shape; + const a = v.arena.allocator(); + const terms = try a.alloc(f64, s.k); + for (0..s.rows) |row| { + const x = v.x[row * s.k ..][0..s.k]; + for (terms, x) |*t, bits| t.* = c.fromBf16(bits); + const mean = contract.rows.sumIn(acc, order, terms) / @as(f64, @floatFromInt(s.k)); + for (terms, x) |*t, bits| t.* = (c.fromBf16(bits) - mean) * (c.fromBf16(bits) - mean); + const variance = contract.rows.sumIn(acc, order, terms) / @as(f64, @floatFromInt(s.k)); + for (0..s.k) |j| out[row * s.k + j] = (c.fromBf16(x[j]) - mean) / @sqrt(variance + 1e-6) * c.fromBf16(v.weight[j]) + c.fromBf16(v.bias[j]); + } + } +}; + +extern "c" fn erf(f64) f64; diff --git a/zig/tests/hip/contract_vl_position.zig b/zig/tests/hip/contract_vl_position.zig new file mode 100644 index 0000000000..c955893f3b --- /dev/null +++ b/zig/tests/hip/contract_vl_position.zig @@ -0,0 +1,268 @@ +//! Position resize and border sampling through independent FP64 geometry, interpolation and accumulation references. +const std = @import("std"); +const hip = @import("hip"); +const core = @import("core"); +const Rig = @import("contract_rig.zig").Rig; +const c = core.contract.spec; +const g = core.contract.gates; + +pub const Which = enum { horizontal, vertical, area, position }; +fn horizontalDepth(s: c.Shape) usize { + return (2 * ((128 + s.n - 1) / s.n) + 2) * 2 + 1; +} +fn verticalDepth(s: c.Shape) usize { + return (2 * ((128 + s.rows - 1) / s.rows) + 2) * 2 + 1; +} +fn areaDepth(_: c.Shape) usize { + return 128 * 128 + 2; +} +fn positionDepth(_: c.Shape) usize { + return 8; +} + +pub fn spec(w: Which) c.Spec { + return .{ + .name = switch (w) { + .horizontal => "hip tf_vl_resize horizontal AA", + .vertical => "hip tf_vl_resize vertical AA", + .area => "hip tf_vl_resize adaptive area", + .position => "hip tf_vl_position content border sampling", + }, + .accum = .f32, + .depth = switch (w) { + .area => areaDepth, + .position => positionDepth, + .horizontal => horizontalDepth, + .vertical => verticalDepth, + }, + .ops = 16, + .mids = if (w == .position) 1 else 0, + .out = if (w == .position) .bf16 else .f32, + .served = switch (w) { + .horizontal => &.{ .{ .rows = 128, .n = 104, .k = 1152 }, .{ .rows = 128, .n = 17, .k = 8 } }, + .vertical => &.{ .{ .rows = 113, .n = 104, .k = 1152 }, .{ .rows = 17, .n = 65, .k = 8 } }, + .area => &.{ .{ .rows = 1, .n = 73, .k = 1152 }, .{ .rows = 113, .n = 1, .k = 8 } }, + .position => &.{ .{ .rows = 128, .n = 104, .k = 1152 }, .{ .rows = 8, .n = 8, .k = 8 }, .{ .rows = 1, .n = 4, .k = 8 } }, + }, + .sums = .{ .exact = w == .area }, + .tiles = .{ .{ .size = 16 }, .{ .size = 4 }, .{ .size = 256 } }, + }; +} + +const Term = struct { value: f64, weight: f64, index: usize }; +pub const Position = struct { + rig: Rig, + kernels: hip.vision.Kernels, + which: Which, + arena: std.heap.ArenaAllocator, + shape: c.Shape = undefined, + input: [2]usize = undefined, + output: [2]usize = undefined, + content: [2]usize = undefined, + values: []f64 = &.{}, + patches: []u16 = &.{}, + classes: []u16 = &.{}, + ptrs: [4]u64 = @splat(0), + stress: bool = false, + const vt = g.hooks(Position); + + pub fn init(a: std.mem.Allocator, r: *const hip.Runtime, stream: hip.Stream, kernels: hip.vision.Kernels, which: Which) Position { + return .{ .rig = .{ .a = a, .r = r, .s = stream }, .kernels = kernels, .which = which, .arena = .init(a) }; + } + pub fn deinit(p: *Position) void { + p.rig.deinit(); + p.arena.deinit(); + } + pub fn kernel(p: *Position) g.Kernel { + return .{ .spec = spec(p.which), .ctx = p, .vt = &vt }; + } + pub fn count(p: *Position) usize { + return (p.output[0] * p.output[1] + @as(usize, if (p.which == .position) 10 else 0)) * p.shape.k; + } + + pub fn prepare(p: *Position, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.rows > 128 or shape.n == 0 or shape.n > 128 or shape.k == 0 or shape.k > 4096) return false; + if (p.which == .area and shape.rows != 1 and shape.n != 1 and fill != .nonfinite) return false; + p.rig.reset(); + _ = p.arena.reset(.retain_capacity); + const a = p.arena.allocator(); + p.shape = shape; + p.stress = fill == .stress; + p.output = .{ if (p.which == .area and shape.n > 1) 1 else shape.rows, shape.n }; + p.content = .{ shape.rows * 16 - 7, shape.n * 16 - 5 }; + p.input = switch (p.which) { + .horizontal => .{ shape.rows, 128 }, + .vertical => .{ 128, shape.n }, + .area => .{ 128, 128 }, + .position => .{ shape.rows, shape.n }, + }; + p.values = try a.alloc(f64, p.input[0] * p.input[1] * shape.k); + var rng = std.Random.DefaultPrng.init(seed); + const random = rng.random(); + for (p.values, 0..) |*v, i| v.* = if (fill == .stress) (if ((i / shape.k / @as(usize, if (p.which == .vertical) p.input[1] else 1)) % 2 == 0) @as(f64, 32768) else -32766) else random.float(f64) * 4 - 2; + // Non-finites affect a known channel only, leaving all other channels finite. + if (fill == .nonfinite) for ([_]f64{ std.math.nan(f64), std.math.inf(f64), -std.math.inf(f64) }, 0..) |v, i| { + p.values[i * shape.k] = v; + }; + const f = try a.alloc(f32, p.values.len); + for (f, p.values) |*v, *wide| { + v.* = @floatCast(wide.*); + wide.* = v.*; + } + if (fill == .stress and (p.which == .position or p.which == .vertical)) { + // Large cancelling midpoint contributions distinguish FP32 from both reduced accumulators. + @memset(p.values, 0); + var terms_buffer: [256]Term = undefined; + const list = p.terms(0, &terms_buffer, true); + for (list, 0..) |t, i| if (t.weight != 0) { + const product: f64 = if (i < list.len / 2) 131583 else -131072; + p.values[t.index] = @as(f32, @floatCast(product / t.weight)); + }; + for (f, p.values) |*v, wide| v.* = @floatCast(wide); + } + p.ptrs[0] = try p.rig.input(std.mem.sliceAsBytes(f)); + p.patches = try a.alloc(u16, shape.rows * shape.n * shape.k); + p.classes = try a.alloc(u16, 10 * shape.k); + @memset(p.patches, 0); + for (p.classes, 0..) |*v, i| v.* = c.toBf16(@as(f32, @floatFromInt(i % 17)) / 8); + p.ptrs[1] = try p.rig.input(std.mem.sliceAsBytes(p.patches)); + p.ptrs[2] = try p.rig.input(std.mem.sliceAsBytes(p.classes)); + p.ptrs[3] = try p.rig.output(p.count() * if (p.which == .position) @as(usize, 2) else 4); + return true; + } + + pub fn run(p: *Position) !void { + if (p.which == .position) { + try p.kernels.position(p.rig.s, .{ .ptr = p.ptrs[1], .kind = .bf16 }, p.ptrs[2], p.ptrs[0], .{ .ptr = p.ptrs[3], .kind = .bf16 }, p.output, p.content, p.shape.k, 10); + } else try p.kernels.resize(p.rig.s, .{ .ptr = p.ptrs[0], .kind = .f32 }, p.ptrs[3], p.input, p.output, p.shape.k, switch (p.which) { + .horizontal => .horizontal, + .vertical => .vertical, + else => .area, + }); + try p.rig.s.synchronize(); + } + + fn rounded(v: f64, active: bool) f64 { + return if (active) @as(f32, @floatCast(v)) else v; + } + + fn coordinate32(index: usize, pixels: usize, size: usize) f64 { + @setFloatMode(.strict); + const start: f32 = @floatFromInt(index * 16); + const extent = std.math.clamp(@as(f32, @floatFromInt(pixels)) - start, 0, 16); + const normalized = @min(start + 0.5 * extent, @as(f32, @floatFromInt(pixels))) / @as(f32, @floatFromInt(pixels)); + const grid = 2 * normalized - 1; + return std.math.clamp(((grid + 1) * @as(f32, @floatFromInt(size)) - 1) * 0.5, 0, @as(f32, @floatFromInt(size - 1))); + } + + fn coordinate(index: usize, pixels: usize, size: usize) f64 { + const start: f64 = @floatFromInt(index * 16); + const remaining = std.math.clamp(@as(f64, @floatFromInt(pixels)) - start, 0, 16); + const normalized = @min(start + remaining / 2, @as(f64, @floatFromInt(pixels))) / @as(f64, @floatFromInt(pixels)); + return std.math.clamp(normalized * @as(f64, @floatFromInt(size)) - 0.5, 0, @as(f64, @floatFromInt(size - 1))); + } + + fn terms(p: *Position, index: usize, buffer: []Term, geometry32: bool) []Term { + const ch = index % p.shape.k; + const x = index / p.shape.k % p.output[1]; + const y = index / p.shape.k / p.output[1]; + var n: usize = 0; + if (p.which == .position) { + const px = if (geometry32) coordinate32(x, p.content[1], p.input[1]) else coordinate(x, p.content[1], p.input[1]); + const py = if (geometry32) coordinate32(y, p.content[0], p.input[0]) else coordinate(y, p.content[0], p.input[0]); + const x0: usize = @intFromFloat(@floor(px)); + const y0: usize = @intFromFloat(@floor(py)); + const sx = px - @as(f64, @floatFromInt(x0)); + const sy = py - @as(f64, @floatFromInt(y0)); + for (0..2) |dy| for (0..2) |dx| { + const at = (@min(y0 + dy, p.input[0] - 1) * p.input[1] + @min(x0 + dx, p.input[1] - 1)) * p.shape.k + ch; + const wx = if (dx == 0) rounded(1 - sx, geometry32) else sx; + const wy = if (dy == 0) rounded(1 - sy, geometry32) else sy; + buffer[n] = .{ .value = p.values[at], .index = at, .weight = rounded(wx * wy, geometry32) }; + n += 1; + }; + } else if (p.which == .area) { + const y0 = y * p.input[0] / p.output[0]; + const y1 = ((y + 1) * p.input[0] + p.output[0] - 1) / p.output[0]; + const x0 = x * p.input[1] / p.output[1]; + const x1 = ((x + 1) * p.input[1] + p.output[1] - 1) / p.output[1]; + for (y0..y1) |yy| for (x0..x1) |xx| { + buffer[n] = .{ .value = p.values[(yy * p.input[1] + xx) * p.shape.k + ch], .index = (yy * p.input[1] + xx) * p.shape.k + ch, .weight = 1 / @as(f64, @floatFromInt((y1 - y0) * (x1 - x0))) }; + n += 1; + }; + } else { + const axis: usize = if (p.which == .horizontal) 1 else 0; + const scale = rounded(@as(f64, @floatFromInt(p.input[axis])) / @as(f64, @floatFromInt(p.output[axis])), geometry32); + const support = @max(scale, 1); + const center = rounded((@as(f64, @floatFromInt(if (axis == 1) x else y)) + 0.5) * scale, geometry32); + const inv = rounded(1 / support, geometry32); + const first: usize = @intFromFloat(@max(@trunc(center - support + 0.5), 0)); + const end: usize = @intFromFloat(@min(@trunc(center + support + 0.5), @as(f64, @floatFromInt(p.input[axis])))); + var sum: f64 = 0; + for (first..end) |j| { + const distance = @abs(rounded(rounded(rounded(@as(f64, @floatFromInt(j)) - center, geometry32) + 0.5, geometry32) * inv, geometry32)); + sum = rounded(sum + @max(rounded(1 - distance, geometry32), 0), geometry32); + } + for (first..end) |j| { + const at = if (axis == 1) (y * p.input[1] + j) else (j * p.input[1] + x); + const distance = @abs(rounded(rounded(rounded(@as(f64, @floatFromInt(j)) - center, geometry32) + 0.5, geometry32) * inv, geometry32)); + buffer[n] = .{ .value = p.values[at * p.shape.k + ch], .index = at * p.shape.k + ch, .weight = rounded(@max(rounded(1 - distance, geometry32), 0) / sum, geometry32) }; + n += 1; + } + } + return buffer[0..n]; + } + + fn outputs(p: *Position, want: []f64, mag: ?[]f64, sim: ?struct { c.Accum, g.Order }) !void { + const a = p.arena.allocator(); + const buffer = try a.alloc(Term, 128 * 128); + const sums = try a.alloc(f64, buffer.len); + for (want, 0..) |*w, i| { + if (p.which == .position and i < p.classes.len) { + w.* = c.fromBf16(p.classes[i]); + if (mag) |m| m[i] = 0; + continue; + } + const terms_ = p.terms(if (p.which == .position) i - p.classes.len else i, buffer, sim != null); + var exact: f64 = 0; + var magnitude: f64 = 0; + for (terms_, 0..) |t, j| { + sums[j] = t.value * t.weight; + exact += sums[j]; + magnitude += @abs(sums[j]); + } + w.* = if (sim) |how| core.contract.rows.sumIn(how[0], how[1], sums[0..terms_.len]) else exact; + if (mag) |m| m[i] = magnitude; + } + } + /// Bound fixed FP32 geometry displacement from ideal coordinates separately from the FP32 dot's error. + pub fn bounds(p: *Position, want: []const f64, mag: []const f64, allow: []f64) !void { + const a = p.arena.allocator(); + const buffer = try a.alloc(Term, 128 * 128); + const s = spec(p.which); + for (allow, want, mag, 0..) |*limit, exact, magnitude, i| { + if (p.which == .area) { + limit.* = if (p.stress) c.boundExact(s, exact) else c.bound(s, p.shape, exact, magnitude); + continue; + } + if (p.which == .position and i < p.classes.len) { + limit.* = 0; + continue; + } + const list = p.terms(if (p.which == .position) i - p.classes.len else i, buffer, true); + var value: f64 = 0; + var absolute: f64 = 0; + for (list) |t| { + value += t.value * t.weight; + absolute += @abs(t.value * t.weight); + } + limit.* = @abs(value - exact) + c.bound(s, p.shape, value, absolute); + } + } + pub fn reference(p: *Position, want: []f64, mag: []f64) !void { + return p.outputs(want, mag, null); + } + pub fn simulate(p: *Position, acc: c.Accum, order: g.Order, out: []f64) !void { + return p.outputs(out, null, .{ acc, order }); + } +};