Nemotron 3.5 Bolt 74B-A9B on CUDA: puzzle layers, LatentMoE, ggml IQ2_XXS/Q8_0 weights, image input, host prompt states (on #853)

build.zig

@@ -12,6 +12,7 @@ const counters_build = @import("zig/build/metal_counters.zig");
 const exl3_build = @import("zig/build/exl3.zig");
 const exlr8_build = @import("zig/build/exlr8.zig");
 const golden_build = @import("zig/build/golden.zig");
+const vision_build = @import("zig/build/vision.zig");
 
 comptime {
     const required = std.mem.trim(u8, @embedFile(".zig-version"), "\r\n");
@@ -267,7 +268,7 @@ fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b
     counters_build.targets(b, target, optimize, metal, mods.sources, test_step);
 
     // The engines registry: the families this platform's binaries serve; the CLI reads it for models/info/pull.
-    const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = mods.lanes }} });
+    const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "lanes", .module = mods.lanes }, .{ .name = "qwen_image", .module = vision_build.module(b, target, optimize) } } });
     const engines = b.createModule(.{
         .root_source_file = b.path("zig/src/native/metal.zig"),
         .target = target,
@@ -572,7 +573,7 @@ pub fn distMetalServer(b: *std.Build, target: std.Build.ResolvedTarget, draft_id
     m.metal.addSystemFrameworkPath(b.graph.cwdRelativePath(b.pathJoin(&.{ sdk, "System/Library/Frameworks" })));
     m.metal.addLibraryPath(b.graph.cwdRelativePath(b.pathJoin(&.{ sdk, "usr/lib" })));
     m.metal.addSystemIncludePath(b.graph.cwdRelativePath(b.pathJoin(&.{ sdk, "usr/include" })));
-    const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = .fast, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = m.lanes }} });
+    const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = .fast, .link_libc = true, .imports = &.{ .{ .name = "lanes", .module = m.lanes }, .{ .name = "qwen_image", .module = vision_build.module(b, target, .ReleaseFast) } } });
     const engines = b.createModule(.{ .root_source_file = b.path("zig/src/native/metal.zig"), .target = target, .optimize = .fast, .link_libc = true, .imports = &.{ .{ .name = "metal", .module = m.metal }, .{ .name = "engine_api", .module = api }, .{ .name = "tensorfold", .module = m.engine } } });
     return nativeServer(b, target, .fast, m.metal, m.engine, api, engines, build_options, null);
 }

tools/zig/bolt_ggml.py

