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/build.zig b/build.zig index e78cd0bbcb..6f47553a25 100644 --- a/build.zig +++ b/build.zig @@ -191,6 +191,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/recipes/README.md b/docs/recipes/README.md index 8d38798aed..fc9b612b81 100644 --- a/docs/recipes/README.md +++ b/docs/recipes/README.md @@ -8,6 +8,7 @@ For agent-assisted contributions, read the [TensorFold contributor skill](SKILL. | Family | Recipe | | --- | --- | | 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..8e15af5e86 --- /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 model-free contract runner registers twelve new operator cases. +`hip-vision-record-objects` builds the actual ops image for gfx1030, +gfx1100, gfx1151 and gfx1201. `hip_vision_records.zig`, with the +`core/contract/hsaco.zig` module imported as `hsaco`, emits their measured +launch records. The runner's `vision` mode is diagnostic; its default +mode runs the complete contract and planted controls. + +## 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. New +launch records match actual compiles on all four recorded chips and the +live gfx1201 device. The complete contract still reports 120 inherited +launch-record mismatches for older kernels. Fresh compiles of the exact +base have identical usage metadata for all 456 pre-existing symbols. +Do not describe the complete contract as passing. + +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/hip_vision_records.zig b/tools/zig/hip_vision_records.zig new file mode 100644 index 0000000000..ff531e2bd6 --- /dev/null +++ b/tools/zig/hip_vision_records.zig @@ -0,0 +1,26 @@ +//! Emit measured records from hip-vision-record-objects; import core/contract/hsaco.zig as the hsaco module. +const std = @import("std"); +const hsaco = @import("hsaco"); + +pub fn main(init: std.process.Init) !void { + const a = init.arena.allocator(); + const args = try init.minimal.args.toSlice(a); + if (args.len != 2) return error.Usage; + const chips = [_][]const u8{ "gfx1030", "gfx1100", "gfx1151", "gfx1201" }; + var objects: [chips.len][]hsaco.Kernel = undefined; + for (chips, 0..) |chip, i| { + const path = try std.fmt.allocPrint(a, "{s}/{s}.hsaco", .{ args[1], chip }); + const bytes = try std.Io.Dir.cwd().readFileAlloc(init.io, path, a, .limited(1 << 24)); + objects[i] = try hsaco.kernels(a, bytes); + } + std.debug.print("//! Measured with hip-vision-record-objects; hipcc flags and caps match the served ops image.\nconst Record = @import(\"hip_records.zig\").Record;\n\npub const records = [_]Record{{\n", .{}); + for ([_][]const u8{ "tf_vl_embed", "tf_vl_linear", "tf_vl_layernorm", "tf_vl_pointwise", "tf_vl_layout", "tf_vl_resize", "tf_vl_position" }) |name| { + std.debug.print(" .{{\n .launch = .{{ .name = \"vision {s}\", .symbol = \"{s}\", .image = \"ops\", .threads = 256 }},\n .usage = .{{\n", .{ name, name }); + for (objects) |object| { + const u = (hsaco.find(object, name) orelse return error.NoSymbol).usage; + std.debug.print(" .{{ .vgprs = {d}, .sgprs = {d}, .lds = {d}, .private = {d}, .bound = {d}, .wave = {d} }},\n", .{ u.vgprs, u.sgprs, u.lds, u.private, u.bound, u.wave }); + } + std.debug.print(" }},\n }},\n", .{}); + } + std.debug.print("}};\n", .{}); +} 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('= (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); + } + } +} + +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 e00e23f6b2..cd77e416bd 100644 --- a/zig/src/core/contract/gates.zig +++ b/zig/src/core/contract/gates.zig @@ -310,7 +310,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.zig b/zig/src/core/contract/hip_records.zig index 45a7ce7e4c..1e31e56812 100644 --- a/zig/src/core/contract/hip_records.zig +++ b/zig/src/core/contract/hip_records.zig @@ -498,7 +498,8 @@ const own_records = [_]Record{ pub const records = own_records ++ @import("hip_records_stream.zig").records ++ @import("hip_records_router.zig").records ++ - @import("hip_records_glm_ocr.zig").records; + @import("hip_records_glm_ocr.zig").records ++ + @import("hip_records_vision.zig").records; /// Gate 5's planted HIP controls: a launch past its __launch_bounds__, and 72 KiB of LDS where a workgroup has 64. pub const controls = [_]Record{ diff --git a/zig/src/core/contract/hip_records_vision.zig b/zig/src/core/contract/hip_records_vision.zig new file mode 100644 index 0000000000..c4ce9fcf8c --- /dev/null +++ b/zig/src/core/contract/hip_records_vision.zig @@ -0,0 +1,68 @@ +//! Measured with hip-vision-record-objects; hipcc flags and caps match the served ops image. +const Record = @import("hip_records.zig").Record; + +pub const records = [_]Record{ + .{ + .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_api.zig b/zig/src/core/engine_api.zig index 4a98de5246..2c413f2e7f 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -254,19 +254,8 @@ pub const Ledger = struct { cost: *const fn (ctx: *anyopaque, prompt: u32, max_tokens: u32) u64, }; -pub const PromptCacheCapacity = struct { budget_bytes: u64, minimum_state_bytes: u64 }; - -pub const PromptCachePlan = struct { - source: enum { physical_footprint, explicit, metal_working_set }, - budget_bytes: u64, - explicit_budget: bool = false, - over_cap: ?bool = null, - ram_bytes: ?u64 = null, - ready_footprint_bytes: ?u64 = null, - cap_bytes: ?u64 = null, - room_bytes: ?u64 = null, - margin_bytes: ?u64 = null, -}; +pub const PromptCacheCapacity = @import("engine_status.zig").PromptCacheCapacity; +pub const PromptCachePlan = @import("engine_status.zig").PromptCachePlan; /// A native family's model type, weight formats and explicit raw-text capability. pub const Family = struct { model_type: []const u8, formats: []const []const u8, raw_text: bool = true }; diff --git a/zig/src/core/engine_status.zig b/zig/src/core/engine_status.zig index 4820e701d2..f03679716c 100644 --- a/zig/src/core/engine_status.zig +++ b/zig/src/core/engine_status.zig @@ -32,3 +32,17 @@ pub const Status = struct { pub const PromptCache = struct { states: u64 = 0, bytes: u64 = 0 }; }; + +pub const PromptCacheCapacity = struct { budget_bytes: u64, minimum_state_bytes: u64 }; + +pub const PromptCachePlan = struct { + source: enum { physical_footprint, explicit, metal_working_set }, + budget_bytes: u64, + explicit_budget: bool = false, + over_cap: ?bool = null, + ram_bytes: ?u64 = null, + ready_footprint_bytes: ?u64 = null, + cap_bytes: ?u64 = null, + room_bytes: ?u64 = null, + margin_bytes: ?u64 = null, +}; 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 25f62d1728..9bd6e6db22 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..6ffb35a0d8 --- /dev/null +++ b/zig/src/families/nemotron_parse/hip.zig @@ -0,0 +1,113 @@ +//! 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. +pub fn flagRefused(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, + vision: api.Vision, + cache: ?*api.prompt_cache.Store, +}; + +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, .cache = if (h.cache.budget > 0) &h.cache else null, .vision = .{ .ctx = own, .prepare = prepare, .image_token = 1, .markers = &.{1}, .image_tokens = own.image_tokens, .max_images = 1 } }; +} + +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(flagRefused(.{ .dir = "d", .model_type = model_type }) != null); + const base: api.Open = .{ .dir = "d", .model_type = model_type, .vision = .{} }; + try std.testing.expect(flagRefused(base) == null); + var other = base; + other.tp = 2; + try std.testing.expect(flagRefused(other) != null); + other = base; + other.mtp_drafts = 2; + try std.testing.expect(flagRefused(other) != null); + other = base; + other.context = 16385; + try std.testing.expect(flagRefused(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..aa502984c2 --- /dev/null +++ b/zig/src/families/nemotron_parse/tower.zig @@ -0,0 +1,78 @@ +//! 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; + +/// `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 block(c: Compute, x: Tensor, stem: []const u8, rows: usize) !void { + const h = config.vision_hidden; + var buf: [256]u8 = undefined; + const norm1 = try c.norm(x, try sub(&buf, stem, "norm1"), rows, h, config.vision_eps); + const qkv = try c.linear(norm1, try sub(&buf, stem, "attn.qkv"), rows, 3 * h, h, true); + const attended = try c.global(qkv, rows, config.vision_heads, h / config.vision_heads); + const projected = try c.linear(attended, try sub(&buf, stem, "attn.proj"), rows, h, h, true); + try c.residual(x, projected, rows * h); + const norm2 = try c.norm(x, try sub(&buf, stem, "norm2"), rows, h, config.vision_eps); + const up = try c.linear(norm2, try sub(&buf, stem, "mlp.fc1"), rows, config.vision_ffn, h, true); + try c.gelu(up, rows * config.vision_ffn); + const down = try c.linear(up, try sub(&buf, stem, "mlp.fc2"), rows, h, config.vision_ffn, true); + try c.residual(x, down, rows * h); +} + +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); +} + +fn sub(buf: []u8, stem: []const u8, end: []const u8) ![]const u8 { + return std.fmt.bufPrint(buf, "{s}.{s}", .{ stem, end }); +} 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/root.zig b/zig/src/hip/root.zig index ea7826b249..3c707e2ad7 100644 --- a/zig/src/hip/root.zig +++ b/zig/src/hip/root.zig @@ -26,6 +26,7 @@ pub const Launcher = launches.Launcher; /// The stream handle a dry run passes: launches do nothing on it. pub const counting = @import("launches/util.zig").counting; pub const raw = @import("raw.zig"); +pub const vision = @import("vision.zig"); pub const ops = @import("ops/ops.zig"); pub const blas = @import("blas.zig"); pub const blas_fixed = @import("blas_fixed.zig"); diff --git a/zig/src/hip/vision.zig b/zig/src/hip/vision.zig new file mode 100644 index 0000000000..2fd4d15fb9 --- /dev/null +++ b/zig/src/hip/vision.zig @@ -0,0 +1,136 @@ +//! Dense, biased LayerNorm and layouts for image towers and dense causal-prefix decoders. +const std = @import("std"); +const Module = @import("module.zig").Module; +const Function = @import("module.zig").Function; +const Stream = @import("stream.zig").Stream; +const launch = @import("launch.zig"); +const types = @import("ops/types.zig"); +const Tensor = types.Tensor; +pub const Error = types.Error; + +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, count: usize, mode: Point, scale: f32) Error!void { + if (x.kind != out.kind or count == 0 or count > 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(count))); + try args.add(@backingInt(x.kind)); + try args.add(@backingInt(mode)); + try args.add(scale); + try flat(k.point_fn, s, count, &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, count: usize, args: *launch.Args) Error!void { + try launch.launch(f, .{ .grid = .{ .x = @intCast((count + 255) / 256) }, .block = .{ .x = 256 } }, s, args); +} 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 44c5d9876b..bb475cfd04 100644 --- a/zig/src/native/hip.zig +++ b/zig/src/native/hip.zig @@ -9,21 +9,22 @@ const flashnext = @import("flashnext"); const nemotron = @import("nemotron"); const serve = @import("hip_serve"); const glm_ocr = @import("glm_ocr"); +const parse = @import("nemotron_parse"); 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 }; +const registry = .{ qwen35.native, qwen35.native_moe, flashnext.native, nemotron.native, glm_ocr.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. pub const std_options: std.Options = .{ .signal_stack_size = null }; /// The serve flags past the common ones this backend takes (server/cli.zig lists and parses them). -pub const serves: []const []const u8 = &.{ "--tp", "--policy", "--mtp-drafts", "--mtp-confidence", "--checkpoint-slots", "--backend=rocm" }; +pub const serves: []const []const u8 = &.{ "--tp", "--policy", "--mtp-drafts", "--mtp-confidence", "--checkpoint-slots", "--backend=rocm", "--vision" }; pub const families: []const api.Family = blk: { var out: [registry.len]api.Family = undefined; - for (registry, 0..) |F, i| out[i] = .{ .model_type = F.model_type, .formats = F.formats }; + for (registry, 0..) |F, i| out[i] = .{ .model_type = F.model_type, .formats = F.formats, .raw_text = if (@hasDecl(F, "raw_text")) F.raw_text else true }; const final = out; break :blk &final; }; @@ -207,6 +208,7 @@ fn reasoning(comptime F: type) bool { fn flagRefused(comptime F: type, o: api.Open) ?[]const u8 { if (o.tp != 1 and @hasDecl(F, "tp_refusal")) return F.tp_refusal; if (o.mtp_drafts) |n| if (n > 0 and @hasDecl(F, "mtp_refusal")) return F.mtp_refusal; + if (@hasDecl(F, "flagRefused")) if (F.flagRefused(o)) |why| return why; if (o.rope_yarn != null) return "--rope-yarn: the native HIP engine has no YaRN rope; serve without --rope-yarn, or on CUDA"; if (o.slide and @hasDecl(F, "slide_refusal")) return F.slide_refusal; if (o.mtp_confidence != null and @hasDecl(F, "confidence_refusal")) return F.confidence_refusal; @@ -282,6 +284,12 @@ pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]c return null; } +/// Under --vision, the start note of a family with no image tower: it serves text and the server refuses images. +fn textOnlyNote(comptime Loaded: type, o: api.Open, buf: []u8) ?[]const u8 { + if (o.vision == null or @hasField(Loaded, "vision")) return null; + return std.fmt.bufPrint(buf, "[tensorfold] --vision: the HIP engine has no image tower for {s}; serving text only, and image requests are refused\n", .{o.model_type}) catch null; +} + fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !?api.Opened { if (flagRefused(F, o)) |why| { problem.* = why; @@ -298,6 +306,8 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O else => |x| return x, }; errdefer loaded.deinit(loaded.ctx); + var note: [256]u8 = undefined; + if (textOnlyNote(@TypeOf(loaded), o, ¬e)) |line| std.debug.print("{s}", .{line}); const h = try gpa.create(Host); errdefer gpa.destroy(h); h.gpa = gpa; @@ -310,12 +320,12 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O h.clock = .{ .io = io }; h.core = lanes.Engine.init(gpa, &h.cfg, loaded.backend, h.clock.clock()); errdefer h.core.deinit(); - h.host = api.LaneHost.init(gpa, io, &h.core, .{ .lanes = @intCast(loaded.streams), .context_window = @intCast(r.options.window), .prefill_step = F.prefill_step, .decode_share = o.decode_share, .policy = h.policy_line, .startup = loaded.startup, .structures = true, .vision = visionOffer(loaded), .reasoning = reasoning(F) }); + h.host = api.LaneHost.init(gpa, io, &h.core, .{ .lanes = @intCast(loaded.streams), .context_window = @intCast(r.options.window), .prefill_step = F.prefill_step, .decode_share = o.decode_share, .policy = h.policy_line, .startup = loaded.startup, .structures = true, .penalties = loaded.backend.penalties, .vision = visionOffer(loaded), .reasoning = reasoning(F) }); h.host.memory = .{ .read = readMemory }; if (@hasField(@TypeOf(loaded), "cache")) { h.host.cache = loaded.cache; h.host.info_.prompt_cache = loaded.cache != null; - h.host.info_.prompt_cache_capacity = loaded.cache_capacity; + if (@hasField(@TypeOf(loaded), "cache_capacity")) h.host.info_.prompt_cache_capacity = loaded.cache_capacity; } try h.host.start(); served = true; @@ -366,6 +376,8 @@ test "every registered family is listed for capabilities" { try std.testing.expectEqualStrings("qwen4_exp", families[2].model_type); try std.testing.expectEqualStrings("nemotron_h", families[3].model_type); try std.testing.expectEqualStrings("glm_ocr", families[4].model_type); + try std.testing.expectEqualStrings("nemotron_parse", families[5].model_type); + try std.testing.expect(!families[5].raw_text); } test "GLM-OCR native flags: MTP and tensor ranks refuse before readiness; copy drafts remain" { @@ -376,6 +388,37 @@ 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)); +} + +/// The engine a family's `open` returns. +fn LoadedOf(comptime F: type) type { + return @typeInfo(@typeInfo(@TypeOf(F.open)).@"fn".return_type.?).error_union.payload; +} + +test "--vision text families: the host prints the text-only note and offers no tower, so image parts get the 400" { + const seen: api.Open = .{ .dir = "d", .model_type = "qwen3_5", .vision = .{ .image_tokens = 64, .max_images = 1 } }; + var buf: [256]u8 = undefined; + inline for (.{ qwen35.native, qwen35.native_moe, flashnext.native, nemotron.native }) |F| { + try std.testing.expectEqualStrings( + "[tensorfold] --vision: the HIP engine has no image tower for qwen3_5; serving text only, and image requests are refused\n", + textOnlyNote(LoadedOf(F), seen, &buf).?, + ); + try std.testing.expect(textOnlyNote(LoadedOf(F), .{ .dir = "d", .model_type = "qwen3_5" }, &buf) == null); + // openWith hands LaneHost null, so messages.zig refuses image parts with a 400 (server/compact_test.zig) + try std.testing.expect(visionOffer(@as(LoadedOf(F), undefined)) == null); + } + try std.testing.expect(textOnlyNote(LoadedOf(glm_ocr.native), seen, &buf) == null); + try std.testing.expect(textOnlyNote(LoadedOf(parse), seen, &buf) == null); +} + +test "--vision refuses no text family on HIP: they start text-only, as on main" { + const seen: api.Open = .{ .dir = "d", .model_type = "m", .vision = .{ .image_tokens = 64, .max_images = 1 } }; + inline for (.{ qwen35.native, qwen35.native_moe, flashnext.native, nemotron.native }) |F| + try std.testing.expect(flagRefused(F, seen) == null); +} + 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 247be44ce3..a9eb2212ac 100644 --- a/zig/src/server/hf_text.zig +++ b/zig/src/server/hf_text.zig @@ -72,7 +72,14 @@ pub const HfText = struct { problem.* = try std.fmt.allocPrint(pa, "cannot read chat template override ({s})", .{@errorName(e)}); return error.Load; }; - } else try t.readTemplates(io, a, dir, config); + } else { + try t.readTemplates(io, a, dir, config); + // The upstream OCR example passes images separately; bridge server media parts to its visual prefix. + if (model_config == .object) if (model_config.object.get("model_type")) |kind| { + if (kind == .string and std.mem.eql(u8, kind.string, "nemotron_parse")) + t.source = @embedFile("fixtures/nemotron_parse.jinja"); + }; + } t.specials = try specialsOf(a, config); t.eos = t.eosIds(a, model_config, generation, config, t.tok.id_to_token.tokens.items.len) catch |e| { if (e == error.OutOfMemory) return e; 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 2b6aadb42f..aaeb3fd410 100644 --- a/zig/tests/hip/contract_main.zig +++ b/zig/tests/hip/contract_main.zig @@ -24,15 +24,18 @@ const kit = @import("contract_kit.zig"); const ocr_ops = @import("glm_ocr_ops").Ops; const ocr_points = @import("contract_glm_ocr_ops.zig"); const ocr_norms = @import("contract_glm_ocr_norm.zig"); +const vl_k = @import("contract_vl.zig"); +const pos_k = @import("contract_vl_position.zig"); const g = core.contract.gates; pub fn main(init: std.process.Init) !u8 { const a = init.gpa; const args = try init.minimal.args.toSlice(init.arena.allocator()); - if (args.len == 1) return if (try run(a, init.io) > 0) 1 else 0; + if (args.len == 1) return if (try run(a, init.io, false) > 0) 1 else 0; + if (args.len == 2 and std.mem.eql(u8, args[1], "vision")) return if (try run(a, init.io, true) > 0) 1 else 0; if (args.len != 3 or !std.mem.eql(u8, args[1], "limits")) { - std.debug.print("usage: tf-hip-kernel-contract [limits DIR]\n", .{}); + std.debug.print("usage: tf-hip-kernel-contract [limits DIR | vision]\n", .{}); return 2; } const n = try gate5.objects(a, init.io, args[2]); @@ -50,7 +53,7 @@ fn ms(io: std.Io) i96 { } /// Gates 1-4 on every registered HIP kernel and each control at its own gate, then gate 5; returns the problems. -fn run(a: std.mem.Allocator, io: std.Io) !usize { +fn run(a: std.mem.Allocator, io: std.Io, only_vision: bool) !usize { var built: [2]usize = .{ hip.kernels.images[0].len, controls.bytes.len }; _ = &built; // runtime-known, so a stub build still compiles all that follows if (built[0] == 0 or built[1] == 0) { @@ -171,8 +174,19 @@ fn run(a: std.mem.Allocator, io: std.Io) !usize { for (&norms, std.enums.values(ocr_norms.Which)) |*op, which_op| op.* = .init(a, &r, s, &ocr, which_op); defer for (&norms) |*op| op.deinit(); for (&norms) |*op| try kernels.append(a, op.kernel()); + // the vision operators go last: a vision-only run takes the tail + const vision = 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, vision, 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, vision, which_pos); + defer for (&pos) |*p| p.deinit(); + for (&pos) |*p| try kernels.append(a, p.kernel()); const t0 = ms(io); - try g.runAll(a, kernels.items, 0x6b63, &t); + if (only_vision) std.debug.print("DIAGNOSTIC: vision operators only; this is not a complete qualification receipt\n", .{}); + try g.runAll(a, if (only_vision) kernels.items[kernels.items.len - vl.len - pos.len ..] else kernels.items, 0x6b63, &t); const t1 = ms(io); const lim = try gate5.table(a) + try gate5.live(a, &r, ctx, arch); const t2 = ms(io); 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 }); + } +};