diff --git a/docs/recipes/README.md b/docs/recipes/README.md index f3e82f73b4..7d717e65dc 100644 --- a/docs/recipes/README.md +++ b/docs/recipes/README.md @@ -7,6 +7,7 @@ the native engine, or the Python engine frozen at 0.6.6. | --- | --- | | Nemotron 3.5 Lightning | [MLX](nemotron-3.5.md), [Intel Arc, experimental](xpu-arc.md) | | Nemotron 3.5 Super VL | [Native CUDA on one or two GB10s, images](nemotron-3.5-super-vl.md) | +| Nemotron 3.5 Bolt | [Native CUDA on one GB10, libbolt, images and video](nemotron-3.5-bolt.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) | | Qwen3.8 Flash Next | [MLX prefill and CUDA](qwen3.8-flash-next.md), [CUDA images](flash-next-vision.md) | diff --git a/docs/recipes/nemotron-3.5-bolt.md b/docs/recipes/nemotron-3.5-bolt.md new file mode 100644 index 0000000000..948847b9d5 --- /dev/null +++ b/docs/recipes/nemotron-3.5-bolt.md @@ -0,0 +1,97 @@ +# Nemotron 3.5 Bolt + +[NVIDIA-Nemotron-3.5-Bolt-74B-A9B-24GB](https://huggingface.co/nvidia/NVIDIA-Nemotron-3.5-Bolt-74B-A9B-24GB) is a +74B-parameter Mamba-2, attention and MoE model (9B active) packed to 24 GB: IQ2_XXS experts, NVFP4 Mamba projections, +Q8_0 attention and latent projections. Its `config.json` is a `nemotron_h_omni` wrapper around a `nemotron_h_puzzle` +text model (hidden 4096, 88 blocks, 240 to 496 routed experts a layer, top-9), so the native CUDA host serves it as its +own family, `nemotron_h_puzzle`, on one GB10 (sm_121). The family is `zig/src/families/nemotron/bolt_native.zig`, a lane +backend over `libbolt.so`; the HTTP server, chat template, round loop, keyed sampler and prompt cache are the engine's. + +## Build and run + +`libbolt.so` is built from its CUDA sources, which this repository does not carry (they are +[MiaAI-Lab's](https://github.com/MiaAI-Lab/Nemotron-3.5-Bolt-TensorFold-DGX-Spark/tree/main/engine)): + +```bash +zig build -Dnvcc=/usr/local/cuda/bin/nvcc -Dbolt-src=PATH/engine -Doptimize=fast install native bolt +tensorfold-native serve NVIDIA-Nemotron-3.5-Bolt-74B-A9B-24GB --parallel 8 --context 262144 +``` + +`native/lib/libbolt.so` is found beside the server (`TENSORFOLD_BOLT_LIB` names another); the backend checks its C ABI +version at load. Images need `--vision`. The drafter, +[LiLiCorr](https://huggingface.co/nvidia/NVIDIA-Nemotron-3.5-Bolt-74B-A9B-24GB-LiLiCorr), is named by `BOLT_DRAFT` +(a folder) or `--no-drafts` serves one row a round. `BOLT_KV=fp8` stores the attention cache as e4m3 with a scale per +token and head (4,224 bytes a token against 8,192), and `BOLT_KV_PIN=1` maps every stream's whole window at load. + +## What it serves + +- Up to 16 streams (`--parallel`) on one copy of the weights; a shared round carries every stream's verify window in one + forward. The window is the model's 262,144 tokens a stream; attention rows are mapped as a stream grows. +- LiLiCorr drafts verified in the lane rounds, seven tokens a window, drawn with the stream's keyed sampler. +- Prompt cache over the Mamba state, attention rows and drafter context (`snap_*`, any row, no lookahead). +- Images through the RADIO tower (`--vision`): each prompt piece encodes the images it crosses into the rows buffer, so a + request holds as many images as fit the window (`--vision-max-images`, up to 64). +- Structured output and forced tool calls: `response_format` (`json_schema`, `json_object`), `guided_*` fields and + `tool_choice` `required` or a named function. The server compiles the grammar (`libtfgrammar`, `zig build grammar`) to + per-row allowed-token bits; the backend redraws the rows of a verify window among those tokens (a masked top-k on the + GPU, `bolt_mask_rows`), hands the grammar its held drafts, and draws the first token again under the first mask. +- Not served: logprobs, speculative trees, int8 or int4 caches, `--slide`. + +## Exactness + +A stream's rounding is the same alone, among other streams and in a verify window of any width. Three things made it so: + +1. **Attention chunks follow the stream's own length.** A verify window merges partial softmaxes over chunks of keys. The + chunk is `win_chunk(base)`: 128 keys, doubled until 32 chunks reach the end of the sequence and its window. Before, the + chunks followed the number of streams in the block (about 32 splits for one stream, 8 for four) and the window's width, + so a stream among others rounded differently from the same stream alone. +2. **Decode GEMMs sum K in one order.** Dense GEMMs for blocks of up to 32 rows split K across four warps (interleaved + partial sums, then added); blocks over 32 rows used a kernel that walks K in one accumulator. A stream's rounding changed + when other streams grew the block past 32 rows. Decode and verify blocks now always use the K-split kernel (64 tokens + a CTA); the prompt pass keeps the faster large-block kernel. +3. **The plain step is a window.** The one-token step (no drafts) runs as a one-row verify window and is kept at once, so + it shares the arithmetic of drafted decoding. + +Checks in the engine's harness (not in this repository), run with 32-token GEMM tiles: the perplexity of four passages +scored alone and in a block of 2, 3, 4, 8 and 16 streams (windows of 1, 5 and 8 rows, up to 128 rows a block) is +identical to five decimals; sixteen sequences decoded in one block equal their solo runs for 768 of 768 tokens, plain +steps and 4-row windows alike. The final 64-token tile changes only the tokens a CTA holds, not the order K is summed +in, and is covered by the `bench_concurrent.py` table below. Perplexity on 1,499 tokens is 5.286. A stream's logits are +bit-identical whatever its neighbours hold (leak test). + +The prompt pass uses different kernels from decode by design; `--alone` compares decode. + +## Receipts + +TensorFold `46a08c18` plus this change; DGX Spark GB10, Ubuntu 24.04.4, kernel 6.17.0-1026-nvidia, driver 580.159.03, +CUDA 13.3.73; checkpoint `a326822af111290f2af5ac59ef4b5040347dabc4`, drafter `4c6a5613b6d887194a3a5493e3a18a5b3930bb86`; +`--parallel 8 --context 262144 --prompt-cache-gib 8 --vision`, bf16 cache with every stream's window mapped at load +(`BOLT_KV_PIN=1`), greedy, 256 tokens, 3 repetitions. + +`tools/bench_concurrent.py --levels 1,4,8 --alone --serial`: + +| Prompt | Streams | Aggregate | Each stream | Max TTFT | Concurrent equals solo (`--alone`) | Drafted equals `"draft": false` (`--serial`) | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | +| code | 1 | 115.1 tok/s | 115.1 | 0.07 s | 3 of 3 | 1 of 1 | +| code | 4 | 207.9 tok/s | 54.8 | 0.45 s | 12 of 12 | 1 of 1 | +| code | 8 | 252.6 tok/s | 35.2 | 1.15 s | 24 of 24 | 1 of 1 | +| chat | 1 | 60.7 tok/s | 60.7 | 0.05 s | 3 of 3 | 1 of 1 | +| chat | 4 | 139.1 tok/s | 35.6 | 0.42 s | 12 of 12 | 1 of 1 | +| chat | 8 | 174.7 tok/s | 23.6 | 1.05 s | 24 of 24 | 1 of 1 | + +Cold prefill with `tools/bench_cache.py --require-cold` (0 cached tokens, `qualified: true`), measured on the build +before the exactness changes, which touch decode blocks only (not re-run since): TTFT 1.44 s at 2,048 tokens, 5.60 s at +8,192, 24.66 s at 32,768, 55.60 s at 65,536. + +Cost of the K-split kernel in the widest rounds (a shared round of N streams x R rows, harness, ms; before the change / +after): 8 x 8 rows 135.5 / 130.4, 16 x 4 rows 133.2 / 141.7, 16 x 8 rows (128 rows) 176.9 / 220.3. Rounds of up to 32 rows +are unchanged. + +## Limits of this change + +`libbolt` is a CUDA runtime library, not fatbins under the kernel contract, and its sources are outside this tree +(4 files, about 4,800 lines, past the lean checker's file and comment rules); they are in the +[recipe repository](https://github.com/MiaAI-Lab/Nemotron-3.5-Bolt-TensorFold-DGX-Spark/tree/main/engine), which is +private for now. It targets sm_121 only and was run on one GB10. Whether the kernels belong in `zig/kernels/cuda` under +the contract, stay a separate repository, or go through another route is the maintainers' call; the Zig side here is the +part that does not depend on that answer. diff --git a/zig/build/bolt.zig b/zig/build/bolt.zig new file mode 100644 index 0000000000..f56d62462c --- /dev/null +++ b/zig/build/bolt.zig @@ -0,0 +1,27 @@ +//! libbolt, Nemotron 3.5 Bolt's forward pass: an optional library nvcc builds beside the server from -Dbolt-src. + +const std = @import("std"); + +const sources = [_][]const u8{ "bolt_kernels.cu", "bolt_mma.cu", "bolt_host.cu", "bolt_draft.cu" }; + +/// `zig build bolt -Dnvcc=NVCC -Dbolt-src=DIR`: native/lib/libbolt.so, which a Bolt checkpoint opens. +pub fn steps(b: *std.Build, nvcc: ?[]const u8) void { + const step = b.step("bolt", "libbolt (Nemotron 3.5 Bolt's engine, sm_121) from -Dbolt-src with -Dnvcc, beside the server"); + const src = b.option([]const u8, "bolt-src", "the directory of libbolt's .cu sources (bolt_host.cu and its siblings)"); + const tool = nvcc orelse { + step.dependOn(&b.addFail("zig build bolt needs -Dnvcc=PATH to nvcc").step); + return; + }; + const dir = src orelse { + step.dependOn(&b.addFail("zig build bolt needs -Dbolt-src=DIR, the folder holding bolt_host.cu").step); + return; + }; + const root: std.Build.LazyPath = .{ .cwd_relative = dir }; + const run = b.addSystemCommand(&.{ tool, "-shared", "-Xcompiler", "-fPIC", "-O3", "-gencode", "arch=compute_121,code=sm_121", "-I" }); + run.addDirectoryArg(root); + run.addArg("-o"); + const lib = run.addOutputFileArg("libbolt.so"); + for (sources) |name| run.addFileArg(root.path(b, name)); + run.addArg("-lcuda"); + step.dependOn(&b.addInstallFile(lib, "native/lib/libbolt.so").step); +} diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 684d6c3cc9..26e993ec9d 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -3,6 +3,7 @@ const std = @import("std"); const exlr8_build = @import("exlr8.zig"); const golden_build = @import("golden.zig"); +const bolt_build = @import("bolt.zig"); /// Each .cu in zig/kernels/cuda (`src`, else `name`) with the flags its Python extension passes in `extra_cuda_cflags`. const Kernel = struct { name: []const u8, flags: []const []const u8, src: ?[]const u8 = null, arch_specific: bool = false }; @@ -111,6 +112,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu const prebuilt = b.option([]const u8, "fatbins", "absolute directory of prebuilt .fatbin files to embed"); const sms = b.option([]const u8, "sm", "SASS targets, comma separated (default 121; 80, 86, 89, 120 and 121 build)") orelse "121"; const strip = b.option(bool, "strip", "No debug info in the CUDA executables: no build machine paths leave with them") orelse false; + bolt_build.steps(b, nvcc); // the compiler's version text is an input of every fatbin, so a new nvcc rebuilds them all const version: ?std.Build.LazyPath = if (prebuilt == null and nvcc != null) blk: { const run = b.addSystemCommand(&.{ nvcc.?, "--version" }); diff --git a/zig/src/cli/hub.zig b/zig/src/cli/hub.zig index 99d6835a49..cccb1f1123 100644 --- a/zig/src/cli/hub.zig +++ b/zig/src/cli/hub.zig @@ -11,6 +11,7 @@ pub const Family = struct { model_type: []const u8, title: []const u8, formats: const titles = [_]struct { model_type: []const u8, title: []const u8 }{ .{ .model_type = "nemotron_h", .title = "Nemotron 3.5 Lightning" }, .{ .model_type = "nemotron_h_omni", .title = "Nemotron 3.5 Super VL" }, + .{ .model_type = "nemotron_h_puzzle", .title = "Nemotron 3.5 Bolt" }, .{ .model_type = "qwen4_exp", .title = "Qwen3.8 Flash Next" }, .{ .model_type = "glm5_next", .title = "GLM-5.3-Flash" }, .{ .model_type = "qwen3_5", .title = "Qwen3.5-2B or Qwen3.8-27B" }, @@ -114,14 +115,18 @@ pub fn configInt(a: Allocator, io: std.Io, dir: []const u8, key: []const u8) ?u6 return if (v == .integer and v.integer >= 0) @intCast(v.integer) else null; } -/// config.json's model_type, "unknown" when unreadable. +/// config.json's model_type, "unknown" when unreadable; an omni wrapper of a puzzle model (Bolt) is nemotron_h_puzzle. pub fn modelType(a: Allocator, io: std.Io, dir: []const u8) []const u8 { const text = readSmall(a, io, pathJoin(a, dir, "config.json") catch return "unknown") orelse return "unknown"; var parsed = std.json.parseFromSlice(std.json.Value, a, text, .{}) catch return "unknown"; defer parsed.deinit(); if (parsed.value != .object) return "unknown"; const t = parsed.value.object.get("model_type") orelse return "unknown"; - return if (t == .string) t.string else "unknown"; + if (t != .string) return "unknown"; + if (std.mem.eql(u8, t.string, "nemotron_h_omni")) if (parsed.value.object.get("llm_config")) |llm| if (llm == .object) { + if (llm.object.get("model_type")) |lt| if (lt == .string and std.mem.eql(u8, lt.string, "nemotron_h_puzzle")) return "nemotron_h_puzzle"; + }; + return a.dupe(u8, t.string) catch "unknown"; } /// config.json's quantization: bits and group size, when the checkpoint is MLX-affine quantized. @@ -196,6 +201,11 @@ test "cache layout helpers read a synthetic cache" { const snapshot = (try cachedSnapshot(a, io, hub, "Org/Name")).?; try std.testing.expect(std.mem.endsWith(u8, snapshot, "0123456789abcdef0123456789abcdef01234567")); try std.testing.expectEqualStrings("qwen4_exp", modelType(a, io, snapshot)); + try w.writeFile(io, .{ .sub_path = try std.fs.path.join(a, &.{ snapshot, "config.json" }), .data = "{\"model_type\": \"nemotron_h_omni\", \"llm_config\": {\"model_type\": \"nemotron_h_puzzle\"}}" }); + try std.testing.expectEqualStrings("nemotron_h_puzzle", modelType(a, io, snapshot)); + try w.writeFile(io, .{ .sub_path = try std.fs.path.join(a, &.{ snapshot, "config.json" }), .data = "{\"model_type\": \"nemotron_h_omni\", \"llm_config\": {\"model_type\": \"nemotron_h\"}}" }); + try std.testing.expectEqualStrings("nemotron_h_omni", modelType(a, io, snapshot)); + try w.writeFile(io, .{ .sub_path = try std.fs.path.join(a, &.{ snapshot, "config.json" }), .data = "{\"model_type\": \"qwen4_exp\", \"quantization\": {\"bits\": 6, \"group_size\": 32}, \"max_position_embeddings\": 65536}" }); const q = quantization(a, io, snapshot).?; try std.testing.expectEqual(@as(u64, 6), q.bits); try std.testing.expectEqual(@as(u64, 32), q.group); diff --git a/zig/src/families/nemotron/bolt_api.zig b/zig/src/families/nemotron/bolt_api.zig new file mode 100644 index 0000000000..5e3b7c271a --- /dev/null +++ b/zig/src/families/nemotron/bolt_api.zig @@ -0,0 +1,95 @@ +//! libbolt's C ABI as the Bolt backend sees it: the symbols, resolved once, and where the library is found. +const std = @import("std"); + +/// The C ABI version of libbolt this backend speaks (`bolt_abi`). +pub const abi_version: c_int = 2; + +extern "c" fn dlopen(path: [*:0]const u8, mode: c_int) ?*anyopaque; +extern "c" fn dlsym(handle: ?*anyopaque, name: [*:0]const u8) ?*anyopaque; +pub extern "c" fn dlerror() ?[*:0]const u8; +extern "c" fn readlink(path: [*:0]const u8, buf: [*]u8, len: usize) isize; + +/// A run of a prompt piece whose embeddings are rows on the device (bolt_prefill_seq_rows). +pub const RowRun = extern struct { pos: c_int, count: c_int, rows: u64 }; + +/// libbolt's C ABI, resolved once. +pub const Api = struct { + open: *const fn ([*:0]const u8, c_int) callconv(.c) ?*anyopaque, + free: *const fn (?*anyopaque) callconv(.c) void, + seqs: *const fn (?*anyopaque) callconv(.c) c_int, + seq_alloc: *const fn (?*anyopaque) callconv(.c) c_int, + seq_release: *const fn (?*anyopaque, c_int) callconv(.c) void, + seq_bytes: *const fn (?*anyopaque) callconv(.c) u64, + prompt_end: *const fn (?*anyopaque, c_int, c_int) callconv(.c) void, + prefill: *const fn (?*anyopaque, c_int, [*]const c_int, c_int, c_int) callconv(.c) c_int, + forward: *const fn (?*anyopaque, c_int, [*]const c_int, c_int, c_int, c_int) callconv(.c) c_int, + forward_multi: *const fn (?*anyopaque, c_int, [*]const c_int, [*]const c_int, [*]const c_int, c_int) callconv(.c) c_int, + candidates: *const fn (?*anyopaque, c_int, c_int, [*]f32, [*]c_int) callconv(.c) void, + keep: *const fn (?*anyopaque, c_int, c_int) callconv(.c) void, + propose: *const fn (?*anyopaque, c_int, c_int, [*]c_int) callconv(.c) c_int, + propose_multi: *const fn (?*anyopaque, c_int, [*]const c_int, [*]const c_int, [*]c_int, [*]c_int) callconv(.c) void, + has_draft: *const fn (?*anyopaque) callconv(.c) c_int, + abi: *const fn () callconv(.c) c_int, + mask_rows: *const fn (?*anyopaque, c_int, [*]const c_int, [*]const u32, c_int, c_int) callconv(.c) c_int, + prefill_rows: *const fn (?*anyopaque, c_int, [*]const c_int, c_int, c_int, [*]const RowRun, c_int) callconv(.c) c_int, + cost: *const fn (?*anyopaque, c_int) callconv(.c) f64, + draft_multi_cost: *const fn (?*anyopaque) callconv(.c) f64, + shared_costs: *const fn (?*anyopaque, [*]c_int, [*]f64, c_int) callconv(.c) c_int, + snap_bytes: *const fn (?*anyopaque, c_int) callconv(.c) u64, + snap_save: *const fn (?*anyopaque, c_int, c_int) callconv(.c) ?*anyopaque, + snap_restore: *const fn (?*anyopaque, c_int, ?*anyopaque) callconv(.c) c_int, + snap_drop: *const fn (?*anyopaque, ?*anyopaque) callconv(.c) void, + + pub fn load(handle: ?*anyopaque) !Api { + return .{ + .open = try sym(*const fn ([*:0]const u8, c_int) callconv(.c) ?*anyopaque, handle, "bolt_open"), + .free = try sym(*const fn (?*anyopaque) callconv(.c) void, handle, "bolt_free"), + .seqs = try sym(*const fn (?*anyopaque) callconv(.c) c_int, handle, "bolt_seqs"), + .seq_alloc = try sym(*const fn (?*anyopaque) callconv(.c) c_int, handle, "bolt_seq_alloc"), + .seq_release = try sym(*const fn (?*anyopaque, c_int) callconv(.c) void, handle, "bolt_seq_release"), + .seq_bytes = try sym(*const fn (?*anyopaque) callconv(.c) u64, handle, "bolt_seq_bytes"), + .prompt_end = try sym(*const fn (?*anyopaque, c_int, c_int) callconv(.c) void, handle, "bolt_prompt_end"), + .prefill = try sym(*const fn (?*anyopaque, c_int, [*]const c_int, c_int, c_int) callconv(.c) c_int, handle, "bolt_prefill_seq"), + .forward = try sym(*const fn (?*anyopaque, c_int, [*]const c_int, c_int, c_int, c_int) callconv(.c) c_int, handle, "bolt_forward_seq"), + .forward_multi = try sym(*const fn (?*anyopaque, c_int, [*]const c_int, [*]const c_int, [*]const c_int, c_int) callconv(.c) c_int, handle, "bolt_forward_multi"), + .candidates = try sym(*const fn (?*anyopaque, c_int, c_int, [*]f32, [*]c_int) callconv(.c) void, handle, "bolt_candidates"), + .keep = try sym(*const fn (?*anyopaque, c_int, c_int) callconv(.c) void, handle, "bolt_keep_seq"), + .propose = try sym(*const fn (?*anyopaque, c_int, c_int, [*]c_int) callconv(.c) c_int, handle, "bolt_propose_seq"), + .propose_multi = try sym(*const fn (?*anyopaque, c_int, [*]const c_int, [*]const c_int, [*]c_int, [*]c_int) callconv(.c) void, handle, "bolt_propose_multi"), + .has_draft = try sym(*const fn (?*anyopaque) callconv(.c) c_int, handle, "bolt_has_draft"), + .abi = try sym(*const fn () callconv(.c) c_int, handle, "bolt_abi"), + .mask_rows = try sym(*const fn (?*anyopaque, c_int, [*]const c_int, [*]const u32, c_int, c_int) callconv(.c) c_int, handle, "bolt_mask_rows"), + .prefill_rows = try sym(*const fn (?*anyopaque, c_int, [*]const c_int, c_int, c_int, [*]const RowRun, c_int) callconv(.c) c_int, handle, "bolt_prefill_seq_rows"), + .cost = try sym(*const fn (?*anyopaque, c_int) callconv(.c) f64, handle, "bolt_cost"), + .draft_multi_cost = try sym(*const fn (?*anyopaque) callconv(.c) f64, handle, "bolt_draft_multi_cost"), + .shared_costs = try sym(*const fn (?*anyopaque, [*]c_int, [*]f64, c_int) callconv(.c) c_int, handle, "bolt_shared_costs"), + .snap_bytes = try sym(*const fn (?*anyopaque, c_int) callconv(.c) u64, handle, "bolt_snap_bytes"), + .snap_save = try sym(*const fn (?*anyopaque, c_int, c_int) callconv(.c) ?*anyopaque, handle, "bolt_snap_save_seq"), + .snap_restore = try sym(*const fn (?*anyopaque, c_int, ?*anyopaque) callconv(.c) c_int, handle, "bolt_snap_restore_seq"), + .snap_drop = try sym(*const fn (?*anyopaque, ?*anyopaque) callconv(.c) void, handle, "bolt_snap_drop"), + }; + } + + fn sym(comptime T: type, handle: ?*anyopaque, name: [*:0]const u8) !T { + const p = dlsym(handle, name) orelse { + std.log.err("bolt: libbolt.so has no {s}", .{name}); + return error.BoltMissing; + }; + return @ptrCast(@alignCast(p)); + } +}; + +/// libbolt.so: TENSORFOLD_BOLT_LIB, else beside the server (native/lib), else the loader's path. +pub fn openLibrary() ?*anyopaque { + if (std.c.getenv("TENSORFOLD_BOLT_LIB")) |p| return dlopen(p, 2); + var exe: [std.fs.max_path_bytes]u8 = undefined; + const n = readlink("/proc/self/exe", &exe, exe.len - 1); + if (n > 0) { + const dir = std.fs.path.dirname(exe[0..@intCast(n)]) orelse ""; + var buf: [std.fs.max_path_bytes]u8 = undefined; + if (std.fmt.bufPrintSentinel(&buf, "{s}/../lib/libbolt.so", .{dir}, 0)) |path| { + if (dlopen(path, 2)) |h| return h; + } else |_| {} + } + return dlopen("libbolt.so", 2); +} diff --git a/zig/src/families/nemotron/bolt_native.zig b/zig/src/families/nemotron/bolt_native.zig new file mode 100644 index 0000000000..d21162f343 --- /dev/null +++ b/zig/src/families/nemotron/bolt_native.zig @@ -0,0 +1,538 @@ +//! Nemotron 3.5 Bolt's lane backend over libbolt.so: prompts, shared verify windows, rollback, drafts, keyed draws. +const std = @import("std"); +const lanes = @import("lanes"); +const reuse = @import("cuda_reuse.zig"); +const cuda = @import("cuda"); +const radio = @import("radio"); +const native = @import("cuda_native.zig"); +const lib = @import("bolt_api.zig"); +pub const vision = native.vision; + +pub const model_type = "nemotron_h_puzzle"; +pub const formats: []const []const u8 = &.{"modelopt-mixed"}; +pub const default_context: i64 = 262144; +pub const max_segments: u32 = 1; +pub const prompt_rows: u32 = 2048; +pub const lone_min_budget: u64 = 0; +pub const explains_open = true; // explain() words open refusals too (--vision on a checkpoint without a tower) +pub const warm_quiet: u64 = 0; +/// No head reads a token ahead of a prompt: a state keys at its own row. +pub const snap_lookahead: u32 = 0; + +pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false, vision: ?vision.Limits = null, lanes: u32 = 1 }; + +/// Rows a verify window carries at most: the pending token and seven drafts. +const window: u32 = 8; +/// Prompt rows a cancel waits for at most, and the smaller piece while other streams decode. +const piece_rows: usize = 2048; +const busy_piece_rows: usize = 512; +/// Streams and rows one shared forward carries (libbolt's SEG_MAX and BATCH_MAX). +const max_segs: usize = 16; +const max_batch: u32 = 128; +/// Image spans a prompt piece may cross, and the rows its images take together (a piece cuts two 3,328-token images). +const max_runs = 64; +const image_rows: u32 = 16384; +/// Candidates kept per row for the keyed draw. +const cand_max: usize = 64; + +/// Acceptance by depth until a stream has its own (greedy replies on the drafter's own text: 0.86 a level). +const prior = [_]f64{ 0.80, 0.78, 0.76, 0.74, 0.72, 0.70, 0.68 }; + +const Timespec = extern struct { sec: i64, nsec: i64 }; +extern "c" fn clock_gettime(clk: c_int, ts: *Timespec) c_int; +fn nowMs() f64 { + var ts: Timespec = undefined; + _ = clock_gettime(1, &ts); + return @as(f64, @floatFromInt(ts.sec)) * 1000.0 + @as(f64, @floatFromInt(ts.nsec)) / 1e6; +} +extern "c" fn setenv(name: [*:0]const u8, value: [*:0]const u8, overwrite: c_int) c_int; + +/// What the backend keeps for one live stream: its libbolt sequence, the prompt pass in progress, and the drafts held. +const Slot = struct { + seq: c_int, + ids: []c_int = &.{}, // the prompt, while its pass runs + at: usize = 0, // prompt tokens done + decoding: bool = false, // the prompt pass is over + held: [window]u32 = @splat(0), // the drafts held for the stream's next window + held_n: u32 = 0, +}; + +/// Draws wait here between the forward and the round loop reading them back. +const draw_ring = 4096; + +const Engine = struct { + gpa: std.mem.Allocator, + api: lib.Api, + model: *anyopaque, + drafts: bool, + seqs: u32, + slots: std.AutoHashMapUnmanaged(*const lanes.Stream, *Slot) = .empty, + ring: [draw_ring]u32 = @splat(0), + drawn: u64 = 0, + decoding: u32 = 0, // streams past their prompt pass + stat_on: bool = false, // BOLT_STATS: log the round loop's time in the backend every 100 windows + stat: struct { rounds: u32 = 0, streams: u32 = 0, rows: u32 = 0, verify: f64 = 0, draw: f64 = 0, draft: f64 = 0, keep: f64 = 0 } = .{}, + target: reuse.Target = undefined, + vision: ?*vision.Vision = null, // --vision: the RADIO tower and a rows buffer one prompt piece uses + stream: ?cuda.Stream = null, // the stream the tower runs on + maskbuf: std.ArrayList(u32) = .empty, // structured output: a verify's masked rows' allowed-token bits, back to back + + /// Candidates the keyed draw needs for this stream's sampling: 1 when greedy. + fn candCount(s: *const lanes.Stream) c_int { + const sm = s.sampling orelse return 1; + if (sm.temperature <= 0) return 1; + const k: usize = if (sm.top_k == 0 or sm.top_k > cand_max) cand_max else sm.top_k; + return if (k <= 20) 20 else @intCast(cand_max); + } + + fn slotOf(e: *Engine, s: *const lanes.Stream) !*Slot { + return e.slots.get(s) orelse error.BoltNoSequence; + } + + /// The token drawn from one row of the last forward at `position`, keyed by the stream's seed. + fn draw(e: *Engine, s: *const lanes.Stream, row: usize, position: u64, k: c_int) !u32 { + var vals: [cand_max]f32 = undefined; + var ids: [cand_max]c_int = undefined; + e.api.candidates(e.model, @intCast(row), k, &vals, &ids); + const sm = s.sampling orelse return @intCast(ids[0]); + if (sm.temperature <= 0) return @intCast(ids[0]); + var n: usize = @intCast(k); + while (n > 1 and vals[n - 1] < -3.0e38) n -= 1; // a masked row with fewer allowed tokens than candidates + var v: [cand_max]f64 = undefined; + var d: [cand_max]u64 = undefined; + for (0..n) |i| { + v[i] = vals[i]; + d[i] = @intCast(ids[i]); + } + return @intCast(try lanes.sampling.choose(e.gpa, v[0..n], d[0..n], position, sm)); + } + + fn remember(e: *Engine, token: u32) u64 { + const h = e.drawn; + e.ring[h % draw_ring] = token; + e.drawn += 1; + return h; + } + + /// Take a sequence for the stream and resume it from a kept prompt state when the cache has one for its prefix. + fn begin(e: *Engine, s: *lanes.Stream) !*Slot { + if (s.images.len > 0 and e.vision == null) return error.NoVision; + const seq = e.api.seq_alloc(e.model); + if (seq < 0) return error.LaneMemory; // every sequence is in use: the stream waits at the queue's head + const slot = try e.gpa.create(Slot); + errdefer e.gpa.destroy(slot); + slot.* = .{ .seq = seq }; + errdefer e.api.seq_release(e.model, seq); + try e.slots.put(e.gpa, s, slot); + errdefer _ = e.slots.remove(s); + const prompt = s.prompt(); + s.cached = 0; + if (s.reuse.saved) |saved| { + if (s.reuse.at < prompt.len and e.api.snap_restore(e.model, seq, saved) == 0) { + slot.at = s.reuse.at; + s.cached = @intCast(slot.at); + } else { + s.reuse_failed = true; // the restore refused before touching the sequence: it starts from 0 + } + } + e.api.prompt_end(e.model, seq, @intCast(prompt.len)); + std.log.info("bolt: prefill {d} tokens ({d} from the cache), sequence {d}", .{ prompt.len, slot.at, slot.seq }); + slot.ids = try e.gpa.alloc(c_int, prompt.len); + for (slot.ids, prompt) |*d, t| d.* = @intCast(t); + return slot; + } + + /// One piece of the prompt pass, cut at the cache's marks, small enough that a cancel lands quickly. + fn prefillStep(ptr: *anyopaque, s: *lanes.Stream) !bool { + const e: *Engine = @ptrCast(@alignCast(ptr)); + const prompt_len = s.prompt().len; + const slot = e.slots.get(s) orelse try e.begin(s); + if (s.isCancelled() or s.yields()) return error.Cancelled; + if (slot.at >= prompt_len) return true; + const rows: usize = if (e.decoding > 0) busy_piece_rows else piece_rows; + var end = @min(prompt_len, slot.at + rows); + for (s.reuse.marks) |m| { + if (m > slot.at and m < end) end = m; + } + const last = end == prompt_len; + const cand: c_int = if (last) candCount(s) else 0; + var runs: [max_runs]lib.RowRun = undefined; + const nruns = if (s.images.len > 0) try e.imageRuns(s, slot.at, end, &runs) else 0; + const rc = if (nruns > 0) + e.api.prefill_rows(e.model, slot.seq, slot.ids.ptr + slot.at, @intCast(end - slot.at), cand, &runs, @intCast(nruns)) + else + e.api.prefill(e.model, slot.seq, slot.ids.ptr + slot.at, @intCast(end - slot.at), cand); + if (rc != 0) return error.BoltForward; + slot.at = end; + if (s.reuse.hook) |hk| { + if (std.mem.indexOfScalar(u32, s.reuse.marks, @intCast(slot.at)) != null) hk.at(hk.ptr, s, @intCast(slot.at)); + } + if (!last) return false; + e.gpa.free(slot.ids); + slot.ids = &.{}; + slot.decoding = true; + e.decoding += 1; + return true; + } + + /// The images rows [lo, hi) cross, encoded into the rows buffer, and the runs libbolt reads (a cut one is redone). + fn imageRuns(e: *Engine, s: *const lanes.Stream, lo: usize, hi: usize, runs: []lib.RowRun) !usize { + const v = e.vision orelse return error.NoVision; + const st = e.stream.?; + const row_bytes = v.rows.hidden * 2; + var top: u32 = 0; + var n: usize = 0; + for (s.images) |img| { + const a: usize = img.at; + const b = a + img.tokens; + if (b <= lo) continue; + if (a >= hi) break; // in prompt order + if (n == runs.len or top + img.tokens > v.rows.limit) return error.ImageRows; + const base = v.rows.buf.ptr + @as(u64, top) * row_bytes; + try v.tower.picture(st, img, base); + const from = @max(lo, a); + const to = @min(hi, b); + runs[n] = .{ .pos = @intCast(from - lo), .count = @intCast(to - from), .rows = base + @as(u64, @intCast(from - a)) * row_bytes }; + n += 1; + top += img.tokens; + } + if (n > 0) try st.synchronize(); // the rows are whole before libbolt copies them + return n; + } + + fn prefill(ptr: *anyopaque, s: *lanes.Stream) !void { + while (!try prefillStep(ptr, s)) {} + } + + fn first(ptr: *anyopaque, s: *lanes.Stream, position: u64) !u64 { + const e: *Engine = @ptrCast(@alignCast(ptr)); + const token = try e.draw(s, 0, position, candCount(s)); + return e.remember(token); + } + + /// The first token again, drawn among the tokens `mask` allows (the prompt pass drew it without the grammar). + fn firstMasked(ptr: *anyopaque, s: *lanes.Stream, position: u64, mask: []const u32) !u64 { + const e: *Engine = @ptrCast(@alignCast(ptr)); + const k = candCount(s); + const row0 = [1]c_int{0}; + if (mask.len == 0 or e.api.mask_rows(e.model, 1, &row0, mask.ptr, @intCast(mask.len), k) != 0) return error.BoltForward; + const token = try e.draw(s, 0, position, k); + return e.remember(token); + } + + /// The drafts held for the stream's next window, as host tokens in window order (a grammar walks them). + fn held(ptr: *anyopaque, s: *lanes.Stream, out: []u32) !void { + const e: *Engine = @ptrCast(@alignCast(ptr)); + const slot = try e.slotOf(s); + if (out.len > slot.held_n) return error.HeldDraftsUnreadable; + @memcpy(out, slot.held[0..out.len]); + } + + fn queue(ptr: *anyopaque, s: *lanes.Stream, feed: lanes.backend.Feed, position: u64) !u64 { + const e: *Engine = @ptrCast(@alignCast(ptr)); + const slot = try e.slotOf(s); + const tok: u32 = switch (feed) { + .handle => |h| e.ring[h % draw_ring], + .value => |v| v, + }; + const ids = [1]c_int{@intCast(tok)}; + const k = candCount(s); + // a one-row window, kept at once: the same attention arithmetic as every verify window, alone or among streams + if (e.api.forward(e.model, slot.seq, &ids, 1, 1, k) != 0) return error.BoltForward; + e.api.keep(e.model, slot.seq, 1); + const next = try e.draw(s, 0, position, k); + return e.remember(next); + } + + fn read(ptr: *anyopaque, handle: u64) !u32 { + const e: *Engine = @ptrCast(@alignCast(ptr)); + return e.ring[handle % draw_ring]; + } + + /// One forward over every window (one per stream), each row drawn at its position with its stream's sampling. + fn verify(ptr: *anyopaque, windows: []const lanes.backend.Window, out: []lanes.backend.Verified) !void { + const e: *Engine = @ptrCast(@alignCast(ptr)); + if (windows.len == 0 or windows.len > max_segs) return error.WindowTooWide; + var seqs: [max_segs]c_int = undefined; + var counts: [max_segs]c_int = undefined; + var toks: [max_batch]c_int = undefined; + var total: usize = 0; + var k: c_int = 1; + var masked: [max_batch]c_int = undefined; // the block rows whose draw is restricted by a grammar, and their mask width + var nmasked: usize = 0; + var mwords: usize = 0; + e.maskbuf.clearRetainingCapacity(); + for (windows, 0..) |w, i| { + const slot = try e.slotOf(w.stream); + const rows = w.rows(); + if (rows > window or w.parents != null or total + rows > max_batch) return error.WindowTooWide; + if (w.masks.len != 0) { + const words = w.masks.len / rows; + if (words == 0 or words * rows != w.masks.len or (mwords != 0 and words != mwords)) return error.BadMask; + mwords = words; + try e.maskbuf.appendSlice(e.gpa, w.masks); + for (0..rows) |r| { + masked[nmasked] = @intCast(total + r); + nmasked += 1; + } + } + toks[total] = @intCast(w.pending); + for (0..w.held) |j| { + toks[total + 1 + j] = @intCast(slot.held[j]); + out[i].drafts[j] = slot.held[j]; + } + for (w.tokens, 0..) |t, j| toks[total + 1 + w.held + j] = @intCast(t); + seqs[i] = slot.seq; + counts[i] = @intCast(rows); + total += rows; + k = @max(k, candCount(w.stream)); + } + if (e.seqs > 1) k = cand_max; // one candidate count keeps the captured forward graphs few + const t0 = nowMs(); + if (e.api.forward_multi(e.model, @intCast(windows.len), &seqs, &counts, &toks, k) != 0) return error.BoltForward; + if (nmasked > 0 and e.api.mask_rows(e.model, @intCast(nmasked), &masked, e.maskbuf.items.ptr, @intCast(mwords), cand_max) != 0) return error.BoltForward; + const t1 = nowMs(); + var row: usize = 0; + for (windows, 0..) |w, i| { + const kw = candCount(w.stream); + for (0..w.rows()) |r| out[i].sampled[r] = try e.draw(w.stream, row + r, w.positions[r], kw); + row += w.rows(); + (try e.slotOf(w.stream)).held_n = 0; + } + if (e.stat_on) { + e.stat.rounds += 1; + e.stat.streams += @intCast(windows.len); + e.stat.rows += @intCast(total); + e.stat.verify += t1 - t0; + e.stat.draw += nowMs() - t1; + if (e.stat.rounds == 100) { + const n: f64 = 100; + std.log.info("bolt stats: 100 rounds, {d:.1} streams {d:.1} rows a round; forward {d:.1} ms, draws {d:.1} ms, drafts {d:.1} ms, keeps {d:.1} ms a round", .{ @as(f64, @floatFromInt(e.stat.streams)) / n, @as(f64, @floatFromInt(e.stat.rows)) / n, e.stat.verify / n, e.stat.draw / n, e.stat.draft / n, e.stat.keep / n }); + e.stat = .{}; + } + } + } + + /// Roll each stream back to its window's kept rows. + fn keep(ptr: *anyopaque, windows: []const lanes.backend.Window, paths: []const []const u32) !void { + const e: *Engine = @ptrCast(@alignCast(ptr)); + const t0 = nowMs(); + for (windows, paths) |w, path| { + const slot = try e.slotOf(w.stream); + e.api.keep(e.model, slot.seq, @intCast(path.len)); + } + e.stat.keep += nowMs() - t0; + } + + /// Hold each stream's drafts for its next window: LiLiCorr over every stream's kept rows in one pass. + fn draft(ptr: *anyopaque, requests: []const lanes.backend.DraftRequest) !void { + const e: *Engine = @ptrCast(@alignCast(ptr)); + const t0 = nowMs(); + defer e.stat.draft += nowMs() - t0; + var seqs: [max_segs]c_int = undefined; + var anchors: [max_segs]c_int = undefined; + var who: [max_segs]*Slot = undefined; + var depths: [max_segs]u32 = undefined; + var n: usize = 0; + for (requests) |r| { + const slot = try e.slotOf(r.stream); + slot.held_n = 0; + // a shared round drafts before it keeps: cut the stream back to its kept rows first (a repeat is a no-op) + if (r.rows) |rows| e.api.keep(e.model, slot.seq, @intCast(rows.len)); + if (!e.drafts or r.depth == 0 or n == max_segs) continue; + const pending: u32 = if (r.follow.len > 0) r.follow[r.follow.len - 1] else switch (r.first orelse continue) { + .handle => |h| e.ring[h % draw_ring], + .value => |v| v, + }; + seqs[n] = slot.seq; + anchors[n] = @intCast(pending); + who[n] = slot; + depths[n] = r.depth; + n += 1; + } + if (n == 0) return; + var out: [max_segs * window]c_int = @splat(0); + var got: [max_segs]c_int = @splat(0); + e.api.propose_multi(e.model, @intCast(n), &seqs, &anchors, &out, &got); + for (0..n) |i| { + const depth: usize = @min(@as(usize, depths[i]), window - 1); + const have: usize = @intCast(@max(0, got[i])); + const pending: u32 = @intCast(anchors[i]); + for (0..depth) |j| who[i].held[j] = if (j < have) @intCast(out[i * window + j]) else pending; + who[i].held_n = @intCast(depth); + } + } + + fn release(ptr: *anyopaque, s: *lanes.Stream) void { + const e: *Engine = @ptrCast(@alignCast(ptr)); + const kv = e.slots.fetchRemove(s) orelse return; + const slot = kv.value; + if (slot.decoding) e.decoding -|= 1; + e.api.seq_release(e.model, slot.seq); + if (slot.ids.len > 0) e.gpa.free(slot.ids); + e.gpa.destroy(slot); + } + + fn backend(e: *Engine) lanes.backend.Backend { + const vt = struct { + var v: lanes.backend.Backend.VTable = .{ + .prefill = Engine.prefill, + .prefill_step = Engine.prefillStep, + .first = Engine.first, + .first_masked = Engine.firstMasked, + .held = Engine.held, + .queue = Engine.queue, + .read = Engine.read, + .verify = Engine.verify, + .keep = Engine.keep, + .draft = Engine.draft, + .release = Engine.release, + }; + }; + return .{ .ptr = e, .vtable = &vt.v }; + } +}; + +pub const Loaded = struct { + backend: lanes.backend.Backend, + facts: lanes.Model, + rows: u32, + stream_bytes: usize, + free_streams: u32 = 1, + ctx: *anyopaque, + deinit: *const fn (*anyopaque) void, + lone: ?@import("cuda_native.zig").LoneRun = null, + target: *reuse.Target, + learns: ?*@import("cuda_engine.zig").Engine = null, + learn_bytes: usize = 0, + vision: ?*vision.Vision = null, +}; + +fn shutdown(p: *anyopaque) void { + const e: *Engine = @ptrCast(@alignCast(p)); + e.api.free(e.model); + e.maskbuf.deinit(e.gpa); + if (e.vision) |v| v.deinit(); + if (e.stream) |*st| st.deinit(); + e.slots.deinit(e.gpa); + e.gpa.destroy(e); +} + +// -- prompt-cache copies: the Mamba state, the attention rows and the drafter's context after `at` prompt tokens ---- + +pub fn snap_bytes(ptr: *anyopaque, at: u32) u64 { + const e: *Engine = @ptrCast(@alignCast(ptr)); + return e.api.snap_bytes(e.model, @intCast(at)); +} + +/// `owner` is the lane core's stream: its sequence is the state to copy. +pub fn snap_save(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!*anyopaque { + const e: *Engine = @ptrCast(@alignCast(ptr)); + const s: *const lanes.Stream = @ptrCast(@alignCast(owner orelse return error.BoltSnapshot)); + const slot = e.slots.get(s) orelse return error.BoltSnapshot; + return e.api.snap_save(e.model, slot.seq, @intCast(at)) orelse error.BoltSnapshot; +} + +pub fn snap_restore(ptr: *anyopaque, owner: ?*anyopaque, saved: *anyopaque) anyerror!void { + const e: *Engine = @ptrCast(@alignCast(ptr)); + const s: *const lanes.Stream = @ptrCast(@alignCast(owner orelse return error.BoltSnapshot)); + const slot = e.slots.get(s) orelse return error.BoltSnapshot; + if (e.api.snap_restore(e.model, slot.seq, saved) != 0) return error.BoltSnapshot; + slot.held_n = 0; +} + +pub fn snap_drop(ptr: *anyopaque, saved: *anyopaque) void { + const e: *Engine = @ptrCast(@alignCast(ptr)); + e.api.snap_drop(e.model, saved); +} + +/// What the loaded libbolt measured: forward ms by window width, and shared forwards by total rows. +var costs: [window]lanes.config.Cost = undefined; +var shared: [32]lanes.config.Cost = undefined; + +pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, o: Options) !Loaded { + _ = kernels; + // --vision: the RADIO tower and its rows buffer, before the language model takes its memory + var tower: ?*vision.Vision = null; + var problem: []const u8 = ""; + if (o.vision) |lim| { + // the rows buffer holds one prompt piece's images; a request's token totals are the server's limit alone + tower = (try vision.Vision.init(gpa, io, ctx.d, dir, .{ .image_tokens = @min(lim.image_tokens, image_rows), .max_images = lim.max_images }, &problem)) orelse { + std.log.err("--vision: {s}", .{problem}); + return error.VisionRefused; + }; + tower.?.limits.image_tokens = lim.image_tokens; + } + errdefer if (tower) |v| v.deinit(); + var stream: ?cuda.Stream = if (tower != null) try cuda.Stream.init(ctx.d, false) else null; + errdefer if (stream) |*st| st.deinit(); + const handle = lib.openLibrary() orelse { + std.log.err("bolt: libbolt.so: {s} (zig build bolt -Dnvcc=NVCC installs it beside the server)", .{if (lib.dlerror()) |m| std.mem.span(m) else "not found"}); + return error.BoltMissing; + }; + const api = try lib.Api.load(handle); + if (api.abi() != lib.abi_version) { + std.log.err("bolt: libbolt.so speaks ABI {d}, this server {d}: build it from the same source", .{ api.abi(), lib.abi_version }); + return error.BoltMissing; + } + // one libbolt sequence a stream (--parallel), unless BOLT_SEQS says otherwise + var seqs_buf: [16]u8 = undefined; + _ = setenv("BOLT_SEQS", std.fmt.bufPrintSentinel(&seqs_buf, "{d}", .{o.lanes}, 0) catch "1", 0); + // no drafter checkpoint is read (or its memory held) when the server runs --no-drafts + if (!o.drafts) _ = setenv("BOLT_DRAFT_OFF", "1", 1); + const path = try gpa.allocSentinel(u8, dir.len, 0); + defer gpa.free(path); + @memcpy(path[0..dir.len], dir); + const model = api.open(path, @intCast(o.context)) orelse return error.BoltLoad; + const e = try gpa.create(Engine); + const nseq: u32 = @intCast(@max(1, api.seqs(model))); + e.* = .{ .gpa = gpa, .api = api, .model = model, .drafts = o.drafts and api.has_draft(model) != 0, .seqs = nseq, .stat_on = std.c.getenv("BOLT_STATS") != null, .vision = tower, .stream = stream }; + for (0..window) |i| costs[i] = .{ .width = @intCast(i + 1), .ms = api.cost(model, @intCast(i + 1)) }; + const propose_ms = api.cost(model, 0); + const propose_multi_ms = api.draft_multi_cost(model); + var rows_buf: [32]c_int = undefined; + var ms_buf: [32]f64 = undefined; + const n_shared: usize = @intCast(@max(0, api.shared_costs(model, &rows_buf, &ms_buf, 32))); + for (0..n_shared) |i| shared[i] = .{ .width = @intCast(rows_buf[i]), .ms = ms_buf[i] }; + std.log.info("bolt: windows ms {d:.1} {d:.1} {d:.1} {d:.1} {d:.1} {d:.1} {d:.1} {d:.1}, one proposal {d:.1} ({d} streams {d:.1}), drafts {s}, {d} sequence{s}", .{ costs[0].ms, costs[1].ms, costs[2].ms, costs[3].ms, costs[4].ms, costs[5].ms, costs[6].ms, costs[7].ms, propose_ms, nseq, propose_multi_ms, if (e.drafts) "on" else "off", nseq, if (nseq == 1) "" else "s" }); + for (shared[0..n_shared]) |c| std.log.info("bolt: shared round of {d} rows {d:.1} ms", .{ c.width, c.ms }); + return .{ + .backend = e.backend(), + .facts = .{ + .exact_width = window, + .gpu_tokens = false, + .mtp = e.drafts, + .speculate = e.drafts, + .speculate_early = false, + .draft_prior = &prior, + .drafts = window - 1, + .window_costs = &costs, + .shared_costs = shared[0..n_shared], + .mtp_step_ms = propose_ms / 3.0, + .draft_streams = nseq > 1, + .hidden_rows = nseq > 1, + .streams_exact = true, + .batch_rows = @min(max_batch, nseq * window), + .max_streams = nseq, + }, + .rows = window, + .stream_bytes = @intCast(api.seq_bytes(model)), + .free_streams = nseq, // their state is already resident + .ctx = e, + .deinit = shutdown, + .target = @ptrCast(@alignCast(e)), // the prompt cache's opaque handle: the engine itself + .vision = tower, + }; +} + +pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { + return switch (err) { + error.BoltLoad => "Bolt could not read this checkpoint", + error.BoltForward => "Bolt's forward failed", + error.VisionRefused => "--vision: this checkpoint has no image tower Bolt reads (the RADIO tower of an omni checkpoint); serve it without --vision", + error.NoVision => "this server reads no images: start it with --vision", + error.WindowTooWide => "Bolt verifies chains of up to eight rows a stream, sixteen streams at a time", + error.BadMask => "a grammar mask does not match the rows of its window", + else => null, + }; +} diff --git a/zig/src/families/nemotron/cuda.zig b/zig/src/families/nemotron/cuda.zig index 12ff32c8c0..5abe69982e 100644 --- a/zig/src/families/nemotron/cuda.zig +++ b/zig/src/families/nemotron/cuda.zig @@ -28,6 +28,7 @@ pub const reuse = @import("cuda_reuse.zig"); pub const slide = @import("cuda_slide.zig"); pub const worker = @import("cuda_worker.zig"); pub const omni = @import("cuda_native_omni.zig"); +pub const bolt = @import("bolt_native.zig"); pub const identity = @import("cuda_identity.zig"); test { diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 254a8799a9..b697b9ee2d 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -15,7 +15,7 @@ const cuda_disk = @import("cuda_disk.zig"); const vision_offer = @import("cuda_vision_offer.zig"); /// CUDA families provide metadata, open their lane backend and explain their refusals. -const registry = .{ nemotron.native, nemotron.omni, glm.native }; +const registry = .{ nemotron.native, nemotron.omni, nemotron.bolt, glm.native }; /// --record opens here (engines.zig refuses it for a backend without this). pub const records = true; @@ -335,6 +335,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O opts.precision = std.meta.stringToEnum(@typeInfo(@TypeOf(opts.precision)).optional.child, @tagName(p)).?; } if (@hasField(F.Options, "room")) opts.room = .{ .ctx = &probe, .read = Probe.read }; + if (@hasField(F.Options, "lanes")) opts.lanes = o.lanes; if (weights > before.room(held0)) { problem.* = try std.fmt.allocPrint(a, "the checkpoint's {d:.1} GiB of weights do not fit the {d:.1} GiB the CUDA memory budget grants ({d:.1} GiB free less a {d:.1} GiB reserve{s}); free device memory or adjust TENSORFOLD_MEMORY_RESERVE_GIB / TENSORFOLD_CUDA_MEMORY_LIMIT_GB", .{ toGib(weights), toGib(before.room(held0)), toGib(before.free), toGib(before.reserve), if (before.limit != null) ", under TENSORFOLD_CUDA_MEMORY_LIMIT_GB" else "" }); return null; diff --git a/zig/src/native/cuda_snaps.zig b/zig/src/native/cuda_snaps.zig index bbaedc2c57..5a294c645a 100644 --- a/zig/src/native/cuda_snaps.zig +++ b/zig/src/native/cuda_snaps.zig @@ -4,7 +4,7 @@ const api = @import("engine_api"); /// Any row, the head one token on (Nemotron's); `snap_grid`'s chunk starts and `warm_quiet`'s warm turns if declared. pub fn cacheRules(comptime F: type) api.prompt_cache.Rules { - var rules: api.prompt_cache.Rules = .{ .lookahead = 1, .min_prompt = 0 }; + var rules: api.prompt_cache.Rules = .{ .lookahead = if (@hasDecl(F, "snap_lookahead")) F.snap_lookahead else 1, .min_prompt = 0 }; if (@hasDecl(F, "warm_quiet")) rules.warm = .waits; if (@hasDecl(F, "snap_grid")) { rules.planned = true; diff --git a/zig/src/server/cli.zig b/zig/src/server/cli.zig index 460acd4599..a8317c98c0 100644 --- a/zig/src/server/cli.zig +++ b/zig/src/server/cli.zig @@ -315,7 +315,7 @@ fn apply(a: Allocator, out: *Args, name: []const u8, value: ?[]const u8, u: *Usa out.vision_video_tokens = @intCast(n); } else if (is(name, "--vision-image-tokens")) { const n = try int(u, a, name, v); - if (n < 16 or n > 16384) return fail(u, a, "argument --vision-image-tokens: expected 16 to 16384 tokens: '{s}'", .{v}); + if (n < 16 or n > 262144) return fail(u, a, "argument --vision-image-tokens: expected 16 to 262144 tokens: '{s}'", .{v}); out.vision_image_tokens = @intCast(n); out.vision_image_tokens_set = true; } diff --git a/zig/src/server/main.zig b/zig/src/server/main.zig index 4803950807..248c45483c 100644 --- a/zig/src/server/main.zig +++ b/zig/src/server/main.zig @@ -160,7 +160,12 @@ fn modelType(a: std.mem.Allocator, io: std.Io, dir: []const u8) []const u8 { const doc = std.json.parseFromSliceLeaky(std.json.Value, a, bytes, .{}) catch return "unknown"; if (doc != .object) return "unknown"; const t = doc.object.get("model_type") orelse return "unknown"; - return if (t == .string) t.string else "unknown"; + if (t != .string) return "unknown"; + // an omni wrapper around a puzzle text model is Bolt, a family of its own + if (std.mem.eql(u8, t.string, "nemotron_h_omni")) if (doc.object.get("llm_config")) |llm| if (llm == .object) { + if (llm.object.get("model_type")) |lt| if (lt == .string and std.mem.eql(u8, lt.string, "nemotron_h_puzzle")) return "nemotron_h_puzzle"; + }; + return t.string; } /// generation_config.json's sampling (``do_sample`` false is greedy), then the serve flags over it. diff --git a/zig/src/server/root.zig b/zig/src/server/root.zig index 381c864ad0..ecba18bfdf 100644 --- a/zig/src/server/root.zig +++ b/zig/src/server/root.zig @@ -46,6 +46,7 @@ test { _ = @import("image_key.zig"); _ = @import("media_fetch.zig"); _ = @import("videos.zig"); + _ = @import("vocab.zig"); _ = @import("stream_preflight_test.zig"); _ = @import("stream_dsml_test.zig"); _ = @import("logprobs_test.zig"); diff --git a/zig/src/server/server.zig b/zig/src/server/server.zig index 5406581774..b0d11b4eb9 100644 --- a/zig/src/server/server.zig +++ b/zig/src/server/server.zig @@ -23,6 +23,7 @@ const live = @import("live.zig"); const slide_graph = @import("slide_graph.zig"); const slide_lesson = @import("slide_lesson.zig"); const cortheon = @import("cortheon/cortheon.zig"); +const vocab_size = @import("vocab.zig"); const Value = json.Value; const Cx = errors.Cx; const Allocator = std.mem.Allocator; @@ -229,7 +230,7 @@ pub const Server = struct { return .{ .kind = kind, .text = text, .after = if (thinking) srv.text.tokenId(srv.markers.close) else null, .compiler = c }; } - /// The model's grammar compiler: tokenizer.json, config.json's vocab_size (text_config first), the eos ids. + /// The model's grammar compiler: tokenizer.json, config.json's vocab_size (the text model's), the eos ids. fn grammarCompiler(srv: *Server) ?*api.grammar.Compiler { srv.grammar_mutex.lockUncancelable(srv.io); defer srv.grammar_mutex.unlock(srv.io); @@ -244,10 +245,7 @@ pub const Server = struct { .ok => |v| v, .err => break :blk "config.json is not JSON", }; - var vocab: ?u32 = null; - for ([_]?json.Value{ cfg.get("text_config"), cfg }) |part| if (vocab == null) if (part) |p| if (p.get("vocab_size")) |v| if (v.int64()) |n| { - vocab = @intCast(n); - }; + const vocab = vocab_size.size(cfg); if (vocab == null) break :blk "structured output needs the checkpoint's config.json vocab_size"; var err: [512]u8 = undefined; const c = a.create(api.grammar.Compiler) catch break :blk "out of memory"; diff --git a/zig/src/server/vocab.zig b/zig/src/server/vocab.zig new file mode 100644 index 0000000000..d5db492d2f --- /dev/null +++ b/zig/src/server/vocab.zig @@ -0,0 +1,29 @@ +//! The vocabulary size a checkpoint declares, for the token masks of structured output. +const std = @import("std"); +const json = @import("json"); + +/// config.json's vocab_size: the text model's (text_config, or an omni checkpoint's llm_config) before the top level's. +pub fn size(cfg: json.Value) ?u32 { + for ([_]?json.Value{ cfg.get("text_config"), cfg.get("llm_config"), cfg }) |part| if (part) |p| if (p.get("vocab_size")) |v| if (v.int64()) |n| { + if (n > 0 and n <= std.math.maxInt(u32)) return @intCast(n); + }; + return null; +} + +fn parse(a: std.mem.Allocator, text: []const u8) !json.Value { + return switch (try json.parse(a, text)) { + .ok => |v| v, + .err => error.TestUnexpectedResult, + }; +} + +test "the vocabulary is the text model's: text_config, llm_config of an omni wrapper, else the top level" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + try std.testing.expectEqual(@as(?u32, 131072), size(try parse(a, "{\"model_type\": \"nemotron_h_omni\", \"llm_config\": {\"vocab_size\": 131072}}"))); + try std.testing.expectEqual(@as(?u32, 151936), size(try parse(a, "{\"text_config\": {\"vocab_size\": 151936}, \"vocab_size\": 1}"))); + try std.testing.expectEqual(@as(?u32, 32000), size(try parse(a, "{\"vocab_size\": 32000}"))); + try std.testing.expectEqual(@as(?u32, null), size(try parse(a, "{\"llm_config\": {\"hidden_size\": 4096}}"))); + try std.testing.expectEqual(@as(?u32, null), size(try parse(a, "{\"vocab_size\": 0}"))); +}