@@ -0,0 +1,87 @@
+"""Reference decode of Bolt's ggml tensors (IQ2_XXS, Q8_0), tables from the CUDA header: MODEL_DIR TENSOR [ROWS]."""
+
+import json
+import re
+import struct
+import sys
+from pathlib import Path
+
+import numpy as np
+
+HEADER = Path(__file__).resolve().parents[2] / "zig/kernels/cuda/ggml_iq2xxs_tables.cuh"
+IQ2_BYTES, IQ2_ROW = 66, 256
+Q8_BYTES, Q8_ROW = 34, 32
+
+
+def _table(text, name):
+    body = text.split(name, 1)[1].split("};", 1)[0]
+    return [int(v, 0) for v in re.findall(r"0x[0-9a-fA-F]+|\b\d+\b", body.split("=", 1)[1])]
+
+
+def tables():
+    text = HEADER.read_text()
+    signs = np.array(_table(text, "tf_ksigns_iq2xs[128]"), dtype=np.uint8)
+    grid = np.array(_table(text, "tf_iq2xxs_grid[256]"), dtype=np.uint64)
+    assert signs.shape == (128,) and grid.shape == (256,)
+    return signs, grid.view(np.uint8).reshape(256, 8)
+
+
+def iq2_xxs(blocks):
+    """[..., 66] uint8 -> [..., 256] float32, as ggml's dequantize_row_iq2_xxs."""
+    signs, grid = tables()
+    lead = blocks.shape[:-1]
+    b = blocks.reshape(-1, IQ2_BYTES)
+    d = b[:, :2].copy().view(np.float16).astype(np.float32)[:, 0]
+    q = b[:, 2:].copy().view(np.uint32).reshape(-1, 8, 2)  # 8 sub-blocks of 32 weights
+    idx = np.ascontiguousarray(q[:, :, 0]).view(np.uint8).reshape(-1, 8, 4)
+    aux = q[:, :, 1]
+    db = d[:, None] * (0.5 + (aux >> 28).astype(np.float32)) * 0.25
+    out = np.empty((b.shape[0], 8, 4, 8), dtype=np.float32)
+    bits = (1 << np.arange(8)).astype(np.uint8)
+    for l in range(4):
+        s = signs[(aux >> (7 * l)) & 127]
+        neg = (s[..., None] & bits) != 0
+        out[:, :, l, :] = grid[idx[:, :, l]] * np.where(neg, -1.0, 1.0) * db[..., None]
+    return out.reshape(*lead, IQ2_ROW)
+
+
+def q8_0(blocks):
+    """[..., 34] uint8 -> [..., 32] float32."""
+    lead = blocks.shape[:-1]
+    b = blocks.reshape(-1, Q8_BYTES)
+    d = b[:, :2].copy().view(np.float16).astype(np.float32)
+    return (b[:, 2:].view(np.int8).astype(np.float32) * d).reshape(*lead, Q8_ROW)
+
+
+def read(model, name):
+    """The raw U8 array and its ggml format (from the export's quantized_layers)."""
+    model = Path(model)
+    index = json.loads((model / "model.safetensors.index.json").read_text())["weight_map"]
+    path = model / index[name]
+    with open(path, "rb") as fh:
+        n = struct.unpack("<Q", fh.read(8))[0]
+        meta = json.loads(fh.read(n))[name]
+        lo, hi = meta["data_offsets"]
+        fh.seek(8 + n + lo)
+        raw = np.frombuffer(fh.read(hi - lo), dtype=np.uint8).reshape(meta["shape"])
+    layers = json.loads((model / "config.json").read_text())["quantization_config"]["quantized_layers"]
+    algo = layers.get(name.removesuffix(".weight"), {}).get("quant_algo")
+    return raw, algo
+
+
+def weight(model, name, rows=None):
+    raw, algo = read(model, name)
+    if rows is not None:
+        raw = raw[:rows]
+    if algo == "IQ2_XXS":
+        assert raw.shape[-1] == IQ2_BYTES, raw.shape
+        return iq2_xxs(raw).reshape(raw.shape[0], -1)
+    if algo == "Q8_0":
+        assert raw.shape[-1] == Q8_BYTES, raw.shape
+        return q8_0(raw).reshape(raw.shape[0], -1)
+    raise SystemExit(f"{name}: {algo} is not a ggml format")
+
+
+if __name__ == "__main__":
+    w = weight(sys.argv[1], sys.argv[2], int(sys.argv[3]) if len(sys.argv) > 3 else None)
+    print(w.shape, f"mean {w.mean():.6g} std {w.std():.6g} absmax {np.abs(w).max():.6g}")

tools/zig/bolt_ggml_fixture.py

@@ -0,0 +1,95 @@
+"""Fixture for `tf-cuda-test ggml`: real Bolt tensors, random rows, float64 references (MODEL_DIR OUT_DIR)."""
+
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+import bolt_ggml  # noqa: E402
+
+PRE = "language_model.backbone.layers."
+EXPERTS = [0, 1, 2, 3, 17, 64, 100, 128, 200, 238, 239]  # of layer 1; the stacked layer's experts 0 .. 10
+SLOTS = 9
+DECODE_ROWS = [1, 3]
+PROMPT_ROWS = [37, 300]
+DENSE = [  # (case, tensor, rows of the weight kept or None)
+    ("shared_up", PRE + "1.mixer.shared_experts.up_proj.weight", None),
+    ("shared_down", PRE + "1.mixer.shared_experts.down_proj.weight", None),
+    ("fc1", PRE + "1.mixer.fc1_latent_proj.weight", None),
+    ("q7", PRE + "7.mixer.q_proj.weight", None),
+    ("lm_head", "language_model.lm_head.weight", 2048),
+]
+DENSE_ROWS = [1, 5, 16, 100]
+
+
+def bf16(a):
+    """float32 -> bf16 bits (round to nearest even) and their float64 values."""
+    u = np.ascontiguousarray(a, dtype=np.float32).view(np.uint32).astype(np.uint64)
+    bits = ((u + 0x7FFF + ((u >> 16) & 1)) >> 16).astype(np.uint16)
+    return bits, (bits.astype(np.uint32) << 16).view(np.float32).astype(np.float64)
+
+
+def main(model, out):
+    out = Path(out)
+    out.mkdir(parents=True, exist_ok=True)
+    arrays, params = {}, {}
+    rng = np.random.default_rng(1234)
+
+    def put(name, a):
+        a = np.ascontiguousarray(a)
+        (out / f"{name}.bin").write_bytes(a.tobytes())
+        arrays[name] = {"file": f"{name}.bin", "dtype": str(a.dtype), "shape": list(a.shape)}
+
+    # grouped experts: up [W, 1024] and down [1024, W] of each, raw blocks back to back
+    ups, downs, wu, wd = [], [], [], []
+    for e in EXPERTS:
+        for kind, raws, ws in (("up", ups, wu), ("down", downs, wd)):
+            raw, algo = bolt_ggml.read(model, f"{PRE}1.mixer.experts.{e}.{kind}_proj.weight")
+            assert algo == "IQ2_XXS"
+            raws.append(raw)
+            ws.append(bolt_ggml.iq2_xxs(raw).reshape(raw.shape[0], -1).astype(np.float64))
+    width, dims = wu[0].shape
+    params.update(experts=len(EXPERTS), width=width, dims=dims, slots=SLOTS)
+    params["decode"] = ",".join(map(str, DECODE_ROWS))
+    params["prompt"] = ",".join(map(str, PROMPT_ROWS))
+    put("up_raw", np.stack(ups))
+    put("down_raw", np.stack(downs))
+    for rows in DECODE_ROWS + PROMPT_ROWS:
+        picks = np.stack([rng.choice(len(EXPERTS), SLOTS, replace=False) for _ in range(rows)]).astype(np.int32)
+        xb, x = bf16(rng.standard_normal((rows, dims)).astype(np.float32))
+        hb, h = bf16((rng.standard_normal((rows * SLOTS, width)) * 0.05).astype(np.float32))
+        flat = picks.reshape(-1)
+        up = np.stack([wu[e] @ x[p // SLOTS] for p, e in enumerate(flat)])
+        down = np.stack([wd[e] @ h[p] for p, e in enumerate(flat)])
+        put(f"picks{rows}", flat)
+        put(f"x{rows}", xb)
+        put(f"h{rows}", hb)
+        put(f"up{rows}", up)
+        put(f"down{rows}", down)
+
+    # dense projections
+    names = []
+    for case, tensor, keep in DENSE:
+        raw, algo = bolt_ggml.read(model, tensor)
+        if keep is not None:
+            raw = raw[:keep]
+        w = bolt_ggml.weight(model, tensor, keep).astype(np.float64)
+        n, k = w.shape
+        names.append(case)
+        params[f"{case}_n"], params[f"{case}_k"] = n, k
+        params[f"{case}_kind"] = algo
+        put(f"{case}_raw", raw)
+        for rows in DENSE_ROWS:
+            xb, x = bf16(rng.standard_normal((rows, k)).astype(np.float32))
+            put(f"{case}_x{rows}", xb)
+            put(f"{case}_y{rows}", x @ w.T)
+    params["dense"] = ",".join(names)
+    params["dense_rows"] = ",".join(map(str, DENSE_ROWS))
+    (out / "manifest.json").write_text(json.dumps({"params": params, "arrays": arrays}, indent=1))
+    print(f"wrote {len(arrays)} arrays to {out}")
+
+
+if __name__ == "__main__":
+    main(sys.argv[1], sys.argv[2])

tools/zig/bolt_image_fixture.py

@@ -0,0 +1,42 @@
+"""Fixtures for nemotron_image.zig: torch bicubic antialiased resizes and CLIP normalization of a synthetic image."""
+
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+from bolt_vision_reference import target_patches  # noqa: E402
+
+MEAN = [0.48145466, 0.4578275, 0.40821073]
+STD = [0.26862954, 0.26130258, 0.27577711]
+CFG = {"patch_size": 16, "min_num_patches": 1024, "max_num_patches": 13312, "downsample_ratio": 0.5,
+       "max_model_len": 16384}
+
+
+def image(w, h):
+    y, x, c = np.meshgrid(np.arange(h), np.arange(w), np.arange(3), indexing="ij")
+    return ((7 * x + 13 * y + 29 * c + (x * y) % 17) % 256).astype(np.uint8)
+
+
+def prep(rgb, ow, oh):
+    t = torch.from_numpy(rgb).permute(2, 0, 1)[None].float()
+    t = F.interpolate(t, size=(oh, ow), mode="bicubic", align_corners=False, antialias=True)
+    mean = torch.tensor(MEAN).view(1, 3, 1, 1)
+    std = torch.tensor(STD).view(1, 3, 1, 1)
+    return ((t / 255.0 - mean) / std)[0].numpy().astype("<f4")
+
+
+def main():
+    rgb = image(37, 23)
+    out = b"".join(prep(rgb, ow, oh).tobytes() for ow, oh in [(48, 32), (16, 16), (24, 8)])
+    Path(sys.argv[1]).write_bytes(out)
+    for w, h in [(2172, 724), (640, 480), (37, 23), (4000, 3000), (16, 16), (100, 5000), (1024, 1024), (512, 512),
+                 (3000, 100), (1920, 1080), (800, 600)]:
+        print(w, h, target_patches(w, h, CFG))
+
+
+if __name__ == "__main__":
+    main()